feat(server-utils): Port SQS, SNS and Lambda aws-sdk extensions with trace propagation

packages/server-utils/src/integrations/tracing-channel/aws-sdk/aws-sdk.types.ts

@@ -0,0 +1,59 @@
+/*
+ * AWS SDK for JavaScript
+ * Copyright 2012-2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * This product includes software developed at
+ * Amazon Web Services, Inc. (http://aws.amazon.com/).
+ */
+
+/*
+  These are slightly modified and simplified versions of the actual SQS/SNS types included
+  in the official distribution:
+  https://github.com/aws/aws-sdk-js/blob/master/clients/sqs.d.ts
+  These are brought here to avoid having users install the `aws-sdk` whenever they
+  require this instrumentation.
+*/
+
+// eslint-disable-next-line @typescript-eslint/no-empty-object-type
+interface Blob {}
+type Binary = Buffer | Uint8Array | Blob | string;
+
+// eslint-disable-next-line @typescript-eslint/no-namespace -- Prefer to contain the types copied over in one location
+export namespace SNS {
+  interface MessageAttributeValue {
+    DataType: string;
+    StringValue?: string;
+    BinaryValue?: Binary;
+  }
+
+  export type MessageAttributeMap = { [key: string]: MessageAttributeValue };
+}
+
+// eslint-disable-next-line @typescript-eslint/no-namespace -- Prefer to contain the types copied over in one location
+export namespace SQS {
+  type StringList = string[];
+  type BinaryList = Binary[];
+  interface MessageAttributeValue {
+    StringValue?: string;
+    BinaryValue?: Binary;
+    StringListValues?: StringList;
+    BinaryListValues?: BinaryList;
+    DataType: string;
+  }
+
+  export type MessageBodyAttributeMap = {
+    [key: string]: MessageAttributeValue;
+  };
+
+  type MessageSystemAttributeMap = { [key: string]: string };
+
+  export interface Message {
+    MessageId?: string;
+    ReceiptHandle?: string;
+    MD5OfBody?: string;
+    Body?: string;
+    Attributes?: MessageSystemAttributeMap;
+    MD5OfMessageAttributes?: string;
+    MessageAttributes?: MessageBodyAttributeMap;
+  }
+}

packages/server-utils/src/integrations/tracing-channel/aws-sdk/constants.ts

@@ -1,11 +1,18 @@
 /**
  * AWS-specific span constants used by the aws-sdk channel integration that are NOT covered by
  * `@sentry/conventions/attributes` (attribute names that exist there are imported from there
- * directly). Per-service files append their own such constants below.
+ * directly). These are either Sentry-specific (the span origin), attribute *values* (not keys), or
+ * obsolete OTel conventions with no `@sentry/conventions` export.
  */
 
 /** The span origin every aws-sdk channel span carries, mirroring the uniform OTel `auto.otel.aws`. */
 export const AWS_SDK_ORIGIN = 'auto.aws.orchestrion.aws_sdk';
 
 /** DynamoDB `db.system` value (an attribute value, not a key, so not covered by conventions). */
 export const DB_SYSTEM_VALUE_DYNAMODB = 'dynamodb';
+
+// Messaging (obsolete OTel convention with no `@sentry/conventions` export, kept for parity)
+// TODO(v11): import from `@sentry/conventions` once a release including it ships (added in
+// getsentry/sentry-conventions#509), and drop this local constant.
+export const ATTR_MESSAGING_DESTINATION_KIND = 'messaging.destination_kind';
+export const MESSAGING_DESTINATION_KIND_VALUE_TOPIC = 'topic';

packages/server-utils/src/integrations/tracing-channel/aws-sdk/services/MessageAttributes.ts

