Hook up peanutbutter as an LPQ backend

.github/actions/setup-sentry/action.yml

@@ -35,6 +35,10 @@ inputs:
     description: 'Is symbolicator required?'
     required: false
     default: 'false'
+  peanutbutter:
+    description: 'Is peanutbutter required?'
+    required: false
+    default: 'false'
   python-version:
     description: 'python version to install'
     required: false
@@ -146,6 +150,7 @@ runs:
         NEED_CHARTCUTERIE: ${{ inputs.chartcuterie }}
         NEED_REDIS_CLUSTER: ${{ inputs.redis_cluster }}
         NEED_SYMBOLICATOR: ${{ inputs.symbolicator }}
+        NEED_PEANUTBUTTER: ${{ inputs.peanutbutter }}
         WORKDIR: ${{ inputs.workdir }}
         PG_VERSION: ${{ inputs.pg-version }}
         ENABLE_AUTORUN_MIGRATION_SEARCH_ISSUES: '1'
@@ -180,6 +185,10 @@ runs:
           services+=(symbolicator)
         fi
 
+        if [ "$NEED_PEANUTBUTTER" = "true" ]; then
+          services+=(peanutbutter)
+        fi
+
         if [ "$NEED_KAFKA" = "true" ]; then
           services+=(kafka)
         fi

.github/workflows/backend.yml

@@ -99,6 +99,7 @@ jobs:
           kafka: true
           snuba: true
           symbolicator: true
+          peanutbutter: true
           # Right now, we run so few bigtable related tests that the
           # overhead of running bigtable in all backend tests
           # is way smaller than the time it would take to run in its own job.

src/sentry/conf/server.py

@@ -2743,6 +2743,11 @@ def custom_parameter_sort(parameter: dict) -> tuple[str, int]:
 # This flag activates Spotlight Sidecar in the development environment
 SENTRY_USE_SPOTLIGHT = False
 
+# This flags enables the `peanutbutter` realtime metrics backend.
+# See https://github.com/getsentry/peanutbutter.
+# We do not want/need this in normal devservices, but we need it for certain tests.
+SENTRY_USE_PEANUTBUTTER = False
+
 # SENTRY_DEVSERVICES = {
 #     "service-name": lambda settings, options: (
 #         {
@@ -3002,6 +3007,14 @@ def custom_parameter_sort(parameter: dict) -> tuple[str, int]:
             "only_if": settings.SENTRY_USE_SPOTLIGHT,
         }
     ),
+    "peanutbutter": lambda settings, options: (
+        {
+            "image": "us.gcr.io/sentryio/peanutbutter:latest",
+            "environment": {},
+            "ports": {"4433/tcp": 4433},
+            "only_if": settings.SENTRY_USE_PEANUTBUTTER,
+        }
+    ),
 }
 
 # Max file size for serialized file uploads in API

src/sentry/processing/realtime_metrics/pb.py

@@ -0,0 +1,75 @@
+import logging
+from collections.abc import Iterable
+from urllib.parse import urljoin
+
+from requests import RequestException
+
+from sentry.net.http import Session
+
+from . import base
+
+logger = logging.getLogger(__name__)
+
+# The timeout for rpc calls, in seconds.
+# We expect these to be very quick, and never want to block more than 2 ms (4 with connect + read).
+RPC_TIMEOUT = 2 / 1000  # timeout in seconds
+
+
+class PbRealtimeMetricsStore(base.RealtimeMetricsStore):
+    def __init__(self, target: str):
+        self.target = target
+        self.session = Session()
+
+    def record_project_duration(self, project_id: int, duration: float) -> None:
+        url = urljoin(self.target, "/record_spending")
+        request = {
+            "config_name": "symbolication-native",
+            "project_id": project_id,
+            "spent": duration,
+        }
+        try:
+            self.session.post(
+                url,
+                timeout=RPC_TIMEOUT,
+                json=request,
+            )
+        except RequestException:
+            pass
+
+    def is_lpq_project(self, project_id: int) -> bool:
+        url = urljoin(self.target, "/exceeds_budget")
+        request = {
+            "config_name": "symbolication-native",
+            "project_id": project_id,
+        }
+        try:
+            response = self.session.post(
+                url,
+                timeout=RPC_TIMEOUT,
+                json=request,
+            )
+            return response.json()["exceeds_budget"]
+        except RequestException:
+            return False
+
+    # NOTE: The functions below are just default impls copy-pasted from `DummyRealtimeMetricsStore`.
+    # They are not used in the actual implementation of recording budget spend,
+    # and checking if a project is within its budget.
+
+    def validate(self) -> None:
+        pass
+
+    def projects(self) -> Iterable[int]:
+        yield from ()
+
+    def get_used_budget_for_project(self, project_id: int) -> float:
+        return 0.0
+
+    def get_lpq_projects(self) -> set[int]:
+        return set()
+
+    def add_project_to_lpq(self, project_id: int) -> bool:
+        return False
+
+    def remove_projects_from_lpq(self, project_ids: set[int]) -> int:
+        return 0

tests/sentry/processing/realtime_metrics/test_pb.py

@@ -0,0 +1,20 @@
+from math import floor
+from random import random
+
+from sentry.processing.realtime_metrics.pb import PbRealtimeMetricsStore
+
+
+def test_invalid_target():
+    # there is no grpc service at that addr
+    store = PbRealtimeMetricsStore(target="http://localhost:12345")
+    store.record_project_duration(1, 123456789)
+    assert not store.is_lpq_project(1)
+
+
+def test_pb_works():
+    store = PbRealtimeMetricsStore(target="http://localhost:4433")
+
+    project_id = floor(random() * (1 << 32))
+    assert not store.is_lpq_project(project_id)
+    store.record_project_duration(project_id, 123456789)
+    assert store.is_lpq_project(project_id)