mocra-cluster 0.4.1

//! Partition ownership (refactor Phase 3): spreads collection work across cluster nodes by
//! partition.
//!
//! Design:
//! - **Partition key**: `hash(account) % N` maps an account onto a fixed number of virtual
//!   partitions (256 by default).
//! - **Ownership**: rendezvous (HRW) hashing — each partition belongs to the node whose
//!   `(partition, node)` combination hashes highest. When members are added or removed, **only
//!   about 1/N of the partitions migrate**, which minimizes rebalancing by construction.
//! - **Consistency**: the member list comes from Raft (strongly consistent) and the hash is
//!   **deterministic** (no randomly seeded `DefaultHasher`), so every node computes exactly the
//!   same ownership and reaches consensus without any extra negotiation.
//! - **Split-brain protection**: node views can diverge briefly at the moment membership changes;
//!   when a strong guarantee is needed, take an ownership lease carrying a monotonic fencing token
//!   with [`RaftControlPlane::acquire_partition`](crate::RaftControlPlane::acquire_partition)
//!   (which reuses the Raft lock) — a stale owner is rejected on its smaller token.
//!
//! For MQs with consumer groups (Kafka/NATS), partition assignment can reuse the consumer group
//! directly; this module is the fallback ownership scheme for backends **without** consumer groups
//! (in-memory).

use crate::raft::NodeId;

/// The default number of virtual partitions. A power of two, which distributes evenly; fine
/// enough to balance load across small and medium clusters.
pub const DEFAULT_PARTITIONS: u32 = 256;

const FNV_OFFSET: u64 = 14695981039346656037;
const FNV_PRIME: u64 = 1099511628211;

#[inline]
fn fnv1a(bytes: &[u8]) -> u64 {
    let mut h = FNV_OFFSET;
    for &b in bytes {
        h ^= b as u64;
        h = h.wrapping_mul(FNV_PRIME);
    }
    h
}

/// Account / session key → partition number. Stable and process-independent.
pub fn partition_of(key: &str, partitions: u32) -> u32 {
    let n = partitions.max(1);
    (fnv1a(key.as_bytes()) % n as u64) as u32
}

/// Rendezvous (HRW) score: a stable hash of the `(partition, node)` combination.
#[inline]
fn hrw_score(partition: u32, node: NodeId) -> u64 {
    let mut buf = [0u8; 12];
    buf[..8].copy_from_slice(&node.to_le_bytes());
    buf[8..].copy_from_slice(&partition.to_le_bytes());
    fnv1a(&buf)
}

/// Computes which node owns a partition for a given member set (the highest score wins).
///
/// Returns `None` when the member set is empty. Ties are broken in favour of the larger `node_id`
/// (a deterministic tie-break).
pub fn owner_of_partition(partition: u32, members: &[NodeId]) -> Option<NodeId> {
    members
        .iter()
        .copied()
        .map(|node| (hrw_score(partition, node), node))
        .max()
        .map(|(_, node)| node)
}

/// The set of partitions a member owns for a given member set.
pub fn partitions_owned_by(me: NodeId, members: &[NodeId], partitions: u32) -> Vec<u32> {
    let n = partitions.max(1);
    (0..n)
        .filter(|&p| owner_of_partition(p, members) == Some(me))
        .collect()
}

/// Whether a given account / session key is handled by `me`.
pub fn owns_key(me: NodeId, key: &str, members: &[NodeId], partitions: u32) -> bool {
    owner_of_partition(partition_of(key, partitions), members) == Some(me)
}

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

    #[test]
    fn assignment_is_deterministic_and_process_independent() {
        // The same input must always yield the same ownership (no reliance on a random seed).
        let members = vec![1u64, 2, 3, 4, 5];
        for p in 0..DEFAULT_PARTITIONS {
            let a = owner_of_partition(p, &members);
            let b = owner_of_partition(p, &members);
            assert_eq!(a, b);
            assert!(a.is_some());
        }
        // Known vector: the partition hash is stable (a regression guard; the value only needs to
        // stay pinned to the algorithm).
        assert_eq!(
            partition_of("account-42", DEFAULT_PARTITIONS),
            partition_of("account-42", DEFAULT_PARTITIONS)
        );
    }

    #[test]
    fn every_partition_owned_by_exactly_one_member() {
        let members = vec![10u64, 20, 30];
        let mut union: BTreeSet<u32> = BTreeSet::new();
        for &m in &members {
            for p in partitions_owned_by(m, &members, DEFAULT_PARTITIONS) {
                // No overlap: the same partition must not be claimed by two members.
                assert!(
                    union.insert(p),
                    "partition {p} owned by more than one member"
                );
            }
        }
        // Full coverage: every partition has an owner.
        assert_eq!(union.len(), DEFAULT_PARTITIONS as usize);
    }

    #[test]
    fn load_is_roughly_balanced() {
        let members = vec![1u64, 2, 3, 4];
        let per = DEFAULT_PARTITIONS as usize / members.len();
        for &m in &members {
            let owned = partitions_owned_by(m, &members, DEFAULT_PARTITIONS).len();
            // Allow ±50% jitter (256 partitions / 4 nodes ≈ 64, i.e. the range [32,96]).
            assert!(
                owned >= per / 2 && owned <= per * 3 / 2,
                "member {m} owns {owned}, expected ~{per}"
            );
        }
    }

    #[test]
    fn membership_change_moves_minimal_partitions() {
        // The key HRW property: adding a node migrates only the partitions that move to the new
        // node; everything else stays put.
        let before = vec![1u64, 2, 3];
        let after = vec![1u64, 2, 3, 4];
        let mut moved = 0;
        for p in 0..DEFAULT_PARTITIONS {
            let o1 = owner_of_partition(p, &before).unwrap();
            let o2 = owner_of_partition(p, &after).unwrap();
            if o1 != o2 {
                moved += 1;
                // The only possible migration target is the newly joined node 4.
                assert_eq!(o2, 4, "partition {p} moved to a non-new node");
            }
        }
        // The number moved should be close to 1/4 (the new node's fair share), far short of the
        // whole set.
        let expected = DEFAULT_PARTITIONS as usize / after.len();
        assert!(
            moved > 0 && moved < DEFAULT_PARTITIONS as usize / 2,
            "moved {moved}, expected ~{expected} (minimal reshuffle)"
        );
    }

    #[test]
    fn single_member_owns_everything() {
        let members = vec![7u64];
        assert_eq!(
            partitions_owned_by(7, &members, DEFAULT_PARTITIONS).len(),
            DEFAULT_PARTITIONS as usize
        );
        assert!(owns_key(7, "any-account", &members, DEFAULT_PARTITIONS));
    }

    #[test]
    fn empty_members_owns_nothing() {
        assert_eq!(owner_of_partition(0, &[]), None);
        assert!(!owns_key(1, "acc", &[], DEFAULT_PARTITIONS));
    }
}