dynamo_runtime/
worker.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! The [Worker] class is a convenience wrapper around the construction of the [Runtime]
5//! and execution of the users application.
6//!
7//! In the future, the [Worker] should probably be moved to a procedural macro similar
8//! to the `#[tokio::main]` attribute, where we might annotate an async main function with
9//! `#[dynamo::main]` or similar.
10//!
11//! The [Worker::execute] method is designed to be called once from main and will block
12//! the calling thread until the application completes or is canceled. The method initialized
13//! the signal handler used to trap `SIGINT` and `SIGTERM` signals and trigger a graceful shutdown.
14//!
15//! On termination, the user application is given a graceful shutdown period of controlled by
16//! the [DYN_WORKER_GRACEFUL_SHUTDOWN_TIMEOUT] environment variable. If the application does not
17//! shutdown in time, the worker will terminate the application with an exit code of 911.
18//!
19//! The default values of [DYN_WORKER_GRACEFUL_SHUTDOWN_TIMEOUT] differ between the development
20//! and release builds. In development, the default is [DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT_DEBUG] and
21//! in release, the default is [DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT_RELEASE].
22
23use super::{CancellationToken, Result, Runtime, RuntimeConfig, error};
24
25use futures::Future;
26use once_cell::sync::OnceCell;
27use parking_lot::Mutex;
28use std::time::Duration;
29use tokio::{signal, task::JoinHandle};
30
31static RT: OnceCell<tokio::runtime::Runtime> = OnceCell::new();
32static RTHANDLE: OnceCell<tokio::runtime::Handle> = OnceCell::new();
33static INIT: OnceCell<Mutex<Option<tokio::task::JoinHandle<Result<()>>>>> = OnceCell::new();
34
35const SHUTDOWN_MESSAGE: &str =
36    "Application received shutdown signal; attempting to gracefully shutdown";
37const SHUTDOWN_TIMEOUT_MESSAGE: &str =
38    "Use DYN_WORKER_GRACEFUL_SHUTDOWN_TIMEOUT to control the graceful shutdown timeout";
39
40/// Environment variable to control the graceful shutdown timeout
41pub const DYN_WORKER_GRACEFUL_SHUTDOWN_TIMEOUT: &str = "DYN_WORKER_GRACEFUL_SHUTDOWN_TIMEOUT";
42
43/// Default graceful shutdown timeout in seconds in debug mode
44pub const DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT_DEBUG: u64 = 5;
45
46/// Default graceful shutdown timeout in seconds in release mode
47pub const DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT_RELEASE: u64 = 30;
48
49#[derive(Debug, Clone)]
50pub struct Worker {
51    runtime: Runtime,
52    config: RuntimeConfig,
53}
54
55impl Worker {
56    /// Create a new [`Worker`] instance from [`RuntimeConfig`] settings which is sourced from the environment
57    pub fn from_settings() -> Result<Worker> {
58        let config = RuntimeConfig::from_settings()?;
59        Worker::from_config(config)
60    }
61
62    /// Create a new [`Worker`] instance from a provided [`RuntimeConfig`]
63    pub fn from_config(config: RuntimeConfig) -> Result<Worker> {
64        // if the runtime is already initialized, return an error
65        if RT.get().is_some() || RTHANDLE.get().is_some() {
66            return Err(error!("Worker already initialized"));
67        }
68
69        // create a new runtime and insert it into the OnceCell
70        // there is still a potential race-condition here, two threads cou have passed the first check
71        // but only one will succeed in inserting the runtime
72        let rt = RT.try_insert(config.create_runtime()?).map_err(|_| {
73            error!("Failed to create worker; Only a single Worker should ever be created")
74        })?;
75
76        let runtime = Runtime::from_handle(rt.handle().clone())?;
77        Ok(Worker { runtime, config })
78    }
79
80    pub fn runtime_from_existing() -> Result<Runtime> {
81        if let Some(rt) = RT.get() {
82            Ok(Runtime::from_handle(rt.handle().clone())?)
83        } else if let Some(rt) = RTHANDLE.get() {
84            Ok(Runtime::from_handle(rt.clone())?)
85        } else {
86            Runtime::from_settings()
87        }
88    }
89
90    pub fn tokio_runtime(&self) -> Result<&'static tokio::runtime::Runtime> {
91        RT.get().ok_or_else(|| error!("Worker not initialized"))
92    }
93
94    pub fn runtime(&self) -> &Runtime {
95        &self.runtime
96    }
97
98    pub fn execute<F, Fut>(self, f: F) -> Result<()>
99    where
100        F: FnOnce(Runtime) -> Fut + Send + 'static,
101        Fut: Future<Output = Result<()>> + Send + 'static,
102    {
103        let runtime = self.runtime.clone();
104        runtime.secondary().block_on(self.execute_internal(f))??;
105        runtime.shutdown();
106        Ok(())
107    }
108
109    pub async fn execute_async<F, Fut>(self, f: F) -> Result<()>
110    where
111        F: FnOnce(Runtime) -> Fut + Send + 'static,
112        Fut: Future<Output = Result<()>> + Send + 'static,
113    {
114        let runtime = self.runtime.clone();
115        let task = self.execute_internal(f);
116        task.await??;
117        runtime.shutdown();
118        Ok(())
119    }
120
121    /// Executes the provided application/closure on the [`Runtime`].
122    /// This is designed to be called once from main and will block the calling thread until the application completes.
123    fn execute_internal<F, Fut>(self, f: F) -> JoinHandle<Result<()>>
124    where
125        F: FnOnce(Runtime) -> Fut + Send + 'static,
126        Fut: Future<Output = Result<()>> + Send + 'static,
127    {
128        let runtime = self.runtime.clone();
129        let primary = runtime.primary();
130        let secondary = runtime.secondary();
131
132        let timeout = std::env::var(DYN_WORKER_GRACEFUL_SHUTDOWN_TIMEOUT)
133            .ok()
134            .and_then(|s| s.parse::<u64>().ok())
135            .unwrap_or({
136                if cfg!(debug_assertions) {
137                    DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT_DEBUG
138                } else {
139                    DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT_RELEASE
140                }
141            });
142
143        INIT.set(Mutex::new(Some(secondary.spawn(async move {
144            // start signal handler
145            tokio::spawn(signal_handler(runtime.cancellation_token.clone()));
146
147            let cancel_token = runtime.child_token();
148            let (mut app_tx, app_rx) = tokio::sync::oneshot::channel::<()>();
149
150            // spawn a task to run the application
151            let task: JoinHandle<Result<()>> = primary.spawn(async move {
152                let _rx = app_rx;
153                f(runtime).await
154            });
155
156            tokio::select! {
157                _ = cancel_token.cancelled() => {
158                    tracing::debug!("{}", SHUTDOWN_MESSAGE);
159                    tracing::debug!("{} {} seconds", SHUTDOWN_TIMEOUT_MESSAGE, timeout);
160                }
161
162                _ = app_tx.closed() => {
163                }
164            };
165
166            let result = tokio::select! {
167                result = task => {
168                    result
169                }
170
171                _ = tokio::time::sleep(tokio::time::Duration::from_secs(timeout)) => {
172                    tracing::debug!("Application did not shutdown in time; terminating");
173                    std::process::exit(911);
174                }
175            }?;
176
177            match &result {
178                Ok(_) => {
179                    tracing::debug!("Application shutdown successfully");
180                }
181                Err(e) => {
182                    tracing::error!("Application shutdown with error: {:?}", e);
183                }
184            }
185
186            result
187        }))))
188        .expect("Failed to spawn application task");
189
190        INIT
191            .get()
192            .expect("Application task not initialized")
193            .lock()
194            .take()
195            .expect("Application initialized; but another thread is awaiting it; Worker.execute() can only be called once")
196    }
197
198    pub fn from_current() -> Result<Worker> {
199        if RT.get().is_some() || RTHANDLE.get().is_some() {
200            return Err(error!("Worker already initialized"));
201        }
202        let runtime = Runtime::from_current()?;
203        let config = RuntimeConfig::from_settings()?;
204        Ok(Worker { runtime, config })
205    }
206}
207
208/// Catch signals and trigger a shutdown
209async fn signal_handler(cancel_token: CancellationToken) -> Result<()> {
210    let ctrl_c = async {
211        signal::ctrl_c().await?;
212        anyhow::Ok(())
213    };
214
215    let sigterm = async {
216        signal::unix::signal(signal::unix::SignalKind::terminate())?
217            .recv()
218            .await;
219        anyhow::Ok(())
220    };
221
222    tokio::select! {
223        _ = ctrl_c => {
224            tracing::info!("Ctrl+C received, starting graceful shutdown");
225        },
226        _ = sigterm => {
227            tracing::info!("SIGTERM received, starting graceful shutdown");
228        },
229        _ = cancel_token.cancelled() => {
230            tracing::debug!("CancellationToken triggered; shutting down");
231        },
232    }
233
234    // trigger a shutdown
235    cancel_token.cancel();
236
237    Ok(())
238}
dynamo_runtime/worker.rs

dynamo_runtime/
worker.rs
