chore: add commit retry logic (#2322)

* chore: add commit retry logic

* refactor test

Author: alkatrivedi

src/transaction.ts

@@ -229,7 +229,8 @@ export type CommitCallback =
 type PrecommitTokenProvider =
   | spannerClient.spanner.v1.ITransaction
   | spannerClient.spanner.v1.IPartialResultSet
-  | spannerClient.spanner.v1.IExecuteBatchDmlResponse;
+  | spannerClient.spanner.v1.IExecuteBatchDmlResponse
+  | spannerClient.spanner.v1.ICommitResponse;
 
 /**
  * @typedef {object} TimestampBounds
@@ -1831,6 +1832,7 @@ export class Transaction extends Dml {
   commitTimestamp?: PreciseDate;
   commitTimestampProto?: spannerClient.protobuf.ITimestamp;
   private _queuedMutations: spannerClient.spanner.v1.Mutation[];
+  private _retryCommit: Boolean;
 
   /**
    * Timestamp at which the transaction was committed. Will be populated once
@@ -1884,6 +1886,7 @@ export class Transaction extends Dml {
     this._options = {readWrite: options};
     this._options.isolationLevel = IsolationLevel.ISOLATION_LEVEL_UNSPECIFIED;
     this.requestOptions = requestOptions;
+    this._retryCommit = false;
   }
 
   /**
@@ -2291,6 +2294,16 @@ export class Transaction extends Dml {
             err: null | Error,
             resp: spannerClient.spanner.v1.ICommitResponse,
           ) => {
+            if (
+              resp &&
+              'MultiplexedSessionRetry' in resp &&
+              !this._retryCommit
+            ) {
+              this._retryCommit = true;
+              this._updatePrecommitToken(resp);
+              return this.commit(options, callback);
+            }
+
             this.end();
 
             if (err) {

test/session-factory.ts

@@ -130,6 +130,10 @@ describe('SessionFactory', () => {
         process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS = 'true';
       });
 
+      after(() => {
+        process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS = 'false';
+      });
+
       it('should create a MultiplexedSession object', () => {
         assert(
           sessionFactory.multiplexedSession_ instanceof MultiplexedSession,
@@ -171,6 +175,11 @@ describe('SessionFactory', () => {
         process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_FOR_RW = 'true';
       });
 
+      after(() => {
+        process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS = 'false';
+        process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_FOR_RW = 'false';
+      });
+
       it('should correctly initialize the isMultiplexedRW field', () => {
         const sessionFactory = new SessionFactory(DATABASE, NAME, POOL_OPTIONS);
         assert.strictEqual(sessionFactory.isMultiplexedRW, true);
@@ -213,6 +222,10 @@ describe('SessionFactory', () => {
         process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS = 'true';
       });
 
+      after(() => {
+        process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS = 'false';
+      });
+
       it('should return the multiplexed session', done => {
         (
           sandbox.stub(
@@ -283,6 +296,11 @@ describe('SessionFactory', () => {
         process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_FOR_RW = 'true';
       });
 
+      after(() => {
+        process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS = 'false';
+        process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_FOR_RW = 'false';
+      });
+
       it('should return the multiplexed session', done => {
         (
           sandbox.stub(
@@ -330,6 +348,10 @@ describe('SessionFactory', () => {
         process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS = 'true';
       });
 
+      after(() => {
+        process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS = 'false';
+      });
+
       it('should not call the release method', () => {
         const releaseStub = sandbox.stub(sessionFactory.pool_, 'release');
         const fakeMuxSession = createMuxSession();

test/transaction.ts

@@ -1879,6 +1879,65 @@ describe('Transaction', () => {
         assert.strictEqual(transaction.commitTimestampProto, fakeTimestamp);
       });
 
+      it('should retry commit only once upon sending precommitToken to read-only participants', () => {
+        const requestStub = sandbox.stub(transaction, 'request');
+
+        const expectedTimestamp = new PreciseDate(0);
+        const fakeTimestamp = {seconds: 0, nanos: 0};
+
+        const fakeResponse = {commitTimestamp: fakeTimestamp};
+        const fakePrecommitToken = {
+          precommitToken: Buffer.from('precommit-token-commit'),
+          seqNum: 1,
+        };
+
+        transaction._latestPreCommitToken = fakePrecommitToken;
+
+        // retry response on commit retry
+        const fakeCommitRetryResponse = {
+          commitTimestamp: null,
+          MultiplexedSessionRetry: 'precommitToken',
+          precommitToken: {
+            precommitToken: Buffer.from('precommit-token-commit-retry'),
+            seqNum: 2,
+          },
+        };
+
+        requestStub.onFirstCall().callsFake((_, callback) => {
+          // assert that the transaction contains the precommit token
+          assert.deepStrictEqual(
+            transaction._latestPreCommitToken,
+            fakePrecommitToken,
+          );
+          // retry commit response
+          callback(null, fakeCommitRetryResponse);
+        });
+
+        requestStub.onSecondCall().callsFake((_, callback) => {
+          // assert that before second commit retry the _latestPreCommitToken
+          // containing the commit retry reponse in the transaction object
+          assert.deepStrictEqual(
+            transaction._latestPreCommitToken,
+            fakeCommitRetryResponse.precommitToken,
+          );
+          callback(null, fakeResponse);
+        });
+
+        transaction.commit((err, resp) => {
+          // assert there is no error
+          assert.ifError(err);
+          // make sure that retry happens only once
+          assert.strictEqual(requestStub.callCount, 2);
+          assert.deepStrictEqual(
+            transaction.commitTimestamp,
+            expectedTimestamp,
+          );
+          assert.strictEqual(transaction.commitTimestampProto, fakeTimestamp);
+          // assert on the successfull commit response
+          assert.deepStrictEqual(resp, fakeResponse);
+        });
+      });
+
       it('should return any errors and the response', () => {
         const requestStub = sandbox.stub(transaction, 'request');