Ensure correct capacity in DefaultDataBuffer

See gh-31873 Closes gh-31979
spring-projects · Jan 9, 2024 · 0c22866 · 0c22866
1 parent 0c6957e
commit 0c22866
Show file tree

Hide file tree

Showing 3 changed files with 87 additions and 24 deletions.
diff --git a/spring-core/src/main/java/org/springframework/core/io/buffer/DataBuffer.java b/spring-core/src/main/java/org/springframework/core/io/buffer/DataBuffer.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2023 the original author or authors.
+ * Copyright 2002-2024 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.
@@ -263,29 +263,35 @@ default DataBuffer ensureCapacity(int capacity) {
 	default DataBuffer write(CharSequence charSequence, Charset charset) {
 		Assert.notNull(charSequence, "CharSequence must not be null");
 		Assert.notNull(charset, "Charset must not be null");
-		if (charSequence.length() > 0) {
+		if (!charSequence.isEmpty()) {
 			CharsetEncoder encoder = charset.newEncoder()
 					.onMalformedInput(CodingErrorAction.REPLACE)
 					.onUnmappableCharacter(CodingErrorAction.REPLACE);
 			CharBuffer src = CharBuffer.wrap(charSequence);
-			int cap = (int) (src.remaining() * encoder.averageBytesPerChar());
+			int averageSize = (int) Math.ceil(src.remaining() * encoder.averageBytesPerChar());
+			ensureWritable(averageSize);
 			while (true) {
-				ensureWritable(cap);
 				CoderResult cr;
-				try (ByteBufferIterator iterator = writableByteBuffers()) {
-					Assert.state(iterator.hasNext(), "No ByteBuffer available");
-					ByteBuffer dest = iterator.next();
-					cr = encoder.encode(src, dest, true);
-					if (cr.isUnderflow()) {
-						cr = encoder.flush(dest);
+				if (src.hasRemaining()) {
+					try (ByteBufferIterator iterator = writableByteBuffers()) {
+						Assert.state(iterator.hasNext(), "No ByteBuffer available");
+						ByteBuffer dest = iterator.next();
+						cr = encoder.encode(src, dest, true);
+						if (cr.isUnderflow()) {
+							cr = encoder.flush(dest);
+						}
+						writePosition(writePosition() + dest.position());
 					}
-					writePosition(dest.position());
+				}
+				else {
+					cr = CoderResult.UNDERFLOW;
 				}
 				if (cr.isUnderflow()) {
 					break;
 				}
-				if (cr.isOverflow()) {
-					cap = 2 * cap + 1;
+				else if (cr.isOverflow()) {
+					int maxSize = (int) Math.ceil(src.remaining() * encoder.maxBytesPerChar());
+					ensureWritable(maxSize);
 				}
 			}
 		}

diff --git a/spring-core/src/main/java/org/springframework/core/io/buffer/DefaultDataBuffer.java b/spring-core/src/main/java/org/springframework/core/io/buffer/DefaultDataBuffer.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2023 the original author or authors.
+ * Copyright 2002-2024 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.
@@ -416,16 +416,15 @@ public void toByteBuffer(int srcPos, ByteBuffer dest, int destPos, int length) {
 
 	@Override
 	public DataBuffer.ByteBufferIterator readableByteBuffers() {
-		ByteBuffer readOnly = this.byteBuffer.asReadOnlyBuffer();
-		readOnly.clear().position(this.readPosition).limit(this.writePosition - this.readPosition);
+		ByteBuffer readOnly = this.byteBuffer.slice(this.readPosition, readableByteCount())
+				.asReadOnlyBuffer();
 		return new ByteBufferIterator(readOnly);
 	}
 
 	@Override
 	public DataBuffer.ByteBufferIterator writableByteBuffers() {
-		ByteBuffer duplicate = this.byteBuffer.duplicate();
-		duplicate.clear().position(this.writePosition).limit(this.capacity - this.writePosition);
-		return new ByteBufferIterator(duplicate);
+		ByteBuffer slice = this.byteBuffer.slice(this.writePosition, writableByteCount());
+		return new ByteBufferIterator(slice);
 	}
 
 	@Override

diff --git a/spring-core/src/test/java/org/springframework/core/io/buffer/DataBufferTests.java b/spring-core/src/test/java/org/springframework/core/io/buffer/DataBufferTests.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2023 the original author or authors.
+ * Copyright 2002-2024 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.
@@ -678,6 +678,38 @@ void toByteBufferDestination(DataBufferFactory bufferFactory) {
 	void readableByteBuffers(DataBufferFactory bufferFactory) throws IOException {
 		super.bufferFactory = bufferFactory;
 
+		DataBuffer dataBuffer = this.bufferFactory.allocateBuffer(3);
+		dataBuffer.write("abc".getBytes(StandardCharsets.UTF_8));
+		dataBuffer.readPosition(1);
+		dataBuffer.writePosition(2);
+
+
+		byte[] result = new byte[1];
+		try (var iterator = dataBuffer.readableByteBuffers()) {
+			assertThat(iterator).hasNext();
+			int i = 0;
+			while (iterator.hasNext()) {
+				ByteBuffer byteBuffer = iterator.next();
+				assertThat(byteBuffer.position()).isEqualTo(0);
+				assertThat(byteBuffer.limit()).isEqualTo(1);
+				assertThat(byteBuffer.capacity()).isEqualTo(1);
+				assertThat(byteBuffer.remaining()).isEqualTo(1);
+
+				byteBuffer.get(result, i, 1);
+
+				assertThat(iterator).isExhausted();
+			}
+		}
+
+		assertThat(result).containsExactly('b');
+
+		release(dataBuffer);
+	}
+
+	@ParameterizedDataBufferAllocatingTest
+	void readableByteBuffersJoined(DataBufferFactory bufferFactory) {
+		super.bufferFactory = bufferFactory;
+
 		DataBuffer dataBuffer = this.bufferFactory.join(Arrays.asList(stringBuffer("a"),
 				stringBuffer("b"), stringBuffer("c")));
 
@@ -703,17 +735,26 @@ void readableByteBuffers(DataBufferFactory bufferFactory) throws IOException {
 	void writableByteBuffers(DataBufferFactory bufferFactory) {
 		super.bufferFactory = bufferFactory;
 
-		DataBuffer dataBuffer = this.bufferFactory.allocateBuffer(1);
+		DataBuffer dataBuffer = this.bufferFactory.allocateBuffer(3);
+		dataBuffer.write("ab".getBytes(StandardCharsets.UTF_8));
+		dataBuffer.readPosition(1);
 
 		try (DataBuffer.ByteBufferIterator iterator = dataBuffer.writableByteBuffers()) {
 			assertThat(iterator).hasNext();
 			ByteBuffer byteBuffer = iterator.next();
-			byteBuffer.put((byte) 'a');
-			dataBuffer.writePosition(1);
+			assertThat(byteBuffer.position()).isEqualTo(0);
+			assertThat(byteBuffer.limit()).isEqualTo(1);
+			assertThat(byteBuffer.capacity()).isEqualTo(1);
+			assertThat(byteBuffer.remaining()).isEqualTo(1);
+
+			byteBuffer.put((byte) 'c');
+			dataBuffer.writePosition(3);
 
 			assertThat(iterator).isExhausted();
 		}
-		assertThat(dataBuffer.read()).isEqualTo((byte) 'a');
+		byte[] result = new byte[2];
+		dataBuffer.read(result);
+		assertThat(result).containsExactly('b', 'c');
 
 		release(dataBuffer);
 	}
@@ -945,4 +986,21 @@ void shouldHonorSourceBuffersReadPosition(DataBufferFactory bufferFactory) {
 		assertThat(StandardCharsets.UTF_8.decode(byteBuffer).toString()).isEqualTo("b");
 	}
 
+	@ParameterizedDataBufferAllocatingTest // gh-31873
+	void repeatedWrites(DataBufferFactory bufferFactory) {
+		super.bufferFactory = bufferFactory;
+
+		DataBuffer buffer = bufferFactory.allocateBuffer(256);
+		String name = "Müller";
+		int repeatCount = 19;
+		for (int i = 0; i < repeatCount; i++) {
+			buffer.write(name, StandardCharsets.UTF_8);
+		}
+		String result = buffer.toString(StandardCharsets.UTF_8);
+		String expected = name.repeat(repeatCount);
+		assertThat(result).isEqualTo(expected);
+
+		release(buffer);
+	}
+
 }