later 0.0.49

use later::topic::{partition_for_key, JobPartition, PartitionKey, TopicConfig};
use later::{backend::Backend, BackgroundJobServer, Config};
use serde::{Deserialize, Serialize};
use std::{collections::HashMap, sync::Arc};
use tokio::sync::Mutex;

/// The fixed partition count for the demo's "orders" topic. A real
/// application would size this for its expected concurrency; changing it
/// later needs an explicit migration because producers must agree on which
/// partition owns a given key.
pub const PARTITION_COUNT: u32 = 4;
/// Name of the one topic registered by this example.
pub const PARTITION_TOPIC: &str = "orders";
/// Number of recent partition job completions kept for the dashboard log.
const PARTITION_LOG_CAPACITY: usize = 200;

#[derive(Clone)]
pub struct AppContext {
    server_name: String,
    attempts: Arc<Mutex<HashMap<String, usize>>>,
    partition_log: Arc<Mutex<Vec<PartitionLogEntry>>>,
}

#[derive(Debug, Deserialize, Serialize)]
pub struct PrintJob {
    pub message: String,
    pub failures_before_success: usize,
}

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

/// A job that must run in strict order relative to every other job enqueued
/// with the same `key`. Later hashes `key` to a partition, the way a Kafka
/// producer key is hashed by the partitioner; the caller never picks a
/// partition number. Different keys usually land in different partitions
/// and run concurrently.
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct PartitionJob {
    pub key: String,
    pub seq: u32,
    /// Simulated handler work, so a slow key's partition can be seen not
    /// blocking a fast one running at the same time.
    pub delay_ms: u32,
}

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

/// One completed [`PartitionJob`], recorded for the dashboard's partition
/// log so a viewer can see the strict per-key order and cross-partition
/// concurrency as it happens.
#[derive(Debug, Clone, Serialize)]
pub struct PartitionLogEntry {
    pub key: String,
    /// The partition `key` hashed to, shown so a viewer can see which keys
    /// share a partition (and therefore share ordering) without having to
    /// compute the hash themselves.
    pub partition: u32,
    pub seq: u32,
    pub delay_ms: u32,
    pub finished_at_ms: u128,
}

later::background_job! {
    struct Jobs {
        noop: NoopJob,
        print: PrintJob,
        #[topic("orders")]
        partition_job: PartitionJob,
    }
}

pub type Server = BackgroundJobServer<AppContext, Jobs<AppContext>>;

async fn partition_job(ctx: JobsContext<AppContext>, job: PartitionJob) -> anyhow::Result<()> {
    if job.delay_ms > 0 {
        tokio::time::sleep(std::time::Duration::from_millis(job.delay_ms.into())).await;
    }

    let finished_at_ms = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default()
        .as_millis();
    let partition = partition_for_key(&PartitionKey::from(job.key.clone()), PARTITION_COUNT).0;
    let mut log = ctx.app.partition_log.lock().await;
    log.push(PartitionLogEntry {
        key: job.key,
        partition,
        seq: job.seq,
        delay_ms: job.delay_ms,
        finished_at_ms,
    });
    let overflow = log.len().saturating_sub(PARTITION_LOG_CAPACITY);
    if overflow > 0 {
        log.drain(0..overflow);
    }

    Ok(())
}

async fn print_job(ctx: JobsContext<AppContext>, job: PrintJob) -> anyhow::Result<()> {
    let attempt = {
        let mut attempts = ctx.app.attempts.lock().await;
        let attempt = attempts.entry(job.message.clone()).or_default();
        *attempt += 1;
        let attempt = *attempt;
        if attempt > job.failures_before_success {
            attempts.remove(&job.message);
        }
        attempt
    };

    if !job.message.starts_with("bulk-") {
        println!(
            "{} handling {} (attempt {})",
            ctx.app.server_name, job.message, attempt
        );
    }
    if attempt <= job.failures_before_success {
        return Err(anyhow::anyhow!(
            "demo failure for {} on attempt {}",
            job.message,
            attempt
        ));
    }

    Ok(())
}

async fn noop_job(_ctx: JobsContext<AppContext>, _job: NoopJob) -> anyhow::Result<()> {
    Ok(())
}

#[allow(dead_code)]
pub async fn start_server(server_name: &str, backend: Box<dyn Backend>) -> anyhow::Result<Server> {
    start_server_with_workers(server_name, backend, 6, Arc::new(Mutex::new(Vec::new()))).await
}

/// Starts the demo server. `partition_log` is shared with the caller so an
/// HTTP handler can read it directly as [`PartitionJob`]s complete.
pub async fn start_server_with_workers(
    server_name: &str,
    backend: Box<dyn Backend>,
    worker_count: u8,
    partition_log: Arc<Mutex<Vec<PartitionLogEntry>>>,
) -> anyhow::Result<Server> {
    let context = AppContext {
        server_name: server_name.to_string(),
        attempts: Arc::new(Mutex::new(HashMap::new())),
        partition_log,
    };

    JobsBuilder::new(
        Config::builder()
            .context(context)
            .backend(backend)
            .worker_count(worker_count)
            .topics(vec![TopicConfig::new(PARTITION_TOPIC, PARTITION_COUNT)
                .expect("valid partition topic configuration")])
            .build(),
    )
    .with_noop_handler(noop_job)
    .with_print_handler(print_job)
    .with_partition_job_handler(partition_job)
    .build()
    .await
}