Fix KeyNum Extraction in Slot Verification (#734)

* add support for keynum extraction is slot verification * add more comments
microsoft · Oct 21, 2024 · 01f17f4 · 01f17f4
1 parent a4b47b4
commit 01f17f4
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 5 deletions.
diff --git a/libs/cluster/Session/SlotVerification/ClusterSlotVerify.cs b/libs/cluster/Session/SlotVerification/ClusterSlotVerify.cs
@@ -157,6 +157,8 @@ ClusterSlotVerificationResult MultiKeySlotVerify(ClusterConfig config, ref Sessi
 
             for (var i = stride; i < csvi.lastKey; i += stride)
             {
+                if (csvi.keyNumOffset == i)
+                    continue;
                 key = ref parseState.GetArgSliceByRef(i);
                 var _slot = ArgSliceUtils.HashSlot(ref key);
                 var _verifyResult = SingleKeySlotVerify(ref config, ref key, csvi.readOnly, csvi.sessionAsking, _slot);

diff --git a/libs/server/Cluster/ClusterSlotVerificationInput.cs b/libs/server/Cluster/ClusterSlotVerificationInput.cs
@@ -29,5 +29,10 @@ public struct ClusterSlotVerificationInput
         /// The step, or increment, between the first key and the position of the next key
         /// </summary>
         public int step;
+
+        /// <summary>
+        /// Offset of key num if any
+        /// </summary>
+        public int keyNumOffset;
     }
 }
diff --git a/libs/server/Resp/RespServerSessionSlotVerify.cs b/libs/server/Resp/RespServerSessionSlotVerify.cs
@@ -40,6 +40,7 @@ bool CanServeSlot(RespCommand cmd)
             if (commandInfo == null)
                 return true;
 
+            csvi.keyNumOffset = -1;
             var specs = commandInfo.KeySpecifications;
             switch (specs.Length)
             {
@@ -62,7 +63,7 @@ bool CanServeSlot(RespCommand cmd)
                             break;
                         case FindKeysUnknown:
                         default:
-                            throw new GarnetException("FindKeysUnknown range");
+                            throw new GarnetException("FindKeys spec not known");
                     }
 
                     break;
@@ -76,26 +77,31 @@ bool CanServeSlot(RespCommand cmd)
                         case FindKeysKeyNum:
                         case FindKeysUnknown:
                         default:
-                            throw new GarnetException("FindKeysUnknown range");
+                            throw new GarnetException("FindKeys spec not known");
                     }
 
+                    var searchIndex1 = (BeginSearchIndex)specs[1].BeginSearch;
                     switch (specs[1].FindKeys)
                     {
                         case FindKeysRange:
-                            var searchIndex1 = (BeginSearchIndex)specs[1].BeginSearch;
                             var findRange = (FindKeysRange)specs[1].FindKeys;
                             csvi.lastKey = findRange.LastKey < 0 ? findRange.LastKey + parseState.Count + 1 : findRange.LastKey + searchIndex1.Index - searchIndex.Index + 1;
                             csvi.step = findRange.KeyStep;
                             break;
                         case FindKeysKeyNum:
+                            var findKeysKeyNum = (FindKeysKeyNum)specs[1].FindKeys;
+                            csvi.keyNumOffset = searchIndex1.Index + findKeysKeyNum.KeyNumIdx - 1;
+                            csvi.lastKey = searchIndex1.Index + parseState.GetInt(csvi.keyNumOffset);
+                            csvi.step = findKeysKeyNum.KeyStep;
+                            break;
                         case FindKeysUnknown:
                         default:
-                            throw new GarnetException("FindKeysUnknown range");
+                            throw new GarnetException("FindKeys spec not known");
                     }
 
                     break;
                 default:
-                    throw new GarnetException("KeySpecification unknown count");
+                    throw new GarnetException("KeySpecification not supported count");
             }
             csvi.readOnly = cmd.IsReadOnly();
             csvi.sessionAsking = SessionAsking;