#![cfg(feature = "sqlite")]
use async_std::sync::{Arc, Mutex};
use later::{backend::SqliteBackend, storage::Sqlite, BackgroundJobServer, Config};
use serde::{Deserialize, Serialize};
use std::{
sync::atomic::{AtomicUsize, Ordering},
time::{Duration, Instant},
};
const JOB_COUNT: usize = 20;
#[cfg(feature = "prometheus")]
fn metric_value(metrics: &str, name: &str, labels: &[(&str, &str)]) -> Option<f64> {
metrics.lines().find_map(|line| {
if !line.starts_with(name)
|| !labels
.iter()
.all(|(key, value)| line.contains(&format!(r#"{key}="{value}""#)))
{
return None;
}
line.split_whitespace().last()?.parse().ok()
})
}
#[derive(Clone)]
struct AppContext {
worker_name: &'static str,
handled_by: Arc<Mutex<Vec<&'static str>>>,
concurrency: Arc<Concurrency>,
handler_delay: Duration,
}
#[derive(Default)]
struct Concurrency {
current: AtomicUsize,
peak: AtomicUsize,
completed: AtomicUsize,
}
#[derive(Debug, Deserialize, Serialize)]
struct DistributedJob;
later::background_job! {
struct DistributedServer {
distributed_job: DistributedJob,
}
}
async fn handle_job(
context: DistributedServerContext<AppContext>,
_job: DistributedJob,
) -> anyhow::Result<()> {
let current = context
.app
.concurrency
.current
.fetch_add(1, Ordering::SeqCst)
+ 1;
context
.app
.concurrency
.peak
.fetch_max(current, Ordering::SeqCst);
context
.app
.handled_by
.lock()
.await
.push(context.app.worker_name);
tokio::time::sleep(context.app.handler_delay).await;
context
.app
.concurrency
.current
.fetch_sub(1, Ordering::SeqCst);
context
.app
.concurrency
.completed
.fetch_add(1, Ordering::SeqCst);
Ok(())
}
async fn start_server(
namespace: &str,
database_url: &str,
context: AppContext,
) -> anyhow::Result<BackgroundJobServer<AppContext, DistributedServer<AppContext>>> {
let storage = Sqlite::new(database_url).await?;
let backend = SqliteBackend::new(namespace, storage)?;
DistributedServerBuilder::new(
Config::builder()
.context(context)
.backend(Box::new(backend))
.worker_count(1)
.build(),
)
.with_distributed_job_handler(handle_job)
.build()
.await
}
#[tokio::test]
async fn servers_in_one_namespace_distribute_jobs() -> anyhow::Result<()> {
let directory = tempfile::tempdir()?;
let database_url = format!("sqlite://{}", directory.path().join("jobs.db").display());
let namespace = format!("distribution-test-{}", later::generate_id());
let handled_by = Arc::new(Mutex::new(Vec::new()));
let concurrency = Arc::new(Concurrency::default());
let first = start_server(
&namespace,
&database_url,
AppContext {
worker_name: "first",
handled_by: handled_by.clone(),
concurrency: concurrency.clone(),
handler_delay: Duration::from_millis(20),
},
)
.await?;
let second = start_server(
&namespace,
&database_url,
AppContext {
worker_name: "second",
handled_by: handled_by.clone(),
concurrency,
handler_delay: Duration::from_millis(20),
},
)
.await?;
for _ in 0..JOB_COUNT {
first.enqueue(DistributedJob).await?;
}
let deadline = Instant::now() + Duration::from_secs(10);
while handled_by.lock().await.len() < JOB_COUNT && Instant::now() < deadline {
tokio::time::sleep(Duration::from_millis(25)).await;
}
let handled_by = handled_by.lock().await.clone();
let first_count = handled_by.iter().filter(|name| **name == "first").count();
let second_count = handled_by.iter().filter(|name| **name == "second").count();
assert_eq!(
(handled_by.len(), first_count > 0, second_count > 0),
(JOB_COUNT, true, true)
);
#[cfg(feature = "prometheus")]
{
let queue = format!("later-{namespace}");
let deadline = Instant::now() + Duration::from_secs(2);
let metrics = loop {
let metrics = first.get_metrics()?;
if metric_value(
&metrics,
"later_job_transitions_total",
&[
("queue", &queue),
("job_type", "distributed_job"),
("stage", "success"),
],
) == Some(JOB_COUNT as f64)
|| Instant::now() >= deadline
{
break metrics;
}
tokio::time::sleep(Duration::from_millis(10)).await;
};
assert_eq!(
(
metric_value(
&metrics,
"later_job_transitions_total",
&[
("queue", &queue),
("job_type", "distributed_job"),
("stage", "success"),
],
),
metric_value(
&metrics,
"later_job_handler_duration_seconds_count",
&[
("queue", &queue),
("worker_id", "0"),
("job_type", "distributed_job"),
("outcome", "success"),
],
),
metric_value(
&metrics,
"later_job_duration_seconds_count",
&[
("queue", &queue),
("job_type", "distributed_job"),
("outcome", "success"),
],
),
metric_value(
&metrics,
"later_job_wait_duration_seconds_count",
&[
("queue", &queue),
("job_type", "distributed_job"),
("outcome", "success"),
],
),
metric_value(
&metrics,
"later_worker_commands_total",
&[
("queue", &queue),
("worker_id", "0"),
("command", "execute_job"),
("outcome", "success"),
],
),
metric_value(
&metrics,
"later_workers",
&[("queue", &queue), ("worker_id", "0"), ("state", "active"),],
),
),
(
Some(JOB_COUNT as f64),
Some(JOB_COUNT as f64),
Some(JOB_COUNT as f64),
Some(JOB_COUNT as f64),
Some(JOB_COUNT as f64),
Some(2.0),
),
);
}
drop(first);
drop(second);
Ok(())
}
#[tokio::test]
async fn workers_can_be_added_and_removed_while_the_server_is_running() -> anyhow::Result<()> {
let directory = tempfile::tempdir()?;
let database_url = format!("sqlite://{}", directory.path().join("workers.db").display());
let namespace = format!("worker-scaling-test-{}", later::generate_id());
let handled_by = Arc::new(Mutex::new(Vec::new()));
let concurrency = Arc::new(Concurrency::default());
let server = start_server(
&namespace,
&database_url,
AppContext {
worker_name: "scalable",
handled_by: handled_by.clone(),
concurrency: concurrency.clone(),
handler_delay: Duration::from_millis(200),
},
)
.await?;
assert_eq!((server.worker_count(), server.add_worker().await?), (1, 2));
for _ in 0..2 {
server.enqueue(DistributedJob).await?;
}
wait_for_running_jobs(&concurrency, 2).await;
assert_eq!(concurrency.peak.load(Ordering::SeqCst), 2);
assert_eq!(server.remove_worker().await?, 1);
assert!(concurrency.completed.load(Ordering::SeqCst) >= 1);
handled_by.lock().await.clear();
concurrency.peak.store(0, Ordering::SeqCst);
for _ in 0..2 {
server.enqueue(DistributedJob).await?;
}
wait_for_handled_jobs(&handled_by, 2).await;
assert_eq!(
(
server.worker_count(),
concurrency.peak.load(Ordering::SeqCst)
),
(1, 1)
);
assert!(server.remove_worker().await.is_err());
Ok(())
}
#[tokio::test]
async fn graceful_shutdown_waits_for_a_running_job_to_commit() -> anyhow::Result<()> {
let directory = tempfile::tempdir()?;
let database_url = format!(
"sqlite://{}",
directory.path().join("shutdown.db").display()
);
let handled_by = Arc::new(Mutex::new(Vec::new()));
let concurrency = Arc::new(Concurrency::default());
let server = start_server(
"graceful-shutdown",
&database_url,
AppContext {
worker_name: "draining",
handled_by,
concurrency: concurrency.clone(),
handler_delay: Duration::from_millis(100),
},
)
.await?;
server.enqueue(DistributedJob).await?;
wait_for_running_jobs(&concurrency, 1).await;
server.shutdown(Duration::from_secs(2)).await?;
assert_eq!(concurrency.completed.load(Ordering::SeqCst), 1);
Ok(())
}
async fn wait_for_handled_jobs(handled_by: &Arc<Mutex<Vec<&'static str>>>, expected: usize) {
let deadline = Instant::now() + Duration::from_secs(10);
while handled_by.lock().await.len() < expected && Instant::now() < deadline {
tokio::time::sleep(Duration::from_millis(10)).await;
}
}
async fn wait_for_running_jobs(concurrency: &Concurrency, expected: usize) {
let deadline = Instant::now() + Duration::from_secs(10);
while concurrency.current.load(Ordering::SeqCst) < expected && Instant::now() < deadline {
tokio::time::sleep(Duration::from_millis(10)).await;
}
}