Fix sorted-set Memory leaks and BITOP epoch tracking (#1752)

## Summary

Two correctness bugs in the sorted-set object store and the BITOP read
loop.

## 1. Sorted-set `IMemoryOwner` leaks in `GEO*STORE` / `ZUNIONSTORE` / `ZINTERSTORE`

Three pooled-buffer leaks where backends (`GeoSearch`, `SortedSetRange`,
internal `ZADD`'s `SortedSetAdd`) wrote replies via `RespMemoryWriter`,
which (with a default `SpanByte`) rents a `MemoryPool<byte>` buffer
(≥512 bytes) and assigns it to `.SpanByteAndMemory.Memory`:

- **`SortedSetGeoOps.cs` (`GEO*STORE`)** — the `searchOutMem.Memory`
  from `GeoSearch` was leaked (only `searchOutHandler` — the
  `MemoryHandle` from `Pin()` — was disposed). The internal `ZADD`
  invocation (`zAddOutput`) was discarded entirely.
- **`SortedSetOps.cs` (`ZUNIONSTORE` / `ZINTERSTORE`)** — same pattern:
  `rangeOutputMem.Memory` from `SortedSetRange` was leaked, and the
  internal `ZADD`'s `zAddOutput.SpanByteAndMemory.Memory` was leaked.

Under heavy `GEO*STORE` / `ZUNIONSTORE` / `ZINTERSTORE` traffic this was
real `MemoryPool` churn and GC pressure.

### Fix

- For each `*STORE`-style internal `ZADD`, wrap the
  `RMWObjectStoreOperationWithOutput` call in a `try` / `finally` that
  disposes `zAddOutput.SpanByteAndMemory.Memory` if `!IsSpanByte`.
- Extend the existing `finally` blocks that dispose `*Handler` (the
  `MemoryHandle`) to also dispose the underlying `*Mem.Memory`
  (the `IMemoryOwner<byte>`).

## 2. BITOP: pending-completion epoch tracking was broken

`StorageSession.HeadAddress` was a `readonly long` field captured at
session-construction time and never updated.
`MainStoreOps.ReadWithUnsafeContext` compared it against itself
(`HeadAddress == localHeadAddress`) to decide whether to set
`epochChanged = true` after pending completion. Two bugs:

1. The field is frozen, so the check was meaningless — the live store
   `HeadAddress` was never consulted.
2. The condition was also **inverted**: it set `epochChanged = true`
   when the addresses were equal (i.e., head did NOT move), the
   opposite of what the comment said.

In addition, `Read` can return synchronously with a pointer into the
**read cache** (a separate log with its own `HeadAddress` that can be
evicted independently of the main log). The original check would not
detect read-cache eviction.

### Fix

- Removed the stale `StorageSession.HeadAddress` field.
- Added `ClientSession.HeadAddress` and `ClientSession.ReadCacheHeadAddress`
  accessors that read the live values from `store.Log.HeadAddress` /
  `store.ReadCache?.HeadAddress`.
- `ReadWithUnsafeContext` now captures both addresses at the start of
  the BITOP loop and, after pending completion, sets `epochChanged = true`
  if **either** has advanced — correctly invalidating any pointers
  captured into either log.

## Files changed

- `libs/storage/Tsavorite/cs/src/core/ClientSession/ClientSession.cs` —
  new `HeadAddress` and `ReadCacheHeadAddress` live accessors
- `libs/server/Storage/Session/StorageSession.cs` — removed stale
  `HeadAddress` field
- `libs/server/Storage/Session/MainStore/MainStoreOps.cs` —
  `ReadWithUnsafeContext` uses live addresses with the corrected
  comparison; signature now also takes `localReadCacheHeadAddress`
- `libs/server/Storage/Session/MainStore/BitmapOps.cs` — captures both
  live head addresses; passes them to `ReadWithUnsafeContext`
- `libs/server/Storage/Session/ObjectStore/SortedSetGeoOps.cs` — leak
  fixes (`searchOutMem` + `zAddOutput`)
- `libs/server/Storage/Session/ObjectStore/SortedSetOps.cs` — leak
  fixes (`rangeOutputMem` + `zAddOutput`)

## Validation

- All 660 sorted-set + geo + bitmap tests pass on `main`
- `dotnet format --verify-no-changes` clean

## Note on related fixes on `dev`

This is a subset of the fixes on the companion `dev` branch
(`badrishc/memory-fixes`). Other fixes from that branch were
intentionally **not** ported here because they do not apply to `main`:

- The BITOP overflow-pointer fix relies on the `ISourceLogRecord` /
  `LogRecord` / `OverflowByteArray` model that exists only on `dev`. On
  `main` the BITOP backend operates on `SpanByte` values that are always
  pinned in log memory, so the use-after-fixed bug fixed on `dev` does
  not exist on `main`.
- The PFCOUNT/PFMERGE bounds-check tightening from `dev` is unnecessary
  on `main` because the `main` backend already validates
  `value.Length <= dst.Length` *before* the `Buffer.MemoryCopy`, so the
  wrong-capacity argument is gated.

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

Author: badrishc

libs/server/Storage/Session/MainStore/BitmapOps.cs

@@ -105,15 +105,16 @@ public unsafe GarnetStatus StringBitOperation(ref RawStringInput input, BitmapOp
             {
                 uc.BeginUnsafe();
             readFromScratch:
-                var localHeadAddress = HeadAddress;
+                var localHeadAddress = uc.Session.HeadAddress;
+                var localReadCacheHeadAddress = uc.Session.ReadCacheHeadAddress;
                 var keysFound = 0;
 
                 for (var i = 1; i < keys.Length; i++)
                 {
                     var srcKey = keys[i];
                     //Read srcKey
                     var outputBitmap = SpanByteAndMemory.FromPinnedSpan(output);
-                    status = ReadWithUnsafeContext(srcKey, ref input, ref outputBitmap, localHeadAddress, out bool epochChanged, ref uc);
+                    status = ReadWithUnsafeContext(srcKey, ref input, ref outputBitmap, localHeadAddress, localReadCacheHeadAddress, out bool epochChanged, ref uc);
                     if (epochChanged)
                     {
                         goto readFromScratch;

libs/server/Storage/Session/MainStore/MainStoreOps.cs

@@ -47,7 +47,7 @@ public GarnetStatus GET<TContext>(ref SpanByte key, ref RawStringInput input, re
             }
         }
 
-        public unsafe GarnetStatus ReadWithUnsafeContext<TContext>(ArgSlice key, ref RawStringInput input, ref SpanByteAndMemory output, long localHeadAddress, out bool epochChanged, ref TContext context)
+        public unsafe GarnetStatus ReadWithUnsafeContext<TContext>(ArgSlice key, ref RawStringInput input, ref SpanByteAndMemory output, long localHeadAddress, long localReadCacheHeadAddress, out bool epochChanged, ref TContext context)
             where TContext : ITsavoriteContext<SpanByte, SpanByte, RawStringInput, SpanByteAndMemory, long, MainSessionFunctions, MainStoreFunctions, MainStoreAllocator>, IUnsafeContext
         {
             var _key = key.SpanByte;
@@ -63,8 +63,12 @@ public unsafe GarnetStatus ReadWithUnsafeContext<TContext>(ArgSlice key, ref Raw
                 CompletePendingForSession(ref status, ref output, ref context);
                 StopPendingMetrics();
                 context.BeginUnsafe();
-                // Start read of pointers from beginning if epoch changed
-                if (HeadAddress == localHeadAddress)
+                // If either the main-log head or the read-cache head advanced while we waited on
+                // pending I/O, any log/read-cache pointers captured by previous reads in this loop
+                // may now reference evicted pages. Tell the caller to re-read all sources from
+                // scratch. (Synchronously-returned reads can have pointers into either log.)
+                if (context.Session.HeadAddress != localHeadAddress
+                    || context.Session.ReadCacheHeadAddress != localReadCacheHeadAddress)
                 {
                     context.EndUnsafe();
                     epochChanged = true;

libs/server/Storage/Session/ObjectStore/SortedSetGeoOps.cs

@@ -207,13 +207,27 @@ public unsafe GarnetStatus GeoSearchStore<TObjectContext>(ArgSlice key, ArgSlice
                     }, ref parseState);
 
                     var zAddOutput = new GarnetObjectStoreOutput();
-                    RMWObjectStoreOperationWithOutput(destinationKey, ref zAddInput, ref objectStoreLockableContext, ref zAddOutput);
+                    try
+                    {
+                        RMWObjectStoreOperationWithOutput(destinationKey, ref zAddInput, ref objectStoreLockableContext, ref zAddOutput);
 
-                    writer.WriteInt32(foundItems);
+                        writer.WriteInt32(foundItems);
+                    }
+                    finally
+                    {
+                        // ZADD backend writes its result via RespMemoryWriter, which allocates a
+                        // MemoryPool buffer when the (default) SpanByte cannot hold the response.
+                        // Dispose to avoid leaking that buffer back to the pool.
+                        if (!zAddOutput.SpanByteAndMemory.IsSpanByte)
+                            zAddOutput.SpanByteAndMemory.Memory?.Dispose();
+                    }
                 }
                 finally
                 {
                     searchOutHandler.Dispose();
+                    // GeoSearch writes via RespMemoryWriter, which (with a default SpanByte) rents a
+                    // MemoryPool buffer and assigns it here. Dispose to release it back to the pool.
+                    searchOutMem.Memory?.Dispose();
                 }
 
                 return GarnetStatus.OK;

libs/server/Storage/Session/ObjectStore/SortedSetOps.cs

@@ -785,13 +785,27 @@ public unsafe GarnetStatus SortedSetRangeStore<TObjectContext>(ArgSlice dstKey,
                         }, ref parseState);
 
                         var zAddOutput = new GarnetObjectStoreOutput();
-                        RMWObjectStoreOperationWithOutput(destinationKey, ref zAddInput, ref objectStoreLockableContext, ref zAddOutput);
-                        itemBroker.HandleCollectionUpdate(destinationKey);
+                        try
+                        {
+                            RMWObjectStoreOperationWithOutput(destinationKey, ref zAddInput, ref objectStoreLockableContext, ref zAddOutput);
+                            itemBroker.HandleCollectionUpdate(destinationKey);
+                        }
+                        finally
+                        {
+                            // ZADD backend writes its result via RespMemoryWriter, which allocates a
+                            // MemoryPool buffer when the (default) SpanByte cannot hold the response.
+                            // Dispose to avoid leaking that buffer back to the pool.
+                            if (!zAddOutput.SpanByteAndMemory.IsSpanByte)
+                                zAddOutput.SpanByteAndMemory.Memory?.Dispose();
+                        }
                     }
                 }
                 finally
                 {
                     rangeOutputHandler.Dispose();
+                    // SortedSetRange writes via RespMemoryWriter, which (with a default SpanByte) rents
+                    // a MemoryPool buffer and assigns it here. Dispose to release it back to the pool.
+                    rangeOutputMem.Memory?.Dispose();
                 }
                 return status;
             }

libs/server/Storage/Session/StorageSession.cs

@@ -22,7 +22,6 @@ sealed partial class StorageSession : IDisposable
     {
         int bitmapBufferSize = 1 << 15;
         SectorAlignedMemory sectorAlignedMemoryBitmap;
-        readonly long HeadAddress;
 
         /// <summary>
         /// Session Contexts for main store
@@ -107,7 +106,6 @@ public StorageSession(StoreWrapper storeWrapper,
             vectorContext = vectorSession.BasicContext;
             vectorLockableContext = vectorSession.LockableContext;
 
-            HeadAddress = db.MainStore.Log.HeadAddress;
             ObjectScanCountLimit = storeWrapper.serverOptions.ObjectScanCountLimit;
         }
 

libs/storage/Tsavorite/cs/src/core/ClientSession/ClientSession.cs

@@ -108,6 +108,23 @@ internal ClientSession(
         /// </summary>
         public long Version => ctx.version;
 
+        /// <summary>
+        /// The current head address of the underlying store's main log. Reads the live value, so it
+        /// reflects any updates from log eviction or page advancement. Callers that compare snapshots
+        /// (e.g. before vs. after a pending I/O completion) should hold epoch protection so that the
+        /// addresses they read remain meaningful.
+        /// </summary>
+        public long HeadAddress => store.Log.HeadAddress;
+
+        /// <summary>
+        /// The current head address of the underlying store's read cache, or 0 if the read cache is
+        /// not configured. Reads the live value. Callers that capture pointers into records returned
+        /// by <c>Read</c> must check this in addition to <see cref="HeadAddress"/>: a synchronously
+        /// returned pointer can live in the read cache (when the record was cached on a prior disk
+        /// read), and that page can be evicted independently of the main log.
+        /// </summary>
+        public long ReadCacheHeadAddress => store.ReadCache?.HeadAddress ?? 0;
+
         /// <summary>
         /// Dispose session
         /// </summary>