hive-router 0.2.0

GraphQL router for Federation, part of the Hive platform
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111

use async_trait::async_trait;
use ntex::rt::{spawn, JoinHandle};
use std::future::Future;
pub use tokio_util::sync::CancellationToken;
use tracing::{debug, info};

use crate::telemetry::logging::targets;

#[async_trait]
pub trait BackgroundTask: Send + Sync {
    fn id(&self) -> &str;
    async fn run(&self, token: CancellationToken);
}

pub struct BackgroundTasksManager {
    cancellation_token: CancellationToken,
    handles: Vec<JoinHandle<()>>,
    graceful_handles: Vec<JoinHandle<()>>,
}

impl Default for BackgroundTasksManager {
    fn default() -> Self {
        Self::new()
    }
}

impl BackgroundTasksManager {
    pub fn new() -> Self {
        Self {
            cancellation_token: CancellationToken::new(),
            handles: Vec::new(),
            graceful_handles: Vec::new(),
        }
    }

    pub fn register_task<T>(&mut self, task: T)
    where
        T: BackgroundTask + 'static,
    {
        info!(
            target: targets::CORE,
            task_id = task.id(),
            "registering background task"
        );

        let child_token = self.cancellation_token.clone();
        let handle = spawn(async move {
            task.run(child_token).await;
        });
        self.handles.push(handle);
    }

    /// Registers a task whose cancellation cleanup must finish before router shutdown completes.
    pub fn register_graceful_task<T>(&mut self, task: T)
    where
        T: BackgroundTask + 'static,
    {
        info!(
            target: targets::CORE,
            task_id = task.id(),
            "registering background task"
        );

        let child_token = self.cancellation_token.clone();
        self.graceful_handles.push(spawn(async move {
            task.run(child_token).await;
        }));
    }

    pub fn register_handle<F>(&mut self, f: F)
    where
        F: Future<Output = ()> + Send + 'static,
    {
        self.handles.push(spawn(f));
    }

    pub fn shutdown(&mut self) {
        debug!(
            target: targets::CORE,
            "shutdown triggered, stopping all background tasks..."
        );
        self.cancellation_token.cancel();
        for handle in self
            .handles
            .drain(..)
            .chain(self.graceful_handles.drain(..))
        {
            handle.cancel();
        }
        info!(target: targets::CORE, "all background tasks have been shut down.");
    }

    /// Stops ordinary tasks immediately and waits for graceful task cleanup to finish.
    pub async fn graceful_shutdown(&mut self) {
        debug!(
            target: targets::CORE,
            "shutdown triggered, stopping all background tasks..."
        );
        self.cancellation_token.cancel();
        for handle in self.handles.drain(..) {
            handle.cancel();
        }
        for handle in self.graceful_handles.drain(..) {
            let _ = handle.await;
        }
        info!(
            target: targets::CORE,
            "all background tasks have been shut down gracefully."
        );
    }
}