relay: install the local forwarder chain on subscribe-created tracks

A track created by its first subscriber left no entry in the publisher
thread's LocalForwarderRegistry, so a subscriber landing on that thread
joined a separate local forwarder instead of the publisher's. First
subscriber setup now calls installPublisherForwarder on the publisher
executor with removeOnEmpty=true, so the entry is there to find and is
dropped with the last subscriber.

The initial state for the requesting subscriber is captured from the
publisher forwarder, so subscribers joining an existing attachment also
get the most recent largest and extensions.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: afrind

src/MoqxRelay.cpp

@@ -458,20 +458,18 @@ MoqxRelay::validatePublishNamespace(const FullTrackName& ftn, RequestID requestI
 // source ends, so a chain over a forwarder that does not hold it has nothing to remove.
 LocalForwarderRegistry::Claim MoqxRelay::installPublisherForwarder(
     const FullTrackName& ftn,
-    const std::shared_ptr<MoQForwarder>& fwd
+    const std::shared_ptr<MoQForwarder>& fwd,
+    bool removeOnEmpty
 ) {
-  // removeOnEmpty=false: the publisher's forwarder must survive subscriber churn.
+  auto& localReg = localRegistry();
   auto relayAdapter = std::make_shared<WeakRelayForwarderCallback>(weak_from_this());
   auto crossExec =
       std::make_shared<CrossExecForwarderCallback>(relayExec_, fwd, std::move(relayAdapter));
-  fwd->setCallback(std::make_shared<LocalForwarderCallback>(
-      tlForwarders_.get(),
-      ftn,
-      std::move(crossExec),
-      /*removeOnEmpty=*/false
-  ));
+  fwd->setCallback(
+      std::make_shared<LocalForwarderCallback>(&localReg, ftn, std::move(crossExec), removeOnEmpty)
+  );
 
-  return tlForwarders_->replace(ftn, fwd);
+  return localReg.replace(ftn, fwd);
 }
 
 // Called from LocalPublishFilter::publish() on publisherExec. Creates the publisher's
@@ -485,12 +483,8 @@ Subscriber::PublishResult MoqxRelay::publishFromPublisherExec(
     return folly::makeUnexpected(std::move(*err));
   }
 
-  if (!tlForwarders_.get()) {
-    tlForwarders_.reset(new LocalForwarderRegistry());
-  }
-
   auto localPubFwd = std::make_shared<MoQForwarder>(pub.fullTrackName);
-  installPublisherForwarder(pub.fullTrackName, localPubFwd)
+  installPublisherForwarder(pub.fullTrackName, localPubFwd, /*removeOnEmpty=*/false)
       .markReady(InitialTrackState{pub.largest, pub.extensions});
 
   // crossExecFilter is a channel subscriber for the relay exec
@@ -1137,6 +1131,10 @@ folly::coro::Task<void> MoqxRelay::addSubscriberAndPublishViaLocalForwarder(
   );
 
   auto [localFwd, isNew, localReg] = acquireLocalForwarder(ftn, initial);
+  if (!localFwd) {
+    // Setup for this track is in flight on this thread; drop the fanout.
+    co_return;
+  }
 
   auto p = startPublish(subscriberSession, localFwd, forward, pinned, nullptr);
   if (!p) {
@@ -1208,8 +1206,12 @@ MoqxRelay::acquireLocalForwarder(const FullTrackName& ftn, const InitialTrackSta
     return {std::move(localFwd), /*isNew=*/true, localReg};
   }
   auto* ready = std::get_if<LocalForwarderRegistry::Ready>(&joined);
-  XCHECK(ready) << "local forwarder entry still pending; a caller failed to resolve its claim: "
-                << ftn;
+  if (!ready) {
+    // A setup on this thread owns the entry and cannot be joined mid-flight. Drop the
+    // fanout rather than publishing over a subscribe that is still claiming the track.
+    XLOG(WARNING) << "local forwarder setup in flight, dropping publish fanout for " << ftn;
+    return {};
+  }
   return {ready->forwarder, /*isNew=*/false, localReg};
 }
 
@@ -1765,7 +1767,7 @@ std::optional<SubscribeError> MoqxRelay::completeUpstreamSubscription(
 // Runs on relayExec_. Registers the subscribe and wires the local forwarder to the
 // publisher; if it's a new subscription,also installs the passive relay chain, issues the
 // upstream subscribe, and completes the setup. The subscriberExec tail reads the
-// returned PublisherAttachment (handles upstreamOk).
+// returned PublisherAttachment.
 folly::coro::Task<MoqxRelay::PublisherAttachment> MoqxRelay::attachNewLocalForwarderOnRelayExec(
     const SubscribeRequest& subReq,
     LocalForwarderRegistry* localReg,
@@ -1809,6 +1811,12 @@ folly::coro::Task<MoqxRelay::PublisherAttachment> MoqxRelay::attachNewLocalForwa
   co_await folly::coro::co_withExecutor(
       folly::getKeepAliveToken(attach.publisherExec),
       [&]() -> folly::coro::Task<void> {
+        LocalForwarderRegistry::Claim publisherClaim;
+        if (sr.firstSetup) {
+          publisherClaim =
+              installPublisherForwarder(ftn, attach.publisherFwd, /*removeOnEmpty=*/true);
+        }
+
         installChannelSubscriber(
             *cbs.channelCb,
             *attach.publisherFwd,
@@ -1837,7 +1845,14 @@ folly::coro::Task<MoqxRelay::PublisherAttachment> MoqxRelay::attachNewLocalForwa
               attach.publisherFwd,
               setup.clientRequestID
           );
+          if (upstreamResult->hasValue()) {
+            const auto& ok = upstreamResult->value();
+            publisherClaim.markReady(InitialTrackState{ok.largest, ok.extensions});
+          }
         }
+        // Capture largest/extensions on publisherExec to set the initial state for the
+        // subscriber.  Captured but ignored when the firstSetup returned an error.
+        attach.initial = InitialTrackState::capture(*attach.publisherFwd);
       }()
   );
   // Back on relayExec_.
@@ -1882,7 +1897,6 @@ folly::coro::Task<MoqxRelay::PublisherAttachment> MoqxRelay::attachNewLocalForwa
     attach.error = folly::makeUnexpected(std::move(*err));
     co_return attach;
   }
-  attach.upstreamOk = std::move(upstreamOk);
   co_return attach;
 }
 
