add onSubscriptionBegin/onSubscriptionEnd gauge stats; fix onPublishDone semantics (#159)

Summary:
- onPublishDone fires only when a PUBLISH_DONE wire frame is received/sent, not for session-close synthesized cases
- Add onSubscriptionBegin/onSubscriptionEnd to MoQStatsCallback (common base) for an active-subscription gauge that can never go negative; fires on both SUBSCRIBE and PUBLISH paths (publisher and subscriber sides)
- onSubscriptionEnd fires at every subscription termination path: subscriber: deliverPublishDoneAndRemove, unsubscribe publisher: sendPublishDone, onUnsubscribe, cleanup
- Switch stats test mocks to NiceMock to suppress uninteresting-call warnings

Co-authored by: akash-a-n

This is an alternative to #158

Pull Request resolved: #159

Reviewed By: sandarsh

Differential Revision: D104433947

Pulled By: afrind

fbshipit-source-id: cdd86716c2bc9634dd37ad55240978a0224605d1

Author: afrind

moxygen/MoQSession.cpp

@@ -1958,6 +1958,8 @@ class MoQSession::SubscribeTrackReceiveState
             << "deliverPublishDoneAndRemove: Delivering PUBLISH_DONE to app; statusCode="
             << folly::to_underlying(pendingPublishDone_->statusCode)
             << " alias=" << alias_ << " requestID=" << requestID_;
+        MOQ_SUBSCRIBER_STATS(
+            session_->subscriberStatsCallback_, onSubscriptionEnd);
         auto token = cancelSource_.getToken();
         auto cb = std::exchange(callback_, nullptr);
         cb->publishDone(std::move(*pendingPublishDone_));
@@ -2195,6 +2197,7 @@ void MoQSession::cleanup() {
     auto requestID = it->first;
     auto pubTrack = std::move(it->second);
     pubTracks_.erase(it);
+    MOQ_PUBLISHER_STATS(publisherStatsCallback_, onSubscriptionEnd);
     pubTrack->terminatePublish(
         PublishDone(
             {requestID,
@@ -3658,6 +3661,7 @@ void MoQSession::onUnsubscribe(Unsubscribe unsubscribe) {
   } else {
     trackPublisher->unsubscribe();
     if (pubTracks_.erase(unsubscribe.requestID)) {
+      MOQ_PUBLISHER_STATS(publisherStatsCallback_, onSubscriptionEnd);
       retireRequestID(/*signalWriteLoop=*/true);
     } // else, the caller invoked publishDone, which isn't needed but fine
   }
@@ -3693,6 +3697,7 @@ void MoQSession::onPublishOk(PublishOk publishOk) {
         std::static_pointer_cast<TrackPublisherImpl>(trackIt->second);
     trackPublisher->onPublishOk(publishOk);
     pendingPublishTracks_.erase(trackIt->second->fullTrackName());
+    MOQ_PUBLISHER_STATS(publisherStatsCallback_, onSubscriptionBegin);
   }
 
   publishPtr->setValue(std::move(publishOk));
@@ -4753,6 +4758,7 @@ Subscriber::PublishResult MoQSession::publish(
 void MoQSession::publishOk(const PublishOk& pubOk, ReplyContext& replyContext) {
   XLOG(DBG1) << __func__ << " reqID=" << pubOk.requestID << " sess=" << this;
   MOQ_SUBSCRIBER_STATS(subscriberStatsCallback_, onPublishOk);
+  MOQ_SUBSCRIBER_STATS(subscriberStatsCallback_, onSubscriptionBegin);
 
   if (logger_) {
     logger_->logPublishOk(pubOk, ControlMessageType::CREATED);
@@ -4875,6 +4881,7 @@ folly::coro::Task<Publisher::SubscribeResult> MoQSession::subscribe(
     co_return folly::makeUnexpected(subscribeResult.error());
   } else {
     MOQ_SUBSCRIBER_STATS(subscriberStatsCallback_, onSubscribeSuccess);
+    MOQ_SUBSCRIBER_STATS(subscriberStatsCallback_, onSubscriptionBegin);
     co_return std::make_shared<ReceiverSubscriptionHandle>(
         std::move(subscribeResult.value()), trackAlias, shared_from_this());
   }
@@ -4883,6 +4890,7 @@ folly::coro::Task<Publisher::SubscribeResult> MoQSession::subscribe(
 void MoQSession::sendSubscribeOk(const SubscribeOk& subOk, ReplyContext& ctx) {
   XLOG(DBG1) << __func__ << " sess=" << this;
   MOQ_PUBLISHER_STATS(publisherStatsCallback_, onSubscribeSuccess);
+  MOQ_PUBLISHER_STATS(publisherStatsCallback_, onSubscriptionBegin);
   auto res = moqFrameWriter_.writeSubscribeOk(ctx.writeBuf(), subOk);
   if (!res) {
     XLOG(ERR) << "writeSubscribeOk failed sess=" << this;
@@ -4941,6 +4949,9 @@ void MoQSession::unsubscribe(const Unsubscribe& unsubscribe) {
              << " sess=" << this;
   // cancel() should send STOP_SENDING on any open streams for this
   // subscription
+  if (trackIt->second->getSubscribeCallback()) {
+    MOQ_SUBSCRIBER_STATS(subscriberStatsCallback_, onSubscriptionEnd);
+  }
   trackIt->second->cancel();
   subTracks_.erase(trackIt);
   reqIdToTrackAlias_.erase(trackAliasIt);
@@ -4969,6 +4980,7 @@ void MoQSession::sendPublishDone(const PublishDone& pubDone) {
               << " sess=" << this;
     return;
   }
+  MOQ_PUBLISHER_STATS(publisherStatsCallback_, onSubscriptionEnd);
   auto* ctx = it->second->replyContext();
   SCOPE_EXIT {
     pubTracks_.erase(it);

moxygen/stats/MoQStats.h

@@ -125,6 +125,25 @@ class MoQStatsCallback {
    * Subscriber: The publisher closed a subscription stream to us
    */
   virtual void onSubscriptionStreamClosed() = 0;
+
+  /*
+   * Publisher: A subscription became active (sent SUBSCRIBE_OK, or had a
+   *   PUBLISH accepted via PUBLISH_OK)
+   * Subscriber: A subscription became active (received SUBSCRIBE_OK, or sent
+   *   PUBLISH_OK accepting a PUBLISH)
+   * Pairs with onSubscriptionEnd to form a gauge of active subscriptions.
+   */
+  virtual void onSubscriptionBegin() = 0;
+
+  /*
+   * Publisher: An active subscription ended (sent PUBLISH_DONE, received
+   *   UNSUBSCRIBE, or session closed while subscription was active)
+   * Subscriber: An active subscription ended (PUBLISH_DONE delivered to app,
+   *   unsubscribe called, or session closed while subscription was active)
+   * Pairs with onSubscriptionBegin; this callback fires exactly once per
+   * onSubscriptionBegin, so the difference is a non-negative gauge.
+   */
+  virtual void onSubscriptionEnd() = 0;
 };
 
 class MoQPublisherStatsCallback : public MoQStatsCallback {

moxygen/test/MoQSessionSubscribeTests.cpp

@@ -806,6 +806,10 @@ CO_TEST_P_X(MoQSessionTest, UnsubscribeWithinPublishDone) {
   auto subscribeHandler = res.value();
 
   EXPECT_CALL(*clientSubscriberStatsCallback_, onUnsubscribe());
+  EXPECT_CALL(
+      *clientSubscriberStatsCallback_,
+      onPublishDone(PublishDoneStatusCode::TRACK_ENDED));
+  EXPECT_CALL(*clientSubscriberStatsCallback_, onSubscriptionEnd());
 
   EXPECT_CALL(*subscribeCallback_, publishDone(_))
       .WillOnce(testing::Invoke([&](const auto&) {
@@ -1020,6 +1024,8 @@ CO_TEST_P_X(MoQSessionTest, UnsubscribeFromWithinPublishDoneHandler) {
   // When publishDone is delivered, immediately call unsubscribe() from
   // handler
   folly::coro::Baton publishDoneInvoked;
+  // onPublishDone fires before this point (wire handler); NiceMock handles it.
+  EXPECT_CALL(*clientSubscriberStatsCallback_, onSubscriptionEnd());
   EXPECT_CALL(*subscribeCallback_, publishDone(_))
       .WillOnce(testing::Invoke([&](const auto&) {
         subscribeHandle->unsubscribe();
@@ -1107,3 +1113,26 @@ CO_TEST_P_X(MoQSessionTest, BeginSubgroupCallbackError) {
   co_await publishDone_;
   clientSession_->close(SessionCloseErrorCode::NO_ERROR);
 }
+
+// Regression: session close fires onSubscriptionEnd for abandoned publishers;
+// onPublishDone must NOT fire — no PUBLISH_DONE wire frame was ever received.
+CO_TEST_P_X(MoQSessionTest, SubscriptionEndStatFiredOnAbandonedPublisher) {
+  co_await setupMoQSession();
+  expectSubscribe([](auto sub, auto) -> TaskSubscribeResult {
+    // Publisher accepts the subscription but intentionally never calls
+    // publishDone, simulating an abandoned / improperly-closed publisher.
+    co_return makeSubscribeOkResult(sub);
+  });
+
+  auto res = co_await clientSession_->subscribe(
+      getSubscribe(kTestTrackName), subscribeCallback_);
+  EXPECT_FALSE(res.hasError());
+
+  EXPECT_CALL(*clientSubscriberStatsCallback_, onPublishDone(_)).Times(0);
+  EXPECT_CALL(*clientSubscriberStatsCallback_, onSubscriptionEnd());
+  EXPECT_CALL(*subscribeCallback_, publishDone(_))
+      .WillOnce(testing::Return(folly::unit));
+
+  clientSession_->close(SessionCloseErrorCode::NO_ERROR);
+  serverSession_->close(SessionCloseErrorCode::NO_ERROR);
+}

moxygen/test/MoQSessionTestCommon.cpp

@@ -194,16 +194,20 @@ void MoQSessionTest::SetUp() {
               folly::Expected<folly::Unit, MoQPublishError>(folly::unit)));
   EXPECT_CALL(*subscribeCallback_, setTrackAlias(_)).Times(testing::AtLeast(0));
 
-  clientSubscriberStatsCallback_ = std::make_shared<MockSubscriberStats>();
+  clientSubscriberStatsCallback_ =
+      std::make_shared<testing::NiceMock<MockSubscriberStats>>();
   clientSession_->setSubscriberStatsCallback(clientSubscriberStatsCallback_);
 
-  clientPublisherStatsCallback_ = std::make_shared<MockPublisherStats>();
+  clientPublisherStatsCallback_ =
+      std::make_shared<testing::NiceMock<MockPublisherStats>>();
   clientSession_->setPublisherStatsCallback(clientPublisherStatsCallback_);
 
-  serverSubscriberStatsCallback_ = std::make_shared<MockSubscriberStats>();
+  serverSubscriberStatsCallback_ =
+      std::make_shared<testing::NiceMock<MockSubscriberStats>>();
   serverSession_->setSubscriberStatsCallback(serverSubscriberStatsCallback_);
 
-  serverPublisherStatsCallback_ = std::make_shared<MockPublisherStats>();
+  serverPublisherStatsCallback_ =
+      std::make_shared<testing::NiceMock<MockPublisherStats>>();
   serverSession_->setPublisherStatsCallback(serverPublisherStatsCallback_);
 
   // For Draft15+, initialize version via ALPN since it's required
@@ -436,7 +440,11 @@ void MoQSessionTest::expectPublishDone(MoQControlCodec::Direction recipient) {
   EXPECT_CALL(
       *getPublisherStatsCallback(oppositeDirection(recipient)),
       onPublishDone(_));
+  EXPECT_CALL(
+      *getPublisherStatsCallback(oppositeDirection(recipient)),
+      onSubscriptionEnd());
   EXPECT_CALL(*getSubscriberStatsCallback(recipient), onPublishDone(_));
+  EXPECT_CALL(*getSubscriberStatsCallback(recipient), onSubscriptionEnd());
   EXPECT_CALL(*subscribeCallback_, publishDone(_)).WillOnce([&] {
     publishDone_.post();
     return folly::unit;

moxygen/test/Mocks.h

@@ -422,6 +422,10 @@ class MockPublisherStats : public MoQPublisherStatsCallback {
 
   MOCK_METHOD(void, onSubscriptionStreamClosed, (), (override));
 
+  MOCK_METHOD(void, onSubscriptionBegin, (), (override));
+
+  MOCK_METHOD(void, onSubscriptionEnd, (), (override));
+
   MOCK_METHOD(void, recordPublishNamespaceLatency, (uint64_t), (override));
 
   MOCK_METHOD(void, recordPublishLatency, (uint64_t), (override));
@@ -477,6 +481,10 @@ class MockSubscriberStats : public MoQSubscriberStatsCallback {
 
   MOCK_METHOD(void, onSubscriptionStreamClosed, (), (override));
 
+  MOCK_METHOD(void, onSubscriptionBegin, (), (override));
+
+  MOCK_METHOD(void, onSubscriptionEnd, (), (override));
+
   MOCK_METHOD(void, recordSubscribeLatency, (uint64_t), (override));
 
   MOCK_METHOD(void, recordFetchLatency, (uint64_t), (override));