ref(attachments): Prepare Kafka Messages for stored Attachments (#5083)

This is the first step towards storing attachments directly from Relay
into objectstore. No such storing is happening yet, but the kafka
message types are prepared to allow for transmitting the `stored_id` in
the future.

---

Ref FS-102

Author: Swatinem

relay-server/src/services/store.rs

@@ -686,17 +686,19 @@ impl StoreService {
     ) -> Result<Option<ChunkedAttachment>, StoreError> {
         let id = Uuid::new_v4().to_string();
 
-        let mut chunk_index = 0;
         let payload = item.payload();
         let size = item.len();
         let max_chunk_size = self.config.attachment_chunk_size();
 
         // When sending individual attachments, and we have a single chunk, we want to send the
         // `data` inline in the `attachment` message.
         // This avoids a needless roundtrip through the attachments cache on the Sentry side.
-        let data = if send_individual_attachments && size < max_chunk_size {
-            (size > 0).then_some(payload)
+        let payload = if size == 0 {
+            AttachmentPayload::Chunked(0)
+        } else if send_individual_attachments && size < max_chunk_size {
+            AttachmentPayload::Inline(payload)
         } else {
+            let mut chunk_index = 0;
             let mut offset = 0;
             // This skips chunks for empty attachments. The consumer does not require chunks for
             // empty attachments. `chunks` will be `0` in this case.
@@ -717,26 +719,25 @@ impl StoreService {
                 offset += chunk_size;
                 chunk_index += 1;
             }
-            None
-        };
 
-        // The chunk_index is incremented after every loop iteration. After we exit the loop, it
-        // is one larger than the last chunk, so it is equal to the number of chunks.
+            // The chunk_index is incremented after every loop iteration. After we exit the loop, it
+            // is one larger than the last chunk, so it is equal to the number of chunks.
+            AttachmentPayload::Chunked(chunk_index)
+        };
 
         let attachment = ChunkedAttachment {
             id,
             name: match item.filename() {
                 Some(name) => name.to_owned(),
                 None => UNNAMED_ATTACHMENT.to_owned(),
             },
+            rate_limited: item.rate_limited(),
             content_type: item
                 .content_type()
                 .map(|content_type| content_type.as_str().to_owned()),
             attachment_type: item.attachment_type().cloned().unwrap_or_default(),
-            chunks: chunk_index,
-            data,
-            size: Some(size),
-            rate_limited: Some(item.rate_limited()),
+            size,
+            payload,
         };
 
         if send_individual_attachments {
@@ -1341,6 +1342,26 @@ impl Service for StoreService {
     }
 }
 
+/// This signifies how the attachment payload is being transfered.
+#[derive(Debug, Serialize)]
+enum AttachmentPayload {
+    /// The payload has been split into multiple chunks.
+    ///
+    /// The individual chunks are being sent as separate [`AttachmentChunkKafkaMessage`] messages.
+    /// If the payload `size == 0`, the number of chunks will also be `0`.
+    #[serde(rename = "chunks")]
+    Chunked(usize),
+
+    /// The payload is inlined here directly, and thus into the [`ChunkedAttachment`].
+    #[serde(rename = "data")]
+    Inline(Bytes),
+
+    /// The attachment has already been stored into the objectstore, with the given Id.
+    #[serde(rename = "stored_id")]
+    #[allow(unused)] // TODO: actually storing it in objectstore first is still WIP
+    Stored(String),
+}
+
 /// Common attributes for both standalone attachments and processing-relevant attachments.
 #[derive(Debug, Serialize)]
 struct ChunkedAttachment {
@@ -1352,6 +1373,14 @@ struct ChunkedAttachment {
     /// File name of the attachment file.
     name: String,
 
+    /// Whether this attachment was rate limited and should be removed after processing.
+    ///
+    /// By default, rate limited attachments are immediately removed from Envelopes. For processing,
+    /// native crash reports still need to be retained. These attachments are marked with the
+    /// `rate_limited` header, which signals to the processing pipeline that the attachment should
+    /// not be persisted after processing.
+    rate_limited: bool,
+
     /// Content type of the attachment payload.
     #[serde(skip_serializing_if = "Option::is_none")]
     content_type: Option<String>,
@@ -1360,27 +1389,12 @@ struct ChunkedAttachment {
     #[serde(serialize_with = "serialize_attachment_type")]
     attachment_type: AttachmentType,
 
-    /// Number of outlined chunks.
-    /// Zero if the attachment has `size: 0`, or there was only a single chunk which has been inlined into `data`.
-    chunks: usize,
-
-    /// The content of the attachment,
-    /// if they are smaller than the configured `attachment_chunk_size`.
-    #[serde(skip_serializing_if = "Option::is_none")]
-    data: Option<Bytes>,
-
     /// The size of the attachment in bytes.
-    #[serde(skip_serializing_if = "Option::is_none")]
-    size: Option<usize>,
+    size: usize,
 
-    /// Whether this attachment was rate limited and should be removed after processing.
-    ///
-    /// By default, rate limited attachments are immediately removed from Envelopes. For processing,
-    /// native crash reports still need to be retained. These attachments are marked with the
-    /// `rate_limited` header, which signals to the processing pipeline that the attachment should
-    /// not be persisted after processing.
-    #[serde(skip_serializing_if = "Option::is_none")]
-    rate_limited: Option<bool>,
+    /// The attachment payload, chunked, inlined, or already stored.
+    #[serde(flatten)]
+    payload: AttachmentPayload,
 }
 
 /// A hack to make rmp-serde behave more like serde-json when serializing enums.
@@ -1823,18 +1837,13 @@ struct ProfileChunkKafkaMessage {
 #[allow(clippy::large_enum_variant)]
 enum KafkaMessage<'a> {
     Event(EventKafkaMessage),
-    Attachment(AttachmentKafkaMessage),
-    AttachmentChunk(AttachmentChunkKafkaMessage),
     UserReport(UserReportKafkaMessage),
     Metric {
         #[serde(skip)]
         headers: BTreeMap<String, String>,
         #[serde(flatten)]
         message: MetricKafkaMessage<'a>,
     },
-    Profile(ProfileKafkaMessage),
-    ReplayEvent(ReplayEventKafkaMessage<'a>),
-    ReplayRecordingNotChunked(ReplayRecordingNotChunkedKafkaMessage<'a>),
     CheckIn(CheckInKafkaMessage),
     Item {
         #[serde(skip)]
@@ -1850,15 +1859,21 @@ enum KafkaMessage<'a> {
         #[serde(flatten)]
         message: SpanKafkaMessage<'a>,
     },
+
+    Attachment(AttachmentKafkaMessage),
+    AttachmentChunk(AttachmentChunkKafkaMessage),
+
+    Profile(ProfileKafkaMessage),
     ProfileChunk(ProfileChunkKafkaMessage),
+
+    ReplayEvent(ReplayEventKafkaMessage<'a>),
+    ReplayRecordingNotChunked(ReplayRecordingNotChunkedKafkaMessage<'a>),
 }
 
 impl Message for KafkaMessage<'_> {
     fn variant(&self) -> &'static str {
         match self {
             KafkaMessage::Event(_) => "event",
-            KafkaMessage::Attachment(_) => "attachment",
-            KafkaMessage::AttachmentChunk(_) => "attachment_chunk",
             KafkaMessage::UserReport(_) => "user_report",
             KafkaMessage::Metric { message, .. } => match message.name.namespace() {
                 MetricNamespace::Sessions => "metric_sessions",
@@ -1868,24 +1883,26 @@ impl Message for KafkaMessage<'_> {
                 MetricNamespace::Stats => "metric_metric_stats",
                 MetricNamespace::Unsupported => "metric_unsupported",
             },
-            KafkaMessage::Profile(_) => "profile",
-            KafkaMessage::ReplayEvent(_) => "replay_event",
-            KafkaMessage::ReplayRecordingNotChunked(_) => "replay_recording_not_chunked",
             KafkaMessage::CheckIn(_) => "check_in",
             KafkaMessage::Span { .. } => "span",
-            KafkaMessage::ProfileChunk(_) => "profile_chunk",
             KafkaMessage::Item { item_type, .. } => item_type.as_str_name(),
+
+            KafkaMessage::Attachment(_) => "attachment",
+            KafkaMessage::AttachmentChunk(_) => "attachment_chunk",
+
+            KafkaMessage::Profile(_) => "profile",
+            KafkaMessage::ProfileChunk(_) => "profile_chunk",
+
+            KafkaMessage::ReplayEvent(_) => "replay_event",
+            KafkaMessage::ReplayRecordingNotChunked(_) => "replay_recording_not_chunked",
         }
     }
 
     /// Returns the partitioning key for this Kafka message determining.
     fn key(&self) -> Option<relay_kafka::Key> {
         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.
@@ -1894,6 +1911,10 @@ impl Message for KafkaMessage<'_> {
             // recieve the routing_key_hint form their envelopes.
             Self::CheckIn(message) => message.routing_key_hint,
 
+            Self::Attachment(message) => Some(message.event_id.0),
+            Self::AttachmentChunk(message) => Some(message.event_id.0),
+            Self::ReplayEvent(message) => Some(message.replay_id.0),
+
             // Random partitioning
             _ => None,
         }

tests/integration/test_attachments.py

@@ -64,12 +64,12 @@ def test_mixed_attachments_with_processing(
     assert attachment == {
         "type": "attachment",
         "attachment": {
-            "attachment_type": "event.attachment",
-            "chunks": attachment_num_chunks[id1],
             "id": id1,
             "name": "foo.txt",
-            "size": len(chunked_contents),
             "rate_limited": False,
+            "attachment_type": "event.attachment",
+            "size": len(chunked_contents),
+            "chunks": attachment_num_chunks[id1],
         },
         "event_id": event_id,
         "project_id": project_id,
@@ -84,12 +84,11 @@ def test_mixed_attachments_with_processing(
     assert attachment == {
         "type": "attachment",
         "attachment": {
-            "attachment_type": "event.attachment",
-            "chunks": 0,
-            "data": b"hell yeah",
             "name": "bar.txt",
-            "size": len(b"hell yeah"),
             "rate_limited": False,
+            "attachment_type": "event.attachment",
+            "size": len(b"hell yeah"),
+            "data": b"hell yeah",
         },
         "event_id": event_id,
         "project_id": project_id,
@@ -106,11 +105,11 @@ def test_mixed_attachments_with_processing(
     assert attachment == {
         "type": "attachment",
         "attachment": {
-            "attachment_type": "event.attachment",
-            "chunks": 0,
             "name": "foobar.txt",
-            "size": 0,
             "rate_limited": False,
+            "attachment_type": "event.attachment",
+            "size": 0,
+            "chunks": 0,
         },
         "event_id": event_id,
         "project_id": project_id,
@@ -473,13 +472,12 @@ def test_view_hierarchy_processing(
     assert attachment == {
         "type": "attachment",
         "attachment": {
-            "attachment_type": "event.view_hierarchy",
-            "chunks": 0,
-            "data": expected_payload,
-            "content_type": "application/json",
             "name": "Unnamed Attachment",
-            "size": len(expected_payload),
             "rate_limited": False,
+            "content_type": "application/json",
+            "attachment_type": "event.view_hierarchy",
+            "size": len(expected_payload),
+            "data": expected_payload,
         },
         "event_id": event_id,
         "project_id": project_id,
@@ -524,12 +522,12 @@ def test_event_with_attachment(
     assert event_message["attachments"][0].pop("id")
     assert list(event_message["attachments"]) == [
         {
-            "attachment_type": "event.attachment",
-            "chunks": 1,
-            "content_type": "application/octet-stream",
             "name": "Unnamed Attachment",
-            "size": len(b"event attachment"),
             "rate_limited": False,
+            "content_type": "application/octet-stream",
+            "attachment_type": "event.attachment",
+            "size": len(b"event attachment"),
+            "chunks": 1,
         }
     ]
 
@@ -547,13 +545,12 @@ def test_event_with_attachment(
     relay.send_envelope(project_id, envelope)
 
     expected_attachment = {
-        "attachment_type": "event.attachment",
-        "chunks": 0,
-        "content_type": "application/octet-stream",
         "name": "Unnamed Attachment",
+        "rate_limited": False,
+        "content_type": "application/octet-stream",
+        "attachment_type": "event.attachment",
         "size": len(b"transaction attachment"),
         "data": b"transaction attachment",
-        "rate_limited": False,
     }
 
     attachment = attachments_consumer.get_individual_attachment()
@@ -596,12 +593,11 @@ def test_form_data_is_rejected(
     assert attachment == {
         "type": "attachment",
         "attachment": {
-            "attachment_type": "event.attachment",
-            "chunks": 0,
-            "data": b"file content",
             "name": "foo.txt",
-            "size": len(b"file content"),
             "rate_limited": False,
+            "attachment_type": "event.attachment",
+            "size": len(b"file content"),
+            "data": b"file content",
         },
         "event_id": event_id,
         "project_id": project_id,

tests/integration/test_minidump.py

@@ -455,11 +455,11 @@ def test_minidump_with_processing(
         {
             "id": attachment_id,
             "name": "minidump.dmp",
+            "rate_limited": rate_limit == "attachment",
             "attachment_type": "event.minidump",
             "content_type": "application/x-dmp",
-            "chunks": num_chunks,
             "size": len(content),
-            "rate_limited": rate_limit == "attachment",
+            "chunks": num_chunks,
         }
     ]
 
@@ -501,11 +501,11 @@ def test_minidump_with_processing_invalid(
         {
             "id": attachment_id,
             "name": "minidump.dmp",
+            "rate_limited": False,
             "content_type": "application/x-dmp",
             "attachment_type": "event.minidump",
-            "chunks": num_chunks,
             "size": len(content),
-            "rate_limited": False,
+            "chunks": num_chunks,
         }
     ]