feat(instrumentation-aws-lambda): support streaming handlers (#2970)

Author: msonnb

packages/instrumentation-aws-lambda/src/instrumentation.ts

@@ -48,6 +48,7 @@ import {
   Callback,
   Context,
   Handler,
+  StreamifyHandler,
 } from 'aws-lambda';
 
 import { AwsLambdaInstrumentationConfig, EventContextExtractor } from './types';
@@ -65,6 +66,10 @@ const headerGetter: TextMapGetter<APIGatewayProxyEventHeaders> = {
 };
 
 export const lambdaMaxInitInMilliseconds = 10_000;
+export const AWS_HANDLER_STREAMING_SYMBOL = Symbol.for(
+  'aws.lambda.runtime.handler.streaming'
+);
+export const AWS_HANDLER_STREAMING_RESPONSE = 'response';
 
 export class AwsLambdaInstrumentation extends InstrumentationBase<AwsLambdaInstrumentationConfig> {
   private declare _traceForceFlusher?: () => Promise<void>;
@@ -171,12 +176,31 @@ export class AwsLambdaInstrumentation extends InstrumentationBase<AwsLambdaInstr
   }
 
   private _getHandler(handlerLoadStartTime: number) {
-    return (original: Handler) => {
-      return this._getPatchHandler(original, handlerLoadStartTime);
+    return (
+      original: Handler | StreamifyHandler
+    ): Handler | StreamifyHandler => {
+      const patchedHandler = this._getPatchHandler(
+        original,
+        handlerLoadStartTime
+      );
+
+      if (this._isStreamingHandler(original)) {
+        // Streaming handlers have special symbols that we need to copy over to the patched handler.
+        for (const symbol of Object.getOwnPropertySymbols(original)) {
+          (patchedHandler as unknown as Record<symbol, unknown>)[symbol] = (
+            original as unknown as Record<symbol, unknown>
+          )[symbol];
+        }
+      }
+
+      return patchedHandler;
     };
   }
 
-  private _getPatchHandler(original: Handler, lambdaStartTime: number) {
+  private _getPatchHandler(
+    original: Handler | StreamifyHandler,
+    lambdaStartTime: number
+  ): Handler | StreamifyHandler {
     diag.debug('patch handler function');
     const plugin = this;
 
@@ -211,6 +235,43 @@ export class AwsLambdaInstrumentation extends InstrumentationBase<AwsLambdaInstr
       }
     }
 
+    if (this._isStreamingHandler(original)) {
+      return function patchedStreamingHandler(
+        this: never,
+        // The event can be a user type, it truly is any.
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        event: any,
+        responseStream: Parameters<StreamifyHandler>[1],
+        context: Context
+      ) {
+        _onRequest();
+
+        const parent = plugin._determineParent(event, context);
+        const span = plugin._createSpanForRequest(
+          event,
+          context,
+          requestIsColdStart,
+          parent
+        );
+        plugin._applyRequestHook(span, event, context);
+
+        return otelContext.with(trace.setSpan(parent, span), () => {
+          const maybePromise = safeExecuteInTheMiddle(
+            () => original.apply(this, [event, responseStream, context]),
+            error => {
+              if (error != null) {
+                // Exception thrown synchronously before resolving promise.
+                plugin._applyResponseHook(span, error);
+                plugin._endSpan(span, error, () => {});
+              }
+            }
+          ) as Promise<{}> | undefined;
+
+          return plugin._handlePromiseResult(span, maybePromise);
+        });
+      };
+    }
+
     return function patchedHandler(
       this: never,
       // The event can be a user type, it truly is any.
@@ -221,43 +282,16 @@ export class AwsLambdaInstrumentation extends InstrumentationBase<AwsLambdaInstr
     ) {
       _onRequest();
 
-      const config = plugin.getConfig();
-      const parent = AwsLambdaInstrumentation._determineParent(
+      const parent = plugin._determineParent(event, context);
+
+      const span = plugin._createSpanForRequest(
         event,
         context,
-        config.eventContextExtractor ||
-          AwsLambdaInstrumentation._defaultEventContextExtractor
-      );
-
-      const name = context.functionName;
-      const span = plugin.tracer.startSpan(
-        name,
-        {
-          kind: SpanKind.SERVER,
-          attributes: {
-            [ATTR_FAAS_EXECUTION]: context.awsRequestId,
-            [ATTR_FAAS_ID]: context.invokedFunctionArn,
-            [ATTR_CLOUD_ACCOUNT_ID]: AwsLambdaInstrumentation._extractAccountId(
-              context.invokedFunctionArn
-            ),
-            [ATTR_FAAS_COLDSTART]: requestIsColdStart,
-            ...AwsLambdaInstrumentation._extractOtherEventFields(event),
-          },
-        },
+        requestIsColdStart,
         parent
       );
 
-      const { requestHook } = config;
-      if (requestHook) {
-        safeExecuteInTheMiddle(
-          () => requestHook(span, { event, context }),
-          e => {
-            if (e)
-              diag.error('aws-lambda instrumentation: requestHook error', e);
-          },
-          true
-        );
-      }
+      plugin._applyRequestHook(span, event, context);
 
       return otelContext.with(trace.setSpan(parent, span), () => {
         // Lambda seems to pass a callback even if handler is of Promise form, so we wrap all the time before calling
@@ -275,27 +309,97 @@ export class AwsLambdaInstrumentation extends InstrumentationBase<AwsLambdaInstr
             }
           }
         ) as Promise<{}> | undefined;
-        if (typeof maybePromise?.then === 'function') {
-          return maybePromise.then(
-            value => {
-              plugin._applyResponseHook(span, null, value);
-              return new Promise(resolve =>
-                plugin._endSpan(span, undefined, () => resolve(value))
-              );
-            },
-            (err: Error | string) => {
-              plugin._applyResponseHook(span, err);
-              return new Promise((resolve, reject) =>
-                plugin._endSpan(span, err, () => reject(err))
-              );
-            }
-          );
-        }
-        return maybePromise;
+
+        return plugin._handlePromiseResult(span, maybePromise);
       });
     };
   }
 
