feat(metrics): Implement metric_bucket rate limits

Author: cleptric

sentry_sdk/transport.py

@@ -144,10 +144,22 @@ def _parse_rate_limits(header, now=None):
 
     for limit in header.split(","):
         try:
-            retry_after, categories, _ = limit.strip().split(":", 2)
+            parameters = limit.strip().split(":")
+            retry_after, categories = parameters[:2]
+
             retry_after = now + timedelta(seconds=int(retry_after))
             for category in categories and categories.split(";") or (None,):
-                yield category, retry_after
+                if category == "metric_bucket":
+                    try:
+                        namespaces = parameters[4].split(";")
+                    except IndexError:
+                        namespaces = []
+
+                    if not namespaces or "custom" in namespaces:
+                        yield category, retry_after
+                
+                else:
+                    yield category, retry_after
         except (LookupError, ValueError):
             continue
 
@@ -336,7 +348,14 @@ def _check_disabled(self, category):
         # type: (str) -> bool
         def _disabled(bucket):
             # type: (Any) -> bool
+
+            # The envelope item type used for metrics is statsd
+            # whereas the rate limit category is metric_bucket
+            if bucket == "statsd":
+                bucket = "metric_bucket"
+
             ts = self._disabled_until.get(bucket)
+
             return ts is not None and ts > datetime_utcnow()
 
         return _disabled(category) or _disabled(None)
@@ -402,7 +421,7 @@ def _send_envelope(
         new_items = []
         for item in envelope.items:
             if self._check_disabled(item.data_category):
-                if item.data_category in ("transaction", "error", "default"):
+                if item.data_category in ("transaction", "error", "default", "statsd"):
                     self.on_dropped_event("self_rate_limits")
                 self.record_lost_event("ratelimit_backoff", item=item)
             else:

tests/test_transport.py

@@ -14,7 +14,7 @@
 from sentry_sdk import Hub, Client, add_breadcrumb, capture_message, Scope
 from sentry_sdk._compat import datetime_utcnow
 from sentry_sdk.transport import KEEP_ALIVE_SOCKET_OPTIONS, _parse_rate_limits
-from sentry_sdk.envelope import Envelope, parse_json
+from sentry_sdk.envelope import Envelope, Item, parse_json
 from sentry_sdk.integrations.logging import LoggingIntegration
 
 try:
@@ -466,3 +466,111 @@ def test_complex_limits_without_data_category(
     client.flush()
 
     assert len(capturing_server.captured) == 0
+
+@pytest.mark.parametrize("response_code", [200, 429])
+def test_metric_bucket_limits(
+    capturing_server, response_code, make_client
+):
+    client = make_client()
+    capturing_server.respond_with(
+        code=response_code,
+        headers={"X-Sentry-Rate-Limits": "4711:metric_bucket:organization:quota_exceeded:custom"},
+    )
+
+    envelope = Envelope()
+    envelope.add_item(Item(payload=b"{}", type="statsd"))
+    client.transport.capture_envelope(envelope)
+    client.flush()
+
+    assert len(capturing_server.captured) == 1
+    assert capturing_server.captured[0].path == "/api/132/envelope/"
+    capturing_server.clear_captured()
+
+    assert set(client.transport._disabled_until) == set(["metric_bucket"])
+
+    client.transport.capture_envelope(envelope)
+    client.capture_event({"type": "transaction"})
+    client.flush()
+
+    assert len(capturing_server.captured) == 2
+
+    envelope = capturing_server.captured[0].envelope
+    assert envelope.items[0].type == "transaction"
+    envelope = capturing_server.captured[1].envelope
+    assert envelope.items[0].type == "client_report"
+    report = parse_json(envelope.items[0].get_bytes())
+    assert report["discarded_events"] == [
+        # Clarify category statsd or metric_bucket
+        {"category": "statsd", "reason": "ratelimit_backoff", "quantity": 1},
+    ]
+
+
+
+@pytest.mark.parametrize("response_code", [200, 429])
+def test_metric_bucket_limits_with_namespace(
+    capturing_server, response_code, make_client
+):
+    client = make_client()
+    capturing_server.respond_with(
+        code=response_code,
+        headers={"X-Sentry-Rate-Limits": "4711:metric_bucket:organization:quota_exceeded:foo"},
+    )
+
+    envelope = Envelope()
+    envelope.add_item(Item(payload=b"{}", type="statsd"))
+    client.transport.capture_envelope(envelope)
+    client.flush()
+
+    assert len(capturing_server.captured) == 1
+    assert capturing_server.captured[0].path == "/api/132/envelope/"
+    capturing_server.clear_captured()
+
+    assert set(client.transport._disabled_until) == set([])
+
+    client.transport.capture_envelope(envelope)
+    client.capture_event({"type": "transaction"})
+    client.flush()
+
+    assert len(capturing_server.captured) == 2
+
+    envelope = capturing_server.captured[0].envelope
+    assert envelope.items[0].type == "statsd"
+    envelope = capturing_server.captured[1].envelope
+    assert envelope.items[0].type == "transaction"
+
+@pytest.mark.parametrize("response_code", [200, 429])
+def test_metric_bucket_limits_with_all_namespaces(
+    capturing_server, response_code, make_client
+):
+    client = make_client()
+    capturing_server.respond_with(
+        code=response_code,
+        headers={"X-Sentry-Rate-Limits": "4711:metric_bucket:organization:quota_exceeded"},
+    )
+
+    envelope = Envelope()
+    envelope.add_item(Item(payload=b"{}", type="statsd"))
+    client.transport.capture_envelope(envelope)
+    client.flush()
+
+    assert len(capturing_server.captured) == 1
+    assert capturing_server.captured[0].path == "/api/132/envelope/"
+    capturing_server.clear_captured()
+
+    assert set(client.transport._disabled_until) == set(["metric_bucket"])
+
+    client.transport.capture_envelope(envelope)
+    client.capture_event({"type": "transaction"})
+    client.flush()
+
+    assert len(capturing_server.captured) == 2
+
+    envelope = capturing_server.captured[0].envelope
+    assert envelope.items[0].type == "transaction"
+    envelope = capturing_server.captured[1].envelope
+    assert envelope.items[0].type == "client_report"
+    report = parse_json(envelope.items[0].get_bytes())
+    assert report["discarded_events"] == [
+        # Clarify category statsd or metric_bucket
+        {"category": "statsd", "reason": "ratelimit_backoff", "quantity": 1},
+    ]