You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by vi...@apache.org on 2026/09/01 07:15:18 UTC
(spark-connect-gateway) branch main updated: [SPARK-59103] Add scg-proxy crate
This is an automated email from the ASF dual-hosted git repository.
viirya pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/spark-connect-gateway.git
The following commit(s) were added to refs/heads/main by this push:
new 52f1dc0 [SPARK-59103] Add scg-proxy crate
52f1dc0 is described below
commit 52f1dc0d7c7cf73db9c5d1fd602931459437d3fe
Author: Liang-Chi Hsieh <vi...@gmail.com>
AuthorDate: Tue Sep 1 00:15:11 2026 -0700
[SPARK-59103] Add scg-proxy crate
### What changes were proposed in this pull request?
Part of importing the Spark Connect Gateway reference implementation into
`apache/spark-connect-gateway` as a series of small, bottom-up PRs (tracked in
the do-not-merge reference PR #1).
This adds the `scg-proxy` crate — the gRPC proxy core of the gateway, and the
crate that ties the previously imported layers together. It contains:
- **handler** — a `SparkConnectService` implementation that forwards every RPC to
a backend chosen by routing + session affinity, stamping the authenticated
identity onto the forwarded request
- **dial** — a per-backend tonic `Channel` cache, so the dial cost isn't paid per
RPC
- **outbound** — the credentials the gateway presents on the gateway→backend hop
- **config_filter** — withholds backend-only secrets from `Config` RPC responses
`scg-proxy` depends on `scg-routing`, `scg-auth`, `scg-genproto`,
`scg-observability`, `scg-tenant`, `scg-ratelimit`, and `scg-audit`. Its examples
and integration tests additionally use `scg-store-memory`, `scg-store-redis`,
`scg-pool-static`, and `scg-healthcheck` as dev-dependencies — all of which are
already in the workspace, which is why this PR lands after them.
It is imported complete: src, all integration tests (forwarding, auth, audit,
tenant, ratelimit, observability/tracing, multi-tenant e2e, clone-session,
outbound-token), and all examples (ha_smoke, drain_smoke, load, e2e_client,
fake_collector).
### Why are the changes needed?
Implements the reference gateway from the approved SPIP (SPARK-58455), imported
one reviewable piece at a time.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
`cargo fmt --all --check`, `cargo clippy --workspace --all-targets -- -D warnings`,
`cargo build --workspace`, and `cargo test --workspace` all pass locally: 201
tests pass, 10 Redis/testcontainer tests are `#[ignore]` so no Docker is needed
in CI. Note `--all-targets` compiles all of this crate's examples and integration
test binaries, which is why its four dev-dependency crates had to land first.
### Was this patch authored or co-authored using generative AI tooling?
Yes, co-authored with Claude Code.
Closes #14 from viirya/import-11-proxy.
Authored-by: Liang-Chi Hsieh <vi...@gmail.com>
Signed-off-by: Liang-Chi Hsieh <li...@databricks.com>
---
Cargo.lock | 51 ++
Cargo.toml | 1 +
crates/proxy/Cargo.toml | 70 ++
crates/proxy/examples/drain_smoke.rs | 316 +++++++
crates/proxy/examples/e2e_client.rs | 74 ++
crates/proxy/examples/fake_collector.rs | 169 ++++
crates/proxy/examples/ha_smoke.rs | 390 +++++++++
crates/proxy/examples/load.rs | 1008 ++++++++++++++++++++++
crates/proxy/src/config_filter.rs | 369 ++++++++
crates/proxy/src/dial.rs | 71 ++
crates/proxy/src/handler.rs | 1041 +++++++++++++++++++++++
crates/proxy/src/lib.rs | 41 +
crates/proxy/src/outbound.rs | 179 ++++
crates/proxy/tests/audit_integration.rs | 458 ++++++++++
crates/proxy/tests/auth_integration.rs | 296 +++++++
crates/proxy/tests/clone_session.rs | 335 ++++++++
crates/proxy/tests/common/mod.rs | 161 ++++
crates/proxy/tests/forwarding.rs | 387 +++++++++
crates/proxy/tests/multitenant_e2e.rs | 616 ++++++++++++++
crates/proxy/tests/multitenant_pool.rs | 305 +++++++
crates/proxy/tests/observability_integration.rs | 321 +++++++
crates/proxy/tests/outbound_token.rs | 577 +++++++++++++
crates/proxy/tests/ratelimit_integration.rs | 358 ++++++++
crates/proxy/tests/tenant_integration.rs | 401 +++++++++
crates/proxy/tests/tracing_integration.rs | 376 ++++++++
25 files changed, 8371 insertions(+)
diff --git a/Cargo.lock b/Cargo.lock
index 9f0b511..8660c5a 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -327,6 +327,12 @@ version = "3.20.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649"
+[[package]]
+name = "byteorder"
+version = "1.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b"
+
[[package]]
name = "bytes"
version = "1.12.1"
@@ -935,6 +941,16 @@ version = "0.17.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a"
+[[package]]
+name = "hdrhistogram"
+version = "7.6.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f49d1053f4708f0af3cf9fc5bffc7e68a914a3c45becb231c80068c9c3f78bea"
+dependencies = [
+ "byteorder",
+ "num-traits",
+]
+
[[package]]
name = "heck"
version = "0.5.0"
@@ -2695,6 +2711,41 @@ dependencies = [
"thiserror 2.0.20",
]
+[[package]]
+name = "scg-proxy"
+version = "0.1.0"
+dependencies = [
+ "async-stream",
+ "async-trait",
+ "futures",
+ "hdrhistogram",
+ "opentelemetry",
+ "opentelemetry-proto",
+ "opentelemetry_sdk",
+ "parking_lot",
+ "prost",
+ "reqwest",
+ "scg-audit",
+ "scg-auth",
+ "scg-genproto",
+ "scg-healthcheck",
+ "scg-observability",
+ "scg-pool-static",
+ "scg-ratelimit",
+ "scg-routing",
+ "scg-store-memory",
+ "scg-store-redis",
+ "scg-tenant",
+ "thiserror 2.0.20",
+ "tokio",
+ "tokio-stream",
+ "tonic",
+ "tower",
+ "tracing",
+ "tracing-opentelemetry",
+ "tracing-subscriber",
+]
+
[[package]]
name = "scg-ratelimit"
version = "0.1.0"
diff --git a/Cargo.toml b/Cargo.toml
index 0528ae0..01e0443 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -29,6 +29,7 @@ members = [
"crates/store-redis",
"crates/healthcheck",
"crates/pool-k8s",
+ "crates/proxy",
]
[workspace.package]
diff --git a/crates/proxy/Cargo.toml b/crates/proxy/Cargo.toml
new file mode 100644
index 0000000..6a4e594
--- /dev/null
+++ b/crates/proxy/Cargo.toml
@@ -0,0 +1,70 @@
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+[package]
+name = "scg-proxy"
+version.workspace = true
+edition.workspace = true
+license.workspace = true
+authors.workspace = true
+repository.workspace = true
+readme.workspace = true
+keywords.workspace = true
+categories.workspace = true
+description = "gRPC SparkConnectService handler that forwards each RPC to a backend chosen by the router."
+publish = false
+
+[dependencies]
+scg-auth = { path = "../auth" }
+scg-genproto = { path = "../genproto" }
+scg-observability = { path = "../observability" }
+scg-routing = { path = "../routing" }
+scg-tenant = { path = "../tenant" }
+scg-ratelimit = { path = "../ratelimit" }
+scg-audit = { path = "../audit" }
+tokio = { workspace = true }
+tokio-stream = { workspace = true }
+tonic = { workspace = true }
+tower = { workspace = true }
+prost = { workspace = true }
+async-trait = { workspace = true }
+futures = { workspace = true }
+parking_lot = { workspace = true }
+tracing = { workspace = true }
+tracing-opentelemetry = { workspace = true }
+opentelemetry = { workspace = true }
+async-stream = "0.3"
+thiserror = { workspace = true }
+
+[dev-dependencies]
+scg-audit = { path = "../audit" }
+scg-store-memory = { path = "../store-memory" }
+scg-store-redis = { path = "../store-redis" }
+scg-pool-static = { path = "../pool-static" }
+scg-auth = { path = "../auth" }
+scg-tenant = { path = "../tenant" }
+scg-ratelimit = { path = "../ratelimit" }
+scg-observability = { path = "../observability", features = ["testing"] }
+tokio = { workspace = true }
+async-stream = "0.3"
+parking_lot = { workspace = true }
+reqwest = { workspace = true }
+tracing-subscriber = { workspace = true }
+opentelemetry = { workspace = true }
+opentelemetry_sdk = { workspace = true, features = ["rt-tokio", "testing"] }
+opentelemetry-proto = { workspace = true }
+scg-healthcheck = { path = "../healthcheck" }
+hdrhistogram = { workspace = true }
+
diff --git a/crates/proxy/examples/drain_smoke.rs b/crates/proxy/examples/drain_smoke.rs
new file mode 100644
index 0000000..95205e3
--- /dev/null
+++ b/crates/proxy/examples/drain_smoke.rs
@@ -0,0 +1,316 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Graceful-shutdown smoke test.
+//!
+//! Exercises the two-step drain logic the gateway main does on
+//! SIGTERM, but in-process (no actual signals — drain is triggered
+//! by sending on a channel the same way `shutdown_signal()` does in
+//! production).
+//!
+//! Verifies:
+//!
+//! 1. Before drain, a long-running ExecutePlan stream is in
+//! progress.
+//! 2. Drain is triggered → readiness flips to not-ready immediately.
+//! 3. The in-flight stream continues to receive messages and
+//! eventually completes normally — it is *not* killed.
+//! 4. After the stream completes, the gRPC server actually shuts
+//! down (i.e. the drain loop didn't deadlock).
+//!
+//! Run with:
+//! ```bash
+//! cargo run -p scg-proxy --example drain_smoke
+//! ```
+//!
+//! No external dependencies (no Redis, no Docker).
+
+use std::pin::Pin;
+use std::sync::Arc;
+use std::time::Duration;
+
+use futures::{Stream, StreamExt};
+use scg_auth::{AnonymousAuthenticator, AuthInterceptor};
+use scg_genproto::pb;
+use scg_observability::{Metrics, ReadinessProbe};
+use scg_pool_static::StaticPool;
+use scg_proxy::{Dialer, SparkConnectProxy};
+use scg_routing::{AffinityStore, Pool, Router};
+use scg_store_memory::MemoryStore;
+use tokio::net::TcpListener;
+use tokio::sync::{oneshot, watch};
+use tokio_stream::wrappers::TcpListenerStream;
+use tonic::transport::{Channel, Endpoint, Server};
+use tonic::{Request, Response, Status};
+
+/// Backend that streams 5 ExecutePlan responses with a small delay
+/// between each, totalling ~1.5s. Long enough that we can:
+/// - start the stream,
+/// - trigger drain mid-stream,
+/// - observe drain *waits* for it,
+/// - then see the stream finish.
+struct SlowBackend;
+
+#[tonic::async_trait]
+impl pb::spark_connect_service_server::SparkConnectService for SlowBackend {
+ type ExecutePlanStream =
+ Pin<Box<dyn Stream<Item = Result<pb::ExecutePlanResponse, Status>> + Send + 'static>>;
+ type ReattachExecuteStream = Self::ExecutePlanStream;
+
+ async fn execute_plan(
+ &self,
+ req: Request<pb::ExecutePlanRequest>,
+ ) -> Result<Response<Self::ExecutePlanStream>, Status> {
+ let body = req.into_inner();
+ let stream = async_stream::stream! {
+ for i in 0..5 {
+ tokio::time::sleep(Duration::from_millis(300)).await;
+ yield Ok(pb::ExecutePlanResponse {
+ session_id: body.session_id.clone(),
+ operation_id: body.operation_id.clone().unwrap_or_default(),
+ response_id: format!("msg-{}", i),
+ ..Default::default()
+ });
+ }
+ };
+ Ok(Response::new(Box::pin(stream)))
+ }
+
+ async fn config(
+ &self,
+ _: Request<pb::ConfigRequest>,
+ ) -> Result<Response<pb::ConfigResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn analyze_plan(
+ &self,
+ _: Request<pb::AnalyzePlanRequest>,
+ ) -> Result<Response<pb::AnalyzePlanResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn artifact_status(
+ &self,
+ _: Request<pb::ArtifactStatusesRequest>,
+ ) -> Result<Response<pb::ArtifactStatusesResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn interrupt(
+ &self,
+ _: Request<pb::InterruptRequest>,
+ ) -> Result<Response<pb::InterruptResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn release_execute(
+ &self,
+ _: Request<pb::ReleaseExecuteRequest>,
+ ) -> Result<Response<pb::ReleaseExecuteResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn release_session(
+ &self,
+ _: Request<pb::ReleaseSessionRequest>,
+ ) -> Result<Response<pb::ReleaseSessionResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn fetch_error_details(
+ &self,
+ _: Request<pb::FetchErrorDetailsRequest>,
+ ) -> Result<Response<pb::FetchErrorDetailsResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn clone_session(
+ &self,
+ _: Request<pb::CloneSessionRequest>,
+ ) -> Result<Response<pb::CloneSessionResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn get_status(
+ &self,
+ _: Request<pb::GetStatusRequest>,
+ ) -> Result<Response<pb::GetStatusResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn reattach_execute(
+ &self,
+ _: Request<pb::ReattachExecuteRequest>,
+ ) -> Result<Response<Self::ReattachExecuteStream>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn add_artifacts(
+ &self,
+ _: Request<tonic::Streaming<pb::AddArtifactsRequest>>,
+ ) -> Result<Response<pb::AddArtifactsResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+}
+
+#[tokio::main]
+async fn main() {
+ // 1. Backend
+ let lis = TcpListener::bind("127.0.0.1:0").await.unwrap();
+ let be_addr = lis.local_addr().unwrap().to_string();
+ let (be_kill_tx, be_kill_rx) = oneshot::channel();
+ tokio::spawn(async move {
+ Server::builder()
+ .add_service(
+ pb::spark_connect_service_server::SparkConnectServiceServer::new(SlowBackend),
+ )
+ .serve_with_incoming_shutdown(TcpListenerStream::new(lis), async {
+ let _ = be_kill_rx.await;
+ })
+ .await
+ .ok();
+ });
+
+ // 2. Gateway with the same drain plumbing main uses.
+ let metrics = Metrics::new().unwrap();
+ let readiness = ReadinessProbe::new(true);
+ let pool: Arc<dyn Pool> = Arc::new(StaticPool::new(vec![be_addr.clone()]).unwrap());
+ let store: Arc<dyn AffinityStore> = Arc::new(MemoryStore::new());
+ let router = Arc::new(Router::single_pool(pool, store));
+ let dialer = Dialer::new();
+ let proxy = SparkConnectProxy::with_auth_and_metrics(
+ router,
+ dialer,
+ AuthInterceptor::new(Arc::new(AnonymousAuthenticator)),
+ metrics.clone(),
+ );
+
+ let gw_lis = TcpListener::bind("127.0.0.1:0").await.unwrap();
+ let gw_addr = gw_lis.local_addr().unwrap();
+
+ // Drain plumbing — mirrors gateway main.
+ let (shutdown_tx, mut shutdown_rx) = watch::channel(false);
+ let readiness_for_drain = readiness.clone();
+ let metrics_for_drain = metrics.clone();
+ let (drain_trigger_tx, drain_trigger_rx) = oneshot::channel();
+ let deadline = Duration::from_secs(5);
+ tokio::spawn(async move {
+ let _ = drain_trigger_rx.await;
+ eprintln!("[drain_smoke] drain triggered; flipping readiness off");
+ readiness_for_drain.mark_not_ready();
+ let drained = tokio::time::timeout(deadline, async {
+ loop {
+ if metrics_for_drain.active_streams_value() <= 0 {
+ return;
+ }
+ tokio::time::sleep(Duration::from_millis(100)).await;
+ }
+ })
+ .await;
+ match drained {
+ Ok(()) => eprintln!("[drain_smoke] drain finished cleanly"),
+ Err(_) => eprintln!("[drain_smoke] drain deadline reached"),
+ }
+ let _ = shutdown_tx.send(true);
+ });
+
+ let server_handle = tokio::spawn(async move {
+ Server::builder()
+ .add_service(pb::spark_connect_service_server::SparkConnectServiceServer::new(proxy))
+ .serve_with_incoming_shutdown(TcpListenerStream::new(gw_lis), async move {
+ let _ = shutdown_rx.changed().await;
+ })
+ .await
+ .ok();
+ });
+
+ // 3. Client opens a long-running stream.
+ let endpoint = Endpoint::from_shared(format!("http://{}", gw_addr)).unwrap();
+ let ch: Channel = endpoint
+ .connect_timeout(Duration::from_secs(2))
+ .connect()
+ .await
+ .unwrap();
+ let mut client = pb::spark_connect_service_client::SparkConnectServiceClient::new(ch);
+
+ let started_at = std::time::Instant::now();
+ let stream = client
+ .execute_plan(Request::new(pb::ExecutePlanRequest {
+ session_id: "drain-test".into(),
+ operation_id: Some("op-drain".into()),
+ ..Default::default()
+ }))
+ .await
+ .expect("ExecutePlan accepted")
+ .into_inner();
+ let mut stream = Box::pin(stream);
+
+ // Pull the first message so we know the stream is genuinely
+ // active before triggering drain.
+ let first = stream
+ .next()
+ .await
+ .expect("first message")
+ .expect("ok response");
+ eprintln!(
+ "[drain_smoke] received first message: response_id={}",
+ first.response_id
+ );
+
+ assert!(readiness.is_ready(), "before drain, /readyz should be 200");
+ assert!(
+ metrics.active_streams_value() >= 1,
+ "active stream should be counted; got {}",
+ metrics.active_streams_value()
+ );
+
+ // 4. Trigger drain. Readiness flips immediately.
+ let _ = drain_trigger_tx.send(());
+ tokio::time::sleep(Duration::from_millis(50)).await;
+ assert!(
+ !readiness.is_ready(),
+ "readiness must flip to not-ready as soon as drain is triggered"
+ );
+
+ // 5. Drain (and the gRPC server) waits for our stream to end.
+ // Pull the rest. Each message takes ~300ms; total should be
+ // around 1.2s after the first.
+ let mut count = 1;
+ while let Some(msg) = stream.next().await {
+ let m = msg.expect("ok response");
+ eprintln!(
+ "[drain_smoke] received message during drain: response_id={}",
+ m.response_id
+ );
+ count += 1;
+ }
+ let stream_duration = started_at.elapsed();
+ assert_eq!(
+ count, 5,
+ "all 5 stream messages must arrive even after drain triggered"
+ );
+ assert!(
+ stream_duration >= Duration::from_millis(1400),
+ "stream should take at least ~1.5s; got {:?}",
+ stream_duration
+ );
+ eprintln!(
+ "[drain_smoke] stream completed in {:?} with {} messages",
+ stream_duration, count
+ );
+
+ // 6. After the stream ended, active_streams should hit 0 and the
+ // server should shut down within ~200ms (the drain loop's
+ // poll interval).
+ let server_ended = tokio::time::timeout(Duration::from_secs(2), server_handle)
+ .await
+ .expect("gateway shut down within timeout");
+ server_ended.expect("gateway task ran to completion");
+ eprintln!("[drain_smoke] gateway shut down cleanly after drain");
+
+ let _ = be_kill_tx.send(());
+ println!("[drain_smoke] all drain invariants passed");
+}
diff --git a/crates/proxy/examples/e2e_client.rs b/crates/proxy/examples/e2e_client.rs
new file mode 100644
index 0000000..4d026a0
--- /dev/null
+++ b/crates/proxy/examples/e2e_client.rs
@@ -0,0 +1,74 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Tiny e2e helper used to validate the OTel exporter path against a
+//! real gateway + real OTLP collector. Sends one Config RPC.
+//!
+//! Run with:
+//! ```
+//! cargo run -p scg-proxy --example e2e_client -- http://127.0.0.1:15003
+//! ```
+//!
+//! The RPC is expected to FAIL (the e2e config points the gateway at an
+//! unreachable backend) — but the gateway's `scg_rpc` span fires
+//! regardless, which is what we want to observe at the collector.
+
+use scg_genproto::pb;
+use std::env;
+use tonic::metadata::MetadataValue;
+use tonic::Request;
+
+#[tokio::main]
+async fn main() {
+ let mut url = "http://127.0.0.1:15003".to_string();
+ let mut traceparent: Option<String> = None;
+ let mut args = env::args().skip(1);
+ while let Some(arg) = args.next() {
+ match arg.as_str() {
+ "--traceparent" => traceparent = args.next(),
+ other if other.starts_with("http") => url = other.to_string(),
+ other => eprintln!("[e2e_client] ignoring arg: {}", other),
+ }
+ }
+ eprintln!(
+ "[e2e_client] connecting to gateway at {} (traceparent={:?})",
+ url, traceparent
+ );
+
+ let ch = tonic::transport::Endpoint::from_shared(url)
+ .unwrap()
+ .connect()
+ .await
+ .expect("connect to gateway");
+ let mut c = pb::spark_connect_service_client::SparkConnectServiceClient::new(ch);
+
+ let mut req = Request::new(pb::ConfigRequest {
+ session_id: "e2e-session".into(),
+ ..Default::default()
+ });
+ if let Some(tp) = traceparent {
+ req.metadata_mut()
+ .insert("traceparent", MetadataValue::try_from(tp).unwrap());
+ }
+ let res = c.config(req).await;
+ match res {
+ Ok(resp) => eprintln!("[e2e_client] OK: {:?}", resp),
+ Err(s) => eprintln!(
+ "[e2e_client] RPC failed (expected): code={:?} message={}",
+ s.code(),
+ s.message()
+ ),
+ }
+}
diff --git a/crates/proxy/examples/fake_collector.rs b/crates/proxy/examples/fake_collector.rs
new file mode 100644
index 0000000..0e20717
--- /dev/null
+++ b/crates/proxy/examples/fake_collector.rs
@@ -0,0 +1,169 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Minimal OTLP/gRPC trace receiver. Listens on `:4317`, accepts
+//! `ExportTraceServiceRequest`, prints span name + selected attributes
+//! for each span, and exits as soon as it has seen `--expect <name>`
+//! at least once (or after `--timeout-secs` if set).
+//!
+//! Used by the e2e harness to validate that the production gateway
+//! actually emits the `scg_rpc` span via OTLP/gRPC.
+
+use std::collections::HashSet;
+use std::env;
+use std::sync::Arc;
+use std::time::Duration;
+
+use opentelemetry_proto::tonic::collector::trace::v1::trace_service_server::{
+ TraceService, TraceServiceServer,
+};
+use opentelemetry_proto::tonic::collector::trace::v1::{
+ ExportTraceServiceRequest, ExportTraceServiceResponse,
+};
+use tokio::sync::Mutex;
+use tonic::{transport::Server, Request, Response, Status};
+
+fn bytes_hex(b: &[u8]) -> String {
+ let mut s = String::with_capacity(b.len() * 2);
+ for byte in b {
+ use std::fmt::Write;
+ write!(&mut s, "{:02x}", byte).unwrap();
+ }
+ s
+}
+
+#[derive(Clone)]
+struct Collector {
+ seen_names: Arc<Mutex<HashSet<String>>>,
+ expect: Option<String>,
+ done_tx: Arc<Mutex<Option<tokio::sync::oneshot::Sender<()>>>>,
+}
+
+#[tonic::async_trait]
+impl TraceService for Collector {
+ async fn export(
+ &self,
+ req: Request<ExportTraceServiceRequest>,
+ ) -> Result<Response<ExportTraceServiceResponse>, Status> {
+ let body = req.into_inner();
+ for rs in &body.resource_spans {
+ // Resource attributes (service.name etc.)
+ if let Some(res) = &rs.resource {
+ let svc = res
+ .attributes
+ .iter()
+ .find(|kv| kv.key == "service.name")
+ .and_then(|kv| kv.value.as_ref())
+ .and_then(|v| match v.value.as_ref() {
+ Some(
+ opentelemetry_proto::tonic::common::v1::any_value::Value::StringValue(
+ s,
+ ),
+ ) => Some(s.as_str()),
+ _ => None,
+ });
+ if let Some(s) = svc {
+ println!("[collector] resource service.name={}", s);
+ }
+ }
+ for ss in &rs.scope_spans {
+ for span in &ss.spans {
+ let attrs: Vec<(String, String)> = span
+ .attributes
+ .iter()
+ .filter_map(|kv| {
+ kv.value.as_ref().and_then(|v| match v.value.as_ref()? {
+ opentelemetry_proto::tonic::common::v1::any_value::Value::StringValue(s) => {
+ Some((kv.key.clone(), s.clone()))
+ }
+ _ => None,
+ })
+ })
+ .collect();
+ let trace_id = bytes_hex(&span.trace_id);
+ let span_id = bytes_hex(&span.span_id);
+ let parent_id = if span.parent_span_id.is_empty() {
+ "(root)".to_string()
+ } else {
+ bytes_hex(&span.parent_span_id)
+ };
+ println!(
+ "[collector] span name={} trace_id={} span_id={} parent={} attrs={:?}",
+ span.name, trace_id, span_id, parent_id, attrs
+ );
+
+ let mut seen = self.seen_names.lock().await;
+ seen.insert(span.name.clone());
+ if let Some(expected) = &self.expect {
+ if seen.contains(expected) {
+ // Trigger shutdown if first time we see it.
+ let mut tx = self.done_tx.lock().await;
+ if let Some(s) = tx.take() {
+ println!(
+ "[collector] saw expected span name={}; shutting down",
+ expected
+ );
+ let _ = s.send(());
+ }
+ }
+ }
+ }
+ }
+ }
+ Ok(Response::new(ExportTraceServiceResponse::default()))
+ }
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+ let mut args = env::args().skip(1);
+ let mut expect: Option<String> = None;
+ let mut timeout_secs: u64 = 30;
+ while let Some(arg) = args.next() {
+ match arg.as_str() {
+ "--expect" => expect = args.next(),
+ "--timeout-secs" => {
+ timeout_secs = args.next().and_then(|s| s.parse().ok()).unwrap_or(30)
+ }
+ other => eprintln!("[collector] ignoring unknown arg: {}", other),
+ }
+ }
+
+ let (done_tx, done_rx) = tokio::sync::oneshot::channel();
+ let svc = Collector {
+ seen_names: Arc::new(Mutex::new(HashSet::new())),
+ expect: expect.clone(),
+ done_tx: Arc::new(Mutex::new(Some(done_tx))),
+ };
+
+ let addr = "0.0.0.0:4317".parse()?;
+ eprintln!(
+ "[collector] listening on {}; expect={:?} timeout_secs={}",
+ addr, expect, timeout_secs
+ );
+
+ let server = Server::builder()
+ .add_service(TraceServiceServer::new(svc))
+ .serve_with_shutdown(addr, async move {
+ tokio::select! {
+ _ = done_rx => {}
+ _ = tokio::time::sleep(Duration::from_secs(timeout_secs)) => {
+ eprintln!("[collector] timeout reached; shutting down");
+ }
+ }
+ });
+ server.await?;
+ Ok(())
+}
diff --git a/crates/proxy/examples/ha_smoke.rs b/crates/proxy/examples/ha_smoke.rs
new file mode 100644
index 0000000..d8901ab
--- /dev/null
+++ b/crates/proxy/examples/ha_smoke.rs
@@ -0,0 +1,390 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Multi-replica HA smoke test.
+//!
+//! Spawns two real SparkConnectProxy gateways, both backed by:
+//! * the same Redis affinity store (`scg-store-redis`),
+//! * the same static pool of two fake Spark Connect backends.
+//!
+//! Then drives a series of RPCs through different replicas and
+//! verifies three HA properties:
+//!
+//! A. **Shared state.** A session bound through replica A resolves
+//! to the *same* backend through replica B.
+//! B. **Failover.** After replica A is killed, the same session
+//! hitting replica B still routes to the original backend
+//! (binding survives the death of the replica that wrote it).
+//! C. **Op-id reverse index across replicas.** An ExecutePlan
+//! started through replica A registers its `operation_id`; a
+//! `ReattachExecute(op_id, session_id="different")` arriving at
+//! replica B (after A is gone) still reaches the backend that
+//! ran the plan.
+//!
+//! Run with a Redis listening on :6399 (or override via REDIS_URL):
+//!
+//! ```bash
+//! redis-server --port 6399 --daemonize yes
+//! cargo run -p scg-proxy --example ha_smoke
+//! ```
+//!
+//! Exits with non-zero status (panic) on any invariant failure so a
+//! CI script can wrap it.
+
+use std::pin::Pin;
+use std::sync::Arc;
+use std::time::Duration;
+
+use futures::Stream;
+use scg_auth::{AnonymousAuthenticator, AuthInterceptor};
+use scg_genproto::pb;
+use scg_observability::Metrics;
+use scg_pool_static::StaticPool;
+use scg_proxy::{Dialer, SparkConnectProxy};
+use scg_routing::{AffinityStore, Pool, Router};
+use scg_store_redis::{RedisStore, RedisStoreConfig};
+use tokio::net::TcpListener;
+use tokio::sync::oneshot;
+use tokio_stream::wrappers::TcpListenerStream;
+use tonic::transport::{Channel, Endpoint, Server};
+use tonic::{Request, Response, Status};
+
+/// Fake backend that tags every response with its own id, so the
+/// driver code can verify which backend handled a given RPC.
+#[derive(Default)]
+struct FakeBackend {
+ id: String,
+}
+
+impl FakeBackend {
+ fn new(id: impl Into<String>) -> Self {
+ Self { id: id.into() }
+ }
+}
+
+#[tonic::async_trait]
+impl pb::spark_connect_service_server::SparkConnectService for FakeBackend {
+ type ExecutePlanStream =
+ Pin<Box<dyn Stream<Item = Result<pb::ExecutePlanResponse, Status>> + Send + 'static>>;
+ type ReattachExecuteStream = Self::ExecutePlanStream;
+
+ async fn config(
+ &self,
+ req: Request<pb::ConfigRequest>,
+ ) -> Result<Response<pb::ConfigResponse>, Status> {
+ let body = req.into_inner();
+ Ok(Response::new(pb::ConfigResponse {
+ session_id: format!("{}@{}", body.session_id, self.id),
+ ..Default::default()
+ }))
+ }
+
+ async fn execute_plan(
+ &self,
+ req: Request<pb::ExecutePlanRequest>,
+ ) -> Result<Response<Self::ExecutePlanStream>, Status> {
+ let body = req.into_inner();
+ let id = self.id.clone();
+ let stream = async_stream::stream! {
+ yield Ok(pb::ExecutePlanResponse {
+ session_id: format!("{}@{}", body.session_id, id),
+ operation_id: body.operation_id.clone().unwrap_or_default(),
+ ..Default::default()
+ });
+ };
+ Ok(Response::new(Box::pin(stream)))
+ }
+
+ async fn reattach_execute(
+ &self,
+ req: Request<pb::ReattachExecuteRequest>,
+ ) -> Result<Response<Self::ReattachExecuteStream>, Status> {
+ let body = req.into_inner();
+ let id = self.id.clone();
+ let stream = async_stream::stream! {
+ yield Ok(pb::ExecutePlanResponse {
+ session_id: format!("{}@{}", body.session_id, id),
+ operation_id: body.operation_id.clone(),
+ ..Default::default()
+ });
+ };
+ Ok(Response::new(Box::pin(stream)))
+ }
+
+ // Unused but trait-required.
+ async fn analyze_plan(
+ &self,
+ _: Request<pb::AnalyzePlanRequest>,
+ ) -> Result<Response<pb::AnalyzePlanResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn artifact_status(
+ &self,
+ _: Request<pb::ArtifactStatusesRequest>,
+ ) -> Result<Response<pb::ArtifactStatusesResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn interrupt(
+ &self,
+ _: Request<pb::InterruptRequest>,
+ ) -> Result<Response<pb::InterruptResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn release_execute(
+ &self,
+ _: Request<pb::ReleaseExecuteRequest>,
+ ) -> Result<Response<pb::ReleaseExecuteResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn release_session(
+ &self,
+ _: Request<pb::ReleaseSessionRequest>,
+ ) -> Result<Response<pb::ReleaseSessionResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn fetch_error_details(
+ &self,
+ _: Request<pb::FetchErrorDetailsRequest>,
+ ) -> Result<Response<pb::FetchErrorDetailsResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn clone_session(
+ &self,
+ _: Request<pb::CloneSessionRequest>,
+ ) -> Result<Response<pb::CloneSessionResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn get_status(
+ &self,
+ _: Request<pb::GetStatusRequest>,
+ ) -> Result<Response<pb::GetStatusResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn add_artifacts(
+ &self,
+ _: Request<tonic::Streaming<pb::AddArtifactsRequest>>,
+ ) -> Result<Response<pb::AddArtifactsResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+}
+
+async fn spawn_backend(id: &'static str) -> (String, oneshot::Sender<()>) {
+ let svc =
+ pb::spark_connect_service_server::SparkConnectServiceServer::new(FakeBackend::new(id));
+ let lis = TcpListener::bind("127.0.0.1:0").await.unwrap();
+ let addr = lis.local_addr().unwrap().to_string();
+ let (tx, rx) = oneshot::channel();
+ tokio::spawn(async move {
+ Server::builder()
+ .add_service(svc)
+ .serve_with_incoming_shutdown(TcpListenerStream::new(lis), async {
+ let _ = rx.await;
+ })
+ .await
+ .ok();
+ });
+ (addr, tx)
+}
+
+async fn spawn_gateway(
+ backends: Vec<String>,
+ store: Arc<dyn AffinityStore>,
+) -> (Channel, oneshot::Sender<()>) {
+ let metrics = Metrics::new().unwrap();
+ let pool: Arc<dyn Pool> = Arc::new(StaticPool::new(backends).unwrap());
+ let router = Arc::new(Router::single_pool(pool, store));
+ let dialer = Dialer::new();
+ let proxy = SparkConnectProxy::with_auth_and_metrics(
+ router,
+ dialer,
+ AuthInterceptor::new(Arc::new(AnonymousAuthenticator)),
+ metrics,
+ );
+
+ let lis = TcpListener::bind("127.0.0.1:0").await.unwrap();
+ let addr = lis.local_addr().unwrap();
+ let (tx, rx) = oneshot::channel();
+ tokio::spawn(async move {
+ Server::builder()
+ .add_service(pb::spark_connect_service_server::SparkConnectServiceServer::new(proxy))
+ .serve_with_incoming_shutdown(TcpListenerStream::new(lis), async {
+ let _ = rx.await;
+ })
+ .await
+ .ok();
+ });
+
+ let endpoint = Endpoint::from_shared(format!("http://{}", addr)).unwrap();
+ let ch = endpoint
+ .connect_timeout(Duration::from_secs(2))
+ .connect()
+ .await
+ .unwrap();
+ (ch, tx)
+}
+
+/// Issue a Config call and return the backend id that handled it
+/// (extracted from the FakeBackend's tagged session_id).
+async fn config_through(ch: &Channel, session_id: &str) -> String {
+ let mut c = pb::spark_connect_service_client::SparkConnectServiceClient::new(ch.clone());
+ let resp = c
+ .config(Request::new(pb::ConfigRequest {
+ session_id: session_id.into(),
+ ..Default::default()
+ }))
+ .await
+ .expect("Config RPC succeeds");
+ let tagged = resp.into_inner().session_id;
+ // FakeBackend formats it as "{client_session_id}@{backend_id}".
+ tagged.rsplit('@').next().unwrap().to_string()
+}
+
+#[tokio::main]
+async fn main() {
+ let redis_url = std::env::var("REDIS_URL").unwrap_or_else(|_| "redis://127.0.0.1:6399".into());
+ eprintln!("[ha_smoke] redis at {}", redis_url);
+
+ // Two fake backends. The pool round-robins, so the first session
+ // hits "be-a", the next "be-b", and so on.
+ let (be_a, _be_a_kill) = spawn_backend("be-a").await;
+ let (be_b, _be_b_kill) = spawn_backend("be-b").await;
+ let backends = vec![be_a.clone(), be_b.clone()];
+ eprintln!("[ha_smoke] backends: {:?}", backends);
+
+ // Unique key prefix per run so re-runs don't see each other's state.
+ let prefix = format!("scg-ha-{}", std::process::id());
+
+ let make_store = || async {
+ let store = RedisStore::connect(RedisStoreConfig {
+ url: redis_url.clone(),
+ key_prefix: prefix.clone(),
+ session_ttl: Duration::from_secs(60),
+ op_ttl: Duration::from_secs(60),
+ })
+ .await
+ .expect("connect redis");
+ Arc::new(store) as Arc<dyn AffinityStore>
+ };
+
+ // Two replicas. Each gets its own RedisStore handle, but they
+ // both point at the same Redis (and same key prefix), so their
+ // affinity tables are the same persistent dataset.
+ let store_a = make_store().await;
+ let store_b = make_store().await;
+ let (ch_a, kill_a) = spawn_gateway(backends.clone(), store_a).await;
+ let (ch_b, kill_b) = spawn_gateway(backends.clone(), store_b).await;
+ eprintln!("[ha_smoke] gateway A and gateway B up");
+
+ // ----- Scenario A: shared state across replicas -------------------
+ let session = "ha-session-1";
+ let chosen_via_a = config_through(&ch_a, session).await;
+ eprintln!(
+ "[ha_smoke] scenario A: session {} bound via replica A -> {}",
+ session, chosen_via_a
+ );
+ let chosen_via_b = config_through(&ch_b, session).await;
+ assert_eq!(
+ chosen_via_a, chosen_via_b,
+ "scenario A: replica B must resolve to the same backend as replica A"
+ );
+ eprintln!("[ha_smoke] ok: scenario A — shared state");
+
+ // ----- Scenario B: failover after replica death ------------------
+ let session_b = "ha-session-2";
+ let chosen_b_pre = config_through(&ch_a, session_b).await;
+ eprintln!(
+ "[ha_smoke] scenario B: session {} bound via replica A -> {}",
+ session_b, chosen_b_pre
+ );
+ // Kill replica A. Drop the channel first so the connection
+ // closes; then send the shutdown.
+ drop(ch_a);
+ let _ = kill_a.send(());
+ // Give the runtime a beat to actually tear the server down.
+ tokio::time::sleep(Duration::from_millis(50)).await;
+
+ let chosen_b_post = config_through(&ch_b, session_b).await;
+ assert_eq!(
+ chosen_b_post, chosen_b_pre,
+ "scenario B: same session through replica B must still hit the original backend"
+ );
+ eprintln!("[ha_smoke] ok: scenario B — failover preserves binding");
+
+ // ----- Scenario C: reattach via the surviving replica ------------
+ // ExecutePlan through *replica B* (replica A is gone). Capture
+ // the backend that handled it via the tagged session_id, then
+ // ReattachExecute(op_id, **different session_id**) and prove the
+ // op-id reverse index sends us to the same backend.
+ let session_c = "ha-session-3";
+ let op_id = "ha-op-3";
+ let mut client_b =
+ pb::spark_connect_service_client::SparkConnectServiceClient::new(ch_b.clone());
+
+ use futures::StreamExt;
+ let mut exec_stream = client_b
+ .execute_plan(Request::new(pb::ExecutePlanRequest {
+ session_id: session_c.into(),
+ operation_id: Some(op_id.into()),
+ ..Default::default()
+ }))
+ .await
+ .expect("ExecutePlan succeeds")
+ .into_inner();
+ let first = exec_stream
+ .next()
+ .await
+ .expect("at least one response")
+ .expect("ok response");
+ let exec_backend = first.session_id.rsplit('@').next().unwrap().to_string();
+ eprintln!(
+ "[ha_smoke] scenario C: ExecutePlan(op={}) routed to {}",
+ op_id, exec_backend
+ );
+ // Drain the rest of the stream so the server-side handler closes
+ // and the op-id binding is durable in Redis.
+ while exec_stream.next().await.is_some() {}
+
+ let mut reattach_stream = client_b
+ .reattach_execute(Request::new(pb::ReattachExecuteRequest {
+ // Deliberately a *different* session id — the op-id index
+ // is what must save us.
+ session_id: "totally-different-session".into(),
+ operation_id: op_id.into(),
+ ..Default::default()
+ }))
+ .await
+ .expect("ReattachExecute succeeds")
+ .into_inner();
+ let reattach_first = reattach_stream
+ .next()
+ .await
+ .expect("at least one reattach response")
+ .expect("ok reattach response");
+ let reattach_backend = reattach_first
+ .session_id
+ .rsplit('@')
+ .next()
+ .unwrap()
+ .to_string();
+ assert_eq!(
+ reattach_backend, exec_backend,
+ "scenario C: reattach via op-id reverse index must land on the same backend"
+ );
+ eprintln!("[ha_smoke] ok: scenario C — op-id reverse index works across replicas");
+
+ // Cleanup.
+ let _ = kill_b.send(());
+ println!("[ha_smoke] all HA invariants passed");
+}
diff --git a/crates/proxy/examples/load.rs b/crates/proxy/examples/load.rs
new file mode 100644
index 0000000..e369369
--- /dev/null
+++ b/crates/proxy/examples/load.rs
@@ -0,0 +1,1008 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Load-test harness.
+//!
+//! Spins up a SparkConnectProxy + N FakeBackends in-process and
+//! drives traffic at it through real gRPC clients. Measures
+//! throughput, latency distribution, and the cost of the optional
+//! features (active health probing, graceful drain).
+//!
+//! NOT a microbenchmark. The interesting question is end-to-end
+//! gateway behaviour under realistic-ish concurrency, not the
+//! per-instruction cost of `Router::resolve_session`. Use criterion
+//! for the latter.
+//!
+//! Usage:
+//!
+//! cargo run -p scg-proxy --example load --release -- \
+//! unary --workers 32 --duration-secs 30
+//!
+//! cargo run -p scg-proxy --example load --release -- \
+//! streaming --concurrency 100 --messages 100
+//!
+//! cargo run -p scg-proxy --example load --release -- \
+//! hc-overhead --workers 32 --duration-secs 20
+//!
+//! cargo run -p scg-proxy --example load --release -- \
+//! drain-under-load --concurrency 100 --deadline-secs 10
+//!
+//! cargo run -p scg-proxy --example load --release -- \
+//! overhead --workers 32 --duration-secs 20
+//!
+//! cargo run -p scg-proxy --example load --release -- \
+//! redis-affinity --workers 32 --duration-secs 20 \
+//! --redis-url redis://127.0.0.1:6379
+
+use std::pin::Pin;
+use std::sync::atomic::{AtomicU64, Ordering};
+use std::sync::Arc;
+use std::time::{Duration, Instant};
+
+use futures::{Stream, StreamExt};
+use hdrhistogram::Histogram;
+use scg_audit::{AuditConfig, AuditLogger};
+use scg_auth::{AnonymousAuthenticator, AuthInterceptor};
+use scg_genproto::pb;
+use scg_healthcheck::{HealthAwarePool, HealthCheckConfig};
+use scg_observability::{Metrics, ReadinessProbe};
+use scg_pool_static::StaticPool;
+use scg_proxy::{Dialer, SparkConnectProxy};
+use scg_ratelimit::{
+ BucketRate, LimiterObserver, MemoryLimiter, RateLimiter, RejectScope, TenantLimits,
+};
+use scg_routing::{AffinityStore, Pool, Router};
+use scg_store_memory::MemoryStore;
+use scg_store_redis::{RedisStore, RedisStoreConfig};
+use scg_tenant::{OnMissing, TenantResolver, TenantResolverConfig, TenantSource};
+use tokio::net::TcpListener;
+use tokio::sync::oneshot;
+use tokio_stream::wrappers::TcpListenerStream;
+use tonic::transport::{Channel, Endpoint, Server};
+use tonic::{Request, Response, Status};
+
+// ---------- Helpers ---------------------------------------------------------
+
+/// No-op limiter observer — the perf rig configures a generous
+/// bucket that never rejects, so we don't care about the on-reject
+/// signal. Matches the noop pattern in the in-memory limiter's own
+/// tests.
+struct NoopLimiterObs;
+impl LimiterObserver for NoopLimiterObs {
+ fn on_reject(&self, _tenant: &str, _scope: RejectScope) {}
+}
+
+// ---------- Fake backend ----------------------------------------------------
+
+#[derive(Clone, Default)]
+struct FakeBackend {
+ /// Per-message delay for streaming RPCs. 0 = as fast as possible.
+ stream_delay_ms: u64,
+ stream_messages: usize,
+}
+
+#[tonic::async_trait]
+impl pb::spark_connect_service_server::SparkConnectService for FakeBackend {
+ type ExecutePlanStream =
+ Pin<Box<dyn Stream<Item = Result<pb::ExecutePlanResponse, Status>> + Send + 'static>>;
+ type ReattachExecuteStream = Self::ExecutePlanStream;
+
+ async fn config(
+ &self,
+ req: Request<pb::ConfigRequest>,
+ ) -> Result<Response<pb::ConfigResponse>, Status> {
+ let body = req.into_inner();
+ Ok(Response::new(pb::ConfigResponse {
+ session_id: body.session_id,
+ ..Default::default()
+ }))
+ }
+
+ async fn execute_plan(
+ &self,
+ req: Request<pb::ExecutePlanRequest>,
+ ) -> Result<Response<Self::ExecutePlanStream>, Status> {
+ let body = req.into_inner();
+ let n = self.stream_messages;
+ let delay_ms = self.stream_delay_ms;
+ let stream = async_stream::stream! {
+ for i in 0..n {
+ if delay_ms > 0 {
+ tokio::time::sleep(Duration::from_millis(delay_ms)).await;
+ }
+ yield Ok(pb::ExecutePlanResponse {
+ session_id: body.session_id.clone(),
+ operation_id: body.operation_id.clone().unwrap_or_default(),
+ response_id: format!("msg-{}", i),
+ ..Default::default()
+ });
+ }
+ };
+ Ok(Response::new(Box::pin(stream)))
+ }
+
+ // The other RPCs are unused by the load harness — return Unimplemented.
+ async fn analyze_plan(
+ &self,
+ _: Request<pb::AnalyzePlanRequest>,
+ ) -> Result<Response<pb::AnalyzePlanResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn artifact_status(
+ &self,
+ _: Request<pb::ArtifactStatusesRequest>,
+ ) -> Result<Response<pb::ArtifactStatusesResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn interrupt(
+ &self,
+ _: Request<pb::InterruptRequest>,
+ ) -> Result<Response<pb::InterruptResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn release_execute(
+ &self,
+ _: Request<pb::ReleaseExecuteRequest>,
+ ) -> Result<Response<pb::ReleaseExecuteResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn release_session(
+ &self,
+ _: Request<pb::ReleaseSessionRequest>,
+ ) -> Result<Response<pb::ReleaseSessionResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn fetch_error_details(
+ &self,
+ _: Request<pb::FetchErrorDetailsRequest>,
+ ) -> Result<Response<pb::FetchErrorDetailsResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn clone_session(
+ &self,
+ _: Request<pb::CloneSessionRequest>,
+ ) -> Result<Response<pb::CloneSessionResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn get_status(
+ &self,
+ _: Request<pb::GetStatusRequest>,
+ ) -> Result<Response<pb::GetStatusResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn reattach_execute(
+ &self,
+ _: Request<pb::ReattachExecuteRequest>,
+ ) -> Result<Response<Self::ReattachExecuteStream>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn add_artifacts(
+ &self,
+ _: Request<tonic::Streaming<pb::AddArtifactsRequest>>,
+ ) -> Result<Response<pb::AddArtifactsResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+}
+
+// ---------- Test rig --------------------------------------------------------
+
+struct Rig {
+ grpc: Channel,
+ metrics: Metrics,
+ readiness: ReadinessProbe,
+ /// Dropped → backends shut down.
+ _be_kills: Vec<oneshot::Sender<()>>,
+ /// Sender → gateway shuts down.
+ gw_kill: Option<oneshot::Sender<()>>,
+}
+
+#[derive(Clone, Debug)]
+struct RigOpts {
+ n_backends: usize,
+ stream_messages: usize,
+ stream_delay_ms: u64,
+ health_check: bool,
+ /// When `Some`, build a `TenantResolver` (always-default source)
+ /// and route every RPC through it. Measures the tenant-resolution
+ /// hot-path cost vs the baseline rig that uses `with_auth_and_metrics`
+ /// (no resolver wired in).
+ tenant_resolver: bool,
+ /// When `Some`, wire in a `MemoryLimiter` with a generous quota
+ /// (high enough that no test request gets throttled) so we can
+ /// measure the per-RPC bucket-check cost in isolation.
+ rate_limit: bool,
+ /// When `true`, attach an enabled `AuditLogger`. Every RPC fires
+ /// at least one tracing event; with `log_successful_rpcs=false`
+ /// the cost is dominated by the session.create event on the
+ /// binding path.
+ audit: bool,
+ /// When `Some(url)`, swap the in-process `MemoryStore` for a
+ /// `RedisStore` pointed at this URL. Measures the affinity-store
+ /// round-trip cost. The Redis must already be reachable; the
+ /// harness fails fast if connect errors.
+ affinity_redis_url: Option<String>,
+}
+
+impl Default for RigOpts {
+ fn default() -> Self {
+ Self {
+ n_backends: 2,
+ stream_messages: 100,
+ stream_delay_ms: 0,
+ health_check: false,
+ tenant_resolver: false,
+ rate_limit: false,
+ audit: false,
+ affinity_redis_url: None,
+ }
+ }
+}
+
+async fn spawn_rig(opts: RigOpts) -> Rig {
+ // Backends.
+ let mut backends = Vec::with_capacity(opts.n_backends);
+ let mut be_kills = Vec::with_capacity(opts.n_backends);
+ for _ in 0..opts.n_backends {
+ let lis = TcpListener::bind("127.0.0.1:0").await.unwrap();
+ let addr = lis.local_addr().unwrap().to_string();
+ let backend = FakeBackend {
+ stream_delay_ms: opts.stream_delay_ms,
+ stream_messages: opts.stream_messages,
+ };
+ let svc = pb::spark_connect_service_server::SparkConnectServiceServer::new(backend.clone());
+ let (tx, rx) = oneshot::channel();
+ tokio::spawn(async move {
+ Server::builder()
+ .add_service(svc)
+ .serve_with_incoming_shutdown(TcpListenerStream::new(lis), async {
+ let _ = rx.await;
+ })
+ .await
+ .ok();
+ });
+ backends.push(addr);
+ be_kills.push(tx);
+ }
+
+ // Gateway.
+ let metrics = Metrics::new().unwrap();
+ let readiness = ReadinessProbe::new(true);
+ let inner: Arc<dyn Pool> = Arc::new(StaticPool::new(backends.clone()).unwrap());
+ let pool: Arc<dyn Pool> = if opts.health_check {
+ let hc_cfg = HealthCheckConfig {
+ interval: Duration::from_secs(5),
+ timeout: Duration::from_secs(2),
+ unhealthy_threshold: 3,
+ healthy_threshold: 2,
+ };
+ let wrapped = HealthAwarePool::new(inner, hc_cfg);
+ let _probe = wrapped.spawn_probe();
+ wrapped
+ } else {
+ inner
+ };
+ let store: Arc<dyn AffinityStore> = if let Some(url) = &opts.affinity_redis_url {
+ // Unique key prefix per rig so consecutive harness runs
+ // don't trip over each other: each run uses fresh backend
+ // ports, so stale bindings from a previous run would point
+ // at dead addresses.
+ let prefix = format!(
+ "scg-perf-{}-{}",
+ std::process::id(),
+ std::time::SystemTime::now()
+ .duration_since(std::time::UNIX_EPOCH)
+ .map(|d| d.as_nanos())
+ .unwrap_or(0),
+ );
+ let cfg = RedisStoreConfig {
+ url: url.clone(),
+ key_prefix: prefix,
+ ..Default::default()
+ };
+ Arc::new(
+ RedisStore::connect(cfg)
+ .await
+ .expect("connect to Redis affinity store"),
+ )
+ } else {
+ Arc::new(MemoryStore::new())
+ };
+ let router = Arc::new(Router::single_pool(pool, store));
+ let dialer = Dialer::new();
+ let auth = AuthInterceptor::new(Arc::new(AnonymousAuthenticator));
+
+ // Decide between the back-compat constructor (matches the
+ // pre-Phase-3 perf baseline) and the full constructor. The
+ // full path runs through `authenticate_and_resolve` + tenant
+ // resolver + optional rate limit + optional audit, so each
+ // toggle has an additive overhead that the `overhead` scenario
+ // measures step by step.
+ let proxy = if opts.tenant_resolver || opts.rate_limit || opts.audit {
+ let tr = TenantResolver::new(TenantResolverConfig {
+ source: TenantSource::AlwaysDefault,
+ on_missing: OnMissing::UseDefault,
+ default_name: "default".into(),
+ });
+ let limiter = if opts.rate_limit {
+ // Generous quota — 1M rps, 1M burst — so the limiter
+ // never actually rejects. We only want to measure the
+ // bucket-check cost on the hot path.
+ let limits = TenantLimits {
+ tenant: BucketRate {
+ rpcs_per_second: 1_000_000.0,
+ burst: 1_000_000,
+ },
+ per_user: BucketRate::disabled(),
+ };
+ Some(RateLimiter::Memory(MemoryLimiter::new(
+ limits,
+ Default::default(),
+ Arc::new(NoopLimiterObs),
+ )))
+ } else {
+ None
+ };
+ let audit = if opts.audit {
+ AuditLogger::new(AuditConfig {
+ enabled: true,
+ log_successful_rpcs: false,
+ })
+ } else {
+ AuditLogger::disabled()
+ };
+ SparkConnectProxy::with_all(router, dialer, auth, metrics.clone(), tr, limiter, audit)
+ } else {
+ SparkConnectProxy::with_auth_and_metrics(router, dialer, auth, metrics.clone())
+ };
+
+ let lis = TcpListener::bind("127.0.0.1:0").await.unwrap();
+ let addr = lis.local_addr().unwrap();
+ let (gw_tx, gw_rx) = oneshot::channel();
+ tokio::spawn(async move {
+ Server::builder()
+ .add_service(pb::spark_connect_service_server::SparkConnectServiceServer::new(proxy))
+ .serve_with_incoming_shutdown(TcpListenerStream::new(lis), async {
+ let _ = gw_rx.await;
+ })
+ .await
+ .ok();
+ });
+
+ let endpoint = Endpoint::from_shared(format!("http://{}", addr)).unwrap();
+ let grpc = endpoint
+ .connect_timeout(Duration::from_secs(2))
+ .connect()
+ .await
+ .unwrap();
+
+ Rig {
+ grpc,
+ metrics,
+ readiness,
+ _be_kills: be_kills,
+ gw_kill: Some(gw_tx),
+ }
+}
+
+// ---------- Latency reporter -----------------------------------------------
+
+fn new_hist() -> Histogram<u64> {
+ // 1µs..60s, 3 significant digits.
+ Histogram::<u64>::new_with_bounds(1, 60_000_000, 3).unwrap()
+}
+
+fn print_summary(label: &str, hist: &Histogram<u64>, rpcs: u64, errors: u64, duration: Duration) {
+ let qps = rpcs as f64 / duration.as_secs_f64();
+ println!("---");
+ println!("[{}]", label);
+ println!(" duration: {:>10.2}s", duration.as_secs_f64());
+ println!(" total RPCs: {:>10}", rpcs);
+ println!(" errors: {:>10}", errors);
+ println!(" QPS: {:>10.0}", qps);
+ println!(
+ " latency (ms) p50={:.3} p95={:.3} p99={:.3} p999={:.3} max={:.3}",
+ hist.value_at_quantile(0.50) as f64 / 1000.0,
+ hist.value_at_quantile(0.95) as f64 / 1000.0,
+ hist.value_at_quantile(0.99) as f64 / 1000.0,
+ hist.value_at_quantile(0.999) as f64 / 1000.0,
+ hist.max() as f64 / 1000.0,
+ );
+}
+
+// ---------- Scenario 1: unary throughput -----------------------------------
+
+async fn scenario_unary(workers: usize, duration: Duration, n_sessions: usize) {
+ println!(
+ "[scenario] unary throughput: workers={} duration={:?} sessions={}",
+ workers, duration, n_sessions
+ );
+ let rig = spawn_rig(RigOpts::default()).await;
+ let stop_at = Instant::now() + duration;
+ let total = Arc::new(AtomicU64::new(0));
+ let errors = Arc::new(AtomicU64::new(0));
+ let hist_handles: Vec<_> = (0..workers)
+ .map(|wid| {
+ let ch = rig.grpc.clone();
+ let total = total.clone();
+ let errors = errors.clone();
+ tokio::spawn(async move {
+ let mut h = new_hist();
+ let mut client =
+ pb::spark_connect_service_client::SparkConnectServiceClient::new(ch);
+ let mut local = 0u64;
+ while Instant::now() < stop_at {
+ let session_id = format!("session-{}", local % n_sessions as u64);
+ let started = Instant::now();
+ let res = client
+ .config(Request::new(pb::ConfigRequest {
+ session_id,
+ ..Default::default()
+ }))
+ .await;
+ let elapsed_us = started.elapsed().as_micros() as u64;
+ h.record(elapsed_us.max(1)).ok();
+ match res {
+ Ok(_) => {
+ total.fetch_add(1, Ordering::Relaxed);
+ }
+ Err(_) => {
+ errors.fetch_add(1, Ordering::Relaxed);
+ }
+ }
+ local += 1;
+ let _ = wid; // silence unused; useful when debugging by worker
+ }
+ h
+ })
+ })
+ .collect();
+ let started = Instant::now();
+ let mut combined = new_hist();
+ for h in hist_handles {
+ let part = h.await.unwrap();
+ combined.add(part).ok();
+ }
+ let elapsed = started.elapsed();
+ print_summary(
+ "unary",
+ &combined,
+ total.load(Ordering::Relaxed),
+ errors.load(Ordering::Relaxed),
+ elapsed,
+ );
+ drop(rig);
+}
+
+// ---------- Scenario 2: streaming concurrency -------------------------------
+
+async fn scenario_streaming(concurrency: usize, messages: usize, n_sessions: usize) {
+ println!(
+ "[scenario] streaming concurrency: concurrency={} msgs/stream={} sessions={}",
+ concurrency, messages, n_sessions
+ );
+ let rig = spawn_rig(RigOpts {
+ stream_messages: messages,
+ ..RigOpts::default()
+ })
+ .await;
+ let total = Arc::new(AtomicU64::new(0));
+ let errors = Arc::new(AtomicU64::new(0));
+ let started = Instant::now();
+ let handles: Vec<_> = (0..concurrency)
+ .map(|i| {
+ let ch = rig.grpc.clone();
+ let total = total.clone();
+ let errors = errors.clone();
+ tokio::spawn(async move {
+ let mut h = new_hist();
+ let mut client =
+ pb::spark_connect_service_client::SparkConnectServiceClient::new(ch);
+ let session_id = format!("session-{}", i % n_sessions);
+ let started = Instant::now();
+ let stream = client
+ .execute_plan(Request::new(pb::ExecutePlanRequest {
+ session_id,
+ operation_id: Some(format!("op-{}", i)),
+ ..Default::default()
+ }))
+ .await;
+ match stream {
+ Ok(resp) => {
+ let mut s = resp.into_inner();
+ let mut count = 0u64;
+ while let Some(item) = s.next().await {
+ if item.is_err() {
+ errors.fetch_add(1, Ordering::Relaxed);
+ break;
+ }
+ count += 1;
+ }
+ let elapsed_us = started.elapsed().as_micros() as u64;
+ h.record(elapsed_us.max(1)).ok();
+ total.fetch_add(count, Ordering::Relaxed);
+ }
+ Err(_) => {
+ errors.fetch_add(1, Ordering::Relaxed);
+ }
+ }
+ h
+ })
+ })
+ .collect();
+ let mut combined = new_hist();
+ let peak_active = rig.metrics.active_streams_value();
+ println!(
+ " peak active_streams observed during fan-out: {}",
+ peak_active
+ );
+ for h in handles {
+ let part = h.await.unwrap();
+ combined.add(part).ok();
+ }
+ let elapsed = started.elapsed();
+ print_summary(
+ "streaming (per-stream completion)",
+ &combined,
+ concurrency as u64,
+ errors.load(Ordering::Relaxed),
+ elapsed,
+ );
+ let total_msgs = total.load(Ordering::Relaxed);
+ println!(
+ " messages forwarded: {} ({:.0} msg/s)",
+ total_msgs,
+ total_msgs as f64 / elapsed.as_secs_f64()
+ );
+ drop(rig);
+}
+
+// ---------- Scenario 3: health-check overhead ------------------------------
+
+async fn scenario_hc_overhead(workers: usize, duration: Duration, n_sessions: usize) {
+ println!(
+ "[scenario] health-check overhead: workers={} duration={:?} (each)",
+ workers, duration
+ );
+ println!(" baseline (no HC):");
+ let baseline = run_unary_internal(workers, duration, n_sessions, RigOpts::default()).await;
+ println!(" with HC:");
+ let with_hc = run_unary_internal(
+ workers,
+ duration,
+ n_sessions,
+ RigOpts {
+ health_check: true,
+ ..RigOpts::default()
+ },
+ )
+ .await;
+ println!("---");
+ println!("[hc-overhead summary]");
+ println!(
+ " p50 (ms): baseline={:.3} with_hc={:.3} delta={:+.3}",
+ baseline.0.value_at_quantile(0.50) as f64 / 1000.0,
+ with_hc.0.value_at_quantile(0.50) as f64 / 1000.0,
+ (with_hc.0.value_at_quantile(0.50) as i64 - baseline.0.value_at_quantile(0.50) as i64)
+ as f64
+ / 1000.0,
+ );
+ println!(
+ " p99 (ms): baseline={:.3} with_hc={:.3} delta={:+.3}",
+ baseline.0.value_at_quantile(0.99) as f64 / 1000.0,
+ with_hc.0.value_at_quantile(0.99) as f64 / 1000.0,
+ (with_hc.0.value_at_quantile(0.99) as i64 - baseline.0.value_at_quantile(0.99) as i64)
+ as f64
+ / 1000.0,
+ );
+ let qps_baseline = baseline.1 as f64 / duration.as_secs_f64();
+ let qps_with_hc = with_hc.1 as f64 / duration.as_secs_f64();
+ println!(
+ " QPS: baseline={:.0} with_hc={:.0} delta={:+.1}%",
+ qps_baseline,
+ qps_with_hc,
+ (qps_with_hc - qps_baseline) / qps_baseline * 100.0
+ );
+}
+
+async fn run_unary_internal(
+ workers: usize,
+ duration: Duration,
+ n_sessions: usize,
+ opts: RigOpts,
+) -> (Histogram<u64>, u64) {
+ let rig = spawn_rig(opts).await;
+ let stop_at = Instant::now() + duration;
+ let total = Arc::new(AtomicU64::new(0));
+ let errors = Arc::new(AtomicU64::new(0));
+ let hist_handles: Vec<_> = (0..workers)
+ .map(|_wid| {
+ let ch = rig.grpc.clone();
+ let total = total.clone();
+ let errors = errors.clone();
+ tokio::spawn(async move {
+ let mut h = new_hist();
+ let mut client =
+ pb::spark_connect_service_client::SparkConnectServiceClient::new(ch);
+ let mut local = 0u64;
+ while Instant::now() < stop_at {
+ let session_id = format!("session-{}", local % n_sessions as u64);
+ let started = Instant::now();
+ let res = client
+ .config(Request::new(pb::ConfigRequest {
+ session_id,
+ ..Default::default()
+ }))
+ .await;
+ let elapsed_us = started.elapsed().as_micros() as u64;
+ h.record(elapsed_us.max(1)).ok();
+ if res.is_err() {
+ errors.fetch_add(1, Ordering::Relaxed);
+ } else {
+ total.fetch_add(1, Ordering::Relaxed);
+ }
+ local += 1;
+ }
+ h
+ })
+ })
+ .collect();
+ let mut combined = new_hist();
+ for h in hist_handles {
+ combined.add(h.await.unwrap()).ok();
+ }
+ let total = total.load(Ordering::Relaxed);
+ let errors = errors.load(Ordering::Relaxed);
+ println!(
+ " rpcs={} errors={} qps={:.0} p50={:.3}ms p99={:.3}ms",
+ total,
+ errors,
+ total as f64 / duration.as_secs_f64(),
+ combined.value_at_quantile(0.50) as f64 / 1000.0,
+ combined.value_at_quantile(0.99) as f64 / 1000.0,
+ );
+ drop(rig);
+ (combined, total)
+}
+
+// ---------- Scenario 5: per-feature hot-path overhead -----------------------
+
+/// Walk through five rig configurations and report per-RPC overhead
+/// added by each multi-tenant feature stacked on top of the bare
+/// baseline. The bucket and the audit logger are configured so that
+/// they never reject and only emit on the binding path, respectively
+/// — the goal is to measure the *check* / *emit* cost on every RPC,
+/// not the cost of the work they protect against.
+async fn scenario_overhead(workers: usize, duration: Duration, n_sessions: usize) {
+ println!(
+ "[scenario] per-feature overhead: workers={} duration={:?} (each)",
+ workers, duration
+ );
+
+ let configs: &[(&str, RigOpts)] = &[
+ ("baseline (with_auth_and_metrics)", RigOpts::default()),
+ (
+ "+ tenant_resolver",
+ RigOpts {
+ tenant_resolver: true,
+ ..RigOpts::default()
+ },
+ ),
+ (
+ "+ tenant_resolver + rate_limit",
+ RigOpts {
+ tenant_resolver: true,
+ rate_limit: true,
+ ..RigOpts::default()
+ },
+ ),
+ (
+ "+ tenant_resolver + audit",
+ RigOpts {
+ tenant_resolver: true,
+ audit: true,
+ ..RigOpts::default()
+ },
+ ),
+ (
+ "+ all three",
+ RigOpts {
+ tenant_resolver: true,
+ rate_limit: true,
+ audit: true,
+ ..RigOpts::default()
+ },
+ ),
+ ];
+
+ let mut results: Vec<(String, Histogram<u64>, u64)> = Vec::with_capacity(configs.len());
+ for (label, opts) in configs {
+ println!(" {}:", label);
+ let (h, total) = run_unary_internal(workers, duration, n_sessions, opts.clone()).await;
+ results.push((label.to_string(), h, total));
+ }
+
+ println!("---");
+ println!("[overhead summary]");
+ let (_, baseline_h, baseline_total) = &results[0];
+ let baseline_p50 = baseline_h.value_at_quantile(0.50) as f64 / 1000.0;
+ let baseline_p99 = baseline_h.value_at_quantile(0.99) as f64 / 1000.0;
+ let baseline_qps = *baseline_total as f64 / duration.as_secs_f64();
+
+ println!(
+ " {:38} p50 p99 QPS Δp50 Δp99 ΔQPS",
+ "config"
+ );
+ for (label, h, total) in &results {
+ let p50 = h.value_at_quantile(0.50) as f64 / 1000.0;
+ let p99 = h.value_at_quantile(0.99) as f64 / 1000.0;
+ let qps = *total as f64 / duration.as_secs_f64();
+ println!(
+ " {:38} {:>5.3} {:>5.3} {:>6.0} {:>+6.3} {:>+6.3} {:>+5.1}%",
+ label,
+ p50,
+ p99,
+ qps,
+ p50 - baseline_p50,
+ p99 - baseline_p99,
+ (qps - baseline_qps) / baseline_qps * 100.0,
+ );
+ }
+}
+
+// ---------- Scenario 6: Redis affinity store round-trip ---------------------
+
+/// Compare in-memory affinity store against a real Redis-backed one.
+/// Requires a Redis reachable at `--redis-url` (default
+/// redis://127.0.0.1:6379). Measures the extra round-trip cost of
+/// every lookup/bind hitting Redis.
+async fn scenario_redis_affinity(
+ workers: usize,
+ duration: Duration,
+ n_sessions: usize,
+ redis_url: String,
+) {
+ println!(
+ "[scenario] redis affinity store: workers={} duration={:?} url={}",
+ workers, duration, redis_url
+ );
+
+ println!(" baseline (memory affinity):");
+ let baseline = run_unary_internal(workers, duration, n_sessions, RigOpts::default()).await;
+ println!(" with Redis affinity:");
+ let with_redis = run_unary_internal(
+ workers,
+ duration,
+ n_sessions,
+ RigOpts {
+ affinity_redis_url: Some(redis_url),
+ ..RigOpts::default()
+ },
+ )
+ .await;
+
+ println!("---");
+ println!("[redis-affinity summary]");
+ let p50_b = baseline.0.value_at_quantile(0.50) as f64 / 1000.0;
+ let p99_b = baseline.0.value_at_quantile(0.99) as f64 / 1000.0;
+ let p50_r = with_redis.0.value_at_quantile(0.50) as f64 / 1000.0;
+ let p99_r = with_redis.0.value_at_quantile(0.99) as f64 / 1000.0;
+ let qps_b = baseline.1 as f64 / duration.as_secs_f64();
+ let qps_r = with_redis.1 as f64 / duration.as_secs_f64();
+ println!(
+ " p50 (ms): memory={:.3} redis={:.3} delta={:+.3}",
+ p50_b,
+ p50_r,
+ p50_r - p50_b
+ );
+ println!(
+ " p99 (ms): memory={:.3} redis={:.3} delta={:+.3}",
+ p99_b,
+ p99_r,
+ p99_r - p99_b
+ );
+ println!(
+ " QPS: memory={:.0} redis={:.0} delta={:+.1}%",
+ qps_b,
+ qps_r,
+ (qps_r - qps_b) / qps_b * 100.0
+ );
+}
+
+// ---------- Scenario 4: drain under load -----------------------------------
+
+async fn scenario_drain_under_load(
+ concurrency: usize,
+ messages: usize,
+ delay_ms: u64,
+ deadline_secs: u64,
+) {
+ println!(
+ "[scenario] drain under load: concurrency={} msgs/stream={} per-msg-delay={}ms deadline={}s",
+ concurrency, messages, delay_ms, deadline_secs
+ );
+ let mut rig = spawn_rig(RigOpts {
+ stream_messages: messages,
+ stream_delay_ms: delay_ms,
+ ..RigOpts::default()
+ })
+ .await;
+
+ let completed = Arc::new(AtomicU64::new(0));
+ let cancelled = Arc::new(AtomicU64::new(0));
+ let started = Instant::now();
+ let stream_handles: Vec<_> = (0..concurrency)
+ .map(|i| {
+ let ch = rig.grpc.clone();
+ let completed = completed.clone();
+ let cancelled = cancelled.clone();
+ tokio::spawn(async move {
+ let mut client =
+ pb::spark_connect_service_client::SparkConnectServiceClient::new(ch);
+ let session_id = format!("session-{}", i);
+ let stream = client
+ .execute_plan(Request::new(pb::ExecutePlanRequest {
+ session_id,
+ operation_id: Some(format!("op-{}", i)),
+ ..Default::default()
+ }))
+ .await;
+ let Ok(resp) = stream else {
+ cancelled.fetch_add(1, Ordering::Relaxed);
+ return;
+ };
+ let mut s = resp.into_inner();
+ let mut got_error = false;
+ while let Some(item) = s.next().await {
+ if item.is_err() {
+ got_error = true;
+ break;
+ }
+ }
+ if got_error {
+ cancelled.fetch_add(1, Ordering::Relaxed);
+ } else {
+ completed.fetch_add(1, Ordering::Relaxed);
+ }
+ })
+ })
+ .collect();
+
+ // Wait until the gateway sees all the streams in flight.
+ let target_active = concurrency as i64;
+ let fanout_deadline = Instant::now() + Duration::from_secs(10);
+ while rig.metrics.active_streams_value() < target_active && Instant::now() < fanout_deadline {
+ tokio::time::sleep(Duration::from_millis(20)).await;
+ }
+ let active_at_drain = rig.metrics.active_streams_value();
+ println!(
+ " active_streams at drain trigger: {} (target {})",
+ active_at_drain, target_active
+ );
+
+ // Trigger drain (mirrors gateway main).
+ let drain_start = Instant::now();
+ rig.readiness.mark_not_ready();
+ let metrics = rig.metrics.clone();
+ let drain_deadline = Duration::from_secs(deadline_secs);
+ let drained = tokio::time::timeout(drain_deadline, async {
+ loop {
+ if metrics.active_streams_value() <= 0 {
+ return;
+ }
+ tokio::time::sleep(Duration::from_millis(50)).await;
+ }
+ })
+ .await;
+ let drain_elapsed = drain_start.elapsed();
+ let drain_outcome = if drained.is_ok() {
+ "clean"
+ } else {
+ "deadline-hit"
+ };
+ let final_active = metrics.active_streams_value();
+
+ // Now actually shut down the gateway so streams that were
+ // cancelled by the deadline see their errors.
+ if let Some(tx) = rig.gw_kill.take() {
+ let _ = tx.send(());
+ }
+ for h in stream_handles {
+ let _ = h.await;
+ }
+ let total_elapsed = started.elapsed();
+
+ println!("---");
+ println!("[drain-under-load summary]");
+ println!(" concurrency: {}", concurrency);
+ println!(" drain outcome: {}", drain_outcome);
+ println!(" drain elapsed: {:?}", drain_elapsed);
+ println!(" active@deadline: {}", final_active);
+ println!(
+ " streams completed: {}",
+ completed.load(Ordering::Relaxed)
+ );
+ println!(
+ " streams cancelled: {}",
+ cancelled.load(Ordering::Relaxed)
+ );
+ println!(" test wall-clock: {:?}", total_elapsed);
+}
+
+// ---------- CLI -------------------------------------------------------------
+
+fn parse_arg<T: std::str::FromStr>(args: &[String], flag: &str, default: T) -> T {
+ let mut i = 0;
+ while i < args.len() {
+ if args[i] == flag {
+ if let Some(v) = args.get(i + 1) {
+ if let Ok(parsed) = v.parse() {
+ return parsed;
+ }
+ }
+ }
+ i += 1;
+ }
+ default
+}
+
+#[tokio::main(flavor = "multi_thread")]
+async fn main() {
+ let argv: Vec<String> = std::env::args().collect();
+ let scenario = argv.get(1).map(String::as_str).unwrap_or("unary");
+
+ let workers = parse_arg(&argv, "--workers", 32usize);
+ let duration_secs = parse_arg(&argv, "--duration-secs", 20u64);
+ let concurrency = parse_arg(&argv, "--concurrency", 64usize);
+ let messages = parse_arg(&argv, "--messages", 100usize);
+ let delay_ms = parse_arg(&argv, "--delay-ms", 0u64);
+ let deadline_secs = parse_arg(&argv, "--deadline-secs", 10u64);
+ let n_sessions = parse_arg(&argv, "--sessions", 64usize);
+
+ let redis_url = parse_arg::<String>(&argv, "--redis-url", "redis://127.0.0.1:6379".to_string());
+
+ match scenario {
+ "unary" => {
+ scenario_unary(workers, Duration::from_secs(duration_secs), n_sessions).await;
+ }
+ "streaming" => {
+ scenario_streaming(concurrency, messages, n_sessions).await;
+ }
+ "hc-overhead" => {
+ scenario_hc_overhead(workers, Duration::from_secs(duration_secs), n_sessions).await;
+ }
+ "drain-under-load" => {
+ scenario_drain_under_load(concurrency, messages, delay_ms, deadline_secs).await;
+ }
+ "overhead" => {
+ scenario_overhead(workers, Duration::from_secs(duration_secs), n_sessions).await;
+ }
+ "redis-affinity" => {
+ scenario_redis_affinity(
+ workers,
+ Duration::from_secs(duration_secs),
+ n_sessions,
+ redis_url,
+ )
+ .await;
+ }
+ other => {
+ eprintln!(
+ "unknown scenario: {} — try one of: unary streaming hc-overhead drain-under-load overhead redis-affinity",
+ other
+ );
+ std::process::exit(2);
+ }
+ }
+}
diff --git a/crates/proxy/src/config_filter.rs b/crates/proxy/src/config_filter.rs
new file mode 100644
index 0000000..2d21f75
--- /dev/null
+++ b/crates/proxy/src/config_filter.rs
@@ -0,0 +1,369 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Withholding backend-only secrets from `Config` RPC responses.
+//!
+//! The backend's `spark.connect.authenticate.token` is the shared
+//! secret between the gateway and its backends: the gateway presents
+//! it so that a client bypassing the gateway is refused. Spark's
+//! `SparkConnectConfigHandler` applies no denylist on the read path
+//! (`handleGet` / `handleGetOption` / `handleGetWithDefault` /
+//! `handleGetAll` all hand back whatever `RuntimeConfig` returns),
+//! and `SQLConf.mergeSparkConf` copies every `SparkConf` entry —
+//! static ones included — into the session config. A client can
+//! therefore ask the backend for the token and read it back.
+//!
+//! In Spark's own direct client-to-server model that discloses
+//! nothing: the client had to present the token to get connected.
+//! Behind a gateway the token means something else — the client
+//! authenticates with its own credential (JWT/OIDC) and is *not*
+//! supposed to hold the gateway↔backend secret — so relaying the
+//! value would hand any authorized user the means to bypass the
+//! gateway entirely.
+//!
+//! The gateway therefore strips the value on the way out, and does
+//! so unconditionally: a proxy cannot assume every backend it
+//! forwards to has been patched.
+//!
+//! Two details make this less trivial than a key comparison:
+//!
+//! * `GetAll { prefix }` returns keys with the **prefix stripped**
+//! (`handleGetAll` does `key.substring(prefix.length)`), so a
+//! request for prefix `spark.connect.` comes back as
+//! `authenticate.token`. Matching has to re-join the prefix, or a
+//! prefixed query walks straight past the filter.
+//! * Responses are filtered rather than requests. Dropping keys from
+//! the request would desynchronise positional expectations in
+//! clients that assume one response pair per requested key.
+//! * "Unset" does not look the same on every operation: `GetAll`
+//! omits the key, `Get` / `GetOption` return a pair with no value,
+//! and `GetWithDefault` resolves to the default the caller passed
+//! in. Substituting the wrong one still withholds the secret, but
+//! makes the gateway's filtering observable.
+
+use std::collections::HashMap;
+
+use scg_genproto::pb;
+
+/// Config keys the gateway never relays to a client. Compared
+/// case-insensitively — Spark's own lookup is case-sensitive, so a
+/// differently-cased request cannot retrieve the value anyway, but
+/// matching loosely costs nothing and removes a class of doubt.
+const WITHHELD_KEYS: &[&str] = &["spark.connect.authenticate.token"];
+
+fn is_withheld(key: &str) -> bool {
+ WITHHELD_KEYS
+ .iter()
+ .any(|blocked| key.eq_ignore_ascii_case(blocked))
+}
+
+/// What a withheld key should look like in the response, so that it
+/// is indistinguishable from a key the backend never had.
+enum Substitute {
+ /// `GetAll` enumerates whatever exists, so an absent key simply
+ /// is not listed. (Dropping the pair is also required, not just
+ /// tidier: the Spark client asserts `hasValue` on every pair
+ /// `GetAll` returns.)
+ Drop,
+ /// `Get` / `GetOption` report an unset key as a pair with no
+ /// value, and callers expect one pair per requested key.
+ NoValue,
+ /// `GetWithDefault` resolves an unset key to the default the
+ /// caller supplied, keyed here by config key. A key the caller
+ /// asked about with no default of its own maps to `None`.
+ Defaults(HashMap<String, Option<String>>),
+}
+
+/// Response-side redaction plan derived from a `ConfigRequest`.
+pub(crate) struct ConfigGuard {
+ /// Prefix `handleGetAll` stripped from the response keys, if the
+ /// operation was a prefixed `GetAll`.
+ prefix: String,
+ substitute: Substitute,
+}
+
+impl ConfigGuard {
+ pub(crate) fn for_request(req: &pb::ConfigRequest) -> Self {
+ use pb::config_request::operation::OpType;
+ match req.operation.as_ref().and_then(|o| o.op_type.as_ref()) {
+ Some(OpType::GetAll(get_all)) => Self {
+ prefix: get_all.prefix.clone().unwrap_or_default(),
+ substitute: Substitute::Drop,
+ },
+ Some(OpType::GetWithDefault(get_with_default)) => {
+ let defaults = get_with_default
+ .pairs
+ .iter()
+ .filter(|pair| is_withheld(&pair.key))
+ .map(|pair| (pair.key.clone(), pair.value.clone()))
+ .collect();
+ Self {
+ prefix: String::new(),
+ substitute: Substitute::Defaults(defaults),
+ }
+ }
+ _ => Self {
+ prefix: String::new(),
+ substitute: Substitute::NoValue,
+ },
+ }
+ }
+
+ /// Strip withheld values from `resp`. Returns the fully-qualified
+ /// keys that were redacted, for audit.
+ pub(crate) fn apply(&self, resp: &mut pb::ConfigResponse) -> Vec<String> {
+ let mut redacted = Vec::new();
+ // `GetAll` strips the prefix from the keys it returns, so the
+ // pair key alone is not the config key.
+ let full_key = |key: &str| format!("{}{}", self.prefix, key);
+
+ match &self.substitute {
+ Substitute::Drop => {
+ resp.pairs.retain(|pair| {
+ let key = full_key(&pair.key);
+ if is_withheld(&key) {
+ redacted.push(key);
+ false
+ } else {
+ true
+ }
+ });
+ }
+ Substitute::NoValue => {
+ for pair in &mut resp.pairs {
+ let key = full_key(&pair.key);
+ if is_withheld(&key) {
+ redacted.push(key);
+ pair.value = None;
+ }
+ }
+ }
+ Substitute::Defaults(defaults) => {
+ for pair in &mut resp.pairs {
+ let key = full_key(&pair.key);
+ if is_withheld(&key) {
+ redacted.push(key.clone());
+ pair.value = defaults.get(&key).cloned().flatten();
+ }
+ }
+ }
+ }
+ redacted
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ fn pair(key: &str, value: Option<&str>) -> pb::KeyValue {
+ pb::KeyValue {
+ key: key.into(),
+ value: value.map(str::to_string),
+ }
+ }
+
+ fn request(op: pb::config_request::operation::OpType) -> pb::ConfigRequest {
+ pb::ConfigRequest {
+ operation: Some(pb::config_request::Operation { op_type: Some(op) }),
+ ..Default::default()
+ }
+ }
+
+ fn get_request(keys: &[&str]) -> pb::ConfigRequest {
+ request(pb::config_request::operation::OpType::Get(
+ pb::config_request::Get {
+ keys: keys.iter().map(|k| k.to_string()).collect(),
+ },
+ ))
+ }
+
+ fn get_with_default_request(pairs: &[(&str, Option<&str>)]) -> pb::ConfigRequest {
+ request(pb::config_request::operation::OpType::GetWithDefault(
+ pb::config_request::GetWithDefault {
+ pairs: pairs.iter().map(|(k, v)| pair(k, *v)).collect(),
+ },
+ ))
+ }
+
+ fn get_all_request(prefix: Option<&str>) -> pb::ConfigRequest {
+ request(pb::config_request::operation::OpType::GetAll(
+ pb::config_request::GetAll {
+ prefix: prefix.map(str::to_string),
+ },
+ ))
+ }
+
+ fn response(pairs: Vec<pb::KeyValue>) -> pb::ConfigResponse {
+ pb::ConfigResponse {
+ pairs,
+ ..Default::default()
+ }
+ }
+
+ #[test]
+ fn get_of_the_token_returns_no_value() {
+ let guard = ConfigGuard::for_request(&get_request(&["spark.connect.authenticate.token"]));
+ let mut resp = response(vec![pair(
+ "spark.connect.authenticate.token",
+ Some("SUPERSECRET"),
+ )]);
+ let redacted = guard.apply(&mut resp);
+
+ assert_eq!(redacted, vec!["spark.connect.authenticate.token"]);
+ // One pair per requested key is preserved; the value is gone.
+ assert_eq!(resp.pairs.len(), 1);
+ assert_eq!(resp.pairs[0].key, "spark.connect.authenticate.token");
+ assert_eq!(resp.pairs[0].value, None);
+ }
+
+ #[test]
+ fn unrelated_keys_pass_through_untouched() {
+ let guard = ConfigGuard::for_request(&get_request(&[
+ "spark.sql.shuffle.partitions",
+ "spark.hadoop.fs.s3a.secret.key",
+ ]));
+ let mut resp = response(vec![
+ pair("spark.sql.shuffle.partitions", Some("200")),
+ // The operator's own secret — not ours to withhold; the
+ // client set it and Spark's redaction policy owns it.
+ pair("spark.hadoop.fs.s3a.secret.key", Some("user-secret")),
+ ]);
+ let redacted = guard.apply(&mut resp);
+
+ assert!(redacted.is_empty());
+ assert_eq!(resp.pairs[0].value.as_deref(), Some("200"));
+ assert_eq!(resp.pairs[1].value.as_deref(), Some("user-secret"));
+ }
+
+ #[test]
+ fn get_all_omits_the_token_entirely() {
+ let guard = ConfigGuard::for_request(&get_all_request(None));
+ let mut resp = response(vec![
+ pair("spark.sql.shuffle.partitions", Some("200")),
+ pair("spark.connect.authenticate.token", Some("SUPERSECRET")),
+ pair("spark.app.name", Some("demo")),
+ ]);
+ let redacted = guard.apply(&mut resp);
+
+ assert_eq!(redacted, vec!["spark.connect.authenticate.token"]);
+ let keys: Vec<&str> = resp.pairs.iter().map(|p| p.key.as_str()).collect();
+ assert_eq!(keys, vec!["spark.sql.shuffle.partitions", "spark.app.name"]);
+ }
+
+ #[test]
+ fn get_all_with_prefix_matches_on_the_rejoined_key() {
+ // `handleGetAll` returns keys with the prefix removed, so the
+ // response carries `authenticate.token`, not the full key.
+ // Matching the bare pair key would miss it entirely — this is
+ // the bypass this test exists to prevent.
+ let guard = ConfigGuard::for_request(&get_all_request(Some("spark.connect.")));
+ let mut resp = response(vec![
+ pair("authenticate.token", Some("SUPERSECRET")),
+ pair("session.planCache.enabled", Some("true")),
+ ]);
+ let redacted = guard.apply(&mut resp);
+
+ assert_eq!(redacted, vec!["spark.connect.authenticate.token"]);
+ let keys: Vec<&str> = resp.pairs.iter().map(|p| p.key.as_str()).collect();
+ assert_eq!(keys, vec!["session.planCache.enabled"]);
+ }
+
+ #[test]
+ fn get_all_with_a_deeper_prefix_still_matches() {
+ let guard = ConfigGuard::for_request(&get_all_request(Some("spark.connect.authenticate.")));
+ let mut resp = response(vec![pair("token", Some("SUPERSECRET"))]);
+ let redacted = guard.apply(&mut resp);
+
+ assert_eq!(redacted, vec!["spark.connect.authenticate.token"]);
+ assert!(resp.pairs.is_empty());
+ }
+
+ #[test]
+ fn get_all_with_unrelated_prefix_is_untouched() {
+ let guard = ConfigGuard::for_request(&get_all_request(Some("spark.sql.")));
+ let mut resp = response(vec![pair("shuffle.partitions", Some("200"))]);
+ let redacted = guard.apply(&mut resp);
+
+ assert!(redacted.is_empty());
+ assert_eq!(resp.pairs[0].value.as_deref(), Some("200"));
+ }
+
+ #[test]
+ fn case_variations_are_withheld_too() {
+ let guard = ConfigGuard::for_request(&get_request(&["SPARK.CONNECT.AUTHENTICATE.TOKEN"]));
+ let mut resp = response(vec![pair(
+ "SPARK.CONNECT.AUTHENTICATE.TOKEN",
+ Some("SUPERSECRET"),
+ )]);
+ let redacted = guard.apply(&mut resp);
+
+ assert_eq!(redacted.len(), 1);
+ assert_eq!(resp.pairs[0].value, None);
+ }
+
+ #[test]
+ fn get_with_default_resolves_to_the_callers_default() {
+ // What an unset key does on this operation: the caller's own
+ // default comes back. Returning no value would still withhold
+ // the secret, but it would tell the caller that something
+ // filtered the response.
+ let guard = ConfigGuard::for_request(&get_with_default_request(&[
+ ("spark.connect.authenticate.token", Some("fallback")),
+ ("spark.sql.shuffle.partitions", Some("100")),
+ ]));
+ let mut resp = response(vec![
+ pair("spark.connect.authenticate.token", Some("SUPERSECRET")),
+ pair("spark.sql.shuffle.partitions", Some("200")),
+ ]);
+ let redacted = guard.apply(&mut resp);
+
+ assert_eq!(redacted, vec!["spark.connect.authenticate.token"]);
+ assert_eq!(resp.pairs[0].value.as_deref(), Some("fallback"));
+ // The backend's answer for everything else is untouched — the
+ // caller's default is not a substitute for a real value.
+ assert_eq!(resp.pairs[1].value.as_deref(), Some("200"));
+ }
+
+ #[test]
+ fn get_with_default_without_a_default_yields_no_value() {
+ // `GetWithDefault` allows a pair with no value; an unset key
+ // then resolves to nothing.
+ let guard = ConfigGuard::for_request(&get_with_default_request(&[(
+ "spark.connect.authenticate.token",
+ None,
+ )]));
+ let mut resp = response(vec![pair(
+ "spark.connect.authenticate.token",
+ Some("SUPERSECRET"),
+ )]);
+ let redacted = guard.apply(&mut resp);
+
+ assert_eq!(redacted.len(), 1);
+ assert_eq!(resp.pairs[0].value, None);
+ }
+
+ #[test]
+ fn get_option_of_an_unset_key_is_unchanged() {
+ // A value-less pair is exactly what Spark returns for an
+ // unset key — the redacted shape must be indistinguishable.
+ let guard = ConfigGuard::for_request(&get_request(&["spark.not.set"]));
+ let mut resp = response(vec![pair("spark.not.set", None)]);
+ let redacted = guard.apply(&mut resp);
+
+ assert!(redacted.is_empty());
+ assert_eq!(resp.pairs[0].value, None);
+ }
+}
diff --git a/crates/proxy/src/dial.rs b/crates/proxy/src/dial.rs
new file mode 100644
index 0000000..716b302
--- /dev/null
+++ b/crates/proxy/src/dial.rs
@@ -0,0 +1,71 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Per-backend tonic Channel cache. We don't want to pay the dial cost on
+//! every RPC, so each unique backend address is opened lazily and reused.
+
+use parking_lot::Mutex;
+use std::collections::HashMap;
+use std::sync::Arc;
+use tonic::transport::{Channel, Endpoint};
+
+#[derive(Debug, thiserror::Error)]
+pub enum DialError {
+ #[error("invalid backend uri {addr}: {source}")]
+ BadUri {
+ addr: String,
+ #[source]
+ source: tonic::transport::Error,
+ },
+}
+
+#[derive(Default)]
+pub struct Dialer {
+ inner: Mutex<HashMap<String, Channel>>,
+}
+
+impl Dialer {
+ pub fn new() -> Arc<Self> {
+ Arc::new(Self::default())
+ }
+
+ /// Returns a (cached) lazy Channel for `addr`. The Channel is created in
+ /// "lazy connect" mode — actual TCP/HTTP/2 setup happens on first use,
+ /// and tonic transparently reconnects on transport failures.
+ pub fn channel(&self, addr: &str) -> Result<Channel, DialError> {
+ if let Some(ch) = self.inner.lock().get(addr).cloned() {
+ return Ok(ch);
+ }
+ let uri = ensure_scheme(addr);
+ let endpoint = Endpoint::from_shared(uri.clone()).map_err(|e| DialError::BadUri {
+ addr: addr.to_string(),
+ source: e,
+ })?;
+ // Lazy connect: tonic returns a Channel immediately and reconnects
+ // automatically on transport errors. This matches Go grpc.NewClient.
+ let channel = endpoint.connect_lazy();
+ self.inner.lock().insert(addr.to_string(), channel.clone());
+ Ok(channel)
+ }
+}
+
+/// Bare `host:port` strings need a scheme for `Endpoint::from_shared`.
+fn ensure_scheme(addr: &str) -> String {
+ if addr.starts_with("http://") || addr.starts_with("https://") {
+ addr.to_string()
+ } else {
+ format!("http://{}", addr)
+ }
+}
diff --git a/crates/proxy/src/handler.rs b/crates/proxy/src/handler.rs
new file mode 100644
index 0000000..a410f56
--- /dev/null
+++ b/crates/proxy/src/handler.rs
@@ -0,0 +1,1041 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! `SparkConnectService` implementation that forwards every RPC to a
+//! backend chosen by the [`Router`].
+
+use std::pin::Pin;
+use std::sync::Arc;
+
+use futures::{Stream, StreamExt};
+use opentelemetry::trace::TraceContextExt;
+use scg_audit::AuditLogger;
+use scg_auth::{AnonymousAuthenticator, AuthInterceptor, Identity};
+use scg_genproto::pb;
+use scg_observability::{
+ extract_parent, inject_context, request_id, Metrics, RpcGuard, REQUEST_ID_HEADER,
+};
+use scg_ratelimit::RateLimiter;
+use scg_routing::{Router, SessionKey};
+use scg_tenant::TenantResolver;
+use tokio_stream::wrappers::ReceiverStream;
+use tonic::metadata::{MetadataMap, MetadataValue};
+use tonic::{Request, Response, Status, Streaming};
+use tracing::{info, warn, Instrument, Span};
+use tracing_opentelemetry::OpenTelemetrySpanExt;
+
+use crate::config_filter::ConfigGuard;
+use crate::dial::Dialer;
+use crate::outbound::BackendTokens;
+
+/// gRPC handler implementing `SparkConnectService` as a forwarding proxy.
+///
+/// On every RPC the proxy:
+///
+/// 1. Starts a per-RPC [`RpcGuard`] that records duration + the final
+/// gRPC code on Drop.
+/// 2. Generates a correlation ID; stamps it on the tracing span and on
+/// the outbound `x-request-id` metadata.
+/// 3. Authenticates the inbound request (via the configured
+/// [`AuthInterceptor`]; defaults to `Anonymous` when auth is
+/// disabled in config).
+/// 4. Overwrites the request's `UserContext.user_id` with the verified
+/// identity — never trusting the value the client supplied.
+/// 5. Resolves a backend through the [`Router`].
+/// 6. Forwards the (rewritten) request and pumps the response back.
+pub struct SparkConnectProxy {
+ router: Arc<Router>,
+ dialer: Arc<Dialer>,
+ auth: AuthInterceptor,
+ /// Metrics handle. Required (we always want metrics): pass a
+ /// freshly built `Metrics::new()?` in tests where you don't care
+ /// about scraping.
+ metrics: Metrics,
+ /// Resolves the tenant for every RPC after auth. The routing key
+ /// is `(tenant, user_id, session_id)`; without a resolver every
+ /// session ends up in `tenant="default"`.
+ tenant_resolver: TenantResolver,
+ /// Per-tenant rate limiter. `None` skips the limiter check
+ /// entirely; a disabled limiter (`Some` with no buckets
+ /// enabled) is also free via [`RateLimiter::is_active`].
+ rate_limiter: Option<RateLimiter>,
+ /// Audit logger. Always present (defaults to disabled) so
+ /// handlers can call into it unconditionally without
+ /// `Option::map` ceremony.
+ audit: AuditLogger,
+ /// Per-pool `authorization: Bearer …` credentials presented on
+ /// the gateway→backend hop, for backends that enforce
+ /// `spark.connect.authenticate.token`. Defaults to none.
+ backend_tokens: BackendTokens,
+}
+
+impl SparkConnectProxy {
+ /// Build a proxy with auth disabled, a fresh throwaway
+ /// `Metrics`, and the default (back-compat) tenant resolver
+ /// (every RPC ends up in `tenant="default"`). Used by tests and
+ /// minimal single-tenant deployments.
+ ///
+ /// Production deployments use [`SparkConnectProxy::builder`] (or
+ /// the older [`with_auth_and_metrics`]) and hand in a `Metrics`
+ /// shared with the admin server.
+ pub fn new(router: Arc<Router>, dialer: Arc<Dialer>) -> Self {
+ let metrics = Metrics::new().expect("Metrics::new() in test scaffolding");
+ Self::with_auth_and_metrics(
+ router,
+ dialer,
+ AuthInterceptor::new(Arc::new(AnonymousAuthenticator)),
+ metrics,
+ )
+ }
+
+ /// Build a proxy with the given auth interceptor and a fresh
+ /// `Metrics`. Convenient for auth-only tests.
+ pub fn with_auth(router: Arc<Router>, dialer: Arc<Dialer>, auth: AuthInterceptor) -> Self {
+ let metrics = Metrics::new().expect("Metrics::new() in test scaffolding");
+ Self::with_auth_and_metrics(router, dialer, auth, metrics)
+ }
+
+ /// Convenience constructor: auth + metrics, default tenant
+ /// resolver, no rate limiter. Equivalent to calling
+ /// [`SparkConnectProxy::with_components`] with
+ /// `TenantResolver::new(TenantResolverConfig::default())` and
+ /// `rate_limiter = None`.
+ pub fn with_auth_and_metrics(
+ router: Arc<Router>,
+ dialer: Arc<Dialer>,
+ auth: AuthInterceptor,
+ metrics: Metrics,
+ ) -> Self {
+ Self::with_components(
+ router,
+ dialer,
+ auth,
+ metrics,
+ TenantResolver::new(scg_tenant::TenantResolverConfig::default()),
+ )
+ }
+
+ /// Production constructor with tenant resolver but no rate
+ /// limiter and a disabled audit logger. Most existing tests
+ /// and examples use this.
+ pub fn with_components(
+ router: Arc<Router>,
+ dialer: Arc<Dialer>,
+ auth: AuthInterceptor,
+ metrics: Metrics,
+ tenant_resolver: TenantResolver,
+ ) -> Self {
+ Self::with_all(
+ router,
+ dialer,
+ auth,
+ metrics,
+ tenant_resolver,
+ None,
+ AuditLogger::disabled(),
+ )
+ }
+
+ /// Full constructor: every component. Used by
+ /// `crates/gateway/main` once the operator's config is loaded.
+ pub fn with_all(
+ router: Arc<Router>,
+ dialer: Arc<Dialer>,
+ auth: AuthInterceptor,
+ metrics: Metrics,
+ tenant_resolver: TenantResolver,
+ rate_limiter: Option<RateLimiter>,
+ audit: AuditLogger,
+ ) -> Self {
+ Self {
+ router,
+ dialer,
+ auth,
+ metrics,
+ tenant_resolver,
+ rate_limiter,
+ audit,
+ backend_tokens: BackendTokens::none(),
+ }
+ }
+
+ /// Attach per-pool backend credentials (chainable). Without this
+ /// the gateway sends no `authorization` header to backends.
+ pub fn with_backend_tokens(mut self, tokens: BackendTokens) -> Self {
+ self.backend_tokens = tokens;
+ self
+ }
+
+ pub fn metrics(&self) -> &Metrics {
+ &self.metrics
+ }
+
+ /// Stamp everything an outbound request needs: correlation ID,
+ /// W3C traceparent, and — when configured — the backend pool's
+ /// `authorization: Bearer …` credential for `tenant`.
+ fn stamp_outbound<T>(&self, req: &mut Request<T>, request_id: &str, tenant: &str) {
+ stamp_propagation(req, request_id);
+ self.backend_tokens.apply(tenant, req);
+ }
+
+ fn client(
+ &self,
+ addr: &str,
+ ) -> Result<
+ pb::spark_connect_service_client::SparkConnectServiceClient<tonic::transport::Channel>,
+ Status,
+ > {
+ let ch = self
+ .dialer
+ .channel(addr)
+ .map_err(|e| Status::unavailable(format!("dial backend {}: {}", addr, e)))?;
+ Ok(pb::spark_connect_service_client::SparkConnectServiceClient::new(ch))
+ }
+
+ /// Authenticate an inbound RPC. Returns the verified `Identity`
+ /// or the `Status::unauthenticated` the proxy should hand back
+ /// to the client. Auth failures bump `scg_auth_failures_total`
+ /// and emit an `auth.failure` audit event.
+ async fn authenticate(
+ &self,
+ metadata: &MetadataMap,
+ rid: &str,
+ rpc: &'static str,
+ ) -> Result<Arc<Identity>, Status> {
+ match self.auth.authenticate(metadata).await {
+ Ok(ext) => Ok(ext.0),
+ Err(status) => {
+ let reason = auth_failure_reason(&status);
+ self.metrics.record_auth_failure(reason);
+ self.audit.auth_failure(rid, rpc, reason);
+ Err(status)
+ }
+ }
+ }
+
+ /// Authenticate and resolve the tenant for an inbound RPC. Used
+ /// by every RPC handler — the tenant becomes the first segment
+ /// of the routing key. A tenant-resolver `Reject` policy on a
+ /// missing tenant surfaces as `Status::unauthenticated`.
+ ///
+ /// When a rate limiter is configured, this also takes a token
+ /// from the (tenant, user) bucket pair. Quota violations bubble
+ /// up as `Status::resource_exhausted` to the client and bump
+ /// `scg_rate_limit_rejected_total{tenant, scope}`.
+ ///
+ /// `rid` and `rpc` are threaded through so auth-failure audit
+ /// events carry the correlation ID and the RPC name. Tenant
+ /// resolution failures (configured-source yielded nothing and
+ /// `on_missing = Reject`) are treated as auth failures for
+ /// audit / metric purposes — the JWT may have been signed
+ /// correctly, but without a tenant the gateway can't route,
+ /// and the operational signal is the same as a bad token:
+ /// "this caller is not authorised to enter."
+ async fn authenticate_and_resolve(
+ &self,
+ metadata: &MetadataMap,
+ rid: &str,
+ rpc: &'static str,
+ ) -> Result<(Arc<Identity>, String), Status> {
+ let identity = self.authenticate(metadata, rid, rpc).await?;
+ let tenant = match self.tenant_resolver.resolve(metadata, &identity) {
+ Ok(t) => t,
+ Err(status) => {
+ self.metrics.record_auth_failure("missing_tenant");
+ self.audit.auth_failure(rid, rpc, "missing_tenant");
+ return Err(status);
+ }
+ };
+ if let Some(limiter) = &self.rate_limiter {
+ limiter.check(&tenant, &identity.user_id).await?;
+ }
+ Ok((identity, tenant))
+ }
+
+ /// Resolve a session through the router and, when the binding
+ /// is fresh, emit a `session.create` audit event. Callers use
+ /// this in place of `router.resolve_session` whenever they
+ /// want audit coverage; the empty-session_id path (zero key)
+ /// never emits the event because there's nothing to record.
+ async fn resolve_session_audited(
+ &self,
+ key: &SessionKey,
+ rid: &str,
+ identity: &Identity,
+ ) -> Result<Option<String>, Status> {
+ let outcome = self.router.resolve_session_detailed(key).await?;
+ if let Some(o) = &outcome {
+ if o.newly_bound {
+ self.audit.session_create(
+ rid,
+ &key.tenant,
+ &identity.user_id,
+ &identity.groups,
+ &key.session_id,
+ &o.addr,
+ );
+ }
+ }
+ Ok(outcome.map(|o| o.addr))
+ }
+}
+
+/// Map an auth Status into a small, fixed-cardinality reason label.
+fn auth_failure_reason(status: &Status) -> &'static str {
+ let m = status.message();
+ if m.contains("missing") {
+ "missing_token"
+ } else if m.contains("expired") {
+ "expired"
+ } else if m.contains("kid") {
+ "unknown_kid"
+ } else if m.contains("invalid") {
+ "invalid_token"
+ } else {
+ "unknown"
+ }
+}
+
+/// Build the routing key from a verified identity, a resolved
+/// tenant, and the request's session id. The identity *replaces*
+/// whatever `user_id` the client claimed in `UserContext`; the
+/// tenant comes from the [`TenantResolver`] (auth claim, gRPC
+/// metadata header, or fixed deployment-wide string).
+fn key_from_identity(tenant: &str, session_id: &str, id: &Identity) -> SessionKey {
+ SessionKey::with_tenant(tenant, id.user_id.as_str(), session_id)
+}
+
+/// Overwrite (or create) `UserContext.user_id` to match the verified
+/// identity. Backend Spark Connect servers consume this value as part
+/// of their SparkSession key, so trusting the client value would let
+/// callers impersonate one another.
+fn stamp_user_context(uc: &mut Option<pb::UserContext>, id: &Identity) {
+ match uc {
+ Some(ctx) => ctx.user_id = id.user_id.clone(),
+ None => {
+ *uc = Some(pb::UserContext {
+ user_id: id.user_id.clone(),
+ ..Default::default()
+ });
+ }
+ }
+}
+
+/// Convert a possibly-empty backend selection into an `Unavailable` Status.
+/// A None is most commonly produced when a dynamic pool (the K8s
+/// Endpoints watcher) hasn't seen any healthy endpoint yet, e.g.
+/// during gateway boot before the watcher's initial list event.
+fn require_addr(addr: Result<Option<String>, Status>) -> Result<String, Status> {
+ // Two failure modes:
+ // * `Err(Status)` — the tenant has no configured pool and the
+ // policy is `Reject`. Forward the `PermissionDenied` to the
+ // client unchanged.
+ // * `Ok(None)` — the tenant's pool has no healthy backend
+ // (K8s discovery during startup, Redis down so we can't
+ // resolve stickiness, etc.). Surface as `Unavailable`.
+ addr?.ok_or_else(|| Status::unavailable("no healthy backend available"))
+}
+
+/// Map a Status code into the small fixed string used in metric labels.
+fn status_code_label(status: &Status) -> &'static str {
+ match status.code() {
+ tonic::Code::Ok => "OK",
+ tonic::Code::Cancelled => "Cancelled",
+ tonic::Code::Unknown => "Unknown",
+ tonic::Code::InvalidArgument => "InvalidArgument",
+ tonic::Code::DeadlineExceeded => "DeadlineExceeded",
+ tonic::Code::NotFound => "NotFound",
+ tonic::Code::AlreadyExists => "AlreadyExists",
+ tonic::Code::PermissionDenied => "PermissionDenied",
+ tonic::Code::ResourceExhausted => "ResourceExhausted",
+ tonic::Code::FailedPrecondition => "FailedPrecondition",
+ tonic::Code::Aborted => "Aborted",
+ tonic::Code::OutOfRange => "OutOfRange",
+ tonic::Code::Unimplemented => "Unimplemented",
+ tonic::Code::Internal => "Internal",
+ tonic::Code::Unavailable => "Unavailable",
+ tonic::Code::DataLoss => "DataLoss",
+ tonic::Code::Unauthenticated => "Unauthenticated",
+ }
+}
+
+/// Inject the gateway's correlation ID onto an outbound request's
+/// metadata so backend logs can correlate.
+fn stamp_request_id<T>(req: &mut Request<T>, request_id: &str) {
+ if let Ok(v) = MetadataValue::try_from(request_id) {
+ req.metadata_mut().insert(REQUEST_ID_HEADER, v);
+ }
+}
+
+/// Build a per-RPC tracing span and parent it to whatever traceparent
+/// the inbound metadata carries (no-op when absent — the span becomes a
+/// fresh trace root).
+///
+/// Returns the span; the caller `Instrument`s the handler future with
+/// it. Span name follows the OpenTelemetry semantic-conventions form
+/// `<service>/<rpc>` so distributed-trace UIs render it nicely.
+fn rpc_span(rpc: &'static str, rid: &str, inbound: &MetadataMap) -> Span {
+ // Field names use underscores rather than dots: Span attribute
+ // names with `.` round-trip through tracing-opentelemetry, but
+ // some feature combos in the dependency graph silently drop them
+ // before they reach the OTel exporter. Use snake_case here and
+ // map to the dotted OTel semantic-conventions name in
+ // post-processing if needed.
+ let span = tracing::info_span!(
+ "scg_rpc",
+ rpc_method = rpc,
+ rpc_system = "grpc",
+ rpc_service = "spark.connect.SparkConnectService",
+ scg_rid = %rid,
+ );
+ let parent_cx = extract_parent(inbound);
+ if parent_cx.span().span_context().is_valid() {
+ // Failure here just means the OTel layer isn't installed
+ // (e.g. tracing is disabled in config) — the span is still
+ // useful for the JSON formatter, so we silently ignore.
+ let _ = span.set_parent(parent_cx);
+ }
+ span
+}
+
+/// Stamp both the gateway's correlation ID and the current span's
+/// W3C traceparent onto an outbound request — backend logs / spans
+/// can now join the same trace as the gateway.
+fn stamp_propagation<T>(req: &mut Request<T>, request_id: &str) {
+ stamp_request_id(req, request_id);
+ let cx = Span::current().context();
+ inject_context(&cx, req.metadata_mut());
+}
+
+/// Update an `RpcGuard` with the final code from a `Status` (or "OK"
+/// when the result is Ok).
+fn finalise_guard<T>(guard: &mut RpcGuard, result: &Result<T, Status>) {
+ match result {
+ Ok(_) => guard.record("OK"),
+ Err(s) => guard.record(status_code_label(s)),
+ }
+}
+
+/// Hold the (tenant, user_id) captured once
+/// `authenticate_and_resolve` succeeded, so audit emissions on the
+/// way out (`rpc.ok`, `rpc.error`) can carry them. `None` means the
+/// handler hit a failure *before* identity was known (typically an
+/// `Unauthenticated` from auth) — those are already covered by
+/// `auth.failure`, so we skip the rpc-level audit event to avoid
+/// duplicate noise.
+#[derive(Default)]
+struct AuditCtx {
+ identity: Option<(String, String)>,
+}
+
+impl AuditCtx {
+ fn set(&mut self, tenant: &str, user_id: &str) {
+ self.identity = Some((tenant.to_string(), user_id.to_string()));
+ }
+}
+
+/// Update the metrics guard *and* emit the matching audit event
+/// (`rpc.ok` if `log_successful_rpcs` is on, otherwise just `rpc.error`
+/// for non-OK results). The Cancelled code is filtered out: clients
+/// dropping streams are noisy and not security-relevant.
+fn finalise_rpc<T>(
+ guard: &mut RpcGuard,
+ result: &Result<T, Status>,
+ audit: &AuditLogger,
+ ctx: &AuditCtx,
+ rid: &str,
+ rpc: &'static str,
+) {
+ finalise_guard(guard, result);
+ let Some((tenant, user_id)) = ctx.identity.as_ref() else {
+ return;
+ };
+ match result {
+ Ok(_) => audit.rpc_ok(rid, rpc, tenant, user_id),
+ Err(s) if s.code() == tonic::Code::Cancelled => {}
+ Err(s) => audit.rpc_error(rid, rpc, tenant, user_id, status_code_label(s), s.message()),
+ }
+}
+
+/// Forward a tonic streaming response to a fresh server-stream sent through
+/// `tx`. We don't pipe directly to the inbound `Streaming` because we need
+/// to return early on the first error and report it as a Status.
+type StreamItem<T> = Result<T, Status>;
+
+#[tonic::async_trait]
+impl pb::spark_connect_service_server::SparkConnectService for SparkConnectProxy {
+ type ExecutePlanStream =
+ Pin<Box<dyn Stream<Item = StreamItem<pb::ExecutePlanResponse>> + Send + 'static>>;
+ type ReattachExecuteStream = Self::ExecutePlanStream;
+
+ // ----- Unary RPCs ----------------------------------------------------
+
+ async fn analyze_plan(
+ &self,
+ req: Request<pb::AnalyzePlanRequest>,
+ ) -> Result<Response<pb::AnalyzePlanResponse>, Status> {
+ let mut guard = self.metrics.rpc_guard("AnalyzePlan");
+ let rid = request_id();
+ let span = rpc_span("AnalyzePlan", &rid, req.metadata());
+ let mut audit_ctx = AuditCtx::default();
+ let result: Result<Response<pb::AnalyzePlanResponse>, Status> = async {
+ let (identity, tenant) = self.authenticate_and_resolve(req.metadata(), &rid, "AnalyzePlan").await?;
+ audit_ctx.set(&tenant, &identity.user_id);
+ let mut body = req.into_inner();
+ stamp_user_context(&mut body.user_context, &identity);
+ let key = key_from_identity(&tenant, &body.session_id, &identity);
+ let addr = require_addr(self.resolve_session_audited(&key, &rid, &identity).await)?;
+ info!(rid = %rid, rpc = "AnalyzePlan", user = %identity.user_id, session = %key.session_id, %addr, "forwarding");
+ let mut c = self.client(&addr)?;
+ let mut outbound = Request::new(body);
+ self.stamp_outbound(&mut outbound, &rid, &tenant);
+ c.analyze_plan(outbound).await
+ }
+ .instrument(span)
+ .await;
+ finalise_rpc(
+ &mut guard,
+ &result,
+ &self.audit,
+ &audit_ctx,
+ &rid,
+ "AnalyzePlan",
+ );
+ result
+ }
+
+ async fn config(
+ &self,
+ req: Request<pb::ConfigRequest>,
+ ) -> Result<Response<pb::ConfigResponse>, Status> {
+ let mut guard = self.metrics.rpc_guard("Config");
+ let rid = request_id();
+ let span = rpc_span("Config", &rid, req.metadata());
+ let mut audit_ctx = AuditCtx::default();
+ let result: Result<Response<pb::ConfigResponse>, Status> = async {
+ let (identity, tenant) = self.authenticate_and_resolve(req.metadata(), &rid, "Config").await?;
+ audit_ctx.set(&tenant, &identity.user_id);
+ let mut body = req.into_inner();
+ stamp_user_context(&mut body.user_context, &identity);
+ let key = key_from_identity(&tenant, &body.session_id, &identity);
+ let addr = require_addr(self.resolve_session_audited(&key, &rid, &identity).await)?;
+ info!(rid = %rid, rpc = "Config", user = %identity.user_id, session = %key.session_id, %addr, "forwarding");
+ // Spark's Config handler will hand back any key it holds,
+ // including the gateway↔backend pre-shared token. Decide
+ // what to withhold before the body is moved.
+ let guard = ConfigGuard::for_request(&body);
+ let mut c = self.client(&addr)?;
+ let mut outbound = Request::new(body);
+ self.stamp_outbound(&mut outbound, &rid, &tenant);
+ let mut resp = c.config(outbound).await?;
+ for withheld in guard.apply(resp.get_mut()) {
+ self.audit
+ .config_redacted(&rid, &tenant, &identity.user_id, &withheld);
+ warn!(
+ rid = %rid,
+ user = %identity.user_id,
+ tenant = %tenant,
+ key = %withheld,
+ "Config: withheld a backend-only key from the client"
+ );
+ }
+ Ok(resp)
+ }
+ .instrument(span)
+ .await;
+ finalise_rpc(&mut guard, &result, &self.audit, &audit_ctx, &rid, "Config");
+ result
+ }
+
+ async fn artifact_status(
+ &self,
+ req: Request<pb::ArtifactStatusesRequest>,
+ ) -> Result<Response<pb::ArtifactStatusesResponse>, Status> {
+ let mut guard = self.metrics.rpc_guard("ArtifactStatus");
+ let rid = request_id();
+ let span = rpc_span("ArtifactStatus", &rid, req.metadata());
+ let mut audit_ctx = AuditCtx::default();
+ let result: Result<Response<pb::ArtifactStatusesResponse>, Status> = async {
+ let (identity, tenant) = self
+ .authenticate_and_resolve(req.metadata(), &rid, "ArtifactStatus")
+ .await?;
+ audit_ctx.set(&tenant, &identity.user_id);
+ let mut body = req.into_inner();
+ stamp_user_context(&mut body.user_context, &identity);
+ let key = key_from_identity(&tenant, &body.session_id, &identity);
+ let addr = require_addr(self.resolve_session_audited(&key, &rid, &identity).await)?;
+ let mut c = self.client(&addr)?;
+ let mut outbound = Request::new(body);
+ self.stamp_outbound(&mut outbound, &rid, &tenant);
+ c.artifact_status(outbound).await
+ }
+ .instrument(span)
+ .await;
+ finalise_rpc(
+ &mut guard,
+ &result,
+ &self.audit,
+ &audit_ctx,
+ &rid,
+ "ArtifactStatus",
+ );
+ result
+ }
+
+ async fn interrupt(
+ &self,
+ req: Request<pb::InterruptRequest>,
+ ) -> Result<Response<pb::InterruptResponse>, Status> {
+ let mut guard = self.metrics.rpc_guard("Interrupt");
+ let rid = request_id();
+ let span = rpc_span("Interrupt", &rid, req.metadata());
+ let mut audit_ctx = AuditCtx::default();
+ let result: Result<Response<pb::InterruptResponse>, Status> = async {
+ let (identity, tenant) = self
+ .authenticate_and_resolve(req.metadata(), &rid, "Interrupt")
+ .await?;
+ audit_ctx.set(&tenant, &identity.user_id);
+ let mut body = req.into_inner();
+ stamp_user_context(&mut body.user_context, &identity);
+ let key = key_from_identity(&tenant, &body.session_id, &identity);
+ // Interrupt may target a specific operation id (one of several
+ // InterruptType variants); when present, route by op id.
+ let op_id = match body.interrupt.as_ref() {
+ Some(pb::interrupt_request::Interrupt::OperationId(id)) => id.clone(),
+ _ => String::new(),
+ };
+ let addr = require_addr(self.router.resolve_op(&op_id, &key).await)?;
+ let mut c = self.client(&addr)?;
+ let mut outbound = Request::new(body);
+ self.stamp_outbound(&mut outbound, &rid, &tenant);
+ c.interrupt(outbound).await
+ }
+ .instrument(span)
+ .await;
+ finalise_rpc(
+ &mut guard,
+ &result,
+ &self.audit,
+ &audit_ctx,
+ &rid,
+ "Interrupt",
+ );
+ result
+ }
+
+ async fn release_execute(
+ &self,
+ req: Request<pb::ReleaseExecuteRequest>,
+ ) -> Result<Response<pb::ReleaseExecuteResponse>, Status> {
+ let mut guard = self.metrics.rpc_guard("ReleaseExecute");
+ let rid = request_id();
+ let span = rpc_span("ReleaseExecute", &rid, req.metadata());
+ let mut audit_ctx = AuditCtx::default();
+ let result: Result<Response<pb::ReleaseExecuteResponse>, Status> = async {
+ let (identity, tenant) = self
+ .authenticate_and_resolve(req.metadata(), &rid, "ReleaseExecute")
+ .await?;
+ audit_ctx.set(&tenant, &identity.user_id);
+ let mut body = req.into_inner();
+ stamp_user_context(&mut body.user_context, &identity);
+ let key = key_from_identity(&tenant, &body.session_id, &identity);
+ let op_id = body.operation_id.clone();
+ let addr = require_addr(self.router.resolve_op(&op_id, &key).await)?;
+ let mut c = self.client(&addr)?;
+ let mut outbound = Request::new(body);
+ self.stamp_outbound(&mut outbound, &rid, &tenant);
+ let resp = c.release_execute(outbound).await?;
+ // On a successful release the server has dropped the operation, so
+ // we drop our reverse-index entry too.
+ self.router.forget_op(&op_id).await;
+ Ok(resp)
+ }
+ .instrument(span)
+ .await;
+ finalise_rpc(
+ &mut guard,
+ &result,
+ &self.audit,
+ &audit_ctx,
+ &rid,
+ "ReleaseExecute",
+ );
+ result
+ }
+
+ async fn release_session(
+ &self,
+ req: Request<pb::ReleaseSessionRequest>,
+ ) -> Result<Response<pb::ReleaseSessionResponse>, Status> {
+ let mut guard = self.metrics.rpc_guard("ReleaseSession");
+ let rid = request_id();
+ let span = rpc_span("ReleaseSession", &rid, req.metadata());
+ let mut audit_ctx = AuditCtx::default();
+ let result: Result<Response<pb::ReleaseSessionResponse>, Status> = async {
+ let (identity, tenant) = self
+ .authenticate_and_resolve(req.metadata(), &rid, "ReleaseSession")
+ .await?;
+ audit_ctx.set(&tenant, &identity.user_id);
+ let mut body = req.into_inner();
+ stamp_user_context(&mut body.user_context, &identity);
+ let key = key_from_identity(&tenant, &body.session_id, &identity);
+ let addr = require_addr(self.resolve_session_audited(&key, &rid, &identity).await)?;
+ let mut c = self.client(&addr)?;
+ let mut outbound = Request::new(body);
+ self.stamp_outbound(&mut outbound, &rid, &tenant);
+ let resp = c.release_session(outbound).await?;
+ self.router.forget_session(&key).await;
+ self.audit.session_release(
+ &rid,
+ &key.tenant,
+ &identity.user_id,
+ &identity.groups,
+ &key.session_id,
+ &addr,
+ );
+ Ok(resp)
+ }
+ .instrument(span)
+ .await;
+ finalise_rpc(
+ &mut guard,
+ &result,
+ &self.audit,
+ &audit_ctx,
+ &rid,
+ "ReleaseSession",
+ );
+ result
+ }
+
+ async fn fetch_error_details(
+ &self,
+ req: Request<pb::FetchErrorDetailsRequest>,
+ ) -> Result<Response<pb::FetchErrorDetailsResponse>, Status> {
+ let mut guard = self.metrics.rpc_guard("FetchErrorDetails");
+ let rid = request_id();
+ let span = rpc_span("FetchErrorDetails", &rid, req.metadata());
+ let mut audit_ctx = AuditCtx::default();
+ let result: Result<Response<pb::FetchErrorDetailsResponse>, Status> = async {
+ let (identity, tenant) = self
+ .authenticate_and_resolve(req.metadata(), &rid, "FetchErrorDetails")
+ .await?;
+ audit_ctx.set(&tenant, &identity.user_id);
+ let mut body = req.into_inner();
+ stamp_user_context(&mut body.user_context, &identity);
+ let key = key_from_identity(&tenant, &body.session_id, &identity);
+ let addr = require_addr(self.resolve_session_audited(&key, &rid, &identity).await)?;
+ let mut c = self.client(&addr)?;
+ let mut outbound = Request::new(body);
+ self.stamp_outbound(&mut outbound, &rid, &tenant);
+ c.fetch_error_details(outbound).await
+ }
+ .instrument(span)
+ .await;
+ finalise_rpc(
+ &mut guard,
+ &result,
+ &self.audit,
+ &audit_ctx,
+ &rid,
+ "FetchErrorDetails",
+ );
+ result
+ }
+
+ async fn clone_session(
+ &self,
+ req: Request<pb::CloneSessionRequest>,
+ ) -> Result<Response<pb::CloneSessionResponse>, Status> {
+ let mut guard = self.metrics.rpc_guard("CloneSession");
+ let rid = request_id();
+ let span = rpc_span("CloneSession", &rid, req.metadata());
+ let mut audit_ctx = AuditCtx::default();
+ let result: Result<Response<pb::CloneSessionResponse>, Status> = async {
+ let (identity, tenant) = self
+ .authenticate_and_resolve(req.metadata(), &rid, "CloneSession")
+ .await?;
+ audit_ctx.set(&tenant, &identity.user_id);
+ let mut body = req.into_inner();
+ stamp_user_context(&mut body.user_context, &identity);
+ let key = key_from_identity(&tenant, &body.session_id, &identity);
+ let addr = require_addr(self.resolve_session_audited(&key, &rid, &identity).await)?;
+ let mut c = self.client(&addr)?;
+ let mut outbound = Request::new(body);
+ self.stamp_outbound(&mut outbound, &rid, &tenant);
+ let resp = c.clone_session(outbound).await?;
+
+ // The cloned session's state lives on the driver that
+ // executed the clone. Bind the new session id to that same
+ // backend so follow-up RPCs route there instead of being
+ // placed on an arbitrary backend by pool selection.
+ let new_sid = resp.get_ref().new_session_id.clone();
+ if !new_sid.is_empty() {
+ let clone_key = key_from_identity(&tenant, &new_sid, &identity);
+ let winner = self.router.remember_session(&clone_key, addr.clone()).await;
+ if winner == addr {
+ self.audit.session_create(
+ &rid,
+ &tenant,
+ &identity.user_id,
+ &identity.groups,
+ &new_sid,
+ &addr,
+ );
+ } else {
+ // A racing RPC already bound the cloned id to a
+ // different backend — the clone's state is on
+ // `addr` but routing will follow `winner`. Surface
+ // loudly; the next RPC on the cloned session will
+ // see single-server unknown-session semantics.
+ warn!(
+ rid = %rid,
+ cloned_session = %new_sid,
+ clone_backend = %addr,
+ bound_backend = %winner,
+ "CloneSession: cloned session id was already bound elsewhere"
+ );
+ }
+ }
+ Ok(resp)
+ }
+ .instrument(span)
+ .await;
+ finalise_rpc(
+ &mut guard,
+ &result,
+ &self.audit,
+ &audit_ctx,
+ &rid,
+ "CloneSession",
+ );
+ result
+ }
+
+ async fn get_status(
+ &self,
+ req: Request<pb::GetStatusRequest>,
+ ) -> Result<Response<pb::GetStatusResponse>, Status> {
+ let mut guard = self.metrics.rpc_guard("GetStatus");
+ let rid = request_id();
+ let span = rpc_span("GetStatus", &rid, req.metadata());
+ let mut audit_ctx = AuditCtx::default();
+ let result: Result<Response<pb::GetStatusResponse>, Status> = async {
+ let (identity, tenant) = self
+ .authenticate_and_resolve(req.metadata(), &rid, "GetStatus")
+ .await?;
+ audit_ctx.set(&tenant, &identity.user_id);
+ let mut body = req.into_inner();
+ stamp_user_context(&mut body.user_context, &identity);
+ let key = key_from_identity(&tenant, &body.session_id, &identity);
+ let addr = require_addr(self.resolve_session_audited(&key, &rid, &identity).await)?;
+ let mut c = self.client(&addr)?;
+ let mut outbound = Request::new(body);
+ self.stamp_outbound(&mut outbound, &rid, &tenant);
+ c.get_status(outbound).await
+ }
+ .instrument(span)
+ .await;
+ finalise_rpc(
+ &mut guard,
+ &result,
+ &self.audit,
+ &audit_ctx,
+ &rid,
+ "GetStatus",
+ );
+ result
+ }
+
+ // ----- Server-streaming RPCs ----------------------------------------
+
+ async fn execute_plan(
+ &self,
+ req: Request<pb::ExecutePlanRequest>,
+ ) -> Result<Response<Self::ExecutePlanStream>, Status> {
+ let mut guard = self.metrics.rpc_guard("ExecutePlan");
+ // The stream guard is moved into the forwarded stream below so it
+ // outlives the response — a streaming RPC's "lifetime" for
+ // scg_active_streams is the lifetime of the *stream*, not of
+ // this handler function.
+ let stream_guard = self.metrics.stream_guard();
+ let rid = request_id();
+ let span = rpc_span("ExecutePlan", &rid, req.metadata());
+ let metrics = self.metrics.clone();
+ let mut audit_ctx = AuditCtx::default();
+ let result: Result<Response<Self::ExecutePlanStream>, Status> = async {
+ let (identity, tenant) = self.authenticate_and_resolve(req.metadata(), &rid, "ExecutePlan").await?;
+ audit_ctx.set(&tenant, &identity.user_id);
+ let mut body = req.into_inner();
+ stamp_user_context(&mut body.user_context, &identity);
+ let key = key_from_identity(&tenant, &body.session_id, &identity);
+ let addr = require_addr(self.resolve_session_audited(&key, &rid, &identity).await)?;
+ info!(rid = %rid, rpc = "ExecutePlan", user = %identity.user_id, session = %key.session_id, %addr, "forwarding stream");
+
+ // Bind operation_id → backend so a follow-up ReattachExecute reaches
+ // the same driver even if its session id is missing or has been
+ // forgotten by the affinity cache.
+ if let Some(op_id) = body.operation_id.clone() {
+ if !op_id.is_empty() {
+ self.router.remember_op(op_id, addr.clone()).await;
+ }
+ }
+
+ let mut c = self.client(&addr)?;
+ let mut outbound = Request::new(body);
+ self.stamp_outbound(&mut outbound, &rid, &tenant);
+ let upstream = c.execute_plan(outbound).await?.into_inner();
+ Ok(Response::new(forward_server_stream(upstream, stream_guard)))
+ }
+ .instrument(span)
+ .await;
+ // If we never built a stream (early error), the guard would be
+ // dropped by the .await above; otherwise the stream owns it now.
+ // Either way active_streams accounting is correct.
+ let _ = metrics;
+ finalise_rpc(
+ &mut guard,
+ &result,
+ &self.audit,
+ &audit_ctx,
+ &rid,
+ "ExecutePlan",
+ );
+ result
+ }
+
+ async fn reattach_execute(
+ &self,
+ req: Request<pb::ReattachExecuteRequest>,
+ ) -> Result<Response<Self::ReattachExecuteStream>, Status> {
+ let mut guard = self.metrics.rpc_guard("ReattachExecute");
+ let stream_guard = self.metrics.stream_guard();
+ let rid = request_id();
+ let span = rpc_span("ReattachExecute", &rid, req.metadata());
+ let mut audit_ctx = AuditCtx::default();
+ let result: Result<Response<Self::ReattachExecuteStream>, Status> = async {
+ let (identity, tenant) = self
+ .authenticate_and_resolve(req.metadata(), &rid, "ReattachExecute")
+ .await?;
+ audit_ctx.set(&tenant, &identity.user_id);
+ let mut body = req.into_inner();
+ stamp_user_context(&mut body.user_context, &identity);
+ let key = key_from_identity(&tenant, &body.session_id, &identity);
+ let addr = require_addr(self.router.resolve_op(&body.operation_id, &key).await)?;
+ let mut c = self.client(&addr)?;
+ let mut outbound = Request::new(body);
+ self.stamp_outbound(&mut outbound, &rid, &tenant);
+ let upstream = c.reattach_execute(outbound).await?.into_inner();
+ Ok(Response::new(forward_server_stream(upstream, stream_guard)))
+ }
+ .instrument(span)
+ .await;
+ finalise_rpc(
+ &mut guard,
+ &result,
+ &self.audit,
+ &audit_ctx,
+ &rid,
+ "ReattachExecute",
+ );
+ result
+ }
+
+ // ----- Client-streaming RPCs ----------------------------------------
+
+ async fn add_artifacts(
+ &self,
+ req: Request<Streaming<pb::AddArtifactsRequest>>,
+ ) -> Result<Response<pb::AddArtifactsResponse>, Status> {
+ let mut guard = self.metrics.rpc_guard("AddArtifacts");
+ let _stream_guard = self.metrics.stream_guard();
+ let rid = request_id();
+ let span = rpc_span("AddArtifacts", &rid, req.metadata());
+ let mut audit_ctx = AuditCtx::default();
+ let result: Result<Response<pb::AddArtifactsResponse>, Status> = async {
+ // Authenticate from the request metadata before we touch the
+ // streaming body — the credential lives in HTTP/2 headers, not
+ // in any of the message frames.
+ let (identity, tenant) = self
+ .authenticate_and_resolve(req.metadata(), &rid, "AddArtifacts")
+ .await?;
+ audit_ctx.set(&tenant, &identity.user_id);
+ let mut inbound = req.into_inner();
+
+ // We need the first message to make the routing decision, then we
+ // forward it plus the remainder to the chosen backend.
+ let mut first = inbound
+ .message()
+ .await?
+ .ok_or_else(|| Status::invalid_argument("AddArtifacts: empty client stream"))?;
+ stamp_user_context(&mut first.user_context, &identity);
+
+ let key = key_from_identity(&tenant, &first.session_id, &identity);
+ let addr = require_addr(self.resolve_session_audited(&key, &rid, &identity).await)?;
+ let mut c = self.client(&addr)?;
+
+ let (tx, rx) = tokio::sync::mpsc::channel::<pb::AddArtifactsRequest>(8);
+ tx.send(first)
+ .await
+ .map_err(|_| Status::cancelled("backend closed"))?;
+
+ tokio::spawn(async move {
+ while let Ok(Some(m)) = inbound.message().await {
+ if tx.send(m).await.is_err() {
+ break;
+ }
+ }
+ });
+
+ let mut outbound = Request::new(ReceiverStream::new(rx));
+ self.stamp_outbound(&mut outbound, &rid, &tenant);
+ let resp = c.add_artifacts(outbound).await?;
+ Ok(resp)
+ }
+ .instrument(span)
+ .await;
+ if let Err(ref e) = result {
+ warn!(rid = %rid, error = %e, "AddArtifacts failed");
+ }
+ finalise_rpc(
+ &mut guard,
+ &result,
+ &self.audit,
+ &audit_ctx,
+ &rid,
+ "AddArtifacts",
+ );
+ result
+ }
+}
+
+fn forward_server_stream<T: Send + 'static>(
+ upstream: Streaming<T>,
+ // The stream guard must outlive the stream, not the handler
+ // function — otherwise scg_active_streams returns to 0 the
+ // moment we hand the Response back, even though the stream is
+ // still flowing. Capture it in the async_stream so it drops
+ // when the *stream* is dropped (i.e. when the client closes,
+ // the stream completes, or it errors).
+ stream_guard: scg_observability::StreamGuard,
+) -> Pin<Box<dyn Stream<Item = StreamItem<T>> + Send + 'static>> {
+ Box::pin(async_stream::stream! {
+ let _guard = stream_guard;
+ let mut upstream = upstream;
+ while let Some(item) = upstream.next().await {
+ yield item;
+ }
+ })
+}
diff --git a/crates/proxy/src/lib.rs b/crates/proxy/src/lib.rs
new file mode 100644
index 0000000..1968a46
--- /dev/null
+++ b/crates/proxy/src/lib.rs
@@ -0,0 +1,41 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! gRPC proxy for the Spark Connect service.
+//!
+//! Request flow:
+//!
+//! 1. Accept inbound `spark.connect.SparkConnectService` traffic via tonic.
+//! 2. For each RPC, derive a [`SessionKey`] (and an `operation_id` where
+//! applicable) from the request.
+//! 3. Ask the [`Router`] for the backend address. The router consults the
+//! affinity store first; on miss, it picks from the pool and records the
+//! decision so subsequent calls stick to the same backend.
+//! 4. Open or reuse a tonic [`Channel`] to that backend (via [`Dialer`]),
+//! forward the request, and pump any response stream back to the client.
+//!
+//! All RPCs in the Spark Connect surface are forwarded explicitly; new
+//! RPCs added by upstream Spark will surface here as `Unimplemented`
+//! until they are wired in. There is no generic tower-level
+//! passthrough.
+
+mod config_filter;
+mod dial;
+mod handler;
+mod outbound;
+
+pub use dial::Dialer;
+pub use handler::SparkConnectProxy;
+pub use outbound::BackendTokens;
diff --git a/crates/proxy/src/outbound.rs b/crates/proxy/src/outbound.rs
new file mode 100644
index 0000000..7b3425f
--- /dev/null
+++ b/crates/proxy/src/outbound.rs
@@ -0,0 +1,179 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Credentials the gateway presents on the gateway→backend hop.
+//!
+//! Backends started with `spark.connect.authenticate.token` (Spark
+//! 4.0+) reject any request whose `authorization` metadata is not
+//! exactly `Bearer <token>`. Holding that token *only* in the
+//! gateway is how a deployment enforces the trust boundary: a
+//! client that dials a backend directly is refused with
+//! `UNAUTHENTICATED`, while gateway-mediated traffic carries the
+//! credential.
+//!
+//! Tokens are per pool. Lookup mirrors pool selection: a tenant
+//! with its own entry uses it; every other tenant (including
+//! unknown tenants routed to the default pool) uses the default
+//! token; no default means no header.
+
+use std::collections::HashMap;
+
+use tonic::metadata::MetadataValue;
+use tonic::Request;
+
+/// A configured token can't be carried in an `authorization` header.
+/// Raised at startup, never per-request. The token value itself is
+/// deliberately absent from the message.
+#[derive(Debug, thiserror::Error)]
+#[error("backend token for pool `{pool}` is not a valid header value (must be visible ASCII)")]
+pub struct InvalidBackendToken {
+ pub pool: String,
+}
+
+/// Pre-encoded `authorization` values, one per pool that has a
+/// token configured. Built once at startup; applied to every
+/// outbound request.
+#[derive(Default)]
+pub struct BackendTokens {
+ default: Option<MetadataValue<tonic::metadata::Ascii>>,
+ per_tenant: HashMap<String, MetadataValue<tonic::metadata::Ascii>>,
+}
+
+impl BackendTokens {
+ /// No tokens anywhere — the gateway sends no backend credential.
+ pub fn none() -> Self {
+ Self::default()
+ }
+
+ /// Build from resolved token strings. `per_tenant` only needs
+ /// entries for tenants whose pool uses a *different* token than
+ /// the default; everything else falls back to `default`.
+ pub fn new(
+ default: Option<String>,
+ per_tenant: HashMap<String, String>,
+ ) -> Result<Self, InvalidBackendToken> {
+ let encode = |pool: &str, token: String| -> Result<_, InvalidBackendToken> {
+ let err = || InvalidBackendToken {
+ pool: pool.to_string(),
+ };
+ // HeaderValue would accept any opaque byte ≥ 0x20, but a
+ // token the backend can never match is a misconfiguration
+ // — Spark compares against `Bearer <token>` as a string.
+ if !token.is_ascii() {
+ return Err(err());
+ }
+ let mut v = MetadataValue::try_from(format!("Bearer {token}")).map_err(|_| err())?;
+ // Keep the credential out of h2 header logging.
+ v.set_sensitive(true);
+ Ok(v)
+ };
+ let default = default.map(|t| encode("default", t)).transpose()?;
+ let per_tenant = per_tenant
+ .into_iter()
+ .map(|(t, tok)| {
+ let v = encode(&t, tok)?;
+ Ok((t, v))
+ })
+ .collect::<Result<HashMap<_, _>, InvalidBackendToken>>()?;
+ Ok(Self {
+ default,
+ per_tenant,
+ })
+ }
+
+ /// True when at least one pool has a token configured.
+ pub fn is_configured(&self) -> bool {
+ self.default.is_some() || !self.per_tenant.is_empty()
+ }
+
+ /// Number of tenant-specific token entries.
+ pub fn tenant_override_count(&self) -> usize {
+ self.per_tenant.len()
+ }
+
+ /// Stamp the `authorization` header for `tenant`'s pool onto an
+ /// outbound request. No-op when neither the tenant nor the
+ /// default has a token.
+ pub fn apply<T>(&self, tenant: &str, req: &mut Request<T>) {
+ let value = self.per_tenant.get(tenant).or(self.default.as_ref());
+ if let Some(v) = value {
+ req.metadata_mut().insert("authorization", v.clone());
+ }
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ fn header<T>(req: &Request<T>) -> Option<&str> {
+ req.metadata()
+ .get("authorization")
+ .map(|v| v.to_str().expect("ascii"))
+ }
+
+ #[test]
+ fn no_tokens_is_a_noop() {
+ let tokens = BackendTokens::none();
+ assert!(!tokens.is_configured());
+ let mut req = Request::new(());
+ tokens.apply("any-tenant", &mut req);
+ assert!(header(&req).is_none());
+ }
+
+ #[test]
+ fn default_token_applies_to_every_tenant() {
+ let tokens = BackendTokens::new(Some("shared".into()), HashMap::new()).unwrap();
+ for tenant in ["default", "team-a", "unknown"] {
+ let mut req = Request::new(());
+ tokens.apply(tenant, &mut req);
+ assert_eq!(header(&req), Some("Bearer shared"));
+ }
+ }
+
+ #[test]
+ fn tenant_entry_wins_over_default() {
+ let per_tenant = HashMap::from([("team-a".to_string(), "a-token".to_string())]);
+ let tokens = BackendTokens::new(Some("shared".into()), per_tenant).unwrap();
+ let mut req = Request::new(());
+ tokens.apply("team-a", &mut req);
+ assert_eq!(header(&req), Some("Bearer a-token"));
+ let mut req = Request::new(());
+ tokens.apply("team-b", &mut req);
+ assert_eq!(header(&req), Some("Bearer shared"));
+ }
+
+ #[test]
+ fn tenant_entry_without_default_leaves_others_bare() {
+ let per_tenant = HashMap::from([("team-a".to_string(), "a-token".to_string())]);
+ let tokens = BackendTokens::new(None, per_tenant).unwrap();
+ let mut req = Request::new(());
+ tokens.apply("team-b", &mut req);
+ assert!(header(&req).is_none());
+ }
+
+ #[test]
+ fn header_is_marked_sensitive() {
+ let tokens = BackendTokens::new(Some("s".into()), HashMap::new()).unwrap();
+ let mut req = Request::new(());
+ tokens.apply("t", &mut req);
+ assert!(req.metadata().get("authorization").unwrap().is_sensitive());
+ }
+
+ #[test]
+ fn non_ascii_token_is_rejected_at_build() {
+ assert!(BackendTokens::new(Some("秘密".into()), HashMap::new()).is_err());
+ }
+}
diff --git a/crates/proxy/tests/audit_integration.rs b/crates/proxy/tests/audit_integration.rs
new file mode 100644
index 0000000..ba91483
--- /dev/null
+++ b/crates/proxy/tests/audit_integration.rs
@@ -0,0 +1,458 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Integration test for structured audit logging.
+//!
+//! Drives RPCs through a real gRPC server with an enabled
+//! [`AuditLogger`] and asserts the four default event types fire
+//! end-to-end:
+//!
+//! * `session.create` — first RPC on a fresh `(tenant, user, session)`
+//! * `session.release` — `ReleaseSession` succeeds
+//! * `auth.failure` — auth interceptor rejects a missing token
+//! * `rpc.error` — handler returns a non-OK Status
+//!
+//! Events are captured via a `tracing_subscriber::Layer` that filters
+//! on `target = "scg::audit"` — the same target the production JSON
+//! formatter would pick up. We don't install the JSON formatter here
+//! because we want structured access to fields, not stringified lines.
+
+mod common;
+
+use std::pin::Pin;
+use std::sync::Arc;
+use std::time::Duration;
+
+use common::AuditCapture;
+use futures::Stream;
+use scg_audit::{AuditConfig, AuditLogger};
+use scg_auth::token::{StaticTokenAuthenticator, TokenEntry};
+use scg_auth::{AnonymousAuthenticator, AuthInterceptor};
+use scg_genproto::pb;
+use scg_observability::Metrics;
+use scg_pool_static::StaticPool;
+use scg_proxy::{Dialer, SparkConnectProxy};
+use scg_routing::{AffinityStore, Pool, Router, TenantRouter};
+use scg_store_memory::MemoryStore;
+use scg_tenant::{OnMissing, TenantResolver, TenantResolverConfig, TenantSource};
+use tokio::net::TcpListener;
+use tokio_stream::wrappers::TcpListenerStream;
+use tonic::metadata::MetadataValue;
+use tonic::transport::{Channel, Endpoint, Server};
+use tonic::{Request, Response, Status};
+
+/// Backend that either OKs the request (Config / ReleaseSession) or
+/// returns a non-OK status the gateway must surface as `rpc.error`.
+#[derive(Default)]
+struct PartialBackend;
+
+#[tonic::async_trait]
+impl pb::spark_connect_service_server::SparkConnectService for PartialBackend {
+ type ExecutePlanStream =
+ Pin<Box<dyn Stream<Item = Result<pb::ExecutePlanResponse, Status>> + Send + 'static>>;
+ type ReattachExecuteStream = Self::ExecutePlanStream;
+
+ async fn config(
+ &self,
+ req: Request<pb::ConfigRequest>,
+ ) -> Result<Response<pb::ConfigResponse>, Status> {
+ let body = req.into_inner();
+ Ok(Response::new(pb::ConfigResponse {
+ session_id: body.session_id,
+ ..Default::default()
+ }))
+ }
+ async fn release_session(
+ &self,
+ req: Request<pb::ReleaseSessionRequest>,
+ ) -> Result<Response<pb::ReleaseSessionResponse>, Status> {
+ let _ = req.into_inner();
+ Ok(Response::new(pb::ReleaseSessionResponse::default()))
+ }
+ async fn analyze_plan(
+ &self,
+ _: Request<pb::AnalyzePlanRequest>,
+ ) -> Result<Response<pb::AnalyzePlanResponse>, Status> {
+ // Deliberately fail so we can verify `rpc.error` fires.
+ Err(Status::internal("boom"))
+ }
+ async fn artifact_status(
+ &self,
+ _: Request<pb::ArtifactStatusesRequest>,
+ ) -> Result<Response<pb::ArtifactStatusesResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn interrupt(
+ &self,
+ _: Request<pb::InterruptRequest>,
+ ) -> Result<Response<pb::InterruptResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn release_execute(
+ &self,
+ _: Request<pb::ReleaseExecuteRequest>,
+ ) -> Result<Response<pb::ReleaseExecuteResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn fetch_error_details(
+ &self,
+ _: Request<pb::FetchErrorDetailsRequest>,
+ ) -> Result<Response<pb::FetchErrorDetailsResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn clone_session(
+ &self,
+ _: Request<pb::CloneSessionRequest>,
+ ) -> Result<Response<pb::CloneSessionResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn get_status(
+ &self,
+ _: Request<pb::GetStatusRequest>,
+ ) -> Result<Response<pb::GetStatusResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn execute_plan(
+ &self,
+ _: Request<pb::ExecutePlanRequest>,
+ ) -> Result<Response<Self::ExecutePlanStream>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn reattach_execute(
+ &self,
+ _: Request<pb::ReattachExecuteRequest>,
+ ) -> Result<Response<Self::ReattachExecuteStream>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn add_artifacts(
+ &self,
+ _: Request<tonic::Streaming<pb::AddArtifactsRequest>>,
+ ) -> Result<Response<pb::AddArtifactsResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+}
+
+struct Rig {
+ channel: Channel,
+ _be_shutdown: tokio::sync::oneshot::Sender<()>,
+ _gw_shutdown: tokio::sync::oneshot::Sender<()>,
+}
+
+async fn spawn_rig(auth: AuthInterceptor) -> Rig {
+ let resolver = TenantResolver::new(TenantResolverConfig {
+ source: TenantSource::FromMetadata {
+ header: "x-tenant".into(),
+ },
+ on_missing: OnMissing::UseDefault,
+ default_name: "default".into(),
+ });
+ spawn_rig_with_resolver(auth, resolver).await
+}
+
+async fn spawn_rig_with_resolver(auth: AuthInterceptor, resolver: TenantResolver) -> Rig {
+ let be_lis = TcpListener::bind("127.0.0.1:0").await.unwrap();
+ let be_addr = be_lis.local_addr().unwrap().to_string();
+ let (be_tx, be_rx) = tokio::sync::oneshot::channel();
+ tokio::spawn(async move {
+ Server::builder()
+ .add_service(
+ pb::spark_connect_service_server::SparkConnectServiceServer::new(PartialBackend),
+ )
+ .serve_with_incoming_shutdown(TcpListenerStream::new(be_lis), async {
+ let _ = be_rx.await;
+ })
+ .await
+ .ok();
+ });
+
+ let metrics = Metrics::new().unwrap();
+ let pool: Arc<dyn Pool> = Arc::new(StaticPool::new(vec![be_addr]).unwrap());
+ let store: Arc<dyn AffinityStore> = Arc::new(MemoryStore::new());
+ let tr = TenantRouter::single(pool);
+ let router = Arc::new(Router::new(tr, store));
+ let dialer = Dialer::new();
+ let audit = AuditLogger::new(AuditConfig {
+ enabled: true,
+ log_successful_rpcs: false,
+ });
+ let proxy = SparkConnectProxy::with_all(router, dialer, auth, metrics, resolver, None, audit);
+
+ let gw_lis = TcpListener::bind("127.0.0.1:0").await.unwrap();
+ let gw_addr = gw_lis.local_addr().unwrap();
+ let (gw_tx, gw_rx) = tokio::sync::oneshot::channel();
+ tokio::spawn(async move {
+ Server::builder()
+ .add_service(pb::spark_connect_service_server::SparkConnectServiceServer::new(proxy))
+ .serve_with_incoming_shutdown(TcpListenerStream::new(gw_lis), async {
+ let _ = gw_rx.await;
+ })
+ .await
+ .ok();
+ });
+
+ let endpoint = Endpoint::from_shared(format!("http://{}", gw_addr)).unwrap();
+ let channel = endpoint
+ .connect_timeout(Duration::from_secs(2))
+ .connect()
+ .await
+ .unwrap();
+ Rig {
+ channel,
+ _be_shutdown: be_tx,
+ _gw_shutdown: gw_tx,
+ }
+}
+
+fn config_request(session: &str, tenant: &str) -> Request<pb::ConfigRequest> {
+ let mut req = Request::new(pb::ConfigRequest {
+ session_id: session.into(),
+ ..Default::default()
+ });
+ req.metadata_mut()
+ .insert("x-tenant", MetadataValue::try_from(tenant).unwrap());
+ req
+}
+
+fn analyze_request(session: &str, tenant: &str) -> Request<pb::AnalyzePlanRequest> {
+ let mut req = Request::new(pb::AnalyzePlanRequest {
+ session_id: session.into(),
+ ..Default::default()
+ });
+ req.metadata_mut()
+ .insert("x-tenant", MetadataValue::try_from(tenant).unwrap());
+ req
+}
+
+fn release_request(session: &str, tenant: &str) -> Request<pb::ReleaseSessionRequest> {
+ let mut req = Request::new(pb::ReleaseSessionRequest {
+ session_id: session.into(),
+ ..Default::default()
+ });
+ req.metadata_mut()
+ .insert("x-tenant", MetadataValue::try_from(tenant).unwrap());
+ req
+}
+
+#[tokio::test]
+async fn session_create_fires_once_per_binding() {
+ let cap = AuditCapture::lease();
+ let auth = AuthInterceptor::new(Arc::new(AnonymousAuthenticator));
+ let rig = spawn_rig(auth).await;
+ let mut c =
+ pb::spark_connect_service_client::SparkConnectServiceClient::new(rig.channel.clone());
+
+ // First Config on (default, anonymous, sess-1) → session.create fires.
+ c.config(config_request("sess-1", "default")).await.unwrap();
+ // Second Config on the same key → router hits the cached binding;
+ // session.create must NOT fire again.
+ c.config(config_request("sess-1", "default")).await.unwrap();
+ // A different session id → fresh binding → another session.create.
+ c.config(config_request("sess-2", "default")).await.unwrap();
+
+ let events = cap.snapshot();
+ assert_eq!(
+ cap.count_events("session.create"),
+ 2,
+ "expected exactly two session.create events (sess-1 + sess-2), got events: {:?}",
+ events.iter().map(|e| &e.fields).collect::<Vec<_>>(),
+ );
+ let first = cap.find_event("session.create").unwrap();
+ assert_eq!(
+ first.fields.get("tenant").map(|s| s.as_str()),
+ Some("default")
+ );
+ assert_eq!(
+ first.fields.get("session_id").map(|s| s.as_str()),
+ Some("sess-1")
+ );
+ // Anonymous auth → empty groups. The field still appears (as ""),
+ // so consumers can rely on its presence regardless of auth shape.
+ assert_eq!(first.fields.get("groups").map(|s| s.as_str()), Some(""));
+}
+
+#[tokio::test]
+async fn session_create_carries_groups_from_token() {
+ // Static-token auth with a token that declares groups. The audit
+ // event must carry them so operators querying the audit stream
+ // can see who has which memberships.
+ let cap = AuditCapture::lease();
+ let auth_inner = StaticTokenAuthenticator::new(vec![TokenEntry {
+ token: "dev-token".into(),
+ user_id: "alice".into(),
+ tenant: None,
+ groups: vec!["devs".into(), "admins".into()],
+ }])
+ .unwrap();
+ let auth = AuthInterceptor::new(Arc::new(auth_inner));
+ let rig = spawn_rig(auth).await;
+ let mut c =
+ pb::spark_connect_service_client::SparkConnectServiceClient::new(rig.channel.clone());
+
+ let mut req = config_request("sess-grp", "default");
+ req.metadata_mut().insert(
+ "authorization",
+ MetadataValue::try_from("Bearer dev-token").unwrap(),
+ );
+ c.config(req).await.unwrap();
+
+ let evt = cap.find_event("session.create").unwrap();
+ assert_eq!(evt.fields.get("user_id").map(|s| s.as_str()), Some("alice"));
+ assert_eq!(
+ evt.fields.get("groups").map(|s| s.as_str()),
+ Some("devs,admins"),
+ );
+}
+
+#[tokio::test]
+async fn release_session_emits_release_event() {
+ let cap = AuditCapture::lease();
+ let auth = AuthInterceptor::new(Arc::new(AnonymousAuthenticator));
+ let rig = spawn_rig(auth).await;
+ let mut c =
+ pb::spark_connect_service_client::SparkConnectServiceClient::new(rig.channel.clone());
+
+ c.config(config_request("sess-rel", "default"))
+ .await
+ .unwrap();
+ c.release_session(release_request("sess-rel", "default"))
+ .await
+ .unwrap();
+
+ assert_eq!(cap.count_events("session.release"), 1);
+ let rel = cap.find_event("session.release").unwrap();
+ assert_eq!(
+ rel.fields.get("session_id").map(|s| s.as_str()),
+ Some("sess-rel")
+ );
+}
+
+#[tokio::test]
+async fn auth_failure_emits_audit_event() {
+ let cap = AuditCapture::lease();
+ // Static-token auth with one valid token; an unauthenticated
+ // client request should fail.
+ let auth_inner = StaticTokenAuthenticator::new(vec![TokenEntry {
+ token: "secret".into(),
+ user_id: "alice".into(),
+ tenant: None,
+ groups: vec![],
+ }])
+ .unwrap();
+ let auth = AuthInterceptor::new(Arc::new(auth_inner));
+ let rig = spawn_rig(auth).await;
+ let mut c =
+ pb::spark_connect_service_client::SparkConnectServiceClient::new(rig.channel.clone());
+
+ // No Authorization metadata → Unauthenticated → auth.failure.
+ let err = c
+ .config(config_request("sess-x", "default"))
+ .await
+ .unwrap_err();
+ assert_eq!(err.code(), tonic::Code::Unauthenticated);
+
+ assert_eq!(cap.count_events("auth.failure"), 1);
+ let auth_evt = cap.find_event("auth.failure").unwrap();
+ assert_eq!(
+ auth_evt.fields.get("rpc").map(|s| s.as_str()),
+ Some("Config")
+ );
+ assert!(auth_evt.fields.contains_key("reason"));
+ // Pre-identity failure: no rpc.error follow-up (auth.failure
+ // already covers this).
+ assert_eq!(cap.count_events("rpc.error"), 0);
+}
+
+#[tokio::test]
+async fn rpc_error_emits_audit_event_with_code() {
+ let cap = AuditCapture::lease();
+ let auth = AuthInterceptor::new(Arc::new(AnonymousAuthenticator));
+ let rig = spawn_rig(auth).await;
+ let mut c =
+ pb::spark_connect_service_client::SparkConnectServiceClient::new(rig.channel.clone());
+
+ // Backend returns Internal("boom") on AnalyzePlan.
+ let err = c
+ .analyze_plan(analyze_request("sess-err", "default"))
+ .await
+ .unwrap_err();
+ assert_eq!(err.code(), tonic::Code::Internal);
+
+ assert_eq!(cap.count_events("rpc.error"), 1);
+ let rpc_err = cap.find_event("rpc.error").unwrap();
+ assert_eq!(
+ rpc_err.fields.get("rpc").map(|s| s.as_str()),
+ Some("AnalyzePlan")
+ );
+ assert_eq!(
+ rpc_err.fields.get("code").map(|s| s.as_str()),
+ Some("Internal")
+ );
+}
+
+/// Auth succeeds (anonymous) but the tenant resolver is configured
+/// for `FromMetadata` + `Reject`, and the inbound RPC carries no
+/// `x-tenant` header. The proxy should reject with `Unauthenticated`
+/// AND emit a single `auth.failure` audit event with
+/// `reason=missing_tenant` — symmetric to how a missing bearer token
+/// emits `reason=missing_token`. Before #105 this rejection path
+/// produced only a `tracing::warn!` line and nothing on the audit
+/// pipeline.
+#[tokio::test]
+async fn missing_tenant_under_reject_emits_audit_event() {
+ let cap = AuditCapture::lease();
+ let resolver = TenantResolver::new(TenantResolverConfig {
+ source: TenantSource::FromMetadata {
+ header: "x-tenant".into(),
+ },
+ on_missing: OnMissing::Reject,
+ default_name: "default".into(),
+ });
+ let auth = AuthInterceptor::new(Arc::new(AnonymousAuthenticator));
+ let rig = spawn_rig_with_resolver(auth, resolver).await;
+ let mut c =
+ pb::spark_connect_service_client::SparkConnectServiceClient::new(rig.channel.clone());
+
+ // No x-tenant metadata → tenant_resolver rejects → Unauthenticated.
+ // Build a Config request without the helper (which would set
+ // x-tenant) so the resolver sees an empty MetadataMap.
+ let bare = Request::new(pb::ConfigRequest {
+ session_id: "sess-no-tenant".into(),
+ ..Default::default()
+ });
+ let err = c.config(bare).await.unwrap_err();
+ assert_eq!(err.code(), tonic::Code::Unauthenticated);
+ assert!(
+ err.message().contains("tenant"),
+ "expected tenant-resolution error, got: {}",
+ err.message()
+ );
+
+ // Exactly one auth.failure event with reason=missing_tenant.
+ assert_eq!(cap.count_events("auth.failure"), 1);
+ let auth_evt = cap.find_event("auth.failure").unwrap();
+ assert_eq!(
+ auth_evt.fields.get("reason").map(|s| s.as_str()),
+ Some("missing_tenant")
+ );
+ assert_eq!(
+ auth_evt.fields.get("rpc").map(|s| s.as_str()),
+ Some("Config")
+ );
+ // Pre-routing failure: no rpc.error follow-up (the auth.failure
+ // event already covers this RPC).
+ assert_eq!(cap.count_events("rpc.error"), 0);
+ // And no session.create, since we never reached the binding step.
+ assert_eq!(cap.count_events("session.create"), 0);
+}
diff --git a/crates/proxy/tests/auth_integration.rs b/crates/proxy/tests/auth_integration.rs
new file mode 100644
index 0000000..83ed683
--- /dev/null
+++ b/crates/proxy/tests/auth_integration.rs
@@ -0,0 +1,296 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Integration test: proxy with auth enabled rejects unauthenticated
+//! callers and stamps the verified user_id onto requests forwarded to
+//! the backend (proving client-supplied user_id is *not* trusted).
+
+use std::pin::Pin;
+use std::sync::atomic::{AtomicI64, Ordering};
+use std::sync::Arc;
+use std::time::Duration;
+
+use futures::Stream;
+use parking_lot::Mutex;
+use scg_auth::{
+ token::{StaticTokenAuthenticator, TokenEntry},
+ AuthInterceptor,
+};
+use scg_genproto::pb;
+use scg_pool_static::StaticPool;
+use scg_proxy::{Dialer, SparkConnectProxy};
+use scg_routing::{AffinityStore, Pool, Router};
+use scg_store_memory::MemoryStore;
+use tokio::net::TcpListener;
+use tokio_stream::wrappers::TcpListenerStream;
+use tonic::metadata::MetadataValue;
+use tonic::transport::{Channel, Endpoint, Server};
+use tonic::{Request, Response, Status};
+
+/// A backend that records each Config request's UserContext so the
+/// test can assert the proxy stamped the right user_id onto it.
+#[derive(Default)]
+struct RecordingBackend {
+ last_user_id: Mutex<Option<String>>,
+ seen: AtomicI64,
+}
+
+#[tonic::async_trait]
+impl pb::spark_connect_service_server::SparkConnectService for RecordingBackend {
+ type ExecutePlanStream =
+ Pin<Box<dyn Stream<Item = Result<pb::ExecutePlanResponse, Status>> + Send + 'static>>;
+ type ReattachExecuteStream = Self::ExecutePlanStream;
+
+ async fn config(
+ &self,
+ req: Request<pb::ConfigRequest>,
+ ) -> Result<Response<pb::ConfigResponse>, Status> {
+ self.seen.fetch_add(1, Ordering::SeqCst);
+ let body = req.into_inner();
+ *self.last_user_id.lock() = body.user_context.as_ref().map(|u| u.user_id.clone());
+ Ok(Response::new(pb::ConfigResponse {
+ session_id: body.session_id,
+ ..Default::default()
+ }))
+ }
+
+ // Stub the rest — only Config is exercised in these tests.
+ async fn analyze_plan(
+ &self,
+ _: Request<pb::AnalyzePlanRequest>,
+ ) -> Result<Response<pb::AnalyzePlanResponse>, Status> {
+ Err(Status::unimplemented("not used"))
+ }
+ async fn artifact_status(
+ &self,
+ _: Request<pb::ArtifactStatusesRequest>,
+ ) -> Result<Response<pb::ArtifactStatusesResponse>, Status> {
+ Err(Status::unimplemented("not used"))
+ }
+ async fn interrupt(
+ &self,
+ _: Request<pb::InterruptRequest>,
+ ) -> Result<Response<pb::InterruptResponse>, Status> {
+ Err(Status::unimplemented("not used"))
+ }
+ async fn release_execute(
+ &self,
+ _: Request<pb::ReleaseExecuteRequest>,
+ ) -> Result<Response<pb::ReleaseExecuteResponse>, Status> {
+ Err(Status::unimplemented("not used"))
+ }
+ async fn release_session(
+ &self,
+ _: Request<pb::ReleaseSessionRequest>,
+ ) -> Result<Response<pb::ReleaseSessionResponse>, Status> {
+ Err(Status::unimplemented("not used"))
+ }
+ async fn fetch_error_details(
+ &self,
+ _: Request<pb::FetchErrorDetailsRequest>,
+ ) -> Result<Response<pb::FetchErrorDetailsResponse>, Status> {
+ Err(Status::unimplemented("not used"))
+ }
+ async fn clone_session(
+ &self,
+ _: Request<pb::CloneSessionRequest>,
+ ) -> Result<Response<pb::CloneSessionResponse>, Status> {
+ Err(Status::unimplemented("not used"))
+ }
+ async fn get_status(
+ &self,
+ _: Request<pb::GetStatusRequest>,
+ ) -> Result<Response<pb::GetStatusResponse>, Status> {
+ Err(Status::unimplemented("not used"))
+ }
+ async fn execute_plan(
+ &self,
+ _: Request<pb::ExecutePlanRequest>,
+ ) -> Result<Response<Self::ExecutePlanStream>, Status> {
+ Err(Status::unimplemented("not used"))
+ }
+ async fn reattach_execute(
+ &self,
+ _: Request<pb::ReattachExecuteRequest>,
+ ) -> Result<Response<Self::ReattachExecuteStream>, Status> {
+ Err(Status::unimplemented("not used"))
+ }
+ async fn add_artifacts(
+ &self,
+ _: Request<tonic::Streaming<pb::AddArtifactsRequest>>,
+ ) -> Result<Response<pb::AddArtifactsResponse>, Status> {
+ Err(Status::unimplemented("not used"))
+ }
+}
+
+struct TestRig {
+ channel: Channel,
+ backend: Arc<RecordingBackend>,
+ _be_shutdown: tokio::sync::oneshot::Sender<()>,
+ _gw_shutdown: tokio::sync::oneshot::Sender<()>,
+}
+
+async fn rig() -> TestRig {
+ // Backend
+ let backend = Arc::new(RecordingBackend::default());
+ let svc =
+ pb::spark_connect_service_server::SparkConnectServiceServer::from_arc(backend.clone());
+ let lis = TcpListener::bind("127.0.0.1:0").await.unwrap();
+ let be_addr = lis.local_addr().unwrap().to_string();
+ let (be_tx, be_rx) = tokio::sync::oneshot::channel();
+ tokio::spawn(async move {
+ Server::builder()
+ .add_service(svc)
+ .serve_with_incoming_shutdown(TcpListenerStream::new(lis), async {
+ let _ = be_rx.await;
+ })
+ .await
+ .ok();
+ });
+
+ // Gateway with static-token auth
+ let pool: Arc<dyn Pool> = Arc::new(StaticPool::new(vec![be_addr]).unwrap());
+ let store: Arc<dyn AffinityStore> = Arc::new(MemoryStore::new());
+ let router = Arc::new(Router::single_pool(pool, store));
+ let dialer = Dialer::new();
+
+ let auth = StaticTokenAuthenticator::new(vec![TokenEntry {
+ token: "alice-secret".into(),
+ user_id: "alice".into(),
+ tenant: None,
+ groups: vec![],
+ }])
+ .unwrap();
+ let interceptor = AuthInterceptor::new(Arc::new(auth));
+ let proxy = SparkConnectProxy::with_auth(router, dialer, interceptor);
+
+ let gw_lis = TcpListener::bind("127.0.0.1:0").await.unwrap();
+ let gw_addr = gw_lis.local_addr().unwrap();
+ let (gw_tx, gw_rx) = tokio::sync::oneshot::channel();
+ tokio::spawn(async move {
+ Server::builder()
+ .add_service(pb::spark_connect_service_server::SparkConnectServiceServer::new(proxy))
+ .serve_with_incoming_shutdown(TcpListenerStream::new(gw_lis), async {
+ let _ = gw_rx.await;
+ })
+ .await
+ .ok();
+ });
+
+ let endpoint = Endpoint::from_shared(format!("http://{}", gw_addr)).unwrap();
+ let channel = endpoint
+ .connect_timeout(Duration::from_secs(2))
+ .connect()
+ .await
+ .unwrap();
+ TestRig {
+ channel,
+ backend,
+ _be_shutdown: be_tx,
+ _gw_shutdown: gw_tx,
+ }
+}
+
+fn client(ch: Channel) -> pb::spark_connect_service_client::SparkConnectServiceClient<Channel> {
+ pb::spark_connect_service_client::SparkConnectServiceClient::new(ch)
+}
+
+#[tokio::test]
+async fn unauthenticated_request_rejected() {
+ let rig = rig().await;
+ let mut c = client(rig.channel.clone());
+ let err = c
+ .config(pb::ConfigRequest {
+ session_id: "s1".into(),
+ ..Default::default()
+ })
+ .await
+ .unwrap_err();
+ assert_eq!(err.code(), tonic::Code::Unauthenticated);
+ assert_eq!(
+ rig.backend.seen.load(Ordering::SeqCst),
+ 0,
+ "backend must not have been called"
+ );
+}
+
+#[tokio::test]
+async fn invalid_token_rejected() {
+ let rig = rig().await;
+ let mut c = client(rig.channel.clone());
+
+ let mut req = Request::new(pb::ConfigRequest {
+ session_id: "s1".into(),
+ ..Default::default()
+ });
+ req.metadata_mut().insert(
+ "authorization",
+ MetadataValue::try_from("Bearer wrong-secret").unwrap(),
+ );
+ let err = c.config(req).await.unwrap_err();
+ assert_eq!(err.code(), tonic::Code::Unauthenticated);
+ assert_eq!(rig.backend.seen.load(Ordering::SeqCst), 0);
+}
+
+#[tokio::test]
+async fn valid_token_authenticates_and_overwrites_user_id() {
+ let rig = rig().await;
+ let mut c = client(rig.channel.clone());
+
+ // Client *claims* to be "evil-impostor" — gateway must overwrite
+ // the user_id with the verified identity ("alice").
+ let mut req = Request::new(pb::ConfigRequest {
+ session_id: "s1".into(),
+ user_context: Some(pb::UserContext {
+ user_id: "evil-impostor".into(),
+ ..Default::default()
+ }),
+ ..Default::default()
+ });
+ req.metadata_mut().insert(
+ "authorization",
+ MetadataValue::try_from("Bearer alice-secret").unwrap(),
+ );
+ let resp = c.config(req).await.unwrap().into_inner();
+ assert_eq!(resp.session_id, "s1");
+ assert_eq!(rig.backend.seen.load(Ordering::SeqCst), 1);
+ let stamped = rig.backend.last_user_id.lock().clone();
+ assert_eq!(
+ stamped.as_deref(),
+ Some("alice"),
+ "gateway must overwrite client-supplied user_id with verified identity, \
+ not trust 'evil-impostor'"
+ );
+}
+
+#[tokio::test]
+async fn missing_user_context_gets_one_stamped() {
+ let rig = rig().await;
+ let mut c = client(rig.channel.clone());
+
+ // Client sends no UserContext at all.
+ let mut req = Request::new(pb::ConfigRequest {
+ session_id: "s1".into(),
+ user_context: None,
+ ..Default::default()
+ });
+ req.metadata_mut().insert(
+ "authorization",
+ MetadataValue::try_from("Bearer alice-secret").unwrap(),
+ );
+ c.config(req).await.unwrap();
+ let stamped = rig.backend.last_user_id.lock().clone();
+ assert_eq!(stamped.as_deref(), Some("alice"));
+}
diff --git a/crates/proxy/tests/clone_session.rs b/crates/proxy/tests/clone_session.rs
new file mode 100644
index 0000000..3909a5c
--- /dev/null
+++ b/crates/proxy/tests/clone_session.rs
@@ -0,0 +1,335 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Integration test for `CloneSession` binding semantics.
+//!
+//! A cloned session's state lives on the driver that executed the
+//! clone, so the gateway must bind the *new* session id to that same
+//! backend. Without the binding, follow-up RPCs on the cloned session
+//! go through pool selection and land on an arbitrary backend — with
+//! two backends in the pool, a coin flip away from a broken session.
+//!
+//! The rig runs two real backend servers so a missing binding has a
+//! real chance of being observable, and asserts the binding directly
+//! against the shared `MemoryStore` rather than sampling routing
+//! behaviour probabilistically.
+
+mod common;
+
+use std::sync::Arc;
+use std::time::Duration;
+
+use common::AuditCapture;
+use scg_audit::{AuditConfig, AuditLogger};
+use scg_auth::{AnonymousAuthenticator, AuthInterceptor};
+use scg_genproto::pb;
+use scg_observability::Metrics;
+use scg_pool_static::StaticPool;
+use scg_proxy::{Dialer, SparkConnectProxy};
+use scg_routing::{AffinityStore, Pool, Router, SessionKey, TenantRouter};
+use scg_store_memory::MemoryStore;
+use scg_tenant::{OnMissing, TenantResolver, TenantResolverConfig, TenantSource};
+use tokio::net::TcpListener;
+use tokio_stream::wrappers::TcpListenerStream;
+use tonic::transport::{Channel, Endpoint, Server};
+use tonic::{Request, Response, Status};
+
+/// Backend that supports `Config` (to establish a parent binding) and
+/// `CloneSession` (honouring a client-chosen id, generating one
+/// otherwise). Everything else is `Unimplemented`.
+#[derive(Default)]
+struct CloneBackend;
+
+#[tonic::async_trait]
+impl pb::spark_connect_service_server::SparkConnectService for CloneBackend {
+ type ExecutePlanStream = std::pin::Pin<
+ Box<dyn futures::Stream<Item = Result<pb::ExecutePlanResponse, Status>> + Send + 'static>,
+ >;
+ type ReattachExecuteStream = Self::ExecutePlanStream;
+
+ async fn config(
+ &self,
+ req: Request<pb::ConfigRequest>,
+ ) -> Result<Response<pb::ConfigResponse>, Status> {
+ let body = req.into_inner();
+ Ok(Response::new(pb::ConfigResponse {
+ session_id: body.session_id,
+ ..Default::default()
+ }))
+ }
+
+ async fn clone_session(
+ &self,
+ req: Request<pb::CloneSessionRequest>,
+ ) -> Result<Response<pb::CloneSessionResponse>, Status> {
+ let body = req.into_inner();
+ let new_id = body
+ .new_session_id
+ .filter(|s| !s.is_empty())
+ .unwrap_or_else(|| format!("{}-clone", body.session_id));
+ Ok(Response::new(pb::CloneSessionResponse {
+ session_id: body.session_id,
+ new_session_id: new_id,
+ ..Default::default()
+ }))
+ }
+
+ async fn release_session(
+ &self,
+ _: Request<pb::ReleaseSessionRequest>,
+ ) -> Result<Response<pb::ReleaseSessionResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn analyze_plan(
+ &self,
+ _: Request<pb::AnalyzePlanRequest>,
+ ) -> Result<Response<pb::AnalyzePlanResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn artifact_status(
+ &self,
+ _: Request<pb::ArtifactStatusesRequest>,
+ ) -> Result<Response<pb::ArtifactStatusesResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn interrupt(
+ &self,
+ _: Request<pb::InterruptRequest>,
+ ) -> Result<Response<pb::InterruptResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn release_execute(
+ &self,
+ _: Request<pb::ReleaseExecuteRequest>,
+ ) -> Result<Response<pb::ReleaseExecuteResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn fetch_error_details(
+ &self,
+ _: Request<pb::FetchErrorDetailsRequest>,
+ ) -> Result<Response<pb::FetchErrorDetailsResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn get_status(
+ &self,
+ _: Request<pb::GetStatusRequest>,
+ ) -> Result<Response<pb::GetStatusResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn execute_plan(
+ &self,
+ _: Request<pb::ExecutePlanRequest>,
+ ) -> Result<Response<Self::ExecutePlanStream>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn reattach_execute(
+ &self,
+ _: Request<pb::ReattachExecuteRequest>,
+ ) -> Result<Response<Self::ReattachExecuteStream>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn add_artifacts(
+ &self,
+ _: Request<tonic::Streaming<pb::AddArtifactsRequest>>,
+ ) -> Result<Response<pb::AddArtifactsResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+}
+
+async fn spawn_backend() -> (String, tokio::sync::oneshot::Sender<()>) {
+ let lis = TcpListener::bind("127.0.0.1:0").await.unwrap();
+ let addr = lis.local_addr().unwrap().to_string();
+ let (tx, rx) = tokio::sync::oneshot::channel();
+ tokio::spawn(async move {
+ Server::builder()
+ .add_service(
+ pb::spark_connect_service_server::SparkConnectServiceServer::new(CloneBackend),
+ )
+ .serve_with_incoming_shutdown(TcpListenerStream::new(lis), async {
+ let _ = rx.await;
+ })
+ .await
+ .ok();
+ });
+ (addr, tx)
+}
+
+struct Rig {
+ channel: Channel,
+ store: Arc<MemoryStore>,
+ _backend_shutdowns: Vec<tokio::sync::oneshot::Sender<()>>,
+ _gw_shutdown: tokio::sync::oneshot::Sender<()>,
+}
+
+async fn spawn_rig_two_backends() -> Rig {
+ let (addr_a, tx_a) = spawn_backend().await;
+ let (addr_b, tx_b) = spawn_backend().await;
+
+ let metrics = Metrics::new().unwrap();
+ let store = Arc::new(MemoryStore::new());
+ let store_dyn: Arc<dyn AffinityStore> = store.clone();
+ let pool: Arc<dyn Pool> = Arc::new(StaticPool::new(vec![addr_a, addr_b]).unwrap());
+ let router = Arc::new(Router::new(TenantRouter::single(pool), store_dyn));
+ let resolver = TenantResolver::new(TenantResolverConfig {
+ source: TenantSource::FromMetadata {
+ header: "x-tenant".into(),
+ },
+ on_missing: OnMissing::UseDefault,
+ default_name: "default".into(),
+ });
+ let audit = AuditLogger::new(AuditConfig {
+ enabled: true,
+ log_successful_rpcs: false,
+ });
+ let proxy = SparkConnectProxy::with_all(
+ router,
+ Dialer::new(),
+ AuthInterceptor::new(Arc::new(AnonymousAuthenticator)),
+ metrics,
+ resolver,
+ None,
+ audit,
+ );
+
+ let gw_lis = TcpListener::bind("127.0.0.1:0").await.unwrap();
+ let gw_addr = gw_lis.local_addr().unwrap();
+ let (gw_tx, gw_rx) = tokio::sync::oneshot::channel();
+ tokio::spawn(async move {
+ Server::builder()
+ .add_service(pb::spark_connect_service_server::SparkConnectServiceServer::new(proxy))
+ .serve_with_incoming_shutdown(TcpListenerStream::new(gw_lis), async {
+ let _ = gw_rx.await;
+ })
+ .await
+ .ok();
+ });
+
+ let channel = Endpoint::from_shared(format!("http://{}", gw_addr))
+ .unwrap()
+ .connect_timeout(Duration::from_secs(2))
+ .connect()
+ .await
+ .unwrap();
+
+ Rig {
+ channel,
+ store,
+ _backend_shutdowns: vec![tx_a, tx_b],
+ _gw_shutdown: gw_tx,
+ }
+}
+
+fn key(session_id: &str) -> SessionKey {
+ // AnonymousAuthenticator yields user_id="anonymous"; the resolver
+ // defaults the tenant to "default" when no x-tenant header is set.
+ SessionKey::with_tenant("default", "anonymous", session_id)
+}
+
+#[tokio::test]
+async fn cloned_session_is_bound_to_parents_backend() {
+ let cap = AuditCapture::lease();
+ let rig = spawn_rig_two_backends().await;
+ let mut c =
+ pb::spark_connect_service_client::SparkConnectServiceClient::new(rig.channel.clone());
+
+ // Establish the parent binding.
+ c.config(pb::ConfigRequest {
+ session_id: "parent-1".into(),
+ ..Default::default()
+ })
+ .await
+ .unwrap();
+ let parent_addr = rig
+ .store
+ .lookup_session(&key("parent-1"))
+ .await
+ .expect("parent session bound");
+
+ // Clone with a server-generated new id.
+ let resp = c
+ .clone_session(pb::CloneSessionRequest {
+ session_id: "parent-1".into(),
+ ..Default::default()
+ })
+ .await
+ .unwrap()
+ .into_inner();
+ let cloned = resp.new_session_id;
+ assert!(!cloned.is_empty());
+
+ // The cloned id must be bound to the same backend as the parent —
+ // not left to pool selection.
+ assert_eq!(
+ rig.store.lookup_session(&key(&cloned)).await.as_deref(),
+ Some(parent_addr.as_str()),
+ "cloned session must be bound to the parent's backend"
+ );
+
+ // And the fresh binding is audited: one session.create for the
+ // parent (from Config), one for the clone, with the same backend.
+ assert_eq!(cap.count_events("session.create"), 2);
+ let clone_evt = cap
+ .snapshot()
+ .into_iter()
+ .find(|e| {
+ e.fields.get("event").map(|s| s.as_str()) == Some("session.create")
+ && e.fields.get("session_id").map(|s| s.as_str()) == Some(cloned.as_str())
+ })
+ .expect("session.create for the cloned session");
+ assert_eq!(
+ clone_evt.fields.get("backend").map(|s| s.as_str()),
+ Some(parent_addr.as_str())
+ );
+}
+
+#[tokio::test]
+async fn client_chosen_clone_id_is_bound_to_parents_backend() {
+ let _cap = AuditCapture::lease();
+ let rig = spawn_rig_two_backends().await;
+ let mut c =
+ pb::spark_connect_service_client::SparkConnectServiceClient::new(rig.channel.clone());
+
+ c.config(pb::ConfigRequest {
+ session_id: "parent-2".into(),
+ ..Default::default()
+ })
+ .await
+ .unwrap();
+ let parent_addr = rig
+ .store
+ .lookup_session(&key("parent-2"))
+ .await
+ .expect("parent session bound");
+
+ // Clone with a client-chosen id, as the proto allows.
+ let resp = c
+ .clone_session(pb::CloneSessionRequest {
+ session_id: "parent-2".into(),
+ new_session_id: Some("my-clone-id".into()),
+ ..Default::default()
+ })
+ .await
+ .unwrap()
+ .into_inner();
+ assert_eq!(resp.new_session_id, "my-clone-id");
+
+ assert_eq!(
+ rig.store
+ .lookup_session(&key("my-clone-id"))
+ .await
+ .as_deref(),
+ Some(parent_addr.as_str()),
+ );
+}
diff --git a/crates/proxy/tests/common/mod.rs b/crates/proxy/tests/common/mod.rs
new file mode 100644
index 0000000..3a28035
--- /dev/null
+++ b/crates/proxy/tests/common/mod.rs
@@ -0,0 +1,161 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// `tests/common/mod.rs` is compiled into every integration-test binary
+// that does `mod common;`, but each binary only uses a subset of the
+// helpers — Rust's dead-code analysis runs per-binary and can't see
+// the cross-binary usage. Suppress here so a helper that's used in
+// only one test file doesn't fire a warning in the other.
+#![allow(dead_code)]
+
+//! Shared test helpers for the proxy integration tests.
+//!
+//! The big-ticket item here is [`AuditCapture`], a process-wide
+//! tracing subscriber that records every `target = "scg::audit"`
+//! event into an in-memory buffer. Tests use it to assert that the
+//! audit pipeline emits what it should.
+//!
+//! ## Why a global subscriber, not `set_default`?
+//!
+//! Earlier versions installed a per-test capture via
+//! `tracing::subscriber::set_default(...)`, which scopes the
+//! subscriber to the current thread. That works when the test's
+//! tokio runtime happens to keep every task on the current thread,
+//! but the gRPC server is `tokio::spawn`'d onto a worker pool and
+//! its audit events frequently fire on a different thread — which
+//! the per-thread subscriber doesn't see. The result was a flaky
+//! test that passed in isolation and failed under workspace
+//! parallelism.
+//!
+//! `tracing::subscriber::set_global_default` is process-wide, so it
+//! sees every thread, but it can only be called once. We wrap it in
+//! a `OnceLock` and serialize the tests that use it with a mutex —
+//! integration tests that need audit capture all hold the same
+//! `AuditCapture::lease`, so they run one-at-a-time within their
+//! file (and across files, since the global state is shared).
+
+use std::collections::HashMap;
+use std::sync::{Mutex, MutexGuard, OnceLock};
+
+use parking_lot::Mutex as PlMutex;
+use tracing_subscriber::layer::SubscriberExt;
+use tracing_subscriber::Layer;
+
+/// A single captured `target = "scg::audit"` event, flattened into a
+/// `name -> value` map. `String` values for both keys and values
+/// keep the assertion API simple — tests compare with string
+/// literals.
+#[derive(Debug, Clone)]
+pub struct CapturedEvent {
+ pub fields: HashMap<String, String>,
+}
+
+/// Shared event buffer. The tracing `Layer` writes into it; tests
+/// read from it via `AuditCapture::snapshot`.
+#[derive(Clone, Default)]
+struct EventSink {
+ events: std::sync::Arc<PlMutex<Vec<CapturedEvent>>>,
+}
+
+impl<S: tracing::Subscriber> Layer<S> for EventSink {
+ fn on_event(
+ &self,
+ event: &tracing::Event<'_>,
+ _ctx: tracing_subscriber::layer::Context<'_, S>,
+ ) {
+ if event.metadata().target() != "scg::audit" {
+ return;
+ }
+ let mut fields = HashMap::new();
+ struct Vis<'a>(&'a mut HashMap<String, String>);
+ impl<'a> tracing::field::Visit for Vis<'a> {
+ fn record_debug(&mut self, f: &tracing::field::Field, v: &dyn std::fmt::Debug) {
+ self.0.insert(f.name().into(), format!("{:?}", v));
+ }
+ fn record_str(&mut self, f: &tracing::field::Field, v: &str) {
+ self.0.insert(f.name().into(), v.into());
+ }
+ }
+ event.record(&mut Vis(&mut fields));
+ self.events.lock().push(CapturedEvent { fields });
+ }
+}
+
+/// Process-wide audit capture handle. Acquire one via
+/// [`AuditCapture::lease`]; the returned guard releases the
+/// global-subscriber lock on drop so other tests can proceed.
+///
+/// Tests do not need to (and cannot) install the subscriber
+/// themselves — the first `lease()` call installs it; later calls
+/// reuse it. Events from prior leases are cleared at the start of
+/// each new lease.
+pub struct AuditCapture<'a> {
+ sink: EventSink,
+ _guard: MutexGuard<'a, ()>,
+}
+
+impl<'a> AuditCapture<'a> {
+ pub fn lease() -> AuditCapture<'static> {
+ let (sink, mutex) = init();
+ let guard = mutex.lock().unwrap_or_else(|p| p.into_inner());
+ sink.events.lock().clear();
+ AuditCapture {
+ sink: sink.clone(),
+ _guard: guard,
+ }
+ }
+
+ pub fn snapshot(&self) -> Vec<CapturedEvent> {
+ self.sink.events.lock().clone()
+ }
+
+ pub fn count_events(&self, event_name: &str) -> usize {
+ self.snapshot()
+ .iter()
+ .filter(|e| e.fields.get("event").map(|s| s.as_str()) == Some(event_name))
+ .count()
+ }
+
+ pub fn find_event(&self, event_name: &str) -> Option<CapturedEvent> {
+ self.snapshot()
+ .into_iter()
+ .find(|e| e.fields.get("event").map(|s| s.as_str()) == Some(event_name))
+ }
+}
+
+/// Lazily install the global subscriber on first use. Returns the
+/// shared sink (so we can snapshot events) and the lease mutex (so
+/// callers serialize). `set_global_default` is called at most once
+/// per process.
+fn init() -> (&'static EventSink, &'static Mutex<()>) {
+ static STATE: OnceLock<(EventSink, Mutex<()>)> = OnceLock::new();
+ let (sink, mutex) = STATE.get_or_init(|| {
+ let sink = EventSink::default();
+ // Build the subscriber and stash it — every thread sees the
+ // same `Layer`, so events from tokio worker threads land in
+ // our shared buffer.
+ let subscriber = tracing_subscriber::registry().with(sink.clone());
+ // Setting the global default can fail if some other code in
+ // the test binary already set one. Tests that care about the
+ // audit stream all go through this helper, so the only way
+ // we lose the race is if a non-capture test in the same
+ // binary also installed a global subscriber. Currently none
+ // do, so we treat a `set_global_default` error as fatal.
+ tracing::subscriber::set_global_default(subscriber)
+ .expect("install audit-capture subscriber: another global subscriber already exists");
+ (sink, Mutex::new(()))
+ });
+ (sink, mutex)
+}
diff --git a/crates/proxy/tests/forwarding.rs b/crates/proxy/tests/forwarding.rs
new file mode 100644
index 0000000..db83813
--- /dev/null
+++ b/crates/proxy/tests/forwarding.rs
@@ -0,0 +1,387 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! In-process integration tests: gateway proxy in front of one or more fake
+//! Spark Connect backends. Validates forwarding semantics without needing a
+//! real Spark distribution.
+//!
+//! Four scenarios:
+//!
+//! 1. Unary RPC forwarding (`Config`).
+//! 2. Server-streaming forwarding emits all upstream messages (`ExecutePlan`).
+//! 3. Session affinity stickiness across many calls (no matter how many
+//! different sessions are interleaved).
+//! 4. Op-id reverse index — `ReattachExecute` with a *different* `session_id`
+//! still routes to the backend that handled the original `ExecutePlan`.
+
+use std::pin::Pin;
+use std::sync::Arc;
+use std::time::Duration;
+
+use futures::Stream;
+use scg_genproto::pb;
+use scg_pool_static::StaticPool;
+use scg_proxy::{Dialer, SparkConnectProxy};
+use scg_routing::{AffinityStore, Pool, Router};
+use scg_store_memory::MemoryStore;
+use tokio::net::TcpListener;
+use tokio_stream::wrappers::TcpListenerStream;
+use tonic::transport::{Channel, Endpoint, Server};
+use tonic::{Request, Response, Status};
+
+// ----- Fake backend -------------------------------------------------------
+
+/// Minimal SparkConnectService implementation. Each backend tags responses
+/// with its `id` so tests can prove which backend served a given RPC.
+#[derive(Default)]
+struct FakeBackend {
+ id: String,
+ config_count: std::sync::atomic::AtomicU64,
+ execute_count: std::sync::atomic::AtomicU64,
+}
+
+impl FakeBackend {
+ fn new(id: impl Into<String>) -> Self {
+ Self {
+ id: id.into(),
+ ..Default::default()
+ }
+ }
+}
+
+#[tonic::async_trait]
+impl pb::spark_connect_service_server::SparkConnectService for FakeBackend {
+ type ExecutePlanStream =
+ Pin<Box<dyn Stream<Item = Result<pb::ExecutePlanResponse, Status>> + Send + 'static>>;
+ type ReattachExecuteStream = Self::ExecutePlanStream;
+
+ async fn config(
+ &self,
+ req: Request<pb::ConfigRequest>,
+ ) -> Result<Response<pb::ConfigResponse>, Status> {
+ self.config_count
+ .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+ let body = req.into_inner();
+ Ok(Response::new(pb::ConfigResponse {
+ session_id: format!("{}@{}", body.session_id, self.id),
+ ..Default::default()
+ }))
+ }
+
+ async fn execute_plan(
+ &self,
+ req: Request<pb::ExecutePlanRequest>,
+ ) -> Result<Response<Self::ExecutePlanStream>, Status> {
+ self.execute_count
+ .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+ let body = req.into_inner();
+ let id = self.id.clone();
+ // Emit 3 messages so we can prove server-stream forwarding emits them all.
+ let stream = async_stream::stream! {
+ for _ in 0..3 {
+ yield Ok(pb::ExecutePlanResponse {
+ session_id: format!("{}@{}", body.session_id, id),
+ operation_id: body.operation_id.clone().unwrap_or_default(),
+ ..Default::default()
+ });
+ }
+ };
+ Ok(Response::new(Box::pin(stream)))
+ }
+
+ async fn reattach_execute(
+ &self,
+ req: Request<pb::ReattachExecuteRequest>,
+ ) -> Result<Response<Self::ReattachExecuteStream>, Status> {
+ self.execute_count
+ .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+ let body = req.into_inner();
+ let id = self.id.clone();
+ let stream = async_stream::stream! {
+ yield Ok(pb::ExecutePlanResponse {
+ session_id: format!("{}@{}", body.session_id, id),
+ operation_id: body.operation_id.clone(),
+ ..Default::default()
+ });
+ };
+ Ok(Response::new(Box::pin(stream)))
+ }
+
+ // The fake doesn't need to implement the rest for these tests. The
+ // default `Unimplemented` is fine and won't be triggered.
+ async fn analyze_plan(
+ &self,
+ _: Request<pb::AnalyzePlanRequest>,
+ ) -> Result<Response<pb::AnalyzePlanResponse>, Status> {
+ Err(Status::unimplemented("not used in tests"))
+ }
+ async fn artifact_status(
+ &self,
+ _: Request<pb::ArtifactStatusesRequest>,
+ ) -> Result<Response<pb::ArtifactStatusesResponse>, Status> {
+ Err(Status::unimplemented("not used in tests"))
+ }
+ async fn interrupt(
+ &self,
+ _: Request<pb::InterruptRequest>,
+ ) -> Result<Response<pb::InterruptResponse>, Status> {
+ Err(Status::unimplemented("not used in tests"))
+ }
+ async fn release_execute(
+ &self,
+ _: Request<pb::ReleaseExecuteRequest>,
+ ) -> Result<Response<pb::ReleaseExecuteResponse>, Status> {
+ Err(Status::unimplemented("not used in tests"))
+ }
+ async fn release_session(
+ &self,
+ _: Request<pb::ReleaseSessionRequest>,
+ ) -> Result<Response<pb::ReleaseSessionResponse>, Status> {
+ Err(Status::unimplemented("not used in tests"))
+ }
+ async fn fetch_error_details(
+ &self,
+ _: Request<pb::FetchErrorDetailsRequest>,
+ ) -> Result<Response<pb::FetchErrorDetailsResponse>, Status> {
+ Err(Status::unimplemented("not used in tests"))
+ }
+ async fn clone_session(
+ &self,
+ _: Request<pb::CloneSessionRequest>,
+ ) -> Result<Response<pb::CloneSessionResponse>, Status> {
+ Err(Status::unimplemented("not used in tests"))
+ }
+ async fn get_status(
+ &self,
+ _: Request<pb::GetStatusRequest>,
+ ) -> Result<Response<pb::GetStatusResponse>, Status> {
+ Err(Status::unimplemented("not used in tests"))
+ }
+ async fn add_artifacts(
+ &self,
+ _: Request<tonic::Streaming<pb::AddArtifactsRequest>>,
+ ) -> Result<Response<pb::AddArtifactsResponse>, Status> {
+ Err(Status::unimplemented("not used in tests"))
+ }
+}
+
+// ----- Test harness -------------------------------------------------------
+
+struct BackendHandle {
+ addr: String,
+ shutdown: Option<tokio::sync::oneshot::Sender<()>>,
+}
+
+impl Drop for BackendHandle {
+ fn drop(&mut self) {
+ if let Some(tx) = self.shutdown.take() {
+ let _ = tx.send(());
+ }
+ }
+}
+
+async fn start_backend(id: &str) -> BackendHandle {
+ let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+ let addr = listener.local_addr().unwrap().to_string();
+ let svc =
+ pb::spark_connect_service_server::SparkConnectServiceServer::new(FakeBackend::new(id));
+ let (tx, rx) = tokio::sync::oneshot::channel();
+ tokio::spawn(async move {
+ Server::builder()
+ .add_service(svc)
+ .serve_with_incoming_shutdown(TcpListenerStream::new(listener), async {
+ let _ = rx.await;
+ })
+ .await
+ .ok();
+ });
+ BackendHandle {
+ addr,
+ shutdown: Some(tx),
+ }
+}
+
+struct GatewayHandle {
+ channel: Channel,
+ shutdown: Option<tokio::sync::oneshot::Sender<()>>,
+}
+
+impl Drop for GatewayHandle {
+ fn drop(&mut self) {
+ if let Some(tx) = self.shutdown.take() {
+ let _ = tx.send(());
+ }
+ }
+}
+
+async fn start_gateway(backends: Vec<String>) -> GatewayHandle {
+ let pool: Arc<dyn Pool> = Arc::new(StaticPool::new(backends).unwrap());
+ let store: Arc<dyn AffinityStore> = Arc::new(MemoryStore::new());
+ let router = Arc::new(Router::single_pool(pool, store));
+ let dialer = Dialer::new();
+ let svc = SparkConnectProxy::new(router, dialer);
+ let server = pb::spark_connect_service_server::SparkConnectServiceServer::new(svc);
+
+ let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+ let addr = listener.local_addr().unwrap();
+ let (tx, rx) = tokio::sync::oneshot::channel();
+ tokio::spawn(async move {
+ Server::builder()
+ .add_service(server)
+ .serve_with_incoming_shutdown(TcpListenerStream::new(listener), async {
+ let _ = rx.await;
+ })
+ .await
+ .ok();
+ });
+
+ // Give the listener a moment to be ready and then dial in.
+ let endpoint = Endpoint::from_shared(format!("http://{}", addr)).unwrap();
+ let channel = endpoint
+ .connect_timeout(Duration::from_secs(2))
+ .connect()
+ .await
+ .unwrap();
+
+ GatewayHandle {
+ channel,
+ shutdown: Some(tx),
+ }
+}
+
+fn client(
+ h: &GatewayHandle,
+) -> pb::spark_connect_service_client::SparkConnectServiceClient<Channel> {
+ pb::spark_connect_service_client::SparkConnectServiceClient::new(h.channel.clone())
+}
+
+// ----- Tests --------------------------------------------------------------
+
+#[tokio::test]
+async fn unary_forward_returns_backend_tag() {
+ let be = start_backend("be1").await;
+ let gw = start_gateway(vec![be.addr.clone()]).await;
+
+ let resp = client(&gw)
+ .config(pb::ConfigRequest {
+ session_id: "sess-1".into(),
+ ..Default::default()
+ })
+ .await
+ .unwrap()
+ .into_inner();
+ assert_eq!(resp.session_id, "sess-1@be1");
+}
+
+#[tokio::test]
+async fn server_stream_emits_all_messages() {
+ let be = start_backend("be1").await;
+ let gw = start_gateway(vec![be.addr.clone()]).await;
+
+ let mut stream = client(&gw)
+ .execute_plan(pb::ExecutePlanRequest {
+ session_id: "sess-1".into(),
+ operation_id: Some("op-1".into()),
+ ..Default::default()
+ })
+ .await
+ .unwrap()
+ .into_inner();
+
+ let mut count = 0;
+ while let Some(_msg) = stream.message().await.unwrap() {
+ count += 1;
+ }
+ assert_eq!(count, 3);
+}
+
+#[tokio::test]
+async fn session_affinity_is_sticky_across_calls() {
+ let a = start_backend("A").await;
+ let b = start_backend("B").await;
+ let gw = start_gateway(vec![a.addr.clone(), b.addr.clone()]).await;
+
+ let mut c = client(&gw);
+ // First call binds sess-1 to whichever backend round-robin picks first.
+ let r1 = c
+ .config(pb::ConfigRequest {
+ session_id: "sess-1".into(),
+ ..Default::default()
+ })
+ .await
+ .unwrap()
+ .into_inner();
+ let first = r1.session_id; // "sess-1@A" or "sess-1@B"
+
+ for i in 0..5 {
+ let r = c
+ .config(pb::ConfigRequest {
+ session_id: "sess-1".into(),
+ ..Default::default()
+ })
+ .await
+ .unwrap()
+ .into_inner();
+ assert_eq!(r.session_id, first, "stickiness broken on call {}", i);
+ }
+}
+
+#[tokio::test]
+async fn reattach_routes_via_op_id_even_with_different_session_id() {
+ let a = start_backend("A").await;
+ let b = start_backend("B").await;
+ let gw = start_gateway(vec![a.addr.clone(), b.addr.clone()]).await;
+
+ let mut c = client(&gw);
+
+ // First, run an ExecutePlan with op-xyz on sess-1. The gateway records
+ // (op-xyz → backend) in the reverse index.
+ let mut s1 = c
+ .execute_plan(pb::ExecutePlanRequest {
+ session_id: "sess-1".into(),
+ operation_id: Some("op-xyz".into()),
+ ..Default::default()
+ })
+ .await
+ .unwrap()
+ .into_inner();
+ let mut original_suffix = String::new();
+ while let Some(msg) = s1.message().await.unwrap() {
+ original_suffix = msg.session_id["sess-1".len()..].to_string(); // "@A" or "@B"
+ }
+
+ // Now reattach with a *different* session id. The gateway must still
+ // route to the backend that handled the original ExecutePlan, because
+ // the op-id reverse index outranks session affinity.
+ let mut s2 = c
+ .reattach_execute(pb::ReattachExecuteRequest {
+ session_id: "different-session".into(),
+ operation_id: "op-xyz".into(),
+ ..Default::default()
+ })
+ .await
+ .unwrap()
+ .into_inner();
+ let mut reattach_suffix = String::new();
+ while let Some(msg) = s2.message().await.unwrap() {
+ reattach_suffix = msg.session_id["different-session".len()..].to_string();
+ }
+
+ assert_eq!(
+ original_suffix, reattach_suffix,
+ "reattach landed on the wrong backend (orig {:?}, reattach {:?})",
+ original_suffix, reattach_suffix
+ );
+}
diff --git a/crates/proxy/tests/multitenant_e2e.rs b/crates/proxy/tests/multitenant_e2e.rs
new file mode 100644
index 0000000..f164038
--- /dev/null
+++ b/crates/proxy/tests/multitenant_e2e.rs
@@ -0,0 +1,616 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Multi-tenant end-to-end integration test.
+//!
+//! One gateway, two tenants (`team-a`, `team-b`), three backends
+//! (`be-a`, `be-b`, `be-default`), and the full multi-tenant stack wired
+//! together: static-token auth (with per-token tenant claim) →
+//! tenant resolver (`FromClaim`) → tenant-routed pools → per-tenant
+//! rate limiter → session affinity → audit log.
+//!
+//! The point of this file is *cross-feature* isolation. Each
+//! individual feature already has its own integration test
+//! (`multitenant_pool.rs`, `ratelimit_integration.rs`,
+//! `audit_integration.rs`, `auth_integration.rs`); 3.9 verifies the
+//! features compose correctly when they're all on at once.
+//!
+//! Four axes of isolation are asserted:
+//!
+//! 1. **Pool + affinity isolation** — Same `session_id` from two
+//! tenants binds to two *different* backends, and the bindings
+//! don't leak across tenant boundaries even on repeat calls.
+//! 2. **Rate-limit isolation** — `team-a`'s tight quota does not
+//! affect `team-b`; rejections increment the metric only on the
+//! tenant whose bucket exhausted.
+//! 3. **Audit tenant labeling** — Every audit event carries the
+//! correct tenant label, taken from the resolved tenant (not the
+//! client-supplied `UserContext.user_id` or anything else
+//! forgeable by the caller).
+//! 4. **Auth-level tenant binding** — A token configured for
+//! `team-a` is rejected when its bearer is the only thing
+//! distinguishing the request from `team-b`. Combined with
+//! `OnMissing::Reject`, the gateway returns PermissionDenied or
+//! Unauthenticated before the request ever reaches a backend.
+
+mod common;
+
+use std::collections::HashMap;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::time::Duration;
+
+use common::{AuditCapture, CapturedEvent};
+use futures::Stream;
+use scg_audit::{AuditConfig, AuditLogger};
+use scg_auth::token::{StaticTokenAuthenticator, TokenEntry};
+use scg_auth::AuthInterceptor;
+use scg_genproto::pb;
+use scg_observability::Metrics;
+use scg_pool_static::StaticPool;
+use scg_proxy::{Dialer, SparkConnectProxy};
+use scg_ratelimit::{BucketRate, LimiterObserver, RateLimiter, RejectScope, TenantLimits};
+use scg_routing::{AffinityStore, PoolEntry, Router, TenantRouter, UnknownTenantPolicy};
+use scg_store_memory::MemoryStore;
+use scg_tenant::{OnMissing, TenantResolver, TenantResolverConfig, TenantSource};
+use tokio::net::TcpListener;
+use tokio_stream::wrappers::TcpListenerStream;
+use tonic::metadata::MetadataValue;
+use tonic::transport::{Channel, Endpoint, Server};
+use tonic::{Request, Response, Status};
+
+/// Backend that tags its `ConfigResponse.session_id` with its own id
+/// so the test driver can identify which backend handled an RPC.
+#[derive(Clone)]
+struct TaggedBackend {
+ id: String,
+}
+
+impl TaggedBackend {
+ fn new(id: impl Into<String>) -> Self {
+ Self { id: id.into() }
+ }
+}
+
+#[tonic::async_trait]
+impl pb::spark_connect_service_server::SparkConnectService for TaggedBackend {
+ type ExecutePlanStream =
+ Pin<Box<dyn Stream<Item = Result<pb::ExecutePlanResponse, Status>> + Send + 'static>>;
+ type ReattachExecuteStream = Self::ExecutePlanStream;
+
+ async fn config(
+ &self,
+ req: Request<pb::ConfigRequest>,
+ ) -> Result<Response<pb::ConfigResponse>, Status> {
+ let body = req.into_inner();
+ Ok(Response::new(pb::ConfigResponse {
+ session_id: format!("{}@{}", body.session_id, self.id),
+ ..Default::default()
+ }))
+ }
+ async fn analyze_plan(
+ &self,
+ _: Request<pb::AnalyzePlanRequest>,
+ ) -> Result<Response<pb::AnalyzePlanResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn artifact_status(
+ &self,
+ _: Request<pb::ArtifactStatusesRequest>,
+ ) -> Result<Response<pb::ArtifactStatusesResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn interrupt(
+ &self,
+ _: Request<pb::InterruptRequest>,
+ ) -> Result<Response<pb::InterruptResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn release_execute(
+ &self,
+ _: Request<pb::ReleaseExecuteRequest>,
+ ) -> Result<Response<pb::ReleaseExecuteResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn release_session(
+ &self,
+ _: Request<pb::ReleaseSessionRequest>,
+ ) -> Result<Response<pb::ReleaseSessionResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn fetch_error_details(
+ &self,
+ _: Request<pb::FetchErrorDetailsRequest>,
+ ) -> Result<Response<pb::FetchErrorDetailsResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn clone_session(
+ &self,
+ _: Request<pb::CloneSessionRequest>,
+ ) -> Result<Response<pb::CloneSessionResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn get_status(
+ &self,
+ _: Request<pb::GetStatusRequest>,
+ ) -> Result<Response<pb::GetStatusResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn execute_plan(
+ &self,
+ _: Request<pb::ExecutePlanRequest>,
+ ) -> Result<Response<Self::ExecutePlanStream>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn reattach_execute(
+ &self,
+ _: Request<pb::ReattachExecuteRequest>,
+ ) -> Result<Response<Self::ReattachExecuteStream>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn add_artifacts(
+ &self,
+ _: Request<tonic::Streaming<pb::AddArtifactsRequest>>,
+ ) -> Result<Response<pb::AddArtifactsResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+}
+
+async fn spawn_backend(id: &'static str) -> (String, tokio::sync::oneshot::Sender<()>) {
+ let svc =
+ pb::spark_connect_service_server::SparkConnectServiceServer::new(TaggedBackend::new(id));
+ let lis = TcpListener::bind("127.0.0.1:0").await.unwrap();
+ let addr = lis.local_addr().unwrap().to_string();
+ let (tx, rx) = tokio::sync::oneshot::channel();
+ tokio::spawn(async move {
+ Server::builder()
+ .add_service(svc)
+ .serve_with_incoming_shutdown(TcpListenerStream::new(lis), async {
+ let _ = rx.await;
+ })
+ .await
+ .ok();
+ });
+ (addr, tx)
+}
+
+/// Metrics-backed rate-limit observer, mirroring how `gateway/main.rs`
+/// wires the limiter to Prometheus counters in production.
+fn metrics_observer(metrics: &Metrics) -> Arc<dyn LimiterObserver> {
+ struct Obs(Metrics);
+ impl LimiterObserver for Obs {
+ fn on_reject(&self, tenant: &str, scope: RejectScope) {
+ self.0.record_rate_limit_reject(tenant, scope.as_str());
+ }
+ }
+ Arc::new(Obs(metrics.clone()))
+}
+
+fn rejected_count(metrics: &Metrics, tenant: &str, scope: &str) -> u64 {
+ for mf in metrics.registry().gather() {
+ if mf.name() != "scg_rate_limit_rejected_total" {
+ continue;
+ }
+ for m in mf.get_metric() {
+ let labels = m.get_label();
+ let t = labels
+ .iter()
+ .find(|l| l.name() == "tenant")
+ .map(|l| l.value());
+ let s = labels
+ .iter()
+ .find(|l| l.name() == "scope")
+ .map(|l| l.value());
+ if t == Some(tenant) && s == Some(scope) {
+ return m.get_counter().value() as u64;
+ }
+ }
+ }
+ 0
+}
+
+struct Rig {
+ channel: Channel,
+ metrics: Metrics,
+ _backends: Vec<tokio::sync::oneshot::Sender<()>>,
+ _gw_shutdown: tokio::sync::oneshot::Sender<()>,
+}
+
+/// Build the multi-tenant rig used by all tests in this file.
+///
+/// `unknown_policy` lets the caller swap between `UseDefault`
+/// (permissive single-tenant fall-through) and `Reject` (strict
+/// SaaS-style isolation). The auth-level test wants the strict
+/// policy.
+async fn spawn_rig(unknown_policy: UnknownTenantPolicy) -> Rig {
+ let (be_a, ka) = spawn_backend("be-a").await;
+ let (be_b, kb) = spawn_backend("be-b").await;
+ let (be_default, kd) = spawn_backend("be-default").await;
+
+ let mut tenants: HashMap<String, PoolEntry> = HashMap::new();
+ tenants.insert(
+ "team-a".into(),
+ PoolEntry::round_robin(Arc::new(StaticPool::new(vec![be_a]).unwrap())),
+ );
+ tenants.insert(
+ "team-b".into(),
+ PoolEntry::round_robin(Arc::new(StaticPool::new(vec![be_b]).unwrap())),
+ );
+ let default_pool = PoolEntry::round_robin(Arc::new(StaticPool::new(vec![be_default]).unwrap()));
+ let tr = TenantRouter::new(tenants, Some(default_pool), unknown_policy);
+
+ let metrics = Metrics::new().unwrap();
+ let store: Arc<dyn AffinityStore> = Arc::new(MemoryStore::new());
+ let router = Arc::new(Router::new(tr, store));
+ let dialer = Dialer::new();
+
+ // Static-token auth with one token per tenant. The token's
+ // `tenant` claim is the only thing distinguishing a request — the
+ // resolver below trusts it via `FromClaim`, so a stolen token
+ // gives access to its tenant and nothing else.
+ let token_auth = StaticTokenAuthenticator::new(vec![
+ TokenEntry {
+ token: "tok-a".into(),
+ user_id: "alice".into(),
+ tenant: Some("team-a".into()),
+ groups: vec![],
+ },
+ TokenEntry {
+ token: "tok-b".into(),
+ user_id: "bob".into(),
+ tenant: Some("team-b".into()),
+ groups: vec![],
+ },
+ ])
+ .unwrap();
+ let auth = AuthInterceptor::new(Arc::new(token_auth));
+
+ let resolver = TenantResolver::new(TenantResolverConfig {
+ source: TenantSource::FromClaim,
+ on_missing: match unknown_policy {
+ UnknownTenantPolicy::UseDefault => OnMissing::UseDefault,
+ UnknownTenantPolicy::Reject => OnMissing::Reject,
+ },
+ default_name: "default".into(),
+ });
+
+ // Tight quota on team-a so the rate-limit assertion can drive it
+ // to ResourceExhausted quickly without waiting for token refill.
+ // team-b uses the default (generous) quota.
+ let mut overrides = HashMap::new();
+ overrides.insert(
+ "team-a".to_string(),
+ TenantLimits {
+ tenant: BucketRate {
+ rpcs_per_second: 1.0,
+ burst: 3,
+ },
+ per_user: BucketRate::disabled(),
+ },
+ );
+ let limiter = RateLimiter::new(
+ TenantLimits::default(),
+ overrides,
+ metrics_observer(&metrics),
+ );
+
+ let audit = AuditLogger::new(AuditConfig {
+ enabled: true,
+ log_successful_rpcs: false,
+ });
+
+ let proxy = SparkConnectProxy::with_all(
+ router,
+ dialer,
+ auth,
+ metrics.clone(),
+ resolver,
+ Some(limiter),
+ audit,
+ );
+
+ let lis = TcpListener::bind("127.0.0.1:0").await.unwrap();
+ let addr = lis.local_addr().unwrap();
+ let (gw_tx, gw_rx) = tokio::sync::oneshot::channel();
+ tokio::spawn(async move {
+ Server::builder()
+ .add_service(pb::spark_connect_service_server::SparkConnectServiceServer::new(proxy))
+ .serve_with_incoming_shutdown(TcpListenerStream::new(lis), async {
+ let _ = gw_rx.await;
+ })
+ .await
+ .ok();
+ });
+
+ let endpoint = Endpoint::from_shared(format!("http://{}", addr)).unwrap();
+ let channel = endpoint
+ .connect_timeout(Duration::from_secs(2))
+ .connect()
+ .await
+ .unwrap();
+ Rig {
+ channel,
+ metrics,
+ _backends: vec![ka, kb, kd],
+ _gw_shutdown: gw_tx,
+ }
+}
+
+/// Issue a Config call using the given bearer token and return the
+/// backend id that handled it (extracted from `session_id@backend`).
+async fn config_as(ch: &Channel, session: &str, token: &str) -> Result<String, Status> {
+ let mut c = pb::spark_connect_service_client::SparkConnectServiceClient::new(ch.clone());
+ let mut req = Request::new(pb::ConfigRequest {
+ session_id: session.into(),
+ ..Default::default()
+ });
+ req.metadata_mut().insert(
+ "authorization",
+ MetadataValue::try_from(format!("Bearer {}", token)).unwrap(),
+ );
+ let resp = c.config(req).await?;
+ Ok(resp
+ .into_inner()
+ .session_id
+ .rsplit('@')
+ .next()
+ .unwrap()
+ .to_string())
+}
+
+/// Issue a Config call with no Authorization header — for the
+/// unauthenticated case.
+async fn config_no_auth(ch: &Channel, session: &str) -> Result<(), Status> {
+ let mut c = pb::spark_connect_service_client::SparkConnectServiceClient::new(ch.clone());
+ let req = Request::new(pb::ConfigRequest {
+ session_id: session.into(),
+ ..Default::default()
+ });
+ c.config(req).await.map(|_| ())
+}
+
+fn audit_events_for_tenant<'a>(
+ events: &'a [CapturedEvent],
+ tenant: &str,
+) -> Vec<&'a CapturedEvent> {
+ events
+ .iter()
+ .filter(|e| e.fields.get("tenant").map(|s| s.as_str()) == Some(tenant))
+ .collect()
+}
+
+// -------------------------------------------------------------------
+// Axis 1: pool routing + session-affinity isolation
+// -------------------------------------------------------------------
+
+#[tokio::test]
+async fn same_session_id_across_tenants_lands_on_different_backends() {
+ // Hold the capture lease for the test's lifetime so audit events
+ // emitted by this test's rig don't pollute another test's
+ // snapshot. The lease serializes tests within this file even
+ // though we don't read events here.
+ let _cap = AuditCapture::lease();
+ let rig = spawn_rig(UnknownTenantPolicy::UseDefault).await;
+
+ // Both clients pick the *same* session_id; the only thing
+ // distinguishing them is the bearer token (→ tenant claim).
+ let be_for_a = config_as(&rig.channel, "shared-sid", "tok-a")
+ .await
+ .unwrap();
+ let be_for_b = config_as(&rig.channel, "shared-sid", "tok-b")
+ .await
+ .unwrap();
+ assert_eq!(be_for_a, "be-a");
+ assert_eq!(be_for_b, "be-b");
+
+ // Repeat both — affinity must hold inside each tenant, and must
+ // not bleed across tenants.
+ let again_a = config_as(&rig.channel, "shared-sid", "tok-a")
+ .await
+ .unwrap();
+ let again_b = config_as(&rig.channel, "shared-sid", "tok-b")
+ .await
+ .unwrap();
+ assert_eq!(again_a, "be-a", "team-a affinity drifted to {}", again_a);
+ assert_eq!(again_b, "be-b", "team-b affinity drifted to {}", again_b);
+}
+
+// -------------------------------------------------------------------
+// Axis 2: per-tenant rate-limit isolation
+// -------------------------------------------------------------------
+
+#[tokio::test]
+async fn one_tenant_exhausting_quota_does_not_affect_the_other() {
+ let _cap = AuditCapture::lease();
+ let rig = spawn_rig(UnknownTenantPolicy::UseDefault).await;
+
+ // team-a is configured with burst=3. The first 3 succeed; the
+ // 4th must be ResourceExhausted.
+ for i in 0..3 {
+ config_as(&rig.channel, &format!("s-{}", i), "tok-a")
+ .await
+ .expect("team-a within burst");
+ }
+ let err = config_as(&rig.channel, "s-overflow", "tok-a")
+ .await
+ .unwrap_err();
+ assert_eq!(err.code(), tonic::Code::ResourceExhausted);
+
+ // Metric increments on team-a, scope=tenant. team-b has zero.
+ assert_eq!(rejected_count(&rig.metrics, "team-a", "tenant"), 1);
+ assert_eq!(rejected_count(&rig.metrics, "team-b", "tenant"), 0);
+
+ // team-b uses the (generous) default quota; with team-a's bucket
+ // wedged, team-b should still sail through.
+ for i in 0..5 {
+ config_as(&rig.channel, &format!("s-b-{}", i), "tok-b")
+ .await
+ .expect("team-b unaffected by team-a's exhaustion");
+ }
+ assert_eq!(rejected_count(&rig.metrics, "team-b", "tenant"), 0);
+}
+
+// -------------------------------------------------------------------
+// Axis 3: audit events carry the resolved tenant
+// -------------------------------------------------------------------
+
+#[tokio::test]
+async fn audit_events_label_the_correct_tenant() {
+ let cap = AuditCapture::lease();
+ let rig = spawn_rig(UnknownTenantPolicy::UseDefault).await;
+
+ config_as(&rig.channel, "sess-a", "tok-a").await.unwrap();
+ config_as(&rig.channel, "sess-b", "tok-b").await.unwrap();
+
+ let events = cap.snapshot();
+
+ // Each tenant got exactly one session.create labeled with itself.
+ let a_events: Vec<_> = audit_events_for_tenant(&events, "team-a")
+ .into_iter()
+ .filter(|e| e.fields.get("event").map(|s| s.as_str()) == Some("session.create"))
+ .collect();
+ let b_events: Vec<_> = audit_events_for_tenant(&events, "team-b")
+ .into_iter()
+ .filter(|e| e.fields.get("event").map(|s| s.as_str()) == Some("session.create"))
+ .collect();
+ assert_eq!(a_events.len(), 1, "team-a session.create count");
+ assert_eq!(b_events.len(), 1, "team-b session.create count");
+
+ // Cross-check the user_id on each event was the verified
+ // identity, not anything else — the test relies on this when
+ // claiming the resolver pulled `tenant` from the same token.
+ assert_eq!(
+ a_events[0].fields.get("user_id").map(|s| s.as_str()),
+ Some("alice"),
+ );
+ assert_eq!(
+ b_events[0].fields.get("user_id").map(|s| s.as_str()),
+ Some("bob"),
+ );
+
+ // session_id labels match what each client supplied — the only
+ // way they would collide here is a routing-key bug.
+ assert_eq!(
+ a_events[0].fields.get("session_id").map(|s| s.as_str()),
+ Some("sess-a"),
+ );
+ assert_eq!(
+ b_events[0].fields.get("session_id").map(|s| s.as_str()),
+ Some("sess-b"),
+ );
+}
+
+// -------------------------------------------------------------------
+// Axis 4: auth-level tenant binding
+// -------------------------------------------------------------------
+
+#[tokio::test]
+async fn unauthenticated_request_is_rejected_before_routing() {
+ let _cap = AuditCapture::lease();
+ let rig = spawn_rig(UnknownTenantPolicy::UseDefault).await;
+ let err = config_no_auth(&rig.channel, "anything").await.unwrap_err();
+ assert_eq!(err.code(), tonic::Code::Unauthenticated);
+}
+
+#[tokio::test]
+async fn unknown_tenant_token_rejected_under_reject_policy() {
+ let _cap = AuditCapture::lease();
+ // Add a third token whose tenant claim is missing — the
+ // resolver's Reject policy must turn that into a hard rejection
+ // before the request reaches a backend.
+ //
+ // We can't extend the rig's auth list without rebuilding, so
+ // construct a one-off authenticator inline. Everything else
+ // (pool topology, resolver, limiter) mirrors `spawn_rig`.
+ let (be_a, _ka) = spawn_backend("be-a").await;
+ let (be_b, _kb) = spawn_backend("be-b").await;
+
+ let mut tenants: HashMap<String, PoolEntry> = HashMap::new();
+ tenants.insert(
+ "team-a".into(),
+ PoolEntry::round_robin(Arc::new(StaticPool::new(vec![be_a]).unwrap())),
+ );
+ tenants.insert(
+ "team-b".into(),
+ PoolEntry::round_robin(Arc::new(StaticPool::new(vec![be_b]).unwrap())),
+ );
+ // No default pool + Reject policy + tenantless token = the gateway
+ // must surface Unauthenticated (from the resolver) and never
+ // attempt to forward.
+ let tr = TenantRouter::new(tenants, None, UnknownTenantPolicy::Reject);
+
+ let metrics = Metrics::new().unwrap();
+ let store: Arc<dyn AffinityStore> = Arc::new(MemoryStore::new());
+ let router = Arc::new(Router::new(tr, store));
+ let dialer = Dialer::new();
+
+ let token_auth = StaticTokenAuthenticator::new(vec![TokenEntry {
+ token: "rogue".into(),
+ user_id: "mallory".into(),
+ tenant: None, // <-- no tenant claim
+ groups: vec![],
+ }])
+ .unwrap();
+ let auth = AuthInterceptor::new(Arc::new(token_auth));
+ let resolver = TenantResolver::new(TenantResolverConfig {
+ source: TenantSource::FromClaim,
+ on_missing: OnMissing::Reject,
+ default_name: "default".into(),
+ });
+ let proxy = SparkConnectProxy::with_components(router, dialer, auth, metrics, resolver);
+
+ let lis = TcpListener::bind("127.0.0.1:0").await.unwrap();
+ let addr = lis.local_addr().unwrap();
+ let (_gw_tx, gw_rx) = tokio::sync::oneshot::channel::<()>();
+ tokio::spawn(async move {
+ Server::builder()
+ .add_service(pb::spark_connect_service_server::SparkConnectServiceServer::new(proxy))
+ .serve_with_incoming_shutdown(TcpListenerStream::new(lis), async {
+ let _ = gw_rx.await;
+ })
+ .await
+ .ok();
+ });
+
+ let endpoint = Endpoint::from_shared(format!("http://{}", addr)).unwrap();
+ let ch = endpoint
+ .connect_timeout(Duration::from_secs(2))
+ .connect()
+ .await
+ .unwrap();
+
+ let mut c = pb::spark_connect_service_client::SparkConnectServiceClient::new(ch);
+ let mut req = Request::new(pb::ConfigRequest {
+ session_id: "doesnt-matter".into(),
+ ..Default::default()
+ });
+ req.metadata_mut().insert(
+ "authorization",
+ MetadataValue::try_from("Bearer rogue").unwrap(),
+ );
+ let err = c.config(req).await.unwrap_err();
+ // Missing-tenant + on_missing=Reject → Unauthenticated from the
+ // resolver. PermissionDenied is also accepted in case a future
+ // policy refinement reclassifies it; either proves "the request
+ // was rejected before it reached a backend".
+ assert!(
+ matches!(
+ err.code(),
+ tonic::Code::Unauthenticated | tonic::Code::PermissionDenied,
+ ),
+ "expected hard rejection, got {:?}",
+ err.code(),
+ );
+}
diff --git a/crates/proxy/tests/multitenant_pool.rs b/crates/proxy/tests/multitenant_pool.rs
new file mode 100644
index 0000000..9561d05
--- /dev/null
+++ b/crates/proxy/tests/multitenant_pool.rs
@@ -0,0 +1,305 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Integration test for per-tenant pool routing.
+//!
+//! Spawns three fake backends (team-a-be, team-b-be, default-be),
+//! wires them into a single gateway with a per-tenant pool map, and
+//! drives RPCs through each tenant. Asserts:
+//!
+//! * `tenant=team-a` lands on `team-a-be` even though `team-b-be`
+//! and `default-be` are reachable.
+//! * `tenant=team-b` lands on `team-b-be`.
+//! * `tenant=unknown` falls back to `default-be` under
+//! `UseDefault` policy.
+//! * Under `Reject` policy, the same `unknown` tenant gets
+//! `PermissionDenied` and the backends never see the call.
+
+use std::collections::HashMap;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::time::Duration;
+
+use futures::Stream;
+use scg_auth::{AnonymousAuthenticator, AuthInterceptor};
+use scg_genproto::pb;
+use scg_observability::Metrics;
+use scg_pool_static::StaticPool;
+use scg_proxy::{Dialer, SparkConnectProxy};
+use scg_routing::{AffinityStore, PoolEntry, Router, TenantRouter, UnknownTenantPolicy};
+use scg_store_memory::MemoryStore;
+use scg_tenant::{OnMissing, TenantResolver, TenantResolverConfig, TenantSource};
+use tokio::net::TcpListener;
+use tokio_stream::wrappers::TcpListenerStream;
+use tonic::metadata::MetadataValue;
+use tonic::transport::{Channel, Endpoint, Server};
+use tonic::{Request, Response, Status};
+
+/// Backend that tags every response with its own id so the driver
+/// can verify which backend handled a given RPC.
+#[derive(Clone, Default)]
+struct TaggedBackend {
+ id: String,
+}
+
+impl TaggedBackend {
+ fn new(id: impl Into<String>) -> Self {
+ Self { id: id.into() }
+ }
+}
+
+#[tonic::async_trait]
+impl pb::spark_connect_service_server::SparkConnectService for TaggedBackend {
+ type ExecutePlanStream =
+ Pin<Box<dyn Stream<Item = Result<pb::ExecutePlanResponse, Status>> + Send + 'static>>;
+ type ReattachExecuteStream = Self::ExecutePlanStream;
+
+ async fn config(
+ &self,
+ req: Request<pb::ConfigRequest>,
+ ) -> Result<Response<pb::ConfigResponse>, Status> {
+ let body = req.into_inner();
+ Ok(Response::new(pb::ConfigResponse {
+ session_id: format!("{}@{}", body.session_id, self.id),
+ ..Default::default()
+ }))
+ }
+
+ async fn analyze_plan(
+ &self,
+ _: Request<pb::AnalyzePlanRequest>,
+ ) -> Result<Response<pb::AnalyzePlanResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn artifact_status(
+ &self,
+ _: Request<pb::ArtifactStatusesRequest>,
+ ) -> Result<Response<pb::ArtifactStatusesResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn interrupt(
+ &self,
+ _: Request<pb::InterruptRequest>,
+ ) -> Result<Response<pb::InterruptResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn release_execute(
+ &self,
+ _: Request<pb::ReleaseExecuteRequest>,
+ ) -> Result<Response<pb::ReleaseExecuteResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn release_session(
+ &self,
+ _: Request<pb::ReleaseSessionRequest>,
+ ) -> Result<Response<pb::ReleaseSessionResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn fetch_error_details(
+ &self,
+ _: Request<pb::FetchErrorDetailsRequest>,
+ ) -> Result<Response<pb::FetchErrorDetailsResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn clone_session(
+ &self,
+ _: Request<pb::CloneSessionRequest>,
+ ) -> Result<Response<pb::CloneSessionResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn get_status(
+ &self,
+ _: Request<pb::GetStatusRequest>,
+ ) -> Result<Response<pb::GetStatusResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn execute_plan(
+ &self,
+ _: Request<pb::ExecutePlanRequest>,
+ ) -> Result<Response<Self::ExecutePlanStream>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn reattach_execute(
+ &self,
+ _: Request<pb::ReattachExecuteRequest>,
+ ) -> Result<Response<Self::ReattachExecuteStream>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn add_artifacts(
+ &self,
+ _: Request<tonic::Streaming<pb::AddArtifactsRequest>>,
+ ) -> Result<Response<pb::AddArtifactsResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+}
+
+async fn spawn_backend(id: &'static str) -> (String, tokio::sync::oneshot::Sender<()>) {
+ let svc =
+ pb::spark_connect_service_server::SparkConnectServiceServer::new(TaggedBackend::new(id));
+ let lis = TcpListener::bind("127.0.0.1:0").await.unwrap();
+ let addr = lis.local_addr().unwrap().to_string();
+ let (tx, rx) = tokio::sync::oneshot::channel();
+ tokio::spawn(async move {
+ Server::builder()
+ .add_service(svc)
+ .serve_with_incoming_shutdown(TcpListenerStream::new(lis), async {
+ let _ = rx.await;
+ })
+ .await
+ .ok();
+ });
+ (addr, tx)
+}
+
+async fn spawn_gateway(tenant_router: TenantRouter) -> (Channel, tokio::sync::oneshot::Sender<()>) {
+ let metrics = Metrics::new().unwrap();
+ let store: Arc<dyn AffinityStore> = Arc::new(MemoryStore::new());
+ let router = Arc::new(Router::new(tenant_router, store));
+ let dialer = Dialer::new();
+ // Use the metadata-based tenant resolver so the test client can
+ // declare which tenant to route through without standing up real
+ // JWT plumbing.
+ let resolver = TenantResolver::new(TenantResolverConfig {
+ source: TenantSource::FromMetadata {
+ header: "x-tenant".into(),
+ },
+ on_missing: OnMissing::UseDefault,
+ default_name: "default".into(),
+ });
+ let proxy = SparkConnectProxy::with_components(
+ router,
+ dialer,
+ AuthInterceptor::new(Arc::new(AnonymousAuthenticator)),
+ metrics,
+ resolver,
+ );
+
+ let lis = TcpListener::bind("127.0.0.1:0").await.unwrap();
+ let addr = lis.local_addr().unwrap();
+ let (tx, rx) = tokio::sync::oneshot::channel();
+ tokio::spawn(async move {
+ Server::builder()
+ .add_service(pb::spark_connect_service_server::SparkConnectServiceServer::new(proxy))
+ .serve_with_incoming_shutdown(TcpListenerStream::new(lis), async {
+ let _ = rx.await;
+ })
+ .await
+ .ok();
+ });
+
+ let endpoint = Endpoint::from_shared(format!("http://{}", addr)).unwrap();
+ let ch = endpoint
+ .connect_timeout(Duration::from_secs(2))
+ .connect()
+ .await
+ .unwrap();
+ (ch, tx)
+}
+
+/// Issue a Config call with a chosen tenant and return the backend
+/// id that processed it (extracted from `session_id@backend_id`).
+async fn config_with_tenant(ch: &Channel, session: &str, tenant: &str) -> String {
+ let mut c = pb::spark_connect_service_client::SparkConnectServiceClient::new(ch.clone());
+ let mut req = Request::new(pb::ConfigRequest {
+ session_id: session.into(),
+ ..Default::default()
+ });
+ req.metadata_mut()
+ .insert("x-tenant", MetadataValue::try_from(tenant).unwrap());
+ let resp = c.config(req).await.expect("Config RPC succeeds");
+ resp.into_inner()
+ .session_id
+ .rsplit('@')
+ .next()
+ .unwrap()
+ .to_string()
+}
+
+#[tokio::test]
+async fn per_tenant_pools_route_to_their_own_backends() {
+ let (be_a, _ka) = spawn_backend("be-a").await;
+ let (be_b, _kb) = spawn_backend("be-b").await;
+ let (be_default, _kd) = spawn_backend("be-default").await;
+
+ let mut tenants: HashMap<String, PoolEntry> = HashMap::new();
+ tenants.insert(
+ "team-a".into(),
+ PoolEntry::round_robin(Arc::new(StaticPool::new(vec![be_a]).unwrap())),
+ );
+ tenants.insert(
+ "team-b".into(),
+ PoolEntry::round_robin(Arc::new(StaticPool::new(vec![be_b]).unwrap())),
+ );
+ let default_pool = PoolEntry::round_robin(Arc::new(StaticPool::new(vec![be_default]).unwrap()));
+ let tr = TenantRouter::new(tenants, Some(default_pool), UnknownTenantPolicy::UseDefault);
+
+ let (ch, _gw) = spawn_gateway(tr).await;
+
+ // Same session_id, different tenants → different backends.
+ let a = config_with_tenant(&ch, "s1", "team-a").await;
+ let b = config_with_tenant(&ch, "s1", "team-b").await;
+ assert_eq!(a, "be-a");
+ assert_eq!(b, "be-b");
+
+ // Unknown tenant under UseDefault → default backend.
+ let d = config_with_tenant(&ch, "s1", "stranger-tenant").await;
+ assert_eq!(d, "be-default");
+}
+
+#[tokio::test]
+async fn reject_policy_returns_permission_denied_for_unknown_tenant() {
+ let (be_a, _ka) = spawn_backend("be-a").await;
+
+ let mut tenants: HashMap<String, PoolEntry> = HashMap::new();
+ tenants.insert(
+ "team-a".into(),
+ PoolEntry::round_robin(Arc::new(StaticPool::new(vec![be_a]).unwrap())),
+ );
+ // Reject policy, no default pool. An unknown tenant must fail
+ // loudly with PermissionDenied.
+ let tr = TenantRouter::new(tenants, None, UnknownTenantPolicy::Reject);
+
+ let (ch, _gw) = spawn_gateway(tr).await;
+ let mut c = pb::spark_connect_service_client::SparkConnectServiceClient::new(ch);
+ let mut req = Request::new(pb::ConfigRequest {
+ session_id: "s1".into(),
+ ..Default::default()
+ });
+ req.metadata_mut()
+ .insert("x-tenant", MetadataValue::try_from("stranger").unwrap());
+ let err = c.config(req).await.unwrap_err();
+ assert_eq!(err.code(), tonic::Code::PermissionDenied);
+}
+
+#[tokio::test]
+async fn back_compat_single_pool_via_default_only() {
+ // Empty `overrides` map: every tenant routes through the
+ // configured default. This is the single-tenant baseline.
+ let (be_default, _kd) = spawn_backend("be-default").await;
+ let default_pool = PoolEntry::round_robin(Arc::new(StaticPool::new(vec![be_default]).unwrap()));
+ let tr = TenantRouter::new(
+ HashMap::new(),
+ Some(default_pool),
+ UnknownTenantPolicy::UseDefault,
+ );
+
+ let (ch, _gw) = spawn_gateway(tr).await;
+ let a = config_with_tenant(&ch, "s1", "team-a").await;
+ let b = config_with_tenant(&ch, "s1", "team-b").await;
+ let d = config_with_tenant(&ch, "s1", "default").await;
+ assert_eq!(a, "be-default");
+ assert_eq!(b, "be-default");
+ assert_eq!(d, "be-default");
+}
diff --git a/crates/proxy/tests/observability_integration.rs b/crates/proxy/tests/observability_integration.rs
new file mode 100644
index 0000000..9888047
--- /dev/null
+++ b/crates/proxy/tests/observability_integration.rs
@@ -0,0 +1,321 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Integration test: a forwarded RPC bumps `scg_rpcs_total`, an
+//! unauthenticated request bumps `scg_auth_failures_total`, and a
+//! Prometheus scrape against the admin server returns both as text.
+
+use std::pin::Pin;
+use std::sync::Arc;
+use std::time::Duration;
+
+use futures::Stream;
+use scg_auth::{
+ token::{StaticTokenAuthenticator, TokenEntry},
+ AuthInterceptor,
+};
+use scg_genproto::pb;
+use scg_observability::{serve_admin, AdminConfig, Metrics, ReadinessProbe};
+use scg_pool_static::StaticPool;
+use scg_proxy::{Dialer, SparkConnectProxy};
+use scg_routing::{AffinityStore, Pool, Router};
+use scg_store_memory::MemoryStore;
+use tokio::net::TcpListener;
+use tokio_stream::wrappers::TcpListenerStream;
+use tonic::metadata::MetadataValue;
+use tonic::transport::{Channel, Endpoint, Server};
+use tonic::{Request, Response, Status};
+
+/// Backend that just echoes session_id for simplicity.
+#[derive(Default)]
+struct EchoBackend;
+
+#[tonic::async_trait]
+impl pb::spark_connect_service_server::SparkConnectService for EchoBackend {
+ type ExecutePlanStream =
+ Pin<Box<dyn Stream<Item = Result<pb::ExecutePlanResponse, Status>> + Send + 'static>>;
+ type ReattachExecuteStream = Self::ExecutePlanStream;
+
+ async fn config(
+ &self,
+ req: Request<pb::ConfigRequest>,
+ ) -> Result<Response<pb::ConfigResponse>, Status> {
+ let body = req.into_inner();
+ Ok(Response::new(pb::ConfigResponse {
+ session_id: body.session_id,
+ ..Default::default()
+ }))
+ }
+ async fn analyze_plan(
+ &self,
+ _: Request<pb::AnalyzePlanRequest>,
+ ) -> Result<Response<pb::AnalyzePlanResponse>, Status> {
+ Err(Status::unimplemented("not used"))
+ }
+ async fn artifact_status(
+ &self,
+ _: Request<pb::ArtifactStatusesRequest>,
+ ) -> Result<Response<pb::ArtifactStatusesResponse>, Status> {
+ Err(Status::unimplemented("not used"))
+ }
+ async fn interrupt(
+ &self,
+ _: Request<pb::InterruptRequest>,
+ ) -> Result<Response<pb::InterruptResponse>, Status> {
+ Err(Status::unimplemented("not used"))
+ }
+ async fn release_execute(
+ &self,
+ _: Request<pb::ReleaseExecuteRequest>,
+ ) -> Result<Response<pb::ReleaseExecuteResponse>, Status> {
+ Err(Status::unimplemented("not used"))
+ }
+ async fn release_session(
+ &self,
+ _: Request<pb::ReleaseSessionRequest>,
+ ) -> Result<Response<pb::ReleaseSessionResponse>, Status> {
+ Err(Status::unimplemented("not used"))
+ }
+ async fn fetch_error_details(
+ &self,
+ _: Request<pb::FetchErrorDetailsRequest>,
+ ) -> Result<Response<pb::FetchErrorDetailsResponse>, Status> {
+ Err(Status::unimplemented("not used"))
+ }
+ async fn clone_session(
+ &self,
+ _: Request<pb::CloneSessionRequest>,
+ ) -> Result<Response<pb::CloneSessionResponse>, Status> {
+ Err(Status::unimplemented("not used"))
+ }
+ async fn get_status(
+ &self,
+ _: Request<pb::GetStatusRequest>,
+ ) -> Result<Response<pb::GetStatusResponse>, Status> {
+ Err(Status::unimplemented("not used"))
+ }
+ async fn execute_plan(
+ &self,
+ _: Request<pb::ExecutePlanRequest>,
+ ) -> Result<Response<Self::ExecutePlanStream>, Status> {
+ Err(Status::unimplemented("not used"))
+ }
+ async fn reattach_execute(
+ &self,
+ _: Request<pb::ReattachExecuteRequest>,
+ ) -> Result<Response<Self::ReattachExecuteStream>, Status> {
+ Err(Status::unimplemented("not used"))
+ }
+ async fn add_artifacts(
+ &self,
+ _: Request<tonic::Streaming<pb::AddArtifactsRequest>>,
+ ) -> Result<Response<pb::AddArtifactsResponse>, Status> {
+ Err(Status::unimplemented("not used"))
+ }
+}
+
+struct Rig {
+ grpc: Channel,
+ admin_url: String,
+ _be_shutdown: tokio::sync::oneshot::Sender<()>,
+ _gw_shutdown: tokio::sync::oneshot::Sender<()>,
+ _admin_shutdown: tokio::sync::oneshot::Sender<()>,
+}
+
+async fn rig_with_metrics() -> (Rig, Metrics) {
+ // Backend
+ let svc = pb::spark_connect_service_server::SparkConnectServiceServer::new(EchoBackend);
+ let lis = TcpListener::bind("127.0.0.1:0").await.unwrap();
+ let be_addr = lis.local_addr().unwrap().to_string();
+ let (be_tx, be_rx) = tokio::sync::oneshot::channel();
+ tokio::spawn(async move {
+ Server::builder()
+ .add_service(svc)
+ .serve_with_incoming_shutdown(TcpListenerStream::new(lis), async {
+ let _ = be_rx.await;
+ })
+ .await
+ .ok();
+ });
+
+ // Gateway with static-token auth + shared Metrics
+ let metrics = Metrics::new().unwrap();
+ let pool: Arc<dyn Pool> = Arc::new(StaticPool::new(vec![be_addr]).unwrap());
+ let store: Arc<dyn AffinityStore> = Arc::new(MemoryStore::new());
+ let router = Arc::new(Router::single_pool(pool, store));
+ let dialer = Dialer::new();
+ let auth = StaticTokenAuthenticator::new(vec![TokenEntry {
+ token: "alice-secret".into(),
+ user_id: "alice".into(),
+ tenant: None,
+ groups: vec![],
+ }])
+ .unwrap();
+ let proxy = SparkConnectProxy::with_auth_and_metrics(
+ router,
+ dialer,
+ AuthInterceptor::new(Arc::new(auth)),
+ metrics.clone(),
+ );
+
+ let gw_lis = TcpListener::bind("127.0.0.1:0").await.unwrap();
+ let gw_addr = gw_lis.local_addr().unwrap();
+ let (gw_tx, gw_rx) = tokio::sync::oneshot::channel();
+ tokio::spawn(async move {
+ Server::builder()
+ .add_service(pb::spark_connect_service_server::SparkConnectServiceServer::new(proxy))
+ .serve_with_incoming_shutdown(TcpListenerStream::new(gw_lis), async {
+ let _ = gw_rx.await;
+ })
+ .await
+ .ok();
+ });
+
+ let endpoint = Endpoint::from_shared(format!("http://{}", gw_addr)).unwrap();
+ let grpc = endpoint
+ .connect_timeout(Duration::from_secs(2))
+ .connect()
+ .await
+ .unwrap();
+
+ // Admin server on a separate ephemeral port
+ let admin_lis = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
+ let admin_addr = admin_lis.local_addr().unwrap();
+ drop(admin_lis);
+ let admin_cfg = AdminConfig {
+ bind_addr: admin_addr,
+ };
+ let metrics_for_admin = metrics.clone();
+ let readiness = ReadinessProbe::new(true);
+ let (admin_tx, admin_rx) = tokio::sync::oneshot::channel();
+ tokio::spawn(async move {
+ let shutdown = async move {
+ let _ = admin_rx.await;
+ };
+ let _ = serve_admin(admin_cfg, metrics_for_admin, readiness, shutdown).await;
+ });
+
+ // Wait briefly for admin server to bind
+ for _ in 0..50 {
+ if std::net::TcpStream::connect(admin_addr).is_ok() {
+ break;
+ }
+ tokio::time::sleep(Duration::from_millis(20)).await;
+ }
+
+ let admin_url = format!("http://{}", admin_addr);
+ (
+ Rig {
+ grpc,
+ admin_url,
+ _be_shutdown: be_tx,
+ _gw_shutdown: gw_tx,
+ _admin_shutdown: admin_tx,
+ },
+ metrics,
+ )
+}
+
+async fn scrape(url: &str) -> String {
+ let body = reqwest::get(format!("{}/metrics", url))
+ .await
+ .unwrap()
+ .text()
+ .await
+ .unwrap();
+ body
+}
+
+#[tokio::test]
+async fn successful_rpc_increments_rpcs_total() {
+ let (rig, _metrics) = rig_with_metrics().await;
+ let mut c = pb::spark_connect_service_client::SparkConnectServiceClient::new(rig.grpc.clone());
+
+ let mut req = Request::new(pb::ConfigRequest {
+ session_id: "s1".into(),
+ ..Default::default()
+ });
+ req.metadata_mut().insert(
+ "authorization",
+ MetadataValue::try_from("Bearer alice-secret").unwrap(),
+ );
+ c.config(req).await.unwrap();
+
+ // /metrics should now contain a non-zero counter for rpc=Config code=OK.
+ let body = scrape(&rig.admin_url).await;
+ assert!(
+ body.contains("scg_rpcs_total"),
+ "/metrics missing scg_rpcs_total:\n{body}"
+ );
+ assert!(
+ body.lines().any(|l| l.starts_with("scg_rpcs_total{")
+ && l.contains("rpc=\"Config\"")
+ && l.contains("code=\"OK\"")),
+ "no Config OK counter line:\n{body}"
+ );
+}
+
+#[tokio::test]
+async fn unauthenticated_request_increments_auth_failures() {
+ let (rig, _metrics) = rig_with_metrics().await;
+ let mut c = pb::spark_connect_service_client::SparkConnectServiceClient::new(rig.grpc.clone());
+
+ let err = c
+ .config(pb::ConfigRequest {
+ session_id: "s1".into(),
+ ..Default::default()
+ })
+ .await
+ .unwrap_err();
+ assert_eq!(err.code(), tonic::Code::Unauthenticated);
+
+ let body = scrape(&rig.admin_url).await;
+ assert!(
+ body.contains("scg_auth_failures_total"),
+ "/metrics missing scg_auth_failures_total:\n{body}"
+ );
+ assert!(
+ body.lines()
+ .any(|l| l.starts_with("scg_auth_failures_total{")
+ && l.contains("reason=\"missing_token\"")),
+ "no missing_token failure line:\n{body}"
+ );
+ // The RPC still gets recorded with code=Unauthenticated.
+ assert!(
+ body.lines()
+ .any(|l| l.starts_with("scg_rpcs_total{") && l.contains("code=\"Unauthenticated\"")),
+ "no Unauthenticated rpc total line:\n{body}"
+ );
+}
+
+#[tokio::test]
+async fn healthz_and_readyz_endpoints_work() {
+ let (rig, _metrics) = rig_with_metrics().await;
+ let body = reqwest::get(format!("{}/healthz", rig.admin_url))
+ .await
+ .unwrap()
+ .text()
+ .await
+ .unwrap();
+ assert_eq!(body, "ok");
+
+ let body = reqwest::get(format!("{}/readyz", rig.admin_url))
+ .await
+ .unwrap()
+ .text()
+ .await
+ .unwrap();
+ assert_eq!(body, "ready");
+}
diff --git a/crates/proxy/tests/outbound_token.rs b/crates/proxy/tests/outbound_token.rs
new file mode 100644
index 0000000..44be294
--- /dev/null
+++ b/crates/proxy/tests/outbound_token.rs
@@ -0,0 +1,577 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Integration tests for the gateway→backend bearer token.
+//!
+//! Backends started with `spark.connect.authenticate.token` require
+//! `authorization: Bearer <token>` on every request and reject
+//! everything else with `UNAUTHENTICATED` (Spark's
+//! `PreSharedKeyAuthenticationInterceptor` compares the exact
+//! string). The mock backend here reproduces that check, so these
+//! tests cover both halves of the trust boundary at the process
+//! level: a direct connection without the token is refused, and the
+//! same request through a token-configured gateway succeeds.
+
+use std::collections::HashMap;
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+
+use scg_genproto::pb;
+use scg_proxy::{BackendTokens, Dialer, SparkConnectProxy};
+use scg_routing::{AffinityStore, Pool, Router, TenantRouter};
+use scg_store_memory::MemoryStore;
+use tokio::net::TcpListener;
+use tokio_stream::wrappers::TcpListenerStream;
+use tonic::transport::{Channel, Endpoint, Server};
+use tonic::{Request, Response, Status};
+
+/// Records the `authorization` metadata of every `Config` call and,
+/// when `expected` is set, enforces it exactly the way Spark's
+/// `PreSharedKeyAuthenticationInterceptor` does.
+struct RecordingBackend {
+ seen: Arc<Mutex<Vec<Option<String>>>>,
+ expected: Option<String>,
+}
+
+impl RecordingBackend {
+ fn check<T>(&self, req: &Request<T>) -> Result<(), Status> {
+ let header = req
+ .metadata()
+ .get("authorization")
+ .and_then(|v| v.to_str().ok())
+ .map(str::to_string);
+ self.seen.lock().unwrap().push(header.clone());
+ match &self.expected {
+ None => Ok(()),
+ Some(want) => match header {
+ None => Err(Status::unauthenticated("No authentication token provided")),
+ Some(got) if &got == want => Ok(()),
+ Some(_) => Err(Status::unauthenticated("Invalid authentication token")),
+ },
+ }
+ }
+}
+
+#[tonic::async_trait]
+impl pb::spark_connect_service_server::SparkConnectService for RecordingBackend {
+ type ExecutePlanStream = std::pin::Pin<
+ Box<dyn futures::Stream<Item = Result<pb::ExecutePlanResponse, Status>> + Send + 'static>,
+ >;
+ type ReattachExecuteStream = Self::ExecutePlanStream;
+
+ async fn config(
+ &self,
+ req: Request<pb::ConfigRequest>,
+ ) -> Result<Response<pb::ConfigResponse>, Status> {
+ self.check(&req)?;
+ let body = req.into_inner();
+ Ok(Response::new(pb::ConfigResponse {
+ session_id: body.session_id,
+ ..Default::default()
+ }))
+ }
+
+ async fn analyze_plan(
+ &self,
+ _: Request<pb::AnalyzePlanRequest>,
+ ) -> Result<Response<pb::AnalyzePlanResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn artifact_status(
+ &self,
+ _: Request<pb::ArtifactStatusesRequest>,
+ ) -> Result<Response<pb::ArtifactStatusesResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn interrupt(
+ &self,
+ _: Request<pb::InterruptRequest>,
+ ) -> Result<Response<pb::InterruptResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn release_execute(
+ &self,
+ _: Request<pb::ReleaseExecuteRequest>,
+ ) -> Result<Response<pb::ReleaseExecuteResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn release_session(
+ &self,
+ _: Request<pb::ReleaseSessionRequest>,
+ ) -> Result<Response<pb::ReleaseSessionResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn fetch_error_details(
+ &self,
+ _: Request<pb::FetchErrorDetailsRequest>,
+ ) -> Result<Response<pb::FetchErrorDetailsResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn get_status(
+ &self,
+ _: Request<pb::GetStatusRequest>,
+ ) -> Result<Response<pb::GetStatusResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn clone_session(
+ &self,
+ _: Request<pb::CloneSessionRequest>,
+ ) -> Result<Response<pb::CloneSessionResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn execute_plan(
+ &self,
+ _: Request<pb::ExecutePlanRequest>,
+ ) -> Result<Response<Self::ExecutePlanStream>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn reattach_execute(
+ &self,
+ _: Request<pb::ReattachExecuteRequest>,
+ ) -> Result<Response<Self::ReattachExecuteStream>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn add_artifacts(
+ &self,
+ _: Request<tonic::Streaming<pb::AddArtifactsRequest>>,
+ ) -> Result<Response<pb::AddArtifactsResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+}
+
+type Seen = Arc<Mutex<Vec<Option<String>>>>;
+
+async fn spawn_backend(expected: Option<&str>) -> (String, Seen, tokio::sync::oneshot::Sender<()>) {
+ let seen: Seen = Arc::default();
+ let backend = RecordingBackend {
+ seen: seen.clone(),
+ expected: expected.map(str::to_string),
+ };
+ let lis = TcpListener::bind("127.0.0.1:0").await.unwrap();
+ let addr = lis.local_addr().unwrap().to_string();
+ let (tx, rx) = tokio::sync::oneshot::channel();
+ tokio::spawn(async move {
+ Server::builder()
+ .add_service(pb::spark_connect_service_server::SparkConnectServiceServer::new(backend))
+ .serve_with_incoming_shutdown(TcpListenerStream::new(lis), async {
+ let _ = rx.await;
+ })
+ .await
+ .ok();
+ });
+ (addr, seen, tx)
+}
+
+struct Rig {
+ channel: Channel,
+ backend_addr: String,
+ seen: Seen,
+ _backend_shutdown: tokio::sync::oneshot::Sender<()>,
+ _gw_shutdown: tokio::sync::oneshot::Sender<()>,
+}
+
+async fn spawn_rig(backend_expected: Option<&str>, tokens: BackendTokens) -> Rig {
+ let (backend_addr, seen, backend_tx) = spawn_backend(backend_expected).await;
+
+ let store: Arc<dyn AffinityStore> = Arc::new(MemoryStore::new());
+ let pool: Arc<dyn Pool> =
+ Arc::new(scg_pool_static::StaticPool::new(vec![backend_addr.clone()]).unwrap());
+ let router = Arc::new(Router::new(TenantRouter::single(pool), store));
+ let proxy = SparkConnectProxy::new(router, Dialer::new()).with_backend_tokens(tokens);
+
+ let gw_lis = TcpListener::bind("127.0.0.1:0").await.unwrap();
+ let gw_addr = gw_lis.local_addr().unwrap();
+ let (gw_tx, gw_rx) = tokio::sync::oneshot::channel();
+ tokio::spawn(async move {
+ Server::builder()
+ .add_service(pb::spark_connect_service_server::SparkConnectServiceServer::new(proxy))
+ .serve_with_incoming_shutdown(TcpListenerStream::new(gw_lis), async {
+ let _ = gw_rx.await;
+ })
+ .await
+ .ok();
+ });
+
+ let channel = Endpoint::from_shared(format!("http://{}", gw_addr))
+ .unwrap()
+ .connect_timeout(Duration::from_secs(2))
+ .connect()
+ .await
+ .unwrap();
+
+ Rig {
+ channel,
+ backend_addr,
+ seen,
+ _backend_shutdown: backend_tx,
+ _gw_shutdown: gw_tx,
+ }
+}
+
+fn client(
+ channel: Channel,
+) -> pb::spark_connect_service_client::SparkConnectServiceClient<Channel> {
+ pb::spark_connect_service_client::SparkConnectServiceClient::new(channel)
+}
+
+fn config_req(session: &str) -> pb::ConfigRequest {
+ pb::ConfigRequest {
+ session_id: session.into(),
+ ..Default::default()
+ }
+}
+
+#[tokio::test]
+async fn default_token_reaches_the_backend() {
+ let tokens = BackendTokens::new(Some("hunter2".into()), HashMap::new()).unwrap();
+ let rig = spawn_rig(None, tokens).await;
+ client(rig.channel.clone())
+ .config(config_req("s1"))
+ .await
+ .unwrap();
+ assert_eq!(
+ rig.seen.lock().unwrap().as_slice(),
+ &[Some("Bearer hunter2".to_string())]
+ );
+}
+
+#[tokio::test]
+async fn no_token_configured_sends_no_header() {
+ let rig = spawn_rig(None, BackendTokens::none()).await;
+ client(rig.channel.clone())
+ .config(config_req("s1"))
+ .await
+ .unwrap();
+ assert_eq!(rig.seen.lock().unwrap().as_slice(), &[None]);
+}
+
+#[tokio::test]
+async fn inbound_client_credential_is_not_forwarded() {
+ // The gateway's own inbound auth may be bearer-based too; the
+ // client's credential must never leak to the backend — the
+ // backend sees the pool token or nothing.
+ let rig = spawn_rig(None, BackendTokens::none()).await;
+ let mut req = Request::new(config_req("s1"));
+ req.metadata_mut()
+ .insert("authorization", "Bearer client-secret".parse().unwrap());
+ client(rig.channel.clone()).config(req).await.unwrap();
+ assert_eq!(rig.seen.lock().unwrap().as_slice(), &[None]);
+}
+
+#[tokio::test]
+async fn enforcing_backend_rejects_direct_but_accepts_gateway() {
+ // The trust-boundary property, at process level: the same
+ // backend that refuses a tokenless direct connection accepts
+ // the gateway's, because only the gateway holds the token.
+ let tokens = BackendTokens::new(Some("gw-only".into()), HashMap::new()).unwrap();
+ let rig = spawn_rig(Some("Bearer gw-only"), tokens).await;
+
+ // Direct to the backend, no token: UNAUTHENTICATED.
+ let direct = Endpoint::from_shared(format!("http://{}", rig.backend_addr))
+ .unwrap()
+ .connect()
+ .await
+ .unwrap();
+ let err = client(direct)
+ .config(config_req("s-direct"))
+ .await
+ .unwrap_err();
+ assert_eq!(err.code(), tonic::Code::Unauthenticated);
+
+ // Through the gateway: OK.
+ client(rig.channel.clone())
+ .config(config_req("s-gw"))
+ .await
+ .unwrap();
+}
+
+/// Backend that answers `Config` the way Spark's
+/// `SparkConnectConfigHandler` does: whatever key you ask for, you
+/// get its value back — including the server's own pre-shared token.
+struct LeakyConfigBackend {
+ token: String,
+}
+
+impl LeakyConfigBackend {
+ fn lookup(&self, key: &str) -> Option<String> {
+ match key {
+ "spark.connect.authenticate.token" => Some(self.token.clone()),
+ "spark.sql.shuffle.partitions" => Some("200".into()),
+ _ => None,
+ }
+ }
+}
+
+#[tonic::async_trait]
+impl pb::spark_connect_service_server::SparkConnectService for LeakyConfigBackend {
+ type ExecutePlanStream = std::pin::Pin<
+ Box<dyn futures::Stream<Item = Result<pb::ExecutePlanResponse, Status>> + Send + 'static>,
+ >;
+ type ReattachExecuteStream = Self::ExecutePlanStream;
+
+ async fn config(
+ &self,
+ req: Request<pb::ConfigRequest>,
+ ) -> Result<Response<pb::ConfigResponse>, Status> {
+ use pb::config_request::operation::OpType;
+ let body = req.into_inner();
+ let mut pairs = Vec::new();
+ match body.operation.as_ref().and_then(|o| o.op_type.as_ref()) {
+ Some(OpType::Get(get)) => {
+ for key in &get.keys {
+ pairs.push(pb::KeyValue {
+ key: key.clone(),
+ value: self.lookup(key),
+ });
+ }
+ }
+ Some(OpType::GetAll(get_all)) => {
+ // Mirrors handleGetAll: filter by prefix, then return
+ // the keys with that prefix stripped off.
+ let prefix = get_all.prefix.clone().unwrap_or_default();
+ for key in [
+ "spark.connect.authenticate.token",
+ "spark.sql.shuffle.partitions",
+ ] {
+ if let Some(stripped) = key.strip_prefix(prefix.as_str()) {
+ pairs.push(pb::KeyValue {
+ key: stripped.to_string(),
+ value: self.lookup(key),
+ });
+ }
+ }
+ }
+ _ => return Err(Status::unimplemented("only Get/GetAll in this fixture")),
+ }
+ Ok(Response::new(pb::ConfigResponse {
+ session_id: body.session_id,
+ pairs,
+ ..Default::default()
+ }))
+ }
+
+ async fn analyze_plan(
+ &self,
+ _: Request<pb::AnalyzePlanRequest>,
+ ) -> Result<Response<pb::AnalyzePlanResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn artifact_status(
+ &self,
+ _: Request<pb::ArtifactStatusesRequest>,
+ ) -> Result<Response<pb::ArtifactStatusesResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn interrupt(
+ &self,
+ _: Request<pb::InterruptRequest>,
+ ) -> Result<Response<pb::InterruptResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn release_execute(
+ &self,
+ _: Request<pb::ReleaseExecuteRequest>,
+ ) -> Result<Response<pb::ReleaseExecuteResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn release_session(
+ &self,
+ _: Request<pb::ReleaseSessionRequest>,
+ ) -> Result<Response<pb::ReleaseSessionResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn fetch_error_details(
+ &self,
+ _: Request<pb::FetchErrorDetailsRequest>,
+ ) -> Result<Response<pb::FetchErrorDetailsResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn get_status(
+ &self,
+ _: Request<pb::GetStatusRequest>,
+ ) -> Result<Response<pb::GetStatusResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn clone_session(
+ &self,
+ _: Request<pb::CloneSessionRequest>,
+ ) -> Result<Response<pb::CloneSessionResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn execute_plan(
+ &self,
+ _: Request<pb::ExecutePlanRequest>,
+ ) -> Result<Response<Self::ExecutePlanStream>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn reattach_execute(
+ &self,
+ _: Request<pb::ReattachExecuteRequest>,
+ ) -> Result<Response<Self::ReattachExecuteStream>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn add_artifacts(
+ &self,
+ _: Request<tonic::Streaming<pb::AddArtifactsRequest>>,
+ ) -> Result<Response<pb::AddArtifactsResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+}
+
+async fn spawn_leaky_rig(
+ token: &str,
+) -> (
+ Channel,
+ tokio::sync::oneshot::Sender<()>,
+ tokio::sync::oneshot::Sender<()>,
+) {
+ let backend = LeakyConfigBackend {
+ token: token.to_string(),
+ };
+ let lis = TcpListener::bind("127.0.0.1:0").await.unwrap();
+ let backend_addr = lis.local_addr().unwrap().to_string();
+ let (btx, brx) = tokio::sync::oneshot::channel();
+ tokio::spawn(async move {
+ Server::builder()
+ .add_service(pb::spark_connect_service_server::SparkConnectServiceServer::new(backend))
+ .serve_with_incoming_shutdown(TcpListenerStream::new(lis), async {
+ let _ = brx.await;
+ })
+ .await
+ .ok();
+ });
+
+ let store: Arc<dyn AffinityStore> = Arc::new(MemoryStore::new());
+ let pool: Arc<dyn Pool> =
+ Arc::new(scg_pool_static::StaticPool::new(vec![backend_addr]).unwrap());
+ let router = Arc::new(Router::new(TenantRouter::single(pool), store));
+ let tokens = BackendTokens::new(Some(token.to_string()), HashMap::new()).unwrap();
+ let proxy = SparkConnectProxy::new(router, Dialer::new()).with_backend_tokens(tokens);
+
+ let gw_lis = TcpListener::bind("127.0.0.1:0").await.unwrap();
+ let gw_addr = gw_lis.local_addr().unwrap();
+ let (gtx, grx) = tokio::sync::oneshot::channel();
+ tokio::spawn(async move {
+ Server::builder()
+ .add_service(pb::spark_connect_service_server::SparkConnectServiceServer::new(proxy))
+ .serve_with_incoming_shutdown(TcpListenerStream::new(gw_lis), async {
+ let _ = grx.await;
+ })
+ .await
+ .ok();
+ });
+
+ let channel = Endpoint::from_shared(format!("http://{}", gw_addr))
+ .unwrap()
+ .connect_timeout(Duration::from_secs(2))
+ .connect()
+ .await
+ .unwrap();
+ (channel, btx, gtx)
+}
+
+fn config_get(keys: &[&str]) -> pb::ConfigRequest {
+ pb::ConfigRequest {
+ session_id: "s1".into(),
+ operation: Some(pb::config_request::Operation {
+ op_type: Some(pb::config_request::operation::OpType::Get(
+ pb::config_request::Get {
+ keys: keys.iter().map(|k| k.to_string()).collect(),
+ },
+ )),
+ }),
+ ..Default::default()
+ }
+}
+
+fn config_get_all(prefix: Option<&str>) -> pb::ConfigRequest {
+ pb::ConfigRequest {
+ session_id: "s1".into(),
+ operation: Some(pb::config_request::Operation {
+ op_type: Some(pb::config_request::operation::OpType::GetAll(
+ pb::config_request::GetAll {
+ prefix: prefix.map(str::to_string),
+ },
+ )),
+ }),
+ ..Default::default()
+ }
+}
+
+#[tokio::test]
+async fn backend_token_is_not_readable_through_the_gateway() {
+ // The end-to-end property: a client the gateway authorized must
+ // not be able to read the gateway↔backend credential, even
+ // though the backend hands it over on request.
+ let (channel, _b, _g) = spawn_leaky_rig("SUPERSECRET").await;
+ let resp = client(channel)
+ .config(config_get(&[
+ "spark.connect.authenticate.token",
+ "spark.sql.shuffle.partitions",
+ ]))
+ .await
+ .unwrap()
+ .into_inner();
+
+ assert_eq!(resp.pairs.len(), 2);
+ assert_eq!(resp.pairs[0].key, "spark.connect.authenticate.token");
+ assert_eq!(resp.pairs[0].value, None, "token must not reach the client");
+ assert_eq!(resp.pairs[1].value.as_deref(), Some("200"));
+}
+
+#[tokio::test]
+async fn backend_token_is_not_readable_via_get_all() {
+ let (channel, _b, _g) = spawn_leaky_rig("SUPERSECRET").await;
+ let resp = client(channel)
+ .config(config_get_all(None))
+ .await
+ .unwrap()
+ .into_inner();
+
+ assert!(
+ !resp
+ .pairs
+ .iter()
+ .any(|p| p.value.as_deref() == Some("SUPERSECRET")),
+ "token leaked through GetAll: {:?}",
+ resp.pairs
+ );
+ assert!(resp
+ .pairs
+ .iter()
+ .any(|p| p.key == "spark.sql.shuffle.partitions"));
+}
+
+#[tokio::test]
+async fn backend_token_is_not_readable_via_prefixed_get_all() {
+ // The prefix form is the subtle one: the backend strips the
+ // prefix from the keys it returns, so the response says
+ // `authenticate.token`, not the full key.
+ let (channel, _b, _g) = spawn_leaky_rig("SUPERSECRET").await;
+ let resp = client(channel)
+ .config(config_get_all(Some("spark.connect.")))
+ .await
+ .unwrap()
+ .into_inner();
+
+ assert!(
+ !resp
+ .pairs
+ .iter()
+ .any(|p| p.value.as_deref() == Some("SUPERSECRET")),
+ "token leaked through prefixed GetAll: {:?}",
+ resp.pairs
+ );
+}
diff --git a/crates/proxy/tests/ratelimit_integration.rs b/crates/proxy/tests/ratelimit_integration.rs
new file mode 100644
index 0000000..0d95a71
--- /dev/null
+++ b/crates/proxy/tests/ratelimit_integration.rs
@@ -0,0 +1,358 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Integration test for per-tenant rate limiting.
+//!
+//! Drives RPCs through a real gRPC server with the rate limiter
+//! enabled and asserts:
+//!
+//! * burst RPCs go through; the next one gets ResourceExhausted
+//! * tenant A exhausting its quota doesn't block tenant B
+//! * per-tenant overrides take precedence over the default
+//! * the metric `scg_rate_limit_rejected_total{tenant, scope}`
+//! increments on rejection
+
+use std::collections::HashMap;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::time::Duration;
+
+use futures::Stream;
+use scg_auth::{AnonymousAuthenticator, AuthInterceptor};
+use scg_genproto::pb;
+use scg_observability::Metrics;
+use scg_pool_static::StaticPool;
+use scg_proxy::{Dialer, SparkConnectProxy};
+use scg_ratelimit::{BucketRate, LimiterObserver, RateLimiter, RejectScope, TenantLimits};
+use scg_routing::{AffinityStore, Pool, Router, TenantRouter};
+use scg_store_memory::MemoryStore;
+use scg_tenant::{OnMissing, TenantResolver, TenantResolverConfig, TenantSource};
+use tokio::net::TcpListener;
+use tokio_stream::wrappers::TcpListenerStream;
+use tonic::metadata::MetadataValue;
+use tonic::transport::{Channel, Endpoint, Server};
+use tonic::{Request, Response, Status};
+
+/// Minimal backend — every Config call succeeds. Tests below count
+/// successful vs ResourceExhausted RPCs at the client side, so the
+/// backend just needs to not crash.
+#[derive(Default)]
+struct OkBackend;
+
+#[tonic::async_trait]
+impl pb::spark_connect_service_server::SparkConnectService for OkBackend {
+ type ExecutePlanStream =
+ Pin<Box<dyn Stream<Item = Result<pb::ExecutePlanResponse, Status>> + Send + 'static>>;
+ type ReattachExecuteStream = Self::ExecutePlanStream;
+
+ async fn config(
+ &self,
+ req: Request<pb::ConfigRequest>,
+ ) -> Result<Response<pb::ConfigResponse>, Status> {
+ let body = req.into_inner();
+ Ok(Response::new(pb::ConfigResponse {
+ session_id: body.session_id,
+ ..Default::default()
+ }))
+ }
+ async fn analyze_plan(
+ &self,
+ _: Request<pb::AnalyzePlanRequest>,
+ ) -> Result<Response<pb::AnalyzePlanResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn artifact_status(
+ &self,
+ _: Request<pb::ArtifactStatusesRequest>,
+ ) -> Result<Response<pb::ArtifactStatusesResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn interrupt(
+ &self,
+ _: Request<pb::InterruptRequest>,
+ ) -> Result<Response<pb::InterruptResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn release_execute(
+ &self,
+ _: Request<pb::ReleaseExecuteRequest>,
+ ) -> Result<Response<pb::ReleaseExecuteResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn release_session(
+ &self,
+ _: Request<pb::ReleaseSessionRequest>,
+ ) -> Result<Response<pb::ReleaseSessionResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn fetch_error_details(
+ &self,
+ _: Request<pb::FetchErrorDetailsRequest>,
+ ) -> Result<Response<pb::FetchErrorDetailsResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn clone_session(
+ &self,
+ _: Request<pb::CloneSessionRequest>,
+ ) -> Result<Response<pb::CloneSessionResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn get_status(
+ &self,
+ _: Request<pb::GetStatusRequest>,
+ ) -> Result<Response<pb::GetStatusResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn execute_plan(
+ &self,
+ _: Request<pb::ExecutePlanRequest>,
+ ) -> Result<Response<Self::ExecutePlanStream>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn reattach_execute(
+ &self,
+ _: Request<pb::ReattachExecuteRequest>,
+ ) -> Result<Response<Self::ReattachExecuteStream>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn add_artifacts(
+ &self,
+ _: Request<tonic::Streaming<pb::AddArtifactsRequest>>,
+ ) -> Result<Response<pb::AddArtifactsResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+}
+
+async fn spawn_rig(
+ limiter: Option<RateLimiter>,
+) -> (
+ Channel,
+ Metrics,
+ tokio::sync::oneshot::Sender<()>,
+ tokio::sync::oneshot::Sender<()>,
+) {
+ let be_lis = TcpListener::bind("127.0.0.1:0").await.unwrap();
+ let be_addr = be_lis.local_addr().unwrap().to_string();
+ let (be_tx, be_rx) = tokio::sync::oneshot::channel();
+ tokio::spawn(async move {
+ Server::builder()
+ .add_service(
+ pb::spark_connect_service_server::SparkConnectServiceServer::new(OkBackend),
+ )
+ .serve_with_incoming_shutdown(TcpListenerStream::new(be_lis), async {
+ let _ = be_rx.await;
+ })
+ .await
+ .ok();
+ });
+
+ let metrics = Metrics::new().unwrap();
+ let pool: Arc<dyn Pool> = Arc::new(StaticPool::new(vec![be_addr]).unwrap());
+ let store: Arc<dyn AffinityStore> = Arc::new(MemoryStore::new());
+ let tr = TenantRouter::single(pool);
+ let router = Arc::new(Router::new(tr, store));
+ let dialer = Dialer::new();
+ let resolver = TenantResolver::new(TenantResolverConfig {
+ source: TenantSource::FromMetadata {
+ header: "x-tenant".into(),
+ },
+ on_missing: OnMissing::UseDefault,
+ default_name: "default".into(),
+ });
+ let proxy = SparkConnectProxy::with_all(
+ router,
+ dialer,
+ AuthInterceptor::new(Arc::new(AnonymousAuthenticator)),
+ metrics.clone(),
+ resolver,
+ limiter,
+ scg_audit::AuditLogger::disabled(),
+ );
+
+ let gw_lis = TcpListener::bind("127.0.0.1:0").await.unwrap();
+ let gw_addr = gw_lis.local_addr().unwrap();
+ let (gw_tx, gw_rx) = tokio::sync::oneshot::channel();
+ tokio::spawn(async move {
+ Server::builder()
+ .add_service(pb::spark_connect_service_server::SparkConnectServiceServer::new(proxy))
+ .serve_with_incoming_shutdown(TcpListenerStream::new(gw_lis), async {
+ let _ = gw_rx.await;
+ })
+ .await
+ .ok();
+ });
+
+ let endpoint = Endpoint::from_shared(format!("http://{}", gw_addr)).unwrap();
+ let ch = endpoint
+ .connect_timeout(Duration::from_secs(2))
+ .connect()
+ .await
+ .unwrap();
+ (ch, metrics, be_tx, gw_tx)
+}
+
+fn metrics_observer(metrics: &Metrics) -> Arc<dyn LimiterObserver> {
+ struct Obs(Metrics);
+ impl LimiterObserver for Obs {
+ fn on_reject(&self, tenant: &str, scope: RejectScope) {
+ self.0.record_rate_limit_reject(tenant, scope.as_str());
+ }
+ }
+ Arc::new(Obs(metrics.clone()))
+}
+
+async fn config_with_tenant(ch: &Channel, session: &str, tenant: &str) -> Result<(), Status> {
+ let mut c = pb::spark_connect_service_client::SparkConnectServiceClient::new(ch.clone());
+ let mut req = Request::new(pb::ConfigRequest {
+ session_id: session.into(),
+ ..Default::default()
+ });
+ req.metadata_mut()
+ .insert("x-tenant", MetadataValue::try_from(tenant).unwrap());
+ c.config(req).await.map(|_| ())
+}
+
+fn rejected_count(metrics: &Metrics, tenant: &str, scope: &str) -> u64 {
+ let mfs = metrics.registry().gather();
+ for mf in mfs {
+ if mf.name() != "scg_rate_limit_rejected_total" {
+ continue;
+ }
+ for m in mf.get_metric() {
+ let labels = m.get_label();
+ let t = labels
+ .iter()
+ .find(|l| l.name() == "tenant")
+ .map(|l| l.value());
+ let s = labels
+ .iter()
+ .find(|l| l.name() == "scope")
+ .map(|l| l.value());
+ if t == Some(tenant) && s == Some(scope) {
+ return m.get_counter().value() as u64;
+ }
+ }
+ }
+ 0
+}
+
+#[tokio::test]
+async fn burst_admits_then_rejects_at_tenant_scope() {
+ // Need a non-trivial Metrics to act as the observer's sink.
+ let observer_metrics = Metrics::new().unwrap();
+ let mut overrides = HashMap::new();
+ // 1 RPS / burst 3 → first 3 succeed, 4th gets ResourceExhausted.
+ overrides.insert(
+ "team-a".to_string(),
+ TenantLimits {
+ tenant: BucketRate {
+ rpcs_per_second: 1.0,
+ burst: 3,
+ },
+ per_user: BucketRate::disabled(),
+ },
+ );
+ let limiter = RateLimiter::new(
+ TenantLimits::default(),
+ overrides,
+ metrics_observer(&observer_metrics),
+ );
+
+ let (ch, _metrics, _be, _gw) = spawn_rig(Some(limiter)).await;
+
+ for i in 0..3 {
+ config_with_tenant(&ch, &format!("s-{}", i), "team-a")
+ .await
+ .unwrap_or_else(|e| panic!("RPC {} should succeed, got {:?}", i, e));
+ }
+ let err = config_with_tenant(&ch, "s-4", "team-a").await.unwrap_err();
+ assert_eq!(err.code(), tonic::Code::ResourceExhausted);
+
+ // Metric incremented at tenant scope.
+ assert!(rejected_count(&observer_metrics, "team-a", "tenant") >= 1);
+}
+
+#[tokio::test]
+async fn one_tenant_exhaustion_does_not_block_others() {
+ let observer_metrics = Metrics::new().unwrap();
+ let mut overrides = HashMap::new();
+ let one_token = TenantLimits {
+ tenant: BucketRate {
+ rpcs_per_second: 0.1, // refill barely matters within the test
+ burst: 1,
+ },
+ per_user: BucketRate::disabled(),
+ };
+ overrides.insert("team-a".to_string(), one_token);
+ overrides.insert("team-b".to_string(), one_token);
+ let limiter = RateLimiter::new(
+ TenantLimits::default(),
+ overrides,
+ metrics_observer(&observer_metrics),
+ );
+
+ let (ch, _metrics, _be, _gw) = spawn_rig(Some(limiter)).await;
+
+ // Exhaust team-a.
+ config_with_tenant(&ch, "s-a-1", "team-a").await.unwrap();
+ let err = config_with_tenant(&ch, "s-a-2", "team-a")
+ .await
+ .unwrap_err();
+ assert_eq!(err.code(), tonic::Code::ResourceExhausted);
+
+ // team-b is unaffected.
+ config_with_tenant(&ch, "s-b-1", "team-b").await.unwrap();
+}
+
+#[tokio::test]
+async fn default_bucket_applies_to_unlisted_tenants() {
+ let observer_metrics = Metrics::new().unwrap();
+ // No explicit overrides; default = 0.1 RPS, burst 2.
+ let limiter = RateLimiter::new(
+ TenantLimits {
+ tenant: BucketRate {
+ rpcs_per_second: 0.1,
+ burst: 2,
+ },
+ per_user: BucketRate::disabled(),
+ },
+ HashMap::new(),
+ metrics_observer(&observer_metrics),
+ );
+
+ let (ch, _metrics, _be, _gw) = spawn_rig(Some(limiter)).await;
+
+ // A random unlisted tenant still gets the default bucket.
+ config_with_tenant(&ch, "s1", "random-tenant")
+ .await
+ .unwrap();
+ config_with_tenant(&ch, "s2", "random-tenant")
+ .await
+ .unwrap();
+ let err = config_with_tenant(&ch, "s3", "random-tenant")
+ .await
+ .unwrap_err();
+ assert_eq!(err.code(), tonic::Code::ResourceExhausted);
+}
+
+#[tokio::test]
+async fn no_limiter_admits_everything() {
+ let (ch, _metrics, _be, _gw) = spawn_rig(None).await;
+ for i in 0..50 {
+ config_with_tenant(&ch, &format!("s-{}", i), "any-tenant")
+ .await
+ .unwrap();
+ }
+}
diff --git a/crates/proxy/tests/tenant_integration.rs b/crates/proxy/tests/tenant_integration.rs
new file mode 100644
index 0000000..a436e4b
--- /dev/null
+++ b/crates/proxy/tests/tenant_integration.rs
@@ -0,0 +1,401 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Integration tests for tenant resolution: verifies that
+//! the routing key seen by the affinity store carries the tenant
+//! the resolver produced, that resolver=Reject failure surfaces as
+//! Unauthenticated at the client, and that two tenants with the
+//! same session_id stay isolated.
+
+use std::pin::Pin;
+use std::sync::Arc;
+use std::time::Duration;
+
+use futures::Stream;
+use parking_lot::Mutex;
+use scg_auth::{
+ token::{StaticTokenAuthenticator, TokenEntry},
+ AnonymousAuthenticator, AuthInterceptor,
+};
+use scg_genproto::pb;
+use scg_observability::Metrics;
+use scg_pool_static::StaticPool;
+use scg_proxy::{Dialer, SparkConnectProxy};
+use scg_routing::{AffinityStore, Pool, Router, SessionKey};
+use scg_store_memory::MemoryStore;
+use scg_tenant::{OnMissing, TenantResolver, TenantResolverConfig, TenantSource};
+use tokio::net::TcpListener;
+use tokio_stream::wrappers::TcpListenerStream;
+use tonic::metadata::MetadataValue;
+use tonic::transport::{Channel, Endpoint, Server};
+use tonic::{Request, Response, Status};
+
+/// AffinityStore wrapper that records every `bind_session_if_absent`
+/// key the proxy hands it. Tests inspect the records to confirm the
+/// routing key carries the expected tenant.
+struct RecordingStore {
+ inner: MemoryStore,
+ binds: Arc<Mutex<Vec<SessionKey>>>,
+}
+
+impl RecordingStore {
+ fn new() -> (Self, Arc<Mutex<Vec<SessionKey>>>) {
+ let binds = Arc::new(Mutex::new(Vec::new()));
+ (
+ Self {
+ inner: MemoryStore::new(),
+ binds: binds.clone(),
+ },
+ binds,
+ )
+ }
+}
+
+#[async_trait::async_trait]
+impl AffinityStore for RecordingStore {
+ async fn lookup_session(&self, key: &SessionKey) -> Option<String> {
+ self.inner.lookup_session(key).await
+ }
+ async fn bind_session_if_absent(&self, key: SessionKey, backend: String) -> String {
+ self.binds.lock().push(key.clone());
+ self.inner.bind_session_if_absent(key, backend).await
+ }
+ async fn forget_session(&self, key: &SessionKey) {
+ self.inner.forget_session(key).await
+ }
+ async fn lookup_op(&self, op_id: &str) -> Option<String> {
+ self.inner.lookup_op(op_id).await
+ }
+ async fn bind_op(&self, op_id: String, backend: String) {
+ self.inner.bind_op(op_id, backend).await
+ }
+ async fn forget_op(&self, op_id: &str) {
+ self.inner.forget_op(op_id).await
+ }
+}
+
+#[derive(Default)]
+struct EchoBackend;
+
+#[tonic::async_trait]
+impl pb::spark_connect_service_server::SparkConnectService for EchoBackend {
+ type ExecutePlanStream =
+ Pin<Box<dyn Stream<Item = Result<pb::ExecutePlanResponse, Status>> + Send + 'static>>;
+ type ReattachExecuteStream = Self::ExecutePlanStream;
+
+ async fn config(
+ &self,
+ req: Request<pb::ConfigRequest>,
+ ) -> Result<Response<pb::ConfigResponse>, Status> {
+ let body = req.into_inner();
+ Ok(Response::new(pb::ConfigResponse {
+ session_id: body.session_id,
+ ..Default::default()
+ }))
+ }
+
+ // Unused but trait-required.
+ async fn analyze_plan(
+ &self,
+ _: Request<pb::AnalyzePlanRequest>,
+ ) -> Result<Response<pb::AnalyzePlanResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn artifact_status(
+ &self,
+ _: Request<pb::ArtifactStatusesRequest>,
+ ) -> Result<Response<pb::ArtifactStatusesResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn interrupt(
+ &self,
+ _: Request<pb::InterruptRequest>,
+ ) -> Result<Response<pb::InterruptResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn release_execute(
+ &self,
+ _: Request<pb::ReleaseExecuteRequest>,
+ ) -> Result<Response<pb::ReleaseExecuteResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn release_session(
+ &self,
+ _: Request<pb::ReleaseSessionRequest>,
+ ) -> Result<Response<pb::ReleaseSessionResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn fetch_error_details(
+ &self,
+ _: Request<pb::FetchErrorDetailsRequest>,
+ ) -> Result<Response<pb::FetchErrorDetailsResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn clone_session(
+ &self,
+ _: Request<pb::CloneSessionRequest>,
+ ) -> Result<Response<pb::CloneSessionResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn get_status(
+ &self,
+ _: Request<pb::GetStatusRequest>,
+ ) -> Result<Response<pb::GetStatusResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn execute_plan(
+ &self,
+ _: Request<pb::ExecutePlanRequest>,
+ ) -> Result<Response<Self::ExecutePlanStream>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn reattach_execute(
+ &self,
+ _: Request<pb::ReattachExecuteRequest>,
+ ) -> Result<Response<Self::ReattachExecuteStream>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+ async fn add_artifacts(
+ &self,
+ _: Request<tonic::Streaming<pb::AddArtifactsRequest>>,
+ ) -> Result<Response<pb::AddArtifactsResponse>, Status> {
+ Err(Status::unimplemented("n/a"))
+ }
+}
+
+async fn rig(
+ auth: AuthInterceptor,
+ resolver: TenantResolver,
+) -> (
+ Channel,
+ Arc<Mutex<Vec<SessionKey>>>,
+ tokio::sync::oneshot::Sender<()>,
+ tokio::sync::oneshot::Sender<()>,
+) {
+ // Backend
+ let be_lis = TcpListener::bind("127.0.0.1:0").await.unwrap();
+ let be_addr = be_lis.local_addr().unwrap().to_string();
+ let (be_tx, be_rx) = tokio::sync::oneshot::channel();
+ tokio::spawn(async move {
+ Server::builder()
+ .add_service(
+ pb::spark_connect_service_server::SparkConnectServiceServer::new(EchoBackend),
+ )
+ .serve_with_incoming_shutdown(TcpListenerStream::new(be_lis), async {
+ let _ = be_rx.await;
+ })
+ .await
+ .ok();
+ });
+
+ let metrics = Metrics::new().unwrap();
+ let pool: Arc<dyn Pool> = Arc::new(StaticPool::new(vec![be_addr]).unwrap());
+ let (recording, binds) = RecordingStore::new();
+ let store: Arc<dyn AffinityStore> = Arc::new(recording);
+ let router = Arc::new(Router::single_pool(pool, store));
+ let dialer = Dialer::new();
+ let proxy = SparkConnectProxy::with_components(router, dialer, auth, metrics, resolver);
+
+ let gw_lis = TcpListener::bind("127.0.0.1:0").await.unwrap();
+ let gw_addr = gw_lis.local_addr().unwrap();
+ let (gw_tx, gw_rx) = tokio::sync::oneshot::channel();
+ tokio::spawn(async move {
+ Server::builder()
+ .add_service(pb::spark_connect_service_server::SparkConnectServiceServer::new(proxy))
+ .serve_with_incoming_shutdown(TcpListenerStream::new(gw_lis), async {
+ let _ = gw_rx.await;
+ })
+ .await
+ .ok();
+ });
+
+ let endpoint = Endpoint::from_shared(format!("http://{}", gw_addr)).unwrap();
+ let ch = endpoint
+ .connect_timeout(Duration::from_secs(2))
+ .connect()
+ .await
+ .unwrap();
+ (ch, binds, be_tx, gw_tx)
+}
+
+fn anon_auth() -> AuthInterceptor {
+ AuthInterceptor::new(Arc::new(AnonymousAuthenticator))
+}
+
+fn token_auth_with_tenant(tenant: Option<&str>) -> AuthInterceptor {
+ let mut entry = TokenEntry {
+ token: "alice-secret".into(),
+ user_id: "alice".into(),
+ tenant: None,
+ groups: vec![],
+ };
+ if let Some(t) = tenant {
+ entry.tenant = Some(t.into());
+ }
+ let inner = StaticTokenAuthenticator::new(vec![entry]).unwrap();
+ AuthInterceptor::new(Arc::new(inner))
+}
+
+fn bearer_req<T>(body: T, tok: &str) -> Request<T> {
+ let mut req = Request::new(body);
+ req.metadata_mut().insert(
+ "authorization",
+ MetadataValue::try_from(format!("Bearer {}", tok)).unwrap(),
+ );
+ req
+}
+
+#[tokio::test]
+async fn from_claim_default_falls_back_to_default_tenant() {
+ // anon auth → Identity.tenant is None → resolver default = "default"
+ let resolver = TenantResolver::new(TenantResolverConfig::default());
+ let (ch, binds, _be, _gw) = rig(anon_auth(), resolver).await;
+ let mut c = pb::spark_connect_service_client::SparkConnectServiceClient::new(ch);
+ c.config(Request::new(pb::ConfigRequest {
+ session_id: "s1".into(),
+ ..Default::default()
+ }))
+ .await
+ .expect("RPC ok");
+ let captured = binds.lock().clone();
+ assert_eq!(captured.len(), 1);
+ assert_eq!(captured[0].tenant, "default");
+ assert_eq!(captured[0].user_id, "anonymous");
+ assert_eq!(captured[0].session_id, "s1");
+}
+
+#[tokio::test]
+async fn from_claim_uses_token_tenant_when_present() {
+ let auth = token_auth_with_tenant(Some("team-a"));
+ let resolver = TenantResolver::new(TenantResolverConfig::default());
+ let (ch, binds, _be, _gw) = rig(auth, resolver).await;
+ let mut c = pb::spark_connect_service_client::SparkConnectServiceClient::new(ch);
+ c.config(bearer_req(
+ pb::ConfigRequest {
+ session_id: "s1".into(),
+ ..Default::default()
+ },
+ "alice-secret",
+ ))
+ .await
+ .expect("RPC ok");
+ let captured = binds.lock().clone();
+ assert_eq!(captured[0].tenant, "team-a");
+}
+
+#[tokio::test]
+async fn from_claim_reject_without_tenant_returns_unauthenticated() {
+ // Token has no tenant → on_missing=Reject → Unauthenticated.
+ let auth = token_auth_with_tenant(None);
+ let resolver = TenantResolver::new(TenantResolverConfig {
+ source: TenantSource::FromClaim,
+ on_missing: OnMissing::Reject,
+ default_name: "default".into(),
+ });
+ let (ch, binds, _be, _gw) = rig(auth, resolver).await;
+ let mut c = pb::spark_connect_service_client::SparkConnectServiceClient::new(ch);
+ let err = c
+ .config(bearer_req(
+ pb::ConfigRequest {
+ session_id: "s1".into(),
+ ..Default::default()
+ },
+ "alice-secret",
+ ))
+ .await
+ .unwrap_err();
+ assert_eq!(err.code(), tonic::Code::Unauthenticated);
+ // Backend never gets called → no bind recorded.
+ assert!(binds.lock().is_empty());
+}
+
+#[tokio::test]
+async fn from_metadata_reads_x_tenant_header() {
+ let resolver = TenantResolver::new(TenantResolverConfig {
+ source: TenantSource::FromMetadata {
+ header: "x-tenant".into(),
+ },
+ on_missing: OnMissing::Reject,
+ default_name: "default".into(),
+ });
+ let (ch, binds, _be, _gw) = rig(anon_auth(), resolver).await;
+ let mut c = pb::spark_connect_service_client::SparkConnectServiceClient::new(ch);
+ let mut req = Request::new(pb::ConfigRequest {
+ session_id: "s1".into(),
+ ..Default::default()
+ });
+ req.metadata_mut()
+ .insert("x-tenant", MetadataValue::try_from("team-b").unwrap());
+ c.config(req).await.expect("RPC ok");
+ let captured = binds.lock().clone();
+ assert_eq!(captured[0].tenant, "team-b");
+}
+
+#[tokio::test]
+async fn always_default_ignores_claim_and_header() {
+ let auth = token_auth_with_tenant(Some("should-be-ignored"));
+ let resolver = TenantResolver::new(TenantResolverConfig {
+ source: TenantSource::AlwaysDefault,
+ on_missing: OnMissing::UseDefault,
+ default_name: "fixed-tenant".into(),
+ });
+ let (ch, binds, _be, _gw) = rig(auth, resolver).await;
+ let mut c = pb::spark_connect_service_client::SparkConnectServiceClient::new(ch);
+ let mut req = bearer_req(
+ pb::ConfigRequest {
+ session_id: "s1".into(),
+ ..Default::default()
+ },
+ "alice-secret",
+ );
+ req.metadata_mut()
+ .insert("x-tenant", MetadataValue::try_from("also-ignored").unwrap());
+ c.config(req).await.expect("RPC ok");
+ let captured = binds.lock().clone();
+ assert_eq!(captured[0].tenant, "fixed-tenant");
+}
+
+#[tokio::test]
+async fn two_tenants_with_same_session_id_are_isolated() {
+ // Two clients with different tenants both use session_id="s1".
+ // The store should see two distinct keys, not one.
+ let resolver = TenantResolver::new(TenantResolverConfig {
+ source: TenantSource::FromMetadata {
+ header: "x-tenant".into(),
+ },
+ on_missing: OnMissing::UseDefault,
+ default_name: "default".into(),
+ });
+ let (ch, binds, _be, _gw) = rig(anon_auth(), resolver).await;
+ let mut c = pb::spark_connect_service_client::SparkConnectServiceClient::new(ch);
+
+ for tenant in ["team-a", "team-b"] {
+ let mut req = Request::new(pb::ConfigRequest {
+ session_id: "s1".into(),
+ ..Default::default()
+ });
+ req.metadata_mut()
+ .insert("x-tenant", MetadataValue::try_from(tenant).unwrap());
+ c.config(req).await.expect("RPC ok");
+ }
+
+ let captured = binds.lock().clone();
+ assert_eq!(captured.len(), 2, "two distinct binds");
+ let tenants: Vec<&str> = captured.iter().map(|k| k.tenant.as_str()).collect();
+ assert!(tenants.contains(&"team-a"));
+ assert!(tenants.contains(&"team-b"));
+ // Both have session_id="s1" but the keys differ because tenants differ.
+ assert_ne!(captured[0], captured[1]);
+}
diff --git a/crates/proxy/tests/tracing_integration.rs b/crates/proxy/tests/tracing_integration.rs
new file mode 100644
index 0000000..a1b0821
--- /dev/null
+++ b/crates/proxy/tests/tracing_integration.rs
@@ -0,0 +1,376 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Integration test: gateway emits OTel spans, propagates W3C
+//! traceparent to backend, and adopts an inbound traceparent as the
+//! parent of its own span.
+//!
+//! Strategy: install a single global tracing subscriber feeding an
+//! `InMemorySpanExporter`. Each scenario clears the buffer and
+//! inspects the spans produced by exactly one RPC.
+//!
+//! All scenarios live in one `#[tokio::test]` because `tracing` and
+//! `opentelemetry::global` are process-wide; multiple tests in the
+//! same binary would race.
+//!
+//! ## Why this test is `#[ignore]` by default
+//!
+//! Cargo unifies feature flags across the workspace's dependency
+//! graph, which means `cargo test --workspace` can pull additional
+//! features into `opentelemetry` (e.g. `spec_unstable_logs_enabled`,
+//! activated transitively when `kube` is in scope) that don't show
+//! up under `cargo test -p scg-proxy --test tracing_integration`.
+//! The resulting binary links a slightly different
+//! `tracing-opentelemetry`, and on this combination the layer
+//! silently drops application spans — only h2 internal spans reach
+//! the in-memory exporter. The production gateway is unaffected
+//! (it always builds the same way), but the test must opt out of
+//! workspace-wide runs to stay reliable. Run it explicitly with:
+//!
+//! ```bash
+//! cargo test -p scg-proxy --test tracing_integration -- --ignored
+//! ```
+
+use std::pin::Pin;
+use std::sync::Arc;
+use std::time::Duration;
+
+use futures::Stream;
+use opentelemetry_sdk::trace::{InMemorySpanExporter, SdkTracerProvider, SimpleSpanProcessor};
+use scg_genproto::pb;
+use scg_observability::{install_test_subscriber, Metrics, TRACEPARENT_HEADER};
+use scg_pool_static::StaticPool;
+use scg_proxy::{Dialer, SparkConnectProxy};
+use scg_routing::{AffinityStore, Pool, Router};
+use scg_store_memory::MemoryStore;
+use tokio::net::TcpListener;
+use tokio_stream::wrappers::TcpListenerStream;
+use tonic::metadata::MetadataValue;
+use tonic::transport::{Channel, Endpoint, Server};
+use tonic::{Request, Response, Status};
+
+/// Backend that captures the inbound metadata so the test can assert
+/// that the gateway forwarded a `traceparent` header.
+#[derive(Clone, Default)]
+struct CapturingBackend {
+ captured: Arc<parking_lot::Mutex<Option<tonic::metadata::MetadataMap>>>,
+}
+
+#[tonic::async_trait]
+impl pb::spark_connect_service_server::SparkConnectService for CapturingBackend {
+ type ExecutePlanStream =
+ Pin<Box<dyn Stream<Item = Result<pb::ExecutePlanResponse, Status>> + Send + 'static>>;
+ type ReattachExecuteStream = Self::ExecutePlanStream;
+
+ async fn config(
+ &self,
+ req: Request<pb::ConfigRequest>,
+ ) -> Result<Response<pb::ConfigResponse>, Status> {
+ *self.captured.lock() = Some(req.metadata().clone());
+ let body = req.into_inner();
+ Ok(Response::new(pb::ConfigResponse {
+ session_id: body.session_id,
+ ..Default::default()
+ }))
+ }
+
+ // The remaining methods aren't exercised by these tests but the
+ // trait requires them.
+ async fn analyze_plan(
+ &self,
+ _: Request<pb::AnalyzePlanRequest>,
+ ) -> Result<Response<pb::AnalyzePlanResponse>, Status> {
+ Err(Status::unimplemented("not used"))
+ }
+ async fn artifact_status(
+ &self,
+ _: Request<pb::ArtifactStatusesRequest>,
+ ) -> Result<Response<pb::ArtifactStatusesResponse>, Status> {
+ Err(Status::unimplemented("not used"))
+ }
+ async fn interrupt(
+ &self,
+ _: Request<pb::InterruptRequest>,
+ ) -> Result<Response<pb::InterruptResponse>, Status> {
+ Err(Status::unimplemented("not used"))
+ }
+ async fn release_execute(
+ &self,
+ _: Request<pb::ReleaseExecuteRequest>,
+ ) -> Result<Response<pb::ReleaseExecuteResponse>, Status> {
+ Err(Status::unimplemented("not used"))
+ }
+ async fn release_session(
+ &self,
+ _: Request<pb::ReleaseSessionRequest>,
+ ) -> Result<Response<pb::ReleaseSessionResponse>, Status> {
+ Err(Status::unimplemented("not used"))
+ }
+ async fn fetch_error_details(
+ &self,
+ _: Request<pb::FetchErrorDetailsRequest>,
+ ) -> Result<Response<pb::FetchErrorDetailsResponse>, Status> {
+ Err(Status::unimplemented("not used"))
+ }
+ async fn clone_session(
+ &self,
+ _: Request<pb::CloneSessionRequest>,
+ ) -> Result<Response<pb::CloneSessionResponse>, Status> {
+ Err(Status::unimplemented("not used"))
+ }
+ async fn get_status(
+ &self,
+ _: Request<pb::GetStatusRequest>,
+ ) -> Result<Response<pb::GetStatusResponse>, Status> {
+ Err(Status::unimplemented("not used"))
+ }
+ async fn execute_plan(
+ &self,
+ _: Request<pb::ExecutePlanRequest>,
+ ) -> Result<Response<Self::ExecutePlanStream>, Status> {
+ Err(Status::unimplemented("not used"))
+ }
+ async fn reattach_execute(
+ &self,
+ _: Request<pb::ReattachExecuteRequest>,
+ ) -> Result<Response<Self::ReattachExecuteStream>, Status> {
+ Err(Status::unimplemented("not used"))
+ }
+ async fn add_artifacts(
+ &self,
+ _: Request<tonic::Streaming<pb::AddArtifactsRequest>>,
+ ) -> Result<Response<pb::AddArtifactsResponse>, Status> {
+ Err(Status::unimplemented("not used"))
+ }
+}
+
+/// Locate the gateway span produced for `rpc_method`. We tag handlers
+/// with a `rpc.method` attribute carrying the unprefixed RPC name.
+fn find_rpc_span<'a>(
+ spans: &'a [opentelemetry_sdk::trace::SpanData],
+ rpc_method: &str,
+) -> Option<&'a opentelemetry_sdk::trace::SpanData> {
+ spans.iter().find(|s| {
+ s.attributes.iter().any(|kv| {
+ kv.key.as_str() == "rpc_method"
+ && kv.value.as_str() == std::borrow::Cow::Borrowed(rpc_method)
+ })
+ })
+}
+
+/// Block until a span tagged `rpc.method=<rpc>` shows up in the
+/// exporter, or `deadline` elapses. The handler-side span closes
+/// shortly after tonic emits the response trailer, so the client's
+/// `.await` completing isn't a synchronization point. Returns the
+/// snapshot the assertion path will inspect.
+async fn wait_for_span(
+ exporter: &InMemorySpanExporter,
+ rpc_method: &str,
+ deadline: Duration,
+) -> Vec<opentelemetry_sdk::trace::SpanData> {
+ let start = std::time::Instant::now();
+ loop {
+ let spans = exporter.get_finished_spans().expect("spans accessible");
+ if find_rpc_span(&spans, rpc_method).is_some() {
+ return spans;
+ }
+ if start.elapsed() >= deadline {
+ return spans;
+ }
+ tokio::time::sleep(Duration::from_millis(20)).await;
+ }
+}
+
+struct Rig {
+ grpc: Channel,
+ captured: Arc<parking_lot::Mutex<Option<tonic::metadata::MetadataMap>>>,
+ _be_shutdown: tokio::sync::oneshot::Sender<()>,
+ _gw_shutdown: tokio::sync::oneshot::Sender<()>,
+}
+
+async fn rig() -> Rig {
+ let backend = CapturingBackend::default();
+ let captured = backend.captured.clone();
+ let svc = pb::spark_connect_service_server::SparkConnectServiceServer::new(backend);
+ let lis = TcpListener::bind("127.0.0.1:0").await.unwrap();
+ let be_addr = lis.local_addr().unwrap().to_string();
+ let (be_tx, be_rx) = tokio::sync::oneshot::channel();
+ tokio::spawn(async move {
+ Server::builder()
+ .add_service(svc)
+ .serve_with_incoming_shutdown(TcpListenerStream::new(lis), async {
+ let _ = be_rx.await;
+ })
+ .await
+ .ok();
+ });
+
+ let metrics = Metrics::new().unwrap();
+ let pool: Arc<dyn Pool> = Arc::new(StaticPool::new(vec![be_addr]).unwrap());
+ let store: Arc<dyn AffinityStore> = Arc::new(MemoryStore::new());
+ let router = Arc::new(Router::single_pool(pool, store));
+ let dialer = Dialer::new();
+ // Auth disabled — these tests focus on tracing propagation, not
+ // identity. The proxy still authenticates (Anonymous) and stamps
+ // a UserContext.
+ let proxy = SparkConnectProxy::with_auth_and_metrics(
+ router,
+ dialer,
+ scg_auth::AuthInterceptor::new(Arc::new(scg_auth::AnonymousAuthenticator)),
+ metrics,
+ );
+
+ let gw_lis = TcpListener::bind("127.0.0.1:0").await.unwrap();
+ let gw_addr = gw_lis.local_addr().unwrap();
+ let (gw_tx, gw_rx) = tokio::sync::oneshot::channel();
+ tokio::spawn(async move {
+ Server::builder()
+ .add_service(pb::spark_connect_service_server::SparkConnectServiceServer::new(proxy))
+ .serve_with_incoming_shutdown(TcpListenerStream::new(gw_lis), async {
+ let _ = gw_rx.await;
+ })
+ .await
+ .ok();
+ });
+
+ let endpoint = Endpoint::from_shared(format!("http://{}", gw_addr)).unwrap();
+ let grpc = endpoint
+ .connect_timeout(Duration::from_secs(2))
+ .connect()
+ .await
+ .unwrap();
+
+ Rig {
+ grpc,
+ captured,
+ _be_shutdown: be_tx,
+ _gw_shutdown: gw_tx,
+ }
+}
+
+#[tokio::test]
+#[ignore = "feature-skewed across workspace builds; run with -- --ignored under -p scg-proxy"]
+async fn tracing_emits_spans_and_propagates_traceparent() {
+ // Install one shared in-memory exporter for the whole test
+ // binary. SimpleSpanProcessor flushes synchronously on export, so
+ // we don't need to wait for batches.
+ let exporter = InMemorySpanExporter::default();
+ let provider = SdkTracerProvider::builder()
+ .with_span_processor(SimpleSpanProcessor::new(exporter.clone()))
+ .build();
+ install_test_subscriber(&provider);
+
+ let rig = rig().await;
+ let mut client =
+ pb::spark_connect_service_client::SparkConnectServiceClient::new(rig.grpc.clone());
+
+ // ---- Scenario 1: no inbound traceparent → fresh root span -----
+ exporter.reset();
+ *rig.captured.lock() = None;
+
+ client
+ .config(Request::new(pb::ConfigRequest {
+ session_id: "s-no-parent".into(),
+ ..Default::default()
+ }))
+ .await
+ .expect("Config RPC succeeds");
+
+ // Server-side span closes when the handler future drops. tonic
+ // drops it slightly after the response trailer is sent, so the
+ // client `.await` returning doesn't guarantee the span is in the
+ // exporter yet — give the runtime a beat.
+ let spans = wait_for_span(&exporter, "Config", Duration::from_secs(2)).await;
+ let config_span =
+ find_rpc_span(&spans, "Config").expect("a span tagged rpc_method=Config was emitted");
+ // No parent → SpanId of the parent is zero / unset.
+ assert!(
+ !config_span
+ .parent_span_id
+ .to_string()
+ .chars()
+ .any(|c| c != '0'),
+ "expected root span (no parent), got parent {:?}",
+ config_span.parent_span_id
+ );
+
+ let captured = rig.captured.lock().clone().expect("backend saw the call");
+ let tp = captured
+ .get(TRACEPARENT_HEADER)
+ .expect("gateway forwarded a traceparent")
+ .to_str()
+ .unwrap();
+ // Format: 00-<traceid>-<spanid>-<flags>
+ let parts: Vec<&str> = tp.split('-').collect();
+ assert_eq!(parts.len(), 4, "malformed traceparent: {tp}");
+ let outbound_trace_id = parts[1];
+ let outbound_span_id = parts[2];
+ // The outbound traceparent's trace id matches the gateway span's trace id.
+ assert_eq!(
+ outbound_trace_id,
+ config_span.span_context.trace_id().to_string()
+ );
+ // The outbound parent span (what the backend treats as its parent) is
+ // the gateway's own span — the gateway span's id, not its parent.
+ assert_eq!(
+ outbound_span_id,
+ config_span.span_context.span_id().to_string()
+ );
+
+ // ---- Scenario 2: inbound traceparent → adopted as parent ------
+ exporter.reset();
+ *rig.captured.lock() = None;
+
+ let inbound_trace = "0af7651916cd43dd8448eb211c80319c";
+ let inbound_parent = "b7ad6b7169203331";
+ let inbound_tp = format!("00-{}-{}-01", inbound_trace, inbound_parent);
+ let mut req = Request::new(pb::ConfigRequest {
+ session_id: "s-with-parent".into(),
+ ..Default::default()
+ });
+ req.metadata_mut().insert(
+ TRACEPARENT_HEADER,
+ MetadataValue::try_from(inbound_tp.clone()).unwrap(),
+ );
+ client.config(req).await.expect("Config RPC succeeds");
+
+ let spans = wait_for_span(&exporter, "Config", Duration::from_secs(2)).await;
+ let gw_span =
+ find_rpc_span(&spans, "Config").expect("Config span emitted under propagated parent");
+ // The gateway span's trace id must equal the inbound trace id.
+ assert_eq!(gw_span.span_context.trace_id().to_string(), inbound_trace);
+ // The gateway span's parent must be the inbound parent span id.
+ assert_eq!(gw_span.parent_span_id.to_string(), inbound_parent);
+
+ // The forwarded traceparent must keep the same trace id (so the
+ // backend joins the same distributed trace) but carry the
+ // gateway's own span id as the new parent.
+ let captured = rig.captured.lock().clone().expect("backend saw the call");
+ let tp = captured
+ .get(TRACEPARENT_HEADER)
+ .expect("gateway forwarded a traceparent")
+ .to_str()
+ .unwrap();
+ let parts: Vec<&str> = tp.split('-').collect();
+ assert_eq!(parts[1], inbound_trace, "trace id changed across hop");
+ assert_eq!(
+ parts[2],
+ gw_span.span_context.span_id().to_string(),
+ "outbound parent should be the gateway span id"
+ );
+
+ // Cleanup so the test process exits cleanly.
+ let _ = provider.shutdown();
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org