1 #![allow(unreachable_pub)]
2 use crate::{
3     runtime::{Handle, BOX_FUTURE_THRESHOLD},
4     task::{JoinHandle, LocalSet},
5     util::trace::SpawnMeta,
6 };
7 use std::{future::Future, io, mem};
8 
9 /// Factory which is used to configure the properties of a new task.
10 ///
11 /// **Note**: This is an [unstable API][unstable]. The public API of this type
12 /// may break in 1.x releases. See [the documentation on unstable
13 /// features][unstable] for details.
14 ///
15 /// Methods can be chained in order to configure it.
16 ///
17 /// Currently, there is only one configuration option:
18 ///
19 /// - [`name`], which specifies an associated name for
20 ///   the task
21 ///
22 /// There are three types of task that can be spawned from a Builder:
23 /// - [`spawn_local`] for executing futures on the current thread
24 /// - [`spawn`] for executing [`Send`] futures on the runtime
25 /// - [`spawn_blocking`] for executing blocking code in the
26 ///   blocking thread pool.
27 ///
28 /// ## Example
29 ///
30 /// ```no_run
31 /// use tokio::net::{TcpListener, TcpStream};
32 ///
33 /// use std::io;
34 ///
35 /// async fn process(socket: TcpStream) {
36 ///     // ...
37 /// # drop(socket);
38 /// }
39 ///
40 /// #[tokio::main]
41 /// async fn main() -> io::Result<()> {
42 ///     let listener = TcpListener::bind("127.0.0.1:8080").await?;
43 ///
44 ///     loop {
45 ///         let (socket, _) = listener.accept().await?;
46 ///
47 ///         tokio::task::Builder::new()
48 ///             .name("tcp connection handler")
49 ///             .spawn(async move {
50 ///                 // Process each socket concurrently.
51 ///                 process(socket).await
52 ///             })?;
53 ///     }
54 /// }
55 /// ```
56 /// [unstable]: crate#unstable-features
57 /// [`name`]: Builder::name
58 /// [`spawn_local`]: Builder::spawn_local
59 /// [`spawn`]: Builder::spawn
60 /// [`spawn_blocking`]: Builder::spawn_blocking
61 #[derive(Default, Debug)]
62 #[cfg_attr(docsrs, doc(cfg(all(tokio_unstable, feature = "tracing"))))]
63 pub struct Builder<'a> {
64     name: Option<&'a str>,
65 }
66 
67 impl<'a> Builder<'a> {
68     /// Creates a new task builder.
new() -> Self69     pub fn new() -> Self {
70         Self::default()
71     }
72 
73     /// Assigns a name to the task which will be spawned.
name(&self, name: &'a str) -> Self74     pub fn name(&self, name: &'a str) -> Self {
75         Self { name: Some(name) }
76     }
77 
78     /// Spawns a task with this builder's settings on the current runtime.
79     ///
80     /// # Panics
81     ///
82     /// This method panics if called outside of a Tokio runtime.
83     ///
84     /// See [`task::spawn`](crate::task::spawn()) for
85     /// more details.
86     #[track_caller]
spawn<Fut>(self, future: Fut) -> io::Result<JoinHandle<Fut::Output>> where Fut: Future + Send + 'static, Fut::Output: Send + 'static,87     pub fn spawn<Fut>(self, future: Fut) -> io::Result<JoinHandle<Fut::Output>>
88     where
89         Fut: Future + Send + 'static,
90         Fut::Output: Send + 'static,
91     {
92         let fut_size = mem::size_of::<Fut>();
93         Ok(if fut_size > BOX_FUTURE_THRESHOLD {
94             super::spawn::spawn_inner(Box::pin(future), SpawnMeta::new(self.name, fut_size))
95         } else {
96             super::spawn::spawn_inner(future, SpawnMeta::new(self.name, fut_size))
97         })
98     }
99 
100     /// Spawn a task with this builder's settings on the provided [runtime
101     /// handle].
102     ///
103     /// See [`Handle::spawn`] for more details.
104     ///
105     /// [runtime handle]: crate::runtime::Handle
106     /// [`Handle::spawn`]: crate::runtime::Handle::spawn
107     #[track_caller]
spawn_on<Fut>(self, future: Fut, handle: &Handle) -> io::Result<JoinHandle<Fut::Output>> where Fut: Future + Send + 'static, Fut::Output: Send + 'static,108     pub fn spawn_on<Fut>(self, future: Fut, handle: &Handle) -> io::Result<JoinHandle<Fut::Output>>
109     where
110         Fut: Future + Send + 'static,
111         Fut::Output: Send + 'static,
112     {
113         let fut_size = mem::size_of::<Fut>();
114         Ok(if fut_size > BOX_FUTURE_THRESHOLD {
115             handle.spawn_named(Box::pin(future), SpawnMeta::new(self.name, fut_size))
116         } else {
117             handle.spawn_named(future, SpawnMeta::new(self.name, fut_size))
118         })
119     }
120 
121     /// Spawns `!Send` a task on the current [`LocalSet`] with this builder's
122     /// settings.
123     ///
124     /// The spawned future will be run on the same thread that called `spawn_local`.
125     /// This may only be called from the context of a [local task set][`LocalSet`].
126     ///
127     /// # Panics
128     ///
129     /// This function panics if called outside of a [local task set][`LocalSet`].
130     ///
131     /// See [`task::spawn_local`] for more details.
132     ///
133     /// [`task::spawn_local`]: crate::task::spawn_local
134     /// [`LocalSet`]: crate::task::LocalSet
135     #[track_caller]
spawn_local<Fut>(self, future: Fut) -> io::Result<JoinHandle<Fut::Output>> where Fut: Future + 'static, Fut::Output: 'static,136     pub fn spawn_local<Fut>(self, future: Fut) -> io::Result<JoinHandle<Fut::Output>>
137     where
138         Fut: Future + 'static,
139         Fut::Output: 'static,
140     {
141         let fut_size = mem::size_of::<Fut>();
142         Ok(if fut_size > BOX_FUTURE_THRESHOLD {
143             super::local::spawn_local_inner(Box::pin(future), SpawnMeta::new(self.name, fut_size))
144         } else {
145             super::local::spawn_local_inner(future, SpawnMeta::new(self.name, fut_size))
146         })
147     }
148 
149     /// Spawns `!Send` a task on the provided [`LocalSet`] with this builder's
150     /// settings.
151     ///
152     /// See [`LocalSet::spawn_local`] for more details.
153     ///
154     /// [`LocalSet::spawn_local`]: crate::task::LocalSet::spawn_local
155     /// [`LocalSet`]: crate::task::LocalSet
156     #[track_caller]
spawn_local_on<Fut>( self, future: Fut, local_set: &LocalSet, ) -> io::Result<JoinHandle<Fut::Output>> where Fut: Future + 'static, Fut::Output: 'static,157     pub fn spawn_local_on<Fut>(
158         self,
159         future: Fut,
160         local_set: &LocalSet,
161     ) -> io::Result<JoinHandle<Fut::Output>>
162     where
163         Fut: Future + 'static,
164         Fut::Output: 'static,
165     {
166         let fut_size = mem::size_of::<Fut>();
167         Ok(if fut_size > BOX_FUTURE_THRESHOLD {
168             local_set.spawn_named(Box::pin(future), SpawnMeta::new(self.name, fut_size))
169         } else {
170             local_set.spawn_named(future, SpawnMeta::new(self.name, fut_size))
171         })
172     }
173 
174     /// Spawns blocking code on the blocking threadpool.
175     ///
176     /// # Panics
177     ///
178     /// This method panics if called outside of a Tokio runtime.
179     ///
180     /// See [`task::spawn_blocking`](crate::task::spawn_blocking)
181     /// for more details.
182     #[track_caller]
spawn_blocking<Function, Output>( self, function: Function, ) -> io::Result<JoinHandle<Output>> where Function: FnOnce() -> Output + Send + 'static, Output: Send + 'static,183     pub fn spawn_blocking<Function, Output>(
184         self,
185         function: Function,
186     ) -> io::Result<JoinHandle<Output>>
187     where
188         Function: FnOnce() -> Output + Send + 'static,
189         Output: Send + 'static,
190     {
191         let handle = Handle::current();
192         self.spawn_blocking_on(function, &handle)
193     }
194 
195     /// Spawns blocking code on the provided [runtime handle]'s blocking threadpool.
196     ///
197     /// See [`Handle::spawn_blocking`] for more details.
198     ///
199     /// [runtime handle]: crate::runtime::Handle
200     /// [`Handle::spawn_blocking`]: crate::runtime::Handle::spawn_blocking
201     #[track_caller]
spawn_blocking_on<Function, Output>( self, function: Function, handle: &Handle, ) -> io::Result<JoinHandle<Output>> where Function: FnOnce() -> Output + Send + 'static, Output: Send + 'static,202     pub fn spawn_blocking_on<Function, Output>(
203         self,
204         function: Function,
205         handle: &Handle,
206     ) -> io::Result<JoinHandle<Output>>
207     where
208         Function: FnOnce() -> Output + Send + 'static,
209         Output: Send + 'static,
210     {
211         use crate::runtime::Mandatory;
212         let fn_size = mem::size_of::<Function>();
213         let (join_handle, spawn_result) = if fn_size > BOX_FUTURE_THRESHOLD {
214             handle.inner.blocking_spawner().spawn_blocking_inner(
215                 Box::new(function),
216                 Mandatory::NonMandatory,
217                 SpawnMeta::new(self.name, fn_size),
218                 handle,
219             )
220         } else {
221             handle.inner.blocking_spawner().spawn_blocking_inner(
222                 function,
223                 Mandatory::NonMandatory,
224                 SpawnMeta::new(self.name, fn_size),
225                 handle,
226             )
227         };
228 
229         spawn_result?;
230         Ok(join_handle)
231     }
232 }
233