use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
use chrono::{DateTime, Utc};
use dashmap::DashMap;
use serde::{Deserialize, Serialize};
use tokio_util::sync::CancellationToken;
use crate::ipc::{KernelMessage, MessagePayload};
use crate::process::Pid;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TimerInfo {
pub id: String,
pub owner_pid: Pid,
pub fire_at: DateTime<Utc>,
pub repeat_interval_ms: Option<u64>,
pub cancelled: bool,
}
pub struct TimerEntry {
pub id: String,
pub owner_pid: Pid,
pub fire_at: DateTime<Utc>,
pub repeat_interval: Option<Duration>,
pub payload: MessagePayload,
pub cancel_token: CancellationToken,
}
impl TimerEntry {
pub fn info(&self) -> TimerInfo {
TimerInfo {
id: self.id.clone(),
owner_pid: self.owner_pid,
fire_at: self.fire_at,
repeat_interval_ms: self.repeat_interval.map(|d| d.as_millis() as u64),
cancelled: self.cancel_token.is_cancelled(),
}
}
}
pub struct TimerService {
timers: DashMap<String, TimerEntry>,
next_id: AtomicU64,
}
impl TimerService {
pub fn new() -> Self {
Self {
timers: DashMap::new(),
next_id: AtomicU64::new(1),
}
}
fn next_timer_id(&self) -> String {
let id = self.next_id.fetch_add(1, Ordering::Relaxed);
format!("timer-{id}")
}
pub fn create_oneshot(
&self,
owner_pid: Pid,
delay: Duration,
payload: MessagePayload,
) -> (String, CancellationToken) {
let id = self.next_timer_id();
let cancel_token = CancellationToken::new();
let fire_at = Utc::now() + chrono::Duration::from_std(delay).unwrap_or_default();
let entry = TimerEntry {
id: id.clone(),
owner_pid,
fire_at,
repeat_interval: None,
payload,
cancel_token: cancel_token.clone(),
};
self.timers.insert(id.clone(), entry);
(id, cancel_token)
}
pub fn create_repeating(
&self,
owner_pid: Pid,
interval: Duration,
payload: MessagePayload,
) -> (String, CancellationToken) {
let id = self.next_timer_id();
let cancel_token = CancellationToken::new();
let fire_at = Utc::now() + chrono::Duration::from_std(interval).unwrap_or_default();
let entry = TimerEntry {
id: id.clone(),
owner_pid,
fire_at,
repeat_interval: Some(interval),
payload,
cancel_token: cancel_token.clone(),
};
self.timers.insert(id.clone(), entry);
(id, cancel_token)
}
pub fn cancel(&self, timer_id: &str) -> bool {
if let Some(entry) = self.timers.get(timer_id) {
entry.cancel_token.cancel();
true
} else {
false
}
}
pub fn cancel_for_pid(&self, pid: Pid) -> usize {
let mut cancelled = 0;
for entry in self.timers.iter() {
if entry.owner_pid == pid {
entry.cancel_token.cancel();
cancelled += 1;
}
}
cancelled
}
pub fn remove(&self, timer_id: &str) -> Option<TimerEntry> {
self.timers.remove(timer_id).map(|(_, e)| e)
}
pub fn cleanup_cancelled(&self) -> usize {
let cancelled: Vec<String> = self
.timers
.iter()
.filter(|e| e.cancel_token.is_cancelled())
.map(|e| e.key().clone())
.collect();
let count = cancelled.len();
for id in cancelled {
self.timers.remove(&id);
}
count
}
pub fn info(&self, timer_id: &str) -> Option<TimerInfo> {
self.timers.get(timer_id).map(|e| e.info())
}
pub fn list(&self) -> Vec<String> {
self.timers.iter().map(|e| e.key().clone()).collect()
}
pub fn list_info(&self) -> Vec<TimerInfo> {
self.timers.iter().map(|e| e.value().info()).collect()
}
pub fn list_for_pid(&self, pid: Pid) -> Vec<TimerInfo> {
self.timers
.iter()
.filter(|e| e.owner_pid == pid)
.map(|e| e.info())
.collect()
}
pub fn len(&self) -> usize {
self.timers.len()
}
pub fn is_empty(&self) -> bool {
self.timers.is_empty()
}
pub fn collect_due(&self) -> Vec<(Pid, KernelMessage)> {
let now = Utc::now();
let mut due = Vec::new();
let mut to_remove = Vec::new();
let mut to_reschedule = Vec::new();
for entry in self.timers.iter() {
if entry.cancel_token.is_cancelled() {
to_remove.push(entry.key().clone());
continue;
}
if entry.fire_at <= now {
let msg = KernelMessage::new(
0, crate::ipc::MessageTarget::Process(entry.owner_pid),
entry.payload.clone(),
);
due.push((entry.owner_pid, msg));
if let Some(interval) = entry.repeat_interval {
to_reschedule.push((entry.key().clone(), interval));
} else {
to_remove.push(entry.key().clone());
}
}
}
for id in &to_remove {
self.timers.remove(id);
}
for (id, interval) in to_reschedule {
if let Some(mut entry) = self.timers.get_mut(&id) {
entry.fire_at =
Utc::now() + chrono::Duration::from_std(interval).unwrap_or_default();
}
}
due
}
}
impl Default for TimerService {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ipc::MessagePayload;
fn text_payload(msg: &str) -> MessagePayload {
MessagePayload::Text(msg.into())
}
#[test]
fn create_oneshot_timer() {
let service = TimerService::new();
let (id, _token) = service.create_oneshot(1, Duration::from_secs(10), text_payload("fire!"));
assert!(service.info(&id).is_some());
assert_eq!(service.len(), 1);
let info = service.info(&id).unwrap();
assert_eq!(info.owner_pid, 1);
assert!(info.repeat_interval_ms.is_none());
assert!(!info.cancelled);
}
#[test]
fn create_repeating_timer() {
let service = TimerService::new();
let (id, _token) =
service.create_repeating(2, Duration::from_millis(500), text_payload("tick"));
let info = service.info(&id).unwrap();
assert_eq!(info.owner_pid, 2);
assert_eq!(info.repeat_interval_ms, Some(500));
}
#[test]
fn cancel_timer() {
let service = TimerService::new();
let (id, _token) = service.create_oneshot(1, Duration::from_secs(10), text_payload("x"));
assert!(service.cancel(&id));
let info = service.info(&id).unwrap();
assert!(info.cancelled);
}
#[test]
fn cancel_nonexistent() {
let service = TimerService::new();
assert!(!service.cancel("nope"));
}
#[test]
fn cancel_for_pid() {
let service = TimerService::new();
service.create_oneshot(1, Duration::from_secs(10), text_payload("a"));
service.create_oneshot(1, Duration::from_secs(20), text_payload("b"));
service.create_oneshot(2, Duration::from_secs(10), text_payload("c"));
let cancelled = service.cancel_for_pid(1);
assert_eq!(cancelled, 2);
let pid2_timers = service.list_for_pid(2);
assert_eq!(pid2_timers.len(), 1);
assert!(!pid2_timers[0].cancelled);
}
#[test]
fn cleanup_cancelled() {
let service = TimerService::new();
let (id, _) = service.create_oneshot(1, Duration::from_secs(10), text_payload("a"));
service.create_oneshot(2, Duration::from_secs(10), text_payload("b"));
service.cancel(&id);
let cleaned = service.cleanup_cancelled();
assert_eq!(cleaned, 1);
assert_eq!(service.len(), 1);
}
#[test]
fn collect_due_oneshot() {
let service = TimerService::new();
service.create_oneshot(5, Duration::from_millis(0), text_payload("now!"));
std::thread::sleep(Duration::from_millis(10));
let due = service.collect_due();
assert_eq!(due.len(), 1);
assert_eq!(due[0].0, 5);
assert!(service.is_empty());
}
#[test]
fn collect_due_repeating_reschedules() {
let service = TimerService::new();
let (id, _) =
service.create_repeating(3, Duration::from_millis(0), text_payload("repeat"));
std::thread::sleep(Duration::from_millis(10));
let due = service.collect_due();
assert_eq!(due.len(), 1);
assert!(service.info(&id).is_some());
assert_eq!(service.len(), 1);
}
#[test]
fn collect_due_skips_cancelled() {
let service = TimerService::new();
let (id, _) = service.create_oneshot(1, Duration::from_millis(0), text_payload("x"));
service.cancel(&id);
std::thread::sleep(Duration::from_millis(10));
let due = service.collect_due();
assert!(due.is_empty());
}
#[test]
fn collect_due_skips_future() {
let service = TimerService::new();
service.create_oneshot(1, Duration::from_secs(3600), text_payload("far future"));
let due = service.collect_due();
assert!(due.is_empty());
assert_eq!(service.len(), 1);
}
#[test]
fn list_and_list_info() {
let service = TimerService::new();
service.create_oneshot(1, Duration::from_secs(10), text_payload("a"));
service.create_oneshot(2, Duration::from_secs(20), text_payload("b"));
assert_eq!(service.list().len(), 2);
assert_eq!(service.list_info().len(), 2);
}
#[test]
fn list_for_pid() {
let service = TimerService::new();
service.create_oneshot(1, Duration::from_secs(10), text_payload("a"));
service.create_oneshot(1, Duration::from_secs(20), text_payload("b"));
service.create_oneshot(2, Duration::from_secs(10), text_payload("c"));
assert_eq!(service.list_for_pid(1).len(), 2);
assert_eq!(service.list_for_pid(2).len(), 1);
assert_eq!(service.list_for_pid(99).len(), 0);
}
#[test]
fn remove_timer() {
let service = TimerService::new();
let (id, _) = service.create_oneshot(1, Duration::from_secs(10), text_payload("x"));
assert_eq!(service.len(), 1);
let entry = service.remove(&id).unwrap();
assert_eq!(entry.owner_pid, 1);
assert!(service.is_empty());
}
#[test]
fn unique_ids() {
let service = TimerService::new();
let (id1, _) = service.create_oneshot(1, Duration::from_secs(10), text_payload("a"));
let (id2, _) = service.create_oneshot(1, Duration::from_secs(10), text_payload("b"));
assert_ne!(id1, id2);
}
#[test]
fn timer_info_serde_roundtrip() {
let info = TimerInfo {
id: "timer-1".into(),
owner_pid: 42,
fire_at: Utc::now(),
repeat_interval_ms: Some(1000),
cancelled: false,
};
let json = serde_json::to_string(&info).unwrap();
let restored: TimerInfo = serde_json::from_str(&json).unwrap();
assert_eq!(restored.id, "timer-1");
assert_eq!(restored.owner_pid, 42);
assert_eq!(restored.repeat_interval_ms, Some(1000));
}
#[test]
fn default_service() {
let service = TimerService::default();
assert!(service.is_empty());
}
}