Allow Upsert to configure value length based on Input, in Tsavorite. (#…

…714) * Allow Upsert to configure value length based on Input, in Tsavorite. * fix to actually call api during creation of new record in upsert * fix test * fix ycsb value length
microsoft · Oct 10, 2024 · e3c8a20 · e3c8a20
1 parent ace7cb6
commit e3c8a20
Show file tree

Hide file tree

Showing 20 changed files with 79 additions and 2 deletions.
diff --git a/libs/server/Storage/Functions/MainStore/VarLenInputMethods.cs b/libs/server/Storage/Functions/MainStore/VarLenInputMethods.cs
@@ -255,5 +255,8 @@ public int GetRMWModifiedValueLength(ref SpanByte t, ref SpanByte input)
 
             return sizeof(int) + input.Length - RespInputHeader.Size;
         }
+
+        public int GetUpsertValueLength(ref SpanByte t, ref SpanByte input)
+            => t.TotalSize;
     }
 }
diff --git a/libs/server/Storage/Functions/ObjectStore/VarLenInputMethods.cs b/libs/server/Storage/Functions/ObjectStore/VarLenInputMethods.cs
@@ -22,5 +22,10 @@ public int GetRMWInitialValueLength(ref ObjectInput input)
         {
             throw new GarnetException("GetRMWInitialValueLength is not available on the object store");
         }
+
+        public int GetUpsertValueLength(ref IGarnetObject value, ref ObjectInput input)
+        {
+            throw new GarnetException("GetUpsertInitialValueLength is not available on the object store");
+        }
     }
 }
