Add more config options to ThriftServerInitialConfig

Summary: TSIA Reviewed By: praihan Differential Revision: D45245991 fbshipit-source-id: f2e4d79ad35d83b54a3596535ed2546bef3bc2ef
facebook · May 16, 2023 · 2304334 · 2304334
1 parent ec20114
commit 2304334
Show file tree

Hide file tree

Showing 3 changed files with 101 additions and 19 deletions.
diff --git a/thrift/lib/cpp2/server/ThriftServerConfig.cpp b/thrift/lib/cpp2/server/ThriftServerConfig.cpp
@@ -35,9 +35,44 @@ ThriftServerConfig::ThriftServerConfig(
   if (auto& [value, isSet] = initialConfig.maxRequests_; isSet) {
     maxRequests_.setDefault(folly::observer::makeStaticObserver(value));
   }
+  if (auto& [value, isSet] = initialConfig.maxConnections_; isSet) {
+    maxConnections_.setDefault(folly::observer::makeStaticObserver(value));
+  }
+  if (auto& [value, isSet] = initialConfig.maxResponseSize_; isSet) {
+    maxResponseSize_.setDefault(folly::observer::makeStaticObserver(value));
+  }
+  if (auto& [value, isSet] = initialConfig.useClientTimeout_; isSet) {
+    useClientTimeout_.setDefault(folly::observer::makeStaticObserver(value));
+  }
+  if (auto& [value, isSet] = initialConfig.taskExpireTimeout_; isSet) {
+    taskExpireTime_.setDefault(folly::observer::makeStaticObserver(value));
+  }
+  if (auto& [value, isSet] = initialConfig.streamExpireTimeout_; isSet) {
+    streamExpireTime_.setDefault(folly::observer::makeStaticObserver(value));
+  }
   if (auto& [value, isSet] = initialConfig.queueTimeout_; isSet) {
     queueTimeout_.setDefault(folly::observer::makeStaticObserver(value));
   }
+  if (auto& [value, isSet] = initialConfig.socketQueueTimeout_; isSet) {
+    socketQueueTimeout_.setDefault(folly::observer::makeStaticObserver(value));
+  }
+  if (auto& [value, isSet] = initialConfig.egressMemoryLimit_; isSet) {
+    egressMemoryLimit_.setDefault(folly::observer::makeStaticObserver(value));
+  }
+  if (auto& [value, isSet] = initialConfig.egressBufferBackpressureThreshold_;
+      isSet) {
+    egressBufferBackpressureThreshold_.setDefault(
+        folly::observer::makeStaticObserver(value));
+  }
+  if (auto& [value, isSet] = initialConfig.ingressMemoryLimit_; isSet) {
+    ingressMemoryLimit_.setDefault(folly::observer::makeStaticObserver(value));
+  }
+  if (auto& [value, isSet] =
+          initialConfig.minPayloadSizeToEnforceIngressMemoryLimit_;
+      isSet) {
+    minPayloadSizeToEnforceIngressMemoryLimit_.setDefault(
+        folly::observer::makeStaticObserver(value));
+  }
 }
 
 std::string ThriftServerConfig::getCPUWorkerThreadName() const {

diff --git a/thrift/lib/cpp2/server/ThriftServerConfig.h b/thrift/lib/cpp2/server/ThriftServerConfig.h
@@ -821,23 +821,50 @@ class ThriftServerInitialConfig {
 
   // to fix oss for now we'll have this as a pair<T, bool> to mimick the
   // behavior of optional
-#define THRIFT_SERVER_INITIAL_CONFIG_DEFINE(TYPE, NAME)        \
- private:                                                      \
-  std::pair<TYPE, bool> NAME##_ = {{}, false};                 \
-                                                               \
- public:                                                       \
-  FOLLY_CONSTEVAL ThriftServerInitialConfig NAME(TYPE value) { \
-    auto initialConfig(*this);                                 \
-    initialConfig.NAME##_.first = value;                       \
-    initialConfig.NAME##_.second = true;                       \
-    return initialConfig;                                      \
+#define THRIFT_SERVER_INITIAL_CONFIG_DEFINE(TYPE, NAME, INIT_VALUE) \
+ private:                                                           \
+  std::pair<TYPE, bool> NAME##_ = {INIT_VALUE, false};              \
+                                                                    \
+ public:                                                            \
+  FOLLY_CONSTEVAL ThriftServerInitialConfig NAME(TYPE value) {      \
+    auto initialConfig(*this);                                      \
+    initialConfig.NAME##_.first = value;                            \
+    initialConfig.NAME##_.second = true;                            \
+    return initialConfig;                                           \
   }
 
   // replace with thrift struct if/when it becomes constexpr friendly
 
-  THRIFT_SERVER_INITIAL_CONFIG_DEFINE(uint32_t, maxRequests)
-  THRIFT_SERVER_INITIAL_CONFIG_DEFINE(std::chrono::milliseconds, queueTimeout)
-
+  THRIFT_SERVER_INITIAL_CONFIG_DEFINE(uint32_t, maxRequests, 0)
+  THRIFT_SERVER_INITIAL_CONFIG_DEFINE(uint32_t, maxConnections, 0)
+  THRIFT_SERVER_INITIAL_CONFIG_DEFINE(uint64_t, maxResponseSize, 0)
+  THRIFT_SERVER_INITIAL_CONFIG_DEFINE(bool, useClientTimeout, true)
+  THRIFT_SERVER_INITIAL_CONFIG_DEFINE(
+      std::chrono::milliseconds, taskExpireTimeout, {})
+  THRIFT_SERVER_INITIAL_CONFIG_DEFINE(
+      std::chrono::milliseconds, streamExpireTimeout, {})
+  THRIFT_SERVER_INITIAL_CONFIG_DEFINE(
+      std::chrono::milliseconds, queueTimeout, {})
+  THRIFT_SERVER_INITIAL_CONFIG_DEFINE(
+      std::chrono::nanoseconds, socketQueueTimeout, {})
+  THRIFT_SERVER_INITIAL_CONFIG_DEFINE(size_t, egressMemoryLimit, 0)
+  THRIFT_SERVER_INITIAL_CONFIG_DEFINE(
+      size_t, egressBufferBackpressureThreshold, 0)
+  THRIFT_SERVER_INITIAL_CONFIG_DEFINE(size_t, ingressMemoryLimit, 0)
+  THRIFT_SERVER_INITIAL_CONFIG_DEFINE(
+      size_t, minPayloadSizeToEnforceIngressMemoryLimit, 0)
+  /*
+   * Running UBSan causes some tests to FATAL. The current theory is that this
+   * is caused by the compiler not generating proper code for FOLLY_CONSTEVAL.
+   * For now, we're adding a dummy field at the end as a hack to workaround this
+   * issue, until its fixed. All new fields should be added before this last
+   * dummy field
+   */
+
+  THRIFT_SERVER_INITIAL_CONFIG_DEFINE(
+      std::chrono::milliseconds,
+      lastUnusedField,
+      {}) // DO NOT ADD FIELDS AFTER THIS
 #undef THRIFT_SERVER_INITIAL_CONFIG_DEFINE
 
  private:

diff --git a/thrift/lib/cpp2/server/test/ThriftServerConfigTest.cpp b/thrift/lib/cpp2/server/test/ThriftServerConfigTest.cpp
@@ -29,13 +29,33 @@ TEST(ThriftServerInitialConfigTest, testSetOneConfigKnob) {
 }
 
 TEST(ThriftServerInitialConfigTest, testSetAllConfigKnobs) {
-  constexpr std::chrono::milliseconds queueTimeout{5};
+  using namespace std::chrono_literals;
   constexpr auto initialConfig =
-      ThriftServerInitialConfig().maxRequests(2048).queueTimeout(queueTimeout);
+      ThriftServerInitialConfig()
+          .maxRequests(1)
+          .maxConnections(2)
+          .maxResponseSize(3)
+          .useClientTimeout(false)
+          .taskExpireTimeout(4ms)
+          .streamExpireTimeout(5ms)
+          .queueTimeout(6ms)
+          .socketQueueTimeout(7ns)
+          .egressMemoryLimit(8)
+          .egressBufferBackpressureThreshold(9)
+          .ingressMemoryLimit(10)
+          .minPayloadSizeToEnforceIngressMemoryLimit(11);
   ThriftServer server(initialConfig);
   folly::observer_detail::ObserverManager::waitForAllUpdates();
-  EXPECT_EQ(detail::getThriftServerConfig(server).getMaxRequests().get(), 2048);
-  EXPECT_EQ(
-      detail::getThriftServerConfig(server).getQueueTimeout().get(),
-      queueTimeout);
+  EXPECT_EQ(server.getMaxRequests(), 1);
+  EXPECT_EQ(server.getMaxConnections(), 2);
+  EXPECT_EQ(server.getMaxResponseSize(), 3);
+  EXPECT_EQ(server.getUseClientTimeout(), false);
+  EXPECT_EQ(server.getTaskExpireTime().count(), 4);
+  EXPECT_EQ(server.getStreamExpireTime().count(), 5);
+  EXPECT_EQ(server.getQueueTimeout().count(), 6);
+  EXPECT_EQ((**server.getSocketQueueTimeout()).count(), 7);
+  EXPECT_EQ(server.getEgressMemoryLimit(), 8);
+  EXPECT_EQ(server.getEgressBufferBackpressureThreshold(), 9);
+  EXPECT_EQ(server.getIngressMemoryLimit(), 10);
+  EXPECT_EQ(server.getMinPayloadSizeToEnforceIngressMemoryLimit(), 11);
 }