@@ -1900,7 +1914,9 @@ MoqxRelay::joinOrPrepareUpstreamSubscription(SubscribeRequest subReq) {
 
   auto firstOrSubsequent = registry_.getOrCreateFromSubscribe(
       ftn,
-      shared_from_this(),
+      // installPublisherForwarder puts the real chain on the forwarder's own exec instead;
+      // a relay-direct callback would run there and touch registry_ off relayExec_.
+      std::shared_ptr<MoQForwarder::Callback>(nullptr),
       [this, &ftn](std::shared_ptr<MoQForwarder> f) { return buildFilterChain(ftn, std::move(f)); }
   );
 
@@ -2055,15 +2071,9 @@ folly::coro::Task<Publisher::SubscribeResult> MoqxRelay::subscribeFromSubscriber
     co_return std::move(*attach.error);
   }
 
-  // Only the first subscriber has an upstream OK. A subsequent one leaves this empty and
-  // marks the entry ready with no largest (a bug fixed in a subsequent commit)
-  InitialTrackState initial;
-  if (attach.upstreamOk) {
-    initial = InitialTrackState{attach.upstreamOk->largest, attach.upstreamOk->extensions};
-    // Also on the subscriber, so a post-SUBSCRIBE_OK joining fetch resolves.
-    initial.applyTo(*sub);
-  }
-  claim.setInitialState(initial);
+  claim.setInitialState(attach.initial);
+  // Also on the subscriber, so a post-SUBSCRIBE_OK joining fetch resolves.
+  attach.initial.applyTo(*sub);
   replayPendingFowarderEvents(localFwd.get(), attach.finalCallback, *pendingCb, forward);
   localFwd->tryProcessNewGroupRequest(subReq.params);
   claim.markReady();

src/MoqxRelay.h

@@ -366,10 +366,12 @@ class MoqxRelay : public moxygen::Publisher,
       std::shared_ptr<moxygen::MoQSession> session
   );
 
-  // Runs on fwd's exec; the returned Claim owes a markReady/fail.
+  // Runs on fwd's exec; the returned Claim owes a markReady/fail. removeOnEmpty=false for a
+  // publish-initiated forwarder, which must survive subscriber churn.
   LocalForwarderRegistry::Claim installPublisherForwarder(
       const moxygen::FullTrackName& ftn,
-      const std::shared_ptr<moxygen::MoQForwarder>& fwd
+      const std::shared_ptr<moxygen::MoQForwarder>& fwd,
+      bool removeOnEmpty
   );
 
   std::optional<moxygen::PublishError>
@@ -491,7 +493,8 @@ class MoqxRelay : public moxygen::Publisher,
     folly::Executor* publisherExec{nullptr};
     bool ownsRelayChain{false}; // firstSetup path installed the passive relay chain
     std::shared_ptr<moxygen::MoQForwarder::Callback> finalCallback;
-    std::optional<UpstreamOk> upstreamOk;
+    // Captured off the publisher forwarder on its own exec, the only race-free place.
+    InitialTrackState initial;
     std::optional<SubscribeResult> error; // set => bail
   };
 

test/MoqxRelayPublishTests.cpp

@@ -451,6 +451,9 @@ TEST_P(MoQRelayTest, PublishReconnectDuringSubscribeScopeGuardCrash) {
   // Relay state mutations must run on the relay executor; doPublishNamespaceDone
   // touches the namespace tree, which publishDone also cleans up via relayEvb_.
   verifyOnRelayExec([&] { relay_->doPublishNamespaceDone(kTestNamespace, publisherSession2); });
+  // The publisher forwarder reaches the relay through relayExec_, so it drops its session
+  // refs a hop after the calls above.
+  driveIfMultiThread();
 }
 
 // Same reconnect scenario but the upstream subscribe returns OK instead of an