outpost: pull Polar's changes instead of holding a WebSocket

sdk/e2e/harness.py

@@ -14,11 +14,19 @@
 
 import anyio
 from anyio.streams.buffered import BufferedByteReceiveStream
-from polar.outpost.client import Configuration, EventInput, Snapshot
+from polar.outpost.client import EventInput, Snapshot
+from polar.outpost.reducer import Reducer
 from redis.asyncio import Redis
 from redis.exceptions import ConnectionError as RedisConnectionError
 
 ROOT = pathlib.Path(__file__).resolve().parent
+
+
+# Polar's configuration as a scenario writes it: the fake Polar adds its organization.
+class ScenarioConfiguration(typing.TypedDict):
+    reducers: list[Reducer]
+
+
 RETRYABLE_READ_ERRORS = ("Customer snapshot is not ready", "Redis is unavailable")
 
 
@@ -153,7 +161,7 @@ class IngestedStep(TimelineStep):
 @dataclasses.dataclass(frozen=True)
 class Scenario:
     name: str
-    configuration: Configuration
+    configuration: ScenarioConfiguration
     snapshots: list[Snapshot]
     steps: list[Step]
 

sdk/e2e/pyproject.toml

@@ -11,7 +11,6 @@ dependencies = [
     "pytest-xdist[psutil]>=3.8.0",
     "starlette>=0.47.0",
     "uvicorn>=0.35.0",
-    "wsproto>=1.3.0",
 ]
 
 [dependency-groups]

sdk/e2e/scenarios.py

@@ -1,6 +1,6 @@
 import itertools
 
-from polar.outpost.client import Configuration, Snapshot
+from polar.outpost.client import Snapshot
 from polar.outpost.reducer import BUCKET_SIZE, Reducer
 
 from harness import (
@@ -15,6 +15,7 @@
     RedisStep,
     RestartStep,
     Scenario,
+    ScenarioConfiguration,
     SnapshotStep,
 )
 
@@ -45,7 +46,7 @@ def credit_reducer(meter_id: str) -> Reducer:
     "filter": {"conjunction": "and", "clauses": []},
     "aggregation": {"func": "sum", "property": "tokens"},
 }
-CONFIGURATION: Configuration = {"reducers": [TOKENS, credit_reducer("meter")]}
+CONFIGURATION: ScenarioConfiguration = {"reducers": [TOKENS, credit_reducer("meter")]}
 _ids = itertools.count(1)
 CUSTOMER: ExternalCustomerID = {"external_customer_id": "alice"}
 METER: ExternalCustomerMeterRequest = {**CUSTOMER, "meter_id": "meter"}

sdk/e2e/snapshot_server.py

@@ -2,13 +2,15 @@
 import socket
 import sys
 import typing
+import uuid
 
 import anyio
 import uvicorn
 from polar.outpost.client import Configuration, Snapshot
 from starlette.applications import Starlette
-from starlette.routing import WebSocketRoute
-from starlette.websockets import WebSocket, WebSocketDisconnect, WebSocketState
+from starlette.requests import Request
+from starlette.responses import JSONResponse, Response
+from starlette.routing import Route
 
 from python_client import read_lines
 
@@ -19,116 +21,117 @@ def customer_key(payload: typing.Mapping[str, object]) -> str:
     return f"external:{payload['external_customer_id']}:{payload['reducer_id']}"
 
 
+# Polar's pull interval here: every read finds the Redis stale and pulls first, so a snapshot
+# reaches the clients as fast as a push did.
+PULL_INTERVAL = 0.05
+
+
 class SnapshotServer:
     def __init__(self) -> None:
-        self.configuration: Configuration = {"reducers": []}
+        # Its own organization: scenarios sharing a Redis keep their keys apart.
+        self.organization_id = uuid.uuid4().hex
+        self.configuration: Configuration = {
+            "organization_id": self.organization_id,
+            "reducers": [],
+        }
         self.snapshots: dict[str, Snapshot] = {}
+        # Polar's changes: a pull takes those after its cursor, the index of the last one taken.
+        self.changes: list[Snapshot] = []
         # The external ids Polar stored: a confirm request is answered with those among them.
         self.ingested: set[str] = set()
