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

* 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

Author: mutianf

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();

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) {
  }
  }
  }

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) {

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);
  }