Ensure database.runTransaction async awaits on caller code + ensure context.active only in Spanner constructor

Author: odeke-em

src/database.ts

@@ -3223,8 +3223,8 @@ class Database extends common.GrpcServiceObject {
           span.addEvent('No session available', {
             'session.id': session?.id,
           });
-          span.end();
           this.runTransaction(options, runFn!);
+          span.end();
           return;
         }
 
@@ -3244,18 +3244,22 @@ class Database extends common.GrpcServiceObject {
 
         const release = () => {
           this.pool_.release(session!);
-          span.end();
+          if (span.isRecording()) {
+            // span.end() might have already been invoked inside
+            // Transactionrunner.
+            span.end();
+          }
         };
 
         const runner = new TransactionRunner(
           session!,
           transaction!,
-          (err, resp) => {
+          async (err, resp) => {
             if (err) {
               setSpanError(span, err!);
             }
+            await runFn!(err, resp);
             span.end();
-            runFn!(err, resp);
           },
           options
         );

src/index.ts

@@ -80,7 +80,10 @@ import {
 import grpcGcpModule = require('grpc-gcp');
 const grpcGcp = grpcGcpModule(grpc);
 import * as v1 from './v1';
-import {ObservabilityOptions} from './instrument';
+import {
+  ObservabilityOptions,
+  ensureInitialContextManagerSet,
+} from './instrument';
 
 // eslint-disable-next-line @typescript-eslint/no-var-requires
 const gcpApiConfig = require('./spanner_grpc_config.json');
@@ -370,6 +373,7 @@ class Spanner extends GrpcService {
     };
     this.directedReadOptions = directedReadOptions;
     this._observabilityOptions = options.observabilityOptions;
+    ensureInitialContextManagerSet();
   }
 
   /**

src/instrument.ts

@@ -117,6 +117,7 @@ function ensureInitialContextManagerSet() {
     context.setGlobalContextManager(contextManager);
   }
 }
+export {ensureInitialContextManagerSet};
 
 const debugTraces = process.env.SPANNER_DEBUG_TRACES === 'true';
 
@@ -138,14 +139,19 @@ export function startTrace<T>(
     config = {} as traceConfig;
   }
 
-  ensureInitialContextManagerSet();
+  const tracer = getTracer(config.opts?.tracerProvider);
+  let parentSpanId: string | undefined;
+  if (debugTraces) {
+    const parentSpan = trace.getActiveSpan();
+    parentSpanId = parentSpan?.spanContext()!.spanId;
+  }
 
-  return getTracer(config.opts?.tracerProvider).startActiveSpan(
+  return tracer.startActiveSpan(
     SPAN_NAMESPACE_PREFIX + '.' + spanNameSuffix,
     {kind: SpanKind.CLIENT},
     span => {
       if (debugTraces) {
-        patchSpanEndForDebugging(span, spanNameSuffix);
+        patchSpanEndForDebugging(span, spanNameSuffix, parentSpanId);
       }
 
       span.setAttribute(SEMATTRS_DB_SYSTEM, 'spanner');
@@ -305,10 +311,19 @@ class noopSpan implements Span {
   }
 }
 
-function patchSpanEndForDebugging(span: Span, spanNameSuffix: string) {
+function patchSpanEndForDebugging(
+  span: Span,
+  spanNameSuffix: string,
+  parentSpanId: string | undefined
+) {
+  console.trace(
+    `\x1b[33m${spanNameSuffix}.start id=${span.spanContext().spanId} with parentSpanId: ${parentSpanId}\x1b[00m`
+  );
   const origSpanEnd = span.end;
   const wrapSpanEnd = function (this: Span) {
-    console.trace(`\x1b[35m${spanNameSuffix}.end()\x1b[00m`);
+    console.trace(
+      `\x1b[35m${spanNameSuffix}.end() id=${span.spanContext().spanId} with parentSpanId: ${parentSpanId}\x1b[00m`
+    );
     return origSpanEnd.apply(this);
   };
   Object.defineProperty(span, 'end', {

src/transaction.ts

@@ -744,6 +744,7 @@ export class Snapshot extends EventEmitter {
           }
         })
         .on('error', err => {
+          setSpanError(span, err);
           const isServiceError =
             err && typeof err === 'object' && 'code' in err;
           if (
@@ -754,9 +755,11 @@ export class Snapshot extends EventEmitter {
               (err as grpc.ServiceError).code === grpc.status.ABORTED
             )
           ) {
+            span.addEvent('Transaction Aborted, retrying begin', {
+              'transaction.id': this.id?.toString(),
+            });
             this.begin();
           }
-          setSpanError(span, err);
         })
         .on('end', err => {
           if (err) {