-        self.clients: dict[WebSocket, anyio.Lock] = {}
-
-    async def websocket(self, websocket: WebSocket) -> None:
-        if websocket.headers.get("authorization") != "Bearer harness-token":
-            await websocket.close(code=1008)
-            return
-        await websocket.accept()
-        self.clients[websocket] = anyio.Lock()
-        try:
-            while True:
-                message = await websocket.receive_json()
-                if message["type"] == "configuration":
-                    await self._send(websocket, "configuration", self.configuration)
-                elif message["type"] == "meter_usage":
-                    payload = message["payload"]
-                    snapshot = self.snapshots.get(customer_key(payload)) or {
-                        "customer_id": payload.get("customer_id"),
-                        "external_customer_id": payload.get("external_customer_id"),
-                        "customer_exists": False,
-                        "reducer_id": payload["reducer_id"],
-                        "version": 0,
-                        "cold_until": 0,
-                        "count": 0,
-                        "sum": 0,
-                        "min": None,
-                        "max": None,
-                    }
-                    await self._send(websocket, "meter_usage", snapshot)
-                elif message["type"] == "confirm":
-                    payload = message["payload"]
-                    snapshot = self.snapshots[customer_key(payload)]
-                    await self._send(
-                        websocket,
-                        "confirmed",
-                        {
-                            "reducer_id": payload["reducer_id"],
-                            "customer_id": payload.get("customer_id"),
-                            "external_customer_id": payload.get("external_customer_id"),
-                            "bucket_start": payload["bucket_start"],
-                            "ids": [
-                                external_id
-                                for external_id in payload["ids"]
-                                if external_id in self.ingested
-                            ],
-                            "snapshot": snapshot,
-                        },
-                    )
-                else:
-                    await websocket.close(code=1008, reason="Unknown message")
-                    return
-        except WebSocketDisconnect:
-            pass
-        except (ValueError, KeyError, TypeError):
-            await websocket.close(code=1008, reason="Invalid message")
-        finally:
-            self.clients.pop(websocket, None)
+        # Polar is unreachable for this many requests.
+        self.dropped = 0
+
+    async def handle(self, request: Request) -> Response:
+        if request.headers.get("authorization") != "Bearer harness-token":
+            return JSONResponse({}, status_code=401)
+        if self.dropped:
+            self.dropped -= 1
+            return JSONResponse({}, status_code=503)
+        route = f"{request.method} {request.url.path}"
+        if route == "GET /v1/outpost/configuration":
+            return JSONResponse(self.configuration)
+        if route == "POST /v1/outpost/meter-usage":
+            payload = await request.json()
+            return JSONResponse([self._snapshot(payload)])
+        if route == "POST /v1/outpost/confirm":
+            payload = await request.json()
+            return JSONResponse(
+                {
+                    "reducer_id": payload["reducer_id"],
+                    "customer_id": payload.get("customer_id"),
+                    "external_customer_id": payload.get("external_customer_id"),
+                    "bucket_start": payload["bucket_start"],
+                    "ids": [
+                        external_id
+                        for external_id in payload["ids"]
+                        if external_id in self.ingested
+                    ],
+                    "snapshot": self.snapshots[customer_key(payload)],
+                }
+            )
+        if route == "GET /v1/outpost/changes":
+            cursor = request.query_params.get("cursor")
+            start = int(cursor.split("-")[0]) if cursor else len(self.changes)
+            return JSONResponse(
+                {
+                    "changes": self.changes[start:],
+                    "cursor": f"{len(self.changes)}-0",
+                    "has_more": False,
+                    "reset": False,
+                    "configuration_changed": False,
+                    "next_poll_in": PULL_INTERVAL,
+                }
+            )
+        return JSONResponse({}, status_code=404)
 
     async def command(self, operation: str, arguments: dict) -> None:
         if operation == "configure":
-            self.configuration = arguments["configuration"]
+            self.configuration = {
+                **arguments["configuration"],
+                "organization_id": self.organization_id,
+            }
             for snapshot in arguments["snapshots"]:
                 self.snapshots[customer_key(snapshot)] = snapshot
         elif operation == "disconnect":
