From b486ccbe5e05e39632551d58fe3afd8fb80b92f4 Mon Sep 17 00:00:00 2001 From: Werner Dijkerman Date: Fri, 29 May 2026 21:37:34 +0200 Subject: [PATCH] Horizontal scaling for very large deployments --- charts/promforecast/templates/configmap.yaml | 24 ++ charts/promforecast/templates/deployment.yaml | 46 +- .../promforecast/templates/networkpolicy.yaml | 9 +- charts/promforecast/templates/rbac.yaml | 7 +- .../promforecast/templates/statefulset.yaml | 104 +++++ charts/promforecast/values.schema.json | 32 ++ charts/promforecast/values.yaml | 43 ++ docs/README.md | 1 + docs/operations/ha.md | 11 +- docs/scaling/sharded-ha.md | 101 +++++ forecaster/pyproject.toml | 1 + forecaster/src/promforecast/config.py | 111 +++++ forecaster/src/promforecast/exporter.py | 51 ++- forecaster/src/promforecast/main.py | 207 ++++++++- forecaster/src/promforecast/point_factory.py | 10 + .../src/promforecast/runner/_scheduler.py | 164 +++++++- forecaster/src/promforecast/sharding.py | 397 ++++++++++++++++++ forecaster/tests/test_chart_sharding.py | 115 +++++ forecaster/tests/test_config_sharding.py | 80 ++++ forecaster/tests/test_leader_assigned_kind.py | 184 ++++++++ forecaster/tests/test_runner_sharding.py | 173 ++++++++ forecaster/tests/test_shard_controller.py | 186 ++++++++ forecaster/tests/test_sharding.py | 279 ++++++++++++ 23 files changed, 2312 insertions(+), 24 deletions(-) create mode 100644 charts/promforecast/templates/statefulset.yaml create mode 100644 docs/scaling/sharded-ha.md create mode 100644 forecaster/src/promforecast/sharding.py create mode 100644 forecaster/tests/test_chart_sharding.py create mode 100644 forecaster/tests/test_config_sharding.py create mode 100644 forecaster/tests/test_leader_assigned_kind.py create mode 100644 forecaster/tests/test_runner_sharding.py create mode 100644 forecaster/tests/test_shard_controller.py create mode 100644 forecaster/tests/test_sharding.py diff --git a/charts/promforecast/templates/configmap.yaml b/charts/promforecast/templates/configmap.yaml index 0bc22d2..b89361b 100644 --- a/charts/promforecast/templates/configmap.yaml +++ b/charts/promforecast/templates/configmap.yaml @@ -27,5 +27,29 @@ data: {{- $_ := set $ha "redis" $redis }} {{- end }} {{- $_ := set $cfg "highAvailability" $ha }} + {{- /* + Sharding follows the same merge philosophy: the chart-level `sharding` + block is the easy switch, and we propagate `enabled`, `mode`, and the + shard count (from `replicaCount`, the single source of truth so the + StatefulSet replica count and the forecaster's view can't drift). The + inner `config.sharding` block stays available for advanced overrides. + For leader_assigned we also seed the query-cache Redis address (the + forecaster reads the membership/assignment store from there when HA is + off) unless the operator already set one. + */}} + {{- $sh := default dict (get $cfg "sharding") }} + {{- if .Values.sharding.enabled }} + {{- $_ := set $sh "enabled" true }} + {{- $_ := set $sh "mode" .Values.sharding.mode }} + {{- $_ := set $sh "replica_count" (.Values.replicaCount | int) }} + {{- if eq .Values.sharding.mode "leader_assigned" }} + {{- $qc := default dict (get $cfg "query_cache") }} + {{- if and (.Values.sharding.redis.address) (not (get $qc "redis_address")) }} + {{- $_ := set $qc "redis_address" .Values.sharding.redis.address }} + {{- $_ := set $cfg "query_cache" $qc }} + {{- end }} + {{- end }} + {{- end }} + {{- $_ := set $cfg "sharding" $sh }} {{- toYaml $cfg | nindent 4 }} {{- end }} diff --git a/charts/promforecast/templates/deployment.yaml b/charts/promforecast/templates/deployment.yaml index f0628db..e48f421 100644 --- a/charts/promforecast/templates/deployment.yaml +++ b/charts/promforecast/templates/deployment.yaml @@ -1,3 +1,10 @@ +{{- /* + consistent_hash sharding renders a StatefulSet (see statefulset.yaml) for + stable per-replica ordinals; the Deployment covers every other case — + default, HA, and leader_assigned sharding (whose assigner distributes by + identity, so pods need no stable ordinal). +*/}} +{{- if not (and .Values.sharding.enabled (eq .Values.sharding.mode "consistent_hash")) }} apiVersion: apps/v1 kind: Deployment metadata: @@ -5,15 +12,16 @@ metadata: labels: {{- include "promforecast.labels" . | nindent 4 }} spec: {{- /* - HA mode uses RollingUpdate so a fresh leader can take over before the - old replica is fully torn down (the elector's best-effort release - shaves another few seconds off failover). Single-replica installs - keep Recreate to honour the stateless invariant — a brief gap during - pod replacement is fine because the next refit reproduces everything. + HA and leader_assigned sharding both run multiple replicas that must roll + without a full outage, so we force RollingUpdate (maxUnavailable: 0) when + the chart's default Recreate would otherwise drop every replica at once. + Single-replica installs keep Recreate to honour the stateless invariant — + a brief gap during pod replacement is fine because the next refit + reproduces everything. */}} replicas: {{ .Values.replicaCount }} strategy: - {{- if and .Values.highAvailability.enabled (eq (.Values.strategy.type | default "Recreate") "Recreate") }} + {{- if and (or .Values.highAvailability.enabled .Values.sharding.enabled) (eq (.Values.strategy.type | default "Recreate") "Recreate") }} type: RollingUpdate rollingUpdate: maxSurge: 1 @@ -83,6 +91,31 @@ spec: key: {{ .Values.highAvailability.redis.secretKey }} {{- end }} {{- end }} + {{- /* + leader_assigned sharding: every replica heartbeats into the + shared Redis and reads the published assignment. The pod name + is the shard identity; the namespace scopes the assigner Lease. + (consistent_hash sharding renders the StatefulSet instead.) + */}} + {{- if and .Values.sharding.enabled (eq .Values.sharding.mode "leader_assigned") }} + - name: PROMFORECAST_SHARDING_ENABLED + value: "true" + - name: PROMFORECAST_POD_NAME + valueFrom: + fieldRef: + fieldPath: metadata.name + - name: PROMFORECAST_NAMESPACE + valueFrom: + fieldRef: + fieldPath: metadata.namespace + {{- if .Values.sharding.redis.existingSecret }} + - name: PROMFORECAST_REDIS_PASSWORD + valueFrom: + secretKeyRef: + name: {{ .Values.sharding.redis.existingSecret }} + key: {{ .Values.sharding.redis.secretKey }} + {{- end }} + {{- end }} {{- with .Values.extraEnv }}{{- toYaml . | nindent 12 }}{{- end }} volumeMounts: - name: config @@ -107,3 +140,4 @@ spec: {{- with .Values.affinity }} affinity: {{- toYaml . | nindent 8 }} {{- end }} +{{- end }} diff --git a/charts/promforecast/templates/networkpolicy.yaml b/charts/promforecast/templates/networkpolicy.yaml index a4fea99..e6c9361 100644 --- a/charts/promforecast/templates/networkpolicy.yaml +++ b/charts/promforecast/templates/networkpolicy.yaml @@ -37,7 +37,14 @@ spec: configured port (default 6379); operators with a non-default Redis port can override via networkPolicy.redisPort. */}} - {{- if .Values.highAvailability.enabled }} + {{- /* + leader_assigned sharding has the same egress needs as HA: the Lease + API on the kube-apiserver plus Redis (membership registry + + assignment). consistent_hash sharding needs neither — its egress is + just DNS + the TSDB (covered by networkPolicy.egress). + */}} + {{- $needsCoordEgress := or .Values.highAvailability.enabled (and .Values.sharding.enabled (eq .Values.sharding.mode "leader_assigned")) }} + {{- if $needsCoordEgress }} - to: - namespaceSelector: {} ports: diff --git a/charts/promforecast/templates/rbac.yaml b/charts/promforecast/templates/rbac.yaml index 043d8ca..d2ca946 100644 --- a/charts/promforecast/templates/rbac.yaml +++ b/charts/promforecast/templates/rbac.yaml @@ -1,4 +1,9 @@ -{{- if and .Values.rbac.create .Values.highAvailability.enabled }} +{{- /* + The Lease API is needed by HA (leader election) and by leader_assigned + sharding (the assigner Lease). consistent_hash sharding needs no Lease. +*/}} +{{- $needsLease := or .Values.highAvailability.enabled (and .Values.sharding.enabled (eq .Values.sharding.mode "leader_assigned")) }} +{{- if and .Values.rbac.create $needsLease }} apiVersion: rbac.authorization.k8s.io/v1 kind: Role metadata: diff --git a/charts/promforecast/templates/statefulset.yaml b/charts/promforecast/templates/statefulset.yaml new file mode 100644 index 0000000..b38b32a --- /dev/null +++ b/charts/promforecast/templates/statefulset.yaml @@ -0,0 +1,104 @@ +{{- /* + consistent_hash sharding needs each replica to have a STABLE shard index + in [0, replicas). A Deployment's pods get random names, so there is no + ordinal to derive an index from; a StatefulSet gives every pod a stable + ordinal name (-0, -1, …). The forecaster reads the trailing + ordinal of PROMFORECAST_POD_NAME as its shard index. This template renders + ONLY for consistent_hash sharding; every other case (default, HA, and + leader_assigned sharding, where the assigner distributes by identity) + renders the Deployment instead. +*/}} +{{- if and .Values.sharding.enabled (eq .Values.sharding.mode "consistent_hash") }} +apiVersion: apps/v1 +kind: StatefulSet +metadata: + name: {{ include "promforecast.fullname" . }} + labels: {{- include "promforecast.labels" . | nindent 4 }} +spec: + # The shard count. Each replica owns hash(group) % replicas == ordinal. + replicas: {{ .Values.replicaCount }} + serviceName: {{ include "promforecast.fullname" . }} + # OrderedReady would gate pod N on pod N-1 becoming Ready — needless here + # since shards are independent. Parallel brings the whole fleet up at once. + podManagementPolicy: Parallel + updateStrategy: + type: RollingUpdate + selector: + matchLabels: {{- include "promforecast.selectorLabels" . | nindent 6 }} + template: + metadata: + labels: + {{- include "promforecast.selectorLabels" . | nindent 8 }} + {{- with .Values.podLabels }}{{- toYaml . | nindent 8 }}{{- end }} + annotations: + checksum/config: {{ include (print $.Template.BasePath "/configmap.yaml") . | sha256sum }} + {{- with .Values.podAnnotations }}{{- toYaml . | nindent 8 }}{{- end }} + spec: + serviceAccountName: {{ include "promforecast.serviceAccountName" . }} + securityContext: {{- toYaml .Values.podSecurityContext | nindent 8 }} + {{- with .Values.image.pullSecrets }} + imagePullSecrets: {{- toYaml . | nindent 8 }} + {{- end }} + containers: + - name: forecaster + image: {{ include "promforecast.image" . }} + imagePullPolicy: {{ .Values.image.pullPolicy }} + args: + - --config=/etc/promforecast/config.yaml + ports: + - name: http + containerPort: 9091 + protocol: TCP + readinessProbe: + httpGet: + path: /healthz + port: http + initialDelaySeconds: 5 + periodSeconds: 10 + livenessProbe: + httpGet: + path: /healthz + port: http + initialDelaySeconds: 30 + periodSeconds: 30 + resources: {{- toYaml .Values.resources | nindent 12 }} + securityContext: {{- toYaml .Values.securityContext | nindent 12 }} + env: + - name: PROMFORECAST_SHARDING_ENABLED + value: "true" + # The pod name carries the StatefulSet ordinal; the forecaster + # parses the trailing -N as its shard index and uses the full name + # as the shard_id label value. + - name: PROMFORECAST_POD_NAME + valueFrom: + fieldRef: + fieldPath: metadata.name + - name: PROMFORECAST_NAMESPACE + valueFrom: + fieldRef: + fieldPath: metadata.namespace + {{- with .Values.extraEnv }}{{- toYaml . | nindent 12 }}{{- end }} + volumeMounts: + - name: config + mountPath: /etc/promforecast + readOnly: true + - name: tmp + mountPath: /tmp + {{- with .Values.extraVolumeMounts }}{{- toYaml . | nindent 12 }}{{- end }} + volumes: + - name: config + configMap: + name: {{ include "promforecast.configMapName" . }} + - name: tmp + emptyDir: {} + {{- with .Values.extraVolumes }}{{- toYaml . | nindent 8 }}{{- end }} + {{- with .Values.nodeSelector }} + nodeSelector: {{- toYaml . | nindent 8 }} + {{- end }} + {{- with .Values.tolerations }} + tolerations: {{- toYaml . | nindent 8 }} + {{- end }} + {{- with .Values.affinity }} + affinity: {{- toYaml . | nindent 8 }} + {{- end }} +{{- end }} diff --git a/charts/promforecast/values.schema.json b/charts/promforecast/values.schema.json index 336b3ac..0a7a3f2 100644 --- a/charts/promforecast/values.schema.json +++ b/charts/promforecast/values.schema.json @@ -50,6 +50,38 @@ } } }, + "sharding": { + "type": "object", + "additionalProperties": false, + "properties": { + "enabled": { "type": "boolean" }, + "mode": { "type": "string", "enum": ["consistent_hash", "leader_assigned"] }, + "redis": { + "type": "object", + "additionalProperties": false, + "properties": { + "address": { "type": "string" }, + "existingSecret": { "type": "string" }, + "secretKey": { "type": "string" } + } + } + }, + "if": { + "properties": { + "enabled": { "const": true }, + "mode": { "const": "leader_assigned" } + }, + "required": ["enabled", "mode"] + }, + "then": { + "properties": { + "redis": { + "properties": { "address": { "minLength": 1 } }, + "required": ["address"] + } + } + } + }, "existingConfigMap": { "type": "string" }, "config": { "type": "object" }, "service": { diff --git a/charts/promforecast/values.yaml b/charts/promforecast/values.yaml index ff65400..b89f130 100644 --- a/charts/promforecast/values.yaml +++ b/charts/promforecast/values.yaml @@ -28,6 +28,29 @@ highAvailability: existingSecret: "" secretKey: password +# Horizontal sharding: run multiple *active* replicas, each fitting a disjoint +# subset of the configured groups so the aggregate fit budget scales with the +# replica count. Distinct from highAvailability (one active leader plus hot +# followers serving a cached snapshot) and mutually exclusive with it. The +# `replicaCount` field above controls the shard count. See +# docs/scaling/sharded-ha.md for the sizing guide and mode trade-offs. +sharding: + enabled: false + # consistent_hash — each replica owns a group iff hash(group) % replicas + # equals its own shard index. With this mode the chart renders a + # StatefulSet so each pod's stable ordinal (-0, -1, …) is its + # shard index; no coordination backend is required. + # leader_assigned — a coordination.k8s.io/v1 Lease elects one replica to + # compute a balanced group→shard assignment over the live membership and + # publish it to Redis; every replica reads it. Rebalances on replica + # join/leave. Requires `redis.address`. Renders a Deployment (pods need no + # stable ordinal — the assigner distributes by identity). + mode: consistent_hash + redis: + address: "" # required for mode=leader_assigned, e.g. redis-master.redis.svc:6379 + existingSecret: "" + secretKey: password + # Forecaster YAML config. Either inline below, or set existingConfigMap. existingConfigMap: "" config: @@ -175,6 +198,26 @@ config: password: "" snapshot_ttl: 5m + # Sharding runtime view of the top-level `sharding` switch. `helm install` + # populates `enabled`, `mode`, and `replica_count` from the chart fields and + # `replicaCount` automatically, so operators usually only edit the top-level + # values. `shard_index` and `shard_id` are left at their derive-from-env + # defaults (-1 / "") — under Kubernetes the StatefulSet pod ordinal becomes + # the index and the pod name becomes the id. Override the inner block to + # tune the leader_assigned heartbeat/Lease or to pin a shard id off-cluster. + sharding: + enabled: false + mode: consistent_hash + replica_count: 1 + shard_index: -1 + shard_id: "" + heartbeat_interval_seconds: 10 + member_ttl_seconds: 30 + lease_name: promforecast-shard-assigner + lease_duration_seconds: 30 + renew_deadline_seconds: 20 + retry_period_seconds: 4 + # Optional OpenTelemetry tracing. Disabled by default; flip # `telemetry.enabled` and set an `otlp.endpoint` to ship traces of the # query → fit → export pipeline to an OTel Collector. The Docker image diff --git a/docs/README.md b/docs/README.md index 9400c01..c2fc4ac 100644 --- a/docs/README.md +++ b/docs/README.md @@ -59,6 +59,7 @@ scheduling, and serving telemetry. - [Cardinality budgeting](operations/cardinality.md) — per-family multipliers and the `--strict-cardinality` CI gate - [Scheduling](operations/scheduling.md) — per-query refresh intervals, max-concurrent-fits cap - [High availability](operations/ha.md) — Lease-based leader election, Redis snapshot cache +- [Sharded HA (horizontal scaling)](scaling/sharded-ha.md) — N active replicas each fitting a disjoint slice of groups, via `consistent_hash` or `leader_assigned`; the `shard_id` label and `forecast_shard_owns_group` gauge; sizing guide - [Query cache](operations/cache.md) — in-memory LRU + optional Redis dedup for overlapping PromQL - [Telemetry](operations/telemetry.md) — internal metrics catalogue + per-series staleness + self-observability (fit success ratio, regressor stability, per-group CPU/memory) - [Warm-up status](operations/warmup.md) — fresh-install "what is still loading?" endpoint and CLI diff --git a/docs/operations/ha.md b/docs/operations/ha.md index be4eff3..d028ef9 100644 --- a/docs/operations/ha.md +++ b/docs/operations/ha.md @@ -48,6 +48,11 @@ config: Defaults match the Kubernetes control plane. The forecaster rejects any config where `renew_deadline_seconds >= lease_duration_seconds`. -## Why not multi-leader sharding? - -True sharding (different replicas fitting different groups) is planned for a future major version. For now, followers carry zero fit cost — the simplest and safest model for v0.x deployments. +## Sharding — when one leader isn't enough + +HA scales to one well-sized leader plus hot followers; the followers carry zero +fit cost. Once a single leader can no longer fit everything within +`refresh_interval`, switch to **[sharded HA](../scaling/sharded-ha.md)**, where +every replica is active and fits a disjoint subset of the groups. HA and +sharding are mutually exclusive — pick HA for zero-downtime failover when one +replica keeps up, sharding for horizontal fit throughput. diff --git a/docs/scaling/sharded-ha.md b/docs/scaling/sharded-ha.md new file mode 100644 index 0000000..de0f6ff --- /dev/null +++ b/docs/scaling/sharded-ha.md @@ -0,0 +1,101 @@ +# Sharded HA — horizontal scaling for very large deployments + +[High availability](../operations/ha.md) gives you one active leader plus hot followers serving a cached snapshot. The ceiling is whatever a single leader can fit within a refresh interval — roughly 50–100k output series. + +**Sharding** lifts that ceiling. Every replica is active and fits a **disjoint subset of groups**. Ownership is disjoint and complete, so the aggregate fit budget scales with the replica count. The `replicas` field becomes a throughput knob. + +Sharding and HA are **mutually exclusive** — HA gates fitting on a single leader; sharding spreads fitting across every replica. The config loader rejects enabling both. + +## When to use sharding + +| Situation | Recommended approach | +|-----------|----------------------| +| Fits comfortably within interval on one replica | Single replica (default) | +| Need zero-downtime failover but one replica suffices | [HA](../operations/ha.md) | +| One replica can no longer fit everything within `refresh_interval` | **Sharding** (this page) | + +Symptom of outgrowing HA: `forecast_run_duration_seconds` approaching `server.refresh_interval` or rising `forecast_slow_fits` on a vertically scaled leader. + +## The two modes + +### `consistent_hash` (default, no coordination) + +Each replica owns a group if `hash(group) % replicas == shard_index`. No shared state — replicas only need to agree on replica count (pinned by the chart). + +Uses a **StatefulSet** so each pod gets a stable ordinal. Forecaster reads trailing ordinal of `PROMFORECAST_POD_NAME` as its `shard_index`. + +```yaml +sharding: + enabled: true + mode: consistent_hash +replicaCount: 4 +``` + +**Pros**: Simple, no Redis, no Lease. +**Cons**: Rebalancing on replica count change reshuffles most groups. Crashed shard's groups wait for pod reschedule. + +### `leader_assigned` (coordinated) + +A `coordination.k8s.io/v1` Lease elects an assigner. Replicas heartbeat into a shared Redis membership registry. The assigner computes a balanced mapping and publishes it. + +Uses a **Deployment** (no stable ordinal needed). Employs rendezvous hashing so membership changes move only ~1/N groups. + +```yaml +sharding: + enabled: true + mode: leader_assigned + redis: + address: redis-master.redis.svc:6379 +replicaCount: 4 +``` + +**Pros**: Automatic rebalancing on join/leave; crashed shard's groups reassigned quickly. +**Cons**: Requires Redis and Lease. Cold start has brief period where replicas own nothing until first assignment. + +## Choosing a mode + +Start with `consistent_hash` — simplest and fewest moving parts. +Move to `leader_assigned` when you need automatic rebalancing on replica changes and already run Redis. + +## Output labels and assignment gauge + +Every output metric for an owned group gains a `shard_id` label (pod name by default). This label is only present when sharding is enabled. + +Assignment gauge for debugging: + +``` +forecast_shard_owns_group{shard_id="...", group="..."} 1 +``` + +Each replica emits one sample per owned group. + +## Sizing guide + +- Aim for **2–3× as many groups as shards** for even distribution. +- Target roughly even series-per-shard, not groups-per-shard. Split oversized groups if needed. +- Reasonable target: ≤ 40–60k output series per shard. + +Detect hot shards — a shard whose busiest group nears the refresh interval, or whose owned series count sits well above the fleet median: + +```promql +# per-shard busiest group run duration +max by (shard_id) (forecast_run_duration_seconds * on(group) group_left(shard_id) forecast_shard_owns_group) + +# per-shard owned series count +sum by (shard_id) (forecast_series_count * on(group) group_left(shard_id) forecast_shard_owns_group) +``` + +When to rebalance: + +- `consistent_hash`: rebalancing means changing `replicas`, which reshuffles most ownership — a deliberate, planned operation. Bump the shard count when the *busiest* shard (not the average) is at its fit ceiling. If a single group dominates, split it rather than adding shards. +- `leader_assigned`: rebalancing is automatic on replica join/leave; raising `replicas` redistributes within a heartbeat interval, moving only the groups that must move. + +## Reload behaviour + +Reloads that add/remove groups are handled per-shard: each replica re-derives ownership and starts/stops loops. In `leader_assigned` mode the assigner republishes the mapping for the new group set on its next tick. + +## See also + +- [High availability](../operations/ha.md) +- [Production-readiness checklist](../production-checklist.md) +- [Cardinality budgeting](../operations/cardinality.md) \ No newline at end of file diff --git a/forecaster/pyproject.toml b/forecaster/pyproject.toml index 7b50885..a937414 100644 --- a/forecaster/pyproject.toml +++ b/forecaster/pyproject.toml @@ -185,4 +185,5 @@ testpaths = ["tests"] asyncio_mode = "auto" markers = [ "benchmark: micro-benchmarks gated out of the default run; opt in with `-m benchmark`", + "kind: integration tests needing a real Kubernetes cluster (kind); skipped unless PROMFORECAST_KIND_TEST is set", ] diff --git a/forecaster/src/promforecast/config.py b/forecaster/src/promforecast/config.py index 2909483..24bc4a0 100644 --- a/forecaster/src/promforecast/config.py +++ b/forecaster/src/promforecast/config.py @@ -1216,6 +1216,89 @@ def _enabled_requires_redis(self) -> HighAvailabilityConfig: return self +class ShardingConfig(BaseModel): + """Horizontal sharding: N *active* replicas, each fitting a disjoint + subset of the configured groups. + + Off by default. This is a different scaling model from + :class:`HighAvailabilityConfig` — HA is "one active leader plus hot + followers serving the leader's cached snapshot", which tops out around + one well-sized leader's fit budget. Sharding instead spreads the fit + work across every replica, so horizontal scale becomes a knob (the + chart's ``replicas`` field). The two are mutually exclusive: HA gates + fitting on leadership, sharding wants every replica fitting. + + Two ownership modes: + + * ``consistent_hash`` — each replica computes which groups it owns from + a stable hash of the group name modulo the replica count, compared + against its own shard index. No cross-replica coordination is needed + beyond agreeing on the replica count and each replica knowing its own + index (the chart runs a StatefulSet so the pod ordinal is the index). + * ``leader_assigned`` — a ``coordination.k8s.io/v1`` Lease elects one + replica as the assigner; it computes a balanced group→shard mapping + over the live membership and publishes it to the shared Redis, and + every replica reads it to learn what it owns. Rebalances when a + replica joins or leaves. Requires a Redis backend (the same one HA + and the query cache use) for the membership registry and the + published assignment. + + Output metrics for the groups a replica owns gain a ``shard_id`` label, + and a ``forecast_shard_owns_group{shard_id, group}`` gauge exposes the + assignment for debugging. The ``shard_id`` label is tool-private (it + describes which replica fit the group, not the underlying signal) so it + is not part of the cross-tool label contract; it is only emitted when + sharding is enabled, leaving single-replica labelsets unchanged. + """ + + model_config = ConfigDict(extra="forbid") + + enabled: bool = False + mode: Literal["consistent_hash", "leader_assigned"] = "consistent_hash" + # Total number of shards. The chart sets this from ``replicas`` so the + # config and the StatefulSet replica count can't drift. ``ge=1`` keeps a + # zero from silently disabling all ownership. + replica_count: int = Field(default=1, ge=1) + # This replica's 0-based shard index for ``consistent_hash`` ownership. + # ``-1`` (the default) means "derive at boot" — from the StatefulSet pod + # ordinal (the trailing ``-N`` of ``PROMFORECAST_POD_NAME``) or the + # explicit ``PROMFORECAST_SHARD_INDEX`` env var. Set a concrete value + # only for local/test runs where neither env var is present. + shard_index: int = Field(default=-1, ge=-1) + # Label value identifying this shard on output metrics. Empty means + # "derive at boot" — the pod name, else ``shard-``. Operators + # rarely set this; it exists so a non-Kubernetes deployment can pin a + # stable id. + shard_id: str = "" + # ``leader_assigned`` membership: each replica writes a heartbeat key + # with this TTL; the assigner lists the live members from those keys. + # The TTL must comfortably exceed the heartbeat interval so a replica + # mid-GC-pause isn't dropped from the membership. + heartbeat_interval_seconds: int = Field(default=10, ge=1) + member_ttl_seconds: int = Field(default=30, ge=2) + # Lease parameters for ``leader_assigned`` (mirrors the HA defaults). + lease_name: str = "promforecast-shard-assigner" + lease_duration_seconds: int = Field(default=30, ge=2) + renew_deadline_seconds: int = Field(default=20, ge=1) + retry_period_seconds: int = Field(default=4, ge=1) + + @model_validator(mode="after") + def _validate(self) -> ShardingConfig: + if self.shard_index >= 0 and self.shard_index >= self.replica_count: + raise ValueError( + f"sharding.shard_index ({self.shard_index}) must be < " + f"replica_count ({self.replica_count})" + ) + if self.member_ttl_seconds <= self.heartbeat_interval_seconds: + raise ValueError( + "sharding.member_ttl_seconds must be > heartbeat_interval_seconds " + f"(got {self.member_ttl_seconds} <= {self.heartbeat_interval_seconds})" + ) + if self.renew_deadline_seconds >= self.lease_duration_seconds: + raise ValueError("sharding.renew_deadline_seconds must be < lease_duration_seconds") + return self + + class OTLPTelemetryConfig(BaseModel): """OTLP exporter target for the optional OpenTelemetry tracer. @@ -1614,6 +1697,7 @@ class Config(BaseModel): high_availability: HighAvailabilityConfig = Field( default_factory=HighAvailabilityConfig, alias="highAvailability" ) + sharding: ShardingConfig = Field(default_factory=ShardingConfig) query_cache: QueryCacheConfig = QueryCacheConfig() telemetry: TelemetryConfig = Field(default_factory=TelemetryConfig) groups: list[GroupConfig] @@ -1636,6 +1720,33 @@ class Config(BaseModel): # see :class:`AnnotationsConfig`. annotations: AnnotationsConfig = Field(default_factory=AnnotationsConfig) + @model_validator(mode="after") + def _validate_scaling_modes(self) -> Config: + # HA (single-active leader) and sharding (every replica active) are + # opposite scaling models — enabling both is a misconfiguration, not + # a richer setup. Reject it at load time so the contradiction can't + # ship to a cluster. + if self.sharding.enabled and self.high_availability.enabled: + raise ValueError( + "sharding.enabled and highAvailability.enabled are mutually exclusive: " + "HA gates fitting on a single leader, sharding spreads fitting across " + "every replica. Pick one." + ) + # leader_assigned needs a Redis backend for the membership registry + # and the published group→shard assignment. Accept the HA Redis or a + # standalone query-cache Redis — main.py reuses whichever is set. + if ( + self.sharding.enabled + and self.sharding.mode == "leader_assigned" + and not (self.high_availability.redis.address or self.query_cache.redis_address) + ): + raise ValueError( + "sharding.mode=leader_assigned requires a Redis backend; set " + "highAvailability.redis.address or query_cache.redis_address so " + "replicas can share the membership registry and assignment map" + ) + return self + def load(path: str | Path) -> Config: """Load and validate a YAML config file. diff --git a/forecaster/src/promforecast/exporter.py b/forecaster/src/promforecast/exporter.py index c035428..211ffdc 100644 --- a/forecaster/src/promforecast/exporter.py +++ b/forecaster/src/promforecast/exporter.py @@ -33,7 +33,7 @@ ) if TYPE_CHECKING: - from collections.abc import Callable, Iterator + from collections.abc import Callable, Iterable, Iterator __all__ = [ "AccuracyPoint", @@ -715,6 +715,14 @@ def __init__(self) -> None: self._leader_identity: str = "" self._is_leader: bool = False self._leader_transitions: int = 0 + # Sharding state (set only when sharding is enabled). ``_shard_id`` is + # this replica's shard identity; ``_shard_owned_groups`` is the set of + # groups it currently owns. Rendered as + # ``forecast_shard_owns_group{shard_id, group} 1`` — one sample per + # owned group. Prometheus unioning every replica's scrape reconstructs + # the full assignment. + self._shard_id: str = "" + self._shard_owned_groups: frozenset[str] = frozenset() # Query cache counters (hits / misses / errors). Always emitted so # dashboards can chart a zero series before the first query runs. self._query_cache_events: dict[str, int] = {} @@ -940,6 +948,17 @@ def set_leader_state(self, *, identity: str, is_leader: bool) -> None: self._leader_transitions += 1 self._is_leader = is_leader + def set_shard_state(self, *, shard_id: str, owned_groups: Iterable[str]) -> None: + """Record which groups this shard owns for the assignment gauge. + + Surfaces as ``forecast_shard_owns_group{shard_id, group} 1`` (one + sample per owned group). Only meaningful when sharding is enabled; a + single-replica deployment never calls this, so the gauge is absent. + """ + with self._lock: + self._shard_id = shard_id + self._shard_owned_groups = frozenset(owned_groups) + def increment_query_cache_event(self, kind: str) -> None: """Bump a query-cache counter for ``kind`` (hit | miss | error).""" with self._lock: @@ -1126,6 +1145,10 @@ def _leader_state(self) -> tuple[str, bool, int]: with self._lock: return self._leader_identity, self._is_leader, self._leader_transitions + def _shard_state(self) -> tuple[str, frozenset[str]]: + with self._lock: + return self._shard_id, self._shard_owned_groups + def _query_cache_state(self) -> dict[str, int]: with self._lock: return dict(self._query_cache_events) @@ -1504,6 +1527,7 @@ def collect(self) -> Iterator[GaugeMetricFamily | CounterMetricFamily]: # noqa: source_retries = self._exporter._source_retries_state() config_hash = self._exporter._config_hash_state() leader_state = self._exporter._leader_state() + shard_state = self._exporter._shard_state() cache_state = self._exporter._query_cache_state() what_if_runs_state = self._exporter._what_if_runs_state() trigger_fires_state = self._exporter._trigger_fires_state() @@ -1551,6 +1575,7 @@ def collect(self) -> Iterator[GaugeMetricFamily | CounterMetricFamily]: # noqa: yield from _remote_write_retries_families(remote_write_retries) yield from _source_retries_families(source_retries) yield from _leader_families(leader_state) + yield from _shard_owns_families(shard_state) yield from _query_cache_families(cache_state) yield from _what_if_runs_families(what_if_runs_state) yield from _trigger_fires_families(trigger_fires_state) @@ -2390,6 +2415,30 @@ def _leader_families( yield counter +def _shard_owns_families( + state: tuple[str, frozenset[str]], +) -> Iterator[GaugeMetricFamily]: + """Render the per-shard group-ownership gauge. + + ``forecast_shard_owns_group{shard_id, group}`` is 1 for every group this + replica owns. Emitted only when sharding is enabled (i.e. once + :meth:`Exporter.set_shard_state` has run with a non-empty shard id), so + single-replica deployments don't carry an empty series. A scrape that + unions every replica reconstructs the full group→shard assignment. + """ + shard_id, owned_groups = state + if not shard_id: + return + gauge = GaugeMetricFamily( + "forecast_shard_owns_group", + "1 for each group owned (fit) by this shard in a sharded deployment.", + labels=["shard_id", "group"], + ) + for group in sorted(owned_groups): + gauge.add_metric([shard_id, group], 1.0) + yield gauge + + def _query_cache_families( state: dict[str, int], ) -> Iterator[CounterMetricFamily]: diff --git a/forecaster/src/promforecast/main.py b/forecaster/src/promforecast/main.py index 4f8320f..5a098a4 100644 --- a/forecaster/src/promforecast/main.py +++ b/forecaster/src/promforecast/main.py @@ -26,6 +26,7 @@ from . import ha_cache as ha_cache_module from . import inspect as inspect_module from . import preview as preview_module +from . import sharding as sharding_module from . import telemetry as telemetry_module from . import triggers as triggers_module from . import tuner as tuner_module @@ -909,9 +910,15 @@ def run(config_path: Path, log_level: str) -> None: else None ) - elector = _build_elector(cfg, exporter, runner) - if elector is None: - # No HA configured — record this replica as a permanent "leader" + # Sharding (every replica active, each owns a slice of the groups) and HA + # (single active leader + cache-serving followers) are mutually exclusive — + # the config validator rejects both at once. When sharding is on it owns + # the runner lifecycle (and, for leader_assigned, its own assigner Lease), + # so the HA elector is not built. + shard_controller = _build_shard_controller(cfg, exporter, runner, redis_client) + elector = None if shard_controller is not None else _build_elector(cfg, exporter, runner) + if shard_controller is None and elector is None: + # No HA, no sharding — record this replica as a permanent "leader" # so the metric is meaningful in single-replica deployments. exporter.set_leader_state(identity=default_identity(), is_leader=True) @@ -928,6 +935,7 @@ def run(config_path: Path, log_level: str) -> None: source=source, snapshot_cache=snapshot_cache, elector=elector, + shard_controller=shard_controller, tracer=tracer, ) host, port = _parse_listen(cfg.server.listen) @@ -951,23 +959,33 @@ async def _lifespan_factory( runner: Runner, watcher: ConfigMapWatcher | None, elector: LeaderElector | None, + shard_controller: _ShardController | None = None, tracer: telemetry_module.Tracer | None = None, ) -> AsyncIterator[None]: - """Start lifecycle, gated on leader election when HA mode is active. + """Start lifecycle, gated on leader election or sharding when active. - Without HA: behaves exactly as before — the runner starts immediately - and stops on shutdown. + Without HA or sharding: behaves exactly as before — the runner starts + immediately and stops on shutdown. With HA: the elector owns the runner. It runs in a background task and calls ``runner.start()`` / ``runner.stop()`` via callbacks set during construction; the lifespan only waits on the elector and the watcher. + With sharding: the shard controller owns the runner. Every replica is an + active fitter; the controller does the initial ownership refresh, starts + the runner (which schedules only this replica's owned groups), and — for + ``leader_assigned`` — runs the assigner Lease plus a periodic rebalance + loop. HA and sharding are mutually exclusive, so ``elector`` is ``None`` + whenever ``shard_controller`` is set. + The OTel tracer (when enabled) is shut down *last* so any spans emitted by the runner's final stop sequence still have a chance to flush via the BatchSpanProcessor before the gRPC channel is torn down. """ elector_task: asyncio.Task[None] | None = None - if elector is None: + if shard_controller is not None: + await shard_controller.start() + elif elector is None: await runner.start() else: elector_task = asyncio.create_task(elector.run()) @@ -978,7 +996,9 @@ async def _lifespan_factory( finally: if watcher is not None: await watcher.stop() - if elector is not None: + if shard_controller is not None: + await shard_controller.stop() + elif elector is not None: await elector.stop() if elector_task is not None: import contextlib # noqa: PLC0415 @@ -1004,11 +1024,12 @@ def _build_app( # noqa: PLR0915 source: PromSource | None = None, snapshot_cache: ha_cache_module.SnapshotCache | None = None, elector: LeaderElector | None = None, + shard_controller: _ShardController | None = None, tracer: telemetry_module.Tracer | None = None, ) -> FastAPI: @asynccontextmanager async def lifespan(_: FastAPI) -> AsyncIterator[None]: - async with _lifespan_factory(runner, watcher, elector, tracer=tracer): + async with _lifespan_factory(runner, watcher, elector, shard_controller, tracer=tracer): yield app = FastAPI(lifespan=lifespan, docs_url=None, redoc_url=None, openapi_url=None) @@ -1357,6 +1378,10 @@ def _redis_needed(cfg: config_module.Config) -> bool: """Does this config want a shared Redis client?""" if cfg.high_availability.enabled: return True + # leader_assigned sharding shares the membership registry + the published + # group→shard assignment through Redis (consistent_hash needs no backend). + if cfg.sharding.enabled and cfg.sharding.mode == "leader_assigned": + return True qc = cfg.query_cache if not qc.enabled: return False @@ -1485,6 +1510,170 @@ async def _on_lost() -> None: ) +class _ShardController: + """Drives horizontal sharding for one replica. + + Owns the runner's lifecycle in sharded mode: refreshes ownership before + starting the runner (so it schedules only its owned groups), and — for + ``leader_assigned`` — runs the assigner Lease and a periodic loop that + re-reads membership / republishes the assignment and reconciles the + running group loops with the new ownership. For ``consistent_hash`` there + is nothing periodic to do: ownership is a pure function of the (static) + replica count and index, so the controller just starts the runner. + """ + + def __init__( + self, + *, + coordinator: sharding_module.ShardCoordinator, + runner: Runner, + elector: LeaderElector | None, + refresh_interval_seconds: float, + ) -> None: + self._coordinator = coordinator + self._runner = runner + self._elector = elector + self._interval = max(1.0, refresh_interval_seconds) + self._stop = asyncio.Event() + self._refresh_task: asyncio.Task[None] | None = None + self._elector_task: asyncio.Task[None] | None = None + + async def start(self) -> None: + if self._elector is not None: + self._elector_task = asyncio.create_task(self._elector.run()) + # Resolve ownership before the runner spawns loops so it schedules the + # right slice from the very first tick. + await self._coordinator.refresh(self._runner.group_names()) + await self._runner.start() + if self._coordinator.needs_periodic_refresh(): + self._refresh_task = asyncio.create_task(self._refresh_loop()) + + async def _refresh_loop(self) -> None: + import contextlib # noqa: PLC0415 + + while not self._stop.is_set(): + # Wake either on the interval (TimeoutError → time to refresh) or + # early when stop() fires. + with contextlib.suppress(TimeoutError): + await asyncio.wait_for(self._stop.wait(), timeout=self._interval) + if self._stop.is_set(): + break + try: + await self._coordinator.refresh(self._runner.group_names()) + await self._runner.reconcile_shard_ownership() + except Exception: + logger.exception("shard_refresh_failed") + + async def stop(self) -> None: + import contextlib # noqa: PLC0415 + + self._stop.set() + if self._refresh_task is not None: + self._refresh_task.cancel() + with contextlib.suppress(asyncio.CancelledError, Exception): + await self._refresh_task + if self._elector is not None: + await self._elector.stop() + if self._elector_task is not None: + with contextlib.suppress(asyncio.CancelledError, Exception): + await self._elector_task + await self._coordinator.close() + + +def _build_shard_controller( + cfg: config_module.Config, + exporter: Exporter, + runner: Runner, + redis_client: object | None, +) -> _ShardController | None: + """Construct the shard controller for sharded mode, or ``None`` if off. + + Resolves this replica's shard index / id (from config or the downward-API + env vars a StatefulSet provides), builds the mode-appropriate coordinator, + installs it on the runner, and — for ``leader_assigned`` — an assigner + Lease elector whose leadership drives the published assignment. + """ + sh = cfg.sharding + if not sh.enabled: + return None + + shard_index = sharding_module.derive_shard_index(sh.shard_index) + shard_id = sharding_module.derive_shard_id(sh.shard_id, shard_index) + + coordinator: sharding_module.ShardCoordinator + elector: LeaderElector | None + if sh.mode == "consistent_hash": + if shard_index < 0: + raise ValueError( + "sharding.mode=consistent_hash could not resolve this replica's shard " + "index. Run the forecaster as a StatefulSet so PROMFORECAST_POD_NAME " + "carries an ordinal, or set sharding.shard_index / PROMFORECAST_SHARD_INDEX." + ) + coordinator = sharding_module.ConsistentHashCoordinator( + shard_index=shard_index, + replica_count=sh.replica_count, + shard_id=shard_id, + ) + # Each shard is an autonomous active fitter; reflect that on the + # leader gauge so a join-on-instance dashboard treats every shard as + # "up and fitting" rather than a follower. + exporter.set_leader_state(identity=shard_id, is_leader=True) + elector = None + else: # leader_assigned + if redis_client is None: # pragma: no cover - guarded by config validation + raise ValueError( + "sharding.mode=leader_assigned requires a Redis backend " + "(highAvailability.redis.address or query_cache.redis_address)" + ) + store = sharding_module.RedisAssignmentStore(redis_client) + namespace = os.environ.get("PROMFORECAST_NAMESPACE", "default") + exporter.set_leader_state(identity=shard_id, is_leader=False) + + async def _on_acquired() -> None: + # Assigner leadership only controls who publishes the mapping; every + # replica keeps fitting, so we do NOT start/stop the runner here. + exporter.set_leader_state(identity=shard_id, is_leader=True) + + async def _on_lost() -> None: + exporter.set_leader_state(identity=shard_id, is_leader=False) + + elector_config = LeaderElectorConfig( + namespace=namespace, + lease_name=sh.lease_name, + identity=shard_id, + lease_duration_seconds=sh.lease_duration_seconds, + renew_deadline_seconds=sh.renew_deadline_seconds, + retry_period_seconds=sh.retry_period_seconds, + ) + elector = LeaderElector( + config=elector_config, + client=KubernetesLeaseClient(), + on_started_leading=_on_acquired, + on_stopped_leading=_on_lost, + ) + coordinator = sharding_module.LeaderAssignedCoordinator( + shard_id=shard_id, + store=store, + is_leader=elector.is_leader, + member_ttl_seconds=sh.member_ttl_seconds, + ) + + runner.set_shard_coordinator(coordinator) + logger.info( + "sharding_enabled", + mode=sh.mode, + shard_id=shard_id, + shard_index=shard_index, + replica_count=sh.replica_count, + ) + return _ShardController( + coordinator=coordinator, + runner=runner, + elector=elector, + refresh_interval_seconds=float(sh.heartbeat_interval_seconds), + ) + + def _build_sink(cfg: config_module.Config) -> ForecastSink | None: """Construct the configured push sink, or ``None`` if disabled. diff --git a/forecaster/src/promforecast/point_factory.py b/forecaster/src/promforecast/point_factory.py index 03f3d56..ed9b239 100644 --- a/forecaster/src/promforecast/point_factory.py +++ b/forecaster/src/promforecast/point_factory.py @@ -123,11 +123,21 @@ def __init__( base_labels: dict[str, str], group: str, levels: list[int], + shard_id: str = "", ) -> None: self._metric_id = metric_id self._group = group self._levels = levels self._base_common = {k: v for k, v in base_labels.items() if k != "__name__"} + # In sharded deployments every output point for this series carries a + # ``shard_id`` label identifying the replica that fit it. Folding it + # into ``_base_common`` propagates it uniformly through every family + # the factory builds (forecast, lower/upper, accuracy, deviation, + # quality, ensemble, contribution, growth, …) without touching each + # ``_build_common_labels`` / ``_common_with_id`` call site. Empty when + # sharding is disabled, so single-replica labelsets are unchanged. + if shard_id: + self._base_common["shard_id"] = shard_id def forecast_points( self, diff --git a/forecaster/src/promforecast/runner/_scheduler.py b/forecaster/src/promforecast/runner/_scheduler.py index f9cd3be..d4a4adf 100644 --- a/forecaster/src/promforecast/runner/_scheduler.py +++ b/forecaster/src/promforecast/runner/_scheduler.py @@ -47,6 +47,7 @@ from .. import revision as revision_mod from .. import right_sizing as right_sizing_mod from .. import seasonality as seasonality_mod +from .. import sharding as sharding_mod from .. import slo as slo_mod from .. import telemetry as telemetry_mod from .. import what_if as what_if_mod @@ -300,7 +301,9 @@ def _trailing_regressor_values(regressors: Any | None) -> dict[str, float]: return out -def _build_threshold_action_cost_points(group: GroupConfig) -> list[ThresholdActionCostPoint]: +def _build_threshold_action_cost_points( + group: GroupConfig, *, shard_id: str = "" +) -> list[ThresholdActionCostPoint]: """Project a group's configured thresholds onto action-cost points. Pure projection from config — the points carry the operator's @@ -309,7 +312,11 @@ def _build_threshold_action_cost_points(group: GroupConfig) -> list[ThresholdAct Independent of fit results: a query whose fit failed (or was skipped this tick) still publishes its action_cost rows so the Grafana "rank by cost" panel doesn't go blank. + + ``shard_id`` (sharded mode only) is attached so the family carries the + same ``shard_id`` label as the rest of this replica's output. """ + shard_label = {"shard_id": shard_id} if shard_id else {} out: list[ThresholdActionCostPoint] = [] for query in group.queries: annotations = capacity_mod.build_threshold_action_costs(query.thresholds) @@ -317,6 +324,7 @@ def _build_threshold_action_cost_points(group: GroupConfig) -> list[ThresholdAct out.append( ThresholdActionCostPoint( labels={ + **shard_label, "id": query.id, "group": group.name, "name": ann.name, @@ -447,6 +455,12 @@ def __init__( # Used by HA mode to publish the rendered /metrics payload to the # shared snapshot cache; ``None`` in single-replica deployments. self._snapshot_publisher: Callable[[], Awaitable[None]] | None = None + # Optional shard coordinator (sharding mode). ``None`` means "own every + # group" — the single-replica / HA default, unchanged behaviour. When + # set, only groups this replica owns get a scheduled loop, and output + # points carry the ``shard_id`` label. + self._shard_coordinator: sharding_mod.ShardCoordinator | None = None + self._shard_id: str = "" self._locks: dict[str, asyncio.Lock] = {g.name: asyncio.Lock() for g in config.groups} # Per-group task: keyed by group.name, holds the long-lived loop. We # key per-group (rather than a flat list) so apply_config can diff and @@ -474,6 +488,41 @@ def set_snapshot_publisher(self, publisher: Callable[[], Awaitable[None]] | None """ self._snapshot_publisher = publisher + def set_shard_coordinator(self, coordinator: sharding_mod.ShardCoordinator | None) -> None: + """Install (or clear) the shard coordinator for horizontal sharding. + + When set, this replica only schedules the groups the coordinator says + it owns, and every output point carries the coordinator's ``shard_id`` + label. Called once at boot from :func:`main.run` before ``start``. + """ + self._shard_coordinator = coordinator + self._shard_id = coordinator.shard_id if coordinator is not None else "" + + def group_names(self) -> list[str]: + """Names of the currently-configured groups (reload-aware). + + Read by the shard refresh loop so the ``leader_assigned`` assigner + always computes its mapping over the live group set, even across a + config reload. + """ + return [g.name for g in self._config.groups] + + def _owns_group(self, group_name: str) -> bool: + """Does this replica own ``group_name``? True for all when unsharded.""" + if self._shard_coordinator is None: + return True + return self._shard_coordinator.owns(group_name) + + def _shard_label(self) -> dict[str, str]: + """``{"shard_id": ...}`` in sharded mode, else ``{}``. + + Spread into the handful of operational metric label dicts that are + built by hand rather than via :class:`SeriesPointFactory` (which folds + the same label in for everything it emits), so the ``shard_id`` label + is present consistently across every per-series output family. + """ + return {"shard_id": self._shard_id} if self._shard_id else {} + async def start(self) -> None: """Run each group once immediately, then on the configured interval.""" # Guard against double-start: HA mode's leader-acquired callback @@ -484,13 +533,109 @@ async def start(self) -> None: self._stopped.clear() for group in self._config.groups: self._spawn_group_loop(group) + self._publish_shard_ownership() def _spawn_group_loop(self, group: GroupConfig) -> None: - """Create the stop event + task for one group. Idempotent per name.""" + """Create the stop event + task for one group. Idempotent per name. + + In sharded mode this is a no-op for groups this replica does not own — + ownership is enforced here so every spawn path (``start``, reload, and + shard reconciliation) funnels through one gate. + """ + if not self._owns_group(group.name): + return ev = asyncio.Event() self._stop_events[group.name] = ev self._tasks[group.name] = asyncio.create_task(self._loop(group, ev)) + def _publish_shard_ownership(self) -> None: + """Refresh the ``forecast_shard_owns_group`` gauge from current owners.""" + if self._shard_coordinator is None: + return + owned = self._shard_coordinator.owned_groups(g.name for g in self._config.groups) + self._exporter.set_shard_state(shard_id=self._shard_id, owned_groups=owned) + + async def reconcile_shard_ownership(self) -> None: + """Align running group loops with current shard ownership. + + Spawns loops for newly-owned groups and drains loops for groups this + replica no longer owns (e.g. after a ``leader_assigned`` rebalance when + a replica joined). Snapshots for relinquished groups are dropped so a + replica stops serving stale forecasts for groups it handed off. A no-op + when sharding is disabled or the scheduler isn't running. + """ + if self._shard_coordinator is None or self._stopped.is_set(): + self._publish_shard_ownership() + return + owned = self._shard_coordinator.owned_groups(g.name for g in self._config.groups) + running = set(self._tasks) + to_stop = running - owned + to_start = owned - running + for name in sorted(to_stop): + await self._drain_group(name) + if to_stop: + self._exporter.retain_groups(set(self._tasks)) + # Drop per-(group, query) tracker state for relinquished groups, + # mirroring the reload prune. Otherwise a group later reassigned + # back to this replica would score drift / revision against a + # forecast made *before* the handoff — stale, because another + # shard produced this group's forecasts in the interim. Prune to + # the (group, query) pairs this replica still owns. + self._prune_caches_to_groups(owned) + groups_by_name = {g.name: g for g in self._config.groups} + for name in sorted(to_start): + group = groups_by_name.get(name) + if group is not None: + self._spawn_group_loop(group) + if to_start or to_stop: + logger.info( + "shard_ownership_reconciled", + shard_id=self._shard_id, + started=sorted(to_start), + stopped=sorted(to_stop), + ) + self._publish_shard_ownership() + + def _prune_caches_to_groups(self, owned: set[str]) -> None: + """Forget per-(group, query) tracker state outside ``owned``. + + Mirrors the reload prune (:meth:`apply_config`) but scoped to the + groups this shard owns, so a ``leader_assigned`` rebalance that hands a + group off doesn't leave a stale drift / revision baseline behind. + """ + live_pairs = { + (g.name, q.id) for g in self._config.groups if g.name in owned for q in g.queries + } + self._drift_cache.prune_to(live_pairs) + self._revision_cache.prune_to(live_pairs) + self._band_coverage.forget(live_pairs) + self._drift_windows = { + key: window + for key, window in self._drift_windows.items() + if (key[0], key[1]) in live_pairs + } + + async def _drain_group(self, name: str) -> None: + """Cooperatively stop a single group's loop and await its drain. + + Same contract as :meth:`_drain_all_groups` but for one group: signal + the stop event, let an in-flight ``run_group`` finish and write its + snapshot, then hard-cancel if it overruns the grace window. + """ + ev = self._stop_events.pop(name, None) + task = self._tasks.pop(name, None) + if ev is not None: + ev.set() + if task is None: + return + grace = self._drain_grace_seconds(self._config) + try: + await asyncio.wait_for(asyncio.gather(task, return_exceptions=True), timeout=grace) + except TimeoutError: + task.cancel() + with contextlib.suppress(asyncio.CancelledError, Exception): + await task + async def apply_config(self, new_config: Config) -> ReloadResult: """Swap in a new validated config, draining in-flight runs first. @@ -666,6 +811,11 @@ async def apply_config(self, new_config: Config) -> ReloadResult: await self._rebuild_integrations(old_config=old) for group in new_config.groups: self._spawn_group_loop(group) + # The group list may have changed which groups this replica owns + # (new groups hash to this shard, removed ones drop off); refresh + # the ownership gauge to match. ``leader_assigned`` republishes its + # mapping for the new group set on the next periodic tick. + self._publish_shard_ownership() return ReloadResult(added=added, removed=removed, changed=changed) @@ -1028,7 +1178,9 @@ async def run_group( # noqa: PLR0912, PLR0915 # /metrics on the next refresh; partial-tick refreshes still # publish the full set because every configured query is # walked here regardless of which subset actually ran. - threshold_action_cost_points = _build_threshold_action_cost_points(group) + threshold_action_cost_points = _build_threshold_action_cost_points( + group, shard_id=self._shard_id + ) self._exporter.update( GroupSnapshot( group=group.name, @@ -1934,6 +2086,7 @@ def _record_series_fit_successes( for model_name in models: labels = { **source_labels, + **self._shard_label(), "id": query_id, "group": group_name, "model": model_name, @@ -2106,6 +2259,7 @@ async def _try_cold_start( base_labels=frame.labels, group=group.name, levels=list(d.confidence_levels), + shard_id=self._shard_id, ) fit_result = _FitResult() forecast_points = [ @@ -2167,6 +2321,7 @@ async def _try_cold_start( # confident the sibling set looked. quality_label = { **{k: v for k, v in frame.labels.items() if k != "__name__"}, + **self._shard_label(), "id": query.id, "group": group.name, "model": cold_start_mod.COLD_START_MODEL_NAME, @@ -2256,6 +2411,7 @@ async def _run_event_pattern_query( points = [ IncidentRecurrencePoint( labels={ + **self._shard_label(), "group": group.name, "id": query.id, "day_of_week": str(bucket.day_of_week), @@ -2443,6 +2599,7 @@ async def _fit_series( # noqa: PLR0912, PLR0915 base_labels=frame.labels, group=group_name, levels=levels, + shard_id=self._shard_id, ) result = _FitResult() @@ -2976,6 +3133,7 @@ def _record_revision_and_disagreement( series_key = _series_key_from_labels(frame.labels) common = { **{k: v for k, v in frame.labels.items() if k != "__name__"}, + **self._shard_label(), "id": query.id, "group": group_name, } diff --git a/forecaster/src/promforecast/sharding.py b/forecaster/src/promforecast/sharding.py new file mode 100644 index 0000000..09ee794 --- /dev/null +++ b/forecaster/src/promforecast/sharding.py @@ -0,0 +1,397 @@ +"""Horizontal sharding: N active replicas each owning a disjoint slice of groups. + +Where :mod:`promforecast.leader` gives a *single* active replica (the leader) +plus hot followers serving a cached snapshot, sharding makes *every* replica an +active fitter that owns a subset of the configured groups. Group ownership is +disjoint (no group is fit twice) and complete (every group is owned by exactly +one live replica), so the aggregate fit budget scales with the replica count. + +Two ownership modes, both reducible to a pure ``owns(group) -> bool`` decision: + +* **consistent_hash** — fully local. Each replica knows its own 0-based shard + index and the total replica count; it owns a group iff + ``stable_bucket(group) % replica_count == shard_index``. No coordination, + no shared state — the only agreement needed is the replica count, which the + chart pins by setting both the StatefulSet ``replicas`` and the config's + ``replica_count`` from one value. The pod ordinal (StatefulSet gives stable + ``-`` names) is the shard index. + +* **leader_assigned** — coordinated. A ``coordination.k8s.io/v1`` Lease elects + one replica as the *assigner*. Every replica heartbeats its identity into a + shared Redis membership registry; the assigner reads the live membership, + computes a balanced group→shard mapping, and publishes it back to Redis. + Every replica (assigner included) reads the published mapping to learn what + it owns. Membership changes (a replica joins or leaves) cause the assigner + to republish on its next tick, so ownership rebalances automatically. The + mapping uses rendezvous (highest-random-weight) hashing so a membership + change moves only ~1/N of the groups rather than reshuffling everything. + +The ownership decision is deliberately split from the transport: the pure +functions and the in-memory coordinators carry no I/O, and the Redis-backed +:class:`RedisAssignmentStore` is the only piece that touches the network. Tests +exercise the logic against an in-memory store. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import time +from collections.abc import Callable, Iterable +from typing import Protocol, runtime_checkable + +import structlog + +logger = structlog.get_logger(__name__) + + +# --------------------------------------------------------------------------- +# Pure ownership logic +# --------------------------------------------------------------------------- + + +def stable_bucket(name: str) -> int: + """Deterministic non-negative hash of ``name``, stable across processes. + + Python's built-in ``hash()`` is salted per-process (``PYTHONHASHSEED``), + so two replicas would disagree on ownership. We hash with SHA-1 and take + the leading 8 bytes — overkill cryptographically, but it gives every + replica (and every test run) the same integer for the same string. + """ + # SHA-1 here is a fast, stable bucketing hash — not a security primitive. + digest = hashlib.sha1(name.encode("utf-8")).digest() + return int.from_bytes(digest[:8], "big") + + +def consistent_hash_owns(group_name: str, shard_index: int, replica_count: int) -> bool: + """True iff ``shard_index`` owns ``group_name`` under modulo hashing. + + ``replica_count`` is clamped to at least 1 so a misconfigured zero can't + divide by zero; a single shard then owns everything, which is the safe + degenerate behaviour (identical to sharding disabled). + """ + count = max(1, replica_count) + if shard_index < 0 or shard_index >= count: + return False + return stable_bucket(group_name) % count == shard_index + + +def consistent_hash_assignment(group_names: Iterable[str], replica_count: int) -> dict[str, int]: + """Full ``group -> shard_index`` map under modulo hashing (for debugging).""" + count = max(1, replica_count) + return {name: stable_bucket(name) % count for name in group_names} + + +def assign_groups(group_names: Iterable[str], members: Iterable[str]) -> dict[str, str]: + """Balanced ``group -> shard_id`` assignment via rendezvous hashing. + + For each group, the owner is the member maximising ``stable_bucket`` of the + ``group\\x00member`` pair. Rendezvous (HRW) hashing is balanced in + expectation and — crucially for ``leader_assigned`` rebalancing — moves + only the groups whose previous owner left (or ~1/N of groups when a member + joins) rather than reshuffling the whole map the way ``index % count`` + round-robin would. Ties (astronomically unlikely with a 64-bit space) break + on the lexically smaller member id for determinism. + + Returns an empty map when there are no members — callers treat "no + assignment yet" as "own nothing" and wait for the next assigner tick. + """ + member_list = sorted(set(members)) + if not member_list: + return {} + mapping: dict[str, str] = {} + for group in group_names: + best_member = member_list[0] + best_weight = stable_bucket(f"{group}\x00{best_member}") + for member in member_list[1:]: + weight = stable_bucket(f"{group}\x00{member}") + if weight > best_weight or (weight == best_weight and member < best_member): + best_member, best_weight = member, weight + mapping[group] = best_member + return mapping + + +def derive_shard_index(configured: int) -> int: + """Resolve this replica's shard index, deriving from the environment. + + Precedence: an explicit non-negative ``configured`` value wins (local/test + runs); otherwise ``PROMFORECAST_SHARD_INDEX``; otherwise the trailing + ordinal of ``PROMFORECAST_POD_NAME`` (``promforecast-3`` → ``3``), which is + what a StatefulSet provides. Returns ``-1`` when none resolve so the caller + can fail loudly rather than silently owning nothing. + """ + if configured >= 0: + return configured + env_index = os.environ.get("PROMFORECAST_SHARD_INDEX") + if env_index is not None and env_index.strip().lstrip("-").isdigit(): + return int(env_index) + pod = os.environ.get("PROMFORECAST_POD_NAME", "") + _, _, suffix = pod.rpartition("-") + if suffix.isdigit(): + return int(suffix) + return -1 + + +def derive_shard_id(configured: str, shard_index: int) -> str: + """Resolve the ``shard_id`` label value. + + Precedence: explicit ``configured`` value, then ``PROMFORECAST_POD_NAME`` + (the natural per-replica identity), then ``shard-``. + """ + if configured: + return configured + pod = os.environ.get("PROMFORECAST_POD_NAME", "") + if pod: + return pod + return f"shard-{shard_index}" + + +# --------------------------------------------------------------------------- +# Coordinators +# --------------------------------------------------------------------------- + + +@runtime_checkable +class ShardCoordinator(Protocol): + """What the runner needs to decide and observe group ownership. + + ``owns`` is synchronous and hot-path cheap (called whenever a group loop + is (re)spawned). ``refresh`` is async and does any I/O (heartbeat, read or + publish the assignment); it is a no-op for fully-local modes. The runner + calls ``refresh`` periodically only when ``needs_periodic_refresh`` is true. + """ + + @property + def shard_id(self) -> str: ... + + def owns(self, group_name: str) -> bool: ... + + def owned_groups(self, group_names: Iterable[str]) -> set[str]: ... + + def needs_periodic_refresh(self) -> bool: ... + + async def refresh(self, group_names: Iterable[str]) -> None: ... + + async def close(self) -> None: ... + + +class ConsistentHashCoordinator: + """Local, coordination-free ownership by ``hash(group) % replica_count``.""" + + def __init__(self, *, shard_index: int, replica_count: int, shard_id: str) -> None: + self._shard_index = shard_index + self._replica_count = max(1, replica_count) + self._shard_id = shard_id + + @property + def shard_id(self) -> str: + return self._shard_id + + def owns(self, group_name: str) -> bool: + return consistent_hash_owns(group_name, self._shard_index, self._replica_count) + + def owned_groups(self, group_names: Iterable[str]) -> set[str]: + return {name for name in group_names if self.owns(name)} + + def needs_periodic_refresh(self) -> bool: + # Ownership is a pure function of (group, index, count) — nothing to + # refresh. The runner reconciles only when the group list changes. + return False + + async def refresh(self, group_names: Iterable[str]) -> None: + return None + + async def close(self) -> None: + return None + + +@runtime_checkable +class AssignmentStore(Protocol): + """Shared membership registry + published assignment for leader_assigned.""" + + async def heartbeat(self, shard_id: str, ttl_seconds: int) -> None: ... + + async def members(self) -> list[str]: ... + + async def publish_assignment(self, mapping: dict[str, str]) -> None: ... + + async def read_assignment(self) -> dict[str, str]: ... + + +class LeaderAssignedCoordinator: + """Ownership read from a Lease-elected assigner's published mapping. + + Every replica heartbeats into the shared store. The replica currently + holding the assigner Lease (``is_leader()`` returns true) recomputes the + balanced mapping over the live membership and publishes it; all replicas + read the published mapping to decide ownership. + + Two windows need care: + + * **Cold start** — no mapping published yet. We own *nothing* and wait for + the assigner (the Lease guarantees one is elected within a lease + duration). Owning everything here would make every replica fit every + group at once — an N-times thundering herd at exactly the moment sharding is + supposed to prevent it. A brief startup gap is the better trade. + * **A just-added group** — the mapping exists but doesn't yet name a group + added by a reload before the assigner republished. For that single group + we fall back to a local rendezvous decision over the membership; because + rendezvous is deterministic given the same membership, every replica + agrees on the owner, so the group is covered by exactly one replica. + """ + + def __init__( + self, + *, + shard_id: str, + store: AssignmentStore, + is_leader: Callable[[], bool], + member_ttl_seconds: int, + ) -> None: + self._shard_id = shard_id + self._store = store + self._is_leader = is_leader + self._member_ttl = member_ttl_seconds + self._assignment: dict[str, str] = {} + self._members: list[str] = [shard_id] + + @property + def shard_id(self) -> str: + return self._shard_id + + def owns(self, group_name: str) -> bool: + owner = self._assignment.get(group_name) + if owner is not None: + return owner == self._shard_id + if not self._assignment: + # Cold start: no mapping at all yet. Own nothing and wait for the + # assigner rather than have every replica fit every group at once. + return False + # Mapping exists but doesn't name this group (a reload added it before + # the assigner republished). Decide locally via rendezvous over the + # membership — deterministic, so every replica agrees on one owner. + fallback = assign_groups([group_name], self._members) + return fallback.get(group_name) == self._shard_id + + def owned_groups(self, group_names: Iterable[str]) -> set[str]: + return {name for name in group_names if self.owns(name)} + + def needs_periodic_refresh(self) -> bool: + return True + + async def refresh(self, group_names: Iterable[str]) -> None: + # Always heartbeat first so a transient read/publish failure below + # doesn't also drop us from the membership. + try: + await self._store.heartbeat(self._shard_id, self._member_ttl) + except Exception as exc: # pragma: no cover - logged, degrades gracefully + logger.warning("shard_heartbeat_failed", shard_id=self._shard_id, error=str(exc)) + try: + members = await self._store.members() + except Exception as exc: + logger.warning("shard_members_read_failed", error=str(exc)) + members = [] + # Keep ourselves in the membership even if the store hasn't observed + # our heartbeat yet (read-after-write lag), so the fallback path never + # decides we own nothing. + self._members = sorted(set(members) | {self._shard_id}) + if self._is_leader(): + mapping = assign_groups(sorted(group_names), self._members) + try: + await self._store.publish_assignment(mapping) + except Exception as exc: + logger.warning("shard_assignment_publish_failed", error=str(exc)) + self._assignment = mapping + logger.info( + "shard_assignment_published", + members=self._members, + groups=len(mapping), + ) + else: + try: + self._assignment = await self._store.read_assignment() + except Exception as exc: + logger.warning("shard_assignment_read_failed", error=str(exc)) + + async def close(self) -> None: + return None + + +# --------------------------------------------------------------------------- +# Redis-backed assignment store +# --------------------------------------------------------------------------- + +_MEMBER_KEY_PREFIX = "promforecast:shard:member:" +_ASSIGNMENT_KEY = "promforecast:shard:assignment" + + +class RedisAssignmentStore: + """:class:`AssignmentStore` over the shared async Redis client. + + Membership is a set of short-TTL per-replica keys (``member:``) so a + dead replica simply expires out — no explicit deregistration needed. The + assignment is a single JSON blob under one key, refreshed by the assigner + every tick (no TTL — a stale-but-present map is better than none, and the + assigner overwrites it continuously). + + Membership is enumerated with ``SCAN`` rather than ``KEYS``: this Redis is + typically shared with the query cache, and ``KEYS`` blocks the server for + an O(whole-keyspace) match on every heartbeat. ``SCAN`` is cursor-based and + non-blocking; the matching set (one key per replica) is tiny, so the cursor + loop is cheap regardless of how large the surrounding keyspace grows. + """ + + def __init__(self, client: object) -> None: + self._client = client + + async def heartbeat(self, shard_id: str, ttl_seconds: int) -> None: + key = f"{_MEMBER_KEY_PREFIX}{shard_id}" + await self._client.set(key, str(time.time()), ex=max(1, ttl_seconds)) # type: ignore[attr-defined] + + async def members(self) -> list[str]: + pattern = f"{_MEMBER_KEY_PREFIX}*" + members: list[str] = [] + cursor = 0 + while True: + cursor, raw_keys = await self._client.scan( # type: ignore[attr-defined] + cursor=cursor, match=pattern, count=100 + ) + for raw in raw_keys or []: + key = raw.decode("utf-8") if isinstance(raw, bytes | bytearray) else str(raw) + members.append(key[len(_MEMBER_KEY_PREFIX) :]) + # redis-py returns the cursor as an int; 0 marks a completed sweep. + if int(cursor) == 0: + break + return members + + async def publish_assignment(self, mapping: dict[str, str]) -> None: + await self._client.set(_ASSIGNMENT_KEY, json.dumps(mapping)) # type: ignore[attr-defined] + + async def read_assignment(self) -> dict[str, str]: + raw = await self._client.get(_ASSIGNMENT_KEY) # type: ignore[attr-defined] + if raw is None: + return {} + text = raw.decode("utf-8") if isinstance(raw, bytes | bytearray) else str(raw) + try: + data = json.loads(text) + except (ValueError, TypeError): + return {} + if not isinstance(data, dict): + return {} + return {str(k): str(v) for k, v in data.items()} + + +__all__ = [ + "AssignmentStore", + "ConsistentHashCoordinator", + "LeaderAssignedCoordinator", + "RedisAssignmentStore", + "ShardCoordinator", + "assign_groups", + "consistent_hash_assignment", + "consistent_hash_owns", + "derive_shard_id", + "derive_shard_index", + "stable_bucket", +] diff --git a/forecaster/tests/test_chart_sharding.py b/forecaster/tests/test_chart_sharding.py new file mode 100644 index 0000000..3a005cc --- /dev/null +++ b/forecaster/tests/test_chart_sharding.py @@ -0,0 +1,115 @@ +"""Chart-template tests for sharding. + +consistent_hash sharding must render a StatefulSet with stable ordinals (and +no Deployment); leader_assigned must render a Deployment plus the Lease RBAC and +flow the Redis address into the forecaster config. The chart values must +round-trip through the forecaster's own loader (mirrors the CI drift check). + +Like the other chart-template tests these shell out to ``helm`` and skip +cleanly when it isn't on PATH. +""" + +from __future__ import annotations + +import shutil +import subprocess +from pathlib import Path + +import pytest +import yaml + +REPO_ROOT = Path(__file__).resolve().parents[2] +STANDALONE_CHART = REPO_ROOT / "charts" / "promforecast" + +helm = shutil.which("helm") +pytestmark = pytest.mark.skipif( + helm is None, reason="`helm` not on PATH — chart-template tests are skipped" +) + + +def _helm_template(*sets: str) -> list[dict]: + cmd = [helm, "template", str(STANDALONE_CHART)] + for s in sets: + cmd.extend(["--set", s]) + result = subprocess.run(cmd, capture_output=True, text=True, check=True) + return [d for d in yaml.safe_load_all(result.stdout) if isinstance(d, dict)] + + +def _kinds(docs: list[dict]) -> set[str]: + return {d.get("kind") for d in docs} + + +def _find(docs: list[dict], kind: str) -> dict | None: + return next((d for d in docs if d.get("kind") == kind), None) + + +def _config_yaml(docs: list[dict]) -> dict: + cm = next( + d for d in docs if d.get("kind") == "ConfigMap" and "config.yaml" in d.get("data", {}) + ) + return yaml.safe_load(cm["data"]["config.yaml"]) + + +def test_default_renders_deployment_not_statefulset() -> None: + docs = _helm_template() + assert "Deployment" in _kinds(docs) + assert "StatefulSet" not in _kinds(docs) + + +def test_consistent_hash_renders_statefulset() -> None: + docs = _helm_template( + "sharding.enabled=true", + "sharding.mode=consistent_hash", + "replicaCount=4", + ) + kinds = _kinds(docs) + assert "StatefulSet" in kinds + assert "Deployment" not in kinds + sts = _find(docs, "StatefulSet") + assert sts is not None + assert sts["spec"]["replicas"] == 4 + assert sts["spec"]["serviceName"] + assert sts["spec"]["podManagementPolicy"] == "Parallel" + # Config carries the resolved sharding block with replica_count from replicas. + cfg = _config_yaml(docs) + assert cfg["sharding"]["enabled"] is True + assert cfg["sharding"]["mode"] == "consistent_hash" + assert cfg["sharding"]["replica_count"] == 4 + # No Lease RBAC for consistent_hash (no coordination backend). + assert "Role" not in kinds + + +def test_consistent_hash_injects_pod_name_env() -> None: + docs = _helm_template("sharding.enabled=true", "sharding.mode=consistent_hash") + sts = _find(docs, "StatefulSet") + assert sts is not None + env = sts["spec"]["template"]["spec"]["containers"][0]["env"] + names = {e["name"] for e in env} + assert "PROMFORECAST_SHARDING_ENABLED" in names + assert "PROMFORECAST_POD_NAME" in names + + +def test_leader_assigned_renders_deployment_and_lease_rbac() -> None: + docs = _helm_template( + "sharding.enabled=true", + "sharding.mode=leader_assigned", + "replicaCount=3", + "sharding.redis.address=redis-master.redis.svc:6379", + ) + kinds = _kinds(docs) + assert "Deployment" in kinds + assert "StatefulSet" not in kinds + assert "Role" in kinds and "RoleBinding" in kinds + dep = _find(docs, "Deployment") + assert dep is not None + assert dep["spec"]["strategy"]["type"] == "RollingUpdate" + # Redis address flows into the forecaster's query-cache config. + cfg = _config_yaml(docs) + assert cfg["sharding"]["mode"] == "leader_assigned" + assert cfg["query_cache"]["redis_address"] == "redis-master.redis.svc:6379" + + +def test_leader_assigned_requires_redis_address() -> None: + # Schema gate: enabling leader_assigned without a redis address fails. + with pytest.raises(subprocess.CalledProcessError): + _helm_template("sharding.enabled=true", "sharding.mode=leader_assigned") diff --git a/forecaster/tests/test_config_sharding.py b/forecaster/tests/test_config_sharding.py new file mode 100644 index 0000000..2faf233 --- /dev/null +++ b/forecaster/tests/test_config_sharding.py @@ -0,0 +1,80 @@ +"""Validation tests for the ``sharding`` config block.""" + +from __future__ import annotations + +import pytest +from pydantic import ValidationError + +from promforecast.config import ( + Config, + DatasourceConfig, + GroupConfig, + HighAvailabilityConfig, + QueryCacheConfig, + QueryConfig, + ShardingConfig, +) + + +def _base_config(**overrides: object) -> Config: + return Config( + datasource=DatasourceConfig(url="http://x"), + groups=[GroupConfig(name="g", queries=[QueryConfig(id="m", promql="up")])], + **overrides, + ) + + +def test_sharding_defaults_off() -> None: + cfg = _base_config() + assert cfg.sharding.enabled is False + assert cfg.sharding.mode == "consistent_hash" + assert cfg.sharding.replica_count == 1 + assert cfg.sharding.shard_index == -1 + + +def test_sharding_consistent_hash_needs_no_redis() -> None: + cfg = _base_config( + sharding=ShardingConfig(enabled=True, mode="consistent_hash", replica_count=4) + ) + assert cfg.sharding.enabled is True + + +def test_sharding_leader_assigned_requires_redis() -> None: + with pytest.raises(ValidationError, match="leader_assigned requires a Redis backend"): + _base_config(sharding=ShardingConfig(enabled=True, mode="leader_assigned", replica_count=3)) + + +def test_sharding_leader_assigned_accepts_query_cache_redis() -> None: + cfg = _base_config( + sharding=ShardingConfig(enabled=True, mode="leader_assigned", replica_count=3), + query_cache=QueryCacheConfig(redis_address="redis:6379"), + ) + assert cfg.sharding.mode == "leader_assigned" + + +def test_sharding_and_ha_mutually_exclusive() -> None: + with pytest.raises(ValidationError, match="mutually exclusive"): + _base_config( + sharding=ShardingConfig(enabled=True, mode="consistent_hash", replica_count=2), + high_availability=HighAvailabilityConfig(enabled=True, redis={"address": "redis:6379"}), + ) + + +def test_sharding_shard_index_must_be_below_replica_count() -> None: + with pytest.raises(ValidationError, match="must be < replica_count"): + ShardingConfig(enabled=True, replica_count=3, shard_index=3) + + +def test_sharding_member_ttl_must_exceed_heartbeat() -> None: + with pytest.raises(ValidationError, match="member_ttl_seconds must be >"): + ShardingConfig(heartbeat_interval_seconds=30, member_ttl_seconds=30) + + +def test_sharding_renew_deadline_below_lease_duration() -> None: + with pytest.raises(ValidationError, match="renew_deadline_seconds must be <"): + ShardingConfig(lease_duration_seconds=20, renew_deadline_seconds=20) + + +def test_sharding_rejects_unknown_mode() -> None: + with pytest.raises(ValidationError): + ShardingConfig(mode="round_robin") # type: ignore[arg-type] diff --git a/forecaster/tests/test_leader_assigned_kind.py b/forecaster/tests/test_leader_assigned_kind.py new file mode 100644 index 0000000..0194da7 --- /dev/null +++ b/forecaster/tests/test_leader_assigned_kind.py @@ -0,0 +1,184 @@ +"""Real-cluster integration test for the ``leader_assigned`` elector path. + +This is the one part of sharding the in-process tests can't reach: the +``KubernetesLeaseClient`` actually talking to a ``coordination.k8s.io/v1`` +Lease, a ``LeaderElector`` acquiring/renewing it, and a +``LeaderAssignedCoordinator`` whose ``is_leader`` is driven by that real +election. We assert the end-to-end invariants against a live API server: + +* exactly one of two replicas becomes the assigner, +* the assigner publishes a complete + disjoint group assignment, +* the follower reads the same assignment, and +* leadership (and the assignment) fail over when the assigner leaves. + +It is **skipped** unless ``PROMFORECAST_KIND_TEST`` is set, so the normal CI +``test`` job (no cluster) is unaffected. To run it locally: + + kind create cluster --name promforecast-shard-test + PROMFORECAST_KIND_TEST=1 uv run pytest tests/test_leader_assigned_kind.py -m kind + kind delete cluster --name promforecast-shard-test + +The membership registry uses an in-memory TTL store (Redis is validated +elsewhere); the point here is the Kubernetes election wiring. +""" + +from __future__ import annotations + +import asyncio +import os +import time +import uuid +from collections.abc import Callable + +import pytest + +from promforecast import sharding +from promforecast.leader import LeaderElector, LeaderElectorConfig + +pytest.importorskip("kubernetes", reason="leader_assigned kind test needs the 'ha' extra") + +pytestmark = [ + pytest.mark.kind, + pytest.mark.skipif( + not os.environ.get("PROMFORECAST_KIND_TEST"), + reason="set PROMFORECAST_KIND_TEST=1 (kubeconfig pointing at a throwaway cluster) to run", + ), +] + +_NAMESPACE = os.environ.get("PROMFORECAST_KIND_NAMESPACE", "default") + + +class _TTLStore: + """In-memory AssignmentStore that honours heartbeat TTLs. + + Mirrors :class:`RedisAssignmentStore` semantics (members expire, single + published map) so the coordinator behaves exactly as it would against + Redis — the cluster dependency here is only the Lease, not the store. + """ + + def __init__(self) -> None: + self._members: dict[str, tuple[float, int]] = {} + self._assignment: dict[str, str] = {} + + async def heartbeat(self, shard_id: str, ttl_seconds: int) -> None: + self._members[shard_id] = (time.monotonic(), ttl_seconds) + + async def members(self) -> list[str]: + now = time.monotonic() + return [s for s, (ts, ttl) in self._members.items() if now - ts <= ttl] + + async def publish_assignment(self, mapping: dict[str, str]) -> None: + self._assignment = dict(mapping) + + async def read_assignment(self) -> dict[str, str]: + return dict(self._assignment) + + +async def _wait_for(predicate: Callable[[], bool], timeout: float = 25.0) -> bool: + loop = asyncio.get_event_loop() + deadline = loop.time() + timeout + while loop.time() < deadline: + if predicate(): + return True + await asyncio.sleep(0.25) + return predicate() + + +def _delete_lease(name: str) -> None: + """Best-effort cleanup so reruns don't collide on a leftover Lease.""" + try: + from kubernetes import client as k8s_client # noqa: PLC0415 + from kubernetes import config as k8s_config # noqa: PLC0415 + + try: + k8s_config.load_incluster_config() + except Exception: + k8s_config.load_kube_config() + k8s_client.CoordinationV1Api().delete_namespaced_lease(name=name, namespace=_NAMESPACE) + except Exception: + pass + + +async def _drive(coord: sharding.LeaderAssignedCoordinator, groups: list[str], rounds: int) -> None: + for _ in range(rounds): + await coord.refresh(groups) + await asyncio.sleep(0.1) + + +async def test_leader_assigned_elects_and_assigns_against_real_lease() -> None: + from promforecast.leader import KubernetesLeaseClient # noqa: PLC0415 + + lease_name = f"promforecast-shard-test-{uuid.uuid4().hex[:8]}" + # Short, valid lease params for a snappy test (renew < duration). + cfg_a = LeaderElectorConfig( + namespace=_NAMESPACE, + lease_name=lease_name, + identity="shard-a", + lease_duration_seconds=5, + renew_deadline_seconds=3, + retry_period_seconds=1, + ) + cfg_b = LeaderElectorConfig( + namespace=_NAMESPACE, + lease_name=lease_name, + identity="shard-b", + lease_duration_seconds=5, + renew_deadline_seconds=3, + retry_period_seconds=1, + ) + elector_a = LeaderElector(config=cfg_a, client=KubernetesLeaseClient()) + elector_b = LeaderElector(config=cfg_b, client=KubernetesLeaseClient()) + store = _TTLStore() + coord_a = sharding.LeaderAssignedCoordinator( + shard_id="shard-a", store=store, is_leader=elector_a.is_leader, member_ttl_seconds=3 + ) + coord_b = sharding.LeaderAssignedCoordinator( + shard_id="shard-b", store=store, is_leader=elector_b.is_leader, member_ttl_seconds=3 + ) + groups = [f"g{i}" for i in range(24)] + + task_a = asyncio.create_task(elector_a.run()) + task_b = asyncio.create_task(elector_b.run()) + tasks: list[asyncio.Task[None]] = [task_a, task_b] + try: + # (1) Real Lease election: exactly one assigner. + assert await _wait_for(lambda: elector_a.is_leader() ^ elector_b.is_leader()), ( + "no single assigner was elected against the real Lease API" + ) + + # (2)+(3) Drive both coordinators; the assigner publishes, the follower + # reads. A few rounds let membership populate and the map propagate. + await asyncio.gather(_drive(coord_a, groups, 4), _drive(coord_b, groups, 4)) + owned_a = coord_a.owned_groups(groups) + owned_b = coord_b.owned_groups(groups) + assert owned_a | owned_b == set(groups), "assignment is not complete" + assert not (owned_a & owned_b), "a group is owned by both shards" + assert owned_a and owned_b, "assignment is not balanced across two live shards" + + # (4) Failover: the assigner leaves. The other replica must take the + # Lease over and, once the departed member's heartbeat TTL lapses, + # reassign every group to itself. + leader_is_a = elector_a.is_leader() + (leaving, leaving_task) = (elector_a, task_a) if leader_is_a else (elector_b, task_b) + survivor = elector_b if leader_is_a else elector_a + survivor_coord = coord_b if leader_is_a else coord_a + await leaving.stop() + await leaving_task + tasks.remove(leaving_task) + + assert await _wait_for(survivor.is_leader), "survivor did not take over the Lease" + # Let the departed shard's membership expire (ttl=3s), then republish. + await asyncio.sleep(3.5) + await _drive(survivor_coord, groups, 4) + assert survivor_coord.owned_groups(groups) == set(groups), ( + "survivor did not reabsorb all groups after the assigner left" + ) + finally: + for elector, task in ((elector_a, task_a), (elector_b, task_b)): + if task in tasks: + await elector.stop() + from contextlib import suppress # noqa: PLC0415 + + with suppress(asyncio.CancelledError, Exception): + await task + _delete_lease(lease_name) diff --git a/forecaster/tests/test_runner_sharding.py b/forecaster/tests/test_runner_sharding.py new file mode 100644 index 0000000..1730924 --- /dev/null +++ b/forecaster/tests/test_runner_sharding.py @@ -0,0 +1,173 @@ +"""Runner-level sharding behaviour: ownership filtering, reconciliation, and +the ``shard_id`` label threading into output points. + +The pure ownership logic is covered in ``test_sharding``; here we assert the +runner *acts* on a coordinator — only owned groups get scheduled loops, an +ownership change reconciles the running set, and the ``shard_id`` label is +folded into the factory's output. +""" + +from __future__ import annotations + +from datetime import datetime + +import pandas as pd +import pytest + +from promforecast import sharding +from promforecast.config import ( + Config, + DatasourceConfig, + DefaultsConfig, + GroupConfig, + QueryConfig, + SafetyConfig, + ServerConfig, +) +from promforecast.exporter import Exporter +from promforecast.point_factory import SeriesPointFactory +from promforecast.runner import Runner +from promforecast.source import PromSource, SeriesFrame + + +class _EmptySource(PromSource): + def __init__(self) -> None: + self.calls: list[str] = [] + + async def query_range( + self, promql: str, start: datetime, end: datetime, step_seconds: int + ) -> list[SeriesFrame]: + self.calls.append(promql) + return [] + + async def aclose(self) -> None: + return None + + +def _multi_group_config(group_names: list[str]) -> Config: + return Config( + datasource=DatasourceConfig(url="http://x"), + server=ServerConfig(), + safety=SafetyConfig(), + defaults=DefaultsConfig(min_points=100, models=["SeasonalNaive"]), + groups=[ + GroupConfig(name=n, queries=[QueryConfig(id=f"m_{n}", promql="up")]) + for n in group_names + ], + ) + + +def test_point_factory_folds_shard_id_into_every_family() -> None: + forecast = pd.DataFrame({"yhat": [1.0, 2.0]}) + factory = SeriesPointFactory( + metric_id="m", + base_labels={"instance": "host-a", "__name__": "raw"}, + group="g", + levels=[80], + shard_id="shard-0", + ) + points = list(factory.forecast_points(forecast, "SeasonalNaive", None)) + assert points and all(p.labels["shard_id"] == "shard-0" for p in points) + # __name__ stripped, source labels preserved. + assert points[0].labels["instance"] == "host-a" + assert "__name__" not in points[0].labels + + +def test_point_factory_omits_shard_id_when_unset() -> None: + forecast = pd.DataFrame({"yhat": [1.0]}) + factory = SeriesPointFactory( + metric_id="m", base_labels={"instance": "a"}, group="g", levels=[80] + ) + points = list(factory.forecast_points(forecast, "SeasonalNaive", None)) + assert points and "shard_id" not in points[0].labels + + +@pytest.mark.asyncio +async def test_start_schedules_only_owned_groups() -> None: + names = [f"grp{i}" for i in range(12)] + cfg = _multi_group_config(names) + coord = sharding.ConsistentHashCoordinator(shard_index=0, replica_count=3, shard_id="s0") + exporter = Exporter() + runner = Runner(config=cfg, source=_EmptySource(), exporter=exporter) + runner.set_shard_coordinator(coord) + try: + await runner.start() + owned = coord.owned_groups(names) + assert owned # shard 0 owns at least one of 12 groups over 3 shards + assert set(runner._tasks) == owned + # The ownership gauge mirrors what was scheduled. + body = exporter.render().decode() + for name in owned: + assert f'forecast_shard_owns_group{{group="{name}",shard_id="s0"}} 1.0' in body + for name in set(names) - owned: + assert f'group="{name}"' not in body + finally: + await runner.stop() + + +class _MutableCoordinator: + """Coordinator whose owned set can be flipped to drive reconciliation.""" + + def __init__(self, owned: set[str]) -> None: + self._owned = owned + + @property + def shard_id(self) -> str: + return "s0" + + def owns(self, group_name: str) -> bool: + return group_name in self._owned + + def owned_groups(self, group_names: object) -> set[str]: + return {n for n in group_names if n in self._owned} # type: ignore[union-attr] + + def needs_periodic_refresh(self) -> bool: + return True + + async def refresh(self, group_names: object) -> None: + return None + + async def close(self) -> None: + return None + + +@pytest.mark.asyncio +async def test_reconcile_starts_and_drains_on_ownership_change() -> None: + names = ["a", "b", "c"] + cfg = _multi_group_config(names) + coord = _MutableCoordinator(owned={"a", "b"}) + exporter = Exporter() + runner = Runner(config=cfg, source=_EmptySource(), exporter=exporter) + runner.set_shard_coordinator(coord) + try: + await runner.start() + assert set(runner._tasks) == {"a", "b"} + # Seed per-(group, query) tracker state for the soon-to-be-relinquished + # group "a" and the retained group "b". + runner._drift_windows[("a", "m_a", "k", "SeasonalNaive")] = object() # type: ignore[assignment] + runner._drift_windows[("b", "m_b", "k", "SeasonalNaive")] = object() # type: ignore[assignment] + # Ownership shifts: lose "a", gain "c". + coord._owned = {"b", "c"} + await runner.reconcile_shard_ownership() + assert set(runner._tasks) == {"b", "c"} + body = exporter.render().decode() + assert 'forecast_shard_owns_group{group="c",shard_id="s0"} 1.0' in body + assert 'group="a"' not in body + # Relinquished group's tracker state pruned; retained group's kept. + drift_groups = {k[0] for k in runner._drift_windows} + assert "a" not in drift_groups + assert "b" in drift_groups + finally: + await runner.stop() + + +@pytest.mark.asyncio +async def test_unsharded_runner_owns_every_group() -> None: + names = ["a", "b", "c"] + cfg = _multi_group_config(names) + runner = Runner(config=cfg, source=_EmptySource(), exporter=Exporter()) + try: + await runner.start() + assert set(runner._tasks) == set(names) + finally: + await runner.stop() diff --git a/forecaster/tests/test_shard_controller.py b/forecaster/tests/test_shard_controller.py new file mode 100644 index 0000000..f04ab86 --- /dev/null +++ b/forecaster/tests/test_shard_controller.py @@ -0,0 +1,186 @@ +"""Unit tests for ``main._ShardController`` orchestration. + +The controller is the assembled wiring that sharding's runtime depends on: +it refreshes ownership before starting the runner, runs the assigner elector +(``leader_assigned``), and drives the periodic refresh→reconcile loop. These +tests exercise that orchestration with fakes — no Kubernetes, no Redis — so +they run in the normal suite. The real-cluster elector path is covered +separately in ``test_leader_assigned_kind`` (skipped without a cluster). +""" + +from __future__ import annotations + +import asyncio +from collections.abc import Iterable +from datetime import datetime + +from promforecast.config import ( + Config, + DatasourceConfig, + DefaultsConfig, + GroupConfig, + QueryConfig, + SafetyConfig, + ServerConfig, +) +from promforecast.exporter import Exporter +from promforecast.main import _ShardController +from promforecast.runner import Runner +from promforecast.source import PromSource, SeriesFrame + + +class _EmptySource(PromSource): + def __init__(self) -> None: + self.calls: list[str] = [] + + async def query_range( + self, promql: str, start: datetime, end: datetime, step_seconds: int + ) -> list[SeriesFrame]: + self.calls.append(promql) + return [] + + async def aclose(self) -> None: + return None + + +def _config(group_names: list[str]) -> Config: + return Config( + datasource=DatasourceConfig(url="http://x"), + server=ServerConfig(), + safety=SafetyConfig(), + defaults=DefaultsConfig(min_points=100, models=["SeasonalNaive"]), + groups=[ + GroupConfig(name=n, queries=[QueryConfig(id=f"m_{n}", promql="up")]) + for n in group_names + ], + ) + + +class _FakeCoordinator: + """Coordinator whose owned set can flip; counts refresh() calls.""" + + def __init__(self, owned: set[str], *, periodic: bool) -> None: + self.owned = owned + self._periodic = periodic + self.refreshed = 0 + + @property + def shard_id(self) -> str: + return "s0" + + def owns(self, group_name: str) -> bool: + return group_name in self.owned + + def owned_groups(self, group_names: Iterable[str]) -> set[str]: + return {n for n in group_names if n in self.owned} + + def needs_periodic_refresh(self) -> bool: + return self._periodic + + async def refresh(self, group_names: Iterable[str]) -> None: + self.refreshed += 1 + + async def close(self) -> None: + return None + + +class _FakeElector: + """Stand-in elector: run() blocks until stop() like the real one.""" + + def __init__(self, *, leader: bool = True) -> None: + self._leader = leader + self._stop = asyncio.Event() + self.ran = False + self.stopped = False + + async def run(self) -> None: + self.ran = True + await self._stop.wait() + + async def stop(self) -> None: + self.stopped = True + self._stop.set() + + def is_leader(self) -> bool: + return self._leader + + +async def _wait_for(predicate, timeout: float = 2.0) -> bool: # type: ignore[no-untyped-def] + deadline = asyncio.get_event_loop().time() + timeout + while asyncio.get_event_loop().time() < deadline: + if predicate(): + return True + await asyncio.sleep(0.02) + return predicate() + + +async def test_controller_starts_runner_with_owned_slice_and_no_loop_when_static() -> None: + # consistent_hash-style: no periodic refresh, just refresh-once + start. + cfg = _config(["a", "b", "c"]) + coord = _FakeCoordinator(owned={"a", "c"}, periodic=False) + runner = Runner(config=cfg, source=_EmptySource(), exporter=Exporter()) + runner.set_shard_coordinator(coord) # type: ignore[arg-type] + controller = _ShardController( + coordinator=coord, # type: ignore[arg-type] + runner=runner, + elector=None, + refresh_interval_seconds=0.05, + ) + try: + await controller.start() + assert coord.refreshed == 1 # refreshed once before runner start + assert set(runner._tasks) == {"a", "c"} # only owned groups scheduled + # No periodic loop for a static coordinator. + assert controller._refresh_task is None + finally: + await controller.stop() + await runner.stop() + + +async def test_controller_runs_elector_and_periodic_reconcile() -> None: + cfg = _config(["a", "b", "c"]) + coord = _FakeCoordinator(owned={"a", "b"}, periodic=True) + elector = _FakeElector(leader=True) + runner = Runner(config=cfg, source=_EmptySource(), exporter=Exporter()) + runner.set_shard_coordinator(coord) # type: ignore[arg-type] + controller = _ShardController( + coordinator=coord, # type: ignore[arg-type] + runner=runner, + elector=elector, # type: ignore[arg-type] + refresh_interval_seconds=0.05, + ) + try: + await controller.start() + assert await _wait_for(lambda: elector.ran) # the assigner elector task launched + assert set(runner._tasks) == {"a", "b"} + assert controller._refresh_task is not None # periodic loop running + + # Ownership rebalances (a replica joined): drop "a", gain "c". + coord.owned = {"b", "c"} + converged = await _wait_for(lambda: set(runner._tasks) == {"b", "c"}) + assert converged, f"reconcile did not converge: {set(runner._tasks)}" + # The periodic loop kept calling refresh(). + assert coord.refreshed >= 2 + finally: + await controller.stop() + await runner.stop() + # stop() tore down the elector and the refresh loop cleanly. + assert elector.stopped + assert controller._refresh_task is not None and controller._refresh_task.done() + + +async def test_controller_stop_is_idempotent_without_elector() -> None: + cfg = _config(["a"]) + coord = _FakeCoordinator(owned={"a"}, periodic=True) + runner = Runner(config=cfg, source=_EmptySource(), exporter=Exporter()) + runner.set_shard_coordinator(coord) # type: ignore[arg-type] + controller = _ShardController( + coordinator=coord, # type: ignore[arg-type] + runner=runner, + elector=None, + refresh_interval_seconds=0.05, + ) + await controller.start() + await controller.stop() + await controller.stop() # second stop must not raise + await runner.stop() diff --git a/forecaster/tests/test_sharding.py b/forecaster/tests/test_sharding.py new file mode 100644 index 0000000..1d8c048 --- /dev/null +++ b/forecaster/tests/test_sharding.py @@ -0,0 +1,279 @@ +"""Unit tests for the sharding ownership logic and coordinators. + +The pure functions (``stable_bucket``, ``consistent_hash_owns``, +``assign_groups``) carry the correctness-critical invariants: ownership must be +disjoint and complete, and stable across processes. The coordinators wrap those +in the ``owns`` / ``refresh`` surface the runner consumes; ``LeaderAssigned`` is +exercised against an in-memory store so no Redis is needed. +""" + +from __future__ import annotations + +import pytest + +from promforecast import sharding +from promforecast.exporter import Exporter + + +def test_stable_bucket_is_deterministic() -> None: + # Must NOT depend on PYTHONHASHSEED — two processes have to agree. + assert sharding.stable_bucket("node_capacity") == sharding.stable_bucket("node_capacity") + assert sharding.stable_bucket("a") != sharding.stable_bucket("b") + + +@pytest.mark.parametrize("replicas", [1, 2, 3, 4, 7]) +def test_consistent_hash_ownership_is_disjoint_and_complete(replicas: int) -> None: + groups = [f"group_{i}" for i in range(50)] + owners_per_group: dict[str, list[int]] = {g: [] for g in groups} + for shard in range(replicas): + for g in groups: + if sharding.consistent_hash_owns(g, shard, replicas): + owners_per_group[g].append(shard) + # Every group owned by exactly one shard. + assert all(len(owners) == 1 for owners in owners_per_group.values()) + + +def test_consistent_hash_owns_rejects_out_of_range_index() -> None: + assert sharding.consistent_hash_owns("g", shard_index=5, replica_count=3) is False + assert sharding.consistent_hash_owns("g", shard_index=-1, replica_count=3) is False + + +def test_consistent_hash_owns_clamps_zero_replica_count() -> None: + # A misconfigured count of 0 must not divide by zero; shard 0 owns all. + assert sharding.consistent_hash_owns("g", shard_index=0, replica_count=0) is True + + +def test_assign_groups_disjoint_complete_and_balanced() -> None: + groups = [f"g{i}" for i in range(60)] + members = ["shard-a", "shard-b", "shard-c"] + mapping = sharding.assign_groups(groups, members) + assert set(mapping) == set(groups) + assert set(mapping.values()) <= set(members) + # Rendezvous hashing is balanced in expectation — no shard starves with + # 60 groups over 3 members. + counts = {m: sum(1 for v in mapping.values() if v == m) for m in members} + assert all(c > 0 for c in counts.values()) + + +def test_assign_groups_empty_without_members() -> None: + assert sharding.assign_groups(["g1", "g2"], []) == {} + + +def test_assign_groups_minimal_churn_on_member_join() -> None: + # Rendezvous hashing's defining property: adding a member moves only the + # groups that land on the newcomer, never reshuffles the rest. + groups = [f"g{i}" for i in range(100)] + before = sharding.assign_groups(groups, ["a", "b", "c"]) + after = sharding.assign_groups(groups, ["a", "b", "c", "d"]) + moved = [g for g in groups if before[g] != after[g]] + # Everything that moved must have moved TO the new member. + assert all(after[g] == "d" for g in moved) + # And the churn is a minority of groups. + assert len(moved) < len(groups) // 2 + + +def test_derive_shard_index_precedence(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("PROMFORECAST_SHARD_INDEX", raising=False) + monkeypatch.delenv("PROMFORECAST_POD_NAME", raising=False) + # Explicit config value wins. + assert sharding.derive_shard_index(2) == 2 + # Then the env override. + monkeypatch.setenv("PROMFORECAST_SHARD_INDEX", "5") + assert sharding.derive_shard_index(-1) == 5 + monkeypatch.delenv("PROMFORECAST_SHARD_INDEX") + # Then the StatefulSet pod ordinal. + monkeypatch.setenv("PROMFORECAST_POD_NAME", "promforecast-3") + assert sharding.derive_shard_index(-1) == 3 + monkeypatch.delenv("PROMFORECAST_POD_NAME") + # Nothing resolvable → -1 (caller fails loudly). + assert sharding.derive_shard_index(-1) == -1 + + +def test_derive_shard_id_precedence(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("PROMFORECAST_POD_NAME", raising=False) + assert sharding.derive_shard_id("custom", 0) == "custom" + monkeypatch.setenv("PROMFORECAST_POD_NAME", "promforecast-2") + assert sharding.derive_shard_id("", 2) == "promforecast-2" + monkeypatch.delenv("PROMFORECAST_POD_NAME") + assert sharding.derive_shard_id("", 2) == "shard-2" + + +def test_consistent_hash_coordinator_owns_and_owned_groups() -> None: + groups = [f"g{i}" for i in range(20)] + coords = [ + sharding.ConsistentHashCoordinator(shard_index=i, replica_count=3, shard_id=f"s{i}") + for i in range(3) + ] + assert not any(c.needs_periodic_refresh() for c in coords) + # Union of owned groups == all groups; pairwise disjoint. + owned_sets = [c.owned_groups(groups) for c in coords] + union: set[str] = set().union(*owned_sets) + assert union == set(groups) + assert sum(len(s) for s in owned_sets) == len(groups) + + +class _FakeStore: + """In-memory :class:`AssignmentStore` for coordinator tests.""" + + def __init__(self) -> None: + self.members_set: set[str] = set() + self.assignment: dict[str, str] = {} + + async def heartbeat(self, shard_id: str, ttl_seconds: int) -> None: + self.members_set.add(shard_id) + + async def members(self) -> list[str]: + return sorted(self.members_set) + + async def publish_assignment(self, mapping: dict[str, str]) -> None: + self.assignment = dict(mapping) + + async def read_assignment(self) -> dict[str, str]: + return dict(self.assignment) + + +@pytest.mark.asyncio +async def test_leader_assigned_coordinator_leader_publishes() -> None: + store = _FakeStore() + # Pretend two peers already heartbeated. + store.members_set.update({"shard-b", "shard-c"}) + coord = sharding.LeaderAssignedCoordinator( + shard_id="shard-a", + store=store, + is_leader=lambda: True, + member_ttl_seconds=30, + ) + groups = [f"g{i}" for i in range(30)] + await coord.refresh(groups) + # The leader published a complete assignment over the live membership. + assert set(store.assignment) == set(groups) + assert set(store.assignment.values()) <= {"shard-a", "shard-b", "shard-c"} + # Its own owned set is the subset assigned to it. + owned = coord.owned_groups(groups) + assert owned == {g for g, s in store.assignment.items() if s == "shard-a"} + + +@pytest.mark.asyncio +async def test_leader_assigned_follower_reads_published_map() -> None: + store = _FakeStore() + store.assignment = {"g1": "shard-a", "g2": "shard-b"} + coord = sharding.LeaderAssignedCoordinator( + shard_id="shard-b", + store=store, + is_leader=lambda: False, + member_ttl_seconds=30, + ) + await coord.refresh(["g1", "g2"]) + assert coord.owns("g2") is True + assert coord.owns("g1") is False + + +@pytest.mark.asyncio +async def test_leader_assigned_owns_nothing_before_assignment_published() -> None: + # Cold start: no published assignment → own nothing and wait for the + # assigner, rather than every replica fitting every group (the N-times herd). + store = _FakeStore() + coord_a = sharding.LeaderAssignedCoordinator( + shard_id="a", store=store, is_leader=lambda: False, member_ttl_seconds=30 + ) + assert coord_a.owns("anything") is False + + +@pytest.mark.asyncio +async def test_leader_assigned_covers_group_missing_from_published_map() -> None: + # A reload added "new_group" before the assigner republished: the mapping + # exists but omits it. Every replica must agree (via rendezvous) on exactly + # one owner so the group is covered, not orphaned and not duplicated. + store = _FakeStore() + members = ["a", "b", "c"] + coords = { + m: sharding.LeaderAssignedCoordinator( + shard_id=m, store=store, is_leader=lambda: False, member_ttl_seconds=30 + ) + for m in members + } + # Simulate the published-but-stale state each replica has cached. + for c in coords.values(): + c._assignment = {"g1": "a", "g2": "b"} + c._members = list(members) + owners = [m for m, c in coords.items() if c.owns("new_group")] + assert len(owners) == 1 # exactly one replica covers the un-mapped group + + +class _FakeRedis: + """Minimal async Redis stand-in for RedisAssignmentStore. + + Implements ``scan`` (cursor-based, paged) rather than ``keys`` so the test + exercises the same non-blocking enumeration the production store uses. + """ + + def __init__(self) -> None: + self.store: dict[str, bytes] = {} + + async def set(self, key: str, value: object, ex: int | None = None) -> None: + self.store[key] = value if isinstance(value, bytes) else str(value).encode() + + async def get(self, key: str) -> bytes | None: + return self.store.get(key) + + async def scan( + self, cursor: int = 0, match: str = "*", count: int = 100 + ) -> tuple[int, list[bytes]]: + prefix = match.rstrip("*") + matched = [k.encode() for k in sorted(self.store) if k.startswith(prefix)] + # Page the results to prove the store's cursor loop terminates. + page = matched[cursor : cursor + count] + next_cursor = cursor + count + if next_cursor >= len(matched): + next_cursor = 0 + return next_cursor, page + + +@pytest.mark.asyncio +async def test_redis_assignment_store_roundtrip() -> None: + client = _FakeRedis() + store = sharding.RedisAssignmentStore(client) + await store.heartbeat("shard-a", ttl_seconds=30) + await store.heartbeat("shard-b", ttl_seconds=30) + assert sorted(await store.members()) == ["shard-a", "shard-b"] + await store.publish_assignment({"g1": "shard-a"}) + assert await store.read_assignment() == {"g1": "shard-a"} + + +@pytest.mark.asyncio +async def test_redis_assignment_store_empty_assignment() -> None: + store = sharding.RedisAssignmentStore(_FakeRedis()) + assert await store.read_assignment() == {} + + +@pytest.mark.asyncio +async def test_redis_assignment_store_members_scan_pages() -> None: + # More members than one SCAN page (count=100) → the cursor loop must walk + # every page and terminate. Guards against a regression to a single KEYS + # call or a cursor loop that never advances. + client = _FakeRedis() + store = sharding.RedisAssignmentStore(client) + for i in range(250): + await store.heartbeat(f"shard-{i:03d}", ttl_seconds=30) + members = await store.members() + assert len(members) == 250 + assert "shard-000" in members and "shard-249" in members + + +def test_exporter_shard_owns_group_gauge() -> None: + exporter = Exporter() + # Absent until set. + assert "forecast_shard_owns_group" not in exporter.render().decode() + exporter.set_shard_state(shard_id="shard-1", owned_groups={"node_cap", "kube_res"}) + body = exporter.render().decode() + assert 'forecast_shard_owns_group{group="node_cap",shard_id="shard-1"} 1.0' in body + assert 'forecast_shard_owns_group{group="kube_res",shard_id="shard-1"} 1.0' in body + + +def test_exporter_shard_state_empty_owned_renders_no_samples() -> None: + exporter = Exporter() + exporter.set_shard_state(shard_id="shard-1", owned_groups=set()) + # Shard id set but nothing owned → the family may appear with no samples; + # crucially no spurious empty-group series. + body = exporter.render().decode() + assert 'group=""' not in body