usagerail 0.1.0__py3-none-any.whl

usagerail/__init__.py

@@ -0,0 +1,43 @@
+"""UsageRail 的 Python SDK：在你自己的 API 服务里发 key、限流、管额度、计量（PRD 5.1）。
+
+FastAPI / Starlette 用 usagerail.asgi，Flask 用 usagerail.flask，Django 用 usagerail.django。
+自己接别的框架时用 create_engine：begin() 拿到结论，请求结束调 end()。
+"""
+
+from .engine import (
+    MAX_TRACK_UNITS,
+    Admission,
+    ConsumerIdentity,
+    Decision,
+    Engine,
+    Pending,
+    Rejection,
+    RejectionCode,
+    RequestHandle,
+    create_engine,
+    read_api_key,
+    rejection_body,
+)
+from .key_store import SnapshotStore
+from .snapshot import FileSnapshotStore
+from .transport import HttpTransport, Transport, TransportError
+
+__all__ = [
+    "MAX_TRACK_UNITS",
+    "Admission",
+    "ConsumerIdentity",
+    "Decision",
+    "Engine",
+    "FileSnapshotStore",
+    "HttpTransport",
+    "Pending",
+    "Rejection",
+    "RejectionCode",
+    "RequestHandle",
+    "SnapshotStore",
+    "Transport",
+    "TransportError",
+    "create_engine",
+    "read_api_key",
+    "rejection_body",
+]

usagerail/_adapters.py

@@ -0,0 +1,14 @@
+"""三个框架适配器共用的一点东西。"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from .engine import Engine, create_engine
+from .snapshot import FileSnapshotStore
+
+
+def engine_for(app_key: str, auth: str, snapshot: bool, options: dict[str, Any]) -> Engine[Any]:
+    """key 快照只对用我们 key 的应用有意义，默认落盘到本地临时目录（PRD 5.3）"""
+    store = FileSnapshotStore(app_key) if snapshot and auth == "api-key" and app_key else None
+    return create_engine(app_key=app_key, auth=auth, snapshot_store=store, **options)

usagerail/_futures.py

@@ -0,0 +1,100 @@
+"""把 concurrent.futures.Future 当成 Node 的 Promise 用的几个小工具。
+
+引擎里走网络的地方（远程校验、领租约、同步、上报）都返回 Future：Flask、Django 的线程里阻塞等它，
+FastAPI 里 await 它。等的时候不占线程，几百个请求同时等租约也不会把线程池耗光（PRD「最多多等 300 毫秒」）。
+回调在完成它的那个线程里跑，已经完成的 Future 加回调会当场跑，所以引擎的状态一律在锁里改。
+"""
+
+from __future__ import annotations
+
+import asyncio
+from concurrent.futures import Future, InvalidStateError
+from typing import Any, Callable, Generator, Generic, TypeVar
+
+T = TypeVar("T")
+U = TypeVar("U")
+
+
+def resolved(value: T) -> Future[T]:
+    future: Future[T] = Future()
+    future.set_result(value)
+    return future
+
+
+def settle(future: Future[T], value: T) -> None:
+    """已经有结果的不再改（超时和完成同时到时，先到的算）"""
+    try:
+        future.set_result(value)
+    except InvalidStateError:
+        pass
+
+
+def fail(future: Future[Any], error: BaseException) -> None:
+    try:
+        future.set_exception(error)
+    except InvalidStateError:
+        pass
+
+
+def error_of(future: Future[Any]) -> BaseException | None:
+    """已经完成的 Future 的异常；取消的当作出错"""
+    if future.cancelled():
+        return RuntimeError("cancelled")
+    return future.exception()
+
+
+def then(
+    future: Future[T],
+    on_value: Callable[[T], U | Future[U]],
+    on_error: Callable[[BaseException], U | Future[U]] | None = None,
+) -> Future[U]:
+    """Promise.then：回调可以返回值，也可以返回 Future（会展开）"""
+    out: Future[U] = Future()
+
+    def run(callback: Callable[[Any], U | Future[U]], argument: Any) -> None:
+        try:
+            result = callback(argument)
+        except BaseException as error:  # noqa: BLE001 - 原样交给下一环
+            fail(out, error)
+            return
+        if isinstance(result, Future):
+            result.add_done_callback(lambda inner: _copy(inner, out))
+        else:
+            settle(out, result)
+
+    def done(completed: Future[T]) -> None:
+        error = error_of(completed)
+        if error is None:
+            run(on_value, completed.result())
+        elif on_error is not None:
+            run(on_error, error)
+        else:
+            fail(out, error)
+
+    future.add_done_callback(done)
+    return out
+
+
+def _copy(source: Future[Any], target: Future[Any]) -> None:
+    error = error_of(source)
+    if error is None:
+        settle(target, source.result())
+    else:
+        fail(target, error)
+
+
+class Pending(Generic[T]):
+    """还要等服务端的结论。WSGI 里 .result() 阻塞等；ASGI 里直接 await"""
+
+    __slots__ = ("future",)
+
+    def __init__(self, future: Future[T]) -> None:
+        self.future = future
+
+    def result(self, timeout: float | None = None) -> T:
+        return self.future.result(timeout)
+
+    def __await__(self) -> Generator[Any, None, T]:
+        # shield：等的一方被取消（客户端走了）不能连带取消底下的 Future。引擎还会给它结论（比如领到的租约），
+        # 取消了结论就没处放，预留一直挂着；调用方自己决定被取消以后怎么收尾
+        return asyncio.shield(asyncio.wrap_future(self.future)).__await__()

usagerail/_time.py

@@ -0,0 +1,24 @@
+"""ISO 8601 时间和毫秒时间戳之间的换算，格式和 JavaScript 的 Date 一致（服务端和 Node SDK 用的就是它）。"""
+
+from __future__ import annotations
+
+from datetime import datetime, timezone
+
+
+def parse_iso_ms(value: object) -> float | None:
+    """读不懂的返回 None。没带时区的按 UTC 算（服务端总是带 Z）"""
+    if not isinstance(value, str) or not value:
+        return None
+    text = value[:-1] + "+00:00" if value.endswith(("Z", "z")) else value
+    try:
+        parsed = datetime.fromisoformat(text)
+    except ValueError:
+        return None
+    if parsed.tzinfo is None:
+        parsed = parsed.replace(tzinfo=timezone.utc)
+    return parsed.timestamp() * 1000
+
+
+def iso_from_ms(ms: float) -> str:
+    """2026-10-20T12:00:10.000Z，和 new Date(ms).toISOString() 一样"""
+    return datetime.fromtimestamp(ms / 1000, tz=timezone.utc).isoformat(timespec="milliseconds").replace("+00:00", "Z")

usagerail/clock.py

@@ -0,0 +1,95 @@
+"""时钟。SDK 里所有「现在」都从这里取，测试注入假时钟，不用真的等（docs/decisions.md D4）。"""
+
+from __future__ import annotations
+
+import heapq
+import itertools
+import os
+import threading
+import time
+from typing import Callable, Protocol
+
+
+class Clock(Protocol):
+    def now(self) -> float:
+        """墙上时间（毫秒）：分钟归属、key 过期、周期边界。这些要和服务端、Stripe 的时间对得上"""
+        ...
+
+    def monotonic(self) -> float:
+        """单调时间（毫秒）：租约本地有效期、令牌桶、退避、各种缓存时长。机器改时间、NTP 校时都不影响（PRD 5.5）"""
+        ...
+
+    def set_timer(self, callback: Callable[[], None], delay_ms: float) -> Callable[[], None]:
+        """定时器，返回取消函数。后台的同步、上报、续租都走它，测试里由假时钟推进"""
+        ...
+
+
+class _Timer:
+    __slots__ = ("callback", "cancelled")
+
+    def __init__(self, callback: Callable[[], None]) -> None:
+        self.callback = callback
+        self.cancelled = False
+
+    def cancel(self) -> None:
+        self.cancelled = True
+
+
+class SystemClock:
+    """所有定时器共用一个后台线程，按到期时间排队跑，不是一个定时器一个线程：等租约的请求每个都有一个定时器。
+    线程是 daemon 的，不拖住进程退出（同 Node 的 unref）。回调要快，走网络的事交给 transport 的线程。"""
+
+    def __init__(self) -> None:
+        self._reset()
+        # gunicorn --preload 这类先建好应用再 fork 的：子进程里没有父进程的线程，锁也可能停在被持有的状态。
+        # 子进程重新来过；父进程排着的定时器属于父进程的引擎，引擎自己会在子进程里重新排（engine.py）
+        # Windows 没有 fork，也没有这个函数
+        if hasattr(os, "register_at_fork"):
+            os.register_at_fork(after_in_child=self._reset)
+
+    def _reset(self) -> None:
+        self._condition = threading.Condition()
+        self._queue: list[tuple[float, int, _Timer]] = []
+        self._order = itertools.count()
+        self._thread: threading.Thread | None = None
+
+    def now(self) -> float:
+        return time.time() * 1000
+
+    def monotonic(self) -> float:
+        return time.monotonic() * 1000
+
+    def set_timer(self, callback: Callable[[], None], delay_ms: float) -> Callable[[], None]:
+        timer = _Timer(callback)
+        due = time.monotonic() + max(0.0, delay_ms) / 1000
+        with self._condition:
+            heapq.heappush(self._queue, (due, next(self._order), timer))
+            if self._thread is None or not self._thread.is_alive():
+                self._thread = threading.Thread(target=self._run, name="usagerail-timers", daemon=True)
+                self._thread.start()
+            self._condition.notify()
+        return timer.cancel
+
+    def _run(self) -> None:
+        while True:
+            with self._condition:
+                while True:
+                    while self._queue and self._queue[0][2].cancelled:
+                        heapq.heappop(self._queue)
+                    if not self._queue:
+                        self._condition.wait()
+                        continue
+                    wait = self._queue[0][0] - time.monotonic()
+                    if wait <= 0:
+                        timer = heapq.heappop(self._queue)[2]
+                        break
+                    self._condition.wait(wait)
+            if timer.cancelled:
+                continue
+            try:
+                timer.callback()
+            except Exception:  # noqa: BLE001 - 回调出错不能让定时器线程死掉，引擎自己会打警告
+                pass
+
+
+system_clock = SystemClock()

usagerail/django.py

@@ -0,0 +1,152 @@
+"""Django 中间件（PRD 5.1）。settings.py 里：
+
+    MIDDLEWARE = [..., "usagerail.django.UsageRailMiddleware", ...]
+    USAGERAIL = {"app_key": os.environ["USAGERAIL_APP_KEY"], "auth": "api-key"}
+
+视图里按自定义单位计量：request.usagerail.track(tokens)
+跑在 uvicorn 这类 ASGI 服务器上时，asgi.py 里用 asgi() 包一层（见 asgi 的说明），不然停机时最后一批用量会丢。
+用自己的鉴权时："auth": "custom", "consumer": 一个函数，或者它的点分路径（"myapp.usage.consumer"），收到 request 返回是谁。
+要用 request.user 的，把这个中间件放在 AuthenticationMiddleware 后面。
+USAGERAIL 里其余的键原样交给 Engine（on_reject 收 (rejection, request)，返回一个 HttpResponse）。
+"""
+
+from __future__ import annotations
+
+import asyncio
+import threading
+import weakref
+from collections.abc import Awaitable, MutableMapping
+from typing import Any, Callable
+
+from django.conf import settings
+from django.http import HttpRequest, HttpResponse, JsonResponse
+from django.urls import Resolver404, resolve
+from django.utils.module_loading import import_string
+
+from . import _adapters
+from ._futures import Pending
+from .engine import Engine, Rejection, read_api_key, rejection_body
+
+# 一个进程一个引擎：Django 每加载一次中间件链就建一个中间件实例（测试里常常不止一次），引擎要共用
+_engines: dict[str, Engine[HttpRequest]] = {}
+_engines_lock = threading.Lock()
+# 中间件实际在用的引擎（包括 USAGERAIL 里直接传进来的），asgi() 停机时逐个关
+_middleware_engines: weakref.WeakSet[Engine[Any]] = weakref.WeakSet()
+
+Scope = MutableMapping[str, Any]
+Message = MutableMapping[str, Any]
+Receive = Callable[[], Awaitable[Message]]
+Send = Callable[[Message], Awaitable[None]]
+ASGIApp = Callable[[Scope, Receive, Send], Awaitable[None]]
+
+
+class UsageRailMiddleware:
+    def __init__(self, get_response: Callable[[HttpRequest], HttpResponse]) -> None:
+        self.get_response = get_response
+        config = dict(getattr(settings, "USAGERAIL", {}))
+        on_reject = config.pop("on_reject", None)
+        self.on_reject: Callable[[Rejection, HttpRequest], HttpResponse] | None = import_string(on_reject) if isinstance(on_reject, str) else on_reject
+        self.engine: Engine[Any] = config.pop("engine", None) or _shared_engine(config)
+        _middleware_engines.add(self.engine)
+
+    def __call__(self, request: HttpRequest) -> HttpResponse:
+        route = _template(request)
+        decision = self.engine.begin(
+            method=request.method or "GET",
+            path=request.path_info,
+            # 先按 Django 的 URL 配置解析一遍：路由模板开始时就知道，re_path 的正则模板也能精确对上卖家的配置
+            route=route,
+            api_key=read_api_key(request.headers.get("Authorization"), request.headers.get("X-Api-Key")),
+            context=request,
+        )
+        if isinstance(decision, Pending):
+            decision = decision.result()
+        if not decision.allowed:
+            return self.on_reject(decision, request) if self.on_reject else _rejection_response(decision)
+        handle = decision.handle
+        request.usagerail = handle
+        try:
+            response = self.get_response(request)
+        except BaseException:
+            # 一般走不到：Django 自己把视图的异常变成 500 响应。别的中间件抛出来的才会到这里
+            self.engine.end(handle, status=500, route=route)
+            raise
+        engine, status = self.engine, response.status_code
+        # SSE 一开始发送就算建立了一个长连接，当场计 1 次（v0.3 5.1），单位数等结束再补
+        if status < 300 and response.get("Content-Type", "").lower().startswith("text/event-stream"):
+            engine.stream_started(handle, route)
+        # 等服务器把响应发完、关掉它（流式响应发完最后一块）再结算。挂在 request.close 上：Django 把它挂进最终那个响应的
+        # 关闭回调（BaseHandler.get_response），外层的中间件把响应换掉了（CommonMiddleware 把 404 换成补斜杠的重定向）
+        # 也照样会跑。挂在这个响应自己的 close 上的话，它一被换掉就永远不跑，用量不记、预留一直挂着。
+        # 换掉的响应按这里看到的状态码记
+        close_request = request.close
+
+        def close_and_end() -> None:
+            try:
+                close_request()
+            finally:
+                engine.end(handle, status=status, route=route)
+
+        request.close = close_and_end
+        return response
+
+
+def _shared_engine(config: dict[str, Any]) -> Engine[HttpRequest]:
+    app_key = config.pop("app_key", "")
+    auth = config.pop("auth", "api-key")
+    snapshot = config.pop("snapshot", True)
+    consumer = config.get("consumer")
+    if isinstance(consumer, str):
+        config["consumer"] = import_string(consumer)
+    with _engines_lock:
+        engine = _engines.get(app_key)
+        if engine is None:
+            engine = _adapters.engine_for(app_key, auth, snapshot, config)
+            _engines[app_key] = engine
+        return engine
+
+
+def _template(request: HttpRequest) -> str | None:
+    """Django 的路由不带开头的斜杠，补上；没匹配到的为 None（结束时记 _unmatched）"""
+    try:
+        match = resolve(request.path_info, getattr(request, "urlconf", None))
+    except Resolver404:
+        return None
+    return "/" + match.route if match.route is not None else None
+
+
+def _rejection_response(rejection: Rejection) -> HttpResponse:
+    response = JsonResponse(rejection_body(rejection), status=rejection.status)
+    for name, value in rejection.headers.items():
+        response[name] = value
+    return response
+
+
+def asgi(application: ASGIApp) -> ASGIApp:
+    """Django 跑在 ASGI 服务器上（uvicorn 等）时包一层，asgi.py 里：
+
+        from django.core.asgi import get_asgi_application
+        from usagerail.django import asgi
+
+        application = asgi(get_asgi_application())
+
+    Django 自己不接 lifespan，uvicorn 收完尾又会把 SIGTERM 再抛一次，进程直接结束、atexit 不跑，不包的话最后一批用量会丢。
+    包了以后由这里应答 lifespan：停机时先把剩下的用量报上去（最多等 5 秒），再告诉服务器停好了，和 FastAPI 的中间件一样。
+    gunicorn 跑 WSGI 用不着：它收到 SIGTERM 正常退出，atexit 会报"""
+
+    async def app(scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] != "lifespan":
+            await application(scope, receive, send)
+            return
+        while True:
+            message = await receive()
+            if message["type"] == "lifespan.startup":
+                await send({"type": "lifespan.startup.complete"})
+            elif message["type"] == "lifespan.shutdown":
+                # 几个引擎同时关、一起等，总共还是最多 5 秒。报不上去也要让服务器停下来，所以不管出错
+                closing = [asyncio.wrap_future(engine.close()) for engine in list(_middleware_engines)]
+                await asyncio.gather(*closing, return_exceptions=True)
+                await send({"type": "lifespan.shutdown.complete"})
+                return
+
+    return app