fix(spanner): avoid overriding host context manager and isolate session creation context (#8441)

Author: alkatrivedi

handwritten/spanner/src/instrument.ts

@@ -97,23 +97,26 @@ const {
   AsyncHooksContextManager,
 } = require('@opentelemetry/context-async-hooks');
 
+let contextManagerInstallAttempted = false;
+
 /*
- * This function ensures that async/await works correctly by
- * checking if context.active() returns an invalid/unset context
- * and if so, sets a global AsyncHooksContextManager otherwise
- * spans resulting from async/await invocations won't be correctly
- * associated in their respective hierarchies.
+ * If no global ContextManager is registered, install an AsyncHooksContextManager
+ * so that async/await trace context propagation works for apps that haven't
+ * configured OpenTelemetry themselves. If the host app has already installed a
+ * ContextManager, leave it alone — tearing down a working manager breaks the
+ * host's baggage and span parentage on the next gRPC call.
+ *
+ * setGlobalContextManager() returns false when a manager is already registered,
+ * which is the documented signal that we shouldn't replace it. The
+ * `contextManagerInstallAttempted` latch makes the call idempotent so we don't
+ * allocate a new AsyncHooksContextManager on each Spanner client construction.
  */
 function ensureInitialContextManagerSet() {
-  if (!context['_contextManager'] || context.active() === ROOT_CONTEXT) {
-    // If no context manager is currently set, or if the active context is the ROOT_CONTEXT,
-    // trace context propagation cannot
-    // function correctly with async/await for OpenTelemetry
-    // See {@link https://opentelemetry.io/docs/languages/js/context/#active-context}
-    context.disable(); // Disable any prior contextManager.
-    const contextManager = new AsyncHooksContextManager();
+  if (contextManagerInstallAttempted) return;
+  contextManagerInstallAttempted = true;
+  const contextManager = new AsyncHooksContextManager();
+  if (context.setGlobalContextManager(contextManager)) {
     contextManager.enable();
-    context.setGlobalContextManager(contextManager);
   }
 }
 

handwritten/spanner/src/multiplexed-session.ts

@@ -124,25 +124,29 @@ export class MultiplexedSession
         opts: this._observabilityOptions,
         dbName: this.database.formattedName_,
       };
-      return startTrace(
-        'MultiplexedSession.createSession',
-        traceConfig,
-        async span => {
-          span.addEvent('Requesting a multiplexed session');
-          try {
-            const [createSessionResponse] = await this.database.createSession({
-              multiplexed: true,
-            });
-            this._multiplexedSession = createSessionResponse;
-            span.addEvent('Created a multiplexed session');
-          } catch (e) {
-            setSpanError(span, e as Error);
-            throw e;
-          } finally {
-            span.end();
-          }
-        },
-      );
+      return context.with(ROOT_CONTEXT, () => {
+        return startTrace(
+          'MultiplexedSession.createSession',
+          traceConfig,
+          async span => {
+            span.addEvent('Requesting a multiplexed session');
+            try {
+              const [createSessionResponse] = await this.database.createSession(
+                {
+                  multiplexed: true,
+                },
+              );
+              this._multiplexedSession = createSessionResponse;
+              span.addEvent('Created a multiplexed session');
+            } catch (e) {
+              setSpanError(span, e as Error);
+              throw e;
+            } finally {
+              span.end();
+            }
+          },
+        );
+      });
     };
 
     // Assign the running task to the shared promise variable, and ensure

handwritten/spanner/src/session-pool.ts

@@ -689,49 +689,55 @@ export class SessionPool extends EventEmitter implements SessionPoolInterface {
       opts: this._observabilityOptions,
       dbName: this.database.formattedName_,
     };
-    return startTrace('SessionPool.createSessions', traceConfig, async span => {
-      span.addEvent(`Requesting ${amount} sessions`);
-
-      // while we can request as many sessions be created as we want, the backend
-      // will return at most 100 at a time, hence the need for a while loop.
-      while (amount > 0) {
-        let sessions: Session[] | null = null;
-
-        span.addEvent(`Creating ${amount} sessions`);
-
-        try {
-          [sessions] = await this.database.batchCreateSessions({
-            count: amount,
-            labels: labels,
-            databaseRole: databaseRole,
-          });
-
-          amount -= sessions.length;
-          nReturned += sessions.length;
-        } catch (e) {
-          this._pending -= amount;
-          this.emit('createError', e);
+    return context.with(ROOT_CONTEXT, () => {
+      return startTrace(
+        'SessionPool.createSessions',
+        traceConfig,
+        async span => {
+          span.addEvent(`Requesting ${amount} sessions`);
+
+          // while we can request as many sessions be created as we want, the backend
+          // will return at most 100 at a time, hence the need for a while loop.
+          while (amount > 0) {
+            let sessions: Session[] | null = null;
+
+            span.addEvent(`Creating ${amount} sessions`);
+
+            try {
+              [sessions] = await this.database.batchCreateSessions({
+                count: amount,
+                labels: labels,
+                databaseRole: databaseRole,
+              });
+
+              amount -= sessions.length;
+              nReturned += sessions.length;
+            } catch (e) {
+              this._pending -= amount;
+              this.emit('createError', e);
+              span.addEvent(
+                `Requested for ${nRequested} sessions returned ${nReturned}`,
+              );
+              setSpanErrorAndException(span, e as Error);
+              span.end();
+              throw e;
+            }
+
+            sessions.forEach((session: Session) => {
+              setImmediate(() => {
+                this._inventory.borrowed.add(session);
+                this._pending -= 1;
+                this.release(session);
+              });
+            });
+          }
+
           span.addEvent(
             `Requested for ${nRequested} sessions returned ${nReturned}`,
           );
-          setSpanErrorAndException(span, e as Error);
           span.end();
-          throw e;
-        }
-
-        sessions.forEach((session: Session) => {
-          setImmediate(() => {
-            this._inventory.borrowed.add(session);
-            this._pending -= 1;
-            this.release(session);
-          });
-        });
-      }
-
-      span.addEvent(
-        `Requested for ${nRequested} sessions returned ${nReturned}`,
+        },
       );
-      span.end();
     });
   }