diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/XGoogSpannerRequestId.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/XGoogSpannerRequestId.java index 7c2006b64a..3b15f974fe 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/XGoogSpannerRequestId.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/XGoogSpannerRequestId.java @@ -198,4 +198,8 @@ XGoogSpannerRequestId withNthClientId(long replacementClientId) { return XGoogSpannerRequestId.of( replacementClientId, this.nthChannelId, this.nthRequest, this.attempt); } + + public static String getRequestIdFromMetadata(Metadata md) { + return md.get(REQUEST_HEADER_KEY); + } } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/HeaderInterceptor.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/HeaderInterceptor.java index 01aa28d8d9..6b82978e9e 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/HeaderInterceptor.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/HeaderInterceptor.java @@ -28,6 +28,7 @@ import com.google.cloud.spanner.CompositeTracer; import com.google.cloud.spanner.SpannerExceptionFactory; import com.google.cloud.spanner.SpannerRpcMetrics; +import com.google.cloud.spanner.XGoogSpannerRequestId; import com.google.common.cache.Cache; import com.google.common.cache.CacheBuilder; import com.google.spanner.admin.database.v1.DatabaseName; @@ -174,6 +175,10 @@ private void processHeader( } if (span != null) { span.setAttribute("gfe_latency", String.valueOf(gfeLatency)); + String reqId = XGoogSpannerRequestId.getRequestIdFromMetadata(metadata); + if (reqId != null) { + span.setAttribute("x_goog_spanner_request_id", reqId); + } } } else { measureMap.put(SPANNER_GFE_HEADER_MISSING_COUNT, 1L).record(tagContext); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java index fa8b5c982f..a8b2bcf365 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java @@ -1417,8 +1417,7 @@ public void testWriteAtLeastOnceAborted() { List commitRequests = mockSpanner.getRequestsOfType(CommitRequest.class); assertEquals(2, commitRequests.size()); - // TODO(@odeke-em): Enable in later PR. - // xGoogReqIdInterceptor.assertIntegrity(); + xGoogReqIdInterceptor.assertIntegrity(); } @Test @@ -2920,9 +2919,7 @@ public void testPartitionedDmlDoesNotTimeout() { "google.spanner.v1.Spanner/ExecuteStreamingSql", new XGoogSpannerRequestId(NON_DETERMINISTIC, channelId, 6, 1)), }; - if (false) { // TODO(@odeke-em): enable in next PRs. - xGoogReqIdInterceptor.checkExpectedStreamingXGoogRequestIds(wantStreamingValues); - } + xGoogReqIdInterceptor.checkExpectedStreamingXGoogRequestIds(wantStreamingValues); XGoogSpannerRequestIdTest.MethodAndRequestId[] wantUnaryValues = { XGoogSpannerRequestIdTest.ofMethodAndRequestId( @@ -2935,9 +2932,7 @@ public void testPartitionedDmlDoesNotTimeout() { "google.spanner.v1.Spanner/ExecuteSql", new XGoogSpannerRequestId(NON_DETERMINISTIC, channelId, 8, 1)), }; - if (false) { // TODO(@odeke-em): enable in next PRs. - xGoogReqIdInterceptor.checkExpectedUnaryXGoogRequestIdsAsSuffixes(wantUnaryValues); - } + xGoogReqIdInterceptor.checkExpectedUnaryXGoogRequestIdsAsSuffixes(wantUnaryValues); } } @@ -3034,9 +3029,7 @@ public void testPartitionedDmlWithHigherTimeout() { new XGoogSpannerRequestId(NON_DETERMINISTIC, channelId, 6, 1)), }; - if (false) { // TODO(@odeke-em): enable in next PRs. - xGoogReqIdInterceptor.checkExpectedStreamingXGoogRequestIds(wantStreamingValues); - } + xGoogReqIdInterceptor.checkExpectedStreamingXGoogRequestIds(wantStreamingValues); XGoogSpannerRequestIdTest.MethodAndRequestId[] wantUnaryValues = { XGoogSpannerRequestIdTest.ofMethodAndRequestId( @@ -3049,9 +3042,7 @@ public void testPartitionedDmlWithHigherTimeout() { "google.spanner.v1.Spanner/ExecuteSql", new XGoogSpannerRequestId(NON_DETERMINISTIC, channelId, 8, 1)), }; - if (false) { // TODO(@odeke-em): enable in next PRs. - xGoogReqIdInterceptor.checkExpectedUnaryXGoogRequestIdsAsSuffixes(wantUnaryValues); - } + xGoogReqIdInterceptor.checkExpectedUnaryXGoogRequestIdsAsSuffixes(wantUnaryValues); } } @@ -5304,8 +5295,7 @@ public void testSelectHasXGoogRequestIdHeader() { assertEquals(1L, resultSet.getLong(0)); assertFalse(resultSet.next()); } finally { - // TODO(@odeke-em): Enable in later PR. - // xGoogReqIdInterceptor.assertIntegrity(); + xGoogReqIdInterceptor.assertIntegrity(); } } @@ -5394,9 +5384,7 @@ public void testSessionPoolExhaustedError_containsStackTraces() { "google.spanner.v1.Spanner/CreateSession", new XGoogSpannerRequestId(NON_DETERMINISTIC, 0, 1, 1)), }; - if (false) { // TODO(@odeke-em): enable in next PRs. - xGoogReqIdInterceptor.checkExpectedUnaryXGoogRequestIdsAsSuffixes(wantUnaryValues); - } + xGoogReqIdInterceptor.checkExpectedUnaryXGoogRequestIdsAsSuffixes(wantUnaryValues); } } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/OpenTelemetrySpanTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/OpenTelemetrySpanTest.java index 90e66526f1..55681d3570 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/OpenTelemetrySpanTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/OpenTelemetrySpanTest.java @@ -748,6 +748,13 @@ public void testTransactionRunnerWithRetryOnBeginTransaction() { beginTransactionSpan.toString(), beginTransactionSpan.getEvents().stream() .anyMatch(event -> event.getName().equals("Starting RPC retry 1"))); + verifyAtLeast1SpanHasXGoogSpannerRequestIdAttribute(finishedSpans); + } + + private void verifyAtLeast1SpanHasXGoogSpannerRequestIdAttribute(List finishedSpans) { + AttributeKey attributeKey = AttributeKey.stringKey("x_goog_spanner_request_id"); + assertTrue( + finishedSpans.stream().anyMatch(span -> !span.getAttributes().get(attributeKey).isEmpty())); } @Test @@ -798,6 +805,7 @@ public void testSingleUseRetryOnExecuteStreamingSql() { executeStreamingQuery.toString(), executeStreamingQuery.getEvents().stream() .anyMatch(event -> event.getName().contains("Stream broken. Safe to retry"))); + verifyAtLeast1SpanHasXGoogSpannerRequestIdAttribute(finishedSpans); } @Test @@ -845,6 +853,7 @@ public void testRetryOnExecuteSql() { executeSqlSpan.toString(), executeSqlSpan.getEvents().stream() .anyMatch(event -> event.getName().equals("Starting RPC retry 1"))); + verifyAtLeast1SpanHasXGoogSpannerRequestIdAttribute(finishedSpans); } @Test @@ -866,12 +875,14 @@ public void testTableAttributes() { } return null; }); + List finishedSpans = spanExporter.getFinishedSpanItems(); SpanData spanData = - spanExporter.getFinishedSpanItems().stream() + finishedSpans.stream() .filter(x -> x.getName().equals("CloudSpannerOperation.ExecuteStreamingRead")) .findFirst() .get(); verifyTableAttributes(spanData); + verifyAtLeast1SpanHasXGoogSpannerRequestIdAttribute(finishedSpans); } private void waitForFinishedSpans(int numExpectedSpans) { diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/XGoogSpannerRequestIdTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/XGoogSpannerRequestIdTest.java index 719b94593b..26b5885b6a 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/XGoogSpannerRequestIdTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/XGoogSpannerRequestIdTest.java @@ -73,11 +73,9 @@ public static class ServerHeaderEnforcer implements ServerInterceptor { private Set checkMethods; ServerHeaderEnforcer(Set checkMethods) { - this.gotValues = new CopyOnWriteArrayList(); - this.unaryResults = - new ConcurrentHashMap>(); - this.streamingResults = - new ConcurrentHashMap>(); + this.gotValues = new CopyOnWriteArrayList<>(); + this.unaryResults = new ConcurrentHashMap<>(); + this.streamingResults = new ConcurrentHashMap<>(); this.checkMethods = checkMethods; } @@ -128,17 +126,16 @@ public String[] accumulatedValues() { } public void assertIntegrity() { - this.unaryResults.forEach( - (String method, CopyOnWriteArrayList values) -> { - assertMonotonicityOfIds(method, values); - }); - this.streamingResults.forEach( - (String method, CopyOnWriteArrayList values) -> { - assertMonotonicityOfIds(method, values); - }); + this.unaryResults.forEach(this::assertMonotonicityOfIds); + this.streamingResults.forEach(this::assertMonotonicityOfIds); } private void assertMonotonicityOfIds(String prefix, List reqIds) { + reqIds.sort( + (id1, id2) -> { + if (id1.equals(id2)) return 0; + return id1.isGreaterThan(id2) ? 1 : -1; + }); int size = reqIds.size(); List violations = new ArrayList<>(); @@ -161,7 +158,7 @@ private void assertMonotonicityOfIds(String prefix, List } public MethodAndRequestId[] accumulatedUnaryValues() { - List accumulated = new ArrayList(); + List accumulated = new ArrayList<>(); this.unaryResults.forEach( (String method, CopyOnWriteArrayList values) -> { for (int i = 0; i < values.size(); i++) { @@ -172,7 +169,7 @@ public MethodAndRequestId[] accumulatedUnaryValues() { } public MethodAndRequestId[] accumulatedStreamingValues() { - List accumulated = new ArrayList(); + List accumulated = new ArrayList<>(); this.streamingResults.forEach( (String method, CopyOnWriteArrayList values) -> { for (int i = 0; i < values.size(); i++) {