S3 Downstream Span Pointers (#503)

Author: nhulston

bottlecap/src/lifecycle/invocation/processor.rs

@@ -361,6 +361,7 @@ impl Processor {
                 header_tags,
                 vec![traces],
                 body_size,
+                self.inferrer.span_pointers.clone(),
             );
 
             if let Err(e) = trace_agent_tx.send(send_data).await {

bottlecap/src/lifecycle/invocation/span_inferrer.rs

@@ -22,6 +22,7 @@ use crate::lifecycle::invocation::{
         Trigger, FUNCTION_TRIGGER_EVENT_SOURCE_ARN_TAG,
     },
 };
+use crate::traces::span_pointers::SpanPointer;
 use crate::traces::{context::SpanContext, propagation::Propagator};
 
 #[derive(Default)]
@@ -39,6 +40,8 @@ pub struct SpanInferrer {
     generated_span_context: Option<SpanContext>,
     // Tags generated from the trigger
     trigger_tags: Option<HashMap<String, String>>,
+    // Span pointers from S3 or DynamoDB streams
+    pub span_pointers: Option<Vec<SpanPointer>>,
 }
 
 impl SpanInferrer {
@@ -52,6 +55,7 @@ impl SpanInferrer {
             carrier: None,
             generated_span_context: None,
             trigger_tags: None,
+            span_pointers: None,
         }
     }
 
@@ -178,6 +182,7 @@ impl SpanInferrer {
         } else if S3Record::is_match(payload_value) {
             if let Some(t) = S3Record::new(payload_value.clone()) {
                 t.enrich_span(&mut inferred_span, &self.service_mapping);
+                self.span_pointers = t.get_span_pointers();
 
                 trigger = Some(Box::new(t));
             }

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

@@ -10,6 +10,7 @@ use crate::lifecycle::invocation::{
     processor::MS_TO_NS,
     triggers::{ServiceNameResolver, Trigger, FUNCTION_TRIGGER_EVENT_SOURCE_TAG},
 };
+use crate::traces::span_pointers::{generate_span_pointer_hash, SpanPointer};
 
 #[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
 pub struct S3Event {
@@ -133,6 +134,28 @@ impl ServiceNameResolver for S3Record {
     }
 }
 
+impl S3Record {
+    pub fn get_span_pointers(&self) -> Option<Vec<SpanPointer>> {
+        let bucket_name = &self.s3.bucket.name;
+        let key = &self.s3.object.key;
+        // The AWS SDK sometimes wraps the S3 eTag in quotes, but sometimes doesn't.
+        let e_tag = self.s3.object.e_tag.trim_matches('"');
+
+        if bucket_name.is_empty() || key.is_empty() || e_tag.is_empty() {
+            debug!("Unable to create span pointer because bucket name, key, or etag is missing.");
+            return None;
+        }
+
+        // https://github.com/DataDog/dd-span-pointer-rules/blob/main/AWS/S3/Object/README.md
+        let hash = generate_span_pointer_hash(&[bucket_name, key, e_tag]);
+
+        Some(vec![SpanPointer {
+            hash,
+            kind: String::from("aws.s3.object"),
+        }])
+    }
+}
+
 #[cfg(test)]
 #[allow(clippy::unwrap_used)]
 mod tests {
@@ -274,4 +297,51 @@ mod tests {
             "generic-service"
         );
     }
+
+    #[test]
+    fn test_get_span_pointers() {
+        let event = S3Record {
+            event_source: String::from("aws:s3"),
+            event_time: Utc::now(),
+            event_name: String::from("ObjectCreated:Put"),
+            s3: S3Entity {
+                bucket: S3Bucket {
+                    name: String::from("test-bucket"),
+                    arn: String::from("arn:aws:s3:::test-bucket"),
+                },
+                object: S3Object {
+                    key: String::from("test/key"),
+                    size: 1024,
+                    e_tag: String::from("0123456789abcdef0123456789abcdef"),
+                },
+            },
+        }; //
+
+        let span_pointers = event.get_span_pointers().expect("Should return Some(vec)");
+        assert_eq!(span_pointers.len(), 1);
+        assert_eq!(span_pointers[0].kind, "aws.s3.object");
+        assert_eq!(span_pointers[0].hash, "40df87dbfdf59f32253a2668c23e51b4");
+    }
+
+    #[test]
+    fn test_get_span_pointers_missing_fields() {
+        let event = S3Record {
+            event_source: String::from("aws:s3"),
+            event_time: Utc::now(),
+            event_name: String::from("ObjectCreated:Put"),
+            s3: S3Entity {
+                bucket: S3Bucket {
+                    name: String::new(), // Empty bucket name
+                    arn: String::from("arn"),
+                },
+                object: S3Object {
+                    key: String::from("key"),
+                    size: 0,
+                    e_tag: String::from("etag"),
+                },
+            },
+        };
+
+        assert!(event.get_span_pointers().is_none());
+    }
 }

bottlecap/src/traces/mod.rs

@@ -3,6 +3,7 @@
 
 pub mod context;
 pub mod propagation;
+pub mod span_pointers;
 pub mod stats_flusher;
 pub mod stats_processor;
 pub mod trace_agent;

bottlecap/src/traces/span_pointers.rs

@@ -0,0 +1,257 @@
+// Copyright 2023-Present Datadog, Inc. https://www.datadoghq.com/
+// SPDX-License-Identifier: Apache-2.0
+
+use datadog_trace_protobuf::pb::SpanLink;
+use sha2::{Digest, Sha256};
+use std::collections::HashMap;
+
+const SPAN_POINTER_HASH_LENGTH: usize = 32;
+
+#[derive(Clone)]
+pub struct SpanPointer {
+    pub hash: String,
+    pub kind: String,
+}
+
+/// Returns the first 32 characters of the SHA-256 hash of the components joined by a '|'.
+/// Used by span pointers to uniquely & deterministically identify an `S3` or `DynamoDB` stream.
+/// <https://github.com/DataDog/dd-span-pointer-rules/blob/main/README.md#General%20Hashing%20Rules>
+#[must_use]
+pub fn generate_span_pointer_hash(components: &[&str]) -> String {
+    let mut hasher = Sha256::new();
+    hasher.update(components.join("|").as_bytes());
+    let result = hasher.finalize();
+    hex::encode(result)[..SPAN_POINTER_HASH_LENGTH].to_string()
+}
+
+pub fn attach_span_pointers_to_meta<S: ::std::hash::BuildHasher>(
+    meta: &mut HashMap<String, String, S>,
+    span_pointers: &Option<Vec<SpanPointer>>,
+) {
+    let Some(span_pointers) = span_pointers.as_ref().filter(|sp| !sp.is_empty()) else {
+        return;
+    };
+
+    let new_span_links: Vec<SpanLink> = span_pointers
+        .iter()
+        .map(|sp| {
+            SpanLink {
+                // We set all these fields as 0 or empty since they're unknown; the frontend
+                // uses `ptr.hash` instead to find the opposite link if it exists.
+                trace_id: 0,
+                span_id: 0,
+                trace_id_high: 0,
+                tracestate: String::new(),
+                flags: 0,
+                attributes: HashMap::from([
+                    ("link.kind".to_string(), "span-pointer".to_string()),
+                    ("ptr.dir".to_string(), "u".to_string()),
+                    ("ptr.hash".to_string(), sp.hash.clone()),
+                    ("ptr.kind".to_string(), sp.kind.clone()),
+                ]),
+            }
+        })
+        .collect();
+
+    let mut all_span_links = meta
+        .get("_dd.span_links")
+        .and_then(|existing| serde_json::from_str::<Vec<SpanLink>>(existing).ok())
+        .unwrap_or_default();
+
+    all_span_links.extend(new_span_links);
+    let _ = serde_json::to_string(&all_span_links)
+        .map(|json| meta.insert("_dd.span_links".to_string(), json));
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use serde_json::json;
+    use std::collections::HashMap;
+
+    #[derive(Debug, Default)]
+    struct TestSpan {
+        pub meta: HashMap<String, String>,
+    }
+
+    struct SpanPointerTestCase {
+        test_name: &'static str,
+        existing_links: Option<serde_json::Value>,
+        span_pointers: Option<Vec<SpanPointer>>,
+        expected_links: Option<serde_json::Value>,
+    }
+
+    #[test]
+    fn test_attach_span_pointers_to_span() {
+        let test_cases = vec![
+            SpanPointerTestCase {
+                test_name: "adds span links to span",
+                existing_links: None,
+                span_pointers: Some(vec![
+                    SpanPointer {
+                        hash: "hash1".to_string(),
+                        kind: "test.kind1".to_string(),
+                    },
+                    SpanPointer {
+                        hash: "hash2".to_string(),
+                        kind: "test.kind2".to_string(),
+                    },
+                ]),
+                expected_links: Some(json!([
+                    {
+                        "attributes": {
+                            "link.kind": "span-pointer",
+                            "ptr.dir": "u",
+                            "ptr.hash": "hash1",
+                            "ptr.kind": "test.kind1"
+                        },
+                        "span_id": 0,
+                        "trace_id": 0,
+                        "trace_id_high": 0,
+                        "tracestate": "",
+                        "flags": 0
+                    },
+                    {
+                        "attributes": {
+                            "link.kind": "span-pointer",
+                            "ptr.dir": "u",
+                            "ptr.hash": "hash2",
+                            "ptr.kind": "test.kind2"
+                        },
+                        "span_id": 0,
+                        "trace_id": 0,
+                        "trace_id_high": 0,
+                        "tracestate": "",
+                        "flags": 0
+                    }
+                ])),
+            },
+            SpanPointerTestCase {
+                test_name: "handles empty span pointers",
+                existing_links: None,
+                span_pointers: Some(vec![]),
+                expected_links: None,
+            },
+            SpanPointerTestCase {
+                test_name: "handles None span pointers",
+                existing_links: None,
+                span_pointers: None,
+                expected_links: None,
+            },
+            SpanPointerTestCase {
+                test_name: "appends to existing span links",
+                existing_links: Some(json!([{
+                    "attributes": {
+                        "link.kind": "span-pointer",
+                        "ptr.dir": "d",
+                        "ptr.hash": "hash1",
+                        "ptr.kind": "test.kind1"
+                    },
+                    "span_id": 0,
+                    "trace_id": 0,
+                    "trace_id_high": 0,
+                    "tracestate": "",
+                    "flags": 0
+                }])),
+                span_pointers: Some(vec![SpanPointer {
+                    hash: "hash2".to_string(),
+                    kind: "test.kind2".to_string(),
+                }]),
+                expected_links: Some(json!([
+                    {
+                        "attributes": {
+                            "link.kind": "span-pointer",
+                            "ptr.dir": "d",
+                            "ptr.hash": "hash1",
+                            "ptr.kind": "test.kind1"
+                        },
+                        "span_id": 0,
+                        "trace_id": 0,
+                        "trace_id_high": 0,
+                        "tracestate": "",
+                        "flags": 0
+                    },
+                    {
+                        "attributes": {
+                            "link.kind": "span-pointer",
+                            "ptr.dir": "u",
+                            "ptr.hash": "hash2",
+                            "ptr.kind": "test.kind2"
+                        },
+                        "span_id": 0,
+                        "trace_id": 0,
+                        "trace_id_high": 0,
+                        "tracestate": "",
+                        "flags": 0
+                    }
+                ])),
+            },
+        ];
+
+        for case in test_cases {
+            let mut test_span = TestSpan {
+                meta: HashMap::new(),
+            };
+
+            // Set up existing links if any
+            if let Some(links) = case.existing_links {
+                test_span
+                    .meta
+                    .insert("_dd.span_links".to_string(), links.to_string());
+            }
+
+            attach_span_pointers_to_meta(&mut test_span.meta, &case.span_pointers);
+
+            match case.expected_links {
+                Some(expected) => {
+                    let span_links = test_span.meta.get("_dd.span_links").unwrap_or_else(|| {
+                        panic!(
+                            "[{}] _dd.span_links should be present in span meta",
+                            case.test_name
+                        )
+                    });
+                    let actual_links: serde_json::Value =
+                        serde_json::from_str(span_links).expect("Should be valid JSON");
+                    assert_eq!(
+                        actual_links, expected,
+                        "Failed test case: {}",
+                        case.test_name
+                    );
+                }
+                None => {
+                    assert!(
+                        !test_span.meta.contains_key("_dd.span_links"),
+                        "Failed test case: {}",
+                        case.test_name
+                    );
+                }
+            }
+        }
+    }
+
+    #[test]
+    fn test_generate_span_pointer_hash() {
+        let test_cases = vec![
+            (
+                "basic values",
+                vec!["some-bucket", "some-key.data", "ab12ef34"],
+                "e721375466d4116ab551213fdea08413",
+            ),
+            (
+                "non-ascii key",
+                vec!["some-bucket", "some-key.你好", "ab12ef34"],
+                "d1333a04b9928ab462b5c6cadfa401f4",
+            ),
+            (
+                "multipart-upload",
+                vec!["some-bucket", "some-key.data", "ab12ef34-5"],
+                "2b90dffc37ebc7bc610152c3dc72af9f",
+            ),
+        ];
+
+        for (name, components, expected_hash) in test_cases {
+            let actual_hash = generate_span_pointer_hash(&components);
+            assert_eq!(actual_hash, expected_hash, "Test case: {name}");
+        }
+    }
+}