Spans for abort transactions

Author: surbhigarg92

observability-test/spanner.ts

@@ -142,6 +142,11 @@ describe('EndToEnd', async () => {
     tracerProvider: tracerProvider,
     enableExtendedTracing: false,
   });
+  let dbCounter = 1;
+
+  function newTestDatabase(): Database {
+    return instance.database(`database-${dbCounter++}`,);
+  }
 
   const server = setupResult.server;
   const spannerMock = setupResult.spannerMock;
@@ -305,6 +310,7 @@ describe('EndToEnd', async () => {
     });
 
     it('runTransactionAsync', async () => {
+      
       await database.runTransactionAsync(async transaction => {
         await transaction!.run('SELECT 1');
       });
@@ -327,6 +333,105 @@ describe('EndToEnd', async () => {
       );
     });
 
+    it.only('runTransaction with abort', done => {
+      let attempts = 0;
+      let rowCount = 0;
+      const database = newTestDatabase();
+      database.runTransaction(async (err, transaction) => {
+        assert.ifError(err);
+        if (!attempts) {
+          spannerMock.abortTransaction(transaction!);
+        }
+        attempts++;
+        transaction!.run(selectSql, (err, rows) => {
+          assert.ifError(err);
+          rows.forEach(() => rowCount++);
+          transaction!
+            .commit()
+            .catch(done)
+            .then(async () => {
+              const expectedSpanNames = [
+                'CloudSpanner.Database.batchCreateSessions',
+                'CloudSpanner.SessionPool.createSessions',
+                'CloudSpanner.Snapshot.runStream',
+                'CloudSpanner.Snapshot.run',
+                'CloudSpanner.Snapshot.runStream',
+                'CloudSpanner.Snapshot.run',
+                'CloudSpanner.Transaction.commit',
+                'CloudSpanner.Snapshot.begin',
+                'CloudSpanner.Database.runTransaction',
+              ];
+              const expectedEventNames = [
+                ...waitingSessionsEvents,
+                'Retrying Transaction',
+                'Starting stream',
+                'exception',
+                'Stream broken. Not safe to retry',
+                'Begin Transaction',
+                'Transaction Creation Done',
+                'Starting stream',
+                'Starting Commit',
+                'Commit Done',
+              ];
+              await verifySpansAndEvents(traceExporter, expectedSpanNames, expectedEventNames)
+              database
+                .close()
+                .catch(done)
+                .then(() => done());
+            });
+        });
+          });
+    });
+
+    it('runTransactionAsync with abort', async () => {
+      let attempts = 0;
+      const database = newTestDatabase();
+      await database.runTransactionAsync((transaction): Promise<number> => {
+        if (!attempts) {
+          spannerMock.abortTransaction(transaction);
+        }
+        attempts++;
+        return transaction.run(selectSql).then(([rows]) => {
+          let count = 0;
+          rows.forEach(() => count++);
+          return transaction.commit().then(() => count);
+        });
+      });
+      assert.strictEqual(attempts, 2);
+      const expectedSpanNames = [
+        'CloudSpanner.Database.batchCreateSessions',
+        'CloudSpanner.SessionPool.createSessions',
+        'CloudSpanner.Snapshot.runStream',
+        'CloudSpanner.Snapshot.run',
+        'CloudSpanner.Snapshot.begin',
+        'CloudSpanner.Snapshot.runStream',
+        'CloudSpanner.Snapshot.run',
+        'CloudSpanner.Transaction.commit',
+        'CloudSpanner.Database.runTransactionAsync',
+      ];
+      const expectedEventNames = [
+        'Requesting 25 sessions',
+        'Creating 25 sessions',
+        'Requested for 25 sessions returned 25',
+        'Starting stream',
+        'exception',
+        'Stream broken. Not safe to retry',
+        'Begin Transaction',
+        'Transaction Creation Done',
+        'Starting stream',
+        'Starting Commit',
+        'Commit Done',
+        ...waitingSessionsEvents,
+        'Retrying transaction',
+      ];
+      await verifySpansAndEvents(
+        traceExporter,
+        expectedSpanNames,
+        expectedEventNames
+      );
+      await database.close();
+    });
+
     it('writeAtLeastOnce', done => {
       const blankMutations = new MutationSet();
       database.writeAtLeastOnce(blankMutations, async (err, response) => {

src/partial-result-stream.ts

@@ -24,7 +24,7 @@ import {Readable, Transform} from 'stream';
 import * as streamEvents from 'stream-events';
 import {grpc, CallOptions} from 'google-gax';
 import {DeadlineError, isRetryableInternalError} from './transaction-runner';
-
+import {getActiveOrNoopSpan, setSpanErrorAndException} from './instrument';
 import {codec, JSONOptions, Json, Field, Value} from './codec';
 import {google} from '../protos/protos';
 import * as stream from 'stream';
@@ -494,6 +494,7 @@ export function partialResultStream(
   let lastRequestStream: Readable;
   const startTime = Date.now();
   const timeout = options?.gaxOptions?.timeout ?? Infinity;
+  const span = getActiveOrNoopSpan();
 
   // mergeStream allows multiple streams to be connected into one. This is good;
   // if we need to retry a request and pipe more data to the user's stream.
@@ -568,6 +569,7 @@ export function partialResultStream(
       // checkpoint stream has queued. After that, we will destroy the
       // user's stream with the same error.
       setImmediate(() => batchAndSplitOnTokenStream.destroy(err));
+      setSpanErrorAndException(span, err as Error);
       return;
     }
 

src/transaction-runner.ts

@@ -23,6 +23,7 @@ import {Session} from './session';
 import {Transaction} from './transaction';
 import {NormalCallback} from './common';
 import {isSessionNotFoundError} from './session-pool';
+import {getActiveOrNoopSpan} from './instrument';
 import {Database} from './database';
 import {google} from '../protos/protos';
 import IRequestOptions = google.spanner.v1.IRequestOptions;
@@ -238,6 +239,7 @@ export abstract class Runner<T> {
         this.session.lastError = e as grpc.ServiceError;
         lastError = e as grpc.ServiceError;
       }
+      const span = getActiveOrNoopSpan();
 
       // Note that if the error is a 'Session not found' error, it will be
       // thrown here. We do this to bubble this error up to the caller who is
@@ -250,7 +252,7 @@ export abstract class Runner<T> {
       }
 
       this.attempts += 1;
-
+      span.addEvent('Retrying transaction');
       const delay = this.getNextDelay(lastError);
       await new Promise(resolve => setTimeout(resolve, delay));
     }
@@ -321,6 +323,7 @@ export class TransactionRunner extends Runner<void> {
           }
 
           stream.unpipe(proxyStream);
+          // proxyStream.emit('error', err);
           reject(err);
         })
         .pipe(proxyStream);