Skip to main content

near_async/
lib.rs

1pub use near_async_derive::{MultiSend, MultiSenderFrom};
2
3mod functional;
4pub mod futures;
5pub mod instrumentation;
6pub mod messaging;
7pub mod multithread;
8pub mod test_loop;
9pub mod test_utils;
10pub mod thread_pool;
11pub mod tokio;
12
13use crate::futures::FutureSpawner;
14use crate::messaging::Actor;
15use crate::multithread::runtime_handle::{MultithreadRuntimeHandle, spawn_multithread_actor};
16use crate::tokio::runtime_handle::{TokioRuntimeBuilder, spawn_tokio_actor};
17use crate::tokio::{CancellableFutureSpawner, TokioRuntimeHandle};
18pub use near_time as time;
19use parking_lot::Mutex;
20use std::any::type_name;
21use std::sync::Arc;
22use std::sync::atomic::AtomicU64;
23use tokio_util::sync::CancellationToken;
24
25/// Sequence number to be shared for all messages, to distinguish messages when logging.
26static MESSAGE_SEQUENCE_NUM: AtomicU64 = AtomicU64::new(0);
27
28pub(crate) fn next_message_sequence_num() -> u64 {
29    MESSAGE_SEQUENCE_NUM.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
30}
31
32// Quick and dirty way of getting the type name without the module path.
33// Does not work for more complex types like std::sync::Arc<std::sync::atomic::AtomicBool<...>>
34// example near_chunks::shards_manager_actor::ShardsManagerActor -> ShardsManagerActor
35// To support using it with "SpanWrapped<>" types, we trim the trailing '>' characters.
36fn pretty_type_name<T>() -> &'static str {
37    type_name::<T>().rsplit("::").next().unwrap().trim_end_matches('>')
38}
39
40/// Actor that doesn't handle any messages and does nothing. It's used to host a runtime that can
41/// run futures only.
42struct EmptyActor;
43impl Actor for EmptyActor {}
44
45/// Represents a collection of actors, so that they can be shutdown together.
46#[derive(Clone)]
47pub struct ActorSystem {
48    /// Cancellation token used to signal shutdown of Tokio runtimes spawned with this actor system.
49    tokio_cancellation_signal: CancellationToken,
50    /// Cancellation signal used to signal shutdown of multithread actors spawned with this actor
51    /// system. To send the cancellation signal, the sender is dropped, which causes the receivers
52    /// to error.
53    multithread_cancellation_signal: Arc<Mutex<Option<crossbeam_channel::Sender<()>>>>,
54    multithread_cancellation_receiver: crossbeam_channel::Receiver<()>,
55}
56
57impl ActorSystem {
58    pub fn new() -> Self {
59        let mut systems = ACTOR_SYSTEMS.lock();
60        let (multithread_cancellation_sender, multithread_cancellation_receiver) =
61            crossbeam_channel::bounded(0);
62        let ret = Self {
63            tokio_cancellation_signal: CancellationToken::new(),
64            multithread_cancellation_signal: Arc::new(Mutex::new(Some(
65                multithread_cancellation_sender,
66            ))),
67            multithread_cancellation_receiver,
68        };
69        systems.push(ret.clone());
70        ret
71    }
72
73    pub fn stop(&self) {
74        tracing::info!("stopping all actors in actor system");
75        self.tokio_cancellation_signal.cancel();
76        self.multithread_cancellation_signal.lock().take();
77    }
78
79    /// Spawns an actor in a single threaded Tokio runtime and returns a handle to it.
80    /// The handle can be used to get the sender and future spawner for the actor.
81    ///
82    /// ```rust, ignore
83    ///
84    /// struct MyActor;
85    ///
86    /// impl Actor for MyActor {}
87    ///
88    /// impl Handler<MyMessage> for MyActor {
89    ///     fn handle(&mut self, msg: MyMessage) {}
90    /// }
91    ///
92    /// // We can use the actor_handle to create senders and future spawners.
93    /// let actor_handle = actor_system.spawn_tokio_actor(MyActor);
94    ///
95    /// let sender: MyAdapter = actor_handle.sender();
96    /// let future_spawner = actor_handle.future_spawner();
97    /// ```
98    ///
99    /// The sender and future spawner can then be passed onto other components that need to send messages
100    /// to the actor or spawn futures in the runtime of the actor.
101    pub fn spawn_tokio_actor<A: messaging::Actor + Send + 'static>(
102        &self,
103        actor: A,
104    ) -> TokioRuntimeHandle<A> {
105        spawn_tokio_actor(
106            actor,
107            std::any::type_name::<A>().to_string(),
108            self.tokio_cancellation_signal.clone(),
109        )
110    }
111
112    /// A more granular way to build a tokio runtime. It allows spawning futures and getting a handle
113    /// before the actor is constructed (so that the actor can be constructed with the handle,
114    /// for sending messages to itself).
115    pub fn new_tokio_builder<A: messaging::Actor + Send + 'static>(
116        &self,
117    ) -> TokioRuntimeBuilder<A> {
118        TokioRuntimeBuilder::new(
119            pretty_type_name::<A>().to_string(),
120            self.tokio_cancellation_signal.clone(),
121        )
122    }
123
124    /// Spawns a multi-threaded actor which handles messages in a synchronous thread pool.
125    /// Used similarly to `spawn_tokio_actor`, but this actor is intended for CPU-bound tasks,
126    /// can run multiple threads, and does not support futures, timers, or delayed messages.
127    pub fn spawn_multithread_actor<A: messaging::Actor + Send + 'static>(
128        &self,
129        num_threads: usize,
130        make_actor_fn: impl Fn() -> A + Sync + Send + 'static,
131    ) -> MultithreadRuntimeHandle<A> {
132        spawn_multithread_actor(
133            num_threads,
134            make_actor_fn,
135            self.multithread_cancellation_receiver.clone(),
136            None,
137        )
138    }
139
140    /// Returns a future spawner for the actor system on an independent Tokio runtime.
141    /// Note: For typical actors, it is recommended we use the future spawner of the
142    /// actor instead.
143    ///
144    /// This is useful for keeping track of spawned futures and their lifetimes.
145    /// Behind the scenes, this builds a new EmptyActor each time.
146    pub fn new_future_spawner(&self, description: &str) -> Box<dyn FutureSpawner> {
147        let handle = spawn_tokio_actor(
148            EmptyActor,
149            description.to_string(),
150            self.tokio_cancellation_signal.clone(),
151        );
152        handle.future_spawner()
153    }
154
155    /// Returns a future spawner for the actor system on an independent multi-threaded Tokio
156    /// runtime.
157    /// Multi-threaded future spawner does not support instrumentation.
158    pub fn new_multi_threaded_future_spawner(&self, description: &str) -> Box<dyn FutureSpawner> {
159        let handle = CancellableFutureSpawner::new(
160            self.tokio_cancellation_signal.clone(),
161            description.to_string(),
162        );
163        handle.future_spawner()
164    }
165}
166
167/// Spawns a future spawner that is NOT owned by any ActorSystem.
168/// Rather, the returned FutureSpawner, when dropped, will stop the runtime.
169pub fn new_owned_future_spawner(description: &str) -> Box<dyn FutureSpawner> {
170    Box::new(OwnedFutureSpawner {
171        handle: spawn_tokio_actor(EmptyActor, description.to_string(), CancellationToken::new()),
172    })
173}
174
175/// Spawns a multithreaded actor which is NOT owned by any ActorSystem.
176/// Rather, the returned handle, when dropped, will stop the actor and its runtime.
177pub fn new_owned_multithread_actor<A: Actor + Send + 'static>(
178    num_threads: usize,
179    make_actor_fn: impl Fn() -> A + Sync + Send + 'static,
180) -> MultithreadRuntimeHandle<A> {
181    let (cancellation_signal, cancellation_receiver) = crossbeam_channel::bounded::<()>(0);
182    spawn_multithread_actor(
183        num_threads,
184        make_actor_fn,
185        cancellation_receiver,
186        Some(cancellation_signal), // never cancelled
187    )
188}
189
190struct OwnedFutureSpawner {
191    handle: TokioRuntimeHandle<EmptyActor>,
192}
193
194impl FutureSpawner for OwnedFutureSpawner {
195    fn spawn_boxed(&self, description: &'static str, f: crate::futures::BoxFuture<'static, ()>) {
196        self.handle.future_spawner().spawn_boxed(description, f);
197    }
198}
199
200impl Drop for OwnedFutureSpawner {
201    fn drop(&mut self) {
202        self.handle.stop();
203    }
204}
205
206/// Used to determine whether shutdown_all_actors is being used properly. If there are multiple
207/// ActorSystems, shutdown_all_actors shall not be used, but instead the test needs to manage
208/// the shutdown of each ActorSystem individually.
209static ACTOR_SYSTEMS: Mutex<Vec<ActorSystem>> = Mutex::new(Vec::new());
210
211/// Shutdown all actors, assuming at most one ActorSystem.
212/// TODO(#14005): Ideally, shutting down actors should not be done by calling a global function.
213pub fn shutdown_all_actors() {
214    {
215        let systems = ACTOR_SYSTEMS.lock();
216        if systems.len() > 1 {
217            panic!("shutdown_all_actors should not be used when there are multiple ActorSystems");
218        }
219        if let Some(system) = systems.first() {
220            system.stop();
221        }
222    }
223}