fix(instrumentation): wait for asyncEnd on promise settlement (#9538)

Author: juan-fernandez

packages/datadog-instrumentations/src/helpers/rewriter/transforms.js

@@ -17,45 +17,77 @@ const { parse, query } = require('./compiler')
 module.exports = { waitForAsyncEnd }
 
 /**
- * Injects a wait for `ctx.asyncEndPromise` into a generated `tracePromise`
- * wrapper's native-Promise fulfillment handler.
+ * Injects settlement-specific asyncEnd waits into a generated `tracePromise`
+ * wrapper's native-Promise handlers.
  *
  * @param {object} _state
  * @param {import('estree').CallExpression} node
  * @returns {void}
  */
 function waitForAsyncEnd (_state, node) {
   const onFulfilled = node.arguments[0]
-  const statements = onFulfilled?.body?.body
+  const onRejected = node.arguments[1]
 
-  if (!statements || query(onFulfilled.body, '[id.name=__apm$asyncEndPromise]').length > 0) {
+  if (!onFulfilled?.body || !onRejected?.body) {
     return
   }
 
-  const returnIndex = statements.findIndex(statement =>
-    statement.type === 'ReturnStatement' && statement.argument
+  injectAsyncEndCallbackWait(onFulfilled.body, 'ReturnStatement', 'resolveCallback')
+  injectAsyncEndCallbackWait(onRejected.body, 'ThrowStatement', 'rejectCallback')
+}
+
+/**
+ * Injects a settlement-specific callback wait before an exit.
+ *
+ * @param {import('estree').BlockStatement} body
+ * @param {'ReturnStatement'|'ThrowStatement'} exitType
+ * @param {'resolveCallback'|'rejectCallback'} callbackProperty
+ * @returns {void}
+ */
+function injectAsyncEndCallbackWait (body, exitType, callbackProperty) {
+  const callbackVariable = `__apm$${callbackProperty}`
+  if (query(body, `[id.name=${callbackVariable}]`).length > 0) {
+    return
+  }
+
+  const exitIndex = body.body.findIndex(statement =>
+    statement.type === exitType && statement.argument
   )
 
-  // The generated fulfillment handler always ends in a return; a miss means the
+  // The generated settlement handlers always end in a return or throw; a miss means the
   // upstream template changed and the caller's try/catch falls back to the
   // unwrapped source.
-  assert(returnIndex !== -1, 'waitForAsyncEnd: no return statement to wait on')
+  assert(exitIndex !== -1, `waitForAsyncEnd: no ${exitType} to wait on`)
 
+  // This runs inside tracePromise's native-Promise settlement handler. The
+  // Promise adapts subscriber callback completion to that existing chain.
   const waitStatements = parse(`
     function wrapper () {
-      const __apm$asyncEndPromise = __apm$ctx.asyncEndPromise;
-      if (__apm$asyncEndPromise && typeof __apm$asyncEndPromise.then === 'function') {
-        return __apm$asyncEndPromise.then(() => __apm$result, () => __apm$result);
+      const ${callbackVariable} = __apm$ctx.${callbackProperty};
+      if (typeof ${callbackVariable} === 'function') {
+        return new Promise(${callbackVariable}).then(() => __apm$result, () => __apm$result);
       }
     }
   `).body[0].body.body
 
-  // Resolve to whatever the fulfillment handler returns (its return argument),
-  // so a subscriber that reassigned `__apm$ctx.result` in `asyncEnd` still wins.
-  const returnArgument = statements[returnIndex].argument
-  const { arguments: onSettled } = waitStatements[1].consequent.body[0].argument
-  onSettled[0].body = clone(returnArgument)
-  onSettled[1].body = clone(returnArgument)
+  const exitArgument = body.body[exitIndex].argument
+  const callbackIf = waitStatements[1]
+  const { arguments: onCallbackSettled } = callbackIf.consequent.body[0].argument
+
+  if (exitType === 'ThrowStatement') {
+    for (const handler of onCallbackSettled) {
+      handler.body = {
+        type: 'BlockStatement',
+        body: [{
+          type: 'ThrowStatement',
+          argument: clone(exitArgument),
+        }],
+      }
+    }
+  } else {
+    onCallbackSettled[0].body = clone(exitArgument)
+    onCallbackSettled[1].body = clone(exitArgument)
+  }
 
-  statements.splice(returnIndex, 0, ...waitStatements)
+  body.body.splice(exitIndex, 0, ...waitStatements)
 }

packages/datadog-instrumentations/src/playwright.js

@@ -1494,7 +1494,8 @@ createRootSuiteCh.subscribe({
 pageGotoCh.subscribe({
   asyncEnd (ctx) {
     // The Page.goto rewriter waits for this so tests closing immediately after navigation still get RUM tags.
-    ctx.asyncEndPromise = handlePageGoto(ctx.self)
+    const rumDetectionPromise = handlePageGoto(ctx.self)
+    ctx.resolveCallback = onDone => rumDetectionPromise.then(onDone, onDone)
   },
 })
 

packages/datadog-instrumentations/test/helpers/rewriter/index.spec.js

@@ -562,19 +562,22 @@ describe('check-require-cache', () => {
     assert.ok(subs.start.called)
   })
 
-  it('should wait for an asyncEnd promise when configured', async () => {
+  it('should wait for the resolve callback only before resolving', async () => {
     const { test } = compileFile('trace-promise-async-end')
     const steps = []
 
     subs = {
       asyncEnd (ctx) {
         steps.push('asyncEnd')
-        ctx.asyncEndPromise = new Promise(resolve => {
+        ctx.resolveCallback = onDone => {
           setImmediate(() => {
-            steps.push('asyncEndPromise')
-            resolve()
+            steps.push('resolveCallback')
+            onDone()
           })
-        })
+        }
+        ctx.rejectCallback = () => {
+          assert.fail('reject callback called for a fulfilled promise')
+        }
       },
     }
 
@@ -593,7 +596,68 @@ describe('check-require-cache', () => {
     const result = await resultPromise
 
     assert.equal(result, 'result')
-    assert.deepStrictEqual(steps, ['asyncEnd', 'asyncEndPromise', 'resolved'])
+    assert.deepStrictEqual(steps, ['asyncEnd', 'resolveCallback', 'resolved'])
+  })
+
+  it('should wait for the reject callback only before preserving a rejection', async () => {
+    const { test } = compileFile('trace-promise-async-end')
+    const error = new Error('test rejection')
+    const steps = []
+
+    subs = {
+      asyncEnd (ctx) {
+        steps.push('asyncEnd')
+        ctx.resolveCallback = () => {
+          assert.fail('resolve callback called for a rejected promise')
+        }
+        ctx.rejectCallback = onDone => {
+          setImmediate(() => {
+            steps.push('rejectCallback')
+            onDone()
+          })
+        }
+      },
+    }
+
+    ch = tracingChannel('orchestrion:test:trace_promise_async_end')
+    ch.subscribe(subs)
+
+    const resultPromise = test(error)
+
+    await Promise.resolve()
+
+    assert.deepStrictEqual(steps, ['asyncEnd'])
+    await assert.rejects(resultPromise, actualError => {
+      steps.push('rejected')
+      return actualError === error
+    })
+    assert.deepStrictEqual(steps, ['asyncEnd', 'rejectCallback', 'rejected'])
+  })
+
+  it('should preserve promise settlement when its callback throws', async () => {
+    const { test } = compileFile('trace-promise-async-end')
+    const error = new Error('test rejection')
+
+    subs = {
+      asyncEnd (ctx) {
+        ctx.resolveCallback = () => {
+          throw new Error('resolve callback error')
+        }
+        ctx.rejectCallback = () => {
+          throw new Error('reject callback error')
+        }
+      },
+    }
+
+    ch = tracingChannel('orchestrion:test:trace_promise_async_end')
+    ch.subscribe(subs)
+
+    const [result] = await Promise.all([
+      test(),
+      assert.rejects(test(error), actualError => actualError === error),
+    ])
+
+    assert.equal(result, 'result')
   })
 
   it('should use import when rewriting esm modules', () => {