changes to S3 instrumentation so that its parent/child relationship works (see examples/parent-child.js case)

We will have to do the same for other aws-sdk, for ES, and evaluate for
other instrumentations. Note this is an improvement over current
behaviour which doesn't handle cases like the parent/child reln above
anyway.

Author: trentm

examples/parent-child.js

@@ -1,6 +1,6 @@
 // vim: set ts=2 sw=2:
 
-var apm = require('./').start({ // elastic-apm-node
+var apm = require('../').start({ // elastic-apm-node
   serviceName: 'parent-child',
   captureExceptions: false,
   logUncaughtExceptions: true,
@@ -14,6 +14,7 @@ var apm = require('./').start({ // elastic-apm-node
   // disableSend: true
 })
 
+const assert = require('assert').strict
 const express = require('express')
 
 const app = express()
@@ -39,10 +40,6 @@ app.get('/a', (req, res) => {
   }, 10)
 })
 
-setTimeout(function () {
-  console.warn('XXX in unrelated 3s timeout: currTx is: ', apm._instrumentation.currTx())
-}, 3000)
-
 app.get('/b', (req, res) => {
   var s1 = apm.startSpan('s1')
   s1.end()
@@ -130,28 +127,79 @@ app.get('/unended-span', (req, res) => {
 //       HTTP span
 //         a-sibling-span
 //
-// Perhaps this is fine?
+// Perhaps this is fine? Nope, it isn't.
+//
 app.get('/dario-1963', (req, res) => {
+  const { Client } = require('@elastic/elasticsearch')
+  const client = new Client({
+    node: 'http://localhost:9200',
+    auth: { username: 'admin', password: 'changeme' }
+  })
+  // Note: works fine if client.search is under a setImmediate for sep context.
+  // setImmediate(function () { ... })
   client.search({
-    index: 'kibana_sample_data_logs',
+    // index: 'kibana_sample_data_logs',
     body: { size: 1, query: { match_all: {} } }
   }, (err, _result) => {
+    console.warn('XXX in client.search cb: %s', apm._instrumentation._runCtxMgr.active())
     if (err) {
       res.send(err)
     } else {
       res.send('ok')
     }
   })
 
+  console.warn('XXX after client.search sync: %s', apm._instrumentation._runCtxMgr.active())
+
   // What if I add this?
-  setImmediate(function () {
+  setImmediate(function aSiblingSpanInHere () {
     var span = apm.startSpan('a-sibling-span')
     setImmediate(function () {
       span.end()
     })
   })
 })
 
+// Want:
+//   transaction "GET /s3"
+//   `- span "span1"
+//     `- span "S3 ListBuckets"
+//       `- span "GET s3.amazonaws.com/"
+//     `- span "span3"
+//     `- span "span2"
+//
+// Eventually the HTTP span should be removed as exit spans are supported.
+app.get('/s3', (req, res) => {
+  const AWS = require('aws-sdk')
+  const s3Client = new AWS.S3({ apiVersion: '2006-03-01' })
+
+  setImmediate(function () {
+    var s1 = apm.startSpan('span1')
+
+    s3Client.listBuckets({}, function (err, _data) {
+      if (err) {
+        res.send(err)
+      } else {
+        res.send('ok')
+      }
+      s1.end()
+    })
+    assert(apm._instrumentation.currSpan() === s1)
+
+    setImmediate(function () {
+      assert(apm._instrumentation.currSpan() === s1)
+      var s2 = apm.startSpan('span2')
+      setImmediate(function () {
+        s2.end()
+      })
+    })
+
+    assert(apm._instrumentation.currSpan() === s1)
+    var s3 = apm.startSpan('span3')
+    s3.end()
+  })
+})
+
 app.listen(port, function () {
   console.log(`listening at http://localhost:${port}`)
 })

lib/instrumentation/generic-span.js

@@ -25,6 +25,7 @@ function GenericSpan (agent, ...args) {
   // console.warn('XXX new GenericSpan traceContext: ', this._traceContext.toTraceParentString(), this._traceContext.toTraceStateString())
 
   // XXX change this var name to _traceContext.
+  console.warn('XXX new GenericSpan: opts.childOf=', opts.childOf && (opts.childOf.constructor.name + ' ' + opts.childOf.name))
   this._context = TraceContext.startOrResume(opts.childOf, agent._conf, opts.tracestate)
 
   this._agent = agent

lib/instrumentation/http-shared.js

@@ -134,7 +134,6 @@ exports.traceOutgoingRequest = function (agent, moduleName, method) {
       //     shimmer to expose access to the original via a method)? Invoke
       //     in a separate no-op RunContext???
       console.warn('XXX traceOutgoingRequest: start')
-      // console.warn('XXX curr span:', agent._instrumentation.currentSpan?.name) // considers 'bindingSpan'
 
       // TODO: See if we can delay the creation of span until the `response`
       // event is fired, while still having it have the correct stack trace
@@ -178,6 +177,7 @@ exports.traceOutgoingRequest = function (agent, moduleName, method) {
       var req = orig.apply(this, newArgs)
       if (!span) return req
 
+      // XXX
       if (getSafeHost(req) === agent._conf.serverHost) {
         agent.logger.debug('ignore %s request to intake API %o', moduleName, { id: id })
         return req

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

@@ -52,76 +52,83 @@ function instrumentationS3 (orig, origArguments, request, AWS, agent, { version,
     name += ' ' + resource
   }
 
-  const span = agent.startSpan(name, TYPE, SUBTYPE, opName)
-  if (span) {
-    request.on('complete', function onComplete (response) {
-      // `response` is an AWS.Response
-      // https://docs.aws.amazon.com/AWSJavaScriptSDK/latest/AWS/Response.html
-
-      // Determining the bucket's region.
-      // `request.httpRequest.region` isn't documented, but the aws-sdk@2
-      // lib/services/s3.js will set it to the bucket's determined region.
-      // This can be asynchronously determined -- e.g. if it differs from the
-      // configured service endpoint region -- so this won't be set until
-      // 'complete'.
-      const region = request.httpRequest && request.httpRequest.region
-
-      // Destination context.
-      // '.httpRequest.endpoint' might differ from '.service.endpoint' if
-      // the bucket is in a different region.
-      const endpoint = request.httpRequest && request.httpRequest.endpoint
-      const destContext = {
-        service: {
-          name: SUBTYPE,
-          type: TYPE
-        }
-      }
-      if (endpoint) {
-        destContext.address = endpoint.hostname
-        destContext.port = endpoint.port
-      }
-      if (resource) {
-        destContext.service.resource = resource
+  const ins = agent._instrumentation
+  const span = ins.currTx().startSpan(name, TYPE, SUBTYPE, opName)
+  if (!span) {
+    return orig.apply(request, origArguments)
+  }
+
+  function onComplete (response) {
+    // `response` is an AWS.Response
+    // https://docs.aws.amazon.com/AWSJavaScriptSDK/latest/AWS/Response.html
+
+    // Determining the bucket's region.
+    // `request.httpRequest.region` isn't documented, but the aws-sdk@2
+    // lib/services/s3.js will set it to the bucket's determined region.
+    // This can be asynchronously determined -- e.g. if it differs from the
+    // configured service endpoint region -- so this won't be set until
+    // 'complete'.
+    const region = request.httpRequest && request.httpRequest.region
+
+    // Destination context.
+    // '.httpRequest.endpoint' might differ from '.service.endpoint' if
+    // the bucket is in a different region.
+    const endpoint = request.httpRequest && request.httpRequest.endpoint
+    const destContext = {
+      service: {
+        name: SUBTYPE,
+        type: TYPE
       }
-      if (region) {
-        destContext.cloud = { region }
+    }
+    if (endpoint) {
+      destContext.address = endpoint.hostname
+      destContext.port = endpoint.port
+    }
+    if (resource) {
+      destContext.service.resource = resource
+    }
+    if (region) {
+      destContext.cloud = { region }
+    }
+    span.setDestinationContext(destContext)
+
+    if (response) {
+      // Follow the spec for HTTP client span outcome.
+      // https://github.com/elastic/apm/blob/master/specs/agents/tracing-instrumentation-http.md#outcome
+      //
+      // For example, a S3 GetObject conditional request (e.g. using the
+      // IfNoneMatch param) will respond with response.error=NotModifed and
+      // statusCode=304. This is a *successful* outcome.
+      const statusCode = response.httpResponse && response.httpResponse.statusCode
+      if (statusCode) {
+        span._setOutcomeFromHttpStatusCode(statusCode)
+      } else {
+        // `statusCode` will be undefined for errors before sending a request, e.g.:
+        //  InvalidConfiguration: Custom endpoint is not compatible with access point ARN
+        span._setOutcomeFromErrorCapture(constants.OUTCOME_FAILURE)
       }
-      span.setDestinationContext(destContext)
-
-      if (response) {
-        // Follow the spec for HTTP client span outcome.
-        // https://github.com/elastic/apm/blob/master/specs/agents/tracing-instrumentation-http.md#outcome
-        //
-        // For example, a S3 GetObject conditional request (e.g. using the
-        // IfNoneMatch param) will respond with response.error=NotModifed and
-        // statusCode=304. This is a *successful* outcome.
-        const statusCode = response.httpResponse && response.httpResponse.statusCode
-        if (statusCode) {
-          span._setOutcomeFromHttpStatusCode(statusCode)
-        } else {
-          // `statusCode` will be undefined for errors before sending a request, e.g.:
-          //  InvalidConfiguration: Custom endpoint is not compatible with access point ARN
-          span._setOutcomeFromErrorCapture(constants.OUTCOME_FAILURE)
-        }
-
-        if (response.error && (!statusCode || statusCode >= 400)) {
-          agent.captureError(response.error, { skipOutcome: true })
-        }
+
+      if (response.error && (!statusCode || statusCode >= 400)) {
+        agent.captureError(response.error, { skipOutcome: true })
       }
+    }
 
-      // Workaround a bug in the agent's handling of `span.sync`.
-      //
-      // The bug: Currently this span.sync is not set `false` because there is
-      // an HTTP span created (for this S3 request) in the same async op. That
-      // HTTP span becomes the "active span" for this async op, and *it* gets
-      // marked as sync=false in `before()` in async-hooks.js.
-      span.sync = false
-
-      span.end()
-    })
+    // Workaround a bug in the agent's handling of `span.sync`.
+    //
+    // The bug: Currently this span.sync is not set `false` because there is
+    // an HTTP span created (for this S3 request) in the same async op. That
+    // HTTP span becomes the "active span" for this async op, and *it* gets
+    // marked as sync=false in `before()` in async-hooks.js.
+    span.sync = false
+
+    span.end()
   }
 
-  return orig.apply(request, origArguments)
+  // Derive a new run context from the current one for this span. Then run
+  // the AWS.Request.send and a 'complete' event handler in that run context.
+  const runContext = ins._runCtxMgr.active().enterSpan(span) // XXX I don't like `enterSpan` name here, perhaps `newWithSpan()`?
+  request.on('complete', ins._runCtxMgr.bind(runContext, onComplete))
+  return ins._runCtxMgr.with(runContext, orig, request, ...origArguments)
 }
 
 module.exports = {

lib/instrumentation/span.js

@@ -17,6 +17,7 @@ util.inherits(Span, GenericSpan)
 
 function Span (transaction, name, ...args) {
   const parent = transaction._agent._instrumentation.currSpan() || transaction
+  console.warn('XXX new Span(name=%s, args=%s): parent=', name, args, parent.constructor.name, parent.name, parent.ended ? '.ended' : '')
   const opts = typeof args[args.length - 1] === 'object'
     ? (args.pop() || {})
     : {}