later 0.0.49

//! Measures end-to-end Later throughput with normal and partitioned jobs.
//!
//! Run SQLite without external services (only the `sqlite` feature is
//! needed):
//!
//! ```text
//! RUSTC_WRAPPER= cargo run -p later --example throughput --features sqlite --release -- --backend sqlite
//! ```
//!
//! Run PostgreSQL after `make init` and `eval "$(make env)"` (only the
//! `postgres` feature is needed):
//!
//! ```text
//! RUSTC_WRAPPER= cargo run -p later --example throughput --features postgres --release -- --backend postgres
//! ```
//!
//! Environment options for the SQLite backend:
//!
//! - `THROUGHPUT_SQLITE_SYNCHRONOUS=normal` runs on an application-style pool
//!   (WAL, `synchronous = NORMAL`) instead of the library's own default, which
//!   is how Later runs in production and what to use when comparing CPU cost.
//! - `POOL=<n>` sets that pool's connection count (default 10).
//! - `CENSUS=1` prints how many SQL statements each job costs, by statement.
//!
//! CPU per job is the figure to watch when tuning: wrap the run in `time` and
//! divide user+system by the job count.
//!
//! The default workload enqueues 10,000 jobs with an even normal/partitioned
//! mix, 64 sequential partitions, eight workers, and eight concurrent
//! producers. Use `--help` for the available tuning options.

use later::{
    backend::Backend,
    topic::{JobPartition, PartitionKey, TopicConfig},
    Config,
};
use serde::{Deserialize, Serialize};
use std::{
    sync::{
        atomic::{AtomicUsize, Ordering},
        Arc,
    },
    time::{Duration, Instant},
};

