Skip to main content

rocketmq_runtime/
service_lifecycle.rs

1// Copyright 2026 The RocketMQ Rust Authors
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Process-level readiness, liveness, and shutdown coordination.
16//!
17//! The lifecycle is intentionally transport-neutral. Kubernetes probes observe
18//! an explicit service state instead of treating an open business port as
19//! health, while pre-stop and operating-system signals converge on the first
20//! absolute shutdown deadline.
21
22mod 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
56/// The health bind addr env constant.
57pub const HEALTH_BIND_ADDR_ENV: &str = "ROCKETMQ_HEALTH_BIND_ADDR";
58/// The shutdown timeout seconds env constant.
59pub const SHUTDOWN_TIMEOUT_SECONDS_ENV: &str = "ROCKETMQ_SHUTDOWN_TIMEOUT_SECONDS";
60/// The liveness stale seconds env constant.
61pub const LIVENESS_STALE_SECONDS_ENV: &str = "ROCKETMQ_LIVENESS_STALE_SECONDS";
62/// HTTP methods accepted by `/drainz`: `POST` (the default) or `GET,POST`.
63pub const HEALTH_DRAIN_METHODS_ENV: &str = "ROCKETMQ_HEALTH_DRAIN_METHODS";
64/// Default process shutdown budget of [`ServiceLifecycleConfig::shutdown_timeout`].
65pub const DEFAULT_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(45);
66/// The default liveness stale after constant.
67pub 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/// Aggregate dependency readiness, independent of maintenance and process shutdown.
77#[derive(Debug, Clone, Copy, PartialEq, Eq)]
78pub enum DependencyReadiness {
79    /// Required dependencies currently satisfy the caller's health policy.
80    Ready,
81    /// Required dependencies are unavailable or their evidence has expired.
82    Degraded,
83}
84
85/// How a handled critical task failure affects the service.
86///
87/// The choice belongs to the business: the runtime only guarantees that
88/// readiness is revoked and that the failure reaches the handler.
89#[derive(Debug, Clone, Copy, PartialEq, Eq)]
90pub enum CriticalFailureRecovery {
91    /// Revoke dependency readiness and keep serving.
92    ///
93    /// Use this when the failed work is not required for the process to answer
94    /// traffic correctly and the caller handles the consequence elsewhere.
95    RevokeReadiness,
96    /// Revoke readiness and mark the service failed.
97    FailService,
98    /// Revoke readiness, mark the service failed, and request an ordered shutdown.
99    ///
100    /// The first request wins, so this cannot extend or replace a shutdown that
101    /// already started.
102    FailAndRequestShutdown,
103}
104
105/// Stable process lifecycle states used by readiness and liveness probes.
106#[derive(Debug, Clone, Copy, PartialEq, Eq)]
107pub enum ServiceLifecycleState {
108    /// Represents the starting case.
109    Starting,
110    /// Represents the ready case.
111    Ready,
112    /// Represents the draining case.
113    Draining,
114    /// Represents the stopped case.
115    Stopped,
116    /// Represents the failed case.
117    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    /// Borrows this value as str.
132    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/// Source of the first shutdown request.
144#[derive(Debug, Clone, Copy, PartialEq, Eq)]
145pub enum ShutdownReason {
146    /// Represents the pre stop case.
147    PreStop,
148    /// Represents the signal case.
149    Signal,
150    /// Represents the internal case.
151    Internal,
152}
153
154impl ShutdownReason {
155    /// Borrows this value as str.
156    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/// Immutable first shutdown request and its shared absolute deadline.
166#[derive(Debug, Clone, Copy, PartialEq, Eq)]
167pub struct ShutdownRequest {
168    /// The reason value.
169    pub reason: ShutdownReason,
170    /// The deadline value.
171    pub deadline: ShutdownDeadline,
172}
173
174/// HTTP methods that the `/drainz` probe route accepts.
175///
176/// A drain request starts the process shutdown, so the default accepts only
177/// `POST`. Kubernetes `preStop.httpGet` hooks can send only `GET`; a deployment
178/// that drains through such a hook opts in with [`Self::GetOrPost`].
179#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
180pub enum DrainRequestMethods {
181    /// Accepts only `POST`. `GET` receives `405 Method Not Allowed`.
182    #[default]
183    PostOnly,
184    /// Accepts `GET` and `POST`.
185    GetOrPost,
186}
187
188impl DrainRequestMethods {
189    /// Parses the value of `ROCKETMQ_HEALTH_DRAIN_METHODS`.
190    ///
191    /// Accepts a comma-separated method list, ignoring case, blanks around
192    /// methods, and order: `POST` or `GET,POST`. Returns `None` for any other
193    /// list, including `GET` alone, since draining must stay reachable by `POST`.
194    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    /// Returns whether `method` may start a drain.
214    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    /// Returns the value of the `Allow` header for a rejected drain request.
222    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/// Versioned runtime configuration for the process lifecycle boundary.
231#[derive(Debug, Clone, PartialEq, Eq)]
232pub struct ServiceLifecycleConfig {
233    /// The service name value.
234    pub service_name: Arc<str>,
235    /// The probe bind addr value.
236    pub probe_bind_addr: Option<SocketAddr>,
237    /// Process shutdown budget.
238    ///
239    /// The first shutdown request fixes a `ShutdownDeadline` this far ahead.
240    /// That deadline is the single source for component shutdown and the
241    /// owner's final shutdown; repeated requests never extend it.
242    pub shutdown_timeout: Duration,
243    /// The liveness stale after value.
244    pub liveness_stale_after: Duration,
245    /// HTTP methods that `/drainz` accepts.
246    pub drain_request_methods: DrainRequestMethods,
247}
248
249impl ServiceLifecycleConfig {
250    /// Reads the optional probe address and bounded lifecycle budgets from the environment.
251    ///
252    /// The probe server is disabled when `ROCKETMQ_HEALTH_BIND_ADDR` is absent, preserving
253    /// non-Kubernetes entrypoint behavior. Shutdown coordination remains active.
254    ///
255    /// # Errors
256    ///
257    /// Returns a configuration runtime failure for a malformed address, non-UTF-8 input,
258    /// zero timeout, a liveness window shorter than two progress intervals, or a
259    /// drain method list other than `POST` or `GET,POST`.
260    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/// One committed process state change. `from = None` is the initial Starting event.
336#[derive(Debug, Clone, Copy, PartialEq, Eq)]
337pub struct ServiceLifecycleTransition {
338    /// Previous committed state, absent only for initial observer attachment.
339    pub from: Option<ServiceLifecycleState>,
340    /// State committed by this transition.
341    pub to: ServiceLifecycleState,
342}
343
344/// Observes committed state changes without owning the lifecycle.
345///
346/// Callbacks run synchronously outside lifecycle locks. They must be short,
347/// nonblocking and panic-free, and must not flush exporters or perform I/O.
348/// Concurrent callers may deliver callbacks out of order; each event carries
349/// the actual atomic transition and is delivered once, without polling.
350pub trait ServiceLifecycleObserver: std::fmt::Debug + Send + Sync {
351    /// Records a committed transition without waiting for external work.
352    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/// Shared process lifecycle and health-probe owner.
404#[derive(Debug, Clone)]
405pub struct ServiceLifecycle {
406    inner: Arc<ServiceLifecycleInner>,
407}
408
409impl ServiceLifecycle {
410    /// Creates a new `ServiceLifecycle`.
411    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    /// Creates a value from env.
432    pub fn from_env(service_name: impl Into<Arc<str>>) -> RuntimeResult<Self> {
433        Ok(Self::new(ServiceLifecycleConfig::from_env(service_name)?))
434    }
435
436    /// Returns the config.
437    pub fn config(&self) -> &ServiceLifecycleConfig {
438        &self.inner.config
439    }
440
441    /// Binds one observer before the first state transition and emits Starting.
442    ///
443    /// # Errors
444    ///
445    /// Returns an error if an observer is already installed or the lifecycle
446    /// has left Starting. Installation is serialized with state transitions.
447    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    /// Starts the progress heartbeat and optional HTTP probe server under a dedicated
482    /// component task group of `service_context`.
483    ///
484    /// # Errors
485    ///
486    /// Returns an error when called twice, when the listener cannot bind, or when either
487    /// owned service task cannot be registered. A failed start cancels and awaits every
488    /// task registered by that attempt before making the lifecycle retryable.
489    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    /// Returns the state.
554    pub fn state(&self) -> ServiceLifecycleState {
555        ServiceLifecycleState::from_u8(self.inner.state.load(Ordering::Acquire))
556    }
557
558    /// Returns the probe local addr.
559    pub fn probe_local_addr(&self) -> Option<SocketAddr> {
560        *self.inner.probe_local_addr.lock()
561    }
562
563    /// Returns whether ready.
564    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    /// Publishes the aggregate dependency decision without changing lifecycle state.
571    /// A healthy dependency cannot undo maintenance, drain, failure, or stop.
572    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    /// Returns the latest aggregate dependency decision independently of lifecycle state.
579    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    /// Applies the handling policy for one critical task failure.
588    ///
589    /// Dependency readiness is revoked first, because a process whose critical
590    /// work failed must stop advertising itself as ready even when it keeps
591    /// making progress. The remaining effects are the caller's policy: this
592    /// method never decides on its own whether the process fails or asks for an
593    /// ordered shutdown.
594    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    /// Spawns a monitor that applies `recovery` when a critical failure is pending.
613    ///
614    /// `owner` must be an owner outside the group of the monitored tasks, because
615    /// a poisoned group can neither spawn nor run its own monitor. The monitor
616    /// takes each pending failure before handling it, so repeated failures are
617    /// handled in sequence rather than replayed.
618    ///
619    /// # Errors
620    ///
621    /// Returns an error when `owner` is shutting down or closed.
622    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    /// Returns whether live.
635    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    /// Records progress.
646    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    /// Marks startup dependencies ready without allowing a drain request to be reversed.
660    ///
661    /// # Errors
662    ///
663    /// Returns an error when shutdown or failure already began.
664    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    /// Temporarily removes this process from readiness during a reversible
681    /// maintenance operation.
682    ///
683    /// This does not alter the process lifecycle state. A real shutdown request
684    /// remains irreversible and prevents a later readiness restore.
685    ///
686    /// # Errors
687    ///
688    /// Returns an error unless the process is ready and shutdown has not begun.
689    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    /// Restores readiness after a reversible maintenance operation.
705    ///
706    /// # Errors
707    ///
708    /// Returns an error when shutdown, failure, or stop has begun.
709    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    /// Executes mark failed.
725    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    /// Executes mark stopped.
734    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    /// Records the first shutdown request and returns its immutable deadline.
746    ///
747    /// Repeated pre-stop or signal delivery is idempotent and cannot extend the budget.
748    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        // Failure and stop can race this request without taking its mutex.
759        // Recheck the terminal states on every compare-and-exchange attempt.
760        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    /// Shuts down request.
774    pub fn shutdown_request(&self) -> Option<ShutdownRequest> {
775        *self.inner.shutdown_request.lock()
776    }
777
778    /// Returns the wait for shutdown.
779    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    /// Waits for pre-stop or an operating-system termination signal.
795    ///
796    /// # Errors
797    ///
798    /// Returns the platform signal registration error without inventing a successful signal.
799    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}