feat(profiling): Route raw profiles through objectstore service

relay-server/src/processing/forward.rs

@@ -57,6 +57,22 @@ impl<'a> StoreHandle<'a> {
         }
     }
 
+    /// Tries to send a message to the [`Objectstore`] service.
+    ///
+    /// Returns the message back if the service is not configured,
+    /// allowing the caller to handle the fallback.
+    pub fn try_send_to_objectstore<M>(&self, message: M) -> Option<M>
+    where
+        Objectstore: FromMessage<M>,
+    {
+        if let Some(objectstore) = self.objectstore {
+            objectstore.send(message);
+            None
+        } else {
+            Some(message)
+        }
+    }
+
     /// Dispatches an envelopes to either the [`Objectstore`] or [`Store`] service.
     pub fn send_envelope(&self, envelope: ManagedEnvelope) {
         use crate::services::store::StoreEnvelope;

relay-server/src/processing/profile_chunks/mod.rs

@@ -230,6 +230,7 @@ impl Forward for ProfileChunkOutput {
         s: processing::forward::StoreHandle<'_>,
         ctx: processing::ForwardContext<'_>,
     ) -> Result<(), Rejected<()>> {
+        use crate::services::objectstore::StoreRawProfile;
         use crate::services::store::StoreProfileChunk;
 
         let expanded = match self {
@@ -243,12 +244,34 @@ impl Forward for ProfileChunkOutput {
         let retention_days = ctx.event_retention().standard;
 
         for chunk in expanded.split(|e| e.chunks) {
-            s.send_to_store(chunk.map(|chunk, _| StoreProfileChunk {
-                retention_days,
-                payload: chunk.payload,
-                quantities: chunk.quantities,
-                raw_profile: chunk.raw_profile,
-            }));
+            if chunk.raw_profile.is_some() {
+                let msg = chunk.map(|chunk, _| {
+                    let raw_profile = chunk.raw_profile.unwrap();
+                    StoreRawProfile {
+                        payload: raw_profile.payload,
+                        content_type: raw_profile.content_type,
+                        store_message: StoreProfileChunk {
+                            retention_days,
+                            payload: chunk.payload,
+                            quantities: chunk.quantities,
+                            raw_profile_object_store_key: None,
+                            raw_profile_content_type: None,
+                        },
+                        retention: retention_days,
+                    }
+                });
+                if let Some(unsent) = s.try_send_to_objectstore(msg) {
+                    s.send_to_store(unsent.map(|profile, _| profile.store_message));
+                }
+            } else {
+                s.send_to_store(chunk.map(|chunk, _| StoreProfileChunk {
+                    retention_days,
+                    payload: chunk.payload,
+                    quantities: chunk.quantities,
+                    raw_profile_object_store_key: None,
+                    raw_profile_content_type: None,
+                }));
+            }
         }
 
         Ok(())

relay-server/src/services/objectstore.rs

@@ -20,13 +20,15 @@ use relay_system::{
 use sentry_protos::snuba::v1::TraceItem;
 
 use crate::constants::DEFAULT_ATTACHMENT_RETENTION;
-use crate::envelope::{Item, ItemType};
+use crate::envelope::{ContentType, Item, ItemType};
 use crate::managed::{
     Counted, Managed, ManagedEnvelope, ManagedResult, OutcomeError, Quantities, Rejected,
 };
 use crate::processing::utils::store::item_id_to_uuid;
 use crate::services::outcome::DiscardReason;
-use crate::services::store::{Store, StoreAttachment, StoreEnvelope, StoreTraceItem};
+use crate::services::store::{
+    Store, StoreAttachment, StoreEnvelope, StoreProfileChunk, StoreTraceItem,
+};
 use crate::services::upload::ByteStream;
 use crate::statsd::{RelayCounters, RelayTimers};
 use crate::utils::{BoundedStream, MeteredStream, RetryableStream, TakeOnce};
@@ -38,6 +40,7 @@ pub enum Objectstore {
     Envelope(StoreEnvelope),
     TraceAttachment(Managed<StoreTraceAttachment>),
     EventAttachment(Managed<StoreAttachment>),
+    RawProfile(Managed<StoreRawProfile>),
     Stream(Stream, Sender<Result<ObjectstoreKey, Error>>),
 }
 
@@ -47,6 +50,7 @@ impl Objectstore {
             Self::Envelope(_) => MessageKind::Envelope,
             Self::TraceAttachment(_) => MessageKind::TraceAttachment,
             Self::EventAttachment(_) => MessageKind::EventAttachment,
+            Self::RawProfile(_) => MessageKind::RawProfile,
             Self::Stream { .. } => MessageKind::Stream,
         }
     }
@@ -60,6 +64,7 @@ impl Objectstore {
                 .count(),
             Self::TraceAttachment(_) => 1,
             Self::EventAttachment(_) => 1,
+            Self::RawProfile(_) => 1,
             Self::Stream { .. } => 1,
         }
     }
@@ -91,12 +96,21 @@ impl FromMessage<Managed<StoreAttachment>> for Objectstore {
     }
 }
 
+impl FromMessage<Managed<StoreRawProfile>> for Objectstore {
+    type Response = NoResponse;
+
+    fn from_message(message: Managed<StoreRawProfile>, _sender: ()) -> Self {
+        Self::RawProfile(message)
+    }
+}
+
 /// A type tag used for logging.
 #[derive(Debug, Clone, Copy)]
 enum MessageKind {
     Envelope,
     EventAttachment,
     TraceAttachment,
+    RawProfile,
     Stream,
 }
 
@@ -106,6 +120,7 @@ impl MessageKind {
             Self::Envelope => "envelope",
             Self::EventAttachment => "attachment",
             Self::TraceAttachment => "attachment_v2",
+            Self::RawProfile => "profiles_raw",
             Self::Stream => "stream",
         }
     }
@@ -143,6 +158,28 @@ impl Counted for StoreTraceAttachment {
     }
 }
 
