test: commit retry logic

Author: alkatrivedi

test/mockserver/mockspanner.ts

@@ -325,7 +325,6 @@ export class MockSpanner {
     string,
     protobuf.Session
   >();
-  private multiplexedSession: boolean | undefined;
   private begin: boolean;
   private transactionCounters: Map<string, number> = new Map<string, number>();
   private transactions: Map<string, protobuf.Transaction> = new Map<
@@ -473,7 +472,6 @@ export class MockSpanner {
       createTime: now(),
     });
     this.sessions.set(name, session);
-    this.multiplexedSession = multiplexed;
     return session;
   }
 
@@ -695,13 +693,14 @@ export class MockSpanner {
           }
         }
         const res = this.statementResults.get(call.request!.sql);
+        const session = this.sessions.get(call.request!.session);
         if (res) {
           if (call.request!.transaction?.begin) {
             const txn = this._updateTransaction(
               call.request!.session,
               call.request!.transaction.begin,
             );
-            const precommitToken = this.multiplexedSession
+            const precommitToken = session?.multiplexed
               ? protobuf.MultiplexedSessionPrecommitToken.create({
                   precommitToken: Buffer.from('mock-precommit-token'),
                   seqNum: randomInt(1, 1000),
@@ -1012,13 +1011,14 @@ export class MockSpanner {
         );
         const key = `${call.request!.table}|${keySet}`;
         const res = this.readRequestResults.get(key);
+        const session = this.sessions.get(call.request!.session);
         if (res) {
           if (call.request!.transaction?.begin) {
             const txn = this._updateTransaction(
               call.request!.session,
               call.request!.transaction.begin,
             );
-            const precommitToken = this.multiplexedSession
+            const precommitToken = session?.multiplexed
               ? protobuf.MultiplexedSessionPrecommitToken.create({
                   precommitToken: Buffer.from('mock-precommit-token'),
                   seqNum: randomInt(1, 1000),

test/spanner.ts

@@ -2099,6 +2099,12 @@ describe('Spanner with mock server', () => {
           'false';
       });
 
+      after(() => {
+        process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS = 'false';
+        process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_PARTITIONED_OPS =
+          'false';
+      });
+
       it('should execute the transaction(database.runPartitionedUpdate) successfully using regular/pool session', done => {
         const database = newTestDatabase({min: 1, max: 1});
         const pool = (database.sessionFactory_ as SessionFactory)
@@ -2127,6 +2133,12 @@ describe('Spanner with mock server', () => {
           'true';
       });
 
+      after(() => {
+        process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS = 'false';
+        process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_PARTITIONED_OPS =
+          'false';
+      });
+
       it('should execute the transaction(database.runPartitionedUpdate) successfully using regular/pool session', done => {
         const database = newTestDatabase({min: 1, max: 1});
         const pool = (database.sessionFactory_ as SessionFactory)
@@ -2154,6 +2166,12 @@ describe('Spanner with mock server', () => {
           'true';
       });
 
+      after(() => {
+        process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS = 'false';
+        process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_PARTITIONED_OPS =
+          'false';
+      });
+
       it('should execute the transaction(database.runPartitionedUpdate) successfully using multiplexed session', done => {
         const database = newTestDatabase({min: 1, max: 1});
         const pool = (database.sessionFactory_ as SessionFactory)
@@ -4452,6 +4470,86 @@ describe('Spanner with mock server', () => {
           });
         });
       });
+
+      // test(s) for commit retry logic
+      describe('Transaction Commit Retry Logic', () => {
+        let commitCallCount = 0;
+        let capturedCommitRequests: any[] = [];
+
+        it('should retry commit only once with a precommit token', async () => {
+          commitCallCount = 0;
+          capturedCommitRequests = [];
+
+          const database = newTestDatabase({min: 1, max: 1});
+          const fakeRetryToken = Buffer.from('mock-retry-token-123');
+
+          const commitRetryResponse = {
+            MultiplexedSessionRetry: 'precommitToken',
+            precommitToken: {
+              precommitToken: fakeRetryToken,
+              seqNum: 1,
+            },
+            commitTimestamp: mock.now(),
+          };
+
+          const commitSuccessResponse = {
+            commitTimestamp: mock.now(),
+          };
+
+          await database.runTransactionAsync(async tx => {
+            // mock commit request
+            tx.request = (config: any, callback: Function) => {
+              const cb = callback as (err: any, response: any) => void;
+
+              if (config.method !== 'commit') return;
+
+              commitCallCount++;
+              capturedCommitRequests.push(config.reqOpts);
+
+              if (commitCallCount === 1) {
+                cb(null, commitRetryResponse);
+              } else {
+                cb(null, commitSuccessResponse);
+              }
+            };
+
+            // perform read
+            await tx!.run(selectSql);
+
+            // perform mutations
+            await tx.upsert('foo', [
+              {id: 1, name: 'One'},
+              {id: 2, name: 'Two'},
+            ]);
+
+            // make a call to commit
+            await tx.commit();
+
+            // assert that retry heppen only once
+            assert.strictEqual(
+              commitCallCount,
+              2,
+              'The mock commit method should have been called exactly twice.',
+            );
+            const firstRequest = capturedCommitRequests[0];
+            // assert that during the first request to commit
+            // the precommitToken was missing
+            assert.ok(
+              !firstRequest.precommitToken,
+              'The first commit request should not have a precommitToken.',
+            );
+            const secondRequest = capturedCommitRequests[1];
+            // assert that during the second request to commit
+            // the precommitToken was present
+            assert.deepStrictEqual(
+              secondRequest.precommitToken,
+              commitRetryResponse.precommitToken,
+              'The second commit request should have the precommitToken from the retry response.',
+            );
+          });
+          await database.close();
+        });
+      });
     });
   });