MB-35061: Ensure Consumer is removed from vbConns when stream is erased

DcpConsumer::{registerStream,removeStream} were introduced in
http://review.couchbase.org/#/c/103845/ . They are complementary
methods to add a stream to the Consumer::streams map and add the
Consumer to the DcpConnMap.vbConns (and vice versa), keeping the two
consistent - a consumer without a (non-dead) stream for a given vb
should not be in vbConns for that vb.

In some cases, Consumer::streams.erase() was called directly, leaving
the Consumer in the vbConns map, despite it no longer holding a
relevant stream.

Methods such as DcpConnMap::seqnoAckVBPassiveStream assume only a
single Consumer will be stored per vbucket, and use the first found.
This fails, as the "old" consumer does not have the associated stream.

Change-Id: Ifedc7a11d0439c2ec9365c277c75e5119bbc991b
Reviewed-on: http://review.couchbase.org/112784
Reviewed-by: Dave Rigby <[email protected]>
Tested-by: Build Bot <[email protected]>

Author: jameseh96

engines/ep/src/connmap.cc

@@ -190,3 +190,19 @@ void ConnMap::removeVBConnByVBId(const void* connCookie, Vbid vbid) {
     std::lock_guard<std::mutex> lh(vbConnLocks[lock_num]);
     removeVBConnByVBId_UNLOCKED(connCookie, vbid);
 }
+
+bool ConnMap::vbConnectionExists(ConnHandler* conn, Vbid vbid) {
+    size_t lock_num = vbid.get() % vbConnLockNum;
+    std::lock_guard<std::mutex> lh(vbConnLocks[lock_num]);
+    auto& connsForVb = vbConns[vbid.get()];
+
+    // Check whether the connhandler already exists in vbConns for the
+    // provided vbid
+    for (const auto& existingConn : connsForVb) {
+        if (conn == existingConn.lock().get()) {
+            return true;
+        }
+    }
+
+    return false;
+}

engines/ep/src/connmap.h

@@ -76,6 +76,13 @@ class ConnMap {
 
     void removeVBConnByVBId(const void* connCookie, Vbid vbid);
 
+    /**
+     * Checks (by pointer comparison) whether a ConnHandler is already
+     * present in vbConns.
+     *
+     */
+    bool vbConnectionExists(ConnHandler* conn, Vbid vbid);
+
     /**
      * Notifies the front-end synchronously on this thread that this paused
      * connection should be re-considered for work.

engines/ep/src/dcp/consumer.cc

@@ -291,7 +291,7 @@ ENGINE_ERROR_CODE DcpConsumer::addStream(uint32_t opaque,
                          vbucket);
             return ENGINE_KEY_EEXISTS;
         } else {
-            streams.erase(vbucket);
+            removeStream(vbucket);
         }
     }
 
@@ -1340,12 +1340,11 @@ void DcpConsumer::closeAllStreams() {
 
 void DcpConsumer::closeStreamDueToVbStateChange(Vbid vbucket,
                                                 vbucket_state_t state) {
-    auto it = streams.erase(vbucket);
-    if (it.second) {
+    auto stream = removeStream(vbucket);
+    if (stream) {
         logger->debug("({}) State changed to {}, closing passive stream!",
                       vbucket,
                       VBucket::toString(state));
-        auto& stream = it.first;
         uint32_t bytesCleared = stream->setDead(END_STREAM_STATE);
         flowControl.incrFreedBytes(bytesCleared);
         scheduleNotifyIfNecessary();
@@ -1743,10 +1742,15 @@ void DcpConsumer::setDisconnect() {
 void DcpConsumer::registerStream(std::shared_ptr<PassiveStream> stream) {
     auto vbid = stream->getVBucket();
     streams.insert({vbid, stream});
-    engine_.getDcpConnMap().addVBConnByVBId(shared_from_this(), vbid);
+    auto& connMap = engine_.getDcpConnMap();
+
+    Expects(!connMap.vbConnectionExists(this, vbid));
+
+    connMap.addVBConnByVBId(shared_from_this(), vbid);
 }
 
-void DcpConsumer::removeStream(Vbid vbid) {
-    streams.erase(vbid);
+std::shared_ptr<PassiveStream> DcpConsumer::removeStream(Vbid vbid) {
+    auto eraseResult = streams.erase(vbid).first;
     engine_.getDcpConnMap().removeVBConnByVBId(getCookie(), vbid);
+    return eraseResult;
 }

engines/ep/src/dcp/consumer.h

@@ -478,9 +478,12 @@ class DcpConsumer : public ConnHandler,
      * Remove a stream from this Consumer and remove the VB-to-Consumer
      * mapping from DcpConnMap.
      *
+     * Returns the removed stream ptr, or an empty shared_ptr if it
+     * was not found
+     *
      * @param vbid The stream to be removed
      */
-    void removeStream(Vbid vbid);
+    std::shared_ptr<PassiveStream> removeStream(Vbid vbid);
 
     /**
      * RAII helper class to update the flowControl object with the number of