1use std::collections::HashMap;
8use std::fmt;
9use std::future::Future;
10use std::pin::Pin;
11use std::sync::{Arc, RwLock as StdRwLock};
12
13use async_trait::async_trait;
14use bamboo_domain::{
15 SessionActivationDisposition, SessionActivationError, SessionActivationPort, SessionInboxPort,
16 SessionInboxWakeCandidate,
17};
18use chrono::{DateTime, Duration, Utc};
19use tokio::sync::{mpsc, watch, Mutex, RwLock};
20
21type AsyncRollback =
22 Box<dyn FnOnce() -> Pin<Box<dyn Future<Output = ()> + Send + 'static>> + Send + 'static>;
23
24pub struct SessionActivationLaunch {
29 pub run_id: String,
30 launch: Option<Box<dyn FnOnce() + Send + 'static>>,
31 rollback: Option<AsyncRollback>,
32 rollback_completion: Option<Arc<RollbackCompletion>>,
33}
34
35struct RollbackCompletion {
36 completed: watch::Sender<bool>,
37}
38
39impl RollbackCompletion {
40 fn new() -> Arc<Self> {
41 let (completed, _receiver) = watch::channel(false);
42 Arc::new(Self { completed })
43 }
44
45 fn complete(&self) {
46 self.completed.send_replace(true);
47 }
48
49 async fn wait(&self) {
50 let mut completed = self.completed.subscribe();
51 while !*completed.borrow() {
52 if completed.changed().await.is_err() {
53 break;
54 }
55 }
56 }
57}
58
59impl SessionActivationLaunch {
60 pub fn new(run_id: impl Into<String>, launch: impl FnOnce() + Send + 'static) -> Self {
61 Self {
62 run_id: run_id.into(),
63 launch: Some(Box::new(launch)),
64 rollback: None,
65 rollback_completion: None,
66 }
67 }
68
69 pub fn new_with_rollback(
75 run_id: impl Into<String>,
76 launch: impl FnOnce() + Send + 'static,
77 rollback: impl FnOnce() + Send + 'static,
78 ) -> Self {
79 Self::new_with_async_rollback(run_id, launch, move || async move {
80 rollback();
81 })
82 }
83
84 pub fn new_with_async_rollback<F, Fut>(
89 run_id: impl Into<String>,
90 launch: impl FnOnce() + Send + 'static,
91 rollback: F,
92 ) -> Self
93 where
94 F: FnOnce() -> Fut + Send + 'static,
95 Fut: Future<Output = ()> + Send + 'static,
96 {
97 let rollback_completion = RollbackCompletion::new();
98 Self {
99 run_id: run_id.into(),
100 launch: Some(Box::new(launch)),
101 rollback: Some(Box::new(move || Box::pin(rollback()))),
102 rollback_completion: Some(rollback_completion),
103 }
104 }
105
106 fn rollback_completion(&self) -> Option<Arc<RollbackCompletion>> {
107 self.rollback_completion.clone()
108 }
109
110 fn launch(mut self) {
111 if let Some(launch) = self.launch.take() {
112 self.rollback = None;
113 self.rollback_completion = None;
114 launch();
115 }
116 }
117}
118
119impl Drop for SessionActivationLaunch {
120 fn drop(&mut self) {
121 if self.launch.is_some() {
122 if let Some(rollback) = self.rollback.take() {
123 let completion = self.rollback_completion.take();
124 if let Ok(runtime) = tokio::runtime::Handle::try_current() {
125 runtime.spawn(async move {
126 rollback().await;
127 if let Some(completion) = completion {
128 completion.complete();
129 }
130 });
131 }
132 }
133 }
134 }
135}
136
137impl fmt::Debug for SessionActivationLaunch {
138 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
139 formatter
140 .debug_struct("SessionActivationLaunch")
141 .field("run_id", &self.run_id)
142 .finish_non_exhaustive()
143 }
144}
145
146#[derive(Debug)]
147pub enum SessionActivationReserveOutcome {
148 Reserved(SessionActivationLaunch),
151 AlreadyRunning { run_id: String },
153 NotFound,
155 NoWork,
157}
158
159#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
165pub enum SessionRunRegistrationError {
166 #[error(
167 "session activation owner collision for {target_session_id}: \
168 existing run {existing_run_id}, attempted run {attempted_run_id}"
169 )]
170 OwnerCollision {
171 target_session_id: String,
172 existing_run_id: String,
173 attempted_run_id: String,
174 },
175}
176
177impl SessionRunRegistrationError {
178 pub fn existing_run_id(&self) -> &str {
180 match self {
181 Self::OwnerCollision {
182 existing_run_id, ..
183 } => existing_run_id,
184 }
185 }
186}
187
188#[async_trait]
194pub trait SessionActivationSpawner: Send + Sync {
195 async fn reserve_activation(
196 &self,
197 target_session_id: &str,
198 inbox_generation: u64,
199 ) -> Result<SessionActivationReserveOutcome, SessionActivationError>;
200}
201
202#[derive(Debug, Clone)]
203struct ActiveOwner {
204 run_id: String,
205 finalizing: bool,
206 registrations: usize,
207 delivery_sink: Option<mpsc::UnboundedSender<u64>>,
211}
212
213#[derive(Debug)]
214struct TargetActivationState {
215 latest_generation: u64,
216 last_dispatched_generation: u64,
223 last_dispatched_coordinator_generation: u64,
227 owner_completion_epoch: u64,
230 last_recovery_attempt: Option<RecoveryAttempt>,
231 owner: Option<ActiveOwner>,
232 activation_reserved: bool,
233 activation_token: u64,
236 activation_epoch: watch::Sender<u64>,
241 notify: watch::Sender<u64>,
242}
243
244impl Default for TargetActivationState {
245 fn default() -> Self {
246 let (activation_epoch, _activation_receiver) = watch::channel(0);
247 let (notify, _receiver) = watch::channel(0);
248 Self {
249 latest_generation: 0,
250 last_dispatched_generation: 0,
251 last_dispatched_coordinator_generation: 0,
252 owner_completion_epoch: 0,
253 last_recovery_attempt: None,
254 owner: None,
255 activation_reserved: false,
256 activation_token: 0,
257 activation_epoch,
258 notify,
259 }
260 }
261}
262
263const INITIAL_RECOVERY_BACKOFF: Duration = Duration::seconds(30);
264const MAX_RECOVERY_BACKOFF: Duration = Duration::minutes(30);
265
266#[derive(Debug, Clone)]
267struct RecoveryAttempt {
268 candidate: SessionInboxWakeCandidate,
269 completion_epoch: u64,
270 backoff: Duration,
271 next_retry_at: DateTime<Utc>,
272 requires_completion: bool,
275}
276
277fn recovery_retry_at(
278 state: &TargetActivationState,
279 candidate: &SessionInboxWakeCandidate,
280 now: DateTime<Utc>,
281) -> Option<Option<DateTime<Utc>>> {
282 let previous = state
283 .last_recovery_attempt
284 .as_ref()
285 .filter(|attempt| attempt.candidate == *candidate)?;
286 if previous.requires_completion && state.owner_completion_epoch <= previous.completion_epoch {
287 return Some(None);
288 }
289 (now < previous.next_retry_at).then_some(Some(previous.next_retry_at))
290}
291
292fn record_recovery_attempt(
293 state: &mut TargetActivationState,
294 candidate: SessionInboxWakeCandidate,
295 now: DateTime<Utc>,
296 requires_completion: bool,
297) {
298 let backoff = state
299 .last_recovery_attempt
300 .as_ref()
301 .filter(|attempt| attempt.candidate == candidate)
302 .map_or(INITIAL_RECOVERY_BACKOFF, |attempt| {
303 (attempt.backoff * 2).min(MAX_RECOVERY_BACKOFF)
304 });
305 state.last_recovery_attempt = Some(RecoveryAttempt {
306 candidate,
307 completion_epoch: state.owner_completion_epoch,
308 backoff,
309 next_retry_at: now + backoff,
310 requires_completion,
311 });
312}
313
314#[derive(Debug, Clone, PartialEq, Eq)]
317pub struct SessionWakeReconcileResult {
318 pub disposition: Option<SessionActivationDisposition>,
319 pub next_due_at: Option<DateTime<Utc>>,
320}
321
322fn reserve_activation_token(state: &mut TargetActivationState) -> u64 {
323 state.activation_reserved = true;
324 state.activation_token = state.activation_token.wrapping_add(1).max(1);
325 state.activation_token
326}
327
328fn release_activation_token(state: &mut TargetActivationState, token: u64) -> bool {
329 if !state.activation_reserved || state.activation_token != token {
330 return false;
331 }
332 state.activation_reserved = false;
333 let next_epoch = (*state.activation_epoch.borrow()).wrapping_add(1);
334 state.activation_epoch.send_replace(next_epoch);
335 true
336}
337
338fn inbox_activation_error(error: bamboo_domain::SessionInboxError) -> SessionActivationError {
339 match error {
340 bamboo_domain::SessionInboxError::TargetNotFound(target) => {
341 SessionActivationError::TargetNotFound(target)
342 }
343 error => SessionActivationError::Internal(error.to_string()),
344 }
345}
346
347struct ActivationReservationLease {
353 router: SessionActivationRouter,
354 target_session_id: String,
355 token: u64,
356 recover_on_drop: bool,
357 armed: bool,
358 rollback_completion: Option<Arc<RollbackCompletion>>,
359}
360
361impl ActivationReservationLease {
362 fn new(
363 router: SessionActivationRouter,
364 target_session_id: &str,
365 token: u64,
366 recover_on_drop: bool,
367 ) -> Self {
368 Self {
369 router,
370 target_session_id: target_session_id.to_string(),
371 token,
372 recover_on_drop,
373 armed: true,
374 rollback_completion: None,
375 }
376 }
377
378 fn wait_for_rollback(&mut self, completion: Option<Arc<RollbackCompletion>>) {
379 self.rollback_completion = completion;
380 }
381
382 fn disarm(&mut self) {
383 self.armed = false;
384 }
385}
386
387impl Drop for ActivationReservationLease {
388 fn drop(&mut self) {
389 if !self.armed {
390 return;
391 }
392 let router = self.router.clone();
393 let target_session_id = self.target_session_id.clone();
394 let token = self.token;
395 let recover_on_drop = self.recover_on_drop;
396 let rollback_completion = self.rollback_completion.take();
397 if let Ok(runtime) = tokio::runtime::Handle::try_current() {
398 runtime.spawn(async move {
399 if let Some(completion) = rollback_completion {
400 completion.wait().await;
401 }
402 router
403 .recover_cancelled_reservation(&target_session_id, token, recover_on_drop)
404 .await;
405 });
406 }
407 }
408}
409
410#[derive(Clone, Default)]
412pub struct SessionActivationRouter {
413 states: Arc<Mutex<HashMap<String, TargetActivationState>>>,
414 spawner: Arc<RwLock<Option<Arc<dyn SessionActivationSpawner>>>>,
415 inbox: Arc<StdRwLock<Option<Arc<dyn SessionInboxPort>>>>,
416}
417
418type AbortCleanup =
419 Box<dyn FnOnce() -> Pin<Box<dyn Future<Output = ()> + Send + 'static>> + Send + 'static>;
420
421pub struct SessionRunRegistration {
429 router: Arc<SessionActivationRouter>,
430 target_session_id: String,
431 run_id: String,
432 notifications: Option<watch::Receiver<u64>>,
433 abort_cleanup: Option<AbortCleanup>,
434 armed: bool,
435}
436
437impl SessionRunRegistration {
438 pub fn notifications_mut(&mut self) -> &mut watch::Receiver<u64> {
439 self.notifications
440 .as_mut()
441 .expect("run registration notifications are live before finalization")
442 }
443
444 pub fn set_abort_cleanup<F, Fut>(&mut self, cleanup: F)
447 where
448 F: FnOnce() -> Fut + Send + 'static,
449 Fut: Future<Output = ()> + Send + 'static,
450 {
451 self.abort_cleanup = Some(Box::new(move || Box::pin(cleanup())));
452 }
453
454 pub async fn begin_finalization(&mut self) {
455 self.router
456 .begin_finalization(&self.target_session_id, &self.run_id)
457 .await;
458 }
459
460 pub async fn abandon(mut self) {
469 self.armed = false;
470 let abort_cleanup = self.abort_cleanup.take();
471 let router = self.router.clone();
472 let target_session_id = self.target_session_id.clone();
473 let run_id = self.run_id.clone();
474 let cleanup_target_session_id = target_session_id.clone();
475 let cleanup_run_id = run_id.clone();
476 let cleanup = tokio::spawn(async move {
477 if let Some(cleanup) = abort_cleanup {
478 cleanup().await;
479 }
480 router
481 .cleanup_abandoned_registration(&cleanup_target_session_id, &cleanup_run_id)
482 .await;
483 });
484 if let Err(error) = cleanup.await {
485 tracing::error!(
486 %target_session_id,
487 %run_id,
488 %error,
489 "detached SessionInbox registration cleanup failed"
490 );
491 }
492 }
493
494 pub async fn finish(
497 mut self,
498 admitted_generation: u64,
499 ) -> Result<Option<SessionActivationDisposition>, SessionActivationError> {
500 drop(self.notifications.take());
503 let result = self
504 .router
505 .finish_finalization(&self.target_session_id, &self.run_id, admitted_generation)
506 .await;
507 if result.is_ok() {
508 self.armed = false;
509 self.abort_cleanup = None;
510 }
511 result
512 }
513}
514
515impl Drop for SessionRunRegistration {
516 fn drop(&mut self) {
517 if !self.armed {
518 return;
519 }
520 let router = self.router.clone();
521 let target_session_id = self.target_session_id.clone();
522 let run_id = self.run_id.clone();
523 let abort_cleanup = self.abort_cleanup.take();
524 if let Ok(runtime) = tokio::runtime::Handle::try_current() {
525 runtime.spawn(async move {
526 if let Some(cleanup) = abort_cleanup {
527 cleanup().await;
528 }
529 router
530 .cleanup_abandoned_registration(&target_session_id, &run_id)
531 .await;
532 });
533 }
534 }
535}
536
537impl SessionActivationRouter {
538 async fn coordinator_generation(&self, target: &str) -> Result<u64, SessionActivationError> {
539 let inbox = self
540 .inbox
541 .read()
542 .unwrap_or_else(std::sync::PoisonError::into_inner)
543 .clone();
544 match inbox {
545 Some(inbox) => inbox
546 .coordinator_activation_generation(target)
547 .await
548 .map_err(inbox_activation_error),
549 None => Ok(0),
550 }
551 }
552
553 async fn released_pending_generation(
554 &self,
555 target: &str,
556 after: u64,
557 ) -> Result<u64, SessionActivationError> {
558 if self.coordinator_generation(target).await? <= after {
559 return Ok(0);
560 }
561 let inbox = self
562 .inbox
563 .read()
564 .unwrap_or_else(std::sync::PoisonError::into_inner)
565 .clone();
566 match inbox {
567 Some(inbox) => {
568 let backlog = inbox
569 .inspect(target)
570 .await
571 .map_err(inbox_activation_error)?;
572 Ok(
573 if backlog.activation_pending() && backlog.coordinator_generation > after {
574 backlog.coordinator_generation
575 } else {
576 0
577 },
578 )
579 }
580 None => Ok(0),
581 }
582 }
583
584 pub fn new() -> Arc<Self> {
585 Arc::new(Self::default())
586 }
587
588 pub async fn set_spawner(&self, spawner: Arc<dyn SessionActivationSpawner>) {
591 *self.spawner.write().await = Some(spawner);
592 }
593
594 pub async fn current_run_id(&self, target_session_id: &str) -> Option<String> {
596 self.states
597 .lock()
598 .await
599 .get(target_session_id)
600 .and_then(|state| state.owner.as_ref().map(|owner| owner.run_id.clone()))
601 }
602
603 pub fn set_inbox(&self, inbox: Arc<dyn SessionInboxPort>) {
605 *self
606 .inbox
607 .write()
608 .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(inbox);
609 }
610
611 pub async fn reconcile_wake(
615 &self,
616 target_session_id: &str,
617 now: DateTime<Utc>,
618 ) -> Result<SessionWakeReconcileResult, SessionActivationError> {
619 let inbox = self
620 .inbox
621 .read()
622 .unwrap_or_else(std::sync::PoisonError::into_inner)
623 .clone()
624 .ok_or_else(|| {
625 SessionActivationError::Internal("session inbox is not configured".into())
626 })?;
627 let readiness = inbox
628 .inspect_wake_readiness(target_session_id, now)
629 .await
630 .map_err(inbox_activation_error)?;
631 let Some(candidate) = readiness.ready else {
632 return Ok(SessionWakeReconcileResult {
633 disposition: None,
634 next_due_at: readiness.next_due_at,
635 });
636 };
637 let mut result = self
638 .request_recovery_activation(target_session_id, candidate, now)
639 .await?;
640 result.next_due_at = match (result.next_due_at, readiness.next_due_at) {
641 (Some(left), Some(right)) => Some(left.min(right)),
642 (left, right) => left.or(right),
643 };
644 Ok(result)
645 }
646
647 async fn request_recovery_activation(
648 &self,
649 target_session_id: &str,
650 candidate: SessionInboxWakeCandidate,
651 now: DateTime<Utc>,
652 ) -> Result<SessionWakeReconcileResult, SessionActivationError> {
653 loop {
654 let (reservation_wait, reservation_token) = {
655 let mut states = self.states.lock().await;
656 let state = states.entry(target_session_id.to_string()).or_default();
657 state.latest_generation = state.latest_generation.max(candidate.generation);
658 if let Some(owner) = state.owner.as_mut() {
659 if !owner.finalizing {
660 state.notify.send_replace(candidate.generation);
661 if owner
662 .delivery_sink
663 .as_ref()
664 .is_some_and(|sink| sink.send(candidate.generation).is_err())
665 {
666 owner.delivery_sink = None;
667 }
668 return Ok(SessionWakeReconcileResult {
669 disposition: Some(SessionActivationDisposition::ActiveNotified),
670 next_due_at: None,
671 });
672 }
673 return Ok(SessionWakeReconcileResult {
674 disposition: Some(SessionActivationDisposition::ActivationCoalesced),
675 next_due_at: None,
676 });
677 }
678 if let Some(next_due_at) = recovery_retry_at(state, &candidate, now) {
679 return Ok(SessionWakeReconcileResult {
680 disposition: None,
681 next_due_at,
682 });
683 } else if state.activation_reserved {
684 (Some(state.activation_epoch.subscribe()), None)
685 } else {
686 let token = reserve_activation_token(state);
687 (None, Some(token))
688 }
689 };
690
691 if let Some(mut reservation_wait) = reservation_wait {
692 if reservation_wait.changed().await.is_err() {
693 return Err(SessionActivationError::TargetNotFound(
694 target_session_id.to_string(),
695 ));
696 }
697 continue;
698 }
699 let token = reservation_token.expect("recovery activation owns reservation token");
700 let disposition = self
701 .dispatch_reserved_with_recovery(
702 target_session_id,
703 candidate.generation,
704 token,
705 false,
706 Some((candidate, now)),
707 )
708 .await?;
709 return Ok(SessionWakeReconcileResult {
710 disposition: Some(disposition),
711 next_due_at: None,
712 });
713 }
714 }
715
716 #[cfg(test)]
717 pub(crate) async fn hold_state_lock_for_test(
718 &self,
719 entered: Arc<tokio::sync::Notify>,
720 release: Arc<tokio::sync::Notify>,
721 ) {
722 let _states = self.states.lock().await;
723 entered.notify_one();
724 release.notified().await;
725 }
726
727 #[cfg(test)]
728 pub(crate) async fn install_owner_placeholder_for_test(
729 &self,
730 target_session_id: &str,
731 run_id: &str,
732 ) {
733 let mut states = self.states.lock().await;
734 states
735 .entry(target_session_id.to_string())
736 .or_default()
737 .owner = Some(ActiveOwner {
738 run_id: run_id.to_string(),
739 finalizing: false,
740 registrations: 0,
741 delivery_sink: None,
742 });
743 }
744
745 pub async fn register_run(
749 self: &Arc<Self>,
750 target_session_id: &str,
751 run_id: &str,
752 ) -> Result<SessionRunRegistration, SessionRunRegistrationError> {
753 loop {
754 let (notifications, reservation_wait) = {
755 let mut states = self.states.lock().await;
756 let state = states.entry(target_session_id.to_string()).or_default();
757 match state.owner.as_mut() {
758 Some(owner)
759 if owner.run_id == run_id
760 && owner.registrations == 0
761 && !owner.finalizing =>
762 {
763 owner.registrations = 1;
767 (Some(state.notify.subscribe()), None)
768 }
769 Some(owner) => {
770 return Err(SessionRunRegistrationError::OwnerCollision {
771 target_session_id: target_session_id.to_string(),
772 existing_run_id: owner.run_id.clone(),
773 attempted_run_id: run_id.to_string(),
774 });
775 }
776 None if state.activation_reserved => {
777 (None, Some(state.activation_epoch.subscribe()))
783 }
784 None => {
785 state.owner = Some(ActiveOwner {
786 run_id: run_id.to_string(),
787 finalizing: false,
788 registrations: 1,
789 delivery_sink: None,
790 });
791 (Some(state.notify.subscribe()), None)
792 }
793 }
794 };
795
796 if let Some(mut reservation_wait) = reservation_wait {
797 let _ = reservation_wait.changed().await;
798 continue;
799 }
800 return Ok(SessionRunRegistration {
801 router: self.clone(),
802 target_session_id: target_session_id.to_string(),
803 run_id: run_id.to_string(),
804 notifications,
805 abort_cleanup: None,
806 armed: true,
807 });
808 }
809 }
810
811 pub async fn attach_delivery_sink(
817 &self,
818 target_session_id: &str,
819 sink: mpsc::UnboundedSender<u64>,
820 ) -> Option<String> {
821 let mut states = self.states.lock().await;
822 let state = states.get_mut(target_session_id)?;
823 let owner = state.owner.as_mut()?;
824 if owner.finalizing {
825 return None;
826 }
827 owner.delivery_sink = Some(sink.clone());
828 let run_id = owner.run_id.clone();
829 if state.latest_generation > 0 && sink.send(state.latest_generation).is_err() {
830 owner.delivery_sink = None;
831 return None;
832 }
833 Some(run_id)
834 }
835
836 pub async fn detach_delivery_sink(&self, target_session_id: &str, run_id: &str) {
839 let mut states = self.states.lock().await;
840 let Some(state) = states.get_mut(target_session_id) else {
841 return;
842 };
843 let Some(owner) = state.owner.as_mut() else {
844 return;
845 };
846 if owner.run_id == run_id {
847 owner.delivery_sink = None;
848 }
849 }
850
851 pub async fn owns_run(&self, target_session_id: &str, run_id: &str) -> bool {
853 let states = self.states.lock().await;
854 states
855 .get(target_session_id)
856 .and_then(|state| state.owner.as_ref())
857 .is_some_and(|owner| owner.run_id == run_id && !owner.finalizing)
858 }
859
860 pub async fn begin_finalization(&self, target_session_id: &str, run_id: &str) {
865 let mut states = self.states.lock().await;
866 let state = states.entry(target_session_id.to_string()).or_default();
867 match state.owner.as_mut() {
868 Some(owner) if owner.run_id == run_id => owner.finalizing = true,
869 Some(_) => {}
870 None => {
871 state.owner = Some(ActiveOwner {
872 run_id: run_id.to_string(),
873 finalizing: true,
874 registrations: 0,
875 delivery_sink: None,
876 });
877 }
878 }
879 }
880
881 async fn cleanup_abandoned_registration(
882 self: &Arc<Self>,
883 target_session_id: &str,
884 run_id: &str,
885 ) {
886 {
887 let mut states = self.states.lock().await;
888 let Some(state) = states.get_mut(target_session_id) else {
889 return;
890 };
891 let Some(owner) = state.owner.as_mut() else {
892 return;
893 };
894 if owner.run_id != run_id {
895 return;
896 }
897 owner.registrations = 0;
898 owner.finalizing = true;
899 }
900
901 let inbox = self
906 .inbox
907 .read()
908 .unwrap_or_else(std::sync::PoisonError::into_inner)
909 .clone();
910 let (durable_generation, released_prefix) = match inbox {
911 Some(inbox) => match inbox.inspect(target_session_id).await {
912 Ok(backlog) if backlog.activation_pending() => (
913 backlog.activation_generation,
914 backlog.coordinator_generation,
915 ),
916 Ok(_) => (0, 0),
917 Err(error) => {
918 tracing::warn!(
919 %target_session_id,
920 %run_id,
921 %error,
922 "failed to reconcile durable SessionInbox generation for abandoned run"
923 );
924 (0, 0)
925 }
926 },
927 None => (0, 0),
928 };
929
930 let reservation_to_dispatch = {
931 let mut states = self.states.lock().await;
932 let Some(state) = states.get_mut(target_session_id) else {
933 return;
934 };
935 if !state.owner.as_ref().is_some_and(|owner| {
936 owner.run_id == run_id && owner.finalizing && owner.registrations == 0
937 }) {
938 return;
939 }
940 state.latest_generation = state.latest_generation.max(durable_generation);
941 state.owner = None;
942 state.owner_completion_epoch = state.owner_completion_epoch.wrapping_add(1);
943 if (state.latest_generation > state.last_dispatched_generation
944 || released_prefix > state.last_dispatched_coordinator_generation)
945 && !state.activation_reserved
946 {
947 let generation = state.latest_generation;
948 let token = reserve_activation_token(state);
949 Some((generation, token))
950 } else {
951 if state.latest_generation == 0
952 && !state.activation_reserved
953 && state.notify.receiver_count() == 0
954 && state.activation_epoch.receiver_count() == 0
955 {
956 states.remove(target_session_id);
957 }
958 None
959 }
960 };
961
962 if let Some((generation, token)) = reservation_to_dispatch {
963 if let Err(error) = self
964 .dispatch_reserved(target_session_id, generation, token)
965 .await
966 {
967 tracing::error!(
968 %target_session_id,
969 %run_id,
970 %error,
971 "failed to reserve successor for abandoned SessionInbox owner"
972 );
973 }
974 }
975 }
976
977 pub async fn finish_finalization(
981 &self,
982 target_session_id: &str,
983 run_id: &str,
984 admitted_generation: u64,
985 ) -> Result<Option<SessionActivationDisposition>, SessionActivationError> {
986 let inbox = self
989 .inbox
990 .read()
991 .unwrap_or_else(std::sync::PoisonError::into_inner)
992 .clone();
993 let durable_pending = match inbox {
994 Some(inbox) => Some(
995 inbox
996 .inspect(target_session_id)
997 .await
998 .map_err(|error| SessionActivationError::Internal(error.to_string()))?,
999 ),
1000 None => None,
1001 };
1002 let reservation_to_dispatch = {
1003 let mut states = self.states.lock().await;
1004 let state = states.entry(target_session_id.to_string()).or_default();
1005 if state
1006 .owner
1007 .as_ref()
1008 .is_some_and(|owner| owner.run_id != run_id)
1009 {
1010 return Ok(None);
1011 }
1012 if state.owner.take().is_some() {
1013 state.owner_completion_epoch = state.owner_completion_epoch.wrapping_add(1);
1014 }
1015 let pending = durable_pending
1016 .as_ref()
1017 .is_some_and(|backlog| backlog.activation_pending());
1018 let released_prefix = durable_pending
1019 .as_ref()
1020 .filter(|_| pending)
1021 .map_or(0, |backlog| backlog.coordinator_generation);
1022 if let Some(backlog) = durable_pending.as_ref() {
1023 state.latest_generation =
1024 state.latest_generation.max(backlog.activation_generation);
1025 }
1026 if (pending || state.latest_generation > admitted_generation)
1027 && (state.latest_generation > state.last_dispatched_generation
1028 || released_prefix > state.last_dispatched_coordinator_generation)
1029 && !state.activation_reserved
1030 {
1031 let generation = state.latest_generation;
1032 let token = reserve_activation_token(state);
1033 Some((generation, token))
1034 } else {
1035 if state.latest_generation <= admitted_generation
1040 && !pending
1041 && state.owner.is_none()
1042 && !state.activation_reserved
1043 && state.notify.receiver_count() == 0
1044 && state.activation_epoch.receiver_count() == 0
1045 {
1046 states.remove(target_session_id);
1047 }
1048 None
1049 }
1050 };
1051
1052 match reservation_to_dispatch {
1053 Some((generation, token)) => self
1054 .dispatch_reserved(target_session_id, generation, token)
1055 .await
1056 .map(Some),
1057 None => Ok(None),
1058 }
1059 }
1060
1061 pub async fn subscribe(&self, target_session_id: &str) -> watch::Receiver<u64> {
1062 let mut states = self.states.lock().await;
1063 states
1064 .entry(target_session_id.to_string())
1065 .or_default()
1066 .notify
1067 .subscribe()
1068 }
1069
1070 async fn recover_cancelled_reservation(
1079 &self,
1080 target_session_id: &str,
1081 reservation_token: u64,
1082 recover_on_drop: bool,
1083 ) {
1084 let previous_prefix = {
1085 let mut states = self.states.lock().await;
1086 let Some(state) = states.get_mut(target_session_id) else {
1087 return;
1088 };
1089 if !release_activation_token(state, reservation_token) || !recover_on_drop {
1090 return;
1091 }
1092 state.last_dispatched_coordinator_generation
1093 };
1094 let released_prefix = self
1095 .released_pending_generation(target_session_id, previous_prefix)
1096 .await
1097 .unwrap_or(0);
1098 let reservation_to_dispatch = {
1099 let mut states = self.states.lock().await;
1100 let Some(state) = states.get_mut(target_session_id) else {
1101 return;
1102 };
1103 if state.owner.is_none()
1104 && (state.latest_generation > state.last_dispatched_generation
1105 || released_prefix > state.last_dispatched_coordinator_generation)
1106 && !state.activation_reserved
1107 {
1108 let generation = state.latest_generation;
1109 let token = reserve_activation_token(state);
1110 Some((generation, token))
1111 } else {
1112 None
1113 }
1114 };
1115
1116 if let Some((generation, token)) = reservation_to_dispatch {
1117 if let Err(error) = self
1118 .dispatch_reserved_with_recovery(target_session_id, generation, token, false, None)
1119 .await
1120 {
1121 tracing::error!(
1122 %target_session_id,
1123 %error,
1124 "failed bounded retry after cancelled SessionInbox activation reservation"
1125 );
1126 }
1127 }
1128 }
1129
1130 async fn dispatch_reserved(
1131 &self,
1132 target_session_id: &str,
1133 generation: u64,
1134 reservation_token: u64,
1135 ) -> Result<SessionActivationDisposition, SessionActivationError> {
1136 self.dispatch_reserved_with_recovery(
1137 target_session_id,
1138 generation,
1139 reservation_token,
1140 true,
1141 None,
1142 )
1143 .await
1144 }
1145
1146 async fn dispatch_reserved_with_recovery(
1147 &self,
1148 target_session_id: &str,
1149 generation: u64,
1150 reservation_token: u64,
1151 recover_on_drop: bool,
1152 recovery: Option<(SessionInboxWakeCandidate, DateTime<Utc>)>,
1153 ) -> Result<SessionActivationDisposition, SessionActivationError> {
1154 let mut lease = ActivationReservationLease::new(
1155 self.clone(),
1156 target_session_id,
1157 reservation_token,
1158 recover_on_drop,
1159 );
1160 let Some(spawner) = self.spawner.read().await.clone() else {
1161 let mut states = self.states.lock().await;
1162 if let Some(state) = states.get_mut(target_session_id) {
1163 release_activation_token(state, reservation_token);
1164 }
1165 lease.disarm();
1166 return Err(SessionActivationError::Internal(
1167 "session activation spawner is not configured".to_string(),
1168 ));
1169 };
1170
1171 let coordinator_generation = self.coordinator_generation(target_session_id).await?;
1174
1175 let outcome = spawner
1176 .reserve_activation(target_session_id, generation)
1177 .await;
1178 if let Ok(SessionActivationReserveOutcome::Reserved(launch)) = &outcome {
1179 lease.wait_for_rollback(launch.rollback_completion());
1180 }
1181 match outcome {
1182 Ok(SessionActivationReserveOutcome::Reserved(launch)) => {
1183 let run_id = launch.run_id.clone();
1184 {
1185 let mut states = self.states.lock().await;
1186 let state = states.entry(target_session_id.to_string()).or_default();
1187 if !release_activation_token(state, reservation_token) {
1188 lease.disarm();
1196 return if state.owner.is_some() {
1197 Ok(SessionActivationDisposition::ActiveNotified)
1198 } else {
1199 Err(SessionActivationError::Internal(
1200 "session activation reservation was superseded".to_string(),
1201 ))
1202 };
1203 }
1204 state.owner = Some(ActiveOwner {
1205 run_id,
1206 finalizing: false,
1207 registrations: 0,
1208 delivery_sink: None,
1209 });
1210 state.last_dispatched_generation =
1211 state.last_dispatched_generation.max(generation);
1212 state.last_dispatched_coordinator_generation = state
1213 .last_dispatched_coordinator_generation
1214 .max(coordinator_generation);
1215 if let Some((candidate, now)) = recovery {
1216 record_recovery_attempt(state, candidate, now, true);
1217 }
1218 }
1219 lease.disarm();
1220 launch.launch();
1223 Ok(SessionActivationDisposition::ActivationReserved)
1224 }
1225 Ok(SessionActivationReserveOutcome::AlreadyRunning { run_id }) => {
1226 let mut states = self.states.lock().await;
1227 let state = states.entry(target_session_id.to_string()).or_default();
1228 if !release_activation_token(state, reservation_token) {
1229 lease.disarm();
1230 return Ok(SessionActivationDisposition::ActiveNotified);
1231 }
1232 state.owner = Some(ActiveOwner {
1233 run_id,
1234 finalizing: false,
1235 registrations: 0,
1236 delivery_sink: None,
1237 });
1238 state.notify.send_replace(generation);
1239 if let Some((candidate, now)) = recovery {
1240 record_recovery_attempt(state, candidate, now, true);
1241 }
1242 lease.disarm();
1243 Ok(SessionActivationDisposition::ActiveNotified)
1244 }
1245 Ok(SessionActivationReserveOutcome::NoWork) => {
1246 let mut states = self.states.lock().await;
1247 if let Some(state) = states.get_mut(target_session_id) {
1248 if release_activation_token(state, reservation_token) {
1249 if let Some((candidate, now)) = recovery {
1250 record_recovery_attempt(state, candidate, now, false);
1251 }
1252 }
1253 }
1254 lease.disarm();
1255 Ok(SessionActivationDisposition::ActivationCoalesced)
1256 }
1257 Ok(SessionActivationReserveOutcome::NotFound) => {
1258 let mut states = self.states.lock().await;
1259 if states
1260 .get(target_session_id)
1261 .is_some_and(|state| state.activation_token == reservation_token)
1262 {
1263 states.remove(target_session_id);
1264 }
1265 lease.disarm();
1266 Err(SessionActivationError::TargetNotFound(
1267 target_session_id.to_string(),
1268 ))
1269 }
1270 Err(error) => {
1271 let mut states = self.states.lock().await;
1272 if let Some(state) = states.get_mut(target_session_id) {
1273 release_activation_token(state, reservation_token);
1274 }
1275 lease.disarm();
1276 Err(error)
1277 }
1278 }
1279 }
1280}
1281
1282#[async_trait]
1283impl SessionActivationPort for SessionActivationRouter {
1284 async fn request_activation(
1285 &self,
1286 target_session_id: &str,
1287 inbox_generation: u64,
1288 ) -> Result<SessionActivationDisposition, SessionActivationError> {
1289 let previous_prefix = {
1290 let states = self.states.lock().await;
1291 states
1292 .get(target_session_id)
1293 .filter(|state| inbox_generation <= state.last_dispatched_generation)
1294 .map(|state| state.last_dispatched_coordinator_generation)
1295 };
1296 let released_prefix = match previous_prefix {
1297 Some(prefix) => {
1298 self.released_pending_generation(target_session_id, prefix)
1299 .await?
1300 }
1301 None => 0,
1302 };
1303 loop {
1304 let (reservation_wait, reservation_token) = {
1305 let mut states = self.states.lock().await;
1306 let state = states.entry(target_session_id.to_string()).or_default();
1307 state.latest_generation = state.latest_generation.max(inbox_generation);
1308 if let Some(owner) = state.owner.as_mut() {
1309 if !owner.finalizing {
1310 state.notify.send_replace(inbox_generation);
1311 if owner
1312 .delivery_sink
1313 .as_ref()
1314 .is_some_and(|sink| sink.send(inbox_generation).is_err())
1315 {
1316 owner.delivery_sink = None;
1317 }
1318 return Ok(SessionActivationDisposition::ActiveNotified);
1319 }
1320 return Ok(SessionActivationDisposition::ActivationCoalesced);
1321 }
1322 if inbox_generation <= state.last_dispatched_generation
1323 && released_prefix <= state.last_dispatched_coordinator_generation
1324 {
1325 return Ok(SessionActivationDisposition::ActivationCoalesced);
1326 }
1327 if state.activation_reserved {
1328 (Some(state.activation_epoch.subscribe()), None)
1329 } else {
1330 let token = reserve_activation_token(state);
1331 (None, Some(token))
1332 }
1333 };
1334
1335 if let Some(mut reservation_wait) = reservation_wait {
1336 if reservation_wait.changed().await.is_err() {
1342 return Err(SessionActivationError::TargetNotFound(
1343 target_session_id.to_string(),
1344 ));
1345 }
1346 continue;
1347 }
1348
1349 let reservation_token =
1350 reservation_token.expect("non-waiting activation owns a reservation token");
1351 return self
1352 .dispatch_reserved_with_recovery(
1353 target_session_id,
1354 inbox_generation,
1355 reservation_token,
1356 false,
1357 None,
1358 )
1359 .await;
1360 }
1361 }
1362}
1363
1364#[cfg(test)]
1365mod tests {
1366 use super::*;
1367 use std::sync::atomic::{AtomicUsize, Ordering};
1368 use tokio::sync::{Barrier, Notify};
1369
1370 #[tokio::test]
1371 async fn wake_recovery_retries_same_candidate_only_after_owner_completion_and_backoff() {
1372 use bamboo_domain::{
1373 Session, SessionActivationPolicy, SessionInboxLimits, SessionMessageEnvelope, Storage,
1374 };
1375 use bamboo_storage::{FileSessionInbox, SessionStoreV2};
1376
1377 let temp = tempfile::TempDir::new().unwrap();
1378 let store = Arc::new(SessionStoreV2::new(temp.path().into()).await.unwrap());
1379 store
1380 .save_session(&Session::new("target", "model"))
1381 .await
1382 .unwrap();
1383 let inbox = Arc::new(FileSessionInbox::new(store, SessionInboxLimits::default()));
1384 inbox
1385 .deliver_with_activation_intent(
1386 &SessionMessageEnvelope::user_input("target", "still pending"),
1387 SessionActivationPolicy::RespectSpecificWait,
1388 None,
1389 )
1390 .await
1391 .unwrap();
1392 let router = SessionActivationRouter::new();
1393 router.set_inbox(inbox);
1394 let spawner = spawner();
1395 router.set_spawner(spawner.clone()).await;
1396 let now = Utc::now();
1397 assert_eq!(
1398 router
1399 .reconcile_wake("target", now)
1400 .await
1401 .unwrap()
1402 .disposition,
1403 Some(SessionActivationDisposition::ActivationReserved)
1404 );
1405 assert_eq!(
1406 router
1407 .reconcile_wake("target", now)
1408 .await
1409 .unwrap()
1410 .disposition,
1411 Some(SessionActivationDisposition::ActiveNotified)
1412 );
1413 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
1414 let mut registration = router.register_run("target", "run-1").await.unwrap();
1415 registration.begin_finalization().await;
1416 assert_eq!(registration.finish(0).await.unwrap(), None);
1417
1418 let deferred = router
1419 .reconcile_wake("target", now + Duration::seconds(29))
1420 .await
1421 .unwrap();
1422 assert_eq!(deferred.disposition, None);
1423 assert_eq!(deferred.next_due_at, Some(now + INITIAL_RECOVERY_BACKOFF));
1424 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
1425 assert_eq!(
1426 router
1427 .reconcile_wake("target", now + INITIAL_RECOVERY_BACKOFF)
1428 .await
1429 .unwrap()
1430 .disposition,
1431 Some(SessionActivationDisposition::ActivationReserved)
1432 );
1433 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
1434 }
1435
1436 #[tokio::test]
1437 async fn wake_recovery_waits_for_exact_owned_lease_expiry() {
1438 use bamboo_domain::{
1439 Session, SessionActivationPolicy, SessionInboxLeaseRequest, SessionInboxLimits,
1440 SessionMessageEnvelope, Storage,
1441 };
1442 use bamboo_storage::{FileSessionInbox, SessionStoreV2};
1443
1444 let temp = tempfile::TempDir::new().unwrap();
1445 let store = Arc::new(SessionStoreV2::new(temp.path().into()).await.unwrap());
1446 store
1447 .save_session(&Session::new("target", "model"))
1448 .await
1449 .unwrap();
1450 let inbox = Arc::new(FileSessionInbox::new(store, SessionInboxLimits::default()));
1451 inbox
1452 .deliver_with_activation_intent(
1453 &SessionMessageEnvelope::user_input("target", "unacknowledged"),
1454 SessionActivationPolicy::RespectSpecificWait,
1455 None,
1456 )
1457 .await
1458 .unwrap();
1459 let router = SessionActivationRouter::new();
1460 router.set_inbox(inbox.clone());
1461 let spawner = spawner();
1462 router.set_spawner(spawner.clone()).await;
1463 let now = Utc::now();
1464 assert_eq!(
1465 router
1466 .reconcile_wake("target", now)
1467 .await
1468 .unwrap()
1469 .disposition,
1470 Some(SessionActivationDisposition::ActivationReserved)
1471 );
1472 let claim = inbox
1473 .claim_owned(
1474 "target",
1475 1,
1476 None,
1477 &SessionInboxLeaseRequest {
1478 consumer: bamboo_domain::SessionInboxConsumerId::new(),
1479 now,
1480 duration: Duration::seconds(10),
1481 },
1482 )
1483 .await
1484 .unwrap()
1485 .remove(0);
1486 let mut registration = router.register_run("target", "run-1").await.unwrap();
1487 registration.begin_finalization().await;
1488 assert_eq!(registration.finish(0).await.unwrap(), None);
1489 let deferred = router
1490 .reconcile_wake("target", claim.lease.expires_at - Duration::seconds(1))
1491 .await
1492 .unwrap();
1493 assert_eq!(deferred.disposition, None);
1494 assert_eq!(deferred.next_due_at, Some(claim.lease.expires_at));
1495 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
1496 assert_eq!(
1497 router
1498 .reconcile_wake("target", claim.lease.expires_at)
1499 .await
1500 .unwrap()
1501 .disposition,
1502 Some(SessionActivationDisposition::ActivationReserved)
1503 );
1504 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
1505 }
1506
1507 #[tokio::test]
1508 async fn older_staged_release_after_immediate_ack_dispatches_once_without_poison_hot_loop() {
1509 use bamboo_domain::{
1510 Session, SessionActivationPolicy, SessionInboxLimits, SessionMessageEnvelope, Storage,
1511 };
1512 use bamboo_storage::{FileSessionInbox, SessionStoreV2};
1513 for release_during_finalization in [false, true] {
1514 let temp = tempfile::TempDir::new().unwrap();
1515 let store = Arc::new(SessionStoreV2::new(temp.path().into()).await.unwrap());
1516 store
1517 .save_session(&Session::new("parent", "model"))
1518 .await
1519 .unwrap();
1520 let inbox = Arc::new(FileSessionInbox::new(
1521 store.clone(),
1522 SessionInboxLimits::default(),
1523 ));
1524 let staged = SessionMessageEnvelope::user_input("parent", "staged result");
1525 let first = inbox.deliver(&staged).await.unwrap();
1526 let immediate = SessionMessageEnvelope::user_input("parent", "immediate steer");
1527 let second = inbox
1528 .deliver_with_activation_intent(
1529 &immediate,
1530 SessionActivationPolicy::RespectSpecificWait,
1531 None,
1532 )
1533 .await
1534 .unwrap();
1535 let reopened_store = Arc::new(SessionStoreV2::new(temp.path().into()).await.unwrap());
1536 let restarted = Arc::new(FileSessionInbox::new(
1537 reopened_store,
1538 SessionInboxLimits::default(),
1539 ));
1540 let router = SessionActivationRouter::new();
1541 router.set_inbox(restarted.clone());
1542 let spawner = spawner();
1543 router.set_spawner(spawner.clone()).await;
1544 let _observer = router.subscribe("parent").await;
1547 let backlog = restarted.inspect("parent").await.unwrap();
1548 assert_eq!(
1549 router
1550 .request_activation("parent", backlog.activation_generation)
1551 .await
1552 .unwrap(),
1553 SessionActivationDisposition::ActivationReserved
1554 );
1555 let claims = restarted.claim("parent", 128).await.unwrap();
1556 assert_eq!(claims.len(), 1);
1557 assert_eq!(claims[0].generation, second.generation);
1558 restarted.ack("parent", &claims[0]).await.unwrap();
1559 let mut registration = router.register_run("parent", "run-1").await.unwrap();
1560 registration.begin_finalization().await;
1561 if release_during_finalization {
1562 restarted
1563 .mark_activation_eligible(
1564 "parent",
1565 first.generation,
1566 SessionActivationPolicy::RespectSpecificWait,
1567 )
1568 .await
1569 .unwrap();
1570 assert_eq!(
1571 router
1572 .request_activation("parent", first.generation)
1573 .await
1574 .unwrap(),
1575 SessionActivationDisposition::ActivationCoalesced
1576 );
1577 assert_eq!(
1578 registration.finish(second.generation).await.unwrap(),
1579 Some(SessionActivationDisposition::ActivationReserved)
1580 );
1581 } else {
1582 assert_eq!(registration.finish(second.generation).await.unwrap(), None);
1583 restarted
1584 .mark_activation_eligible(
1585 "parent",
1586 first.generation,
1587 SessionActivationPolicy::RespectSpecificWait,
1588 )
1589 .await
1590 .unwrap();
1591 assert_eq!(
1592 router
1593 .request_activation("parent", first.generation)
1594 .await
1595 .unwrap(),
1596 SessionActivationDisposition::ActivationReserved
1597 );
1598 }
1599 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
1600 let claims = restarted.claim("parent", 128).await.unwrap();
1601 assert_eq!(claims.len(), 1);
1602 assert_eq!(claims[0].envelope.id, staged.id);
1603 let mut registration = router.register_run("parent", "run-2").await.unwrap();
1606 registration.begin_finalization().await;
1607 assert_eq!(registration.finish(second.generation).await.unwrap(), None);
1608 assert_eq!(
1609 router
1610 .request_activation("parent", first.generation)
1611 .await
1612 .unwrap(),
1613 SessionActivationDisposition::ActivationCoalesced
1614 );
1615 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
1616 restarted.ack("parent", &claims[0]).await.unwrap();
1617 }
1618 }
1619
1620 struct RecordingSpawner {
1621 reservations: AtomicUsize,
1622 launches: Arc<AtomicUsize>,
1623 entered: Option<Arc<Barrier>>,
1624 release: Option<Arc<Notify>>,
1625 }
1626
1627 #[async_trait]
1628 impl SessionActivationSpawner for RecordingSpawner {
1629 async fn reserve_activation(
1630 &self,
1631 _target_session_id: &str,
1632 _inbox_generation: u64,
1633 ) -> Result<SessionActivationReserveOutcome, SessionActivationError> {
1634 let ordinal = self.reservations.fetch_add(1, Ordering::SeqCst) + 1;
1635 if let Some(entered) = &self.entered {
1636 entered.wait().await;
1637 }
1638 if let Some(release) = &self.release {
1639 release.notified().await;
1640 }
1641 let launches = self.launches.clone();
1642 Ok(SessionActivationReserveOutcome::Reserved(
1643 SessionActivationLaunch::new(format!("run-{ordinal}"), move || {
1644 launches.fetch_add(1, Ordering::SeqCst);
1645 }),
1646 ))
1647 }
1648 }
1649
1650 fn spawner() -> Arc<RecordingSpawner> {
1651 Arc::new(RecordingSpawner {
1652 reservations: AtomicUsize::new(0),
1653 launches: Arc::new(AtomicUsize::new(0)),
1654 entered: None,
1655 release: None,
1656 })
1657 }
1658
1659 struct BlockingInspectInbox {
1660 entered: Arc<Notify>,
1661 release: Arc<Notify>,
1662 drained: Arc<std::sync::atomic::AtomicBool>,
1663 fail_once: Option<Arc<std::sync::atomic::AtomicBool>>,
1664 }
1665
1666 #[async_trait]
1667 impl SessionInboxPort for BlockingInspectInbox {
1668 async fn coordinator_activation_generation(
1669 &self,
1670 _target_session_id: &str,
1671 ) -> Result<u64, bamboo_domain::SessionInboxError> {
1672 Ok(1)
1673 }
1674
1675 async fn deliver(
1676 &self,
1677 _envelope: &bamboo_domain::SessionMessageEnvelope,
1678 ) -> Result<bamboo_domain::SessionInboxReceipt, bamboo_domain::SessionInboxError> {
1679 unreachable!("router cleanup only inspects the durable inbox")
1680 }
1681
1682 async fn mark_activation_eligible(
1683 &self,
1684 _target_session_id: &str,
1685 _generation: u64,
1686 _policy: bamboo_domain::SessionActivationPolicy,
1687 ) -> Result<(), bamboo_domain::SessionInboxError> {
1688 unreachable!("router cleanup only inspects the durable inbox")
1689 }
1690
1691 async fn claim(
1692 &self,
1693 _target_session_id: &str,
1694 _limit: usize,
1695 ) -> Result<Vec<bamboo_domain::SessionInboxClaim>, bamboo_domain::SessionInboxError>
1696 {
1697 unreachable!("router cleanup only inspects the durable inbox")
1698 }
1699
1700 async fn was_admitted(
1701 &self,
1702 _target_session_id: &str,
1703 _id: &bamboo_domain::SessionMessageId,
1704 ) -> Result<bool, bamboo_domain::SessionInboxError> {
1705 unreachable!("router cleanup only inspects the durable inbox")
1706 }
1707
1708 async fn ack(
1709 &self,
1710 _target_session_id: &str,
1711 _claim: &bamboo_domain::SessionInboxClaim,
1712 ) -> Result<(), bamboo_domain::SessionInboxError> {
1713 unreachable!("router cleanup only inspects the durable inbox")
1714 }
1715
1716 async fn inspect(
1717 &self,
1718 _target_session_id: &str,
1719 ) -> Result<bamboo_domain::SessionInboxBacklog, bamboo_domain::SessionInboxError> {
1720 if self
1721 .fail_once
1722 .as_ref()
1723 .is_some_and(|flag| flag.swap(false, Ordering::SeqCst))
1724 {
1725 return Err(bamboo_domain::SessionInboxError::Storage(
1726 "injected inspect failure".into(),
1727 ));
1728 }
1729 if self.drained.load(Ordering::SeqCst) {
1730 return Ok(bamboo_domain::SessionInboxBacklog {
1731 generation: 1,
1732 activation_generation: 1,
1733 ..Default::default()
1734 });
1735 }
1736 self.entered.notify_one();
1737 self.release.notified().await;
1738 Ok(bamboo_domain::SessionInboxBacklog {
1739 pending: 1,
1740 claimed: 0,
1741 generation: 1,
1742 activation_generation: 1,
1743 coordinator_generation: 1,
1744 interrupt_generation: 1,
1745 oldest_generation: Some(1),
1746 })
1747 }
1748 }
1749
1750 #[tokio::test]
1751 async fn rollback_completion_is_retained_before_waiter_subscribes() {
1752 let completion = RollbackCompletion::new();
1753 completion.complete();
1754 tokio::time::timeout(std::time::Duration::from_millis(100), completion.wait())
1755 .await
1756 .expect("completion-before-wait must not lose its wake");
1757 }
1758
1759 #[derive(Clone, Copy)]
1760 enum InjectedFirstOutcome {
1761 NoWork,
1762 Error,
1763 }
1764
1765 struct RetrySpawner {
1766 reservations: AtomicUsize,
1767 launches: Arc<AtomicUsize>,
1768 first_outcome: InjectedFirstOutcome,
1769 first_entered: Arc<Barrier>,
1770 release_first: Arc<Notify>,
1771 }
1772
1773 struct CancelFirstSpawner {
1774 reservations: AtomicUsize,
1775 launches: Arc<AtomicUsize>,
1776 first_entered: Arc<Notify>,
1777 }
1778
1779 struct PanicFirstSpawner {
1780 reservations: AtomicUsize,
1781 launches: Arc<AtomicUsize>,
1782 }
1783
1784 struct AlreadyRunningThenReserveSpawner {
1785 reservations: AtomicUsize,
1786 launches: Arc<AtomicUsize>,
1787 }
1788
1789 #[async_trait]
1790 impl SessionActivationSpawner for AlreadyRunningThenReserveSpawner {
1791 async fn reserve_activation(
1792 &self,
1793 _target_session_id: &str,
1794 _inbox_generation: u64,
1795 ) -> Result<SessionActivationReserveOutcome, SessionActivationError> {
1796 let ordinal = self.reservations.fetch_add(1, Ordering::SeqCst) + 1;
1797 if ordinal == 1 {
1798 return Ok(SessionActivationReserveOutcome::AlreadyRunning {
1799 run_id: "adopted-run".to_string(),
1800 });
1801 }
1802 let launches = self.launches.clone();
1803 Ok(SessionActivationReserveOutcome::Reserved(
1804 SessionActivationLaunch::new("successor-run", move || {
1805 launches.fetch_add(1, Ordering::SeqCst);
1806 }),
1807 ))
1808 }
1809 }
1810
1811 #[async_trait]
1812 impl SessionActivationSpawner for CancelFirstSpawner {
1813 async fn reserve_activation(
1814 &self,
1815 _target_session_id: &str,
1816 _inbox_generation: u64,
1817 ) -> Result<SessionActivationReserveOutcome, SessionActivationError> {
1818 let ordinal = self.reservations.fetch_add(1, Ordering::SeqCst) + 1;
1819 if ordinal == 1 {
1820 self.first_entered.notify_one();
1821 std::future::pending::<()>().await;
1822 }
1823 let launches = self.launches.clone();
1824 Ok(SessionActivationReserveOutcome::Reserved(
1825 SessionActivationLaunch::new(format!("run-{ordinal}"), move || {
1826 launches.fetch_add(1, Ordering::SeqCst);
1827 }),
1828 ))
1829 }
1830 }
1831
1832 #[async_trait]
1833 impl SessionActivationSpawner for PanicFirstSpawner {
1834 async fn reserve_activation(
1835 &self,
1836 _target_session_id: &str,
1837 _inbox_generation: u64,
1838 ) -> Result<SessionActivationReserveOutcome, SessionActivationError> {
1839 let ordinal = self.reservations.fetch_add(1, Ordering::SeqCst) + 1;
1840 if ordinal == 1 {
1841 panic!("injected first activation-spawner panic");
1842 }
1843 let launches = self.launches.clone();
1844 Ok(SessionActivationReserveOutcome::Reserved(
1845 SessionActivationLaunch::new(format!("run-{ordinal}"), move || {
1846 launches.fetch_add(1, Ordering::SeqCst);
1847 }),
1848 ))
1849 }
1850 }
1851
1852 struct RollbackSpawner {
1853 entered: Arc<Barrier>,
1854 release: Arc<Notify>,
1855 launch_ready: Arc<Notify>,
1856 rollbacks: Arc<AtomicUsize>,
1857 }
1858
1859 struct RealRegistryRollbackSpawner {
1860 runners: Arc<
1861 tokio::sync::RwLock<std::collections::HashMap<String, crate::execution::AgentRunner>>,
1862 >,
1863 senders: Arc<
1864 tokio::sync::RwLock<
1865 std::collections::HashMap<
1866 String,
1867 tokio::sync::broadcast::Sender<bamboo_agent_core::AgentEvent>,
1868 >,
1869 >,
1870 >,
1871 sender: tokio::sync::broadcast::Sender<bamboo_agent_core::AgentEvent>,
1872 reservations: AtomicUsize,
1873 first_reserved: Arc<tokio::sync::Notify>,
1874 allow_first_return: Arc<tokio::sync::Notify>,
1875 first_returning: Arc<tokio::sync::Notify>,
1876 rollback_started: Arc<tokio::sync::Notify>,
1877 allow_rollback: Arc<tokio::sync::Notify>,
1878 launches: Arc<AtomicUsize>,
1879 }
1880
1881 #[async_trait]
1882 impl SessionActivationSpawner for RealRegistryRollbackSpawner {
1883 async fn reserve_activation(
1884 &self,
1885 target_session_id: &str,
1886 _inbox_generation: u64,
1887 ) -> Result<SessionActivationReserveOutcome, SessionActivationError> {
1888 let ordinal = self.reservations.fetch_add(1, Ordering::SeqCst) + 1;
1889 match crate::execution::reserve_runner_core(
1890 &self.runners,
1891 &self.senders,
1892 target_session_id,
1893 &self.sender,
1894 )
1895 .await
1896 {
1897 crate::execution::ReserveOutcome::AlreadyRunning(run_id) => {
1898 Ok(SessionActivationReserveOutcome::AlreadyRunning { run_id })
1899 }
1900 crate::execution::ReserveOutcome::Reserved(reservation) => {
1901 let run_id = reservation.run_id.clone();
1902 if ordinal == 1 {
1903 self.first_reserved.notify_one();
1904 self.allow_first_return.notified().await;
1905 self.first_returning.notify_one();
1906 }
1907 let rollback_runners = self.runners.clone();
1908 let rollback_session_id = target_session_id.to_string();
1909 let rollback_run_id = run_id.clone();
1910 let rollback_started = self.rollback_started.clone();
1911 let allow_rollback = self.allow_rollback.clone();
1912 let launches = self.launches.clone();
1913 Ok(SessionActivationReserveOutcome::Reserved(
1914 SessionActivationLaunch::new_with_async_rollback(
1915 run_id,
1916 move || {
1917 launches.fetch_add(1, Ordering::SeqCst);
1918 },
1919 move || async move {
1920 rollback_started.notify_one();
1921 allow_rollback.notified().await;
1922 let mut runners = rollback_runners.write().await;
1923 if runners
1924 .get(&rollback_session_id)
1925 .is_some_and(|runner| runner.run_id == rollback_run_id)
1926 {
1927 runners.remove(&rollback_session_id);
1928 }
1929 },
1930 ),
1931 ))
1932 }
1933 }
1934 }
1935 }
1936
1937 #[async_trait]
1938 impl SessionActivationSpawner for RollbackSpawner {
1939 async fn reserve_activation(
1940 &self,
1941 _target_session_id: &str,
1942 _inbox_generation: u64,
1943 ) -> Result<SessionActivationReserveOutcome, SessionActivationError> {
1944 self.entered.wait().await;
1945 self.release.notified().await;
1946 let rollbacks = self.rollbacks.clone();
1947 let outcome = SessionActivationReserveOutcome::Reserved(
1948 SessionActivationLaunch::new_with_rollback(
1949 "reserved-run",
1950 || {},
1951 move || {
1952 rollbacks.fetch_add(1, Ordering::SeqCst);
1953 },
1954 ),
1955 );
1956 self.launch_ready.notify_one();
1957 Ok(outcome)
1958 }
1959 }
1960
1961 #[async_trait]
1962 impl SessionActivationSpawner for RetrySpawner {
1963 async fn reserve_activation(
1964 &self,
1965 _target_session_id: &str,
1966 _inbox_generation: u64,
1967 ) -> Result<SessionActivationReserveOutcome, SessionActivationError> {
1968 let ordinal = self.reservations.fetch_add(1, Ordering::SeqCst) + 1;
1969 if ordinal == 1 {
1970 self.first_entered.wait().await;
1971 self.release_first.notified().await;
1972 return match self.first_outcome {
1973 InjectedFirstOutcome::NoWork => Ok(SessionActivationReserveOutcome::NoWork),
1974 InjectedFirstOutcome::Error => Err(SessionActivationError::Internal(
1975 "injected first reservation failure".to_string(),
1976 )),
1977 };
1978 }
1979 let launches = self.launches.clone();
1980 Ok(SessionActivationReserveOutcome::Reserved(
1981 SessionActivationLaunch::new(format!("run-{ordinal}"), move || {
1982 launches.fetch_add(1, Ordering::SeqCst);
1983 }),
1984 ))
1985 }
1986 }
1987
1988 async fn assert_newer_generation_retries_after(
1989 first_outcome: InjectedFirstOutcome,
1990 ) -> Result<SessionActivationDisposition, SessionActivationError> {
1991 let router = SessionActivationRouter::new();
1992 let first_entered = Arc::new(Barrier::new(2));
1993 let release_first = Arc::new(Notify::new());
1994 let spawner = Arc::new(RetrySpawner {
1995 reservations: AtomicUsize::new(0),
1996 launches: Arc::new(AtomicUsize::new(0)),
1997 first_outcome,
1998 first_entered: first_entered.clone(),
1999 release_first: release_first.clone(),
2000 });
2001 router.set_spawner(spawner.clone()).await;
2002
2003 let first_router = router.clone();
2004 let first =
2005 tokio::spawn(async move { first_router.request_activation("session", 1).await });
2006 first_entered.wait().await;
2007
2008 let second_router = router.clone();
2009 let second =
2010 tokio::spawn(async move { second_router.request_activation("session", 2).await });
2011 tokio::task::yield_now().await;
2012 assert_eq!(
2013 spawner.reservations.load(Ordering::SeqCst),
2014 1,
2015 "generation 2 must coalesce while generation 1 owns the reservation"
2016 );
2017
2018 release_first.notify_one();
2019 let first_result = first.await.unwrap();
2020 assert_eq!(
2021 second.await.unwrap().unwrap(),
2022 SessionActivationDisposition::ActivationReserved,
2023 "the coalesced newer delivery must take the released reservation"
2024 );
2025 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
2026 assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
2027 first_result
2028 }
2029
2030 #[tokio::test]
2031 async fn active_owner_is_notified_without_second_reservation() {
2032 let router = SessionActivationRouter::new();
2033 let spawner = spawner();
2034 router.set_spawner(spawner.clone()).await;
2035 let mut registration = router.register_run("session", "run-live").await.unwrap();
2036
2037 assert_eq!(
2038 router.request_activation("session", 7).await.unwrap(),
2039 SessionActivationDisposition::ActiveNotified
2040 );
2041 registration.notifications_mut().changed().await.unwrap();
2042 assert_eq!(*registration.notifications_mut().borrow(), 7);
2043 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 0);
2044 registration.begin_finalization().await;
2045 assert_eq!(registration.finish(7).await.unwrap(), None);
2046 }
2047
2048 #[tokio::test]
2049 async fn different_run_registration_never_overwrites_live_owner() {
2050 let router = SessionActivationRouter::new();
2051 let mut first = router
2052 .register_run("session", "run-1")
2053 .await
2054 .expect("first run owns the logical session");
2055
2056 let error = router
2057 .register_run("session", "run-2")
2058 .await
2059 .err()
2060 .expect("a second exact run must collide");
2061 assert_eq!(
2062 error,
2063 SessionRunRegistrationError::OwnerCollision {
2064 target_session_id: "session".to_string(),
2065 existing_run_id: "run-1".to_string(),
2066 attempted_run_id: "run-2".to_string(),
2067 }
2068 );
2069 assert!(router.owns_run("session", "run-1").await);
2070 assert!(!router.owns_run("session", "run-2").await);
2071
2072 first.begin_finalization().await;
2073 assert_eq!(first.finish(0).await.unwrap(), None);
2074 }
2075
2076 #[tokio::test]
2077 async fn direct_registration_waits_for_activation_to_publish_its_exact_owner() {
2078 let router = SessionActivationRouter::new();
2079 let reservation_entered = Arc::new(Barrier::new(2));
2080 let allow_reservation = Arc::new(Notify::new());
2081 let launches = Arc::new(AtomicUsize::new(0));
2082 let spawner = Arc::new(RecordingSpawner {
2083 reservations: AtomicUsize::new(0),
2084 launches: launches.clone(),
2085 entered: Some(reservation_entered.clone()),
2086 release: Some(allow_reservation.clone()),
2087 });
2088 router.set_spawner(spawner.clone()).await;
2089
2090 let activation_router = router.clone();
2091 let activation =
2092 tokio::spawn(async move { activation_router.request_activation("session", 1).await });
2093 reservation_entered.wait().await;
2094
2095 let direct_router = router.clone();
2096 let direct =
2097 tokio::spawn(async move { direct_router.register_run("session", "manual-run").await });
2098 tokio::task::yield_now().await;
2099 assert!(
2100 !direct.is_finished(),
2101 "manual registration must not supersede an in-flight activation token"
2102 );
2103 assert!(!router.owns_run("session", "manual-run").await);
2104
2105 allow_reservation.notify_one();
2106 assert_eq!(
2107 activation.await.unwrap().unwrap(),
2108 SessionActivationDisposition::ActivationReserved
2109 );
2110 let collision = match direct.await.unwrap() {
2111 Ok(_) => panic!("manual registration must collide with the published activation"),
2112 Err(error) => error,
2113 };
2114 assert_eq!(
2115 collision,
2116 SessionRunRegistrationError::OwnerCollision {
2117 target_session_id: "session".to_string(),
2118 existing_run_id: "run-1".to_string(),
2119 attempted_run_id: "manual-run".to_string(),
2120 }
2121 );
2122 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
2123 assert_eq!(launches.load(Ordering::SeqCst), 1);
2124 assert!(router.owns_run("session", "run-1").await);
2125 assert!(!router.owns_run("session", "manual-run").await);
2126
2127 let mut registration = router.register_run("session", "run-1").await.unwrap();
2128 registration.begin_finalization().await;
2129 assert_eq!(registration.finish(1).await.unwrap(), None);
2130 }
2131
2132 #[tokio::test]
2133 async fn cancelled_explicit_abandon_still_finishes_owned_registration_cleanup() {
2134 let router = SessionActivationRouter::new();
2135 let cleanup_entered = Arc::new(Notify::new());
2136 let cleanup_release = Arc::new(Notify::new());
2137 let cleanup_completed = Arc::new(AtomicUsize::new(0));
2138 let mut registration = router.register_run("session", "run-abandon").await.unwrap();
2139 let entered = cleanup_entered.clone();
2140 let release = cleanup_release.clone();
2141 let completed = cleanup_completed.clone();
2142 registration.set_abort_cleanup(move || async move {
2143 entered.notify_one();
2144 release.notified().await;
2145 completed.fetch_add(1, Ordering::SeqCst);
2146 });
2147
2148 let abandon = tokio::spawn(registration.abandon());
2149 tokio::time::timeout(
2150 std::time::Duration::from_secs(1),
2151 cleanup_entered.notified(),
2152 )
2153 .await
2154 .expect("explicit abandon must start its owned cleanup");
2155 abandon.abort();
2156 assert!(abandon.await.unwrap_err().is_cancelled());
2157 cleanup_release.notify_one();
2158
2159 tokio::time::timeout(std::time::Duration::from_secs(1), async {
2160 loop {
2161 if cleanup_completed.load(Ordering::SeqCst) == 1
2162 && !router.owns_run("session", "run-abandon").await
2163 {
2164 break;
2165 }
2166 tokio::task::yield_now().await;
2167 }
2168 })
2169 .await
2170 .expect("caller cancellation must not cancel detached registration cleanup");
2171 }
2172
2173 #[tokio::test]
2174 async fn delayed_same_run_cannot_adopt_an_owner_during_abandoned_cleanup() {
2175 let router = SessionActivationRouter::new();
2176 let spawner = spawner();
2177 router.set_spawner(spawner.clone()).await;
2178 let inspect_entered = Arc::new(Notify::new());
2179 let inspect_release = Arc::new(Notify::new());
2180 let drained = Arc::new(std::sync::atomic::AtomicBool::new(false));
2181 router.set_inbox(Arc::new(BlockingInspectInbox {
2182 entered: inspect_entered.clone(),
2183 release: inspect_release.clone(),
2184 drained: drained.clone(),
2185 fail_once: None,
2186 }));
2187
2188 let registration = router
2189 .register_run("session", "run-abandoned")
2190 .await
2191 .unwrap();
2192 drop(registration);
2193 tokio::time::timeout(
2194 std::time::Duration::from_secs(1),
2195 inspect_entered.notified(),
2196 )
2197 .await
2198 .expect("abandoned cleanup must reach durable inbox reconciliation");
2199
2200 let error = router
2201 .register_run("session", "run-abandoned")
2202 .await
2203 .err()
2204 .expect("a finalizing abandoned owner is not an adoptable launch placeholder");
2205 assert_eq!(error.existing_run_id(), "run-abandoned");
2206
2207 inspect_release.notify_one();
2208 tokio::time::timeout(std::time::Duration::from_secs(1), async {
2209 while spawner.launches.load(Ordering::SeqCst) == 0 {
2210 tokio::task::yield_now().await;
2211 }
2212 })
2213 .await
2214 .expect("durable pending work must launch one successor");
2215 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
2216 assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
2217 assert!(!router.owns_run("session", "run-abandoned").await);
2218
2219 let mut successor = router.register_run("session", "run-1").await.unwrap();
2220 drained.store(true, Ordering::SeqCst);
2221 successor.begin_finalization().await;
2222 assert_eq!(successor.finish(1).await.unwrap(), None);
2223 }
2224
2225 #[tokio::test]
2226 async fn delayed_registry_owner_collision_rolls_back_only_its_exact_runner() {
2227 let router = SessionActivationRouter::new();
2228 let runners = Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new()));
2229 let senders = Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new()));
2230 let (sender, _receiver) = tokio::sync::broadcast::channel(8);
2231 let reserved =
2232 match crate::execution::reserve_runner_core(&runners, &senders, "session", &sender)
2233 .await
2234 {
2235 crate::execution::ReserveOutcome::Reserved(reservation) => reservation,
2236 crate::execution::ReserveOutcome::AlreadyRunning(_) => {
2237 panic!("fixture must reserve a fresh server runner")
2238 }
2239 };
2240
2241 let mut direct = router.register_run("session", "sdk-direct").await.unwrap();
2245 let collision = router
2246 .register_run("session", &reserved.run_id)
2247 .await
2248 .err()
2249 .expect("delayed server registration must collide");
2250 let collision_result = Err(bamboo_agent_core::AgentError::LLM(collision.to_string()));
2251 assert!(
2252 crate::execution::finalize_runner_exact(
2253 &runners,
2254 "session",
2255 &reserved.run_id,
2256 &collision_result,
2257 )
2258 .await
2259 );
2260 assert!(router.owns_run("session", "sdk-direct").await);
2261 assert!(matches!(
2262 runners
2263 .read()
2264 .await
2265 .get("session")
2266 .map(|runner| &runner.status),
2267 Some(crate::execution::AgentStatus::Error(_))
2268 ));
2269
2270 direct.begin_finalization().await;
2271 assert_eq!(direct.finish(0).await.unwrap(), None);
2272 }
2273
2274 #[tokio::test]
2275 async fn concurrent_idle_deliveries_coalesce_into_one_launch() {
2276 let router = SessionActivationRouter::new();
2277 let entered = Arc::new(Barrier::new(2));
2278 let release = Arc::new(Notify::new());
2279 let spawner = Arc::new(RecordingSpawner {
2280 reservations: AtomicUsize::new(0),
2281 launches: Arc::new(AtomicUsize::new(0)),
2282 entered: Some(entered.clone()),
2283 release: Some(release.clone()),
2284 });
2285 router.set_spawner(spawner.clone()).await;
2286
2287 let first_router = router.clone();
2288 let first =
2289 tokio::spawn(async move { first_router.request_activation("session", 1).await });
2290 entered.wait().await;
2291 let second_router = router.clone();
2292 let second =
2293 tokio::spawn(async move { second_router.request_activation("session", 2).await });
2294 release.notify_one();
2295 assert_eq!(
2296 first.await.unwrap().unwrap(),
2297 SessionActivationDisposition::ActivationReserved
2298 );
2299 assert_eq!(
2300 second.await.unwrap().unwrap(),
2301 SessionActivationDisposition::ActiveNotified
2302 );
2303 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
2304 assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
2305 }
2306
2307 #[tokio::test]
2308 async fn newer_generation_re_reserves_after_coalesced_no_work() {
2309 assert_eq!(
2310 assert_newer_generation_retries_after(InjectedFirstOutcome::NoWork)
2311 .await
2312 .unwrap(),
2313 SessionActivationDisposition::ActivationCoalesced
2314 );
2315 }
2316
2317 #[tokio::test]
2318 async fn newer_generation_re_reserves_after_coalesced_error() {
2319 let error = assert_newer_generation_retries_after(InjectedFirstOutcome::Error)
2320 .await
2321 .unwrap_err();
2322 assert!(matches!(error, SessionActivationError::Internal(_)));
2323 }
2324
2325 async fn assert_cancelled_finalization_recovers_without_redelivery(registered_owner: bool) {
2326 let router = SessionActivationRouter::new();
2327 let first_entered = Arc::new(Notify::new());
2328 let spawner = Arc::new(CancelFirstSpawner {
2329 reservations: AtomicUsize::new(0),
2330 launches: Arc::new(AtomicUsize::new(0)),
2331 first_entered: first_entered.clone(),
2332 });
2333 router.set_spawner(spawner.clone()).await;
2334
2335 let finish = if registered_owner {
2336 let mut registration = router.register_run("session", "run-old").await.unwrap();
2337 registration.begin_finalization().await;
2338 assert_eq!(
2339 router.request_activation("session", 1).await.unwrap(),
2340 SessionActivationDisposition::ActivationCoalesced
2341 );
2342 tokio::spawn(async move { registration.finish(0).await })
2343 } else {
2344 router.begin_finalization("session", "run-old").await;
2348 assert_eq!(
2349 router.request_activation("session", 1).await.unwrap(),
2350 SessionActivationDisposition::ActivationCoalesced
2351 );
2352 let finish_router = router.clone();
2353 tokio::spawn(async move {
2354 finish_router
2355 .finish_finalization("session", "run-old", 0)
2356 .await
2357 })
2358 };
2359 first_entered.notified().await;
2360 finish.abort();
2361 assert!(finish.await.unwrap_err().is_cancelled());
2362
2363 tokio::time::timeout(std::time::Duration::from_secs(1), async {
2364 while spawner.launches.load(Ordering::SeqCst) == 0 {
2365 tokio::task::yield_now().await;
2366 }
2367 })
2368 .await
2369 .expect("cancelled finalization must retry pending work without redelivery");
2370 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
2371 assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
2372 assert!(router.owns_run("session", "run-2").await);
2373
2374 let mut successor = router.register_run("session", "run-2").await.unwrap();
2375 successor.begin_finalization().await;
2376 assert_eq!(successor.finish(1).await.unwrap(), None);
2377 }
2378
2379 #[tokio::test]
2380 async fn cancelled_registered_finalization_retries_without_another_delivery() {
2381 assert_cancelled_finalization_recovers_without_redelivery(true).await;
2382 }
2383
2384 #[tokio::test]
2385 async fn cancelled_reserved_child_finalization_retries_without_another_delivery() {
2386 assert_cancelled_finalization_recovers_without_redelivery(false).await;
2387 }
2388
2389 #[tokio::test]
2390 async fn panicking_finalization_spawner_gets_one_bounded_retry() {
2391 let router = SessionActivationRouter::new();
2392 let spawner = Arc::new(PanicFirstSpawner {
2393 reservations: AtomicUsize::new(0),
2394 launches: Arc::new(AtomicUsize::new(0)),
2395 });
2396 router.set_spawner(spawner.clone()).await;
2397 router.begin_finalization("session", "run-old").await;
2398 assert_eq!(
2399 router.request_activation("session", 1).await.unwrap(),
2400 SessionActivationDisposition::ActivationCoalesced
2401 );
2402
2403 let finish_router = router.clone();
2404 let finish = tokio::spawn(async move {
2405 finish_router
2406 .finish_finalization("session", "run-old", 0)
2407 .await
2408 });
2409 assert!(finish.await.unwrap_err().is_panic());
2410
2411 tokio::time::timeout(std::time::Duration::from_secs(1), async {
2412 while spawner.launches.load(Ordering::SeqCst) == 0 {
2413 tokio::task::yield_now().await;
2414 }
2415 })
2416 .await
2417 .expect("router Drop recovery must survive one spawner panic");
2418 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
2419 assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
2420
2421 let mut successor = router.register_run("session", "run-2").await.unwrap();
2422 successor.begin_finalization().await;
2423 assert_eq!(successor.finish(1).await.unwrap(), None);
2424 }
2425
2426 #[tokio::test]
2427 async fn cancelled_reservation_releases_token_and_newer_generation_launches() {
2428 let router = SessionActivationRouter::new();
2429 let first_entered = Arc::new(Notify::new());
2430 let spawner = Arc::new(CancelFirstSpawner {
2431 reservations: AtomicUsize::new(0),
2432 launches: Arc::new(AtomicUsize::new(0)),
2433 first_entered: first_entered.clone(),
2434 });
2435 router.set_spawner(spawner.clone()).await;
2436
2437 let first_router = router.clone();
2438 let first =
2439 tokio::spawn(async move { first_router.request_activation("session", 1).await });
2440 first_entered.notified().await;
2441 let second_router = router.clone();
2442 let second =
2443 tokio::spawn(async move { second_router.request_activation("session", 2).await });
2444 tokio::task::yield_now().await;
2445 first.abort();
2446 assert!(first.await.unwrap_err().is_cancelled());
2447
2448 assert_eq!(
2449 tokio::time::timeout(std::time::Duration::from_secs(1), second)
2450 .await
2451 .expect("coalesced caller must not hang behind a cancelled reservation")
2452 .unwrap()
2453 .unwrap(),
2454 SessionActivationDisposition::ActivationReserved
2455 );
2456 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
2457 assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
2458 }
2459
2460 #[tokio::test]
2461 async fn cancellation_after_external_reservation_rolls_back_unlaunched_slot() {
2462 let router = SessionActivationRouter::new();
2463 let entered = Arc::new(Barrier::new(2));
2464 let release = Arc::new(Notify::new());
2465 let launch_ready = Arc::new(Notify::new());
2466 let rollbacks = Arc::new(AtomicUsize::new(0));
2467 let spawner = Arc::new(RollbackSpawner {
2468 entered: entered.clone(),
2469 release: release.clone(),
2470 launch_ready: launch_ready.clone(),
2471 rollbacks: rollbacks.clone(),
2472 });
2473 router.set_spawner(spawner).await;
2474
2475 let task_router = router.clone();
2476 let task = tokio::spawn(async move { task_router.request_activation("session", 1).await });
2477 entered.wait().await;
2478 let states = router.states.lock().await;
2482 release.notify_one();
2483 launch_ready.notified().await;
2484 task.abort();
2485 assert!(task.await.unwrap_err().is_cancelled());
2486 drop(states);
2487
2488 tokio::time::timeout(std::time::Duration::from_secs(1), async {
2489 while rollbacks.load(Ordering::SeqCst) == 0 {
2490 tokio::task::yield_now().await;
2491 }
2492 })
2493 .await
2494 .expect("unlaunched reservation rollback");
2495 assert_eq!(rollbacks.load(Ordering::SeqCst), 1);
2496 }
2497
2498 #[tokio::test]
2499 async fn redelivery_waits_for_real_runner_rollback_before_reserving() {
2500 let router = SessionActivationRouter::new();
2501 let runners = Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new()));
2502 let senders = Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new()));
2503 let (sender, _receiver) = tokio::sync::broadcast::channel(8);
2504 let first_reserved = Arc::new(tokio::sync::Notify::new());
2505 let allow_first_return = Arc::new(tokio::sync::Notify::new());
2506 let rollback_started = Arc::new(tokio::sync::Notify::new());
2507 let first_returning = Arc::new(tokio::sync::Notify::new());
2508 let allow_rollback = Arc::new(tokio::sync::Notify::new());
2509 let launches = Arc::new(AtomicUsize::new(0));
2510 let spawner = Arc::new(RealRegistryRollbackSpawner {
2511 runners: runners.clone(),
2512 senders,
2513 sender,
2514 reservations: AtomicUsize::new(0),
2515 first_reserved: first_reserved.clone(),
2516 allow_first_return: allow_first_return.clone(),
2517 first_returning: first_returning.clone(),
2518 rollback_started: rollback_started.clone(),
2519 allow_rollback: allow_rollback.clone(),
2520 launches: launches.clone(),
2521 });
2522 router.set_spawner(spawner.clone()).await;
2523
2524 let first_router = router.clone();
2525 let first =
2526 tokio::spawn(async move { first_router.request_activation("session", 1).await });
2527 first_reserved.notified().await;
2528
2529 let state_guard = router.states.lock().await;
2532 allow_first_return.notify_one();
2533 first_returning.notified().await;
2534 first.abort();
2535 assert!(first.await.unwrap_err().is_cancelled());
2536 drop(state_guard);
2537 rollback_started.notified().await;
2538
2539 let second_router = router.clone();
2540 let second =
2541 tokio::spawn(async move { second_router.request_activation("session", 2).await });
2542 tokio::task::yield_now().await;
2543 assert_eq!(
2544 spawner.reservations.load(Ordering::SeqCst),
2545 1,
2546 "redelivery must remain coalesced until exact external rollback completes"
2547 );
2548
2549 allow_rollback.notify_one();
2550 assert_eq!(
2551 tokio::time::timeout(std::time::Duration::from_secs(1), second)
2552 .await
2553 .expect("redelivery must reserve after rollback")
2554 .unwrap()
2555 .unwrap(),
2556 SessionActivationDisposition::ActivationReserved
2557 );
2558 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
2559 assert_eq!(launches.load(Ordering::SeqCst), 1);
2560 let current = runners
2561 .read()
2562 .await
2563 .get("session")
2564 .map(|runner| runner.run_id.clone())
2565 .expect("second real registry reservation remains");
2566 assert_ne!(
2567 current, "",
2568 "the redelivery must never adopt an empty/stale run id"
2569 );
2570 }
2571
2572 #[tokio::test]
2573 async fn delivery_during_finalization_launches_one_successor() {
2574 let router = SessionActivationRouter::new();
2575 let spawner = spawner();
2576 router.set_spawner(spawner.clone()).await;
2577 let mut registration = router.register_run("session", "run-old").await.unwrap();
2578 registration.begin_finalization().await;
2579
2580 assert_eq!(
2581 router.request_activation("session", 11).await.unwrap(),
2582 SessionActivationDisposition::ActivationCoalesced
2583 );
2584 let disposition = registration.finish(10).await.unwrap();
2585 assert_eq!(
2586 disposition,
2587 Some(SessionActivationDisposition::ActivationReserved)
2588 );
2589 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
2590 assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
2591 }
2592
2593 #[tokio::test]
2594 async fn finalization_does_not_restart_when_cursor_caught_up() {
2595 let router = SessionActivationRouter::new();
2596 let spawner = spawner();
2597 router.set_spawner(spawner.clone()).await;
2598 let mut registration = router.register_run("session", "run-old").await.unwrap();
2599 assert_eq!(
2600 router.request_activation("session", 4).await.unwrap(),
2601 SessionActivationDisposition::ActiveNotified
2602 );
2603 registration.begin_finalization().await;
2604 assert_eq!(registration.finish(4).await.unwrap(), None);
2605 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 0);
2606 assert!(
2607 router.states.lock().await.get("session").is_none(),
2608 "a caught-up completed session must not leak routing state"
2609 );
2610 }
2611
2612 #[tokio::test]
2613 async fn successor_finalization_removes_state_after_raced_generation_catches_up() {
2614 let router = SessionActivationRouter::new();
2615 let spawner = spawner();
2616 router.set_spawner(spawner).await;
2617 let mut registration = router.register_run("session", "run-old").await.unwrap();
2618 registration.begin_finalization().await;
2619
2620 assert_eq!(
2621 router.request_activation("session", 11).await.unwrap(),
2622 SessionActivationDisposition::ActivationCoalesced
2623 );
2624 assert_eq!(
2625 registration.finish(10).await.unwrap(),
2626 Some(SessionActivationDisposition::ActivationReserved)
2627 );
2628 let mut successor_registration = router.register_run("session", "run-1").await.unwrap();
2629 successor_registration.begin_finalization().await;
2630 assert_eq!(successor_registration.finish(11).await.unwrap(), None);
2631 assert!(
2632 router.states.lock().await.get("session").is_none(),
2633 "the successor's caught-up terminal state must be compacted"
2634 );
2635 }
2636
2637 #[tokio::test]
2638 async fn failed_terminal_inspection_keeps_the_exact_owner_cleanup_armed() {
2639 let router = SessionActivationRouter::new();
2640 let spawner = spawner();
2641 router.set_spawner(spawner.clone()).await;
2642 let inspect_release = Arc::new(Notify::new());
2643 inspect_release.notify_one();
2644 router.set_inbox(Arc::new(BlockingInspectInbox {
2645 entered: Arc::new(Notify::new()),
2646 release: inspect_release,
2647 drained: Arc::new(std::sync::atomic::AtomicBool::new(false)),
2648 fail_once: Some(Arc::new(std::sync::atomic::AtomicBool::new(true))),
2649 }));
2650 let mut registration = router
2651 .register_run("session", "failed-inspect-run")
2652 .await
2653 .unwrap();
2654 router.request_activation("session", 1).await.unwrap();
2655 registration.begin_finalization().await;
2656 assert!(registration.finish(1).await.is_err());
2657 tokio::time::timeout(std::time::Duration::from_secs(1), async {
2658 while spawner.launches.load(Ordering::SeqCst) == 0 {
2659 tokio::task::yield_now().await;
2660 }
2661 })
2662 .await
2663 .expect("failed inspection must not leave the old run owning the inbox");
2664 assert!(!router.owns_run("session", "failed-inspect-run").await);
2665 assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
2666 }
2667
2668 #[tokio::test]
2669 async fn poison_generation_launches_only_one_successor_until_new_work_arrives() {
2670 let router = SessionActivationRouter::new();
2671 let spawner = spawner();
2672 router.set_spawner(spawner.clone()).await;
2673 let inspect_release = Arc::new(Notify::new());
2674 router.set_inbox(Arc::new(BlockingInspectInbox {
2675 entered: Arc::new(Notify::new()),
2676 release: inspect_release.clone(),
2677 drained: Arc::new(std::sync::atomic::AtomicBool::new(false)),
2678 fail_once: None,
2679 }));
2680 let mut registration = router
2681 .register_run("session", "run-original")
2682 .await
2683 .unwrap();
2684 assert_eq!(
2685 router.request_activation("session", 7).await.unwrap(),
2686 SessionActivationDisposition::ActiveNotified
2687 );
2688 registration.begin_finalization().await;
2689 inspect_release.notify_one();
2690 assert_eq!(
2691 registration.finish(0).await.unwrap(),
2692 Some(SessionActivationDisposition::ActivationReserved)
2693 );
2694 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
2695
2696 let mut successor_registration = router.register_run("session", "run-1").await.unwrap();
2700 successor_registration.begin_finalization().await;
2701 inspect_release.notify_one();
2702 assert_eq!(successor_registration.finish(0).await.unwrap(), None);
2703 assert_eq!(
2704 router.request_activation("session", 7).await.unwrap(),
2705 SessionActivationDisposition::ActivationCoalesced
2706 );
2707 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 1);
2708
2709 assert_eq!(
2711 router.request_activation("session", 8).await.unwrap(),
2712 SessionActivationDisposition::ActivationReserved
2713 );
2714 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
2715 }
2716
2717 #[tokio::test]
2718 async fn adopted_terminal_run_preserves_generation_for_one_real_successor() {
2719 let router = SessionActivationRouter::new();
2720 let spawner = Arc::new(AlreadyRunningThenReserveSpawner {
2721 reservations: AtomicUsize::new(0),
2722 launches: Arc::new(AtomicUsize::new(0)),
2723 });
2724 router.set_spawner(spawner.clone()).await;
2725
2726 assert_eq!(
2727 router.request_activation("session", 12).await.unwrap(),
2728 SessionActivationDisposition::ActiveNotified
2729 );
2730 router.begin_finalization("session", "adopted-run").await;
2731 assert_eq!(
2732 router
2733 .finish_finalization("session", "adopted-run", 0)
2734 .await
2735 .unwrap(),
2736 Some(SessionActivationDisposition::ActivationReserved)
2737 );
2738 assert_eq!(spawner.reservations.load(Ordering::SeqCst), 2);
2739 assert_eq!(spawner.launches.load(Ordering::SeqCst), 1);
2740 }
2741}