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, HashSet},
14    sync::atomic::{AtomicU64, Ordering},
15};
16
17use crate::{consumer::Consumer, error::ConsumerError, poll::COMMITTED_SENTINEL};
18
19/// Ownership ids of manual assignments. They start far above the ids of
20/// group assignments, and they never repeat, so a commit that waits across
21/// an `assign` never takes a new ownership for the old one.
22static MANUAL_OWNERSHIP_IDS: AtomicU64 = AtomicU64::new(1 << 62);
23
24impl Consumer {
25    /// Fetch `partitions`, in place of the current manual assignment. Kafka's
26    /// `KafkaConsumer.assign`.
27    ///
28    /// A partition that stays assigned keeps its position. A new partition
29    /// starts at the committed offset of the group, or by `auto_offset_reset`.
30    /// An empty slice gives up the assignment, as Kafka's `assign` of an empty
31    /// collection calls `unsubscribe`.
32    ///
33    /// # Errors
34    ///
35    /// Returns [`ConsumerError::InvalidArgument`] for an empty topic name, and
36    /// [`ConsumerError::IllegalState`] while the consumer has a subscription.
37    pub async fn assign(&self, partitions: &[(String, i32)]) -> Result<(), ConsumerError> {
38        if partitions.iter().any(|(topic, _)| topic.is_empty()) {
39            return Err(ConsumerError::InvalidArgument(
40                "Topic partitions to assign to cannot have null or empty topic".to_owned(),
41            ));
42        }
43        if let Some((topic, partition)) = partitions.iter().find(|(_, partition)| *partition < 0) {
44            return Err(ConsumerError::InvalidArgument(format!(
45                "Partition index of {topic}-{partition} is negative"
46            )));
47        }
48        if !self.subscription.borrow().is_none() {
49            return Err(ConsumerError::IllegalState(
50                "Subscription to topics, partitions and pattern are mutually exclusive".to_owned(),
51            ));
52        }
53        let wanted: Vec<(String, i32)> = partitions
54            .iter()
55            .cloned()
56            .collect::<BTreeSet<_>>()
57            .into_iter()
58            .collect();
59        let wanted_set: HashSet<(String, i32)> = wanted.iter().cloned().collect();
60        let added: Vec<(String, i32)> = {
61            // Lock order: `assigned`, then `commit_identity`, as the
62            // coordinator task publishes an assignment.
63            let mut assigned = self.assigned.lock().await;
64            let mut identity = self.commit_identity.lock().await;
65            let added = wanted
66                .iter()
67                .filter(|partition| !assigned.contains(*partition))
68                .cloned()
69                .collect();
70            identity
71                .ownership_ids
72                .retain(|partition, _| wanted_set.contains(partition));
73            for partition in &wanted {
74                identity
75                    .ownership_ids
76                    .entry(partition.clone())
77                    .or_insert_with(|| MANUAL_OWNERSHIP_IDS.fetch_add(1, Ordering::Relaxed));
78            }
79            identity.generation = -1;
80            identity.member_id.clear();
81            identity.rejoin_on_poll = false;
82            assigned.clone_from(&wanted);
83            added
84        };
85        {
86            let mut offsets = self.next_offsets.lock().await;
87            let mut positions = self.positions.lock().await;
88            offsets.retain(|partition, _| wanted_set.contains(partition));
89            positions.retain(|partition, _| wanted_set.contains(partition));
90            for partition in added {
91                offsets.insert(partition.clone(), COMMITTED_SENTINEL);
92                positions.entry(partition).or_default();
93            }
94        }
95        self.end_offsets
96            .lock()
97            .await
98            .retain(|partition, _| wanted_set.contains(partition));
99        self.current_generation.store(-1, Ordering::Release);
100        // Kafka's `assignFromUser` changes the metadata topics to the topics of
101        // the assignment.
102        self.client.metadata_topics().set(
103            wanted
104                .iter()
105                .map(|(topic, _)| topic.clone())
106                .collect::<BTreeSet<_>>(),
107        );
108        self.subscription.send_if_modified(|subscription| {
109            let manual = !wanted.is_empty();
110            let changed = subscription.manual_assignment != manual;
111            subscription.manual_assignment = manual;
112            changed
113        });
114        self.assignment_changed.notify_waiters();
115        Ok(())
116    }
117
118    /// Fail a call that needs a group id. Kafka's
119    /// `throwIfGroupIdNotDefined`.
120    pub(crate) fn require_group_id(&self) -> Result<(), ConsumerError> {
121        if self.group_id.is_empty() {
122            return Err(ConsumerError::InvalidGroupId);
123        }
124        Ok(())
125    }
126
127    /// Fail a call that needs the group membership of a coordinator task.
128    pub(crate) fn require_group_membership(&self) -> Result<(), ConsumerError> {
129        self.require_group_id()?;
130        if self.coordinator_handle.is_none() {
131            return Err(ConsumerError::IllegalState(
132                "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(),
133            ));
134        }
135        Ok(())
136    }
137}
krabka_client_consumer/assignment.rs

krabka_client_consumer/
assignment.rs
