MB-47318: [BP] Push UpdateFlowControl through to PassiveStream

Pass the UpdateFlowControl object which owns the correct buffer
ack value. If the PassiveStream has to buffer the DcpResponse
for later processing it is in charge of releasing the ack bytes
and will in the next patch save the value.

Change-Id: Ifcc3aade4d53b1780fee3000b11b3f8ef14c6266
Reviewed-on: https://review.couchbase.org/c/kv_engine/+/166227
Well-Formed: Restriction Checker
Tested-by: Build Bot <[email protected]>
Reviewed-by: Dave Rigby <[email protected]>

Author: jimwwalker

engines/ep/src/dcp/consumer.cc

@@ -561,7 +561,8 @@ ENGINE_ERROR_CODE DcpConsumer::deletion(uint32_t opaque,
                                         cb::const_byte_buffer meta,
                                         uint32_t deleteTime,
                                         IncludeDeleteTime includeDeleteTime,
-                                        DeleteSource deletionCause) {
+                                        DeleteSource deletionCause,
+                                        UpdateFlowControl& ufc) {
     lastMessageTime = ep_current_time();
 
     if (doDisconnect()) {
@@ -646,15 +647,16 @@ ENGINE_ERROR_CODE DcpConsumer::deletion(uint32_t opaque,
 
     try {
         err = stream->messageReceived(std::make_unique<MutationConsumerMessage>(
-                item,
-                opaque,
-                IncludeValue::Yes,
-                IncludeXattrs::Yes,
-                includeDeleteTime,
-                IncludeDeletedUserXattrs::Yes,
-                key.getEncoding(),
-                emd.release(),
-                cb::mcbp::DcpStreamId{}));
+                                              item,
+                                              opaque,
+                                              IncludeValue::Yes,
+                                              IncludeXattrs::Yes,
+                                              includeDeleteTime,
+                                              IncludeDeletedUserXattrs::Yes,
+                                              key.getEncoding(),
+                                              emd.release(),
+                                              cb::mcbp::DcpStreamId{}),
+                                      ufc);
     } catch (const std::bad_alloc&) {
         err = ENGINE_ENOMEM;
     }
@@ -749,12 +751,12 @@ ENGINE_ERROR_CODE DcpConsumer::toMainDeletion(DeleteType origin,
                         meta,
                         deleteTime,
                         includeDeleteTime,
-                        deleteSource);
+                        deleteSource,
+                        ufc);
 
     // TMPFAIL means the stream has buffered the message for later processing
     // so skip flowControl, success or any other error, we still need to ack
     if (err == ENGINE_TMPFAIL) {
-        ufc.release();
         // Mask the TMPFAIL
         return ENGINE_SUCCESS;
     }
