From e54cf7d61128049cea2817b409a8524636b4a6cf Mon Sep 17 00:00:00 2001
From: Mattie Fu <mattiefu@google.com>
Date: Tue, 25 Oct 2022 16:54:11 -0400
Subject: [PATCH] fix: fix attempt status tag for metrics (#1477)

---
 .../data/v2/stub/metrics/BuiltinMetricsTracer.java       | 9 +++++++++
 .../bigtable/data/v2/stub/metrics/MetricsTracer.java     | 8 ++++++++
 2 files changed, 17 insertions(+)

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
index 5e29065860..4704f04a4d 100644
--- 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
index 3c63b1b5f7..3b6b1b40ae 100644
--- 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(