Reduce fsdb server queuing for stats subscriptions

Summary:
In fsdb server, limit stats subscription queue size to 8 in order to
limit server-side memory usage due to subscription serve queue
in case of slow subscribers.

Rationale for choice of default value: queue size 8 provides server-side
queueing to hold messages enqueued in 1 min interval (6 stats updates
and 2 heartbeats). If subscriber is backlogged for longer than 1 min,
fsdb server will disconnect the subscription with disconnect reason
SLOW_SUBSCRIBER_QUEUE_FULL

In prod the p99 latency observed for NSDB stats processing is 37 seconds,
which is well within the latency bound indicated by this queue size change.
So this change is not expected to cause disconnects for subscriptions that
are normally working as of now, while protecting fsdb server from memory
pressure due to slow stats subscriber.

Differential Revision: D85967896

fbshipit-source-id: ae99ddc24ad2de6bfb2e92091e7cd85974029505

Author: Priyank Warkhede

fboss/fsdb/server/ServiceHandler.cpp

@@ -39,6 +39,16 @@ DEFINE_int32(
     10000,
     "Interval at which stats subscriptions are served");
 
+// queue size for serving stats subscriptions is chosen
+// to accommodate pending updates generated over 1 min interval
+// (6 updates + 2 heartbeats).
+// This limits memory usage on server due to subscription serve
+// queue while not disconnecting subscriber aggressively.
+DEFINE_int32(
+    statsSubscriptionServeQueueSize,
+    8,
+    "Max stats subscription updates to hold in server queue");
+
 DEFINE_int32(
     statsSubscriptionHeartbeat_s,
     30,
@@ -272,7 +282,10 @@ ServiceHandler::ServiceHandler(
               "stats",
               options_.serveIdPathSubs,
               true,
-              true)) {
+              true,
+              true /* serveGetRequestsWithLastPublishedState */,
+              FLAGS_statsSubscriptionServeQueueSize /* pathSubscriptionServeQueueSize */,
+              FLAGS_statsSubscriptionServeQueueSize /* defaultSubscriptionServeQueueSize */)) {
   num_instances_.incrementValue(1);
 
   initPerStreamCounters();

fboss/fsdb/server/ServiceHandler.h

@@ -19,6 +19,8 @@
 #include "fboss/lib/ThreadHeartbeat.h"
 #include "re2/re2.h"
 
+DECLARE_int32(statsSubscriptionServeQueueSize);
+
 DECLARE_bool(checkSubscriberConfig);
 DECLARE_bool(enforceSubscriberConfig);
 DECLARE_bool(checkOperOwnership);

fboss/fsdb/tests/client/BUCK

@@ -39,6 +39,7 @@ cpp_unittest(
         "//fboss/fsdb/client:fsdb_stream_client",
         "//fboss/fsdb/if:fsdb_model",
         "//fboss/fsdb/oper:extended_path_builder",
+        "//fboss/fsdb/server:handler",
         "//fboss/fsdb/tests/utils:fsdb_test_server",
         "//fboss/lib:common_utils",
         "//fboss/lib/thrift_service_client:thrift-service-client",

fboss/fsdb/tests/client/FsdbPubSubTest.cpp

@@ -2,6 +2,7 @@
 
 #include <gtest/gtest.h>
 #include "fboss/fsdb/oper/ExtendedPathBuilder.h"
+#include "fboss/fsdb/server/ServiceHandler.h"
 #include "fboss/fsdb/tests/client/FsdbTestClients.h"
 #include "fboss/fsdb/tests/utils/FsdbTestServer.h"
 #include "fboss/lib/CommonUtils.h"
@@ -18,6 +19,7 @@ namespace {
 
 const uint32_t kStateServeIntervalMs = 50;
 const uint32_t kStatsServeIntervalMs = 50;
+const uint32_t kSubscriptionServeQueueSize = 10;
 auto constexpr kSubscriberId = "fsdb_test_subscriber";
 auto constexpr kPublisherId = "fsdb_test_publisher";
 auto constexpr kUnknownPublisherId = "publisher_unknown";
@@ -45,7 +47,11 @@ class FsdbPubSubTest : public ::testing::Test {
     folly::LoggerDB::get().setLevel("fboss.fsdb", folly::LogLevel::DBG4);
     auto config = getFsdbConfig();
     fsdbTestServer_ = std::make_unique<FsdbTestServer>(
-        std::move(config), 0, kStateServeIntervalMs, kStatsServeIntervalMs);
+        std::move(config),
+        0,
+        kStateServeIntervalMs,
+        kStatsServeIntervalMs,
+        kSubscriptionServeQueueSize);
     publisherStreamEvbThread_ =
         std::make_unique<folly::ScopedEventBaseThread>();
     subscriberStreamEvbThread_ =
@@ -442,9 +448,91 @@ TYPED_TEST(FsdbPubSubTest, dupSubscriber) {
   EXPECT_NO_THROW({ auto res2 = this->subscribe(req); });
 }
 
-TYPED_TEST(FsdbPubSubTest, slowSubscriber) {
-  FLAGS_subscriptionServeQueueSize = 2;
+TYPED_TEST(FsdbPubSubTest, slowSubscriberDisconnectThreshold) {
+  // verify threshold for number of pending updates for slow subscriber
+  // disconnect
+
+  uint32_t queueSize = this->pubSubStats()
+      ? FLAGS_statsSubscriptionServeQueueSize
+      : kSubscriptionServeQueueSize;
+  uint32_t updatesPublished = 100 + queueSize;
+  uint32_t subscriptionServeIntervalMs =
+      this->pubSubStats() ? kStatsServeIntervalMs : kStateServeIntervalMs;
+  uint32_t publishIntervalMs = subscriptionServeIntervalMs + 20;
+  this->setupConnection(*this->publisher_, false);
+  this->checkPublishing({this->publisher_->clientId()});
 
+  // pause subscriber on initial sync long enough for all updates to be
+  // published and served so that queue builds up.
+  folly::Baton<> waitForInitialSync, waitForDisconnect;
+  folly::Baton<> resumeDataCb, resumeReconnect;
+  bool initialSyncOnce{false}, disconnectOnce{false};
+  auto slowSub = this->createSubscriber(
+      "fsdb_slow_subscriber",
+      [&waitForInitialSync, &resumeDataCb, &initialSyncOnce]() {
+        if (!initialSyncOnce) {
+          initialSyncOnce = true;
+          waitForInitialSync.post();
+          resumeDataCb.wait();
+        }
+      },
+      [&waitForDisconnect, &resumeReconnect, &disconnectOnce]() {
+        if (!disconnectOnce) {
+          disconnectOnce = true;
+          waitForDisconnect.post();
+          resumeReconnect.wait();
+        }
+      });
+  this->setupConnection(*slowSub);
+  this->checkSubscribed({slowSub->clientId()});
+  int updateNum{0};
+  for (; updateNum < updatesPublished; updateNum++) {
+    if (this->pubSubStats()) {
+      this->publishPortStats(makePortStats(updateNum));
+    } else {
+      std::string testStr = folly::to<std::string>("bar", updateNum);
+      this->publishAgentConfig(makeAgentConfig({{"foo", testStr}}));
+    }
+    std::this_thread::sleep_for(std::chrono::milliseconds(publishIntervalMs));
+  }
+
+  // validate server does not disconnect subscriber yet
+  waitForInitialSync.wait();
+  SubscriptionInfo info = slowSub->getInfo();
+  ASSERT_EQ(
+      info.state, FsdbStreamClient::ReconnectingThriftClient::State::CONNECTED);
+
+  // post another update and validate that server disconnects the slow
+  // subscriber
+  if (this->pubSubStats()) {
+    this->publishPortStats(makePortStats(updateNum));
+  } else {
+    std::string testStr = folly::to<std::string>("bar", updateNum);
+    this->publishAgentConfig(makeAgentConfig({{"foo", testStr}}));
+  }
+  std::this_thread::sleep_for(std::chrono::milliseconds(publishIntervalMs));
+
+  // resume subscriber data callback after all updates are published
+  resumeDataCb.post();
+
+  waitForDisconnect.wait();
+  info = slowSub->getInfo();
+  ASSERT_EQ(
+      info.disconnectReason, FsdbErrorCode::SUBSCRIPTION_SERVE_QUEUE_FULL);
+
+  WITH_RETRIES_N(90, {
+    fb303::ThreadCachedServiceData::get()->publishStats();
+    auto counterName = folly::sformat(
+        "{}.subscriber.{}.disconnects.slow_subscriber.count",
+        this->pubSubStats() ? "stats" : "fsdb",
+        "unspecified");
+    EXPECT_EVENTUALLY_GT(fb303::ServiceData::get()->getCounter(counterName), 0);
+  });
+
+  resumeReconnect.post();
+}
+
+TYPED_TEST(FsdbPubSubTest, slowSubscriber) {
   // publishInterval: wait for subscriptionServeIntervalMs+delta to prevent
   // published updates from being coalesced
   uint32_t updatesPublished = 120;
@@ -486,7 +574,6 @@ TYPED_TEST(FsdbPubSubTest, slowSubscriber) {
     std::this_thread::sleep_for(std::chrono::milliseconds(publishIntervalMs));
   }
 
-  // TODO: validate subscription serve queue watermark counter
   // resume subscriber data callback after all updates are published
   resumeDataCb.post();
 
@@ -513,7 +600,6 @@ TYPED_TEST(FsdbPubSubTest, slowSubscriber) {
 }
 
 TYPED_TEST(FsdbPubSubTest, slowSubscriberQueueWatermark) {
-  FLAGS_subscriptionServeQueueSize = 100;
   FLAGS_forceCloseSlowSubscriber = false;
 
   // publishInterval: wait for subscriptionServeIntervalMs+delta to prevent

fboss/fsdb/tests/utils/FsdbTestServer.cpp

@@ -12,7 +12,21 @@ FsdbTestServer::FsdbTestServer(
     std::shared_ptr<FsdbConfig> config,
     uint16_t port,
     uint32_t stateSubscriptionServe_ms,
-    uint32_t statsSubscriptionServe_ms) {
+    uint32_t statsSubscriptionServe_ms,
+    uint32_t subscriptionServeQueueSize,
+    uint32_t statsSubscriptionServeQueueSize) {
+  auto queueSize = std::to_string(subscriptionServeQueueSize);
+  gflags::SetCommandLineOptionWithMode(
+      "subscriptionServeQueueSize",
+      queueSize.c_str(),
+      gflags::SET_FLAG_IF_DEFAULT);
+
+  auto statsQueueSize = std::to_string(statsSubscriptionServeQueueSize);
+  gflags::SetCommandLineOptionWithMode(
+      "statsSubscriptionServeQueueSize",
+      statsQueueSize.c_str(),
+      gflags::SET_FLAG_IF_DEFAULT);
+
   // Run tests faster
   auto stateServeInterval = std::to_string(stateSubscriptionServe_ms);
   auto statsServeInterval = std::to_string(statsSubscriptionServe_ms);

fboss/fsdb/tests/utils/FsdbTestServer.h

@@ -27,7 +27,9 @@ class FsdbTestServer {
       std::shared_ptr<FsdbConfig> config,
       uint16_t port = 0,
       uint32_t stateSubscriptionServe_ms = 50,
-      uint32_t statsSubscriptionServe_ms = 1000);
+      uint32_t statsSubscriptionServe_ms = 1000,
+      uint32_t subscriptionServeQueueSize = 1000,
+      uint32_t statsSubscriptionServeQueueSize = 8);
   ~FsdbTestServer();
 
   uint16_t getFsdbPort() const {