fix: The metadata could be returned in trailer or header depends on i… (#1337)

mutianf · web-flow · commit c4b8c03ece7b · 2022-08-03T11:52:31.000-04:00
* fix: The metadata could be returned in trailer or header depends on if sidecar is enabled. Check both for now.

* fix

* fix npe

* fix NPE when metadata is null
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/BigtableTracerStreamingCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/BigtableTracerStreamingCallable.java
@@ -104,11 +104,24 @@ public void onError(Throwable t) {
       Long latency = Util.getGfeLatency(metadata);
       tracer.recordGfeMetadata(latency, t);
       try {
-        byte[] trailers =
-            metadata.get(Metadata.Key.of(Util.RESPONSE_PRAMS_KEY, Metadata.BINARY_BYTE_MARSHALLER));
-        ResponseParams decodedTrailers = ResponseParams.parseFrom(trailers);
-        tracer.setLocations(decodedTrailers.getZoneId(), decodedTrailers.getClusterId());
-      } catch (NullPointerException | InvalidProtocolBufferException e) {
+        // Check both headers and trailers because in different environments the metadata
+        // could be returned in headers or trailers
+        if (metadata != null) {
+          byte[] trailers = metadata.get(Util.METADATA_KEY);
+          if (trailers == null) {
+            Metadata trailingMetadata = responseMetadata.getTrailingMetadata();
+            if (trailingMetadata != null) {
+              trailers = trailingMetadata.get(Util.METADATA_KEY);
+            }
+          }
+          // If the response is terminated abnormally and we didn't get location information in
+          // trailers or headers, skip setting the locations
+          if (trailers != null) {
+            ResponseParams decodedTrailers = ResponseParams.parseFrom(trailers);
+            tracer.setLocations(decodedTrailers.getZoneId(), decodedTrailers.getClusterId());
+          }
+        }
+      } catch (InvalidProtocolBufferException e) {
       }
 
       outerObserver.onError(t);
@@ -120,11 +133,24 @@ public void onComplete() {
       Long latency = Util.getGfeLatency(metadata);
       tracer.recordGfeMetadata(latency, null);
       try {
-        byte[] trailers =
-            metadata.get(Metadata.Key.of(Util.RESPONSE_PRAMS_KEY, Metadata.BINARY_BYTE_MARSHALLER));
-        ResponseParams decodedTrailers = ResponseParams.parseFrom(trailers);
-        tracer.setLocations(decodedTrailers.getZoneId(), decodedTrailers.getClusterId());
-      } catch (NullPointerException | InvalidProtocolBufferException e) {
+        // Check both headers and trailers because in different environments the metadata
+        // could be returned in headers or trailers
+        if (metadata != null) {
+          byte[] trailers = metadata.get(Util.METADATA_KEY);
+          if (trailers == null) {
+            Metadata trailingMetadata = responseMetadata.getTrailingMetadata();
+            if (trailingMetadata != null) {
+              trailers = trailingMetadata.get(Util.METADATA_KEY);
+            }
+          }
+          // If the response is terminated abnormally and we didn't get location information in
+          // trailers or headers, skip setting the locations
+          if (trailers != null) {
+            ResponseParams decodedTrailers = ResponseParams.parseFrom(trailers);
+            tracer.setLocations(decodedTrailers.getZoneId(), decodedTrailers.getClusterId());
+          }
+        }
+      } catch (InvalidProtocolBufferException e) {
       }
 
       outerObserver.onComplete();
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/BigtableTracerUnaryCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/BigtableTracerUnaryCallable.java
@@ -83,11 +83,24 @@ public void onFailure(Throwable throwable) {
       Long latency = Util.getGfeLatency(metadata);
       tracer.recordGfeMetadata(latency, throwable);
       try {
-        byte[] trailers =
-            metadata.get(Metadata.Key.of(Util.RESPONSE_PRAMS_KEY, Metadata.BINARY_BYTE_MARSHALLER));
-        ResponseParams decodedTrailers = ResponseParams.parseFrom(trailers);
-        tracer.setLocations(decodedTrailers.getZoneId(), decodedTrailers.getClusterId());
-      } catch (NullPointerException | InvalidProtocolBufferException e) {
+        // Check both headers and trailers because in different environments the metadata
+        // could be returned in headers or trailers
+        if (metadata != null) {
+          byte[] trailers = metadata.get(Util.METADATA_KEY);
+          if (trailers == null) {
+            Metadata trailingMetadata = responseMetadata.getTrailingMetadata();
+            if (trailingMetadata != null) {
+              trailers = trailingMetadata.get(Util.METADATA_KEY);
+            }
+          }
+          // If the response is terminated abnormally and we didn't get location information in
+          // trailers or headers, skip setting the locations
+          if (trailers != null) {
+            ResponseParams decodedTrailers = ResponseParams.parseFrom(trailers);
+            tracer.setLocations(decodedTrailers.getZoneId(), decodedTrailers.getClusterId());
+          }
+        }
+      } catch (InvalidProtocolBufferException e) {
       }
     }
 
@@ -97,11 +110,24 @@ public void onSuccess(ResponseT response) {
       Long latency = Util.getGfeLatency(metadata);
       tracer.recordGfeMetadata(latency, null);
       try {
-        byte[] trailers =
-            metadata.get(Metadata.Key.of(Util.RESPONSE_PRAMS_KEY, Metadata.BINARY_BYTE_MARSHALLER));
-        ResponseParams decodedTrailers = ResponseParams.parseFrom(trailers);
-        tracer.setLocations(decodedTrailers.getZoneId(), decodedTrailers.getClusterId());
-      } catch (NullPointerException | InvalidProtocolBufferException e) {
+        // Check both headers and trailers because in different environments the metadata
+        // could be returned in headers or trailers
+        if (metadata != null) {
+          byte[] trailers = metadata.get(Util.METADATA_KEY);
+          if (trailers == null) {
+            Metadata trailingMetadata = responseMetadata.getTrailingMetadata();
+            if (trailingMetadata != null) {
+              trailers = trailingMetadata.get(Util.METADATA_KEY);
+            }
+          }
+          // If the response is terminated abnormally and we didn't get location information in
+          // trailers or headers, skip setting the locations
+          if (trailers != null) {
+            ResponseParams decodedTrailers = ResponseParams.parseFrom(trailers);
+            tracer.setLocations(decodedTrailers.getZoneId(), decodedTrailers.getClusterId());
+          }
+        }
+      } catch (InvalidProtocolBufferException e) {
       }
     }
   }
diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/Util.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/metrics/Util.java
@@ -57,8 +57,8 @@ public class Util {
   private static final Metadata.Key<String> SERVER_TIMING_HEADER_KEY =
       Metadata.Key.of("server-timing", Metadata.ASCII_STRING_MARSHALLER);
   private static final Pattern SERVER_TIMING_HEADER_PATTERN = Pattern.compile(".*dur=(?<dur>\\d+)");
-
-  static final String RESPONSE_PRAMS_KEY = "x-goog-ext-425905942-bin";
+  static final Metadata.Key<byte[]> METADATA_KEY =
+      Metadata.Key.of("x-goog-ext-425905942-bin", Metadata.BINARY_BYTE_MARSHALLER);
 
   /** Convert an exception into a value that can be used to create an OpenCensus tag value. */
   static String extractStatus(@Nullable Throwable error) {
diff --git a/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/BuiltinMetricsTracerTest.java b/google-cloud-bigtable/src/test/java/com/google/cloud/bigtable/data/v2/stub/metrics/BuiltinMetricsTracerTest.java
@@ -127,9 +127,7 @@ public void sendHeaders(Metadata headers) {
                     ResponseParams params =
                         ResponseParams.newBuilder().setZoneId(ZONE).setClusterId(CLUSTER).build();
                     byte[] byteArray = params.toByteArray();
-                    headers.put(
-                        Metadata.Key.of(Util.RESPONSE_PRAMS_KEY, Metadata.BINARY_BYTE_MARSHALLER),
-                        byteArray);
+                    headers.put(Util.METADATA_KEY, byteArray);
 
                     super.sendHeaders(headers);
                   }
