ref(kafka): Make routing key optional, instead of randomizing it

Author: Dav1dde

Cargo.lock

@@ -23,7 +23,6 @@ serde_json = { workspace = true, optional = true }
 thiserror = { workspace = true }
 sentry-kafka-schemas = { workspace = true, default-features = false, optional = true }
 parking_lot = { workspace = true }
-uuid = { workspace = true }
 hashbrown = { workspace = true }
 
 [dev-dependencies]

relay-kafka/Cargo.toml

@@ -11,7 +11,6 @@ use rdkafka::message::Header;
 use rdkafka::producer::{BaseRecord, Producer as _};
 use relay_statsd::metric;
 use thiserror::Error;
-use uuid::Uuid;
 
 use crate::config::{KafkaParams, KafkaTopic};
 use crate::debounced::Debounced;
@@ -77,7 +76,7 @@ pub enum ClientError {
 /// Describes the type which can be sent using kafka producer provided by this crate.
 pub trait Message {
     /// Returns the partitioning key for this kafka message determining.
-    fn key(&self) -> [u8; 16];
+    fn key(&self) -> Option<[u8; 16]>;
 
     /// Returns the type of the message.
     fn variant(&self) -> &'static str;
@@ -188,7 +187,7 @@ impl KafkaClient {
     pub fn send(
         &self,
         topic: KafkaTopic,
-        key: [u8; 16],
+        key: Option<[u8; 16]>,
         headers: Option<&BTreeMap<String, String>>,
         variant: &str,
         payload: &[u8],
@@ -305,7 +304,7 @@ impl Producer {
     /// Sends the payload to the correct producer for the current topic.
     fn send(
         &self,
-        key: [u8; 16],
+        key: Option<[u8; 16]>,
         headers: Option<&BTreeMap<String, String>>,
         variant: &str,
         payload: &[u8],
@@ -328,28 +327,37 @@ impl Producer {
             })
             .collect::<KafkaHeaders>();
 
-        let mut key = key;
-        if let Some(ref limiter) = self.rate_limiter {
-            if limiter.try_increment(now, key, 1) < 1 {
-                metric!(
-                    counter(KafkaCounters::ProducerPartitionKeyRateLimit) += 1,
-                    variant = variant,
-                    topic = topic_name,
-                );
-
-                key = Uuid::new_v4().into_bytes();
-                headers.insert(Header {
-                    key: "sentry-reshuffled",
-                    value: Some("1"),
-                });
+        let key = match (key, self.rate_limiter.as_ref()) {
+            (Some(key), Some(limiter)) => {
+                let is_limited = limiter.try_increment(now, key, 1) < 1;
+
+                if is_limited {
+                    metric!(
+                        counter(KafkaCounters::ProducerPartitionKeyRateLimit) += 1,
+                        variant = variant,
+                        topic = topic_name,
+                    );
+
+                    headers.insert(Header {
+                        key: "sentry-reshuffled",
+                        value: Some("1"),
+                    });
+
+                    None
+                } else {
+                    Some(key)
+                }
             }
-        }
-
-        let mut record = BaseRecord::to(topic_name).key(&key).payload(payload);
+            (key, _) => key,
+        };
 
+        let mut record = BaseRecord::to(topic_name).payload(payload);
         if let Some(headers) = headers.into_inner() {
             record = record.headers(headers);
         }
+        if let Some(key) = key.as_ref() {
+            record = record.key(key);
+        }
 
         self.metrics.debounce(now, || {
             metric!(

relay-kafka/src/producer/mod.rs

@@ -1151,10 +1151,13 @@ impl OutcomeBroker {
             KafkaTopic::Outcomes
         };
 
-        let result =
-            producer
-                .client
-                .send(topic, key.into_bytes(), None, "outcome", payload.as_bytes());
+        let result = producer.client.send(
+            topic,
+            Some(key.into_bytes()),
+            None,
+            "outcome",
+            payload.as_bytes(),
+        );
 
         match result {
             Ok(_) => Ok(()),

relay-server/src/services/outcome.rs

@@ -1595,34 +1595,30 @@ impl Message for KafkaMessage<'_> {
     }
 
     /// Returns the partitioning key for this kafka message determining.
-    fn key(&self) -> [u8; 16] {
-        let mut uuid = match self {
-            Self::Event(message) => message.event_id.0,
-            Self::Attachment(message) => message.event_id.0,
-            Self::AttachmentChunk(message) => message.event_id.0,
-            Self::UserReport(message) => message.event_id.0,
-            Self::ReplayEvent(message) => message.replay_id.0,
-            Self::Span { message, .. } => message.trace_id.0,
+    fn key(&self) -> Option<[u8; 16]> {
+        match self {
+            Self::Event(message) => Some(message.event_id.0),
+            Self::Attachment(message) => Some(message.event_id.0),
+            Self::AttachmentChunk(message) => Some(message.event_id.0),
+            Self::UserReport(message) => Some(message.event_id.0),
+            Self::ReplayEvent(message) => Some(message.replay_id.0),
+            Self::Span { message, .. } => Some(message.trace_id.0),
 
             // Monitor check-ins use the hinted UUID passed through from the Envelope.
             //
             // XXX(epurkhiser): In the future it would be better if all KafkaMessage's would
             // recieve the routing_key_hint form their envelopes.
-            Self::CheckIn(message) => message.routing_key_hint.unwrap_or_else(Uuid::nil),
+            Self::CheckIn(message) => message.routing_key_hint,
 
             // Random partitioning
             Self::Profile(_)
             | Self::Log { .. }
             | Self::ReplayRecordingNotChunked(_)
             | Self::ProfileChunk(_)
-            | Self::Metric { .. } => Uuid::nil(),
-        };
-
-        if uuid.is_nil() {
-            uuid = Uuid::new_v4();
+            | Self::Metric { .. } => None,
         }
-
-        *uuid.as_bytes()
+        .filter(|uuid| !uuid.is_nil())
+        .map(|uuid| uuid.into_bytes())
     }
 
     fn headers(&self) -> Option<&BTreeMap<String, String>> {
@@ -1719,7 +1715,7 @@ mod tests {
         for topic in [KafkaTopic::Outcomes, KafkaTopic::OutcomesBilling] {
             let res = producer
                 .client
-                .send(topic, *b"0123456789abcdef", None, "foo", b"");
+                .send(topic, Some(*b"0123456789abcdef"), None, "foo", b"");
 
             assert!(matches!(res, Err(ClientError::InvalidTopicName)));
         }