1use std::collections::hash_map::Entry::*;
2use std::sync::Arc;
3
4use futures::{FutureExt, StreamExt};
5use mm1_address::address::Address;
6use mm1_common::log;
7use mm1_core::actor_exit::ActorExit;
8use tokio::sync::{Mutex, mpsc};
9
10use super::MainActorOutcome;
11use crate::rt::actor_fn::ActorFn;
12use crate::rt::{
13 ActorEntry, ActorTaskOutcome, AddressLease, Event, EventKind, ForkTaskOutcome, RuntimeError,
14 RuntimeShared, TaskKey, TestContext, TestRuntime,
15};
16
17impl<R> TestRuntime<R> {
18 pub fn new() -> Self {
19 let (queries_tx, queries_rx) = mpsc::unbounded_channel();
20 let shared = Arc::new(Mutex::new(RuntimeShared {
21 queries_rx,
22 tasks: Default::default(),
23 entries: Default::default(),
24 }));
25 Self { queries_tx, shared }
26 }
27
28 pub fn new_context(
29 &self,
30 task_key: TaskKey,
31 address_lease: Option<AddressLease>,
32 ) -> TestContext<R> {
33 let queries_tx = self.queries_tx.clone();
34 TestContext {
35 task_key,
36 address_lease,
37 queries_tx,
38 }
39 }
40
41 pub async fn add_actor<A>(
42 &self,
43 actor_address: Address,
44 address_lease: Option<AddressLease>,
45 actor: A,
46 ) -> Result<&Self, RuntimeError>
47 where
48 R: Send + 'static,
49 A: for<'a> ActorFn<'a, R>,
50 {
51 let task_key = TaskKey::actor(actor_address);
52 let context = self.new_context(task_key, None);
53 let fut = async move {
54 let mut context = context;
55 let actor_exit = actor.run(&mut context).await;
56 match actor_exit.exit(&mut context).await {}
57 };
58 self.add_task(task_key, address_lease, fut).await?;
59 Ok(self)
60 }
61
62 pub async fn with_actor<A>(
63 self,
64 actor_address: Address,
65 address_lease: Option<AddressLease>,
66 actor: A,
67 ) -> Result<Self, RuntimeError>
68 where
69 R: Send + 'static,
70 A: for<'a> ActorFn<'a, R>,
71 {
72 self.add_actor(actor_address, address_lease, actor).await?;
73 Ok(self)
74 }
75
76 pub async fn add_task<F>(
77 &self,
78 task_key: TaskKey,
79 address_lease: Option<AddressLease>,
80 fut: F,
81 ) -> Result<(), RuntimeError>
82 where
83 F: Future<Output = ()> + Send + 'static,
84 {
85 self.shared
86 .lock()
87 .await
88 .add_task(task_key, address_lease, fut)
89 .await
90 }
91
92 pub async fn expect_next_event(&self) -> Event<R, EventKind<R>> {
93 self.next_event()
94 .await
95 .expect("error fetching next_event")
96 .expect("no more events")
97 }
98
99 pub async fn next_event(&self) -> Result<Option<Event<R, EventKind<R>>>, RuntimeError> {
100 let runtime = self.shared.clone();
101 let mut shared = self.shared.lock().await;
102 let RuntimeShared {
103 queries_rx, tasks, ..
104 } = &mut *shared;
105
106 if tasks.is_empty() {
107 let event_opt = queries_rx.try_recv().ok().map(|query| {
108 Event {
109 runtime,
110 kind: EventKind::Query(query),
111 }
112 });
113 drop(shared);
114 return Ok(event_opt)
115 }
116
117 let next_query = queries_rx.recv();
118 let task_done = tasks.next();
119
120 let event_opt = tokio::select! {
121 done = task_done => {
122 if let Some(done) = done {
123 let event = Event {runtime, kind: EventKind::Done(done)};
124 Some(event)
125 } else {
126 None
127 }},
128 query = next_query => {
129 if let Some(query) = query {
130 let event = Event {runtime, kind: EventKind::Query(query)};
131 Some(event)
132 } else {
133 None
134 }},
135 };
136
137 Ok(event_opt)
138 }
139}
140
141impl<R> Default for TestRuntime<R> {
142 fn default() -> Self {
143 Self::new()
144 }
145}
146
147impl<R> RuntimeShared<R> {
148 pub(crate) async fn add_task<F>(
149 &mut self,
150 task_key: TaskKey,
151 address_lease: Option<AddressLease>,
152 fut: F,
153 ) -> Result<(), RuntimeError>
154 where
155 F: Future<Output = ()> + Send + 'static,
156 {
157 match (
158 self.entries.entry(task_key.actor),
159 task_key.actor == task_key.context,
160 ) {
161 (Vacant(_), false) => Err(RuntimeError::NoActor(task_key.actor)),
162 (Vacant(vacant), true) => {
163 log::trace!("adding main actor [key: {:?}]", task_key);
164
165 let entry = ActorEntry {
166 forks: Default::default(),
167 actor_canceled: Default::default(),
168 actor_done: Default::default(),
169 };
170 let actor_canceled = entry.actor_canceled.clone();
171 let actor_done = entry.actor_done.clone();
172 let fut = async move {
173 tokio::select! {
174 _actor_canceled = actor_canceled.notified() => log::trace!("actor canceled [key: {:?}]", task_key),
175 _actor_completed = fut => log::trace!("actor completed [key: {:?}]", task_key),
176 }
177 actor_done.notify_waiters();
178 ActorTaskOutcome::Main(MainActorOutcome {
179 address: task_key.actor,
180 address_lease,
181 })
182 }
183 .boxed();
184
185 vacant.insert(entry);
186 self.tasks.push(fut);
187
188 Ok(())
189 },
190 (Occupied(_), true) => Err(RuntimeError::Duplicate(task_key.actor)),
191 (Occupied(occupied), false) => {
192 log::trace!("adding fork task [key: {:?}]", task_key);
193 let entry = occupied.into_mut();
194 let true = entry.forks.insert(task_key.context) else {
195 return Err(RuntimeError::Duplicate(task_key.context))
196 };
197 let actor_done = entry.actor_done.clone();
198 let fut = async move {
199 tokio::select! {
200 _fork_done = fut => log::trace!("fork task done [key: {:?}]", task_key),
201 _actor_done = actor_done.notified() => log::trace!("main actor notified completion [key: {:?}]", task_key),
202 };
203 ActorTaskOutcome::Fork(ForkTaskOutcome { task_key , address_lease, })
204 }
205 .boxed();
206 self.tasks.push(fut);
207
208 Ok(())
209 },
210 }
211 }
212}