later 0.0.48

//! Named topics and Kafka-like partitions for strictly ordered jobs.
//!
//! Jobs enqueued without a topic keep Later's default unordered,
//! competing-worker delivery. A job type opts into ordering by implementing
//! [`crate::topic::JobPartition`] and being declared with a topic in
//! [`crate::background_job!`]. Only
//! [`crate::BackgroundJobServerPublisher::enqueue_to_partition`] and its
//! sibling methods can request partition ordering, and only for message types
//! that implement both [`crate::topic::JobPartition`] and the topic marker the
//! macro generates.

use std::fmt::Display;

/// Reserved topic name sequential recurring jobs share.
///
/// Every occurrence of a sequential recurring job — and every job chained as
/// a continuation of one — is enqueued into this topic, partitioned by the
/// recurring job's own identifier (see
/// [`crate::BackgroundJobServerPublisher::enqueue_recurring_sequential`]).
/// It is always registered at server startup; application code never
/// enqueues to it directly.
pub(crate) const RECURRING_SEQUENTIAL_TOPIC: &str = "__later_recurring__";

/// The name of a configured topic.
///
/// Topic names are namespace-scoped and must be registered with
/// [`crate::Config`] before jobs can be enqueued to them.
#[derive(Clone, Debug, Eq, PartialEq, Hash, serde::Serialize, serde::Deserialize)]
pub struct TopicName(pub String);

impl Display for TopicName {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(&self.0)
    }
}

impl From<&str> for TopicName {
    fn from(value: &str) -> Self {
        Self(value.to_owned())
    }
}

impl From<String> for TopicName {
    fn from(value: String) -> Self {
        Self(value)
    }
}

/// An arbitrary key used to route a message to a partition, the way a Kafka
/// producer key does.
///
/// Later hashes this key to pick a partition (see [`partition_for_key`]), so
/// the same key always lands on the same partition for a given topic's
/// partition count, and the caller never chooses a partition number
/// directly.
#[derive(Clone, Debug, Eq, PartialEq, Hash)]
pub struct PartitionKey(pub Vec<u8>);

impl From<&str> for PartitionKey {
    fn from(value: &str) -> Self {
        Self(value.as_bytes().to_vec())
    }
}

impl From<String> for PartitionKey {
    fn from(value: String) -> Self {
        Self(value.into_bytes())
    }
}

impl From<&[u8]> for PartitionKey {
    fn from(value: &[u8]) -> Self {
        Self(value.to_vec())
    }
}

impl From<Vec<u8>> for PartitionKey {
    fn from(value: Vec<u8>) -> Self {
        Self(value)
    }
}

macro_rules! impl_partition_key_from_int {
    ($($ty:ty),+ $(,)?) => {
        $(
            impl From<$ty> for PartitionKey {
                fn from(value: $ty) -> Self {
                    Self(value.to_be_bytes().to_vec())
                }
            }
        )+
    };
}

impl_partition_key_from_int!(u8, u16, u32, u64, u128, i8, i16, i32, i64, i128);

/// A partition identifier within a topic, in the range `0..partition_count`.
///
/// Application code enqueuing jobs never constructs this directly;
/// [`partition_for_key`] is the only way one is produced, from a
/// [`PartitionKey`] and a topic's configured partition count. A custom
/// [`crate::backend::JobCommitter`] implementation may still need to read or
/// construct one directly, which is why the field stays public.
#[derive(
    Clone, Copy, Debug, Eq, PartialEq, Hash, PartialOrd, Ord, serde::Serialize, serde::Deserialize,
)]
pub struct PartitionId(pub u32);

impl Display for PartitionId {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.0)
    }
}

/// Deterministically maps a [`PartitionKey`] to a partition number.
///
/// Uses a stable hash (BLAKE3) so every process and every server computes the
/// same partition for the same key and partition count, the same way a
/// Kafka producer's default partitioner routes by key rather than letting
/// the caller pick a partition number. `partition_count` must be positive;
/// zero is treated as one to avoid a division by zero.
pub fn partition_for_key(key: &PartitionKey, partition_count: u32) -> PartitionId {
    let partition_count = u64::from(partition_count.max(1));
    let hash = blake3::hash(&key.0);
    let bucket = u64::from_le_bytes(
        hash.as_bytes()[0..8]
            .try_into()
            .expect("BLAKE3 digest is at least 8 bytes"),
    );
    PartitionId((bucket % partition_count) as u32)
}

/// Returns the key used to route this message to a partition.
///
/// [`crate::background_job!`] requires this trait for any message declared
/// with a topic. Later hashes the returned key with [`partition_for_key`] to
/// pick the partition, the way a Kafka producer key is hashed by the
/// partitioner rather than naming a partition directly.
pub trait JobPartition {
    /// Returns the partition key that must route this message.
    ///
    /// Messages with the same key, in the same topic, always run in the
    /// same partition and therefore in strict relative order.
    fn partition_key(&self) -> PartitionKey;
}

/// The topic a message type is declared under.
///
/// [`crate::background_job!`] generates this implementation for any message
/// declared with `#[topic("name")]`. It is rarely implemented by hand.
pub trait JobTopic {
    /// Returns the stable topic name used to route and order this message.
    fn topic_name() -> &'static str;
}

/// A topic's fixed partition count, configured once for a namespace.
///
/// Changing a topic's partition count after jobs have been enqueued needs an
/// explicit migration, because old and new producers must agree on which
/// partition a given key hashes to.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct TopicConfig {
    pub(crate) name: String,
    pub(crate) partition_count: u32,
}