+/// A raw profile (e.g. Perfetto trace) ready for objectstore upload.
+///
+/// After upload, the [`StoreProfileChunk`] is forwarded to the Store service
+/// with the objectstore key set, so the Kafka message carries a reference
+/// instead of the full binary blob.
+pub struct StoreRawProfile {
+    /// The raw binary profile payload to upload.
+    pub payload: Bytes,
+    /// Content type of the raw profile.
+    pub content_type: ContentType,
+    /// The profile chunk message to forward to Kafka after upload.
+    pub store_message: StoreProfileChunk,
+    /// Data retention in days.
+    pub retention: u16,
+}
+
+impl Counted for StoreRawProfile {
+    fn quantities(&self) -> Quantities {
+        self.store_message.quantities()
+    }
+}
+
 #[derive(Debug, thiserror::Error)]
 #[error("objectstore upload failed")]
 pub struct Error {
@@ -302,12 +339,15 @@ impl ObjectstoreService {
             .with_expiration_policy(ExpirationPolicy::TimeToLive(DEFAULT_ATTACHMENT_RETENTION));
         let trace_attachments = Usecase::new("trace_attachments")
             .with_expiration_policy(ExpirationPolicy::TimeToLive(DEFAULT_ATTACHMENT_RETENTION));
+        let profiles = Usecase::new("profiles_raw")
+            .with_expiration_policy(ExpirationPolicy::TimeToLive(DEFAULT_ATTACHMENT_RETENTION));
 
         let inner = ObjectstoreServiceInner {
             store,
             objectstore_client,
             event_attachments,
             trace_attachments,
+            profiles,
             timeout: Duration::from_secs(*timeout),
             stream_timeout: Duration::from_secs(*stream_timeout),
             retry_interval: Duration::from_secs_f64(*retry_delay),
@@ -344,6 +384,11 @@ impl LoadShed<Objectstore> for ObjectstoreService {
             Objectstore::TraceAttachment(managed) => {
                 let _ = managed.reject_err(error);
             }
+            Objectstore::RawProfile(managed) => {
+                self.inner
+                    .store
+                    .send(managed.map(|profile, _| profile.store_message));
+            }
             Objectstore::Stream(_, sender) => {
                 sender.send(Err(error));
             }
@@ -357,6 +402,7 @@ struct ObjectstoreServiceInner {
     objectstore_client: Client,
     event_attachments: Usecase,
     trace_attachments: Usecase,
+    profiles: Usecase,
     timeout: Duration,
     stream_timeout: Duration,
     retry_interval: Duration,
@@ -381,6 +427,9 @@ impl ObjectstoreServiceInner {
             Objectstore::EventAttachment(attachment) => {
                 self.handle_event_attachment(attachment).await;
             }
+            Objectstore::RawProfile(profile) => {
+                self.handle_raw_profile(profile).await;
+            }
             Objectstore::Stream(stream, sender) => {
                 let result = self.handle_stream(stream).await;
                 if let Err(error) = &result {
@@ -424,6 +473,7 @@ impl ObjectstoreServiceInner {
                             attachment.payload(),
                             retention,
                             None,
+                            None,
                         )
                         .await;
 
@@ -469,6 +519,7 @@ impl ObjectstoreServiceInner {
                         attachment.attachment.payload(),
                         attachment.retention,
                         None,
+                        None,
                     )
                     .await;
 
@@ -532,6 +583,7 @@ impl ObjectstoreServiceInner {
                     body,
                     retention,
                     Some(key),
+                    None,
                 )
                 .await
                 .reject(&trace_item)?;
@@ -546,6 +598,53 @@ impl ObjectstoreServiceInner {
         Ok(())
     }
 
+    async fn handle_raw_profile(&self, managed: Managed<StoreRawProfile>) {
+        let scoping = managed.scoping();
+        let session = self
+            .profiles
+            .for_project(scoping.organization_id.value(), scoping.project_id.value())
+            .session(&self.objectstore_client);
+
+        let payload = managed.payload.clone();
+        let content_type = managed.content_type;
+        let retention = managed.retention;
+
+        let mut store_message = managed.map(|profile, _| profile.store_message);
+
+        match session {
+            Err(error) => Error::from(error).log(MessageKind::RawProfile),
+            Ok(session) if !payload.is_empty() => {
+                let result = self
+                    .upload_bytes(
+                        MessageKind::RawProfile,
+                        &session,
+                        payload,
+                        retention,
+                        None,
+                        Some(content_type),
+                    )
+                    .await;
+
+                match result {
+                    Ok(stored_key) => {
+                        store_message.modify(|msg, _| {
+                            msg.raw_profile_object_store_key = Some(stored_key.into_inner());
+                            msg.raw_profile_content_type = Some(content_type);
+                        });
+                    }
+                    Err(error) => {
+                        error.log(MessageKind::RawProfile);
+                    }
+                }
+            }
+            Ok(_) => {}
+        }
+
+        // Always forward to store even if the raw profile upload failed,
+        // to ensure the kafka message is produced.
+        self.store.send(store_message);
+    }
+
     async fn handle_stream(&self, stream: Stream) -> Result<ObjectstoreKey, Error> {
         let Stream {
             organization_id,
@@ -564,6 +663,7 @@ impl ObjectstoreServiceInner {
             Some(key),
             Body::Stream(TakeOnce::new(stream)),
             None,
+            None,
         )
         .await
     }
@@ -575,10 +675,18 @@ impl ObjectstoreServiceInner {
         payload: Bytes,
         retention: u16,
         key: Option<String>,
+        content_type: Option<ContentType>,
     ) -> Result<ObjectstoreKey, Error> {
         let retention_hours = retention.checked_mul(24);
-        self.upload(kind, session, key, Body::Bytes(payload), retention_hours)
-            .await
+        self.upload(
+            kind,
+            session,
+            key,
+            Body::Bytes(payload),
+            retention_hours,
+            content_type,
+        )
+        .await
     }
 
     async fn upload(
@@ -588,6 +696,7 @@ impl ObjectstoreServiceInner {
         key: Option<String>,
         body: Body,
         retention_hours: Option<u16>,
+        content_type: Option<ContentType>,
     ) -> Result<ObjectstoreKey, Error> {
         let mut attempts = 0;
         let timeout = match &body {
@@ -602,8 +711,15 @@ impl ObjectstoreServiceInner {
                 };
                 attempts += 1;
                 result.replace(
-                    self.attempt_upload(kind, session, key.clone(), body, retention_hours)
-                        .await,
+                    self.attempt_upload(
+                        kind,
+                        session,
+                        key.clone(),
+                        body,
+                        retention_hours,
+                        content_type,
+                    )
+                    .await,
                 );
 
                 if attempts < self.max_attempts.get()
@@ -642,12 +758,16 @@ impl ObjectstoreServiceInner {
         key: Option<String>,
         body: BodyAttempt,
         retention_hours: Option<u16>,
+        content_type: Option<ContentType>,
     ) -> Result<ObjectstoreKey, objectstore_client::Error> {
         let mut request = match body {
             BodyAttempt::Bytes(bytes) => session.put(bytes),
             BodyAttempt::Stream(stream) => session.put_stream(stream.boxed()),
         };
 
+        if let Some(content_type) = content_type {
+            request = request.content_type(content_type.as_str());
+        }
         if let Some(retention_hours) = retention_hours {
             request = request.expiration_policy(ExpirationPolicy::TimeToLive(
                 Duration::from_hours(retention_hours.into()),