refactor: move all destruction-side heap accounting into Tsavorite

Tsavorite now handles ALL heap-size decrements internally via logSizeTracker:
- OnDispose(Deleted): decrements value heap before ClearHeapFields
- EvictRecordsInRange: decrements key overflow for ALL records (including
  tombstoned), decrements value heap for non-tombstoned records
- CopyUpdated: decrements value-object heap (already internalized earlier)

GarnetRecordTriggers.OnDispose and OnEvict are now no-ops for accounting.
CallOnEvict returns false. The trigger callbacks remain available for
app-level resource cleanup (e.g. IDisposable.Dispose on value objects
that hold external resources), but Garnet does not use them.

The heap-tracking contract is now cleanly split:
- Session functions: emit +value at creation sites (only the app knows
  the heap size of a newly created value)
- Tsavorite: emit -key and -value at all destruction/eviction sites
  (the record is in hand, Tsavorite can read the sizes directly)

EvictRecordsInRange is now called whenever logSizeTracker is set OR
CallOnEvict is true, ensuring internal accounting runs even when the
application opts out of the OnEvict trigger.

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

Author: badrishc

libs/server/Storage/Functions/GarnetRecordTriggers.cs

@@ -31,42 +31,22 @@ public GarnetRecordTriggers(CacheSizeTracker cacheSizeTracker)
         public bool CallOnFlush => false;
 
         /// <inheritdoc/>
-        // Drives per-record heap-size decrement on page eviction. Mirrors the work the
-        // legacy SubscribeEvictions → LogSizeTracker.OnNext observer path used to perform
-        // (see CacheSizeTracker.Initialize for the wiring change). Gated per source so
-        // that enabling only a main-log or only a read-cache memory budget does not
-        // force a per-record eviction walk on the other allocator.
-        public bool CallOnEvict(EvictionSource source)
-        {
-            if (cacheSizeTracker is null)
-                return false;
-            return source == EvictionSource.ReadCache
-                ? cacheSizeTracker.readCacheTracker is not null
-                : cacheSizeTracker.mainLogTracker is not null;
-        }
+        // Heap-size accounting on eviction is now handled internally by Tsavorite's
+        // EvictRecordsInRange via logSizeTracker. The OnEvict trigger is only needed
+        // for app-level resource cleanup, which Garnet does not require.
+        public bool CallOnEvict(EvictionSource source) => false;
 
         /// <inheritdoc/>
         public bool CallOnDiskRead => false;
 
         /// <inheritdoc/>
         public void OnDispose(ref LogRecord logRecord, DisposeReason reason)
         {
-            if (cacheSizeTracker is null)
-                return;
-
-            if (reason == DisposeReason.Deleted)
-            {
-                // Decrement only the VALUE-side heap (value overflow bytes or value-object heap).
-                // Key overflow is NOT subtracted here because ClearHeapFields(clearKey=false) keeps the
-                // key alive for hash-chain traversal. Key overflow is handled separately:
-                //  - Mutable delete: record is tombstoned after OnDispose → OnEvict skips it → key overflow
-                //    is a bounded phantom (freed by GC when the page is freed, same as pre-IRecordTriggers).
-                //  - Immutable delete: source is sealed (not tombstoned) → OnEvict visits it and subtracts
-                //    the remaining key overflow via CalculateHeapMemorySize.
-                var size = logRecord.GetValueHeapMemorySize();
-                if (size != 0)
-                    cacheSizeTracker.AddHeapSize(-size);
-            }
+            // Heap-size accounting for destruction is handled internally by Tsavorite
+            // (logSizeTracker). This trigger is available for app-level resource cleanup
+            // (e.g. calling IDisposable.Dispose on value objects that hold external resources).
+            // Garnet's IHeapObject implementations (Hash/List/Set/SortedSet) hold no external
+            // resources, so this is a no-op.
         }
 
         /// <inheritdoc/>
