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, WatchdogCallback,
10 },
11 schedule::{OrdinaryDirective, ScheduleError, TimerCadence, TimerSchedule},
12 snapshot::{
13 DeclarationLifetime, MemoryPageExtent, TimerCompletion, TimerControlFailure, TimerEpoch,
14 TimerIdentity, TimerInventorySnapshot, TimerPolicy, TimerSnapshot, WatchdogRunResult,
15 },
16};
17use std::{cell::RefCell, future::Future, rc::Rc, time::Duration};
18use thiserror::Error;
19
20thread_local! {
21 static RUNTIME: RefCell<Option<TimerRegistry>> = const { RefCell::new(None) };
22}
23
24#[non_exhaustive]
26#[derive(Debug, Error)]
27pub enum TimerError {
28 #[error("timer runtime is not initialized")]
30 NotInitialized,
31 #[error("timer runtime is already borrowed")]
33 RuntimeBusy,
34 #[error(transparent)]
36 Register(#[from] RegisterError),
37 #[error(transparent)]
39 Schedule(#[from] ScheduleError),
40 #[error("timer registration is no longer authoritative")]
42 RegistrationExpired,
43 #[error("timer control failed: {0:?}")]
45 ControlFailure(TimerControlFailure),
46 #[error("timer runtime ownership invariant failed")]
48 OwnershipInvariant,
49 #[error("timer lifecycle reconciliation conflicts with the canonical declaration")]
51 ReconciliationConflict,
52}
53
54impl From<RegistryError> for TimerError {
55 fn from(value: RegistryError) -> Self {
56 match value {
57 RegistryError::UnknownRegistration
58 | RegistryError::StaleRegistration
59 | RegistryError::StaleCallback => Self::RegistrationExpired,
60 RegistryError::PolicyMismatch { .. } | RegistryError::ProviderHandleAlreadyOwned => {
61 Self::OwnershipInvariant
62 }
63 RegistryError::Schedule(error) => Self::Schedule(error),
64 }
65 }
66}
67
68pub fn initialize_runtime() -> Result<TimerEpoch, TimerError> {
73 RUNTIME.with(|runtime| {
74 let mut runtime = runtime
75 .try_borrow_mut()
76 .map_err(|_| TimerError::RuntimeBusy)?;
77 let registry = runtime.get_or_insert_with(|| {
78 TimerRegistry::new(TimerEpoch::new(
79 platform::canister_version(),
80 platform::time_ns(),
81 ))
82 });
83 Ok(registry.epoch())
84 })
85}
86
87pub struct OnceContext {
93 token: CallbackToken,
94}
95
96impl OnceContext {
97 const fn new(token: CallbackToken) -> Self {
98 Self { token }
99 }
100
101 #[must_use]
103 pub const fn identity(&self) -> &TimerIdentity {
104 self.token.identity()
105 }
106
107 pub fn ensure_scheduled(&self, schedule: TimerSchedule) -> Result<(), TimerError> {
112 ensure_once_claim(self.token.claim(), Some(&self.token), schedule)
113 }
114
115 pub fn reconcile_schedule(&self, schedule: Option<TimerSchedule>) -> Result<(), TimerError> {
125 reconcile_ordinary_claim(self.token.claim(), Some(&self.token), schedule)
126 }
127
128 pub fn cancel(&self) -> Result<(), TimerError> {
136 cancel_claim(self.token.claim(), Some(&self.token))
137 }
138}
139
140pub struct AfterCompletionContext {
147 token: CallbackToken,
148}
149
150impl AfterCompletionContext {
151 const fn new(token: CallbackToken) -> Self {
152 Self { token }
153 }
154
155 #[must_use]
157 pub const fn identity(&self) -> &TimerIdentity {
158 self.token.identity()
159 }
160
161 pub fn ensure_scheduled(&self) -> Result<(), TimerError> {
163 ensure_recurring_claim(self.token.claim(), Some(&self.token))
164 }
165
166 pub fn reconcile_schedule(&self, schedule: Option<TimerSchedule>) -> Result<(), TimerError> {
177 reconcile_ordinary_claim(self.token.claim(), Some(&self.token), schedule)
178 }
179
180 pub fn cancel(&self) -> Result<(), TimerError> {
188 cancel_claim(self.token.claim(), Some(&self.token))
189 }
190}
191
192pub struct WatchdogContext {
197 token: CallbackToken,
198}
199
200impl WatchdogContext {
201 const fn new(token: CallbackToken) -> Self {
202 Self { token }
203 }
204
205 #[must_use]
207 pub const fn identity(&self) -> &TimerIdentity {
208 self.token.identity()
209 }
210
211 pub fn ensure_scheduled(&self) -> Result<(), TimerError> {
217 ensure_recurring_claim(self.token.claim(), Some(&self.token))
218 }
219
220 pub fn ensure_scheduled_immediately(&self) -> Result<(), TimerError> {
227 ensure_watchdog_immediately_claim(self.token.claim(), Some(&self.token))
228 }
229
230 pub fn reconcile_schedule(&self, schedule: Option<TimerSchedule>) -> Result<(), TimerError> {
236 reconcile_watchdog_claim(self.token.claim(), Some(&self.token), schedule)
237 }
238
239 pub fn cancel(&self) -> Result<(), TimerError> {
247 cancel_claim(self.token.claim(), Some(&self.token))
248 }
249}
250
251#[must_use = "retain the registration claim so the timer remains controllable"]
253pub struct OnceRegistration {
254 claim: RegistrationClaim,
255}
256
257impl OnceRegistration {
258 #[must_use]
260 pub const fn identity(&self) -> &TimerIdentity {
261 self.claim.identity()
262 }
263
264 pub fn has_armed_wakeup(&self) -> Result<bool, TimerError> {
270 has_armed_wakeup_claim(&self.claim)
271 }
272
273 pub fn ensure_scheduled(&self, schedule: TimerSchedule) -> Result<(), TimerError> {
275 ensure_once_claim(&self.claim, None, schedule)
276 }
277
278 pub fn reconcile_schedule(&self, schedule: Option<TimerSchedule>) -> Result<(), TimerError> {
288 reconcile_ordinary_claim(&self.claim, None, schedule)
289 }
290
291 pub fn cancel(&self) -> Result<(), TimerError> {
298 cancel_claim(&self.claim, None)
299 }
300
301 pub fn unregister(self) -> Result<(), TimerError> {
306 unregister_claim(&self.claim)
307 }
308}
309
310#[must_use = "retain the registration claim so the timer remains controllable"]
313pub struct AfterCompletionRegistration {
314 claim: RegistrationClaim,
315}
316
317#[must_use = "retain the registration claim so the timer remains controllable"]
319pub struct WatchdogRegistration {
320 claim: RegistrationClaim,
321}
322
323trait RegistrationClaimOwner {
324 fn registration_claim(&self) -> &RegistrationClaim;
325}
326
327impl RegistrationClaimOwner for OnceRegistration {
328 fn registration_claim(&self) -> &RegistrationClaim {
329 &self.claim
330 }
331}
332
333impl RegistrationClaimOwner for AfterCompletionRegistration {
334 fn registration_claim(&self) -> &RegistrationClaim {
335 &self.claim
336 }
337}
338
339impl RegistrationClaimOwner for WatchdogRegistration {
340 fn registration_claim(&self) -> &RegistrationClaim {
341 &self.claim
342 }
343}
344
345#[derive(Clone, Copy, Debug, Eq, PartialEq)]
347pub enum TimerReconcileState {
348 Inactive,
350 Scheduled,
352}
353
354#[derive(Clone, Copy, Debug, Eq, PartialEq)]
356pub enum WatchdogReconcileState {
357 Inactive,
359 Scheduled,
361 ScheduledImmediately,
366 ScheduledAt(u64),
368}
369
370impl WatchdogRegistration {
371 #[must_use]
373 pub const fn identity(&self) -> &TimerIdentity {
374 self.claim.identity()
375 }
376
377 pub fn has_armed_wakeup(&self) -> Result<bool, TimerError> {
384 has_armed_wakeup_claim(&self.claim)
385 }
386
387 pub fn ensure_scheduled(&self) -> Result<(), TimerError> {
391 ensure_recurring_claim(&self.claim, None)
392 }
393
394 pub fn ensure_scheduled_immediately(&self) -> Result<(), TimerError> {
402 ensure_watchdog_immediately_claim(&self.claim, None)
403 }
404
405 pub fn reconcile_schedule(&self, schedule: Option<TimerSchedule>) -> Result<(), TimerError> {
413 reconcile_watchdog_claim(&self.claim, None, schedule)
414 }
415
416 pub fn cancel(&self) -> Result<(), TimerError> {
424 cancel_claim(&self.claim, None)
425 }
426
427 pub fn unregister(self) -> Result<(), TimerError> {
432 unregister_claim(&self.claim)
433 }
434}
435
436impl AfterCompletionRegistration {
437 #[must_use]
439 pub const fn identity(&self) -> &TimerIdentity {
440 self.claim.identity()
441 }
442
443 pub fn has_armed_wakeup(&self) -> Result<bool, TimerError> {
449 has_armed_wakeup_claim(&self.claim)
450 }
451
452 pub fn ensure_scheduled(&self) -> Result<(), TimerError> {
455 ensure_recurring_claim(&self.claim, None)
456 }
457
458 pub fn reconcile_schedule(&self, schedule: Option<TimerSchedule>) -> Result<(), TimerError> {
467 reconcile_ordinary_claim(&self.claim, None, schedule)
468 }
469
470 pub fn cancel(&self) -> Result<(), TimerError> {
477 cancel_claim(&self.claim, None)
478 }
479
480 pub fn unregister(self) -> Result<(), TimerError> {
485 unregister_claim(&self.claim)
486 }
487}
488
489pub fn register_once<F, Fut>(
503 identity: TimerIdentity,
504 lifetime: DeclarationLifetime,
505 callback: F,
506) -> Result<OnceRegistration, TimerError>
507where
508 F: FnMut(OnceContext) -> Fut + 'static,
509 Fut: Future<Output = OnceRunResult> + 'static,
510{
511 let callback = erase_ordinary_callback(callback, OnceContext::new, OrdinaryRunResult::from);
512 let claim = with_registry_mut(|registry| {
513 registry
514 .register_once_with_callback(identity, lifetime, Rc::clone(&callback))
515 .map_err(TimerError::from)
516 });
517 drop(callback);
519 Ok(OnceRegistration { claim: claim? })
520}
521
522pub fn register_after_completion<F, Fut>(
536 identity: TimerIdentity,
537 cadence: TimerCadence,
538 lifetime: DeclarationLifetime,
539 callback: F,
540) -> Result<AfterCompletionRegistration, TimerError>
541where
542 F: FnMut(AfterCompletionContext) -> Fut + 'static,
543 Fut: Future<Output = AfterCompletionRunResult> + 'static,
544{
545 let callback = erase_ordinary_callback(
546 callback,
547 AfterCompletionContext::new,
548 OrdinaryRunResult::from,
549 );
550 let claim = with_registry_mut(|registry| {
551 registry
552 .register_after_completion_with_callback(
553 identity,
554 cadence,
555 lifetime,
556 Rc::clone(&callback),
557 )
558 .map_err(TimerError::from)
559 });
560 drop(callback);
562 Ok(AfterCompletionRegistration { claim: claim? })
563}
564
565pub fn register_watchdog<F>(
572 identity: TimerIdentity,
573 cadence: TimerCadence,
574 lifetime: DeclarationLifetime,
575 callback: F,
576) -> Result<WatchdogRegistration, TimerError>
577where
578 F: FnMut(WatchdogContext) -> WatchdogRunResult + 'static,
579{
580 let mut callback = callback;
581 let callback: WatchdogCallback = Rc::new(RefCell::new(Box::new(move |token| {
582 callback(WatchdogContext::new(token))
583 })));
584 let claim = with_registry_mut(|registry| {
585 registry
586 .register_watchdog_with_callback(identity, cadence, lifetime, Rc::clone(&callback))
587 .map_err(TimerError::from)
588 });
589 drop(callback);
591 Ok(WatchdogRegistration { claim: claim? })
592}
593
594pub fn reconcile_once<F, Fut>(
614 registration: &mut Option<OnceRegistration>,
615 identity: &TimerIdentity,
616 desired: Option<TimerSchedule>,
617 callback: F,
618) -> Result<(), TimerError>
619where
620 F: FnMut(OnceContext) -> Fut + 'static,
621 Fut: Future<Output = OnceRunResult> + 'static,
622{
623 let registration = reconcile_registration(registration, identity, TimerPolicy::Once, || {
624 register_once(identity.clone(), DeclarationLifetime::Retained, callback)
625 })?;
626 registration.reconcile_schedule(desired)
627}
628
629pub fn reconcile_after_completion<F, Fut>(
652 registration: &mut Option<AfterCompletionRegistration>,
653 identity: &TimerIdentity,
654 cadence: TimerCadence,
655 desired: TimerReconcileState,
656 callback: F,
657) -> Result<(), TimerError>
658where
659 F: FnMut(AfterCompletionContext) -> Fut + 'static,
660 Fut: Future<Output = AfterCompletionRunResult> + 'static,
661{
662 let registration = reconcile_registration(
663 registration,
664 identity,
665 TimerPolicy::AfterCompletion { cadence },
666 || {
667 register_after_completion(
668 identity.clone(),
669 cadence,
670 DeclarationLifetime::Retained,
671 callback,
672 )
673 },
674 )?;
675 match desired {
676 TimerReconcileState::Inactive => registration.cancel(),
677 TimerReconcileState::Scheduled => registration.ensure_scheduled(),
678 }
679}
680
681pub fn reconcile_watchdog<F>(
691 registration: &mut Option<WatchdogRegistration>,
692 identity: &TimerIdentity,
693 cadence: TimerCadence,
694 desired: WatchdogReconcileState,
695 callback: F,
696) -> Result<(), TimerError>
697where
698 F: FnMut(WatchdogContext) -> WatchdogRunResult + 'static,
699{
700 let registration = reconcile_registration(
701 registration,
702 identity,
703 TimerPolicy::Watchdog { cadence },
704 || {
705 register_watchdog(
706 identity.clone(),
707 cadence,
708 DeclarationLifetime::Retained,
709 callback,
710 )
711 },
712 )?;
713 match desired {
714 WatchdogReconcileState::Inactive => registration.cancel(),
715 WatchdogReconcileState::Scheduled => registration.ensure_scheduled(),
716 WatchdogReconcileState::ScheduledImmediately => registration.ensure_scheduled_immediately(),
717 WatchdogReconcileState::ScheduledAt(deadline_ns) => {
718 registration.reconcile_schedule(Some(TimerSchedule::At(deadline_ns)))
719 }
720 }
721}
722
723fn reconcile_registration<'a, Registration>(
724 registration: &'a mut Option<Registration>,
725 identity: &TimerIdentity,
726 policy: TimerPolicy,
727 register: impl FnOnce() -> Result<Registration, TimerError>,
728) -> Result<&'a Registration, TimerError>
729where
730 Registration: RegistrationClaimOwner,
731{
732 if registration.is_none() {
733 *registration = Some(register()?);
734 }
735 let registration = registration
736 .as_ref()
737 .ok_or(TimerError::ReconciliationConflict)?;
738 let claim = registration.registration_claim();
739 if claim.identity() != identity {
740 return Err(TimerError::ReconciliationConflict);
741 }
742 with_registry(|registry| {
743 registry
744 .declaration_matches(claim, policy, DeclarationLifetime::Retained)
745 .map_err(TimerError::from)?
746 .then_some(())
747 .ok_or(TimerError::ReconciliationConflict)
748 })?;
749 Ok(registration)
750}
751
752pub fn timer_snapshot(identity: &TimerIdentity) -> Result<Option<TimerSnapshot>, TimerError> {
754 with_registry(|registry| Ok(registry.snapshot(identity)))
755}
756
757pub fn timer_inventory() -> Result<TimerInventorySnapshot, TimerError> {
762 with_registry(|registry| Ok(registry.inventory()))
763}
764
765pub fn consecutive_expected_failures(identity: &TimerIdentity) -> Result<Option<u64>, TimerError> {
767 with_registry(|registry| Ok(registry.consecutive_expected_failures(identity)))
768}
769
770fn has_armed_wakeup_claim(claim: &RegistrationClaim) -> Result<bool, TimerError> {
771 with_registry(|registry| registry.has_armed_wakeup(claim).map_err(TimerError::from))
772}
773
774fn erase_ordinary_callback<Context: 'static, Output: 'static, F, Fut>(
775 mut callback: F,
776 context: fn(CallbackToken) -> Context,
777 erase_result: fn(Output) -> OrdinaryRunResult,
778) -> OrdinaryCallback
779where
780 F: FnMut(Context) -> Fut + 'static,
781 Fut: Future<Output = Output> + 'static,
782{
783 Rc::new(RefCell::new(Box::new(move |token| {
784 let future = callback(context(token));
785 Box::pin(async move { erase_result(future.await) })
786 })))
787}
788
789fn apply_claim_transition(
790 claim: &RegistrationClaim,
791 context: Option<&CallbackToken>,
792 operation: impl FnOnce(&mut TimerRegistry) -> Result<RegistryTransition, RegistryError>,
793) -> Result<(), TimerError> {
794 let (transition, handles) = with_registry_mut(|registry| {
795 if let Some(token) = context {
796 registry
797 .validate_running_context(token)
798 .map_err(TimerError::from)?;
799 }
800 registry
801 .transition_claim(claim, operation)
802 .map_err(TimerError::from)
803 })?;
804 match finish_transition(transition, handles) {
805 result @ (Ok(()) | Err(TimerError::ControlFailure(_))) => result,
806 Err(error) => retire_failed_claim(claim, error),
807 }
808}
809
810fn ensure_once_claim(
811 claim: &RegistrationClaim,
812 context: Option<&CallbackToken>,
813 schedule: TimerSchedule,
814) -> Result<(), TimerError> {
815 apply_claim_transition(claim, context, |registry| {
816 registry.ensure_once(claim, platform::time_ns(), schedule)
817 })
818}
819
820fn reconcile_ordinary_claim(
821 claim: &RegistrationClaim,
822 context: Option<&CallbackToken>,
823 schedule: Option<TimerSchedule>,
824) -> Result<(), TimerError> {
825 apply_claim_transition(claim, context, |registry| {
826 registry.reconcile_ordinary(claim, platform::time_ns(), schedule)
827 })
828}
829
830fn ensure_recurring_claim(
831 claim: &RegistrationClaim,
832 context: Option<&CallbackToken>,
833) -> Result<(), TimerError> {
834 apply_claim_transition(claim, context, |registry| {
835 registry.ensure_recurring(claim, platform::time_ns())
836 })
837}
838
839fn ensure_watchdog_immediately_claim(
840 claim: &RegistrationClaim,
841 context: Option<&CallbackToken>,
842) -> Result<(), TimerError> {
843 apply_claim_transition(claim, context, |registry| {
844 registry.ensure_watchdog_immediately(claim, platform::time_ns())
845 })
846}
847
848fn reconcile_watchdog_claim(
849 claim: &RegistrationClaim,
850 context: Option<&CallbackToken>,
851 schedule: Option<TimerSchedule>,
852) -> Result<(), TimerError> {
853 apply_claim_transition(claim, context, |registry| {
854 registry.reconcile_watchdog_schedule(claim, platform::time_ns(), schedule)
855 })
856}
857
858fn cancel_claim(
859 claim: &RegistrationClaim,
860 context: Option<&CallbackToken>,
861) -> Result<(), TimerError> {
862 apply_claim_transition(claim, context, |registry| registry.cancel(claim))
863}
864
865fn unregister_claim(claim: &RegistrationClaim) -> Result<(), TimerError> {
866 apply_claim_transition(claim, None, |registry| registry.unregister(claim))
867}
868
869fn retire_failed_claim(claim: &RegistrationClaim, error: TimerError) -> Result<(), TimerError> {
870 match fail_claim_provider_binding(claim) {
871 Ok(()) | Err(TimerError::RegistrationExpired) => Err(error),
872 Err(cleanup_error) => Err(cleanup_error),
873 }
874}
875
876fn finish_transition(
877 transition: RegistryTransition,
878 handles: ProviderHandles,
879) -> Result<(), TimerError> {
880 let failure = transition.failure();
881 let applied = apply_effect(transition.effect(), handles);
882 drop(transition);
885 applied?;
886 failure.map_or(Ok(()), |failure| Err(TimerError::ControlFailure(failure)))
887}
888
889fn apply_effect(effect: &RegistryEffect, mut handles: ProviderHandles) -> Result<(), TimerError> {
890 if !effect.has_valid_shape() {
891 clear_provider_handles(handles);
892 return Err(TimerError::OwnershipInvariant);
893 }
894 match effect {
895 RegistryEffect::None => restore_provider_handles(handles),
896 RegistryEffect::ArmWakeup {
897 token,
898 delay_ns,
899 arm,
900 } => {
901 if arm.replaces_existing() {
902 let replaced = take_detached_or_owned_handle(handles.take_wakeup(), |registry| {
903 registry.take_wakeup_handle(token.identity())
904 })?;
905 if let Some(replaced) = replaced {
906 clear_provider_handle(replaced);
907 }
908 }
909 restore_provider_handles(handles)?;
910 let handle = platform::set_timer(
911 Duration::from_nanos(*delay_ns),
912 dispatch_wakeup(token.clone()),
913 );
914 bind_provider_handle(token, handle)?;
915 if let Err(error) = confirm_effect(effect) {
916 let handle = with_registry_mut(|registry| {
917 registry
918 .take_wakeup_handle(token.identity())
919 .ok_or(TimerError::OwnershipInvariant)
920 })?;
921 clear_provider_handle(handle);
922 return Err(error);
923 }
924 Ok(())
925 }
926 RegistryEffect::ClearCallbacks {
927 identity,
928 handles: selected,
929 } => {
930 if selected.includes_wakeup() {
931 let wakeup = take_detached_or_owned_handle(handles.take_wakeup(), |registry| {
932 registry.take_wakeup_handle(identity)
933 })?;
934 if let Some(wakeup) = wakeup {
935 clear_provider_handle(wakeup);
936 }
937 }
938 if selected.includes_work() {
939 let work = take_detached_or_owned_handle(handles.take_work(), |registry| {
940 registry.take_work_handle(identity)
941 })?;
942 if let Some(work) = work {
943 clear_provider_handle(work);
944 }
945 }
946 restore_provider_handles(handles)
947 }
948 RegistryEffect::DispatchWatchdog {
949 successor,
950 successor_delay_ns,
951 work,
952 } => {
953 if let Some(wakeup) = handles.take_wakeup() {
954 clear_provider_handle(wakeup);
955 }
956 let replaced_work = take_detached_or_owned_handle(handles.take_work(), |registry| {
957 registry.take_work_handle(successor.identity())
958 })?;
959 if let Some(replaced_work) = replaced_work {
960 clear_provider_handle(replaced_work);
961 }
962 let successor_token = successor.clone();
963 let successor_handle =
964 platform::set_timer(Duration::from_nanos(*successor_delay_ns), async move {
965 dispatch_watchdog_scheduler(&successor_token);
966 });
967 bind_provider_handle(successor, successor_handle)?;
968
969 let work_token = work.clone();
970 let work_handle = platform::set_timer(Duration::ZERO, async move {
971 dispatch_watchdog_work(&work_token);
972 });
973 if let Err(error) =
974 bind_provider_handle(work, work_handle).and_then(|()| confirm_effect(effect))
975 {
976 clear_entry_provider_handles(successor.identity())?;
977 return Err(error);
978 }
979 Ok(())
980 }
981 }
982}
983
984fn take_detached_or_owned_handle(
985 detached: Option<ProviderHandle>,
986 take_owned: impl FnOnce(&mut TimerRegistry) -> Option<ProviderHandle>,
987) -> Result<Option<ProviderHandle>, TimerError> {
988 detached.map_or_else(
989 || with_registry_mut(|registry| Ok(take_owned(registry))),
990 |handle| Ok(Some(handle)),
991 )
992}
993
994fn bind_provider_handle(token: &CallbackToken, handle: TimerHandle) -> Result<(), TimerError> {
996 RUNTIME
997 .with(|runtime| {
998 #[cfg(test)]
999 if take_provider_install_fault() {
1000 return Err((TimerError::OwnershipInvariant, handle));
1001 }
1002 let Ok(mut runtime) = runtime.try_borrow_mut() else {
1003 return Err((TimerError::RuntimeBusy, handle));
1004 };
1005 let Some(registry) = runtime.as_mut() else {
1006 return Err((TimerError::NotInitialized, handle));
1007 };
1008 match registry.install_provider_handle(token, handle) {
1009 Ok(()) => Ok(()),
1010 Err((error, handle)) => Err((TimerError::from(error), handle)),
1011 }
1012 })
1013 .map_err(|(error, handle)| {
1014 platform::clear_timer(handle);
1016 error
1017 })
1018}
1019
1020fn confirm_effect(effect: &RegistryEffect) -> Result<(), TimerError> {
1021 #[cfg(test)]
1022 if take_provider_confirmation_fault() {
1023 return Err(TimerError::OwnershipInvariant);
1024 }
1025 with_registry_mut(|registry| {
1026 registry
1027 .confirm_effect_applied(effect)
1028 .map_err(TimerError::from)
1029 })
1030}
1031
1032fn restore_provider_handles(mut handles: ProviderHandles) -> Result<(), TimerError> {
1033 let wakeup_failure = handles
1036 .take_wakeup()
1037 .and_then(|handle| restore_provider_handle(handle).err());
1038 let work_failure = handles
1039 .take_work()
1040 .and_then(|handle| restore_provider_handle(handle).err());
1041 wakeup_failure.or(work_failure).map_or(Ok(()), Err)
1042}
1043
1044fn restore_provider_handle(handle: ProviderHandle) -> Result<(), TimerError> {
1045 let (token, handle) = handle.into_parts();
1046 bind_provider_handle(&token, handle)
1047}
1048
1049fn clear_provider_handle(handle: ProviderHandle) {
1050 let (_, handle) = handle.into_parts();
1051 platform::clear_timer(handle);
1052}
1053
1054fn clear_provider_handles(mut handles: ProviderHandles) {
1055 if let Some(wakeup) = handles.take_wakeup() {
1056 clear_provider_handle(wakeup);
1057 }
1058 if let Some(work) = handles.take_work() {
1059 clear_provider_handle(work);
1060 }
1061}
1062
1063fn clear_entry_provider_handles(identity: &TimerIdentity) -> Result<(), TimerError> {
1064 let handles = with_registry_mut(|registry| Ok(registry.take_provider_handles(identity)))?;
1065 clear_provider_handles(handles);
1066 Ok(())
1067}
1068
1069#[expect(
1070 clippy::future_not_send,
1071 reason = "IC callbacks and canister-local state are single-threaded."
1072)]
1073fn dispatch_wakeup(token: CallbackToken) -> impl Future<Output = ()> {
1074 let mut delivery = OrdinaryDelivery {
1077 token,
1078 completed: false,
1079 };
1080 async move {
1081 match delivery.token.role() {
1082 CallbackRole::OrdinaryWork => dispatch_ordinary(&delivery.token).await,
1083 CallbackRole::WatchdogScheduler => dispatch_watchdog_scheduler(&delivery.token),
1084 CallbackRole::WatchdogWork => {}
1085 }
1086 delivery.completed = true;
1087 drop(delivery);
1089 }
1090}
1091
1092struct OrdinaryDelivery {
1093 token: CallbackToken,
1094 completed: bool,
1095}
1096
1097impl Drop for OrdinaryDelivery {
1098 fn drop(&mut self) {
1099 if self.completed || self.token.role() != CallbackRole::OrdinaryWork {
1100 return;
1101 }
1102 let removed = RUNTIME.with(|runtime| {
1103 let Ok(mut runtime) = runtime.try_borrow_mut() else {
1104 let runtime = runtime.try_borrow().map_err(|_| TimerError::RuntimeBusy)?;
1108 return if runtime
1109 .as_ref()
1110 .is_some_and(|registry| registry.owns_ordinary_delivery(&self.token))
1111 {
1112 Err(TimerError::RuntimeBusy)
1113 } else {
1114 Ok(None)
1115 };
1116 };
1117 Ok(runtime
1118 .as_mut()
1119 .and_then(|registry| registry.abandon_ordinary(&self.token, platform::time_ns())))
1120 });
1121 let removed = removed
1122 .unwrap_or_else(|error| trap_callback_failure("ordinary delivery abandonment", &error));
1123 drop(removed);
1124 }
1125}
1126
1127#[expect(
1128 clippy::future_not_send,
1129 reason = "IC callbacks and canister-local state are single-threaded."
1130)]
1131async fn dispatch_ordinary(token: &CallbackToken) {
1132 let measurement = CallbackMeasurementStart::capture();
1133 let callback = with_registry_mut(|registry| Ok(registry.begin_ordinary(token)));
1134 let callback = match callback {
1135 Ok(Some(callback)) => callback,
1136 Ok(None) => return,
1137 Err(error) => trap_callback_failure("ordinary callback acceptance", &error),
1138 };
1139 let future = {
1140 let Ok(mut callback) = callback.try_borrow_mut() else {
1141 finish_ordinary_callback(
1142 token,
1143 OrdinaryRunResult::new(
1144 TimerCompletion::invariant_failure(0),
1145 OrdinaryDirective::Stop,
1146 ),
1147 );
1148 return;
1149 };
1150 callback(token.clone())
1151 };
1152 let result = future.await;
1153 finish_ordinary_callback(token, result);
1154 record_callback_measurements(token, measurement.finish());
1155}
1156
1157fn finish_ordinary_callback(token: &CallbackToken, result: OrdinaryRunResult) {
1158 let transition = with_registry_mut(|registry| {
1159 registry
1160 .complete_ordinary(token, platform::time_ns(), result)
1161 .map_err(TimerError::from)
1162 });
1163 let transition = transition
1164 .unwrap_or_else(|error| trap_callback_failure("ordinary callback completion", &error));
1165 finish_callback_transition(token, transition, ProviderHandles::default());
1166}
1167
1168fn dispatch_watchdog_scheduler(token: &CallbackToken) {
1169 let measurement = CallbackMeasurementStart::capture();
1170 let dispatched = with_registry_mut(|registry| {
1171 let handles = registry
1172 .take_watchdog_scheduler_handles(token)
1173 .map_err(TimerError::from)?;
1174 Ok((
1175 registry.begin_watchdog_scheduler(token, platform::time_ns()),
1176 handles,
1177 ))
1178 });
1179 let (transition, handles) = dispatched
1180 .unwrap_or_else(|error| trap_callback_failure("watchdog scheduler transition", &error));
1181 let accepted = !matches!(transition.effect(), RegistryEffect::None);
1182 finish_callback_transition(token, transition, handles);
1183 if accepted {
1184 record_callback_measurements(token, measurement.finish());
1185 }
1186}
1187
1188fn dispatch_watchdog_work(token: &CallbackToken) {
1189 let measurement = CallbackMeasurementStart::capture();
1190 let callback = with_registry_mut(|registry| Ok(registry.begin_watchdog_work(token)));
1191 let callback = match callback {
1192 Ok(Some(callback)) => callback,
1193 Ok(None) => return,
1194 Err(error) => trap_callback_failure("watchdog work acceptance", &error),
1195 };
1196 let result = {
1197 let Ok(mut callback) = callback.try_borrow_mut() else {
1198 trap_callback_failure(
1199 "watchdog callback ownership",
1200 &TimerError::OwnershipInvariant,
1201 );
1202 };
1203 callback(token.clone())
1204 };
1205 finish_watchdog_dispatch(token, result);
1206 record_callback_measurements(token, measurement.finish());
1207}
1208
1209fn finish_watchdog_dispatch(token: &CallbackToken, result: WatchdogRunResult) {
1210 let completed = with_registry_mut(|registry| {
1215 let handles = registry
1216 .take_provider_handles_for_claim(token.claim())
1217 .map_err(TimerError::from)?;
1218 #[cfg(test)]
1219 {
1220 if take_watchdog_completion_fault() {
1221 return Err(TimerError::OwnershipInvariant);
1222 }
1223 }
1224 let transition = registry
1225 .complete_watchdog_work(token, platform::time_ns(), result)
1226 .map_err(TimerError::from)?;
1227 Ok((transition, handles))
1228 });
1229 let (transition, handles) =
1230 completed.unwrap_or_else(|error| trap_callback_failure("watchdog work completion", &error));
1231 finish_callback_transition(token, transition, handles);
1232}
1233
1234fn finish_callback_transition(
1235 token: &CallbackToken,
1236 transition: RegistryTransition,
1237 handles: ProviderHandles,
1238) {
1239 match finish_transition(transition, handles) {
1242 Ok(()) | Err(TimerError::ControlFailure(_)) => {}
1243 Err(error) => {
1244 if token.role() == CallbackRole::WatchdogWork {
1245 trap_callback_failure("watchdog provider-handle completion", &error);
1246 }
1247 fail_claim_provider_binding(token.claim()).unwrap_or_else(|binding_error| {
1248 trap_callback_failure("provider-binding failure cleanup", &binding_error)
1249 });
1250 }
1251 }
1252}
1253
1254fn fail_claim_provider_binding(claim: &RegistrationClaim) -> Result<(), TimerError> {
1255 let failed = with_registry_mut(|registry| {
1256 registry
1257 .fail_registration(claim, TimerControlFailure::ProviderBindingFailed)
1258 .map_err(TimerError::from)
1259 });
1260 let (handles, transition) = failed?;
1261 clear_provider_handles(handles);
1262 drop(transition);
1263 Ok(())
1264}
1265
1266#[derive(Clone, Copy)]
1267struct CallbackMeasurementStart {
1268 instructions_before: u64,
1269 memory_start: MemoryPageExtent,
1270}
1271
1272impl CallbackMeasurementStart {
1273 fn capture() -> Self {
1274 let memory_start = platform::memory_pages();
1276 let instructions_before = platform::instruction_counter();
1277 Self {
1278 instructions_before,
1279 memory_start,
1280 }
1281 }
1282
1283 fn finish(self) -> CallbackMeasurement {
1284 let instructions = platform::instruction_counter().saturating_sub(self.instructions_before);
1286 let memory_end = platform::memory_pages();
1287 CallbackMeasurement {
1288 instructions,
1289 memory_start: self.memory_start,
1290 memory_end,
1291 }
1292 }
1293}
1294
1295#[derive(Clone, Copy)]
1296struct CallbackMeasurement {
1297 instructions: u64,
1298 memory_start: MemoryPageExtent,
1299 memory_end: MemoryPageExtent,
1300}
1301
1302fn record_callback_measurements(token: &CallbackToken, measurement: CallbackMeasurement) {
1303 with_registry_mut(|registry| {
1304 registry
1305 .record_callback_measurements(
1306 token,
1307 measurement.instructions,
1308 measurement.memory_start,
1309 measurement.memory_end,
1310 )
1311 .map_err(TimerError::from)
1312 })
1313 .unwrap_or_else(|error| trap_callback_failure("callback measurement accounting", &error));
1314}
1315
1316fn trap_callback_failure(context: &str, error: &TimerError) -> ! {
1317 platform::trap(&format!("ic-timers {context} failed: {error}"))
1318}
1319
1320fn with_registry<T>(
1321 operation: impl FnOnce(&TimerRegistry) -> Result<T, TimerError>,
1322) -> Result<T, TimerError> {
1323 RUNTIME.with(|runtime| {
1324 let runtime = runtime.try_borrow().map_err(|_| TimerError::RuntimeBusy)?;
1325 let registry = runtime.as_ref().ok_or(TimerError::NotInitialized)?;
1326 operation(registry)
1327 })
1328}
1329
1330fn with_registry_mut<T>(
1331 operation: impl FnOnce(&mut TimerRegistry) -> Result<T, TimerError>,
1332) -> Result<T, TimerError> {
1333 RUNTIME.with(|runtime| {
1334 let mut runtime = runtime
1335 .try_borrow_mut()
1336 .map_err(|_| TimerError::RuntimeBusy)?;
1337 let registry = runtime.as_mut().ok_or(TimerError::NotInitialized)?;
1338 operation(registry)
1339 })
1340}
1341
1342#[cfg(test)]
1343fn reset_for_test(now_ns: u64, canister_version: u64) {
1344 platform::reset(now_ns, canister_version);
1345 WATCHDOG_COMPLETION_FAULT.with(|fault| fault.set(false));
1346 PROVIDER_INSTALL_FAULT_AFTER.with(|fault| fault.set(None));
1347 PROVIDER_CONFIRMATION_FAULT.with(|fault| fault.set(false));
1348 RUNTIME.with(|runtime| {
1349 *runtime.borrow_mut() = None;
1350 });
1351}
1352
1353#[cfg(test)]
1354thread_local! {
1355 static WATCHDOG_COMPLETION_FAULT: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
1356 static PROVIDER_INSTALL_FAULT_AFTER: std::cell::Cell<Option<u64>> = const { std::cell::Cell::new(None) };
1357 static PROVIDER_CONFIRMATION_FAULT: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
1358}
1359
1360#[cfg(test)]
1361fn inject_watchdog_completion_fault() {
1362 WATCHDOG_COMPLETION_FAULT.with(|fault| fault.set(true));
1363}
1364
1365#[cfg(test)]
1366fn take_watchdog_completion_fault() -> bool {
1367 WATCHDOG_COMPLETION_FAULT.with(|fault| fault.replace(false))
1368}
1369
1370#[cfg(test)]
1371fn inject_provider_install_fault() {
1372 inject_provider_install_fault_after(0);
1373}
1374
1375#[cfg(test)]
1376fn inject_provider_install_fault_after(successful_installs: u64) {
1377 PROVIDER_INSTALL_FAULT_AFTER.with(|fault| fault.set(Some(successful_installs)));
1378}
1379
1380#[cfg(test)]
1381fn take_provider_install_fault() -> bool {
1382 PROVIDER_INSTALL_FAULT_AFTER.with(|fault| match fault.get() {
1383 Some(0) => {
1384 fault.set(None);
1385 true
1386 }
1387 Some(remaining) => {
1388 fault.set(Some(remaining - 1));
1389 false
1390 }
1391 None => false,
1392 })
1393}
1394
1395#[cfg(test)]
1396fn inject_provider_confirmation_fault() {
1397 PROVIDER_CONFIRMATION_FAULT.with(|fault| fault.set(true));
1398}
1399
1400#[cfg(test)]
1401fn take_provider_confirmation_fault() -> bool {
1402 PROVIDER_CONFIRMATION_FAULT.with(|fault| fault.replace(false))
1403}
1404
1405#[cfg(test)]
1406mod tests;