//! Supervised, generation-fenced grid-session lifecycle. #![allow(clippy::missing_errors_doc)] use crate::types::{BoundedText, MAX_IDENTIFIER_BYTES}; use libremetaverse_types::compat::{CancellationToken, CancellationTokenSource}; use std::collections::{BTreeSet, VecDeque}; use std::error::Error; use std::fmt; use std::future::Future; use std::pin::Pin; use std::sync::atomic::{AtomicU8, AtomicU64, Ordering}; use std::sync::{Arc, Mutex}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use tokio::sync::{mpsc, oneshot, watch}; use tokio::task::JoinHandle; use tokio::time::Instant; const MAX_OFFLINE_WORK: usize = 1_024; const MAX_SEEN_WORK: usize = 4_096; const MAX_BACKOFF: Duration = Duration::from_hours(1); const MAX_STABLE_RESET: Duration = Duration::from_hours(24); const MAX_READINESS_TIMEOUT: Duration = Duration::from_mins(10); const MAX_SHUTDOWN: Duration = Duration::from_mins(1); pub type SessionFuture<'a, T> = Pin + Send + 'a>>; /// Generation-scoped gate between transport readiness and autonomous agent readiness. pub struct FrameworkReadiness { generation: watch::Sender, } impl Default for FrameworkReadiness { fn default() -> Self { let (generation, _) = watch::channel(0); Self { generation } } } impl FrameworkReadiness { pub fn mark_ready(&self, generation: u64) { if generation != 0 { self.generation.send_replace(generation); } } pub async fn wait_ready(&self, generation: u64, cancellation: CancellationToken) -> bool { let mut ready = self.generation.subscribe(); loop { if *ready.borrow() == generation { return true; } tokio::select! { () = cancellation.cancelled() => return false, changed = ready.changed() => if changed.is_err() { return false; }, } } } } #[derive(Clone, Copy, Debug, Eq, PartialEq)] #[repr(u8)] pub enum SessionState { Stopped = 0, Connecting = 1, Online = 2, Degraded = 3, Backoff = 4, AuthenticationBlocked = 5, Paused = 6, ShuttingDown = 7, } impl SessionState { #[must_use] pub const fn as_str(self) -> &'static str { match self { Self::Stopped => "stopped", Self::Connecting => "connecting", Self::Online => "online", Self::Degraded => "degraded", Self::Backoff => "backoff", Self::AuthenticationBlocked => "authentication-blocked", Self::Paused => "paused", Self::ShuttingDown => "shutting-down", } } const fn from_u8(value: u8) -> Self { match value { 1 => Self::Connecting, 2 => Self::Online, 3 => Self::Degraded, 4 => Self::Backoff, 5 => Self::AuthenticationBlocked, 6 => Self::Paused, 7 => Self::ShuttingDown, _ => Self::Stopped, } } } #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub enum SessionReason { Startup, LoginSucceeded, AgentReady, ReadinessLost, TransientTransport, Maintenance, Kicked, SimulatorDisconnected, ServerFailure, InvalidCredentials, InvalidConfiguration, StableSessionReset, OperatorPause, OperatorResume, OperatorForceReconnect, OperatorLogout, ShutdownRequested, ShutdownComplete, } #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub enum SessionFailureKind { InvalidCredentials, InvalidConfiguration, TransientTransport, Maintenance, Kicked, SimulatorDisconnected, ServerFailure, } /// Secret-free failure classification. Raw login/server text is deliberately /// not representable at this boundary. #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub struct SessionFailure { kind: SessionFailureKind, retry_after: Option, } impl SessionFailure { #[must_use] pub const fn new(kind: SessionFailureKind) -> Self { Self { kind, retry_after: None, } } #[must_use] pub const fn with_retry_after(kind: SessionFailureKind, retry_after: Duration) -> Self { Self { kind, retry_after: Some(retry_after), } } #[must_use] pub const fn kind(self) -> SessionFailureKind { self.kind } #[must_use] pub const fn retry_after(self) -> Option { self.retry_after } const fn retryable(self) -> bool { !matches!( self.kind, SessionFailureKind::InvalidCredentials | SessionFailureKind::InvalidConfiguration ) } const fn reason(self) -> SessionReason { match self.kind { SessionFailureKind::InvalidCredentials => SessionReason::InvalidCredentials, SessionFailureKind::InvalidConfiguration => SessionReason::InvalidConfiguration, SessionFailureKind::TransientTransport => SessionReason::TransientTransport, SessionFailureKind::Maintenance => SessionReason::Maintenance, SessionFailureKind::Kicked => SessionReason::Kicked, SessionFailureKind::SimulatorDisconnected => SessionReason::SimulatorDisconnected, SessionFailureKind::ServerFailure => SessionReason::ServerFailure, } } } impl fmt::Display for SessionFailure { fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { write!(formatter, "grid session failure: {:?}", self.kind) } } impl Error for SessionFailure {} #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub enum SessionSignal { Ready, Degraded, Disconnected(SessionFailure), } /// One connected generation. Dropping its subscriptions and joining its /// workers is the responsibility of `logout`, which the supervisor calls once. pub trait GridSession: Send { fn generation(&self) -> u64; fn next_signal(&mut self, cancellation: CancellationToken) -> SessionFuture<'_, SessionSignal>; fn logout( self: Box, cancellation: CancellationToken, ) -> SessionFuture<'static, Result<(), SessionFailure>>; } /// Injectable login owner. A production implementation must reuse the native /// `NetworkManager`; deterministic tests supply a fake implementation. pub trait GridSessionBackend: Send + Sync + 'static { fn login( &self, generation: u64, cancellation: CancellationToken, ) -> SessionFuture<'_, Result, SessionFailure>>; fn flush_audit( &self, cancellation: CancellationToken, ) -> SessionFuture<'_, Result<(), SessionFailure>>; } #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub struct ReconnectPolicy { pub initial_delay: Duration, pub maximum_delay: Duration, pub readiness_timeout: Duration, pub stable_reset_after: Duration, pub shutdown_deadline: Duration, pub jitter_basis_points: u16, pub instance_seed: u64, pub offline_work_capacity: usize, } impl Default for ReconnectPolicy { fn default() -> Self { Self { initial_delay: Duration::from_secs(1), maximum_delay: Duration::from_mins(1), readiness_timeout: Duration::from_mins(2), stable_reset_after: Duration::from_mins(2), shutdown_deadline: Duration::from_secs(10), jitter_basis_points: 2_000, instance_seed: runtime_jitter_seed(), offline_work_capacity: 128, } } } fn runtime_jitter_seed() -> u64 { static NEXT_SEED: AtomicU64 = AtomicU64::new(1); let time = SystemTime::now() .duration_since(UNIX_EPOCH) .map_or(0, |duration| { u64::try_from(duration.as_nanos()).unwrap_or(u64::MAX) }); mix64(time ^ u64::from(std::process::id()) ^ NEXT_SEED.fetch_add(1, Ordering::Relaxed)) } impl ReconnectPolicy { pub fn validate(self) -> Result { if self.initial_delay.is_zero() || self.initial_delay > self.maximum_delay || self.maximum_delay > MAX_BACKOFF || self.readiness_timeout.is_zero() || self.readiness_timeout > MAX_READINESS_TIMEOUT || self.stable_reset_after.is_zero() || self.stable_reset_after > MAX_STABLE_RESET || self.shutdown_deadline.is_zero() || self.shutdown_deadline > MAX_SHUTDOWN || self.jitter_basis_points > 5_000 || !(1..=MAX_OFFLINE_WORK).contains(&self.offline_work_capacity) { return Err(SessionSupervisorError::UnsafePolicy); } Ok(self) } #[must_use] pub fn retry_delay(self, failure_count: u32, hint: Option) -> Duration { let shift = failure_count.saturating_sub(1).min(31); let factor = 1_u128 << shift; let base_ms = self.initial_delay.as_millis().saturating_mul(factor); let cap_ms = self.maximum_delay.as_millis(); let hinted_ms = hint.map_or(0, |value| value.as_millis().min(cap_ms)); let unclamped = base_ms.max(hinted_ms).min(cap_ms); let spread = unclamped.saturating_mul(u128::from(self.jitter_basis_points)) / 10_000; let random = mix64(self.instance_seed ^ u64::from(failure_count)); let width = spread.saturating_mul(2).saturating_add(1); let offset = if width == 0 { 0 } else { u128::from(random) % width }; let jittered = unclamped .saturating_sub(spread) .saturating_add(offset) .max(hinted_ms); Duration::from_millis(u64::try_from(jittered.min(cap_ms)).unwrap_or(u64::MAX)) } } fn mix64(mut value: u64) -> u64 { value ^= value >> 30; value = value.wrapping_mul(0xbf58_476d_1ce4_e5b9); value ^= value >> 27; value = value.wrapping_mul(0x94d0_49bb_1331_11eb); value ^ (value >> 31) } #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub struct SessionStatus { pub state: SessionState, pub generation: u64, pub transport_connected: bool, pub agent_ready: bool, pub consecutive_failures: u32, } impl Default for SessionStatus { fn default() -> Self { Self { state: SessionState::Stopped, generation: 0, transport_connected: false, agent_ready: false, consecutive_failures: 0, } } } #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub enum WorkKind { ReadOnly, IdempotentMutation, NonIdempotentMutation, } #[derive(Clone, Debug, Eq, PartialEq)] pub struct SessionWork { pub id: BoundedText, pub kind: WorkKind, } impl SessionWork { pub fn new(id: impl Into, kind: WorkKind) -> Result { Ok(Self { id: BoundedText::new("session.work_id", id) .map_err(|_| SessionSupervisorError::InvalidWork)?, kind, }) } } #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub enum WorkDisposition { Accepted { generation: u64 }, Queued, RejectedOfflineMutation, RejectedQueueFull, RejectedDuplicate, RejectedShuttingDown, } #[derive(Clone, Debug, Eq, PartialEq)] pub enum SessionObservation { Transition { status: SessionStatus, reason: SessionReason, retry_in: Option, }, Work { id: BoundedText, disposition: WorkDisposition, }, } #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub enum SessionControl { Pause, Resume, ForceReconnect, Logout, Shutdown, } #[derive(Clone, Debug, Eq, PartialEq)] pub enum SessionSupervisorError { UnsafePolicy, InvalidWork, ControlClosed, ObservationClosed, TaskPanicked, ShutdownTimedOut, } impl fmt::Display for SessionSupervisorError { fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { match self { Self::UnsafePolicy => formatter.write_str("unsafe reconnect policy"), Self::InvalidWork => formatter.write_str("invalid session work envelope"), Self::ControlClosed => formatter.write_str("session control queue closed"), Self::ObservationClosed => formatter.write_str("session observation queue closed"), Self::TaskPanicked => formatter.write_str("session supervisor task panicked"), Self::ShutdownTimedOut => formatter.write_str("session shutdown deadline exceeded"), } } } impl Error for SessionSupervisorError {} enum Command { Control(SessionControl), Submit(SessionWork, oneshot::Sender), } #[derive(Debug)] struct GenerationFence { generation: AtomicU64, source: Mutex>, } impl GenerationFence { fn rotate(&self, generation: u64) -> CancellationToken { let mut source = lock(&self.source); if let Some(previous) = source.take() { previous.cancel(); } let current = CancellationTokenSource::new(); let token = current.token(); *source = Some(current); self.generation.store(generation, Ordering::Release); token } fn invalidate(&self) { if let Some(source) = lock(&self.source).take() { source.cancel(); } self.generation.store(0, Ordering::Release); } fn token(&self, generation: u64) -> Option { if self.generation.load(Ordering::Acquire) != generation { return None; } lock(&self.source) .as_ref() .map(CancellationTokenSource::token) } } pub struct SessionSupervisor { backend: Arc, policy: ReconnectPolicy, control_capacity: usize, observation_capacity: usize, } impl SessionSupervisor { pub fn new( backend: Arc, policy: ReconnectPolicy, control_capacity: usize, observation_capacity: usize, ) -> Result { let policy = policy.validate()?; if control_capacity == 0 || control_capacity > 256 || observation_capacity == 0 || observation_capacity > 8_192 { return Err(SessionSupervisorError::UnsafePolicy); } Ok(Self { backend, policy, control_capacity, observation_capacity, }) } #[must_use] pub fn start(self) -> SessionSupervisorHandle { let (commands, receiver) = mpsc::channel(self.control_capacity); let (observations, observation_receiver) = mpsc::channel(self.observation_capacity); let cancellation = CancellationTokenSource::new(); let state = Arc::new(AtomicU8::new(SessionState::Stopped as u8)); let status = Arc::new(Mutex::new(SessionStatus::default())); let fence = Arc::new(GenerationFence { generation: AtomicU64::new(0), source: Mutex::new(None), }); let runtime = Runtime { backend: self.backend, policy: self.policy, commands: receiver, observations, cancellation: cancellation.clone(), state: Arc::clone(&state), status: Arc::clone(&status), fence: Arc::clone(&fence), queued: VecDeque::new(), queued_ids: BTreeSet::new(), seen: VecDeque::new(), seen_ids: BTreeSet::new(), next_generation: 1, failures: 0, }; let task = tokio::spawn(runtime.run()); SessionSupervisorHandle { commands, observations: observation_receiver, cancellation, state, status, fence, task: Some(task), shutdown_deadline: self.policy.shutdown_deadline, } } } pub struct SessionSupervisorHandle { commands: mpsc::Sender, observations: mpsc::Receiver, cancellation: CancellationTokenSource, state: Arc, status: Arc>, fence: Arc, task: Option>, shutdown_deadline: Duration, } impl SessionSupervisorHandle { #[must_use] pub fn state(&self) -> SessionState { SessionState::from_u8(self.state.load(Ordering::Acquire)) } #[must_use] pub fn status(&self) -> SessionStatus { *lock(&self.status) } pub async fn control(&self, control: SessionControl) -> Result<(), SessionSupervisorError> { self.commands .send(Command::Control(control)) .await .map_err(|_| SessionSupervisorError::ControlClosed) } pub async fn submit( &self, work: SessionWork, ) -> Result { let (sender, receiver) = oneshot::channel(); self.commands .send(Command::Submit(work, sender)) .await .map_err(|_| SessionSupervisorError::ControlClosed)?; receiver .await .map_err(|_| SessionSupervisorError::ControlClosed) } pub async fn next_observation(&mut self) -> Option { self.observations.recv().await } /// Returns the cancellation token shared by inference, scheduled work, and /// tools for exactly one live generation. #[must_use] pub fn generation_token(&self, generation: u64) -> Option { self.fence.token(generation) } /// Late LLM/tool results are accepted only while their generation is fully ready. #[must_use] pub fn accepts_result(&self, generation: u64) -> bool { let status = self.status(); status.agent_ready && status.generation == generation } pub async fn shutdown(&mut self) -> Result<(), SessionSupervisorError> { self.state .store(SessionState::ShuttingDown as u8, Ordering::Release); self.cancellation.cancel(); let Some(task) = self.task.as_mut() else { return Ok(()); }; match tokio::time::timeout(self.shutdown_deadline, &mut *task).await { Ok(Ok(())) => { self.task = None; Ok(()) } Ok(Err(_)) => { self.task = None; Err(SessionSupervisorError::TaskPanicked) } Err(_) => { task.abort(); let _ = (&mut *task).await; self.task = None; Err(SessionSupervisorError::ShutdownTimedOut) } } } } impl Drop for SessionSupervisorHandle { fn drop(&mut self) { self.cancellation.cancel(); if let Some(task) = &self.task { task.abort(); } } } #[derive(Clone, Copy, Debug, Eq, PartialEq)] enum DesiredState { Running, Paused, LoggedOut, AuthenticationBlocked, Shutdown, } enum LoginEvent { Cancelled, Command(Option), Complete(Result, SessionFailure>), } enum ActiveEvent { Cancelled, Command(Option), ReadinessTimedOut, Signal(SessionSignal), } struct Runtime { backend: Arc, policy: ReconnectPolicy, commands: mpsc::Receiver, observations: mpsc::Sender, cancellation: CancellationTokenSource, state: Arc, status: Arc>, fence: Arc, queued: VecDeque, queued_ids: BTreeSet, seen: VecDeque, seen_ids: BTreeSet, next_generation: u64, failures: u32, } impl Runtime { #[allow(clippy::too_many_lines)] async fn run(mut self) { self.transition(SessionState::Stopped, SessionReason::Startup, None, 0); let mut desired = DesiredState::Running; let mut session: Option> = None; let mut ready = false; let mut readiness_deadline = None; let mut stable_since = None; loop { if self.cancellation.token().is_cancellation_requested() { desired = DesiredState::Shutdown; } match desired { DesiredState::Shutdown => break, DesiredState::Paused | DesiredState::LoggedOut | DesiredState::AuthenticationBlocked => { let command = tokio::select! { () = self.cancellation.token().cancelled() => { desired = DesiredState::Shutdown; continue; } command = self.commands.recv() => command, }; desired = self.handle_idle_command(command, desired); } DesiredState::Running if session.is_none() => { if self.next_generation == u64::MAX { desired = DesiredState::Shutdown; continue; } let generation = self.next_generation; self.next_generation += 1; let token = self.fence.rotate(generation); self.transition( SessionState::Connecting, if self.failures == 0 { SessionReason::OperatorResume } else { SessionReason::TransientTransport }, None, generation, ); let login_event = { let backend = Arc::clone(&self.backend); let login = backend.login(generation, token); tokio::pin!(login); loop { let event = tokio::select! { () = self.cancellation.token().cancelled() => LoginEvent::Cancelled, command = self.commands.recv() => LoginEvent::Command(command), result = &mut login => LoginEvent::Complete(result), }; if let LoginEvent::Command(Some(Command::Submit(work, response))) = event { let disposition = self.handle_work(work.clone(), false, 0); let _ = response.send(disposition); self.emit_work(work.id, disposition); } else { break event; } } }; let result = match login_event { LoginEvent::Cancelled => { desired = DesiredState::Shutdown; self.fence.invalidate(); continue; } LoginEvent::Command(command) => { self.fence.invalidate(); desired = self.handle_idle_command(command, DesiredState::Running); continue; } LoginEvent::Complete(result) => result, }; match result { Ok(connected) if connected.generation() == generation => { session = Some(connected); ready = false; readiness_deadline = Some(Instant::now() + self.policy.readiness_timeout); stable_since = None; self.transition( SessionState::Degraded, SessionReason::LoginSucceeded, None, generation, ); } Ok(connected) => { self.fence.invalidate(); self.close_session(connected).await; self.failures = self.failures.saturating_add(1); desired = self .backoff(SessionFailure::new(SessionFailureKind::ServerFailure)) .await; } Err(failure) if !failure.retryable() => { self.fence.invalidate(); self.transition( SessionState::AuthenticationBlocked, failure.reason(), None, generation, ); desired = DesiredState::AuthenticationBlocked; } Err(failure) => { self.fence.invalidate(); self.failures = self.failures.saturating_add(1); desired = self.backoff(failure).await; } } } DesiredState::Running => { let generation = self.status_snapshot().generation; let Some(token) = self.fence.token(generation) else { session = None; continue; }; let active_event = { let Some(active) = session.as_mut() else { continue; }; let signal = active.next_signal(token); tokio::pin!(signal); let readiness = async { if let Some(deadline) = readiness_deadline { tokio::time::sleep_until(deadline).await; } else { std::future::pending::<()>().await; } }; tokio::pin!(readiness); tokio::select! { () = self.cancellation.token().cancelled() => ActiveEvent::Cancelled, command = self.commands.recv() => ActiveEvent::Command(command), () = &mut readiness => ActiveEvent::ReadinessTimedOut, next = &mut signal => ActiveEvent::Signal(next), } }; match active_event { ActiveEvent::Cancelled => { desired = DesiredState::Shutdown; } ActiveEvent::ReadinessTimedOut => { ready = false; readiness_deadline = None; stable_since = None; self.fence.invalidate(); if let Some(unready) = session.take() { self.close_session(unready).await; } self.failures = self.failures.saturating_add(1); desired = self .backoff(SessionFailure::new( SessionFailureKind::TransientTransport, )) .await; } ActiveEvent::Command(command) => match command { Some(Command::Submit(work, response)) => { let disposition = self.handle_work(work.clone(), ready, generation); let _ = response.send(disposition); self.emit_work(work.id, disposition); } None => { let (next, reconnect) = self.handle_active_control(None); desired = next; if reconnect { self.failures = 0; } if desired != DesiredState::Running || reconnect { self.fence.invalidate(); } } Some(Command::Control(control)) => { let (next, reconnect) = self.handle_active_control(Some(control)); desired = next; if reconnect { self.failures = 0; } if desired != DesiredState::Running || reconnect { self.fence.invalidate(); } } }, ActiveEvent::Signal(next) => match next { SessionSignal::Ready => { ready = true; readiness_deadline = None; stable_since.get_or_insert_with(Instant::now); self.transition( SessionState::Online, SessionReason::AgentReady, None, generation, ); self.release_queued(generation); } SessionSignal::Degraded => { ready = false; readiness_deadline = Some(Instant::now() + self.policy.readiness_timeout); stable_since = None; self.transition( SessionState::Degraded, SessionReason::ReadinessLost, None, generation, ); } SessionSignal::Disconnected(failure) => { ready = false; self.fence.invalidate(); if stable_since.is_some_and(|start| { start.elapsed() >= self.policy.stable_reset_after }) { self.failures = 0; self.transition( SessionState::Degraded, SessionReason::StableSessionReset, None, generation, ); } self.failures = self.failures.saturating_add(1); if let Some(disconnected) = session.take() { self.close_session(disconnected).await; } desired = self.backoff(failure).await; stable_since = None; } }, } if session.is_some() && (desired != DesiredState::Running || self.fence.token(generation).is_none()) { if let Some(active) = session.take() { self.close_session(active).await; } ready = false; readiness_deadline = None; stable_since = None; } } } } self.transition( SessionState::ShuttingDown, SessionReason::ShutdownRequested, None, self.status_snapshot().generation, ); self.fence.invalidate(); if let Some(active) = session.take() { self.close_session(active).await; } self.reject_queued(); let cleanup = CancellationTokenSource::new(); let _ = tokio::time::timeout( self.cleanup_timeout(), self.backend.flush_audit(cleanup.token()), ) .await; cleanup.cancel(); self.transition( SessionState::Stopped, SessionReason::ShutdownComplete, None, 0, ); } fn handle_idle_command( &mut self, command: Option, current: DesiredState, ) -> DesiredState { match command { None | Some(Command::Control(SessionControl::Shutdown)) => DesiredState::Shutdown, Some(Command::Control(SessionControl::Pause)) => { self.transition(SessionState::Paused, SessionReason::OperatorPause, None, 0); DesiredState::Paused } Some(Command::Control(SessionControl::Resume | SessionControl::ForceReconnect)) => { self.failures = 0; DesiredState::Running } Some(Command::Control(SessionControl::Logout)) => { self.transition( SessionState::Stopped, SessionReason::OperatorLogout, None, 0, ); DesiredState::LoggedOut } Some(Command::Submit(work, response)) => { let disposition = self.handle_work(work.clone(), false, 0); let _ = response.send(disposition); self.emit_work(work.id, disposition); current } } } fn handle_active_control(&mut self, command: Option) -> (DesiredState, bool) { match command { None | Some(SessionControl::Shutdown) => (DesiredState::Shutdown, false), Some(SessionControl::Pause) => { self.transition(SessionState::Paused, SessionReason::OperatorPause, None, 0); (DesiredState::Paused, false) } Some(SessionControl::Resume) => (DesiredState::Running, false), Some(SessionControl::ForceReconnect) => { self.transition( SessionState::Connecting, SessionReason::OperatorForceReconnect, None, 0, ); (DesiredState::Running, true) } Some(SessionControl::Logout) => { self.transition( SessionState::Stopped, SessionReason::OperatorLogout, None, 0, ); (DesiredState::LoggedOut, false) } } } async fn backoff(&mut self, failure: SessionFailure) -> DesiredState { let delay = self .policy .retry_delay(self.failures, failure.retry_after()); self.transition(SessionState::Backoff, failure.reason(), Some(delay), 0); let sleep = tokio::time::sleep(delay); tokio::pin!(sleep); loop { tokio::select! { () = self.cancellation.token().cancelled() => return DesiredState::Shutdown, () = &mut sleep => return DesiredState::Running, command = self.commands.recv() => match command { Some(Command::Submit(work, response)) => { let disposition = self.handle_work(work.clone(), false, 0); let _ = response.send(disposition); self.emit_work(work.id, disposition); } other => return self.handle_idle_command(other, DesiredState::Running), } } } } async fn close_session(&self, session: Box) { let cleanup = CancellationTokenSource::new(); let _ = tokio::time::timeout(self.cleanup_timeout(), session.logout(cleanup.token())).await; cleanup.cancel(); } fn cleanup_timeout(&self) -> Duration { (self.policy.shutdown_deadline / 3).max(Duration::from_millis(1)) } fn handle_work(&mut self, work: SessionWork, ready: bool, generation: u64) -> WorkDisposition { if self.cancellation.token().is_cancellation_requested() { return WorkDisposition::RejectedShuttingDown; } if self.seen_ids.contains(work.id.as_str()) || self.queued_ids.contains(work.id.as_str()) { return WorkDisposition::RejectedDuplicate; } if ready { self.remember_work(work.id.as_str().to_owned()); return WorkDisposition::Accepted { generation }; } if work.kind != WorkKind::ReadOnly { return WorkDisposition::RejectedOfflineMutation; } if self.queued.len() == self.policy.offline_work_capacity { return WorkDisposition::RejectedQueueFull; } self.queued_ids.insert(work.id.as_str().to_owned()); self.queued.push_back(work); WorkDisposition::Queued } fn release_queued(&mut self, generation: u64) { while let Some(work) = self.queued.pop_front() { self.queued_ids.remove(work.id.as_str()); self.remember_work(work.id.as_str().to_owned()); self.emit_work(work.id, WorkDisposition::Accepted { generation }); } } fn remember_work(&mut self, id: String) { if self.seen.len() == MAX_SEEN_WORK && let Some(expired) = self.seen.pop_front() { self.seen_ids.remove(&expired); } self.seen_ids.insert(id.clone()); self.seen.push_back(id); } fn reject_queued(&mut self) { while let Some(work) = self.queued.pop_front() { self.queued_ids.remove(work.id.as_str()); self.emit_work(work.id, WorkDisposition::RejectedShuttingDown); } } fn transition( &self, state: SessionState, reason: SessionReason, retry_in: Option, generation: u64, ) { self.state.store(state as u8, Ordering::Release); let status = SessionStatus { state, generation, transport_connected: matches!(state, SessionState::Online | SessionState::Degraded), agent_ready: state == SessionState::Online, consecutive_failures: self.failures, }; *lock(&self.status) = status; let _ = self.observations.try_send(SessionObservation::Transition { status, reason, retry_in, }); } fn emit_work(&self, id: BoundedText, disposition: WorkDisposition) { let _ = self .observations .try_send(SessionObservation::Work { id, disposition }); } fn status_snapshot(&self) -> SessionStatus { *lock(&self.status) } } fn lock(mutex: &Mutex) -> std::sync::MutexGuard<'_, T> { mutex .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) }