Apply attributes to the underlying HTTP request

See gh-29958

spring-web/src/main/java/org/springframework/http/client/reactive/AbstractClientHttpRequest.java

@@ -17,7 +17,10 @@
 package org.springframework.http.client.reactive;
 
 import java.util.ArrayList;
+import java.util.Collections;
+import java.util.LinkedHashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Supplier;
 
@@ -55,6 +58,10 @@ public abstract class AbstractClientHttpRequest implements ClientHttpRequest {
 
 	private final MultiValueMap<String, HttpCookie> cookies;
 
+	private final Map<String, Object> attributes;
+
+	private final boolean applyAttributes;
+
 	private final AtomicReference<State> state = new AtomicReference<>(State.NEW);
 
 	private final List<Supplier<? extends Publisher<Void>>> commitActions = new ArrayList<>(4);
@@ -64,13 +71,19 @@ public abstract class AbstractClientHttpRequest implements ClientHttpRequest {
 
 
 	public AbstractClientHttpRequest() {
-		this(new HttpHeaders());
+		this(new HttpHeaders(), false);
 	}
 
-	public AbstractClientHttpRequest(HttpHeaders headers) {
+	public AbstractClientHttpRequest(boolean applyAttributes) {
+		this(new HttpHeaders(), applyAttributes);
+	}
+
+	public AbstractClientHttpRequest(HttpHeaders headers, boolean applyAttributes) {
 		Assert.notNull(headers, "HttpHeaders must not be null");
 		this.headers = headers;
 		this.cookies = new LinkedMultiValueMap<>();
+		this.attributes = new LinkedHashMap<>();
+		this.applyAttributes = applyAttributes;
 	}
 
 
@@ -106,6 +119,14 @@ public abstract class AbstractClientHttpRequest implements ClientHttpRequest {
 		return this.cookies;
 	}
 
+	@Override
+	public Map<String, Object> getAttributes() {
+		if (State.COMMITTED.equals(this.state.get())) {
+			return Collections.unmodifiableMap(this.attributes);
+		}
+		return this.attributes;
+	}
+
 	@Override
 	public void beforeCommit(Supplier<? extends Mono<Void>> action) {
 		Assert.notNull(action, "Action must not be null");
@@ -140,6 +161,9 @@ public abstract class AbstractClientHttpRequest implements ClientHttpRequest {
 				Mono.fromRunnable(() -> {
 					applyHeaders();
 					applyCookies();
+					if (this.applyAttributes) {
+						applyAttributes();
+					}
 					this.state.set(State.COMMITTED);
 				}));
 
@@ -168,4 +192,10 @@ public abstract class AbstractClientHttpRequest implements ClientHttpRequest {
 	 */
 	protected abstract void applyCookies();
 
+	/**
+	 * Add additional attributes from {@link #getAttributes()} to the underlying request.
+	 * This method is called once only.
+	 */
+	protected abstract void applyAttributes();
+
 }

spring-web/src/main/java/org/springframework/http/client/reactive/ClientHttpConnector.java

@@ -48,4 +48,14 @@ public interface ClientHttpConnector {
 	Mono<ClientHttpResponse> connect(HttpMethod method, URI uri,
 			Function<? super ClientHttpRequest, Mono<Void>> requestCallback);
 
+	/**
+	 * Set whether or not attributes should be applied to the underlying http-client library request.
+	 */
+	void setApplyAttributes(boolean applyAttributes);
+
+	/**
+	 * Whether or not attributes should be applied to the underlying http-client library request.
+	 */
+	boolean getApplyAttributes();
+
 }

spring-web/src/main/java/org/springframework/http/client/reactive/ClientHttpRequest.java

@@ -17,6 +17,7 @@
 package org.springframework.http.client.reactive;
 
 import java.net.URI;
+import java.util.Map;
 
 import org.springframework.http.HttpCookie;
 import org.springframework.http.HttpMethod;
@@ -54,4 +55,9 @@ public interface ClientHttpRequest extends ReactiveHttpOutputMessage {
 	 */
 	<T> T getNativeRequest();
 
+	/**
+	 * Return a mutable map of the request attributes.
+	 */
+	Map<String, Object> getAttributes();
+
 }

spring-web/src/main/java/org/springframework/http/client/reactive/ClientHttpRequestDecorator.java

@@ -17,6 +17,7 @@
 package org.springframework.http.client.reactive;
 
 import java.net.URI;
+import java.util.Map;
 import java.util.function.Supplier;
 
 import org.reactivestreams.Publisher;
@@ -85,6 +86,11 @@ public class ClientHttpRequestDecorator implements ClientHttpRequest {
 		return this.delegate.getNativeRequest();
 	}
 
+	@Override
+	public Map<String, Object> getAttributes() {
+		return this.delegate.getAttributes();
+	}
+
 	@Override
 	public void beforeCommit(Supplier<? extends Mono<Void>> action) {
 		this.delegate.beforeCommit(action);

spring-web/src/main/java/org/springframework/http/client/reactive/HttpComponentsClientHttpConnector.java

@@ -59,6 +59,7 @@ public class HttpComponentsClientHttpConnector implements ClientHttpConnector, C
 
 	private DataBufferFactory dataBufferFactory = DefaultDataBufferFactory.sharedInstance;
 
+	private boolean applyAttributes = true;
 
 	/**
 	 * Default constructor that creates and starts a new instance of {@link CloseableHttpAsyncClient}.
@@ -67,6 +68,7 @@ public class HttpComponentsClientHttpConnector implements ClientHttpConnector, C
 		this(HttpAsyncClients.createDefault());
 	}
 
+
 	/**
 	 * Constructor with a pre-configured {@link CloseableHttpAsyncClient} instance.
 	 * @param client the client to use
@@ -111,11 +113,22 @@ public class HttpComponentsClientHttpConnector implements ClientHttpConnector, C
 			context.setCookieStore(new BasicCookieStore());
 		}
 
-		HttpComponentsClientHttpRequest request =
-				new HttpComponentsClientHttpRequest(method, uri, context, this.dataBufferFactory);
+		HttpComponentsClientHttpRequest request = new HttpComponentsClientHttpRequest(
+				method, uri, context, this.dataBufferFactory, this.applyAttributes);
+
 		return requestCallback.apply(request).then(Mono.defer(() -> execute(request, context)));
 	}
 
+	@Override
+	public void setApplyAttributes(boolean applyAttributes) {
+		this.applyAttributes = applyAttributes;
+	}
+
+	@Override
+	public boolean getApplyAttributes() {
+		return this.applyAttributes;
+	}
+
 	private Mono<ClientHttpResponse> execute(HttpComponentsClientHttpRequest request, HttpClientContext context) {
 		AsyncRequestProducer requestProducer = request.toRequestProducer();
 

spring-web/src/main/java/org/springframework/http/client/reactive/HttpComponentsClientHttpRequest.java

@@ -66,8 +66,8 @@ class HttpComponentsClientHttpRequest extends AbstractClientHttpRequest {
 
 
 	public HttpComponentsClientHttpRequest(HttpMethod method, URI uri, HttpClientContext context,
-			DataBufferFactory dataBufferFactory) {
-
+			DataBufferFactory dataBufferFactory, boolean applyAttributes) {
+		super(applyAttributes);
 		this.context = context;
 		this.httpRequest = new BasicHttpRequest(method.name(), uri);
 		this.dataBufferFactory = dataBufferFactory;
@@ -157,6 +157,18 @@ class HttpComponentsClientHttpRequest extends AbstractClientHttpRequest {
 				});
 	}
 
+	/**
+	 * Applies the attributes to the {@link HttpClientContext}.
+	 */
+	@Override
+	protected void applyAttributes() {
+		getAttributes().forEach((key, value) -> {
+			if(this.context.getAttribute(key) == null) {
+				this.context.setAttribute(key, value);
+			}
+		});
+	}
+
 	@Override
 	protected HttpHeaders initReadOnlyHeaders() {
 		return HttpHeaders.readOnlyHttpHeaders(new HttpComponentsHeadersAdapter(this.httpRequest));

spring-web/src/main/java/org/springframework/http/client/reactive/JdkClientHttpConnector.java

@@ -96,7 +96,7 @@ public class JdkClientHttpConnector implements ClientHttpConnector {
 	public Mono<ClientHttpResponse> connect(
 			HttpMethod method, URI uri, Function<? super ClientHttpRequest, Mono<Void>> requestCallback) {
 
-		JdkClientHttpRequest jdkClientHttpRequest = new JdkClientHttpRequest(method, uri, this.bufferFactory);
+		JdkClientHttpRequest jdkClientHttpRequest = new JdkClientHttpRequest(method, uri, this.bufferFactory, getApplyAttributes());
 
 		return requestCallback.apply(jdkClientHttpRequest).then(Mono.defer(() -> {
 			HttpRequest httpRequest = jdkClientHttpRequest.getNativeRequest();
@@ -109,4 +109,19 @@ public class JdkClientHttpConnector implements ClientHttpConnector {
 		}));
 	}
 
+	/**
+	 * Sets nothing, since {@link JdkClientHttpConnector} does not offer any possibility to add attributes.
+	 */
+	@Override
+	public void setApplyAttributes(boolean applyAttributes) {
+	}
+
+	/**
+	 * Returns false, since {@link JdkClientHttpConnector} does not offer any possibility to add attributes.
+	 */
+	@Override
+	public boolean getApplyAttributes() {
+		return false;
+	}
+
 }

spring-web/src/main/java/org/springframework/http/client/reactive/JdkClientHttpRequest.java

@@ -56,7 +56,8 @@ class JdkClientHttpRequest extends AbstractClientHttpRequest {
 	private final HttpRequest.Builder builder;
 
 
-	public JdkClientHttpRequest(HttpMethod httpMethod, URI uri, DataBufferFactory bufferFactory) {
+	public JdkClientHttpRequest(HttpMethod httpMethod, URI uri, DataBufferFactory bufferFactory, boolean applyAttributes) {
+		super(applyAttributes);
 		Assert.notNull(httpMethod, "HttpMethod is required");
 		Assert.notNull(uri, "URI is required");
 		Assert.notNull(bufferFactory, "DataBufferFactory is required");
@@ -112,6 +113,15 @@ class JdkClientHttpRequest extends AbstractClientHttpRequest {
 				.flatMap(List::stream).map(HttpCookie::toString).collect(Collectors.joining(";")));
 	}
 
+	/**
+	 * Not implemented, since {@link HttpRequest} does not offer any possibility to add request attributes.
+	 */
+	@Override
+	protected void applyAttributes() {
+		// TODO
+		throw new RuntimeException(String.format("Using attributes is not available for %s", HttpRequest.class.getName()));
+	}
+
 	@Override
 	public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
 		return doCommit(() -> {

spring-web/src/main/java/org/springframework/http/client/reactive/JettyClientHttpConnector.java

@@ -52,6 +52,8 @@ public class JettyClientHttpConnector implements ClientHttpConnector {
 
 	private DataBufferFactory bufferFactory = DefaultDataBufferFactory.sharedInstance;
 
+	private boolean applyAttributes = true;
+
 
 	/**
 	 * Default constructor that creates a new instance of {@link HttpClient}.
@@ -126,11 +128,21 @@ public class JettyClientHttpConnector implements ClientHttpConnector {
 		}
 
 		Request jettyRequest = this.httpClient.newRequest(uri).method(method.toString());
-		JettyClientHttpRequest request = new JettyClientHttpRequest(jettyRequest, this.bufferFactory);
+		JettyClientHttpRequest request = new JettyClientHttpRequest(jettyRequest, this.bufferFactory, getApplyAttributes());
 
 		return requestCallback.apply(request).then(execute(request));
 	}
 
+	@Override
+	public void setApplyAttributes(boolean applyAttributes) {
+		this.applyAttributes = applyAttributes;
+	}
+
+	@Override
+	public boolean getApplyAttributes() {
+		return this.applyAttributes;
+	}
+
 	private Mono<ClientHttpResponse> execute(JettyClientHttpRequest request) {
 		return Mono.fromDirect(request.toReactiveRequest()
 				.response((reactiveResponse, chunkPublisher) -> {

spring-web/src/main/java/org/springframework/http/client/reactive/JettyClientHttpRequest.java

@@ -55,7 +55,8 @@ class JettyClientHttpRequest extends AbstractClientHttpRequest {
 	private final ReactiveRequest.Builder builder;
 
 
-	public JettyClientHttpRequest(Request jettyRequest, DataBufferFactory bufferFactory) {
+	public JettyClientHttpRequest(Request jettyRequest, DataBufferFactory bufferFactory, boolean applyAttributes) {
+		super(applyAttributes);
 		this.jettyRequest = jettyRequest;
 		this.bufferFactory = bufferFactory;
 		this.builder = ReactiveRequest.newBuilder(this.jettyRequest).abortOnCancel(true);
@@ -138,6 +139,14 @@ class JettyClientHttpRequest extends AbstractClientHttpRequest {
 				.forEach(this.jettyRequest::cookie);
 	}
 
+	/**
+	 * Applies the attributes to {@link Request#getAttributes()}.
+	 */
+	@Override
+	protected void applyAttributes() {
+		getAttributes().forEach(this.jettyRequest::attribute);
+	}
+
 	@Override
 	protected void applyHeaders() {
 		HttpHeaders headers = getHeaders();

spring-web/src/main/java/org/springframework/http/client/reactive/ReactorClientHttpConnector.java

@@ -66,6 +66,7 @@ public class ReactorClientHttpConnector implements ClientHttpConnector, SmartLif
 
 	private final Object lifecycleMonitor = new Object();
 
+	private boolean applyAttributes = true;
 
 	/**
 	 * Default constructor. Initializes {@link HttpClient} via:
@@ -170,10 +171,20 @@ public class ReactorClientHttpConnector implements ClientHttpConnector, SmartLif
 		return requestSender.uri(uri.toString());
 	}
 
+	@Override
+	public void setApplyAttributes(boolean applyAttributes) {
+		this.applyAttributes = applyAttributes;
+	}
+
+	@Override
+	public boolean getApplyAttributes() {
+		return this.applyAttributes;
+	}
+
 	private ReactorClientHttpRequest adaptRequest(HttpMethod method, URI uri, HttpClientRequest request,
 			NettyOutbound nettyOutbound) {
 
-		return new ReactorClientHttpRequest(method, uri, request, nettyOutbound);
+		return new ReactorClientHttpRequest(method, uri, request, nettyOutbound, this.applyAttributes);
 	}
 
 	@Override

spring-web/src/main/java/org/springframework/http/client/reactive/ReactorClientHttpRequest.java

@@ -18,13 +18,16 @@ package org.springframework.http.client.reactive;
 
 import java.net.URI;
 import java.nio.file.Path;
+import java.util.Map;
 
 import io.netty.buffer.ByteBuf;
 import io.netty.handler.codec.http.cookie.DefaultCookie;
+import io.netty.util.AttributeKey;
 import org.reactivestreams.Publisher;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.netty.NettyOutbound;
+import reactor.netty.channel.ChannelOperations;
 import reactor.netty.http.client.HttpClientRequest;
 
 import org.springframework.core.io.buffer.DataBuffer;
@@ -45,6 +48,8 @@ import org.springframework.http.support.Netty4HeadersAdapter;
  */
 class ReactorClientHttpRequest extends AbstractClientHttpRequest implements ZeroCopyHttpOutputMessage {
 
+	public static final String ATTRIBUTES_CHANNEL_KEY = "attributes";
+
 	private final HttpMethod httpMethod;
 
 	private final URI uri;
@@ -56,7 +61,8 @@ class ReactorClientHttpRequest extends AbstractClientHttpRequest implements Zero
 	private final NettyDataBufferFactory bufferFactory;
 
 
-	public ReactorClientHttpRequest(HttpMethod method, URI uri, HttpClientRequest request, NettyOutbound outbound) {
+	public ReactorClientHttpRequest(HttpMethod method, URI uri, HttpClientRequest request, NettyOutbound outbound, boolean applyAttributes) {
+		super(applyAttributes);
 		this.httpMethod = method;
 		this.uri = uri;
 		this.request = request;
@@ -135,6 +141,17 @@ class ReactorClientHttpRequest extends AbstractClientHttpRequest implements Zero
 		}));
 	}
 
+	/**
+	 * Applies the request attributes to the {@link reactor.netty.http.client.HttpClientRequest} by setting
+	 * a single {@link Map} into the {@link reactor.netty.channel.ChannelOperations#channel()},
+	 * with {@link io.netty5.util.AttributeKey#name()} equal to {@link #ATTRIBUTES_CHANNEL_KEY}.
+	 */
+	@Override
+	protected void applyAttributes() {
+		((ChannelOperations<?, ?>) this.request)
+				.channel().attr(AttributeKey.valueOf(ATTRIBUTES_CHANNEL_KEY)).set(getAttributes());
+	}
+
 	@Override
 	protected HttpHeaders initReadOnlyHeaders() {
 		return HttpHeaders.readOnlyHttpHeaders(new Netty4HeadersAdapter(this.request.requestHeaders()));

spring-web/src/main/java/org/springframework/http/client/reactive/ReactorNetty2ClientHttpConnector.java

@@ -46,6 +46,8 @@ public class ReactorNetty2ClientHttpConnector implements ClientHttpConnector {
 
 	private final HttpClient httpClient;
 
+	private boolean applyAttributes = true;
+
 
 	/**
 	 * Default constructor. Initializes {@link HttpClient} via:
@@ -126,10 +128,20 @@ public class ReactorNetty2ClientHttpConnector implements ClientHttpConnector {
 				});
 	}
 
+	@Override
+	public void setApplyAttributes(boolean applyAttributes) {
+		this.applyAttributes = applyAttributes;
+	}
+
+	@Override
+	public boolean getApplyAttributes() {
+		return this.applyAttributes;
+	}
+
 	private ReactorNetty2ClientHttpRequest adaptRequest(HttpMethod method, URI uri, HttpClientRequest request,
 			NettyOutbound nettyOutbound) {
 
-		return new ReactorNetty2ClientHttpRequest(method, uri, request, nettyOutbound);
+		return new ReactorNetty2ClientHttpRequest(method, uri, request, nettyOutbound, getApplyAttributes());
 	}
 
 }

spring-web/src/main/java/org/springframework/http/client/reactive/ReactorNetty2ClientHttpRequest.java

@@ -18,13 +18,16 @@ package org.springframework.http.client.reactive;
 
 import java.net.URI;
 import java.nio.file.Path;
+import java.util.Map;
 
 import io.netty5.buffer.Buffer;
 import io.netty5.handler.codec.http.headers.DefaultHttpCookiePair;
+import io.netty5.util.AttributeKey;
 import org.reactivestreams.Publisher;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.netty5.NettyOutbound;
+import reactor.netty5.channel.ChannelOperations;
 import reactor.netty5.http.client.HttpClientRequest;
 
 import org.springframework.core.io.buffer.DataBuffer;
@@ -46,6 +49,8 @@ import org.springframework.http.support.Netty5HeadersAdapter;
  */
 class ReactorNetty2ClientHttpRequest extends AbstractClientHttpRequest implements ZeroCopyHttpOutputMessage {
 
+	public static final String ATTRIBUTES_CHANNEL_KEY = "attributes";
+
 	private final HttpMethod httpMethod;
 
 	private final URI uri;
@@ -57,7 +62,8 @@ class ReactorNetty2ClientHttpRequest extends AbstractClientHttpRequest implement
 	private final Netty5DataBufferFactory bufferFactory;
 
 
-	public ReactorNetty2ClientHttpRequest(HttpMethod method, URI uri, HttpClientRequest request, NettyOutbound outbound) {
+	public ReactorNetty2ClientHttpRequest(HttpMethod method, URI uri, HttpClientRequest request, NettyOutbound outbound, boolean applyAttributes) {
+		super(applyAttributes);
 		this.httpMethod = method;
 		this.uri = uri;
 		this.request = request;
@@ -136,6 +142,17 @@ class ReactorNetty2ClientHttpRequest extends AbstractClientHttpRequest implement
 		}));
 	}
 
+	/**
+	 * Applies the request attributes to the {@link reactor.netty.http.client.HttpClientRequest} by setting
+	 * a single {@link Map} into the {@link reactor.netty.channel.ChannelOperations#channel()},
+	 * with {@link AttributeKey#name()} equal to {@link #ATTRIBUTES_CHANNEL_KEY}.
+	 */
+	@Override
+	protected void applyAttributes() {
+		((ChannelOperations<?, ?>) this.request)
+				.channel().attr(AttributeKey.valueOf(ATTRIBUTES_CHANNEL_KEY)).set(getAttributes());
+	}
+
 	@Override
 	protected HttpHeaders initReadOnlyHeaders() {
 		return HttpHeaders.readOnlyHttpHeaders(new Netty5HeadersAdapter(this.request.requestHeaders()));