fix: run context handling for aws-sdk DynamoDB instrumentation (#2504)

Refs: #2430

.tav.yml

@@ -479,6 +479,8 @@ aws-sdk:
   # Test v2.858.0, every N=31 of 159 releases, and current latest.
   versions: '2.858.0 || 2.889.0 || 2.920.0 || 2.951.0 || 2.982.0 || 2.1013.0 || 2.1016.0 || >2.1016.0 <3'
   commands:
+    - node test/instrumentation/modules/aws-sdk/aws4-retries.test.js
     - node test/instrumentation/modules/aws-sdk/s3.test.js
+    - node test/instrumentation/modules/aws-sdk/sns.test.js
     - node test/instrumentation/modules/aws-sdk/sqs.test.js
     - node test/instrumentation/modules/aws-sdk/dynamodb.test.js

CHANGELOG.asciidoc

@@ -28,6 +28,23 @@ Notes:
 [[release-notes-3.x]]
 === Node.js Agent version 3.x
 
+==== Unreleased
+
+[float]
+===== Breaking changes
+
+[float]
+===== Features
+
+[float]
+===== Bug fixes
+
+* Fixes for run context handling for DynamoDB instrumentation ('aws-sdk'
+  package) so that a span created after an AWS client command (in the same
+  tick, in the command callback, or promise) is not a child of the automatic
+  AWS span. This change also ensures captured errors from failing client
+  commands are a child of the AWS span. ({issues}2430[#2430])
+
 
 [[release-notes-3.26.0]]
 ==== 3.26.0 2021/12/07

lib/instrumentation/modules/aws-sdk/dynamodb.js

@@ -1,5 +1,5 @@
 'use strict'
-const constants = require('../../../constants')
+
 const TYPE = 'db'
 const SUBTYPE = 'dynamodb'
 const ACTION = 'query'
@@ -63,12 +63,9 @@ function instrumentationDynamoDb (orig, origArguments, request, AWS, agent, { ve
     return orig.apply(request, origArguments)
   }
 
-  const type = TYPE
-  const subtype = SUBTYPE
-  const action = ACTION
-
+  const ins = agent._instrumentation
   const name = getSpanNameFromRequest(request)
-  const span = agent.startSpan(name, type, subtype, action)
+  const span = ins.createSpan(name, TYPE, SUBTYPE, ACTION)
   if (!span) {
     return orig.apply(request, origArguments)
   }
@@ -91,19 +88,23 @@ function instrumentationDynamoDb (orig, origArguments, request, AWS, agent, { ve
     }
   })
 
-  request.on('complete', function (response) {
+  const onComplete = function (response) {
     if (response && response.error) {
-      const errOpts = {
-        skipOutcome: true
-      }
-      agent.captureError(response.error, errOpts)
-      span._setOutcomeFromErrorCapture(constants.OUTCOME_FAILURE)
+      agent.captureError(response.error)
     }
-
     span.end()
-  })
-
-  return orig.apply(request, origArguments)
+  }
+  // Bind onComplete to the span's run context so that `captureError` picks
+  // up the correct currentSpan.
+  const parentRunContext = ins.currRunContext()
+  const spanRunContext = parentRunContext.enterSpan(span)
+  request.on('complete', ins.bindFunctionToRunContext(spanRunContext, onComplete))
+
+  const cb = origArguments[origArguments.length - 1]
+  if (typeof cb === 'function') {
+    origArguments[origArguments.length - 1] = ins.bindFunctionToRunContext(parentRunContext, cb)
+  }
+  return ins.withRunContext(spanRunContext, orig, request, ...origArguments)
 }
 
 module.exports = {

test/instrumentation/modules/aws-sdk/dynamodb.test.js

@@ -1,3 +1,5 @@
+'use strict'
+
 const agent = require('../../../..').start({
   serviceName: 'test',
   secretToken: 'test',
@@ -167,10 +169,12 @@ tape.test('AWS DynamoDB: End to End Test', function (test) {
         }
       }
       ddb.query(params, function (err, data) {
+        t.ok(agent.currentSpan === null, 'no currentSpan in ddb.query callback')
         t.error(err)
         agent.endTransaction()
         listener.close()
       })
+      t.ok(agent.currentSpan === null, 'no currentSpan in sync code after ddb.query(...)')
     })
   })
 
@@ -196,10 +200,12 @@ tape.test('AWS DynamoDB: End to End Test', function (test) {
       agent.startTransaction('myTransaction')
       var ddb = new AWS.DynamoDB({ apiVersion: '2012-08-10' })
       ddb.listTables(function (err, data) {
+        t.ok(agent.currentSpan === null, 'no currentSpan in ddb.listTables callback')
         t.error(err)
         agent.endTransaction()
         listener.close()
       })
+      t.ok(agent.currentSpan === null, 'no currentSpan in sync code after ddb.listTables(...)')
     })
   })
 
@@ -209,10 +215,12 @@ tape.test('AWS DynamoDB: End to End Test', function (test) {
     )
     const listener = app.listen(0, function () {
       resetAgent(function (data) {
-        t.equals(data.errors.length, 1, 'expect captured error')
         const span = data.spans.filter((span) => span.type === 'db').pop()
         t.ok(span, 'expect a db span')
         t.equals(span.outcome, 'failure', 'expect db span to have failure outcome')
+        t.equals(data.errors.length, 1, 'expect captured error')
+        const error = data.errors[0]
+        t.equals(error.parent_id, span.id, 'error is a child of the failing span')
         t.end()
       })
       const port = listener.address().port
@@ -229,10 +237,12 @@ tape.test('AWS DynamoDB: End to End Test', function (test) {
         }
       }
       ddb.query(params, function (err, data) {
+        t.ok(agent.currentSpan === null, 'no currentSpan in ddb.query callback')
         t.ok(err, 'expect error')
         agent.endTransaction()
         listener.close()
       })
+      t.ok(agent.currentSpan === null, 'no currentSpan in sync code after ddb.query(...)')
     })
   })