Browse Source

Add ability to shift time buckets based on organization id. (#5530)

James Cunningham 7 years ago
parent
commit
702a75482d
1 changed files with 14 additions and 7 deletions
  1. 14 7
      src/sentry/quotas/redis.py

+ 14 - 7
src/sentry/quotas/redis.py

@@ -53,8 +53,12 @@ class RedisQuota(Quota):
         except Exception as e:
             raise InvalidConfiguration(six.text_type(e))
 
-    def __get_redis_key(self, key, timestamp, interval):
-        return '{}:{}:{}'.format(self.namespace, key, int(timestamp // interval))
+    def __get_redis_key(self, key, timestamp, interval, shift):
+        return '{}:{}:{}'.format(
+            self.namespace,
+            key,
+            int((timestamp - shift) // interval),
+        )
 
     def get_quotas(self, project, key=None):
         if key:
@@ -98,6 +102,7 @@ class RedisQuota(Quota):
                     quota.key,
                     timestamp,
                     quota.window,
+                    organization_id % quota.window
                 ),
             )
 
@@ -138,15 +143,16 @@ class RedisQuota(Quota):
         if not quotas:
             return NotRateLimited()
 
-        def get_next_period_start(interval):
+        def get_next_period_start(interval, shift):
             """Return the timestamp when the next rate limit period begins for an interval."""
-            return ((timestamp // interval) + 1) * interval
+            return ((timestamp // interval) + 1) * interval - shift
 
         keys = []
         args = []
         for quota in quotas:
-            keys.append(self.__get_redis_key(quota.key, timestamp, quota.window))
-            expiry = get_next_period_start(quota.window) + self.grace
+            shift = project.organization_id % quota.window
+            keys.append(self.__get_redis_key(quota.key, timestamp, quota.window, shift))
+            expiry = get_next_period_start(quota.window, shift) + self.grace
             args.extend((quota.limit, int(expiry)))
 
         client = self.cluster.get_local_client_for_key(six.text_type(project.organization.pk))
@@ -159,7 +165,8 @@ class RedisQuota(Quota):
                     continue
                 if quota.enforce:
                     enforce = True
-                    delay = get_next_period_start(quota.window) - timestamp
+                    shift = project.organization_id % quota.window
+                    delay = get_next_period_start(quota.window, shift) - timestamp
                     if delay > worst_case[0]:
                         worst_case = (delay, quota.reason_code)
             if enforce: