ic_canister_kit/functions/
schedule.rs

1use std::cell::Cell;
2
3use crate::types::DurationNanos;
4
5const MIN_SCHEDULE_INTERVAL_NANOS: u128 = 1_000_000_000;
6
7thread_local! {
8    static SCHEDULE_TASK_RUNNING: Cell<bool> = const { Cell::new(false) };
9}
10
11/// 定时任务执行锁。
12///
13/// 锁离开作用域时会自动释放，使后续定时任务可以继续执行。
14#[must_use = "dropping the guard immediately releases the schedule task lock"]
15#[non_exhaustive]
16pub struct ScheduleTaskGuard;
17
18impl Drop for ScheduleTaskGuard {
19    fn drop(&mut self) {
20        SCHEDULE_TASK_RUNNING.with(|running| running.set(false));
21    }
22}
23
24/// 尝试取得定时任务执行锁，防止自动任务和手动触发并发执行。
25pub fn try_schedule_task_guard() -> Result<ScheduleTaskGuard, String> {
26    SCHEDULE_TASK_RUNNING.with(|running| {
27        if running.get() {
28            return Err("Schedule task is already running.".to_string());
29        }
30        running.set(true);
31        Ok(ScheduleTaskGuard)
32    })
33}
34
35/// 验证定时任务间隔是否能够被 IC 定时器安全执行。
36///
37/// 已启用的任务间隔不得少于一秒，也不得超过定时器的 `u64` 纳秒范围或导致当前 Canister 时间溢出。
38pub fn validate_schedule(schedule: Option<DurationNanos>) -> Result<Option<DurationNanos>, String> {
39    if let Some(interval) = schedule {
40        let nanos = interval.into_inner();
41        if nanos < MIN_SCHEDULE_INTERVAL_NANOS {
42            return Err(format!(
43                "Schedule interval must be at least {MIN_SCHEDULE_INTERVAL_NANOS} nanoseconds."
44            ));
45        }
46        let nanos = u64::try_from(nanos)
47            .map_err(|_| "Schedule interval exceeds the timer's u64 nanosecond range.".to_string())?;
48        if ic_cdk::api::time().checked_add(nanos).is_none() {
49            return Err("Schedule interval is too large for the current canister time.".to_string());
50        }
51    }
52    Ok(schedule)
53}
54
55// ================== 异步执行代码 ==================
56// 不知道和 ic_cdk::spawn(future) 区别在哪里
57
58/// 异步执行代码
59#[inline]
60pub fn async_execute<Task>(task: Task) -> ic_cdk_timers::TimerId
61where
62    Task: Future<Output = ()> + 'static,
63{
64    ic_cdk_timers::set_timer(std::time::Duration::ZERO, task)
65}
66
67// ================== 功能 ==================
68
69pub use ic_cdk_timers::TimerId;
70
71/// 定时任务功能
72pub trait Schedulable {
73    /// 查询
74    fn schedule_find(&self) -> Option<DurationNanos>;
75    /// 修改
76    fn schedule_replace(&mut self, schedule: Option<DurationNanos>);
77}
78
79/// 停止任务
80#[inline]
81pub fn schedule_stop(timer_id: Option<TimerId>) {
82    if let Some(timer_id) = timer_id {
83        ic_cdk_timers::clear_timer(timer_id)
84    }
85}
86
87/// 启动任务
88#[inline]
89pub fn schedule_start<F>(schedule: &Option<DurationNanos>, task: impl FnMut() -> F + 'static) -> Option<TimerId>
90where
91    F: Future<Output = ()> + 'static,
92{
93    schedule.map(|interval| {
94        ic_cdk_timers::set_timer_interval(std::time::Duration::from_nanos(interval.into_inner() as u64), task)
95    })
96}
97
98// ================== 简单实现 ==================
99
100/// 定时任务简单实现
101pub mod basic {
102    use candid::CandidType;
103    use serde::{Deserialize, Serialize};
104
105    use crate::{functions::types::Schedulable, types::DurationNanos};
106
107    #[cfg(feature = "schedule")]
108    mod schedule {
109        use std::cell::RefCell;
110
111        use ic_cdk_timers::TimerId;
112
113        use crate::types::DurationNanos;
114
115        thread_local! {
116            static SCHEDULE: RefCell<Option<TimerId>> = RefCell::default(); // 定时任务 id 记录
117        }
118
119        /// 停止定时任务
120        #[inline]
121        pub fn stop_schedule() {
122            SCHEDULE.with_borrow_mut(|timer_id| crate::functions::schedule::schedule_stop(std::mem::take(timer_id)));
123        }
124
125        /// 启动定时任务
126        #[inline]
127        pub fn start_schedule<F>(schedule: &Option<DurationNanos>, task: impl FnMut() -> F + 'static)
128        where
129            F: Future<Output = ()> + 'static,
130        {
131            stop_schedule();
132            let new_timer_id = crate::functions::schedule::schedule_start(schedule, task);
133            SCHEDULE.with_borrow_mut(|timer_id| *timer_id = new_timer_id);
134        }
135    }
136    #[cfg(feature = "schedule")]
137    pub use schedule::*;
138
139    /// 周期定时任务
140    #[derive(CandidType, Serialize, Deserialize, Debug, Clone, Default)]
141    pub struct Schedule(Option<DurationNanos>);
142
143    impl Schedulable for Schedule {
144        // 查询
145        fn schedule_find(&self) -> Option<DurationNanos> {
146            self.0
147        }
148        // 修改
149        fn schedule_replace(&mut self, schedule: Option<DurationNanos>) {
150            self.0 = schedule
151        }
152    }
153}
154
155#[cfg(feature = "schedule")]
156pub use basic::{start_schedule, stop_schedule};
157
158#[cfg(test)]
159mod tests {
160    use super::*;
161
162    #[test]
163    fn schedule_task_guard_prevents_reentry_until_dropped() {
164        let first = try_schedule_task_guard().expect("first schedule task should acquire the guard");
165        assert!(try_schedule_task_guard().is_err());
166        drop(first);
167        assert!(try_schedule_task_guard().is_ok());
168    }
169
170    #[test]
171    fn schedule_interval_rejects_unsafe_values_before_reading_canister_time() {
172        assert!(validate_schedule(Some(0_u128.into())).is_err());
173        assert!(validate_schedule(Some((u64::MAX as u128 + 1).into())).is_err());
174    }
175}
ic_canister_kit/functions/schedule.rs

ic_canister_kit/functions/
schedule.rs
