diff --git a/spring-integration-core/src/main/java/org/springframework/integration/channel/FluxMessageChannel.java b/spring-integration-core/src/main/java/org/springframework/integration/channel/FluxMessageChannel.java index fa45245787..b2f38a7424 100644 --- a/spring-integration-core/src/main/java/org/springframework/integration/channel/FluxMessageChannel.java +++ b/spring-integration-core/src/main/java/org/springframework/integration/channel/FluxMessageChannel.java @@ -21,7 +21,6 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicReference; import java.util.concurrent.locks.LockSupport; -import io.micrometer.context.ContextSnapshotFactory; import org.reactivestreams.Publisher; import org.reactivestreams.Subscriber; import reactor.core.Disposable; @@ -31,17 +30,16 @@ import reactor.core.publisher.Mono; import reactor.core.publisher.Sinks; import reactor.core.scheduler.Scheduler; import reactor.core.scheduler.Schedulers; -import reactor.util.context.Context; import reactor.util.context.ContextView; import org.springframework.core.log.LogMessage; import org.springframework.integration.IntegrationMessageHeaderAccessor; import org.springframework.integration.StaticMessageHeaderAccessor; import org.springframework.integration.support.MutableMessageBuilder; +import org.springframework.integration.util.IntegrationReactiveUtils; import org.springframework.messaging.Message; import org.springframework.messaging.MessageDeliveryException; import org.springframework.util.Assert; -import org.springframework.util.ClassUtils; /** * The {@link AbstractMessageChannel} implementation for the @@ -56,9 +54,6 @@ import org.springframework.util.ClassUtils; public class FluxMessageChannel extends AbstractMessageChannel implements Publisher>, ReactiveStreamsSubscribableChannel { - private static final boolean isContextPropagationPresent = ClassUtils.isPresent( - "io.micrometer.context.ContextSnapshot", FluxMessageChannel.class.getClassLoader()); - private final Scheduler scheduler = Schedulers.boundedElastic(); private final Sinks.Many> sink = Sinks.many().multicast().onBackpressureBuffer(1, false); @@ -91,8 +86,8 @@ public class FluxMessageChannel extends AbstractMessageChannel private boolean tryEmitMessage(Message message) { Message messageToEmit = message; - if (isContextPropagationPresent) { - ContextView contextView = ContextSnapshotHelper.captureContext(); + if (IntegrationReactiveUtils.isContextPropagationPresent) { + ContextView contextView = IntegrationReactiveUtils.captureReactorContext(); if (!contextView.isEmpty()) { messageToEmit = MutableMessageBuilder.fromMessage(message) .setHeader(IntegrationMessageHeaderAccessor.REACTOR_CONTEXT, contextView) @@ -196,14 +191,4 @@ public class FluxMessageChannel extends AbstractMessageChannel super.destroy(); } - private static final class ContextSnapshotHelper { - - private static final ContextSnapshotFactory CONTEXT_SNAPSHOT_FACTORY = ContextSnapshotFactory.builder().build(); - - static ContextView captureContext() { - return CONTEXT_SNAPSHOT_FACTORY.captureAll().updateContext(Context.empty()); - } - - } - } diff --git a/spring-integration-core/src/main/java/org/springframework/integration/endpoint/ReactiveStreamsConsumer.java b/spring-integration-core/src/main/java/org/springframework/integration/endpoint/ReactiveStreamsConsumer.java index fc5a3ade22..b9a8a41970 100644 --- a/spring-integration-core/src/main/java/org/springframework/integration/endpoint/ReactiveStreamsConsumer.java +++ b/spring-integration-core/src/main/java/org/springframework/integration/endpoint/ReactiveStreamsConsumer.java @@ -19,7 +19,6 @@ package org.springframework.integration.endpoint; import java.util.function.Consumer; import java.util.function.Function; -import io.micrometer.context.ContextSnapshotFactory; import org.reactivestreams.Publisher; import org.reactivestreams.Subscriber; import org.reactivestreams.Subscription; @@ -31,6 +30,7 @@ import reactor.util.context.ContextView; import org.springframework.context.Lifecycle; import org.springframework.integration.IntegrationMessageHeaderAccessor; +import org.springframework.integration.StaticMessageHeaderAccessor; import org.springframework.integration.channel.ChannelUtils; import org.springframework.integration.channel.NullChannel; import org.springframework.integration.core.MessageProducer; @@ -44,7 +44,6 @@ import org.springframework.messaging.MessageChannel; import org.springframework.messaging.MessageHandler; import org.springframework.messaging.ReactiveMessageHandler; import org.springframework.util.Assert; -import org.springframework.util.ClassUtils; import org.springframework.util.ErrorHandler; @@ -59,9 +58,6 @@ import org.springframework.util.ErrorHandler; */ public class ReactiveStreamsConsumer extends AbstractEndpoint implements IntegrationConsumer { - private static final boolean isContextPropagationPresent = ClassUtils.isPresent( - "io.micrometer.context.ContextSnapshot", ReactiveStreamsConsumer.class.getClassLoader()); - private final MessageChannel inputChannel; private final Publisher> publisher; @@ -189,7 +185,9 @@ public class ReactiveStreamsConsumer extends AbstractEndpoint implements Integra if (this.reactiveMessageHandler != null) { this.subscription = fluxFromChannel - .flatMap(this.reactiveMessageHandler::handleMessage) + .flatMap((message) -> + this.reactiveMessageHandler.handleMessage(message) + .contextWrite(StaticMessageHeaderAccessor.getReactorContext(message))) .onErrorContinue((ex, data) -> this.errorHandler.handleError(ex)) .subscribe(); } @@ -302,7 +300,7 @@ public class ReactiveStreamsConsumer extends AbstractEndpoint implements Integra protected void hookOnNext(Message message) { Message messageToDeliver = message; - if (isContextPropagationPresent) { + if (IntegrationReactiveUtils.isContextPropagationPresent) { ContextView reactorContext = message.getHeaders() .get(IntegrationMessageHeaderAccessor.REACTOR_CONTEXT, ContextView.class); @@ -312,7 +310,7 @@ public class ReactiveStreamsConsumer extends AbstractEndpoint implements Integra .removeHeader(IntegrationMessageHeaderAccessor.REACTOR_CONTEXT) .build(); - try (AutoCloseable scope = ContextSnapshotHelper.setContext(reactorContext)) { + try (AutoCloseable scope = IntegrationReactiveUtils.setThreadLocalsFromReactorContext(reactorContext)) { this.delegate.onNext(messageToDeliver); } catch (Exception ex) { @@ -337,14 +335,4 @@ public class ReactiveStreamsConsumer extends AbstractEndpoint implements Integra } - private static final class ContextSnapshotHelper { - - private static final ContextSnapshotFactory CONTEXT_SNAPSHOT_FACTORY = ContextSnapshotFactory.builder().build(); - - static AutoCloseable setContext(ContextView context) { - return CONTEXT_SNAPSHOT_FACTORY.setThreadLocalsFrom(context); - } - - } - } diff --git a/spring-integration-core/src/main/java/org/springframework/integration/util/IntegrationReactiveUtils.java b/spring-integration-core/src/main/java/org/springframework/integration/util/IntegrationReactiveUtils.java index a27b98aa98..0a57d70c4f 100644 --- a/spring-integration-core/src/main/java/org/springframework/integration/util/IntegrationReactiveUtils.java +++ b/spring-integration-core/src/main/java/org/springframework/integration/util/IntegrationReactiveUtils.java @@ -1,5 +1,5 @@ /* - * Copyright 2020-2022 the original author or authors. + * Copyright 2020-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. @@ -19,6 +19,7 @@ package org.springframework.integration.util; import java.time.Duration; import java.util.concurrent.locks.LockSupport; +import io.micrometer.context.ContextSnapshotFactory; import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; import org.reactivestreams.Publisher; @@ -26,17 +27,22 @@ import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; import reactor.core.publisher.Sinks; import reactor.core.scheduler.Schedulers; +import reactor.util.context.Context; +import reactor.util.context.ContextView; import reactor.util.retry.Retry; +import org.springframework.integration.IntegrationMessageHeaderAccessor; import org.springframework.integration.StaticMessageHeaderAccessor; import org.springframework.integration.acks.AckUtils; import org.springframework.integration.core.MessageSource; +import org.springframework.integration.support.MutableMessageBuilder; import org.springframework.messaging.Message; import org.springframework.messaging.MessageChannel; import org.springframework.messaging.MessageHandler; import org.springframework.messaging.MessagingException; import org.springframework.messaging.PollableChannel; import org.springframework.messaging.SubscribableChannel; +import org.springframework.util.ClassUtils; /** * Utilities for adapting integration components to/from reactive types. @@ -60,9 +66,40 @@ public final class IntegrationReactiveUtils { */ public static final Duration DEFAULT_DELAY_WHEN_EMPTY = Duration.ofSeconds(1); + /** + * The indicator that {@code io.micrometer:context-propagation} library is on classpath. + * @since 6.2.5 + */ + public static final boolean isContextPropagationPresent = ClassUtils.isPresent( + "io.micrometer.context.ContextSnapshot", IntegrationReactiveUtils.class.getClassLoader()); + + private static final ContextSnapshotFactory CONTEXT_SNAPSHOT_FACTORY = ContextSnapshotFactory.builder().build(); + private IntegrationReactiveUtils() { } + /** + * Capture a Reactor {@link ContextView} from the current thread local state + * according to the {@link ContextSnapshotFactory} logic. + * @return the Reactor {@link ContextView} from the current thread local state. + * @since 6.2.5 + */ + public static ContextView captureReactorContext() { + return CONTEXT_SNAPSHOT_FACTORY.captureAll().updateContext(Context.empty()); + } + + /** + * Populate thread local variables from the provided Reactor {@link ContextView} + * according to the {@link ContextSnapshotFactory} logic. + * @param context the Reactor {@link ContextView} to populate from. + * @return the {@link io.micrometer.context.ContextSnapshot.Scope} as a {@link AutoCloseable} + * to not pollute the target classpath. Can be cast if necessary. + * @since 6.2.5 + */ + public static AutoCloseable setThreadLocalsFromReactorContext(ContextView context) { + return CONTEXT_SNAPSHOT_FACTORY.setThreadLocalsFrom(context); + } + /** * Wrap a provided {@link MessageSource} into a {@link Flux} for pulling the on demand. * When {@link MessageSource#receive()} returns {@code null}, the source {@link Mono} @@ -137,8 +174,17 @@ public final class IntegrationReactiveUtils { return Flux.defer(() -> { Sinks.Many> sink = Sinks.many().unicast().onBackpressureError(); MessageHandler messageHandler = (message) -> { + Message messageToEmit = message; + if (IntegrationReactiveUtils.isContextPropagationPresent) { + ContextView contextView = IntegrationReactiveUtils.captureReactorContext(); + if (!contextView.isEmpty()) { + messageToEmit = MutableMessageBuilder.fromMessage(message) + .setHeader(IntegrationMessageHeaderAccessor.REACTOR_CONTEXT, contextView) + .build(); + } + } while (true) { - switch (sink.tryEmitNext((Message) message)) { + switch (sink.tryEmitNext((Message) messageToEmit)) { case FAIL_NON_SERIALIZED: case FAIL_OVERFLOW: LockSupport.parkNanos(1000); // NOSONAR diff --git a/spring-integration-core/src/test/java/org/springframework/integration/support/management/observation/IntegrationObservabilityZipkinTests.java b/spring-integration-core/src/test/java/org/springframework/integration/support/management/observation/IntegrationObservabilityZipkinTests.java index 530eafdf7a..c6842d3cd4 100644 --- a/spring-integration-core/src/test/java/org/springframework/integration/support/management/observation/IntegrationObservabilityZipkinTests.java +++ b/spring-integration-core/src/test/java/org/springframework/integration/support/management/observation/IntegrationObservabilityZipkinTests.java @@ -18,25 +18,31 @@ package org.springframework.integration.support.management.observation; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; import io.micrometer.common.KeyValues; import io.micrometer.core.tck.MeterRegistryAssert; +import io.micrometer.observation.Observation; import io.micrometer.observation.ObservationRegistry; import io.micrometer.tracing.Span; import io.micrometer.tracing.test.SampleTestRunner; import io.micrometer.tracing.test.simple.SpansAssert; +import reactor.core.publisher.Mono; import org.springframework.beans.factory.annotation.Qualifier; import org.springframework.context.annotation.AnnotationConfigApplicationContext; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; +import org.springframework.integration.annotation.BridgeTo; import org.springframework.integration.annotation.EndpointId; import org.springframework.integration.annotation.Poller; import org.springframework.integration.annotation.ServiceActivator; import org.springframework.integration.channel.NullChannel; +import org.springframework.integration.channel.PublishSubscribeChannel; import org.springframework.integration.channel.QueueChannel; import org.springframework.integration.config.EnableIntegration; import org.springframework.integration.config.EnableIntegrationManagement; +import org.springframework.integration.dsl.IntegrationFlow; import org.springframework.integration.gateway.MessagingGatewaySupport; import org.springframework.integration.handler.BridgeHandler; import org.springframework.integration.handler.advice.HandleMessageAdvice; @@ -46,6 +52,7 @@ import org.springframework.messaging.PollableChannel; import org.springframework.messaging.support.GenericMessage; import static org.assertj.core.api.Assertions.assertThat; +import static org.awaitility.Awaitility.await; /** * @author Artem Bilan @@ -93,6 +100,8 @@ public class IntegrationObservabilityZipkinTests extends SampleTestRunner { assertThat(receive).isNull(); assertThat(configuration.observedHandlerLatch.await(10, TimeUnit.SECONDS)).isTrue(); + + await().untilAsserted(() -> assertThat(configuration.observationReference.get()).isNotNull()); } SpansAssert.assertThat(bb.getFinishedSpans()) @@ -110,7 +119,7 @@ public class IntegrationObservabilityZipkinTests extends SampleTestRunner { .hasTag(IntegrationObservation.ProducerTags.COMPONENT_NAME.asString(), "queueChannel") .hasTag(IntegrationObservation.ProducerTags.COMPONENT_TYPE.asString(), "producer") .hasKindEqualTo(Span.Kind.PRODUCER)) - .hasSize(3); + .hasSize(4); MeterRegistryAssert.assertThat(getMeterRegistry()) .hasTimerWithNameAndTags("spring.integration.handler", @@ -126,7 +135,7 @@ public class IntegrationObservabilityZipkinTests extends SampleTestRunner { @EnableIntegration @EnableIntegrationManagement( observationPatterns = { - "${spring.integration.management.observation-patterns:testInboundGateway,skippedObservationInboundGateway,queueChannel,observedEndpoint}", + "${spring.integration.management.observation-patterns:testInboundGateway,skippedObservationInboundGateway,queueChannel,observedEndpoint,publishSubscribeChannel}", "${spring.integration.management.observation-patterns:}" }) public static class ObservationIntegrationTestConfiguration { @@ -168,8 +177,10 @@ public class IntegrationObservabilityZipkinTests extends SampleTestRunner { @ServiceActivator(inputChannel = "queueChannel", poller = @Poller(fixedDelay = "100"), adviceChain = "observedHandlerAdvice") - BridgeHandler bridgeHandler() { - return new BridgeHandler(); + BridgeHandler bridgeHandler(PublishSubscribeChannel publishSubscribeChannel) { + BridgeHandler bridgeHandler = new BridgeHandler(); + bridgeHandler.setOutputChannel(publishSubscribeChannel); + return bridgeHandler; } @Bean @@ -184,6 +195,26 @@ public class IntegrationObservabilityZipkinTests extends SampleTestRunner { }; } + @Bean + @BridgeTo + PublishSubscribeChannel publishSubscribeChannel() { + return new PublishSubscribeChannel(); + } + + AtomicReference observationReference = new AtomicReference<>(); + + @Bean + IntegrationFlow handleReactiveFlow(PublishSubscribeChannel publishSubscribeChannel, + ObservationRegistry observationRegistry) { + + return IntegrationFlow.from(publishSubscribeChannel) + .handleReactive(m -> + Mono.just("Hi There") + .doOnSuccess(val -> + observationReference.set(observationRegistry.getCurrentObservation())) + .then()); + } + } private static class TestMessagingGatewaySupport extends MessagingGatewaySupport {