diff --git a/libs/server/Resp/ArrayCommands.cs b/libs/server/Resp/ArrayCommands.cs
index 9b57b425b..ac96bc214 100644
--- a/libs/server/Resp/ArrayCommands.cs
+++ b/libs/server/Resp/ArrayCommands.cs
@@ -7,6 +7,7 @@
using System.IO;
using System.Linq;
using System.Runtime.CompilerServices;
+using System.Runtime.InteropServices;
using System.Text;
using Garnet.common;
using Garnet.server.Custom;
@@ -22,6 +23,8 @@ internal sealed unsafe partial class RespServerSession : ServerSessionBase
int opsDone = 0;
int keysDeleted = 0;
+ byte[] db_id = [(byte)'0'];
+
///
/// MGET
///
@@ -649,8 +652,10 @@ private bool NetworkSELECT(byte* ptr)
readHead = (int)(ptr - recvBufferPtr);
- if (string.Equals(result, "0"))
+ if (!string.Equals(result, "16"))
{
+ this.db_id = Encoding.UTF8.GetBytes(result);
+
while (!RespWriteUtils.WriteResponse(CmdStrings.RESP_OK, ref dcurr, dend))
SendAndReset();
}
@@ -684,7 +689,21 @@ private bool NetworkKEYS(int count, byte* ptr, ref TGarnetApi storag
if (!RespReadUtils.ReadPtrWithLengthHeader(ref pattern, ref psize, ref ptr, recvBufferPtr + bytesRead))
return false;
- var patternAS = new ArgSlice(pattern, psize);
+ ArgSlice patternAS = default;
+
+ if (this.db_id.Length == 1 && this.db_id[0] == '0')
+ {
+ patternAS.ptr = (byte*)pattern;
+ patternAS.length = psize;
+ }
+ else
+ {
+ this.AddKeyPrefix(pattern, psize, (p, s) =>
+ {
+ patternAS.ptr = (byte*)p;
+ patternAS.length = s;
+ });
+ }
var keys = storageApi.GetDbKeys(patternAS);
@@ -697,8 +716,16 @@ private bool NetworkKEYS(int count, byte* ptr, ref TGarnetApi storag
// Write the keys matching the pattern
foreach (var item in keys)
{
- while (!RespWriteUtils.WriteBulkString(item, ref dcurr, dend))
- SendAndReset();
+ if (this.db_id.Length == 1 && this.db_id[0] == '0')
+ {
+ while (!RespWriteUtils.WriteBulkString(item, ref dcurr, dend))
+ SendAndReset();
+ }
+ else
+ {
+ while (!RespWriteUtils.WriteBulkString(this.RemoveKeyPrefix(item), ref dcurr, dend))
+ SendAndReset();
+ }
}
}
else
@@ -881,5 +908,36 @@ private void WriteOutputForScan(long cursorValue, List keys, ref byte* c
SendAndReset();
}
}
+
+ private void AddKeyPrefix(byte* source, int sourceSize, Action action)
+ {
+ int newSize = sourceSize + db_id.Length + 1;
+ fixed (byte* newptr = new byte[newSize])
+ {
+
+ Marshal.Copy(db_id, 0, (nint)newptr, db_id.Length);
+ *(newptr + db_id.Length) = (byte)':';
+
+ byte* sourceStart = newptr + db_id.Length + 1;
+
+ for (int i = 0; i < sourceSize; i++)
+ {
+ *(sourceStart + i) = *(source + i);
+ }
+
+ action((nint)newptr, newSize);
+ }
+ }
+
+ private byte[] RemoveKeyPrefix(byte[] bytes)
+ {
+ byte[] result = new byte[bytes.Length - (db_id.Length + 1)];
+ Array.Copy(bytes, db_id.Length + 1, result, 0, result.Length);
+ return result;
+ }
+ //private void RemoveKeyPrefix(byte* ptr)
+ //{
+
+ //}
}
}
\ No newline at end of file
diff --git a/libs/server/Resp/BasicCommands.cs b/libs/server/Resp/BasicCommands.cs
index d5ccbdc42..1ab7ed33b 100644
--- a/libs/server/Resp/BasicCommands.cs
+++ b/libs/server/Resp/BasicCommands.cs
@@ -4,6 +4,7 @@
using System;
using System.Diagnostics;
using System.Runtime.CompilerServices;
+using System.Runtime.InteropServices;
using Garnet.common;
using Tsavorite.core;
@@ -255,12 +256,31 @@ private bool NetworkSET(byte* ptr, ref TGarnetApi storageApi)
if (NetworkSingleKeySlotVerify(keyPtr, ksize, false))
return true;
+
+ var temp = Marshal.AllocHGlobal(sizeof(int) + this.db_id.Length + 1 + ksize);
+
+ var newkeyptr = temp + sizeof(int);
+
+
+ Unsafe.CopyBlock((byte*)newkeyptr, (byte*)keyPtr, (uint)ksize);
+
+ this.AddKeyPrefix((byte*)newkeyptr, ksize, (p, s) =>
+ {
+ Unsafe.CopyBlock((byte*)newkeyptr, (byte*)p, (uint)s);
+
+ keyPtr = (byte*)newkeyptr;
+ ksize = s;
+ });
+
keyPtr -= sizeof(int);
valPtr -= sizeof(int);
*(int*)keyPtr = ksize;
*(int*)valPtr = vsize;
-
var status = storageApi.SET(ref Unsafe.AsRef(keyPtr), ref Unsafe.AsRef(valPtr));
+
+ Marshal.FreeHGlobal(temp);
+
+
while (!RespWriteUtils.WriteResponse(CmdStrings.RESP_OK, ref dcurr, dend))
SendAndReset();