1use crate::{
4 platform::{self, TimerHandle},
5 registry::{
6 CallbackAcceptance, CallbackRole, CallbackToken, OrdinaryCallback, ProviderHandle,
7 ProviderHandles, RegisterError, RegistrationClaim, RegistryEffect, RegistryError,
8 RegistryTransition, TimerRegistry, WatchdogCallback,
9 },
10 schedule::{ScheduleError, TimerCadence, TimerDirective, TimerSchedule},
11 snapshot::{
12 DeclarationLifetime, TimerCompletion, TimerControlFailure, TimerEpoch, TimerIdentity,
13 TimerInventorySnapshot, TimerPolicy, TimerRunResult, TimerSnapshot, WatchdogRunResult,
14 },
15};
16use std::{cell::RefCell, future::Future, rc::Rc, time::Duration};
17use thiserror::Error;
18
19thread_local! {
20 static RUNTIME: RefCell<Option<TimerRegistry>> = const { RefCell::new(None) };
21}
22
23#[non_exhaustive]
25#[derive(Debug, Error)]
26pub enum TimerError {
27 #[error("timer runtime is not initialized")]
29 NotInitialized,
30 #[error("timer runtime is already borrowed")]
32 RuntimeBusy,
33 #[error(transparent)]
35 Register(#[from] RegisterError),
36 #[error(transparent)]
38 Schedule(#[from] ScheduleError),
39 #[error("timer registration is no longer authoritative")]
41 RegistrationExpired,
42 #[error("timer control failed: {0:?}")]
44 ControlFailure(TimerControlFailure),
45 #[error("timer runtime ownership invariant failed")]
47 OwnershipInvariant,
48 #[error("timer lifecycle reconciliation conflicts with the canonical declaration")]
50 ReconciliationConflict,
51}
52
53impl From<RegistryError> for TimerError {
54 fn from(value: RegistryError) -> Self {
55 match value {
56 RegistryError::UnknownRegistration
57 | RegistryError::StaleRegistration
58 | RegistryError::StaleCallback => Self::RegistrationExpired,
59 RegistryError::PolicyMismatch { .. } | RegistryError::ProviderHandleAlreadyOwned => {
60 Self::OwnershipInvariant
61 }
62 RegistryError::Schedule(error) => Self::Schedule(error),
63 }
64 }
65}
66
67pub fn initialize_runtime() -> Result<TimerEpoch, TimerError> {
72 let epoch = TimerEpoch::new(platform::canister_version(), platform::time_ns());
73 RUNTIME.with(|runtime| {
74 let mut runtime = runtime
75 .try_borrow_mut()
76 .map_err(|_| TimerError::RuntimeBusy)?;
77 if let Some(registry) = runtime.as_ref() {
78 return Ok(registry.epoch());
79 }
80 *runtime = Some(TimerRegistry::new(epoch));
81 Ok(epoch)
82 })
83}
84
85struct CallbackContext {
86 token: CallbackToken,
87}
88
89impl CallbackContext {
90 const fn new(token: CallbackToken) -> Self {
91 Self { token }
92 }
93
94 fn claim(&self) -> RegistrationClaim {
95 RegistrationClaim::from_callback(&self.token)
96 }
97
98 const fn identity(&self) -> &TimerIdentity {
99 self.token.identity()
100 }
101
102 fn cancel(&self) -> Result<(), TimerError> {
103 cancel_claim(&self.claim(), Some(&self.token))
104 }
105
106 fn schedule_once(&self, schedule: TimerSchedule) -> Result<(), TimerError> {
107 ensure_once_claim(&self.claim(), Some(&self.token), schedule)
108 }
109
110 fn schedule_recurring(&self) -> Result<(), TimerError> {
111 ensure_recurring_claim(&self.claim(), Some(&self.token))
112 }
113
114 fn schedule_watchdog_immediately(&self) -> Result<(), TimerError> {
115 ensure_watchdog_immediately_claim(&self.claim(), Some(&self.token))
116 }
117
118 fn reconcile_watchdog(&self, schedule: Option<TimerSchedule>) -> Result<(), TimerError> {
119 reconcile_watchdog_claim(&self.claim(), Some(&self.token), schedule)
120 }
121
122 fn reconcile_ordinary(&self, schedule: Option<TimerSchedule>) -> Result<(), TimerError> {
123 reconcile_ordinary_claim(&self.claim(), Some(&self.token), schedule)
124 }
125}
126
127pub struct OnceContext {
133 inner: CallbackContext,
134}
135
136impl OnceContext {
137 const fn new(token: CallbackToken) -> Self {
138 Self {
139 inner: CallbackContext::new(token),
140 }
141 }
142
143 #[must_use]
145 pub const fn identity(&self) -> &TimerIdentity {
146 self.inner.identity()
147 }
148
149 pub fn ensure_scheduled(&self, schedule: TimerSchedule) -> Result<(), TimerError> {
154 self.inner.schedule_once(schedule)
155 }
156
157 pub fn reconcile_schedule(&self, schedule: Option<TimerSchedule>) -> Result<(), TimerError> {
167 self.inner.reconcile_ordinary(schedule)
168 }
169
170 pub fn cancel(&self) -> Result<(), TimerError> {
177 self.inner.cancel()
178 }
179}
180
181pub struct AfterCompletionContext {
187 inner: CallbackContext,
188}
189
190impl AfterCompletionContext {
191 const fn new(token: CallbackToken) -> Self {
192 Self {
193 inner: CallbackContext::new(token),
194 }
195 }
196
197 #[must_use]
199 pub const fn identity(&self) -> &TimerIdentity {
200 self.inner.identity()
201 }
202
203 pub fn ensure_scheduled(&self) -> Result<(), TimerError> {
205 self.inner.schedule_recurring()
206 }
207
208 pub fn reconcile_schedule(&self, schedule: Option<TimerSchedule>) -> Result<(), TimerError> {
219 self.inner.reconcile_ordinary(schedule)
220 }
221
222 pub fn cancel(&self) -> Result<(), TimerError> {
229 self.inner.cancel()
230 }
231}
232
233pub struct WatchdogContext {
238 inner: CallbackContext,
239}
240
241impl WatchdogContext {
242 const fn new(token: CallbackToken) -> Self {
243 Self {
244 inner: CallbackContext::new(token),
245 }
246 }
247
248 #[must_use]
250 pub const fn identity(&self) -> &TimerIdentity {
251 self.inner.identity()
252 }
253
254 pub fn ensure_scheduled(&self) -> Result<(), TimerError> {
260 self.inner.schedule_recurring()
261 }
262
263 pub fn ensure_scheduled_immediately(&self) -> Result<(), TimerError> {
270 self.inner.schedule_watchdog_immediately()
271 }
272
273 pub fn reconcile_schedule(&self, schedule: Option<TimerSchedule>) -> Result<(), TimerError> {
279 self.inner.reconcile_watchdog(schedule)
280 }
281
282 pub fn cancel(&self) -> Result<(), TimerError> {
289 self.inner.cancel()
290 }
291}
292
293#[must_use = "retain the registration claim so the timer remains controllable"]
295pub struct OnceRegistration {
296 claim: RegistrationClaim,
297}
298
299impl OnceRegistration {
300 #[must_use]
302 pub const fn identity(&self) -> &TimerIdentity {
303 self.claim.identity()
304 }
305
306 pub fn has_armed_wakeup(&self) -> Result<bool, TimerError> {
312 has_armed_wakeup_claim(&self.claim)
313 }
314
315 pub fn ensure_scheduled(&self, schedule: TimerSchedule) -> Result<(), TimerError> {
317 ensure_once_claim(&self.claim, None, schedule)
318 }
319
320 pub fn reconcile_schedule(&self, schedule: Option<TimerSchedule>) -> Result<(), TimerError> {
330 reconcile_ordinary_claim(&self.claim, None, schedule)
331 }
332
333 pub fn cancel(&self) -> Result<(), TimerError> {
339 cancel_claim(&self.claim, None)
340 }
341
342 pub fn unregister(self) -> Result<(), TimerError> {
347 unregister_claim(&self.claim)
348 }
349}
350
351#[must_use = "retain the registration claim so the timer remains controllable"]
354pub struct AfterCompletionRegistration {
355 claim: RegistrationClaim,
356}
357
358#[must_use = "retain the registration claim so the timer remains controllable"]
360pub struct WatchdogRegistration {
361 claim: RegistrationClaim,
362}
363
364trait RegistrationClaimOwner {
365 fn registration_claim(&self) -> &RegistrationClaim;
366}
367
368impl RegistrationClaimOwner for OnceRegistration {
369 fn registration_claim(&self) -> &RegistrationClaim {
370 &self.claim
371 }
372}
373
374impl RegistrationClaimOwner for AfterCompletionRegistration {
375 fn registration_claim(&self) -> &RegistrationClaim {
376 &self.claim
377 }
378}
379
380impl RegistrationClaimOwner for WatchdogRegistration {
381 fn registration_claim(&self) -> &RegistrationClaim {
382 &self.claim
383 }
384}
385
386#[derive(Clone, Copy, Debug, Eq, PartialEq)]
388pub enum TimerReconcileState {
389 Inactive,
391 Scheduled,
393}
394
395#[derive(Clone, Copy, Debug, Eq, PartialEq)]
397pub enum WatchdogReconcileState {
398 Inactive,
400 Scheduled,
402 ScheduledImmediately,
407 ScheduledAt(u64),
409}
410
411impl WatchdogRegistration {
412 #[must_use]
414 pub const fn identity(&self) -> &TimerIdentity {
415 self.claim.identity()
416 }
417
418 pub fn has_armed_wakeup(&self) -> Result<bool, TimerError> {
425 has_armed_wakeup_claim(&self.claim)
426 }
427
428 pub fn ensure_scheduled(&self) -> Result<(), TimerError> {
432 ensure_recurring_claim(&self.claim, None)
433 }
434
435 pub fn ensure_scheduled_immediately(&self) -> Result<(), TimerError> {
443 ensure_watchdog_immediately_claim(&self.claim, None)
444 }
445
446 pub fn reconcile_schedule(&self, schedule: Option<TimerSchedule>) -> Result<(), TimerError> {
454 reconcile_watchdog_claim(&self.claim, None, schedule)
455 }
456
457 pub fn cancel(&self) -> Result<(), TimerError> {
464 cancel_claim(&self.claim, None)
465 }
466
467 pub fn unregister(self) -> Result<(), TimerError> {
472 unregister_claim(&self.claim)
473 }
474}
475
476impl AfterCompletionRegistration {
477 #[must_use]
479 pub const fn identity(&self) -> &TimerIdentity {
480 self.claim.identity()
481 }
482
483 pub fn has_armed_wakeup(&self) -> Result<bool, TimerError> {
489 has_armed_wakeup_claim(&self.claim)
490 }
491
492 pub fn ensure_scheduled(&self) -> Result<(), TimerError> {
495 ensure_recurring_claim(&self.claim, None)
496 }
497
498 pub fn reconcile_schedule(&self, schedule: Option<TimerSchedule>) -> Result<(), TimerError> {
506 reconcile_ordinary_claim(&self.claim, None, schedule)
507 }
508
509 pub fn cancel(&self) -> Result<(), TimerError> {
515 cancel_claim(&self.claim, None)
516 }
517
518 pub fn unregister(self) -> Result<(), TimerError> {
523 unregister_claim(&self.claim)
524 }
525}
526
527pub fn register_once<F, Fut>(
529 identity: TimerIdentity,
530 lifetime: DeclarationLifetime,
531 callback: F,
532) -> Result<OnceRegistration, TimerError>
533where
534 F: FnMut(OnceContext) -> Fut + 'static,
535 Fut: Future<Output = TimerRunResult> + 'static,
536{
537 let callback = erase_ordinary_callback(callback, OnceContext::new);
538 let claim = with_registry_mut(|registry| {
539 registry
540 .register_once_with_callback(identity, lifetime, callback)
541 .map_err(TimerError::from)
542 })?;
543 Ok(OnceRegistration { claim })
544}
545
546pub fn register_after_completion<F, Fut>(
548 identity: TimerIdentity,
549 cadence: TimerCadence,
550 lifetime: DeclarationLifetime,
551 callback: F,
552) -> Result<AfterCompletionRegistration, TimerError>
553where
554 F: FnMut(AfterCompletionContext) -> Fut + 'static,
555 Fut: Future<Output = TimerRunResult> + 'static,
556{
557 let callback = erase_ordinary_callback(callback, AfterCompletionContext::new);
558 let claim = with_registry_mut(|registry| {
559 registry
560 .register_after_completion_with_callback(identity, cadence, lifetime, callback)
561 .map_err(TimerError::from)
562 })?;
563 Ok(AfterCompletionRegistration { claim })
564}
565
566pub fn register_watchdog<F>(
573 identity: TimerIdentity,
574 cadence: TimerCadence,
575 lifetime: DeclarationLifetime,
576 callback: F,
577) -> Result<WatchdogRegistration, TimerError>
578where
579 F: FnMut(WatchdogContext) -> WatchdogRunResult + 'static,
580{
581 let mut callback = callback;
582 let callback: WatchdogCallback = Rc::new(RefCell::new(Box::new(move |token| {
583 callback(WatchdogContext::new(token))
584 })));
585 let claim = with_registry_mut(|registry| {
586 registry
587 .register_watchdog_with_callback(identity, cadence, lifetime, callback)
588 .map_err(TimerError::from)
589 })?;
590 Ok(WatchdogRegistration { claim })
591}
592
593pub fn reconcile_once<F, Fut>(
601 registration: &mut Option<OnceRegistration>,
602 identity: &TimerIdentity,
603 desired: Option<TimerSchedule>,
604 callback: F,
605) -> Result<(), TimerError>
606where
607 F: FnMut(OnceContext) -> Fut + 'static,
608 Fut: Future<Output = TimerRunResult> + 'static,
609{
610 let registration = reconcile_registration(registration, identity, TimerPolicy::Once, || {
611 register_once(identity.clone(), DeclarationLifetime::Retained, callback)
612 })?;
613 registration.reconcile_schedule(desired)
614}
615
616pub fn reconcile_after_completion<F, Fut>(
625 registration: &mut Option<AfterCompletionRegistration>,
626 identity: &TimerIdentity,
627 cadence: TimerCadence,
628 desired: TimerReconcileState,
629 callback: F,
630) -> Result<(), TimerError>
631where
632 F: FnMut(AfterCompletionContext) -> Fut + 'static,
633 Fut: Future<Output = TimerRunResult> + 'static,
634{
635 let registration = reconcile_registration(
636 registration,
637 identity,
638 TimerPolicy::AfterCompletion { cadence },
639 || {
640 register_after_completion(
641 identity.clone(),
642 cadence,
643 DeclarationLifetime::Retained,
644 callback,
645 )
646 },
647 )?;
648 match desired {
649 TimerReconcileState::Inactive => registration.cancel(),
650 TimerReconcileState::Scheduled => registration.ensure_scheduled(),
651 }
652}
653
654pub fn reconcile_watchdog<F>(
664 registration: &mut Option<WatchdogRegistration>,
665 identity: &TimerIdentity,
666 cadence: TimerCadence,
667 desired: WatchdogReconcileState,
668 callback: F,
669) -> Result<(), TimerError>
670where
671 F: FnMut(WatchdogContext) -> WatchdogRunResult + 'static,
672{
673 let registration = reconcile_registration(
674 registration,
675 identity,
676 TimerPolicy::Watchdog { cadence },
677 || {
678 register_watchdog(
679 identity.clone(),
680 cadence,
681 DeclarationLifetime::Retained,
682 callback,
683 )
684 },
685 )?;
686 match desired {
687 WatchdogReconcileState::Inactive => registration.cancel(),
688 WatchdogReconcileState::Scheduled => registration.ensure_scheduled(),
689 WatchdogReconcileState::ScheduledImmediately => registration.ensure_scheduled_immediately(),
690 WatchdogReconcileState::ScheduledAt(deadline_ns) => {
691 registration.reconcile_schedule(Some(TimerSchedule::At(deadline_ns)))
692 }
693 }
694}
695
696fn reconcile_registration<'a, Registration>(
697 registration: &'a mut Option<Registration>,
698 identity: &TimerIdentity,
699 policy: TimerPolicy,
700 register: impl FnOnce() -> Result<Registration, TimerError>,
701) -> Result<&'a Registration, TimerError>
702where
703 Registration: RegistrationClaimOwner,
704{
705 if registration.is_none() {
706 *registration = Some(register()?);
707 }
708 let registration = registration
709 .as_ref()
710 .ok_or(TimerError::ReconciliationConflict)?;
711 verify_declaration(registration.registration_claim(), identity, policy)?;
712 Ok(registration)
713}
714
715fn verify_declaration(
716 claim: &RegistrationClaim,
717 identity: &TimerIdentity,
718 policy: TimerPolicy,
719) -> Result<(), TimerError> {
720 if claim.identity() != identity {
721 return Err(TimerError::ReconciliationConflict);
722 }
723 with_registry(|registry| {
724 registry
725 .declaration_matches(claim, policy, DeclarationLifetime::Retained)
726 .map_err(TimerError::from)?
727 .then_some(())
728 .ok_or(TimerError::ReconciliationConflict)
729 })
730}
731
732pub fn timer_snapshot(identity: &TimerIdentity) -> Result<Option<TimerSnapshot>, TimerError> {
734 with_registry(|registry| Ok(registry.snapshot(identity)))
735}
736
737pub fn timer_inventory() -> Result<TimerInventorySnapshot, TimerError> {
742 with_registry(|registry| Ok(registry.inventory()))
743}
744
745pub fn consecutive_expected_failures(identity: &TimerIdentity) -> Result<Option<u64>, TimerError> {
747 with_registry(|registry| Ok(registry.consecutive_expected_failures(identity)))
748}
749
750fn has_armed_wakeup_claim(claim: &RegistrationClaim) -> Result<bool, TimerError> {
751 with_registry(|registry| registry.has_armed_wakeup(claim).map_err(TimerError::from))
752}
753
754fn erase_ordinary_callback<Context: 'static, F, Fut>(
755 mut callback: F,
756 context: fn(CallbackToken) -> Context,
757) -> OrdinaryCallback
758where
759 F: FnMut(Context) -> Fut + 'static,
760 Fut: Future<Output = TimerRunResult> + 'static,
761{
762 Rc::new(RefCell::new(Box::new(move |token| {
763 Box::pin(callback(context(token)))
764 })))
765}
766
767fn apply_claim_transition(
768 claim: &RegistrationClaim,
769 context: Option<&CallbackToken>,
770 operation: impl FnOnce(&mut TimerRegistry) -> Result<RegistryTransition, RegistryError>,
771) -> Result<(), TimerError> {
772 let (handles, transition) = with_registry_mut(|registry| {
773 validate_context(registry, context)?;
774 let handles = registry
777 .take_provider_handles_for_claim(claim)
778 .map_err(TimerError::from)?;
779 let transition = operation(registry).map_err(TimerError::from);
780 Ok((handles, transition))
781 })?;
782 finish_detached_claim_transition(claim, handles, transition)
783}
784
785fn ensure_once_claim(
786 claim: &RegistrationClaim,
787 context: Option<&CallbackToken>,
788 schedule: TimerSchedule,
789) -> Result<(), TimerError> {
790 apply_claim_transition(claim, context, |registry| {
791 registry.ensure_once(claim, platform::time_ns(), schedule)
792 })
793}
794
795fn reconcile_ordinary_claim(
796 claim: &RegistrationClaim,
797 context: Option<&CallbackToken>,
798 schedule: Option<TimerSchedule>,
799) -> Result<(), TimerError> {
800 apply_claim_transition(claim, context, |registry| {
801 registry.reconcile_ordinary(claim, platform::time_ns(), schedule)
802 })
803}
804
805fn ensure_recurring_claim(
806 claim: &RegistrationClaim,
807 context: Option<&CallbackToken>,
808) -> Result<(), TimerError> {
809 apply_claim_transition(claim, context, |registry| {
810 registry.ensure_recurring(claim, platform::time_ns())
811 })
812}
813
814fn ensure_watchdog_immediately_claim(
815 claim: &RegistrationClaim,
816 context: Option<&CallbackToken>,
817) -> Result<(), TimerError> {
818 apply_claim_transition(claim, context, |registry| {
819 registry.ensure_watchdog_immediately(claim, platform::time_ns())
820 })
821}
822
823fn reconcile_watchdog_claim(
824 claim: &RegistrationClaim,
825 context: Option<&CallbackToken>,
826 schedule: Option<TimerSchedule>,
827) -> Result<(), TimerError> {
828 apply_claim_transition(claim, context, |registry| {
829 registry.reconcile_watchdog_schedule(claim, platform::time_ns(), schedule)
830 })
831}
832
833fn cancel_claim(
834 claim: &RegistrationClaim,
835 context: Option<&CallbackToken>,
836) -> Result<(), TimerError> {
837 apply_claim_transition(claim, context, |registry| registry.cancel(claim))
838}
839
840fn validate_context(
841 registry: &TimerRegistry,
842 context: Option<&CallbackToken>,
843) -> Result<(), TimerError> {
844 context.map_or(Ok(()), |token| {
845 registry
846 .validate_running_context(token)
847 .map_err(TimerError::from)
848 })
849}
850
851fn unregister_claim(claim: &RegistrationClaim) -> Result<(), TimerError> {
852 apply_claim_transition(claim, None, |registry| registry.unregister(claim))
853}
854
855fn finish_detached_claim_transition(
856 claim: &RegistrationClaim,
857 handles: ProviderHandles,
858 transition: Result<RegistryTransition, TimerError>,
859) -> Result<(), TimerError> {
860 match transition {
861 Ok(transition) => finish_claim_transition(claim, transition, handles),
862 Err(error) => match restore_provider_handles(handles) {
863 Ok(()) => Err(error),
864 Err(restoration_error) => retire_failed_claim(claim, restoration_error),
865 },
866 }
867}
868
869fn finish_claim_transition(
870 claim: &RegistrationClaim,
871 transition: RegistryTransition,
872 handles: ProviderHandles,
873) -> Result<(), TimerError> {
874 match finish_transition(transition, handles) {
875 result @ (Ok(()) | Err(TimerError::ControlFailure(_))) => result,
876 Err(error) => retire_failed_claim(claim, error),
877 }
878}
879
880fn retire_failed_claim(claim: &RegistrationClaim, error: TimerError) -> Result<(), TimerError> {
881 match fail_claim_provider_binding(claim) {
882 Ok(()) | Err(TimerError::RegistrationExpired) => Err(error),
883 Err(cleanup_error) => Err(cleanup_error),
884 }
885}
886
887fn finish_transition(
888 transition: RegistryTransition,
889 handles: ProviderHandles,
890) -> Result<(), TimerError> {
891 let failure = transition.failure();
892 let effect = transition.into_effect();
893 apply_effect(&effect, handles)?;
894 failure.map_or(Ok(()), |failure| Err(TimerError::ControlFailure(failure)))
895}
896
897fn apply_effect(effect: &RegistryEffect, mut handles: ProviderHandles) -> Result<(), TimerError> {
898 if !effect.has_valid_shape() {
899 clear_provider_handles(handles);
900 return Err(TimerError::OwnershipInvariant);
901 }
902 match effect {
903 RegistryEffect::None => restore_provider_handles(handles),
904 RegistryEffect::ArmWakeup { token, arm, .. } => {
905 if arm.replaces_existing() {
906 let replaced = take_detached_or_owned_handle(handles.take_wakeup(), |registry| {
907 registry.take_wakeup_handle(token.identity())
908 })?;
909 if let Some(replaced) = replaced {
910 clear_provider_handle(replaced);
911 }
912 }
913 restore_provider_handles(handles)?;
914 arm_wakeup(effect)
915 }
916 RegistryEffect::ClearCallbacks {
917 identity,
918 handles: selected,
919 } => {
920 if selected.includes_wakeup() {
921 let wakeup = take_detached_or_owned_handle(handles.take_wakeup(), |registry| {
922 registry.take_wakeup_handle(identity)
923 })?;
924 if let Some(wakeup) = wakeup {
925 clear_provider_handle(wakeup);
926 }
927 }
928 if selected.includes_work() {
929 let work = take_detached_or_owned_handle(handles.take_work(), |registry| {
930 registry.take_work_handle(identity)
931 })?;
932 if let Some(work) = work {
933 clear_provider_handle(work);
934 }
935 }
936 restore_provider_handles(handles)
937 }
938 RegistryEffect::DispatchWatchdog { successor, .. } => {
939 if let Some(wakeup) = handles.take_wakeup() {
940 clear_provider_handle(wakeup);
941 }
942 let replaced_work = take_detached_or_owned_handle(handles.take_work(), |registry| {
943 registry.take_work_handle(successor.identity())
944 })?;
945 if let Some(replaced_work) = replaced_work {
946 clear_provider_handle(replaced_work);
947 }
948 dispatch_watchdog_effect(effect)
949 }
950 }
951}
952
953fn take_detached_or_owned_handle(
954 detached: Option<ProviderHandle>,
955 take_owned: impl FnOnce(&mut TimerRegistry) -> Option<ProviderHandle>,
956) -> Result<Option<ProviderHandle>, TimerError> {
957 detached.map_or_else(
958 || with_registry_mut(|registry| Ok(take_owned(registry))),
959 |handle| Ok(Some(handle)),
960 )
961}
962
963fn arm_wakeup(effect: &RegistryEffect) -> Result<(), TimerError> {
964 let RegistryEffect::ArmWakeup {
965 token, delay_ns, ..
966 } = effect
967 else {
968 return Err(TimerError::OwnershipInvariant);
969 };
970 let task_token = token.clone();
971 let handle = platform::set_timer(Duration::from_nanos(*delay_ns), async move {
972 dispatch_wakeup(task_token).await;
973 });
974 bind_provider_handle(token, handle)?;
975 if let Err(error) = confirm_effect(effect) {
976 let handle = with_registry_mut(|registry| {
977 registry
978 .take_wakeup_handle(token.identity())
979 .ok_or(TimerError::OwnershipInvariant)
980 })?;
981 clear_provider_handle(handle);
982 return Err(error);
983 }
984 Ok(())
985}
986
987fn dispatch_watchdog_effect(effect: &RegistryEffect) -> Result<(), TimerError> {
988 let RegistryEffect::DispatchWatchdog {
989 successor,
990 successor_delay_ns,
991 work,
992 ..
993 } = effect
994 else {
995 return Err(TimerError::OwnershipInvariant);
996 };
997 let successor_token = successor.clone();
998 let successor_handle =
999 platform::set_timer(Duration::from_nanos(*successor_delay_ns), async move {
1000 dispatch_watchdog_scheduler(&successor_token);
1001 });
1002 bind_provider_handle(successor, successor_handle)?;
1003
1004 let work_token = work.clone();
1005 let work_handle = platform::set_timer(Duration::ZERO, async move {
1006 dispatch_watchdog_work(&work_token);
1007 });
1008 if let Err(error) =
1009 bind_provider_handle(work, work_handle).and_then(|()| confirm_effect(effect))
1010 {
1011 clear_entry_provider_handles(successor.identity())?;
1012 return Err(error);
1013 }
1014 Ok(())
1015}
1016
1017fn bind_provider_handle(token: &CallbackToken, handle: TimerHandle) -> Result<(), TimerError> {
1019 RUNTIME
1020 .with(|runtime| {
1021 #[cfg(test)]
1022 if take_provider_install_fault() {
1023 return Err((TimerError::OwnershipInvariant, handle));
1024 }
1025 let Ok(mut runtime) = runtime.try_borrow_mut() else {
1026 return Err((TimerError::RuntimeBusy, handle));
1027 };
1028 let Some(registry) = runtime.as_mut() else {
1029 return Err((TimerError::NotInitialized, handle));
1030 };
1031 match registry.install_provider_handle(token, handle) {
1032 Ok(()) => Ok(()),
1033 Err((error, handle)) => Err((TimerError::from(error), handle)),
1034 }
1035 })
1036 .map_err(|(error, handle)| {
1037 platform::clear_timer(handle);
1039 error
1040 })
1041}
1042
1043fn confirm_effect(effect: &RegistryEffect) -> Result<(), TimerError> {
1044 #[cfg(test)]
1045 if take_provider_confirmation_fault() {
1046 return Err(TimerError::OwnershipInvariant);
1047 }
1048 with_registry_mut(|registry| {
1049 registry
1050 .confirm_effect_applied(effect)
1051 .map_err(TimerError::from)
1052 })
1053}
1054
1055fn restore_provider_handles(mut handles: ProviderHandles) -> Result<(), TimerError> {
1056 let wakeup_failure = handles
1059 .take_wakeup()
1060 .and_then(|handle| restore_provider_handle(handle).err());
1061 let work_failure = handles
1062 .take_work()
1063 .and_then(|handle| restore_provider_handle(handle).err());
1064 wakeup_failure.or(work_failure).map_or(Ok(()), Err)
1065}
1066
1067fn restore_provider_handle(handle: ProviderHandle) -> Result<(), TimerError> {
1068 let (token, handle) = handle.into_parts();
1069 bind_provider_handle(&token, handle)
1070}
1071
1072fn clear_provider_handle(handle: ProviderHandle) {
1073 let (_, handle) = handle.into_parts();
1074 platform::clear_timer(handle);
1075}
1076
1077fn clear_provider_handles(mut handles: ProviderHandles) {
1078 if let Some(wakeup) = handles.take_wakeup() {
1079 clear_provider_handle(wakeup);
1080 }
1081 if let Some(work) = handles.take_work() {
1082 clear_provider_handle(work);
1083 }
1084}
1085
1086fn clear_entry_provider_handles(identity: &TimerIdentity) -> Result<(), TimerError> {
1087 let handles = with_registry_mut(|registry| {
1088 Ok(ProviderHandles::from_parts(
1089 registry.take_wakeup_handle(identity),
1090 registry.take_work_handle(identity),
1091 ))
1092 })?;
1093 clear_provider_handles(handles);
1094 Ok(())
1095}
1096
1097#[expect(
1098 clippy::future_not_send,
1099 reason = "IC callbacks and canister-local state are single-threaded."
1100)]
1101async fn dispatch_wakeup(token: CallbackToken) {
1102 match token.role() {
1103 CallbackRole::OrdinaryWork => dispatch_ordinary(token).await,
1104 CallbackRole::WatchdogScheduler => dispatch_watchdog_scheduler(&token),
1105 CallbackRole::WatchdogWork => {}
1106 }
1107}
1108
1109#[expect(
1110 clippy::future_not_send,
1111 reason = "IC callbacks and canister-local state are single-threaded."
1112)]
1113async fn dispatch_ordinary(token: CallbackToken) {
1114 let measurement = CallbackMeasurementStart::capture();
1115 let accepted = with_registry_mut(|registry| {
1116 registry.consume_provider_handle(&token);
1117 Ok(registry.begin_ordinary(&token))
1118 });
1119 match accepted {
1120 Ok(CallbackAcceptance::Accepted) => {}
1121 Ok(CallbackAcceptance::Stale) => return,
1122 Err(error) => trap_callback_failure("ordinary callback acceptance", &error),
1123 }
1124
1125 let callback = match with_registry(|registry| {
1126 registry.ordinary_callback(&token).map_err(TimerError::from)
1127 }) {
1128 Ok(callback) => callback,
1129 Err(TimerError::OwnershipInvariant) => {
1130 fail_ordinary_dispatch(&token);
1131 return;
1132 }
1133 Err(error) => trap_callback_failure("ordinary callback lookup", &error),
1134 };
1135 let future = {
1136 let Ok(mut callback) = callback.try_borrow_mut() else {
1137 fail_ordinary_dispatch(&token);
1138 return;
1139 };
1140 callback(token.clone())
1141 };
1142 let result = future.await;
1143 let transition = with_registry_mut(|registry| {
1144 registry
1145 .complete_ordinary(&token, platform::time_ns(), result)
1146 .map_err(TimerError::from)
1147 });
1148 let transition = transition
1149 .unwrap_or_else(|error| trap_callback_failure("ordinary callback completion", &error));
1150 finish_callback_transition(&token, transition, ProviderHandles::default());
1151 record_callback_measurements(&token, measurement.finish());
1152}
1153
1154fn fail_ordinary_dispatch(token: &CallbackToken) {
1155 let transition = with_registry_mut(|registry| {
1156 registry
1157 .complete_ordinary(
1158 token,
1159 platform::time_ns(),
1160 TimerRunResult::new(TimerCompletion::invariant_failure(0), TimerDirective::Stop),
1161 )
1162 .map_err(TimerError::from)
1163 });
1164 let transition = transition.unwrap_or_else(|error| {
1165 trap_callback_failure("ordinary invariant-failure completion", &error)
1166 });
1167 finish_callback_transition(token, transition, ProviderHandles::default());
1168}
1169
1170fn dispatch_watchdog_scheduler(token: &CallbackToken) {
1171 let measurement = CallbackMeasurementStart::capture();
1172 let dispatched = with_registry_mut(|registry| {
1173 registry.consume_provider_handle(token);
1174 let claim = RegistrationClaim::from_callback(token);
1175 let handles = match registry.take_provider_handles_for_claim(&claim) {
1176 Ok(handles) => handles,
1177 Err(RegistryError::UnknownRegistration | RegistryError::StaleRegistration) => {
1179 ProviderHandles::default()
1180 }
1181 Err(error) => return Err(TimerError::from(error)),
1182 };
1183 Ok((
1184 registry.begin_watchdog_scheduler(token, platform::time_ns()),
1185 handles,
1186 ))
1187 });
1188 let (transition, handles) = dispatched
1189 .unwrap_or_else(|error| trap_callback_failure("watchdog scheduler transition", &error));
1190 let accepted = !matches!(transition.effect(), RegistryEffect::None);
1191 finish_callback_transition(token, transition, handles);
1192 if accepted {
1193 record_callback_measurements(token, measurement.finish());
1194 }
1195}
1196
1197fn dispatch_watchdog_work(token: &CallbackToken) {
1198 let measurement = CallbackMeasurementStart::capture();
1199 let accepted = with_registry_mut(|registry| {
1200 registry.consume_provider_handle(token);
1201 Ok(registry.begin_watchdog_work(token))
1202 });
1203 match accepted {
1204 Ok(CallbackAcceptance::Accepted) => {}
1205 Ok(CallbackAcceptance::Stale) => return,
1206 Err(error) => trap_callback_failure("watchdog work acceptance", &error),
1207 }
1208
1209 let callback =
1210 match with_registry(|registry| registry.watchdog_callback(token).map_err(TimerError::from))
1211 {
1212 Ok(callback) => callback,
1213 Err(error) => trap_callback_failure("watchdog callback lookup", &error),
1214 };
1215 let result = {
1216 let Ok(mut callback) = callback.try_borrow_mut() else {
1217 trap_callback_failure(
1218 "watchdog callback ownership",
1219 &TimerError::OwnershipInvariant,
1220 );
1221 };
1222 callback(token.clone())
1223 };
1224 finish_watchdog_dispatch(token, result);
1225 record_callback_measurements(token, measurement.finish());
1226}
1227
1228fn finish_watchdog_dispatch(token: &CallbackToken, result: WatchdogRunResult) {
1229 let claim = RegistrationClaim::from_callback(token);
1230 let completed = with_registry_mut(|registry| {
1235 let handles = registry
1236 .take_provider_handles_for_claim(&claim)
1237 .map_err(TimerError::from)?;
1238 #[cfg(test)]
1239 {
1240 if take_watchdog_completion_fault() {
1241 return Err(TimerError::OwnershipInvariant);
1242 }
1243 }
1244 let transition = registry
1245 .complete_watchdog_work(token, platform::time_ns(), result)
1246 .map_err(TimerError::from)?;
1247 Ok((transition, handles))
1248 });
1249 let (transition, handles) =
1250 completed.unwrap_or_else(|error| trap_callback_failure("watchdog work completion", &error));
1251 finish_callback_transition(token, transition, handles);
1252}
1253
1254fn finish_callback_transition(
1255 token: &CallbackToken,
1256 transition: RegistryTransition,
1257 handles: ProviderHandles,
1258) {
1259 match finish_transition(transition, handles) {
1262 Ok(()) | Err(TimerError::ControlFailure(_)) => {}
1263 Err(error) => {
1264 if token.role() == CallbackRole::WatchdogWork {
1265 trap_callback_failure("watchdog provider-handle completion", &error);
1266 }
1267 fail_provider_binding(token).unwrap_or_else(|binding_error| {
1268 trap_callback_failure("provider-binding failure cleanup", &binding_error)
1269 });
1270 }
1271 }
1272}
1273
1274fn fail_provider_binding(token: &CallbackToken) -> Result<(), TimerError> {
1275 let claim = RegistrationClaim::from_callback(token);
1276 fail_claim_provider_binding(&claim)
1277}
1278
1279fn fail_claim_provider_binding(claim: &RegistrationClaim) -> Result<(), TimerError> {
1280 let failed = with_registry_mut(|registry| {
1281 registry
1282 .fail_registration(claim, TimerControlFailure::ProviderBindingFailed)
1283 .map_err(TimerError::from)
1284 });
1285 clear_provider_handles(failed?);
1286 Ok(())
1287}
1288
1289#[derive(Clone, Copy)]
1290struct CallbackMeasurementStart {
1291 instructions_before: u64,
1292 memory_start: platform::MemoryPages,
1293}
1294
1295impl CallbackMeasurementStart {
1296 fn capture() -> Self {
1297 let memory_start = platform::memory_pages();
1299 let instructions_before = platform::instruction_counter();
1300 Self {
1301 instructions_before,
1302 memory_start,
1303 }
1304 }
1305
1306 fn finish(self) -> CallbackMeasurement {
1307 let instructions = platform::instruction_counter().saturating_sub(self.instructions_before);
1309 let memory_end = platform::memory_pages();
1310 CallbackMeasurement {
1311 instructions,
1312 memory_start: self.memory_start,
1313 memory_end,
1314 }
1315 }
1316}
1317
1318#[derive(Clone, Copy)]
1319struct CallbackMeasurement {
1320 instructions: u64,
1321 memory_start: platform::MemoryPages,
1322 memory_end: platform::MemoryPages,
1323}
1324
1325fn record_callback_measurements(token: &CallbackToken, measurement: CallbackMeasurement) {
1326 with_registry_mut(|registry| {
1327 registry
1328 .record_callback_measurements(
1329 token,
1330 measurement.instructions,
1331 measurement.memory_start,
1332 measurement.memory_end,
1333 )
1334 .map_err(TimerError::from)
1335 })
1336 .unwrap_or_else(|error| trap_callback_failure("callback measurement accounting", &error));
1337}
1338
1339fn trap_callback_failure(context: &str, error: &TimerError) -> ! {
1340 platform::trap(&format!("ic-timers {context} failed: {error}"))
1341}
1342
1343fn with_registry<T>(
1344 operation: impl FnOnce(&TimerRegistry) -> Result<T, TimerError>,
1345) -> Result<T, TimerError> {
1346 RUNTIME.with(|runtime| {
1347 let runtime = runtime.try_borrow().map_err(|_| TimerError::RuntimeBusy)?;
1348 let registry = runtime.as_ref().ok_or(TimerError::NotInitialized)?;
1349 operation(registry)
1350 })
1351}
1352
1353fn with_registry_mut<T>(
1354 operation: impl FnOnce(&mut TimerRegistry) -> Result<T, TimerError>,
1355) -> Result<T, TimerError> {
1356 RUNTIME.with(|runtime| {
1357 let mut runtime = runtime
1358 .try_borrow_mut()
1359 .map_err(|_| TimerError::RuntimeBusy)?;
1360 let registry = runtime.as_mut().ok_or(TimerError::NotInitialized)?;
1361 operation(registry)
1362 })
1363}
1364
1365#[cfg(test)]
1366fn reset_for_test(now_ns: u64, canister_version: u64) {
1367 platform::reset(now_ns, canister_version);
1368 WATCHDOG_COMPLETION_FAULT.with(|fault| fault.set(false));
1369 PROVIDER_INSTALL_FAULT_AFTER.with(|fault| fault.set(None));
1370 PROVIDER_CONFIRMATION_FAULT.with(|fault| fault.set(false));
1371 RUNTIME.with(|runtime| {
1372 *runtime.borrow_mut() = None;
1373 });
1374}
1375
1376#[cfg(test)]
1377thread_local! {
1378 static WATCHDOG_COMPLETION_FAULT: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
1379 static PROVIDER_INSTALL_FAULT_AFTER: std::cell::Cell<Option<u64>> = const { std::cell::Cell::new(None) };
1380 static PROVIDER_CONFIRMATION_FAULT: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
1381}
1382
1383#[cfg(test)]
1384fn inject_watchdog_completion_fault() {
1385 WATCHDOG_COMPLETION_FAULT.with(|fault| fault.set(true));
1386}
1387
1388#[cfg(test)]
1389fn take_watchdog_completion_fault() -> bool {
1390 WATCHDOG_COMPLETION_FAULT.with(|fault| fault.replace(false))
1391}
1392
1393#[cfg(test)]
1394fn inject_provider_install_fault() {
1395 inject_provider_install_fault_after(0);
1396}
1397
1398#[cfg(test)]
1399fn inject_provider_install_fault_after(successful_installs: u64) {
1400 PROVIDER_INSTALL_FAULT_AFTER.with(|fault| fault.set(Some(successful_installs)));
1401}
1402
1403#[cfg(test)]
1404fn take_provider_install_fault() -> bool {
1405 PROVIDER_INSTALL_FAULT_AFTER.with(|fault| match fault.get() {
1406 Some(0) => {
1407 fault.set(None);
1408 true
1409 }
1410 Some(remaining) => {
1411 fault.set(Some(remaining - 1));
1412 false
1413 }
1414 None => false,
1415 })
1416}
1417
1418#[cfg(test)]
1419fn inject_provider_confirmation_fault() {
1420 PROVIDER_CONFIRMATION_FAULT.with(|fault| fault.set(true));
1421}
1422
1423#[cfg(test)]
1424fn take_provider_confirmation_fault() -> bool {
1425 PROVIDER_CONFIRMATION_FAULT.with(|fault| fault.replace(false))
1426}
1427
1428#[cfg(test)]
1429mod tests;