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 { .. } => Self::OwnershipInvariant,
60 RegistryError::Schedule(error) => Self::Schedule(error),
61 RegistryError::MissingCallback | RegistryError::ProviderHandleAlreadyOwned => {
62 Self::OwnershipInvariant
63 }
64 }
65 }
66}
67
68pub fn initialize_runtime() -> Result<TimerEpoch, TimerError> {
73 let epoch = TimerEpoch::new(platform::canister_version(), platform::time_ns());
74 RUNTIME.with(|runtime| {
75 let mut runtime = runtime
76 .try_borrow_mut()
77 .map_err(|_| TimerError::RuntimeBusy)?;
78 if let Some(registry) = runtime.as_ref() {
79 return Ok(registry.epoch());
80 }
81 *runtime = Some(TimerRegistry::new(epoch));
82 Ok(epoch)
83 })
84}
85
86struct CallbackContext {
87 token: CallbackToken,
88}
89
90impl CallbackContext {
91 const fn new(token: CallbackToken) -> Self {
92 Self { token }
93 }
94
95 fn claim(&self) -> RegistrationClaim {
96 RegistrationClaim::from_callback(&self.token)
97 }
98
99 const fn identity(&self) -> &TimerIdentity {
100 self.token.identity()
101 }
102
103 fn cancel(&self) -> Result<(), TimerError> {
104 cancel_claim(&self.claim(), Some(&self.token))
105 }
106
107 fn schedule_once(&self, schedule: TimerSchedule) -> Result<(), TimerError> {
108 ensure_once_claim(&self.claim(), Some(&self.token), schedule)
109 }
110
111 fn schedule_recurring(&self) -> Result<(), TimerError> {
112 ensure_recurring_claim(&self.claim(), Some(&self.token))
113 }
114
115 fn schedule_watchdog_immediately(&self) -> Result<(), TimerError> {
116 ensure_watchdog_immediately_claim(&self.claim(), Some(&self.token))
117 }
118
119 fn reconcile_watchdog(&self, schedule: Option<TimerSchedule>) -> Result<(), TimerError> {
120 reconcile_watchdog_claim(&self.claim(), Some(&self.token), schedule)
121 }
122
123 fn reconcile_ordinary(&self, schedule: Option<TimerSchedule>) -> Result<(), TimerError> {
124 reconcile_ordinary_claim(&self.claim(), Some(&self.token), schedule)
125 }
126}
127
128pub struct OnceContext {
134 inner: CallbackContext,
135}
136
137impl OnceContext {
138 const fn new(token: CallbackToken) -> Self {
139 Self {
140 inner: CallbackContext::new(token),
141 }
142 }
143
144 #[must_use]
146 pub const fn identity(&self) -> &TimerIdentity {
147 self.inner.identity()
148 }
149
150 pub fn ensure_scheduled(&self, schedule: TimerSchedule) -> Result<(), TimerError> {
155 self.inner.schedule_once(schedule)
156 }
157
158 pub fn reconcile_schedule(&self, schedule: Option<TimerSchedule>) -> Result<(), TimerError> {
165 self.inner.reconcile_ordinary(schedule)
166 }
167
168 pub fn cancel(&self) -> Result<(), TimerError> {
175 self.inner.cancel()
176 }
177}
178
179pub struct AfterCompletionContext {
185 inner: CallbackContext,
186}
187
188impl AfterCompletionContext {
189 const fn new(token: CallbackToken) -> Self {
190 Self {
191 inner: CallbackContext::new(token),
192 }
193 }
194
195 #[must_use]
197 pub const fn identity(&self) -> &TimerIdentity {
198 self.inner.identity()
199 }
200
201 pub fn ensure_scheduled(&self) -> Result<(), TimerError> {
203 self.inner.schedule_recurring()
204 }
205
206 pub fn reconcile_schedule(&self, schedule: Option<TimerSchedule>) -> Result<(), TimerError> {
214 self.inner.reconcile_ordinary(schedule)
215 }
216
217 pub fn cancel(&self) -> Result<(), TimerError> {
224 self.inner.cancel()
225 }
226}
227
228pub struct WatchdogContext {
233 inner: CallbackContext,
234}
235
236impl WatchdogContext {
237 const fn new(token: CallbackToken) -> Self {
238 Self {
239 inner: CallbackContext::new(token),
240 }
241 }
242
243 #[must_use]
245 pub const fn identity(&self) -> &TimerIdentity {
246 self.inner.identity()
247 }
248
249 pub fn ensure_scheduled(&self) -> Result<(), TimerError> {
255 self.inner.schedule_recurring()
256 }
257
258 pub fn ensure_scheduled_immediately(&self) -> Result<(), TimerError> {
265 self.inner.schedule_watchdog_immediately()
266 }
267
268 pub fn reconcile_schedule(&self, schedule: Option<TimerSchedule>) -> Result<(), TimerError> {
274 self.inner.reconcile_watchdog(schedule)
275 }
276
277 pub fn cancel(&self) -> Result<(), TimerError> {
284 self.inner.cancel()
285 }
286}
287
288#[must_use = "retain the registration claim so the timer remains controllable"]
290pub struct OnceRegistration {
291 claim: RegistrationClaim,
292}
293
294impl OnceRegistration {
295 #[must_use]
297 pub const fn identity(&self) -> &TimerIdentity {
298 self.claim.identity()
299 }
300
301 pub fn has_armed_wakeup(&self) -> Result<bool, TimerError> {
307 has_armed_wakeup_claim(&self.claim)
308 }
309
310 pub fn ensure_scheduled(&self, schedule: TimerSchedule) -> Result<(), TimerError> {
312 ensure_once_claim(&self.claim, None, schedule)
313 }
314
315 pub fn reconcile_schedule(&self, schedule: Option<TimerSchedule>) -> Result<(), TimerError> {
322 reconcile_ordinary_claim(&self.claim, None, schedule)
323 }
324
325 pub fn cancel(&self) -> Result<(), TimerError> {
331 cancel_claim(&self.claim, None)
332 }
333
334 pub fn unregister(self) -> Result<(), TimerError> {
339 unregister_claim(&self.claim)
340 }
341}
342
343#[must_use = "retain the registration claim so the timer remains controllable"]
346pub struct AfterCompletionRegistration {
347 claim: RegistrationClaim,
348}
349
350#[must_use = "retain the registration claim so the timer remains controllable"]
352pub struct WatchdogRegistration {
353 claim: RegistrationClaim,
354}
355
356trait RegistrationClaimOwner {
357 fn registration_claim(&self) -> &RegistrationClaim;
358}
359
360impl RegistrationClaimOwner for OnceRegistration {
361 fn registration_claim(&self) -> &RegistrationClaim {
362 &self.claim
363 }
364}
365
366impl RegistrationClaimOwner for AfterCompletionRegistration {
367 fn registration_claim(&self) -> &RegistrationClaim {
368 &self.claim
369 }
370}
371
372impl RegistrationClaimOwner for WatchdogRegistration {
373 fn registration_claim(&self) -> &RegistrationClaim {
374 &self.claim
375 }
376}
377
378#[derive(Clone, Copy, Debug, Eq, PartialEq)]
380pub enum TimerReconcileState {
381 Inactive,
383 Scheduled,
385}
386
387#[derive(Clone, Copy, Debug, Eq, PartialEq)]
389pub enum WatchdogReconcileState {
390 Inactive,
392 Scheduled,
394 ScheduledImmediately,
399 ScheduledAt(u64),
401}
402
403impl WatchdogRegistration {
404 #[must_use]
406 pub const fn identity(&self) -> &TimerIdentity {
407 self.claim.identity()
408 }
409
410 pub fn has_armed_wakeup(&self) -> Result<bool, TimerError> {
417 has_armed_wakeup_claim(&self.claim)
418 }
419
420 pub fn ensure_scheduled(&self) -> Result<(), TimerError> {
424 ensure_recurring_claim(&self.claim, None)
425 }
426
427 pub fn ensure_scheduled_immediately(&self) -> Result<(), TimerError> {
435 ensure_watchdog_immediately_claim(&self.claim, None)
436 }
437
438 pub fn reconcile_schedule(&self, schedule: Option<TimerSchedule>) -> Result<(), TimerError> {
446 reconcile_watchdog_claim(&self.claim, None, schedule)
447 }
448
449 pub fn cancel(&self) -> Result<(), TimerError> {
456 cancel_claim(&self.claim, None)
457 }
458
459 pub fn unregister(self) -> Result<(), TimerError> {
464 unregister_claim(&self.claim)
465 }
466}
467
468impl AfterCompletionRegistration {
469 #[must_use]
471 pub const fn identity(&self) -> &TimerIdentity {
472 self.claim.identity()
473 }
474
475 pub fn has_armed_wakeup(&self) -> Result<bool, TimerError> {
481 has_armed_wakeup_claim(&self.claim)
482 }
483
484 pub fn ensure_scheduled(&self) -> Result<(), TimerError> {
487 ensure_recurring_claim(&self.claim, None)
488 }
489
490 pub fn reconcile_schedule(&self, schedule: Option<TimerSchedule>) -> Result<(), TimerError> {
495 reconcile_ordinary_claim(&self.claim, None, schedule)
496 }
497
498 pub fn cancel(&self) -> Result<(), TimerError> {
504 cancel_claim(&self.claim, None)
505 }
506
507 pub fn unregister(self) -> Result<(), TimerError> {
512 unregister_claim(&self.claim)
513 }
514}
515
516pub fn register_once<F, Fut>(
518 identity: TimerIdentity,
519 lifetime: DeclarationLifetime,
520 callback: F,
521) -> Result<OnceRegistration, TimerError>
522where
523 F: FnMut(OnceContext) -> Fut + 'static,
524 Fut: Future<Output = TimerRunResult> + 'static,
525{
526 let callback = erase_ordinary_callback(callback, OnceContext::new);
527 let claim = with_registry_mut(|registry| {
528 registry
529 .register_once_with_callback(identity, lifetime, callback)
530 .map_err(TimerError::from)
531 })?;
532 Ok(OnceRegistration { claim })
533}
534
535pub fn register_after_completion<F, Fut>(
537 identity: TimerIdentity,
538 cadence: TimerCadence,
539 lifetime: DeclarationLifetime,
540 callback: F,
541) -> Result<AfterCompletionRegistration, TimerError>
542where
543 F: FnMut(AfterCompletionContext) -> Fut + 'static,
544 Fut: Future<Output = TimerRunResult> + 'static,
545{
546 let callback = erase_ordinary_callback(callback, AfterCompletionContext::new);
547 let claim = with_registry_mut(|registry| {
548 registry
549 .register_after_completion_with_callback(identity, cadence, lifetime, callback)
550 .map_err(TimerError::from)
551 })?;
552 Ok(AfterCompletionRegistration { claim })
553}
554
555pub fn register_watchdog<F>(
562 identity: TimerIdentity,
563 cadence: TimerCadence,
564 lifetime: DeclarationLifetime,
565 callback: F,
566) -> Result<WatchdogRegistration, TimerError>
567where
568 F: FnMut(WatchdogContext) -> WatchdogRunResult + 'static,
569{
570 let mut callback = callback;
571 let callback: WatchdogCallback = Rc::new(RefCell::new(Box::new(move |token| {
572 callback(WatchdogContext::new(token))
573 })));
574 let claim = with_registry_mut(|registry| {
575 registry
576 .register_watchdog_with_callback(identity, cadence, lifetime, callback)
577 .map_err(TimerError::from)
578 })?;
579 Ok(WatchdogRegistration { claim })
580}
581
582pub fn reconcile_once<F, Fut>(
590 registration: &mut Option<OnceRegistration>,
591 identity: &TimerIdentity,
592 desired: Option<TimerSchedule>,
593 callback: F,
594) -> Result<(), TimerError>
595where
596 F: FnMut(OnceContext) -> Fut + 'static,
597 Fut: Future<Output = TimerRunResult> + 'static,
598{
599 let registration = reconcile_registration(registration, identity, TimerPolicy::Once, || {
600 register_once(identity.clone(), DeclarationLifetime::Retained, callback)
601 })?;
602 registration.reconcile_schedule(desired)
603}
604
605pub fn reconcile_after_completion<F, Fut>(
614 registration: &mut Option<AfterCompletionRegistration>,
615 identity: &TimerIdentity,
616 cadence: TimerCadence,
617 desired: TimerReconcileState,
618 callback: F,
619) -> Result<(), TimerError>
620where
621 F: FnMut(AfterCompletionContext) -> Fut + 'static,
622 Fut: Future<Output = TimerRunResult> + 'static,
623{
624 let registration = reconcile_registration(
625 registration,
626 identity,
627 TimerPolicy::AfterCompletion { cadence },
628 || {
629 register_after_completion(
630 identity.clone(),
631 cadence,
632 DeclarationLifetime::Retained,
633 callback,
634 )
635 },
636 )?;
637 match desired {
638 TimerReconcileState::Inactive => registration.cancel(),
639 TimerReconcileState::Scheduled => registration.ensure_scheduled(),
640 }
641}
642
643pub fn reconcile_watchdog<F>(
653 registration: &mut Option<WatchdogRegistration>,
654 identity: &TimerIdentity,
655 cadence: TimerCadence,
656 desired: WatchdogReconcileState,
657 callback: F,
658) -> Result<(), TimerError>
659where
660 F: FnMut(WatchdogContext) -> WatchdogRunResult + 'static,
661{
662 let registration = reconcile_registration(
663 registration,
664 identity,
665 TimerPolicy::Watchdog { cadence },
666 || {
667 register_watchdog(
668 identity.clone(),
669 cadence,
670 DeclarationLifetime::Retained,
671 callback,
672 )
673 },
674 )?;
675 match desired {
676 WatchdogReconcileState::Inactive => registration.cancel(),
677 WatchdogReconcileState::Scheduled => registration.ensure_scheduled(),
678 WatchdogReconcileState::ScheduledImmediately => registration.ensure_scheduled_immediately(),
679 WatchdogReconcileState::ScheduledAt(deadline_ns) => {
680 registration.reconcile_schedule(Some(TimerSchedule::At(deadline_ns)))
681 }
682 }
683}
684
685fn reconcile_registration<'a, Registration>(
686 registration: &'a mut Option<Registration>,
687 identity: &TimerIdentity,
688 policy: TimerPolicy,
689 register: impl FnOnce() -> Result<Registration, TimerError>,
690) -> Result<&'a Registration, TimerError>
691where
692 Registration: RegistrationClaimOwner,
693{
694 if registration.is_none() {
695 *registration = Some(register()?);
696 }
697 let registration = registration
698 .as_ref()
699 .ok_or(TimerError::ReconciliationConflict)?;
700 verify_declaration(registration.registration_claim(), identity, policy)?;
701 Ok(registration)
702}
703
704fn verify_declaration(
705 claim: &RegistrationClaim,
706 identity: &TimerIdentity,
707 policy: TimerPolicy,
708) -> Result<(), TimerError> {
709 if claim.identity() != identity {
710 return Err(TimerError::ReconciliationConflict);
711 }
712 with_registry(|registry| {
713 registry
714 .declaration_matches(claim, policy, DeclarationLifetime::Retained)
715 .map_err(TimerError::from)?
716 .then_some(())
717 .ok_or(TimerError::ReconciliationConflict)
718 })
719}
720
721pub fn timer_snapshot(identity: &TimerIdentity) -> Result<Option<TimerSnapshot>, TimerError> {
723 with_registry(|registry| Ok(registry.snapshot(identity)))
724}
725
726pub fn timer_inventory() -> Result<TimerInventorySnapshot, TimerError> {
731 with_registry(|registry| Ok(registry.inventory()))
732}
733
734pub fn consecutive_expected_failures(identity: &TimerIdentity) -> Result<Option<u64>, TimerError> {
736 with_registry(|registry| Ok(registry.consecutive_expected_failures(identity)))
737}
738
739fn has_armed_wakeup_claim(claim: &RegistrationClaim) -> Result<bool, TimerError> {
740 with_registry(|registry| registry.has_armed_wakeup(claim).map_err(TimerError::from))
741}
742
743fn erase_ordinary_callback<Context: 'static, F, Fut>(
744 mut callback: F,
745 context: fn(CallbackToken) -> Context,
746) -> OrdinaryCallback
747where
748 F: FnMut(Context) -> Fut + 'static,
749 Fut: Future<Output = TimerRunResult> + 'static,
750{
751 Rc::new(RefCell::new(Box::new(move |token| {
752 Box::pin(callback(context(token)))
753 })))
754}
755
756fn apply_claim_transition(
757 claim: &RegistrationClaim,
758 context: Option<&CallbackToken>,
759 operation: impl FnOnce(&mut TimerRegistry) -> Result<RegistryTransition, RegistryError>,
760) -> Result<(), TimerError> {
761 let transition = with_registry_mut(|registry| {
762 validate_context(registry, context)?;
763 operation(registry).map_err(TimerError::from)
764 })?;
765 finish_claim_transition(claim, transition, ProviderHandles::default())
766}
767
768fn ensure_once_claim(
769 claim: &RegistrationClaim,
770 context: Option<&CallbackToken>,
771 schedule: TimerSchedule,
772) -> Result<(), TimerError> {
773 apply_claim_transition(claim, context, |registry| {
774 registry.ensure_once(claim, platform::time_ns(), schedule)
775 })
776}
777
778fn reconcile_ordinary_claim(
779 claim: &RegistrationClaim,
780 context: Option<&CallbackToken>,
781 schedule: Option<TimerSchedule>,
782) -> Result<(), TimerError> {
783 if schedule.is_none() {
784 let (handles, transition) = with_registry_mut(|registry| {
785 validate_context(registry, context)?;
786 registry
787 .validate_ordinary_claim(claim)
788 .map_err(TimerError::from)?;
789 let handles = registry
790 .take_provider_handles_for_claim(claim)
791 .map_err(TimerError::from)?;
792 let transition = registry
793 .reconcile_ordinary(claim, platform::time_ns(), None)
794 .map_err(TimerError::from);
795 Ok((handles, transition))
796 })?;
797 return finish_detached_claim_transition(claim, handles, transition);
798 }
799 apply_claim_transition(claim, context, |registry| {
800 registry.reconcile_ordinary(claim, platform::time_ns(), schedule)
801 })
802}
803
804fn ensure_recurring_claim(
805 claim: &RegistrationClaim,
806 context: Option<&CallbackToken>,
807) -> Result<(), TimerError> {
808 apply_claim_transition(claim, context, |registry| {
809 registry.ensure_recurring(claim, platform::time_ns())
810 })
811}
812
813fn ensure_watchdog_immediately_claim(
814 claim: &RegistrationClaim,
815 context: Option<&CallbackToken>,
816) -> Result<(), TimerError> {
817 apply_claim_transition(claim, context, |registry| {
818 registry.ensure_watchdog_immediately(claim, platform::time_ns())
819 })
820}
821
822fn reconcile_watchdog_claim(
823 claim: &RegistrationClaim,
824 context: Option<&CallbackToken>,
825 schedule: Option<TimerSchedule>,
826) -> Result<(), TimerError> {
827 if schedule.is_none() {
828 return apply_detached_claim_transition(claim, context, |registry| {
829 registry.reconcile_watchdog_schedule(claim, platform::time_ns(), None)
830 });
831 }
832 apply_claim_transition(claim, context, |registry| {
833 registry.reconcile_watchdog_schedule(claim, platform::time_ns(), schedule)
834 })
835}
836
837fn cancel_claim(
838 claim: &RegistrationClaim,
839 context: Option<&CallbackToken>,
840) -> Result<(), TimerError> {
841 apply_detached_claim_transition(claim, context, |registry| registry.cancel(claim))
842}
843
844fn validate_context(
845 registry: &TimerRegistry,
846 context: Option<&CallbackToken>,
847) -> Result<(), TimerError> {
848 context.map_or(Ok(()), |token| {
849 registry
850 .validate_running_context(token)
851 .map_err(TimerError::from)
852 })
853}
854
855fn unregister_claim(claim: &RegistrationClaim) -> Result<(), TimerError> {
856 apply_detached_claim_transition(claim, None, |registry| registry.unregister(claim))
857}
858
859fn apply_detached_claim_transition(
860 claim: &RegistrationClaim,
861 context: Option<&CallbackToken>,
862 operation: impl FnOnce(&mut TimerRegistry) -> Result<RegistryTransition, RegistryError>,
863) -> Result<(), TimerError> {
864 let (handles, transition) = with_registry_mut(|registry| {
865 validate_context(registry, context)?;
866 let handles = registry
867 .take_provider_handles_for_claim(claim)
868 .map_err(TimerError::from)?;
869 let transition = operation(registry).map_err(TimerError::from);
870 Ok((handles, transition))
871 })?;
872 finish_detached_claim_transition(claim, handles, transition)
873}
874
875fn finish_detached_claim_transition(
876 claim: &RegistrationClaim,
877 handles: ProviderHandles,
878 transition: Result<RegistryTransition, TimerError>,
879) -> Result<(), TimerError> {
880 match transition {
881 Ok(transition) => finish_claim_transition(claim, transition, handles),
882 Err(error) => match restore_provider_handles(handles) {
883 Ok(()) => Err(error),
884 Err(restoration_error) => retire_failed_claim(claim, restoration_error),
885 },
886 }
887}
888
889fn finish_claim_transition(
890 claim: &RegistrationClaim,
891 transition: RegistryTransition,
892 handles: ProviderHandles,
893) -> Result<(), TimerError> {
894 match finish_transition(transition, handles) {
895 result @ (Ok(()) | Err(TimerError::ControlFailure(_))) => result,
896 Err(error) => retire_failed_claim(claim, error),
897 }
898}
899
900fn retire_failed_claim(claim: &RegistrationClaim, error: TimerError) -> Result<(), TimerError> {
901 match fail_claim_provider_binding(claim) {
902 Ok(()) | Err(TimerError::RegistrationExpired) => Err(error),
903 Err(cleanup_error) => Err(cleanup_error),
904 }
905}
906
907fn finish_transition(
908 transition: RegistryTransition,
909 handles: ProviderHandles,
910) -> Result<(), TimerError> {
911 let failure = transition.failure();
912 let effect = transition.into_effect();
913 apply_effect(&effect, handles)?;
914 failure.map_or(Ok(()), |failure| Err(TimerError::ControlFailure(failure)))
915}
916
917fn apply_effect(effect: &RegistryEffect, mut handles: ProviderHandles) -> Result<(), TimerError> {
918 if !effect.has_valid_shape() {
919 clear_provider_handles(handles);
920 return Err(TimerError::OwnershipInvariant);
921 }
922 match effect {
923 RegistryEffect::None => restore_provider_handles(handles),
924 RegistryEffect::ArmWakeup { token, arm, .. } => {
925 if arm.replaces_existing() {
926 let replaced = take_detached_or_owned_handle(handles.take_wakeup(), |registry| {
927 registry.take_wakeup_handle(token.identity())
928 })?;
929 if let Some(replaced) = replaced {
930 clear_provider_handle(replaced);
931 }
932 }
933 restore_provider_handles(handles)?;
934 arm_wakeup(effect)
935 }
936 RegistryEffect::ClearCallbacks {
937 identity,
938 handles: selected,
939 } => {
940 if selected.includes_wakeup() {
941 let wakeup = take_detached_or_owned_handle(handles.take_wakeup(), |registry| {
942 registry.take_wakeup_handle(identity)
943 })?;
944 if let Some(wakeup) = wakeup {
945 clear_provider_handle(wakeup);
946 }
947 }
948 if selected.includes_work() {
949 let work = take_detached_or_owned_handle(handles.take_work(), |registry| {
950 registry.take_work_handle(identity)
951 })?;
952 if let Some(work) = work {
953 clear_provider_handle(work);
954 }
955 }
956 restore_provider_handles(handles)
957 }
958 RegistryEffect::DispatchWatchdog { successor, .. } => {
959 if let Some(wakeup) = handles.take_wakeup() {
960 clear_provider_handle(wakeup);
961 }
962 let replaced_work = take_detached_or_owned_handle(handles.take_work(), |registry| {
963 registry.take_work_handle(successor.identity())
964 })?;
965 if let Some(replaced_work) = replaced_work {
966 clear_provider_handle(replaced_work);
967 }
968 dispatch_watchdog_effect(effect)
969 }
970 }
971}
972
973fn take_detached_or_owned_handle(
974 detached: Option<ProviderHandle>,
975 take_owned: impl FnOnce(&mut TimerRegistry) -> Option<ProviderHandle>,
976) -> Result<Option<ProviderHandle>, TimerError> {
977 detached.map_or_else(
978 || with_registry_mut(|registry| Ok(take_owned(registry))),
979 |handle| Ok(Some(handle)),
980 )
981}
982
983fn arm_wakeup(effect: &RegistryEffect) -> Result<(), TimerError> {
984 let RegistryEffect::ArmWakeup {
985 token, delay_ns, ..
986 } = effect
987 else {
988 return Err(TimerError::OwnershipInvariant);
989 };
990 let task_token = token.clone();
991 let handle = platform::set_timer(Duration::from_nanos(*delay_ns), async move {
992 dispatch_wakeup(task_token).await;
993 });
994 if let Err((error, handle)) = install_provider_handle(token, handle) {
995 platform::clear_timer(handle);
996 return Err(error);
997 }
998 if let Err(error) = confirm_effect(effect) {
999 let handle = with_registry_mut(|registry| {
1000 registry
1001 .take_wakeup_handle(token.identity())
1002 .ok_or(TimerError::OwnershipInvariant)
1003 })?;
1004 clear_provider_handle(handle);
1005 return Err(error);
1006 }
1007 Ok(())
1008}
1009
1010fn dispatch_watchdog_effect(effect: &RegistryEffect) -> Result<(), TimerError> {
1011 let RegistryEffect::DispatchWatchdog {
1012 successor,
1013 successor_delay_ns,
1014 work,
1015 ..
1016 } = effect
1017 else {
1018 return Err(TimerError::OwnershipInvariant);
1019 };
1020 let successor_token = successor.clone();
1021 let successor_handle =
1022 platform::set_timer(Duration::from_nanos(*successor_delay_ns), async move {
1023 dispatch_watchdog_scheduler(&successor_token);
1024 });
1025 if let Err((error, handle)) = install_provider_handle(successor, successor_handle) {
1026 platform::clear_timer(handle);
1027 return Err(error);
1028 }
1029
1030 let work_token = work.clone();
1031 let work_handle = platform::set_timer(Duration::ZERO, async move {
1032 dispatch_watchdog_work(&work_token);
1033 });
1034 if let Err((error, handle)) = install_provider_handle(work, work_handle) {
1035 platform::clear_timer(handle);
1036 clear_entry_provider_handles(successor.identity())?;
1037 return Err(error);
1038 }
1039 if let Err(error) = confirm_effect(effect) {
1040 clear_entry_provider_handles(successor.identity())?;
1041 return Err(error);
1042 }
1043 Ok(())
1044}
1045
1046fn install_provider_handle(
1047 token: &CallbackToken,
1048 handle: TimerHandle,
1049) -> Result<(), (TimerError, TimerHandle)> {
1050 #[cfg(test)]
1051 if take_provider_install_fault() {
1052 return Err((TimerError::OwnershipInvariant, handle));
1053 }
1054 RUNTIME.with(|runtime| {
1055 let Ok(mut runtime) = runtime.try_borrow_mut() else {
1056 return Err((TimerError::RuntimeBusy, handle));
1057 };
1058 let Some(registry) = runtime.as_mut() else {
1059 return Err((TimerError::NotInitialized, handle));
1060 };
1061 match registry.install_provider_handle(token, handle) {
1062 Ok(()) => Ok(()),
1063 Err((error, handle)) => Err((TimerError::from(error), handle)),
1064 }
1065 })
1066}
1067
1068fn confirm_effect(effect: &RegistryEffect) -> Result<(), TimerError> {
1069 #[cfg(test)]
1070 if take_provider_confirmation_fault() {
1071 return Err(TimerError::OwnershipInvariant);
1072 }
1073 with_registry_mut(|registry| {
1074 registry
1075 .confirm_effect_applied(effect)
1076 .map_err(TimerError::from)
1077 })
1078}
1079
1080fn restore_provider_handles(mut handles: ProviderHandles) -> Result<(), TimerError> {
1081 let wakeup_failure = handles
1084 .take_wakeup()
1085 .and_then(|handle| restore_provider_handle(handle).err());
1086 let work_failure = handles
1087 .take_work()
1088 .and_then(|handle| restore_provider_handle(handle).err());
1089 wakeup_failure.or(work_failure).map_or(Ok(()), Err)
1090}
1091
1092fn restore_provider_handle(handle: ProviderHandle) -> Result<(), TimerError> {
1093 let (token, handle) = handle.into_parts();
1094 match install_provider_handle(&token, handle) {
1095 Ok(()) => Ok(()),
1096 Err((error, handle)) => {
1097 platform::clear_timer(handle);
1098 Err(error)
1099 }
1100 }
1101}
1102
1103fn clear_provider_handle(handle: ProviderHandle) {
1104 let (_, handle) = handle.into_parts();
1105 platform::clear_timer(handle);
1106}
1107
1108fn clear_provider_handles(mut handles: ProviderHandles) {
1109 if let Some(wakeup) = handles.take_wakeup() {
1110 clear_provider_handle(wakeup);
1111 }
1112 if let Some(work) = handles.take_work() {
1113 clear_provider_handle(work);
1114 }
1115}
1116
1117fn clear_entry_provider_handles(identity: &TimerIdentity) -> Result<(), TimerError> {
1118 let handles = with_registry_mut(|registry| {
1119 Ok(ProviderHandles::from_parts(
1120 registry.take_wakeup_handle(identity),
1121 registry.take_work_handle(identity),
1122 ))
1123 })?;
1124 clear_provider_handles(handles);
1125 Ok(())
1126}
1127
1128#[allow(clippy::future_not_send)] async fn dispatch_wakeup(token: CallbackToken) {
1130 match token.role() {
1131 CallbackRole::OrdinaryWork => dispatch_ordinary(token).await,
1132 CallbackRole::WatchdogScheduler => dispatch_watchdog_scheduler(&token),
1133 CallbackRole::WatchdogWork => {}
1134 }
1135}
1136
1137#[allow(clippy::future_not_send)] async fn dispatch_ordinary(token: CallbackToken) {
1139 let measurement = CallbackMeasurementStart::capture();
1140 let accepted = with_registry_mut(|registry| {
1141 registry.consume_provider_handle(&token);
1142 Ok(registry.begin_ordinary(&token))
1143 });
1144 match accepted {
1145 Ok(CallbackAcceptance::Accepted) => {}
1146 Ok(CallbackAcceptance::Stale) => return,
1147 Err(error) => trap_callback_failure("ordinary callback acceptance", &error),
1148 }
1149
1150 let callback = match with_registry(|registry| {
1151 registry.ordinary_callback(&token).map_err(TimerError::from)
1152 }) {
1153 Ok(callback) => callback,
1154 Err(TimerError::OwnershipInvariant) => {
1155 fail_ordinary_dispatch(&token);
1156 return;
1157 }
1158 Err(error) => trap_callback_failure("ordinary callback lookup", &error),
1159 };
1160 let future = {
1161 let Ok(mut callback) = callback.try_borrow_mut() else {
1162 fail_ordinary_dispatch(&token);
1163 return;
1164 };
1165 callback(token.clone())
1166 };
1167 let result = future.await;
1168 let transition = with_registry_mut(|registry| {
1169 registry
1170 .complete_ordinary(&token, platform::time_ns(), result)
1171 .map_err(TimerError::from)
1172 });
1173 let transition = transition
1174 .unwrap_or_else(|error| trap_callback_failure("ordinary callback completion", &error));
1175 finish_callback_transition(&token, transition, ProviderHandles::default());
1176 record_callback_measurements(&token, measurement.finish());
1177}
1178
1179fn fail_ordinary_dispatch(token: &CallbackToken) {
1180 let transition = with_registry_mut(|registry| {
1181 registry
1182 .complete_ordinary(
1183 token,
1184 platform::time_ns(),
1185 TimerRunResult::new(TimerCompletion::invariant_failure(0), TimerDirective::Stop),
1186 )
1187 .map_err(TimerError::from)
1188 });
1189 let transition = transition.unwrap_or_else(|error| {
1190 trap_callback_failure("ordinary invariant-failure completion", &error)
1191 });
1192 finish_callback_transition(token, transition, ProviderHandles::default());
1193}
1194
1195fn dispatch_watchdog_scheduler(token: &CallbackToken) {
1196 let measurement = CallbackMeasurementStart::capture();
1197 let transition = with_registry_mut(|registry| {
1198 registry.consume_provider_handle(token);
1199 Ok(registry.begin_watchdog_scheduler(token, platform::time_ns()))
1200 });
1201 let transition = transition
1202 .unwrap_or_else(|error| trap_callback_failure("watchdog scheduler transition", &error));
1203 let accepted = !matches!(transition.effect(), RegistryEffect::None);
1204 finish_callback_transition(token, transition, ProviderHandles::default());
1205 if accepted {
1206 record_callback_measurements(token, measurement.finish());
1207 }
1208}
1209
1210fn dispatch_watchdog_work(token: &CallbackToken) {
1211 let measurement = CallbackMeasurementStart::capture();
1212 let accepted = with_registry_mut(|registry| {
1213 registry.consume_provider_handle(token);
1214 Ok(registry.begin_watchdog_work(token))
1215 });
1216 match accepted {
1217 Ok(CallbackAcceptance::Accepted) => {}
1218 Ok(CallbackAcceptance::Stale) => return,
1219 Err(error) => trap_callback_failure("watchdog work acceptance", &error),
1220 }
1221
1222 let callback =
1223 match with_registry(|registry| registry.watchdog_callback(token).map_err(TimerError::from))
1224 {
1225 Ok(callback) => callback,
1226 Err(error) => trap_callback_failure("watchdog callback lookup", &error),
1227 };
1228 let result = {
1229 let Ok(mut callback) = callback.try_borrow_mut() else {
1230 trap_callback_failure(
1231 "watchdog callback ownership",
1232 &TimerError::OwnershipInvariant,
1233 );
1234 };
1235 callback(token.clone())
1236 };
1237 finish_watchdog_dispatch(token, result);
1238 record_callback_measurements(token, measurement.finish());
1239}
1240
1241fn finish_watchdog_dispatch(token: &CallbackToken, result: WatchdogRunResult) {
1242 let claim = RegistrationClaim::from_callback(token);
1243 let completed = with_registry_mut(|registry| {
1248 let handles = registry
1249 .take_provider_handles_for_claim(&claim)
1250 .map_err(TimerError::from)?;
1251 #[cfg(test)]
1252 {
1253 if take_watchdog_completion_fault() {
1254 return Err(TimerError::OwnershipInvariant);
1255 }
1256 }
1257 let transition = registry
1258 .complete_watchdog_work(token, platform::time_ns(), result)
1259 .map_err(TimerError::from)?;
1260 Ok((transition, handles))
1261 });
1262 let (transition, handles) =
1263 completed.unwrap_or_else(|error| trap_callback_failure("watchdog work completion", &error));
1264 finish_callback_transition(token, transition, handles);
1265}
1266
1267fn finish_callback_transition(
1268 token: &CallbackToken,
1269 transition: RegistryTransition,
1270 handles: ProviderHandles,
1271) {
1272 match finish_transition(transition, handles) {
1273 Ok(()) | Err(TimerError::ControlFailure(_)) => {}
1274 Err(
1275 error @ (TimerError::NotInitialized
1276 | TimerError::RuntimeBusy
1277 | TimerError::Register(_)
1278 | TimerError::Schedule(_)
1279 | TimerError::RegistrationExpired
1280 | TimerError::OwnershipInvariant
1281 | TimerError::ReconciliationConflict),
1282 ) => {
1283 if token.role() == CallbackRole::WatchdogWork {
1284 trap_callback_failure("watchdog provider-handle completion", &error);
1285 }
1286 fail_provider_binding(token).unwrap_or_else(|binding_error| {
1287 trap_callback_failure("provider-binding failure cleanup", &binding_error)
1288 });
1289 }
1290 }
1291}
1292
1293fn fail_provider_binding(token: &CallbackToken) -> Result<(), TimerError> {
1294 let claim = RegistrationClaim::from_callback(token);
1295 fail_claim_provider_binding(&claim)
1296}
1297
1298fn fail_claim_provider_binding(claim: &RegistrationClaim) -> Result<(), TimerError> {
1299 let failed = with_registry_mut(|registry| {
1300 registry
1301 .fail_registration(claim, TimerControlFailure::ProviderBindingFailed)
1302 .map_err(TimerError::from)
1303 });
1304 clear_provider_handles(failed?);
1305 Ok(())
1306}
1307
1308#[derive(Clone, Copy)]
1309struct CallbackMeasurementStart {
1310 instructions_before: u64,
1311 memory_start: platform::MemoryPages,
1312}
1313
1314impl CallbackMeasurementStart {
1315 fn capture() -> Self {
1316 let memory_start = platform::memory_pages();
1318 let instructions_before = platform::instruction_counter();
1319 Self {
1320 instructions_before,
1321 memory_start,
1322 }
1323 }
1324
1325 fn finish(self) -> CallbackMeasurement {
1326 let instructions = platform::instruction_counter().saturating_sub(self.instructions_before);
1328 let memory_end = platform::memory_pages();
1329 CallbackMeasurement {
1330 instructions,
1331 memory_start: self.memory_start,
1332 memory_end,
1333 }
1334 }
1335}
1336
1337#[derive(Clone, Copy)]
1338struct CallbackMeasurement {
1339 instructions: u64,
1340 memory_start: platform::MemoryPages,
1341 memory_end: platform::MemoryPages,
1342}
1343
1344fn record_callback_measurements(token: &CallbackToken, measurement: CallbackMeasurement) {
1345 with_registry_mut(|registry| {
1346 registry
1347 .record_callback_measurements(
1348 token,
1349 measurement.instructions,
1350 measurement.memory_start,
1351 measurement.memory_end,
1352 )
1353 .map_err(TimerError::from)
1354 })
1355 .unwrap_or_else(|error| trap_callback_failure("callback measurement accounting", &error));
1356}
1357
1358fn trap_callback_failure(context: &str, error: &TimerError) -> ! {
1359 platform::trap(&format!("ic-timers {context} failed: {error}"))
1360}
1361
1362fn with_registry<T>(
1363 operation: impl FnOnce(&TimerRegistry) -> Result<T, TimerError>,
1364) -> Result<T, TimerError> {
1365 RUNTIME.with(|runtime| {
1366 let runtime = runtime.try_borrow().map_err(|_| TimerError::RuntimeBusy)?;
1367 let registry = runtime.as_ref().ok_or(TimerError::NotInitialized)?;
1368 operation(registry)
1369 })
1370}
1371
1372fn with_registry_mut<T>(
1373 operation: impl FnOnce(&mut TimerRegistry) -> Result<T, TimerError>,
1374) -> Result<T, TimerError> {
1375 RUNTIME.with(|runtime| {
1376 let mut runtime = runtime
1377 .try_borrow_mut()
1378 .map_err(|_| TimerError::RuntimeBusy)?;
1379 let registry = runtime.as_mut().ok_or(TimerError::NotInitialized)?;
1380 operation(registry)
1381 })
1382}
1383
1384#[cfg(test)]
1385fn reset_for_test(now_ns: u64, canister_version: u64) {
1386 platform::reset(now_ns, canister_version);
1387 WATCHDOG_COMPLETION_FAULT.with(|fault| fault.set(false));
1388 PROVIDER_INSTALL_FAULT_AFTER.with(|fault| fault.set(None));
1389 PROVIDER_CONFIRMATION_FAULT.with(|fault| fault.set(false));
1390 RUNTIME.with(|runtime| {
1391 *runtime.borrow_mut() = None;
1392 });
1393}
1394
1395#[cfg(test)]
1396thread_local! {
1397 static WATCHDOG_COMPLETION_FAULT: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
1398 static PROVIDER_INSTALL_FAULT_AFTER: std::cell::Cell<Option<u64>> = const { std::cell::Cell::new(None) };
1399 static PROVIDER_CONFIRMATION_FAULT: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
1400}
1401
1402#[cfg(test)]
1403fn inject_watchdog_completion_fault() {
1404 WATCHDOG_COMPLETION_FAULT.with(|fault| fault.set(true));
1405}
1406
1407#[cfg(test)]
1408fn take_watchdog_completion_fault() -> bool {
1409 WATCHDOG_COMPLETION_FAULT.with(|fault| fault.replace(false))
1410}
1411
1412#[cfg(test)]
1413fn inject_provider_install_fault() {
1414 inject_provider_install_fault_after(0);
1415}
1416
1417#[cfg(test)]
1418fn inject_provider_install_fault_after(successful_installs: u64) {
1419 PROVIDER_INSTALL_FAULT_AFTER.with(|fault| fault.set(Some(successful_installs)));
1420}
1421
1422#[cfg(test)]
1423fn take_provider_install_fault() -> bool {
1424 PROVIDER_INSTALL_FAULT_AFTER.with(|fault| match fault.get() {
1425 Some(0) => {
1426 fault.set(None);
1427 true
1428 }
1429 Some(remaining) => {
1430 fault.set(Some(remaining - 1));
1431 false
1432 }
1433 None => false,
1434 })
1435}
1436
1437#[cfg(test)]
1438fn inject_provider_confirmation_fault() {
1439 PROVIDER_CONFIRMATION_FAULT.with(|fault| fault.set(true));
1440}
1441
1442#[cfg(test)]
1443fn take_provider_confirmation_fault() -> bool {
1444 PROVIDER_CONFIRMATION_FAULT.with(|fault| fault.replace(false))
1445}
1446
1447#[cfg(test)]
1448mod tests;