Write CharSequence instances to DataBuffers

Prior to this commit, one could write a `CharSequence` to an existing
`DataBuffer` instance by turning it into a byte array or `ByteBuffer`
first. This had the following disadvantages:

1. Memory allocation was not efficient (not leveraging pooled memory
when available)
2. Dealing with `CharsetEncoder` is not always easy
3. `DataBuffer` implementations, like `NettyDataBuffer` can use
optimized implementations in some cases

This commit adds a new `DataBuffer#write(CharSequence, Charset)` method
for those cases and also an `ensureCapacity` method useful for checking
that the current buffer has enough capacity to write to it..

Issue: SPR-17558

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

@@ -19,8 +19,15 @@ package org.springframework.core.io.buffer;
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.nio.ByteBuffer;
+import java.nio.CharBuffer;
+import java.nio.charset.Charset;
+import java.nio.charset.CharsetEncoder;
+import java.nio.charset.CoderResult;
+import java.nio.charset.CodingErrorAction;
 import java.util.function.IntPredicate;
 
+import org.springframework.util.Assert;
+
 /**
  * Basic abstraction over byte buffers.
  *
@@ -45,6 +52,7 @@ import java.util.function.IntPredicate;
  * can also be used on non-Netty platforms (i.e. Servlet containers).
  *
  * @author Arjen Poutsma
+ * @author Brian Clozel
  * @since 5.0
  * @see DataBufferFactory
  */
