MB-31978: DCP close-stream support for stream-ID

Allow the client to encode (in flex-frame) a stream-ID
so that a single stream can be closed.

Change-Id: I22fde5e5696b9a951b598c241399d215bdf9588a
Reviewed-on: http://review.couchbase.org/102723
Reviewed-by: Ben Huddleston <[email protected]>
Tested-by: Build Bot <[email protected]>

Author: jimwwalker

daemon/mcbp_validators.cc

@@ -406,7 +406,7 @@ Status McbpValidator::verify_header(Cookie& cookie,
                 return false;
             }
         case cb::mcbp::request::FrameInfoId::DcpStreamId:
-            if (data.size() != sizeof(cb::mcbp::DcpStreamIdFrameInfo)) {
+            if (data.size() != sizeof(cb::mcbp::DcpStreamId)) {
                 status = Status::Einval;
                 cookie.setErrorContext("DcpStreamId invalid size:" +
                                        std::to_string(data.size()));

daemon/protocol/mcbp/dcp_close_stream_executor.cc

@@ -28,10 +28,21 @@ void dcp_close_stream_executor(Cookie& cookie) {
     auto& connection = cookie.getConnection();
     if (ret == ENGINE_SUCCESS) {
         const auto& header = cookie.getHeader().getRequest();
-        ret = dcpCloseStream(cookie,
-                             header.getOpaque(),
-                             header.getVBucket(),
-                             {/*@todo read a DcpStreamId from flex-frame*/});
+        cb::mcbp::DcpStreamId dcpStreamId; // Initialises to 'none'
+        header.parseFrameExtras([&dcpStreamId](
+                                        cb::mcbp::request::FrameInfoId id,
+                                        cb::const_byte_buffer data) -> bool {
+            if (id == cb::mcbp::request::FrameInfoId::DcpStreamId) {
+                // Data is the u16 stream-ID in network byte-order
+                dcpStreamId = cb::mcbp::DcpStreamId(
+                        ntohs(*reinterpret_cast<const uint16_t*>(data.data())));
+                return false;
+            }
+            return true;
+        });
+
+        ret = dcpCloseStream(
+                cookie, header.getOpaque(), header.getVBucket(), dcpStreamId);
     }
 
     ret = connection.remapErrorCode(ret);

engines/ep/src/dcp/active_stream.h

@@ -129,6 +129,10 @@ class ActiveStream : public Stream,
         return cursor;
     }
 
+    bool compareStreamId(cb::mcbp::DcpStreamId id) const override {
+        return id == sid;
+    }
+
 protected:
     /**
      * @param vb reference to the associated vbucket

engines/ep/src/dcp/producer.cc

@@ -631,6 +631,18 @@ ENGINE_ERROR_CODE DcpProducer::step(struct dcp_message_producers* producers) {
                     se->getVbucket(),
                     mapEndStreamStatus(getCookie(), se->getFlags()),
                     resp->getStreamId());
+
+            if (resp->getStreamId() && sendStreamEndOnClientStreamClose) {
+                if (!closeStreamInner(
+                             se->getVbucket(), resp->getStreamId(), true)
+                             .first) {
+                    throw std::logic_error(
+                            "DcpProducer::step(StreamEnd): no stream was found "
+                            "for " +
+                            se->getVbucket().to_string() + " " +
+                            resp->getStreamId().to_string());
+                }
+            }
             break;
         }
         case DcpResponse::Event::Commit: {
@@ -1053,6 +1065,42 @@ bool DcpProducer::handleResponse(const protocol_binary_response_header* resp) {
     return false;
 }
 
+std::pair<std::shared_ptr<Stream>, bool> DcpProducer::closeStreamInner(
+        Vbid vbucket, cb::mcbp::DcpStreamId sid, bool eraseFromMapIfFound) {
+    std::shared_ptr<Stream> stream;
+    bool vbFound = false;
+
+    // Obtain exclusive access to the streams map and see if the vbucket is
+    // mapped.
+    std::lock_guard<StreamsMap> guard(streams);
+    auto rv = streams.find(vbucket, guard);
+    if (rv.second) {
+        vbFound = true;
+        // Vbucket is mapped, get exclusive access to the StreamContainer
+        auto handle = rv.first->wlock();
+        // Try and locate a matching stream
+        for (; !handle.end(); handle.next()) {
+            if (handle.get()->compareStreamId(sid)) {
+                stream = handle.get();
+                break;
+            }
+        }
+
+        if (eraseFromMapIfFound && stream) {
+            // Need to tidy up the map, we first call erase on the handle,
+            // which will erase the current element from the container
+            handle.erase();
+
+            // If the container is now empty, remove it from the map, the
+            // shared_ptr (held by rv) will do the real destruction.
+            if (handle.empty()) {
+                streams.erase(vbucket, guard);
+            }
+        }
+    }
+    return {stream, vbFound};
+}
+
 ENGINE_ERROR_CODE DcpProducer::closeStream(uint32_t opaque,
                                            Vbid vbucket,
                                            cb::mcbp::DcpStreamId sid) {
@@ -1078,55 +1126,30 @@ ENGINE_ERROR_CODE DcpProducer::closeStream(uint32_t opaque,
     }
 
     /* We should not remove the stream from the streams map if we have to
-       send the "STREAM_END" response asynchronously to the consumer */
-    std::shared_ptr<Stream> stream;
-
-    {
-        // Obtain exclusive access to the streams map and see if the vbucket is
-        // mapped.
-        std::lock_guard<StreamsMap> guard(streams);
-        auto rv = streams.find(vbucket, guard);
-        if (rv.second) {
-            // Vbucket is mapped, get exclusive access to the StreamContainer
-            auto handle = rv.first->wlock();
-            // Try and locate a matching stream
-            for (; !handle.end(); handle.next()) {
-                if (handle.get()->compareStreamId(sid)) {
-                    stream = handle.get();
-                    break;
-                }
-            }
-
-            if (!sendStreamEndOnClientStreamClose) {
-                // Need to tidy up the map, we first call erase on the handle,
-                // which will erase the current element from the container
-                handle.erase();
-
-                // If the container is now empty, remove it from the map, the
-                // shared_ptr (held by rv) will do the real destruction.
-                if (handle.empty()) {
-                    streams.erase(vbucket, guard);
-                }
-            }
-        }
-    } // end streams lock scope
+       send the "STREAM_END" response asynchronously to the consumer, so
+       use the value of sendStreamEndOnClientStreamClose to determine if the
+       stream should be removed if found*/
+    auto rv = closeStreamInner(vbucket, sid, !sendStreamEndOnClientStreamClose);
 
     ENGINE_ERROR_CODE ret;
-    if (!stream) {
+    if (!rv.first) {
         logger->warn(
                 "({}) Cannot close stream because no "
-                "stream exists for this vbucket",
-                vbucket);
-        return ENGINE_KEY_ENOENT;
+                "stream exists for this vbucket {}",
+                vbucket,
+                sid);
+        return sid && rv.second ? ENGINE_DCP_STREAMID_INVALID
+                                : ENGINE_KEY_ENOENT;
     } else {
-        if (!stream->isActive()) {
+        if (!rv.first->isActive()) {
             logger->warn(
                     "({}) Cannot close stream because "
-                    "stream is already marked as dead",
-                    vbucket);
+                    "stream is already marked as dead {}",
+                    vbucket,
+                    sid);
             ret = ENGINE_KEY_ENOENT;
         } else {
-            stream->setDead(END_STREAM_CLOSED);
+            rv.first->setDead(END_STREAM_CLOSED);
             ret = ENGINE_SUCCESS;
         }
         if (!sendStreamEndOnClientStreamClose) {

engines/ep/src/dcp/producer.h

@@ -448,6 +448,29 @@ class DcpProducer : public ConnHandler,
                           cb::mcbp::DcpStreamId sid,
                           std::shared_ptr<Stream>& stream);
 
+    /**
+     * Attempt to locate a stream associated with the vbucket/stream-id and
+     * return it (this function is dedicated to the closeStream path)
+     * The function returns a pair because in the case the shared_ptr is null
+     * it could be because 1) the vbucket has no streams or 2) the vbucket
+     * has streams, but none that matched the given sid. If the vbucket does
+     * have streams, the pair.second will be true.
+     * The function is invoked in two places in the case that closeStream wants
+     * to send a CLOSE_STREAM message we leave the stream in the map... if a
+     * new stream is created, we would replace that stream... however with
+     * stream-ID enabled, it's possible a new stream is created with a new-ID
+     * leaking the dead stream, hence a second path now frees the dead stream
+     * by using this function and forcing erasure from the map.
+     *
+     * @param vbucket look for a stream associated with this vbucket
+     * @param sid and with a stream-ID matching sid
+     * @param eraseFromMapIfFound remove the shared_ptr from the streamsMap
+     * @return a pair the shared_ptr (can be null if not found) and a bool which
+     *  allows the caller to determine if the vbucket or sid caused not found
+     */
+    std::pair<std::shared_ptr<Stream>, bool> closeStreamInner(
+            Vbid vbucket, cb::mcbp::DcpStreamId sid, bool eraseFromMapIfFound);
+
     // stash response for retry if E2BIG was hit
     std::unique_ptr<DcpResponse> rejectResp;
 

engines/ep/src/dcp/stream.h

@@ -134,8 +134,8 @@ class Stream {
 
     void clear();
 
-    bool compareStreamId(cb::mcbp::DcpStreamId id) const {
-        return id == this->id;
+    virtual bool compareStreamId(cb::mcbp::DcpStreamId id) const {
+        return id == cb::mcbp::DcpStreamId(0);
     }
 
 protected:
@@ -196,8 +196,6 @@ class Stream {
 
     Cursor noCursor;
 
-    cb::mcbp::DcpStreamId id;
-
 private:
     /* readyQueueMemory tracks the memory occupied by elements
      * in the readyQ.  It is an atomic because otherwise

engines/ep/tests/mock/mock_dcp_producer.cc

@@ -78,4 +78,20 @@ std::shared_ptr<Stream> MockDcpProducer::findStream(Vbid vbid) {
         return handle.get();
     }
     return nullptr;
+}
+
+std::pair<std::shared_ptr<Stream>, bool> MockDcpProducer::findStream(
+        Vbid vbid, cb::mcbp::DcpStreamId sid) {
+    auto rv = streams.find(vbid);
+    if (rv.second) {
+        auto handle = rv.first->rlock();
+        // Try and locate a matching stream
+        for (; !handle.end(); handle.next()) {
+            if (handle.get()->compareStreamId(sid)) {
+                return {handle.get(), true};
+            }
+        }
+        return {nullptr, true};
+    }
+    return {nullptr, false};
 }

engines/ep/tests/mock/mock_dcp_producer.h

@@ -119,6 +119,14 @@ class MockDcpProducer : public DcpProducer {
      */
     std::shared_ptr<Stream> findStream(Vbid vbid);
 
+    /**
+     * Finds the stream for a given vbucket/sid
+     * @returns a pair where second indicates if the VB has no entries at all
+     *          first is the stream (or null if no stream)
+     */
+    std::pair<std::shared_ptr<Stream>, bool> findStream(
+            Vbid vbid, cb::mcbp::DcpStreamId sid);
+
     /**
      * Sets the backfill buffer size (max limit) to a particular value
      */
@@ -181,4 +189,8 @@ class MockDcpProducer : public DcpProducer {
     void enableMultipleStreamRequests() {
         multipleStreamRequests = MultipleStreamRequests::Yes;
     }
+
+    void enableStreamEndOnClientStreamClose() {
+        sendStreamEndOnClientStreamClose = true;
+    }
 };

engines/ep/tests/module_tests/collections/evp_store_collections_dcp_stream_test.cc

@@ -37,6 +37,8 @@ class CollectionsDcpStreamsTest : public CollectionsDcpTest {
         CollectionsDcpTest::consumer = std::make_shared<MockDcpConsumer>(
                 *engine, cookieC, "test_consumer");
     }
+
+    void close_stream_by_id_test(bool enableStreamEnd, bool manyStreams);
 };
 
 TEST_F(CollectionsDcpStreamsTest, request_validation) {
@@ -104,6 +106,71 @@ TEST_F(CollectionsDcpStreamsTest, close_stream_validation2) {
               producer->closeStream(0, vbid, cb::mcbp::DcpStreamId(99)));
 }
 
+// Core test method for close with a stream ID.
+// Allows a number of cases to be covered where we have many streams mapped to
+// the vbid and stream_end message enabled/disabled.
+void CollectionsDcpStreamsTest::close_stream_by_id_test(bool enableStreamEnd,
+                                                        bool manyStreams) {
+    CollectionsManifest cm;
+    cm.add(CollectionEntry::fruit);
+    auto vb = store->getVBucket(vbid);
+    vb->updateFromManifest({cm});
+    if (enableStreamEnd) {
+        producer->enableStreamEndOnClientStreamClose();
+    }
+    producer->enableMultipleStreamRequests();
+    createDcpStream({{R"({"collections":["9"], "sid":99})"}}, vbid);
+
+    if (manyStreams) {
+        createDcpStream({{R"({"collections":["9", "0"], "sid":101})"}}, vbid);
+        createDcpStream({{R"({"collections":["0"], "sid":555})"}}, vbid);
+    }
+
+    EXPECT_EQ(ENGINE_SUCCESS,
+              producer->closeStream(0, vbid, cb::mcbp::DcpStreamId(99)));
+    {
+        auto rv = producer->findStream(vbid, cb::mcbp::DcpStreamId(99));
+        // enableStreamEnd:true we must still find the stream in the producer
+        EXPECT_EQ(enableStreamEnd, rv.first != nullptr);
+        // manyStreams:true we always expect to find streams
+        // manyStreams:false only expect streams if enableStreamEnd:true
+        EXPECT_EQ(enableStreamEnd | manyStreams, rv.second);
+    }
+
+    if (enableStreamEnd) {
+        EXPECT_EQ(
+                ENGINE_SUCCESS,
+                producer->stepAndExpect(producers.get(),
+                                        cb::mcbp::ClientOpcode::DcpStreamEnd));
+        EXPECT_EQ(cb::mcbp::DcpStreamId(99), producers->last_stream_id);
+    }
+
+    {
+        auto rv = producer->findStream(vbid, cb::mcbp::DcpStreamId(99));
+        // Always expect at this point to not find stream for sid 99
+        EXPECT_FALSE(rv.first);
+        // However the vb may still have streams for manyStreams test
+        EXPECT_EQ(manyStreams, rv.second);
+    }
+}
+
+TEST_F(CollectionsDcpStreamsTest, close_stream_by_id_with_end_stream) {
+    close_stream_by_id_test(true, false);
+}
+
+TEST_F(CollectionsDcpStreamsTest, close_stream_by_id) {
+    close_stream_by_id_test(false, false);
+}
+
+TEST_F(CollectionsDcpStreamsTest,
+       close_stream_by_id_with_end_stream_and_many_streams) {
+    close_stream_by_id_test(true, true);
+}
+
+TEST_F(CollectionsDcpStreamsTest, close_stream_by_id_and_many_streams) {
+    close_stream_by_id_test(false, true);
+}
+
 TEST_F(CollectionsDcpStreamsTest, two_streams) {
     CollectionsManifest cm;
     cm.add(CollectionEntry::fruit);

include/mcbp/protocol/request.h

@@ -83,7 +83,9 @@ class Request {
 
     ServerOpcode getServerOpcode() const {
         if (is_client_magic(getMagic())) {
-            throw std::logic_error("getServerOpcode: magic != server request");
+            throw std::logic_error(
+                    "getServerOpcode: magic != server request: " +
+                    std::to_string(uint8_t(getMagic())));
         }
         return ServerOpcode(opcode);
     }