assignment

package
v2.10.2 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Jul 17, 2026 License: Apache-2.0 Imports: 24 Imported by: 0

Documentation

Overview

Package assignment provides partition assignment calculation and distribution.

The assignment package implements the leader-based assignment coordination system. The Calculator runs only on the leader worker and is responsible for:

  • Monitoring worker health via heartbeat tracking
  • Calculating partition assignments using configured strategies
  • Publishing assignments to NATS KV for worker discovery
  • Handling rebalancing on worker join/leave events
  • Applying rebalance cooldowns to prevent thrashing

Design Overview

The assignment system uses a leader-based approach where one worker (the leader) is responsible for calculating and distributing assignments to all workers:

  1. Leader monitors worker heartbeats in NATS KV
  2. Leader detects worker join/leave by tracking heartbeat changes
  3. Leader calculates new assignments using AssignmentStrategy
  4. Leader publishes assignments via the refs-always commit protocol (three protocol keys: _commit, _commit_log.<V>, _payload.<hex>)
  5. Workers watch assignment._commit and fetch their payload by key

Calculator Lifecycle

The Calculator should only be started on the leader worker:

  1. Create calculator with NewCalculator(&Config{...})
  2. Start calculator with Start(ctx) (performs initial assignment)
  3. Calculator monitors workers in background
  4. Stop calculator with Stop() when stepping down from leadership

Example: 

// Leader worker creates and starts calculator
calc, err := assignment.NewCalculator(&assignment.Config{
    AssignmentKV:         assignmentKV,
    HeartbeatKV:          heartbeatKV,
    AssignmentPrefix:     "assignment",
    Source:               partitionSource,
    Strategy:             assignmentStrategy,
    HeartbeatPrefix:      "worker-hb",
    HeartbeatTTL:         15 * time.Second,
    EmergencyGracePeriod: 7500 * time.Millisecond,
    Cooldown:             10 * time.Second,
    ColdStartWindow:      30 * time.Second,
    PlannedScaleWindow:   10 * time.Second,
})
if err != nil {
    log.Fatal(err)
}

// Start assignment calculation
err = calc.Start(ctx)
if err != nil {
    log.Fatal(err)
}
defer calc.Stop(ctx)

// Calculator now monitors workers and triggers rebalancing

Rebalancing Strategy

The calculator uses stabilization windows to balance responsiveness and stability. The window is chosen structurally by detectRebalanceType, not by a ratio:

  • Cold start: no previously-known workers (the first non-empty scan)
  • Cold start window: 30 seconds (wait for most workers to start)
  • Planned scale: any change once workers are already known
  • Planned scale window: 10 seconds (quick response to membership changes)
  • Rebalance cooldown: 10 seconds (prevent excessive rebalancing)

Assignment Distribution

Assignments are published to NATS KV using a three-key commit-publisher protocol (refs-always model):

  • "{prefix}._commit" — current commit object mapping each worker ID to the content-addressable payload key that holds its assignment slice.
  • "{prefix}._commit_log.<V>" — append-only commit-log entry per assignment version V; used by the GC to determine which payload keys are live.
  • "{prefix}._payload.<hex(sha256)>" — content-addressable assignment blobs. Identical assignment slices for different workers share a single payload key.

Workers watch "{prefix}._commit" and fetch the payload keys referenced by their entry. A background GC pass (CommitGC) reaps payload keys that are not referenced by any recent commit-log entry.

Worker Health Monitoring

The calculator monitors worker health by:

  • Listing all heartbeat keys (format: "{hbPrefix}.{workerID}")
  • Extracting worker IDs from heartbeat keys
  • Detecting changes in the worker set (joins/leaves)
  • Triggering rebalancing when worker set changes

Workers are considered active if their heartbeat key exists in NATS KV. The heartbeat TTL ensures dead workers are automatically removed.

Thread Safety

The Calculator is thread-safe and can be accessed concurrently from multiple goroutines. All public methods use proper synchronization to protect internal state.

Configuration Options

The calculator supports several configuration options:

  • Cooldown: Minimum time between rebalances (default: 10s)
  • ColdStartWindow: Stabilization time for cold starts (default: 30s)
  • PlannedScaleWindow: Stabilization time for planned scaling (default: 10s)

Integration with Manager

The Calculator is integrated into the Manager's leader election logic:

  1. Worker wins election and becomes leader
  2. Manager creates and starts Calculator
  3. Calculator performs initial assignment
  4. Calculator monitors workers and handles rebalancing
  5. Worker loses election and steps down
  6. Manager stops Calculator

Only the leader worker runs the Calculator. Follower workers watch assignment._commit in NATS KV and fetch their payload by key on each update.

Index

Constants

View Source 
const (
	// DefaultPayloadGCInterval is the default cadence at which a CommitGC
	// instance triggers a sweep when run via Start.
	DefaultPayloadGCInterval = 5 * time.Minute

	// DefaultPayloadGCRetention is the default time window during which a
	// payload key is retained even if it is unreferenced. Keys created less
	// than this duration ago are never deleted, providing forensic margin.
	DefaultPayloadGCRetention = 24 * time.Hour

	// DefaultPayloadGCKeepCommits is the default count of recent commit logs
	// scanned to compute the live payload set. Defaults to 10.
	DefaultPayloadGCKeepCommits = 10
)

GC defaults per §3.9 of the partition assignment robustness plan.

Variables

View Source 
var (
	// ErrCommitPayloadFetch indicates the payload key referenced by the
	// commit could not be fetched from KV.
	ErrCommitPayloadFetch = errors.New("commit payload fetch failed")

	// ErrCommitPayloadDecompress indicates the payload bytes failed gzip
	// decompression.
	ErrCommitPayloadDecompress = errors.New("commit payload gzip decompression failed")

	// ErrCommitPayloadHashMismatch indicates the sha256 of the
	// decompressed bytes did not match ref.PayloadHash. This is either a
	// hash collision (extraordinary) or KV corruption.
	ErrCommitPayloadHashMismatch = errors.New("commit payload hash mismatch")

	// ErrCommitPayloadDecode indicates the canonical bytes failed JSON
	// decoding into AssignmentPayload.
	ErrCommitPayloadDecode = errors.New("commit payload JSON decode failed")

	// ErrCommitPayloadDigestMismatch indicates the set digest computed
	// over the decoded partitions did not match ref.SetDigest. The
	// PayloadHash already matched, so this is an internal-consistency
	// failure rather than a transport corruption.
	ErrCommitPayloadDigestMismatch = errors.New("commit payload set-digest mismatch")
)

Typed sentinels for the worker-side commit-path payload verification (§3.6 case (c)). Callers should errors.Is-switch on these to pick stage-specific metrics.

Functions

func FetchAndVerifyCommitPayload added in v2.4.0 

func FetchAndVerifyCommitPayload(
	ctx context.Context,
	kv jetstream.KeyValue,
	ref types.AssignmentPayloadRef,
) (types.AssignmentPayload, error)

FetchAndVerifyCommitPayload fetches the payload key referenced by ref from kv, gzip-decompresses, sha256-verifies the bytes against ref.PayloadHash, JSON-decodes into AssignmentPayload, and verifies the computed set digest against ref.SetDigest.

On success, returns the decoded payload. On failure, returns a wrapped typed sentinel from this package — callers should errors.Is-switch on the sentinels to pick stage-specific metrics. After a failure the caller should leave lastSeenLeaderRevision and any pending state unchanged (case (c) retry-on-next-tick semantics).

Parameters:

  • ctx: Context for the KV Get operation.
  • kv: Assignment KV bucket carrying the payload key.
  • ref: AssignmentPayloadRef from a commit's Payloads map.

Returns:

  • types.AssignmentPayload: Decoded and verified payload on success.
  • error: Wrapped typed sentinel on any verification failure.

Types

type AssignmentPublisher 

type AssignmentPublisher struct {
	// contains filtered or unexported fields
}

AssignmentPublisher publishes partition assignments using the refs-always commit model.

On every Publish, the publisher:

  1. Writes content-addressable AssignmentPayload keys ("assignment._payload.<hex(sha256)>") via kv.Create with hash-verify on ErrKeyExists (immutable, dedupe-safe). An adopted (reused) key is CAS-touched to a fresh revision before the adoption is treated as final, fencing it against a racing GC delete (see createOrAdoptPayload).
  2. Reads heartbeats to classify legacy (CapAckV1=0) workers in the batch.
  3. Brackets the legacy alias barrier with two leadership rechecks (pre-alias and post-alias).
  4. Writes mandatory "assignment.<W>" legacy aliases for legacy workers with bounded retries (failure aborts the batch).
  5. CAS-writes the AssignmentCommit to "assignment._commit" — this is the single atomic decision point.
  6. Best-effort writes the AssignmentCommitLog and compat-noise legacy aliases for commit-capable workers.
  7. Best-effort triggers GC.

The publisher is intentionally agnostic to the strategy and source layout; the calculator computes the assignment map and source snapshot and passes them in via Publish.

func NewAssignmentPublisher 

func NewAssignmentPublisher(cfg PublisherConfig) *AssignmentPublisher

NewAssignmentPublisher creates a publisher wired to the assignment KV bucket.

LeaderCheckFn and HeartbeatKV may be nil only in narrow test contexts; production calls MUST supply both. When LeaderCheckFn is nil, the pre-alias and post-alias leadership rechecks become no-ops (every Publish is treated as if leadership is held). When HeartbeatKV is nil, the legacy alias barrier degrades into "treat all workers as commit-capable" — the publisher will not write mandatory aliases and the cluster cannot safely host pre-CapAckV1 workers.

func (*AssignmentPublisher) AssignmentKV added in v2.4.0 

func (p *AssignmentPublisher) AssignmentKV() jetstream.KeyValue

AssignmentKV returns the KV bucket the publisher writes into. Exposed so the GC loop and tests can list and inspect protocol keys without re-opening the bucket.

func (*AssignmentPublisher) BootstrapLastCommit added in v2.4.0 

func (p *AssignmentPublisher) BootstrapLastCommit(ctx context.Context) error

BootstrapLastCommit seeds the publisher's in-memory lastCommit cache from the live "<prefix>._commit" KV entry, if one exists. Called once at calculator startup so the audit loop can run from t=0 without waiting for the publisher to issue its own commit.

Best-effort: KV access failures, missing entries, and unmarshal errors all leave lastCommit nil and return nil. Errors are logged for ops visibility but do not propagate.

Parameters:

  • ctx: Context for the KV Get operation.

Returns:

  • error: Always nil; the bootstrap path is intentionally non-fatal.

func (*AssignmentPublisher) CleanupAllAssignments 

func (p *AssignmentPublisher) CleanupAllAssignments(ctx context.Context) error

CleanupAllAssignments removes ALL legacy per-worker assignment aliases from KV, including aliases for workers that may still be active in the cluster.

Protocol keys (assignment._commit, assignment._commit_log.*, assignment._payload.*) are NEVER deleted by this method — they are the authoritative cluster state, and a successor leader needs them to continue the CAS chain and to GC payloads safely.

Intended use: admin tooling / whole-cluster teardown where every parti instance is being decommissioned together. NOT safe to call on leader-step-down or single-node shutdown: it will yank aliases for workers that are still serving traffic on peer leaders, which would then have to wait for the next rebalance to be re-published. There is currently no production caller; the method exists for tests and operator scripts.

A leader-step-down-safe variant would need a separate CleanupInactiveAssignments(activeWorkers []string) method that preserves aliases of currently-live workers — out of scope here.

func (*AssignmentPublisher) CurrentVersion 

func (p *AssignmentPublisher) CurrentVersion() int64

CurrentVersion returns the highest assignment version this publisher has observed (either via DiscoverHighestVersion or via a successful commit).

This method is safe to call concurrently.

func (*AssignmentPublisher) DiscoverHighestVersion 

func (p *AssignmentPublisher) DiscoverHighestVersion(ctx context.Context) ([]string, error)

DiscoverHighestVersion scans the assignment KV for the highest legacy assignment version and returns the set of legacy worker IDs that still have alias keys.

Protocol keys (any sub-component starting with "_": _commit, _commit_log.*, _payload.*) are filtered out and never interpreted as worker IDs. This is the bootstrap path used by §3.7 "new-leader recovery" before the first commit lands; once a commit exists, the publisher seeds currentVersion from it instead.

Returns:

  • workerIDs: sorted list of legacy worker IDs found
  • error: KV access failure (non-fatal in many bootstrap paths)

func (*AssignmentPublisher) LastCommit added in v2.4.0 

func (p *AssignmentPublisher) LastCommit() *types.AssignmentCommit

LastCommit returns a defensive copy of the most recently observed AssignmentCommit, populated either by a successful CAS or by BootstrapLastCommit.

Returns nil when no commit has been observed yet (pre-first-commit bootstrap path against an empty assignment bucket).

The returned struct is a deep copy: callers may iterate or mutate Workers and Payloads without affecting publisher state.

Safe to call concurrently with Publish.

Returns:

  • *types.AssignmentCommit: Defensive copy, or nil if no commit observed.

func (*AssignmentPublisher) LastCommitObservedAt added in v2.4.0 

func (p *AssignmentPublisher) LastCommitObservedAt() time.Time

LastCommitObservedAt returns the local monotonic-clock instant at which this publisher observed the current lastCommit. Returns the zero time when lastCommit is nil.

The leader-side audit uses this for grace-window math instead of the wire-clock commit.PublishedAt so cross-leader wall-clock skew (e.g., after takeover or container suspend/resume) cannot suppress or prematurely trigger audit_repair escalation.

Callers use time.Since on the return value; subtracting wall-clock times against this value is meaningless and will silently lose the monotonic reading.

func (*AssignmentPublisher) LastCommitRev added in v2.4.0 

func (p *AssignmentPublisher) LastCommitRev() uint64

LastCommitRev returns the KV revision of the most recently CAS-written assignment._commit. Used by the GC loop to identify the live commit.

This method is safe to call concurrently.

func (*AssignmentPublisher) LastRebalanceTime 

func (p *AssignmentPublisher) LastRebalanceTime() time.Time

LastRebalanceTime returns the time of the last successful commit.

This method is safe to call concurrently.

func (*AssignmentPublisher) LiveRefs added in v2.4.0 

func (p *AssignmentPublisher) LiveRefs() []string

LiveRefs returns a snapshot of the payload keys this publisher has selected for an in-progress publish (after step 4 verify-back, before step 9 commit CAS returns). The GC consults this set to avoid deleting a payload that an in-flight commit is about to reference (§3.9 / P0-2).

The returned slice is a point-in-time copy: callers may iterate without holding any lock. An empty slice is a normal steady-state value when no publish is in flight.

This method is safe to call concurrently with Publish.

func (*AssignmentPublisher) Prefix added in v2.4.0 

func (p *AssignmentPublisher) Prefix() string

Prefix returns the assignment key prefix. Exposed for the GC loop.

func (*AssignmentPublisher) Publish 

func (p *AssignmentPublisher) Publish(ctx context.Context, in PublishInput) error

Publish runs the 12-step refs-always commit flow described in §3.5 of the partition-assignment robustness plan.

The flow is:

  1. Verify set-equality coverage of Assignments against SourcePartitions.
  2. Build per-worker AssignmentPayloads, Create their content-addressable keys, and on ErrKeyExists verify-back via sha256 of the canonical uncompressed bytes.
  3. Pre-alias leadership recheck.
  4. Heartbeat-aware legacy alias barrier (mandatory pre-commit aliases for CapAckV1=0 workers in the batch, with bounded retries).
  5. Post-alias leadership recheck.
  6. CAS-write AssignmentCommit (THE commit point).
  7. Best-effort: write AssignmentCommitLog, write compat-noise aliases for commit-capable workers, sweep WorkersToRemove, schedule GC if wired.

Returns:

  • error: a typed sentinel from types/errors.go (ErrCoverageMismatch, ErrLeadershipLostPreAlias, ErrAliasBarrierFailed, ErrLeadershipLostPostAlias, ErrCommitCASFailed, ErrPayloadHashCollisionOrCorruption) wrapped with context where useful.

type Calculator 

type Calculator struct {
	Config
	// contains filtered or unexported fields
}

Calculator manages partition assignment calculation and distribution.

The calculator runs on the leader worker and orchestrates three focused components:

  • WorkerMonitor: Detects worker health changes via NATS KV heartbeats
  • StateMachine: Manages state transitions (Idle, Scaling, Rebalancing, Emergency)
  • AssignmentPublisher: Publishes partition assignments to NATS KV

The calculator handles rebalancing logic and coordinates these components. It does NOT run on follower workers.

func NewCalculator 

func NewCalculator(cfg *Config) (*Calculator, error)

NewCalculator creates a calculator with validated configuration.

This constructor provides clear, self-documenting configuration and validation of required fields.

Parameters:

  • cfg: Calculator configuration (required fields must be set)

Returns:

  • *Calculator: New calculator instance ready to start
  • error: Validation error if required fields are missing

Example: 

calc, err := assignment.NewCalculator(&assignment.Config{
    AssignmentKV:     assignKV,
    HeartbeatKV:      heartbeatKV,
    Source:           source,
    Strategy:         strategy,
    AssignmentPrefix: "assignment",
    HeartbeatPrefix:  "heartbeat",
    HeartbeatTTL:     3 * time.Second,
    // Optional fields use sensible defaults
    Logger:           logger,
})
if err != nil {
    log.Fatal(err)
}

func (*Calculator) CurrentVersion 

func (c *Calculator) CurrentVersion() int64

CurrentVersion returns the current assignment version.

func (*Calculator) GetScalingReason 

func (c *Calculator) GetScalingReason() string

GetScalingReason returns the reason for the current or last scaling operation.

Returns:

  • string: Scaling reason ("cold_start", "planned_scale", "emergency", "restart") or empty string if idle

func (*Calculator) GetState 

func (c *Calculator) GetState() types.CalculatorState

GetState returns the current calculator state.

Returns:

  • types.CalculatorState: Current calculator state (type-safe enum)

func (*Calculator) IsStarted 

func (c *Calculator) IsStarted() bool

IsStarted returns true if the calculator is currently running.

func (*Calculator) LabelSnapshot added in v2.10.1 

func (c *Calculator) LabelSnapshot() (pools, parked map[string]int, ok bool)

LabelSnapshot returns copies of the retained label read model, or ok=false when no rebalance has published since this calculator started. Lock-free; safe concurrently with Stop (the caller may observe either the snapshot or the cleared state, never a partial one).

func (*Calculator) Start 

func (c *Calculator) Start(ctx context.Context) error

Start begins monitoring workers and calculating assignments.

This method should only be called on the leader worker. It:

  1. Discovers highest version from existing assignments
  2. Starts background monitoring for worker changes
  3. Performs initial assignment asynchronously (with stabilization window)
  4. Triggers rebalancing when workers join/leave

The initial assignment runs in a background goroutine, allowing Start() to return immediately without blocking on the stabilization window (10-30 seconds). This enables:

  • Fast manager startup (milliseconds instead of seconds)
  • Concurrent worker initialization across all instances
  • Leader calculates assignment with all workers visible from the start

Parameters:

  • ctx: Context for cancellation and timeout

Returns:

  • error: Start error (e.g., already started, KV operation failed during setup)

Note: Errors during background initial assignment are logged but not returned. Callers should wait for assignment via Manager.waitForAssignment() or similar mechanism.

func (*Calculator) Stop 

func (c *Calculator) Stop(ctx context.Context) error

Stop stops the calculator and waits for background goroutines to finish.

This method performs a clean shutdown sequence:

  1. Signals stop to all components
  2. Cleans up assignments from KV (provides clean slate for new leader)
  3. Stops worker monitor
  4. Waits for state machine shutdown

The assignment cleanup is best-effort and won't fail the Stop() operation. If cleanup fails, the new leader will discover existing versions and maintain version monotonicity via DiscoverHighestVersion().

Parameters:

  • ctx: Context for cleanup timeout control (typically 5s)

Returns:

  • error: Stop error (e.g., not started)

func (*Calculator) SubscribeToStateChanges 

func (c *Calculator) SubscribeToStateChanges() (<-chan types.CalculatorState, func())

SubscribeToStateChanges returns a channel for state updates and a function to unsubscribe.

func (*Calculator) TriggerRebalance 

func (c *Calculator) TriggerRebalance(ctx context.Context) error

TriggerRebalance forces an immediate rebalance, bypassing cooldown.

This is useful when partitions are added/removed dynamically and you want to redistribute them immediately without waiting for the next worker change.

Parameters:

  • ctx: Context for operation timeout

Returns:

  • error: Rebalance error

type CalculatorAndAssignmentMetrics 

type CalculatorAndAssignmentMetrics interface {
	types.CalculatorMetrics
	types.AssignmentMetrics
	types.PublisherMetrics
	types.GCMetrics
	types.AuditMetrics
}

