Preview Vector Set API support

Directory.Packages.props

@@ -30,5 +30,6 @@
     <PackageVersion Include="System.Numerics.Tensors" Version="9.0.9" />
     <PackageVersion Include="Microsoft.Extensions.Hosting" Version="9.0.8" />
     <PackageVersion Include="Microsoft.Extensions.Hosting.WindowsServices" Version="9.0.8" />
+    <PackageVersion Include="diskann-garnet" Version="1.0.15" />
   </ItemGroup>
 </Project>

libs/client/ClientSession/GarnetClientSessionMigrationExtensions.cs

@@ -25,6 +25,7 @@ public sealed unsafe partial class GarnetClientSession : IServerHook, IMessageCo
 
         static ReadOnlySpan<byte> MAIN_STORE => "SSTORE"u8;
         static ReadOnlySpan<byte> OBJECT_STORE => "OSTORE"u8;
+        static ReadOnlySpan<byte> VECTOR_STORE => "VSTORE"u8;
         static ReadOnlySpan<byte> T => "T"u8;
         static ReadOnlySpan<byte> F => "F"u8;
 
@@ -170,14 +171,30 @@ public Task<string> SetSlotRange(Memory<byte> state, string nodeid, List<(int, i
         /// <param name="sourceNodeId"></param>
         /// <param name="replace"></param>
         /// <param name="isMainStore"></param>
-        public void SetClusterMigrateHeader(string sourceNodeId, bool replace, bool isMainStore)
+        public void SetClusterMigrateHeader(string sourceNodeId, bool replace, bool isMainStore, bool isVectorSets)
         {
             currTcsIterationTask = new TaskCompletionSource<string>(TaskCreationOptions.RunContinuationsAsynchronously);
             tcsQueue.Enqueue(currTcsIterationTask);
             curr = offset;
             this.isMainStore = isMainStore;
             this.ist = IncrementalSendType.MIGRATE;
-            var storeType = isMainStore ? MAIN_STORE : OBJECT_STORE;
+            ReadOnlySpan<byte> storeType;
+            if (isMainStore)
+            {
+                if (isVectorSets)
+                {
+                    storeType = VECTOR_STORE;
+                }
+                else
+                {
+                    storeType = MAIN_STORE;
+                }
+            }
+            else
+            {
+                storeType = OBJECT_STORE;
+            }
+
             var replaceOption = replace ? T : F;
 
             var arraySize = 6;
@@ -249,7 +266,7 @@ public void SetClusterMigrateHeader(string sourceNodeId, bool replace, bool isMa
         /// <returns></returns>
         public Task<string> CompleteMigrate(string sourceNodeId, bool replace, bool isMainStore)
         {
-            SetClusterMigrateHeader(sourceNodeId, replace, isMainStore);
+            SetClusterMigrateHeader(sourceNodeId, replace, isMainStore, isVectorSets: false);
 
             Debug.Assert(end - curr >= 2);
             *curr++ = (byte)'\r';

libs/cluster/Server/ClusterManager.cs

@@ -239,22 +239,27 @@ public string GetInfo()
         public static string GetRange(int[] slots)
         {
             var range = "> ";
-            var start = slots[0];
-            var end = slots[0];
-            for (var i = 1; i < slots.Length + 1; i++)
+            if (slots.Length >= 1)
             {
-                if (i < slots.Length && slots[i] == end + 1)
-                    end = slots[i];
-                else
+
+                var start = slots[0];
+                var end = slots[0];
+                for (var i = 1; i < slots.Length + 1; i++)
                 {
-                    range += $"{start}-{end} ";
-                    if (i < slots.Length)
-                    {
-                        start = slots[i];
+                    if (i < slots.Length && slots[i] == end + 1)
                         end = slots[i];
+                    else
+                    {
+                        range += $"{start}-{end} ";
+                        if (i < slots.Length)
+                        {
+                            start = slots[i];
+                            end = slots[i];
+                        }
                     }
                 }
             }
+
             return range;
         }
 

libs/cluster/Server/ClusterManagerSlotState.cs

@@ -17,7 +17,10 @@ namespace Garnet.cluster
             SpanByteAllocator<StoreFunctions<SpanByte, SpanByte, SpanByteComparer, SpanByteRecordDisposer>>>,
         BasicContext<byte[], IGarnetObject, ObjectInput, GarnetObjectStoreOutput, long, ObjectSessionFunctions,
             /* ObjectStoreFunctions */ StoreFunctions<byte[], IGarnetObject, ByteArrayKeyComparer, DefaultRecordDisposer<byte[], IGarnetObject>>,
-            GenericAllocator<byte[], IGarnetObject, StoreFunctions<byte[], IGarnetObject, ByteArrayKeyComparer, DefaultRecordDisposer<byte[], IGarnetObject>>>>>;
+            GenericAllocator<byte[], IGarnetObject, StoreFunctions<byte[], IGarnetObject, ByteArrayKeyComparer, DefaultRecordDisposer<byte[], IGarnetObject>>>>,
+        BasicContext<SpanByte, SpanByte, VectorInput, SpanByte, long, VectorSessionFunctions,
+            /* VectorStoreFunctions */ StoreFunctions<SpanByte, SpanByte, SpanByteComparer, SpanByteRecordDisposer>,
+            SpanByteAllocator<StoreFunctions<SpanByte, SpanByte, SpanByteComparer, SpanByteRecordDisposer>>>>;
 
     /// <summary>
     /// Cluster manager

libs/cluster/Server/ClusterProvider.cs

@@ -15,12 +15,21 @@
 
 namespace Garnet.cluster
 {
+    using BasicContext = BasicContext<SpanByte, SpanByte, RawStringInput, SpanByteAndMemory, long, MainSessionFunctions,
+            /* MainStoreFunctions */ StoreFunctions<SpanByte, SpanByte, SpanByteComparer, SpanByteRecordDisposer>,
+            SpanByteAllocator<StoreFunctions<SpanByte, SpanByte, SpanByteComparer, SpanByteRecordDisposer>>>;
+
     using BasicGarnetApi = GarnetApi<BasicContext<SpanByte, SpanByte, RawStringInput, SpanByteAndMemory, long, MainSessionFunctions,
             /* MainStoreFunctions */ StoreFunctions<SpanByte, SpanByte, SpanByteComparer, SpanByteRecordDisposer>,
             SpanByteAllocator<StoreFunctions<SpanByte, SpanByte, SpanByteComparer, SpanByteRecordDisposer>>>,
         BasicContext<byte[], IGarnetObject, ObjectInput, GarnetObjectStoreOutput, long, ObjectSessionFunctions,
             /* ObjectStoreFunctions */ StoreFunctions<byte[], IGarnetObject, ByteArrayKeyComparer, DefaultRecordDisposer<byte[], IGarnetObject>>,
-            GenericAllocator<byte[], IGarnetObject, StoreFunctions<byte[], IGarnetObject, ByteArrayKeyComparer, DefaultRecordDisposer<byte[], IGarnetObject>>>>>;
+            GenericAllocator<byte[], IGarnetObject, StoreFunctions<byte[], IGarnetObject, ByteArrayKeyComparer, DefaultRecordDisposer<byte[], IGarnetObject>>>>,
+        BasicContext<SpanByte, SpanByte, VectorInput, SpanByte, long, VectorSessionFunctions,
+            /* VectorStoreFunctions */ StoreFunctions<SpanByte, SpanByte, SpanByteComparer, SpanByteRecordDisposer>,
+            SpanByteAllocator<StoreFunctions<SpanByte, SpanByte, SpanByteComparer, SpanByteRecordDisposer>>>>;
+
+    using VectorContext = BasicContext<SpanByte, SpanByte, VectorInput, SpanByte, long, VectorSessionFunctions, StoreFunctions<SpanByte, SpanByte, SpanByteComparer, SpanByteRecordDisposer>, SpanByteAllocator<StoreFunctions<SpanByte, SpanByte, SpanByteComparer, SpanByteRecordDisposer>>>;
 
     /// <summary>
     /// Cluster provider
@@ -100,8 +109,8 @@ public void Start()
         }
 
         /// <inheritdoc />
-        public IClusterSession CreateClusterSession(TransactionManager txnManager, IGarnetAuthenticator authenticator, UserHandle userHandle, GarnetSessionMetrics garnetSessionMetrics, BasicGarnetApi basicGarnetApi, INetworkSender networkSender, ILogger logger = null)
-            => new ClusterSession(this, txnManager, authenticator, userHandle, garnetSessionMetrics, basicGarnetApi, networkSender, logger);
+        public IClusterSession CreateClusterSession(TransactionManager txnManager, IGarnetAuthenticator authenticator, UserHandle userHandle, GarnetSessionMetrics garnetSessionMetrics, BasicGarnetApi basicGarnetApi, BasicContext basicContext, VectorContext vectorContext, INetworkSender networkSender, ILogger logger = null)
+            => new ClusterSession(this, txnManager, authenticator, userHandle, garnetSessionMetrics, basicGarnetApi, basicContext, vectorContext, networkSender, logger);
 
         /// <inheritdoc />
         public void UpdateClusterAuth(string clusterUsername, string clusterPassword)

libs/cluster/Server/Migration/MigrateOperation.cs

@@ -2,9 +2,11 @@
 // Licensed under the MIT license.
 
 using System;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
 using Garnet.client;
 using Garnet.server;
+using Microsoft.Extensions.Logging;
 using Tsavorite.core;
 
 namespace Garnet.cluster
@@ -18,16 +20,28 @@ internal sealed partial class MigrateOperation
             public MainStoreScan mss;
             public ObjectStoreScan oss;
 
+            public readonly ConcurrentDictionary<byte[], byte[]> vectorSetsIndexKeysToMigrate;
+#if NET9_0_OR_GREATER
+            private readonly ConcurrentDictionary<byte[], byte[]>.AlternateLookup<ReadOnlySpan<byte>> vectorSetsIndexKeysToMigrateLookup;
+#endif
+
             readonly MigrateSession session;
             readonly GarnetClientSession gcs;
             readonly LocalServerSession localServerSession;
 
             public GarnetClientSession Client => gcs;
 
+            public IEnumerable<KeyValuePair<byte[], byte[]>> VectorSets => vectorSetsIndexKeysToMigrate;
+
             public void ThrowIfCancelled() => session._cts.Token.ThrowIfCancellationRequested();
 
             public bool Contains(int slot) => session._sslots.Contains(slot);
 
+            public bool ContainsNamespace(ulong ns) => session._namespaces?.Contains(ns) ?? false;
+
+            public void EncounteredVectorSet(byte[] key, byte[] value)
+            => vectorSetsIndexKeysToMigrate.TryAdd(key, value);
+
             public MigrateOperation(MigrateSession session, Sketch sketch = null, int batchSize = 1 << 18)
             {
                 this.session = session;
@@ -37,6 +51,10 @@ public MigrateOperation(MigrateSession session, Sketch sketch = null, int batchS
                 mss = new MainStoreScan(this);
                 oss = new ObjectStoreScan(this);
                 keysToDelete = [];
+                vectorSetsIndexKeysToMigrate = new(ByteArrayComparer.Instance);
+#if NET9_0_OR_GREATER
+                vectorSetsIndexKeysToMigrateLookup = vectorSetsIndexKeysToMigrate.GetAlternateLookup<ReadOnlySpan<byte>>();
+#endif
             }
 
             public bool Initialize()
@@ -72,7 +90,7 @@ public void Scan(StoreType storeType, ref long currentAddress, long endAddress)
             /// </summary>
             /// <param name="storeType"></param>
             /// <returns></returns>
-            public bool TrasmitSlots(StoreType storeType)
+            public bool TransmitSlots(StoreType storeType)
             {
                 var bufferSize = 1 << 10;
                 SectorAlignedMemory buffer = new(bufferSize, 1);
@@ -87,7 +105,7 @@ public bool TrasmitSlots(StoreType storeType)
                     {
                         foreach (var key in sketch.argSliceVector)
                         {
-                            var spanByte = key.SpanByte;
+                            var spanByte = key;
                             if (!session.WriteOrSendMainStoreKeyValuePair(gcs, localServerSession, ref spanByte, ref input, ref o, out _))
                                 return false;
 
@@ -117,7 +135,10 @@ public bool TrasmitSlots(StoreType storeType)
                 return true;
             }
 
-            public bool TransmitKeys(StoreType storeType)
+            /// <summary>
+            /// Move keys in sketch out of the given store, UNLESS they are also in <paramref name="vectorSetKeysToIgnore"/>.
+            /// </summary>
+            public bool TransmitKeys(StoreType storeType, Dictionary<byte[], byte[]> vectorSetKeysToIgnore)
             {
                 var bufferSize = 1 << 10;
                 SectorAlignedMemory buffer = new(bufferSize, 1);
@@ -131,12 +152,30 @@ public bool TransmitKeys(StoreType storeType)
                     var keys = sketch.Keys;
                     if (storeType == StoreType.Main)
                     {
+#if NET9_0_OR_GREATER
+                        var ignoreLookup = vectorSetKeysToIgnore.GetAlternateLookup<ReadOnlySpan<byte>>();
+#endif
+
                         for (var i = 0; i < keys.Count; i++)
                         {
                             if (keys[i].Item2)
                                 continue;
 
                             var spanByte = keys[i].Item1.SpanByte;
+
+                            // Don't transmit if a Vector Set
+                            var isVectorSet =
+                                vectorSetKeysToIgnore.Count > 0 &&
+#if NET9_0_OR_GREATER
+                                ignoreLookup.ContainsKey(spanByte.AsReadOnlySpan());
+#else
+                                vectorSetKeysToIgnore.ContainsKey(spanByte.ToByteArray());
+#endif
+                            if (isVectorSet)
+                            {
+                                continue;
+                            }
+
                             if (!session.WriteOrSendMainStoreKeyValuePair(gcs, localServerSession, ref spanByte, ref input, ref o, out var status))
                                 return false;
 
@@ -158,8 +197,8 @@ public bool TransmitKeys(StoreType storeType)
                             if (keys[i].Item2)
                                 continue;
 
-                            var argSlice = keys[i].Item1;
-                            if (!session.WriteOrSendObjectStoreKeyValuePair(gcs, localServerSession, ref argSlice, out var status))
+                            var spanByte = keys[i].Item1.SpanByte;
+                            if (!session.WriteOrSendObjectStoreKeyValuePair(gcs, localServerSession, ref spanByte, out var status))
                                 return false;
 
                             // Skip if key NOTFOUND
@@ -182,6 +221,54 @@ public bool TransmitKeys(StoreType storeType)
                 return true;
             }
 
+            /// <summary>
+            /// Transmit data in namespaces during a MIGRATE ... KEYS operation.
+            /// 
+            /// Doesn't delete anything, just scans and transmits.
+            /// </summary>
+            public bool TransmitKeysNamespaces(ILogger logger)
+            {
+                var migrateOperation = this;
+
+                if (!migrateOperation.Initialize())
+                    return false;
+
+                var workerStartAddress = migrateOperation.session.clusterProvider.storeWrapper.store.Log.BeginAddress;
+                var workerEndAddress = migrateOperation.session.clusterProvider.storeWrapper.store.Log.TailAddress;
+
+                var cursor = workerStartAddress;
+                logger?.LogWarning("<MainStore> migrate keys (namespaces) scan range [{workerStartAddress}, {workerEndAddress}]", workerStartAddress, workerEndAddress);
+                while (true)
+                {
+                    var current = cursor;
+                    // Build Sketch
+                    migrateOperation.sketch.SetStatus(SketchStatus.INITIALIZING);
+                    migrateOperation.Scan(StoreType.Main, ref current, workerEndAddress);
+
+                    // Stop if no keys have been found
+                    if (migrateOperation.sketch.argSliceVector.IsEmpty) break;
+
+                    logger?.LogWarning("Scan from {cursor} to {current} and discovered {count} keys", cursor, current, migrateOperation.sketch.argSliceVector.Count);
+
+                    // Transition EPSM to MIGRATING
+                    migrateOperation.sketch.SetStatus(SketchStatus.TRANSMITTING);
+                    migrateOperation.session.WaitForConfigPropagation();
+
+                    // Transmit all keys gathered
+                    migrateOperation.TransmitSlots(StoreType.Main);
+
+                    // Transition EPSM to DELETING
+                    migrateOperation.sketch.SetStatus(SketchStatus.DELETING);
+                    migrateOperation.session.WaitForConfigPropagation();
+
+                    // Clear keys from buffer
+                    migrateOperation.sketch.Clear();
+                    cursor = current;
+                }
+
+                return true;
+            }
+
             /// <summary>
             /// Delete keys after migration if copyOption is not set
             /// </summary>
@@ -193,7 +280,13 @@ public void DeleteKeys()
                 {
                     foreach (var key in sketch.argSliceVector)
                     {
-                        var spanByte = key.SpanByte;
+                        if (key.MetadataSize == 1)
+                        {
+                            // Namespace'd keys are not deleted here, but when migration finishes
+                            continue;
+                        }
+
+                        var spanByte = key;
                         _ = localServerSession.BasicGarnetApi.DELETE(ref spanByte);
                     }
                 }
@@ -209,6 +302,19 @@ public void DeleteKeys()
                     }
                 }
             }
+
+            /// <summary>
+            /// Delete a Vector Set after migration if _copyOption is not set.
+            /// </summary>
+            public void DeleteVectorSet(ref SpanByte key)
+            {
+                if (session._copyOption)
+                    return;
+
+                var delRes = localServerSession.BasicGarnetApi.DELETE(ref key);
+
+                session.logger?.LogDebug("Deleting Vector Set {key} after migration: {delRes}", System.Text.Encoding.UTF8.GetString(key.AsReadOnlySpan()), delRes);
+            }
         }
     }
 }

libs/cluster/Server/Migration/MigrateScanFunctions.cs

@@ -36,10 +36,34 @@ public unsafe bool SingleReader(ref SpanByte key, ref SpanByte value, RecordMeta
                 if (ClusterSession.Expired(ref value))
                     return true;
 
-                var s = HashSlotUtils.HashSlot(ref key);
-                // Check if key belongs to slot that is being migrated and if it can be added to our buffer
-                if (mss.Contains(s) && !mss.sketch.TryHashAndStore(key.AsSpan()))
-                    return false;
+                // TODO: Some other way to detect namespaces
+                if (key.MetadataSize == 1)
+                {
+                    var ns = key.GetNamespaceInPayload();
+
+                    if (mss.ContainsNamespace(ns) && !mss.sketch.TryHashAndStore(ns, key.AsSpan()))
+                        return false;
+                }
+                else
+                {
+                    var s = HashSlotUtils.HashSlot(ref key);
+
+                    // Check if key belongs to slot that is being migrated...
+                    if (mss.Contains(s))
+                    {
+                        if (recordMetadata.RecordInfo.VectorSet)
+                        {
+                            // We can't delete the vector set _yet_ nor can we migrate it, 
+                            // we just need to remember it to migrate once the associated namespaces are all moved over
+                            mss.EncounteredVectorSet(key.ToByteArray(), value.ToByteArray());
+                        }
+                        else if (!mss.sketch.TryHashAndStore(key.AsSpan()))
+                        {
+                            // Out of space, end scan for now
+                            return false;
+                        }
+                    }
+                }
 
                 return true;
             }