krabka_replicator/
admin_util.rs

1//! Shared admin/IO helpers used by replicator engine tasks.
2//!
3//! All functions return `Result<_, String>` and map client errors with
4//! `.map_err(|e| e.to_string())`, so callers stay wire-error-agnostic.
5
6use std::collections::BTreeMap;
7
8use bytes::Bytes;
9use krabka_client_admin::{AdminClient, CreateTopicSpec, TopicMutationOptions};
10use krabka_client_consumer::{AutoOffsetReset, Consumer, IsolationLevel};
11use krabka_client_core::security::ClientSecurity;
12
13use crate::config::{ClientResourcePolicy, ReplicationFactor, ReplicatorRuntimePolicy};
14
15/// Kafka error code: the topic already exists.
16const TOPIC_ALREADY_EXISTS: i16 = 36;
17
18/// Ensure `topic` exists with the given parameters. An already-exists response
19/// counts as success.
20#[tracing::instrument(
21    level = "info",
22    skip_all,
23    fields(topic = %topic, partitions, bootstrap = %bootstrap),
24    err,
25)]
26/// # Errors
27/// Returns an error when configuration is invalid, protocol encoding fails, the broker rejects the request, or transport I/O fails.
28pub async fn ensure_topic(
29    bootstrap: &str,
30    topic: &str,
31    partitions: i32,
32    security: Option<ClientSecurity>,
33) -> Result<(), String> {
34    ensure_topic_with_policy(
35        bootstrap,
36        topic,
37        partitions,
38        security,
39        ClientResourcePolicy::default(),
40    )
41    .await
42}
43
44/// Ensure a topic using the deployment's client resource policy.
45///
46/// # Errors
47///
48/// Returns an error when configuration is invalid, protocol encoding fails,
49/// the broker rejects the request, or transport I/O fails.
50pub async fn ensure_topic_with_policy(
51    bootstrap: &str,
52    topic: &str,
53    partitions: i32,
54    security: Option<ClientSecurity>,
55    client_resource_policy: ClientResourcePolicy,
56) -> Result<(), String> {
57    let runtime_policy = ReplicatorRuntimePolicy::default();
58    ensure_topic_with_runtime_policy(
59        bootstrap,
60        topic,
61        partitions,
62        security,
63        client_resource_policy,
64        &runtime_policy,
65        runtime_policy.data_topic_replication_factor,
66    )
67    .await
68}
69
70/// Ensure a topic with an explicit replication factor.
71///
72/// # Errors
73///
74/// Returns an error when the admin client cannot create or inspect the topic.
75pub async fn ensure_topic_with_replication_factor(
76    bootstrap: &str,
77    topic: &str,
78    partitions: i32,
79    security: Option<ClientSecurity>,
80    replication_factor: ReplicationFactor,
81) -> Result<(), String> {
82    let runtime_policy = ReplicatorRuntimePolicy {
83        data_topic_replication_factor: replication_factor,
84        ..ReplicatorRuntimePolicy::default()
85    };
86    ensure_topic_with_runtime_policy(
87        bootstrap,
88        topic,
89        partitions,
90        security,
91        ClientResourcePolicy::default(),
92        &runtime_policy,
93        replication_factor,
94    )
95    .await
96}
97
98/// Ensure a topic using explicit process and replication policy.
99pub(crate) async fn ensure_topic_with_runtime_policy(
100    bootstrap: &str,
101    topic: &str,
102    partitions: i32,
103    security: Option<ClientSecurity>,
104    client_resource_policy: ClientResourcePolicy,
105    runtime_policy: &ReplicatorRuntimePolicy,
106    replication_factor: ReplicationFactor,
107) -> Result<(), String> {
108    let mut admin = AdminClient::connect_with_options(
109        &[bootstrap.to_string()],
110        admin_options(security, client_resource_policy, runtime_policy)?,
111    )
112    .await
113    .map_err(|e| e.to_string())?;
114
115    let outcomes = admin
116        .create_topics(
117            &[CreateTopicSpec {
118                name: topic.to_string(),
119                partitions,
120                replicas: i32::from(replication_factor.get()),
121                configs: BTreeMap::new(),
122                ..Default::default()
123            }],
124            TopicMutationOptions::with_timeout(runtime_policy.topic_create_timeout),
125        )
126        .await
127        .map_err(|e| e.to_string())?;
128
129    for outcome in &outcomes {
130        if let Some(ref err) = outcome.error
131            && err.code != TOPIC_ALREADY_EXISTS
132        {
133            return Err(format!(
134                "create_topic {topic}: error code {} ({}): {:?}",
135                err.code, err.name, err.message
136            ));
137        }
138    }
139
140    Ok(())
141}
142
143/// Ensure a compacted topic exists with one partition and the default replica count.
144#[tracing::instrument(
145    level = "info",
146    skip_all,
147    fields(topic = %topic, bootstrap = %bootstrap),
148    err,
149)]
150/// # Errors
151/// Returns an error when configuration is invalid, protocol encoding fails, the broker rejects the request, or transport I/O fails.
152pub async fn ensure_compacted_topic(
153    bootstrap: &str,
154    topic: &str,
155    security: Option<ClientSecurity>,
156) -> Result<(), String> {
157    ensure_compacted_topic_with_policy(bootstrap, topic, security, ClientResourcePolicy::default())
158        .await
159}
160
161/// Ensure a compacted topic using the deployment's client resource policy.
162///
163/// # Errors
164///
165/// Returns an error when configuration is invalid, protocol encoding fails,
166/// the broker rejects the request, or transport I/O fails.
167pub async fn ensure_compacted_topic_with_policy(
168    bootstrap: &str,
169    topic: &str,
170    security: Option<ClientSecurity>,
171    client_resource_policy: ClientResourcePolicy,
172) -> Result<(), String> {
173    let runtime_policy = ReplicatorRuntimePolicy::default();
174    ensure_compacted_topic_with_runtime_policy(
175        bootstrap,
176        topic,
177        security,
178        client_resource_policy,
179        &runtime_policy,
180    )
181    .await
182}
183
184/// Ensure a compacted topic with an explicit replication factor.
185///
186/// # Errors
187///
188/// Returns an error when the admin client cannot create or inspect the topic.
189pub async fn ensure_compacted_topic_with_replication_factor(
190    bootstrap: &str,
191    topic: &str,
192    security: Option<ClientSecurity>,
193    replication_factor: ReplicationFactor,
194) -> Result<(), String> {
195    let runtime_policy = ReplicatorRuntimePolicy {
196        internal_topic_replication_factor: replication_factor,
197        ..ReplicatorRuntimePolicy::default()
198    };
199    ensure_compacted_topic_with_runtime_policy(
200        bootstrap,
201        topic,
202        security,
203        ClientResourcePolicy::default(),
204        &runtime_policy,
205    )
206    .await
207}
208
209pub(crate) async fn ensure_compacted_topic_with_runtime_policy(
210    bootstrap: &str,
211    topic: &str,
212    security: Option<ClientSecurity>,
213    client_resource_policy: ClientResourcePolicy,
214    runtime_policy: &ReplicatorRuntimePolicy,
215) -> Result<(), String> {
216    let mut admin = AdminClient::connect_with_options(
217        &[bootstrap.to_string()],
218        admin_options(security, client_resource_policy, runtime_policy)?,
219    )
220    .await
221    .map_err(|e| e.to_string())?;
222
223    let mut configs = BTreeMap::new();
224    configs.insert("cleanup.policy".to_string(), "compact".to_string());
225
226    let outcomes = admin
227        .create_topics(
228            &[CreateTopicSpec {
229                name: topic.to_string(),
230                partitions: 1,
231                replicas: i32::from(runtime_policy.internal_topic_replication_factor.get()),
232                configs,
233                ..Default::default()
234            }],
235            TopicMutationOptions::with_timeout(runtime_policy.topic_create_timeout),
236        )
237        .await
238        .map_err(|e| e.to_string())?;
239
240    for outcome in &outcomes {
241        if let Some(ref err) = outcome.error
242            && err.code != TOPIC_ALREADY_EXISTS
243        {
244            return Err(format!(
245                "ensure_compacted_topic {topic}: error code {} ({}): {:?}",
246                err.code, err.name, err.message
247            ));
248        }
249    }
250
251    Ok(())
252}
253
254fn admin_options(
255    security: Option<ClientSecurity>,
256    policy: ClientResourcePolicy,
257    runtime_policy: &ReplicatorRuntimePolicy,
258) -> Result<krabka_client_core::ConnectionOptions, String> {
259    Ok(krabka_client_core::ConnectionOptions {
260        dns_timeout: krabka_client_core::ClientDnsTimeout::new(runtime_policy.client_dns_timeout)?,
261        socket_connection_setup_timeout: runtime_policy.client_connect_timeout,
262        socket_connection_setup_timeout_max: runtime_policy.client_connect_timeout,
263        request_timeout: runtime_policy.client_request_timeout,
264        client_id: "krabka-operator".to_owned(),
265        dispatch_queue_capacity: policy.dispatch_queue_capacity,
266        frame_max: policy.frame_max,
267        security: security.map(Box::new),
268        ..Default::default()
269    })
270}
271
272/// Build a drain consumer for the given topic. This function threads security
273/// through conditionally. `bon` wraps bare `T` setters in `Some`, so a `None`
274/// value means that the function omits the `.security()` call.
275async fn build_drain_consumer(
276    bootstrap: &str,
277    group_id: String,
278    topic: &str,
279    security: Option<ClientSecurity>,
280    client_resource_policy: ClientResourcePolicy,
281) -> Result<Consumer, krabka_client_consumer::ConsumerError> {
282    if let Some(sec) = security {
283        Consumer::builder()
284            .bootstrap(bootstrap)
285            .group_id(group_id)
286            .client_id("krabka-replicator-util")
287            .dispatch_queue_capacity(client_resource_policy.dispatch_queue_capacity.get())
288            .frame_max(client_resource_policy.frame_max.size())
289            .subscribe(vec![topic.to_string()])
290            .auto_offset_reset(AutoOffsetReset::Earliest)
291            .enable_auto_commit(false)
292            .allow_auto_create_topics(false)
293            .isolation_level(IsolationLevel::ReadCommitted)
294            .security(sec)
295            .build()
296            .await
297    } else {
298        Consumer::builder()
299            .bootstrap(bootstrap)
300            .group_id(group_id)
301            .client_id("krabka-replicator-util")
302            .dispatch_queue_capacity(client_resource_policy.dispatch_queue_capacity.get())
303            .frame_max(client_resource_policy.frame_max.size())
304            .subscribe(vec![topic.to_string()])
305            .auto_offset_reset(AutoOffsetReset::Earliest)
306            .enable_auto_commit(false)
307            .allow_auto_create_topics(false)
308            .isolation_level(IsolationLevel::ReadCommitted)
309            .build()
310            .await
311    }
312}
313
314/// Drain all records from `topic` from the earliest offset. Returns
315/// `(key, value)` pairs in order.
316///
317/// The function uses the consecutive-empty threshold of the runtime policy as
318/// the drain sentinel. It silently treats a poll error for a topic that does
319/// not exist yet as empty.
320pub type RawRecord = (Option<Bytes>, Option<Bytes>);
321
322#[tracing::instrument(
323    level = "debug",
324    skip_all,
325    fields(topic = %topic, drained = tracing::field::Empty),
326    err,
327)]
328/// # Errors
329/// Returns an error when configuration is invalid, protocol encoding fails, the broker rejects the request, or transport I/O fails.
330pub async fn read_all(
331    bootstrap: &str,
332    topic: &str,
333    security: Option<ClientSecurity>,
334) -> Result<Vec<RawRecord>, String> {
335    read_all_with_policy(bootstrap, topic, security, ClientResourcePolicy::default()).await
336}
337
338/// Drain a topic using the deployment's client resource policy.
339///
340/// # Errors
341///
342/// Returns an error when configuration is invalid, protocol encoding fails,
343/// the broker rejects the request, or transport I/O fails.
344pub async fn read_all_with_policy(
345    bootstrap: &str,
346    topic: &str,
347    security: Option<ClientSecurity>,
348    client_resource_policy: ClientResourcePolicy,
349) -> Result<Vec<RawRecord>, String> {
350    let runtime_policy = ReplicatorRuntimePolicy::default();
351    read_all_with_runtime_policy(
352        bootstrap,
353        topic,
354        security,
355        client_resource_policy,
356        &runtime_policy,
357    )
358    .await
359}
360
361pub(crate) async fn read_all_with_runtime_policy(
362    bootstrap: &str,
363    topic: &str,
364    security: Option<ClientSecurity>,
365    client_resource_policy: ClientResourcePolicy,
366    runtime_policy: &ReplicatorRuntimePolicy,
367) -> Result<Vec<RawRecord>, String> {
368    let group_id = format!("krabka-replicator-reader-{topic}");
369
370    let mut consumer =
371        match build_drain_consumer(bootstrap, group_id, topic, security, client_resource_policy)
372            .await
373        {
374            Ok(c) => c,
375            Err(e) => {
376                let msg = e.to_string();
377                if is_unknown_topic_error(&msg) {
378                    return Ok(Vec::new());
379                }
380                return Err(msg);
381            }
382        };
383
384    let mut records = Vec::new();
385    let mut consecutive_empty = 0usize;
386
387    loop {
388        match consumer
389            .poll(runtime_policy.internal_drain_poll_timeout)
390            .await
391        {
392            Ok(batch) => {
393                if batch.is_empty() {
394                    consecutive_empty += 1;
395                    if consecutive_empty >= runtime_policy.internal_drain_empty_polls.get() {
396                        break;
397                    }
398                } else {
399                    consecutive_empty = 0;
400                    for r in batch {
401                        records.push((r.key, r.value));
402                    }
403                }
404            }
405            Err(e) => {
406                let msg = e.to_string();
407                if is_unknown_topic_error(&msg) {
408                    break;
409                }
410                let _ = consumer.close().await;
411                return Err(msg);
412            }
413        }
414    }
415
416    let _ = consumer.close().await;
417    tracing::Span::current().record("drained", records.len());
418    Ok(records)
419}
420
421/// Return the value bytes of the last record whose key equals `key`.
422///
423/// If `key` is empty, returns the last record overall regardless of key.
424///
425/// # Errors
426/// Returns an error when configuration is invalid, protocol encoding fails, the broker rejects the request, or transport I/O fails.
427pub async fn read_last_value_for_key(
428    bootstrap: &str,
429    topic: &str,
430    key: &[u8],
431    security: Option<ClientSecurity>,
432) -> Result<Option<Vec<u8>>, String> {
433    read_last_value_for_key_with_policy(
434        bootstrap,
435        topic,
436        key,
437        security,
438        ClientResourcePolicy::default(),
439    )
440    .await
441}
442
443/// Read a key using the deployment's client resource policy.
444///
445/// # Errors
446///
447/// Returns an error when configuration is invalid, protocol encoding fails,
448/// the broker rejects the request, or transport I/O fails.
449pub async fn read_last_value_for_key_with_policy(
450    bootstrap: &str,
451    topic: &str,
452    key: &[u8],
453    security: Option<ClientSecurity>,
454    client_resource_policy: ClientResourcePolicy,
455) -> Result<Option<Vec<u8>>, String> {
456    let runtime_policy = ReplicatorRuntimePolicy::default();
457    read_last_value_for_key_with_runtime_policy(
458        bootstrap,
459        topic,
460        key,
461        security,
462        client_resource_policy,
463        &runtime_policy,
464    )
465    .await
466}
467
468pub(crate) async fn read_last_value_for_key_with_runtime_policy(
469    bootstrap: &str,
470    topic: &str,
471    key: &[u8],
472    security: Option<ClientSecurity>,
473    client_resource_policy: ClientResourcePolicy,
474    runtime_policy: &ReplicatorRuntimePolicy,
475) -> Result<Option<Vec<u8>>, String> {
476    let all = read_all_with_runtime_policy(
477        bootstrap,
478        topic,
479        security,
480        client_resource_policy,
481        runtime_policy,
482    )
483    .await?;
484
485    let matched = if key.is_empty() {
486        all.into_iter().last()
487    } else {
488        all.into_iter()
489            .filter(|(k, _)| k.as_deref() == Some(key))
490            .last()
491    };
492
493    Ok(matched.and_then(|(_, v)| v.map(|b| b.to_vec())))
494}
495
496/// Returns `true` if the error message shows that the topic does not exist.
497pub(crate) fn is_unknown_topic_error(msg: &str) -> bool {
498    msg.contains("UNKNOWN_TOPIC_OR_PARTITION")
499        || msg.contains("unknown topic")
500        || msg.contains("UnknownTopicOrPartition")
501        || msg.contains("NotSubscribed")
502}
503
504#[cfg(test)]
505mod tests {
506    use krabka_units::prelude::{TimeExt as _, secs};
507
508    #[test]
509    fn create_topics_timeout_reaches_the_wire_as_int32_millis() {
510        // Kafka's `CreateTopics.timeoutMs` is an int32 count of milliseconds, so
511        // the configured extent must cross the seam as 10000 — a seconds-valued
512        // 10 would have the broker give up almost immediately.
513        let policy = crate::config::ReplicatorRuntimePolicy::default();
514        assert2::assert!(policy.topic_create_timeout.millis_i32() == 10_000);
515    }
516
517    #[test]
518    fn drain_poll_timeout_is_half_a_second() {
519        let policy = crate::config::ReplicatorRuntimePolicy::default();
520        assert2::assert!(policy.internal_drain_poll_timeout == krabka_units::millis(500));
521    }
522
523    #[test]
524    fn unknown_topic_error_matches_each_substring() {
525        // Each positive exercises exactly one of the OR'd substrings, so the
526        // single-substring match must hold on its own (kills `||`→`&&`, which
527        // would require all four, and the `-> true`/`-> false` constants).
528        for (msg, want) in [
529            ("x UNKNOWN_TOPIC_OR_PARTITION y", true),
530            ("unknown topic foo", true),
531            ("UnknownTopicOrPartition", true),
532            ("err: NotSubscribed", true),
533            // A message with none of the substrings must not match (kills `-> true`).
534            ("connection refused", false),
535        ] {
536            assert2::assert!(super::is_unknown_topic_error(msg) == want);
537        }
538    }
539
540    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
541    async fn reading_a_missing_topic_preserves_explicit_partition_creation() {
542        let dir = tempfile::TempDir::new().unwrap();
543        let broker = krabka_broker::Broker::start(krabka_broker::BrokerConfig::for_tests(
544            dir.path().to_path_buf(),
545        ))
546        .await
547        .unwrap();
548        let bootstrap = broker.listen_addr().to_string();
549        let topic = "explicit-three-partitions";
550
551        let records = super::read_all(&bootstrap, topic, None).await.unwrap();
552        assert2::assert!(records.is_empty());
553
554        super::ensure_topic(&bootstrap, topic, 3, None)
555            .await
556            .unwrap();
557        let mut admin = krabka_client_admin::AdminClient::connect(&[bootstrap])
558            .await
559            .unwrap();
560        let metadata = admin.metadata(&[topic]).await.unwrap();
561        assert2::assert!(metadata.topics[0].partition_count == 3);
562    }
563
564    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
565    async fn ensure_produce_read_roundtrip() {
566        let dir = tempfile::TempDir::new().unwrap();
567        let broker = krabka_broker::Broker::start(krabka_broker::BrokerConfig::for_tests(
568            dir.path().to_path_buf(),
569        ))
570        .await
571        .unwrap();
572        let b = broker.listen_addr().to_string();
573
574        crate::test_util::create_topic(&b, "t", 1).await;
575        crate::test_util::produce(&b, "t", b"k1", b"v1").await;
576        crate::test_util::produce(&b, "t", b"k1", b"v2").await;
577
578        assert2::assert!(crate::test_util::topic_record_count(&b, "t").await == 2);
579
580        let last = super::read_last_value_for_key(&b, "t", b"k1", None)
581            .await
582            .unwrap();
583        assert2::assert!(last.as_deref() == Some(b"v2".as_slice()));
584
585        super::ensure_compacted_topic(&b, "state", None)
586            .await
587            .unwrap();
588
589        crate::test_util::await_topic_count(&b, "t", 2, secs(5)).await;
590    }
591}
krabka_replicator/admin_util.rs

krabka_replicator/
admin_util.rs
