Internalize heap-size tracking into Tsavorite

libs/cluster/Session/RespClusterMigrateCommands.cs

@@ -185,7 +185,9 @@ void Process(BasicGarnetApi basicGarnetApi, byte[] input, bool replaceOption, bo
                                     if (replaceOption || !Exists(keySlice))
                                         _ = basicGarnetApi.SET(in diskLogRecord);
 
+                                    storeWrapper.storeFunctions.OnDisposeDiskRecord(ref diskLogRecord, DisposeReason.DeserializedFromDisk);
                                     diskLogRecord.Dispose();
+                                    diskLogRecord = default; // prevent double-trigger in finally
                                 }
                                 else
                                 {
@@ -198,7 +200,11 @@ void Process(BasicGarnetApi basicGarnetApi, byte[] input, bool replaceOption, bo
                     }
                     finally
                     {
-                        diskLogRecord.Dispose();
+                        if (diskLogRecord.IsSet)
+                        {
+                            storeWrapper.storeFunctions.OnDisposeDiskRecord(ref diskLogRecord, DisposeReason.DeserializedFromDisk);
+                            diskLogRecord.Dispose();
+                        }
                     }
                 }
             }

libs/cluster/Session/RespClusterReplicationCommands.cs

@@ -536,7 +536,9 @@ private bool NetworkClusterSync(out bool invalidParameters)
 
                         diskLogRecord = DiskLogRecord.Deserialize(recordSpan, storeWrapper.GarnetObjectSerializer, transientObjectIdMap, storeWrapper.storeFunctions);
                         _ = basicGarnetApi.SET(in diskLogRecord);
+                        storeWrapper.storeFunctions.OnDisposeDiskRecord(ref diskLogRecord, DisposeReason.DeserializedFromDisk);
                         diskLogRecord.Dispose();
