Added DataBufferAllocator parameter to Encoder.encode, and use that instead of AbstractAllocatingEncoder base class.

spring-web-reactive/src/main/java/org/springframework/core/codec/Encoder.java

@@ -16,7 +16,6 @@
 
 package org.springframework.core.codec;
 
-import java.nio.ByteBuffer;
 import java.util.List;
 
 import org.reactivestreams.Publisher;
@@ -24,6 +23,7 @@ import reactor.core.publisher.Flux;
 
 import org.springframework.core.ResolvableType;
 import org.springframework.core.io.buffer.DataBuffer;
+import org.springframework.core.io.buffer.DataBufferAllocator;
 import org.springframework.util.MimeType;
 
 /**
@@ -44,14 +44,16 @@ public interface Encoder<T> {
 	boolean canEncode(ResolvableType type, MimeType mimeType, Object... hints);
 
 	/**
-	 * Encode an input stream of {@code T} to an output {@link ByteBuffer} stream.
+	 * Encode an input stream of {@code T} to an output {@link DataBuffer} stream.
 	 * @param inputStream the input stream to process.
+	 * @param allocator a buffer allocator used to create the output
 	 * @param type the stream element type to process.
 	 * @param mimeType the mime type to process.
 	 * @param hints Additional information about how to do decode, optional.
 	 * @return the output stream
 	 */
