From 494179bf671df8e50b6c00ae8bd6f8f016abb34d Mon Sep 17 00:00:00 2001 From: MoonBoi9001 Date: Mon, 24 Aug 2026 16:56:07 +0300 Subject: [PATCH 1/4] feat(producer): vendor Studio's subgraph indexing request event schema Studio announces how many indexers a subgraph deployment wants over Redpanda; the dipper will consume those messages. This vendors Studio's protobuf schema, generates Rust bindings for it, and adds one field, the network the subgraph indexes, which the consumer needs and Studio has. --- dipper-producer/README.md | 12 +-- dipper-producer/build.rs | 8 +- .../subgraph-indexing-request-events.proto | 88 +++++++++++++++++++ ...r.subgraph.indexing.agreement.events.v1.rs | 2 +- dipper-producer/src/proto/mod.rs | 9 +- ...io.subgraph.indexing.requests.events.v1.rs | 67 ++++++++++++++ justfile | 8 +- 7 files changed, 182 insertions(+), 12 deletions(-) create mode 100644 dipper-producer/proto/subgraph-indexing-request-events.proto create mode 100644 dipper-producer/src/proto/studio.subgraph.indexing.requests.events.v1.rs diff --git a/dipper-producer/README.md b/dipper-producer/README.md index 32ee5100..6e6b5ccf 100644 --- a/dipper-producer/README.md +++ b/dipper-producer/README.md @@ -1,16 +1,18 @@ # dipper-producer -This crate contains protobuf definitions for dipper indexer agreement event streaming. -The generated Rust bindings are committed to the repository and only need to be regenerated when the `.proto` file changes. +This crate contains the protobuf definitions and Kafka plumbing for dipper event streaming: the producer side for the agreement lifecycle events the dipper emits, and the consumer side for the subgraph indexing request events Studio emits. The generated Rust bindings are committed to the repository and only need to be regenerated when a `.proto` file changes. ## Protobuf Generation -The build script uses a configuration flag `gen_event_proto` that enables protobuf code generation via `prost-build`. When enabled, the build script compiles `proto/indexing-agreement-events.proto` into Rust types under `src/proto/`. +The build script uses a configuration flag `gen_event_proto` that enables protobuf code generation via `prost-build`. When enabled, the build script compiles the schemas under `proto/` into Rust types under `src/proto/`: + +- `proto/indexing-agreement-events.proto`, owned by this repo, generates `src/proto/dipper.subgraph.indexing.agreement.events.v1.rs`. +- `proto/subgraph-indexing-request-events.proto`, vendored from the subgraph-studio repo (`packages/shared/src/helpers/dips/proto/SubgraphIndexingRequest.proto`), generates `src/proto/studio.subgraph.indexing.requests.events.v1.rs`. When Studio's copy changes, re-vendor it here and regenerate. To regenerate protobuf bindings, run: ```bash -just gen-indexing-agreement-events-proto +just gen-event-protos ``` Or using the full `cargo` command: @@ -18,5 +20,3 @@ Or using the full `cargo` command: ```bash RUSTFLAGS="--cfg gen_event_proto" cargo check -p dipper-producer ``` - -This will regenerate `src/proto/dipper.subgraph.indexing.agreement.events.v1.rs` from `proto/indexing-agreement-events.proto`. diff --git a/dipper-producer/build.rs b/dipper-producer/build.rs index c894d68b..c4c5f531 100644 --- a/dipper-producer/build.rs +++ b/dipper-producer/build.rs @@ -7,7 +7,13 @@ fn main() -> Result<(), Box> { config.out_dir("src/proto"); config.protoc_arg("--experimental_allow_proto3_optional"); - config.compile_protos(&["proto/indexing-agreement-events.proto"], &["proto/"])?; + config.compile_protos( + &[ + "proto/indexing-agreement-events.proto", + "proto/subgraph-indexing-request-events.proto", + ], + &["proto/"], + )?; // Instruct cargo to rerun this build script if any of the proto files change println!("cargo:rerun-if-changed=proto"); diff --git a/dipper-producer/proto/subgraph-indexing-request-events.proto b/dipper-producer/proto/subgraph-indexing-request-events.proto new file mode 100644 index 00000000..d9bc7233 --- /dev/null +++ b/dipper-producer/proto/subgraph-indexing-request-events.proto @@ -0,0 +1,88 @@ +// Subgraph Indexing Requests Events Protocol Buffer Schema +// +// Vendored copy of the schema owned by Subgraph Studio: +// repo: edgeandnode/subgraph-studio +// path: packages/shared/src/helpers/dips/proto/SubgraphIndexingRequest.proto +// commit: a8140613 +// +// Studio produces these messages on a Redpanda topic; the dipper consumes them +// and turns each propose event into a set_indexing_target_candidates call. +// When Studio's copy changes, re-vendor it here and regenerate the bindings +// (see the crate README). +// +// Local divergence from the vendored source, kept intentionally: +// - SubgraphIndexingRequestPropose.indexed_network_caip2id (field 2) is the +// network the subgraph indexes, which the dipper needs to key the indexing +// request and Studio has at every call site. It is a backward-compatible +// proto3 addition agreed as the preferred fix for that gap; drop this note +// once Studio's copy carries the field. +// +// Event flow: +// 1. subgraph.indexing.request.propose - Emitted when the developer sends a request to have their Subgraph indexed +// 2. subgraph.indexing.agreements.terminate - Emitted when the developer sends a request to terminate active Indexing agreements +// +// Partition key format: {the_graph_network_caip2id}/{subgraph_deployment_qm_hash}/{request/terminate} +// +// - the_graph_network_caip2id -> CAIP2 ID of The Graph network where the Subgraph is published +// - subgraph_deployment_qm_hash -> Qm hash of the Subgraph deployment requesting to be indexed +// Example: eip155:42161/QmTXzATwNfgGVukV1fX2T6xw9f6LAYRVWpsdXyRWzUR2H9/request + +syntax = "proto3"; + +package studio.subgraph.indexing.requests.events.v1; + +// SubgraphIndexingRequestEvent is the envelope that wraps all the Subgraph Indexing Request event types. +// +// Provides: +// - event_id -> unique identifier (uuid) +// - event_type -> event type discrimination +// - event_version -> versioning for schema evolution awareness to the consumer +// - timestamp -> when the event occurred, useful for ordering/debugging +// - subgraph_deployment_qm_hash -> Qm hash of the Subgraph deployment +// - the_graph_network_caip2id -> CAIP-2 chain id of The Graph protocol network the Subgraph Deployment is published to (e.g. eip155:42161) +// - payload -> the event payload, determined by the event_type +message SubgraphIndexingRequestEvent { + // Unique event identifier. + // Format: UUID v7 (time-ordered) for natural chronological sorting. + // Example: "01912345-6789-7abc-def0-123456789abc" + string event_id = 1; + + // Event type discriminator for routing and filtering. + // Values: "subgraph.indexing.request.propose", "subgraph.indexing.agreements.terminate" + string event_type = 2; + + // Schema version for forward compatibility. + // Consumers should handle unknown fields gracefully. + // Current version: "1.0" + string event_version = 3; + + // Event timestamp in RFC 3339 format. + // Example: "2024-01-15T10:30:00.123Z" + string timestamp = 4; + + // Qm hash of the Subgraph deployment with a submitted indexing agreement event. + string subgraph_deployment_qm_hash = 5; + + // CAIP-2 chain id of The Graph protocol network where the Subgraph was published to. + // Format: "eip155:{chain_id}". Examples: "eip155:42161" (arbitrum), "eip155:421614" (arbitrum-sepolia) + string the_graph_network_caip2id = 6; + + // Event payload - exactly one of the specific event types. + // Use event_type field to determine which payload is present. + oneof payload { + SubgraphIndexingRequestPropose subgraph_indexing_request_propose = 7; + } +} + +// SubgraphIndexingRequestPropose is emitted when the Subgraph developer initiates a request to have the Subgraph indexed +// +// Event type: subgraph.indexing.request.propose +message SubgraphIndexingRequestPropose { + // Number of requested Indexing agreements to find + int32 indexing_agreements_requested = 1; + + // CAIP-2 chain id of the network the Subgraph indexes (its data source), + // distinct from the_graph_network_caip2id on the envelope. + // Format: "eip155:{chain_id}". Example: "eip155:1" for a subgraph indexing Ethereum mainnet. + string indexed_network_caip2id = 2; +} diff --git a/dipper-producer/src/proto/dipper.subgraph.indexing.agreement.events.v1.rs b/dipper-producer/src/proto/dipper.subgraph.indexing.agreement.events.v1.rs index be209e64..100b909d 100644 --- a/dipper-producer/src/proto/dipper.subgraph.indexing.agreement.events.v1.rs +++ b/dipper-producer/src/proto/dipper.subgraph.indexing.agreement.events.v1.rs @@ -178,7 +178,7 @@ pub struct SubgraphIndexingAgreementTerminated { /// the one just terminated. /// 0 -> the Subgraph no longer has any active, accepted indexing agreements. /// -1 -> the count was unavailable when the event was emitted (e.g. a transient - /// datastore error); treat as unknown, not as zero. + /// datastore error); treat as unknown, not as zero. #[prost(int32, tag = "5")] pub remaining_accepted_indexing_agreements: i32, } diff --git a/dipper-producer/src/proto/mod.rs b/dipper-producer/src/proto/mod.rs index 75c540c5..7709ba0c 100644 --- a/dipper-producer/src/proto/mod.rs +++ b/dipper-producer/src/proto/mod.rs @@ -1,4 +1,11 @@ // Generated protobuf types for indexing agreement events // -// To regenerate, run: just gen-indexing-agreement-events-proto +// To regenerate, run: just gen-event-protos include!("dipper.subgraph.indexing.agreement.events.v1.rs"); + +/// Subgraph indexing request events that Studio produces and the dipper consumes. +/// Schema vendored in `proto/subgraph-indexing-request-events.proto`. +/// To regenerate, run: just gen-event-protos +pub mod studio { + include!("studio.subgraph.indexing.requests.events.v1.rs"); +} diff --git a/dipper-producer/src/proto/studio.subgraph.indexing.requests.events.v1.rs b/dipper-producer/src/proto/studio.subgraph.indexing.requests.events.v1.rs new file mode 100644 index 00000000..760abe2a --- /dev/null +++ b/dipper-producer/src/proto/studio.subgraph.indexing.requests.events.v1.rs @@ -0,0 +1,67 @@ +// This file is @generated by prost-build. +/// SubgraphIndexingRequestEvent is the envelope that wraps all the Subgraph Indexing Request event types. +/// +/// Provides: +/// - event_id -> unique identifier (uuid) +/// - event_type -> event type discrimination +/// - event_version -> versioning for schema evolution awareness to the consumer +/// - timestamp -> when the event occurred, useful for ordering/debugging +/// - subgraph_deployment_qm_hash -> Qm hash of the Subgraph deployment +/// - the_graph_network_caip2id -> CAIP-2 chain id of The Graph protocol network the Subgraph Deployment is published to (e.g. eip155:42161) +/// - payload -> the event payload, determined by the event_type +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] +pub struct SubgraphIndexingRequestEvent { + /// Unique event identifier. + /// Format: UUID v7 (time-ordered) for natural chronological sorting. + /// Example: "01912345-6789-7abc-def0-123456789abc" + #[prost(string, tag = "1")] + pub event_id: ::prost::alloc::string::String, + /// Event type discriminator for routing and filtering. + /// Values: "subgraph.indexing.request.propose", "subgraph.indexing.agreements.terminate" + #[prost(string, tag = "2")] + pub event_type: ::prost::alloc::string::String, + /// Schema version for forward compatibility. + /// Consumers should handle unknown fields gracefully. + /// Current version: "1.0" + #[prost(string, tag = "3")] + pub event_version: ::prost::alloc::string::String, + /// Event timestamp in RFC 3339 format. + /// Example: "2024-01-15T10:30:00.123Z" + #[prost(string, tag = "4")] + pub timestamp: ::prost::alloc::string::String, + /// Qm hash of the Subgraph deployment with a submitted indexing agreement event. + #[prost(string, tag = "5")] + pub subgraph_deployment_qm_hash: ::prost::alloc::string::String, + /// CAIP-2 chain id of The Graph protocol network where the Subgraph was published to. + /// Format: "eip155:{chain_id}". Examples: "eip155:42161" (arbitrum), "eip155:421614" (arbitrum-sepolia) + #[prost(string, tag = "6")] + pub the_graph_network_caip2id: ::prost::alloc::string::String, + /// Event payload - exactly one of the specific event types. + /// Use event_type field to determine which payload is present. + #[prost(oneof = "subgraph_indexing_request_event::Payload", tags = "7")] + pub payload: ::core::option::Option, +} +/// Nested message and enum types in `SubgraphIndexingRequestEvent`. +pub mod subgraph_indexing_request_event { + /// Event payload - exactly one of the specific event types. + /// Use event_type field to determine which payload is present. + #[derive(Clone, PartialEq, Eq, Hash, ::prost::Oneof)] + pub enum Payload { + #[prost(message, tag = "7")] + SubgraphIndexingRequestPropose(super::SubgraphIndexingRequestPropose), + } +} +/// SubgraphIndexingRequestPropose is emitted when the Subgraph developer initiates a request to have the Subgraph indexed +/// +/// Event type: subgraph.indexing.request.propose +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] +pub struct SubgraphIndexingRequestPropose { + /// Number of requested Indexing agreements to find + #[prost(int32, tag = "1")] + pub indexing_agreements_requested: i32, + /// CAIP-2 chain id of the network the Subgraph indexes (its data source), + /// distinct from the_graph_network_caip2id on the envelope. + /// Format: "eip155:{chain_id}". Example: "eip155:1" for a subgraph indexing Ethereum mainnet. + #[prost(string, tag = "2")] + pub indexed_network_caip2id: ::prost::alloc::string::String, +} diff --git a/justfile b/justfile index 70778005..92220050 100644 --- a/justfile +++ b/justfile @@ -113,7 +113,9 @@ remove-git-hooks: # Remove the pre-commit hooks pre-commit uninstall --config .github/pre-commit-config.yaml -# Generate job events protobuf bindings (RUSTFLAGS="--cfg gen_event_proto" cargo check) +# Generate event protobuf bindings (RUSTFLAGS="--cfg gen_event_proto" cargo check) [group: 'codegen'] -gen-indexing-agreement-events-proto: - RUSTFLAGS="--cfg gen_event_proto" cargo check -p dipper-producer \ No newline at end of file +gen-event-protos: + RUSTFLAGS="--cfg gen_event_proto" cargo check -p dipper-producer + +alias gen-indexing-agreement-events-proto := gen-event-protos \ No newline at end of file From 6114c84cf4f67d56dca1882bb1b65eae3e3e2864 Mon Sep 17 00:00:00 2001 From: MoonBoi9001 Date: Mon, 24 Aug 2026 17:00:37 +0300 Subject: [PATCH 2/4] feat(producer): add a Kafka consumer for Studio's indexing requests The dipper could only send Kafka events, not read the indexing requests Studio publishes. The new consumer finds the topic's partitions from broker metadata, refuses to start if the topic is missing, and leaves offset tracking to the caller; SASL/TLS setup is now shared. --- Cargo.lock | 1 + dipper-producer/Cargo.toml | 3 + dipper-producer/src/kafka.rs | 30 +-- dipper-producer/src/kafka/connection.rs | 216 ++++++++++++++++++ dipper-producer/src/kafka/consumer.rs | 287 ++++++++++++++++++++++++ dipper-producer/src/kafka/producer.rs | 240 ++------------------ 6 files changed, 534 insertions(+), 243 deletions(-) create mode 100644 dipper-producer/src/kafka/connection.rs create mode 100644 dipper-producer/src/kafka/consumer.rs diff --git a/Cargo.lock b/Cargo.lock index 2ab2aa3a..b53ab07e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2456,6 +2456,7 @@ dependencies = [ "rustls", "rustls-pemfile 2.2.0", "serde", + "serde_json", "thegraph-core", "thiserror 2.0.18", "tokio", diff --git a/dipper-producer/Cargo.toml b/dipper-producer/Cargo.toml index 489f170b..c8a5432f 100644 --- a/dipper-producer/Cargo.toml +++ b/dipper-producer/Cargo.toml @@ -20,6 +20,9 @@ tracing.workspace = true uuid = { workspace = true, features = ["v7"] } webpki-roots = "1.0.8" +[dev-dependencies] +serde_json.workspace = true + # Protobuf code generation dependencies # Run with: RUSTFLAGS="--cfg gen_event_proto" cargo check -p dipper-producer [target.'cfg(gen_event_proto)'.build-dependencies] diff --git a/dipper-producer/src/kafka.rs b/dipper-producer/src/kafka.rs index 25714154..a6dfe341 100644 --- a/dipper-producer/src/kafka.rs +++ b/dipper-producer/src/kafka.rs @@ -1,27 +1,11 @@ -//! Kafka client for subgraph indexing agreement event streaming. -//! -//! This module provides a Kafka producer for emitting subgraph indexing agreement lifecycle events -//! to a kafka topic -//! -//! Events are encoded using Protocol Buffers for compact, schema-enforced messages. -//! -//! # Example -//! -//! ```ignore -//! use dipper_producer::kafka::{KafkaConfig, KafkaProducer}; -//! -//! let config = KafkaConfig { -//! brokers: vec!["localhost:9092".to_string()], -//! topic: "dipper.subgraph.indexing.agreement.events".to_string(), -//! partitions: 16, -//! }; -//! -//! let producer = KafkaProducer::new(&config).await?; -//! -//! // Send an event with partition key and protobuf payload -//! producer.send("QmT329Bej8AwSLahmgnmi6fdYkj3rorYAcCes45gDv9aJ4", &encoded_event).await?; -//! ``` +//! Kafka clients for dipper event streaming: a producer for the agreement +//! lifecycle events the dipper emits and a consumer for the indexing request +//! events Studio emits. Events are Protocol Buffers encoded. +mod connection; +mod consumer; mod producer; +pub use connection::ConnectionError; +pub use consumer::{ConsumerError, KafkaConsumer, KafkaConsumerConfig}; pub use producer::{Error, KafkaConfig, KafkaProducer}; diff --git a/dipper-producer/src/kafka/connection.rs b/dipper-producer/src/kafka/connection.rs new file mode 100644 index 00000000..caf89359 --- /dev/null +++ b/dipper-producer/src/kafka/connection.rs @@ -0,0 +1,216 @@ +//! Broker connection helpers shared by the Kafka producer and consumer: +//! SASL/TLS setup and the client bootstrap built from them. + +use std::{ + path::Path, + sync::{Arc, Once}, +}; + +use rskafka::client::{Client, ClientBuilder, Credentials, SaslConfig}; +use rustls::ClientConfig; + +static RUSTLS_CRYPTO_PROVIDER: Once = Once::new(); + +/// Broker connection parameters, borrowed from the producer or consumer config. +pub(crate) struct ConnectOptions<'a> { + pub brokers: &'a [String], + pub sasl_mechanism: Option<&'a str>, + pub sasl_username: Option<&'a str>, + pub sasl_password: Option<&'a str>, + pub tls_enabled: bool, + pub tls_ca_cert_path: Option<&'a Path>, +} + +/// Connects a Kafka client with the given SASL/TLS settings. +pub(crate) async fn connect(opts: ConnectOptions<'_>) -> Result { + let mut builder = ClientBuilder::new(opts.brokers.to_vec()); + + if let Some(mechanism_str) = opts.sasl_mechanism { + let mechanism: SaslMechanism = mechanism_str.parse()?; + let sasl_config = build_sasl_config(mechanism, opts.sasl_username, opts.sasl_password)?; + builder = builder.sasl_config(sasl_config); + } + + if opts.tls_enabled { + let tls_config = build_tls_config(opts.tls_ca_cert_path)?; + builder = builder.tls_config(tls_config); + } + + builder.build().await.map_err(ConnectionError::Connection) +} + +/// Builds SASL configuration from the provided mechanism and credentials. +pub(crate) fn build_sasl_config( + mechanism: SaslMechanism, + username: Option<&str>, + password: Option<&str>, +) -> Result { + let username = username.ok_or(ConnectionError::MissingSaslUsername)?; + let password = password.ok_or(ConnectionError::MissingSaslPassword)?; + + let credentials = Credentials::new(username.to_string(), password.to_string()); + + Ok(match mechanism { + SaslMechanism::Plain => SaslConfig::Plain(credentials), + SaslMechanism::ScramSha256 => SaslConfig::ScramSha256(credentials), + SaslMechanism::ScramSha512 => SaslConfig::ScramSha512(credentials), + }) +} + +/// Builds TLS configuration. A custom CA certificate path makes the client +/// trust that CA for broker verification; otherwise system roots are used. +pub(crate) fn build_tls_config( + ca_cert_path: Option<&Path>, +) -> Result, ConnectionError> { + install_rustls_crypto_provider(); + + let root_store = match ca_cert_path { + Some(path) => { + let ca_pem = + fs_err::read(path).map_err(|e| ConnectionError::TlsCaCert { source: e })?; + let mut reader = std::io::BufReader::new(&ca_pem[..]); + let certs: Vec<_> = rustls_pemfile::certs(&mut reader) + .collect::>() + .map_err(|e| ConnectionError::TlsCaCert { source: e })?; + + let mut store = rustls::RootCertStore::empty(); + for cert in certs { + store.add(cert).map_err(|e| ConnectionError::TlsCaCert { + source: std::io::Error::new(std::io::ErrorKind::InvalidData, e), + })?; + } + store + } + None => rustls::RootCertStore { + roots: webpki_roots::TLS_SERVER_ROOTS.to_vec(), + }, + }; + + let tls_config = ClientConfig::builder() + .with_root_certificates(root_store) + .with_no_client_auth(); + + Ok(Arc::new(tls_config)) +} + +fn install_rustls_crypto_provider() { + RUSTLS_CRYPTO_PROVIDER.call_once(|| { + // Necessary for the Kafka client: it builds a Rustls TLS config directly, + // so install a provider before `ClientConfig::builder()` tries to infer one. + let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); + }); +} + +/// Errors that can occur while establishing a broker connection. +#[derive(Debug, thiserror::Error)] +pub enum ConnectionError { + /// Failed to connect to Kafka brokers + #[error("failed to connect to Kafka brokers")] + Connection(#[source] rskafka::client::error::Error), + + /// Unsupported SASL mechanism + #[error("unsupported SASL mechanism '{0}', supported: PLAIN, SCRAM-SHA-256, SCRAM-SHA-512")] + UnsupportedSaslMechanism(String), + + /// Missing SASL username + #[error("sasl_username is required when sasl_mechanism is set")] + MissingSaslUsername, + + /// Missing SASL password + #[error("sasl_password is required when sasl_mechanism is set")] + MissingSaslPassword, + + /// Failed to load TLS CA certificate + #[error("failed to load TLS CA certificate")] + TlsCaCert { + #[source] + source: std::io::Error, + }, +} + +/// Supported SASL authentication mechanisms. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum SaslMechanism { + Plain, + ScramSha256, + ScramSha512, +} + +impl std::str::FromStr for SaslMechanism { + type Err = ConnectionError; + + fn from_str(s: &str) -> Result { + match s.to_uppercase().as_str() { + "PLAIN" => Ok(Self::Plain), + "SCRAM-SHA-256" => Ok(Self::ScramSha256), + "SCRAM-SHA-512" => Ok(Self::ScramSha512), + _ => Err(ConnectionError::UnsupportedSaslMechanism(s.to_string())), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn should_parse_sasl_mechanism_correctly() { + // Case-insensitive parsing + assert_eq!( + "PLAIN".parse::().unwrap(), + SaslMechanism::Plain + ); + assert_eq!( + "plain".parse::().unwrap(), + SaslMechanism::Plain + ); + assert_eq!( + "SCRAM-SHA-256".parse::().unwrap(), + SaslMechanism::ScramSha256 + ); + assert_eq!( + "scram-sha-256".parse::().unwrap(), + SaslMechanism::ScramSha256 + ); + assert_eq!( + "SCRAM-SHA-512".parse::().unwrap(), + SaslMechanism::ScramSha512 + ); + + // Unsupported mechanism + assert!("GSSAPI".parse::().is_err()); + } + + #[test] + fn should_build_sasl_plain_sasl_config() { + let result = build_sasl_config(SaslMechanism::Plain, Some("user"), Some("pass")); + assert!(matches!(result, Ok(SaslConfig::Plain(_)))); + } + + #[test] + fn should_build_scram_sha_256_sasl_config() { + let result = build_sasl_config(SaslMechanism::ScramSha256, Some("user"), Some("pass")); + assert!(matches!(result, Ok(SaslConfig::ScramSha256(_)))); + } + + #[test] + fn should_build_sasl_sha_512_sasl_config() { + let result = build_sasl_config(SaslMechanism::ScramSha512, Some("user"), Some("pass")); + assert!(matches!(result, Ok(SaslConfig::ScramSha512(_)))); + } + + #[test] + fn should_throw_err_on_missing_credentials() { + // Missing username + assert!(matches!( + build_sasl_config(SaslMechanism::Plain, None, Some("pass")), + Err(ConnectionError::MissingSaslUsername) + )); + + // Missing password + assert!(matches!( + build_sasl_config(SaslMechanism::Plain, Some("user"), None), + Err(ConnectionError::MissingSaslPassword) + )); + } +} diff --git a/dipper-producer/src/kafka/consumer.rs b/dipper-producer/src/kafka/consumer.rs new file mode 100644 index 00000000..eea727ca --- /dev/null +++ b/dipper-producer/src/kafka/consumer.rs @@ -0,0 +1,287 @@ +//! Kafka consumer for the subgraph indexing request events Studio produces. +//! Deliberately a thin fetch layer: rskafka has no consumer groups, so offset +//! tracking belongs to the caller (the dipper persists offsets in its own DB). + +use std::{collections::BTreeMap, path::PathBuf, sync::Arc, time::Duration}; + +use rskafka::{ + client::partition::{OffsetAt, PartitionClient, UnknownTopicHandling}, + record::RecordAndOffset, +}; + +use super::connection::{self, ConnectOptions, ConnectionError}; + +/// Kafka consumer configuration. +#[derive(Clone, serde::Deserialize)] +#[serde(deny_unknown_fields)] +pub struct KafkaConsumerConfig { + /// Kafka broker addresses. + pub brokers: Vec, + /// Kafka topic to consume. Deliberately has no default: a consumer pointed + /// at a missing or misnamed topic must fail at startup, not idle on nothing. + pub topic: String, + /// SASL authentication mechanism (e.g., "PLAIN", "SCRAM-SHA-256", "SCRAM-SHA-512"). + #[serde(default)] + pub sasl_mechanism: Option, + /// SASL username. + #[serde(default)] + pub sasl_username: Option, + /// SASL password. + #[serde(default)] + pub sasl_password: Option, + /// Enable TLS encryption. + #[serde(default)] + pub tls_enabled: bool, + /// Path to a PEM-encoded CA certificate file for TLS verification. + #[serde(default)] + pub tls_ca_cert_path: Option, +} + +// Manual impl instead of derive: the service logs the whole config with Debug +// formatting at startup, so the SASL password must never reach the output. +impl std::fmt::Debug for KafkaConsumerConfig { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("KafkaConsumerConfig") + .field("brokers", &self.brokers) + .field("topic", &self.topic) + .field("sasl_mechanism", &self.sasl_mechanism) + .field("sasl_username", &self.sasl_username) + .field( + "sasl_password", + &self.sasl_password.as_ref().map(|_| ""), + ) + .field("tls_enabled", &self.tls_enabled) + .field("tls_ca_cert_path", &self.tls_ca_cert_path) + .finish() + } +} + +/// Kafka consumer bound to a single topic, with a partition client per +/// discovered partition. Thread-safe; share across tasks via `Arc`. +pub struct KafkaConsumer { + topic: String, + partition_clients: BTreeMap>, +} + +impl KafkaConsumer { + /// Margin added on top of a fetch's `max_wait_ms` before the whole call is + /// abandoned, covering broker round-trip and retry time. + const FETCH_TIMEOUT_MARGIN: Duration = Duration::from_secs(30); + + /// Timeout for offset queries, which carry no server-side wait. + const OFFSET_TIMEOUT: Duration = Duration::from_secs(30); + + /// Connects to the brokers and binds to the configured topic. Partitions + /// are discovered from broker metadata, so there is no partition count to + /// configure; a topic the credentials cannot see is an error. + pub async fn connect(config: &KafkaConsumerConfig) -> Result { + let client = connection::connect(ConnectOptions { + brokers: &config.brokers, + sasl_mechanism: config.sasl_mechanism.as_deref(), + sasl_username: config.sasl_username.as_deref(), + sasl_password: config.sasl_password.as_deref(), + tls_enabled: config.tls_enabled, + tls_ca_cert_path: config.tls_ca_cert_path.as_deref(), + }) + .await?; + + let topics = client + .list_topics() + .await + .map_err(ConsumerError::Metadata)?; + let topic = topics + .into_iter() + .find(|t| t.name == config.topic) + .ok_or_else(|| ConsumerError::TopicNotFound { + topic: config.topic.clone(), + })?; + + let mut partition_clients = BTreeMap::new(); + for partition in topic.partitions { + let partition_client = client + .partition_client(&config.topic, partition, UnknownTopicHandling::Error) + .await + .map_err(ConsumerError::PartitionClient)?; + partition_clients.insert(partition, Arc::new(partition_client)); + } + + Ok(Self { + topic: config.topic.clone(), + partition_clients, + }) + } + + /// The topic this consumer is bound to. + pub fn topic(&self) -> &str { + &self.topic + } + + /// The partition ids discovered for the topic, in ascending order. + pub fn partitions(&self) -> Vec { + self.partition_clients.keys().copied().collect() + } + + /// Fetches records from one partition starting at `offset`, waiting up to + /// `max_wait_ms` for data to arrive. Returns the records (with their + /// offsets) and the partition's current high watermark. + pub async fn fetch( + &self, + partition: i32, + offset: i64, + max_bytes: i32, + max_wait_ms: i32, + ) -> Result<(Vec, i64), ConsumerError> { + let partition_client = self.partition_client(partition)?; + let timeout = Duration::from_millis(max_wait_ms.max(0) as u64) + Self::FETCH_TIMEOUT_MARGIN; + + tokio::time::timeout( + timeout, + partition_client.fetch_records(offset, 1..max_bytes, max_wait_ms), + ) + .await + .map_err(|_| ConsumerError::Timeout)? + .map_err(ConsumerError::Fetch) + } + + /// Queries one partition's earliest or latest offset. + pub async fn offset(&self, partition: i32, at: OffsetAt) -> Result { + let partition_client = self.partition_client(partition)?; + + tokio::time::timeout(Self::OFFSET_TIMEOUT, partition_client.get_offset(at)) + .await + .map_err(|_| ConsumerError::Timeout)? + .map_err(ConsumerError::Offset) + } + + fn partition_client(&self, partition: i32) -> Result<&Arc, ConsumerError> { + self.partition_clients + .get(&partition) + .ok_or(ConsumerError::UnknownPartition { + partition, + topic: self.topic.clone(), + }) + } +} + +/// Errors that can occur when working with the Kafka consumer. +#[derive(Debug, thiserror::Error)] +pub enum ConsumerError { + /// Failed to establish the broker connection (SASL, TLS, or bootstrap) + #[error(transparent)] + Connection(#[from] ConnectionError), + + /// Failed to list topics from the broker + #[error("failed to list topics from the broker")] + Metadata(#[source] rskafka::client::error::Error), + + /// The configured topic does not exist (or is not visible to the credentials) + #[error("topic '{topic}' does not exist on the broker or is not visible to the credentials")] + TopicNotFound { topic: String }, + + /// Failed to get partition client + #[error("failed to get partition client")] + PartitionClient(#[source] rskafka::client::error::Error), + + /// The requested partition is not part of the bound topic + #[error("partition {partition} is not part of topic '{topic}'")] + UnknownPartition { partition: i32, topic: String }, + + /// Failed to fetch records from Kafka + #[error("failed to fetch records from Kafka")] + Fetch(#[source] rskafka::client::error::Error), + + /// Failed to query a partition offset + #[error("failed to query partition offset")] + Offset(#[source] rskafka::client::error::Error), + + /// Kafka operation timed out + #[error("Kafka operation timed out")] + Timeout, +} + +impl ConsumerError { + /// Whether this is the broker refusing a fetch offset outside its retained + /// range. Callers recover by re-anchoring to the earliest available offset + /// (records below it were deleted by retention). + pub fn is_offset_out_of_range(&self) -> bool { + matches!( + self, + Self::Fetch(rskafka::client::error::Error::ServerError { + protocol_error: rskafka::client::error::ProtocolError::OffsetOutOfRange, + .. + }) + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn config_deserializes_with_topic_and_defaults() { + let json = r#"{ + "brokers": ["localhost:9092"], + "topic": "studio.subgraph.indexing.requests" + }"#; + + let config: KafkaConsumerConfig = serde_json::from_str(json).expect("valid config"); + assert_eq!(config.brokers, vec!["localhost:9092".to_string()]); + assert_eq!(config.topic, "studio.subgraph.indexing.requests"); + assert_eq!(config.sasl_mechanism, None); + assert_eq!(config.sasl_username, None); + assert_eq!(config.sasl_password, None); + assert!(!config.tls_enabled); + assert_eq!(config.tls_ca_cert_path, None); + } + + #[test] + fn config_without_topic_is_rejected() { + // No default topic on purpose: a missing name must fail configuration + // loading, not silently consume from nowhere. + let json = r#"{ "brokers": ["localhost:9092"] }"#; + + let err = serde_json::from_str::(json).unwrap_err(); + assert!( + err.to_string().contains("topic"), + "error should name the missing field: {err}" + ); + } + + #[test] + fn config_with_unknown_fields_is_rejected() { + let json = r#"{ + "brokers": ["localhost:9092"], + "topic": "t", + "partitions": 16 + }"#; + + let err = serde_json::from_str::(json).unwrap_err(); + assert!( + err.to_string().contains("partitions"), + "error should name the unknown field: {err}" + ); + } + + #[test] + fn debug_output_redacts_the_sasl_password() { + let config = KafkaConsumerConfig { + brokers: vec!["localhost:9092".to_string()], + topic: "test".to_string(), + sasl_mechanism: Some("PLAIN".to_string()), + sasl_username: Some("user".to_string()), + sasl_password: Some("hunter2".to_string()), + tls_enabled: false, + tls_ca_cert_path: None, + }; + let rendered = format!("{config:?}"); + assert!( + !rendered.contains("hunter2"), + "debug output must not contain the password: {rendered}" + ); + assert!( + rendered.contains(""), + "debug output should mark the password as redacted: {rendered}" + ); + } +} diff --git a/dipper-producer/src/kafka/producer.rs b/dipper-producer/src/kafka/producer.rs index 07ddecbd..d2f78226 100644 --- a/dipper-producer/src/kafka/producer.rs +++ b/dipper-producer/src/kafka/producer.rs @@ -1,21 +1,13 @@ //! Kafka producer for sending dipper events on a configured topic -use std::{ - path::{Path, PathBuf}, - sync::{Arc, Once}, - time::Duration, -}; +use std::{path::PathBuf, sync::Arc, time::Duration}; use rskafka::{ - client::{ - ClientBuilder, Credentials, SaslConfig, - partition::{Compression, PartitionClient, UnknownTopicHandling}, - }, + client::partition::{Compression, PartitionClient, UnknownTopicHandling}, record::Record, }; -use rustls::ClientConfig; -static RUSTLS_CRYPTO_PROVIDER: Once = Once::new(); +use super::connection::{self, ConnectOptions, ConnectionError}; /// Kafka producer configuration. #[derive(Clone, serde::Deserialize)] @@ -95,22 +87,15 @@ impl KafkaProducer { return Err(Error::InvalidPartitionCount); } - let mut builder = ClientBuilder::new(config.brokers.clone()); - - // Configure SASL authentication if mechanism is specified - if let Some(mechanism_str) = &config.sasl_mechanism { - let mechanism: SaslMechanism = mechanism_str.parse()?; - let sasl_config = Self::build_sasl_config(mechanism, config)?; - builder = builder.sasl_config(sasl_config); - } - - // Configure TLS if enabled - if config.tls_enabled { - let tls_config = Self::build_tls_config(config.tls_ca_cert_path.as_deref())?; - builder = builder.tls_config(tls_config); - } - - let client = builder.build().await.map_err(Error::Connection)?; + let client = connection::connect(ConnectOptions { + brokers: &config.brokers, + sasl_mechanism: config.sasl_mechanism.as_deref(), + sasl_username: config.sasl_username.as_deref(), + sasl_password: config.sasl_password.as_deref(), + tls_enabled: config.tls_enabled, + tls_ca_cert_path: config.tls_ca_cert_path.as_deref(), + }) + .await?; let client = Arc::new(client); let mut partition_clients = Vec::with_capacity(config.partitions as usize); @@ -130,68 +115,8 @@ impl KafkaProducer { }) } - /// Builds SASL configuration from the provided mechanism and credentials. - fn build_sasl_config( - mechanism: SaslMechanism, - config: &KafkaConfig, - ) -> Result { - let username = config - .sasl_username - .clone() - .ok_or(Error::MissingSaslUsername)?; - let password = config - .sasl_password - .clone() - .ok_or(Error::MissingSaslPassword)?; - - let credentials = Credentials::new(username, password); - - Ok(match mechanism { - SaslMechanism::Plain => SaslConfig::Plain(credentials), - SaslMechanism::ScramSha256 => SaslConfig::ScramSha256(credentials), - SaslMechanism::ScramSha512 => SaslConfig::ScramSha512(credentials), - }) - } - - /// Builds TLS configuration. - /// - /// If a custom CA certificate path is provided, the client will trust that CA - /// for verifying broker connections. Otherwise, system root certificates are used. - fn build_tls_config(ca_cert_path: Option<&Path>) -> Result, Error> { - install_rustls_crypto_provider(); - - let root_store = match ca_cert_path { - Some(path) => { - let ca_pem = fs_err::read(path).map_err(|e| Error::TlsCaCert { source: e })?; - let mut reader = std::io::BufReader::new(&ca_pem[..]); - let certs: Vec<_> = rustls_pemfile::certs(&mut reader) - .collect::>() - .map_err(|e| Error::TlsCaCert { source: e })?; - - let mut store = rustls::RootCertStore::empty(); - for cert in certs { - store.add(cert).map_err(|e| Error::TlsCaCert { - source: std::io::Error::new(std::io::ErrorKind::InvalidData, e), - })?; - } - store - } - None => rustls::RootCertStore { - roots: webpki_roots::TLS_SERVER_ROOTS.to_vec(), - }, - }; - - let tls_config = ClientConfig::builder() - .with_root_certificates(root_store) - .with_no_client_auth(); - - Ok(Arc::new(tls_config)) - } - - /// Sends an event to Kafka. - /// - /// Events are partitioned by the partition key (table discriminator) before being written to - /// Kafka. The produce attempt times out after 30 seconds. + /// Sends an event to Kafka, partitioned by the partition key (table + /// discriminator). The produce attempt times out after 30 seconds. pub async fn send(&self, partition_key: &str, payload: &[u8]) -> Result<(), Error> { let partition = self.partition_for_key(partition_key); let partition_client = &self.partition_clients[partition as usize]; @@ -213,12 +138,9 @@ impl KafkaProducer { .map(|_| ()) } - /// Computes the partition for a given key. - /// - /// Uses a deterministic FNV-1a hash modulo the partition count so that a given - /// key always maps to the same partition across restarts and instances, - /// preserving per-key ordering. The partition count is configured via - /// `KafkaConfig::partitions`. + /// Computes the partition for a key: deterministic FNV-1a hash modulo + /// `KafkaConfig::partitions`, so a key maps to the same partition across + /// restarts and instances, preserving per-key ordering. fn partition_for_key(&self, key: &str) -> i32 { // FNV-1a (32-bit): order-dependent and well-distributed, unlike a byte sum. const FNV_OFFSET_BASIS: u32 = 0x811c_9dc5; @@ -231,20 +153,12 @@ impl KafkaProducer { } } -fn install_rustls_crypto_provider() { - RUSTLS_CRYPTO_PROVIDER.call_once(|| { - // Necessary for the Kafka client: it builds a Rustls TLS config directly, - // so install a provider before `ClientConfig::builder()` tries to infer one. - let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); - }); -} - /// Errors that can occur when working with the Kafka producer. #[derive(Debug, thiserror::Error)] pub enum Error { - /// Failed to connect to Kafka brokers - #[error("failed to connect to Kafka brokers")] - Connection(#[source] rskafka::client::error::Error), + /// Failed to establish the broker connection (SASL, TLS, or bootstrap) + #[error(transparent)] + Connection(#[from] ConnectionError), /// Failed to get partition client #[error("failed to get partition client")] @@ -261,126 +175,12 @@ pub enum Error { /// Partition count must be greater than zero #[error("partitions must be greater than zero")] InvalidPartitionCount, - - /// Unsupported SASL mechanism - #[error("unsupported SASL mechanism '{0}', supported: PLAIN, SCRAM-SHA-256, SCRAM-SHA-512")] - UnsupportedSaslMechanism(String), - - /// Missing SASL username - #[error("sasl_username is required when sasl_mechanism is set")] - MissingSaslUsername, - - /// Missing SASL password - #[error("sasl_password is required when sasl_mechanism is set")] - MissingSaslPassword, - - /// Failed to load TLS CA certificate - #[error("failed to load TLS CA certificate")] - TlsCaCert { - #[source] - source: std::io::Error, - }, -} - -/// Supported SASL authentication mechanisms. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum SaslMechanism { - Plain, - ScramSha256, - ScramSha512, -} - -impl std::str::FromStr for SaslMechanism { - type Err = Error; - - fn from_str(s: &str) -> Result { - match s.to_uppercase().as_str() { - "PLAIN" => Ok(Self::Plain), - "SCRAM-SHA-256" => Ok(Self::ScramSha256), - "SCRAM-SHA-512" => Ok(Self::ScramSha512), - _ => Err(Error::UnsupportedSaslMechanism(s.to_string())), - } - } } #[cfg(test)] mod tests { use super::*; - #[test] - fn should_parse_sasl_mechanism_correctly() { - // Case-insensitive parsing - assert_eq!( - "PLAIN".parse::().unwrap(), - SaslMechanism::Plain - ); - assert_eq!( - "plain".parse::().unwrap(), - SaslMechanism::Plain - ); - assert_eq!( - "SCRAM-SHA-256".parse::().unwrap(), - SaslMechanism::ScramSha256 - ); - assert_eq!( - "scram-sha-256".parse::().unwrap(), - SaslMechanism::ScramSha256 - ); - assert_eq!( - "SCRAM-SHA-512".parse::().unwrap(), - SaslMechanism::ScramSha512 - ); - - // Unsupported mechanism - assert!("GSSAPI".parse::().is_err()); - } - - #[test] - fn should_build_sasl_plain_sasl_config() { - let config = make_kafka_config(Some("PLAIN"), Some("user".into()), Some("pass".into())); - let result = KafkaProducer::build_sasl_config(SaslMechanism::Plain, &config); - assert!(matches!(result, Ok(SaslConfig::Plain(_)))); - } - - #[test] - fn should_build_scram_sha_256_sasl_config() { - let config = make_kafka_config( - Some("SCRAM-SHA-256"), - Some("user".into()), - Some("pass".into()), - ); - let result = KafkaProducer::build_sasl_config(SaslMechanism::ScramSha256, &config); - assert!(matches!(result, Ok(SaslConfig::ScramSha256(_)))); - } - - #[test] - fn should_build_sasl_sha_512_sasl_config() { - let config = make_kafka_config( - Some("SCRAM-SHA-512"), - Some("user".into()), - Some("pass".into()), - ); - let result = KafkaProducer::build_sasl_config(SaslMechanism::ScramSha512, &config); - assert!(matches!(result, Ok(SaslConfig::ScramSha512(_)))); - } - - #[test] - fn should_throw_err_on_missing_credentials() { - // Missing username - let config = make_kafka_config(Some("PLAIN"), None, Some("pass".into())); - assert!(matches!( - KafkaProducer::build_sasl_config(SaslMechanism::Plain, &config), - Err(Error::MissingSaslUsername) - )); - - // Missing password - let config = make_kafka_config(Some("PLAIN"), Some("user".into()), None); - assert!(matches!( - KafkaProducer::build_sasl_config(SaslMechanism::Plain, &config), - Err(Error::MissingSaslPassword) - )); - } - #[test] fn debug_output_redacts_the_sasl_password() { let config = make_kafka_config(Some("PLAIN"), Some("user".into()), Some("hunter2".into())); From d893d69091df18faca7f3cd7a6f27408d9392c28 Mon Sep 17 00:00:00 2001 From: MoonBoi9001 Date: Mon, 24 Aug 2026 18:25:50 +0300 Subject: [PATCH 3/4] fix(producer): bound Kafka connects that would otherwise hang The Kafka client library retries an unreachable broker with no deadline, so connecting while Redpanda is down waited endlessly, dodging every retry and shutdown path around it. Connects now time out (60s default), reject an empty broker list, and metadata can be re-read on demand. --- ...raph_indexing_agreements_events_emitter.rs | 1 + dipper-producer/src/kafka/connection.rs | 8 ++ dipper-producer/src/kafka/consumer.rs | 126 ++++++++++++++++-- dipper-producer/src/kafka/producer.rs | 20 ++- 4 files changed, 140 insertions(+), 15 deletions(-) diff --git a/dipper-producer/src/events/subgraph_indexing_agreements_events_emitter.rs b/dipper-producer/src/events/subgraph_indexing_agreements_events_emitter.rs index 2c1a45a1..7dd2ed9e 100644 --- a/dipper-producer/src/events/subgraph_indexing_agreements_events_emitter.rs +++ b/dipper-producer/src/events/subgraph_indexing_agreements_events_emitter.rs @@ -711,6 +711,7 @@ mod tests { sasl_password: None, tls_enabled: false, tls_ca_cert_path: None, + connect_timeout_secs: 60, } } diff --git a/dipper-producer/src/kafka/connection.rs b/dipper-producer/src/kafka/connection.rs index caf89359..db407626 100644 --- a/dipper-producer/src/kafka/connection.rs +++ b/dipper-producer/src/kafka/connection.rs @@ -23,6 +23,10 @@ pub(crate) struct ConnectOptions<'a> { /// Connects a Kafka client with the given SASL/TLS settings. pub(crate) async fn connect(opts: ConnectOptions<'_>) -> Result { + if opts.brokers.is_empty() { + return Err(ConnectionError::MissingBrokers); + } + let mut builder = ClientBuilder::new(opts.brokers.to_vec()); if let Some(mechanism_str) = opts.sasl_mechanism { @@ -108,6 +112,10 @@ pub enum ConnectionError { #[error("failed to connect to Kafka brokers")] Connection(#[source] rskafka::client::error::Error), + /// The brokers list is empty + #[error("brokers must list at least 1 broker address")] + MissingBrokers, + /// Unsupported SASL mechanism #[error("unsupported SASL mechanism '{0}', supported: PLAIN, SCRAM-SHA-256, SCRAM-SHA-512")] UnsupportedSaslMechanism(String), diff --git a/dipper-producer/src/kafka/consumer.rs b/dipper-producer/src/kafka/consumer.rs index eea727ca..c3ef9786 100644 --- a/dipper-producer/src/kafka/consumer.rs +++ b/dipper-producer/src/kafka/consumer.rs @@ -5,7 +5,10 @@ use std::{collections::BTreeMap, path::PathBuf, sync::Arc, time::Duration}; use rskafka::{ - client::partition::{OffsetAt, PartitionClient, UnknownTopicHandling}, + client::{ + Client, + partition::{OffsetAt, PartitionClient, UnknownTopicHandling}, + }, record::RecordAndOffset, }; @@ -35,6 +38,16 @@ pub struct KafkaConsumerConfig { /// Path to a PEM-encoded CA certificate file for TLS verification. #[serde(default)] pub tls_ca_cert_path: Option, + /// Seconds allowed for the initial connect and topic discovery (default: + /// 60). Load-bearing: the underlying client retries an unreachable broker + /// forever, so without this bound `connect` would never return. + #[serde(default = "default_connect_timeout_secs")] + pub connect_timeout_secs: u64, +} + +/// Default number of seconds allowed for connect and topic discovery. +pub fn default_connect_timeout_secs() -> u64 { + 60 } // Manual impl instead of derive: the service logs the whole config with Debug @@ -52,6 +65,7 @@ impl std::fmt::Debug for KafkaConsumerConfig { ) .field("tls_enabled", &self.tls_enabled) .field("tls_ca_cert_path", &self.tls_ca_cert_path) + .field("connect_timeout_secs", &self.connect_timeout_secs) .finish() } } @@ -59,6 +73,7 @@ impl std::fmt::Debug for KafkaConsumerConfig { /// Kafka consumer bound to a single topic, with a partition client per /// discovered partition. Thread-safe; share across tasks via `Arc`. pub struct KafkaConsumer { + client: Client, topic: String, partition_clients: BTreeMap>, } @@ -73,8 +88,19 @@ impl KafkaConsumer { /// Connects to the brokers and binds to the configured topic. Partitions /// are discovered from broker metadata, so there is no partition count to - /// configure; a topic the credentials cannot see is an error. + /// configure; a topic the credentials cannot see is an error. The whole + /// sequence is bounded by `connect_timeout_secs`, since the underlying + /// client would otherwise retry an unreachable broker forever. pub async fn connect(config: &KafkaConsumerConfig) -> Result { + tokio::time::timeout( + Duration::from_secs(config.connect_timeout_secs), + Self::connect_inner(config), + ) + .await + .map_err(|_| ConsumerError::Timeout)? + } + + async fn connect_inner(config: &KafkaConsumerConfig) -> Result { let client = connection::connect(ConnectOptions { brokers: &config.brokers, sasl_mechanism: config.sasl_mechanism.as_deref(), @@ -85,19 +111,10 @@ impl KafkaConsumer { }) .await?; - let topics = client - .list_topics() - .await - .map_err(ConsumerError::Metadata)?; - let topic = topics - .into_iter() - .find(|t| t.name == config.topic) - .ok_or_else(|| ConsumerError::TopicNotFound { - topic: config.topic.clone(), - })?; + let partitions = discover_partitions(&client, &config.topic).await?; let mut partition_clients = BTreeMap::new(); - for partition in topic.partitions { + for partition in partitions { let partition_client = client .partition_client(&config.topic, partition, UnknownTopicHandling::Error) .await @@ -106,6 +123,7 @@ impl KafkaConsumer { } Ok(Self { + client, topic: config.topic.clone(), partition_clients, }) @@ -121,6 +139,19 @@ impl KafkaConsumer { self.partition_clients.keys().copied().collect() } + /// Re-reads broker metadata and returns the topic's current partition + /// count, so callers can notice a topic growing partitions after connect + /// (this consumer keeps serving only the set discovered at connect). + pub async fn current_partition_count(&self) -> Result { + tokio::time::timeout( + Self::OFFSET_TIMEOUT, + discover_partitions(&self.client, &self.topic), + ) + .await + .map_err(|_| ConsumerError::Timeout)? + .map(|partitions| partitions.len()) + } + /// Fetches records from one partition starting at `offset`, waiting up to /// `max_wait_ms` for data to arrive. Returns the records (with their /// offsets) and the partition's current high watermark. @@ -134,9 +165,11 @@ impl KafkaConsumer { let partition_client = self.partition_client(partition)?; let timeout = Duration::from_millis(max_wait_ms.max(0) as u64) + Self::FETCH_TIMEOUT_MARGIN; + // rskafka encodes the range's exclusive end minus 1 as the wire-level + // max_bytes, so widen by 1 to request the full `max_bytes` budget. tokio::time::timeout( timeout, - partition_client.fetch_records(offset, 1..max_bytes, max_wait_ms), + partition_client.fetch_records(offset, 1..max_bytes.saturating_add(1), max_wait_ms), ) .await .map_err(|_| ConsumerError::Timeout)? @@ -163,6 +196,22 @@ impl KafkaConsumer { } } +/// Lists the topic's partitions from broker metadata; a topic the broker does +/// not report (missing, or invisible to the credentials) is an error. +async fn discover_partitions( + client: &Client, + topic: &str, +) -> Result, ConsumerError> { + let topics = client.list_topics().await.map_err(ConsumerError::Metadata)?; + topics + .into_iter() + .find(|t| t.name == topic) + .map(|t| t.partitions) + .ok_or_else(|| ConsumerError::TopicNotFound { + topic: topic.to_string(), + }) +} + /// Errors that can occur when working with the Kafka consumer. #[derive(Debug, thiserror::Error)] pub enum ConsumerError { @@ -233,6 +282,54 @@ mod tests { assert_eq!(config.sasl_password, None); assert!(!config.tls_enabled); assert_eq!(config.tls_ca_cert_path, None); + assert_eq!(config.connect_timeout_secs, 60); + } + + #[tokio::test] + async fn connect_gives_up_after_the_configured_timeout() { + // Port 1 refuses connections, which the underlying client retries + // forever; the configured bound must turn that into an error. + let config = KafkaConsumerConfig { + brokers: vec!["127.0.0.1:1".to_string()], + topic: "test".to_string(), + sasl_mechanism: None, + sasl_username: None, + sasl_password: None, + tls_enabled: false, + tls_ca_cert_path: None, + connect_timeout_secs: 1, + }; + let started = std::time::Instant::now(); + let result = KafkaConsumer::connect(&config).await; + assert!( + matches!(result, Err(ConsumerError::Timeout)), + "expected a timeout, got {:?}", + result.err() + ); + assert!( + started.elapsed() < Duration::from_secs(30), + "connect must give up promptly" + ); + } + + #[tokio::test] + async fn connect_rejects_an_empty_brokers_list() { + let config = KafkaConsumerConfig { + brokers: Vec::new(), + topic: "test".to_string(), + sasl_mechanism: None, + sasl_username: None, + sasl_password: None, + tls_enabled: false, + tls_ca_cert_path: None, + connect_timeout_secs: 60, + }; + assert!(matches!( + KafkaConsumer::connect(&config).await, + Err(ConsumerError::Connection( + super::super::connection::ConnectionError::MissingBrokers + )) + )); } #[test] @@ -273,6 +370,7 @@ mod tests { sasl_password: Some("hunter2".to_string()), tls_enabled: false, tls_ca_cert_path: None, + connect_timeout_secs: 60, }; let rendered = format!("{config:?}"); assert!( diff --git a/dipper-producer/src/kafka/producer.rs b/dipper-producer/src/kafka/producer.rs index d2f78226..802d8ac7 100644 --- a/dipper-producer/src/kafka/producer.rs +++ b/dipper-producer/src/kafka/producer.rs @@ -38,6 +38,11 @@ pub struct KafkaConfig { /// Path to a PEM-encoded CA certificate file for TLS verification. #[serde(default)] pub tls_ca_cert_path: Option, + /// Seconds allowed for the initial connect and partition binding (default: + /// 60). Load-bearing: the underlying client retries an unreachable broker + /// forever, so without this bound `new` would never return. + #[serde(default = "super::consumer::default_connect_timeout_secs")] + pub connect_timeout_secs: u64, } // Manual impl instead of derive: the service logs the whole config with Debug @@ -56,6 +61,7 @@ impl std::fmt::Debug for KafkaConfig { ) .field("tls_enabled", &self.tls_enabled) .field("tls_ca_cert_path", &self.tls_ca_cert_path) + .field("connect_timeout_secs", &self.connect_timeout_secs) .finish() } } @@ -81,8 +87,19 @@ pub struct KafkaProducer { impl KafkaProducer { const PRODUCE_TIMEOUT: Duration = Duration::from_secs(30); - /// Creates a new Kafka producer with the given configuration. + /// Creates a new Kafka producer with the given configuration. Bounded by + /// `connect_timeout_secs`, since the underlying client retries an + /// unreachable broker forever. pub async fn new(config: &KafkaConfig) -> Result { + tokio::time::timeout( + Duration::from_secs(config.connect_timeout_secs), + Self::new_inner(config), + ) + .await + .map_err(|_| Error::Timeout)? + } + + async fn new_inner(config: &KafkaConfig) -> Result { if config.partitions == 0 { return Err(Error::InvalidPartitionCount); } @@ -222,6 +239,7 @@ mod tests { sasl_password: sasl_pass, tls_enabled: false, tls_ca_cert_path: None, + connect_timeout_secs: 60, } } } From 15b0b87339900d3c49d155994c0ae1ab59f3d488 Mon Sep 17 00:00:00 2001 From: MoonBoi9001 Date: Mon, 24 Aug 2026 18:42:20 +0300 Subject: [PATCH 4/4] chore(producer): apply nightly rustfmt to the consumer module --- dipper-producer/src/kafka/consumer.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/dipper-producer/src/kafka/consumer.rs b/dipper-producer/src/kafka/consumer.rs index c3ef9786..c578f4a7 100644 --- a/dipper-producer/src/kafka/consumer.rs +++ b/dipper-producer/src/kafka/consumer.rs @@ -202,7 +202,10 @@ async fn discover_partitions( client: &Client, topic: &str, ) -> Result, ConsumerError> { - let topics = client.list_topics().await.map_err(ConsumerError::Metadata)?; + let topics = client + .list_topics() + .await + .map_err(ConsumerError::Metadata)?; topics .into_iter() .find(|t| t.name == topic)