delegate time bucket naming to the backend; ensure same-slot hashing

Author: mikeroll

dramatiq/rate_limits/backend.py

@@ -65,6 +65,19 @@ def decr(self, key, amount, minimum, ttl):  # pragma: no cover
         """
         raise NotImplementedError
 
+    def format_key_variants(self, key, variants):  # pragma: no cover
+        """Build a list of related key names from a "base" key name and and a list of
+        distinct "variants" that can act as e.g. name suffixes.
+
+        Parameters:
+          key(str): The base key name.
+          variants(list[str]): Distinct values to incorporate into the resulting names.
+
+        Returns:
+          list[str]: The list of resulting key names.
+        """
+        return [f"{key}@{variant}" for variant in variants]
+
     def incr_and_sum(self, key, keys, amount, maximum, ttl):  # pragma: no cover
         """Atomically increment a key unless the sum of keys is greater
         than the given maximum.

dramatiq/rate_limits/backends/redis.py

@@ -63,6 +63,13 @@ def decr(self, key, amount, minimum, ttl):
         decr_down_to = self.scripts["decr_down_to"]
         return decr_down_to([key], [amount, minimum, ttl]) == 1
 
+    def format_key_variants(self, key, variants):
+        # NOTE the extra { } around the key - this is to make use of Hash tags [0]
+        # to make sure that multi-key commands execute on the keys in same hash slots.
+        # This helps avoid a ClusterCrossSlotError in case this redis is a Redis Cluster.
+        # [0]: https://redis.io/docs/latest/operate/oss_and_stack/reference/cluster-spec/#keys-hash-tags
+        return [f"{{{key}}}@{variant}" for variant in variants]
+
     def incr_and_sum(self, key, keys, amount, maximum, ttl):
         # TODO: Drop non-callable keys in Dramatiq v2.
         keys_list = keys() if callable(keys) else keys

dramatiq/rate_limits/window.py

@@ -52,7 +52,7 @@ def __init__(self, backend, key, *, limit=1, window=1):
 
     def _get_keys(self):
         timestamp = int(time.time())
-        return ["%s@%s" % (self.key, timestamp - i) for i in range(self.window)]
+        return self.backend.format_key_variants(self.key, [str(timestamp - i) for i in range(self.window)])
 
     def _acquire(self):
         keys = self._get_keys()