@@ -106,6 +114,16 @@ public interface DataBuffer {
 	 */
 	DataBuffer capacity(int capacity);
 
+	/**
+	 * Ensure that the current buffer has enough {@link #writableByteCount()}
+	 * to write the amount of data given as an argument. If not, the missing
+	 * capacity will be added to the buffer.
+	 * @param capacity the writable capacity to check for
+	 * @return this buffer
+	 * @since 5.1.4
+	 */
+	DataBuffer ensureCapacity(int capacity);
+
 	/**
 	 * Return the position from which this buffer will read.
 	 * @return the read position
@@ -181,7 +199,7 @@ public interface DataBuffer {
 	DataBuffer write(byte b);
 
 	/**
-	 * Write the given source into this buffer, startin at the current writing position
+	 * Write the given source into this buffer, starting at the current writing position
 	 * of this buffer.
 	 * @param source the bytes to be written into this buffer
 	 * @return this buffer
@@ -215,6 +233,44 @@ public interface DataBuffer {
 	 */
 	DataBuffer write(ByteBuffer... buffers);
 
+	/**
+	 * Write the given {@code CharSequence} using the given {@code Charset},
+	 * starting at the current writing position.
+	 * @param charSequence the char sequence to write into this buffer
+	 * @param charset the charset to encode the char sequence with
+	 * @return this buffer
+	 * @since 5.1.4
+	 */
+	default DataBuffer write(CharSequence charSequence, Charset charset) {
+		Assert.notNull(charSequence, "'charSequence' must not be null");
+		Assert.notNull(charset, "'charset' must not be null");
+		CharsetEncoder charsetEncoder = charset.newEncoder()
+				.onMalformedInput(CodingErrorAction.REPLACE)
+				.onUnmappableCharacter(CodingErrorAction.REPLACE);
+		CharBuffer inBuffer = CharBuffer.wrap(charSequence);
+		int estimatedSize = (int) (inBuffer.remaining() * charsetEncoder.averageBytesPerChar());
+		ByteBuffer outBuffer = ensureCapacity(estimatedSize)
+				.asByteBuffer(writePosition(), writableByteCount());
+		for (; ; ) {
+			CoderResult cr = inBuffer.hasRemaining() ?
+					charsetEncoder.encode(inBuffer, outBuffer, true) : CoderResult.UNDERFLOW;
+			if (cr.isUnderflow()) {
+				cr = charsetEncoder.flush(outBuffer);
+			}
+			if (cr.isUnderflow()) {
+				break;
+			}
+			if (cr.isOverflow()) {
+				writePosition(outBuffer.position());
+				int maximumSize = (int) (inBuffer.remaining() * charsetEncoder.maxBytesPerChar());
+				ensureCapacity(maximumSize);
+				outBuffer = asByteBuffer(writePosition(), writableByteCount());
+			}
+		}
+		writePosition(outBuffer.position());
+		return this;
+	}
+
 	/**
 	 * Create a new {@code DataBuffer} whose contents is a shared subsequence of this
 	 * data buffer's content.  Data between this data buffer and the returned buffer is

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

@@ -215,6 +215,15 @@ public class DefaultDataBuffer implements DataBuffer {
 		return this;
 	}
 
+	@Override
+	public DataBuffer ensureCapacity(int length) {
+		if (length > writableByteCount()) {
+			int newCapacity = calculateCapacity(this.writePosition + length);
+			capacity(newCapacity);
+		}
+		return this;
+	}
+
 	private static ByteBuffer allocate(int capacity, boolean direct) {
 		return direct ? ByteBuffer.allocateDirect(capacity) : ByteBuffer.allocate(capacity);
 	}
@@ -369,13 +378,6 @@ public class DefaultDataBuffer implements DataBuffer {
 		return new DefaultDataBufferOutputStream();
 	}
 
-	private void ensureCapacity(int length) {
-		if (length <= writableByteCount()) {
-			return;
-		}
-		int newCapacity = calculateCapacity(this.writePosition + length);
-		capacity(newCapacity);
-	}
 
 	/**
 	 * Calculate the capacity of the buffer.

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

@@ -19,11 +19,14 @@ package org.springframework.core.io.buffer;
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.nio.ByteBuffer;
+import java.nio.charset.Charset;
+import java.nio.charset.StandardCharsets;
 import java.util.function.IntPredicate;
 
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.ByteBufInputStream;
 import io.netty.buffer.ByteBufOutputStream;
+import io.netty.buffer.ByteBufUtil;
 
 import org.springframework.util.Assert;
 import org.springframework.util.ObjectUtils;
@@ -138,6 +141,12 @@ public class NettyDataBuffer implements PooledDataBuffer {
 		return this;
 	}
 
+	@Override
+	public DataBuffer ensureCapacity(int capacity) {
+		this.byteBuf.ensureWritable(capacity);
+		return this;
+	}
+
 	@Override
 	public byte read() {
 		return this.byteBuf.readByte();
@@ -178,14 +187,14 @@ public class NettyDataBuffer implements PooledDataBuffer {
 		if (!ObjectUtils.isEmpty(buffers)) {
 			if (hasNettyDataBuffers(buffers)) {
 				ByteBuf[] nativeBuffers = new ByteBuf[buffers.length];
-				for (int i = 0 ; i < buffers.length; i++) {
+				for (int i = 0; i < buffers.length; i++) {
 					nativeBuffers[i] = ((NettyDataBuffer) buffers[i]).getNativeBuffer();
 				}
 				write(nativeBuffers);
 			}
 			else {
 				ByteBuffer[] byteBuffers = new ByteBuffer[buffers.length];
-				for (int i = 0 ; i < buffers.length; i++) {
+				for (int i = 0; i < buffers.length; i++) {
 					byteBuffers[i] = buffers[i].asByteBuffer();
 
 				}
@@ -229,6 +238,22 @@ public class NettyDataBuffer implements PooledDataBuffer {
 		return this;
 	}
 
+	@Override
+	public DataBuffer write(CharSequence charSequence, Charset charset) {
+		Assert.notNull(charSequence, "'charSequence' must not be null");
+		Assert.notNull(charset, "'charset' must not be null");
+		if (StandardCharsets.UTF_8.equals(charset)) {
+			ByteBufUtil.writeUtf8(this.byteBuf, charSequence);
+		}
+		else if (StandardCharsets.US_ASCII.equals(charset)) {
+			ByteBufUtil.writeAscii(this.byteBuf, charSequence);
+		}
+		else {
+			return PooledDataBuffer.super.write(charSequence, charset);
+		}
+		return this;
+	}
+
 	@Override
 	public NettyDataBuffer slice(int index, int length) {
 		ByteBuf slice = this.byteBuf.slice(index, length);