rivven_core/

transaction.rs

//! Native Transaction Support (KIP-98)
//!
//! Provides exactly-once semantics with cross-topic atomic writes.
//!
//! ## Transaction Protocol
//!
//! ```text
//! Producer                     Transaction Coordinator            Partitions
//!    │                                   │                            │
//!    │─── InitProducerId ───────────────>│                            │
//!    │<── PID=123, Epoch=0 ──────────────│                            │
//!    │                                   │                            │
//!    │─── BeginTransaction(TxnId) ──────>│                            │
//!    │<── OK ────────────────────────────│                            │
//!    │                                   │                            │
//!    │─── AddPartitionsToTxn(p1,p2) ────>│                            │
//!    │<── OK ────────────────────────────│                            │
//!    │                                   │                            │
//!    │─── Produce(p1, PID, Seq) ──────────────────────────────────────>│
//!    │<── OK ───────────────────────────────────────────────────────────│
//!    │                                   │                            │
//!    │─── Produce(p2, PID, Seq) ──────────────────────────────────────>│
//!    │<── OK ───────────────────────────────────────────────────────────│
//!    │                                   │                            │
//!    │─── CommitTransaction(TxnId) ─────>│                            │
//!    │                                   │─── WriteTxnMarker(COMMIT) ─>│
//!    │                                   │<── OK ─────────────────────│
//!    │<── OK ────────────────────────────│                            │
//! ```
//!
//! ## Transaction States
//!
//! ```text
//! Empty ──────> Ongoing ──────> PrepareCommit ──────> CompleteCommit
//!                  │                  │                     │
//!                  │                  v                     v
//!                  └───────> PrepareAbort ───────> CompleteAbort
//!                                    │                     │
//!                                    └─────────────────────┘
//! ```
//!
//! ## Exactly-Once Guarantees
//!
//! 1. **Atomic Writes**: All messages in a transaction are committed or aborted together
//! 2. **Consumer Isolation**: Consumers only see committed messages (read_committed)
//! 3. **Fencing**: Old producer instances are fenced via epoch
//! 4. **Durability**: Transaction state is persisted before acknowledgment
//!

use crate::idempotent::{ProducerEpoch, ProducerId};
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::RwLock;
use std::time::{Duration, Instant, SystemTime};

/// Unique identifier for a transaction
pub type TransactionId = String;

/// Transaction timeout default (1 minute)
pub const DEFAULT_TRANSACTION_TIMEOUT: Duration = Duration::from_secs(60);

/// Maximum pending transactions per producer
pub const MAX_PENDING_TRANSACTIONS: usize = 5;

/// Transaction state machine
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum TransactionState {
    /// No active transaction
    Empty,

    /// Transaction in progress, accepting writes
    Ongoing,

    /// Preparing to commit (2PC phase 1)
    PrepareCommit,

    /// Preparing to abort (2PC phase 1)
    PrepareAbort,

    /// Commit complete (2PC phase 2)
    CompleteCommit,

    /// Abort complete (2PC phase 2)
    CompleteAbort,

    /// Transaction has expired without completion
    Dead,
}

impl TransactionState {
    /// Check if transaction is in a terminal state
    pub fn is_terminal(&self) -> bool {
        matches!(
            self,
            TransactionState::Empty
                | TransactionState::CompleteCommit
                | TransactionState::CompleteAbort
                | TransactionState::Dead
        )
    }

    /// Check if transaction is still active (can accept writes)
    pub fn is_active(&self) -> bool {
        matches!(self, TransactionState::Ongoing)
    }

    /// Check if transaction can transition to commit
    pub fn can_commit(&self) -> bool {
        matches!(self, TransactionState::Ongoing)
    }

    /// Check if transaction can transition to abort
    pub fn can_abort(&self) -> bool {
        matches!(
            self,
            TransactionState::Ongoing
                | TransactionState::PrepareCommit
                | TransactionState::PrepareAbort
        )
    }
}

/// Result of a transaction operation
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TransactionResult {
    /// Operation succeeded
    Ok,

    /// Transaction ID is invalid or not found
    InvalidTransactionId,

    /// Transaction is in wrong state for this operation
    InvalidTransactionState {
        current: TransactionState,
        expected: &'static str,
    },

    /// Producer ID/epoch mismatch
    ProducerFenced {
        expected_epoch: ProducerEpoch,
        received_epoch: ProducerEpoch,
    },

    /// Transaction has timed out
    TransactionTimeout,

    /// Too many pending transactions
    TooManyTransactions,

    /// Concurrent modification detected
    ConcurrentTransaction,

    /// Partition not part of transaction
    PartitionNotInTransaction { topic: String, partition: u32 },
}

