use crate::{TimerNameLike, TimerRegistration, TimerResult, TimerTarget, TimerWrapper, timer_iter};
use anyhow::{Result, anyhow};
use framework_core::current_millis;
use futures_util::FutureExt;
use std::any::Any;
use std::collections::HashMap;
use std::panic::AssertUnwindSafe;
use std::sync::{Arc, OnceLock};
use tokio::sync::Semaphore;
use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel};
use tokio::task::{Id, JoinSet};
use tokio::time::{Duration, sleep};
static WAKE_SENDER: OnceLock<UnboundedSender<String>> = OnceLock::new();
#[derive(Clone, Copy, Debug, Default)]
pub struct TimerOption {
pub max_concurrent_tasks: Option<usize>,
pub retry_immediately_on_panic: bool,
}
struct TimerState {
name: String,
registration: &'static TimerRegistration,
next_run_at: i64,
running: bool,
}
pub fn start(wrapper: Option<TimerWrapper>, option: TimerOption) {
let semaphore = option.max_concurrent_tasks.map(|count| {
assert!(count > 0, "max_concurrent_tasks 必须大于 0");
Arc::new(Semaphore::new(count))
});
let states = timer_iter()
.flat_map(|registration| {
(1..=registration.count).map(move |index| TimerState {
name: if registration.count == 1 {
registration.name.to_owned()
} else {
format!("{}_{index}", registration.name)
},
registration,
next_run_at: 0,
running: false,
})
})
.collect::<Vec<_>>();
let (wake_sender, wake_receiver) = unbounded_channel();
WAKE_SENDER
.set(wake_sender)
.expect("定时任务运行时只能启动一次");
tokio::spawn(run_scheduler(
states,
wake_receiver,
wrapper,
option,
semaphore,
));
}
pub fn wake(name: impl TimerNameLike) {
let Some(sender) = WAKE_SENDER.get() else {
return;
};
let _ = sender.send(name.timer_name().to_owned());
}
async fn run_scheduler(
mut states: Vec<TimerState>,
mut wake_receiver: UnboundedReceiver<String>,
wrapper: Option<TimerWrapper>,
option: TimerOption,
semaphore: Option<Arc<Semaphore>>,
) {
let mut tasks = JoinSet::new();
let mut task_indices = HashMap::new();
loop {
let wait_millis = dispatch_due_tasks(
&mut states,
&mut tasks,
&mut task_indices,
&wrapper,
semaphore.as_ref(),
current_millis().expect("获取当前毫秒时间戳失败"),
);
match wait_millis {
Some(wait_millis) => {
tokio::select! {
Some(name) = wake_receiver.recv() => wake_timer(&mut states, name),
result = tasks.join_next_with_id(), if !tasks.is_empty() => {
if let Some(result) = result {
handle_task_result(result, &mut task_indices, &mut states, option);
}
}
_ = sleep(Duration::from_millis(wait_millis)) => {}
else => return,
}
}
None => {
tokio::select! {
Some(name) = wake_receiver.recv() => wake_timer(&mut states, name),
result = tasks.join_next_with_id() => {
if let Some(result) = result {
handle_task_result(result, &mut task_indices, &mut states, option);
}
}
else => return,
}
}
}
}
}
fn wake_timer(states: &mut [TimerState], name: String) {
if let Some(state) = states
.iter_mut()
.find(|state| !state.running && state.registration.name == name)
{
state.next_run_at = -10086;
}
}
fn dispatch_due_tasks(
states: &mut [TimerState],
tasks: &mut JoinSet<Result<TimerResult>>,
task_indices: &mut HashMap<Id, usize>,
wrapper: &Option<TimerWrapper>,
semaphore: Option<&Arc<Semaphore>>,
now: i64,
) -> Option<u64> {
let mut next_run_at: Option<i64> = None;
for (index, state) in states.iter_mut().enumerate() {
if state.running {
continue;
}
if state.next_run_at > now {
next_run_at =
Some(next_run_at.map_or(state.next_run_at, |next| next.min(state.next_run_at)));
continue;
}
state.running = true;
let target = state.registration.target;
let wrapper = wrapper.clone();
let semaphore = semaphore.cloned();
let task = tasks.spawn(async move {
execute_timer(target, wrapper.as_ref(), semaphore.as_ref()).await
});
task_indices.insert(task.id(), index);
}
next_run_at.map(|next_run_at| next_run_at.saturating_sub(now).max(0) as u64)
}
fn handle_task_result(
result: Result<(Id, Result<TimerResult>), tokio::task::JoinError>,
task_indices: &mut HashMap<Id, usize>,
states: &mut [TimerState],
option: TimerOption,
) {
let (task_id, result) = match result {
Ok((task_id, result)) => (task_id, result),
Err(error) => {
log::error!("定时任务 Tokio 子任务执行失败：{error:?}");
(
error.id(),
Err(anyhow!("定时任务 Tokio 子任务执行失败：{error:?}")),
)
}
};
let index = task_indices
.remove(&task_id)
.expect("定时任务缺少任务实例索引");
let state = &mut states[index];
state.running = false;
let now = current_millis().expect("获取当前毫秒时间戳失败");
let next = match result {
Ok(result) => {
if let Some(error) = result.error {
log::error!("定时任务 {} 执行失败：{error:?}", state.name);
}
match result.next {
None => now.saturating_add(state.registration.interval),
Some(v) => v,
}
}
Err(error) => {
log::error!("定时任务 {} 发生 panic: {error:?}", state.name);
if state
.registration
.retry_immediately_on_panic
.unwrap_or(option.retry_immediately_on_panic)
{
0
} else {
now.saturating_add(state.registration.interval)
}
}
};
state.next_run_at = next
}
async fn execute_timer(
target: TimerTarget,
wrapper: Option<&TimerWrapper>,
semaphore: Option<&Arc<Semaphore>>,
) -> Result<TimerResult> {
let _permit = match semaphore {
Some(semaphore) => Some(
semaphore
.acquire()
.await
.map_err(|_| anyhow!("定时任务并发控制 Semaphore 已关闭"))?,
),
None => None,
};
let future = match wrapper {
Some(wrapper) => wrapper(target),
None => target(),
};
AssertUnwindSafe(future)
.catch_unwind()
.await
.map_err(panic_to_error)
}
fn panic_to_error(payload: Box<dyn Any + Send>) -> anyhow::Error {
let message = if let Some(message) = payload.downcast_ref::<&str>() {
*message
} else if let Some(message) = payload.downcast_ref::<String>() {
message.as_str()
} else {
"未知 panic payload"
};
anyhow!("定时任务执行发生 panic: {message}")
}