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 index 6f636bf55d..c7f09c4db1 100644 --- 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 index 0efc99fe4d..50d24ecbaf 100644 --- 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 index 7487703fc0..0e356ebaf9 100644 --- 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 SERVER_TIMING_HEADER_KEY = Metadata.Key.of("server-timing", Metadata.ASCII_STRING_MARSHALLER); private static final Pattern SERVER_TIMING_HEADER_PATTERN = Pattern.compile(".*dur=(?\\d+)"); - - static final String RESPONSE_PRAMS_KEY = "x-goog-ext-425905942-bin"; + static final Metadata.Key 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 index b9bd2a926c..9ea222d012 100644 --- 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); }