Use early static registration of EventPublishingContextWrapper in tests

Add `OpenTelemetryEventPublisherBeansTestExecutionListener` JUnit
class to automatically trigger early addition of the `ContextStorage`
wrapper. The listener has also been updated with a static `addWrapper()`
method that can be called directly for other test frameworks.

Closes gh-42005

spring-boot-project/spring-boot-actuator-autoconfigure/build.gradle

@@ -106,6 +106,7 @@ dependencies {
 	optional("org.hibernate.orm:hibernate-micrometer")
 	optional("org.hibernate.validator:hibernate-validator")
 	optional("org.influxdb:influxdb-java")
+	optional("org.junit.platform:junit-platform-launcher")
 	optional("org.liquibase:liquibase-core") {
 		exclude group: "javax.xml.bind", module: "jaxb-api"
 	}

spring-boot-project/spring-boot-actuator-autoconfigure/src/main/java/org/springframework/boot/actuate/autoconfigure/tracing/OpenTelemetryEventPublisherBeansApplicationListener.java

@@ -18,7 +18,6 @@ package org.springframework.boot.actuate.autoconfigure.tracing;
 
 import java.util.List;
 import java.util.concurrent.atomic.AtomicBoolean;
-import java.util.function.UnaryOperator;
 
 import io.micrometer.tracing.otel.bridge.EventPublishingContextWrapper;
 import io.micrometer.tracing.otel.bridge.OtelTracer.EventPublisher;
@@ -42,11 +41,17 @@ import org.springframework.util.MultiValueMap;
 /**
  * {@link ApplicationListener} to add an OpenTelemetry {@link ContextStorage} wrapper for
  * {@link EventPublisher} bean support. A single {@link ContextStorage} wrapper is added
- * as early as possible then updated with {@link EventPublisher} beans as needed.
+ * on the {@link ApplicationStartingEvent} then updated with {@link EventPublisher} beans
+ * as needed.
+ * <p>
+ * The {@link #addWrapper()} method may also be called directly if the
+ * {@link ApplicationStartingEvent} isn't called early enough or isn't fired.
  *
  * @author Phillip Webb
+ * @since 3.4.0
+ * @see OpenTelemetryEventPublisherBeansTestExecutionListener
  */
-class OpenTelemetryEventPublisherApplicationListener implements GenericApplicationListener {
+public class OpenTelemetryEventPublisherBeansApplicationListener implements GenericApplicationListener {
 
 	private static final boolean OTEL_CONTEXT_PRESENT = ClassUtils.isPresent("io.opentelemetry.context.ContextStorage",
 			null);
@@ -54,6 +59,8 @@ class OpenTelemetryEventPublisherApplicationListener implements GenericApplicati
 	private static final boolean MICROMETER_OTEL_PRESENT = ClassUtils
 		.isPresent("io.micrometer.tracing.otel.bridge.OtelTracer", null);
 
+	private static final AtomicBoolean added = new AtomicBoolean();
+
 	@Override
 	public int getOrder() {
 		return Ordered.HIGHEST_PRECEDENCE;
@@ -69,11 +76,11 @@ class OpenTelemetryEventPublisherApplicationListener implements GenericApplicati
 
 	@Override
 	public void onApplicationEvent(ApplicationEvent event) {
-		if (!OTEL_CONTEXT_PRESENT || !MICROMETER_OTEL_PRESENT) {
+		if (!isInstallable()) {
 			return;
 		}
 		if (event instanceof ApplicationStartingEvent) {
-			EventPublisherBeansContextWrapper.addWrapperIfNecessary();
+			addWrapper();
 		}
 		if (event instanceof ContextRefreshedEvent contextRefreshedEvent) {
 			ApplicationContext applicationContext = contextRefreshedEvent.getApplicationContext();
@@ -83,92 +90,105 @@ class OpenTelemetryEventPublisherApplicationListener implements GenericApplicati
 				.stream()
 				.map(EventPublishingContextWrapper::new)
 				.toList();
-			EventPublisherBeansContextWrapper.instance.put(applicationContext, publishers);
+			Wrapper.instance.put(applicationContext, publishers);
 		}
 		if (event instanceof ContextClosedEvent contextClosedEvent) {
-			EventPublisherBeansContextWrapper.instance.remove(contextClosedEvent.getApplicationContext());
+			Wrapper.instance.remove(contextClosedEvent.getApplicationContext());
 		}
 	}
 
 	/**
-	 * The single {@link ContextStorage} wrapper that delegates to {@link EventPublisher}
+	 * {@link ContextStorage#addWrapper(java.util.function.Function) Add} the
+	 * {@link ContextStorage} wrapper to ensure that {@link EventPublisher} are propagated
+	 * correctly.
+	 */
+	public static void addWrapper() {
+		if (isInstallable() && added.compareAndSet(false, true)) {
+			Wrapper.instance.addWrapper();
+		}
+	}
+
+	private static boolean isInstallable() {
+		return OTEL_CONTEXT_PRESENT && MICROMETER_OTEL_PRESENT;
+	}
+
+	/**
+	 * Single instance class used to add the wrapper and manage the {@link EventPublisher}
 	 * beans.
 	 */
-	static class EventPublisherBeansContextWrapper implements UnaryOperator<ContextStorage> {
+	static final class Wrapper {
 
-		private static final AtomicBoolean added = new AtomicBoolean();
+		static Wrapper instance = new Wrapper();
 
-		private static final EventPublisherBeansContextWrapper instance = new EventPublisherBeansContextWrapper();
+		private final MultiValueMap<ApplicationContext, EventPublishingContextWrapper> beans = new LinkedMultiValueMap<>();
 
-		private final MultiValueMap<ApplicationContext, EventPublishingContextWrapper> publishers = new LinkedMultiValueMap<>();
+		private volatile ContextStorage storageDelegate;
 
-		private volatile ContextStorage delegate;
-
-		static void addWrapperIfNecessary() {
-			if (added.compareAndSet(false, true)) {
-				ContextStorage.addWrapper(instance);
-			}
+		private Wrapper() {
 		}
 
-		@Override
-		public ContextStorage apply(ContextStorage contextStorage) {
-			return new EventPublisherBeansContextStorage(contextStorage);
+		private void addWrapper() {
+			ContextStorage.addWrapper(Storage::new);
 		}
 
 		void put(ApplicationContext applicationContext, List<EventPublishingContextWrapper> publishers) {
 			synchronized (this) {
-				this.publishers.addAll(applicationContext, publishers);
-				this.delegate = null;
+				this.beans.addAll(applicationContext, publishers);
+				this.storageDelegate = null;
 			}
 		}
 
 		void remove(ApplicationContext applicationContext) {
 			synchronized (this) {
-				this.publishers.remove(applicationContext);
-				this.delegate = null;
+				this.beans.remove(applicationContext);
+				this.storageDelegate = null;
 			}
 		}
 
-		private ContextStorage getDelegate(ContextStorage parent) {
-			ContextStorage delegate = this.delegate;
+		ContextStorage getStorageDelegate(ContextStorage parent) {
+			ContextStorage delegate = this.storageDelegate;
 			if (delegate == null) {
 				synchronized (this) {
 					delegate = parent;
-					for (List<EventPublishingContextWrapper> publishers : this.publishers.values()) {
+					for (List<EventPublishingContextWrapper> publishers : this.beans.values()) {
 						for (EventPublishingContextWrapper publisher : publishers) {
 							delegate = publisher.apply(delegate);
 						}
 					}
+					this.storageDelegate = delegate;
 				}
 			}
 			return delegate;
 		}
 
 		/**
-		 * The wrapped {@link ContextStorage} that delegates to the
-		 * {@link EventPublisherBeansContextWrapper}.
+		 * {@link ContextStorage} that delegates to the {@link EventPublisher} beans.
 		 */
-		class EventPublisherBeansContextStorage implements ContextStorage {
+		class Storage implements ContextStorage {
 
 			private final ContextStorage parent;
 
-			EventPublisherBeansContextStorage(ContextStorage wrapped) {
-				this.parent = wrapped;
+			Storage(ContextStorage parent) {
+				this.parent = parent;
 			}
 
 			@Override
 			public Scope attach(Context toAttach) {
-				return getDelegate(this.parent).attach(toAttach);
+				return getDelegate().attach(toAttach);
 			}
 
 			@Override
 			public Context current() {
-				return getDelegate(this.parent).current();
+				return getDelegate().current();
 			}
 
 			@Override
 			public Context root() {
-				return getDelegate(this.parent).root();
+				return getDelegate().root();
+			}
+
+			private ContextStorage getDelegate() {
+				return getStorageDelegate(this.parent);
 			}
 
 		}

spring-boot-project/spring-boot-actuator-autoconfigure/src/main/java/org/springframework/boot/actuate/autoconfigure/tracing/OpenTelemetryEventPublisherBeansTestExecutionListener.java

@@ -0,0 +1,38 @@
+/*
+ * Copyright 2012-2024 the original author or authors.
+ *
+ * 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
+ *
+ *      https://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 org.springframework.boot.actuate.autoconfigure.tracing;
+
+import org.junit.platform.launcher.TestExecutionListener;
+import org.junit.platform.launcher.TestIdentifier;
+
+/**
+ * JUnit {@link TestExecutionListener} to ensure
+ * {@link OpenTelemetryEventPublisherBeansApplicationListener#addWrapper()} is called as
+ * early as possible.
+ *
+ * @author Phillip Webb
+ * @since 3.4.0
+ * @see OpenTelemetryEventPublisherBeansApplicationListener
+ */
+public class OpenTelemetryEventPublisherBeansTestExecutionListener implements TestExecutionListener {
+
+	@Override
+	public void executionStarted(TestIdentifier testIdentifier) {
+		OpenTelemetryEventPublisherBeansApplicationListener.addWrapper();
+	}
+
+}

spring-boot-project/spring-boot-actuator-autoconfigure/src/main/resources/META-INF/services/org.junit.platform.launcher.TestExecutionListener

@@ -0,0 +1 @@
+org.springframework.boot.actuate.autoconfigure.tracing.OpenTelemetryEventPublisherBeansTestExecutionListener

spring-boot-project/spring-boot-actuator-autoconfigure/src/main/resources/META-INF/spring.factories

@@ -13,4 +13,4 @@ org.springframework.boot.actuate.autoconfigure.tracing.LogCorrelationEnvironment
 
 # Application Listeners
 org.springframework.context.ApplicationListener=\
-org.springframework.boot.actuate.autoconfigure.tracing.OpenTelemetryEventPublisherApplicationListener
+org.springframework.boot.actuate.autoconfigure.tracing.OpenTelemetryEventPublisherBeansApplicationListener

spring-boot-project/spring-boot-actuator-autoconfigure/src/test/java/org/springframework/boot/actuate/autoconfigure/tracing/BaggagePropagationIntegrationTests.java

@@ -31,7 +31,6 @@ import org.junit.jupiter.params.provider.EnumSource;
 import org.slf4j.MDC;
 
 import org.springframework.boot.actuate.autoconfigure.opentelemetry.OpenTelemetryAutoConfiguration;
-import org.springframework.boot.actuate.autoconfigure.tracing.OpenTelemetryEventPublisherApplicationListener.EventPublisherBeansContextWrapper;
 import org.springframework.boot.autoconfigure.AutoConfigurations;
 import org.springframework.boot.test.context.runner.ApplicationContextRunner;
 import org.springframework.boot.testsupport.classpath.ForkedClassPath;
@@ -58,7 +57,7 @@ class BaggagePropagationIntegrationTests {
 	@BeforeEach
 	@AfterEach
 	void setup() {
-		EventPublisherBeansContextWrapper.addWrapperIfNecessary();
+		OpenTelemetryEventPublisherBeansApplicationListener.addWrapper();
 		MDC.clear();
 	}
 
@@ -291,7 +290,7 @@ class BaggagePropagationIntegrationTests {
 
 		@Override
 		public void initialize(ConfigurableApplicationContext applicationContext) {
-			applicationContext.addApplicationListener(new OpenTelemetryEventPublisherApplicationListener());
+			applicationContext.addApplicationListener(new OpenTelemetryEventPublisherBeansApplicationListener());
 		}
 
 	}

spring-boot-project/spring-boot-actuator-autoconfigure/src/test/java/org/springframework/boot/actuate/autoconfigure/tracing/OpenTelemetryEventPublishingContextWrapperBeansTestExecutionListenerIntegrationTests.java

@@ -0,0 +1,52 @@
+/*
+ * Copyright 2012-2024 the original author or authors.
+ *
+ * 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
+ *
+ *      https://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 org.springframework.boot.actuate.autoconfigure.tracing;
+
+import java.lang.reflect.Method;
+import java.util.List;
+import java.util.function.Function;
+
+import io.opentelemetry.context.ContextStorage;
+import org.junit.jupiter.api.Test;
+
+import org.springframework.boot.actuate.autoconfigure.tracing.OpenTelemetryEventPublisherBeansApplicationListener.Wrapper.Storage;
+import org.springframework.boot.testsupport.classpath.ForkedClassPath;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.mock;
+
+/**
+ * Integration tests for {@link OpenTelemetryEventPublisherBeansTestExecutionListener}.
+ *
+ * @author Phillip Webb
+ */
+@ForkedClassPath
+class OpenTelemetryEventPublishingContextWrapperBeansTestExecutionListenerIntegrationTests {
+
+	private final ContextStorage parent = mock(ContextStorage.class);
+
+	@Test
+	@SuppressWarnings({ "unchecked", "rawtypes" })
+	void wrapperIsInstalled() throws Exception {
+		Class<?> wrappersClass = Class.forName("io.opentelemetry.context.ContextStorageWrappers");
+		Method getWrappersMethod = wrappersClass.getDeclaredMethod("getWrappers");
+		getWrappersMethod.setAccessible(true);
+		List<Function> wrappers = (List<Function>) getWrappersMethod.invoke(null);
+		assertThat(wrappers).anyMatch((function) -> function.apply(this.parent) instanceof Storage);
+	}
+
+}

spring-boot-project/spring-boot-actuator-autoconfigure/src/test/java/org/springframework/boot/actuate/autoconfigure/tracing/OpenTelemetryTracingAutoConfigurationTests.java

@@ -54,14 +54,12 @@ import io.opentelemetry.sdk.trace.data.SpanData;
 import io.opentelemetry.sdk.trace.export.SpanExporter;
 import io.opentelemetry.sdk.trace.samplers.Sampler;
 import org.assertj.core.api.InstanceOfAssertFactories;
-import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
 import org.mockito.Mockito;
 
 import org.springframework.boot.actuate.autoconfigure.observation.ObservationAutoConfiguration;
-import org.springframework.boot.actuate.autoconfigure.tracing.OpenTelemetryEventPublisherApplicationListener.EventPublisherBeansContextWrapper;
 import org.springframework.boot.autoconfigure.AutoConfigurations;
 import org.springframework.boot.context.annotation.Configurations;
 import org.springframework.boot.test.context.FilteredClassLoader;
@@ -93,11 +91,6 @@ class OpenTelemetryTracingAutoConfigurationTests {
 				org.springframework.boot.actuate.autoconfigure.opentelemetry.OpenTelemetryAutoConfiguration.class,
 				OpenTelemetryTracingAutoConfiguration.class));
 
-	@BeforeAll
-	static void addWrapper() {
-		EventPublisherBeansContextWrapper.addWrapperIfNecessary();
-	}
-
 	@Test
 	void shouldSupplyBeans() {
 		this.contextRunner.run((context) -> {
@@ -354,7 +347,7 @@ class OpenTelemetryTracingAutoConfigurationTests {
 	}
 
 	private void initializeOpenTelemetry(ConfigurableApplicationContext context) {
-		context.addApplicationListener(new OpenTelemetryEventPublisherApplicationListener());
+		context.addApplicationListener(new OpenTelemetryEventPublisherBeansApplicationListener());
 		Span.current();
 	}