worker

package
v0.18.23 Latest Latest
Warning

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

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

Documentation

Overview

Package worker implements the distributed task execution worker.

Index

Constants

View Source
const StaleInputAttemptMarker = "eager-dispatch stale input attempt"

manifestStreamSource is the eager-consumer input source (docs/design/eager-consumer-dispatch.md §3.2): it consumes one producer stage's shuffle files as each producer TASK completes, instead of a frozen file list built after the whole stage drained.

Contract:

  • The candidate set (spec.ProducerTaskIDs) is fixed at task build; which files exist, and where, streams in as ProducerTaskManifests (Replay for tasks completed before dispatch, NATS for the rest — subscribe happens in Init, before any wait, so nothing is missed; duplicates are idempotent).
  • Files outside [PartitionStart, PartitionEnd] are ignored.
  • Reads go through the standard tiered fetch (LocalStageCache → peer → S3) by delegating each resolved manifest's in-range files to an inner cachedFileStreamSource; manifest PeerAddr hints are registered so the peer tier works for files the task spec could not have hinted.
  • Attempt fencing (§5): consuming any file of producer task T pins T's attempt. A later manifest for T with a higher attempt poisons the source; Next returns errStaleInputAttempt and the task fails loudly (coordinator retries it against the stable attempt set).

StaleInputAttemptMarker tags the poison error so the coordinator's result classification can retry the consumer task without burning the generic failure path's diagnostics.

Variables

View Source
var DeclaredSchemaStrict = optswitch.Register("declared-schema-strict", "WADJET_DECLARED_SCHEMA_STRICT",
	"refuse a base-table parquet read that arrives with no declared schema instead of typing it from the file")

DeclaredSchemaStrict gates the refusal above. Registered in the kill-switch registry so the optimization-invariance oracle runs the corpus with it off as well as on: a healthy fixture declares every base-table read, so the two configurations must answer identically, and a divergence would mean some path is already reading a file the catalog disagrees with.

View Source
var ForcedMorselCollapses atomic.Int64

ForcedMorselCollapses counts collapses taken because of the knob rather than because of memory pressure. A gate asserts it moved: a forcing knob that silently stopped engaging turns the gate it arms into a no-op.

View Source
var StageUploadsRefused atomic.Int64

StageUploadsRefused counts stage-output uploads refused because the query was already tombstoned. It is the M3 half of #625: the coordinator's per-query cleanup is a one-shot LIST+DELETE, and a straggler task that finishes after it lands recreates the prefix that was just reclaimed — then nothing revisits it until the TTL sweep. The async upload path has refused tombstoned roots since the q22-R2 stall (queryState returns nil); the SYNCHRONOUS stage uploads did not, and a synchronous .wshf landing after ExecuteSQL returned is exactly what round-0 measured.

Functions

func CompressShuffleData

func CompressShuffleData(data []byte) []byte

CompressShuffleData compresses raw WSHF data into the upload envelope: "WSHC" + s2 stream by default, "WSHZ" + zstd stream under WADJET_EXCHANGE_ZSTD=1. Uses streaming (not block) format to support arbitrarily large payloads. If the compressed output is not smaller, the original WSHF data is returned.

func CompressShuffleFile

func CompressShuffleFile(srcPath, dstPath string) (compressedSize int64, useCompressed bool, err error)

CompressShuffleFile streams srcPath through S2 into dstPath, prefixed by the WSHC magic. Returns (compressedSize, useCompressed, error). useCompressed is true when the compressed output is ≥10 % smaller than the source, matching CompressShuffleData's heuristic. When useCompressed is false the caller should drop dst and upload src.

Heap cost is bounded by the s2.Writer's internal block buffer (~64 KB) regardless of file size. This is the file-streaming counterpart to CompressShuffleData, used by the shuffle path to avoid materialising whole partition files in heap at SF100+ scale (project_per_task_share_landed_2026-05-20 followup — executeShuffle's os.ReadFile was the dominant 62 % of heap at the Q05 stall, per pprof on 2026-05-21).

func DecompressShuffleData

func DecompressShuffleData(data []byte) ([]byte, error)

DecompressShuffleData detects and decompresses a WSHC or WSHZ payload back to raw WSHF. If the data is already plain WSHF (or non-shuffle), it is returned unchanged.

func ForceMorselCollapseEvery added in v0.18.23

func ForceMorselCollapseEvery(n int64) int64

ForceMorselCollapseEvery arms (n > 0) or disarms (n <= 0) forced collapse for every morsel-parallel breaker in this process, returning the previous setting so a test can restore it. TEST ONLY.

func MmapReliefThreshold

func MmapReliefThreshold() int64

MmapReliefThreshold returns the current relief threshold in bytes (0 if unset).

func SetMmapRelief

func SetMmapRelief(enabled bool, thresholdBytes int64)

SetMmapRelief configures the relief mechanism from worker config. thresholdMB <= 0 leaves the default. Called once before any task runs.

