feat: Add client side metrics for readModifyWriteRow calls

src/client-side-metrics/gcp-metrics-handler.ts

@@ -31,6 +31,7 @@ const {
 } = require('@opentelemetry/sdk-metrics');
 import * as os from 'os';
 import * as crypto from 'crypto';
+import {MethodName} from './client-side-metrics-attributes';
 
 /**
  * Generates a unique client identifier string.
@@ -251,10 +252,16 @@ export class GCPMetricsHandler implements IMetricsHandler {
       status: data.status,
       ...commonAttributes,
     });
-    otelInstruments.firstResponseLatencies.record(data.firstResponseLatency, {
-      status: data.status,
-      ...commonAttributes,
-    });
+    if (
+      data.metricsCollectorData.method === MethodName.READ_ROWS ||
+      data.metricsCollectorData.method === MethodName.READ_ROW
+    ) {
+      otelInstruments.firstResponseLatencies.record(data.firstResponseLatency, {
+        status: data.status,
+        ...commonAttributes,
+      });
+    }
+
     if (data.applicationLatency) {
       otelInstruments.applicationBlockingLatencies.record(
         data.applicationLatency,

src/interceptor.ts → src/client-side-metrics/metric-interceptor.ts

@@ -13,11 +13,11 @@
 // limitations under the License.
 
 import {CallOptions} from 'google-gax';
-import {OperationMetricsCollector} from './client-side-metrics/operation-metrics-collector';
+import {OperationMetricsCollector} from './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: {}};
@@ -49,16 +49,21 @@
             collector.onStatusMetadataReceived(
               status as unknown as ServerStatus,
             );
+            collector.onAttemptComplete(status.code);
             nextStat(status);
           },
         };
         next(metadata, newListener);
       },
+      sendMessage: function (message, next) {
+        collector.onAttemptStart();
+        next(message);
+      },
     });
   };
 }
 
-export function withInterceptors(
+export function withMetricInterceptors(
   gaxOptions: CallOptions,
   metricsCollector?: OperationMetricsCollector,
 ) {

src/client-side-metrics/operation-metrics-collector.ts

@@ -181,10 +181,7 @@ export class OperationMetricsCollector {
         }) => {
           this.onStatusMetadataReceived(status);
         },
-      )
-      .on('data', () => {
-        this.onResponse();
-      });
+      );
   }
 
   /**
@@ -285,11 +282,23 @@ export class OperationMetricsCollector {
    * @param {grpc.status} finalOperationStatus Information about the completed operation.
    * @param {number} applicationLatency The application latency measurement.
    */
