Repository: getsentry/sentry
Pull request: feat(relocation): Add throttling
closes https://github.com/getsentry/sentry/issues/60694

Problem statement:
Throttling for relocations
### Problem Statement

the broad goal here is to have global throttles for each of `small` (<10MB), `medium` (>=10MB && <100MB), and `large` (>100MB) relocation uploads. These should respectively be controlled by 3 new options: `relocation:daily-limit-small`, `relocation:daily-limit-medium`, and `relocation:daily-limit-large`, each of which is an integer which, you guessed it, controls the number of relocations of that size that we allow on a single (UTC) day. For example, we might start with a `large` limit of 2, a `medium` limit of 10, and a `small` limit of 50. Let's set the defaults of each to 0 as well - we'll want to turn these on manually.

Other notes:
- I *think* the endpoints themselves should have access to the size of the file being uploaded, so we can do all of this handling in the POST endpoint itself, rather than the tasks.
- Best HTTP error code is probably 429.
- We should make sure not to tell the user what the actual limit is - just tell them that relocations of their size are throttled, and that they should either try again tomorrow or contact support.

### Solution Brainstorm

_No response_

### Product Area

Unknown

Hints from issue discussion:
Routing to @getsentry/product-owners-org-relocation for [triage](https://develop.sentry.dev/processing-tickets/#3-triage) ⏲️
CC @hubertdeng123

Commit message: add line breaks back

Patch to review:
diff --git a/src/sentry/api/endpoints/relocations/index.py b/src/sentry/api/endpoints/relocations/index.py
index 1da3d64314a2b2..33bef6dc9b6226 100644
--- a/src/sentry/api/endpoints/relocations/index.py
+++ b/src/sentry/api/endpoints/relocations/index.py
@@ -1,6 +1,9 @@
 import logging
+from datetime import timedelta
+from functools import reduce
 
 from django.db import router
+from django.utils import timezone
 from rest_framework import serializers
 from rest_framework.request import Request
 from rest_framework.response import Response
@@ -14,6 +17,7 @@
 from sentry.models.files.file import File
 from sentry.models.relocation import Relocation, RelocationFile
 from sentry.models.user import MAX_USERNAME_LENGTH
+from sentry.options import get
 from sentry.services.hybrid_cloud.user.service import user_service
 from sentry.tasks.relocation import uploading_complete
 from sentry.utils.db import atomic_transaction
@@ -23,6 +27,35 @@
 
 logger = logging.getLogger(__name__)
 
+RELOCATION_FILE_SIZE_SMALL = 10 * 1024**2
+RELOCATION_FILE_SIZE_MEDIUM = 100 * 1024**2
+
+
+def get_relocation_bucket(size) -> str:
+    if size < RELOCATION_FILE_SIZE_SMALL:
+        return "small"
+    elif size < RELOCATION_FILE_SIZE_MEDIUM:
+        return "medium"
+    return "large"
+
+
+def should_throttle_relocation(relocation_bucket_size) -> bool:
+    recent_relocation_files = RelocationFile.objects.filter(
+        date_added__gte=(timezone.now() - timedelta(days=1))
+    )
+    num_recent_same_size_relocation_files = reduce(
+        lambda acc, relocation_file: acc + 1
+        if get_relocation_bucket(relocation_file.file.size) == relocation_bucket_size
+        else acc,
+        recent_relocation_files,
+        0,
+    )
+    if num_recent_same_size_relocation_files < get(
+        f"relocation.daily-limit-{relocation_bucket_size}"
+    ):
+        return False
+    return True
+
 
 class RelocationPostSerializer(serializers.Serializer):
     file = serializers.FileField(required=True)
@@ -83,8 +116,12 @@ def post(self, request: Request) -> Response:
         ).exists():
             return Response({"detail": ERR_DUPLICATE_RELOCATION}, status=409)
 
-        # TODO(getsentry/team-ospo#216): check import size, and maybe throttle based on that
-        # information.
+        relocation_bucket_size = get_relocation_bucket(fileobj.size)
+        if should_throttle_relocation(relocation_bucket_size):
+            return Response(
+                {"detail": f"Already reached daily limit of {relocation_bucket_size} relocations"},
+                status=429,
+            )
 
         file = File.objects.create(name="raw-relocation-data.tar", type=RELOCATION_FILE_TYPE)
         file.putfile(fileobj, blob_size=RELOCATION_BLOB_SIZE, logger=logger)
