src: invoke await for operations inisde some traces to ensure they return before span.end()

Author: odeke-em

src/database.ts

@@ -3214,7 +3214,7 @@ class Database extends common.GrpcServiceObject {
         : {};
 
     startTrace('Database.runTransaction', this._traceConfig, span => {
-      this.pool_.getSession((err, session?, transaction?) => {
+      this.pool_.getSession(async (err, session?, transaction?) => {
         if (err) {
           setSpanError(span, err);
         }
@@ -3223,7 +3223,7 @@ class Database extends common.GrpcServiceObject {
           span.addEvent('No session available', {
             'session.id': session?.id,
           });
-          this.runTransaction(options, runFn!);
+          await this.runTransaction(options, runFn!);
           span.end();
           return;
         }
@@ -3259,7 +3259,6 @@ class Database extends common.GrpcServiceObject {
               setSpanError(span, err!);
             }
             await runFn!(err, resp);
-            span.end();
           },
           options
         );

src/transaction.ts

@@ -1320,7 +1320,7 @@ export class Snapshot extends EventEmitter {
             this._update(response.metadata!.transaction);
           }
         })
-        .on('error', err => {
+        .on('error', async err => {
           setSpanError(span, err as Error);
           const isServiceError =
             err && typeof err === 'object' && 'code' in err;
@@ -1332,7 +1332,7 @@ export class Snapshot extends EventEmitter {
               (err as grpc.ServiceError).code === grpc.status.ABORTED
             )
           ) {
-            this.begin();
+            await this.begin();
           }
         })
         .on('end', err => {