feat: Add plumbing to support unary calls for client side metric collection (#1631)

danieljbruce · web-flow · commit c267ede0140a · 2025-07-23T16:04:18.000Z
## Description

This PR adds source code changes with plumbing so that client side metric collection can be done for unary calls. No client side metrics are actually collected yet for readModifyWriteRow or checkAndMutate calls, but future PRs will use this plumbing to collect metrics for those calls. A test is provided that makes a fake method on the table to demonstrate how a method would collect these metrics to serve as an example.

## Impact

Doesn't change the way the client library works. Just provides support for when we open a PR that collects client side metrics for unary calls.

## Testing

A unit test is created that adds a fake method to a table that does a basic unary readModifyWriteRow call. It ensures that the metrics make it to the test metrics handler correctly.

## Additional Information

**Changes:**

`src/interceptor.ts`: This file contains code for attaching interceptors to a call that will call the correct metrics collector methods when data arrives

`src/client-side-metrics/operation-metrics-collector.ts`: The types are cleaned up in this file. None of the functionality has changed in this file

`system-test/read-modify-write-row-interceptors.ts`: This adds a test that ensures the plumbing for collecting metadata trailers and headers works as intended for unary calls.
diff --git a/src/interceptor.ts b/src/interceptor.ts
@@ -0,0 +1,86 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import {CallOptions} from 'google-gax';
+import {OperationMetricsCollector} from './client-side-metrics/operation-metrics-collector';
+
+// Mock Server Implementation
+import * as grpcJs from '@grpc/grpc-js';
+import {status as GrpcStatus} from '@grpc/grpc-js';
+
+export type ServerStatus = {
+  metadata: {internalRepr: Map<string, Uint8Array[]>; options: {}};
+  code: number;
+  details: string;
+};
+
+// Helper to create interceptor provider for OperationMetricsCollector
+function createMetricsInterceptorProvider(
+  collector: OperationMetricsCollector,
+) {
+  return (options: grpcJs.InterceptorOptions, nextCall: grpcJs.NextCall) => {
+    // savedReceiveMetadata and savedReceiveStatus are not strictly needed here anymore for the interceptor's own state
+    // OperationStart and AttemptStart will be called by the calling code (`fakeReadModifyWriteRow`)
+    return new grpcJs.InterceptingCall(nextCall(options), {
+      start: (metadata, listener, next) => {
+        // AttemptStart is called by the orchestrating code
+        const newListener: grpcJs.Listener = {
+          onReceiveMetadata: (metadata, nextMd) => {
+            collector.onMetadataReceived(
+              metadata as unknown as {
+                internalRepr: Map<string, string[]>;
+                options: {};
+              },
+            );
+            nextMd(metadata);
+          },
+          onReceiveStatus: (status, nextStat) => {
+            collector.onStatusMetadataReceived(
+              status as unknown as ServerStatus,
+            );
+            nextStat(status);
+          },
+        };
+        next(metadata, newListener);
+      },
+    });
+  };
+}
+
+export function withInterceptors(
+  gaxOptions: CallOptions,
+  metricsCollector?: OperationMetricsCollector,
+) {
+  if (metricsCollector) {
+    const interceptor = createMetricsInterceptorProvider(metricsCollector);
+    if (!gaxOptions.otherArgs) {
+      gaxOptions.otherArgs = {};
+    }
+    if (!gaxOptions.otherArgs.options) {
+      gaxOptions.otherArgs.options = {};
+    }
+    if (!gaxOptions.otherArgs.options.interceptors) {
+      gaxOptions.otherArgs.options.interceptors = [interceptor];
+    } else {
+      if (Array.isArray(gaxOptions.otherArgs.options.interceptors)) {
+        // We check that interceptors is an array so that the code has no
+        // chance of throwing an error.
+        // Then, if the interceptors is an array, make sure it also includes the
+        // client side metrics interceptor.
+        gaxOptions.otherArgs.options.interceptors.push(interceptor);
+      }
+    }
+  }
+  return gaxOptions;
+}
diff --git a/system-test/read-modify-write-row-interceptors.ts b/system-test/read-modify-write-row-interceptors.ts
@@ -0,0 +1,294 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import {describe, it, before, after} from 'mocha';
+import {Bigtable} from '../src';
+import {ServiceError} from 'google-gax';
+import {ClientSideMetricsConfigManager} from '../src/client-side-metrics/metrics-config-manager';
+import {TestMetricsHandler} from '../test-common/test-metrics-handler';
+import {
+  OnAttemptCompleteData,
+  OnOperationCompleteData,
+} from '../src/client-side-metrics/metrics-handler';
+import {OperationMetricsCollector} from '../src/client-side-metrics/operation-metrics-collector';
+import {
+  MethodName,
+  StreamingState,
+} from '../src/client-side-metrics/client-side-metrics-attributes';
+import * as assert from 'assert';
+import {status as GrpcStatus} from '@grpc/grpc-js';
+import {withInterceptors} from '../src/interceptor';
+
+const INSTANCE_ID = 'isolated-rmw-instance';
+const TABLE_ID = 'isolated-rmw-table';
+const ZONE = 'us-central1-a';
+const CLUSTER = 'fake-cluster';
+const COLUMN_FAMILY = 'traits';
+const COLUMN_FAMILIES = [COLUMN_FAMILY];
+const ROW_KEY = 'gwashington';
+const COLUMN = 'teeth';
+
+/**
+ * Creates a Bigtable instance if it does not already exist.
+ *
+ * @param bigtable - The Bigtable client.
+ * @param instanceId - The ID of the instance to create.
+ * @param clusterId - The ID of the initial cluster in the instance.
+ * @param locationId - The location (region) for the initial cluster.
+ * @returns The created instance object if successful, otherwise logs a message and returns the existing instance.
+ */
+async function createInstance(
+  bigtable: Bigtable,
+  instanceId: string,
+  clusterId: string,
+  locationId: string,
+) {
+  const instance = bigtable.instance(instanceId);
+
+  const [exists] = await instance.exists();
+  if (exists) {
+    console.log(`Instance ${instanceId} already exists.`);
+    return instance;
+  }
+
+  const [i, operation] = await instance.create({
+    clusters: [
+      {
+        id: clusterId,
+        location: locationId,
+        nodes: 3,
+      },
+    ],
+    labels: {
+      time_created: Date.now(),
+    },
+  });
+  await operation.promise();
+  console.log(`Created instance ${instanceId}`);
+  return i;
+}
+
+/**
+ * Creates a Bigtable table if it does not already exist.
+ *
+ * @param bigtable - The Bigtable client.
+ * @param instanceId - The ID of the instance containing the table.
+ * @param tableId - The ID of the table to create.
+ * @param families - An array of column family names to create in the table.
+ * @returns A promise that resolves with the created Table object.
+ */
+async function createTable(
+  bigtable: Bigtable,
+  instanceId: string,
+  tableId: string,
+  families: string[],
+) {
+  const instance = bigtable.instance(instanceId);
+  const table = instance.table(tableId);
+
+  const [exists] = await table.exists();
+  if (exists) {
+    console.log(`Table ${tableId} already exists.`);
+    return table;
+  }
+
+  const [t] = await table.create({
+    families: families,
+  });
+  const row = table.row(ROW_KEY);
+  await row.save({
+    [COLUMN_FAMILY]: {
+      [COLUMN]: 'shiny',
+    },
+  });
+  console.log(`Created table ${tableId}`);
+  return t;
+}
+
+/**
+ * Creates and returns a TestMetricsHandler instance for testing purposes.
+ *
+ * @returns A TestMetricsHandler instance with the projectId set to 'test-project-id'.
+ */
+function getTestMetricsHandler() {
+  const testMetricsHandler = new TestMetricsHandler();
+  testMetricsHandler.projectId = 'test-project-id';
+  return testMetricsHandler;
+}
+
+/**
+ * Asynchronously retrieves the project ID associated with the Bigtable client.
+ *
+ * @param bigtable - The Bigtable client instance.
+ * @returns A promise that resolves with the project ID as a string.
+ * @throws An error if the project ID cannot be retrieved.
+ */
+async function getProjectIdFromClient(bigtable: Bigtable): Promise<string> {
+  return new Promise((resolve, reject) => {
+    bigtable.getProjectId_((err, projectId) => {
+      if (err) {
+        reject(err);
+      } else {
+        resolve(projectId!);
+      }
+    });
+  });
+}
+
+describe('Bigtable/ReadModifyWriteRowInterceptorMetrics', () => {
+  let bigtable: Bigtable;
+  let testMetricsHandler: TestMetricsHandler;
+
+  before(async () => {
+    bigtable = new Bigtable();
+    await getProjectIdFromClient(bigtable);
+    await createInstance(bigtable, INSTANCE_ID, CLUSTER, ZONE);
+    await createTable(bigtable, INSTANCE_ID, TABLE_ID, COLUMN_FAMILIES);
+    testMetricsHandler = getTestMetricsHandler();
+    bigtable._metricsConfigManager = new ClientSideMetricsConfigManager([
+      testMetricsHandler,
+    ]);
+  });
+
+  after(async () => {
+    const instance = bigtable.instance(INSTANCE_ID);
+    await instance.delete();
+  });
+
+  it('should record and export correct metrics for ReadModifyWriteRow via interceptors', async () => {
+    const instance = bigtable.instance(INSTANCE_ID);
+
+    const table = instance.table(TABLE_ID);
+
+    /*
+    fakeReadModifyWriteRowMethod is just a fake method on a table that makes a
+    call to the readWriteModifyRow grpc endpoint. It demonstrates what a method
+    might look like when trying to make a unary call while extracting
+    information from the headers and trailers that the server returns so that
+    the extracted information can be recorded in client side metrics.
+     */
+    (table as any).fakeReadModifyWriteRowMethod = async () => {
+      // 1. Create a metrics collector.
+      const metricsCollector = new OperationMetricsCollector(
+        table,
+        MethodName.READ_MODIFY_WRITE_ROW,
+        StreamingState.UNARY,
+        (table as any).bigtable._metricsConfigManager!.metricsHandlers,
+      );
+      // 2. Tell the metrics collector an attempt has been started.
+      metricsCollector.onOperationStart();
+      metricsCollector.onAttemptStart();
+      // 3. Make a unary call with gax options that include interceptors. The
+      // interceptors are built from a method that hooks them up to the
+      // metrics collector
+      const responseArray = await new Promise((resolve, reject) => {
+        bigtable.request(
+          {
+            client: 'BigtableClient',
+            method: 'readModifyWriteRow',
+            reqOpts: {
+              tableName: table.name,
+              rowKey: Buffer.from(ROW_KEY),
+              rules: [
+                {
+                  familyName: COLUMN_FAMILY,
+                  columnQualifier: Buffer.from(COLUMN),
+                  appendValue: Buffer.from('-wood'),
+                },
+              ],
+              appProfileId: undefined,
+            },
+            gaxOpts: withInterceptors({}, metricsCollector),
+          },
+          (err: ServiceError | null, resp?: any) => {
+            if (err) {
+              reject(err);
+            } else {
+              resolve(resp);
+            }
+          },
+        );
+      });
+      // 4. Tell the metrics collector the attempt is over
+      metricsCollector.onAttemptComplete(GrpcStatus.OK);
+      metricsCollector.onOperationComplete(GrpcStatus.OK);
+      // 5. Return results of method call to the user
+      return responseArray;
+    };
+
+    await (table as any).fakeReadModifyWriteRowMethod();
+
+    assert.strictEqual(testMetricsHandler.requestsHandled.length, 2);
+
+    const attemptCompleteData = testMetricsHandler.requestsHandled.find(
+      m => (m as {attemptLatency?: number}).attemptLatency !== undefined,
+    ) as OnAttemptCompleteData | undefined;
+    const operationCompleteData = testMetricsHandler.requestsHandled.find(
+      m => (m as {operationLatency?: number}).operationLatency !== undefined,
+    ) as OnOperationCompleteData | undefined;
+
+    assert.ok(attemptCompleteData, 'OnAttemptCompleteData should be present');
+    assert.ok(
+      operationCompleteData,
+      'OnOperationCompleteData should be present',
+    );
+    if (!attemptCompleteData || !operationCompleteData) {
+      throw new Error('Metrics data is missing'); // Should be caught by asserts above
+    }
+    assert.strictEqual(
+      attemptCompleteData.metricsCollectorData.method,
+      MethodName.READ_MODIFY_WRITE_ROW,
+    );
+    assert.strictEqual(attemptCompleteData.status, '0');
+    assert.strictEqual(
+      attemptCompleteData.metricsCollectorData.table,
+      TABLE_ID,
+    );
+    assert.strictEqual(
+      attemptCompleteData.metricsCollectorData.instanceId,
+      INSTANCE_ID,
+    );
+    assert.ok(attemptCompleteData.attemptLatency >= 0);
+    assert(attemptCompleteData.serverLatency);
+    assert.ok(attemptCompleteData.serverLatency >= 0);
+    assert.strictEqual(attemptCompleteData.metricsCollectorData.zone, ZONE);
+    assert.strictEqual(
+      attemptCompleteData.metricsCollectorData.cluster,
+      CLUSTER,
+    );
+    assert.strictEqual(attemptCompleteData.streaming, StreamingState.UNARY);
+
+    assert.strictEqual(
+      operationCompleteData.metricsCollectorData.method,
+      MethodName.READ_MODIFY_WRITE_ROW,
+    );
+    assert.strictEqual(operationCompleteData.status, '0');
+    assert.strictEqual(
+      operationCompleteData.metricsCollectorData.table,
+      TABLE_ID,
+    );
+    assert.strictEqual(
+      operationCompleteData.metricsCollectorData.instanceId,
+      INSTANCE_ID,
+    );
+    assert.ok(operationCompleteData.operationLatency >= 0);
+    assert.strictEqual(operationCompleteData.retryCount, 0);
+    assert.strictEqual(operationCompleteData.metricsCollectorData.zone, ZONE);
+    assert.strictEqual(
+      operationCompleteData.metricsCollectorData.cluster,
+      CLUSTER,
+    );
+    assert.strictEqual(operationCompleteData.streaming, StreamingState.UNARY);
+  });
+});