@@ -81,22 +61,10 @@ public void OnDisposeDiskRecord(ref DiskLogRecord logRecord, DisposeReason reaso
         /// <inheritdoc/>
         public void OnEvict(ref LogRecord logRecord, EvictionSource source)
         {
-            if (cacheSizeTracker is null)
-                return;
-
-            // Decrement heap size by this record's heap contribution. Uses the same sizing
-            // helper that LogSizeTracker.OnNext used to sum over an evicted iterator. Routes
-            // through the standard AddHeapSize/AddReadCacheHeapSize path so the assertion
-            // guarding against negative totals remains in force; creation sites on the main
-            // log (RMW PostInitialUpdater/PostCopyUpdater and in-place grow/shrink) must emit
-            // a matching positive bump so the account stays balanced.
-            var size = MemoryUtils.CalculateHeapMemorySize(in logRecord);
-            if (size == 0)
-                return;
-            if (source == EvictionSource.ReadCache)
-                cacheSizeTracker.AddReadCacheHeapSize(-size);
-            else
-                cacheSizeTracker.AddHeapSize(-size);
+            // Heap-size accounting for eviction is handled internally by Tsavorite
+            // (logSizeTracker decrements both key overflow and value heap in EvictRecordsInRange).
+            // This trigger is available for app-level resource cleanup on eviction.
+            // Garnet has no app-level eviction cleanup, so this is a no-op.
         }
     }
 }

libs/server/Storage/Functions/MainStore/RMWMethods.cs

@@ -399,9 +399,9 @@ public readonly void PostInitialUpdater(ref LogRecord logRecord, in RecordSizeIn
                 rmwInfo.UserData |= NeedAofLog; // Mark that we need to write to AOF
             }
 
-            // Account for any heap contribution (key/value overflow) of the newly created record;
-            // balances the decrement emitted by GarnetRecordTriggers.OnEvict on page eviction.
-            var heap = logRecord.CalculateHeapMemorySize();
+            // Account for value-side overflow heap of the newly created record.
+            // Key overflow is tracked internally by Tsavorite via logSizeTracker.
+            var heap = logRecord.GetValueHeapMemorySize();
             if (heap != 0)
                 functionsState.cacheSizeTracker?.AddHeapSize(heap);
         }
@@ -1712,15 +1712,9 @@ public readonly bool PostCopyUpdater<TSourceLogRecord>(in TSourceLogRecord srcLo
             if (functionsState.appendOnlyFile != null)
                 rmwInfo.UserData |= NeedAofLog; // Mark that we need to write to AOF
 
-            // Account for the new record's heap contribution so that the matching OnEvict decrement
-            // balances. We do NOT subtract srcLogRecord's heap: the source may be a main-log record
-            // (whose heap is tracked in mainLogTracker and will leak upward only by a bounded sealed-
-            // source amount — parity with ObjectStore's ClearSource=false branch and the legacy
-            // SubscribeEvictions path), a read-cache record (heap lives in readCacheTracker, unrelated
-            // to this write), or a pending-IO DiskLogRecord (heap never counted in any tracker).
-            // Subtracting unconditionally would undercount in the last two cases and could drive the
-            // counter negative.
-            var newHeap = dstLogRecord.CalculateHeapMemorySize();
+            // Account for the new record's value-side heap contribution.
+            // Key overflow is tracked internally by Tsavorite via logSizeTracker.
+            var newHeap = dstLogRecord.GetValueHeapMemorySize();
             if (newHeap != 0)
                 functionsState.cacheSizeTracker?.AddHeapSize(newHeap);
             return true;

libs/server/Storage/Functions/MainStore/UpsertMethods.cs

@@ -46,9 +46,9 @@ public void PostInitialWriter(ref LogRecord logRecord, in RecordSizeInfo sizeInf
             if (functionsState.appendOnlyFile != null)
                 upsertInfo.UserData |= NeedAofLog; // Mark that we need to write to AOF
 
-            // Account for overflow key/value heap on the freshly-inserted record so the matching
-            // OnEvict decrement balances.
-            var heap = logRecord.CalculateHeapMemorySize();
+            // Account for value-side overflow heap on the freshly-inserted record.
+            // Key overflow is tracked internally by Tsavorite via logSizeTracker.
+            var heap = logRecord.GetValueHeapMemorySize();
             if (heap != 0)
                 functionsState.cacheSizeTracker?.AddHeapSize(heap);
         }
@@ -68,7 +68,7 @@ public void PostInitialWriter<TSourceLogRecord>(ref LogRecord logRecord, in Reco
                 upsertInfo.UserData |= NeedAofLog; // Mark that we need to write to AOF
             }
 
-            var heap = logRecord.CalculateHeapMemorySize();
+            var heap = logRecord.GetValueHeapMemorySize();
             if (heap != 0)
                 functionsState.cacheSizeTracker?.AddHeapSize(heap);
         }

