Defer subacks to lazy subscriptions until the remote side acknowledged

When a client places subscriptions that match lazy subscriptions that
need to be expanded, we need to wait with sending the SUBACK until the
other end(s) have sent their SUBACKS. This is so that when clients
trigger subsequent actions on it, the subscription will actually exist.

Author: halfgaar

FlashMQTests/tst_maintests.cpp

@@ -979,20 +979,20 @@ void MainTests::testSavingSessions()
         const std::string topic1 = "one/two/three";
         std::vector<std::string> subtopics;
         subtopics = splitTopic(topic1);
-        store->addSubscription(c1->getSession(), subtopics, 0, true, false, "", 0);
+        store->addSubscription(c1->getSession(), 0, subtopics, 0, true, false, "", 0);
 
         const std::string topic2 = "four/five/six";
         subtopics = splitTopic(topic2);
-        store->addSubscription(c2->getSession(), subtopics, 0, false, true, "", 0);
-        store->addSubscription(c1->getSession(), subtopics, 0, false, false, "", 94612);
+        store->addSubscription(c2->getSession(), 0, subtopics, 0, false, true, "", 0);
+        store->addSubscription(c1->getSession(), 0, subtopics, 0, false, false, "", 94612);
 
         const std::string topic3 = "";
         subtopics = splitTopic(topic3);
-        store->addSubscription(c2->getSession(), subtopics, 0, false, false, "", 0);
+        store->addSubscription(c2->getSession(), 0, subtopics, 0, false, false, "", 0);
 
         const std::string topic4 = "#";
         subtopics = splitTopic(topic4);
-        store->addSubscription(c2->getSession(), subtopics, 0, false, false, "", 0);
+        store->addSubscription(c2->getSession(), 0, subtopics, 0, false, false, "", 0);
 
         Publish publish("a/b/c", "Hello Barry", 1);
         publish.client_id = "ClientIdFromFakePublisher";
@@ -3083,7 +3083,7 @@ void MainTests::testTopicMatchingInSubscriptionTreeHelper(const std::string &sub
     client->setClientProperties(ProtocolVersion::Mqtt5, "mytestclient", {}, "myusername", true, 60);
     store.registerClientAndKickExistingOne(client);
 
-    store.addSubscription(client->getSession(), subscribe_subtopics, 0, false, false, "", 0);
+    store.addSubscription(client->getSession(), 0, subscribe_subtopics, 0, false, false, "", 0);
 
     std::vector<ReceivingSubscriber> receivers;
     store.publishRecursively(publish_subtopics.begin(), publish_subtopics.end(), store.root.get(), receivers, "fakeclientid", {});

client.cpp

@@ -699,6 +699,43 @@ void Client::clearRegistrationData()
     this->registrationData.reset();
 }
 
+void Client::stageOrSendSubAck(const std::shared_ptr<Client> &this_client, SubAckAction &&action, const size_t expandedCount)
+{
+    assert(this == this_client.get());
+
+    if (expandedCount == 0)
+    {
+        sendSubAck(action);
+        return;
+    }
+
+    const uint16_t packid = action.mPacketId;
+
+    auto result = stagedSubAcks.insert({packid, std::move(action)});
+    if (!std::get<bool>(result))
+    {
+        logger->log(LOG_WARNING)
+            << "stageOrSendSubAck: staged SUBACK with id " << action.mPacketId << " is already exists (id reused). This is probabaly a bug in the client. "
+            << "Replying directly instead of staging.";
+        sendSubAck(action);
+        return;
+    }
+}
+
+void Client::sendStagedSuback(const uint16_t packetId)
+{
+    assert(this->threadData.lock()->thread_id == pthread_self());
+
+    auto pos = stagedSubAcks.find(packetId);
+
+    if (pos == stagedSubAcks.end())
+        return;
+
+    const SubAckAction x(std::move(pos->second));
+    stagedSubAcks.erase(pos);
+    sendSubAck(x);
+}
+
 /**
  * @brief Client::stageConnack saves the success connack for later use.
  * @param c
@@ -936,6 +973,28 @@ void Client::setAddr(const std::string &address)
     addr.setAddress(address);
 }
 
+void Client::sendSubAck(const SubAckAction &action)
+{
+    auto store = globals->subscriptionStore;
+
+    /*
+     * Writing the suback here, as opposed to before adding subscriptions way back in the calling context, means the client may
+     * have already received packets that match the subscription. That should not be a problem, because the spec says:
+     *
+     * "The Server is permitted to start sending PUBLISH packets matching the Subscription before the Server sends the SUBACK packet."
+     */
+    SubAck subAck(this->protocolVersion, action.mPacketId, action.mResponseCodes);
+    MqttPacket response(subAck);
+    writeMqttPacket(response);
+
+    auto session = getSession();
+
+    for (const DeferredRetainedSending &d : action.mRetainedSending)
+    {
+        store->giveClientRetainedMessages(session, d.mSubtopics, d.mQos, d.mSubscriptionIdentifier);
+    }
+}
+
 bool Client::tryAcmeRedirect()
 {
     if (authenticated)

client.h

@@ -118,6 +118,8 @@ class Client
     int disconnectReasonLogLevel = LOG_NOTICE;
     std::chrono::time_point<std::chrono::steady_clock> lastActivity = std::chrono::steady_clock::now();
 
+    std::unordered_map<uint16_t, SubAckAction> stagedSubAcks;
+
     std::string ssl_version;
     std::string clientid;
     std::string username;
@@ -161,6 +163,7 @@ class Client
     void setReadyForWriting(bool val, MutexLocked<WriteBuf> &writebuf);
     void setReadyForReading(bool val);
     void setAddr(const std::string &address);
+    void sendSubAck(const SubAckAction &action);
 
 public:
     uint8_t preAuthPacketCounter = 0;
@@ -254,6 +257,8 @@ class Client
     const std::unique_ptr<StowedClientRegistrationData> &getRegistrationData() const;
     void clearRegistrationData();
 
+    void stageOrSendSubAck(const std::shared_ptr<Client> &this_client, SubAckAction &&action, const size_t expandedCount);
+    void sendStagedSuback(const uint16_t packetId);
     void stageConnack(std::unique_ptr<ConnAck> &&c);
     void sendConnackSuccess();
     void sendConnackDeny(ReasonCodes reason);

enums.h

@@ -60,4 +60,11 @@ enum class HaProxyMode
     HaProxyClientVerficiationWithAuthn
 };
 
