diff --git a/libs/cluster/Server/Migration/MigrateOperation.cs b/libs/cluster/Server/Migration/MigrateOperation.cs index bebf5a4637b..7e1e4570f51 100644 --- a/libs/cluster/Server/Migration/MigrateOperation.cs +++ b/libs/cluster/Server/Migration/MigrateOperation.cs @@ -103,9 +103,10 @@ public async Task TransmitSlotsAsync() var input = new UnifiedInput(RespCommand.MIGRATE); input.arg1 = session.NetworkBufferSettings.sendBufferSize - common.NetworkBufferSettings.SendBufferOverheadReserve; - VectorInput vectorInput = new(); - vectorInput.AlignmentExpected = true; // We're moving DiskANN sourced data, so alignment is expected - vectorInput.MaxMigrationHeapAllocationSize = session.NetworkBufferSettings.sendBufferSize - common.NetworkBufferSettings.SendBufferOverheadReserve; + VectorInput vectorInput = new() + { + IsMigrationRead = true, + }; foreach (var (ns, key, hasNs) in sketch.argSliceVector) { diff --git a/libs/server/InputHeader.cs b/libs/server/InputHeader.cs index deebfb6aa94..c2d32313e1d 100644 --- a/libs/server/InputHeader.cs +++ b/libs/server/InputHeader.cs @@ -3,7 +3,6 @@ using System; using System.Diagnostics; -using System.Diagnostics.CodeAnalysis; using System.Runtime.CompilerServices; using System.Runtime.InteropServices; using Tsavorite.core; @@ -609,24 +608,46 @@ public unsafe int DeserializeFrom(byte* src) /// /// Header for Garnet Main Store inputs but for Vector element r/w/d ops /// - public struct VectorInput : IStoreInput + public readonly struct VectorInput : IStoreInput { public int SerializedLength => throw new NotImplementedException(); - public int ReadDesiredSize { get; set; } + /// + /// True if the read value might not fit in the provided output buffer. + /// + /// If false, the output buffer is guaranteed to be correctly sized. + /// + public bool VariableSizedRead { get; init; } - public int WriteDesiredSize { get; set; } + /// + /// Size of write in bytes. + /// + /// If negative, this is an append to an existing value. + /// If positive, this is a create of a new value or an overwrite of an existing value. + /// + public int WriteDesiredSize { get; init; } - public int Index { get; set; } - public nint CallbackContext { get; set; } - public nint Callback { get; set; } + /// + /// If part of a batch operation, the zero-based index of that operation. + /// + public int Index { get; init; } - public bool AlignmentExpected { get; set; } + /// + /// Context to pass to , if any. + /// + /// This value is opaque to Garnet and should not be modified. + /// + public nint CallbackContext { get; init; } - [MemberNotNullWhen(returnValue: true, member: nameof(MaxMigrationHeapAllocationSize))] - public bool IsMigrationRead => MaxMigrationHeapAllocationSize != null; + /// + /// The native callback to invoke, if any, on the inplace record data. + /// + public nint Callback { get; init; } - public int? MaxMigrationHeapAllocationSize { get; set; } + /// + /// True if the read being performed is part of a migration operation. + /// + public bool IsMigrationRead { get; init; } public VectorInput() { diff --git a/libs/server/Resp/Vector/VectorManager.Callbacks.cs b/libs/server/Resp/Vector/VectorManager.Callbacks.cs index 301f159b504..a66af464304 100644 --- a/libs/server/Resp/Vector/VectorManager.Callbacks.cs +++ b/libs/server/Resp/Vector/VectorManager.Callbacks.cs @@ -107,6 +107,8 @@ private void AdvanceTo(int i) currentLen = *(int*)currentPtr; currentIndex = 0; + Debug.Assert((currentLen % 4) == 0, "Keys must be 4-byte aligned to preserve value alignment"); + if (i == 0) { return; @@ -144,10 +146,12 @@ public readonly void GetInput(int i, out VectorInput input) { Debug.Assert(i >= 0 && i < Count, "Trying to advance out of bounds"); - input = default; - input.CallbackContext = callbackContext; - input.Callback = (nint)callback; - input.Index = i; + input = new() + { + CallbackContext = callbackContext, + Callback = (nint)callback, + Index = i, + }; } /// @@ -222,10 +226,11 @@ nint dataCallbackContext [UnmanagedCallersOnly(CallConvs = [typeof(CallConvCdecl)])] private static unsafe byte WriteCallbackUnmanaged(ulong context, nint keyData, nuint keyLength, nint writeData, nuint writeLength) { + Debug.Assert((keyLength % 4) == 0, "Key must be 4-byte aligned to preserve value alignment"); + var keyWithNamespace = MakeVectorElementKey(context, keyData, keyLength); ref var ctx = ref ActiveThreadSession.vectorBasicContext; VectorInput input = new(); - input.AlignmentExpected = true; var valueSpan = SpanByte.FromPinnedPointer((byte*)writeData, (int)writeLength); VectorOutput outputSpan = new(); @@ -241,6 +246,8 @@ private static unsafe byte WriteCallbackUnmanaged(ulong context, nint keyData, n [UnmanagedCallersOnly(CallConvs = [typeof(CallConvCdecl)])] private static byte DeleteCallbackUnmanaged(ulong context, nint keyData, nuint keyLength) { + Debug.Assert((keyLength % 4) == 0, "Key must be 4-byte aligned to preserve value alignment"); + var keyWithNamespace = MakeVectorElementKey(context, keyData, keyLength); ref var ctx = ref ActiveThreadSession.vectorBasicContext; @@ -254,14 +261,18 @@ private static byte DeleteCallbackUnmanaged(ulong context, nint keyData, nuint k [UnmanagedCallersOnly(CallConvs = [typeof(CallConvCdecl)])] private static byte ReadModifyWriteCallbackUnmanaged(ulong context, nint keyData, nuint keyLength, nuint writeLength, nint dataCallback, nint dataCallbackContext) { + Debug.Assert((keyLength % 4) == 0, "Key must be 4-byte aligned to preserve value alignment"); + var keyWithNamespace = MakeVectorElementKey(context, keyData, keyLength); ref var ctx = ref ActiveThreadSession.vectorBasicContext; - VectorInput input = default; - input.Callback = dataCallback; - input.CallbackContext = dataCallbackContext; - input.WriteDesiredSize = (int)writeLength; + VectorInput input = new() + { + Callback = dataCallback, + CallbackContext = dataCallbackContext, + WriteDesiredSize = (int)writeLength, + }; var status = ctx.RMW(keyWithNamespace, ref input); if (status.IsPending) @@ -274,8 +285,10 @@ private static byte ReadModifyWriteCallbackUnmanaged(ulong context, nint keyData return status.IsCompletedSuccessfully ? (byte)1 : default; } - private static unsafe bool ReadSizeUnknown(ulong context, bool forceAlignment, ReadOnlySpan key, ref SpanByteAndMemory value) + private static unsafe bool ReadSizeUnknown(ulong context, ReadOnlySpan key, ref SpanByteAndMemory value) { + // We explicitly DO NOT check alignment here because we're always in managed code which doesn't care + #pragma warning disable IDE0302 // [...]-style collection initialization doesn't actually _guarantee_ stackalloc (or inline arrays), which we need here ReadOnlySpan nsBytes = stackalloc byte[1] { (byte)context }; #pragma warning restore IDE0302 @@ -286,20 +299,19 @@ private static unsafe bool ReadSizeUnknown(ulong context, bool forceAlignment, R while (true) { - VectorInput input = new(); - input.ReadDesiredSize = -1; + VectorInput input = new() + { + VariableSizedRead = true, + }; - // Sometimes we read DiskANN written data from the .NET side - // If that's the case, we need to pad for alignment even though .NET doesn't require it - input.AlignmentExpected = forceAlignment; fixed (byte* ptr = value.Span) { - VectorOutput asSpanByte = new(ptr, value.Length); + VectorOutput output = new(ptr, value.Length); - var status = ctx.Read(keyWithNamespace, ref input, ref asSpanByte); + var status = ctx.Read(keyWithNamespace, ref input, ref output); if (status.IsPending) { - CompletePending(ref status, ref input, ref asSpanByte, ref ctx); + CompletePending(ref status, ref output, ref ctx); } if (!status.Found) @@ -308,15 +320,16 @@ private static unsafe bool ReadSizeUnknown(ulong context, bool forceAlignment, R return false; } - if (input.ReadDesiredSize > asSpanByte.SpanByteAndMemory.Length) + var updateReadDesiredSize = output.UpdatedReadDesiredSize.GetValueOrDefault(-1); + if (updateReadDesiredSize > output.SpanByteAndMemory.Length) { value.Memory?.Dispose(); - var newAlloc = MemoryPool.Shared.Rent(input.ReadDesiredSize); + var newAlloc = MemoryPool.Shared.Rent(updateReadDesiredSize); value = new(newAlloc, newAlloc.Memory.Length); continue; } - value.Length = asSpanByte.SpanByteAndMemory.Length; + value.Length = output.SpanByteAndMemory.Length; return true; } } diff --git a/libs/server/Resp/Vector/VectorManager.ContextMetadata.cs b/libs/server/Resp/Vector/VectorManager.ContextMetadata.cs index 1dfb23075ff..e82b8c17f51 100644 --- a/libs/server/Resp/Vector/VectorManager.ContextMetadata.cs +++ b/libs/server/Resp/Vector/VectorManager.ContextMetadata.cs @@ -508,12 +508,15 @@ private void UpdateContextMetadata(ref VectorBasicContext ctx) // empty key is context metadata VectorElementKey key = new(nsBytes, []); - VectorInput input = default; - input.Callback = 0; - input.WriteDesiredSize = ContextMetadata.Size; + VectorInput input; unsafe { - input.CallbackContext = (nint)Unsafe.AsPointer(ref MemoryMarshal.GetReference(dataSpan)); + input = new() + { + Callback = 0, + WriteDesiredSize = ContextMetadata.Size, + CallbackContext = (nint)Unsafe.AsPointer(ref MemoryMarshal.GetReference(dataSpan)), + }; } var status = ctx.RMW(key, ref input); diff --git a/libs/server/Resp/Vector/VectorManager.Migration.cs b/libs/server/Resp/Vector/VectorManager.Migration.cs index a1188054f38..66dee4c5468 100644 --- a/libs/server/Resp/Vector/VectorManager.Migration.cs +++ b/libs/server/Resp/Vector/VectorManager.Migration.cs @@ -48,7 +48,6 @@ ReadOnlySpan value #endif VectorInput input = default; - input.AlignmentExpected = true; VectorOutput outputSpan = new(new SpanByteAndMemory()); VectorElementKey key = new(namespaceBytes[0..1], keyWithoutNamespace); diff --git a/libs/server/Resp/Vector/VectorManager.cs b/libs/server/Resp/Vector/VectorManager.cs index 7affdc610b9..4942a709a09 100644 --- a/libs/server/Resp/Vector/VectorManager.cs +++ b/libs/server/Resp/Vector/VectorManager.cs @@ -350,27 +350,6 @@ private static void CompletePending(ref Status status, ref VectorOutput output, completedOutputs.Dispose(); } - /// - /// As , but also propagates the - /// completed back to the caller. - /// - /// This is required for the unknown-size read path (): the Reader records the actual - /// value size on , and on the pending (disk) path that update lives on the - /// pending context's copy of the input. Without propagating it back the caller keeps its stale value and skips the - /// grow-and-retry, returning a truncated/uninitialized buffer. - /// - private static void CompletePending(ref Status status, ref VectorInput input, ref VectorOutput output, ref VectorBasicContext ctx) - { - _ = ctx.CompletePendingWithOutputs(out var completedOutputs, wait: true); - var more = completedOutputs.Next(); - Debug.Assert(more); - status = completedOutputs.Current.Status; - output = completedOutputs.Current.Output; - input = completedOutputs.Current.Input; - Debug.Assert(!completedOutputs.Next()); - completedOutputs.Dispose(); - } - private static void CompletePending(ref Status status, ref StringBasicContext ctx) { _ = ctx.CompletePendingWithOutputs(out var completedOutputs, wait: true); @@ -821,7 +800,7 @@ internal VectorManagerResult FetchSingleVectorElementAttributes(ReadOnlySpan indexValue, ReadOnlySpan var internalIdBytes = SpanByteAndMemory.FromPinnedSpan(internalId); try { - if (!ReadSizeUnknown(context | DiskANNService.InternalIdMap, forceAlignment: true, element, ref internalIdBytes)) + if (!ReadSizeUnknown(context | DiskANNService.InternalIdMap, element, ref internalIdBytes)) { return false; } @@ -964,7 +943,7 @@ internal bool TryGetEmbedding(ReadOnlySpan indexValue, ReadOnlySpan var asBytes = SpanByteAndMemory.FromPinnedSpan(asBytesSpan); try { - if (!ReadSizeUnknown(context | DiskANNService.FullVector, forceAlignment: true, internalId, ref asBytes)) + if (!ReadSizeUnknown(context | DiskANNService.FullVector, internalId, ref asBytes)) { return false; } diff --git a/libs/server/Storage/Functions/VectorStore/VectorSessionFunctions.cs b/libs/server/Storage/Functions/VectorStore/VectorSessionFunctions.cs index c67423ce2aa..b3f9006b670 100644 --- a/libs/server/Storage/Functions/VectorStore/VectorSessionFunctions.cs +++ b/libs/server/Storage/Functions/VectorStore/VectorSessionFunctions.cs @@ -17,8 +17,6 @@ namespace Garnet.server /// public readonly struct VectorSessionFunctions : ISessionFunctions { - private const int ValueAlignmentBytes = 4; - private readonly FunctionsState functionsState; private readonly ReadSessionState readSessionState; @@ -41,80 +39,83 @@ public readonly bool Reader(in TSourceLogRecord srcLogRecord, Debug.Assert(srcLogRecord.HasNamespace, "Should never write a non-namespaced value with VectorSessionFunctions"); Debug.Assert(srcLogRecord.NamespaceBytes.Length == 1, "Variable length namespaces not supported"); - var value = AlignOrPin(in srcLogRecord, ref input, out var pin); - try + var value = srcLogRecord.ValueSpan; + if (input.IsMigrationRead) { - if (input.IsMigrationRead) - { - Debug.Assert(input.Callback == 0, "No callback expected"); + Debug.Assert(input.Callback == 0, "No callback expected"); - // We can't ship the log record over because of alignment shenanigans - // TODO: When alignment is handled at the Tsavorite level, we CAN start shipping the log over like everything else + // Format for migration, including space for namespace - var neededSpace = - sizeof(int) + srcLogRecord.NamespaceBytes.Length + - sizeof(int) + srcLogRecord.KeyBytes.Length + - sizeof(int) + value.Length; + var neededSpace = + sizeof(int) + srcLogRecord.NamespaceBytes.Length + + sizeof(int) + srcLogRecord.KeyBytes.Length + + sizeof(int) + value.Length; - output.SpanByteAndMemory.EnsureHeapMemorySize(neededSpace); + output.SpanByteAndMemory.EnsureHeapMemorySize(neededSpace); - var writeTo = output.SpanByteAndMemory.Span; + var writeTo = output.SpanByteAndMemory.Span; - BinaryPrimitives.WriteInt32LittleEndian(writeTo, srcLogRecord.NamespaceBytes.Length); - writeTo = writeTo[sizeof(int)..]; - srcLogRecord.NamespaceBytes.CopyTo(writeTo); - writeTo = writeTo[srcLogRecord.NamespaceBytes.Length..]; + BinaryPrimitives.WriteInt32LittleEndian(writeTo, srcLogRecord.NamespaceBytes.Length); + writeTo = writeTo[sizeof(int)..]; + srcLogRecord.NamespaceBytes.CopyTo(writeTo); + writeTo = writeTo[srcLogRecord.NamespaceBytes.Length..]; - BinaryPrimitives.WriteInt32LittleEndian(writeTo, srcLogRecord.KeyBytes.Length); - writeTo = writeTo[sizeof(int)..]; - srcLogRecord.KeyBytes.CopyTo(writeTo); - writeTo = writeTo[srcLogRecord.KeyBytes.Length..]; + BinaryPrimitives.WriteInt32LittleEndian(writeTo, srcLogRecord.KeyBytes.Length); + writeTo = writeTo[sizeof(int)..]; + srcLogRecord.KeyBytes.CopyTo(writeTo); + writeTo = writeTo[srcLogRecord.KeyBytes.Length..]; - // Move value over _without_ any padding for alignment - BinaryPrimitives.WriteInt32LittleEndian(writeTo, value.Length); - writeTo = writeTo[sizeof(int)..]; - value.CopyTo(writeTo); + // Move value over _without_ any padding for alignment + BinaryPrimitives.WriteInt32LittleEndian(writeTo, value.Length); + writeTo = writeTo[sizeof(int)..]; + value.CopyTo(writeTo); - return true; - } + return true; + } - unsafe + unsafe + { + if (input.Callback != 0) { - if (input.Callback != 0) - { - var callback = (delegate* unmanaged[Cdecl, SuppressGCTransition])input.Callback; + var callback = (delegate* unmanaged[Cdecl, SuppressGCTransition])input.Callback; - var dataPtr = (nint)Unsafe.AsPointer(ref MemoryMarshal.GetReference(value)); - var dataLen = (nuint)value.Length; + var dataLen = (nuint)value.Length; + if (srcLogRecord.IsPinnedValue) + { + var dataPtr = (nint)Unsafe.AsPointer(ref MemoryMarshal.GetReference(value)); callback(input.Index, input.CallbackContext, dataPtr, dataLen); - return true; } - } - - if (input.ReadDesiredSize > 0) - { - Debug.Assert(output.SpanByteAndMemory.Length >= value.Length, "Should always have space for vector point reads"); - - output.SpanByteAndMemory.Length = value.Length; - value.CopyTo(output.SpanByteAndMemory.Span); - } - else - { - input.ReadDesiredSize = value.Length; - if (output.SpanByteAndMemory.Length >= value.Length) + else { - value.CopyTo(output.SpanByteAndMemory.Span); - output.SpanByteAndMemory.Length = value.Length; + fixed (byte* dataPtr = value) + { + callback(input.Index, input.CallbackContext, (nint)dataPtr, dataLen); + } } + + return true; } + } - return true; + if (!input.VariableSizedRead) + { + Debug.Assert(output.SpanByteAndMemory.Length >= value.Length, "Should always have space for vector point reads"); + + output.SpanByteAndMemory.Length = value.Length; + value.CopyTo(output.SpanByteAndMemory.Span); } - finally + else { - pin?.Free(); + output.UpdatedReadDesiredSize = value.Length; + if (output.SpanByteAndMemory.Length >= value.Length) + { + value.CopyTo(output.SpanByteAndMemory.Span); + output.SpanByteAndMemory.Length = value.Length; + } } + + return true; } /// @@ -130,17 +131,10 @@ public readonly bool InitialWriter(ref LogRecord logRecord, in RecordSizeInfo si Debug.Assert(logRecord.HasNamespace, "Should never write a non-namespaced value with VectorSessionFunctions"); Debug.Assert(logRecord.NamespaceBytes.Length == 1, "Variable length namespaces not supported"); - var value = AlignOrPin(in logRecord, ref input, out var pin); - try - { - srcValue.CopyTo(value); + var value = logRecord.ValueSpan; + srcValue.CopyTo(value); - return logRecord.TrySetContentLengths(logRecord.ValueSpan.Length, in sizeInfo); - } - finally - { - pin?.Free(); - } + return logRecord.TrySetContentLengths(logRecord.ValueSpan.Length, in sizeInfo); } /// @@ -158,17 +152,10 @@ public readonly bool InPlaceWriter(ref LogRecord logRecord, ref VectorInput inpu Debug.Assert(logRecord.HasNamespace, "Should never write a non-namespaced value with VectorSessionFunctions"); Debug.Assert(logRecord.NamespaceBytes.Length == 1, "Variable length namespaces not supported"); - var value = AlignOrPin(in logRecord, ref input, out var pin); - try - { - newValue.CopyTo(value); + var value = logRecord.ValueSpan; + newValue.CopyTo(value); - return true; - } - finally - { - pin?.Free(); - } + return true; } /// @@ -187,6 +174,8 @@ public readonly bool InPlaceWriter(ref LogRecord logRecord, re public readonly RecordFieldInfo GetRMWModifiedFieldInfo(in TSourceLogRecord srcLogRecord, ref VectorInput input) where TSourceLogRecord : ISourceLogRecord { + Debug.Assert((srcLogRecord.KeyBytes.Length % 4) == 0, "Keys must be 4-byte aligned to preserve value alignment"); + var value = srcLogRecord.ValueSpan; if (input.WriteDesiredSize < 0) @@ -195,17 +184,7 @@ public readonly RecordFieldInfo GetRMWModifiedFieldInfo(in TSo return new RecordFieldInfo() { KeySize = srcLogRecord.Key.Length, ValueSize = value.Length + (-input.WriteDesiredSize) }; } - var needsAlignmentPadding = input.AlignmentExpected || input.Callback != 0; - - // Constant size indicated - if (needsAlignmentPadding) - { - return new RecordFieldInfo() { KeySize = srcLogRecord.Key.Length, ValueSize = input.WriteDesiredSize + ValueAlignmentBytes }; - } - else - { - return new RecordFieldInfo() { KeySize = srcLogRecord.Key.Length, ValueSize = input.WriteDesiredSize }; - } + return new RecordFieldInfo() { KeySize = srcLogRecord.Key.Length, ValueSize = input.WriteDesiredSize }; } /// Initial expected length of value object when populated by RMW using given input @@ -215,23 +194,16 @@ public readonly RecordFieldInfo GetRMWInitialFieldInfo(TKey key, ref Vecto , allows ref struct #endif { - var effectiveWriteDesiredSize = input.WriteDesiredSize; + Debug.Assert((key.KeyBytes.Length % 4) == 0, "Keys must be 4-byte aligned to preserve value alignment"); - var needsAlignmentPadding = input.AlignmentExpected || input.Callback != 0; + var effectiveWriteDesiredSize = input.WriteDesiredSize; if (effectiveWriteDesiredSize < 0) { effectiveWriteDesiredSize = -effectiveWriteDesiredSize; } - if (!needsAlignmentPadding) - { - return new() { KeySize = key.KeyBytes.Length, ValueSize = effectiveWriteDesiredSize }; - } - else - { - return new() { KeySize = key.KeyBytes.Length, ValueSize = effectiveWriteDesiredSize + ValueAlignmentBytes }; - } + return new() { KeySize = key.KeyBytes.Length, ValueSize = effectiveWriteDesiredSize }; } /// Length of value object, when populated by Upsert using given value and input @@ -240,7 +212,11 @@ public readonly RecordFieldInfo GetUpsertFieldInfo(TKey key, ReadOnlySpan< #if NET9_0_OR_GREATER , allows ref struct #endif - => new() { KeySize = key.KeyBytes.Length, ValueSize = value.Length + ValueAlignmentBytes }; + { + Debug.Assert((key.KeyBytes.Length % 4) == 0, "Keys must be 4-byte aligned to preserve value alignment"); + + return new() { KeySize = key.KeyBytes.Length, ValueSize = value.Length }; + } /// Length of value object, when populated by Upsert using given value and input public readonly RecordFieldInfo GetUpsertFieldInfo(TKey key, IHeapObject value, ref VectorInput input) @@ -257,7 +233,11 @@ public readonly RecordFieldInfo GetUpsertFieldInfo(TKey , allows ref struct #endif where TSourceLogRecord : ISourceLogRecord - => new() { KeySize = key.KeyBytes.Length, ValueSize = inputLogRecord.ValueSpan.Length }; + { + Debug.Assert((key.KeyBytes.Length % 4) == 0, "Keys must be 4-byte aligned to preserve value alignment"); + + return new() { KeySize = key.KeyBytes.Length, ValueSize = inputLogRecord.ValueSpan.Length }; + } #endregion Variable Length #region InitialUpdater @@ -282,52 +262,54 @@ public readonly bool InitialUpdater(ref LogRecord logRecord, in RecordSizeInfo s Debug.Assert(logRecord.NamespaceBytes.Length == 1, "Variable length namespaces not supported"); var key = logRecord.Key; - var alignedValue = AlignOrPin(in logRecord, ref input, out var pin); - - try + var value = logRecord.ValueSpan; + if (input.Callback == 0) { + Debug.Assert(logRecord.NamespaceBytes.Length == 1 && logRecord.NamespaceBytes[0] == VectorManager.MetadataNamespace, "Should never write a non-namespaced value with VectorSessionFunctions"); + Debug.Assert(key.Length == 0, "Shouldn't have a non-zero key, expected to working on ContextMetadata"); - if (input.Callback == 0) + // Operating on ContextMetadata + + PinnedSpanByte newMetadataValue; + unsafe { - Debug.Assert(logRecord.NamespaceBytes.Length == 1 && logRecord.NamespaceBytes[0] == VectorManager.MetadataNamespace, "Should never write a non-namespaced value with VectorSessionFunctions"); - Debug.Assert(key.Length == 0, "Shouldn't have a non-zero key, expected to working on ContextMetadata"); + newMetadataValue = PinnedSpanByte.FromPinnedPointer((byte*)input.CallbackContext, VectorManager.ContextMetadata.Size); + } - // Operating on ContextMetadata + newMetadataValue.CopyTo(value); - PinnedSpanByte newMetadataValue; - unsafe - { - newMetadataValue = PinnedSpanByte.FromPinnedPointer((byte*)input.CallbackContext, VectorManager.ContextMetadata.Size); - } + return logRecord.TrySetContentLengths(logRecord.ValueSpan.Length, in sizeInfo); + } + else + { + Debug.Assert(input.WriteDesiredSize <= value.Length, "Insufficient space for initial update, this should never happen"); - newMetadataValue.CopyTo(alignedValue); + // Must explicitly 0 before passing if we're doing an initial update + value.Clear(); - return logRecord.TrySetContentLengths(logRecord.ValueSpan.Length, in sizeInfo); - } - else + unsafe { - Debug.Assert(input.WriteDesiredSize <= alignedValue.Length, "Insufficient space for initial update, this should never happen"); + // Callback takes: dataCallbackContext, dataPtr, dataLength + var callback = (delegate* unmanaged[Cdecl, SuppressGCTransition])input.Callback; - // Must explicitly 0 before passing if we're doing an initial update - alignedValue.Clear(); + var dataLen = (nuint)input.WriteDesiredSize; - unsafe + if (logRecord.IsPinnedValue) { - // Callback takes: dataCallbackContext, dataPtr, dataLength - var callback = (delegate* unmanaged[Cdecl, SuppressGCTransition])input.Callback; - - var dataPtr = (nint)Unsafe.AsPointer(ref MemoryMarshal.GetReference(alignedValue)); - var dataLen = (nuint)input.WriteDesiredSize; + var dataPtr = (nint)Unsafe.AsPointer(ref MemoryMarshal.GetReference(value)); callback(input.CallbackContext, dataPtr, dataLen); - - return logRecord.TrySetContentLengths(logRecord.ValueSpan.Length, in sizeInfo); } + else + { + fixed (byte* dataPtr = value) + { + callback(input.CallbackContext, (nint)dataPtr, dataLen); + } + } + + return logRecord.TrySetContentLengths(logRecord.ValueSpan.Length, in sizeInfo); } } - finally - { - pin?.Free(); - } } #endregion InitialUpdater @@ -346,68 +328,72 @@ public readonly bool CopyUpdater(in TSourceLogRecord srcLogRec var key = srcLogRecord.Key; - var oldValueAligned = AlignOrPin(in srcLogRecord, ref input, out var srcPin); - var newValueAligned = AlignOrPin(in dstLogRecord, ref input, out var dstPin); + var oldValue = srcLogRecord.ValueSpan; + var newValue = dstLogRecord.ValueSpan; - try + if (input.Callback == 0) { - if (input.Callback == 0) - { - // We're doing a Metadata or InProgressDelete update + // We're doing a Metadata or InProgressDelete update - Debug.Assert(srcLogRecord.NamespaceBytes[0] == VectorManager.MetadataNamespace, "Should be operating on special namespace"); - Debug.Assert(key.Length == 0, "Shouldn't have a non-zero key, expected to working on ContextMetadata"); + Debug.Assert(srcLogRecord.NamespaceBytes[0] == VectorManager.MetadataNamespace, "Should be operating on special namespace"); + Debug.Assert(key.Length == 0, "Shouldn't have a non-zero key, expected to working on ContextMetadata"); - // Doing a Metadata update - Debug.Assert(srcLogRecord.ValueSpan.Length == VectorManager.ContextMetadata.Size, "Should be ContextMetadata"); - Debug.Assert(dstLogRecord.ValueSpan.Length == VectorManager.ContextMetadata.Size, "Should be ContextMetadata"); - Debug.Assert(input.CallbackContext != 0, "Should have data on VectorInput"); + // Doing a Metadata update + Debug.Assert(srcLogRecord.ValueSpan.Length == VectorManager.ContextMetadata.Size, "Should be ContextMetadata"); + Debug.Assert(dstLogRecord.ValueSpan.Length == VectorManager.ContextMetadata.Size, "Should be ContextMetadata"); + Debug.Assert(input.CallbackContext != 0, "Should have data on VectorInput"); - ref readonly var oldMetadata = ref MemoryMarshal.Cast(oldValueAligned)[0]; - - PinnedSpanByte newMetadataValue; - unsafe - { - newMetadataValue = PinnedSpanByte.FromPinnedPointer((byte*)input.CallbackContext, VectorManager.ContextMetadata.Size); - } + ref readonly var oldMetadata = ref MemoryMarshal.Cast(oldValue)[0]; - ref readonly var newMetadata = ref MemoryMarshal.Cast(newMetadataValue.ReadOnlySpan)[0]; + PinnedSpanByte newMetadataValue; + unsafe + { + newMetadataValue = PinnedSpanByte.FromPinnedPointer((byte*)input.CallbackContext, VectorManager.ContextMetadata.Size); + } - if (newMetadata.Version < oldMetadata.Version) - { - rmwInfo.Action = RMWAction.CancelOperation; - return false; - } + ref readonly var newMetadata = ref MemoryMarshal.Cast(newMetadataValue.ReadOnlySpan)[0]; - newMetadataValue.CopyTo(newValueAligned); - return dstLogRecord.TrySetContentLengths(srcLogRecord.ValueSpan.Length, in sizeInfo); - } - else + if (newMetadata.Version < oldMetadata.Version) { - Debug.Assert(input.WriteDesiredSize <= newValueAligned.Length, "Insufficient space for copy update, this should never happen"); - Debug.Assert(input.WriteDesiredSize <= oldValueAligned.Length, "Insufficient space for copy update, this should never happen"); + rmwInfo.Action = RMWAction.CancelOperation; + return false; + } - oldValueAligned.CopyTo(newValueAligned); + newMetadataValue.CopyTo(newValue); + return dstLogRecord.TrySetContentLengths(srcLogRecord.ValueSpan.Length, in sizeInfo); + } + else + { + Debug.Assert(input.WriteDesiredSize <= newValue.Length, "Insufficient space for copy update, this should never happen"); + Debug.Assert(input.WriteDesiredSize <= oldValue.Length, "Insufficient space for copy update, this should never happen"); - unsafe - { - // Callback takes: dataCallbackContext, dataPtr, dataLength - var callback = (delegate* unmanaged[Cdecl, SuppressGCTransition])input.Callback; + oldValue.CopyTo(newValue); + + unsafe + { + // Callback takes: dataCallbackContext, dataPtr, dataLength + var callback = (delegate* unmanaged[Cdecl, SuppressGCTransition])input.Callback; - var dataPtr = (nint)Unsafe.AsPointer(ref MemoryMarshal.GetReference(newValueAligned)); - var dataLen = (nuint)input.WriteDesiredSize; + var dataLen = (nuint)input.WriteDesiredSize; + if (dstLogRecord.IsPinnedValue) + { + var dataPtr = (nint)Unsafe.AsPointer(ref MemoryMarshal.GetReference(newValue)); callback(input.CallbackContext, dataPtr, dataLen); } + else + { + fixed (byte* dataPtr = newValue) + { + callback(input.CallbackContext, (nint)dataPtr, dataLen); + } + } - return true; } - } - finally - { - srcPin?.Free(); - dstPin?.Free(); + + return true; + } } #endregion CopyUpdater @@ -421,85 +407,87 @@ public readonly bool InPlaceUpdater(ref LogRecord logRecord, ref VectorInput inp var key = logRecord.Key; - var alignedValue = AlignOrPin(in logRecord, ref input, out var pin); - try + var value = logRecord.ValueSpan; + if (input.Callback == 0) { - if (input.Callback == 0) - { - // We're doing a Metadata or InProgressDelete update - - Debug.Assert(logRecord.NamespaceBytes.Length == 1 && logRecord.NamespaceBytes[0] == VectorManager.MetadataNamespace, "Should be operating on special namespace"); - - if (key.Length == 0) - { - // Doing a Metadata update - Debug.Assert(alignedValue.Length >= VectorManager.ContextMetadata.Size, "Should be ContextMetadata"); - Debug.Assert(input.CallbackContext != 0, "Should have data on VectorInput"); + // We're doing a Metadata or InProgressDelete update - ref readonly var oldMetadata = ref MemoryMarshal.Cast(alignedValue)[0]; - - PinnedSpanByte newMetadataValue; - unsafe - { - newMetadataValue = PinnedSpanByte.FromPinnedPointer((byte*)input.CallbackContext, VectorManager.ContextMetadata.Size); - } + Debug.Assert(logRecord.NamespaceBytes.Length == 1 && logRecord.NamespaceBytes[0] == VectorManager.MetadataNamespace, "Should be operating on special namespace"); - ref readonly var newMetadata = ref MemoryMarshal.Cast(newMetadataValue.ReadOnlySpan)[0]; + if (key.Length == 0) + { + // Doing a Metadata update + Debug.Assert(value.Length >= VectorManager.ContextMetadata.Size, "Should be ContextMetadata"); + Debug.Assert(input.CallbackContext != 0, "Should have data on VectorInput"); - if (newMetadata.Version < oldMetadata.Version) - { - rmwInfo.Action = RMWAction.CancelOperation; - return false; - } + ref readonly var oldMetadata = ref MemoryMarshal.Cast(value)[0]; - newMetadataValue.CopyTo(alignedValue); - return true; - } - else + PinnedSpanByte newMetadataValue; + unsafe { - // Doing an InProgressDelete update - Debug.Assert(input.CallbackContext != 0, "Should have data on VectorInput"); - Debug.Assert(key.Length == 1 && key[0] == 1, "Should be working on InProgressDeletes"); - - Span inProgressDeleteUpdateData; - bool adding; + newMetadataValue = PinnedSpanByte.FromPinnedPointer((byte*)input.CallbackContext, VectorManager.ContextMetadata.Size); + } - unsafe - { - var len = BinaryPrimitives.ReadInt32LittleEndian(new Span((byte*)input.CallbackContext + sizeof(long), sizeof(int))); - adding = len > 0; - if (!adding) - { - len = -len; - } - - inProgressDeleteUpdateData = new Span((byte*)input.CallbackContext, sizeof(ulong) + sizeof(int) + len); - } + ref readonly var newMetadata = ref MemoryMarshal.Cast(newMetadataValue.ReadOnlySpan)[0]; - return true; + if (newMetadata.Version < oldMetadata.Version) + { + rmwInfo.Action = RMWAction.CancelOperation; + return false; } + + newMetadataValue.CopyTo(value); + return true; } else { - Debug.Assert(input.WriteDesiredSize <= alignedValue.Length, "Insufficient space for inplace update, this should never happen"); + // Doing an InProgressDelete update + Debug.Assert(input.CallbackContext != 0, "Should have data on VectorInput"); + Debug.Assert(key.Length == 1 && key[0] == 1, "Should be working on InProgressDeletes"); + + Span inProgressDeleteUpdateData; + bool adding; unsafe { - // Callback takes: dataCallbackContext, dataPtr, dataLength - var callback = (delegate* unmanaged[Cdecl, SuppressGCTransition])input.Callback; - - var dataPtr = (nint)Unsafe.AsPointer(ref MemoryMarshal.GetReference(alignedValue)); - var dataLen = (nuint)input.WriteDesiredSize; + var len = BinaryPrimitives.ReadInt32LittleEndian(new Span((byte*)input.CallbackContext + sizeof(long), sizeof(int))); + adding = len > 0; + if (!adding) + { + len = -len; + } - callback(input.CallbackContext, dataPtr, dataLen); + inProgressDeleteUpdateData = new Span((byte*)input.CallbackContext, sizeof(ulong) + sizeof(int) + len); } return true; } } - finally + else { - pin?.Free(); + Debug.Assert(input.WriteDesiredSize <= value.Length, "Insufficient space for inplace update, this should never happen"); + + unsafe + { + // Callback takes: dataCallbackContext, dataPtr, dataLength + var callback = (delegate* unmanaged[Cdecl, SuppressGCTransition])input.Callback; + + var dataLen = (nuint)input.WriteDesiredSize; + if (logRecord.IsPinnedValue) + { + var dataPtr = (nint)Unsafe.AsPointer(ref MemoryMarshal.GetReference(value)); + callback(input.CallbackContext, dataPtr, dataLen); + } + else + { + fixed (byte* dataPtr = value) + { + callback(input.CallbackContext, (nint)dataPtr, dataLen); + } + } + } + + return true; } } #endregion InPlaceUpdater @@ -540,72 +528,6 @@ private static TReturn ObjectOperationsNotExpected([CallerMemberName] s private static TReturn LogRecordOperationsNotExpected([CallerMemberName] string callerName = null, [CallerLineNumber] int lineNum = -1) => throw new InvalidOperationException($"LogRecord related operations are not expected, was: {callerName} on {lineNum}"); - // TODO: Remove all this alignment hackery when Tsavorite can enforce it - - [MethodImpl(MethodImplOptions.AggressiveInlining)] - internal static unsafe Span AlignOrPin(in TSourceLogRecord logRecord, ref VectorInput input, out GCHandle? pin) - where TSourceLogRecord : ISourceLogRecord - { - var maybeUnaligned = logRecord.ValueSpan; - - // Alignment is expected if we're passing to DiskANN or Garnet code explicitly requested it - var inputRequiresAligment = input.AlignmentExpected || input.Callback != 0; - - if (inputRequiresAligment) - { - if (logRecord.IsPinnedValue) - { - // LogRecord itself is in POH, but value might not be aligned so we need to do some checking - - Span ret; - - var leading = (nint)Unsafe.AsPointer(ref MemoryMarshal.GetReference(maybeUnaligned)) % 4; - if (leading == 0) - { - ret = maybeUnaligned[..^ValueAlignmentBytes]; - } - else - { - var skip = (int)(ValueAlignmentBytes - leading); - var tail = ValueAlignmentBytes - skip; - ret = maybeUnaligned[skip..^tail]; - } - - AssertAlignment(ret); - - pin = null; - return ret; - } - else - { - // Value isn't in log record, it's on the (presumably unpinned) heap as a byte[] - // - // This guarantees it's aligned, but it might move during any callback so pin - - pin = logRecord.ValueOverflow.Pin(); - - // We over allocated (we don't know how Tsavorite is going to place the value in advance) so trim the extra allocation off the end. - var ret = maybeUnaligned[..^ValueAlignmentBytes]; - - AssertAlignment(ret); - - return ret; - } - } - else - { - pin = null; - return maybeUnaligned; - } - } - - [Conditional("DEBUG")] - private static unsafe void AssertAlignment(ReadOnlySpan aligned) - { - var ptr = (nint)Unsafe.AsPointer(ref MemoryMarshal.GetReference(aligned)); - Debug.Assert((ptr % ValueAlignmentBytes) == 0, "Must guarantee 4-byte alignment before invoking callback"); - } - #region Post operation callbacks /// public readonly void PostInitialWriter(ref LogRecord logRecord, in RecordSizeInfo sizeInfo, ref VectorInput input, ReadOnlySpan srcValue, ref VectorOutput output, ref UpsertInfo upsertInfo) diff --git a/libs/server/VectorOutput.cs b/libs/server/VectorOutput.cs index fd9df75d6fe..9718f3e5361 100644 --- a/libs/server/VectorOutput.cs +++ b/libs/server/VectorOutput.cs @@ -18,6 +18,11 @@ public struct VectorOutput /// public SpanByteAndMemory SpanByteAndMemory; + /// + /// If a call needs a larger output buffer than was provided, it is stored here. + /// + public int? UpdatedReadDesiredSize { get; set; } + public VectorOutput() => SpanByteAndMemory = new(null); public VectorOutput(SpanByteAndMemory span) => SpanByteAndMemory = span; diff --git a/test/standalone/Garnet.test.vectorset/RespVectorSetTests.cs b/test/standalone/Garnet.test.vectorset/RespVectorSetTests.cs index cfb6b50dc4c..821b60c1b4c 100644 --- a/test/standalone/Garnet.test.vectorset/RespVectorSetTests.cs +++ b/test/standalone/Garnet.test.vectorset/RespVectorSetTests.cs @@ -1389,9 +1389,11 @@ public unsafe void VectorReadBatchVariants() { // Single key, 4 byte keys { - VectorInput input = default; - input.Callback = 5678; - input.CallbackContext = 9012; + VectorInput input = new() + { + Callback = 5678, + CallbackContext = 9012, + }; ReadOnlySpan namespaceBytes = stackalloc byte[1] { 64 }; @@ -1439,9 +1441,11 @@ public unsafe void VectorReadBatchVariants() // Multiple keys, 4 byte keys { - VectorInput input = default; - input.Callback = 5678; - input.CallbackContext = 9012; + VectorInput input = new() + { + Callback = 5678, + CallbackContext = 9012, + }; ReadOnlySpan namespaceBytes = stackalloc byte[1] { 32 }; @@ -1493,9 +1497,11 @@ public unsafe void VectorReadBatchVariants() // Multiple keys, 4 byte keys, random order { - VectorInput input = default; - input.Callback = 5678; - input.CallbackContext = 9012; + VectorInput input = new() + { + Callback = 5678, + CallbackContext = 9012, + }; ReadOnlySpan namespaceBytes = stackalloc byte[1] { 16 }; @@ -1546,9 +1552,11 @@ public unsafe void VectorReadBatchVariants() // Single key, variable length { - VectorInput input = default; - input.Callback = 5678; - input.CallbackContext = 9012; + VectorInput input = new() + { + Callback = 5678, + CallbackContext = 9012, + }; ReadOnlySpan namespaceBytes = stackalloc byte[1] { 8 }; @@ -1616,9 +1624,11 @@ public unsafe void VectorReadBatchVariants() // Multiple keys, variable length { - VectorInput input = default; - input.Callback = 5678; - input.CallbackContext = 9012; + VectorInput input = new() + { + Callback = 5678, + CallbackContext = 9012, + }; ReadOnlySpan namespaceBytes = stackalloc byte[1] { 4 }; @@ -1749,20 +1759,22 @@ public unsafe void VectorReadBatchVariants() // Multiple keys, variable length, random access { - VectorInput input = default; - input.Callback = 5678; - input.CallbackContext = 9012; + VectorInput input = new() + { + Callback = 5678, + CallbackContext = 9012, + }; ReadOnlySpan namespaceBytes = stackalloc byte[1] { 2 }; - var key0 = "hello"u8.ToArray(); + var key0 = "buzz"u8.ToArray(); var key1 = "fizz"u8.ToArray(); - var key2 = "the quick brown fox jumps over the lazy dog"u8.ToArray(); + var key2 = "the quick brown fox jumps over the lazy dog."u8.ToArray(); var key3 = "CF29E323-E376-4BC4-AB63-FCFD371EB445"u8.ToArray(); var key4 = Array.Empty(); var key5 = new byte[] { 1 }; - var key6 = new byte[] { 2, 3 }; - var key7 = new byte[] { 4, 5, 6 }; + var key6 = new byte[] { 2, 3, 4, 5 }; + var key7 = new byte[] { 6, 7, 8, 9, 10, 11, 12, 13 }; var data = MemoryMarshal.Cast([key0.Length]) .ToArray()