libs/server/Storage/Functions/ObjectStore/UpsertMethods.cs

@@ -50,7 +50,7 @@ public void PostInitialWriter(ref LogRecord logRecord, in RecordSizeInfo sizeInf
                 upsertInfo.UserData |= NeedAofLog; // Mark that we need to write to AOF
 
             if (logRecord.Info.RecordHasObjects)
-                functionsState.cacheSizeTracker?.AddHeapSize(MemoryUtils.CalculateHeapMemorySize(in logRecord));
+                functionsState.cacheSizeTracker?.AddHeapSize(logRecord.GetValueHeapMemorySize());
         }
 
         /// <inheritdoc />
@@ -62,7 +62,7 @@ public void PostInitialWriter(ref LogRecord logRecord, in RecordSizeInfo sizeInf
                 upsertInfo.UserData |= NeedAofLog; // Mark that we need to write to AOF
 
             if (logRecord.Info.RecordHasObjects)
-                functionsState.cacheSizeTracker?.AddHeapSize(MemoryUtils.CalculateHeapMemorySize(in logRecord));
+                functionsState.cacheSizeTracker?.AddHeapSize(logRecord.GetValueHeapMemorySize());
         }
 
         /// <inheritdoc />
@@ -75,7 +75,7 @@ public void PostInitialWriter<TSourceLogRecord>(ref LogRecord logRecord, in Reco
                 upsertInfo.UserData |= NeedAofLog; // Mark that we need to write to AOF
             }
             if (logRecord.Info.RecordHasObjects)
-                functionsState.cacheSizeTracker?.AddHeapSize(MemoryUtils.CalculateHeapMemorySize(in logRecord));
+                functionsState.cacheSizeTracker?.AddHeapSize(logRecord.GetValueHeapMemorySize());
         }
 
         /// <inheritdoc />

libs/server/Storage/Functions/UnifiedStore/UpsertMethods.cs

@@ -66,7 +66,7 @@ public void PostInitialWriter(ref LogRecord logRecord, in RecordSizeInfo sizeInf
                 upsertInfo.UserData |= NeedAofLog; // Mark that we need to write to AOF
 
             if (logRecord.Info.RecordHasObjects)
-                functionsState.cacheSizeTracker?.AddHeapSize(MemoryUtils.CalculateHeapMemorySize(in logRecord));
+                functionsState.cacheSizeTracker?.AddHeapSize(logRecord.GetValueHeapMemorySize());
         }
 
         /// <inheritdoc />
@@ -79,7 +79,7 @@ public void PostInitialWriter(ref LogRecord logRecord, in RecordSizeInfo sizeInf
                 upsertInfo.UserData |= NeedAofLog; // Mark that we need to write to AOF
 
             if (logRecord.Info.RecordHasObjects)
-                functionsState.cacheSizeTracker?.AddHeapSize(MemoryUtils.CalculateHeapMemorySize(in logRecord));
+                functionsState.cacheSizeTracker?.AddHeapSize(logRecord.GetValueHeapMemorySize());
         }
 
         /// <inheritdoc />
