Ensure correct capacity in DefaultDataBuffer

Closes gh-31873
This commit is contained in:
Arjen Poutsma 2024-01-08 14:45:31 +01:00
parent bb1cdb6b48
commit 3452354a11
3 changed files with 87 additions and 24 deletions

View File

@ -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 @@ public interface DataBuffer {
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);
}
}
}

View File

@ -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 class DefaultDataBuffer implements DataBuffer {
@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

View File

@ -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.
@ -677,6 +677,38 @@ class DataBufferTests extends AbstractDataBufferAllocatingTests {
void readableByteBuffers(DataBufferFactory bufferFactory) {
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")));
@ -702,17 +734,26 @@ class DataBufferTests extends AbstractDataBufferAllocatingTests {
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);
}
@ -944,4 +985,21 @@ class DataBufferTests extends AbstractDataBufferAllocatingTests {
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);
}
}