More span.end() checks

Author: odeke-em

observability-test/spanner.ts

@@ -2177,7 +2177,6 @@ describe('Traces for ExecuteStream broken stream retries', () => {
       await tx!.runUpdate(insertSql, err => {
         assert.ok(err, 'Missing expected error');
         assert.strictEqual(err!.code, grpc.status.INVALID_ARGUMENT);
-        console.log('Got expected error');
         assert.strictEqual(attempts, 1);
         tx!
           .commit()

src/database.ts

@@ -2101,6 +2101,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);
               });
@@ -3223,7 +3226,8 @@ class Database extends common.GrpcServiceObject {
         txn!.once('end', () => {
           span.end();
         });
-        txn!.once('error', () => {
+        txn!.once('error', err => {
+          setSpanError(span, err!);
           span.end();
         });
         runFn!(null, txn!);
@@ -3245,8 +3249,6 @@ class Database extends common.GrpcServiceObject {
         }
 
         if (err) {
-          // This code path failed to even begin running the transaction
-          // hence end the span immediately and return it.
           span.end();
           runFn!(err as grpc.ServiceError);
           return;
@@ -3553,6 +3555,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();
@@ -3652,10 +3659,17 @@ 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) {
           setSpanError(span, err);
           span.end();