1mod probe;
23
24use std::env;
25use std::net::SocketAddr;
26use std::sync::atomic::AtomicBool;
27use std::sync::atomic::AtomicU64;
28use std::sync::atomic::AtomicU8;
29use std::sync::atomic::Ordering;
30use std::sync::Arc;
31use std::time::Duration;
32use std::time::Instant;
33
34use parking_lot::Mutex;
35#[cfg(test)]
36use tokio::io::AsyncReadExt;
37#[cfg(test)]
38use tokio::io::AsyncWriteExt;
39#[cfg(test)]
40use tokio::net::TcpListener;
41#[cfg(test)]
42use tokio::net::TcpStream;
43use tokio::sync::watch;
44
45use crate::critical::CriticalFailure;
46use crate::critical::CriticalFailureState;
47use crate::task_group::TaskId;
48use crate::task_spawner::TaskSpawner;
49use crate::wait_for_signal_result;
50use crate::ChildServiceContext;
51use crate::RuntimeError;
52use crate::RuntimeResult;
53use crate::ShutdownDeadline;
54use crate::TaskGroup;
55
56pub const HEALTH_BIND_ADDR_ENV: &str = "ROCKETMQ_HEALTH_BIND_ADDR";
58pub const SHUTDOWN_TIMEOUT_SECONDS_ENV: &str = "ROCKETMQ_SHUTDOWN_TIMEOUT_SECONDS";
60pub const LIVENESS_STALE_SECONDS_ENV: &str = "ROCKETMQ_LIVENESS_STALE_SECONDS";
62pub const HEALTH_DRAIN_METHODS_ENV: &str = "ROCKETMQ_HEALTH_DRAIN_METHODS";
64pub const DEFAULT_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(45);
66pub const DEFAULT_LIVENESS_STALE_AFTER: Duration = Duration::from_secs(30);
68const PROGRESS_INTERVAL: Duration = Duration::from_secs(1);
69
70const STATE_STARTING: u8 = 0;
71const STATE_READY: u8 = 1;
72const STATE_DRAINING: u8 = 2;
73const STATE_STOPPED: u8 = 3;
74const STATE_FAILED: u8 = 4;
75
76#[derive(Debug, Clone, Copy, PartialEq, Eq)]
78pub enum DependencyReadiness {
79 Ready,
81 Degraded,
83}
84
85#[derive(Debug, Clone, Copy, PartialEq, Eq)]
90pub enum CriticalFailureRecovery {
91 RevokeReadiness,
96 FailService,
98 FailAndRequestShutdown,
103}
104
105#[derive(Debug, Clone, Copy, PartialEq, Eq)]
107pub enum ServiceLifecycleState {
108 Starting,
110 Ready,
112 Draining,
114 Stopped,
116 Failed,
118}
119
120impl ServiceLifecycleState {
121 const fn from_u8(value: u8) -> Self {
122 match value {
123 STATE_READY => Self::Ready,
124 STATE_DRAINING => Self::Draining,
125 STATE_STOPPED => Self::Stopped,
126 STATE_FAILED => Self::Failed,
127 _ => Self::Starting,
128 }
129 }
130
131 pub const fn as_str(self) -> &'static str {
133 match self {
134 Self::Starting => "starting",
135 Self::Ready => "ready",
136 Self::Draining => "draining",
137 Self::Stopped => "stopped",
138 Self::Failed => "failed",
139 }
140 }
141}
142
143#[derive(Debug, Clone, Copy, PartialEq, Eq)]
145pub enum ShutdownReason {
146 PreStop,
148 Signal,
150 Internal,
152}
153
154impl ShutdownReason {
155 pub const fn as_str(self) -> &'static str {
157 match self {
158 Self::PreStop => "pre_stop",
159 Self::Signal => "signal",
160 Self::Internal => "internal",
161 }
162 }
163}
164
165#[derive(Debug, Clone, Copy, PartialEq, Eq)]
167pub struct ShutdownRequest {
168 pub reason: ShutdownReason,
170 pub deadline: ShutdownDeadline,
172}
173
174#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
180pub enum DrainRequestMethods {
181 #[default]
183 PostOnly,
184 GetOrPost,
186}
187
188impl DrainRequestMethods {
189 pub fn parse(raw: &str) -> Option<Self> {
195 let mut get = false;
196 let mut post = false;
197 for method in raw.split(',').map(str::trim) {
198 if method.eq_ignore_ascii_case("GET") && !get {
199 get = true;
200 } else if method.eq_ignore_ascii_case("POST") && !post {
201 post = true;
202 } else {
203 return None;
204 }
205 }
206 match (get, post) {
207 (false, true) => Some(Self::PostOnly),
208 (true, true) => Some(Self::GetOrPost),
209 (_, false) => None,
210 }
211 }
212
213 pub fn allows(self, method: &str) -> bool {
215 match self {
216 Self::PostOnly => method == "POST",
217 Self::GetOrPost => method == "GET" || method == "POST",
218 }
219 }
220
221 pub const fn allow_header(self) -> &'static str {
223 match self {
224 Self::PostOnly => "POST",
225 Self::GetOrPost => "GET, POST",
226 }
227 }
228}
229
230#[derive(Debug, Clone, PartialEq, Eq)]
232pub struct ServiceLifecycleConfig {
233 pub service_name: Arc<str>,
235 pub probe_bind_addr: Option<SocketAddr>,
237 pub shutdown_timeout: Duration,
243 pub liveness_stale_after: Duration,
245 pub drain_request_methods: DrainRequestMethods,
247}
248
249impl ServiceLifecycleConfig {
250 pub fn from_env(service_name: impl Into<Arc<str>>) -> RuntimeResult<Self> {
261 let probe_bind_addr = optional_env(HEALTH_BIND_ADDR_ENV)?
262 .map(|raw| parse_socket_addr(HEALTH_BIND_ADDR_ENV, &raw))
263 .transpose()?;
264 let shutdown_timeout = parse_duration_env(SHUTDOWN_TIMEOUT_SECONDS_ENV, DEFAULT_SHUTDOWN_TIMEOUT, 1, 300)?;
265 let liveness_stale_after =
266 parse_duration_env(LIVENESS_STALE_SECONDS_ENV, DEFAULT_LIVENESS_STALE_AFTER, 2, 300)?;
267 let drain_request_methods = match optional_env(HEALTH_DRAIN_METHODS_ENV)? {
268 None => DrainRequestMethods::default(),
269 Some(raw) => DrainRequestMethods::parse(&raw)
270 .ok_or_else(|| RuntimeError::configuration(crate::RuntimeOperation::ServiceLifecycleDrainMethods))?,
271 };
272 Ok(Self {
273 service_name: service_name.into(),
274 probe_bind_addr,
275 shutdown_timeout,
276 liveness_stale_after,
277 drain_request_methods,
278 })
279 }
280}
281
282fn optional_env(name: &'static str) -> RuntimeResult<Option<String>> {
283 match env::var(name) {
284 Ok(value) => Ok(Some(value)),
285 Err(env::VarError::NotPresent) => Ok(None),
286 Err(error) => Err(RuntimeError::configuration_failure(
287 crate::RuntimeOperation::ServiceLifecycleEnvironment,
288 error,
289 )),
290 }
291}
292
293fn parse_socket_addr(_name: &'static str, raw: &str) -> RuntimeResult<SocketAddr> {
294 raw.parse::<SocketAddr>().map_err(|error| {
295 RuntimeError::configuration_failure(crate::RuntimeOperation::ServiceLifecycleProbeAddress, error)
296 })
297}
298
299fn parse_duration_env(
300 name: &'static str,
301 default: Duration,
302 minimum_seconds: u64,
303 maximum_seconds: u64,
304) -> RuntimeResult<Duration> {
305 let Some(raw) = optional_env(name)? else {
306 return Ok(default);
307 };
308 let seconds = raw.parse::<u64>().map_err(|error| {
309 RuntimeError::configuration_failure(crate::RuntimeOperation::ServiceLifecycleDuration, error)
310 })?;
311 if !(minimum_seconds..=maximum_seconds).contains(&seconds) {
312 return Err(RuntimeError::configuration(
313 crate::RuntimeOperation::ServiceLifecycleDurationRange,
314 ));
315 }
316 Ok(Duration::from_secs(seconds))
317}
318
319#[derive(Debug)]
320struct ServiceLifecycleInner {
321 config: ServiceLifecycleConfig,
322 state: AtomicU8,
323 maintenance_readiness_suspended: AtomicBool,
324 dependencies_ready: AtomicBool,
325 started_at: Instant,
326 last_progress_millis: AtomicU64,
327 shutdown_request: Mutex<Option<ShutdownRequest>>,
328 shutdown_tx: watch::Sender<Option<ShutdownRequest>>,
329 started: AtomicBool,
330 lifecycle_tasks: Mutex<Option<TaskGroup>>,
331 probe_local_addr: Mutex<Option<SocketAddr>>,
332 observer: Mutex<Option<Arc<dyn ServiceLifecycleObserver>>>,
333}
334
335#[derive(Debug, Clone, Copy, PartialEq, Eq)]
337pub struct ServiceLifecycleTransition {
338 pub from: Option<ServiceLifecycleState>,
340 pub to: ServiceLifecycleState,
342}
343
344pub trait ServiceLifecycleObserver: std::fmt::Debug + Send + Sync {
351 fn on_transition(&self, transition: ServiceLifecycleTransition);
353}
354
355struct LifecycleNotification {
356 observer: Arc<dyn ServiceLifecycleObserver>,
357 transition: ServiceLifecycleTransition,
358}
359
360impl LifecycleNotification {
361 fn deliver(self) {
362 self.observer.on_transition(self.transition);
363 }
364}
365
366struct ServiceLifecycleStartAttempt<'a> {
367 inner: &'a ServiceLifecycleInner,
368 cancellation: Option<tokio_util::sync::CancellationToken>,
369 committed: bool,
370}
371
372impl<'a> ServiceLifecycleStartAttempt<'a> {
373 fn new(inner: &'a ServiceLifecycleInner) -> Self {
374 Self {
375 inner,
376 cancellation: None,
377 committed: false,
378 }
379 }
380
381 fn own(&mut self, cancellation: tokio_util::sync::CancellationToken) {
382 self.cancellation = Some(cancellation);
383 }
384
385 fn commit(mut self) {
386 self.committed = true;
387 }
388}
389
390impl Drop for ServiceLifecycleStartAttempt<'_> {
391 fn drop(&mut self) {
392 if self.committed {
393 return;
394 }
395 if let Some(cancellation) = self.cancellation.take() {
396 cancellation.cancel();
397 }
398 self.inner.probe_local_addr.lock().take();
399 self.inner.started.store(false, Ordering::Release);
400 }
401}
402
403#[derive(Debug, Clone)]
405pub struct ServiceLifecycle {
406 inner: Arc<ServiceLifecycleInner>,
407}
408
409impl ServiceLifecycle {
410 pub fn new(config: ServiceLifecycleConfig) -> Self {
412 let (shutdown_tx, _shutdown_rx) = watch::channel(None);
413 Self {
414 inner: Arc::new(ServiceLifecycleInner {
415 config,
416 state: AtomicU8::new(STATE_STARTING),
417 maintenance_readiness_suspended: AtomicBool::new(false),
418 dependencies_ready: AtomicBool::new(true),
419 started_at: Instant::now(),
420 last_progress_millis: AtomicU64::new(0),
421 shutdown_request: Mutex::new(None),
422 shutdown_tx,
423 started: AtomicBool::new(false),
424 lifecycle_tasks: Mutex::new(None),
425 probe_local_addr: Mutex::new(None),
426 observer: Mutex::new(None),
427 }),
428 }
429 }
430
431 pub fn from_env(service_name: impl Into<Arc<str>>) -> RuntimeResult<Self> {
433 Ok(Self::new(ServiceLifecycleConfig::from_env(service_name)?))
434 }
435
436 pub fn config(&self) -> &ServiceLifecycleConfig {
438 &self.inner.config
439 }
440
441 pub fn set_observer(&self, observer: Arc<dyn ServiceLifecycleObserver>) -> RuntimeResult<()> {
448 let mut slot = self.inner.observer.lock();
449 if slot.is_some() || self.state() != ServiceLifecycleState::Starting {
450 return Err(RuntimeError::context_unavailable(
451 crate::RuntimeOperation::StartServiceLifecycle,
452 ));
453 }
454 *slot = Some(Arc::clone(&observer));
455 drop(slot);
456 observer.on_transition(ServiceLifecycleTransition {
457 from: None,
458 to: ServiceLifecycleState::Starting,
459 });
460 Ok(())
461 }
462
463 fn transition(&self, target: u8, mut allowed: impl FnMut(u8) -> bool) -> Result<Option<LifecycleNotification>, u8> {
464 let observer = self.inner.observer.lock();
465 self.inner
466 .state
467 .try_update(Ordering::AcqRel, Ordering::Acquire, |state| {
468 (state != target && allowed(state)).then_some(target)
469 })
470 .map(|from| {
471 observer.as_ref().map(|observer| LifecycleNotification {
472 observer: Arc::clone(observer),
473 transition: ServiceLifecycleTransition {
474 from: Some(ServiceLifecycleState::from_u8(from)),
475 to: ServiceLifecycleState::from_u8(target),
476 },
477 })
478 })
479 }
480
481 pub async fn start(&self, service_context: &ChildServiceContext) -> RuntimeResult<()> {
490 if self.inner.started.swap(true, Ordering::AcqRel) {
491 return Err(RuntimeError::internal_failure(
492 crate::RuntimeOperation::StartServiceLifecycle,
493 ));
494 }
495
496 let mut start_attempt = ServiceLifecycleStartAttempt::new(&self.inner);
497 let lifecycle_context = service_context.component("service-lifecycle");
498 let lifecycle_tasks = lifecycle_context.task_group().clone();
499 let lifecycle_state = lifecycle_tasks.lifecycle_state();
500 if lifecycle_state != crate::TaskGroupLifecycleState::Open {
501 return Err(lifecycle_state.admission_error(crate::RuntimeOperation::ServiceLifecycleTaskGroup));
502 }
503 start_attempt.own(lifecycle_tasks.cancellation_token());
504
505 let probe = match self.bind_probe_listener().await {
506 Ok(probe) => probe,
507 Err(error) => return self.rollback_failed_start(lifecycle_tasks, error).await,
508 };
509
510 self.record_progress();
511 let heartbeat = self.clone();
512 let heartbeat_cancellation = lifecycle_tasks.cancellation_token();
513 if let Err(error) = lifecycle_tasks.spawn_service("service-lifecycle.progress", async move {
514 let mut interval = tokio::time::interval(PROGRESS_INTERVAL);
515 loop {
516 tokio::select! {
517 _ = heartbeat_cancellation.cancelled() => break,
518 _ = interval.tick() => heartbeat.record_progress(),
519 }
520 }
521 }) {
522 return self.rollback_failed_start(lifecycle_tasks, error).await;
523 }
524
525 if let Some((listener, local_addr)) = probe {
526 let lifecycle = self.clone();
527 let connection_tasks = lifecycle_tasks.clone();
528 if let Err(error) = lifecycle_tasks.spawn_service("service-lifecycle.probe-server", async move {
529 lifecycle.serve_probe_requests(listener, connection_tasks).await;
530 }) {
531 return self.rollback_failed_start(lifecycle_tasks, error).await;
532 }
533 *self.inner.probe_local_addr.lock() = Some(local_addr);
534 tracing::info!(
535 service = %self.inner.config.service_name,
536 bind = %local_addr,
537 "service lifecycle probe server listening"
538 );
539 }
540
541 *self.inner.lifecycle_tasks.lock() = Some(lifecycle_tasks);
542 start_attempt.commit();
543 Ok(())
544 }
545
546 async fn rollback_failed_start(&self, lifecycle_tasks: TaskGroup, error: RuntimeError) -> RuntimeResult<()> {
547 let deadline = ShutdownDeadline::after(self.inner.config.shutdown_timeout);
548 let report = lifecycle_tasks.shutdown_until(deadline).await;
549 report.log_if_unhealthy();
550 Err(error)
551 }
552
553 pub fn state(&self) -> ServiceLifecycleState {
555 ServiceLifecycleState::from_u8(self.inner.state.load(Ordering::Acquire))
556 }
557
558 pub fn probe_local_addr(&self) -> Option<SocketAddr> {
560 *self.inner.probe_local_addr.lock()
561 }
562
563 pub fn is_ready(&self) -> bool {
565 self.state() == ServiceLifecycleState::Ready
566 && !self.inner.maintenance_readiness_suspended.load(Ordering::Acquire)
567 && self.inner.dependencies_ready.load(Ordering::Acquire)
568 }
569
570 pub fn set_dependency_readiness(&self, readiness: DependencyReadiness) {
573 self.inner
574 .dependencies_ready
575 .store(matches!(readiness, DependencyReadiness::Ready), Ordering::Release);
576 }
577
578 pub fn dependency_readiness(&self) -> DependencyReadiness {
580 if self.inner.dependencies_ready.load(Ordering::Acquire) {
581 DependencyReadiness::Ready
582 } else {
583 DependencyReadiness::Degraded
584 }
585 }
586
587 pub fn handle_critical_failure(&self, failure: &CriticalFailure, recovery: CriticalFailureRecovery) {
595 tracing::warn!(
596 kind = failure.kind().as_str(),
597 task_kind = ?failure.task_kind(),
598 sequence = failure.sequence(),
599 "critical task failure"
600 );
601 self.set_dependency_readiness(DependencyReadiness::Degraded);
602 match recovery {
603 CriticalFailureRecovery::RevokeReadiness => {}
604 CriticalFailureRecovery::FailService => self.mark_failed(),
605 CriticalFailureRecovery::FailAndRequestShutdown => {
606 self.mark_failed();
607 self.request_shutdown(ShutdownReason::Internal);
608 }
609 }
610 }
611
612 pub fn spawn_critical_failure_monitor(
623 &self,
624 owner: &TaskSpawner,
625 failures: &CriticalFailureState,
626 recovery: CriticalFailureRecovery,
627 ) -> RuntimeResult<TaskId> {
628 let lifecycle = self.clone();
629 failures.spawn_monitor(owner, "service-lifecycle.critical-failures", move |failure| {
630 lifecycle.handle_critical_failure(&failure, recovery);
631 })
632 }
633
634 pub fn is_live(&self) -> bool {
636 if matches!(
637 self.state(),
638 ServiceLifecycleState::Stopped | ServiceLifecycleState::Failed
639 ) {
640 return false;
641 }
642 self.progress_age() <= self.inner.config.liveness_stale_after
643 }
644
645 pub fn record_progress(&self) {
647 let elapsed = self.inner.started_at.elapsed().as_millis();
648 self.inner
649 .last_progress_millis
650 .store(u64::try_from(elapsed).unwrap_or(u64::MAX), Ordering::Release);
651 }
652
653 fn progress_age(&self) -> Duration {
654 let last = self.inner.last_progress_millis.load(Ordering::Acquire);
655 let now = u64::try_from(self.inner.started_at.elapsed().as_millis()).unwrap_or(u64::MAX);
656 Duration::from_millis(now.saturating_sub(last))
657 }
658
659 pub fn mark_ready(&self) -> RuntimeResult<()> {
665 match self.transition(STATE_READY, |state| state == STATE_STARTING) {
666 Ok(notification) => {
667 self.record_progress();
668 if let Some(notification) = notification {
669 notification.deliver();
670 }
671 Ok(())
672 }
673 Err(STATE_READY) => Ok(()),
674 Err(_state) => Err(RuntimeError::internal_failure(
675 crate::RuntimeOperation::MarkServiceReady,
676 )),
677 }
678 }
679
680 pub fn suspend_readiness_for_maintenance(&self) -> RuntimeResult<()> {
690 let shutdown_request = self.inner.shutdown_request.lock();
691 if shutdown_request.is_some() || self.state() != ServiceLifecycleState::Ready {
692 return Err(RuntimeError::internal_failure(
693 crate::RuntimeOperation::SuspendServiceReadiness,
694 ));
695 }
696 self.inner
697 .maintenance_readiness_suspended
698 .store(true, Ordering::Release);
699 drop(shutdown_request);
700 self.record_progress();
701 Ok(())
702 }
703
704 pub fn restore_readiness_after_maintenance(&self) -> RuntimeResult<()> {
710 let shutdown_request = self.inner.shutdown_request.lock();
711 if shutdown_request.is_some() || self.state() != ServiceLifecycleState::Ready {
712 return Err(RuntimeError::internal_failure(
713 crate::RuntimeOperation::RestoreServiceReadiness,
714 ));
715 }
716 self.inner
717 .maintenance_readiness_suspended
718 .store(false, Ordering::Release);
719 drop(shutdown_request);
720 self.record_progress();
721 Ok(())
722 }
723
724 pub fn mark_failed(&self) {
726 let notification = self.transition(STATE_FAILED, |_| true).ok().flatten();
727 self.record_progress();
728 if let Some(notification) = notification {
729 notification.deliver();
730 }
731 }
732
733 pub fn mark_stopped(&self) {
735 let notification = self
736 .transition(STATE_STOPPED, |state| state != STATE_FAILED)
737 .ok()
738 .flatten();
739 self.record_progress();
740 if let Some(notification) = notification {
741 notification.deliver();
742 }
743 }
744
745 pub fn request_shutdown(&self, reason: ShutdownReason) -> ShutdownRequest {
749 let mut request = self.inner.shutdown_request.lock();
750 if let Some(existing) = *request {
751 return existing;
752 }
753 let first = ShutdownRequest {
754 reason,
755 deadline: ShutdownDeadline::after(self.inner.config.shutdown_timeout),
756 };
757 *request = Some(first);
758 let notification = self
761 .transition(STATE_DRAINING, |state| state != STATE_FAILED && state != STATE_STOPPED)
762 .ok()
763 .flatten();
764 self.record_progress();
765 self.inner.shutdown_tx.send_replace(Some(first));
766 drop(request);
767 if let Some(notification) = notification {
768 notification.deliver();
769 }
770 first
771 }
772
773 pub fn shutdown_request(&self) -> Option<ShutdownRequest> {
775 *self.inner.shutdown_request.lock()
776 }
777
778 pub async fn wait_for_shutdown(&self) -> ShutdownRequest {
780 if let Some(request) = self.shutdown_request() {
781 return request;
782 }
783 let mut receiver = self.inner.shutdown_tx.subscribe();
784 loop {
785 if let Some(request) = *receiver.borrow_and_update() {
786 return request;
787 }
788 if receiver.changed().await.is_err() {
789 return self.request_shutdown(ShutdownReason::Internal);
790 }
791 }
792 }
793
794 pub async fn wait_for_shutdown_signal(&self) -> RuntimeResult<ShutdownRequest> {
800 tokio::select! {
801 request = self.wait_for_shutdown() => Ok(request),
802 signal = wait_for_signal_result() => {
803 signal?;
804 Ok(self.request_shutdown(ShutdownReason::Signal))
805 }
806 }
807 }
808}
809
810#[cfg(test)]
811mod tests {
812 use std::error::Error;
813
814 use super::*;
815 use crate::RuntimeContext;
816
817 #[derive(Debug)]
818 struct RecordingObserver {
819 lifecycle: std::sync::Weak<ServiceLifecycleInner>,
820 events: Mutex<Vec<ServiceLifecycleTransition>>,
821 }
822
823 impl ServiceLifecycleObserver for RecordingObserver {
824 fn on_transition(&self, transition: ServiceLifecycleTransition) {
825 let inner = self.lifecycle.upgrade().unwrap();
826 assert!(
827 inner.shutdown_request.try_lock().is_some(),
828 "shutdown lock held by callback"
829 );
830 assert!(inner.observer.try_lock().is_some(), "observer lock held by callback");
831 self.events.lock().push(transition);
832 }
833 }
834
835 #[test]
836 fn observers_receive_only_committed_transitions_outside_lifecycle_locks() {
837 let lifecycle = ServiceLifecycle::new(config(None));
838 let observer = Arc::new(RecordingObserver {
839 lifecycle: Arc::downgrade(&lifecycle.inner),
840 events: Mutex::new(Vec::new()),
841 });
842 lifecycle.set_observer(observer.clone()).unwrap();
843 assert!(lifecycle.set_observer(observer.clone()).is_err());
844 lifecycle.mark_ready().unwrap();
845 lifecycle.mark_ready().unwrap();
846 let first = lifecycle.request_shutdown(ShutdownReason::Internal);
847 let repeated = lifecycle.request_shutdown(ShutdownReason::Signal);
848 assert_eq!(first.deadline, repeated.deadline);
849 lifecycle.mark_stopped();
850 lifecycle.mark_stopped();
851 let states: Vec<_> = observer.events.lock().iter().map(|event| event.to).collect();
852 assert_eq!(
853 states,
854 [
855 ServiceLifecycleState::Starting,
856 ServiceLifecycleState::Ready,
857 ServiceLifecycleState::Draining,
858 ServiceLifecycleState::Stopped,
859 ]
860 );
861 assert_eq!(observer.events.lock()[2].from, Some(ServiceLifecycleState::Ready));
862
863 let failed = ServiceLifecycle::new(config(None));
864 let observer = Arc::new(RecordingObserver {
865 lifecycle: Arc::downgrade(&failed.inner),
866 events: Mutex::new(Vec::new()),
867 });
868 failed.set_observer(observer.clone()).unwrap();
869 failed.mark_failed();
870 failed.mark_failed();
871 failed.request_shutdown(ShutdownReason::Internal);
872 failed.mark_stopped();
873 assert!(failed.mark_ready().is_err());
874 assert_eq!(observer.events.lock().len(), 2);
875 assert_eq!(observer.events.lock()[1].to, ServiceLifecycleState::Failed);
876 }
877
878 fn config(probe_bind_addr: Option<SocketAddr>) -> ServiceLifecycleConfig {
879 ServiceLifecycleConfig {
880 service_name: Arc::from("test-service"),
881 probe_bind_addr,
882 shutdown_timeout: Duration::from_secs(45),
883 liveness_stale_after: Duration::from_secs(30),
884 drain_request_methods: DrainRequestMethods::PostOnly,
885 }
886 }
887
888 async fn request(addr: SocketAddr, path: &str) -> String {
889 request_with_method(addr, "GET", path).await
890 }
891
892 async fn request_with_method(addr: SocketAddr, method: &str, path: &str) -> String {
893 let mut stream = TcpStream::connect(addr).await.expect("connect lifecycle probe");
894 stream
895 .write_all(format!("{method} {path} HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n").as_bytes())
896 .await
897 .expect("write lifecycle probe request");
898 let mut response = Vec::new();
899 stream
900 .read_to_end(&mut response)
901 .await
902 .expect("read lifecycle probe response");
903 String::from_utf8(response).expect("probe response is UTF-8")
904 }
905
906 #[test]
907 fn repeated_shutdown_requests_reuse_the_first_deadline() {
908 let lifecycle = ServiceLifecycle::new(config(None));
909 lifecycle.mark_ready().unwrap();
910 let first = lifecycle.request_shutdown(ShutdownReason::PreStop);
911 let second = lifecycle.request_shutdown(ShutdownReason::Signal);
912 assert_eq!(first, second);
913 assert_eq!(second.reason, ShutdownReason::PreStop);
914 assert_eq!(lifecycle.state(), ServiceLifecycleState::Draining);
915 assert!(
916 lifecycle.mark_ready().is_err(),
917 "drain must not transition back to ready"
918 );
919 }
920
921 #[test]
922 fn maintenance_readiness_is_reversible_until_shutdown_begins() {
923 let lifecycle = ServiceLifecycle::new(config(None));
924 lifecycle.mark_ready().unwrap();
925
926 lifecycle.suspend_readiness_for_maintenance().unwrap();
927 assert_eq!(lifecycle.state(), ServiceLifecycleState::Ready);
928 assert!(!lifecycle.is_ready());
929
930 lifecycle.restore_readiness_after_maintenance().unwrap();
931 assert!(lifecycle.is_ready());
932
933 lifecycle.suspend_readiness_for_maintenance().unwrap();
934 lifecycle.request_shutdown(ShutdownReason::Internal);
935 assert!(lifecycle.restore_readiness_after_maintenance().is_err());
936 assert!(!lifecycle.is_ready());
937 assert_eq!(lifecycle.state(), ServiceLifecycleState::Draining);
938 }
939
940 #[test]
941 fn dependency_recovery_cannot_restore_maintenance_or_shutdown_readiness() {
942 let lifecycle = ServiceLifecycle::new(config(None));
943 lifecycle.set_dependency_readiness(DependencyReadiness::Degraded);
944 lifecycle.mark_ready().unwrap();
945 assert!(!lifecycle.is_ready());
946 assert!(lifecycle.is_live());
947 lifecycle.set_dependency_readiness(DependencyReadiness::Ready);
948 assert!(lifecycle.is_ready());
949 lifecycle.suspend_readiness_for_maintenance().unwrap();
950 lifecycle.set_dependency_readiness(DependencyReadiness::Degraded);
951 lifecycle.set_dependency_readiness(DependencyReadiness::Ready);
952 assert!(!lifecycle.is_ready());
953 lifecycle.restore_readiness_after_maintenance().unwrap();
954 assert!(lifecycle.is_ready());
955 lifecycle.request_shutdown(ShutdownReason::Internal);
956 lifecycle.set_dependency_readiness(DependencyReadiness::Ready);
957 assert!(!lifecycle.is_ready());
958 }
959
960 #[test]
961 fn failed_is_not_overwritten_by_normal_stop_completion() {
962 let lifecycle = ServiceLifecycle::new(config(None));
963 lifecycle.mark_failed();
964 lifecycle.mark_stopped();
965
966 assert_eq!(lifecycle.state(), ServiceLifecycleState::Failed);
967 }
968
969 #[test]
970 fn concurrent_shutdown_cannot_overwrite_a_terminal_state() {
971 for fail in [false, true] {
972 for _ in 0..64 {
973 let lifecycle = ServiceLifecycle::new(config(None));
974 lifecycle.mark_ready().unwrap();
975 let barrier = std::sync::Barrier::new(2);
976 std::thread::scope(|scope| {
977 scope.spawn(|| {
978 barrier.wait();
979 lifecycle.request_shutdown(ShutdownReason::Internal);
980 });
981 barrier.wait();
982 if fail {
983 lifecycle.mark_failed();
984 } else {
985 lifecycle.mark_stopped();
986 }
987 });
988 assert_eq!(
989 lifecycle.state(),
990 if fail {
991 ServiceLifecycleState::Failed
992 } else {
993 ServiceLifecycleState::Stopped
994 }
995 );
996 let first = lifecycle.shutdown_request().unwrap();
997 assert_eq!(lifecycle.request_shutdown(ShutdownReason::Signal), first);
998 }
999 }
1000 }
1001
1002 #[tokio::test]
1003 async fn probe_server_reports_state_and_pre_stop_starts_drain() {
1004 let context = RuntimeContext::from_current("service-lifecycle-probe-test");
1005 let service = context.service_context("service-lifecycle-probe-test");
1006 let lifecycle = ServiceLifecycle::new(config(Some(SocketAddr::from(([127, 0, 0, 1], 0)))));
1007 lifecycle.start(&service).await.unwrap();
1008 let addr = lifecycle.probe_local_addr().expect("bound probe address");
1009
1010 let starting = request(addr, "/readyz").await;
1011 assert!(starting.starts_with("HTTP/1.1 503"), "{starting}");
1012 assert!(request(addr, "/livez").await.starts_with("HTTP/1.1 200"));
1013
1014 lifecycle.mark_ready().unwrap();
1015 assert!(request(addr, "/readyz").await.starts_with("HTTP/1.1 200"));
1016 assert!(request_with_method(addr, "POST", "/drainz")
1017 .await
1018 .starts_with("HTTP/1.1 200"));
1019 assert!(!lifecycle.is_ready());
1020 assert_eq!(lifecycle.wait_for_shutdown().await.reason, ShutdownReason::PreStop);
1021
1022 let report = context.shutdown_tasks(Duration::from_secs(1)).await;
1023 assert!(report.is_healthy(), "{}", report.to_json());
1024 }
1025
1026 #[tokio::test]
1027 async fn failed_probe_bind_rolls_back_owned_tasks_and_allows_retry() {
1028 let context = RuntimeContext::from_current("service-lifecycle-retry-test");
1029 let service = context.service_context("service-lifecycle-retry-test");
1030 let occupied_listener = TcpListener::bind(SocketAddr::from(([127, 0, 0, 1], 0)))
1031 .await
1032 .expect("reserve probe address");
1033 let probe_addr = occupied_listener.local_addr().expect("inspect reserved address");
1034 let lifecycle = ServiceLifecycle::new(config(Some(probe_addr)));
1035
1036 let error = lifecycle
1037 .start(&service)
1038 .await
1039 .expect_err("occupied probe address must fail startup");
1040 assert_eq!(error.operation(), crate::RuntimeOperation::BindServiceHealthProbe);
1041 assert!(error.source().is_some());
1042 assert_eq!(lifecycle.state(), ServiceLifecycleState::Starting);
1043 assert_eq!(lifecycle.probe_local_addr(), None);
1044 assert!(!lifecycle.inner.started.load(Ordering::Acquire));
1045 assert_eq!(service.task_group().task_count(), 0);
1046 assert_eq!(service.task_group().component_count(), 0);
1047
1048 drop(occupied_listener);
1049 lifecycle
1050 .start(&service)
1051 .await
1052 .expect("failed startup must be retryable");
1053 assert_eq!(lifecycle.probe_local_addr(), Some(probe_addr));
1054 assert!(lifecycle.inner.started.load(Ordering::Acquire));
1055 assert_eq!(
1056 lifecycle
1057 .inner
1058 .lifecycle_tasks
1059 .lock()
1060 .as_ref()
1061 .expect("successful startup owns a task group")
1062 .task_count(),
1063 2
1064 );
1065 assert_eq!(service.task_group().component_count(), 1);
1066 assert!(request(probe_addr, "/livez").await.starts_with("HTTP/1.1 200"));
1067
1068 let report = context.shutdown_tasks(Duration::from_secs(1)).await;
1069 assert!(report.is_healthy(), "{}", report.to_json());
1070 let lifecycle_tasks = lifecycle.inner.lifecycle_tasks.lock();
1071 let lifecycle_tasks = lifecycle_tasks
1072 .as_ref()
1073 .expect("successful startup retains task ownership");
1074 assert_eq!(lifecycle_tasks.task_count(), 0);
1075 assert_eq!(
1076 lifecycle_tasks.lifecycle_state(),
1077 crate::TaskGroupLifecycleState::ShutdownCompleted
1078 );
1079 }
1080
1081 #[test]
1082 fn cancelled_start_attempt_resets_admission_and_cancels_owned_tasks() {
1083 let lifecycle = ServiceLifecycle::new(config(None));
1084 lifecycle.inner.started.store(true, Ordering::Release);
1085 *lifecycle.inner.probe_local_addr.lock() = Some(SocketAddr::from(([127, 0, 0, 1], 8088)));
1086 let cancellation = tokio_util::sync::CancellationToken::new();
1087
1088 {
1089 let mut attempt = ServiceLifecycleStartAttempt::new(&lifecycle.inner);
1090 attempt.own(cancellation.clone());
1091 }
1092
1093 assert!(cancellation.is_cancelled());
1094 assert!(!lifecycle.inner.started.load(Ordering::Acquire));
1095 assert_eq!(lifecycle.probe_local_addr(), None);
1096 }
1097
1098 #[tokio::test]
1099 async fn liveness_rejects_stale_progress_and_terminal_states() {
1100 let mut stale_config = config(None);
1101 stale_config.liveness_stale_after = Duration::from_millis(5);
1102 let lifecycle = ServiceLifecycle::new(stale_config);
1103 lifecycle.record_progress();
1104 tokio::time::sleep(Duration::from_millis(10)).await;
1105 assert!(!lifecycle.is_live());
1106 lifecycle.record_progress();
1107 assert!(lifecycle.is_live());
1108 lifecycle.mark_failed();
1109 assert!(!lifecycle.is_live());
1110 }
1111}