refactor: move tombstone after OnDispose so triggers see pre-tombstone state

Session functions (ISessionFunctions) should not be responsible for heap-size
tracking on destruction paths — that's IRecordTriggers' job. Two places were
setting Tombstone BEFORE OnDispose(Deleted) fired, causing CalculateHeapMemorySize
to return 0 and preventing the trigger from decrementing the tracked heap:

1. SessionFunctionsWrapper.InPlaceDeleter: set Tombstone+Dirty before returning
   to InternalDelete, which then called OnDispose(Deleted). Fixed by moving
   SetTombstone+SetDirtyAndModified to InternalDelete after OnDispose.

2. SessionFunctionsWrapper.InPlaceUpdater (ExpireAndStop): set Tombstone before
   returning to InternalRMW, which then called OnDispose(Deleted). Fixed by
   moving SetTombstone+SetDirtyAndModified to InternalRMW after OnDispose.

This removes the manual cacheSizeTracker decrement from MainStore.InPlaceDeleter
(it was a workaround for the wrong ordering) and makes GarnetRecordTriggers.OnDispose
the single destruction-side decrement path for all record types.

Also tightened PageEvictionFiresOnEvictForEveryLiveRecord test to use precise
counts by observing TailAddress movement to distinguish mutable vs immutable deletes.

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

Author: badrishc

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