mod census {
    use std::{collections::HashMap, sync::Mutex};
    use tracing::{
        field::{Field, Visit},
        span, Event, Metadata, Subscriber,
    };
    pub static COUNTS: Mutex<Option<HashMap<String, usize>>> = Mutex::new(None);
    struct V(Option<String>);
    impl Visit for V {
        fn record_debug(&mut self, f: &Field, v: &dyn std::fmt::Debug) {
            if f.name() == "db.statement" {
                self.0 = Some(format!("{v:?}"));
            }
        }
        fn record_str(&mut self, f: &Field, v: &str) {
            if f.name() == "db.statement" {
                self.0 = Some(v.to_string());
            }
        }
    }
    pub struct Census;
    impl Subscriber for Census {
        fn enabled(&self, m: &Metadata<'_>) -> bool {
            m.target() == "sqlx::query"
        }
        fn new_span(&self, _: &span::Attributes<'_>) -> span::Id {
            span::Id::from_u64(1)
        }
        fn record(&self, _: &span::Id, _: &span::Record<'_>) {}
        fn record_follows_from(&self, _: &span::Id, _: &span::Id) {}
        fn event(&self, e: &Event<'_>) {
            let mut v = V(None);
            e.record(&mut v);
            if let Some(sql) = v.0 {
                let key: String = sql
                    .split_whitespace()
                    .collect::<Vec<_>>()
                    .join(" ")
                    .chars()
                    .take(110)
                    .collect();
                *COUNTS
                    .lock()
                    .unwrap()
                    .get_or_insert_with(HashMap::new)
                    .entry(key)
                    .or_insert(0) += 1;
            }
        }
        fn enter(&self, _: &span::Id) {}
        fn exit(&self, _: &span::Id) {}
    }
}

/// The leaked SQLite directory of this run, removed once it finishes.
static BENCH_DIRECTORY: std::sync::OnceLock<std::path::PathBuf> = std::sync::OnceLock::new();

const TOPIC: &str = "throughput";

#[derive(Debug, Deserialize, Serialize)]
struct NormalJob;

#[derive(Debug, Deserialize, Serialize)]
struct SequentialJob {
    key: String,
}

impl JobPartition for SequentialJob {
    fn partition_key(&self) -> PartitionKey {
        PartitionKey::from(self.key.clone())
    }
}

later::background_job! {
    struct BenchmarkJobs {
        normal: NormalJob,
        #[topic("throughput")]
        sequential: SequentialJob,
    }
}

struct BenchmarkContext {
    normal_completed: AtomicUsize,
    sequential_completed: AtomicUsize,
    completed: tokio::sync::Notify,
}

async fn normal(
    ctx: BenchmarkJobsContext<Arc<BenchmarkContext>>,
    _job: NormalJob,
) -> anyhow::Result<()> {
    ctx.app.normal_completed.fetch_add(1, Ordering::Relaxed);
    ctx.app.completed.notify_waiters();
    Ok(())
}

async fn sequential(
    ctx: BenchmarkJobsContext<Arc<BenchmarkContext>>,
    _job: SequentialJob,
) -> anyhow::Result<()> {
    ctx.app.sequential_completed.fetch_add(1, Ordering::Relaxed);
    ctx.app.completed.notify_waiters();
    Ok(())
}

#[derive(Clone, Copy)]
enum BackendKind {
    Sqlite,
    Postgres,
}

struct Options {
    backend: BackendKind,
    normal_jobs: usize,
    sequential_jobs: usize,
    partitions: u32,
    workers: u8,
    producer_concurrency: usize,
}

impl Options {
    fn parse() -> anyhow::Result<Self> {
        let mut backend = BackendKind::Sqlite;
        let mut normal_jobs = 5_000;
        let mut sequential_jobs = 5_000;
        let mut partitions = 64;
        let mut workers = 8;
        let mut producer_concurrency = 8;
        let mut arguments = std::env::args().skip(1);

        while let Some(argument) = arguments.next() {
            match argument.as_str() {
                "--backend" => {
                    let value = arguments
                        .next()
                        .ok_or_else(|| anyhow::anyhow!("--backend needs sqlite or postgres"))?;
                    backend = match value.as_str() {
                        "sqlite" => BackendKind::Sqlite,
                        "postgres" => BackendKind::Postgres,
                        _ => return Err(anyhow::anyhow!("--backend needs sqlite or postgres")),
                    };
                }
                "--normal-jobs" => {
                    normal_jobs = parse_number(&mut arguments, "--normal-jobs")?;
                }
                "--sequential-jobs" => {
                    sequential_jobs = parse_number(&mut arguments, "--sequential-jobs")?;
                }
                "--partitions" => {
                    partitions = parse_number(&mut arguments, "--partitions")?;
                }
                "--workers" => {
                    workers = parse_number(&mut arguments, "--workers")?;
                }
                "--producer-concurrency" => {
                    producer_concurrency = parse_number(&mut arguments, "--producer-concurrency")?;
                }
                "--help" | "-h" => {
                    print_help();
                    std::process::exit(0);
                }
                _ => return Err(anyhow::anyhow!("unknown argument: {argument}")),
            }
        }

        if normal_jobs + sequential_jobs == 0 {
            return Err(anyhow::anyhow!("at least one job is required"));
        }
        if sequential_jobs > 0 && partitions == 0 {
            return Err(anyhow::anyhow!("--partitions must be at least 1"));
        }
        if workers == 0 {
            return Err(anyhow::anyhow!("--workers must be at least 1"));
        }
        if producer_concurrency == 0 {
            return Err(anyhow::anyhow!("--producer-concurrency must be at least 1"));
        }

        Ok(Self {
            backend,
            normal_jobs,
            sequential_jobs,
            partitions,
            workers,
            producer_concurrency,
        })
    }
}

fn parse_number<T>(
    arguments: &mut impl Iterator<Item = String>,
    argument: &str,
) -> anyhow::Result<T>
where
    T: std::str::FromStr,
    T::Err: std::fmt::Display,
{
    let value = arguments
        .next()
        .ok_or_else(|| anyhow::anyhow!("{argument} needs a number"))?;
    value
        .parse()
        .map_err(|error| anyhow::anyhow!("{argument} needs a number: {error}"))
}

fn print_help() {
    println!(
        "Usage: throughput [--backend sqlite|postgres] [--normal-jobs N] \\\n+         [--sequential-jobs N] [--partitions N] [--workers N] \\\n+         [--producer-concurrency N]"
    );
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let options = Options::parse()?;
    if std::env::var("CENSUS").is_ok() {
        tracing::subscriber::set_global_default(census::Census)?;
    }
    let backend = match options.backend {
        BackendKind::Sqlite => create_sqlite_backend().await?,
        BackendKind::Postgres => create_postgres_backend().await?,
    };
    let jobs = (options.normal_jobs + options.sequential_jobs) as f64;
    let result = run(backend, options).await;
    if let Some(directory) = BENCH_DIRECTORY.get() {
        let _ = std::fs::remove_dir_all(directory);
    }
    result?;
    if let Some(counts) = census::COUNTS.lock().unwrap().take() {
        let mut rows: Vec<_> = counts.into_iter().collect();
        rows.sort_by(|a, b| b.1.cmp(&a.1));
        let total: usize = rows.iter().map(|r| r.1).sum();
        println!("statements per job: {:.1}", total as f64 / jobs);
        for (sql, n) in rows.iter().take(40) {
            println!("{:6.2}  {}", *n as f64 / jobs, sql);
        }
    }
    Ok(())
}

#[cfg(feature = "sqlite")]
async fn create_sqlite_backend() -> anyhow::Result<Box<dyn Backend>> {
    // Intentionally leaked: the pool opens further connections to this file
    // for as long as the benchmark runs, well after this function returns,
    // so the directory must outlive it rather than being deleted here.
    let database_directory = tempfile::tempdir()?.into_path();
    // Removed when the run ends (see `main`): left behind, each run kept a
    // database file and a write-ahead log of a gigabyte or more in /tmp.
    let _ = BENCH_DIRECTORY.set(database_directory.clone());
    let database_url = format!("sqlite://{}", database_directory.join("later.db").display());
    // `THROUGHPUT_SQLITE_SYNCHRONOUS=normal` mirrors an application that
    // owns its pool with WAL's usual `synchronous = NORMAL` (no fsync per
    // commit), which is how Later is run in production. The default keeps
    // the library's own `Sqlite::new` behaviour so results stay comparable.
    if std::env::var("THROUGHPUT_SQLITE_SYNCHRONOUS").as_deref() == Ok("normal") {
        use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqliteSynchronous};
        use std::str::FromStr;
        let options = SqliteConnectOptions::from_str(&database_url)?
            .create_if_missing(true)
            .foreign_keys(true)
            .busy_timeout(std::time::Duration::from_secs(30))
            .journal_mode(SqliteJournalMode::Wal)
            .synchronous(SqliteSynchronous::Normal);
        let pool = sqlx::sqlite::SqlitePoolOptions::new()
            .max_connections(
                std::env::var("POOL")
                    .ok()
                    .and_then(|v| v.parse().ok())
                    .unwrap_or(10),
            )
            .connect_with(options)
            .await?;
        let backend = later::backend::SqliteBackend::from_pool("throughput", pool).await?;
        return Ok(Box::new(backend));
    }
    let backend = later::backend::SqliteBackend::connect("throughput", &database_url).await?;
    Ok(Box::new(backend))
}

