diff --git a/bson/src/main/org/bson/RawBsonDocument.java b/bson/src/main/org/bson/RawBsonDocument.java index eb672bcef8..9f9fa650d5 100644 --- a/bson/src/main/org/bson/RawBsonDocument.java +++ b/bson/src/main/org/bson/RawBsonDocument.java @@ -35,7 +35,11 @@ import java.io.StringWriter; import java.nio.ByteBuffer; import java.nio.ByteOrder; +import java.util.AbstractMap; +import java.util.ArrayList; import java.util.Collection; +import java.util.LinkedHashSet; +import java.util.List; import java.util.Map; import java.util.NoSuchElementException; import java.util.Set; @@ -225,17 +229,42 @@ public int size() { @Override public Set> entrySet() { - return toBaseBsonDocument().entrySet(); + List> entries = new ArrayList<>(); + try (BsonBinaryReader bsonReader = createReader()) { + bsonReader.readStartDocument(); + while (bsonReader.readBsonType() != BsonType.END_OF_DOCUMENT) { + String key = bsonReader.readName(); + BsonValue value = RawBsonValueHelper.decode(bytes, bsonReader); + entries.add(new AbstractMap.SimpleImmutableEntry<>(key, value)); + } + } + return new LinkedHashSet<>(entries); } @Override public Collection values() { - return toBaseBsonDocument().values(); + List values = new ArrayList<>(); + try (BsonBinaryReader bsonReader = createReader()) { + bsonReader.readStartDocument(); + while (bsonReader.readBsonType() != BsonType.END_OF_DOCUMENT) { + bsonReader.skipName(); + values.add(RawBsonValueHelper.decode(bytes, bsonReader)); + } + } + return new LinkedHashSet<>(values); } @Override public Set keySet() { - return toBaseBsonDocument().keySet(); + List keys = new ArrayList<>(); + try (BsonBinaryReader bsonReader = createReader()) { + bsonReader.readStartDocument(); + while (bsonReader.readBsonType() != BsonType.END_OF_DOCUMENT) { + keys.add(bsonReader.readName()); + bsonReader.skipValue(); + } + } + return new LinkedHashSet<>(keys); } @Override @@ -318,12 +347,19 @@ public String toJson(final JsonWriterSettings settings) { @Override public boolean equals(final Object o) { - return toBaseBsonDocument().equals(o); + return toBsonDocument().equals(o); } @Override public int hashCode() { - return toBaseBsonDocument().hashCode(); + return toBsonDocument().hashCode(); + } + + @Override + public BsonDocument toBsonDocument() { + try (BsonBinaryReader bsonReader = createReader()) { + return new BsonDocumentCodec().decode(bsonReader, DecoderContext.builder().build()); + } } @Override @@ -335,13 +371,6 @@ private BsonBinaryReader createReader() { return new BsonBinaryReader(new ByteBufferBsonInput(getByteBuffer())); } - // Transform to an org.bson.BsonDocument instance - private BsonDocument toBaseBsonDocument() { - try (BsonBinaryReader bsonReader = createReader()) { - return new BsonDocumentCodec().decode(bsonReader, DecoderContext.builder().build()); - } - } - /** * Write the replacement object. * diff --git a/driver-core/src/main/com/mongodb/internal/connection/ByteBufBsonDocument.java b/driver-core/src/main/com/mongodb/internal/connection/ByteBufBsonDocument.java index 4d4ebcc116..639bf7654e 100644 --- a/driver-core/src/main/com/mongodb/internal/connection/ByteBufBsonDocument.java +++ b/driver-core/src/main/com/mongodb/internal/connection/ByteBufBsonDocument.java @@ -24,13 +24,15 @@ import com.mongodb.lang.Nullable; import org.bson.BsonArray; import org.bson.BsonBinaryReader; +import org.bson.BsonBinaryWriter; import org.bson.BsonDocument; +import org.bson.BsonElement; import org.bson.BsonReader; import org.bson.BsonType; import org.bson.BsonValue; import org.bson.ByteBuf; -import org.bson.codecs.BsonDocumentCodec; -import org.bson.codecs.DecoderContext; +import org.bson.RawBsonDocument; +import org.bson.io.BasicOutputBuffer; import org.bson.io.ByteBufferBsonInput; import org.bson.json.JsonMode; import org.bson.json.JsonWriterSettings; @@ -41,13 +43,9 @@ import java.io.ObjectInputStream; import java.io.UnsupportedEncodingException; import java.nio.charset.StandardCharsets; -import java.util.AbstractCollection; -import java.util.AbstractMap; -import java.util.AbstractSet; import java.util.ArrayList; import java.util.Collection; import java.util.HashMap; -import java.util.Iterator; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; @@ -186,7 +184,6 @@ public final class ByteBufBsonDocument extends BsonDocument implements Closeable */ private transient boolean closed; - /** * Creates a {@code ByteBufBsonDocument} from an OP_MSG command message. * @@ -382,60 +379,17 @@ public String getFirstKey() { @Override public Set> entrySet() { - ensureOpen(); - if (cachedDocument != null) { - return cachedDocument.entrySet(); - } - return new AbstractSet>() { - @Override - public Iterator> iterator() { - // Combine body entries with sequence entries - return new CombinedIterator<>(createBodyIterator(IteratorMode.ENTRIES), createSequenceEntryIterator()); - } - - @Override - public int size() { - return ByteBufBsonDocument.this.size(); - } - }; + return toBsonDocument().entrySet(); } @Override public Collection values() { - ensureOpen(); - if (cachedDocument != null) { - return cachedDocument.values(); - } - return new AbstractCollection() { - @Override - public Iterator iterator() { - return new CombinedIterator<>(createBodyIterator(IteratorMode.VALUES), createSequenceValueIterator()); - } - - @Override - public int size() { - return ByteBufBsonDocument.this.size(); - } - }; + return toBsonDocument().values(); } @Override public Set keySet() { - ensureOpen(); - if (cachedDocument != null) { - return cachedDocument.keySet(); - } - return new AbstractSet() { - @Override - public Iterator iterator() { - return new CombinedIterator<>(createBodyIterator(IteratorMode.KEYS), sequenceFields.keySet().iterator()); - } - - @Override - public int size() { - return ByteBufBsonDocument.this.size(); - } - }; + return toBsonDocument().keySet(); } // ==================== Conversion Methods ==================== @@ -452,9 +406,8 @@ public BsonReader asBsonReader() { * *

After this method is called:

*
    - *
  • The result is cached for future calls
    • + *
  • The result is cached as a {@link RawBsonDocument} for future calls
    • *
  • All underlying byte buffers are released
    • - *
  • Sequence field documents are hydrated to regular {@code BsonDocument} instances
    • *
  • All subsequent read operations use the cached document
    • *
* @@ -464,21 +417,29 @@ public BsonReader asBsonReader() { public BsonDocument toBsonDocument() { ensureOpen(); if (cachedDocument == null) { - ByteBuf dup = bodyByteBuf.duplicate(); - try (BsonBinaryReader reader = new BsonBinaryReader(new ByteBufferBsonInput(dup))) { - // Decode body document - BsonDocument doc = new BsonDocumentCodec().decode(reader, DecoderContext.builder().build()); - // Add hydrated sequence fields - for (Map.Entry entry : sequenceFields.entrySet()) { - doc.put(entry.getKey(), entry.getValue().toHydratedArray()); + if (sequenceFields.isEmpty()) { + // No sequence fields: clone body bytes directly + byte[] clonedBytes = new byte[bodyByteBuf.remaining()]; + bodyByteBuf.get(bodyByteBuf.position(), clonedBytes); + cachedDocument = new RawBsonDocument(clonedBytes); + } else { + // With sequence fields: pipe body + extra elements + BasicOutputBuffer buffer = new BasicOutputBuffer(); + ByteBuf dup = bodyByteBuf.duplicate(); + try (BsonBinaryWriter writer = new BsonBinaryWriter(buffer); + BsonBinaryReader reader = new BsonBinaryReader(new ByteBufferBsonInput(dup))) { + List extraElements = new ArrayList<>(); + for (Map.Entry entry : sequenceFields.entrySet()) { + extraElements.add(new BsonElement(entry.getKey(), entry.getValue().asArray())); + } + writer.pipe(reader, extraElements); + } finally { + dup.release(); } - cachedDocument = doc; - closed = true; - // Release buffers since we no longer need them - releaseResources(); - } finally { - dup.release(); + cachedDocument = new RawBsonDocument(buffer.getInternalBuffer(), 0, buffer.getPosition()); } + closed = true; + releaseResources(); } return cachedDocument; } @@ -504,7 +465,7 @@ public String toString() { @Override public BsonDocument clone() { ensureOpen(); - return toBsonDocument().clone(); + return toBsonDocument(); } @SuppressWarnings("EqualsDoesntCheckParameterClass") @@ -645,8 +606,8 @@ private BsonValue getValueFromBody(final String key) { } /** - * Gets the first key from the body, or from sequence fields if body is empty. - * Throws NoSuchElementException if the document is completely empty. + * Gets the first key from the body. + * Throws NoSuchElementException if the body document is completely empty. */ private String getFirstKeyFromBody() { ByteBuf dup = bodyByteBuf.duplicate(); @@ -655,10 +616,6 @@ private String getFirstKeyFromBody() { if (reader.readBsonType() != BsonType.END_OF_DOCUMENT) { return reader.readName(); } - // Body is empty, try sequence fields - if (!sequenceFields.isEmpty()) { - return sequenceFields.keySet().iterator().next(); - } throw new NoSuchElementException(); } finally { dup.release(); @@ -697,147 +654,6 @@ private int countBodyFields() { } } - // ==================== Iterator Support ==================== - - /** - * Mode for the body iterator, determining what type of elements it produces. - */ - private enum IteratorMode { ENTRIES, KEYS, VALUES } - - /** - * Creates an iterator over the body document fields. - * - *

The iterator creates a duplicated ByteBuf that is temporarily tracked for safety. - * When iteration completes normally, the buffer is released immediately and removed from tracking. - * This prevents accumulation of finished iterator buffers while ensuring cleanup if the parent - * document is closed before iteration completes.

- * - * @param mode Determines whether to return entries, keys, or values - * @return An iterator of the appropriate type - */ - @SuppressWarnings("unchecked") - private Iterator createBodyIterator(final IteratorMode mode) { - return new Iterator() { - private final Closeable resourceHandle; - private ByteBuf duplicatedByteBuf; - private BsonBinaryReader reader; - private boolean started; - private boolean finished; - - { - // Create duplicate buffer for iteration and track it temporarily - duplicatedByteBuf = bodyByteBuf.duplicate(); - resourceHandle = () -> { - if (duplicatedByteBuf != null) { - try { - if (reader != null) { - reader.close(); - } - } catch (Exception e) { - // Ignore - } - duplicatedByteBuf.release(); - duplicatedByteBuf = null; - reader = null; - } - }; - trackedResources.add(resourceHandle); - reader = new BsonBinaryReader(new ByteBufferBsonInput(duplicatedByteBuf)); - } - - @Override - public boolean hasNext() { - if (finished) { - return false; - } - ensureOpen(); - if (!started) { - reader.readStartDocument(); - reader.readBsonType(); - started = true; - } - boolean hasNext = reader.getCurrentBsonType() != BsonType.END_OF_DOCUMENT; - if (!hasNext) { - cleanup(); - } - return hasNext; - } - - @Override - public T next() { - if (!hasNext()) { - throw new NoSuchElementException(); - } - ensureOpen(); - String key = reader.readName(); - BsonValue value = readBsonValue(duplicatedByteBuf, reader, trackedResources); - reader.readBsonType(); - - switch (mode) { - case ENTRIES: - return (T) new AbstractMap.SimpleImmutableEntry<>(key, value); - case KEYS: - return (T) key; - case VALUES: - return (T) value; - default: - throw new IllegalStateException("Unknown iterator mode: " + mode); - } - } - - private void cleanup() { - if (!finished) { - finished = true; - // Remove from tracked resources since we're cleaning up immediately - trackedResources.remove(resourceHandle); - try { - resourceHandle.close(); - } catch (Exception e) { - // Ignore - } - } - } - }; - } - - /** - * Creates an iterator over sequence fields as map entries. - * Each entry contains the field name and its array value. - */ - private Iterator> createSequenceEntryIterator() { - Iterator> iter = sequenceFields.entrySet().iterator(); - return new Iterator>() { - @Override - public boolean hasNext() { - return iter.hasNext(); - } - - @Override - public Entry next() { - Entry entry = iter.next(); - return new AbstractMap.SimpleImmutableEntry<>(entry.getKey(), entry.getValue().asArray()); - } - }; - } - - /** - * Creates an iterator over sequence field values (arrays). - */ - private Iterator createSequenceValueIterator() { - Iterator iter = sequenceFields.values().iterator(); - return new Iterator() { - @Override - public boolean hasNext() { - return iter.hasNext(); - } - - @Override - public BsonValue next() { - return iter.next().asArray(); - } - }; - } - // ==================== Resource Management Helpers ==================== /** @@ -977,69 +793,6 @@ boolean containsValue(final Object value) { return value instanceof BsonValue && asArray().asArray().contains(value); } - /** - * Converts this sequence to a BsonArray of regular BsonDocument instances. - * - *

Used by {@link ByteBufBsonDocument#toBsonDocument()} to fully hydrate the document. - * Unlike {@link #asArray()}, this creates regular BsonDocument instances, not - * ByteBufBsonDocument wrappers.

- * - * @return A BsonArray containing fully deserialized BsonDocument instances - */ - BsonArray toHydratedArray() { - ByteBuf dup = sequenceByteBuf.duplicate(); - try { - List hydratedDocs = new ArrayList<>(); - while (dup.hasRemaining()) { - int docStart = dup.position(); - int docSize = dup.getInt(); - int docEnd = docStart + docSize; - ByteBuf docBuf = dup.duplicate().position(docStart).limit(docEnd); - try (BsonBinaryReader reader = new BsonBinaryReader(new ByteBufferBsonInput(docBuf))) { - hydratedDocs.add(new BsonDocumentCodec().decode(reader, DecoderContext.builder().build())); - } finally { - docBuf.release(); - } - dup.position(docEnd); - } - return new BsonArray(hydratedDocs); - } finally { - dup.release(); - } - } } - /** - * An iterator that combines two iterators sequentially. - * - *

Used to merge body field iteration with sequence field iteration, - * presenting a unified view of all document fields.

- * - * @param The type of elements returned by the iterator - */ - private static final class CombinedIterator implements Iterator { - private final Iterator primary; - private final Iterator secondary; - - CombinedIterator(final Iterator primary, final Iterator secondary) { - this.primary = primary; - this.secondary = secondary; - } - - @Override - public boolean hasNext() { - return primary.hasNext() || secondary.hasNext(); - } - - @Override - public T next() { - if (primary.hasNext()) { - return primary.next(); - } - if (secondary.hasNext()) { - return secondary.next(); - } - throw new NoSuchElementException(); - } - } } diff --git a/driver-core/src/test/unit/com/mongodb/internal/connection/ByteBufBsonArrayTest.java b/driver-core/src/test/unit/com/mongodb/internal/connection/ByteBufBsonArrayTest.java index 637f89cb34..4b05607a56 100644 --- a/driver-core/src/test/unit/com/mongodb/internal/connection/ByteBufBsonArrayTest.java +++ b/driver-core/src/test/unit/com/mongodb/internal/connection/ByteBufBsonArrayTest.java @@ -49,8 +49,6 @@ import org.junit.jupiter.api.DisplayName; import org.junit.jupiter.api.Test; -import java.io.ByteArrayOutputStream; -import java.io.IOException; import java.nio.ByteBuffer; import java.util.Date; import java.util.Iterator; @@ -351,17 +349,19 @@ void testAllBsonTypes() { } static ByteBufBsonArray fromBsonValues(final List values) { - BsonDocument document = new BsonDocument() - .append("a", new BsonArray(values)); + BsonDocument document = new BsonDocument("a", new BsonArray(values)); BasicOutputBuffer buffer = new BasicOutputBuffer(); new BsonDocumentCodec().encode(new BsonBinaryWriter(buffer), document, EncoderContext.builder().build()); - ByteArrayOutputStream baos = new ByteArrayOutputStream(); - try { - buffer.pipe(baos); - } catch (IOException e) { - throw new RuntimeException("impossible!"); - } - ByteBuf documentByteBuf = new ByteBufNIO(ByteBuffer.wrap(baos.toByteArray())); - return (ByteBufBsonArray) new ByteBufBsonDocument(documentByteBuf).entrySet().iterator().next().getValue(); + byte[] bytes = new byte[buffer.getPosition()]; + System.arraycopy(buffer.getInternalBuffer(), 0, bytes, 0, bytes.length); + // Skip past the outer document header to the array value bytes. + // Document format: [4-byte size][type byte (0x04)][field name "a\0"][array bytes...][0x00] + int arrayOffset = 4 + 1 + 2; // doc size + type byte + "a" + null terminator + int arraySize = (bytes[arrayOffset] & 0xFF) + | ((bytes[arrayOffset + 1] & 0xFF) << 8) + | ((bytes[arrayOffset + 2] & 0xFF) << 16) + | ((bytes[arrayOffset + 3] & 0xFF) << 24); + ByteBuf arrayByteBuf = new ByteBufNIO(ByteBuffer.wrap(bytes, arrayOffset, arraySize)); + return new ByteBufBsonArray(arrayByteBuf); } } diff --git a/driver-core/src/test/unit/com/mongodb/internal/connection/ByteBufBsonDocumentTest.java b/driver-core/src/test/unit/com/mongodb/internal/connection/ByteBufBsonDocumentTest.java index 1f61f309d1..7db52c1631 100644 --- a/driver-core/src/test/unit/com/mongodb/internal/connection/ByteBufBsonDocumentTest.java +++ b/driver-core/src/test/unit/com/mongodb/internal/connection/ByteBufBsonDocumentTest.java @@ -438,52 +438,32 @@ void deeplyNestedClosedRecursively() { } @Test - @DisplayName("Iteration tracks resources correctly") - void iterationTracksResources() { + @DisplayName("Iterators work as expected") + void iteratorsWorksAsExpected() { BsonDocument doc = new BsonDocument() .append("doc1", new BsonDocument("a", new BsonInt32(1))) .append("arr1", new BsonArray(asList(new BsonInt32(2), new BsonInt32(3)))) .append("primitive", new BsonString("test")); - ByteBuf buf = createByteBufFromDocument(doc); - ByteBufBsonDocument byteBufDoc = new ByteBufBsonDocument(buf); - - int count = 0; - for (Map.Entry entry : byteBufDoc.entrySet()) { - assertNotNull(entry.getKey()); - assertNotNull(entry.getValue()); - count++; - } - assertEquals(3, count); - - byteBufDoc.close(); - assertThrows(IllegalStateException.class, byteBufDoc::size); - } - - @Test - @DisplayName("Iterators ensure the resource is still open") - void iteratorsEnsureResourceIsStillOpen() { - BsonDocument doc = new BsonDocument() - .append("doc1", new BsonDocument("a", new BsonInt32(1))) - .append("arr1", new BsonArray(asList(new BsonInt32(2), new BsonInt32(3)))) - .append("primitive", new BsonString("test")); - - ByteBuf buf = createByteBufFromDocument(doc); - ByteBufBsonDocument byteBufDoc = new ByteBufBsonDocument(buf); + try (ByteBufBsonDocument byteBufDoc = new ByteBufBsonDocument(createByteBufFromDocument(doc))) { - Iterator keysIterator = byteBufDoc.keySet().iterator(); - assertDoesNotThrow(keysIterator::hasNext); + int count = 0; + for (Map.Entry entry : byteBufDoc.entrySet()) { + assertNotNull(entry.getKey()); + assertNotNull(entry.getValue()); + count++; + } + assertEquals(3, count); - Iterator nestedKeysIterator = byteBufDoc.getDocument("doc1").keySet().iterator(); - assertDoesNotThrow(nestedKeysIterator::hasNext); + Iterator keysIterator = byteBufDoc.keySet().iterator(); + assertDoesNotThrow(keysIterator::hasNext); - Iterator arrayIterator = byteBufDoc.getArray("arr1").iterator(); - assertDoesNotThrow(arrayIterator::hasNext); + Iterator nestedKeysIterator = byteBufDoc.getDocument("doc1").keySet().iterator(); + assertDoesNotThrow(nestedKeysIterator::hasNext); - byteBufDoc.close(); - assertThrows(IllegalStateException.class, keysIterator::hasNext); - assertThrows(IllegalStateException.class, nestedKeysIterator::hasNext); - assertThrows(IllegalStateException.class, arrayIterator::hasNext); + Iterator arrayIterator = byteBufDoc.getArray("arr1").iterator(); + assertDoesNotThrow(arrayIterator::hasNext); + } } @Test