From e05445817056cd13887b9784956b887a16a6b4ce Mon Sep 17 00:00:00 2001 From: Markus Unterwaditzer Date: Mon, 14 Sep 2026 16:44:38 +0200 Subject: [PATCH 1/5] wip --- Cargo.toml | 4 + rust-arroyo/examples/reloading_producer.rs | 112 ++ rust-arroyo/src/backends/kafka/config.rs | 5 + rust-arroyo/src/backends/kafka/config_blob.rs | 362 ++++ rust-arroyo/src/backends/kafka/mod.rs | 2 + rust-arroyo/src/backends/kafka/producer.rs | 10 + .../src/backends/kafka/reloading_producer.rs | 1478 +++++++++++++++++ 7 files changed, 1973 insertions(+) create mode 100644 rust-arroyo/examples/reloading_producer.rs create mode 100644 rust-arroyo/src/backends/kafka/config_blob.rs create mode 100644 rust-arroyo/src/backends/kafka/reloading_producer.rs diff --git a/Cargo.toml b/Cargo.toml index d8a20001..28f038e2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -45,3 +45,7 @@ path = "rust-arroyo/examples/transform_and_produce.rs" [[example]] name = "base_consumer" path = "rust-arroyo/examples/base_consumer.rs" + +[[example]] +name = "reloading_producer" +path = "rust-arroyo/examples/reloading_producer.rs" diff --git a/rust-arroyo/examples/reloading_producer.rs b/rust-arroyo/examples/reloading_producer.rs new file mode 100644 index 00000000..5e918eae --- /dev/null +++ b/rust-arroyo/examples/reloading_producer.rs @@ -0,0 +1,112 @@ +// How Relay wires up a self-reloading producer. Relay passes opaque bytes from +// options-automator at startup and on changes; Arroyo validates and applies them. +// +// Run with: cargo run --example reloading_producer +extern crate sentry_arroyo; + +use std::time::Duration; + +use sentry_arroyo::backends::kafka::config_blob::ProducerSelector; +use sentry_arroyo::backends::kafka::reloading_producer::{ + ReloadConfig, ReloadOutcome, ReloadingKafkaProducer, +}; +use sentry_arroyo::backends::kafka::types::KafkaPayload; +use sentry_arroyo::backends::Producer; +use sentry_arroyo::types::{Topic, TopicOrPartition}; + +/// Stands in for the opaque bytes options-automator hands to Relay. +fn config_blob(acks: &str) -> Vec { + format!( + r#"{{ + "clusters": {{ + "events": {{ + "config": {{"bootstrap.servers": "127.0.0.1:9092"}}, + "producer_config": {{"compression.type": "lz4"}} + }} + }}, + "topics": {{ + "ingest-events": {{ + "cluster": "events", + "logical_topic": "ingest-events", + "producer_config": {{"acks": "{acks}"}}, + "producer_config_overrides": {{"relay": {{"linger.ms": "5"}}}} + }} + }} + }}"# + ) + .into_bytes() +} + +#[tokio::main] +async fn main() { + tracing_subscriber::fmt::init(); + + // `app` picks which producer_config_overrides entry applies. Asserting the + // logical topic makes Arroyo reject a blob that repoints this topic at + // something Relay does not expect. + let selector = + ProducerSelector::new("ingest-events", "relay").expect_logical_topic("ingest-events"); + + let settings = ReloadConfig { + // How long a keyed produce waits for the old client to drain before + // giving up on ordering. + drain_timeout: Duration::from_secs(10), + // Spread swaps across the fleet so pods do not all stall at once. + jitter: Duration::from_secs(2), + // Require the new client to reach a broker before swapping it in, so a + // typo'd address cannot retire a working producer. + probe_timeout: Some(Duration::from_secs(5)), + probe_retry_interval: Duration::from_secs(5), + }; + + let producer = ReloadingKafkaProducer::new(&config_blob("all"), selector, settings) + .expect("initial config should be valid"); + + let destination = TopicOrPartition::Topic(Topic::new("ingest-events")); + + // Clones share one client and all see reloads. + let worker = producer.clone(); + tokio::task::spawn_blocking(move || { + for i in 0..20 { + // Keyed messages wait for drains to preserve order; unkeyed ones do not. + let key = (i % 2 == 0).then(|| b"some-key".to_vec()); + let payload = KafkaPayload::new(key, None, Some(format!("message-{i}").into_bytes())); + + match worker.produce(&destination, payload) { + Ok(()) => {} + Err(error) => tracing::error!(%error, "produce failed"), + } + + std::thread::sleep(Duration::from_millis(100)); + } + }); + + // Simulate Relay's options-automator callback. + for (attempt, blob) in [ + config_blob("all"), // same as what is running: no swap + config_blob("1"), // a real change: swap + b"{ not json".to_vec(), // rejected, old client keeps running + ] + .into_iter() + .enumerate() + { + tokio::time::sleep(Duration::from_millis(500)).await; + + // Returns immediately: the worker does the probing and swapping. + match producer.push_config(&blob) { + Ok(ReloadOutcome::Accepted) => { + tracing::info!(attempt, "config accepted, rolling out in the background") + } + Ok(ReloadOutcome::Unchanged) => tracing::info!(attempt, "config unchanged"), + // A bad config leaves the current producer running. + Err(error) => tracing::error!(attempt, %error, "config rejected, still on old config"), + } + } + + // Give the worker time to finish rolling out the accepted config. + tokio::time::sleep(Duration::from_secs(6)).await; + tracing::info!( + generation = producer.generation(), + "final config generation" + ); +} diff --git a/rust-arroyo/src/backends/kafka/config.rs b/rust-arroyo/src/backends/kafka/config.rs index 004f06d4..424b5a56 100644 --- a/rust-arroyo/src/backends/kafka/config.rs +++ b/rust-arroyo/src/backends/kafka/config.rs @@ -93,6 +93,11 @@ impl KafkaConfig { pub fn get_config_value(&self, key: &str) -> Option<&String> { self.config_map.get(key) } + + /// All configuration parameters, for comparing two configs. + pub fn config_map(&self) -> &HashMap { + &self.config_map + } } impl From for RdKafkaConfig { diff --git a/rust-arroyo/src/backends/kafka/config_blob.rs b/rust-arroyo/src/backends/kafka/config_blob.rs new file mode 100644 index 00000000..02c80cf2 --- /dev/null +++ b/rust-arroyo/src/backends/kafka/config_blob.rs @@ -0,0 +1,362 @@ +//! Parsing of the central Kafka configmap blob. +//! +//! The host passes bytes from options-automator at startup and on changes; +//! Arroyo never reads the blob from disk. Arroyo owns this Sentry-specific +//! schema so hosts can treat it as opaque. + +use std::collections::HashMap; + +use serde::Deserialize; +use thiserror::Error; + +use super::config::KafkaConfig; + +/// Failure to resolve a config blob. Rejection leaves the current config intact. +#[non_exhaustive] +#[derive(Error, Debug)] +pub enum ConfigBlobError { + #[error("config blob is not valid JSON")] + Malformed(#[from] serde_json::Error), + + #[error("config blob has no entry for topic {topic:?}")] + UnknownTopic { topic: String }, + + #[error("topic {topic:?} refers to undefined cluster {cluster:?}")] + UnknownCluster { topic: String, cluster: String }, + + #[error("environment variable {variable:?} for config key {key:?} is not set")] + MissingEnvVar { key: String, variable: String }, + + #[error("topic {topic:?} has logical_topic {actual:?}, application expected {expected:?}")] + UnexpectedLogicalTopic { + topic: String, + actual: String, + expected: String, + }, +} + +/// The whole configmap, covering every cluster and topic in a region. +#[derive(Debug, Clone, Deserialize)] +pub struct ConfigBlob { + #[serde(default)] + clusters: HashMap, + #[serde(default)] + topics: HashMap, +} + +#[derive(Debug, Clone, Default, Deserialize)] +struct ClusterEntry { + /// Config applied to producers and consumers alike. + #[serde(default)] + config: HashMap, + /// Config mapping keys to environment variables, keeping secrets out of the blob. + #[serde(default)] + env_config: HashMap, + #[serde(default)] + producer_config: HashMap, + #[serde(default)] + producer_env_config: HashMap, +} + +#[derive(Debug, Clone, Default, Deserialize)] +struct TopicEntry { + /// Which cluster this physical topic lives on. + #[serde(default)] + cluster: Option, + /// Informational, and optionally asserted by the application. + #[serde(default)] + logical_topic: Option, + #[serde(default)] + producer_config: HashMap, + /// Per-application overrides, keyed by application name such as `relay`. + #[serde(default)] + producer_config_overrides: HashMap>, +} + +/// Where to read environment variables from when resolving `env_config`. +/// +/// Exists so tests can resolve a blob without touching process state. +pub trait EnvSource { + fn get(&self, key: &str) -> Option; +} + +/// Reads from the real process environment. +pub struct ProcessEnv; + +impl EnvSource for ProcessEnv { + fn get(&self, key: &str) -> Option { + std::env::var(key).ok() + } +} + +impl EnvSource for HashMap { + fn get(&self, key: &str) -> Option { + HashMap::get(self, key).cloned() + } +} + +/// Identifies which producer config to pull out of a blob. +#[derive(Debug, Clone)] +pub struct ProducerSelector { + topic: String, + app: String, + expected_logical_topic: Option, +} + +impl ProducerSelector { + /// Selects the producer config for a physical topic, as seen by `app`. + /// + /// `app` decides which `producer_config_overrides` entry applies. + pub fn new(topic: impl Into, app: impl Into) -> Self { + Self { + topic: topic.into(), + app: app.into(), + expected_logical_topic: None, + } + } + + /// Rejects a topic entry with an unexpected `logical_topic`. + pub fn expect_logical_topic(mut self, logical_topic: impl Into) -> Self { + self.expected_logical_topic = Some(logical_topic.into()); + self + } + + pub fn topic(&self) -> &str { + &self.topic + } + + pub fn app(&self) -> &str { + &self.app + } +} + +impl ConfigBlob { + /// Parses a blob. The bytes come from the host application verbatim. + pub fn parse(blob: &[u8]) -> Result { + Ok(serde_json::from_slice(blob)?) + } + + /// Resolves the producer config for one topic, reading secrets from the + /// process environment. + pub fn producer_config( + &self, + selector: &ProducerSelector, + ) -> Result { + self.producer_config_with_env(selector, &ProcessEnv) + } + + /// Same as [`ConfigBlob::producer_config`], with an explicit environment. + pub fn producer_config_with_env( + &self, + selector: &ProducerSelector, + env: &impl EnvSource, + ) -> Result { + let topic_name = &selector.topic; + let topic = self + .topics + .get(topic_name) + .ok_or_else(|| ConfigBlobError::UnknownTopic { + topic: topic_name.clone(), + })?; + + if let Some(expected) = &selector.expected_logical_topic { + let actual = topic.logical_topic.as_deref().unwrap_or(""); + if actual != expected { + return Err(ConfigBlobError::UnexpectedLogicalTopic { + topic: topic_name.clone(), + actual: actual.to_owned(), + expected: expected.clone(), + }); + } + } + + let mut params = HashMap::new(); + + // Later sources win: cluster-wide settings first, then producer + // settings, then the topic's own, then the per-app override. + if let Some(cluster_name) = &topic.cluster { + let cluster = + self.clusters + .get(cluster_name) + .ok_or_else(|| ConfigBlobError::UnknownCluster { + topic: topic_name.clone(), + cluster: cluster_name.clone(), + })?; + + params.extend(cluster.config.clone()); + resolve_env_into(&mut params, &cluster.env_config, env)?; + params.extend(cluster.producer_config.clone()); + resolve_env_into(&mut params, &cluster.producer_env_config, env)?; + } + + params.extend(topic.producer_config.clone()); + if let Some(overrides) = topic.producer_config_overrides.get(&selector.app) { + params.extend(overrides.clone()); + } + + // `bootstrap.servers` is just another key in the blob, so hand the + // whole map over as overrides rather than splitting it out. + Ok(KafkaConfig::new_producer_config(Vec::new(), Some(params))) + } + + /// The logical topic recorded for a physical topic, if any. Useful as a + /// metric tag. + pub fn logical_topic(&self, topic: &str) -> Option<&str> { + self.topics.get(topic)?.logical_topic.as_deref() + } +} + +fn resolve_env_into( + params: &mut HashMap, + env_config: &HashMap, + env: &impl EnvSource, +) -> Result<(), ConfigBlobError> { + for (key, variable) in env_config { + // An unset variable is an error, not an empty value that builds a broken client. + let value = env + .get(variable) + .ok_or_else(|| ConfigBlobError::MissingEnvVar { + key: key.clone(), + variable: variable.clone(), + })?; + params.insert(key.clone(), value); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + const BLOB: &[u8] = br#"{ + "clusters": { + "events": { + "config": {"bootstrap.servers": "broker-1:9092"}, + "env_config": {"sasl.username": "KAFKA_USERNAME"}, + "producer_config": {"acks": "all"} + } + }, + "topics": { + "ingest-events": { + "cluster": "events", + "logical_topic": "ingest-events", + "producer_config": {"compression.type": "lz4"}, + "producer_config_overrides": {"relay": {"acks": "1"}} + }, + "orphan-topic": {"cluster": "does-not-exist"} + } + }"#; + + fn env() -> HashMap { + HashMap::from([("KAFKA_USERNAME".to_string(), "ingest".to_string())]) + } + + fn resolve(topic: &str, app: &str) -> Result { + ConfigBlob::parse(BLOB) + .unwrap() + .producer_config_with_env(&ProducerSelector::new(topic, app), &env()) + } + + #[test] + fn test_layers_cluster_topic_and_app_config() { + let config = resolve("ingest-events", "relay").unwrap(); + + assert_eq!( + config.get_config_value("bootstrap.servers"), + Some(&"broker-1:9092".to_string()) + ); + assert_eq!( + config.get_config_value("compression.type"), + Some(&"lz4".to_string()) + ); + // The per-app override beats the cluster-wide value. + assert_eq!(config.get_config_value("acks"), Some(&"1".to_string())); + } + + #[test] + fn test_app_without_override_keeps_cluster_value() { + let config = resolve("ingest-events", "snuba").unwrap(); + assert_eq!(config.get_config_value("acks"), Some(&"all".to_string())); + } + + #[test] + fn test_env_config_is_resolved_from_environment() { + let config = resolve("ingest-events", "relay").unwrap(); + assert_eq!( + config.get_config_value("sasl.username"), + Some(&"ingest".to_string()) + ); + } + + #[test] + fn test_missing_env_var_is_rejected() { + let result = ConfigBlob::parse(BLOB).unwrap().producer_config_with_env( + &ProducerSelector::new("ingest-events", "relay"), + &HashMap::new(), + ); + assert!(matches!( + result, + Err(ConfigBlobError::MissingEnvVar { variable, .. }) if variable == "KAFKA_USERNAME" + )); + } + + #[test] + fn test_unknown_topic_is_rejected() { + assert!(matches!( + resolve("no-such-topic", "relay"), + Err(ConfigBlobError::UnknownTopic { .. }) + )); + } + + #[test] + fn test_unknown_cluster_is_rejected() { + assert!(matches!( + resolve("orphan-topic", "relay"), + Err(ConfigBlobError::UnknownCluster { .. }) + )); + } + + #[test] + fn test_malformed_blob_is_rejected() { + assert!(matches!( + ConfigBlob::parse(b"not json"), + Err(ConfigBlobError::Malformed(..)) + )); + } + + #[test] + fn test_unexpected_logical_topic_is_rejected() { + let selector = + ProducerSelector::new("ingest-events", "relay").expect_logical_topic("ingest-spans"); + let result = ConfigBlob::parse(BLOB) + .unwrap() + .producer_config_with_env(&selector, &env()); + assert!(matches!( + result, + Err(ConfigBlobError::UnexpectedLogicalTopic { .. }) + )); + } + + #[test] + fn test_matching_logical_topic_is_accepted() { + let selector = + ProducerSelector::new("ingest-events", "relay").expect_logical_topic("ingest-events"); + assert!(ConfigBlob::parse(BLOB) + .unwrap() + .producer_config_with_env(&selector, &env()) + .is_ok()); + } + + #[test] + fn test_unknown_fields_are_tolerated() { + // Consumer keys and future additions must not break producers. + let blob = br#"{ + "topics": {"t": {"consumer_config": {"max.poll.interval.ms": "1"}, "future": 1}}, + "unrelated": {} + }"#; + assert!(ConfigBlob::parse(blob) + .unwrap() + .producer_config_with_env(&ProducerSelector::new("t", "relay"), &env()) + .is_ok()); + } +} diff --git a/rust-arroyo/src/backends/kafka/mod.rs b/rust-arroyo/src/backends/kafka/mod.rs index d896e8ff..8a08d2f3 100644 --- a/rust-arroyo/src/backends/kafka/mod.rs +++ b/rust-arroyo/src/backends/kafka/mod.rs @@ -26,8 +26,10 @@ use std::sync::Arc; use std::time::Duration; pub mod config; +pub mod config_blob; mod errors; pub mod producer; +pub mod reloading_producer; pub mod types; #[derive(Eq, Hash, PartialEq)] diff --git a/rust-arroyo/src/backends/kafka/producer.rs b/rust-arroyo/src/backends/kafka/producer.rs index d630953b..365ae642 100644 --- a/rust-arroyo/src/backends/kafka/producer.rs +++ b/rust-arroyo/src/backends/kafka/producer.rs @@ -112,6 +112,16 @@ where pub fn in_flight_count(&self) -> i32 { self.producer.in_flight_count() } + + /// Blocks until every queued message has been delivered, or the timeout + /// expires. + /// + /// On timeout, messages are left in the queue. Note that dropping the + /// producer purges them rather than delivering them, so a timeout here + /// means those messages are lost unless the caller flushes again. + pub fn flush_for(&self, timeout: Duration) -> Result<(), KafkaError> { + self.producer.flush(timeout) + } } impl ArroyoProducer for KafkaProducer diff --git a/rust-arroyo/src/backends/kafka/reloading_producer.rs b/rust-arroyo/src/backends/kafka/reloading_producer.rs new file mode 100644 index 00000000..8290620c --- /dev/null +++ b/rust-arroyo/src/backends/kafka/reloading_producer.rs @@ -0,0 +1,1478 @@ +//! A producer that swaps its underlying Kafka client when config changes. +//! +//! The host pushes an opaque config blob at startup and on changes. Arroyo +//! validates it, detects changes, and drains and replaces the client. +//! +//! # Ordering across a swap +//! +//! Messages already queued on the old client have not been written yet. If a +//! keyed message went to the new client while an older message with the same +//! key sat in the old client's queue, the two could reach the partition out of +//! order. So keyed produce calls wait for the old client to drain, while +//! unkeyed produce calls, which have no ordering requirement, go straight to +//! the new client. +//! +//! If the drain outlasts [`ReloadConfig::drain_timeout`], keyed produce calls +//! stop waiting rather than blocking the application indefinitely, accepting +//! that ordering may break. +//! +//! # Messages lost on a slow drain +//! +//! rdkafka purges queued and in-flight messages when a producer is dropped, so +//! a drain timeout loses messages rather than delivering them late. Size +//! [`ReloadConfig::drain_timeout`] for the queue depth at reload time, and alert +//! on `arroyo.producer.config_reload_purged_messages`. +//! +//! # Scope +//! +//! This wraps [`KafkaProducer`] only. There is deliberately no reloading +//! counterpart for [`AsyncKafkaProducer`]: the ordering guarantee above is +//! built on a blocking condvar wait made while holding a lock, which an async +//! produce path cannot reuse. An async version would be a sibling type sharing +//! the config and settings here, not a generalization of this one. +//! +//! [`AsyncKafkaProducer`]: super::producer::AsyncKafkaProducer + +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use parking_lot::{Condvar, Mutex, RwLock}; +use rdkafka::error::KafkaError; + +use super::config::KafkaConfig; +use super::config_blob::{ConfigBlob, ConfigBlobError, ProducerSelector}; +use super::producer::KafkaProducer; +use super::types::KafkaPayload; +use crate::backends::{Producer as ArroyoProducer, ProducerError}; +use crate::types::TopicOrPartition; + +/// Knobs for how reloads are carried out. +#[derive(Debug, Clone)] +pub struct ReloadConfig { + /// How long to drain the old client and how long keyed produces wait for it. + pub drain_timeout: Duration, + /// Upper bound on the random delay before a swap. Spreads swaps across a + /// fleet so a config change does not stall every pod at once. + pub jitter: Duration, + /// How long to wait for the new client to reach a broker before swapping + /// it in. `None` skips the check. + /// + /// Building a client does no I/O, so a config naming an unreachable broker + /// constructs fine. Without this check such a config would retire a + /// working producer and install a dead one. The check runs off the + /// caller's thread, so it delays a swap rather than blocking anyone. + pub probe_timeout: Option, + /// How long to wait before probing again after a failed probe. + /// + /// Probing retries until it succeeds or a newer config supersedes it, + /// since an unreachable broker is as likely to be a blip as a bad config. + pub probe_retry_interval: Duration, +} + +impl Default for ReloadConfig { + fn default() -> Self { + Self { + drain_timeout: Duration::from_secs(10), + jitter: Duration::ZERO, + probe_timeout: Some(Duration::from_secs(5)), + probe_retry_interval: Duration::from_secs(5), + } + } +} + +/// What a call to [`ReloadingKafkaProducer::push_config`] did. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ReloadOutcome { + /// The config is valid and differs from what is running, so it has been + /// handed to the reload worker. + /// + /// The swap happens in the background once the config clears the broker + /// probe and jitter, so the new client is not live yet when this + /// returns. Watch [`ReloadingKafkaProducer::generation`] or the + /// `arroyo.producer.config_reload_applied` metric for that. + Accepted, + /// The blob parsed, but produced the same config as the running client. + /// Any config queued but not yet applied is dropped. + Unchanged, +} + +struct Current { + producer: Arc, + config: KafkaConfig, +} + +/// A producer whose configuration can be replaced while it is running. +/// +/// Cheap to clone; clones share one underlying client and all see a reload. +#[derive(Clone)] +pub struct ReloadingKafkaProducer { + inner: Arc, + /// Held only for its `Drop`: when the last handle goes away, this stops + /// the reload worker. The worker holds a `Weak` instead of a handle, + /// so it does not keep this alive. + _shutdown: Arc, +} + +/// Stops the reload worker once every user-facing handle is dropped. +struct ShutdownOnDrop { + inner: Arc, +} + +impl Drop for ShutdownOnDrop { + fn drop(&mut self) { + self.inner.reload.lock().shutdown = true; + self.inner.reload_wake.notify_all(); + } +} + +struct Inner { + /// Read on every produce, write only during a swap. + current: RwLock, + /// Number of drains in progress, normally zero or one. Keyed produce calls + /// wait for it to reach zero. Separate from `current` so produce calls are + /// not blocked merely to read the pointer. + drains_in_flight: Mutex, + /// Signalled when `drains_in_flight` reaches zero. + drain_done: Condvar, + /// Serializes reloads and holds the state a swap needs. + reload: Mutex, + /// Signalled when `reload.desired` or `reload.shutdown` changes. + reload_wake: Condvar, + selector: ProducerSelector, + settings: ReloadConfig, + generation: AtomicU64, + /// Test-only: keyed enqueues without the ordering proof. + #[cfg(test)] + enqueued_during_drain: AtomicU64, + /// Test-only: keyed enqueues with the ordering proof. + #[cfg(test)] + keyed_enqueues: AtomicU64, + /// Test-only: blocks the worker inside `probe` so a test can make a push + /// land while a rollout is provably in flight, instead of racing a sleep + /// against a probe that finishes in microseconds. + #[cfg(test)] + probe_gate: Mutex>>, + /// Test-only: counts probe attempts, so a test can tell whether the worker + /// abandoned a config or is still retrying it. + #[cfg(test)] + probe_attempts: AtomicU64, + /// Test-only: parks the worker after the probe and jitter, immediately + /// before the install decision, so a test can land a push inside the exact + /// window the atomic re-check in `swap` guards. In production that window + /// is the jitter sleep; the gate makes it wide deterministically where + /// jitter makes it wide randomly. + /// + /// This must sit between the jitter sleep and the `reload` lock. Moving it + /// inside the critical section would quietly weaken the test that depends + /// on it into a duplicate of `test_revert_cancels_an_in_flight_rollout`. + #[cfg(test)] + swap_gate: Mutex>>, + /// Test-only: rollouts the worker has finished, however they ended, so a + /// test can wait for the outcome to be settled instead of sleeping. + #[cfg(test)] + rollouts_finished: AtomicU64, +} + +struct ReloadState { + /// The config the host wants instead of the current one. + /// + /// Each push overwrites it; a revert clears it. The worker clones it and + /// leaves it in place while working, so either change cancels an in-flight + /// rollout before it can swap. + desired: Option, + /// Set when every handle is dropped, to stop the worker. + shutdown: bool, +} + +impl ReloadingKafkaProducer { + /// Builds a producer from an initial config blob. + /// + /// The blob is opaque to the caller: pass the bytes from options-automator + /// straight through. `selector` says which topic and application to resolve + /// config for. + pub fn new( + blob: &[u8], + selector: ProducerSelector, + settings: ReloadConfig, + ) -> Result { + let config = ConfigBlob::parse(blob)?.producer_config(&selector)?; + Self::from_config(config, selector, settings) + } + + /// Builds a producer from an already-resolved config. + /// + /// For applications not yet on the central configmap that still want + /// reloading. Once a blob is pushed in, it takes over. + pub fn from_config( + config: KafkaConfig, + selector: ProducerSelector, + settings: ReloadConfig, + ) -> Result { + let producer = Arc::new(KafkaProducer::new(config.clone())?); + tracing::info!( + topic = %selector.topic(), + app = %selector.app(), + "loaded initial producer config" + ); + + let inner = Arc::new(Inner { + current: RwLock::new(Current { producer, config }), + drains_in_flight: Mutex::new(0), + drain_done: Condvar::new(), + reload: Mutex::new(ReloadState { + desired: None, + shutdown: false, + }), + reload_wake: Condvar::new(), + selector, + settings, + generation: AtomicU64::new(0), + #[cfg(test)] + enqueued_during_drain: AtomicU64::new(0), + #[cfg(test)] + keyed_enqueues: AtomicU64::new(0), + #[cfg(test)] + probe_gate: Mutex::new(None), + #[cfg(test)] + probe_attempts: AtomicU64::new(0), + #[cfg(test)] + swap_gate: Mutex::new(None), + #[cfg(test)] + rollouts_finished: AtomicU64::new(0), + }); + + // A weak reference lets the worker exit after the last handle drops. + let weak = Arc::downgrade(&inner); + std::thread::Builder::new() + .name("arroyo-producer-reload".to_owned()) + .spawn(move || { + // If every handle was dropped before the worker started, there + // is nothing to do. + let Some(inner) = weak.upgrade() else { return }; + inner.run_worker(); + }) + .map_err(ReloadError::WorkerSpawn)?; + + Ok(Self { + inner: inner.clone(), + _shutdown: Arc::new(ShutdownOnDrop { inner }), + }) + } + + /// Hands a new config blob to the producer. + /// + /// Call when the blob changes. Unchanged blobs do nothing; invalid blobs + /// leave the current client untouched. + /// + /// This parses and hands off without blocking. [`ReloadOutcome::Accepted`] + /// means rollout started, not that the config is live; watch + /// [`ReloadingKafkaProducer::generation`]. Broker work happens on the worker. + pub fn push_config(&self, blob: &[u8]) -> Result { + let config = ConfigBlob::parse(blob)?.producer_config(&self.inner.selector)?; + Ok(self.inner.set_desired(config)) + } + + /// Number of swaps so far. Useful as a metric tag to tell which config a + /// message was produced under. + pub fn generation(&self) -> u64 { + self.inner.generation.load(Ordering::SeqCst) + } + + /// Messages accepted but not yet delivered by the live client. + pub fn in_flight_count(&self) -> i32 { + self.inner.current.read().producer.in_flight_count() + } + + /// Pretends a drain is in progress, so tests can observe how produce calls + /// behave during one without racing a real flush. + #[cfg(test)] + fn block_drain_for_test(&self) -> DrainGuard { + DrainGuard::new(self.inner.clone()) + } +} + +impl Inner { + /// Records the config the host wants running and wakes the worker. + fn set_desired(self: &Arc, config: KafkaConfig) -> ReloadOutcome { + let mut reload = self.reload.lock(); + + // Lock order is `reload` then `current`, matching `swap`. + if config_eq(&config, &self.current.read().config) { + // Clearing `desired` cancels any rollout back to this current config. + reload.desired = None; + self.reload_wake.notify_all(); + metrics::counter!( + "arroyo.producer.config_reload_unchanged", + "topic" => self.selector.topic().to_owned(), + ) + .increment(1); + return ReloadOutcome::Unchanged; + } + + // Overwrite so only the newest config remains. An identical push leaves + // an in-flight rollout undisturbed. + reload.desired = Some(config); + self.reload_wake.notify_all(); + + metrics::counter!( + "arroyo.producer.config_reload_accepted", + "topic" => self.selector.topic().to_owned(), + ) + .increment(1); + + ReloadOutcome::Accepted + } + + /// Runs the reload worker until every handle is dropped. + /// + /// Owns everything that must not block the application: the broker probe, + /// jitter, the swap, and draining the old client. + fn run_worker(self: &Arc) { + loop { + let config = { + let mut reload = self.reload.lock(); + + loop { + if reload.shutdown { + return; + } + + // Clone rather than take: a newer push or revert can change + // `desired` and supersede this rollout. + match reload.desired.clone() { + Some(config) => break config, + None => self.reload_wake.wait(&mut reload), + } + } + }; + + self.reload_to(&config); + + #[cfg(test)] + self.rollouts_finished.fetch_add(1, Ordering::SeqCst); + + // Done with this config, whether it went live or was abandoned. + // Clear it so the worker can go back to sleep, unless a push + // replaced it while we worked, in which case leave the newer one + // for the next iteration. + let mut reload = self.reload.lock(); + if reload + .desired + .as_ref() + .is_some_and(|d| config_eq(d, &config)) + { + reload.desired = None; + } + } + } + + /// Probes a config and swaps to it, retrying until it is healthy or the + /// host stops wanting it. + fn reload_to(self: &Arc, config: &KafkaConfig) { + // Build once outside the retry loop. Rebuilding would spawn fresh + // rdkafka threads to connect-spam a dead broker on every retry. + // Invalid client config never reaches the probe or affects the current client. + let candidate = match KafkaProducer::new(config.clone()) { + Ok(producer) => Arc::new(producer), + Err(error) => { + metrics::counter!( + "arroyo.producer.config_reload_rejected", + "topic" => self.selector.topic().to_owned(), + "reason" => "client_build", + ) + .increment(1); + tracing::error!( + topic = %self.selector.topic(), + %error, + "kafka rejected the new config, staying on the running one" + ); + return; + } + }; + + loop { + match self.probe(&candidate) { + Ok(()) => { + // `swap` re-checks that the config is still wanted itself, + // atomically with installing the client. + self.swap(candidate.clone(), config); + return; + } + Err(error) => { + metrics::counter!( + "arroyo.producer.config_reload_probe_failed", + "topic" => self.selector.topic().to_owned(), + ) + .increment(1); + tracing::warn!( + topic = %self.selector.topic(), + %error, + "new config cannot reach a broker, staying on the running one" + ); + } + } + + // Retry until healthy or superseded. An unreachable broker is as + // likely to be a blip as a bad config, and the running producer is + // unaffected while we wait. + let mut reload = self.reload.lock(); + let interval = self.settings.probe_retry_interval; + + // Check before waiting too: a push that arrived while we were + // blocked in the probe has already been notified, so waiting the + // full interval first would just add latency. + if superseded(&reload, config) { + return; + } + + self.reload_wake.wait_for(&mut reload, interval); + + if superseded(&reload, config) { + // The host wants something else, or we are shutting down. + // Abandon this one; the worker loop picks up what is desired. + return; + } + } + } + + /// Checks that a newly built client can actually reach a broker. + /// + /// Client construction does no I/O; this prevents swapping in an + /// unreachable producer. + fn probe(&self, candidate: &KafkaProducer) -> Result<(), KafkaError> { + let Some(timeout) = self.settings.probe_timeout else { + return Ok(()); + }; + + #[cfg(test)] + self.probe_attempts.fetch_add(1, Ordering::Relaxed); + + // Test-only: wait until the test releases the gate, so the rollout is + // observably in flight while the test pushes a new config. + #[cfg(test)] + { + let gate = self.probe_gate.lock().take(); + if let Some(gate) = gate { + let _ = gate.recv(); + } + } + + let started = Instant::now(); + // Metadata for the topic we produce to: one round-trip, no message + // written, and it also catches a topic that does not exist. + let result = + candidate.validate_topic(crate::types::Topic::new(self.selector.topic()), timeout); + + metrics::histogram!( + "arroyo.producer.config_reload_probe_ms", + "topic" => self.selector.topic().to_owned(), + ) + .record(started.elapsed().as_millis() as f64); + + result + } + + /// Installs an already-probed client and drains the old one. + fn swap(self: &Arc, new_producer: Arc, config: &KafkaConfig) { + // Jitter before claiming anything, so a fleet does not swap in unison. + if !self.settings.jitter.is_zero() { + let jitter = rand::random::() * self.settings.jitter.as_secs_f64(); + std::thread::sleep(Duration::from_secs_f64(jitter)); + } + + // Test-only: park here, after the jitter and before the install + // decision, so a test can land a push inside the guarded window. + #[cfg(test)] + { + let gate = self.swap_gate.lock().take(); + if let Some(gate) = gate { + let _ = gate.recv(); + } + } + + // After jitter, hold `reload` across the superseded check and install; + // otherwise a push between them could install an unwanted client. + // Lock order is `reload` then `current`, matching `set_desired`; never + // reverse it. + let reload = self.reload.lock(); + + if superseded(&reload, config) { + metrics::counter!( + "arroyo.producer.config_reload_superseded", + "topic" => self.selector.topic().to_owned(), + ) + .increment(1); + tracing::info!( + topic = %self.selector.topic(), + "config superseded before it went live, not swapping to it" + ); + return; + } + + // Count the drain before the swap becomes visible. A keyed produce + // checks this counter while holding a read guard on `current`, so + // anything that sees the new client also sees the drain and waits. + // + // Decremented when this guard is dropped on the drain thread, even if + // that thread panics. Leaking a count would make every later keyed + // produce wait out the full drain timeout, forever. + let drain_guard = DrainGuard::new(self.clone()); + + let old = { + let mut current = self.current.write(); + std::mem::replace( + &mut *current, + Current { + producer: new_producer, + config: config.clone(), + }, + ) + }; + + drop(reload); + + self.generation.fetch_add(1, Ordering::SeqCst); + metrics::counter!( + "arroyo.producer.config_reload_applied", + "topic" => self.selector.topic().to_owned(), + ) + .increment(1); + + let timeout = self.settings.drain_timeout; + let selector_topic = self.selector.topic().to_owned(); + + // Drain inline on the worker to keep `drains_in_flight` at one. Thus a + // keyed produce waits for at most one drain; unkeyed produce is unaffected. + { + let _drain_guard = drain_guard; + + let started = Instant::now(); + let result = old.producer.flush_for(timeout); + let elapsed = started.elapsed(); + + metrics::histogram!( + "arroyo.producer.config_reload_drain_ms", + "topic" => selector_topic.clone(), + ) + .record(elapsed.as_millis() as f64); + + match result { + Ok(()) => { + tracing::info!( + topic = %selector_topic, + drain_ms = elapsed.as_millis() as u64, + "drained producer after config reload" + ); + } + Err(error) => { + // Drop purges queued and in-flight messages, so a timeout + // loses them rather than delivering them late. Their callbacks report + // purge errors in `arroyo.producer.produce_status`; alert on + // `arroyo.producer.config_reload_purged_messages`. + let queued = old.producer.in_flight_count(); + metrics::counter!( + "arroyo.producer.config_reload_drain_timeout", + "topic" => selector_topic.clone(), + ) + .increment(1); + metrics::counter!( + "arroyo.producer.config_reload_purged_messages", + "topic" => selector_topic.clone(), + ) + .increment(queued.max(0) as u64); + tracing::error!( + topic = %selector_topic, + %error, + queued, + "producer drain timed out after config reload, queued messages will be dropped" + ); + } + } + + drop(old); + } + } + + /// Blocks until no drain is in progress, or until the drain timeout + /// expires. Returns `false` if it gave up waiting. + /// + /// On timeout we fail open: producing matters more than ordering, so the + /// caller proceeds and we record that ordering may have been broken. + fn wait_for_drain(&self, deadline: Instant) -> bool { + let mut in_flight = self.drains_in_flight.lock(); + + while *in_flight > 0 { + if self + .drain_done + .wait_until(&mut in_flight, deadline) + .timed_out() + { + metrics::counter!( + "arroyo.producer.config_reload_produce_fail_open", + "topic" => self.selector.topic().to_owned(), + ) + .increment(1); + tracing::warn!( + topic = %self.selector.topic(), + "keyed produce proceeding before drain finished, ordering may be broken" + ); + return false; + } + } + + true + } + + /// Records whether a keyed enqueue observed no drain under its `current` + /// guard. Checking first or enqueueing after releasing the guard records no + /// proof; only the deliberate timeout path may do so. + #[cfg(test)] + fn record_keyed_enqueue(&self, ordered: bool) { + if ordered { + self.keyed_enqueues.fetch_add(1, Ordering::Relaxed); + } else { + self.enqueued_during_drain.fetch_add(1, Ordering::Relaxed); + } + } +} + +/// Releases the drain count on drop, including during panic. +struct DrainGuard { + inner: Arc, +} + +impl DrainGuard { + /// Registers a drain; `Drop` guarantees the matching decrement. + fn new(inner: Arc) -> Self { + *inner.drains_in_flight.lock() += 1; + Self { inner } + } +} + +impl Drop for DrainGuard { + fn drop(&mut self) { + let mut in_flight = self.inner.drains_in_flight.lock(); + *in_flight -= 1; + if *in_flight == 0 { + self.inner.drain_done.notify_all(); + } + } +} + +impl ArroyoProducer for ReloadingKafkaProducer { + fn produce( + &self, + destination: &TopicOrPartition, + payload: KafkaPayload, + ) -> Result<(), ProducerError> { + // Unkeyed messages have no ordering constraint, so they never wait for + // a drain and go straight to whichever client is live. + if payload.key().is_none() { + let current = self.inner.current.read(); + return current.producer.produce(destination, payload); + } + + // Take `current` first and hold it through the counter check and enqueue. + // `swap` counts the drain before taking the write lock, so we either see + // the drain and wait, or block the swap until enqueueing on the old client. + // Checking the counter first would let a swap land before the `current` + // read, enqueueing this message on the new client ahead of older ones. + let deadline = Instant::now() + self.inner.settings.drain_timeout; + + loop { + let current = self.inner.current.read(); + let draining = *self.inner.drains_in_flight.lock() > 0; + + if !draining { + #[cfg(test)] + self.inner.record_keyed_enqueue(true); + + // Enqueue under the read guard so `swap` cannot install first. + return current.producer.produce(destination, payload); + } + + // Release the read guard before waiting so the swap can proceed. + drop(current); + + if !self.inner.wait_for_drain(deadline) { + // Timed out. Fail open: produce against whatever is live now. + // The only sanctioned way to enqueue a keyed message without + // the ordering proof. + #[cfg(test)] + self.inner.record_keyed_enqueue(false); + + let current = self.inner.current.read(); + return current.producer.produce(destination, payload); + } + } + } +} + +/// Something went wrong while loading or applying a config. +#[non_exhaustive] +#[derive(Debug, thiserror::Error)] +pub enum ReloadError { + #[error("invalid config blob")] + Blob(#[from] ConfigBlobError), + + #[error("kafka rejected the new configuration")] + Kafka(#[from] KafkaError), + + #[error("could not spawn the config reload worker")] + WorkerSpawn(#[source] std::io::Error), +} + +/// Whether shutdown, a newer push, or a revert clearing `desired` cancels this config. +fn superseded(reload: &ReloadState, config: &KafkaConfig) -> bool { + reload.shutdown + || !reload + .desired + .as_ref() + .is_some_and(|desired| config_eq(desired, config)) +} + +/// Compares only the parameter map. `offset_reset_config` is consumer-only and +/// never set by the producer paths, so it cannot differ here. +fn config_eq(left: &KafkaConfig, right: &KafkaConfig) -> bool { + left.config_map() == right.config_map() +} + +#[cfg(test)] +mod tests { + use super::*; + use rdkafka::mocking::MockCluster; + use std::sync::mpsc; + + fn blob(servers: &str, acks: &str) -> Vec { + format!( + r#"{{ + "clusters": {{"events": {{"config": {{"bootstrap.servers": "{servers}"}}}}}}, + "topics": {{ + "ingest-events": {{ + "cluster": "events", + "producer_config": {{"acks": "{acks}"}} + }} + }} + }}"# + ) + .into_bytes() + } + + fn settings() -> ReloadConfig { + ReloadConfig { + drain_timeout: Duration::from_secs(5), + jitter: Duration::ZERO, + // MockCluster answers metadata, so probing stays on in tests. + probe_timeout: Some(Duration::from_secs(5)), + probe_retry_interval: Duration::from_millis(50), + } + } + + /// Waits for the reload worker to reach `generation`, since swaps are no + /// longer synchronous with `push_config`. + #[track_caller] + fn await_generation(producer: &ReloadingKafkaProducer, generation: u64) { + let deadline = Instant::now() + Duration::from_secs(15); + while producer.generation() < generation { + assert!( + Instant::now() < deadline, + "timed out waiting for generation {generation}, at {}", + producer.generation() + ); + std::thread::sleep(Duration::from_millis(10)); + } + } + + /// Blocks until the worker is inside `probe`, which it signals by taking + /// the gate receiver. Deterministic, unlike sleeping and hoping. + #[track_caller] + fn wait_until_probing(producer: &ReloadingKafkaProducer) { + let deadline = Instant::now() + Duration::from_secs(5); + while producer.inner.probe_gate.lock().is_some() { + assert!(Instant::now() < deadline, "worker never reached the probe"); + std::thread::sleep(Duration::from_millis(5)); + } + } + + /// Asserts no swap happens within a short window. + #[track_caller] + fn assert_stays_at_generation(producer: &ReloadingKafkaProducer, generation: u64) { + std::thread::sleep(Duration::from_millis(300)); + assert_eq!(producer.generation(), generation); + } + + fn selector() -> ProducerSelector { + ProducerSelector::new("ingest-events", "relay") + } + + fn producer(servers: &str) -> ReloadingKafkaProducer { + ReloadingKafkaProducer::new(&blob(servers, "all"), selector(), settings()).unwrap() + } + + #[test] + fn test_reload_replaces_the_client() { + let cluster = MockCluster::new(1).unwrap(); + let producer = producer(&cluster.bootstrap_servers()); + assert_eq!(producer.generation(), 0); + + let outcome = producer + .push_config(&blob(&cluster.bootstrap_servers(), "1")) + .unwrap(); + + // The swap is handed to the worker, not done inline. + assert_eq!(outcome, ReloadOutcome::Accepted); + await_generation(&producer, 1); + } + + #[test] + fn test_identical_config_does_not_swap() { + let cluster = MockCluster::new(1).unwrap(); + let servers = cluster.bootstrap_servers(); + let producer = producer(&servers); + + let outcome = producer.push_config(&blob(&servers, "all")).unwrap(); + + assert_eq!(outcome, ReloadOutcome::Unchanged); + assert_stays_at_generation(&producer, 0); + } + + #[test] + fn test_invalid_blob_keeps_the_running_client() { + let cluster = MockCluster::new(1).unwrap(); + let producer = producer(&cluster.bootstrap_servers()); + + assert!(producer.push_config(b"{not json").is_err()); + assert!(producer.push_config(br#"{"topics": {}}"#).is_err()); + + assert_stays_at_generation(&producer, 0); + // Still usable. + producer + .produce( + &TopicOrPartition::Topic(crate::types::Topic::new("ingest-events")), + KafkaPayload::new(None, None, Some(b"payload".to_vec())), + ) + .unwrap(); + } + + #[test] + fn test_config_rejected_by_rdkafka_keeps_the_running_client() { + let cluster = MockCluster::new(1).unwrap(); + let servers = cluster.bootstrap_servers(); + let producer = producer(&servers); + + let bad = String::from_utf8(blob(&servers, "all")).unwrap().replace( + r#""acks": "all""#, + r#""message.timeout.ms": "not-a-number""#, + ); + + // The blob is well-formed, so it is accepted; rdkafka rejects it when + // the worker builds the client, and no swap happens. + assert_eq!( + producer.push_config(bad.as_bytes()).unwrap(), + ReloadOutcome::Accepted + ); + assert_stays_at_generation(&producer, 0); + + // Still producing on the original config. + producer + .produce( + &TopicOrPartition::Topic(crate::types::Topic::new("ingest-events")), + KafkaPayload::new(None, None, Some(b"payload".to_vec())), + ) + .unwrap(); + } + + /// A failing probe retries until a reachable config supersedes it. + #[test] + fn test_failed_probe_is_superseded_by_a_good_config() { + let cluster = MockCluster::new(1).unwrap(); + let servers = cluster.bootstrap_servers(); + let producer = ReloadingKafkaProducer::new( + &blob(&servers, "all"), + selector(), + ReloadConfig { + probe_timeout: Some(Duration::from_millis(300)), + probe_retry_interval: Duration::from_millis(50), + ..settings() + }, + ) + .unwrap(); + + producer.push_config(&blob("127.0.0.1:1", "all")).unwrap(); + assert_stays_at_generation(&producer, 0); + + // A reachable config replaces the one stuck retrying. + producer.push_config(&blob(&servers, "1")).unwrap(); + await_generation(&producer, 1); + assert_eq!( + producer + .inner + .current + .read() + .config + .get_config_value("acks"), + Some(&"1".to_string()) + ); + } + + /// A burst of pushes must cost one swap, not one per push. + /// + /// There is no debounce timer: `desired` is a single slot, so pushes that + /// arrive while the worker is busy overwrite each other and only the + /// newest is ever taken. The probe gate holds the worker busy for the + /// duration of the burst, making this deterministic. + #[test] + fn test_burst_collapses_to_one_swap() { + let cluster = MockCluster::new(1).unwrap(); + let servers = cluster.bootstrap_servers(); + let producer = producer(&servers); + + // Hold the worker inside the probe so the whole burst lands while it + // is busy. + let (release, gate) = mpsc::channel(); + *producer.inner.probe_gate.lock() = Some(gate); + + producer.push_config(&blob(&servers, "1")).unwrap(); + wait_until_probing(&producer); + + for acks in ["0", "2", "3"] { + assert_eq!( + producer.push_config(&blob(&servers, acks)).unwrap(), + ReloadOutcome::Accepted + ); + } + + release.send(()).unwrap(); + + // The first config was superseded and the burst collapsed, so exactly + // one swap happens and it carries the newest config. + await_generation(&producer, 1); + assert_stays_at_generation(&producer, 1); + assert_eq!( + producer + .inner + .current + .read() + .config + .get_config_value("acks"), + Some(&"3".to_string()) + ); + } + + /// A pending config replaced before the worker reaches it is dropped, and + /// re-pushing the config already rolling out does not disturb it. + #[test] + fn test_pending_config_replaced_before_pickup_is_dropped() { + let cluster = MockCluster::new(1).unwrap(); + let servers = cluster.bootstrap_servers(); + let producer = producer(&servers); + + let (release, gate) = mpsc::channel(); + *producer.inner.probe_gate.lock() = Some(gate); + + // Roll out "1" and hold the worker in its probe. + producer.push_config(&blob(&servers, "1")).unwrap(); + wait_until_probing(&producer); + + // Queue a further change, then revert to the config being rolled out. + producer.push_config(&blob(&servers, "0")).unwrap(); + producer.push_config(&blob(&servers, "1")).unwrap(); + + release.send(()).unwrap(); + await_generation(&producer, 1); + + // Settles on "1" with no second swap: the "0" was replaced before the + // worker ever took it. + assert_stays_at_generation(&producer, 1); + assert_eq!( + producer + .inner + .current + .read() + .config + .get_config_value("acks"), + Some(&"1".to_string()) + ); + } + + #[test] + fn test_produce_works_across_a_reload() { + let cluster = MockCluster::new(1).unwrap(); + let servers = cluster.bootstrap_servers(); + let producer = producer(&servers); + let destination = TopicOrPartition::Topic(crate::types::Topic::new("ingest-events")); + + for i in 0..5 { + let keyed = i % 2 == 0; + producer + .produce( + &destination, + KafkaPayload::new( + keyed.then(|| b"key".to_vec()), + None, + Some(format!("msg-{i}").into_bytes()), + ), + ) + .unwrap(); + + if i == 2 { + producer.push_config(&blob(&servers, "1")).unwrap(); + } + } + + await_generation(&producer, 1); + } + + /// The ordering guarantee: a keyed message must never be enqueued while an + /// older client is still draining. + /// + /// The test holds a drain open because MockCluster drains in microseconds, + /// then races a keyed produce with a real swap. It fails if `produce` checks + /// the counter before `current`, which can enqueue on the new client first. + #[test] + fn test_keyed_produce_never_races_a_swap() { + let cluster = MockCluster::new(1).unwrap(); + let servers = cluster.bootstrap_servers(); + let destination = TopicOrPartition::Topic(crate::types::Topic::new("ingest-events")); + + for round in 0..40 { + let producer = producer(&servers); + + // Hold a drain open so the window is guaranteed to be wide. + let held = producer.block_drain_for_test(); + + // A keyed produce that must wait for that drain. + let (tx, rx) = mpsc::channel(); + let writer = producer.clone(); + std::thread::spawn(move || { + let _ = tx.send(writer.produce( + &destination, + KafkaPayload::new(Some(b"key".to_vec()), None, Some(b"keyed".to_vec())), + )); + }); + + // Let it reach the wait, then swap underneath it. + std::thread::sleep(Duration::from_millis(5)); + producer + .push_config(&blob(&servers, &format!("{round}"))) + .unwrap(); + + // Release the drain and let the produce complete. + drop(held); + rx.recv_timeout(Duration::from_secs(5)) + .expect("keyed produce never completed") + .expect("keyed produce failed"); + + assert_eq!( + producer.inner.enqueued_during_drain.load(Ordering::Relaxed), + 0, + "keyed produce landed on a client while an older one was draining" + ); + assert!( + producer.inner.keyed_enqueues.load(Ordering::Relaxed) > 0, + "no keyed enqueues were recorded" + ); + } + } + + /// Unkeyed produce must stay non-blocking even under constant swapping, + /// since it has no ordering constraint to preserve. + #[test] + fn test_unkeyed_produce_never_blocks_under_swaps() { + let cluster = MockCluster::new(1).unwrap(); + let servers = cluster.bootstrap_servers(); + let producer = producer(&servers); + let destination = TopicOrPartition::Topic(crate::types::Topic::new("ingest-events")); + + let stop = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let worst = Arc::new(Mutex::new(Duration::ZERO)); + + let writer = { + let producer = producer.clone(); + let stop = stop.clone(); + let worst = worst.clone(); + std::thread::spawn(move || { + while !stop.load(std::sync::atomic::Ordering::Relaxed) { + let started = Instant::now(); + producer + .produce( + &destination, + KafkaPayload::new(None, None, Some(b"payload".to_vec())), + ) + .unwrap(); + let elapsed = started.elapsed(); + + let mut worst = worst.lock(); + if elapsed > *worst { + *worst = elapsed; + } + } + }) + }; + + for i in 0..25 { + producer + .push_config(&blob(&servers, &format!("{}", i % 2))) + .unwrap(); + std::thread::sleep(Duration::from_millis(2)); + } + + stop.store(true, std::sync::atomic::Ordering::Relaxed); + writer.join().unwrap(); + + // Nowhere near drain_timeout (5s); it never waited on a drain. + let worst = *worst.lock(); + assert!( + worst < Duration::from_millis(500), + "unkeyed produce blocked for {worst:?}, it should never wait for a drain" + ); + } + + /// A panicking drain must not strand the counter, which would make every + /// later keyed produce wait out the full timeout. + #[test] + fn test_drain_guard_releases_on_panic() { + let cluster = MockCluster::new(1).unwrap(); + let producer = producer(&cluster.bootstrap_servers()); + + let inner = producer.inner.clone(); + std::thread::spawn(move || { + let _guard = DrainGuard::new(inner); + panic!("drain exploded"); + }) + .join() + .unwrap_err(); + + assert_eq!(*producer.inner.drains_in_flight.lock(), 0); + + // Keyed produce still proceeds immediately. + let started = Instant::now(); + producer + .produce( + &TopicOrPartition::Topic(crate::types::Topic::new("ingest-events")), + KafkaPayload::new(Some(b"key".to_vec()), None, Some(b"payload".to_vec())), + ) + .unwrap(); + assert!(started.elapsed() < Duration::from_secs(1)); + } + + /// Dropping every handle must stop the reload worker, otherwise each + /// producer leaks a thread for the life of the process. + #[test] + fn test_worker_exits_when_handles_are_dropped() { + let cluster = MockCluster::new(1).unwrap(); + let servers = cluster.bootstrap_servers(); + + let weak = { + let producer = producer(&servers); + let clone = producer.clone(); + producer.push_config(&blob(&servers, "1")).unwrap(); + await_generation(&producer, 1); + drop(clone); + Arc::downgrade(&producer.inner) + }; + + // The worker holds only a weak reference, so once it observes the + // shutdown flag and exits, `Inner` is freed. + let deadline = Instant::now() + Duration::from_secs(10); + while weak.upgrade().is_some() { + assert!( + Instant::now() < deadline, + "reload worker still holds the producer alive" + ); + std::thread::sleep(Duration::from_millis(10)); + } + } + + /// A push that reverts to the running config must cancel a rollout that is + /// already in flight. + /// + /// The worker must clone `desired` and leave it in place so a revert can + /// clear it before the swap. The probe is gated instead of raced against a + /// sleep because MockCluster answers in microseconds and this path requires + /// a successful probe. + #[test] + fn test_revert_cancels_an_in_flight_rollout() { + let cluster = MockCluster::new(1).unwrap(); + let servers = cluster.bootstrap_servers(); + let producer = producer(&servers); + + // Arm the gate before pushing, so the worker cannot get past the probe + // before the revert lands. + let (release, gate) = mpsc::channel(); + *producer.inner.probe_gate.lock() = Some(gate); + + // Reachable config: this probe will succeed once released. + producer.push_config(&blob(&servers, "9")).unwrap(); + + // The worker is now inside the probe, working on this config while it + // remains recorded as the desired one. + wait_until_probing(&producer); + + // The rollout is provably in flight. Revert to what is running. + assert_eq!( + producer.push_config(&blob(&servers, "all")).unwrap(), + ReloadOutcome::Unchanged + ); + + // Let the probe finish successfully. + release.send(()).unwrap(); + + // The reverted-away config must not go live. + assert_stays_at_generation(&producer, 0); + assert_eq!( + producer + .inner + .current + .read() + .config + .get_config_value("acks"), + Some(&"all".to_string()) + ); + } + + /// A revert that lands after the probe passed but before the install must + /// cancel the swap. + /// + /// The superseded check and install must be atomic after jitter, or a revert + /// in between is lost. This test sets jitter to zero and parks the worker at + /// the gate between sleep and lock acquisition, so the revert cannot lose a + /// wall-clock race. + #[test] + fn test_revert_between_probe_and_install_cancels_the_swap() { + let cluster = MockCluster::new(1).unwrap(); + let servers = cluster.bootstrap_servers(); + let producer = producer(&servers); + + // Arm before pushing, so the worker cannot get past the gate first. + let (release, gate) = mpsc::channel(); + *producer.inner.swap_gate.lock() = Some(gate); + + producer.push_config(&blob(&servers, "9")).unwrap(); + + // Worker parked: probe passed, install not reached. + let deadline = Instant::now() + Duration::from_secs(5); + while producer.inner.swap_gate.lock().is_some() { + assert!(Instant::now() < deadline, "worker never reached the swap"); + std::thread::sleep(Duration::from_millis(5)); + } + + // The revert lands inside the window. `Unchanged` is guaranteed here, + // not raced for: the worker is parked before the install. + assert_eq!( + producer.push_config(&blob(&servers, "all")).unwrap(), + ReloadOutcome::Unchanged + ); + + release.send(()).unwrap(); + + // Wait until the rollout ended one way or the other, then check which. + // Both a correct and a broken implementation reach this point, so the + // generation assertion below is what tells them apart. + let deadline = Instant::now() + Duration::from_secs(5); + while producer.inner.rollouts_finished.load(Ordering::SeqCst) == 0 { + assert!( + Instant::now() < deadline, + "worker never finished the rollout" + ); + std::thread::sleep(Duration::from_millis(5)); + } + + assert_eq!( + producer.generation(), + 0, + "a config reverted before install was swapped in anyway" + ); + assert_eq!( + producer + .inner + .current + .read() + .config + .get_config_value("acks"), + Some(&"all".to_string()) + ); + } + + /// After a revert, the worker must stop retrying the reverted-away config. + /// + /// This pins the liveness half of the fix. If the supersede check only + /// looked for a *newer* config rather than comparing against `desired`, a + /// revert would leave nothing to compare against and the worker would keep + /// probing a config the host rolled back. If that broker later recovers, + /// the config goes live long after the fact. + #[test] + fn test_revert_stops_the_probe_retry_loop() { + let cluster = MockCluster::new(1).unwrap(); + let servers = cluster.bootstrap_servers(); + let producer = ReloadingKafkaProducer::new( + &blob(&servers, "all"), + selector(), + ReloadConfig { + probe_timeout: Some(Duration::from_millis(50)), + probe_retry_interval: Duration::from_millis(20), + ..settings() + }, + ) + .unwrap(); + + // Unreachable broker, so the worker keeps retrying. + producer.push_config(&blob("127.0.0.1:1", "9")).unwrap(); + + let deadline = Instant::now() + Duration::from_secs(10); + while producer.inner.probe_attempts.load(Ordering::Relaxed) < 2 { + assert!(Instant::now() < deadline, "worker never retried the probe"); + std::thread::sleep(Duration::from_millis(10)); + } + + // Revert to the running config. + assert_eq!( + producer.push_config(&blob(&servers, "all")).unwrap(), + ReloadOutcome::Unchanged + ); + let after_revert = producer.inner.probe_attempts.load(Ordering::Relaxed); + + // At most one more probe may be in flight when the revert lands; after + // that the worker must be idle rather than still retrying. + std::thread::sleep(Duration::from_millis(400)); + let settled = producer.inner.probe_attempts.load(Ordering::Relaxed); + assert!( + settled <= after_revert + 1, + "worker kept retrying a reverted config: {after_revert} -> {settled} probes" + ); + assert_eq!(producer.generation(), 0); + } + + /// Pushing the same config twice while the first is rolling out must not + /// cost a second drain and swap. + #[test] + fn test_duplicate_push_does_not_swap_twice() { + let cluster = MockCluster::new(1).unwrap(); + let servers = cluster.bootstrap_servers(); + let producer = producer(&servers); + + // Hold the worker in the probe for the first rollout, then push the + // same config again. The duplicate is queued behind a rollout that is + // about to make it a no-op, which is the case the dedup check in + // `reload_to` exists for: without it the worker takes the duplicate + // after the swap and pays for a second build, probe, swap and drain. + let (release, gate) = mpsc::channel(); + *producer.inner.probe_gate.lock() = Some(gate); + + producer.push_config(&blob(&servers, "1")).unwrap(); + wait_until_probing(&producer); + producer.push_config(&blob(&servers, "1")).unwrap(); + + release.send(()).unwrap(); + await_generation(&producer, 1); + + // The duplicate resolved to the config now running, so it was dropped + // rather than triggering another swap. + assert_stays_at_generation(&producer, 1); + } + + #[test] + fn test_clones_share_one_client() { + let cluster = MockCluster::new(1).unwrap(); + let servers = cluster.bootstrap_servers(); + let producer = producer(&servers); + let clone = producer.clone(); + + clone.push_config(&blob(&servers, "1")).unwrap(); + + // The reload is visible through the original handle. + await_generation(&producer, 1); + } + + #[test] + fn test_unkeyed_produce_does_not_wait_for_drain() { + let cluster = MockCluster::new(1).unwrap(); + let servers = cluster.bootstrap_servers(); + let producer = producer(&servers); + let destination = TopicOrPartition::Topic(crate::types::Topic::new("ingest-events")); + + // Stand in for a drain in progress. + let held = producer.block_drain_for_test(); + + let (tx, rx) = mpsc::channel(); + let unkeyed = producer.clone(); + let dest = destination; + std::thread::spawn(move || { + let result = unkeyed.produce( + &dest, + KafkaPayload::new(None, None, Some(b"unkeyed".to_vec())), + ); + let _ = tx.send(result); + }); + + // Returns while the drain is still running. + let result = rx + .recv_timeout(Duration::from_secs(2)) + .expect("unkeyed produce blocked on the drain"); + assert!(result.is_ok()); + drop(held); + } + + #[test] + fn test_keyed_produce_waits_for_drain() { + let cluster = MockCluster::new(1).unwrap(); + let servers = cluster.bootstrap_servers(); + let producer = producer(&servers); + let destination = TopicOrPartition::Topic(crate::types::Topic::new("ingest-events")); + + let held = producer.block_drain_for_test(); + + let (tx, rx) = mpsc::channel(); + let keyed = producer.clone(); + let dest = destination; + std::thread::spawn(move || { + let result = keyed.produce( + &dest, + KafkaPayload::new(Some(b"key".to_vec()), None, Some(b"keyed".to_vec())), + ); + let _ = tx.send(result); + }); + + // Blocked while the drain is in progress. + assert!(rx.recv_timeout(Duration::from_millis(300)).is_err()); + + drop(held); + + // Proceeds once the drain finishes. + let result = rx + .recv_timeout(Duration::from_secs(2)) + .expect("keyed produce did not resume after the drain"); + assert!(result.is_ok()); + } + + #[test] + fn test_keyed_produce_fails_open_when_drain_is_slow() { + let cluster = MockCluster::new(1).unwrap(); + let servers = cluster.bootstrap_servers(); + let producer = ReloadingKafkaProducer::new( + &blob(&servers, "all"), + selector(), + ReloadConfig { + drain_timeout: Duration::from_millis(200), + ..settings() + }, + ) + .unwrap(); + let destination = TopicOrPartition::Topic(crate::types::Topic::new("ingest-events")); + + // A drain that outlasts drain_timeout. + let held = producer.block_drain_for_test(); + + let started = Instant::now(); + let result = producer.produce( + &destination, + KafkaPayload::new(Some(b"key".to_vec()), None, Some(b"keyed".to_vec())), + ); + + // Waited for the timeout, then produced anyway rather than erroring. + assert!(result.is_ok()); + assert!(started.elapsed() >= Duration::from_millis(200)); + drop(held); + } +} From dc7d7d04694b6b9430a443d819ec25a785624ce0 Mon Sep 17 00:00:00 2001 From: Markus Unterwaditzer Date: Mon, 14 Sep 2026 16:57:46 +0200 Subject: [PATCH 2/5] add ignore_key_ordeirng --- .../src/backends/kafka/reloading_producer.rs | 56 ++++++++++++++++++- 1 file changed, 54 insertions(+), 2 deletions(-) diff --git a/rust-arroyo/src/backends/kafka/reloading_producer.rs b/rust-arroyo/src/backends/kafka/reloading_producer.rs index 8290620c..0ce98d58 100644 --- a/rust-arroyo/src/backends/kafka/reloading_producer.rs +++ b/rust-arroyo/src/backends/kafka/reloading_producer.rs @@ -12,6 +12,10 @@ //! unkeyed produce calls, which have no ordering requirement, go straight to //! the new client. //! +//! Applications that key only to spread load, and do not care about the order +//! of same-key messages, can set [`ReloadConfig::ignore_key_ordering`] to skip +//! the wait entirely. +//! //! If the drain outlasts [`ReloadConfig::drain_timeout`], keyed produce calls //! stop waiting rather than blocking the application indefinitely, accepting //! that ordering may break. @@ -68,6 +72,13 @@ pub struct ReloadConfig { /// Probing retries until it succeeds or a newer config supersedes it, /// since an unreachable broker is as likely to be a blip as a bad config. pub probe_retry_interval: Duration, + /// Treat keyed messages like unkeyed ones: never wait for a drain. + /// + /// Set this when the key is only there to spread load across partitions + /// and nothing downstream depends on same-key messages arriving in order. + /// Produce calls then never block on a reload, at the cost of the ordering + /// guarantee described in the module docs. + pub ignore_key_ordering: bool, } impl Default for ReloadConfig { @@ -77,6 +88,7 @@ impl Default for ReloadConfig { jitter: Duration::ZERO, probe_timeout: Some(Duration::from_secs(5)), probe_retry_interval: Duration::from_secs(5), + ignore_key_ordering: false, } } } @@ -667,8 +679,9 @@ impl ArroyoProducer for ReloadingKafkaProducer { payload: KafkaPayload, ) -> Result<(), ProducerError> { // Unkeyed messages have no ordering constraint, so they never wait for - // a drain and go straight to whichever client is live. - if payload.key().is_none() { + // a drain and go straight to whichever client is live. With + // `ignore_key_ordering` the caller says the same holds for keyed ones. + if payload.key().is_none() || self.inner.settings.ignore_key_ordering { let current = self.inner.current.read(); return current.producer.produce(destination, payload); } @@ -766,6 +779,7 @@ mod tests { // MockCluster answers metadata, so probing stays on in tests. probe_timeout: Some(Duration::from_secs(5)), probe_retry_interval: Duration::from_millis(50), + ignore_key_ordering: false, } } @@ -1446,6 +1460,44 @@ mod tests { assert!(result.is_ok()); } + #[test] + fn test_keyed_produce_does_not_wait_when_ordering_is_ignored() { + let cluster = MockCluster::new(1).unwrap(); + let servers = cluster.bootstrap_servers(); + let producer = ReloadingKafkaProducer::new( + &blob(&servers, "all"), + selector(), + ReloadConfig { + // Long enough that waiting would show up as a timeout below. + drain_timeout: Duration::from_secs(30), + ignore_key_ordering: true, + ..settings() + }, + ) + .unwrap(); + let destination = TopicOrPartition::Topic(crate::types::Topic::new("ingest-events")); + + let held = producer.block_drain_for_test(); + + let (tx, rx) = mpsc::channel(); + let keyed = producer.clone(); + let dest = destination; + std::thread::spawn(move || { + let result = keyed.produce( + &dest, + KafkaPayload::new(Some(b"key".to_vec()), None, Some(b"keyed".to_vec())), + ); + let _ = tx.send(result); + }); + + // Returns while the drain is still running, like an unkeyed produce. + let result = rx + .recv_timeout(Duration::from_secs(2)) + .expect("keyed produce waited for the drain despite ignore_key_ordering"); + assert!(result.is_ok()); + drop(held); + } + #[test] fn test_keyed_produce_fails_open_when_drain_is_slow() { let cluster = MockCluster::new(1).unwrap(); From edf2bb1710973ca356eed8b0441f013d1dbb0132 Mon Sep 17 00:00:00 2001 From: Markus Unterwaditzer Date: Mon, 14 Sep 2026 20:27:35 +0200 Subject: [PATCH 3/5] fix example for ignore_key_ordering --- rust-arroyo/examples/reloading_producer.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/rust-arroyo/examples/reloading_producer.rs b/rust-arroyo/examples/reloading_producer.rs index 5e918eae..d56e6a40 100644 --- a/rust-arroyo/examples/reloading_producer.rs +++ b/rust-arroyo/examples/reloading_producer.rs @@ -57,6 +57,9 @@ async fn main() { // typo'd address cannot retire a working producer. probe_timeout: Some(Duration::from_secs(5)), probe_retry_interval: Duration::from_secs(5), + // Relay keys messages to spread them across partitions, not to order + // them, so produce calls need not wait for a drain. + ignore_key_ordering: true, }; let producer = ReloadingKafkaProducer::new(&config_blob("all"), selector, settings) From b1e02ec0acb5e33c8fd50273945e792558f509df Mon Sep 17 00:00:00 2001 From: Markus Unterwaditzer Date: Mon, 21 Sep 2026 11:27:36 +0200 Subject: [PATCH 4/5] implement buffer, add metrics --- rust-arroyo/examples/reloading_producer.rs | 10 +- rust-arroyo/src/backends/kafka/producer.rs | 23 +- .../src/backends/kafka/reloading_producer.rs | 1209 +++++++++-------- 3 files changed, 702 insertions(+), 540 deletions(-) diff --git a/rust-arroyo/examples/reloading_producer.rs b/rust-arroyo/examples/reloading_producer.rs index d56e6a40..efc93d99 100644 --- a/rust-arroyo/examples/reloading_producer.rs +++ b/rust-arroyo/examples/reloading_producer.rs @@ -48,9 +48,12 @@ async fn main() { ProducerSelector::new("ingest-events", "relay").expect_logical_topic("ingest-events"); let settings = ReloadConfig { - // How long a keyed produce waits for the old client to drain before - // giving up on ordering. + // How long the reload worker flushes the old client. Off the produce + // path entirely. drain_timeout: Duration::from_secs(10), + // Cap on messages held during a drain. Unused here since + // `ignore_key_ordering` is set, so nothing is ever buffered. + max_buffered_messages: 10_000, // Spread swaps across the fleet so pods do not all stall at once. jitter: Duration::from_secs(2), // Require the new client to reach a broker before swapping it in, so a @@ -58,7 +61,7 @@ async fn main() { probe_timeout: Some(Duration::from_secs(5)), probe_retry_interval: Duration::from_secs(5), // Relay keys messages to spread them across partitions, not to order - // them, so produce calls need not wait for a drain. + // them, so the swap need not drain the old client first. ignore_key_ordering: true, }; @@ -71,7 +74,6 @@ async fn main() { let worker = producer.clone(); tokio::task::spawn_blocking(move || { for i in 0..20 { - // Keyed messages wait for drains to preserve order; unkeyed ones do not. let key = (i % 2 == 0).then(|| b"some-key".to_vec()); let payload = KafkaPayload::new(key, None, Some(format!("message-{i}").into_bytes())); diff --git a/rust-arroyo/src/backends/kafka/producer.rs b/rust-arroyo/src/backends/kafka/producer.rs index 365ae642..199e7820 100644 --- a/rust-arroyo/src/backends/kafka/producer.rs +++ b/rust-arroyo/src/backends/kafka/producer.rs @@ -14,7 +14,7 @@ use rdkafka::producer::{ ThreadedProducer, }; use rdkafka::{Message, Statistics}; -use std::time::Duration; +use std::time::{Duration, Instant}; mod statistics; @@ -113,14 +113,23 @@ where self.producer.in_flight_count() } - /// Blocks until every queued message has been delivered, or the timeout - /// expires. + /// Blocks until the queue is delivered or the timeout expires. Messages + /// remain queued on timeout. /// - /// On timeout, messages are left in the queue. Note that dropping the - /// producer purges them rather than delivering them, so a timeout here - /// means those messages are lost unless the caller flushes again. + /// Retries against a wall-clock deadline because rdkafka before 0.38 + /// budgets `flush` by poll count, so early-returning polls can exhaust its + /// timeout before the requested time passes. pub fn flush_for(&self, timeout: Duration) -> Result<(), KafkaError> { - self.producer.flush(timeout) + let deadline = Instant::now() + timeout; + + loop { + let remaining = deadline.saturating_duration_since(Instant::now()); + match self.producer.flush(remaining) { + Ok(()) => return Ok(()), + Err(error) if remaining.is_zero() => return Err(error), + Err(_) => {} + } + } } } diff --git a/rust-arroyo/src/backends/kafka/reloading_producer.rs b/rust-arroyo/src/backends/kafka/reloading_producer.rs index 0ce98d58..aa8ffb56 100644 --- a/rust-arroyo/src/backends/kafka/reloading_producer.rs +++ b/rust-arroyo/src/backends/kafka/reloading_producer.rs @@ -1,41 +1,23 @@ -//! A producer that swaps its underlying Kafka client when config changes. -//! -//! The host pushes an opaque config blob at startup and on changes. Arroyo -//! validates it, detects changes, and drains and replaces the client. +//! A producer that swaps its Kafka client when config changes. //! //! # Ordering across a swap //! -//! Messages already queued on the old client have not been written yet. If a -//! keyed message went to the new client while an older message with the same -//! key sat in the old client's queue, the two could reach the partition out of -//! order. So keyed produce calls wait for the old client to drain, while -//! unkeyed produce calls, which have no ordering requirement, go straight to -//! the new client. -//! -//! Applications that key only to spread load, and do not care about the order -//! of same-key messages, can set [`ReloadConfig::ignore_key_ordering`] to skip -//! the wait entirely. -//! -//! If the drain outlasts [`ReloadConfig::drain_timeout`], keyed produce calls -//! stop waiting rather than blocking the application indefinitely, accepting -//! that ordering may break. +//! A keyed message sent through the new client could overtake an older message +//! queued on the old client. During a swap, produce calls therefore append to a +//! buffer while the old client drains. The new client is then installed and +//! the buffer is replayed under the write lock. Produce does not wait for the +//! drain. //! -//! # Messages lost on a slow drain +//! Unkeyed messages also buffer. Sending them to the old client would prevent +//! its queue from emptying, while sending them to the new client would bypass +//! the buffered keyed messages. //! -//! rdkafka purges queued and in-flight messages when a producer is dropped, so -//! a drain timeout loses messages rather than delivering them late. Size -//! [`ReloadConfig::drain_timeout`] for the queue depth at reload time, and alert -//! on `arroyo.producer.config_reload_purged_messages`. +//! [`ReloadConfig::ignore_key_ordering`] skips the buffer, installs the new +//! client immediately, and drains the old client in the background. //! -//! # Scope -//! -//! This wraps [`KafkaProducer`] only. There is deliberately no reloading -//! counterpart for [`AsyncKafkaProducer`]: the ordering guarantee above is -//! built on a blocking condvar wait made while holding a lock, which an async -//! produce path cannot reuse. An async version would be a sibling type sharing -//! the config and settings here, not a generalization of this one. -//! -//! [`AsyncKafkaProducer`]: super::producer::AsyncKafkaProducer +//! The buffer is capped by [`ReloadConfig::max_buffered_messages`]. A timed-out +//! drain drops messages left on the old client because rdkafka purges them when +//! the client is dropped. use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; @@ -54,30 +36,25 @@ use crate::types::TopicOrPartition; /// Knobs for how reloads are carried out. #[derive(Debug, Clone)] pub struct ReloadConfig { - /// How long to drain the old client and how long keyed produces wait for it. + /// How long to flush the old client before dropping it and purging what + /// remains. Produce calls buffer during this time. pub drain_timeout: Duration, + /// Maximum messages held while the old client drains. Produce calls fail + /// once this is reached. Keep it below `queue.buffering.max.messages`. + pub max_buffered_messages: usize, /// Upper bound on the random delay before a swap. Spreads swaps across a /// fleet so a config change does not stall every pod at once. pub jitter: Duration, - /// How long to wait for the new client to reach a broker before swapping - /// it in. `None` skips the check. - /// - /// Building a client does no I/O, so a config naming an unreachable broker - /// constructs fine. Without this check such a config would retire a - /// working producer and install a dead one. The check runs off the - /// caller's thread, so it delays a swap rather than blocking anyone. + /// How long the worker waits for the new client to reach a broker before + /// swapping it in. `None` skips the check. pub probe_timeout: Option, - /// How long to wait before probing again after a failed probe. - /// - /// Probing retries until it succeeds or a newer config supersedes it, - /// since an unreachable broker is as likely to be a blip as a bad config. + /// Delay between failed broker probes. Probes continue until one succeeds + /// or the config is superseded. pub probe_retry_interval: Duration, - /// Treat keyed messages like unkeyed ones: never wait for a drain. + /// Install the new client immediately and drain the old one afterwards, + /// with no buffering. /// - /// Set this when the key is only there to spread load across partitions - /// and nothing downstream depends on same-key messages arriving in order. - /// Produce calls then never block on a reload, at the cost of the ordering - /// guarantee described in the module docs. + /// Set this only when same-key message order does not matter. pub ignore_key_ordering: bool, } @@ -85,6 +62,7 @@ impl Default for ReloadConfig { fn default() -> Self { Self { drain_timeout: Duration::from_secs(10), + max_buffered_messages: 10_000, jitter: Duration::ZERO, probe_timeout: Some(Duration::from_secs(5)), probe_retry_interval: Duration::from_secs(5), @@ -96,13 +74,8 @@ impl Default for ReloadConfig { /// What a call to [`ReloadingKafkaProducer::push_config`] did. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum ReloadOutcome { - /// The config is valid and differs from what is running, so it has been - /// handed to the reload worker. - /// - /// The swap happens in the background once the config clears the broker - /// probe and jitter, so the new client is not live yet when this - /// returns. Watch [`ReloadingKafkaProducer::generation`] or the - /// `arroyo.producer.config_reload_applied` metric for that. + /// The valid, changed config was handed to the worker. It is not live yet; + /// watch [`ReloadingKafkaProducer::generation`] for the swap. Accepted, /// The blob parsed, but produced the same config as the running client. /// Any config queued but not yet applied is dropped. @@ -112,6 +85,12 @@ pub enum ReloadOutcome { struct Current { producer: Arc, config: KafkaConfig, + /// Set while `producer` is draining for a swap. Produce calls append here + /// instead of enqueueing on a client that is on its way out. + /// + /// The worker changes this only under the write lock, so a produce call + /// sees either a usable client or the buffer. + buffer: Option>>, } /// A producer whose configuration can be replaced while it is running. @@ -120,9 +99,7 @@ struct Current { #[derive(Clone)] pub struct ReloadingKafkaProducer { inner: Arc, - /// Held only for its `Drop`: when the last handle goes away, this stops - /// the reload worker. The worker holds a `Weak` instead of a handle, - /// so it does not keep this alive. + /// Its final `Drop` stops the worker, which holds only a weak reference. _shutdown: Arc, } @@ -139,14 +116,9 @@ impl Drop for ShutdownOnDrop { } struct Inner { - /// Read on every produce, write only during a swap. + /// Read on every produce, taken for writing only to start buffering and + /// to install. current: RwLock, - /// Number of drains in progress, normally zero or one. Keyed produce calls - /// wait for it to reach zero. Separate from `current` so produce calls are - /// not blocked merely to read the pointer. - drains_in_flight: Mutex, - /// Signalled when `drains_in_flight` reaches zero. - drain_done: Condvar, /// Serializes reloads and holds the state a swap needs. reload: Mutex, /// Signalled when `reload.desired` or `reload.shutdown` changes. @@ -154,45 +126,43 @@ struct Inner { selector: ProducerSelector, settings: ReloadConfig, generation: AtomicU64, - /// Test-only: keyed enqueues without the ordering proof. + /// `client.id`, for metric tags. Resolved once: the selector this producer + /// serves is fixed, so the name is too. + producer_name: String, + /// Test-only queue depth when the old client was swapped out. #[cfg(test)] - enqueued_during_drain: AtomicU64, - /// Test-only: keyed enqueues with the ordering proof. + queue_depth_at_swap: std::sync::atomic::AtomicI32, + /// Test-only messages purged when the retired client was dropped. #[cfg(test)] - keyed_enqueues: AtomicU64, - /// Test-only: blocks the worker inside `probe` so a test can make a push - /// land while a rollout is provably in flight, instead of racing a sleep - /// against a probe that finishes in microseconds. + purged_at_drop: AtomicU64, + /// Test-only calls to `drain`. + #[cfg(test)] + drains: AtomicU64, + /// Test-only gate that holds the worker inside `probe`. #[cfg(test)] probe_gate: Mutex>>, - /// Test-only: counts probe attempts, so a test can tell whether the worker - /// abandoned a config or is still retrying it. + /// Test-only probe attempt count. #[cfg(test)] probe_attempts: AtomicU64, - /// Test-only: parks the worker after the probe and jitter, immediately - /// before the install decision, so a test can land a push inside the exact - /// window the atomic re-check in `swap` guards. In production that window - /// is the jitter sleep; the gate makes it wide deterministically where - /// jitter makes it wide randomly. + /// Test-only gate immediately before `install`. /// - /// This must sit between the jitter sleep and the `reload` lock. Moving it - /// inside the critical section would quietly weaken the test that depends - /// on it into a duplicate of `test_revert_cancels_an_in_flight_rollout`. + /// This must sit after the drain and before `install` takes `reload`. + /// Moving it earlier lets `swap` catch the push and weakens the test into a + /// duplicate of `test_revert_cancels_an_in_flight_rollout`. #[cfg(test)] swap_gate: Mutex>>, - /// Test-only: rollouts the worker has finished, however they ended, so a - /// test can wait for the outcome to be settled instead of sleeping. + /// Test-only completed rollout count. #[cfg(test)] rollouts_finished: AtomicU64, } struct ReloadState { - /// The config the host wants instead of the current one. - /// - /// Each push overwrites it; a revert clears it. The worker clones it and - /// leaves it in place while working, so either change cancels an in-flight - /// rollout before it can swap. + /// Desired config. A push overwrites it and a revert clears it, which can + /// cancel an in-flight rollout. desired: Option, + /// When `desired` was last set, to measure how long a config took to go + /// live. Reset on overwrite, so it always belongs to the current `desired`. + desired_since: Option, /// Set when every handle is dropped, to stop the worker. shutdown: bool, } @@ -200,15 +170,33 @@ struct ReloadState { impl ReloadingKafkaProducer { /// Builds a producer from an initial config blob. /// - /// The blob is opaque to the caller: pass the bytes from options-automator - /// straight through. `selector` says which topic and application to resolve - /// config for. + /// Pass the opaque config bytes through. `selector` chooses the topic and + /// application config. pub fn new( blob: &[u8], selector: ProducerSelector, settings: ReloadConfig, ) -> Result { - let config = ConfigBlob::parse(blob)?.producer_config(&selector)?; + let parsed = ConfigBlob::parse(blob).and_then(|blob| blob.producer_config(&selector)); + + let config = match parsed { + Ok(config) => config, + Err(error) => { + // Tagged `initial` because a blob that is bad at startup and one + // that goes bad later are different problems: the first fails a + // deploy, the second leaves a pod running stale config. + metrics::counter!( + "arroyo.producer.config_reload_rejected", + "topic" => selector.topic().to_owned(), + "producer_name" => UNNAMED_PRODUCER, + "source" => "initial", + "reason" => blob_error_reason(&error), + ) + .increment(1); + return Err(error.into()); + } + }; + Self::from_config(config, selector, settings) } @@ -221,29 +209,44 @@ impl ReloadingKafkaProducer { selector: ProducerSelector, settings: ReloadConfig, ) -> Result { + let producer_name = producer_name_of(&config); let producer = Arc::new(KafkaProducer::new(config.clone())?); tracing::info!( topic = %selector.topic(), app = %selector.app(), + producer_name = %producer_name, "loaded initial producer config" ); + metrics::counter!( + "arroyo.producer.config_reload_loaded", + "topic" => selector.topic().to_owned(), + "producer_name" => producer_name.clone(), + ) + .increment(1); + let inner = Arc::new(Inner { - current: RwLock::new(Current { producer, config }), - drains_in_flight: Mutex::new(0), - drain_done: Condvar::new(), + current: RwLock::new(Current { + producer, + config, + buffer: None, + }), reload: Mutex::new(ReloadState { desired: None, + desired_since: None, shutdown: false, }), reload_wake: Condvar::new(), selector, settings, generation: AtomicU64::new(0), + producer_name, + #[cfg(test)] + queue_depth_at_swap: std::sync::atomic::AtomicI32::new(0), #[cfg(test)] - enqueued_during_drain: AtomicU64::new(0), + purged_at_drop: AtomicU64::new(0), #[cfg(test)] - keyed_enqueues: AtomicU64::new(0), + drains: AtomicU64::new(0), #[cfg(test)] probe_gate: Mutex::new(None), #[cfg(test)] @@ -274,14 +277,36 @@ impl ReloadingKafkaProducer { /// Hands a new config blob to the producer. /// - /// Call when the blob changes. Unchanged blobs do nothing; invalid blobs - /// leave the current client untouched. - /// - /// This parses and hands off without blocking. [`ReloadOutcome::Accepted`] - /// means rollout started, not that the config is live; watch - /// [`ReloadingKafkaProducer::generation`]. Broker work happens on the worker. + /// Parses and hands a changed blob to the worker. Invalid blobs leave the + /// current client untouched. [`ReloadOutcome::Accepted`] does not mean the + /// config is live; watch [`ReloadingKafkaProducer::generation`]. pub fn push_config(&self, blob: &[u8]) -> Result { - let config = ConfigBlob::parse(blob)?.producer_config(&self.inner.selector)?; + let parsed = + ConfigBlob::parse(blob).and_then(|blob| blob.producer_config(&self.inner.selector)); + + let config = match parsed { + Ok(config) => config, + Err(error) => { + // Named after the running client: a blob that does not parse + // has no name of its own, and what an operator needs to know is + // which producer is now stuck on stale config. + metrics::counter!( + "arroyo.producer.config_reload_rejected", + "topic" => self.inner.selector.topic().to_owned(), + "producer_name" => self.inner.producer_name.clone(), + "source" => "push", + "reason" => blob_error_reason(&error), + ) + .increment(1); + tracing::error!( + topic = %self.inner.selector.topic(), + %error, + "could not resolve the pushed config blob, staying on the running config" + ); + return Err(error.into()); + } + }; + Ok(self.inner.set_desired(config)) } @@ -291,20 +316,34 @@ impl ReloadingKafkaProducer { self.inner.generation.load(Ordering::SeqCst) } - /// Messages accepted but not yet delivered by the live client. + /// Messages accepted but not yet delivered, including any held while a + /// reload drains the old client. pub fn in_flight_count(&self) -> i32 { - self.inner.current.read().producer.in_flight_count() - } + let current = self.inner.current.read(); + let buffered = current + .buffer + .as_ref() + .map_or(0, |buffer| buffer.lock().len()); - /// Pretends a drain is in progress, so tests can observe how produce calls - /// behave during one without racing a real flush. - #[cfg(test)] - fn block_drain_for_test(&self) -> DrainGuard { - DrainGuard::new(self.inner.clone()) + current + .producer + .in_flight_count() + .saturating_add(buffered.try_into().unwrap_or(i32::MAX)) } } impl Inner { + /// Records whether a rollout is pending. Call under `reload` after each + /// change to `desired` so the gauge cannot drift. + fn record_pending(&self, reload: &ReloadState) { + metrics::gauge!( + "arroyo.producer.config_reload_pending", + "topic" => self.selector.topic().to_owned(), + "producer_name" => self.producer_name.clone(), + ) + .set(u8::from(reload.desired.is_some())); + } + /// Records the config the host wants running and wakes the worker. fn set_desired(self: &Arc, config: KafkaConfig) -> ReloadOutcome { let mut reload = self.reload.lock(); @@ -313,33 +352,35 @@ impl Inner { if config_eq(&config, &self.current.read().config) { // Clearing `desired` cancels any rollout back to this current config. reload.desired = None; + reload.desired_since = None; self.reload_wake.notify_all(); + self.record_pending(&reload); metrics::counter!( "arroyo.producer.config_reload_unchanged", "topic" => self.selector.topic().to_owned(), + "producer_name" => self.producer_name.clone(), ) .increment(1); return ReloadOutcome::Unchanged; } - // Overwrite so only the newest config remains. An identical push leaves - // an in-flight rollout undisturbed. + // Overwrite so only the newest config remains. reload.desired = Some(config); + reload.desired_since = Some(Instant::now()); self.reload_wake.notify_all(); + self.record_pending(&reload); metrics::counter!( "arroyo.producer.config_reload_accepted", "topic" => self.selector.topic().to_owned(), + "producer_name" => self.producer_name.clone(), ) .increment(1); ReloadOutcome::Accepted } - /// Runs the reload worker until every handle is dropped. - /// - /// Owns everything that must not block the application: the broker probe, - /// jitter, the swap, and draining the old client. + /// Runs broker probes and swaps until every handle is dropped. fn run_worker(self: &Arc) { loop { let config = { @@ -364,33 +405,38 @@ impl Inner { #[cfg(test)] self.rollouts_finished.fetch_add(1, Ordering::SeqCst); - // Done with this config, whether it went live or was abandoned. - // Clear it so the worker can go back to sleep, unless a push - // replaced it while we worked, in which case leave the newer one - // for the next iteration. + // Clear the handled config, but preserve a newer one. + // + // Known gap: the lock is released between abandoning a rollout and + // getting here, and a config re-pushed inside that window is + // indistinguishable from the one just handled, so it gets cleared + // and never rolls out despite `push_config` returning `Accepted`. + // Needs a sequence number on `desired` to fix properly. let mut reload = self.reload.lock(); if reload .desired .as_ref() - .is_some_and(|d| config_eq(d, &config)) + .is_some_and(|desired| config_eq(desired, &config)) { reload.desired = None; + reload.desired_since = None; } + self.record_pending(&reload); } } /// Probes a config and swaps to it, retrying until it is healthy or the /// host stops wanting it. fn reload_to(self: &Arc, config: &KafkaConfig) { - // Build once outside the retry loop. Rebuilding would spawn fresh - // rdkafka threads to connect-spam a dead broker on every retry. - // Invalid client config never reaches the probe or affects the current client. + // Build once; rebuilding would spawn rdkafka threads on every retry. let candidate = match KafkaProducer::new(config.clone()) { Ok(producer) => Arc::new(producer), Err(error) => { metrics::counter!( "arroyo.producer.config_reload_rejected", "topic" => self.selector.topic().to_owned(), + "producer_name" => self.producer_name.clone(), + "source" => "push", "reason" => "client_build", ) .increment(1); @@ -403,9 +449,20 @@ impl Inner { } }; + let mut attempts = 0u64; + loop { + attempts += 1; + match self.probe(&candidate) { Ok(()) => { + metrics::histogram!( + "arroyo.producer.config_reload_probe_attempts", + "topic" => self.selector.topic().to_owned(), + "producer_name" => self.producer_name.clone(), + ) + .record(attempts as f64); + // `swap` re-checks that the config is still wanted itself, // atomically with installing the client. self.swap(candidate.clone(), config); @@ -415,19 +472,18 @@ impl Inner { metrics::counter!( "arroyo.producer.config_reload_probe_failed", "topic" => self.selector.topic().to_owned(), + "producer_name" => self.producer_name.clone(), ) .increment(1); tracing::warn!( topic = %self.selector.topic(), %error, + attempts, "new config cannot reach a broker, staying on the running one" ); } } - // Retry until healthy or superseded. An unreachable broker is as - // likely to be a blip as a bad config, and the running producer is - // unaffected while we wait. let mut reload = self.reload.lock(); let interval = self.settings.probe_retry_interval; @@ -441,8 +497,6 @@ impl Inner { self.reload_wake.wait_for(&mut reload, interval); if superseded(&reload, config) { - // The host wants something else, or we are shutting down. - // Abandon this one; the worker loop picks up what is desired. return; } } @@ -471,30 +525,55 @@ impl Inner { } let started = Instant::now(); - // Metadata for the topic we produce to: one round-trip, no message - // written, and it also catches a topic that does not exist. + // Topic metadata probes the broker without producing a message. let result = candidate.validate_topic(crate::types::Topic::new(self.selector.topic()), timeout); metrics::histogram!( "arroyo.producer.config_reload_probe_ms", "topic" => self.selector.topic().to_owned(), + "producer_name" => self.producer_name.clone(), ) .record(started.elapsed().as_millis() as f64); result } - /// Installs an already-probed client and drains the old one. + /// Drains the old client and installs an already-probed one in its place. fn swap(self: &Arc, new_producer: Arc, config: &KafkaConfig) { - // Jitter before claiming anything, so a fleet does not swap in unison. + // Spread swaps across instances. if !self.settings.jitter.is_zero() { let jitter = rand::random::() * self.settings.jitter.as_secs_f64(); std::thread::sleep(Duration::from_secs_f64(jitter)); } - // Test-only: park here, after the jitter and before the install - // decision, so a test can land a push inside the guarded window. + // Bail out before doing any work if the host already wants something + // else. `install` checks again, atomically with the install itself; + // this one just avoids buffering and draining for nothing. + if superseded(&self.reload.lock(), config) { + self.record_superseded(0); + return; + } + + // Start buffering, so the old client stops receiving messages and its + // queue can actually reach zero. Produce calls append to the buffer + // from here until `install` takes it away again. + // + // Skipped when ordering does not matter: the new client then goes in + // straight away and the old one drains behind it. + if !self.settings.ignore_key_ordering { + let mut current = self.current.write(); + current.buffer = Some(Mutex::new(Vec::new())); + let draining = current.producer.clone(); + drop(current); + + // Nothing feeds this client now, so the flush can converge. + self.drain(&draining, self.settings.drain_timeout); + } + + // Test-only: park here, with the buffer live and the old client + // drained, so a test can land a push in the window that the atomic + // re-check inside `install` guards. #[cfg(test)] { let gate = self.swap_gate.lock().take(); @@ -503,171 +582,213 @@ impl Inner { } } - // After jitter, hold `reload` across the superseded check and install; - // otherwise a push between them could install an unwanted client. - // Lock order is `reload` then `current`, matching `set_desired`; never - // reverse it. - let reload = self.reload.lock(); + let Some(old) = self.install(new_producer, config) else { + return; + }; + + // Only the unordered path still has a client to drain. On the ordered + // path it was drained before the install, and anything the drain did + // not get to has to be purged rather than delivered: the buffer has + // just gone out on the new client, so delivering older same-key + // messages now is the reorder this whole design exists to prevent. + if self.settings.ignore_key_ordering { + self.drain(&old, self.settings.drain_timeout); + } - if superseded(&reload, config) { + let purged = old.in_flight_count().max(0) as u64; + + #[cfg(test)] + self.purged_at_drop.store(purged, Ordering::SeqCst); + + if purged > 0 { metrics::counter!( - "arroyo.producer.config_reload_superseded", + "arroyo.producer.config_reload_purged_messages", "topic" => self.selector.topic().to_owned(), + "producer_name" => self.producer_name.clone(), ) - .increment(1); - tracing::info!( + .increment(purged); + tracing::error!( topic = %self.selector.topic(), - "config superseded before it went live, not swapping to it" + purged, + "dropping the old client with messages still queued, they are lost" ); - return; } - // Count the drain before the swap becomes visible. A keyed produce - // checks this counter while holding a read guard on `current`, so - // anything that sees the new client also sees the drain and waits. - // - // Decremented when this guard is dropped on the drain thread, even if - // that thread panics. Leaking a count would make every later keyed - // produce wait out the full drain timeout, forever. - let drain_guard = DrainGuard::new(self.clone()); + drop(old); + } + + /// Installs a drained client and returns the one it replaced, or `None` if + /// the config went stale, in which case the buffer goes back to the client + /// that stays. + fn install( + self: &Arc, + new_producer: Arc, + config: &KafkaConfig, + ) -> Option> { + // Hold `reload` across the superseded check and the install; otherwise + // a push between them could install an unwanted client. Lock order is + // `reload` then `current`, matching `set_desired`; never reverse it. + let reload = self.reload.lock(); + let superseded = superseded(&reload, config); + let waited_for = reload.desired_since.map(|since| since.elapsed()); - let old = { + let (old, buffered) = { let mut current = self.current.write(); - std::mem::replace( - &mut *current, - Current { - producer: new_producer, - config: config.clone(), - }, - ) + + // Taking and replaying the buffer under the same guard as the + // install is what keeps the order intact: nothing can enqueue in + // between, so the buffer goes out ahead of anything produced after. + let buffer = current.buffer.take(); + + let old = (!superseded).then(|| { + #[cfg(test)] + self.queue_depth_at_swap + .store(current.producer.in_flight_count(), Ordering::SeqCst); + + current.config = config.clone(); + std::mem::replace(&mut current.producer, new_producer) + }); + + let buffered = match buffer { + Some(buffer) => self.replay(buffer.into_inner(), ¤t.producer), + None => 0, + }; + + (old, buffered) }; drop(reload); - self.generation.fetch_add(1, Ordering::SeqCst); + let Some(old) = old else { + self.record_superseded(buffered); + return None; + }; + + let generation = self.generation.fetch_add(1, Ordering::SeqCst) + 1; metrics::counter!( "arroyo.producer.config_reload_applied", "topic" => self.selector.topic().to_owned(), + "producer_name" => self.producer_name.clone(), ) .increment(1); - let timeout = self.settings.drain_timeout; - let selector_topic = self.selector.topic().to_owned(); - - // Drain inline on the worker to keep `drains_in_flight` at one. Thus a - // keyed produce waits for at most one drain; unkeyed produce is unaffected. - { - let _drain_guard = drain_guard; - - let started = Instant::now(); - let result = old.producer.flush_for(timeout); - let elapsed = started.elapsed(); + metrics::gauge!( + "arroyo.producer.config_reload_generation", + "topic" => self.selector.topic().to_owned(), + "producer_name" => self.producer_name.clone(), + ) + .set(generation as f64); + if let Some(waited_for) = waited_for { metrics::histogram!( - "arroyo.producer.config_reload_drain_ms", - "topic" => selector_topic.clone(), + "arroyo.producer.config_reload_latency_ms", + "topic" => self.selector.topic().to_owned(), + "producer_name" => self.producer_name.clone(), ) - .record(elapsed.as_millis() as f64); + .record(waited_for.as_millis() as f64); + } - match result { - Ok(()) => { - tracing::info!( - topic = %selector_topic, - drain_ms = elapsed.as_millis() as u64, - "drained producer after config reload" - ); - } - Err(error) => { - // Drop purges queued and in-flight messages, so a timeout - // loses them rather than delivering them late. Their callbacks report - // purge errors in `arroyo.producer.produce_status`; alert on - // `arroyo.producer.config_reload_purged_messages`. - let queued = old.producer.in_flight_count(); - metrics::counter!( - "arroyo.producer.config_reload_drain_timeout", - "topic" => selector_topic.clone(), - ) - .increment(1); - metrics::counter!( - "arroyo.producer.config_reload_purged_messages", - "topic" => selector_topic.clone(), - ) - .increment(queued.max(0) as u64); - tracing::error!( - topic = %selector_topic, - %error, - queued, - "producer drain timed out after config reload, queued messages will be dropped" - ); - } - } + tracing::info!( + topic = %self.selector.topic(), + producer_name = %self.producer_name, + generation, + latency_ms = waited_for.map(|waited| waited.as_millis() as u64), + buffered, + "config reload applied" + ); - drop(old); - } + Some(old) } - /// Blocks until no drain is in progress, or until the drain timeout - /// expires. Returns `false` if it gave up waiting. - /// - /// On timeout we fail open: producing matters more than ordering, so the - /// caller proceeds and we record that ordering may have been broken. - fn wait_for_drain(&self, deadline: Instant) -> bool { - let mut in_flight = self.drains_in_flight.lock(); - - while *in_flight > 0 { - if self - .drain_done - .wait_until(&mut in_flight, deadline) - .timed_out() - { + fn record_superseded(&self, buffered: usize) { + metrics::counter!( + "arroyo.producer.config_reload_superseded", + "topic" => self.selector.topic().to_owned(), + "producer_name" => self.producer_name.clone(), + ) + .increment(1); + tracing::info!( + topic = %self.selector.topic(), + buffered, + "config superseded before it went live, not swapping to it" + ); + } + + /// Hands buffered messages to `producer`, in the order they were produced. + /// Returns how many there were. + fn replay( + &self, + buffer: Vec<(TopicOrPartition, KafkaPayload)>, + producer: &KafkaProducer, + ) -> usize { + let buffered = buffer.len(); + + metrics::histogram!( + "arroyo.producer.config_reload_buffered_messages", + "topic" => self.selector.topic().to_owned(), + "producer_name" => self.producer_name.clone(), + ) + .record(buffered as f64); + + for (destination, payload) in buffer { + // The produce call that buffered this already returned Ok, so + // there is nobody left to report the error to. + if let Err(error) = producer.produce(&destination, payload) { metrics::counter!( - "arroyo.producer.config_reload_produce_fail_open", + "arroyo.producer.config_reload_replay_failed", "topic" => self.selector.topic().to_owned(), + "producer_name" => self.producer_name.clone(), ) .increment(1); - tracing::warn!( + tracing::error!( topic = %self.selector.topic(), - "keyed produce proceeding before drain finished, ordering may be broken" + %error, + "could not replay a buffered message onto the new client" ); - return false; } } - true + buffered } - /// Records whether a keyed enqueue observed no drain under its `current` - /// guard. Checking first or enqueueing after releasing the guard records no - /// proof; only the deliberate timeout path may do so. - #[cfg(test)] - fn record_keyed_enqueue(&self, ordered: bool) { - if ordered { - self.keyed_enqueues.fetch_add(1, Ordering::Relaxed); - } else { - self.enqueued_during_drain.fetch_add(1, Ordering::Relaxed); - } - } -} + /// Flushes a client that no longer accepts messages. + fn drain(&self, producer: &KafkaProducer, timeout: Duration) { + #[cfg(test)] + self.drains.fetch_add(1, Ordering::SeqCst); -/// Releases the drain count on drop, including during panic. -struct DrainGuard { - inner: Arc, -} + let started = Instant::now(); + let result = producer.flush_for(timeout); + let elapsed = started.elapsed(); -impl DrainGuard { - /// Registers a drain; `Drop` guarantees the matching decrement. - fn new(inner: Arc) -> Self { - *inner.drains_in_flight.lock() += 1; - Self { inner } - } -} + metrics::histogram!( + "arroyo.producer.config_reload_drain_ms", + "topic" => self.selector.topic().to_owned(), + "producer_name" => self.producer_name.clone(), + ) + .record(elapsed.as_millis() as f64); -impl Drop for DrainGuard { - fn drop(&mut self) { - let mut in_flight = self.inner.drains_in_flight.lock(); - *in_flight -= 1; - if *in_flight == 0 { - self.inner.drain_done.notify_all(); + match result { + Ok(()) => tracing::info!( + topic = %self.selector.topic(), + producer_name = %self.producer_name, + drain_ms = elapsed.as_millis() as u64, + "drained old producer for config reload" + ), + Err(error) => { + metrics::counter!( + "arroyo.producer.config_reload_drain_timeout", + "topic" => self.selector.topic().to_owned(), + "producer_name" => self.producer_name.clone(), + ) + .increment(1); + tracing::error!( + topic = %self.selector.topic(), + producer_name = %self.producer_name, + %error, + queued = producer.in_flight_count(), + "old producer did not drain in time, its queued messages will be dropped" + ); + } } } } @@ -678,47 +799,33 @@ impl ArroyoProducer for ReloadingKafkaProducer { destination: &TopicOrPartition, payload: KafkaPayload, ) -> Result<(), ProducerError> { - // Unkeyed messages have no ordering constraint, so they never wait for - // a drain and go straight to whichever client is live. With - // `ignore_key_ordering` the caller says the same holds for keyed ones. - if payload.key().is_none() || self.inner.settings.ignore_key_ordering { - let current = self.inner.current.read(); - return current.producer.produce(destination, payload); - } - - // Take `current` first and hold it through the counter check and enqueue. - // `swap` counts the drain before taking the write lock, so we either see - // the drain and wait, or block the swap until enqueueing on the old client. - // Checking the counter first would let a swap land before the `current` - // read, enqueueing this message on the new client ahead of older ones. - let deadline = Instant::now() + self.inner.settings.drain_timeout; - - loop { - let current = self.inner.current.read(); - let draining = *self.inner.drains_in_flight.lock() > 0; + // The read guard is the whole of the coordination with a reload: the + // swap takes the write lock to start buffering and again to install and + // replay, so this either enqueues on a client that is still live or + // lands in the buffer, never in between. + let current = self.inner.current.read(); - if !draining { - #[cfg(test)] - self.inner.record_keyed_enqueue(true); + let Some(buffer) = ¤t.buffer else { + return current.producer.produce(destination, payload); + }; - // Enqueue under the read guard so `swap` cannot install first. - return current.producer.produce(destination, payload); - } + let mut buffer = buffer.lock(); - // Release the read guard before waiting so the swap can proceed. - drop(current); + if buffer.len() >= self.inner.settings.max_buffered_messages { + metrics::counter!( + "arroyo.producer.config_reload_buffer_full", + "topic" => self.inner.selector.topic().to_owned(), + "producer_name" => self.inner.producer_name.clone(), + ) + .increment(1); + return Err(ProducerError::ProducerFailure { + error: "reload buffer is full".to_owned(), + }); + } - if !self.inner.wait_for_drain(deadline) { - // Timed out. Fail open: produce against whatever is live now. - // The only sanctioned way to enqueue a keyed message without - // the ordering proof. - #[cfg(test)] - self.inner.record_keyed_enqueue(false); + buffer.push((*destination, payload)); - let current = self.inner.current.read(); - return current.producer.produce(destination, payload); - } - } + Ok(()) } } @@ -736,7 +843,29 @@ pub enum ReloadError { WorkerSpawn(#[source] std::io::Error), } -/// Whether shutdown, a newer push, or a revert clearing `desired` cancels this config. +/// Default `client.id` before config is resolved. +const UNNAMED_PRODUCER: &str = "unknown"; + +/// Returns the producer name used in metric tags. +fn producer_name_of(config: &KafkaConfig) -> String { + config + .get_config_value("client.id") + .cloned() + .unwrap_or_else(|| UNNAMED_PRODUCER.to_owned()) +} + +/// Metric tag for why a pushed blob was rejected. Low cardinality on purpose: +/// the offending topic or variable goes in the log line, not the tag. +fn blob_error_reason(error: &ConfigBlobError) -> &'static str { + match error { + ConfigBlobError::Malformed(_) => "malformed", + ConfigBlobError::UnknownTopic { .. } => "unknown_topic", + ConfigBlobError::UnknownCluster { .. } => "unknown_cluster", + ConfigBlobError::MissingEnvVar { .. } => "missing_env_var", + ConfigBlobError::UnexpectedLogicalTopic { .. } => "unexpected_logical_topic", + } +} + fn superseded(reload: &ReloadState, config: &KafkaConfig) -> bool { reload.shutdown || !reload @@ -772,9 +901,26 @@ mod tests { .into_bytes() } + /// Builds a config whose queue does not drain by itself during tests. + fn lingering_blob(servers: &str, acks: &str) -> Vec { + format!( + r#"{{ + "clusters": {{"events": {{"config": {{"bootstrap.servers": "{servers}"}}}}}}, + "topics": {{ + "ingest-events": {{ + "cluster": "events", + "producer_config": {{"acks": "{acks}", "linger.ms": "1000"}} + }} + }} + }}"# + ) + .into_bytes() + } + fn settings() -> ReloadConfig { ReloadConfig { drain_timeout: Duration::from_secs(5), + max_buffered_messages: 10_000, jitter: Duration::ZERO, // MockCluster answers metadata, so probing stays on in tests. probe_timeout: Some(Duration::from_secs(5)), @@ -783,8 +929,6 @@ mod tests { } } - /// Waits for the reload worker to reach `generation`, since swaps are no - /// longer synchronous with `push_config`. #[track_caller] fn await_generation(producer: &ReloadingKafkaProducer, generation: u64) { let deadline = Instant::now() + Duration::from_secs(15); @@ -798,8 +942,16 @@ mod tests { } } - /// Blocks until the worker is inside `probe`, which it signals by taking - /// the gate receiver. Deterministic, unlike sleeping and hoping. + /// Waits for post-install drain and drop work to finish. + #[track_caller] + fn await_rollout(producer: &ReloadingKafkaProducer, rollouts: u64) { + let deadline = Instant::now() + Duration::from_secs(15); + while producer.inner.rollouts_finished.load(Ordering::SeqCst) < rollouts { + assert!(Instant::now() < deadline, "timed out waiting for a rollout"); + std::thread::sleep(Duration::from_millis(5)); + } + } + #[track_caller] fn wait_until_probing(producer: &ReloadingKafkaProducer) { let deadline = Instant::now() + Duration::from_secs(5); @@ -809,7 +961,6 @@ mod tests { } } - /// Asserts no swap happens within a short window. #[track_caller] fn assert_stays_at_generation(producer: &ReloadingKafkaProducer, generation: u64) { std::thread::sleep(Duration::from_millis(300)); @@ -860,7 +1011,6 @@ mod tests { assert!(producer.push_config(br#"{"topics": {}}"#).is_err()); assert_stays_at_generation(&producer, 0); - // Still usable. producer .produce( &TopicOrPartition::Topic(crate::types::Topic::new("ingest-events")), @@ -888,7 +1038,6 @@ mod tests { ); assert_stays_at_generation(&producer, 0); - // Still producing on the original config. producer .produce( &TopicOrPartition::Topic(crate::types::Topic::new("ingest-events")), @@ -916,7 +1065,6 @@ mod tests { producer.push_config(&blob("127.0.0.1:1", "all")).unwrap(); assert_stays_at_generation(&producer, 0); - // A reachable config replaces the one stuck retrying. producer.push_config(&blob(&servers, "1")).unwrap(); await_generation(&producer, 1); assert_eq!( @@ -985,19 +1133,15 @@ mod tests { let (release, gate) = mpsc::channel(); *producer.inner.probe_gate.lock() = Some(gate); - // Roll out "1" and hold the worker in its probe. producer.push_config(&blob(&servers, "1")).unwrap(); wait_until_probing(&producer); - // Queue a further change, then revert to the config being rolled out. producer.push_config(&blob(&servers, "0")).unwrap(); producer.push_config(&blob(&servers, "1")).unwrap(); release.send(()).unwrap(); await_generation(&producer, 1); - // Settles on "1" with no second swap: the "0" was replaced before the - // worker ever took it. assert_stays_at_generation(&producer, 1); assert_eq!( producer @@ -1038,137 +1182,131 @@ mod tests { await_generation(&producer, 1); } - /// The ordering guarantee: a keyed message must never be enqueued while an - /// older client is still draining. - /// - /// The test holds a drain open because MockCluster drains in microseconds, - /// then races a keyed produce with a real swap. It fails if `produce` checks - /// the counter before `current`, which can enqueue on the new client first. + /// The old queue must be empty at the swap or a same-key message on the + /// new client can overtake it. The lingering config makes draining the only + /// way to empty the queue. #[test] - fn test_keyed_produce_never_races_a_swap() { + fn test_old_queue_is_empty_when_the_swap_happens() { let cluster = MockCluster::new(1).unwrap(); let servers = cluster.bootstrap_servers(); - let destination = TopicOrPartition::Topic(crate::types::Topic::new("ingest-events")); - - for round in 0..40 { - let producer = producer(&servers); - - // Hold a drain open so the window is guaranteed to be wide. - let held = producer.block_drain_for_test(); - - // A keyed produce that must wait for that drain. - let (tx, rx) = mpsc::channel(); - let writer = producer.clone(); - std::thread::spawn(move || { - let _ = tx.send(writer.produce( - &destination, - KafkaPayload::new(Some(b"key".to_vec()), None, Some(b"keyed".to_vec())), - )); - }); - - // Let it reach the wait, then swap underneath it. - std::thread::sleep(Duration::from_millis(5)); - producer - .push_config(&blob(&servers, &format!("{round}"))) + let producer = + ReloadingKafkaProducer::new(&lingering_blob(&servers, "all"), selector(), settings()) .unwrap(); - - // Release the drain and let the produce complete. - drop(held); - rx.recv_timeout(Duration::from_secs(5)) - .expect("keyed produce never completed") - .expect("keyed produce failed"); - - assert_eq!( - producer.inner.enqueued_during_drain.load(Ordering::Relaxed), - 0, - "keyed produce landed on a client while an older one was draining" - ); - assert!( - producer.inner.keyed_enqueues.load(Ordering::Relaxed) > 0, - "no keyed enqueues were recorded" - ); - } - } - - /// Unkeyed produce must stay non-blocking even under constant swapping, - /// since it has no ordering constraint to preserve. - #[test] - fn test_unkeyed_produce_never_blocks_under_swaps() { - let cluster = MockCluster::new(1).unwrap(); - let servers = cluster.bootstrap_servers(); - let producer = producer(&servers); let destination = TopicOrPartition::Topic(crate::types::Topic::new("ingest-events")); let stop = Arc::new(std::sync::atomic::AtomicBool::new(false)); - let worst = Arc::new(Mutex::new(Duration::ZERO)); - let writer = { let producer = producer.clone(); let stop = stop.clone(); - let worst = worst.clone(); std::thread::spawn(move || { while !stop.load(std::sync::atomic::Ordering::Relaxed) { - let started = Instant::now(); producer .produce( &destination, - KafkaPayload::new(None, None, Some(b"payload".to_vec())), + KafkaPayload::new(Some(b"key".to_vec()), None, Some(b"keyed".to_vec())), ) .unwrap(); - let elapsed = started.elapsed(); - - let mut worst = worst.lock(); - if elapsed > *worst { - *worst = elapsed; - } + std::thread::sleep(Duration::from_millis(1)); } }) }; - for i in 0..25 { - producer - .push_config(&blob(&servers, &format!("{}", i % 2))) - .unwrap(); - std::thread::sleep(Duration::from_millis(2)); - } + producer + .push_config(&lingering_blob(&servers, "1")) + .unwrap(); + await_generation(&producer, 1); stop.store(true, std::sync::atomic::Ordering::Relaxed); writer.join().unwrap(); - // Nowhere near drain_timeout (5s); it never waited on a drain. - let worst = *worst.lock(); - assert!( - worst < Duration::from_millis(500), - "unkeyed produce blocked for {worst:?}, it should never wait for a drain" + assert_eq!( + producer.inner.queue_depth_at_swap.load(Ordering::SeqCst), + 0, + "swapped out a client that still had messages queued" ); } - /// A panicking drain must not strand the counter, which would make every - /// later keyed produce wait out the full timeout. + /// Produce must not stall during a drain. An unreachable broker forces + /// this drain to run for the full timeout. #[test] - fn test_drain_guard_releases_on_panic() { + fn test_produce_does_not_wait_for_a_drain_that_never_finishes() { let cluster = MockCluster::new(1).unwrap(); - let producer = producer(&cluster.bootstrap_servers()); - - let inner = producer.inner.clone(); - std::thread::spawn(move || { - let _guard = DrainGuard::new(inner); - panic!("drain exploded"); - }) - .join() - .unwrap_err(); + let servers = cluster.bootstrap_servers(); - assert_eq!(*producer.inner.drains_in_flight.lock(), 0); + let producer = ReloadingKafkaProducer::new( + &blob("127.0.0.1:1", "all"), + selector(), + ReloadConfig { + drain_timeout: Duration::from_secs(30), + ..settings() + }, + ) + .unwrap(); + let destination = TopicOrPartition::Topic(crate::types::Topic::new("ingest-events")); - // Keyed produce still proceeds immediately. - let started = Instant::now(); producer .produce( - &TopicOrPartition::Topic(crate::types::Topic::new("ingest-events")), - KafkaPayload::new(Some(b"key".to_vec()), None, Some(b"payload".to_vec())), + &destination, + KafkaPayload::new(Some(b"key".to_vec()), None, Some(b"stuck".to_vec())), ) .unwrap(); - assert!(started.elapsed() < Duration::from_secs(1)); + + producer.push_config(&blob(&servers, "1")).unwrap(); + + // Wait for buffering to come on, so the count below is not polluted by + // produce calls that legitimately went straight to the old client while + // the worker was still probing. + let deadline = Instant::now() + Duration::from_secs(5); + while producer.inner.current.read().buffer.is_none() { + assert!( + Instant::now() < deadline, + "worker never turned buffering on" + ); + std::thread::sleep(Duration::from_millis(5)); + } + + let deadline = Instant::now() + Duration::from_secs(3); + let mut worst = Duration::ZERO; + let mut produced = 0; + + while Instant::now() < deadline { + let started = Instant::now(); + producer + .produce( + &destination, + KafkaPayload::new(Some(b"key".to_vec()), None, Some(b"keyed".to_vec())), + ) + .unwrap(); + worst = worst.max(started.elapsed()); + produced += 1; + std::thread::sleep(Duration::from_millis(1)); + } + + assert!( + worst < Duration::from_millis(500), + "produce blocked for {worst:?} while a drain was in progress" + ); + + // Not blocking is only half of it: those messages must have gone into + // the buffer, not onto the client that is being retired. Without this + // the test passes against an implementation that never buffers at all. + let current = producer.inner.current.read(); + assert_eq!( + current + .buffer + .as_ref() + .expect("buffering went away") + .lock() + .len(), + produced + ); + assert_eq!( + current.producer.in_flight_count(), + 1, + "produce reached the draining client instead of the buffer" + ); + + assert_eq!(producer.generation(), 0); } /// Dropping every handle must stop the reload worker, otherwise each @@ -1199,13 +1337,8 @@ mod tests { } } - /// A push that reverts to the running config must cancel a rollout that is - /// already in flight. - /// - /// The worker must clone `desired` and leave it in place so a revert can - /// clear it before the swap. The probe is gated instead of raced against a - /// sleep because MockCluster answers in microseconds and this path requires - /// a successful probe. + /// Reverting to the running config must cancel an in-flight rollout. The + /// probe gate makes the timing deterministic. #[test] fn test_revert_cancels_an_in_flight_rollout() { let cluster = MockCluster::new(1).unwrap(); @@ -1217,23 +1350,17 @@ mod tests { let (release, gate) = mpsc::channel(); *producer.inner.probe_gate.lock() = Some(gate); - // Reachable config: this probe will succeed once released. producer.push_config(&blob(&servers, "9")).unwrap(); - // The worker is now inside the probe, working on this config while it - // remains recorded as the desired one. wait_until_probing(&producer); - // The rollout is provably in flight. Revert to what is running. assert_eq!( producer.push_config(&blob(&servers, "all")).unwrap(), ReloadOutcome::Unchanged ); - // Let the probe finish successfully. release.send(()).unwrap(); - // The reverted-away config must not go live. assert_stays_at_generation(&producer, 0); assert_eq!( producer @@ -1247,17 +1374,19 @@ mod tests { } /// A revert that lands after the probe passed but before the install must - /// cancel the swap. + /// cancel the swap, and the messages buffered in the meantime must go to + /// the client that stays rather than be dropped with the abandoned one. /// - /// The superseded check and install must be atomic after jitter, or a revert - /// in between is lost. This test sets jitter to zero and parks the worker at - /// the gate between sleep and lock acquisition, so the revert cannot lose a - /// wall-clock race. + /// The superseded check and the install must be atomic, or a revert in + /// between is lost. The worker is parked at the gate after the drain, which + /// is the window that check guards, so the revert cannot lose a wall-clock + /// race. #[test] fn test_revert_between_probe_and_install_cancels_the_swap() { let cluster = MockCluster::new(1).unwrap(); let servers = cluster.bootstrap_servers(); let producer = producer(&servers); + let destination = TopicOrPartition::Topic(crate::types::Topic::new("ingest-events")); // Arm before pushing, so the worker cannot get past the gate first. let (release, gate) = mpsc::channel(); @@ -1265,13 +1394,40 @@ mod tests { producer.push_config(&blob(&servers, "9")).unwrap(); - // Worker parked: probe passed, install not reached. + // Worker parked: probe passed and the old client is drained, but the + // install has not run. let deadline = Instant::now() + Duration::from_secs(5); while producer.inner.swap_gate.lock().is_some() { assert!(Instant::now() < deadline, "worker never reached the swap"); std::thread::sleep(Duration::from_millis(5)); } + // Buffering is on, so these land in the buffer rather than on a client. + for i in 0..4 { + producer + .produce( + &destination, + KafkaPayload::new( + Some(b"key".to_vec()), + None, + Some(format!("msg-{i}").into_bytes()), + ), + ) + .unwrap(); + } + assert_eq!( + producer + .inner + .current + .read() + .buffer + .as_ref() + .expect("worker did not turn buffering on") + .lock() + .len(), + 4 + ); + // The revert lands inside the window. `Unchanged` is guaranteed here, // not raced for: the worker is parked before the install. assert_eq!( @@ -1283,7 +1439,7 @@ mod tests { // Wait until the rollout ended one way or the other, then check which. // Both a correct and a broken implementation reach this point, so the - // generation assertion below is what tells them apart. + // assertions below are what tell them apart. let deadline = Instant::now() + Duration::from_secs(5); while producer.inner.rollouts_finished.load(Ordering::SeqCst) == 0 { assert!( @@ -1298,24 +1454,17 @@ mod tests { 0, "a config reverted before install was swapped in anyway" ); - assert_eq!( - producer - .inner - .current - .read() - .config - .get_config_value("acks"), - Some(&"all".to_string()) + + let current = producer.inner.current.read(); + assert_eq!(current.config.get_config_value("acks"), Some(&"all".into())); + assert!( + current.buffer.is_none(), + "buffering was left on after the rollout was abandoned" ); } - /// After a revert, the worker must stop retrying the reverted-away config. - /// - /// This pins the liveness half of the fix. If the supersede check only - /// looked for a *newer* config rather than comparing against `desired`, a - /// revert would leave nothing to compare against and the worker would keep - /// probing a config the host rolled back. If that broker later recovers, - /// the config goes live long after the fact. + /// After a revert, the worker must stop retrying the old config or it + /// could go live if its broker later recovers. #[test] fn test_revert_stops_the_probe_retry_loop() { let cluster = MockCluster::new(1).unwrap(); @@ -1331,7 +1480,6 @@ mod tests { ) .unwrap(); - // Unreachable broker, so the worker keeps retrying. producer.push_config(&blob("127.0.0.1:1", "9")).unwrap(); let deadline = Instant::now() + Duration::from_secs(10); @@ -1340,7 +1488,6 @@ mod tests { std::thread::sleep(Duration::from_millis(10)); } - // Revert to the running config. assert_eq!( producer.push_config(&blob(&servers, "all")).unwrap(), ReloadOutcome::Unchanged @@ -1369,7 +1516,7 @@ mod tests { // Hold the worker in the probe for the first rollout, then push the // same config again. The duplicate is queued behind a rollout that is // about to make it a no-op, which is the case the dedup check in - // `reload_to` exists for: without it the worker takes the duplicate + // `run_worker` exists for: without it the worker takes the duplicate // after the swap and pays for a second build, probe, swap and drain. let (release, gate) = mpsc::channel(); *producer.inner.probe_gate.lock() = Some(gate); @@ -1399,132 +1546,136 @@ mod tests { await_generation(&producer, 1); } - #[test] - fn test_unkeyed_produce_does_not_wait_for_drain() { - let cluster = MockCluster::new(1).unwrap(); - let servers = cluster.bootstrap_servers(); - let producer = producer(&servers); - let destination = TopicOrPartition::Topic(crate::types::Topic::new("ingest-events")); - - // Stand in for a drain in progress. - let held = producer.block_drain_for_test(); - - let (tx, rx) = mpsc::channel(); - let unkeyed = producer.clone(); - let dest = destination; - std::thread::spawn(move || { - let result = unkeyed.produce( - &dest, - KafkaPayload::new(None, None, Some(b"unkeyed".to_vec())), - ); - let _ = tx.send(result); - }); - - // Returns while the drain is still running. - let result = rx - .recv_timeout(Duration::from_secs(2)) - .expect("unkeyed produce blocked on the drain"); - assert!(result.is_ok()); - drop(held); + fn start_buffering(producer: &ReloadingKafkaProducer) { + producer.inner.current.write().buffer = Some(Mutex::new(Vec::new())); } + /// Produce during a drain must return at once and use the buffer. #[test] - fn test_keyed_produce_waits_for_drain() { + fn test_produce_buffers_while_the_old_client_drains() { let cluster = MockCluster::new(1).unwrap(); - let servers = cluster.bootstrap_servers(); - let producer = producer(&servers); + let producer = producer(&cluster.bootstrap_servers()); let destination = TopicOrPartition::Topic(crate::types::Topic::new("ingest-events")); - let held = producer.block_drain_for_test(); + start_buffering(&producer); - let (tx, rx) = mpsc::channel(); - let keyed = producer.clone(); - let dest = destination; - std::thread::spawn(move || { - let result = keyed.produce( - &dest, - KafkaPayload::new(Some(b"key".to_vec()), None, Some(b"keyed".to_vec())), - ); - let _ = tx.send(result); - }); - - // Blocked while the drain is in progress. - assert!(rx.recv_timeout(Duration::from_millis(300)).is_err()); - - drop(held); + let started = Instant::now(); + for i in 0..3 { + producer + .produce( + &destination, + KafkaPayload::new( + Some(b"key".to_vec()), + None, + Some(format!("msg-{i}").into_bytes()), + ), + ) + .unwrap(); + } + assert!(started.elapsed() < Duration::from_millis(100)); - // Proceeds once the drain finishes. - let result = rx - .recv_timeout(Duration::from_secs(2)) - .expect("keyed produce did not resume after the drain"); - assert!(result.is_ok()); + let current = producer.inner.current.read(); + let buffer = current.buffer.as_ref().expect("buffer went away"); + assert_eq!(buffer.lock().len(), 3); + assert_eq!( + current.producer.in_flight_count(), + 0, + "messages reached the draining client instead of the buffer" + ); } + /// A full buffer must return an error instead of growing without bound. #[test] - fn test_keyed_produce_does_not_wait_when_ordering_is_ignored() { + fn test_buffer_full_is_reported_to_the_caller() { let cluster = MockCluster::new(1).unwrap(); let servers = cluster.bootstrap_servers(); let producer = ReloadingKafkaProducer::new( &blob(&servers, "all"), selector(), ReloadConfig { - // Long enough that waiting would show up as a timeout below. - drain_timeout: Duration::from_secs(30), - ignore_key_ordering: true, + max_buffered_messages: 2, ..settings() }, ) .unwrap(); let destination = TopicOrPartition::Topic(crate::types::Topic::new("ingest-events")); - let held = producer.block_drain_for_test(); + start_buffering(&producer); - let (tx, rx) = mpsc::channel(); - let keyed = producer.clone(); - let dest = destination; - std::thread::spawn(move || { - let result = keyed.produce( - &dest, + let mut produce = || { + producer.produce( + &destination, KafkaPayload::new(Some(b"key".to_vec()), None, Some(b"keyed".to_vec())), - ); - let _ = tx.send(result); - }); + ) + }; - // Returns while the drain is still running, like an unkeyed produce. - let result = rx - .recv_timeout(Duration::from_secs(2)) - .expect("keyed produce waited for the drain despite ignore_key_ordering"); - assert!(result.is_ok()); - drop(held); + assert!(produce().is_ok()); + assert!(produce().is_ok()); + assert!(produce().is_err()); } + /// A swap must drain the client it retires exactly once, before the + /// install. + /// + /// Draining it again afterwards would give whatever the first drain could + /// not deliver a second chance, by which point the buffer has already gone + /// out on the new client. Older same-key messages would then land behind + /// newer ones, which is the reorder the buffer exists to prevent; they have + /// to be purged instead. + #[test] + fn test_retired_client_is_drained_exactly_once() { + let cluster = MockCluster::new(1).unwrap(); + let servers = cluster.bootstrap_servers(); + let producer = producer(&servers); + producer.push_config(&blob(&servers, "1")).unwrap(); + await_rollout(&producer, 1); + + assert_eq!(producer.inner.drains.load(Ordering::SeqCst), 1); + } + + /// `ignore_key_ordering` skips the buffer but must still drain the retired + /// client. The lingering config keeps a message queued until that drain. #[test] - fn test_keyed_produce_fails_open_when_drain_is_slow() { + fn test_ignoring_ordering_still_drains_the_retired_client() { let cluster = MockCluster::new(1).unwrap(); let servers = cluster.bootstrap_servers(); let producer = ReloadingKafkaProducer::new( - &blob(&servers, "all"), + &lingering_blob(&servers, "all"), selector(), ReloadConfig { - drain_timeout: Duration::from_millis(200), + ignore_key_ordering: true, ..settings() }, ) .unwrap(); let destination = TopicOrPartition::Topic(crate::types::Topic::new("ingest-events")); - // A drain that outlasts drain_timeout. - let held = producer.block_drain_for_test(); + producer + .produce( + &destination, + KafkaPayload::new(Some(b"key".to_vec()), None, Some(b"keyed".to_vec())), + ) + .unwrap(); - let started = Instant::now(); - let result = producer.produce( - &destination, - KafkaPayload::new(Some(b"key".to_vec()), None, Some(b"keyed".to_vec())), + producer + .push_config(&lingering_blob(&servers, "1")) + .unwrap(); + + // The drain and the drop happen after the install, so waiting for the + // generation alone would read the counter before it is written. + await_rollout(&producer, 1); + + assert_eq!( + producer.inner.purged_at_drop.load(Ordering::SeqCst), + 0, + "retired the old client without draining it, its queued messages were purged" ); - // Waited for the timeout, then produced anyway rather than erroring. - assert!(result.is_ok()); - assert!(started.elapsed() >= Duration::from_millis(200)); - drop(held); + producer + .produce( + &destination, + KafkaPayload::new(Some(b"key".to_vec()), None, Some(b"keyed".to_vec())), + ) + .unwrap(); } } From 9267fba03e989a1ab5bca1bd8c17bbe9eb88a4bb Mon Sep 17 00:00:00 2001 From: Markus Unterwaditzer Date: Mon, 21 Sep 2026 11:31:26 +0200 Subject: [PATCH 5/5] make probe_timeout non-optiona;l --- rust-arroyo/examples/reloading_producer.rs | 2 +- .../src/backends/kafka/reloading_producer.rs | 30 +++++++++---------- 2 files changed, 16 insertions(+), 16 deletions(-) diff --git a/rust-arroyo/examples/reloading_producer.rs b/rust-arroyo/examples/reloading_producer.rs index efc93d99..05fb95e9 100644 --- a/rust-arroyo/examples/reloading_producer.rs +++ b/rust-arroyo/examples/reloading_producer.rs @@ -58,7 +58,7 @@ async fn main() { jitter: Duration::from_secs(2), // Require the new client to reach a broker before swapping it in, so a // typo'd address cannot retire a working producer. - probe_timeout: Some(Duration::from_secs(5)), + probe_timeout: Duration::from_secs(5), probe_retry_interval: Duration::from_secs(5), // Relay keys messages to spread them across partitions, not to order // them, so the swap need not drain the old client first. diff --git a/rust-arroyo/src/backends/kafka/reloading_producer.rs b/rust-arroyo/src/backends/kafka/reloading_producer.rs index aa8ffb56..f663429d 100644 --- a/rust-arroyo/src/backends/kafka/reloading_producer.rs +++ b/rust-arroyo/src/backends/kafka/reloading_producer.rs @@ -45,9 +45,12 @@ pub struct ReloadConfig { /// Upper bound on the random delay before a swap. Spreads swaps across a /// fleet so a config change does not stall every pod at once. pub jitter: Duration, - /// How long the worker waits for the new client to reach a broker before - /// swapping it in. `None` skips the check. - pub probe_timeout: Option, + /// Deadline for one attempt at reaching a broker with the new client. + /// + /// A client that fails is never installed: the worker retries every + /// `probe_retry_interval` until it passes or the config is superseded. So + /// this bounds a single round trip, not the rollout. + pub probe_timeout: Duration, /// Delay between failed broker probes. Probes continue until one succeeds /// or the config is superseded. pub probe_retry_interval: Duration, @@ -64,7 +67,7 @@ impl Default for ReloadConfig { drain_timeout: Duration::from_secs(10), max_buffered_messages: 10_000, jitter: Duration::ZERO, - probe_timeout: Some(Duration::from_secs(5)), + probe_timeout: Duration::from_secs(5), probe_retry_interval: Duration::from_secs(5), ignore_key_ordering: false, } @@ -507,10 +510,6 @@ impl Inner { /// Client construction does no I/O; this prevents swapping in an /// unreachable producer. fn probe(&self, candidate: &KafkaProducer) -> Result<(), KafkaError> { - let Some(timeout) = self.settings.probe_timeout else { - return Ok(()); - }; - #[cfg(test)] self.probe_attempts.fetch_add(1, Ordering::Relaxed); @@ -526,8 +525,10 @@ impl Inner { let started = Instant::now(); // Topic metadata probes the broker without producing a message. - let result = - candidate.validate_topic(crate::types::Topic::new(self.selector.topic()), timeout); + let result = candidate.validate_topic( + crate::types::Topic::new(self.selector.topic()), + self.settings.probe_timeout, + ); metrics::histogram!( "arroyo.producer.config_reload_probe_ms", @@ -922,8 +923,7 @@ mod tests { drain_timeout: Duration::from_secs(5), max_buffered_messages: 10_000, jitter: Duration::ZERO, - // MockCluster answers metadata, so probing stays on in tests. - probe_timeout: Some(Duration::from_secs(5)), + probe_timeout: Duration::from_secs(5), probe_retry_interval: Duration::from_millis(50), ignore_key_ordering: false, } @@ -1055,7 +1055,7 @@ mod tests { &blob(&servers, "all"), selector(), ReloadConfig { - probe_timeout: Some(Duration::from_millis(300)), + probe_timeout: Duration::from_millis(300), probe_retry_interval: Duration::from_millis(50), ..settings() }, @@ -1473,7 +1473,7 @@ mod tests { &blob(&servers, "all"), selector(), ReloadConfig { - probe_timeout: Some(Duration::from_millis(50)), + probe_timeout: Duration::from_millis(50), probe_retry_interval: Duration::from_millis(20), ..settings() }, @@ -1602,7 +1602,7 @@ mod tests { start_buffering(&producer); - let mut produce = || { + let produce = || { producer.produce( &destination, KafkaPayload::new(Some(b"key".to_vec()), None, Some(b"keyed".to_vec())),