[GROW-2938] recover from SlotNotCoveredError

Author: zach-iee

redis/cluster.py

@@ -858,41 +858,49 @@ def set_response_callback(self, command, callback):
     def _determine_nodes(self, *args, **kwargs) -> List["ClusterNode"]:
         # Determine which nodes should be executed the command on.
         # Returns a list of target nodes.
-        command = args[0].upper()
-        if len(args) >= 2 and f"{args[0]} {args[1]}".upper() in self.command_flags:
-            command = f"{args[0]} {args[1]}".upper()
-
-        nodes_flag = kwargs.pop("nodes_flag", None)
-        if nodes_flag is not None:
-            # nodes flag passed by the user
-            command_flag = nodes_flag
-        else:
-            # get the nodes group for this command if it was predefined
-            command_flag = self.command_flags.get(command)
-        if command_flag == self.__class__.RANDOM:
-            # return a random node
-            return [self.get_random_node()]
-        elif command_flag == self.__class__.PRIMARIES:
-            # return all primaries
-            return self.get_primaries()
-        elif command_flag == self.__class__.REPLICAS:
-            # return all replicas
-            return self.get_replicas()
-        elif command_flag == self.__class__.ALL_NODES:
-            # return all nodes
-            return self.get_nodes()
-        elif command_flag == self.__class__.DEFAULT_NODE:
-            # return the cluster's default node
-            return [self.nodes_manager.default_node]
-        elif command in self.__class__.SEARCH_COMMANDS[0]:
-            return [self.nodes_manager.default_node]
-        else:
-            # get the node that holds the key's slot
-            slot = self.determine_slot(*args)
-            node = self.nodes_manager.get_node_from_slot(
-                slot, self.read_from_replicas and command in READ_COMMANDS
-            )
-            return [node]
+        try:
+            command = args[0].upper()
+            if len(args) >= 2 and f"{args[0]} {args[1]}".upper() in self.command_flags:
+                command = f"{args[0]} {args[1]}".upper()
+
+            nodes_flag = kwargs.pop("nodes_flag", None)
+            if nodes_flag is not None:
+                # nodes flag passed by the user
+                command_flag = nodes_flag
+            else:
+                # get the nodes group for this command if it was predefined
+                command_flag = self.command_flags.get(command)
+            if command_flag == self.__class__.RANDOM:
+                # return a random node
+                return [self.get_random_node()]
+            elif command_flag == self.__class__.PRIMARIES:
+                # return all primaries
+                return self.get_primaries()
+            elif command_flag == self.__class__.REPLICAS:
+                # return all replicas
+                return self.get_replicas()
+            elif command_flag == self.__class__.ALL_NODES:
+                # return all nodes
+                return self.get_nodes()
+            elif command_flag == self.__class__.DEFAULT_NODE:
+                # return the cluster's default node
+                return [self.nodes_manager.default_node]
+            elif command in self.__class__.SEARCH_COMMANDS[0]:
+                return [self.nodes_manager.default_node]
+            else:
+                # get the node that holds the key's slot
+                slot = self.determine_slot(*args)
+                node = self.nodes_manager.get_node_from_slot(
+                    slot, self.read_from_replicas and command in READ_COMMANDS
+                )
+                return [node]
+        except SlotNotCoveredError as e:
+            self.reinitialize_counter += 1
+            if self._should_reinitialized():
+                self.nodes_manager.initialize()
+                # Reset the counter
+                self.reinitialize_counter = 0
+            raise e
 
     def _should_reinitialized(self):
         # To reinitialize the cluster on every MOVED error,
@@ -1166,6 +1174,13 @@ def _execute_command(self, target_node, *args, **kwargs):
                 else:
                     self.nodes_manager.update_moved_exception(e)
                 moved = True
+            except SlotNotCoveredError as e:
+                self.reinitialize_counter += 1
+                if self._should_reinitialized():
+                    self.nodes_manager.initialize()
+                    # Reset the counter
+                    self.reinitialize_counter = 0
+                raise e
             except TryAgainError:
                 if ttl < self.RedisClusterRequestTTL / 2:
                     time.sleep(0.05)

tests/test_cluster.py

@@ -4,6 +4,7 @@
 import socket
 import socketserver
 import threading
+import uuid
 import warnings
 from queue import LifoQueue, Queue
 from time import sleep
@@ -35,7 +36,7 @@
     RedisClusterException,
     RedisError,
     ResponseError,
-    TimeoutError,
+    TimeoutError, SlotNotCoveredError,
 )
 from redis.retry import Retry
 from redis.utils import str_if_bytes
@@ -922,6 +923,19 @@ def raise_error(target_node, *args, **kwargs):
                     rc.get("bar")
                 assert compute.call_count == rc.cluster_error_retry_attempts
 
+    @pytest.mark.parametrize('reinitialize_steps', [2, 10, 99])
+    def test_recover_slot_not_covered_error(self, request, reinitialize_steps):
+        rc = _get_client(RedisCluster, request, reinitialize_steps=reinitialize_steps)
+        key = uuid.uuid4().hex
+
+        rc.nodes_manager.slots_cache[rc.keyslot(key)] = []
+
+        for _ in range(0, reinitialize_steps):
+            with pytest.raises(SlotNotCoveredError):
+                rc.get(key)
+
+        rc.get(key)
+
 
 @pytest.mark.onlycluster
 class TestClusterRedisCommands: