Three knobs to rebalance Compression/Encryption: partial mock-ZSTD, server-ZSTD env gate, driver↔server TLS (facebookresearch#746)

Summary:
Pull Request resolved: facebookresearch#746

The t41 9-cell sweep (single×autoscale × rpc_fanout_scale on CPL/BGM/GRC, classified with the new v2 hot-func rules in D108210357) showed two persistent gaps vs prod multifeed_aggregator:

- **Compression over-target**: bench 5.9-10.0% across cells vs prod 3.9-5.2%. Bench runs ZSTD on every outbound mock_services channel (`MOCK_COMPRESS_ZSTD=1`) AND on every server response (`compressThrift`), so the total compression CPU footprint is roughly 2× prod.
- **Encryption under-target**: bench 0.9-1.8% on x86 vs prod 3.3-3.6%. The mock_services channel uses TLS (closes part of the gap) but the driver↔server channel runs plaintext over a raw `AsyncSocket`/libevent `bufferevent_socket`, leaving the bulk of the bench unencrypted.

This diff adds three independent knobs that let us experiment with the rebalance without touching the production-aligned defaults. All knobs are env-var-driven (LeafNodeRank uses gengetopt and rejects unknown CLI flags); all default to current behavior so unset = no change.

**Knob 1 — `MOCK_ZSTD_FRAC` (replaces all-or-nothing `MOCK_COMPRESS_ZSTD`):**
Float in [0.0, 1.0] that controls the fraction of `MockServicesClient` instances that enable per-channel ZSTD `CompressionConfig`. Each leaf builds one client per `SREventBase` thread (~88-176 per process); a process-global atomic counter assigns each new client a monotonic index, and the decision `(idx % 100) < int(frac * 100)` gives deterministic per-channel selection. Matches prod's mix-of-compressed-and-plaintext downstream services. Backwards compat: when `MOCK_ZSTD_FRAC` is unset, falls back to legacy `MOCK_COMPRESS_ZSTD` (1→1.0, 0→0.0); default when both unset = 1.0.

**Knob 2 — `FEEDSIM_SERVER_ZSTD` (gate for server-side response compression):**
Boolean env var, default true. When "0", both `compressThrift()` and `compressPayload()` in `LeafNodeRank.cc` return the input unchanged (passthrough), bypassing the ZSTD codec entirely. Reads the env once via `static const bool` so request-time cost is a single branch.

**Knob 3 — `FEEDSIM_DRIVER_TLS` (TLS on the driver↔server channel):**
Boolean env var, default false. When "1":
  - `FeedSimServer` reads `FEEDSIM_TLS_CERT` / `FEEDSIM_TLS_KEY` env vars (wired by `run.sh` to the existing bench cert/key at `${FEEDSIM_ROOT}/certs/example.{crt,key}`), creates a `folly::SSLContext`, and wraps each accepted fd in `AsyncSSLSocket(ctx, evb, fd, true /*server*/)` via the `AcceptCallback` constructor.
  - `FeedSimDriver` lazy-inits a process-global `SSL_CTX` (TLS_client_method, `SSL_VERIFY_NONE`) on first connection, and wraps each `DriverConnection`'s libevent `bufferevent` with `bufferevent_openssl_socket_new(base, fd, SSL_new(ctx), BUFFEREVENT_SSL_CONNECTING, BEV_OPT_CLOSE_ON_FREE)`. Requires `libevent_openssl` (added to `CMakeLists.txt` via `find_library(LIBEVENT_OPENSSL_LIB NAMES event_openssl ...)` and appended to `ranking-cpp2`'s link line).
Independent of `FEEDSIM_TLS`, which only covers the mock_services channel. No ALPN — the driver↔server protocol is custom binary, not Rocket.

