1use crate::{
4 callback::{AfterCompletionRunResult, OnceRunResult, OrdinaryRunResult},
5 platform::{self, TimerHandle},
6 registry::{
7 CallbackRole, CallbackToken, OrdinaryCallback, ProviderHandle, ProviderHandles,
8 RegisterError, RegistrationClaim, RegistryEffect, RegistryError, RegistryTransition,
9 TimerRegistry, WakeupToken, WatchdogCallback,
10 },
11 schedule::{OrdinaryDirective, ScheduleError, TimerCadence, TimerSchedule},
12 snapshot::{
13 DeclarationLifetime, MemoryPageExtent, TimerCompletion, TimerControlFailure, TimerEpoch,
14 TimerIdentity, TimerInventorySnapshot, TimerPolicy, TimerSnapshot, WatchdogRunResult,
15 },
16};
17#[cfg(test)]
18use std::time::Duration;
19use std::{cell::RefCell, future::Future, rc::Rc};
20use thiserror::Error;
21
22mod effects;
23#[cfg(test)]
24use effects::{
25 apply_effect, bind_provider_handle, clear_entry_provider_handles, clear_provider_handle,
26 restore_provider_handles,
27};
28use effects::{clear_provider_handles, finish_transition};
29
30thread_local! {
31 static RUNTIME: RefCell<Option<TimerRegistry>> = const { RefCell::new(None) };
32}
33
34#[derive(Clone, Copy, Debug, Eq, PartialEq)]
47pub enum TimerStart {
48 Retained {
50 initial: Option<TimerSchedule>,
52 },
53 RemoveWhenStopped {
55 initial: TimerSchedule,
57 },
58}
59
60impl TimerStart {
61 const fn into_parts(self) -> (DeclarationLifetime, Option<TimerSchedule>) {
62 match self {
63 Self::Retained { initial } => (DeclarationLifetime::Retained, initial),
64 Self::RemoveWhenStopped { initial } => {
65 (DeclarationLifetime::RemoveWhenStopped, Some(initial))
66 }
67 }
68 }
69}
70
71#[non_exhaustive]
73#[derive(Debug, Error)]
74pub enum TimerError {
75 #[error("timer runtime is not initialized")]
77 NotInitialized,
78 #[error("timer runtime is already borrowed")]
80 RuntimeBusy,
81 #[error(transparent)]
83 Register(#[from] RegisterError),
84 #[error(transparent)]
86 Schedule(#[from] ScheduleError),
87 #[error("timer registration is no longer authoritative")]
89 RegistrationExpired,
90 #[error("timer control failed: {0:?}")]
92 ControlFailure(TimerControlFailure),
93 #[error("timer runtime ownership invariant failed")]
95 OwnershipInvariant,
96 #[error("timer backend pending capacity exhausted")]
98 BackendCapacity,
99 #[error("timer backend callback identities exhausted")]
101 BackendIdentityExhausted,
102 #[error("timer lifecycle reconciliation conflicts with the canonical declaration")]
104 ReconciliationConflict,
105}
106
107impl From<platform::AdmissionError> for TimerError {
108 fn from(error: platform::AdmissionError) -> Self {
109 match error {
110 platform::AdmissionError::Capacity => Self::BackendCapacity,
111 platform::AdmissionError::IdentityExhausted => Self::BackendIdentityExhausted,
112 platform::AdmissionError::DeadlineOverflow => {
113 Self::Schedule(ScheduleError::DeadlineOverflow)
114 }
115 }
116 }
117}
118
119impl From<RegistryError> for TimerError {
120 fn from(value: RegistryError) -> Self {
121 match value {
122 RegistryError::UnknownRegistration
123 | RegistryError::StaleRegistration
124 | RegistryError::StaleCallback => Self::RegistrationExpired,
125 RegistryError::PolicyMismatch { .. } | RegistryError::ProviderHandleAlreadyOwned => {
126 Self::OwnershipInvariant
127 }
128 RegistryError::Schedule(error) => Self::Schedule(error),
129 }
130 }
131}
132
133pub fn initialize_runtime() -> Result<TimerEpoch, TimerError> {
138 RUNTIME.with(|runtime| {
139 let mut runtime = runtime
140 .try_borrow_mut()
141 .map_err(|_| TimerError::RuntimeBusy)?;
142 let registry = runtime.get_or_insert_with(|| {
143 TimerRegistry::new(TimerEpoch::new(
144 platform::canister_version(),
145 platform::time_ns(),
146 ))
147 });
148 Ok(registry.epoch())
149 })
150}
151
152pub struct OnceContext {
158 token: CallbackToken,
159}
160
161impl OnceContext {
162 const fn new(token: CallbackToken) -> Self {
163 Self { token }
164 }
165
166 #[must_use]
168 pub const fn identity(&self) -> &TimerIdentity {
169 self.token.identity()
170 }
171
172 pub fn ensure_scheduled(&self, schedule: TimerSchedule) -> Result<(), TimerError> {
177 ensure_once_claim(self.token.claim(), Some(&self.token), schedule)
178 }
179
180 pub fn reconcile_schedule(&self, schedule: Option<TimerSchedule>) -> Result<(), TimerError> {
190 reconcile_ordinary_claim(self.token.claim(), Some(&self.token), schedule)
191 }
192
193 pub fn cancel(&self) -> Result<(), TimerError> {
201 cancel_claim(self.token.claim(), Some(&self.token))
202 }
203}
204
205pub struct AfterCompletionContext {
212 token: CallbackToken,
213}
214
215impl AfterCompletionContext {
216 const fn new(token: CallbackToken) -> Self {
217 Self { token }
218 }
219
220 #[must_use]
222 pub const fn identity(&self) -> &TimerIdentity {
223 self.token.identity()
224 }
225
226 pub fn ensure_scheduled(&self) -> Result<(), TimerError> {
228 ensure_recurring_claim(self.token.claim(), Some(&self.token))
229 }
230
231 pub fn reconcile_schedule(&self, schedule: Option<TimerSchedule>) -> Result<(), TimerError> {
242 reconcile_ordinary_claim(self.token.claim(), Some(&self.token), schedule)
243 }
244
245 pub fn cancel(&self) -> Result<(), TimerError> {
253 cancel_claim(self.token.claim(), Some(&self.token))
254 }
255}
256
257pub struct WatchdogContext {
262 token: CallbackToken,
263}
264
265impl WatchdogContext {
266 const fn new(token: CallbackToken) -> Self {
267 Self { token }
268 }
269
270 #[must_use]
272 pub const fn identity(&self) -> &TimerIdentity {
273 self.token.identity()
274 }
275
276 pub fn ensure_scheduled(&self) -> Result<(), TimerError> {
282 ensure_recurring_claim(self.token.claim(), Some(&self.token))
283 }
284
285 pub fn ensure_scheduled_immediately(&self) -> Result<(), TimerError> {
292 ensure_watchdog_immediately_claim(self.token.claim(), Some(&self.token))
293 }
294
295 pub fn reconcile_schedule(&self, schedule: Option<TimerSchedule>) -> Result<(), TimerError> {
301 reconcile_watchdog_claim(self.token.claim(), Some(&self.token), schedule)
302 }
303
304 pub fn cancel(&self) -> Result<(), TimerError> {
312 cancel_claim(self.token.claim(), Some(&self.token))
313 }
314}
315
316#[must_use = "retain the registration claim so the timer remains controllable"]
323pub struct OnceRegistration {
324 claim: RegistrationClaim,
325}
326
327impl OnceRegistration {
328 #[must_use]
330 pub const fn identity(&self) -> &TimerIdentity {
331 self.claim.identity()
332 }
333
334 pub fn has_armed_wakeup(&self) -> Result<bool, TimerError> {
340 has_armed_wakeup_claim(&self.claim)
341 }
342
343 pub fn ensure_scheduled(&self, schedule: TimerSchedule) -> Result<(), TimerError> {
345 ensure_once_claim(&self.claim, None, schedule)
346 }
347
348 pub fn reconcile_schedule(&self, schedule: Option<TimerSchedule>) -> Result<(), TimerError> {
358 reconcile_ordinary_claim(&self.claim, None, schedule)
359 }
360
361 pub fn cancel(&self) -> Result<(), TimerError> {
368 cancel_claim(&self.claim, None)
369 }
370
371 pub fn unregister(self) -> Result<(), TimerError> {
376 unregister_claim(&self.claim)
377 }
378}
379
380#[must_use = "retain the registration claim so the timer remains controllable"]
388pub struct AfterCompletionRegistration {
389 claim: RegistrationClaim,
390}
391
392#[must_use = "retain the registration claim so the timer remains controllable"]
399pub struct WatchdogRegistration {
400 claim: RegistrationClaim,
401}
402
403#[derive(Clone, Copy, Debug, Eq, PartialEq)]
405pub enum TimerReconcileState {
406 Inactive,
408 Scheduled,
410}
411
412#[derive(Clone, Copy, Debug, Eq, PartialEq)]
414pub enum WatchdogReconcileState {
415 Inactive,
417 Scheduled,
419 ScheduledImmediately,
424 ScheduledAt(u64),
426}
427
428impl WatchdogRegistration {
429 #[must_use]
431 pub const fn identity(&self) -> &TimerIdentity {
432 self.claim.identity()
433 }
434
435 pub fn has_armed_wakeup(&self) -> Result<bool, TimerError> {
442 has_armed_wakeup_claim(&self.claim)
443 }
444
445 pub fn ensure_scheduled(&self) -> Result<(), TimerError> {
449 ensure_recurring_claim(&self.claim, None)
450 }
451
452 pub fn ensure_scheduled_immediately(&self) -> Result<(), TimerError> {
460 ensure_watchdog_immediately_claim(&self.claim, None)
461 }
462
463 pub fn reconcile_schedule(&self, schedule: Option<TimerSchedule>) -> Result<(), TimerError> {
471 reconcile_watchdog_claim(&self.claim, None, schedule)
472 }
473
474 pub fn cancel(&self) -> Result<(), TimerError> {
482 cancel_claim(&self.claim, None)
483 }
484
485 pub fn unregister(self) -> Result<(), TimerError> {
490 unregister_claim(&self.claim)
491 }
492}
493
494impl AfterCompletionRegistration {
495 #[must_use]
497 pub const fn identity(&self) -> &TimerIdentity {
498 self.claim.identity()
499 }
500
501 pub fn has_armed_wakeup(&self) -> Result<bool, TimerError> {
507 has_armed_wakeup_claim(&self.claim)
508 }
509
510 pub fn ensure_scheduled(&self) -> Result<(), TimerError> {
513 ensure_recurring_claim(&self.claim, None)
514 }
515
516 pub fn reconcile_schedule(&self, schedule: Option<TimerSchedule>) -> Result<(), TimerError> {
525 reconcile_ordinary_claim(&self.claim, None, schedule)
526 }
527
528 pub fn cancel(&self) -> Result<(), TimerError> {
535 cancel_claim(&self.claim, None)
536 }
537
538 pub fn unregister(self) -> Result<(), TimerError> {
543 unregister_claim(&self.claim)
544 }
545}
546
547pub fn register_once<F, Fut>(
567 identity: TimerIdentity,
568 start: TimerStart,
569 callback: F,
570) -> Result<OnceRegistration, TimerError>
571where
572 F: FnMut(OnceContext) -> Fut + 'static,
573 Fut: Future<Output = OnceRunResult> + 'static,
574{
575 let callback = erase_ordinary_callback(callback, OnceContext::new, OrdinaryRunResult::from);
576 let claim = register_initial(start, |registry, lifetime| {
577 registry.register_once_with_callback(identity, lifetime, Rc::clone(&callback))
578 });
579 drop(callback);
581 Ok(OnceRegistration { claim: claim? })
582}
583
584pub fn register_after_completion<F, Fut>(
602 identity: TimerIdentity,
603 cadence: TimerCadence,
604 start: TimerStart,
605 callback: F,
606) -> Result<AfterCompletionRegistration, TimerError>
607where
608 F: FnMut(AfterCompletionContext) -> Fut + 'static,
609 Fut: Future<Output = AfterCompletionRunResult> + 'static,
610{
611 let callback = erase_ordinary_callback(
612 callback,
613 AfterCompletionContext::new,
614 OrdinaryRunResult::from,
615 );
616 let claim = register_initial(start, |registry, lifetime| {
617 registry.register_after_completion_with_callback(
618 identity,
619 cadence,
620 lifetime,
621 Rc::clone(&callback),
622 )
623 });
624 drop(callback);
626 Ok(AfterCompletionRegistration { claim: claim? })
627}
628
629pub fn register_watchdog<F>(
640 identity: TimerIdentity,
641 cadence: TimerCadence,
642 start: TimerStart,
643 callback: F,
644) -> Result<WatchdogRegistration, TimerError>
645where
646 F: FnMut(WatchdogContext) -> WatchdogRunResult + 'static,
647{
648 let mut callback = callback;
649 let callback: WatchdogCallback = Rc::new(RefCell::new(move |token| {
650 callback(WatchdogContext::new(token))
651 }));
652 let claim = register_initial(start, |registry, lifetime| {
653 registry.register_watchdog_with_callback(identity, cadence, lifetime, Rc::clone(&callback))
654 });
655 drop(callback);
657 Ok(WatchdogRegistration { claim: claim? })
658}
659
660fn register_initial(
663 start: TimerStart,
664 register: impl FnOnce(
665 &mut TimerRegistry,
666 DeclarationLifetime,
667 ) -> Result<RegistrationClaim, RegisterError>,
668) -> Result<RegistrationClaim, TimerError> {
669 let (lifetime, initial) = start.into_parts();
670 register_initialized(
671 |registry| {
672 let now_ns = platform::time_ns();
673 let requested = initial
674 .map(|schedule| schedule.resolve(now_ns))
675 .transpose()?;
676 Ok((register(registry, lifetime)?, (requested, now_ns)))
677 },
678 |registry, claim, (requested, now_ns)| {
679 requested
680 .map(|schedule| {
681 registry
682 .transition_claim(claim, |registry| {
683 registry.reconcile_initial(claim, now_ns, schedule)
684 })
685 .map_err(TimerError::from)
686 })
687 .transpose()
688 },
689 )
690}
691
692fn register_initialized<Prepared>(
695 register: impl FnOnce(&mut TimerRegistry) -> Result<(RegistrationClaim, Prepared), TimerError>,
696 initialize: impl FnOnce(
697 &mut TimerRegistry,
698 &RegistrationClaim,
699 Prepared,
700 ) -> Result<Option<(RegistryTransition, ProviderHandles)>, TimerError>,
701) -> Result<RegistrationClaim, TimerError> {
702 let (claim, initialized) = with_registry_mut(|registry| {
703 let (claim, prepared) = register(registry)?;
704 let initialized = initialize(registry, &claim, prepared);
705 Ok((claim, initialized))
706 })?;
707 let applied = initialized.and_then(|transition| {
708 transition.map_or(Ok(()), |(transition, handles)| {
709 finish_transition(transition, handles)
710 })
711 });
712 applied.map_err(|error| cleanup_creation_error(&claim, error))?;
713 Ok(claim)
714}
715
716fn cleanup_creation_error(claim: &RegistrationClaim, error: TimerError) -> TimerError {
717 match unregister_claim(claim) {
718 Ok(()) | Err(TimerError::RegistrationExpired) => error,
719 Err(cleanup_error) => cleanup_error,
720 }
721}
722
723pub fn reconcile_once<F, Fut>(
746 registration: &mut Option<OnceRegistration>,
747 identity: &TimerIdentity,
748 desired: Option<TimerSchedule>,
749 callback: F,
750) -> Result<(), TimerError>
751where
752 F: FnMut(OnceContext) -> Fut + 'static,
753 Fut: Future<Output = OnceRunResult> + 'static,
754{
755 reconcile_registration(
756 registration,
757 identity,
758 TimerPolicy::Once,
759 |current: &OnceRegistration| ¤t.claim,
760 || {
761 let callback =
762 erase_ordinary_callback(callback, OnceContext::new, OrdinaryRunResult::from);
763 let claim = register_initialized(
764 |registry| {
765 Ok((
766 registry.register_once_with_callback(
767 identity.clone(),
768 DeclarationLifetime::Retained,
769 Rc::clone(&callback),
770 )?,
771 (),
772 ))
773 },
774 |registry, claim, ()| {
775 registry
776 .transition_claim(claim, |registry| {
777 registry.reconcile_ordinary(claim, platform::time_ns(), desired)
778 })
779 .map(Some)
780 .map_err(TimerError::from)
781 },
782 );
783 drop(callback);
784 Ok(OnceRegistration { claim: claim? })
785 },
786 |registration| registration.reconcile_schedule(desired),
787 )
788}
789
790pub fn reconcile_after_completion<F, Fut>(
815 registration: &mut Option<AfterCompletionRegistration>,
816 identity: &TimerIdentity,
817 cadence: TimerCadence,
818 desired: TimerReconcileState,
819 callback: F,
820) -> Result<(), TimerError>
821where
822 F: FnMut(AfterCompletionContext) -> Fut + 'static,
823 Fut: Future<Output = AfterCompletionRunResult> + 'static,
824{
825 reconcile_registration(
826 registration,
827 identity,
828 TimerPolicy::AfterCompletion { cadence },
829 |current: &AfterCompletionRegistration| ¤t.claim,
830 || {
831 let callback = erase_ordinary_callback(
832 callback,
833 AfterCompletionContext::new,
834 OrdinaryRunResult::from,
835 );
836 let claim = register_initialized(
837 |registry| {
838 Ok((
839 registry.register_after_completion_with_callback(
840 identity.clone(),
841 cadence,
842 DeclarationLifetime::Retained,
843 Rc::clone(&callback),
844 )?,
845 (),
846 ))
847 },
848 |registry, claim, ()| {
849 registry
850 .transition_claim(claim, |registry| match desired {
851 TimerReconcileState::Inactive => registry.cancel(claim),
852 TimerReconcileState::Scheduled => {
853 registry.ensure_recurring(claim, platform::time_ns())
854 }
855 })
856 .map(Some)
857 .map_err(TimerError::from)
858 },
859 );
860 drop(callback);
861 Ok(AfterCompletionRegistration { claim: claim? })
862 },
863 |registration| match desired {
864 TimerReconcileState::Inactive => registration.cancel(),
865 TimerReconcileState::Scheduled => registration.ensure_scheduled(),
866 },
867 )
868}
869
870pub fn reconcile_watchdog<F>(
883 registration: &mut Option<WatchdogRegistration>,
884 identity: &TimerIdentity,
885 cadence: TimerCadence,
886 desired: WatchdogReconcileState,
887 callback: F,
888) -> Result<(), TimerError>
889where
890 F: FnMut(WatchdogContext) -> WatchdogRunResult + 'static,
891{
892 reconcile_registration(
893 registration,
894 identity,
895 TimerPolicy::Watchdog { cadence },
896 |current: &WatchdogRegistration| ¤t.claim,
897 || {
898 let mut callback = callback;
899 let callback: WatchdogCallback = Rc::new(RefCell::new(move |token| {
900 callback(WatchdogContext::new(token))
901 }));
902 let claim = register_initialized(
903 |registry| {
904 Ok((
905 registry.register_watchdog_with_callback(
906 identity.clone(),
907 cadence,
908 DeclarationLifetime::Retained,
909 Rc::clone(&callback),
910 )?,
911 (),
912 ))
913 },
914 |registry, claim, ()| {
915 registry
916 .transition_claim(claim, |registry| match desired {
917 WatchdogReconcileState::Inactive => registry.cancel(claim),
918 WatchdogReconcileState::Scheduled => {
919 registry.ensure_recurring(claim, platform::time_ns())
920 }
921 WatchdogReconcileState::ScheduledImmediately => {
922 registry.ensure_watchdog_immediately(claim, platform::time_ns())
923 }
924 WatchdogReconcileState::ScheduledAt(deadline_ns) => registry
925 .reconcile_watchdog_schedule(
926 claim,
927 platform::time_ns(),
928 Some(TimerSchedule::At(deadline_ns)),
929 ),
930 })
931 .map(Some)
932 .map_err(TimerError::from)
933 },
934 );
935 drop(callback);
936 Ok(WatchdogRegistration { claim: claim? })
937 },
938 |registration| match desired {
939 WatchdogReconcileState::Inactive => registration.cancel(),
940 WatchdogReconcileState::Scheduled => registration.ensure_scheduled(),
941 WatchdogReconcileState::ScheduledImmediately => {
942 registration.ensure_scheduled_immediately()
943 }
944 WatchdogReconcileState::ScheduledAt(deadline_ns) => {
945 registration.reconcile_schedule(Some(TimerSchedule::At(deadline_ns)))
946 }
947 },
948 )
949}
950
951fn reconcile_registration<Registration>(
952 registration: &mut Option<Registration>,
953 identity: &TimerIdentity,
954 policy: TimerPolicy,
955 claim: impl FnOnce(&Registration) -> &RegistrationClaim,
956 register: impl FnOnce() -> Result<Registration, TimerError>,
957 reconcile: impl FnOnce(&Registration) -> Result<(), TimerError>,
958) -> Result<(), TimerError> {
959 let Some(current) = registration.as_ref() else {
960 *registration = Some(register()?);
961 return Ok(());
962 };
963 let claim = claim(current);
964 if claim.identity() != identity {
965 return Err(TimerError::ReconciliationConflict);
966 }
967 with_registry(|registry| {
968 registry
969 .declaration_matches(claim, policy, DeclarationLifetime::Retained)
970 .map_err(TimerError::from)?
971 .then_some(())
972 .ok_or(TimerError::ReconciliationConflict)
973 })?;
974 drop(register);
977 reconcile(current)
978}
979
980pub fn timer_snapshot(identity: &TimerIdentity) -> Result<Option<TimerSnapshot>, TimerError> {
982 with_registry(|registry| Ok(registry.snapshot(identity)))
983}
984
985pub fn timer_inventory() -> Result<TimerInventorySnapshot, TimerError> {
990 with_registry(|registry| Ok(registry.inventory()))
991}
992
993pub fn consecutive_expected_failures(identity: &TimerIdentity) -> Result<Option<u64>, TimerError> {
995 with_registry(|registry| Ok(registry.consecutive_expected_failures(identity)))
996}
997
998fn has_armed_wakeup_claim(claim: &RegistrationClaim) -> Result<bool, TimerError> {
999 with_registry(|registry| registry.has_armed_wakeup(claim).map_err(TimerError::from))
1000}
1001
1002fn erase_ordinary_callback<Context: 'static, Output: 'static, F, Fut>(
1003 mut callback: F,
1004 context: fn(CallbackToken) -> Context,
1005 erase_result: fn(Output) -> OrdinaryRunResult,
1006) -> OrdinaryCallback
1007where
1008 F: FnMut(Context) -> Fut + 'static,
1009 Fut: Future<Output = Output> + 'static,
1010{
1011 Rc::new(RefCell::new(move |token| {
1012 let future = callback(context(token));
1013 Box::pin(async move { erase_result(future.await) }) as _
1014 }))
1015}
1016
1017fn apply_claim_transition(
1018 claim: &RegistrationClaim,
1019 context: Option<&CallbackToken>,
1020 operation: impl FnOnce(&mut TimerRegistry) -> Result<RegistryTransition, RegistryError>,
1021) -> Result<(), TimerError> {
1022 let (transition, handles) = with_registry_mut(|registry| {
1023 if let Some(token) = context {
1024 registry
1025 .validate_running_context(token)
1026 .map_err(TimerError::from)?;
1027 }
1028 registry
1029 .transition_claim(claim, operation)
1030 .map_err(TimerError::from)
1031 })?;
1032 match finish_transition(transition, handles) {
1033 result @ (Ok(()) | Err(TimerError::ControlFailure(_))) => result,
1034 Err(error) => retire_failed_claim(claim, error),
1035 }
1036}
1037
1038fn ensure_once_claim(
1039 claim: &RegistrationClaim,
1040 context: Option<&CallbackToken>,
1041 schedule: TimerSchedule,
1042) -> Result<(), TimerError> {
1043 apply_claim_transition(claim, context, |registry| {
1044 registry.ensure_once(claim, platform::time_ns(), schedule)
1045 })
1046}
1047
1048fn reconcile_ordinary_claim(
1049 claim: &RegistrationClaim,
1050 context: Option<&CallbackToken>,
1051 schedule: Option<TimerSchedule>,
1052) -> Result<(), TimerError> {
1053 apply_claim_transition(claim, context, |registry| {
1054 registry.reconcile_ordinary(claim, platform::time_ns(), schedule)
1055 })
1056}
1057
1058fn ensure_recurring_claim(
1059 claim: &RegistrationClaim,
1060 context: Option<&CallbackToken>,
1061) -> Result<(), TimerError> {
1062 apply_claim_transition(claim, context, |registry| {
1063 registry.ensure_recurring(claim, platform::time_ns())
1064 })
1065}
1066
1067fn ensure_watchdog_immediately_claim(
1068 claim: &RegistrationClaim,
1069 context: Option<&CallbackToken>,
1070) -> Result<(), TimerError> {
1071 apply_claim_transition(claim, context, |registry| {
1072 registry.ensure_watchdog_immediately(claim, platform::time_ns())
1073 })
1074}
1075
1076fn reconcile_watchdog_claim(
1077 claim: &RegistrationClaim,
1078 context: Option<&CallbackToken>,
1079 schedule: Option<TimerSchedule>,
1080) -> Result<(), TimerError> {
1081 apply_claim_transition(claim, context, |registry| {
1082 registry.reconcile_watchdog_schedule(claim, platform::time_ns(), schedule)
1083 })
1084}
1085
1086fn cancel_claim(
1087 claim: &RegistrationClaim,
1088 context: Option<&CallbackToken>,
1089) -> Result<(), TimerError> {
1090 apply_claim_transition(claim, context, |registry| registry.cancel(claim))
1091}
1092
1093fn unregister_claim(claim: &RegistrationClaim) -> Result<(), TimerError> {
1094 apply_claim_transition(claim, None, |registry| registry.unregister(claim))
1095}
1096
1097fn retire_failed_claim(claim: &RegistrationClaim, error: TimerError) -> Result<(), TimerError> {
1098 match fail_claim_provider_binding(claim) {
1099 Ok(()) | Err(TimerError::RegistrationExpired) => Err(error),
1100 Err(cleanup_error) => Err(cleanup_error),
1101 }
1102}
1103
1104#[expect(
1105 clippy::future_not_send,
1106 reason = "IC callbacks and canister-local state are single-threaded."
1107)]
1108fn dispatch_wakeup(token: WakeupToken) -> impl Future<Output = ()> {
1109 let delivery = match token {
1112 WakeupToken::Ordinary(token) => WakeupDelivery::Ordinary(OrdinaryDelivery {
1113 token,
1114 completed: false,
1115 }),
1116 WakeupToken::Scheduler(token) => WakeupDelivery::Scheduler(token),
1117 };
1118 async move {
1119 match delivery {
1120 WakeupDelivery::Ordinary(mut delivery) => {
1121 dispatch_ordinary(&delivery.token).await;
1122 delivery.completed = true;
1123 drop(delivery);
1125 }
1126 WakeupDelivery::Scheduler(token) => dispatch_watchdog_scheduler(&token),
1127 }
1128 }
1129}
1130
1131enum WakeupDelivery {
1132 Ordinary(OrdinaryDelivery),
1133 Scheduler(CallbackToken),
1134}
1135
1136struct OrdinaryDelivery {
1137 token: CallbackToken,
1138 completed: bool,
1139}
1140
1141impl Drop for OrdinaryDelivery {
1142 fn drop(&mut self) {
1143 if self.completed {
1144 return;
1145 }
1146 let removed = RUNTIME.with(|runtime| {
1147 let Ok(mut runtime) = runtime.try_borrow_mut() else {
1148 let runtime = runtime.try_borrow().map_err(|_| TimerError::RuntimeBusy)?;
1152 return if runtime
1153 .as_ref()
1154 .is_some_and(|registry| registry.owns_ordinary_delivery(&self.token))
1155 {
1156 Err(TimerError::RuntimeBusy)
1157 } else {
1158 Ok(None)
1159 };
1160 };
1161 Ok(runtime
1162 .as_mut()
1163 .and_then(|registry| registry.abandon_ordinary(&self.token, platform::time_ns())))
1164 });
1165 let removed = removed
1166 .unwrap_or_else(|error| trap_callback_failure("ordinary delivery abandonment", &error));
1167 drop(removed);
1168 }
1169}
1170
1171#[expect(
1172 clippy::future_not_send,
1173 reason = "IC callbacks and canister-local state are single-threaded."
1174)]
1175async fn dispatch_ordinary(token: &CallbackToken) {
1176 let measurement = CallbackMeasurementStart::capture();
1177 let callback = with_registry_mut(|registry| Ok(registry.begin_ordinary(token)));
1178 let callback = match callback {
1179 Ok(Some(callback)) => callback,
1180 Ok(None) => return,
1181 Err(error) => trap_callback_failure("ordinary callback acceptance", &error),
1182 };
1183 let future = {
1184 let Ok(mut callback) = callback.try_borrow_mut() else {
1185 finish_ordinary_callback(
1186 token,
1187 OrdinaryRunResult::new(
1188 TimerCompletion::invariant_failure(0),
1189 OrdinaryDirective::Stop,
1190 ),
1191 );
1192 return;
1193 };
1194 callback(token.clone())
1195 };
1196 let result = future.await;
1197 finish_ordinary_callback(token, result);
1198 measurement.record(token);
1199}
1200
1201fn finish_ordinary_callback(token: &CallbackToken, result: OrdinaryRunResult) {
1202 let transition = with_registry_mut(|registry| {
1203 registry
1204 .complete_ordinary(token, platform::time_ns(), result)
1205 .map_err(TimerError::from)
1206 });
1207 let transition = transition
1208 .unwrap_or_else(|error| trap_callback_failure("ordinary callback completion", &error));
1209 finish_callback_transition(token, transition, ProviderHandles::default());
1210}
1211
1212fn dispatch_watchdog_scheduler(token: &CallbackToken) {
1213 let measurement = CallbackMeasurementStart::capture();
1214 let dispatched = with_registry_mut(|registry| {
1215 let handles = registry
1216 .take_watchdog_scheduler_handles(token)
1217 .map_err(TimerError::from)?;
1218 Ok((
1219 registry.begin_watchdog_scheduler(token, platform::time_ns()),
1220 handles,
1221 ))
1222 });
1223 let (transition, handles) = dispatched
1224 .unwrap_or_else(|error| trap_callback_failure("watchdog scheduler transition", &error));
1225 let accepted = !matches!(transition.effect(), RegistryEffect::None);
1226 finish_callback_transition(token, transition, handles);
1227 if accepted {
1228 measurement.record(token);
1229 }
1230}
1231
1232fn dispatch_watchdog_work(token: &CallbackToken) {
1233 let measurement = CallbackMeasurementStart::capture();
1234 let callback = with_registry_mut(|registry| Ok(registry.begin_watchdog_work(token)));
1235 let callback = match callback {
1236 Ok(Some(callback)) => callback,
1237 Ok(None) => return,
1238 Err(error) => trap_callback_failure("watchdog work acceptance", &error),
1239 };
1240 let result = {
1241 let Ok(mut callback) = callback.try_borrow_mut() else {
1242 trap_callback_failure(
1243 "watchdog callback ownership",
1244 &TimerError::OwnershipInvariant,
1245 );
1246 };
1247 callback(token.clone())
1248 };
1249 finish_watchdog_dispatch(token, result);
1250 measurement.record(token);
1251}
1252
1253fn finish_watchdog_dispatch(token: &CallbackToken, result: WatchdogRunResult) {
1254 let completed = with_registry_mut(|registry| {
1259 let handles = registry
1260 .take_provider_handles_for_claim(token.claim())
1261 .map_err(TimerError::from)?;
1262 #[cfg(test)]
1263 {
1264 if take_watchdog_completion_fault() {
1265 return Err(TimerError::OwnershipInvariant);
1266 }
1267 }
1268 let transition = registry
1269 .complete_watchdog_work(token, platform::time_ns(), result)
1270 .map_err(TimerError::from)?;
1271 Ok((transition, handles))
1272 });
1273 let (transition, handles) =
1274 completed.unwrap_or_else(|error| trap_callback_failure("watchdog work completion", &error));
1275 finish_callback_transition(token, transition, handles);
1276}
1277
1278fn finish_callback_transition(
1279 token: &CallbackToken,
1280 transition: RegistryTransition,
1281 handles: ProviderHandles,
1282) {
1283 match finish_transition(transition, handles) {
1286 Ok(()) | Err(TimerError::ControlFailure(_)) => {}
1287 Err(error) => {
1288 if token.role() == CallbackRole::WatchdogWork {
1289 trap_callback_failure("watchdog provider-handle completion", &error);
1290 }
1291 fail_claim_provider_binding(token.claim()).unwrap_or_else(|binding_error| {
1292 trap_callback_failure("provider-binding failure cleanup", &binding_error)
1293 });
1294 }
1295 }
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 let (handles, transition) = failed?;
1305 clear_provider_handles(handles);
1306 drop(transition);
1307 Ok(())
1308}
1309
1310#[derive(Clone, Copy)]
1311struct CallbackMeasurementStart {
1312 instructions_before: u64,
1313 memory_start: MemoryPageExtent,
1314}
1315
1316impl CallbackMeasurementStart {
1317 fn capture() -> Self {
1318 let memory_start = platform::memory_pages();
1320 let instructions_before = platform::instruction_counter();
1321 Self {
1322 instructions_before,
1323 memory_start,
1324 }
1325 }
1326
1327 fn record(self, token: &CallbackToken) {
1328 let instructions = platform::instruction_counter().saturating_sub(self.instructions_before);
1330 let memory_end = platform::memory_pages();
1331 with_registry_mut(|registry| {
1332 registry
1333 .record_callback_measurements(token, instructions, self.memory_start, memory_end)
1334 .map_err(TimerError::from)
1335 })
1336 .unwrap_or_else(|error| trap_callback_failure("callback measurement accounting", &error));
1337 }
1338}
1339
1340fn trap_callback_failure(context: &str, error: &TimerError) -> ! {
1341 platform::trap(&format!("ic-timers {context} failed: {error}"))
1342}
1343
1344fn with_registry<T>(
1345 operation: impl FnOnce(&TimerRegistry) -> Result<T, TimerError>,
1346) -> Result<T, TimerError> {
1347 RUNTIME.with(|runtime| {
1348 let runtime = runtime.try_borrow().map_err(|_| TimerError::RuntimeBusy)?;
1349 let registry = runtime.as_ref().ok_or(TimerError::NotInitialized)?;
1350 operation(registry)
1351 })
1352}
1353
1354fn with_registry_mut<T>(
1355 operation: impl FnOnce(&mut TimerRegistry) -> Result<T, TimerError>,
1356) -> Result<T, TimerError> {
1357 RUNTIME.with(|runtime| {
1358 let mut runtime = runtime
1359 .try_borrow_mut()
1360 .map_err(|_| TimerError::RuntimeBusy)?;
1361 let registry = runtime.as_mut().ok_or(TimerError::NotInitialized)?;
1362 operation(registry)
1363 })
1364}
1365
1366#[cfg(test)]
1367fn reset_for_test(now_ns: u64, canister_version: u64) {
1368 platform::reset(now_ns, canister_version);
1369 platform::ADMISSION_FAULT.set(None);
1370 WATCHDOG_COMPLETION_FAULT.with(|fault| fault.set(false));
1371 PROVIDER_INSTALL_FAULT_AFTER.with(|fault| fault.set(None));
1372 RUNTIME.with(|runtime| {
1373 *runtime.borrow_mut() = None;
1374 });
1375}
1376
1377#[cfg(test)]
1378thread_local! {
1379 static WATCHDOG_COMPLETION_FAULT: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
1380 static PROVIDER_INSTALL_FAULT_AFTER: std::cell::Cell<Option<u64>> = const { std::cell::Cell::new(None) };
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)]
1419mod tests;