Correlate data buffers to request log messages

HttpMessageWriter implementations now attach the request log prefix
as a hint to created data buffers when the logger associated with
the writer is at DEBUG level.

Closes gh-26230

Author: rstoyanchev

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

@@ -21,6 +21,8 @@
 
 import org.apache.commons.logging.Log;
 
+import org.springframework.core.io.buffer.DataBuffer;
+import org.springframework.core.io.buffer.DataBufferUtils;
 import org.springframework.lang.Nullable;
 import org.springframework.util.CollectionUtils;
 
@@ -148,4 +150,22 @@ public static Map<String, Object> merge(Map<String, Object> hints, String hintNa
 		}
 	}
 
+	/**
+	 * If the hints contain a {@link #LOG_PREFIX_HINT} and the given logger has
+	 * DEBUG level enabled, apply the log prefix as a hint to the given buffer
+	 * via {@link DataBufferUtils#touch(DataBuffer, Object)}.
+	 * @param buffer the buffer to touch
+	 * @param hints the hints map to check for a log prefix
+	 * @param logger the logger whose level to check
+	 * @since 5.3.2
+	 */
+	public static void touchDataBuffer(DataBuffer buffer, @Nullable Map<String, Object> hints, Log logger) {
+		if (logger.isDebugEnabled() && hints != null) {
+			Object logPrefix = hints.get(LOG_PREFIX_HINT);
+			if (logPrefix != null) {
+				DataBufferUtils.touch(buffer, logPrefix);
+			}
+		}
+	}
+
 }

spring-core/src/main/java/org/springframework/core/codec/ResourceRegionEncoder.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.
@@ -130,6 +130,9 @@ private Flux<DataBuffer> writeResourceRegion(
 		}
 
 		Flux<DataBuffer> in = DataBufferUtils.read(resource, position, bufferFactory, this.bufferSize);
+		if (logger.isDebugEnabled()) {
+			in = in.doOnNext(buffer -> Hints.touchDataBuffer(buffer, hints, logger));
+		}
 		return DataBufferUtils.takeUntilByteCount(in, count);
 	}
 

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

@@ -487,6 +487,24 @@ public static <T extends DataBuffer> T retain(T dataBuffer) {
 		}
 	}
 
+	/**
+	 * Associate the given hint with the data buffer if it is a pooled buffer
+	 * and supports leak tracking.
+	 * @param dataBuffer the data buffer to attach the hint to
+	 * @param hint the hint to attach
+	 * @return the input buffer
+	 * @since 5.3.2
+	 */
+	@SuppressWarnings("unchecked")
+	public static <T extends DataBuffer> T touch(T dataBuffer, Object hint) {
+		if (dataBuffer instanceof PooledDataBuffer) {
+			return (T) ((PooledDataBuffer) dataBuffer).touch(hint);
+		}
+		else {
+			return dataBuffer;
+		}
+	}
+
 	/**
 	 * Release the given data buffer, if it is a {@link PooledDataBuffer} and
 	 * has been {@linkplain PooledDataBuffer#isAllocated() allocated}.

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

@@ -315,6 +315,12 @@ public PooledDataBuffer retain() {
 		return new NettyDataBuffer(this.byteBuf.retain(), this.dataBufferFactory);
 	}
 
+	@Override
+	public PooledDataBuffer touch(Object hint) {
+		this.byteBuf.touch(hint);
+		return this;
+	}
+
 	@Override
 	public boolean release() {
 		return this.byteBuf.release();

spring-core/src/main/java/org/springframework/core/io/buffer/PooledDataBuffer.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.
@@ -38,6 +38,13 @@ public interface PooledDataBuffer extends DataBuffer {
 	 */
 	PooledDataBuffer retain();
 
