Skip to main content

rocketmq_runtime/task/
service_task.rs

1// Copyright 2023 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
15use std::future::Future;
16use std::sync::atomic::AtomicBool;
17use std::sync::atomic::AtomicU64;
18use std::sync::atomic::Ordering;
19use std::sync::Arc;
20use std::time::Duration;
21use std::time::Instant;
22
23use parking_lot::Mutex;
24use tokio::sync::Notify;
25use tokio::time::timeout;
26use tracing::info;
27use tracing::warn;
28
29use crate::config::DEFAULT_SHUTDOWN_TIMEOUT;
30use crate::shutdown_deadline::ABORT_CONFIRMATION_TIMEOUT;
31use crate::RuntimeError;
32use crate::RuntimeResult;
33use crate::ShutdownDeadline;
34use crate::ShutdownReport;
35use crate::TaskGroup;
36use crate::TaskId;
37
38/// Lifecycle state of a [`ServiceManager`] and the loop it runs.
39#[derive(Debug, Clone, Copy, PartialEq, Eq)]
40pub enum ServiceTaskState {
41    /// The service has never been started.
42    NotStarted,
43    /// `start` accepted the request and is spawning the loop.
44    Starting,
45    /// The loop is running.
46    Running,
47    /// A stop was requested and the loop has not finished yet.
48    Stopping,
49    /// The loop finished, or shutdown completed.
50    Stopped,
51}
52
53impl ServiceTaskState {
54    const fn as_bits(self) -> u64 {
55        match self {
56            Self::NotStarted => 0,
57            Self::Starting => 1,
58            Self::Running => 2,
59            Self::Stopping => 3,
60            Self::Stopped => 4,
61        }
62    }
63
64    const fn from_bits(bits: u64) -> Self {
65        match bits {
66            0 => Self::NotStarted,
67            1 => Self::Starting,
68            2 => Self::Running,
69            3 => Self::Stopping,
70            _ => Self::Stopped,
71        }
72    }
73}
74
75const STATE_BITS: u32 = 8;
76const STATE_MASK: u64 = (1 << STATE_BITS) - 1;
77const GENERATION_MASK: u64 = u64::MAX >> STATE_BITS;
78
79fn pack(generation: u64, state: ServiceTaskState) -> u64 {
80    (generation << STATE_BITS) | state.as_bits()
81}
82
83fn generation_of(word: u64) -> u64 {
84    word >> STATE_BITS
85}
86
87fn state_of(word: u64) -> ServiceTaskState {
88    ServiceTaskState::from_bits(word & STATE_MASK)
89}
90
91/// Control state shared by a [`ServiceManager`] and the loop it runs.
92///
93/// One atomic word holds the state together with the generation of the
94/// current run. `start` begins a new generation, so a loop left over from an
95/// earlier run that outlived its shutdown sees itself as stopped and cannot
96/// overwrite the state of the run that replaced it.
97#[derive(Debug)]
98struct ServiceSignals {
99    word: AtomicU64,
100    has_notified: AtomicBool,
101    wait_point: Notify,
102}
103
104impl ServiceSignals {
105    fn new() -> Self {
106        Self {
107            word: AtomicU64::new(pack(0, ServiceTaskState::NotStarted)),
108            has_notified: AtomicBool::new(false),
109            wait_point: Notify::new(),
110        }
111    }
112
113    fn state(&self) -> ServiceTaskState {
114        state_of(self.word.load(Ordering::Acquire))
115    }
116
117    /// Starts a new generation from `NotStarted` or `Stopped`.
118    ///
119    /// Returns the replaced word and the new generation, or the current state
120    /// when a run is still active.
121    fn begin_start(&self) -> Result<(u64, u64), ServiceTaskState> {
122        let mut word = self.word.load(Ordering::Acquire);
123        loop {
124            let state = state_of(word);
125            if !matches!(state, ServiceTaskState::NotStarted | ServiceTaskState::Stopped) {
126                return Err(state);
127            }
128            let generation = generation_of(word).wrapping_add(1) & GENERATION_MASK;
129            match self.word.compare_exchange_weak(
130                word,
131                pack(generation, ServiceTaskState::Starting),
132                Ordering::AcqRel,
133                Ordering::Acquire,
134            ) {
135                Ok(_) => return Ok((word, generation)),
136                Err(actual) => word = actual,
137            }
138        }
139    }
140
141    /// Moves `generation` from `from` to `to`; does nothing if either changed.
142    fn transition(&self, generation: u64, from: ServiceTaskState, to: ServiceTaskState) -> bool {
143        self.word
144            .compare_exchange(
145                pack(generation, from),
146                pack(generation, to),
147                Ordering::AcqRel,
148                Ordering::Acquire,
149            )
150            .is_ok()
151    }
152
153    /// Marks `generation` stopped unless a newer run already replaced it.
154    fn finish(&self, generation: u64) -> bool {
155        self.word
156            .fetch_update(Ordering::AcqRel, Ordering::Acquire, |word| {
157                (generation_of(word) == generation).then(|| pack(generation, ServiceTaskState::Stopped))
158            })
159            .is_ok()
160    }
161
162    /// Requests a stop of the current run and returns its state and generation.
163    fn request_stop(&self) -> (ServiceTaskState, u64) {
164        let mut word = self.word.load(Ordering::Acquire);
165        loop {
166            let state = state_of(word);
167            let generation = generation_of(word);
168            if !matches!(state, ServiceTaskState::Starting | ServiceTaskState::Running) {
169                return (state, generation);
170            }
171            match self.word.compare_exchange_weak(
172                word,
173                pack(generation, ServiceTaskState::Stopping),
174                Ordering::AcqRel,
175                Ordering::Acquire,
176            ) {
177                Ok(_) => return (state, generation),
178                Err(actual) => word = actual,
179            }
180        }
181    }
182
183    fn is_stopped_for(&self, generation: u64) -> bool {
184        let word = self.word.load(Ordering::Acquire);
185        generation_of(word) != generation
186            || matches!(state_of(word), ServiceTaskState::Stopping | ServiceTaskState::Stopped)
187    }
188
189    fn wakeup(&self) {
190        if self
191            .has_notified
192            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
193            .is_ok()
194        {
195            self.wait_point.notify_one();
196        }
197    }
198}
199
200/// Stops the run it belongs to when the loop returns or is aborted.
201struct StopOnExit {
202    signals: Arc<ServiceSignals>,
203    generation: u64,
204}
205
206impl Drop for StopOnExit {
207    fn drop(&mut self) {
208        if self.signals.finish(self.generation) {
209            self.signals.has_notified.store(false, Ordering::Release);
210        }
211    }
212}
213
214/// Control handle passed to [`ServiceTask::run`].
215pub struct ServiceTaskContext {
216    signals: Arc<ServiceSignals>,
217    generation: u64,
218}
219
220impl ServiceTaskContext {
221    /// Creates a context that is not stopped and has no pending wakeup.
222    ///
223    /// [`ServiceManager`] creates the context it passes to [`ServiceTask::run`];
224    /// this constructor lets a test drive a service loop directly.
225    pub fn new() -> Self {
226        Self {
227            signals: Arc::new(ServiceSignals::new()),
228            generation: 0,
229        }
230    }
231
232    /// Returns whether this run was asked to stop, or was replaced by a newer run.
233    pub fn is_stopped(&self) -> bool {
234        self.signals.is_stopped_for(self.generation)
235    }
236
237    /// Waits until woken up or until `interval` elapses.
238    ///
239    /// Returns at once when a wakeup is already pending. Always returns
240    /// `true`, after which the loop does its `on_wait_end` work.
241    pub async fn wait_for_running(&self, interval: Duration) -> bool {
242        if self
243            .signals
244            .has_notified
245            .compare_exchange(true, false, Ordering::AcqRel, Ordering::Acquire)
246            .is_ok()
247        {
248            return true;
249        }
250
251        // A timeout is the normal way to leave the wait.
252        let _ = timeout(interval, self.signals.wait_point.notified()).await;
253        self.signals.has_notified.store(false, Ordering::Release);
254        true
255    }
256
257    /// Wakes the loop if it is waiting, or makes its next wait return at once.
258    pub fn wakeup(&self) {
259        self.signals.wakeup();
260    }
261}
262
263impl Default for ServiceTaskContext {
264    fn default() -> Self {
265        Self::new()
266    }
267}
268
269/// Defines service task behavior.
270pub trait ServiceTask: Sync + Send {
271    /// Get the service name
272    fn get_service_name(&self) -> String;
273
274    /// implement the service logic here
275    fn run(&self, context: &ServiceTaskContext) -> impl ::core::future::Future<Output = ()> + Send;
276
277    /// override for custom behavior
278    fn on_wait_end(&self) -> impl ::core::future::Future<Output = ()> + Send {
279        async {
280            // Default implementation does nothing
281        }
282    }
283}
284
285/// Runs a [`ServiceTask`] loop as a service task of an owned task group.
286///
287/// Each run is spawned in a `rocketmq.service-manager` child of the task group
288/// given at construction, so the owner's shutdown report and diagnostics
289/// include it. A stopped manager can be started again.
290///
291/// Shutdown waits no longer than the earliest of the deadline passed to
292/// [`Self::shutdown_until`] and the deadline installed on the parent task
293/// group by an owner that is shutting down. When neither exists, it uses the
294/// runtime's default shutdown budget. An interrupting shutdown is the one
295/// exception: it gives the aborted loop at least one second to confirm its
296/// destruction, so with less than a second left it can return up to one
297/// second after that deadline.
298pub struct ServiceManager<T: ServiceTask + 'static> {
299    service: Arc<T>,
300    signals: Arc<ServiceSignals>,
301    task_handle: Mutex<Option<ServiceTaskHandle>>,
302    last_task_group_shutdown_report: Mutex<Option<ShutdownReport>>,
303    parent_task_group: TaskGroup,
304}
305
306struct ServiceTaskHandle {
307    task_id: TaskId,
308    task_group: TaskGroup,
309}
310
311impl ServiceTaskHandle {
312    async fn shutdown(self, deadline: ShutdownDeadline, interrupt: bool, service_name: &str) -> ShutdownReport {
313        let report = if interrupt {
314            // The abort needs a window to confirm even after the deadline passed.
315            let confirmation = deadline.remaining().max(ABORT_CONFIRMATION_TIMEOUT);
316            if !self.task_group.abort_task_and_wait(self.task_id, confirmation).await {
317                warn!(
318                    "Service thread {} interrupt did not finish before timeout",
319                    service_name
320                );
321            }
322            self.task_group.shutdown(Duration::ZERO).await
323        } else if self.task_group.wait_task(self.task_id, deadline.remaining()).await {
324            self.task_group.shutdown_until(deadline).await
325        } else {
326            warn!("Service thread {} shutdown timeout", service_name);
327            self.task_group.shutdown(Duration::ZERO).await
328        };
329        if !report.is_healthy() {
330            warn!(
331                "Service thread {} shutdown report is unhealthy: {}",
332                service_name,
333                report.to_json()
334            );
335        }
336        report
337    }
338}
339
340fn spawn_service_task<F>(
341    parent_task_group: &TaskGroup,
342    task_name: String,
343    future: F,
344) -> RuntimeResult<ServiceTaskHandle>
345where
346    F: Future<Output = ()> + Send + 'static,
347{
348    let task_group = parent_task_group.component("rocketmq.service-manager");
349    let task_id = task_group
350        .spawn_service(task_name, future)
351        .map_err(|error| RuntimeError::within(crate::RuntimeOperation::SpawnServiceTask, error))?;
352    Ok(ServiceTaskHandle { task_id, task_group })
353}
354
355impl<T: ServiceTask> AsRef<T> for ServiceManager<T> {
356    fn as_ref(&self) -> &T {
357        &self.service
358    }
359}
360
361impl<T: ServiceTask + 'static> ServiceManager<T> {
362    /// Creates a manager whose runs are owned by `parent_task_group`.
363    pub fn new_with_task_group(service: T, parent_task_group: TaskGroup) -> Self {
364        Self::new_arc_with_task_group(Arc::new(service), parent_task_group)
365    }
366
367    /// Creates a manager for a shared service whose runs are owned by
368    /// `parent_task_group`.
369    pub fn new_arc_with_task_group(service: Arc<T>, parent_task_group: TaskGroup) -> Self {
370        Self {
371            service,
372            signals: Arc::new(ServiceSignals::new()),
373            task_handle: Mutex::new(None),
374            last_task_group_shutdown_report: Mutex::new(None),
375            parent_task_group,
376        }
377    }
378
379    /// Starts a run of the service loop.
380    ///
381    /// Starting a manager whose loop is still active only logs a warning.
382    ///
383    /// # Errors
384    ///
385    /// Returns an error if the parent task group no longer accepts work.
386    pub async fn start(&self) -> RuntimeResult<()> {
387        let service_name = self.service.get_service_name();
388        let (replaced, generation) = match self.signals.begin_start() {
389            Ok(started) => started,
390            Err(state) => {
391                warn!(
392                    "Service thread {} is already started, current_state: {:?}",
393                    service_name, state
394                );
395                return Ok(());
396            }
397        };
398        info!("Try to start service thread: {}", service_name);
399
400        let service = self.service.clone();
401        let signals = self.signals.clone();
402        // The service loop is created on the heap by the worker's first poll. An
403        // inline loop would make this future as large as the service's state,
404        // and unoptimized builds copy it at every hop down to the task group.
405        let future = async move {
406            Box::pin(Self::run_internal(service, signals, generation)).await;
407        };
408        let handle = match spawn_service_task(&self.parent_task_group, service_name.clone(), future) {
409            Ok(handle) => handle,
410            Err(error) => {
411                // Nothing was spawned, so restore the state this start replaced.
412                let _ = self.signals.word.compare_exchange(
413                    pack(generation, ServiceTaskState::Starting),
414                    replaced,
415                    Ordering::AcqRel,
416                    Ordering::Acquire,
417                );
418                return Err(error);
419            }
420        };
421
422        // A run that ended on its own leaves its task group behind; release it.
423        let previous = self.task_handle.lock().replace(handle);
424        if let Some(previous) = previous {
425            let _ = previous.task_group.shutdown_now();
426        }
427        self.signals
428            .transition(generation, ServiceTaskState::Starting, ServiceTaskState::Running);
429
430        info!("Started service thread: {}", service_name);
431        Ok(())
432    }
433
434    async fn run_internal(service: Arc<T>, signals: Arc<ServiceSignals>, generation: u64) {
435        let service_name = service.get_service_name();
436        let _stop_on_exit = StopOnExit {
437            signals: signals.clone(),
438            generation,
439        };
440        info!("Service thread {} is running", service_name);
441        signals.transition(generation, ServiceTaskState::Starting, ServiceTaskState::Running);
442
443        let context = ServiceTaskContext { signals, generation };
444        service.run(&context).await;
445
446        info!("Service thread {} has stopped", service_name);
447    }
448
449    /// Stops the service and waits for its loop.
450    ///
451    /// The wait ends at the deadline installed on the parent task group, or
452    /// after the runtime's default shutdown budget when there is none. A loop
453    /// still running then is aborted.
454    ///
455    /// # Errors
456    ///
457    /// Currently always succeeds; the result is kept for callers that
458    /// propagate shutdown failures.
459    pub async fn shutdown(&self) -> RuntimeResult<()> {
460        self.shutdown_inner(false, None).await
461    }
462
463    /// Stops the service like [`Self::shutdown`], aborting the loop at once
464    /// when `interrupt` is set.
465    ///
466    /// # Errors
467    ///
468    /// Currently always succeeds; the result is kept for callers that
469    /// propagate shutdown failures.
470    pub async fn shutdown_with_interrupt(&self, interrupt: bool) -> RuntimeResult<()> {
471        self.shutdown_inner(interrupt, None).await
472    }
473
474    /// Stops the service and waits for its loop no later than `deadline`.
475    ///
476    /// # Errors
477    ///
478    /// Currently always succeeds; the result is kept for callers that
479    /// propagate shutdown failures.
480    pub async fn shutdown_until(&self, deadline: ShutdownDeadline) -> RuntimeResult<()> {
481        self.shutdown_inner(false, Some(deadline)).await
482    }
483
484    /// Stops the service no later than `deadline`, aborting the loop at once
485    /// when `interrupt` is set.
486    ///
487    /// An interrupted loop is given at least one second to confirm the abort,
488    /// so when less than a second of `deadline` remains the call can return up
489    /// to one second after it.
490    ///
491    /// # Errors
492    ///
493    /// Currently always succeeds; the result is kept for callers that
494    /// propagate shutdown failures.
495    pub async fn shutdown_with_interrupt_until(
496        &self,
497        interrupt: bool,
498        deadline: ShutdownDeadline,
499    ) -> RuntimeResult<()> {
500        self.shutdown_inner(interrupt, Some(deadline)).await
501    }
502
503    async fn shutdown_inner(&self, interrupt: bool, requested: Option<ShutdownDeadline>) -> RuntimeResult<()> {
504        let service_name = self.service.get_service_name();
505        let handle = self.task_handle.lock().take();
506        let (state, generation) = self.signals.request_stop();
507        if handle.is_none()
508            && !matches!(
509                state,
510                ServiceTaskState::Starting | ServiceTaskState::Running | ServiceTaskState::Stopping
511            )
512        {
513            warn!("Service thread {} is not running", service_name);
514            return Ok(());
515        }
516
517        info!("Shutdown thread[{}] interrupt={}", service_name, interrupt);
518        self.signals.wakeup();
519
520        let deadline = self.shutdown_deadline(requested);
521        let begin_time = Instant::now();
522        if let Some(handle) = handle {
523            let report = handle.shutdown(deadline, interrupt, &service_name).await;
524            *self.last_task_group_shutdown_report.lock() = Some(report);
525        }
526        info!(
527            "Join thread[{}], elapsed time: {}ms",
528            service_name,
529            begin_time.elapsed().as_millis()
530        );
531
532        self.signals.finish(generation);
533        Ok(())
534    }
535
536    fn shutdown_deadline(&self, requested: Option<ShutdownDeadline>) -> ShutdownDeadline {
537        match (requested, self.parent_task_group.shutdown_deadline()) {
538            (Some(requested), Some(installed)) => requested.earliest(installed),
539            (Some(deadline), None) | (None, Some(deadline)) => deadline,
540            (None, None) => ShutdownDeadline::after(DEFAULT_SHUTDOWN_TIMEOUT),
541        }
542    }
543
544    /// Asks the loop to stop without waiting for it.
545    pub fn make_stop(&self) {
546        let (state, _) = self.signals.request_stop();
547        if matches!(state, ServiceTaskState::Starting | ServiceTaskState::Running) {
548            info!("Make stop thread[{}]", self.service.get_service_name());
549        }
550    }
551
552    /// Wakes the loop if it is waiting, or makes its next wait return at once.
553    pub fn wakeup(&self) {
554        self.signals.wakeup();
555    }
556
557    /// Returns the current lifecycle state.
558    pub fn get_lifecycle_state(&self) -> ServiceTaskState {
559        self.signals.state()
560    }
561
562    /// Returns the shutdown report of the task group used by the last stopped run.
563    pub fn last_task_group_shutdown_report(&self) -> Option<ShutdownReport> {
564        self.last_task_group_shutdown_report.lock().clone()
565    }
566}
567
568#[cfg(test)]
569mod tests {
570    use std::future;
571    use std::sync::atomic::AtomicBool;
572    use std::sync::atomic::AtomicUsize;
573    use std::sync::atomic::Ordering;
574    use std::sync::Arc;
575
576    use tokio::sync::Notify;
577    use tokio::time::Duration;
578
579    use super::*;
580    use crate::RuntimeContext;
581
582    /// Counts loop iterations and signals each one.
583    struct CountingService {
584        iterations: Arc<AtomicUsize>,
585        iterated: Arc<Notify>,
586    }
587
588    impl ServiceTask for CountingService {
589        fn get_service_name(&self) -> String {
590            "counting-service".to_string()
591        }
592
593        async fn run(&self, context: &ServiceTaskContext) {
594            while !context.is_stopped() {
595                context.wait_for_running(Duration::from_secs(60)).await;
596                self.on_wait_end().await;
597            }
598        }
599
600        async fn on_wait_end(&self) {
601            self.iterations.fetch_add(1, Ordering::AcqRel);
602            self.iterated.notify_one();
603        }
604    }
605
606    struct PendingService {
607        entered: Arc<Notify>,
608        dropped: Arc<AtomicBool>,
609    }
610
611    struct DropMarker {
612        dropped: Arc<AtomicBool>,
613    }
614
615    impl Drop for DropMarker {
616        fn drop(&mut self) {
617            self.dropped.store(true, Ordering::Release);
618        }
619    }
620
621    impl ServiceTask for PendingService {
622        fn get_service_name(&self) -> String {
623            "pending-service".to_string()
624        }
625
626        async fn run(&self, _context: &ServiceTaskContext) {
627            let _marker = DropMarker {
628                dropped: Arc::clone(&self.dropped),
629            };
630            self.entered.notify_one();
631            future::pending::<()>().await;
632        }
633    }
634
635    fn counting_service() -> (CountingService, Arc<AtomicUsize>, Arc<Notify>) {
636        let iterations = Arc::new(AtomicUsize::new(0));
637        let iterated = Arc::new(Notify::new());
638        (
639            CountingService {
640                iterations: iterations.clone(),
641                iterated: iterated.clone(),
642            },
643            iterations,
644            iterated,
645        )
646    }
647
648    /// Returns a service that never stops on its own, with its entry signal
649    /// and a flag set when its loop future is dropped.
650    fn pending_service() -> (PendingService, Arc<Notify>, Arc<AtomicBool>) {
651        let entered = Arc::new(Notify::new());
652        let dropped = Arc::new(AtomicBool::new(false));
653        (
654            PendingService {
655                entered: entered.clone(),
656                dropped: dropped.clone(),
657            },
658            entered,
659            dropped,
660        )
661    }
662
663    #[tokio::test]
664    async fn service_runs_under_its_parent_task_group() {
665        let context = RuntimeContext::from_current("service-manager-parent-test");
666        let service_context = context.service_context("service-manager-service");
667        let (service, _, _) = counting_service();
668        let manager = ServiceManager::new_with_task_group(service, service_context.task_group().clone());
669
670        manager.start().await.unwrap();
671        assert_eq!(service_context.task_group().component_count(), 1);
672        manager.shutdown().await.unwrap();
673        let child_report = manager
674            .last_task_group_shutdown_report()
675            .expect("service manager shutdown report should exist");
676        assert_eq!(child_report.name, "rocketmq.service-manager");
677        assert!(child_report.is_healthy(), "{}", child_report.to_json());
678
679        let report = service_context.task_group().shutdown(Duration::from_secs(1)).await;
680        assert!(report.is_healthy(), "{}", report.to_json());
681        assert!(report.children.is_empty(), "{}", report.to_json());
682    }
683
684    #[tokio::test]
685    async fn lifecycle_moves_through_running_and_stopped() {
686        let context = RuntimeContext::from_current("service-manager-lifecycle-test");
687        let service_context = context.service_context("service-manager-lifecycle");
688        let (service, iterations, iterated) = counting_service();
689        let manager = ServiceManager::new_with_task_group(service, service_context.task_group().clone());
690        assert_eq!(manager.get_lifecycle_state(), ServiceTaskState::NotStarted);
691
692        manager.start().await.unwrap();
693        assert_eq!(manager.get_lifecycle_state(), ServiceTaskState::Running);
694
695        let iteration = iterated.notified();
696        manager.wakeup();
697        iteration.await;
698        assert!(iterations.load(Ordering::Acquire) >= 1);
699
700        manager.shutdown().await.unwrap();
701        assert_eq!(manager.get_lifecycle_state(), ServiceTaskState::Stopped);
702    }
703
704    #[tokio::test]
705    async fn second_start_while_running_is_ignored() {
706        let context = RuntimeContext::from_current("service-manager-multi-start-test");
707        let service_context = context.service_context("service-manager-multi-start");
708        let (service, _, _) = counting_service();
709        let manager = ServiceManager::new_with_task_group(service, service_context.task_group().clone());
710
711        manager.start().await.unwrap();
712        manager.start().await.unwrap();
713        assert_eq!(service_context.task_group().component_count(), 1);
714        assert_eq!(manager.get_lifecycle_state(), ServiceTaskState::Running);
715
716        manager.shutdown().await.unwrap();
717        assert_eq!(manager.get_lifecycle_state(), ServiceTaskState::Stopped);
718    }
719
720    #[tokio::test]
721    async fn stopped_manager_can_start_again() {
722        let context = RuntimeContext::from_current("service-manager-restart-test");
723        let service_context = context.service_context("service-manager-restart");
724        let (service, _, _) = counting_service();
725        let manager = ServiceManager::new_with_task_group(service, service_context.task_group().clone());
726
727        manager.start().await.unwrap();
728        manager.shutdown().await.unwrap();
729        manager.start().await.unwrap();
730        assert_eq!(manager.get_lifecycle_state(), ServiceTaskState::Running);
731        manager.shutdown().await.unwrap();
732        assert_eq!(manager.get_lifecycle_state(), ServiceTaskState::Stopped);
733
734        let report = service_context.task_group().shutdown(Duration::from_secs(1)).await;
735        assert!(report.is_healthy(), "{}", report.to_json());
736    }
737
738    #[tokio::test]
739    async fn make_stop_ends_the_loop_without_shutdown() {
740        let context = RuntimeContext::from_current("service-manager-make-stop-test");
741        let service_context = context.service_context("service-manager-make-stop");
742        let (service, _, _) = counting_service();
743        let manager = ServiceManager::new_with_task_group(service, service_context.task_group().clone());
744
745        manager.start().await.unwrap();
746        manager.make_stop();
747        assert_eq!(manager.get_lifecycle_state(), ServiceTaskState::Stopping);
748        manager.wakeup();
749
750        let report = service_context.task_group().shutdown(Duration::from_secs(1)).await;
751        assert!(report.is_healthy(), "{}", report.to_json());
752        assert_eq!(manager.get_lifecycle_state(), ServiceTaskState::Stopped);
753    }
754
755    #[tokio::test]
756    async fn shutdown_deadline_aborts_a_loop_that_does_not_stop() {
757        let context = RuntimeContext::from_current("service-manager-deadline-test");
758        let service_context = context.service_context("service-manager-deadline");
759        let (service, entered, dropped) = pending_service();
760        let manager = ServiceManager::new_with_task_group(service, service_context.task_group().clone());
761
762        manager.start().await.unwrap();
763        entered.notified().await;
764        let started = std::time::Instant::now();
765        manager
766            .shutdown_until(ShutdownDeadline::after(Duration::from_millis(20)))
767            .await
768            .unwrap();
769
770        assert!(
771            started.elapsed() < Duration::from_secs(5),
772            "shutdown should stop at the deadline instead of a fixed join time"
773        );
774        assert!(
775            dropped.load(Ordering::Acquire),
776            "service future should be dropped after the deadline aborts the task"
777        );
778        assert_eq!(manager.get_lifecycle_state(), ServiceTaskState::Stopped);
779        let report = manager
780            .last_task_group_shutdown_report()
781            .expect("service task group shutdown report should be recorded");
782        assert!(!report.is_healthy(), "{}", report.to_json());
783        assert_eq!(report.aborted, 1, "{}", report.to_json());
784        assert_eq!(report.timed_out, 1, "{}", report.to_json());
785    }
786
787    #[tokio::test]
788    async fn shutdown_follows_the_deadline_installed_by_the_parent_owner() {
789        let context = RuntimeContext::from_current("service-manager-installed-deadline-test");
790        let service_context = context.service_context("service-manager-installed-deadline");
791        let (service, entered, dropped) = pending_service();
792        let manager = ServiceManager::new_with_task_group(service, service_context.task_group().clone());
793        manager.start().await.unwrap();
794        entered.notified().await;
795
796        let started = std::time::Instant::now();
797        // The owner installs its deadline when the shutdown call is made, so
798        // the manager's unbounded shutdown below inherits it.
799        let owner_shutdown = service_context.task_group().shutdown(Duration::from_millis(20));
800        let (report, stopped) = tokio::join!(owner_shutdown, manager.shutdown());
801        stopped.unwrap();
802
803        assert!(
804            started.elapsed() < Duration::from_secs(5),
805            "an unbounded shutdown call should inherit the owner's deadline"
806        );
807        assert!(dropped.load(Ordering::Acquire));
808        assert!(!report.is_healthy(), "{}", report.to_json());
809    }
810
811    #[tokio::test]
812    async fn interrupt_aborts_the_loop_at_once() {
813        let context = RuntimeContext::from_current("service-manager-interrupt-test");
814        let service_context = context.service_context("service-manager-interrupt");
815        let (service, entered, dropped) = pending_service();
816        let manager = ServiceManager::new_with_task_group(service, service_context.task_group().clone());
817
818        manager.start().await.unwrap();
819        entered.notified().await;
820        manager.shutdown_with_interrupt(true).await.unwrap();
821
822        assert!(dropped.load(Ordering::Acquire));
823        assert_eq!(manager.get_lifecycle_state(), ServiceTaskState::Stopped);
824    }
825
826    #[test]
827    fn a_replaced_run_sees_itself_stopped() {
828        let signals = Arc::new(ServiceSignals::new());
829        let (_, first) = signals.begin_start().expect("first run starts");
830        let old = ServiceTaskContext {
831            signals: signals.clone(),
832            generation: first,
833        };
834        signals.transition(first, ServiceTaskState::Starting, ServiceTaskState::Running);
835        assert!(!old.is_stopped());
836
837        signals.request_stop();
838        assert!(signals.finish(first));
839        let (_, second) = signals.begin_start().expect("second run starts");
840        signals.transition(second, ServiceTaskState::Starting, ServiceTaskState::Running);
841
842        assert!(old.is_stopped(), "a loop from the first run must not keep running");
843        assert!(!signals.finish(first), "the first run must not stop the second");
844        assert_eq!(signals.state(), ServiceTaskState::Running);
845    }
846
847    #[tokio::test]
848    async fn standalone_context_returns_at_once_for_a_pending_wakeup() {
849        let context = ServiceTaskContext::new();
850        assert!(!context.is_stopped());
851        context.wakeup();
852        assert!(context.wait_for_running(Duration::from_secs(60)).await);
853    }
854}