CalculatorAndAssignmentMetrics combines the metrics interfaces needed by the calculator and its sub-components (e.g., AssignmentPublisher and the payload GC loop).

The calculator itself records CalculatorMetrics (rebalance durations, worker changes, etc.) and delegates AssignmentMetrics, PublisherMetrics, and GCMetrics to the AssignmentPublisher / GC.

type CommitGC added in v2.4.0 

type CommitGC struct {
	// contains filtered or unexported fields
}

CommitGC reaps orphan content-addressable payload keys.

GC is conservative and never participates in correctness:

  • A payload key is "live" if it appears in either the current assignment._commit or any of the last KeepCommits commit_log entries.
  • Live keys are never deleted.
  • Non-live keys older than Retention are eligible for deletion. Failures are non-fatal and surface via the IncrementPayloadDeleteErrors metric.

CommitGC is safe to call concurrently with Publisher.Publish, including across processes: GC's delete is conditioned on the payload key's revision (jetstream.LastRevision) as observed immediately before the delete, and a publish adopting an existing key (Create → ErrKeyExists → verify-back) CAS-touches it to advance that same revision before treating the adoption as final (see createOrAdoptPayload in assignment_publisher.go). Whichever side's conditioned write reaches the server first wins deterministically: a GC delete that lands first makes the adopter's touch fail and retry (recreating the key), and an adopter's touch that lands first makes GC's delete fail and the key survives.

Example: 

gc := assignment.NewCommitGC(assignment.CommitGCConfig{
    Publisher: pub,
    Logger:    logger,
    Metrics:   metrics,
})
gc.Start(ctx) // runs in background
defer gc.Stop()

func NewCommitGC added in v2.4.0 

func NewCommitGC(cfg CommitGCConfig) *CommitGC

NewCommitGC constructs a CommitGC, applying defaults to optional fields.

Returns nil if cfg.Publisher is nil — the publisher is the only required dependency.

func (*CommitGC) RunOnce added in v2.4.0 

func (g *CommitGC) RunOnce(ctx context.Context) error

RunOnce performs a single synchronous GC sweep and returns the number of keys deleted (and the count of delete errors observed). Tests use this to drive deterministic GC behavior.

Failures during the sweep are NOT fatal: errors are logged and surfaced via the IncrementPayloadDeleteErrors metric, and RunOnce continues to the next candidate.

func (*CommitGC) Start added in v2.4.0 

func (g *CommitGC) Start(ctx context.Context) error

Start launches the background GC loop. Call Stop to terminate it.

Returns an error if already started.

func (*CommitGC) Stop added in v2.4.0 

func (g *CommitGC) Stop()

Stop terminates the background loop and waits for the in-flight pass (if any) to return. Safe to call multiple times.

func (*CommitGC) Trigger added in v2.4.0 

func (g *CommitGC) Trigger()

Trigger requests an immediate GC sweep. The notification is non-blocking and coalesces — a backlog of triggers collapses into a single sweep, so it is safe to call after every successful publish without throttling.

No-op when the GC has not been Started.

type CommitGCConfig added in v2.4.0 

type CommitGCConfig struct {
	// Publisher provides the KV bucket and key prefix to operate on. The GC
	// only ever reads protocol keys (assignment._commit, _commit_log.<V>,
	// _payload.<hash>) and deletes orphan _payload.<hash> keys.
	Publisher *AssignmentPublisher

	// LiveRefsProvider is consulted on every sweep to obtain the publisher's
	// in-flight payload-ref set, which the GC must never delete (P0-2 / §3.9).
	// When nil, defaults to the Publisher itself; tests may inject a fake.
	LiveRefsProvider LiveRefsProvider

	// Interval is the cadence at which Start triggers a sweep. Defaults to
	// DefaultPayloadGCInterval.
	Interval time.Duration

	// Retention is the minimum age a payload key must reach before it is
	// eligible for deletion (even if otherwise unreferenced). Defaults to
	// DefaultPayloadGCRetention.
	Retention time.Duration

	// KeepCommits is the number of recent commit_log.<V> entries scanned to
	// compute the live payload set. Defaults to DefaultPayloadGCKeepCommits.
	KeepCommits int

	// Logger / Metrics: optional. The publisher's logger/metrics are reused if
	// these are nil.
	Logger  types.Logger
	Metrics types.GCMetrics

	// Now is a clock injection for tests. Defaults to time.Now.
	Now func() time.Time
}

CommitGCConfig configures a CommitGC instance.

All fields are optional except Publisher, which provides the KV bucket and prefix. Zero-valued fields are filled in with the Default* constants.

type Config 