@@ -92,7 +92,7 @@ public void PostInitialWriter<TSourceLogRecord>(ref LogRecord logRecord, in Reco
                 upsertInfo.UserData |= NeedAofLog; // Mark that we need to write to AOF
 
             if (logRecord.Info.RecordHasObjects)
-                functionsState.cacheSizeTracker?.AddHeapSize(MemoryUtils.CalculateHeapMemorySize(in logRecord));
+                functionsState.cacheSizeTracker?.AddHeapSize(logRecord.GetValueHeapMemorySize());
         }
 
         /// <inheritdoc />

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

@@ -1419,21 +1419,19 @@ internal void EvictPageForRecovery(long page)
             var end = GetLogicalAddressOfStartOfPage(page + 1);
 
             var source = IsReadCache ? EvictionSource.ReadCache : EvictionSource.MainLog;
-            if (storeFunctions.CallOnEvict(source))
+
+            // Per-record eviction walk handles internal heap accounting (key + value via
+            // logSizeTracker) and optionally notifies the application via OnEvict.
+            if (logSizeTracker is not null || storeFunctions.CallOnEvict(source))
             {
-                // New per-record path: OnEvict handles heap-size decrement per record. Parity
-                // with the runtime eviction path in OnPagesClosedWorker.
                 _wrapper.EvictRecordsInRange(start, end, source);
             }
-            else if (logSizeTracker is not null)
+            else if (onEvictionObserver is not null)
             {
-                // Legacy observer path: materialize an iterator and push heap-size decrement to
-                // the LogSizeTracker.OnNext observer (kept for consumers still using SubscribeEvictions).
-                MemoryPageScan(start, end, logSizeTracker);
+                // Legacy observer path for consumers still using SubscribeEvictions.
+                MemoryPageScan(start, end, onEvictionObserver);
             }
 
-            // TODO: Currently we don't call OnDispose or OnDisposeValueObject on eviction; we defer to the OnEvictionObserver
-            // and do nothing if that is not supplied. Should we add our own observer if they don't supply one?
             _wrapper.FreePage(page);
         }
 
@@ -1510,9 +1508,11 @@ private void OnPagesClosedWorker()
                     if (onEvictionObserver is not null)
                         MemoryPageScan(start, end, onEvictionObserver);
 
-                    // Notify application of records being evicted — allows cleanup of external resources.
+                    // Per-record eviction walk: handles internal heap-size accounting (key overflow
+                    // and value heap via logSizeTracker) and optionally notifies the application
+                    // via OnEvict for app-level cleanup.
                     var evictSource = IsReadCache ? EvictionSource.ReadCache : EvictionSource.MainLog;
-                    if (storeFunctions.CallOnEvict(evictSource))
+                    if (logSizeTracker is not null || storeFunctions.CallOnEvict(evictSource))
                         _wrapper.EvictRecordsInRange(start, end, evictSource);
 
                     // If we are using a null storage device, we must also shift BeginAddress (leave it in-memory)

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

