feat(store): Introduce a sample rate to rollout json spans (#5101)

Adds two options that allow to completely override what's set in
`produce_json`
and `produce_protobuf`. The rollout forces JSON and disables protobuf.

Closes STREAM-413

Author: jan-auer

relay-config/src/config.rs

@@ -1197,6 +1197,14 @@ impl Default for Processing {
 #[derive(Debug, Serialize, Deserialize)]
 #[serde(default)]
 pub struct SpanProducers {
+    /// Random sample of organizations to produce json and no protobuf.
+    ///
+    /// Overrides both `produce_json` and `produce_protobuf` for matching orgs.
+    pub produce_json_sample_rate: Option<f32>,
+    /// List of organization IDs to produce JSON spans for.
+    ///
+    /// Overrides both `produce_json` and `produce_protobuf` for matching orgs.
+    pub produce_json_orgs: Vec<u64>,
     /// Send JSON spans to `ingest-spans`.
     pub produce_json: bool,
     /// Send Protobuf (TraceItem) to `snuba-items`.
@@ -1206,6 +1214,8 @@ pub struct SpanProducers {
 impl Default for SpanProducers {
     fn default() -> Self {
         Self {
+            produce_json_sample_rate: None,
+            produce_json_orgs: vec![],
             produce_json: false,
             produce_protobuf: true,
         }
@@ -2614,14 +2624,9 @@ impl Config {
         forward.unwrap_or_else(|| !self.processing_enabled())
     }
 
-    /// Returns `true` if we should produce TraceItem spans on `snuba-items`.
-    pub fn produce_protobuf_spans(&self) -> bool {
-        self.values.processing.span_producers.produce_protobuf
-    }
-
-    /// Returns `true` if we should produce JSON spans on `ingest-spans`.
-    pub fn produce_json_spans(&self) -> bool {
-        self.values.processing.span_producers.produce_json
+    /// Returns the configuration for span producers.
+    pub fn span_producers(&self) -> &SpanProducers {
+        &self.values.processing.span_producers
     }
 }
 

relay-server/src/services/store.rs

@@ -45,7 +45,7 @@ use crate::services::global_config::GlobalConfigHandle;
 use crate::services::outcome::{DiscardItemType, DiscardReason, Outcome, TrackOutcome};
 use crate::services::processor::Processed;
 use crate::statsd::{RelayCounters, RelayGauges, RelayTimers};
-use crate::utils::FormDataIter;
+use crate::utils::{self, FormDataIter, PickResult};
 
 /// Fallback name used for attachment items without a `filename` header.
 const UNNAMED_ATTACHMENT: &str = "Unnamed Attachment";
@@ -1055,7 +1055,8 @@ impl StoreService {
         span.start_timestamp_ms = (span.start_timestamp_precise * 1e3) as u64;
         span.key_id = scoping.key_id;
 
-        if self.config.produce_protobuf_spans() {
+        let spans_target = self.spans_target(span.organization_id);
+        if spans_target.is_protobuf() {
             self.inner_produce_protobuf_span(
                 scoping,
                 received_at,
@@ -1075,13 +1076,32 @@ impl StoreService {
             });
         }
 
-        if self.config.produce_json_spans() {
+        if spans_target.is_json() {
             self.inner_produce_json_span(scoping, span)?;
         }
 
         Ok(())
     }
 
+    fn spans_target(&self, org_id: u64) -> SpansTarget {
+        let config = self.config.span_producers();
+        if config.produce_json_orgs.contains(&org_id) {
+            return SpansTarget::Json;
+        } else if let Some(rate) = config.produce_json_sample_rate {
+            return match utils::is_rolled_out(org_id, rate) {
+                PickResult::Keep => SpansTarget::Json,
+                PickResult::Discard => SpansTarget::Protobuf,
+            };
+        }
+
+        match (config.produce_json, config.produce_protobuf) {
+            (true, true) => SpansTarget::Both,
+            (true, false) => SpansTarget::Json,
+            (false, true) => SpansTarget::Protobuf,
+            (false, false) => SpansTarget::default(),
+        }
+    }
+
     fn inner_produce_json_span(
         &self,
         scoping: Scoping,
@@ -2005,6 +2025,24 @@ fn safe_timestamp(timestamp: DateTime<Utc>) -> u64 {
     Utc::now().timestamp() as u64
 }
 
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
+enum SpansTarget {
+    #[default]
+    Protobuf,
+    Json,
+    Both,
+}
+
+impl SpansTarget {
+    fn is_protobuf(&self) -> bool {
+        matches!(self, SpansTarget::Protobuf | SpansTarget::Both)
+    }
+
+    fn is_json(&self) -> bool {
+        matches!(self, SpansTarget::Json | SpansTarget::Both)
+    }
+}
+
 #[cfg(test)]
 mod tests {