Skip to content

Commit 98e0761

Browse files
rozzavbabanin
andauthored
Optimize RawBsonDocument encode and decode by eliminating intermediate allocations (#1988)
- Add BsonInput.pipe(BsonOutput, int) to remove the temporary byte[] copy in BsonBinaryWriter.pipeDocument() on both encode and decode paths - Add public getByteBacking(), getByteOffset(), getByteLength() on RawBsonDocument to expose the backing byte array JAVA-6133 --------- Co-authored-by: slav.babanin <slav.babanin@mongodb.com>
1 parent 40ee8d5 commit 98e0761

9 files changed

Lines changed: 459 additions & 20 deletions

File tree

bson/src/main/org/bson/BsonBinaryWriter.java

Lines changed: 42 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -334,6 +334,26 @@ public void pipe(final BsonReader reader) {
334334
pipeDocument(reader, null);
335335
}
336336

337+
/**
338+
* Pipes an encoded BSON document from the given byte array to this writer.
339+
*
340+
* @param bytes the byte array containing the encoded BSON document
341+
* @param offset the offset into the byte array
342+
* @param length the length of the encoded BSON document
343+
* @since 5.8
344+
*/
345+
public void pipe(final byte[] bytes, final int offset, final int length) {
346+
notNull("bytes", bytes);
347+
checkMinDocumentSize(length);
348+
if (getState() == State.VALUE) {
349+
bsonOutput.writeByte(BsonType.DOCUMENT.getValue());
350+
writeCurrentName();
351+
}
352+
int pipedDocumentStartPosition = bsonOutput.getPosition();
353+
bsonOutput.writeBytes(bytes, offset, length);
354+
completePipeDocument(pipedDocumentStartPosition);
355+
}
356+
337357
@Override
338358
public void pipe(final BsonReader reader, final List<BsonElement> extraElements) {
339359
notNull("reader", reader);
@@ -350,14 +370,10 @@ private void pipeDocument(final BsonReader reader, final List<BsonElement> extra
350370
}
351371
BsonInput bsonInput = binaryReader.getBsonInput();
352372
int size = bsonInput.readInt32();
353-
if (size < 5) {
354-
throw new BsonSerializationException("Document size must be at least 5");
355-
}
373+
checkMinDocumentSize(size);
356374
int pipedDocumentStartPosition = bsonOutput.getPosition();
357375
bsonOutput.writeInt32(size);
358-
byte[] bytes = new byte[size - 4];
359-
bsonInput.readBytes(bytes);
360-
bsonOutput.writeBytes(bytes);
376+
bsonInput.pipe(bsonOutput, size - 4);
361377

362378
binaryReader.setState(AbstractBsonReader.State.TYPE);
363379

@@ -371,24 +387,27 @@ private void pipeDocument(final BsonReader reader, final List<BsonElement> extra
371387
setContext(getContext().getParentContext());
372388
}
373389

374-
if (getContext() == null) {
375-
setState(State.DONE);
376-
} else {
377-
if (getContext().getContextType() == BsonContextType.JAVASCRIPT_WITH_SCOPE) {
378-
backpatchSize(); // size of the JavaScript with scope value
379-
setContext(getContext().getParentContext());
380-
}
381-
setState(getNextState());
382-
}
383-
384-
validateSize(bsonOutput.getPosition() - pipedDocumentStartPosition);
390+
completePipeDocument(pipedDocumentStartPosition);
385391
} else if (extraElements != null) {
386392
super.pipe(reader, extraElements);
387393
} else {
388394
super.pipe(reader);
389395
}
390396
}
391397

