feat: Added support for response streaming Lambda functions (#2981)

Author: mrickard

lib/serverless/aws-lambda.js

@@ -24,6 +24,7 @@ const EVENT_SOURCE_TYPE_KEY = `${EVENT_SOURCE_PREFIX}.eventType`
 const NAMES = require('../metrics/names')
 
 const EVENT_SOURCE_INFO = require('./event-sources')
+const HANDLER_STREAMING = Symbol.for('aws.lambda.runtime.handler.streaming')
 
 // A function with no references used to stub out closures
 function cleanClosure() {}
@@ -35,7 +36,7 @@ let transactionEnders = []
 // the invocation transaction.
 let uncaughtException = null
 
-// Tracking the first time patchLambdaHandler is called for one off functionality
+// Tracking the first time patchLambdaHandler is called for one-off functionality
 let patchCalled = false
 let coldStartRecorded = false
 
@@ -105,6 +106,45 @@ class AwsLambda {
     })
   }
 
+  /**
+   * Response-streaming handlers are identified by symbol properties on the function.
+   * We propagate any symbols if they're present, so that the handler keeps its signatue for any AWS features that rely on symbols
+   * @param handler
+   * @param nrHandler
+   */
+  propagateSymbols(handler, nrHandler) {
+    for (const symbol of Object.getOwnPropertySymbols(handler)) {
+      logger.trace(`Setting symbol ${symbol.toString()} on handler`)
+      nrHandler[symbol] = handler[symbol]
+    }
+  }
+
+  createSegment({ event, context, transaction, recorder }) {
+    const shim = this.shim
+    const functionName = context.functionName
+    const group = NAMES.FUNCTION.PREFIX
+    const transactionName = group + functionName
+
+    const activeSegment = shim.tracer.getSegment()
+
+    transaction.setPartialName(transactionName)
+    const txnEnder = endTransaction.bind(null, transaction, transactionEnders.length)
+
+    transactionEnders.push(txnEnder)
+    const segment = shim.createSegment(functionName, recorder, activeSegment)
+    transaction.baseSegment = segment
+    const awsAttributes = this._getAwsAgentAttributes(event, context)
+    transaction.trace.attributes.addAttributes(ATTR_DEST.TRANS_COMMON, awsAttributes)
+
+    shim.agent.setLambdaArn(context.invokedFunctionArn)
+
+    shim.agent.setLambdaFunctionVersion(context.functionVersion)
+    segment.addSpanAttributes(awsAttributes)
+
+    segment.start()
+    return { segment, txnEnder }
+  }
+
   patchLambdaHandler(handler) {
     const awsLambda = this
     const shim = this.shim
@@ -114,6 +154,11 @@ class AwsLambda {
       return handler
     }
 
+    const isStreamHandler = handler[HANDLER_STREAMING] === 'response'
+    if (isStreamHandler) {
+      this.agent.recordSupportability('Nodejs/Serverless/Lambda/ResponseStreaming')
+    }
+
     if (!patchCalled) {
       // Only wrap emit on process the first time patch is called.
       patchCalled = true
@@ -122,52 +167,87 @@ class AwsLambda {
       this.wrapFatal()
     }
 
-    return shim.bindCreateTransaction(wrappedHandler, new specs.TransactionSpec({ type: shim.BG }))
+    const wrapper = isStreamHandler ? wrappedStreamHandler : wrappedHandler
+    const nrHandler = shim.bindCreateTransaction(wrapper, new specs.TransactionSpec({ type: shim.BG }))
+    awsLambda.propagateSymbols(handler, nrHandler)
+
+    return nrHandler
+
+    /**
+     * Wraps a response streaming lambda handler.
+     *
+     * Creates and applies segment based on function name, assigns attributes to transaction trace,
+     * listen when stream errors(log error), ends(end transaction)
+     *
+     * **Note**: AWS doesn't support response streaming with API gateway invoked lambdas.
+     * This means we do not handle that as it would require intercepting the stream to parse
+     * the response code and headers.
+     */
+    function wrappedStreamHandler() {
+      const transaction = shim.tracer.getTransaction()
+      if (!transaction) {
+        logger.trace('No active transaction, not wrapping streaming handler')
+        return handler.apply(this, arguments)
+      }
 
-    function wrappedHandler() {
       const args = shim.argsToArray.apply(shim, arguments)
-
       const event = args[0]
-      const context = args[1]
+      const context = args[2]
+      logger.trace('In stream handler, lambda function name', context?.functionName)
+      const { segment, txnEnder } = awsLambda.createSegment({ context, event, transaction, recorder: recordBackground })
+      args[1] = awsLambda.wrapStreamAndCaptureError(
+        transaction,
+        txnEnder,
+        args[1]
+      )
 
-      const functionName = context.functionName
-      const group = NAMES.FUNCTION.PREFIX
-      const transactionName = group + functionName
+      let res
+      try {
+        res = shim.applySegment(handler, segment, false, this, args)
+      } catch (err) {
+        uncaughtException = err
+        txnEnder()
+        throw err
+      }
+
+      return res
+    }
 
+    /**
+     * Wraps a non response streaming lambda handler.
+     *
+     * Creates and applies segment based on function name, assigns attributes to transaction trace,
+     * adds handlers if api gateway to wrap request/response
+     * wraps the callback(if present), wraps the context `done`, `succeed`, `fail methods`, intercepts promise
+     * and properly ends transaction
+     */
+    function wrappedHandler() {
       const transaction = shim.tracer.getTransaction()
       if (!transaction) {
+        logger.trace('No active transaction, not wrapping handler')
         return handler.apply(this, arguments)
       }
-      const activeSegment = shim.tracer.getSegment()
-
-      transaction.setPartialName(transactionName)
+      const args = shim.argsToArray.apply(shim, arguments)
 
+      const event = args[0]
+      const context = args[1]
+      logger.trace('Lambda function name', context?.functionName)
       const isApiGatewayLambdaProxy = apiGateway.isLambdaProxyEvent(event)
+      logger.trace('Is this Lambda event an API Gateway or ALB web proxy?', isApiGatewayLambdaProxy)
+      logger.trace('Lambda event keys', Object.keys(event))
       const segmentRecorder = isApiGatewayLambdaProxy ? recordWeb : recordBackground
-      const segment = shim.createSegment(functionName, segmentRecorder, activeSegment)
-      transaction.baseSegment = segment
+      const { segment, txnEnder } = awsLambda.createSegment({ context, event, transaction, recorder: segmentRecorder })
+
       // resultProcessor is used to execute additional logic based on the
       // payload supplied to the callback.
       let resultProcessor
 
-      logger.trace('Is this Lambda event an API Gateway or ALB web proxy?', isApiGatewayLambdaProxy)
-      logger.trace('Lambda event keys', Object.keys(event))
-
       if (isApiGatewayLambdaProxy) {
         const webRequest = new apiGateway.LambdaProxyWebRequest(event)
         setWebRequest(shim, transaction, webRequest)
         resultProcessor = getApiGatewayLambdaProxyResultProcessor(transaction)
       }
-
       const cbIndex = args.length - 1
-
-      // Add transaction ending closure to the list of functions to be called on
-      // beforeExit (i.e. in the case that context.{done,fail,succeed} or callback
-      // were not called).
-      const txnEnder = endTransaction.bind(null, transaction, transactionEnders.length)
-
-      transactionEnders.push(txnEnder)
-
       args[cbIndex] = awsLambda.wrapCallbackAndCaptureError(
         transaction,
         txnEnder,
@@ -186,16 +266,6 @@ class AwsLambda {
         }
       })
 
-      const awsAttributes = awsLambda._getAwsAgentAttributes(event, context)
-      transaction.trace.attributes.addAttributes(ATTR_DEST.TRANS_COMMON, awsAttributes)
-
-      shim.agent.setLambdaArn(context.invokedFunctionArn)
-
-      shim.agent.setLambdaFunctionVersion(context.functionVersion)
-      segment.addSpanAttributes(awsAttributes)
-
-      segment.start()
-
       let res
       try {
         res = shim.applySegment(handler, segment, false, this, args)
@@ -251,6 +321,18 @@ class AwsLambda {
     }
   }
 
+  wrapStreamAndCaptureError(transaction, txnEnder, stream) {
+    const shim = this.shim
+    stream.on('error', (error) => {
+      shim.agent.errors.add(transaction, error)
+    })
+
+    stream.on('close', () => {
+      txnEnder()
+    })
+    return stream
+  }
+
   _getAwsAgentAttributes(event, context) {
     const attributes = {
       'aws.lambda.arn': context.invokedFunctionArn,
@@ -372,7 +454,7 @@ function lowercaseObjectKeys(original) {
 }
 
 function endTransaction(transaction, enderIndex) {
-  if (transactionEnders[enderIndex] === cleanClosure) {
+  if (transactionEnders.length === 0 || transactionEnders[enderIndex] === cleanClosure) {
     // In the case where we have already been called, we return early. There may be a
     // case where this is called more than once, given the lambda is left in a dirty
     // state after thread suspension (e.g. timeouts)
@@ -411,7 +493,7 @@ function setWebResponse(transaction, response) {
 
     // We are adding http.statusCode to base segment as
     // we found in testing async invoked lambdas, the
-    // active segement is not available at this point.
+    // active segment is not available at this point.
     const segment = transaction.baseSegment
 
     segment.addSpanAttribute('http.statusCode', responseCode)