diff --git a/libs/storage/Tsavorite/cs/benchmark/YCSB.benchmark/SessionFunctions.cs b/libs/storage/Tsavorite/cs/benchmark/YCSB.benchmark/SessionFunctions.cs
@@ -82,6 +82,7 @@ public void PostInitialUpdater(ref Key key, ref Input input, ref Value value, re
 
         public int GetRMWModifiedValueLength(ref Value value, ref Input input) => 0;
         public int GetRMWInitialValueLength(ref Input input) => 0;
+        public int GetUpsertValueLength(ref Value value, ref Input input) => Value.Size;
 
         public void PostSingleDeleter(ref Key key, ref DeleteInfo deleteInfo) { }
 

diff --git a/libs/storage/Tsavorite/cs/benchmark/YCSB.benchmark/Value.cs b/libs/storage/Tsavorite/cs/benchmark/YCSB.benchmark/Value.cs
@@ -12,6 +12,8 @@ namespace Tsavorite.benchmark
     [StructLayout(LayoutKind.Explicit, Size = 8)]
     public struct Value
     {
+        public const int Size = 8;
+
         [FieldOffset(0)]
         public long value;
     }

diff --git a/libs/storage/Tsavorite/cs/src/core/Allocator/AllocatorScan.cs b/libs/storage/Tsavorite/cs/src/core/Allocator/AllocatorScan.cs
@@ -336,6 +336,7 @@ public void RMWCompletionCallback(ref TKey key, ref TInput input, ref TOutput ou
 
             public int GetRMWModifiedValueLength(ref TValue value, ref TInput input) => 0;
             public int GetRMWInitialValueLength(ref TInput input) => 0;
+            public int GetUpsertValueLength(ref TValue value, ref TInput input) => 0;
 
             public void ConvertOutputToHeap(ref TInput input, ref TOutput output) { }
         }

diff --git a/libs/storage/Tsavorite/cs/src/core/Allocator/BlittableAllocator.cs b/libs/storage/Tsavorite/cs/src/core/Allocator/BlittableAllocator.cs
@@ -104,6 +104,12 @@ public readonly (int actualSize, int allocatedSize, int keySize) GetRMWInitialRe
             where TSessionFunctionsWrapper : IVariableLengthInput<TValue, TInput>
             => BlittableAllocatorImpl<TKey, TValue, TStoreFunctions>.GetRMWInitialRecordSize(ref key, ref input, sessionFunctions);
 
+        /// <inheritdoc/>
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public readonly (int actualSize, int allocatedSize, int keySize) GetUpsertRecordSize<TInput, TSessionFunctionsWrapper>(ref TKey key, ref TValue value, ref TInput input, TSessionFunctionsWrapper sessionFunctions)
+            where TSessionFunctionsWrapper : IVariableLengthInput<TValue, TInput>
+            => BlittableAllocatorImpl<TKey, TValue, TStoreFunctions>.GetUpsertRecordSize(ref key, ref value, ref input, sessionFunctions);
+
         /// <inheritdoc/>
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public readonly (int actualSize, int allocatedSize, int keySize) GetRecordSize(ref TKey key, ref TValue value)

diff --git a/libs/storage/Tsavorite/cs/src/core/Allocator/BlittableAllocatorImpl.cs b/libs/storage/Tsavorite/cs/src/core/Allocator/BlittableAllocatorImpl.cs
@@ -102,6 +102,10 @@ public static (int actualSize, int allocatedSize, int keySize) GetRMWInitialReco
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public static (int actualSize, int allocatedSize, int keySize) GetRecordSize(ref TKey key, ref TValue value) => (RecordSize, RecordSize, KeySize);
 
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public static (int actualSize, int allocatedSize, int keySize) GetUpsertRecordSize<TInput, TSessionFunctionsWrapper>(ref TKey key, ref TValue value, ref TInput input, TSessionFunctionsWrapper sessionFunctions)
+            => (RecordSize, RecordSize, KeySize);
+
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public static int GetValueLength(ref TValue value) => ValueSize;
 

diff --git a/libs/storage/Tsavorite/cs/src/core/Allocator/GenericAllocator.cs b/libs/storage/Tsavorite/cs/src/core/Allocator/GenericAllocator.cs
@@ -97,6 +97,12 @@ public readonly (int actualSize, int allocatedSize, int keySize) GetRMWInitialRe
             where TSessionFunctionsWrapper : IVariableLengthInput<TValue, TInput>
             => _this.GetRMWInitialRecordSize(ref key, ref input, sessionFunctions);
 
+        /// <inheritdoc/>
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public readonly (int actualSize, int allocatedSize, int keySize) GetUpsertRecordSize<TInput, TSessionFunctionsWrapper>(ref TKey key, ref TValue value, ref TInput input, TSessionFunctionsWrapper sessionFunctions)
+            where TSessionFunctionsWrapper : IVariableLengthInput<TValue, TInput>
+            => _this.GetUpsertRecordSize(ref key, ref value, ref input, sessionFunctions);
+
         /// <inheritdoc/>
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public readonly (int actualSize, int allocatedSize, int keySize) GetRecordSize(ref TKey key, ref TValue value) => _this.GetRecordSize(ref key, ref value);

diff --git a/libs/storage/Tsavorite/cs/src/core/Allocator/GenericAllocatorImpl.cs b/libs/storage/Tsavorite/cs/src/core/Allocator/GenericAllocatorImpl.cs
@@ -151,6 +151,9 @@ internal ref TValue GetValue(long physicalAddress)
 
         internal (int actualSize, int allocatedSize, int keySize) GetRecordSize(ref TKey key, ref TValue value) => (RecordSize, RecordSize, KeySize);
 
+        internal (int actualSize, int allocatedSize, int keySize) GetUpsertRecordSize<TInput, TSessionFunctionsWrapper>(ref TKey key, ref TValue value, ref TInput input, TSessionFunctionsWrapper sessionFunctions)
+            => (RecordSize, RecordSize, KeySize);
+
         internal override bool TryComplete()
         {
             var b1 = objectLogDevice.TryComplete();

diff --git a/libs/storage/Tsavorite/cs/src/core/Allocator/IAllocator.cs b/libs/storage/Tsavorite/cs/src/core/Allocator/IAllocator.cs
@@ -31,6 +31,10 @@ AllocatorBase<TKey, TValue, TStoreFunctions, TAllocator> GetBase<TAllocator>()
         (int actualSize, int allocatedSize, int keySize) GetRMWInitialRecordSize<TInput, TSessionFunctionsWrapper>(ref TKey key, ref TInput input, TSessionFunctionsWrapper sessionFunctions)
             where TSessionFunctionsWrapper : IVariableLengthInput<TValue, TInput>;
 
+        /// <summary>Get record size required for the given <paramref name="key"/>, <paramref name="value"/>, and <paramref name="input"/></summary>
+        (int actualSize, int allocatedSize, int keySize) GetUpsertRecordSize<TInput, TSessionFunctionsWrapper>(ref TKey key, ref TValue value, ref TInput input, TSessionFunctionsWrapper sessionFunctions)
+            where TSessionFunctionsWrapper : IVariableLengthInput<TValue, TInput>;
+
         /// <summary>Get record size required for the given <paramref name="key"/> and <paramref name="value"/></summary>
         (int actualSize, int allocatedSize, int keySize) GetRecordSize(ref TKey key, ref TValue value);
 

diff --git a/libs/storage/Tsavorite/cs/src/core/Allocator/SpanByteAllocator.cs b/libs/storage/Tsavorite/cs/src/core/Allocator/SpanByteAllocator.cs
@@ -98,6 +98,12 @@ public readonly (int actualSize, int allocatedSize, int keySize) GetRMWInitialRe
             where TSessionFunctionsWrapper : IVariableLengthInput<SpanByte, TInput>
             => _this.GetRMWInitialRecordSize(ref key, ref input, sessionFunctions);
 
+        /// <inheritdoc/>
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public readonly (int actualSize, int allocatedSize, int keySize) GetUpsertRecordSize<TInput, TSessionFunctionsWrapper>(ref SpanByte key, ref SpanByte value, ref TInput input, TSessionFunctionsWrapper sessionFunctions)
+            where TSessionFunctionsWrapper : IVariableLengthInput<SpanByte, TInput>
+            => _this.GetUpsertRecordSize(ref key, ref value, ref input, sessionFunctions);
+
         /// <inheritdoc/>
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public readonly (int actualSize, int allocatedSize, int keySize) GetRecordSize(ref SpanByte key, ref SpanByte value) => _this.GetRecordSize(ref key, ref value);

diff --git a/libs/storage/Tsavorite/cs/src/core/Allocator/SpanByteAllocatorImpl.cs b/libs/storage/Tsavorite/cs/src/core/Allocator/SpanByteAllocatorImpl.cs
@@ -176,6 +176,14 @@ public int GetRequiredRecordSize(long physicalAddress, int availableBytes)
             return (actualSize, RoundUp(actualSize, Constants.kRecordAlignment), keySize);
         }
 
+        public (int actualSize, int allocatedSize, int keySize) GetUpsertRecordSize<TInput, TSessionFunctionsWrapper>(ref SpanByte key, ref SpanByte value, ref TInput input, TSessionFunctionsWrapper sessionFunctions)
+            where TSessionFunctionsWrapper : IVariableLengthInput<SpanByte, TInput>
+        {
+            int keySize = key.TotalSize;
+            var actualSize = RecordInfo.GetLength() + RoundUp(keySize, Constants.kRecordAlignment) + sessionFunctions.GetUpsertValueLength(ref value, ref input);
+            return (actualSize, RoundUp(actualSize, Constants.kRecordAlignment), keySize);
+        }
+
         public (int actualSize, int allocatedSize, int keySize) GetRecordSize(ref SpanByte key, ref SpanByte value)
         {
             int keySize = key.TotalSize;

diff --git a/libs/storage/Tsavorite/cs/src/core/ClientSession/SessionFunctionsWrapper.cs b/libs/storage/Tsavorite/cs/src/core/ClientSession/SessionFunctionsWrapper.cs
@@ -197,6 +197,9 @@ public void UnlockTransientShared(ref TKey key, ref OperationStackContext<TKey,
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public int GetRMWModifiedValueLength(ref TValue t, ref TInput input) => _clientSession.functions.GetRMWModifiedValueLength(ref t, ref input);
 
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public int GetUpsertValueLength(ref TValue t, ref TInput input) => _clientSession.functions.GetUpsertValueLength(ref t, ref input);
+
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public IHeapContainer<TInput> GetHeapContainer(ref TInput input)
         {

diff --git a/libs/storage/Tsavorite/cs/src/core/Compaction/LogCompactionFunctions.cs b/libs/storage/Tsavorite/cs/src/core/Compaction/LogCompactionFunctions.cs
@@ -52,6 +52,7 @@ public void RMWCompletionCallback(ref TKey key, ref TInput input, ref TOutput ou
 
         public int GetRMWModifiedValueLength(ref TValue value, ref TInput input) => 0;
         public int GetRMWInitialValueLength(ref TInput input) => 0;
+        public int GetUpsertValueLength(ref TValue value, ref TInput input) => _functions.GetUpsertValueLength(ref value, ref input);
 
         /// <summary>
         /// No reads during compaction

diff --git a/libs/storage/Tsavorite/cs/src/core/Index/Interfaces/ISessionFunctions.cs b/libs/storage/Tsavorite/cs/src/core/Index/Interfaces/ISessionFunctions.cs
@@ -189,6 +189,11 @@ public interface ISessionFunctions<TKey, TValue, TInput, TOutput, TContext>
         /// Initial expected length of value object when populated by RMW using given input
         /// </summary>
         int GetRMWInitialValueLength(ref TInput input);
+
+        /// <summary>
+        /// Length of resulting value object when performing Upsert of value using given input
+        /// </summary>
+        int GetUpsertValueLength(ref TValue value, ref TInput input);
         #endregion Variable-length value size
 
         /// <summary>

diff --git a/libs/storage/Tsavorite/cs/src/core/Index/Interfaces/SessionFunctionsBase.cs b/libs/storage/Tsavorite/cs/src/core/Index/Interfaces/SessionFunctionsBase.cs
@@ -57,6 +57,8 @@ public virtual void RMWCompletionCallback(ref TKey key, ref TInput input, ref TO
         public virtual int GetRMWModifiedValueLength(ref TValue value, ref TInput input) => throw new TsavoriteException("GetRMWModifiedValueLength is only available for SpanByte Functions");
         /// <inheritdoc/>
         public virtual int GetRMWInitialValueLength(ref TInput input) => throw new TsavoriteException("GetRMWInitialValueLength is only available for SpanByte Functions");
+        /// <inheritdoc/>
+        public virtual int GetUpsertValueLength(ref TValue value, ref TInput input) => throw new TsavoriteException("GetUpsertValueLength is only available for SpanByte Functions");
 
         /// <inheritdoc/>
         public virtual void ConvertOutputToHeap(ref TInput input, ref TOutput output) { }

diff --git a/libs/storage/Tsavorite/cs/src/core/Index/Tsavorite/Implementation/InternalUpsert.cs b/libs/storage/Tsavorite/cs/src/core/Index/Tsavorite/Implementation/InternalUpsert.cs
@@ -306,7 +306,7 @@ private OperationStatus CreateNewRecordUpsert<TInput, TOutput, TContext, TSessio
                                                                                              ref OperationStackContext<TKey, TValue, TStoreFunctions, TAllocator> stackCtx, ref RecordInfo srcRecordInfo)
             where TSessionFunctionsWrapper : ISessionFunctionsWrapper<TKey, TValue, TInput, TOutput, TContext, TStoreFunctions, TAllocator>
         {
-            var (actualSize, allocatedSize, keySize) = hlog.GetRecordSize(ref key, ref value);   // Input is not included in record-length calculations for Upsert
+            var (actualSize, allocatedSize, keySize) = hlog.GetUpsertRecordSize(ref key, ref value, ref input, sessionFunctions);
             AllocateOptions allocOptions = new()
             {
                 Recycle = true,

diff --git a/libs/storage/Tsavorite/cs/src/core/VarLen/IVariableLengthInput.cs b/libs/storage/Tsavorite/cs/src/core/VarLen/IVariableLengthInput.cs
@@ -19,5 +19,13 @@ public interface IVariableLengthInput<TValue, TInput>
         /// <param name="input"></param>
         /// <returns></returns>
         int GetRMWInitialValueLength(ref TInput input);
+
+        /// <summary>
+        /// Length of value object, when populated by Upsert using given value and input
+        /// </summary>
+        /// <param name="value"></param>
+        /// <param name="input"></param>
+        /// <returns></returns>
+        int GetUpsertValueLength(ref TValue value, ref TInput input);
     }
 }
diff --git a/libs/storage/Tsavorite/cs/src/core/VarLen/SpanByteFunctions.cs b/libs/storage/Tsavorite/cs/src/core/VarLen/SpanByteFunctions.cs
@@ -3,7 +3,6 @@
 
 using System.Buffers;
 using System.Runtime.CompilerServices;
-using static Tsavorite.core.Utility;
 
 namespace Tsavorite.core
 {
@@ -125,5 +124,12 @@ public override int GetRMWModifiedValueLength(ref SpanByte t, ref SpanByte input
 
         /// <inheritdoc/>
         public override int GetRMWInitialValueLength(ref SpanByte input) => input.TotalSize;
+
+        /// <summary>
+        /// Length of resulting object when doing Upsert with given value and input. Here we set the length to the
+        /// length of the provided value, ignoring input. You can provide a custom implementation for other cases.
+        /// </summary>
+        public override int GetUpsertValueLength(ref SpanByte t, ref SpanByte input)
+            => t.TotalSize;
     }
 }
diff --git a/libs/storage/Tsavorite/cs/test/ExpirationTests.cs b/libs/storage/Tsavorite/cs/test/ExpirationTests.cs
@@ -475,6 +475,9 @@ public override void ReadCompletionCallback(ref SpanByte key, ref ExpirationInpu
             /// <inheritdoc/>
             public override int GetRMWInitialValueLength(ref ExpirationInput input) => MinValueLen;
 
+            /// <inheritdoc/>
+            public override int GetUpsertValueLength(ref SpanByte value, ref ExpirationInput input) => value.TotalSize;
+
             // Read functions
             public override unsafe bool SingleReader(ref SpanByte key, ref ExpirationInput input, ref SpanByte value, ref ExpirationOutput output, ref ReadInfo readInfo)
             {