Support TrackStatus in moxygen relay (#107)

Summary:
Added track_status to MoQRelay
issue: #97

co-authored by akash-a-n

Pull Request resolved: #107

Reviewed By: sandarsh

Differential Revision: D95152978

Pulled By: afrind

fbshipit-source-id: bd8296e628f78f332355e4cae192477c58d09a3a

Author: afrind

moxygen/relay/MoQRelay.cpp

@@ -988,6 +988,76 @@ folly::coro::Task<Publisher::FetchResult> MoQRelay::fetch(
       fetch, std::move(consumer), std::move(upstreamSession));
 }
 
+folly::coro::Task<Publisher::TrackStatusResult> MoQRelay::trackStatus(
+    TrackStatus trackStatus) {
+  XLOG(DBG1) << __func__ << " ftn=" << trackStatus.fullTrackName;
+
+  if (trackStatus.fullTrackName.trackNamespace.empty()) {
+    co_return folly::makeUnexpected(TrackStatusError(
+        {trackStatus.requestID,
+         TrackStatusErrorCode::TRACK_NOT_EXIST,
+         "namespace required"}));
+  }
+
+  auto subscriptionIt = subscriptions_.find(trackStatus.fullTrackName);
+  if (subscriptionIt != subscriptions_.end() &&
+      subscriptionIt->second.forwarder->numForwardingSubscribers() > 0) {
+    // We have active subscription - answer directly from local forwarder state
+    auto& subscription = subscriptionIt->second;
+    auto& forwarder = subscription.forwarder;
+
+    TrackStatusCode statusCode = TrackStatusCode::TRACK_NOT_STARTED;
+    // forwarder->largest() being set means: we have actually
+    // received at least one object for this track.
+    // subscription.handle being non-null means: the relay still has a
+    // live upstream Publisher::SubscriptionHandle for this track
+    if (forwarder->largest()) {
+      if (subscription.handle) {
+        statusCode = TrackStatusCode::IN_PROGRESS;
+      } else {
+        statusCode = TrackStatusCode::UNKNOWN;
+      }
+    }
+
+    TrackStatusOk trackStatusOk;
+    trackStatusOk.requestID = trackStatus.requestID;
+    trackStatusOk.groupOrder = forwarder->groupOrder();
+    trackStatusOk.largest = forwarder->largest();
+    trackStatusOk.fullTrackName = trackStatus.fullTrackName;
+    trackStatusOk.statusCode = statusCode;
+
+    XLOG(DBG1) << "Returning local track status for "
+               << trackStatus.fullTrackName
+               << " statusCode=" << (uint32_t)statusCode;
+    co_return trackStatusOk;
+  } else {
+    // No subscription - forward to upstream
+    auto upstreamSession =
+        findPublishNamespaceSession(trackStatus.fullTrackName.trackNamespace);
+
+    if (!upstreamSession) {
+      XLOG(DBG1) << "No upstream session for track: "
+                 << trackStatus.fullTrackName;
+      co_return folly::makeUnexpected(
+          TrackStatusError{
+              trackStatus.requestID,
+              TrackStatusErrorCode::TRACK_NOT_EXIST,
+              "no such namespace or track"});
+    }
+
+    // Forward the trackStatus request to the upstream publisher session
+    auto result = co_await upstreamSession->trackStatus(trackStatus);
+
+    if (result.hasError()) {
+      XLOG(DBG1) << "Upstream trackStatus failed: "
+                 << result.error().reasonPhrase;
+    } else {
+      XLOG(DBG1) << "Upstream trackStatus succeeded";
+    }
+    co_return result;
+  }
+}
+
 void MoQRelay::onEmpty(MoQForwarder* forwarder) {
   auto subscriptionIt = subscriptions_.find(forwarder->fullTrackName());
   if (subscriptionIt == subscriptions_.end()) {

moxygen/relay/MoQRelay.h

@@ -57,6 +57,9 @@ class MoQRelay : public Publisher,
     XLOG(INFO) << "Processing goaway uri=" << goaway.newSessionUri;
   }
 
+  folly::coro::Task<Publisher::TrackStatusResult> trackStatus(
+      TrackStatus req) override;
+
   std::shared_ptr<MoQSession> findPublishNamespaceSession(
       const TrackNamespace& ns);
 

moxygen/relay/test/MoQRelayTest.cpp

@@ -2176,4 +2176,103 @@ TEST_F(MoQRelayTest, ExactNamespaceSubscriberReceivesPublishNamespace) {
   removeSession(subscriber);
 }
 
+// Test: TrackStatus on non-existent track
+TEST_F(MoQRelayTest, TrackStatusNonExistentTrack) {
+  auto clientSession = createMockSession();
+
+  // Request trackStatus for a track that doesn't exist
+  TrackStatus trackStatus;
+  trackStatus.fullTrackName = kTestTrackName;
+  trackStatus.requestID = RequestID(1);
+
+  withSessionContext(clientSession, [&]() {
+    auto task = relay_->trackStatus(trackStatus);
+    auto res = folly::coro::blockingWait(std::move(task), exec_.get());
+
+    // Should return error indicating track not found
+    EXPECT_FALSE(res.hasValue());
+    EXPECT_EQ(res.error().errorCode, TrackStatusErrorCode::TRACK_NOT_EXIST);
+    EXPECT_FALSE(res.error().reasonPhrase.empty());
+  });
+
+  removeSession(clientSession);
+}
+
+// Test: TrackStatus on existing track - returns forwarder state (no upstream
+// call)
+TEST_F(MoQRelayTest, TrackStatusSuccessfulForward) {
+  auto publisherSession = createMockSession();
+  auto clientSession = createMockSession();
+
+  doPublish(publisherSession, kTestTrackName);
+
+  auto consumer = createMockConsumer();
+  subscribeToTrack(clientSession, kTestTrackName, consumer, RequestID(1));
+
+  TrackStatus trackStatus;
+  trackStatus.fullTrackName = kTestTrackName;
+  trackStatus.requestID = RequestID(2);
+
+  withSessionContext(clientSession, [&]() {
+    auto task = relay_->trackStatus(trackStatus);
+    auto res = folly::coro::blockingWait(std::move(task), exec_.get());
+
+    // Should return status from local forwarder
+    // Since no data was sent, statusCode should be TRACK_NOT_STARTED
+    EXPECT_TRUE(res.hasValue());
+    EXPECT_EQ(res.value().statusCode, TrackStatusCode::TRACK_NOT_STARTED);
+    EXPECT_EQ(res.value().fullTrackName, kTestTrackName);
+  });
+
+  removeSession(clientSession);
+  exec_->drive();
+  removeSession(publisherSession);
+}
+
+// Test: TrackStatus using namespace prefix matching (no exact subscription)
+// Verifies that when there's no exact subscription but a publisher has
+// published a matching namespace prefix, the relay correctly routes
+// TRACK_STATUS upstream using prefix matching
+TEST_F(MoQRelayTest, TrackStatusViaPrefixMatching) {
+  auto publisher = createMockSession();
+  auto requester = createMockSession();
+
+  // Publisher publishes namespace but NOT the specific track
+  doPublishNamespace(publisher, kTestNamespace);
+
+  // No exact subscription exists for kTestTrackName, so trackStatus should
+  // use prefix matching to find the publisher
+
+  // Mock the upstream trackStatus call
+  TrackStatusOk statusOk;
+  statusOk.requestID = RequestID(1);
+  statusOk.trackAlias = TrackAlias(0);
+  statusOk.largest = AbsoluteLocation{50, 25};
+
+  EXPECT_CALL(*publisher, trackStatus(_)).WillOnce([statusOk](auto /*ts*/) {
+    return folly::coro::makeTask<Publisher::TrackStatusResult>(statusOk);
+  });
+
+  // Execute trackStatus from requester's perspective
+  TrackStatus trackStatus;
+  trackStatus.requestID = RequestID(1);
+  trackStatus.fullTrackName = kTestTrackName;
+
+  withSessionContext(requester, [&]() {
+    auto task = relay_->trackStatus(trackStatus);
+    auto result = folly::coro::blockingWait(std::move(task), exec_.get());
+
+    // Should successfully forward via prefix matching and return the result
+    EXPECT_TRUE(result.hasValue())
+        << "TrackStatus via namespace prefix matching should succeed";
+    EXPECT_EQ(result.value().requestID, RequestID(1));
+    EXPECT_TRUE(result.value().largest.has_value());
+    EXPECT_EQ(result.value().largest->group, 50);
+    EXPECT_EQ(result.value().largest->object, 25);
+  });
+
+  removeSession(publisher);
+  removeSession(requester);
+}
+
 } // namespace moxygen::test

moxygen/test/MockMoQSession.h

@@ -61,6 +61,12 @@ class MockMoQSession : public MoQSession {
       (),
       (const, override));
 
+  MOCK_METHOD(
+      folly::coro::Task<Publisher::TrackStatusResult>,
+      trackStatus,
+      (TrackStatus),
+      (override));
+
   RequestID peekNextRequestID() {
     return RequestID(nextRequestID_++);
   }