@@ -0,0 +1,69 @@
+import type { SerializedTraceData } from '@sentry/core';
+import { debug, uniq } from '@sentry/core';
+import { DEBUG_BUILD } from '../../../../debug-build';
+import type { SNS, SQS } from '../aws-sdk.types';
+
+// https://docs.aws.amazon.com/AWSSimpleQueueService/latest/SQSDeveloperGuide/sqs-quotas.html
+export const MAX_MESSAGE_ATTRIBUTES = 10;
+
+// Sentry trace-propagation headers written into / read from AWS message attributes.
+const SENTRY_TRACE_HEADER = 'sentry-trace';
+const BAGGAGE_HEADER = 'baggage';
+const PROPAGATION_FIELDS = [SENTRY_TRACE_HEADER, BAGGAGE_HEADER];
+
+export interface AwsSdkContextObject {
+  [key: string]: {
+    StringValue?: string;
+    Value?: string;
+  };
+}
+
+/**
+ * Inject trace-propagation headers (from `getTraceData({ span })`) into an SQS/SNS message-attribute
+ * map, so the consumer can continue the trace. Respects the SQS 10-attribute quota. Mirrors the OTel
+ * integration's `injectPropagationContext`, but writes Sentry's `sentry-trace`/`baggage` instead of
+ * W3C headers. Callers pass the precomputed headers so batch sends serialize them only once.
+ */
+export function injectPropagationContext(
+  attributesMap: SQS.MessageBodyAttributeMap | SNS.MessageAttributeMap | undefined,
+  traceData: SerializedTraceData,
+): SQS.MessageBodyAttributeMap | SNS.MessageAttributeMap {
+  const attributes = attributesMap ?? {};
+  const headerKeys = Object.keys(traceData) as (keyof SerializedTraceData)[];
+
+  if (Object.keys(attributes).length + headerKeys.length <= MAX_MESSAGE_ATTRIBUTES) {
+    for (const key of headerKeys) {
+      const value = traceData[key];
+      if (value) {
+        // Index-assigning into the SQS/SNS map union needs one concrete map type; the written value
+        // shape is valid for both.
+        (attributes as SQS.MessageBodyAttributeMap)[key] = { DataType: 'String', StringValue: value };
+      }
+    }
+  } else {
+    DEBUG_BUILD &&
+      debug.warn(
+        '[orchestrion:aws-sdk] cannot set trace propagation on SQS/SNS message due to maximum amount of MessageAttributes',
+      );
+  }
+  return attributes;
+}
+
+/** Read the propagation headers back off a received SQS message, if present. */
+export function extractPropagationHeaders(
+  message: SQS.Message,
+): { sentryTrace?: string; baggage?: string } | undefined {
+  const carrier = (message.MessageAttributes ?? {}) as AwsSdkContextObject;
+  const sentryTrace = carrier[SENTRY_TRACE_HEADER]?.StringValue ?? carrier[SENTRY_TRACE_HEADER]?.Value;
+  if (!sentryTrace) {
+    return undefined;
+  }
+  return {
+    sentryTrace,
+    baggage: carrier[BAGGAGE_HEADER]?.StringValue ?? carrier[BAGGAGE_HEADER]?.Value,
+  };
+}
+
+export function addPropagationFieldsToAttributeNames(messageAttributeNames: string[] = []): string[] {
+  return uniq([...messageAttributeNames, ...PROPAGATION_FIELDS]);
+}

packages/server-utils/src/integrations/tracing-channel/aws-sdk/services/ServicesExtensions.ts

@@ -2,9 +2,12 @@ import type { Span } from '@sentry/core';
 import type { NormalizedRequest, NormalizedResponse, RequestMetadata } from '../types';
 import { DynamodbServiceExtension } from './dynamodb';
 import { KinesisServiceExtension } from './kinesis';
+import { LambdaServiceExtension } from './lambda';
 import { S3ServiceExtension } from './s3';
 import { SecretsManagerServiceExtension } from './secretsmanager';
 import type { ServiceExtension } from './ServiceExtension';