type Config struct {
	// Required dependencies
	AssignmentKV jetstream.KeyValue // NATS KV bucket for assignments
	HeartbeatKV  jetstream.KeyValue // NATS KV bucket for heartbeats
	Source       types.PartitionSource
	Strategy     types.AssignmentStrategy

	// Required configuration
	AssignmentPrefix string        // Key prefix for assignments (e.g., "assignment")
	HeartbeatPrefix  string        // Key prefix for heartbeats (e.g., "heartbeat")
	HeartbeatTTL     time.Duration // Heartbeat TTL for worker health detection

	// Optional configuration (with defaults)
	EmergencyGracePeriod time.Duration // Minimum time before emergency rebalance (default: 5s)
	Cooldown             time.Duration // Minimum time between rebalances (default: 10s)
	ColdStartWindow      time.Duration // Stabilization window for cold start (default: 30s)
	PlannedScaleWindow   time.Duration // Stabilization window for planned scale (default: 10s)

	// Now returns the current time and backs the cooldown gate
	// (time-since-last-rebalance vs Cooldown) and the publisher's
	// lastRebalance stamp. When nil, time.Now is used. Tests inject a
	// controllable clock so cooldown timing is deterministic under load;
	// production leaves it nil.
	Now func() time.Time

	// Partition-input credibility (see
	// errSuspiciousPartitionObservation Godoc in calculator.go):
	//
	// PartitionShrinkConfirmationCount is the number of consecutive
	// suspicious partition-source observations the calculator requires
	// before trusting an empty / sharply-shrunk shape. Default: 3.
	// Setting this to 1 disables the confirmation window (every
	// observation is acted on immediately).
	PartitionShrinkConfirmationCount int

	// PartitionShrinkConfirmationThresholdPct defines "sharply shrunk".
	// A new partition count where
	//   observed * 100 < lastKnownPartitionCount * Pct
	// is suspicious. Default: 50 (a >=50% drop in one poll is
	// suspicious). An empty observation is always suspicious
	// regardless of this threshold.
	PartitionShrinkConfirmationThresholdPct int

	// Worker-input credibility (see errSuspiciousWorkerObservation Godoc
	// in calculator.go):
	//
	// WorkerShrinkConfirmationCount is the number of consecutive
	// suspicious worker-set observations the calculator requires before
	// trusting a sharply-shrunk heartbeat scan. Default: 2. Setting this
	// to 1 disables the confirmation window (every observation is acted
	// on immediately). The default is lower than
	// PartitionShrinkConfirmationCount because heartbeat scans run on
	// every monitor poll (~HeartbeatTTL/2) whereas partition-source
	// observations are watcher-driven and arrive only on source change —
	// a smaller worker window converges faster on a real shrink without
	// stranding correctness.
	WorkerShrinkConfirmationCount int

	// WorkerShrinkConfirmationThresholdPct defines "sharply shrunk" for
	// the heartbeat scan. A new worker count where
	//   observed * 100 < lastKnownWorkerCount * Pct
	// is suspicious. Default: 50 (a >=50% drop in one scan is
	// suspicious). The defense fires only when lastKnownWorkerCount > 0
	// (the first ever scan is always trusted).
	WorkerShrinkConfirmationThresholdPct int

	// RebalanceGraceDrainInterval is the period at which monitorPartitions
	// checks for a partition-source update that was deferred because the
	// leader was in recovery grace. When grace lifts, the deferred update
	// is drained on the next tick. Default: min(Cooldown, 30s), capped
	// below at 1s.
	RebalanceGraceDrainInterval time.Duration

	// ApplyGracePeriod is the time after THIS leader observed the current
	// commit (via successful CAS or BootstrapLastCommit) before the audit
	// loop emits retry-pressure metrics for behind workers. Measured against
	// a monotonic clock so cross-leader wall-clock skew does not influence
	// the grace window. Default: 2 × HeartbeatTTL.
	ApplyGracePeriod time.Duration

	// ExtendedApplyGracePeriod is the time after THIS leader observed the
	// current commit before the audit may escalate via two-phase handoff
	// (audit_repair rebalance). Measured against a monotonic clock; see
	// ApplyGracePeriod. Default: 5 × HeartbeatTTL.
	ExtendedApplyGracePeriod time.Duration

	// AuditInterval is the period between audit passes. Default: HeartbeatTTL.
	AuditInterval time.Duration

	// EnableTwoPhaseHandoff mirrors the manager's flag so the audit can skip
	// escalation when two-phase mode is off. The audit only escalates when
	// this flag is true AND the full safety chain is present on both the
	// behind worker and at least one target.
	EnableTwoPhaseHandoff bool

	// Optional dependencies
	Metrics       CalculatorAndAssignmentMetrics // Metrics collector (default: no-op)
	Logger        types.Logger                   // Logger (default: no-op)
	StateProvider types.StateProvider            // Manager state provider for degraded mode checks (default: nil)
	// LeaderRevision, if set, is called before each rebalance to obtain the current
	// leader revision (NATS KV revision of the leader key). The value is embedded in
	// every published assignment so workers can detect assignments from a former
	// leader. When nil, LeaderRevision defaults to 0 in published assignments.
	LeaderRevision func() uint64

	// LeaderCheck, if set, is invoked by the publisher's pre-alias and
	// post-alias leadership fences (publish steps 5 and 7 of §3.5). It must
	// perform a LIVE verification of the leader-key revision (e.g., a KV read
	// of the election key) and return nil iff the live revision matches the
	// claimed value. Returning a wrapped types.ErrLeadershipRevisionMismatch
	// signals a takeover; any other error is treated as transient/abort.
	//
	// When nil, the publisher's leadership fences are no-ops (test-only). In
	// production, the manager wires this to its election agent's CheckLeadership
	// method so a former leader cannot pass the fence after another worker has
	// taken over.
	LeaderCheck func(ctx context.Context, claimed uint64) error

	// OnEnumerationError, if set, is invoked when worker enumeration (the
	// heartbeat Keys scan in WorkerMonitor.GetActiveWorkers) has failed
	// EnumerationFailureThreshold times in a row with a non-connectivity error —
	// notably a context.DeadlineExceeded on the stream-wide scan while single-key
	// heartbeat Puts still succeed (NP-10). The manager wires this to a DIRECT
	// degrade (NOT the transient KV-error window, which a succeeding heartbeat Put
	// would clear). When nil, sustained enumeration failure remains logged-only
	// (back-compat / test default). Must be safe for concurrent use.
	OnEnumerationError func(err error)

	// OnEnumerationSuccess, if set, is invoked on every SUCCESSFUL worker
	// enumeration — i.e. whenever the heartbeat Keys scan (GetActiveWorkers)
	// returns a nil error — REGARDLESS of the F10-A worker-credibility decision
	// (it fires even when a sharply-shrunk result is treated as suspicious). The
	// manager stamps it as the "enumeration recovered" signal its recovery exit
	// gate keys on. Must be safe for concurrent use.
	OnEnumerationSuccess func()

	// EnumerationFailureThreshold is the number of consecutive non-connectivity
	// enumeration failures before OnEnumerationError fires. Default: 3 (also
	// applied to any value < 1, so a zero/negative config cannot fire on the
	// first failure). Set explicitly to 1 to fire on the first failure.
	EnumerationFailureThreshold int

	// UnlabeledPartitionPolicy: "dedicated" (default when empty) routes
	// unlabeled partitions to unlabeled workers only (falling back to all
	// workers when none are live); "shared" routes them to all workers.
	UnlabeledPartitionPolicy string

	// LabelSpillGrace is how long a label pool must be continuously empty
	// before its partitions spill. 0 = immediate spill.
	LabelSpillGrace time.Duration

	// OnLabelReadBroadFailure, if set, is invoked when a rebalance aborts
	// because worker label reads failed broadly (connectivity/degrading-
	// JetStream class, or above the isolated-failure cap). The manager
	// wires this to its KV-error recorder so sustained heartbeat-read
	// trouble trips the degraded circuit (spec §6). Mirrors the
	// OnEnumerationError seam. Must be safe for concurrent use.
	OnLabelReadBroadFailure func(err error)
}

Config holds calculator configuration.

Use NewCalculatorWithConfig(cfg) to create a calculator with validated configuration and sensible defaults for optional fields.

Required fields must be set before calling NewCalculatorWithConfig. Optional fields will be set to sensible defaults if zero-valued.

func (*Config) SetDefaults 

func (c *Config) SetDefaults()

SetDefaults applies default values for optional fields.

This method is called automatically by NewCalculatorWithConfig. Fields that are already set (non-zero) are not overwritten.

func (*Config) Validate 

func (c *Config) Validate() error

Validate checks configuration validity.

Returns an error if any required field is missing or invalid.

type EmergencyDetector 

type EmergencyDetector struct {
	// contains filtered or unexported fields
}

EmergencyDetector tracks worker disappearances with hysteresis to prevent false positives from transient network issues.

Workers must remain disappeared for the grace period before triggering an emergency rebalance. This prevents flapping during brief connectivity loss.

func NewEmergencyDetector 

func NewEmergencyDetector(gracePeriod time.Duration) *EmergencyDetector

NewEmergencyDetector creates a new emergency detector with specified grace period.

The grace period prevents false positives from transient network issues by requiring workers to remain disappeared for the full duration before triggering emergency rebalancing.

