From d23b68d615775b6cc11062e6eb20bc064ee2882e Mon Sep 17 00:00:00 2001
From: yunhanw-google <yunhanw@google.com>
Date: Thu, 28 Jul 2022 00:38:16 -0700
Subject: [PATCH] Add mNumMessagesToAllowBeforeDropping (#21321)

* Add mNumMessagesToAllowBeforeDropping

* Fix typo in comment.

Co-authored-by: Boris Zbarsky <bzbarsky@apple.com>
---
 src/transport/raw/tests/NetworkTestHelpers.h | 48 +++++++++++++-------
 1 file changed, 31 insertions(+), 17 deletions(-)

diff --git a/src/transport/raw/tests/NetworkTestHelpers.h b/src/transport/raw/tests/NetworkTestHelpers.h
index bbf3fcb936310e..cde07460556eb5 100644
--- a/src/transport/raw/tests/NetworkTestHelpers.h
+++ b/src/transport/raw/tests/NetworkTestHelpers.h
@@ -64,7 +64,8 @@ class LoopbackTransportDelegate
 public:
     virtual ~LoopbackTransportDelegate() {}
 
-    // Called by the loopback transport when it drops a message due to a nonzero mNumMessagesToDrop.
+    // Called by the loopback transport when it drops one of a configurable number of messages (mDroppedMessageCount) after a
+    // configurable allowed number of messages (mNumMessagesToAllowBeforeDropping)
     virtual void OnMessageDropped() {}
 };
 
@@ -102,19 +103,30 @@ class LoopbackTransport : public Transport::Base
     {
         ReturnErrorOnFailure(mMessageSendError);
         mSentMessageCount++;
-
-        if (mNumMessagesToDrop == 0)
+        bool dropMessage = false;
+        if (mNumMessagesToAllowBeforeDropping > 0)
         {
-            System::PacketBufferHandle receivedMessage = msgBuf.CloneData();
-            mPendingMessageQueue.push(PendingMessageItem(address, std::move(receivedMessage)));
-            mSystemLayer->ScheduleWork(OnMessageReceived, this);
+            --mNumMessagesToAllowBeforeDropping;
         }
-        else
+        else if (mNumMessagesToDrop > 0)
+        {
+            dropMessage = true;
+            --mNumMessagesToDrop;
+        }
+
+        if (dropMessage)
         {
-            mNumMessagesToDrop--;
             mDroppedMessageCount++;
             if (mDelegate != nullptr)
+            {
                 mDelegate->OnMessageDropped();
+            }
+        }
+        else
+        {
+            System::PacketBufferHandle receivedMessage = msgBuf.CloneData();
+            mPendingMessageQueue.push(PendingMessageItem(address, std::move(receivedMessage)));
+            mSystemLayer->ScheduleWork(OnMessageReceived, this);
         }
 
         return CHIP_NO_ERROR;
@@ -124,10 +136,11 @@ class LoopbackTransport : public Transport::Base
 
     void Reset()
     {
-        mNumMessagesToDrop   = 0;
-        mDroppedMessageCount = 0;
-        mSentMessageCount    = 0;
-        mMessageSendError    = CHIP_NO_ERROR;
+        mNumMessagesToDrop                = 0;
+        mDroppedMessageCount              = 0;
+        mSentMessageCount                 = 0;
+        mNumMessagesToAllowBeforeDropping = 0;
+        mMessageSendError                 = CHIP_NO_ERROR;
     }
 
     struct PendingMessageItem
@@ -143,11 +156,12 @@ class LoopbackTransport : public Transport::Base
     System::Layer * mSystemLayer = nullptr;
     std::queue<PendingMessageItem> mPendingMessageQueue;
     Transport::PeerAddress mTxAddress;
-    uint32_t mNumMessagesToDrop           = 0;
-    uint32_t mDroppedMessageCount         = 0;
-    uint32_t mSentMessageCount            = 0;
-    CHIP_ERROR mMessageSendError          = CHIP_NO_ERROR;
-    LoopbackTransportDelegate * mDelegate = nullptr;
+    uint32_t mNumMessagesToDrop                = 0;
+    uint32_t mDroppedMessageCount              = 0;
+    uint32_t mSentMessageCount                 = 0;
+    uint32_t mNumMessagesToAllowBeforeDropping = 0;
+    CHIP_ERROR mMessageSendError               = CHIP_NO_ERROR;
+    LoopbackTransportDelegate * mDelegate      = nullptr;
 };
 
 } // namespace Test