DataBuffer fixes in Protobuf codecs

Closes gh-22731

spring-web/src/test/java/org/springframework/http/codec/CancelWithoutDemandCodecTests.java

@@ -20,6 +20,8 @@ import java.util.Collections;
 import java.util.List;
 import java.util.function.Supplier;
 
+import com.google.protobuf.Message;
+import org.junit.After;
 import org.junit.Test;
 import org.reactivestreams.Publisher;
 import org.reactivestreams.Subscription;
@@ -38,18 +40,28 @@ import org.springframework.http.ReactiveHttpOutputMessage;
 import org.springframework.http.client.MultipartBodyBuilder;
 import org.springframework.http.codec.json.Jackson2JsonEncoder;
 import org.springframework.http.codec.multipart.MultipartHttpMessageWriter;
+import org.springframework.http.codec.protobuf.ProtobufDecoder;
+import org.springframework.http.codec.protobuf.ProtobufEncoder;
 import org.springframework.http.codec.xml.Jaxb2XmlEncoder;
+import org.springframework.protobuf.Msg;
+import org.springframework.protobuf.SecondMsg;
+import org.springframework.util.MimeType;
 
 /**
  * Test scenarios for data buffer leaks.
  * @author Rossen Stoyanchev
- * @since 5.2
  */
-public class CodecDataBufferLeakTests {
+public class CancelWithoutDemandCodecTests {
 
 	private final LeakAwareDataBufferFactory bufferFactory = new LeakAwareDataBufferFactory();
 
 
+	@After
+	public void tearDown() throws Exception {
+		this.bufferFactory.checkForLeaks();
+	}
+
+
 	@Test // gh-22107
 	public void cancelWithEncoderHttpMessageWriterAndSingleValue() {
 		CharSequenceEncoder encoder = CharSequenceEncoder.allMimeTypes();
@@ -58,8 +70,6 @@ public class CodecDataBufferLeakTests {
 
 		writer.write(Mono.just("foo"), ResolvableType.forType(String.class), MediaType.TEXT_PLAIN,
 				outputMessage, Collections.emptyMap()).block(Duration.ofSeconds(5));
-
-		this.bufferFactory.checkForLeaks();
 	}
 
 	@Test // gh-22107
@@ -73,8 +83,6 @@ public class CodecDataBufferLeakTests {
 		BaseSubscriber<DataBuffer> subscriber = new ZeroDemandSubscriber();
 		flux.subscribe(subscriber); // Assume sync execution (e.g. encoding with Flux.just)..
 		subscriber.cancel();
-
-		this.bufferFactory.checkForLeaks();
 	}
 
 	@Test // gh-22107
@@ -88,8 +96,39 @@ public class CodecDataBufferLeakTests {
 		BaseSubscriber<DataBuffer> subscriber = new ZeroDemandSubscriber();
 		flux.subscribe(subscriber); // Assume sync execution (e.g. encoding with Flux.just)..
 		subscriber.cancel();
+	}
 
-		this.bufferFactory.checkForLeaks();
+	@Test // gh-22543
+	public void cancelWithProtobufEncoder() {
+		ProtobufEncoder encoder = new ProtobufEncoder();
+		Msg msg = Msg.newBuilder().setFoo("Foo").setBlah(SecondMsg.newBuilder().setBlah(123).build()).build();
+
+		Flux<DataBuffer> flux = encoder.encode(Mono.just(msg),
+				this.bufferFactory, ResolvableType.forClass(Msg.class),
+				new MimeType("application", "x-protobuf"), Collections.emptyMap());
+
+		BaseSubscriber<DataBuffer> subscriber = new ZeroDemandSubscriber();
+		flux.subscribe(subscriber); // Assume sync execution (e.g. encoding with Flux.just)..
+		subscriber.cancel();
+	}
+
+	@Test // gh-22731
+	public void cancelWithProtobufDecoder() throws InterruptedException {
+		ProtobufDecoder decoder = new ProtobufDecoder();
+
+		Mono<DataBuffer> input = Mono.fromCallable(() -> {
+			Msg msg = Msg.newBuilder().setFoo("Foo").build();
+			byte[] bytes = msg.toByteArray();
+			DataBuffer buffer = this.bufferFactory.allocateBuffer(bytes.length);
+			buffer.write(bytes);
+			return buffer;
+		});
+
+		Flux<Message> messages = decoder.decode(input, ResolvableType.forType(Msg.class),
+				new MimeType("application", "x-protobuf"), Collections.emptyMap());
+		ZeroDemandMessageSubscriber subscriber = new ZeroDemandMessageSubscriber();
+		messages.subscribe(subscriber);
+		subscriber.cancel();
 	}
 
 	@Test // gh-22107
@@ -104,8 +143,6 @@ public class CodecDataBufferLeakTests {
 
 		writer.write(Mono.just(builder.build()), null, MediaType.MULTIPART_FORM_DATA,
 				outputMessage, Collections.emptyMap()).block(Duration.ofSeconds(5));
-
-		this.bufferFactory.checkForLeaks();
 	}
 
 	@Test // gh-22107
@@ -116,8 +153,6 @@ public class CodecDataBufferLeakTests {
 
 		writer.write(Mono.just(event), ResolvableType.forClass(ServerSentEvent.class), MediaType.TEXT_EVENT_STREAM,
 				outputMessage, Collections.emptyMap()).block(Duration.ofSeconds(5));
-
-		this.bufferFactory.checkForLeaks();
 	}
 
 
@@ -183,4 +218,13 @@ public class CodecDataBufferLeakTests {
 			// Just subscribe without requesting
 		}
 	}
+
+
+	private static class ZeroDemandMessageSubscriber extends BaseSubscriber<Message> {
+
+		@Override
+		protected void hookOnSubscribe(Subscription subscription) {
+			// Just subscribe without requesting
+		}
+	}
 }

spring-web/src/test/java/org/springframework/http/codec/protobuf/ProtobufDecoderTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2018 the original author or authors.
+ * Copyright 2002-2019 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.
@@ -35,10 +35,10 @@ import org.springframework.protobuf.Msg;
 import org.springframework.protobuf.SecondMsg;
 import org.springframework.util.MimeType;
 
-import static java.util.Collections.emptyMap;
+import static java.util.Collections.*;
 import static org.junit.Assert.*;
-import static org.springframework.core.ResolvableType.forClass;
-import static org.springframework.core.io.buffer.DataBufferUtils.release;
+import static org.springframework.core.ResolvableType.*;
+import static org.springframework.core.io.buffer.DataBufferUtils.*;
 
 /**
  * Unit tests for {@link ProtobufDecoder}.
@@ -223,11 +223,11 @@ public class ProtobufDecoderTests extends AbstractDecoderTestCase<ProtobufDecode
 	}
 
 	private Mono<DataBuffer> dataBuffer(Msg msg) {
-		return Mono.defer(() -> {
+		return Mono.fromCallable(() -> {
 			byte[] bytes = msg.toByteArray();
 			DataBuffer buffer = this.bufferFactory.allocateBuffer(bytes.length);
 			buffer.write(bytes);
-			return Mono.just(buffer);
+			return buffer;
 		});
 	}