Fix server IO->CPU dispatch so rpc_num_cpu_worker_threads > 1 works

Summary:
Setting `rpc_num_cpu_worker_threads > 1` made the server go completely idle (~2% CPU, requests never processed) during a benchmark, so the flag was unusable and the prod-like "IO -> CPU -> IO" dispatch pattern it is meant to exercise never actually ran.

Root cause: the Thrift handlers are `async_eb_*`, so they must complete their `HandlerCallback` on the callback's own EventBase. `ucacheBenchOnRequestCommon` offloaded work to a `folly::CPUThreadPoolExecutor` (from `getCpuPool()`, only created when the flag is > 1) and then invoked the handler — which called `cb->result(...)` — directly on the CPU-pool thread. Those pool threads are not EventBase threads and are unknown to the per-EventBase handler map, so the reply was never delivered. In-flight requests never drained, clients stopped sending, and the server sat idle. With the flag == 1, `getCpuPool()` returns `nullptr`, everything runs inline on the IO thread, and it works — which is why the bug only showed up for > 1.

Fix: split compute from delivery. The per-request `handler` now COMPUTES and returns the reply; `ucacheBenchOnRequestCommon` owns delivering it via `callback->result(...)`. On the CPU-pool path it captures the callback's EventBase up front, runs the compute on the pool, then bounces the completion back with `evb->runInEventBaseThread(...)` so the async_eb callback is completed on its own EventBase. The inline and fiber paths are behaviorally unchanged for the default `rpc_num_cpu_worker_threads == 1`.

Verified: with the fix, `rpc_num_cpu_worker_threads=64` drives real traffic and the run completes cleanly ("All clients finished benchmark"), where before it hung the server at ~2% CPU. Note: the dispatch path trades throughput for the extra IO->CPU->IO context switches (as intended, to match production's tao:slow scheduling); it is a fidelity knob, not a throughput/utilization lever.

Reviewed By: excelle08

Differential Revision: D110506501

Author: charles-typ

packages/ucache_bench/server/UcacheBenchOnRequest.cpp

@@ -21,26 +21,26 @@ void UcacheBenchOnRequest::onRequestThrift(
     apache::thrift::HandlerCallbackPtr<UcbGetReply> callback,
     UcbGetRequest&& request) {
   ucacheBenchOnRequestCommon(
-      std::move(callback), std::move(request), [this](auto&& cb, auto&& req) {
-        cb->result(server_->processUcbGetSync(req));
+      std::move(callback), std::move(request), [this](auto&& req) {
+        return server_->processUcbGetSync(req);
       });
 }
 
 void UcacheBenchOnRequest::onRequestThrift(
     apache::thrift::HandlerCallbackPtr<UcbSetReply> callback,
     UcbSetRequest&& request) {
   ucacheBenchOnRequestCommon(
-      std::move(callback), std::move(request), [this](auto&& cb, auto&& req) {
-        cb->result(server_->processUcbSetSync(req));
+      std::move(callback), std::move(request), [this](auto&& req) {
+        return server_->processUcbSetSync(req);
       });
 }
 
 void UcacheBenchOnRequest::onRequestThrift(
     apache::thrift::HandlerCallbackPtr<UcbDeleteReply> callback,
     UcbDeleteRequest&& request) {
   ucacheBenchOnRequestCommon(
-      std::move(callback), std::move(request), [this](auto&& cb, auto&& req) {
-        cb->result(server_->processUcbDeleteSync(req));
+      std::move(callback), std::move(request), [this](auto&& req) {
+        return server_->processUcbDeleteSync(req);
       });
 }
 
@@ -49,12 +49,12 @@ void UcacheBenchOnRequest::onRequestThrift(
         callback,
     facebook::memcache::McVersionRequest&& request) {
   ucacheBenchOnRequestCommon(
-      std::move(callback), std::move(request), [](auto&& cb, auto&& /* req */) {
+      std::move(callback), std::move(request), [](auto&& /* req */) {
         facebook::memcache::McVersionReply reply;
         reply.result() = carbon::Result::FOUND;
         reply.value() =
             *folly::IOBuf::copyBuffer("UcacheBench 1.0 (with Fiber support)");
-        cb->result(std::move(reply));
+        return reply;
       });
 }
 

packages/ucache_bench/server/UcacheBenchRequestCommon.h

@@ -4,6 +4,7 @@
 
 #include <fmt/format.h>
 #include <folly/executors/CPUThreadPoolExecutor.h>
+#include <folly/io/async/EventBase.h>
 #include <folly/portability/GFlags.h>
 #include <mcrouter/lib/carbon/Result.h>
 #include "UcacheBenchIOThreadContext.h"
@@ -30,7 +31,16 @@ inline folly::CPUThreadPoolExecutor* getCpuPool() {
 }
 
 /**
- * Common entry point to run request with fiber management
+ * Common entry point to run a request.
+ *
+ * `handler` COMPUTES and returns the reply (it must not complete the callback
+ * itself). This function owns delivering the reply via `callback->result(...)`.
+ * That split matters for the IO→CPU→IO path: these are async_eb handlers, so
+ * the callback MUST be completed on its own EventBase. When work is offloaded
+ * to the CPU pool we compute the reply there, then bounce the completion back
+ * to the callback's EventBase via runInEventBaseThread(). Completing an
+ * async_eb callback directly from a CPU-pool thread never delivers the reply —
+ * in-flight requests never drain and the server goes idle.
  */
 template <class Callback, class Request, class Handler>
 void ucacheBenchOnRequestCommon(
@@ -41,27 +51,33 @@ void ucacheBenchOnRequestCommon(
 {
   auto* cpuPool = getCpuPool();
   if (cpuPool) {
+    auto* evb = callback->getEventBase();
     cpuPool->add([handler,
                   cb = std::forward<Callback>(callback),
-                  req = std::forward<Request>(request)]() mutable {
-      handler(std::move(cb), std::move(req));
+                  req = std::forward<Request>(request),
+                  evb]() mutable {
+      auto reply = handler(req);
+      evb->runInEventBaseThread(
+          [cb = std::move(cb), reply = std::move(reply)]() mutable {
+            cb->result(std::move(reply));
+          });
     });
     return;
   }
 
-  // If fibers are disabled, execute directly
+  // If fibers are disabled, execute directly on the IO thread.
   if (!FLAGS_enable_fibers ||
       !UcacheBenchIOThreadContext::isInitializedForCurrentThread()) {
-    handler(std::forward<Callback>(callback), std::forward<Request>(request));
+    callback->result(handler(request));
     return;
   }
 
-  // Execute the handler in a fiber
+  // Execute the handler in a fiber (so simulateIOLatency can yield).
   UcacheBenchIOThreadContext::tlInstance().fm().addTaskEager(
       [handler,
-       callbackFiber = std::forward<Callback>(callback),
-       requestFiber = std::forward<Request>(request)]() mutable {
-        handler(std::move(callbackFiber), std::move(requestFiber));
+       cb = std::forward<Callback>(callback),
+       req = std::forward<Request>(request)]() mutable {
+        cb->result(handler(req));
       });
 }