398+
private void completePipeDocument(final int pipedDocumentStartPosition) {
399+
if (getContext() == null) {
400+
setState(State.DONE);
401+
} else {
402+
if (getContext().getContextType() == BsonContextType.JAVASCRIPT_WITH_SCOPE) {
403+
backpatchSize(); // size of the JavaScript with scope value
404+
setContext(getContext().getParentContext());
405+
}
406+
setState(getNextState());
407+
}
408+
validateSize(bsonOutput.getPosition() - pipedDocumentStartPosition);
409+
}
410+
392411
/**
393412
* Sets a maximum size for documents from this point.
394413
*
@@ -426,6 +445,12 @@ public void reset() {
426445
mark = null;
427446
}
428447

448+
private static void checkMinDocumentSize(final int size) {
449+
if (size < 5) {
450+
throw new BsonSerializationException("Document size must be at least 5");
451+
}
452+
}
453+
429454
private void writeCurrentName() {
430455
if (getContext().getContextType() == BsonContextType.ARRAY) {
431456
int index = getContext().index++;

bson/src/main/org/bson/RawBsonDocument.java

Lines changed: 35 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@
4444
import static org.bson.assertions.Assertions.notNull;
4545

4646
/**
47-
* An immutable BSON document that is represented using only the raw bytes.
47+
* A BSON document that is represented using only the raw bytes.
4848
*
4949
* @since 3.0
5050
*/
@@ -144,6 +144,40 @@ public ByteBuf getByteBuffer() {
144144
return new ByteBufNIO(buffer);
145145
}
146146

147+
/**
148+
* Returns the byte array backing this document. The returned array may be larger than the BSON document itself;
149+
* only the range from {@link #getByteOffset()} to {@code getByteOffset() + }{@link #getByteLength()} contains
150+
* valid document bytes. Changes to the returned array will be reflected in this document.
151+
*
152+
* @return the backing byte array
153+
* @since 5.8
154+
* @see #getByteOffset()
155+
* @see #getByteLength()
156+
*/
157+
public byte[] getBackingArray() {
158+
return bytes;
159+
}
160+
161+
/**
162+
* Returns the offset into the {@linkplain #getBackingArray() backing byte array} where this document starts.
163+
*
164+
* @return the offset
165+
* @since 5.8
166+
*/
167+
public int getByteOffset() {
168+
return offset;
169+
}
170+
171+
/**
172+
* Returns the length of this document within the {@linkplain #getBackingArray() backing byte array}.
173+
*
174+
* @return the length
175+
* @since 5.8
176+
*/
177+
public int getByteLength() {
178+
return length;
179+
}
180+
147181
/**
148182
* Decode this into a document.
149183
*

bson/src/main/org/bson/codecs/RawBsonDocumentCodec.java

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -40,8 +40,17 @@ public RawBsonDocumentCodec() {
4040

4141
@Override
4242
public void encode(final BsonWriter writer, final RawBsonDocument value, final EncoderContext encoderContext) {
43-
try (BsonBinaryReader reader = new BsonBinaryReader(new ByteBufferBsonInput(value.getByteBuffer()))) {
44-
writer.pipe(reader);
43+
if (writer instanceof BsonBinaryWriter) {
44+
// Fast path. The pipe method should ideally exist on BsonWriter, but adding it as
45+
// abstract would be a breaking change, and adding it as a default method would force
46+
// BsonWriter to depend on BsonBinaryReader/ByteBufferBsonInput, violating the
47+
// interface's abstraction.
48+
// TODO JAVA-6211 move pipe(byte[], int, int) to BsonWriter to remove this instanceof.
49+
((BsonBinaryWriter) writer).pipe(value.getBackingArray(), value.getByteOffset(), value.getByteLength());
50+
} else {
51+
try (BsonBinaryReader reader = new BsonBinaryReader(new ByteBufferBsonInput(value.getByteBuffer()))) {
52+
writer.pipe(reader);
53+
}
4554
}
4655
}
4756

bson/src/main/org/bson/io/BsonInput.java

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -127,6 +127,19 @@ public interface BsonInput extends Closeable {
127127
*/
128128
boolean hasRemaining();
129129

