From db4d6735b94e23488928011501d1c15771f91c64 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=B0=8F=E9=BB=91=E5=85=94?= <ywj6792341@126.com>
Date: Wed, 27 Mar 2024 03:10:52 +0800
Subject: [PATCH] =?UTF-8?q?=E6=9B=B4=E6=96=B0=20KEYS=E3=80=81SELECT?=
 =?UTF-8?q?=E3=80=81SET=20=E5=91=BD=E4=BB=A4=E4=BB=A5=20=E6=A8=A1=E6=8B=9F?=
 =?UTF-8?q?=E6=94=AF=E6=8C=81=20=E5=A4=9ADB=E5=8A=9F=E8=83=BD?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 libs/server/Resp/ArrayCommands.cs | 66 +++++++++++++++++++++++++++++--
 libs/server/Resp/BasicCommands.cs | 22 ++++++++++-
 2 files changed, 83 insertions(+), 5 deletions(-)

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'];
+
         /// <summary>
         /// MGET
         /// </summary>
@@ -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<TGarnetApi>(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<TGarnetApi>(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<byte[]> keys, ref byte* c
                     SendAndReset();
             }
         }
+
+        private void AddKeyPrefix(byte* source, int sourceSize, Action<nint, int> 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<TGarnetApi>(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<SpanByte>(keyPtr), ref Unsafe.AsRef<SpanByte>(valPtr));
+
+            Marshal.FreeHGlobal(temp);
+
+
             while (!RespWriteUtils.WriteResponse(CmdStrings.RESP_OK, ref dcurr, dend))
                 SendAndReset();