Fix Garnet RMW InitialUpdater to work with Revivification (In-Chain) (#1307)

* Rmw and Reviv fix

* WIP

* wip

* fmt

* fmt

* WIP

* fmt

* fmt

* stupid indentation

* fix double setting of filler breaking debug assertion

* wip

* cleanup

* try a much cleaner way

* wip

* fmt

* fix test

* fix up setting of spanbyte lengths

* clean used value length setting

* is needed after all

* Update libs/storage/Tsavorite/cs/src/core/Index/Tsavorite/Implementation/Revivification/RecordLengths.cs

Co-authored-by: Copilot <[email protected]>

* fix missing copy

* uno mas comment

---------

Co-authored-by: Hamdaan Khalid <[email protected]>
Co-authored-by: Copilot <[email protected]>

Author: 3 people

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

@@ -92,23 +92,19 @@ public bool InitialUpdater(ref SpanByte key, ref RawStringInput input, ref SpanB
                     var length = sbSrcHLL.Length;
                     var srcHLL = sbSrcHLL.ToPointer();
                     var dstHLL = value.ToPointer();
-
                     value.ShrinkSerializedLength(length);
                     Buffer.MemoryCopy(srcHLL, dstHLL, value.Length, value.Length);
                     break;
 
                 case RespCommand.SETIFGREATER:
                 case RespCommand.SETIFMATCH:
                     int spaceForEtag = this.functionsState.etagState.etagOffsetForVarlen;
-                    // Copy input to value
                     var newInputValue = input.parseState.GetArgSliceByRef(0).ReadOnlySpan;
                     var metadataSize = input.arg1 == 0 ? 0 : sizeof(long);
                     value.ShrinkSerializedLength(newInputValue.Length + metadataSize + spaceForEtag);
                     value.ExtraMetadata = input.arg1;
                     newInputValue.CopyTo(value.AsSpan(spaceForEtag));
-
                     long clientSentEtag = input.parseState.GetLong(1);
-
                     if (cmd == RespCommand.SETIFMATCH)
                         clientSentEtag++;
 
@@ -183,11 +179,9 @@ public bool InitialUpdater(ref SpanByte key, ref RawStringInput input, ref SpanB
                 case RespCommand.SETBIT:
                     var bOffset = input.arg1;
                     var bSetVal = (byte)(input.parseState.GetArgSliceByRef(1).ReadOnlySpan[0] - '0');
-
                     value.ShrinkSerializedLength(BitmapManager.Length(bOffset));
-
-                    // Always return 0 at initial updater because previous value was 0
                     BitmapManager.UpdateBitmap(value.ToPointer(), bOffset, bSetVal);
+                    // Always return 0 at initial updater because previous value was 0
                     CopyDefaultResp(CmdStrings.RESP_RETURN_VAL_0, ref output);
                     break;
 
@@ -219,9 +213,7 @@ public bool InitialUpdater(ref SpanByte key, ref RawStringInput input, ref SpanB
                 case RespCommand.APPEND:
                     var appendValue = input.parseState.GetArgSliceByRef(0);
                     value.ShrinkSerializedLength(appendValue.Length);
-                    // Copy value to be appended to the newly allocated value buffer
                     appendValue.ReadOnlySpan.CopyTo(value.AsSpan());
-
                     CopyValueLengthToOutput(ref value, ref output, 0);
                     break;
                 case RespCommand.INCR:
@@ -326,15 +318,15 @@ public bool InPlaceUpdater(ref SpanByte key, ref RawStringInput input, ref SpanB
         // NOTE: In the below control flow if you decide to add a new command or modify a command such that it will now do an early return with TRUE, you must make sure you must reset etagState in FunctionState
         private bool InPlaceUpdaterWorker(ref SpanByte key, ref RawStringInput input, ref SpanByte value, ref SpanByteAndMemory output, ref RMWInfo rmwInfo, ref RecordInfo recordInfo)
         {
+            RespCommand cmd = input.header.cmd;
             // Expired data
             if (value.MetadataSize > 0 && input.header.CheckExpiry(value.ExtraMetadata))
             {
-                rmwInfo.Action = RMWAction.ExpireAndResume;
+                rmwInfo.Action = cmd is RespCommand.DELIFEXPIM ? RMWAction.ExpireAndStop : RMWAction.ExpireAndResume;
                 recordInfo.ClearHasETag();
                 return false;
             }
 
-            RespCommand cmd = input.header.cmd;
             bool hadRecordPreMutation = recordInfo.ETag;
             bool shouldUpdateEtag = hadRecordPreMutation;
             if (shouldUpdateEtag)
@@ -815,15 +807,7 @@ private bool InPlaceUpdaterWorker(ref SpanByte key, ref RawStringInput input, re
 
                     return false;
                 case RespCommand.DELIFEXPIM:
-                    // Only if the key has expired, will we delete it.
-                    if (value.MetadataSize > 0 && input.header.CheckExpiry(value.ExtraMetadata))
-                    {
-                        // setting the action and returning false will tombstone this record
-                        rmwInfo.Action = RMWAction.ExpireAndStop;
-                        // reset etag state that may have been initialized earlier,
-                        EtagState.ResetState(ref functionsState.etagState);
-                        return false;
-                    }
+                    // this is the case where it isn't expired
                     shouldUpdateEtag = false;
                     break;
                 default:

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

@@ -79,7 +79,7 @@ public GarnetStatus RMW_MainStore<TContext>(ref SpanByte key, ref RawStringInput
             if (status.IsPending)
                 CompletePendingForSession(ref status, ref output, ref context);
 
-            if (status.Found || status.Record.Created)
+            if (status.Found || status.Record.Created || status.Record.InPlaceUpdated)
                 return GarnetStatus.OK;
             else
                 return GarnetStatus.NOTFOUND;

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

@@ -285,7 +285,7 @@ private bool TryRevivifyInChain<TInput, TOutput, TContext, TSessionFunctionsWrap
                         // RMW uses GetInitialRecordSize because it has only the initial Input, not a Value
                         var (requiredSize, _, _) = hlog.GetRMWInitialRecordSize(ref key, ref input, sessionFunctions);
                         (ok, rmwInfo.UsedValueLength) = TryReinitializeTombstonedValue<TInput, TOutput, TContext, TSessionFunctionsWrapper>(sessionFunctions,
-                                ref srcRecordInfo, ref key, ref recordValue, requiredSize, recordLengths);
+                                ref srcRecordInfo, ref key, ref recordValue, requiredSize, recordLengths, stackCtx.recSrc.PhysicalAddress);
                     }
 
                     ref RevivificationStats stats = ref sessionFunctions.Ctx.RevivificationStats;
@@ -295,7 +295,8 @@ private bool TryRevivifyInChain<TInput, TOutput, TContext, TSessionFunctionsWrap
                         MarkPage(stackCtx.recSrc.LogicalAddress, sessionFunctions.Ctx);
                         pendingContext.recordInfo = srcRecordInfo;
                         pendingContext.logicalAddress = stackCtx.recSrc.LogicalAddress;
-                        status = OperationStatusUtils.AdvancedOpCode(OperationStatus.SUCCESS, StatusCode.InPlaceUpdatedRecord);
+                        // We "IPU'd" because we reused a tombstone, but since the record we have reused did not logically exist, we must also bubble up that the original key was not found (logically). OperationStatus.NOTFOUND bubbles up success but also indicates that the record was not found in the database.
+                        status = OperationStatusUtils.AdvancedOpCode(OperationStatus.NOTFOUND, StatusCode.InPlaceUpdatedRecord);
                         stats.inChainSuccesses++;
                         return true;
                     }

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

@@ -233,7 +233,7 @@ private bool TryRevivifyInChain<TInput, TOutput, TContext, TSessionFunctionsWrap
                         // Input is not included in record-length calculations for Upsert
                         var (requiredSize, _, _) = hlog.GetRecordSize(ref key, ref value);
                         (ok, upsertInfo.UsedValueLength) = TryReinitializeTombstonedValue<TInput, TOutput, TContext, TSessionFunctionsWrapper>(sessionFunctions,
-                                ref srcRecordInfo, ref key, ref recordValue, requiredSize, recordLengths);
+                                ref srcRecordInfo, ref key, ref recordValue, requiredSize, recordLengths, stackCtx.recSrc.PhysicalAddress);
                     }
 
                     ref RevivificationStats stats = ref sessionFunctions.Ctx.RevivificationStats;
@@ -244,7 +244,8 @@ private bool TryRevivifyInChain<TInput, TOutput, TContext, TSessionFunctionsWrap
                         MarkPage(stackCtx.recSrc.LogicalAddress, sessionFunctions.Ctx);
                         pendingContext.recordInfo = srcRecordInfo;
                         pendingContext.logicalAddress = stackCtx.recSrc.LogicalAddress;
-                        status = OperationStatusUtils.AdvancedOpCode(OperationStatus.SUCCESS, StatusCode.InPlaceUpdatedRecord);
+                        // Return NOTFOUND OperationStatus to indicate that the operation was successful but a previous record was not found.
+                        status = OperationStatusUtils.AdvancedOpCode(OperationStatus.NOTFOUND, StatusCode.InPlaceUpdatedRecord);
                         stats.inChainSuccesses++;
                         return true;
                     }

libs/storage/Tsavorite/cs/src/core/Index/Tsavorite/Implementation/Revivification/RecordLengths.cs

@@ -208,7 +208,7 @@ internal void SetTombstoneAndExtraValueLength(ref TValue recordValue, ref Record
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         internal (bool ok, int usedValueLength) TryReinitializeTombstonedValue<TInput, TOutput, TContext, TSessionFunctionsWrapper>(TSessionFunctionsWrapper sessionFunctions,
-                ref RecordInfo srcRecordInfo, ref TKey key, ref TValue recordValue, int requiredSize, (int usedValueLength, int fullValueLength, int allocatedSize) recordLengths)
+                ref RecordInfo srcRecordInfo, ref TKey key, ref TValue recordValue, int requiredSize, (int usedValueLength, int fullValueLength, int allocatedSize) recordLengths, long physicalAddress)
             where TSessionFunctionsWrapper : ISessionFunctionsWrapper<TKey, TValue, TInput, TOutput, TContext, TStoreFunctions, TAllocator>
         {
             if (RevivificationManager.IsFixedLength || recordLengths.allocatedSize < requiredSize)
@@ -219,13 +219,21 @@ internal void SetTombstoneAndExtraValueLength(ref TValue recordValue, ref Record
             var requiredValueLength = requiredSize - valueOffset;
             var minValueLength = requiredValueLength < recordLengths.usedValueLength ? requiredValueLength : recordLengths.usedValueLength;
 
+            // clears out the minimum space possible. So let's say we are shrinking our usage from 8 bytes to 3 bytes. This will clear only the bytes 4-8.
+            // if we are expanding this will not clear anything.
             ClearExtraValueSpace(ref srcRecordInfo, ref recordValue, minValueLength, recordLengths.fullValueLength);
-            storeFunctions.DisposeRecord(ref key, ref recordValue, DisposeReason.RevivificationFreeList);
 
             srcRecordInfo.ClearTombstone();
 
-            SetExtraValueLength(ref recordValue, ref srcRecordInfo, recordLengths.usedValueLength, recordLengths.fullValueLength);
-            return (true, hlog.GetValueLength(ref recordValue));
+            // for SpanByte, this will set the new length (payload + metadata).
+            hlog.GetAndInitializeValue(physicalAddress, physicalAddress + requiredSize);
+            // since the above sets the Length, we can use the below to get TotalSize which represents the UsedLength of a record
+            var newUsedValueLength = hlog.GetValueLength(ref recordValue);
+
+            // potentially sets filler, if the used value length is going to be under the full length by more than 4 bytes.
+            SetExtraValueLength(ref recordValue, ref srcRecordInfo, usedValueLength: newUsedValueLength, recordLengths.fullValueLength);
+
+            return (true, newUsedValueLength);
         }
 
         #endregion TombstonedRecords

libs/storage/Tsavorite/cs/src/core/VarLen/SpanByte.cs

@@ -394,28 +394,19 @@ public void CopyTo(ref SpanByte dst, long metadata = 0)
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public bool TrySafeCopyTo(ref SpanByte dst, int fullDestSize, long metadata = 0)
         {
-            // Need to account for extra metadata if current value does not have any.
+            // If the incoming caller wants to addMetadata and the destination does not already have metadata, the new length needs to account for it.
             var addMetadata = metadata > 0 && MetadataSize == 0;
 
             var newTotalSize = addMetadata ? TotalSize + sizeof(long) : TotalSize;
             if (fullDestSize < newTotalSize)
                 return false;
 
             var newLength = addMetadata ? Length + sizeof(long) : Length;
-            if (dst.Length < newLength)
-            {
-                // dst is shorter than src, but we have already verified there is enough extra value space to grow dst to store src.
-                dst.Length = newLength;
-                CopyTo(ref dst, metadata);
-            }
-            else
-            {
-                // dst length is equal or longer than src. We can adjust the length header on the serialized log, if we wish (here, we do).
-                // This method will also zero out the extra space to retain log scan correctness.
-                dst.ShrinkSerializedLength(newLength);
-                CopyTo(ref dst, metadata);
-                dst.Length = newLength;
-            }
+            dst.ShrinkSerializedLength(newLength);
+            // Note: If dst is shorter than src we have already verified there is enough extra value space to grow dst to store src.
+            dst.Length = newLength;
+            CopyTo(ref dst, metadata);
+
             return true;
         }