Don't restore the root span for withContext

The ScopeManager.withContext only temporary sets the current active
span. It then restores the active span afterwards.

But it also restores the root span, even though that should not have
changed. At least not by anything we've done in this withContext
function.

If the user has set another rootSpan, that is their decision, but it's
not up to us to restore it.

packages/nodejs/.changesets/do-not-restore-root-span-after-withcontext-callback-has-finished.md

@@ -0,0 +1,16 @@
+---
+bump: "patch"
+type: "change"
+---
+
+Do not restore root span after `withSpan` callback has finished. Previously the root span was restored to the original root span before the `withSpan` helper was called. This has been changed, because the `withSpan` helper is only about changing the active span, not the root span. If a new root span has been set within a `withSpan` helper callback, the root span will no longer be restored. We recommend setting a new root span before calling `withSpan` instead.
+
+```js
+const rootSpan = tracer.rootSpan()
+const span = tracer.createSpan(...)
+tracer.withSpan(span, function(span) {
+  tracer.createRootSpan(...)
+});
+// No longer match
+rootSpan != tracer.rootSpan()
+```

packages/nodejs/src/__tests__/scope.test.ts

@@ -211,17 +211,19 @@ describe("ScopeManager", () => {
       })
     })
 
-    it("restores the previous active span and root span", () => {
+    it("restores the previous active span", () => {
       const outerRootSpan = new RootSpan()
       const outerChildSpan = new ChildSpan(outerRootSpan)
       const innerChildSpan = new ChildSpan(outerChildSpan)
-      const innerRootSpan = new RootSpan()
 
       scopeManager.setRoot(outerRootSpan)
+      expect(scopeManager.active()).toBe(outerRootSpan)
+      expect(scopeManager.root()).toBe(outerRootSpan)
 
       scopeManager.withContext(outerChildSpan, () => {
         scopeManager.withContext(innerChildSpan, () => {
-          scopeManager.setRoot(innerRootSpan)
+          expect(scopeManager.root()).toBe(outerRootSpan)
+          expect(scopeManager.active()).toBe(innerChildSpan)
         })
 
         expect(scopeManager.active()).toBe(outerChildSpan)
@@ -231,6 +233,26 @@ describe("ScopeManager", () => {
       expect(scopeManager.active()).toBe(outerRootSpan)
       expect(scopeManager.root()).toBe(outerRootSpan)
     })
+
+    it("does not restore the root span if it was changed within withContext", () => {
+      const outerRootSpan = new RootSpan()
+      const outerChildSpan = new ChildSpan(outerRootSpan)
+      const innerRootSpan = new RootSpan()
+
+      scopeManager.setRoot(outerRootSpan)
+      expect(scopeManager.active()).toBe(outerRootSpan)
+      expect(scopeManager.root()).toBe(outerRootSpan)
+
+      scopeManager.withContext(outerChildSpan, () => {
+        expect(scopeManager.active()).toBe(outerChildSpan)
+        expect(scopeManager.root()).toBe(outerRootSpan)
+        scopeManager.setRoot(innerRootSpan)
+        expect(scopeManager.root()).toBe(innerRootSpan)
+      })
+
+      expect(scopeManager.active()).toBe(outerRootSpan)
+      expect(scopeManager.root()).toBe(innerRootSpan)
+    })
   })
 
   describe(".bindContext()", () => {

packages/nodejs/src/scope.ts

@@ -181,7 +181,6 @@ export class ScopeManager {
   public withContext<T>(span: Span, fn: (s: Span) => T): T {
     const uid = asyncHooks.executionAsyncId()
     const oldScope = this.active()
-    const rootSpan = this.root()
 
     if (span.open) {
       this.#scopes.set(uid, span)
@@ -192,15 +191,14 @@ export class ScopeManager {
     try {
       return fn(span)
     } catch (err) {
-      rootSpan?.setError(err)
+      this.root()?.setError(err)
       throw err
     } finally {
       // revert to the previous span
       if (oldScope === undefined) {
-        this.removeSpanForUid(uid)
+        this.#scopes.delete(uid)
       } else {
         this.#scopes.set(uid, oldScope)
-        this.#roots.set(uid, rootSpan)
       }
     }
   }