From 6c77f1de96d9e82ba856cc4884f04081d78b25c0 Mon Sep 17 00:00:00 2001 From: jimit-j-shah Date: Fri, 8 Mar 2024 16:18:23 -0800 Subject: [PATCH 01/34] Adding first e2e client-tracing test w/ Custom Root Span --- google-cloud-firestore/pom.xml | 13 + .../cloud/firestore/it/ITE2ETracingTest.java | 242 ++++++++++++++++++ 2 files changed, 255 insertions(+) create mode 100644 google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java diff --git a/google-cloud-firestore/pom.xml b/google-cloud-firestore/pom.xml index 942e38ad5..d29308f87 100644 --- a/google-cloud-firestore/pom.xml +++ b/google-cloud-firestore/pom.xml @@ -105,6 +105,7 @@ com.google.protobuf protobuf-java-util + io.opentelemetry @@ -213,6 +214,18 @@ ${opentelemetry.version} test + + com.google.cloud.opentelemetry + exporter-trace + 0.15.0 + test + + + io.opentelemetry + opentelemetry-exporter-logging + 1.29.0 + test + diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java new file mode 100644 index 000000000..348c150c0 --- /dev/null +++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java @@ -0,0 +1,242 @@ +/* + * Copyright 2024 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.firestore.it; + +import static io.opentelemetry.semconv.resource.attributes.ResourceAttributes.SERVICE_NAME; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.fail; + +import com.google.api.gax.rpc.NotFoundException; +import com.google.cloud.firestore.Firestore; +import com.google.cloud.firestore.FirestoreOpenTelemetryOptions; +import com.google.cloud.firestore.FirestoreOptions; +import com.google.cloud.opentelemetry.trace.TraceConfiguration; +import com.google.cloud.opentelemetry.trace.TraceExporter; +import com.google.cloud.trace.v1.TraceServiceClient; +import com.google.common.base.Preconditions; +import com.google.devtools.cloudtrace.v1.Trace; +import com.google.devtools.cloudtrace.v1.TraceSpan; +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.SpanContext; +import io.opentelemetry.api.trace.TraceFlags; +import io.opentelemetry.api.trace.TraceState; +import io.opentelemetry.api.trace.Tracer; +import io.opentelemetry.context.Context; +import io.opentelemetry.context.Scope; +import io.opentelemetry.sdk.OpenTelemetrySdk; +import io.opentelemetry.sdk.common.CompletableResultCode; +import io.opentelemetry.sdk.resources.Resource; +import io.opentelemetry.sdk.trace.SdkTracerProvider; +import io.opentelemetry.sdk.trace.export.BatchSpanProcessor; +import io.opentelemetry.sdk.trace.samplers.Sampler; +import java.io.IOException; +import java.util.Random; +import java.util.concurrent.TimeUnit; +import java.util.logging.Level; +import java.util.logging.Logger; +import org.junit.After; +import org.junit.AfterClass; +import org.junit.Before; +import org.junit.BeforeClass; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class ITE2ETracingTest { + + private static final Logger logger = + Logger.getLogger(ITBaseTest.class.getName()); + + // TODO(jimit) Generalize project ID, accept from config + private static final String PROJECT_ID = "jimit-test"; + + private static final int NUM_TRACE_ID_BYTES = 32; + + private static final int NUM_SPAN_ID_BYTES = 16; + + private static final int GET_TRACE_RETRY_COUNT = 5; + + private static final int GET_TRACE_RETRY_BACKOFF_MILLIS = 1000; + + private static final int TRACE_FORCE_FLUSH_MILLIS = 1000; + private static final int TRACE_PROVIDER_SHUTDOWN_MILLIS = 1000; + + // Random int generator for trace ID and span ID + private static Random random; + + private static TraceExporter traceExporter; + + // Required for reading back traces from Cloud Trace for validation + private static TraceServiceClient traceClient_v1; + + private static String rootSpanName; + private static Tracer tracer; + + // Required to set custom-root span + private static OpenTelemetrySdk openTelemetrySdk; + + private static Firestore firestore; + + @BeforeClass + public static void setup() throws IOException { + // Set up OTel SDK + Resource resource = + Resource.getDefault().merge(Resource.builder().put(SERVICE_NAME, "Sparky").build()); + + // TODO(jimit) Make it re-usable w/ InMemorySpanExporter + traceExporter = TraceExporter.createWithConfiguration( + TraceConfiguration.builder().setProjectId(PROJECT_ID).build()); + + openTelemetrySdk = OpenTelemetrySdk.builder() + .setTracerProvider( + SdkTracerProvider.builder() + .setResource(resource) + .addSpanProcessor(BatchSpanProcessor.builder(traceExporter).build()) + .setSampler(Sampler.alwaysOn()) + .build()).build(); + + traceClient_v1 = TraceServiceClient.create(); + + // Initialize the Firestore DB w/ the OTel SDK + FirestoreOptions.Builder optionsBuilder = + FirestoreOptions.newBuilder() + .setOpenTelemetryOptions( + FirestoreOpenTelemetryOptions.newBuilder() + .setOpenTelemetry(openTelemetrySdk) + .setTracingEnabled(true).build()); + + String namedDb = System.getProperty("FIRESTORE_NAMED_DATABASE"); + if (namedDb != null) { + logger.log(Level.INFO, "Integration test using named database " + namedDb); + optionsBuilder = optionsBuilder.setDatabaseId(namedDb); + } else { + logger.log(Level.INFO, "Integration test using default database."); + } + firestore = optionsBuilder.build().getService(); + Preconditions.checkNotNull( + firestore, + "Error instantiating Firestore. Check that the service account credentials " + + "were properly set."); + random = new Random(); + } + @Before + public void before() throws Exception { + rootSpanName = + String.format("%s%d", this.getClass().getSimpleName(), System.currentTimeMillis()); + tracer = + firestore.getOptions().getOpenTelemetryOptions().getOpenTelemetry().getTracer(rootSpanName); + } + + @After + public void after() throws Exception { + rootSpanName = null; + tracer = null; + } + + @AfterClass + public static void teardown() { + traceClient_v1.close(); + firestore.shutdown(); + } + + // Generates a random 32-byte hex string + private String generateRandomHexString(int numBytes) { + StringBuffer newTraceId = new StringBuffer(); + while (newTraceId.length() < numBytes) { + newTraceId.append(Integer.toHexString(random.nextInt())); + } + return newTraceId.substring(0, numBytes); + } + + protected String generateNewTraceId() { + return generateRandomHexString(NUM_TRACE_ID_BYTES); + } + + // Generates a random 16-byte hex string + protected String generateNewSpanId() { + return generateRandomHexString(NUM_SPAN_ID_BYTES); + } + + // Cloud Trace indexes traces w/ eventual consistency, even when indexing traceId, therefore the + // test must retry a few times before the trace is available. + protected Trace getTraceWithRetry(String project, String traceId) + throws InterruptedException, RuntimeException { + int retryCount = GET_TRACE_RETRY_COUNT; + + while (retryCount-- > 0) { + try { + Trace t = traceClient_v1.getTrace(project, traceId); + return t; + } catch (NotFoundException notFound) { + logger.warning("Trace not found, retrying in " + GET_TRACE_RETRY_BACKOFF_MILLIS + " ms"); + Thread.sleep(GET_TRACE_RETRY_BACKOFF_MILLIS); + } + } + throw new RuntimeException( + "Trace " + traceId + " for project " + project + " not found in Cloud Trace."); + } + + // Generates a new SpanContext w/ random traceId,spanId + protected SpanContext getNewSpanContext() { + String traceId = generateNewTraceId(); + String spanId = generateNewSpanId(); + logger.info("traceId=" + traceId + ", spanId=" + spanId); + + return SpanContext.create( + traceId, + spanId, + TraceFlags.getSampled(), + TraceState.getDefault()); + } + + protected void waitForTracesToComplete() throws Exception { + CompletableResultCode completableResultCode = + openTelemetrySdk.getSdkTracerProvider().forceFlush(); + completableResultCode.join(TRACE_FORCE_FLUSH_MILLIS, TimeUnit.MILLISECONDS); + + // We need to call `firestore.close()` because that will also close the + // gRPC channel and hence force the gRPC instrumentation library to flush + // its spans. + firestore.close(); + } + + @Test + // Trace an Aggregation.Get request + public void basicTraceTestWithCustomRootSpan() throws Exception { + SpanContext newCtx = getNewSpanContext(); + + // Execute the DB operation in the context of the custom root span. + Span rootSpan = tracer.spanBuilder(rootSpanName) + .setParent(Context.root().with(Span.wrap(newCtx))) + .startSpan(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore.collection("col").count().get().get(); + } finally { + rootSpan.end(); + } + + // Ingest traces to Cloud Trace + waitForTracesToComplete(); + + String traceId = newCtx.getTraceId(); + Trace t = getTraceWithRetry(PROJECT_ID, traceId); + assertEquals(t.getTraceId(), traceId); + assertEquals(t.getSpans(0).getName(),rootSpanName); + assertEquals(t.getSpans(1).getName(), "AggregationQuery.Get"); + } +} From 911a40c26012a1651e481f3097b51848afca2465 Mon Sep 17 00:00:00 2001 From: jimit-j-shah Date: Fri, 8 Mar 2024 16:18:23 -0800 Subject: [PATCH 02/34] test: Adding first e2e client-tracing test w/ Custom Root Span --- google-cloud-firestore/pom.xml | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/google-cloud-firestore/pom.xml b/google-cloud-firestore/pom.xml index d29308f87..ca3ac894c 100644 --- a/google-cloud-firestore/pom.xml +++ b/google-cloud-firestore/pom.xml @@ -226,6 +226,18 @@ 1.29.0 test + + com.google.cloud.opentelemetry + exporter-trace + 0.15.0 + test + + + io.opentelemetry + opentelemetry-exporter-logging + 1.29.0 + test + From 9e925c9c5e08cbb0b1c2a5d0f7f157cf8c2d043f Mon Sep 17 00:00:00 2001 From: jimit-j-shah Date: Mon, 18 Mar 2024 11:55:15 -0700 Subject: [PATCH 03/34] Fixing test dependencies and use default GCP testing project. Fixing --- google-cloud-firestore/pom.xml | 16 ++++++++++++---- .../cloud/firestore/it/ITE2ETracingTest.java | 19 ++++++++++++------- 2 files changed, 24 insertions(+), 11 deletions(-) diff --git a/google-cloud-firestore/pom.xml b/google-cloud-firestore/pom.xml index ca3ac894c..e0d7d5774 100644 --- a/google-cloud-firestore/pom.xml +++ b/google-cloud-firestore/pom.xml @@ -232,13 +232,21 @@ 0.15.0 test + + - io.opentelemetry - opentelemetry-exporter-logging - 1.29.0 + com.google.api.grpc + proto-google-cloud-trace-v1 + 1.3.0 test - + + com.google.cloud + google-cloud-trace + 1.3.0 + test + + diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java index 348c150c0..c6118fa45 100644 --- a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java +++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java @@ -57,19 +57,16 @@ import org.junit.runners.JUnit4; @RunWith(JUnit4.class) -public class ITE2ETracingTest { +public class ITE2ETracingTest extends ITBaseTest { private static final Logger logger = Logger.getLogger(ITBaseTest.class.getName()); - // TODO(jimit) Generalize project ID, accept from config - private static final String PROJECT_ID = "jimit-test"; - private static final int NUM_TRACE_ID_BYTES = 32; private static final int NUM_SPAN_ID_BYTES = 16; - private static final int GET_TRACE_RETRY_COUNT = 5; + private static final int GET_TRACE_RETRY_COUNT = 10; private static final int GET_TRACE_RETRY_BACKOFF_MILLIS = 1000; @@ -90,17 +87,22 @@ public class ITE2ETracingTest { // Required to set custom-root span private static OpenTelemetrySdk openTelemetrySdk; + private static String projectId; + private static Firestore firestore; @BeforeClass public static void setup() throws IOException { + projectId = FirestoreOptions.getDefaultProjectId(); + logger.info("projectId:" + projectId); + // Set up OTel SDK Resource resource = Resource.getDefault().merge(Resource.builder().put(SERVICE_NAME, "Sparky").build()); // TODO(jimit) Make it re-usable w/ InMemorySpanExporter traceExporter = TraceExporter.createWithConfiguration( - TraceConfiguration.builder().setProjectId(PROJECT_ID).build()); + TraceConfiguration.builder().setProjectId(projectId).build()); openTelemetrySdk = OpenTelemetrySdk.builder() .setTracerProvider( @@ -151,6 +153,9 @@ public void after() throws Exception { @AfterClass public static void teardown() { traceClient_v1.close(); + CompletableResultCode completableResultCode = + openTelemetrySdk.getSdkTracerProvider().shutdown(); + completableResultCode.join(TRACE_PROVIDER_SHUTDOWN_MILLIS, TimeUnit.MILLISECONDS); firestore.shutdown(); } @@ -234,7 +239,7 @@ public void basicTraceTestWithCustomRootSpan() throws Exception { waitForTracesToComplete(); String traceId = newCtx.getTraceId(); - Trace t = getTraceWithRetry(PROJECT_ID, traceId); + Trace t = getTraceWithRetry(projectId, traceId); assertEquals(t.getTraceId(), traceId); assertEquals(t.getSpans(0).getName(),rootSpanName); assertEquals(t.getSpans(1).getName(), "AggregationQuery.Get"); From 1edb8267e7882161ca94e97f1ee4e663e0fcd668 Mon Sep 17 00:00:00 2001 From: jimit-j-shah Date: Mon, 18 Mar 2024 11:55:15 -0700 Subject: [PATCH 04/34] Fixing test dependencies and use default GCP testing project. --- google-cloud-firestore/pom.xml | 8 +++++--- .../com/google/cloud/firestore/it/ITE2ETracingTest.java | 3 +-- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/google-cloud-firestore/pom.xml b/google-cloud-firestore/pom.xml index e0d7d5774..becd396d4 100644 --- a/google-cloud-firestore/pom.xml +++ b/google-cloud-firestore/pom.xml @@ -220,10 +220,12 @@ 0.15.0 test + + - io.opentelemetry - opentelemetry-exporter-logging - 1.29.0 + com.google.api.grpc + proto-google-cloud-trace-v1 + 1.3.0 test diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java index c6118fa45..3349f9c33 100644 --- a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java +++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java @@ -18,7 +18,6 @@ import static io.opentelemetry.semconv.resource.attributes.ResourceAttributes.SERVICE_NAME; import static org.junit.Assert.assertEquals; -import static org.junit.Assert.fail; import com.google.api.gax.rpc.NotFoundException; import com.google.cloud.firestore.Firestore; @@ -29,7 +28,6 @@ import com.google.cloud.trace.v1.TraceServiceClient; import com.google.common.base.Preconditions; import com.google.devtools.cloudtrace.v1.Trace; -import com.google.devtools.cloudtrace.v1.TraceSpan; import io.opentelemetry.api.trace.Span; import io.opentelemetry.api.trace.SpanContext; import io.opentelemetry.api.trace.TraceFlags; @@ -71,6 +69,7 @@ public class ITE2ETracingTest extends ITBaseTest { private static final int GET_TRACE_RETRY_BACKOFF_MILLIS = 1000; private static final int TRACE_FORCE_FLUSH_MILLIS = 1000; + private static final int TRACE_PROVIDER_SHUTDOWN_MILLIS = 1000; // Random int generator for trace ID and span ID From a7039882b7c9f029379c7aabb037016ce8d51d51 Mon Sep 17 00:00:00 2001 From: jimit-j-shah Date: Mon, 18 Mar 2024 13:27:56 -0700 Subject: [PATCH 05/34] Fixing formatting --- .../cloud/firestore/it/ITE2ETracingTest.java | 50 ++++++++++--------- 1 file changed, 26 insertions(+), 24 deletions(-) diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java index 3349f9c33..32f3cf936 100644 --- a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java +++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java @@ -57,8 +57,7 @@ @RunWith(JUnit4.class) public class ITE2ETracingTest extends ITBaseTest { - private static final Logger logger = - Logger.getLogger(ITBaseTest.class.getName()); + private static final Logger logger = Logger.getLogger(ITBaseTest.class.getName()); private static final int NUM_TRACE_ID_BYTES = 32; @@ -100,16 +99,19 @@ public static void setup() throws IOException { Resource.getDefault().merge(Resource.builder().put(SERVICE_NAME, "Sparky").build()); // TODO(jimit) Make it re-usable w/ InMemorySpanExporter - traceExporter = TraceExporter.createWithConfiguration( - TraceConfiguration.builder().setProjectId(projectId).build()); - - openTelemetrySdk = OpenTelemetrySdk.builder() - .setTracerProvider( - SdkTracerProvider.builder() - .setResource(resource) - .addSpanProcessor(BatchSpanProcessor.builder(traceExporter).build()) - .setSampler(Sampler.alwaysOn()) - .build()).build(); + traceExporter = + TraceExporter.createWithConfiguration( + TraceConfiguration.builder().setProjectId(projectId).build()); + + openTelemetrySdk = + OpenTelemetrySdk.builder() + .setTracerProvider( + SdkTracerProvider.builder() + .setResource(resource) + .addSpanProcessor(BatchSpanProcessor.builder(traceExporter).build()) + .setSampler(Sampler.alwaysOn()) + .build()) + .build(); traceClient_v1 = TraceServiceClient.create(); @@ -119,7 +121,8 @@ public static void setup() throws IOException { .setOpenTelemetryOptions( FirestoreOpenTelemetryOptions.newBuilder() .setOpenTelemetry(openTelemetrySdk) - .setTracingEnabled(true).build()); + .setTracingEnabled(true) + .build()); String namedDb = System.getProperty("FIRESTORE_NAMED_DATABASE"); if (namedDb != null) { @@ -131,10 +134,11 @@ public static void setup() throws IOException { firestore = optionsBuilder.build().getService(); Preconditions.checkNotNull( firestore, - "Error instantiating Firestore. Check that the service account credentials " + - "were properly set."); + "Error instantiating Firestore. Check that the service account credentials " + + "were properly set."); random = new Random(); } + @Before public void before() throws Exception { rootSpanName = @@ -201,11 +205,7 @@ protected SpanContext getNewSpanContext() { String spanId = generateNewSpanId(); logger.info("traceId=" + traceId + ", spanId=" + spanId); - return SpanContext.create( - traceId, - spanId, - TraceFlags.getSampled(), - TraceState.getDefault()); + return SpanContext.create(traceId, spanId, TraceFlags.getSampled(), TraceState.getDefault()); } protected void waitForTracesToComplete() throws Exception { @@ -225,9 +225,11 @@ public void basicTraceTestWithCustomRootSpan() throws Exception { SpanContext newCtx = getNewSpanContext(); // Execute the DB operation in the context of the custom root span. - Span rootSpan = tracer.spanBuilder(rootSpanName) - .setParent(Context.root().with(Span.wrap(newCtx))) - .startSpan(); + Span rootSpan = + tracer + .spanBuilder(rootSpanName) + .setParent(Context.root().with(Span.wrap(newCtx))) + .startSpan(); try (Scope ss = rootSpan.makeCurrent()) { firestore.collection("col").count().get().get(); } finally { @@ -240,7 +242,7 @@ public void basicTraceTestWithCustomRootSpan() throws Exception { String traceId = newCtx.getTraceId(); Trace t = getTraceWithRetry(projectId, traceId); assertEquals(t.getTraceId(), traceId); - assertEquals(t.getSpans(0).getName(),rootSpanName); + assertEquals(t.getSpans(0).getName(), rootSpanName); assertEquals(t.getSpans(1).getName(), "AggregationQuery.Get"); } } From 5210eb7c1c2cdeb68c3af26c59a506f7614e10ea Mon Sep 17 00:00:00 2001 From: jimit-j-shah Date: Mon, 18 Mar 2024 15:53:27 -0700 Subject: [PATCH 06/34] Add aggregationQueryGet Test --- .../cloud/firestore/it/ITE2ETracingTest.java | 127 ++++++++++++++++-- 1 file changed, 113 insertions(+), 14 deletions(-) diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java index 32f3cf936..b717d91bf 100644 --- a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java +++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java @@ -18,6 +18,7 @@ import static io.opentelemetry.semconv.resource.attributes.ResourceAttributes.SERVICE_NAME; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNull; import com.google.api.gax.rpc.NotFoundException; import com.google.cloud.firestore.Firestore; @@ -42,6 +43,8 @@ import io.opentelemetry.sdk.trace.export.BatchSpanProcessor; import io.opentelemetry.sdk.trace.samplers.Sampler; import java.io.IOException; +import java.util.Arrays; +import java.util.List; import java.util.Random; import java.util.concurrent.TimeUnit; import java.util.logging.Level; @@ -79,6 +82,9 @@ public class ITE2ETracingTest extends ITBaseTest { // Required for reading back traces from Cloud Trace for validation private static TraceServiceClient traceClient_v1; + // Trace received + private static Trace retrievedTrace; + private static String rootSpanName; private static Tracer tracer; @@ -145,12 +151,14 @@ public void before() throws Exception { String.format("%s%d", this.getClass().getSimpleName(), System.currentTimeMillis()); tracer = firestore.getOptions().getOpenTelemetryOptions().getOpenTelemetry().getTracer(rootSpanName); + assertNull(retrievedTrace); } @After public void after() throws Exception { rootSpanName = null; tracer = null; + retrievedTrace = null; } @AfterClass @@ -208,6 +216,13 @@ protected SpanContext getNewSpanContext() { return SpanContext.create(traceId, spanId, TraceFlags.getSampled(), TraceState.getDefault()); } + protected Span getNewRootSpanWithContext(SpanContext spanContext) { + // Execute the DB operation in the context of the custom root span. + return tracer.spanBuilder(rootSpanName) + .setParent(Context.root().with(Span.wrap(spanContext))) + .startSpan(); + } + protected void waitForTracesToComplete() throws Exception { CompletableResultCode completableResultCode = openTelemetrySdk.getSdkTracerProvider().forceFlush(); @@ -219,30 +234,114 @@ protected void waitForTracesToComplete() throws Exception { firestore.close(); } + // Validates `retrievedTrace` + protected void fetchAndValidateTraces(String traceId, String... spanNames) + throws InterruptedException { + // Fetch traces + retrievedTrace = getTraceWithRetry(projectId, traceId); + List spanNameList = Arrays.asList(spanNames); + + // Validate trace spans + assertEquals(retrievedTrace.getTraceId(), traceId); + assertEquals(retrievedTrace.getSpans(0).getName(), rootSpanName); + for (int i = 0; i < spanNameList.size() ; ++i) { + assertEquals(retrievedTrace.getSpans(i+1).getName(), spanNameList.get(i)); + } + } + @Test // Trace an Aggregation.Get request - public void basicTraceTestWithCustomRootSpan() throws Exception { + public void aggregateQueryGet() throws Exception { + // Create custom span context for trace ID injection SpanContext newCtx = getNewSpanContext(); - // Execute the DB operation in the context of the custom root span. - Span rootSpan = - tracer - .spanBuilder(rootSpanName) - .setParent(Context.root().with(Span.wrap(newCtx))) - .startSpan(); + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(newCtx); try (Scope ss = rootSpan.makeCurrent()) { firestore.collection("col").count().get().get(); } finally { rootSpan.end(); } - - // Ingest traces to Cloud Trace waitForTracesToComplete(); - String traceId = newCtx.getTraceId(); - Trace t = getTraceWithRetry(projectId, traceId); - assertEquals(t.getTraceId(), traceId); - assertEquals(t.getSpans(0).getName(), rootSpanName); - assertEquals(t.getSpans(1).getName(), "AggregationQuery.Get"); + // Read and validate traces + fetchAndValidateTraces(newCtx.getTraceId(), "AggregationQuery.Get"); + } + + public void bulkWriterCommit() throws Exception { + } + + @Test + public void partitionQuery() throws Exception { } + + @Test + public void collectionListDocuments() throws Exception { + } + + public void docRefCreate() throws Exception { + } + + @Test + public void docRefCreate2() throws Exception { + } + + @Test + public void docRefSet() throws Exception {} + + @Test + public void docRefSet2() throws Exception {} + + @Test + public void docRefSet3() throws Exception {} + + public void docRefSet4() throws Exception {} + + @Test + public void docRefUpdate() throws Exception {} + + @Test + public void docRefUpdate2() throws Exception {} + + @Test + public void docRefUpdate3() throws Exception {} + + @Test + public void docRefUpdate4() throws Exception {} + + @Test + public void docRefUpdate5() throws Exception {} + + @Test + public void docRefUpdate6() throws Exception {} + + @Test + public void docRefDelete() throws Exception {} + + @Test + public void docRefDelete2() throws Exception {} + + @Test + public void docRefGet() throws Exception {} + + @Test + public void docRefGet2() throws Exception {} + + @Test + public void docListCollections() throws Exception {} + + @Test + public void getAll() throws Exception {} + + @Test + public void queryGet() throws Exception {} + + @Test + public void transaction() throws Exception {} + + @Test + public void transactionRollback() throws Exception {} + + @Test + public void writeBatch() throws Exception {} } From a72ecedf22451ba6685f11d655a7de659f12648a Mon Sep 17 00:00:00 2001 From: jimit-j-shah Date: Mon, 18 Mar 2024 16:08:05 -0700 Subject: [PATCH 07/34] Add bulkWriterCommitTrace Test --- .../cloud/firestore/it/ITE2ETracingTest.java | 45 +++++++++++++++++-- 1 file changed, 42 insertions(+), 3 deletions(-) diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java index b717d91bf..52ee9d352 100644 --- a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java +++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java @@ -16,14 +16,18 @@ package com.google.cloud.firestore.it; +import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_BULK_WRITER_COMMIT; import static io.opentelemetry.semconv.resource.attributes.ResourceAttributes.SERVICE_NAME; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNull; import com.google.api.gax.rpc.NotFoundException; +import com.google.cloud.firestore.BulkWriter; +import com.google.cloud.firestore.BulkWriterOptions; import com.google.cloud.firestore.Firestore; import com.google.cloud.firestore.FirestoreOpenTelemetryOptions; import com.google.cloud.firestore.FirestoreOptions; +import com.google.cloud.firestore.telemetry.TraceUtil; import com.google.cloud.opentelemetry.trace.TraceConfiguration; import com.google.cloud.opentelemetry.trace.TraceExporter; import com.google.cloud.trace.v1.TraceServiceClient; @@ -44,8 +48,11 @@ import io.opentelemetry.sdk.trace.samplers.Sampler; import java.io.IOException; import java.util.Arrays; +import java.util.Collections; import java.util.List; import java.util.Random; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import java.util.logging.Level; import java.util.logging.Logger; @@ -62,6 +69,10 @@ public class ITE2ETracingTest extends ITBaseTest { private static final Logger logger = Logger.getLogger(ITBaseTest.class.getName()); + private static final String SERVICE = "google.firestore.v1.Firestore/"; + + private static final String BATCH_WRITE_RPC_NAME = "BatchWrite"; + private static final int NUM_TRACE_ID_BYTES = 32; private static final int NUM_SPAN_ID_BYTES = 16; @@ -223,6 +234,10 @@ protected Span getNewRootSpanWithContext(SpanContext spanContext) { .startSpan(); } + protected String grpcSpanName(String rpcName) { + return "Sent." + SERVICE + rpcName; + } + protected void waitForTracesToComplete() throws Exception { CompletableResultCode completableResultCode = openTelemetrySdk.getSdkTracerProvider().forceFlush(); @@ -245,13 +260,13 @@ protected void fetchAndValidateTraces(String traceId, String... spanNames) assertEquals(retrievedTrace.getTraceId(), traceId); assertEquals(retrievedTrace.getSpans(0).getName(), rootSpanName); for (int i = 0; i < spanNameList.size() ; ++i) { - assertEquals(retrievedTrace.getSpans(i+1).getName(), spanNameList.get(i)); + assertEquals(spanNameList.get(i), retrievedTrace.getSpans(i+1).getName()); } } @Test // Trace an Aggregation.Get request - public void aggregateQueryGet() throws Exception { + public void aggregateQueryGetTraceTest() throws Exception { // Create custom span context for trace ID injection SpanContext newCtx = getNewSpanContext(); @@ -268,7 +283,31 @@ public void aggregateQueryGet() throws Exception { fetchAndValidateTraces(newCtx.getTraceId(), "AggregationQuery.Get"); } - public void bulkWriterCommit() throws Exception { + @Test + public void bulkWriterCommitTraceTest() throws Exception { + // Create custom span context for trace ID injection + SpanContext newCtx = getNewSpanContext(); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(newCtx); + try (Scope ss = rootSpan.makeCurrent()) { + ScheduledExecutorService bulkWriterExecutor = Executors.newSingleThreadScheduledExecutor(); + BulkWriter bulkWriter = + firestore.bulkWriter(BulkWriterOptions.builder().setExecutor(bulkWriterExecutor).build()); + bulkWriter.set( + firestore.collection("col").document("foo"), + Collections.singletonMap("bulk-foo", "bulk-bar")); + bulkWriter.close(); + bulkWriterExecutor.awaitTermination(100, TimeUnit.MILLISECONDS); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + // Read and validate traces + fetchAndValidateTraces(newCtx.getTraceId(), + SPAN_NAME_BULK_WRITER_COMMIT, + grpcSpanName(BATCH_WRITE_RPC_NAME)); } @Test From bdf4d79b40a59a1762a3017212332c3ec8da0b4c Mon Sep 17 00:00:00 2001 From: jimit-j-shah Date: Mon, 18 Mar 2024 16:25:07 -0700 Subject: [PATCH 08/34] Fixing running multiple-tests --- .../cloud/firestore/it/ITE2ETracingTest.java | 43 +++++++++++-------- 1 file changed, 26 insertions(+), 17 deletions(-) diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java index 52ee9d352..44f28335b 100644 --- a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java +++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java @@ -19,6 +19,7 @@ import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_BULK_WRITER_COMMIT; import static io.opentelemetry.semconv.resource.attributes.ResourceAttributes.SERVICE_NAME; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertNull; import com.google.api.gax.rpc.NotFoundException; @@ -93,7 +94,10 @@ public class ITE2ETracingTest extends ITBaseTest { // Required for reading back traces from Cloud Trace for validation private static TraceServiceClient traceClient_v1; - // Trace received + // Custom SpanContext for each test, required for TraceID injection + private static SpanContext customSpanContext; + + // Trace read back from Cloud Trace using traceClient_v1 for verification private static Trace retrievedTrace; private static String rootSpanName; @@ -162,6 +166,10 @@ public void before() throws Exception { String.format("%s%d", this.getClass().getSimpleName(), System.currentTimeMillis()); tracer = firestore.getOptions().getOpenTelemetryOptions().getOpenTelemetry().getTracer(rootSpanName); + + // Get up a new SpanContext (ergo TraceId) for each test + customSpanContext = getNewSpanContext(); + assertNotNull(customSpanContext); assertNull(retrievedTrace); } @@ -170,14 +178,16 @@ public void after() throws Exception { rootSpanName = null; tracer = null; retrievedTrace = null; + customSpanContext = null; } @AfterClass - public static void teardown() { + public static void teardown() throws Exception { traceClient_v1.close(); CompletableResultCode completableResultCode = openTelemetrySdk.getSdkTracerProvider().shutdown(); completableResultCode.join(TRACE_PROVIDER_SHUTDOWN_MILLIS, TimeUnit.MILLISECONDS); + firestore.close(); firestore.shutdown(); } @@ -227,10 +237,10 @@ protected SpanContext getNewSpanContext() { return SpanContext.create(traceId, spanId, TraceFlags.getSampled(), TraceState.getDefault()); } - protected Span getNewRootSpanWithContext(SpanContext spanContext) { + protected Span getNewRootSpanWithContext() { // Execute the DB operation in the context of the custom root span. return tracer.spanBuilder(rootSpanName) - .setParent(Context.root().with(Span.wrap(spanContext))) + .setParent(Context.root().with(Span.wrap(customSpanContext))) .startSpan(); } @@ -239,14 +249,10 @@ protected String grpcSpanName(String rpcName) { } protected void waitForTracesToComplete() throws Exception { + logger.info("Flushing traces..."); CompletableResultCode completableResultCode = openTelemetrySdk.getSdkTracerProvider().forceFlush(); completableResultCode.join(TRACE_FORCE_FLUSH_MILLIS, TimeUnit.MILLISECONDS); - - // We need to call `firestore.close()` because that will also close the - // gRPC channel and hence force the gRPC instrumentation library to flush - // its spans. - firestore.close(); } // Validates `retrievedTrace` @@ -267,12 +273,13 @@ protected void fetchAndValidateTraces(String traceId, String... spanNames) @Test // Trace an Aggregation.Get request public void aggregateQueryGetTraceTest() throws Exception { - // Create custom span context for trace ID injection - SpanContext newCtx = getNewSpanContext(); + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); // Inject new trace ID - Span rootSpan = getNewRootSpanWithContext(newCtx); + Span rootSpan = getNewRootSpanWithContext(); try (Scope ss = rootSpan.makeCurrent()) { + // Execute the Firestore SDK op firestore.collection("col").count().get().get(); } finally { rootSpan.end(); @@ -280,17 +287,18 @@ public void aggregateQueryGetTraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces(newCtx.getTraceId(), "AggregationQuery.Get"); + fetchAndValidateTraces(customSpanContext.getTraceId(), "AggregationQuery.Get"); } @Test public void bulkWriterCommitTraceTest() throws Exception { - // Create custom span context for trace ID injection - SpanContext newCtx = getNewSpanContext(); + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); // Inject new trace ID - Span rootSpan = getNewRootSpanWithContext(newCtx); + Span rootSpan = getNewRootSpanWithContext(); try (Scope ss = rootSpan.makeCurrent()) { + // Execute the Firestore SDK op ScheduledExecutorService bulkWriterExecutor = Executors.newSingleThreadScheduledExecutor(); BulkWriter bulkWriter = firestore.bulkWriter(BulkWriterOptions.builder().setExecutor(bulkWriterExecutor).build()); @@ -305,13 +313,14 @@ public void bulkWriterCommitTraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces(newCtx.getTraceId(), + fetchAndValidateTraces(customSpanContext.getTraceId(), SPAN_NAME_BULK_WRITER_COMMIT, grpcSpanName(BATCH_WRITE_RPC_NAME)); } @Test public void partitionQuery() throws Exception { + } @Test From 6012b377df765248f97a34dfaacbf770e406e9bd Mon Sep 17 00:00:00 2001 From: jimit-j-shah Date: Mon, 18 Mar 2024 16:30:42 -0700 Subject: [PATCH 09/34] Add partitionQuery Test --- .../cloud/firestore/it/ITE2ETracingTest.java | 22 +++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java index 44f28335b..ac9c22814 100644 --- a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java +++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java @@ -17,6 +17,7 @@ package com.google.cloud.firestore.it; import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_BULK_WRITER_COMMIT; +import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_PARTITION_QUERY; import static io.opentelemetry.semconv.resource.attributes.ResourceAttributes.SERVICE_NAME; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNotNull; @@ -25,6 +26,7 @@ import com.google.api.gax.rpc.NotFoundException; import com.google.cloud.firestore.BulkWriter; import com.google.cloud.firestore.BulkWriterOptions; +import com.google.cloud.firestore.CollectionGroup; import com.google.cloud.firestore.Firestore; import com.google.cloud.firestore.FirestoreOpenTelemetryOptions; import com.google.cloud.firestore.FirestoreOptions; @@ -78,7 +80,7 @@ public class ITE2ETracingTest extends ITBaseTest { private static final int NUM_SPAN_ID_BYTES = 16; - private static final int GET_TRACE_RETRY_COUNT = 10; + private static final int GET_TRACE_RETRY_COUNT = 15; private static final int GET_TRACE_RETRY_BACKOFF_MILLIS = 1000; @@ -319,8 +321,24 @@ public void bulkWriterCommitTraceTest() throws Exception { } @Test - public void partitionQuery() throws Exception { + public void partitionQueryTraceTest() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + CollectionGroup collectionGroup = firestore.collectionGroup("col"); + collectionGroup.getPartitions(3).get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + // Read and validate traces + fetchAndValidateTraces(customSpanContext.getTraceId(), + SPAN_NAME_PARTITION_QUERY, + grpcSpanName(SPAN_NAME_PARTITION_QUERY)); } @Test From da16c9e582fe11deaac17bf3503ab0db972379ab Mon Sep 17 00:00:00 2001 From: jimit-j-shah Date: Mon, 18 Mar 2024 16:33:32 -0700 Subject: [PATCH 10/34] Add collectionListDocumentsTrace Test --- .../cloud/firestore/it/ITE2ETracingTest.java | 20 ++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java index ac9c22814..6f5c1390d 100644 --- a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java +++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java @@ -17,6 +17,7 @@ package com.google.cloud.firestore.it; import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_BULK_WRITER_COMMIT; +import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_COL_REF_LIST_DOCUMENTS; import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_PARTITION_QUERY; import static io.opentelemetry.semconv.resource.attributes.ResourceAttributes.SERVICE_NAME; import static org.junit.Assert.assertEquals; @@ -76,6 +77,8 @@ public class ITE2ETracingTest extends ITBaseTest { private static final String BATCH_WRITE_RPC_NAME = "BatchWrite"; + private static final String LIST_DOCUMENTS_RPC_NAME = "ListDocuments"; + private static final int NUM_TRACE_ID_BYTES = 32; private static final int NUM_SPAN_ID_BYTES = 16; @@ -342,7 +345,22 @@ public void partitionQueryTraceTest() throws Exception { } @Test - public void collectionListDocuments() throws Exception { + public void collectionListDocumentsTraceTest() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore.collection("col").listDocuments(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + // Read and validate traces + fetchAndValidateTraces(customSpanContext.getTraceId(), + SPAN_NAME_COL_REF_LIST_DOCUMENTS, grpcSpanName(LIST_DOCUMENTS_RPC_NAME)); } public void docRefCreate() throws Exception { From d94262110a53dd3ac33f1a359c6c86c591858d88 Mon Sep 17 00:00:00 2001 From: jimit-j-shah Date: Tue, 19 Mar 2024 10:50:50 -0700 Subject: [PATCH 11/34] Add docRef*Trace Tests --- .../cloud/firestore/it/ITE2ETracingTest.java | 122 +++++++++++++++++- 1 file changed, 116 insertions(+), 6 deletions(-) diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java index 6f5c1390d..42a774049 100644 --- a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java +++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java @@ -16,8 +16,11 @@ package com.google.cloud.firestore.it; +import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_BATCH_COMMIT; import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_BULK_WRITER_COMMIT; import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_COL_REF_LIST_DOCUMENTS; +import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_DOC_REF_CREATE; +import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_DOC_REF_SET; import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_PARTITION_QUERY; import static io.opentelemetry.semconv.resource.attributes.ResourceAttributes.SERVICE_NAME; import static org.junit.Assert.assertEquals; @@ -31,6 +34,8 @@ import com.google.cloud.firestore.Firestore; import com.google.cloud.firestore.FirestoreOpenTelemetryOptions; import com.google.cloud.firestore.FirestoreOptions; +import com.google.cloud.firestore.SetOptions; +import com.google.cloud.firestore.it.ITTracingTest.Pojo; import com.google.cloud.firestore.telemetry.TraceUtil; import com.google.cloud.opentelemetry.trace.TraceConfiguration; import com.google.cloud.opentelemetry.trace.TraceExporter; @@ -48,6 +53,7 @@ import io.opentelemetry.sdk.common.CompletableResultCode; import io.opentelemetry.sdk.resources.Resource; import io.opentelemetry.sdk.trace.SdkTracerProvider; +import io.opentelemetry.sdk.trace.data.SpanData; import io.opentelemetry.sdk.trace.export.BatchSpanProcessor; import io.opentelemetry.sdk.trace.samplers.Sampler; import java.io.IOException; @@ -64,6 +70,7 @@ import org.junit.AfterClass; import org.junit.Before; import org.junit.BeforeClass; +import org.junit.Ignore; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -77,6 +84,8 @@ public class ITE2ETracingTest extends ITBaseTest { private static final String BATCH_WRITE_RPC_NAME = "BatchWrite"; + private static final String COMMIT_RPC_NAME = "Commit"; + private static final String LIST_DOCUMENTS_RPC_NAME = "ListDocuments"; private static final int NUM_TRACE_ID_BYTES = 32; @@ -227,6 +236,8 @@ protected Trace getTraceWithRetry(String project, String traceId) } catch (NotFoundException notFound) { logger.warning("Trace not found, retrying in " + GET_TRACE_RETRY_BACKOFF_MILLIS + " ms"); Thread.sleep(GET_TRACE_RETRY_BACKOFF_MILLIS); + } catch (Exception e) { + logger.severe(e.getStackTrace().toString()); } } throw new RuntimeException( @@ -363,23 +374,122 @@ public void collectionListDocumentsTraceTest() throws Exception { SPAN_NAME_COL_REF_LIST_DOCUMENTS, grpcSpanName(LIST_DOCUMENTS_RPC_NAME)); } - public void docRefCreate() throws Exception { + @Test + public void docRefCreateTraceTest() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore.collection("col").document().create(Collections.singletonMap("foo", "bar")).get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + // Read and validate traces + fetchAndValidateTraces(customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_CREATE, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME)); + } + + @Test + public void docRefCreate2TraceTest() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore.collection("col").document().create(new Pojo(1)).get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + // Read and validate traces + fetchAndValidateTraces(customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_CREATE, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME)); } @Test - public void docRefCreate2() throws Exception { + public void docRefSetTraceTest() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore.collection("col").document("foo").set(Collections.singletonMap("foo", "bar")).get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + // Read and validate traces + fetchAndValidateTraces(customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_SET, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME)); } @Test - public void docRefSet() throws Exception {} + public void docRefSet2TraceTest() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore + .collection("col") + .document("foo") + .set(Collections.singletonMap("foo", "bar"), SetOptions.merge()) + .get(); } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + // Read and validate traces + fetchAndValidateTraces(customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_SET, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME)); + } @Test - public void docRefSet2() throws Exception {} + public void docRefSet3TraceTest() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore.collection("col").document("foo").set(new Pojo(1)).get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + // Read and validate traces + fetchAndValidateTraces(customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_SET, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME)); + } @Test - public void docRefSet3() throws Exception {} + public void docRefSet4TraceTest() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); - public void docRefSet4() throws Exception {} + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore.collection("col").document("foo").set(new Pojo(1), SetOptions.merge()).get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + // Read and validate traces + fetchAndValidateTraces(customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_SET, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME)); + } @Test public void docRefUpdate() throws Exception {} From 94e48d05a9e7338568e1aad0ef6842d417f65227 Mon Sep 17 00:00:00 2001 From: jimit-j-shah Date: Tue, 19 Mar 2024 11:00:59 -0700 Subject: [PATCH 12/34] Add docRefUpdate*Trace and docRefDelete*Trace Tests --- .../cloud/firestore/it/ITE2ETracingTest.java | 259 +++++++++++++++--- 1 file changed, 228 insertions(+), 31 deletions(-) diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java index 42a774049..9fa300f11 100644 --- a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java +++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java @@ -20,7 +20,9 @@ import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_BULK_WRITER_COMMIT; import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_COL_REF_LIST_DOCUMENTS; import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_DOC_REF_CREATE; +import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_DOC_REF_DELETE; import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_DOC_REF_SET; +import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_DOC_REF_UPDATE; import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_PARTITION_QUERY; import static io.opentelemetry.semconv.resource.attributes.ResourceAttributes.SERVICE_NAME; import static org.junit.Assert.assertEquals; @@ -31,12 +33,13 @@ import com.google.cloud.firestore.BulkWriter; import com.google.cloud.firestore.BulkWriterOptions; import com.google.cloud.firestore.CollectionGroup; +import com.google.cloud.firestore.FieldPath; import com.google.cloud.firestore.Firestore; import com.google.cloud.firestore.FirestoreOpenTelemetryOptions; import com.google.cloud.firestore.FirestoreOptions; +import com.google.cloud.firestore.Precondition; import com.google.cloud.firestore.SetOptions; import com.google.cloud.firestore.it.ITTracingTest.Pojo; -import com.google.cloud.firestore.telemetry.TraceUtil; import com.google.cloud.opentelemetry.trace.TraceConfiguration; import com.google.cloud.opentelemetry.trace.TraceExporter; import com.google.cloud.trace.v1.TraceServiceClient; @@ -53,7 +56,6 @@ import io.opentelemetry.sdk.common.CompletableResultCode; import io.opentelemetry.sdk.resources.Resource; import io.opentelemetry.sdk.trace.SdkTracerProvider; -import io.opentelemetry.sdk.trace.data.SpanData; import io.opentelemetry.sdk.trace.export.BatchSpanProcessor; import io.opentelemetry.sdk.trace.samplers.Sampler; import java.io.IOException; @@ -70,7 +72,6 @@ import org.junit.AfterClass; import org.junit.Before; import org.junit.BeforeClass; -import org.junit.Ignore; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -255,7 +256,8 @@ protected SpanContext getNewSpanContext() { protected Span getNewRootSpanWithContext() { // Execute the DB operation in the context of the custom root span. - return tracer.spanBuilder(rootSpanName) + return tracer + .spanBuilder(rootSpanName) .setParent(Context.root().with(Span.wrap(customSpanContext))) .startSpan(); } @@ -281,8 +283,8 @@ protected void fetchAndValidateTraces(String traceId, String... spanNames) // Validate trace spans assertEquals(retrievedTrace.getTraceId(), traceId); assertEquals(retrievedTrace.getSpans(0).getName(), rootSpanName); - for (int i = 0; i < spanNameList.size() ; ++i) { - assertEquals(spanNameList.get(i), retrievedTrace.getSpans(i+1).getName()); + for (int i = 0; i < spanNameList.size(); ++i) { + assertEquals(spanNameList.get(i), retrievedTrace.getSpans(i + 1).getName()); } } @@ -329,7 +331,8 @@ public void bulkWriterCommitTraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces(customSpanContext.getTraceId(), + fetchAndValidateTraces( + customSpanContext.getTraceId(), SPAN_NAME_BULK_WRITER_COMMIT, grpcSpanName(BATCH_WRITE_RPC_NAME)); } @@ -350,7 +353,8 @@ public void partitionQueryTraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces(customSpanContext.getTraceId(), + fetchAndValidateTraces( + customSpanContext.getTraceId(), SPAN_NAME_PARTITION_QUERY, grpcSpanName(SPAN_NAME_PARTITION_QUERY)); } @@ -370,8 +374,10 @@ public void collectionListDocumentsTraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces(customSpanContext.getTraceId(), - SPAN_NAME_COL_REF_LIST_DOCUMENTS, grpcSpanName(LIST_DOCUMENTS_RPC_NAME)); + fetchAndValidateTraces( + customSpanContext.getTraceId(), + SPAN_NAME_COL_REF_LIST_DOCUMENTS, + grpcSpanName(LIST_DOCUMENTS_RPC_NAME)); } @Test @@ -389,8 +395,11 @@ public void docRefCreateTraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces(customSpanContext.getTraceId(), - SPAN_NAME_DOC_REF_CREATE, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME)); + fetchAndValidateTraces( + customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_CREATE, + SPAN_NAME_BATCH_COMMIT, + grpcSpanName(COMMIT_RPC_NAME)); } @Test @@ -408,8 +417,11 @@ public void docRefCreate2TraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces(customSpanContext.getTraceId(), - SPAN_NAME_DOC_REF_CREATE, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME)); + fetchAndValidateTraces( + customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_CREATE, + SPAN_NAME_BATCH_COMMIT, + grpcSpanName(COMMIT_RPC_NAME)); } @Test @@ -427,8 +439,11 @@ public void docRefSetTraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces(customSpanContext.getTraceId(), - SPAN_NAME_DOC_REF_SET, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME)); + fetchAndValidateTraces( + customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_SET, + SPAN_NAME_BATCH_COMMIT, + grpcSpanName(COMMIT_RPC_NAME)); } @Test @@ -443,14 +458,18 @@ public void docRefSet2TraceTest() throws Exception { .collection("col") .document("foo") .set(Collections.singletonMap("foo", "bar"), SetOptions.merge()) - .get(); } finally { + .get(); + } finally { rootSpan.end(); } waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces(customSpanContext.getTraceId(), - SPAN_NAME_DOC_REF_SET, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME)); + fetchAndValidateTraces( + customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_SET, + SPAN_NAME_BATCH_COMMIT, + grpcSpanName(COMMIT_RPC_NAME)); } @Test @@ -468,8 +487,11 @@ public void docRefSet3TraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces(customSpanContext.getTraceId(), - SPAN_NAME_DOC_REF_SET, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME)); + fetchAndValidateTraces( + customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_SET, + SPAN_NAME_BATCH_COMMIT, + grpcSpanName(COMMIT_RPC_NAME)); } @Test @@ -487,33 +509,208 @@ public void docRefSet4TraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces(customSpanContext.getTraceId(), - SPAN_NAME_DOC_REF_SET, SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME)); + fetchAndValidateTraces( + customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_SET, + SPAN_NAME_BATCH_COMMIT, + grpcSpanName(COMMIT_RPC_NAME)); } @Test - public void docRefUpdate() throws Exception {} + public void docRefUpdate() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore + .collection("col") + .document("foo") + .update(Collections.singletonMap("foo", "bar")) + .get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + // Read and validate traces + fetchAndValidateTraces( + customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_UPDATE, + SPAN_NAME_BATCH_COMMIT, + grpcSpanName(COMMIT_RPC_NAME)); + } @Test - public void docRefUpdate2() throws Exception {} + public void docRefUpdate2() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore + .collection("col") + .document("foo") + .update(Collections.singletonMap("foo", "bar"), Precondition.NONE) + .get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + // Read and validate traces + fetchAndValidateTraces( + customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_UPDATE, + SPAN_NAME_BATCH_COMMIT, + grpcSpanName(COMMIT_RPC_NAME)); + } @Test - public void docRefUpdate3() throws Exception {} + public void docRefUpdate3() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore.collection("col").document("foo").update("key", "value", "key2", "value2").get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + // Read and validate traces + fetchAndValidateTraces( + customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_UPDATE, + SPAN_NAME_BATCH_COMMIT, + grpcSpanName(COMMIT_RPC_NAME)); + } @Test - public void docRefUpdate4() throws Exception {} + public void docRefUpdate4() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore + .collection("col") + .document("foo") + .update(FieldPath.of("key"), "value", FieldPath.of("key2"), "value2") + .get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + // Read and validate traces + fetchAndValidateTraces( + customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_UPDATE, + SPAN_NAME_BATCH_COMMIT, + grpcSpanName(COMMIT_RPC_NAME)); + } @Test - public void docRefUpdate5() throws Exception {} + public void docRefUpdate5() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore + .collection("col") + .document("foo") + .update(Precondition.NONE, "key", "value", "key2", "value2") + .get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + // Read and validate traces + fetchAndValidateTraces( + customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_UPDATE, + SPAN_NAME_BATCH_COMMIT, + grpcSpanName(COMMIT_RPC_NAME)); + } @Test - public void docRefUpdate6() throws Exception {} + public void docRefUpdate6() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore + .collection("col") + .document("foo") + .update(Precondition.NONE, FieldPath.of("key"), "value", FieldPath.of("key2"), "value2") + .get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + // Read and validate traces + fetchAndValidateTraces( + customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_UPDATE, + SPAN_NAME_BATCH_COMMIT, + grpcSpanName(COMMIT_RPC_NAME)); + } @Test - public void docRefDelete() throws Exception {} + public void docRefDelete() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore.collection("col").document("doc0").delete().get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + // Read and validate traces + fetchAndValidateTraces( + customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_DELETE, + SPAN_NAME_BATCH_COMMIT, + grpcSpanName(COMMIT_RPC_NAME)); + } @Test - public void docRefDelete2() throws Exception {} + public void docRefDelete2() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore.collection("col").document("doc0").delete(Precondition.NONE).get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + // Read and validate traces + fetchAndValidateTraces( + customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_DELETE, + SPAN_NAME_BATCH_COMMIT, + grpcSpanName(COMMIT_RPC_NAME)); + } @Test public void docRefGet() throws Exception {} From 686d34f01cf32514acc1070951d41aef34deb253 Mon Sep 17 00:00:00 2001 From: jimit-j-shah Date: Tue, 19 Mar 2024 14:13:30 -0700 Subject: [PATCH 13/34] Fixing Trace fetching using retries for missing or incomplete traces due to eventual consistency of Cloud Trace --- .../cloud/firestore/it/ITE2ETracingTest.java | 138 ++++++++++++------ 1 file changed, 93 insertions(+), 45 deletions(-) diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java index 9fa300f11..1e84214be 100644 --- a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java +++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java @@ -21,6 +21,7 @@ import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_COL_REF_LIST_DOCUMENTS; import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_DOC_REF_CREATE; import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_DOC_REF_DELETE; +import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_DOC_REF_GET; import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_DOC_REF_SET; import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_DOC_REF_UPDATE; import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_PARTITION_QUERY; @@ -33,6 +34,7 @@ import com.google.cloud.firestore.BulkWriter; import com.google.cloud.firestore.BulkWriterOptions; import com.google.cloud.firestore.CollectionGroup; +import com.google.cloud.firestore.FieldMask; import com.google.cloud.firestore.FieldPath; import com.google.cloud.firestore.Firestore; import com.google.cloud.firestore.FirestoreOpenTelemetryOptions; @@ -59,9 +61,9 @@ import io.opentelemetry.sdk.trace.export.BatchSpanProcessor; import io.opentelemetry.sdk.trace.samplers.Sampler; import java.io.IOException; +import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; -import java.util.List; import java.util.Random; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; @@ -83,12 +85,16 @@ public class ITE2ETracingTest extends ITBaseTest { private static final String SERVICE = "google.firestore.v1.Firestore/"; + private static final String BATCH_GET_DOCUMENTS_RPC_NAME = "BatchGetDocuments"; + private static final String BATCH_WRITE_RPC_NAME = "BatchWrite"; private static final String COMMIT_RPC_NAME = "Commit"; private static final String LIST_DOCUMENTS_RPC_NAME = "ListDocuments"; + private static final String RUN_AGGREGATION_QUERY_RPC_NAME = "RunAggregationQuery"; + private static final int NUM_TRACE_ID_BYTES = 32; private static final int NUM_SPAN_ID_BYTES = 16; @@ -97,7 +103,7 @@ public class ITE2ETracingTest extends ITBaseTest { private static final int GET_TRACE_RETRY_BACKOFF_MILLIS = 1000; - private static final int TRACE_FORCE_FLUSH_MILLIS = 1000; + private static final int TRACE_FORCE_FLUSH_MILLIS = 3000; private static final int TRACE_PROVIDER_SHUTDOWN_MILLIS = 1000; @@ -151,7 +157,14 @@ public static void setup() throws IOException { traceClient_v1 = TraceServiceClient.create(); - // Initialize the Firestore DB w/ the OTel SDK + random = new Random(); + } + + @Before + public void before() throws Exception { + // Initialize the Firestore DB w/ the OTel SDK. Ideally we'd do this is the @BeforeAll method + // but because gRPC traces need to be deterministically force-flushed, firestore.shutdown() + // must be called in @After for each test. FirestoreOptions.Builder optionsBuilder = FirestoreOptions.newBuilder() .setOpenTelemetryOptions( @@ -172,11 +185,6 @@ public static void setup() throws IOException { firestore, "Error instantiating Firestore. Check that the service account credentials " + "were properly set."); - random = new Random(); - } - - @Before - public void before() throws Exception { rootSpanName = String.format("%s%d", this.getClass().getSimpleName(), System.currentTimeMillis()); tracer = @@ -190,6 +198,7 @@ public void before() throws Exception { @After public void after() throws Exception { + firestore.shutdown(); rootSpanName = null; tracer = null; retrievedTrace = null; @@ -202,8 +211,6 @@ public static void teardown() throws Exception { CompletableResultCode completableResultCode = openTelemetrySdk.getSdkTracerProvider().shutdown(); completableResultCode.join(TRACE_PROVIDER_SHUTDOWN_MILLIS, TimeUnit.MILLISECONDS); - firestore.close(); - firestore.shutdown(); } // Generates a random 32-byte hex string @@ -224,27 +231,6 @@ protected String generateNewSpanId() { return generateRandomHexString(NUM_SPAN_ID_BYTES); } - // Cloud Trace indexes traces w/ eventual consistency, even when indexing traceId, therefore the - // test must retry a few times before the trace is available. - protected Trace getTraceWithRetry(String project, String traceId) - throws InterruptedException, RuntimeException { - int retryCount = GET_TRACE_RETRY_COUNT; - - while (retryCount-- > 0) { - try { - Trace t = traceClient_v1.getTrace(project, traceId); - return t; - } catch (NotFoundException notFound) { - logger.warning("Trace not found, retrying in " + GET_TRACE_RETRY_BACKOFF_MILLIS + " ms"); - Thread.sleep(GET_TRACE_RETRY_BACKOFF_MILLIS); - } catch (Exception e) { - logger.severe(e.getStackTrace().toString()); - } - } - throw new RuntimeException( - "Trace " + traceId + " for project " + project + " not found in Cloud Trace."); - } - // Generates a new SpanContext w/ random traceId,spanId protected SpanContext getNewSpanContext() { String traceId = generateNewTraceId(); @@ -273,19 +259,42 @@ protected void waitForTracesToComplete() throws Exception { completableResultCode.join(TRACE_FORCE_FLUSH_MILLIS, TimeUnit.MILLISECONDS); } - // Validates `retrievedTrace` + // Validates `retrievedTrace`. Cloud Trace indexes traces w/ eventual consistency, even when + // indexing traceId, therefore the + // test must retry a few times before the complete trace is available. protected void fetchAndValidateTraces(String traceId, String... spanNames) throws InterruptedException { - // Fetch traces - retrievedTrace = getTraceWithRetry(projectId, traceId); - List spanNameList = Arrays.asList(spanNames); - - // Validate trace spans - assertEquals(retrievedTrace.getTraceId(), traceId); - assertEquals(retrievedTrace.getSpans(0).getName(), rootSpanName); - for (int i = 0; i < spanNameList.size(); ++i) { - assertEquals(spanNameList.get(i), retrievedTrace.getSpans(i + 1).getName()); - } + int numRetries = GET_TRACE_RETRY_COUNT; + do { + try { + // Fetch traces + retrievedTrace = traceClient_v1.getTrace(projectId, traceId); + ArrayList spanNameList = new ArrayList(Arrays.asList(spanNames)); + spanNameList.add(0, rootSpanName); + // Validate trace spans + assertEquals(traceId, retrievedTrace.getTraceId()); + for (int i = 0; i < spanNameList.size(); ++i) { + assertEquals(spanNameList.get(i), retrievedTrace.getSpans(i).getName()); + } + assertEquals(spanNameList.size(), retrievedTrace.getSpansCount()); + return; + } catch (NotFoundException notFound) { + logger.info("Trace not found, retrying in " + GET_TRACE_RETRY_BACKOFF_MILLIS + " ms"); + Thread.sleep(GET_TRACE_RETRY_BACKOFF_MILLIS); + } catch (IndexOutOfBoundsException outOfBoundsException) { + logger.info( + "Expected # spans: " + + (spanNames.length + 1) + + "Actual # spans: " + + retrievedTrace.getSpansCount()); + Thread.sleep(GET_TRACE_RETRY_BACKOFF_MILLIS); + } + } while (numRetries-- > 0); + throw new RuntimeException( + "Expected spans: " + + spanNames.toString() + + ", Actual spans: " + + (retrievedTrace != null ? retrievedTrace.getSpansList().toString() : "null")); } @Test @@ -305,7 +314,10 @@ public void aggregateQueryGetTraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces(customSpanContext.getTraceId(), "AggregationQuery.Get"); + fetchAndValidateTraces( + customSpanContext.getTraceId(), + "AggregationQuery.Get", + grpcSpanName(RUN_AGGREGATION_QUERY_RPC_NAME)); } @Test @@ -713,10 +725,46 @@ public void docRefDelete2() throws Exception { } @Test - public void docRefGet() throws Exception {} + public void docRefGet() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore.collection("col").document("doc0").get().get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + // Read and validate traces + fetchAndValidateTraces( + customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_GET, + grpcSpanName(BATCH_GET_DOCUMENTS_RPC_NAME)); + } @Test - public void docRefGet2() throws Exception {} + public void docRefGet2() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore.collection("col").document("doc0").get(FieldMask.of("foo")).get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + // Read and validate traces + fetchAndValidateTraces( + customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_GET, + grpcSpanName(BATCH_GET_DOCUMENTS_RPC_NAME)); + } @Test public void docListCollections() throws Exception {} From 95452e72cd9d11b32a20d34029ca0b81b3e7ef21 Mon Sep 17 00:00:00 2001 From: jimit-j-shah Date: Tue, 19 Mar 2024 15:49:52 -0700 Subject: [PATCH 14/34] Add get/query Trace Tests --- .../cloud/firestore/it/ITE2ETracingTest.java | 64 ++++++++++++++++++- 1 file changed, 61 insertions(+), 3 deletions(-) diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java index 1e84214be..3431c8833 100644 --- a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java +++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java @@ -22,9 +22,11 @@ import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_DOC_REF_CREATE; import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_DOC_REF_DELETE; import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_DOC_REF_GET; +import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_DOC_REF_LIST_COLLECTIONS; import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_DOC_REF_SET; import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_DOC_REF_UPDATE; import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_PARTITION_QUERY; +import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_QUERY_GET; import static io.opentelemetry.semconv.resource.attributes.ResourceAttributes.SERVICE_NAME; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNotNull; @@ -34,6 +36,7 @@ import com.google.cloud.firestore.BulkWriter; import com.google.cloud.firestore.BulkWriterOptions; import com.google.cloud.firestore.CollectionGroup; +import com.google.cloud.firestore.DocumentReference; import com.google.cloud.firestore.FieldMask; import com.google.cloud.firestore.FieldPath; import com.google.cloud.firestore.Firestore; @@ -91,10 +94,14 @@ public class ITE2ETracingTest extends ITBaseTest { private static final String COMMIT_RPC_NAME = "Commit"; + private static final String LIST_COLLECTIONS_RPC_NAME = "ListCollectionIds"; + private static final String LIST_DOCUMENTS_RPC_NAME = "ListDocuments"; private static final String RUN_AGGREGATION_QUERY_RPC_NAME = "RunAggregationQuery"; + private static final String RUN_QUERY_RPC_NAME = "RunQuery"; + private static final int NUM_TRACE_ID_BYTES = 32; private static final int NUM_SPAN_ID_BYTES = 16; @@ -767,13 +774,64 @@ public void docRefGet2() throws Exception { } @Test - public void docListCollections() throws Exception {} + public void docListCollections() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore.collection("col").document("doc0").listCollections(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + // Read and validate traces + fetchAndValidateTraces( + customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_LIST_COLLECTIONS, + grpcSpanName(LIST_COLLECTIONS_RPC_NAME)); + } @Test - public void getAll() throws Exception {} + public void getAll() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + DocumentReference docRef0 = firestore.collection("col").document(); + DocumentReference docRef1 = firestore.collection("col").document(); + DocumentReference[] docs = {docRef0, docRef1}; + firestore.getAll(docs).get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + fetchAndValidateTraces( + customSpanContext.getTraceId(), grpcSpanName(BATCH_GET_DOCUMENTS_RPC_NAME)); + } @Test - public void queryGet() throws Exception {} + public void queryGet() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore.collection("col").whereEqualTo("foo", "my_non_existent_value").get().get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + fetchAndValidateTraces( + customSpanContext.getTraceId(), SPAN_NAME_QUERY_GET, grpcSpanName(RUN_QUERY_RPC_NAME)); + } @Test public void transaction() throws Exception {} From c11678c46d8782e5ebde50df94eae2a3e1684436 Mon Sep 17 00:00:00 2001 From: jimit-j-shah Date: Wed, 20 Mar 2024 14:29:43 -0700 Subject: [PATCH 15/34] Add Transaction test --- .../cloud/firestore/it/ITE2ETracingTest.java | 48 ++++++++++++++++++- 1 file changed, 47 insertions(+), 1 deletion(-) diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java index 3431c8833..fbf04dd08 100644 --- a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java +++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java @@ -43,6 +43,7 @@ import com.google.cloud.firestore.FirestoreOpenTelemetryOptions; import com.google.cloud.firestore.FirestoreOptions; import com.google.cloud.firestore.Precondition; +import com.google.cloud.firestore.Query; import com.google.cloud.firestore.SetOptions; import com.google.cloud.firestore.it.ITTracingTest.Pojo; import com.google.cloud.opentelemetry.trace.TraceConfiguration; @@ -50,6 +51,7 @@ import com.google.cloud.trace.v1.TraceServiceClient; import com.google.common.base.Preconditions; import com.google.devtools.cloudtrace.v1.Trace; +import com.google.devtools.cloudtrace.v1.TraceSpan; import io.opentelemetry.api.trace.Span; import io.opentelemetry.api.trace.SpanContext; import io.opentelemetry.api.trace.TraceFlags; @@ -67,6 +69,7 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; +import java.util.List; import java.util.Random; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; @@ -834,7 +837,50 @@ public void queryGet() throws Exception { } @Test - public void transaction() throws Exception {} + public void transaction() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore + .runTransaction( + transaction -> { + Query q = firestore.collection("col").whereGreaterThan("bla", ""); + DocumentReference d = firestore.collection("col").document("foo"); + DocumentReference[] docList = {d, d}; + // Document Query. + transaction.get(q).get(); + + // Aggregation Query. + transaction.get(q.count()); + + // Get multiple documents. + transaction.getAll(d, d).get(); + + // Commit 2 documents. + transaction.set( + firestore.collection("foo").document("bar"), + Collections.singletonMap("foo", "bar")); + transaction.set( + firestore.collection("foo").document("bar2"), + Collections.singletonMap("foo2", "bar2")); + return 0; + }) + .get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + Thread.sleep(10000); + Trace trace = traceClient_v1.getTrace(projectId, customSpanContext.getTraceId()); + List spanList = trace.getSpansList(); + for (TraceSpan span : spanList) { + logger.info("spanName: " + span.getName() + ", spanId: " + span.getSpanId() + ", parentSpanId: " + span.getParentSpanId()); + } + } @Test public void transactionRollback() throws Exception {} From c9226cf5c2201ab5b37d1608928f101015c85785 Mon Sep 17 00:00:00 2001 From: jimit-j-shah Date: Tue, 26 Mar 2024 14:13:54 -0700 Subject: [PATCH 16/34] Added TraceContainer to be able to test transaction test-cases --- .../cloud/firestore/it/ITE2ETracingTest.java | 126 +++++++++++++++++- 1 file changed, 125 insertions(+), 1 deletion(-) diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java index fbf04dd08..94605ff59 100644 --- a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java +++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java @@ -29,8 +29,10 @@ import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_QUERY_GET; import static io.opentelemetry.semconv.resource.attributes.ResourceAttributes.SERVICE_NAME; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; import com.google.api.gax.rpc.NotFoundException; import com.google.cloud.firestore.BulkWriter; @@ -70,7 +72,10 @@ import java.util.Arrays; import java.util.Collections; import java.util.List; +import java.util.Map; import java.util.Random; +import java.util.Stack; +import java.util.TreeMap; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; @@ -87,6 +92,74 @@ @RunWith(JUnit4.class) public class ITE2ETracingTest extends ITBaseTest { + // Helper class to track call-stacks in a trace + protected class TraceContainer { + + // Maps Span ID to TraceSpan + private Map idSpanMap; + + // Maps Parent Span ID to a list of Child SpanIDs, useful for top-down traversal + private Map> parentChildIdMap; + + // Tracks the Root Span ID + private long rootId; + + public TraceContainer(String rootSpanName, Trace trace) { + idSpanMap = new TreeMap(); + parentChildIdMap = new TreeMap>(); + for (TraceSpan span : trace.getSpansList()) { + long spanId = span.getSpanId(); + idSpanMap.put(spanId, span); + if (rootSpanName.equals(span.getName())) { + rootId = span.getSpanId(); + } + + // Add self as a child of the parent span + if (!parentChildIdMap.containsKey(span.getParentSpanId())) { + parentChildIdMap.put(span.getParentSpanId(), new ArrayList()); + } + parentChildIdMap.get(span.getParentSpanId()).add(spanId); + } + } + + String spanName(long spanId) { + return idSpanMap.get(spanId).getName(); + } + + List childSpans(long spanId) { + return parentChildIdMap.get(spanId); + } + + boolean containsCallStack(List callStack) throws RuntimeException { + if (callStack == null || callStack.isEmpty()) { + throw new RuntimeException("Input callStack is empty"); + } + // tracks the callStack + int iter = 0; + Stack dfsStack = new Stack(); + dfsStack.push(rootId); + long topId = -1; + while ((iter < callStack.size()) && !dfsStack.isEmpty()) { + topId = dfsStack.pop(); + if (spanName(topId).equals(callStack.get(iter++))) { + List childSpans = childSpans(topId); + if (childSpans != null) { + for (Long id : childSpans) { + dfsStack.push(id); + } + } + } + } + + if (iter == callStack.size()) { + if (childSpans(topId) == null && spanName(topId).equals(callStack.get(iter - 1))) { + return true; + } + } + return false; + } + } + private static final Logger logger = Logger.getLogger(ITBaseTest.class.getName()); private static final String SERVICE = "google.firestore.v1.Firestore/"; @@ -307,6 +380,51 @@ protected void fetchAndValidateTraces(String traceId, String... spanNames) + (retrievedTrace != null ? retrievedTrace.getSpansList().toString() : "null")); } + @Test + public void testTraceContainer() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore.collection("col").whereEqualTo("foo", "my_non_existent_value").get().get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + Thread.sleep(4000); + Trace traceResp = traceClient_v1.getTrace(projectId, customSpanContext.getTraceId()); + TraceContainer traceCont = new TraceContainer(rootSpanName, traceResp); + List callStack = new ArrayList(); + + // Contains exact path + callStack.add(rootSpanName); + callStack.add(SPAN_NAME_QUERY_GET); + callStack.add(grpcSpanName(RUN_QUERY_RPC_NAME)); + assertTrue(traceCont.containsCallStack(callStack)); + + // Top-level mismatch + callStack.clear(); + callStack.add(SPAN_NAME_QUERY_GET); + callStack.add(grpcSpanName(RUN_QUERY_RPC_NAME)); + assertFalse(traceCont.containsCallStack(callStack)); + + // Mid-level match + callStack.clear(); + callStack.add(rootSpanName); + callStack.add(SPAN_NAME_QUERY_GET); + assertFalse(traceCont.containsCallStack(callStack)); + + // Leaf-level mismatch/missing + callStack.clear(); + callStack.add(rootSpanName); + callStack.add(SPAN_NAME_QUERY_GET); + callStack.add(grpcSpanName(RUN_AGGREGATION_QUERY_RPC_NAME)); + assertFalse(traceCont.containsCallStack(callStack)); + } + @Test // Trace an Aggregation.Get request public void aggregateQueryGetTraceTest() throws Exception { @@ -878,7 +996,13 @@ public void transaction() throws Exception { Trace trace = traceClient_v1.getTrace(projectId, customSpanContext.getTraceId()); List spanList = trace.getSpansList(); for (TraceSpan span : spanList) { - logger.info("spanName: " + span.getName() + ", spanId: " + span.getSpanId() + ", parentSpanId: " + span.getParentSpanId()); + logger.info( + "spanName: " + + span.getName() + + ", spanId: " + + span.getSpanId() + + ", parentSpanId: " + + span.getParentSpanId()); } } From 70707a8bef16518fbf6bae5c53c8fc6b195bb9b8 Mon Sep 17 00:00:00 2001 From: jimit-j-shah Date: Wed, 27 Mar 2024 16:21:01 -0700 Subject: [PATCH 17/34] test: Adding Transaction tests --- .../cloud/firestore/it/ITE2ETracingTest.java | 302 ++++++++++++++---- 1 file changed, 247 insertions(+), 55 deletions(-) diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java index 94605ff59..e6079ee57 100644 --- a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java +++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java @@ -27,6 +27,13 @@ import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_DOC_REF_UPDATE; import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_PARTITION_QUERY; import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_QUERY_GET; +import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_BEGIN; +import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_COMMIT; +import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_GET_AGGREGATION_QUERY; +import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_GET_DOCUMENTS; +import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_GET_QUERY; +import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_ROLLBACK; +import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN; import static io.opentelemetry.semconv.resource.attributes.ResourceAttributes.SERVICE_NAME; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; @@ -47,6 +54,7 @@ import com.google.cloud.firestore.Precondition; import com.google.cloud.firestore.Query; import com.google.cloud.firestore.SetOptions; +import com.google.cloud.firestore.WriteBatch; import com.google.cloud.firestore.it.ITTracingTest.Pojo; import com.google.cloud.opentelemetry.trace.TraceConfiguration; import com.google.cloud.opentelemetry.trace.TraceExporter; @@ -74,7 +82,6 @@ import java.util.List; import java.util.Map; import java.util.Random; -import java.util.Stack; import java.util.TreeMap; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; @@ -89,6 +96,27 @@ import org.junit.runner.RunWith; import org.junit.runners.JUnit4; +// This End-to-End test verifies Client-side Tracing Functionality instrumented using the +// OpenTelemetry API. +// The test depends on the following external APIs/Services: +// 1. Java OpenTelemetry SDK +// 2. Cloud Trace Exporter +// 3. TraceServiceClient from Cloud Trace API v1. +// +// Permissions required to run this test (https://cloud.google.com/trace/docs/iam#trace-roles): +// 1. gcloud auth application-default login must be run with the test user. +// 2. To write traces, test user must have one of roles/cloudtrace.[admin|agent|user] roles. +// 3. To read traces, test user must have one of roles/cloudtrace.[admin|user] roles. +// +// Each test-case has the following workflow: +// 1. OpenTelemetry SDK is initialized with Cloud Trace Exporter and 100% Trace Sampling +// 2. On initialization, Firestore client is provided the OpenTelemetry SDK object from (1) +// 3. A custom TraceID is generated and injected using a custom SpanContext +// 4. Firestore operations are run inside a root TraceSpan created using the custom SpanContext from +// (3). +// 5. Traces are read-back using TraceServiceClient and verified against expected Call Stacks. +// TODO In the future it would be great to have a single test-driver for this test and +// ITTracingTest. @RunWith(JUnit4.class) public class ITE2ETracingTest extends ITBaseTest { @@ -130,30 +158,52 @@ List childSpans(long spanId) { return parentChildIdMap.get(spanId); } - boolean containsCallStack(List callStack) throws RuntimeException { - if (callStack == null || callStack.isEmpty()) { + // This method only works for matching call stacks with traces which have children of distinct + // type at all + // levels. This is good enough as the intention is to validate if the e2e path is WAI - the + // intention is not to validate Cloud Trace's correctness w.r.t. durability of all kinds of + // traces. + boolean containsCallStack(String... callStack) throws RuntimeException { + ArrayList expectedCallStack = new ArrayList(); + for (String call : callStack) { + expectedCallStack.add(call); + } + if (expectedCallStack.isEmpty()) { throw new RuntimeException("Input callStack is empty"); } - // tracks the callStack - int iter = 0; - Stack dfsStack = new Stack(); - dfsStack.push(rootId); - long topId = -1; - while ((iter < callStack.size()) && !dfsStack.isEmpty()) { - topId = dfsStack.pop(); - if (spanName(topId).equals(callStack.get(iter++))) { - List childSpans = childSpans(topId); - if (childSpans != null) { - for (Long id : childSpans) { - dfsStack.push(id); + return dfsContainsCallStack(rootId, expectedCallStack); + } + + // Depth-first check for call stack in the trace + private boolean dfsContainsCallStack(long spanId, List expectedCallStack) { + logger.info( + "span=" + + spanName(spanId) + + ", expectedCallStack[0]=" + + (expectedCallStack.isEmpty() ? "null" : expectedCallStack.get(0))); + if (!expectedCallStack.isEmpty() && spanName(spanId).equals(expectedCallStack.get(0))) { + if (childSpans(spanId) == null) { + logger.info("No more chilren for " + spanName(spanId)); + return true; + } else { + for (Long childSpan : childSpans(spanId)) { + int callStackListSize = expectedCallStack.size(); + logger.info( + "childSpan=" + + spanName(childSpan) + + ", expectedCallStackSize=" + + callStackListSize); + if (dfsContainsCallStack( + childSpan, + expectedCallStack.subList( + /*fromIndexInclusive=*/ 1, /*toIndexExclusive*/ callStackListSize))) { + return true; } } } - } - - if (iter == callStack.size()) { - if (childSpans(topId) == null && spanName(topId).equals(callStack.get(iter - 1))) { - return true; + } else { + if (!expectedCallStack.isEmpty()) { + logger.warning(spanName(spanId) + " didn't match " + expectedCallStack.get(0)); } } return false; @@ -168,12 +218,16 @@ boolean containsCallStack(List callStack) throws RuntimeException { private static final String BATCH_WRITE_RPC_NAME = "BatchWrite"; + private static final String BEGIN_TRANSACTION_RPC_NAME = "BeginTransaction"; + private static final String COMMIT_RPC_NAME = "Commit"; private static final String LIST_COLLECTIONS_RPC_NAME = "ListCollectionIds"; private static final String LIST_DOCUMENTS_RPC_NAME = "ListDocuments"; + private static final String ROLLBACK_RPC_NAME = "Rollback"; + private static final String RUN_AGGREGATION_QUERY_RPC_NAME = "RunAggregationQuery"; private static final String RUN_QUERY_RPC_NAME = "RunQuery"; @@ -342,6 +396,57 @@ protected void waitForTracesToComplete() throws Exception { completableResultCode.join(TRACE_FORCE_FLUSH_MILLIS, TimeUnit.MILLISECONDS); } + // Validates `retrievedTrace`. Cloud Trace indexes traces w/ eventual consistency, even when + // indexing traceId, therefore the + // test must retry a few times before the complete trace is available. + protected void fetchAndValidateTransactionTrace( + String traceId, int numExpectedSpans, String... callStack) throws InterruptedException { + // Large enough count to accommodate eventually consistent Cloud Trace backend + int numRetries = GET_TRACE_RETRY_COUNT; + do { + try { + // Fetch traces + retrievedTrace = traceClient_v1.getTrace(projectId, traceId); + assertEquals(traceId, retrievedTrace.getTraceId()); + + ArrayList expectedCallStack = new ArrayList(Arrays.asList(callStack)); + + // numExpectedSpans should account for rootSpanName (which is not passed in callStack) + expectedCallStack.add(0, rootSpanName); + numExpectedSpans++; + + System.out.println( + "expectedSpanCount=" + + numExpectedSpans + + ", retrievedSpanCount=" + + retrievedTrace.getSpansCount()); + // *Maybe* the full trace was returned + if (retrievedTrace.getSpansCount() == numExpectedSpans) { + System.out.println("Checking if TraceContainer containsCallStack"); + TraceContainer traceContainer = new TraceContainer(rootSpanName, retrievedTrace); + String[] temp = new String[expectedCallStack.size()]; + if (traceContainer.containsCallStack(expectedCallStack.toArray(temp))) { + return; + } + logger.severe("CallStack not found in TraceContainer."); + } // else the trace may not have been fully committed to Cloud Trace storage + } catch (NotFoundException notFound) { + logger.info("Trace not found, retrying in " + GET_TRACE_RETRY_BACKOFF_MILLIS + " ms"); + Thread.sleep(GET_TRACE_RETRY_BACKOFF_MILLIS); + // getTrace could fail + } catch (IndexOutOfBoundsException outOfBoundsException) { + logger.info("Call stack not found in trace. Retrying."); + Thread.sleep(GET_TRACE_RETRY_BACKOFF_MILLIS); + } + // Retrieved trace doesn't have expected number of spans + } while (numRetries-- > 0); + throw new RuntimeException( + "Expected spans: " + + callStack.toString() + + ", Actual spans: " + + (retrievedTrace != null ? retrievedTrace.getSpansList().toString() : "null")); + } + // Validates `retrievedTrace`. Cloud Trace indexes traces w/ eventual consistency, even when // indexing traceId, therefore the // test must retry a few times before the complete trace is available. @@ -381,7 +486,7 @@ protected void fetchAndValidateTraces(String traceId, String... spanNames) } @Test - public void testTraceContainer() throws Exception { + public void traceContainerTest() throws Exception { // Make sure the test has a new SpanContext (and TraceId for injection) assertNotNull(customSpanContext); @@ -394,35 +499,40 @@ public void testTraceContainer() throws Exception { } waitForTracesToComplete(); - Thread.sleep(4000); - Trace traceResp = traceClient_v1.getTrace(projectId, customSpanContext.getTraceId()); + Trace traceResp = null; + int expectedSpanCount = 3; + + int numRetries = GET_TRACE_RETRY_COUNT; + do { + try { + traceResp = traceClient_v1.getTrace(projectId, customSpanContext.getTraceId()); + if (traceResp.getSpansCount() == expectedSpanCount) { + break; + } + } catch (NotFoundException notFoundException) { + Thread.sleep(GET_TRACE_RETRY_BACKOFF_MILLIS); + logger.info("Trace not found, retrying in " + GET_TRACE_RETRY_BACKOFF_MILLIS + " ms"); + } + numRetries--; + } while (numRetries > 0); + TraceContainer traceCont = new TraceContainer(rootSpanName, traceResp); - List callStack = new ArrayList(); // Contains exact path - callStack.add(rootSpanName); - callStack.add(SPAN_NAME_QUERY_GET); - callStack.add(grpcSpanName(RUN_QUERY_RPC_NAME)); - assertTrue(traceCont.containsCallStack(callStack)); + assertTrue( + traceCont.containsCallStack( + rootSpanName, SPAN_NAME_QUERY_GET, grpcSpanName(RUN_QUERY_RPC_NAME))); // Top-level mismatch - callStack.clear(); - callStack.add(SPAN_NAME_QUERY_GET); - callStack.add(grpcSpanName(RUN_QUERY_RPC_NAME)); - assertFalse(traceCont.containsCallStack(callStack)); + assertFalse(traceCont.containsCallStack(SPAN_NAME_QUERY_GET, RUN_QUERY_RPC_NAME)); // Mid-level match - callStack.clear(); - callStack.add(rootSpanName); - callStack.add(SPAN_NAME_QUERY_GET); - assertFalse(traceCont.containsCallStack(callStack)); + assertFalse(traceCont.containsCallStack(rootSpanName, SPAN_NAME_QUERY_GET)); // Leaf-level mismatch/missing - callStack.clear(); - callStack.add(rootSpanName); - callStack.add(SPAN_NAME_QUERY_GET); - callStack.add(grpcSpanName(RUN_AGGREGATION_QUERY_RPC_NAME)); - assertFalse(traceCont.containsCallStack(callStack)); + assertFalse( + traceCont.containsCallStack( + rootSpanName, SPAN_NAME_QUERY_GET, RUN_AGGREGATION_QUERY_RPC_NAME)); } @Test @@ -992,23 +1102,105 @@ public void transaction() throws Exception { } waitForTracesToComplete(); - Thread.sleep(10000); - Trace trace = traceClient_v1.getTrace(projectId, customSpanContext.getTraceId()); - List spanList = trace.getSpansList(); - for (TraceSpan span : spanList) { - logger.info( - "spanName: " - + span.getName() - + ", spanId: " - + span.getSpanId() - + ", parentSpanId: " - + span.getParentSpanId()); - } + fetchAndValidateTransactionTrace( + customSpanContext.getTraceId(), + /*numExpectedSpans=*/ 11, + SPAN_NAME_TRANSACTION_RUN, + SPAN_NAME_TRANSACTION_BEGIN, + grpcSpanName(BEGIN_TRANSACTION_RPC_NAME)); + + fetchAndValidateTransactionTrace( + customSpanContext.getTraceId(), + /*numExpectedSpans=*/ 11, + SPAN_NAME_TRANSACTION_RUN, + SPAN_NAME_TRANSACTION_GET_QUERY, + grpcSpanName(RUN_QUERY_RPC_NAME)); + + fetchAndValidateTransactionTrace( + customSpanContext.getTraceId(), + /*numExpectedSpans=*/ 11, + SPAN_NAME_TRANSACTION_RUN, + SPAN_NAME_TRANSACTION_GET_AGGREGATION_QUERY, + grpcSpanName(RUN_AGGREGATION_QUERY_RPC_NAME)); + + fetchAndValidateTransactionTrace( + customSpanContext.getTraceId(), + /*numExpectedSpans=*/ 11, + SPAN_NAME_TRANSACTION_RUN, + SPAN_NAME_TRANSACTION_GET_DOCUMENTS, + grpcSpanName(BATCH_GET_DOCUMENTS_RPC_NAME)); + + fetchAndValidateTransactionTrace( + customSpanContext.getTraceId(), + /*numExpectedSpans=*/ 11, + SPAN_NAME_TRANSACTION_RUN, + SPAN_NAME_TRANSACTION_COMMIT, + grpcSpanName(COMMIT_RPC_NAME)); } @Test - public void transactionRollback() throws Exception {} + public void transactionRollback() throws Exception { + String myErrorMessage = "My error message."; + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore + .runTransaction( + transaction -> { + if (true) { + throw (new Exception(myErrorMessage)); + } + return 0; + }) + .get(); + } catch (Exception e) { + // Catch and move on. + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + fetchAndValidateTransactionTrace( + customSpanContext.getTraceId(), + /*numExpectedSpans=*/ 5, + SPAN_NAME_TRANSACTION_RUN, + SPAN_NAME_TRANSACTION_BEGIN, + grpcSpanName(BEGIN_TRANSACTION_RPC_NAME)); + + fetchAndValidateTransactionTrace( + customSpanContext.getTraceId(), + /*numExpectedSpans=*/ 5, + SPAN_NAME_TRANSACTION_RUN, + SPAN_NAME_TRANSACTION_ROLLBACK, + grpcSpanName(ROLLBACK_RPC_NAME)); + } @Test - public void writeBatch() throws Exception {} + public void writeBatch() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + WriteBatch batch = firestore.batch(); + DocumentReference docRef = firestore.collection("foo").document(); + batch.create(docRef, Collections.singletonMap("foo", "bar")); + batch.update(docRef, Collections.singletonMap("foo", "bar")); + batch.delete(docRef); + batch.commit().get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + fetchAndValidateTransactionTrace( + customSpanContext.getTraceId(), + /*numExpectedSpans=*/ 2, + SPAN_NAME_BATCH_COMMIT, + grpcSpanName(COMMIT_RPC_NAME)); + } } From 04a4d29e17d2ea358a3a331956a91a1eb5636ebb Mon Sep 17 00:00:00 2001 From: jimit-j-shah Date: Wed, 20 Mar 2024 14:29:43 -0700 Subject: [PATCH 18/34] test: Adding Transaction tests From cfd1b5f62facb6267feb00acf67bcf2f5f85bef2 Mon Sep 17 00:00:00 2001 From: jimit-j-shah Date: Thu, 28 Mar 2024 10:57:27 -0700 Subject: [PATCH 19/34] test: Adding TestParameterInjector to run the test for global and non-global opentelemetry SDK instances --- google-cloud-firestore/pom.xml | 6 ++ .../cloud/firestore/it/ITE2ETracingTest.java | 91 ++++++++++++++----- 2 files changed, 72 insertions(+), 25 deletions(-) diff --git a/google-cloud-firestore/pom.xml b/google-cloud-firestore/pom.xml index becd396d4..742b2ee6a 100644 --- a/google-cloud-firestore/pom.xml +++ b/google-cloud-firestore/pom.xml @@ -248,6 +248,12 @@ 1.3.0 test + + com.google.testparameterinjector + test-parameter-injector + 1.15 + test + diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java index e6079ee57..a5b463288 100644 --- a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java +++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java @@ -62,6 +62,9 @@ import com.google.common.base.Preconditions; import com.google.devtools.cloudtrace.v1.Trace; import com.google.devtools.cloudtrace.v1.TraceSpan; +import com.google.testing.junit.testparameterinjector.TestParameter; +import com.google.testing.junit.testparameterinjector.TestParameterInjector; +import io.opentelemetry.api.GlobalOpenTelemetry; import io.opentelemetry.api.trace.Span; import io.opentelemetry.api.trace.SpanContext; import io.opentelemetry.api.trace.TraceFlags; @@ -94,7 +97,6 @@ import org.junit.BeforeClass; import org.junit.Test; import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; // This End-to-End test verifies Client-side Tracing Functionality instrumented using the // OpenTelemetry API. @@ -117,9 +119,13 @@ // 5. Traces are read-back using TraceServiceClient and verified against expected Call Stacks. // TODO In the future it would be great to have a single test-driver for this test and // ITTracingTest. -@RunWith(JUnit4.class) +@RunWith(TestParameterInjector.class) public class ITE2ETracingTest extends ITBaseTest { + protected boolean isUsingGlobalOpenTelemetrySDK() { + return useGlobalOpenTelemetrySDK; + } + // Helper class to track call-stacks in a trace protected class TraceContainer { @@ -268,30 +274,18 @@ private boolean dfsContainsCallStack(long spanId, List expectedCallStack private static Firestore firestore; + @TestParameter boolean useGlobalOpenTelemetrySDK; + @BeforeClass public static void setup() throws IOException { projectId = FirestoreOptions.getDefaultProjectId(); logger.info("projectId:" + projectId); - // Set up OTel SDK - Resource resource = - Resource.getDefault().merge(Resource.builder().put(SERVICE_NAME, "Sparky").build()); - // TODO(jimit) Make it re-usable w/ InMemorySpanExporter traceExporter = TraceExporter.createWithConfiguration( TraceConfiguration.builder().setProjectId(projectId).build()); - openTelemetrySdk = - OpenTelemetrySdk.builder() - .setTracerProvider( - SdkTracerProvider.builder() - .setResource(resource) - .addSpanProcessor(BatchSpanProcessor.builder(traceExporter).build()) - .setSampler(Sampler.alwaysOn()) - .build()) - .build(); - traceClient_v1 = TraceServiceClient.create(); random = new Random(); @@ -299,16 +293,50 @@ public static void setup() throws IOException { @Before public void before() throws Exception { + // Set up OTel SDK + Resource resource = + Resource.getDefault().merge(Resource.builder().put(SERVICE_NAME, "Sparky").build()); + + if (isUsingGlobalOpenTelemetrySDK()) { + openTelemetrySdk = + OpenTelemetrySdk.builder() + .setTracerProvider( + SdkTracerProvider.builder() + .setResource(resource) + .addSpanProcessor(BatchSpanProcessor.builder(traceExporter).build()) + .setSampler(Sampler.alwaysOn()) + .build()) + .buildAndRegisterGlobal(); + } else { + openTelemetrySdk = + OpenTelemetrySdk.builder() + .setTracerProvider( + SdkTracerProvider.builder() + .setResource(resource) + .addSpanProcessor(BatchSpanProcessor.builder(traceExporter).build()) + .setSampler(Sampler.alwaysOn()) + .build()) + .build(); + } + // Initialize the Firestore DB w/ the OTel SDK. Ideally we'd do this is the @BeforeAll method // but because gRPC traces need to be deterministically force-flushed, firestore.shutdown() // must be called in @After for each test. - FirestoreOptions.Builder optionsBuilder = - FirestoreOptions.newBuilder() - .setOpenTelemetryOptions( - FirestoreOpenTelemetryOptions.newBuilder() - .setOpenTelemetry(openTelemetrySdk) - .setTracingEnabled(true) - .build()); + FirestoreOptions.Builder optionsBuilder; + if (isUsingGlobalOpenTelemetrySDK()) { + optionsBuilder = + FirestoreOptions.newBuilder() + .setOpenTelemetryOptions( + FirestoreOpenTelemetryOptions.newBuilder().setTracingEnabled(true).build()); + } else { + optionsBuilder = + FirestoreOptions.newBuilder() + .setOpenTelemetryOptions( + FirestoreOpenTelemetryOptions.newBuilder() + .setOpenTelemetry(openTelemetrySdk) + .setTracingEnabled(true) + .build()); + } String namedDb = System.getProperty("FIRESTORE_NAMED_DATABASE"); if (namedDb != null) { @@ -322,10 +350,20 @@ public void before() throws Exception { firestore, "Error instantiating Firestore. Check that the service account credentials " + "were properly set."); + + // Set up the tracer for custom TraceID injection rootSpanName = String.format("%s%d", this.getClass().getSimpleName(), System.currentTimeMillis()); - tracer = - firestore.getOptions().getOpenTelemetryOptions().getOpenTelemetry().getTracer(rootSpanName); + if (isUsingGlobalOpenTelemetrySDK()) { + tracer = GlobalOpenTelemetry.getTracer(rootSpanName); + } else { + tracer = + firestore + .getOptions() + .getOpenTelemetryOptions() + .getOpenTelemetry() + .getTracer(rootSpanName); + } // Get up a new SpanContext (ergo TraceId) for each test customSpanContext = getNewSpanContext(); @@ -335,6 +373,9 @@ public void before() throws Exception { @After public void after() throws Exception { + if (isUsingGlobalOpenTelemetrySDK()) { + GlobalOpenTelemetry.resetForTest(); + } firestore.shutdown(); rootSpanName = null; tracer = null; From b6291d85552d1b2ebca5078ce6b07fb39326f34a Mon Sep 17 00:00:00 2001 From: jimit-j-shah Date: Thu, 28 Mar 2024 12:54:32 -0700 Subject: [PATCH 20/34] test: formatting and cleanup --- .../cloud/firestore/it/ITE2ETracingTest.java | 92 ++++++++++--------- 1 file changed, 48 insertions(+), 44 deletions(-) diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java index a5b463288..cf93bda4b 100644 --- a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java +++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java @@ -438,8 +438,11 @@ protected void waitForTracesToComplete() throws Exception { } // Validates `retrievedTrace`. Cloud Trace indexes traces w/ eventual consistency, even when - // indexing traceId, therefore the - // test must retry a few times before the complete trace is available. + // indexing traceId, therefore the test must retry a few times before the complete trace is + // available. + // For Transaction traces, there may be more spans than in the trace than specified in + // `callStack`. So `numExpectedSpans` is the expected total number of spans (and not just the + // spans in `callStack`) protected void fetchAndValidateTransactionTrace( String traceId, int numExpectedSpans, String... callStack) throws InterruptedException { // Large enough count to accommodate eventually consistent Cloud Trace backend @@ -489,9 +492,10 @@ protected void fetchAndValidateTransactionTrace( } // Validates `retrievedTrace`. Cloud Trace indexes traces w/ eventual consistency, even when - // indexing traceId, therefore the - // test must retry a few times before the complete trace is available. - protected void fetchAndValidateTraces(String traceId, String... spanNames) + // indexing traceId, therefore the test must retry a few times before the complete trace is + // available. + // For Non-Transaction traces, there is a 1:1 ratio of spans in `spanNames` and in the trace. + protected void fetchAndValidateNonTransactionTrace(String traceId, String... spanNames) throws InterruptedException { int numRetries = GET_TRACE_RETRY_COUNT; do { @@ -593,7 +597,7 @@ public void aggregateQueryGetTraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), "AggregationQuery.Get", grpcSpanName(RUN_AGGREGATION_QUERY_RPC_NAME)); @@ -622,7 +626,7 @@ public void bulkWriterCommitTraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_BULK_WRITER_COMMIT, grpcSpanName(BATCH_WRITE_RPC_NAME)); @@ -644,7 +648,7 @@ public void partitionQueryTraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_PARTITION_QUERY, grpcSpanName(SPAN_NAME_PARTITION_QUERY)); @@ -665,7 +669,7 @@ public void collectionListDocumentsTraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_COL_REF_LIST_DOCUMENTS, grpcSpanName(LIST_DOCUMENTS_RPC_NAME)); @@ -686,7 +690,7 @@ public void docRefCreateTraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_CREATE, SPAN_NAME_BATCH_COMMIT, @@ -708,7 +712,7 @@ public void docRefCreate2TraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_CREATE, SPAN_NAME_BATCH_COMMIT, @@ -730,7 +734,7 @@ public void docRefSetTraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_SET, SPAN_NAME_BATCH_COMMIT, @@ -756,7 +760,7 @@ public void docRefSet2TraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_SET, SPAN_NAME_BATCH_COMMIT, @@ -778,7 +782,7 @@ public void docRefSet3TraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_SET, SPAN_NAME_BATCH_COMMIT, @@ -800,7 +804,7 @@ public void docRefSet4TraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_SET, SPAN_NAME_BATCH_COMMIT, @@ -808,7 +812,7 @@ public void docRefSet4TraceTest() throws Exception { } @Test - public void docRefUpdate() throws Exception { + public void docRefUpdateTraceTest() throws Exception { // Make sure the test has a new SpanContext (and TraceId for injection) assertNotNull(customSpanContext); @@ -826,7 +830,7 @@ public void docRefUpdate() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_UPDATE, SPAN_NAME_BATCH_COMMIT, @@ -834,7 +838,7 @@ public void docRefUpdate() throws Exception { } @Test - public void docRefUpdate2() throws Exception { + public void docRefUpdate2TraceTest() throws Exception { // Make sure the test has a new SpanContext (and TraceId for injection) assertNotNull(customSpanContext); @@ -852,7 +856,7 @@ public void docRefUpdate2() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_UPDATE, SPAN_NAME_BATCH_COMMIT, @@ -860,7 +864,7 @@ public void docRefUpdate2() throws Exception { } @Test - public void docRefUpdate3() throws Exception { + public void docRefUpdate3TraceTest() throws Exception { // Make sure the test has a new SpanContext (and TraceId for injection) assertNotNull(customSpanContext); @@ -874,7 +878,7 @@ public void docRefUpdate3() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_UPDATE, SPAN_NAME_BATCH_COMMIT, @@ -882,7 +886,7 @@ public void docRefUpdate3() throws Exception { } @Test - public void docRefUpdate4() throws Exception { + public void docRefUpdate4TraceTest() throws Exception { // Make sure the test has a new SpanContext (and TraceId for injection) assertNotNull(customSpanContext); @@ -900,7 +904,7 @@ public void docRefUpdate4() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_UPDATE, SPAN_NAME_BATCH_COMMIT, @@ -908,7 +912,7 @@ public void docRefUpdate4() throws Exception { } @Test - public void docRefUpdate5() throws Exception { + public void docRefUpdate5TraceTest() throws Exception { // Make sure the test has a new SpanContext (and TraceId for injection) assertNotNull(customSpanContext); @@ -926,7 +930,7 @@ public void docRefUpdate5() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_UPDATE, SPAN_NAME_BATCH_COMMIT, @@ -934,7 +938,7 @@ public void docRefUpdate5() throws Exception { } @Test - public void docRefUpdate6() throws Exception { + public void docRefUpdate6TraceTest() throws Exception { // Make sure the test has a new SpanContext (and TraceId for injection) assertNotNull(customSpanContext); @@ -952,7 +956,7 @@ public void docRefUpdate6() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_UPDATE, SPAN_NAME_BATCH_COMMIT, @@ -960,7 +964,7 @@ public void docRefUpdate6() throws Exception { } @Test - public void docRefDelete() throws Exception { + public void docRefDeleteTraceTest() throws Exception { // Make sure the test has a new SpanContext (and TraceId for injection) assertNotNull(customSpanContext); @@ -974,7 +978,7 @@ public void docRefDelete() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_DELETE, SPAN_NAME_BATCH_COMMIT, @@ -982,7 +986,7 @@ public void docRefDelete() throws Exception { } @Test - public void docRefDelete2() throws Exception { + public void docRefDelete2TraceTest() throws Exception { // Make sure the test has a new SpanContext (and TraceId for injection) assertNotNull(customSpanContext); @@ -996,7 +1000,7 @@ public void docRefDelete2() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_DELETE, SPAN_NAME_BATCH_COMMIT, @@ -1004,7 +1008,7 @@ public void docRefDelete2() throws Exception { } @Test - public void docRefGet() throws Exception { + public void docRefGetTraceTest() throws Exception { // Make sure the test has a new SpanContext (and TraceId for injection) assertNotNull(customSpanContext); @@ -1018,14 +1022,14 @@ public void docRefGet() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_GET, grpcSpanName(BATCH_GET_DOCUMENTS_RPC_NAME)); } @Test - public void docRefGet2() throws Exception { + public void docRefGet2TraceTest() throws Exception { // Make sure the test has a new SpanContext (and TraceId for injection) assertNotNull(customSpanContext); @@ -1039,14 +1043,14 @@ public void docRefGet2() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_GET, grpcSpanName(BATCH_GET_DOCUMENTS_RPC_NAME)); } @Test - public void docListCollections() throws Exception { + public void docListCollectionsTraceTest() throws Exception { // Make sure the test has a new SpanContext (and TraceId for injection) assertNotNull(customSpanContext); @@ -1060,14 +1064,14 @@ public void docListCollections() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_LIST_COLLECTIONS, grpcSpanName(LIST_COLLECTIONS_RPC_NAME)); } @Test - public void getAll() throws Exception { + public void getAllTraceTest() throws Exception { // Make sure the test has a new SpanContext (and TraceId for injection) assertNotNull(customSpanContext); @@ -1083,12 +1087,12 @@ public void getAll() throws Exception { } waitForTracesToComplete(); - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), grpcSpanName(BATCH_GET_DOCUMENTS_RPC_NAME)); } @Test - public void queryGet() throws Exception { + public void queryGetTraceTest() throws Exception { // Make sure the test has a new SpanContext (and TraceId for injection) assertNotNull(customSpanContext); @@ -1101,12 +1105,12 @@ public void queryGet() throws Exception { } waitForTracesToComplete(); - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_QUERY_GET, grpcSpanName(RUN_QUERY_RPC_NAME)); } @Test - public void transaction() throws Exception { + public void transactionTraceTest() throws Exception { // Make sure the test has a new SpanContext (and TraceId for injection) assertNotNull(customSpanContext); @@ -1180,7 +1184,7 @@ public void transaction() throws Exception { } @Test - public void transactionRollback() throws Exception { + public void transactionRollbackTraceTest() throws Exception { String myErrorMessage = "My error message."; // Make sure the test has a new SpanContext (and TraceId for injection) assertNotNull(customSpanContext); @@ -1220,7 +1224,7 @@ public void transactionRollback() throws Exception { } @Test - public void writeBatch() throws Exception { + public void writeBatchTraceTest() throws Exception { // Make sure the test has a new SpanContext (and TraceId for injection) assertNotNull(customSpanContext); From 1f984f12accaa2e8312c211e802f12e2e0fbabcb Mon Sep 17 00:00:00 2001 From: jimit-j-shah Date: Fri, 8 Mar 2024 16:18:23 -0800 Subject: [PATCH 21/34] test: Adding first e2e client-tracing test w/ Custom Root Span --- google-cloud-firestore/pom.xml | 35 +++ .../cloud/firestore/it/ITE2ETracingTest.java | 248 ++++++++++++++++++ 2 files changed, 283 insertions(+) create mode 100644 google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java diff --git a/google-cloud-firestore/pom.xml b/google-cloud-firestore/pom.xml index 942e38ad5..becd396d4 100644 --- a/google-cloud-firestore/pom.xml +++ b/google-cloud-firestore/pom.xml @@ -105,6 +105,7 @@ com.google.protobuf protobuf-java-util + io.opentelemetry @@ -213,7 +214,41 @@ ${opentelemetry.version} test + + com.google.cloud.opentelemetry + exporter-trace + 0.15.0 + test + + + + + com.google.api.grpc + proto-google-cloud-trace-v1 + 1.3.0 + test + + + com.google.cloud.opentelemetry + exporter-trace + 0.15.0 + test + + + + com.google.api.grpc + proto-google-cloud-trace-v1 + 1.3.0 + test + + + com.google.cloud + google-cloud-trace + 1.3.0 + test + + diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java new file mode 100644 index 000000000..32f3cf936 --- /dev/null +++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java @@ -0,0 +1,248 @@ +/* + * Copyright 2024 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.firestore.it; + +import static io.opentelemetry.semconv.resource.attributes.ResourceAttributes.SERVICE_NAME; +import static org.junit.Assert.assertEquals; + +import com.google.api.gax.rpc.NotFoundException; +import com.google.cloud.firestore.Firestore; +import com.google.cloud.firestore.FirestoreOpenTelemetryOptions; +import com.google.cloud.firestore.FirestoreOptions; +import com.google.cloud.opentelemetry.trace.TraceConfiguration; +import com.google.cloud.opentelemetry.trace.TraceExporter; +import com.google.cloud.trace.v1.TraceServiceClient; +import com.google.common.base.Preconditions; +import com.google.devtools.cloudtrace.v1.Trace; +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.SpanContext; +import io.opentelemetry.api.trace.TraceFlags; +import io.opentelemetry.api.trace.TraceState; +import io.opentelemetry.api.trace.Tracer; +import io.opentelemetry.context.Context; +import io.opentelemetry.context.Scope; +import io.opentelemetry.sdk.OpenTelemetrySdk; +import io.opentelemetry.sdk.common.CompletableResultCode; +import io.opentelemetry.sdk.resources.Resource; +import io.opentelemetry.sdk.trace.SdkTracerProvider; +import io.opentelemetry.sdk.trace.export.BatchSpanProcessor; +import io.opentelemetry.sdk.trace.samplers.Sampler; +import java.io.IOException; +import java.util.Random; +import java.util.concurrent.TimeUnit; +import java.util.logging.Level; +import java.util.logging.Logger; +import org.junit.After; +import org.junit.AfterClass; +import org.junit.Before; +import org.junit.BeforeClass; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class ITE2ETracingTest extends ITBaseTest { + + private static final Logger logger = Logger.getLogger(ITBaseTest.class.getName()); + + private static final int NUM_TRACE_ID_BYTES = 32; + + private static final int NUM_SPAN_ID_BYTES = 16; + + private static final int GET_TRACE_RETRY_COUNT = 10; + + private static final int GET_TRACE_RETRY_BACKOFF_MILLIS = 1000; + + private static final int TRACE_FORCE_FLUSH_MILLIS = 1000; + + private static final int TRACE_PROVIDER_SHUTDOWN_MILLIS = 1000; + + // Random int generator for trace ID and span ID + private static Random random; + + private static TraceExporter traceExporter; + + // Required for reading back traces from Cloud Trace for validation + private static TraceServiceClient traceClient_v1; + + private static String rootSpanName; + private static Tracer tracer; + + // Required to set custom-root span + private static OpenTelemetrySdk openTelemetrySdk; + + private static String projectId; + + private static Firestore firestore; + + @BeforeClass + public static void setup() throws IOException { + projectId = FirestoreOptions.getDefaultProjectId(); + logger.info("projectId:" + projectId); + + // Set up OTel SDK + Resource resource = + Resource.getDefault().merge(Resource.builder().put(SERVICE_NAME, "Sparky").build()); + + // TODO(jimit) Make it re-usable w/ InMemorySpanExporter + traceExporter = + TraceExporter.createWithConfiguration( + TraceConfiguration.builder().setProjectId(projectId).build()); + + openTelemetrySdk = + OpenTelemetrySdk.builder() + .setTracerProvider( + SdkTracerProvider.builder() + .setResource(resource) + .addSpanProcessor(BatchSpanProcessor.builder(traceExporter).build()) + .setSampler(Sampler.alwaysOn()) + .build()) + .build(); + + traceClient_v1 = TraceServiceClient.create(); + + // Initialize the Firestore DB w/ the OTel SDK + FirestoreOptions.Builder optionsBuilder = + FirestoreOptions.newBuilder() + .setOpenTelemetryOptions( + FirestoreOpenTelemetryOptions.newBuilder() + .setOpenTelemetry(openTelemetrySdk) + .setTracingEnabled(true) + .build()); + + String namedDb = System.getProperty("FIRESTORE_NAMED_DATABASE"); + if (namedDb != null) { + logger.log(Level.INFO, "Integration test using named database " + namedDb); + optionsBuilder = optionsBuilder.setDatabaseId(namedDb); + } else { + logger.log(Level.INFO, "Integration test using default database."); + } + firestore = optionsBuilder.build().getService(); + Preconditions.checkNotNull( + firestore, + "Error instantiating Firestore. Check that the service account credentials " + + "were properly set."); + random = new Random(); + } + + @Before + public void before() throws Exception { + rootSpanName = + String.format("%s%d", this.getClass().getSimpleName(), System.currentTimeMillis()); + tracer = + firestore.getOptions().getOpenTelemetryOptions().getOpenTelemetry().getTracer(rootSpanName); + } + + @After + public void after() throws Exception { + rootSpanName = null; + tracer = null; + } + + @AfterClass + public static void teardown() { + traceClient_v1.close(); + CompletableResultCode completableResultCode = + openTelemetrySdk.getSdkTracerProvider().shutdown(); + completableResultCode.join(TRACE_PROVIDER_SHUTDOWN_MILLIS, TimeUnit.MILLISECONDS); + firestore.shutdown(); + } + + // Generates a random 32-byte hex string + private String generateRandomHexString(int numBytes) { + StringBuffer newTraceId = new StringBuffer(); + while (newTraceId.length() < numBytes) { + newTraceId.append(Integer.toHexString(random.nextInt())); + } + return newTraceId.substring(0, numBytes); + } + + protected String generateNewTraceId() { + return generateRandomHexString(NUM_TRACE_ID_BYTES); + } + + // Generates a random 16-byte hex string + protected String generateNewSpanId() { + return generateRandomHexString(NUM_SPAN_ID_BYTES); + } + + // Cloud Trace indexes traces w/ eventual consistency, even when indexing traceId, therefore the + // test must retry a few times before the trace is available. + protected Trace getTraceWithRetry(String project, String traceId) + throws InterruptedException, RuntimeException { + int retryCount = GET_TRACE_RETRY_COUNT; + + while (retryCount-- > 0) { + try { + Trace t = traceClient_v1.getTrace(project, traceId); + return t; + } catch (NotFoundException notFound) { + logger.warning("Trace not found, retrying in " + GET_TRACE_RETRY_BACKOFF_MILLIS + " ms"); + Thread.sleep(GET_TRACE_RETRY_BACKOFF_MILLIS); + } + } + throw new RuntimeException( + "Trace " + traceId + " for project " + project + " not found in Cloud Trace."); + } + + // Generates a new SpanContext w/ random traceId,spanId + protected SpanContext getNewSpanContext() { + String traceId = generateNewTraceId(); + String spanId = generateNewSpanId(); + logger.info("traceId=" + traceId + ", spanId=" + spanId); + + return SpanContext.create(traceId, spanId, TraceFlags.getSampled(), TraceState.getDefault()); + } + + protected void waitForTracesToComplete() throws Exception { + CompletableResultCode completableResultCode = + openTelemetrySdk.getSdkTracerProvider().forceFlush(); + completableResultCode.join(TRACE_FORCE_FLUSH_MILLIS, TimeUnit.MILLISECONDS); + + // We need to call `firestore.close()` because that will also close the + // gRPC channel and hence force the gRPC instrumentation library to flush + // its spans. + firestore.close(); + } + + @Test + // Trace an Aggregation.Get request + public void basicTraceTestWithCustomRootSpan() throws Exception { + SpanContext newCtx = getNewSpanContext(); + + // Execute the DB operation in the context of the custom root span. + Span rootSpan = + tracer + .spanBuilder(rootSpanName) + .setParent(Context.root().with(Span.wrap(newCtx))) + .startSpan(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore.collection("col").count().get().get(); + } finally { + rootSpan.end(); + } + + // Ingest traces to Cloud Trace + waitForTracesToComplete(); + + String traceId = newCtx.getTraceId(); + Trace t = getTraceWithRetry(projectId, traceId); + assertEquals(t.getTraceId(), traceId); + assertEquals(t.getSpans(0).getName(), rootSpanName); + assertEquals(t.getSpans(1).getName(), "AggregationQuery.Get"); + } +} From cd77a48afe5323fc8d6437b851f36438e05af61f Mon Sep 17 00:00:00 2001 From: jimit-j-shah Date: Mon, 18 Mar 2024 15:53:27 -0700 Subject: [PATCH 22/34] test: Add aggregationQueryGet Test --- .../cloud/firestore/it/ITE2ETracingTest.java | 127 ++++++++++++++++-- 1 file changed, 113 insertions(+), 14 deletions(-) diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java index 32f3cf936..b717d91bf 100644 --- a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java +++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java @@ -18,6 +18,7 @@ import static io.opentelemetry.semconv.resource.attributes.ResourceAttributes.SERVICE_NAME; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNull; import com.google.api.gax.rpc.NotFoundException; import com.google.cloud.firestore.Firestore; @@ -42,6 +43,8 @@ import io.opentelemetry.sdk.trace.export.BatchSpanProcessor; import io.opentelemetry.sdk.trace.samplers.Sampler; import java.io.IOException; +import java.util.Arrays; +import java.util.List; import java.util.Random; import java.util.concurrent.TimeUnit; import java.util.logging.Level; @@ -79,6 +82,9 @@ public class ITE2ETracingTest extends ITBaseTest { // Required for reading back traces from Cloud Trace for validation private static TraceServiceClient traceClient_v1; + // Trace received + private static Trace retrievedTrace; + private static String rootSpanName; private static Tracer tracer; @@ -145,12 +151,14 @@ public void before() throws Exception { String.format("%s%d", this.getClass().getSimpleName(), System.currentTimeMillis()); tracer = firestore.getOptions().getOpenTelemetryOptions().getOpenTelemetry().getTracer(rootSpanName); + assertNull(retrievedTrace); } @After public void after() throws Exception { rootSpanName = null; tracer = null; + retrievedTrace = null; } @AfterClass @@ -208,6 +216,13 @@ protected SpanContext getNewSpanContext() { return SpanContext.create(traceId, spanId, TraceFlags.getSampled(), TraceState.getDefault()); } + protected Span getNewRootSpanWithContext(SpanContext spanContext) { + // Execute the DB operation in the context of the custom root span. + return tracer.spanBuilder(rootSpanName) + .setParent(Context.root().with(Span.wrap(spanContext))) + .startSpan(); + } + protected void waitForTracesToComplete() throws Exception { CompletableResultCode completableResultCode = openTelemetrySdk.getSdkTracerProvider().forceFlush(); @@ -219,30 +234,114 @@ protected void waitForTracesToComplete() throws Exception { firestore.close(); } + // Validates `retrievedTrace` + protected void fetchAndValidateTraces(String traceId, String... spanNames) + throws InterruptedException { + // Fetch traces + retrievedTrace = getTraceWithRetry(projectId, traceId); + List spanNameList = Arrays.asList(spanNames); + + // Validate trace spans + assertEquals(retrievedTrace.getTraceId(), traceId); + assertEquals(retrievedTrace.getSpans(0).getName(), rootSpanName); + for (int i = 0; i < spanNameList.size() ; ++i) { + assertEquals(retrievedTrace.getSpans(i+1).getName(), spanNameList.get(i)); + } + } + @Test // Trace an Aggregation.Get request - public void basicTraceTestWithCustomRootSpan() throws Exception { + public void aggregateQueryGet() throws Exception { + // Create custom span context for trace ID injection SpanContext newCtx = getNewSpanContext(); - // Execute the DB operation in the context of the custom root span. - Span rootSpan = - tracer - .spanBuilder(rootSpanName) - .setParent(Context.root().with(Span.wrap(newCtx))) - .startSpan(); + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(newCtx); try (Scope ss = rootSpan.makeCurrent()) { firestore.collection("col").count().get().get(); } finally { rootSpan.end(); } - - // Ingest traces to Cloud Trace waitForTracesToComplete(); - String traceId = newCtx.getTraceId(); - Trace t = getTraceWithRetry(projectId, traceId); - assertEquals(t.getTraceId(), traceId); - assertEquals(t.getSpans(0).getName(), rootSpanName); - assertEquals(t.getSpans(1).getName(), "AggregationQuery.Get"); + // Read and validate traces + fetchAndValidateTraces(newCtx.getTraceId(), "AggregationQuery.Get"); + } + + public void bulkWriterCommit() throws Exception { + } + + @Test + public void partitionQuery() throws Exception { } + + @Test + public void collectionListDocuments() throws Exception { + } + + public void docRefCreate() throws Exception { + } + + @Test + public void docRefCreate2() throws Exception { + } + + @Test + public void docRefSet() throws Exception {} + + @Test + public void docRefSet2() throws Exception {} + + @Test + public void docRefSet3() throws Exception {} + + public void docRefSet4() throws Exception {} + + @Test + public void docRefUpdate() throws Exception {} + + @Test + public void docRefUpdate2() throws Exception {} + + @Test + public void docRefUpdate3() throws Exception {} + + @Test + public void docRefUpdate4() throws Exception {} + + @Test + public void docRefUpdate5() throws Exception {} + + @Test + public void docRefUpdate6() throws Exception {} + + @Test + public void docRefDelete() throws Exception {} + + @Test + public void docRefDelete2() throws Exception {} + + @Test + public void docRefGet() throws Exception {} + + @Test + public void docRefGet2() throws Exception {} + + @Test + public void docListCollections() throws Exception {} + + @Test + public void getAll() throws Exception {} + + @Test + public void queryGet() throws Exception {} + + @Test + public void transaction() throws Exception {} + + @Test + public void transactionRollback() throws Exception {} + + @Test + public void writeBatch() throws Exception {} } From c26ee57458c3f366c1ce08fa881b426e410ce737 Mon Sep 17 00:00:00 2001 From: jimit-j-shah Date: Mon, 18 Mar 2024 16:08:05 -0700 Subject: [PATCH 23/34] test: Add bulkWriterCommitTrace Test and fixed running multiple-tests --- .../cloud/firestore/it/ITE2ETracingTest.java | 80 +++++++++++++++---- 1 file changed, 64 insertions(+), 16 deletions(-) diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java index b717d91bf..44f28335b 100644 --- a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java +++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java @@ -16,14 +16,19 @@ package com.google.cloud.firestore.it; +import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_BULK_WRITER_COMMIT; import static io.opentelemetry.semconv.resource.attributes.ResourceAttributes.SERVICE_NAME; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertNull; import com.google.api.gax.rpc.NotFoundException; +import com.google.cloud.firestore.BulkWriter; +import com.google.cloud.firestore.BulkWriterOptions; import com.google.cloud.firestore.Firestore; import com.google.cloud.firestore.FirestoreOpenTelemetryOptions; import com.google.cloud.firestore.FirestoreOptions; +import com.google.cloud.firestore.telemetry.TraceUtil; import com.google.cloud.opentelemetry.trace.TraceConfiguration; import com.google.cloud.opentelemetry.trace.TraceExporter; import com.google.cloud.trace.v1.TraceServiceClient; @@ -44,8 +49,11 @@ import io.opentelemetry.sdk.trace.samplers.Sampler; import java.io.IOException; import java.util.Arrays; +import java.util.Collections; import java.util.List; import java.util.Random; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import java.util.logging.Level; import java.util.logging.Logger; @@ -62,6 +70,10 @@ public class ITE2ETracingTest extends ITBaseTest { private static final Logger logger = Logger.getLogger(ITBaseTest.class.getName()); + private static final String SERVICE = "google.firestore.v1.Firestore/"; + + private static final String BATCH_WRITE_RPC_NAME = "BatchWrite"; + private static final int NUM_TRACE_ID_BYTES = 32; private static final int NUM_SPAN_ID_BYTES = 16; @@ -82,7 +94,10 @@ public class ITE2ETracingTest extends ITBaseTest { // Required for reading back traces from Cloud Trace for validation private static TraceServiceClient traceClient_v1; - // Trace received + // Custom SpanContext for each test, required for TraceID injection + private static SpanContext customSpanContext; + + // Trace read back from Cloud Trace using traceClient_v1 for verification private static Trace retrievedTrace; private static String rootSpanName; @@ -151,6 +166,10 @@ public void before() throws Exception { String.format("%s%d", this.getClass().getSimpleName(), System.currentTimeMillis()); tracer = firestore.getOptions().getOpenTelemetryOptions().getOpenTelemetry().getTracer(rootSpanName); + + // Get up a new SpanContext (ergo TraceId) for each test + customSpanContext = getNewSpanContext(); + assertNotNull(customSpanContext); assertNull(retrievedTrace); } @@ -159,14 +178,16 @@ public void after() throws Exception { rootSpanName = null; tracer = null; retrievedTrace = null; + customSpanContext = null; } @AfterClass - public static void teardown() { + public static void teardown() throws Exception { traceClient_v1.close(); CompletableResultCode completableResultCode = openTelemetrySdk.getSdkTracerProvider().shutdown(); completableResultCode.join(TRACE_PROVIDER_SHUTDOWN_MILLIS, TimeUnit.MILLISECONDS); + firestore.close(); firestore.shutdown(); } @@ -216,22 +237,22 @@ protected SpanContext getNewSpanContext() { return SpanContext.create(traceId, spanId, TraceFlags.getSampled(), TraceState.getDefault()); } - protected Span getNewRootSpanWithContext(SpanContext spanContext) { + protected Span getNewRootSpanWithContext() { // Execute the DB operation in the context of the custom root span. return tracer.spanBuilder(rootSpanName) - .setParent(Context.root().with(Span.wrap(spanContext))) + .setParent(Context.root().with(Span.wrap(customSpanContext))) .startSpan(); } + protected String grpcSpanName(String rpcName) { + return "Sent." + SERVICE + rpcName; + } + protected void waitForTracesToComplete() throws Exception { + logger.info("Flushing traces..."); CompletableResultCode completableResultCode = openTelemetrySdk.getSdkTracerProvider().forceFlush(); completableResultCode.join(TRACE_FORCE_FLUSH_MILLIS, TimeUnit.MILLISECONDS); - - // We need to call `firestore.close()` because that will also close the - // gRPC channel and hence force the gRPC instrumentation library to flush - // its spans. - firestore.close(); } // Validates `retrievedTrace` @@ -245,19 +266,20 @@ protected void fetchAndValidateTraces(String traceId, String... spanNames) assertEquals(retrievedTrace.getTraceId(), traceId); assertEquals(retrievedTrace.getSpans(0).getName(), rootSpanName); for (int i = 0; i < spanNameList.size() ; ++i) { - assertEquals(retrievedTrace.getSpans(i+1).getName(), spanNameList.get(i)); + assertEquals(spanNameList.get(i), retrievedTrace.getSpans(i+1).getName()); } } @Test // Trace an Aggregation.Get request - public void aggregateQueryGet() throws Exception { - // Create custom span context for trace ID injection - SpanContext newCtx = getNewSpanContext(); + public void aggregateQueryGetTraceTest() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); // Inject new trace ID - Span rootSpan = getNewRootSpanWithContext(newCtx); + Span rootSpan = getNewRootSpanWithContext(); try (Scope ss = rootSpan.makeCurrent()) { + // Execute the Firestore SDK op firestore.collection("col").count().get().get(); } finally { rootSpan.end(); @@ -265,14 +287,40 @@ public void aggregateQueryGet() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces(newCtx.getTraceId(), "AggregationQuery.Get"); + fetchAndValidateTraces(customSpanContext.getTraceId(), "AggregationQuery.Get"); } - public void bulkWriterCommit() throws Exception { + @Test + public void bulkWriterCommitTraceTest() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + // Execute the Firestore SDK op + ScheduledExecutorService bulkWriterExecutor = Executors.newSingleThreadScheduledExecutor(); + BulkWriter bulkWriter = + firestore.bulkWriter(BulkWriterOptions.builder().setExecutor(bulkWriterExecutor).build()); + bulkWriter.set( + firestore.collection("col").document("foo"), + Collections.singletonMap("bulk-foo", "bulk-bar")); + bulkWriter.close(); + bulkWriterExecutor.awaitTermination(100, TimeUnit.MILLISECONDS); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + // Read and validate traces + fetchAndValidateTraces(customSpanContext.getTraceId(), + SPAN_NAME_BULK_WRITER_COMMIT, + grpcSpanName(BATCH_WRITE_RPC_NAME)); } @Test public void partitionQuery() throws Exception { + } @Test From 7fb283ac8fc3f50e1e091f0eef6c2e3ffe21e796 Mon Sep 17 00:00:00 2001 From: jimit-j-shah Date: Mon, 18 Mar 2024 16:30:42 -0700 Subject: [PATCH 24/34] test: Add partitionQuery Test --- .../cloud/firestore/it/ITE2ETracingTest.java | 22 +++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java index 44f28335b..ac9c22814 100644 --- a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java +++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java @@ -17,6 +17,7 @@ package com.google.cloud.firestore.it; import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_BULK_WRITER_COMMIT; +import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_PARTITION_QUERY; import static io.opentelemetry.semconv.resource.attributes.ResourceAttributes.SERVICE_NAME; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNotNull; @@ -25,6 +26,7 @@ import com.google.api.gax.rpc.NotFoundException; import com.google.cloud.firestore.BulkWriter; import com.google.cloud.firestore.BulkWriterOptions; +import com.google.cloud.firestore.CollectionGroup; import com.google.cloud.firestore.Firestore; import com.google.cloud.firestore.FirestoreOpenTelemetryOptions; import com.google.cloud.firestore.FirestoreOptions; @@ -78,7 +80,7 @@ public class ITE2ETracingTest extends ITBaseTest { private static final int NUM_SPAN_ID_BYTES = 16; - private static final int GET_TRACE_RETRY_COUNT = 10; + private static final int GET_TRACE_RETRY_COUNT = 15; private static final int GET_TRACE_RETRY_BACKOFF_MILLIS = 1000; @@ -319,8 +321,24 @@ public void bulkWriterCommitTraceTest() throws Exception { } @Test - public void partitionQuery() throws Exception { + public void partitionQueryTraceTest() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + CollectionGroup collectionGroup = firestore.collectionGroup("col"); + collectionGroup.getPartitions(3).get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + // Read and validate traces + fetchAndValidateTraces(customSpanContext.getTraceId(), + SPAN_NAME_PARTITION_QUERY, + grpcSpanName(SPAN_NAME_PARTITION_QUERY)); } @Test From 47c64ae32cf20572edf392f9b590a42f803a43c4 Mon Sep 17 00:00:00 2001 From: jimit-j-shah Date: Mon, 18 Mar 2024 16:33:32 -0700 Subject: [PATCH 25/34] test: Add collectionListDocumentsTrace Test --- .../cloud/firestore/it/ITE2ETracingTest.java | 20 ++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java index ac9c22814..6f5c1390d 100644 --- a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java +++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java @@ -17,6 +17,7 @@ package com.google.cloud.firestore.it; import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_BULK_WRITER_COMMIT; +import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_COL_REF_LIST_DOCUMENTS; import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_PARTITION_QUERY; import static io.opentelemetry.semconv.resource.attributes.ResourceAttributes.SERVICE_NAME; import static org.junit.Assert.assertEquals; @@ -76,6 +77,8 @@ public class ITE2ETracingTest extends ITBaseTest { private static final String BATCH_WRITE_RPC_NAME = "BatchWrite"; + private static final String LIST_DOCUMENTS_RPC_NAME = "ListDocuments"; + private static final int NUM_TRACE_ID_BYTES = 32; private static final int NUM_SPAN_ID_BYTES = 16; @@ -342,7 +345,22 @@ public void partitionQueryTraceTest() throws Exception { } @Test - public void collectionListDocuments() throws Exception { + public void collectionListDocumentsTraceTest() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore.collection("col").listDocuments(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + // Read and validate traces + fetchAndValidateTraces(customSpanContext.getTraceId(), + SPAN_NAME_COL_REF_LIST_DOCUMENTS, grpcSpanName(LIST_DOCUMENTS_RPC_NAME)); } public void docRefCreate() throws Exception { From 02805f0b95e0f11e948d9e0870e724852c604cad Mon Sep 17 00:00:00 2001 From: jimit-j-shah Date: Tue, 19 Mar 2024 10:50:50 -0700 Subject: [PATCH 26/34] test: Add docRefUpdate*Trace and docRefDelete*Trace Tests and fixed Trace fetching using retries for missing or incomplete traces due to eventual consistency of Cloud Trace --- .../cloud/firestore/it/ITE2ETracingTest.java | 481 +++++++++++++++--- 1 file changed, 418 insertions(+), 63 deletions(-) diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java index 6f5c1390d..1e84214be 100644 --- a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java +++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java @@ -16,8 +16,14 @@ package com.google.cloud.firestore.it; +import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_BATCH_COMMIT; import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_BULK_WRITER_COMMIT; import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_COL_REF_LIST_DOCUMENTS; +import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_DOC_REF_CREATE; +import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_DOC_REF_DELETE; +import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_DOC_REF_GET; +import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_DOC_REF_SET; +import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_DOC_REF_UPDATE; import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_PARTITION_QUERY; import static io.opentelemetry.semconv.resource.attributes.ResourceAttributes.SERVICE_NAME; import static org.junit.Assert.assertEquals; @@ -28,10 +34,14 @@ import com.google.cloud.firestore.BulkWriter; import com.google.cloud.firestore.BulkWriterOptions; import com.google.cloud.firestore.CollectionGroup; +import com.google.cloud.firestore.FieldMask; +import com.google.cloud.firestore.FieldPath; import com.google.cloud.firestore.Firestore; import com.google.cloud.firestore.FirestoreOpenTelemetryOptions; import com.google.cloud.firestore.FirestoreOptions; -import com.google.cloud.firestore.telemetry.TraceUtil; +import com.google.cloud.firestore.Precondition; +import com.google.cloud.firestore.SetOptions; +import com.google.cloud.firestore.it.ITTracingTest.Pojo; import com.google.cloud.opentelemetry.trace.TraceConfiguration; import com.google.cloud.opentelemetry.trace.TraceExporter; import com.google.cloud.trace.v1.TraceServiceClient; @@ -51,9 +61,9 @@ import io.opentelemetry.sdk.trace.export.BatchSpanProcessor; import io.opentelemetry.sdk.trace.samplers.Sampler; import java.io.IOException; +import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; -import java.util.List; import java.util.Random; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; @@ -75,10 +85,16 @@ public class ITE2ETracingTest extends ITBaseTest { private static final String SERVICE = "google.firestore.v1.Firestore/"; + private static final String BATCH_GET_DOCUMENTS_RPC_NAME = "BatchGetDocuments"; + private static final String BATCH_WRITE_RPC_NAME = "BatchWrite"; + private static final String COMMIT_RPC_NAME = "Commit"; + private static final String LIST_DOCUMENTS_RPC_NAME = "ListDocuments"; + private static final String RUN_AGGREGATION_QUERY_RPC_NAME = "RunAggregationQuery"; + private static final int NUM_TRACE_ID_BYTES = 32; private static final int NUM_SPAN_ID_BYTES = 16; @@ -87,7 +103,7 @@ public class ITE2ETracingTest extends ITBaseTest { private static final int GET_TRACE_RETRY_BACKOFF_MILLIS = 1000; - private static final int TRACE_FORCE_FLUSH_MILLIS = 1000; + private static final int TRACE_FORCE_FLUSH_MILLIS = 3000; private static final int TRACE_PROVIDER_SHUTDOWN_MILLIS = 1000; @@ -141,7 +157,14 @@ public static void setup() throws IOException { traceClient_v1 = TraceServiceClient.create(); - // Initialize the Firestore DB w/ the OTel SDK + random = new Random(); + } + + @Before + public void before() throws Exception { + // Initialize the Firestore DB w/ the OTel SDK. Ideally we'd do this is the @BeforeAll method + // but because gRPC traces need to be deterministically force-flushed, firestore.shutdown() + // must be called in @After for each test. FirestoreOptions.Builder optionsBuilder = FirestoreOptions.newBuilder() .setOpenTelemetryOptions( @@ -162,11 +185,6 @@ public static void setup() throws IOException { firestore, "Error instantiating Firestore. Check that the service account credentials " + "were properly set."); - random = new Random(); - } - - @Before - public void before() throws Exception { rootSpanName = String.format("%s%d", this.getClass().getSimpleName(), System.currentTimeMillis()); tracer = @@ -180,6 +198,7 @@ public void before() throws Exception { @After public void after() throws Exception { + firestore.shutdown(); rootSpanName = null; tracer = null; retrievedTrace = null; @@ -192,8 +211,6 @@ public static void teardown() throws Exception { CompletableResultCode completableResultCode = openTelemetrySdk.getSdkTracerProvider().shutdown(); completableResultCode.join(TRACE_PROVIDER_SHUTDOWN_MILLIS, TimeUnit.MILLISECONDS); - firestore.close(); - firestore.shutdown(); } // Generates a random 32-byte hex string @@ -214,25 +231,6 @@ protected String generateNewSpanId() { return generateRandomHexString(NUM_SPAN_ID_BYTES); } - // Cloud Trace indexes traces w/ eventual consistency, even when indexing traceId, therefore the - // test must retry a few times before the trace is available. - protected Trace getTraceWithRetry(String project, String traceId) - throws InterruptedException, RuntimeException { - int retryCount = GET_TRACE_RETRY_COUNT; - - while (retryCount-- > 0) { - try { - Trace t = traceClient_v1.getTrace(project, traceId); - return t; - } catch (NotFoundException notFound) { - logger.warning("Trace not found, retrying in " + GET_TRACE_RETRY_BACKOFF_MILLIS + " ms"); - Thread.sleep(GET_TRACE_RETRY_BACKOFF_MILLIS); - } - } - throw new RuntimeException( - "Trace " + traceId + " for project " + project + " not found in Cloud Trace."); - } - // Generates a new SpanContext w/ random traceId,spanId protected SpanContext getNewSpanContext() { String traceId = generateNewTraceId(); @@ -244,7 +242,8 @@ protected SpanContext getNewSpanContext() { protected Span getNewRootSpanWithContext() { // Execute the DB operation in the context of the custom root span. - return tracer.spanBuilder(rootSpanName) + return tracer + .spanBuilder(rootSpanName) .setParent(Context.root().with(Span.wrap(customSpanContext))) .startSpan(); } @@ -260,19 +259,42 @@ protected void waitForTracesToComplete() throws Exception { completableResultCode.join(TRACE_FORCE_FLUSH_MILLIS, TimeUnit.MILLISECONDS); } - // Validates `retrievedTrace` + // Validates `retrievedTrace`. Cloud Trace indexes traces w/ eventual consistency, even when + // indexing traceId, therefore the + // test must retry a few times before the complete trace is available. protected void fetchAndValidateTraces(String traceId, String... spanNames) throws InterruptedException { - // Fetch traces - retrievedTrace = getTraceWithRetry(projectId, traceId); - List spanNameList = Arrays.asList(spanNames); - - // Validate trace spans - assertEquals(retrievedTrace.getTraceId(), traceId); - assertEquals(retrievedTrace.getSpans(0).getName(), rootSpanName); - for (int i = 0; i < spanNameList.size() ; ++i) { - assertEquals(spanNameList.get(i), retrievedTrace.getSpans(i+1).getName()); - } + int numRetries = GET_TRACE_RETRY_COUNT; + do { + try { + // Fetch traces + retrievedTrace = traceClient_v1.getTrace(projectId, traceId); + ArrayList spanNameList = new ArrayList(Arrays.asList(spanNames)); + spanNameList.add(0, rootSpanName); + // Validate trace spans + assertEquals(traceId, retrievedTrace.getTraceId()); + for (int i = 0; i < spanNameList.size(); ++i) { + assertEquals(spanNameList.get(i), retrievedTrace.getSpans(i).getName()); + } + assertEquals(spanNameList.size(), retrievedTrace.getSpansCount()); + return; + } catch (NotFoundException notFound) { + logger.info("Trace not found, retrying in " + GET_TRACE_RETRY_BACKOFF_MILLIS + " ms"); + Thread.sleep(GET_TRACE_RETRY_BACKOFF_MILLIS); + } catch (IndexOutOfBoundsException outOfBoundsException) { + logger.info( + "Expected # spans: " + + (spanNames.length + 1) + + "Actual # spans: " + + retrievedTrace.getSpansCount()); + Thread.sleep(GET_TRACE_RETRY_BACKOFF_MILLIS); + } + } while (numRetries-- > 0); + throw new RuntimeException( + "Expected spans: " + + spanNames.toString() + + ", Actual spans: " + + (retrievedTrace != null ? retrievedTrace.getSpansList().toString() : "null")); } @Test @@ -292,7 +314,10 @@ public void aggregateQueryGetTraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces(customSpanContext.getTraceId(), "AggregationQuery.Get"); + fetchAndValidateTraces( + customSpanContext.getTraceId(), + "AggregationQuery.Get", + grpcSpanName(RUN_AGGREGATION_QUERY_RPC_NAME)); } @Test @@ -318,7 +343,8 @@ public void bulkWriterCommitTraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces(customSpanContext.getTraceId(), + fetchAndValidateTraces( + customSpanContext.getTraceId(), SPAN_NAME_BULK_WRITER_COMMIT, grpcSpanName(BATCH_WRITE_RPC_NAME)); } @@ -339,7 +365,8 @@ public void partitionQueryTraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces(customSpanContext.getTraceId(), + fetchAndValidateTraces( + customSpanContext.getTraceId(), SPAN_NAME_PARTITION_QUERY, grpcSpanName(SPAN_NAME_PARTITION_QUERY)); } @@ -359,57 +386,385 @@ public void collectionListDocumentsTraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces(customSpanContext.getTraceId(), - SPAN_NAME_COL_REF_LIST_DOCUMENTS, grpcSpanName(LIST_DOCUMENTS_RPC_NAME)); + fetchAndValidateTraces( + customSpanContext.getTraceId(), + SPAN_NAME_COL_REF_LIST_DOCUMENTS, + grpcSpanName(LIST_DOCUMENTS_RPC_NAME)); } - public void docRefCreate() throws Exception { + @Test + public void docRefCreateTraceTest() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore.collection("col").document().create(Collections.singletonMap("foo", "bar")).get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + // Read and validate traces + fetchAndValidateTraces( + customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_CREATE, + SPAN_NAME_BATCH_COMMIT, + grpcSpanName(COMMIT_RPC_NAME)); } @Test - public void docRefCreate2() throws Exception { + public void docRefCreate2TraceTest() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore.collection("col").document().create(new Pojo(1)).get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + // Read and validate traces + fetchAndValidateTraces( + customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_CREATE, + SPAN_NAME_BATCH_COMMIT, + grpcSpanName(COMMIT_RPC_NAME)); } @Test - public void docRefSet() throws Exception {} + public void docRefSetTraceTest() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore.collection("col").document("foo").set(Collections.singletonMap("foo", "bar")).get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + // Read and validate traces + fetchAndValidateTraces( + customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_SET, + SPAN_NAME_BATCH_COMMIT, + grpcSpanName(COMMIT_RPC_NAME)); + } @Test - public void docRefSet2() throws Exception {} + public void docRefSet2TraceTest() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore + .collection("col") + .document("foo") + .set(Collections.singletonMap("foo", "bar"), SetOptions.merge()) + .get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + // Read and validate traces + fetchAndValidateTraces( + customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_SET, + SPAN_NAME_BATCH_COMMIT, + grpcSpanName(COMMIT_RPC_NAME)); + } @Test - public void docRefSet3() throws Exception {} + public void docRefSet3TraceTest() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore.collection("col").document("foo").set(new Pojo(1)).get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); - public void docRefSet4() throws Exception {} + // Read and validate traces + fetchAndValidateTraces( + customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_SET, + SPAN_NAME_BATCH_COMMIT, + grpcSpanName(COMMIT_RPC_NAME)); + } @Test - public void docRefUpdate() throws Exception {} + public void docRefSet4TraceTest() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore.collection("col").document("foo").set(new Pojo(1), SetOptions.merge()).get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + // Read and validate traces + fetchAndValidateTraces( + customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_SET, + SPAN_NAME_BATCH_COMMIT, + grpcSpanName(COMMIT_RPC_NAME)); + } @Test - public void docRefUpdate2() throws Exception {} + public void docRefUpdate() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore + .collection("col") + .document("foo") + .update(Collections.singletonMap("foo", "bar")) + .get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + // Read and validate traces + fetchAndValidateTraces( + customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_UPDATE, + SPAN_NAME_BATCH_COMMIT, + grpcSpanName(COMMIT_RPC_NAME)); + } @Test - public void docRefUpdate3() throws Exception {} + public void docRefUpdate2() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore + .collection("col") + .document("foo") + .update(Collections.singletonMap("foo", "bar"), Precondition.NONE) + .get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + // Read and validate traces + fetchAndValidateTraces( + customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_UPDATE, + SPAN_NAME_BATCH_COMMIT, + grpcSpanName(COMMIT_RPC_NAME)); + } @Test - public void docRefUpdate4() throws Exception {} + public void docRefUpdate3() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore.collection("col").document("foo").update("key", "value", "key2", "value2").get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + // Read and validate traces + fetchAndValidateTraces( + customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_UPDATE, + SPAN_NAME_BATCH_COMMIT, + grpcSpanName(COMMIT_RPC_NAME)); + } @Test - public void docRefUpdate5() throws Exception {} + public void docRefUpdate4() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore + .collection("col") + .document("foo") + .update(FieldPath.of("key"), "value", FieldPath.of("key2"), "value2") + .get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + // Read and validate traces + fetchAndValidateTraces( + customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_UPDATE, + SPAN_NAME_BATCH_COMMIT, + grpcSpanName(COMMIT_RPC_NAME)); + } @Test - public void docRefUpdate6() throws Exception {} + public void docRefUpdate5() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore + .collection("col") + .document("foo") + .update(Precondition.NONE, "key", "value", "key2", "value2") + .get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + // Read and validate traces + fetchAndValidateTraces( + customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_UPDATE, + SPAN_NAME_BATCH_COMMIT, + grpcSpanName(COMMIT_RPC_NAME)); + } @Test - public void docRefDelete() throws Exception {} + public void docRefUpdate6() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore + .collection("col") + .document("foo") + .update(Precondition.NONE, FieldPath.of("key"), "value", FieldPath.of("key2"), "value2") + .get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + // Read and validate traces + fetchAndValidateTraces( + customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_UPDATE, + SPAN_NAME_BATCH_COMMIT, + grpcSpanName(COMMIT_RPC_NAME)); + } @Test - public void docRefDelete2() throws Exception {} + public void docRefDelete() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore.collection("col").document("doc0").delete().get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + // Read and validate traces + fetchAndValidateTraces( + customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_DELETE, + SPAN_NAME_BATCH_COMMIT, + grpcSpanName(COMMIT_RPC_NAME)); + } @Test - public void docRefGet() throws Exception {} + public void docRefDelete2() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore.collection("col").document("doc0").delete(Precondition.NONE).get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + // Read and validate traces + fetchAndValidateTraces( + customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_DELETE, + SPAN_NAME_BATCH_COMMIT, + grpcSpanName(COMMIT_RPC_NAME)); + } @Test - public void docRefGet2() throws Exception {} + public void docRefGet() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore.collection("col").document("doc0").get().get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + // Read and validate traces + fetchAndValidateTraces( + customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_GET, + grpcSpanName(BATCH_GET_DOCUMENTS_RPC_NAME)); + } + + @Test + public void docRefGet2() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore.collection("col").document("doc0").get(FieldMask.of("foo")).get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + // Read and validate traces + fetchAndValidateTraces( + customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_GET, + grpcSpanName(BATCH_GET_DOCUMENTS_RPC_NAME)); + } @Test public void docListCollections() throws Exception {} From cc17d6621db3cd099a798d130e9675ceaec6b489 Mon Sep 17 00:00:00 2001 From: jimit-j-shah Date: Tue, 19 Mar 2024 15:49:52 -0700 Subject: [PATCH 27/34] test: Add get/query Trace Tests --- .../cloud/firestore/it/ITE2ETracingTest.java | 64 ++++++++++++++++++- 1 file changed, 61 insertions(+), 3 deletions(-) diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java index 1e84214be..3431c8833 100644 --- a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java +++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java @@ -22,9 +22,11 @@ import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_DOC_REF_CREATE; import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_DOC_REF_DELETE; import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_DOC_REF_GET; +import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_DOC_REF_LIST_COLLECTIONS; import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_DOC_REF_SET; import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_DOC_REF_UPDATE; import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_PARTITION_QUERY; +import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_QUERY_GET; import static io.opentelemetry.semconv.resource.attributes.ResourceAttributes.SERVICE_NAME; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNotNull; @@ -34,6 +36,7 @@ import com.google.cloud.firestore.BulkWriter; import com.google.cloud.firestore.BulkWriterOptions; import com.google.cloud.firestore.CollectionGroup; +import com.google.cloud.firestore.DocumentReference; import com.google.cloud.firestore.FieldMask; import com.google.cloud.firestore.FieldPath; import com.google.cloud.firestore.Firestore; @@ -91,10 +94,14 @@ public class ITE2ETracingTest extends ITBaseTest { private static final String COMMIT_RPC_NAME = "Commit"; + private static final String LIST_COLLECTIONS_RPC_NAME = "ListCollectionIds"; + private static final String LIST_DOCUMENTS_RPC_NAME = "ListDocuments"; private static final String RUN_AGGREGATION_QUERY_RPC_NAME = "RunAggregationQuery"; + private static final String RUN_QUERY_RPC_NAME = "RunQuery"; + private static final int NUM_TRACE_ID_BYTES = 32; private static final int NUM_SPAN_ID_BYTES = 16; @@ -767,13 +774,64 @@ public void docRefGet2() throws Exception { } @Test - public void docListCollections() throws Exception {} + public void docListCollections() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore.collection("col").document("doc0").listCollections(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + // Read and validate traces + fetchAndValidateTraces( + customSpanContext.getTraceId(), + SPAN_NAME_DOC_REF_LIST_COLLECTIONS, + grpcSpanName(LIST_COLLECTIONS_RPC_NAME)); + } @Test - public void getAll() throws Exception {} + public void getAll() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + DocumentReference docRef0 = firestore.collection("col").document(); + DocumentReference docRef1 = firestore.collection("col").document(); + DocumentReference[] docs = {docRef0, docRef1}; + firestore.getAll(docs).get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + fetchAndValidateTraces( + customSpanContext.getTraceId(), grpcSpanName(BATCH_GET_DOCUMENTS_RPC_NAME)); + } @Test - public void queryGet() throws Exception {} + public void queryGet() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore.collection("col").whereEqualTo("foo", "my_non_existent_value").get().get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + fetchAndValidateTraces( + customSpanContext.getTraceId(), SPAN_NAME_QUERY_GET, grpcSpanName(RUN_QUERY_RPC_NAME)); + } @Test public void transaction() throws Exception {} From 1363df0bb698e483b55ff61fa1e46f4935a4d7cc Mon Sep 17 00:00:00 2001 From: jimit-j-shah Date: Wed, 20 Mar 2024 14:29:43 -0700 Subject: [PATCH 28/34] test: Added Transaction tests using TraceContainer to verify traces for Transaction ops (BeginTransaction, Rollback etc) --- .../cloud/firestore/it/ITE2ETracingTest.java | 368 +++++++++++++++++- 1 file changed, 365 insertions(+), 3 deletions(-) diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java index 3431c8833..e6079ee57 100644 --- a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java +++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java @@ -27,10 +27,19 @@ import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_DOC_REF_UPDATE; import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_PARTITION_QUERY; import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_QUERY_GET; +import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_BEGIN; +import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_COMMIT; +import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_GET_AGGREGATION_QUERY; +import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_GET_DOCUMENTS; +import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_GET_QUERY; +import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_ROLLBACK; +import static com.google.cloud.firestore.telemetry.TraceUtil.SPAN_NAME_TRANSACTION_RUN; import static io.opentelemetry.semconv.resource.attributes.ResourceAttributes.SERVICE_NAME; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; import com.google.api.gax.rpc.NotFoundException; import com.google.cloud.firestore.BulkWriter; @@ -43,13 +52,16 @@ import com.google.cloud.firestore.FirestoreOpenTelemetryOptions; import com.google.cloud.firestore.FirestoreOptions; import com.google.cloud.firestore.Precondition; +import com.google.cloud.firestore.Query; import com.google.cloud.firestore.SetOptions; +import com.google.cloud.firestore.WriteBatch; import com.google.cloud.firestore.it.ITTracingTest.Pojo; import com.google.cloud.opentelemetry.trace.TraceConfiguration; import com.google.cloud.opentelemetry.trace.TraceExporter; import com.google.cloud.trace.v1.TraceServiceClient; import com.google.common.base.Preconditions; import com.google.devtools.cloudtrace.v1.Trace; +import com.google.devtools.cloudtrace.v1.TraceSpan; import io.opentelemetry.api.trace.Span; import io.opentelemetry.api.trace.SpanContext; import io.opentelemetry.api.trace.TraceFlags; @@ -67,7 +79,10 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; +import java.util.List; +import java.util.Map; import java.util.Random; +import java.util.TreeMap; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; @@ -81,9 +96,120 @@ import org.junit.runner.RunWith; import org.junit.runners.JUnit4; +// This End-to-End test verifies Client-side Tracing Functionality instrumented using the +// OpenTelemetry API. +// The test depends on the following external APIs/Services: +// 1. Java OpenTelemetry SDK +// 2. Cloud Trace Exporter +// 3. TraceServiceClient from Cloud Trace API v1. +// +// Permissions required to run this test (https://cloud.google.com/trace/docs/iam#trace-roles): +// 1. gcloud auth application-default login must be run with the test user. +// 2. To write traces, test user must have one of roles/cloudtrace.[admin|agent|user] roles. +// 3. To read traces, test user must have one of roles/cloudtrace.[admin|user] roles. +// +// Each test-case has the following workflow: +// 1. OpenTelemetry SDK is initialized with Cloud Trace Exporter and 100% Trace Sampling +// 2. On initialization, Firestore client is provided the OpenTelemetry SDK object from (1) +// 3. A custom TraceID is generated and injected using a custom SpanContext +// 4. Firestore operations are run inside a root TraceSpan created using the custom SpanContext from +// (3). +// 5. Traces are read-back using TraceServiceClient and verified against expected Call Stacks. +// TODO In the future it would be great to have a single test-driver for this test and +// ITTracingTest. @RunWith(JUnit4.class) public class ITE2ETracingTest extends ITBaseTest { + // Helper class to track call-stacks in a trace + protected class TraceContainer { + + // Maps Span ID to TraceSpan + private Map idSpanMap; + + // Maps Parent Span ID to a list of Child SpanIDs, useful for top-down traversal + private Map> parentChildIdMap; + + // Tracks the Root Span ID + private long rootId; + + public TraceContainer(String rootSpanName, Trace trace) { + idSpanMap = new TreeMap(); + parentChildIdMap = new TreeMap>(); + for (TraceSpan span : trace.getSpansList()) { + long spanId = span.getSpanId(); + idSpanMap.put(spanId, span); + if (rootSpanName.equals(span.getName())) { + rootId = span.getSpanId(); + } + + // Add self as a child of the parent span + if (!parentChildIdMap.containsKey(span.getParentSpanId())) { + parentChildIdMap.put(span.getParentSpanId(), new ArrayList()); + } + parentChildIdMap.get(span.getParentSpanId()).add(spanId); + } + } + + String spanName(long spanId) { + return idSpanMap.get(spanId).getName(); + } + + List childSpans(long spanId) { + return parentChildIdMap.get(spanId); + } + + // This method only works for matching call stacks with traces which have children of distinct + // type at all + // levels. This is good enough as the intention is to validate if the e2e path is WAI - the + // intention is not to validate Cloud Trace's correctness w.r.t. durability of all kinds of + // traces. + boolean containsCallStack(String... callStack) throws RuntimeException { + ArrayList expectedCallStack = new ArrayList(); + for (String call : callStack) { + expectedCallStack.add(call); + } + if (expectedCallStack.isEmpty()) { + throw new RuntimeException("Input callStack is empty"); + } + return dfsContainsCallStack(rootId, expectedCallStack); + } + + // Depth-first check for call stack in the trace + private boolean dfsContainsCallStack(long spanId, List expectedCallStack) { + logger.info( + "span=" + + spanName(spanId) + + ", expectedCallStack[0]=" + + (expectedCallStack.isEmpty() ? "null" : expectedCallStack.get(0))); + if (!expectedCallStack.isEmpty() && spanName(spanId).equals(expectedCallStack.get(0))) { + if (childSpans(spanId) == null) { + logger.info("No more chilren for " + spanName(spanId)); + return true; + } else { + for (Long childSpan : childSpans(spanId)) { + int callStackListSize = expectedCallStack.size(); + logger.info( + "childSpan=" + + spanName(childSpan) + + ", expectedCallStackSize=" + + callStackListSize); + if (dfsContainsCallStack( + childSpan, + expectedCallStack.subList( + /*fromIndexInclusive=*/ 1, /*toIndexExclusive*/ callStackListSize))) { + return true; + } + } + } + } else { + if (!expectedCallStack.isEmpty()) { + logger.warning(spanName(spanId) + " didn't match " + expectedCallStack.get(0)); + } + } + return false; + } + } + private static final Logger logger = Logger.getLogger(ITBaseTest.class.getName()); private static final String SERVICE = "google.firestore.v1.Firestore/"; @@ -92,12 +218,16 @@ public class ITE2ETracingTest extends ITBaseTest { private static final String BATCH_WRITE_RPC_NAME = "BatchWrite"; + private static final String BEGIN_TRANSACTION_RPC_NAME = "BeginTransaction"; + private static final String COMMIT_RPC_NAME = "Commit"; private static final String LIST_COLLECTIONS_RPC_NAME = "ListCollectionIds"; private static final String LIST_DOCUMENTS_RPC_NAME = "ListDocuments"; + private static final String ROLLBACK_RPC_NAME = "Rollback"; + private static final String RUN_AGGREGATION_QUERY_RPC_NAME = "RunAggregationQuery"; private static final String RUN_QUERY_RPC_NAME = "RunQuery"; @@ -266,6 +396,57 @@ protected void waitForTracesToComplete() throws Exception { completableResultCode.join(TRACE_FORCE_FLUSH_MILLIS, TimeUnit.MILLISECONDS); } + // Validates `retrievedTrace`. Cloud Trace indexes traces w/ eventual consistency, even when + // indexing traceId, therefore the + // test must retry a few times before the complete trace is available. + protected void fetchAndValidateTransactionTrace( + String traceId, int numExpectedSpans, String... callStack) throws InterruptedException { + // Large enough count to accommodate eventually consistent Cloud Trace backend + int numRetries = GET_TRACE_RETRY_COUNT; + do { + try { + // Fetch traces + retrievedTrace = traceClient_v1.getTrace(projectId, traceId); + assertEquals(traceId, retrievedTrace.getTraceId()); + + ArrayList expectedCallStack = new ArrayList(Arrays.asList(callStack)); + + // numExpectedSpans should account for rootSpanName (which is not passed in callStack) + expectedCallStack.add(0, rootSpanName); + numExpectedSpans++; + + System.out.println( + "expectedSpanCount=" + + numExpectedSpans + + ", retrievedSpanCount=" + + retrievedTrace.getSpansCount()); + // *Maybe* the full trace was returned + if (retrievedTrace.getSpansCount() == numExpectedSpans) { + System.out.println("Checking if TraceContainer containsCallStack"); + TraceContainer traceContainer = new TraceContainer(rootSpanName, retrievedTrace); + String[] temp = new String[expectedCallStack.size()]; + if (traceContainer.containsCallStack(expectedCallStack.toArray(temp))) { + return; + } + logger.severe("CallStack not found in TraceContainer."); + } // else the trace may not have been fully committed to Cloud Trace storage + } catch (NotFoundException notFound) { + logger.info("Trace not found, retrying in " + GET_TRACE_RETRY_BACKOFF_MILLIS + " ms"); + Thread.sleep(GET_TRACE_RETRY_BACKOFF_MILLIS); + // getTrace could fail + } catch (IndexOutOfBoundsException outOfBoundsException) { + logger.info("Call stack not found in trace. Retrying."); + Thread.sleep(GET_TRACE_RETRY_BACKOFF_MILLIS); + } + // Retrieved trace doesn't have expected number of spans + } while (numRetries-- > 0); + throw new RuntimeException( + "Expected spans: " + + callStack.toString() + + ", Actual spans: " + + (retrievedTrace != null ? retrievedTrace.getSpansList().toString() : "null")); + } + // Validates `retrievedTrace`. Cloud Trace indexes traces w/ eventual consistency, even when // indexing traceId, therefore the // test must retry a few times before the complete trace is available. @@ -304,6 +485,56 @@ protected void fetchAndValidateTraces(String traceId, String... spanNames) + (retrievedTrace != null ? retrievedTrace.getSpansList().toString() : "null")); } + @Test + public void traceContainerTest() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore.collection("col").whereEqualTo("foo", "my_non_existent_value").get().get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + Trace traceResp = null; + int expectedSpanCount = 3; + + int numRetries = GET_TRACE_RETRY_COUNT; + do { + try { + traceResp = traceClient_v1.getTrace(projectId, customSpanContext.getTraceId()); + if (traceResp.getSpansCount() == expectedSpanCount) { + break; + } + } catch (NotFoundException notFoundException) { + Thread.sleep(GET_TRACE_RETRY_BACKOFF_MILLIS); + logger.info("Trace not found, retrying in " + GET_TRACE_RETRY_BACKOFF_MILLIS + " ms"); + } + numRetries--; + } while (numRetries > 0); + + TraceContainer traceCont = new TraceContainer(rootSpanName, traceResp); + + // Contains exact path + assertTrue( + traceCont.containsCallStack( + rootSpanName, SPAN_NAME_QUERY_GET, grpcSpanName(RUN_QUERY_RPC_NAME))); + + // Top-level mismatch + assertFalse(traceCont.containsCallStack(SPAN_NAME_QUERY_GET, RUN_QUERY_RPC_NAME)); + + // Mid-level match + assertFalse(traceCont.containsCallStack(rootSpanName, SPAN_NAME_QUERY_GET)); + + // Leaf-level mismatch/missing + assertFalse( + traceCont.containsCallStack( + rootSpanName, SPAN_NAME_QUERY_GET, RUN_AGGREGATION_QUERY_RPC_NAME)); + } + @Test // Trace an Aggregation.Get request public void aggregateQueryGetTraceTest() throws Exception { @@ -834,11 +1065,142 @@ public void queryGet() throws Exception { } @Test - public void transaction() throws Exception {} + public void transaction() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore + .runTransaction( + transaction -> { + Query q = firestore.collection("col").whereGreaterThan("bla", ""); + DocumentReference d = firestore.collection("col").document("foo"); + DocumentReference[] docList = {d, d}; + // Document Query. + transaction.get(q).get(); + + // Aggregation Query. + transaction.get(q.count()); + + // Get multiple documents. + transaction.getAll(d, d).get(); + + // Commit 2 documents. + transaction.set( + firestore.collection("foo").document("bar"), + Collections.singletonMap("foo", "bar")); + transaction.set( + firestore.collection("foo").document("bar2"), + Collections.singletonMap("foo2", "bar2")); + return 0; + }) + .get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + fetchAndValidateTransactionTrace( + customSpanContext.getTraceId(), + /*numExpectedSpans=*/ 11, + SPAN_NAME_TRANSACTION_RUN, + SPAN_NAME_TRANSACTION_BEGIN, + grpcSpanName(BEGIN_TRANSACTION_RPC_NAME)); + + fetchAndValidateTransactionTrace( + customSpanContext.getTraceId(), + /*numExpectedSpans=*/ 11, + SPAN_NAME_TRANSACTION_RUN, + SPAN_NAME_TRANSACTION_GET_QUERY, + grpcSpanName(RUN_QUERY_RPC_NAME)); + + fetchAndValidateTransactionTrace( + customSpanContext.getTraceId(), + /*numExpectedSpans=*/ 11, + SPAN_NAME_TRANSACTION_RUN, + SPAN_NAME_TRANSACTION_GET_AGGREGATION_QUERY, + grpcSpanName(RUN_AGGREGATION_QUERY_RPC_NAME)); + + fetchAndValidateTransactionTrace( + customSpanContext.getTraceId(), + /*numExpectedSpans=*/ 11, + SPAN_NAME_TRANSACTION_RUN, + SPAN_NAME_TRANSACTION_GET_DOCUMENTS, + grpcSpanName(BATCH_GET_DOCUMENTS_RPC_NAME)); + + fetchAndValidateTransactionTrace( + customSpanContext.getTraceId(), + /*numExpectedSpans=*/ 11, + SPAN_NAME_TRANSACTION_RUN, + SPAN_NAME_TRANSACTION_COMMIT, + grpcSpanName(COMMIT_RPC_NAME)); + } @Test - public void transactionRollback() throws Exception {} + public void transactionRollback() throws Exception { + String myErrorMessage = "My error message."; + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + firestore + .runTransaction( + transaction -> { + if (true) { + throw (new Exception(myErrorMessage)); + } + return 0; + }) + .get(); + } catch (Exception e) { + // Catch and move on. + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + fetchAndValidateTransactionTrace( + customSpanContext.getTraceId(), + /*numExpectedSpans=*/ 5, + SPAN_NAME_TRANSACTION_RUN, + SPAN_NAME_TRANSACTION_BEGIN, + grpcSpanName(BEGIN_TRANSACTION_RPC_NAME)); + + fetchAndValidateTransactionTrace( + customSpanContext.getTraceId(), + /*numExpectedSpans=*/ 5, + SPAN_NAME_TRANSACTION_RUN, + SPAN_NAME_TRANSACTION_ROLLBACK, + grpcSpanName(ROLLBACK_RPC_NAME)); + } @Test - public void writeBatch() throws Exception {} + public void writeBatch() throws Exception { + // Make sure the test has a new SpanContext (and TraceId for injection) + assertNotNull(customSpanContext); + + // Inject new trace ID + Span rootSpan = getNewRootSpanWithContext(); + try (Scope ss = rootSpan.makeCurrent()) { + WriteBatch batch = firestore.batch(); + DocumentReference docRef = firestore.collection("foo").document(); + batch.create(docRef, Collections.singletonMap("foo", "bar")); + batch.update(docRef, Collections.singletonMap("foo", "bar")); + batch.delete(docRef); + batch.commit().get(); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + fetchAndValidateTransactionTrace( + customSpanContext.getTraceId(), + /*numExpectedSpans=*/ 2, + SPAN_NAME_BATCH_COMMIT, + grpcSpanName(COMMIT_RPC_NAME)); + } } From cb49e2190264bcaaec329ccaa0f0a7b80937612d Mon Sep 17 00:00:00 2001 From: jimit-j-shah Date: Thu, 28 Mar 2024 10:57:27 -0700 Subject: [PATCH 29/34] test: Adding TestParameterInjector to run the test for global and non-global opentelemetry SDK instances --- google-cloud-firestore/pom.xml | 6 ++ .../cloud/firestore/it/ITE2ETracingTest.java | 91 ++++++++++++++----- 2 files changed, 72 insertions(+), 25 deletions(-) diff --git a/google-cloud-firestore/pom.xml b/google-cloud-firestore/pom.xml index becd396d4..742b2ee6a 100644 --- a/google-cloud-firestore/pom.xml +++ b/google-cloud-firestore/pom.xml @@ -248,6 +248,12 @@ 1.3.0 test + + com.google.testparameterinjector + test-parameter-injector + 1.15 + test + diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java index e6079ee57..a5b463288 100644 --- a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java +++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java @@ -62,6 +62,9 @@ import com.google.common.base.Preconditions; import com.google.devtools.cloudtrace.v1.Trace; import com.google.devtools.cloudtrace.v1.TraceSpan; +import com.google.testing.junit.testparameterinjector.TestParameter; +import com.google.testing.junit.testparameterinjector.TestParameterInjector; +import io.opentelemetry.api.GlobalOpenTelemetry; import io.opentelemetry.api.trace.Span; import io.opentelemetry.api.trace.SpanContext; import io.opentelemetry.api.trace.TraceFlags; @@ -94,7 +97,6 @@ import org.junit.BeforeClass; import org.junit.Test; import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; // This End-to-End test verifies Client-side Tracing Functionality instrumented using the // OpenTelemetry API. @@ -117,9 +119,13 @@ // 5. Traces are read-back using TraceServiceClient and verified against expected Call Stacks. // TODO In the future it would be great to have a single test-driver for this test and // ITTracingTest. -@RunWith(JUnit4.class) +@RunWith(TestParameterInjector.class) public class ITE2ETracingTest extends ITBaseTest { + protected boolean isUsingGlobalOpenTelemetrySDK() { + return useGlobalOpenTelemetrySDK; + } + // Helper class to track call-stacks in a trace protected class TraceContainer { @@ -268,30 +274,18 @@ private boolean dfsContainsCallStack(long spanId, List expectedCallStack private static Firestore firestore; + @TestParameter boolean useGlobalOpenTelemetrySDK; + @BeforeClass public static void setup() throws IOException { projectId = FirestoreOptions.getDefaultProjectId(); logger.info("projectId:" + projectId); - // Set up OTel SDK - Resource resource = - Resource.getDefault().merge(Resource.builder().put(SERVICE_NAME, "Sparky").build()); - // TODO(jimit) Make it re-usable w/ InMemorySpanExporter traceExporter = TraceExporter.createWithConfiguration( TraceConfiguration.builder().setProjectId(projectId).build()); - openTelemetrySdk = - OpenTelemetrySdk.builder() - .setTracerProvider( - SdkTracerProvider.builder() - .setResource(resource) - .addSpanProcessor(BatchSpanProcessor.builder(traceExporter).build()) - .setSampler(Sampler.alwaysOn()) - .build()) - .build(); - traceClient_v1 = TraceServiceClient.create(); random = new Random(); @@ -299,16 +293,50 @@ public static void setup() throws IOException { @Before public void before() throws Exception { + // Set up OTel SDK + Resource resource = + Resource.getDefault().merge(Resource.builder().put(SERVICE_NAME, "Sparky").build()); + + if (isUsingGlobalOpenTelemetrySDK()) { + openTelemetrySdk = + OpenTelemetrySdk.builder() + .setTracerProvider( + SdkTracerProvider.builder() + .setResource(resource) + .addSpanProcessor(BatchSpanProcessor.builder(traceExporter).build()) + .setSampler(Sampler.alwaysOn()) + .build()) + .buildAndRegisterGlobal(); + } else { + openTelemetrySdk = + OpenTelemetrySdk.builder() + .setTracerProvider( + SdkTracerProvider.builder() + .setResource(resource) + .addSpanProcessor(BatchSpanProcessor.builder(traceExporter).build()) + .setSampler(Sampler.alwaysOn()) + .build()) + .build(); + } + // Initialize the Firestore DB w/ the OTel SDK. Ideally we'd do this is the @BeforeAll method // but because gRPC traces need to be deterministically force-flushed, firestore.shutdown() // must be called in @After for each test. - FirestoreOptions.Builder optionsBuilder = - FirestoreOptions.newBuilder() - .setOpenTelemetryOptions( - FirestoreOpenTelemetryOptions.newBuilder() - .setOpenTelemetry(openTelemetrySdk) - .setTracingEnabled(true) - .build()); + FirestoreOptions.Builder optionsBuilder; + if (isUsingGlobalOpenTelemetrySDK()) { + optionsBuilder = + FirestoreOptions.newBuilder() + .setOpenTelemetryOptions( + FirestoreOpenTelemetryOptions.newBuilder().setTracingEnabled(true).build()); + } else { + optionsBuilder = + FirestoreOptions.newBuilder() + .setOpenTelemetryOptions( + FirestoreOpenTelemetryOptions.newBuilder() + .setOpenTelemetry(openTelemetrySdk) + .setTracingEnabled(true) + .build()); + } String namedDb = System.getProperty("FIRESTORE_NAMED_DATABASE"); if (namedDb != null) { @@ -322,10 +350,20 @@ public void before() throws Exception { firestore, "Error instantiating Firestore. Check that the service account credentials " + "were properly set."); + + // Set up the tracer for custom TraceID injection rootSpanName = String.format("%s%d", this.getClass().getSimpleName(), System.currentTimeMillis()); - tracer = - firestore.getOptions().getOpenTelemetryOptions().getOpenTelemetry().getTracer(rootSpanName); + if (isUsingGlobalOpenTelemetrySDK()) { + tracer = GlobalOpenTelemetry.getTracer(rootSpanName); + } else { + tracer = + firestore + .getOptions() + .getOpenTelemetryOptions() + .getOpenTelemetry() + .getTracer(rootSpanName); + } // Get up a new SpanContext (ergo TraceId) for each test customSpanContext = getNewSpanContext(); @@ -335,6 +373,9 @@ public void before() throws Exception { @After public void after() throws Exception { + if (isUsingGlobalOpenTelemetrySDK()) { + GlobalOpenTelemetry.resetForTest(); + } firestore.shutdown(); rootSpanName = null; tracer = null; From aaa9efedb26541413936e20210dfeb8e00e01cdf Mon Sep 17 00:00:00 2001 From: jimit-j-shah Date: Thu, 28 Mar 2024 12:54:32 -0700 Subject: [PATCH 30/34] test: Formatting and cleanup --- .../cloud/firestore/it/ITE2ETracingTest.java | 92 ++++++++++--------- 1 file changed, 48 insertions(+), 44 deletions(-) diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java index a5b463288..cf93bda4b 100644 --- a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java +++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java @@ -438,8 +438,11 @@ protected void waitForTracesToComplete() throws Exception { } // Validates `retrievedTrace`. Cloud Trace indexes traces w/ eventual consistency, even when - // indexing traceId, therefore the - // test must retry a few times before the complete trace is available. + // indexing traceId, therefore the test must retry a few times before the complete trace is + // available. + // For Transaction traces, there may be more spans than in the trace than specified in + // `callStack`. So `numExpectedSpans` is the expected total number of spans (and not just the + // spans in `callStack`) protected void fetchAndValidateTransactionTrace( String traceId, int numExpectedSpans, String... callStack) throws InterruptedException { // Large enough count to accommodate eventually consistent Cloud Trace backend @@ -489,9 +492,10 @@ protected void fetchAndValidateTransactionTrace( } // Validates `retrievedTrace`. Cloud Trace indexes traces w/ eventual consistency, even when - // indexing traceId, therefore the - // test must retry a few times before the complete trace is available. - protected void fetchAndValidateTraces(String traceId, String... spanNames) + // indexing traceId, therefore the test must retry a few times before the complete trace is + // available. + // For Non-Transaction traces, there is a 1:1 ratio of spans in `spanNames` and in the trace. + protected void fetchAndValidateNonTransactionTrace(String traceId, String... spanNames) throws InterruptedException { int numRetries = GET_TRACE_RETRY_COUNT; do { @@ -593,7 +597,7 @@ public void aggregateQueryGetTraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), "AggregationQuery.Get", grpcSpanName(RUN_AGGREGATION_QUERY_RPC_NAME)); @@ -622,7 +626,7 @@ public void bulkWriterCommitTraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_BULK_WRITER_COMMIT, grpcSpanName(BATCH_WRITE_RPC_NAME)); @@ -644,7 +648,7 @@ public void partitionQueryTraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_PARTITION_QUERY, grpcSpanName(SPAN_NAME_PARTITION_QUERY)); @@ -665,7 +669,7 @@ public void collectionListDocumentsTraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_COL_REF_LIST_DOCUMENTS, grpcSpanName(LIST_DOCUMENTS_RPC_NAME)); @@ -686,7 +690,7 @@ public void docRefCreateTraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_CREATE, SPAN_NAME_BATCH_COMMIT, @@ -708,7 +712,7 @@ public void docRefCreate2TraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_CREATE, SPAN_NAME_BATCH_COMMIT, @@ -730,7 +734,7 @@ public void docRefSetTraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_SET, SPAN_NAME_BATCH_COMMIT, @@ -756,7 +760,7 @@ public void docRefSet2TraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_SET, SPAN_NAME_BATCH_COMMIT, @@ -778,7 +782,7 @@ public void docRefSet3TraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_SET, SPAN_NAME_BATCH_COMMIT, @@ -800,7 +804,7 @@ public void docRefSet4TraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_SET, SPAN_NAME_BATCH_COMMIT, @@ -808,7 +812,7 @@ public void docRefSet4TraceTest() throws Exception { } @Test - public void docRefUpdate() throws Exception { + public void docRefUpdateTraceTest() throws Exception { // Make sure the test has a new SpanContext (and TraceId for injection) assertNotNull(customSpanContext); @@ -826,7 +830,7 @@ public void docRefUpdate() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_UPDATE, SPAN_NAME_BATCH_COMMIT, @@ -834,7 +838,7 @@ public void docRefUpdate() throws Exception { } @Test - public void docRefUpdate2() throws Exception { + public void docRefUpdate2TraceTest() throws Exception { // Make sure the test has a new SpanContext (and TraceId for injection) assertNotNull(customSpanContext); @@ -852,7 +856,7 @@ public void docRefUpdate2() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_UPDATE, SPAN_NAME_BATCH_COMMIT, @@ -860,7 +864,7 @@ public void docRefUpdate2() throws Exception { } @Test - public void docRefUpdate3() throws Exception { + public void docRefUpdate3TraceTest() throws Exception { // Make sure the test has a new SpanContext (and TraceId for injection) assertNotNull(customSpanContext); @@ -874,7 +878,7 @@ public void docRefUpdate3() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_UPDATE, SPAN_NAME_BATCH_COMMIT, @@ -882,7 +886,7 @@ public void docRefUpdate3() throws Exception { } @Test - public void docRefUpdate4() throws Exception { + public void docRefUpdate4TraceTest() throws Exception { // Make sure the test has a new SpanContext (and TraceId for injection) assertNotNull(customSpanContext); @@ -900,7 +904,7 @@ public void docRefUpdate4() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_UPDATE, SPAN_NAME_BATCH_COMMIT, @@ -908,7 +912,7 @@ public void docRefUpdate4() throws Exception { } @Test - public void docRefUpdate5() throws Exception { + public void docRefUpdate5TraceTest() throws Exception { // Make sure the test has a new SpanContext (and TraceId for injection) assertNotNull(customSpanContext); @@ -926,7 +930,7 @@ public void docRefUpdate5() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_UPDATE, SPAN_NAME_BATCH_COMMIT, @@ -934,7 +938,7 @@ public void docRefUpdate5() throws Exception { } @Test - public void docRefUpdate6() throws Exception { + public void docRefUpdate6TraceTest() throws Exception { // Make sure the test has a new SpanContext (and TraceId for injection) assertNotNull(customSpanContext); @@ -952,7 +956,7 @@ public void docRefUpdate6() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_UPDATE, SPAN_NAME_BATCH_COMMIT, @@ -960,7 +964,7 @@ public void docRefUpdate6() throws Exception { } @Test - public void docRefDelete() throws Exception { + public void docRefDeleteTraceTest() throws Exception { // Make sure the test has a new SpanContext (and TraceId for injection) assertNotNull(customSpanContext); @@ -974,7 +978,7 @@ public void docRefDelete() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_DELETE, SPAN_NAME_BATCH_COMMIT, @@ -982,7 +986,7 @@ public void docRefDelete() throws Exception { } @Test - public void docRefDelete2() throws Exception { + public void docRefDelete2TraceTest() throws Exception { // Make sure the test has a new SpanContext (and TraceId for injection) assertNotNull(customSpanContext); @@ -996,7 +1000,7 @@ public void docRefDelete2() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_DELETE, SPAN_NAME_BATCH_COMMIT, @@ -1004,7 +1008,7 @@ public void docRefDelete2() throws Exception { } @Test - public void docRefGet() throws Exception { + public void docRefGetTraceTest() throws Exception { // Make sure the test has a new SpanContext (and TraceId for injection) assertNotNull(customSpanContext); @@ -1018,14 +1022,14 @@ public void docRefGet() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_GET, grpcSpanName(BATCH_GET_DOCUMENTS_RPC_NAME)); } @Test - public void docRefGet2() throws Exception { + public void docRefGet2TraceTest() throws Exception { // Make sure the test has a new SpanContext (and TraceId for injection) assertNotNull(customSpanContext); @@ -1039,14 +1043,14 @@ public void docRefGet2() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_GET, grpcSpanName(BATCH_GET_DOCUMENTS_RPC_NAME)); } @Test - public void docListCollections() throws Exception { + public void docListCollectionsTraceTest() throws Exception { // Make sure the test has a new SpanContext (and TraceId for injection) assertNotNull(customSpanContext); @@ -1060,14 +1064,14 @@ public void docListCollections() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_LIST_COLLECTIONS, grpcSpanName(LIST_COLLECTIONS_RPC_NAME)); } @Test - public void getAll() throws Exception { + public void getAllTraceTest() throws Exception { // Make sure the test has a new SpanContext (and TraceId for injection) assertNotNull(customSpanContext); @@ -1083,12 +1087,12 @@ public void getAll() throws Exception { } waitForTracesToComplete(); - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), grpcSpanName(BATCH_GET_DOCUMENTS_RPC_NAME)); } @Test - public void queryGet() throws Exception { + public void queryGetTraceTest() throws Exception { // Make sure the test has a new SpanContext (and TraceId for injection) assertNotNull(customSpanContext); @@ -1101,12 +1105,12 @@ public void queryGet() throws Exception { } waitForTracesToComplete(); - fetchAndValidateTraces( + fetchAndValidateNonTransactionTrace( customSpanContext.getTraceId(), SPAN_NAME_QUERY_GET, grpcSpanName(RUN_QUERY_RPC_NAME)); } @Test - public void transaction() throws Exception { + public void transactionTraceTest() throws Exception { // Make sure the test has a new SpanContext (and TraceId for injection) assertNotNull(customSpanContext); @@ -1180,7 +1184,7 @@ public void transaction() throws Exception { } @Test - public void transactionRollback() throws Exception { + public void transactionRollbackTraceTest() throws Exception { String myErrorMessage = "My error message."; // Make sure the test has a new SpanContext (and TraceId for injection) assertNotNull(customSpanContext); @@ -1220,7 +1224,7 @@ public void transactionRollback() throws Exception { } @Test - public void writeBatch() throws Exception { + public void writeBatchTraceTest() throws Exception { // Make sure the test has a new SpanContext (and TraceId for injection) assertNotNull(customSpanContext); From 98569a234511dfed9279a3e3424a02d03b0003fb Mon Sep 17 00:00:00 2001 From: jimit-j-shah Date: Mon, 1 Apr 2024 23:30:05 -0700 Subject: [PATCH 31/34] test: review comments --- .../cloud/firestore/it/ITE2ETracingTest.java | 26 ++++++++++--------- 1 file changed, 14 insertions(+), 12 deletions(-) diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java index cf93bda4b..ebe95936c 100644 --- a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java +++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java @@ -165,10 +165,9 @@ List childSpans(long spanId) { } // This method only works for matching call stacks with traces which have children of distinct - // type at all - // levels. This is good enough as the intention is to validate if the e2e path is WAI - the - // intention is not to validate Cloud Trace's correctness w.r.t. durability of all kinds of - // traces. + // type at all levels. This is good enough as the intention is to validate if the e2e path is + // WAI - the intention is not to validate Cloud Trace's correctness w.r.t. durability of all + // kinds of traces. boolean containsCallStack(String... callStack) throws RuntimeException { ArrayList expectedCallStack = new ArrayList(); for (String call : callStack) { @@ -187,11 +186,16 @@ private boolean dfsContainsCallStack(long spanId, List expectedCallStack + spanName(spanId) + ", expectedCallStack[0]=" + (expectedCallStack.isEmpty() ? "null" : expectedCallStack.get(0))); - if (!expectedCallStack.isEmpty() && spanName(spanId).equals(expectedCallStack.get(0))) { + if (expectedCallStack.isEmpty()) { + throw new RuntimeException("Input callStack is empty"); + } + if (spanName(spanId).equals(expectedCallStack.get(0))) { + // Recursion termination if (childSpans(spanId) == null) { logger.info("No more chilren for " + spanName(spanId)); return true; } else { + // Examine the child spans for (Long childSpan : childSpans(spanId)) { int callStackListSize = expectedCallStack.size(); logger.info( @@ -208,15 +212,13 @@ private boolean dfsContainsCallStack(long spanId, List expectedCallStack } } } else { - if (!expectedCallStack.isEmpty()) { - logger.warning(spanName(spanId) + " didn't match " + expectedCallStack.get(0)); - } + logger.info(spanName(spanId) + " didn't match " + expectedCallStack.get(0)); } return false; } } - private static final Logger logger = Logger.getLogger(ITBaseTest.class.getName()); + private static final Logger logger = Logger.getLogger(ITE2ETracingTest.class.getName()); private static final String SERVICE = "google.firestore.v1.Firestore/"; @@ -391,7 +393,7 @@ public static void teardown() throws Exception { completableResultCode.join(TRACE_PROVIDER_SHUTDOWN_MILLIS, TimeUnit.MILLISECONDS); } - // Generates a random 32-byte hex string + // Generates a random hex string of length `numBytes` private String generateRandomHexString(int numBytes) { StringBuffer newTraceId = new StringBuffer(); while (newTraceId.length() < numBytes) { @@ -459,14 +461,14 @@ protected void fetchAndValidateTransactionTrace( expectedCallStack.add(0, rootSpanName); numExpectedSpans++; - System.out.println( + logger.info( "expectedSpanCount=" + numExpectedSpans + ", retrievedSpanCount=" + retrievedTrace.getSpansCount()); // *Maybe* the full trace was returned if (retrievedTrace.getSpansCount() == numExpectedSpans) { - System.out.println("Checking if TraceContainer containsCallStack"); + logger.info("Checking if TraceContainer containsCallStack"); TraceContainer traceContainer = new TraceContainer(rootSpanName, retrievedTrace); String[] temp = new String[expectedCallStack.size()]; if (traceContainer.containsCallStack(expectedCallStack.toArray(temp))) { From ea4b9d6df935bc55689806127fa6ca591baf1995 Mon Sep 17 00:00:00 2001 From: jimit-j-shah Date: Mon, 1 Apr 2024 23:36:15 -0700 Subject: [PATCH 32/34] test: fixing dfs to handle case where the compareTo callstack may be shorter than the trace callstack - don't need to throw an exception in that case --- .../java/com/google/cloud/firestore/it/ITE2ETracingTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java index ebe95936c..ada1cee05 100644 --- a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java +++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java @@ -187,7 +187,7 @@ private boolean dfsContainsCallStack(long spanId, List expectedCallStack + ", expectedCallStack[0]=" + (expectedCallStack.isEmpty() ? "null" : expectedCallStack.get(0))); if (expectedCallStack.isEmpty()) { - throw new RuntimeException("Input callStack is empty"); + return false; } if (spanName(spanId).equals(expectedCallStack.get(0))) { // Recursion termination From 2b0a1f71c25a1de3a0f6ec371a1b9902933a5a98 Mon Sep 17 00:00:00 2001 From: jimit-j-shah Date: Tue, 2 Apr 2024 16:16:54 -0700 Subject: [PATCH 33/34] test: Consolidating verification methods --- .../cloud/firestore/it/ITE2ETracingTest.java | 238 ++++++++---------- 1 file changed, 110 insertions(+), 128 deletions(-) diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java index ada1cee05..ad5200c28 100644 --- a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java +++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java @@ -445,91 +445,81 @@ protected void waitForTracesToComplete() throws Exception { // For Transaction traces, there may be more spans than in the trace than specified in // `callStack`. So `numExpectedSpans` is the expected total number of spans (and not just the // spans in `callStack`) - protected void fetchAndValidateTransactionTrace( - String traceId, int numExpectedSpans, String... callStack) throws InterruptedException { + protected void fetchAndValidateTrace( + String traceId, int numExpectedSpans, List> callStackList) + throws InterruptedException { // Large enough count to accommodate eventually consistent Cloud Trace backend int numRetries = GET_TRACE_RETRY_COUNT; + // Account for rootSpanName + numExpectedSpans++; + + // Fetch traces do { try { - // Fetch traces retrievedTrace = traceClient_v1.getTrace(projectId, traceId); assertEquals(traceId, retrievedTrace.getTraceId()); - ArrayList expectedCallStack = new ArrayList(Arrays.asList(callStack)); - - // numExpectedSpans should account for rootSpanName (which is not passed in callStack) - expectedCallStack.add(0, rootSpanName); - numExpectedSpans++; - logger.info( "expectedSpanCount=" + numExpectedSpans + ", retrievedSpanCount=" + retrievedTrace.getSpansCount()); - // *Maybe* the full trace was returned - if (retrievedTrace.getSpansCount() == numExpectedSpans) { - logger.info("Checking if TraceContainer containsCallStack"); - TraceContainer traceContainer = new TraceContainer(rootSpanName, retrievedTrace); - String[] temp = new String[expectedCallStack.size()]; - if (traceContainer.containsCallStack(expectedCallStack.toArray(temp))) { - return; - } - logger.severe("CallStack not found in TraceContainer."); - } // else the trace may not have been fully committed to Cloud Trace storage } catch (NotFoundException notFound) { logger.info("Trace not found, retrying in " + GET_TRACE_RETRY_BACKOFF_MILLIS + " ms"); - Thread.sleep(GET_TRACE_RETRY_BACKOFF_MILLIS); - // getTrace could fail } catch (IndexOutOfBoundsException outOfBoundsException) { logger.info("Call stack not found in trace. Retrying."); + } + if (retrievedTrace == null || numExpectedSpans != retrievedTrace.getSpansCount()) { Thread.sleep(GET_TRACE_RETRY_BACKOFF_MILLIS); } - // Retrieved trace doesn't have expected number of spans - } while (numRetries-- > 0); - throw new RuntimeException( - "Expected spans: " - + callStack.toString() - + ", Actual spans: " - + (retrievedTrace != null ? retrievedTrace.getSpansList().toString() : "null")); + } while (numRetries-- > 0 + && (retrievedTrace == null || numExpectedSpans != retrievedTrace.getSpansCount())); + + if (retrievedTrace == null || numExpectedSpans != retrievedTrace.getSpansCount()) { + throw new RuntimeException( + "Expected number of spans: " + + numExpectedSpans + + ", Actual number of spans: " + + (retrievedTrace != null + ? retrievedTrace.getSpansList().toString() + : "Trace NOT_FOUND")); + } + + TraceContainer traceContainer = new TraceContainer(rootSpanName, retrievedTrace); + + for (List callStack : callStackList) { + // Update all call stacks to be rooted at rootSpanName + ArrayList expectedCallStack = new ArrayList<>(callStack); + + // numExpectedSpans should account for rootSpanName (not passed in callStackList) + expectedCallStack.add(0, rootSpanName); + + // *May be* the full trace was returned + logger.info("Checking if TraceContainer contains the callStack"); + String[] expectedCallList = new String[expectedCallStack.size()]; + if (!traceContainer.containsCallStack(expectedCallStack.toArray(expectedCallList))) { + throw new RuntimeException( + "Expected spans: " + + expectedCallList.toString() + + ", Actual spans: " + + (retrievedTrace != null + ? retrievedTrace.getSpansList().toString() + : "Trace NOT_FOUND")); + } + logger.severe("CallStack not found in TraceContainer."); + } } // Validates `retrievedTrace`. Cloud Trace indexes traces w/ eventual consistency, even when // indexing traceId, therefore the test must retry a few times before the complete trace is // available. // For Non-Transaction traces, there is a 1:1 ratio of spans in `spanNames` and in the trace. - protected void fetchAndValidateNonTransactionTrace(String traceId, String... spanNames) + protected void fetchAndValidateTrace(String traceId, String... spanNames) throws InterruptedException { int numRetries = GET_TRACE_RETRY_COUNT; - do { - try { - // Fetch traces - retrievedTrace = traceClient_v1.getTrace(projectId, traceId); - ArrayList spanNameList = new ArrayList(Arrays.asList(spanNames)); - spanNameList.add(0, rootSpanName); - // Validate trace spans - assertEquals(traceId, retrievedTrace.getTraceId()); - for (int i = 0; i < spanNameList.size(); ++i) { - assertEquals(spanNameList.get(i), retrievedTrace.getSpans(i).getName()); - } - assertEquals(spanNameList.size(), retrievedTrace.getSpansCount()); - return; - } catch (NotFoundException notFound) { - logger.info("Trace not found, retrying in " + GET_TRACE_RETRY_BACKOFF_MILLIS + " ms"); - Thread.sleep(GET_TRACE_RETRY_BACKOFF_MILLIS); - } catch (IndexOutOfBoundsException outOfBoundsException) { - logger.info( - "Expected # spans: " - + (spanNames.length + 1) - + "Actual # spans: " - + retrievedTrace.getSpansCount()); - Thread.sleep(GET_TRACE_RETRY_BACKOFF_MILLIS); - } - } while (numRetries-- > 0); - throw new RuntimeException( - "Expected spans: " - + spanNames.toString() - + ", Actual spans: " - + (retrievedTrace != null ? retrievedTrace.getSpansList().toString() : "null")); + int numSpans = spanNames.length; + + fetchAndValidateTrace(traceId, numSpans, Arrays.asList(Arrays.asList(spanNames))); } @Test @@ -599,7 +589,7 @@ public void aggregateQueryGetTraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateNonTransactionTrace( + fetchAndValidateTrace( customSpanContext.getTraceId(), "AggregationQuery.Get", grpcSpanName(RUN_AGGREGATION_QUERY_RPC_NAME)); @@ -628,7 +618,7 @@ public void bulkWriterCommitTraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateNonTransactionTrace( + fetchAndValidateTrace( customSpanContext.getTraceId(), SPAN_NAME_BULK_WRITER_COMMIT, grpcSpanName(BATCH_WRITE_RPC_NAME)); @@ -650,7 +640,7 @@ public void partitionQueryTraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateNonTransactionTrace( + fetchAndValidateTrace( customSpanContext.getTraceId(), SPAN_NAME_PARTITION_QUERY, grpcSpanName(SPAN_NAME_PARTITION_QUERY)); @@ -671,7 +661,7 @@ public void collectionListDocumentsTraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateNonTransactionTrace( + fetchAndValidateTrace( customSpanContext.getTraceId(), SPAN_NAME_COL_REF_LIST_DOCUMENTS, grpcSpanName(LIST_DOCUMENTS_RPC_NAME)); @@ -692,7 +682,7 @@ public void docRefCreateTraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateNonTransactionTrace( + fetchAndValidateTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_CREATE, SPAN_NAME_BATCH_COMMIT, @@ -714,7 +704,7 @@ public void docRefCreate2TraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateNonTransactionTrace( + fetchAndValidateTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_CREATE, SPAN_NAME_BATCH_COMMIT, @@ -736,7 +726,7 @@ public void docRefSetTraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateNonTransactionTrace( + fetchAndValidateTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_SET, SPAN_NAME_BATCH_COMMIT, @@ -762,7 +752,7 @@ public void docRefSet2TraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateNonTransactionTrace( + fetchAndValidateTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_SET, SPAN_NAME_BATCH_COMMIT, @@ -784,7 +774,7 @@ public void docRefSet3TraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateNonTransactionTrace( + fetchAndValidateTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_SET, SPAN_NAME_BATCH_COMMIT, @@ -806,7 +796,7 @@ public void docRefSet4TraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateNonTransactionTrace( + fetchAndValidateTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_SET, SPAN_NAME_BATCH_COMMIT, @@ -832,7 +822,7 @@ public void docRefUpdateTraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateNonTransactionTrace( + fetchAndValidateTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_UPDATE, SPAN_NAME_BATCH_COMMIT, @@ -858,7 +848,7 @@ public void docRefUpdate2TraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateNonTransactionTrace( + fetchAndValidateTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_UPDATE, SPAN_NAME_BATCH_COMMIT, @@ -880,7 +870,7 @@ public void docRefUpdate3TraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateNonTransactionTrace( + fetchAndValidateTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_UPDATE, SPAN_NAME_BATCH_COMMIT, @@ -906,7 +896,7 @@ public void docRefUpdate4TraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateNonTransactionTrace( + fetchAndValidateTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_UPDATE, SPAN_NAME_BATCH_COMMIT, @@ -932,7 +922,7 @@ public void docRefUpdate5TraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateNonTransactionTrace( + fetchAndValidateTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_UPDATE, SPAN_NAME_BATCH_COMMIT, @@ -958,7 +948,7 @@ public void docRefUpdate6TraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateNonTransactionTrace( + fetchAndValidateTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_UPDATE, SPAN_NAME_BATCH_COMMIT, @@ -980,7 +970,7 @@ public void docRefDeleteTraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateNonTransactionTrace( + fetchAndValidateTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_DELETE, SPAN_NAME_BATCH_COMMIT, @@ -1002,7 +992,7 @@ public void docRefDelete2TraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateNonTransactionTrace( + fetchAndValidateTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_DELETE, SPAN_NAME_BATCH_COMMIT, @@ -1024,7 +1014,7 @@ public void docRefGetTraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateNonTransactionTrace( + fetchAndValidateTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_GET, grpcSpanName(BATCH_GET_DOCUMENTS_RPC_NAME)); @@ -1045,7 +1035,7 @@ public void docRefGet2TraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateNonTransactionTrace( + fetchAndValidateTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_GET, grpcSpanName(BATCH_GET_DOCUMENTS_RPC_NAME)); @@ -1066,7 +1056,7 @@ public void docListCollectionsTraceTest() throws Exception { waitForTracesToComplete(); // Read and validate traces - fetchAndValidateNonTransactionTrace( + fetchAndValidateTrace( customSpanContext.getTraceId(), SPAN_NAME_DOC_REF_LIST_COLLECTIONS, grpcSpanName(LIST_COLLECTIONS_RPC_NAME)); @@ -1089,7 +1079,7 @@ public void getAllTraceTest() throws Exception { } waitForTracesToComplete(); - fetchAndValidateNonTransactionTrace( + fetchAndValidateTrace( customSpanContext.getTraceId(), grpcSpanName(BATCH_GET_DOCUMENTS_RPC_NAME)); } @@ -1107,7 +1097,7 @@ public void queryGetTraceTest() throws Exception { } waitForTracesToComplete(); - fetchAndValidateNonTransactionTrace( + fetchAndValidateTrace( customSpanContext.getTraceId(), SPAN_NAME_QUERY_GET, grpcSpanName(RUN_QUERY_RPC_NAME)); } @@ -1149,40 +1139,34 @@ public void transactionTraceTest() throws Exception { } waitForTracesToComplete(); - fetchAndValidateTransactionTrace( - customSpanContext.getTraceId(), - /*numExpectedSpans=*/ 11, - SPAN_NAME_TRANSACTION_RUN, - SPAN_NAME_TRANSACTION_BEGIN, - grpcSpanName(BEGIN_TRANSACTION_RPC_NAME)); - - fetchAndValidateTransactionTrace( - customSpanContext.getTraceId(), - /*numExpectedSpans=*/ 11, - SPAN_NAME_TRANSACTION_RUN, - SPAN_NAME_TRANSACTION_GET_QUERY, - grpcSpanName(RUN_QUERY_RPC_NAME)); - - fetchAndValidateTransactionTrace( - customSpanContext.getTraceId(), - /*numExpectedSpans=*/ 11, - SPAN_NAME_TRANSACTION_RUN, - SPAN_NAME_TRANSACTION_GET_AGGREGATION_QUERY, - grpcSpanName(RUN_AGGREGATION_QUERY_RPC_NAME)); - - fetchAndValidateTransactionTrace( - customSpanContext.getTraceId(), - /*numExpectedSpans=*/ 11, - SPAN_NAME_TRANSACTION_RUN, - SPAN_NAME_TRANSACTION_GET_DOCUMENTS, - grpcSpanName(BATCH_GET_DOCUMENTS_RPC_NAME)); - - fetchAndValidateTransactionTrace( + fetchAndValidateTrace( customSpanContext.getTraceId(), /*numExpectedSpans=*/ 11, - SPAN_NAME_TRANSACTION_RUN, - SPAN_NAME_TRANSACTION_COMMIT, - grpcSpanName(COMMIT_RPC_NAME)); + Arrays.asList( + Arrays.asList( + SPAN_NAME_TRANSACTION_RUN, + SPAN_NAME_TRANSACTION_BEGIN, + grpcSpanName(BEGIN_TRANSACTION_RPC_NAME)), + Arrays.asList( + SPAN_NAME_TRANSACTION_RUN, + SPAN_NAME_TRANSACTION_BEGIN, + grpcSpanName(BEGIN_TRANSACTION_RPC_NAME)), + Arrays.asList( + SPAN_NAME_TRANSACTION_RUN, + SPAN_NAME_TRANSACTION_GET_QUERY, + grpcSpanName(RUN_QUERY_RPC_NAME)), + Arrays.asList( + SPAN_NAME_TRANSACTION_RUN, + SPAN_NAME_TRANSACTION_GET_AGGREGATION_QUERY, + grpcSpanName(RUN_AGGREGATION_QUERY_RPC_NAME)), + Arrays.asList( + SPAN_NAME_TRANSACTION_RUN, + SPAN_NAME_TRANSACTION_GET_DOCUMENTS, + grpcSpanName(BATCH_GET_DOCUMENTS_RPC_NAME)), + Arrays.asList( + SPAN_NAME_TRANSACTION_RUN, + SPAN_NAME_TRANSACTION_COMMIT, + grpcSpanName(COMMIT_RPC_NAME)))); } @Test @@ -1210,19 +1194,18 @@ public void transactionRollbackTraceTest() throws Exception { } waitForTracesToComplete(); - fetchAndValidateTransactionTrace( + fetchAndValidateTrace( customSpanContext.getTraceId(), /*numExpectedSpans=*/ 5, - SPAN_NAME_TRANSACTION_RUN, - SPAN_NAME_TRANSACTION_BEGIN, - grpcSpanName(BEGIN_TRANSACTION_RPC_NAME)); - - fetchAndValidateTransactionTrace( - customSpanContext.getTraceId(), - /*numExpectedSpans=*/ 5, - SPAN_NAME_TRANSACTION_RUN, - SPAN_NAME_TRANSACTION_ROLLBACK, - grpcSpanName(ROLLBACK_RPC_NAME)); + Arrays.asList( + Arrays.asList( + SPAN_NAME_TRANSACTION_RUN, + SPAN_NAME_TRANSACTION_BEGIN, + grpcSpanName(BEGIN_TRANSACTION_RPC_NAME)), + Arrays.asList( + SPAN_NAME_TRANSACTION_RUN, + SPAN_NAME_TRANSACTION_ROLLBACK, + grpcSpanName(ROLLBACK_RPC_NAME)))); } @Test @@ -1244,10 +1227,9 @@ public void writeBatchTraceTest() throws Exception { } waitForTracesToComplete(); - fetchAndValidateTransactionTrace( + fetchAndValidateTrace( customSpanContext.getTraceId(), /*numExpectedSpans=*/ 2, - SPAN_NAME_BATCH_COMMIT, - grpcSpanName(COMMIT_RPC_NAME)); + Arrays.asList(Arrays.asList(SPAN_NAME_BATCH_COMMIT, grpcSpanName(COMMIT_RPC_NAME)))); } } From 1a7f67e810122226f82262eb4c5d49d812ad7e5d Mon Sep 17 00:00:00 2001 From: jimit-j-shah Date: Mon, 1 Apr 2024 23:30:05 -0700 Subject: [PATCH 34/34] test: review comments --- .../java/com/google/cloud/firestore/it/ITE2ETracingTest.java | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java index ad5200c28..3fb8496ff 100644 --- a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java +++ b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITE2ETracingTest.java @@ -169,10 +169,7 @@ List childSpans(long spanId) { // WAI - the intention is not to validate Cloud Trace's correctness w.r.t. durability of all // kinds of traces. boolean containsCallStack(String... callStack) throws RuntimeException { - ArrayList expectedCallStack = new ArrayList(); - for (String call : callStack) { - expectedCallStack.add(call); - } + List expectedCallStack = Arrays.asList(callStack); if (expectedCallStack.isEmpty()) { throw new RuntimeException("Input callStack is empty"); }