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 #[cfg(not(reifydb_single_threaded))]
159 pub fn schedule_repeat_fn<F: Fn() -> M + Send + Sync + 'static>(
160 &self,
161 interval: impl Into<Duration>,
162 factory: F,
163 ) -> TimerHandle {
164 let interval = interval.into().to_std();
165 let actor_ref = self.self_ref.clone();
166 self.system.scheduler().schedule_repeat(interval, move || Repeat::after_send(actor_ref.send(factory())))
167 }
168
169 #[cfg(all(reifydb_single_threaded, not(reifydb_dst)))]
170 pub fn schedule_repeat_fn<F: Fn() -> M + Send + Sync + 'static>(
171 &self,
172 interval: impl Into<Duration>,
173 factory: F,
174 ) -> TimerHandle {
175 let interval = interval.into().to_std();
176 schedule_repeat_fn(self.self_ref.clone(), interval, factory)
177 }
178
179 #[cfg(reifydb_dst)]
180 pub fn schedule_repeat_fn<F: Fn() -> M + Send + Sync + 'static>(
181 &self,
182 interval: impl Into<Duration>,
183 factory: F,
184 ) -> TimerHandle {
185 let interval = interval.into().to_std();
186 dst_timers::schedule_repeat_fn(
187 self.system.timer_heap(),
188 self.system.mock_clock(),
189 self.self_ref.clone(),
190 interval,
191 factory,
192 )
193 }
194
195 pub fn schedule_tick<F: Fn(u64) -> M + Send + Sync + 'static>(
196 &self,
197 interval: impl Into<Duration>,
198 factory: F,
199 ) -> TimerHandle {
200 let interval = interval.into().to_std();
201 let actor_ref = self.self_ref.clone();
202 let clock = self.system.clock().clone();
203
204 #[cfg(not(reifydb_single_threaded))]
205 {
206 self.system.scheduler().schedule_repeat(interval, move || {
207 let now = clock.now().to_nanos();
208 Repeat::after_send(actor_ref.send(factory(now)))
209 })
210 }
211
212 #[cfg(all(reifydb_single_threaded, not(reifydb_dst)))]
213 {
214 schedule_repeat_fn(actor_ref, interval, move || {
215 let now = clock.now().to_nanos();
216 factory(now)
217 })
218 }
219
220 #[cfg(reifydb_dst)]
221 {
222 dst_timers::schedule_repeat_fn(
223 self.system.timer_heap(),
224 self.system.mock_clock(),
225 actor_ref,
226 interval,
227 move || {
228 let now = clock.now().to_nanos();
229 factory(now)
230 },
231 )
232 }
233 }
234}
235
236impl<M> Clone for Context<M> {
237 fn clone(&self) -> Self {
238 Self {
239 self_ref: self.self_ref.clone(),
240 system: self.system.clone(),
241 cancel: self.cancel.clone(),
242 }
243 }
244}