Merge branch 'main' into joey/handle_msk

Author: joeyzhao2018

bottlecap/src/otlp/agent.rs

@@ -19,7 +19,7 @@ use tracing::{debug, error};
 use crate::{
     config::Config,
     http::{extract_request_body, handler_not_found},
-    otlp::processor::Processor as OtlpProcessor,
+    otlp::processor::{OtlpEncoding, Processor as OtlpProcessor},
     tags::provider,
     traces::{
         stats_generator::StatsGenerator, trace_aggregator::SendDataBuilderInfo,
@@ -141,24 +141,33 @@ impl Agent {
         >,
         request: Request,
     ) -> Response {
+        let content_type = request
+            .headers()
+            .get(header::CONTENT_TYPE)
+            .and_then(|v| v.to_str().ok())
+            .map(String::from);
+        let encoding = OtlpEncoding::from_content_type(content_type.as_deref());
+
         let (parts, body) = match extract_request_body(request).await {
             Ok(r) => r,
             Err(e) => {
                 error!("OTLP | Failed to extract request body: {e}");
                 return Self::otlp_error_response(
                     StatusCode::BAD_REQUEST,
                     format!("Failed to extract request body: {e}"),
+                    encoding,
                 );
             }
         };
 
-        let traces = match processor.process(&body) {
+        let traces = match processor.process(&body, encoding) {
             Ok(traces) => traces,
             Err(e) => {
                 error!("OTLP | Failed to process request: {:?}", e);
                 return Self::otlp_error_response(
                     StatusCode::INTERNAL_SERVER_ERROR,
                     format!("Failed to process request: {e}"),
+                    encoding,
                 );
             }
         };
@@ -170,6 +179,7 @@ impl Agent {
             return Self::otlp_error_response(
                 StatusCode::INTERNAL_SERVER_ERROR,
                 "Not sending traces, processor returned empty data".to_string(),
+                encoding,
             );
         }
 
@@ -193,6 +203,7 @@ impl Agent {
                     return Self::otlp_error_response(
                         StatusCode::INTERNAL_SERVER_ERROR,
                         format!("Error sending traces to the trace aggregator: {err}"),
+                        encoding,
                     );
                 }
             }
@@ -208,49 +219,88 @@ impl Agent {
             error!("OTLP | Error sending traces to the stats concentrator: {err}");
         }
 
-        Self::otlp_success_response()
+        Self::otlp_success_response(encoding)
     }
 
-    fn otlp_error_response(status_code: StatusCode, message: String) -> Response {
-        let status = Status {
-            code: i32::from(status_code.as_u16()),
-            message: message.clone(),
-            details: vec![],
+    fn otlp_error_response(
+        status_code: StatusCode,
+        message: String,
+        encoding: OtlpEncoding,
+    ) -> Response {
+        let body = match encoding {
+            OtlpEncoding::Json => {
+                let json_error = serde_json::json!({
+                    "code": status_code.as_u16(),
+                    "message": message
+                });
+                match serde_json::to_vec(&json_error) {
+                    Ok(buf) => buf,
+                    Err(e) => {
+                        error!("OTLP | Failed to encode error response as JSON: {e}");
+                        return (status_code, [(header::CONTENT_TYPE, "text/plain")], message)
+                            .into_response();
+                    }
+                }
+            }
+            OtlpEncoding::Protobuf => {
+                let status = Status {
+                    code: i32::from(status_code.as_u16()),
+                    message: message.clone(),
+                    details: vec![],
+                };
+                let mut buf = Vec::new();
+                if let Err(e) = status.encode(&mut buf) {
+                    error!("OTLP | Failed to encode error response as Protobuf: {e}");
+                    return (status_code, [(header::CONTENT_TYPE, "text/plain")], message)
+                        .into_response();
+                }
+                buf
+            }
         };
 
-        let mut buf = Vec::new();
-        if let Err(e) = status.encode(&mut buf) {
-            error!("OTLP | Failed to encode error response: {e}");
-            return (status_code, [(header::CONTENT_TYPE, "text/plain")], message).into_response();
-        }
-
         (
             status_code,
-            [(header::CONTENT_TYPE, "application/x-protobuf")],
-            buf,
+            [(header::CONTENT_TYPE, encoding.content_type())],
+            body,
         )
             .into_response()
     }
 
-    fn otlp_success_response() -> Response {
+    fn otlp_success_response(encoding: OtlpEncoding) -> Response {
         let response = ExportTraceServiceResponse {
             partial_success: None,
         };
 
-        let mut buf = Vec::new();
-        if let Err(e) = response.encode(&mut buf) {
-            error!("OTLP | Failed to encode success response: {e}");
-            return (
-                StatusCode::INTERNAL_SERVER_ERROR,
-                "Failed to encode response".to_string(),
-            )
-                .into_response();
-        }
+        let body = match encoding {
+            OtlpEncoding::Json => match serde_json::to_vec(&response) {
+                Ok(buf) => buf,
+                Err(e) => {
+                    error!("OTLP | Failed to encode success response as JSON: {e}");
+                    return (
+                        StatusCode::INTERNAL_SERVER_ERROR,
+                        "Failed to encode response".to_string(),
+                    )
+                        .into_response();
+                }
+            },
+            OtlpEncoding::Protobuf => {
+                let mut buf = Vec::new();
+                if let Err(e) = response.encode(&mut buf) {
+                    error!("OTLP | Failed to encode success response as Protobuf: {e}");
+                    return (
+                        StatusCode::INTERNAL_SERVER_ERROR,
+                        "Failed to encode response".to_string(),
+                    )
+                        .into_response();
+                }
+                buf
+            }
+        };
 
         (
             StatusCode::OK,
-            [(header::CONTENT_TYPE, "application/x-protobuf")],
-            buf,
+            [(header::CONTENT_TYPE, encoding.content_type())],
+            body,
         )
             .into_response()
     }

bottlecap/src/otlp/processor.rs

@@ -5,6 +5,30 @@ use std::{error::Error, sync::Arc};
 
 use crate::{config::Config, otlp::transform::otel_resource_spans_to_dd_spans};
 
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub enum OtlpEncoding {
+    Protobuf,
+    Json,
+}
+
+impl OtlpEncoding {
+    #[must_use]
+    pub fn from_content_type(content_type: Option<&str>) -> Self {
+        match content_type {
+            Some(ct) if ct.starts_with("application/json") => OtlpEncoding::Json,
+            _ => OtlpEncoding::Protobuf,
+        }
+    }
+
+    #[must_use]
+    pub fn content_type(&self) -> &'static str {
+        match self {
+            OtlpEncoding::Json => "application/json",
+            OtlpEncoding::Protobuf => "application/x-protobuf",
+        }
+    }
+}
+
 #[derive(Clone)]
 pub struct Processor {
     config: Arc<Config>,
@@ -16,9 +40,15 @@ impl Processor {
         Self { config }
     }
 
-    pub fn process(&self, body: &[u8]) -> Result<Vec<Vec<DatadogSpan>>, Box<dyn Error>> {
-        // Decode the OTLP HTTP request
-        let request = ExportTraceServiceRequest::decode(body)?;
+    pub fn process(
+        &self,
+        body: &[u8],
+        encoding: OtlpEncoding,
+    ) -> Result<Vec<Vec<DatadogSpan>>, Box<dyn Error>> {
+        let request = match encoding {
+            OtlpEncoding::Json => serde_json::from_slice::<ExportTraceServiceRequest>(body)?,
+            OtlpEncoding::Protobuf => ExportTraceServiceRequest::decode(body)?,
+        };
 
         let mut spans: Vec<Vec<DatadogSpan>> = Vec::new();
         for resource_spans in &request.resource_spans {

integration-tests/lambda/otlp-node/package-lock.json

@@ -5,6 +5,7 @@
   "main": "index.js",
   "dependencies": {
     "@opentelemetry/api": "^1.9.0",
+    "@opentelemetry/exporter-trace-otlp-http": "^0.54.2",
     "@opentelemetry/exporter-trace-otlp-proto": "^0.54.2",
     "@opentelemetry/otlp-transformer": "^0.54.2",
     "@opentelemetry/resources": "^1.28.0",

integration-tests/lambda/otlp-node/package.json

@@ -1,122 +1,48 @@
+const { NodeTracerProvider } = require('@opentelemetry/sdk-trace-node');
+const { SimpleSpanProcessor } = require('@opentelemetry/sdk-trace-base');
+const { OTLPTraceExporter: OTLPTraceExporterHttp } = require('@opentelemetry/exporter-trace-otlp-http');
+const { OTLPTraceExporter: OTLPTraceExporterProto } = require('@opentelemetry/exporter-trace-otlp-proto');
 const { Resource } = require('@opentelemetry/resources');
-const { BasicTracerProvider, InMemorySpanExporter } = require('@opentelemetry/sdk-trace-base');
-const { createExportTraceServiceRequest } = require('@opentelemetry/otlp-transformer');
-const { trace, SpanKind, SpanStatusCode } = require('@opentelemetry/api');
-
-const root = require('@opentelemetry/otlp-transformer/build/esm/generated/root');
-const ExportTraceServiceRequest = root.opentelemetry.proto.collector.trace.v1.ExportTraceServiceRequest;
-const ExportTraceServiceResponse = root.opentelemetry.proto.collector.trace.v1.ExportTraceServiceResponse;
+const { ATTR_SERVICE_NAME } = require('@opentelemetry/semantic-conventions');
 
 exports.handler = async (event, context) => {
-  console.log('Starting OTLP response validation...');
-  const result = await validateOtlpResponse();
-  return {
-    statusCode: result.success ? 200 : 500,
-    body: JSON.stringify(result, null, 2),
-  };
-};
-
-async function validateOtlpResponse() {
-  try {
-    const requestBuffer = createOtlpTraceRequest();
-    const response = await fetch('http://localhost:4318/v1/traces', {
-      method: 'POST',
-      headers: {
-        'Content-Type': 'application/x-protobuf',
-      },
-      body: requestBuffer,
-    });
-
-    if (response.status !== 200) {
-      return {
-        success: false,
-        error: `Expected status 200, got ${response.status}`,
-        statusCode: response.status,
-      };
-    }
-
-    const responseBuffer = Buffer.from(await response.arrayBuffer());
-
-    const contentType = response.headers.get('content-type');
-    if (contentType !== 'application/x-protobuf') {
-      return {
-        success: false,
-        error: `Expected Content-Type 'application/x-protobuf', got '${contentType}'`,
-        statusCode: response.status,
-        contentType,
-      };
-    }
-
-    let decodedResponse;
-    try {
-      decodedResponse = ExportTraceServiceResponse.decode(responseBuffer);
-    } catch (decodeError) {
-      return {
-        success: false,
-        error: `Failed to decode response as ExportTraceServiceResponse: ${decodeError.message}`,
-        statusCode: response.status,
-        contentType,
-        decodeError: decodeError.message,
-      };
-    }
+  const serviceName = context.functionName;
+  const requestId = context.awsRequestId;
+  const endpoint = process.env.OTEL_EXPORTER_OTLP_ENDPOINT || 'http://localhost:4318';
 
-    // Verify that partial_success is not set (should be None for success)
-    if (decodedResponse.partialSuccess) {
-      return {
-        success: false,
-        error: 'Unexpected partial_success in response - should only be success or failure',
-        statusCode: response.status,
-        contentType,
-        partialSuccess: decodedResponse.partialSuccess,
-      };
-    }
+  const jsonProvider = new NodeTracerProvider({
+    resource: new Resource({ [ATTR_SERVICE_NAME]: serviceName }),
+  });
+  const jsonExporter = new OTLPTraceExporterHttp({
+    url: endpoint + '/v1/traces',
+    headers: { 'Content-Type': 'application/json' },
+  });
+  jsonProvider.addSpanProcessor(new SimpleSpanProcessor(jsonExporter));
+  jsonProvider.register();
 
-    return {
-      success: true,
-      statusCode: response.status,
-      contentType,
-      response: decodedResponse,
-      message: 'OTLP response is properly formatted and decodable',
-    };
+  const jsonSpan = jsonProvider.getTracer('json-test').startSpan('test-span-json');
+  jsonSpan.setAttribute('request_id', requestId);
+  jsonSpan.setAttribute('encoding', 'json');
+  jsonSpan.end();
 
-  } catch (error) {
-    console.error('Validation error:', error);
-    return {
-      success: false,
-      error: `Unexpected error: ${error.message}`,
-      stack: error.stack,
-    };
-  }
-}
+  await jsonProvider.forceFlush();
+  await jsonProvider.shutdown();
 
-function createOtlpTraceRequest() {
-  const resource = new Resource({
-    'service.name': process.env.DD_SERVICE,
+  const protoProvider = new NodeTracerProvider({
+    resource: new Resource({ [ATTR_SERVICE_NAME]: serviceName }),
   });
-
-  const provider = new BasicTracerProvider({ resource });
-  const exporter = new InMemorySpanExporter();
-  provider.register();
-
-  const tracer = provider.getTracer('validation-tracer', '1.0.0');
-  const span = tracer.startSpan('test-span', {
-    kind: SpanKind.INTERNAL,
-    attributes: {
-      'test': 'validation',
-    },
+  const protoExporter = new OTLPTraceExporterProto({
+    url: endpoint + '/v1/traces',
   });
-  span.setStatus({ code: SpanStatusCode.OK });
-  span.end();
+  protoProvider.addSpanProcessor(new SimpleSpanProcessor(protoExporter));
 
-  provider.forceFlush();
+  const protoSpan = protoProvider.getTracer('proto-test').startSpan('test-span-protobuf');
+  protoSpan.setAttribute('request_id', requestId);
+  protoSpan.setAttribute('encoding', 'protobuf');
+  protoSpan.end();
 
-  const finishedSpans = exporter.getFinishedSpans();
-  const otlpRequest = createExportTraceServiceRequest(finishedSpans, {
-    useHex: true,
-    useLongBits: false,
-  });
+  await protoProvider.forceFlush();
+  await protoProvider.shutdown();
 
-  const message = ExportTraceServiceRequest.create(otlpRequest);
-  const buffer = ExportTraceServiceRequest.encode(message).finish();
-  return Buffer.from(buffer);
-}
+  return { statusCode: 200, body: JSON.stringify({ message: 'Success' }) };
+};