near_async/multithread/
runtime_handle.rs

1use crate::instrumentation::queue::InstrumentedQueue;
2use crate::instrumentation::writer::InstrumentedThreadWriterSharedPart;
3use crate::messaging::Actor;
4use crate::pretty_type_name;
5use std::sync::Arc;
6use std::thread;
7use std::time::Duration;
8
9/// MultithreadRuntimeMessage is a type alias for a boxed function that can be sent to the multithread runtime,
10/// as well as a description for debugging purposes.
11pub(super) struct MultithreadRuntimeMessage<A> {
12    pub(super) seq: u64,
13    pub(super) enqueued_time_ns: u64,
14    pub(super) name: &'static str,
15    pub(super) function: Box<dyn FnOnce(&mut A) + Send>,
16}
17
18/// Allows sending messages to a multithreaded actor runtime. Implements CanSend and CanSendAsync traits
19/// for the messages that the actor can handle.
20pub struct MultithreadRuntimeHandle<A> {
21    pub(super) sender: crossbeam_channel::Sender<MultithreadRuntimeMessage<A>>,
22    /// This is used in the case where the handle controls the lifetime of the runtime,
23    /// dropping (all of) which automatically stops the runtime, as an alterative of having
24    /// the ActorSystem control it.
25    cancellation_signal_holder: Option<crossbeam_channel::Sender<()>>,
26    pub(super) instrumentation: Arc<InstrumentedThreadWriterSharedPart>,
27}
28
29impl<A> Clone for MultithreadRuntimeHandle<A> {
30    fn clone(&self) -> Self {
31        Self {
32            sender: self.sender.clone(),
33            cancellation_signal_holder: self.cancellation_signal_holder.clone(),
34            instrumentation: self.instrumentation.clone(),
35        }
36    }
37}
38
39impl<A> MultithreadRuntimeHandle<A>
40where
41    A: 'static,
42{
43    pub fn sender(&self) -> Arc<MultithreadRuntimeHandle<A>> {
44        Arc::new(self.clone())
45    }
46}
47
48impl<A> MultithreadRuntimeHandle<A> {
49    pub(super) fn send_message(
50        &self,
51        message: MultithreadRuntimeMessage<A>,
52    ) -> Result<(), crossbeam_channel::SendError<MultithreadRuntimeMessage<A>>> {
53        let name = message.name;
54        self.sender.send(message).map(|_| {
55            // Only increment the queue if the message was successfully sent.
56            self.instrumentation.queue().enqueue(name);
57        })
58    }
59}
60
61/// See ActorSystem::spawn_multithread_actor.
62///
63/// The `cancellation_signal_holder` is an optional sender that can be used to disable
64/// system-wide cancellation. If this sender is used, it is just the other side of the
65/// `cancellation_signal`.
66pub(crate) fn spawn_multithread_actor<A>(
67    num_threads: usize,
68    make_actor_fn: impl Fn() -> A + Sync + Send + 'static,
69    cancellation_signal: crossbeam_channel::Receiver<()>,
70    cancellation_signal_holder: Option<crossbeam_channel::Sender<()>>,
71) -> MultithreadRuntimeHandle<A>
72where
73    A: Actor + Send + 'static,
74{
75    let actor_name = pretty_type_name::<A>();
76    tracing::debug!(
77        target: "multithread_runtime",
78        actor_name,
79        num_threads,
80        "starting multithread actor",
81    );
82    let (sender, receiver) = crossbeam_channel::unbounded::<MultithreadRuntimeMessage<A>>();
83    let instrumented_queue = InstrumentedQueue::new(actor_name);
84    let shared_instrumentation =
85        InstrumentedThreadWriterSharedPart::new(actor_name.to_string(), instrumented_queue.clone());
86    let handle = MultithreadRuntimeHandle {
87        sender,
88        cancellation_signal_holder,
89        instrumentation: shared_instrumentation,
90    };
91    let make_actor_fn = Arc::new(make_actor_fn);
92
93    // Spawn num_threads OS-level threads
94    for thread_id in 0..num_threads {
95        let receiver = receiver.clone();
96        let cancellation_signal = cancellation_signal.clone();
97        let instrumented_queue = instrumented_queue.clone();
98        let handle = handle.clone();
99        let make_actor_fn = make_actor_fn.clone();
100
101        thread::spawn(move || {
102            let mut instrumentation =
103                handle.instrumentation.new_writer_with_global_registration(Some(thread_id));
104            let mut actor = make_actor_fn();
105            let window_update_ticker = crossbeam_channel::tick(Duration::from_secs(1));
106            loop {
107                crossbeam_channel::select! {
108                    recv(cancellation_signal) -> _ => {
109                        tracing::info!(target: "multithread_runtime", actor_name, thread_id, "cancellation received, exiting loop");
110                        return;
111                    }
112                    recv(window_update_ticker) -> _ => {
113                        tracing::trace!(target: "multithread_runtime", actor_name, thread_id, "updating instrumentation window");
114                        instrumentation.advance_window_if_needed();
115                    }
116                    recv(receiver) -> message => {
117                        let Ok(message) = message else {
118                            tracing::warn!(target: "multithread_runtime", actor_name, thread_id, "message queue closed, exiting event loop");
119                            return;
120                        };
121                        instrumented_queue.dequeue(message.name);
122                        let seq = message.seq;
123                        let dequeue_time_ns = handle.instrumentation.current_time().saturating_sub(message.enqueued_time_ns);
124                        instrumentation.start_event(message.name, dequeue_time_ns);
125                        tracing::trace!(target: "multithread_runtime", seq, "executing message");
126                        (message.function)(&mut actor);
127                        instrumentation.end_event(message.name);
128                    }
129                }
130            }
131        });
132    }
133    handle
134}
near_async/multithread/runtime_handle.rs

near_async/multithread/
runtime_handle.rs
