mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-31 06:56:59 +00:00
feat(accounting): add bounded transactional usage persistence (#49190)
* feat(accounting): add bounded transactional usage persistence * fix(accounting): preserve usage flush and shutdown on Python 3.10 * test(accounting): cover batch rejection and partial shutdown recovery * fix(accounting): batch durable usage on a five-second interval * fix(accounting): drain queued full batches without timer delays * fix(accounting): isolate invalid producer submissions within usage batches * fix(accounting): await commit prerequisites within bounded usage admission --------- Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
This commit is contained in:
joshua-berri
GitHub
parent
e6b4774ac6
commit
02a6a96dec
7 changed files with 876 additions and 2 deletions
|
|
@ -31,6 +31,13 @@ class AccountingConflictError(ValueError):
|
|||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class UsageWrite:
|
||||
event: InboundUsageEvent | OutboundUsageEvent
|
||||
envelope: str
|
||||
targets: tuple[DeliveryTarget, ...] = ()
|
||||
|
||||
|
||||
_Transaction: TypeAlias = AbstractAsyncContextManager["Prisma"]
|
||||
|
||||
|
||||
|
|
@ -107,6 +114,35 @@ SELECT event_id FROM "LiteLLM_AccountingEvent" WHERE kind = 'cost' AND source_ev
|
|||
_LOCK_EVENT: Final = """
|
||||
SELECT payload_hash FROM "LiteLLM_AccountingEvent" WHERE event_id = $1 FOR UPDATE
|
||||
"""
|
||||
_INSERT_USAGE_BATCH: Final = """
|
||||
INSERT INTO "LiteLLM_AccountingEvent"
|
||||
(event_id, kind, schema_version, payload_hash, occurred_at, payload)
|
||||
SELECT event_id, kind::"AccountingEventKind", 1, payload_hash, occurred_at::timestamp, payload
|
||||
FROM jsonb_to_recordset($1::jsonb) AS records(
|
||||
event_id text, kind text, payload_hash text, occurred_at text, payload jsonb)
|
||||
ORDER BY event_id
|
||||
ON CONFLICT (event_id) DO NOTHING RETURNING event_id
|
||||
"""
|
||||
_LOCK_USAGE_BATCH: Final = """
|
||||
SELECT event_id, payload_hash FROM "LiteLLM_AccountingEvent"
|
||||
WHERE event_id IN (SELECT value FROM jsonb_array_elements_text($1::jsonb))
|
||||
ORDER BY event_id FOR UPDATE
|
||||
"""
|
||||
_LOAD_USAGE_BATCH: Final = """
|
||||
SELECT event.payload, ARRAY(
|
||||
SELECT target::text FROM "LiteLLM_AccountingDelivery" delivery
|
||||
WHERE delivery.event_id = event.event_id ORDER BY target::text
|
||||
) AS targets
|
||||
FROM "LiteLLM_AccountingEvent" event
|
||||
WHERE event.event_id IN (SELECT value FROM jsonb_array_elements_text($1::jsonb))
|
||||
AND event.kind = 'outbound_usage' AND event.schema_version = 1
|
||||
"""
|
||||
_INSERT_USAGE_DELIVERIES: Final = """
|
||||
INSERT INTO "LiteLLM_AccountingDelivery" (event_id, target, version)
|
||||
SELECT records.event_id, destination.value::"AccountingDeliveryTarget", 1
|
||||
FROM jsonb_to_recordset($1::jsonb) AS records(event_id text, targets jsonb)
|
||||
CROSS JOIN LATERAL jsonb_array_elements_text(records.targets) AS destination
|
||||
"""
|
||||
_COMPLETED_FAMILY: Final = """
|
||||
root.source_event_id IS NULL AND root.kind IN ('inbound_usage', 'outbound_usage')
|
||||
AND root.recorded_at < $1::timestamp
|
||||
|
|
@ -199,6 +235,53 @@ class AccountingLedger:
|
|||
|
||||
return cls(transaction)
|
||||
|
||||
async def insert_usage_batch(self, writes: tuple[UsageWrite, ...]) -> None:
|
||||
records: Final = tuple(_prepare(write.event, write.targets, write.envelope, 1, 1) for write in writes)
|
||||
if not records:
|
||||
return
|
||||
unique: Final = {record.event_id: record for record in records}
|
||||
if any(record.payload_hash != unique[record.event_id].payload_hash for record in records):
|
||||
raise AccountingConflictError("Batch contains conflicting copies of a usage identity")
|
||||
batch: Final = _json(
|
||||
tuple(
|
||||
{
|
||||
"event_id": record.event_id,
|
||||
"kind": record.kind.value,
|
||||
"payload_hash": record.payload_hash,
|
||||
"occurred_at": record.occurred_at.astimezone(timezone.utc).isoformat(),
|
||||
"payload": _JSON.validate_json(record.payload),
|
||||
}
|
||||
for record in unique.values()
|
||||
)
|
||||
)
|
||||
async with self._transaction() as db:
|
||||
inserted: Final = _STRINGS.validate_python(await db.query_raw(_INSERT_USAGE_BATCH, batch))
|
||||
accepted: Final = _STRINGS.validate_python(await db.query_raw(_LOCK_USAGE_BATCH, _json(tuple(unique))))
|
||||
if len(accepted) != len(unique) or any(
|
||||
row["payload_hash"] != unique[row["event_id"]].payload_hash for row in accepted
|
||||
):
|
||||
raise AccountingConflictError("Usage identity conflicts with accepted content or delivery set")
|
||||
deliveries: Final = _json(
|
||||
tuple({"event_id": row["event_id"], "targets": unique[row["event_id"]].targets} for row in inserted)
|
||||
)
|
||||
await db.execute_raw(_INSERT_USAGE_DELIVERIES, deliveries)
|
||||
|
||||
async def load_usage_batch(self, event_ids: tuple[str, ...]) -> tuple[UsageWrite, ...]:
|
||||
if not event_ids:
|
||||
return ()
|
||||
async with self._transaction() as db:
|
||||
rows: Final = TypeAdapter(tuple[dict[str, JsonValue], ...]).validate_python(
|
||||
await db.query_raw(_LOAD_USAGE_BATCH, _json(event_ids))
|
||||
)
|
||||
return tuple(
|
||||
UsageWrite(
|
||||
OutboundUsageEvent.model_validate(_JSON.validate_python(row["payload"])["event"]),
|
||||
_json(row["payload"]),
|
||||
TypeAdapter(tuple[DeliveryTarget, ...]).validate_python(row["targets"]),
|
||||
)
|
||||
for row in rows
|
||||
)
|
||||
|
||||
async def insert_event(
|
||||
self,
|
||||
event: AccountingEvent,
|
||||
|
|
|
|||
|
|
@ -72,6 +72,22 @@ def _exception_types(*candidates: object) -> tuple[type[BaseException], ...]:
|
|||
return tuple(c for c in candidates if isinstance(c, type) and issubclass(c, BaseException))
|
||||
|
||||
|
||||
def _prisma_error_code(e: Exception) -> str | None:
|
||||
import prisma.engine.errors
|
||||
|
||||
if isinstance(e, _exception_types(prisma.errors.DataError)):
|
||||
code: Final[object] = getattr(e, "code", None)
|
||||
return code if isinstance(code, str) else None
|
||||
if not isinstance(e, _exception_types(prisma.engine.errors.EngineRequestError)):
|
||||
return None
|
||||
try:
|
||||
body: Final = _DATABASE_ERROR_META.validate_json(str(e).partition(": ")[2])
|
||||
except ValidationError:
|
||||
return None
|
||||
engine_code: Final = body.get("error_code")
|
||||
return engine_code if isinstance(engine_code, str) else None
|
||||
|
||||
|
||||
class PrismaDBExceptionHandler:
|
||||
"""
|
||||
Class to handle DB Exceptions or Connection Errors
|
||||
|
|
@ -206,6 +222,8 @@ class PrismaDBExceptionHandler:
|
|||
return True
|
||||
if PrismaDBExceptionHandler.is_database_capacity_error(e):
|
||||
return False
|
||||
if _prisma_error_code(e) == "P1017":
|
||||
return True
|
||||
if isinstance(e, _exception_types(prisma.errors.PrismaError)):
|
||||
error_message: Final = str(e).lower()
|
||||
connection_keywords: Final = (
|
||||
|
|
|
|||
268
litellm/proxy/spend_tracking/usage_meter.py
Normal file
268
litellm/proxy/spend_tracking/usage_meter.py
Normal file
|
|
@ -0,0 +1,268 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections import deque
|
||||
from collections.abc import Awaitable, Callable
|
||||
from concurrent.futures import Future
|
||||
from dataclasses import dataclass
|
||||
from itertools import accumulate, chain, count, takewhile
|
||||
from threading import Lock
|
||||
from typing import Final
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.accounting import (
|
||||
DeliveryTarget,
|
||||
Dimension,
|
||||
DimensionName,
|
||||
InboundUsageEvent,
|
||||
OutboundUsageEvent,
|
||||
)
|
||||
from litellm.proxy.db.accounting_ledger import AccountingConflictError, UsageWrite
|
||||
|
||||
_DAILY_TARGETS: Final = (
|
||||
(DimensionName.KEY, DeliveryTarget.DAILY_USER_SPEND),
|
||||
(DimensionName.USER, DeliveryTarget.DAILY_USER_SPEND),
|
||||
(DimensionName.TEAM, DeliveryTarget.DAILY_TEAM_SPEND),
|
||||
(DimensionName.ORGANIZATION, DeliveryTarget.DAILY_ORGANIZATION_SPEND),
|
||||
(DimensionName.END_USER, DeliveryTarget.DAILY_END_USER_SPEND),
|
||||
(DimensionName.AGENT, DeliveryTarget.DAILY_AGENT_SPEND),
|
||||
(DimensionName.TAG, DeliveryTarget.DAILY_TAG_SPEND),
|
||||
(DimensionName.MODEL, DeliveryTarget.DAILY_MODEL_USAGE),
|
||||
(DimensionName.TOOL, DeliveryTarget.DAILY_TOOL_SPEND),
|
||||
)
|
||||
|
||||
|
||||
def usage_delivery_targets(dimensions: tuple[Dimension, ...]) -> tuple[DeliveryTarget, ...]:
|
||||
names: Final = frozenset(dimension.name for dimension in dimensions)
|
||||
return tuple(dict.fromkeys(target for dimension, target in _DAILY_TARGETS if dimension in names))
|
||||
|
||||
|
||||
class UsageCaptureSettings(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
batch_size: int = Field(default=1000, ge=2, le=1000)
|
||||
flush_interval_seconds: float = Field(default=5.0, gt=0, le=10)
|
||||
max_pending_bytes: int = Field(default=16 * 1024 * 1024, ge=1024)
|
||||
max_pending_records: int = Field(default=10000, ge=2)
|
||||
|
||||
|
||||
class UsageCaptureError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Pending:
|
||||
writes: tuple[UsageWrite, ...]
|
||||
size: int
|
||||
committed: Future[None]
|
||||
finished: asyncio.Event
|
||||
prerequisites: tuple[asyncio.Future[None], ...]
|
||||
|
||||
|
||||
async def _wait_for_batch(wakeup: asyncio.Event, timeout: float) -> None:
|
||||
try:
|
||||
await asyncio.wait_for(wakeup.wait(), timeout)
|
||||
except asyncio.TimeoutError:
|
||||
return
|
||||
|
||||
|
||||
def _report_failure(receipt: Future[None]) -> None:
|
||||
if not receipt.cancelled() and receipt.exception() is not None:
|
||||
verbose_proxy_logger.error("Usage capture did not acknowledge commit", exc_info=receipt.exception())
|
||||
|
||||
|
||||
def _retryable_persistence_failure(error: Exception) -> bool:
|
||||
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler, exception_chain
|
||||
|
||||
outage: Final = PrismaDBExceptionHandler.find_database_service_unavailable_error_in_chain(error)
|
||||
return (outage is not None and not PrismaDBExceptionHandler.is_permanent_database_fault(outage)) or any(
|
||||
isinstance(cause, Exception)
|
||||
and (
|
||||
PrismaDBExceptionHandler.is_deadlock_error(cause)
|
||||
or PrismaDBExceptionHandler.is_lock_timeout_error(cause)
|
||||
or PrismaDBExceptionHandler.postgres_sqlstate(cause) == "40001"
|
||||
)
|
||||
for cause in exception_chain(error)
|
||||
)
|
||||
|
||||
|
||||
def _isolatable_persistence_failure(error: Exception) -> bool:
|
||||
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler, exception_chain
|
||||
|
||||
return isinstance(error, AccountingConflictError) or any(
|
||||
isinstance(cause, Exception) and (PrismaDBExceptionHandler.postgres_sqlstate(cause) or "").startswith("22")
|
||||
for cause in exception_chain(error)
|
||||
)
|
||||
|
||||
|
||||
class DurableUsageMeter:
|
||||
def __init__(
|
||||
self,
|
||||
write_batch: Callable[[tuple[UsageWrite, ...]], Awaitable[None]],
|
||||
settings: UsageCaptureSettings,
|
||||
*,
|
||||
wait_for_batch: Callable[[asyncio.Event, float], Awaitable[None]] = _wait_for_batch,
|
||||
retry_delay: Callable[[float], Awaitable[None]] = asyncio.sleep,
|
||||
) -> None:
|
||||
self._write_batch: Final = write_batch
|
||||
self.settings: Final = settings
|
||||
self._wait_for_batch: Final = wait_for_batch
|
||||
self._retry_delay: Final = retry_delay
|
||||
self._loop: Final = asyncio.get_running_loop()
|
||||
self._admission: Final = Lock()
|
||||
self._queue: Final[deque[_Pending]] = deque() # mutable-ok: one bounded producer/consumer queue
|
||||
self._wakeup: Final = asyncio.Event()
|
||||
self._task: asyncio.Task[None] | None = None
|
||||
self._closed = False
|
||||
self.pending_bytes = 0
|
||||
self.pending_records = 0
|
||||
|
||||
def submit(
|
||||
self, writes: tuple[UsageWrite, ...], *, prerequisites: tuple[asyncio.Future[None], ...] = ()
|
||||
) -> asyncio.Future[None]:
|
||||
size: Final = sum(len(write.envelope.encode("utf-8")) for write in writes)
|
||||
with self._admission:
|
||||
if (
|
||||
self._closed
|
||||
or not writes
|
||||
or len(writes) > self.settings.batch_size
|
||||
or self.pending_bytes + size > self.settings.max_pending_bytes
|
||||
or self.pending_records + len(writes) > self.settings.max_pending_records
|
||||
):
|
||||
raise UsageCaptureError("Usage capture admission rejected: closed, empty or capacity exhausted")
|
||||
receipt: Final = Future[None]()
|
||||
receipt.set_running_or_notify_cancel()
|
||||
receipt.add_done_callback(_report_failure)
|
||||
self._queue.append(_Pending(writes, size, receipt, asyncio.Event(), prerequisites))
|
||||
self.pending_bytes += size
|
||||
self.pending_records += len(writes)
|
||||
self._loop.call_soon_threadsafe(self._ensure_worker)
|
||||
wrapped: Final = asyncio.wrap_future(receipt, loop=self._loop)
|
||||
wrapped.add_done_callback(lambda done: None if done.cancelled() else done.exception())
|
||||
return wrapped
|
||||
|
||||
def _ensure_worker(self) -> None:
|
||||
if self.pending_records >= self.settings.batch_size or self._queue and self._queue[0].prerequisites:
|
||||
self._wakeup.set()
|
||||
if self._task is None or self._task.done():
|
||||
self._task = asyncio.create_task(self._run(), name="accounting usage capture")
|
||||
|
||||
async def record(self, event: InboundUsageEvent | OutboundUsageEvent) -> None:
|
||||
envelope: Final = '{"event":' + event.model_dump_json() + "}"
|
||||
await asyncio.shield(self.submit((UsageWrite(event, envelope, usage_delivery_targets(event.dimensions)),)))
|
||||
|
||||
async def flush(self) -> None:
|
||||
self._ensure_worker()
|
||||
self._wakeup.set()
|
||||
if self._task is not None:
|
||||
await asyncio.shield(self._task)
|
||||
|
||||
async def close(self) -> None:
|
||||
with self._admission:
|
||||
self._closed = True
|
||||
await self.flush()
|
||||
|
||||
async def stop(self, timeout: float) -> None:
|
||||
with self._admission:
|
||||
self._closed = True
|
||||
try:
|
||||
await asyncio.wait_for(self.close(), timeout)
|
||||
except asyncio.TimeoutError:
|
||||
if self._task is not None:
|
||||
self._task.cancel()
|
||||
await asyncio.gather(self._task, return_exceptions=True)
|
||||
while self._queue:
|
||||
self._finish(
|
||||
self._take_batch(), UsageCaptureError("Usage capture stopped before commit acknowledgement")
|
||||
)
|
||||
verbose_proxy_logger.error("Usage capture shutdown deadline exhausted; commit acknowledgement incomplete")
|
||||
|
||||
def _take_batch(self) -> tuple[_Pending, ...]:
|
||||
with self._admission:
|
||||
if self._queue[0].prerequisites:
|
||||
return (self._queue.popleft(),)
|
||||
independent: Final = takewhile(lambda item: not item.prerequisites, self._queue)
|
||||
sizes: Final = accumulate(len(item.writes) for item in independent)
|
||||
count: Final = sum(1 for _ in takewhile(lambda size: size <= self.settings.batch_size, sizes))
|
||||
return tuple(self._queue.popleft() for _ in range(count))
|
||||
|
||||
async def _persist(self, writes: tuple[UsageWrite, ...]) -> None:
|
||||
for attempt in count():
|
||||
try:
|
||||
await self._write_batch(writes)
|
||||
except Exception as error: # noqa: BLE001 # retry only classified database failures, preserving frozen writes
|
||||
if not _retryable_persistence_failure(error):
|
||||
raise
|
||||
verbose_proxy_logger.warning("Usage persistence unavailable; retaining admitted batch for retry")
|
||||
await self._retry_delay(min(10.0, 0.2 * 2 ** min(attempt, 6)))
|
||||
else:
|
||||
return
|
||||
|
||||
async def _commit(self, batch: tuple[_Pending, ...]) -> None:
|
||||
writes: Final = tuple(chain.from_iterable(item.writes for item in batch))
|
||||
try:
|
||||
if batch[0].prerequisites:
|
||||
outcomes: Final = await asyncio.gather(
|
||||
*(asyncio.shield(receipt) for receipt in batch[0].prerequisites), return_exceptions=True
|
||||
)
|
||||
for outcome in outcomes:
|
||||
if isinstance(outcome, BaseException):
|
||||
raise UsageCaptureError("Usage dependency did not acknowledge commit") from outcome
|
||||
await self._persist(writes)
|
||||
except asyncio.CancelledError:
|
||||
self._finish(batch, UsageCaptureError("Usage commit acknowledgement interrupted"))
|
||||
raise
|
||||
except Exception as error: # noqa: BLE001 # isolate data faults, surface systemic failures to all producers
|
||||
if len(batch) > 1 and _isolatable_persistence_failure(error):
|
||||
try:
|
||||
for item in batch:
|
||||
await self._commit((item,))
|
||||
except asyncio.CancelledError:
|
||||
self._finish(batch, UsageCaptureError("Usage commit acknowledgement interrupted"))
|
||||
raise
|
||||
return
|
||||
self._finish(
|
||||
batch,
|
||||
UsageCaptureError("Usage identity conflicts with accepted content")
|
||||
if isinstance(error, AccountingConflictError)
|
||||
else error,
|
||||
)
|
||||
else:
|
||||
self._finish(batch, None)
|
||||
|
||||
def _finish(self, batch: tuple[_Pending, ...], error: Exception | None) -> None:
|
||||
for item in batch:
|
||||
if item.finished.is_set():
|
||||
continue
|
||||
item.finished.set()
|
||||
with self._admission:
|
||||
self.pending_bytes -= item.size
|
||||
self.pending_records -= len(item.writes)
|
||||
if item.committed.done():
|
||||
continue
|
||||
if error is None:
|
||||
item.committed.set_result(None)
|
||||
else:
|
||||
item.committed.set_exception(error)
|
||||
|
||||
async def _run(self) -> None:
|
||||
try:
|
||||
while self._queue:
|
||||
await self._wait_for_batch(self._wakeup, self.settings.flush_interval_seconds)
|
||||
self._wakeup.clear()
|
||||
await self._commit(self._take_batch())
|
||||
if (
|
||||
self._closed
|
||||
or self.pending_records >= self.settings.batch_size
|
||||
or self._queue
|
||||
and self._queue[0].prerequisites
|
||||
):
|
||||
self._wakeup.set()
|
||||
except (Exception, asyncio.CancelledError): # noqa: BLE001 # reject queued receipts if the writer exits abnormally
|
||||
while self._queue:
|
||||
self._finish(
|
||||
self._take_batch(), UsageCaptureError("Usage capture stopped before commit acknowledgement")
|
||||
)
|
||||
raise
|
||||
|
|
@ -32,7 +32,7 @@ from litellm.accounting import (
|
|||
PricingSnapshot,
|
||||
)
|
||||
from litellm.proxy.db import accounting_ledger
|
||||
from litellm.proxy.db.accounting_ledger import AccountingConflictError, AccountingEvent, AccountingLedger
|
||||
from litellm.proxy.db.accounting_ledger import AccountingConflictError, AccountingEvent, AccountingLedger, UsageWrite
|
||||
from litellm.types.utils import CallTypes
|
||||
from tests.integration._support.database import scratch_database
|
||||
from tests.integration._support.prisma_clients import connected_proxy_prisma
|
||||
|
|
@ -532,3 +532,74 @@ def test_family_cleanup_lookups_use_the_kind_source_index(ledger_database: str,
|
|||
name == "LiteLLM_AccountingEvent_kind_source_event_id_idx" and "source_event_id" in str(condition)
|
||||
for name, condition in conditions
|
||||
), conditions
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("failure", ["before_commit", "after_commit", "conflict"])
|
||||
async def test_usage_batch_is_atomic_and_replay_preserves_receipts(
|
||||
writer: _Writer, ledger_database: str, failure: str
|
||||
) -> None:
|
||||
client, store = writer
|
||||
events: Final = (_usage("caller", inbound=True), _usage("attempt"))
|
||||
writes: Final = tuple(UsageWrite(event, _envelope(event), _TARGETS) for event in events)
|
||||
|
||||
@asynccontextmanager
|
||||
async def interrupted() -> AsyncGenerator[Prisma]:
|
||||
async with client.writer_db.tx() as tx:
|
||||
yield tx
|
||||
if failure == "before_commit":
|
||||
raise RuntimeError("injected pre-commit failure")
|
||||
if failure == "after_commit":
|
||||
raise RuntimeError("injected lost acknowledgement")
|
||||
|
||||
if failure == "conflict":
|
||||
await store.insert_usage_batch((writes[1],))
|
||||
_complete_and_age(ledger_database)
|
||||
rejected: Final = (writes[0], UsageWrite(events[1], _envelope(events[1], context={"changed": True}), _TARGETS))
|
||||
with pytest.raises(AccountingConflictError):
|
||||
await store.insert_usage_batch(rejected)
|
||||
assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent"') == [("attempt",)]
|
||||
else:
|
||||
with pytest.raises(RuntimeError, match="injected"):
|
||||
await AccountingLedger(interrupted).insert_usage_batch(writes)
|
||||
count: Final = 2 if failure == "after_commit" else 0
|
||||
assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(count,)]
|
||||
assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(count * len(_TARGETS),)]
|
||||
await store.insert_usage_batch(writes)
|
||||
recovered: Final = await store.load_usage_batch(tuple(event.event_id for event in events))
|
||||
assert len(recovered) == 1
|
||||
assert recovered[0].event == events[1]
|
||||
assert json.loads(recovered[0].envelope) == json.loads(writes[1].envelope)
|
||||
assert set(recovered[0].targets) == set(_TARGETS)
|
||||
_complete_and_age(ledger_database)
|
||||
before: Final = _rows(ledger_database, 'SELECT * FROM "LiteLLM_AccountingDelivery" ORDER BY event_id, target')
|
||||
await store.insert_usage_batch((*writes, *writes))
|
||||
assert _rows(ledger_database, 'SELECT * FROM "LiteLLM_AccountingDelivery" ORDER BY event_id, target') == before
|
||||
assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingEvent"') == [(2,)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalid_jsonb_submission_cannot_reject_another_producers_commit(
|
||||
writer: _Writer, ledger_database: str
|
||||
) -> None:
|
||||
from prisma.errors import DataError
|
||||
|
||||
from litellm.proxy.spend_tracking.usage_meter import DurableUsageMeter, UsageCaptureSettings
|
||||
|
||||
_, store = writer
|
||||
invalid: Final = _usage("invalid")
|
||||
valid: Final = _usage("valid")
|
||||
meter: Final = DurableUsageMeter(store.insert_usage_batch, UsageCaptureSettings(batch_size=2))
|
||||
rejected: Final = meter.submit(
|
||||
(UsageWrite(invalid, _envelope(invalid, context={"parent": "bad\x00text"}), _TARGETS),)
|
||||
)
|
||||
accepted: Final = meter.submit((UsageWrite(valid, _envelope(valid), _TARGETS),))
|
||||
try:
|
||||
with pytest.raises(DataError):
|
||||
await rejected
|
||||
await accepted
|
||||
finally:
|
||||
await meter.close()
|
||||
assert _rows(ledger_database, 'SELECT event_id FROM "LiteLLM_AccountingEvent"') == [("valid",)]
|
||||
assert _rows(ledger_database, 'SELECT count(*) FROM "LiteLLM_AccountingDelivery"') == [(len(_TARGETS),)]
|
||||
assert (meter.pending_bytes, meter.pending_records) == (0, 0)
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ from typing import Final, Literal, NoReturn
|
|||
import pytest
|
||||
|
||||
from litellm.accounting import InboundUsageEvent, Metric, MetricName, Outcome
|
||||
from litellm.proxy.db.accounting_ledger import AccountingLedger
|
||||
from litellm.proxy.db.accounting_ledger import AccountingConflictError, AccountingLedger, UsageWrite
|
||||
from litellm.types.utils import CallTypes
|
||||
|
||||
|
||||
|
|
@ -101,3 +101,18 @@ async def test_unknown_delivery_target_is_rejected_before_database_access(event:
|
|||
("unknown",), # pyright: ignore[reportArgumentType] # Exercise an unknown runtime target
|
||||
envelope=envelope,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_usage_batches_do_not_open_a_transaction() -> None:
|
||||
ledger: Final = AccountingLedger(_no_transaction)
|
||||
await ledger.insert_usage_batch(())
|
||||
assert await ledger.load_usage_batch(()) == ()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_conflicting_usage_within_one_batch_is_rejected_before_database_access(event: InboundUsageEvent) -> None:
|
||||
original: Final = UsageWrite(event, json.dumps({"event": event.model_dump(mode="json")}))
|
||||
changed: Final = UsageWrite(event, json.dumps({"event": event.model_dump(mode="json"), "pricing": {"rate": 2}}))
|
||||
with pytest.raises(AccountingConflictError, match="conflicting copies"):
|
||||
await AccountingLedger(_no_transaction).insert_usage_batch((original, changed))
|
||||
|
|
|
|||
|
|
@ -951,3 +951,33 @@ def test_is_lock_timeout_error_recognises_a_statement_cancelled_under_lock_timeo
|
|||
)
|
||||
def test_is_lock_timeout_error_excludes_other_failures(error: Exception) -> None:
|
||||
assert PrismaDBExceptionHandler.is_lock_timeout_error(error) is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("body", "retryable"),
|
||||
[
|
||||
('{"error_code":"P1017","message":"Server has closed the connection."}', True),
|
||||
('{"error_code":"P1000","message":"Authentication failed."}', False),
|
||||
('{"error_code":"P2021","message":"Table does not exist."}', False),
|
||||
('{"error_code":"P9999","message":"Unknown engine failure."}', False),
|
||||
("invalid engine response", False),
|
||||
],
|
||||
)
|
||||
def test_engine_connection_closure_is_recoverable_without_allowing_unverified_auth(
|
||||
body: str,
|
||||
retryable: bool,
|
||||
) -> None:
|
||||
from prisma.http import Response
|
||||
|
||||
error: Final = EngineRequestError(Response(httpx.Response(500)), body)
|
||||
assert PrismaDBExceptionHandler.is_database_transport_error(error) is retryable
|
||||
assert PrismaDBExceptionHandler.is_permanent_database_fault(error) is not retryable
|
||||
assert PrismaDBExceptionHandler.is_database_connection_error(error) is False
|
||||
|
||||
|
||||
def test_connection_closure_code_survives_prisma_data_error_wrapping() -> None:
|
||||
error: Final = DataError(
|
||||
data={"user_facing_error": {"error_code": "P1017", "message": "Server has closed the connection."}}
|
||||
)
|
||||
assert PrismaDBExceptionHandler.is_database_transport_error(error) is True
|
||||
assert PrismaDBExceptionHandler.is_permanent_database_fault(error) is False
|
||||
|
|
|
|||
389
tests/unit/proxy/spend_tracking/test_usage_meter.py
Normal file
389
tests/unit/proxy/spend_tracking/test_usage_meter.py
Normal file
|
|
@ -0,0 +1,389 @@
|
|||
import asyncio
|
||||
import json
|
||||
from datetime import datetime, timezone
|
||||
from decimal import Decimal
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.accounting import InboundUsageEvent, Metric, MetricName, Outcome
|
||||
from litellm.proxy.db.accounting_ledger import AccountingConflictError, UsageWrite
|
||||
from litellm.proxy.spend_tracking.usage_meter import DurableUsageMeter, UsageCaptureError, UsageCaptureSettings
|
||||
from litellm.types.utils import CallTypes
|
||||
|
||||
|
||||
def _write(identity: str) -> UsageWrite:
|
||||
event: Final = InboundUsageEvent(
|
||||
event_id=identity,
|
||||
request_id=identity,
|
||||
call_type=CallTypes.acompletion,
|
||||
started_at=datetime(2026, 1, 1, tzinfo=timezone.utc),
|
||||
ended_at=datetime(2026, 1, 1, tzinfo=timezone.utc),
|
||||
outcome=Outcome.SUCCEEDED,
|
||||
dimensions=(),
|
||||
metrics=(Metric(name=MetricName.REQUESTS, value=Decimal(1)),),
|
||||
)
|
||||
return UsageWrite(event, json.dumps({"event": event.model_dump(mode="json")}))
|
||||
|
||||
|
||||
class CommitBarrier:
|
||||
def __init__(self, *, conflict: str | None = None, failure: Exception | None = None) -> None:
|
||||
self.entered: Final = asyncio.Event()
|
||||
self.release: Final = asyncio.Event()
|
||||
self.committed: tuple[str, ...] = ()
|
||||
self.transactions = 0
|
||||
self.conflict: Final = conflict
|
||||
self.failure: Final = failure
|
||||
|
||||
async def write(self, writes: tuple[UsageWrite, ...]) -> None:
|
||||
self.transactions += 1
|
||||
self.entered.set()
|
||||
await self.release.wait()
|
||||
if self.failure is not None:
|
||||
raise self.failure
|
||||
identities: Final = tuple(write.event.event_id for write in writes)
|
||||
if self.conflict in identities:
|
||||
raise AccountingConflictError("Conflicting test event")
|
||||
self.committed += identities
|
||||
|
||||
|
||||
async def _manual_flush(wakeup: asyncio.Event, timeout: float) -> None:
|
||||
await wakeup.wait()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_receipts_wait_for_one_batched_commit_and_keep_inflight_capacity_reserved() -> None:
|
||||
database: Final = CommitBarrier()
|
||||
meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=2, max_pending_records=2))
|
||||
first: Final = meter.submit((_write("one"),))
|
||||
second: Final = meter.submit((_write("two"),))
|
||||
await database.entered.wait()
|
||||
assert not first.done() and not second.done()
|
||||
assert meter.pending_records == 2
|
||||
with pytest.raises(UsageCaptureError, match="capacity"):
|
||||
meter.submit((_write("three"),))
|
||||
database.release.set()
|
||||
await asyncio.gather(first, second)
|
||||
await meter.close()
|
||||
assert database.committed == ("one", "two")
|
||||
assert database.transactions == 1
|
||||
assert (meter.pending_bytes, meter.pending_records) == (0, 0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_byte_limit_rejects_new_work_without_evicting_admitted_work() -> None:
|
||||
database: Final = CommitBarrier()
|
||||
write: Final = _write("large" * 300)
|
||||
size: Final = len(write.envelope.encode())
|
||||
meter: Final = DurableUsageMeter(
|
||||
database.write, UsageCaptureSettings(max_pending_bytes=size), wait_for_batch=_manual_flush
|
||||
)
|
||||
receipt: Final = meter.submit((write,))
|
||||
with pytest.raises(UsageCaptureError, match="capacity"):
|
||||
meter.submit((_write("extra"),))
|
||||
assert meter.pending_bytes == size
|
||||
database.release.set()
|
||||
await meter.close()
|
||||
await receipt
|
||||
assert database.committed == (write.event.event_id,)
|
||||
with pytest.raises(UsageCaptureError, match="closed"):
|
||||
meter.submit((_write("after-close"),))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_conflicting_group_rolls_back_and_valid_neighbor_still_commits() -> None:
|
||||
database: Final = CommitBarrier(conflict="conflict")
|
||||
meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=3))
|
||||
rejected: Final = meter.submit((_write("caller"), _write("conflict")))
|
||||
accepted: Final = meter.submit((_write("valid"),))
|
||||
await database.entered.wait()
|
||||
database.release.set()
|
||||
with pytest.raises(UsageCaptureError, match="identity"):
|
||||
await rejected
|
||||
await accepted
|
||||
await meter.close()
|
||||
assert database.committed == ("valid",)
|
||||
assert meter.pending_records == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_permanent_persistence_failure_reaches_each_receipt_and_releases_capacity() -> None:
|
||||
failure: Final = ValueError("invalid envelope")
|
||||
database: Final = CommitBarrier(failure=failure)
|
||||
meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=2))
|
||||
first: Final = meter.submit((_write("one"),))
|
||||
second: Final = meter.submit((_write("two"),))
|
||||
await database.entered.wait()
|
||||
database.release.set()
|
||||
results: Final = await asyncio.gather(first, second, return_exceptions=True)
|
||||
await meter.close()
|
||||
assert results == [failure, failure]
|
||||
assert meter.pending_bytes == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancelling_record_waiter_does_not_cancel_persistence() -> None:
|
||||
database: Final = CommitBarrier()
|
||||
meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=2))
|
||||
waiting: Final = asyncio.create_task(meter.record(_write("caller").event))
|
||||
await database.entered.wait()
|
||||
waiting.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await waiting
|
||||
database.release.set()
|
||||
await meter.close()
|
||||
assert database.committed == ("caller",)
|
||||
assert meter.pending_records == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_shutdown_deadline_rejects_unacknowledged_work_and_releases_capacity() -> None:
|
||||
database: Final = CommitBarrier()
|
||||
meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=2))
|
||||
first: Final = meter.submit((_write("in-flight"), _write("in-flight-two")))
|
||||
await database.entered.wait()
|
||||
queued: Final = meter.submit((_write("queued"),))
|
||||
await meter.stop(0)
|
||||
results: Final = await asyncio.gather(first, queued, return_exceptions=True)
|
||||
assert all(isinstance(result, UsageCaptureError) for result in results)
|
||||
assert meter.pending_records == 0
|
||||
assert meter.pending_bytes == 0
|
||||
assert database.committed == ()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_shutdown_during_database_backoff_releases_reserved_capacity() -> None:
|
||||
retrying: Final = asyncio.Event()
|
||||
blocked: Final = asyncio.Event()
|
||||
|
||||
async def unavailable(writes: tuple[UsageWrite, ...]) -> None:
|
||||
raise ConnectionError("database unavailable")
|
||||
|
||||
async def backoff(delay: float) -> None:
|
||||
retrying.set()
|
||||
await blocked.wait()
|
||||
|
||||
meter: Final = DurableUsageMeter(unavailable, UsageCaptureSettings(batch_size=2), retry_delay=backoff)
|
||||
receipt: Final = meter.submit((_write("one"), _write("two")))
|
||||
await retrying.wait()
|
||||
await meter.stop(0)
|
||||
with pytest.raises(UsageCaptureError, match="interrupted"):
|
||||
await receipt
|
||||
assert (meter.pending_records, meter.pending_bytes) == (0, 0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("wrapped", [False, True])
|
||||
async def test_closed_pooled_connection_retains_batch_until_automatic_recovery(wrapped: bool) -> None:
|
||||
import httpx
|
||||
from prisma.engine.errors import EngineRequestError
|
||||
from prisma.http import Response
|
||||
|
||||
database: Final = CommitBarrier()
|
||||
attempts: Final = asyncio.Queue[tuple[UsageWrite, ...]]()
|
||||
retrying: Final = asyncio.Event()
|
||||
recovered: Final = asyncio.Event()
|
||||
closure: Final = EngineRequestError(
|
||||
Response(httpx.Response(500)), '{"error_code":"P1017","message":"Server has closed the connection."}'
|
||||
)
|
||||
failure: Final = RuntimeError("ledger transaction failed") if wrapped else closure
|
||||
if wrapped:
|
||||
failure.__cause__ = closure
|
||||
|
||||
async def persist(writes: tuple[UsageWrite, ...]) -> None:
|
||||
attempts.put_nowait(writes)
|
||||
if attempts.qsize() == 1:
|
||||
raise failure
|
||||
await database.write(writes)
|
||||
|
||||
async def backoff(delay: float) -> None:
|
||||
retrying.set()
|
||||
await recovered.wait()
|
||||
|
||||
meter: Final = DurableUsageMeter(
|
||||
persist, UsageCaptureSettings(batch_size=2, max_pending_records=2), retry_delay=backoff
|
||||
)
|
||||
writes: Final = (_write("caller"), _write("attempt"))
|
||||
receipt: Final = meter.submit(writes)
|
||||
retry_signal: Final = asyncio.create_task(retrying.wait())
|
||||
reserved: Final = meter.pending_bytes
|
||||
try:
|
||||
await asyncio.wait((receipt, retry_signal), return_when=asyncio.FIRST_COMPLETED)
|
||||
assert not receipt.done(), "A dropped pooled connection must retain admitted usage for retry"
|
||||
assert (meter.pending_records, meter.pending_bytes) == (2, reserved)
|
||||
with pytest.raises(UsageCaptureError, match="capacity"):
|
||||
meter.submit((_write("overflow"),))
|
||||
database.release.set()
|
||||
recovered.set()
|
||||
await receipt
|
||||
assert attempts.get_nowait() == attempts.get_nowait() == writes
|
||||
assert database.committed == ("caller", "attempt")
|
||||
assert (meter.pending_records, meter.pending_bytes) == (0, 0)
|
||||
finally:
|
||||
retry_signal.cancel()
|
||||
await asyncio.gather(retry_signal, return_exceptions=True)
|
||||
await meter.stop(0)
|
||||
await asyncio.gather(receipt, return_exceptions=True)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("first_commits", [False, True])
|
||||
async def test_shutdown_during_conflict_isolation_completes_every_receipt(first_commits: bool) -> None:
|
||||
isolating: Final = asyncio.Event()
|
||||
blocked: Final = asyncio.Event()
|
||||
|
||||
async def persist(writes: tuple[UsageWrite, ...]) -> None:
|
||||
if len(writes) > 1:
|
||||
raise AccountingConflictError("one producer conflicts")
|
||||
if first_commits and writes[0].event.event_id == "one":
|
||||
return
|
||||
isolating.set()
|
||||
await blocked.wait()
|
||||
|
||||
meter: Final = DurableUsageMeter(persist, UsageCaptureSettings(batch_size=2))
|
||||
first: Final = meter.submit((_write("one"),))
|
||||
second: Final = meter.submit((_write("two"),))
|
||||
await isolating.wait()
|
||||
await meter.stop(0)
|
||||
results: Final = await asyncio.gather(first, second, return_exceptions=True)
|
||||
assert results[0] is None if first_commits else isinstance(results[0], UsageCaptureError)
|
||||
assert isinstance(results[1], UsageCaptureError)
|
||||
assert (meter.pending_records, meter.pending_bytes) == (0, 0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_elapsed_flush_deadline_commits_a_partial_batch() -> None:
|
||||
from litellm.proxy.spend_tracking.usage_meter import _wait_for_batch
|
||||
|
||||
database: Final = CommitBarrier()
|
||||
database.release.set()
|
||||
|
||||
async def expired_flush(wakeup: asyncio.Event, timeout: float) -> None:
|
||||
assert timeout == 5.0
|
||||
await _wait_for_batch(wakeup, 0)
|
||||
|
||||
meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=10), wait_for_batch=expired_flush)
|
||||
receipt: Final = meter.submit((_write("partial-batch"),))
|
||||
try:
|
||||
await receipt
|
||||
assert database.committed == ("partial-batch",)
|
||||
assert database.transactions == 1
|
||||
assert (meter.pending_records, meter.pending_bytes) == (0, 0)
|
||||
finally:
|
||||
await meter.stop(0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_burst_drains_full_batches_without_waiting_for_partial_batch_deadline() -> None:
|
||||
database: Final = CommitBarrier()
|
||||
database.release.set()
|
||||
waits: Final = asyncio.Queue[bool]()
|
||||
|
||||
async def observe_wait(wakeup: asyncio.Event, timeout: float) -> None:
|
||||
waits.put_nowait(wakeup.is_set())
|
||||
await wakeup.wait()
|
||||
|
||||
meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=2), wait_for_batch=observe_wait)
|
||||
receipts: Final = tuple(meter.submit((_write(str(index)),)) for index in range(5))
|
||||
try:
|
||||
assert await waits.get() is True
|
||||
await asyncio.gather(*receipts[:2])
|
||||
assert await waits.get() is True, "A full queued batch must not wait for the flush deadline"
|
||||
await asyncio.gather(*receipts[2:4])
|
||||
assert await waits.get() is False
|
||||
assert database.committed == ("0", "1", "2", "3")
|
||||
assert meter.pending_records == 1
|
||||
assert not receipts[-1].done()
|
||||
await meter.close()
|
||||
await receipts[-1]
|
||||
assert database.committed == ("0", "1", "2", "3", "4")
|
||||
assert database.transactions == 3
|
||||
assert (meter.pending_records, meter.pending_bytes) == (0, 0)
|
||||
finally:
|
||||
await meter.stop(0)
|
||||
await asyncio.gather(*receipts, return_exceptions=True)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("sqlstate", ["22P05", "22003", "42P01"])
|
||||
async def test_invalid_submission_isolated_without_retrying_systemic_faults(sqlstate: str) -> None:
|
||||
from prisma.errors import DataError
|
||||
|
||||
database: Final = CommitBarrier()
|
||||
database.release.set()
|
||||
attempts: Final = asyncio.Queue[tuple[str, ...]]()
|
||||
|
||||
async def persist(writes: tuple[UsageWrite, ...]) -> None:
|
||||
identities: Final = tuple(write.event.event_id for write in writes)
|
||||
attempts.put_nowait(identities)
|
||||
if "invalid" in identities or sqlstate == "42P01":
|
||||
raise DataError(
|
||||
data={"user_facing_error": {"message": "database rejected input", "meta": {"code": sqlstate}}}
|
||||
)
|
||||
await database.write(writes)
|
||||
|
||||
meter: Final = DurableUsageMeter(persist, UsageCaptureSettings(batch_size=3))
|
||||
rejected: Final = meter.submit((_write("caller"), _write("invalid")))
|
||||
neighbor: Final = meter.submit((_write("neighbor"),))
|
||||
try:
|
||||
with pytest.raises(DataError):
|
||||
await rejected
|
||||
if sqlstate == "42P01":
|
||||
with pytest.raises(DataError):
|
||||
await neighbor
|
||||
assert attempts.qsize() == 1
|
||||
else:
|
||||
await neighbor
|
||||
assert database.committed == ("neighbor",)
|
||||
assert attempts.qsize() == 3
|
||||
finally:
|
||||
await meter.close()
|
||||
assert (meter.pending_records, meter.pending_bytes) == (0, 0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("dependency_outcome", ["commit", "failure", "cancel"])
|
||||
async def test_commit_dependencies_preserve_capacity_and_isolate_failed_dependents(dependency_outcome: str) -> None:
|
||||
database: Final = CommitBarrier()
|
||||
meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=3, max_pending_records=4))
|
||||
prerequisite: Final[asyncio.Future[None]] = asyncio.get_running_loop().create_future()
|
||||
dependent: Final = meter.submit((_write("parent"), _write("final")), prerequisites=(prerequisite,))
|
||||
neighbor: Final = meter.submit((_write("neighbor"),))
|
||||
assert meter.pending_records == 3
|
||||
assert not dependent.done()
|
||||
with pytest.raises(UsageCaptureError, match="capacity"):
|
||||
meter.submit((_write("excess-one"), _write("excess-two")))
|
||||
if dependency_outcome == "commit":
|
||||
prerequisite.set_result(None)
|
||||
elif dependency_outcome == "failure":
|
||||
prerequisite.set_exception(ValueError("dependency failed"))
|
||||
else:
|
||||
prerequisite.cancel()
|
||||
database.release.set()
|
||||
await meter.close()
|
||||
await neighbor
|
||||
if dependency_outcome == "commit":
|
||||
await dependent
|
||||
assert database.committed == ("parent", "final", "neighbor")
|
||||
else:
|
||||
with pytest.raises(UsageCaptureError, match="dependency"):
|
||||
await dependent
|
||||
assert database.committed == ("neighbor",)
|
||||
assert meter.pending_records == 0
|
||||
assert meter.pending_bytes == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_shutdown_does_not_cancel_dependency_receipt_or_leak_reserved_capacity() -> None:
|
||||
database: Final = CommitBarrier()
|
||||
meter: Final = DurableUsageMeter(database.write, UsageCaptureSettings(batch_size=2))
|
||||
prerequisite: Final[asyncio.Future[None]] = asyncio.get_running_loop().create_future()
|
||||
dependent: Final = meter.submit((_write("parent"), _write("final")), prerequisites=(prerequisite,))
|
||||
await meter.stop(0)
|
||||
with pytest.raises(UsageCaptureError, match="acknowledgement"):
|
||||
await dependent
|
||||
assert not prerequisite.done()
|
||||
assert database.committed == ()
|
||||
assert meter.pending_records == 0
|
||||
assert meter.pending_bytes == 0
|
||||
prerequisite.set_result(None)
|
||||
Loading…
Add table
Reference in a new issue