Skip to main content

krabka_client_core/
metadata_topics.rs

1//! The topics that a client names in its `Metadata` requests.
2
3use std::{
4    collections::BTreeMap,
5    sync::{
6        Mutex, MutexGuard, PoisonError,
7        atomic::{AtomicBool, Ordering},
8    },
9    time::Duration,
10};
11
12use krabka_protocol::owned::metadata_request::{MetadataRequest, MetadataRequestTopic};
13use tokio::time::Instant;
14
15/// Which topics [`crate::Client::refresh_metadata`] asks for.
16#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
17pub enum MetadataScope {
18    /// Every topic of the cluster (`MetadataRequest.Builder.allTopics`). A
19    /// full-cluster request never makes the broker create a topic. Kafka's
20    /// consumer uses it only for a client-side pattern subscription.
21    #[default]
22    AllTopics,
23    /// The topics in [`MetadataTopics`], as Kafka's `ProducerMetadata` and
24    /// `ConsumerMetadata` name them. With `allow_auto_topic_creation`, a broker
25    /// with `auto.create.topics.enable=true` creates a named topic that does
26    /// not exist.
27    Topics { allow_auto_topic_creation: bool },
28}
29
30/// The topics of a client's metadata requests, with the last time that each
31/// topic was used.
32///
33/// The lock is never held across an `.await`.
34#[derive(Debug, Default)]
35pub struct MetadataTopics {
36    scope: MetadataScope,
37    topics: Mutex<BTreeMap<String, Instant>>,
38    /// When a request of this scope last succeeded.
39    refreshed: Mutex<Option<Instant>>,
40    /// Whether a caller asked for a refresh before the metadata expires,
41    /// Kafka's `Metadata.requestUpdate`.
42    update_requested: AtomicBool,
43}
44
45impl MetadataTopics {
46    /// An empty topic set with `scope`.
47    #[must_use]
48    pub fn new(scope: MetadataScope) -> Self {
49        Self {
50            scope,
51            topics: Mutex::default(),
52            refreshed: Mutex::default(),
53            update_requested: AtomicBool::new(false),
54        }
55    }
56
57    /// Record a successful response to [`Self::request`], as Kafka's
58    /// `Metadata.update` sets `lastSuccessfulRefreshMs` for a full update. It
59    /// answers any [`request_update`](Self::request_update) made before it.
60    pub fn mark_refreshed(&self) {
61        self.update_requested.store(false, Ordering::SeqCst);
62        *self
63            .refreshed
64            .lock()
65            .unwrap_or_else(PoisonError::into_inner) = Some(Instant::now());
66    }
67
68    /// Ask for a refresh before the metadata expires, as Kafka's
69    /// `Metadata.requestUpdate` does after a fetch names a stale leader or an
70    /// unknown topic.
71    pub fn request_update(&self) {
72        self.update_requested.store(true, Ordering::SeqCst);
73    }
74
75    /// Whether a refresh was asked for since the last successful one.
76    #[must_use]
77    pub fn update_requested(&self) -> bool {
78        self.update_requested.load(Ordering::SeqCst)
79    }
80
81    /// When a response to [`Self::request`] last succeeded, or `None` before
82    /// the first one.
83    #[must_use]
84    pub fn last_refreshed(&self) -> Option<Instant> {
85        *self
86            .refreshed
87            .lock()
88            .unwrap_or_else(PoisonError::into_inner)
89    }
90
91    fn topics(&self) -> MutexGuard<'_, BTreeMap<String, Instant>> {
92        self.topics.lock().unwrap_or_else(PoisonError::into_inner)
93    }
94
95    /// The scope of the requests.
96    #[must_use]
97    pub const fn scope(&self) -> MetadataScope {
98        self.scope
99    }
100
101    /// Add `topic`, or mark it as used now. Kafka's producer calls
102    /// `ProducerMetadata.add` for each send. Returns whether the topic is new.
103    pub fn add(&self, topic: &str) -> bool {
104        self.topics()
105            .insert(topic.to_owned(), Instant::now())
106            .is_none()
107    }
108
109    /// Replace the topics with `topics`, as a consumer subscription change
110    /// does (`SubscriptionState.metadataTopics`). A change of the topic set
111    /// asks for a refresh, as Kafka's `ConsumerMetadata` requests an update
112    /// for new topics.
113    pub fn set<I, S>(&self, topics: I)
114    where
115        I: IntoIterator<Item = S>,
116        S: Into<String>,
117    {
118        let now = Instant::now();
119        let topics: BTreeMap<String, Instant> = topics
120            .into_iter()
121            .map(|topic| (topic.into(), now))
122            .collect();
123        let mut current = self.topics();
124        if !current.keys().eq(topics.keys()) {
125            self.request_update();
126        }
127        *current = topics;
128    }
129
130    /// Remove the topics that were not used for `idle` and return their
131    /// names, as Kafka's `ProducerMetadata.retainTopic` drops a topic after
132    /// `metadata.max.idle.ms`.
133    pub fn remove_idle(&self, idle: Duration) -> Vec<String> {
134        let now = Instant::now();
135        let mut removed = Vec::new();
136        self.topics().retain(|topic, used| {
137            let keep = now.saturating_duration_since(*used) < idle;
138            if !keep {
139                removed.push(topic.clone());
140            }
141            keep
142        });
143        removed
144    }
145
146    /// The topic names, in order.
147    #[must_use]
148    pub fn names(&self) -> Vec<String> {
149        self.topics().keys().cloned().collect()
150    }
151
152    /// The `Metadata` request of the scope and the current topics.
153    #[must_use]
154    pub fn request(&self) -> MetadataRequest {
155        match self.scope {
156            MetadataScope::AllTopics => all_topics_request(),
157            MetadataScope::Topics {
158                allow_auto_topic_creation,
159            } => topics_request(self.names(), allow_auto_topic_creation),
160        }
161    }
162}
163
164/// The `Metadata` request for every topic of the cluster
165/// (`MetadataRequest.Builder.allTopics`): a null topic list.
166///
167/// `MetadataRequest::default()` is not this request. Its `topics` is the
168/// schema default, an empty list, which asks for no topic.
169#[must_use]
170pub fn all_topics_request() -> MetadataRequest {
171    MetadataRequest {
172        topics: None,
173        allow_auto_topic_creation: true,
174        ..Default::default()
175    }
176}
177
178/// A `Metadata` request that names `topics` (`MetadataRequest.Builder`
179/// `forTopicNames`). An empty list asks for no topic.
180#[must_use]
181pub fn topics_request(
182    topics: impl IntoIterator<Item = String>,
183    allow_auto_topic_creation: bool,
184) -> MetadataRequest {
185    MetadataRequest {
186        topics: Some(
187            topics
188                .into_iter()
189                .map(|name| MetadataRequestTopic {
190                    name: Some(name),
191                    ..Default::default()
192                })
193                .collect(),
194        ),
195        allow_auto_topic_creation,
196        ..Default::default()
197    }
198}
199
200#[cfg(test)]
201mod tests {
202    use assert2::check;
203    use tokio::time::advance;
204
205    use super::*;
206
207    /// A change of the topic set or an explicit request asks for a refresh,
208    /// and a successful refresh answers it. Setting the same topics again does
209    /// not.
210    #[test]
211    fn a_topic_change_or_a_request_asks_for_an_update_until_a_refresh() {
212        // (name, topics set after the first refresh, explicit request, expected)
213        let cases: [(&str, &[&str], bool, bool); 4] = [
214            ("same topics", &["t1"], false, false),
215            ("a new topic", &["t1", "t2"], false, true),
216            ("a removed topic", &[], false, true),
217            ("an explicit request", &["t1"], true, true),
218        ];
219        for (name, topics, request, expected) in cases {
220            let metadata = MetadataTopics::new(MetadataScope::Topics {
221                allow_auto_topic_creation: false,
222            });
223            metadata.set(["t1"]);
224            check!(metadata.update_requested(), "{name}: the first topics");
225            metadata.mark_refreshed();
226            check!(!metadata.update_requested(), "{name}");
227            metadata.set(topics.iter().copied());
228            if request {
229                metadata.request_update();
230            }
231            check!(metadata.update_requested() == expected, "{name}");
232            metadata.mark_refreshed();
233            check!(!metadata.update_requested(), "{name}: refreshed");
234        }
235    }
236
237    #[tokio::test(start_paused = true)]
238    async fn requests_follow_the_scope_and_the_topics() {
239        let named =
240            |names: &[&str], allow| topics_request(names.iter().map(|&n| n.to_owned()), allow);
241        let producer = MetadataScope::Topics {
242            allow_auto_topic_creation: true,
243        };
244        let no_auto_create = MetadataScope::Topics {
245            allow_auto_topic_creation: false,
246        };
247        for (name, scope, added, expected) in [
248            (
249                "all topics",
250                MetadataScope::AllTopics,
251                vec!["t1"],
252                MetadataRequest {
253                    topics: None,
254                    allow_auto_topic_creation: true,
255                    ..Default::default()
256                },
257            ),
258            ("no topic yet", producer, vec![], named(&[], true)),
259            (
260                "two sends",
261                producer,
262                vec!["t1", "t2", "t1"],
263                named(&["t1", "t2"], true),
264            ),
265            (
266                "auto creation off",
267                no_auto_create,
268                vec!["a"],
269                named(&["a"], false),
270            ),
271        ] {
272            let topics = MetadataTopics::new(scope);
273            for topic in added {
274                topics.add(topic);
275            }
276            check!(topics.request() == expected, "{name}");
277        }
278    }
279
280    #[tokio::test(start_paused = true)]
281    async fn set_replaces_and_remove_idle_drops_unused_topics() {
282        let topics = MetadataTopics::new(MetadataScope::Topics {
283            allow_auto_topic_creation: true,
284        });
285        topics.set(["a", "b"]);
286        check!(topics.names() == vec!["a".to_owned(), "b".to_owned()]);
287        advance(Duration::from_secs(200)).await;
288        check!(!topics.add("b"));
289        check!(topics.add("c"));
290        advance(Duration::from_secs(101)).await;
291        check!(topics.remove_idle(Duration::from_mins(5)) == vec!["a".to_owned()]);
292        check!(topics.names() == vec!["b".to_owned(), "c".to_owned()]);
293    }
294}