+                        diskLogRecord = default; // prevent double-trigger in catch
                     }
                     else
                     {
@@ -549,7 +551,11 @@ private bool NetworkClusterSync(out bool invalidParameters)
             catch
             {
                 // Dispose the diskLogRecord if there was an exception in SET
-                diskLogRecord.Dispose();
+                if (diskLogRecord.IsSet)
+                {
+                    storeWrapper.storeFunctions.OnDisposeDiskRecord(ref diskLogRecord, DisposeReason.DeserializedFromDisk);
+                    diskLogRecord.Dispose();
+                }
                 throw;
             }
 

libs/host/GarnetServer.cs

@@ -372,7 +372,7 @@ private TsavoriteKV<StoreFunctions, StoreAllocator> CreateStore(int dbId, IClust
             var store = new TsavoriteKV<StoreFunctions, StoreAllocator>(kvSettings
                 , Tsavorite.core.StoreFunctions.Create(new GarnetKeyComparer(),
                     () => new GarnetObjectSerializer(customCommandManager),
-                    new GarnetRecordTriggers(cacheSizeTracker))
+                    new GarnetRecordTriggers())
                 , (allocatorSettings, storeFunctions) => new(allocatorSettings, storeFunctions));
 
             if (kvSettings.LogMemorySize > 0 || kvSettings.ReadCacheMemorySize > 0)

libs/server/Custom/CustomObjectBase.cs

@@ -75,10 +75,8 @@ public sealed override void DoSerialize(BinaryWriter writer)
 
         /// <inheritdoc />
         public sealed override bool Operate(ref ObjectInput input, ref ObjectOutput output,
-                                            byte respProtocolVersion, out long sizeChange)
+                                            byte respProtocolVersion)
         {
-            sizeChange = 0;
-
             switch (input.header.cmd)
             {
                 // Scan Command

libs/server/Objects/Hash/HashObject.cs

@@ -182,10 +182,8 @@ public override void Dispose() { }
         public override GarnetObjectBase Clone() => new HashObject(hash, expirationTimes, expirationQueue, HeapMemorySize);
 
         /// <inheritdoc />
-        public override bool Operate(ref ObjectInput input, ref ObjectOutput output, byte respProtocolVersion, out long memorySizeChange)
+        public override bool Operate(ref ObjectInput input, ref ObjectOutput output, byte respProtocolVersion)
         {
-            memorySizeChange = 0;
-
             if (input.header.type != GarnetObjectType.Hash)
             {
                 //Indicates when there is an incorrect type 
@@ -194,7 +192,6 @@ public override bool Operate(ref ObjectInput input, ref ObjectOutput output, byt
                 return true;
             }
 
-            var previousMemorySize = HeapMemorySize;
             switch (input.header.HashOp)
             {
                 case HashOperation.HSET:
@@ -261,8 +258,6 @@ public override bool Operate(ref ObjectInput input, ref ObjectOutput output, byt
                     throw new GarnetException($"Unsupported operation {input.header.HashOp} in HashObject.Operate");
             }
 
-            memorySizeChange = HeapMemorySize - previousMemorySize;
-
             if (hash.Count == 0)
                 output.OutputFlags |= ObjectOutputFlags.RemoveKey;
 

libs/server/Objects/List/ListObject.cs

@@ -128,10 +128,8 @@ public override void Dispose() { }
 
         /// <inheritdoc />
         public override bool Operate(ref ObjectInput input, ref ObjectOutput output,
-                                     byte respProtocolVersion, out long memorySizeChange)
+                                     byte respProtocolVersion)
         {
-            memorySizeChange = 0;
-
             if (input.header.type != GarnetObjectType.List)
             {
                 // Indicates an incorrect type of key
@@ -140,7 +138,6 @@ public override bool Operate(ref ObjectInput input, ref ObjectOutput output,
                 return true;
             }
 
-            var previousMemorySize = HeapMemorySize;
             switch (input.header.ListOp)
             {
                 case ListOperation.LPUSH:
@@ -186,8 +183,6 @@ public override bool Operate(ref ObjectInput input, ref ObjectOutput output,
                     throw new GarnetException($"Unsupported operation {input.header.ListOp} in ListObject.Operate");
             }
 
-            memorySizeChange = HeapMemorySize - previousMemorySize;
-
             if (list.Count == 0)
                 output.OutputFlags |= ObjectOutputFlags.RemoveKey;
 

libs/server/Objects/Set/SetObject.cs

@@ -122,10 +122,8 @@ public override void Dispose() { }
 
         /// <inheritdoc />
         public override bool Operate(ref ObjectInput input, ref ObjectOutput output,
-                                     byte respProtocolVersion, out long memorySizeChange)
+                                     byte respProtocolVersion)
         {
-            memorySizeChange = 0;
-
             if (input.header.type != GarnetObjectType.Set)
             {
                 // Indicates an incorrect type of key
@@ -134,7 +132,6 @@ public override bool Operate(ref ObjectInput input, ref ObjectOutput output,
                 return true;
             }
 
-            var prevMemorySize = HeapMemorySize;
             switch (input.header.SetOp)
             {
                 case SetOperation.SADD:
@@ -168,8 +165,6 @@ public override bool Operate(ref ObjectInput input, ref ObjectOutput output,
                     throw new GarnetException($"Unsupported operation {input.header.SetOp} in SetObject.Operate");
             }
 
-            memorySizeChange = HeapMemorySize - prevMemorySize;
-
             if (Set.Count == 0)
                 output.OutputFlags |= ObjectOutputFlags.RemoveKey;
 

libs/server/Objects/SortedSet/SortedSetObject.cs

@@ -316,10 +316,8 @@ public override void Dispose() { }
 
         /// <inheritdoc />
         public override bool Operate(ref ObjectInput input, ref ObjectOutput output,
-                                     byte respProtocolVersion, out long memorySizeChange)
+                                     byte respProtocolVersion)
         {
-            memorySizeChange = 0;
-
             var header = input.header;
             if (header.type != GarnetObjectType.SortedSet)
             {
@@ -329,7 +327,6 @@ public override bool Operate(ref ObjectInput input, ref ObjectOutput output,
                 return true;
             }
 
-            var prevMemorySize = HeapMemorySize;
             var op = header.SortedSetOp;
             switch (op)
             {
@@ -415,8 +412,6 @@ public override bool Operate(ref ObjectInput input, ref ObjectOutput output,
                     throw new GarnetException($"Unsupported operation {op} in SortedSetObject.Operate");
             }
 
-            memorySizeChange = HeapMemorySize - prevMemorySize;
-
             if (sortedSetDict.Count == 0)
                 output.OutputFlags |= ObjectOutputFlags.RemoveKey;
 

libs/server/Objects/Types/GarnetObjectBase.cs

@@ -29,7 +29,7 @@ protected GarnetObjectBase(BinaryReader reader, long heapMemorySize)
         }
 
         /// <inheritdoc />
-        public abstract bool Operate(ref ObjectInput input, ref ObjectOutput output, byte respProtocolVersion, out long sizeChange);
+        public abstract bool Operate(ref ObjectInput input, ref ObjectOutput output, byte respProtocolVersion);
 
         /// <summary>
         /// Serialize to given writer

libs/server/Objects/Types/IGarnetObject.cs

@@ -21,9 +21,8 @@ public interface IGarnetObject : IHeapObject
         /// </summary>
         /// <param name="input"></param>
         /// <param name="output"></param>
-        /// <param name="sizeChange"></param>
         /// <returns></returns>
-        bool Operate(ref ObjectInput input, ref ObjectOutput output, byte respProtocolVersion, out long sizeChange);
+        bool Operate(ref ObjectInput input, ref ObjectOutput output, byte respProtocolVersion);
 
         /// <summary>
         /// Scan the items of the collection

libs/server/Storage/Functions/GarnetRecordTriggers.cs

@@ -6,48 +6,12 @@
 namespace Garnet.server
 {
     /// <summary>
-    /// Record lifecycle triggers for Garnet's unified store. Handles per-record cleanup
-    /// on delete via <see cref="IRecordTriggers.OnDispose"/>.
+    /// Record lifecycle triggers for Garnet's unified store. Available for app-level
+    /// resource cleanup (e.g. <see cref="System.IDisposable.Dispose"/> on value objects
+    /// holding external resources). Currently a no-op — Garnet's IHeapObject implementations
+    /// (Hash/List/Set/SortedSet) hold no external resources.
     /// </summary>
     public readonly struct GarnetRecordTriggers : IRecordTriggers
     {
-        /// <summary>
-        /// Cache size tracker for heap size accounting on delete.
-        /// Created before the store and initialized after via <see cref="CacheSizeTracker.Initialize"/>.
-        /// </summary>
-        internal readonly CacheSizeTracker cacheSizeTracker;
-
-        /// <summary>
-        /// Creates a GarnetRecordTriggers with a cache size tracker.
-        /// </summary>
-        public GarnetRecordTriggers(CacheSizeTracker cacheSizeTracker)
-        {
-            this.cacheSizeTracker = cacheSizeTracker;
-        }
-
-        /// <inheritdoc/>
-        public bool CallOnFlush => false;
-
-        /// <inheritdoc/>
-        public bool CallOnEvict => false;
-
-        /// <inheritdoc/>
-        public bool CallOnDiskRead => false;
-
-        /// <inheritdoc/>
-        public void OnDisposeValueObject(IHeapObject valueObject, DisposeReason reason)
-        {
-            // Heap object disposal is handled by ClearHeapFields in ObjectAllocatorImpl
-        }
-
-        /// <inheritdoc/>
-        public void OnDispose(ref LogRecord logRecord, DisposeReason reason)
-        {
-            // Handle heap objects: update cache size tracker on delete
-            if (logRecord.Info.ValueIsObject && reason == DisposeReason.Deleted)
-            {
-                cacheSizeTracker?.AddHeapSize(-logRecord.ValueObject.HeapMemorySize);
-            }
-        }
     }
 }

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

@@ -1701,6 +1701,7 @@ public readonly bool PostCopyUpdater<TSourceLogRecord>(in TSourceLogRecord srcLo
             functionsState.watchVersionMap.IncrementVersion(rmwInfo.KeyHash);
             if (functionsState.appendOnlyFile != null)
                 rmwInfo.UserData |= NeedAofLog; // Mark that we need to write to AOF
+
             return true;
         }
 

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

@@ -50,7 +50,7 @@ public bool InitialUpdater(ref LogRecord logRecord, in RecordSizeInfo sizeInfo,
             if ((byte)type < CustomCommandManager.CustomTypeIdStartOffset)
             {
                 value = GarnetObject.Create(type);
-                _ = value.Operate(ref input, ref output, functionsState.respProtocolVersion, out _);
+                _ = value.Operate(ref input, ref output, functionsState.respProtocolVersion);
                 _ = logRecord.TrySetValueObjectAndPrepareOptionals(value, in sizeInfo);
                 return true;
             }
@@ -82,8 +82,6 @@ public void PostInitialUpdater(ref LogRecord dstLogRecord, in RecordSizeInfo siz
                 input.header.SetExpiredFlag();
                 rmwInfo.UserData |= NeedAofLog; // Mark that we need to write to AOF
             }
-
-            functionsState.cacheSizeTracker?.AddHeapSize(dstLogRecord.ValueObject.HeapMemorySize);
         }
 
         /// <inheritdoc />
@@ -96,34 +94,29 @@ public bool InPlaceUpdater(ref LogRecord logRecord, ref ObjectInput input, ref O
                 return true;
             }
 
-            if (InPlaceUpdaterWorker(ref logRecord, ref input, ref output, ref rmwInfo, out long sizeChange))
+            if (InPlaceUpdaterWorker(ref logRecord, ref input, ref output, ref rmwInfo))
             {
                 if (!logRecord.Info.Modified)
                     functionsState.watchVersionMap.IncrementVersion(rmwInfo.KeyHash);
                 if (functionsState.appendOnlyFile != null)
                     rmwInfo.UserData |= NeedAofLog; // Mark that we need to write to AOF
-                functionsState.cacheSizeTracker?.AddHeapSize(sizeChange);
                 return true;
             }
             return false;
         }
 
-        bool InPlaceUpdaterWorker(ref LogRecord logRecord, ref ObjectInput input, ref ObjectOutput output, ref RMWInfo rmwInfo, out long sizeChange)
+        bool InPlaceUpdaterWorker(ref LogRecord logRecord, ref ObjectInput input, ref ObjectOutput output, ref RMWInfo rmwInfo)
         {
-            sizeChange = 0;
-
             // Expired data
             if (logRecord.Info.HasExpiration && input.header.CheckExpiry(logRecord.Expiration))
             {
-                // Heap disposal and cache size tracking are handled by
-                // OnDispose(Deleted) in InternalRMW when processing ExpireAndResume.
                 rmwInfo.Action = RMWAction.ExpireAndResume;
                 return false;
             }
 
             if ((byte)input.header.type < CustomCommandManager.CustomTypeIdStartOffset)
             {
-                var operateSuccessful = ((IGarnetObject)logRecord.ValueObject).Operate(ref input, ref output, functionsState.respProtocolVersion, out sizeChange);
+                var operateSuccessful = ((IGarnetObject)logRecord.ValueObject).Operate(ref input, ref output, functionsState.respProtocolVersion);
                 if (output.HasWrongType)
                     return true;
                 if (output.HasRemoveKey)
@@ -132,8 +125,6 @@ bool InPlaceUpdaterWorker(ref LogRecord logRecord, ref ObjectInput input, ref Ob
                         functionsState.watchVersionMap.IncrementVersion(rmwInfo.KeyHash);
                     if (functionsState.appendOnlyFile != null)
                         rmwInfo.UserData |= NeedAofLog;
-                    // Heap disposal and cache size tracking are handled by
-                    // OnDispose(Deleted) in the ExpireAndStop path of InternalRMW.
                     rmwInfo.Action = RMWAction.ExpireAndStop;
                     return false;
                 }
@@ -190,11 +181,11 @@ public bool CopyUpdater<TSourceLogRecord>(in TSourceLogRecord srcLogRecord, ref
         public bool PostCopyUpdater<TSourceLogRecord>(in TSourceLogRecord srcLogRecord, ref LogRecord dstLogRecord, in RecordSizeInfo sizeInfo, ref ObjectInput input, ref ObjectOutput output, ref RMWInfo rmwInfo)
             where TSourceLogRecord : ISourceLogRecord
         {
-            // We're performing the object update here (and not in CopyUpdater) so that we are guaranteed that
-            // the record was CASed into the hash chain before it gets modified
-            var value = Unsafe.As<IGarnetObject>(srcLogRecord.ValueObject.Clone());
-            var oldValueSize = srcLogRecord.ValueObject.HeapMemorySize;
-            _ = dstLogRecord.TrySetValueObject(value);
+            // We perform the object update here (and not in CopyUpdater) so that we are guaranteed that the record
+            // was CASed into the hash chain before it gets modified. Tsavorite's CacheSerializedObjectData (called
+            // by InternalRMW right before PCU, for both memory and disk sources) has already cloned src.ValueObject
+            // into dstLogRecord. Reuse it directly — do not clone again.
+            var value = Unsafe.As<IGarnetObject>(dstLogRecord.ValueObject);
 
             // Do not set actually set dstLogRecord.Expiration until we know it is a command for which we allocated length in the LogRecord for it.
             // TODO: Object store ETags
@@ -203,7 +194,7 @@ public bool PostCopyUpdater<TSourceLogRecord>(in TSourceLogRecord srcLogRecord,
 
             if ((byte)input.header.type < CustomCommandManager.CustomTypeIdStartOffset)
             {
-                value.Operate(ref input, ref output, functionsState.respProtocolVersion, out _);
+                value.Operate(ref input, ref output, functionsState.respProtocolVersion);
                 if (output.HasWrongType)
                     return true;
                 if (output.HasRemoveKey)
@@ -241,10 +232,6 @@ public bool PostCopyUpdater<TSourceLogRecord>(in TSourceLogRecord srcLogRecord,
 
             sizeInfo.AssertOptionalsIfSet(dstLogRecord.Info);
 
-            // If oldValue has been set to null, subtract its size from the tracked heap size
-            var sizeAdjustment = rmwInfo.ClearSourceValueObject ? value.HeapMemorySize - oldValueSize : value.HeapMemorySize;
-            functionsState.cacheSizeTracker?.AddHeapSize(sizeAdjustment);
-
             if (functionsState.appendOnlyFile != null)
                 rmwInfo.UserData |= NeedAofLog; // Mark that we need to write to AOF
             return true;

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

@@ -34,7 +34,7 @@ public bool Reader<TSourceLogRecord>(in TSourceLogRecord srcLogRecord, ref Objec
                 var garnetObject = (IGarnetObject)srcLogRecord.ValueObject;
                 if ((byte)input.header.type < CustomCommandManager.CustomTypeIdStartOffset)
                 {
-                    var opResult = garnetObject.Operate(ref input, ref output, functionsState.respProtocolVersion, out _);
+                    var opResult = garnetObject.Operate(ref input, ref output, functionsState.respProtocolVersion);
                     if (output.HasWrongType)
                         return true;