+enum class AddSubscriptionType
+{
+    Invalid,
+    NewSubscription,
+    ExistingSubscription
+};
+
 #endif // ENUMS_H

flashmq_plugin.cpp

@@ -138,7 +138,7 @@ bool flashmq_plugin_add_subscription(
     std::string topicDummy;
     parseSubscriptionShare(subtopics, shareName, topicDummy);
 
-    const AddSubscriptionType result = store->addSubscription(session_locked, subtopics, qos, noLocal, retainAsPublished, shareName, subscriptionIdentifier);
+    const auto [result, expand_count] = store->addSubscription(session_locked, 0, subtopics, qos, noLocal, retainAsPublished, shareName, subscriptionIdentifier);
     return result == AddSubscriptionType::Invalid ? false : true;
 }
 

lazysubscriptions.cpp

@@ -279,29 +279,31 @@ void LazySubscriptions::collectClients(
     }
 }
 
-void LazySubscriptions::expandLazySubscriptions(
-    TrackedSubscriptionMutationTask task, const std::shared_ptr<Session> originating_session,
+size_t LazySubscriptions::expandLazySubscriptions(
+    TrackedSubscriptionMutationTask task, const std::shared_ptr<Session> &originating_session, const uint16_t originating_packetId,
     const std::vector<std::string> &subtopics, const uint8_t qos)
 {
     assert(subtopics.size() == 0 || subtopics.at(0) != "$share");
 
     // Don't match the bridge's own internal subscriptions to ourself (the 'publish' lines in the config).
     if (originating_session->getClientType() == ClientType::LocalBridge)
-        return;
+        return 0;
 
     std::vector<ReceivingLazySubscriber> collected_receivers;
 
     {
         auto data = root.shared_lock();
 
         if ((*data)->children.empty())
-            return;
+            return 0;
 
-        collectClients(subtopics.begin(), subtopics.end(), data->get(), data->get()->node_hash, collected_receivers, -1, -1);
+        collectClients(subtopics.begin(), subtopics.end(), data->get(), 0, collected_receivers, -1, -1);
     }
 
     if (collected_receivers.empty())
-        return;
+        return 0;
+
+    size_t count = 0;
 
     for(ReceivingLazySubscriber &receiver : collected_receivers)
     {
@@ -321,13 +323,17 @@ void LazySubscriptions::expandLazySubscriptions(
             continue;
 
         const bool doWakeup = tracked_subs->addTrackedSubscriptionMutation(
-            TrackedSubscriptionMutation(pattern, effective_qos, originating_session->getClientId(), originating_session, task));
+            TrackedSubscriptionMutation(pattern, effective_qos, originating_session->getClientId(), originating_session, originating_packetId, task));
+
+        count++;
 
         if (doWakeup)
         {
             receiver.thread.lock()->queueProcessTrackedSubscriptionMutations(bridgeState, ProcessTrackedSubscriptionMutationsModifier::FirstFinishResending);
         }
     }
+
+    return count;
 }
 
 void registerLazySubscriptions(std::shared_ptr<BridgeState> &bridgeState)

lazysubscriptions.h

@@ -73,8 +73,8 @@ class LazySubscriptions
     void addSubscription(
         const std::shared_ptr<BridgeState> &bridgeState, const std::string &pattern,
         uint8_t qos, const std::string &distribution_group_name);
-    void expandLazySubscriptions(
-        TrackedSubscriptionMutationTask task, const std::shared_ptr<Session> originating_session,
+    size_t expandLazySubscriptions(
+        TrackedSubscriptionMutationTask task, const std::shared_ptr<Session> &originating_session, const uint16_t originating_packetId,
         const std::vector<std::string> &subtopics, const uint8_t qos);
 };
 

lockedweakptr.h

@@ -26,6 +26,12 @@ class LockedWeakPtr
         return p.lock();
     }
 