diff --git a/src/sentry/options/defaults.py b/src/sentry/options/defaults.py
index 2875081d5e0b1e..f4441185053242 100644
--- a/src/sentry/options/defaults.py
+++ b/src/sentry/options/defaults.py
@@ -1747,3 +1747,10 @@
     default=False,
     flags=FLAG_MODIFIABLE_BOOL | FLAG_AUTOMATOR_MODIFIABLE,
 )
+
+# Relocation
+#
+# Throttling limits for relocation requests
+register("relocation.daily-limit-small", default=0)
+register("relocation.daily-limit-medium", default=0)
+register("relocation.daily-limit-large", default=0)
diff --git a/tests/sentry/api/endpoints/relocations/test_index.py b/tests/sentry/api/endpoints/relocations/test_index.py
index fcca75edbbaa25..364c57650db797 100644
--- a/tests/sentry/api/endpoints/relocations/test_index.py
+++ b/tests/sentry/api/endpoints/relocations/test_index.py
@@ -1,4 +1,5 @@
 import tempfile
+from datetime import timedelta
 from pathlib import Path
 from typing import Tuple
 from unittest.mock import patch
@@ -12,6 +13,7 @@
 from sentry.testutils.cases import APITestCase
 from sentry.testutils.factories import get_fixture_path
 from sentry.testutils.helpers.backups import generate_rsa_key_pair
+from sentry.testutils.helpers.datetime import freeze_time
 from sentry.testutils.silo import region_silo_test
 from sentry.utils import json
 
@@ -51,7 +53,9 @@ def test_success_simple(self, uploading_complete_mock):
 
         with tempfile.TemporaryDirectory() as tmp_dir:
             (_, tmp_pub_key_path) = self.tmp_keys(tmp_dir)
