krabka_client_consumer/
assignment.rs

1//! Manual partition assignment: Kafka's `KafkaConsumer.assign`.
2//!
3//! A consumer without a subscription fetches the partitions that the
4//! application assigns. It joins no group. With a group id it can commit the
5//! offsets, with generation `-1` and an empty member id, as Kafka's
6//! consumer does for a manual assignment. A new partition gets its position
7//! at the next position update: the committed offset when the consumer has a
8//! group id and the group committed one, otherwise `auto_offset_reset`
9//! (Kafka's `ClassicKafkaConsumer.updateFetchPositions`,
10//! `refreshCommittedOffsetsIfNeeded` and `resetInitializingPositions`).
11
12use std::{
13    collections::{BTreeSet, HashMap, HashSet},
14    sync::atomic::{AtomicU64, Ordering},
15};
16
17use crate::{
18    consumer::{CommitIdentity, Consumer},
19    error::ConsumerError,
20    poll::COMMITTED_SENTINEL,
21    position::PartitionPosition,
22};
23
24/// Ownership ids of manual assignments. They start far above the ids of
25/// group assignments, and they never repeat, so a commit that waits across
26/// an `assign` never takes a new ownership for the old one.
27static MANUAL_OWNERSHIP_IDS: AtomicU64 = AtomicU64::new(1 << 62);
28
29impl Consumer {
30    /// Fetch `partitions`, in place of the current manual assignment. Kafka's
31    /// `KafkaConsumer.assign`.
32    ///
33    /// A partition that stays assigned keeps its position. A new partition
34    /// starts at the committed offset of the group, or by `auto_offset_reset`.
35    /// An empty slice gives up the assignment, as Kafka's `assign` of an empty
36    /// collection calls `unsubscribe`.
37    ///
38    /// # Errors
39    ///
40    /// Returns [`ConsumerError::InvalidArgument`] for an empty topic name, and
41    /// [`ConsumerError::IllegalState`] while the consumer has a subscription.
42    pub async fn assign(&self, partitions: &[(String, i32)]) -> Result<(), ConsumerError> {
43        if partitions.iter().any(|(topic, _)| topic.is_empty()) {
44            return Err(ConsumerError::InvalidArgument(
45                "Topic partitions to assign to cannot have null or empty topic".to_owned(),
46            ));
47        }
48        if let Some((topic, partition)) = partitions.iter().find(|(_, partition)| *partition < 0) {
49            return Err(ConsumerError::InvalidArgument(format!(
50                "Partition index of {topic}-{partition} is negative"
51            )));
52        }
53        if !self.subscription.borrow().is_none() {
54            return Err(ConsumerError::IllegalState(
55                "Subscription to topics, partitions and pattern are mutually exclusive".to_owned(),
56            ));
57        }
58        let wanted: Vec<(String, i32)> = partitions
59            .iter()
60            .cloned()
61            .collect::<BTreeSet<_>>()
62            .into_iter()
63            .collect();
64        let wanted_set: HashSet<(String, i32)> = wanted.iter().cloned().collect();
65        let added: Vec<(String, i32)> = {
66            // Lock order: `assigned`, then `commit_identity`, as the
67            // coordinator task publishes an assignment.
68            let mut assigned = self.assigned.lock().await;
69            let mut identity = self.commit_identity.lock().await;
70            replace_manual_assignment(&mut assigned, &mut identity, &wanted, &wanted_set)
71        };
72        {
73            let mut offsets = self.next_offsets.lock().await;
74            let mut positions = self.positions.lock().await;
75            reset_manual_positions(&mut offsets, &mut positions, &wanted_set, added);
76        }
77        self.end_offsets
78            .lock()
79            .await
80            .retain(|partition, _| wanted_set.contains(partition));
81        self.current_generation.store(-1, Ordering::Release);
82        // Kafka's `assignFromUser` changes the metadata topics to the topics of
83        // the assignment.
84        self.client
85            .metadata_topics()
86            .set(wanted.iter().map(|(topic, _)| topic.clone()));
87        self.subscription.send_if_modified(|subscription| {
88            let manual = !wanted.is_empty();
89            let changed = subscription.manual_assignment != manual;
90            subscription.manual_assignment = manual;
91            changed
92        });
93        self.assignment_changed.notify_waiters();
94        Ok(())
95    }
96
97    /// Fail a call that needs a group id. Kafka's
98    /// `throwIfGroupIdNotDefined`.
99    pub(crate) fn require_group_id(&self) -> Result<(), ConsumerError> {
100        if self.group_id.is_empty() {
101            return Err(ConsumerError::InvalidGroupId);
102        }
103        Ok(())
104    }
105
106    /// Fail a call that needs the group membership of a coordinator task.
107    pub(crate) fn require_group_membership(&self) -> Result<(), ConsumerError> {
108        self.require_group_id()?;
109        if self.coordinator_handle.is_none() {
110            return Err(ConsumerError::IllegalState(
111                "the consumer was built without a subscription, so it has no group membership; build it with subscribe or subscribe_pattern to join a group".to_owned(),
112            ));
113        }
114        Ok(())
115    }
116}
117
118/// Replace the manual assignment `assigned` with `wanted`, and give `identity`
119/// the generation `-1` ownership of `wanted`. The caller holds both locks for
120/// the whole change. Returns the partitions that `wanted` adds.
121fn replace_manual_assignment(
122    assigned: &mut Vec<(String, i32)>,
123    identity: &mut CommitIdentity,
124    wanted: &[(String, i32)],
125    wanted_set: &HashSet<(String, i32)>,
126) -> Vec<(String, i32)> {
127    let added = wanted
128        .iter()
129        .filter(|partition| !assigned.contains(*partition))
130        .cloned()
131        .collect();
132    identity
133        .ownership_ids
134        .retain(|partition, _| wanted_set.contains(partition));
135    for partition in wanted {
136        identity
137            .ownership_ids
138            .entry(partition.clone())
139            .or_insert_with(|| MANUAL_OWNERSHIP_IDS.fetch_add(1, Ordering::Relaxed));
140    }
141    identity.generation = -1;
142    identity.member_id.clear();
143    identity.rejoin_on_poll = false;
144    assigned.clear();
145    assigned.extend_from_slice(wanted);
146    added
147}
148
149/// Drop the offsets and positions of the partitions outside `wanted_set`, and
150/// start each `added` partition at its committed offset. The caller holds both
151/// locks for the whole change.
152fn reset_manual_positions(
153    offsets: &mut HashMap<(String, i32), i64>,
154    positions: &mut HashMap<(String, i32), PartitionPosition>,
155    wanted_set: &HashSet<(String, i32)>,
156    added: Vec<(String, i32)>,
157) {
158    offsets.retain(|partition, _| wanted_set.contains(partition));
159    positions.retain(|partition, _| wanted_set.contains(partition));
160    for partition in added {
161        offsets.insert(partition.clone(), COMMITTED_SENTINEL);
162        positions.entry(partition).or_default();
163    }
164}
krabka_client_consumer/assignment.rs

krabka_client_consumer/
assignment.rs