func WSHZStats

func WSHZStats() (files, bytes int64)

WSHZStats reports how many uploads chose the WSHZ envelope and their total compressed bytes.

Types

type CachedStore

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

CachedStore wraps an objstore.Store with the worker's LRU cache so that scanners (which use cat.Store() directly) benefit from cross-query file caching. Without this, each query re-reads all Parquet files from S3.

CachedStore implements both objstore.Store and objstore.ReaderAtStore. For ReaderAt requests, it downloads the full file into the cache and serves random-access reads from memory — this is faster than S3 range requests for files that will be accessed across multiple queries.

func NewCachedStore

func NewCachedStore(inner objstore.Store, cache *LRUCache, tracker *memory.Tracker) *CachedStore

NewCachedStore creates a store that checks the LRU cache before delegating to the underlying store. Files read from the inner store are automatically cached for subsequent queries. tracker (nil-safe) accounts each download buffer for the window between allocation and cache insertion — before the Put the bytes are invisible to admission and spill pressure; after it the LRU cache's reservoir owns them.

func (*CachedStore) BucketExists

func (s *CachedStore) BucketExists(ctx context.Context, bucket string) (bool, error)

func (*CachedStore) Delete

func (s *CachedStore) Delete(ctx context.Context, bucket, key string) error

func (*CachedStore) Get

func (s *CachedStore) Get(ctx context.Context, bucket, key string) (io.ReadCloser, objstore.ObjectInfo, error)

func (*CachedStore) GetReaderAt

func (s *CachedStore) GetReaderAt(ctx context.Context, bucket, key string) (objstore.ReaderAtCloser, int64, error)

GetReaderAt serves random-access reads from the cache. On cache hit, it returns a zero-latency in-memory reader. On cache miss, it delegates to the inner store's ReaderAt (S3 range reads) to avoid downloading full files when only a few columns are needed. The cache is populated by Get() calls so subsequent pipeline queries on the same worker benefit from cached full files.

func (*CachedStore) Head

func (s *CachedStore) Head(ctx context.Context, bucket, key string) (objstore.ObjectInfo, error)

func (*CachedStore) List

func (s *CachedStore) List(ctx context.Context, bucket string, opts objstore.ListOptions) ([]objstore.ObjectInfo, error)

func (*CachedStore) MakeBucket

func (s *CachedStore) MakeBucket(ctx context.Context, bucket string) error

func (*CachedStore) Put

func (s *CachedStore) Put(ctx context.Context, bucket, key string, r io.Reader, size int64, contentType string) (string, error)

func (*CachedStore) PutIfMatch

func (s *CachedStore) PutIfMatch(ctx context.Context, bucket, key string, r io.Reader, size int64, contentType string, expectedETag string) (string, error)

type Config

