chore: add mutation key heuristic for mutation only transaction (#2345)

Author: alkatrivedi

src/transaction.ts

@@ -311,6 +311,7 @@ export class Snapshot extends EventEmitter {
   _observabilityOptions?: ObservabilityOptions;
   _traceConfig: traceConfig;
   protected _dbName?: string;
+  protected _mutationKey: spannerClient.spanner.v1.Mutation | null;
 
   /**
    * The transaction ID.
@@ -383,6 +384,7 @@ export class Snapshot extends EventEmitter {
       dbName: this._dbName,
     };
     this._latestPreCommitToken = null;
+    this._mutationKey = null;
   }
 
   protected _updatePrecommitToken(resp: PrecommitTokenProvider): void {
@@ -395,6 +397,102 @@ export class Snapshot extends EventEmitter {
     }
   }
 
+  /**
+   * Selects a single representative mutation from a list to be used as the
+   * transaction's `mutationKey`.
+   *
+   * This key is required by Spanner and is sent in the `BeginTransactionRequest`
+   * for read-write transactions that only contain mutations. The selection follows
+   * a two-tiered heuristic to choose the most significant mutation.
+   *
+   * The selection heuristic is as follows:
+   *
+   * 1. Priority of Operation Type: High-priority mutations (`delete`, `update`,
+   * `replace`, `insertOrUpdate`) are always chosen over low-priority
+   * (`insert`) mutations.
+   *
+   * 2. Selection Strategy:
+   * - If any high-priority mutations exist, one is chosen randomly from
+   * that group, ignoring the number of rows.
+   * - If only `insert` mutations exist, the one(s) with the largest number
+   * of rows are identified, and one is chosen randomly from that subset.
+   *
+   * @protected
+   * @param mutations The list of mutations from which to select the key.
+   */
+  protected _setMutationKey(mutations: spannerClient.spanner.v1.Mutation[]) {
+    // return if the list is empty
+    if (mutations.length === 0) {
+      return;
+    }
+
+    // maintain a set of high priority keys
+    const HIGH_PRIORITY_KEYS = new Set([
+      'delete',
+      'update',
+      'replace',
+      'insertOrUpdate',
+    ]);
+
+    // maintain a variable for low priority key
+    const LOW_PRIORITY_KEY = 'insert';
+
+    // Partition mutations into high and low priority groups.
+    const [highPriority, lowPriority] = mutations.reduce(
+      (acc, mutation) => {
+        const key = Object.keys(mutation)[0] as keyof typeof mutation;
+        if (HIGH_PRIORITY_KEYS.has(key)) {
+          acc[0].push(mutation);
+        } else if (key === LOW_PRIORITY_KEY) {
+          acc[1].push(mutation);
+        }
+        // return accumulated mutations list
+        return acc;
+      },
+      [[], []] as [
+        spannerClient.spanner.v1.Mutation[],
+        spannerClient.spanner.v1.Mutation[],
+      ],
+    );
+
+    // Apply the selection logic based on the rules.
+    if (highPriority.length > 0) {
+      // RULE 1: If high-priority keys exist, pick one randomly.
+      const randomIndex = Math.floor(Math.random() * highPriority.length);
+      this._mutationKey = highPriority[randomIndex];
+    } else if (lowPriority.length > 0) {
+      // RULE 2: If only 'insert' key(s) exist, find the one with
+      // highest number of values
+      const {bestCandidates} = lowPriority.reduce(
+        (acc, mutation) => {
+          const size = mutation.insert?.values?.length || 0;
+
+          if (size > acc.maxSize) {
+            // New largest size found, start a new list
+            return {maxSize: size, bestCandidates: [mutation]};
+          }
+          if (size === acc.maxSize) {
+            // Same size as current max, add to list
+            acc.bestCandidates.push(mutation);
+          }
+          // return accumulated mutations list
+          return acc;
+        },
+        {
+          maxSize: -1,
+          bestCandidates: [] as spannerClient.spanner.v1.Mutation[],
+        },
+      );
+
+      // Pick randomly from the largest 'insert' mutation(s).
+      const randomIndex = Math.floor(Math.random() * bestCandidates.length);
+      this._mutationKey = bestCandidates[randomIndex];
+    } else {
+      // No mutations to select from.
+      this._mutationKey = null;
+    }
+  }
+
   /**
    * Modifies transaction selector to include the multiplexed session previous
    * transaction id.
@@ -486,6 +584,10 @@ export class Snapshot extends EventEmitter {
       options,
     };
 
+    if (this._mutationKey) {
+      reqOpts.mutationKey = this._mutationKey;
+    }
+
     // Only hand crafted read-write transactions will be able to set a
     // transaction tag for the BeginTransaction RPC. Also, this.requestOptions
     // is only set in the constructor of Transaction, which is the constructor
@@ -2286,6 +2388,9 @@ export class Transaction extends Dml {
         } else if (!this._useInRunner) {
           reqOpts.singleUseTransaction = this._options;
         } else {
+          if ((this.session.parent as Database).isMuxEnabledForRW_) {
+            this._setMutationKey(mutations);
+          }
           this.begin().then(
             () => {
               this.commit(options, (err, resp) => {

test/mockserver/mockspanner.ts

@@ -31,6 +31,7 @@ import Any = google.protobuf.Any;
 import QueryMode = google.spanner.v1.ExecuteSqlRequest.QueryMode;
 import NullValue = google.protobuf.NullValue;
 import {ExecuteSqlRequest, ReadRequest} from '../../src/transaction';
+import {randomInt} from 'crypto';
 
 const PROTO_PATH = 'spanner.proto';
 const IMPORT_PATH = __dirname + '/../../../protos';
@@ -1232,9 +1233,17 @@ export class MockSpanner {
     const transactionId = id.toString().padStart(12, '0');
     const fullTransactionId = session.name + '/transactions/' + transactionId;
     const readTimestamp = options && options.readOnly ? now() : undefined;
+    const precommitToken =
+      session.multiplexed && options?.readWrite
+        ? {
+            precommitToken: Buffer.from('mock-precommit-token'),
+            seqNum: randomInt(1, 1000),
+          }
+        : null;
     const transaction = protobuf.Transaction.create({
       id: Buffer.from(transactionId),
       readTimestamp,
+      precommitToken,
     });
     this.transactions.set(fullTransactionId, transaction);
     this.transactionOptions.set(fullTransactionId, options);

test/spanner.ts

@@ -81,6 +81,7 @@ import {SessionFactory} from '../src/session-factory';
 import {MultiplexedSession} from '../src/multiplexed-session';
 import {WriteAtLeastOnceOptions} from '../src/database';
 import {MetricsTracerFactory} from '../src/metrics/metrics-tracer-factory';
+import {randomUUID} from 'crypto';
 
 const {
   AlwaysOnSampler,
@@ -3917,6 +3918,124 @@ describe('Spanner with mock server', () => {
         await database.close();
       });
     });
+
+    // TODO: enable when mux session support is available in public methods
+    describe.skip('when multiplexed session is enabled for R/W', () => {
+      before(() => {
+        process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS = 'true';
+        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 select the insertOrUpdate(upsert)/delete(deleteRows) mutation key over insert', async () => {
+        const database = newTestDatabase();
+        await database.runTransactionAsync(async tx => {
+          tx.upsert('foo', [
+            {id: 1, name: 'One'},
+            {id: 2, name: 'Two'},
+          ]);
+          tx.insert('foo', [{id: 3, name: 'Three'}]);
+          tx.insert('foo', [{id: 4, name: 'Four'}]);
+          tx.deleteRows('foo', ['3', '4']);
+          await tx.commit();
+        });
+
+        const beginTransactionRequest = spannerMock
+          .getRequests()
+          .filter(val => {
+            return (val as v1.BeginTransactionRequest).mutationKey;
+          }) as v1.BeginTransactionRequest[];
+
+        // assert on begin transaction request
+        assert.strictEqual(beginTransactionRequest.length, 1);
+
+        // selected mutation key
+        const selectedMutationKey = beginTransactionRequest[0]!.mutationKey;
+
+        // assert that mutation key have been selected
+        assert.ok(
+          selectedMutationKey,
+          'A mutation key should have been selected',
+        );
+
+        // get the type of mutation key
+        const mutationType = Object.keys(selectedMutationKey!)[0];
+
+        // assert that mutation key is not insert
+        assert.notStrictEqual(
+          mutationType,
+          'insert',
+          'The selected mutation key should not be "insert"',
+        );
+
+        // assert that mutation key is either insertOrUpdate or delete
+        assert.ok(
+          ['insertOrUpdate', 'delete'].includes(mutationType),
+          "Expected either 'insertOrUpdate' or 'delete' key.",
+        );
+
+        const commitRequest = spannerMock.getRequests().filter(val => {
+          return (val as v1.CommitRequest).precommitToken;
+        }) as v1.CommitRequest[];
+
+        // assert on commit request
+        assert.strictEqual(commitRequest.length, 1);
+        await database.close();
+      });
+
+      it('should select the mutation key with highest number of values when insert key(s) are present', async () => {
+        const database = newTestDatabase();
+        await database.runTransactionAsync(async tx => {
+          tx.insert('foo', [
+            {id: randomUUID(), name: 'One'},
+            {id: randomUUID(), name: 'Two'},
+            {id: randomUUID(), name: 'Three'},
+          ]);
+          tx.insert('foo', {id: randomUUID(), name: 'Four'});
+          await tx.commit();
+        });
+
+        const beginTransactionRequest = spannerMock
+          .getRequests()
+          .filter(val => {
+            return (val as v1.BeginTransactionRequest).mutationKey;
+          }) as v1.BeginTransactionRequest[];
+
+        // assert on begin transaction request
+        assert.strictEqual(beginTransactionRequest.length, 1);
+
+        // selected mutation key
+        const selectedMutationKey = beginTransactionRequest[0]!.mutationKey;
+
+        // assert that mutation key have been selected
+        assert.ok(
+          selectedMutationKey,
+          'A mutation key should have been selected',
+        );
+
+        // assert that mutation key is insert
+        const mutationType = Object.keys(selectedMutationKey!)[0];
+        assert.ok(
+          ['insert'].includes(mutationType),
+          'insert key must have been selected',
+        );
+
+        // assert that insert mutation key with highest number of rows has been selected
+        assert.strictEqual(selectedMutationKey.insert?.values?.length, 3);
+
+        const commitRequest = spannerMock.getRequests().filter(val => {
+          return (val as v1.CommitRequest).precommitToken;
+        }) as v1.CommitRequest[];
+
+        // assert on commit request
+        assert.strictEqual(commitRequest.length, 1);
+        await database.close();
+      });
+    });
   });
 
   describe('hand-crafted transaction', () => {
@@ -5013,6 +5132,68 @@ describe('Spanner with mock server', () => {
 
       await database.close();
     });
+
+    // TODO: enable when mux session support is available in public methods
+    describe.skip('when multiplexed session is enabled for R/W', () => {
+      before(() => {
+        process.env.GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS = 'true';
+        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 pass the mutation key in begin transaction request in case of mutations only transactions', async () => {
+        const database = newTestDatabase();
+        await database.table('foo').upsert({id: 1, name: randomUUID()});
+        await database.table('foo').insert({id: 2, name: randomUUID()});
+        await database.table('foo').deleteRows(['2']);
+
+        const beginTransactionRequest = spannerMock
+          .getRequests()
+          .filter(val => {
+            return (val as v1.BeginTransactionRequest).mutationKey;
+          }) as v1.BeginTransactionRequest[];
+
+        // assert on begin transaction request
+        assert.strictEqual(beginTransactionRequest.length, 3);
+
+        // assert that on first begin transaction request insertOrUpdate is being selected as mutation key
+        assert.ok(
+          ['insertOrUpdate'].includes(
+            Object.keys(beginTransactionRequest[0]!.mutationKey!)[0],
+          ),
+          'insertOrUpdate key must have been selected',
+        );
+
+        // assert that on second begin transaction request insert is being selected as mutation key
+        assert.ok(
+          ['insert'].includes(
+            Object.keys(beginTransactionRequest[1]!.mutationKey!)[0],
+          ),
+          'insert key must have been selected',
+        );
+
+        // assert that on third begin transaction request delete is being selected as mutation key
+        assert.ok(
+          ['delete'].includes(
+            Object.keys(beginTransactionRequest[2]!.mutationKey!)[0],
+          ),
+          'delete key must have been selected',
+        );
+
+        const commitRequest = spannerMock.getRequests().filter(val => {
+          return (val as v1.CommitRequest).precommitToken;
+        }) as v1.CommitRequest[];
+
+        // assert on commit request
+        assert.strictEqual(commitRequest.length, 3);
+
+        await database.close();
+      });
+    });
   });
 
   describe('chunking', () => {