Function wrapStreamingHandler

packages/aws-serverless/src/sdk.ts:263–315

(
  handler: StreamifyHandler<TEvent, TResult>,
  options: WrapperOptions,
  startTime: number,
)

}

function wrapStreamingHandler<TEvent, TResult>(
  handler: StreamifyHandler<TEvent, TResult>,
  options: WrapperOptions,
  startTime: number,
): StreamifyHandler<TEvent, TResult> {
  let timeoutWarningTimer: NodeJS.Timeout | undefined;

  const wrappedHandler = async (
    event: TEvent,
    responseStream: Parameters<StreamifyHandler<TEvent, TResult>>[1],
    context: Context,
  ): Promise<TResult> => {
    context.callbackWaitsForEmptyEventLoop = options.callbackWaitsForEmptyEventLoop;

    timeoutWarningTimer = setupTimeoutWarning(context, options);

    async function processStreamingResult(): Promise<TResult> {
      const scope = getCurrentScope();

      try {
        enhanceScopeWithEnvironmentData(scope, context, startTime);

        responseStream.on('error', error => {
          captureException(error, scope => markEventUnhandled(scope, 'auto.function.aws_serverless.stream'));
        });

        return await handler(event, responseStream, context);
      } catch (e) {
        // Errors should already be captured in the AwsLambda instrumentation's error handler,
        // we capture them here just to be safe. Double captures are deduplicated by the SDK.
        captureException(e, scope => markEventUnhandled(scope, 'auto.function.aws_serverless.handler'));
        throw e;
      } finally {
        if (timeoutWarningTimer) {
          clearTimeout(timeoutWarningTimer);
        }
        await flush(options.flushTimeout).catch(e => {
          DEBUG_BUILD && debug.error(e);
        });
      }
    }

    return withScope(() => processStreamingResult());
  };

  const handlerWithSymbols = handler as unknown as Record<symbol, unknown>;
  (wrappedHandler as unknown as Record<symbol, unknown>)[AWS_HANDLER_STREAMING_SYMBOL] =
    handlerWithSymbols[AWS_HANDLER_STREAMING_SYMBOL];
  (wrappedHandler as unknown as Record<symbol, unknown>)[AWS_HANDLER_HIGHWATERMARK_SYMBOL] =
    handlerWithSymbols[AWS_HANDLER_HIGHWATERMARK_SYMBOL];

  return wrappedHandler;
}