+    std::weak_ptr<T> get_weak()
+    {
+        std::lock_guard<std::mutex> l(m);
+        return p;
+    }
+
     bool expired()
     {
         std::lock_guard<std::mutex> l(m);

mainapp.cpp

@@ -623,10 +623,10 @@ void MainApp::start()
             websocketsubscriber->setFakeUpgraded();
             subscriptionStore->registerClientAndKickExistingOne(websocketsubscriber);
             subtopics = splitTopic("#");
-            subscriptionStore->addSubscription(websocketsubscriber->getSession(), subtopics, 0, false, false, empty, 0);
+            subscriptionStore->addSubscription(websocketsubscriber->getSession(), 0, subtopics, 0, false, false, empty, 0);
 
             subscriptionStore->registerClientAndKickExistingOne(subscriber);
-            subscriptionStore->addSubscription(subscriber->getSession(), subtopics, 0, false, false, empty, 0);
+            subscriptionStore->addSubscription(subscriber->getSession(), 0, subtopics, 0, false, false, empty, 0);
 
             if (connectionProtocol == ConnectionProtocol::WebsocketMqtt && strContains(fuzzFilePathLower, "upgrade"))
             {

mqttpacket.cpp

@@ -1388,7 +1388,7 @@ void MqttPacket::handleConnAck(std::shared_ptr<Client> &sender)
         std::string _;
         parseSubscriptionShare(subtopics, shareName, _);
         const bool no_local = shareName.empty(); // See above about no-local.
-        store->addSubscription(session, subtopics, pub.qos, no_local, true, shareName, 0);
+        store->addSubscription(session, 0, subtopics, pub.qos, no_local, true, shareName, 0);
     }
 
     ThreadGlobals::getThreadData()->publishBridgeState(bridgeState, true, {});
@@ -1794,11 +1794,10 @@ void MqttPacket::handleSubscribe(std::shared_ptr<Client> &sender)
         throw ProtocolError("No topics specified to subscribe to.", ReasonCodes::MalformedPacket);
     }
 
-    SubAck subAck(this->protocolVersion, packet_id, subs_reponse_codes);
-    MqttPacket response(subAck);
-    sender->writeMqttPacket(response);
-
     std::shared_ptr<Session> session = sender->getSession();
+    std::vector<DeferredRetainedSending> retainedSending;
+
+    size_t expandedCount = 0;
 
     // Adding the subscription will also send publishes for retained messages, so that's why we're doing it at the end.
     for(const SubscriptionTuple &tup : deferredSubscribes)
@@ -1808,18 +1807,22 @@ void MqttPacket::handleSubscribe(std::shared_ptr<Client> &sender)
 
         auto store = globals->subscriptionStore;
 
-        const AddSubscriptionType add_type = store->addSubscription(
-            session, tup.subtopics, tup.qos, tup.noLocal, tup.retainAsPublished, tup.shareName, tup.subscriptionIdentifier);
+        auto [addType, oneExpandedCount] = store->addSubscription(
+            session, packet_id, tup.subtopics, tup.qos, tup.noLocal, tup.retainAsPublished, tup.shareName, tup.subscriptionIdentifier);
+
+        expandedCount += oneExpandedCount;
 
         if (tup.authResult == AuthResult::success && tup.shareName.empty())
         {
             if ((tup.retainHandling == RetainHandling::SendRetainedMessagesAtSubscribe) ||
-                (tup.retainHandling == RetainHandling::SendRetainedMessagesAtNewSubscribeOnly && add_type == AddSubscriptionType::NewSubscription) )
+                (tup.retainHandling == RetainHandling::SendRetainedMessagesAtNewSubscribeOnly && addType == AddSubscriptionType::NewSubscription) )
             {
-                store->giveClientRetainedMessages(session, tup.subtopics, tup.qos, tup.subscriptionIdentifier);
+                retainedSending.emplace_back(tup.subtopics, tup.qos, tup.subscriptionIdentifier);
             }
         }
     }
+
+    sender->stageOrSendSubAck(sender, {std::move(retainedSending), std::move(subs_reponse_codes), packet_id}, expandedCount);
 }
 
 void MqttPacket::handleSubAck(std::shared_ptr<Client> &sender)
@@ -1847,7 +1850,7 @@ void MqttPacket::handleSubAck(std::shared_ptr<Client> &sender)
 
     if (tracked_subs)
     {
-        tracked_subs->removeMatchingInFlightTrackedSubscriptions(data.packet_id);
+        tracked_subs->handledSubackActions(data.packet_id);
 
         if (tracked_subs->requiresProcessingTrackedSubscriptions())
         {