1616
1717package com .mongodb .internal .connection ;
1818
19+ import com .mongodb .annotations .Sealed ;
20+ import com .mongodb .internal .ResourceUtil ;
21+ import com .mongodb .internal .VisibleForTesting ;
1922import org .bson .BsonSerializationException ;
2023import org .bson .ByteBuf ;
2124import org .bson .io .OutputBuffer ;
2831
2932import static com .mongodb .assertions .Assertions .assertTrue ;
3033import static com .mongodb .assertions .Assertions .notNull ;
34+ import static com .mongodb .internal .VisibleForTesting .AccessModifier .PRIVATE ;
3135import static java .lang .String .format ;
3236
3337/**
38+ * A BSON output implementation that uses pooled {@link ByteBuf} instances for efficient memory management.
39+ *
40+ * <h2>ByteBuf Ownership and Lifecycle</h2>
41+ * <p>This class manages the lifecycle of {@link ByteBuf} instances obtained from the {@link BufferProvider}.
42+ * The ownership model is as follows:</p>
43+ * <ul>
44+ * <li>Internal buffers are owned by this output and released when {@link #close()} is called or
45+ * when {@link #truncateToPosition(int)} removes them.</li>
46+ * <li>Methods that return {@link ByteBuf} instances (e.g., {@link #getByteBuffers()}) return
47+ * duplicates with their own reference counts. <strong>Callers are responsible for releasing
48+ * these buffers</strong> to prevent memory leaks.</li>
49+ * <li>The {@link Branch} subclass merges its buffers into the parent on close, transferring
50+ * ownership by retaining buffers before the branch releases them.</li>
51+ * </ul>
52+ *
3453 * <p>This class is not part of the public API and may be removed or changed at any time</p>
3554 */
55+ @ Sealed
3656public class ByteBufferBsonOutput extends OutputBuffer {
3757
3858 private static final int MAX_SHIFT = 31 ;
@@ -50,6 +70,9 @@ public class ByteBufferBsonOutput extends OutputBuffer {
5070 /**
5171 * Construct an instance that uses the given buffer provider to allocate byte buffers as needs as it grows.
5272 *
73+ * <p>The buffer provider is used to allocate new {@link ByteBuf} instances as the output grows.
74+ * All allocated buffers are owned by this output and will be released when {@link #close()} is called.</p>
75+ *
5376 * @param bufferProvider the non-null buffer provider
5477 */
5578 public ByteBufferBsonOutput (final BufferProvider bufferProvider ) {
@@ -63,6 +86,10 @@ public ByteBufferBsonOutput(final BufferProvider bufferProvider) {
6386 * If multiple branches are created, they are merged in the order they are {@linkplain ByteBufferBsonOutput.Branch#close() closed}.
6487 * {@linkplain #close() Closing} this {@link ByteBufferBsonOutput} does not {@linkplain ByteBufferBsonOutput.Branch#close() close} the branch.
6588 *
89+ * <p><strong>ByteBuf Ownership:</strong> The branch allocates its own buffers. When the branch is closed,
90+ * ownership of these buffers is transferred to the parent by retaining them before the branch releases
91+ * its references. The parent then becomes responsible for releasing these buffers when it is closed.</p>
92+ *
6693 * @return A new {@link ByteBufferBsonOutput.Branch}.
6794 */
6895 public ByteBufferBsonOutput .Branch branch () {
@@ -223,17 +250,46 @@ protected void write(final int absolutePosition, final int value) {
223250 byteBuffer .put (bufferPositionPair .position ++, (byte ) value );
224251 }
225252
253+ /**
254+ * Returns a list of duplicated byte buffers containing the written data, flipped for reading.
255+ *
256+ * <p><strong>ByteBuf Ownership:</strong> The returned buffers are duplicates with their own
257+ * reference counts (each starts with a reference count of 1). <strong>The caller is responsible
258+ * for releasing each buffer</strong> when done to prevent memory leaks. Example usage:</p>
259+ * <pre>{@code
260+ * List<ByteBuf> buffers = output.getByteBuffers();
261+ * try {
262+ * // use buffers
263+ * } finally {
264+ * ResourceUtil.release(buffers);
265+ * }
266+ * }</pre>
267+ * <p><strong>Note:</strong> These buffers must be released before this {@code ByteBufferBsonOutput} is closed.
268+ * Otherwise there is a risk of the buffers being released back to the bufferProvider and data corruption.</p>
269+ *
270+ * @return a list of duplicated buffers, flipped for reading
271+ */
226272 @ Override
227273 public List <ByteBuf > getByteBuffers () {
228274 ensureOpen ();
229-
230275 List <ByteBuf > buffers = new ArrayList <>(bufferList .size ());
231276 for (final ByteBuf cur : bufferList ) {
232277 buffers .add (cur .duplicate ().order (ByteOrder .LITTLE_ENDIAN ).flip ());
233278 }
234279 return buffers ;
235280 }
236281
282+ /**
283+ * Returns a list of duplicated byte buffers without flipping them.
284+ *
285+ * <p><strong>ByteBuf Ownership:</strong> The returned buffers are duplicates with their own
286+ * reference counts (each starts with a reference count of 1). <strong>The caller is responsible
287+ * for releasing each buffer</strong> when done to prevent memory leaks.</p>
288+ *
289+ * @return a list of duplicated buffers
290+ * @see #getByteBuffers()
291+ */
292+ @ VisibleForTesting (otherwise = PRIVATE )
237293 public List <ByteBuf > getDuplicateByteBuffers () {
238294 ensureOpen ();
239295
@@ -245,6 +301,13 @@ public List<ByteBuf> getDuplicateByteBuffers() {
245301 }
246302
247303
304+ /**
305+ * {@inheritDoc}
306+ *
307+ * <p><strong>ByteBuf Management:</strong> This method obtains duplicated buffers via
308+ * {@link #getByteBuffers()} and releases them after writing to the output stream,
309+ * ensuring no buffer leaks occur.</p>
310+ */
248311 @ Override
249312 public int pipe (final OutputStream out ) throws IOException {
250313 ensureOpen ();
@@ -263,11 +326,20 @@ public int pipe(final OutputStream out) throws IOException {
263326 total += cur .limit ();
264327 }
265328 } finally {
266- byteBuffers . forEach ( ByteBuf :: release );
329+ ResourceUtil . release ( byteBuffers );
267330 }
268331 return total ;
269332 }
270333
334+ /**
335+ * Truncates this output to the specified position, releasing any buffers that are no longer needed.
336+ *
337+ * <p><strong>ByteBuf Management:</strong> Any buffers beyond the new position are removed from
338+ * the internal buffer list and released. This ensures no memory leaks when truncating.</p>
339+ *
340+ * @param newPosition the new position to truncate to
341+ * @throws IllegalArgumentException if newPosition is negative or greater than the current position
342+ */
271343 @ Override
272344 public void truncateToPosition (final int newPosition ) {
273345 ensureOpen ();
@@ -306,13 +378,15 @@ public final void flush() throws IOException {
306378 * {@inheritDoc}
307379 * <p>
308380 * Idempotent.</p>
381+ *
382+ * <p><strong>ByteBuf Management:</strong> Releases internal buffers and clears the buffer list.
383+ * After this method returns, all buffers that were allocated by this output will have been fully released
384+ * back to the buffer provider.</p>
309385 */
310386 @ Override
311387 public void close () {
312388 if (isOpen ()) {
313- for (final ByteBuf cur : bufferList ) {
314- cur .release ();
315- }
389+ ResourceUtil .release (bufferList );
316390 currentByteBuffer = null ;
317391 bufferList .clear ();
318392 closed = true ;
@@ -345,7 +419,14 @@ boolean isOpen() {
345419 }
346420
347421 /**
348- * @see #branch()
422+ * Merges a branch's buffers into this output.
423+ *
424+ * <p><strong>ByteBuf Ownership:</strong> This method retains each buffer from the branch before
425+ * adding it to this output's buffer list. This is necessary because the branch will release its
426+ * references when it closes. The retain ensures the buffers remain valid and are now owned by
427+ * this output.</p>
428+ *
429+ * @param branch the branch to merge
349430 */
350431 private void merge (final ByteBufferBsonOutput branch ) {
351432 assertTrue (branch instanceof ByteBufferBsonOutput .Branch );
@@ -356,6 +437,20 @@ private void merge(final ByteBufferBsonOutput branch) {
356437 currentByteBuffer = null ;
357438 }
358439
440+ /**
441+ * A branch of a {@link ByteBufferBsonOutput} that can be merged back into its parent.
442+ *
443+ * <p><strong>ByteBuf Ownership:</strong> A branch allocates its own buffers independently.
444+ * When {@link #close()} is called:</p>
445+ * <ol>
446+ * <li>The parent's {@link ByteBufferBsonOutput#merge(ByteBufferBsonOutput)} method is called,
447+ * which retains all buffers in this branch.</li>
448+ * <li>Then {@code super.close()} is called, which releases the branch's references to the buffers.</li>
449+ * </ol>
450+ * <p>The retain/release sequence ensures buffers are safely transferred to the parent without leaks.</p>
451+ *
452+ * @see #branch()
453+ */
359454 public static final class Branch extends ByteBufferBsonOutput {
360455 private final ByteBufferBsonOutput parent ;
361456
@@ -365,6 +460,16 @@ private Branch(final ByteBufferBsonOutput parent) {
365460 }
366461
367462 /**
463+ * Closes this branch and merges its data into the parent output.
464+ *
465+ * <p><strong>ByteBuf Ownership:</strong> On close, this branch's buffers are transferred
466+ * to the parent. The parent retains the buffers (incrementing reference counts), and then
467+ * this branch releases only its own single reference. The parent
468+ * becomes the sole owner of the buffers and is responsible for releasing them.</p>
469+ *
470+ * <p>Idempotent. If already closed, this method does nothing.</p>
471+ *
472+ * @throws AssertionError if the parent has been closed before this branch
368473 * @see #branch()
369474 */
370475 @ Override
0 commit comments