Skip to content

Commit 7981dcc

Browse files
excelle08meta-codesync[bot]
excelle08
authored and
meta-codesync[bot]
committed
Enable TLS-over-Rocket on mock_services with ALPN rs (facebookresearch#710)
Summary: Pull Request resolved: facebookresearch#710 Closes the −9.9pp `RPC-AsyncIO` and −3.3pp `Encryption` CPU gaps observed in t31v5 BGM (vs prod `multifeed/aggregator_main`). The prior attempt in this stack omitted TLS because `folly::AsyncSSLSocket` + `RocketClientChannel::newChannel(...)` was failing the probe RPC with timeout; root cause was missing ALPN `"rs"` advertisement (required by ThriftServer with `setSSLPolicy::REQUIRED` to route the connection into the Rocket transport — without it, the server rejects the Rocket setup frame because the TLS-layer protocol selection picked the wrong path). This change adds the missing ALPN bits on both sides and matches the canonical pattern at `thrift/lib/cpp2/test/server/ThriftServerTest.cpp` `RocketOverSSLNoALPN` in the pinned fbthrift v2026.01.05.00. Server (`MockServiceMain.cc`): - New `--tls_cert` / `--tls_key` gflags. When both set, build a `wangle::SSLContextConfig`, set the cert/key, set `clientVerification = DO_NOT_REQUEST` (no peer cert needed), and **call `setNextProtocols({"rs"})`** so the server advertises Rocket-over-TLS via ALPN. - `server->setSSLConfig(sslCfg)` + `server->setSSLPolicy(REQUIRED)`. Client (`MockServicesClient.cc`): - Read `MOCK_TLS` env var (LeafNodeRank uses gengetopt for CLI, rejects unknown flags, so env-var is the only viable injection path). - When set: build a `folly::SSLContext` with `authenticate(false, false)` + `setVerificationOption(NO_VERIFY)` (self-signed cert) + **`setAdvertisedNextProtocols({"rs"})`**. - Construct `folly::AsyncSSLSocket::UniquePtr` from that context, call `ssl_sock->connect(nullptr, addr)`, then `socket = std::move(ssl_sock)` so the UniquePtr conversion to `AsyncSocket::UniquePtr` preserves the SSL type. Hand to `RocketClientChannel::newChannel(...)`. AsyncSSLSocket buffers writes until the TLS handshake completes, so Rocket's setup frame goes out only after TLS is up. Shell-script plumbing: - `run-feedsim-multi.sh`: when `FEEDSIM_TLS=1`, pass `--tls_cert=${FEEDSIM_ROOT}/certs/example.crt --tls_key=...` to each mock_services instance. Fail fast if certs are missing. - `run.sh`: when `FEEDSIM_TLS=1`, export `MOCK_TLS=1` so LeafNodeRank's child env inherits it. Cert plumbing was already in place from the prior diff in this stack (install_feedsim.sh extracts `packages/common/certs.tar.gz` with `--strip-components=1` to `${FEEDSIM_ROOT}/certs/example.{crt,key}`). Reviewed By: YifanYuan3 Differential Revision: D107327606
1 parent dc1f316 commit 7981dcc

5 files changed

Lines changed: 125 additions & 48 deletions

File tree

‎packages/feedsim/run-feedsim-multi.sh‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -164,7 +164,21 @@ function start_mock_services() {
164164
local offset_us="${MOCK_LATENCY_OFFSET_US:-0}"
165165
local skip_us="${MOCK_LATENCY_SKIP_THRESHOLD_US:-100}"
166166

167-
echo "Starting mock_services on port ${port} (cores=${core_range}, io_threads=${io_threads}, cap_us=${cap_us}, offset_us=${offset_us}, skip_us=${skip_us}, silesia=${SILESIA_DIR_ABS})"
167+
# TLS opt-in: FEEDSIM_TLS=1 enables TLS on both server and client.
168+
# Server reads --tls_cert / --tls_key; client picks up MOCK_TLS env var
169+
# (LeafNodeRank uses gengetopt and rejects unknown CLI flags).
170+
local tls_opts=""
171+
if [ "${FEEDSIM_TLS:-0}" = "1" ]; then
172+
local cert_dir="${FEEDSIM_ROOT}/certs"
173+
if [ ! -r "${cert_dir}/example.crt" ] || [ ! -r "${cert_dir}/example.key" ]; then
174+
echo "ERROR: FEEDSIM_TLS=1 but ${cert_dir}/example.{crt,key} not found" >&2
175+
exit 1
176+
fi
177+
tls_opts="--tls_cert=${cert_dir}/example.crt --tls_key=${cert_dir}/example.key"
178+
fi
179+
180+
echo "Starting mock_services on port ${port} (cores=${core_range}, io_threads=${io_threads}, cap_us=${cap_us}, offset_us=${offset_us}, skip_us=${skip_us}, tls=${FEEDSIM_TLS:-0}, silesia=${SILESIA_DIR_ABS})"
181+
# shellcheck disable=SC2086
168182
taskset --cpu-list "$core_range" \
169183
"$MOCK_SERVICES_BIN" \
170184
--port="$port" \
@@ -173,6 +187,7 @@ function start_mock_services() {
173187
--latency_offset_us="$offset_us" \
174188
--latency_skip_threshold_us="$skip_us" \
175189
--silesia_dir="$SILESIA_DIR_ABS" \
190+
$tls_opts \
176191
> "$log_path" 2>&1 &
177192
local pid=$!
178193
MOCK_SERVICES_PIDS+=("$pid")

‎packages/feedsim/run.sh‎

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -729,12 +729,16 @@ main() {
729729
mock_services_opts="$mock_services_opts --mock_keepalive_interval_ms=${MOCK_KEEPALIVE_INTERVAL_MS}"
730730
echo "MockServicesClient keepalive: ENABLED (interval=${MOCK_KEEPALIVE_INTERVAL_MS} ms)"
731731
fi
732-
# Wire compression on the outbound MockServicesClient channel.
733-
# LeafNodeRank uses gengetopt (rejects unknown flags), so the knob is
734-
# plumbed via the MOCK_COMPRESS_ZSTD env var read inside
735-
# MockServicesClient.cc — not as a CLI flag. FEEDSIM_NO_RPC_ZSTD=1 turns
736-
# ZSTD off (binary default is on); leave unset for default-on parity with
737-
# prod's wire-compressed channels.
732+
# TLS + wire compression on the outbound MockServicesClient channel.
733+
# LeafNodeRank uses gengetopt (rejects unknown CLI flags), so these knobs
734+
# are plumbed via env vars MOCK_TLS / MOCK_COMPRESS_ZSTD read inside
735+
# MockServicesClient.cc. FEEDSIM_TLS=1 must match the server-side
736+
# --tls_cert/--tls_key wiring in run-feedsim-multi.sh. FEEDSIM_NO_RPC_ZSTD=1
737+
# disables ZSTD (binary default is on); leave unset for prod-parity.
738+
if [ "${FEEDSIM_TLS:-0}" = "1" ]; then
739+
export MOCK_TLS=1
740+
echo "MockServicesClient TLS: ENABLED (via MOCK_TLS env)"
741+
fi
738742
if [ "${FEEDSIM_NO_RPC_ZSTD:-0}" = "1" ]; then
739743
export MOCK_COMPRESS_ZSTD=0
740744
echo "MockServicesClient ZSTD: DISABLED (via MOCK_COMPRESS_ZSTD env)"

‎packages/feedsim/third_party/src/workloads/ranking/MockServicesClient.cc‎

Lines changed: 31 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,9 @@
2424
#include <cstring>
2525

2626
#include <folly/SocketAddress.h>
27+
#include <folly/io/async/AsyncSSLSocket.h>
2728
#include <folly/io/async/AsyncSocket.h>
29+
#include <folly/io/async/SSLContext.h>
2830

2931
#include <thrift/lib/cpp2/async/RocketClientChannel.h>
3032
#include <thrift/lib/thrift/gen-cpp2/RpcMetadata_types.h>
@@ -33,8 +35,10 @@
3335

3436
namespace {
3537
// LeafNodeRank uses gengetopt (not gflags) for CLI parsing, so we cannot add
36-
// CLI flags here. Read knobs from env vars instead. Set MOCK_COMPRESS_ZSTD=0
37-
// to disable per-channel ZSTD negotiation (default on).
38+
// CLI flags here. Read knobs from env vars instead. Set MOCK_TLS=1 to use
39+
// AsyncSSLSocket (with ALPN "rs" so the server routes the connection into
40+
// the Rocket transport). Set MOCK_COMPRESS_ZSTD=0 to disable per-channel
41+
// ZSTD negotiation (default on).
3842
bool envBoolTrue(const char* name, bool default_value) {
3943
const char* v = std::getenv(name);
4044
if (v == nullptr) {
@@ -118,11 +122,33 @@ MockServicesClient::MockServicesClient(
118122
// RocketClientChannel must be created on the EventBase thread. Use
119123
// runInEventBaseThreadAndWait so this constructor remains usable from
120124
// any thread (typically the main thread during ThreadStartup).
125+
const bool use_tls = envBoolTrue("MOCK_TLS", false);
121126
const bool use_zstd = envBoolTrue("MOCK_COMPRESS_ZSTD", true);
122-
evb_->runInEventBaseThreadAndWait([this, &host, port, use_zstd]() {
127+
evb_->runInEventBaseThreadAndWait([this, &host, port, use_tls, use_zstd]() {
123128
folly::SocketAddress addr(host, port, /*allowNameLookup=*/true);
124-
folly::AsyncSocket::UniquePtr socket(
125-
new folly::AsyncSocket(evb_, addr));
129+
folly::AsyncSocket::UniquePtr socket;
130+
if (use_tls) {
131+
auto ssl_ctx = std::make_shared<folly::SSLContext>();
132+
// Benchmark-only: skip peer cert verification so example certs work.
133+
ssl_ctx->authenticate(/*checkPeerCert=*/false, /*checkPeerName=*/false);
134+
ssl_ctx->setVerificationOption(
135+
folly::SSLContext::SSLVerifyPeerEnum::NO_VERIFY);
136+
// ALPN "rs" tells the server to route this TLS connection into the
137+
// Rocket transport (matches `RocketOverSSLNoALPN` test pattern in
138+
// pinned fbthrift v2026.01.05.00). Without ALPN, the server may
139+
// reject the connection or fall back to the header-upgrade path.
140+
ssl_ctx->setAdvertisedNextProtocols({"rs"});
141+
folly::AsyncSSLSocket::UniquePtr ssl_sock(
142+
new folly::AsyncSSLSocket(ssl_ctx, evb_));
143+
// AsyncSSLSocket buffers writes until the TLS handshake completes,
144+
// so RocketClientChannel's setup frame goes out only AFTER TLS is
145+
// negotiated. Passing nullptr connect-callback is the canonical
146+
// fbthrift pattern (see RocketOverSSLNoALPN in ThriftServerTest.cpp).
147+
ssl_sock->connect(/*callback=*/nullptr, addr);
148+
socket = std::move(ssl_sock);
149+
} else {
150+
socket.reset(new folly::AsyncSocket(evb_, addr));
151+
}
126152
auto channel =
127153
apache::thrift::RocketClientChannel::newChannel(std::move(socket));
128154
if (use_zstd) {

‎packages/feedsim/third_party/src/workloads/ranking/mock_services/MockServiceHandler.cc‎

Lines changed: 28 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -29,9 +29,34 @@
2929

3030
#include "LatencyHistogram.h"
3131

32-
DECLARE_int32(latency_cap_us);
33-
DECLARE_int32(latency_offset_us);
34-
DECLARE_int32(latency_skip_threshold_us);
32+
// rpc_dist.json contains very long-tail latencies (p99 ~8s, max ~28s)
33+
// because the production aggregator's outbound RPCs occasionally hit
34+
// retries / queueing / GC pauses; reproducing those raw values inside
35+
// the mock makes the bench unusable (every fanout blocks on the
36+
// slowest call). The flags below let us shape the simulated delay so
37+
// it stays within a useful operating range without losing the
38+
// per-method latency mix. Defined here (where they're consumed) so the
39+
// mock_service_handler cpp_library links cleanly on its own; declared
40+
// in MockServiceMain.cc for the startup LOG line.
41+
DEFINE_int32(
42+
latency_cap_us,
43+
200000,
44+
"Cap the requested per-call latency at this value (us). Defaults to "
45+
"200ms so the slowest per-RPC sleep is bounded; set to 0 to disable.");
46+
DEFINE_int32(
47+
latency_offset_us,
48+
0,
49+
"Subtract this many us from the requested latency to compensate for "
50+
"intrinsic RPC-stack overhead the client already pays for. Defaults "
51+
"to 0; tune empirically based on the gap between the leaf-side "
52+
"dispatch_per_rpc histogram and the mock-handler actual latency.");
53+
DEFINE_int32(
54+
latency_skip_threshold_us,
55+
100,
56+
"If the effective latency (after cap + offset) is below this many us, "
57+
"skip the spin/sleep entirely and respond immediately. Avoids paying "
58+
"for spin-wait jitter on requests whose modeled budget is already "
59+
"comparable to the natural RPC round trip.");
3560

3661
namespace mock_services {
3762

‎packages/feedsim/third_party/src/workloads/ranking/mock_services/MockServiceMain.cc‎

Lines changed: 40 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,8 @@
2121
#include <glog/logging.h>
2222

2323
#include <folly/init/Init.h>
24+
#include <folly/io/async/SSLContext.h>
25+
#include <wangle/ssl/SSLContextConfig.h>
2426

2527
#include "thrift/lib/cpp2/server/ThriftServer.h"
2628

@@ -45,38 +47,23 @@ DEFINE_string(
4547
"",
4648
"Required. Path to the Silesia corpus directory used for response bytes.");
4749

48-
// rpc_dist.json contains very long-tail latencies (p99 ~8s, max ~28s)
49-
// because the production aggregator's outbound RPCs occasionally hit
50-
// retries / queueing / GC pauses; reproducing those raw values inside
51-
// the mock makes the bench unusable (every fanout blocks on the
52-
// slowest call). The flags below let us shape the simulated delay so
53-
// it stays within a useful operating range without losing the
54-
// per-method latency mix.
55-
DEFINE_int32(
56-
latency_cap_us,
57-
200000,
58-
"Cap the requested per-call latency at this value (us). Defaults to "
59-
"200ms so the slowest per-RPC sleep is bounded; set to 0 to disable.");
60-
DEFINE_int32(
61-
latency_offset_us,
62-
0,
63-
"Subtract this many us from the requested latency to compensate for "
64-
"intrinsic RPC-stack overhead the client already pays for. Defaults "
65-
"to 0; tune empirically based on the gap between the leaf-side "
66-
"dispatch_per_rpc histogram and the mock-handler actual latency.");
67-
DEFINE_int32(
68-
latency_skip_threshold_us,
69-
100,
70-
"If the effective latency (after cap + offset) is below this many us, "
71-
"skip the spin/sleep entirely and respond immediately. Avoids paying "
72-
"for spin-wait jitter on requests whose modeled budget is already "
73-
"comparable to the natural RPC round trip.");
74-
75-
// NOTE: TLS support intentionally not wired here. The open-source FBThrift
76-
// v2026.01.05.00 Rocket transport requires a fizz-based AsyncTransport (not
77-
// plain AsyncSSLSocket) on the client side; without it, TLS-REQUIRED servers
78-
// reject the first Rocket setup frame before SSL handshake even begins. See
79-
// the Encryption gap analysis in the t31 memo for the closure plan.
50+
// Latency-shaping flags are DEFINE'd in MockServiceHandler.cc (where they're
51+
// consumed); declare them here so the main-side LOG line can read them.
52+
DECLARE_int32(latency_cap_us);
53+
DECLARE_int32(latency_offset_us);
54+
DECLARE_int32(latency_skip_threshold_us);
55+
56+
// TLS knob. When --tls_cert and --tls_key are both set, the server requires
57+
// TLS on every inbound connection and advertises ALPN "rs" so RocketClient
58+
// transports can negotiate Rocket-over-TLS at handshake time (matches the
59+
// canonical pattern at thrift/lib/cpp2/test/server/ThriftServerTest.cpp
60+
// `RocketOverSSLNoALPN` in fbthrift v2026.01.05.00). Closes prod
61+
// multifeed/aggregator_main's Encryption CPU footprint (~3% on BGM).
62+
DEFINE_string(
63+
tls_cert,
64+
"",
65+
"Path to TLS cert PEM. Empty disables TLS (plaintext sockets).");
66+
DEFINE_string(tls_key, "", "Path to TLS key PEM. Required when --tls_cert set.");
8067

8168
int main(int argc, char** argv) {
8269
folly::Init init(&argc, &argv);
@@ -102,11 +89,31 @@ int main(int argc, char** argv) {
10289
server->setPort(FLAGS_port);
10390
server->setNumIOWorkerThreads(io_threads);
10491

92+
bool tls_enabled = false;
93+
if (!FLAGS_tls_cert.empty() && !FLAGS_tls_key.empty()) {
94+
auto sslCfg = std::make_shared<wangle::SSLContextConfig>();
95+
sslCfg->setCertificate(FLAGS_tls_cert, FLAGS_tls_key, "");
96+
sslCfg->clientVerification =
97+
folly::SSLContext::VerifyClientCertificate::DO_NOT_REQUEST;
98+
// Advertise ALPN "rs" so RocketClient transports negotiate
99+
// Rocket-over-TLS in the handshake. Without this, REQUIRED servers
100+
// can reject the connection or fall back to the header-upgrade path
101+
// which never speaks Rocket on TLS.
102+
sslCfg->setNextProtocols({"rs"});
103+
server->setSSLConfig(sslCfg);
104+
server->setSSLPolicy(apache::thrift::SSLPolicy::REQUIRED);
105+
tls_enabled = true;
106+
} else if (!FLAGS_tls_cert.empty() || !FLAGS_tls_key.empty()) {
107+
LOG(ERROR) << "Both --tls_cert and --tls_key must be set together";
108+
return EXIT_FAILURE;
109+
}
110+
105111
LOG(INFO) << "mock_services listening on port " << FLAGS_port
106112
<< " with " << io_threads << " IO worker threads"
107113
<< "; Silesia corpus from " << FLAGS_silesia_dir
108114
<< " (" << silesia->numFiles() << " files, "
109-
<< (silesia->totalSize() / (1024 * 1024)) << " MB)";
115+
<< (silesia->totalSize() / (1024 * 1024)) << " MB)"
116+
<< "; tls=" << (tls_enabled ? "on" : "off");
110117

111118
LOG(INFO) << "latency shaping: cap_us=" << FLAGS_latency_cap_us
112119
<< " offset_us=" << FLAGS_latency_offset_us

0 commit comments

Comments
 (0)