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