Revert "ref(kafka): Make routing key optional, instead of randomizing it" (#4763)

Reverts #4738

Author: onewland

Cargo.lock

@@ -23,6 +23,7 @@ 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,6 +11,7 @@ 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;
@@ -76,7 +77,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) -> Option<[u8; 16]>;
+    fn key(&self) -> [u8; 16];
 
     /// Returns the type of the message.
     fn variant(&self) -> &'static str;
@@ -187,7 +188,7 @@ impl KafkaClient {
     pub fn send(
         &self,
         topic: KafkaTopic,
-        key: Option<[u8; 16]>,
+        key: [u8; 16],
         headers: Option<&BTreeMap<String, String>>,
         variant: &str,
         payload: &[u8],
@@ -304,7 +305,7 @@ impl Producer {
     /// Sends the payload to the correct producer for the current topic.
     fn send(
         &self,
-        key: Option<[u8; 16]>,
+        key: [u8; 16],
         headers: Option<&BTreeMap<String, String>>,
         variant: &str,
         payload: &[u8],
@@ -327,37 +328,28 @@ impl Producer {
             })
             .collect::<KafkaHeaders>();
 
-        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 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"),
+                });
             }
-            (key, _) => key,
-        };
+        }
+
+        let mut record = BaseRecord::to(topic_name).key(&key).payload(payload);
 
-        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,13 +1151,10 @@ impl OutcomeBroker {
             KafkaTopic::Outcomes
         };
 
-        let result = producer.client.send(
-            topic,
-            Some(key.into_bytes()),
-            None,
-            "outcome",
-            payload.as_bytes(),
-        );
+        let result =
+            producer
+                .client
+                .send(topic, key.into_bytes(), None, "outcome", payload.as_bytes());
 
         match result {
             Ok(_) => Ok(()),

relay-server/src/services/outcome.rs

@@ -1595,30 +1595,34 @@ impl Message for KafkaMessage<'_> {
     }
 
     /// Returns the partitioning key for this kafka message determining.
-    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),
+    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,
 
             // 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,
+            Self::CheckIn(message) => message.routing_key_hint.unwrap_or_else(Uuid::nil),
 
             // Random partitioning
             Self::Profile(_)
             | Self::Log { .. }
             | Self::ReplayRecordingNotChunked(_)
             | Self::ProfileChunk(_)
-            | Self::Metric { .. } => None,
+            | Self::Metric { .. } => Uuid::nil(),
+        };
+
+        if uuid.is_nil() {
+            uuid = Uuid::new_v4();
         }
-        .filter(|uuid| !uuid.is_nil())
-        .map(|uuid| uuid.into_bytes())
+
+        *uuid.as_bytes()
     }
 
     fn headers(&self) -> Option<&BTreeMap<String, String>> {
@@ -1715,7 +1719,7 @@ mod tests {
         for topic in [KafkaTopic::Outcomes, KafkaTopic::OutcomesBilling] {
             let res = producer
                 .client
-                .send(topic, Some(*b"0123456789abcdef"), None, "foo", b"");
+                .send(topic, *b"0123456789abcdef", None, "foo", b"");
 
             assert!(matches!(res, Err(ClientError::InvalidTopicName)));
         }