use std::collections::HashMap;
use std::time::Duration;
use once_cell::sync::Lazy;
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use tokio::sync::{mpsc, Mutex};
use tokio::task;
use tracing::{error, info};
use crate::core::Result;
use super::engine::FlowEngine;
use super::id;
use super::queue_disk;
use super::types::{GraphId, GraphStatus};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QueueMessage {
pub id: String,
pub worker: String,
#[serde(default)]
pub input: Value,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EnqueueResponse {
pub message_id: String,
pub queue: String,
}
static WAKES: Lazy<RwLock<HashMap<String, mpsc::Sender<()>>>> =
Lazy::new(|| RwLock::new(HashMap::new()));
static STARTED: Lazy<Mutex<bool>> = Lazy::new(|| Mutex::new(false));
fn poll_interval_ms() -> u64 {
std::env::var("RESUMA_QUEUE_POLL_MS")
.ok()
.and_then(|s| s.parse().ok())
.unwrap_or(200)
}
pub fn configure_disk(root: impl AsRef<std::path::Path>) {
queue_disk::configure(root);
}
pub async fn start_consumers() {
let mut started = STARTED.lock().await;
if *started {
return;
}
*started = true;
register_queue("default");
for name in queue_disk::list_queues() {
if name != "default" {
register_queue(&name);
}
}
}
pub fn register_queue(name: &str) {
let mut wakes = WAKES.write();
if wakes.contains_key(name) {
return;
}
let (tx, rx) = mpsc::channel(512);
wakes.insert(name.to_string(), tx);
drop(wakes);
let qname = name.to_string();
task::spawn(async move {
disk_consumer_loop(&qname, rx).await;
});
}
pub async fn enqueue(queue: &str, worker: &str, input: Value) -> Result<EnqueueResponse> {
register_queue(queue);
let msg = QueueMessage {
id: format!("m_{}", id::next_id()),
worker: worker.to_string(),
input,
};
let mid = msg.id.clone();
queue_disk::persist_pending(queue, &msg)?;
wake(queue);
Ok(EnqueueResponse {
message_id: mid,
queue: queue.to_string(),
})
}
pub fn queue_stats(queue: &str) -> queue_disk::QueueStats {
queue_disk::stats(queue)
}
fn wake(queue: &str) {
if let Some(tx) = WAKES.read().get(queue) {
let _ = tx.try_send(());
}
}
async fn disk_consumer_loop(queue: &str, mut wake_rx: mpsc::Receiver<()>) {
let recovered = queue_disk::recover_processing(queue);
if recovered > 0 {
info!(queue = %queue, recovered, "re-queued stuck processing jobs");
}
info!(queue = %queue, "resuma disk queue consumer started (multi-process)");
let poll = Duration::from_millis(poll_interval_ms());
loop {
drain_claimed(queue).await;
tokio::select! {
_ = wake_rx.recv() => {}
_ = tokio::time::sleep(poll) => {}
}
}
}
async fn drain_claimed(queue: &str) {
while let Some(msg) = queue_disk::claim_next(queue) {
process_message(queue, msg).await;
}
}
async fn process_message(queue: &str, msg: QueueMessage) {
info!(
queue = %queue,
worker = %msg.worker,
message_id = %msg.id,
"claimed queued job from disk"
);
match FlowEngine::start(&msg.worker, msg.input.clone()).await {
Ok(started) => {
info!(
graph_id = %started.graph_id.0,
worker = %msg.worker,
"queued worker started"
);
let queue = queue.to_string();
let msg_id = msg.id.clone();
let graph_id = started.graph_id;
tokio::spawn(async move {
let success = wait_graph_terminal(&graph_id).await;
if let Err(e) = queue_disk::complete(&queue, &msg_id, success) {
tracing::error!(
error = %e,
queue = %queue,
message_id = %msg_id,
"failed to complete queued job"
);
}
});
}
Err(e) => {
error!(error = %e, worker = %msg.worker, "queued worker failed to start");
let _ = queue_disk::complete(queue, &msg.id, false);
}
}
}
async fn wait_graph_terminal(graph_id: &GraphId) -> bool {
let poll = Duration::from_millis(100);
let timeout_secs = std::env::var("RESUMA_QUEUE_GRAPH_TIMEOUT_SECS")
.ok()
.and_then(|v| v.parse().ok())
.unwrap_or(3600);
let deadline = tokio::time::Instant::now() + Duration::from_secs(timeout_secs);
loop {
match FlowEngine::snapshot(graph_id) {
Some(snap) => match snap.status {
GraphStatus::Done => return true,
GraphStatus::Failed => return false,
GraphStatus::Running | GraphStatus::Paused => {}
},
None => return false,
}
if tokio::time::Instant::now() >= deadline {
tracing::warn!(
graph_id = %graph_id.0,
timeout_secs,
"queued job graph wait timed out"
);
return false;
}
tokio::time::sleep(poll).await;
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[allow(clippy::await_holding_lock)]
#[tokio::test]
async fn enqueue_persists_for_disk_consumer() {
let _guard = super::queue_disk::test_queue_lock().lock();
let root = std::env::temp_dir().join(format!("resuma-qe-{}", id::next_id()));
configure_disk(&root);
register_queue("test");
let resp = enqueue("test", "missing_worker", json!({}))
.await
.expect("enqueue");
let s = queue_stats("test");
assert!(s.pending >= 1 || s.processing >= 1);
assert!(!resp.message_id.is_empty());
}
}