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

GH-8889: FluxMCh: propagate context per message

Browse Source 
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`**
This commit is contained in:
Artem Bilan
2024-02-07 09:48:54 -05:00
parent 8b877fc887
commit dbd46dfe3d
2 changed files with 74 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

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

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

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