From 28c0f19d5f7f5b0637646542fce1864f2038b67a Mon Sep 17 00:00:00 2001
From: Richard Ortenberg <richard.ortenberg@sentry.io>
Date: Fri, 23 Jun 2023 09:00:36 -0700
Subject: [PATCH] feat(monitors): Add trace context to CheckIns (#2241)

Adds `contexts` to the CheckIn payload in order to enable tracing.
---
 py/CHANGELOG.md                    |  1 +
 relay-monitors/src/lib.rs          | 30 ++++++++++++++++++++++++++++--
 relay-server/src/endpoints/cron.rs |  1 +
 3 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/py/CHANGELOG.md b/py/CHANGELOG.md
index 3ad16b5a6c..6f2858f789 100644
--- a/py/CHANGELOG.md
+++ b/py/CHANGELOG.md
@@ -4,6 +4,7 @@
 
 - Add filter based on transaction names. ([#2118](https://github.com/getsentry/relay/pull/2118))
 - Add `lock` attribute to the frame protocol. ([#2171](https://github.com/getsentry/relay/pull/2171))
+- Add trace context to CheckIns. ([#2241](https://github.com/getsentry/relay/pull/2241))
 
 ## 0.8.25
 
diff --git a/relay-monitors/src/lib.rs b/relay-monitors/src/lib.rs
index a5b556c4cc..824ffb9702 100644
--- a/relay-monitors/src/lib.rs
+++ b/relay-monitors/src/lib.rs
@@ -77,7 +77,7 @@ enum IntervalName {
     Minute,
 }
 
-/// The monitor configuration playload for upserting monitors during check-in
+/// The monitor configuration payload for upserting monitors during check-in
 #[derive(Debug, Deserialize, Serialize)]
 pub struct MonitorConfig {
     /// The monitor schedule configuration
@@ -98,6 +98,22 @@ pub struct MonitorConfig {
     timezone: Option<String>,
 }
 
+/// The trace context sent with a check-in.
+#[derive(Debug, Deserialize, Serialize)]
+pub struct CheckInTrace {
+    /// Trace-ID of the check-in.
+    #[serde(serialize_with = "uuid_simple")]
+    trace_id: Uuid,
+}
+
+/// Any contexts sent in the check-in payload.
+#[derive(Debug, Deserialize, Serialize)]
+pub struct CheckInContexts {
+    /// Trace context sent with a check-in.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    trace: Option<CheckInTrace>,
+}
+
 /// The monitor check-in payload.
 #[derive(Debug, Deserialize, Serialize)]
 pub struct CheckIn {
@@ -122,6 +138,11 @@ pub struct CheckIn {
     /// monitor configuration to support upserts.
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub monitor_config: Option<MonitorConfig>,
+
+    /// Contexts describing the associated environment of the job run.
+    /// Only supports trace for now.
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub contexts: Option<CheckInContexts>,
 }
 
 /// Normalizes a monitor check-in payload.
@@ -172,7 +193,12 @@ mod tests {
   "monitor_slug": "my-monitor",
   "status": "in_progress",
   "environment": "production",
-  "duration": 21.0
+  "duration": 21.0,
+  "contexts": {
+    "trace": {
+      "trace_id": "8f431b7aa08441bbbd5a0100fd91f9fe"
+    }
+  }
 }"#;
 
         let check_in = serde_json::from_str::<CheckIn>(json).unwrap();
diff --git a/relay-server/src/endpoints/cron.rs b/relay-server/src/endpoints/cron.rs
index 21fc002933..456a6e3af8 100644
--- a/relay-server/src/endpoints/cron.rs
+++ b/relay-server/src/endpoints/cron.rs
@@ -51,6 +51,7 @@ impl CronParams {
                 environment: query.environment,
                 duration: query.duration,
                 monitor_config: None,
+                contexts: None,
             })
             .map_err(BadStoreRequest::InvalidJson)?,
         );