fix: fix attempt status tag for metrics (#1477)

mutianf · web-flow · commit e54cf7d61128 · 2022-10-25T16:54:11.000-04:00
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/BuiltinMetricsTracer.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/BuiltinMetricsTracer.java
@@ -17,6 +17,7 @@
 
 import static com.google.api.gax.tracing.ApiTracerFactory.OperationType;
 
+import com.google.api.gax.retrying.ServerStreamingAttemptException;
 import com.google.api.gax.tracing.SpanName;
 import com.google.cloud.bigtable.stats.StatsRecorderWrapper;
 import com.google.common.annotations.VisibleForTesting;
@@ -256,6 +257,14 @@ private void recordAttemptCompletion(@Nullable Throwable status) {
         serverLatencyTimerIsRunning = false;
       }
     }
+
+    // Patch the status until it's fixed in gax. When an attempt failed,
+    // it'll throw a ServerStreamingAttemptException. Unwrap the exception
+    // so it could get processed by extractStatus
+    if (status instanceof ServerStreamingAttemptException) {
+      status = status.getCause();
+    }
+
     recorder.putAttemptLatencies(attemptTimer.elapsed(TimeUnit.MILLISECONDS));
     recorder.recordAttempt(Util.extractStatus(status), tableId, zone, cluster);
   }
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/MetricsTracer.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/MetricsTracer.java
@@ -15,6 +15,7 @@
  */
 package com.google.cloud.bigtable.data.v2.stub.metrics;
 
+import com.google.api.gax.retrying.ServerStreamingAttemptException;
 import com.google.api.gax.tracing.ApiTracerFactory.OperationType;
 import com.google.api.gax.tracing.SpanName;
 import com.google.common.base.Stopwatch;
@@ -166,6 +167,13 @@ private void recordAttemptCompletion(@Nullable Throwable throwable) {
                 RpcMeasureConstants.BIGTABLE_ATTEMPT_LATENCY,
                 attemptTimer.elapsed(TimeUnit.MILLISECONDS));
 
+    // Patch the throwable until it's fixed in gax. When an attempt failed,
+    // it'll throw a ServerStreamingAttemptException. Unwrap the exception
+    // so it could get processed by extractStatus
+    if (throwable instanceof ServerStreamingAttemptException) {
+      throwable = throwable.getCause();
+    }
+
     TagContextBuilder tagCtx =
         newTagCtxBuilder()
             .putLocal(