impl TopicConfig {
    /// Declares a topic with a fixed, positive partition count.
    pub fn new(name: impl Into<String>, partition_count: u32) -> anyhow::Result<Self> {
        let name = name.into();
        if name.trim().is_empty() {
            return Err(anyhow::anyhow!("topic name cannot be empty"));
        }
        if partition_count == 0 {
            return Err(anyhow::anyhow!(
                "topic '{name}' must have a positive partition count"
            ));
        }
        Ok(Self {
            name,
            partition_count,
        })
    }

    /// Returns the topic name.
    pub fn name(&self) -> &str {
        &self.name
    }

    /// Returns the configured partition count.
    pub fn partition_count(&self) -> u32 {
        self.partition_count
    }
}

/// Looks up `topic` and hashes `key` to the partition it must run in.
///
/// Returns an error when the topic is not registered, without computing a
/// partition.
pub(crate) fn resolve_partition(
    topics: &[TopicConfig],
    topic: &str,
    key: &PartitionKey,
) -> anyhow::Result<PartitionId> {
    let config = topics
        .iter()
        .find(|candidate| candidate.name == topic)
        .ok_or_else(|| anyhow::anyhow!("unknown topic '{topic}'"))?;
    Ok(partition_for_key(key, config.partition_count))
}

/// Picks the stable owner for `(topic, partition)` from `live_workers` using
/// rendezvous (highest random weight) hashing.
///
/// Every worker computes a score by hashing `(topic, partition, owner)` and
/// the owner with the highest score wins. Because the score depends only on
/// the candidate set, a small membership change (one worker added or
/// removed) only moves the partitions whose winner changes, rather than
/// reshuffling every assignment the way `hash(key) % worker_count` would.
/// Returns `None` when `live_workers` is empty.
pub(crate) fn rendezvous_owner<'a>(
    topic: &str,
    partition: PartitionId,
    live_workers: &'a [String],
) -> Option<&'a str> {
    live_workers
        .iter()
        .max_by_key(|owner| {
            let mut input = Vec::with_capacity(topic.len() + 1 + 4 + 1 + owner.len());
            input.extend_from_slice(topic.as_bytes());
            input.push(0);
            input.extend_from_slice(&partition.0.to_be_bytes());
            input.push(0);
            input.extend_from_slice(owner.as_bytes());
            *blake3::hash(&input).as_bytes()
        })
        .map(String::as_str)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn topic_config_rejects_empty_name_or_zero_partitions() {
        assert!(TopicConfig::new("", 1).is_err());
        assert!(TopicConfig::new("orders", 0).is_err());
        assert!(TopicConfig::new("orders", 4).is_ok());
    }

    #[test]
    fn resolve_partition_rejects_unknown_topic() {
        let topics = vec![TopicConfig::new("orders", 4).unwrap()];

        assert!(resolve_partition(&topics, "unknown", &PartitionKey::from("a")).is_err());
        assert!(resolve_partition(&topics, "orders", &PartitionKey::from("a")).is_ok());
    }

    #[test]
    fn partition_for_key_is_deterministic_and_in_range() {
        let key = PartitionKey::from("customer-42");

        let first = partition_for_key(&key, 8);
        let second = partition_for_key(&key, 8);

        assert_eq!(first, second);
        assert!(first.0 < 8);
    }

    #[test]
    fn partition_for_key_spreads_distinct_keys_across_partitions() {
        let partitions: std::collections::HashSet<u32> = (0..64)
            .map(|n| partition_for_key(&PartitionKey::from(n as u64), 8).0)
            .collect();

        assert!(
            partitions.len() > 1,
            "64 distinct keys all hashed to the same partition"
        );
    }

    #[test]
    fn partition_for_key_treats_zero_partition_count_as_one() {
        assert_eq!(
            partition_for_key(&PartitionKey::from("any"), 0),
            PartitionId(0)
        );
    }

    #[test]
    fn rendezvous_owner_is_stable_and_deterministic() {
        let workers = vec!["a".to_string(), "b".to_string(), "c".to_string()];

        let first = rendezvous_owner("orders", PartitionId(3), &workers);
        let second = rendezvous_owner("orders", PartitionId(3), &workers);
        assert_eq!(first, second);
        assert!(first.is_some());
    }

    #[test]
    fn rendezvous_owner_is_none_with_no_live_workers() {
        assert_eq!(rendezvous_owner("orders", PartitionId(0), &[]), None);
    }

    #[test]
    fn rendezvous_owner_moves_only_some_partitions_when_membership_changes() {
        let before: Vec<String> = vec!["a".into(), "b".into(), "c".into()];
        let after: Vec<String> = vec!["a".into(), "b".into(), "c".into(), "d".into()];

        let partitions = 64u32;
        let mut moved = 0;
        for partition in 0..partitions {
            let before_owner = rendezvous_owner("orders", PartitionId(partition), &before);
            let after_owner = rendezvous_owner("orders", PartitionId(partition), &after);
            if before_owner != after_owner {
                moved += 1;
            }
        }

        // Adding one worker to three should only reassign roughly a
        // quarter of partitions (to the new worker), not all of them.
        assert!(
            moved > 0 && moved < partitions,
            "expected some but not all partitions to move, moved {moved} of {partitions}"
        );
    }
}