zenoh_task/
lib.rs

1//
2// Copyright (c) 2024 ZettaScale Technology
3//
4// This program and the accompanying materials are made available under the
5// terms of the Eclipse Public License 2.0 which is available at
6// http://www.eclipse.org/legal/epl-2.0, or the Apache License, Version 2.0
7// which is available at https://www.apache.org/licenses/LICENSE-2.0.
8//
9// SPDX-License-Identifier: EPL-2.0 OR Apache-2.0
10//
11// Contributors:
12//   ZettaScale Zenoh Team, <zenoh@zettascale.tech>
13//
14
15//! ⚠️ WARNING ⚠️
16//!
17//! This module is intended for Zenoh's internal use.
18//!
19//! [Click here for Zenoh's documentation](https://docs.rs/zenoh/latest/zenoh)
20
21use std::{future::Future, time::Duration};
22
23use futures::future::FutureExt;
24use tokio::task::JoinHandle;
25use tokio_util::{sync::CancellationToken, task::TaskTracker};
26use zenoh_core::{ResolveFuture, Wait};
27use zenoh_runtime::ZRuntime;
28
29#[derive(Clone)]
30pub struct TaskController {
31    tracker: TaskTracker,
32    token: CancellationToken,
33}
34
35impl std::fmt::Debug for TaskController {
36    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
37        f.debug_struct("TaskController")
38            .field("is_cancelled", &self.token.is_cancelled())
39            .finish_non_exhaustive()
40    }
41}
42
43impl Default for TaskController {
44    fn default() -> Self {
45        TaskController {
46            tracker: TaskTracker::new(),
47            token: CancellationToken::new(),
48        }
49    }
50}
51
52impl TaskController {
53    /// Converts a task to abortable one, which can later be terminated by call to [`TaskController::terminate_all()`].
54    pub fn into_abortable<'a, F, T>(&self, future: F) -> impl Future<Output = Option<T>> + Send + 'a
55    where
56        F: Future<Output = T> + Send + 'a,
57        T: Send + 'static,
58    {
59        self.token.child_token().run_until_cancelled_owned(future)
60    }
61
62    /// Spawns a task that can be later terminated by call to [`TaskController::terminate_all()`].
63    /// Task output is ignored.
64    pub fn spawn_abortable<F, T>(&self, future: F) -> JoinHandle<Option<T>>
65    where
66        F: Future<Output = T> + Send + 'static,
67        T: Send + 'static,
68    {
69        #[cfg(feature = "tracing-instrument")]
70        let future = tracing::Instrument::instrument(future, tracing::Span::current());
71
72        self.tracker.spawn(self.into_abortable(future))
73    }
74
75    /// Spawns a task using a specified runtime that can be later terminated by call to [`TaskController::terminate_all()`].
76    pub fn spawn_abortable_with_rt<F, T>(&self, rt: ZRuntime, future: F) -> JoinHandle<Option<T>>
77    where
78        F: Future<Output = T> + Send + 'static,
79        T: Send + 'static,
80    {
81        #[cfg(feature = "tracing-instrument")]
82        let future = tracing::Instrument::instrument(future, tracing::Span::current());
83
84        self.tracker.spawn_on(self.into_abortable(future), &rt)
85    }
86
87    pub fn get_cancellation_token(&self) -> CancellationToken {
88        self.token.child_token()
89    }
90
91    /// Spawns a task that can be cancelled cancellation of a token obtained by [`TaskController::get_cancellation_token()`],
92    /// was created via [`TaskController::into_abortable()`],
93    /// or can run to completion in finite amount of time, using a specified runtime.
94    /// It can be later terminated by call to [`TaskController::terminate_all()`].
95    pub fn spawn<F, T>(&self, future: F) -> JoinHandle<T>
96    where
97        F: Future<Output = T> + Send + 'static,
98        T: Send + 'static,
99    {
100        #[cfg(feature = "tracing-instrument")]
101        let future = tracing::Instrument::instrument(future, tracing::Span::current());
102
103        self.tracker.spawn(future)
104    }
105
106    /// Spawns a task which can be cancelled via cancellation of a token obtained by [`TaskController::get_cancellation_token()`],
107    /// was created via [`TaskController::into_abortable()`],
108    /// or can run to completion in finite amount of time, using a specified runtime.
109    /// It can be later aborted by call to [`TaskController::terminate_all()`].
110    pub fn spawn_with_rt<F, T>(&self, rt: ZRuntime, future: F) -> JoinHandle<T>
111    where
112        F: Future<Output = T> + Send + 'static,
113        T: Send + 'static,
114    {
115        #[cfg(feature = "tracing-instrument")]
116        let future = tracing::Instrument::instrument(future, tracing::Span::current());
117
118        self.tracker.spawn_on(future, &rt)
119    }
120
121    /// Attempts tp terminate all previously spawned tasks
122    /// The caller must ensure that all tasks spawned with [`TaskController::spawn()`]
123    /// or [`TaskController::spawn_with_rt()`] can yield in finite amount of time either because they will run to completion
124    /// or due to cancellation of token acquired via [`TaskController::get_cancellation_token()`].
125    /// Tasks spawned with [`TaskController::spawn_abortable()`] or [`TaskController::spawn_abortable_with_rt()`] will be aborted (i.e. terminated upon next await call).
126    /// The call blocks until all tasks yield or timeout duration expires.
127    /// Returns 0 in case of success, number of non terminated tasks otherwise.
128    pub fn terminate_all(&self, timeout: Duration) -> usize {
129        ResolveFuture::new(async move {
130            if tokio::time::timeout(timeout, self.terminate_all_async())
131                .await
132                .is_err()
133            {
134                tracing::error!("Failed to terminate {} tasks", self.tracker.len());
135            }
136            self.tracker.len()
137        })
138        .wait()
139    }
140
141    /// Async version of [`TaskController::terminate_all()`].
142    pub async fn terminate_all_async(&self) {
143        self.tracker.close();
144        self.token.cancel();
145        self.tracker.wait().await
146    }
147}
148
149pub struct TerminatableTask {
150    handle: Option<JoinHandle<()>>,
151    token: CancellationToken,
152}
153
154impl std::fmt::Debug for TerminatableTask {
155    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
156        f.debug_struct("TerminatableTask")
157            .field("has_handle", &self.handle.is_some())
158            .field("is_cancelled", &self.token.is_cancelled())
159            .finish()
160    }
161}
162
163impl Drop for TerminatableTask {
164    fn drop(&mut self) {
165        self.terminate(std::time::Duration::from_secs(10));
166    }
167}
168
169impl TerminatableTask {
170    pub fn create_cancellation_token() -> CancellationToken {
171        CancellationToken::new()
172    }
173
174    /// Spawns a task that can be later terminated by [`TerminatableTask::terminate()`].
175    /// Prior to termination attempt the specified cancellation token will be cancelled.
176    pub fn spawn<F, T>(rt: ZRuntime, future: F, token: CancellationToken) -> TerminatableTask
177    where
178        F: Future<Output = T> + Send + 'static,
179        T: Send + 'static,
180    {
181        TerminatableTask {
182            handle: Some(rt.spawn(future.map(|_f| ()))),
183            token,
184        }
185    }
186
187    /// Spawns a task that can be later aborted by [`TerminatableTask::terminate()`].
188    pub fn spawn_abortable<F, T>(rt: ZRuntime, future: F) -> TerminatableTask
189    where
190        F: Future<Output = T> + Send + 'static,
191        T: Send + 'static,
192    {
193        let token = CancellationToken::new();
194        let token2 = token.clone();
195        let task = async move {
196            tokio::select! {
197                _ = token2.cancelled() => {},
198                _ = future => {}
199            }
200        };
201
202        TerminatableTask {
203            handle: Some(rt.spawn(task)),
204            token,
205        }
206    }
207
208    /// Attempts to terminate the task.
209    /// Returns true if task completed / aborted within timeout duration, false otherwise.
210    pub fn terminate(&mut self, timeout: Duration) -> bool {
211        ResolveFuture::new(async move {
212            if tokio::time::timeout(timeout, self.terminate_async())
213                .await
214                .is_err()
215            {
216                tracing::error!("Failed to terminate the task");
217                return false;
218            };
219            true
220        })
221        .wait()
222    }
223
224    /// Async version of [`TerminatableTask::terminate()`].
225    pub async fn terminate_async(&mut self) {
226        self.token.cancel();
227        if let Some(handle) = self.handle.take() {
228            let _ = handle.await;
229        }
230    }
231}
zenoh_task/lib.rs

zenoh_task/
lib.rs
