Clean up established FETCH receivers on session close

Summary:
This commit and the next one make it so that we send STOP_SENDING on the unidirectional stream (on which we are receiving the FETCH data) when the session is being cleaned up.

Here's the flow:
* `MoQSession::cleanup` calls `FetchTrackReceiveState::sessionClosed`
* `FetchTrackReceiveState::sessionClosed` (defined in the next commit) calls `cancelSource_.requestCancellation`
* The cancellation callback leads to sending a STOP_SENDING

In addition to this, we call `reset` on the `FetchConsumer` in order to inform the application.

___

Differential Revision: D116635422

fbshipit-source-id: 68939c0365819aaa8cb0e1cb9d62594dd083092b

Author: Aman Sharma

moxygen/MoQSession.cpp

@@ -2224,6 +2224,10 @@ class MoQSession::FetchTrackReceiveState
     return fetchGroupOrder_;
   }
 
+  uint32_t dataStreamCancelCode() const {
+    return dataStreamCancelCode_;
+  }
+
   void resetFetchCallback(MoQSession* session) {
     callback_.reset();
     if (fetchOkAndAllDataReceived()) {
@@ -2246,8 +2250,25 @@ class MoQSession::FetchTrackReceiveState
     resetFetchCallback(session);
   }
 
+  void sessionClosed() {
+    if (!fetchEstablished_) {
+      return;
+    }
+    auto callback = std::exchange(callback_, nullptr);
+    if (bidiControl_) {
+      bidiControl_->cancel(ResetStreamErrorCode::SESSION_CLOSED);
+    }
+    dataStreamCancelCode_ =
+        folly::to_underlying(ResetStreamErrorCode::SESSION_CLOSED);
+    cancelSource_.requestCancellation();
+    if (callback) {
+      callback->reset(ResetStreamErrorCode::SESSION_CLOSED);
+    }
+  }
+
   void fetchOK(FetchOk ok) {
     XLOG(DBG1) << __func__ << " trackReceiveState=" << this;
+    fetchEstablished_ = true;
     promise_.setValue(std::move(ok));
   }
 
@@ -2280,6 +2301,8 @@ class MoQSession::FetchTrackReceiveState
   GroupOrder fetchGroupOrder_;
   folly::coro::Promise<FetchResult> promise_;
   uint64_t currentStreamId_{0};
+  uint32_t dataStreamCancelCode_{0};
+  bool fetchEstablished_{false};
 };
 
 const std::shared_ptr<BidiStreamControl>&