type Config struct {
	NATSUrl       string
	WorkerID      string
	ClusterID     string // cluster this worker belongs to (for federated routing)
	MaxConcurrent int    // max concurrent tasks
	// MaxTaskDuration caps how long a single task may run before its context
	// is cancelled, guaranteeing its concurrency slot is eventually released
	// even if it wedges. The coordinator's per-task DeadlineUnixNano (the
	// query deadline) takes precedence when supplied; this is the floor for
	// the gRPC dispatch path when no deadline rides the dispatch. 0 = no cap.
	MaxTaskDuration  time.Duration
	CacheBytes       int64  // local LRU cache size
	MemoryBudget     int64  // per-task memory budget in bytes (0 = unlimited, no spill); used as legacy fallback when SharedPoolBudget is unset
	SharedPoolBudget int64  // worker-wide memory pool in bytes (0 = derived as MemoryBudget*MaxConcurrent). All concurrent tasks Reserve against this pool; spill triggers fire on cumulative worker pressure.
	SpillDir         string // directory for spill files (default: os temp dir)
	// DrainTimeout bounds Drain(): how long to wait for in-flight tasks and
	// then pending stage-output uploads before escalating to a hard stop.
	// 0 = unbounded (the platform's kill timeout — e.g. the Kubernetes
	// termination grace period — is the backstop).
	DrainTimeout     time.Duration
	ResultStoreBytes int64 // in-memory result store capacity (0 = disabled, results go to S3)

	// Reservoirs is the Phase-3 system-reservoir registry. When non-nil, the
	// worker registers its cache + result-store reservoirs against it and wires
	// it into the shared SpillManager for ACCOUNTING (Available()/drift). nil =
	// the legacy static-budget path.
	Reservoirs *memory.ReservoirRegistry
	// FloatingBudgetActive enables the deploy-gated floating spill threshold.
	// Default false: ShouldSpillFor stays on the tuned static 40%/90% path even
	// with reservoirs wired (the floating budget is mmap-blind until Phase-4
	// RSS-sampling).
	FloatingBudgetActive bool

	// PeerWireCompression s2-compresses raw WSHF payloads on outgoing
	// peer-exchange streams (docs/design/peer-wire-compression.md): the
	// wire carries a standard WSHC envelope every consumer already
	// decodes. ~20% fewer peer-stream bytes for ~1 core-GB/s of producer
	// CPU. Default false pending SF100 validation.
	PeerWireCompression bool
	// AsyncScratchPurge defers per-query stage-cache file deletion to a
	// paced background janitor instead of unlinking inline on the
	// query-complete broadcast handler (docs/design/async-scratch-purge.md
	// — the SF100 Q22/Q14/Q11 straggler-tail fix). False = inline unlinks
	// (pre-fix behavior, A/B kill switch).
	AsyncScratchPurge bool
	// StreamingShuffleRead decodes WSHF/WSHC exchange inputs directly from
	// the peer/S3 byte stream instead of staging whole files to NVMe +
	// mmap first (docs/design/exchange-streaming-consumption.md §3 D1).
	// Default on at the CLI (SF100-validated); false is the kill switch.
	StreamingShuffleRead bool

	// ScanDecodeAhead decodes parquet row groups ahead of scan consumption
	// with a bounded window instead of one group per pull
	// (docs/design/scan-decode-pipelining.md). Default false pending SF100
	// validation. ScanDecodeAheadBytes bounds decoded-but-unconsumed bytes
	// per scan source; <= 0 selects the engine default (256 MiB).
	ScanDecodeAhead      bool
	ScanDecodeAheadBytes int64

	// ShuffleDecodeAhead fans WSHF chunk decode out to CPU-token-budgeted
	// workers behind the streaming reader's scanner, with strict in-order
	// delivery (docs/design/shuffle-decode-ahead.md — the q08/q09 probe
	// width-plateau fix). Default on at the CLI; false is the kill switch
	// restoring the serial streaming reader.
	ShuffleDecodeAhead bool

	// DecodedCacheBytes bounds the worker-lifetime decoded-chunk cache:
	// decoded base-table parquet column chunks reused across queries and
	// runs instead of re-paying zstd decompress + decode
	// (docs/design/decoded-rowgroup-cache.md). 0 (the default) disables
	// the cache entirely — the kill switch. Registered as a hard system
	// reservoir and as a relief target (evicts before operators spill).
	DecodedCacheBytes int64

	// MmapRelief enables the Phase-5 MADV_DONTNEED relief of cold mmap'd cache
	// files under heap pressure. Default false: fully dormant (no region
	// tracking, no per-Next cost, no syscall). Deploy-gated.
	MmapRelief bool
	// MmapReliefThresholdMB is the TOTAL process RSS ceiling in MB at/above
	// which relief MADV_DONTNEEDs the coldest mmap to bring RSS back down. Tune
	// below the worker's cgroup memory.max so relief has headroom (e.g. ~16000
	// on a ~20 GB per-proc envelope). 0 leaves the default. Only meaningful when
	// MmapRelief is true.
	MmapReliefThresholdMB int64

	// BoundedDirtyWrites enables windowed sync_file_range writeback (plus
	// FADV_DONTNEED on spill-class files) for all large sequential disk
	// writes, capping each writer's dirty page-cache footprint so kernel
	// reclaim stops evicting the mmap'd cache pages concurrent tasks are
	// walking (see internal/engine/diskio). Default false: writes rely on
	// kernel writeback as before. Deploy-gated.
	BoundedDirtyWrites bool

	// PeerListenAddr enables the streaming-exchange PeerExchange server
	// (worker→worker shuffle fetches, docs/design/streaming-exchange.md)
	// on the given listen address (e.g. ":9095"; ":0" in tests). Empty
	// (the zero value) = no peer server: this worker serves no fetches
	// and advertises no peer address, so coordinators never hint at it.
	PeerListenAddr string
	// PeerAdvertiseAddr overrides the peer address carried in heartbeats.
	// Empty = derived from the bound listener (specific IP, else the first
	// non-loopback unicast IPv4).
	PeerAdvertiseAddr string

	// MorselWorkers controls intra-fragment parallel pipeline consumers per
	// task (morsel-driven execution, docs/design/morsel-execution.md). 0
	// (zero value) and 1 = serial, today's behavior — dormant-safe. -1 =
	// auto: width adapts to fragment input size and idle CPU tokens. N>1 =
	// fixed width of N (bypasses the size gate; testing/benchmark knob).
	// Extra consumers are bounded worker-wide by a CPU-token pool sized
	// GOMAXPROCS−2 so concurrent tasks cannot oversubscribe cores.
	MorselWorkers int
}

Config holds worker configuration.

func DefaultConfig

func DefaultConfig() Config

type Executor

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

Executor dispatches task types to the appropriate execution logic.

func NewExecutor

func NewExecutor(store objstore.Store, cache *LRUCache, js jetstream.JetStream) *Executor

NewExecutor creates a new task executor.

func (*Executor) CPUTokenAdmissionStats

