chore: clean up Util to return Status.Code instead of string (#2805)

This in prep for migrating to the strongly typed metrics

Author: igorbernstein2

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/BuiltinMetricsTracer.java

@@ -40,6 +40,7 @@
 import com.google.common.base.Stopwatch;
 import com.google.common.math.IntMath;
 import io.grpc.Deadline;
+import io.grpc.Status;
 import io.opentelemetry.api.common.Attributes;
 import io.opentelemetry.api.metrics.DoubleGauge;
 import io.opentelemetry.api.metrics.DoubleHistogram;
@@ -336,9 +337,9 @@ public void disableFlowControl() {
     flowControlIsDisabled = true;
   }
 
-  private void recordOperationCompletion(@Nullable Throwable status) {
+  private void recordOperationCompletion(@Nullable Throwable throwable) {
     if (operationFinishedEarly.get()) {
-      status = null; // force an ok
+      throwable = null; // force an ok
     }
 
     if (!opFinished.compareAndSet(false, true)) {
@@ -347,7 +348,7 @@ private void recordOperationCompletion(@Nullable Throwable status) {
     long operationLatencyNano = operationTimer.elapsed(TimeUnit.NANOSECONDS);
 
     boolean isStreaming = operationType == OperationType.ServerStreaming;
-    String statusStr = extractStatus(status);
+    Status.Code code = extractStatus(throwable);
 
     // Publish metric data with all the attributes. The attributes get filtered in
     // BuiltinMetricsConstants when we construct the views.
@@ -359,7 +360,7 @@ private void recordOperationCompletion(@Nullable Throwable status) {
             .put(METHOD_KEY, spanName.toString())
             .put(CLIENT_NAME_KEY, NAME)
             .put(STREAMING_KEY, isStreaming)
-            .put(STATUS_KEY, statusStr)
+            .put(STATUS_KEY, code.name())
             .build();
 
     // Only record when retry count is greater than 0 so the retry
@@ -381,9 +382,9 @@ private void recordOperationCompletion(@Nullable Throwable status) {
     }
   }
 
-  private void recordAttemptCompletion(@Nullable Throwable status) {
+  private void recordAttemptCompletion(@Nullable Throwable throwable) {
     if (operationFinishedEarly.get()) {
-      status = null; // force an ok
+      throwable = null; // force an ok
     }
     // If the attempt failed, the time spent in retry should be counted in application latency.
     // Stop the stopwatch and decrement requestLeft.
@@ -397,14 +398,14 @@ private void recordAttemptCompletion(@Nullable Throwable status) {
 
     boolean isStreaming = operationType == OperationType.ServerStreaming;
 
-    // Patch the status until it's fixed in gax. When an attempt failed,
+    // 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 (status instanceof ServerStreamingAttemptException) {
-      status = status.getCause();
+    if (throwable instanceof ServerStreamingAttemptException) {
+      throwable = throwable.getCause();
     }
 
-    String statusStr = extractStatus(status);
+    Status.Code code = extractStatus(throwable);
 
     Attributes attributes =
         baseAttributes.toBuilder()
@@ -414,7 +415,7 @@ private void recordAttemptCompletion(@Nullable Throwable status) {
             .put(METHOD_KEY, spanName.toString())
             .put(CLIENT_NAME_KEY, NAME)
             .put(STREAMING_KEY, isStreaming)
-            .put(STATUS_KEY, statusStr)
+            .put(STATUS_KEY, code.name())
             .build();
 
     totalClientBlockingTime.addAndGet(grpcMessageSentDelay.get());
@@ -477,11 +478,11 @@ public void setBatchWriteFlowControlTargetQps(double targetQps) {
 
   @Override
   public void addBatchWriteFlowControlFactor(
-      double factor, @Nullable Throwable status, boolean applied) {
+      double factor, @Nullable Throwable throwable, boolean applied) {
     Attributes attributes =
         baseAttributes.toBuilder()
             .put(METHOD_KEY, spanName.toString())
-            .put(STATUS_KEY, extractStatus(status))
+            .put(STATUS_KEY, extractStatus(throwable).name())
             .put(APPLIED_KEY, applied)
             .build();
 

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/MetricsTracer.java

@@ -133,7 +133,7 @@ private void recordOperationCompletion(@Nullable Throwable throwable) {
         newTagCtxBuilder()
             .putLocal(
                 RpcMeasureConstants.BIGTABLE_STATUS,
-                TagValue.create(Util.extractStatus(throwable)));
+                TagValue.create(Util.extractStatus(throwable).name()));
 
     measures.record(tagCtx.build());
   }
@@ -216,7 +216,7 @@ private void recordAttemptCompletion(@Nullable Throwable throwable) {
         newTagCtxBuilder()
             .putLocal(
                 RpcMeasureConstants.BIGTABLE_STATUS,
-                TagValue.create(Util.extractStatus(throwable)));
+                TagValue.create(Util.extractStatus(throwable).name()));
 
     measures.record(tagCtx.build());
   }

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/Util.java

@@ -17,10 +17,9 @@
 
 import com.google.api.core.InternalApi;
 import com.google.api.gax.grpc.GaxGrpcProperties;
+import com.google.api.gax.grpc.GrpcStatusCode;
 import com.google.api.gax.rpc.ApiCallContext;
 import com.google.api.gax.rpc.ApiException;
-import com.google.api.gax.rpc.StatusCode;
-import com.google.api.gax.rpc.StatusCode.Code;
 import com.google.api.gax.tracing.ApiTracerFactory;
 import com.google.api.gax.tracing.OpencensusTracerFactory;
 import com.google.auth.Credentials;
@@ -44,8 +43,6 @@
 import com.google.common.collect.ImmutableMap;
 import io.grpc.Metadata;
 import io.grpc.Status;
-import io.grpc.StatusException;
-import io.grpc.StatusRuntimeException;
 import io.opencensus.stats.StatsRecorder;
 import io.opencensus.tags.TagKey;
 import io.opencensus.tags.TagValue;
@@ -79,25 +76,26 @@ public class Util {
   static final Metadata.Key<String> ATTEMPT_EPOCH_KEY =
       Metadata.Key.of("bigtable-client-attempt-epoch-usec", Metadata.ASCII_STRING_MARSHALLER);
 
-  /** Convert an exception into a value that can be used to create an OpenCensus tag value. */
-  public static String extractStatus(@Nullable Throwable error) {
-    final String statusString;
-
+  public static Status.Code extractStatus(@Nullable Throwable error) {
     if (error == null) {
-      return StatusCode.Code.OK.toString();
-    } else if (error instanceof CancellationException) {
-      statusString = Status.Code.CANCELLED.toString();
-    } else if (error instanceof ApiException) {
-      statusString = ((ApiException) error).getStatusCode().getCode().toString();
-    } else if (error instanceof StatusRuntimeException) {
-      statusString = ((StatusRuntimeException) error).getStatus().getCode().toString();
-    } else if (error instanceof StatusException) {
-      statusString = ((StatusException) error).getStatus().getCode().toString();
-    } else {
-      statusString = Code.UNKNOWN.toString();
+      return Status.Code.OK;
+    }
+    // Handle java CancellationException as if it was a gax CancelledException
+    if (error instanceof CancellationException) {
+      return Status.Code.CANCELLED;
+    }
+    if (error instanceof ApiException) {
+      ApiException apiException = (ApiException) error;
+      if (apiException.getStatusCode() instanceof GrpcStatusCode) {
+        return ((GrpcStatusCode) apiException.getStatusCode()).getTransportCode();
+      }
     }
 
-    return statusString;
+    Status s = Status.fromThrowable(error);
+    if (s != null) {
+      return s.getCode();
+    }
+    return Status.Code.UNKNOWN;
   }
 
   static String extractTableId(Object request) {

google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/UtilTest.java

@@ -29,7 +29,7 @@
 public class UtilTest {
   @Test
   public void testOk() {
-    TagValue tagValue = TagValue.create(Util.extractStatus((Throwable) null));
+    TagValue tagValue = TagValue.create(Util.extractStatus(null).name());
     assertThat(tagValue.asString()).isEqualTo("OK");
   }
 
@@ -38,7 +38,7 @@ public void testError() {
     DeadlineExceededException error =
         new DeadlineExceededException(
             "Deadline exceeded", null, GrpcStatusCode.of(Status.Code.DEADLINE_EXCEEDED), true);
-    TagValue tagValue = TagValue.create(Util.extractStatus(error));
+    TagValue tagValue = TagValue.create(Util.extractStatus(error).name());
     assertThat(tagValue.asString()).isEqualTo("DEADLINE_EXCEEDED");
   }
 }