1use 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#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
17pub enum MetadataScope {
18 #[default]
22 AllTopics,
23 Topics { allow_auto_topic_creation: bool },
28}
29
30#[derive(Debug, Default)]
35pub struct MetadataTopics {
36 scope: MetadataScope,
37 topics: Mutex<BTreeMap<String, Instant>>,
38 refreshed: Mutex<Option<Instant>>,
40 update_requested: AtomicBool,
43}
44
45impl MetadataTopics {
46 #[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 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 pub fn request_update(&self) {
72 self.update_requested.store(true, Ordering::SeqCst);
73 }
74
75 #[must_use]
77 pub fn update_requested(&self) -> bool {
78 self.update_requested.load(Ordering::SeqCst)
79 }
80
81 #[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 #[must_use]
97 pub const fn scope(&self) -> MetadataScope {
98 self.scope
99 }
100
101 pub fn add(&self, topic: &str) -> bool {
104 self.topics()
105 .insert(topic.to_owned(), Instant::now())
106 .is_none()
107 }
108
109 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 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 #[must_use]
148 pub fn names(&self) -> Vec<String> {
149 self.topics().keys().cloned().collect()
150 }
151
152 #[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#[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#[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 #[test]
211 fn a_topic_change_or_a_request_asks_for_an_update_until_a_refresh() {
212 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}