Skip to main content

ic_timers/runtime/
mod.rs

1//! Canister-local owner and live timer execution.
2
3use 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/// Initial schedule and declaration retention for a direct registration.
35///
36/// Retained declarations may start inactive. A declaration removed when stopped
37/// always starts with a schedule; all three registration policies use this same
38/// configuration. Durations and deadlines are checked when registering.
39///
40/// A remove-on-stop declaration cannot start inactive:
41///
42/// ```compile_fail,E0308
43/// use ic_timers::TimerStart;
44/// let start = TimerStart::RemoveWhenStopped { initial: None };
45/// ```
46#[derive(Clone, Copy, Debug, Eq, PartialEq)]
47pub enum TimerStart {
48    /// Keep callback authority after stopping, optionally starting scheduled.
49    Retained {
50        /// Initial wake-up, or `None` for an inactive declaration.
51        initial: Option<TimerSchedule>,
52    },
53    /// Start scheduled and remove callback authority after stopping.
54    RemoveWhenStopped {
55        /// Required initial wake-up.
56        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/// Failure from the canister-local timer runtime API.
72#[non_exhaustive]
73#[derive(Debug, Error)]
74pub enum TimerError {
75    /// The lifecycle owner has not initialized the volatile runtime.
76    #[error("timer runtime is not initialized")]
77    NotInitialized,
78    /// A nested internal borrow indicates unsupported re-entrancy.
79    #[error("timer runtime is already borrowed")]
80    RuntimeBusy,
81    /// Claiming one canonical timer identity failed.
82    #[error(transparent)]
83    Register(#[from] RegisterError),
84    /// Cadence or deadline validation failed.
85    #[error(transparent)]
86    Schedule(#[from] ScheduleError),
87    /// The logical registration was removed or superseded.
88    #[error("timer registration is no longer authoritative")]
89    RegistrationExpired,
90    /// Pure checked control reached a terminal failure after effects were applied.
91    #[error("timer control failed: {0:?}")]
92    ControlFailure(TimerControlFailure),
93    /// Canonical callback or provider-handle ownership was internally inconsistent.
94    #[error("timer runtime ownership invariant failed")]
95    OwnershipInvariant,
96    /// The bounded pending-callback store is full.
97    #[error("timer backend pending capacity exhausted")]
98    BackendCapacity,
99    /// The backend cannot allocate another non-reusable callback identity.
100    #[error("timer backend callback identities exhausted")]
101    BackendIdentityExhausted,
102    /// A retained lifecycle claim no longer matches its canonical declaration.
103    #[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
133/// Initialize the volatile canister-local runtime once for this Wasm instance.
134///
135/// Repeated calls are idempotent and return the original epoch. This function
136/// exports no lifecycle hook; the canister's existing lifecycle owner calls it.
137pub 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
152/// Delegated control capability scoped to one exact `Once` work attempt.
153///
154/// Identity remains inspectable after work returns or is abandoned, but mutation
155/// methods then return [`TimerError::RegistrationExpired`]. Retain the
156/// [`OnceRegistration`] for longer-lived ownership.
157pub struct OnceContext {
158    token: CallbackToken,
159}
160
161impl OnceContext {
162    const fn new(token: CallbackToken) -> Self {
163        Self { token }
164    }
165
166    /// Return the logical timer identity executing this work.
167    #[must_use]
168    pub const fn identity(&self) -> &TimerIdentity {
169        self.token.identity()
170    }
171
172    /// Schedule a `Once` declaration while this exact work attempt is active.
173    ///
174    /// A context retained after completion or delivery abandonment is expired
175    /// and returns [`TimerError::RegistrationExpired`].
176    pub fn ensure_scheduled(&self, schedule: TimerSchedule) -> Result<(), TimerError> {
177        ensure_once_claim(self.token.claim(), Some(&self.token), schedule)
178    }
179
180    /// Reconcile the executing declaration to one exact schedule.
181    ///
182    /// A winning exact schedule replaces the returned scheduling proposal
183    /// before validation. Explicit consumer invariant failure still stops.
184    ///
185    /// `None` requests inactive state at normal completion. Retained callback
186    /// authority remains; a remove-on-stop declaration is removed when that
187    /// cancellation wins arbitration. A stored context cannot mutate the
188    /// registration after its exact work attempt ends.
189    pub fn reconcile_schedule(&self, schedule: Option<TimerSchedule>) -> Result<(), TimerError> {
190        reconcile_ordinary_claim(self.token.claim(), Some(&self.token), schedule)
191    }
192
193    /// Request cancellation while this exact work attempt is active.
194    /// Cancellation allocates no callback generation.
195    ///
196    /// Cancellation does not interrupt the current invocation; normal
197    /// completion applies it before any callback successor is retained.
198    /// A retained declaration becomes inactive. A remove-on-stop declaration
199    /// is removed when cancellation wins arbitration.
200    pub fn cancel(&self) -> Result<(), TimerError> {
201        cancel_claim(self.token.claim(), Some(&self.token))
202    }
203}
204
205/// Delegated control capability scoped to one exact after-completion work
206/// attempt.
207///
208/// Mutation authority expires when the callback completes or its delivery is
209/// abandoned. Retain the [`AfterCompletionRegistration`] for longer-lived
210/// ownership.
211pub struct AfterCompletionContext {
212    token: CallbackToken,
213}
214
215impl AfterCompletionContext {
216    const fn new(token: CallbackToken) -> Self {
217        Self { token }
218    }
219
220    /// Return the logical timer identity executing this work.
221    #[must_use]
222    pub const fn identity(&self) -> &TimerIdentity {
223        self.token.identity()
224    }
225
226    /// Request configured recurrence after this exact work attempt.
227    pub fn ensure_scheduled(&self) -> Result<(), TimerError> {
228        ensure_recurring_claim(self.token.claim(), Some(&self.token))
229    }
230
231    /// Reconcile the executing declaration to one exact schedule without
232    /// changing its configured after-completion cadence.
233    ///
234    /// A winning exact schedule replaces the returned scheduling proposal
235    /// before validation. Explicit consumer invariant failure still stops.
236    ///
237    /// `None` requests inactive state at normal completion. A retained
238    /// declaration keeps its callback authority; a remove-on-stop declaration
239    /// is removed when cancellation wins arbitration. A stored context cannot
240    /// mutate the registration after its exact work attempt ends.
241    pub fn reconcile_schedule(&self, schedule: Option<TimerSchedule>) -> Result<(), TimerError> {
242        reconcile_ordinary_claim(self.token.claim(), Some(&self.token), schedule)
243    }
244
245    /// Request cancellation while this exact work attempt is active.
246    /// Cancellation allocates no callback generation.
247    ///
248    /// Cancellation does not interrupt the current invocation; normal
249    /// completion applies it before any callback successor is retained. A
250    /// retained declaration becomes inactive; a remove-on-stop declaration is
251    /// removed when cancellation wins arbitration.
252    pub fn cancel(&self) -> Result<(), TimerError> {
253        cancel_claim(self.token.claim(), Some(&self.token))
254    }
255}
256
257/// Delegated control capability scoped to one exact watchdog work attempt.
258///
259/// Mutation authority expires when the callback completes. Retain the
260/// [`WatchdogRegistration`] for longer-lived ownership.
261pub struct WatchdogContext {
262    token: CallbackToken,
263}
264
265impl WatchdogContext {
266    const fn new(token: CallbackToken) -> Self {
267        Self { token }
268    }
269
270    /// Return the logical timer identity executing this work.
271    #[must_use]
272    pub const fn identity(&self) -> &TimerIdentity {
273        self.token.identity()
274    }
275
276    /// Request that the pre-armed cadence successor remain scheduled.
277    ///
278    /// This is idempotent unless it supersedes a nested cancellation request.
279    /// It never arms an additional Watchdog successor and cannot move a
280    /// pending immediate successor later.
281    pub fn ensure_scheduled(&self) -> Result<(), TimerError> {
282        ensure_recurring_claim(self.token.claim(), Some(&self.token))
283    }
284
285    /// Request that this work attempt's pre-armed successor run immediately.
286    ///
287    /// Repeated requests coalesce, and a cadence ensure cannot move the
288    /// pending immediate deadline later. Normal completion replaces the exact
289    /// cadence successor with a zero-delay scheduler callback. Consumer work
290    /// still runs only in the scheduler's separately queued later message.
291    pub fn ensure_scheduled_immediately(&self) -> Result<(), TimerError> {
292        ensure_watchdog_immediately_claim(self.token.claim(), Some(&self.token))
293    }
294
295    /// Replace this attempt's successor with an exact schedule, or cancel it.
296    ///
297    /// Applied on normal completion; a trap preserves the committed cadence
298    /// successor. The latest reconciliation wins over prior scheduling demand,
299    /// but unregistration and invariant failure remain terminal.
300    pub fn reconcile_schedule(&self, schedule: Option<TimerSchedule>) -> Result<(), TimerError> {
301        reconcile_watchdog_claim(self.token.claim(), Some(&self.token), schedule)
302    }
303
304    /// Request cancellation while this exact work attempt is active.
305    /// Cancellation allocates no callback generation.
306    ///
307    /// Cancellation does not interrupt the current invocation; normal
308    /// completion clears the already-armed successor when cancellation wins.
309    /// A retained declaration becomes inactive; a remove-on-stop declaration
310    /// is removed.
311    pub fn cancel(&self) -> Result<(), TimerError> {
312        cancel_claim(self.token.claim(), Some(&self.token))
313    }
314}
315
316/// Opaque non-clone claim for one registered `Once` callback.
317///
318/// Dropping this capability does not cancel or unregister its declaration. Keep
319/// it in volatile owner state and call `unregister` for explicit removal. Losing
320/// a retained claim can leave a declaration occupying registry capacity until
321/// the next Wasm instance; dropping the claim does not release callback captures.
322#[must_use = "retain the registration claim so the timer remains controllable"]
323pub struct OnceRegistration {
324    claim: RegistrationClaim,
325}
326
327impl OnceRegistration {
328    /// Return the claimed logical identity.
329    #[must_use]
330    pub const fn identity(&self) -> &TimerIdentity {
331        self.claim.identity()
332    }
333
334    /// Return whether this exact claim currently owns an armed provider wake-up.
335    ///
336    /// This is a volatile observation, not durable scheduling authority or a
337    /// delivery guarantee. Call [`Self::ensure_scheduled`] unconditionally when
338    /// a wake-up is required rather than using this value as a scheduling guard.
339    pub fn has_armed_wakeup(&self) -> Result<bool, TimerError> {
340        has_armed_wakeup_claim(&self.claim)
341    }
342
343    /// Ensure one invocation is armed or retained as the running work's successor.
344    pub fn ensure_scheduled(&self, schedule: TimerSchedule) -> Result<(), TimerError> {
345        ensure_once_claim(&self.claim, None, schedule)
346    }
347
348    /// Reconcile to one exact desired schedule, replacing a later or earlier
349    /// live deadline as necessary.
350    ///
351    /// During running work, a winning exact schedule replaces the callback
352    /// scheduling proposal before validation; explicit invariant failure stops.
353    ///
354    /// `None` leaves a retained declaration inactive after any running work
355    /// completes or is abandoned. A remove-on-stop declaration is removed and
356    /// this claim expires when the transition finalizes.
357    pub fn reconcile_schedule(&self, schedule: Option<TimerSchedule>) -> Result<(), TimerError> {
358        reconcile_ordinary_claim(&self.claim, None, schedule)
359    }
360
361    /// Cancel the armed callback or the running work's successor.
362    /// Cancellation allocates no callback generation, even at the counter limit.
363    ///
364    /// Consumer work already running is not interrupted.
365    /// A retained declaration keeps callback authority. A remove-on-stop
366    /// declaration and this claim expire when cancellation finalizes.
367    pub fn cancel(&self) -> Result<(), TimerError> {
368        cancel_claim(&self.claim, None)
369    }
370
371    /// Consume the claim and unregister its callback authority.
372    ///
373    /// When called from running work, removal is deferred until that invocation
374    /// completes normally or its confirmed delivery is abandoned.
375    pub fn unregister(self) -> Result<(), TimerError> {
376        unregister_claim(&self.claim)
377    }
378}
379
380/// Opaque non-clone claim for one callback with configured
381/// after-completion recurrence.
382///
383/// Dropping this capability does not cancel or unregister its declaration. Keep
384/// it in volatile owner state and call `unregister` for explicit removal. Losing
385/// a retained claim can leave a declaration occupying registry capacity until
386/// the next Wasm instance; dropping the claim does not release callback captures.
387#[must_use = "retain the registration claim so the timer remains controllable"]
388pub struct AfterCompletionRegistration {
389    claim: RegistrationClaim,
390}
391
392/// Opaque non-clone claim for one pre-armed watchdog callback.
393///
394/// Dropping this capability does not cancel or unregister its declaration. Keep
395/// it in volatile owner state and call `unregister` for explicit removal. Losing
396/// a retained claim can leave a declaration occupying registry capacity until
397/// the next Wasm instance; dropping the claim does not release callback captures.
398#[must_use = "retain the registration claim so the timer remains controllable"]
399pub struct WatchdogRegistration {
400    claim: RegistrationClaim,
401}
402
403/// Desired volatile after-completion state during lifecycle reconciliation.
404#[derive(Clone, Copy, Debug, Eq, PartialEq)]
405pub enum TimerReconcileState {
406    /// Request inactive state immediately unless consumer work is running.
407    Inactive,
408    /// Preserve scheduling demand now or through the running work's successor.
409    Scheduled,
410}
411
412/// Desired volatile watchdog state during synchronous lifecycle reconciliation.
413#[derive(Clone, Copy, Debug, Eq, PartialEq)]
414pub enum WatchdogReconcileState {
415    /// Request inactive state immediately unless consumer work is running.
416    Inactive,
417    /// Preserve or arm one scheduler wake-up at the configured cadence.
418    Scheduled,
419    /// Preserve an earlier wake-up or arm/move one scheduler wake-up to now.
420    ///
421    /// The zero-delay provider callback executes as a later replicated
422    /// scheduler message and does not invoke consumer work synchronously.
423    ScheduledImmediately,
424    /// Reconcile the scheduler to an exact absolute IC timestamp.
425    ScheduledAt(u64),
426}
427
428impl WatchdogRegistration {
429    /// Return the claimed logical identity.
430    #[must_use]
431    pub const fn identity(&self) -> &TimerIdentity {
432        self.claim.identity()
433    }
434
435    /// Return whether this exact claim currently owns an armed scheduler wake-up.
436    ///
437    /// The separately queued work callback is not itself a wake-up. During
438    /// watchdog work this returns `true` only because the scheduler has already
439    /// committed and installed the cadence successor. This is a volatile
440    /// observation, not a delivery guarantee.
441    pub fn has_armed_wakeup(&self) -> Result<bool, TimerError> {
442        has_armed_wakeup_claim(&self.claim)
443    }
444
445    /// Synchronously ensure one watchdog scheduler wake-up is authoritative.
446    ///
447    /// An existing earlier immediate deadline is retained and never delayed.
448    pub fn ensure_scheduled(&self) -> Result<(), TimerError> {
449        ensure_recurring_claim(&self.claim, None)
450    }
451
452    /// Synchronously ensure one watchdog scheduler wake-up is due now.
453    ///
454    /// An inactive declaration arms one zero-delay scheduler. A later cadence
455    /// deadline is replaced in place; an already earlier or equivalent
456    /// deadline and repeated immediate requests coalesce. If work is already
457    /// dispatched, that zero-delay work satisfies the request. If work is
458    /// running, the request applies to that attempt's pre-armed successor.
459    pub fn ensure_scheduled_immediately(&self) -> Result<(), TimerError> {
460        ensure_watchdog_immediately_claim(&self.claim, None)
461    }
462
463    /// Reconcile to one exact scheduler deadline, or cancel with `None`.
464    ///
465    /// Unlike ensure, this can move a wake-up later. Identical deadlines
466    /// coalesce. Dispatched or running work is not postponed: reconciliation
467    /// selects its successor on normal completion. Until then, the committed
468    /// cadence successor remains the recovery wake-up. A subsequent recovery
469    /// attempt retires the interrupted attempt's pending proposal.
470    pub fn reconcile_schedule(&self, schedule: Option<TimerSchedule>) -> Result<(), TimerError> {
471        reconcile_watchdog_claim(&self.claim, None, schedule)
472    }
473
474    /// Cancel the scheduler and any work callback that has not started.
475    /// Cancellation allocates no callback generation, even at the counter limits.
476    ///
477    /// Consumer work already running is not interrupted; normal completion
478    /// clears its pre-armed successor.
479    /// A retained declaration keeps callback authority. A remove-on-stop
480    /// declaration and this claim expire when cancellation finalizes.
481    pub fn cancel(&self) -> Result<(), TimerError> {
482        cancel_claim(&self.claim, None)
483    }
484
485    /// Consume the claim and unregister its callback authority.
486    ///
487    /// When called from running work, removal is deferred until that invocation
488    /// completes normally.
489    pub fn unregister(self) -> Result<(), TimerError> {
490        unregister_claim(&self.claim)
491    }
492}
493
494impl AfterCompletionRegistration {
495    /// Return the claimed logical identity.
496    #[must_use]
497    pub const fn identity(&self) -> &TimerIdentity {
498        self.claim.identity()
499    }
500
501    /// Return whether this exact claim currently owns an armed provider wake-up.
502    ///
503    /// A callback currently running without an installed successor returns
504    /// `false`. This is a volatile observation, not durable scheduling authority
505    /// or a delivery guarantee.
506    pub fn has_armed_wakeup(&self) -> Result<bool, TimerError> {
507        has_armed_wakeup_claim(&self.claim)
508    }
509
510    /// Ensure configured recurrence is armed or retained as the running work's
511    /// successor.
512    pub fn ensure_scheduled(&self) -> Result<(), TimerError> {
513        ensure_recurring_claim(&self.claim, None)
514    }
515
516    /// Reconcile to one exact desired schedule without changing the configured
517    /// after-completion cadence. `None` makes a retained declaration inactive
518    /// after any running work completes or is abandoned; it removes a
519    /// remove-on-stop declaration and expires this claim when the transition
520    /// finalizes.
521    ///
522    /// During running work, a winning exact schedule replaces the callback
523    /// scheduling proposal before validation; explicit invariant failure stops.
524    pub fn reconcile_schedule(&self, schedule: Option<TimerSchedule>) -> Result<(), TimerError> {
525        reconcile_ordinary_claim(&self.claim, None, schedule)
526    }
527
528    /// Cancel the armed callback or the running work's successor.
529    /// Cancellation allocates no callback generation, even at the counter limit.
530    ///
531    /// Consumer work already running is not interrupted.
532    /// A retained declaration keeps callback authority. A remove-on-stop
533    /// declaration and this claim expire when cancellation finalizes.
534    pub fn cancel(&self) -> Result<(), TimerError> {
535        cancel_claim(&self.claim, None)
536    }
537
538    /// Consume the claim and unregister its callback authority.
539    ///
540    /// When called from running work, removal is deferred until that invocation
541    /// completes normally or its confirmed delivery is abandoned.
542    pub fn unregister(self) -> Result<(), TimerError> {
543        unregister_claim(&self.claim)
544    }
545}
546
547/// Register one asynchronous `Once` callback with its initial schedule.
548///
549/// [`TimerStart::Retained`] permits an inactive declaration;
550/// [`TimerStart::RemoveWhenStopped`] always supplies its initial schedule.
551/// Success returns installed callback authority and,
552/// when requested, a confirmed provider wakeup. A returned error retires this
553/// creation's declaration and handles before releasing callback captures.
554///
555/// The future must return [`OnceRunResult`]; a recurring policy's result cannot
556/// be used here, even when its decision is Stop:
557///
558/// ```compile_fail,E0271
559/// use ic_timers::{register_once, TimerIdentity, TimerStart,
560///     TimerCompletion, AfterCompletionDecision, AfterCompletionRunResult};
561/// let identity = TimerIdentity::try_new("app", "jobs", "once").unwrap();
562/// let _ = register_once(identity, TimerStart::Retained { initial: None }, |_| async {
563///     AfterCompletionRunResult::new(TimerCompletion::no_work(), AfterCompletionDecision::Stop)
564/// });
565/// ```
566pub 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    // Release factory captures only after initial registry/provider cleanup.
580    drop(callback);
581    Ok(OnceRegistration { claim: claim? })
582}
583
584/// Register one asynchronous callback with configured after-completion recurrence.
585///
586/// The explicit initial schedule selects the first deadline independently of
587/// the recurrence cadence. [`TimerStart::Retained`] may start inactive.
588/// Initial scheduling failure removes the new declaration and provider handles.
589///
590/// The future must return [`AfterCompletionRunResult`]:
591///
592/// ```compile_fail,E0271
593/// use ic_timers::{register_after_completion, TimerIdentity, TimerCadence,
594///     TimerStart, TimerCompletion, OnceDecision, OnceRunResult};
595/// let identity = TimerIdentity::try_new("app", "jobs", "recurring").unwrap();
596/// let cadence = TimerCadence::from_nanos(5).unwrap();
597/// let _ = register_after_completion(identity, cadence, TimerStart::Retained { initial: None }, |_| async {
598///     OnceRunResult::new(TimerCompletion::no_work(), OnceDecision::Stop)
599/// });
600/// ```
601pub 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    // Release factory captures only after initial registry/provider cleanup.
625    drop(callback);
626    Ok(AfterCompletionRegistration { claim: claim? })
627}
628
629/// Register one synchronous pre-armed watchdog callback with its initial schedule.
630///
631/// The initial schedule selects the first scheduler message, never synchronous
632/// consumer work. [`TimerStart::Retained`] may start inactive. Initial
633/// scheduling failure removes the new declaration and provider handles.
634///
635/// The callback cannot be async: it runs only in the work message after a
636/// separate scheduler message has armed the next cadence successor. Its
637/// `WatchdogDecision` retains the cadence successor, replaces it with a
638/// deadline of now or an exact timestamp, or clears it.
639pub 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    // Release factory captures only after initial registry/provider cleanup.
656    drop(callback);
657    Ok(WatchdogRegistration { claim: claim? })
658}
659
660/// Finish creation before exposing a claim. The callback factory retains its
661/// captures until this returns, including through rejected provider binding.
662fn 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
692/// Allocate and initialize under one borrow; publish authority only after effects.
693/// Fresh initialization failures share direct registration's retirement owner.
694fn 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
723/// Reconstruct or reconcile one `Once` declaration synchronously.
724///
725/// `Some(schedule)` is authoritative and may move an existing deadline in
726/// either direction. `None` retains an inactive declaration in the canonical
727/// inventory, including on a fresh heap. Lifecycle reconciliation always owns
728/// a [`TimerStart::Retained`] declaration; transient
729/// `RemoveWhenStopped` callbacks use [`register_once`] directly.
730/// An empty owner slot is filled only after initial reconciliation succeeds.
731/// A returned error cleans up the new declaration and captures, leaving that
732/// slot empty. Later errors preserve an existing registration's authority.
733///
734/// Lifecycle reconstruction retains the same policy-specific result boundary:
735///
736/// ```compile_fail,E0271
737/// use ic_timers::{reconcile_once, TimerIdentity, TimerCompletion,
738///     AfterCompletionDecision, AfterCompletionRunResult};
739/// let identity = TimerIdentity::try_new("app", "jobs", "once").unwrap();
740/// let mut registration = None;
741/// let _ = reconcile_once(&mut registration, &identity, None, |_| async {
742///     AfterCompletionRunResult::new(TimerCompletion::no_work(), AfterCompletionDecision::Stop)
743/// });
744/// ```
745pub 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| &current.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
790/// Reconstruct or reconcile one after-completion declaration synchronously.
791///
792/// The consumer owns `registration` in volatile state. A fresh Wasm heap has
793/// `None`, so this function creates callback authority and publishes it only
794/// after initial active or inactive reconciliation succeeds. A returned error
795/// cleans up a fresh declaration, leaving the slot empty. Later errors retain
796/// existing authority. A repeated call reuses the exact claim and does not
797/// replace its callback. The installed declaration is always retained;
798/// transient `RemoveWhenStopped` recurrence uses
799/// [`register_after_completion`] directly.
800///
801/// Lifecycle reconstruction also requires [`AfterCompletionRunResult`]:
802///
803/// ```compile_fail,E0271
804/// use ic_timers::{reconcile_after_completion, TimerIdentity, TimerCadence,
805///     TimerReconcileState, TimerCompletion, OnceDecision, OnceRunResult};
806/// let identity = TimerIdentity::try_new("app", "jobs", "recurring").unwrap();
807/// let cadence = TimerCadence::from_nanos(5).unwrap();
808/// let mut registration = None;
809/// let _ = reconcile_after_completion(&mut registration, &identity, cadence,
810///     TimerReconcileState::Inactive, |_| async {
811///         OnceRunResult::new(TimerCompletion::no_work(), OnceDecision::Stop)
812///     });
813/// ```
814pub 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| &current.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
870/// Reconstruct or reconcile one watchdog declaration synchronously.
871///
872/// Durable readiness remains consumer-owned. Fresh inactive authority still
873/// installs an observable retained declaration. `ScheduledImmediately` arms
874/// its first scheduler at deadline now without synchronous consumer work.
875/// `ScheduledAt` reconciles an exact deadline, including moving it later.
876/// This helper owns no lifecycle export and persists no policy, generation,
877/// provider handle, or callback.
878/// Transient `RemoveWhenStopped` watchdogs use [`register_watchdog`] directly.
879/// An empty owner slot is filled only after its initial state succeeds. A
880/// returned error cleans up the new declaration and captures, leaving the slot
881/// empty. Later errors retain existing authority and its failure observations.
882pub 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| &current.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    // Repeated calls retain the original callback. Release the unused factory
975    // outside the registry borrow, before applying control as in prior calls.
976    drop(register);
977    reconcile(current)
978}
979
980/// Return one coherent inert snapshot by identity.
981pub fn timer_snapshot(identity: &TimerIdentity) -> Result<Option<TimerSnapshot>, TimerError> {
982    with_registry(|registry| Ok(registry.snapshot(identity)))
983}
984
985/// Return one atomic bounded inventory with its volatile runtime epoch.
986///
987/// Timer snapshots are ordered deterministically by identity. An initialized
988/// empty registry still returns its epoch and an empty timer slice.
989pub fn timer_inventory() -> Result<TimerInventorySnapshot, TimerError> {
990    with_registry(|registry| Ok(registry.inventory()))
991}
992
993/// Return functional expected-failure state by identity without a full inventory.
994pub 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    // Construct outside the async body: a provider rejection can discard the
1110    // future before its first poll, after rolling registry acceptance back.
1111    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                // Capture the whole guard until normal completion.
1124                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                // Rejected binding may clear a detached scheduled handle
1149                // while another shared borrow exists. Its transition owns
1150                // cleanup; do not turn that typed failure into a Drop trap.
1151                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    // Unlike synchronous public control, an unexpected callback-completion
1255    // failure must trap. IC message rollback restores these temporarily
1256    // detached heap capabilities while the previously committed successor
1257    // remains armed by the scheduler message.
1258    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    // Registry control failures are already finalized. Other binding failures
1284    // follow the callback role's cleanup or rollback rule.
1285    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        // Keep page observation outside the established instruction interval.
1319        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        // Close the instruction interval before taking its paired end extent.
1329        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;