From dbd46dfe3dad7c1fc61bb83daaf7d53959237358 Mon Sep 17 00:00:00 2001 From: Artem Bilan Date: Wed, 7 Feb 2024 09:48:54 -0500 Subject: [PATCH] GH-8889: FluxMCh: propagate context per message Fixes: #8889 * Gather the current thread local state for the message into `IntegrationMessageHeaderAccessor.REACTOR_CONTEXT` header before `this.sink.tryEmitNext(messageToEmit)` * Restore thread locals for the current message in the `ReactiveStreamsConsumer.SubscriberDecorator` just before `this.delegate.onNext(messageToDeliver)` **Cherry-pick to `6.2.x`** --- .../channel/FluxMessageChannel.java | 34 ++++++++++++-- .../endpoint/ReactiveStreamsConsumer.java | 47 +++++++++++++++++-- 2 files changed, 74 insertions(+), 7 deletions(-) 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 34d80b52d5..fa45245787 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 @@ -1,5 +1,5 @@ /* - * Copyright 2015-2023 the original author or authors. + * Copyright 2015-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. @@ -21,6 +21,7 @@ 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; @@ -30,14 +31,17 @@ 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.MessageBuilder; +import org.springframework.integration.support.MutableMessageBuilder; 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 @@ -52,6 +56,9 @@ import org.springframework.util.Assert; 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); @@ -83,7 +90,16 @@ public class FluxMessageChannel extends AbstractMessageChannel } private boolean tryEmitMessage(Message message) { - return switch (this.sink.tryEmitNext(message)) { + Message messageToEmit = message; + if (isContextPropagationPresent) { + ContextView contextView = ContextSnapshotHelper.captureContext(); + if (!contextView.isEmpty()) { + messageToEmit = MutableMessageBuilder.fromMessage(message) + .setHeader(IntegrationMessageHeaderAccessor.REACTOR_CONTEXT, contextView) + .build(); + } + } + return switch (this.sink.tryEmitNext(messageToEmit)) { case OK -> true; case FAIL_NON_SERIALIZED, FAIL_OVERFLOW -> false; case FAIL_ZERO_SUBSCRIBER -> @@ -154,7 +170,7 @@ public class FluxMessageChannel extends AbstractMessageChannel // We have just restored Reactor context, so no need in a header anymore. if (messageToSend.getHeaders().containsKey(IntegrationMessageHeaderAccessor.REACTOR_CONTEXT)) { messageToSend = - MessageBuilder.fromMessage(message) + MutableMessageBuilder.fromMessage(message) .removeHeader(IntegrationMessageHeaderAccessor.REACTOR_CONTEXT) .build(); } @@ -180,4 +196,14 @@ 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 e36b8ce332..fc5a3ade22 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 @@ -1,5 +1,5 @@ /* - * Copyright 2016-2022 the original author or authors. + * Copyright 2016-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.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; @@ -26,13 +27,16 @@ import reactor.core.CoreSubscriber; import reactor.core.Disposable; import reactor.core.publisher.BaseSubscriber; import reactor.core.publisher.Flux; +import reactor.util.context.ContextView; import org.springframework.context.Lifecycle; +import org.springframework.integration.IntegrationMessageHeaderAccessor; import org.springframework.integration.channel.ChannelUtils; import org.springframework.integration.channel.NullChannel; import org.springframework.integration.core.MessageProducer; import org.springframework.integration.handler.ReactiveMessageHandlerAdapter; import org.springframework.integration.router.MessageRouter; +import org.springframework.integration.support.MutableMessageBuilder; import org.springframework.integration.util.IntegrationReactiveUtils; import org.springframework.lang.Nullable; import org.springframework.messaging.Message; @@ -40,6 +44,7 @@ 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; @@ -54,6 +59,9 @@ 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; @@ -289,10 +297,33 @@ public class ReactiveStreamsConsumer extends AbstractEndpoint implements Integra this.delegate.onSubscribe(subscription); } + @SuppressWarnings("try") @Override - protected void hookOnNext(Message value) { + protected void hookOnNext(Message message) { + Message messageToDeliver = message; + + if (isContextPropagationPresent) { + ContextView reactorContext = message.getHeaders() + .get(IntegrationMessageHeaderAccessor.REACTOR_CONTEXT, ContextView.class); + + if (reactorContext != null) { + messageToDeliver = + MutableMessageBuilder.fromMessage(message) + .removeHeader(IntegrationMessageHeaderAccessor.REACTOR_CONTEXT) + .build(); + + try (AutoCloseable scope = ContextSnapshotHelper.setContext(reactorContext)) { + this.delegate.onNext(messageToDeliver); + } + catch (Exception ex) { + this.errorHandler.handleError(ex); + } + return; + } + } + try { - this.delegate.onNext(value); + this.delegate.onNext(messageToDeliver); } catch (Exception ex) { this.errorHandler.handleError(ex); @@ -306,4 +337,14 @@ 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); + } + + } + }