Plumb more with await to ensure correct ending of callbacks before span.end

Author: odeke-em

src/database.ts

@@ -3229,8 +3229,8 @@ class Database extends common.GrpcServiceObject {
         }
 
         if (err) {
+          await runFn!(err as grpc.ServiceError);
           span.end();
-          runFn!(err as grpc.ServiceError);
           return;
         }
 
@@ -3244,11 +3244,6 @@ class Database extends common.GrpcServiceObject {
 
         const release = () => {
           this.pool_.release(session!);
-          if (span.isRecording()) {
-            // span.end() might have already been invoked inside
-            // Transactionrunner.
-            span.end();
-          }
         };
 
         const runner = new TransactionRunner(
@@ -3259,26 +3254,27 @@ class Database extends common.GrpcServiceObject {
               setSpanError(span, err!);
             }
             await runFn!(err, resp);
-            span.end();
           },
           options
         );
 
-        runner.run().then(release, err => {
+        await runner.run().then(release, async err => {
           setSpanError(span, err);
 
           if (isSessionNotFoundError(err)) {
             span.addEvent('No session available', {
               'session.id': session?.id,
             });
             release();
-            this.runTransaction(options, runFn!);
+            await this.runTransaction(options, runFn!);
           } else {
             span.addEvent('Using Session', {'session.id': session!.id});
             setImmediate(runFn!, err);
             release();
           }
         });
+
+        span.end();
       });
     });
   }

src/table.ts

@@ -29,6 +29,7 @@ import {
   CommitCallback,
 } from './transaction';
 import {google as databaseAdmin} from '../protos/protos';
+import {google} from '../protos/protos';
 import {Schema, LongRunningCallback} from './common';
 import IRequestOptions = databaseAdmin.spanner.v1.IRequestOptions;
 import {
@@ -1105,22 +1106,23 @@ class Table {
           requestOptions: requestOptions,
           excludeTxnFromChangeStreams: excludeTxnFromChangeStreams,
         },
-        (err, transaction) => {
+        async (err, transaction) => {
           if (err) {
             setSpanError(span, err);
+            await callback(err);
             span.end();
-            callback(err);
             return;
           }
 
-          transaction![method](this.name, rows as Key[]);
-          transaction!.commit(options, (err, resp) => {
-            if (err) {
-              setSpanError(span, err);
-            }
+          try {
+            transaction![method](this.name, rows as Key[]);
+            const resp = await transaction!.commit(options);
+            await callback(err, resp as google.spanner.v1.ICommitResponse);
+          } catch (e) {
+            await callback(err, null);
+          } finally {
             span.end();
-            callback(err, resp);
-          });
+          }
         }
       );
     });

src/transaction.ts

@@ -456,19 +456,18 @@ export class Snapshot extends EventEmitter {
           gaxOpts,
           headers: headers,
         },
-        (
+        async (
           err: null | grpc.ServiceError,
           resp: spannerClient.spanner.v1.ITransaction
         ) => {
           if (err) {
             setSpanError(span, err);
-            span.end();
-            callback!(err, resp);
-            return;
+          } else {
+            this._update(resp);
           }
-          this._update(resp);
+
+          await callback!(err, resp);
           span.end();
-          callback!(null, resp);
         }
       );
     });
@@ -2149,35 +2148,25 @@ export class Transaction extends Dml {
       opts: this._observabilityOptions,
       dbName: this._dbName!,
     };
-    return startTrace('Transaction.commit', traceConfig, span => {
+    startTrace('Transaction.commit', traceConfig, async span => {
       if (this.id) {
         reqOpts.transactionId = this.id as Uint8Array;
       } else if (!this._useInRunner) {
         reqOpts.singleUseTransaction = this._options;
       } else {
-        this.begin()
-          .then(
-            () => {
-              this.commit(options, (err, resp) => {
-                if (err) {
-                  setSpanError(span, err);
-                }
-                span.end();
-                callback(err, resp);
-              });
-            },
-            err => {
-              setSpanError(span, err);
-              callback(err);
-              span.end();
-            }
-          )
-          .catch(err => {
-            setSpanErrorAndException(span, err as Error);
-            span.end();
-            // Re-throw the exception after recording it.
-            throw err;
-          });
+        try {
+          await this.begin();
+          const resp = await this.commit(options);
+          span.end();
+          await callback(
+            null,
+            resp as spannerClient.spanner.v1.ICommitResponse
+          );
+        } catch (err) {
+          setSpanErrorAndException(span, err as Error);
+          span.end();
+          await callback(err as ServiceError, null);
+        }
         return;
       }
 
@@ -2215,7 +2204,10 @@ export class Transaction extends Dml {
           gaxOpts: gaxOpts,
           headers: headers,
         },
-        (err: null | Error, resp: spannerClient.spanner.v1.ICommitResponse) => {
+        async (
+          err: null | Error,
+          resp: spannerClient.spanner.v1.ICommitResponse
+        ) => {
           this.end();
 
           if (err) {