feat(observability): fix bugs found from product review + negative cases

This change adds recording of retry span annotations,
catching cases in which exceptions where thrown but spans
were not ended while testing out and visually confirming the results.

Author: odeke-em

observability-test/spanner.ts

@@ -479,7 +479,7 @@ describe('Transaction', () => {
             `span names mismatch:\n\tGot:  ${actualSpanNames}\n\tWant: ${expectedSpanNames}`
           );
 
-          const expectedEventNames = ['exception'];
+          const expectedEventNames = ['Starting stream', 'exception'];
           assert.deepStrictEqual(
             actualEventNames,
             expectedEventNames,

observability-test/transaction.ts

@@ -1130,6 +1130,8 @@ class Database extends common.GrpcServiceObject {
       } catch (e) {
         setSpanErrorAndException(span, e as Error);
         this.emit('error', e);
+      } finally {
+        span.end();
       }
     });
   }
@@ -2101,6 +2103,9 @@ class Database extends common.GrpcServiceObject {
               session!.lastError = err;
               this.pool_.release(session!);
               this.getSnapshot(options, (err, snapshot) => {
+                if (err) {
+                  setSpanError(span, err);
+                }
                 span.end();
                 callback!(err, snapshot);
               });
@@ -2815,6 +2820,7 @@ class Database extends common.GrpcServiceObject {
       this.runStream(query, options)
         .on('error', err => {
           setSpanError(span, err);
+          span.end();
           callback!(err as grpc.ServiceError, rows, stats, metadata);
         })
         .on('response', response => {
@@ -3054,7 +3060,6 @@ class Database extends common.GrpcServiceObject {
         let dataStream = snapshot.runStream(query);
 
         const endListener = () => {
-          span.end();
           snapshot.end();
         };
         dataStream
@@ -3083,6 +3088,11 @@ class Database extends common.GrpcServiceObject {
               // Create a new data stream and add it to the end user stream.
               dataStream = this.runStream(query, options);
               dataStream.pipe(proxyStream);
+              dataStream.on('end', () => span.end());
+              dataStream.on('error', err => {
+                setSpanError(span, err);
+                span.end();
+              });
             } else {
               proxyStream.destroy(err);
               snapshot.end();
@@ -3098,7 +3108,9 @@ class Database extends common.GrpcServiceObject {
         if (err) {
           setSpanError(span, err);
         }
-        span.end();
+        if (span.isRecording()) {
+          span.end();
+        }
       });
 
       return proxyStream as PartialResultStream;
@@ -3212,6 +3224,26 @@ class Database extends common.GrpcServiceObject {
         ? (optionsOrRunFn as RunTransactionOptions)
         : {};
 
+    const retry = (span: Span) => {
+      this.runTransaction(options, (err, txn) => {
+        if (err) {
+          setSpanError(span, err);
+          span.end();
+          runFn!(err, null);
+          return;
+        }
+
+        txn!.once('end', () => {
+          span.end();
+        });
+        txn!.once('error', err => {
+          setSpanError(span, err!);
+          span.end();
+        });
+        runFn!(null, txn!);
+      });
+    };
+
     startTrace('Database.runTransaction', this._traceConfig, span => {
       this.pool_.getSession((err, session?, transaction?) => {
         if (err) {
@@ -3222,8 +3254,7 @@ class Database extends common.GrpcServiceObject {
           span.addEvent('No session available', {
             'session.id': session?.id,
           });
-          this.runTransaction(options, runFn!);
-          span.end();
+          retry(span);
           return;
         }
 
@@ -3241,41 +3272,42 @@ class Database extends common.GrpcServiceObject {
           transaction!.excludeTxnFromChangeStreams();
         }
 
-        const release = () => {
+        // Our span should only be ended if the
+        // transaction either errored or was ended.
+        transaction!.once('error', err => {
+          setSpanError(span, err!);
+          span.end();
+        });
+
+        transaction!.once('end', err => {
+          setSpanError(span, err!);
           span.end();
+        });
+
+        const release = () => {
           this.pool_.release(session!);
         };
 
         const runner = new TransactionRunner(
           session!,
           transaction!,
-          (err, resp) => {
-            if (err) {
-              setSpanError(span, err!);
-            }
-            span.end();
-            runFn!(err, resp);
-          },
+          runFn!,
           options
         );
 
         runner.run().then(release, err => {
-          if (err) {
-            setSpanError(span, err!);
-          }
+          setSpanError(span, err);
 
           if (isSessionNotFoundError(err)) {
             span.addEvent('No session available', {
               'session.id': session?.id,
             });
             release();
-            this.runTransaction(options, runFn!);
+            retry(span);
           } else {
-            if (!err) {
-              span.addEvent('Using Session', {'session.id': session!.id});
-            }
             setImmediate(runFn!, err);
             release();
+            span.end();
           }
         });
       });
@@ -3536,6 +3568,11 @@ class Database extends common.GrpcServiceObject {
                   mutationGroups,
                   options
                 );
+                dataStream.once('end', () => span.end());
+                dataStream.once('error', err => {
+                  setSpanError(span, err!);
+                  span.end();
+                });
                 dataStream.pipe(proxyStream);
               } else {
                 span.end();
@@ -3635,8 +3672,14 @@ class Database extends common.GrpcServiceObject {
           span.addEvent('No session available', {
             'session.id': session?.id,
           });
-          this.writeAtLeastOnce(mutations, options, cb!);
-          span.end();
+          // Retry this method.
+          this.writeAtLeastOnce(mutations, options, (err, resp) => {
+            if (err) {
+              setSpanError(span, err);
+            }
+            span.end();
+            cb!(err, resp);
+          });
           return;
         }
         if (err) {

src/database.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/index.ts

@@ -117,6 +117,7 @@ function ensureInitialContextManagerSet() {
     context.setGlobalContextManager(contextManager);
   }
 }
+export {ensureInitialContextManagerSet};
 
 /**
  * startTrace begins an active span in the current active context
@@ -136,8 +137,6 @@ export function startTrace<T>(
     config = {} as traceConfig;
   }
 
-  ensureInitialContextManagerSet();
-
   return getTracer(config.opts?.tracerProvider).startActiveSpan(
     SPAN_NAMESPACE_PREFIX + '.' + spanNameSuffix,
     {kind: SpanKind.CLIENT},
@@ -165,11 +164,20 @@ export function startTrace<T>(
         }
       }
 
-      if (config.that) {
-        const fn = cb.bind(config.that);
-        return fn(span);
-      } else {
-        return cb(span);
+      // If at all the invoked function throws an exception,
+      // record the exception and then end this span.
+      try {
+        if (config.that) {
+          const fn = cb.bind(config.that);
+          return fn(span);
+        } else {
+          return cb(span);
+        }
+      } catch (e) {
+        setSpanErrorAndException(span, e as Error);
+        span.end();
+        // Finally re-throw the exception.
+        throw e;
       }
     }
   );

src/instrument.ts

@@ -35,6 +35,7 @@ import {
   ObservabilityOptions,
   startTrace,
   setSpanError,
+  setSpanErrorAndException,
   traceConfig,
 } from './instrument';