Files:
- `MockServicesClient.cc`: add `envFloat()` helper + `resolveZstdFraction()` + atomic per-instance counter; replace `use_zstd = envBoolTrue(...)` with deterministic per-channel decision based on resolved fraction. `LOG_FIRST_N` logs the active fraction once per process.
- `LeafNodeRank.cc`: gate `compressThrift()` and `compressPayload()` on `FEEDSIM_SERVER_ZSTD`. Each function reads the env once via `static const bool`.
- `FeedSimServer.{cc,h}`: add `std::shared_ptr<folly::SSLContext>` member to `Impl`; in `run()` read `FEEDSIM_TLS_CERT` / `FEEDSIM_TLS_KEY` env vars and build the context if both set; pass the shared_ptr through `AcceptCallback`'s new ctor arg; in `connectionAccepted`, when ctx is set, wrap the fd in `AsyncSSLSocket` instead of plain `AsyncSocket`. Falls back to plain `AsyncSocket` when ctx is null.
- `FeedSimDriver.cc`: include `<event2/bufferevent_ssl.h>` + `<openssl/ssl.h>`; add `getDriverSslCtxOrNull()` with lazy static SSL_CTX init from `FEEDSIM_DRIVER_TLS` env; in `DriverConnection` ctor, when ctx is non-null, wrap fd in `bufferevent_openssl_socket_new` instead of plain `bufferevent_socket_new`.
- `CMakeLists.txt`: `find_library(LIBEVENT_OPENSSL_LIB NAMES event_openssl ...)` after the `find_program(GENGETOPT_EXECUTABLE)` block; append `${LIBEVENT_OPENSSL_LIB}` to `ranking-cpp2` next to `${LIBEVENT_LIB}`.
- `run.sh`: pass through `MOCK_ZSTD_FRAC` and `FEEDSIM_SERVER_ZSTD` env vars; when `FEEDSIM_DRIVER_TLS=1`, also export `FEEDSIM_TLS_CERT` and `FEEDSIM_TLS_KEY` pointing at `${FEEDSIM_ROOT}/certs/example.{crt,key}` (with readability check that errors out if missing).

Reviewed By: charles-typ

Differential Revision: D109352301

Author: excelle08

packages/feedsim/run.sh

