fix: ensure correct run-context for instrumented aws-sdk client usage

This fixes aws-sdk (S3, SQS, SNS) instrumentation to ensure the
automatically created AWS spans do not spill into user code. This also
ensures that any captured APM error for a failed AWS command is a child
of the AWS span.

Refs: #2430

CHANGELOG.asciidoc

@@ -39,6 +39,12 @@ Notes:
 [float]
 ===== Bug fixes
 
+* Fixes for run context handling for 'aws-sdk' instrumentation (S3, SQS, SNS)
+  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.
+
 * Fixes for 'ioredis' instrumentation ({pull}2460[#2460]):
 +
 **  Fix run-context so that a span created in the same tick as an ioredis

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

@@ -54,11 +54,6 @@ function instrumentationS3 (orig, origArguments, request, AWS, agent, { version,
 
   const ins = agent._instrumentation
 
-  // `instrumentationS3` is called *synchronously* for an S3 client method call.
-  // We must use `ins.createSpan` rather than `ins.startSpan`, otherwise we
-  // would impact apm.currentSpan in the calling code. For example:
-  //      s3Client.listBuckets({}, function (...) { ... })
-  //      // The "S3 ListBuckets" span should *not* be apm.currentSpan here.
   const span = ins.createSpan(name, TYPE, SUBTYPE, opName)
   if (!span) {
     return orig.apply(request, origArguments)

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

@@ -1,3 +1,5 @@
+'use strict'
+
 const constants = require('../../../constants')
 
 const TYPE = 'messaging'
@@ -105,36 +107,37 @@ function shouldIgnoreRequest (request, agent) {
 }
 
 function instrumentationSns (orig, origArguments, request, AWS, agent, { version, enabled }) {
-  const type = TYPE
-  const subtype = SUBTYPE
-  const action = ACTION
-
   if (shouldIgnoreRequest(request, agent)) {
     return orig.apply(request, origArguments)
   }
 
+  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)
   }
 
   span.setDestinationContext(getMessageDestinationContextFromRequest(request))
   span.setMessageContext(getMessageContextFromRequest(request))
 
-  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 = {

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

@@ -1,6 +1,8 @@
 'use strict'
+
 const { URL } = require('url')
 const constants = require('../../../constants')
+
 const OPERATIONS_TO_ACTIONS = {
   deleteMessage: 'delete',
   deleteMessageBatch: 'delete_batch',
@@ -145,35 +147,39 @@ function instrumentationSqs (orig, origArguments, request, AWS, agent, { version
     return orig.apply(request, origArguments)
   }
 
-  const type = TYPE
-  const subtype = SUBTYPE
+  const ins = agent._instrumentation
   const action = getActionFromRequest(request)
   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)
   }
 
   span.setDestinationContext(getMessageDestinationContextFromRequest(request))
   span.setMessageContext(getMessageContextFromRequest(request))
 
-  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()
 
     if (request.operation === 'receiveMessage' && response && response.data) {
       recordMetrics(getQueueNameFromRequest(request), response.data, agent)
     }
-  })
-
-  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/fixtures/sqs.js

@@ -55,7 +55,7 @@ module.exports = {
         MessageBody: 'Information about current NY Times fiction bestseller for week of 12/11/2016.'
       }]
     },
-    respnse: `<?xml version="1.0"?>
+    response: `<?xml version="1.0"?>
       <SendMessageBatchResponse xmlns="http://queue.amazonaws.com/doc/2012-11-05/">
         <SendMessageBatchResult>
           <SendMessageBatchResultEntry>

test/instrumentation/modules/aws-sdk/fixtures/use-s3.js

@@ -175,19 +175,41 @@ function useS3 (s3Client, bucketName, cb) {
           Bucket: bucketName,
           Key: key
         }).promise()
+        assert(apm.currentSpan === null,
+          'S3 span should NOT be the currentSpan after .promise()')
 
         req.then(
           function onResolve (data) {
             log.info({ data }, 'getObject using Promise, resolve')
+            assert(apm.currentSpan === null,
+              'S3 span should NOT be the currentSpan in promise resolve')
             next()
           },
           function onReject (err) {
             log.info({ err }, 'getObject using Promise, reject')
+            assert(apm.currentSpan === null,
+              'S3 span should NOT be the currentSpan in promise reject')
             next(err)
           }
         )
       },
 
