feat: Azure Function HTTP streaming support (newrelic#3070)

Author: amychisholm03

lib/instrumentation/@azure/functions.js

@@ -9,6 +9,7 @@ const defaultLogger = require('../../logger').child({ component: 'azure-function
 const urltils = require('../../util/urltils')
 const headerProcessing = require('../../header-processing')
 const synthetics = require('../../synthetics')
+const { Transform } = require('node:stream')
 
 const backgroundRecorder = require('../../metrics/recorders/other.js')
 const recordWeb = require('../../metrics/recorders/http')
@@ -134,24 +135,29 @@ function wrapAzureHttpMethods(shim, appMethod) {
 
       const newContext = ctx.enterSegment({ segment })
       const boundHandler = tracer.bindFunction(handler, newContext)
-
       const result = await boundHandler(...args)
+
       // Responses should have a shape as described at:
       // https://learn.microsoft.com/en-us/azure/azure-functions/functions-reference-node?tabs=javascript%2Cwindows%2Cazure-cli&pivots=nodejs-model-v4#http-response
-      if (result.status) {
+      if (result?.status) {
         tx.trace.attributes.addAttribute(
           DESTS.TRANS_COMMON,
           'http.statusCode',
           result.status
         )
       }
-
       if (coldStart === true) {
         tx.trace.attributes.addAttribute(DESTS.TRANS_COMMON, 'faas.coldStart', true)
         coldStart = false
       }
+      if (result?.body instanceof Transform) {
+        result.body.on('close', () => {
+          tx.end()
+        })
+      } else {
+        tx.end()
+      }
 
-      tx.end()
       return result
     })
 

lib/instrumentation/core/http.js

@@ -167,6 +167,7 @@ function instrumentedFinish(segment, transaction) {
 
   // Naming must happen before the segment and transaction are ended,
   // because metrics recording depends on naming's side effects.
+
   transaction.finalizeNameFromUri(transaction.parsedUrl, this.statusCode)
 
   if (this) {
@@ -444,6 +445,11 @@ module.exports = function initialize(agent, http, moduleName) {
     return false
   }
 
+  if (process.env.FUNCTIONS_WORKER_RUNTIME) {
+    logger.debug('In azure functions environment, disabling core http instrumentation in favor of @azure/functions')
+    return false
+  }
+
   const IS_HTTPS = moduleName === 'https'
 
   // FIXME: will this ever not be called?

test/unit/instrumentation/http/http.test.js

@@ -45,6 +45,7 @@ test('built-in http module instrumentation', async (t) => {
 
     t.afterEach((ctx) => {
       helper.unloadAgent(ctx.nr.agent)
+      process.env.FUNCTIONS_WORKER_RUNTIME = ''
     })
 
     await t.test('when passed no module', (t) => {
@@ -56,6 +57,23 @@ test('built-in http module instrumentation', async (t) => {
       const { agent, initialize } = t.nr
       assert.doesNotThrow(() => initialize(agent, {}, 'http', new Shim(agent, 'http')))
     })
+
+    await t.test('should not instrument if azure functions environment detected', (t) => {
+      const { agent, initialize } = t.nr
+      process.env.FUNCTIONS_WORKER_RUNTIME = 'node'
+      const http = {
+        request: function request(options) {
+          const requested = new EventEmitter()
+          requested.path = '/TEST'
+          if (options.path) {
+            requested.path = options.path
+          }
+
+          return requested
+        }
+      }
+      assert.equal(initialize(agent, http, 'http', new Shim(agent, 'http')), false)
+    })
   })
 
   await t.test('after loading', async (t) => {

test/versioned/azure-functions/http.test.js

@@ -10,6 +10,7 @@
 const test = require('node:test')
 const assert = require('node:assert')
 const { once } = require('node:events')
+const { Transform, Readable } = require('node:stream')
 
 const helper = require('../../lib/agent_helper.js')
 const { removeMatchedModules } = require('../../lib/cache-buster.js')
@@ -352,3 +353,48 @@ test('uses port provided in url', async (t) => {
   assert.ok(tx)
   assert.equal(tx.port, '8080')
 })
+
+test('ends transaction on stream close', async (t) => {
+  bootstrapModule({ t })
+  const { agent, initialize, mockApi, shim } = t.nr
+  initialize(agent, mockApi, MODULE_NAME, shim)
+
+  const handler = async function () {
+    const response = new AzureFunctionHttpResponse()
+    const stream = new Readable({
+      read() {
+        this.push('streamed data')
+        this.push(null) // End the stream
+      }
+    })
+    response.body = stream.pipe(new Transform({
+      transform(chunk, encoding, callback) {
+        this.push(chunk.toString())
+        callback()
+      }
+    }))
+    response.status = 200
+    return response
+  }
+  const options = { handler }
+
+  const txFinished = once(agent, 'transactionFinished')
+  mockApi.app.get('a-test', options)
+  const response = await mockApi.httpRequest('get')
+
+  response.body.on('data', (data) => {
+    assert.equal(data.toString(), 'streamed data')
+  })
+  await new Promise((resolve, reject) => {
+    response.body.on('close', async () => {
+      try {
+        const [tx] = await txFinished
+        assert.ok(tx)
+        assert.equal(tx.baseSegment.name, 'WebTransaction/AzureFunction/test-func')
+        resolve()
+      } catch (err) {
+        reject(err)
+      }
+    })
+  })
+})