Drain deferred DoReadPage callbacks before disposing scan iterator

BufferAndLoad defers page reads via BumpCurrentEpoch(() => DoReadPage(...)).
These drain callbacks capture 'this' and access instance state (frame memory,
loadCompletionEvents, etc.). For read-ahead pages, the callback may still be
in the epoch drain list when the scan completes and the iterator is disposed.
The callback then accesses freed native frame memory (AccessViolationException)
or null arrays (NullReferenceException), and the resulting exception from
within a drain callback leaves the epoch in a held state, causing cascading
'Trying to acquire protected epoch' assertion failures.

Fix: add a drain phase at the start of Dispose() that detects pending deferred
reads (nextLoadedPages[i] > loadedPages[i]) and drains them by acquiring the
epoch and calling ProtectAndDrain until the callback executes. This ensures
loadCompletionEvents[i] is set, so the existing Phase 2 wait can then handle
async I/O completion. Only after both phases is it safe to free the frame.

Also reverts the insufficient null-check guard in DoReadPage from the prior
fix, since the drain-before-dispose approach makes it unnecessary.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: badrishc

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

@@ -208,9 +208,10 @@ protected bool BufferAndLoad(long currentIterationAddress, long currentPage, lon
 
                         void DoReadPage(int frameIndex)
                         {
-                            // The drain callback may execute after the iterator has been disposed (loadCompletionEvents set to null),
-                            // because the callback is deferred via BumpCurrentEpoch and only runs when SafeToReclaimEpoch advances.
-                            // This can happen for read-ahead pages (frameIndex > 0) when the scan completes before the callback runs.
+                            // Guard against the callback running after the iterator has been disposed.
+                            // Dispose drains pending callbacks (Phase 1), but there is a small race window
+                            // where BumpCurrentEpoch on another thread executes this callback before our
+                            // drain loop catches it. In that case, loadCompletionEvents is null.
                             var events = loadCompletionEvents;
                             if (events is null)
                                 return;
@@ -322,6 +323,43 @@ private bool WaitForFrameLoad(long currentAddress, long currentFrame)
         /// </summary>
         public virtual void Dispose()
         {
+            // BufferAndLoad uses BumpCurrentEpoch(() => DoReadPage(...)) to defer page reads as epoch
+            // drain callbacks. These callbacks capture 'this' (accessing loadCompletionEvents, frame
+            // memory, etc.). If the scan completes before a deferred callback executes — common for
+            // read-ahead pages — the callback would access disposed state: NRE on null arrays or AV
+            // on freed native frame memory.
+            //
+            // We detect pending callbacks by comparing nextLoadedPages[i] (set by CAS in BufferAndLoad
+            // when the read is *registered*) vs loadedPages[i] (set at the end of DoReadPage when the
+            // callback *executes* and issues the async I/O). If nextLoadedPages is ahead, the callback
+            // hasn't run yet.
+            //
+            // Phase 1 below drains pending callbacks so that loadCompletionEvents[i] gets set.
+            // Phase 2 (existing loop) then waits on loadCompletionEvents[i] for the async I/O to finish.
+            // Only after both phases is it safe to free the frame and other resources.
+
+            // Phase 1: Drain any pending deferred DoReadPage callbacks.
+            if (epoch != null && nextLoadedPages != null)
+            {
+                for (var i = 0; i < frameSize; i++)
+                {
+                    if (nextLoadedPages[i] > loadedPages[i])
+                    {
+                        epoch.Resume();
+                        try
+                        {
+                            while (Volatile.Read(ref loadedPages[i]) < nextLoadedPages[i])
+                                epoch.ProtectAndDrain();
+                        }
+                        finally
+                        {
+                            epoch.Suspend();
+                        }
+                    }
+                }
+            }
+
+            // Phase 2: Wait for async I/O completion and dispose resources.
             for (var i = 0; i < frameSize; i++)
             {
                 // Wait for ongoing reads to complete/fail; if the wait throws (e.g. due to cancellation), we still