use std::collections::HashMap;
use std::fmt;
use std::future::Future;
use std::pin::Pin;
use std::sync::{Arc, RwLock as StdRwLock};
use async_trait::async_trait;
use bamboo_domain::{
SessionActivationDisposition, SessionActivationError, SessionActivationPort, SessionInboxPort,
};
use tokio::sync::{mpsc, watch, Mutex, RwLock};
pub struct SessionActivationLaunch {
pub run_id: String,
launch: Option<Box<dyn FnOnce() + Send + 'static>>,
rollback: Option<
Box<dyn FnOnce() -> Pin<Box<dyn Future<Output = ()> + Send + 'static>> + Send + 'static>,
>,
rollback_completion: Option<Arc<RollbackCompletion>>,
}
struct RollbackCompletion {
completed: watch::Sender<bool>,
}
impl RollbackCompletion {
fn new() -> Arc<Self> {
let (completed, _receiver) = watch::channel(false);
Arc::new(Self { completed })
}
fn complete(&self) {
self.completed.send_replace(true);
}
async fn wait(&self) {
let mut completed = self.completed.subscribe();
while !*completed.borrow() {
if completed.changed().await.is_err() {
break;
}
}
}
}
impl SessionActivationLaunch {
pub fn new(run_id: impl Into<String>, launch: impl FnOnce() + Send + 'static) -> Self {
Self {
run_id: run_id.into(),
launch: Some(Box::new(launch)),
rollback: None,
rollback_completion: None,
}
}
pub fn new_with_rollback(
run_id: impl Into<String>,
launch: impl FnOnce() + Send + 'static,
rollback: impl FnOnce() + Send + 'static,
) -> Self {
Self::new_with_async_rollback(run_id, launch, move || async move {
rollback();
})
}
pub fn new_with_async_rollback<F, Fut>(
run_id: impl Into<String>,
launch: impl FnOnce() + Send + 'static,
rollback: F,
) -> Self
where
F: FnOnce() -> Fut + Send + 'static,
Fut: Future<Output = ()> + Send + 'static,
{
let rollback_completion = RollbackCompletion::new();
Self {
run_id: run_id.into(),
launch: Some(Box::new(launch)),
rollback: Some(Box::new(move || Box::pin(rollback()))),
rollback_completion: Some(rollback_completion),
}
}
fn rollback_completion(&self) -> Option<Arc<RollbackCompletion>> {
self.rollback_completion.clone()
}
fn launch(mut self) {
if let Some(launch) = self.launch.take() {
self.rollback = None;
self.rollback_completion = None;
launch();
}
}
}
impl Drop for SessionActivationLaunch {
fn drop(&mut self) {
if self.launch.is_some() {
if let Some(rollback) = self.rollback.take() {
let completion = self.rollback_completion.take();
if let Ok(runtime) = tokio::runtime::Handle::try_current() {
runtime.spawn(async move {
rollback().await;
if let Some(completion) = completion {
completion.complete();
}
});
}
}
}
}
}
impl fmt::Debug for SessionActivationLaunch {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("SessionActivationLaunch")
.field("run_id", &self.run_id)
.finish_non_exhaustive()
}
}
#[derive(Debug)]
pub enum SessionActivationReserveOutcome {
Reserved(SessionActivationLaunch),
AlreadyRunning { run_id: String },
NotFound,
NoWork,
}
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum SessionRunRegistrationError {
#[error(
"session activation owner collision for {target_session_id}: \
existing run {existing_run_id}, attempted run {attempted_run_id}"
)]
OwnerCollision {
target_session_id: String,
existing_run_id: String,
attempted_run_id: String,
},
}
impl SessionRunRegistrationError {
pub fn existing_run_id(&self) -> &str {
match self {
Self::OwnerCollision {
existing_run_id, ..
} => existing_run_id,
}
}
}
#[async_trait]
pub trait SessionActivationSpawner: Send + Sync {
async fn reserve_activation(
&self,
target_session_id: &str,
inbox_generation: u64,
) -> Result<SessionActivationReserveOutcome, SessionActivationError>;
}
#[derive(Debug, Clone)]
struct ActiveOwner {
run_id: String,
finalizing: bool,
registrations: usize,
delivery_sink: Option<mpsc::UnboundedSender<u64>>,
}
#[derive(Debug)]
struct TargetActivationState {
latest_generation: u64,
last_dispatched_generation: u64,
owner: Option<ActiveOwner>,
activation_reserved: bool,
activation_token: u64,
activation_epoch: watch::Sender<u64>,
notify: watch::Sender<u64>,
}
impl Default for TargetActivationState {
fn default() -> Self {
let (activation_epoch, _activation_receiver) = watch::channel(0);
let (notify, _receiver) = watch::channel(0);
Self {
latest_generation: 0,
last_dispatched_generation: 0,
owner: None,
activation_reserved: false,
activation_token: 0,
activation_epoch,
notify,
}
}
}
fn reserve_activation_token(state: &mut TargetActivationState) -> u64 {
state.activation_reserved = true;
state.activation_token = state.activation_token.wrapping_add(1).max(1);
state.activation_token
}
fn release_activation_token(state: &mut TargetActivationState, token: u64) -> bool {
if !state.activation_reserved || state.activation_token != token {
return false;
}
state.activation_reserved = false;
let next_epoch = (*state.activation_epoch.borrow()).wrapping_add(1);
state.activation_epoch.send_replace(next_epoch);
true
}
struct ActivationReservationLease {
router: SessionActivationRouter,
target_session_id: String,
token: u64,
recover_on_drop: bool,
armed: bool,
rollback_completion: Option<Arc<RollbackCompletion>>,
}
impl ActivationReservationLease {
fn new(
router: SessionActivationRouter,
target_session_id: &str,
token: u64,
recover_on_drop: bool,
) -> Self {
Self {
router,
target_session_id: target_session_id.to_string(),
token,
recover_on_drop,
armed: true,
rollback_completion: None,
}
}
fn wait_for_rollback(&mut self, completion: Option<Arc<RollbackCompletion>>) {
self.rollback_completion = completion;
}
fn disarm(&mut self) {
self.armed = false;
}
}
impl Drop for ActivationReservationLease {
fn drop(&mut self) {
if !self.armed {
return;
}
let router = self.router.clone();
let target_session_id = self.target_session_id.clone();
let token = self.token;
let recover_on_drop = self.recover_on_drop;
let rollback_completion = self.rollback_completion.take();
if let Ok(runtime) = tokio::runtime::Handle::try_current() {
runtime.spawn(async move {
if let Some(completion) = rollback_completion {
completion.wait().await;
}
router
.recover_cancelled_reservation(&target_session_id, token, recover_on_drop)
.await;
});
}
}
}
#[derive(Clone, Default)]
pub struct SessionActivationRouter {
states: Arc<Mutex<HashMap<String, TargetActivationState>>>,
spawner: Arc<RwLock<Option<Arc<dyn SessionActivationSpawner>>>>,
inbox: Arc<StdRwLock<Option<Arc<dyn SessionInboxPort>>>>,
}
type AbortCleanup =
Box<dyn FnOnce() -> Pin<Box<dyn Future<Output = ()> + Send + 'static>> + Send + 'static>;
pub struct SessionRunRegistration {
router: Arc<SessionActivationRouter>,
target_session_id: String,
run_id: String,
notifications: Option<watch::Receiver<u64>>,
abort_cleanup: Option<AbortCleanup>,
armed: bool,
}
impl SessionRunRegistration {
pub fn notifications_mut(&mut self) -> &mut watch::Receiver<u64> {
self.notifications
.as_mut()
.expect("run registration notifications are live before finalization")
}
pub fn set_abort_cleanup<F, Fut>(&mut self, cleanup: F)
where
F: FnOnce() -> Fut + Send + 'static,
Fut: Future<Output = ()> + Send + 'static,
{
self.abort_cleanup = Some(Box::new(move || Box::pin(cleanup())));
}
pub async fn begin_finalization(&mut self) {
self.router
.begin_finalization(&self.target_session_id, &self.run_id)
.await;
}
pub async fn abandon(mut self) {
self.armed = false;
let abort_cleanup = self.abort_cleanup.take();
let router = self.router.clone();
let target_session_id = self.target_session_id.clone();
let run_id = self.run_id.clone();
let cleanup_target_session_id = target_session_id.clone();
let cleanup_run_id = run_id.clone();
let cleanup = tokio::spawn(async move {
if let Some(cleanup) = abort_cleanup {
cleanup().await;
}
router
.cleanup_abandoned_registration(&cleanup_target_session_id, &cleanup_run_id)
.await;
});
if let Err(error) = cleanup.await {
tracing::error!(
%target_session_id,
%run_id,
%error,
"detached SessionInbox registration cleanup failed"
);
}
}
pub async fn finish(
mut self,
admitted_generation: u64,
) -> Result<Option<SessionActivationDisposition>, SessionActivationError> {
drop(self.notifications.take());
let result = self
.router
.finish_finalization(&self.target_session_id, &self.run_id, admitted_generation)
.await;
self.armed = false;
self.abort_cleanup = None;
result
}
}
impl Drop for SessionRunRegistration {
fn drop(&mut self) {
if !self.armed {
return;
}
let router = self.router.clone();
let target_session_id = self.target_session_id.clone();
let run_id = self.run_id.clone();
let abort_cleanup = self.abort_cleanup.take();
if let Ok(runtime) = tokio::runtime::Handle::try_current() {
runtime.spawn(async move {
if let Some(cleanup) = abort_cleanup {
cleanup().await;
}
router
.cleanup_abandoned_registration(&target_session_id, &run_id)
.await;
});
}
}
}
impl SessionActivationRouter {
pub fn new() -> Arc<Self> {
Arc::new(Self::default())
}
pub async fn set_spawner(&self, spawner: Arc<dyn SessionActivationSpawner>) {
*self.spawner.write().await = Some(spawner);
}
pub fn set_inbox(&self, inbox: Arc<dyn SessionInboxPort>) {
*self
.inbox
.write()
.unwrap_or_else(std::sync::PoisonError::into_inner) = Some(inbox);
}
#[cfg(test)]
pub(crate) async fn hold_state_lock_for_test(
&self,
entered: Arc<tokio::sync::Notify>,
release: Arc<tokio::sync::Notify>,
) {
let _states = self.states.lock().await;
entered.notify_one();
release.notified().await;
}
#[cfg(test)]
pub(crate) async fn install_owner_placeholder_for_test(
&self,
target_session_id: &str,
run_id: &str,
) {
let mut states = self.states.lock().await;
states
.entry(target_session_id.to_string())
.or_default()
.owner = Some(ActiveOwner {
run_id: run_id.to_string(),
finalizing: false,
registrations: 0,
delivery_sink: None,
});
}
pub async fn register_run(
self: &Arc<Self>,
target_session_id: &str,
run_id: &str,
) -> Result<SessionRunRegistration, SessionRunRegistrationError> {
loop {
let (notifications, reservation_wait) = {
let mut states = self.states.lock().await;
let state = states.entry(target_session_id.to_string()).or_default();
match state.owner.as_mut() {
Some(owner)
if owner.run_id == run_id
&& owner.registrations == 0
&& !owner.finalizing =>
{
owner.registrations = 1;
(Some(state.notify.subscribe()), None)
}
Some(owner) => {
return Err(SessionRunRegistrationError::OwnerCollision {
target_session_id: target_session_id.to_string(),
existing_run_id: owner.run_id.clone(),
attempted_run_id: run_id.to_string(),
});
}
None if state.activation_reserved => {
(None, Some(state.activation_epoch.subscribe()))
}
None => {
state.owner = Some(ActiveOwner {
run_id: run_id.to_string(),
finalizing: false,
registrations: 1,
delivery_sink: None,
});
(Some(state.notify.subscribe()), None)
}
}
};
if let Some(mut reservation_wait) = reservation_wait {
let _ = reservation_wait.changed().await;
continue;
}
return Ok(SessionRunRegistration {
router: self.clone(),
target_session_id: target_session_id.to_string(),
run_id: run_id.to_string(),
notifications,
abort_cleanup: None,
armed: true,
});
}
}
pub async fn attach_delivery_sink(
&self,
target_session_id: &str,
sink: mpsc::UnboundedSender<u64>,
) -> Option<String> {
let mut states = self.states.lock().await;
let state = states.get_mut(target_session_id)?;
let owner = state.owner.as_mut()?;
if owner.finalizing {
return None;
}
owner.delivery_sink = Some(sink.clone());
let run_id = owner.run_id.clone();
if state.latest_generation > 0 && sink.send(state.latest_generation).is_err() {
owner.delivery_sink = None;
return None;
}
Some(run_id)
}
pub async fn detach_delivery_sink(&self, target_session_id: &str, run_id: &str) {
let mut states = self.states.lock().await;
let Some(state) = states.get_mut(target_session_id) else {
return;
};
let Some(owner) = state.owner.as_mut() else {
return;
};
if owner.run_id == run_id {
owner.delivery_sink = None;
}
}
pub async fn owns_run(&self, target_session_id: &str, run_id: &str) -> bool {
let states = self.states.lock().await;
states
.get(target_session_id)
.and_then(|state| state.owner.as_ref())
.is_some_and(|owner| owner.run_id == run_id && !owner.finalizing)
}
pub async fn begin_finalization(&self, target_session_id: &str, run_id: &str) {
let mut states = self.states.lock().await;
let state = states.entry(target_session_id.to_string()).or_default();
match state.owner.as_mut() {
Some(owner) if owner.run_id == run_id => owner.finalizing = true,
Some(_) => {}
None => {
state.owner = Some(ActiveOwner {
run_id: run_id.to_string(),
finalizing: true,
registrations: 0,
delivery_sink: None,
});
}
}
}
async fn cleanup_abandoned_registration(
self: &Arc<Self>,
target_session_id: &str,
run_id: &str,
) {
{
let mut states = self.states.lock().await;
let Some(state) = states.get_mut(target_session_id) else {
return;
};
let Some(owner) = state.owner.as_mut() else {
return;
};
if owner.run_id != run_id {
return;
}
owner.registrations = 0;
owner.finalizing = true;
}
let inbox = self
.inbox
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.clone();
let durable_generation = match inbox {
Some(inbox) => match inbox.inspect(target_session_id).await {
Ok(backlog) if backlog.activation_pending() => backlog.activation_generation,
Ok(_) => 0,
Err(error) => {
tracing::warn!(
%target_session_id,
%run_id,
%error,
"failed to reconcile durable SessionInbox generation for abandoned run"
);
0
}
},
None => 0,
};
let reservation_to_dispatch = {
let mut states = self.states.lock().await;
let Some(state) = states.get_mut(target_session_id) else {
return;
};
if !state.owner.as_ref().is_some_and(|owner| {
owner.run_id == run_id && owner.finalizing && owner.registrations == 0
}) {
return;
}
state.latest_generation = state.latest_generation.max(durable_generation);
state.owner = None;
if state.latest_generation > state.last_dispatched_generation
&& !state.activation_reserved
{
let generation = state.latest_generation;
let token = reserve_activation_token(state);
Some((generation, token))
} else {
if state.latest_generation == 0
&& !state.activation_reserved
&& state.notify.receiver_count() == 0
&& state.activation_epoch.receiver_count() == 0
{
states.remove(target_session_id);
}
None
}
};
if let Some((generation, token)) = reservation_to_dispatch {
if let Err(error) = self
.dispatch_reserved(target_session_id, generation, token)
.await
{
tracing::error!(
%target_session_id,
%run_id,
%error,
"failed to reserve successor for abandoned SessionInbox owner"
);
}
}
}
pub async fn finish_finalization(
&self,
target_session_id: &str,
run_id: &str,
admitted_generation: u64,
) -> Result<Option<SessionActivationDisposition>, SessionActivationError> {
let reservation_to_dispatch = {
let mut states = self.states.lock().await;
let state = states.entry(target_session_id.to_string()).or_default();
if state
.owner
.as_ref()
.is_some_and(|owner| owner.run_id != run_id)
{
return Ok(None);
}
state.owner = None;
if state.latest_generation > admitted_generation
&& state.latest_generation > state.last_dispatched_generation
&& !state.activation_reserved
{
let generation = state.latest_generation;
let token = reserve_activation_token(state);
Some((generation, token))
} else {
if state.latest_generation <= admitted_generation
&& state.owner.is_none()
&& !state.activation_reserved
&& state.notify.receiver_count() == 0
&& state.activation_epoch.receiver_count() == 0
{
states.remove(target_session_id);
}
None
}
};
match reservation_to_dispatch {
Some((generation, token)) => self
.dispatch_reserved(target_session_id, generation, token)
.await
.map(Some),
None => Ok(None),
}
}
pub async fn subscribe(&self, target_session_id: &str) -> watch::Receiver<u64> {
let mut states = self.states.lock().await;
states
.entry(target_session_id.to_string())
.or_default()
.notify
.subscribe()
}
async fn recover_cancelled_reservation(
&self,
target_session_id: &str,
reservation_token: u64,
recover_on_drop: bool,
) {
let reservation_to_dispatch = {
let mut states = self.states.lock().await;
let Some(state) = states.get_mut(target_session_id) else {
return;
};
if !release_activation_token(state, reservation_token) || !recover_on_drop {
return;
}
if state.owner.is_none()
&& state.latest_generation > state.last_dispatched_generation
&& !state.activation_reserved
{
let generation = state.latest_generation;
let token = reserve_activation_token(state);
Some((generation, token))
} else {
None
}
};
if let Some((generation, token)) = reservation_to_dispatch {
if let Err(error) = self
.dispatch_reserved_with_recovery(target_session_id, generation, token, false)
.await
{
tracing::error!(
%target_session_id,
%error,
"failed bounded retry after cancelled SessionInbox activation reservation"
);
}
}
}
async fn dispatch_reserved(
&self,
target_session_id: &str,
generation: u64,
reservation_token: u64,
) -> Result<SessionActivationDisposition, SessionActivationError> {
self.dispatch_reserved_with_recovery(target_session_id, generation, reservation_token, true)
.await
}
async fn dispatch_reserved_with_recovery(
&self,
target_session_id: &str,
generation: u64,
reservation_token: u64,
recover_on_drop: bool,
) -> Result<SessionActivationDisposition, SessionActivationError> {
let mut lease = ActivationReservationLease::new(
self.clone(),
target_session_id,
reservation_token,
recover_on_drop,
);
let Some(spawner) = self.spawner.read().await.clone() else {
let mut states = self.states.lock().await;
if let Some(state) = states.get_mut(target_session_id) {
release_activation_token(state, reservation_token);
}
lease.disarm();
return Err(SessionActivationError::Internal(
"session activation spawner is not configured".to_string(),
));
};
let outcome = spawner
.reserve_activation(target_session_id, generation)
.await;
if let Ok(SessionActivationReserveOutcome::Reserved(launch)) = &outcome {
lease.wait_for_rollback(launch.rollback_completion());
}
match outcome {
Ok(SessionActivationReserveOutcome::Reserved(launch)) => {
let run_id = launch.run_id.clone();
{
let mut states = self.states.lock().await;
let state = states.entry(target_session_id.to_string()).or_default();
if !release_activation_token(state, reservation_token) {
lease.disarm();
return if state.owner.is_some() {
Ok(SessionActivationDisposition::ActiveNotified)
} else {
Err(SessionActivationError::Internal(
"session activation reservation was superseded".to_string(),
))
};
}
state.owner = Some(ActiveOwner {
run_id,
finalizing: false,
registrations: 0,
delivery_sink: None,
});
state.last_dispatched_generation =
state.last_dispatched_generation.max(generation);
}
lease.disarm();
launch.launch();
Ok(SessionActivationDisposition::ActivationReserved)
}
Ok(SessionActivationReserveOutcome::AlreadyRunning { run_id }) => {
let mut states = self.states.lock().await;
let state = states.entry(target_session_id.to_string()).or_default();
if !release_activation_token(state, reservation_token) {
lease.disarm();
return Ok(SessionActivationDisposition::ActiveNotified);
}
state.owner = Some(ActiveOwner {
run_id,
finalizing: false,
registrations: 0,
delivery_sink: None,
});
state.notify.send_replace(generation);
lease.disarm();
Ok(SessionActivationDisposition::ActiveNotified)
}
Ok(SessionActivationReserveOutcome::NoWork) => {
let mut states = self.states.lock().await;
if let Some(state) = states.get_mut(target_session_id) {
release_activation_token(state, reservation_token);
}
lease.disarm();
Ok(SessionActivationDisposition::ActivationCoalesced)
}
Ok(SessionActivationReserveOutcome::NotFound) => {
let mut states = self.states.lock().await;
if states
.get(target_session_id)
.is_some_and(|state| state.activation_token == reservation_token)
{
states.remove(target_session_id);
}
lease.disarm();
Err(SessionActivationError::TargetNotFound(
target_session_id.to_string(),
))
}
Err(error) => {
let mut states = self.states.lock().await;
if let Some(state) = states.get_mut(target_session_id) {
release_activation_token(state, reservation_token);
}
lease.disarm();
Err(error)
}
}
}
}
#[async_trait]
impl SessionActivationPort for SessionActivationRouter {
async fn request_activation(
&self,
target_session_id: &str,
inbox_generation: u64,
) -> Result<SessionActivationDisposition, SessionActivationError> {
loop {
let (reservation_wait, reservation_token) = {
let mut states = self.states.lock().await;
let state = states.entry(target_session_id.to_string()).or_default();
state.latest_generation = state.latest_generation.max(inbox_generation);
if let Some(owner) = state.owner.as_mut() {
if !owner.finalizing {
state.notify.send_replace(inbox_generation);
if owner
.delivery_sink
.as_ref()
.is_some_and(|sink| sink.send(inbox_generation).is_err())
{
owner.delivery_sink = None;
}
return Ok(SessionActivationDisposition::ActiveNotified);
}
return Ok(SessionActivationDisposition::ActivationCoalesced);
}
if inbox_generation <= state.last_dispatched_generation {
return Ok(SessionActivationDisposition::ActivationCoalesced);
}
if state.activation_reserved {
(Some(state.activation_epoch.subscribe()), None)
} else {
let token = reserve_activation_token(state);
(None, Some(token))
}
};
if let Some(mut reservation_wait) = reservation_wait {
if reservation_wait.changed().await.is_err() {
return Err(SessionActivationError::TargetNotFound(
target_session_id.to_string(),
));
}
continue;
}
let reservation_token =
reservation_token.expect("non-waiting activation owns a reservation token");
return self
.dispatch_reserved_with_recovery(
target_session_id,
inbox_generation,
reservation_token,
false,
)
.await;
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::{AtomicUsize, Ordering};
use tokio::sync::{Barrier, Notify};
struct RecordingSpawner {
reservations: AtomicUsize,
launches: Arc<AtomicUsize>,
entered: Option<Arc<Barrier>>,
release: Option<Arc<Notify>>,
}
#[async_trait]
impl SessionActivationSpawner for RecordingSpawner {
async fn reserve_activation(
&self,
_target_session_id: &str,
_inbox_generation: u64,
) -> Result<SessionActivationReserveOutcome, SessionActivationError> {
let ordinal = self.reservations.fetch_add(1, Ordering::SeqCst) + 1;
if let Some(entered) = &self.entered {
entered.wait().await;
}
if let Some(release) = &self.release {
release.notified().await;
}
let launches = self.launches.clone();
Ok(SessionActivationReserveOutcome::Reserved(
SessionActivationLaunch::new(format!("run-{ordinal}"), move || {
launches.fetch_add(1, Ordering::SeqCst);
}),
))
}
}
fn spawner() -> Arc<RecordingSpawner> {
Arc::new(RecordingSpawner {
reservations: AtomicUsize::new(0),
launches: Arc::new(AtomicUsize::new(0)),
entered: None,
release: None,
})
}
struct BlockingInspectInbox {
entered: Arc<Notify>,
release: Arc<Notify>,
}
#[async_trait]
impl SessionInboxPort for BlockingInspectInbox {
async fn deliver(
&self,
_envelope: &bamboo_domain::SessionMessageEnvelope,
) -> Result<bamboo_domain::SessionInboxReceipt, bamboo_domain::SessionInboxError> {
unreachable!("router cleanup only inspects the durable inbox")
}
async fn mark_activation_eligible(
&self,
_target_session_id: &str,
_generation: u64,
_policy: bamboo_domain::SessionActivationPolicy,
) -> Result<(), bamboo_domain::SessionInboxError> {
unreachable!("router cleanup only inspects the durable inbox")
}
async fn claim(
&self,
_target_session_id: &str,
_limit: usize,
) -> Result<Vec<bamboo_domain::SessionInboxClaim>, bamboo_domain::SessionInboxError>
{
unreachable!("router cleanup only inspects the durable inbox")
}
async fn was_admitted(
&self,
_target_session_id: &str,
_id: &bamboo_domain::SessionMessageId,
) -> Result<bool, bamboo_domain::SessionInboxError> {
unreachable!("router cleanup only inspects the durable inbox")
}
async fn ack(
&self,
_target_session_id: &str,
_claim: &bamboo_domain::SessionInboxClaim,
) -> Result<(), bamboo_domain::SessionInboxError> {
unreachable!("router cleanup only inspects the durable inbox")
}
async fn inspect(
&self,
_target_session_id: &str,
) -> Result<bamboo_domain::SessionInboxBacklog, bamboo_domain::SessionInboxError> {
self.entered.notify_one();
self.release.notified().await;
Ok(bamboo_domain::SessionInboxBacklog {
pending: 1,
claimed: 0,
generation: 1,
activation_generation: 1,
interrupt_generation: 1,
oldest_generation: Some(1),
})
}
}
#[tokio::test]
async fn rollback_completion_is_retained_before_waiter_subscribes() {
let completion = RollbackCompletion::new();
completion.complete();
tokio::time::timeout(std::time::Duration::from_millis(100), completion.wait())
.await
.expect("completion-before-wait must not lose its wake");
}
#[derive(Clone, Copy)]
enum InjectedFirstOutcome {
NoWork,
Error,
}
struct RetrySpawner {
reservations: AtomicUsize,
launches: Arc<AtomicUsize>,
first_outcome: InjectedFirstOutcome,
first_entered: Arc<Barrier>,
release_first: Arc<Notify>,
}
struct CancelFirstSpawner {
reservations: AtomicUsize,
launches: Arc<AtomicUsize>,
first_entered: Arc<Notify>,
}
struct PanicFirstSpawner {
reservations: AtomicUsize,
launches: Arc<AtomicUsize>,
}
struct AlreadyRunningThenReserveSpawner {
reservations: AtomicUsize,
launches: Arc<AtomicUsize>,
}
#[async_trait]
impl SessionActivationSpawner for AlreadyRunningThenReserveSpawner {
async fn reserve_activation(
&self,
_target_session_id: &str,
_inbox_generation: u64,
) -> Result<SessionActivationReserveOutcome, SessionActivationError> {
let ordinal = self.reservations.fetch_add(1, Ordering::SeqCst) + 1;
if ordinal == 1 {
return Ok(SessionActivationReserveOutcome::AlreadyRunning {
run_id: "adopted-run".to_string(),
});
}
let launches = self.launches.clone();
Ok(SessionActivationReserveOutcome::Reserved(
SessionActivationLaunch::new("successor-run", move || {
launches.fetch_add(1, Ordering::SeqCst);
}),
))
}
}
#[async_trait]
impl SessionActivationSpawner for CancelFirstSpawner {
async fn reserve_activation(
&self,
_target_session_id: &str,
_inbox_generation: u64,
) -> Result<SessionActivationReserveOutcome, SessionActivationError> {
let ordinal = self.reservations.fetch_add(1, Ordering::SeqCst) + 1;
if ordinal == 1 {
self.first_entered.notify_one();
std::future::pending::<()>().await;
}
let launches = self.launches.clone();
Ok(SessionActivationReserveOutcome::Reserved(
SessionActivationLaunch::new(format!("run-{ordinal}"), move || {
launches.fetch_add(1, Ordering::SeqCst);
}),
))
}
}
#[async_trait]
impl SessionActivationSpawner for PanicFirstSpawner {
async fn reserve_activation(
&self,
_target_session_id: &str,
_inbox_generation: u64,
) -> Result<SessionActivationReserveOutcome, SessionActivationError> {
let ordinal = self.reservations.fetch_add(1, Ordering::SeqCst) + 1;
if ordinal == 1 {
panic!("injected first activation-spawner panic");
}
let launches = self.launches.clone();
Ok(SessionActivationReserveOutcome::Reserved(
SessionActivationLaunch::new(format!("run-{ordinal}"), move || {
launches.fetch_add(1, Ordering::SeqCst);
}),
))
}
}
struct RollbackSpawner {
entered: Arc<Barrier>,
release: Arc<Notify>,
launch_ready: Arc<Notify>,
rollbacks: Arc<AtomicUsize>,
}
struct RealRegistryRollbackSpawner {
runners: Arc<
tokio::sync::RwLock<std::collections::HashMap<String, crate::execution::AgentRunner>>,
>,
senders: Arc<
tokio::sync::RwLock<
std::collections::HashMap<
String,
tokio::sync::broadcast::Sender<bamboo_agent_core::AgentEvent>,
>,
>,
>,
sender: tokio::sync::broadcast::Sender<bamboo_agent_core::AgentEvent>,
reservations: AtomicUsize,
first_reserved: Arc<tokio::sync::Notify>,
allow_first_return: Arc<tokio::sync::Notify>,
first_returning: Arc<tokio::sync::Notify>,
rollback_started: Arc<tokio::sync::Notify>,
allow_rollback: Arc<tokio::sync::Notify>,
launches: Arc<AtomicUsize>,
}
#[async_trait]
impl SessionActivationSpawner for RealRegistryRollbackSpawner {
async fn reserve_activation(
&self,
target_session_id: &str,
_inbox_generation: u64,
) -> Result<SessionActivationReserveOutcome, SessionActivationError> {
let ordinal = self.reservations.fetch_add(1, Ordering::SeqCst) + 1;
match crate::execution::reserve_runner_core(
&self.runners,
&self.senders,
target_session_id,
&self.sender,
)
.await
{
crate::execution::ReserveOutcome::AlreadyRunning(run_id) => {
Ok(SessionActivationReserveOutcome::AlreadyRunning { run_id })
}
crate::execution::ReserveOutcome::Reserved(reservation) => {
let run_id = reservation.run_id.clone();
if ordinal == 1 {
self.first_reserved.notify_one();
self.allow_first_return.notified().await;
self.first_returning.notify_one();
}
let rollback_runners = self.runners.clone();
let rollback_session_id = target_session_id.to_string();
let rollback_run_id = run_id.clone();
let rollback_started = self.rollback_started.clone();
let allow_rollback = self.allow_rollback.clone();
let launches = self.launches.clone();
Ok(SessionActivationReserveOutcome::Reserved(
SessionActivationLaunch::new_with_async_rollback(
run_id,
move || {
launches.fetch_add(1, Ordering::SeqCst);
},
move || async move {
rollback_started.notify_one();
allow_rollback.notified().await;
let mut runners = rollback_runners.write().await;
if runners
.get(&rollback_session_id)
.is_some_and(|runner| runner.run_id == rollback_run_id)
{
runners.remove(&rollback_session_id);
}
},
),
))
}
}
}
}
#[async_trait]
impl SessionActivationSpawner for RollbackSpawner {
async fn reserve_activation(
&self,
_target_session_id: &str,
_inbox_generation: u64,
) -> Result<SessionActivationReserveOutcome, SessionActivationError> {
self.entered.wait().await;
self.release.notified().await;
let rollbacks = self.rollbacks.clone();
let outcome = SessionActivationReserveOutcome::Reserved(
SessionActivationLaunch::new_with_rollback(
"reserved-run",
|| {},
move || {
rollbacks.fetch_add(1, Ordering::SeqCst);
},
),
);
self.launch_ready.notify_one();
Ok(outcome)
}
}
#[async_trait]
impl SessionActivationSpawner for RetrySpawner {
async fn reserve_activation(
&self,
_target_session_id: &str,
_inbox_generation: u64,
) -> Result<SessionActivationReserveOutcome, SessionActivationError> {
let ordinal = self.reservations.fetch_add(1, Ordering::SeqCst) + 1;
if ordinal == 1 {
self.first_entered.wait().await;
self.release_first.notified().await;
return match self.first_outcome {
InjectedFirstOutcome::NoWork => Ok(SessionActivationReserveOutcome::NoWork),
InjectedFirstOutcome::Error => Err(SessionActivationError::Internal(
"injected first reservation failure".to_string(),
)),
};
}
let launches = self.launches.clone();
Ok(SessionActivationReserveOutcome::Reserved(
SessionActivationLaunch::new(format!("run-{ordinal}"), move || {
launches.fetch_add(1, Ordering::SeqCst);
}),
))
}
}
async fn assert_newer_generation_retries_after(
first_outcome: InjectedFirstOutcome,
) -> Result<SessionActivationDisposition, SessionActivationError> {
let router = SessionActivationRouter::new();
let first_entered = Arc::new(Barrier::new(2));
let release_first = Arc::new(Notify::new());
let spawner = Arc::new(RetrySpawner {
reservations: AtomicUsize::new(0),
launches: Arc::new(AtomicUsize::new(0)),
first_outcome,
first_entered: first_entered.clone(),
release_first: release_first.clone(),
});
router.set_spawner(spawner.clone()).await;
let first_router = router.clone();
let first =
tokio::spawn(async move { first_router.request_activation("session", 1).await });
first_entered.wait().await;
let second_router = router.clone();
let second =
tokio::spawn(async move { second_router.request_activation("session", 2).await });
tokio::task::yield_now().await;
assert_eq!(
spawner.reservations.load(Ordering::SeqCst),
1,
"generation 2 must coalesce while generation 1 owns the reservation"
);
release_first.notify_one();
let first_result = first.await.unwrap();
assert_eq!(
second.await.unwrap().unwrap(),
SessionActivationDisposition::ActivationReserved,
"the coalesced newer delivery must take the released reservation"
);
assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
first_result
}
#[tokio::test]
async fn active_owner_is_notified_without_second_reservation() {
let router = SessionActivationRouter::new();
let spawner = spawner();
router.set_spawner(spawner.clone()).await;
let mut registration = router.register_run("session", "run-live").await.unwrap();
assert_eq!(
router.request_activation("session", 7).await.unwrap(),
SessionActivationDisposition::ActiveNotified
);
registration.notifications_mut().changed().await.unwrap();
assert_eq!(*registration.notifications_mut().borrow(), 7);
assert_eq!(spawner.reservations.load(Ordering::SeqCst), 0);
registration.begin_finalization().await;
assert_eq!(registration.finish(7).await.unwrap(), None);
}
#[tokio::test]
async fn different_run_registration_never_overwrites_live_owner() {
let router = SessionActivationRouter::new();
let mut first = router
.register_run("session", "run-1")
.await
.expect("first run owns the logical session");
let error = router
.register_run("session", "run-2")
.await
.err()
.expect("a second exact run must collide");
assert_eq!(
error,
SessionRunRegistrationError::OwnerCollision {
target_session_id: "session".to_string(),
existing_run_id: "run-1".to_string(),
attempted_run_id: "run-2".to_string(),
}
);
assert!(router.owns_run("session", "run-1").await);
assert!(!router.owns_run("session", "run-2").await);
first.begin_finalization().await;
assert_eq!(first.finish(0).await.unwrap(), None);
}
#[tokio::test]
async fn direct_registration_waits_for_activation_to_publish_its_exact_owner() {
let router = SessionActivationRouter::new();
let reservation_entered = Arc::new(Barrier::new(2));
let allow_reservation = Arc::new(Notify::new());
let launches = Arc::new(AtomicUsize::new(0));
let spawner = Arc::new(RecordingSpawner {
reservations: AtomicUsize::new(0),
launches: launches.clone(),
entered: Some(reservation_entered.clone()),
release: Some(allow_reservation.clone()),
});
router.set_spawner(spawner.clone()).await;
let activation_router = router.clone();
let activation =
tokio::spawn(async move { activation_router.request_activation("session", 1).await });
reservation_entered.wait().await;
let direct_router = router.clone();
let direct =
tokio::spawn(async move { direct_router.register_run("session", "manual-run").await });
tokio::task::yield_now().await;
assert!(
!direct.is_finished(),
"manual registration must not supersede an in-flight activation token"
);
assert!(!router.owns_run("session", "manual-run").await);
allow_reservation.notify_one();
assert_eq!(
activation.await.unwrap().unwrap(),
SessionActivationDisposition::ActivationReserved
);
let collision = match direct.await.unwrap() {
Ok(_) => panic!("manual registration must collide with the published activation"),
Err(error) => error,
};
assert_eq!(
collision,
SessionRunRegistrationError::OwnerCollision {
target_session_id: "session".to_string(),
existing_run_id: "run-1".to_string(),
attempted_run_id: "manual-run".to_string(),
}
);
assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
assert_eq!(launches.load(Ordering::SeqCst), 1);
assert!(router.owns_run("session", "run-1").await);
assert!(!router.owns_run("session", "manual-run").await);
let mut registration = router.register_run("session", "run-1").await.unwrap();
registration.begin_finalization().await;
assert_eq!(registration.finish(1).await.unwrap(), None);
}
#[tokio::test]
async fn cancelled_explicit_abandon_still_finishes_owned_registration_cleanup() {
let router = SessionActivationRouter::new();
let cleanup_entered = Arc::new(Notify::new());
let cleanup_release = Arc::new(Notify::new());
let cleanup_completed = Arc::new(AtomicUsize::new(0));
let mut registration = router.register_run("session", "run-abandon").await.unwrap();
let entered = cleanup_entered.clone();
let release = cleanup_release.clone();
let completed = cleanup_completed.clone();
registration.set_abort_cleanup(move || async move {
entered.notify_one();
release.notified().await;
completed.fetch_add(1, Ordering::SeqCst);
});
let abandon = tokio::spawn(registration.abandon());
tokio::time::timeout(
std::time::Duration::from_secs(1),
cleanup_entered.notified(),
)
.await
.expect("explicit abandon must start its owned cleanup");
abandon.abort();
assert!(abandon.await.unwrap_err().is_cancelled());
cleanup_release.notify_one();
tokio::time::timeout(std::time::Duration::from_secs(1), async {
loop {
if cleanup_completed.load(Ordering::SeqCst) == 1
&& !router.owns_run("session", "run-abandon").await
{
break;
}
tokio::task::yield_now().await;
}
})
.await
.expect("caller cancellation must not cancel detached registration cleanup");
}
#[tokio::test]
async fn delayed_same_run_cannot_adopt_an_owner_during_abandoned_cleanup() {
let router = SessionActivationRouter::new();
let spawner = spawner();
router.set_spawner(spawner.clone()).await;
let inspect_entered = Arc::new(Notify::new());
let inspect_release = Arc::new(Notify::new());
router.set_inbox(Arc::new(BlockingInspectInbox {
entered: inspect_entered.clone(),
release: inspect_release.clone(),
}));
let registration = router
.register_run("session", "run-abandoned")
.await
.unwrap();
drop(registration);
tokio::time::timeout(
std::time::Duration::from_secs(1),
inspect_entered.notified(),
)
.await
.expect("abandoned cleanup must reach durable inbox reconciliation");
let error = router
.register_run("session", "run-abandoned")
.await
.err()
.expect("a finalizing abandoned owner is not an adoptable launch placeholder");
assert_eq!(error.existing_run_id(), "run-abandoned");
inspect_release.notify_one();
tokio::time::timeout(std::time::Duration::from_secs(1), async {
while spawner.launches.load(Ordering::SeqCst) == 0 {
tokio::task::yield_now().await;
}
})
.await
.expect("durable pending work must launch one successor");
assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
assert!(!router.owns_run("session", "run-abandoned").await);
let mut successor = router.register_run("session", "run-1").await.unwrap();
successor.begin_finalization().await;
assert_eq!(successor.finish(1).await.unwrap(), None);
}
#[tokio::test]
async fn delayed_registry_owner_collision_rolls_back_only_its_exact_runner() {
let router = SessionActivationRouter::new();
let runners = Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new()));
let senders = Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new()));
let (sender, _receiver) = tokio::sync::broadcast::channel(8);
let reserved =
match crate::execution::reserve_runner_core(&runners, &senders, "session", &sender)
.await
{
crate::execution::ReserveOutcome::Reserved(reservation) => reservation,
crate::execution::ReserveOutcome::AlreadyRunning(_) => {
panic!("fixture must reserve a fresh server runner")
}
};
let mut direct = router.register_run("session", "sdk-direct").await.unwrap();
let collision = router
.register_run("session", &reserved.run_id)
.await
.err()
.expect("delayed server registration must collide");
let collision_result = Err(bamboo_agent_core::AgentError::LLM(collision.to_string()));
assert!(
crate::execution::finalize_runner_exact(
&runners,
"session",
&reserved.run_id,
&collision_result,
)
.await
);
assert!(router.owns_run("session", "sdk-direct").await);
assert!(matches!(
runners
.read()
.await
.get("session")
.map(|runner| &runner.status),
Some(crate::execution::AgentStatus::Error(_))
));
direct.begin_finalization().await;
assert_eq!(direct.finish(0).await.unwrap(), None);
}
#[tokio::test]
async fn concurrent_idle_deliveries_coalesce_into_one_launch() {
let router = SessionActivationRouter::new();
let entered = Arc::new(Barrier::new(2));
let release = Arc::new(Notify::new());
let spawner = Arc::new(RecordingSpawner {
reservations: AtomicUsize::new(0),
launches: Arc::new(AtomicUsize::new(0)),
entered: Some(entered.clone()),
release: Some(release.clone()),
});
router.set_spawner(spawner.clone()).await;
let first_router = router.clone();
let first =
tokio::spawn(async move { first_router.request_activation("session", 1).await });
entered.wait().await;
let second_router = router.clone();
let second =
tokio::spawn(async move { second_router.request_activation("session", 2).await });
release.notify_one();
assert_eq!(
first.await.unwrap().unwrap(),
SessionActivationDisposition::ActivationReserved
);
assert_eq!(
second.await.unwrap().unwrap(),
SessionActivationDisposition::ActiveNotified
);
assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn newer_generation_re_reserves_after_coalesced_no_work() {
assert_eq!(
assert_newer_generation_retries_after(InjectedFirstOutcome::NoWork)
.await
.unwrap(),
SessionActivationDisposition::ActivationCoalesced
);
}
#[tokio::test]
async fn newer_generation_re_reserves_after_coalesced_error() {
let error = assert_newer_generation_retries_after(InjectedFirstOutcome::Error)
.await
.unwrap_err();
assert!(matches!(error, SessionActivationError::Internal(_)));
}
async fn assert_cancelled_finalization_recovers_without_redelivery(registered_owner: bool) {
let router = SessionActivationRouter::new();
let first_entered = Arc::new(Notify::new());
let spawner = Arc::new(CancelFirstSpawner {
reservations: AtomicUsize::new(0),
launches: Arc::new(AtomicUsize::new(0)),
first_entered: first_entered.clone(),
});
router.set_spawner(spawner.clone()).await;
let finish = if registered_owner {
let mut registration = router.register_run("session", "run-old").await.unwrap();
registration.begin_finalization().await;
assert_eq!(
router.request_activation("session", 1).await.unwrap(),
SessionActivationDisposition::ActivationCoalesced
);
tokio::spawn(async move { registration.finish(0).await })
} else {
router.begin_finalization("session", "run-old").await;
assert_eq!(
router.request_activation("session", 1).await.unwrap(),
SessionActivationDisposition::ActivationCoalesced
);
let finish_router = router.clone();
tokio::spawn(async move {
finish_router
.finish_finalization("session", "run-old", 0)
.await
})
};
first_entered.notified().await;
finish.abort();
assert!(finish.await.unwrap_err().is_cancelled());
tokio::time::timeout(std::time::Duration::from_secs(1), async {
while spawner.launches.load(Ordering::SeqCst) == 0 {
tokio::task::yield_now().await;
}
})
.await
.expect("cancelled finalization must retry pending work without redelivery");
assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
assert!(router.owns_run("session", "run-2").await);
let mut successor = router.register_run("session", "run-2").await.unwrap();
successor.begin_finalization().await;
assert_eq!(successor.finish(1).await.unwrap(), None);
}
#[tokio::test]
async fn cancelled_registered_finalization_retries_without_another_delivery() {
assert_cancelled_finalization_recovers_without_redelivery(true).await;
}
#[tokio::test]
async fn cancelled_reserved_child_finalization_retries_without_another_delivery() {
assert_cancelled_finalization_recovers_without_redelivery(false).await;
}
#[tokio::test]
async fn panicking_finalization_spawner_gets_one_bounded_retry() {
let router = SessionActivationRouter::new();
let spawner = Arc::new(PanicFirstSpawner {
reservations: AtomicUsize::new(0),
launches: Arc::new(AtomicUsize::new(0)),
});
router.set_spawner(spawner.clone()).await;
router.begin_finalization("session", "run-old").await;
assert_eq!(
router.request_activation("session", 1).await.unwrap(),
SessionActivationDisposition::ActivationCoalesced
);
let finish_router = router.clone();
let finish = tokio::spawn(async move {
finish_router
.finish_finalization("session", "run-old", 0)
.await
});
assert!(finish.await.unwrap_err().is_panic());
tokio::time::timeout(std::time::Duration::from_secs(1), async {
while spawner.launches.load(Ordering::SeqCst) == 0 {
tokio::task::yield_now().await;
}
})
.await
.expect("router Drop recovery must survive one spawner panic");
assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
let mut successor = router.register_run("session", "run-2").await.unwrap();
successor.begin_finalization().await;
assert_eq!(successor.finish(1).await.unwrap(), None);
}
#[tokio::test]
async fn cancelled_reservation_releases_token_and_newer_generation_launches() {
let router = SessionActivationRouter::new();
let first_entered = Arc::new(Notify::new());
let spawner = Arc::new(CancelFirstSpawner {
reservations: AtomicUsize::new(0),
launches: Arc::new(AtomicUsize::new(0)),
first_entered: first_entered.clone(),
});
router.set_spawner(spawner.clone()).await;
let first_router = router.clone();
let first =
tokio::spawn(async move { first_router.request_activation("session", 1).await });
first_entered.notified().await;
let second_router = router.clone();
let second =
tokio::spawn(async move { second_router.request_activation("session", 2).await });
tokio::task::yield_now().await;
first.abort();
assert!(first.await.unwrap_err().is_cancelled());
assert_eq!(
tokio::time::timeout(std::time::Duration::from_secs(1), second)
.await
.expect("coalesced caller must not hang behind a cancelled reservation")
.unwrap()
.unwrap(),
SessionActivationDisposition::ActivationReserved
);
assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn cancellation_after_external_reservation_rolls_back_unlaunched_slot() {
let router = SessionActivationRouter::new();
let entered = Arc::new(Barrier::new(2));
let release = Arc::new(Notify::new());
let launch_ready = Arc::new(Notify::new());
let rollbacks = Arc::new(AtomicUsize::new(0));
let spawner = Arc::new(RollbackSpawner {
entered: entered.clone(),
release: release.clone(),
launch_ready: launch_ready.clone(),
rollbacks: rollbacks.clone(),
});
router.set_spawner(spawner).await;
let task_router = router.clone();
let task = tokio::spawn(async move { task_router.request_activation("session", 1).await });
entered.wait().await;
let states = router.states.lock().await;
release.notify_one();
launch_ready.notified().await;
task.abort();
assert!(task.await.unwrap_err().is_cancelled());
drop(states);
tokio::time::timeout(std::time::Duration::from_secs(1), async {
while rollbacks.load(Ordering::SeqCst) == 0 {
tokio::task::yield_now().await;
}
})
.await
.expect("unlaunched reservation rollback");
assert_eq!(rollbacks.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn redelivery_waits_for_real_runner_rollback_before_reserving() {
let router = SessionActivationRouter::new();
let runners = Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new()));
let senders = Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new()));
let (sender, _receiver) = tokio::sync::broadcast::channel(8);
let first_reserved = Arc::new(tokio::sync::Notify::new());
let allow_first_return = Arc::new(tokio::sync::Notify::new());
let rollback_started = Arc::new(tokio::sync::Notify::new());
let first_returning = Arc::new(tokio::sync::Notify::new());
let allow_rollback = Arc::new(tokio::sync::Notify::new());
let launches = Arc::new(AtomicUsize::new(0));
let spawner = Arc::new(RealRegistryRollbackSpawner {
runners: runners.clone(),
senders,
sender,
reservations: AtomicUsize::new(0),
first_reserved: first_reserved.clone(),
allow_first_return: allow_first_return.clone(),
first_returning: first_returning.clone(),
rollback_started: rollback_started.clone(),
allow_rollback: allow_rollback.clone(),
launches: launches.clone(),
});
router.set_spawner(spawner.clone()).await;
let first_router = router.clone();
let first =
tokio::spawn(async move { first_router.request_activation("session", 1).await });
first_reserved.notified().await;
let state_guard = router.states.lock().await;
allow_first_return.notify_one();
first_returning.notified().await;
first.abort();
assert!(first.await.unwrap_err().is_cancelled());
drop(state_guard);
rollback_started.notified().await;
let second_router = router.clone();
let second =
tokio::spawn(async move { second_router.request_activation("session", 2).await });
tokio::task::yield_now().await;
assert_eq!(
spawner.reservations.load(Ordering::SeqCst),
1,
"redelivery must remain coalesced until exact external rollback completes"
);
allow_rollback.notify_one();
assert_eq!(
tokio::time::timeout(std::time::Duration::from_secs(1), second)
.await
.expect("redelivery must reserve after rollback")
.unwrap()
.unwrap(),
SessionActivationDisposition::ActivationReserved
);
assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
assert_eq!(launches.load(Ordering::SeqCst), 1);
let current = runners
.read()
.await
.get("session")
.map(|runner| runner.run_id.clone())
.expect("second real registry reservation remains");
assert_ne!(
current, "",
"the redelivery must never adopt an empty/stale run id"
);
}
#[tokio::test]
async fn delivery_during_finalization_launches_one_successor() {
let router = SessionActivationRouter::new();
let spawner = spawner();
router.set_spawner(spawner.clone()).await;
let mut registration = router.register_run("session", "run-old").await.unwrap();
registration.begin_finalization().await;
assert_eq!(
router.request_activation("session", 11).await.unwrap(),
SessionActivationDisposition::ActivationCoalesced
);
let disposition = registration.finish(10).await.unwrap();
assert_eq!(
disposition,
Some(SessionActivationDisposition::ActivationReserved)
);
assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn finalization_does_not_restart_when_cursor_caught_up() {
let router = SessionActivationRouter::new();
let spawner = spawner();
router.set_spawner(spawner.clone()).await;
let mut registration = router.register_run("session", "run-old").await.unwrap();
assert_eq!(
router.request_activation("session", 4).await.unwrap(),
SessionActivationDisposition::ActiveNotified
);
registration.begin_finalization().await;
assert_eq!(registration.finish(4).await.unwrap(), None);
assert_eq!(spawner.reservations.load(Ordering::SeqCst), 0);
assert!(
router.states.lock().await.get("session").is_none(),
"a caught-up completed session must not leak routing state"
);
}
#[tokio::test]
async fn successor_finalization_removes_state_after_raced_generation_catches_up() {
let router = SessionActivationRouter::new();
let spawner = spawner();
router.set_spawner(spawner).await;
let mut registration = router.register_run("session", "run-old").await.unwrap();
registration.begin_finalization().await;
assert_eq!(
router.request_activation("session", 11).await.unwrap(),
SessionActivationDisposition::ActivationCoalesced
);
assert_eq!(
registration.finish(10).await.unwrap(),
Some(SessionActivationDisposition::ActivationReserved)
);
let mut successor_registration = router.register_run("session", "run-1").await.unwrap();
successor_registration.begin_finalization().await;
assert_eq!(successor_registration.finish(11).await.unwrap(), None);
assert!(
router.states.lock().await.get("session").is_none(),
"the successor's caught-up terminal state must be compacted"
);
}
#[tokio::test]
async fn poison_generation_launches_only_one_successor_until_new_work_arrives() {
let router = SessionActivationRouter::new();
let spawner = spawner();
router.set_spawner(spawner.clone()).await;
let mut registration = router
.register_run("session", "run-original")
.await
.unwrap();
assert_eq!(
router.request_activation("session", 7).await.unwrap(),
SessionActivationDisposition::ActiveNotified
);
registration.begin_finalization().await;
assert_eq!(
registration.finish(0).await.unwrap(),
Some(SessionActivationDisposition::ActivationReserved)
);
assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
let mut successor_registration = router.register_run("session", "run-1").await.unwrap();
successor_registration.begin_finalization().await;
assert_eq!(successor_registration.finish(0).await.unwrap(), None);
assert_eq!(
router.request_activation("session", 7).await.unwrap(),
SessionActivationDisposition::ActivationCoalesced
);
assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
assert_eq!(
router.request_activation("session", 8).await.unwrap(),
SessionActivationDisposition::ActivationReserved
);
assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
}
#[tokio::test]
async fn adopted_terminal_run_preserves_generation_for_one_real_successor() {
let router = SessionActivationRouter::new();
let spawner = Arc::new(AlreadyRunningThenReserveSpawner {
reservations: AtomicUsize::new(0),
launches: Arc::new(AtomicUsize::new(0)),
});
router.set_spawner(spawner.clone()).await;
assert_eq!(
router.request_activation("session", 12).await.unwrap(),
SessionActivationDisposition::ActiveNotified
);
router.begin_finalization("session", "adopted-run").await;
assert_eq!(
router
.finish_finalization("session", "adopted-run", 0)
.await
.unwrap(),
Some(SessionActivationDisposition::ActivationReserved)
);
assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
}
}