Partition dataflow workers by tenant

This commit is contained in:
2026-08-04 09:29:35 +02:00
parent 93dfda0e10
commit 42f87ddc88
3 changed files with 47 additions and 22 deletions
+3 -1
View File
@@ -266,7 +266,9 @@ DOCUMENTATION = (
"Publication to a governed Datasource uses forward recovery: an unknown provider outcome is reconciled "
"before retry so output is not duplicated. Staging and production promotion is explicit and does not "
"rewrite a revision. Cancellation is best effort once external work has started; the final evidence "
"states whether work stopped, completed, failed, or requires operator reconciliation."
"states whether work stopped, completed, failed, or requires operator reconciliation. Scheduled, event, "
"and queued execution is partitioned by tenant module entitlement before a run is claimed. Disabling "
"Dataflow stops new admission and leaves accepted runs available for an explicit operator decision."
),
layer="available",
documentation_types=("admin", "user"),
+42 -21
View File
@@ -52,6 +52,7 @@ def dispatch_pending_runs(
session: Session,
*,
registry: object | None,
tenant_id: str | None = None,
now: datetime | None = None,
limit: int = 10,
worker_id: str | None = None,
@@ -65,6 +66,7 @@ def dispatch_pending_runs(
recovered, recovered_outcome_unknown = _recover_expired_leases(
session,
now=current,
tenant_id=tenant_id,
)
session.commit()
summary: dict[str, object] = {
@@ -82,6 +84,7 @@ def dispatch_pending_runs(
session,
now=current,
worker_id=resolved_worker_id,
tenant_id=tenant_id,
)
if run_id is None:
session.rollback()
@@ -105,21 +108,23 @@ def dispatch_pending_runs(
def purge_expired_runs(
session: Session,
*,
tenant_id: str | None = None,
now: datetime | None = None,
limit: int = 500,
) -> dict[str, object]:
current = _as_utc(now or utcnow())
clauses = [
DataflowRun.status.in_(("succeeded", "failed", "cancelled")),
DataflowRun.retention_until.is_not(None),
DataflowRun.retention_until <= current,
DataflowRun.purged_at.is_(None),
]
if tenant_id:
clauses.append(DataflowRun.tenant_id == tenant_id)
runs = list(
session.scalars(
select(DataflowRun)
.where(
DataflowRun.status.in_(
("succeeded", "failed", "cancelled")
),
DataflowRun.retention_until.is_not(None),
DataflowRun.retention_until <= current,
DataflowRun.purged_at.is_(None),
)
.where(*clauses)
.order_by(DataflowRun.retention_until, DataflowRun.id)
.limit(max(1, min(int(limit), 5_000)))
.with_for_update(skip_locked=True)
@@ -204,15 +209,19 @@ def _recover_expired_leases(
session: Session,
*,
now: datetime,
tenant_id: str | None = None,
) -> tuple[int, int]:
clauses = [
DataflowRun.status == "running",
DataflowRun.lease_expires_at.is_not(None),
DataflowRun.lease_expires_at < now,
]
if tenant_id:
clauses.append(DataflowRun.tenant_id == tenant_id)
runs = list(
session.scalars(
select(DataflowRun)
.where(
DataflowRun.status == "running",
DataflowRun.lease_expires_at.is_not(None),
DataflowRun.lease_expires_at < now,
)
.where(*clauses)
.with_for_update(skip_locked=True)
)
)
@@ -276,17 +285,21 @@ def _claim_next_run(
*,
now: datetime,
worker_id: str,
tenant_id: str | None = None,
) -> str | None:
clauses = [
DataflowRun.status.in_(("queued", "retrying")),
or_(
DataflowRun.available_at.is_(None),
DataflowRun.available_at <= now,
),
]
if tenant_id:
clauses.append(DataflowRun.tenant_id == tenant_id)
candidates = list(
session.scalars(
select(DataflowRun)
.where(
DataflowRun.status.in_(("queued", "retrying")),
or_(
DataflowRun.available_at.is_(None),
DataflowRun.available_at <= now,
),
)
.where(*clauses)
.order_by(DataflowRun.available_at, DataflowRun.created_at)
.limit(40)
.with_for_update(skip_locked=True)
@@ -723,6 +736,7 @@ class SqlDataflowRunWorker:
self,
session: object,
*,
tenant_id: str | None = None,
now: datetime | None = None,
limit: int = 10,
worker_id: str | None = None,
@@ -732,6 +746,7 @@ class SqlDataflowRunWorker:
return dispatch_pending_runs(
session,
registry=self._registry,
tenant_id=tenant_id,
now=now,
limit=limit,
worker_id=worker_id,
@@ -741,12 +756,18 @@ class SqlDataflowRunWorker:
self,
session: object,
*,
tenant_id: str | None = None,
now: datetime | None = None,
limit: int = 500,
) -> Mapping[str, object]:
if not isinstance(session, Session):
raise TypeError("Dataflow retention requires a Session.")
return purge_expired_runs(session, now=now, limit=limit)
return purge_expired_runs(
session,
tenant_id=tenant_id,
now=now,
limit=limit,
)
__all__ = [
@@ -966,6 +966,7 @@ class SqlDataflowTriggerDispatcher:
self,
session: object,
*,
tenant_id: str | None = None,
now: datetime | None = None,
limit: int = 50,
) -> Mapping[str, object]:
@@ -974,6 +975,7 @@ class SqlDataflowTriggerDispatcher:
return dispatch_due_triggers(
session,
registry=self._registry,
tenant_id=tenant_id,
now=now,
limit=limit,
).model_dump(mode="json")