Spring/spring-integration
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

GH-9110: propagate Reactor context to ReactiveMessageHandler

Browse Source 
Fixes: #9110

* Move context propagation utilities to the `IntegrationReactiveUtils`
* Capture context into message header in the `IntegrationReactiveUtils.adaptSubscribableChannelToPublisher()`
before `sink.tryEmitNext()`
* Restore the context from message header in the `flatMap()` for `ReactiveStreamsConsumer.reactiveMessageHandler`

(cherry picked from commit 7abf1a53a7)
This commit is contained in:
Artem Bilan
2024-05-03 16:57:21 -04:00
committed by Spring Builds
parent c4ebec1a6a
commit 154e1ec795
4 changed files with 92 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
spring-integration-core/src/main/java/org/springframework/integration/channel/FluxMessageChannel.java
View File

@@ -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<Message<?>>, 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<Message<?>> 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());
}
}
}

24
spring-integration-core/src/main/java/org/springframework/integration/endpoint/ReactiveStreamsConsumer.java
View File

@@ -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<Message<Object>> 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);
}
}
}

50
spring-integration-core/src/main/java/org/springframework/integration/util/IntegrationReactiveUtils.java
View File

@@ -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<Message<T>> 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<T>) message)) {
switch (sink.tryEmitNext((Message<T>) messageToEmit)) {
case FAIL_NON_SERIALIZED:
case FAIL_OVERFLOW:
LockSupport.parkNanos(1000); // NOSONAR

39
spring-integration-core/src/test/java/org/springframework/integration/support/management/observation/IntegrationObservabilityZipkinTests.java
View File

@@ -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<Observation> 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 {