-  onOperationComplete(
+  onOperationAndAttemptComplete(
     finalOperationStatus: grpc.status,
     applicationLatency?: number,
   ) {
     this.onAttemptComplete(finalOperationStatus);
+    this.onOperationComplete(finalOperationStatus, applicationLatency);
+  }
+
+  /**
+   * Called when an operation completes (successfully or unsuccessfully).
+   * Records operation latencies, retry counts, and connectivity error counts.
+   * @param {grpc.status} finalOperationStatus Information about the completed operation.
+   */
+  onOperationComplete(
+    finalOperationStatus: grpc.status,
+    applicationLatency?: number,
+  ) {
     withMetricsDebug(() => {
       checkState(this.state, [
         MetricsCollectorState.OPERATION_STARTED_ATTEMPT_NOT_IN_PROGRESS,
@@ -310,7 +319,7 @@ export class OperationMetricsCollector {
                 client_name: `nodejs-bigtable/${version}`,
                 operationLatency: totalMilliseconds,
                 retryCount: this.attemptCount - 1,
-                firstResponseLatency: this.firstResponseLatency ?? undefined,
+                firstResponseLatency: this.firstResponseLatency ?? 0,
                 applicationLatency: applicationLatency ?? 0,
               });
             }

src/row-data-utils.ts

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+import {OperationMetricsCollector} from './client-side-metrics/operation-metrics-collector';
+
 const dotProp = require('dot-prop');
 import {Filter, RawFilter} from './filter';
 import {
@@ -30,6 +32,11 @@ import {TabularApiSurface} from './tabular-api-surface';
 import arrify = require('arrify');
 import {Bigtable} from './index';
 import {CallOptions} from 'google-gax';
+import {
+  MethodName,
+  StreamingState,
+} from './client-side-metrics/client-side-metrics-attributes';
+import {withMetricInterceptors} from './client-side-metrics/metric-interceptor';
 
 interface TabularApiSurfaceRequest {
   tableName?: string;
@@ -193,14 +200,31 @@ class RowDataUtils {
       properties.reqOpts,
     );
     properties.requestData.data = {};
+    // 1. Create a metrics collector.
+    const metricsCollector = new OperationMetricsCollector(
+      properties.requestData.table,
+      MethodName.READ_MODIFY_WRITE_ROW,
+      StreamingState.UNARY,
+      (
+        properties.requestData.table as any
+      ).bigtable._metricsConfigManager!.metricsHandlers,
+    );
+    // 2. Tell the metrics collector an attempt has been started.
+    metricsCollector.onOperationStart();
+    // 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
     properties.requestData.bigtable.request<google.bigtable.v2.IReadModifyWriteRowResponse>(
       {
         client: 'BigtableClient',
         method: 'readModifyWriteRow',
         reqOpts,
-        gaxOpts: gaxOptions,
+        gaxOpts: withMetricInterceptors(gaxOptions, metricsCollector),
+      },
+      (err, ...args) => {
+        metricsCollector.onOperationComplete(err ? err.code : 0);
+        callback(err, ...args);
       },
-      callback,
     );
   }
 

src/tabular-api-surface.ts

@@ -358,7 +358,7 @@ Please use the format 'prezzy' or '${instance.name}/tables/prezzy'.`);
       // communicate a code to the metrics collector.
       //
       const code = originalError ? originalError.code : 0;
-      metricsCollector.onOperationComplete(code);
+      metricsCollector.onOperationAndAttemptComplete(code);
       callback(err, apiResponse);
     };
 

src/utils/createReadStreamInternal.ts

@@ -26,10 +26,6 @@ import {
 } from '../chunktransformer';
 import {TableUtils} from './table';
 import {Duplex, PassThrough, Transform} from 'stream';
-import {
-  MethodName,
-  StreamingState,
-} from '../client-side-metrics/client-side-metrics-attributes';
 import {google} from '../../protos/protos';
 const pumpify = require('pumpify');
 import {grpc, ServiceError} from 'google-gax';
@@ -324,6 +320,10 @@ export function createReadStreamInternal(
       gaxOpts,
       retryOpts,
     });
+    requestStream.on('data', () => {
+      // This handler is necessary for recording firstResponseLatencies.
+      metricsCollector.onResponse();
+    });
 
     activeRequestStream = requestStream!;
 
@@ -363,6 +363,9 @@ export function createReadStreamInternal(
     rowStream = pumpify.obj([requestStream, chunkTransformer, toRowStream]);
 
     metricsCollector.wrapRequest(requestStream);
+    requestStream.on('data', () => {
+      metricsCollector.onResponse();
+    });
     rowStream
       .on('error', (error: ServiceError) => {
         rowStreamUnpipe(rowStream, userStream);
@@ -371,7 +374,7 @@ export function createReadStreamInternal(
           // We ignore the `cancelled` "error", since we are the ones who cause
           // it when the user calls `.abort()`.
           userStream.end();
-          metricsCollector.onOperationComplete(
+          metricsCollector.onOperationAndAttemptComplete(
             error.code,
             userStream.getTotalDurationMs(),
           );
@@ -406,7 +409,7 @@ export function createReadStreamInternal(
             //
             error.code = grpc.status.CANCELLED;
           }
-          metricsCollector.onOperationComplete(
+          metricsCollector.onOperationAndAttemptComplete(
             error.code,
             userStream.getTotalDurationMs(),
           );
@@ -420,8 +423,7 @@ export function createReadStreamInternal(
       })
       .on('end', () => {
         activeRequestStream = null;
-        const applicationLatency = userStream.getTotalDurationMs();
-        metricsCollector.onOperationComplete(
+        metricsCollector.onOperationAndAttemptComplete(
           grpc.status.OK,
           userStream.getTotalDurationMs(),
         );