130+
/**
131+
* Pipes the specified number of bytes from {@linkplain BsonInput this} input to the given {@linkplain BsonOutput output}.
132+
*
133+
* @param output the output to pipe to
134+
* @param numBytes the number of bytes to pipe
135+
* @since 5.8
136+
*/
137+
default void pipe(BsonOutput output, int numBytes) {
138+
byte[] bytes = new byte[numBytes];
139+
readBytes(bytes);
140+
output.writeBytes(bytes);
141+
}
142+
130143
@Override
131144
void close();
132145
}

bson/src/main/org/bson/io/ByteBufferBsonInput.java

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -275,6 +275,24 @@ public boolean hasRemaining() {
275275
return buffer.hasRemaining();
276276
}
277277

278+
@Override
279+
public void pipe(final BsonOutput output, final int numBytes) {
280+
ensureOpen();
281+
ensureAvailable(numBytes);
282+
283+
if (buffer.isBackedByArray()) {
284+
int position = buffer.position();
285+
int arrayOffset = buffer.arrayOffset();
286+
output.writeBytes(buffer.array(), arrayOffset + position, numBytes);
287+
buffer.position(position + numBytes);
288+
} else {
289+
// Fallback: use temporary buffer for non-array-backed buffers
290+
byte[] temp = new byte[numBytes];
291+
buffer.get(temp);
292+
output.writeBytes(temp);
293+
}
294+
}
295+
278296
@Override
279297
public void close() {
280298
buffer.release();

bson/src/test/unit/org/bson/BsonBinaryWriterTest.java

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -802,7 +802,34 @@ public void testPipeOfDocumentWithInvalidSize() {
802802
// expected
803803
}
804804
}
805+
}
806+
807+
@Test
808+
public void testPipeOfRawBytes() {
809+
BasicOutputBuffer sourceBuffer = new BasicOutputBuffer();
810+
try (BsonBinaryWriter sourceWriter = new BsonBinaryWriter(sourceBuffer)) {
811+
sourceWriter.writeStartDocument();
812+
sourceWriter.writeBoolean("a", true);
813+
sourceWriter.writeEndDocument();
814+
}
815+
byte[] documentBytes = sourceBuffer.toByteArray();
805816

817+
BasicOutputBuffer destBuffer = new BasicOutputBuffer();
818+
try (BsonBinaryWriter destWriter = new BsonBinaryWriter(destBuffer)) {
819+
destWriter.pipe(documentBytes, 0, documentBytes.length);
820+
}
821+
822+
assertArrayEquals(documentBytes, destBuffer.toByteArray());
823+
}
824+
825+
@Test
826+
public void testPipeOfRawBytesWithInvalidSize() {
827+
byte[] bytes = {4, 0, 0, 0}; // minimum document size is 5
828+
829+
BasicOutputBuffer newBuffer = new BasicOutputBuffer();
830+
try (BsonBinaryWriter newWriter = new BsonBinaryWriter(newBuffer)) {
831+
assertThrows(BsonSerializationException.class, () -> newWriter.pipe(bytes, 0, bytes.length));
832+
}
806833
}
807834