@@ -28,14 +28,6 @@ public void PostInitialDeleter(ref LogRecord logRecord, ref DeleteInfo deleteInf
         /// <inheritdoc />
         public bool InPlaceDeleter(ref LogRecord logRecord, ref DeleteInfo deleteInfo)
         {
-            // Decrement here — before the wrapper sets Tombstone. After Tombstone is set,
-            // CalculateHeapMemorySize short-circuits to zero and EvictRecordsInRange skips
-            // the record, so any matching decrement would be lost and the creation-side
-            // increments in PostInitialUpdater/PostInitialWriter/InPlaceUpdater would leak.
-            var heap = logRecord.CalculateHeapMemorySize();
-            if (heap != 0)
-                functionsState.cacheSizeTracker?.AddHeapSize(-heap);
-
             logRecord.ClearOptionals();
             if (!logRecord.Info.Modified)
                 functionsState.watchVersionMap.IncrementVersion(deleteInfo.KeyHash);

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

@@ -193,12 +193,14 @@ public bool InPlaceUpdater(ref LogRecord logRecord, in RecordSizeInfo sizeInfo,
             }
             else if (rmwInfo.Action == RMWAction.ExpireAndStop)
             {
-                logRecord.InfoRef.SetTombstone();
+                // Tombstone is set by the caller (InternalRMW) AFTER OnDispose(Deleted),
+                // so that CalculateHeapMemorySize sees the pre-tombstone state.
                 status = OperationStatusUtils.AdvancedOpCode(OperationStatus.SUCCESS, StatusCode.InPlaceUpdatedRecord | StatusCode.Expired);
             }
             else if (rmwInfo.Action == RMWAction.WrongType)
             {
-                logRecord.InfoRef.SetTombstone();
+                // WrongType means the operation was rejected — the record must NOT be modified.
+                // Do not set Tombstone; the key should remain intact for correct-type operations.
                 status = OperationStatusUtils.AdvancedOpCode(OperationStatus.NOTFOUND, StatusCode.WrongType);
             }
             else
@@ -238,8 +240,9 @@ public bool InPlaceDeleter(ref LogRecord logRecord, ref DeleteInfo deleteInfo)
         {
             if (!_clientSession.functions.InPlaceDeleter(ref logRecord, ref deleteInfo))
                 return false;
-            logRecord.InfoRef.SetTombstone();
-            logRecord.InfoRef.SetDirtyAndModified();
+            // Tombstone and Dirty/Modified are set by the caller (InternalDelete) AFTER
+            // OnDispose(Deleted), so that CalculateHeapMemorySize sees the pre-tombstone
+            // state and IRecordTriggers can correctly account for the full heap contribution.
             return true;
         }
 

libs/storage/Tsavorite/cs/src/core/Index/Tsavorite/Implementation/InternalDelete.cs

@@ -124,10 +124,13 @@ internal OperationStatus InternalDelete<TKey, TInput, TOutput, TContext, TSessio
                     // DeleteInfo's lengths are filled in and GetRecordLengths and SetDeletedValueLength are called inside InPlaceDeleter.
                     if (sessionFunctions.InPlaceDeleter(ref srcLogRecord, ref deleteInfo))
                     {
-                        // Immediately dispose all resources: app callback (storeFunctions.OnDispose)
-                        // + ClearHeapFields + ClearOptionals — all via hlog.OnDispose.
+                        // Dispose all resources BEFORE setting Tombstone, so that
+                        // IRecordTriggers.OnDispose sees the pre-tombstone state and
+                        // CalculateHeapMemorySize returns the full heap contribution.
                         OnDispose(ref srcLogRecord, DisposeReason.Deleted);
 
+                        srcLogRecord.InfoRef.SetTombstone();
+                        srcLogRecord.InfoRef.SetDirtyAndModified();
                         MarkPage(stackCtx.recSrc.LogicalAddress, sessionFunctions.Ctx);
 
                         // Try to transfer the record from the tag chain to the free record pool iff previous address points to invalid address.

libs/storage/Tsavorite/cs/src/core/Index/Tsavorite/Implementation/InternalRMW.cs

@@ -156,18 +156,18 @@ internal OperationStatus InternalRMW<TKey, TInput, TOutput, TContext, TSessionFu
                     if (rmwInfo.Action == RMWAction.ExpireAndStop)
                     {
                         MarkPage(stackCtx.recSrc.LogicalAddress, sessionFunctions.Ctx);
-                        // Tombstone has been set by SessionFunctionsWrapper.InPlaceUpdater
-                        srcLogRecord.InfoRef.SetDirtyAndModified();
 
                         pendingContext.logicalAddress = stackCtx.recSrc.LogicalAddress;
                         pendingContext.eTag = srcLogRecord.ETag;
 
-                        // ExpireAndStop means to override default Delete handling (which is to go to InitialUpdater) by leaving the tombstoned record as current.
-                        // Our SessionFunctionsWrapper.InPlaceUpdater implementation has already reinitialized-in-place or set Tombstone as appropriate and marked the record.
-
-                        // Immediately dispose all resources at delete site.
+                        // Dispose all resources BEFORE setting Tombstone, so that
+                        // IRecordTriggers.OnDispose sees the pre-tombstone state and
+                        // CalculateHeapMemorySize returns the full heap contribution.
                         OnDispose(ref srcLogRecord, DisposeReason.Deleted);
 
+                        srcLogRecord.InfoRef.SetTombstone();
+                        srcLogRecord.InfoRef.SetDirtyAndModified();
+
                         // Try to transfer the record from the tag chain to the free record pool iff previous address points to invalid address.
                         // Otherwise an earlier record for this key could be reachable again.
                         if (CanElide<TInput, TOutput, TContext, TSessionFunctionsWrapper>(sessionFunctions, ref stackCtx, srcLogRecord.Info))

libs/storage/Tsavorite/cs/test/RecordLifecycleTests.cs

@@ -459,53 +459,52 @@ public void PendingReadFromDiskFiresOnDisposeDiskRecordOnce()
         /// <summary>
         /// Filling the log well beyond its mutable window forces page eviction. OnEvict must fire for
         /// every non-tombstoned, non-invalid record evicted past HeadAddress — including sealed source
-        /// records from immutable-region deletes or CopyUpdates. Tombstoned records are skipped (heap
-        /// was decremented at the delete site). Invalid/elided records are skipped (heap was already
-        /// cleaned up). Sealed records are visited because they may still own overflow key/value bytes
-        /// or value objects (checkpoint path).
+        /// records from immutable-region deletes. Tombstoned records are skipped (heap was decremented
+        /// at the delete site). Invalid/elided records are skipped (already cleaned up).
         /// </summary>
         [Test, Category("TsavoriteKV")]
         public void PageEvictionFiresOnEvictForEveryLiveRecord()
         {
             const int n = 500; // enough to push pages past HeadAddress given 32KB memory / 1KB pages
             for (int i = 0; i < n; i++) Upsert(i, i);
 
-            // Delete ~10% of records to create tombstones. OnDispose(Deleted) fires for each.
+            // Delete ~10% of records. Track which go through the immutable path (TailAddress moves)
+            // vs in-place mutable path (TailAddress stays).
             tracker.Reset();
             const int deleted = 50;
+            int mutableDeletes = 0;
             using (var s = NewSession())
             {
                 for (int i = 0; i < deleted; i++)
+                {
+                    var tailBefore = store.Log.TailAddress;
                     _ = s.BasicContext.Delete(new TestObjectKey { key = i });
+                    if (store.Log.TailAddress == tailBefore)
+                        mutableDeletes++;
+                }
             }
+            var immutableDeletes = deleted - mutableDeletes;
+
             ClassicAssert.AreEqual(deleted, tracker.DisposeCount(DisposeReason.Deleted),
                 "Each Delete should fire OnDispose(Deleted) exactly once");
             var deletedDisposeCountBeforeEvict = tracker.DisposeCount(DisposeReason.Deleted);
 
-            // Force all records out to disk. This triggers OnEvict for every non-tombstoned,
-            // non-invalid record on each evicted page.
+            // Force all records out to disk.
             store.Log.FlushAndEvict(wait: true);
 
-            // OnEvict fires for: (1) the n live records that were never deleted, plus (2) sealed
-            // source records from immutable-region deletes (these are not tombstoned — the tombstone
-            // is on a NEW record at the tail). In-place (mutable) deletes set tombstone directly on
-            // the source and are skipped. The exact split depends on the mutable-vs-readonly boundary
-            // at delete time, so we assert bounds rather than an exact count:
-            //  - Lower bound: at least n - deleted (the live records)
-            //  - Upper bound: at most n (if all deletes hit the readonly region, producing sealed sources)
-            var mainLogEvicts = tracker.EvictCount(EvictionSource.MainLog);
-            ClassicAssert.GreaterOrEqual(mainLogEvicts, n - deleted,
-                $"OnEvict(MainLog) must fire for at least the {n - deleted} non-deleted records");
-            ClassicAssert.LessOrEqual(mainLogEvicts, n,
-                $"OnEvict(MainLog) must not fire more than {n} times (n original + sealed sources, minus tombstones)");
-
+            // Precise count:
+            //  - (n - deleted) live records: visited by OnEvict.
+            //  - mutableDeletes records: tombstoned in-place, skipped by OnEvict.
+            //  - immutableDeletes sealed source records: NOT tombstoned, visited by OnEvict.
+            //  - immutableDeletes new tombstone records at tail: tombstoned, skipped by OnEvict.
+            // Total = (n - deleted) + immutableDeletes = n - mutableDeletes.
+            ClassicAssert.AreEqual(n - mutableDeletes, tracker.EvictCount(EvictionSource.MainLog),
+                $"OnEvict(MainLog) must fire exactly {n - mutableDeletes} times: " +
+                $"{n - deleted} live + {immutableDeletes} sealed sources, skipping {mutableDeletes} in-place tombstones");
             ClassicAssert.AreEqual(0, tracker.EvictCount(EvictionSource.ReadCache),
                 "No read cache is configured, OnEvict(ReadCache) must never fire");
-            // Tombstoned records must NOT produce a paired OnDispose(Deleted) during eviction; the
-            // OnDispose(Deleted) count should be unchanged from the pre-eviction snapshot.
             ClassicAssert.AreEqual(deletedDisposeCountBeforeEvict, tracker.DisposeCount(DisposeReason.Deleted),
                 "Page eviction must not re-fire OnDispose(Deleted) for tombstoned records");
-            // And OnDisposeDiskRecord must not fire during eviction — eviction uses OnEvict only.
             ClassicAssert.AreEqual(0, tracker.TotalDisposeDisk(),
                 "Page eviction must not route through OnDisposeDiskRecord");
         }