+      // Get a non-existant object to test APM captureError.
+      function getObjNonExistantObject (_, next) {
+        const nonExistantKey = key + '-does-not-exist'
+        s3Client.getObject({
+          Bucket: bucketName,
+          Key: nonExistantKey
+        }, function (err, data) {
+          log.info({ err, data }, 'getObject non-existant key, expect error')
+          if (err) {
+            next()
+          } else {
+            next(new Error(`did not get an error from getObject(${nonExistantKey})`))
+          }
+        })
+      },
+
       // https://docs.aws.amazon.com/AWSJavaScriptSDK/latest/AWS/S3.html#deleteObject-property
       function deleteTheObj (_, next) {
         s3Client.deleteObject({

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

@@ -59,25 +59,25 @@ tape.test('simple S3 usage scenario', function (t) {
         // First the transaction.
         t.ok(events[0].transaction, 'got the transaction')
         const tx = events.shift().transaction
-        t.equal(events.filter(e => e.span).length, events.length,
-          'all remaining events are spans')
+        const errors = events.filter(e => e.error).map(e => e.error)
 
         // Currently HTTP spans under each S3 span are included. Eventually
         // those will be excluded. Filter those out for now.
         // https://github.com/elastic/apm-agent-nodejs/issues/2125
-        const spans = events.map(e => e.span).filter(e => e.subtype !== 'http')
+        const spans = events.filter(e => e.span)
+          .map(e => e.span)
+          .filter(e => e.subtype !== 'http')
 
         // Compare some common fields across all spans.
         t.equal(spans.filter(s => s.trace_id === tx.trace_id).length,
           spans.length, 'all spans have the same trace_id')
         t.equal(spans.filter(s => s.transaction_id === tx.id).length,
           spans.length, 'all spans have the same transaction_id')
-        t.equal(spans.filter(s => s.outcome === 'success').length,
-          spans.length, 'all spans have outcome="success"')
         t.equal(spans.filter(s => s.sync === false).length,
           spans.length, 'all spans have sync=false')
         t.equal(spans.filter(s => s.sample_rate === 1).length,
           spans.length, 'all spans have sample_rate=1')
+        const failingSpanId = spans[8].id // index of `getObjNonExistantObject`
         spans.forEach(s => {
           // Remove variable and common fields to facilitate t.deepEqual below.
           delete s.id
@@ -86,7 +86,6 @@ tape.test('simple S3 usage scenario', function (t) {
           delete s.trace_id
           delete s.timestamp
           delete s.duration
-          delete s.outcome
           delete s.sync
           delete s.sample_rate
         })
@@ -105,7 +104,8 @@ tape.test('simple S3 usage scenario', function (t) {
               service: { name: 's3', type: 'storage' },
               cloud: { region: 'us-east-2' }
             }
-          }
+          },
+          outcome: 'success'
         }, 'listAllBuckets produced expected span')
 
         t.deepEqual(spans.shift(), {
@@ -124,7 +124,8 @@ tape.test('simple S3 usage scenario', function (t) {
               },
               cloud: { region: 'us-east-2' }
             }
-          }
+          },
+          outcome: 'success'
         }, 'createTheBucketIfNecessary produced expected span')
 
         t.deepEqual(spans.shift(), {
@@ -143,7 +144,8 @@ tape.test('simple S3 usage scenario', function (t) {
               },
               cloud: { region: 'us-east-2' }
             }
-          }
+          },
+          outcome: 'success'
         }, 'waitForBucketToExist produced expected span')
 
         t.deepEqual(spans.shift(), {
@@ -162,7 +164,8 @@ tape.test('simple S3 usage scenario', function (t) {
               },
               cloud: { region: 'us-east-2' }
             }
-          }
+          },
+          outcome: 'success'
         }, 'createObj produced expected span')
 
         t.deepEqual(spans.shift(), {
@@ -181,7 +184,8 @@ tape.test('simple S3 usage scenario', function (t) {
               },
               cloud: { region: 'us-east-2' }
             }
-          }
+          },
+          outcome: 'success'
         }, 'waitForObjectToExist produced expected span')
 
         t.deepEqual(spans.shift(), {
@@ -200,7 +204,8 @@ tape.test('simple S3 usage scenario', function (t) {
               },
               cloud: { region: 'us-east-2' }
             }
-          }
+          },
+          outcome: 'success'
         }, 'getObj produced expected span')
 
         t.deepEqual(spans.shift(), {
@@ -219,7 +224,8 @@ tape.test('simple S3 usage scenario', function (t) {
               },
               cloud: { region: 'us-east-2' }
             }
-          }
+          },
+          outcome: 'success'
         }, 'getObjConditionalGet produced expected span')
 
         t.deepEqual(spans.shift(), {
@@ -238,9 +244,35 @@ tape.test('simple S3 usage scenario', function (t) {
               },
               cloud: { region: 'us-east-2' }
             }
-          }
+          },
+          outcome: 'success'
         }, 'getObjUsingPromise produced expected span')
 
+        // This is the GetObject to a non-existant-key, so we expect a failure.
+        t.deepEqual(spans.shift(), {
+          name: 'S3 GetObject elasticapmtest-bucket-1',
+          type: 'storage',
+          subtype: 's3',
+          action: 'GetObject',
+          context: {
+            destination: {
+              address: LOCALSTACK_HOST,
+              port: 4566,
+              service: {
+                name: 's3',
+                type: 'storage',
+                resource: 'elasticapmtest-bucket-1'
+              },
+              cloud: { region: 'us-east-2' }
+            }
+          },
+          outcome: 'failure'
+        }, 'getObjNonExistantObject produced expected span')
+        t.equal(errors.length, 1, 'got 1 error')
+        t.equal(errors[0].parent_id, failingSpanId, 'error is a child of the failing span from getObjNonExistantObject')
+        t.equal(errors[0].transaction_id, tx.id, 'error.transaction_id')
+        t.equal(errors[0].exception.type, 'NoSuchKey', 'error.exception.type')
+
         t.deepEqual(spans.shift(), {
           name: 'S3 DeleteObject elasticapmtest-bucket-1',
           type: 'storage',
@@ -257,7 +289,8 @@ tape.test('simple S3 usage scenario', function (t) {
               },
               cloud: { region: 'us-east-2' }
             }
-          }
+          },
+          outcome: 'success'
         }, 'deleteTheObj produced expected span')
 
         t.deepEqual(spans.shift(), {
@@ -276,7 +309,8 @@ tape.test('simple S3 usage scenario', function (t) {
               },
               cloud: { region: 'us-east-2' }
             }
-          }
+          },
+          outcome: 'success'
         }, 'deleteTheBucketIfCreatedIt produced expected span')
 
         t.equal(spans.length, 0, 'all spans accounted for')