func (e *Executor) CPUTokenAdmissionStats() (capacity, reserve, admits, bypasses, holdbacks int64)

CPUTokenAdmissionStats reports the worker pool's decode-class admission counters: pool capacity, the reserved decode floor, decode-class tokens granted, how many of those were granted while morsel consumers were queued (the admissions the old strict-FIFO policy refused — the direct measure of the fix), and how often a release was held back to keep the floor reachable. Read alongside scan/shuffle token_stall_ms, which is what these are meant to move.

func (*Executor) DegradedTaskBudget

func (e *Executor) DegradedTaskBudget() int64

DegradedTaskBudget returns the reduced spill budget a poison-suspect retry runs under on this executor, or 0 when no pool is configured (no degradation possible — the quarantine rung still applies).

func (*Executor) Execute

func (e *Executor) Execute(ctx context.Context, task distributed.Task, workerID string) distributed.ResultNotification

Execute runs a task and returns the result notification.

func (*Executor) HeapDrift

func (e *Executor) HeapDrift(heapInuse int64) int64

HeapDrift returns the Phase-4 accounting drift in bytes for the supplied HeapInuse sample — HeapInuse − (operator owned + reservoir actual) — or 0 when no shared spill manager is configured. Observability only.

func (*Executor) PeerFetchFallthroughs

func (e *Executor) PeerFetchFallthroughs() int64

PeerFetchFallthroughs returns how many hinted fetches failed over to S3.

func (*Executor) PeerFetchHits

func (e *Executor) PeerFetchHits() int64

PeerFetchHits returns how many input files were served via peer fetch.

func (*Executor) ResolveShuffleFile

func (e *Executor) ResolveShuffleFile(ctx context.Context, _, key, token string) (string, error)

ResolveShuffleFile implements dataplane.ShuffleFileResolver: validate the fetch token against the one this worker's own tasks for the root query carried, then look the key up in the LocalStageCache. Called by the PeerServer for every incoming FetchShuffle. The request's queryID is advisory — the key's "queries/<id>/" prefix is the identity.

A worker that never executed a task for the root query has no token recorded and denies the fetch — it also could not hold the file, so the consumer loses nothing by falling through to S3.

func (*Executor) ScanBackingStats

func (e *Executor) ScanBackingStats() (hits, misses, claimed int64)

ScanBackingStats returns the row-group backing-reuse counters (docs/design/scan-output-backing-reuse.md): decodes served from a released backing, decodes that minted one, and releases refused by a Detach claim.

func (*Executor) ScanDecodeAheadDecodeSpans

func (e *Executor) ScanDecodeAheadDecodeSpans() (ns, bytes int64)

ScanDecodeAheadDecodeSpans returns the total wall time (ns) inside row-group decode calls and the projected compressed bytes decoded — the inline-fault discriminator the stall counters cannot see.

func (*Executor) ScanDecodeAheadPrunedGroups

func (e *Executor) ScanDecodeAheadPrunedGroups() int64

ScanDecodeAheadPrunedGroups returns the row groups skipped by dynamic- filter pruning at the iterator layer (bloom + range combined).

func (*Executor) ScanDecodeAheadStallNs

func (e *Executor) ScanDecodeAheadStallNs() (windowFullNs, pressureNs, tokenNs, ledgerNs int64)

ScanDecodeAheadStallNs returns the total blocked time (ns) behind the four stall counters of ScanDecodeAheadStats, same order.

func (*Executor) ScanDecodeAheadStats

func (e *Executor) ScanDecodeAheadStats() (groups, windowFulls, pressureStalls, tokenStalls, ledgerStalls int64)

ScanDecodeAheadStats returns the decode-ahead counters: row groups decoded ahead, worker stalls on a full window, admissions refused under heap pressure, per-group admissions deferred for lack of a cpu token, and admissions denied by the shared memory pool ledger (each affected group still decodes, serially at worst).

func (*Executor) SetBaseTableCache

func (e *Executor) SetBaseTableCache(c *objstore.BaseTableCache)

SetBaseTableCache attaches the store stack's base-table cache layer so incoming peer fetches can be served from it. nil (default) rejects base-table fetches with NotFound.

func (*Executor) SetBaseTableOwnership

func (e *Executor) SetBaseTableOwnership(owns func(key string) bool)

SetBaseTableOwnership wires the rendezvous-ownership check the owner read-through path consults (worker wiring, pre-Start — read without synchronization on every base-table fetch). owns receives the bare object key, matching the hash the coordinator and the consumer tier use for placement. nil (default) disables read-through.

func (*Executor) SetDataPlaneClient

func (e *Executor) SetDataPlaneClient(c *dataplane.Client)

SetDataPlaneClient enables gRPC result streaming for gather sinks constructed by this executor. When set and Connected, each gatherReplySink prefers the gRPC stream over NATS Publish. nil reverts to NATS-only delivery.

func (*Executor) SetDecodedCache