#[cfg(not(feature = "sqlite"))]
async fn create_sqlite_backend() -> anyhow::Result<Box<dyn Backend>> {
    Err(anyhow::anyhow!("SQLite needs the sqlite Cargo feature"))
}

#[cfg(feature = "postgres")]
async fn create_postgres_backend() -> anyhow::Result<Box<dyn Backend>> {
    let database_url = std::env::var("DATABASE_URL")
        .map_err(|_| anyhow::anyhow!("DATABASE_URL is required for the PostgreSQL benchmark"))?;
    let namespace = format!("throughput-{}", later::generate_id());
    let backend = later::backend::PostgresBackend::connect(namespace, &database_url).await?;
    Ok(Box::new(backend))
}

#[cfg(not(feature = "postgres"))]
async fn create_postgres_backend() -> anyhow::Result<Box<dyn Backend>> {
    Err(anyhow::anyhow!(
        "PostgreSQL needs the postgres Cargo feature"
    ))
}

async fn run(backend: Box<dyn Backend>, options: Options) -> anyhow::Result<()> {
    let context = Arc::new(BenchmarkContext {
        normal_completed: AtomicUsize::new(0),
        sequential_completed: AtomicUsize::new(0),
        completed: tokio::sync::Notify::new(),
    });
    let topics = if options.sequential_jobs == 0 {
        Vec::new()
    } else {
        vec![TopicConfig::new(TOPIC, options.partitions)?]
    };
    let server = Arc::new(
        BenchmarkJobsBuilder::new(
            Config::builder()
                .context(context.clone())
                .backend(backend)
                .worker_count(options.workers)
                .topics(topics)
                .build(),
        )
        .with_normal_handler(normal)
        .with_sequential_handler(sequential)
        .build()
        .await?,
    );

    let total_jobs = options.normal_jobs + options.sequential_jobs;
    let started = Instant::now();
    let mut producers = Vec::with_capacity(options.producer_concurrency);
    for producer in 0..options.producer_concurrency {
        let server = server.clone();
        let sequential_jobs = options.sequential_jobs;
        let partitions = options.partitions;
        producers.push(tokio::spawn(async move {
            for job_index in (producer..total_jobs).step_by(options.producer_concurrency) {
                let sequential_before = job_index.saturating_mul(sequential_jobs) / total_jobs;
                let sequential_after = (job_index + 1).saturating_mul(sequential_jobs) / total_jobs;
                if sequential_after > sequential_before {
                    let key = format!("key-{}", sequential_before % usize::try_from(partitions)?);
                    server.enqueue_to_partition(SequentialJob { key }).await?;
                } else {
                    server.enqueue(NormalJob).await?;
                }
            }
            Ok::<(), anyhow::Error>(())
        }));
    }
    for producer in producers {
        producer.await??;
    }
    let enqueue_elapsed = started.elapsed();

    let completion_timeout = Duration::from_secs(60 + (total_jobs as u64 / 50));
    tokio::time::timeout(completion_timeout, async {
        loop {
            let completed = context.normal_completed.load(Ordering::Relaxed)
                + context.sequential_completed.load(Ordering::Relaxed);
            if completed == total_jobs {
                return;
            }
            context.completed.notified().await;
        }
    })
    .await
    .map_err(|_| anyhow::anyhow!("timed out waiting for {total_jobs} jobs to complete"))?;

    let total_elapsed = started.elapsed();
    let normal_completed = context.normal_completed.load(Ordering::Relaxed);
    let sequential_completed = context.sequential_completed.load(Ordering::Relaxed);
    println!("backend: {}", backend_name(options.backend));
    println!(
        "workers: {}, producers: {}, partitions: {}",
        options.workers, options.producer_concurrency, options.partitions
    );
    println!("completed: {normal_completed} normal + {sequential_completed} sequential");
    println!(
        "enqueue: {:.0} jobs/s ({:.2}s)",
        total_jobs as f64 / enqueue_elapsed.as_secs_f64(),
        enqueue_elapsed.as_secs_f64()
    );
    println!(
        "end-to-end: {:.0} jobs/s ({:.2}s)",
        total_jobs as f64 / total_elapsed.as_secs_f64(),
        total_elapsed.as_secs_f64()
    );
    Ok(())
}

fn backend_name(backend: BackendKind) -> &'static str {
    match backend {
        BackendKind::Sqlite => "sqlite",
        BackendKind::Postgres => "postgres",
    }
}