relay_kafka/producer/
mod.rs

1//! This module contains the Kafka producer related code.
2
3use std::borrow::Cow;
4use std::collections::{BTreeMap, HashMap};
5use std::fmt;
6use std::sync::Arc;
7use std::sync::atomic::{AtomicU32, Ordering};
8use std::time::{Duration, Instant};
9
10use rdkafka::message::Header;
11use rdkafka::producer::BaseRecord;
12use relay_statsd::metric;
13use sentry_arroyo::backends::ProducerError;
14use sentry_arroyo::backends::kafka::config_blob::ProducerSelector;
15use sentry_arroyo::backends::kafka::reloading_producer::{
16    ReloadConfig, ReloadError, ReloadingKafkaProducer,
17};
18use sentry_arroyo::types::Topic;
19use thiserror::Error;
20
21use crate::KafkaTopicConfig;
22use crate::config::{KafkaParams, KafkaTopic};
23use crate::debounced::Debounced;
24use crate::limits::KafkaRateLimits;
25use crate::producer::utils::KafkaHeaders;
26use crate::statsd::{KafkaCounters, KafkaDistributions, KafkaGauges};
27
28mod utils;
29use utils::Context;
30
31#[cfg(debug_assertions)]
32mod schemas;
33
34const REPORT_FREQUENCY_SECS: u64 = 1;
35const KAFKA_FETCH_METADATA_TIMEOUT: Duration = Duration::from_secs(30);
36
37/// Key type used for partitioning.
38pub type Key = u128;
39
40/// Kafka producer errors.
41#[derive(Error, Debug)]
42pub enum ClientError {
43    /// Failed to send a Kafka message using Arroyo.
44    #[error("failed to send kafka message")]
45    SendFailed(#[source] ProducerError),
46
47    /// Failed to find configured producer for the requested kafka topic.
48    #[error("failed to find producer for the requested kafka topic")]
49    InvalidTopicName,
50
51    /// Failed to serialize the message.
52    #[error("failed to serialize kafka message")]
53    InvalidMsgPack(#[source] rmp_serde::encode::Error),
54
55    /// Failed to serialize the json message using serde.
56    #[error("failed to serialize json message")]
57    InvalidJson(#[from] serde_json::Error),
58
59    /// Failed to run schema validation on message.
60    #[cfg(debug_assertions)]
61    #[error("schema validation failed")]
62    SchemaValidationFailed(#[source] schemas::SchemaError),
63
64    /// Configuration is wrong and it cannot be used to identify the number of a shard.
65    #[error("no kafka configuration for topic")]
66    MissingTopic,
67
68    /// Failed to validate the topic using Arroyo.
69    #[error("failed to validate the topic with name {0}: {1}")]
70    TopicError(String, #[source] rdkafka::error::KafkaError),
71
72    /// Failed to create a kafka producer, for example because the config blob does not resolve
73    /// for the topic or the resolved config is invalid.
74    #[error("failed to create kafka producer")]
75    Reload(#[from] ReloadError),
76
77    /// Failed to encode the protobuf into the buffer
78    /// because the buffer is too small.
79    #[error("failed to encode protobuf because the buffer is too small")]
80    ProtobufEncodingFailed,
81}
82
83/// Describes the type which can be sent using kafka producer provided by this crate.
84pub trait Message {
85    /// Returns the partitioning key for this kafka message determining.
86    fn key(&self) -> Option<Key>;
87
88    /// Returns the type of the message.
89    fn variant(&self) -> &'static str;
90
91    /// Return the list of headers to be provided when payload is sent to Kafka.
92    fn headers(&self) -> Option<&BTreeMap<String, String>>;
93
94    /// Serializes the message into its binary format.
95    ///
96    /// # Errors
97    /// Returns the [`ClientError::InvalidMsgPack`], [`ClientError::InvalidJson`] or [`ClientError::ProtobufEncodingFailed`]  if the
98    /// serialization failed.
99    fn serialize(&self) -> Result<SerializationOutput<'_>, ClientError>;
100}
101
102/// The output of serializing a message for kafka.
103#[derive(Debug, Clone)]
104pub enum SerializationOutput<'a> {
105    /// Serialized as Json.
106    Json(Cow<'a, [u8]>),
107
108    /// Serialized as MsgPack.
109    MsgPack(Cow<'a, [u8]>),
110
111    /// Serialized as Protobuf.
112    Protobuf(Cow<'a, [u8]>),
113}
114
115impl SerializationOutput<'_> {
116    /// Return the serialized bytes.
117    pub fn as_bytes(&self) -> &[u8] {
118        match self {
119            SerializationOutput::Json(cow) => cow,
120            SerializationOutput::MsgPack(cow) => cow,
121            SerializationOutput::Protobuf(cow) => cow,
122        }
123    }
124}
125
126struct TopicProducers {
127    /// All configured producers.
128    producers: Vec<TopicProducer>,
129}
130
131impl TopicProducers {
132    fn new() -> Self {
133        Self {
134            producers: Vec::new(),
135        }
136    }
137
138    fn select(&self, key: u128) -> Option<&TopicProducer> {
139        debug_assert!(!self.producers.is_empty());
140
141        if self.producers.is_empty() {
142            return None;
143        } else if self.producers.len() == 1 {
144            return self.producers.first();
145        }
146
147        let select = (key % self.producers.len() as u128) as usize;
148        self.producers.get(select)
149    }
150
151    /// Validates the topic by fetching the metadata of the topic directly from Kafka.
152    fn validate_topic(&self) -> Result<(), ClientError> {
153        for tp in &self.producers {
154            tp.producer
155                .validate_topic(Topic::new(&tp.topic_name), KAFKA_FETCH_METADATA_TIMEOUT)
156                .map_err(|error| ClientError::TopicError(tp.topic_name.clone(), error))?;
157        }
158
159        Ok(())
160    }
161}
162
163struct TopicProducer {
164    pub topic_name: String,
165    pub producer: Arc<ReloadingKafkaProducer<Context>>,
166    pub rate_limiter: Option<KafkaRateLimits>,
167}
168
169/// Where producers take their initial Kafka config from, and how they reload it.
170#[derive(Debug, Clone, Copy)]
171pub struct KafkaConfigSource<'a> {
172    /// The Kafka config blob the initial clients are resolved from.
173    pub blob: &'a [u8],
174    /// Which application's overrides in the blob apply. For Relay, `"relay"`.
175    pub app: &'a str,
176    /// Upper bound on the random delay before a swap, to spread swaps across the fleet.
177    pub jitter: Duration,
178}
179
180/// How Relay wants reloads carried out.
181fn reload_config(jitter: Duration) -> ReloadConfig {
182    ReloadConfig {
183        // Relay keys messages to spread them evenly across partitions, not to
184        // order them: an absent key is replaced by a generated one, see
185        // `Producer::next_key`. Nothing downstream depends on two messages
186        // with the same key keeping their order, so buffering produce calls
187        // and draining the old client before installing the new one would
188        // stall produce for no benefit.
189        ignore_key_ordering: true,
190        jitter,
191        ..ReloadConfig::default()
192    }
193}
194
195/// Single kafka producer config with assigned topic.
196struct Producer {
197    /// Topic to producer and rate limiter mappings for sharding.
198    topic_producers: TopicProducers,
199    /// Debouncer for metrics.
200    metrics: Debounced,
201    next_key: AtomicU32,
202}
203
204impl Producer {
205    fn new(topic_producers: TopicProducers) -> Self {
206        Self {
207            topic_producers,
208            metrics: Debounced::new(REPORT_FREQUENCY_SECS),
209            next_key: AtomicU32::new(0),
210        }
211    }
212}
213
214impl Producer {
215    /// Sends the payload to the correct producer for the current topic.
216    fn send(
217        &self,
218        key: Option<Key>,
219        headers: Option<&BTreeMap<String, String>>,
220        variant: &str,
221        payload: &[u8],
222    ) -> Result<&str, ClientError> {
223        let now = Instant::now();
224
225        // Always generate a key to force a slightly more equal distribution across partitions,
226        // see also the documentation for `Self::next_key`.
227        let mut key = key.unwrap_or_else(|| self.next_key());
228
229        let Some(TopicProducer {
230            topic_name,
231            producer,
232            rate_limiter,
233        }) = self.topic_producers.select(key)
234        else {
235            return Err(ClientError::MissingTopic);
236        };
237
238        relay_log::configure_scope(|s| s.set_tag("topic", topic_name));
239
240        let producer_name = producer.context().producer_name();
241
242        metric!(
243            distribution(KafkaDistributions::KafkaMessageSize) = payload.len() as u64,
244            variant = variant,
245            topic = topic_name,
246            producer_name = producer_name
247        );
248
249        let mut headers = headers
250            .unwrap_or(&BTreeMap::new())
251            .iter()
252            .map(|(key, value)| Header {
253                key,
254                value: Some(value),
255            })
256            .collect::<KafkaHeaders>();
257
258        // Always rate limit if there is a rate limiter defined.
259        // Defining a rate limiter for a topic which does not have a consistent routing key is just
260        // a misconfiguration.
261        if let Some(limiter) = &rate_limiter {
262            let is_limited = limiter.try_increment(now, key, 1) < 1;
263
264            if is_limited {
265                metric!(
266                    counter(KafkaCounters::ProducerPartitionKeyRateLimit) += 1,
267                    variant = variant,
268                    topic = topic_name,
269                    producer_name = producer_name
270                );
271
272                headers.insert(Header {
273                    key: "sentry-reshuffled",
274                    value: Some("1"),
275                });
276
277                // Force a 'random' partition, instead the originally assigned partition.
278                key = self.next_key();
279            }
280        }
281
282        let key = u128::to_be_bytes(key);
283        let mut record = BaseRecord::to(topic_name).payload(payload).key(&key);
284        if let Some(headers) = headers.into_inner() {
285            record = record.headers(headers);
286        }
287
288        self.metrics.debounce(now, || {
289            metric!(
290                gauge(KafkaGauges::InFlightCount) = producer.in_flight_count() as u64,
291                variant = variant,
292                topic = topic_name,
293                producer_name = producer_name
294            );
295        });
296
297        producer
298            .produce_record(record)
299            .map_err(ClientError::SendFailed)
300            .inspect_err(|_| {
301                metric!(
302                    counter(KafkaCounters::ProducerEnqueueError) += 1,
303                    variant = variant,
304                    topic = topic_name,
305                    producer_name = producer_name
306                );
307            })?;
308
309        Ok(topic_name)
310    }
311
312    /// Returns a newly generated key.
313    ///
314    /// Keys are created in sequence, they can be used as a Kafka partition key to force an equal
315    /// distribution across partitions.
316    ///
317    /// Producing a Kafka message without a 'random' key, produces slightly uneven batches to
318    /// partitions. We've seen that this pattern does not play well with our Arroyo consumers
319    /// and leads to higher partition lags.
320    fn next_key(&self) -> u128 {
321        // Overflowing is fine, as it would wrap. The only use for the atomic is to generate
322        // a different key to the one before.
323        self.next_key.fetch_add(1, Ordering::Relaxed) as u128
324    }
325}
326
327impl fmt::Debug for Producer {
328    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
329        let topic_names: Vec<&String> = self
330            .topic_producers
331            .producers
332            .iter()
333            .map(|tp| &tp.topic_name)
334            .collect();
335        f.debug_struct("Producer")
336            .field("topic_names", &topic_names)
337            .field("producers", &"<KafkaProducers>")
338            .finish_non_exhaustive()
339    }
340}
341
342/// Keeps all the configured Kafka producers and responsible for the routing of the messages.
343#[derive(Debug)]
344pub struct KafkaClient {
345    producers: HashMap<KafkaTopic, Producer>,
346    #[cfg(debug_assertions)]
347    schema_validator: schemas::Validator,
348}
349
350impl KafkaClient {
351    /// Creates a Kafka client builder.
352    pub fn builder() -> KafkaClientBuilder {
353        KafkaClientBuilder::new()
354    }
355
356    /// Sends message to the provided Kafka topic.
357    ///
358    /// Returns the name of the Kafka topic to which the message was produced.
359    pub fn send_message(
360        &self,
361        topic: KafkaTopic,
362        message: &impl Message,
363    ) -> Result<&str, ClientError> {
364        let serialized = message.serialize()?;
365
366        #[cfg(debug_assertions)]
367        if let SerializationOutput::Json(ref bytes) = serialized {
368            self.schema_validator
369                .validate_message_schema(topic, bytes)
370                .map_err(ClientError::SchemaValidationFailed)?;
371        }
372        self.send(
373            topic,
374            message.key(),
375            message.headers(),
376            message.variant(),
377            serialized.as_bytes(),
378        )
379    }
380
381    /// Sends the payload to the correct producer for the current topic.
382    ///
383    /// Returns the name of the Kafka topic to which the message was produced.
384    fn send(
385        &self,
386        topic: KafkaTopic,
387        key: Option<Key>,
388        headers: Option<&BTreeMap<String, String>>,
389        variant: &str,
390        payload: &[u8],
391    ) -> Result<&str, ClientError> {
392        let producer = self
393            .producers
394            .get(&topic)
395            .ok_or_else(|| ClientError::InvalidTopicName)?;
396
397        producer.send(key, headers, variant, payload)
398    }
399
400    /// Hands a new Kafka config blob to every topic's producer.
401    ///
402    /// A blob that does not resolve for a topic is logged and counted rather
403    /// than returned; that topic keeps the client it is running.
404    ///
405    /// A successful push only means the rollout started. Arroyo probes a
406    /// broker before swapping; watch `arroyo.producer.config_reload_generation` for the swap.
407    pub fn push_kafka_config(&self, blob: &[u8]) {
408        for producer in self.producers.values() {
409            for tp in &producer.topic_producers.producers {
410                let producer_name = tp.producer.context().producer_name();
411
412                match tp.producer.push_config(blob) {
413                    Ok(_) => metric!(
414                        counter(KafkaCounters::ConfigReloadPushed) += 1,
415                        topic = &tp.topic_name,
416                        producer_name = producer_name
417                    ),
418                    Err(error) => {
419                        relay_log::error!(
420                            error = &error as &dyn std::error::Error,
421                            tags.topic = tp.topic_name.as_str(),
422                            "could not apply the kafka config, staying on the running one",
423                        );
424                        metric!(
425                            counter(KafkaCounters::ConfigReloadRejected) += 1,
426                            topic = &tp.topic_name,
427                            producer_name = producer_name
428                        );
429                    }
430                }
431            }
432        }
433    }
434}
435
436/// Helper structure responsible for building the actual [`KafkaClient`].
437#[derive(Default)]
438pub struct KafkaClientBuilder {
439    /// Producers by physical topic name.
440    ///
441    /// A producer resolves its config out of the blob by physical topic, so it can only be
442    /// shared by logical topics that produce to that same physical topic.
443    reused_producers: BTreeMap<String, Arc<ReloadingKafkaProducer<Context>>>,
444    producers: HashMap<KafkaTopic, Producer>,
445}
446
447impl KafkaClientBuilder {
448    /// Creates an empty Kafka client builder.
449    pub fn new() -> Self {
450        Self::default()
451    }
452
453    /// Adds topic configuration to the current [`KafkaClientBuilder`], which in return assigns
454    /// dedicates producer to the topic which can will be used to send the messages.
455    ///
456    /// Each physical topic's client is resolved from the config blob in `source`, and is
457    /// replaced whenever a newer blob is pushed in through [`KafkaClient::push_kafka_config`].
458    ///
459    /// # Errors
460    /// Returns [`ClientError::Reload`] if a producer could not be created, for example because
461    /// the blob has no entry for one of the topic's physical topics.
462    pub fn add_kafka_topic_config(
463        mut self,
464        topic: KafkaTopic,
465        topic_config: &KafkaTopicConfig<'_>,
466        validate_topic: bool,
467        source: &KafkaConfigSource<'_>,
468    ) -> Result<Self, ClientError> {
469        let mut topic_producers = TopicProducers::new();
470
471        // Process each shard configuration (one KafkaParams per shard)
472        // We must preserve the original order from the configuration
473        // because hash-based routing depends on shard index positions
474        for params in topic_config.topics() {
475            let KafkaParams {
476                topic_name,
477                config_name,
478                params: config_params,
479                key_rate_limit,
480            } = params;
481
482            let rate_limiter = key_rate_limit.map(|limit| {
483                KafkaRateLimits::new(
484                    limit.limit_per_window,
485                    Duration::from_secs(limit.window_secs),
486                )
487            });
488
489            // Get or create the producer for this physical topic
490            let kafka_producer = if let Some(producer) = self.reused_producers.get(topic_name) {
491                Arc::clone(producer)
492            } else {
493                // Extract producer name from client.id, fallback to config name, then "unknown"
494                let producer_name = config_params
495                    .iter()
496                    .find(|p| p.name == "client.id")
497                    .map(|p| p.value.clone())
498                    .or_else(|| config_name.map(str::to_owned))
499                    .unwrap_or_else(|| "unknown".to_owned());
500
501                let producer = Arc::new(ReloadingKafkaProducer::new_with_context(
502                    source.blob,
503                    ProducerSelector::new(topic_name.clone(), source.app),
504                    reload_config(source.jitter),
505                    Context::new(producer_name),
506                )?);
507
508                self.reused_producers
509                    .insert(topic_name.clone(), Arc::clone(&producer));
510
511                producer
512            };
513
514            topic_producers.producers.push(TopicProducer {
515                topic_name: topic_name.clone(),
516                producer: kafka_producer,
517                rate_limiter,
518            });
519        }
520
521        let producer = Producer::new(topic_producers);
522        if validate_topic {
523            producer.topic_producers.validate_topic()?;
524        }
525        self.producers.insert(topic, producer);
526
527        Ok(self)
528    }
529
530    /// Consumes self and returns the built [`KafkaClient`].
531    pub fn build(self) -> KafkaClient {
532        KafkaClient {
533            producers: self.producers,
534            #[cfg(debug_assertions)]
535            schema_validator: schemas::Validator::default(),
536        }
537    }
538}
539
540impl fmt::Debug for KafkaClientBuilder {
541    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
542        f.debug_struct("KafkaClientBuilder")
543            .field("reused_producers", &"<CachedProducers>")
544            .field("producers", &self.producers)
545            .finish()
546    }
547}
548
549#[cfg(test)]
550mod tests {
551    use rdkafka::mocking::MockCluster;
552
553    use super::*;
554    use crate::config::{KafkaConfigParam, TopicAssignment};
555
556    /// A blob covering the topics the tests produce to, in the shape
557    /// [`ConfigSnapshot::kafka_config_blob`] emits.
558    fn blob(servers: &str, topics: &[&str]) -> Vec<u8> {
559        let topics: serde_json::Map<_, _> = topics
560            .iter()
561            .map(|topic| {
562                (
563                    (*topic).to_owned(),
564                    serde_json::json!({"physical_topics": [{"cluster": "default"}]}),
565                )
566            })
567            .collect();
568
569        serde_json::to_vec(&serde_json::json!({
570            "clusters": {"default": {"config": {"bootstrap.servers": servers}}},
571            "topics": topics,
572        }))
573        .unwrap()
574    }
575
576    fn source(blob: &[u8]) -> KafkaConfigSource<'_> {
577        KafkaConfigSource {
578            blob,
579            app: "relay",
580            jitter: Duration::ZERO,
581        }
582    }
583
584    fn params(servers: &str) -> Vec<KafkaConfigParam> {
585        vec![KafkaConfigParam {
586            name: "bootstrap.servers".to_owned(),
587            value: servers.to_owned(),
588        }]
589    }
590
591    fn producers(
592        client: &KafkaClient,
593        topic: KafkaTopic,
594    ) -> Vec<&Arc<ReloadingKafkaProducer<Context>>> {
595        client.producers[&topic]
596            .topic_producers
597            .producers
598            .iter()
599            .map(|tp| &tp.producer)
600            .collect()
601    }
602
603    /// Every shard needs its own producer, since each is a separate physical
604    /// topic with its own entry in the blob.
605    #[test]
606    fn test_sharded_topic_gets_one_producer_per_shard() {
607        let cluster = MockCluster::new(1).unwrap();
608        let servers = cluster.bootstrap_servers();
609        let assignment: TopicAssignment =
610            serde_yaml::from_str(r#"["ingest-events-1", "ingest-events-2"]"#).unwrap();
611        let blob = blob(&servers, &["ingest-events-1", "ingest-events-2"]);
612        let params = params(&servers);
613        let secondary = BTreeMap::new();
614        let configs = assignment.kafka_configs(&params, &secondary).unwrap();
615
616        let client = KafkaClient::builder()
617            .add_kafka_topic_config(KafkaTopic::Events, &configs, false, &source(&blob))
618            .unwrap()
619            .build();
620
621        let producers = producers(&client, KafkaTopic::Events);
622        assert_eq!(producers.len(), 2);
623        // Sharding routes by index, so two shards sharing one producer would
624        // send both to whichever topic that producer resolved.
625        assert!(
626            !Arc::ptr_eq(producers[0], producers[1]),
627            "shards share a producer, so one of them produces to the wrong topic"
628        );
629    }
630
631    /// Logical topics producing to the same physical topic share a producer,
632    /// while different physical topics on the same cluster do not.
633    #[test]
634    fn test_producers_are_shared_by_physical_topic() {
635        let cluster = MockCluster::new(1).unwrap();
636        let servers = cluster.bootstrap_servers();
637        let params = params(&servers);
638        let secondary = BTreeMap::new();
639        let blob = blob(&servers, &["outcomes", "ingest-events"]);
640        let outcomes = TopicAssignment::from("outcomes".to_owned());
641        let events = TopicAssignment::from("ingest-events".to_owned());
642        let outcomes = outcomes.kafka_configs(&params, &secondary).unwrap();
643        let events = events.kafka_configs(&params, &secondary).unwrap();
644
645        let client = KafkaClient::builder()
646            .add_kafka_topic_config(KafkaTopic::Outcomes, &outcomes, false, &source(&blob))
647            .unwrap()
648            .add_kafka_topic_config(
649                KafkaTopic::OutcomesBilling,
650                &outcomes,
651                false,
652                &source(&blob),
653            )
654            .unwrap()
655            .add_kafka_topic_config(KafkaTopic::Events, &events, false, &source(&blob))
656            .unwrap()
657            .build();
658
659        let outcomes = producers(&client, KafkaTopic::Outcomes)[0];
660        let billing = producers(&client, KafkaTopic::OutcomesBilling)[0];
661        let events = producers(&client, KafkaTopic::Events)[0];
662
663        assert!(Arc::ptr_eq(outcomes, billing));
664        assert!(!Arc::ptr_eq(outcomes, events));
665    }
666
667    /// A topic missing from the blob cannot get a producer, so building fails
668    /// instead of producing with a config nobody asked for.
669    #[test]
670    fn test_topic_missing_from_blob_is_an_error() {
671        let cluster = MockCluster::new(1).unwrap();
672        let servers = cluster.bootstrap_servers();
673        let blob = blob(&servers, &["ingest-events"]);
674        let params = params(&servers);
675        let secondary = BTreeMap::new();
676        let assignment = TopicAssignment::from("ingest-transactions".to_owned());
677        let configs = assignment.kafka_configs(&params, &secondary).unwrap();
678
679        let result = KafkaClient::builder().add_kafka_topic_config(
680            KafkaTopic::Transactions,
681            &configs,
682            false,
683            &source(&blob),
684        );
685
686        assert!(matches!(result, Err(ClientError::Reload(_))));
687    }
688}
relay_kafka/producer/mod.rs

relay_kafka/producer/
mod.rs
