use crate::config::ScheduleError;
use crate::scheduler::BeatScheduler;
use chrono::{DateTime, Utc};
pub const DEFAULT_DISPATCH_LOCK_TTL_SECS: u64 = 30;
pub fn dispatch_lock_key(entry_name: &str, scheduled_instant: DateTime<Utc>) -> String {
format!(
"celers-beat:dispatch:{}@{}",
entry_name,
scheduled_instant.to_rfc3339()
)
}
impl BeatScheduler {
pub async fn try_acquire_dispatch_lock(
&mut self,
entry_name: &str,
scheduled_instant: DateTime<Utc>,
ttl_secs: u64,
) -> Result<bool, ScheduleError> {
let key = dispatch_lock_key(entry_name, scheduled_instant);
self.try_acquire_distributed_lock(&key, ttl_secs).await
}
pub async fn release_dispatch_lock(
&mut self,
entry_name: &str,
scheduled_instant: DateTime<Utc>,
) -> Result<bool, ScheduleError> {
let key = dispatch_lock_key(entry_name, scheduled_instant);
self.release_distributed_lock(&key).await
}
pub async fn tick_with_locks(
&mut self,
lock_ttl_secs: u64,
) -> Result<Vec<String>, ScheduleError> {
let role = self.heartbeat_tick().await?;
let is_leader = match role {
Some(crate::heartbeat::BeatRole::Leader) => true,
Some(_) => false,
None => true, };
if !is_leader {
self.update_heartbeat_info().await;
return Ok(Vec::new());
}
let due_fires: Vec<(String, DateTime<Utc>)> = self
.get_due_tasks_by_priority()
.iter()
.filter_map(|task| {
task.next_run_time()
.ok()
.map(|instant| (task.name.clone(), instant))
})
.collect();
let mut dispatched = Vec::new();
for (name, instant) in due_fires {
let acquired = self
.try_acquire_dispatch_lock(&name, instant, lock_ttl_secs)
.await?;
if !acquired {
continue;
}
let _ = self.mark_task_run(&name);
dispatched.push(name.clone());
}
self.update_heartbeat_info().await;
let _ = self.save_state();
Ok(dispatched)
}
}
#[cfg(test)]
mod tests {
use super::*;
use chrono::TimeZone;
#[test]
fn key_is_stable_for_same_entry_and_instant() {
let instant = Utc
.with_ymd_and_hms(2026, 6, 13, 9, 30, 0)
.single()
.expect("valid instant");
let a = dispatch_lock_key("report", instant);
let b = dispatch_lock_key("report", instant);
assert_eq!(a, b);
}
#[test]
fn key_differs_for_different_instants() {
let first = Utc
.with_ymd_and_hms(2026, 6, 13, 9, 30, 0)
.single()
.expect("valid instant");
let second = first + chrono::Duration::seconds(60);
assert_ne!(
dispatch_lock_key("report", first),
dispatch_lock_key("report", second)
);
}
#[test]
fn key_differs_for_different_entries() {
let instant = Utc
.with_ymd_and_hms(2026, 6, 13, 9, 30, 0)
.single()
.expect("valid instant");
assert_ne!(
dispatch_lock_key("entry_a", instant),
dispatch_lock_key("entry_b", instant)
);
}
}