Fix object-log write buffer leaks in ObjectAllocator read-only and recovery flushes (#2046)

* Fix object-log write buffer leaks in ObjectAllocator read-only and recovery flushes

ObjectAllocator flushes rent pooled object-log write buffers via
CircularDiskWriteBuffer, which are returned to the pool only on Dispose(). The
read-only flush and snapshot-region recovery copy paths never disposed theirs,
leaking a buffer set per flush range / per recovered page.

- Dispose the shared buffers after the read-only flush loop.
- Dispose each recovery page's buffer after its WriteAsync.
- Skip renting the buffer for inline-only read-only pages (objectIdMap empty),
  routing them through WriteInlinePageAsync.
- Harden CircularDiskWriteBuffer.Dispose() memory ordering, since both paths now
  dispose while device writes may still be in flight.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Copilot-Session: f45128a9-8d17-4a2e-b84d-6da1afd3cc73

* Make flush-buffer disposal exception-safe

Wrap the read-only flush loop and the per-page recovery WriteAsync in
try/finally so the rented CircularDiskWriteBuffer is returned to the pool even
if a page write throws.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Copilot-Session: f45128a9-8d17-4a2e-b84d-6da1afd3cc73

---------

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Badrish Chandramouli <badrishc@microsoft.com>
Copilot-Session: f45128a9-8d17-4a2e-b84d-6da1afd3cc73

Author: 3 people

libs/storage/Tsavorite/cs/src/core/Allocator/AllocatorBase.cs

@@ -1948,6 +1948,10 @@ public void AsyncFlushPagesForRecovery<TContext>(long scanFromAddress, long flus
             {
                 var pageStartAddress = GetLogicalAddressOfStartOfPage(flushPage);
                 var flushFromAddress = Math.Max(scanFromAddress, pageStartAddress);
+
+                // When copying snapshot object bytes into the main object-log, rent this page's object-log write buffers; the hybrid-log region
+                // (which reuses the stored lengths/positions without writing object bytes) needs none.
+                var flushBuffers = copyObjects ? CreateCircularFlushBuffers(objectLogDevice: null, logger) : null;
                 var asyncResult = new PageAsyncFlushResult<TContext>()
                 {
                     page = flushPage,
@@ -1963,12 +1967,22 @@ public void AsyncFlushPagesForRecovery<TContext>(long scanFromAddress, long flus
                     flushRequestState = FlushRequestState.Recovery,
                     recoverySnapshotObjectLogDevice = snapshotObjectLogDevice,
                     recoveryFormerFlushedUntilAddress = formerFlushedUntilAddress,
-                    flushBuffers = copyObjects ? CreateCircularFlushBuffers(objectLogDevice: null, logger) : null
+                    flushBuffers = flushBuffers
                 };
 
-                // For the snapshot region (records at/above formerFlushedUntilAddress) we copy object bytes from the snapshot object-log to the main
-                // object-log using flushBuffers; otherwise (hybrid-log region) we reuse the stored lengths/positions without writing object bytes.
-                WriteAsync(flushPage, callback, asyncResult);
+                try
+                {
+                    // For the snapshot region (records at/above formerFlushedUntilAddress) we copy object bytes from the snapshot object-log to the main
+                    // object-log using flushBuffers; otherwise (hybrid-log region) we reuse the stored lengths/positions without writing object bytes.
+                    WriteAsync(flushPage, callback, asyncResult);
+                }
+                finally
+                {
+                    // WriteAsync issues all of this page's object-log and main-log device writes synchronously before returning, and the recovery completion
+                    // callback (AsyncFlushPageCallbackForRecovery) never reuses flushBuffers, so dispose it here, even if WriteAsync throws. Dispose defers the
+                    // actual return-to-pool (ClearBuffers) until any still-in-flight writes complete.
+                    flushBuffers?.Dispose();
+                }
             }
         }
 

libs/storage/Tsavorite/cs/src/core/Allocator/ObjectAllocatorImpl.cs

@@ -619,25 +619,37 @@ internal override void AsyncFlushPagesForReadOnly(long fromAddress, long untilAd
             // the current TailAddress is, but for normal flush operations we do set it to page alignment to eliminate concerns about rewriting partial sectors.
             GetFlushPageRange(fromAddress, untilAddress, out var startPage, out var numPages);
 
-            // Create the buffers we will use for all ranges of the flush. This calls our callback and disposes itself when the last write of a range completes.
+            // Create the buffers we will use for all ranges of the flush. Each page that has out-of-line data rents pooled object-log write
+            // buffers from this instance; pages that are entirely inline skip it (see WriteAsync). ObjectAllocator flushes are page-aligned and
+            // never use the PendingFlush chaining path, so once this loop has issued every page's write, no further writes will reference these
+            // buffers and it is safe to Dispose below.
             var flushBuffers = CreateCircularFlushBuffers(objectLogDevice: null, logger);
 
-            // Write each page (or partial page) in the range.
-            for (var flushPage = startPage; flushPage < (startPage + numPages); flushPage++)
+            try
             {
-                // The result from PrepareFlushAsyncResult indicates whether we are to perform an actual flush--but asyncResult will be set anyway.
-                if (PrepareFlushAsyncResult(fromAddress, untilAddress, noFlush, flushPage, out var asyncResult))
+                // Write each page (or partial page) in the range.
+                for (var flushPage = startPage; flushPage < (startPage + numPages); flushPage++)
                 {
-                    asyncResult.flushBuffers = flushBuffers;
+                    // The result from PrepareFlushAsyncResult indicates whether we are to perform an actual flush--but asyncResult will be set anyway.
+                    if (PrepareFlushAsyncResult(fromAddress, untilAddress, noFlush, flushPage, out var asyncResult))
+                    {
+                        asyncResult.flushBuffers = flushBuffers;
 
-                    // TsavoriteKV using ObjectAllocator always moves ReadOnlyAddress in page alignment, so if we have a partial first page, it can be written
-                    // in the same loop as full pages, because there are no adjacent fragments. Write the entire page up to asyncResult.untilAddress.
-                    Debug.Assert(PendingFlush[GetPageIndexForAddress(asyncResult.fromAddress)].list.Count == 0,
-                        $"Expected PendingFlush count {PendingFlush[GetPageIndexForAddress(asyncResult.fromAddress)].list.Count} to be 0 for ObjectAllocator");
+                        // TsavoriteKV using ObjectAllocator always moves ReadOnlyAddress in page alignment, so if we have a partial first page, it can be written
+                        // in the same loop as full pages, because there are no adjacent fragments. Write the entire page up to asyncResult.untilAddress.
+                        Debug.Assert(PendingFlush[GetPageIndexForAddress(asyncResult.fromAddress)].list.Count == 0,
+                            $"Expected PendingFlush count {PendingFlush[GetPageIndexForAddress(asyncResult.fromAddress)].list.Count} to be 0 for ObjectAllocator");
 
-                    WriteAsync(flushPage, AsyncFlushPageCallback, asyncResult);
+                        WriteAsync(flushPage, AsyncFlushPageCallback, asyncResult);
+                    }
                 }
             }
+            finally
+            {
+                // Dispose the shared flush buffers so their pooled object-log write buffers are returned to the pool, even if a page write throws;
+                // Dispose defers the actual return (ClearBuffers) until any still-in-flight device writes complete.
+                flushBuffers?.Dispose();
+            }
         }
 
         protected override void WriteAsync<TContext>(long flushPage, DeviceIOCompletionCallback callback, PageAsyncFlushResult<TContext> asyncResult)
@@ -708,8 +720,16 @@ private void WriteAsync<TContext>(long flushPage, ulong alignedMainLogFlushPageA
             if (isFirstRecordOnPage)
                 ((PageHeader*)logPagePointer)->SetLowestObjectLogPosition(objectLogTail);
 
-            // Short circuit if we are not using flushBuffers and not in recovery (e.g. using ObjectAllocator for string-only purposes).
-            if (asyncResult.flushBuffers is null)
+            // A ReadOnly flush of a page whose records are entirely inline (no Overflow keys/values and no Object values, i.e. the page's
+            // objectIdMap is empty) has nothing to serialize to the object log, so take the cheaper WriteInlinePageAsync path and skip renting
+            // an object-log write buffer. This is restricted to ReadOnly flushes: a Recovery flush does not populate objectIdMap (it reuses
+            // the on-disk lengths/positions), and a Snapshot flush may still need to invalidate v+1 records in the disk-image copy.
+            var objectIdMap = objectPages[flushPage % BufferSize].objectIdMap;
+            var pageHasNoObjectsToFlush = asyncResult.flushRequestState == FlushRequestState.ReadOnly && objectIdMap.Count == 0;
+
+            // Short circuit if we are not using flushBuffers and not in recovery (e.g. using ObjectAllocator for string-only purposes), or if a
+            // ReadOnly flush of this page has no out-of-line data to write to the object log.
+            if (asyncResult.flushBuffers is null || pageHasNoObjectsToFlush)
             {
                 if (asyncResult.flushRequestState != FlushRequestState.Recovery)
                 {
@@ -792,8 +812,6 @@ private void WriteAsync<TContext>(long flushPage, ulong alignedMainLogFlushPageA
                 // not change record sizes, so the logicalAddress space is unchanged. Also, we will not advance HeadAddress until this flush is complete
                 // and has updated FlushedUntilAddress, so we don't have to worry about the page being yanked out from underneath us (and Objects
                 // won't be disposed before we're done). TODO: Loop on successive subsets of the page's records to make this initial copy buffer smaller.
-                var objectIdMap = objectPages[flushPage % BufferSize].objectIdMap;
-
                 srcBuffer = bufferPool.Get(alignedBufferSize);
                 asyncResult.freeBuffer1 = srcBuffer;
 

libs/storage/Tsavorite/cs/src/core/Allocator/ObjectSerialization/CircularDiskWriteBuffer.cs

@@ -254,7 +254,13 @@ public void Dispose()
             // any further writes. For this class, "disposed" means "we're done issuing writes". And filePosition must be preserved; checkpoints will retrieve it later,
             // and chained partial flushes will append to it.
             disposed = true;
-            if (numInFlightWrites == 0)
+
+            // Full fence so the 'disposed' store above is globally visible before we read numInFlightWrites below (and, symmetrically, so this read is
+            // not hoisted above the store). FlushToDeviceCallback decrements numInFlightWrites (a full-fence Interlocked op) and then reads 'disposed';
+            // without this barrier the store and load could reorder and both sides could observe the pre-decrement / pre-dispose values, so neither would
+            // call ClearBuffers and the pooled buffers would leak. This matters now that the ReadOnly flush path Disposes while writes may be in flight.
+            Interlocked.MemoryBarrier();
+            if (Interlocked.Read(ref numInFlightWrites) == 0)
                 ClearBuffers();
         }