Add allDataReceived callback to ObjectReceiver

Summary:
The `allDataReceived` callback that I'm adding in is going to be called when:
1. SUBSCRIBE_DONE has been received
2. All subgroups have been closed

Reviewed By: afrind

Differential Revision: D88871876

fbshipit-source-id: abf9feebe7b9db4a6c9cc0735fc3ddd26d3639c3

Author: Aman Sharma

moxygen/ObjectReceiver.h

@@ -26,14 +26,21 @@ class ObjectReceiverCallback {
   virtual void onEndOfStream() = 0;
   virtual void onError(ResetStreamErrorCode) = 0;
   virtual void onSubscribeDone(SubscribeDone done) = 0;
+  // Called when SUBSCRIBE_DONE has arrived AND all outstanding subgroup
+  // streams have closed. Only fires for subscriptions, not fetches.
+  virtual void onAllDataReceived() {}
 };
 
+class ObjectReceiver;
+
 class ObjectSubgroupReceiver : public SubgroupConsumer {
   std::shared_ptr<ObjectReceiverCallback> callback_{nullptr};
+  std::shared_ptr<ObjectReceiver> parent_{nullptr};
   StreamType streamType_;
   ObjectHeader header_;
   folly::IOBufQueue payload_{folly::IOBufQueue::cacheChainLength()};
   folly::Optional<TrackAlias> trackAlias_;
+  bool finished_{false};
 
  public:
   explicit ObjectSubgroupReceiver(
@@ -47,6 +54,10 @@ class ObjectSubgroupReceiver : public SubgroupConsumer {
         header_(groupID, subgroupID, 0, priority),
         trackAlias_(trackAlias) {}
 
+  void setParent(std::shared_ptr<ObjectReceiver> parent) {
+    parent_ = std::move(parent);
+  }
+
   void setFetchGroupAndSubgroup(uint64_t groupID, uint64_t subgroupID) {
     streamType_ = StreamType::FETCH_HEADER;
     header_.group = groupID;
@@ -138,18 +149,29 @@ class ObjectSubgroupReceiver : public SubgroupConsumer {
 
   folly::Expected<folly::Unit, MoQPublishError> endOfSubgroup() override {
     callback_->onEndOfStream();
+    notifyParentFinished();
     return folly::unit;
   }
 
   void reset(ResetStreamErrorCode error) override {
     callback_->onError(error);
+    notifyParentFinished();
   }
+
+ private:
+  void notifyParentFinished();
 };
 
-class ObjectReceiver : public TrackConsumer, public FetchConsumer {
+class ObjectReceiver : public TrackConsumer,
+                       public FetchConsumer,
+                       public std::enable_shared_from_this<ObjectReceiver> {
   std::shared_ptr<ObjectReceiverCallback> callback_{nullptr};
   folly::Optional<ObjectSubgroupReceiver> fetchPublisher_;
   folly::Optional<TrackAlias> trackAlias_;
+  // Tracking for onAllDataReceived callback (subscription mode only)
+  size_t openSubgroups_{0};
+  bool subscribeDoneDelivered_{false};
+  bool allDataCallbackSent_{false};
 
  public:
   enum Type { SUBSCRIBE, FETCH };
@@ -171,8 +193,30 @@ class ObjectReceiver : public TrackConsumer, public FetchConsumer {
   folly::Expected<std::shared_ptr<SubgroupConsumer>, MoQPublishError>
   beginSubgroup(uint64_t groupID, uint64_t subgroupID, Priority priority)
       override {
-    return std::make_shared<ObjectSubgroupReceiver>(
+    ++openSubgroups_;
+    auto receiver = std::make_shared<ObjectSubgroupReceiver>(
         callback_, trackAlias_, groupID, subgroupID, priority);
+    receiver->setParent(shared_from_this());
+    return receiver;
+  }
+
+  // Called when a subgroup stream finishes (via endOfSubgroup or reset)
+  void onSubgroupFinished() {
+    if (openSubgroups_ > 0) {
+      --openSubgroups_;
+    }
+    maybeFireAllDataReceived();
+  }
+
+  // Fire onAllDataReceived callback once when both conditions are met:
+  // 1. SUBSCRIBE_DONE has been received
+  // 2. All subgroup streams have closed
+  void maybeFireAllDataReceived() {
+    if (!allDataCallbackSent_ && subscribeDoneDelivered_ &&
+        openSubgroups_ == 0) {
+      allDataCallbackSent_ = true;
+      callback_->onAllDataReceived();
+    }
   }
 
   folly::Expected<folly::SemiFuture<folly::Unit>, MoQPublishError>
@@ -222,6 +266,8 @@ class ObjectReceiver : public TrackConsumer, public FetchConsumer {
   folly::Expected<folly::Unit, MoQPublishError> subscribeDone(
       SubscribeDone subDone) override {
     callback_->onSubscribeDone(std::move(subDone));
+    subscribeDoneDelivered_ = true;
+    maybeFireAllDataReceived();
     return folly::unit;
   }
 
@@ -311,4 +357,13 @@ class ObjectReceiver : public TrackConsumer, public FetchConsumer {
   }
 };
 
+// Definition of ObjectSubgroupReceiver::notifyParentFinished() - needs full
+// ObjectReceiver definition
+inline void ObjectSubgroupReceiver::notifyParentFinished() {
+  if (parent_ && !finished_) {
+    finished_ = true;
+    parent_->onSubgroupFinished();
+  }
+}
+
 } // namespace moxygen

moxygen/test/ObjectReceiverTest.cpp

@@ -0,0 +1,93 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <folly/portability/GMock.h>
+#include <folly/portability/GTest.h>
+#include <moxygen/ObjectReceiver.h>
+
+using namespace moxygen;
+using namespace testing;
+
+namespace {
+
+class MockObjectReceiverCallback : public ObjectReceiverCallback {
+ public:
+  MOCK_METHOD(
+      FlowControlState,
+      onObject,
+      (folly::Optional<TrackAlias>, const ObjectHeader&, Payload),
+      (override));
+  MOCK_METHOD(
+      void,
+      onObjectStatus,
+      (folly::Optional<TrackAlias>, const ObjectHeader&),
+      (override));
+  MOCK_METHOD(void, onEndOfStream, (), (override));
+  MOCK_METHOD(void, onError, (ResetStreamErrorCode), (override));
+  MOCK_METHOD(void, onSubscribeDone, (SubscribeDone), (override));
+  MOCK_METHOD(void, onAllDataReceived, (), (override));
+};
+
+Payload makePayload(const std::string& str) {
+  return folly::IOBuf::copyBuffer(str);
+}
+
+} // namespace
+
+class ObjectReceiverTest : public Test {
+ protected:
+  void SetUp() override {
+    callback_ = std::make_shared<MockObjectReceiverCallback>();
+  }
+
+  std::shared_ptr<MockObjectReceiverCallback> callback_;
+};
+
+TEST_F(ObjectReceiverTest, SubscribeDoneDelivery) {
+  auto receiver = std::make_shared<ObjectReceiver>(
+      ObjectReceiver::Type::SUBSCRIBE, callback_);
+
+  SubscribeDone done;
+  done.requestID = RequestID(1);
+  done.statusCode = SubscribeDoneStatusCode::SUBSCRIPTION_ENDED;
+
+  EXPECT_CALL(*callback_, onSubscribeDone(_)).Times(1);
+  // onAllDataReceived should be called since no subgroups are open
+  EXPECT_CALL(*callback_, onAllDataReceived()).Times(1);
+
+  auto result = receiver->subscribeDone(std::move(done));
+  EXPECT_TRUE(result.hasValue());
+}
+
+TEST_F(ObjectReceiverTest, AllDataReceivedAfterSubgroupClose) {
+  auto receiver = std::make_shared<ObjectReceiver>(
+      ObjectReceiver::Type::SUBSCRIBE, callback_);
+  receiver->setTrackAlias(TrackAlias(1));
+
+  // Start a subgroup
+  EXPECT_CALL(*callback_, onEndOfStream()).Times(1);
+  auto subgroupResult =
+      receiver->beginSubgroup(/*groupID=*/0, /*subgroupID=*/0, /*priority=*/0);
+  ASSERT_TRUE(subgroupResult.hasValue());
+  auto subgroup = *subgroupResult;
+
+  // Deliver subscribeDone while subgroup is open
+  SubscribeDone done;
+  done.requestID = RequestID(1);
+  done.statusCode = SubscribeDoneStatusCode::SUBSCRIPTION_ENDED;
+
+  EXPECT_CALL(*callback_, onSubscribeDone(_)).Times(1);
+  // onAllDataReceived should NOT be called yet because subgroup is open
+  EXPECT_CALL(*callback_, onAllDataReceived()).Times(0);
+
+  auto subDoneResult = receiver->subscribeDone(std::move(done));
+  EXPECT_TRUE(subDoneResult.hasValue());
+
+  // Now close the subgroup - onAllDataReceived should fire
+  EXPECT_CALL(*callback_, onAllDataReceived()).Times(1);
+  auto endResult = subgroup->endOfSubgroup();
+  EXPECT_TRUE(endResult.hasValue());
+}