Merge branch 'main' into Address-final-comments-on-single-row-mutate-PR

Author: danieljbruce

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 …lient-side-metrics/metric-interceptor.tssrc/interceptor.ts renamed to src/client-side-metrics/metric-interceptor.ts src/interceptor.ts renamed to src/client-side-metrics/metric-interceptor.ts

@@ -13,7 +13,7 @@
 // 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';
@@ -49,16 +49,21 @@ function createMetricsInterceptorProvider(
             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();
-      });
+      );
   }
 
   /**
@@ -289,8 +286,15 @@ export class OperationMetricsCollector {
     finalOperationStatus: grpc.status,
     applicationLatency?: number,
   ) {
-    this.onAttemptComplete(finalOperationStatus);
     withMetricsDebug(() => {
+      if (
+        this.state ===
+          MetricsCollectorState.OPERATION_STARTED_ATTEMPT_IN_PROGRESS_NO_ROWS_YET ||
+        this.state ===
+          MetricsCollectorState.OPERATION_STARTED_ATTEMPT_IN_PROGRESS_SOME_ROWS_RECEIVED
+      ) {
+        this.onAttemptComplete(finalOperationStatus);
+      }
       checkState(this.state, [
         MetricsCollectorState.OPERATION_STARTED_ATTEMPT_NOT_IN_PROGRESS,
       ]);
@@ -310,7 +314,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/utils/createReadStreamInternal.ts

@@ -320,6 +320,10 @@ export function createReadStreamInternal(
       gaxOpts,
       retryOpts,
     });
+    requestStream.on('data', () => {
+      // This handler is necessary for recording firstResponseLatencies.
+      metricsCollector.onResponse();
+    });
 
     activeRequestStream = requestStream!;
 
@@ -416,7 +420,6 @@ export function createReadStreamInternal(
       })
       .on('end', () => {
         activeRequestStream = null;
-        const applicationLatency = userStream.getTotalDurationMs();
         metricsCollector.onOperationComplete(
           grpc.status.OK,
           userStream.getTotalDurationMs(),