coordinator

package
v0.18.52 Latest Latest
Warning

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

 Go to latest Published: Sep 6, 2026 License: AGPL-3.0 Imports: 49 Imported by: 0

Documentation

Overview

Package coordinator manages query planning, scheduling, and lifecycle tracking.

Index

Constants

View Source 
const DefaultLocalFastPathBytes = 64 << 20

DefaultLocalFastPathBytes is the default routing threshold: a query whose total post-pruning catalog scan bytes stay under this executes in-process on the coordinator instead of as a distributed stage DAG. The DAG's fixed costs (task dispatch + object-store materialization per stage boundary) are independent of data size, so small queries pay a latency floor the local pipeline doesn't have.

Variables

View Source 
var EagerChainedEdgesPlanned atomic.Int64

EagerChainedEdgesPlanned counts the subset of EagerEdgesPlanned whose consumer is a stage-chain-fused join (§13/A2): the fused chain cleared on its primary probe/build feeds while its chained-build deps completed at the barrier.

View Source 
var EagerEdgesPlanned atomic.Int64

EagerEdgesPlanned counts consumer stages that cleared dispatch on an eager feed instead of the done barrier — the memo §8 activation marker for the SF100 pair (grep "eager dispatch: consumer cleared early" / this counter's log line in benchmark.log).

View Source 
var EagerManifestsPublished atomic.Int64

EagerManifestsPublished counts manifests published across all queries — the mechanism marker proving the eager path engaged in a benchmark log (the SortMergeJoinsPlanned/DynamicFiltersPlanned observability precedent).

View Source 
var FastPathStrict = optswitch.Register("fastpath-strict", "WADJET_FASTPATH_STRICT",
	"report a local fast-path execution failure instead of retrying it on the DAG")

FastPathStrict makes a local execution FAILURE final instead of retrying the query on the distributed DAG (#308). Retry is kept only for failures that say nothing about the query's meaning — the deliberate result-budget bail-out, a local memory budget a worker fleet would not hit, and object-store unavailability. Every other execution error is deterministic: re-running it on the DAG either reproduces it (wasted work) or, when the two paths disagree, silently returns rows the local path considered impossible to produce. That second case is precisely the class of defect #312 was — a fallback would have shipped its wrong answer instead of surfacing anything. Kill switch: WADJET_FASTPATH_STRICT=0 restores the unconditional fallback.

Routing and PLANNING failures are unaffected: no query semantics have been evaluated at that point, and the DAG legitimately covers plan shapes the local pipeline declines.

View Source 
var SkewSplitsPlanned atomic.Int64

SkewSplitsPlanned counts hot-group split decisions (one per group per stage dispatch). The A/B mechanism marker: a wall-clock delta without this counter moving (and its "skew split planned" Info log) is window drift, not skew-split signal.

Functions

func IsInputLostErr 

func IsInputLostErr(err error) bool

IsInputLostErr reports whether err (anywhere in its message chain) is a streaming-exchange input-lost failure.

func QueryIDFromPath 

func QueryIDFromPath(path string) string

QueryIDFromPath extracts the query ID from a result file path.

func ReadResultFiles 

func ReadResultFiles(ctx context.Context, store objstore.Store, bucket string, paths []string) ([]*batch.RecordBatch, []string, int64, error)

ReadResultFiles reads result Parquet files from S3 and returns columnar batches. This is intended for callers (tpch-bench, CLI) that need full result data — they pull from S3 directly instead of routing through the coordinator's heap.

Types

type BatchStream 

type BatchStream interface {
	Next(ctx context.Context) (*batch.RecordBatch, error)
	Close() error
}

BatchStream is a lazy, single-consumer iterator over result batches. Next returns (nil, nil) when exhausted. Implementations drop their reference to each batch as it is handed out, so a consumer that sends and releases batch-by-batch keeps peak residency at one batch plus whatever the producer still buffers.

Close releases everything still held — buffered batches and, for spill-backed streams, replay scratch on disk. It is idempotent and must be called by whoever owns the stream, including on error and early-exit paths: an undrained spill-backed stream pins its scratch file until Close runs.

func NewSliceStream 

func NewSliceStream(batches []*batch.RecordBatch) BatchStream

NewSliceStream adapts a materialized batch slice to BatchStream. The stream takes ownership of the slice (entries are nil'd as yielded).

type BuildStats 

type BuildStats struct {
	FilterID  string
	KeyType   string
	HasRange  bool
	Min, Max  int64
	RowCount  int64
	Bloom     []uint64
	BloomMask uint64
	// StagedBucket/Key set when the unioned bloom exceeded the inline
	// threshold and was uploaded; Bloom is nil in that case.
	StagedBucket string
	StagedKey    string
}

BuildStats is the coordinator-merged dynamic-filter artifact for one FilterID across all build-scan tasks of one stage. Computed by mergeBuildStatsFromPartials after the build-scan stage completes.

type Config 

type Config struct {
	NATSUrl        string
	ResultBucket   string
	MaxInflight    int           // max concurrent queries, 0 = default (64)
	QueryTimeout   time.Duration // max time for a query to complete, 0 = default (30m)
	WorkerStaleTTL time.Duration // time after which a silent worker is reaped, 0 = default (30s)
	DynamicFilters bool          // Trino-style semi-join dynamic-filter pushdown (off by default in v1)
	// GatherResultBudget caps the decoded result bytes a single query may
	// hold in coordinator heap (gather receiver). 0 = derive from
	// GOMEMLIMIT (half of it) or fall back to 2 GiB; negative = uncapped.
	// Past the budget the remaining gather payload spills as raw frames to
	// local scratch and the SQLResult replays them lazily disk→wire; only
	// a scratch-write failure fails the query (cleanly — the coordinator
	// process never OOMs for one query's result size).
	GatherResultBudget int64
	// LocalFastPathBytes routes queries whose total post-pruning catalog
	// scan bytes stay under this threshold onto the coordinator-local
	// single-process pipeline, skipping task dispatch and per-stage
	// object-store materialization. <=0 = disabled (every query runs the
	// distributed DAG — the zero value keeps library/test semantics
	// unchanged). `wadjet serve` enables it by default via its flag
	// (DefaultLocalFastPathBytes).
	LocalFastPathBytes int64
	// BroadcastBytesOverride, when non-zero, replaces the cluster-derived
	// broadcast threshold (broadcastThresholdFromCluster): >0 = fixed byte
	// threshold, <0 = never broadcast (every join takes the hash-shuffle /
	// sort-merge path). 0 = derive from worker pool budget (default).
	// Primarily a benchmarking/debugging surface: forcing the shuffled-join
	// path at small scale is how the sort-merge join gate gets exercised
	// before data is big enough to defeat broadcast on its own.
	BroadcastBytesOverride int64
	// SortMergeJoinBytes routes inner equi-joins whose sides BOTH exceed
	// this estimated size through the sort-merge join instead of the hash
	// join (docs/design/sort-merge-join.md), on the local fast path and in
	// the distributed stage DAG alike (the join stage swaps operator; its
	// exchange children are identical). 0 = disabled (default, dormant).
	SortMergeJoinBytes int64
	// IntermediateTTL is how long a queries/<id>/* prefix that the
	// per-query cleanup missed survives before the periodic sweep reclaims
	// it. 0 = 1 hour. It is a backstop, not the primary reclamation: every
	// terminal path (completion, error, cancel) runs cleanupQuery
	// (ADR-0028), so a prefix reaching this TTL means one of them was
	// prevented — typically by an open delete breaker.
	IntermediateTTL time.Duration
	// LateMaterialization emits inner/left hash-join output as view
	// (dictionary) columns with the gather deferred to first touch
	// (docs/design/late-materialization.md), on the local fast path and in
	// worker fragments alike (rides the join-probe OpSpec). Off by default.
	LateMaterialization bool
	// SkewSplit enables adaptive skew-aware task layout for shuffled hash
	// joins (docs/design/skew-aware-shuffle.md): hot partition groups —
	// detected from the worker-reported per-partition shuffle output bytes
	// — split into k sub-tasks that divide the group's probe files and
	// replicate its build files, bounding the straggler task's input and
	// memory footprint. Decision logic in skew_split.go. The wadjet CLI
	// and tpch-bench default this ON (2026-07-11 SF10 A/B: −41% straggler
	// wall on the hot-key fixture, plan-identical on uniform workloads via
	// the ratio gate); this struct field's zero value stays false so
	// embedded/test constructors opt in explicitly.
	SkewSplit bool
	// AggPartialSplit enables the round-robin partial-aggregate fan-out
	// (aggregatePartialSplit in execute_stage_dag.go): partial "aggregate"
	// stages over a non-trivial multi-file upstream split into at most
	// workerCount tasks aggregating disjoint file slices. The wadjet CLI
	// and tpch-bench default this ON; --agg-partial-split=false is the
	// kill switch. Zero value stays false so embedded/test constructors
	// opt in explicitly (mirrors SkewSplit).
	AggPartialSplit bool
	// EagerDispatch enables eager consumer dispatch (docs/design/
	// eager-consumer-dispatch.md): the coordinator republishes per-
	// producer-task file manifests on EagerManifestSubject as shuffle
	// tasks complete, and (Phase C1) dispatches eligible consumer stages
	// before their producer stage fully drains. Zero value false —
	// default off until SF100 validation; requires StreamingExchange.
	EagerDispatch bool
	// StreamingExchange annotates dispatched tasks with peer-location
	// hints (Task.InputLocations) and per-query fetch tokens so consumers
	// stream stage outputs from the producing workers' local disk instead
	// of S3 (Phase A, docs/design/streaming-exchange.md). Purely additive:
	// hints only reference workers that advertise a PeerAddr, every fetch
	// failure falls through to the unchanged S3 read path, and the write
	// path (synchronous upload before stage completion) is untouched.
	// Default false: dormant, no hints, no tokens.
	StreamingExchange bool
	// ShuffleDurability is the stage-output upload policy stamped on
	// dispatched stage/shuffle tasks (docs/design/shuffle-durability.md).
	// Eager (zero value) starts every background S3 upload immediately —
	// the pre-knob behavior. Lazy queues uploads unstarted on the workers
	// and releases them only on demand (a consumer missing an input whose
	// producer is alive, a coordinator-side stage read, or worker drain);
	// scratch a query finishes without ever needing durably is elided.
	// Off never uploads scratch: producer death degrades to the one-shot
	// streaming-disabled re-execution (the ErrInputLost fallback), and
	// draining a worker mid-query loses its outputs the same way.
	// Stages whose outputs the coordinator itself reads (scalar-subquery
	// producers) always stay eager — the coordinator has no peer tier.
	// Only meaningful with StreamingExchange (the peer tier is what makes
	// the durable copy optional).
	ShuffleDurability distributed.UploadPolicy
	// LocalityPlacement places a task whose peer-location hints all point
	// at one connected worker onto that worker (docs/design/locality-
	// placement.md): 1:1 stage chains read their whole input set via
	// same-worker mmap instead of peer gRPC streams. Requires
	// StreamingExchange (the hint source) and the gRPC data plane
	// (targeted dispatch). Zero value false — default off until SF100
	// validation (mirrors EagerDispatch).
	LocalityPlacement bool
}

Config holds coordinator configuration.

type Coordinator 

type Coordinator struct {

	// BuildCacheThreshold overrides the default build cache threshold (bytes).
	// Zero means use the default (2GB). Exported for testing with small datasets.
	BuildCacheThreshold int64
	// contains filtered or unexported fields
}

Coordinator accepts queries, plans them, dispatches tasks, and tracks results.

func New 

func New(cfg Config, cat *catalog.Catalog, nc *nats.Conn, js jetstream.JetStream, logger *slog.Logger) *Coordinator

New creates a new Coordinator.

func (*Coordinator) CancelQuery 

func (c *Coordinator) CancelQuery(queryID string) error

CancelQuery cancels a running query.

func (*Coordinator) Cleaner 

func (c *Coordinator) Cleaner(store objstore.Store, bucket string) *ResultCleaner

Cleaner returns the result cleaner, creating it if needed. Its TTL is the sweep backstop for prefixes the per-query cleanup did not reclaim; Config.IntermediateTTL (--query-intermediate-ttl) sets it, 0 = 1 hour.

func (*Coordinator) CorrelatedLocalRoutes 

func (c *Coordinator) CorrelatedLocalRoutes() int64

CorrelatedLocalRoutes reports how many refused correlated-subquery plans were routed to the coordinator-local pipeline. Exposed for tests and observability — a distributed suite asserting DAG engagement uses it to prove the refusal fires ONLY for the correlated shapes.

func (*Coordinator) DistinctLocalRoutes added in v0.18.1 

func (c *Coordinator) DistinctLocalRoutes() int64

DistinctLocalRoutes reports how many plans refused for an unstageable DISTINCT were routed to the coordinator-local pipeline. Separate from CorrelatedLocalRoutes so a suite can assert which refusal fired.

func (*Coordinator) EnforcesABAC 

func (c *Coordinator) EnforcesABAC() bool

EnforcesABAC reports whether ExecuteSQL enforces access policies itself. pgwire uses this to decide that routing authed connections through the coordinator is safe (canBypassDB).

func (*Coordinator) ExecuteSQL 

func (c *Coordinator) ExecuteSQL(ctx context.Context, sql string) (res *SQLResult, err error)

ExecuteSQL parses SQL, plans, distributes across workers, and collects results.

func (*Coordinator) GetQueryResults 

func (c *Coordinator) GetQueryResults(ctx context.Context, queryID string) (*SQLResult, error)

GetQueryResults retrieves the final results for a completed query.

func (*Coordinator) GetQueryStatus 

func (c *Coordinator) GetQueryStatus(queryID string) (*QueryStatus, error)

GetQueryStatus returns the current status of a query.

func (*Coordinator) GroupKeyLocalRoutes added in v0.18.9 

func (c *Coordinator) GroupKeyLocalRoutes() int64

GroupKeyLocalRoutes reports how many plans refused for a GROUP BY key whose published name a stage cannot carry were routed to the coordinator-local pipeline (#736). Separate from the others so a suite can assert WHICH refusal fired — and, just as importantly, that an ordinary computed key did NOT fire it and stayed distributed.

func (*Coordinator) GroupingSetsLocalRoutes added in v0.18.9 

func (c *Coordinator) GroupingSetsLocalRoutes() int64

GroupingSetsLocalRoutes reports how many plans refused for GROUPING SETS / ROLLUP / CUBE were routed to the coordinator-local pipeline. Separate from the others so a suite can assert WHICH refusal fired — a gate that only checked the rows would pass just as happily if the DAG had started answering them by accident (#778).

func (*Coordinator) InSubqueryLocalRoutes added in v0.18.3 

func (c *Coordinator) InSubqueryLocalRoutes() int64

InSubqueryLocalRoutes reports how many plans refused for an unmaterializable IN-subquery were routed to the coordinator-local pipeline. Separate from the other two counters so a suite can assert WHICH refusal fired.

func (*Coordinator) ListQueries 

func (c *Coordinator) ListQueries() []QueryStatus

ListQueries returns recent query statuses.

func (*Coordinator) LocalFastPathBails 

func (c *Coordinator) LocalFastPathBails() int64

LocalFastPathBails reports how many local executions bailed out over the result budget and re-dispatched as DAG queries.

func (*Coordinator) LocalFastPathHits 

func (c *Coordinator) LocalFastPathHits() int64

LocalFastPathHits reports how many queries executed on the local fast path. Exposed for tests and observability.

func (*Coordinator) LocalFastPathStrictFailures 

func (c *Coordinator) LocalFastPathStrictFailures() int64

LocalFastPathStrictFailures reports how many local executions were reported to the client instead of retried on the DAG.

func (*Coordinator) MaybeRestoreCatalog 

func (c *Coordinator) MaybeRestoreCatalog(ctx context.Context, forceTS string) error

MaybeRestoreCatalog restores the catalog from S3 if:

  • Snapshot options are configured (Store != nil)
  • Either forceTS is set, or the catalog is fresh: no KV at all, or a KV that holds no tables

Otherwise it is a no-op. Errors during restore are propagated unchanged; the caller decides whether to fatal.

Freshness cannot be IsKVEmpty alone: every production caller runs Catalog.Init first, and Init writes the <cluster>.meta key, so the KV is never empty by the time this runs. Gating on it made the whole path dead code — `wadjet serve --catalog-snapshot-s3-prefix=...` against a bucket full of data silently restored nothing and came up with an empty table list. Restore overwrites meta with the snapshot's own copy, so the Init-written stub is fine to clobber. (cmd/tpch-bench reached the same conclusion and worked around it locally; this is that fix, upstream.)

func (*Coordinator) NullAwareAntiLocalRoutes added in v0.18.40 

func (c *Coordinator) NullAwareAntiLocalRoutes() int64

NullAwareAntiLocalRoutes reports how many plans were refused because a null-aware anti join's build could not be shown to reach every task whole, and answered on the coordinator-local pipeline instead.

func (*Coordinator) RecoverQueries 

func (c *Coordinator) RecoverQueries(ctx context.Context) error

RecoverQueries is called when this coordinator becomes leader after a failover. It reads active query states from the store and logs them for manual or automated recovery.

func (*Coordinator) ScalarProjectionLocalRoutes added in v0.18.5 

func (c *Coordinator) ScalarProjectionLocalRoutes() int64

ScalarProjectionLocalRoutes reports how many plans refused for a SELECT-list subquery were routed to the coordinator-local pipeline (#659).

func (*Coordinator) SetAlertsEnabled 

func (c *Coordinator) SetAlertsEnabled(on bool)

SetAlertsEnabled toggles the feature flag. When false, StartAlertScheduler is a no-op. DDL-level rejection is added in Task 13.

func (*Coordinator) SetAuthProvider 

func (c *Coordinator) SetAuthProvider(p *auth.Provider) error

SetAuthProvider wires ABAC enforcement into ExecuteSQL: with a provider set, every query is policy-checked at plan level (table denial, row filters, column deny/mask) for the identity in the request context — the same auth.EnforcePlanPolicies the embedded engine applies. Call before serving traffic (same contract as the other Set<X>-before-Start setters). It also BINDS the policy set's names to the catalog, because attaching a set to a catalog is exactly what this call is — see wadjet.DB.SetAuthProvider for why binding belongs at the attach and not at two `serve` call sites. The error is returned and remembered; an ignored one still refuses at enforcement through Provider.BindError.

func (*Coordinator) SetCatalogSnapshotInterval 

func (c *Coordinator) SetCatalogSnapshotInterval(d time.Duration)

SetCatalogSnapshotInterval configures the periodic snapshot cadence. Zero disables periodic snapshots (explicit-only).

func (*Coordinator) SetCatalogSnapshotOptions 

func (c *Coordinator) SetCatalogSnapshotOptions(opts catalog.SnapshotOptions)

SetCatalogSnapshotOptions configures the S3 target for catalog snapshots. Call once during coordinator setup. If Store is nil, snapshot/restore functionality is disabled.

func (*Coordinator) SetDataPlaneServer 

func (c *Coordinator) SetDataPlaneServer(srv *dataplane.Server)

SetDataPlaneServer enables gRPC result delivery, task dispatch, and progress signaling. When set:

  • gather receivers register as ResultHandlers so workers can stream ResultBatch messages directly instead of via NATS (Phase B),
  • the scheduler routes TaskDispatch over the gRPC stream instead of NATS publish (Phase C),
  • the WorkerRegistry installs a global TaskProgress handler that treats every progress arrival as a liveness signal (Phase E), and per-query stage bridges register for stage-progress fanout (also Phase E, in newStageProgressBridge).

Must be called before any query runs. Pass nil (or skip) to use the NATS-only path.

func (*Coordinator) SetLeaderElection 

func (c *Coordinator) SetLeaderElection(le *LeaderElection)

SetLeaderElection attaches a leader election instance to the coordinator. When set, the coordinator will only accept queries if it is the current leader. If nil (default), the coordinator is always considered leader (standalone mode).

func (*Coordinator) SetQueryLimits added in v0.18.15 

func (c *Coordinator) SetQueryLimits(global *config.QueryLimits, perRole map[string]*config.QueryLimits)

SetQueryLimits wires the configured cost guard into every plan the coordinator builds: the global limits and the per-role overrides (role name -> limits; a role present here uses ITS limits and not the global ones, including a nil entry meaning unlimited).

This is what makes `query_limits:` a control rather than a parsed-and- discarded config section. The guard itself has always been implemented in physical.Planner.enforceQueryLimits and reached only the HTTP server's no-coordinator path; in every `wadjet serve` mode the coordinator answers, and its planners were built without limits, so nothing bounded a query on any protocol (#803).

Call before serving traffic (same contract as SetAuthProvider).

func (*Coordinator) SetQueryStateStore 

func (c *Coordinator) SetQueryStateStore(qs *QueryStateStore)

SetQueryStateStore attaches a query state store for HA persistence. When set, query state transitions are persisted to NATS KV so a new leader can recover in-flight queries after failover.

func (*Coordinator) SetTelemetry 

func (c *Coordinator) SetTelemetry(tp *telemetry.Provider)

SetTelemetry enables OpenTelemetry tracing on the coordinator.

func (*Coordinator) StageReadTierCounts 

func (c *Coordinator) StageReadTierCounts() (kv, peer, s3, peerMisses int64)

StageReadTierCounts reports how many coordinator-side stage-output reads each tier served, plus how many peer attempts fell through to the durable copy. Observability + test assertions; the per-read tier also rides the "scalar substitution" log line.

func (*Coordinator) StartAlertScheduler 

func (c *Coordinator) StartAlertScheduler(parent context.Context)

StartAlertScheduler begins scheduling alerts. Must only be called while this coordinator holds leadership. Safe to call multiple times; a running scheduler is stopped and replaced.

func (*Coordinator) StartCatalogSnapshotLoop 

func (c *Coordinator) StartCatalogSnapshotLoop(parent context.Context)

StartCatalogSnapshotLoop begins periodic catalog snapshots. Safe to call only while this coordinator holds leadership. No-op when snapshot options are not configured or the interval is zero.

func (*Coordinator) StartLeaderWatch 

func (c *Coordinator) StartLeaderWatch(ctx context.Context)

StartLeaderWatch starts a background goroutine that watches for leadership changes and triggers recovery when this coordinator becomes leader.

func (*Coordinator) StartQueryActiveHandler 

func (c *Coordinator) StartQueryActiveHandler()

StartQueryReaper starts a background goroutine that periodically removes StartQueryActiveHandler subscribes to query-active check requests from workers. Workers ask "is query X still active?" before executing tasks pulled from JetStream, preventing wasted work on queries killed by the watchdog.

func (*Coordinator) StartQueryReaper 

func (c *Coordinator) StartQueryReaper(ctx context.Context)

StartQueryReaper periodically removes old entries for completed, failed, and cancelled queries from the tracker and queryMetas maps. This prevents unbounded memory growth from accumulated query metadata.

func (*Coordinator) StopAlertScheduler 

func (c *Coordinator) StopAlertScheduler()

StopAlertScheduler cancels the running scheduler and waits for it to exit. Safe to call when no scheduler is running.

func (*Coordinator) StopCatalogSnapshotLoop 

func (c *Coordinator) StopCatalogSnapshotLoop()

StopCatalogSnapshotLoop cancels the running loop and waits for it to exit. Safe to call when no loop is running.

func (*Coordinator) SubmitSQL 

func (c *Coordinator) SubmitSQL(ctx context.Context, sql string) (queryID string, planStr string, err error)

SubmitSQL parses, plans, and dispatches a query without blocking for results. Returns the query ID and plan string immediately.

func (*Coordinator) SubmitScanQuery 

func (c *Coordinator) SubmitScanQuery(ctx context.Context, tableName string, columns []string, partFilter map[string]string) (*QueryResult, error)

SubmitScanQuery submits a simple scan query for distributed execution. This is the primary entry point before the SQL planner is available.

func (*Coordinator) TableLessLocalRoutes added in v0.18.16 

func (c *Coordinator) TableLessLocalRoutes() int64

TableLessLocalRoutes reports how many plans refused for a table-less SELECT were routed to the coordinator-local pipeline (#806). Separate from the other counters so a suite can assert WHICH refusal fired.

func (*Coordinator) Tracker 

func (c *Coordinator) Tracker() *QueryTracker

Tracker returns the query tracker (for inspection).

func (*Coordinator) UnbuildableStageLocalRoutes added in v0.18.21 

func (c *Coordinator) UnbuildableStageLocalRoutes() int64

UnbuildableStageLocalRoutes reports how many plans refused for an undispatchable stage were routed to the coordinator-local pipeline (#812).

func (*Coordinator) UnreachableOutputLocalRoutes added in v0.18.5 

func (c *Coordinator) UnreachableOutputLocalRoutes() int64

UnreachableOutputLocalRoutes reports how many plans refused for an uncomputed SELECT list were routed to the coordinator-local pipeline.

func (*Coordinator) Workers 

func (c *Coordinator) Workers() *WorkerRegistry

Workers returns the worker registry for inspecting active workers.

type DLQ 

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

DLQ provides access to the dead-letter queue for failed task inspection.

func NewDLQ 

func NewDLQ(js jetstream.JetStream) *DLQ

NewDLQ creates a DLQ reader from a JetStream connection.

func (*DLQ) Count 

func (d *DLQ) Count(ctx context.Context) (uint64, error)

Count returns the number of entries in the DLQ.

func (*DLQ) Get 

func (d *DLQ) Get(ctx context.Context, entryID string) (*distributed.DLQEntry, error)

Get retrieves a specific DLQ entry by its entry ID.

func (*DLQ) List 

func (d *DLQ) List(ctx context.Context, limit int) ([]distributed.DLQEntry, error)

List returns all DLQ entries, most recent first. limit controls the maximum number of entries returned (0 = all, up to stream max).

func (*DLQ) Purge 

func (d *DLQ) Purge(ctx context.Context) error

Purge deletes all DLQ entries.

type LeaderElection 

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

LeaderElection implements active-passive leader election for coordinators using NATS KV with TTL-based leasing. One coordinator is leader (accepts queries), the other is standby (warm, ready to take over).

The leader writes its ID + timestamp, refreshing every 2s with a 5s TTL. CAS (revision-based) updates prevent split-brain. Standbys poll the key and attempt acquisition when it expires.

NOT YET WIRED INTO ANY SERVER MODE. `NewLeaderElection` has no non-test caller and `Coordinator.SetLeaderElection` has no caller at all, so no `wadjet serve` path instantiates this and no deployment has ever run it. The behaviour described here — and every defect fixed in it — is therefore forward-looking: it is what will be true of the first mode that wires it up, not a description of anything a running server does today.

func NewLeaderElection 

func NewLeaderElection(js jetstream.JetStream, id string, logger *slog.Logger) (*LeaderElection, error)

NewLeaderElection creates a new leader election instance. The id should be unique across coordinators (e.g., hostname + pid or UUID).

func (*LeaderElection) CurrentLeader 

func (le *LeaderElection) CurrentLeader(ctx context.Context) string

CurrentLeader returns the ID of the current leader, or empty string if none.

func (*LeaderElection) IsLeader 

func (le *LeaderElection) IsLeader() bool

IsLeader returns true if this coordinator is the current leader.

func (*LeaderElection) LeaderChanged 

func (le *LeaderElection) LeaderChanged() <-chan bool

LeaderChanged returns a channel that receives true when this instance becomes leader, or false when it loses leadership.

func (*LeaderElection) LeaseReclaims added in v0.18.23 

func (le *LeaderElection) LeaseReclaims() int64

LeaseReclaims reports how many times a failed refresh was resolved by consulting the store instead of standing down. Exposed so a gate can tell "the refresh simply succeeded" from "the recovery path ran and held", which the leadership flag alone cannot.

func (*LeaderElection) Start 

func (le *LeaderElection) Start(ctx context.Context) error

Start begins the leader election process. A single run loop goroutine handles both leader and standby states, switching between them as needed.

func (*LeaderElection) Stop 

func (le *LeaderElection) Stop()

Stop gracefully stops the leader election. If this instance is leader, the key will expire after TTL, allowing another coordinator to take over.

type OutputKind 

type OutputKind int

OutputKind describes how a stage's output is distributed across S3 keys. 

const (
	// OutputPartitioned: Files[p] holds the keys for partition p, for p in
	// [0, NumPartitions). Produced by StageExchangeRepartition; consumed by
	// hash-join/aggregate downstream stages that bind one worker per
	// contiguous partition range.
	OutputPartitioned OutputKind = iota
	// OutputReplicated: Files[0] holds broadcast keys. Produced by
	// StageExchangeReplicate; every downstream worker reads the same file.
	OutputReplicated
	// OutputSinglePart: Files[0] holds the full, unpartitioned output.
	// Produced by pipeline stages feeding a Gather (single output stream).
	OutputSinglePart
)

type PersistentQueryState 

type PersistentQueryState struct {
	ID              string    `json:"id"`
	SQL             string    `json:"sql"`
	CompletedStages []string  `json:"completed_stages"`
	Status          string    `json:"status"` // "planning", "executing", "merging", "done", "failed"
	LeaderID        string    `json:"leader_id"`
	StartedAt       time.Time `json:"started_at"`
}

PersistentQueryState represents the persisted state of an in-flight query. This is stored in NATS KV so that a standby coordinator can recover queries after a failover.

type QueryInfo 

type QueryInfo struct {
	QueryID     string
	SQL         string
	State       QueryState
	Stages      map[string]*StageInfo
	StageOrder  []string // topological order
	StartTime   time.Time
	EndTime     time.Time
	Error       string
	ResultFiles []string
	TotalRows   int64
}

QueryInfo tracks the full state of a query.

type QueryResult 

type QueryResult struct {
	QueryID     string        `json:"query_id"`
	State       string        `json:"state"`
	ResultFiles []string      `json:"result_files,omitempty"`
	TotalRows   int64         `json:"total_rows"`
	Elapsed     time.Duration `json:"elapsed"`
	Error       string        `json:"error,omitempty"`
}

QueryResult represents the outcome of a query execution.

type QueryState 

type QueryState int

QueryState represents the lifecycle state of a query. 

const (
	QueryStatePending QueryState = iota
	QueryStateRunning
	QueryStateCompleted
	QueryStateFailed
	QueryStateCancelled
)

func (QueryState) String 

func (s QueryState) String() string

type QueryStateStore 

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

QueryStateStore persists in-flight query state to NATS KV for HA recovery. When a coordinator fails, the new leader can read active query states and resume or clean up in-flight queries.

func NewQueryStateStore 

func NewQueryStateStore(js jetstream.JetStream, logger *slog.Logger) (*QueryStateStore, error)

NewQueryStateStore creates a query state store backed by NATS KV.

func (*QueryStateStore) Delete 

func (s *QueryStateStore) Delete(ctx context.Context, queryID string) error

Delete removes a query state. Called when a query completes successfully.

func (*QueryStateStore) Get 

func (s *QueryStateStore) Get(ctx context.Context, queryID string) (*PersistentQueryState, error)

Get retrieves a query state by ID. Returns nil if not found.

func (*QueryStateStore) ListActive 

func (s *QueryStateStore) ListActive(ctx context.Context) ([]*PersistentQueryState, error)

ListActive returns all query states that are not in a terminal state (i.e., not "done" or "failed"). Used during failover recovery.

func (*QueryStateStore) Save 

func (s *QueryStateStore) Save(ctx context.Context, state *PersistentQueryState) error

Save persists a query state to NATS KV. Overwrites any existing entry for the same query ID.

type QueryStatus 

type QueryStatus struct {
	QueryID   string        `json:"query_id"`
	SQL       string        `json:"sql"`
	State     string        `json:"state"`
	Stages    []StageStatus `json:"stages,omitempty"`
	Elapsed   time.Duration `json:"elapsed"`
	TotalRows int64         `json:"total_rows"`
	Error     string        `json:"error,omitempty"`
}

QueryStatus represents the current status of an async query.

type QueryTracker 

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

QueryTracker manages per-query lifecycle state.

func NewQueryTracker 

func NewQueryTracker() *QueryTracker

NewQueryTracker creates a new query tracker.

func (*QueryTracker) ActiveQueryIDs 

func (qt *QueryTracker) ActiveQueryIDs() map[string]struct{}

ActiveQueryIDs returns the set of query IDs in a non-terminal state (pending or running). Used by ResultCleaner to avoid deleting intermediate files for in-flight queries.

func (*QueryTracker) Cancel 

func (qt *QueryTracker) Cancel(queryID string)

Cancel marks a query as cancelled.

func (*QueryTracker) ClearResults 

func (qt *QueryTracker) ClearResults(queryID string)

ClearResults drops all stored result data (InlineData) for a query. Called after results have been read and merged to free memory immediately rather than waiting for the reaper TTL.

func (*QueryTracker) CollectResultPaths 

func (qt *QueryTracker) CollectResultPaths(queryID string) map[string][]string

CollectResultPaths gathers all result file paths from completed stages under a single lock acquisition. This avoids the race where a shallow copy's stage pointers are read without synchronization.

func (*QueryTracker) Complete 

func (qt *QueryTracker) Complete(queryID string)

Complete marks a query as completed.

func (*QueryTracker) Delete 

func (qt *QueryTracker) Delete(queryID string)

Delete removes a query from the tracker entirely.

func (*QueryTracker) Fail 

func (qt *QueryTracker) Fail(queryID string, err string)

Fail marks a query as failed.

func (*QueryTracker) Get 

func (qt *QueryTracker) Get(queryID string) *QueryInfo

Get returns a snapshot of a query's state, safe to read without qt.mu held. See snapshotStages for why a plain `copy := *q` is not enough (#514).

func (*QueryTracker) GetReadyStages 

func (qt *QueryTracker) GetReadyStages(queryID string) []string

GetReadyStages returns stages whose dependencies are all complete. Takes the exclusive lock, not RLock: it marks each returned stage Scheduled and stamps ScheduledAt, mutating *StageInfo state in place. RLock only serializes against Lock, not against other concurrent RLock holders, so those writes under RLock would race both with themselves (two overlapping GetReadyStages calls) and with any other RLock reader of the same fields (StalledStages) — the same class of bug as #514, just write-write instead of write-read.

func (*QueryTracker) IsComplete 

func (qt *QueryTracker) IsComplete(queryID string) bool

IsComplete returns true if all stages of the query are done.

func (*QueryTracker) List 

func (qt *QueryTracker) List() []*QueryInfo

List returns all tracked queries as snapshots (see snapshotStages) — no exported caller currently reads a listed query's Stages, but List has the exact same shallow-copy shape Get had, so it gets the same fix rather than waiting for its own #514-shaped report.

func (*QueryTracker) MarkScheduled 

func (qt *QueryTracker) MarkScheduled(queryID, stageID string)

MarkScheduled marks a stage as already dispatched, preventing GetReadyStages from re-scheduling it.

func (*QueryTracker) ReapCompleted 

func (qt *QueryTracker) ReapCompleted(maxAge time.Duration) []string

ReapCompleted returns and removes query IDs in a terminal state (completed, failed, cancelled) whose EndTime is older than maxAge.

func (*QueryTracker) RecordResult 

func (qt *QueryTracker) RecordResult(result distributed.ResultNotification) (stageComplete bool)

RecordResult records a task result for a query stage. Idempotent — duplicate results from worker retries are silently ignored. Returns true if the stage is now complete.

func (*QueryTracker) Register 

func (qt *QueryTracker) Register(queryID, sql string, stages map[string]*StageInfo, stageOrder []string)

Register registers a new query.

func (*QueryTracker) SetStageTasks 

func (qt *QueryTracker) SetStageTasks(queryID, stageID string, count int)

SetStageTasks updates the total task count for a stage (used when intermediate stage tasks are created dynamically after dependencies complete).

func (*QueryTracker) StageFailed 

func (qt *QueryTracker) StageFailed(queryID, stageID string) string

StageFailed returns the error message if a completed stage has all tasks failed (no successful tasks). Returns "" if the stage succeeded or is still in progress.

func (*QueryTracker) StageResults 

func (qt *QueryTracker) StageResults(queryID, stageID string) []distributed.ResultNotification

StageResults returns the result notifications for a given stage.

func (*QueryTracker) StalledStages 

func (qt *QueryTracker) StalledStages(queryID string, threshold time.Duration) []string

StalledStages returns stages that have been scheduled for longer than the given threshold but still have unreported tasks. Used by the coordinator watchdog to detect and fail queries where result notifications were lost.

func (*QueryTracker) Start 

func (qt *QueryTracker) Start(queryID string)

Start marks a query as running.

func (*QueryTracker) UpdateResultPath 

func (qt *QueryTracker) UpdateResultPath(queryID, stageID, taskID, path string)

UpdateResultPath updates the result path for a specific task result. Used when inline results are materialized to S3 for downstream stages.

type ResultCleaner 

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

ResultCleaner manages cleanup of query result files in object storage.

func NewResultCleaner 

func NewResultCleaner(store objstore.Store, bucket string, ttl time.Duration, logger *slog.Logger) *ResultCleaner

NewResultCleaner creates a result cleaner.

func (*ResultCleaner) CleanAll 

func (rc *ResultCleaner) CleanAll(ctx context.Context) (int, error)

CleanAll removes every queries/ object regardless of TTL or active set. Intended for shutdown — when the coordinator is exiting, in-flight queries are already dying, so leftover intermediates would never be read again and just leak. A normal-running coordinator should use CleanQuery for individual completions and CleanStale for periodic GC.

func (*ResultCleaner) CleanQuery 

func (rc *ResultCleaner) CleanQuery(ctx context.Context, queryID string) (int, error)

CleanQuery removes all result files for a specific query.

It reports what it actually reclaimed. The loop used to log every failure and continue with a nil return, so a caller was told the cleanup succeeded after deleting nothing (measured: deleted=0, err=<nil> while the store's breaker was open) and nothing ever retried. It also kept issuing deletes after the caller's deadline had expired, manufacturing one instant DeadlineExceeded per remaining object — a burst of consecutive failures that opened the breaker by itself (#820, and the producer half of #798). The first ctx error ends the loop; anything left is named in the error.

func (*ResultCleaner) CleanStale 

func (rc *ResultCleaner) CleanStale(ctx context.Context) (int, error)

CleanStale removes result files older than the TTL, skipping files that belong to queries currently in-flight.

func (*ResultCleaner) SetActiveQueriesFunc 

func (rc *ResultCleaner) SetActiveQueriesFunc(fn func() map[string]struct{})

SetActiveQueriesFunc registers a callback that returns the set of query IDs currently in-flight. CleanStale will skip files belonging to these queries regardless of age.

func (*ResultCleaner) StartPeriodicCleanup 

func (rc *ResultCleaner) StartPeriodicCleanup(ctx context.Context, interval time.Duration)

StartPeriodicCleanup runs cleanup on a schedule.

type SQLResult 

type SQLResult struct {
	QueryID     string
	Columns     []string
	Batches     []*batch.RecordBatch
	ResultFiles []string
	TotalRows   int64
	Elapsed     time.Duration
	Plan        string
	Error       string

	// Schema is the declared type of each output column, in Columns order.
	// Only needed for a result whose Batches are not set — a lazy stream —
	// since OutputSchema() otherwise reads it off the first batch.
	Schema []parquet.Column

	// WireUnconstrainedDecimal names the DECIMAL columns in Schema whose
	// PostgreSQL wire typmod must say "unconstrained" (-1) even though
	// Schema itself carries their real (p,s) — an aggregate function call,
	// unlike a bare column reference, on live PostgreSQL. Unlike Schema's
	// zero-row-only role above, this is a PLAN property consulted for every
	// result (FIX 2, #457/#458 fold-in; see
	// physical.declaredWireUnconstrainedDecimal).
	WireUnconstrainedDecimal map[string]bool
	// StringLength names the output columns in Schema whose declaration
	// carries a string LENGTH — `CAST(x AS VARCHAR(4))` — and what it is, so
	// the wire can send PostgreSQL's `character varying(n)` rather than an
	// unconstrained text (#838). Plan-time, like the map above.
	StringLength map[string]int
	// contains filtered or unexported fields
}

SQLResult holds the result of a distributed SQL query. Results are kept columnar (as RecordBatches) to avoid materializing per-row map[string]any which causes massive heap pressure at SF10+.

Results arrive in one of two forms: fully materialized in Batches, or as a lazy stream (gather results that exceeded the coordinator budget and spilled to local scratch — see gatherReceiver). Consumers should iterate via Stream(), which handles both forms; whoever receives an SQLResult owns it and must either drain the stream or call Close, or spill scratch leaks until process exit.

func (*SQLResult) Close 

func (r *SQLResult) Close() error

Close releases whatever the result still holds — the lazy stream's buffered batches and spill scratch, or the materialized slice. Idempotent and nil-safe.

func (*SQLResult) OutputSchema 

func (r *SQLResult) OutputSchema() []parquet.Column

OutputSchema returns the declared type of each result column, or nil when the result carries no schema (introspection and error results).

The DAG's gather hands back typed RecordBatches and the local fast path runs a typed pipeline, but SQLResult used to publish only Columns — so pgwire's coord path had no ColumnMetas and fell back to declaring OID 25 (text) for every column, while the embedded API's identical query declared real OIDs from deriveColumnMetas. Same statement, two different RowDescriptions depending on which entry point answered.

Read before Stream(): the first Stream() call detaches Batches.

func (*SQLResult) Rows 

func (r *SQLResult) Rows() ([]map[string]any, error)

Rows materializes the result batches into row-oriented maps. This is expensive for large results — prefer Stream(). On a materialized result it is repeatable (Batches are left in place, matching the historical behavior); on a lazy result it consumes the stream, and a second call returns nothing.

func (*SQLResult) Stream 

func (r *SQLResult) Stream() BatchStream

Stream returns a consuming iterator over the result batches. The first call detaches the result's batches (lazy stream or materialized slice); subsequent calls return an empty stream. The caller owns the returned stream and must drain it or call Close.

type Scheduler 

type Scheduler struct {

	// PolicedQuery answers, for a query ID, whether a row or column security
	// policy shaped that query's plan. The SQL-text guard in PublishTasks
	// asks it rather than the CONTEXT, because a dispatcher may publish under
	// a context of its own: SubmitSQL's async publish runs on a ctx derived
	// from context.Background(), which drops any mark set at planning time,
	// so a ctx-based guard was inert exactly on the door it was written for
	// (#859 round 3). Nil = nothing is policed.
	PolicedQuery func(queryID string) bool
	// contains filtered or unexported fields
}

Scheduler publishes tasks to NATS and tracks stage DAGs.

func NewScheduler 

func NewScheduler(nc *nats.Conn, logger *slog.Logger) *Scheduler

NewScheduler creates a new task scheduler.

func (*Scheduler) PublishTasks 

func (s *Scheduler) PublishTasks(ctx context.Context, tasks []distributed.Task) error

PublishTasks publishes a set of tasks for worker consumption. Routes via gRPC TaskDispatch (data-plane server) when configured, otherwise falls back to NATS JetStream publish. Tasks are serialized in batch before publishing to minimize time spent holding the transport.

func (*Scheduler) SetDataPlaneServer 

func (s *Scheduler) SetDataPlaneServer(srv *dataplane.Server)

SetDataPlaneServer enables gRPC task dispatch. When set, PublishTasks pushes through the data-plane server instead of NATS publish.

func (*Scheduler) SetLocalityPlacement 

func (s *Scheduler) SetLocalityPlacement(on bool)

SetLocalityPlacement toggles input-locality placement. Call before any query runs (same contract as the other Set<X> setters).

func (*Scheduler) SetTaskAnnotator 

func (s *Scheduler) SetTaskAnnotator(fn func(*distributed.Task))

SetTaskAnnotator installs the per-task pre-dispatch hook. Call before any query runs (same contract as the other Set<X> setters).

func (*Scheduler) SetWorkerRegistry 

func (s *Scheduler) SetWorkerRegistry(wr *WorkerRegistry)

SetWorkerRegistry enables memory-aware placement for gRPC dispatch: tasks carrying an EstimatedBytes go to the worker with the most free pool (heartbeat budget − heartbeat used − in-flight estimates) instead of round-robin.

func (*Scheduler) TaskDone 

func (s *Scheduler) TaskDone(taskID string)

TaskDone clears the task's in-flight admission estimate. Called by the stage result subscriptions on every ResultNotification; the TTL covers lost results.

type ShuffleLayout 

type ShuffleLayout struct {
	BuildAlias    string
	ProbeAlias    string
	NumPartitions int
	// BuildShardFiles[p] contains the S3 keys for build partition p.
	// May be nil/empty for a partition if no build rows hashed there.
	BuildShardFiles [][]string
	// ProbeShardFiles[p] contains the S3 keys for probe partition p.
	ProbeShardFiles [][]string
	// Per-partition on-disk bytes for each side, reduced across the final
	// surviving shuffle task set (worker-reported; nil when unreported).
	// Recorded for observability — the legacy pipeline path built from this
	// layout stays non-adaptive; skew splitting lives on the native-DAG path.
	BuildPartitionBytes []int64
	ProbePartitionBytes []int64
}

ShuffleLayout describes the shard file layout produced by the two-sided shuffle stages. The caller (coordinator routing path) constructs probe pipeline tasks from this layout by assigning contiguous partition slices to each worker and populating PreScannedInputs with the corresponding shard files.

type StageInfo 

type StageInfo struct {
	StageID      string
	Type         distributed.TaskType
	TotalTasks   int
	DoneTasks    int
	FailedTasks  int
	Scheduled    bool      // true once tasks have been dispatched (prevents re-scheduling)
	ScheduledAt  time.Time // when tasks were dispatched (zero if not yet scheduled)
	Dependencies []string  // stage IDs that must complete before this stage
	Results      []distributed.ResultNotification
	SeenTasks    map[string]bool // dedup: task IDs already recorded (handles retries)
}

StageInfo tracks the state of a single stage within a query.

type StageOutput 

type StageOutput struct {
	Kind          OutputKind
	NumPartitions int        // valid when Kind == OutputPartitioned
	Files         [][]string // Partitioned: Files[p]; Replicated/SinglePart: Files[0]

	// Bytes is the total on-disk size of Files as reported by the workers
	// that wrote them (ResultNotification.SizeBytes), or the catalog size
	// for pass-through leaf scans. Feeds downstream tasks' EstimatedBytes
	// for memory-aware admission. 0 = unknown (legacy worker, empty
	// output) — downstream estimates degrade to unknown, never to wrong.
	// Note these are file bytes (compressed), the same unit the scan-stage
	// estimates use; in-memory footprint is larger by the codec ratio.
	Bytes int64

	// PartitionRows/PartitionBytes: per-partition output totals reduced
	// across the producing stage's tasks (final surviving attempts only),
	// indexed by partition id. Valid when Kind == OutputPartitioned; nil
	// when the workers didn't report (legacy build) — downstream skew
	// detection degrades to off, never to wrong. Bytes are on-disk
	// uncompressed .wshf sizes, the unit per-task admission estimates use.
	PartitionRows  []int64
	PartitionBytes []int64

	// BuildStats, populated when this stage's Stage.EmitDynamicFilters was
	// non-empty. One entry per emit, keyed by FilterID. Downstream consumer
	// stages reach in via the stat-dep edge to pull the materialized stats.
	BuildStats map[string]*BuildStats

	// DynamicFilters, populated when this stage's Stage.ConsumeDynamicFilters
	// resolved against an upstream build-scan's BuildStats. Downstream
	// dispatchers (shuffle, compute) thread these into their task
	// descriptors so workers apply the bloom at scan time. Carries forward
	// across pass-through leaf scans that don't dispatch tasks themselves.
	DynamicFilters []distributed.DynamicFilterSpec

	// ScanFileSizes aligns 1:1 with Files[0] on the pass-through leaf
	// scan path (catalog SizeBytes, copied from Stage.ScanFileSizes) so
	// the downstream shuffle dispatcher's synthetic source stage can
	// byte-balance its affine fan-out. Nil elsewhere.
	ScanFileSizes []int64

	// ScanTable/ScanColumns identify a pass-through leaf scan's relation
	// and (sanitized) projection. Set ONLY on the no-task pass-through
	// path, where downstream shuffle tasks read base parquet directly —
	// dispatchShuffleStage copies them onto its synthetic source stage so
	// prunedScanColumns can narrow the shuffle payload. Without them the
	// synthetic stage has no column list and every scan-absorbed shuffle
	// leg ships full-width rows (85.69 GB vs ~12 GB needed on Q21's l2
	// leg; docs/design/exchange-reuse.md §2 A2). Empty on all other
	// outputs — WSHF inputs ignore column projection.
	ScanTable   string
	ScanColumns []string
	// ScanSchema is the CATALOG's declared schema for that same relation,
	// and it rides the pass-through for the same reason ScanColumns does:
	// the consumer really is reading base-table parquet, so it needs to be
	// told what the file cannot say about itself. Nine of this engine's
	// types have no parquet annotation, so a file written before the
	// declared-schema footer key existed hands the consumer an INT64 where
	// the catalog says IPv4 — 167772165 instead of 10.0.0.5, on the DAG
	// only (#423). Empty on every non-pass-through output, where the input
	// is WSHF and carries its own types.
	ScanSchema []parquet.Column
	// contains filtered or unexported fields
}

StageOutput is the materialized output of one stage in the native-DAG executor. Produced by each dispatchXxxStage helper; consumed by downstream stages' collectInputs.

type StageStatus 

type StageStatus struct {
	StageID     string `json:"stage_id"`
	Type        string `json:"type"`
	TotalTasks  int    `json:"total_tasks"`
	DoneTasks   int    `json:"done_tasks"`
	FailedTasks int    `json:"failed_tasks"`
}

StageStatus represents the progress of a single query stage.

type TaskLiveness 

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

TaskLiveness tracks when each in-flight task was last reported active. Updated from worker heartbeats AND from per-task TaskProgress messages — either signal is proof the worker process is alive (see WorkerRegistry's progress subscription for why we need both). Used to detect stuck tasks.

It also records which worker each gRPC-dispatched task was sent to (Assign). Assignment alone does NOT make a task eligible for the stuck-task sweep — a task queued behind a busy worker's slots is mentioned in no heartbeat and must never be falsely re-dispatched. The binding exists for ExpireWorker: when the registry reaps a dead worker, every task dispatched to it — including tasks it never got to report before dying — enters the stuck set at once. Without this, a task that died with its worker before ever appearing in a heartbeat fell through both recovery nets and wedged its stage until the query timeout (2026-08-11 Q09/Q21-R2 coordinator wedge).

func NewTaskLiveness 

func NewTaskLiveness() *TaskLiveness

func (*TaskLiveness) Assign 

func (tl *TaskLiveness) Assign(taskID, workerID string)

Assign records which worker a task was dispatched to. Re-dispatches overwrite the binding with the new target. Deliberately does not touch the task's liveness clock (see type comment).

func (*TaskLiveness) ExpireWorker 

func (tl *TaskLiveness) ExpireWorker(workerID string) []string

ExpireWorker moves every task assigned to workerID into the stuck set immediately (zero last-seen), returning the affected task IDs. Called when the registry reaps a dead worker so its dispatched-but-unreported tasks re-dispatch on the next stuck sweep instead of never.

func (*TaskLiveness) Remove 

func (tl *TaskLiveness) Remove(taskID string)

Remove stops tracking the given task (completed or failed).

func (*TaskLiveness) StuckTasks 

func (tl *TaskLiveness) StuckTasks(threshold time.Duration) []string

StuckTasks returns task IDs that haven't been reported in a heartbeat for longer than the given threshold.

func (*TaskLiveness) Update 

func (tl *TaskLiveness) Update(taskIDs []string, now time.Time)

Update records that the given task IDs are actively running.

type WorkerInfo 

type WorkerInfo struct {
	WorkerID      string
	ClusterID     string
	MaxConcurrent int // effective task-slot count reported in heartbeat; 0 = legacy worker (assume default)
	MemoryUsed    int64
	MemoryTotal   int64
	PoolUsed      int64    // shared memory pool bytes Reserved
	PoolBudget    int64    // shared memory pool capacity in bytes; pressure = PoolUsed / PoolBudget
	ActiveTaskIDs []string // task IDs in flight per the most recent heartbeat
	Draining      bool
	PeerAddr      string // dialable peer-exchange address; "" = worker serves no peer fetches
	LastSeen      time.Time
}

WorkerInfo tracks the state of a registered worker.

type WorkerRegistry 

type WorkerRegistry struct {
	Liveness *TaskLiveness // per-task progress tracking from heartbeats AND TaskProgress messages

	// PendingNonDurable reports how many of the worker's recorded outputs
	// still lack a durable copy. Set before StartReaper (never after — the
	// reaper goroutine reads it unlocked); nil disables the grace.
	PendingNonDurable func(workerID string) int
	// contains filtered or unexported fields
}

WorkerRegistry tracks active workers from heartbeats.

func NewWorkerRegistry 

func NewWorkerRegistry(nc *nats.Conn, logger *slog.Logger, staleTTL time.Duration) *WorkerRegistry

NewWorkerRegistry creates a worker registry that subscribes to heartbeats. staleTTL controls how long since the last heartbeat before a worker is considered dead. Pass 0 to use the default (90s).

func (*WorkerRegistry) ActiveWorkers 

func (wr *WorkerRegistry) ActiveWorkers() []*WorkerInfo

ActiveWorkers returns workers that have sent a heartbeat recently and are not draining. Draining workers are finishing in-flight tasks and will not accept new work.

func (*WorkerRegistry) Close 

func (wr *WorkerRegistry) Close()

Close unsubscribes from heartbeats and task progress.

func (*WorkerRegistry) ClusterCapacity 

func (wr *WorkerRegistry) ClusterCapacity() int

ClusterCapacity returns the sum of effective task-slot counts across active (non-draining, non-stale) workers, taken from the most recent heartbeat each worker reported. Workers running pre-MaxConcurrent-in- heartbeat builds report 0; those are skipped here so the result reflects only workers we have honest capacity data for. Returns 0 when no worker has reported MaxConcurrent yet — callers should fall back to a conservative static cap (e.g. 2 * Count).

func (*WorkerRegistry) ClusterPoolPressure 

func (wr *WorkerRegistry) ClusterPoolPressure() float64

ClusterPoolPressure returns the cluster-wide shared memory pool pressure as a value in [0.0, 1.0], computed from the most recent heartbeat each active worker reported. Returns 0 when no worker has reported pool stats (legacy workers, or worker without --shared-pool-budget). Observability only — earlier per-stage dispatch gating on this value caused a deadlock in broadcast_join chains (the probe-side stage that would free build-side memory was itself gated on memory dropping). Worker-side pool-pressure spill is the load-shedding primitive now.

func (*WorkerRegistry) Count 

func (wr *WorkerRegistry) Count() int

Count returns the number of active workers.

func (*WorkerRegistry) IsAlive 

func (wr *WorkerRegistry) IsAlive(workerID string) bool

IsAlive reports whether the worker has been seen within the stale TTL (draining still counts as alive — a draining worker serves peer fetches and finishes uploads).

func (*WorkerRegistry) MarkWorkerSeen 

func (wr *WorkerRegistry) MarkWorkerSeen(workerID string)

MarkWorkerSeen records that the given worker is alive as of `now`. Multi-signal liveness: any worker→coord NATS message (TaskProgress, ResultNotification, GatherBatchMsg) is proof the worker process is running and able to publish, even if its heartbeat goroutine is starved or its heartbeat NATS publish is lagging. Coord still requires a real heartbeat for initial registration (cluster ID, pool stats), but post-registration any of these signals counts toward LastSeen.

No-op when workerID is empty (some workers don't stamp messages until the worker ID is plumbed end-to-end).

func (*WorkerRegistry) MayRecover 

func (wr *WorkerRegistry) MayRecover(workerID string) bool

MayRecover reports whether the worker is alive OR still registered inside the reap-grace window (silent for less than stale+grace). The input-lost classifiers use it in place of IsAlive so a grace-deferred producer keeps missing-input failures retryable — its background uploads may still land, or the peer path may resume, before the grace expires. Once the reaper removes the worker this returns false.

func (*WorkerRegistry) MinWorkerPoolBudget 

func (wr *WorkerRegistry) MinWorkerPoolBudget() int64

MinWorkerPoolBudget returns the smallest non-zero shared-pool budget reported by any active worker. Used by the planner to bound broadcast-vs- shuffle decisions: a build that wouldn't fit in the smallest worker's memory pool must NOT be broadcast (every worker would pay the full duplication). Returns 0 when no worker has reported pool stats yet (legacy workers, or pre-heartbeat phase) — callers should treat 0 as "no information; use absolute defaults."

func (*WorkerRegistry) PeerAddr 

func (wr *WorkerRegistry) PeerAddr(workerID string) string

PeerAddr returns the peer-exchange address of a live worker, or "" when the worker is unknown, stale, or doesn't serve peer fetches. Draining workers still serve fetches for files they already hold, so unlike ActiveWorkers this only requires freshness.

func (*WorkerRegistry) ReapStale 

func (wr *WorkerRegistry) ReapStale() int

ReapStale removes workers that haven't sent a heartbeat recently. Reaped workers receive a drain message so they stop pulling new tasks from the JetStream consumer. If the worker is truly dead, the drain is a no-op; if it's alive but GC-stalled, it will stop accepting work.

func (*WorkerRegistry) SetDataPlaneServer 

func (wr *WorkerRegistry) SetDataPlaneServer(srv *dataplane.Server)

SetDataPlaneServer installs the registry's global TaskProgress handler on the data-plane server. When workers send TaskProgress over gRPC (Phase E), this handler treats each arrival as proof of life for the emitting worker — same liveness semantics as the NATS progressSub above. Idempotent; a later SetDataPlaneServer(nil) installs a no-op handler. Must be called after NewWorkerRegistry.

func (*WorkerRegistry) StartReaper 

func (wr *WorkerRegistry) StartReaper(ctx context.Context)

StartReaper starts a background goroutine that periodically removes stale workers.

func (*WorkerRegistry) StartSubStatsLogger 

func (wr *WorkerRegistry) StartSubStatsLogger(ctx context.Context)

StartSubStatsLogger periodically logs PendingMsgs/PendingBytes/Dropped/ Delivered for the long-lived subscriptions (heartbeat + TaskProgress). 2026-05-03 SF10 mass-reap was consistent with NATS server-side slow- consumer drops on these subs but we had no instrumentation to confirm. With this running, drops surface in the coord log as they happen.

Cadence: 10s — same as worker heartbeat cadence so a stat-tick reliably sits between consecutive heartbeats. Stats are O(1) to read; the only cost is one log line per tick.

Source Files

View all Source files

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.