Parameters:

  • gracePeriod: Minimum time workers must be missing (recommended: 1.5 * HeartbeatInterval)

Returns:

  • *EmergencyDetector: Initialized detector ready for use

func (*EmergencyDetector) CheckEmergency 

func (d *EmergencyDetector) CheckEmergency(
	prev, curr map[string]bool,
) (emergency bool, confirmed []string, pending bool)

CheckEmergency reconciles detector state with a fresh poll observation. Atomic under the detector mutex.

Phases:

  1. Clear by curr — any heartbeat-visible worker is alive.
  2. Track newly missing workers in prev (firstSeen preserved if already tracked).
  3. Safety valve — drop stranded (!prev) entries older than 10*gracePeriod to bound map growth under pathological churn.
  4. Confirm — entries in prev whose firstSeen exceeds gracePeriod.

Parameters:

  • prev: Previous set of active worker IDs.
  • curr: Current set of active worker IDs.

Returns:

  • emergency: true if at least one worker's grace period has expired.
  • confirmed: workers whose grace period has expired (empty if none).
  • pending: true if at least one tracked entry is in prev (informational — allows callers to suppress planned_scale while a disappearance is mid-grace).

func (*EmergencyDetector) ObserveAlive added in v2.4.1 

func (d *EmergencyDetector) ObserveAlive(alive []string)

ObserveAlive clears tracking for every worker observed alive in alive.

Intended to be called by every code path that performs a fresh live-worker scan, not only CheckEmergency: rebalance(), audit-repair flows, partition- lifecycle rebalances, manual TriggerRebalance — all of them perform an independent KV scan and would otherwise update c.lastWorkers (eventually, via handleRebalance) without the detector observing the freshly-seen workers. This invariant must hold: 

"firstSeen[A] is the moment since which A has been continuously absent
 from any leader observation of the live worker set."

ObserveAlive is the side-channel that maintains it for non-poll observations.

Parameters:

  • alive: Worker IDs observed alive in the most recent fresh scan.

type LabelReadModel added in v2.10.1 

type LabelReadModel struct {
	PoolSizes map[string]int
	Parked    map[string]int
}

LabelReadModel is the retained label snapshot behind Manager.LabelState: per-label live pool sizes and parked partition counts from the last successfully published rebalance. Keys carry exact parity with the recordLabelMetrics gauge pass (topo.SortedLabels — labeled pools only).

type LeaderCheckFunc added in v2.4.0 

type LeaderCheckFunc func(ctx context.Context, claimed uint64) error

LeaderCheckFunc is the contract for the publisher's pre-alias and post-alias leadership fences (publish steps 5 and 7 of §3.5).

Production implementations MUST perform a live verification (e.g., a kv.Get against the election leader key and a revision comparison) and return a wrapped types.ErrLeadershipRevisionMismatch when the live revision differs from claimed. A cached/stale check is NOT acceptable: a former leader whose local term-revision has not yet been cleared could otherwise pass the fence after another worker has taken over.

Returning any non-nil error aborts the batch; the publisher wraps the returned error with the appropriate sentinel (types.ErrLeadershipLostPreAlias or types.ErrLeadershipLostPostAlias).

type LiveRefsProvider added in v2.4.0 

type LiveRefsProvider interface {
	// LiveRefs returns a point-in-time snapshot of payload keys the publisher
	// has selected for an in-progress publish. The GC must NOT delete any
	// returned key during this sweep, even if it appears unreferenced and is
	// older than Retention.
	LiveRefs() []string
}

LiveRefsProvider is the contract the GC uses to consult the publisher's in-flight payload-ref set. The default implementation (AssignmentPublisher) returns the keys of every payload it has selected for an in-progress publish; tests can inject a fake provider to drive race scenarios deterministically.

type NopCalculator 

type NopCalculator struct{}

NopCalculator implements a no-op assignment calculator. It is used when the manager is not the leader or during initialization.

func NewNopCalculator 

func NewNopCalculator() *NopCalculator

NewNopCalculator creates a new no-op calculator.

func (*NopCalculator) GetState added in v2.4.1 

func (n *NopCalculator) GetState() types.CalculatorState

GetState implements the Calculator interface. Always reports Idle so the manager's reconcile arm is a no-op when no real calculator is wired.

func (*NopCalculator) Start 

func (n *NopCalculator) Start(ctx context.Context) error

Start implements the Calculator interface.

func (*NopCalculator) Stop 

func (n *NopCalculator) Stop(ctx context.Context) error

Stop implements the Calculator interface.

func (*NopCalculator) SubscribeToStateChanges 

func (n *NopCalculator) SubscribeToStateChanges() (<-chan types.CalculatorState, func())

SubscribeToStateChanges implements the Calculator interface.

func (*NopCalculator) TriggerRebalance 

func (n *NopCalculator) TriggerRebalance(ctx context.Context) error

TriggerRebalance implements the Calculator interface.

type PublishInput added in v2.4.0 

type PublishInput struct {
	// Workers is the active worker set the assignments cover. Determines
	// AssignmentCommit.Workers (sorted internally).
	Workers []string

	// Assignments maps worker ID → its assigned partition slice.
	Assignments map[string][]types.Partition

	// SourcePartitions is the full sorted-or-unsorted partition list from the
	// source. The publisher uses it for the strict set-equality coverage check
	// at publish step 3 (see §3.8). Pass exactly what was returned from the
	// source snapshot — the publisher canonicalizes internally.
	SourcePartitions []types.Partition

	// SourceRevision and SourceRevisionKnown come from
	// types.RevisionedPartitionSource.Snapshot when available. SourceRevision=0
	// + SourceRevisionKnown=false is the legitimate non-revisioned-source
	// state.
	SourceRevision      uint64
	SourceRevisionKnown bool

	// LeaderRevision is the publisher's claimed leader epoch. Steps 5 and 7
	// re-read the live election state and abort if it disagrees.
	LeaderRevision uint64

	// Lifecycle is the rebalance reason. Diagnostic only.
	Lifecycle string

	// WorkersToRemove is the list of legacy assignment.<W> aliases the
	// publisher should sweep AFTER a successful commit. Workers not in
	// AssignmentCommit.Workers are implicitly revoked at the commit level;
	// alias sweeping is rolling-upgrade hygiene only.
	WorkersToRemove []string

	// ParkedPartitions is the set deliberately left unassigned this batch.
	// Coverage becomes: assigned ∪ parked == source AND assigned ∩ parked == ∅.
	// Nil (the default) degenerates to the pre-label set-equality check.
	ParkedPartitions []types.Partition

	// WorkerLabels maps workerID → labels-of-record for payload stamping.
	// Workers absent from the map get WorkerLabels=nil (still Known=true).
	WorkerLabels map[string][]string
}

PublishInput captures the per-batch inputs to Publish that have to come from the calculator (because the publisher does not own the strategy or the source snapshot path). Group them in a struct so future additions don't re-break callers.

type PublisherConfig added in v2.4.0 

