1132 lines
39 KiB
Rust
1132 lines
39 KiB
Rust
//! 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<Box<dyn Future<Output = T> + Send + 'a>>;
|
|
|
|
/// Generation-scoped gate between transport readiness and autonomous agent readiness.
|
|
pub struct FrameworkReadiness {
|
|
generation: watch::Sender<u64>,
|
|
}
|
|
|
|
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<Duration>,
|
|
}
|
|
|
|
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<Duration> {
|
|
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<Self>,
|
|
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<Box<dyn GridSession>, 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<Self, SessionSupervisorError> {
|
|
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>) -> 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<MAX_IDENTIFIER_BYTES>,
|
|
pub kind: WorkKind,
|
|
}
|
|
|
|
impl SessionWork {
|
|
pub fn new(id: impl Into<String>, kind: WorkKind) -> Result<Self, SessionSupervisorError> {
|
|
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<Duration>,
|
|
},
|
|
Work {
|
|
id: BoundedText<MAX_IDENTIFIER_BYTES>,
|
|
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<WorkDisposition>),
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
struct GenerationFence {
|
|
generation: AtomicU64,
|
|
source: Mutex<Option<CancellationTokenSource>>,
|
|
}
|
|
|
|
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<CancellationToken> {
|
|
if self.generation.load(Ordering::Acquire) != generation {
|
|
return None;
|
|
}
|
|
lock(&self.source)
|
|
.as_ref()
|
|
.map(CancellationTokenSource::token)
|
|
}
|
|
}
|
|
|
|
pub struct SessionSupervisor {
|
|
backend: Arc<dyn GridSessionBackend>,
|
|
policy: ReconnectPolicy,
|
|
control_capacity: usize,
|
|
observation_capacity: usize,
|
|
}
|
|
|
|
impl SessionSupervisor {
|
|
pub fn new(
|
|
backend: Arc<dyn GridSessionBackend>,
|
|
policy: ReconnectPolicy,
|
|
control_capacity: usize,
|
|
observation_capacity: usize,
|
|
) -> Result<Self, SessionSupervisorError> {
|
|
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<Command>,
|
|
observations: mpsc::Receiver<SessionObservation>,
|
|
cancellation: CancellationTokenSource,
|
|
state: Arc<AtomicU8>,
|
|
status: Arc<Mutex<SessionStatus>>,
|
|
fence: Arc<GenerationFence>,
|
|
task: Option<JoinHandle<()>>,
|
|
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<WorkDisposition, SessionSupervisorError> {
|
|
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<SessionObservation> {
|
|
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<CancellationToken> {
|
|
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<Command>),
|
|
Complete(Result<Box<dyn GridSession>, SessionFailure>),
|
|
}
|
|
|
|
enum ActiveEvent {
|
|
Cancelled,
|
|
Command(Option<Command>),
|
|
ReadinessTimedOut,
|
|
Signal(SessionSignal),
|
|
}
|
|
|
|
struct Runtime {
|
|
backend: Arc<dyn GridSessionBackend>,
|
|
policy: ReconnectPolicy,
|
|
commands: mpsc::Receiver<Command>,
|
|
observations: mpsc::Sender<SessionObservation>,
|
|
cancellation: CancellationTokenSource,
|
|
state: Arc<AtomicU8>,
|
|
status: Arc<Mutex<SessionStatus>>,
|
|
fence: Arc<GenerationFence>,
|
|
queued: VecDeque<SessionWork>,
|
|
queued_ids: BTreeSet<String>,
|
|
seen: VecDeque<String>,
|
|
seen_ids: BTreeSet<String>,
|
|
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<Box<dyn GridSession>> = 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<Command>,
|
|
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<SessionControl>) -> (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<dyn GridSession>) {
|
|
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<Duration>,
|
|
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<MAX_IDENTIFIER_BYTES>, disposition: WorkDisposition) {
|
|
let _ = self
|
|
.observations
|
|
.try_send(SessionObservation::Work { id, disposition });
|
|
}
|
|
|
|
fn status_snapshot(&self) -> SessionStatus {
|
|
*lock(&self.status)
|
|
}
|
|
}
|
|
|
|
fn lock<T>(mutex: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
|
|
mutex
|
|
.lock()
|
|
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
|
}
|