Files
OpenShell/crates/openshell-server/tests/supervisor_relay_integration.rs
Emilien Macchianddivesh 8b3cc3fdc0 feat(server): add gateway capacity metrics and optional HPA (#3978)
* feat(server): add gateway capacity metrics and optional HPA

Expose per-replica supervisor sessions, pending relay capacity, relay
rejections and claim latency, and outbound peer request outcomes and
latency. Use bounded labels and Prometheus histograms for the new latency
metrics while preserving existing summary metrics.

Add an optional Helm HPA with external-database and resource validation,
conservative scale-down defaults, and support for custom metrics. Keep
certificate hook pods outside gateway workload selectors.

Document per-pod scraping, scaling limits, upgrade behavior, and
PostgreSQL connection sizing.

Part of #3528

Signed-off-by: Emilien Macchi <emacchi@redhat.com>

* feat(server): unify routing metrics and clarify replica capacity

Combine local relay setup and outbound peer requests in one counter, labeled by operation, route, target, outcome, and gRPC status. Preserve peer latency metrics.

Rename the rejection reason from global_capacity to replica_capacity to reflect the per-replica relay budget. Keep capacity limits unchanged.

Update tests and documentation.

Signed-off-by: divesh <dgude@nvidia.com>

* fix(server): align routed request metric labels and outcomes

Count a local relay as successful only when its supervisor claims it, the
same event that answers a peer relay on the owner, and rename the outcomes
to local_error and remote_error so they say where an attempt failed. Label
the peer latency histogram by operation, like the routed request counter,
and rename the target label to relay_kind so it does not read as the
Prometheus scrape target.

Rename RelayCapacity.global to per_replica, drop the per-sandbox relay
capacity gauge, which no per-sandbox series can pair with, describe the
latency histogram as peer-only, and restore the note that unavailable
spikes are expected during rollouts.

Part of #3528

Signed-off-by: Emilien Macchi <emacchi@redhat.com>

---------

Signed-off-by: Emilien Macchi <emacchi@redhat.com>
Signed-off-by: divesh <dgude@nvidia.com>
Co-authored-by: divesh <dgude@nvidia.com>
2026-10-05 15:38:40 +00:00

1069 lines
41 KiB
Rust

// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
//! Integration tests for the supervisor relay gRPC path.
//!
//! Stands up an in-process tonic server hosting the real `handle_relay_stream`
//! handler, plus a mock "supervisor" client that calls `relay_stream` over a
//! real `Channel`. Exercises the wire contract (typed `RelayFrame { Init | Data }`),
//! `SupervisorSessionRegistry::open_relay` → `claim_relay` pairing, and the
//! bidirectional byte bridge inside the handler.
//!
//! These tests complement the unit tests in `supervisor_session.rs` (which
//! exercise registry state only) and the live cluster tests (which exercise
//! the full CLI → gateway → sandbox path). They catch regressions in the gRPC
//! wire layer that unit tests can't see and that are expensive to catch in
//! E2E.
use std::sync::Arc;
use std::time::Duration;
use hyper_util::{
rt::{TokioExecutor, TokioIo},
server::conn::auto::Builder,
};
use metrics::LocalRecorderGuard;
use metrics_exporter_prometheus::{PrometheusBuilder, PrometheusHandle, PrometheusRecorder};
use openshell_core::proto::{
GatewayMessage, PeerRelayFrame, RelayFrame, RelayInit, SupervisorMessage, TcpForwardFrame,
open_shell_client::OpenShellClient,
open_shell_server::{OpenShell, OpenShellServer},
};
use openshell_server::supervisor_session::SupervisorSessionRegistry;
use openshell_server::{MultiplexedService, Store, health_router};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;
use tokio::sync::mpsc;
use tokio_stream::StreamExt;
use tokio_stream::wrappers::ReceiverStream;
use tonic::transport::{Channel, Endpoint};
use tonic::{Response, Status};
// ---------------------------------------------------------------------------
// Gateway service: only relay_stream does real work; everything else stubs.
// ---------------------------------------------------------------------------
#[derive(Clone)]
struct RelayGateway {
registry: Arc<SupervisorSessionRegistry>,
}
#[tonic::async_trait]
impl OpenShell for RelayGateway {
async fn peer_report_provider_readiness(
&self,
_request: tonic::Request<openshell_core::proto::ReportProviderReadinessRequest>,
) -> Result<Response<openshell_core::proto::ReportProviderReadinessResponse>, Status> {
Err(Status::unimplemented("not used by this test server"))
}
async fn peer_report_endpoint_status(
&self,
_request: tonic::Request<openshell_core::proto::ReportEndpointStatusRequest>,
) -> Result<Response<openshell_core::proto::ReportEndpointStatusResponse>, Status> {
Err(Status::unimplemented("not used by this test server"))
}
async fn peer_get_sandbox_provider_status(
&self,
_request: tonic::Request<openshell_core::proto::GetSandboxProviderStatusRequest>,
) -> Result<Response<openshell_core::proto::GetSandboxProviderStatusResponse>, Status> {
Err(Status::unimplemented("not used by this test server"))
}
async fn report_endpoint_status(
&self,
_request: tonic::Request<openshell_core::proto::ReportEndpointStatusRequest>,
) -> Result<Response<openshell_core::proto::ReportEndpointStatusResponse>, Status> {
Ok(Response::new(
openshell_core::proto::ReportEndpointStatusResponse {},
))
}
async fn begin_rootfs_tar_staging(
&self,
_request: tonic::Request<openshell_core::proto::BeginRootfsTarStagingRequest>,
) -> Result<Response<openshell_core::proto::BeginRootfsTarStagingResponse>, Status> {
Err(Status::unimplemented("not used by this test server"))
}
async fn report_main_process_exit(
&self,
_request: tonic::Request<openshell_core::proto::ReportMainProcessExitRequest>,
) -> Result<Response<openshell_core::proto::ReportMainProcessExitResponse>, Status> {
Err(Status::unimplemented("not used by this test server"))
}
async fn finalize_main_process_exit(
&self,
_request: tonic::Request<openshell_core::proto::FinalizeMainProcessExitRequest>,
) -> Result<Response<openshell_core::proto::FinalizeMainProcessExitResponse>, Status> {
Err(Status::unimplemented("not used by this test server"))
}
async fn get_current_user(
&self,
_request: tonic::Request<openshell_core::proto::GetCurrentUserRequest>,
) -> Result<Response<openshell_core::proto::GetCurrentUserResponse>, Status> {
Err(Status::unimplemented("not used by this test server"))
}
type RelayStreamStream = std::pin::Pin<
Box<dyn tokio_stream::Stream<Item = Result<RelayFrame, Status>> + Send + 'static>,
>;
async fn relay_stream(
&self,
request: tonic::Request<tonic::Streaming<RelayFrame>>,
) -> Result<Response<Self::RelayStreamStream>, Status> {
openshell_server::supervisor_session::handle_relay_stream(&self.registry, request).await
}
// ------ unused stubs ------
async fn create_sandbox_template(
&self,
_request: tonic::Request<openshell_core::proto::CreateSandboxTemplateRequest>,
) -> Result<Response<openshell_core::proto::SandboxTemplateResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn get_sandbox_template(
&self,
_request: tonic::Request<openshell_core::proto::GetSandboxTemplateRequest>,
) -> Result<Response<openshell_core::proto::SandboxTemplateResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn list_sandbox_templates(
&self,
_request: tonic::Request<openshell_core::proto::ListSandboxTemplatesRequest>,
) -> Result<Response<openshell_core::proto::ListSandboxTemplatesResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn delete_sandbox_template(
&self,
_request: tonic::Request<openshell_core::proto::DeleteSandboxTemplateRequest>,
) -> Result<Response<openshell_core::proto::DeleteSandboxTemplateResponse>, Status> {
Err(Status::unimplemented("unused"))
}
type PeerRelayStream =
std::pin::Pin<Box<dyn tokio_stream::Stream<Item = Result<PeerRelayFrame, Status>> + Send>>;
async fn peer_relay(
&self,
_: tonic::Request<tonic::Streaming<PeerRelayFrame>>,
) -> Result<Response<Self::PeerRelayStream>, Status> {
Err(Status::unimplemented("unused"))
}
type ConnectSupervisorStream = ReceiverStream<Result<GatewayMessage, Status>>;
async fn connect_supervisor(
&self,
_: tonic::Request<tonic::Streaming<SupervisorMessage>>,
) -> Result<Response<Self::ConnectSupervisorStream>, Status> {
Err(Status::unimplemented("unused"))
}
type WatchSandboxStream =
ReceiverStream<Result<openshell_core::proto::SandboxStreamEvent, Status>>;
async fn watch_sandbox(
&self,
_: tonic::Request<openshell_core::proto::WatchSandboxRequest>,
) -> Result<Response<Self::WatchSandboxStream>, Status> {
Err(Status::unimplemented("unused"))
}
type ExecSandboxStream =
ReceiverStream<Result<openshell_core::proto::ExecSandboxEvent, Status>>;
async fn exec_sandbox(
&self,
_: tonic::Request<openshell_core::proto::ExecSandboxRequest>,
) -> Result<Response<Self::ExecSandboxStream>, Status> {
Err(Status::unimplemented("unused"))
}
type ForwardTcpStream =
std::pin::Pin<Box<dyn tokio_stream::Stream<Item = Result<TcpForwardFrame, Status>> + Send>>;
async fn forward_tcp(
&self,
_: tonic::Request<tonic::Streaming<TcpForwardFrame>>,
) -> Result<Response<Self::ForwardTcpStream>, Status> {
Err(Status::unimplemented("unused"))
}
type ExecSandboxInteractiveStream =
ReceiverStream<Result<openshell_core::proto::ExecSandboxEvent, Status>>;
async fn exec_sandbox_interactive(
&self,
_: tonic::Request<tonic::Streaming<openshell_core::proto::ExecSandboxInput>>,
) -> Result<Response<Self::ExecSandboxInteractiveStream>, Status> {
Err(Status::unimplemented("unused"))
}
async fn health(
&self,
_: tonic::Request<openshell_core::proto::HealthRequest>,
) -> Result<Response<openshell_core::proto::HealthResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn get_gateway_info(
&self,
_: tonic::Request<openshell_core::proto::GetGatewayInfoRequest>,
) -> Result<Response<openshell_core::proto::GetGatewayInfoResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn create_sandbox(
&self,
_: tonic::Request<openshell_core::proto::CreateSandboxRequest>,
) -> Result<Response<openshell_core::proto::SandboxResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn stop_sandbox(
&self,
_: tonic::Request<openshell_core::proto::StopSandboxRequest>,
) -> Result<Response<openshell_core::proto::SandboxResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn start_sandbox(
&self,
_: tonic::Request<openshell_core::proto::StartSandboxRequest>,
) -> Result<Response<openshell_core::proto::SandboxResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn get_sandbox(
&self,
_: tonic::Request<openshell_core::proto::GetSandboxRequest>,
) -> Result<Response<openshell_core::proto::SandboxResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn list_sandboxes(
&self,
_: tonic::Request<openshell_core::proto::ListSandboxesRequest>,
) -> Result<Response<openshell_core::proto::ListSandboxesResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn list_sandbox_providers(
&self,
_: tonic::Request<openshell_core::proto::ListSandboxProvidersRequest>,
) -> Result<Response<openshell_core::proto::ListSandboxProvidersResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn attach_sandbox_provider(
&self,
_: tonic::Request<openshell_core::proto::AttachSandboxProviderRequest>,
) -> Result<Response<openshell_core::proto::AttachSandboxProviderResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn detach_sandbox_provider(
&self,
_: tonic::Request<openshell_core::proto::DetachSandboxProviderRequest>,
) -> Result<Response<openshell_core::proto::DetachSandboxProviderResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn delete_sandbox(
&self,
_: tonic::Request<openshell_core::proto::DeleteSandboxRequest>,
) -> Result<Response<openshell_core::proto::DeleteSandboxResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn get_sandbox_config(
&self,
_: tonic::Request<openshell_core::proto::GetSandboxConfigRequest>,
) -> Result<Response<openshell_core::proto::GetSandboxConfigResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn get_gateway_config(
&self,
_: tonic::Request<openshell_core::proto::GetGatewayConfigRequest>,
) -> Result<Response<openshell_core::proto::GetGatewayConfigResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn get_sandbox_provider_status(
&self,
_request: tonic::Request<openshell_core::proto::GetSandboxProviderStatusRequest>,
) -> Result<Response<openshell_core::proto::GetSandboxProviderStatusResponse>, Status> {
Err(Status::unimplemented(
"provider readiness is not exercised by this mock",
))
}
async fn report_provider_readiness(
&self,
_request: tonic::Request<openshell_core::proto::ReportProviderReadinessRequest>,
) -> Result<Response<openshell_core::proto::ReportProviderReadinessResponse>, Status> {
Err(Status::unimplemented("provider readiness"))
}
async fn get_sandbox_provider_environment(
&self,
_: tonic::Request<openshell_core::proto::GetSandboxProviderEnvironmentRequest>,
) -> Result<Response<openshell_core::proto::GetSandboxProviderEnvironmentResponse>, Status>
{
Err(Status::unimplemented("unused"))
}
async fn create_ssh_session(
&self,
_: tonic::Request<openshell_core::proto::CreateSshSessionRequest>,
) -> Result<Response<openshell_core::proto::CreateSshSessionResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn expose_service(
&self,
_: tonic::Request<openshell_core::proto::ExposeServiceRequest>,
) -> Result<Response<openshell_core::proto::ServiceEndpointResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn get_service(
&self,
_: tonic::Request<openshell_core::proto::GetServiceRequest>,
) -> Result<Response<openshell_core::proto::ServiceEndpointResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn list_services(
&self,
_: tonic::Request<openshell_core::proto::ListServicesRequest>,
) -> Result<Response<openshell_core::proto::ListServicesResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn delete_service(
&self,
_: tonic::Request<openshell_core::proto::DeleteServiceRequest>,
) -> Result<Response<openshell_core::proto::DeleteServiceResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn revoke_ssh_session(
&self,
_: tonic::Request<openshell_core::proto::RevokeSshSessionRequest>,
) -> Result<Response<openshell_core::proto::RevokeSshSessionResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn exchange_provider_subject_token(
&self,
_: tonic::Request<openshell_core::proto::ExchangeProviderSubjectTokenRequest>,
) -> Result<Response<openshell_core::proto::ExchangeProviderSubjectTokenResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn create_provider(
&self,
_: tonic::Request<openshell_core::proto::CreateProviderRequest>,
) -> Result<Response<openshell_core::proto::ProviderResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn update_provider(
&self,
_: tonic::Request<openshell_core::proto::UpdateProviderRequest>,
) -> Result<Response<openshell_core::proto::ProviderResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn get_provider(
&self,
_: tonic::Request<openshell_core::proto::GetProviderRequest>,
) -> Result<Response<openshell_core::proto::ProviderResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn list_providers(
&self,
_: tonic::Request<openshell_core::proto::ListProvidersRequest>,
) -> Result<Response<openshell_core::proto::ListProvidersResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn list_provider_profiles(
&self,
_: tonic::Request<openshell_core::proto::ListProviderProfilesRequest>,
) -> Result<Response<openshell_core::proto::ListProviderProfilesResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn get_provider_profile(
&self,
_: tonic::Request<openshell_core::proto::GetProviderProfileRequest>,
) -> Result<Response<openshell_core::proto::ProviderProfileResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn import_provider_profiles(
&self,
_: tonic::Request<openshell_core::proto::ImportProviderProfilesRequest>,
) -> Result<Response<openshell_core::proto::ImportProviderProfilesResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn update_provider_profiles(
&self,
_: tonic::Request<openshell_core::proto::UpdateProviderProfilesRequest>,
) -> Result<Response<openshell_core::proto::UpdateProviderProfilesResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn lint_provider_profiles(
&self,
_: tonic::Request<openshell_core::proto::LintProviderProfilesRequest>,
) -> Result<Response<openshell_core::proto::LintProviderProfilesResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn delete_provider_profile(
&self,
_: tonic::Request<openshell_core::proto::DeleteProviderProfileRequest>,
) -> Result<Response<openshell_core::proto::DeleteProviderProfileResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn get_provider_refresh_status(
&self,
_: tonic::Request<openshell_core::proto::GetProviderRefreshStatusRequest>,
) -> Result<Response<openshell_core::proto::GetProviderRefreshStatusResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn configure_provider_refresh(
&self,
_: tonic::Request<openshell_core::proto::ConfigureProviderRefreshRequest>,
) -> Result<Response<openshell_core::proto::ConfigureProviderRefreshResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn rotate_provider_credential(
&self,
_: tonic::Request<openshell_core::proto::RotateProviderCredentialRequest>,
) -> Result<Response<openshell_core::proto::RotateProviderCredentialResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn delete_provider_refresh(
&self,
_: tonic::Request<openshell_core::proto::DeleteProviderRefreshRequest>,
) -> Result<Response<openshell_core::proto::DeleteProviderRefreshResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn delete_provider(
&self,
_: tonic::Request<openshell_core::proto::DeleteProviderRequest>,
) -> Result<Response<openshell_core::proto::DeleteProviderResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn update_config(
&self,
_: tonic::Request<openshell_core::proto::UpdateConfigRequest>,
) -> Result<Response<openshell_core::proto::UpdateConfigResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn get_sandbox_policy_status(
&self,
_: tonic::Request<openshell_core::proto::GetSandboxPolicyStatusRequest>,
) -> Result<Response<openshell_core::proto::GetSandboxPolicyStatusResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn list_sandbox_policies(
&self,
_: tonic::Request<openshell_core::proto::ListSandboxPoliciesRequest>,
) -> Result<Response<openshell_core::proto::ListSandboxPoliciesResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn report_sandbox_configuration(
&self,
_: tonic::Request<openshell_core::proto::ReportSandboxConfigurationRequest>,
) -> Result<Response<openshell_core::proto::ReportSandboxConfigurationResponse>, Status> {
Err(Status::unimplemented("not implemented in test"))
}
async fn report_policy_status(
&self,
_: tonic::Request<openshell_core::proto::ReportPolicyStatusRequest>,
) -> Result<Response<openshell_core::proto::ReportPolicyStatusResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn get_sandbox_logs(
&self,
_: tonic::Request<openshell_core::proto::GetSandboxLogsRequest>,
) -> Result<Response<openshell_core::proto::GetSandboxLogsResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn push_sandbox_logs(
&self,
_: tonic::Request<tonic::Streaming<openshell_core::proto::PushSandboxLogsRequest>>,
) -> Result<Response<openshell_core::proto::PushSandboxLogsResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn submit_policy_analysis(
&self,
_: tonic::Request<openshell_core::proto::SubmitPolicyAnalysisRequest>,
) -> Result<Response<openshell_core::proto::SubmitPolicyAnalysisResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn get_draft_policy(
&self,
_: tonic::Request<openshell_core::proto::GetDraftPolicyRequest>,
) -> Result<Response<openshell_core::proto::GetDraftPolicyResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn approve_draft_chunk(
&self,
_: tonic::Request<openshell_core::proto::ApproveDraftChunkRequest>,
) -> Result<Response<openshell_core::proto::ApproveDraftChunkResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn reject_draft_chunk(
&self,
_: tonic::Request<openshell_core::proto::RejectDraftChunkRequest>,
) -> Result<Response<openshell_core::proto::RejectDraftChunkResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn approve_all_draft_chunks(
&self,
_: tonic::Request<openshell_core::proto::ApproveAllDraftChunksRequest>,
) -> Result<Response<openshell_core::proto::ApproveAllDraftChunksResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn edit_draft_chunk(
&self,
_: tonic::Request<openshell_core::proto::EditDraftChunkRequest>,
) -> Result<Response<openshell_core::proto::EditDraftChunkResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn undo_draft_chunk(
&self,
_: tonic::Request<openshell_core::proto::UndoDraftChunkRequest>,
) -> Result<Response<openshell_core::proto::UndoDraftChunkResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn clear_draft_chunks(
&self,
_: tonic::Request<openshell_core::proto::ClearDraftChunksRequest>,
) -> Result<Response<openshell_core::proto::ClearDraftChunksResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn get_draft_history(
&self,
_: tonic::Request<openshell_core::proto::GetDraftHistoryRequest>,
) -> Result<Response<openshell_core::proto::GetDraftHistoryResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn issue_sandbox_token(
&self,
_: tonic::Request<openshell_core::proto::IssueSandboxTokenRequest>,
) -> Result<Response<openshell_core::proto::IssueSandboxTokenResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn refresh_sandbox_token(
&self,
_: tonic::Request<openshell_core::proto::RefreshSandboxTokenRequest>,
) -> Result<Response<openshell_core::proto::RefreshSandboxTokenResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn create_workspace(
&self,
_: tonic::Request<openshell_core::proto::CreateWorkspaceRequest>,
) -> Result<Response<openshell_core::proto::CreateWorkspaceResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn get_workspace(
&self,
_: tonic::Request<openshell_core::proto::GetWorkspaceRequest>,
) -> Result<Response<openshell_core::proto::GetWorkspaceResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn list_workspaces(
&self,
_: tonic::Request<openshell_core::proto::ListWorkspacesRequest>,
) -> Result<Response<openshell_core::proto::ListWorkspacesResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn delete_workspace(
&self,
_: tonic::Request<openshell_core::proto::DeleteWorkspaceRequest>,
) -> Result<Response<openshell_core::proto::DeleteWorkspaceResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn add_workspace_member(
&self,
_: tonic::Request<openshell_core::proto::AddWorkspaceMemberRequest>,
) -> Result<Response<openshell_core::proto::AddWorkspaceMemberResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn remove_workspace_member(
&self,
_: tonic::Request<openshell_core::proto::RemoveWorkspaceMemberRequest>,
) -> Result<Response<openshell_core::proto::RemoveWorkspaceMemberResponse>, Status> {
Err(Status::unimplemented("unused"))
}
async fn list_workspace_members(
&self,
_: tonic::Request<openshell_core::proto::ListWorkspaceMembersRequest>,
) -> Result<Response<openshell_core::proto::ListWorkspaceMembersResponse>, Status> {
Err(Status::unimplemented("unused"))
}
}
// ---------------------------------------------------------------------------
// Test harness
// ---------------------------------------------------------------------------
async fn spawn_gateway(registry: Arc<SupervisorSessionRegistry>) -> Channel {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
let grpc = OpenShellServer::new(RelayGateway { registry });
let service = MultiplexedService::new(grpc, health_router(test_health_store().await));
tokio::spawn(async move {
loop {
let Ok((stream, _)) = listener.accept().await else {
continue;
};
let svc = service.clone();
tokio::spawn(async move {
let _ = Builder::new(TokioExecutor::new())
.serve_connection(TokioIo::new(stream), svc)
.await;
});
}
});
Endpoint::from_shared(format!("http://{addr}"))
.unwrap()
.connect()
.await
.expect("client connect")
}
fn register_session(
registry: &SupervisorSessionRegistry,
sandbox_id: &str,
) -> mpsc::Receiver<GatewayMessage> {
register_session_with_capacity(registry, sandbox_id, 8)
}
fn register_session_with_capacity(
registry: &SupervisorSessionRegistry,
sandbox_id: &str,
capacity: usize,
) -> mpsc::Receiver<GatewayMessage> {
let (tx, rx) = mpsc::channel(capacity);
let (shutdown_tx, _shutdown_rx) = tokio::sync::oneshot::channel::<()>();
registry.register(
sandbox_id.to_string(),
"sess-1".to_string(),
tx,
shutdown_tx,
);
rx
}
/// Captures metrics recorded on this test's thread. `#[tokio::test]` is current-thread, so the
/// in-process gateway tasks record here too. Do not pass it into async helper fns (it is !Send).
struct MetricsCapture {
handle: PrometheusHandle,
_guard: LocalRecorderGuard<'static>,
}
impl MetricsCapture {
fn install() -> Self {
// Leaked (test only) so the guard can borrow the recorder for 'static.
let recorder: &'static PrometheusRecorder =
Box::leak(Box::new(PrometheusBuilder::new().build_recorder()));
let handle = recorder.handle();
let guard = metrics::set_default_local_recorder(recorder);
Self {
handle,
_guard: guard,
}
}
fn render(&self) -> String {
self.handle.render()
}
/// Integer value of one exact series, or `None` when it was never emitted.
fn value(&self, series: &str) -> Option<i64> {
self.render()
.lines()
.find_map(|line| line.strip_prefix(series)?.strip_prefix(' ')?.parse().ok())
}
}
/// Mock supervisor that opens a `RelayStream`, sends `Init`, then echoes every
/// data frame it receives. Returns when the gateway drops the stream or when
/// the supervisor's own outbound channel closes.
async fn run_echo_supervisor(channel: Channel, channel_id: String) {
let mut client = OpenShellClient::new(channel);
let (out_tx, out_rx) = mpsc::channel::<RelayFrame>(16);
let outbound = ReceiverStream::new(out_rx);
out_tx
.send(RelayFrame {
payload: Some(openshell_core::proto::relay_frame::Payload::Init(
RelayInit { channel_id },
)),
})
.await
.expect("send init");
let response = client
.relay_stream(outbound)
.await
.expect("relay_stream rpc");
let mut inbound = response.into_inner();
while let Some(msg) = inbound.next().await {
let Ok(frame) = msg else { break };
let Some(openshell_core::proto::relay_frame::Payload::Data(data)) = frame.payload else {
continue;
};
let echoed = RelayFrame {
payload: Some(openshell_core::proto::relay_frame::Payload::Data(data)),
};
if out_tx.send(echoed).await.is_err() {
break;
}
}
}
// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------
#[tokio::test]
async fn relay_round_trips_bytes() {
let registry = Arc::new(SupervisorSessionRegistry::new());
let channel = spawn_gateway(Arc::clone(&registry)).await;
let mut session_rx = register_session(&registry, "sbx");
let (channel_id, relay_rx) = registry
.open_relay("sbx", Duration::from_secs(2))
.await
.expect("open_relay");
let opened = match session_rx.recv().await.expect("RelayOpen").payload {
Some(openshell_core::proto::gateway_message::Payload::RelayOpen(r)) => r.channel_id,
other => panic!("expected RelayOpen, got {other:?}"),
};
assert_eq!(opened, channel_id);
tokio::spawn(run_echo_supervisor(channel, channel_id));
let relay = relay_rx.await.expect("relay result").expect("relay duplex");
let (mut read_half, mut write_half) = tokio::io::split(relay);
write_half.write_all(b"hello relay").await.expect("write");
write_half.flush().await.expect("flush");
let mut buf = [0u8; 11];
read_half.read_exact(&mut buf).await.expect("read echoed");
assert_eq!(&buf, b"hello relay");
}
#[tokio::test]
async fn relay_closes_cleanly_when_gateway_drops() {
let registry = Arc::new(SupervisorSessionRegistry::new());
let channel = spawn_gateway(Arc::clone(&registry)).await;
let mut session_rx = register_session(&registry, "sbx");
let (channel_id, relay_rx) = registry
.open_relay("sbx", Duration::from_secs(2))
.await
.expect("open_relay");
let _ = session_rx.recv().await.expect("RelayOpen");
let supervisor = tokio::spawn(run_echo_supervisor(channel, channel_id));
let relay = relay_rx.await.expect("relay result").expect("relay duplex");
drop(relay);
// The supervisor's inbound stream should terminate shortly after the
// gateway side drops — its echo loop exits and the task finishes.
tokio::time::timeout(Duration::from_secs(5), supervisor)
.await
.expect("supervisor should terminate after gateway drop")
.expect("supervisor task");
}
#[tokio::test]
async fn relay_sees_eof_when_supervisor_closes() {
let registry = Arc::new(SupervisorSessionRegistry::new());
let channel = spawn_gateway(Arc::clone(&registry)).await;
let mut session_rx = register_session(&registry, "sbx");
let (channel_id, relay_rx) = registry
.open_relay("sbx", Duration::from_secs(2))
.await
.expect("open_relay");
let _ = session_rx.recv().await.expect("RelayOpen");
// Supervisor sends init, then drops its outbound sender → gateway reader
// should see EOF.
let supervisor = {
let channel_id = channel_id.clone();
tokio::spawn(async move {
let mut client = OpenShellClient::new(channel);
let (out_tx, out_rx) = mpsc::channel::<RelayFrame>(4);
let outbound = ReceiverStream::new(out_rx);
out_tx
.send(RelayFrame {
payload: Some(openshell_core::proto::relay_frame::Payload::Init(
RelayInit { channel_id },
)),
})
.await
.expect("send init");
let _response = client.relay_stream(outbound).await.expect("rpc");
drop(out_tx);
tokio::time::sleep(Duration::from_millis(200)).await;
})
};
let relay = relay_rx.await.expect("relay result").expect("relay duplex");
let (mut read_half, _write_half) = tokio::io::split(relay);
let mut buf = [0u8; 16];
let n = tokio::time::timeout(Duration::from_secs(5), read_half.read(&mut buf))
.await
.expect("read should complete")
.expect("read ok");
assert_eq!(n, 0, "gateway-side read should see EOF");
supervisor.await.expect("supervisor task");
}
#[tokio::test]
async fn open_relay_times_out_when_no_session() {
let registry = Arc::new(SupervisorSessionRegistry::new());
let _channel = spawn_gateway(Arc::clone(&registry)).await;
let err = registry
.open_relay("missing", Duration::from_millis(100))
.await
.expect_err("should time out");
assert_eq!(err.code(), tonic::Code::Unavailable);
}
#[tokio::test]
async fn concurrent_relays_multiplex_independently() {
let registry = Arc::new(SupervisorSessionRegistry::new());
let channel = spawn_gateway(Arc::clone(&registry)).await;
let mut session_rx = register_session(&registry, "sbx");
let (id_a, rx_a) = registry
.open_relay("sbx", Duration::from_secs(2))
.await
.expect("open_relay a");
let _ = session_rx.recv().await.expect("RelayOpen a");
let (id_b, rx_b) = registry
.open_relay("sbx", Duration::from_secs(2))
.await
.expect("open_relay b");
let _ = session_rx.recv().await.expect("RelayOpen b");
assert_ne!(id_a, id_b);
tokio::spawn(run_echo_supervisor(channel.clone(), id_a));
tokio::spawn(run_echo_supervisor(channel, id_b));
let relay_a = rx_a.await.expect("relay a result").expect("relay a");
let relay_b = rx_b.await.expect("relay b result").expect("relay b");
let (mut ra, mut wa) = tokio::io::split(relay_a);
let (mut rb, mut wb) = tokio::io::split(relay_b);
wa.write_all(b"stream-A").await.unwrap();
wb.write_all(b"stream-B").await.unwrap();
wa.flush().await.unwrap();
wb.flush().await.unwrap();
let mut buf_a = [0u8; 8];
let mut buf_b = [0u8; 8];
ra.read_exact(&mut buf_a).await.unwrap();
rb.read_exact(&mut buf_b).await.unwrap();
assert_eq!(&buf_a, b"stream-A");
assert_eq!(&buf_b, b"stream-B");
}
/// Bursts more `open_relay` calls than the per-sandbox cap allows in parallel
/// and asserts the registry enforces the ceiling cleanly. A well-behaved
/// caller inside the cap still succeeds; overflow calls return `ResourceExhausted`
/// rather than racing the pending map into an inconsistent state.
#[tokio::test]
async fn open_relay_enforces_per_sandbox_cap_under_concurrent_burst() {
let metrics = MetricsCapture::install();
let registry = Arc::new(SupervisorSessionRegistry::new());
let _channel = spawn_gateway(Arc::clone(&registry)).await;
// Oversized mpsc so the session doesn't backpressure the burst — the cap,
// not the channel, is what we're testing.
let _session_rx = register_session_with_capacity(&registry, "sbx", 256);
// Fire 64 concurrent opens. Per-sandbox cap is 32, global cap is 256,
// so exactly 32 should succeed and 32 should be rejected with
// `ResourceExhausted` carrying the per-sandbox message.
let mut handles = Vec::with_capacity(64);
for _ in 0..64 {
let r = Arc::clone(&registry);
handles.push(tokio::spawn(async move {
r.open_relay("sbx", Duration::from_secs(1)).await
}));
}
let mut ok = 0usize;
let mut exhausted = 0usize;
for h in handles {
match h.await.expect("task joined") {
Ok(_pair) => ok += 1,
Err(status) if status.code() == tonic::Code::ResourceExhausted => {
assert!(
status.message().contains("per-sandbox relay limit"),
"expected per-sandbox error message, got: {}",
status.message()
);
exhausted += 1;
}
Err(other) => panic!("unexpected open_relay error: {other:?}"),
}
}
assert_eq!(ok, 32, "exactly per-sandbox cap should succeed");
assert_eq!(exhausted, 32, "overflow should be rejected, not dropped");
// The successful receivers were dropped, so their entries stay pending
// until they are claimed or reaped.
assert_eq!(metrics.value("openshell_server_relay_pending"), Some(32));
assert_eq!(
metrics.value("openshell_server_relay_rejected_total{reason=\"sandbox_capacity\"}"),
Some(32)
);
assert_eq!(
metrics
.value("openshell_server_relay_rejected_total{reason=\"replica_capacity\"}")
.unwrap_or(0),
0
);
assert_eq!(
metrics.value("openshell_server_supervisor_sessions"),
Some(1)
);
// A different sandbox still has headroom — the per-sandbox cap doesn't
// leak onto unrelated tenants.
let _other_rx = register_session_with_capacity(&registry, "sbx-other", 8);
registry
.open_relay("sbx-other", Duration::from_secs(1))
.await
.expect("other sandbox should not be affected by sbx cap");
assert_eq!(metrics.value("openshell_server_relay_pending"), Some(33));
assert_eq!(
metrics.value("openshell_server_supervisor_sessions"),
Some(2)
);
}
/// Bursts more `open_relay` calls than the replica cap allows, spread so that
/// no sandbox reaches its own cap, and asserts the replica ceiling and its
/// metrics.
#[tokio::test]
async fn open_relay_enforces_replica_cap_under_concurrent_burst() {
let metrics = MetricsCapture::install();
let registry = Arc::new(SupervisorSessionRegistry::new());
let _channel = spawn_gateway(Arc::clone(&registry)).await;
let sandbox_ids: Vec<String> = (0..9).map(|i| format!("sbx-{i}")).collect();
let _session_rxs: Vec<_> = sandbox_ids
.iter()
.map(|id| register_session_with_capacity(&registry, id, 64))
.collect();
// 9 x 32 = 288 opens against a replica cap of 256. No sandbox gets more
// than its cap of 32 attempts and the replica check runs first, so exactly
// 32 opens hit the replica cap.
let mut handles = Vec::with_capacity(288);
for id in &sandbox_ids {
for _ in 0..32 {
let r = Arc::clone(&registry);
let id = id.clone();
handles.push(tokio::spawn(async move {
r.open_relay(&id, Duration::from_secs(1)).await
}));
}
}
let mut ok = 0usize;
let mut exhausted = 0usize;
for h in handles {
match h.await.expect("task joined") {
Ok(_pair) => ok += 1,
Err(status) if status.code() == tonic::Code::ResourceExhausted => {
assert!(
status.message().contains("gateway relay capacity"),
"expected replica capacity error message, got: {}",
status.message()
);
exhausted += 1;
}
Err(other) => panic!("unexpected open_relay error: {other:?}"),
}
}
assert_eq!(ok, 256, "exactly the replica cap should succeed");
assert_eq!(exhausted, 32, "overflow should be rejected, not dropped");
assert_eq!(metrics.value("openshell_server_relay_pending"), Some(256));
assert_eq!(
metrics.value("openshell_server_relay_rejected_total{reason=\"replica_capacity\"}"),
Some(32)
);
assert_eq!(
metrics
.value("openshell_server_relay_rejected_total{reason=\"sandbox_capacity\"}")
.unwrap_or(0),
0
);
assert_eq!(
metrics.value("openshell_server_supervisor_sessions"),
Some(9)
);
let rendered = metrics.render();
for forbidden in ["sandbox_id=", "channel_id=", "sbx-0", "endpoint="] {
assert!(
!rendered.contains(forbidden),
"metric labels must not carry identifiers ({forbidden})"
);
}
}
#[tokio::test]
async fn relay_claim_releases_pending_slot_and_records_claim_latency() {
let metrics = MetricsCapture::install();
let registry = Arc::new(SupervisorSessionRegistry::new());
let channel = spawn_gateway(Arc::clone(&registry)).await;
let _session_rx = register_session(&registry, "sbx");
let (channel_id, relay_rx) = registry
.open_relay("sbx", Duration::from_secs(2))
.await
.expect("open_relay");
assert_eq!(metrics.value("openshell_server_relay_pending"), Some(1));
tokio::spawn(run_echo_supervisor(channel, channel_id));
let _relay = relay_rx.await.expect("relay result").expect("relay duplex");
// The claim releases the pending slot under the pending lock, before it
// wakes the waiter.
assert_eq!(metrics.value("openshell_server_relay_pending"), Some(0));
assert_eq!(
metrics.value("openshell_server_relay_claim_duration_seconds_count"),
Some(1)
);
}
/// Build an in-memory store sufficient for wiring `health_router` in tests
/// where the persistence layer itself is not under test.
async fn test_health_store() -> Arc<Store> {
Arc::new(
Store::connect("sqlite::memory:")
.await
.expect("connect in-memory sqlite store for tests"),
)
}