//! Session typestate, single-owner registry synchronization, and lifecycle transitions.
#![forbid(unsafe_code)]
use anyhow::{bail, ensure, Context, Result};
use klbr_runtime::{
behavior::Release,
protocol::{OperationId, SessionId},
session::SessionRuntime,
store::{Receipt, Store},
worker::WorkerConfig,
};
use serde::{Deserialize, Serialize};
use std::{
collections::HashMap,
path::PathBuf,
sync::{
atomic::{AtomicBool, Ordering},
Arc, Mutex as StdMutex,
},
};
use tokio::sync::{Mutex, RwLock};
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub struct Epoch(pub i64);
impl From<i64> for Epoch {
fn from(epoch: i64) -> Self {
Self(epoch)
}
}
impl From<Epoch> for i64 {
fn from(epoch: Epoch) -> Self {
epoch.0
}
}
impl std::fmt::Display for Epoch {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.0)
}
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum SessionState {
Ready,
Executing {
epoch: i64,
execution_id: OperationId,
},
Waiting {
receipt: Option<Receipt>,
wake_at_ms: Option<i64>,
reason: Option<String>,
},
Draining,
Stopped,
}
impl SessionState {
#[must_use]
pub fn can_admit(&self) -> bool {
matches!(
self,
Self::Ready | Self::Executing { .. } | Self::Waiting { .. }
)
}
pub fn transition_to_executing(&mut self, epoch: i64, execution_id: OperationId) -> Result<()> {
match self {
Self::Ready => {
*self = Self::Executing {
epoch,
execution_id,
};
Ok(())
}
_ => bail!("invalid state transition: cannot transition to Executing from {self:?}"),
}
}
pub fn transition_to_ready(&mut self) -> Result<()> {
match self {
Self::Executing { .. } | Self::Waiting { .. } => {
*self = Self::Ready;
Ok(())
}
_ => bail!("invalid state transition: cannot transition to Ready from {self:?}"),
}
}
pub fn transition_to_waiting(
&mut self,
receipt: Option<Receipt>,
wake_at_ms: Option<i64>,
reason: Option<String>,
) -> Result<()> {
match self {
Self::Executing { .. } => {
*self = Self::Waiting {
receipt,
wake_at_ms,
reason,
};
Ok(())
}
_ => bail!("invalid state transition: cannot transition to Waiting from {self:?}"),
}
}
pub fn transition_to_draining(&mut self) -> Result<()> {
match self {
Self::Ready | Self::Executing { .. } | Self::Waiting { .. } => {
*self = Self::Draining;
Ok(())
}
_ => bail!("invalid state transition: cannot transition to Draining from {self:?}"),
}
}
pub fn transition_to_stopped(&mut self) -> Result<()> {
*self = Self::Stopped;
Ok(())
}
}
// lock order invariant: always acquire active_runner first, then state.
// never acquire active_runner while holding state. keeps lock lifetimes
// short and never holds either lock across runtime awaits.
/// A registered runner: the operation it owns, the handle that interrupts it,
/// and, when attached, the join handle that proves it has fully settled.
struct RunnerRegistration {
op_id: OperationId,
abort: futures::future::AbortHandle,
task: Option<tokio::task::JoinHandle<()>>,
/// The attempt currently holding this session's work claim. A slice mints its
/// own attempt identity, so releasing claims by the runner's operation id
/// alone would orphan the interrupted slice's claim.
active_attempt: Option<OperationId>,
}
/// Coarse owner class for service-level background tasks tracked by the
/// registry so `shutdown_all` can join every coordinator before returning.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum ServiceTaskGroup {
/// periodic wake / pending-work driver loop
Driver,
/// one-shot startup recovery scan
Recovery,
/// per-client connection handler
Client,
/// any other service-level background task
Other,
}
#[derive(Clone)]
pub struct SessionEntry {
pub id: SessionId,
pub runtime: Arc<SessionRuntime>,
pub state: Arc<RwLock<SessionState>>,
pub started_at_ms: i64,
active_runner: Arc<Mutex<Option<RunnerRegistration>>>,
retained_runners: Arc<Mutex<Vec<tokio::task::JoinHandle<()>>>>,
quiescing: Arc<AtomicBool>,
}
impl SessionEntry {
pub fn new(id: SessionId, runtime: Arc<SessionRuntime>, started_at_ms: i64) -> Self {
Self {
id,
runtime,
state: Arc::new(RwLock::new(SessionState::Ready)),
started_at_ms,
active_runner: Arc::new(Mutex::new(None)),
retained_runners: Arc::new(Mutex::new(Vec::new())),
quiescing: Arc::new(AtomicBool::new(false)),
}
}
async fn drain_retained_runners(&self) {
let tasks = {
let mut guard = self.retained_runners.lock().await;
std::mem::take(&mut *guard)
};
let curr_id = tokio::task::try_id();
let mut keep = Vec::new();
for task in tasks {
if curr_id == Some(task.id()) {
keep.push(task);
} else {
let _ = task.await;
}
}
if !keep.is_empty() {
let mut guard = self.retained_runners.lock().await;
guard.extend(keep);
}
}
pub async fn register_runner(&self, handle: futures::future::AbortHandle) {
let old_task = {
let mut guard = self.active_runner.lock().await;
let old = guard.take();
*guard = Some(RunnerRegistration {
op_id: OperationId::fresh(),
abort: handle,
task: None,
active_attempt: None,
});
old.and_then(|mut r| r.task.take())
};
if let Some(t) = old_task {
if tokio::task::try_id() == Some(t.id()) {
self.retained_runners.lock().await.push(t);
} else {
let _ = t.await;
}
}
self.drain_retained_runners().await;
}
pub async fn register_runner_with_id(
&self,
op_id: OperationId,
handle: futures::future::AbortHandle,
) {
let old_task = {
let mut guard = self.active_runner.lock().await;
let old = guard.take();
*guard = Some(RunnerRegistration {
op_id,
abort: handle,
task: None,
active_attempt: None,
});
old.and_then(|mut r| r.task.take())
};
if let Some(t) = old_task {
if tokio::task::try_id() == Some(t.id()) {
self.retained_runners.lock().await.push(t);
} else {
let _ = t.await;
}
}
self.drain_retained_runners().await;
}
pub async fn try_start_runner(
&self,
op_id: OperationId,
handle: futures::future::AbortHandle,
) -> Result<bool> {
if self.quiescing.load(Ordering::SeqCst) {
// service teardown owns this entry; refuse new owners
return Ok(false);
}
// Await any retained predecessor runner tasks before admitting a successor
self.drain_retained_runners().await;
// read epoch before acquiring locks so we never hold a lock across runtime awaits
let epoch = self.runtime.hook_epoch().await;
// lock order: active_runner first, then state
let mut runner_guard = self.active_runner.lock().await;
if runner_guard.is_some() {
// predecessor retirement is still pending; refuse to start successor
return Ok(false);
}
let mut state_guard = self.state.write().await;
match &*state_guard {
SessionState::Draining | SessionState::Stopped => {
bail!(
"session {} is in state {:?}, cannot run turn",
self.id,
*state_guard
);
}
SessionState::Executing { .. } => {
return Ok(false);
}
SessionState::Waiting { .. } => {
state_guard.transition_to_ready()?;
}
SessionState::Ready => {}
}
state_guard.transition_to_executing(epoch, op_id.clone())?;
*runner_guard = Some(RunnerRegistration {
op_id,
abort: handle,
task: None,
active_attempt: None,
});
Ok(true)
}
/// Record which attempt currently holds this session's work claim.
///
/// The runner mints a fresh attempt identity per slice, so anything that must
/// find the live claim (quiesce releasing a claim) has to be told about it
/// rather than assuming the runner's operation id.
pub async fn set_active_attempt(&self, op_id: &OperationId, attempt: OperationId) -> bool {
let mut guard = self.active_runner.lock().await;
match guard.as_mut() {
Some(reg) if reg.op_id == *op_id => {
reg.active_attempt = Some(attempt);
true
}
_ => false,
}
}
/// Atomically retire the active runner only if its operation ID matches.
/// An old runner cannot retire its successor. Coherently resets state if still Executing.
/// Retains and awaits the runner task to full exit so no runner outlives retirement.
pub async fn retire_runner(&self, expected_op_id: &OperationId) -> bool {
let task = {
// lock order: active_runner first, then state
let mut guard = self.active_runner.lock().await;
if let Some(reg) = guard.as_mut() {
if reg.op_id == *expected_op_id {
let task = reg.task.take();
*guard = None;
let mut state_guard = self.state.write().await;
if let SessionState::Executing { execution_id, .. } = &*state_guard {
if execution_id == expected_op_id {
*state_guard = SessionState::Ready;
}
}
task
} else {
return false;
}
} else {
return false;
}
};
if let Some(task) = task {
if tokio::task::try_id() == Some(task.id()) {
self.retained_runners.lock().await.push(task);
} else {
let _ = task.await;
}
}
self.drain_retained_runners().await;
true
}
/// Attach the spawned runner task's join handle so a service quiesce can
/// await the owner's own retirement instead of sleeping. Only the runner
/// still registered under `expected_op_id` accepts the handle.
pub async fn attach_runner_task(
&self,
expected_op_id: &OperationId,
task: tokio::task::JoinHandle<()>,
) -> bool {
let mut guard = self.active_runner.lock().await;
if let Some(reg) = guard.as_mut() {
if reg.op_id == *expected_op_id {
reg.task = Some(task);
return true;
}
}
// Active runner slot cleared, aborted, or op_id mismatched; retain and abort task
task.abort();
self.retained_runners.lock().await.push(task);
false
}
pub fn set_quiescing(&self) {
self.quiescing.store(true, Ordering::SeqCst);
}
pub async fn clear_runner(&self) {
let task = {
let mut guard = self.active_runner.lock().await;
guard.take().and_then(|mut r| r.task.take())
};
if let Some(task) = task {
if tokio::task::try_id() == Some(task.id()) {
self.retained_runners.lock().await.push(task);
} else {
let _ = task.await;
}
}
self.drain_retained_runners().await;
}
pub async fn active_runner_id(&self) -> Option<OperationId> {
self.active_runner
.lock()
.await
.as_ref()
.map(|reg| reg.op_id.clone())
}
pub async fn abort_runner(&self) {
// lock order: active_runner first, then state
let task = {
let mut guard = self.active_runner.lock().await;
if let Some(mut reg) = guard.take() {
let op_id = reg.op_id;
reg.abort.abort();
let mut state_guard = self.state.write().await;
if let SessionState::Executing { execution_id, .. } = &*state_guard {
if *execution_id == op_id {
*state_guard = SessionState::Ready;
}
}
reg.task.take()
} else {
None
}
};
// interrupt runtime outside of any locks
self.runtime.interrupt().await;
if let Some(task) = task {
if tokio::task::try_id() == Some(task.id()) {
self.retained_runners.lock().await.push(task);
} else {
let _ = task.await;
}
}
self.drain_retained_runners().await;
}
pub async fn state(&self) -> SessionState {
self.state.read().await.clone()
}
pub async fn can_admit(&self) -> bool {
!self.quiescing.load(Ordering::SeqCst) && self.state.read().await.can_admit()
}
pub async fn transition_to_executing(
&self,
epoch: i64,
execution_id: OperationId,
) -> Result<()> {
self.state
.write()
.await
.transition_to_executing(epoch, execution_id)
}
pub async fn transition_to_ready(&self) -> Result<()> {
self.state.write().await.transition_to_ready()
}
pub async fn transition_to_ready_with_store(&self, store: &Store) -> Result<()> {
self.state.write().await.transition_to_ready()?;
store.set_session_state(&self.id, "ready").await?;
Ok(())
}
pub async fn transition_to_waiting(
&self,
receipt: Option<Receipt>,
wake_at_ms: Option<i64>,
reason: Option<String>,
) -> Result<()> {
self.state
.write()
.await
.transition_to_waiting(receipt, wake_at_ms, reason)
}
pub async fn transition_to_waiting_with_store(
&self,
store: &Store,
receipt: Option<Receipt>,
wake_at_ms: Option<i64>,
reason: Option<String>,
) -> Result<()> {
self.state
.write()
.await
.transition_to_waiting(receipt, wake_at_ms, reason)?;
store.wait_session(&self.id, wake_at_ms).await?;
Ok(())
}
/// Persist a waiting state for a task the owner has stopped working on.
///
/// `Waiting` is only reachable from `Executing`, so if a concurrent path left
/// the entry `Ready` while this runner still owns the session, the owner is
/// recorded as executing again first. That is truthful - it holds the runner
/// registration - and it keeps the state machine's transition rules intact
/// instead of reaching around them.
pub async fn wait_as_owner(
&self,
store: &Store,
op_id: &OperationId,
reason: Option<String>,
) -> Result<()> {
let needs_claim = {
let mut state = self.state.write().await;
if matches!(*state, SessionState::Ready) {
let epoch = self.runtime.hook_epoch().await;
state.transition_to_executing(epoch, op_id.clone())?;
true
} else {
false
}
};
if needs_claim {
tracing::info!(session = %self.id, "re-claimed an idle session to record why the task stopped");
}
self.transition_to_waiting_with_store(store, None, None, reason)
.await
}
pub async fn transition_to_draining(&self) -> Result<()> {
self.state.write().await.transition_to_draining()
}
pub async fn transition_to_stopped(&self) -> Result<()> {
self.state.write().await.transition_to_stopped()
}
/// Service quiescence: tear this entry down WITHOUT persisting an operator
/// Stop. No new work is admitted, the in-flight attempt is interrupted and
/// awaited until its owner retires, its durable work claims are released,
/// and truthful continuation state (including a waiting deadline) is
/// persisted. An already-explicit Stopped session stays Stopped.
pub async fn quiesce(&self, store: &Store) -> Result<()> {
self.set_quiescing();
let taken = { self.active_runner.lock().await.take() };
if let Some(mut reg) = taken {
reg.abort.abort();
// The interrupted attempt may be the runner itself or the slice it was
// running, so both identities are released.
let mut owners = vec![reg.op_id.clone()];
if let Some(attempt) = reg.active_attempt.clone() {
if attempt != reg.op_id {
owners.push(attempt);
}
}
for owner in owners {
match store
.release_attempt_claims(&owner, "service quiesced; attempt interrupted")
.await
{
Ok(ids) if !ids.is_empty() => {
tracing::info!(session = %self.id, attempt = %owner, released = ids.len(), "released interrupted attempt claims");
}
Ok(_) => {}
Err(e) => {
tracing::warn!(session = %self.id, attempt = %owner, err = %e, "failed to release attempt claims during quiesce");
}
}
}
if let Some(task) = reg.task.take() {
if tokio::task::try_id() != Some(task.id()) {
let _ = task.await;
} else {
self.retained_runners.lock().await.push(task);
}
}
// the joined owner's tail could not retire what we already took;
// finish that reset here so memory matches the persisted state
let mut state_guard = self.state.write().await;
if let SessionState::Executing { execution_id, .. } = &*state_guard {
if *execution_id == reg.op_id {
*state_guard = SessionState::Ready;
}
}
}
self.runtime.interrupt().await;
self.drain_retained_runners().await;
// quiesce persists continuation truth; an explicit Stop is never erased
match self.state().await {
SessionState::Stopped => {
store.set_session_state(&self.id, "stopped").await?;
}
SessionState::Draining => {
store.set_session_state(&self.id, "draining").await?;
}
SessionState::Waiting { wake_at_ms, .. } => {
store.wait_session(&self.id, wake_at_ms).await?;
}
SessionState::Ready | SessionState::Executing { .. } => {
match store.session_state(&self.id).await? {
Some((ref st, wake)) if st == "waiting" => {
store.wait_session(&self.id, wake).await?;
}
_ => {
store.set_session_state(&self.id, "ready").await?;
}
}
}
}
self.runtime.shutdown().await;
Ok(())
}
}
fn ensure_configured_matches_durable(
configured: &klbr_runtime::session::Placement,
durable: &klbr_runtime::target::SessionPlacement,
) -> anyhow::Result<()> {
let configured_target = match configured {
Placement::Remote { target, .. } => Some(target.id.to_string()),
Placement::Local => None,
};
let durable_target = match &durable.kernel_target {
KernelTarget::Remote(remote) => Some(remote.id.to_string()),
KernelTarget::Local => None,
};
anyhow::ensure!(
configured_target == durable_target,
"durable placement names target {:?}; the configured placement names {:?}",
durable_target,
configured_target
);
Ok(())
}
use klbr_runtime::{session::Placement, target::KernelTarget};
/// Placement source for the session descriptor. Resolved once per session
/// open; None means the local default. Durable remote placement must come
/// from here - a restarted remote session must never silently fall back to
/// local execution.
pub trait PlacementSource: Send + Sync {
fn placement(
&self,
session: &SessionId,
) -> futures::future::BoxFuture<'_, anyhow::Result<Option<klbr_runtime::session::Placement>>>;
}
/// Placement resolved from the service's packaged configuration.
///
/// The registry never invents a target: whatever this source returns is what
/// a session opens on, and a session the configuration does not claim opens
/// locally. A configured entry that cannot be built (no workspace, no relay)
/// fails the open instead of falling back.
pub struct ConfiguredPlacementSource {
config: klbr_runtime::remote::placement::PlacementConfig,
}
impl ConfiguredPlacementSource {
#[must_use]
pub fn new(config: klbr_runtime::remote::placement::PlacementConfig) -> Self {
Self { config }
}
/// Resolve once at boot, so a broken declaration fails the daemon rather
/// than the first session that happens to touch it.
pub fn build(config: klbr_runtime::remote::placement::PlacementConfig) -> anyhow::Result<Self> {
config.validate()?;
Ok(Self::new(config))
}
}
impl PlacementSource for ConfiguredPlacementSource {
fn placement(
&self,
session: &SessionId,
) -> futures::future::BoxFuture<'_, anyhow::Result<Option<klbr_runtime::session::Placement>>>
{
let entry = self.config.entry_for(session.as_str()).cloned();
Box::pin(async move {
match entry {
Some(entry) => Ok(Some(entry.placement()?)),
None => Ok(None),
}
})
}
}
#[derive(Clone)]
pub struct SessionRegistry {
pub store: Store,
pub worker_config: WorkerConfig,
pub sessions: Arc<RwLock<HashMap<SessionId, Arc<SessionEntry>>>>,
open_locks: Arc<Mutex<HashMap<SessionId, Arc<Mutex<()>>>>>,
service_tasks: Arc<StdMutex<HashMap<ServiceTaskGroup, Vec<tokio::task::JoinHandle<()>>>>>,
placement: Option<Arc<dyn PlacementSource>>,
pub recovery_policy: crate::observation::RecoveryObservationPolicy,
/// Token budget available to REQUIRED prompt guidance.
///
/// Derived from the operator's configured `watermark_tokens` rather than a
/// constant of its own, so there is one context-budget number in the system.
/// A second independent constant would silently disagree the moment config
/// changed, with no compile error and no failing test.
guidance_budget_tokens: usize,
task_budget: klbr_core::runtime_turn::budget::TaskBudget,
/// Conversation events kept raw below an installed context frontier. Derived from the
/// operator's configured session window for the same reason as the guidance budget: one
/// window in the system, not two numbers that silently disagree.
history_window: usize,
pub peer_host: Arc<std::sync::RwLock<Option<Arc<dyn klbr_runtime::protocol::PeerHost>>>>,
}
impl SessionRegistry {
/// Steps one slice may use, and steps the whole task may use. A task never
/// grows a new allowance by being resliced.
pub fn task_budget(&self) -> klbr_core::runtime_turn::budget::TaskBudget {
self.task_budget
}
/// Set the task budget. A test uses a tiny one to make exhaustion ordinary.
#[must_use]
pub fn with_task_budget(
mut self,
task_budget: klbr_core::runtime_turn::budget::TaskBudget,
) -> Self {
self.task_budget = task_budget;
self
}
pub fn new(store: Store, worker_config: WorkerConfig) -> Self {
Self {
store,
worker_config,
sessions: Arc::new(RwLock::new(HashMap::new())),
open_locks: Arc::new(Mutex::new(HashMap::new())),
service_tasks: Arc::new(StdMutex::new(HashMap::new())),
placement: None,
recovery_policy: crate::observation::RecoveryObservationPolicy::default(),
guidance_budget_tokens: klbr_core::context_receipt::guidance_budget_tokens(
klbr_core::config::Config::default().watermark_tokens,
),
task_budget: klbr_core::runtime_turn::budget::TaskBudget::default(),
history_window: klbr_runtime::frontier::DEFAULT_KEEP_EVENTS,
peer_host: Arc::new(std::sync::RwLock::new(None)),
}
}
/// The operator's session window, which is also how much raw history stays below a frontier.
#[must_use]
pub fn with_history_window(mut self, history_window: usize) -> Self {
self.history_window = history_window.max(1);
self
}
#[must_use]
pub fn history_window(&self) -> usize {
self.history_window
}
/// Derives the required-guidance budget from the operator's configured watermark.
///
/// Callers pass `config.watermark_tokens`; the fraction reserved for history and
/// the turn itself lives in `guidance_budget_tokens` so there is exactly one place
/// that decides it.
#[must_use]
pub fn with_guidance_budget_from_watermark(mut self, watermark_tokens: usize) -> Self {
self.guidance_budget_tokens =
klbr_core::context_receipt::guidance_budget_tokens(watermark_tokens);
self
}
/// Token budget available to required prompt guidance for this daemon.
#[must_use]
pub fn guidance_budget_tokens(&self) -> usize {
self.guidance_budget_tokens
}
/// Register the host-owned placement resolver used at session open.
pub fn with_placement(mut self, placement: Arc<dyn PlacementSource>) -> Self {
self.placement = Some(placement);
self
}
pub fn with_peer_host(self, host: Arc<dyn klbr_runtime::protocol::PeerHost>) -> Self {
*self.peer_host.write().unwrap() = Some(host);
self
}
pub fn set_peer_host(&self, host: Arc<dyn klbr_runtime::protocol::PeerHost>) {
*self.peer_host.write().unwrap() = Some(host);
}
pub fn enable_default_peer_host(
self: &Arc<Self>,
llm: Option<klbr_core::models::LlmClient>,
release: Option<Release>,
v2_tx: Option<tokio::sync::broadcast::Sender<klbr_ipc::v2::V2ServerMsg>>,
) {
let host = Arc::new(crate::peer::DaemonPeerHost::with_components(
Arc::downgrade(self),
llm,
release,
v2_tx,
));
self.set_peer_host(host);
}
pub fn with_recovery_policy(
mut self,
policy: crate::observation::RecoveryObservationPolicy,
) -> Self {
self.recovery_policy = policy;
self
}
pub fn recovery_policy(&self) -> &crate::observation::RecoveryObservationPolicy {
&self.recovery_policy
}
pub fn store(&self) -> &Store {
&self.store
}
pub fn worker_config(&self) -> &WorkerConfig {
&self.worker_config
}
pub async fn get(&self, id: &SessionId) -> Option<Arc<SessionEntry>> {
let sessions = self.sessions.read().await;
sessions.get(id).cloned()
}
/// The worker configuration this session actually runs under.
///
/// A peer session is admitted with its own workspace, and that workspace is
/// part of the contract: the recipient must run in the directory its peer
/// record names, not in whichever directory the daemon was started from. A
/// recorded workspace that is no longer there is a typed refusal, never a
/// silent fallback into the daemon's own tree. The peer record is durable, so
/// the same rule holds after a restart.
async fn worker_config_for(&self, id: &SessionId) -> Result<WorkerConfig> {
let peer = self
.store
.get_peer_session(&id.to_string())
.await
.context("failed to read peer session record")?;
let workspace = match peer.and_then(|record| record.workspace) {
Some(workspace) if !workspace.is_empty() => workspace,
_ => return Ok(self.worker_config.clone()),
};
let path = PathBuf::from(&workspace);
ensure!(
path.is_dir(),
"session {id} is a peer whose recorded workspace {workspace} does not exist; refusing to run it in the daemon's own tree"
);
Ok(WorkerConfig {
cwd: path,
..self.worker_config.clone()
})
}
/// Reopen a session the operator stopped, as a fresh runtime.
///
/// `stop` is durable and admission refuses a stopped session, so without this
/// the session stays unreachable for the life of the database: no input can be
/// admitted, so no runner can start, so nothing ever clears the state. Only an
/// explicit operator reset calls this; nothing automatic reopens what an
/// operator closed.
pub async fn reopen(
&self,
id: &SessionId,
initial_release: &Release,
) -> Result<Arc<SessionEntry>> {
self.store.set_session_state(id, "ready").await?;
if let Some(previous) = self.sessions.write().await.remove(id) {
previous.abort_runner().await;
}
self.get_or_open(id, initial_release).await
}
pub async fn get_or_open(
&self,
id: &SessionId,
initial_release: &Release,
) -> Result<Arc<SessionEntry>> {
// Fast path: if session already exists in registry
{
let sessions = self.sessions.read().await;
if let Some(entry) = sessions.get(id) {
return Ok(entry.clone());
}
}
// Single-owner synchronization: acquire per-session lock so only one task
// initializes a given SessionId.
let session_mutex = {
let mut locks = self.open_locks.lock().await;
locks
.entry(id.clone())
.or_insert_with(|| Arc::new(Mutex::new(())))
.clone()
};
let _guard = session_mutex.lock().await;
// Re-check sessions after acquiring per-session lock
{
let sessions = self.sessions.read().await;
if let Some(entry) = sessions.get(id) {
return Ok(entry.clone());
}
}
// We are the single owner responsible for starting the SessionRuntime
// Durable placement first: a persisted REMOTE placement is still the
// contract. A restart without a configured backend for it is a typed
// refusal, never a silent local fallback.
let worker_config = self.worker_config_for(id).await?;
let durable = self
.store
.session_placement(id)
.await
.context("failed to read persisted session placement")?;
let runtime = match (&self.placement, durable.clone()) {
(Some(source), _) => {
let placement = source
.placement(id)
.await
.context("placement resolution failed")?;
match placement {
Some(placement) => {
if let Some(durable) = &durable {
ensure_configured_matches_durable(&placement, durable)
.context("configured placement contradicts the persisted placement")?;
}
SessionRuntime::start_with_placement(
self.store.clone(),
id.clone(),
worker_config.clone(),
initial_release.clone(),
placement,
)
.await
}
None => match durable.map(|d| d.kernel_target) {
Some(KernelTarget::Remote(_)) => bail!(
"session {id} is durably placed on a remote target but this service has no backend configured for it; refusing silent local fallback"
),
_ => {
SessionRuntime::start(
self.store.clone(),
id.clone(),
worker_config.clone(),
initial_release.clone(),
)
.await
}
},
}
}
(None, Some(durable)) if matches!(durable.kernel_target, KernelTarget::Remote(_)) => {
bail!(
"session {id} is durably placed on a remote target but no placement source is configured; refusing silent local fallback"
)
}
_ => {
SessionRuntime::start(
self.store.clone(),
id.clone(),
worker_config.clone(),
initial_release.clone(),
)
.await
}
}
.context("failed to start session runtime")?;
runtime.set_orientation_host(Arc::new(crate::maintenance::DaemonOrientationHost));
if let Some(peer_host) = self.peer_host.read().unwrap().clone() {
runtime.set_peer_host(peer_host);
}
let now_ms = chrono::Utc::now().timestamp_millis();
let initial_state = match self.store.session_state(id).await? {
Some((ref st, _)) if st == "stopped" => SessionState::Stopped,
Some((ref st, _)) if st == "draining" => SessionState::Draining,
Some((ref st, wake_at)) if st == "waiting" => SessionState::Waiting {
receipt: None,
wake_at_ms: wake_at,
reason: None,
},
_ => SessionState::Ready,
};
let mut entry_val = SessionEntry::new(id.clone(), runtime, now_ms);
entry_val.state = Arc::new(RwLock::new(initial_state));
let entry = Arc::new(entry_val);
{
let mut sessions = self.sessions.write().await;
sessions.insert(id.clone(), entry.clone());
}
// Prune the open_locks map entry if nobody else is waiting on this lock
{
let mut locks = self.open_locks.lock().await;
if Arc::strong_count(&session_mutex) <= 2 {
locks.remove(id);
}
}
Ok(entry)
}
pub async fn list(&self) -> Vec<(SessionId, SessionState)> {
let sessions = self.sessions.read().await;
let mut list = Vec::with_capacity(sessions.len());
for (id, entry) in sessions.iter() {
let state = entry.state().await;
list.push((id.clone(), state));
}
list
}
pub async fn stop(&self, id: &SessionId) -> Result<()> {
let entry = self
.get(id)
.await
.with_context(|| format!("session not found: {id}"))?;
entry.transition_to_stopped().await?;
entry.abort_runner().await;
self.store.set_session_state(id, "stopped").await?;
entry.runtime.shutdown().await;
Ok(())
}
pub async fn drain(&self, id: &SessionId) -> Result<()> {
let entry = self
.get(id)
.await
.with_context(|| format!("session not found: {id}"))?;
// If not already draining, transition to draining
if entry.state().await != SessionState::Draining {
entry.transition_to_draining().await?;
self.store.set_session_state(id, "draining").await?;
}
let _ = tokio::time::timeout(
std::time::Duration::from_secs(5),
entry.runtime.wait_for_idle(),
)
.await;
entry.abort_runner().await;
entry.transition_to_stopped().await?;
self.store.set_session_state(id, "stopped").await?;
entry.runtime.shutdown().await;
Ok(())
}
pub async fn activate(
&self,
id: &SessionId,
expected_epoch: i64,
release: Release,
) -> Result<i64> {
let entry = self
.get(id)
.await
.with_context(|| format!("session not found: {id}"))?;
entry.runtime.activate_hooks(expected_epoch, release).await
}
/// Record a service-level background task (driver, recovery scan, client
/// handler) so `shutdown_all` owns its stop-and-join boundary.
pub fn track_service_task(&self, group: ServiceTaskGroup, task: tokio::task::JoinHandle<()>) {
let mut map = self.service_tasks.lock().expect("service task map");
map.entry(group).or_default().push(task);
}
/// Owned service teardown: stop admitting work, abort and join every
/// tracked coordinator and client task, interrupt and await each session's
/// in-flight attempt, release its durable claims, and persist truthful
/// continuation state. This is service quiescence, NOT an operator Stop:
/// waiting deadlines survive, and sessions explicitly stopped via
/// `stop()` remain stopped.
pub async fn shutdown_all(&self) {
// 1. stop coordinators first so nothing schedules work during quiesce
let tracked: Vec<tokio::task::JoinHandle<()>> = {
let mut map = self.service_tasks.lock().expect("service task map");
map.drain().flat_map(|(_, v)| v).collect()
};
for task in &tracked {
task.abort();
}
// 2. quiesce every live session without holding registry locks
let entries: Vec<Arc<SessionEntry>> = {
let sessions = self.sessions.read().await;
sessions.values().cloned().collect()
};
for entry in &entries {
if let Err(e) = entry.quiesce(&self.store).await {
tracing::error!(session = %entry.id, err = %e, "session quiesce failed");
}
}
// 3. join the coordinators; no previous owner outlives this method
for task in tracked {
let _ = task.await;
}
// 4. an aborted task does not stop an in-flight blocking write, and every
// such write holds a share of the store's ownership lock. Quiesce the jobs
// started before this point. This is a step, not the end of the story: a
// caller that can still start work - the accepting loop of a daemon, for
// instance - must quiesce again once it has stopped, because a job started
// after this call is not covered by it. Ownership itself is held by the
// Store handles, so the file is handed over when the last one drops.
if let Err(e) = self.store.quiesce().await {
tracing::error!(err = %e, "store quiesce failed; a restart may be refused ownership");
}
}
}