trending/OpenShell
1
0
Fork 0
mirror of https://github.com/NVIDIA/OpenShell.git synced 2026-11-08 09:04:26 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(config): schedule configuration delivery without drops

Browse Source 
Replace the bounded delivery queue with a synchronous scheduler that never
rejects admission. Every publication takes a sequence number and a key is
rebuilt only while its last build predates the publication, so repeated
publications coalesce. Fanouts are cursors over sorted recipients; direct and
fanout builds share build slots with a direct reserve and a fanout floor, and a
separate route lane handles owner lookups and peer hints.

Provider instance changes and credential refreshes publish a targeted provider
scope that fans out only to sandboxes attached to that provider, filtered by a
monotonic attachment cache. Add the PeerConfigProviderTarget hint scope and
remove the unreleased queue_full field.

Push-mode sessions get one latest-wins slot per component, drained by an
outbound stream that sends control frames first. Sessions register with the
scheduler only after setup succeeds and only while still current, and recipient
epochs fence stale tickets.

Signed-off-by: Piotr Mlocek <pmlocek@nvidia.com>
This commit is contained in:
Piotr Mlocek
2026-10-06 08:50:57 -07:00
parent f1e85038a0
commit b907685b3f
13 changed files with 3890 additions and 2230 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/openshell-server/src/config_delivery.rs
-1417
View File
File diff suppressed because it is too large Load Diff
crates/openshell-server/src/config_delivery/mod.rs
+1431
View File
File diff suppressed because it is too large Load Diff
crates/openshell-server/src/config_delivery/scheduler.rs
+1342
View File
File diff suppressed because it is too large Load Diff
crates/openshell-server/src/config_delivery/session_slots.rs
+178
View File
@@ -0,0 +1,178 @@
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
//! Latest-wins configuration slots for one supervisor session stream.
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};
use futures::task::AtomicWaker;
use openshell_core::proto::GatewayMessage;
use tokio::sync::mpsc;
use tonic::Status;
use super::ConfigComponentKind;
/// One unsent snapshot per component. A newer snapshot replaces an unsent one,
/// so a slow stream holds at most one message per component.
#[derive(Debug, Default)]
pub struct ConfigSlots {
pending: Mutex<PendingSlots>,
waker: AtomicWaker,
}
#[derive(Debug, Default)]
struct PendingSlots {
next_order: u64,
sandbox_config: Option<(u64, GatewayMessage)>,
provider_environment: Option<(u64, GatewayMessage)>,
}
impl ConfigSlots {
/// Returns true when an unsent snapshot was replaced.
pub(crate) fn replace(&self, component: ConfigComponentKind, message: GatewayMessage) -> bool {
let replaced = {
let mut pending = self.pending.lock().unwrap();
let order = pending.next_order;
pending.next_order += 1;
let slot = match component {
ConfigComponentKind::SandboxConfig => &mut pending.sandbox_config,
ConfigComponentKind::ProviderEnvironment => &mut pending.provider_environment,
};
slot.replace((order, message)).is_some()
};
self.waker.wake();
replaced
}
fn take_oldest(&self) -> Option<GatewayMessage> {
let mut pending = self.pending.lock().unwrap();
let sandbox_first = match (&pending.sandbox_config, &pending.provider_environment) {
(Some((sandbox, _)), Some((provider, _))) => sandbox < provider,
(Some(_), None) => true,
(None, Some(_)) => false,
(None, None) => return None,
};
let slot = if sandbox_first {
pending.sandbox_config.take()
} else {
pending.provider_environment.take()
};
slot.map(|(_, message)| message)
}
}
/// Session stream that sends queued control frames before configuration.
///
/// `SessionAccepted` is queued on the control channel before the session can
/// receive configuration, so it always stays first.
pub struct SessionOutbound {
control: mpsc::Receiver<GatewayMessage>,
config: Arc<ConfigSlots>,
}
impl SessionOutbound {
pub(crate) fn new(control: mpsc::Receiver<GatewayMessage>, config: Arc<ConfigSlots>) -> Self {
Self { control, config }
}
}
impl tokio_stream::Stream for SessionOutbound {
type Item = Result<GatewayMessage, Status>;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let this = self.get_mut();
match this.control.poll_recv(cx) {
Poll::Ready(Some(message)) => return Poll::Ready(Some(Ok(message))),
Poll::Ready(None) => return Poll::Ready(None),
Poll::Pending => {}
}
this.config.waker.register(cx.waker());
this.config
.take_oldest()
.map_or(Poll::Pending, |message| Poll::Ready(Some(Ok(message))))
}
}
#[cfg(test)]
mod tests {
use super::*;
use openshell_core::proto::{ConfigUpdate, RelayOpen, gateway_message};
use tokio_stream::StreamExt;
fn config(update_id: &str) -> GatewayMessage {
GatewayMessage {
payload: Some(gateway_message::Payload::ConfigUpdate(ConfigUpdate {
update_id: update_id.into(),
..Default::default()
})),
}
}
fn relay_open() -> GatewayMessage {
GatewayMessage {
payload: Some(gateway_message::Payload::RelayOpen(RelayOpen::default())),
}
}
fn update_id(message: GatewayMessage) -> String {
match message.payload {
Some(gateway_message::Payload::ConfigUpdate(update)) => update.update_id,
other => panic!("expected a configuration update, got {other:?}"),
}
}
#[tokio::test]
async fn control_frames_go_before_queued_configuration() {
let (tx, rx) = mpsc::channel(4);
let slots = Arc::new(ConfigSlots::default());
let mut outbound = SessionOutbound::new(rx, Arc::clone(&slots));
assert!(!slots.replace(ConfigComponentKind::SandboxConfig, config("config")));
tx.send(relay_open()).await.unwrap();
let first = outbound.next().await.unwrap().unwrap();
assert!(matches!(
first.payload,
Some(gateway_message::Payload::RelayOpen(_))
));
assert_eq!(update_id(outbound.next().await.unwrap().unwrap()), "config");
}
#[tokio::test]
async fn newer_snapshot_replaces_an_unsent_one() {
let (_tx, rx) = mpsc::channel(4);
let slots = Arc::new(ConfigSlots::default());
let mut outbound = SessionOutbound::new(rx, Arc::clone(&slots));
assert!(!slots.replace(ConfigComponentKind::SandboxConfig, config("old")));
assert!(!slots.replace(ConfigComponentKind::ProviderEnvironment, config("provider")));
assert!(slots.replace(ConfigComponentKind::SandboxConfig, config("new")));
assert_eq!(
update_id(outbound.next().await.unwrap().unwrap()),
"provider"
);
assert_eq!(update_id(outbound.next().await.unwrap().unwrap()), "new");
}
#[tokio::test]
async fn waiting_stream_wakes_for_configuration() {
let (_tx, rx) = mpsc::channel(4);
let slots = Arc::new(ConfigSlots::default());
let mut outbound = SessionOutbound::new(rx, Arc::clone(&slots));
let next = tokio::spawn(async move { outbound.next().await.unwrap().unwrap() });
tokio::task::yield_now().await;
slots.replace(ConfigComponentKind::ProviderEnvironment, config("late"));
assert_eq!(update_id(next.await.unwrap()), "late");
}
#[tokio::test]
async fn stream_ends_with_the_control_channel() {
let (tx, rx) = mpsc::channel(4);
let slots = Arc::new(ConfigSlots::default());
let mut outbound = SessionOutbound::new(rx, Arc::clone(&slots));
slots.replace(ConfigComponentKind::SandboxConfig, config("discarded"));
drop(tx);
assert!(outbound.next().await.is_none());
}
}
crates/openshell-server/src/grpc/mod.rs
+1 -1
View File
@@ -895,7 +895,7 @@ impl OpenShell for OpenShellService {
&self,
request: Request<PeerConfigUpdateHintRequest>,
) -> Result<Response<PeerConfigUpdateHintResponse>, Status> {
crate::config_delivery::handle_peer_config_update_hint(&self.state, request).await
crate::config_delivery::handle_peer_config_update_hint(&self.state, request)
}
}
crates/openshell-server/src/grpc/policy.rs
+6 -9
View File
@@ -12746,14 +12746,8 @@ mod tests {
sandbox.spec.as_mut().unwrap().policy = None;
state.store.put_message(&sandbox).await.unwrap();
let (tx, mut rx) = tokio::sync::mpsc::channel(4);
let (shutdown_tx, _shutdown_rx) = tokio::sync::oneshot::channel();
state.supervisor_sessions.register(
"sb-published-policy".to_string(),
"session-1".to_string(),
tx,
shutdown_tx,
);
let mut session =
crate::config_delivery::register_test_push_session(&state, &sandbox, "session-1");
let owner_index = crate::supervisor_owner::SupervisorOwnerIndex::new(
Arc::clone(&state.store),
crate::supervisor_owner::OWNER_TTL,
@@ -12790,7 +12784,10 @@ mod tests {
.expect("committed policy");
let snapshot = tokio::time::timeout(std::time::Duration::from_secs(2), async {
loop {
let message = rx.recv().await.expect("configuration channel closed");
let message = tokio_stream::StreamExt::next(&mut session.outbound)
.await
.expect("configuration stream closed")
.unwrap();
let Some(openshell_core::proto::gateway_message::Payload::ConfigUpdate(update)) =
message.payload
else {
crates/openshell-server/src/grpc/provider.rs
+22 -11
View File
@@ -2558,7 +2558,18 @@ async fn authorize_and_resolve_profile_workspace(
}
}
fn publish_provider_change(state: &Arc<ServerState>, workspace: &str) {
/// Publish a provider record change to the sandboxes that attach it.
fn publish_provider_change(state: &Arc<ServerState>, workspace: &str, provider_name: &str) {
crate::config_delivery::publish_provider_components(
state,
workspace,
provider_name,
crate::config_delivery::ConfigComponents::ALL,
);
}
/// A profile change affects every provider of its type in scope.
fn publish_provider_profile_change(state: &Arc<ServerState>, workspace: &str) {
if workspace.is_empty() {
crate::config_delivery::publish_all_connected(
state,
@@ -2642,7 +2653,7 @@ pub(super) async fn handle_create_provider(
LifecycleOperation::Create,
TelemetryOutcome::Success,
);
publish_provider_change(state, &workspace);
publish_provider_change(state, &workspace, provider.object_name());
Ok(Response::new(ProviderResponse {
provider: Some(provider),
..Default::default()
@@ -2912,7 +2923,7 @@ pub(super) async fn handle_import_provider_profiles(
stored.profile.unwrap_or_default(),
resource_version,
));
publish_provider_change(state, &workspace);
publish_provider_profile_change(state, &workspace);
}
Ok(Response::new(ImportProviderProfilesResponse {
@@ -3048,7 +3059,7 @@ pub(super) async fn handle_update_provider_profiles(
replay_facts.resource(&stored)?;
let resource_version = stored_profile_resource_version(&stored);
let profile = profile_response_payload(stored.profile.unwrap_or_default(), resource_version);
publish_provider_change(state, &workspace);
publish_provider_profile_change(state, &workspace);
Ok(Response::new(UpdateProviderProfilesResponse {
diagnostics: Vec::new(),
@@ -3143,7 +3154,7 @@ pub(super) async fn handle_delete_provider_profile(
.delete(StoredProviderProfile::object_type(), existing.object_id())
.await
.map_err(|e| Status::internal(format!("delete provider profile failed: {e}")))?;
publish_provider_change(state, &workspace);
publish_provider_profile_change(state, &workspace);
Ok(Response::new(DeleteProviderProfileResponse {
outcome: openshell_core::proto::DeletionOutcome::Completed.into(),
@@ -3996,7 +4007,7 @@ pub(super) async fn handle_update_provider(
LifecycleOperation::Update,
TelemetryOutcome::Success,
);
publish_provider_change(state, &workspace);
publish_provider_change(state, &workspace, provider.object_name());
Ok(Response::new(ProviderResponse {
provider: Some(provider),
target_receipts,
@@ -4964,10 +4975,10 @@ pub(super) async fn handle_configure_provider_refresh(
updated,
)
.await;
publish_provider_change(state, &workspace);
publish_provider_change(state, &workspace, provider_name);
result?;
} else {
publish_provider_change(state, &workspace);
publish_provider_change(state, &workspace, provider_name);
}
replay_facts.refresh(&state_record)?;
@@ -5013,7 +5024,7 @@ pub(super) async fn handle_rotate_provider_credential(
credential_key,
)
.await?;
publish_provider_change(state, &workspace);
publish_provider_change(state, &workspace, provider_name);
replay_facts.refresh(&refresh_state)?;
Ok(Response::new(RotateProviderCredentialResponse {
@@ -5126,7 +5137,7 @@ pub(super) async fn handle_delete_provider_refresh(
})
.map(|_| ())?;
}
publish_provider_change(state, &workspace);
publish_provider_change(state, &workspace, provider_name);
Ok(Response::new(DeleteProviderRefreshResponse {
outcome: openshell_core::proto::DeletionOutcome::Completed.into(),
@@ -5168,7 +5179,7 @@ pub(super) async fn handle_delete_provider(
outcome,
);
if deleted {
publish_provider_change(state, &workspace);
publish_provider_change(state, &workspace, &name);
}
Ok(Response::new(DeleteProviderResponse {
outcome: super::deletion_outcome(deleted, req.allow_missing, "provider")?,
crates/openshell-server/src/lib.rs
+5 -5
View File
@@ -304,8 +304,8 @@ pub struct ServerState {
/// Set once graceful gateway shutdown begins so stream handlers can
/// distinguish expected transport closes from runtime failures.
pub(crate) gateway_shutting_down: AtomicBool,
/// Per-sandbox scheduler for coalesced supervisor configuration delivery.
pub(crate) config_delivery_queue: config_delivery::ConfigDeliveryQueue,
/// Coalescing scheduler for supervisor configuration delivery.
pub(crate) config_delivery: config_delivery::ConfigDelivery,
/// Routing boundary for local or remote supervisor configuration delivery.
pub(crate) supervisor_config_router: Arc<dyn config_delivery::SupervisorConfigRouter>,
@@ -435,8 +435,8 @@ impl ServerState {
let supervisor_config_router: Arc<dyn config_delivery::SupervisorConfigRouter> = Arc::new(
config_delivery::LocalSupervisorConfigRouter::new(Arc::clone(&supervisor_sessions)),
);
let config_delivery_queue =
config_delivery::ConfigDeliveryQueue::for_db_connections(store.max_connections());
let config_delivery =
config_delivery::ConfigDelivery::for_db_connections(store.max_connections());
Self {
config,
store,
@@ -451,7 +451,7 @@ impl ServerState {
settings_mutex: tokio::sync::Mutex::new(()),
supervisor_sessions,
gateway_shutting_down: AtomicBool::new(false),
config_delivery_queue,
config_delivery,
supervisor_config_router,
replica_id,
peer_endpoint,
crates/openshell-server/src/provider_refresh.rs
+29 -12
View File
@@ -1925,11 +1925,12 @@ pub fn spawn_refresh_worker(state: std::sync::Arc<crate::ServerState>, interval:
))
.await
{
Ok(workspaces) => {
for workspace in workspaces {
crate::config_delivery::publish_workspace_components(
Ok(providers) => {
for (workspace, provider_name) in providers {
crate::config_delivery::publish_provider_components(
&state,
&workspace,
&provider_name,
crate::config_delivery::ConfigComponents::ALL,
);
}
@@ -1942,12 +1943,13 @@ pub fn spawn_refresh_worker(state: std::sync::Arc<crate::ServerState>, interval:
});
}
/// Returns the `(workspace, provider)` pairs whose provider records changed.
async fn run_refresh_worker_tick(
store: &Store,
credentials: Option<&crate::credentials::CredentialRuntime>,
compute: Option<&crate::compute::ComputeRuntime>,
) -> Result<std::collections::HashSet<String>, Status> {
let changed_workspaces = std::collections::HashSet::new();
) -> Result<std::collections::HashSet<(String, String)>, Status> {
let changed_providers = std::collections::HashSet::new();
let now_ms = current_time_ms();
let states = list_all_refresh_states(store).await?;
let watched_count = states.len();
@@ -1967,7 +1969,7 @@ async fn run_refresh_worker_tick(
.iter()
.any(|state| refresh_state_has_work(state, now_ms))
{
return Ok(changed_workspaces);
return Ok(changed_providers);
}
let span = tracing::info_span!(
"refresh",
@@ -1998,8 +2000,8 @@ async fn refresh_states(
compute: Option<&crate::compute::ComputeRuntime>,
states: Vec<StoredProviderCredentialRefreshState>,
now_ms: i64,
) -> std::collections::HashSet<String> {
let mut changed_workspaces = std::collections::HashSet::new();
) -> std::collections::HashSet<(String, String)> {
let mut changed_providers = std::collections::HashSet::new();
for state in states {
if state
.metadata
@@ -2030,7 +2032,7 @@ async fn refresh_states(
"failed to finalize tombstoned provider refresh; retrying on the next sweep"
);
} else {
changed_workspaces.insert(state.object_workspace().to_string());
changed_providers.insert(changed_provider(&state));
}
continue;
}
@@ -2115,10 +2117,17 @@ async fn refresh_states(
"provider credential refresh failed"
);
} else {
changed_workspaces.insert(state.object_workspace().to_string());
changed_providers.insert(changed_provider(&state));
}
}
changed_workspaces
changed_providers
}
fn changed_provider(state: &StoredProviderCredentialRefreshState) -> (String, String) {
(
state.object_workspace().to_string(),
state.provider_name.clone(),
)
}
#[cfg(test)]
@@ -3581,10 +3590,18 @@ mod tests {
put_refresh_state(&store, &state).await.unwrap();
assert_eq!(credentials.stored_credential_count(), Some(1));
Box::pin(run_refresh_worker_tick(&store, Some(&credentials), None))
let changed = Box::pin(run_refresh_worker_tick(&store, Some(&credentials), None))
.await
.unwrap();
assert_eq!(
changed,
std::collections::HashSet::from([(
"default".to_string(),
"tombstoned-refresh".to_string()
)]),
"only the refreshed provider's attached sandboxes need an update"
);
assert!(
get_refresh_state(
&store,
crates/openshell-server/src/storage_proto.rs
+2 -2
View File
@@ -130,7 +130,7 @@ mod tests {
// Service authorization also extends both schemas additively. Legacy
// payloads retain the safe Strip default.
const PUBLIC_RPC_SCHEMA_SHA256: &str =
"73d5454909634839151b6d47fe4b62793d5426a783517541e91972766cf81aef";
"95d4b5caeecb2fd482e2657ec264a5093f1180adee0f5a85b716bef3fdb48ad9";
const DURABLE_SCHEMA_SHA256: &str =
"38165d9d76f49fcfe98a12f241e032838a2376c1d1a87ea2796fd33b9b1a3541";
const PUBLIC_DURABLE_OVERLAP_SHA256: &str =
@@ -614,7 +614,7 @@ mod tests {
overlap_hash.as_str(),
),
(
(321, 28),
(322, 28),
(93, 21),
(81, 21),
PUBLIC_RPC_SCHEMA_SHA256,
crates/openshell-server/src/supervisor_session.rs
+164 -133
View File
@@ -1,7 +1,7 @@
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
use std::collections::HashMap;
use std::collections::{HashMap, HashSet};
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
@@ -33,12 +33,13 @@ use openshell_core::transport_errors::is_expected_transport_close_status;
use crate::ServerState;
use crate::auth::principal::Principal;
use crate::config_delivery::{
DeliveryDisposition, MAX_SUPERVISOR_CONFIG_MESSAGE_BYTES, SupervisorConfigMessage,
ConfigComponentKind, ConfigSlots, DeliveryDisposition, MAX_SUPERVISOR_CONFIG_MESSAGE_BYTES,
SessionOutbound, SupervisorConfigMessage,
};
#[cfg(test)]
use crate::config_delivery::{LocalSupervisorConfigRouter, SupervisorConfigRouter};
use crate::grpc::provider_readiness::ProviderReadinessEvidence;
use crate::persistence::ObjectId;
use crate::persistence::{ObjectId, ObjectWorkspace};
use crate::supervisor_owner::{OWNER_TTL, OwnerError, OwnerGuard, SupervisorOwnerIndex};
const HEARTBEAT_INTERVAL_SECS: u32 = 15;
@@ -283,7 +284,8 @@ struct LiveSession {
/// Uniquely identifies this session instance. Used by cleanup to avoid
/// removing a session that has since been superseded by a reconnect.
session_id: String,
config_push_capable: bool,
/// Present only for push-capable sessions on a gateway in push mode.
config_slots: Option<Arc<ConfigSlots>>,
tx: mpsc::Sender<GatewayMessage>,
config_sequences: ConfigSequences,
/// Fires when this session is superseded by a reconnect so the old session
@@ -422,16 +424,16 @@ impl SupervisorSessionRegistry {
tx: mpsc::Sender<GatewayMessage>,
shutdown: oneshot::Sender<()>,
) -> bool {
self.register_with_config_push(sandbox_id, session_id, tx, shutdown, true)
self.register_with_config_slots(sandbox_id, session_id, tx, shutdown, None)
}
pub fn register_with_config_push(
pub(crate) fn register_with_config_slots(
&self,
sandbox_id: String,
session_id: String,
tx: mpsc::Sender<GatewayMessage>,
shutdown: oneshot::Sender<()>,
config_push_capable: bool,
config_slots: Option<Arc<ConfigSlots>>,
) -> bool {
let mut sessions = self.sessions.lock().unwrap();
let previous = sessions.remove(&sandbox_id);
@@ -440,7 +442,7 @@ impl SupervisorSessionRegistry {
LiveSession {
sandbox_id,
session_id,
config_push_capable,
config_slots,
tx,
config_sequences: ConfigSequences::default(),
shutdown,
@@ -558,49 +560,36 @@ impl SupervisorSessionRegistry {
true
}
pub(crate) fn connected_sandbox_ids(&self) -> Vec<String> {
self.sessions
.lock()
.unwrap()
.iter()
.filter(|(_, session)| session.config_push_capable)
.map(|(sandbox_id, _)| sandbox_id.clone())
.collect()
}
pub(crate) fn is_config_push_capable(&self, sandbox_id: &str) -> bool {
self.sessions
.lock()
.unwrap()
.get(sandbox_id)
.is_some_and(|session| session.config_push_capable)
}
/// Hand a snapshot to the session it was built for. A replacement session
/// received its own bootstrap and must not receive an older build.
pub(crate) fn deliver_config(
&self,
sandbox_id: &str,
session_id: &str,
message: SupervisorConfigMessage,
) -> DeliveryDisposition {
let component_name = message.component_name();
let component = message.component();
let mut sessions = self.sessions.lock().unwrap();
let Some(session) = sessions.get_mut(sandbox_id) else {
let Some(session) = sessions
.get_mut(sandbox_id)
.filter(|session| session.session_id == session_id)
else {
return DeliveryDisposition::NoActiveSession;
};
if !session.config_push_capable {
let Some(slots) = session.config_slots.clone() else {
return DeliveryDisposition::UnsupportedSession;
}
let sequence = match &message {
SupervisorConfigMessage::SandboxConfig(_) => {
&mut session.config_sequences.sandbox_config
}
SupervisorConfigMessage::ProviderEnvironment(_) => {
};
let sequence = match component {
ConfigComponentKind::SandboxConfig => &mut session.config_sequences.sandbox_config,
ConfigComponentKind::ProviderEnvironment => {
&mut session.config_sequences.provider_environment
}
};
// Replacing an unsent snapshot leaves a gap in the sequence.
*sequence = sequence.saturating_add(1);
let component_sequence = *sequence;
let component = match message {
let snapshot = match message {
SupervisorConfigMessage::SandboxConfig(snapshot) => {
config_update::Component::SandboxConfig(*snapshot)
}
@@ -612,25 +601,19 @@ impl SupervisorSessionRegistry {
payload: Some(gateway_message::Payload::ConfigUpdate(ConfigUpdate {
update_id: Uuid::new_v4().to_string(),
component_sequence,
component: Some(component),
component: Some(snapshot),
})),
};
if gateway_message.encoded_len() > MAX_SUPERVISOR_CONFIG_MESSAGE_BYTES {
return DeliveryDisposition::PayloadTooLarge;
}
drop(sessions);
match session.tx.try_send(gateway_message) {
Ok(()) => DeliveryDisposition::Enqueued,
Err(mpsc::error::TrySendError::Full(_)) => {
warn!(
sandbox_id = %sandbox_id,
component = component_name,
"supervisor configuration queue is full"
);
DeliveryDisposition::QueueFull
}
Err(mpsc::error::TrySendError::Closed(_)) => DeliveryDisposition::SessionClosed,
if slots.replace(component, gateway_message) {
DeliveryDisposition::Replaced
} else {
DeliveryDisposition::Queued
}
}
@@ -1989,12 +1972,24 @@ pub async fn handle_connect_supervisor(
if let Some(principal) = principal.as_ref() {
crate::auth::guard::ensure_sandbox_principal_scope(principal, &sandbox_id)?;
}
// Captured before reading any bootstrap input, including the sandbox
// record, so registration detects publications the bootstrap may miss.
let captured_seq = (state.config.config_delivery_mode
== openshell_core::config::ConfigDeliveryMode::Push
&& hello.protocol_revision == SUPERVISOR_PROTOCOL_REVISION)
.then(|| state.config_delivery.current_seq());
let sandbox = require_persisted_sandbox(&state.store, &sandbox_id).await?;
let bootstrap = if state.config.config_delivery_mode
== openshell_core::config::ConfigDeliveryMode::Poll
|| hello.protocol_revision == LEGACY_SUPERVISOR_PROTOCOL_REVISION
{
let config_push = captured_seq.map(|captured_seq| ConfigPushSetup {
captured_seq,
workspace: sandbox.object_workspace().to_string(),
providers: sandbox
.spec
.as_ref()
.map(|spec| spec.providers.iter().cloned().collect())
.unwrap_or_default(),
});
let bootstrap = if config_push.is_none() {
None
} else {
match crate::config_delivery::build_config_bootstrap(state, &sandbox).await {
@@ -2034,6 +2029,7 @@ pub async fn handle_connect_supervisor(
inbound,
hello,
bootstrap,
config_push,
provider_readiness,
session_lifetime,
))
@@ -2041,11 +2037,18 @@ pub async fn handle_connect_supervisor(
.map_err(|error| Status::internal(format!("supervisor session setup failed: {error}")))?
}
struct ConfigPushSetup {
captured_seq: u64,
workspace: String,
providers: HashSet<String>,
}
async fn establish_supervisor_session(
state: Arc<ServerState>,
mut inbound: tonic::Streaming<SupervisorMessage>,
hello: SupervisorHello,
bootstrap: Option<openshell_core::proto::ConfigBootstrap>,
config_push: Option<ConfigPushSetup>,
provider_readiness: ProviderReadinessEvidence,
session_lifetime: OwnedRwLockReadGuard<()>,
) -> Result<
@@ -2109,6 +2112,16 @@ async fn establish_supervisor_session(
};
accepted_payload.bootstrap = None;
}
let bootstrap_built = matches!(
&accepted.payload,
Some(gateway_message::Payload::SessionAccepted(SessionAccepted {
bootstrap: Some(_),
..
}))
);
let config_slots = config_push
.as_ref()
.map(|_| Arc::new(ConfigSlots::default()));
if tx.send(accepted).await.is_err() {
if let Err(err) = owner_index.release_if_current(&owner_guard).await {
warn!(sandbox_id = %sandbox_id, session_id = %session_id, error = %err, "supervisor session: failed to release owner after accept queue failure");
@@ -2116,12 +2129,12 @@ async fn establish_supervisor_session(
return Err(Status::internal("failed to send session accepted"));
}
let superseded = state.supervisor_sessions.register_with_config_push(
let superseded = state.supervisor_sessions.register_with_config_slots(
sandbox_id.clone(),
session_id.clone(),
tx.clone(),
shutdown_tx,
hello.protocol_revision == SUPERVISOR_PROTOCOL_REVISION,
config_slots.clone(),
);
if superseded {
info!(
@@ -2210,6 +2223,20 @@ async fn establish_supervisor_session(
.replay_pending_relays(&sandbox_id, &session_id, &tx)
.await;
if let Some(setup) = config_push {
crate::config_delivery::register_session(
&state,
crate::config_delivery::Registration {
sandbox_id: sandbox_id.clone(),
session_id: session_id.clone(),
workspace: setup.workspace,
providers: setup.providers,
captured_seq: setup.captured_seq,
bootstrap_built,
},
);
}
// Step 4: Spawn the session loop that reads inbound messages.
let state_clone = Arc::clone(&state);
let sandbox_id_clone = sandbox_id.clone();
@@ -2229,6 +2256,8 @@ async fn establish_supervisor_session(
let terminal_finalized = state_clone
.supervisor_sessions
.remove_if_current(&sandbox_id_clone, &session_id);
// Every removal path ends this task, including supersede and disconnect.
crate::config_delivery::unregister_session(&state_clone, &sandbox_id_clone, &session_id);
// Release only this exact ownership record. A newer supervisor session
// may already have published replacement ownership on another replica.
let owner_index = SupervisorOwnerIndex::new(state_clone.store.clone(), OWNER_TTL);
@@ -2264,10 +2293,12 @@ async fn establish_supervisor_session(
});
// Return the outbound stream.
let stream = ReceiverStream::new(rx);
let stream: Pin<
Box<dyn tokio_stream::Stream<Item = Result<GatewayMessage, Status>> + Send + 'static>,
> = Box::pin(tokio_stream::StreamExt::map(stream, Ok));
> = match config_slots {
Some(slots) => Box::pin(SessionOutbound::new(rx, slots)),
None => Box::pin(tokio_stream::StreamExt::map(ReceiverStream::new(rx), Ok)),
};
Ok(Response::new(stream))
}
@@ -2699,7 +2730,7 @@ mod tests {
assert_eq!(status.code(), tonic::Code::FailedPrecondition);
assert!(status.message().contains("revision mismatch"));
assert!(state.supervisor_sessions.connected_sandbox_ids().is_empty());
assert!(!state.supervisor_sessions.has_session("sb-future"));
}
#[test]
@@ -2710,6 +2741,36 @@ mod tests {
assert!(error.message().contains("revision mismatch"));
}
fn register_push_session(
registry: &SupervisorSessionRegistry,
sandbox_id: &str,
session_id: &str,
) -> (mpsc::Sender<GatewayMessage>, SessionOutbound) {
let (tx, rx) = mpsc::channel(4);
let slots = Arc::new(ConfigSlots::default());
registry.register_with_config_slots(
sandbox_id.into(),
session_id.into(),
tx.clone(),
make_shutdown(),
Some(Arc::clone(&slots)),
);
(tx, SessionOutbound::new(rx, slots))
}
fn component_sequence(message: GatewayMessage) -> (bool, u64) {
match message.payload {
Some(gateway_message::Payload::ConfigUpdate(update)) => (
matches!(
update.component,
Some(config_update::Component::SandboxConfig(_))
),
update.component_sequence,
),
other => panic!("expected config update, got {other:?}"),
}
}
#[tokio::test]
async fn config_router_reports_missing_session() {
let router = LocalSupervisorConfigRouter::new(Arc::new(SupervisorSessionRegistry::new()));
@@ -2717,6 +2778,7 @@ mod tests {
router
.deliver(
"missing",
"session",
SupervisorConfigMessage::SandboxConfig(Box::default()),
)
.await,
@@ -2725,136 +2787,102 @@ mod tests {
}
#[tokio::test]
async fn config_router_assigns_sequences_per_component() {
async fn config_router_keeps_the_latest_unsent_snapshot_per_component() {
use tokio_stream::StreamExt as _;
let registry = Arc::new(SupervisorSessionRegistry::new());
let router = LocalSupervisorConfigRouter::new(Arc::clone(&registry));
let (tx, mut rx) = mpsc::channel(4);
let (shutdown_tx, _shutdown_rx) = oneshot::channel();
registry.register("sb-1".into(), "session-1".into(), tx, shutdown_tx);
let (_tx, mut outbound) = register_push_session(&registry, "sb-1", "session-1");
for expected in [DeliveryDisposition::Queued, DeliveryDisposition::Replaced] {
assert_eq!(
router
.deliver(
"sb-1",
"session-1",
SupervisorConfigMessage::SandboxConfig(Box::default()),
)
.await,
expected
);
}
assert_eq!(
router
.deliver(
"sb-1",
SupervisorConfigMessage::SandboxConfig(Box::default()),
)
.await,
DeliveryDisposition::Enqueued
);
assert_eq!(
router
.deliver(
"sb-1",
SupervisorConfigMessage::SandboxConfig(Box::default()),
)
.await,
DeliveryDisposition::Enqueued
);
assert_eq!(
router
.deliver(
"sb-1",
"session-1",
SupervisorConfigMessage::ProviderEnvironment(
ProviderEnvironmentSnapshot::default(),
),
)
.await,
DeliveryDisposition::Enqueued
DeliveryDisposition::Queued
);
let first = rx.recv().await.expect("first config update");
let second = rx.recv().await.expect("second config update");
let third = rx.recv().await.expect("provider config update");
let sequence = |message: GatewayMessage| match message.payload {
Some(gateway_message::Payload::ConfigUpdate(update)) => update.component_sequence,
other => panic!("expected config update, got {other:?}"),
};
assert_eq!(sequence(first), 1);
assert_eq!(sequence(second), 2);
assert_eq!(sequence(third), 1);
let first = outbound.next().await.unwrap().unwrap();
let second = outbound.next().await.unwrap().unwrap();
assert_eq!(component_sequence(first), (true, 2));
assert_eq!(component_sequence(second), (false, 1));
}
#[tokio::test]
async fn config_router_uses_replacement_session() {
async fn config_router_never_crosses_into_a_replacement_session() {
use tokio_stream::StreamExt as _;
let registry = Arc::new(SupervisorSessionRegistry::new());
let router = LocalSupervisorConfigRouter::new(Arc::clone(&registry));
let (old_tx, mut old_rx) = mpsc::channel(2);
let (old_shutdown_tx, _old_shutdown_rx) = oneshot::channel();
registry.register("sb-1".into(), "old-session".into(), old_tx, old_shutdown_tx);
let (_old_tx, _old_outbound) = register_push_session(&registry, "sb-1", "old-session");
let (_new_tx, mut new_outbound) = register_push_session(&registry, "sb-1", "new-session");
assert_eq!(
router
.deliver(
"sb-1",
"old-session",
SupervisorConfigMessage::SandboxConfig(Box::default()),
)
.await,
DeliveryDisposition::Enqueued
DeliveryDisposition::NoActiveSession
);
let old_update = old_rx.recv().await.expect("old-session config update");
let Some(gateway_message::Payload::ConfigUpdate(old_update)) = old_update.payload else {
panic!("expected config update");
};
assert_eq!(old_update.component_sequence, 1);
let (new_tx, mut new_rx) = mpsc::channel(1);
let (new_shutdown_tx, _new_shutdown_rx) = oneshot::channel();
assert!(registry.register("sb-1".into(), "new-session".into(), new_tx, new_shutdown_tx,));
assert_eq!(
router
.deliver(
"sb-1",
"new-session",
SupervisorConfigMessage::SandboxConfig(Box::default()),
)
.await,
DeliveryDisposition::Enqueued
DeliveryDisposition::Queued
);
assert!(old_rx.try_recv().is_err());
let new_update = new_rx.try_recv().expect("new-session config update");
let Some(gateway_message::Payload::ConfigUpdate(new_update)) = new_update.payload else {
panic!("expected config update");
};
assert_eq!(new_update.component_sequence, 1);
let update = new_outbound.next().await.unwrap().unwrap();
assert_eq!(component_sequence(update), (true, 1));
}
#[tokio::test]
async fn config_router_reports_full_and_closed_queues() {
async fn config_router_skips_sessions_without_push() {
let registry = Arc::new(SupervisorSessionRegistry::new());
let router = LocalSupervisorConfigRouter::new(Arc::clone(&registry));
let (tx, rx) = mpsc::channel(1);
tx.try_send(GatewayMessage::default()).unwrap();
let (shutdown_tx, _shutdown_rx) = oneshot::channel();
registry.register("sb-1".into(), "session-1".into(), tx, shutdown_tx);
let (tx, _rx) = mpsc::channel(1);
registry.register("sb-1".into(), "session-1".into(), tx, make_shutdown());
assert_eq!(
router
.deliver(
"sb-1",
"session-1",
SupervisorConfigMessage::SandboxConfig(Box::default()),
)
.await,
DeliveryDisposition::QueueFull
);
drop(rx);
assert_eq!(
router
.deliver(
"sb-1",
SupervisorConfigMessage::SandboxConfig(Box::default()),
)
.await,
DeliveryDisposition::SessionClosed
DeliveryDisposition::UnsupportedSession
);
}
#[tokio::test]
async fn config_router_rejects_oversized_messages() {
use tokio_stream::StreamExt as _;
let registry = Arc::new(SupervisorSessionRegistry::new());
let router = LocalSupervisorConfigRouter::new(Arc::clone(&registry));
let (tx, mut rx) = mpsc::channel(1);
let (shutdown_tx, _shutdown_rx) = oneshot::channel();
registry.register("sb-1".into(), "session-1".into(), tx, shutdown_tx);
let (tx, mut outbound) = register_push_session(&registry, "sb-1", "session-1");
let snapshot = ProviderEnvironmentSnapshot {
values: vec![ProviderEnvironmentValue {
@@ -2868,12 +2896,15 @@ mod tests {
router
.deliver(
"sb-1",
"session-1",
SupervisorConfigMessage::ProviderEnvironment(snapshot),
)
.await,
DeliveryDisposition::PayloadTooLarge
);
assert!(rx.try_recv().is_err());
drop(tx);
registry.remove_if_current("sb-1", "session-1");
assert!(outbound.next().await.is_none());
}
use tokio::io::{AsyncReadExt, AsyncWriteExt};
proto/openshell.proto
+7 -2
View File
@@ -3225,11 +3225,18 @@ message PeerConfigSandboxTarget {
string session_id = 2;
}
// Sandboxes in a workspace that attach the named provider.
message PeerConfigProviderTarget {
string workspace = 1;
string name = 2;
}
message PeerConfigUpdateHintRequest {
oneof scope {
PeerConfigSandboxTarget sandbox = 1;
string workspace = 2;
bool all_connected = 3;
PeerConfigProviderTarget provider = 6;
}
bool sandbox_config = 4;
bool provider_environment = 5;
@@ -3238,8 +3245,6 @@ message PeerConfigUpdateHintRequest {
message PeerConfigUpdateHintResponse {
// The targeted session moved before the receiving replica accepted the hint.
bool stale_owner = 1;
// The bounded local delivery or fanout queue could not accept the hint.
bool queue_full = 2;
}
// Supervisor reports the result of a relay open request.
sdk/go/proto/openshellv1/openshell.pb.go
+703 -638
View File
File diff suppressed because it is too large Load Diff