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();