changed queues ws executor

Summary: landing some more experimental ws things

Reviewed By: pranavtbhat

Differential Revision: D59291732

fbshipit-source-id: 0bc6d68b845a590f011d040328fc960990f32bd2

thrift/conformance/stresstest/server/StressTestServer.cpp

@@ -22,6 +22,11 @@
 
 #include <thrift/conformance/stresstest/util/IoUringUtil.h>
 
+#include <rroeser/scripts/src/executor/WorkStealingExecutor.h>
+#include <folly/executors/CPUThreadPoolExecutor.h>
+#include <thrift/lib/cpp2/server/ParallelConcurrencyController.h>
+#include <thrift/lib/cpp2/server/ThriftServer.h>
+
 DEFINE_int32(port, 5000, "Server port");
 DEFINE_int32(io_threads, 0, "Number of IO threads (0 == number of cores)");
 DEFINE_int32(cpu_threads, 0, "Number of CPU threads (0 == number of cores)");
@@ -30,6 +35,8 @@ DEFINE_int32(
     -1,
     "Configures max requests, 0 will disable max request limit");
 DEFINE_bool(io_uring, false, "Enables io_uring if available when set to true");
+DEFINE_bool(no_resource_pile_or_cc, false, "Disable resource pile and cc");
+DEFINE_bool(work_stealing, false, "Enable work stealing test");
 DEFINE_string(
     certPath,
     "folly/io/async/test/certs/tests-cert.pem",
@@ -42,6 +49,10 @@ DEFINE_string(
     caPath,
     "folly/io/async/test/certs/ca-cert.pem",
     "Path to client trusted CA file");
+DEFINE_bool(enable_overload_checker, false, "Enable overload checker");
+DEFINE_bool(enable_resource_pools, false, "Enable resource pools");
+DEFINE_bool(
+    disable_active_request_tracking, false, "Disabled Active Request Tracking");
 
 namespace apache {
 namespace thrift {
@@ -117,6 +128,12 @@ std::shared_ptr<ThriftServer> createStressTestServer(
       getIOThreadPool("thrift_eventbase", FLAGS_io_threads));
   server->setNumCPUWorkerThreads(sanitizeNumThreads(FLAGS_cpu_threads));
 
+  LOG(INFO) << "Active Request Tracking Disabled: "
+            << FLAGS_disable_active_request_tracking;
+  if (FLAGS_disable_active_request_tracking) {
+    server->disableActiveRequestsTracking();
+  }
+
   if (FLAGS_max_requests > -1) {
     LOG(INFO) << "Setting max server requests: " << FLAGS_max_requests;
     server->setMaxRequests(FLAGS_max_requests);
@@ -126,6 +143,50 @@ std::shared_ptr<ThriftServer> createStressTestServer(
       !FLAGS_caPath.empty()) {
     server->setSSLConfig(getSSLConfig());
   }
+  if (FLAGS_no_resource_pile_or_cc) {
+    LOG(INFO) << "Resource Pile and CC disabled";
+    server->resourcePoolSet().setResourcePool(
+        apache::thrift::ResourcePoolHandle::defaultAsync(),
+        nullptr,
+        std::make_shared<folly::CPUThreadPoolExecutor>(
+            sanitizeNumThreads(FLAGS_cpu_threads)),
+        nullptr);
+
+    server->ensureResourcePools();
+    server->requireResourcePools();
+  }
+
+  if (FLAGS_enable_resource_pools) {
+    server->ensureResourcePools();
+    server->requireResourcePools();
+  }
+
+  if (FLAGS_work_stealing) {
+    LOG(INFO) << "Work stealing enabled";
+    auto t = sanitizeNumThreads(FLAGS_cpu_threads);
+
+    auto executor = std::make_shared<folly::WorkStealingExecutor>(t, 4 * t);
+
+    RoundRobinRequestPile::Options options;
+
+    std::unique_ptr<RequestPileInterface> requestPile{
+        std::make_unique<RoundRobinRequestPile>(std::move(options))};
+
+    auto theConcurrencyController =
+        std::make_unique<ParallelConcurrencyController>(
+            *requestPile.get(), *executor.get());
+
+    server->resourcePoolSet().setResourcePool(
+        apache::thrift::ResourcePoolHandle::defaultAsync(),
+        std::move(requestPile),
+        executor,
+        std::move(theConcurrencyController),
+        concurrency::PRIORITY::NORMAL,
+        false);
+
+    server->ensureResourcePools();
+    server->requireResourcePools();
+  }
   return server;
 }