1use std::collections::HashMap;
8use std::future::Future;
9use std::pin::Pin;
10use std::sync::{Arc, Mutex};
11use std::time::{Duration, SystemTime, UNIX_EPOCH};
12
13use blvm_node::module::ipc::protocol::{InvocationResultMessage, InvocationResultPayload};
14use blvm_node::module::timers::types::{
15 AckMode, Clock, CreateMode, First, JitterShape, Missed, Overlap, Pacing, Persist, Randomness,
16 Repeat, TimerAck, TimerFire, TimerId, TimerSpec, TimerTarget, monotonic_interval,
17};
18use blvm_node::module::traits::{ModuleError, NodeAPI};
19use tokio::sync::{Semaphore, mpsc, oneshot};
20
21pub use blvm_node::module::timers::types::TimerInfo;
22
23type FireFut = Pin<Box<dyn Future<Output = Result<TimerAck, String>> + Send>>;
24type FireFn = Arc<dyn Fn(TimerFire) -> FireFut + Send + Sync>;
25
26#[derive(Clone, Debug, PartialEq, Eq)]
28pub enum Exec {
29 Serial,
31 Task,
33 Blocking,
35 Dedicated(String),
37}
38
39struct Slot {
40 exec: Exec,
41 ack: AckMode,
42 run: FireFn,
43}
44
45pub struct TimerRegistry {
47 slots: Mutex<HashMap<TimerId, Slot>>,
48 blocking: Arc<Semaphore>,
49 dedicated: Mutex<HashMap<String, mpsc::UnboundedSender<DedicatedJob>>>,
50}
51
52struct DedicatedJob {
53 fire: TimerFire,
54 run: FireFn,
55 done: oneshot::Sender<Result<TimerAck, String>>,
56}
57
58impl TimerRegistry {
59 pub fn new() -> Arc<Self> {
60 Arc::new(Self {
61 slots: Mutex::new(HashMap::new()),
62 blocking: Arc::new(Semaphore::new(4)),
63 dedicated: Mutex::new(HashMap::new()),
64 })
65 }
66
67 pub fn register(&self, id: TimerId, exec: Exec, ack: AckMode, run: FireFn) {
68 self.slots
69 .lock()
70 .unwrap_or_else(|err| err.into_inner())
71 .insert(id, Slot { exec, ack, run });
72 }
73
74 pub fn is_detached(&self, id: &TimerId) -> bool {
75 self.slots
76 .lock()
77 .unwrap_or_else(|err| err.into_inner())
78 .get(id)
79 .is_some_and(|slot| !matches!(slot.exec, Exec::Serial))
80 }
81
82 pub fn start_detached(
84 self: &Arc<Self>,
85 fire: TimerFire,
86 correlation_id: u64,
87 result_tx: oneshot::Sender<InvocationResultMessage>,
88 ) {
89 let Some(slot) = self.slot(&fire.timer_id) else {
90 let _ = result_tx.send(timer_failure(correlation_id, "unknown timer"));
91 return;
92 };
93 if slot.ack == AckMode::Receipt {
94 let _ = result_tx.send(timer_ok(correlation_id, TimerAck::Continue));
95 self.launch(slot, fire, None);
96 return;
97 }
98 self.launch(slot, fire, Some((correlation_id, result_tx)));
99 }
100
101 pub async fn run_serial(
104 &self,
105 correlation_id: u64,
106 fire: &TimerFire,
107 result_tx: oneshot::Sender<InvocationResultMessage>,
108 ) {
109 let Some(slot) = self.slot(&fire.timer_id) else {
110 let _ = result_tx.send(timer_failure(correlation_id, "unknown timer"));
111 return;
112 };
113 if slot.ack == AckMode::Receipt {
114 let _ = result_tx.send(timer_ok(correlation_id, TimerAck::Continue));
115 let _ = (slot.run)(fire.clone()).await;
116 return;
117 }
118 let message = match (slot.run)(fire.clone()).await {
119 Ok(ack) => timer_ok(correlation_id, ack),
120 Err(err) => timer_failure(correlation_id, &err),
121 };
122 let _ = result_tx.send(message);
123 }
124
125 fn slot(&self, id: &TimerId) -> Option<Slot> {
126 let slots = self.slots.lock().unwrap_or_else(|err| err.into_inner());
127 let slot = slots.get(id)?;
128 Some(Slot {
129 exec: slot.exec.clone(),
130 ack: slot.ack,
131 run: Arc::clone(&slot.run),
132 })
133 }
134
135 fn launch(
136 self: &Arc<Self>,
137 slot: Slot,
138 fire: TimerFire,
139 reply: Option<(u64, oneshot::Sender<InvocationResultMessage>)>,
140 ) {
141 let send = |result: Result<TimerAck, String>| {
142 if let Some((correlation_id, tx)) = reply {
143 let message = match result {
144 Ok(ack) => timer_ok(correlation_id, ack),
145 Err(err) => timer_failure(correlation_id, &err),
146 };
147 let _ = tx.send(message);
148 }
149 };
150 match slot.exec {
151 Exec::Serial => {
152 let run = slot.run;
153 tokio::spawn(async move {
154 send(run(fire).await);
155 });
156 }
157 Exec::Task => {
158 let run = slot.run;
159 tokio::spawn(async move {
160 let joined = tokio::spawn(run(fire)).await;
161 send(flatten_join(joined));
162 });
163 }
164 Exec::Blocking => {
165 let run = slot.run;
166 let blocking = Arc::clone(&self.blocking);
167 tokio::spawn(async move {
168 let Ok(_permit) = blocking.acquire().await else {
169 send(Err("blocking pool closed".into()));
170 return;
171 };
172 let joined = tokio::task::spawn_blocking(move || {
173 let runtime = match tokio::runtime::Builder::new_current_thread()
174 .enable_all()
175 .build()
176 {
177 Ok(runtime) => runtime,
178 Err(err) => return Err(err.to_string()),
179 };
180 runtime.block_on(run(fire))
181 })
182 .await;
183 send(flatten_join(joined));
184 });
185 }
186 Exec::Dedicated(name) => {
187 let run = slot.run;
188 let (done_tx, done_rx) = oneshot::channel();
189 if let Err(err) = self.enqueue_dedicated(name, fire, run, done_tx) {
190 send(Err(err));
191 return;
192 }
193 tokio::spawn(async move {
194 match done_rx.await {
195 Ok(result) => send(result),
196 Err(_) => send(Err("dedicated thread dropped the fire".into())),
197 }
198 });
199 }
200 }
201 }
202
203 fn enqueue_dedicated(
204 &self,
205 name: String,
206 fire: TimerFire,
207 run: FireFn,
208 done: oneshot::Sender<Result<TimerAck, String>>,
209 ) -> Result<(), String> {
210 let mut threads = self.dedicated.lock().unwrap_or_else(|err| err.into_inner());
211 if let Some(tx) = threads.get(&name) {
212 return tx
213 .send(DedicatedJob { fire, run, done })
214 .map_err(|_| "dedicated thread stopped".to_string());
215 }
216 if threads.len() >= 4 {
217 return Err("dedicated thread limit is 4".into());
218 }
219 let (tx, mut rx) = mpsc::unbounded_channel();
220 tx.send(DedicatedJob { fire, run, done })
221 .map_err(|_| "dedicated thread stopped".to_string())?;
222 threads.insert(name.clone(), tx);
223 std::thread::Builder::new()
224 .name(format!("blvm-timer-{name}"))
225 .spawn(move || {
226 let runtime = match tokio::runtime::Builder::new_current_thread()
227 .enable_all()
228 .build()
229 {
230 Ok(runtime) => runtime,
231 Err(_) => return,
232 };
233 let local = tokio::task::LocalSet::new();
234 local.block_on(&runtime, async move {
235 while let Some(job) = rx.recv().await {
236 let DedicatedJob { fire, run, done } = job;
237 let outcome = match tokio::task::spawn_local(run(fire)).await {
238 Ok(result) => result,
239 Err(err) => Err(err.to_string()),
240 };
241 let _ = done.send(outcome);
242 }
243 });
244 })
245 .map_err(|err| err.to_string())?;
246 Ok(())
247 }
248}
249
250fn flatten_join(
251 joined: Result<Result<TimerAck, String>, tokio::task::JoinError>,
252) -> Result<TimerAck, String> {
253 match joined {
254 Ok(result) => result,
255 Err(err) => Err(err.to_string()),
256 }
257}
258
259pub fn timer_ok(correlation_id: u64, ack: TimerAck) -> InvocationResultMessage {
260 InvocationResultMessage {
261 correlation_id,
262 success: true,
263 payload: Some(InvocationResultPayload::Timer(ack)),
264 error: None,
265 }
266}
267
268pub fn timer_failure(correlation_id: u64, err: &str) -> InvocationResultMessage {
269 InvocationResultMessage {
270 correlation_id,
271 success: false,
272 payload: None,
273 error: Some(err.to_string()),
274 }
275}
276
277pub fn sdk_tick_spec() -> TimerSpec {
281 let mut spec = monotonic_interval(Duration::from_secs(30), Pacing::FixedRate);
282 spec.first = First::Delay(Duration::from_secs(30));
283 spec.overlap = Overlap::Skip;
284 spec.key = Some("sdk.tick".into());
285 spec.clock = Clock::Monotonic;
286 spec
287}
288
289fn next_daily_utc_ms(hour: u32, minute: u32) -> u64 {
290 let now_ms = SystemTime::now()
291 .duration_since(UNIX_EPOCH)
292 .unwrap_or_default()
293 .as_millis() as u64;
294 let day = 86_400_000u64;
295 let today = now_ms - (now_ms % day);
296 let at = (hour as u64) * 3_600_000 + (minute as u64) * 60_000;
297 let target = today + at;
298 if target > now_ms {
299 target
300 } else {
301 target + day
302 }
303}
304
305pub struct TimerBuilder {
307 spec: TimerSpec,
308 exec: Exec,
309 api: Option<Arc<dyn NodeAPI>>,
310 registry: Option<Arc<TimerRegistry>>,
311 chain_depth: u32,
312 exp_floor: Option<Duration>,
313 exp_cap: Option<Duration>,
314}
315
316impl TimerBuilder {
317 pub fn new() -> Self {
318 Self {
319 spec: monotonic_interval(Duration::from_secs(10), Pacing::FixedRate),
320 exec: Exec::Serial,
321 api: None,
322 registry: None,
323 chain_depth: 0,
324 exp_floor: None,
325 exp_cap: None,
326 }
327 }
328
329 pub fn with_runtime(api: Arc<dyn NodeAPI>, registry: Arc<TimerRegistry>) -> Self {
330 Self {
331 api: Some(api),
332 registry: Some(registry),
333 ..Self::new()
334 }
335 }
336
337 pub fn every(mut self, period: Duration) -> Self {
338 if let blvm_node::module::timers::types::Repeat::Interval {
339 period: ref mut slot,
340 ..
341 } = self.spec.repeat
342 {
343 *slot = period;
344 }
345 self
346 }
347
348 pub fn once(mut self, delay: Duration) -> Self {
349 self.spec.first = First::Delay(delay);
350 self.spec.repeat = blvm_node::module::timers::types::Repeat::Once;
351 self
352 }
353
354 pub fn fixed_rate(mut self) -> Self {
355 if let blvm_node::module::timers::types::Repeat::Interval { ref mut pacing, .. } =
356 self.spec.repeat
357 {
358 *pacing = Pacing::FixedRate;
359 }
360 self
361 }
362
363 pub fn on_overlap(mut self, overlap: Overlap) -> Self {
364 self.spec.overlap = overlap;
365 self
366 }
367
368 pub fn key(mut self, key: impl Into<String>) -> Self {
369 self.spec.key = Some(key.into());
370 self
371 }
372
373 pub fn exec(mut self, exec: Exec) -> Self {
374 self.exec = exec;
375 self
376 }
377
378 pub fn jitter(mut self, max: Duration, shape: JitterShape) -> Self {
379 self.spec.randomness = match self.spec.randomness {
380 Randomness::Spread { max: spread_max }
381 | Randomness::JitterAndSpread { spread_max, .. } => Randomness::JitterAndSpread {
382 jitter_max: max,
383 shape,
384 spread_max,
385 },
386 _ => Randomness::Jitter { max, shape },
387 };
388 self
389 }
390
391 pub fn spread(mut self, max: Duration) -> Self {
392 self.spec.randomness = match self.spec.randomness {
393 Randomness::Jitter {
394 max: jitter_max,
395 shape,
396 }
397 | Randomness::JitterAndSpread {
398 jitter_max, shape, ..
399 } => Randomness::JitterAndSpread {
400 jitter_max,
401 shape,
402 spread_max: max,
403 },
404 _ => Randomness::Spread { max },
405 };
406 self
407 }
408
409 pub fn max_fires(mut self, n: u32) -> Self {
410 match &mut self.spec.repeat {
411 Repeat::Interval { max_fires, .. }
412 | Repeat::Blocks { max_fires, .. }
413 | Repeat::Cron { max_fires, .. } => *max_fires = Some(n),
414 Repeat::Once => {}
415 }
416 self
417 }
418
419 pub fn missed(mut self, missed: Missed) -> Self {
420 self.spec.missed = missed;
421 self
422 }
423
424 pub fn persist_durable(mut self, key: impl Into<String>) -> Self {
425 self.spec.key = Some(key.into());
426 self.spec.persist = Persist::Durable { audit: false };
427 self
428 }
429
430 pub fn daily_at_utc(mut self, hour: u32, minute: u32) -> Self {
432 self.spec.clock = Clock::Wall;
433 self.spec.first = First::AtUnixMs(next_daily_utc_ms(hour, minute));
434 self.spec.repeat = Repeat::Interval {
435 period: Duration::from_secs(24 * 60 * 60),
436 pacing: Pacing::FixedRate,
437 max_fires: None,
438 until: None,
439 };
440 self
441 }
442
443 pub fn after_blocks(mut self, blocks: u32) -> Self {
444 self.spec.clock = Clock::Chain;
445 self.spec.first = First::AfterBlocks {
446 blocks,
447 depth: self.chain_depth,
448 };
449 self.spec.repeat = Repeat::Once;
450 self
451 }
452
453 pub fn depth(mut self, depth: u32) -> Self {
454 self.chain_depth = depth;
455 match &mut self.spec.first {
456 First::AfterBlocks { depth: slot, .. } | First::AtHeight { depth: slot, .. } => {
457 *slot = depth;
458 }
459 _ => {}
460 }
461 self
462 }
463
464 pub fn every_blocks(mut self, every: u32) -> Self {
465 self.spec.clock = Clock::Chain;
466 self.spec.first = First::AfterBlocks {
467 blocks: 0,
468 depth: self.chain_depth,
469 };
470 self.spec.repeat = Repeat::Blocks {
471 every,
472 max_fires: None,
473 until: None,
474 };
475 self
476 }
477
478 pub fn poisson(mut self, mean: Duration) -> Self {
479 let floor = self.exp_floor.unwrap_or(Duration::from_millis(1));
480 let cap = self.exp_cap.unwrap_or(mean);
481 if let Repeat::Interval {
482 ref mut period,
483 ref mut pacing,
484 ..
485 } = self.spec.repeat
486 {
487 *period = mean;
488 *pacing = Pacing::FixedRate;
489 }
490 self.spec.randomness = Randomness::Exponential { mean, floor, cap };
491 self
492 }
493
494 pub fn bounds(mut self, floor: Duration, cap: Duration) -> Self {
495 self.exp_floor = Some(floor);
496 self.exp_cap = Some(cap);
497 if let Randomness::Exponential {
498 floor: ref mut slot_floor,
499 cap: ref mut slot_cap,
500 ..
501 } = self.spec.randomness
502 {
503 *slot_floor = floor;
504 *slot_cap = cap;
505 }
506 self
507 }
508
509 pub fn spec(&self) -> TimerSpec {
510 self.spec.clone()
511 }
512
513 pub async fn spawn<F, Fut>(self, handler: F) -> Result<TimerHandle, ModuleError>
514 where
515 F: Fn(TimerFire) -> Fut + Send + Sync + 'static,
516 Fut: Future<Output = Result<TimerAck, String>> + Send + 'static,
517 {
518 let api = self
519 .api
520 .clone()
521 .ok_or_else(|| ModuleError::OperationError("timer builder has no node api".into()))?;
522 let registry = self
523 .registry
524 .clone()
525 .ok_or_else(|| ModuleError::OperationError("timer builder has no registry".into()))?;
526 let mode = if self.spec.key.is_some() {
527 CreateMode::Ensure
528 } else {
529 CreateMode::CreateOnly
530 };
531 let info = api.timer_create(self.spec.clone(), mode).await?;
532 let run: FireFn = Arc::new(move |fire| Box::pin(handler(fire)));
533 registry.register(info.timer_id, self.exec, self.spec.delivery.ack, run);
534 Ok(TimerHandle {
535 id: info.timer_id,
536 api,
537 })
538 }
539}
540
541impl Default for TimerBuilder {
542 fn default() -> Self {
543 Self::new()
544 }
545}
546
547pub struct TimerHandle {
549 id: TimerId,
550 api: Arc<dyn NodeAPI>,
551}
552
553impl TimerHandle {
554 pub async fn cancel(&self) -> Result<(), ModuleError> {
555 self.api.timer_cancel(TimerTarget::Id(self.id)).await
556 }
557
558 pub async fn pause(&self) -> Result<(), ModuleError> {
559 self.api.timer_pause(TimerTarget::Id(self.id)).await
560 }
561
562 pub async fn resume(&self) -> Result<(), ModuleError> {
563 self.api.timer_resume(TimerTarget::Id(self.id)).await
564 }
565
566 pub async fn trigger(&self) -> Result<(), ModuleError> {
567 self.api.timer_trigger(TimerTarget::Id(self.id)).await
568 }
569
570 pub async fn info(&self) -> Result<TimerInfo, ModuleError> {
571 self.api.timer_get(TimerTarget::Id(self.id)).await
572 }
573
574 pub fn id(&self) -> TimerId {
575 self.id
576 }
577}
578
579#[cfg(test)]
580mod tests {
581 use super::*;
582 use std::sync::atomic::{AtomicUsize, Ordering};
583 use std::time::Duration;
584
585 fn fire(id: TimerId) -> TimerFire {
586 TimerFire {
587 timer_id: id,
588 key: None,
589 payload: Vec::new(),
590 seq: 1,
591 attempt: 1,
592 clock: Clock::Monotonic,
593 scheduled_unix_ms: 0,
594 planned_unix_ms: 0,
595 fired_unix_ms: 0,
596 lateness_ms: 0,
597 missed: 0,
598 coalesced: 0,
599 skipped_overlap: 0,
600 clock_step_ms: 0,
601 tip_height: None,
602 tip_hash: None,
603 tip_mtp: None,
604 is_last: false,
605 }
606 }
607
608 #[test]
609 fn the_tick_spec_is_a_fixed_rate_skip_timer() {
610 let spec = sdk_tick_spec();
611 assert_eq!(spec.clock, Clock::Monotonic);
612 assert_eq!(spec.overlap, Overlap::Skip);
613 assert_eq!(spec.key.as_deref(), Some("sdk.tick"));
614 assert_eq!(spec.first, First::Delay(Duration::from_secs(30)));
615 }
616
617 #[test]
618 fn the_documented_builders_match_timer_validation() {
619 use blvm_node::module::timers::types::{TimerConfig, ValidateCtx, validate_spec};
620
621 let ctx = ValidateCtx {
622 now_unix_ms: 0,
623 tip_height: Some(100),
624 };
625 let limits = TimerConfig::default();
626 let poll = TimerBuilder::new()
627 .every(Duration::from_secs(10))
628 .fixed_rate()
629 .jitter(Duration::from_secs(2), JitterShape::Symmetric)
630 .spread(Duration::from_secs(10))
631 .max_fires(8640)
632 .on_overlap(Overlap::Skip)
633 .key("upstream.poll")
634 .exec(Exec::Task)
635 .spec();
636 validate_spec(&poll, &limits, &ctx).unwrap();
637 assert!(matches!(
638 poll.randomness,
639 Randomness::JitterAndSpread { .. }
640 ));
641
642 let daily = TimerBuilder::new()
643 .daily_at_utc(3, 0)
644 .persist_durable("housekeeping.compact")
645 .missed(Missed::FireOnce)
646 .exec(Exec::Blocking)
647 .spec();
648 validate_spec(&daily, &limits, &ctx).unwrap();
649 assert_eq!(daily.clock, Clock::Wall);
650 assert_eq!(daily.missed, Missed::FireOnce);
651 assert!(matches!(daily.persist, Persist::Durable { audit: false }));
652
653 let after = TimerBuilder::new().after_blocks(6).depth(0).spec();
654 validate_spec(&after, &limits, &ctx).unwrap();
655 assert_eq!(
656 after.first,
657 First::AfterBlocks {
658 blocks: 6,
659 depth: 0
660 }
661 );
662
663 let every = TimerBuilder::new()
664 .every_blocks(1)
665 .on_overlap(Overlap::Skip)
666 .spec();
667 validate_spec(&every, &limits, &ctx).unwrap();
668
669 let poisson = TimerBuilder::new()
670 .poisson(Duration::from_secs(5))
671 .bounds(Duration::from_millis(500), Duration::from_secs(60))
672 .spec();
673 validate_spec(&poisson, &limits, &ctx).unwrap();
674 }
675
676 #[tokio::test]
677 async fn blocking_handler_sleeps_off_the_caller() {
678 let registry = TimerRegistry::new();
679 let id = TimerId([8; 16]);
680 registry.register(
681 id,
682 Exec::Blocking,
683 AckMode::Completion,
684 Arc::new(|_fire| {
685 Box::pin(async move {
686 tokio::time::sleep(Duration::from_millis(30)).await;
687 Ok(TimerAck::Continue)
688 })
689 }),
690 );
691 let (tx, rx) = oneshot::channel();
692 let started = std::time::Instant::now();
693 registry.start_detached(fire(id), 3, tx);
694 let message = rx.await.unwrap();
695 assert!(message.success);
696 assert!(started.elapsed() >= Duration::from_millis(30));
697 }
698
699 #[tokio::test]
700 async fn a_dedicated_thread_runs_during_a_serial_handler() {
701 let registry = TimerRegistry::new();
702 let serial_id = TimerId([3; 16]);
703 let dedicated_id = TimerId([4; 16]);
704 let release = Arc::new(tokio::sync::Notify::new());
705 let saw = Arc::new(AtomicUsize::new(0));
706 let release_h = Arc::clone(&release);
707 registry.register(
708 serial_id,
709 Exec::Serial,
710 AckMode::Completion,
711 Arc::new(move |_fire| {
712 let release_h = Arc::clone(&release_h);
713 Box::pin(async move {
714 release_h.notified().await;
715 Ok(TimerAck::Continue)
716 })
717 }),
718 );
719 let saw_d = Arc::clone(&saw);
720 registry.register(
721 dedicated_id,
722 Exec::Dedicated("poll".into()),
723 AckMode::Completion,
724 Arc::new(move |_fire| {
725 let saw_d = Arc::clone(&saw_d);
726 Box::pin(async move {
727 saw_d.fetch_add(1, Ordering::SeqCst);
728 Ok(TimerAck::Continue)
729 })
730 }),
731 );
732 let run = registry.slot(&serial_id).unwrap();
733 let serial = tokio::spawn(async move {
734 let _ = (run.run)(fire(serial_id)).await;
735 });
736 tokio::time::sleep(Duration::from_millis(20)).await;
737 let (tx, rx) = oneshot::channel();
738 registry.start_detached(fire(dedicated_id), 1, tx);
739 tokio::time::timeout(Duration::from_secs(2), async {
740 while saw.load(Ordering::SeqCst) == 0 {
741 tokio::task::yield_now().await;
742 }
743 })
744 .await
745 .expect("dedicated handler waited for the serial handler");
746 release.notify_one();
747 serial.await.unwrap();
748 let message = rx.await.unwrap();
749 assert!(message.success);
750 }
751}