fix: different-x-ray-trace-ids-for-incoming-http-api-requests (#40)

Author: mosheshaham-dash0

integration-tests/iac/lib/python-tracing-scenarios-stack.ts

@@ -10,6 +10,7 @@ import * as kinesis from 'aws-cdk-lib/aws-kinesis';
 import * as s3 from 'aws-cdk-lib/aws-s3';
 import * as s3n from 'aws-cdk-lib/aws-s3-notifications';
 import * as apigateway from 'aws-cdk-lib/aws-apigateway';
+import * as apigatewayv2 from 'aws-cdk-lib/aws-apigatewayv2';
 import * as events from 'aws-cdk-lib/aws-events';
 import * as events_targets from 'aws-cdk-lib/aws-events-targets';
 import * as lambda_event_sources from 'aws-cdk-lib/aws-lambda-event-sources';
@@ -298,6 +299,65 @@ export class PythonTracingScenariosStack extends cdk.NestedStack {
         },
       });
 
+      // Scenario 6b: Lambda > HTTP API Gateway > Lambda
+      const httpApiConsumer = new lambda.Function(this, `HttpApiConsumerLambda-${runtimeName}`, {
+        functionName: `${prefix}tracing-httpapi-consumer-${runtimeName}`,
+        runtime,
+        handler: 'consumer.handler',
+        code: pythonCode,
+        layers: [props.layer],
+        role,
+        timeout: cdk.Duration.seconds(10),
+        logGroup: props.logGroup,
+        environment: baseEnvironment,
+      });
+
+      const httpApi = new apigatewayv2.CfnApi(this, `TracingTestHttpApi-${runtimeName}`, {
+        name: `${prefix}tracing-test-http-api-${runtimeName}`,
+        protocolType: 'HTTP',
+      });
+
+      const httpApiIntegration = new apigatewayv2.CfnIntegration(this, `TracingTestHttpApiIntegration-${runtimeName}`, {
+        apiId: httpApi.ref,
+        integrationType: 'AWS_PROXY',
+        integrationUri: httpApiConsumer.functionArn,
+        payloadFormatVersion: '2.0',
+      });
+
+      new apigatewayv2.CfnRoute(this, `TracingTestHttpApiRoute-${runtimeName}`, {
+        apiId: httpApi.ref,
+        routeKey: 'POST /',
+        target: `integrations/${httpApiIntegration.ref}`,
+      });
+
+      new apigatewayv2.CfnStage(this, `TracingTestHttpApiStage-${runtimeName}`, {
+        apiId: httpApi.ref,
+        stageName: '$default',
+        autoDeploy: true,
+      });
+
+      httpApiConsumer.addPermission(`HttpApiInvokePermission-${runtimeName}`, {
+        principal: new iam.ServicePrincipal('apigateway.amazonaws.com'),
+        sourceArn: `arn:aws:execute-api:${this.region}:${this.account}:${httpApi.ref}/*/*`,
+      });
+
+      const httpApiUrl = `https://${httpApi.ref}.execute-api.${this.region}.amazonaws.com/`;
+
+      const httpApiProducer = new lambda.Function(this, `HttpApiProducerLambda-${runtimeName}`, {
+        functionName: `${prefix}tracing-httpapi-producer-${runtimeName}`,
+        runtime,
+        handler: 'apigateway_producer.handler',
+        code: pythonCode,
+        layers: [props.layer],
+        role,
+        timeout: cdk.Duration.seconds(10),
+        logGroup: props.logGroup,
+        environment: {
+          ...baseEnvironment,
+          API_URL: httpApiUrl,
+        },
+      });
+
       // Scenario 7: Lambda > S3 > Lambda
       const s3Bucket = new s3.Bucket(this, `TracingTestS3Bucket-${runtimeName}`, {
         bucketName: `${prefix}tracing-test-s3-bucket-${runtimeName}`,

integration-tests/tests/src/test-tracing-scenarios-general.test.ts

@@ -11,6 +11,7 @@ const nodeRuntimes = ['nodejs20-x', 'nodejs22-x', 'nodejs24-x'];
 const scenarios = [
     { name: 'eventbridge', producerPrefix: `${RESOURCE_PREFIX}tracing-eventbridge-producer`, consumerPrefix: `${RESOURCE_PREFIX}tracing-eventbridge-consumer`, runtimes: [...pythonRuntimes, ...nodeRuntimes] },
     { name: 'apigateway', producerPrefix: `${RESOURCE_PREFIX}tracing-apigateway-producer`, consumerPrefix: `${RESOURCE_PREFIX}tracing-apigateway-consumer`, runtimes: [...pythonRuntimes, ...nodeRuntimes] },
+    { name: 'httpapi', producerPrefix: `${RESOURCE_PREFIX}tracing-httpapi-producer`, consumerPrefix: `${RESOURCE_PREFIX}tracing-httpapi-consumer`, runtimes: [...pythonRuntimes] },
     { name: 's3', producerPrefix: `${RESOURCE_PREFIX}tracing-s3-producer`, consumerPrefix: `${RESOURCE_PREFIX}tracing-s3-consumer`, runtimes: [...pythonRuntimes, ...nodeRuntimes] },
     { name: 'lambda', producerPrefix: `${RESOURCE_PREFIX}tracing-lambda-invoker`, consumerPrefix: `${RESOURCE_PREFIX}tracing-lambda-consumer`, runtimes: [...pythonRuntimes, ...nodeRuntimes] },
 ] as const;

src/otlp/span_creation.rs

@@ -48,15 +48,12 @@ fn create_root_span(
     let span_id = hex::decode(root_span_id).ok()?;
     let trace_id = hex::decode(trace_id_hex).ok()?;
 
-    let parent_span_id = if data.sampled || !crate::config::user::is_remove_lambda_parent_span() {
-        data.parent_span_id
-            .as_deref()
-            .filter(|s| !s.is_empty())
-            .and_then(|p| hex::decode(p).ok())
-            .unwrap_or_default()
-    } else {
-        Vec::new()
-    };
+    let parent_span_id = data
+        .parent_span_id
+        .as_deref()
+        .filter(|s| !s.is_empty())
+        .and_then(|p| hex::decode(p).ok())
+        .unwrap_or_default();
 
     let start_nanos = ((data.start_time - data.init_duration) * 1_000_000.0) as u64;
     let end_nanos = (data.end_time * 1_000_000.0) as u64;

src/otlp/span_mutations.rs

@@ -455,6 +455,12 @@ fn add_event_payload_to_span(span: &mut Span, invocation_id: &str) {
 }
 
 fn reparent_to_root_span(span: &mut Span, invocation_id: &str) {
+    tracing::info!(
+        "[{}] reparenting span for invocation_id={}: original parent_span_id={}",
+        crate::log_prefix(),
+        invocation_id,
+        hex::encode(&span.parent_span_id),
+    );
     let root_span_id = invocation_entry::get_or_create_root_span_id(invocation_id);
     if let Ok(bytes) = hex::decode(&root_span_id) {
         span.parent_span_id = bytes;
@@ -482,14 +488,17 @@ fn store_span_ids(span: &Span, invocation_id: &str) {
         .map(|b| format!("{:02x}", b))
         .collect::<String>();
     invocation_entry::update(invocation_id, |entry| {
+        let is_same_trace = entry
+            .trace_id
+            .as_ref()
+            .map_or(false, |existing| existing == &trace_id_hex);
         entry.trace_id = Some(trace_id_hex);
         entry.span_id = Some(span_id_hex);
+        // we update parent span only if the lambda instrumentation extracted the context not from _X_AMZN_TRACE_ID
+        if !span.parent_span_id.is_empty() && !is_same_trace {
+            entry.parent_span_id = Some(hex::encode(&span.parent_span_id));
+        }
     });
-    tracing::debug!(
-        "[{}] stored trace/span id for invocation_id={}",
-        crate::log_prefix(),
-        invocation_id
-    );
 }
 
 /// Process a decoded trace request by adding event payloads, return payloads, and storing invocation span IDs.
@@ -522,8 +531,8 @@ pub fn process_trace_request(
                 span.name = "handler".to_string();
                 span.kind = SpanKind::Internal as i32;
                 add_event_payload_to_span(span, &invocation_id);
-                reparent_to_root_span(span, &invocation_id);
                 store_span_ids(span, &invocation_id);
+                reparent_to_root_span(span, &invocation_id);
                 store_handler_span_data(span, &invocation_id);
             }
         }