func (e *Executor) SetDecodedCache(c *scan.DecodedChunkCache)

SetDecodedCache attaches the worker-lifetime decoded-chunk cache (docs/design/decoded-rowgroup-cache.md) consulted by parquet scan sources. When the shared SpillManager exists it also registers the cache as an AccountedOperator so RequestRelief can shed cached bytes — eviction, the cheapest relief in the process, runs before any operator pays a real spill. Same setter-order self-healing as SetReservoirs. Call before any task executes; nil leaves scans uncached.

func (*Executor) SetFloatingBudgetActive

func (e *Executor) SetFloatingBudgetActive(active bool)

SetFloatingBudgetActive propagates the deploy-gated floating-threshold flag to the shared SpillManager. Default false keeps ShouldSpillFor static.

func (*Executor) SetLocalStageCache

func (e *Executor) SetLocalStageCache(c *LocalStageCache)

SetLocalStageCache attaches a same-worker local-disk stage-output cache. Producers register their local spill files in it after upload succeeds; consumers consult it before falling back to KV/S3. Lifecycle is driven by query-complete / cancel signals from the coordinator.

func (*Executor) SetLogger

func (e *Executor) SetLogger(l *slog.Logger)

SetLogger sets the executor's logger.

func (*Executor) SetMemoryBudget

func (e *Executor) SetMemoryBudget(budget int64, spillDir string)

SetMemoryBudget configures the per-task memory budget and the spill directory. For backward compatibility it also initializes a shared pool of the same size, so existing callers that pass a single budget continue to get cooperative spill across concurrent tasks. Callers that want a different pool size (typically larger than per-task budget) should call SetSharedPoolBudget afterward to override.

func (*Executor) SetMetrics

func (e *Executor) SetMetrics(m *metrics.Metrics)

SetMetrics attaches Prometheus metrics for spill/memory tracking.

func (*Executor) SetMorselWorkers

func (e *Executor) SetMorselWorkers(n int)