@@ -780,6 +780,43 @@ main() {
         echo "MockServicesClient ZSTD: DISABLED (via MOCK_COMPRESS_ZSTD env)"
     fi
 
+    # t43 knobs (2026-06-10): three independent knobs for the bench-vs-prod
+    # Compression / Encryption rebalance. See plan doc t41/t43 progress logs.
+    #
+    # MOCK_ZSTD_FRAC: float in [0.0, 1.0]. Fraction of MockServicesClient
+    # channels that enable per-channel ZSTD. Replaces all-or-nothing
+    # MOCK_COMPRESS_ZSTD with prod-realistic partial enablement (some
+    # downstream services compress, others don't).
+    if [ -n "${MOCK_ZSTD_FRAC:-}" ]; then
+        export MOCK_ZSTD_FRAC
+        echo "MockServicesClient ZSTD fraction: ${MOCK_ZSTD_FRAC} (overrides MOCK_COMPRESS_ZSTD)"
+    fi
+    # FEEDSIM_SERVER_ZSTD: 0 disables server-side response ZSTD
+    # (compressThrift / compressPayload return passthrough). Default 1
+    # preserves current behavior. Use to reduce the bench's Compression
+    # CPU share when over-target.
+    if [ "${FEEDSIM_SERVER_ZSTD:-1}" != "1" ]; then
+        export FEEDSIM_SERVER_ZSTD=0
+        echo "Server-side response ZSTD: DISABLED (FEEDSIM_SERVER_ZSTD=0)"
+    fi
+    # FEEDSIM_DRIVER_TLS: 1 enables TLS on the driver↔server channel
+    # (DriverNodeRank ↔ LeafNodeRank). Server reads FEEDSIM_TLS_CERT /
+    # FEEDSIM_TLS_KEY env vars (set here to the existing bench cert/key
+    # under ${FEEDSIM_ROOT}/certs/); driver reads FEEDSIM_DRIVER_TLS
+    # directly. Closes the bench's Encryption CPU undershoot (prod
+    # 3.3-3.6% vs bench 0.9-1.5% in t41). Independent of FEEDSIM_TLS,
+    # which only covers the mock_services channel.
+    if [ "${FEEDSIM_DRIVER_TLS:-0}" = "1" ]; then
+        export FEEDSIM_DRIVER_TLS=1
+        export FEEDSIM_TLS_CERT="${FEEDSIM_ROOT}/certs/example.crt"
+        export FEEDSIM_TLS_KEY="${FEEDSIM_ROOT}/certs/example.key"
+        if [ ! -r "${FEEDSIM_TLS_CERT}" ] || [ ! -r "${FEEDSIM_TLS_KEY}" ]; then
+            echo "ERROR: FEEDSIM_DRIVER_TLS=1 but ${FEEDSIM_TLS_CERT} or .key not readable" >&2
+            exit 1
+        fi
+        echo "Driver↔Server TLS: ENABLED (cert=${FEEDSIM_TLS_CERT})"
+    fi
+
     # OMP_NUM_THREADS=1: cap PyTorch's OpenMP parallel backend pool to
     # 1 thread. at::set_num_threads(1) only affects libtorch's native
     # parallel backend; OpenMP-backed builds (which Meta's internal

packages/feedsim/third_party/src/workloads/ranking/CMakeLists.txt

@@ -38,6 +38,18 @@ find_package(Snappy REQUIRED)
 find_package(LibLZMA REQUIRED)
 find_program(GENGETOPT_EXECUTABLE gengetopt REQUIRED)
 
+# libevent_openssl: required by FeedSimDriver.cc's bufferevent_openssl_socket_new
+# (used when FEEDSIM_DRIVER_TLS=1). libevent ships this as a separate library
+# alongside the core libevent. The .pc is libevent_openssl.pc and the lib file
+# is libevent_openssl.{so,a}.
+find_library(LIBEVENT_OPENSSL_LIB
+    NAMES event_openssl
+    PATHS ${CMAKE_INSTALL_PREFIX}/lib /usr/local/lib /usr/lib /usr/lib64
+)
+if(NOT LIBEVENT_OPENSSL_LIB)
+    message(WARNING "libevent_openssl not found; FEEDSIM_DRIVER_TLS will fail at runtime")
+endif()
+
 include(if/CMakeLists.txt)
 add_dependencies(ranking-cpp2-target fbthrift)
 set_target_properties(
@@ -64,6 +76,7 @@ target_link_libraries(ranking-cpp2
     ${SNAPPY_LIBRARY}
     ${FOLLY_LIBRARIES}
     ${LIBEVENT_LIB}
+    ${LIBEVENT_OPENSSL_LIB}
     ${JEMALLOC_LIB}
     ${LIBLZMA_LIBRARIES}
     ${LIBSODIUM_LIBRARIES}

packages/feedsim/third_party/src/workloads/ranking/FeedSimDriver.cc

@@ -11,9 +11,12 @@
 #include <event2/event.h>
 #include <event2/buffer.h>
 #include <event2/bufferevent.h>
+#include <event2/bufferevent_ssl.h>
 #include <event2/thread.h>
 #include <netdb.h>
 #include <netinet/tcp.h>
+#include <openssl/err.h>
+#include <openssl/ssl.h>
 #include <pthread.h>
 #include <signal.h>
 #include <string.h>
@@ -218,6 +221,39 @@ void DriverStats::printStats(uint32_t type, double elapsed_secs) const {
   }
 }
 
+// FEEDSIM_DRIVER_TLS env gate: when "1", every DriverConnection wraps its
+// libevent bufferevent in a TLS bufferevent via OpenSSL. Bench-only — peer
+// cert verification is disabled (SSL_VERIFY_NONE). The SSL_CTX is process-
+// global (one per process), constructed lazily on first use to avoid paying
+// the OpenSSL init cost when TLS is off. Closes the bench's Encryption CPU
+// undershoot on the driver↔server channel (paired with FeedSimServer's
+// FEEDSIM_TLS_CERT / FEEDSIM_TLS_KEY env vars). See t41 progress log.
+namespace {
+SSL_CTX* getDriverSslCtxOrNull() {
+  static SSL_CTX* s_ctx = []() -> SSL_CTX* {
+    const char* env = std::getenv("FEEDSIM_DRIVER_TLS");
+    if (env == nullptr || std::strcmp(env, "1") != 0) {
+      return nullptr;
+    }
+    // OpenSSL >= 1.1 self-initializes; keep these calls as no-ops on
+    // older libs for forward compat.
+    SSL_library_init();
+    SSL_load_error_strings();
+    OpenSSL_add_all_algorithms();
+    SSL_CTX* ctx = SSL_CTX_new(TLS_client_method());
+    if (ctx == nullptr) {
+      std::cerr << "FeedSimDriver: SSL_CTX_new failed" << std::endl;
+      return nullptr;
+    }
+    SSL_CTX_set_verify(ctx, SSL_VERIFY_NONE, nullptr);
+    std::cout << "FeedSimDriver: TLS enabled via FEEDSIM_DRIVER_TLS=1"
+              << std::endl;
+    return ctx;
+  }();
+  return s_ctx;
+}
+} // namespace
+
 // ─── DriverConnection: one TCP connection to the server ─────────────────────
 
 class DriverConnection {
@@ -238,7 +274,33 @@ class DriverConnection {
       setsockopt(sockfd, IPPROTO_TCP, TCP_NODELAY, &optval, sizeof(optval));
     }
     evutil_make_socket_nonblocking(sockfd);
-    bev_ = bufferevent_socket_new(base, sockfd, BEV_OPT_CLOSE_ON_FREE);
+    SSL_CTX* ssl_ctx = getDriverSslCtxOrNull();
+    if (ssl_ctx != nullptr) {
+      // Per-connection SSL object owned by the bufferevent
+      // (BEV_OPT_CLOSE_ON_FREE will SSL_free it). State is CONNECTING
+      // because we already have an open TCP socket and need libevent to
+      // drive the TLS handshake from the client side.
+      SSL* ssl = SSL_new(ssl_ctx);
+      if (ssl == nullptr) {
+        std::cerr << "Error: SSL_new failed" << std::endl;
+        abort();
+      }
+      bev_ = bufferevent_openssl_socket_new(
+          base,
+          sockfd,
+          ssl,
+          BUFFEREVENT_SSL_CONNECTING,
+          BEV_OPT_CLOSE_ON_FREE);
+      if (bev_ == nullptr) {
+        std::cerr << "Error: bufferevent_openssl_socket_new failed"
+                  << std::endl;
+        SSL_free(ssl);
+        ::close(sockfd);
+        abort();
+      }
+    } else {
+      bev_ = bufferevent_socket_new(base, sockfd, BEV_OPT_CLOSE_ON_FREE);
+    }
   }
 
   ~DriverConnection() {

packages/feedsim/third_party/src/workloads/ranking/FeedSimServer.cc

@@ -27,11 +27,13 @@
 
 #include <folly/MPMCQueue.h>
 #include <folly/container/F14Map.h>
+#include <folly/io/async/AsyncSSLSocket.h>
 #include <folly/io/async/AsyncServerSocket.h>
 #include <folly/io/async/AsyncSocket.h>
 #include <folly/io/async/AsyncTransport.h>
 #include <folly/io/async/EventBase.h>
 #include <folly/io/async/EventBaseManager.h>
+#include <folly/io/async/SSLContext.h>
 
 namespace feedsim {
 
@@ -241,9 +243,11 @@ class AcceptCallback : public folly::AsyncServerSocket::AcceptCallback {
  public:
   AcceptCallback(
       std::vector<std::unique_ptr<WorkerThread>>& workers,
-      const folly::F14FastMap<uint32_t, QueryCallback>& callbacks)
+      const folly::F14FastMap<uint32_t, QueryCallback>& callbacks,
+      std::shared_ptr<folly::SSLContext> ssl_ctx)
       : workers_(workers),
         callbacks_(callbacks),
+        ssl_ctx_(std::move(ssl_ctx)),
         next_worker_(0) {}
 
   void connectionAccepted(
@@ -260,12 +264,29 @@ class AcceptCallback : public folly::AsyncServerSocket::AcceptCallback {
     auto& worker = workers_[next_worker_];
     next_worker_ = (next_worker_ + 1) % workers_.size();
 
-    // Create connection on the worker's EventBase
+    // Create connection on the worker's EventBase. When TLS is enabled
+    // (FEEDSIM_TLS_CERT/FEEDSIM_TLS_KEY env vars set at server startup),
+    // wrap the accepted fd in folly::AsyncSSLSocket so the TLS handshake
+    // is performed on the worker's EventBase before any wire reads. This
+    // closes the bench's Encryption CPU undershoot (prod ~3.3-3.6% vs
+    // bench ~0.9-1.5% in t41). Plain AsyncSocket preserves the original
+    // no-TLS behavior when the env vars are unset.
+    auto ssl_ctx = ssl_ctx_;
     worker->evb->runInEventBaseThread(
         [fd, thread_id = worker->thread_id, &callbacks = callbacks_,
-         evb = worker->evb.get()]() {
-          auto socket = folly::AsyncSocket::newSocket(
-              evb, folly::NetworkSocket::fromFd(fd));
+         evb = worker->evb.get(), ssl_ctx]() {
+          folly::AsyncSocket::UniquePtr socket;
+          if (ssl_ctx) {
+            // AsyncSSLSocket server-side: pass true for the server flag.
+            // The handshake is initiated lazily on first read/write,
+            // matching the existing client's connect-then-write pattern.
+            folly::AsyncSSLSocket::UniquePtr ssl_sock(new folly::AsyncSSLSocket(
+                ssl_ctx, evb, folly::NetworkSocket::fromFd(fd), true));
+            socket.reset(ssl_sock.release());
+          } else {
+            socket = folly::AsyncSocket::newSocket(
+                evb, folly::NetworkSocket::fromFd(fd));
+          }
           // ServerConnection self-manages its lifetime
           new ServerConnection(std::move(socket), thread_id, callbacks);
         });
@@ -278,6 +299,7 @@ class AcceptCallback : public folly::AsyncServerSocket::AcceptCallback {
  private:
   std::vector<std::unique_ptr<WorkerThread>>& workers_;
   const folly::F14FastMap<uint32_t, QueryCallback>& callbacks_;
+  std::shared_ptr<folly::SSLContext> ssl_ctx_;
   size_t next_worker_;
 };
 
@@ -298,6 +320,11 @@ struct FeedSimServer::Impl {
   std::unique_ptr<AcceptCallback> accept_cb;
   std::vector<std::unique_ptr<WorkerThread>> workers;
 
+  // Optional SSL context — populated at startup from FEEDSIM_TLS_CERT /
+  // FEEDSIM_TLS_KEY env vars. When set, accepted connections are wrapped
+  // in AsyncSSLSocket. Driver-side env: FEEDSIM_DRIVER_TLS=1.
+  std::shared_ptr<folly::SSLContext> ssl_ctx;
+
   std::atomic<bool> running{false};
 };
 
@@ -390,8 +417,38 @@ void FeedSimServer::run() {
   impl_->main_evb = std::make_unique<folly::EventBase>();
   impl_->server_socket = folly::AsyncServerSocket::newSocket(impl_->main_evb.get());
 
+  // Optional TLS: when FEEDSIM_TLS_CERT and FEEDSIM_TLS_KEY are both set,
+  // create an SSL context and pass it through to AcceptCallback, which
+  // wraps each accepted fd in AsyncSSLSocket. Reuses the same cert/key
+  // pair as mock_services (FEEDSIM_TLS_CERT defaults to
+  // ${FEEDSIM_ROOT}/certs/example.crt via run.sh wiring). Bench-only:
+  // peer cert verification is not configured here; the client side
+  // (FeedSimDriver) skips verify via SSL_VERIFY_NONE.
+  {
+    const char* cert_env = std::getenv("FEEDSIM_TLS_CERT");
+    const char* key_env = std::getenv("FEEDSIM_TLS_KEY");
+    if (cert_env != nullptr && key_env != nullptr && cert_env[0] != '\0' &&
+        key_env[0] != '\0') {
+      try {
+        auto ctx = std::make_shared<folly::SSLContext>();
+        ctx->loadCertificate(cert_env);
+        ctx->loadPrivateKey(key_env);
+        // No ALPN — FeedSim uses its own custom binary protocol over the
+        // TLS-wrapped socket, not Rocket. The client similarly does not
+        // advertise ALPN.
+        impl_->ssl_ctx = std::move(ctx);
+        std::cout << "FeedSimServer: TLS enabled (cert=" << cert_env
+                  << " key=" << key_env << ")" << std::endl;
+      } catch (const std::exception& e) {
+        std::cerr << "FeedSimServer: failed to init SSL context: " << e.what()
+                  << " — falling back to plaintext" << std::endl;
+        impl_->ssl_ctx.reset();
+      }
+    }
+  }
+
   impl_->accept_cb = std::make_unique<AcceptCallback>(
-      impl_->workers, impl_->query_callbacks);
+      impl_->workers, impl_->query_callbacks, impl_->ssl_ctx);
 
   impl_->server_socket->addAcceptCallback(
       impl_->accept_cb.get(), nullptr);

packages/feedsim/third_party/src/workloads/ranking/LeafNodeRank.cc

@@ -809,6 +809,17 @@ std::string compressPayload(const std::string& data, int /*result*/) {
   folly::StringPiece output(
       data.data(),
       std::min(args.compression_data_size_arg, args.random_data_size_arg));
+  // Honor FEEDSIM_SERVER_ZSTD passthrough (see comment above
+  // compressThrift). For symmetry; this path is hit on the random-payload
+  // workload variant. kServerZstdLocal mirrors the global once-per-process
+  // read in compressThrift's kServerZstd.
+  static const bool kServerZstdLocal = []() {
+    const char* v = std::getenv("FEEDSIM_SERVER_ZSTD");
+    return v == nullptr || std::strcmp(v, "0") != 0;
+  }();
+  if (!kServerZstdLocal) {
+    return std::string(output.data(), output.size());
+  }
 #ifdef BENCHPRESS_INTERNAL
   return getRandomStringCodec()->compress(output);
 #else
@@ -830,8 +841,27 @@ std::string decompressPayload(const std::string& data) {
 #endif
 }
 
+// FEEDSIM_SERVER_ZSTD env gate: when "0", server-side response compression
+// is bypassed (passthrough). Default behavior (ZSTD on) preserved. Lets the
+// bench match prod's lighter Compression CPU share (prod 3.9-5.2% vs bench
+// 6-10% in t41) without ripping out the entire ZSTD path. Decision is read
+// once per process to avoid getenv() on every request — kServerZstd is
+// shared with compressPayload below.
+namespace {
+bool readServerZstdEnv() {
+  const char* v = std::getenv("FEEDSIM_SERVER_ZSTD");
+  return v == nullptr || std::strcmp(v, "0") != 0;
+}
+} // namespace
+
 std::unique_ptr<folly::IOBuf> compressThrift(
     std::unique_ptr<folly::IOBuf> buf) {
+  static const bool kServerZstd = readServerZstdEnv();
+  LOG_FIRST_N(INFO, 1) << "Server-side response ZSTD: "
+                       << (kServerZstd ? "ON" : "OFF (passthrough)");
+  if (!kServerZstd) {
+    return buf;
+  }
 #ifdef BENCHPRESS_INTERNAL
   return getThriftPayloadCodec()->compress(buf.get());
 #else