ByteBufBsonDocument & ByteBufBsonArray refactorings (#1874)

* ByteBufBsonDocument & ByteBufBsonArray refactoring

* Now implement `Closeable` to track and manage lifecycle with try-with-resources
* `ByteBufBsonDocument`: Added resource tracking, OP_MSG parsing, caching strategy
* `ByteBufBsonArray`: Added resource tracking and cleanup

CommandMessage Changes:

* `getCommandDocument()` returns `ByteBufBsonDocument` (was `BsonDocument`)
* Delegates document composition to `ByteBufBsonDocument`
* Simplified `OP_MSG` document sequence parsing

JAVA-6010

* Nit fixes and usability improvements

If the document is hydrated allow continued use after resource closing.

* PR updates

* PR updates - ensure iterators track open resources and normalize the tests

Author: rozza

driver-core/src/main/com/mongodb/internal/connection/ByteBufBsonArray.java

@@ -16,13 +16,17 @@
 
 package com.mongodb.internal.connection;
 
+import com.mongodb.annotations.NotThreadSafe;
+import com.mongodb.internal.diagnostics.logging.Logger;
+import com.mongodb.internal.diagnostics.logging.Loggers;
 import org.bson.BsonArray;
 import org.bson.BsonBinaryReader;
 import org.bson.BsonType;
 import org.bson.BsonValue;
 import org.bson.ByteBuf;
 import org.bson.io.ByteBufferBsonInput;
 
+import java.io.Closeable;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Iterator;
@@ -33,20 +37,33 @@
 
 import static com.mongodb.internal.connection.ByteBufBsonHelper.readBsonValue;
 
-final class ByteBufBsonArray extends BsonArray {
+@NotThreadSafe
+final class ByteBufBsonArray extends BsonArray implements Closeable {
+    private static final Logger LOGGER = Loggers.getLogger("connection");
     private final ByteBuf byteBuf;
 
+    /**
+     * List of resources that need to be closed when this array is closed.
+     * Tracks the main ByteBuf and iterator duplicates. Iterator buffers are automatically
+     * removed and released when iteration completes normally to prevent memory accumulation.
+     */
+    private final List<Closeable> trackedResources = new ArrayList<>();
+    private boolean closed;
+
     ByteBufBsonArray(final ByteBuf byteBuf) {
         this.byteBuf = byteBuf;
+        trackedResources.add(byteBuf::release);
     }
 
     @Override
     public Iterator<BsonValue> iterator() {
+        ensureOpen();
         return new ByteBufBsonArrayIterator();
     }
 
     @Override
     public List<BsonValue> getValues() {
+        ensureOpen();
         List<BsonValue> values = new ArrayList<>();
         for (BsonValue cur: this) {
             //noinspection UseBulkOperation
@@ -59,6 +76,7 @@ public List<BsonValue> getValues() {
 
     @Override
     public int size() {
+        ensureOpen();
         int size = 0;
         for (BsonValue ignored : this) {
             size++;
@@ -68,11 +86,13 @@ public int size() {
 
     @Override
     public boolean isEmpty() {
+        ensureOpen();
         return !iterator().hasNext();
     }
 
     @Override
     public boolean equals(final Object o) {
+        ensureOpen();
         if (o == this) {
             return true;
         }
@@ -91,6 +111,7 @@ public boolean equals(final Object o) {
 
     @Override
     public int hashCode() {
+        ensureOpen();
         int hashCode = 1;
         for (BsonValue cur : this) {
             hashCode = 31 * hashCode + (cur == null ? 0 : cur.hashCode());
@@ -100,6 +121,7 @@ public int hashCode() {
 
     @Override
     public boolean contains(final Object o) {
+        ensureOpen();
         for (BsonValue cur : this) {
             if (Objects.equals(o, cur)) {
                 return true;
@@ -111,6 +133,7 @@ public boolean contains(final Object o) {
 
     @Override
     public Object[] toArray() {
+        ensureOpen();
         Object[] retVal = new Object[size()];
         Iterator<BsonValue> it = iterator();
         for (int i = 0; i < retVal.length; i++) {
@@ -122,6 +145,7 @@ public Object[] toArray() {
     @Override
     @SuppressWarnings("unchecked")
     public <T> T[] toArray(final T[] a) {
+        ensureOpen();
         int size = size();
         T[] retVal = a.length >= size ? a : (T[]) java.lang.reflect.Array.newInstance(a.getClass().getComponentType(), size);
         Iterator<BsonValue> it = iterator();
@@ -133,6 +157,7 @@ public <T> T[] toArray(final T[] a) {
 
     @Override
     public boolean containsAll(final Collection<?> c) {
+        ensureOpen();
         for (Object e : c) {
             if (!contains(e)) {
                 return false;
@@ -143,6 +168,7 @@ public boolean containsAll(final Collection<?> c) {
 
     @Override
     public BsonValue get(final int index) {
+        ensureOpen();
         if (index < 0) {
             throw new IndexOutOfBoundsException("Index out of range: " + index);
         }
@@ -159,6 +185,7 @@ public BsonValue get(final int index) {
 
     @Override
     public int indexOf(final Object o) {
+        ensureOpen();
         int i = 0;
         for (BsonValue cur : this) {
             if (Objects.equals(o, cur)) {
@@ -172,6 +199,7 @@ public int indexOf(final Object o) {
 
     @Override
     public int lastIndexOf(final Object o) {
+        ensureOpen();
         ListIterator<BsonValue> listIterator = listIterator(size());
         while (listIterator.hasPrevious()) {
             if (Objects.equals(o, listIterator.previous())) {
@@ -183,17 +211,20 @@ public int lastIndexOf(final Object o) {
 
     @Override
     public ListIterator<BsonValue> listIterator() {
+        ensureOpen();
         return listIterator(0);
     }
 
     @Override
     public ListIterator<BsonValue> listIterator(final int index) {
+        ensureOpen();
         // Not the most efficient way to do this, but unlikely anyone will notice in practice
         return new ArrayList<>(this).listIterator(index);
     }
 
     @Override
     public List<BsonValue> subList(final int fromIndex, final int toIndex) {
+        ensureOpen();
         if (fromIndex < 0) {
             throw new IndexOutOfBoundsException("fromIndex = " + fromIndex);
         }
@@ -234,6 +265,7 @@ public boolean addAll(final Collection<? extends BsonValue> c) {
 
     @Override
     public boolean addAll(final int index, final Collection<? extends BsonValue> c) {
+        ensureOpen();
         throw new UnsupportedOperationException(READ_ONLY_MESSAGE);
     }
 
@@ -267,34 +299,97 @@ public BsonValue remove(final int index) {
         throw new UnsupportedOperationException(READ_ONLY_MESSAGE);
     }
 
+    @Override
+    public void close(){
+        if (!closed) {
+            for (Closeable closeable : trackedResources) {
+                try {
+                    closeable.close();
+                } catch (Exception e) {
+                    // Log and continue closing other resources
+                    LOGGER.error("Error closing resource", e);
+                }
+            }
+            trackedResources.clear();
+            closed = true;
+        }
+    }
+
+    private void ensureOpen() {
+        if (closed) {
+            throw new IllegalStateException("The BsonArray resources have been released.");
+        }
+    }
+
     private class ByteBufBsonArrayIterator implements Iterator<BsonValue> {
-        private final ByteBuf duplicatedByteBuf = byteBuf.duplicate();
-        private final BsonBinaryReader bsonReader;
+        private ByteBuf duplicatedByteBuf;
+        private BsonBinaryReader reader;
+        private Closeable resourceHandle;
+        private boolean finished;
 
         {
-            bsonReader = new BsonBinaryReader(new ByteBufferBsonInput(duplicatedByteBuf));
+            ensureOpen();
+            // Create duplicate buffer for iteration and track it temporarily
+            duplicatedByteBuf = byteBuf.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));
             // While one might expect that this would be a call to BsonReader#readStartArray that doesn't work because BsonBinaryReader
             // expects to be positioned at the start at the beginning of a document, not an array.  Fortunately, a BSON array has exactly
             // the same structure as a BSON document (the keys are just the array indices converted to a strings).  So it works fine to
             // call BsonReader#readStartDocument here, and just skip all the names via BsonReader#skipName below.
-            bsonReader.readStartDocument();
-            bsonReader.readBsonType();
+            reader.readStartDocument();
+            reader.readBsonType();
         }
 
         @Override
         public boolean hasNext() {
-            return bsonReader.getCurrentBsonType() != BsonType.END_OF_DOCUMENT;
+            if (finished) {
+                return false;
+            }
+            ensureOpen();
+            boolean hasNext = reader.getCurrentBsonType() != BsonType.END_OF_DOCUMENT;
+            if (!hasNext) {
+                cleanup();
+            }
+            return hasNext;
         }
 
         @Override
         public BsonValue next() {
             if (!hasNext()) {
                 throw new NoSuchElementException();
             }
-            bsonReader.skipName();
-            BsonValue value = readBsonValue(duplicatedByteBuf, bsonReader);
-            bsonReader.readBsonType();
+            reader.skipName();
+            BsonValue value = readBsonValue(duplicatedByteBuf, reader, trackedResources);
+            reader.readBsonType();
             return value;
         }
+
+        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
+                }
+            }
+        }
     }
 }