@@ -1097,7 +1097,11 @@ public readonly bool TryCopyFrom<TSourceLogRecord>(in TSourceLogRecord srcLogRec
                 }
                 else
                 {
-                    // TODOobjDispose: make sure Object isn't disposed by the source, to avoid use-after-Dispose. Maybe this (and DiskLogRecord remapping to TransientOIDMap) needs Clone()
+                    // TODO: Clone the value object here so source and destination have independent
+                    // HeapMemorySize fields. Currently both records share the same IHeapObject instance,
+                    // which means mutations on the destination affect the source's reported heap size
+                    // at eviction time, causing accounting drift in logSizeTracker. A naive Clone()
+                    // here causes CanDoHashExpireLTM to crash — needs investigation in a follow-up.
                     Debug.Assert(srcLogRecord.ValueObject is not null, "Expected srcLogRecord.ValueObject to be set (or deserialized) already");
                     if (!TrySetValueObjectAndPrepareOptionals(srcLogRecord.ValueObject, in sizeInfo))
                         return false;

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

@@ -328,6 +328,16 @@ internal void OnDispose(ref LogRecord logRecord, DisposeReason disposeReason)
         {
             if (logRecord.IsSet)
             {
+                // For Deleted records, decrement the value-side heap from the tracker before clearing.
+                // Key overflow is handled at eviction time in EvictRecordsInRange. Tombstone has not
+                // been set yet when this is called, so GetValueHeapMemorySize returns the correct value.
+                if (disposeReason == DisposeReason.Deleted)
+                {
+                    var valueHeap = logRecord.GetValueHeapMemorySize();
+                    if (valueHeap != 0)
+                        logSizeTracker?.IncrementSize(-valueHeap);
+                }
+
                 storeFunctions.OnDispose(ref logRecord, disposeReason);
 
                 logRecord.ClearHeapFields(disposeReason != DisposeReason.Deleted);
@@ -354,8 +364,6 @@ internal void OnDisposeDiskRecord(ref DiskLogRecord logRecord, DisposeReason dis
         /// </summary>
         internal void EvictRecordsInRange(long startAddress, long endAddress, EvictionSource source)
         {
-            // Clip to a single page: we don't cross page boundaries in one call, but be defensive in case the caller
-            // ever passes a multi-page range. The start-page PageHeader is skipped via GetFirstValidLogicalAddressOnPage.
             var startPage = GetPage(startAddress);
             var firstValidAddress = GetFirstValidLogicalAddressOnPage(startPage);
             var address = startAddress < firstValidAddress ? firstValidAddress : startAddress;
@@ -368,35 +376,38 @@ internal void EvictRecordsInRange(long startAddress, long endAddress, EvictionSo
                 var logRecord = new LogRecord(physicalAddress, objectPages[GetPageIndexForAddress(address)].objectIdMap);
                 var allocatedSize = logRecord.AllocatedSize;
 
-                // Guard against corrupt / zero-size records causing infinite loops. Also stop if the record would
-                // straddle the end of the page; ObjectScanIterator skips to the next page in that case, but our caller
-                // constrains the range to a single page so there is nothing further to visit here.
                 if (allocatedSize <= 0)
                     break;
                 var offset = GetOffsetOnPage(address);
                 if (offset == 0 || offset + allocatedSize > PageSize)
                     break;
 
-                // Skip null, invalid, and tombstoned records:
-                //  - IsNull: empty record slot, nothing to account for.
-                //  - Invalid: record was elided from the hash chain and already disposed
-                //    (OnDispose(Deleted/Elided) or transferred to the revivification freelist).
-                //  - Tombstone: heap was already decremented at the delete site via
-                //    OnDispose(Deleted).
-                //
-                // Sealed-but-Valid records are NOT skipped. A Sealed source record from a
-                // mutable-region CopyUpdate still owns its overflow key/value bytes (and
-                // possibly its value object if a checkpoint prevented eager clearing via
-                // ClearSourceValueObject). Skipping them would leak their heap contribution
-                // in the tracker.
-                if (logRecord.Info.IsNull || logRecord.Info.Invalid || logRecord.Info.Tombstone)
+                // Skip null and invalid records (elided/disposed, heap already cleaned up).
+                if (logRecord.Info.IsNull || logRecord.Info.Invalid)
                 {
                     address += allocatedSize;
                     continue;
                 }
 
-                // Notify the application that this record is being evicted from memory.
-                storeFunctions.OnEvict(ref logRecord, source);
+                // Decrement the record's heap contribution in a single call.
+                // For non-tombstoned records, CalculateHeapMemorySize returns key + value heap.
+                // For tombstoned records, it returns 0 (by design), but the key overflow is still
+                // alive — value was already decremented at the delete site via OnDispose.
+                long heapSize;
+                if (!logRecord.Info.Tombstone)
+                {
+                    heapSize = logRecord.CalculateHeapMemorySize();
+
+                    if (storeFunctions.CallOnEvict(source))
+                        storeFunctions.OnEvict(ref logRecord, source);
+                }
+                else
+                {
+                    heapSize = logRecord.Info.KeyIsOverflow ? logRecord.KeyOverflow.HeapMemorySize : 0;
+                }
+
+                if (heapSize != 0)
+                    logSizeTracker?.IncrementSize(-heapSize);
 
                 address += allocatedSize;
             }