feat(profiles): Support profiler_id in context (#3714)

getsentry · Jun 26, 2024 · a6bc0a2 · a6bc0a2
1 parent 9ec42ed
commit a6bc0a2
Show file tree

Hide file tree

Showing 5 changed files with 122 additions and 5 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -11,6 +11,7 @@
 **Internal**:
 
 - Aggregate metrics before rate limiting. ([#3746](https://github.com/getsentry/relay/pull/3746))
+- Accept profiler_id in the profile context. ([#3714](https://github.com/getsentry/relay/pull/3714))
 
 ## 24.6.0
 

diff --git a/relay-event-schema/src/protocol/contexts/profile.rs b/relay-event-schema/src/protocol/contexts/profile.rs
@@ -10,8 +10,10 @@ use crate::protocol::EventId;
 #[cfg_attr(feature = "jsonschema", derive(JsonSchema))]
 pub struct ProfileContext {
     /// The profile ID.
-    #[metastructure(required = "true")]
     pub profile_id: Annotated<EventId>,
+
+    /// The profiler ID.
+    pub profiler_id: Annotated<EventId>,
 }
 
 impl super::DefaultContext for ProfileContext {
@@ -60,6 +62,7 @@ mod tests {
             profile_id: Annotated::new(EventId(
                 "4c79f60c11214eb38604f4ae0781bfb2".parse().unwrap(),
             )),
+            ..ProfileContext::default()
         })));
 
         assert_eq!(context, Annotated::from_json(json).unwrap());
@@ -76,6 +79,20 @@ mod tests {
             profile_id: Annotated::new(EventId(
                 "4c79f60c11214eb38604f4ae0781bfb2".parse().unwrap(),
             )),
+            ..ProfileContext::default()
+        })));
+
+        assert_eq!(context, Annotated::from_json(json).unwrap());
+    }
+
+    #[test]
+    fn context_with_profiler_id() {
+        let json = r#"{"profiler_id": "4C79F60C11214EB38604F4AE0781BFB2", "type": "profile"}"#;
+        let context = Annotated::new(Context::Profile(Box::new(ProfileContext {
+            profiler_id: Annotated::new(EventId(
+                "4c79f60c11214eb38604f4ae0781bfb2".parse().unwrap(),
+            )),
+            ..ProfileContext::default()
         })));
 
         assert_eq!(context, Annotated::from_json(json).unwrap());

diff --git a/relay-event-schema/src/protocol/event.rs b/relay-event-schema/src/protocol/event.rs
@@ -1156,6 +1156,7 @@ mod tests {
                     profile_id: Annotated::new(EventId(uuid!(
                         "abadcade-feed-dead-beef-8addadfeedaa"
                     ))),
+                    ..ProfileContext::default()
                 });
                 contexts
             }),

diff --git a/relay-server/src/services/processor/profile.rs b/relay-server/src/services/processor/profile.rs
@@ -76,11 +76,14 @@ pub fn transfer_id(
             let contexts = event.contexts.get_or_insert_with(Contexts::new);
             contexts.add(ProfileContext {
                 profile_id: Annotated::new(profile_id),
+                ..ProfileContext::default()
             });
         }
         None => {
             if let Some(contexts) = event.contexts.value_mut() {
-                contexts.remove::<ProfileContext>();
+                if let Some(profile_context) = contexts.get_mut::<ProfileContext>() {
+                    profile_context.profile_id = Annotated::empty();
+                }
             }
         }
     }
@@ -239,6 +242,7 @@ mod tests {
             profile_id: EventId(
                 012d836b-15bb-49d7-bbf9-9e64295d995b,
             ),
+            profiler_id: ~,
         }
         "###);
     }
@@ -293,4 +297,88 @@ mod tests {
         let envelope_response = processor.process(message).unwrap();
         assert!(envelope_response.envelope.is_none());
     }
+
+    #[tokio::test]
+    async fn test_profile_id_removed_profiler_id_kept() {
+        // Setup
+        let processor = create_test_processor(Default::default());
+        let event_id = EventId::new();
+        let dsn = "https://e12d836b15bb49d7bbf99e64295d995b:@sentry.io/42"
+            .parse()
+            .unwrap();
+        let request_meta = RequestMeta::new(dsn);
+        let mut envelope = Envelope::from_request(Some(event_id), request_meta);
+
+        // Add a valid transaction item.
+        envelope.add_item({
+            let mut item = Item::new(ItemType::Transaction);
+
+            item.set_payload(
+                ContentType::Json,
+                r#"{
+                "type": "transaction",
+                "transaction": "/foo/",
+                "timestamp": 946684810.0,
+                "start_timestamp": 946684800.0,
+                "contexts": {
+                    "trace": {
+                        "trace_id": "4c79f60c11214eb38604f4ae0781bfb2",
+                        "span_id": "fa90fdead5f74053",
+                        "op": "http.server",
+                        "type": "trace"
+                    },
+                    "profile": {
+                        "profile_id": "4c79f60c11214eb38604f4ae0781bfb2",
+                        "profiler_id": "4c79f60c11214eb38604f4ae0781bfb2",
+                        "type": "profile"
+                    }
+                },
+                "transaction_info": {
+                    "source": "url"
+                }
+            }"#,
+            );
+            item
+        });
+
+        let mut project_state = ProjectState::allowed();
+        project_state.config.features.0.insert(Feature::Profiling);
+
+        let mut envelopes = ProcessingGroup::split_envelope(*envelope);
+        assert_eq!(envelopes.len(), 1);
+
+        let (group, envelope) = envelopes.pop().unwrap();
+        let envelope = ManagedEnvelope::standalone(envelope, Addr::dummy(), Addr::dummy(), group);
+
+        let message = ProcessEnvelope {
+            envelope,
+            project_state: Arc::new(project_state),
+            sampling_project_state: None,
+            reservoir_counters: ReservoirCounters::default(),
+        };
+
+        let envelope_response = processor.process(message).unwrap();
+        let ctx = envelope_response.envelope.unwrap();
+        let new_envelope = ctx.envelope();
+
+        // Get the re-serialized context.
+        let item = new_envelope
+            .get_item_by(|item| item.ty() == &ItemType::Transaction)
+            .unwrap();
+        let transaction = Annotated::<Event>::from_json_bytes(&item.payload()).unwrap();
+        let context = transaction
+            .value()
+            .unwrap()
+            .context::<ProfileContext>()
+            .unwrap();
+
+        assert_debug_snapshot!(context, @r###"
+        ProfileContext {
+            profile_id: ~,
+            profiler_id: EventId(
+                4c79f60c-1121-4eb3-8604-f4ae0781bfb2,
+            ),
+        }
+        "###);
+    }
 }
diff --git a/relay-server/tests/snapshots/test_fixtures__event_schema.snap b/relay-server/tests/snapshots/test_fixtures__event_schema.snap
@@ -2988,12 +2988,22 @@ expression: "relay_event_schema::protocol::event_json_schema()"
       "anyOf": [
         {
           "type": "object",
-          "required": [
-            "profile_id"
-          ],
           "properties": {
             "profile_id": {
               "description": " The profile ID.",
+              "default": null,
+              "anyOf": [
+                {
+                  "$ref": "#/definitions/EventId"
+                },
+                {
+                  "type": "null"
+                }
+              ]
+            },
+            "profiler_id": {
+              "description": " The profiler ID.",
+              "default": null,
               "anyOf": [
                 {
                   "$ref": "#/definitions/EventId"