/// A partition involved in a transaction
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct TransactionPartition {
    pub topic: String,
    pub partition: u32,
}

impl TransactionPartition {
    pub fn new(topic: impl Into<String>, partition: u32) -> Self {
        Self {
            topic: topic.into(),
            partition,
        }
    }
}

/// Pending write in a transaction (not yet committed)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PendingWrite {
    /// Target partition
    pub partition: TransactionPartition,

    /// Sequence number for this write
    pub sequence: i32,

    /// Offset assigned by the partition leader
    pub offset: u64,

    /// Write timestamp
    #[serde(with = "crate::serde_utils::system_time")]
    pub timestamp: SystemTime,
}

/// Consumer offset to be committed with the transaction
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TransactionOffsetCommit {
    /// Consumer group
    pub group_id: String,

    /// Topic-partition-offset triples
    pub offsets: Vec<(TransactionPartition, i64)>,
}

/// Active transaction state
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Transaction {
    /// Transaction ID (unique per producer)
    pub txn_id: TransactionId,

    /// Producer ID owning this transaction
    pub producer_id: ProducerId,

    /// Producer epoch (for fencing)
    pub producer_epoch: ProducerEpoch,

    /// Current state
    pub state: TransactionState,

    /// Partitions involved in this transaction
    pub partitions: HashSet<TransactionPartition>,

    /// Pending writes (not yet committed)
    pub pending_writes: Vec<PendingWrite>,

    /// Consumer offsets to commit with this transaction
    pub offset_commits: Vec<TransactionOffsetCommit>,

    /// Transaction start time
    #[serde(with = "crate::serde_utils::system_time")]
    pub started_at: SystemTime,

    /// Transaction timeout
    #[serde(with = "crate::serde_utils::duration")]
    pub timeout: Duration,

    /// Last activity timestamp
    #[serde(skip)]
    pub last_activity: Option<Instant>,
}

impl Transaction {
    /// Create a new transaction
    pub fn new(
        txn_id: TransactionId,
        producer_id: ProducerId,
        producer_epoch: ProducerEpoch,
        timeout: Duration,
    ) -> Self {
        Self {
            txn_id,
            producer_id,
            producer_epoch,
            state: TransactionState::Ongoing,
            partitions: HashSet::new(),
            pending_writes: Vec::new(),
            offset_commits: Vec::new(),
            started_at: SystemTime::now(),
            timeout,
            last_activity: Some(Instant::now()),
        }
    }

    /// Check if transaction has timed out
    pub fn is_timed_out(&self) -> bool {
        self.last_activity
            .map(|t| t.elapsed() > self.timeout)
            .unwrap_or(true)
    }

    /// Update last activity timestamp
    pub fn touch(&mut self) {
        self.last_activity = Some(Instant::now());
    }

    /// Add a partition to the transaction
    pub fn add_partition(&mut self, partition: TransactionPartition) {
        self.partitions.insert(partition);
        self.touch();
    }

    /// Record a pending write
    pub fn add_write(&mut self, partition: TransactionPartition, sequence: i32, offset: u64) {
        self.pending_writes.push(PendingWrite {
            partition,
            sequence,
            offset,
            timestamp: SystemTime::now(),
        });
        self.touch();
    }

    /// Add consumer offset commit
    pub fn add_offset_commit(
        &mut self,
        group_id: String,
        offsets: Vec<(TransactionPartition, i64)>,
    ) {
        self.offset_commits
            .push(TransactionOffsetCommit { group_id, offsets });
        self.touch();
    }

    /// Get total number of writes
    pub fn write_count(&self) -> usize {
        self.pending_writes.len()
    }

    /// Get all affected partitions
    pub fn affected_partitions(&self) -> impl Iterator<Item = &TransactionPartition> {
        self.partitions.iter()
    }
}

/// Transaction marker type written to partition logs
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum TransactionMarker {
    /// Transaction committed
    Commit,

    /// Transaction aborted
    Abort,
}

/// Statistics for transaction coordinator
#[derive(Debug, Default)]
pub struct TransactionStats {
    /// Total transactions initiated
    transactions_started: AtomicU64,

    /// Total transactions committed
    transactions_committed: AtomicU64,

    /// Total transactions aborted
    transactions_aborted: AtomicU64,

    /// Total transactions timed out
    transactions_timed_out: AtomicU64,

    /// Currently active transactions
    active_transactions: AtomicU64,
}

impl TransactionStats {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn record_start(&self) {
        self.transactions_started.fetch_add(1, Ordering::Relaxed);
        self.active_transactions.fetch_add(1, Ordering::Relaxed);
    }

    pub fn record_commit(&self) {
        self.transactions_committed.fetch_add(1, Ordering::Relaxed);
        self.active_transactions.fetch_sub(1, Ordering::Relaxed);
    }

    pub fn record_abort(&self) {
        self.transactions_aborted.fetch_add(1, Ordering::Relaxed);
        self.active_transactions.fetch_sub(1, Ordering::Relaxed);
    }

    pub fn record_timeout(&self) {
        self.transactions_timed_out.fetch_add(1, Ordering::Relaxed);
        self.active_transactions.fetch_sub(1, Ordering::Relaxed);
    }

    pub fn transactions_started(&self) -> u64 {
        self.transactions_started.load(Ordering::Relaxed)
    }

    pub fn transactions_committed(&self) -> u64 {
        self.transactions_committed.load(Ordering::Relaxed)
    }

    pub fn transactions_aborted(&self) -> u64 {
        self.transactions_aborted.load(Ordering::Relaxed)
    }

    pub fn transactions_timed_out(&self) -> u64 {
        self.transactions_timed_out.load(Ordering::Relaxed)
    }

    pub fn active_transactions(&self) -> u64 {
        self.active_transactions.load(Ordering::Relaxed)
    }
}

/// Snapshot of transaction stats for serialization
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TransactionStatsSnapshot {
    pub transactions_started: u64,
    pub transactions_committed: u64,
    pub transactions_aborted: u64,
    pub transactions_timed_out: u64,
    pub active_transactions: u64,
}

impl From<&TransactionStats> for TransactionStatsSnapshot {
    fn from(stats: &TransactionStats) -> Self {
        Self {
            transactions_started: stats.transactions_started(),
            transactions_committed: stats.transactions_committed(),
            transactions_aborted: stats.transactions_aborted(),
            transactions_timed_out: stats.transactions_timed_out(),
            active_transactions: stats.active_transactions(),
        }
    }
}

/// Transaction coordinator manages all active transactions
///
/// This is a per-broker component that tracks transactions for producers
/// assigned to this broker as their transaction coordinator.
pub struct TransactionCoordinator {
    /// Active transactions by (producer_id, txn_id)
    transactions: RwLock<HashMap<(ProducerId, TransactionId), Transaction>>,

    /// Producer to transaction mapping (for single-txn-per-producer enforcement)
    producer_transactions: RwLock<HashMap<ProducerId, TransactionId>>,

    /// Default transaction timeout
    default_timeout: Duration,

    /// Statistics
    stats: TransactionStats,
}

impl Default for TransactionCoordinator {
    fn default() -> Self {
        Self::new()
    }
}

impl TransactionCoordinator {
    /// Create a new transaction coordinator
    pub fn new() -> Self {
        Self {
            transactions: RwLock::new(HashMap::new()),
            producer_transactions: RwLock::new(HashMap::new()),
            default_timeout: DEFAULT_TRANSACTION_TIMEOUT,
            stats: TransactionStats::new(),
        }
    }

    /// Create with custom default timeout
    pub fn with_timeout(timeout: Duration) -> Self {
        Self {
            transactions: RwLock::new(HashMap::new()),
            producer_transactions: RwLock::new(HashMap::new()),
            default_timeout: timeout,
            stats: TransactionStats::new(),
        }
    }

    /// Get statistics
    pub fn stats(&self) -> &TransactionStats {
        &self.stats
    }

    /// Begin a new transaction
    pub fn begin_transaction(
        &self,
        txn_id: TransactionId,
        producer_id: ProducerId,
        producer_epoch: ProducerEpoch,
        timeout: Option<Duration>,
    ) -> TransactionResult {
        // Check if producer already has an active transaction
        {
            let producer_txns = self
                .producer_transactions
                .read()
                .expect("transaction manager lock poisoned");
            if let Some(existing_txn_id) = producer_txns.get(&producer_id) {
                if existing_txn_id != &txn_id {
                    return TransactionResult::ConcurrentTransaction;
                }
                // Same txn_id - check if we're resuming
                let transactions = self
                    .transactions
                    .read()
                    .expect("transaction manager lock poisoned");
                if let Some(txn) = transactions.get(&(producer_id, txn_id.clone())) {
                    if txn.producer_epoch != producer_epoch {
                        return TransactionResult::ProducerFenced {
                            expected_epoch: txn.producer_epoch,
                            received_epoch: producer_epoch,
                        };
                    }
                    if txn.state.is_active() {
                        return TransactionResult::Ok; // Already active
                    }
                }
            }
        }

        // Create new transaction
        let txn = Transaction::new(
            txn_id.clone(),
            producer_id,
            producer_epoch,
            timeout.unwrap_or(self.default_timeout),
        );

        {
            let mut transactions = self
                .transactions
                .write()
                .expect("transaction manager lock poisoned");
            let mut producer_txns = self
                .producer_transactions
                .write()
                .expect("transaction manager lock poisoned");

            transactions.insert((producer_id, txn_id.clone()), txn);
            producer_txns.insert(producer_id, txn_id);
        }

        self.stats.record_start();
        TransactionResult::Ok
    }

    /// Add partitions to an active transaction
    pub fn add_partitions_to_transaction(
        &self,
        txn_id: &TransactionId,
        producer_id: ProducerId,
        producer_epoch: ProducerEpoch,
        partitions: Vec<TransactionPartition>,
    ) -> TransactionResult {
        let mut transactions = self
            .transactions
            .write()
            .expect("transaction manager lock poisoned");

        let txn = match transactions.get_mut(&(producer_id, txn_id.clone())) {
            Some(t) => t,
            None => return TransactionResult::InvalidTransactionId,
        };

        // Validate epoch
        if txn.producer_epoch != producer_epoch {
            return TransactionResult::ProducerFenced {
                expected_epoch: txn.producer_epoch,
                received_epoch: producer_epoch,
            };
        }

        // Check state
        if !txn.state.is_active() {
            return TransactionResult::InvalidTransactionState {
                current: txn.state,
                expected: "Ongoing",
            };
        }

        // Check timeout
        if txn.is_timed_out() {
            txn.state = TransactionState::Dead;
            self.stats.record_timeout();
            return TransactionResult::TransactionTimeout;
        }

        // Add partitions
        for partition in partitions {
            txn.add_partition(partition);
        }

        TransactionResult::Ok
    }

    /// Record a write within a transaction
    pub fn add_write_to_transaction(
        &self,
        txn_id: &TransactionId,
        producer_id: ProducerId,
        producer_epoch: ProducerEpoch,
        partition: TransactionPartition,
        sequence: i32,
        offset: u64,
    ) -> TransactionResult {
        let mut transactions = self
            .transactions
            .write()
            .expect("transaction manager lock poisoned");

        let txn = match transactions.get_mut(&(producer_id, txn_id.clone())) {
            Some(t) => t,
            None => return TransactionResult::InvalidTransactionId,
        };

        // Validate epoch
        if txn.producer_epoch != producer_epoch {
            return TransactionResult::ProducerFenced {
                expected_epoch: txn.producer_epoch,
                received_epoch: producer_epoch,
            };
        }

        // Check state
        if !txn.state.is_active() {
            return TransactionResult::InvalidTransactionState {
                current: txn.state,
                expected: "Ongoing",
            };
        }

        // Check timeout
        if txn.is_timed_out() {
            txn.state = TransactionState::Dead;
            self.stats.record_timeout();
            return TransactionResult::TransactionTimeout;
        }

        // Verify partition is part of transaction
        if !txn.partitions.contains(&partition) {
            return TransactionResult::PartitionNotInTransaction {
                topic: partition.topic,
                partition: partition.partition,
            };
        }

        // Record the write
        txn.add_write(partition, sequence, offset);

        TransactionResult::Ok
    }

    /// Add consumer offset commit to transaction (for exactly-once consume-transform-produce)
    pub fn add_offsets_to_transaction(
        &self,
        txn_id: &TransactionId,
        producer_id: ProducerId,
        producer_epoch: ProducerEpoch,
        group_id: String,
        offsets: Vec<(TransactionPartition, i64)>,
    ) -> TransactionResult {
        let mut transactions = self
            .transactions
            .write()
            .expect("transaction manager lock poisoned");

        let txn = match transactions.get_mut(&(producer_id, txn_id.clone())) {
            Some(t) => t,
            None => return TransactionResult::InvalidTransactionId,
        };

        // Validate epoch
        if txn.producer_epoch != producer_epoch {
            return TransactionResult::ProducerFenced {
                expected_epoch: txn.producer_epoch,
                received_epoch: producer_epoch,
            };
        }

        // Check state
        if !txn.state.is_active() {
            return TransactionResult::InvalidTransactionState {
                current: txn.state,
                expected: "Ongoing",
            };
        }

        // Check timeout
        if txn.is_timed_out() {
            txn.state = TransactionState::Dead;
            self.stats.record_timeout();
            return TransactionResult::TransactionTimeout;
        }

        // Add offset commit
        txn.add_offset_commit(group_id, offsets);

        TransactionResult::Ok
    }

    /// Prepare to commit a transaction (2PC phase 1)
    ///
    /// Returns the transaction data needed for committing to partitions
    pub fn prepare_commit(
        &self,
        txn_id: &TransactionId,
        producer_id: ProducerId,
        producer_epoch: ProducerEpoch,
    ) -> Result<Transaction, TransactionResult> {
        let mut transactions = self
            .transactions
            .write()
            .expect("transaction manager lock poisoned");

        let txn = match transactions.get_mut(&(producer_id, txn_id.clone())) {
            Some(t) => t,
            None => return Err(TransactionResult::InvalidTransactionId),
        };

        // Validate epoch
        if txn.producer_epoch != producer_epoch {
            return Err(TransactionResult::ProducerFenced {
                expected_epoch: txn.producer_epoch,
                received_epoch: producer_epoch,
            });
        }

        // Check state
        if !txn.state.can_commit() {
            return Err(TransactionResult::InvalidTransactionState {
                current: txn.state,
                expected: "Ongoing",
            });
        }

        // Check timeout
        if txn.is_timed_out() {
            txn.state = TransactionState::Dead;
            self.stats.record_timeout();
            return Err(TransactionResult::TransactionTimeout);
        }

        // Transition to PrepareCommit
        txn.state = TransactionState::PrepareCommit;
        txn.touch();

        Ok(txn.clone())
    }

    /// Complete the commit (2PC phase 2)
    pub fn complete_commit(
        &self,
        txn_id: &TransactionId,
        producer_id: ProducerId,
    ) -> TransactionResult {
        let mut transactions = self
            .transactions
            .write()
            .expect("transaction manager lock poisoned");
        let mut producer_txns = self
            .producer_transactions
            .write()
            .expect("transaction manager lock poisoned");

        let txn = match transactions.get_mut(&(producer_id, txn_id.clone())) {
            Some(t) => t,
            None => return TransactionResult::InvalidTransactionId,
        };

        if txn.state != TransactionState::PrepareCommit {
            return TransactionResult::InvalidTransactionState {
                current: txn.state,
                expected: "PrepareCommit",
            };
        }

        txn.state = TransactionState::CompleteCommit;

        // Clean up
        transactions.remove(&(producer_id, txn_id.clone()));
        producer_txns.remove(&producer_id);

        self.stats.record_commit();
        TransactionResult::Ok
    }

    /// Prepare to abort a transaction (2PC phase 1)
    pub fn prepare_abort(
        &self,
        txn_id: &TransactionId,
        producer_id: ProducerId,
        producer_epoch: ProducerEpoch,
    ) -> Result<Transaction, TransactionResult> {
        let mut transactions = self
            .transactions
            .write()
            .expect("transaction manager lock poisoned");

        let txn = match transactions.get_mut(&(producer_id, txn_id.clone())) {
            Some(t) => t,
            None => return Err(TransactionResult::InvalidTransactionId),
        };

        // Validate epoch
        if txn.producer_epoch != producer_epoch {
            return Err(TransactionResult::ProducerFenced {
                expected_epoch: txn.producer_epoch,
                received_epoch: producer_epoch,
            });
        }

        // Check state - abort is allowed from more states than commit
        if !txn.state.can_abort() {
            return Err(TransactionResult::InvalidTransactionState {
                current: txn.state,
                expected: "Ongoing or PrepareCommit",
            });
        }

        // Transition to PrepareAbort
        txn.state = TransactionState::PrepareAbort;
        txn.touch();

        Ok(txn.clone())
    }

    /// Complete the abort (2PC phase 2)
    pub fn complete_abort(
        &self,
        txn_id: &TransactionId,
        producer_id: ProducerId,
    ) -> TransactionResult {
        let mut transactions = self
            .transactions
            .write()
            .expect("transaction manager lock poisoned");
        let mut producer_txns = self
            .producer_transactions
            .write()
            .expect("transaction manager lock poisoned");

        let txn = match transactions.get_mut(&(producer_id, txn_id.clone())) {
            Some(t) => t,
            None => return TransactionResult::InvalidTransactionId,
        };

        if txn.state != TransactionState::PrepareAbort {
            return TransactionResult::InvalidTransactionState {
                current: txn.state,
                expected: "PrepareAbort",
            };
        }

        txn.state = TransactionState::CompleteAbort;

        // Clean up
        transactions.remove(&(producer_id, txn_id.clone()));
        producer_txns.remove(&producer_id);

        self.stats.record_abort();
        TransactionResult::Ok
    }

    /// Get current transaction state for a producer
    pub fn get_transaction(
        &self,
        txn_id: &TransactionId,
        producer_id: ProducerId,
    ) -> Option<Transaction> {
        let transactions = self
            .transactions
            .read()
            .expect("transaction manager lock poisoned");
        transactions.get(&(producer_id, txn_id.clone())).cloned()
    }

    /// Check if a producer has an active transaction
    pub fn has_active_transaction(&self, producer_id: ProducerId) -> bool {
        let producer_txns = self
            .producer_transactions
            .read()
            .expect("transaction manager lock poisoned");
        producer_txns.contains_key(&producer_id)
    }

    /// Get active transaction ID for a producer
    pub fn get_active_transaction_id(&self, producer_id: ProducerId) -> Option<TransactionId> {
        let producer_txns = self
            .producer_transactions
            .read()
            .expect("transaction manager lock poisoned");
        producer_txns.get(&producer_id).cloned()
    }

    /// Clean up timed-out transactions
    pub fn cleanup_timed_out_transactions(&self) -> Vec<Transaction> {
        let mut timed_out = Vec::new();
        let mut transactions = self
            .transactions
            .write()
            .expect("transaction manager lock poisoned");
        let mut producer_txns = self
            .producer_transactions
            .write()
            .expect("transaction manager lock poisoned");

        let keys_to_remove: Vec<_> = transactions
            .iter()
            .filter(|(_, txn)| txn.is_timed_out() && !txn.state.is_terminal())
            .map(|(k, _)| k.clone())
            .collect();

        for key in keys_to_remove {
            if let Some(mut txn) = transactions.remove(&key) {
                txn.state = TransactionState::Dead;
                producer_txns.remove(&txn.producer_id);
                self.stats.record_timeout();
                timed_out.push(txn);
            }
        }

        timed_out
    }

    /// Get number of active transactions
    pub fn active_count(&self) -> usize {
        let transactions = self
            .transactions
            .read()
            .expect("transaction manager lock poisoned");
        transactions
            .values()
            .filter(|t| !t.state.is_terminal())
            .count()
    }
}

