handling subtleties with how span.end() impacts the run context

For example, this handles the problem case from
#1964

lib/instrumentation/index.js

@@ -314,11 +314,13 @@ Instrumentation.prototype.addEndedSpan = function (span) {
     return
   }
 
-  const rc = this._runCtxMgr.active()
-  if (rc.topSpan() === span) {
-    // Replace the active run context with this span popped off the stack,
-    // i.e. this span is no longer active.
-    this._runCtxMgr.replaceActive(rc.exitSpan())
+  // Replace the active run context with this span removed. Typically this
+  // span is the top of stack (i.e. is the current span). However, it is
+  // possible to have out-of-order span.end(), in which case the ended span
+  // might not.
+  const newRc = this._runCtxMgr.active().exitSpan(span)
+  if (newRc) {
+    this._runCtxMgr.replaceActive(newRc)
   }
   this._log.debug({ ctxmgr: this._runCtxMgr.toString() }, 'addEndedSpan(%s)', span.name)
 

lib/run-context/BasicRunContextManager.js

@@ -55,33 +55,8 @@ class RunContext {
     return !this.tx
   }
 
-  // Returns the currently active span, if any.
-  //
-  // Because the `startSpan()/endSpan()` API allows (a) affecting the current
-  // run context and (b) out of order start/end, the "currently active span"
-  // must skip over ended spans.
+  // Returns the currently active span, if any, otherwise null.
   currSpan () {
-    // XXX revisit this with .ended change in opinion
-    // XXX enabling this makes ls-promises.js fail. This may get hairy.
-    // if (true) {
-    //   if (this.spans.length > 0) {
-    //     return this.spans[this.spans.length - 1]
-    //   } else {
-    //     return null
-    //   }
-    // }
-
-    for (let i = this.spans.length - 1; i >= 0; i--) {
-      const span = this.spans[i]
-      if (!span.ended) {
-        return span
-      }
-    }
-    return null
-  }
-
-  // This returns the top span in the span stack (even if it is ended).
-  topSpan () {
     if (this.spans.length > 0) {
       return this.spans[this.spans.length - 1]
     } else {
@@ -93,38 +68,33 @@ class RunContext {
   // stack.
   enterSpan (span) {
     const newSpans = this.spans.slice()
-
-    // Any ended spans at the top of the stack are cruft -- remove them.
-    while (newSpans.length > 0 && newSpans[newSpans.length - 1].ended) {
-      newSpans.pop()
-    }
-
     newSpans.push(span)
     return new RunContext(this.tx, newSpans)
   }
 
-  exitSpan () {
-    const newSpans = this.spans.slice(0, this.spans.length - 1)
-
-    // XXX revisit this with .ended change in opinion; TODO: sync and async test cases for this!
-    // Pop all ended spans.  It is possible that spans lower in the stack have
-    // already been ended. For example, in this code:
-    //    var t1 = apm.startSpan('t1')
-    //    var s2 = apm.startSpan('s2')
-    //    var s3 = apm.startSpan('s3')
-    //    var s4 = apm.startSpan('s4')
-    //    s3.end() // out of order
-    //    s4.end()
-    // when `s4.end()` is called, the current run context will be:
-    //    RC(tx=t1, spans=[s2, s3.ended, s4])
-    // The final result should be:
-    //    RC(tx=t1, spans=[s2])
-    // so that `s2` becomes the current/active span.
-    while (newSpans.length > 0 && newSpans[newSpans.length - 1].ended) {
-      newSpans.pop()
+  // Return a new RunContext with the given span removed, or null if there is
+  // no change (the given span isn't part of the run context).
+  //
+  // Typically this span is the top of stack (i.e. it is the current span).
+  // However, it is possible to have out-of-order span.end() or even end a span
+  // that isn't part of the current run context stack at all.
+  // (See test/run-context/fixtures/end-non-current-spans.js for examples.)
+  exitSpan (span) {
+    let newRc = null
+    let newSpans
+    const lastSpan = this.spans[this.spans.length - 1]
+    if (lastSpan && lastSpan.id === span.id) {
+      // Fast path for common case: `span` is top of stack.
+      newSpans = this.spans.slice(0, this.spans.length - 1)
+      newRc = new RunContext(this.tx, newSpans)
+    } else {
+      const stackIdx = this.spans.findIndex(s => s.id === span.id)
+      if (stackIdx !== -1) {
+        newSpans = this.spans.slice(0, stackIdx).concat(this.spans.slice(stackIdx + 1))
+        newRc = new RunContext(this.tx, newSpans)
+      }
     }
-
-    return new RunContext(this.tx, newSpans)
+    return newRc
   }
 
   toString () {
@@ -312,12 +282,10 @@ class BasicRunContextManager {
   // XXX s/_enterContext/_enterRunContext/ et al
   _enterContext (runContext) {
     this._stack.push(runContext)
-    this._log.trace({ ctxmgr: this.toString() }, '_enterContext %s', runContext)
   }
 
   _exitContext () {
-    var popped = this._stack.pop()
-    this._log.trace({ ctxmgr: this.toString() }, '_exitContext %s', popped)
+    this._stack.pop()
   }
 
   // ---- Additional public API added to support startTransaction/startSpan API.
@@ -330,7 +298,7 @@ class BasicRunContextManager {
   //    names: `stompRunContext`, `hardEnterRunContext`, `squatRunContext`
   //      `occupyRunContext`, `replacingEnterRunContext`
   //    I like `stompRunContext`.
-  // XXX This impl could just be `_exitContext(); _enterContext(rc)` right?
+  // XXX This impl could just be `_exitContext(); _enterContext(rc)` right? Do that, if so.
   replaceActive (runContext) {
     if (this._stack.length > 0) {
       this._stack[this._stack.length - 1] = runContext
@@ -339,7 +307,6 @@ class BasicRunContextManager {
       //     context for startTransaction/startSpan only if there isn't one
       this._stack.push(runContext)
     }
-    this._log.trace({ ctxmgr: this.toString() }, 'replaceActive %s', runContext)
   }
 }
 

test/run-context/fixtures/end-non-current-spans.js

@@ -0,0 +1,41 @@
+// This test case shows that `span.end()` impacts the current run context's
+// span stack, even if the ended span is not the current one. When s3 and s2
+// are ended below, they are not the current span.
+//
+// Expected:
+//   transaction "t0"
+//   `- span "s1"
+//     `- span "s2"
+//     `- span "s3"
+//       `- span "s4"
+
+const apm = require('../../../').start({ // elastic-apm-node
+  captureExceptions: false,
+  captureSpanStackTraces: false,
+  metricsInterval: 0,
+  cloudProvider: 'none',
+  centralConfig: false,
+  // ^^ Boilerplate config above this line is to focus on just tracing.
+  serviceName: 'run-context-end-non-current-spans'
+})
+
+const assert = require('assert').strict
+
+const t0 = apm.startTransaction('t0')
+const s1 = apm.startSpan('s1')
+setImmediate(function () {
+  const s3 = apm.startSpan('s3')
+  setImmediate(function () {
+    const s4 = apm.startSpan('s4')
+    // Ending a span removes it from the current run context, even if it is
+    // not top of stack, or not even part of this run context.
+    s3.end() // out of order
+    s2.end() // not in this run context
+    s4.end()
+    assert(apm.currentSpan === s1)
+    s1.end()
+    t0.end()
+  })
+})
+
+const s2 = apm.startSpan('s2')

test/run-context/fixtures/ls-promises.js

@@ -20,12 +20,13 @@ let t1
 
 function getCwd () {
   var s2 = apm.startSpan('cwd')
-  return Promise.resolve(process.cwd())
-    .finally(() => {
-      assert(apm.currentTransaction === t1)
-      assert(apm.currentSpan === s2)
-      s2.end()
-    })
+  try {
+    return Promise.resolve(process.cwd())
+  } finally {
+    assert(apm.currentTransaction === t1)
+    assert(apm.currentSpan === s2)
+    s2.end()
+  }
 }
 
 function main () {

test/run-context/fixtures/parentage-with-ended-span.js

@@ -6,8 +6,6 @@
 //    - span "s1"
 //      - span "s3"
 //    - span "s2"
-//
-// XXX TODO Also have a more complex case that ends a span in the stack that isn't the top of stack.
 
 const apm = require('../../../').start({ // elastic-apm-node
   captureExceptions: false,
@@ -16,7 +14,7 @@ const apm = require('../../../').start({ // elastic-apm-node
   cloudProvider: 'none',
   centralConfig: false,
   // ^^ Boilerplate config above this line is to focus on just tracing.
-  serviceName: 'run-context-child-of-ended-span'
+  serviceName: 'run-context-parentage-with-ended-span'
 })
 
 const assert = require('assert').strict

test/run-context/run-context.test.js

@@ -59,7 +59,7 @@ const cases = [
       const s2 = findObjInArray(events, 'span.name', 'cwd')
       const s3 = findObjInArray(events, 'span.name', 'readdir')
       t.equal(s2.parent_id, t1.id, 's2 is a child of t1')
-      t.equal(s3.parent_id, t1.id, 's3 is a child of t1')
+      t.equal(s3 && s3.parent_id, t1.id, 's3 is a child of t1')
       // XXX check sync for the spans
     }
   },
@@ -98,7 +98,29 @@ const cases = [
       t.equal(s3.parent_id, s1.id, 's3 is a child of s1')
       // XXX could check that s3 start time is after s1 end time
     }
-  }
+  },
+  {
+    script: 'end-non-current-spans.js',
+    check: (t, events) => {
+      // Expected:
+      //   transaction "t0"
+      //   `- span "s1"
+      //     `- span "s2"
+      //     `- span "s3"
+      //       `- span "s4"
+      t.ok(events[0].metadata, 'APM server got event metadata object')
+      t.equal(events.length, 6, 'exactly 6 events')
+      const t0 = findObjInArray(events, 'transaction.name', 't0')
+      const s1 = findObjInArray(events, 'span.name', 's1')
+      const s2 = findObjInArray(events, 'span.name', 's2')
+      const s3 = findObjInArray(events, 'span.name', 's3')
+      const s4 = findObjInArray(events, 'span.name', 's4')
+      t.equal(s1.parent_id, t0.id, 's1 is a child of t0')
+      t.equal(s2.parent_id, s1.id, 's2 is a child of s1')
+      t.equal(s3.parent_id, s1.id, 's3 is a child of s1')
+      t.equal(s4.parent_id, s3.id, 's4 is a child of s3')
+    }
+  },
 ]
 
 cases.forEach(c => {