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 " "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 " "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 " "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", layer="available",
documentation_types=("admin", "user"), documentation_types=("admin", "user"),
+42 -21
View File
@@ -52,6 +52,7 @@ def dispatch_pending_runs(
session: Session, session: Session,
*, *,
registry: object | None, registry: object | None,
tenant_id: str | None = None,
now: datetime | None = None, now: datetime | None = None,
limit: int = 10, limit: int = 10,
worker_id: str | None = None, worker_id: str | None = None,
@@ -65,6 +66,7 @@ def dispatch_pending_runs(
recovered, recovered_outcome_unknown = _recover_expired_leases( recovered, recovered_outcome_unknown = _recover_expired_leases(
session, session,
now=current, now=current,
tenant_id=tenant_id,
) )
session.commit() session.commit()
summary: dict[str, object] = { summary: dict[str, object] = {
@@ -82,6 +84,7 @@ def dispatch_pending_runs(
session, session,
now=current, now=current,
worker_id=resolved_worker_id, worker_id=resolved_worker_id,
tenant_id=tenant_id,
) )
if run_id is None: if run_id is None:
session.rollback() session.rollback()
@@ -105,21 +108,23 @@ def dispatch_pending_runs(
def purge_expired_runs( def purge_expired_runs(
session: Session, session: Session,
*, *,
tenant_id: str | None = None,
now: datetime | None = None, now: datetime | None = None,
limit: int = 500, limit: int = 500,
) -> dict[str, object]: ) -> dict[str, object]:
current = _as_utc(now or utcnow()) 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( runs = list(
session.scalars( session.scalars(
select(DataflowRun) select(DataflowRun)
.where( .where(*clauses)
DataflowRun.status.in_(
("succeeded", "failed", "cancelled")
),
DataflowRun.retention_until.is_not(None),
DataflowRun.retention_until <= current,
DataflowRun.purged_at.is_(None),
)
.order_by(DataflowRun.retention_until, DataflowRun.id) .order_by(DataflowRun.retention_until, DataflowRun.id)
.limit(max(1, min(int(limit), 5_000))) .limit(max(1, min(int(limit), 5_000)))
.with_for_update(skip_locked=True) .with_for_update(skip_locked=True)
@@ -204,15 +209,19 @@ def _recover_expired_leases(
session: Session, session: Session,
*, *,
now: datetime, now: datetime,
tenant_id: str | None = None,
) -> tuple[int, int]: ) -> 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( runs = list(
session.scalars( session.scalars(
select(DataflowRun) select(DataflowRun)
.where( .where(*clauses)
DataflowRun.status == "running",
DataflowRun.lease_expires_at.is_not(None),
DataflowRun.lease_expires_at < now,
)
.with_for_update(skip_locked=True) .with_for_update(skip_locked=True)
) )
) )
@@ -276,17 +285,21 @@ def _claim_next_run(
*, *,
now: datetime, now: datetime,
worker_id: str, worker_id: str,
tenant_id: str | None = None,
) -> str | 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( candidates = list(
session.scalars( session.scalars(
select(DataflowRun) select(DataflowRun)
.where( .where(*clauses)
DataflowRun.status.in_(("queued", "retrying")),
or_(
DataflowRun.available_at.is_(None),
DataflowRun.available_at <= now,
),
)
.order_by(DataflowRun.available_at, DataflowRun.created_at) .order_by(DataflowRun.available_at, DataflowRun.created_at)
.limit(40) .limit(40)
.with_for_update(skip_locked=True) .with_for_update(skip_locked=True)
@@ -723,6 +736,7 @@ class SqlDataflowRunWorker:
self, self,
session: object, session: object,
*, *,
tenant_id: str | None = None,
now: datetime | None = None, now: datetime | None = None,
limit: int = 10, limit: int = 10,
worker_id: str | None = None, worker_id: str | None = None,
@@ -732,6 +746,7 @@ class SqlDataflowRunWorker:
return dispatch_pending_runs( return dispatch_pending_runs(
session, session,
registry=self._registry, registry=self._registry,
tenant_id=tenant_id,
now=now, now=now,
limit=limit, limit=limit,
worker_id=worker_id, worker_id=worker_id,
@@ -741,12 +756,18 @@ class SqlDataflowRunWorker:
self, self,
session: object, session: object,
*, *,
tenant_id: str | None = None,
now: datetime | None = None, now: datetime | None = None,
limit: int = 500, limit: int = 500,
) -> Mapping[str, object]: ) -> Mapping[str, object]:
if not isinstance(session, Session): if not isinstance(session, Session):
raise TypeError("Dataflow retention requires a 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__ = [ __all__ = [
@@ -966,6 +966,7 @@ class SqlDataflowTriggerDispatcher:
self, self,
session: object, session: object,
*, *,
tenant_id: str | None = None,
now: datetime | None = None, now: datetime | None = None,
limit: int = 50, limit: int = 50,
) -> Mapping[str, object]: ) -> Mapping[str, object]:
@@ -974,6 +975,7 @@ class SqlDataflowTriggerDispatcher:
return dispatch_due_triggers( return dispatch_due_triggers(
session, session,
registry=self._registry, registry=self._registry,
tenant_id=tenant_id,
now=now, now=now,
limit=limit, limit=limit,
).model_dump(mode="json") ).model_dump(mode="json")