Also run strict context check retry for junit tests
MODIFY
testing-common/src/main/groovy/io/opentelemetry/instrumentation/test/InstrumentationSpecification.groovy
testing-common/src/main/groovy/io/opentelemetry/instrumentation/test/InstrumentationSpecification.groovy
package io.opentelemetry.instrumentation.test
-import static org.awaitility.Awaitility.await
-
import groovy.transform.stc.ClosureParams
import groovy.transform.stc.SimpleType
import io.opentelemetry.api.OpenTelemetry
@@ -14,12 +12,11 @@ import io.opentelemetry.api.trace.Span
import io.opentelemetry.context.ContextStorage
import io.opentelemetry.instrumentation.test.asserts.InMemoryExporterAssert
import io.opentelemetry.instrumentation.testing.InstrumentationTestRunner
+import io.opentelemetry.instrumentation.testing.util.ContextStorageCloser
import io.opentelemetry.instrumentation.testing.util.TelemetryDataUtil
import io.opentelemetry.instrumentation.testing.util.ThrowingSupplier
import io.opentelemetry.sdk.metrics.data.MetricData
import io.opentelemetry.sdk.trace.data.SpanData
-import java.lang.reflect.Field
-import java.util.concurrent.ConcurrentHashMap
import org.junit.Rule
import org.junit.rules.Timeout
import spock.lang.Specification
@@ -50,105 +47,8 @@ abstract class InstrumentationSpecification extends Specification {
}
def cleanup() {
- ContextRestorer restorer = ContextRestorer.create()
- if (restorer == null) {
- tryCleanup()
- return
- }
-
- // If close is called before all request processing threads have completed we might get false
- // positive notifications for leaked scopes when strict context checks are enabled. Here we
- // retry close when scope leak was reported.
- await()
- .ignoreException(AssertionError)
- .atMost(15, TimeUnit.SECONDS)
- .until({
- restorer.runWithRestore { tryCleanup() }
- true
- })
- }
-
- def tryCleanup() {
ContextStorage storage = ContextStorage.get()
- if (storage instanceof AutoCloseable) {
- ((AutoCloseable) storage).close()
- }
- }
-
- // Helper class that allows for retrying ContextStorage close operation by restoring ContextStorage
- // to the state where it was before close was called.
- static abstract class ContextRestorer {
- abstract void restore()
-
- void runWithRestore(Closure target) {
- try {
- target.run()
- } catch (AssertionError assertionError) {
- restore()
- throw assertionError
- }
- }
-
- static ContextRestorer create() {
- def strictContextStorage = getStrictContextStorage()
- if (strictContextStorage == null) {
- return null
- }
- def pendingScopes = getStrictContextStoragePendingScopes(strictContextStorage)
-
- def pendingScopesClass = Class.forName("io.opentelemetry.javaagent.shaded.io.opentelemetry.context.StrictContextStorage\$PendingScopes")
- Field mapField = pendingScopesClass.getDeclaredField("map")
- mapField.setAccessible(true)
- ConcurrentHashMap map = mapField.get(pendingScopes)
- Map copy = new HashMap(map)
-
- return new ContextRestorer() {
- @Override
- void restore() {
- map.putAll(copy)
- }
- }
- }
-
- static getStrictContextStoragePendingScopes(def strictContextStorage) {
- def strictContextStorageClass = Class.forName("io.opentelemetry.javaagent.shaded.io.opentelemetry.context.StrictContextStorage")
- Field field = strictContextStorageClass.getDeclaredField("pendingScopes")
- field.setAccessible(true)
- return field.get(strictContextStorage)
- }
-
- static getStrictContextStorage() {
- def contextStorage = getAgentContextStorage()
- if (contextStorage == null) {
- return null
- }
- contextStorage = unwrapStrictContextStressor(contextStorage)
- def contextStorageClass = contextStorage.getClass()
- if (contextStorageClass.getName().contains("StrictContextStorage")) {
- return contextStorage
- }
- return null
- }
-
- static getAgentContextStorage() {
- try {
- def contextStorageClass = Class.forName("io.opentelemetry.javaagent.shaded.io.opentelemetry.context.ContextStorage")
- def method = contextStorageClass.getDeclaredMethod("get")
- return method.invoke(null)
- } catch (Exception exception) {
- return null
- }
- }
-
- static unwrapStrictContextStressor(def contextStorage) {
- Class<?> contextStorageClass = contextStorage.getClass()
- if (contextStorageClass.getName().contains("StrictContextStressor")) {
- Field field = contextStorageClass.getDeclaredField("contextStorage")
- field.setAccessible(true)
- return field.get(contextStorage)
- }
- return contextStorage
- }
+ ContextStorageCloser.close(storage)
}
def cleanupSpec() {

MODIFY
testing-common/src/main/java/io/opentelemetry/instrumentation/testing/junit/InstrumentationExtension.java
testing-common/src/main/java/io/opentelemetry/instrumentation/testing/junit/InstrumentationExtension.java
@@ -11,6 +11,7 @@ import static org.awaitility.Awaitility.await;
import io.opentelemetry.api.OpenTelemetry;
import io.opentelemetry.context.ContextStorage;
import io.opentelemetry.instrumentation.testing.InstrumentationTestRunner;
+import io.opentelemetry.instrumentation.testing.util.ContextStorageCloser;
import io.opentelemetry.instrumentation.testing.util.ThrowingRunnable;
import io.opentelemetry.instrumentation.testing.util.ThrowingSupplier;
import io.opentelemetry.sdk.metrics.data.MetricData;
@@ -27,7 +28,6 @@ import org.junit.jupiter.api.extension.ExtensionContext;
public abstract class InstrumentationExtension
implements BeforeAllCallback, BeforeEachCallback, AfterAllCallback, AfterEachCallback {
- private static final long DEFAULT_TRACE_WAIT_TIMEOUT_SECONDS = 20;
private final InstrumentationTestRunner testRunner;
@@ -48,9 +48,7 @@ public abstract class InstrumentationExtension
@Override
public void afterEach(ExtensionContext context) throws Exception {
ContextStorage storage = ContextStorage.get();
- if (storage instanceof AutoCloseable) {
- ((AutoCloseable) storage).close();
- }
+ ContextStorageCloser.close(storage);
}
@Override

ADD
testing-common/src/main/java/io/opentelemetry/instrumentation/testing/util/ContextStorageCloser.java
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.testing.util;
+
+import static org.awaitility.Awaitility.await;
+
+import io.opentelemetry.context.ContextStorage;
+import java.lang.reflect.Field;
+import java.lang.reflect.Method;
+import java.time.Duration;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+/**
+ * Helper class for closing context storage with retry in case there are strict context check
+ * failures.
+ */
+public final class ContextStorageCloser {
+
+ private ContextStorageCloser() {}
+
+ public static void close(ContextStorage storage) throws Exception {
+ if (storage instanceof AutoCloseable) {
+ cleanup((AutoCloseable) storage);
+ }
+ }
+
+ private static void cleanup(AutoCloseable storage) throws Exception {
+ ContextRestorer restorer = ContextRestorer.create();
+ if (restorer == null) {
+ storage.close();
+ return;
+ }
+
+ // If close is called before all request processing threads have completed we might get false
+ // positive notifications for leaked scopes when strict context checks are enabled. Here we
+ // retry close when scope leak was reported.
+ await()
+ .ignoreException(AssertionError.class)
+ .atMost(Duration.ofSeconds(15))
+ .until(() -> restorer.runWithRestore(storage));
+ }
+
+ // Helper class that allows for retrying ContextStorage close operation by restoring
+ // ContextStorage to the state where it was before close was called.
+ private abstract static class ContextRestorer {
+ abstract void restore();
+
+ boolean runWithRestore(AutoCloseable target) {
+ try {
+ target.close();
+ return true;
+ } catch (Throwable throwable) {
+ restore();
+ if (throwable instanceof AssertionError) {
+ throw (AssertionError) throwable;
+ }
+ throw new IllegalStateException(throwable);
+ }
+ }
+
+ static ContextRestorer create()
+ throws NoSuchFieldException, IllegalAccessException, ClassNotFoundException {
+ Object strictContextStorage = getStrictContextStorage();
+ if (strictContextStorage == null) {
+ return null;
+ }
+ Object pendingScopes = getStrictContextStoragePendingScopes(strictContextStorage);
+
+ Class<?> pendingScopesClass =
+ Class.forName(
+ "io.opentelemetry.javaagent.shaded.io.opentelemetry.context.StrictContextStorage$PendingScopes");
+ Field mapField = pendingScopesClass.getDeclaredField("map");
+ mapField.setAccessible(true);
+ @SuppressWarnings("unchecked")
+ ConcurrentHashMap<Object, Object> map =
+ (ConcurrentHashMap<Object, Object>) mapField.get(pendingScopes);
+ Map<Object, Object> copy = new HashMap<>(map);
+
+ return new ContextRestorer() {
+ @Override
+ void restore() {
+ map.putAll(copy);
+ }
+ };
+ }
+
+ private static Object getStrictContextStoragePendingScopes(Object strictContextStorage)
+ throws ClassNotFoundException, NoSuchFieldException, IllegalAccessException {
+ Class<?> strictContextStorageClass =
+ Class.forName(
+ "io.opentelemetry.javaagent.shaded.io.opentelemetry.context.StrictContextStorage");
+ Field field = strictContextStorageClass.getDeclaredField("pendingScopes");
+ field.setAccessible(true);
+ return field.get(strictContextStorage);
+ }
+
+ private static Object getStrictContextStorage()
+ throws NoSuchFieldException, IllegalAccessException {
+ Object contextStorage = getAgentContextStorage();
+ if (contextStorage == null) {
+ return null;
+ }
+ contextStorage = unwrapStrictContextStressor(contextStorage);
+ Class<?> contextStorageClass = contextStorage.getClass();
+ if (contextStorageClass.getName().contains("StrictContextStorage")) {
+ return contextStorage;
+ }
+ return null;
+ }
+
+ private static Object getAgentContextStorage() {
+ try {
+ Class<?> contextStorageClass =
+ Class.forName(
+ "io.opentelemetry.javaagent.shaded.io.opentelemetry.context.ContextStorage");
+ Method method = contextStorageClass.getDeclaredMethod("get");
+ return method.invoke(null);
+ } catch (Exception exception) {
+ return null;
+ }
+ }
+
+ private static Object unwrapStrictContextStressor(Object contextStorage)
+ throws NoSuchFieldException, IllegalAccessException {
+ Class<?> contextStorageClass = contextStorage.getClass();
+ if (contextStorageClass.getName().contains("StrictContextStressor")) {
+ Field field = contextStorageClass.getDeclaredField("contextStorage");
+ field.setAccessible(true);
+ return field.get(contextStorage);
+ }
+ return contextStorage;
+ }
+ }
+}

Also run strict context check retry for junit tests (#4380)