+  private _createSpanForRequest(
+    event: any,
+    context: Context,
+    requestIsColdStart: boolean,
+    parent: OtelContext
+  ): Span {
+    const name = context.functionName;
+    return this.tracer.startSpan(
+      name,
+      {
+        kind: SpanKind.SERVER,
+        attributes: {
+          [ATTR_FAAS_EXECUTION]: context.awsRequestId,
+          [ATTR_FAAS_ID]: context.invokedFunctionArn,
+          [ATTR_CLOUD_ACCOUNT_ID]: AwsLambdaInstrumentation._extractAccountId(
+            context.invokedFunctionArn
+          ),
+          [ATTR_FAAS_COLDSTART]: requestIsColdStart,
+          ...AwsLambdaInstrumentation._extractOtherEventFields(event),
+        },
+      },
+      parent
+    );
+  }
+
+  private _applyRequestHook(span: Span, event: any, context: Context): void {
+    const { requestHook } = this.getConfig();
+    if (requestHook) {
+      safeExecuteInTheMiddle(
+        () => requestHook(span, { event, context }),
+        e => {
+          if (e) diag.error('aws-lambda instrumentation: requestHook error', e);
+        },
+        true
+      );
+    }
+  }
+
+  private _handlePromiseResult(
+    span: Span,
+    maybePromise: Promise<{}> | undefined
+  ): Promise<{}> | undefined {
+    if (typeof maybePromise?.then === 'function') {
+      return maybePromise.then(
+        value => {
+          this._applyResponseHook(span, null, value);
+          return new Promise(resolve =>
+            this._endSpan(span, undefined, () => resolve(value))
+          );
+        },
+        (err: Error | string) => {
+          this._applyResponseHook(span, err);
+          return new Promise((resolve, reject) =>
+            this._endSpan(span, err, () => reject(err))
+          );
+        }
+      );
+    }
+
+    // Handle synchronous return values by ending the span and applying response hook
+    this._applyResponseHook(span, null, maybePromise);
+    this._endSpan(span, undefined, () => {});
+    return maybePromise;
+  }
+
+  private _determineParent(event: any, context: Context): OtelContext {
+    const config = this.getConfig();
+    return AwsLambdaInstrumentation._determineParent(
+      event,
+      context,
+      config.eventContextExtractor ||
+        AwsLambdaInstrumentation._defaultEventContextExtractor
+    );
+  }
+
+  private _isStreamingHandler<TEvent, TResult>(
+    handler: Handler<TEvent, TResult> | StreamifyHandler<TEvent, TResult>
+  ): handler is StreamifyHandler<TEvent, TResult> {
+    return (
+      (handler as unknown as Record<symbol, unknown>)[
+        AWS_HANDLER_STREAMING_SYMBOL
+      ] === AWS_HANDLER_STREAMING_RESPONSE
+    );
+  }
+
   override setTracerProvider(tracerProvider: TracerProvider) {
     super.setTracerProvider(tracerProvider);
     this._traceForceFlusher = this._traceForceFlush(tracerProvider);

packages/instrumentation-aws-lambda/src/internal-types.ts

@@ -13,6 +13,6 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-import { Handler } from 'aws-lambda';
+import { Handler, StreamifyHandler } from 'aws-lambda';
 
-export type LambdaModule = Record<string, Handler>;
+export type LambdaModule = Record<string, Handler | StreamifyHandler>;