DUPLICATE_SUBSCRIPTION for PUBLISH received during pending SUBSCRIBE_OK

Summary: This addresses a part of the PR moq-wg/moq-transport#1341

Reviewed By: afrind

Differential Revision: D95634475

fbshipit-source-id: de071715d07ee72f7bb078a6bcf05d60a184b070

Author: Aman Sharma

moxygen/MoQSession.cpp

@@ -2184,6 +2184,7 @@ void MoQSession::cleanup() {
   }
   pendingRequests_.clear();
   pendingPublishTracks_.clear();
+  pendingSubscribeTracks_.clear();
   if (!cancellationSource_.isCancellationRequested()) {
     XLOG(DBG1) << "requestCancellation from cleanup sess=" << this;
     cancellationSource_.requestCancellation();
@@ -3652,6 +3653,11 @@ void MoQSession::onRequestError(RequestError error, FrameType frameType) {
   if (it != pendingRequests_.end()) {
     auto pendingState = std::move(it->second);
     pendingRequests_.erase(it);
+    // Remove from pending subscribe tracks if this was a subscribe
+    auto* trackPtr = pendingState->tryGetSubscribeTrack();
+    if (trackPtr) {
+      pendingSubscribeTracks_.erase((*trackPtr)->fullTrackName());
+    }
     if (getDraftMajorVersion(*getNegotiatedVersion()) > 14) {
       // determine real frame type from pendingRequest
       frameType = pendingState->getErrorFrameType();
@@ -3730,6 +3736,7 @@ void MoQSession::onSubscribeOk(SubscribeOk subOk) {
   }
   auto trackReceiveState = std::move(*trackPtr);
   pendingRequests_.erase(it);
+  pendingSubscribeTracks_.erase(trackReceiveState->fullTrackName());
 
   auto res = reqIdToTrackAlias_.try_emplace(subOk.requestID, subOk.trackAlias);
   if (!res.second) {
@@ -3886,6 +3893,19 @@ void MoQSession::onPublish(PublishRequest publish) {
     return;
   }
 
+  // Check for duplicate subscription: if there's already a pending outgoing
+  // SUBSCRIBE for the same track, reject the PUBLISH
+  if (pendingSubscribeTracks_.count(publish.fullTrackName)) {
+    XLOG(DBG1) << "Duplicate subscription for track with pending subscribe"
+               << " ftn=" << publish.fullTrackName << " sess=" << this;
+    publishError(
+        PublishError{
+            publish.requestID,
+            PublishErrorCode::DUPLICATE_SUBSCRIPTION,
+            "duplicate subscription"});
+    return;
+  }
+
   auto publishHandle = std::make_shared<ReceiverSubscriptionHandle>(
       SubscribeOk{publish.requestID}, publish.trackAlias, shared_from_this());
 
@@ -4719,6 +4739,7 @@ folly::coro::Task<Publisher::SubscribeResult> MoQSession::subscribe(
       fullTrackName, reqID, callback, this, trackAlias, logger_);
   pendingRequests_.emplace(
       reqID, PendingRequestState::makeSubscribeTrack(trackReceiveState));
+  pendingSubscribeTracks_.insert(fullTrackName);
   auto subscribeResultTry =
       co_await co_awaitTry(trackReceiveState->subscribeFuture());
   if (subscribeResultTry.hasException()) {

moxygen/MoQSession.h

@@ -896,6 +896,7 @@ class MoQSession : public Subscriber,
   folly::F14FastMap<RequestID, std::shared_ptr<PublisherImpl>, RequestID::hash>
       pubTracks_;
   folly::F14FastSet<FullTrackName, FullTrackName::hash> pendingPublishTracks_;
+  folly::F14FastSet<FullTrackName, FullTrackName::hash> pendingSubscribeTracks_;
   folly::F14FastMap<TrackAlias, std::list<Payload>, TrackAlias::hash>
       bufferedDatagrams_;
   folly::F14FastMap<TrackAlias, std::list<TimedBaton*>, TrackAlias::hash>

moxygen/test/MoQSessionPublishTests.cpp

@@ -938,10 +938,13 @@ CO_TEST_P_X(MoQSessionTest, SubscribeDuplicatesPendingPublish) {
 
   // Use a baton to delay the PUBLISH_OK so the publish stays pending
   folly::coro::Baton publishOkBaton;
+  bool publishMockCalled = false;
   EXPECT_CALL(*serverSubscriber, publish(_, _))
       .WillOnce(
-          [&publishOkBaton](
-              PublishRequest actualPub, auto) -> Subscriber::PublishResult {
+          [&publishOkBaton, &publishMockCalled](
+              const PublishRequest& actualPub,
+              auto) -> Subscriber::PublishResult {
+            publishMockCalled = true;
             auto consumer =
                 std::make_shared<testing::NiceMock<MockTrackConsumer>>();
             ON_CALL(*consumer, setTrackAlias(_))
@@ -972,7 +975,12 @@ CO_TEST_P_X(MoQSessionTest, SubscribeDuplicatesPendingPublish) {
   auto publishResult =
       clientSession_->publish(std::move(pub), makePublishHandle());
   EXPECT_TRUE(publishResult.hasValue());
-  co_await folly::coro::co_reschedule_on_current_executor;
+
+  // Wait for the server to process the PUBLISH (handlePublish is dispatched
+  // asynchronously via .start(), so multiple event loop iterations may be
+  // needed)
+  co_await rescheduleN(5);
+  EXPECT_TRUE(publishMockCalled);
 
   // Server subscribes to client for the same track while publish is pending
   auto subRes = co_await serverSession_->subscribe(
@@ -988,3 +996,64 @@ CO_TEST_P_X(MoQSessionTest, SubscribeDuplicatesPendingPublish) {
 
   clientSession_->close(SessionCloseErrorCode::NO_ERROR);
 }
+CO_TEST_P_X(MoQSessionTest, PublishDuplicatesPendingSubscribe) {
+  co_await setupMoQSessionForPublish(initialMaxRequestID_);
+
+  FullTrackName ftn{TrackNamespace{{"test"}}, "test-track"};
+
+  // Use a baton to delay the SUBSCRIBE_OK so the subscribe stays pending
+  folly::coro::Baton subscribeOkBaton;
+  bool subscribeMockCalled = false;
+  expectSubscribe(
+      [&subscribeOkBaton, &subscribeMockCalled](
+          auto sub, auto) -> TaskSubscribeResult {
+        subscribeMockCalled = true;
+        co_await subscribeOkBaton;
+        co_return makeSubscribeOkResult(sub);
+      });
+
+  // Client subscribes — start eagerly so it populates pendingSubscribeTracks_
+  auto subscribeConsumer =
+      std::make_shared<testing::NiceMock<MockTrackConsumer>>();
+  ON_CALL(*subscribeConsumer, setTrackAlias(_))
+      .WillByDefault(
+          testing::Return(
+              folly::Expected<folly::Unit, MoQPublishError>(folly::unit)));
+  ON_CALL(*subscribeConsumer, publishDone(_))
+      .WillByDefault(testing::Return(folly::unit));
+  auto subscribeSf =
+      clientSession_->subscribe(getSubscribe(ftn), subscribeConsumer)
+          .scheduleOn(&eventBase_)
+          .start();
+
+  // Wait for the SUBSCRIBE to be delivered and processed by the server
+  co_await rescheduleN(5);
+  EXPECT_TRUE(subscribeMockCalled);
+
+  // Server sends PUBLISH for the same track while subscribe is pending.
+  // Client should reject with DUPLICATE_SUBSCRIPTION.
+  auto handle = makePublishHandle();
+  auto publishResult = serverSession_->publish(
+      PublishRequest{
+          RequestID(0),
+          ftn,
+          TrackAlias(100),
+          GroupOrder::Default,
+          AbsoluteLocation{0, 100},
+          true,
+      },
+      handle);
+  EXPECT_TRUE(publishResult.hasValue());
+
+  auto replyRes = co_await std::move(publishResult.value().reply);
+  EXPECT_TRUE(replyRes.hasError());
+  EXPECT_EQ(
+      replyRes.error().errorCode, PublishErrorCode::DUPLICATE_SUBSCRIPTION);
+
+  // Let the SUBSCRIBE_OK through and complete the subscribe
+  subscribeOkBaton.post();
+  auto subRes = co_await std::move(subscribeSf).via(&eventBase_);
+  EXPECT_TRUE(subRes.hasValue());
+
+  clientSession_->close(SessionCloseErrorCode::NO_ERROR);
+}

moxygen/test/MoQSessionTestCommon.cpp

@@ -480,4 +480,10 @@ uint8_t MoQSessionTest::getRequestIDMultiplier() const {
   return 2;
 }
 
+folly::coro::Task<void> MoQSessionTest::rescheduleN(int n) {
+  for (int i = 0; i < n; ++i) {
+    co_await folly::coro::co_reschedule_on_current_executor;
+  }
+}
+
 }} // namespace moxygen::test

moxygen/test/MoQSessionTestCommon.h

@@ -227,6 +227,8 @@ class MoQSessionTest : public testing::TestWithParam<VersionParams>,
       std::shared_ptr<MockSubgroupConsumer> sgc)>;
   folly::coro::Task<void> publishValidationTest(TestLogicFn testLogic);
 
+  folly::coro::Task<void> rescheduleN(int n);
+
   folly::EventBase eventBase_;
   std::shared_ptr<MoQFollyExecutorImpl> MoQExecutor_;
   std::unique_ptr<proxygen::test::FakeSharedWebTransport> clientWt_;