+	/**
+	 * Associate the given hint with the data buffer for debugging purposes.
+	 * @return this buffer
+	 * @since 5.3.2
+	 */
+	PooledDataBuffer touch(Object hint);
+
 	/**
 	 * Decrease the reference count for this buffer by one,
 	 * and deallocate it once the count reaches zero.

spring-core/src/testFixtures/java/org/springframework/core/testfixture/io/buffer/LeakAwareDataBuffer.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.
@@ -17,6 +17,7 @@
 package org.springframework.core.testfixture.io.buffer;
 
 import org.springframework.core.io.buffer.DataBuffer;
+import org.springframework.core.io.buffer.DataBufferUtils;
 import org.springframework.core.io.buffer.DataBufferWrapper;
 import org.springframework.core.io.buffer.PooledDataBuffer;
 import org.springframework.util.Assert;
@@ -67,19 +68,19 @@ public boolean isAllocated() {
 
 	@Override
 	public PooledDataBuffer retain() {
-		DataBuffer delegate = dataBuffer();
-		if (delegate instanceof PooledDataBuffer) {
-			((PooledDataBuffer) delegate).retain();
-		}
+		DataBufferUtils.retain(dataBuffer());
+		return this;
+	}
+
+	@Override
+	public PooledDataBuffer touch(Object hint) {
+		DataBufferUtils.touch(dataBuffer(), hint);
 		return this;
 	}
 
 	@Override
 	public boolean release() {
-		DataBuffer delegate = dataBuffer();
-		if (delegate instanceof PooledDataBuffer) {
-			((PooledDataBuffer) delegate).release();
-		}
+		DataBufferUtils.release(dataBuffer());
 		return isAllocated();
 	}
 

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

@@ -57,6 +57,9 @@
  */
 public class EncoderHttpMessageWriter<T> implements HttpMessageWriter<T> {
 
+	private static final Log logger = HttpLogging.forLogName(EncoderHttpMessageWriter.class);
+
+
 	private final Encoder<T> encoder;
 
 	private final List<MediaType> mediaTypes;
@@ -125,17 +128,23 @@ public Mono<Void> write(Publisher<? extends T> inputStream, ResolvableType eleme
 						return message.setComplete().then(Mono.empty());
 					}))
 					.flatMap(buffer -> {
+						Hints.touchDataBuffer(buffer, hints, logger);
 						message.getHeaders().setContentLength(buffer.readableByteCount());
 						return message.writeWith(Mono.just(buffer)
 								.doOnDiscard(PooledDataBuffer.class, DataBufferUtils::release));
 					});
 		}
 
 		if (isStreamingMediaType(contentType)) {
-			return message.writeAndFlushWith(body.map(buffer ->
-					Mono.just(buffer).doOnDiscard(PooledDataBuffer.class, DataBufferUtils::release)));
+			return message.writeAndFlushWith(body.map(buffer -> {
+				Hints.touchDataBuffer(buffer, hints, logger);
+				return Mono.just(buffer).doOnDiscard(PooledDataBuffer.class, DataBufferUtils::release);
+			}));
 		}
 
+		if (logger.isDebugEnabled()) {
+			body = body.doOnNext(buffer -> Hints.touchDataBuffer(buffer, hints, logger));
+		}
 		return message.writeWith(body);
 	}
 
@@ -166,6 +175,9 @@ private static MediaType addDefaultCharset(MediaType main, @Nullable MediaType d
 		return main;
 	}
 
+	private static void touch(DataBuffer buffer, Map<String, Object> hints) {
+	}
+
 	private boolean isStreamingMediaType(@Nullable MediaType mediaType) {
 		if (mediaType == null || !(this.encoder instanceof HttpMessageEncoder)) {
 			return false;

spring-web/src/main/java/org/springframework/http/codec/ResourceHttpMessageWriter.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.
@@ -132,6 +132,9 @@ private Mono<Void> writeResource(Resource resource, ResolvableType type, @Nullab
 					Mono<Resource> input = Mono.just(resource);
 					DataBufferFactory factory = message.bufferFactory();
 					Flux<DataBuffer> body = this.encoder.encode(input, factory, type, resourceMediaType, hints);
+					if (logger.isDebugEnabled()) {
+						body = body.doOnNext(buffer -> Hints.touchDataBuffer(buffer, hints, logger));
+					}
 					return message.writeWith(body);
 				});
 	}

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

@@ -23,6 +23,7 @@
 import java.util.List;
 import java.util.Map;
 
+import org.apache.commons.logging.Log;
 import org.reactivestreams.Publisher;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
@@ -35,6 +36,7 @@
 import org.springframework.core.io.buffer.DataBufferFactory;
 import org.springframework.core.io.buffer.DataBufferUtils;
 import org.springframework.core.io.buffer.PooledDataBuffer;
+import org.springframework.http.HttpLogging;
 import org.springframework.http.MediaType;
 import org.springframework.http.ReactiveHttpOutputMessage;
 import org.springframework.http.server.reactive.ServerHttpRequest;
@@ -57,6 +59,8 @@ public class ServerSentEventHttpMessageWriter implements HttpMessageWriter<Objec
 
 	private static final List<MediaType> WRITABLE_MEDIA_TYPES = Collections.singletonList(MediaType.TEXT_EVENT_STREAM);
 
+	private static final Log logger = HttpLogging.forLogName(ServerSentEventHttpMessageWriter.class);
+
 
 	@Nullable
 	private final Encoder<?> encoder;
@@ -167,9 +171,11 @@ private <T> Flux<DataBuffer> encodeEvent(StringBuilder eventContent, T data, Res
 		if (this.encoder == null) {
 			throw new CodecException("No SSE encoder configured and the data is not String.");
 		}
+		DataBuffer buffer = ((Encoder<T>) this.encoder).encodeValue(data, factory, dataType, mediaType, hints);
+		Hints.touchDataBuffer(buffer, hints, logger);
 		return Flux.just(factory.join(Arrays.asList(
 				encodeText(eventContent, mediaType, factory),
-				((Encoder<T>) this.encoder).encodeValue(data, factory, dataType, mediaType, hints),
+				buffer,
 				encodeText("\n\n", mediaType, factory))));
 	}
 

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

@@ -199,6 +199,10 @@ private Mono<Void> writeMultipart(MultiValueMap<String, ?> map,
 				.concatWith(generateLastLine(boundary, bufferFactory))
 				.doOnDiscard(PooledDataBuffer.class, DataBufferUtils::release);
 
+		if (logger.isDebugEnabled()) {
+			body = body.doOnNext(buffer -> Hints.touchDataBuffer(buffer, hints, logger));
+		}
+
 		return outputMessage.writeWith(body);
 	}