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"); * Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License. * you may not use this file except in compliance with the License.
@ -263,15 +263,16 @@ public interface DataBuffer {
default DataBuffer write(CharSequence charSequence, Charset charset) { default DataBuffer write(CharSequence charSequence, Charset charset) {
Assert.notNull(charSequence, "CharSequence must not be null"); Assert.notNull(charSequence, "CharSequence must not be null");
Assert.notNull(charset, "Charset must not be null"); Assert.notNull(charset, "Charset must not be null");
if (charSequence.length() > 0) { if (!charSequence.isEmpty()) {
CharsetEncoder encoder = charset.newEncoder() CharsetEncoder encoder = charset.newEncoder()
.onMalformedInput(CodingErrorAction.REPLACE) .onMalformedInput(CodingErrorAction.REPLACE)
.onUnmappableCharacter(CodingErrorAction.REPLACE); .onUnmappableCharacter(CodingErrorAction.REPLACE);
CharBuffer src = CharBuffer.wrap(charSequence); 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) { while (true) {
ensureWritable(cap);
CoderResult cr; CoderResult cr;
if (src.hasRemaining()) {
try (ByteBufferIterator iterator = writableByteBuffers()) { try (ByteBufferIterator iterator = writableByteBuffers()) {
Assert.state(iterator.hasNext(), "No ByteBuffer available"); Assert.state(iterator.hasNext(), "No ByteBuffer available");
ByteBuffer dest = iterator.next(); ByteBuffer dest = iterator.next();
@ -279,13 +280,18 @@ public interface DataBuffer {
if (cr.isUnderflow()) { if (cr.isUnderflow()) {
cr = encoder.flush(dest); cr = encoder.flush(dest);
} }
writePosition(dest.position()); writePosition(writePosition() + dest.position());
}
}
else {
cr = CoderResult.UNDERFLOW;
} }
if (cr.isUnderflow()) { if (cr.isUnderflow()) {
break; break;
} }
if (cr.isOverflow()) { else if (cr.isOverflow()) {
cap = 2 * cap + 1; 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"); * Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License. * you may not use this file except in compliance with the License.
@ -416,16 +416,15 @@ public class DefaultDataBuffer implements DataBuffer {
@Override @Override
public DataBuffer.ByteBufferIterator readableByteBuffers() { public DataBuffer.ByteBufferIterator readableByteBuffers() {
ByteBuffer readOnly = this.byteBuffer.asReadOnlyBuffer(); ByteBuffer readOnly = this.byteBuffer.slice(this.readPosition, readableByteCount())
readOnly.clear().position(this.readPosition).limit(this.writePosition - this.readPosition); .asReadOnlyBuffer();
return new ByteBufferIterator(readOnly); return new ByteBufferIterator(readOnly);
} }
@Override @Override
public DataBuffer.ByteBufferIterator writableByteBuffers() { public DataBuffer.ByteBufferIterator writableByteBuffers() {
ByteBuffer duplicate = this.byteBuffer.duplicate(); ByteBuffer slice = this.byteBuffer.slice(this.writePosition, writableByteCount());
duplicate.clear().position(this.writePosition).limit(this.capacity - this.writePosition); return new ByteBufferIterator(slice);
return new ByteBufferIterator(duplicate);
} }
@Override @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"); * Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with 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) { void readableByteBuffers(DataBufferFactory bufferFactory) {
super.bufferFactory = 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"), DataBuffer dataBuffer = this.bufferFactory.join(Arrays.asList(stringBuffer("a"),
stringBuffer("b"), stringBuffer("c"))); stringBuffer("b"), stringBuffer("c")));
@ -702,17 +734,26 @@ class DataBufferTests extends AbstractDataBufferAllocatingTests {
void writableByteBuffers(DataBufferFactory bufferFactory) { void writableByteBuffers(DataBufferFactory bufferFactory) {
super.bufferFactory = 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()) { try (DataBuffer.ByteBufferIterator iterator = dataBuffer.writableByteBuffers()) {
assertThat(iterator).hasNext(); assertThat(iterator).hasNext();
ByteBuffer byteBuffer = iterator.next(); ByteBuffer byteBuffer = iterator.next();
byteBuffer.put((byte) 'a'); assertThat(byteBuffer.position()).isEqualTo(0);
dataBuffer.writePosition(1); 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(iterator).isExhausted();
} }
assertThat(dataBuffer.read()).isEqualTo((byte) 'a'); byte[] result = new byte[2];
dataBuffer.read(result);
assertThat(result).containsExactly('b', 'c');
release(dataBuffer); release(dataBuffer);
} }
@ -944,4 +985,21 @@ class DataBufferTests extends AbstractDataBufferAllocatingTests {
assertThat(StandardCharsets.UTF_8.decode(byteBuffer).toString()).isEqualTo("b"); 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);
}
} }