krabka_client_producer/
metadata_age.rs

1//! The periodic metadata refresh of the producer.
2//!
3//! Kafka's producer refreshes its metadata when the last successful refresh
4//! is older than `metadata.max.age.ms` (`Metadata.timeToNextUpdate`), and it
5//! stops asking for a topic that had no send for `metadata.max.idle.ms`
6//! (`ProducerMetadata.retainTopic`). The refresh gives the producer a larger
7//! partition count and moved leaders without a send error.
8
9use std::{
10    sync::{
11        Arc, Weak,
12        atomic::{AtomicUsize, Ordering},
13    },
14    time::Duration,
15};
16
17use dashmap::DashMap;
18use krabka_client_core::Client;
19use tokio::time::{self, Instant};
20use tokio_util::sync::CancellationToken;
21
22use crate::{
23    accumulator::AccumulatorMap, builder::ProducerRetryPolicy, error::ProducerError,
24    metadata_wait::MetadataRefresh, producer::MetadataCache, sender::adopt_metadata,
25};
26
27/// Kafka's producer `metadata.max.age.ms` default (5 minutes).
28pub const DEFAULT_PRODUCER_METADATA_MAX_AGE: Duration = Duration::from_mins(5);
29
30/// Kafka's producer `metadata.max.idle.ms` default (5 minutes).
31pub const DEFAULT_PRODUCER_METADATA_MAX_IDLE: Duration = Duration::from_mins(5);
32
33/// Kafka's lower bound of `metadata.max.idle.ms` (`ProducerConfig`
34/// `atLeast(5000)`).
35pub const MIN_PRODUCER_METADATA_MAX_IDLE: Duration = Duration::from_secs(5);
36
37/// The state that the periodic refresh reads and updates.
38pub struct MetadataAge {
39    client: Client,
40    max_age: Duration,
41    max_idle: Duration,
42    retry_backoff: Duration,
43    retry_backoff_max: Duration,
44    metadata_cache: Arc<MetadataCache>,
45    partition_leaders: Arc<DashMap<(String, i32), i32>>,
46    metadata_refresh: Weak<MetadataRefresh>,
47    /// Cancelled when the producer drops its `MetadataRefresh`.
48    producer_dropped: CancellationToken,
49    accumulators: AccumulatorMap,
50    in_flight: Arc<AtomicUsize>,
51}
52
53/// Validate `metadata_max_idle`.
54///
55/// # Errors
56/// Returns [`ProducerError::InvalidConfig`] for a value below 5 s, as Kafka's
57/// `ProducerConfig` defines `metadata.max.idle.ms` with `atLeast(5000)`.
58pub fn validate_metadata_max_idle(max_idle: Duration) -> Result<(), ProducerError> {
59    if max_idle < MIN_PRODUCER_METADATA_MAX_IDLE {
60        return Err(ProducerError::InvalidConfig(format!(
61            "metadata_max_idle must be at least 5 s, got {max_idle:?}"
62        )));
63    }
64    Ok(())
65}
66
67impl MetadataAge {
68    /// A refresh for `client` with Kafka's default ages and the retry backoff
69    /// of `retry_policy`. [`validate_metadata_max_idle`] checks the idle time
70    /// before the producer connects.
71    pub(crate) fn new(client: &Client, retry_policy: &ProducerRetryPolicy) -> Self {
72        Self {
73            client: client.clone(),
74            max_age: DEFAULT_PRODUCER_METADATA_MAX_AGE,
75            max_idle: DEFAULT_PRODUCER_METADATA_MAX_IDLE,
76            retry_backoff: retry_policy.retry_backoff(),
77            retry_backoff_max: retry_policy.retry_backoff_max(),
78            metadata_cache: Arc::default(),
79            partition_leaders: Arc::default(),
80            metadata_refresh: Weak::new(),
81            producer_dropped: CancellationToken::new(),
82            accumulators: AccumulatorMap::default(),
83            in_flight: Arc::default(),
84        }
85    }
86
87    /// Use `metadata.max.age.ms` and `metadata.max.idle.ms`.
88    pub(crate) const fn with_ages(mut self, max_age: Duration, max_idle: Duration) -> Self {
89        self.max_age = max_age;
90        self.max_idle = max_idle;
91        self
92    }
93
94    /// Use the caches of the producer.
95    pub(crate) fn with_caches(
96        mut self,
97        metadata_cache: &Arc<MetadataCache>,
98        partition_leaders: &Arc<DashMap<(String, i32), i32>>,
99        metadata_refresh: &Arc<MetadataRefresh>,
100    ) -> Self {
101        self.metadata_cache = Arc::clone(metadata_cache);
102        self.partition_leaders = Arc::clone(partition_leaders);
103        self.metadata_refresh = Arc::downgrade(metadata_refresh);
104        self.producer_dropped = metadata_refresh.dropped();
105        self
106    }
107
108    /// Use the record queues of the producer. A topic with queued or in-flight
109    /// records does not expire.
110    pub(crate) fn with_queues(
111        mut self,
112        (accumulators, in_flight): (&AccumulatorMap, &Arc<AtomicUsize>),
113    ) -> Self {
114        self.accumulators = Arc::clone(accumulators);
115        self.in_flight = Arc::clone(in_flight);
116        self
117    }
118
119    /// Run [`Self::run`] on a task.
120    pub(crate) fn spawn(self, shutdown: CancellationToken) {
121        tokio::spawn(self.run(shutdown));
122    }
123
124    /// Refresh the metadata each time it becomes older than `max_age`, until
125    /// `shutdown`.
126    async fn run(self, shutdown: CancellationToken) {
127        let started = Instant::now();
128        let mut backoff = self.retry_backoff;
129        loop {
130            let topics = self.client.metadata_topics();
131            let last = topics.last_refreshed().unwrap_or(started);
132            tokio::select! {
133                () = shutdown.cancelled() => return,
134                () = self.producer_dropped.cancelled() => return,
135                () = time::sleep_until(last + self.max_age) => {}
136            }
137            if topics
138                .last_refreshed()
139                .is_some_and(|refreshed| refreshed > last)
140            {
141                // Another refresh came while this task slept.
142                continue;
143            }
144            self.expire_idle_topics().await;
145            match self.client.refresh_metadata().await {
146                Ok(response) => {
147                    adopt_metadata(&response, &self.metadata_cache, &self.partition_leaders);
148                    backoff = self.retry_backoff;
149                }
150                Err(error) => {
151                    tracing::debug!(error = %error, "periodic producer metadata refresh failed");
152                    tokio::select! {
153                        () = shutdown.cancelled() => return,
154                        () = time::sleep(backoff) => {}
155                    }
156                    backoff = backoff.saturating_mul(2).min(self.retry_backoff_max);
157                }
158            }
159        }
160    }
161
162    /// Whether an accumulator of `topic` holds a record.
163    async fn has_queued_records(&self, topic: &str) -> bool {
164        let queues = self
165            .accumulators
166            .iter()
167            .filter(|entry| entry.key().0 == topic)
168            .map(|entry| Arc::clone(entry.value()))
169            .collect::<Vec<_>>();
170        for queue in queues {
171            if queue.lock().await.queue_size() > 0 {
172                return true;
173            }
174        }
175        false
176    }
177
178    /// Drop the topics with no send for `max_idle` from the requests and the
179    /// caches. The next send to such a topic waits for its metadata again.
180    async fn expire_idle_topics(&self) {
181        // Kafka's `Sender.sendProducerData` adds a topic with batched records
182        // to the metadata again, so such a topic never leaves it. The same
183        // holds here for queued records and for records in flight.
184        let topics = self.client.metadata_topics();
185        if self.in_flight.load(Ordering::Acquire) > 0 {
186            return;
187        }
188        for topic in topics.names() {
189            if self.has_queued_records(&topic).await {
190                topics.add(&topic);
191            }
192        }
193        let expired = topics.remove_idle(self.max_idle);
194        if expired.is_empty() {
195            return;
196        }
197        tracing::debug!(topics = ?expired, "producer metadata topics expired");
198        if let Some(refresh) = self.metadata_refresh.upgrade() {
199            refresh.forget(&expired);
200        }
201        let mut cache = self.metadata_cache.lock();
202        for topic in &expired {
203            cache.remove(topic);
204        }
205        drop(cache);
206        self.partition_leaders
207            .retain(|(topic, _), _| !expired.contains(topic));
208    }
209}
210
211#[cfg(test)]
212mod tests {
213    use std::sync::{
214        Mutex as StdMutex,
215        atomic::{AtomicI32, AtomicU16, Ordering},
216    };
217
218    use assert2::check;
219    use bytes::BytesMut;
220    use krabka_client_core::MockBroker;
221    use krabka_protocol::{
222        Decode, Encode,
223        owned::{
224            api_versions_request,
225            api_versions_response::{ApiVersion, ApiVersionsResponse},
226            metadata_request::{self, MetadataRequest},
227            metadata_response::{
228                self, MetadataResponse, MetadataResponseBroker, MetadataResponsePartition,
229                MetadataResponseTopic,
230            },
231        },
232        primitives::uuid::Uuid,
233    };
234    use tokio::sync::Mutex;
235
236    use super::*;
237    use crate::{Producer, accumulator::Accumulator, producer::TopicMetadata};
238
239    const TOPIC: &str = "orders";
240    const CLIENT_ID: &str = "metadata-age-test";
241
242    /// A broker whose partition count and partition 0 leader a test changes.
243    /// It records the topics of each Metadata request.
244    struct ChangingBroker {
245        mock: MockBroker,
246        partitions: Arc<AtomicI32>,
247        leader: Arc<AtomicI32>,
248        requests: Arc<StdMutex<Vec<MetadataRequest>>>,
249    }
250
251    async fn changing_broker() -> ChangingBroker {
252        let port = Arc::new(AtomicU16::new(0));
253        let partitions = Arc::new(AtomicI32::new(4));
254        let leader = Arc::new(AtomicI32::new(1));
255        let requests = Arc::new(StdMutex::new(Vec::new()));
256        let (h_port, h_partitions, h_leader, h_requests) = (
257            Arc::clone(&port),
258            Arc::clone(&partitions),
259            Arc::clone(&leader),
260            Arc::clone(&requests),
261        );
262        let mock = MockBroker::start(move |api_key, version, _corr_id, body| {
263            if api_key == api_versions_request::API_KEY {
264                let mut buf = BytesMut::new();
265                ApiVersionsResponse {
266                    api_keys: vec![ApiVersion {
267                        api_key: metadata_request::API_KEY,
268                        min_version: 0,
269                        max_version: 12,
270                        ..Default::default()
271                    }],
272                    ..Default::default()
273                }
274                .encode(&mut buf, 0)
275                .unwrap();
276                return Some(buf.to_vec());
277            }
278            if api_key != metadata_request::API_KEY {
279                return None;
280            }
281            let header_len =
282                2 + CLIENT_ID.len() + usize::from(version >= metadata_request::FLEXIBLE_MIN);
283            let mut request_body = &body[header_len..];
284            let request = MetadataRequest::decode(&mut request_body, version).unwrap();
285            let named = request.topics.as_ref().is_some_and(|topics| {
286                topics
287                    .iter()
288                    .any(|topic| topic.name.as_deref() == Some(TOPIC))
289            });
290            h_requests.lock().unwrap().push(request);
291            let port = i32::from(h_port.load(Ordering::SeqCst));
292            let response = MetadataResponse {
293                brokers: [1, 2]
294                    .into_iter()
295                    .map(|node_id| MetadataResponseBroker {
296                        node_id,
297                        host: "127.0.0.1".into(),
298                        port,
299                        ..Default::default()
300                    })
301                    .collect(),
302                topics: named
303                    .then(|| MetadataResponseTopic {
304                        name: Some(TOPIC.into()),
305                        partitions: (0..h_partitions.load(Ordering::SeqCst))
306                            .map(|partition_index| MetadataResponsePartition {
307                                partition_index,
308                                leader_id: if partition_index == 0 {
309                                    h_leader.load(Ordering::SeqCst)
310                                } else {
311                                    1
312                                },
313                                ..Default::default()
314                            })
315                            .collect(),
316                        ..Default::default()
317                    })
318                    .into_iter()
319                    .collect(),
320                ..Default::default()
321            };
322            let mut buf = BytesMut::new();
323            if version >= metadata_response::FLEXIBLE_MIN {
324                buf.extend_from_slice(&[0]);
325            }
326            response.encode(&mut buf, version).unwrap();
327            Some(buf.to_vec())
328        })
329        .await
330        .unwrap();
331        port.store(mock.addr.port(), Ordering::SeqCst);
332        ChangingBroker {
333            mock,
334            partitions,
335            leader,
336            requests,
337        }
338    }
339
340    #[derive(Clone, Copy, Debug)]
341    enum Change {
342        Partitions(i32),
343        Leader(i32),
344    }
345
346    /// What the producer holds for the topic after the wait.
347    #[derive(Debug, PartialEq, Eq)]
348    struct Held {
349        partitions: Option<i32>,
350        leader_of_partition_0: Option<i32>,
351    }
352
353    /// Kafka's `Metadata.timeToNextUpdate`: the producer refreshes when the
354    /// last successful refresh is older than `metadata.max.age.ms`. The test
355    /// uses a 400 ms age in place of Kafka's 5 minutes.
356    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
357    async fn the_producer_adopts_broker_changes_after_metadata_max_age() {
358        let max_age = Duration::from_millis(400);
359        let before_age = Duration::from_millis(150);
360        let after_age = Duration::from_millis(700);
361        for (name, change, wait, expected) in [
362            (
363                "partitions 4 to 8, after the age",
364                Change::Partitions(8),
365                after_age,
366                Held {
367                    partitions: Some(8),
368                    leader_of_partition_0: Some(1),
369                },
370            ),
371            (
372                "partitions 4 to 8, before the age",
373                Change::Partitions(8),
374                before_age,
375                Held {
376                    partitions: Some(4),
377                    leader_of_partition_0: Some(1),
378                },
379            ),
380            (
381                "leader of partition 0 moves to broker 2, after the age",
382                Change::Leader(2),
383                after_age,
384                Held {
385                    partitions: Some(4),
386                    leader_of_partition_0: Some(2),
387                },
388            ),
389        ] {
390            let broker = changing_broker().await;
391            let producer = Producer::builder()
392                .bootstrap(broker.mock.addr.to_string())
393                .client_id(CLIENT_ID)
394                .enable_idempotence(false)
395                .metadata_max_age(max_age)
396                .build()
397                .await
398                .expect("producer connects");
399            producer
400                .partition_count(TOPIC, None)
401                .await
402                .expect("topic metadata");
403            // Start the age at a refresh of the full topic list.
404            producer.client.refresh_metadata().await.unwrap();
405            match change {
406                Change::Partitions(count) => broker.partitions.store(count, Ordering::SeqCst),
407                Change::Leader(leader) => broker.leader.store(leader, Ordering::SeqCst),
408            }
409            time::sleep(wait).await;
410            let held = Held {
411                partitions: producer
412                    .metadata_cache
413                    .lock()
414                    .get(TOPIC)
415                    .map(|topic| topic.num_partitions),
416                leader_of_partition_0: producer
417                    .partition_leaders
418                    .get(&(TOPIC.to_owned(), 0))
419                    .map(|leader| *leader),
420            };
421            broker.mock.stop();
422            drop(producer);
423            check!(held == expected, "{name}");
424        }
425    }
426
427    /// What the idle expiry left.
428    #[derive(Debug, PartialEq, Eq)]
429    struct AfterIdle {
430        last_request: Option<MetadataRequest>,
431        cached: bool,
432        leaders: usize,
433        task_finished: bool,
434    }
435
436    /// How the producer holds the topic while it is idle.
437    #[derive(Clone, Copy, Debug)]
438    enum Idle {
439        /// No record of the topic waits.
440        NoRecords,
441        /// A record waits in an accumulator.
442        QueuedRecord,
443        /// A batch is in flight.
444        InFlight,
445    }
446
447    /// An accumulator that holds one queued record.
448    fn accumulator_with_one_record() -> Accumulator {
449        let mut accumulator = Accumulator::new(1024);
450        let _ = accumulator.append(None, Some("v".into()), Vec::new(), 0, None, None);
451        accumulator
452    }
453
454    /// Run a periodic refresh with a 300 ms age and a 200 ms idle time for
455    /// 500 ms, then drop the producer's `MetadataRefresh`.
456    async fn idle_expiry(idle: Idle) -> AfterIdle {
457        let broker = changing_broker().await;
458        let client = Client::builder()
459            .bootstrap(broker.mock.addr.to_string())
460            .client_id(CLIENT_ID)
461            .metadata_scope(krabka_client_core::MetadataScope::Topics {
462                allow_auto_topic_creation: true,
463            })
464            .build()
465            .await
466            .unwrap();
467        client.metadata_topics().add(TOPIC);
468        let metadata_cache = Arc::new(MetadataCache::from(
469            maplit::hashmap! {TOPIC.to_owned() => TopicMetadata {
470                num_partitions: 4,
471                topic_id: Uuid::ZERO,
472            }},
473        ));
474        let partition_leaders = Arc::new(DashMap::from_iter([((TOPIC.to_owned(), 0), 1)]));
475        let accumulators = AccumulatorMap::default();
476        let in_flight = Arc::new(AtomicUsize::new(0));
477        match idle {
478            Idle::NoRecords => {}
479            Idle::QueuedRecord => {
480                accumulators.insert(
481                    (TOPIC.to_owned(), 0),
482                    Arc::new(Mutex::new(accumulator_with_one_record())),
483                );
484            }
485            Idle::InFlight => in_flight.store(1, Ordering::SeqCst),
486        }
487        let retry_policy = ProducerRetryPolicy::default();
488        let metadata_refresh = Arc::new(MetadataRefresh::default());
489        let task = tokio::spawn(
490            MetadataAge::new(&client, &retry_policy)
491                .with_ages(Duration::from_millis(300), Duration::from_millis(200))
492                .with_caches(&metadata_cache, &partition_leaders, &metadata_refresh)
493                .with_queues((&accumulators, &in_flight))
494                .run(CancellationToken::new()),
495        );
496        time::sleep(Duration::from_millis(500)).await;
497        // The producer drops its `MetadataRefresh` when it drops.
498        drop(metadata_refresh);
499        let task_finished = time::timeout(Duration::from_secs(1), task).await.is_ok();
500        let last_request = broker.requests.lock().unwrap().last().cloned();
501        let observed = AfterIdle {
502            last_request,
503            cached: metadata_cache.lock().contains_key(TOPIC),
504            leaders: partition_leaders.len(),
505            task_finished,
506        };
507        broker.mock.stop();
508        observed
509    }
510
511    /// Kafka's `ProducerMetadata.retainTopic`: a topic with no send for
512    /// `metadata.max.idle.ms` leaves the requests and the caches. Kafka's
513    /// `Sender` adds a topic with batched records again, so a topic with
514    /// queued or in-flight records stays. The task ends when the producer
515    /// drops.
516    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
517    async fn an_idle_topic_leaves_the_requests_and_the_caches_unless_records_wait() {
518        let named = |topics: &[&str]| {
519            Some(krabka_client_core::topics_request(
520                topics.iter().map(|&topic| topic.to_owned()),
521                true,
522            ))
523        };
524        for (name, idle, expected) in [
525            (
526                "no records",
527                Idle::NoRecords,
528                AfterIdle {
529                    last_request: named(&[]),
530                    cached: false,
531                    leaders: 0,
532                    task_finished: true,
533                },
534            ),
535            (
536                "a queued record",
537                Idle::QueuedRecord,
538                AfterIdle {
539                    last_request: named(&[TOPIC]),
540                    cached: true,
541                    leaders: 4,
542                    task_finished: true,
543                },
544            ),
545            (
546                "a batch in flight",
547                Idle::InFlight,
548                AfterIdle {
549                    last_request: named(&[TOPIC]),
550                    cached: true,
551                    leaders: 4,
552                    task_finished: true,
553                },
554            ),
555        ] {
556            check!(idle_expiry(idle).await == expected, "{name}");
557        }
558    }
559
560    #[tokio::test]
561    async fn metadata_max_idle_below_five_seconds_is_rejected() {
562        let error = Producer::builder()
563            .bootstrap("unused.invalid:9092")
564            .metadata_max_idle(Duration::from_millis(4999))
565            .build()
566            .await
567            .err();
568        check!(matches!(
569            error,
570            Some(ProducerError::InvalidConfig(message))
571                if message == "metadata_max_idle must be at least 5 s, got 4.999s"
572        ));
573    }
574}
krabka_client_producer/metadata_age.rs

krabka_client_producer/
metadata_age.rs