SetMorselWorkers configures intra-fragment parallel pipeline consumers (morsel-driven execution, docs/design/morsel-execution.md). 0 and 1 = serial (today's behavior); -1 = auto (width adapts to fragment input size and idle CPU tokens); N>1 = fixed width of N. Must be called before the worker starts executing tasks.

func (*Executor) SetNATSConn

func (e *Executor) SetNATSConn(nc *nats.Conn)

SetNATSConn attaches a NATS connection used by Gather tasks to stream batches back to the coordinator's reply subject (and by the async upload manager for UploadComplete notifications).

func (*Executor) SetPeerClient

func (e *Executor) SetPeerClient(c *dataplane.PeerClient)

SetPeerClient attaches the outbound fetch client used by the Tier-1.5 peer read path. nil (default) disables peer fetches — hints are ignored.

func (*Executor) SetReservoirs

func (e *Executor) SetReservoirs(rr *memory.ReservoirRegistry)

SetReservoirs wires the system-reservoir registry into the executor and the shared SpillManager for ACCOUNTING. It does NOT activate the floating spill threshold (see SetFloatingBudgetActive). Call before any task executes.

func (*Executor) SetResultKV

func (e *Executor) SetResultKV(kv jetstream.KeyValue)

SetResultKV attaches a NATS KV store for cross-worker inter-stage result transfer. Results below natsKVResultThreshold are stored here instead of S3, reducing inter-stage latency from ~500ms (S3 round-trip) to ~10ms (NATS KV).

func (*Executor) SetResultStore

func (e *Executor) SetResultStore(rs *ResultStore)

SetResultStore attaches an in-memory result store for inter-stage result passing. When a reservoir registry is wired, it registers the result store as a hard reservoir backed by the store's live UsedBytes accessor so its occupancy (≈496 MB at SF100) feeds Available()/drift. Requires SetReservoirs to have run first.

func (*Executor) SetScanDecodeAhead

func (e *Executor) SetScanDecodeAhead(on bool, windowBytes int64)

SetScanDecodeAhead enables decode-ahead on parquet scan sources (--scan-decode-ahead). windowBytes <= 0 selects the default. Call before Worker.Start.

func (*Executor) SetSharedPoolBudget

func (e *Executor) SetSharedPoolBudget(budget int64)

SetSharedPoolBudget creates the worker-wide memory pool that all concurrent tasks Reserve against. Operators (HashJoin build, sort run accumulation, hash aggregate state) cooperatively spill when the pool fills, regardless of which task is holding the bytes. Matches the Trino MemoryPool / Spark ExecutionMemoryPool model.

Pool budget should be the FULL worker envelope (after cache reservation), not a per-task slice. With 32GB physical RAM and a 24GB GOMEMLIMIT, pool budget is roughly 21GB (envelope − cache).

Calling this with budget<=0 disables the shared pool and falls back to per-task tracking via SetMemoryBudget.

func (*Executor) SetShuffleDecodeAhead

func (e *Executor) SetShuffleDecodeAhead(on bool)

SetShuffleDecodeAhead enables chunk-parallel decode on WSHF shuffle inputs (--shuffle-decode-ahead). Call before Worker.Start.

func (*Executor) SetStreamingShuffleRead

func (e *Executor) SetStreamingShuffleRead(on bool)

SetStreamingShuffleRead enables streaming decode of shuffle inputs (--streaming-shuffle-read). Call before Worker.Start.

func (*Executor) SharedPoolStats

func (e *Executor) SharedPoolStats() (used, budget int64)

SharedPoolStats returns (used, budget) bytes for the worker-wide memory pool, or (0, 0) if no pool is configured. Used by the worker heartbeat loop to publish pool pressure for coord-side dispatch backpressure.

func (*Executor) ShuffleDecodeAheadStats

func (e *Executor) ShuffleDecodeAheadStats() (chunks, windowFullNs, tokenNs, pressureNs, stageNs, decodeNs, donated, preadNs, indexedFiles int64)

ShuffleDecodeAheadStats returns the shuffle decode-ahead markers: chunks decoded ahead, the parked/spent spans (ns) per class — window-full, token, pressure on the admission side; stage (the serial scanner walk, the structural floor) and decode (worker time) on the throughput side — and tokens accepted via producer donation (§2.2; expect token stalls to fall as this rises).

func (*Executor) ShuffleFilePreadStats

func (e *Executor) ShuffleFilePreadStats() (files, bytes int64)

ShuffleFilePreadStats returns the read-staged local shuffle open counters: files opened via the WSHF-pread path and their total bytes.

func (*Executor) ShuffleIOStats

func (e *Executor) ShuffleIOStats() ShuffleIOSnapshot

ShuffleIOStats returns the per-tier shuffle-read transfer counters.

func (*Executor) ShuffleStreamStats

func (e *Executor) ShuffleStreamStats() (reads, fallbacks, skipResumes int64)

ShuffleStreamStats returns the streaming-shuffle-read counters: streaming opens, staged fallbacks, and batches skipped by fallbacks.

type LRUCache

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

LRUCache is a bounded cache for Parquet footers and column chunks.

func NewLRUCache

func NewLRUCache(maxBytes int64) *LRUCache

NewLRUCache creates a new LRU cache with the given byte limit.

func (*LRUCache) Get

func (c *LRUCache) Get(key string) ([]byte, bool)

Get retrieves a copy of the cached value. The returned slice is safe to mutate (e.g., decompress) without corrupting the cached data.

func (*LRUCache) GetRef

func (c *LRUCache) GetRef(key string) ([]byte, bool)

GetRef retrieves the cached value without copying. The caller MUST NOT modify the returned slice. Safe for read-only use (e.g., Parquet decoding) because Go's GC keeps the backing array alive even if the cache evicts the entry while the caller still holds a reference.

func (*LRUCache) Len

func (c *LRUCache) Len() int

Len returns the number of entries in the cache.

func (*LRUCache) Put

func (c *LRUCache) Put(key string, data []byte)

Put adds a value to the cache, evicting old entries if needed.

func (*LRUCache) Size

func (c *LRUCache) Size() int64

Size returns the current cache size in bytes.

type LocalStageCache

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

LocalStageCache maps (queryID, key) → localPath for stage outputs the worker wrote locally and can serve back without an S3 round-trip when a downstream task lands on the same worker.

In single-node / standalone mode every stage transition is same-worker, so hits are common and each one saves an upload + download of a partition shuffle file (~1–3s per partition at SF10). In distributed mode cache hits happen only when the JetStream consumer load-balances a downstream task onto the same worker — when it doesn't, the consumer simply misses and falls through to the existing KV/S3 path.

Producers register entries via Adopt(): the producer's local file is renamed into the cache's own per-query directory (so the producer's task-scoped spill dir can still be RemoveAll'd by deferred cleanup). Consumers consult the cache before falling through to KV/S3. CleanupQuery drops entries and unlinks files when the coordinator reports a query as complete or cancelled.

func NewLocalStageCache

func NewLocalStageCache(rootDir string) *LocalStageCache

NewLocalStageCache returns an empty cache that stores adopted files under rootDir/<queryID>/. Each producer's spill file is moved into this tree so it survives the producing task's spill cleanup.

Any rootDir contents from a prior worker process are wiped at construction time — they're necessarily orphaned (this process has no entries pointing at them), and leaving them would slowly fill the spill volume on workers that crash before publishing query-complete signals.

func (*LocalStageCache) Adopt

func (c *LocalStageCache) Adopt(queryID, key, srcPath string) string

Adopt moves srcPath into the cache's per-query directory and registers it under (queryID, key). The producer no longer owns the file after a successful Adopt — the cache will unlink it on CleanupQuery.

Returns the new path on success, or an empty string on failure (the caller should leave srcPath where it is and proceed without registering — the downstream consumer will simply fall through to S3/KV).

func (*LocalStageCache) CleanupQuery

func (c *LocalStageCache) CleanupQuery(queryID string) int

CleanupQuery drops all entries for queryID and disposes of the files. With asyncPurge the per-query directory is renamed into .trash (instant — open fds and mmaps stay valid across rename and unlink) and the janitor deletes it with pacing; otherwise files are unlinked inline, blocking the caller for the whole storm. Safe to call multiple times and against unknown query IDs.

func (*LocalStageCache) Count

func (c *LocalStageCache) Count() int

Count returns the total number of cached entries across all queries. Observability/test helper.

func (*LocalStageCache) Get

func (c *LocalStageCache) Get(queryID, key string) string

Get returns the local path for (queryID, key), or "" if absent.

func (*LocalStageCache) SetAsyncPurge

func (c *LocalStageCache) SetAsyncPurge(on bool)

SetAsyncPurge enables deferred scratch deletion (janitor goroutine, started lazily on first use). Call before the worker serves traffic (same contract as the other Set<X> setters); false = inline unlinks on the cleanup caller, the pre-2026-07-24 behavior and the A/B kill switch.

func (*LocalStageCache) SetLogger

func (c *LocalStageCache) SetLogger(l *slog.Logger)

SetLogger attaches the worker's structured logger for the purge ledger lines. nil-safe; without it slog.Default() is used.

type ResultStore

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

ResultStore holds intermediate stage results in memory to avoid S3 round-trips when stages execute on the same worker. Results are keyed by their S3 path (the path that would have been used) so the executor can transparently check the store before falling back to S3.

This is the primary performance optimization for standalone mode and co-located workers where all stages run on the same node. For a 3-stage query, this eliminates 4-6 S3 round-trips (serialize→upload→download→deserialize per stage).

Memory is bounded by maxBytes. When exceeded, new results spill to S3 as usual. Results are cleaned up per-query when the query completes.

func NewResultStore

func NewResultStore(maxBytes int64) *ResultStore

NewResultStore creates an in-memory result store with the given capacity. A maxBytes of 0 disables in-memory result passing.

func (*ResultStore) CleanupQuery

func (rs *ResultStore) CleanupQuery(queryID string)

CleanupQuery removes all results for a terminal query — completed OR cancelled. Both handlers call it: a cancelled query will never read its own stage outputs either.

func (*ResultStore) Count

func (rs *ResultStore) Count() int

Count returns the number of stored results.

func (*ResultStore) Evicted added in v0.18.17

func (rs *ResultStore) Evicted() int64

Evicted returns how many queries the TTL dropped — i.e. how many times the coordinator's cleanup broadcast never arrived. A non-zero and rising value is the signal that something is not publishing a terminal message.

func (*ResultStore) Get

func (rs *ResultStore) Get(path string) ([]byte, bool)

Get retrieves result data from memory. Returns nil, false if not found (caller should read from S3).

func (*ResultStore) MaxBytes

func (rs *ResultStore) MaxBytes() int64

MaxBytes returns the configured capacity (0 disables the store). Used as the resultstore reservoir's cap.

func (*ResultStore) Put

func (rs *ResultStore) Put(queryID, path string, data []byte) bool

Put stores result data in memory if there's capacity. Returns true if stored, false if the caller should write to S3 instead.

func (*ResultStore) UsedBytes

func (rs *ResultStore) UsedBytes() int64

UsedBytes returns the current memory usage.

type ShuffleIOSnapshot

type ShuffleIOSnapshot struct {
	LocalFiles, LocalBytes int64
	KVFiles, KVBytes       int64
	PeerFiles, PeerBytes   int64
	S3Files, S3Bytes       int64
}

ShuffleIOSnapshot is one point-in-time reading of the per-tier shuffle-read ledger.

type TaskProgress

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

TaskProgress is the worker-side counter operators bump from their hot loop to signal forward progress. The counter is read by the per-task heartbeat goroutine to (a) decide whether to extend JetStream AckWait and (b) publish TaskProgress messages to coord.

Operators reach this via exec.ProgressReporterFromContext (the engine layer doesn't depend on worker types). The exec interface has AddRows / AddBytes; this struct satisfies it.

func (*TaskProgress) AddBytes

func (p *TaskProgress) AddBytes(n int64)

AddBytes reports that n more bytes have been processed by this task. Used for IO-bound stages (scan, shuffle write) where bytes are a truer measure of forward progress than row count.

func (*TaskProgress) AddRows

func (p *TaskProgress) AddRows(n int64)

AddRows reports that n more rows have been processed by this task. Updates the lastUpdate timestamp; the heartbeat goroutine reads these values without needing to lock.

type Worker

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

func New

func New(cfg Config, store objstore.Store, nc *nats.Conn, js jetstream.JetStream, logger *slog.Logger) *Worker

New creates a new Worker.

func (*Worker) BeginDrain

func (w *Worker) BeginDrain()

BeginDrain marks the worker as draining without blocking: intake loops stop pulling new tasks and the next heartbeat advertises Draining=true so the coordinator excludes this worker from dispatch. Idempotent. The full shutdown sequence is Drain(); process owners typically select on DrainRequested() and call Drain when it fires.

func (*Worker) Drain

func (w *Worker) Drain()

Drain performs a graceful shutdown: stop pulling new tasks, let in-flight tasks run to completion, flush pending stage-output uploads to the object store (so consumers of this worker's outputs keep their durable fallback once the peer-exchange server goes away), then stop background services. Use this for zero-downtime rolling updates and Kubernetes pod termination (SIGTERM). Config.DrainTimeout bounds the whole sequence; on timeout the remaining work is aborted exactly like Stop.

func (*Worker) DrainRequested

func (w *Worker) DrainRequested() <-chan struct{}

DrainRequested returns a channel that is closed once drain has been requested from ANY source: BeginDrain, the NATS drain subject (the coordinator sends it on reap), or the /drain admin endpoint.

func (*Worker) Draining

func (w *Worker) Draining() bool

Draining returns true if the worker is in drain mode.

func (*Worker) PeerExchangeAddr

func (w *Worker) PeerExchangeAddr() string

PeerExchangeAddr exposes the advertised peer-exchange address ("" when peer serving is disabled or the listener isn't bound yet). Tests use it to bootstrap coordinator heartbeats ahead of the first real one.

func (*Worker) PeerFetchStats

func (w *Worker) PeerFetchStats() (hits, fallthroughs int64)

PeerFetchStats returns the streaming-exchange read-tier counters: inputs served via peer fetch, and hinted fetches that fell through to S3.

func (*Worker) ResultStore added in v0.18.17

func (w *Worker) ResultStore() *ResultStore

ResultStore returns the worker's in-memory stage-result store, or nil when --result-store is 0. Exported for the reclamation gate: what the store still holds after a terminal broadcast is the only observable that catches the CANCEL handler's missing cleanup (#818, ADR-0028).

func (*Worker) SetControlConn

func (w *Worker) SetControlConn(nc *nats.Conn)

SetControlConn provides a dedicated NATS connection for the heartbeat publish path. When set, heartbeats go through this connection instead of the data-plane nc. Callers should pass a separately-dialed *nats.Conn (NOT the same one as nc). Optional — heartbeat falls back to nc when unset, preserving the historical single-connection topology used by in-process tests.

func (*Worker) SetDataPlaneClient

func (w *Worker) SetDataPlaneClient(c *dataplane.Client)

SetDataPlaneClient enables the gRPC data plane for this worker. When set:

  • gather sinks prefer the gRPC stream over the NATS reply subject (Phase B behavior; safe even if dispatch stays on NATS),
  • the worker skips its JetStream Fetch loop and executes tasks pushed via TaskDispatch on the same stream (Phase C).

nil = legacy NATS-only path (default). Must be called before Start.

func (*Worker) SetMetrics

func (w *Worker) SetMetrics(m *metrics.Metrics)

SetMetrics attaches Prometheus metrics for spill/memory tracking.

func (*Worker) SetTelemetry

func (w *Worker) SetTelemetry(tp *telemetry.Provider)

SetTelemetry enables OpenTelemetry tracing on the worker.

func (*Worker) ShuffleStreamStats

func (w *Worker) ShuffleStreamStats() (reads, fallbacks, skipResumes int64)

ShuffleStreamStats returns the streaming-shuffle-read counters: streaming opens, staged fallbacks, batches skipped by fallbacks.

func (*Worker) Start

func (w *Worker) Start(ctx context.Context) error

Start begins the worker task loop and heartbeat.

func (*Worker) Stop

func (w *Worker) Stop()

Stop gracefully stops the worker. The shared per-cluster consumer is left in place so other workers can continue pulling tasks.

func (*Worker) UploadStats

func (w *Worker) UploadStats() (completed, cancelled, failed int64)

UploadStats returns the Phase-B background-upload counters: files landed, files cancelled by query completion, files abandoned after retries.

type WorkerProfile

type WorkerProfile struct {
	WorkerID  string `json:"worker_id"`
	CPU       []byte `json:"cpu,omitempty"`       // pprof CPU profile (gzip-compressed)
	Heap      []byte `json:"heap,omitempty"`      // pprof heap profile (gzip-compressed)
	Block     []byte `json:"block,omitempty"`     // pprof block profile; empty unless WADJET_BLOCK_PROFILE_RATE set
	Mutex     []byte `json:"mutex,omitempty"`     // pprof mutex profile; empty unless WADJET_MUTEX_PROFILE_FRACTION set
	Goroutine []byte `json:"goroutine,omitempty"` // pprof goroutine profile; always present — no runtime sampler required
}

WorkerProfile is the JSON envelope for profile data sent over NATS.

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL