krabka_log/log/
verification.rs

1//! KIP-890 part 1: the per-producer verification state that ties a
2//! transactional append to a transaction the coordinator knows.
3//!
4//! Before a transactional batch starts a transaction on this partition, the
5//! broker asks the transaction coordinator to verify (or, with transaction
6//! version 2, to add) the partition. The log hands out a
7//! [`VerificationGuard`] when that check starts and keeps it until the next
8//! transactional append or end marker of the producer. An append must present
9//! the same guard. A marker that lands between the check and the append clears
10//! the guard, so a stale check can never start a transaction the coordinator
11//! already ended.
12//!
13//! This is Kafka's `VerificationStateEntry`, `VerificationGuard`,
14//! `UnifiedLog.maybeStartTransactionVerification` and the
15//! `batchMissingRequiredVerification` check in
16//! `UnifiedLog.analyzeAndValidateProducerState`.
17
18use {
19    super::Log,
20    krabka_ids::ProducerId,
21    std::sync::atomic::{AtomicU64, Ordering},
22};
23
24/// A token for one verification of one producer on one partition.
25///
26/// [`VerificationGuard::SENTINEL`] verifies nothing. Kafka's
27/// `VerificationGuard.SENTINEL` has the value 0 and every real guard takes the
28/// next value of a process-wide counter.
29#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
30pub struct VerificationGuard(u64);
31
32impl VerificationGuard {
33    /// The guard that verifies nothing.
34    pub const SENTINEL: Self = Self(0);
35
36    fn next() -> Self {
37        static NEXT: AtomicU64 = AtomicU64::new(1);
38        Self(NEXT.fetch_add(1, Ordering::Relaxed))
39    }
40
41    /// Whether `presented` is this guard and is not the sentinel. Kafka's
42    /// `VerificationGuard.verify`.
43    #[must_use]
44    pub fn verifies(self, presented: Self) -> bool {
45        presented != Self::SENTINEL && presented == self
46    }
47}
48
49/// The verification state of one producer. Kafka's `VerificationStateEntry`.
50#[derive(Debug, Clone, Copy, PartialEq, Eq)]
51pub(super) struct VerificationState {
52    pub(super) guard: VerificationGuard,
53    pub(super) created_ms: i64,
54    pub(super) epoch: i16,
55    pub(super) lowest_sequence: i32,
56    pub(super) supports_epoch_bump: bool,
57}
58
59/// Why a transactional produce may not start or continue its append.
60#[derive(Debug, Clone, Copy, PartialEq, Eq)]
61pub enum TransactionAppendRefusal {
62    /// The batch epoch is below the producer's epoch on this partition. Kafka
63    /// throws `InvalidProducerEpochException`.
64    StaleProducerEpoch,
65    /// A transactional batch has no open transaction and no matching guard, or
66    /// a non-transactional batch comes from a producer with an open
67    /// transaction. Kafka throws `InvalidTxnStateException`.
68    InvalidTransactionState,
69    /// The first sequence of a verified batch does not fit the verification
70    /// state. Kafka throws `OutOfOrderSequenceException` from
71    /// `ProducerAppendInfo.checkSequence`.
72    OutOfOrderSequence,
73}
74
75/// One batch as the transaction checks see it.
76#[derive(Debug, Clone, Copy, PartialEq, Eq)]
77pub struct TransactionalBatch {
78    pub producer_id: ProducerId,
79    pub producer_epoch: i16,
80    pub base_sequence: i32,
81    pub is_transactional: bool,
82    pub is_control: bool,
83}
84
85impl Log {
86    /// Whether `producer_id` has an open transaction on this partition at
87    /// `producer_epoch`. Kafka's `UnifiedLog.hasOngoingTransaction`.
88    #[must_use]
89    pub fn has_ongoing_transaction(&self, producer_id: ProducerId, producer_epoch: i16) -> bool {
90        self.producer_state.get(&producer_id).is_some_and(|entry| {
91            entry.current_txn_first_offset.is_some() && entry.producer_epoch == producer_epoch
92        })
93    }
94
95    /// Start a verification for a transactional batch, and return its guard.
96    ///
97    /// A producer with an open transaction at the batch epoch needs no
98    /// verification, and gets [`VerificationGuard::SENTINEL`]. Otherwise the
99    /// producer's guard is created, or kept with its lowest sequence moved as
100    /// Kafka's `maybeUpdateLowestSequenceAndEpoch` moves it. The call also
101    /// drops verification state older than `expiration_ms`, as Kafka's
102    /// `ProducerStateManager.removeExpiredProducers` does.
103    ///
104    /// # Errors
105    ///
106    /// Returns [`TransactionAppendRefusal::StaleProducerEpoch`] when the batch
107    /// epoch is below the producer's epoch on this partition.
108    pub fn maybe_start_transaction_verification(
109        &mut self,
110        batch: TransactionalBatch,
111        supports_epoch_bump: bool,
112        (now_ms, expiration_ms): (i64, i64),
113    ) -> Result<VerificationGuard, TransactionAppendRefusal> {
114        let TransactionalBatch {
115            producer_id,
116            producer_epoch,
117            base_sequence,
118            ..
119        } = batch;
120        if self
121            .producer_state
122            .get(&producer_id)
123            .is_some_and(|entry| producer_epoch < entry.producer_epoch)
124        {
125            return Err(TransactionAppendRefusal::StaleProducerEpoch);
126        }
127        if self.has_ongoing_transaction(producer_id, producer_epoch) {
128            return Ok(VerificationGuard::SENTINEL);
129        }
130        self.verification_states
131            .retain(|_, state| now_ms.saturating_sub(state.created_ms) < expiration_ms);
132        let state = self
133            .verification_states
134            .entry(producer_id)
135            .or_insert_with(|| VerificationState {
136                guard: VerificationGuard::next(),
137                created_ms: now_ms,
138                epoch: producer_epoch,
139                lowest_sequence: base_sequence,
140                supports_epoch_bump,
141            });
142        if producer_epoch == state.epoch && base_sequence < state.lowest_sequence {
143            state.lowest_sequence = base_sequence;
144        }
145        if producer_epoch > state.epoch {
146            state.epoch = producer_epoch;
147            state.lowest_sequence = base_sequence;
148        }
149        Ok(state.guard)
150    }
151
152    /// Check a batch against the producer's transaction state just before it
153    /// appends, under the same lock as the append.
154    ///
155    /// A transactional data batch with no open transaction at its epoch must
156    /// present the producer's current guard. A non-transactional data batch
157    /// from a producer with an open transaction is refused. This is Kafka's
158    /// check in `UnifiedLog.analyzeAndValidateProducerState` and in
159    /// `ProducerAppendInfo.appendDataBatch`.
160    ///
161    /// # Errors
162    ///
163    /// Returns the refusal that Kafka's append throws.
164    pub fn check_transactional_append(
165        &self,
166        batch: TransactionalBatch,
167        presented: VerificationGuard,
168    ) -> Result<(), TransactionAppendRefusal> {
169        if batch.producer_id.get() < 0 || batch.is_control {
170            return Ok(());
171        }
172        let entry = self.producer_state.get(&batch.producer_id);
173        if !batch.is_transactional {
174            return if entry.is_some_and(|entry| entry.current_txn_first_offset.is_some()) {
175                Err(TransactionAppendRefusal::InvalidTransactionState)
176            } else {
177                Ok(())
178            };
179        }
180        if self.has_ongoing_transaction(batch.producer_id, batch.producer_epoch) {
181            return Ok(());
182        }
183        if entry.is_some_and(|entry| batch.producer_epoch < entry.producer_epoch) {
184            return Err(TransactionAppendRefusal::StaleProducerEpoch);
185        }
186        let Some(state) = self
187            .verification_states
188            .get(&batch.producer_id)
189            .filter(|state| state.guard.verifies(presented))
190        else {
191            return Err(TransactionAppendRefusal::InvalidTransactionState);
192        };
193        // Kafka's `ProducerAppendInfo.checkSequence` with a verification
194        // state: a transaction version 2 producer with no retained batch on
195        // the partition starts at 0, and no batch may start above the lowest
196        // sequence a verification saw.
197        let has_batch = entry.is_some_and(|entry| entry.last_offset.0 >= 0);
198        if (state.supports_epoch_bump && batch.base_sequence != 0 && !has_batch)
199            || batch.base_sequence > state.lowest_sequence
200        {
201            return Err(TransactionAppendRefusal::OutOfOrderSequence);
202        }
203        Ok(())
204    }
205
206    /// Drop the verification state after a transactional append, as Kafka's
207    /// `UnifiedLog.updateProducers` does. A data batch now has an open
208    /// transaction, and a marker ended one. The one exception is a transaction
209    /// version 2 marker at the verified epoch: the producer already started
210    /// its next transaction at that epoch, so its guard stays.
211    pub(super) fn clear_verification_after_append(
212        &mut self,
213        producer_id: ProducerId,
214        producer_epoch: i16,
215        is_control: bool,
216    ) {
217        let next_transaction_started =
218            self.verification_states
219                .get(&producer_id)
220                .is_some_and(|state| {
221                    state.supports_epoch_bump && is_control && producer_epoch == state.epoch
222                });
223        if !next_transaction_started {
224            self.verification_states.remove(&producer_id);
225        }
226    }
227}
228
229#[cfg(test)]
230mod tests;
krabka_log/log/verification.rs

krabka_log/log/
verification.rs
