diff --git a/java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java b/java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java index 63a80fe1913f..dfaa2aff6d76 100644 --- a/java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java +++ b/java-datastore/google-cloud-datastore/src/test/java/com/google/cloud/datastore/it/ITE2ETracingTest.java @@ -227,14 +227,12 @@ private boolean dfsContainsCallStack(long spanId, List expectedCallStack 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_COUNT = 120; 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; private static Key KEY2; @@ -246,22 +244,24 @@ private boolean dfsContainsCallStack(long spanId, List expectedCallStack // Random int generator for trace ID and span ID private static Random random; - private static TraceExporter traceExporter; + private static Credentials credentials; + + private TraceExporter traceExporter; // Required for reading back traces from Cloud Trace for validation private static TraceServiceClient traceClient; // Custom SpanContext for each test, required for TraceID injection - private static SpanContext customSpanContext; + private SpanContext customSpanContext; // Trace read back from Cloud Trace using traceClient for verification - private static Trace retrievedTrace; + private Trace retrievedTrace; - private static String rootSpanName; - private static Tracer tracer; + private String rootSpanName; + private Tracer tracer; // Required to set custom-root span - private static OpenTelemetrySdk openTelemetrySdk; + private OpenTelemetrySdk openTelemetrySdk; private static String projectId; @@ -287,14 +287,7 @@ public static void setup() throws IOException { // Share the same credentials used by Datastore client with the TraceExporter and // TraceServiceClient to ensure consistency and avoid auth issues in environments // where default ADC resolution might fail for the exporter. - Credentials credentials = DatastoreOptions.getDefaultInstance().getCredentials(); - - TraceConfiguration.Builder traceConfigurationBuilder = - TraceConfiguration.builder().setProjectId(projectId); - if (credentials != null) { - traceConfigurationBuilder.setCredentials(credentials); - } - traceExporter = TraceExporter.createWithConfiguration(traceConfigurationBuilder.build()); + credentials = DatastoreOptions.getDefaultInstance().getCredentials(); TraceServiceSettings.Builder clientBuilder = TraceServiceSettings.newBuilder(); if (credentials != null) { @@ -310,6 +303,13 @@ public void before() throws Exception { Resource resource = Resource.getDefault().merge(Resource.builder().put(SERVICE_NAME, "Sparky").build()); + TraceConfiguration.Builder traceConfigurationBuilder = + TraceConfiguration.builder().setProjectId(projectId); + if (credentials != null) { + traceConfigurationBuilder.setCredentials(credentials); + } + traceExporter = TraceExporter.createWithConfiguration(traceConfigurationBuilder.build()); + if (isUsingGlobalOpenTelemetrySDK()) { openTelemetrySdk = OpenTelemetrySdk.builder() @@ -393,7 +393,22 @@ public void after() throws Exception { tracer = null; retrievedTrace = null; customSpanContext = null; + try { + if (openTelemetrySdk != null) { + openTelemetrySdk.close(); + } + } finally { + if (traceExporter != null) { + try { + // Attempt to shut down traceExporter. + traceExporter.shutdown(); + } catch (Exception e) { + logger.log(Level.WARNING, "Failed to shut down traceExporter", e); + } + } + } openTelemetrySdk = null; + traceExporter = null; } @AfterClass @@ -437,10 +452,17 @@ protected Span getNewRootSpanWithContext() { } protected void waitForTracesToComplete() throws Exception { + if (openTelemetrySdk == null) { + logger.warning("OpenTelemetrySdk is null, cannot flush traces"); + return; + } logger.info("Flushing traces..."); CompletableResultCode completableResultCode = openTelemetrySdk.getSdkTracerProvider().forceFlush(); completableResultCode.join(TRACE_FORCE_FLUSH_MILLIS, TimeUnit.MILLISECONDS); + if (!completableResultCode.isSuccess()) { + logger.warning("Force flush did not complete successfully"); + } } // Validates `retrievedTrace`. Cloud Trace indexes traces w/ eventual consistency, even when