type PublisherConfig struct {
	AssignmentKV    jetstream.KeyValue
	HeartbeatKV     jetstream.KeyValue
	Prefix          string // e.g. "assignment"
	HeartbeatPrefix string // e.g. "heartbeat"
	LeaderCheckFn   LeaderCheckFunc
	GCTriggerFn     func()
	Logger          types.Logger
	Metrics         PublisherDependentMetrics

	// IsShuttingDown, when non-nil, is consulted immediately before the
	// commit CAS (step 9 of §3.5). If it returns true, Publish aborts with
	// errShuttingDown and no `_commit` write is attempted. Calculator wires
	// this to detect that its stopCh has been closed. nil means no gate
	// (test-only).
	IsShuttingDown func() bool

	// Now returns the current time and backs the lastRebalance stamp. When
	// nil, time.Now is used. Tests inject a controllable clock so the
	// calculator's cooldown gate is deterministic under load.
	Now func() time.Time
}

PublisherConfig captures the publisher's construction-time dependencies.

LeaderCheckFn is required for production correctness: it backs the pre-alias and post-alias leadership rechecks (publish steps 5 and 7) and MUST perform a live verification (see LeaderCheckFunc). HeartbeatKV is required for the legacy alias barrier (publish step 6). GCTriggerFn is optional and, if set, is called after a successful commit to wake the GC loop (publish step 12).

type PublisherDependentMetrics added in v2.4.0 

type PublisherDependentMetrics interface {
	types.AssignmentMetrics
	types.PublisherMetrics
	types.GCMetrics
}

PublisherDependentMetrics is the metrics surface the publisher (and GC) write to.

The interface composes AssignmentMetrics (for legacy-style change tracking), PublisherMetrics (refs-always counters), and GCMetrics (payload reaping). Concrete MetricsCollector implementations satisfy all three.

type StateMachine 

type StateMachine struct {
	// contains filtered or unexported fields
}

StateMachine manages calculator state transitions.

Implements a validated state machine with these states:

  • Idle: Ready for rebalancing
  • Scaling: Waiting for stabilization window
  • Rebalancing: Computing/publishing assignments
  • Emergency: Immediate rebalancing (no window)

Valid transitions are enforced to prevent invalid states.

func NewStateMachine 

func NewStateMachine(
	logger types.Logger,
	metrics types.CalculatorMetrics,
	onRebalance func(ctx context.Context, reason string) error,
	stopCh chan struct{},
) *StateMachine

NewStateMachine creates a new state machine.

Parameters:

  • logger: Logger for state transitions
  • metrics: Metrics collector for calculator operations
  • onRebalance: Callback invoked when rebalancing should occur
  • stopCh: Channel to signal shutdown (for canceling scaling timers)

Returns:

  • *StateMachine: A new state machine instance starting in Idle state

func (*StateMachine) EnterEmergency 

func (sm *StateMachine) EnterEmergency(ctx context.Context)

EnterEmergency transitions to emergency state for immediate rebalancing.

Backward-compatible wrapper around TryClaimEmergency + RunClaimedRebalance. If the claim fails (Rebalancing or Emergency already in flight), the call is deferred — the next poll cycle will detect the topology change.

Parameters:

  • ctx: Context for the rebalance operation

func (*StateMachine) EnterRebalancing 

func (sm *StateMachine) EnterRebalancing(ctx context.Context)

EnterRebalancing transitions to rebalancing state and triggers the rebalance callback.

This method must be called from the scaling-timer goroutine after the stabilization window expires. It uses a strict-source CAS from Scaling to Rebalancing: if any other transition (e.g., a concurrent emergency claim) has moved the state out of Scaling, the call is a no-op.

On success or error, the state is returned to Idle to allow the next change to be picked up.

Parameters:

  • ctx: Context for the rebalance operation

func (*StateMachine) EnterScaling 

func (sm *StateMachine) EnterScaling(ctx context.Context, reason string, window time.Duration)

EnterScaling transitions to scaling state and starts a stabilization timer.

This method enforces that the transition only occurs from Idle state. If the current state is not Idle, the transition is rejected and the original state is preserved.

Parameters:

  • ctx: Context for the scaling timer goroutine
  • reason: Reason for scaling ("cold_start", "planned_scale", "restart")
  • window: Stabilization window duration before rebalancing

func (*StateMachine) GetScalingReason 

func (sm *StateMachine) GetScalingReason() string

GetScalingReason returns the reason for the current scaling operation.

Returns:

  • string: Scaling reason ("cold_start", "planned_scale", "emergency", "restart", or "")

func (*StateMachine) GetState 

func (sm *StateMachine) GetState() types.CalculatorState

GetState returns the current calculator state.

This method is thread-safe and can be called concurrently.

Returns:

  • types.CalculatorState: Current state (Idle, Scaling, Rebalancing, or Emergency)

func (*StateMachine) ReturnToIdle 

func (sm *StateMachine) ReturnToIdle()

ReturnToIdle transitions the state machine back to idle after rebalancing completes.

This method clears the scaling reason and notifies all subscribers of the state change.

func (*StateMachine) RunClaimedRebalance added in v2.4.1 

func (sm *StateMachine) RunClaimedRebalance(ctx context.Context, reason string)

RunClaimedRebalance runs the rebalance callback for a previously-claimed lifecycle and returns the state machine to Idle.

Must be called after a successful TryClaimEmergency or after a successful strict-source CAS into Rebalancing (see EnterRebalancing).

Parameters:

  • ctx: Context for the rebalance operation
  • reason: Lifecycle reason to pass to the callback ("emergency", "cold_start", etc.)

func (*StateMachine) RunClaimedRebalanceErr added in v2.4.1 

func (sm *StateMachine) RunClaimedRebalanceErr(ctx context.Context, reason string) error

RunClaimedRebalanceErr runs the partition-rebalance callback for a previously-claimed lifecycle and returns the callback's error to the caller. The FSM is returned to Idle whether the callback succeeded, failed, or returned errShuttingDown.

Distinct from RunClaimedRebalance: this variant gives the caller the callback error so the partition-lifecycle path can run restorePendingOnGraceBail when the rebalance bails on a recovery-grace re-check. The general RunClaimedRebalance preserves the existing void contract for emergency/scaling callers whose callback swallows errShuttingDown.

Must be called after a successful TryClaimRebalancing.

func (*StateMachine) Subscribe 

func (sm *StateMachine) Subscribe() (<-chan types.CalculatorState, func())

Subscribe returns a channel that receives state change notifications.

The returned channel is buffered (size 4) to allow for rapid state transitions without blocking the state machine. The subscriber receives the current state immediately upon subscription.

Returns:

  • <-chan types.CalculatorState: Channel that receives state updates
  • func(): Unsubscribe function to clean up resources

Example: 

ch, unsubscribe := sm.Subscribe()
defer unsubscribe()
for state := range ch {
    fmt.Printf("State changed to: %s\n", state)
}

func (*StateMachine) TryClaimEmergency added in v2.4.1 

func (sm *StateMachine) TryClaimEmergency(_ context.Context) bool

TryClaimEmergency attempts to atomically claim the emergency lifecycle.