-            await self.close(code=1012)
+            self.dropped = 1
         elif operation == "snapshot":
             snapshot = typing.cast(Snapshot, arguments)
             self.snapshots[customer_key(snapshot)] = snapshot
-            for websocket in list(self.clients):
-                await self._send(websocket, "meter_usage", snapshot)
+            self.changes.append(snapshot)
         elif operation == "ingested":
             self.ingested.update(event["external_id"] for event in arguments["events"])
         else:
             raise ValueError(f"Unknown operation: {operation}")
 
-    async def close(self, code: int = 1001) -> None:
-        for websocket, lock in list(self.clients.items()):
-            async with lock:
-                if websocket.application_state == WebSocketState.CONNECTED:
-                    await websocket.close(code=code)
-
-    async def _send(
-        self, websocket: WebSocket, type: str, payload: typing.Mapping[str, object]
-    ) -> None:
-        lock = self.clients.get(websocket)
-        if lock is None:
-            return
-        try:
-            async with lock:
-                if websocket.application_state == WebSocketState.CONNECTED:
-                    await websocket.send_json({"type": type, "payload": payload})
-        except (WebSocketDisconnect, OSError):
-            pass
+    def _snapshot(self, payload: typing.Mapping[str, typing.Any]) -> Snapshot:
+        return self.snapshots.get(customer_key(payload)) or {
+            "customer_id": payload.get("customer_id"),
+            "external_customer_id": payload.get("external_customer_id"),
+            "customer_exists": False,
+            "reducer_id": payload["reducer_id"],
+            "version": 0,
+            "cold_until": 0,
+            "count": 0,
+            "sum": 0,
+            "min": None,
+            "max": None,
+        }
 
 
 async def main() -> None:
     fixture = SnapshotServer()
-    app = Starlette(routes=[WebSocketRoute("/v1/outpost/", fixture.websocket)])
+    app = Starlette(
+        routes=[
+            Route("/v1/outpost/{path:path}", fixture.handle, methods=["GET", "POST"])
+        ]
+    )
     server = uvicorn.Server(
         uvicorn.Config(
             app,
             log_config=None,
             log_level="error",
             access_log=False,
-            ws="wsproto",
             timeout_graceful_shutdown=5,
         )
     )
@@ -160,7 +163,6 @@ async def serve() -> None:
                     except (ValueError, KeyError, TypeError) as error:
                         print(json.dumps({"error": str(error)}), flush=True)
             finally:
-                await fixture.close()
                 server.should_exit = True
 
 

sdk/generator/generator/emitter.py

@@ -3,6 +3,7 @@
 import pathlib
 import re
 import subprocess
+import tomllib
 import typing
 from typing import Literal
 
@@ -11,7 +12,8 @@
 from generator.ir import APIIR, APIVersion, LiteralType
 
 _LOOP_DIRECTORY_PATTERN = re.compile(r"^\[\[([\w\.]+)\]\]$")
-# The Outpost Lua scripts, which every SDK embeds in a string literal.
+# The Outpost Lua scripts, which every SDK embeds in a string literal, and the constants
+# every SDK shares.
 OUTPOST_SCRIPTS_DIRECTORY = pathlib.Path(__file__).parent.parent / "outpost"
 _LITERAL_BREAKERS = ("`", "${", '"""', "\\")
 
@@ -28,6 +30,10 @@ def load_outpost_scripts() -> dict[str, str]:
     return scripts
 
 
+def load_outpost_constants() -> dict[str, typing.Any]:
+    return tomllib.loads((OUTPOST_SCRIPTS_DIRECTORY / "constants.toml").read_text())
+
+
 PrereleaseLabel = Literal["alpha", "beta", "rc"]
 _PRERELEASE_PATTERN = re.compile(r"^(alpha|beta|rc)\.(\d+)$")
 
@@ -103,6 +109,7 @@ def get_context(self) -> dict[str, typing.Any]:
             "ir": self.ir,
             "version": self.format_version(),
             "outpost_scripts": load_outpost_scripts(),
+            "outpost_constants": load_outpost_constants(),
         }
 
     def get_version_context(self, api: APIVersion) -> dict[str, typing.Any]: