feat(bottlecap): add SQS inferred span

bottlecap/src/lifecycle/invocation/processor.rs

@@ -205,14 +205,15 @@ impl Processor {
         self.span.trace_id = 0;
         self.span.parent_id = 0;
         self.span.span_id = 0;
+        self.extracted_span_context = None;
 
         let payload_value = match serde_json::from_slice::<Value>(&payload) {
             Ok(value) => value,
             Err(_) => json!({}),
         };
 
-        self.extracted_span_context = self.extract_span_context(&headers, &payload_value);
         self.inferrer.infer_span(&payload_value, &self.aws_config);
+        self.extracted_span_context = self.extract_span_context(&headers, &payload_value);
 
         if let Some(sc) = &self.extracted_span_context {
             self.span.trace_id = sc.trace_id;
@@ -264,19 +265,46 @@ impl Processor {
     ///
     pub fn on_invocation_end(
         &mut self,
-        trace_id: u64,
-        span_id: u64,
-        parent_id: u64,
+        headers: HashMap<String, String>,
         status_code: Option<String>,
     ) {
-        self.span.trace_id = trace_id;
-        self.span.span_id = span_id;
-
+        self.update_span_context(headers);
         if self.inferrer.inferred_span.is_some() {
             if let Some(status_code) = status_code {
                 self.inferrer.set_status_code(status_code);
             }
-        } else {
+        }
+    }
+
+    fn update_span_context(&mut self, headers: HashMap<String, String>) {
+        // todo: fix this, code is a copy of the existing logic in Go, not accounting
+        // when a 128 bit trace id exist
+        let mut trace_id = 0;
+        let mut span_id = 0;
+        let mut parent_id = 0;
+
+        // If we have a trace context, update the span context
+        if let Some(sc) = &mut self.extracted_span_context {
+            trace_id = sc.trace_id;
+            span_id = sc.span_id;
+        }
+
+        if let Some(header) = headers.get("x-datadog-trace-id") {
+            trace_id = header.parse::<u64>().unwrap_or(0);
+        }
+
+        if let Some(header) = headers.get("x-datadog-span-id") {
+            span_id = header.parse::<u64>().unwrap_or(0);
+        }
+
+        if let Some(header) = headers.get("x-datadog-parent-id") {
+            parent_id = header.parse::<u64>().unwrap_or(0);
+        }
+
+        self.span.trace_id = trace_id;
+        self.span.span_id = span_id;
+
+        if self.inferrer.inferred_span.is_none() {
             self.span.parent_id = parent_id;
         }
     }

bottlecap/src/lifecycle/invocation/span_inferrer.rs

@@ -9,13 +9,10 @@ use crate::config::AwsConfig;
 
 use crate::lifecycle::invocation::triggers::{
     api_gateway_http_event::APIGatewayHttpEvent, api_gateway_rest_event::APIGatewayRestEvent,
-    Trigger,
+    sqs_event::SqsRecord, Trigger, FUNCTION_TRIGGER_EVENT_SOURCE_ARN_TAG,
 };
 use crate::tags::lambda::tags::{INIT_TYPE, SNAP_START_VALUE};
 
-const FUNCTION_TRIGGER_EVENT_SOURCE_TAG: &str = "function_trigger.event_source";
-const FUNCTION_TRIGGER_EVENT_SOURCE_ARN_TAG: &str = "function_trigger.event_source_arn";
-
 pub struct SpanInferrer {
     pub inferred_span: Option<Span>,
     is_async_span: bool,
@@ -46,6 +43,10 @@ impl SpanInferrer {
     ///
     pub fn infer_span(&mut self, payload_value: &Value, aws_config: &AwsConfig) {
         self.inferred_span = None;
+        self.is_async_span = false;
+        self.carrier = None;
+        self.trigger_tags = None;
+
         if APIGatewayHttpEvent::is_match(payload_value) {
             if let Some(t) = APIGatewayHttpEvent::new(payload_value.clone()) {
                 let mut span = Span {
@@ -54,19 +55,14 @@ impl SpanInferrer {
                 };
 
                 t.enrich_span(&mut span);
-                span.meta.extend([
-                    (
-                        FUNCTION_TRIGGER_EVENT_SOURCE_TAG.to_string(),
-                        "api_gateway".to_string(),
-                    ),
-                    (
-                        FUNCTION_TRIGGER_EVENT_SOURCE_ARN_TAG.to_string(),
-                        t.get_arn(&aws_config.region),
-                    ),
-                ]);
+                let mut tt = t.get_tags();
+                tt.extend([(
+                    FUNCTION_TRIGGER_EVENT_SOURCE_ARN_TAG.to_string(),
+                    t.get_arn(&aws_config.region),
+                )]);
 
                 self.carrier = Some(t.get_carrier());
-                self.trigger_tags = Some(t.get_tags());
+                self.trigger_tags = Some(tt);
                 self.is_async_span = t.is_async();
                 self.inferred_span = Some(span);
             }
@@ -78,24 +74,38 @@ impl SpanInferrer {
                 };
 
                 t.enrich_span(&mut span);
-                span.meta.extend([
-                    (
-                        FUNCTION_TRIGGER_EVENT_SOURCE_TAG.to_string(),
-                        "api_gateway".to_string(),
-                    ),
-                    (
-                        FUNCTION_TRIGGER_EVENT_SOURCE_ARN_TAG.to_string(),
-                        t.get_arn(&aws_config.region),
-                    ),
-                ]);
+                let mut tt = t.get_tags();
+                tt.extend([(
+                    FUNCTION_TRIGGER_EVENT_SOURCE_ARN_TAG.to_string(),
+                    t.get_arn(&aws_config.region),
+                )]);
+
+                self.carrier = Some(t.get_carrier());
+                self.trigger_tags = Some(tt);
+                self.is_async_span = t.is_async();
+                self.inferred_span = Some(span);
+            }
+        } else if SqsRecord::is_match(payload_value) {
+            if let Some(t) = SqsRecord::new(payload_value.clone()) {
+                let mut span = Span {
+                    span_id: Self::generate_span_id(),
+                    ..Default::default()
+                };
+
+                t.enrich_span(&mut span);
+                let mut tt = t.get_tags();
+                tt.extend([(
+                    FUNCTION_TRIGGER_EVENT_SOURCE_ARN_TAG.to_string(),
+                    t.get_arn(&aws_config.region),
+                )]);
 
                 self.carrier = Some(t.get_carrier());
-                self.trigger_tags = Some(t.get_tags());
+                self.trigger_tags = Some(tt);
                 self.is_async_span = t.is_async();
                 self.inferred_span = Some(span);
             }
         } else {
-            debug!("Unable to infer span from payload");
+            debug!("Unable to infer span from payload: no matching trigger found");
         }
     }
 
@@ -124,7 +134,8 @@ impl SpanInferrer {
     pub fn complete_inferred_span(&mut self, invocation_span: &Span) {
         if let Some(s) = &mut self.inferred_span {
             if self.is_async_span {
-                if s.duration != 0 {
+                // SNS to SQS span duration will be set
+                if s.duration == 0 {
                     let duration = invocation_span.start - s.start;
                     s.duration = duration;
                 }

bottlecap/src/lifecycle/invocation/triggers/api_gateway_http_event.rs

@@ -6,7 +6,9 @@ use tracing::debug;
 
 use crate::lifecycle::invocation::{
     processor::MS_TO_NS,
-    triggers::{get_aws_partition_by_region, lowercase_key, Trigger},
+    triggers::{
+        get_aws_partition_by_region, lowercase_key, Trigger, FUNCTION_TRIGGER_EVENT_SOURCE_TAG,
+    },
 };
 
 #[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
@@ -114,10 +116,7 @@ impl Trigger for APIGatewayHttpEvent {
                 "request_id".to_string(),
                 self.request_context.request_id.clone(),
             ),
-            ("resource_names".to_string(), resource),
         ]));
-
-        // todo: update global(? IsAsync if event payload is `Event`
     }
 
     fn get_tags(&self) -> HashMap<String, String> {
@@ -140,6 +139,10 @@ impl Trigger for APIGatewayHttpEvent {
                 "http.method".to_string(),
                 self.request_context.http.method.clone(),
             ),
+            (
+                FUNCTION_TRIGGER_EVENT_SOURCE_TAG.to_string(),
+                "api-gateway".to_string(),
+            ),
         ]);
         // route is parameterized
         // /users/{id}/profile
@@ -287,7 +290,6 @@ mod tests {
                 ("http.user_agent".to_string(), "curl/7.64.1".to_string()),
                 ("operation_name".to_string(), "aws.httpapi".to_string()),
                 ("request_id".to_string(), "FaHnXjKCGjQEJ7A=".to_string()),
-                ("resource_names".to_string(), "GET /httpapi/get".to_string()),
             ])
         );
     }
@@ -311,6 +313,10 @@ mod tests {
             ("http.method".to_string(), "GET".to_string()),
             ("http.route".to_string(), "/httpapi/get".to_string()),
             ("http.user_agent".to_string(), "curl/7.64.1".to_string()),
+            (
+                "function_trigger.event_source".to_string(),
+                "api-gateway".to_string(),
+            ),
         ]);
 
         assert_eq!(tags, expected);
@@ -345,7 +351,6 @@ mod tests {
                 ("http.user_agent".to_string(), "curl/8.1.2".to_string()),
                 ("operation_name".to_string(), "aws.httpapi".to_string()),
                 ("request_id".to_string(), "Ur2JtjEfGjQEPOg=".to_string()),
-                ("resource_names".to_string(), "GET /user/{id}".to_string()),
             ])
         );
     }
@@ -367,6 +372,10 @@ mod tests {
             ("http.method".to_string(), "GET".to_string()),
             ("http.route".to_string(), "/user/{id}".to_string()),
             ("http.user_agent".to_string(), "curl/8.1.2".to_string()),
+            (
+                "function_trigger.event_source".to_string(),
+                "api-gateway".to_string(),
+            ),
         ]);
         assert_eq!(tags, expected);
     }

bottlecap/src/lifecycle/invocation/triggers/api_gateway_rest_event.rs

@@ -6,7 +6,9 @@ use tracing::debug;
 
 use crate::lifecycle::invocation::{
     processor::MS_TO_NS,
-    triggers::{get_aws_partition_by_region, lowercase_key, Trigger},
+    triggers::{
+        get_aws_partition_by_region, lowercase_key, Trigger, FUNCTION_TRIGGER_EVENT_SOURCE_TAG,
+    },
 };
 
 #[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
@@ -109,7 +111,6 @@ impl Trigger for APIGatewayRestEvent {
                 "request_id".to_string(),
                 self.request_context.request_id.clone(),
             ),
-            ("resource_names".to_string(), resource.clone()),
             (
                 "http.route".to_string(),
                 self.request_context.resource_path.clone(),
@@ -143,6 +144,10 @@ impl Trigger for APIGatewayRestEvent {
                 "http.user_agent".to_string(),
                 self.request_context.identity.user_agent.to_string(),
             ),
+            (
+                FUNCTION_TRIGGER_EVENT_SOURCE_TAG.to_string(),
+                "api-gateway".to_string(),
+            ),
         ]);
 
         if let Some(referer) = self.headers.get("referer") {
@@ -256,7 +261,6 @@ mod tests {
                 ("http.route".to_string(), "/path".to_string()),
                 ("operation_name".to_string(), "aws.apigateway".to_string()),
                 ("request_id".to_string(), "id=".to_string()),
-                ("resource_names".to_string(), "GET /path".to_string()),
             ])
         );
     }
@@ -278,6 +282,10 @@ mod tests {
             ("http.method".to_string(), "GET".to_string()),
             ("http.route".to_string(), "/path".to_string()),
             ("http.user_agent".to_string(), "user-agent".to_string()),
+            (
+                "function_trigger.event_source".to_string(),
+                "api-gateway".to_string(),
+            ),
         ]);
 
         assert_eq!(tags, expected);
@@ -314,7 +322,6 @@ mod tests {
                 "request_id".to_string(),
                 "e16399f7-e984-463a-9931-745ba021a27f".to_string(),
             ),
-            ("resource_names".to_string(), "GET /user/{id}".to_string()),
         ]);
         assert_eq!(span.meta, expected);
     }
@@ -342,6 +349,10 @@ mod tests {
                 ("http.method".to_string(), "GET".to_string()),
                 ("http.route".to_string(), "/user/{id}".to_string()),
                 ("http.user_agent".to_string(), "curl/8.1.2".to_string()),
+                (
+                    "function_trigger.event_source".to_string(),
+                    "api-gateway".to_string()
+                ),
             ])
         );
     }

bottlecap/src/lifecycle/invocation/triggers/mod.rs

@@ -6,6 +6,11 @@ use serde_json::Value;
 
 pub mod api_gateway_http_event;
 pub mod api_gateway_rest_event;
+pub mod sqs_event;
+
+pub const DATADOG_CARRIER_KEY: &str = "_datadog";
+pub const FUNCTION_TRIGGER_EVENT_SOURCE_TAG: &str = "function_trigger.event_source";
+pub const FUNCTION_TRIGGER_EVENT_SOURCE_ARN_TAG: &str = "function_trigger.event_source_arn";
 
 pub trait Trigger: Sized {
     fn new(payload: Value) -> Option<Self>;