UT[mqbs]: proxy VS removal must not release refs of unpushed Apps

src/groups/mqb/mqbs/mqbs_inmemorystorage.t.cpp

@@ -44,6 +44,7 @@
 #include <ball_severity.h>
 #include <bdlbb_blob.h>
 #include <bdlbb_blobutil.h>
+#include <bdlbb_pooledblobbufferfactory.h>
 #include <bsl_algorithm.h>
 #include <bsl_ostream.h>
 #include <bsl_string.h>
@@ -87,6 +88,7 @@
 // - capacityMeter_limitMessages
 //   capacityMeter_limitBytes
 // - garbageCollect
+// - proxyRemoveVirtualStorage_refCountOfUnpushedApp
 // - addQueueOpRecordHandle
 //-----------------------------------------------------------------------------
 
@@ -1434,7 +1436,75 @@
     BMQTST_ASSERT_EQ(storage.numBytes(mqbu::StorageKey::k_NULL_KEY), 0);
 }
 
+BMQTST_TEST(proxyRemoveVirtualStorage_refCountOfUnpushedApp)
+// ------------------------------------------------------------------------
+// Proxy has 3 Apps.  A PUSH arrives for app1 and app2 only (refCount 2),
+// as done by 'RelayQueueEngine::storePushIfProxy'.  Removing app3's
+// virtual storage must not release a reference app3 never held.
+//
+// Every App slot of a new 'DataStreamMessage' starts as 'e_PUT', which
+// counts as pending, so app3's untouched slot looks like it holds a ref.
+// ------------------------------------------------------------------------
+{
+    bmqtst::TestHelper::printTestName(
+        "proxyRemoveVirtualStorage_refCountOfUnpushedApp");
+
+    bmqu::MemOutStream errDescription(bmqtst::TestHelperUtil::allocator());
+
+    Tester tester(k_PROXY_PARTITION_ID, bmqtst::TestHelperUtil::allocator());
+    tester.configure(k_INT64_MAX, k_INT64_MAX);
+
+    mqbs::ReplicatedStorage& storage = tester.storage();
+
+    storage.addVirtualStorage(errDescription, k_APP_ID1, k_APP_KEY1);
+    storage.addVirtualStorage(errDescription, k_APP_ID2, k_APP_KEY2);
+    storage.addVirtualStorage(errDescription, k_APP_ID3, k_APP_KEY3);
+
+    bdlbb::PooledBlobBufferFactory bufferFactory(
+        1024,
+        bmqtst::TestHelperUtil::allocator());
+    const bmqt::MessageGUID            guid = generateRandomGUID();
+    const int                          data = 0;
+    const bsl::shared_ptr<bdlbb::Blob> appData(
+        new (*bmqtst::TestHelperUtil::allocator())
+            bdlbb::Blob(&bufferFactory, bmqtst::TestHelperUtil::allocator()),
+        bmqtst::TestHelperUtil::allocator());
+    bdlbb::BlobUtil::append(appData.get(),
+                            reinterpret_cast<const char*>(&data),
+                            static_cast<int>(sizeof(int)));
+
+    mqbi::StorageMessageAttributes attributes(
+        0,
+        2,  // refCount: app1 and app2
+        static_cast<unsigned int>(appData->length()),
+        bmqp::MessagePropertiesInfo::makeNoSchema(),
+        bmqt::CompressionAlgorithmType::e_NONE);
+
+    mqbi::DataStreamMessage* dataStreamMessage = 0;
+    BMQTST_ASSERT_EQ(storage.put(&attributes,
+                                 guid,
+                                 appData,
+                                 bsl::shared_ptr<bdlbb::Blob>(),
+                                 &dataStreamMessage),
+                     mqbi::StorageResult::e_SUCCESS);
+    BMQTST_ASSERT(dataStreamMessage);
+
+    // Ordinals follow 'addVirtualStorage' order: app1 is 0, app2 is 1.
+    // app3's slot (ordinal 2) stays 'e_PUT'.
+    dataStreamMessage->app(0).setPushState();
+    dataStreamMessage->app(1).setPushState();
+
+    // app3 never received the PUSH; its last reader goes away.
+    BMQTST_ASSERT(storage.removeVirtualStorage(k_APP_KEY3, false));
+
+    // app2 still holds the message, so app1's CONFIRM must not reach zero.
+    BMQTST_ASSERT_EQ(storage.confirm(guid, k_APP_KEY1, 0),
+                     mqbi::StorageResult::e_NON_ZERO_REFERENCES);
+    BMQTST_ASSERT_EQ(storage.confirm(guid, k_APP_KEY2, 0),
+                     mqbi::StorageResult::e_ZERO_REFERENCES);
+}
+
 BMQTST_TEST_F(Test, addQueueOpRecordHandle)
 {
     // CONSTANTS
     const bsls::Types::Uint64    k_SEQUENCE_NUM     = 1024;

src/groups/mqb/mqbs/mqbs_virtualstoragecatalog.cpp

@@ -472,8 +472,15 @@ VirtualStorageCatalog::purgeImpl(VirtualStorage*     vs,
         // Must call 'seek' before 'purge'
         setup(dataStreamMessage);
 
-        if (vs->remove(dataStreamMessage, replacingOrdinal)) {
-            // The 'data' was not already removed or confirmed.
+        // First, check whether this App is counted in the message's refCount.
+        // Then, call 'remove' for every message (even if this App does not
+        // hold a reference), which overwrites this App's state.
+        const bool holdsRef = holdsReference(
+            appMessageView(*dataStreamMessage, vs->ordinal()));
+
+        vs->remove(dataStreamMessage, replacingOrdinal);
+
+        if (holdsRef) {
             result = d_storage_p->releaseRef(msgGUID, asPrimary);
         }
 

src/groups/mqb/mqbs/mqbs_virtualstoragecatalog.h

@@ -390,6 +390,12 @@ class VirtualStorageCatalog BSLS_KEYWORD_FINAL {
     const mqbi::AppMessage&
     appMessageView(const mqbi::DataStreamMessage& dataStreamMessage,
                    unsigned int                   ordinal) const;
+
+    /// Return `true` if the specified `appMessage` accounts for a reference
+    /// in the refCount of its message, and `false` otherwise.  A proxy
+    /// counts only the Apps a message was pushed to, while other nodes count
+    /// every App with a pending state.
+    bool holdsReference(const mqbi::AppMessage& appMessage) const;
 };
 
 // ============================================================================
@@ -469,6 +475,12 @@ inline mqbi::Queue* VirtualStorageCatalog::queue() const
     return d_isProxy ? 0 : d_queue_p;
 }
 
+inline bool
+VirtualStorageCatalog::holdsReference(const mqbi::AppMessage& appMessage) const
+{
+    return d_isProxy ? appMessage.isPushing() : appMessage.isPending();
+}
+
 }  // close package namespace
 }  // close enterprise namespace