@@ -1818,15 +1820,14 @@ ENGINE_ERROR_CODE DcpConsumer::lookupStreamAndDispatchMessage(
     // Pass the message to the associated stream.
     ENGINE_ERROR_CODE err;
     try {
-        err = stream->messageReceived(std::move(msg));
+        err = stream->messageReceived(std::move(msg), ufc);
     } catch (const std::bad_alloc&) {
         return ENGINE_ENOMEM;
     }
 
     // The item was buffered and will be processed later
     if (err == ENGINE_TMPFAIL) {
         notifyVbucketReady(vbucket);
-        ufc.release();
         return ENGINE_SUCCESS;
     }
 

engines/ep/src/dcp/consumer.h

@@ -36,6 +36,7 @@
 class DcpResponse;
 class PassiveStream;
 class StreamEndResponse;
+class UpdateFlowControl;
 
 /**
  * A DCP Consumer object represents a DCP connection which receives streams
@@ -461,7 +462,8 @@ class DcpConsumer : public ConnHandler,
                                cb::const_byte_buffer meta,
                                uint32_t deleteTime,
                                IncludeDeleteTime includeDeleteTime,
-                               DeleteSource deletionCause);
+                               DeleteSource deletionCause,
+                               UpdateFlowControl& ufc);
 
     /**
      * Helper function for mutation() and prepare() messages as they are handled
@@ -515,39 +517,6 @@ class DcpConsumer : public ConnHandler,
      */
     std::shared_ptr<PassiveStream> removeStream(Vbid vbid);
 
-    /**
-     * RAII helper class to update the flowControl object with the number of
-     * bytes to free and trigger the consumer notify
-     */
-    class UpdateFlowControl {
-    public:
-        UpdateFlowControl(DcpConsumer& consumer, uint32_t bytes)
-            : consumer(consumer), bytes(bytes) {
-            if (bytes == 0) {
-                throw std::invalid_argument("UpdateFlowControl given 0 bytes");
-            }
-        }
-
-        ~UpdateFlowControl() {
-            if (bytes) {
-                consumer.flowControl.incrFreedBytes(bytes);
-                consumer.scheduleNotifyIfNecessary();
-            }
-        }
-
-        /**
-         * If the user no longer wants this instance to perform the update
-         * calling release() means this instance will skip the update.
-         */
-        void release() {
-            bytes = 0;
-        }
-
-    private:
-        DcpConsumer& consumer;
-        uint32_t bytes;
-    };
-
     /**
      * Helper method to lookup the correct stream for the given
      * vbid / opaque pair, and then dispatch the message to that stream.
@@ -659,6 +628,8 @@ class DcpConsumer : public ConnHandler,
 
     bool alwaysBufferOperations{false};
 
+    friend UpdateFlowControl;
+
     static const std::string noopCtrlMsg;
     static const std::string noopIntervalCtrlMsg;
     static const std::string connBufferCtrlMsg;
@@ -671,6 +642,39 @@ class DcpConsumer : public ConnHandler,
     static const std::string enableOpcodeExpiryCtrlMsg;
 };
 
+/**
+ * RAII helper class to update the flowControl object with the number of
+ * bytes to free and trigger the consumer notify
+ */
+class UpdateFlowControl {
+public:
+    UpdateFlowControl(DcpConsumer& consumer, uint32_t bytes)
+        : consumer(consumer), bytes(bytes) {
+        if (bytes == 0) {
+            throw std::invalid_argument("UpdateFlowControl given 0 bytes");
+        }
+    }
+
+    ~UpdateFlowControl() {
+        if (bytes) {
+            consumer.flowControl.incrFreedBytes(bytes);
+            consumer.scheduleNotifyIfNecessary();
+        }
+    }
+
+    /**
+     * Calling release means that this object will not update the FlowControl
+     * instance when destructed.
+     */
+    void release() {
+        bytes = 0;
+    }
+
+private:
+    DcpConsumer& consumer;
+    uint32_t bytes{0};
+};
+
 /*
  * Task that orchestrates rollback on Consumer,
  * runs in background.

engines/ep/src/dcp/passive_stream.cc

@@ -271,7 +271,7 @@ void PassiveStream::reconnectStream(VBucketPtr& vb,
 }
 
 ENGINE_ERROR_CODE PassiveStream::messageReceived(
-        std::unique_ptr<DcpResponse> dcpResponse) {
+        std::unique_ptr<DcpResponse> dcpResponse, UpdateFlowControl& ufc) {
     if (!dcpResponse) {
         return ENGINE_EINVAL;
     }
@@ -406,6 +406,7 @@ ENGINE_ERROR_CODE PassiveStream::messageReceived(
 
     // Only buffer if the stream is not dead
     if (isActive()) {
+        ufc.release(); // @todo save the value along side the buffered response
         buffer.push(std::move(dcpResponse));
     }
     return ENGINE_TMPFAIL;
@@ -1303,7 +1304,7 @@ bool PassiveStream::Buffer::empty() const {
     return messages.empty();
 }
 
-void PassiveStream::Buffer::push(std::unique_ptr<DcpResponse> message) {
+void PassiveStream::Buffer::push(PassiveStream::Buffer::BufferType message) {
     std::lock_guard<std::mutex> lg(bufMutex);
     bytes += message->getMessageSize();
     messages.push_back(std::move(message));
@@ -1318,7 +1319,7 @@ void PassiveStream::Buffer::pop_front(std::unique_lock<std::mutex>& lh,
     bytes -= bytesPopped;
 }
 
-std::unique_ptr<DcpResponse>& PassiveStream::Buffer::front(
+PassiveStream::Buffer::BufferType& PassiveStream::Buffer::front(
         std::unique_lock<std::mutex>& lh) {
     return messages.front();
 }

engines/ep/src/dcp/passive_stream.h

@@ -34,6 +34,7 @@ class DropScopeEvent;
 class EventuallyPersistentEngine;
 class MutationConsumerMessage;
 class SystemEventMessage;
+class UpdateFlowControl;
 
 class PassiveStream : public Stream {
 public:
@@ -94,11 +95,14 @@ class PassiveStream : public Stream {
     /*
      * Calls the appropriate function to process the message.
      *
-     * @params response The dcp message that needs to be processed.
+     * @param response The dcp message that needs to be processed.
+     * @param ackSize the value to use when DCP acking - this may differ from
+     *        DcpResponse::getMessageSize if for example an Item value was
+     *        decompressed
      * @returns the error code from processing the message.
      */
-    virtual ENGINE_ERROR_CODE messageReceived(
-            std::unique_ptr<DcpResponse> response);
+    ENGINE_ERROR_CODE messageReceived(
+            std::unique_ptr<DcpResponse> response, UpdateFlowControl& ackSize);
 
     void addStats(const AddStatFn& add_stat, const void* c) override;
 
@@ -306,7 +310,9 @@ class PassiveStream : public Stream {
 
         bool empty() const;
 
-        void push(std::unique_ptr<DcpResponse> message);
+        using BufferType = std::unique_ptr<DcpResponse>;
+
+        void push(BufferType message);
 
         /*
          * Caller must of locked bufMutex and pass as lh (not asserted)
@@ -317,13 +323,13 @@ class PassiveStream : public Stream {
          * Return a reference to the item at the front.
          * The user must pass a lock to bufMutex.
          */
-        std::unique_ptr<DcpResponse>& front(std::unique_lock<std::mutex>& lh);
+        BufferType& front(std::unique_lock<std::mutex>& lh);
 
-        size_t bytes;
+        size_t bytes{0};
         /* Lock ordering w.r.t to streamMutex:
            First acquire bufMutex and then streamMutex */
         mutable std::mutex bufMutex;
-        std::deque<std::unique_ptr<DcpResponse> > messages;
+        std::deque<BufferType> messages;
     } buffer;
 
     /*

engines/ep/tests/mock/mock_stream.cc

@@ -17,6 +17,7 @@
 
 #include "mock_stream.h"
 #include "checkpoint_manager.h"
+#include "dcp/consumer.h"
 #include "dcp/response.h"
 #include "mock_dcp_producer.h"
 #include "vbucket.h"
@@ -91,7 +92,14 @@ void MockActiveStream::consumeAllBackfillItems() {
 
 ENGINE_ERROR_CODE MockPassiveStream::messageReceived(
         std::unique_ptr<DcpResponse> dcpResponse) {
-    return PassiveStream::messageReceived(std::move(dcpResponse));
+    auto consumer = consumerPtr.lock();
+    if (!consumer) {
+        throw std::logic_error(
+                "MockPassiveStream::messageReceived cannot proceed without the "
+                "DcpConsumer");
+    }
+    UpdateFlowControl ufc(*consumer, dcpResponse->getMessageSize());
+    return PassiveStream::messageReceived(std::move(dcpResponse), ufc);
 }
 
 std::unique_ptr<DcpResponse> MockPassiveStream::public_popFromReadyQ() {

engines/ep/tests/mock/mock_stream.h

@@ -288,8 +288,7 @@ class MockPassiveStream : public PassiveStream {
         transitionState(StreamState::Dead);
     }
 
-    ENGINE_ERROR_CODE messageReceived(
-            std::unique_ptr<DcpResponse> dcpResponse) override;
+    ENGINE_ERROR_CODE messageReceived(std::unique_ptr<DcpResponse> dcpResponse);
 
     void processMarker(SnapshotMarker* marker) override {
         PassiveStream::processMarker(marker);

engines/ep/tests/module_tests/dcp_reflection_test.cc

@@ -152,7 +152,7 @@ class DCPLoopbackStreamTest : public SingleThreadedKVBucketTest {
 
         void transferResponseMessage();
 
-        std::pair<ActiveStream*, PassiveStream*> getStreams();
+        std::pair<ActiveStream*, MockPassiveStream*> getStreams();
 
         Vbid vbid;
         EventuallyPersistentEngine* producerNode;
@@ -400,11 +400,12 @@ DCPLoopbackStreamTest::DcpRoute::getNextProducerMsg(ActiveStream* stream) {
     return producerMsg;
 }
 
-std::pair<ActiveStream*, PassiveStream*>
+std::pair<ActiveStream*, MockPassiveStream*>
 DCPLoopbackStreamTest::DcpRoute::getStreams() {
     auto* pStream =
             dynamic_cast<ActiveStream*>(producer->findStream(vbid).get());
-    auto* cStream = consumer->getVbucketStream(vbid).get();
+    auto* cStream = dynamic_cast<MockPassiveStream*>(
+            consumer->getVbucketStream(vbid).get());
     EXPECT_TRUE(pStream);
     EXPECT_TRUE(cStream);
     return {pStream, cStream};