-	Flux<DataBuffer> encode(Publisher<? extends T> inputStream, ResolvableType type,
+	Flux<DataBuffer> encode(Publisher<? extends T> inputStream,
+			DataBufferAllocator allocator, ResolvableType type,
 			MimeType mimeType, Object... hints);
 
 	/**

spring-web-reactive/src/main/java/org/springframework/core/codec/support/AbstractAllocatingEncoder.java

@@ -1,42 +0,0 @@
-/*
- * Copyright 2002-2016 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.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.springframework.core.codec.support;
-
-import org.springframework.core.io.buffer.DataBufferAllocator;
-import org.springframework.util.Assert;
-import org.springframework.util.MimeType;
-
-/**
- * @author Arjen Poutsma
- */
-public abstract class AbstractAllocatingEncoder<T> extends AbstractEncoder<T> {
-
-	private final DataBufferAllocator allocator;
-
-	public AbstractAllocatingEncoder(DataBufferAllocator allocator,
-			MimeType... supportedMimeTypes) {
-		super(supportedMimeTypes);
-		Assert.notNull(allocator, "'allocator' must not be null");
-
-		this.allocator = allocator;
-	}
-
-	public DataBufferAllocator allocator() {
-		return allocator;
-	}
-
-}

spring-web-reactive/src/main/java/org/springframework/core/codec/support/ByteBufferEncoder.java

@@ -30,10 +30,10 @@ import org.springframework.util.MimeTypeUtils;
 /**
  * @author Sebastien Deleuze
  */
-public class ByteBufferEncoder extends AbstractAllocatingEncoder<ByteBuffer> {
+public class ByteBufferEncoder extends AbstractEncoder<ByteBuffer> {
 
-	public ByteBufferEncoder(DataBufferAllocator allocator) {
-		super(allocator, MimeTypeUtils.ALL);
+	public ByteBufferEncoder() {
+		super(MimeTypeUtils.ALL);
 	}
 
 
@@ -45,12 +45,11 @@ public class ByteBufferEncoder extends AbstractAllocatingEncoder<ByteBuffer> {
 
 	@Override
 	public Flux<DataBuffer> encode(Publisher<? extends ByteBuffer> inputStream,
-			ResolvableType type,
-			MimeType mimeType, Object... hints) {
-
+			DataBufferAllocator allocator, ResolvableType type, MimeType mimeType,
+			Object... hints) {
 		//noinspection unchecked
 		return Flux.from(inputStream).map(byteBuffer -> {
-			DataBuffer dataBuffer = allocator().allocateBuffer(byteBuffer.remaining());
+			DataBuffer dataBuffer = allocator.allocateBuffer(byteBuffer.remaining());
 			dataBuffer.write(byteBuffer);
 			return dataBuffer;
 		});

spring-web-reactive/src/main/java/org/springframework/core/codec/support/JacksonJsonEncoder.java

@@ -30,6 +30,7 @@ import org.springframework.core.codec.CodecException;
 import org.springframework.core.codec.Encoder;
 import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.core.io.buffer.DataBufferAllocator;
+import org.springframework.core.io.buffer.DefaultDataBufferAllocator;
 import org.springframework.util.MimeType;
 
 /**
@@ -38,41 +39,46 @@ import org.springframework.util.MimeType;
  * @author Sebastien Deleuze
  * @see JacksonJsonDecoder
  */
-public class JacksonJsonEncoder extends AbstractAllocatingEncoder<Object> {
+public class JacksonJsonEncoder extends AbstractEncoder<Object> {
 
 	private final ObjectMapper mapper;
 
 	private Encoder<DataBuffer> postProcessor;
 
-	public JacksonJsonEncoder(DataBufferAllocator allocator) {
-		this(allocator, new ObjectMapper(), null);
+	public JacksonJsonEncoder() {
+		this(new ObjectMapper(), null);
 	}
 
-	public JacksonJsonEncoder(DataBufferAllocator allocator,
-			Encoder<DataBuffer> postProcessor) {
-		this(allocator, new ObjectMapper(), postProcessor);
+	public JacksonJsonEncoder(Encoder<DataBuffer> postProcessor) {
+		this(new ObjectMapper(), postProcessor);
 	}
 
-	public JacksonJsonEncoder(DataBufferAllocator allocator, ObjectMapper mapper,
+	public JacksonJsonEncoder(ObjectMapper mapper,
 			Encoder<DataBuffer> postProcessor) {
-		super(allocator, new MimeType("application", "json", StandardCharsets.UTF_8),
+		super(new MimeType("application", "json", StandardCharsets.UTF_8),
 				new MimeType("application", "*+json", StandardCharsets.UTF_8));
 		this.mapper = mapper;
 		this.postProcessor = postProcessor;
 	}
 
 	@Override
-	public Flux<DataBuffer> encode(Publisher<? extends Object> inputStream,
-			ResolvableType type, MimeType mimeType, Object... hints) {
+	public Flux<DataBuffer> encode(Publisher<?> inputStream,
+			DataBufferAllocator allocator, ResolvableType type, MimeType mimeType,
+			Object... hints) {
 
 		Publisher<DataBuffer> stream = (inputStream instanceof Mono ?
-				((Mono<?>)inputStream).map(this::serialize) :
-				Flux.from(inputStream).map(this::serialize));
-		return (this.postProcessor == null ? Flux.from(stream) : this.postProcessor.encode(stream, type, mimeType, hints));
+				((Mono<?>) inputStream).map(value -> serialize(value, allocator)) :
+				Flux.from(inputStream).map(value -> serialize(value, allocator)));
+		// TODO: figure out why using the parameter allocator for the postprocessor
+		// commits the response too early
+		DefaultDataBufferAllocator tempAllocator = new DefaultDataBufferAllocator();
+
+		return (this.postProcessor == null ? Flux.from(stream) :
+				this.postProcessor.encode(stream, tempAllocator, type, mimeType, hints));
 	}
 
-	private DataBuffer serialize(Object value) {
-		DataBuffer buffer = allocator().allocateBuffer();
+	private DataBuffer serialize(Object value, DataBufferAllocator allocator) {
+		DataBuffer buffer = allocator.allocateBuffer();
 		OutputStream outputStream = buffer.asOutputStream();
 		try {
 			this.mapper.writeValue(outputStream, value);

spring-web-reactive/src/main/java/org/springframework/core/codec/support/Jaxb2Encoder.java

@@ -43,23 +43,22 @@ import org.springframework.util.MimeTypeUtils;
  * @author Sebastien Deleuze
  * @see Jaxb2Decoder
  */
-public class Jaxb2Encoder extends AbstractAllocatingEncoder<Object> {
+public class Jaxb2Encoder extends AbstractEncoder<Object> {
 
 	private final ConcurrentMap<Class<?>, JAXBContext> jaxbContexts = new ConcurrentHashMap<>(64);
 
-	public Jaxb2Encoder(DataBufferAllocator allocator) {
-		super(allocator, MimeTypeUtils.APPLICATION_XML, MimeTypeUtils.TEXT_XML);
+	public Jaxb2Encoder() {
+		super(MimeTypeUtils.APPLICATION_XML, MimeTypeUtils.TEXT_XML);
 	}
 
-
 	@Override
-	public Flux<DataBuffer> encode(Publisher<? extends Object> messageStream,
-			ResolvableType type,
-			MimeType mimeType, Object... hints) {
+	public Flux<DataBuffer> encode(Publisher<?> inputStream,
+			DataBufferAllocator allocator, ResolvableType type, MimeType mimeType,
+			Object... hints) {
 
-		return Flux.from(messageStream).map(value -> {
+		return Flux.from(inputStream).map(value -> {
 			try {
-				DataBuffer buffer = allocator().allocateBuffer(1024);
+				DataBuffer buffer = allocator.allocateBuffer(1024);
 				OutputStream outputStream = buffer.asOutputStream();
 				Class<?> clazz = ClassUtils.getUserClass(value);
 				Marshaller marshaller = createMarshaller(clazz);

spring-web-reactive/src/main/java/org/springframework/core/codec/support/JsonObjectEncoder.java

@@ -42,21 +42,22 @@ import org.springframework.util.MimeType;
  *
  * @see JsonObjectDecoder
  */
-public class JsonObjectEncoder extends AbstractAllocatingEncoder<DataBuffer> {
+public class JsonObjectEncoder extends AbstractEncoder<DataBuffer> {
 
-	public JsonObjectEncoder(DataBufferAllocator allocator) {
-		super(allocator, new MimeType("application", "json", StandardCharsets.UTF_8),
+	public JsonObjectEncoder() {
+		super(new MimeType("application", "json", StandardCharsets.UTF_8),
 				new MimeType("application", "*+json", StandardCharsets.UTF_8));
 	}
 
 	@Override
 	public Flux<DataBuffer> encode(Publisher<? extends DataBuffer> inputStream,
+			DataBufferAllocator allocator,
 			ResolvableType type, MimeType mimeType, Object... hints) {
 		if (inputStream instanceof Mono) {
 			return Flux.from(inputStream);
 		}
 		return Flux.from(inputStream)
-				.lift(s -> new JsonArrayEncoderBarrier(s, allocator()));
+				.lift(s -> new JsonArrayEncoderBarrier(s, allocator));
 	}
 
 	private static class JsonArrayEncoderBarrier

spring-web-reactive/src/main/java/org/springframework/core/codec/support/StringEncoder.java

@@ -33,12 +33,12 @@ import org.springframework.util.MimeType;
  * @author Sebastien Deleuze
  * @see StringDecoder
  */
-public class StringEncoder extends AbstractAllocatingEncoder<String> {
+public class StringEncoder extends AbstractEncoder<String> {
 
 	public static final Charset DEFAULT_CHARSET = StandardCharsets.UTF_8;
 
-	public StringEncoder(DataBufferAllocator allocator) {
-		super(allocator, new MimeType("text", "plain", DEFAULT_CHARSET));
+	public StringEncoder() {
+		super(new MimeType("text", "plain", DEFAULT_CHARSET));
 	}
 
 
@@ -49,9 +49,9 @@ public class StringEncoder extends AbstractAllocatingEncoder<String> {
 	}
 
 	@Override
-	public Flux<DataBuffer> encode(Publisher<? extends String> elementStream,
-			ResolvableType type, MimeType mimeType, Object... hints) {
-
+	public Flux<DataBuffer> encode(Publisher<? extends String> inputStream,
+			DataBufferAllocator allocator, ResolvableType type, MimeType mimeType,
+			Object... hints) {
 		Charset charset;
 		if (mimeType != null && mimeType.getCharSet() != null) {
 			charset = mimeType.getCharSet();
@@ -59,9 +59,9 @@ public class StringEncoder extends AbstractAllocatingEncoder<String> {
 		else {
 			 charset = DEFAULT_CHARSET;
 		}
-		return Flux.from(elementStream).map(s -> {
+		return Flux.from(inputStream).map(s -> {
 			byte[] bytes = s.getBytes(charset);
-			DataBuffer dataBuffer = allocator().allocateBuffer(bytes.length);
+			DataBuffer dataBuffer = allocator.allocateBuffer(bytes.length);
 			dataBuffer.write(bytes);
 			return dataBuffer;
 		});

spring-web-reactive/src/test/java/org/springframework/core/codec/support/ByteBufferEncoderTests.java

@@ -23,13 +23,13 @@ import org.junit.Before;
 import org.junit.Test;
 import org.reactivestreams.Publisher;
 import reactor.core.publisher.Flux;
+import reactor.core.test.TestSubscriber;
 
 import org.springframework.core.ResolvableType;
 import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.http.MediaType;
 
 import static org.junit.Assert.*;
-import reactor.core.test.TestSubscriber;
 
 /**
  * @author Sebastien Deleuze
@@ -40,7 +40,7 @@ public class ByteBufferEncoderTests extends AbstractAllocatingTestCase {
 
 	@Before
 	public void createEncoder() {
-		encoder = new ByteBufferEncoder(allocator);
+		encoder = new ByteBufferEncoder();
 	}
 
 	@Test
@@ -57,7 +57,7 @@ public class ByteBufferEncoderTests extends AbstractAllocatingTestCase {
 		Flux<ByteBuffer> source =
 				Flux.just(ByteBuffer.wrap(fooBytes), ByteBuffer.wrap(barBytes));
 
-		Flux<DataBuffer> output = encoder.encode(source,
+		Flux<DataBuffer> output = encoder.encode(source, allocator,
 				ResolvableType.forClassWithGenerics(Publisher.class, ByteBuffer.class),
 				null);
 		TestSubscriber<DataBuffer> testSubscriber = new TestSubscriber<>();

spring-web-reactive/src/test/java/org/springframework/core/codec/support/JacksonJsonEncoderTests.java

@@ -21,11 +21,12 @@ import java.nio.charset.StandardCharsets;
 import org.junit.Before;
 import org.junit.Test;
 import reactor.core.publisher.Flux;
+import reactor.core.test.TestSubscriber;
 
 import org.springframework.http.MediaType;
 
-import static org.junit.Assert.*;
-import reactor.core.test.TestSubscriber;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
 
 /**
  * @author Sebastien Deleuze
@@ -36,7 +37,7 @@ public class JacksonJsonEncoderTests extends AbstractAllocatingTestCase {
 
 	@Before
 	public void createEncoder() {
-		encoder = new JacksonJsonEncoder(allocator);
+		encoder = new JacksonJsonEncoder();
 	}
 
 	@Test
@@ -48,7 +49,7 @@ public class JacksonJsonEncoderTests extends AbstractAllocatingTestCase {
 	@Test
 	public void write() {
 		Flux<Pojo> source = Flux.just(new Pojo("foofoo", "barbar"), new Pojo("foofoofoo", "barbarbar"));
-		Flux<String> output = encoder.encode(source, null, null).map(chunk -> {
+		Flux<String> output = encoder.encode(source, allocator, null, null).map(chunk -> {
 			byte[] b = new byte[chunk.readableByteCount()];
 			chunk.read(b);
 			return new String(b, StandardCharsets.UTF_8);

spring-web-reactive/src/test/java/org/springframework/core/codec/support/Jaxb2EncoderTests.java

@@ -21,11 +21,12 @@ import java.nio.charset.StandardCharsets;
 import org.junit.Before;
 import org.junit.Test;
 import reactor.core.publisher.Flux;
+import reactor.core.test.TestSubscriber;
 
 import org.springframework.http.MediaType;
 
-import static org.junit.Assert.*;
-import reactor.core.test.TestSubscriber;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
 
 /**
  * @author Sebastien Deleuze
@@ -36,7 +37,7 @@ public class Jaxb2EncoderTests extends AbstractAllocatingTestCase {
 
 	@Before
 	public void createEncoder() {
-		encoder = new Jaxb2Encoder(allocator);
+		encoder = new Jaxb2Encoder();
 	}
 
 	@Test
@@ -49,7 +50,7 @@ public class Jaxb2EncoderTests extends AbstractAllocatingTestCase {
 	@Test
 	public void encode() {
 		Flux<Pojo> source = Flux.just(new Pojo("foofoo", "barbar"), new Pojo("foofoofoo", "barbarbar"));
-		Flux<String> output = encoder.encode(source, null, null).map(chunk -> {
+		Flux<String> output = encoder.encode(source, allocator, null, null).map(chunk -> {
 			byte[] b = new byte[chunk.readableByteCount()];
 			chunk.read(b);
 			return new String(b, StandardCharsets.UTF_8);

spring-web-reactive/src/test/java/org/springframework/core/codec/support/JsonObjectEncoderTests.java

@@ -22,11 +22,10 @@ import org.junit.Before;
 import org.junit.Test;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
+import reactor.core.test.TestSubscriber;
 
 import org.springframework.core.io.buffer.DataBuffer;
 
-import reactor.core.test.TestSubscriber;
-
 /**
  * @author Sebastien Deleuze
  */
@@ -36,14 +35,15 @@ public class JsonObjectEncoderTests extends AbstractAllocatingTestCase {
 
 	@Before
 	public void createEncoder() {
-		encoder = new JsonObjectEncoder(allocator);
+		encoder = new JsonObjectEncoder();
 	}
 
 	@Test
 	public void encodeSingleElementFlux() throws InterruptedException {
 		Flux<DataBuffer> source =
 				Flux.just(stringBuffer("{\"foo\": \"foofoo\", \"bar\": \"barbar\"}"));
-		Flux<String> output = Flux.from(encoder.encode(source, null, null)).map(chunk -> {
+		Flux<String> output =
+				Flux.from(encoder.encode(source, allocator, null, null)).map(chunk -> {
 			byte[] b = new byte[chunk.readableByteCount()];
 			chunk.read(b);
 			return new String(b, StandardCharsets.UTF_8);
@@ -57,7 +57,8 @@ public class JsonObjectEncoderTests extends AbstractAllocatingTestCase {
 	public void encodeSingleElementMono() throws InterruptedException {
 		Mono<DataBuffer> source =
 				Mono.just(stringBuffer("{\"foo\": \"foofoo\", \"bar\": \"barbar\"}"));
-		Flux<String> output = Flux.from(encoder.encode(source, null, null)).map(chunk -> {
+		Flux<String> output =
+				Flux.from(encoder.encode(source, allocator, null, null)).map(chunk -> {
 			byte[] b = new byte[chunk.readableByteCount()];
 			chunk.read(b);
 			return new String(b, StandardCharsets.UTF_8);
@@ -72,7 +73,8 @@ public class JsonObjectEncoderTests extends AbstractAllocatingTestCase {
 		Flux<DataBuffer> source =
 				Flux.just(stringBuffer("{\"foo\": \"foofoo\", \"bar\": \"barbar\"}"),
 						stringBuffer("{\"foo\": \"foofoofoo\", \"bar\": \"barbarbar\"}"));
-		Flux<String> output = Flux.from(encoder.encode(source, null, null)).map(chunk -> {
+		Flux<String> output =
+				Flux.from(encoder.encode(source, allocator, null, null)).map(chunk -> {
 			byte[] b = new byte[chunk.readableByteCount()];
 			chunk.read(b);
 			return new String(b, StandardCharsets.UTF_8);
@@ -91,7 +93,8 @@ public class JsonObjectEncoderTests extends AbstractAllocatingTestCase {
 						stringBuffer("{\"foo\": \"foofoofoo\", \"bar\": \"barbarbar\"}"),
 						stringBuffer("{\"foo\": \"foofoofoofoo\", \"bar\": \"barbarbarbar\"}")
 		);
-		Flux<String> output = Flux.from(encoder.encode(source, null, null)).map(chunk -> {
+		Flux<String> output =
+				Flux.from(encoder.encode(source, allocator, null, null)).map(chunk -> {
 			byte[] b = new byte[chunk.readableByteCount()];
 			chunk.read(b);
 			return new String(b, StandardCharsets.UTF_8);

spring-web-reactive/src/test/java/org/springframework/core/codec/support/StringEncoderTests.java

@@ -23,12 +23,13 @@ import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.Parameterized;
 import reactor.core.publisher.Flux;
+import reactor.core.test.TestSubscriber;
 
 import org.springframework.core.ResolvableType;
 import org.springframework.http.MediaType;
 
-import static org.junit.Assert.*;
-import reactor.core.test.TestSubscriber;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
 
 /**
  * @author Sebastien Deleuze
@@ -40,7 +41,7 @@ public class StringEncoderTests extends AbstractAllocatingTestCase {
 
 	@Before
 	public void createEncoder() {
-		encoder = new StringEncoder(allocator);
+		encoder = new StringEncoder();
 	}
 
 	@Test
@@ -52,7 +53,9 @@ public class StringEncoderTests extends AbstractAllocatingTestCase {
 
 	@Test
 	public void write() throws InterruptedException {
-		Flux<String> output = Flux.from(encoder.encode(Flux.just("foo"), null, null)).map(chunk -> {
+		Flux<String> output =
+				Flux.from(encoder.encode(Flux.just("foo"), allocator, null, null))
+						.map(chunk -> {
 			byte[] b = new byte[chunk.readableByteCount()];
 			chunk.read(b);
 			return new String(b, StandardCharsets.UTF_8);