// ============================================================================
// Tests
// ============================================================================

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

    #[test]
    fn test_transaction_state_transitions() {
        // Test terminal states
        assert!(TransactionState::Empty.is_terminal());
        assert!(TransactionState::CompleteCommit.is_terminal());
        assert!(TransactionState::CompleteAbort.is_terminal());
        assert!(TransactionState::Dead.is_terminal());

        // Test active states
        assert!(!TransactionState::Ongoing.is_terminal());
        assert!(!TransactionState::PrepareCommit.is_terminal());
        assert!(!TransactionState::PrepareAbort.is_terminal());

        // Test can_commit
        assert!(TransactionState::Ongoing.can_commit());
        assert!(!TransactionState::Empty.can_commit());
        assert!(!TransactionState::PrepareCommit.can_commit());

        // Test can_abort
        assert!(TransactionState::Ongoing.can_abort());
        assert!(TransactionState::PrepareCommit.can_abort());
        assert!(TransactionState::PrepareAbort.can_abort());
        assert!(!TransactionState::Empty.can_abort());
    }

    #[test]
    fn test_begin_transaction() {
        let coordinator = TransactionCoordinator::new();

        // Begin first transaction
        let result = coordinator.begin_transaction("txn-1".to_string(), 1, 0, None);
        assert_eq!(result, TransactionResult::Ok);

        // Verify transaction exists
        let txn = coordinator.get_transaction(&"txn-1".to_string(), 1);
        assert!(txn.is_some());
        let txn = txn.unwrap();
        assert_eq!(txn.state, TransactionState::Ongoing);
        assert_eq!(txn.producer_id, 1);
        assert_eq!(txn.producer_epoch, 0);

        // Stats
        assert_eq!(coordinator.stats().transactions_started(), 1);
        assert_eq!(coordinator.stats().active_transactions(), 1);
    }

    #[test]
    fn test_concurrent_transaction_rejection() {
        let coordinator = TransactionCoordinator::new();

        // Begin first transaction
        coordinator.begin_transaction("txn-1".to_string(), 1, 0, None);

        // Try to begin another transaction for same producer
        let result = coordinator.begin_transaction("txn-2".to_string(), 1, 0, None);
        assert_eq!(result, TransactionResult::ConcurrentTransaction);
    }

    #[test]
    fn test_add_partitions_to_transaction() {
        let coordinator = TransactionCoordinator::new();
        coordinator.begin_transaction("txn-1".to_string(), 1, 0, None);

        // Add partitions
        let result = coordinator.add_partitions_to_transaction(
            &"txn-1".to_string(),
            1,
            0,
            vec![
                TransactionPartition::new("topic-1", 0),
                TransactionPartition::new("topic-1", 1),
                TransactionPartition::new("topic-2", 0),
            ],
        );
        assert_eq!(result, TransactionResult::Ok);

        // Verify partitions added
        let txn = coordinator
            .get_transaction(&"txn-1".to_string(), 1)
            .unwrap();
        assert_eq!(txn.partitions.len(), 3);
    }

    #[test]
    fn test_add_write_to_transaction() {
        let coordinator = TransactionCoordinator::new();
        coordinator.begin_transaction("txn-1".to_string(), 1, 0, None);

        let partition = TransactionPartition::new("topic-1", 0);
        coordinator.add_partitions_to_transaction(
            &"txn-1".to_string(),
            1,
            0,
            vec![partition.clone()],
        );

        // Record write
        let result =
            coordinator.add_write_to_transaction(&"txn-1".to_string(), 1, 0, partition, 0, 100);
        assert_eq!(result, TransactionResult::Ok);

        // Verify write recorded
        let txn = coordinator
            .get_transaction(&"txn-1".to_string(), 1)
            .unwrap();
        assert_eq!(txn.pending_writes.len(), 1);
        assert_eq!(txn.pending_writes[0].offset, 100);
        assert_eq!(txn.pending_writes[0].sequence, 0);
    }

    #[test]
    fn test_write_to_non_registered_partition() {
        let coordinator = TransactionCoordinator::new();
        coordinator.begin_transaction("txn-1".to_string(), 1, 0, None);

        // Try to write to partition not added to transaction
        let result = coordinator.add_write_to_transaction(
            &"txn-1".to_string(),
            1,
            0,
            TransactionPartition::new("topic-1", 0),
            0,
            100,
        );

        assert!(matches!(
            result,
            TransactionResult::PartitionNotInTransaction { .. }
        ));
    }

    #[test]
    fn test_commit_transaction() {
        let coordinator = TransactionCoordinator::new();
        coordinator.begin_transaction("txn-1".to_string(), 1, 0, None);

        let partition = TransactionPartition::new("topic-1", 0);
        coordinator.add_partitions_to_transaction(
            &"txn-1".to_string(),
            1,
            0,
            vec![partition.clone()],
        );
        coordinator.add_write_to_transaction(&"txn-1".to_string(), 1, 0, partition, 0, 100);

        // Prepare commit
        let txn = coordinator.prepare_commit(&"txn-1".to_string(), 1, 0);
        assert!(txn.is_ok());
        let txn = txn.unwrap();
        assert_eq!(txn.state, TransactionState::PrepareCommit);

        // Complete commit
        let result = coordinator.complete_commit(&"txn-1".to_string(), 1);
        assert_eq!(result, TransactionResult::Ok);

        // Transaction should be removed
        assert!(coordinator
            .get_transaction(&"txn-1".to_string(), 1)
            .is_none());
        assert!(!coordinator.has_active_transaction(1));

        // Stats
        assert_eq!(coordinator.stats().transactions_committed(), 1);
        assert_eq!(coordinator.stats().active_transactions(), 0);
    }

    #[test]
    fn test_abort_transaction() {
        let coordinator = TransactionCoordinator::new();
        coordinator.begin_transaction("txn-1".to_string(), 1, 0, None);

        let partition = TransactionPartition::new("topic-1", 0);
        coordinator.add_partitions_to_transaction(
            &"txn-1".to_string(),
            1,
            0,
            vec![partition.clone()],
        );
        coordinator.add_write_to_transaction(&"txn-1".to_string(), 1, 0, partition, 0, 100);

        // Prepare abort
        let txn = coordinator.prepare_abort(&"txn-1".to_string(), 1, 0);
        assert!(txn.is_ok());

        // Complete abort
        let result = coordinator.complete_abort(&"txn-1".to_string(), 1);
        assert_eq!(result, TransactionResult::Ok);

        // Transaction should be removed
        assert!(coordinator
            .get_transaction(&"txn-1".to_string(), 1)
            .is_none());

        // Stats
        assert_eq!(coordinator.stats().transactions_aborted(), 1);
    }

    #[test]
    fn test_producer_fencing() {
        let coordinator = TransactionCoordinator::new();
        coordinator.begin_transaction("txn-1".to_string(), 1, 0, None);

        // Try with wrong epoch
        let result = coordinator.add_partitions_to_transaction(
            &"txn-1".to_string(),
            1,
            1, // Wrong epoch
            vec![TransactionPartition::new("topic-1", 0)],
        );

        assert!(matches!(
            result,
            TransactionResult::ProducerFenced {
                expected_epoch: 0,
                received_epoch: 1
            }
        ));
    }

    #[test]
    fn test_transaction_timeout() {
        // Create coordinator with very short timeout
        let coordinator = TransactionCoordinator::with_timeout(Duration::from_millis(1));
        coordinator.begin_transaction("txn-1".to_string(), 1, 0, None);

        // Wait for timeout
        std::thread::sleep(Duration::from_millis(5));

        // Try to add partitions - should fail with timeout
        let result = coordinator.add_partitions_to_transaction(
            &"txn-1".to_string(),
            1,
            0,
            vec![TransactionPartition::new("topic-1", 0)],
        );

        assert_eq!(result, TransactionResult::TransactionTimeout);
    }

    #[test]
    fn test_cleanup_timed_out_transactions() {
        let coordinator = TransactionCoordinator::with_timeout(Duration::from_millis(1));

        coordinator.begin_transaction("txn-1".to_string(), 1, 0, None);
        coordinator.begin_transaction("txn-2".to_string(), 2, 0, None);

        // Wait for timeout
        std::thread::sleep(Duration::from_millis(5));

        // Cleanup
        let timed_out = coordinator.cleanup_timed_out_transactions();
        assert_eq!(timed_out.len(), 2);

        // Transactions should be gone
        assert_eq!(coordinator.active_count(), 0);
        assert_eq!(coordinator.stats().transactions_timed_out(), 2);
    }

    #[test]
    fn test_add_offsets_to_transaction() {
        let coordinator = TransactionCoordinator::new();
        coordinator.begin_transaction("txn-1".to_string(), 1, 0, None);

        // Add consumer offsets
        let result = coordinator.add_offsets_to_transaction(
            &"txn-1".to_string(),
            1,
            0,
            "consumer-group-1".to_string(),
            vec![
                (TransactionPartition::new("input-topic", 0), 42),
                (TransactionPartition::new("input-topic", 1), 100),
            ],
        );
        assert_eq!(result, TransactionResult::Ok);

        // Verify
        let txn = coordinator
            .get_transaction(&"txn-1".to_string(), 1)
            .unwrap();
        assert_eq!(txn.offset_commits.len(), 1);
        assert_eq!(txn.offset_commits[0].group_id, "consumer-group-1");
        assert_eq!(txn.offset_commits[0].offsets.len(), 2);
    }

    #[test]
    fn test_invalid_state_transitions() {
        let coordinator = TransactionCoordinator::new();
        coordinator.begin_transaction("txn-1".to_string(), 1, 0, None);

        // Prepare commit
        coordinator
            .prepare_commit(&"txn-1".to_string(), 1, 0)
            .unwrap();

        // Try to add partitions after prepare - should fail
        let result = coordinator.add_partitions_to_transaction(
            &"txn-1".to_string(),
            1,
            0,
            vec![TransactionPartition::new("topic-1", 0)],
        );
        assert!(matches!(
            result,
            TransactionResult::InvalidTransactionState { .. }
        ));
    }

    #[test]
    fn test_abort_from_prepare_commit() {
        let coordinator = TransactionCoordinator::new();
        coordinator.begin_transaction("txn-1".to_string(), 1, 0, None);

        // Prepare commit
        coordinator
            .prepare_commit(&"txn-1".to_string(), 1, 0)
            .unwrap();

        // Abort should still be allowed from PrepareCommit
        let result = coordinator.prepare_abort(&"txn-1".to_string(), 1, 0);
        assert!(result.is_ok());

        let result = coordinator.complete_abort(&"txn-1".to_string(), 1);
        assert_eq!(result, TransactionResult::Ok);
    }

    #[test]
    fn test_transaction_partition_hash() {
        let p1 = TransactionPartition::new("topic", 0);
        let p2 = TransactionPartition::new("topic", 0);
        let p3 = TransactionPartition::new("topic", 1);

        assert_eq!(p1, p2);
        assert_ne!(p1, p3);

        let mut set = HashSet::new();
        set.insert(p1.clone());
        set.insert(p2); // Should not add (duplicate)
        set.insert(p3);
        assert_eq!(set.len(), 2);
    }

    #[test]
    fn test_resume_same_transaction() {
        let coordinator = TransactionCoordinator::new();

        // Begin transaction
        coordinator.begin_transaction("txn-1".to_string(), 1, 0, None);

        // Try to begin same transaction again - should succeed (idempotent)
        let result = coordinator.begin_transaction("txn-1".to_string(), 1, 0, None);
        assert_eq!(result, TransactionResult::Ok);

        // Only one transaction should exist
        assert_eq!(coordinator.active_count(), 1);
        assert_eq!(coordinator.stats().transactions_started(), 1);
    }

    #[test]
    fn test_stats_snapshot() {
        let coordinator = TransactionCoordinator::new();
        coordinator.begin_transaction("txn-1".to_string(), 1, 0, None);
        coordinator
            .prepare_commit(&"txn-1".to_string(), 1, 0)
            .unwrap();
        coordinator.complete_commit(&"txn-1".to_string(), 1);

        let snapshot: TransactionStatsSnapshot = coordinator.stats().into();
        assert_eq!(snapshot.transactions_started, 1);
        assert_eq!(snapshot.transactions_committed, 1);
        assert_eq!(snapshot.active_transactions, 0);
    }
}