+import { SnsServiceExtension } from './sns';
+import { SqsServiceExtension } from './sqs';
 import { StepFunctionsServiceExtension } from './stepfunctions';
 
 export class ServicesExtensions implements ServiceExtension {
@@ -13,7 +16,10 @@ export class ServicesExtensions implements ServiceExtension {
   private _services: Map<string, ServiceExtension> = new Map<string, ServiceExtension>([
     ['SecretsManager', new SecretsManagerServiceExtension()],
     ['SFN', new StepFunctionsServiceExtension()],
+    ['SQS', new SqsServiceExtension()],
+    ['SNS', new SnsServiceExtension()],
     ['DynamoDB', new DynamodbServiceExtension()],
+    ['Lambda', new LambdaServiceExtension()],
     ['S3', new S3ServiceExtension()],
     ['Kinesis', new KinesisServiceExtension()],
   ]);

packages/server-utils/src/integrations/tracing-channel/aws-sdk/services/lambda.ts

@@ -0,0 +1,85 @@
+import type { Span } from '@sentry/core';
+import { debug, getTraceData, SPAN_KIND } from '@sentry/core';
+import {
+  FAAS_EXECUTION as ATTR_FAAS_EXECUTION,
+  FAAS_INVOKED_NAME as ATTR_FAAS_INVOKED_NAME,
+  FAAS_INVOKED_PROVIDER as ATTR_FAAS_INVOKED_PROVIDER,
+  FAAS_INVOKED_REGION as ATTR_FAAS_INVOKED_REGION,
+} from '@sentry/conventions/attributes';
+import { DEBUG_BUILD } from '../../../../debug-build';
+import type { NormalizedRequest, NormalizedResponse } from '../types';
+import type { RequestMetadata, ServiceExtension } from './ServiceExtension';
+
+const INVOKE_COMMAND = 'Invoke';
+
+export class LambdaServiceExtension implements ServiceExtension {
+  public requestPreSpanHook(request: NormalizedRequest): RequestMetadata {
+    const functionName = request.commandInput?.FunctionName;
+
+    const spanAttributes: Record<string, unknown> = {};
+    let spanName: string | undefined;
+
+    if (request.commandName === INVOKE_COMMAND) {
+      spanAttributes[ATTR_FAAS_INVOKED_NAME] = functionName;
+      spanAttributes[ATTR_FAAS_INVOKED_PROVIDER] = 'aws';
+      spanName = `${functionName} ${INVOKE_COMMAND}`;
+    }
+
+    return {
+      spanAttributes,
+      spanKind: SPAN_KIND.CLIENT,
+      spanName,
+    };
+  }
+
+  public requestPostSpanHook(request: NormalizedRequest, span: Span): void {
+    if (request.commandName === INVOKE_COMMAND && request.commandInput) {
+      request.commandInput.ClientContext = injectLambdaPropagationContext(request.commandInput.ClientContext, span);
+    }
+  }
+
+  public responseHook(response: NormalizedResponse, span: Span): void {
+    if (response.request.commandName === INVOKE_COMMAND) {
+      // oxlint-disable-next-line typescript/no-deprecated -- old-semconv faas.execution, matched to the OTel aws-sdk integration
+      span.setAttribute(ATTR_FAAS_EXECUTION, response.requestId);
+      // Region resolves asynchronously after `requestPreSpanHook`, so it's backfilled onto the
+      // normalized request and read here (same timing as `cloud.region`).
+      if (response.request.region) {
+        span.setAttribute(ATTR_FAAS_INVOKED_REGION, response.request.region);
+      }
+    }
+  }
+}
+
+function injectLambdaPropagationContext(clientContext: string | undefined, span: Span): string | undefined {
+  try {
+    const propagatedContext = getTraceData({ span });
+
+    const parsedClientContext = clientContext ? JSON.parse(Buffer.from(clientContext, 'base64').toString('utf8')) : {};
+
+    const updatedClientContext = {
+      ...parsedClientContext,
+      custom: {
+        ...parsedClientContext.custom,
+        ...propagatedContext,
+      },
+    };
+
+    const encodedClientContext = Buffer.from(JSON.stringify(updatedClientContext)).toString('base64');
+
+    // The length of client context is capped at 3583 bytes of base64 encoded data
+    // (https://docs.aws.amazon.com/lambda/latest/dg/API_Invoke.html#API_Invoke_RequestSyntax)
+    if (encodedClientContext.length > 3583) {
+      DEBUG_BUILD &&
+        debug.warn(
+          '[orchestrion:aws-sdk] cannot set trace propagation on lambda invoke parameters due to ClientContext length limitations.',
+        );
+      return clientContext;
+    }
+
+    return encodedClientContext;
+  } catch (e) {
+    DEBUG_BUILD && debug.log('[orchestrion:aws-sdk] failed to set trace propagation on lambda ClientContext', e);
+    return clientContext;
+  }
+}

packages/server-utils/src/integrations/tracing-channel/aws-sdk/services/sns.ts

@@ -0,0 +1,75 @@
+import type { Span, SpanKindValue } from '@sentry/core';
+import { getTraceData, SPAN_KIND } from '@sentry/core';
+import {
+  AWS_SNS_TOPIC_ARN as ATTR_AWS_SNS_TOPIC_ARN,
+  MESSAGING_DESTINATION as ATTR_MESSAGING_DESTINATION,
+  MESSAGING_DESTINATION_NAME,
+  MESSAGING_SYSTEM,
+} from '@sentry/conventions/attributes';
+import { ATTR_MESSAGING_DESTINATION_KIND, MESSAGING_DESTINATION_KIND_VALUE_TOPIC } from '../constants';
+import type { NormalizedRequest, NormalizedResponse } from '../types';
+import { injectPropagationContext } from './MessageAttributes';
+import type { RequestMetadata, ServiceExtension } from './ServiceExtension';
+
+export class SnsServiceExtension implements ServiceExtension {
+  public requestPreSpanHook(request: NormalizedRequest): RequestMetadata {
+    let spanKind: SpanKindValue = SPAN_KIND.CLIENT;
+    let spanName = `SNS ${request.commandName}`;
+    const spanAttributes: Record<string, unknown> = {
+      [MESSAGING_SYSTEM]: 'aws.sns',
+    };
+
+    if (request.commandName === 'Publish') {
+      spanKind = SPAN_KIND.PRODUCER;
+
+      spanAttributes[ATTR_MESSAGING_DESTINATION_KIND] = MESSAGING_DESTINATION_KIND_VALUE_TOPIC;
+      const { TopicArn, TargetArn, PhoneNumber } = request.commandInput;
+      const destinationName = extractDestinationName(TopicArn, TargetArn, PhoneNumber);
+      // oxlint-disable-next-line typescript/no-deprecated -- old-semconv messaging.destination, matched to the OTel aws-sdk integration
+      spanAttributes[ATTR_MESSAGING_DESTINATION] = destinationName;
+      spanAttributes[MESSAGING_DESTINATION_NAME] = TopicArn || TargetArn || PhoneNumber || 'unknown';
+
+      spanName = `${PhoneNumber ? 'phone_number' : destinationName} send`;
+    }
+
+    const topicArn = request.commandInput?.TopicArn;
+    if (topicArn) {
+      spanAttributes[ATTR_AWS_SNS_TOPIC_ARN] = topicArn;
+    }
+
+    return {
+      spanAttributes,
+      spanKind,
+      spanName,
+    };
+  }
+
+  public requestPostSpanHook(request: NormalizedRequest, span: Span): void {
+    if (request.commandName === 'Publish') {
+      const origMessageAttributes = request.commandInput.MessageAttributes ?? {};
+      request.commandInput.MessageAttributes = injectPropagationContext(origMessageAttributes, getTraceData({ span }));
+    }
+  }
+
+  public responseHook(response: NormalizedResponse, span: Span): void {
+    const topicArn = response.data?.TopicArn;
+    if (topicArn) {
+      span.setAttribute(ATTR_AWS_SNS_TOPIC_ARN, topicArn);
+    }
+  }
+}
+
+function extractDestinationName(topicArn: string, targetArn: string, phoneNumber: string): string {
+  if (topicArn || targetArn) {
+    const arn = topicArn ?? targetArn;
+    try {
+      return arn.substring(arn.lastIndexOf(':') + 1);
+    } catch {
+      return arn;
+    }
+  } else if (phoneNumber) {
+    return phoneNumber;
+  } else {
+    return 'unknown';
+  }
+}