From af564e621f572e0b19f9fadd62e15bf977cb8582 Mon Sep 17 00:00:00 2001 From: jimit-j-shah <57637300+jimit-j-shah@users.noreply.github.com> Date: Fri, 10 May 2024 16:27:28 -0700 Subject: [PATCH 01/17] feat: Adding Lookup RPC OpenTelemetry Tracing - Removed OpenCensus Tracing - Added E2E tests with Global and Local OTel SDK --- google-cloud-datastore/pom.xml | 22 + .../google/cloud/datastore/DatastoreImpl.java | 27 +- .../telemetry/DisabledTraceUtil.java | 5 + .../datastore/telemetry/EnabledTraceUtil.java | 6 + .../cloud/datastore/telemetry/TraceUtil.java | 4 + .../cloud/datastore/it/ITE2ETracingTest.java | 549 ++++++++++++++++++ .../it/ITE2ETracingTestGlobalOtel.java | 27 + .../it/ITE2ETracingTestNonGlobalOtel.java | 27 + 8 files changed, 662 insertions(+), 5 deletions(-) create mode 100644 google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java create mode 100644 google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTestGlobalOtel.java create mode 100644 google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTestNonGlobalOtel.java diff --git a/google-cloud-datastore/pom.xml b/google-cloud-datastore/pom.xml index ced4d4069..e6bc6bcb1 100644 --- a/google-cloud-datastore/pom.xml +++ b/google-cloud-datastore/pom.xml @@ -178,12 +178,34 @@ 1.4.2 test + io.opentelemetry opentelemetry-sdk ${opentelemetry.version} test + + io.opentelemetry + opentelemetry-sdk-testing + ${opentelemetry.version} + test + + + + + com.google.api.grpc + proto-google-cloud-trace-v1 + 1.3.0 + test + + + com.google.cloud.opentelemetry + exporter-trace + 0.15.0 + test + + diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java index a3bfb3796..401e3ed54 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java @@ -29,6 +29,7 @@ import com.google.common.base.Preconditions; import com.google.common.collect.AbstractIterator; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; import com.google.common.collect.Iterables; import com.google.common.collect.Sets; import com.google.datastore.v1.ExplainOptions; @@ -61,6 +62,9 @@ final class DatastoreImpl extends BaseService implements Datas TransactionOperationExceptionHandler.build(); private final TraceUtil traceUtil = TraceUtil.getInstance(); + private final com.google.cloud.datastore.telemetry.TraceUtil otelTraceUtil = + getOptions().getTraceUtil(); + private final ReadOptionProtoPreparer readOptionProtoPreparer; private final AggregationQueryExecutor aggregationQueryExecutor; @@ -450,13 +454,26 @@ protected Entity computeNext() { com.google.datastore.v1.LookupResponse lookup( final com.google.datastore.v1.LookupRequest requestPb) { - Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_LOOKUP); - try (Scope scope = traceUtil.getTracer().withSpan(span)) { + com.google.cloud.datastore.telemetry.TraceUtil.Span span = + otelTraceUtil.startSpan(com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_LOOKUP); + final ReadOptions readOptions = requestPb.getReadOptions(); + span.setAttribute( + "isTransactional", (readOptions.hasTransaction() || readOptions.hasNewTransaction())); + + try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = span.makeCurrent()) { return RetryHelper.runWithRetries( new Callable() { @Override public com.google.datastore.v1.LookupResponse call() throws DatastoreException { - return datastoreRpc.lookup(requestPb); + com.google.datastore.v1.LookupResponse response = datastoreRpc.lookup(requestPb); + span.addEvent( + com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_LOOKUP + ": Completed", + new ImmutableMap.Builder() + .put("Received", response.getFoundCount()) + .put("Missing", response.getMissingCount()) + .put("Deferred", response.getDeferredCount()) + .build()); + return response; } }, retrySettings, @@ -465,10 +482,10 @@ public com.google.datastore.v1.LookupResponse call() throws DatastoreException { : TRANSACTION_OPERATION_EXCEPTION_HANDLER, getOptions().getClock()); } catch (RetryHelperException e) { - span.setStatus(Status.UNKNOWN.withDescription(e.getMessage())); + span.end(e); throw DatastoreException.translateAndThrow(e); } finally { - span.end(TraceUtil.END_SPAN_OPTIONS); + span.end(); } } diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/DisabledTraceUtil.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/DisabledTraceUtil.java index 21321897d..a4f25813a 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/DisabledTraceUtil.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/DisabledTraceUtil.java @@ -61,6 +61,11 @@ public TraceUtil.Span setAttribute(String key, String value) { return this; } + @Override + public TraceUtil.Span setAttribute(String key, boolean value) { + return this; + } + @Override public Scope makeCurrent() { return new Scope(); diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/EnabledTraceUtil.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/EnabledTraceUtil.java index 8a5b0b36e..3bf6a7466 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/EnabledTraceUtil.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/EnabledTraceUtil.java @@ -175,6 +175,12 @@ public TraceUtil.Span setAttribute(String key, String value) { return this; } + @Override + public TraceUtil.Span setAttribute(String key, boolean value) { + span.setAttribute(ATTRIBUTE_SERVICE_PREFIX + key, value); + return this; + } + @Override public Scope makeCurrent() { try (io.opentelemetry.context.Scope scope = span.makeCurrent()) { diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java index 4e55a1ff6..c867a5525 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java @@ -32,6 +32,7 @@ public interface TraceUtil { static final String ENABLE_TRACING_ENV_VAR = "DATASTORE_ENABLE_TRACING"; static final String LIBRARY_NAME = "com.google.cloud.datastore"; + static final String SPAN_NAME_LOOKUP = "Lookup"; /** * Creates and returns an instance of the TraceUtil class. * @@ -80,6 +81,9 @@ interface Span { /** Adds the given attribute to this span. */ Span setAttribute(String key, String value); + /** Adds the given attribute to this span. */ + Span setAttribute(String key, boolean value); + /** Marks this span as the current span. */ Scope makeCurrent(); diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java new file mode 100644 index 000000000..c1bf76250 --- /dev/null +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java @@ -0,0 +1,549 @@ +/* + * 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.datastore.it; + +import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_LOOKUP; +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.datastore.Datastore; +import com.google.cloud.datastore.DatastoreOpenTelemetryOptions; +import com.google.cloud.datastore.DatastoreOptions; +import com.google.cloud.datastore.Entity; +import com.google.cloud.datastore.Key; +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.GlobalOpenTelemetry; +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.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Map; +import java.util.Random; +import java.util.TreeMap; +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; + +// 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, Datastore client is provided the OpenTelemetry SDK object from (1) +// 3. A custom TraceID is generated and injected using a custom SpanContext +// 4. Datastore 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. +public abstract class ITE2ETracingTest { + + protected abstract boolean isUsingGlobalOpenTelemetrySDK(); + + // Helper class to track call-stacks in a trace + protected static class TraceContainer { + + // Maps Span ID to TraceSpan + private final Map idSpanMap; + + // Maps Parent Span ID to a list of Child SpanIDs, useful for top-down traversal + private final 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 { + List expectedCallStack = Arrays.asList(callStack); + 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()) { + return false; + } + if (spanName(spanId).equals(expectedCallStack.get(0))) { + // Recursion termination + if (childSpans(spanId) == null) { + logger.info("No more children for " + spanName(spanId)); + return expectedCallStack.size() <= 1; + } else { + // Examine the child spans + 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 { + logger.info(spanName(spanId) + " didn't match " + expectedCallStack.get(0)); + } + return false; + } + } + + private static final Logger logger = Logger.getLogger(ITE2ETracingTest.class.getName()); + + private static final String SERVICE = "google.datastore.v1.Datastore/"; + + 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; + + private static final int GET_TRACE_RETRY_COUNT = 60; + + private static final int GET_TRACE_RETRY_BACKOFF_MILLIS = 1000; + + private static final int TRACE_FORCE_FLUSH_MILLIS = 5000; + + private static final int TRACE_PROVIDER_SHUTDOWN_MILLIS = 1000; + + private static Key KEY1; + + // 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; + + // 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; + private static Tracer tracer; + + // Required to set custom-root span + private static OpenTelemetrySdk openTelemetrySdk; + + private static String projectId; + + private static DatastoreOptions options; + + private static Datastore datastore; + + @BeforeClass + public static void setup() throws IOException { + projectId = DatastoreOptions.getDefaultProjectId(); + logger.info("projectId:" + projectId); + + // TODO(jimit) Make it re-usable w/ InMemorySpanExporter + traceExporter = + TraceExporter.createWithConfiguration( + TraceConfiguration.builder().setProjectId(projectId).build()); + + traceClient_v1 = TraceServiceClient.create(); + + random = new Random(); + } + + @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 Datastore DB w/ the OTel SDK. Ideally we'd do this is the @BeforeAll method + // but because gRPC traces need to be deterministically force-flushed, datastore.shutdown() + // must be called in @After for each test. + DatastoreOptions.Builder optionsBuilder; + if (isUsingGlobalOpenTelemetrySDK()) { + optionsBuilder = + DatastoreOptions.newBuilder() + .setOpenTelemetryOptions( + DatastoreOpenTelemetryOptions.newBuilder().setTracingEnabled(true).build()); + } else { + optionsBuilder = + DatastoreOptions.newBuilder() + .setOpenTelemetryOptions( + DatastoreOpenTelemetryOptions.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."); + } + optionsBuilder.setDatabaseId(namedDb); + options = optionsBuilder.build(); + datastore = options.getService(); + Preconditions.checkNotNull( + datastore, + "Error instantiating Datastore. Check that the service account credentials " + + "were properly set."); + + String PROJECT_ID = options.getProjectId(); + String KIND1 = "kind1"; + KEY1 = Key.newBuilder(PROJECT_ID, KIND1, "name", options.getDatabaseId()).build(); + + // Set up the tracer for custom TraceID injection + rootSpanName = + String.format("%s%d", this.getClass().getSimpleName(), System.currentTimeMillis()); + if (isUsingGlobalOpenTelemetrySDK()) { + tracer = GlobalOpenTelemetry.getTracer(rootSpanName); + } else { + tracer = + datastore + .getOptions() + .getOpenTelemetryOptions() + .getOpenTelemetry() + .getTracer(rootSpanName); + } + + // Get up a new SpanContext (ergo TraceId) for each test + customSpanContext = getNewSpanContext(); + assertNotNull(customSpanContext); + assertNull(retrievedTrace); + } + + @After + public void after() throws Exception { + if (isUsingGlobalOpenTelemetrySDK()) { + GlobalOpenTelemetry.resetForTest(); + } + rootSpanName = null; + tracer = null; + retrievedTrace = null; + customSpanContext = null; + } + + @AfterClass + public static void teardown() throws Exception { + traceClient_v1.close(); + CompletableResultCode completableResultCode = + openTelemetrySdk.getSdkTracerProvider().shutdown(); + completableResultCode.join(TRACE_PROVIDER_SHUTDOWN_MILLIS, TimeUnit.MILLISECONDS); + } + + // Generates a random hex string of length `numBytes` + 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); + } + + // 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 Span getNewRootSpanWithContext() { + // Execute the DB operation in the context of the custom root span. + return tracer + .spanBuilder(rootSpanName) + .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); + } + + // 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 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 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 { + retrievedTrace = traceClient_v1.getTrace(projectId, traceId); + assertEquals(traceId, retrievedTrace.getTraceId()); + + logger.info( + "expectedSpanCount=" + + numExpectedSpans + + ", retrievedSpanCount=" + + retrievedTrace.getSpansCount()); + } catch (NotFoundException notFound) { + logger.info("Trace not found, retrying in " + GET_TRACE_RETRY_BACKOFF_MILLIS + " ms"); + } 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); + } + } 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: " + + Arrays.toString(expectedCallList) + + ", 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 fetchAndValidateTrace(String traceId, String... spanNames) + throws InterruptedException { + fetchAndValidateTrace(traceId, spanNames.length, Arrays.asList(Arrays.asList(spanNames))); + } + + @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 ignored = rootSpan.makeCurrent()) { + Entity entity = datastore.get(KEY1); + assertNull(entity); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + Trace traceResp = null; + int expectedSpanCount = 2; + + int numRetries = GET_TRACE_RETRY_COUNT; + do { + try { + traceResp = traceClient_v1.getTrace(projectId, customSpanContext.getTraceId()); + if (traceResp.getSpansCount() == expectedSpanCount) { + logger.info("Success: Got " + expectedSpanCount + " spans."); + break; + } + } catch (NotFoundException notFoundException) { + Thread.sleep(GET_TRACE_RETRY_BACKOFF_MILLIS); + logger.info("Trace not found, retrying in " + GET_TRACE_RETRY_BACKOFF_MILLIS + " ms"); + } + logger.info( + "Trace Found. The trace did not contain " + + expectedSpanCount + + " spans. Going to retry."); + numRetries--; + } while (numRetries > 0); + + // Make sure we got as many spans as we expected. + assertNotNull(traceResp); + assertEquals(expectedSpanCount, traceResp.getSpansCount()); + + TraceContainer traceCont = new TraceContainer(rootSpanName, traceResp); + + // Contains exact path + assertTrue(traceCont.containsCallStack(rootSpanName, SPAN_NAME_LOOKUP)); + + // Top-level mismatch + assertFalse(traceCont.containsCallStack(SPAN_NAME_LOOKUP, RUN_QUERY_RPC_NAME)); + + // Leaf-level mismatch/missing + assertFalse( + traceCont.containsCallStack( + rootSpanName, SPAN_NAME_LOOKUP, RUN_AGGREGATION_QUERY_RPC_NAME)); + } + + @Test + public void lookupTraceTest() 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 ignored = rootSpan.makeCurrent()) { + Entity entity = datastore.get(KEY1); + assertNull(entity); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + fetchAndValidateTrace(customSpanContext.getTraceId(), SPAN_NAME_LOOKUP); + } +} diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTestGlobalOtel.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTestGlobalOtel.java new file mode 100644 index 000000000..032d4b971 --- /dev/null +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTestGlobalOtel.java @@ -0,0 +1,27 @@ +/* + * 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.datastore.it; + +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class ITE2ETracingTestGlobalOtel extends ITE2ETracingTest { + @Override + protected boolean isUsingGlobalOpenTelemetrySDK() { + return true; + } +} diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTestNonGlobalOtel.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTestNonGlobalOtel.java new file mode 100644 index 000000000..b4f7e9c8d --- /dev/null +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTestNonGlobalOtel.java @@ -0,0 +1,27 @@ +/* + * 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.datastore.it; + +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class ITE2ETracingTestNonGlobalOtel extends ITE2ETracingTest { + @Override + protected boolean isUsingGlobalOpenTelemetrySDK() { + return false; + } +} From 89940418b8697743de5c4c3e9b697b1a69bafff5 Mon Sep 17 00:00:00 2001 From: jimit-j-shah <57637300+jimit-j-shah@users.noreply.github.com> Date: Fri, 10 May 2024 16:38:05 -0700 Subject: [PATCH 02/17] Fixing dependencies --- google-cloud-datastore/pom.xml | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/google-cloud-datastore/pom.xml b/google-cloud-datastore/pom.xml index e6bc6bcb1..b43d6d097 100644 --- a/google-cloud-datastore/pom.xml +++ b/google-cloud-datastore/pom.xml @@ -187,10 +187,21 @@ io.opentelemetry - opentelemetry-sdk-testing + opentelemetry-sdk-common ${opentelemetry.version} test + + io.opentelemetry + opentelemetry-sdk-trace + ${opentelemetry.version} + + + io.opentelemetry + opentelemetry-semconv + 1.1.0-alpha + test + @@ -205,6 +216,12 @@ 0.15.0 test + + com.google.cloud + google-cloud-trace + 1.3.0 + test + From 6f6ad7b1a1a45f0170abeb82aec773601cc06326 Mon Sep 17 00:00:00 2001 From: jimit-j-shah <57637300+jimit-j-shah@users.noreply.github.com> Date: Fri, 10 May 2024 16:27:28 -0700 Subject: [PATCH 03/17] feat: Adding Lookup RPC OpenTelemetry Tracing - Removed OpenCensus Tracing - Added E2E tests with Global and Local OTel SDK --- google-cloud-datastore/pom.xml | 39 ++ .../google/cloud/datastore/DatastoreImpl.java | 27 +- .../telemetry/DisabledTraceUtil.java | 5 + .../datastore/telemetry/EnabledTraceUtil.java | 6 + .../cloud/datastore/telemetry/TraceUtil.java | 4 + .../cloud/datastore/it/ITE2ETracingTest.java | 549 ++++++++++++++++++ .../it/ITE2ETracingTestGlobalOtel.java | 27 + .../it/ITE2ETracingTestNonGlobalOtel.java | 27 + 8 files changed, 679 insertions(+), 5 deletions(-) create mode 100644 google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java create mode 100644 google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTestGlobalOtel.java create mode 100644 google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTestNonGlobalOtel.java diff --git a/google-cloud-datastore/pom.xml b/google-cloud-datastore/pom.xml index ced4d4069..b43d6d097 100644 --- a/google-cloud-datastore/pom.xml +++ b/google-cloud-datastore/pom.xml @@ -178,12 +178,51 @@ 1.4.2 test + io.opentelemetry opentelemetry-sdk ${opentelemetry.version} test + + io.opentelemetry + opentelemetry-sdk-common + ${opentelemetry.version} + test + + + io.opentelemetry + opentelemetry-sdk-trace + ${opentelemetry.version} + + + io.opentelemetry + opentelemetry-semconv + 1.1.0-alpha + 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.cloud + google-cloud-trace + 1.3.0 + test + + diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java index a3bfb3796..401e3ed54 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java @@ -29,6 +29,7 @@ import com.google.common.base.Preconditions; import com.google.common.collect.AbstractIterator; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; import com.google.common.collect.Iterables; import com.google.common.collect.Sets; import com.google.datastore.v1.ExplainOptions; @@ -61,6 +62,9 @@ final class DatastoreImpl extends BaseService implements Datas TransactionOperationExceptionHandler.build(); private final TraceUtil traceUtil = TraceUtil.getInstance(); + private final com.google.cloud.datastore.telemetry.TraceUtil otelTraceUtil = + getOptions().getTraceUtil(); + private final ReadOptionProtoPreparer readOptionProtoPreparer; private final AggregationQueryExecutor aggregationQueryExecutor; @@ -450,13 +454,26 @@ protected Entity computeNext() { com.google.datastore.v1.LookupResponse lookup( final com.google.datastore.v1.LookupRequest requestPb) { - Span span = traceUtil.startSpan(TraceUtil.SPAN_NAME_LOOKUP); - try (Scope scope = traceUtil.getTracer().withSpan(span)) { + com.google.cloud.datastore.telemetry.TraceUtil.Span span = + otelTraceUtil.startSpan(com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_LOOKUP); + final ReadOptions readOptions = requestPb.getReadOptions(); + span.setAttribute( + "isTransactional", (readOptions.hasTransaction() || readOptions.hasNewTransaction())); + + try (com.google.cloud.datastore.telemetry.TraceUtil.Scope ignored = span.makeCurrent()) { return RetryHelper.runWithRetries( new Callable() { @Override public com.google.datastore.v1.LookupResponse call() throws DatastoreException { - return datastoreRpc.lookup(requestPb); + com.google.datastore.v1.LookupResponse response = datastoreRpc.lookup(requestPb); + span.addEvent( + com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_LOOKUP + ": Completed", + new ImmutableMap.Builder() + .put("Received", response.getFoundCount()) + .put("Missing", response.getMissingCount()) + .put("Deferred", response.getDeferredCount()) + .build()); + return response; } }, retrySettings, @@ -465,10 +482,10 @@ public com.google.datastore.v1.LookupResponse call() throws DatastoreException { : TRANSACTION_OPERATION_EXCEPTION_HANDLER, getOptions().getClock()); } catch (RetryHelperException e) { - span.setStatus(Status.UNKNOWN.withDescription(e.getMessage())); + span.end(e); throw DatastoreException.translateAndThrow(e); } finally { - span.end(TraceUtil.END_SPAN_OPTIONS); + span.end(); } } diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/DisabledTraceUtil.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/DisabledTraceUtil.java index 21321897d..a4f25813a 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/DisabledTraceUtil.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/DisabledTraceUtil.java @@ -61,6 +61,11 @@ public TraceUtil.Span setAttribute(String key, String value) { return this; } + @Override + public TraceUtil.Span setAttribute(String key, boolean value) { + return this; + } + @Override public Scope makeCurrent() { return new Scope(); diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/EnabledTraceUtil.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/EnabledTraceUtil.java index 8a5b0b36e..3bf6a7466 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/EnabledTraceUtil.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/EnabledTraceUtil.java @@ -175,6 +175,12 @@ public TraceUtil.Span setAttribute(String key, String value) { return this; } + @Override + public TraceUtil.Span setAttribute(String key, boolean value) { + span.setAttribute(ATTRIBUTE_SERVICE_PREFIX + key, value); + return this; + } + @Override public Scope makeCurrent() { try (io.opentelemetry.context.Scope scope = span.makeCurrent()) { diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java index 4e55a1ff6..c867a5525 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TraceUtil.java @@ -32,6 +32,7 @@ public interface TraceUtil { static final String ENABLE_TRACING_ENV_VAR = "DATASTORE_ENABLE_TRACING"; static final String LIBRARY_NAME = "com.google.cloud.datastore"; + static final String SPAN_NAME_LOOKUP = "Lookup"; /** * Creates and returns an instance of the TraceUtil class. * @@ -80,6 +81,9 @@ interface Span { /** Adds the given attribute to this span. */ Span setAttribute(String key, String value); + /** Adds the given attribute to this span. */ + Span setAttribute(String key, boolean value); + /** Marks this span as the current span. */ Scope makeCurrent(); diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java new file mode 100644 index 000000000..c1bf76250 --- /dev/null +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java @@ -0,0 +1,549 @@ +/* + * 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.datastore.it; + +import static com.google.cloud.datastore.telemetry.TraceUtil.SPAN_NAME_LOOKUP; +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.datastore.Datastore; +import com.google.cloud.datastore.DatastoreOpenTelemetryOptions; +import com.google.cloud.datastore.DatastoreOptions; +import com.google.cloud.datastore.Entity; +import com.google.cloud.datastore.Key; +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.GlobalOpenTelemetry; +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.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Map; +import java.util.Random; +import java.util.TreeMap; +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; + +// 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, Datastore client is provided the OpenTelemetry SDK object from (1) +// 3. A custom TraceID is generated and injected using a custom SpanContext +// 4. Datastore 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. +public abstract class ITE2ETracingTest { + + protected abstract boolean isUsingGlobalOpenTelemetrySDK(); + + // Helper class to track call-stacks in a trace + protected static class TraceContainer { + + // Maps Span ID to TraceSpan + private final Map idSpanMap; + + // Maps Parent Span ID to a list of Child SpanIDs, useful for top-down traversal + private final 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 { + List expectedCallStack = Arrays.asList(callStack); + 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()) { + return false; + } + if (spanName(spanId).equals(expectedCallStack.get(0))) { + // Recursion termination + if (childSpans(spanId) == null) { + logger.info("No more children for " + spanName(spanId)); + return expectedCallStack.size() <= 1; + } else { + // Examine the child spans + 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 { + logger.info(spanName(spanId) + " didn't match " + expectedCallStack.get(0)); + } + return false; + } + } + + private static final Logger logger = Logger.getLogger(ITE2ETracingTest.class.getName()); + + private static final String SERVICE = "google.datastore.v1.Datastore/"; + + 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; + + private static final int GET_TRACE_RETRY_COUNT = 60; + + private static final int GET_TRACE_RETRY_BACKOFF_MILLIS = 1000; + + private static final int TRACE_FORCE_FLUSH_MILLIS = 5000; + + private static final int TRACE_PROVIDER_SHUTDOWN_MILLIS = 1000; + + private static Key KEY1; + + // 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; + + // 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; + private static Tracer tracer; + + // Required to set custom-root span + private static OpenTelemetrySdk openTelemetrySdk; + + private static String projectId; + + private static DatastoreOptions options; + + private static Datastore datastore; + + @BeforeClass + public static void setup() throws IOException { + projectId = DatastoreOptions.getDefaultProjectId(); + logger.info("projectId:" + projectId); + + // TODO(jimit) Make it re-usable w/ InMemorySpanExporter + traceExporter = + TraceExporter.createWithConfiguration( + TraceConfiguration.builder().setProjectId(projectId).build()); + + traceClient_v1 = TraceServiceClient.create(); + + random = new Random(); + } + + @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 Datastore DB w/ the OTel SDK. Ideally we'd do this is the @BeforeAll method + // but because gRPC traces need to be deterministically force-flushed, datastore.shutdown() + // must be called in @After for each test. + DatastoreOptions.Builder optionsBuilder; + if (isUsingGlobalOpenTelemetrySDK()) { + optionsBuilder = + DatastoreOptions.newBuilder() + .setOpenTelemetryOptions( + DatastoreOpenTelemetryOptions.newBuilder().setTracingEnabled(true).build()); + } else { + optionsBuilder = + DatastoreOptions.newBuilder() + .setOpenTelemetryOptions( + DatastoreOpenTelemetryOptions.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."); + } + optionsBuilder.setDatabaseId(namedDb); + options = optionsBuilder.build(); + datastore = options.getService(); + Preconditions.checkNotNull( + datastore, + "Error instantiating Datastore. Check that the service account credentials " + + "were properly set."); + + String PROJECT_ID = options.getProjectId(); + String KIND1 = "kind1"; + KEY1 = Key.newBuilder(PROJECT_ID, KIND1, "name", options.getDatabaseId()).build(); + + // Set up the tracer for custom TraceID injection + rootSpanName = + String.format("%s%d", this.getClass().getSimpleName(), System.currentTimeMillis()); + if (isUsingGlobalOpenTelemetrySDK()) { + tracer = GlobalOpenTelemetry.getTracer(rootSpanName); + } else { + tracer = + datastore + .getOptions() + .getOpenTelemetryOptions() + .getOpenTelemetry() + .getTracer(rootSpanName); + } + + // Get up a new SpanContext (ergo TraceId) for each test + customSpanContext = getNewSpanContext(); + assertNotNull(customSpanContext); + assertNull(retrievedTrace); + } + + @After + public void after() throws Exception { + if (isUsingGlobalOpenTelemetrySDK()) { + GlobalOpenTelemetry.resetForTest(); + } + rootSpanName = null; + tracer = null; + retrievedTrace = null; + customSpanContext = null; + } + + @AfterClass + public static void teardown() throws Exception { + traceClient_v1.close(); + CompletableResultCode completableResultCode = + openTelemetrySdk.getSdkTracerProvider().shutdown(); + completableResultCode.join(TRACE_PROVIDER_SHUTDOWN_MILLIS, TimeUnit.MILLISECONDS); + } + + // Generates a random hex string of length `numBytes` + 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); + } + + // 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 Span getNewRootSpanWithContext() { + // Execute the DB operation in the context of the custom root span. + return tracer + .spanBuilder(rootSpanName) + .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); + } + + // 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 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 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 { + retrievedTrace = traceClient_v1.getTrace(projectId, traceId); + assertEquals(traceId, retrievedTrace.getTraceId()); + + logger.info( + "expectedSpanCount=" + + numExpectedSpans + + ", retrievedSpanCount=" + + retrievedTrace.getSpansCount()); + } catch (NotFoundException notFound) { + logger.info("Trace not found, retrying in " + GET_TRACE_RETRY_BACKOFF_MILLIS + " ms"); + } 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); + } + } 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: " + + Arrays.toString(expectedCallList) + + ", 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 fetchAndValidateTrace(String traceId, String... spanNames) + throws InterruptedException { + fetchAndValidateTrace(traceId, spanNames.length, Arrays.asList(Arrays.asList(spanNames))); + } + + @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 ignored = rootSpan.makeCurrent()) { + Entity entity = datastore.get(KEY1); + assertNull(entity); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + Trace traceResp = null; + int expectedSpanCount = 2; + + int numRetries = GET_TRACE_RETRY_COUNT; + do { + try { + traceResp = traceClient_v1.getTrace(projectId, customSpanContext.getTraceId()); + if (traceResp.getSpansCount() == expectedSpanCount) { + logger.info("Success: Got " + expectedSpanCount + " spans."); + break; + } + } catch (NotFoundException notFoundException) { + Thread.sleep(GET_TRACE_RETRY_BACKOFF_MILLIS); + logger.info("Trace not found, retrying in " + GET_TRACE_RETRY_BACKOFF_MILLIS + " ms"); + } + logger.info( + "Trace Found. The trace did not contain " + + expectedSpanCount + + " spans. Going to retry."); + numRetries--; + } while (numRetries > 0); + + // Make sure we got as many spans as we expected. + assertNotNull(traceResp); + assertEquals(expectedSpanCount, traceResp.getSpansCount()); + + TraceContainer traceCont = new TraceContainer(rootSpanName, traceResp); + + // Contains exact path + assertTrue(traceCont.containsCallStack(rootSpanName, SPAN_NAME_LOOKUP)); + + // Top-level mismatch + assertFalse(traceCont.containsCallStack(SPAN_NAME_LOOKUP, RUN_QUERY_RPC_NAME)); + + // Leaf-level mismatch/missing + assertFalse( + traceCont.containsCallStack( + rootSpanName, SPAN_NAME_LOOKUP, RUN_AGGREGATION_QUERY_RPC_NAME)); + } + + @Test + public void lookupTraceTest() 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 ignored = rootSpan.makeCurrent()) { + Entity entity = datastore.get(KEY1); + assertNull(entity); + } finally { + rootSpan.end(); + } + waitForTracesToComplete(); + + fetchAndValidateTrace(customSpanContext.getTraceId(), SPAN_NAME_LOOKUP); + } +} diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTestGlobalOtel.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTestGlobalOtel.java new file mode 100644 index 000000000..032d4b971 --- /dev/null +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTestGlobalOtel.java @@ -0,0 +1,27 @@ +/* + * 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.datastore.it; + +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class ITE2ETracingTestGlobalOtel extends ITE2ETracingTest { + @Override + protected boolean isUsingGlobalOpenTelemetrySDK() { + return true; + } +} diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTestNonGlobalOtel.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTestNonGlobalOtel.java new file mode 100644 index 000000000..b4f7e9c8d --- /dev/null +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTestNonGlobalOtel.java @@ -0,0 +1,27 @@ +/* + * 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.datastore.it; + +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class ITE2ETracingTestNonGlobalOtel extends ITE2ETracingTest { + @Override + protected boolean isUsingGlobalOpenTelemetrySDK() { + return false; + } +} From e6f5ae9bf0b32eabf6f9dc8be0969344b9f2f5b3 Mon Sep 17 00:00:00 2001 From: jimit-j-shah <57637300+jimit-j-shah@users.noreply.github.com> Date: Sat, 11 May 2024 00:21:43 -0700 Subject: [PATCH 04/17] Fixing depenency --- google-cloud-datastore/pom.xml | 1 + 1 file changed, 1 insertion(+) diff --git a/google-cloud-datastore/pom.xml b/google-cloud-datastore/pom.xml index b43d6d097..8aefcfe8f 100644 --- a/google-cloud-datastore/pom.xml +++ b/google-cloud-datastore/pom.xml @@ -195,6 +195,7 @@ io.opentelemetry opentelemetry-sdk-trace ${opentelemetry.version} + test io.opentelemetry From 766199f223ed9e135278fff09378c1c7641cea6b Mon Sep 17 00:00:00 2001 From: jimit-j-shah <57637300+jimit-j-shah@users.noreply.github.com> Date: Mon, 13 May 2024 09:44:08 -0700 Subject: [PATCH 05/17] Code review: Add BOM as dependency --- google-cloud-datastore/pom.xml | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/google-cloud-datastore/pom.xml b/google-cloud-datastore/pom.xml index 8aefcfe8f..b92a8798a 100644 --- a/google-cloud-datastore/pom.xml +++ b/google-cloud-datastore/pom.xml @@ -18,6 +18,18 @@ google-cloud-datastore 1.37.0 + + + + + com.google.cloud + gapic-libraries-bom + 1.37.0 + pom + import + + + com.google.api.grpc @@ -136,7 +148,6 @@ test-jar test - com.google.guava guava-testlib @@ -205,12 +216,6 @@ - - com.google.api.grpc - proto-google-cloud-trace-v1 - 1.3.0 - test - com.google.cloud.opentelemetry exporter-trace From 5a836d09133cb24f54beb39bd7dec06c86751451 Mon Sep 17 00:00:00 2001 From: jimit-j-shah <57637300+jimit-j-shah@users.noreply.github.com> Date: Mon, 13 May 2024 09:46:48 -0700 Subject: [PATCH 06/17] Code Review: s/FIRESTORE/DATASTORE --- .../java/com/google/cloud/datastore/it/ITE2ETracingTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java index c1bf76250..ca2441b9f 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java @@ -290,7 +290,7 @@ public void before() throws Exception { .build()); } - String namedDb = System.getProperty("FIRESTORE_NAMED_DATABASE"); + String namedDb = System.getProperty("DATASTORE_NAMED_DATABASE"); if (namedDb != null) { logger.log(Level.INFO, "Integration test using named database " + namedDb); optionsBuilder = optionsBuilder.setDatabaseId(namedDb); From 83dc603aa1949976281f9580d828dd474a484d79 Mon Sep 17 00:00:00 2001 From: jimit-j-shah <57637300+jimit-j-shah@users.noreply.github.com> Date: Mon, 13 May 2024 10:24:49 -0700 Subject: [PATCH 07/17] Code review: rename ITE2ETracingTest.java AbstractITE2ETracingTest.java --- ...{ITE2ETracingTest.java => AbstractITE2ETracingTest.java} | 6 ++---- .../cloud/datastore/it/ITE2ETracingTestGlobalOtel.java | 2 +- .../cloud/datastore/it/ITE2ETracingTestNonGlobalOtel.java | 2 +- 3 files changed, 4 insertions(+), 6 deletions(-) rename google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/{ITE2ETracingTest.java => AbstractITE2ETracingTest.java} (98%) diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/AbstractITE2ETracingTest.java similarity index 98% rename from google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java rename to google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/AbstractITE2ETracingTest.java index ca2441b9f..ff114cfad 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/AbstractITE2ETracingTest.java @@ -85,9 +85,7 @@ // 4. Datastore 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. -public abstract class ITE2ETracingTest { +public abstract class AbstractITE2ETracingTest { protected abstract boolean isUsingGlobalOpenTelemetrySDK(); @@ -180,7 +178,7 @@ private boolean dfsContainsCallStack(long spanId, List expectedCallStack } } - private static final Logger logger = Logger.getLogger(ITE2ETracingTest.class.getName()); + private static final Logger logger = Logger.getLogger(AbstractITE2ETracingTest.class.getName()); private static final String SERVICE = "google.datastore.v1.Datastore/"; diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTestGlobalOtel.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTestGlobalOtel.java index 032d4b971..65fc3d7b6 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTestGlobalOtel.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTestGlobalOtel.java @@ -19,7 +19,7 @@ import org.junit.runners.JUnit4; @RunWith(JUnit4.class) -public class ITE2ETracingTestGlobalOtel extends ITE2ETracingTest { +public class ITE2ETracingTestGlobalOtel extends AbstractITE2ETracingTest { @Override protected boolean isUsingGlobalOpenTelemetrySDK() { return true; diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTestNonGlobalOtel.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTestNonGlobalOtel.java index b4f7e9c8d..02ba88e21 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTestNonGlobalOtel.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTestNonGlobalOtel.java @@ -19,7 +19,7 @@ import org.junit.runners.JUnit4; @RunWith(JUnit4.class) -public class ITE2ETracingTestNonGlobalOtel extends ITE2ETracingTest { +public class ITE2ETracingTestNonGlobalOtel extends AbstractITE2ETracingTest { @Override protected boolean isUsingGlobalOpenTelemetrySDK() { return false; From 8d6d2dad0774e36ec67fa00eda4063e237fa4a33 Mon Sep 17 00:00:00 2001 From: jimit-j-shah <57637300+jimit-j-shah@users.noreply.github.com> Date: Mon, 13 May 2024 10:56:01 -0700 Subject: [PATCH 08/17] Code review: fixing pom.xml --- google-cloud-datastore/pom.xml | 6 ------ 1 file changed, 6 deletions(-) diff --git a/google-cloud-datastore/pom.xml b/google-cloud-datastore/pom.xml index b92a8798a..50ef88e83 100644 --- a/google-cloud-datastore/pom.xml +++ b/google-cloud-datastore/pom.xml @@ -222,12 +222,6 @@ 0.15.0 test - - com.google.cloud - google-cloud-trace - 1.3.0 - test - From 9cfdfb602d5dd691192833a58af1c692407eff9a Mon Sep 17 00:00:00 2001 From: jimit-j-shah <57637300+jimit-j-shah@users.noreply.github.com> Date: Mon, 13 May 2024 11:13:42 -0700 Subject: [PATCH 09/17] Code review: fixing dependencies --- google-cloud-datastore/pom.xml | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/google-cloud-datastore/pom.xml b/google-cloud-datastore/pom.xml index 50ef88e83..f89b0d81a 100644 --- a/google-cloud-datastore/pom.xml +++ b/google-cloud-datastore/pom.xml @@ -222,6 +222,14 @@ 0.15.0 test + + com.google.cloud + google-cloud-trace + + + com.google.api.grpc + proto-google-cloud-trace-v1 + From 4442c8132f8e9694fc729a0ce1174a20e68d235c Mon Sep 17 00:00:00 2001 From: jimit-j-shah <57637300+jimit-j-shah@users.noreply.github.com> Date: Mon, 13 May 2024 11:21:46 -0700 Subject: [PATCH 10/17] Code review: pom dependency scope --- google-cloud-datastore/pom.xml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/google-cloud-datastore/pom.xml b/google-cloud-datastore/pom.xml index f89b0d81a..a8f135b23 100644 --- a/google-cloud-datastore/pom.xml +++ b/google-cloud-datastore/pom.xml @@ -225,10 +225,12 @@ com.google.cloud google-cloud-trace + test com.google.api.grpc proto-google-cloud-trace-v1 + test From ce973ddb88f7081d5de992f79626641e154d52b1 Mon Sep 17 00:00:00 2001 From: jimit-j-shah <57637300+jimit-j-shah@users.noreply.github.com> Date: Mon, 13 May 2024 13:38:10 -0700 Subject: [PATCH 11/17] Code review: Parameterize ITE2ETracingTest.java --- google-cloud-datastore/pom.xml | 6 +++++ ...TracingTest.java => ITE2ETracingTest.java} | 21 ++++++++------- .../it/ITE2ETracingTestGlobalOtel.java | 27 ------------------- .../it/ITE2ETracingTestNonGlobalOtel.java | 27 ------------------- 4 files changed, 18 insertions(+), 63 deletions(-) rename google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/{AbstractITE2ETracingTest.java => ITE2ETracingTest.java} (97%) delete mode 100644 google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTestGlobalOtel.java delete mode 100644 google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTestNonGlobalOtel.java diff --git a/google-cloud-datastore/pom.xml b/google-cloud-datastore/pom.xml index a8f135b23..a73c87251 100644 --- a/google-cloud-datastore/pom.xml +++ b/google-cloud-datastore/pom.xml @@ -189,6 +189,12 @@ 1.4.2 test + + com.google.testparameterinjector + test-parameter-injector + 1.16 + test + io.opentelemetry diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/AbstractITE2ETracingTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java similarity index 97% rename from google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/AbstractITE2ETracingTest.java rename to google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java index ff114cfad..4a513b644 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/AbstractITE2ETracingTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java @@ -36,6 +36,8 @@ 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; @@ -65,6 +67,7 @@ import org.junit.Before; import org.junit.BeforeClass; import org.junit.Test; +import org.junit.runner.RunWith; // This End-to-End test verifies Client-side Tracing Functionality instrumented using the // OpenTelemetry API. @@ -85,9 +88,11 @@ // 4. Datastore 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. -public abstract class AbstractITE2ETracingTest { - - protected abstract boolean isUsingGlobalOpenTelemetrySDK(); +@RunWith(TestParameterInjector.class) +public class ITE2ETracingTest { + protected boolean isUsingGlobalOpenTelemetrySDK() { + return useGlobalOpenTelemetrySDK; + } // Helper class to track call-stacks in a trace protected static class TraceContainer { @@ -178,7 +183,7 @@ private boolean dfsContainsCallStack(long spanId, List expectedCallStack } } - private static final Logger logger = Logger.getLogger(AbstractITE2ETracingTest.class.getName()); + private static final Logger logger = Logger.getLogger(ITE2ETracingTest.class.getName()); private static final String SERVICE = "google.datastore.v1.Datastore/"; @@ -226,6 +231,8 @@ private boolean dfsContainsCallStack(long spanId, List expectedCallStack private static Datastore datastore; + @TestParameter boolean useGlobalOpenTelemetrySDK; + @BeforeClass public static void setup() throws IOException { projectId = DatastoreOptions.getDefaultProjectId(); @@ -348,7 +355,7 @@ public static void teardown() throws Exception { // Generates a random hex string of length `numBytes` private String generateRandomHexString(int numBytes) { - StringBuffer newTraceId = new StringBuffer(); + StringBuilder newTraceId = new StringBuilder(); while (newTraceId.length() < numBytes) { newTraceId.append(Integer.toHexString(random.nextInt())); } @@ -381,10 +388,6 @@ protected Span getNewRootSpanWithContext() { .startSpan(); } - protected String grpcSpanName(String rpcName) { - return "Sent." + SERVICE + rpcName; - } - protected void waitForTracesToComplete() throws Exception { logger.info("Flushing traces..."); CompletableResultCode completableResultCode = diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTestGlobalOtel.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTestGlobalOtel.java deleted file mode 100644 index 65fc3d7b6..000000000 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTestGlobalOtel.java +++ /dev/null @@ -1,27 +0,0 @@ -/* - * 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.datastore.it; - -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; - -@RunWith(JUnit4.class) -public class ITE2ETracingTestGlobalOtel extends AbstractITE2ETracingTest { - @Override - protected boolean isUsingGlobalOpenTelemetrySDK() { - return true; - } -} diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTestNonGlobalOtel.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTestNonGlobalOtel.java deleted file mode 100644 index 02ba88e21..000000000 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTestNonGlobalOtel.java +++ /dev/null @@ -1,27 +0,0 @@ -/* - * 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.datastore.it; - -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; - -@RunWith(JUnit4.class) -public class ITE2ETracingTestNonGlobalOtel extends AbstractITE2ETracingTest { - @Override - protected boolean isUsingGlobalOpenTelemetrySDK() { - return false; - } -} From e37c54195aac54b39c1e8c66c97f4599cd86f9d6 Mon Sep 17 00:00:00 2001 From: jimit-j-shah <57637300+jimit-j-shah@users.noreply.github.com> Date: Tue, 14 May 2024 10:46:13 -0700 Subject: [PATCH 12/17] Code review: parameterizing db names --- .../google/cloud/datastore/it/ITE2ETracingTest.java | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java index 4a513b644..924e924c4 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java @@ -94,6 +94,10 @@ protected boolean isUsingGlobalOpenTelemetrySDK() { return useGlobalOpenTelemetrySDK; } + protected String datastoreNamedDatabase() { + return datastoreNamedDatabase; + } + // Helper class to track call-stacks in a trace protected static class TraceContainer { @@ -233,6 +237,8 @@ private boolean dfsContainsCallStack(long spanId, List expectedCallStack @TestParameter boolean useGlobalOpenTelemetrySDK; + @TestParameter({"default", "test-db"}) String datastoreNamedDatabase; + @BeforeClass public static void setup() throws IOException { projectId = DatastoreOptions.getDefaultProjectId(); @@ -295,8 +301,8 @@ public void before() throws Exception { .build()); } - String namedDb = System.getProperty("DATASTORE_NAMED_DATABASE"); - if (namedDb != null) { + String namedDb = datastoreNamedDatabase(); + if (!namedDb.equals("default")) { logger.log(Level.INFO, "Integration test using named database " + namedDb); optionsBuilder = optionsBuilder.setDatabaseId(namedDb); } else { From dbb510835809241e8255b8f26f72fac112e092b5 Mon Sep 17 00:00:00 2001 From: jimit-j-shah <57637300+jimit-j-shah@users.noreply.github.com> Date: Tue, 14 May 2024 10:51:58 -0700 Subject: [PATCH 13/17] formatting --- .../java/com/google/cloud/datastore/it/ITE2ETracingTest.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java index 924e924c4..550ceda4d 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java @@ -237,7 +237,8 @@ private boolean dfsContainsCallStack(long spanId, List expectedCallStack @TestParameter boolean useGlobalOpenTelemetrySDK; - @TestParameter({"default", "test-db"}) String datastoreNamedDatabase; + @TestParameter({"default", "test-db"}) + String datastoreNamedDatabase; @BeforeClass public static void setup() throws IOException { From 153b11fb4840f56f11aaf94d8de20d70711dbd71 Mon Sep 17 00:00:00 2001 From: jimit-j-shah <57637300+jimit-j-shah@users.noreply.github.com> Date: Tue, 14 May 2024 22:07:41 -0700 Subject: [PATCH 14/17] code review: style --- .../com/google/cloud/datastore/it/ITE2ETracingTest.java | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java index 550ceda4d..a25600335 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java @@ -317,9 +317,9 @@ public void before() throws Exception { "Error instantiating Datastore. Check that the service account credentials " + "were properly set."); - String PROJECT_ID = options.getProjectId(); - String KIND1 = "kind1"; - KEY1 = Key.newBuilder(PROJECT_ID, KIND1, "name", options.getDatabaseId()).build(); + String projectId = options.getProjectId(); + String kind1 = "kind1"; + KEY1 = Key.newBuilder(projectId, kind1, "name", options.getDatabaseId()).build(); // Set up the tracer for custom TraceID injection rootSpanName = From 4660b0b53af0136825f9909c4da88ad61ab9fa51 Mon Sep 17 00:00:00 2001 From: jimit-j-shah <57637300+jimit-j-shah@users.noreply.github.com> Date: Wed, 15 May 2024 11:28:43 -0700 Subject: [PATCH 15/17] Code review: Add common functionality to RemoteDatastoreHelper --- google-cloud-datastore/pom.xml | 11 +++--- .../testing/RemoteDatastoreHelper.java | 29 ++++++++++---- .../cloud/datastore/it/ITE2ETracingTest.java | 38 ++++--------------- 3 files changed, 35 insertions(+), 43 deletions(-) diff --git a/google-cloud-datastore/pom.xml b/google-cloud-datastore/pom.xml index a73c87251..171688902 100644 --- a/google-cloud-datastore/pom.xml +++ b/google-cloud-datastore/pom.xml @@ -129,6 +129,11 @@ + + io.opentelemetry + opentelemetry-sdk + ${opentelemetry.version} + io.opentelemetry opentelemetry-api @@ -196,12 +201,6 @@ test - - io.opentelemetry - opentelemetry-sdk - ${opentelemetry.version} - test - io.opentelemetry opentelemetry-sdk-common diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/testing/RemoteDatastoreHelper.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/testing/RemoteDatastoreHelper.java index 596ce96d8..6167cedca 100644 --- a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/testing/RemoteDatastoreHelper.java +++ b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/testing/RemoteDatastoreHelper.java @@ -19,13 +19,16 @@ import com.google.api.core.InternalApi; import com.google.api.gax.retrying.RetrySettings; import com.google.cloud.datastore.Datastore; +import com.google.cloud.datastore.DatastoreOpenTelemetryOptions; import com.google.cloud.datastore.DatastoreOptions; import com.google.cloud.datastore.Key; import com.google.cloud.datastore.Query; import com.google.cloud.datastore.QueryResults; import com.google.cloud.datastore.StructuredQuery; import com.google.cloud.http.HttpTransportOptions; +import io.opentelemetry.sdk.OpenTelemetrySdk; import java.util.UUID; +import javax.annotation.Nullable; import org.threeten.bp.Duration; /** @@ -38,13 +41,13 @@ * RetrySettings#getTotalTimeout()} is {@code 120000} and {@link * RetrySettings#getInitialRetryDelay()} is {@code 250}. {@link * HttpTransportOptions#getConnectTimeout()} and {@link HttpTransportOptions#getReadTimeout()} are - * both both set to {@code 60000}. + * both both set to {@code 60000}. If an OpenTelemetrySdk object is passed in, OpenTelemetry Trace + * collection will be enabled for the Client application. * *

Internal testing use only */ @InternalApi public class RemoteDatastoreHelper { - private final DatastoreOptions options; private final Datastore datastore; private final String namespace; @@ -78,18 +81,30 @@ public static RemoteDatastoreHelper create() { } /** Creates a {@code RemoteStorageHelper} object. */ - public static RemoteDatastoreHelper create(String databaseId) { + public static RemoteDatastoreHelper create( + String databaseId, @Nullable OpenTelemetrySdk openTelemetrySdk) { HttpTransportOptions transportOptions = DatastoreOptions.getDefaultHttpTransportOptions(); transportOptions = transportOptions.toBuilder().setConnectTimeout(60000).setReadTimeout(60000).build(); - DatastoreOptions datastoreOption = + DatastoreOptions.Builder datastoreOptionBuilder = DatastoreOptions.newBuilder() .setDatabaseId(databaseId) .setNamespace(UUID.randomUUID().toString()) .setRetrySettings(retrySettings()) - .setTransportOptions(transportOptions) - .build(); - return new RemoteDatastoreHelper(datastoreOption); + .setTransportOptions(transportOptions); + + if (openTelemetrySdk != null) { + datastoreOptionBuilder.setOpenTelemetryOptions( + DatastoreOpenTelemetryOptions.newBuilder() + .setOpenTelemetry(openTelemetrySdk) + .setTracingEnabled(true) + .build()); + } + return new RemoteDatastoreHelper(datastoreOptionBuilder.build()); + } + + public static RemoteDatastoreHelper create(String databaseId) { + return create(databaseId, /*openTelemetrySdk=*/ null); } private static RetrySettings retrySettings() { diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java index a25600335..46a3522c0 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java @@ -26,10 +26,10 @@ import com.google.api.gax.rpc.NotFoundException; import com.google.cloud.datastore.Datastore; -import com.google.cloud.datastore.DatastoreOpenTelemetryOptions; import com.google.cloud.datastore.DatastoreOptions; import com.google.cloud.datastore.Entity; import com.google.cloud.datastore.Key; +import com.google.cloud.datastore.testing.RemoteDatastoreHelper; import com.google.cloud.opentelemetry.trace.TraceConfiguration; import com.google.cloud.opentelemetry.trace.TraceExporter; import com.google.cloud.trace.v1.TraceServiceClient; @@ -189,8 +189,6 @@ private boolean dfsContainsCallStack(long spanId, List expectedCallStack private static final Logger logger = Logger.getLogger(ITE2ETracingTest.class.getName()); - private static final String SERVICE = "google.datastore.v1.Datastore/"; - private static final String RUN_AGGREGATION_QUERY_RPC_NAME = "RunAggregationQuery"; private static final String RUN_QUERY_RPC_NAME = "RunQuery"; @@ -237,7 +235,7 @@ private boolean dfsContainsCallStack(long spanId, List expectedCallStack @TestParameter boolean useGlobalOpenTelemetrySDK; - @TestParameter({"default", "test-db"}) + @TestParameter({"jimit-test-datastore" /*, "default", "test-db"*/}) String datastoreNamedDatabase; @BeforeClass @@ -284,34 +282,14 @@ public void before() throws Exception { } // Initialize the Datastore DB w/ the OTel SDK. Ideally we'd do this is the @BeforeAll method - // but because gRPC traces need to be deterministically force-flushed, datastore.shutdown() - // must be called in @After for each test. - DatastoreOptions.Builder optionsBuilder; - if (isUsingGlobalOpenTelemetrySDK()) { - optionsBuilder = - DatastoreOptions.newBuilder() - .setOpenTelemetryOptions( - DatastoreOpenTelemetryOptions.newBuilder().setTracingEnabled(true).build()); - } else { - optionsBuilder = - DatastoreOptions.newBuilder() - .setOpenTelemetryOptions( - DatastoreOpenTelemetryOptions.newBuilder() - .setOpenTelemetry(openTelemetrySdk) - .setTracingEnabled(true) - .build()); - } - + // but because gRPC traces need to be deterministically force-flushed for every test String namedDb = datastoreNamedDatabase(); - if (!namedDb.equals("default")) { - logger.log(Level.INFO, "Integration test using named database " + namedDb); - optionsBuilder = optionsBuilder.setDatabaseId(namedDb); - } else { - logger.log(Level.INFO, "Integration test using default database."); - } - optionsBuilder.setDatabaseId(namedDb); - options = optionsBuilder.build(); + logger.log(Level.INFO, "Integration test using named database " + namedDb); + RemoteDatastoreHelper remoteDatastoreHelper = + RemoteDatastoreHelper.create(namedDb, openTelemetrySdk); + options = remoteDatastoreHelper.getOptions(); datastore = options.getService(); + Preconditions.checkNotNull( datastore, "Error instantiating Datastore. Check that the service account credentials " From c7b80dd1dd8f3b3cf317fd05df252eebc0ba54e1 Mon Sep 17 00:00:00 2001 From: jimit-j-shah <57637300+jimit-j-shah@users.noreply.github.com> Date: Wed, 15 May 2024 12:06:00 -0700 Subject: [PATCH 16/17] Code review: clean up personal changes --- .../java/com/google/cloud/datastore/it/ITE2ETracingTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java index 46a3522c0..37b1be723 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java @@ -235,7 +235,7 @@ private boolean dfsContainsCallStack(long spanId, List expectedCallStack @TestParameter boolean useGlobalOpenTelemetrySDK; - @TestParameter({"jimit-test-datastore" /*, "default", "test-db"*/}) + @TestParameter({"default", "test-db"}) String datastoreNamedDatabase; @BeforeClass From 89d28d1aec793c678f1ef39339bcc35697ffe3ec Mon Sep 17 00:00:00 2001 From: jimit-j-shah <57637300+jimit-j-shah@users.noreply.github.com> Date: Wed, 15 May 2024 13:20:13 -0700 Subject: [PATCH 17/17] Code review: clean up --- .../java/com/google/cloud/datastore/it/ITE2ETracingTest.java | 5 ----- 1 file changed, 5 deletions(-) diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java index 37b1be723..24a9a7149 100644 --- a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java +++ b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java @@ -241,15 +241,10 @@ private boolean dfsContainsCallStack(long spanId, List expectedCallStack @BeforeClass public static void setup() throws IOException { projectId = DatastoreOptions.getDefaultProjectId(); - logger.info("projectId:" + projectId); - - // TODO(jimit) Make it re-usable w/ InMemorySpanExporter traceExporter = TraceExporter.createWithConfiguration( TraceConfiguration.builder().setProjectId(projectId).build()); - traceClient_v1 = TraceServiceClient.create(); - random = new Random(); }