@@ -2409,13 +2432,14 @@ MoQSession::~MoQSession() {
 
 void MoQSession::cleanup() {
   cancelGoawayTimeout();
-  // Each loop disarms its entry's bidi control before tearing it down so a
-  // peer close mid-shutdown can't race the canonical error delivery.
-  // fetches_ has no per-entry destroy loop, so it needs an upfront pass.
-  for (auto& [reqID, fetch] : fetches_) {
+  while (!fetches_.empty()) {
+    auto it = fetches_.begin();
+    auto fetch = std::move(it->second);
+    fetches_.erase(it);
     if (const auto& control = fetch->bidiControl()) {
       control->disarmOnPeerTermination();
     }
+    fetch->sessionClosed();
   }
   while (!pubTracks_.empty()) {
     auto it = pubTracks_.begin();
@@ -2453,10 +2477,6 @@ void MoQSession::cleanup() {
   subTracks_.clear();
   // We parse a publishDone after cleanup
   reqIdToTrackAlias_.clear();
-  // TODO: there needs to be a way to queue an error in TrackReceiveState,
-  // both from here, when close races the FETCH stream, and from readLoop
-  // where we get a reset.
-  fetches_.clear();
   for (auto& [reqID, pendingState] : pendingRequests_) {
     if (const auto& control = pendingState->bidiControl()) {
       control->disarmOnPeerTermination();
@@ -3735,6 +3755,7 @@ folly::coro::Task<void> MoQSession::dataStreamReadLoop(
   auto rhToken = readHandle->getCancelToken();
   XLOG(DBG1) << __func__ << " id=" << id << " sess=" << this;
   bool isSubscriptionStream = false;
+  std::shared_ptr<FetchTrackReceiveState> fetchState;
   std::optional<folly::CancellationCallback> fetchCancelCb;
   auto g = folly::makeGuard([func = __func__, this, id, &isSubscriptionStream] {
     XLOG(DBG1) << "exit " << func << " id=" << id << " sess=" << this;
@@ -3824,14 +3845,16 @@ folly::coro::Task<void> MoQSession::dataStreamReadLoop(
 
   // Lambda for onFetch
   auto onFetchFunc =
-      [this, &token, &codec, &fetchCancelCb, &readHandle](RequestID requestID) {
+      [this, &token, &codec, &fetchState, &fetchCancelCb, &readHandle](
+          RequestID requestID) {
         auto state = getFetchTrackReceiveState(requestID);
         if (state) {
           // FetchTrackReceiveState lifecycle now controls read loop
+          fetchState = state;
           token = state->getCancelToken();
-          fetchCancelCb.emplace(token, [&readHandle] {
-            if (readHandle) {
-              readHandle->stopSending(0);
+          fetchCancelCb.emplace(token, [&readHandle, &fetchState] {
+            if (readHandle && fetchState) {
+              readHandle->stopSending(fetchState->dataStreamCancelCode());
               readHandle = nullptr;
             }
           });

moxygen/test/MoQSessionFetchTests.cpp

@@ -410,6 +410,7 @@ CO_TEST_P_X(MoQSessionTest, FetchBadLength) {
       co_await folly::coro::timeout(
           std::move(contract.second), std::chrono::milliseconds(100)),
       folly::FutureTimeout);
+  EXPECT_CALL(*fetchCallback_, reset(ResetStreamErrorCode::SESSION_CLOSED));
   clientSession_->close(SessionCloseErrorCode::NO_ERROR);
 }
 CO_TEST_P_X(PreDraft18Test, FetchOverLimit) {
@@ -436,6 +437,8 @@ CO_TEST_P_X(PreDraft18Test, FetchOverLimit) {
   EXPECT_CALL(
       *clientSubscriberStatsCallback_,
       onFetchError(FetchErrorCode::INTERNAL_ERROR));
+  EXPECT_CALL(*fetchCallback1, reset(ResetStreamErrorCode::SESSION_CLOSED));
+  EXPECT_CALL(*fetchCallback2, reset(ResetStreamErrorCode::SESSION_CLOSED));
   res = co_await clientSession_->fetch(fetch, fetchCallback3);
   EXPECT_TRUE(res.hasError());
 }
@@ -470,6 +473,8 @@ CO_TEST_P_X(Draft18Test, FetchOverBidiStreamLimit) {
   EXPECT_TRUE(blocked.hasError());
   EXPECT_EQ(blocked.error().errorCode, FetchErrorCode::INTERNAL_ERROR);
 
+  EXPECT_CALL(*fetchCb, reset(ResetStreamErrorCode::SESSION_CLOSED))
+      .Times(kLimit);
   clientSession_->close(SessionCloseErrorCode::NO_ERROR);
 }
 CO_TEST_P_X(MoQSessionTest, FetchOutOfOrder) {
@@ -691,6 +696,7 @@ CO_TEST_P_X(Draft18Test, FetchBidiStreamStopSending) {
   clientWt_->readHandles.at(0)->stopSending(0);
   co_await cancelBaton;
 
+  EXPECT_CALL(*fetchCallback_, reset(ResetStreamErrorCode::SESSION_CLOSED));
   clientSession_->close(SessionCloseErrorCode::NO_ERROR);
 }
 
@@ -752,6 +758,71 @@ CO_TEST_P_X(Draft18Test, FetchBidiStreamFetchCancel) {
   clientSession_->close(SessionCloseErrorCode::NO_ERROR);
 }
 
+CO_TEST_P_X(Draft18Test, FetchSessionCleanupBeforeDataStreamOpens) {
+  co_await setupMoQSession();
+
+  expectFetch([](Fetch fetch, auto /*fetchPub*/) -> TaskFetchResult {
+    co_return makeFetchOkResult(fetch, AbsoluteLocation{100, 100});
+  });
+  expectFetchSuccess();
+  EXPECT_CALL(*clientSubscriberStatsCallback_, recordFetchLatency(_));
+  auto res =
+      co_await clientSession_->fetch(getFetch({0, 0}, {0, 1}), fetchCallback_);
+  EXPECT_TRUE(res.hasValue());
+
+  folly::coro::Baton resetBaton;
+  EXPECT_CALL(*fetchCallback_, reset(ResetStreamErrorCode::SESSION_CLOSED))
+      .WillOnce([&] { resetBaton.post(); });
+
+  clientSession_->close(SessionCloseErrorCode::NO_ERROR);
+  co_await resetBaton;
+}
+
+CO_TEST_P_X(Draft18Test, FetchSessionCleanupWithOpenDataStream) {
+  co_await setupMoQSession();
+
+  expectFetch([this](Fetch fetch, auto fetchPub) -> TaskFetchResult {
+    eventBase_.add([fetchPub = std::move(fetchPub)] {
+      EXPECT_TRUE(
+          fetchPub
+              ->object(
+                  0, 0, 0, moxygen::test::makeBuf(100), noExtensions(), false)
+              .hasValue());
+    });
+    co_return makeFetchOkResult(fetch, AbsoluteLocation{100, 100});
+  });
+
+  folly::coro::Baton objectReceived;
+  EXPECT_CALL(
+      *fetchCallback_, object(0, 0, 0, HasChainDataLengthOf(100), _, false, _))
+      .WillOnce([&] {
+        objectReceived.post();
+        return folly::unit;
+      });
+  expectFetchSuccess();
+  EXPECT_CALL(*clientSubscriberStatsCallback_, recordFetchLatency(_));
+  auto res =
+      co_await clientSession_->fetch(getFetch({0, 0}, {0, 1}), fetchCallback_);
+  EXPECT_TRUE(res.hasValue());
+  co_await objectReceived;
+
+  folly::coro::Baton resetBaton;
+  EXPECT_CALL(*fetchCallback_, reset(ResetStreamErrorCode::SESSION_CLOSED))
+      .WillOnce([&] { resetBaton.post(); });
+  auto dataStream = serverWt_->writeHandles.at(serverObjectStreamId());
+
+  clientSession_->close(SessionCloseErrorCode::NO_ERROR);
+  co_await resetBaton;
+
+  auto* dataException = dataStream->writeException();
+  EXPECT_NE(dataException, nullptr);
+  if (dataException) {
+    EXPECT_EQ(
+        dataException->error,
+        folly::to_underlying(ResetStreamErrorCode::SESSION_CLOSED));
+  }
+}
+
 // Once the FETCH data stream FINs, peer STOP_SENDING on the bidi is
 // informational and must not re-fire fetchCancel on the torn-down handle.
 CO_TEST_P_X(Draft18Test, NoFetchCancelAfterFetchComplete) {

moxygen/test/MoQSessionRequestUpdateTests.cpp

@@ -1329,8 +1329,13 @@ CO_TEST_P_X(MoQSessionTest, FetchRequestUpdateNotSupported) {
   EXPECT_EQ(updateResult.error().errorCode, RequestErrorCode::NOT_SUPPORTED);
 
   // Complete the fetch
-  EXPECT_CALL(*fetchCallback_, endOfFetch());
+  folly::coro::Baton fetchComplete;
+  EXPECT_CALL(*fetchCallback_, endOfFetch()).WillOnce([&] {
+    fetchComplete.post();
+    return folly::unit;
+  });
   fetchPubCaptured->endOfFetch();
+  co_await fetchComplete;
 
   clientSession_->close(SessionCloseErrorCode::NO_ERROR);
 }
@@ -1524,6 +1529,7 @@ CO_TEST_P_X(MoQSessionTest, FetchRequestUpdateNullSessionAfterAwait) {
   // fetchComplete -> session_ = null.  The cancellation token is also
   // triggered, but co_awaitTry catches the OperationCancelled — execution
   // continues to the dereference of session_ without a null check.
+  EXPECT_CALL(*fetchCallback_, reset(ResetStreamErrorCode::SESSION_CLOSED));
   serverSession_->close(SessionCloseErrorCode::NO_ERROR);
 
   // Post the baton — the coroutine resumes.  Without the fix, this crashes