ref(processing): Pass store handle to forward_store (#5425)

Author: jjbayer

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

@@ -112,13 +112,13 @@ impl Forward for CheckInsOutput {
     #[cfg(feature = "processing")]
     fn forward_store(
         self,
-        s: &relay_system::Addr<crate::services::store::Store>,
+        s: processing::StoreHandle<'_>,
         ctx: processing::ForwardContext<'_>,
     ) -> Result<(), Rejected<()>> {
         let envelope = self.serialize_envelope(ctx)?;
         let envelope = ManagedEnvelope::from(envelope).into_processed();
 
-        s.send(crate::services::store::StoreEnvelope { envelope });
+        s.store(crate::services::store::StoreEnvelope { envelope });
 
         Ok(())
     }

relay-server/src/processing/common.rs

@@ -1,6 +1,8 @@
 use crate::Envelope;
 use crate::managed::{Managed, Rejected};
 use crate::processing::ForwardContext;
+#[cfg(feature = "processing")]
+use crate::processing::StoreHandle;
 use crate::processing::check_ins::CheckInsProcessor;
 use crate::processing::logs::LogsProcessor;
 use crate::processing::sessions::SessionsProcessor;
@@ -30,7 +32,7 @@ macro_rules! outputs {
             #[cfg(feature = "processing")]
             fn forward_store(
                 self,
-                s: &relay_system::Addr<crate::services::store::Store>,
+                s: StoreHandle<'_>,
                 ctx: ForwardContext<'_>,
             ) -> Result<(), Rejected<()>> {
                 match self {

relay-server/src/processing/forward.rs

@@ -2,10 +2,52 @@ use relay_config::Config;
 use relay_dynamic_config::GlobalConfig;
 #[cfg(feature = "processing")]
 use relay_dynamic_config::{RetentionConfig, RetentionsConfig};
+#[cfg(feature = "processing")]
+use relay_system::{Addr, FromMessage};
 
 use crate::Envelope;
 use crate::managed::{Managed, Rejected};
 use crate::services::projects::project::ProjectInfo;
+#[cfg(feature = "processing")]
+use crate::services::store::Store;
+#[cfg(feature = "processing")]
+use crate::services::upload::Upload;
+
+/// A transparent handle that dispatches between store-like services.
+#[cfg(feature = "processing")]
+#[derive(Debug, Clone, Copy)]
+pub struct StoreHandle<'a> {
+    store: &'a Addr<Store>,
+    upload: Option<&'a Addr<Upload>>,
+}
+
+#[cfg(feature = "processing")]
+impl<'a> StoreHandle<'a> {
+    pub fn new(store: &'a Addr<Store>, upload: Option<&'a Addr<Upload>>) -> Self {
+        Self { store, upload }
+    }
+
+    /// Sends a message to the [`Store`] service.
+    pub fn store<M>(&self, message: M)
+    where
+        Store: FromMessage<M>,
+    {
+        self.store.send(message);
+    }
+
+    /// Sends a message to the [`Upload`] service.
+    #[expect(unused)]
+    pub fn upload<M>(&self, message: M)
+    where
+        Upload: FromMessage<M>,
+    {
+        if let Some(upload) = self.upload {
+            upload.send(message);
+        } else {
+            relay_log::error!("Upload service not configured. Dropping message.");
+        }
+    }
+}
 
 /// A processor output which can be forwarded to a different destination.
 pub trait Forward {
@@ -21,11 +63,8 @@ pub trait Forward {
     ///
     /// This function must only be called when Relay is configured to be in processing mode.
     #[cfg(feature = "processing")]
-    fn forward_store(
-        self,
-        s: &relay_system::Addr<crate::services::store::Store>,
-        ctx: ForwardContext<'_>,
-    ) -> Result<(), Rejected<()>>;
+    fn forward_store(self, s: StoreHandle<'_>, ctx: ForwardContext<'_>)
+    -> Result<(), Rejected<()>>;
 }
 
 /// Context passed to [`Forward`].
@@ -80,11 +119,7 @@ impl Forward for Nothing {
     }
 
     #[cfg(feature = "processing")]
-    fn forward_store(
-        self,
-        _: &relay_system::Addr<crate::services::store::Store>,
-        _: ForwardContext<'_>,
-    ) -> Result<(), Rejected<()>> {
+    fn forward_store(self, _: StoreHandle<'_>, _: ForwardContext<'_>) -> Result<(), Rejected<()>> {
         match self {}
     }
 }

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

@@ -175,7 +175,7 @@ impl Forward for LogOutput {
     #[cfg(feature = "processing")]
     fn forward_store(
         self,
-        s: &relay_system::Addr<crate::services::store::Store>,
+        s: processing::StoreHandle<'_>,
         ctx: processing::ForwardContext<'_>,
     ) -> Result<(), Rejected<()>> {
         let Self(logs) = self;
@@ -188,7 +188,7 @@ impl Forward for LogOutput {
 
         for log in logs.split(|logs| logs.logs) {
             if let Ok(log) = log.try_map(|log, _| store::convert(log, &ctx)) {
-                s.send(log)
+                s.store(log)
             };
         }
 

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

@@ -123,7 +123,7 @@ impl Forward for SessionsOutput {
     #[cfg(feature = "processing")]
     fn forward_store(
         self,
-        _: &relay_system::Addr<crate::services::store::Store>,
+        _: processing::forward::StoreHandle<'_>,
         _: processing::ForwardContext<'_>,
     ) -> Result<(), Rejected<()>> {
         let SessionsOutput(sessions) = self;

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

@@ -226,7 +226,7 @@ impl Forward for SpanOutput {
     #[cfg(feature = "processing")]
     fn forward_store(
         self,
-        s: &relay_system::Addr<crate::services::store::Store>,
+        s: processing::forward::StoreHandle<'_>,
         ctx: processing::ForwardContext<'_>,
     ) -> Result<(), Rejected<()>> {
         let spans = match self {
@@ -255,7 +255,7 @@ impl Forward for SpanOutput {
 
         for span in spans.split(|spans| spans.into_indexed_spans()) {
             if let Ok(span) = span.try_map(|span, _| store::convert(span, &ctx)) {
-                s.send(span)
+                s.store(span)
             };
         }
 

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

@@ -161,7 +161,7 @@ impl Forward for TraceMetricOutput {
     #[cfg(feature = "processing")]
     fn forward_store(
         self,
-        s: &relay_system::Addr<crate::services::store::Store>,
+        s: processing::forward::StoreHandle<'_>,
         ctx: processing::ForwardContext<'_>,
     ) -> Result<(), Rejected<()>> {
         let Self(metrics) = self;
@@ -174,7 +174,7 @@ impl Forward for TraceMetricOutput {
 
         for metric in metrics.split(|metrics| metrics.metrics) {
             if let Ok(metric) = metric.try_map(|metric, _| store::convert(metric, &ctx)) {
-                s.send(metric);
+                s.store(metric);
             }
         }
 

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

@@ -19,6 +19,10 @@ use crate::envelope::{ContentType, EnvelopeHeaders, Item, ItemType, Items};
 use crate::managed::{
     Counted, Managed, ManagedEnvelope, ManagedResult, OutcomeError, Quantities, Rejected,
 };
+#[cfg(feature = "processing")]
+use crate::processing;
+#[cfg(feature = "processing")]
+use crate::processing::forward::StoreHandle;
 use crate::processing::transactions::profile::{Profile, ProfileWithHeaders};
 use crate::processing::utils::event::{
     EventFullyNormalized, EventMetricsExtracted, FiltersStatus, SpansExtracted, event_type,
@@ -592,13 +596,13 @@ impl Forward for TransactionOutput {
     #[cfg(feature = "processing")]
     fn forward_store(
         self,
-        s: &relay_system::Addr<crate::services::store::Store>,
+        s: StoreHandle<'_>,
         ctx: ForwardContext<'_>,
     ) -> Result<(), Rejected<()>> {
         // TODO: split out spans into a separate message.
         let envelope: ManagedEnvelope = self.serialize_envelope(ctx)?.into();
 
-        s.send(StoreEnvelope {
+        s.store(StoreEnvelope {
             envelope: envelope.into_processed(),
         });
 

relay-server/src/services/processor.rs

@@ -2305,6 +2305,9 @@ impl EnvelopeProcessorService {
         if self.inner.config.processing_enabled()
             && let Some(store_forwarder) = &self.inner.addrs.store_forwarder
         {
+            use crate::processing::StoreHandle;
+
+            let upload = self.inner.addrs.upload.as_ref();
             match submit {
                 Submit::Envelope(envelope) => {
                     let envelope_has_attachments = envelope
@@ -2328,7 +2331,7 @@ impl EnvelopeProcessorService {
                     }
                 }
                 Submit::Output { output, ctx } => output
-                    .forward_store(store_forwarder, ctx)
+                    .forward_store(StoreHandle::new(store_forwarder, upload), ctx)
                     .unwrap_or_else(|err| err.into_inner()),
             }
             return;