-            with self.feature("relocation:enabled"), open(FRESH_INSTALL_PATH) as f:
+            with self.feature("relocation:enabled"), self.options(
+                {"relocation.daily-limit-small": 1}
+            ), open(FRESH_INSTALL_PATH) as f:
                 data = json.load(f)
                 with open(tmp_pub_key_path, "rb") as p:
                     response = self.client.post(
@@ -95,7 +99,9 @@ def test_success_relocation_for_same_owner_already_completed(self):
 
         with tempfile.TemporaryDirectory() as tmp_dir:
             (_, tmp_pub_key_path) = self.tmp_keys(tmp_dir)
-            with self.feature("relocation:enabled"), open(FRESH_INSTALL_PATH) as f:
+            with self.feature("relocation:enabled"), self.options(
+                {"relocation.daily-limit-small": 1}
+            ), open(FRESH_INSTALL_PATH) as f:
                 data = json.load(f)
                 with open(tmp_pub_key_path, "rb") as p:
                     response = self.client.post(
@@ -147,7 +153,9 @@ def test_fail_feature_disabled(self):
     def test_fail_missing_file(self):
         with tempfile.TemporaryDirectory() as tmp_dir:
             (_, tmp_pub_key_path) = self.tmp_keys(tmp_dir)
-            with self.feature("relocation:enabled"):
+            with self.feature("relocation:enabled"), self.options(
+                {"relocation.daily-limit-small": 1}
+            ):
                 response = self.client.post(
                     reverse(self.endpoint),
                     {
@@ -164,7 +172,9 @@ def test_fail_missing_file(self):
     def test_fail_missing_orgs(self):
         with tempfile.TemporaryDirectory() as tmp_dir:
             (_, tmp_pub_key_path) = self.tmp_keys(tmp_dir)
-            with self.feature("relocation:enabled"), open(FRESH_INSTALL_PATH) as f:
+            with self.feature("relocation:enabled"), self.options(
+                {"relocation.daily-limit-small": 1}
+            ), open(FRESH_INSTALL_PATH) as f:
                 data = json.load(f)
                 with open(tmp_pub_key_path, "rb") as p:
                     response = self.client.post(
@@ -189,7 +199,9 @@ def test_fail_missing_orgs(self):
     def test_fail_missing_owner(self):
         with tempfile.TemporaryDirectory() as tmp_dir:
             (_, tmp_pub_key_path) = self.tmp_keys(tmp_dir)
-            with self.feature("relocation:enabled"), open(FRESH_INSTALL_PATH) as f:
+            with self.feature("relocation:enabled"), self.options(
+                {"relocation.daily-limit-small": 1}
+            ), open(FRESH_INSTALL_PATH) as f:
                 data = json.load(f)
                 with open(tmp_pub_key_path, "rb") as p:
                     response = self.client.post(
@@ -214,7 +226,9 @@ def test_fail_missing_owner(self):
     def test_fail_nonexistent_owner(self):
         with tempfile.TemporaryDirectory() as tmp_dir:
             (_, tmp_pub_key_path) = self.tmp_keys(tmp_dir)
-            with self.feature("relocation:enabled"), open(FRESH_INSTALL_PATH) as f:
+            with self.feature("relocation:enabled"), self.options(
+                {"relocation.daily-limit-small": 1}
+            ), open(FRESH_INSTALL_PATH) as f:
                 data = json.load(f)
                 with open(tmp_pub_key_path, "rb") as p:
                     response = self.client.post(
@@ -247,7 +261,9 @@ def test_fail_relocation_for_same_owner_already_in_progress(self):
 
         with tempfile.TemporaryDirectory() as tmp_dir:
             (_, tmp_pub_key_path) = self.tmp_keys(tmp_dir)
-            with self.feature("relocation:enabled"), open(FRESH_INSTALL_PATH) as f:
+            with self.feature("relocation:enabled"), self.options(
+                {"relocation.daily-limit-small": 1}
+            ), open(FRESH_INSTALL_PATH) as f:
                 data = json.load(f)
                 with open(tmp_pub_key_path, "rb") as p:
                     simple_file = SimpleUploadedFile(
@@ -267,3 +283,178 @@ def test_fail_relocation_for_same_owner_already_in_progress(self):
                     )
 
         assert response.status_code == 409
+
+    def test_fail_throttle_if_daily_limit_reached(self):
+        relocation_count = Relocation.objects.count()
+        relocation_file_count = RelocationFile.objects.count()
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            (_, tmp_pub_key_path) = self.tmp_keys(tmp_dir)
+            with self.feature("relocation:enabled"), self.options(
+                {"relocation.daily-limit-small": 1}
+            ), open(FRESH_INSTALL_PATH) as f:
+                data = json.load(f)
+                with open(tmp_pub_key_path, "rb") as p:
+                    initial_response = self.client.post(
+                        reverse(self.endpoint),
+                        {
+                            "owner": self.owner.username,
+                            "file": SimpleUploadedFile(
+                                "export.tar",
+                                create_encrypted_export_tarball(
+                                    data, LocalFileEncryptor(p)
+                                ).getvalue(),
+                                content_type="application/tar",
+                            ),
+                            "orgs": ["testing", "foo"],
+                        },
+                        format="multipart",
+                    )
+
+                    # Simulate completion of relocation job
+                    relocation = Relocation.objects.all()[0]
+                    relocation.status = Relocation.Status.SUCCESS.value
+                    relocation.save()
+                    relocation.refresh_from_db()
+
+                with open(tmp_pub_key_path, "rb") as p:
+                    throttled_response = self.client.post(
+                        reverse(self.endpoint),
+                        {
+                            "owner": self.owner.username,
+                            "file": SimpleUploadedFile(
+                                "export.tar",
+                                create_encrypted_export_tarball(
+                                    data, LocalFileEncryptor(p)
+                                ).getvalue(),
+                                content_type="application/tar",
+                            ),
+                            "orgs": ["testing", "foo"],
+                        },
+                        format="multipart",
+                    )
+
+        assert initial_response.status_code == 201
+        assert throttled_response.status_code == 429
+        assert throttled_response.data.get("detail") is not None
+        assert (
+            throttled_response.data.get("detail")
+            == "Already reached daily limit of small relocations"
+        )
+        assert Relocation.objects.count() == relocation_count + 1
+        assert RelocationFile.objects.count() == relocation_file_count + 1
+
+    def test_success_no_throttle_different_bucket_relocations(self):
+        relocation_count = Relocation.objects.count()
+        relocation_file_count = RelocationFile.objects.count()
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            (_, tmp_pub_key_path) = self.tmp_keys(tmp_dir)
+            with self.feature("relocation:enabled"), self.options(
+                {"relocation.daily-limit-small": 1, "relocation.daily-limit-medium": 1}
+            ), open(FRESH_INSTALL_PATH) as f:
+                data = json.load(f)
+                with open(tmp_pub_key_path, "rb") as p:
+                    initial_response = self.client.post(
+                        reverse(self.endpoint),
+                        {
+                            "owner": self.owner.username,
+                            "file": SimpleUploadedFile(
+                                "export.tar",
+                                create_encrypted_export_tarball(
+                                    data, LocalFileEncryptor(p)
+                                ).getvalue(),
+                                content_type="application/tar",
+                            ),
+                            "orgs": ["testing", "foo"],
+                        },
+                        format="multipart",
+                    )
+
+                    # Simulate completion of relocation job
+                    relocation = Relocation.objects.all()[0]
+                    relocation.status = Relocation.Status.SUCCESS.value
+                    relocation.save()
+                    relocation.refresh_from_db()
+
+                with open(tmp_pub_key_path, "rb") as p:
+                    throttled_response = self.client.post(
+                        reverse(self.endpoint),
+                        {
+                            "owner": self.owner.username,
+                            "file": SimpleUploadedFile(
+                                "export.tar",
+                                create_encrypted_export_tarball(
+                                    data, LocalFileEncryptor(p)
+                                ).getvalue()
+                                * 1000,
+                                content_type="application/tar",
+                            ),
+                            "orgs": ["testing", "foo"],
+                        },
+                        format="multipart",
+                    )
+
+        assert initial_response.status_code == 201
+        assert throttled_response.status_code == 201
+        assert Relocation.objects.count() == relocation_count + 2
+        assert RelocationFile.objects.count() == relocation_file_count + 2
+
+    def test_success_no_throttle_relocation_over_multiple_days(self):
+        relocation_count = Relocation.objects.count()
+        relocation_file_count = RelocationFile.objects.count()
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            (_, tmp_pub_key_path) = self.tmp_keys(tmp_dir)
+            with self.feature("relocation:enabled"), self.options(
+                {"relocation.daily-limit-small": 1}
+            ), open(FRESH_INSTALL_PATH) as f, freeze_time("2023-11-28 00:00:00") as frozen_time:
+                data = json.load(f)
+                with open(tmp_pub_key_path, "rb") as p:
+                    initial_response = self.client.post(
+                        reverse(self.endpoint),
+                        {
+                            "owner": self.owner.username,
+                            "file": SimpleUploadedFile(
+                                "export.tar",
+                                create_encrypted_export_tarball(
+                                    data, LocalFileEncryptor(p)
+                                ).getvalue(),
+                                content_type="application/tar",
+                            ),
+                            "orgs": ["testing", "foo"],
+                        },
+                        format="multipart",
+                    )
+
+                    # Simulate completion of relocation job
+                    relocation = Relocation.objects.all()[0]
+                    relocation.status = Relocation.Status.SUCCESS.value
+                    relocation.save()
+                    relocation.refresh_from_db()
+
+                frozen_time.shift(timedelta(days=1, minutes=1))
+
+                # Relogin since session has expired
+                self.login_as(user=self.superuser, superuser=True)
+                with open(tmp_pub_key_path, "rb") as p:
+                    throttled_response = self.client.post(
+                        reverse(self.endpoint),
+                        {
+                            "owner": self.owner.username,
+                            "file": SimpleUploadedFile(
+                                "export.tar",
+                                create_encrypted_export_tarball(
+                                    data, LocalFileEncryptor(p)
+                                ).getvalue(),
+                                content_type="application/tar",
+                            ),
+                            "orgs": ["testing", "foo"],
+                        },
+                        format="multipart",
+                    )
+
+        assert initial_response.status_code == 201
+        assert throttled_response.status_code == 201
+        assert Relocation.objects.count() == relocation_count + 2
+        assert RelocationFile.objects.count() == relocation_file_count + 2


Code review comments:
diff_hunk: @@ -83,8 +116,12 @@ def post(self, request: Request) -> Response:
         ).exists():
             return Response({"detail": ERR_DUPLICATE_RELOCATION}, status=409)
 
-        # TODO(getsentry/team-ospo#216): check import size, and maybe throttle based on that
-        # information.
+        relocation_bucket_size = get_relocation_bucket(fileobj.size)
+        if should_throttle_relocation(relocation_bucket_size):
+            return Response(
+                {"detail": f"Already reached daily limit of {relocation_bucket_size} relocations"},
path: src/sentry/api/endpoints/relocations/index.py
text: @user1:
I think we should be even more vague. Just say something like "we've reached our daily limit of relocations - please try again tomorrow or contact support."

@author:
done

diff_hunk: @@ -23,6 +27,35 @@
 
 logger = logging.getLogger(__name__)
 
+RELOCATION_FILE_SIZE_SMALL = 10 * 1024**2
+RELOCATION_FILE_SIZE_MEDIUM = 100 * 1024**2
+
+
+def get_relocation_bucket(size) -> str:
path: src/sentry/api/endpoints/relocations/index.py
text: @user1:
nit: Maybe avoid the word `bucket`, since it may get confused with GCP buckets easily. Maybe something like `get_relocation_size_category` or similar?

@author:
done