Strict-source CAS from either Idle or Scaling to Emergency. Returns true if the claim succeeded; the caller is then responsible for invoking RunClaimedRebalance to execute the rebalance and return to Idle.

When the claim succeeds, scalingReason is set BEFORE the state-change notification is fanned out to subscribers, so subscribers observing the Emergency state see the correct reason on their first read.

Parameters:

  • ctx: Context (currently unused but kept for symmetry with other Enter* methods)

Returns:

  • bool: true if the emergency lifecycle was successfully claimed; false if another lifecycle (Rebalancing or Emergency) is already in flight.

func (*StateMachine) TryClaimRebalancing added in v2.4.1 

func (sm *StateMachine) TryClaimRebalancing(_ context.Context, reason string) bool

TryClaimRebalancing attempts to atomically claim the rebalancing lifecycle from Idle. Strict-source CAS from Idle to Rebalancing.

Designed for partition-lifecycle callers (monitorPartitions) that need to drive a rebalance but did NOT first enter Scaling. Without this primitive monitorPartitions would bypass the FSM entirely, violating the public contract that Rebalancing means "active partition rebalance in progress".

scalingReason is recorded BEFORE the state-change notification fans out so subscribers see the reason on their first read of the Rebalancing state.

Returns true if the claim succeeded; the caller is then responsible for invoking RunClaimedRebalanceErr to execute the rebalance and return to Idle.

func (*StateMachine) WaitForShutdown 

func (sm *StateMachine) WaitForShutdown()

WaitForShutdown waits for all scaling timer goroutines to complete.

This should be called during shutdown after closing the stopCh to ensure all goroutines have exited cleanly.

type WorkerMonitor 

type WorkerMonitor struct {
	// contains filtered or unexported fields
}

WorkerMonitor handles worker health detection via NATS KV heartbeats.

It provides hybrid monitoring:

  • Watcher (primary): Fast detection <100ms via NATS KV Watch. Routine heartbeat refreshes are classified against a session-local lastSeen map and suppressed (no check, no Keys() scan); joins, graceful leaves, and sweep-detected silent expiries trigger checks.
  • Polling (fallback): Reliable detection ~hbTTL/2 via periodic KV scan; the ground-truth path for silent TTL expiry.

The monitor runs in a background goroutine and invokes a callback when worker topology changes are detected.

func NewWorkerMonitor 

func NewWorkerMonitor(
	heartbeatKV jetstream.KeyValue,
	hbPrefix string,
	hbTTL time.Duration,
	onChange func(ctx context.Context) error,
	logger types.Logger,
) *WorkerMonitor

NewWorkerMonitor creates a new worker monitor.

Parameters:

  • heartbeatKV: NATS KV bucket for worker heartbeats
  • hbPrefix: Prefix for heartbeat keys (e.g., "worker")
  • hbTTL: Heartbeat TTL duration
  • onChange: Callback invoked when worker changes are detected
  • logger: Logger for monitoring events

Returns:

  • *WorkerMonitor: A new worker monitor instance

func (*WorkerMonitor) GetActiveWorkers 

func (m *WorkerMonitor) GetActiveWorkers(ctx context.Context) ([]string, error)

GetActiveWorkers retrieves the list of workers with active heartbeats.

This method scans the heartbeat KV bucket for keys matching the configured prefix and extracts worker IDs from the key names.

Parameters:

  • ctx: Context for cancellation

Returns:

  • []string: List of active worker IDs
  • error: Nil on success, error on KV access failure

func (*WorkerMonitor) GetHeartbeats added in v2.4.0 

func (m *WorkerMonitor) GetHeartbeats(ctx context.Context) (map[string]types.Heartbeat, error)

GetHeartbeats returns the decoded heartbeats for every worker with an active heartbeat key. The map is keyed by worker ID.

Decoding accepts both v1 JSON heartbeats (new workers) and legacy RFC3339 timestamp strings (pre-v1 workers); see types.DecodeHeartbeat. Workers whose payload fails to decode are silently omitted — a malformed heartbeat is logged at debug level but does not fail the entire scan.

Every per-key Get is bounded by the same pass deadline (opCtx) that guards the Keys() scan — mirroring GetHeartbeatsFor's per-iteration opCtx.Err() guard — so a single wedged Get cannot hang the whole pass past that deadline. Once the deadline has expired anywhere in the pass, the pass aborts and returns the partial map collected so far alongside a wrapped error. That includes an empty Keys() scan: nats.go surfaces a context that closes before any entry is yielded as its no-keys error, so an empty result is accepted as an authoritative empty worker set only while the pass deadline is still live — an expired empty scan aborts with the wrapped error instead.

Returns:

  • map[string]types.Heartbeat: Decoded heartbeats collected before any abort, keyed by worker ID
  • error: Non-nil on KV access failure that prevents listing keys, or when the pass deadline expires at any point in the pass — including an empty Keys() scan whose deadline had already expired (partial map still returned on mid-scan aborts)

func (*WorkerMonitor) GetHeartbeatsFor added in v2.9.0 

func (m *WorkerMonitor) GetHeartbeatsFor(ctx context.Context, workerIDs []string) (map[string]types.Heartbeat, map[string]error, error)

GetHeartbeatsFor returns decoded heartbeats for exactly the given worker IDs, keyed by worker ID. Unlike GetHeartbeats it does NOT run a Keys() scan — callers that already hold the active worker list (the rebalance path) use this to avoid a second stream-wide enumeration.

Per-worker Get or decode failures omit that worker from the heartbeat map (logged at debug) but populate the error map under that worker ID with the original error, so the caller can classify the failure (connectivity/degrading-JetStream vs not) rather than have it masquerade as a bare "unknown labels" omission. The only non-nil returned error is context cancellation/deadline.

func (*WorkerMonitor) SetOnLabelChange added in v2.9.0 

func (m *WorkerMonitor) SetOnLabelChange(fn func())

SetOnLabelChange registers a callback invoked when a heartbeat PUT's labels differ from the retained fingerprint for that key — a worker-label change behind a live worker ID (e.g. a stale-ID takeover) that the worker-set change path cannot detect. The callback is coalesced by the caller (the calculator routes it through requestLabelRecheck). It must be called before Start, like the constructor's onChange callback; the field is read only by the watcher goroutine afterward.

func (*WorkerMonitor) Start 

func (m *WorkerMonitor) Start(ctx context.Context) error

Start begins monitoring workers in a background goroutine.

The monitor uses a hybrid approach:

  1. NATS KV watcher for fast detection (~100ms)
  2. Periodic polling as fallback (every hbTTL/2)

Parameters:

  • ctx: Context for cancellation (affects watcher lifetime)

Returns:

  • error: Error if already started or already stopped

func (*WorkerMonitor) Stop 

func (m *WorkerMonitor) Stop() error

Stop stops the worker monitor and waits for cleanup.

This method blocks until all monitoring goroutines have exited. It is safe to call Stop multiple times - subsequent calls will return immediately.

Returns:

  • error: Error if Stop called before Start, nil otherwise

Source Files

View all Source files

Directories

Path Synopsis

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.