krabka_connect_worker/
lib.rs

1#![doc = include_str!("../README.md")]
2
3mod health;
4mod metrics;
5
6pub mod config;
7pub mod kafka;
8
9use std::sync::Arc;
10
11use anyhow::Context as _;
12use bytes::Bytes;
13pub use config::WorkerConfig;
14pub use kafka::{CHECKPOINT_TOPIC, KafkaCheckpointStore, KafkaSink};
15use krabka_connect::{ConnectorRuntime, RuntimeState};
16use krabka_connect_postgres::PostgresWalSource;
17use tokio_util::sync::CancellationToken;
18
19use crate::{kafka::KafkaClientConfig, metrics::WorkerMetrics};
20
21/// Run one Postgres-to-Kafka connector until Ctrl-C and drain it gracefully.
22///
23/// # Errors
24///
25/// Returns an error for invalid configuration, startup failures, a connector
26/// runtime failure, health-server failure, or signal-listener failure.
27pub async fn run(config: WorkerConfig) -> anyhow::Result<()> {
28    config.validate().map_err(anyhow::Error::msg)?;
29    let security = config.client_security().map_err(anyhow::Error::msg)?;
30    let checkpoint_key = config.checkpoint_key();
31    let client = KafkaClientConfig {
32        bootstrap: config.kafka_bootstrap.clone(),
33        security,
34        replication_factor: config.replication_factor,
35        dispatch_queue_capacity: config.client_dispatch_queue_capacity,
36        frame_max_bytes: config.client_frame_max_bytes,
37    };
38    let metrics = WorkerMetrics::new();
39    let shutdown = CancellationToken::new();
40
41    let source = PostgresWalSource::connect(config.postgres_source())
42        .await
43        .context("start PostgreSQL source")?;
44    let sink =
45        KafkaSink::start_with_config(client.clone(), config.topic_prefix.clone(), metrics.clone())
46            .await
47            .context("connect Kafka sink")?;
48    let checkpoints = Arc::new(
49        KafkaCheckpointStore::start_with_config(client, checkpoint_key, metrics.clone())
50            .await
51            .context("connect Kafka checkpoint store")?,
52    );
53    metrics.set_live(true);
54    let mut health_task = health::start(
55        config.health_listen,
56        metrics.clone(),
57        shutdown.child_token(),
58    )
59    .await?;
60    let handle = ConnectorRuntime::<Bytes, Bytes>::new()
61        .add_source(source)
62        .add_sink(sink)
63        .checkpoint_store(checkpoints)
64        .max_batch(config.batch_size)
65        .commit_interval(config.commit_interval())
66        .poll_backoff(config.poll_backoff())
67        .run();
68
69    tracing::info!(connector.id = %config.connector_id, "connector worker running");
70    let mut terminal_error = None;
71    let mut tick = tokio::time::interval(std::time::Duration::from_millis(100));
72    loop {
73        let state = handle.state();
74        metrics.set_ready(matches!(
75            state,
76            RuntimeState::Running | RuntimeState::Paused
77        ));
78        if matches!(state, RuntimeState::Failed | RuntimeState::Stopped) {
79            terminal_error = Some(format!("connector runtime stopped in state {state:?}"));
80            break;
81        }
82        tokio::select! {
83            signal = shutdown_signal() => {
84                signal.context("listen for shutdown signal")?;
85                break;
86            }
87            _ = tick.tick() => {}
88            result = &mut health_task => {
89                terminal_error = Some(match result {
90                    Ok(Ok(())) => "health server stopped unexpectedly".to_owned(),
91                    Ok(Err(error)) => error.to_string(),
92                    Err(error) => format!("health server task failed: {error}"),
93                });
94                break;
95            }
96        }
97    }
98
99    metrics.set_ready(false);
100    if let Err(error) = handle.shutdown().await {
101        metrics.record_error();
102        terminal_error.get_or_insert_with(|| error.to_string());
103    }
104    metrics.set_live(false);
105    shutdown.cancel();
106    if !health_task.is_finished() {
107        health_task.await.context("join health server task")??;
108    }
109    if let Some(error) = terminal_error {
110        anyhow::bail!(error);
111    }
112    tracing::info!(connector.id = %config.connector_id, "connector worker stopped");
113    Ok(())
114}
115
116#[cfg(unix)]
117async fn shutdown_signal() -> std::io::Result<()> {
118    use tokio::signal::unix::{SignalKind, signal};
119
120    let mut terminate = signal(SignalKind::terminate())?;
121    tokio::select! {
122        result = tokio::signal::ctrl_c() => result,
123        _ = terminate.recv() => Ok(()),
124    }
125}
126
127#[cfg(not(unix))]
128async fn shutdown_signal() -> std::io::Result<()> {
129    tokio::signal::ctrl_c().await
130}
krabka_connect_worker/lib.rs

krabka_connect_worker/
lib.rs
