1use std::any::Any;
55use std::collections::BTreeMap;
56use std::fmt;
57use std::io;
58use std::num::NonZeroU32;
59use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
60use std::sync::{Arc, Mutex, MutexGuard};
61use std::time::{Duration, Instant};
62
63use rightkit_child_process::{
64 AdoptedIdentity, Generation, IdentityCheck, OwnedChild, RestartOutcome, RestartPolicy,
65 RestartTracker, WaitOutcome,
66};
67use serde_json::Value;
68
69use crate::client::{Client, ClientConfig};
70use crate::discovery::{
71 record_process_matches, try_adopt_or_spawn, AdoptOrSpawn, SpawnClaim, SpawnClaimError,
72};
73use crate::error::{ConnectError, ServiceError, ERROR_UNAVAILABLE};
74use crate::host_identity::{HostExpectation, HostIdentityError, HostIdentityReason};
75use crate::record::{try_read_record, ServiceRecord};
76
77pub const ERROR_STALE_GENERATION: &str = "stale_generation";
80pub const ERROR_NOT_OWNED: &str = "not_owned";
82pub const ERROR_RESTART_EXHAUSTED: &str = "restart_exhausted";
84pub const ERROR_SESSION_RELEASED: &str = "session_released";
86pub const ERROR_LAUNCH_FAILED: &str = "launch_failed";
88
89fn lock<T>(m: &Mutex<T>) -> MutexGuard<'_, T> {
90 m.lock().unwrap_or_else(|e| e.into_inner())
91}
92
93#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
95pub enum Provenance {
96 Owned,
98 Attached,
100}
101
102impl Provenance {
103 pub fn as_str(self) -> &'static str {
104 match self {
105 Provenance::Owned => "owned",
106 Provenance::Attached => "attached",
107 }
108 }
109
110 pub fn may_stop_host(self) -> bool {
112 matches!(self, Provenance::Owned)
113 }
114}
115
116#[derive(Debug, Clone, PartialEq, Eq, Hash)]
124pub struct HostGeneration {
125 pub service: String,
126 pub endpoint: String,
127 pub pid: u32,
128 pub process_started_ms: u64,
130 pub started_ms: u64,
132}
133
134impl HostGeneration {
135 pub fn of(record: &ServiceRecord) -> Self {
137 Self {
138 service: record.service.clone(),
139 endpoint: record.endpoint.clone(),
140 pid: record.pid,
141 process_started_ms: record.process_started_ms,
142 started_ms: record.started_ms,
143 }
144 }
145}
146
147#[derive(Debug)]
149#[non_exhaustive]
150pub enum LifecycleError {
151 Service(ServiceError),
153 HostIdentity(HostIdentityError),
156 NoHost,
158 Launch(io::Error),
160 LaunchTimeout,
162 RestartExhausted,
164 NotOwned,
166 Released,
168 StaleGeneration { issued: u64, current: u64 },
172 Unconfirmed,
175 Process(io::Error),
177}
178
179impl LifecycleError {
180 pub fn code(&self) -> &str {
182 match self {
183 LifecycleError::Service(e) => &e.code,
184 LifecycleError::HostIdentity(_) => crate::error::ERROR_HOST_IDENTITY,
185 LifecycleError::NoHost => ERROR_UNAVAILABLE,
186 LifecycleError::Launch(_) | LifecycleError::LaunchTimeout => ERROR_LAUNCH_FAILED,
187 LifecycleError::RestartExhausted => ERROR_RESTART_EXHAUSTED,
188 LifecycleError::NotOwned => ERROR_NOT_OWNED,
189 LifecycleError::Released => ERROR_SESSION_RELEASED,
190 LifecycleError::StaleGeneration { .. } | LifecycleError::Unconfirmed => {
191 ERROR_STALE_GENERATION
192 }
193 LifecycleError::Process(_) => "io",
194 }
195 }
196
197 pub fn is_retryable(&self) -> bool {
200 match self {
201 LifecycleError::Service(e) => e.retryable,
202 LifecycleError::NoHost => true,
203 _ => false,
204 }
205 }
206
207 pub fn into_service_error(self) -> ServiceError {
209 let text = self.to_string();
210 match self {
211 LifecycleError::Service(e) => e,
212 LifecycleError::HostIdentity(e) => e.into(),
213 LifecycleError::NoHost => ServiceError::retryable(ERROR_UNAVAILABLE, text),
214 other => ServiceError::new(other.code().to_string(), text),
215 }
216 }
217}
218
219impl fmt::Display for LifecycleError {
220 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
221 match self {
222 LifecycleError::Service(e) => fmt::Display::fmt(e, f),
223 LifecycleError::HostIdentity(e) => fmt::Display::fmt(e, f),
224 LifecycleError::NoHost => f.write_str("no live host and no launcher to start one"),
225 LifecycleError::Launch(e) => write!(f, "the host could not be launched: {e}"),
226 LifecycleError::LaunchTimeout => {
227 f.write_str("the launched host did not publish itself in time; it was stopped")
228 }
229 LifecycleError::RestartExhausted => {
230 f.write_str("the owned host exited and the restart budget is spent")
231 }
232 LifecycleError::NotOwned => f.write_str(
233 "this host was not launched by this lifecycle; attaching never implies the right to stop it",
234 ),
235 LifecycleError::Released => f.write_str("the session or the lifecycle was released"),
236 LifecycleError::StaleGeneration { issued, current } => write!(
237 f,
238 "the host generation changed (bound to epoch {issued}, now {current}); the request may have run on the previous host and is not replayed"
239 ),
240 LifecycleError::Unconfirmed => f.write_str(
241 "the host generation could not be confirmed; the reply is not delivered",
242 ),
243 LifecycleError::Process(e) => write!(f, "process operation failed: {e}"),
244 }
245 }
246}
247
248impl std::error::Error for LifecycleError {
249 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
250 match self {
251 LifecycleError::Service(e) => Some(e),
252 LifecycleError::HostIdentity(e) => Some(e),
253 LifecycleError::Launch(e) | LifecycleError::Process(e) => Some(e),
254 _ => None,
255 }
256 }
257}
258
259impl From<ServiceError> for LifecycleError {
260 fn from(e: ServiceError) -> Self {
261 LifecycleError::Service(e)
262 }
263}
264
265fn unavailable(e: io::Error) -> LifecycleError {
266 LifecycleError::Service(ServiceError::retryable(
267 ERROR_UNAVAILABLE,
268 format!("host discovery failed (try again): {e}"),
269 ))
270}
271
272#[derive(Debug, Clone, Copy, PartialEq, Eq)]
278pub struct Stale {
279 pub issued: u64,
280 pub current: u64,
281}
282
283impl From<Stale> for LifecycleError {
284 fn from(s: Stale) -> Self {
285 LifecycleError::StaleGeneration {
286 issued: s.issued,
287 current: s.current,
288 }
289 }
290}
291
292#[derive(Debug, Clone, Copy, PartialEq, Eq)]
294pub struct Ticket {
295 epoch: u64,
296 request: u64,
297}
298
299impl Ticket {
300 pub fn epoch(&self) -> u64 {
301 self.epoch
302 }
303
304 pub fn request(&self) -> u64 {
305 self.request
306 }
307}
308
309#[derive(Debug, Default)]
318pub struct GenerationFence {
319 epoch: AtomicU64,
320 next_request: AtomicU64,
321}
322
323impl GenerationFence {
324 pub fn new() -> Self {
325 Self::default()
326 }
327
328 pub fn epoch(&self) -> u64 {
330 self.epoch.load(Ordering::SeqCst)
331 }
332
333 pub fn advance(&self) -> u64 {
336 self.epoch.fetch_add(1, Ordering::SeqCst) + 1
337 }
338
339 pub fn issue(&self, bound_epoch: u64) -> Result<Ticket, Stale> {
342 let current = self.epoch();
343 if bound_epoch != current {
344 return Err(Stale {
345 issued: bound_epoch,
346 current,
347 });
348 }
349 Ok(Ticket {
350 epoch: bound_epoch,
351 request: self.next_request.fetch_add(1, Ordering::SeqCst) + 1,
352 })
353 }
354
355 pub fn admit_reply(&self, ticket: &Ticket) -> Result<(), Stale> {
358 let current = self.epoch();
359 if ticket.epoch != current {
360 return Err(Stale {
361 issued: ticket.epoch,
362 current,
363 });
364 }
365 Ok(())
366 }
367
368 pub fn admit_reply_from(&self, ticket: &Ticket, reply_epoch: u64) -> Result<(), Stale> {
371 if reply_epoch != ticket.epoch {
372 return Err(Stale {
373 issued: ticket.epoch,
374 current: reply_epoch,
375 });
376 }
377 self.admit_reply(ticket)
378 }
379}
380
381pub trait HostLink: Send + Sync {
387 fn call(&self, method: &str, params: Value) -> Result<Value, ServiceError>;
388 fn is_closed(&self) -> bool;
390 fn close(&self);
392}
393
394pub enum SpawnPermit {
396 Granted(Box<dyn Any + Send>),
398 Contended,
400}
401
402pub trait HostBackend: Send + Sync + 'static {
405 fn discover(&self) -> io::Result<Option<HostGeneration>>;
408
409 fn connect(&self, host: &HostGeneration) -> Result<Box<dyn HostLink>, LifecycleError>;
413
414 fn is_current(&self, host: &HostGeneration) -> io::Result<bool> {
417 Ok(self.discover()?.as_ref() == Some(host))
418 }
419
420 fn claim_spawn(&self) -> io::Result<SpawnPermit>;
422}
423
424pub struct LocalBackend {
427 client: ClientConfig,
428 expect: HostExpectation,
429}
430
431impl LocalBackend {
432 pub fn new(client: ClientConfig, expect: HostExpectation) -> Self {
436 Self { client, expect }
437 }
438
439 fn suite(&self) -> io::Result<crate::suite::Suite> {
440 self.client.resolve_suite().map_err(io::Error::other)
441 }
442}
443
444struct LocalLink(Client);
445
446impl HostLink for LocalLink {
447 fn call(&self, method: &str, params: Value) -> Result<Value, ServiceError> {
448 self.0.call(method, params)
449 }
450
451 fn is_closed(&self) -> bool {
452 self.0.is_closed()
453 }
454
455 fn close(&self) {
456 self.0.abort();
457 }
458}
459
460impl HostBackend for LocalBackend {
461 fn discover(&self) -> io::Result<Option<HostGeneration>> {
462 let suite = self.suite()?;
463 Ok(match try_adopt_or_spawn(&suite, self.client.service())? {
464 AdoptOrSpawn::Adopt(record) => Some(HostGeneration::of(&record)),
465 AdoptOrSpawn::Spawn => None,
466 })
467 }
468
469 fn connect(&self, host: &HostGeneration) -> Result<Box<dyn HostLink>, LifecycleError> {
470 let client =
471 Client::connect_expecting(&self.client, &self.expect).map_err(|e| match e {
472 ConnectError::HostIdentity(identity) => LifecycleError::HostIdentity(identity),
473 other => LifecycleError::Service(other.into_service_error()),
474 })?;
475 if HostGeneration::of(client.record()) != *host {
477 client.abort();
478 return Err(LifecycleError::Service(ServiceError::retryable(
479 ERROR_UNAVAILABLE,
480 "the host changed while connecting; discover it again",
481 )));
482 }
483 let identity_error = |found_pid: Option<u32>,
485 found_executable: Option<std::path::PathBuf>,
486 reason: HostIdentityReason| {
487 LifecycleError::HostIdentity(HostIdentityError {
488 expected: self.expect.executables().to_vec(),
489 found_pid,
490 found_executable,
491 reason,
492 })
493 };
494 match client.verified_host() {
495 Some(process) if process.pid == host.pid => {}
496 Some(process) => {
497 let error = identity_error(
498 Some(process.pid),
499 Some(process.executable.clone()),
500 HostIdentityReason::HostChanged(format!(
501 "the endpoint is served by pid {} but service.json names pid {}",
502 process.pid, host.pid
503 )),
504 );
505 client.abort();
506 return Err(error);
507 }
508 None => {
509 client.abort();
510 return Err(identity_error(
511 None,
512 None,
513 HostIdentityReason::PidUnavailable,
514 ));
515 }
516 }
517 Ok(Box::new(LocalLink(client)))
518 }
519
520 fn is_current(&self, host: &HostGeneration) -> io::Result<bool> {
521 let suite = self.suite()?;
522 Ok(try_read_record(&suite, self.client.service())?
523 .is_some_and(|r| record_process_matches(&r) && HostGeneration::of(&r) == *host))
524 }
525
526 fn claim_spawn(&self) -> io::Result<SpawnPermit> {
527 let suite = self.suite()?;
528 match SpawnClaim::acquire(&suite, self.client.service()) {
529 Ok(claim) => Ok(SpawnPermit::Granted(Box::new(claim))),
530 Err(SpawnClaimError::AlreadyHeld(_)) => Ok(SpawnPermit::Contended),
531 Err(SpawnClaimError::Io(e)) => Err(e),
532 }
533 }
534}
535
536pub type Launcher = Box<dyn FnMut() -> io::Result<OwnedChild> + Send>;
544
545pub type GracefulStop = Arc<dyn Fn(&dyn HostLink) -> Result<(), ServiceError> + Send + Sync>;
549
550#[derive(Clone)]
552pub struct LifecycleConfig {
553 startup_timeout: Duration,
554 poll_interval: Duration,
555 stop_grace: Duration,
556 restart: RestartPolicy,
557 graceful: Option<GracefulStop>,
558}
559
560impl Default for LifecycleConfig {
561 fn default() -> Self {
562 Self {
563 startup_timeout: Duration::from_secs(30),
564 poll_interval: Duration::from_millis(50),
565 stop_grace: Duration::from_secs(5),
566 restart: RestartPolicy::default(),
567 graceful: None,
568 }
569 }
570}
571
572impl LifecycleConfig {
573 pub fn new() -> Self {
574 Self::default()
575 }
576
577 #[must_use]
579 pub fn startup_timeout(mut self, timeout: Duration) -> Self {
580 self.startup_timeout = timeout;
581 self
582 }
583
584 #[must_use]
586 pub fn poll_interval(mut self, interval: Duration) -> Self {
587 self.poll_interval = interval.max(Duration::from_millis(1));
588 self
589 }
590
591 #[must_use]
594 pub fn stop_grace(mut self, grace: Duration) -> Self {
595 self.stop_grace = grace;
596 self
597 }
598
599 #[must_use]
602 pub fn restart_policy(mut self, policy: RestartPolicy) -> Self {
603 self.restart = policy;
604 self
605 }
606
607 #[must_use]
609 pub fn graceful_stop<F>(mut self, stop: F) -> Self
610 where
611 F: Fn(&dyn HostLink) -> Result<(), ServiceError> + Send + Sync + 'static,
612 {
613 self.graceful = Some(Arc::new(stop));
614 self
615 }
616}
617
618struct OwnedHost {
623 child: OwnedChild,
624 identity: AdoptedIdentity,
626 host: Option<HostGeneration>,
628}
629
630struct Slow {
631 owned: Option<OwnedHost>,
632 tracker: RestartTracker,
633 origin: Instant,
634 generation: Option<Generation>,
635 pending: Option<RestartOutcome>,
637 exhausted: bool,
638}
639
640#[derive(Clone)]
641struct Current {
642 host: HostGeneration,
643 provenance: Provenance,
644 epoch: u64,
645}
646
647struct Inner {
648 backend: Arc<dyn HostBackend>,
649 cfg: LifecycleConfig,
650 launcher: Mutex<Option<Launcher>>,
651 slow: Mutex<Slow>,
654 current: Mutex<Option<Current>>,
655 sessions: Mutex<BTreeMap<u64, String>>,
656 fence: GenerationFence,
657 next_session: AtomicU64,
658 stopped: AtomicBool,
659}
660
661pub struct DaemonLifecycle {
663 inner: Arc<Inner>,
664}
665
666#[derive(Debug, Clone, PartialEq, Eq)]
668pub struct ReleaseReport {
669 pub label: String,
670 pub provenance: Provenance,
671 pub remaining_sessions: usize,
673}
674
675#[derive(Debug, Clone, PartialEq, Eq)]
677pub struct ShutdownReport {
678 pub graceful_requested: bool,
680 pub graceful_acknowledged: bool,
682 pub killed: bool,
684 pub exit_code: Option<i32>,
686}
687
688impl DaemonLifecycle {
689 pub fn attach_or_launch(
695 backend: Arc<dyn HostBackend>,
696 launcher: Option<Launcher>,
697 config: LifecycleConfig,
698 ) -> Result<Self, LifecycleError> {
699 let tracker = RestartTracker::new(config.restart);
700 let inner = Arc::new(Inner {
701 backend,
702 cfg: config,
703 launcher: Mutex::new(launcher),
704 slow: Mutex::new(Slow {
705 owned: None,
706 tracker,
707 origin: Instant::now(),
708 generation: None,
709 pending: None,
710 exhausted: false,
711 }),
712 current: Mutex::new(None),
713 sessions: Mutex::new(BTreeMap::new()),
714 fence: GenerationFence::new(),
715 next_session: AtomicU64::new(1),
716 stopped: AtomicBool::new(false),
717 });
718 inner.ensure_host()?;
719 Ok(Self { inner })
720 }
721
722 pub fn attach_or_launch_local(
724 client: ClientConfig,
725 expect: HostExpectation,
726 launcher: Option<Launcher>,
727 config: LifecycleConfig,
728 ) -> Result<Self, LifecycleError> {
729 Self::attach_or_launch(
730 Arc::new(LocalBackend::new(client, expect)),
731 launcher,
732 config,
733 )
734 }
735
736 pub fn provenance(&self) -> Option<Provenance> {
738 lock(&self.inner.current).as_ref().map(|c| c.provenance)
739 }
740
741 pub fn host(&self) -> Option<HostGeneration> {
743 lock(&self.inner.current).as_ref().map(|c| c.host.clone())
744 }
745
746 pub fn epoch(&self) -> u64 {
748 self.inner.fence.epoch()
749 }
750
751 pub fn session_count(&self) -> usize {
753 lock(&self.inner.sessions).len()
754 }
755
756 pub fn reset_restart_budget(&self) {
758 let mut slow = lock(&self.inner.slow);
759 slow.tracker.reset_budget();
760 slow.pending = None;
761 slow.exhausted = false;
762 }
763
764 pub fn acquire_session(&self, label: &str) -> Result<Session, LifecycleError> {
768 let current = self.inner.ensure_host()?;
769 let link = self.inner.backend.connect(¤t.host)?;
770 match self.inner.backend.is_current(¤t.host) {
772 Ok(true) => {}
773 Ok(false) => {
774 link.close();
775 self.inner.host_left(¤t.host);
776 return Err(LifecycleError::StaleGeneration {
777 issued: current.epoch,
778 current: self.inner.fence.epoch(),
779 });
780 }
781 Err(_) => {
782 link.close();
783 return Err(LifecycleError::Unconfirmed);
784 }
785 }
786 let id = self.inner.next_session.fetch_add(1, Ordering::SeqCst);
787 lock(&self.inner.sessions).insert(id, label.to_string());
788 Ok(Session {
789 inner: self.inner.clone(),
790 id,
791 label: label.to_string(),
792 provenance: current.provenance,
793 host: current.host,
794 epoch: current.epoch,
795 link,
796 released: AtomicBool::new(false),
797 })
798 }
799
800 pub fn release_session(&self, session: Session) -> ReleaseReport {
804 session.finish()
805 }
806
807 pub fn shutdown_owned(&self) -> Result<ShutdownReport, LifecycleError> {
816 let inner = &self.inner;
817 let mut slow = lock(&inner.slow);
818 let identity_state = match slow.owned.as_ref() {
820 None => return Err(LifecycleError::NotOwned),
821 Some(owned) => owned.identity.verify().map_err(LifecycleError::Process)?,
822 };
823 if identity_state == IdentityCheck::Replaced {
824 return Err(LifecycleError::Process(io::Error::other(
825 "the launched pid now names a different process; nothing was signalled",
826 )));
827 }
828 let Some(mut owned) = slow.owned.take() else {
829 return Err(LifecycleError::NotOwned);
830 };
831 inner.stopped.store(true, Ordering::SeqCst);
832 slow.tracker.shutdown();
833 slow.generation = None;
834
835 let mut graceful_requested = false;
836 let mut graceful_acknowledged = false;
837 if let (Some(stop), Some(host)) = (inner.cfg.graceful.as_ref(), owned.host.as_ref()) {
838 if matches!(inner.backend.is_current(host), Ok(true)) {
840 if let Ok(link) = inner.backend.connect(host) {
841 graceful_requested = true;
842 graceful_acknowledged = stop(&*link).is_ok();
843 link.close();
844 }
845 }
846 }
847 let outcome = owned
848 .child
849 .wait_then_kill(inner.cfg.stop_grace)
850 .map_err(LifecycleError::Process)?;
851 let (killed, exit_code) = match outcome {
852 WaitOutcome::Exited(status) => (false, status.code()),
853 WaitOutcome::Terminated(status) => (true, status.code()),
854 };
855 drop(owned);
856 drop(slow);
857 let mut current = lock(&inner.current);
859 if current.take().is_some() {
860 inner.fence.advance();
861 }
862 Ok(ShutdownReport {
863 graceful_requested,
864 graceful_acknowledged,
865 killed,
866 exit_code,
867 })
868 }
869}
870
871impl Inner {
872 fn bind(&self, to: Option<(HostGeneration, Provenance)>) -> Option<Current> {
873 let mut current = lock(&self.current);
874 let unchanged = match (current.as_ref(), to.as_ref()) {
875 (Some(c), Some((host, provenance))) => c.host == *host && c.provenance == *provenance,
876 (None, None) => true,
877 _ => false,
878 };
879 if unchanged {
880 return current.clone();
881 }
882 let epoch = self.fence.advance();
883 *current = to.map(|(host, provenance)| Current {
884 host,
885 provenance,
886 epoch,
887 });
888 current.clone()
889 }
890
891 fn host_left(&self, host: &HostGeneration) {
894 let mut current = lock(&self.current);
895 if current.as_ref().is_some_and(|c| c.host == *host) {
896 *current = None;
897 self.fence.advance();
898 }
899 }
900
901 fn release(&self, id: u64) -> usize {
902 let mut sessions = lock(&self.sessions);
903 sessions.remove(&id);
904 sessions.len()
905 }
906
907 fn owned_exited(&self, host: &HostGeneration) -> bool {
912 let mut slow = match self.slow.try_lock() {
913 Ok(guard) => guard,
914 Err(std::sync::TryLockError::Poisoned(e)) => e.into_inner(),
915 Err(std::sync::TryLockError::WouldBlock) => return false,
916 };
917 match slow.owned.as_mut() {
918 Some(owned) if owned.host.as_ref() == Some(host) => {
919 !matches!(owned.child.try_wait(), Ok(None))
920 }
921 _ => false,
922 }
923 }
924
925 fn ensure_host(&self) -> Result<Current, LifecycleError> {
927 let mut slow = lock(&self.slow);
928 if self.stopped.load(Ordering::SeqCst) {
929 return Err(LifecycleError::Released);
930 }
931 let deadline = Instant::now() + self.cfg.startup_timeout;
932 self.reap_exited_owned(&mut slow);
933 match self.backend.discover().map_err(unavailable)? {
934 Some(host) => self.adopt_found(&mut slow, host),
935 None => {
936 self.bind(None);
937 self.launch(&mut slow, deadline)
938 }
939 }
940 }
941
942 fn reap_exited_owned(&self, slow: &mut Slow) {
944 let exited = match slow.owned.as_mut() {
945 Some(owned) => match owned.child.try_wait() {
946 Ok(None) => false,
947 Ok(Some(_)) => true,
948 Err(_) => {
949 let _ = owned.child.terminate_tree();
951 true
952 }
953 },
954 None => false,
955 };
956 if exited {
957 slow.owned = None;
958 if let Some(generation) = slow.generation.take() {
959 let at = slow.origin.elapsed();
960 slow.pending = Some(slow.tracker.exited(generation, at));
961 }
962 }
963 }
964
965 fn adopt_found(
967 &self,
968 slow: &mut Slow,
969 host: HostGeneration,
970 ) -> Result<Current, LifecycleError> {
971 if let Some(current) = lock(&self.current).clone() {
972 if current.host == host {
973 return Ok(current);
974 }
975 }
976 let link = self.backend.connect(&host)?;
978 link.close();
979 let ours = slow
982 .owned
983 .as_ref()
984 .is_some_and(|owned| owned.child.id() == host.pid);
985 if ours {
986 if let Some(owned) = slow.owned.as_mut() {
987 owned.host = Some(host.clone());
988 }
989 } else if slow.owned.is_some() {
990 self.retire_owned(slow);
993 }
994 let provenance = if ours {
995 Provenance::Owned
996 } else {
997 Provenance::Attached
998 };
999 self.bind(Some((host, provenance)))
1000 .ok_or(LifecycleError::NoHost)
1001 }
1002
1003 fn retire_owned(&self, slow: &mut Slow) {
1005 if let Some(mut owned) = slow.owned.take() {
1006 stop_launched(&mut owned, self.cfg.stop_grace);
1007 }
1008 slow.generation = None;
1009 }
1010
1011 fn fail_owned(&self, slow: &mut Slow) {
1013 if let Some(mut owned) = slow.owned.take() {
1014 stop_launched(&mut owned, self.cfg.stop_grace);
1015 }
1016 let at = slow.origin.elapsed();
1017 if let Some(generation) = slow.generation.take() {
1018 slow.pending = Some(slow.tracker.exited(generation, at));
1019 } else {
1020 slow.pending = Some(slow.tracker.request_restart(at));
1021 }
1022 }
1023
1024 fn charge_failed_launch(&self, slow: &mut Slow) {
1025 let at = slow.origin.elapsed();
1026 slow.pending = Some(slow.tracker.request_restart(at));
1027 }
1028
1029 fn launch(&self, slow: &mut Slow, deadline: Instant) -> Result<Current, LifecycleError> {
1030 if slow.exhausted {
1031 return Err(LifecycleError::RestartExhausted);
1032 }
1033 if lock(&self.launcher).is_none() {
1034 return Err(LifecycleError::NoHost);
1035 }
1036 if let Some(outcome) = slow.pending.take() {
1037 self.honour(slow, outcome, deadline)?;
1038 }
1039 let mut claim: Option<Box<dyn Any + Send>> = None;
1040 if slow.owned.is_none() {
1041 loop {
1042 match self.backend.claim_spawn().map_err(LifecycleError::Launch)? {
1043 SpawnPermit::Granted(held) => {
1044 claim = Some(held);
1045 break;
1046 }
1047 SpawnPermit::Contended => {
1048 if let Some(host) = self.backend.discover().map_err(unavailable)? {
1049 return self.adopt_found(slow, host);
1050 }
1051 if Instant::now() >= deadline {
1052 return Err(LifecycleError::LaunchTimeout);
1053 }
1054 std::thread::sleep(self.cfg.poll_interval);
1055 }
1056 }
1057 }
1058 if let Some(host) = self.backend.discover().map_err(unavailable)? {
1060 return self.adopt_found(slow, host);
1061 }
1062 self.spawn_child(slow)?;
1063 }
1064 let result = self.await_publication(slow, deadline);
1065 drop(claim);
1066 result
1067 }
1068
1069 fn honour(
1071 &self,
1072 slow: &mut Slow,
1073 outcome: RestartOutcome,
1074 deadline: Instant,
1075 ) -> Result<(), LifecycleError> {
1076 let at = slow.origin.elapsed();
1077 let outcome = match outcome {
1078 RestartOutcome::AwaitingDemand => slow.tracker.request_restart(at),
1080 other => other,
1081 };
1082 match outcome {
1083 RestartOutcome::Restart(permit) => {
1084 let delay = permit.delay();
1085 if delay > deadline.saturating_duration_since(Instant::now()) {
1086 slow.pending = Some(RestartOutcome::Restart(permit));
1087 return Err(LifecycleError::Service(ServiceError::retryable(
1088 "restart_backoff",
1089 format!(
1090 "the next start of the owned host waits {} ms (restart backoff)",
1091 delay.as_millis()
1092 ),
1093 )));
1094 }
1095 std::thread::sleep(delay);
1096 if slow.tracker.permit_valid(&permit, slow.origin.elapsed()) {
1097 Ok(())
1098 } else {
1099 Err(LifecycleError::NoHost)
1100 }
1101 }
1102 RestartOutcome::Exhausted => {
1103 slow.exhausted = true;
1104 Err(LifecycleError::RestartExhausted)
1105 }
1106 RestartOutcome::AwaitingDemand
1107 | RestartOutcome::Stopped
1108 | RestartOutcome::IgnoredStale => Err(LifecycleError::NoHost),
1109 }
1110 }
1111
1112 fn spawn_child(&self, slow: &mut Slow) -> Result<(), LifecycleError> {
1113 let launched = {
1114 let mut launcher = lock(&self.launcher);
1115 match launcher.as_mut() {
1116 Some(launch) => launch(),
1117 None => return Err(LifecycleError::NoHost),
1118 }
1119 };
1120 let mut child = match launched {
1121 Ok(child) => child,
1122 Err(e) => {
1123 self.charge_failed_launch(slow);
1124 return Err(LifecycleError::Launch(e));
1125 }
1126 };
1127 let Some(pid) = NonZeroU32::new(child.id()) else {
1128 let _ = child.terminate_tree();
1129 self.charge_failed_launch(slow);
1130 return Err(LifecycleError::Launch(io::Error::other(
1131 "the launcher returned a child without a process id",
1132 )));
1133 };
1134 let identity = match AdoptedIdentity::capture(pid) {
1135 Ok(Some(identity)) => identity,
1136 Ok(None) => {
1137 let _ = child.terminate_tree();
1138 self.charge_failed_launch(slow);
1139 return Err(LifecycleError::Launch(io::Error::other(
1140 "the launched process exited before its identity could be recorded",
1141 )));
1142 }
1143 Err(e) => {
1144 let _ = child.terminate_tree();
1145 self.charge_failed_launch(slow);
1146 return Err(LifecycleError::Process(e));
1147 }
1148 };
1149 slow.generation = Some(slow.tracker.started(slow.origin.elapsed()));
1150 slow.owned = Some(OwnedHost {
1151 child,
1152 identity,
1153 host: None,
1154 });
1155 Ok(())
1156 }
1157
1158 fn await_publication(
1160 &self,
1161 slow: &mut Slow,
1162 deadline: Instant,
1163 ) -> Result<Current, LifecycleError> {
1164 loop {
1165 let pid = match slow.owned.as_mut() {
1166 Some(owned) => {
1167 if !matches!(owned.child.try_wait(), Ok(None)) {
1168 self.fail_owned(slow);
1169 return Err(LifecycleError::Launch(io::Error::other(
1170 "the launched host exited before it published itself",
1171 )));
1172 }
1173 owned.child.id()
1174 }
1175 None => return Err(LifecycleError::NoHost),
1176 };
1177 match self.backend.discover().map_err(unavailable)? {
1178 Some(host) if host.pid == pid => {
1179 if let Err(e) = self.backend.connect(&host).map(|link| link.close()) {
1180 self.fail_owned(slow);
1182 return Err(e);
1183 }
1184 if let Some(owned) = slow.owned.as_mut() {
1185 owned.host = Some(host.clone());
1186 }
1187 if let Some(generation) = slow.generation {
1188 slow.tracker.ready(generation);
1189 }
1190 return self
1191 .bind(Some((host, Provenance::Owned)))
1192 .ok_or(LifecycleError::NoHost);
1193 }
1194 Some(host) => {
1195 self.retire_owned(slow);
1198 return self.adopt_found(slow, host);
1199 }
1200 None => {}
1201 }
1202 if Instant::now() >= deadline {
1203 self.fail_owned(slow);
1204 return Err(LifecycleError::LaunchTimeout);
1205 }
1206 std::thread::sleep(self.cfg.poll_interval);
1207 }
1208 }
1209}
1210
1211fn stop_launched(owned: &mut OwnedHost, grace: Duration) {
1214 if let Ok(IdentityCheck::Replaced) = owned.identity.verify() {
1215 return;
1216 }
1217 let _ = owned.child.wait_then_kill(grace);
1218}
1219
1220pub struct Session {
1223 inner: Arc<Inner>,
1224 id: u64,
1225 label: String,
1226 provenance: Provenance,
1227 host: HostGeneration,
1228 epoch: u64,
1229 link: Box<dyn HostLink>,
1230 released: AtomicBool,
1231}
1232
1233impl Session {
1234 pub fn provenance(&self) -> Provenance {
1235 self.provenance
1236 }
1237
1238 pub fn host(&self) -> &HostGeneration {
1239 &self.host
1240 }
1241
1242 pub fn epoch(&self) -> u64 {
1244 self.epoch
1245 }
1246
1247 pub fn label(&self) -> &str {
1248 &self.label
1249 }
1250
1251 pub fn is_valid(&self) -> bool {
1254 !self.released.load(Ordering::SeqCst)
1255 && !self.inner.stopped.load(Ordering::SeqCst)
1256 && !self.link.is_closed()
1257 && self.inner.fence.epoch() == self.epoch
1258 }
1259
1260 pub fn call(&self, method: &str, params: Value) -> Result<Value, LifecycleError> {
1265 if self.released.load(Ordering::SeqCst) || self.inner.stopped.load(Ordering::SeqCst) {
1266 return Err(LifecycleError::Released);
1267 }
1268 self.confirm_current()?;
1269 let ticket = self.inner.fence.issue(self.epoch)?;
1270 let reply = self.link.call(method, params);
1271 self.confirm_current()?;
1273 self.inner.fence.admit_reply(&ticket)?;
1274 reply.map_err(LifecycleError::Service)
1275 }
1276
1277 fn confirm_current(&self) -> Result<(), LifecycleError> {
1278 let left = if self.provenance == Provenance::Owned && self.inner.owned_exited(&self.host) {
1279 true
1280 } else {
1281 match self.inner.backend.is_current(&self.host) {
1282 Ok(true) => false,
1283 Ok(false) => true,
1284 Err(_) => return Err(LifecycleError::Unconfirmed),
1285 }
1286 };
1287 if left {
1288 self.inner.host_left(&self.host);
1289 return Err(LifecycleError::StaleGeneration {
1290 issued: self.epoch,
1291 current: self.inner.fence.epoch(),
1292 });
1293 }
1294 Ok(())
1295 }
1296
1297 fn finish(&self) -> ReleaseReport {
1298 let first = !self.released.swap(true, Ordering::SeqCst);
1299 let remaining = if first {
1300 self.inner.release(self.id)
1301 } else {
1302 lock(&self.inner.sessions).len()
1303 };
1304 self.link.close();
1305 ReleaseReport {
1306 label: self.label.clone(),
1307 provenance: self.provenance,
1308 remaining_sessions: remaining,
1309 }
1310 }
1311}
1312
1313impl Drop for Session {
1314 fn drop(&mut self) {
1315 let _ = self.finish();
1316 }
1317}
1318
1319#[cfg(test)]
1320mod tests {
1321 use super::*;
1322 use rightkit_child_process::RestartMode;
1323 use serde_json::json;
1324 use std::sync::atomic::AtomicUsize;
1325
1326 fn generation(pid: u32, started_ms: u64) -> HostGeneration {
1327 HostGeneration {
1328 service: "svc".into(),
1329 endpoint: "endpoint".into(),
1330 pid,
1331 process_started_ms: 1_000,
1332 started_ms,
1333 }
1334 }
1335
1336 #[derive(Default)]
1337 struct FakeState {
1338 host: Option<HostGeneration>,
1339 refuse_identity: bool,
1340 contended: bool,
1341 calls: usize,
1342 stop_requests: usize,
1343 replace_during_call: Option<Option<HostGeneration>>,
1345 }
1346
1347 struct Fake {
1348 state: Arc<Mutex<FakeState>>,
1349 }
1350
1351 struct FakeLink {
1352 state: Arc<Mutex<FakeState>>,
1353 host: HostGeneration,
1354 closed: AtomicBool,
1355 }
1356
1357 impl Fake {
1358 fn new(host: Option<HostGeneration>) -> Arc<Self> {
1359 Arc::new(Self {
1360 state: Arc::new(Mutex::new(FakeState {
1361 host,
1362 ..FakeState::default()
1363 })),
1364 })
1365 }
1366
1367 fn publish(&self, host: Option<HostGeneration>) {
1368 lock(&self.state).host = host;
1369 }
1370 }
1371
1372 impl HostLink for FakeLink {
1373 fn call(&self, method: &str, _params: Value) -> Result<Value, ServiceError> {
1374 let mut state = lock(&self.state);
1375 state.calls += 1;
1376 if method == "stop" {
1377 state.stop_requests += 1;
1378 }
1379 if let Some(next) = state.replace_during_call.take() {
1380 state.host = next;
1381 }
1382 Ok(json!({"pid": self.host.pid}))
1383 }
1384
1385 fn is_closed(&self) -> bool {
1386 self.closed.load(Ordering::SeqCst)
1387 }
1388
1389 fn close(&self) {
1390 self.closed.store(true, Ordering::SeqCst);
1391 }
1392 }
1393
1394 impl HostBackend for Fake {
1395 fn discover(&self) -> io::Result<Option<HostGeneration>> {
1396 Ok(lock(&self.state).host.clone())
1397 }
1398
1399 fn connect(&self, host: &HostGeneration) -> Result<Box<dyn HostLink>, LifecycleError> {
1400 let state = lock(&self.state);
1401 if state.refuse_identity {
1402 return Err(LifecycleError::HostIdentity(HostIdentityError {
1403 expected: Vec::new(),
1404 found_pid: Some(host.pid),
1405 found_executable: None,
1406 reason: HostIdentityReason::ExecutableMismatch,
1407 }));
1408 }
1409 if state.host.as_ref() != Some(host) {
1410 return Err(LifecycleError::Service(ServiceError::retryable(
1411 ERROR_UNAVAILABLE,
1412 "gone",
1413 )));
1414 }
1415 Ok(Box::new(FakeLink {
1416 state: self.state.clone(),
1417 host: host.clone(),
1418 closed: AtomicBool::new(false),
1419 }))
1420 }
1421
1422 fn claim_spawn(&self) -> io::Result<SpawnPermit> {
1423 if lock(&self.state).contended {
1424 Ok(SpawnPermit::Contended)
1425 } else {
1426 Ok(SpawnPermit::Granted(Box::new(())))
1427 }
1428 }
1429 }
1430
1431 fn attach(fake: &Arc<Fake>) -> DaemonLifecycle {
1432 DaemonLifecycle::attach_or_launch(fake.clone(), None, LifecycleConfig::new()).unwrap()
1433 }
1434
1435 #[test]
1436 fn an_attached_host_is_never_owned_and_never_stopped() {
1437 let fake = Fake::new(Some(generation(10, 1)));
1438 let life = attach(&fake);
1439 assert_eq!(life.provenance(), Some(Provenance::Attached));
1440 assert!(!Provenance::Attached.may_stop_host());
1441 let session = life.acquire_session("shell").unwrap();
1442 assert_eq!(session.provenance(), Provenance::Attached);
1443 let error = life.shutdown_owned().unwrap_err();
1445 assert!(matches!(error, LifecycleError::NotOwned), "{error}");
1446 assert_eq!(error.code(), ERROR_NOT_OWNED);
1447 assert_eq!(lock(&fake.state).stop_requests, 0);
1448 assert_eq!(lock(&fake.state).host, Some(generation(10, 1)));
1449 assert_eq!(session.call("echo", json!({})).unwrap(), json!({"pid": 10}));
1450 }
1451
1452 #[test]
1453 fn a_host_that_fails_the_identity_check_is_not_attached() {
1454 let fake = Fake::new(Some(generation(10, 1)));
1455 lock(&fake.state).refuse_identity = true;
1456 let error = DaemonLifecycle::attach_or_launch(fake.clone(), None, LifecycleConfig::new())
1457 .err()
1458 .expect("an impostor is refused");
1459 assert!(matches!(error, LifecycleError::HostIdentity(_)), "{error}");
1460 assert!(!error.is_retryable());
1461 assert_eq!(lock(&fake.state).calls, 0, "nothing was sent to it");
1462 }
1463
1464 #[test]
1465 fn no_host_and_no_launcher_is_a_retryable_no_host() {
1466 let fake = Fake::new(None);
1467 let error = DaemonLifecycle::attach_or_launch(fake, None, LifecycleConfig::new())
1468 .err()
1469 .expect("nothing to attach to");
1470 assert!(matches!(error, LifecycleError::NoHost));
1471 assert!(error.is_retryable());
1472 assert_eq!(error.into_service_error().code, ERROR_UNAVAILABLE);
1473 }
1474
1475 #[test]
1476 fn two_shells_attach_and_one_detaching_leaves_the_other_running() {
1477 let fake = Fake::new(Some(generation(10, 1)));
1478 let (a, b) = (attach(&fake), attach(&fake));
1479 let session_a = a.acquire_session("shell-a").unwrap();
1480 let session_b = b.acquire_session("shell-b").unwrap();
1481 assert_eq!(a.session_count(), 1);
1482 assert_eq!(b.session_count(), 1);
1483 assert_eq!(a.host(), b.host());
1484 let report = a.release_session(session_a);
1485 assert_eq!(report.provenance, Provenance::Attached);
1486 assert_eq!(report.remaining_sessions, 0);
1487 assert_eq!(lock(&fake.state).host, Some(generation(10, 1)));
1489 assert_eq!(
1490 session_b.call("echo", json!({})).unwrap(),
1491 json!({"pid": 10})
1492 );
1493 assert_eq!(b.session_count(), 1);
1494 }
1495
1496 #[test]
1497 fn a_dropped_session_leaves_the_count() {
1498 let fake = Fake::new(Some(generation(10, 1)));
1499 let life = attach(&fake);
1500 {
1501 let _held = life.acquire_session("shell").unwrap();
1502 assert_eq!(life.session_count(), 1);
1503 }
1504 assert_eq!(life.session_count(), 0);
1505 }
1506
1507 #[test]
1508 fn when_the_owner_exits_the_session_is_stale_and_there_is_no_host() {
1509 let fake = Fake::new(Some(generation(10, 1)));
1510 let life = attach(&fake);
1511 let session = life.acquire_session("shell").unwrap();
1512 let bound = session.epoch();
1513 fake.publish(None);
1514 match session.call("echo", json!({})) {
1515 Err(LifecycleError::StaleGeneration { issued, current }) => {
1516 assert_eq!(issued, bound);
1517 assert!(current > issued);
1518 }
1519 other => panic!("expected a stale generation, got {other:?}"),
1520 }
1521 assert!(!session.is_valid());
1522 let error = life
1523 .acquire_session("shell")
1524 .err()
1525 .expect("no host any more");
1526 assert!(matches!(error, LifecycleError::NoHost));
1527 assert_eq!(life.provenance(), None);
1528 }
1529
1530 #[test]
1531 fn a_replacement_host_is_a_new_generation_and_the_old_session_stays_stale() {
1532 let fake = Fake::new(Some(generation(10, 1)));
1533 let life = attach(&fake);
1534 let old = life.acquire_session("shell").unwrap();
1535 fake.publish(Some(generation(11, 2)));
1536 let new = life.acquire_session("shell").unwrap();
1537 assert_eq!(new.host().pid, 11);
1538 assert!(new.epoch() > old.epoch());
1539 assert!(matches!(
1540 old.call("echo", json!({})),
1541 Err(LifecycleError::StaleGeneration { .. })
1542 ));
1543 assert_eq!(new.call("echo", json!({})).unwrap(), json!({"pid": 11}));
1544 assert_eq!(life.provenance(), Some(Provenance::Attached));
1545 }
1546
1547 #[test]
1548 fn a_reply_that_arrives_after_the_host_was_replaced_is_not_delivered() {
1549 let fake = Fake::new(Some(generation(10, 1)));
1550 let life = attach(&fake);
1551 let session = life.acquire_session("shell").unwrap();
1552 lock(&fake.state).replace_during_call = Some(Some(generation(11, 2)));
1554 let outcome = session.call("echo", json!({}));
1555 assert!(
1556 matches!(outcome, Err(LifecycleError::StaleGeneration { .. })),
1557 "{outcome:?}"
1558 );
1559 assert_eq!(lock(&fake.state).calls, 1);
1561 }
1562
1563 #[test]
1564 fn the_fence_refuses_a_stale_ticket_and_a_stale_reply_epoch() {
1565 let fence = GenerationFence::new();
1566 let first = fence.advance();
1567 let old = fence.issue(first).unwrap();
1568 assert!(fence.admit_reply(&old).is_ok());
1569 let second = fence.advance();
1570 assert!(second > first);
1571 assert_eq!(
1573 fence.admit_reply(&old),
1574 Err(Stale {
1575 issued: first,
1576 current: second
1577 })
1578 );
1579 assert!(fence.issue(first).is_err());
1580 let new = fence.issue(second).unwrap();
1581 assert_ne!(old.request(), new.request());
1582 assert!(fence.admit_reply_from(&new, first).is_err());
1584 assert!(fence.admit_reply_from(&new, second).is_ok());
1585 }
1586
1587 #[test]
1588 fn a_launcher_that_keeps_failing_spends_the_restart_budget() {
1589 let fake = Fake::new(Some(generation(10, 1)));
1590 let attempts = Arc::new(AtomicUsize::new(0));
1591 let counter = attempts.clone();
1592 let launcher: Launcher = Box::new(move || {
1593 counter.fetch_add(1, Ordering::SeqCst);
1594 Err(io::Error::other("no such engine"))
1595 });
1596 let policy = RestartPolicy::new(
1597 RestartMode::Automatic,
1598 2,
1599 Duration::from_secs(60),
1600 Duration::from_millis(1),
1601 Duration::from_millis(2),
1602 )
1603 .unwrap();
1604 let config = LifecycleConfig::new()
1605 .restart_policy(policy)
1606 .startup_timeout(Duration::from_secs(5))
1607 .poll_interval(Duration::from_millis(1));
1608 let life = DaemonLifecycle::attach_or_launch(fake.clone(), Some(launcher), config).unwrap();
1609 fake.publish(None);
1611 for _ in 0..3 {
1612 let error = life.acquire_session("shell").err().expect("launch fails");
1613 assert!(matches!(error, LifecycleError::Launch(_)), "{error}");
1614 }
1615 assert_eq!(attempts.load(Ordering::SeqCst), 3);
1616 let error = life.acquire_session("shell").err().expect("budget spent");
1617 assert!(matches!(error, LifecycleError::RestartExhausted), "{error}");
1618 assert_eq!(error.code(), ERROR_RESTART_EXHAUSTED);
1619 let again = life.acquire_session("shell").err().expect("still spent");
1621 assert!(matches!(again, LifecycleError::RestartExhausted));
1622 assert_eq!(attempts.load(Ordering::SeqCst), 3);
1623 life.reset_restart_budget();
1625 let after = life
1626 .acquire_session("shell")
1627 .err()
1628 .expect("launch fails again");
1629 assert!(matches!(after, LifecycleError::Launch(_)), "{after}");
1630 assert_eq!(attempts.load(Ordering::SeqCst), 4);
1631 }
1632
1633 #[test]
1634 fn a_contended_claim_waits_for_the_winner_and_attaches() {
1635 let fake = Fake::new(None);
1636 lock(&fake.state).contended = true;
1637 let launcher: Launcher = Box::new(|| Err(io::Error::other("must not launch")));
1638 let publisher = {
1639 let fake = fake.clone();
1640 std::thread::spawn(move || {
1641 std::thread::sleep(Duration::from_millis(50));
1642 fake.publish(Some(generation(20, 5)));
1643 })
1644 };
1645 let config = LifecycleConfig::new()
1646 .startup_timeout(Duration::from_secs(10))
1647 .poll_interval(Duration::from_millis(5));
1648 let life = DaemonLifecycle::attach_or_launch(fake, Some(launcher), config).unwrap();
1649 publisher.join().unwrap();
1650 assert_eq!(life.provenance(), Some(Provenance::Attached));
1651 assert_eq!(life.host().map(|h| h.pid), Some(20));
1652 }
1653
1654 #[test]
1655 fn errors_collapse_to_wire_codes() {
1656 assert_eq!(
1657 LifecycleError::StaleGeneration {
1658 issued: 1,
1659 current: 2
1660 }
1661 .into_service_error()
1662 .code,
1663 ERROR_STALE_GENERATION
1664 );
1665 assert_eq!(
1666 LifecycleError::NotOwned.into_service_error().code,
1667 ERROR_NOT_OWNED
1668 );
1669 assert_eq!(
1670 LifecycleError::Unconfirmed.into_service_error().code,
1671 ERROR_STALE_GENERATION
1672 );
1673 assert!(!LifecycleError::Unconfirmed.is_retryable());
1674 }
1675}