Merge branch '5.2.x'

Author: rstoyanchev

spring-core/src/main/java/org/springframework/core/codec/AbstractSingleValueEncoder.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2019 the original author or authors.
+ * Copyright 2002-2020 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.
@@ -24,6 +24,7 @@
 import org.springframework.core.ResolvableType;
 import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.core.io.buffer.DataBufferFactory;
+import org.springframework.core.io.buffer.DataBufferUtils;
 import org.springframework.core.io.buffer.PooledDataBuffer;
 import org.springframework.lang.Nullable;
 import org.springframework.util.MimeType;
@@ -51,7 +52,7 @@ public final Flux<DataBuffer> encode(Publisher<? extends T> inputStream, DataBuf
 		return Flux.from(inputStream)
 				.take(1)
 				.concatMap(value -> encode(value, bufferFactory, elementType, mimeType, hints))
-				.doOnDiscard(PooledDataBuffer.class, PooledDataBuffer::release);
+				.doOnDiscard(PooledDataBuffer.class, DataBufferUtils::release);
 	}
 
 	/**

spring-core/src/main/java/org/springframework/core/io/buffer/DataBufferUtils.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2019 the original author or authors.
+ * Copyright 2002-2020 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.
@@ -39,6 +39,9 @@
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 
+import io.netty.util.IllegalReferenceCountException;
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
 import org.reactivestreams.Publisher;
 import org.reactivestreams.Subscription;
 import reactor.core.publisher.BaseSubscriber;
@@ -60,6 +63,8 @@
  */
 public abstract class DataBufferUtils {
 
+	private final static Log logger = LogFactory.getLog(DataBufferUtils.class);
+
 	private static final Consumer<DataBuffer> RELEASE_CONSUMER = DataBufferUtils::release;
 
 
@@ -494,7 +499,15 @@ public static boolean release(@Nullable DataBuffer dataBuffer) {
 		if (dataBuffer instanceof PooledDataBuffer) {
 			PooledDataBuffer pooledDataBuffer = (PooledDataBuffer) dataBuffer;
 			if (pooledDataBuffer.isAllocated()) {
-				return pooledDataBuffer.release();
+				try {
+					return pooledDataBuffer.release();
+				}
+				catch (IllegalReferenceCountException ex) {
+					if (logger.isDebugEnabled()) {
+						logger.debug("RefCount already at 0", ex);
+					}
+					return false;
+				}
 			}
 		}
 		return false;
@@ -523,7 +536,6 @@ public static Consumer<DataBuffer> releaseConsumer() {
 	 * @return a buffer that is composed from the {@code dataBuffers} argument
 	 * @since 5.0.3
 	 */
-	@SuppressWarnings("unchecked")
 	public static Mono<DataBuffer> join(Publisher<? extends DataBuffer> dataBuffers) {
 		return join(dataBuffers, -1);
 	}

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2018 the original author or authors.
+ * Copyright 2002-2020 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.
@@ -17,15 +17,13 @@
 package org.springframework.http.client.reactive;
 
 import java.net.URI;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
 
-import io.netty.buffer.ByteBufAllocator;
 import reactor.core.publisher.Mono;
-import reactor.netty.NettyInbound;
 import reactor.netty.NettyOutbound;
 import reactor.netty.http.client.HttpClient;
 import reactor.netty.http.client.HttpClientRequest;
-import reactor.netty.http.client.HttpClientResponse;
 import reactor.netty.resources.ConnectionProvider;
 import reactor.netty.resources.LoopResources;
 
@@ -104,12 +102,23 @@ public Mono<ClientHttpResponse> connect(HttpMethod method, URI uri,
 			return Mono.error(new IllegalArgumentException("URI is not absolute: " + uri));
 		}
 
+		AtomicReference<ReactorClientHttpResponse> responseRef = new AtomicReference<>();
+
 		return this.httpClient
 				.request(io.netty.handler.codec.http.HttpMethod.valueOf(method.name()))
 				.uri(uri.toString())
 				.send((request, outbound) -> requestCallback.apply(adaptRequest(method, uri, request, outbound)))
-				.responseConnection((res, con) -> Mono.just(adaptResponse(res, con.inbound(), con.outbound().alloc())))
-				.next();
+				.responseConnection((response, connection) -> {
+					responseRef.set(new ReactorClientHttpResponse(response, connection));
+					return Mono.just((ClientHttpResponse) responseRef.get());
+				})
+				.next()
+				.doOnCancel(() -> {
+					ReactorClientHttpResponse response = responseRef.get();
+					if (response != null && response.bodyNotSubscribed()) {
+						response.getConnection().dispose();
+					}
+				});
 	}
 
 	private ReactorClientHttpRequest adaptRequest(HttpMethod method, URI uri, HttpClientRequest request,
@@ -118,10 +127,4 @@ private ReactorClientHttpRequest adaptRequest(HttpMethod method, URI uri, HttpCl
 		return new ReactorClientHttpRequest(method, uri, request, nettyOutbound);
 	}
 
-	private ClientHttpResponse adaptResponse(HttpClientResponse response, NettyInbound nettyInbound,
-			ByteBufAllocator allocator) {
-
-		return new ReactorClientHttpResponse(response, nettyInbound, allocator);
-	}
-
 }

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

@@ -17,10 +17,11 @@
 package org.springframework.http.client.reactive;
 
 import java.util.Collection;
-import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.BiFunction;
 
-import io.netty.buffer.ByteBufAllocator;
 import reactor.core.publisher.Flux;
+import reactor.netty.Connection;
 import reactor.netty.NettyInbound;
 import reactor.netty.http.client.HttpClientResponse;
 
@@ -48,16 +49,24 @@ class ReactorClientHttpResponse implements ClientHttpResponse {
 
 	private final NettyDataBufferFactory bufferFactory;
 
+	private final Connection connection;
+
 	private final HttpHeaders headers;
 
-	private final AtomicBoolean rejectSubscribers = new AtomicBoolean();
+	// 0 - not subscribed, 1 - subscribed, 2 - cancelled
+	private final AtomicInteger state = new AtomicInteger(0);
 
 
-	public ReactorClientHttpResponse(HttpClientResponse response, NettyInbound inbound, ByteBufAllocator alloc) {
+	/**
+	 * Constructor that matches the inputs from
+	 * {@link reactor.netty.http.client.HttpClient.ResponseReceiver#responseConnection(BiFunction)}.
+	 * @since 5.3
+	 */
+	public ReactorClientHttpResponse(HttpClientResponse response, Connection connection) {
 		this.response = response;
-		this.inbound = inbound;
-		this.bufferFactory = new NettyDataBufferFactory(alloc);
-
+		this.inbound = connection.inbound();
+		this.bufferFactory = new NettyDataBufferFactory(connection.outbound().alloc());
+		this.connection = connection;
 		MultiValueMap<String, String> adapter = new NettyHeadersAdapter(response.responseHeaders());
 		this.headers = HttpHeaders.readOnlyHttpHeaders(adapter);
 	}
@@ -67,17 +76,17 @@ public ReactorClientHttpResponse(HttpClientResponse response, NettyInbound inbou
 	public Flux<DataBuffer> getBody() {
 		return this.inbound.receive()
 				.doOnSubscribe(s -> {
-					if (this.rejectSubscribers.get()) {
-						throw new IllegalStateException("The client response body can only be consumed once.");
+					if (!this.state.compareAndSet(0, 1)) {
+						// https://github.com/reactor/reactor-netty/issues/503
+						// FluxReceive rejects multiple subscribers, but not after a cancel().
+						// Subsequent subscribers after cancel() will not be rejected, but will hang instead.
+						// So we need to reject once in cancelled state.
+						if (this.state.get() == 2) {
+							throw new IllegalStateException("The client response body can only be consumed once.");
+						}
 					}
 				})
-				.doOnCancel(() ->
-					// https://github.com/reactor/reactor-netty/issues/503
-					// FluxReceive rejects multiple subscribers, but not after a cancel().
-					// Subsequent subscribers after cancel() will not be rejected, but will hang instead.
-					// So we need to intercept and reject them in that case.
-					this.rejectSubscribers.set(true)
-				)
+				.doOnCancel(() -> this.state.compareAndSet(1, 2))
 				.map(byteBuf -> {
 					byteBuf.retain();
 					return this.bufferFactory.wrap(byteBuf);
@@ -114,6 +123,20 @@ public MultiValueMap<String, ResponseCookie> getCookies() {
 		return CollectionUtils.unmodifiableMultiValueMap(result);
 	}
 
+	/**
+	 * For use by {@link ReactorClientHttpConnector}.
+	 */
+	boolean bodyNotSubscribed() {
+		return this.state.get() == 0;
+	}
+
+	/**
+	 * For use by {@link ReactorClientHttpConnector}.
+	 */
+	Connection getConnection() {
+		return this.connection;
+	}
+
 	@Override
 	public String toString() {
 		return "ReactorClientHttpResponse{" +

spring-web/src/main/java/org/springframework/http/codec/EncoderHttpMessageWriter.java

@@ -29,6 +29,7 @@
 import org.springframework.core.codec.Encoder;
 import org.springframework.core.codec.Hints;
 import org.springframework.core.io.buffer.DataBuffer;
+import org.springframework.core.io.buffer.DataBufferUtils;
 import org.springframework.core.io.buffer.PooledDataBuffer;
 import org.springframework.http.HttpLogging;
 import org.springframework.http.MediaType;
@@ -126,13 +127,13 @@ public Mono<Void> write(Publisher<? extends T> inputStream, ResolvableType eleme
 					.flatMap(buffer -> {
 						message.getHeaders().setContentLength(buffer.readableByteCount());
 						return message.writeWith(Mono.just(buffer)
-								.doOnDiscard(PooledDataBuffer.class, PooledDataBuffer::release));
+								.doOnDiscard(PooledDataBuffer.class, DataBufferUtils::release));
 					});
 		}
 
 		if (isStreamingMediaType(contentType)) {
 			return message.writeAndFlushWith(body.map(buffer ->
-					Mono.just(buffer).doOnDiscard(PooledDataBuffer.class, PooledDataBuffer::release)));
+					Mono.just(buffer).doOnDiscard(PooledDataBuffer.class, DataBufferUtils::release)));
 		}
 
 		return message.writeWith(body);

spring-web/src/main/java/org/springframework/http/codec/multipart/MultipartHttpMessageWriter.java

@@ -38,6 +38,7 @@
 import org.springframework.core.io.Resource;
 import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.core.io.buffer.DataBufferFactory;
+import org.springframework.core.io.buffer.DataBufferUtils;
 import org.springframework.core.io.buffer.PooledDataBuffer;
 import org.springframework.core.log.LogFormatUtils;
 import org.springframework.http.HttpEntity;
@@ -206,7 +207,7 @@ private Mono<Void> writeMultipart(MultiValueMap<String, ?> map,
 		Flux<DataBuffer> body = Flux.fromIterable(map.entrySet())
 				.concatMap(entry -> encodePartValues(boundary, entry.getKey(), entry.getValue(), bufferFactory))
 				.concatWith(generateLastLine(boundary, bufferFactory))
-				.doOnDiscard(PooledDataBuffer.class, PooledDataBuffer::release);
+				.doOnDiscard(PooledDataBuffer.class, DataBufferUtils::release);
 
 		return outputMessage.writeWith(body);
 	}