use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use std::time::Duration;
use crate::engine::FlowEngine;
use crate::error::Result;
use crate::model::{ScheduledWakeupKind, WorkflowRunSuspension};
use crate::worker::{FlowTask, FlowTaskDispatcher};
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct FlowSchedulerTick {
pub due_waits: Vec<(String, String)>,
pub due_retries: Vec<(String, String)>,
pub enqueued_tasks: usize,
}
impl FlowSchedulerTick {
pub fn has_due_work(&self) -> bool {
!self.due_waits.is_empty() || !self.due_retries.is_empty()
}
}
#[derive(Clone)]
pub struct FlowScheduler {
engine: FlowEngine,
dispatcher: Arc<dyn FlowTaskDispatcher>,
}
impl FlowScheduler {
pub fn new(engine: FlowEngine, dispatcher: Arc<dyn FlowTaskDispatcher>) -> Self {
Self { engine, dispatcher }
}
pub fn engine(&self) -> &FlowEngine {
&self.engine
}
pub fn dispatcher(&self) -> Arc<dyn FlowTaskDispatcher> {
Arc::clone(&self.dispatcher)
}
#[deprecated(since = "0.4.4", note = "use dispatcher()")]
pub fn queue(&self) -> Arc<dyn FlowTaskDispatcher> {
self.dispatcher()
}
pub async fn next_wakeup(&self, now: DateTime<Utc>) -> Result<Option<WorkflowRunSuspension>> {
self.engine.next_wakeup(now).await
}
pub async fn next_wakeup_delay(&self, now: DateTime<Utc>) -> Result<Option<Duration>> {
let Some(wakeup) = self.next_wakeup(now).await? else {
return Ok(None);
};
let Some(scheduled_at) = wakeup.scheduled_at() else {
return Ok(None);
};
Ok(Some(
scheduled_at
.signed_duration_since(now)
.to_std()
.unwrap_or(Duration::ZERO),
))
}
pub async fn enqueue_due_work(&self, now: DateTime<Utc>) -> Result<FlowSchedulerTick> {
let due = self.engine.list_due_wakeups(now).await?;
let due_waits = due
.iter()
.filter(|wakeup| wakeup.kind == ScheduledWakeupKind::Wait)
.map(|wakeup| (wakeup.run_id.clone(), wakeup.subject_id.clone()))
.collect::<Vec<_>>();
let due_retries = due
.into_iter()
.filter(|wakeup| wakeup.kind == ScheduledWakeupKind::Retry)
.map(|wakeup| (wakeup.run_id, wakeup.subject_id))
.collect::<Vec<_>>();
let mut enqueued_tasks = 0usize;
if !due_waits.is_empty() {
self.dispatcher
.dispatch(FlowTask::ResumeDueWaits { now })
.await?;
enqueued_tasks += 1;
}
if !due_retries.is_empty() {
self.dispatcher
.dispatch(FlowTask::ResumeDueRetries { now })
.await?;
enqueued_tasks += 1;
}
Ok(FlowSchedulerTick {
due_waits,
due_retries,
enqueued_tasks,
})
}
}