808835
// CHECKSTYLE:OFF
Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
1+
/*
2+
* Copyright 2008-present MongoDB, Inc.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.bson;
18+
19+
import org.bson.codecs.BsonDocumentCodec;
20+
import org.bson.codecs.EncoderContext;
21+
import org.bson.io.BasicOutputBuffer;
22+
import org.junit.jupiter.params.ParameterizedTest;
23+
import org.junit.jupiter.api.Named;
24+
import org.junit.jupiter.params.provider.Arguments;
25+
import org.junit.jupiter.params.provider.MethodSource;
26+
27+
import java.util.Arrays;
28+
import java.util.stream.Stream;
29+
30+
import static org.junit.jupiter.api.Assertions.assertArrayEquals;
31+
import static org.junit.jupiter.api.Assertions.assertEquals;
32+
33+
class RawBsonDocumentTest {
34+
35+
private static final BsonDocument DOCUMENT = new BsonDocument()
36+
.append("a", new BsonInt32(1))
37+
.append("b", new BsonInt32(2))
38+
.append("c", new BsonDocument("x", BsonBoolean.TRUE))
39+
.append("d", new BsonArray(Arrays.asList(
40+
new BsonDocument("y", BsonBoolean.FALSE),
41+
new BsonArray(Arrays.asList(new BsonInt32(1))))));
42+
43+
private static final byte[] DOCUMENT_BYTES = encodeDocument();
44+
45+
static Stream<Arguments> backingArrayAccessors() {
46+
int documentLength = DOCUMENT_BYTES.length;
47+
48+
Stream.Builder<Arguments> builder = Stream.builder();
49+
builder.add(Arguments.of(createFromDocument(), 0, documentLength));
50+
builder.add(Arguments.of(createFromByteArray(), 0, documentLength));
51+
52+
for (int padding = 1; padding <= 2; padding++) {
53+
builder.add(Arguments.of(createPaddedBefore(padding), padding, documentLength));
54+
builder.add(Arguments.of(createPaddedAfter(padding), 0, documentLength));
55+
builder.add(Arguments.of(createPaddedBoth(padding), padding, documentLength));
56+
}
57+
58+
return builder.build();
59+
}
60+
61+
@ParameterizedTest(name = "{0}, expectedOffset={1}, expectedLength={2}")
62+
@MethodSource("backingArrayAccessors")
63+
void shouldExposeBackingArrayOffsetAndLength(final RawBsonDocument rawDocument,
64+
final int expectedOffset,
65+
final int expectedLength) {
66+
assertEquals(expectedOffset, rawDocument.getByteOffset());
67+
assertEquals(expectedLength, rawDocument.getByteLength());
68+
assertArrayEquals(DOCUMENT_BYTES,
69+
Arrays.copyOfRange(
70+
rawDocument.getBackingArray(),
71+
rawDocument.getByteOffset(),
72+
rawDocument.getByteOffset() + rawDocument.getByteLength()));
73+
}
74+
75+
private static Named<RawBsonDocument> createFromDocument() {
76+
return Named.of("from document", new RawBsonDocument(DOCUMENT, new BsonDocumentCodec()));
77+
}
78+
79+
private static Named<RawBsonDocument> createFromByteArray() {
80+
return Named.of("from byte array", new RawBsonDocument(DOCUMENT_BYTES));
81+
}
82+
83+
private static Named<RawBsonDocument> createPaddedBefore(final int padding) {
84+
byte[] padded = new byte[DOCUMENT_BYTES.length + padding];
85+
System.arraycopy(DOCUMENT_BYTES, 0, padded, padding, DOCUMENT_BYTES.length);
86+
return Named.of("padded before " + padding, new RawBsonDocument(padded, padding, DOCUMENT_BYTES.length));
87+
}
88+
89+
private static Named<RawBsonDocument> createPaddedAfter(final int padding) {
90+
byte[] padded = new byte[DOCUMENT_BYTES.length + padding];
91+
System.arraycopy(DOCUMENT_BYTES, 0, padded, 0, DOCUMENT_BYTES.length);
92+
return Named.of("padded after " + padding, new RawBsonDocument(padded, 0, DOCUMENT_BYTES.length));
93+
}
94+
95+
private static Named<RawBsonDocument> createPaddedBoth(final int padding) {
96+
byte[] padded = new byte[DOCUMENT_BYTES.length + padding * 2];
97+
System.arraycopy(DOCUMENT_BYTES, 0, padded, padding, DOCUMENT_BYTES.length);
98+
return Named.of("padded both " + padding, new RawBsonDocument(padded, padding, DOCUMENT_BYTES.length));
99+
}
100+
101+
private static byte[] encodeDocument() {
102+
BasicOutputBuffer buffer = new BasicOutputBuffer();
103+
new BsonDocumentCodec().encode(new BsonBinaryWriter(buffer), DOCUMENT, EncoderContext.builder().build());
104+
return Arrays.copyOf(buffer.getInternalBuffer(), buffer.getPosition());
105+
}
106+
}

0 commit comments

Comments
 (0)