1use std::sync::{
5 Arc,
6 atomic::{AtomicBool, Ordering},
7};
8
9use reifydb_value::value::duration::Duration;
10
11#[cfg(not(reifydb_single_threaded))]
12use crate::actor::timers::Repeat;
13#[cfg(reifydb_dst)]
14use crate::actor::timers::dst as dst_timers;
15#[cfg(reifydb_target = "wasi")]
16use crate::actor::timers::wasi::{schedule_once_fn, schedule_repeat, schedule_repeat_fn};
17#[cfg(reifydb_target = "wasm")]
18use crate::actor::timers::wasm::{schedule_once_fn, schedule_repeat, schedule_repeat_fn};
19use crate::actor::{mailbox::ActorRef, system::ActorSystem, timers::TimerHandle};
20
21#[derive(Clone)]
22pub struct CancellationToken {
23 cancelled: Arc<AtomicBool>,
24 parent: Option<Arc<AtomicBool>>,
25}
26
27impl CancellationToken {
28 pub fn new() -> Self {
29 Self {
30 cancelled: Arc::new(AtomicBool::new(false)),
31 parent: None,
32 }
33 }
34
35 pub fn child_token(&self) -> Self {
36 Self {
37 cancelled: Arc::new(AtomicBool::new(false)),
38 parent: Some(Arc::clone(&self.cancelled)),
39 }
40 }
41
42 pub fn cancel(&self) {
43 self.cancelled.store(true, Ordering::SeqCst);
44 }
45
46 pub fn is_cancelled(&self) -> bool {
47 self.cancelled.load(Ordering::SeqCst) || self.parent.as_ref().is_some_and(|p| p.load(Ordering::SeqCst))
48 }
49}
50
51impl Default for CancellationToken {
52 fn default() -> Self {
53 Self::new()
54 }
55}
56
57pub struct Context<M> {
58 self_ref: ActorRef<M>,
59 system: ActorSystem,
60 cancel: CancellationToken,
61}
62
63impl<M: Send + 'static> Context<M> {
64 pub fn new(self_ref: ActorRef<M>, system: ActorSystem, cancel: CancellationToken) -> Self {
65 Self {
66 self_ref,
67 system,
68 cancel,
69 }
70 }
71
72 pub fn self_ref(&self) -> ActorRef<M> {
73 self.self_ref.clone()
74 }
75
76 pub fn system(&self) -> &ActorSystem {
77 &self.system
78 }
79
80 pub fn is_cancelled(&self) -> bool {
81 self.cancel.is_cancelled()
82 }
83
84 pub fn cancellation_token(&self) -> CancellationToken {
85 self.cancel.clone()
86 }
87}
88
89impl<M: Send + 'static> Context<M> {
90 #[cfg(not(reifydb_single_threaded))]
91 pub fn schedule_once<F: FnOnce() -> M + Send + 'static>(
92 &self,
93 delay: impl Into<Duration>,
94 factory: F,
95 ) -> TimerHandle {
96 let delay = delay.into().to_std();
97 let actor_ref = self.self_ref.clone();
98 self.system.scheduler().schedule_once(delay, move || {
99 let _ = actor_ref.send(factory());
100 })
101 }
102
103 #[cfg(all(reifydb_single_threaded, not(reifydb_dst)))]
104 pub fn schedule_once<F: FnOnce() -> M + Send + 'static>(
105 &self,
106 delay: impl Into<Duration>,
107 factory: F,
108 ) -> TimerHandle {
109 let delay = delay.into().to_std();
110 schedule_once_fn(self.self_ref.clone(), delay, factory)
111 }
112
113 #[cfg(reifydb_dst)]
114 pub fn schedule_once<F: FnOnce() -> M + Send + 'static>(
115 &self,
116 delay: impl Into<Duration>,
117 factory: F,
118 ) -> TimerHandle {
119 let delay = delay.into().to_std();
120 dst_timers::schedule_once_fn(
121 self.system.timer_heap(),
122 self.system.mock_clock(),
123 self.self_ref.clone(),
124 delay,
125 factory,
126 )
127 }
128}
129
130impl<M: Send + Sync + Clone + 'static> Context<M> {
131 #[cfg(not(reifydb_single_threaded))]
132 pub fn schedule_repeat(&self, interval: impl Into<Duration>, msg: M) -> TimerHandle {
133 let interval = interval.into().to_std();
134 let actor_ref = self.self_ref.clone();
135 self.system
136 .scheduler()
137 .schedule_repeat(interval, move || Repeat::after_send(actor_ref.send(msg.clone())))
138 }
139
140 #[cfg(all(reifydb_single_threaded, not(reifydb_dst)))]
141 pub fn schedule_repeat(&self, interval: impl Into<Duration>, msg: M) -> TimerHandle {
142 let interval = interval.into().to_std();
143 schedule_repeat(self.self_ref.clone(), interval, msg)
144 }
145
146 #[cfg(reifydb_dst)]
147 pub fn schedule_repeat(&self, interval: impl Into<Duration>, msg: M) -> TimerHandle {
148 let interval = interval.into().to_std();
149 dst_timers::schedule_repeat(
150 self.system.timer_heap(),
151 self.system.mock_clock(),
152 self.self_ref.clone(),
153 interval,
154 msg,
155 )
156 }
157}
158
159impl<M: Send + 'static> Context<M> {
160 #[cfg(not(reifydb_single_threaded))]
161 pub fn schedule_repeat_fn<F: Fn() -> M + Send + Sync + 'static>(
162 &self,
163 interval: impl Into<Duration>,
164 factory: F,
165 ) -> TimerHandle {
166 let interval = interval.into().to_std();
167 let actor_ref = self.self_ref.clone();
168 self.system.scheduler().schedule_repeat(interval, move || Repeat::after_send(actor_ref.send(factory())))
169 }
170
171 #[cfg(all(reifydb_single_threaded, not(reifydb_dst)))]
172 pub fn schedule_repeat_fn<F: Fn() -> M + Send + Sync + 'static>(
173 &self,
174 interval: impl Into<Duration>,
175 factory: F,
176 ) -> TimerHandle {
177 let interval = interval.into().to_std();
178 schedule_repeat_fn(self.self_ref.clone(), interval, factory)
179 }
180
181 #[cfg(reifydb_dst)]
182 pub fn schedule_repeat_fn<F: Fn() -> M + Send + Sync + 'static>(
183 &self,
184 interval: impl Into<Duration>,
185 factory: F,
186 ) -> TimerHandle {
187 let interval = interval.into().to_std();
188 dst_timers::schedule_repeat_fn(
189 self.system.timer_heap(),
190 self.system.mock_clock(),
191 self.self_ref.clone(),
192 interval,
193 factory,
194 )
195 }
196
197 pub fn schedule_tick<F: Fn(u64) -> M + Send + Sync + 'static>(
198 &self,
199 interval: impl Into<Duration>,
200 factory: F,
201 ) -> TimerHandle {
202 let interval = interval.into().to_std();
203 let actor_ref = self.self_ref.clone();
204 let clock = self.system.clock().clone();
205
206 #[cfg(not(reifydb_single_threaded))]
207 {
208 self.system.scheduler().schedule_repeat(interval, move || {
209 let now = clock.now().to_nanos();
210 Repeat::after_send(actor_ref.send(factory(now)))
211 })
212 }
213
214 #[cfg(all(reifydb_single_threaded, not(reifydb_dst)))]
215 {
216 schedule_repeat_fn(actor_ref, interval, move || {
217 let now = clock.now().to_nanos();
218 factory(now)
219 })
220 }
221
222 #[cfg(reifydb_dst)]
223 {
224 dst_timers::schedule_repeat_fn(
225 self.system.timer_heap(),
226 self.system.mock_clock(),
227 actor_ref,
228 interval,
229 move || {
230 let now = clock.now().to_nanos();
231 factory(now)
232 },
233 )
234 }
235 }
236}
237
238impl<M> Clone for Context<M> {
239 fn clone(&self) -> Self {
240 Self {
241 self_ref: self.self_ref.clone(),
242 system: self.system.clone(),
243 cancel: self.cancel.clone(),
244 }
245 }
246}