diff --git a/Makefile b/Makefile index 38b896af73d73..967042d9fe2e6 100644 --- a/Makefile +++ b/Makefile @@ -1315,11 +1315,13 @@ install-static-check-tools: @go install github.com/apache/skywalking-eyes/cmd/license-eye@v0.4.0 .PHONY: static-check +GOLANGCI_LINT_CONCURRENCY ?= +GOLANGCI_LINT_CONCURRENCY_FLAG := $(if $(strip $(GOLANGCI_LINT_CONCURRENCY)),--concurrency $(strip $(GOLANGCI_LINT_CONCURRENCY))) static-check: config err-check $(CGO_OPTS) go vet $(GO_MODULE_MODE) -vettool=`which molint` ./... $(CGO_OPTS) license-eye -c .licenserc.yml header check $(CGO_OPTS) license-eye -c .licenserc.yml dep check - $(CGO_OPTS) golangci-lint run -v -c .golangci.yml ./... + $(CGO_OPTS) golangci-lint run -v $(GOLANGCI_LINT_CONCURRENCY_FLAG) -c .golangci.yml ./... fmtErrs := $(shell grep -onr 'fmt.Errorf' pkg/ --exclude-dir=.git --exclude-dir=vendor \ --exclude=*.pb.go --exclude=*_test.go --exclude=system_vars.go --exclude=Makefile) diff --git a/docs/design/analyze_stats_publication.md b/docs/design/analyze_stats_publication.md new file mode 100644 index 0000000000000..cef8392cc36d9 --- /dev/null +++ b/docs/design/analyze_stats_publication.md @@ -0,0 +1,488 @@ +# ANALYZE Statistics Publication and Plan-Cache Freshness + +- Status: mandatory design review pending (stateful cache/concurrency lifecycle) +- Tracking issue: [matrixorigin/matrixone#27728](https://github.com/matrixorigin/matrixone/issues/27728) +- Implementation PR: [matrixorigin/matrixone#27758](https://github.com/matrixorigin/matrixone/pull/27758) +- Last updated: 2026-08-28 + +## 1. Problem and evidence + +`ANALYZE TABLE` currently computes its SQL-visible result through a derived +aggregate query. Optimizer statistics, however, are owned by disttae's +process-local `GlobalStats` cache and are refreshed independently. A successful +ANALYZE therefore does not establish a boundary after which later statements on +the same CN must plan with the newly collected optimizer statistics. + +The observable failure is a long-lived session that continues to reuse both its +three-second statistics cache and a cached logical plan after ANALYZE. In the +reported Q35-shaped workload, a plan built before ANALYZE kept the stale +non-shuffle topology. Reconnecting appeared to fix the query only because a new +session discarded both caches; reconnecting is not a valid publication +contract. + +The implementation must also reject partial refreshes. Object statistics are +collected concurrently from S3 metadata. A missing/corrupt object, cancellation, +or executor shutdown must not be logged and then treated as a successful scan, +because that would publish a partially accumulated `StatsInfo` and invalidate +plans in favor of worse data. + +## 2. Scope + +This design covers: + +- current, physical tables analyzed outside an already-active user transaction; +- synchronous publication into the local disttae optimizer-statistics cache; +- invalidation of dependent session statistics and plan-cache entries on the + same CN process; +- serialization with automatic logtail-driven refreshes for the same physical + table; +- bounded process/session metadata and executor shutdown/error behavior. + +It intentionally does not provide: + +- cluster-wide or cross-CN invalidation; +- publication for historical snapshots, views, temporary tables, subscriptions, + or other relations that do not own current persistent optimizer statistics; +- publication of a derived query's uncommitted workspace view into the + committed, process-global statistics cache; +- a persistent generation, wire-protocol change, catalog change, or on-disk + format change; +- a guarantee that prepared statements whose physical compile lifecycle is + independently retained will be rebuilt by this first phase. + +## 3. Required invariants + +### 3.1 Safety + +1. A successful ANALYZE publication is all-or-nothing: every admitted object + task completes successfully before the engine cache is replaced. +2. For one physical `(account_id, table_id)`, explicit and automatic refreshes + publish in one serialized order. +3. The frontend generation advances only after the engine cache replacement + succeeds, and the issuing session is tagged with that exact generation. +4. A plan is reusable only if every statistics generation captured while + building it still equals the current generation. +5. Work that observed generation N cannot repopulate a session cache after N+1 + has been published. +6. Physical ownership is resolved before cache lookup. Tenant, system/cluster, + and publisher identities must not alias solely because table IDs match. +7. A failed automatic refresh cannot replace a previously published statistics + value. It may install a nil completion sentinel only when the table has no + prior cache entry, so first-read waiters terminate without losing last-good + state. + +### 3.2 Liveness and ownership + +1. Every frontend and engine admission token has exactly one effective release + owner on success, error, and cancellation. +2. Every concurrent object task is completed by exactly one owner: a worker + executes it, or executor shutdown rejects it. A rejected queued task must + still release the caller's completion barrier. +3. Admission and executor submission observe caller cancellation. The shared + traversal context also observes executor lifecycle cancellation, so shutdown + cannot leave a producer blocked on a full queue, a running S3 task using an + orphaned request, or a caller waiting for abandoned queued work. Async + cancellation callbacks are notification mechanisms, never terminal + predicates: refresh admission, zero-object fast paths, joined traversal, and + cache publication synchronously re-read the owning lifecycle context before + reporting or publishing success. Subscription and object I/O use a + refresh-local context that preserves request values/deadlines and is canceled + by either the request or the owner lifecycle. That linked context is created + only after fixed-stripe admission, bounding simultaneously registered owner + callbacks by the 64 refresh-admission tokens. +4. Unrelated tables normally remain parallel. A bounded hash-stripe collision + may serialize refresh control work but cannot affect query execution. +5. A synchronous first-read waiter may sleep only while the exact refresh + generation it enqueued is still owned by the current table subscription and + has at least one queued or running producer. Cache completion, context + cancellation, worker-lifecycle shutdown, producer exhaustion, and generation + removal are durable predicates checked while holding `GlobalStats.mu`; + `cond.Wait` atomically releases that same mutex. A notification is never + treated as the predicate. +6. A producer may create or capture a refresh generation only while a live + subscription owns its eventual cleanup. Failed queue admission may leave an + idle record for that live subscription, but it cannot create process-lifetime + metadata for an unsubscribed table and no waiter may depend on a rejected job. + +### 3.3 Boundedness + +1. Frontend publication admission uses 64 fixed stripes. +2. Engine refresh admission uses 64 fixed stripes. +3. The process-local table-generation registry retains at most 64K explicit + keys. Compaction advances a reset generation before clearing the map, so all + older missing-key labels become conservatively stale. +4. Session statistics remain bounded by the existing `StatsCache` policy; its + generation tags reset atomically with that cache. +5. Plan dependencies are bounded by the plan-cache capacity and by the number + of physical tables referenced by each cached plan. +6. Concurrent object work remains bounded by the existing worker count and + 2,048-entry executor queue. The fix must not add a channel/future allocation + per object on the successful collection path. +7. Engine statistics and refresh-scheduling metadata share the table cleanup + boundary. Once an unsubscribed table reaches `RemoveTid`, neither + `statsInfoMap` nor `updatingMu.updating` retains any key for that table ID, + and a late automatic or explicit refresh cannot recreate either entry or + write into a replacement generation. Every refresh, including the first + queued request, carries a non-nil scheduling-record pointer as its lifetime + token. Logtail's cache-existence check and token capture use the same cleanup + lock order, so removal cannot fall between those two producer steps. +8. An explicit refresh creates its table-lifetime token only after subscription + and catalog resolution succeed, while holding the subscription lifecycle + read lock. A failed subscription therefore retains no scheduling entry, and + unsubscribe cleanup cannot fall between validation and token capture. +9. A first statistics read whose subscription fails returns without enqueueing + automatic work. Retrying through the worker queue would create cache and + scheduling state before any subscription lifetime can own its cleanup. +10. Prefetch and synchronous first-read producers capture their scheduling + token while holding the subscription lifecycle read lock. A first-read also + requires the exact `subEntry` observed after subscription, so cleanup and a + replacement subscription cannot silently retarget old work. + +## 4. Identity and visibility + +The generation key is the physical owner: + +```text +(account_id, table_id) +``` + +The account is resolved using the same precedence as relation resolution: + +1. snapshot tenant, when applicable; +2. cluster/system ownership rules; +3. publication ownership where relation resolution selects the publisher; +4. explicit system-account overrides. + +Current publication is skipped for historical AS OF references, publication +consumers, temporary tables, views, and non-persistent relation kinds. These +paths preserve the legacy derived ANALYZE result and do not claim a current +engine-cache boundary. + +An ANALYZE whose statement starts with an active transaction also preserves the +legacy result without global publication. The derived SQL can see the +transaction workspace, whereas `GlobalStats` is committed-object state. Mixing +those visibility domains would make uncommitted data process-global. + +## 5. State and publication model + +### 5.1 End-to-end order + +For a publishable table, success follows this order: + +```text +derived ANALYZE query succeeds + -> acquire frontend table stripe + -> acquire engine table stripe + -> subscribe current partition and resolve current table definition + -> submit visible-object tasks + -> wait for every task result or rejection + -> atomically replace GlobalStats entry and wake engine waiters + -> commit engine refresh scheduling metadata + -> release engine stripe + -> advance frontend table generation + -> cache returned StatsInfo in issuing session under that generation + -> release frontend stripe + -> return ANALYZE success +``` + +There are two related linearization points: + +- the engine data publication point is replacement of `statsInfoMap[key]` after + all object tasks succeed; +- the frontend reuse boundary is advancement of the table generation while the + frontend publication stripe is still held. + +The frontend point intentionally follows the engine point. Before generation +advancement, an old plan may still run against the old generation. After it, +new cache admission and later cache hits must reject that old generation. + +### 5.2 Synchronous first-read wait protocol + +`GlobalStats.Get(sync=true)` uses a level-triggered predicate rather than an +edge-triggered wakeup contract: + +```text +subscribe and capture exact subEntry + -> under subscription RLock, capture/create exact updateRecord + -> register and enqueue a job carrying that updateRecord + -> lock GlobalStats.mu + -> if cache key exists (value or nil sentinel), return it + -> if caller context is done, return nil + -> under GlobalStats.mu -> updatingMu, if updateRecord is no longer current + or has no queued/running producer, return nil + -> cond.Wait (atomically publish waiter and release GlobalStats.mu) + -> re-evaluate all predicates +``` + +`RemoveTid` takes `GlobalStats.mu -> updatingMu`, removes both the cache entry +and generation, then broadcasts before releasing `GlobalStats.mu`. Therefore +cleanup either happens before the waiter checks the generation, in which case +the durable generation predicate terminates it, or after `cond.Wait` has +atomically registered the waiter, in which case the broadcast wakes it. There +is no broadcast-before-wait gap. A stale worker may reject its generation +silently; progress never depends on that stale worker issuing another wakeup. +Each enqueue attempt increments the generation's queued-producer count before +it can become visible to a waiter; a forced sender blocked on a full queue is +therefore also a live producer with caller and lifecycle cancellation. Worker +admission atomically transfers that count to `inProgress` or consumes a +coalesced/rejected job. Queue rollback and +every path that removes the final running producer broadcast after persisting +the zero-producer predicate. Thus a waiter cannot depend on a job that was +accepted by the channel but later abandoned by debounce, coalescing, or a +different caller's cancellation. + +Context cancellation uses the same rule: its callback acquires +`GlobalStats.mu` before broadcasting, so cancellation cannot fall between the +predicate check and waiter registration. A forced queue submission that is +rejected or canceled is never followed by a wait. +The `GlobalStats` lifecycle context is an independent terminal predicate for +both a forced queue submission and an already parked waiter. Stopping update +workers therefore cannot strand a request whose caller context remains live. + +### 5.3 Automatic refresh interaction + +Automatic logtail refresh and explicit ANALYZE share the engine stripe. An +automatic refresh commits its statistics entry and object-count/sampling +baseline before releasing that stripe. This prevents an older refresh from +overwriting the scheduling metadata of a newer explicit refresh. + +The automatic-refresh cache transition is last-good preserving. Success +replaces the cached value; failure leaves an existing entry untouched. When the +first automatic attempt fails, it installs a nil completion sentinel and wakes +synchronous waiters, preserving the existing `GlobalStats.Get` termination +contract without representing the failure as a newer publication. + +Frontend publication needs a separate stripe because it serializes the larger +engine-publication-plus-generation transaction. Without it, two concurrent +ANALYZE statements could publish engine results A then B but advance/cache their +frontend generations in the opposite order. + +Hash collisions deliberately trade rare refresh serialization for fixed memory. +They do not merge table identity: engine maps, generations, session tags, and +plan dependencies remain keyed by the full physical key. + +### 5.4 Plan and session-cache admission + +Planning records the first generation observed for each physical table. A +generation change during repeated reads makes the completed plan ineligible for +cache admission. Cache lookup compares all recorded dependencies against the +current process registry; only dependent plans are removed. + +The session statistics cache uses a `(physical key, generation)` tag in addition +to its historical table-ID lookup. A slow storage read can cache its result only +if the generation is unchanged on completion. Publication installs the exact +engine result into the issuing session under the newly advanced generation. + +## 6. Failure, cancellation, and shutdown + +The terminal behavior is: + +| Failure phase | Engine cache | Engine metadata | Frontend generation/session cache | SQL result | +| --- | --- | --- | --- | --- | +| derived query fails | unchanged | unchanged | unchanged | error | +| frontend admission canceled | unchanged | unchanged | unchanged | cancellation | +| explicit subscribe/catalog resolution fails | unchanged | no generation retained before a cleanup owner exists | unchanged | error | +| initial statistics-read subscription fails | unchanged | no automatic work or generation admitted | unchanged | no statistics | +| prefetch outside a live subscription | unchanged | no generation admitted | unchanged | rejected | +| forced first-read queue admission canceled/rejected | unchanged | live-subscription record may remain for reuse and eventual cleanup; no waiter admitted | unchanged | no statistics | +| shared automatic producer canceled before publication | unchanged | queued/running producer count reaches zero and wakes all coalesced waiters | unchanged | no statistics | +| statistics worker lifecycle stops with queued work | unchanged | waiter terminates on lifecycle predicate; process-owned record dies with `GlobalStats` | unchanged | no statistics | +| subscription cleanup before first-read waiter parks | removed | exact generation removed; durable predicate terminates waiter | unchanged | no statistics | +| subscription cleanup after first-read waiter parks | removed | exact generation removed; cleanup broadcast terminates waiter | unchanged | no statistics | +| automatic subscribe/catalog resolution fails | last-good entry retained; nil completion sentinel only when absent | failed generation closed | unchanged | not an ANALYZE result | +| task submission canceled/rejected | unchanged | failed generation closed | unchanged | error | +| object task fails or is canceled | unchanged; local partial object discarded | failed generation closed | unchanged | error | +| engine cache publication succeeds | replaced | committed before engine release | generation must then advance while frontend token is held | success only after remaining steps | + +The shared object executor owns queued tasks only after successful admission. +On normal operation, a worker removes a task and executes its callback. During +executor shutdown, new submissions fail, the executor lifecycle cancels the +shared context used by already-running callbacks, workers join those callbacks, +and a shutdown owner rejects every task left in the queue. Execution and +rejection both invoke the caller-provided completion callback exactly once. + +The per-refresh error accumulator records the first non-nil task or rejection +error and waits for all admitted work before returning. Waiting is required +because callbacks mutate a refresh-local accumulator; returning early would let +old work race a discarded accumulator. Callback I/O receives the request +context, so cancellation terminates the expensive work without polling or +sleeps. The executor lifecycle context is the durable predicate; the shared task +context is only its cancellation-delivery path. Traversal uses a +check/register/check sequence around that delivery callback and re-reads both +contexts after joining all admitted work. Executor shutdown therefore remains a +failed traversal even if callback dispatch is delayed or a running callback +ignored cancellation and returned `nil`. The enclosing refresh applies the same +owner-lifecycle predicate at admission and cache publication, including the +zero-object path that does not traverse objects at all. A successful refresh +linearizes at the final lifecycle and generation checks performed while the +publication lock is held. + +Cancellation and deadline errors remain cancellation/deadline errors at the +public refresh boundary. Other object/metadata failures may be wrapped with +table-refresh context, but must remain a failed publication. + +Retry starts a fresh refresh attempt. The engine reserves only an internal +table-lifetime token before work starts; the frontend reuse generation is not +advanced until engine publication succeeds, so a failed retry creates no +consumer-visible generation gap. + +## 7. Compatibility and operations + +`engine.StatsRefresher` is an optional in-process capability. Engines that do +not implement it preserve legacy ANALYZE behavior. There is no persisted or wire +state, so mixed binaries do not need protocol negotiation: each CN invalidates +only its own caches and a restart naturally starts a new local generation epoch +with empty sessions. + +Rollback removes the optional publication call and process-local generation +tracking. The disttae statistics cache remains compatible with its existing +automatic refresh path. No data/catalog migration, backup/restore action, or +downgrade procedure is required. + +Operational diagnosis should distinguish: + +- derived ANALYZE query duration; +- synchronous disttae refresh duration and S3 request counts; +- refresh failures by subscribe, catalog, object I/O, cancellation, or executor + shutdown; +- plan-cache invalidations caused by statistics generations; +- generation-registry compactions and refresh-stripe wait time. + +The initial implementation can reuse current refresh logging and metrics, but +the above dimensions are the required observability target before enabling a +cross-CN extension. + +## 8. Performance and capacity model + +The common TP path with no recorded statistics dependency performs no new +generation-map read. A dependent cache hit takes one process-local read lock and +O(number of referenced physical tables) comparisons, with no allocations. The +local Apple M4 evidence at implementation revision `cb2327fd90` measured +1.773-1.785 ns for zero dependencies, 34.95-36.45 ns for one, 47.20-48.10 ns +for four, and 121.5-123.1 ns for sixteen, all with zero allocations. The waiter +repair does not touch that path; its producer accounting adds one integer under +the existing scheduling mutex only when a statistics refresh is enqueued. +These values are directional microbenchmark evidence, not a production latency +SLO. + +ANALYZE adds a synchronous disttae object-metadata scan after its derived query. +This increases ANALYZE latency and S3 reads but moves the cost off the normal TP +execution path. The reported 10M-row Q35-shaped case completed that refresh in +approximately 181 ms and changed the next plan from stale non-shuffle to a +16-way hash plan. Production acceptance must continue to compare object count, +sampling mode, S3 requests, and end-to-end ANALYZE latency. + +The executor repair keeps the existing closure and wait-group shape. It adds a +constant number of lifecycle reads per refresh, a success-path branch and error +accumulator, plus executor lifecycle bookkeeping; it does not add a check, result +channel, allocation, or goroutine per object. The refresh-local owner callback is +registered only after fixed-stripe admission, so at most 64 are active. Focused +benchmarks or allocation tests are required if the final implementation changes +that property. + +## 9. Alternatives + +### A. Keep TTL-only session caching + +This is the status quo. It is simple and cross-CN-neutral, but cannot define a +synchronous ANALYZE boundary and leaves cached plans stale indefinitely. It is +rejected for correctness. + +### B. Flush every session and plan cache on ANALYZE + +A process-wide flush is easy to reason about but turns one table's maintenance +into unrelated TP plan rebuilds and requires enumerating/live-coordinating all +sessions. It has a larger latency and availability blast radius and is rejected. + +### C. Rely only on the engine statistics-map replacement + +This fixes fresh statistics reads but not a session's three-second cache or +logical plan cache. It also cannot fence a slow old read from repopulating a +session cache. It is insufficient. + +### D. Broadcast generations across CNs + +Cross-CN invalidation is the desired broader semantic but requires a distributed +ordering, retry, restart, compatibility, and bounded replay design. Adding it to +this repair would materially enlarge the failure surface. The selected first +phase is explicitly CN-local and leaves this as a separately designed feature. + +### E. Selected: CN-local per-table generation after synchronous engine refresh + +This gives a precise local boundary, invalidates only dependent state, requires +no persistent/wire change, and keeps ANALYZE cost out of the normal TP path. Its +tradeoffs are local-only scope, two bounded admission layers, and dependency +checks on affected plan-cache hits. + +## 10. Verification map + +| Contract | Deterministic evidence | +| --- | --- | +| successful publication and exact returned object | engine publication-boundary UT | +| failed engine refresh does not advance/cache | frontend publisher failure UT | +| one successful and one failed object task rejects partial stats | concurrent visible-object UT | +| pre-canceled, in-flight canceled, and shutdown-rejected work terminates | executor/visible-object cancellation UT | +| shutdown cannot become success when an in-flight callback returns nil | executor-lifecycle traversal UT | +| same-table refresh order; unrelated-table concurrency | frontend and engine admission race UT | +| failed automatic refresh preserves last-good stats and completes an absent first generation | injected subscribe-failure state-transition UT | +| table cleanup reclaims both statistics and refresh-scheduling entries; first-queued, late automatic, and late explicit work cannot recreate them or target a replacement lifetime | `RemoveTid` ownership/generation UT | +| failed explicit subscription creates no ownerless scheduling generation | injected subscribe-failure UT | +| failed initial-read subscription queues no ownerless automatic generation | injected subscribe-failure UT | +| synchronous first read returns the exact value published by its accepted producer | producer-publication UT | +| prefetch outside a live subscription creates no ownerless generation | focused ownership UT | +| cleanup before synchronous waiter registration cannot lose a wake | queue/admission phase-barrier UT | +| cleanup after synchronous waiter registration terminates the wait | condition-registration phase-barrier UT | +| caller cancellation after synchronous waiter registration terminates the wait | condition-registration phase-barrier UT | +| coalesced waiter terminates when another caller's producer is canceled at refresh admission | producer-transfer phase-barrier UT | +| synchronous waiter terminates when the statistics worker lifecycle stops | worker-lifecycle phase-barrier UT | +| replacement subscription rejects work captured from the old `subEntry` | subscription-generation UT | +| slow generation-N read cannot overwrite N+1 | session-cache race UT | +| plan build spanning publication is not cached | plan-cache generation UT | +| physical account/view/temporary/transaction rules | focused frontend table-driven UT and ANALYZE BVT | +| no-dependency and 1/4/16-dependency cache-hit cost | allocation/latency benchmarks | +| SQL-visible existing-session plan changes after ANALYZE | recorded real-service validation and explain-plan assertion | + +Every new concurrency test uses explicit phase barriers and an outer timeout +only as a hang guard. The changed disttae and frontend packages require focused +normal tests, focused adaptive race stress, and one owning-package race run when +the local CGo environment can execute their linked test binaries. + +## 11. Rollout and decision log + +The initial rollout is behavior-on for disttae engines implementing the optional +capability. It fails closed: inability to complete the refresh returns an error +and leaves previous generations/caches valid. Legacy engines and excluded table +classes retain their previous behavior. + +Decision log: + +- Use physical `(account_id, table_id)` identity; names are lookup inputs, not + generation identity. +- Publish only committed current-table state; do not globalize an active + transaction's derived result. +- Keep the first phase CN-local; cross-CN publication requires another design. +- Use fixed stripes and conservative 64K compaction rather than unbounded + per-table locks/generations. +- Wait for all admitted object work after the first failure so no callback can + outlive its refresh-local accumulator. +- Preserve no per-object future/channel allocation in the successful scan path. +- Preserve the last successful automatic-refresh value on failure; use a nil + sentinel only to complete an otherwise absent first generation. +- Make `RemoveTid` the common lifetime owner for published engine statistics + and per-table refresh-scheduling metadata; every automatic and explicit + refresh requires a non-nil update record as a lifetime token, so old work + cannot recreate cleanup-owned state or target a replacement lifetime. +- Treat condition-variable broadcasts only as hints. The cache/context/exact + generation predicates are checked under the condition mutex, and every + producer generation is captured under a live subscription cleanup owner. +- Treat lifecycle callbacks only as cancellation delivery. Admission, joined + traversal, zero-object completion, and publication re-read the request and + owner contexts as durable predicates; linked owner callbacks are created only + after fixed-stripe admission. +- Give each admitted automatic refresh one terminal owner that commits the + cache, advances metadata only when that commit succeeds, and finally releases + admission in that order. + +Open approval item: an independent reviewer must approve this exact design +revision before the implementation is considered deliverable. There are no +known unresolved correctness decisions in the proposal itself. diff --git a/pkg/frontend/compiler_context.go b/pkg/frontend/compiler_context.go index 55c70200a86e7..07ef95e2ac76d 100644 --- a/pkg/frontend/compiler_context.go +++ b/pkg/frontend/compiler_context.go @@ -162,6 +162,23 @@ func (tcc *TxnCompilerContext) GetStatsCache() *plan2.StatsCache { return tcc.execCtx.ses.GetStatsCache() } +func (tcc *TxnCompilerContext) getStatsCacheVersion( + key optimizerStatsTableKey, +) (*Session, *plan2.StatsCache, uint64) { + tcc.mu.Lock() + feSes := tcc.execCtx.ses + txnWrapper, _ := tcc.tcw.(*TxnComputationWrapper) + tcc.mu.Unlock() + if ses, ok := feSes.(*Session); ok { + cache, version := ses.getStatsCacheWithVersion(key) + if txnWrapper != nil { + txnWrapper.recordOptimizerStatsVersion(key, version) + } + return ses, cache, version + } + return nil, feSes.GetStatsCache(), 0 +} + func InitTxnCompilerContext(db string) *TxnCompilerContext { return &TxnCompilerContext{dbName: db} } @@ -252,32 +269,61 @@ func (tcc *TxnCompilerContext) GetAccountId() (uint32, error) { // ResolveViewDependencyAccount returns the account whose catalog namespace was // used to resolve a View dependency. Keep the override order aligned with -// getRelation: snapshot tenant, subscription publisher, then relations that -// are always read from the system account. +// getRelation: snapshot tenant, cluster-table name override, subscription +// publisher, then relations that are always read from the system account. func (tcc *TxnCompilerContext) ResolveViewDependencyAccount( obj *plan2.ObjectRef, tableDef *plan2.TableDef, snapshot *plan2.Snapshot, ) (uint32, error) { + return tcc.resolvePhysicalObjectAccount(obj, tableDef, snapshot), nil +} + +// resolvePhysicalObjectAccount keeps statistics and view dependencies aligned +// with the account context used by getRelation. The identity must be resolved +// before consulting any cache so cached data and its generation share one key. +func (tcc *TxnCompilerContext) resolvePhysicalObjectAccount( + obj *plan2.ObjectRef, + tableDef *plan2.TableDef, + snapshot *plan2.Snapshot, +) uint32 { accountID := tcc.execCtx.ses.GetAccountId() if snapshot != nil && snapshot.Tenant != nil { accountID = snapshot.Tenant.TenantID } - if obj.PubInfo != nil { - accountID = uint32(obj.PubInfo.TenantId) - } - dbName, tableName := obj.SchemaName, obj.ObjName - if dbName == "" { + var dbName, tableName string + if obj != nil { + dbName, tableName = obj.SchemaName, obj.ObjName + } + if dbName == "" && tableDef != nil { dbName = tableDef.DbName } - if tableName == "" { + if tableName == "" && tableDef != nil { tableName = tableDef.Name } - if isClusterTable(dbName, tableName) || ShouldSwitchToSysAccount(dbName, tableName) { + if isClusterTable(dbName, tableName) { accountID = sysAccountID } - return accountID, nil + // getRelation applies publication ownership after the generic cluster-table + // name rule, so the publisher remains the physical owner in that overlap. + if obj != nil && obj.PubInfo != nil { + accountID = uint32(obj.PubInfo.TenantId) + } + if ShouldSwitchToSysAccount(dbName, tableName) { + accountID = sysAccountID + } + return accountID +} + +func (tcc *TxnCompilerContext) optimizerStatsKey( + obj *plan2.ObjectRef, + snapshot *plan2.Snapshot, +) optimizerStatsTableKey { + return optimizerStatsTableKey{ + accountID: tcc.resolvePhysicalObjectAccount(obj, nil, snapshot), + tableID: uint64(obj.Obj), + } } func (tcc *TxnCompilerContext) GetAccountName() string { @@ -1125,10 +1171,12 @@ func (tcc *TxnCompilerContext) Stats(obj *plan2.ObjectRef, snapshot *plan2.Snaps }() tableID := uint64(obj.Obj) + statsKey := tcc.optimizerStatsKey(obj, snapshot) + ses, statsCache, statsVersion := tcc.getStatsCacheVersion(statsKey) // Fast path: return cached result if visited within 3 seconds AND stats is valid // Stats is valid if AccurateObjectNumber > 0 (meaning we have real data) - if w := tcc.GetStatsCache().Get(tableID); w.Exists() { + if w := statsCache.Get(tableID); w.Exists() { if time.Now().Unix()-w.GetLastVisit() < 3 { s := w.GetStats() if s != nil && s.AccurateObjectNumber > 0 { @@ -1144,8 +1192,13 @@ func (tcc *TxnCompilerContext) Stats(obj *plan2.ObjectRef, snapshot *plan2.Snaps return nil, err } - // Cache the result - tcc.GetStatsCache().Set(tableID, result) + // A refresh may have completed while the slow path was reading storage. Do + // not let work from the old generation repopulate the new session cache. + if ses == nil { + statsCache.Set(tableID, result) + } else { + ses.cacheStatsIfCurrent(statsKey, statsVersion, result) + } return result, nil } diff --git a/pkg/frontend/compiler_context_test.go b/pkg/frontend/compiler_context_test.go index 101b70b622347..14bc4f571dd5e 100644 --- a/pkg/frontend/compiler_context_test.go +++ b/pkg/frontend/compiler_context_test.go @@ -103,6 +103,7 @@ func TestResolveViewDependencyAccount(t *testing.T) { for _, test := range []struct { name string obj *pbplan.ObjectRef + tableDef *pbplan.TableDef snapshot *pbplan.Snapshot want uint32 }{ @@ -115,6 +116,11 @@ func TestResolveViewDependencyAccount(t *testing.T) { PubInfo: &pbplan.PubInfo{TenantId: 9}}, snapshot: &pbplan.Snapshot{Tenant: &pbplan.SnapshotTenant{TenantID: 8}}, want: 9}, {name: "cluster table", obj: &pbplan.ObjectRef{SchemaName: catalog.MO_CATALOG, ObjName: "cluster_table"}, want: 0}, + {name: "relation kind alone keeps tenant context", obj: &pbplan.ObjectRef{SchemaName: "db", ObjName: "cluster_table"}, + tableDef: &pbplan.TableDef{TableType: catalog.SystemClusterRel}, want: 7}, + {name: "publication overrides generic cluster name", obj: &pbplan.ObjectRef{ + SchemaName: catalog.MO_CATALOG, ObjName: "cluster_table", + PubInfo: &pbplan.PubInfo{TenantId: 9}}, want: 9}, {name: "statement info", obj: &pbplan.ObjectRef{SchemaName: catalog.MO_SYSTEM, ObjName: catalog.MO_STATEMENT}, want: 0}, {name: "system relation overrides publisher", obj: &pbplan.ObjectRef{SchemaName: catalog.MO_SYSTEM, ObjName: catalog.MO_STATEMENT, PubInfo: &pbplan.PubInfo{TenantId: 9}}, want: 0}, @@ -122,7 +128,11 @@ func TestResolveViewDependencyAccount(t *testing.T) { {name: "sql statement cu", obj: &pbplan.ObjectRef{SchemaName: catalog.MO_SYSTEM_METRICS, ObjName: catalog.MO_SQL_STMT_CU}, want: 0}, } { t.Run(test.name, func(t *testing.T) { - got, err := tcc.ResolveViewDependencyAccount(test.obj, &pbplan.TableDef{}, test.snapshot) + tableDef := test.tableDef + if tableDef == nil { + tableDef = &pbplan.TableDef{} + } + got, err := tcc.ResolveViewDependencyAccount(test.obj, tableDef, test.snapshot) require.NoError(t, err) require.Equal(t, test.want, got) }) diff --git a/pkg/frontend/computation_wrapper.go b/pkg/frontend/computation_wrapper.go index 2bad517f1e8d1..693a4cdfc1a96 100644 --- a/pkg/frontend/computation_wrapper.go +++ b/pkg/frontend/computation_wrapper.go @@ -116,9 +116,11 @@ type TxnComputationWrapper struct { hasPreparedSchedulingSQLMode bool preparedSchedulingSQL string - // protocolVersion is captured when plan is built. The session plan cache - // uses it instead of the version observed later when execution completes. - protocolVersion int64 + // protocolVersion and optimizerStatsVersions are captured when the plan is + // built. The session plan cache uses them instead of values observed later + // when execution completes. + protocolVersion int64 + optimizerStatsVersions map[optimizerStatsTableKey]uint64 // A reusable logical plan and its generation snapshot are one immutable // binding. cachedPlan* identifies the session-cache slot so a definition @@ -245,6 +247,7 @@ func (cwft *TxnComputationWrapper) Clear() { cwft.preparedSchedulingSQLMode = "" cwft.hasPreparedSchedulingSQLMode = false cwft.preparedSchedulingSQL = "" + cwft.optimizerStatsVersions = nil cwft.planSnapshotTS = timestamp.Timestamp{} cwft.hasPlanSnapshotTS = false cwft.planGenerationReused = false @@ -254,6 +257,17 @@ func (cwft *TxnComputationWrapper) Clear() { cwft.schedulingTrace.Reset() } +func (cwft *TxnComputationWrapper) recordOptimizerStatsVersion(key optimizerStatsTableKey, version uint64) { + if cwft.optimizerStatsVersions == nil { + cwft.optimizerStatsVersions = make(map[optimizerStatsTableKey]uint64) + } + // Keep the first observed version. If publication happens between repeated + // reads, admission against the newer current version will reject the plan. + if _, exists := cwft.optimizerStatsVersions[key]; !exists { + cwft.optimizerStatsVersions[key] = version + } +} + func (cwft *TxnComputationWrapper) ParamVals() []any { return cwft.paramVals } @@ -363,6 +377,7 @@ func (cwft *TxnComputationWrapper) Compile(any any, fill func(*batch.Batch, *per cacheHit := cwft.plan != nil if !cacheHit { cwft.protocolVersion = currentProtocolVersion(cwft.proc) + clear(cwft.optimizerStatsVersions) cwft.plan, err = buildPlanWithPrepareMode( execCtx.reqCtx, cwft.ses, @@ -725,6 +740,7 @@ func (cwft *TxnComputationWrapper) completeCompileExecution( cwft.cachedPlanGeneration, cwft.plan, cwft.planSnapshotTS, + cwft.optimizerStatsVersions, ) } if !updated { diff --git a/pkg/frontend/mysql_cmd_executor.go b/pkg/frontend/mysql_cmd_executor.go index 98bd2ca25318c..98424806e1e8b 100644 --- a/pkg/frontend/mysql_cmd_executor.go +++ b/pkg/frontend/mysql_cmd_executor.go @@ -55,6 +55,7 @@ import ( "github.com/matrixorigin/matrixone/pkg/logutil" "github.com/matrixorigin/matrixone/pkg/pb/metadata" "github.com/matrixorigin/matrixone/pkg/pb/plan" + pbstats "github.com/matrixorigin/matrixone/pkg/pb/statsinfo" "github.com/matrixorigin/matrixone/pkg/pb/timestamp" pbtxn "github.com/matrixorigin/matrixone/pkg/pb/txn" "github.com/matrixorigin/matrixone/pkg/perfcounter" @@ -2085,12 +2086,125 @@ func handleAnalyzeStmt(ses *Session, execCtx *ExecCtx, stmt *tree.AnalyzeStmt) e if err != nil { return err } + if err := refreshAnalyzeTableStats(ses, execCtx, entry); err != nil { + return err + } results = append(results, result) } execCtx.results = results return nil } +func refreshAnalyzeTableStats(ses *Session, execCtx *ExecCtx, entry *tree.AnalyzeTableEntry) error { + // The derived ANALYZE query observes the transaction workspace. The engine + // statistics cache is process-global and observes only committed catalog and + // object state, so publishing while a user transaction was already active + // would mix two visibility domains. Preserve the legacy derived result and + // leave global publication to an ANALYZE statement outside that transaction. + if !analyzeStatsPublicationAllowed(execCtx) { + return nil + } + ctx := execCtx.reqCtx + if entry == nil || entry.Table == nil || entry.Table.AtTsExpr != nil { + return nil + } + + refresher, ok := getPu(ses.GetService()).StorageEngine.(engine.StatsRefresher) + if !ok { + // Engines without persistent optimizer statistics retain the legacy + // ANALYZE result behavior. + return nil + } + + tcc := ses.GetTxnCompileCtx() + dbName := resolveAnalyzeDatabase(tcc, entry.Table) + if dbName == "" { + return moerr.NewNoDB(ctx) + } + obj, tableDef, err := tcc.Resolve(dbName, string(entry.Table.Name()), nil) + if err != nil { + return err + } + if obj == nil || tableDef == nil { + return moerr.NewNoSuchTable(ctx, dbName, string(entry.Table.Name())) + } + // Historical snapshots and publication-backed tables do not own the current + // local engine statistics generation. Non-physical relations keep the + // legacy derived-query result without asking disttae to subscribe to them. + if obj.PubInfo != nil || !analyzeTableOwnsPersistentStats(tableDef) { + return nil + } + + accountID := tcc.resolvePhysicalObjectAccount(obj, tableDef, nil) + databaseID := tableDef.DbId + if databaseID == 0 { + databaseID, err = tcc.GetDatabaseId(obj.SchemaName, nil) + if err != nil { + return err + } + } + key := pbstats.StatsInfoKey{ + AccId: accountID, + DatabaseID: databaseID, + TableID: uint64(obj.Obj), + DbName: obj.SchemaName, + TableName: obj.ObjName, + } + return publishAnalyzeTableStats(ses, ctx, key, refresher) +} + +func analyzeStatsPublicationAllowed(execCtx *ExecCtx) bool { + return execCtx != nil && + execCtx.txnOpt.activeTxnAtStartKnown && + !execCtx.txnOpt.activeTxnAtStart +} + +func analyzeTableOwnsPersistentStats(tableDef *plan.TableDef) bool { + if tableDef == nil || tableDef.IsTemporary || tableDef.ViewSql != nil { + return false + } + switch tableDef.TableType { + case "", + catalog.SystemOrdinaryRel, + catalog.SystemIndexRel, + catalog.SystemMaterializedRel, + catalog.SystemClusterRel, + catalog.SystemPartitionRel: + return true + default: + return false + } +} + +func publishAnalyzeTableStats( + ses *Session, + ctx context.Context, + key pbstats.StatsInfoKey, + refresher engine.StatsRefresher, +) error { + tableKey := optimizerStatsTableKey{accountID: key.AccId, tableID: key.TableID} + release, err := acquireOptimizerStatsPublisher(ctx, ses.GetService(), tableKey) + if err != nil { + return err + } + defer release() + + stats, err := refresher.RefreshTableStats(ctx, key) + if err != nil { + return err + } + if stats == nil { + return moerr.NewInternalErrorf(ctx, "ANALYZE TABLE did not publish statistics for %s.%s", key.DbName, key.TableName) + } + + // The engine cache swap above is the data publication boundary. Advancing + // this table's version invalidates only dependent session entries; unrelated + // table statistics and plans remain reusable. + version := advanceOptimizerStatsVersion(ses.GetService(), tableKey) + ses.cachePublishedStats(tableKey, version, stats) + return nil +} + func inheritAnalyzeRewriteHint(outerSQL, derivedSQL string) string { content, ok := leadingHintContent(outerSQL) if !ok || !strings.HasPrefix(strings.TrimSpace(content), "{") { @@ -5775,8 +5889,18 @@ func doComQuery(ses *Session, execCtx *ExecCtx, input *UserInput) (retErr error) if ses.isCached(cacheKey) { return nil } + for _, cw := range cws { + if tcw, ok := cw.(*TxnComputationWrapper); ok && tcw.cachedPlanSQL == cacheKey { + // A publication or failed generation replacement made the entry stale + // while these wrappers still borrowed its AST. Do not republish the + // just-executed old plan without rebuilding its statistics dependencies. + // Wrapper cleanup runs first; the next lookup then evicts the stale owner. + return nil + } + } cacheProtocolVersion := currentProtocolVersion(proc) + planStatsVersions := make([]map[optimizerStatsTableKey]uint64, len(cws)) planSnapshotTS := make([]timestamp.Timestamp, len(cws)) for i, cw := range cws { tcw, ok := cw.(*TxnComputationWrapper) @@ -5788,6 +5912,7 @@ func doComQuery(ses *Session, execCtx *ExecCtx, input *UserInput) (retErr error) if !hasPlanSnapshotTS { return nil } + planStatsVersions[i] = tcw.optimizerStatsVersions } plans := make([]*plan.Plan, len(cws)) @@ -5802,7 +5927,8 @@ func doComQuery(ses *Session, execCtx *ExecCtx, input *UserInput) (retErr error) cw.Clear() } Cached = true - ses.cachePlanWithSnapshots(cacheKey, stmts, plans, planSnapshotTS, cacheProtocolVersion) + ses.cachePlanWithSnapshotsAndStatsVersions( + cacheKey, stmts, plans, planSnapshotTS, planStatsVersions, cacheProtocolVersion) return nil } diff --git a/pkg/frontend/mysql_cmd_executor_test.go b/pkg/frontend/mysql_cmd_executor_test.go index 3f0030a3a854e..e4dc4057cee51 100644 --- a/pkg/frontend/mysql_cmd_executor_test.go +++ b/pkg/frontend/mysql_cmd_executor_test.go @@ -50,6 +50,7 @@ import ( "github.com/matrixorigin/matrixone/pkg/geo" "github.com/matrixorigin/matrixone/pkg/pb/metadata" plan0 "github.com/matrixorigin/matrixone/pkg/pb/plan" + pbstats "github.com/matrixorigin/matrixone/pkg/pb/statsinfo" "github.com/matrixorigin/matrixone/pkg/pb/timestamp" "github.com/matrixorigin/matrixone/pkg/pb/txn" "github.com/matrixorigin/matrixone/pkg/perfcounter" @@ -4843,6 +4844,456 @@ func TestHandleAnalyzeStmtRestoresOuterExecCtxOnError(t *testing.T) { require.Same(t, outerExecCtx, ses.GetTxnCompileCtx().execCtx) } +type analyzeStatsRefresherFunc func(context.Context, pbstats.StatsInfoKey) (*pbstats.StatsInfo, error) + +func (f analyzeStatsRefresherFunc) RefreshTableStats( + ctx context.Context, key pbstats.StatsInfoKey, +) (*pbstats.StatsInfo, error) { + return f(ctx, key) +} + +func TestAnalyzeTableOwnsPersistentStats(t *testing.T) { + for _, test := range []struct { + name string + tableDef *plan0.TableDef + want bool + }{ + {name: "missing definition"}, + {name: "ordinary", tableDef: &plan0.TableDef{TableType: catalog.SystemOrdinaryRel}, want: true}, + {name: "legacy physical kind", tableDef: &plan0.TableDef{}, want: true}, + {name: "index", tableDef: &plan0.TableDef{TableType: catalog.SystemIndexRel}, want: true}, + {name: "materialized", tableDef: &plan0.TableDef{TableType: catalog.SystemMaterializedRel}, want: true}, + {name: "cluster", tableDef: &plan0.TableDef{TableType: catalog.SystemClusterRel}, want: true}, + {name: "partition", tableDef: &plan0.TableDef{TableType: catalog.SystemPartitionRel}, want: true}, + {name: "view kind", tableDef: &plan0.TableDef{TableType: catalog.SystemViewRel}}, + {name: "view definition", tableDef: &plan0.TableDef{ + TableType: catalog.SystemOrdinaryRel, + ViewSql: &plan0.ViewDef{View: "select 1"}, + }}, + {name: "external", tableDef: &plan0.TableDef{TableType: catalog.SystemExternalRel}}, + {name: "sequence", tableDef: &plan0.TableDef{TableType: catalog.SystemSequenceRel}}, + {name: "removed source", tableDef: &plan0.TableDef{TableType: catalog.SystemSourceRel}}, + {name: "temporary session table", tableDef: &plan0.TableDef{ + TableType: catalog.SystemOrdinaryRel, + IsTemporary: true, + }}, + {name: "legacy temporary kind", tableDef: &plan0.TableDef{TableType: catalog.SystemTemporaryTable}}, + } { + t.Run(test.name, func(t *testing.T) { + require.Equal(t, test.want, analyzeTableOwnsPersistentStats(test.tableDef)) + }) + } +} + +func TestAnalyzeStatsPublicationRequiresStatementOwnedTransaction(t *testing.T) { + for _, test := range []struct { + name string + execCtx *ExecCtx + want bool + }{ + {name: "missing execution context"}, + {name: "unknown transaction owner", execCtx: &ExecCtx{}}, + {name: "statement-owned transaction", execCtx: &ExecCtx{txnOpt: FeTxnOption{ + activeTxnAtStartKnown: true, + }}, want: true}, + {name: "pre-existing user transaction", execCtx: &ExecCtx{txnOpt: FeTxnOption{ + activeTxnAtStartKnown: true, + activeTxnAtStart: true, + }}}, + } { + t.Run(test.name, func(t *testing.T) { + require.Equal(t, test.want, analyzeStatsPublicationAllowed(test.execCtx)) + }) + } +} + +func isolateOptimizerStatsTest(t *testing.T, sessions ...*Session) { + t.Helper() + service := "optimizer-stats-" + t.Name() + InitServerLevelVars(service) + previous := make([]string, len(sessions)) + for i, ses := range sessions { + previous[i] = ses.GetService() + ses.feSessionImpl.service = service + } + t.Cleanup(func() { + for i, ses := range sessions { + ses.feSessionImpl.service = previous[i] + } + serverVarsMap.Delete(service) + }) +} + +func optimizerStatsTestPlan(tableIDs ...uint64) *plan0.Plan { + nodes := make([]*plan0.Node, 0, len(tableIDs)) + for _, tableID := range tableIDs { + nodes = append(nodes, &plan0.Node{TableDef: &plan0.TableDef{TblId: tableID}}) + } + return &plan0.Plan{Plan: &plan0.Plan_Query{Query: &plan0.Query{Nodes: nodes}}} +} + +func optimizerStatsVersionsForTest( + ses *Session, + tableIDs ...uint64, +) map[optimizerStatsTableKey]uint64 { + keys := make([]optimizerStatsTableKey, 0, len(tableIDs)) + for _, tableID := range tableIDs { + keys = append(keys, ses.optimizerStatsKey(tableID)) + } + return optimizerStatsVersionsForKeysTest(ses, keys...) + +} + +func optimizerStatsVersionsForKeysTest( + ses *Session, + keys ...optimizerStatsTableKey, +) map[optimizerStatsTableKey]uint64 { + versions := make(map[optimizerStatsTableKey]uint64, len(keys)) + for _, key := range keys { + versions[key] = currentOptimizerStatsVersion(ses.GetService(), key) + } + return versions +} + +func cacheOptimizerPlanForTest(ses *Session, sql string, tableIDs ...uint64) { + keys := make([]optimizerStatsTableKey, 0, len(tableIDs)) + for _, tableID := range tableIDs { + keys = append(keys, ses.optimizerStatsKey(tableID)) + } + cacheOptimizerPlanForKeysTest(ses, sql, keys...) +} + +func cacheOptimizerPlanForKeysTest(ses *Session, sql string, keys ...optimizerStatsTableKey) { + tableIDs := make([]uint64, 0, len(keys)) + for _, key := range keys { + tableIDs = append(tableIDs, key.tableID) + } + ses.cachePlanWithStatsVersions( + sql, + []tree.Statement{&tree.Select{}}, + []*plan0.Plan{optimizerStatsTestPlan(tableIDs...)}, + optimizerStatsVersionsForKeysTest(ses, keys...), + ) +} + +func cacheOptimizerStatsForTest(t *testing.T, ses *Session, tableID uint64, stats *pbstats.StatsInfo) uint64 { + return cacheOptimizerStatsForKeyTest(t, ses, ses.optimizerStatsKey(tableID), stats) +} + +func cacheOptimizerStatsForKeyTest( + t *testing.T, + ses *Session, + key optimizerStatsTableKey, + stats *pbstats.StatsInfo, +) uint64 { + t.Helper() + version := currentOptimizerStatsVersion(ses.GetService(), key) + require.True(t, ses.cacheStatsIfCurrent(key, version, stats)) + return version +} + +func TestCompilerContextRecordsTheStatsVersionActuallyRead(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + ses, execCtx := newAnalyzeHandlerTestSession(t, ctrl) + isolateOptimizerStatsTest(t, ses) + ses.SetAccountId(7) + + const tableID = uint64(42) + wrapper := InitTxnComputationWrapper(ses, &tree.Select{}, execCtx.proc) + tcc := ses.GetTxnCompileCtx() + tcc.SetExecCtx(execCtx) + tcc.tcw = wrapper + + key := tcc.optimizerStatsKey(&plan0.ObjectRef{ + Obj: int64(tableID), + SchemaName: catalog.MO_SYSTEM, + ObjName: catalog.MO_STATEMENT, + }, nil) + require.Equal(t, optimizerStatsTableKey{ + accountID: catalog.System_Account, + tableID: tableID, + }, key) + _, firstVersion := ses.getStatsCacheWithVersion(key) + require.NotContains(t, ses.statsCacheVersions, tableID, + "a failed or uncached stats read must not grow session version metadata") + _, _, recordedVersion := tcc.getStatsCacheVersion(key) + require.Equal(t, firstVersion, recordedVersion) + require.Equal(t, firstVersion, wrapper.optimizerStatsVersions[key]) + + advanceOptimizerStatsVersion(ses.GetService(), key) + _, _, currentVersion := tcc.getStatsCacheVersion(key) + require.NotEqual(t, firstVersion, currentVersion) + require.Equal(t, firstVersion, wrapper.optimizerStatsVersions[key], + "a plan that read both sides of publication must retain its stale dependency and be rejected") +} + +func TestSessionStatsCacheDoesNotAliasSameTableIDAcrossAccounts(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + ses, _ := newAnalyzeHandlerTestSession(t, ctrl) + isolateOptimizerStatsTest(t, ses) + + const tableID = uint64(42) + tenantKey := optimizerStatsTableKey{accountID: 7, tableID: tableID} + systemKey := optimizerStatsTableKey{accountID: catalog.System_Account, tableID: tableID} + tenantStats := plan.NewStatsInfo() + tenantStats.TableCnt = 7 + systemStats := plan.NewStatsInfo() + systemStats.TableCnt = 70 + + tenantVersion := currentOptimizerStatsVersion(ses.GetService(), tenantKey) + require.True(t, ses.cacheStatsIfCurrent(tenantKey, tenantVersion, tenantStats)) + cache, _ := ses.getStatsCacheWithVersion(tenantKey) + wrapper := cache.Get(tableID) + require.Same(t, tenantStats, wrapper.GetStats()) + + cache, systemVersion := ses.getStatsCacheWithVersion(systemKey) + wrapper = cache.Get(tableID) + require.False(t, wrapper.Exists(), + "the physical owner is part of the cache identity") + require.True(t, ses.cacheStatsIfCurrent(systemKey, systemVersion, systemStats)) + wrapper = cache.Get(tableID) + require.Same(t, systemStats, wrapper.GetStats()) + + cache, _ = ses.getStatsCacheWithVersion(tenantKey) + wrapper = cache.Get(tableID) + require.False(t, wrapper.Exists(), + "switching back must not expose statistics from the system account") +} + +func TestOptimizerStatsVersionsCompactWithoutRevalidatingOldEntries(t *testing.T) { + vars := &ServerLevelVariables{ + optimizerStatsVersions: make(map[optimizerStatsTableKey]uint64), + } + first := optimizerStatsTableKey{accountID: 1, tableID: 10} + second := optimizerStatsTableKey{accountID: 1, tableID: 20} + third := optimizerStatsTableKey{accountID: 1, tableID: 30} + + firstVersion := advanceOptimizerStatsVersionLocked(vars, first, 2) + secondVersion := advanceOptimizerStatsVersionLocked(vars, second, 2) + require.Equal(t, uint64(1), firstVersion) + require.Equal(t, uint64(2), secondVersion) + + thirdVersion := advanceOptimizerStatsVersionLocked(vars, third, 2) + require.Equal(t, uint64(4), thirdVersion) + require.Len(t, vars.optimizerStatsVersions, 1) + require.Equal(t, uint64(3), currentOptimizerStatsVersionLocked(vars, first)) + require.Equal(t, uint64(3), currentOptimizerStatsVersionLocked(vars, second)) + require.NotEqual(t, firstVersion, currentOptimizerStatsVersionLocked(vars, first)) + require.NotEqual(t, secondVersion, currentOptimizerStatsVersionLocked(vars, second)) +} + +func TestPublishAnalyzeTableStatsDefinesCacheBoundary(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + ses, execCtx := newAnalyzeHandlerTestSession(t, ctrl) + otherSes, _ := newAnalyzeHandlerTestSession(t, ctrl) + otherTenantSes, _ := newAnalyzeHandlerTestSession(t, ctrl) + crossAccountSes, _ := newAnalyzeHandlerTestSession(t, ctrl) + isolateOptimizerStatsTest(t, ses, otherSes, otherTenantSes, crossAccountSes) + otherTenantSes.SetAccountId(7) + crossAccountSes.SetAccountId(7) + + const ( + tableID = uint64(42) + otherTableID = uint64(84) + ) + key := pbstats.StatsInfoKey{ + AccId: catalog.System_Account, + DatabaseID: 7, + TableID: tableID, + DbName: "db", + TableName: "events", + } + oldStats := plan.NewStatsInfo() + oldStats.NdvMap["url"] = 1 + otherStats := plan.NewStatsInfo() + otherStats.NdvMap["id"] = 84 + oldVersion := cacheOptimizerStatsForTest(t, otherSes, tableID, oldStats) + cacheOptimizerStatsForTest(t, otherSes, otherTableID, otherStats) + cacheOptimizerStatsForTest(t, ses, tableID, oldStats) + cacheOptimizerStatsForTest(t, otherTenantSes, tableID, oldStats) + physicalKey := optimizerStatsTableKey{accountID: catalog.System_Account, tableID: tableID} + crossAccountOldVersion := cacheOptimizerStatsForKeyTest(t, crossAccountSes, physicalKey, oldStats) + + dependentPlan := optimizerStatsTestPlan(tableID) + compileVersions := optimizerStatsVersionsForTest(otherSes, tableID) + cacheOptimizerPlanForTest(otherSes, "select url from events", tableID) + cacheOptimizerPlanForTest(otherSes, "select id from other_table", otherTableID) + cacheOptimizerPlanForTest(otherTenantSes, "select url from tenant_events", tableID) + cacheOptimizerPlanForKeysTest(crossAccountSes, "select url from system.statement_info", physicalKey) + + freshStats := plan.NewStatsInfo() + freshStats.AccurateObjectNumber = 8 + freshStats.NdvMap["url"] = 1_000_000 + var gotKey pbstats.StatsInfoKey + refresher := analyzeStatsRefresherFunc(func(_ context.Context, key pbstats.StatsInfoKey) (*pbstats.StatsInfo, error) { + gotKey = key + return freshStats, nil + }) + + require.NoError(t, publishAnalyzeTableStats(ses, execCtx.reqCtx, key, refresher)) + require.Equal(t, key, gotKey) + cache, _ := ses.getStatsCacheWithVersion(physicalKey) + wrapper := cache.Get(tableID) + require.Same(t, freshStats, wrapper.GetStats()) + otherSes.cachePlanWithStatsVersions("compiled before analyze completed", + []tree.Statement{&tree.Select{}}, []*plan0.Plan{dependentPlan}, compileVersions) + require.Nil(t, otherSes.getCachedPlan("compiled before analyze completed"), + "a plan compiled across the publication boundary must not enter the cache") + require.Nil(t, otherSes.getCachedPlan("select url from events")) + require.NotNil(t, otherSes.getCachedPlan("select id from other_table"), + "an unrelated table publication must not flush the session plan cache") + require.NotNil(t, otherTenantSes.getCachedPlan("select url from tenant_events"), + "the same table ID in another account must keep its plan") + require.Nil(t, crossAccountSes.getCachedPlan("select url from system.statement_info"), + "a tenant plan must validate the system account generation that owns its statistics") + require.False(t, otherSes.cacheStatsIfCurrent(otherSes.optimizerStatsKey(tableID), oldVersion, oldStats), + "a stats read started before publication must not repopulate the table entry") + require.False(t, crossAccountSes.cacheStatsIfCurrent(physicalKey, crossAccountOldVersion, oldStats), + "a cross-account stats read must not repopulate the old physical generation") + otherCache, currentVersion := otherSes.getStatsCacheWithVersion(otherSes.optimizerStatsKey(tableID)) + otherWrapper := otherCache.Get(tableID) + require.False(t, otherWrapper.Exists()) + otherTableWrapper := otherCache.Get(otherTableID) + require.Same(t, otherStats, otherTableWrapper.GetStats(), + "invalidating one table must retain unrelated statistics") + require.True(t, otherSes.cacheStatsIfCurrent( + otherSes.optimizerStatsKey(tableID), currentVersion, freshStats)) + currentWrapper := otherCache.Get(tableID) + require.Same(t, freshStats, currentWrapper.GetStats()) +} + +func TestPublishAnalyzeTableStatsDoesNotExposeFailedRefresh(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + ses, execCtx := newAnalyzeHandlerTestSession(t, ctrl) + isolateOptimizerStatsTest(t, ses) + + const tableID = uint64(42) + key := pbstats.StatsInfoKey{TableID: tableID, DbName: "db", TableName: "events"} + oldStats := plan.NewStatsInfo() + cacheOptimizerStatsForTest(t, ses, tableID, oldStats) + cacheOptimizerPlanForTest(ses, "select url from events", tableID) + wantErr := moerr.NewInternalError(execCtx.reqCtx, "refresh failed") + refresher := analyzeStatsRefresherFunc(func(context.Context, pbstats.StatsInfoKey) (*pbstats.StatsInfo, error) { + return nil, wantErr + }) + + err := publishAnalyzeTableStats(ses, execCtx.reqCtx, key, refresher) + require.ErrorIs(t, err, wantErr) + cache, _ := ses.getStatsCacheWithVersion(ses.optimizerStatsKey(tableID)) + wrapper := cache.Get(tableID) + require.Same(t, oldStats, wrapper.GetStats()) + require.NotNil(t, ses.getCachedPlan("select url from events")) + + tableKey := optimizerStatsTableKey{accountID: key.AccId, tableID: key.TableID} + admission := getOptimizerStatsVars(ses.GetService()). + optimizerStatsPublish[optimizerStatsPublisherStripe(tableKey)] + select { + case admission <- struct{}{}: + <-admission + default: + t.Fatal("a failed refresh leaked same-table publication admission") + } + freshStats := plan.NewStatsInfo() + require.NoError(t, publishAnalyzeTableStats(ses, execCtx.reqCtx, key, + analyzeStatsRefresherFunc(func(context.Context, pbstats.StatsInfoKey) (*pbstats.StatsInfo, error) { + return freshStats, nil + })), "a failed refresh must release same-table publication admission") +} + +func TestPublishAnalyzeTableStatsRejectsMissingRefreshResult(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + ses, execCtx := newAnalyzeHandlerTestSession(t, ctrl) + isolateOptimizerStatsTest(t, ses) + + const tableID = uint64(42) + key := pbstats.StatsInfoKey{TableID: tableID, DbName: "db", TableName: "events"} + oldStats := plan.NewStatsInfo() + cacheOptimizerStatsForTest(t, ses, tableID, oldStats) + cacheOptimizerPlanForTest(ses, "select url from events", tableID) + version := currentOptimizerStatsVersion(ses.GetService(), ses.optimizerStatsKey(tableID)) + clock := currentOptimizerStatsClock(ses.GetService()) + refresher := analyzeStatsRefresherFunc(func(context.Context, pbstats.StatsInfoKey) (*pbstats.StatsInfo, error) { + return nil, nil + }) + + err := publishAnalyzeTableStats(ses, execCtx.reqCtx, key, refresher) + require.Error(t, err) + require.Equal(t, version, + currentOptimizerStatsVersion(ses.GetService(), ses.optimizerStatsKey(tableID))) + require.Equal(t, clock, currentOptimizerStatsClock(ses.GetService())) + cache, _ := ses.getStatsCacheWithVersion(ses.optimizerStatsKey(tableID)) + wrapper := cache.Get(tableID) + require.Same(t, oldStats, wrapper.GetStats()) + require.NotNil(t, ses.getCachedPlan("select url from events")) +} + +func TestPublishAnalyzeTableStatsSerializesAndCancelsAdmission(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + firstSes, firstExecCtx := newAnalyzeHandlerTestSession(t, ctrl) + secondSes, secondExecCtx := newAnalyzeHandlerTestSession(t, ctrl) + isolateOptimizerStatsTest(t, firstSes, secondSes) + + key := pbstats.StatsInfoKey{ + AccId: catalog.System_Account, TableID: 42, DbName: "db", TableName: "events", + } + entered := make(chan struct{}) + unblock := make(chan struct{}, 1) + releaseFirst := func() { + select { + case unblock <- struct{}{}: + default: + } + } + t.Cleanup(releaseFirst) + firstDone := make(chan error, 1) + go func() { + firstDone <- publishAnalyzeTableStats(firstSes, firstExecCtx.reqCtx, key, + analyzeStatsRefresherFunc(func(context.Context, pbstats.StatsInfoKey) (*pbstats.StatsInfo, error) { + close(entered) + <-unblock + return plan.NewStatsInfo(), nil + })) + }() + <-entered + + // A publication for a different table must not queue behind this table. + otherKey := key + otherKey.TableID = 43 + otherKey.TableName = "other_events" + var otherCalled atomic.Bool + require.NoError(t, publishAnalyzeTableStats(secondSes, secondExecCtx.reqCtx, otherKey, + analyzeStatsRefresherFunc(func(context.Context, pbstats.StatsInfoKey) (*pbstats.StatsInfo, error) { + otherCalled.Store(true) + return plan.NewStatsInfo(), nil + }))) + require.True(t, otherCalled.Load()) + + secondCtx, cancel := context.WithCancel(secondExecCtx.reqCtx) + cancel() + var secondCalled atomic.Bool + err := publishAnalyzeTableStats(secondSes, secondCtx, key, + analyzeStatsRefresherFunc(func(context.Context, pbstats.StatsInfoKey) (*pbstats.StatsInfo, error) { + secondCalled.Store(true) + return plan.NewStatsInfo(), nil + })) + require.ErrorIs(t, err, context.Canceled) + require.False(t, secondCalled.Load()) + + releaseFirst() + select { + case err = <-firstDone: + require.NoError(t, err) + case <-time.After(5 * time.Second): + t.Fatal("first statistics publication did not finish") + } +} + func TestSetExecCtxClearsPreviousStatementViews(t *testing.T) { tcc := &TxnCompilerContext{} tcc.SetViews([]string{"db#stale_view"}) diff --git a/pkg/frontend/plan_cache.go b/pkg/frontend/plan_cache.go index 48c9fe397776c..05e7c9d469ebf 100644 --- a/pkg/frontend/plan_cache.go +++ b/pkg/frontend/plan_cache.go @@ -23,12 +23,14 @@ import ( ) type cachedPlan struct { - sql string - stmts []tree.Statement - plans []*plan.Plan - planSnapshotTS []timestamp.Timestamp - protocolVersion int64 - invalid bool + sql string + stmts []tree.Statement + plans []*plan.Plan + planSnapshotTS []timestamp.Timestamp + protocolVersion int64 + statsVersions map[optimizerStatsTableKey]uint64 + planStatsVersions []map[optimizerStatsTableKey]uint64 + invalid bool } // planCache uses LRU to cache plan for the same sql @@ -57,15 +59,17 @@ func freeStmts(stmts []tree.Statement) { func (pc *planCache) cache(sql string, stmts []tree.Statement, plans []*plan.Plan, versions ...int64) { // Legacy internal callers get a conservative oldest-possible binding. The // production cache path always supplies the actual generation snapshot. - pc.cacheWithPlanSnapshots( - sql, stmts, plans, make([]timestamp.Timestamp, len(plans)), versions...) + pc.cacheWithPlanSnapshotsAndStatsVersions( + sql, stmts, plans, make([]timestamp.Timestamp, len(plans)), + make([]map[optimizerStatsTableKey]uint64, len(plans)), versions...) } -func (pc *planCache) cacheWithPlanSnapshots( +func (pc *planCache) cacheWithPlanSnapshotsAndStatsVersions( sql string, stmts []tree.Statement, plans []*plan.Plan, planSnapshotTS []timestamp.Timestamp, + planStatsVersions []map[optimizerStatsTableKey]uint64, versions ...int64, ) { protocolVersion := currentProtocolVersion(nil) @@ -76,7 +80,13 @@ func (pc *planCache) cacheWithPlanSnapshots( pc.cachePool = make(map[string]*list.Element) pc.lruList = list.New() } - if len(stmts) != len(plans) || len(planSnapshotTS) != len(plans) { + if len(stmts) != len(plans) || len(planSnapshotTS) != len(plans) || + len(planStatsVersions) != len(plans) { + freeStmts(stmts) + return + } + statsVersions, versionsConsistent := aggregatePlanStatsVersions(planStatsVersions) + if !versionsConsistent { freeStmts(stmts) return } @@ -90,21 +100,25 @@ func (pc *planCache) cacheWithPlanSnapshots( if element, ok := pc.cachePool[sql]; ok { freeStmts(element.Value.(*cachedPlan).stmts) element.Value = &cachedPlan{ - sql: sql, - stmts: stmts, - plans: plans, - planSnapshotTS: planSnapshotTS, - protocolVersion: protocolVersion, + sql: sql, + stmts: stmts, + plans: plans, + planSnapshotTS: planSnapshotTS, + protocolVersion: protocolVersion, + statsVersions: statsVersions, + planStatsVersions: clonePlanStatsVersions(planStatsVersions), } pc.lruList.MoveToFront(element) return } element := pc.lruList.PushFront(&cachedPlan{ - sql: sql, - stmts: stmts, - plans: plans, - planSnapshotTS: planSnapshotTS, - protocolVersion: protocolVersion, + sql: sql, + stmts: stmts, + plans: plans, + planSnapshotTS: planSnapshotTS, + protocolVersion: protocolVersion, + statsVersions: statsVersions, + planStatsVersions: clonePlanStatsVersions(planStatsVersions), }) pc.cachePool[sql] = element if pc.lruList.Len() > pc.capacity { @@ -115,6 +129,66 @@ func (pc *planCache) cacheWithPlanSnapshots( } } +func cloneStatsVersions(versions map[optimizerStatsTableKey]uint64) map[optimizerStatsTableKey]uint64 { + if len(versions) == 0 { + return nil + } + cloned := make(map[optimizerStatsTableKey]uint64, len(versions)) + for key, version := range versions { + cloned[key] = version + } + return cloned +} + +func planStatsVersionsFromAggregate( + planCount int, + versions map[optimizerStatsTableKey]uint64, +) []map[optimizerStatsTableKey]uint64 { + perPlan := make([]map[optimizerStatsTableKey]uint64, planCount) + if planCount > 0 { + perPlan[0] = versions + } + return perPlan +} + +func clonePlanStatsVersions( + versions []map[optimizerStatsTableKey]uint64, +) []map[optimizerStatsTableKey]uint64 { + cloned := make([]map[optimizerStatsTableKey]uint64, len(versions)) + for i := range versions { + cloned[i] = cloneStatsVersions(versions[i]) + } + return cloned +} + +func aggregatePlanStatsVersions( + versions []map[optimizerStatsTableKey]uint64, +) (map[optimizerStatsTableKey]uint64, bool) { + var aggregated map[optimizerStatsTableKey]uint64 + for _, planVersions := range versions { + if len(planVersions) == 0 { + continue + } + if aggregated == nil { + aggregated = make(map[optimizerStatsTableKey]uint64) + } + if !mergeOptimizerStatsVersions(aggregated, planVersions) { + return nil, false + } + } + return aggregated, true +} + +func mergeOptimizerStatsVersions(dst, src map[optimizerStatsTableKey]uint64) bool { + for key, version := range src { + if prior, exists := dst[key]; exists && prior != version { + return false + } + dst[key] = version + } + return true +} + func (pc *planCache) remove(sql string) { if pc.cachePool == nil { return @@ -135,7 +209,8 @@ func (pc *planCache) get(sql string) *cachedPlan { } if element, ok := pc.cachePool[sql]; ok { cp := element.Value.(*cachedPlan) - if cp.invalid || len(cp.planSnapshotTS) != len(cp.plans) { + if cp.invalid || len(cp.planSnapshotTS) != len(cp.plans) || + len(cp.planStatsVersions) != len(cp.plans) { pc.remove(sql) return nil } @@ -154,7 +229,8 @@ func (pc *planCache) isCached(sql string) bool { return false } cached := element.Value.(*cachedPlan) - return !cached.invalid && len(cached.planSnapshotTS) == len(cached.plans) + return !cached.invalid && len(cached.planSnapshotTS) == len(cached.plans) && + len(cached.planStatsVersions) == len(cached.plans) } func (pc *planCache) updatePlanGeneration( @@ -163,6 +239,7 @@ func (pc *planCache) updatePlanGeneration( expectedPlan *plan.Plan, newPlan *plan.Plan, planSnapshotTS timestamp.Timestamp, + statsVersions map[optimizerStatsTableKey]uint64, ) bool { if pc.cachePool == nil || newPlan == nil { return false @@ -174,11 +251,20 @@ func (pc *planCache) updatePlanGeneration( cached := element.Value.(*cachedPlan) if cached.invalid || index < 0 || index >= len(cached.plans) || len(cached.planSnapshotTS) != len(cached.plans) || + len(cached.planStatsVersions) != len(cached.plans) || cached.plans[index] != expectedPlan { return false } + updatedPlanStatsVersions := clonePlanStatsVersions(cached.planStatsVersions) + updatedPlanStatsVersions[index] = cloneStatsVersions(statsVersions) + aggregated, versionsConsistent := aggregatePlanStatsVersions(updatedPlanStatsVersions) + if !versionsConsistent { + return false + } cached.plans[index] = newPlan cached.planSnapshotTS[index] = planSnapshotTS + cached.statsVersions = aggregated + cached.planStatsVersions = updatedPlanStatsVersions pc.lruList.MoveToFront(element) return true } diff --git a/pkg/frontend/plan_cache_test.go b/pkg/frontend/plan_cache_test.go index b7657e4614244..347598b737809 100644 --- a/pkg/frontend/plan_cache_test.go +++ b/pkg/frontend/plan_cache_test.go @@ -140,20 +140,39 @@ func TestPlanCacheUpdatesPlanAndSnapshotAsOneGeneration(t *testing.T) { firstTS := timestamp.Timestamp{PhysicalTime: 10} oldTS := timestamp.Timestamp{PhysicalTime: 20} newTS := timestamp.Timestamp{PhysicalTime: 30} - pc.cacheWithPlanSnapshots( + firstStats := optimizerStatsTableKey{accountID: 1, tableID: 10} + secondStats := optimizerStatsTableKey{accountID: 2, tableID: 20} + rebuiltStats := optimizerStatsTableKey{accountID: 2, tableID: 21} + pc.cacheWithPlanSnapshotsAndStatsVersions( "sql", []tree.Statement{firstStmt, secondStmt}, []*plan.Plan{firstPlan, oldPlan}, []timestamp.Timestamp{firstTS, oldTS}, + []map[optimizerStatsTableKey]uint64{ + {firstStats: 1}, + {secondStats: 2}, + }, ) - require.False(t, pc.updatePlanGeneration("sql", 1, firstPlan, newPlan, newTS)) - require.True(t, pc.updatePlanGeneration("sql", 1, oldPlan, newPlan, newTS)) + require.False(t, pc.updatePlanGeneration("sql", 1, firstPlan, newPlan, newTS, nil)) + require.False(t, pc.updatePlanGeneration( + "sql", 1, oldPlan, newPlan, newTS, + map[optimizerStatsTableKey]uint64{firstStats: 3}), + "a replacement cannot combine two versions of one dependency") + require.True(t, pc.updatePlanGeneration( + "sql", 1, oldPlan, newPlan, newTS, + map[optimizerStatsTableKey]uint64{rebuiltStats: 3})) cached := pc.get("sql") require.Same(t, firstPlan, cached.plans[0]) require.Equal(t, firstTS, cached.planSnapshotTS[0]) require.Same(t, newPlan, cached.plans[1]) require.Equal(t, newTS, cached.planSnapshotTS[1]) + require.Equal(t, map[optimizerStatsTableKey]uint64{ + firstStats: 1, + rebuiltStats: 3, + }, cached.statsVersions) + require.Equal(t, map[optimizerStatsTableKey]uint64{firstStats: 1}, cached.planStatsVersions[0]) + require.Equal(t, map[optimizerStatsTableKey]uint64{rebuiltStats: 3}, cached.planStatsVersions[1]) pc.invalidatePlanGeneration("sql", 1, oldPlan) require.True(t, pc.isCached("sql"), "an obsolete generation cannot invalidate its replacement") @@ -315,6 +334,77 @@ func Test_SessionAccessorsWithNilPlanCache(t *testing.T) { require.NotPanics(t, func() { ses.releasePlanCache() }) } +func TestMergeOptimizerStatsVersionsRejectsMixedGenerations(t *testing.T) { + first := optimizerStatsTableKey{accountID: 7, tableID: 1} + second := optimizerStatsTableKey{accountID: 8, tableID: 1} + versions := map[optimizerStatsTableKey]uint64{first: 10} + require.True(t, mergeOptimizerStatsVersions(versions, + map[optimizerStatsTableKey]uint64{first: 10, second: 20})) + require.Equal(t, map[optimizerStatsTableKey]uint64{first: 10, second: 20}, versions) + require.False(t, mergeOptimizerStatsVersions(versions, + map[optimizerStatsTableKey]uint64{first: 11})) + require.Equal(t, uint64(10), versions[first]) +} + +func TestSessionReportsStaleStatsWithoutFreeingBorrowedCachedAST(t *testing.T) { + const service = "stale-stats-borrowed-ast" + InitServerLevelVars(service) + t.Cleanup(func() { serverVarsMap.Delete(service) }) + + key := optimizerStatsTableKey{accountID: 7, tableID: 11} + stmt := &trackedStatement{} + ses := &Session{ + feSessionImpl: feSessionImpl{service: service}, + planCache: newPlanCache(1), + } + ses.cachePlanWithStatsVersions( + "cached", []tree.Statement{stmt}, []*plan.Plan{{}}, + map[optimizerStatsTableKey]uint64{key: 0}) + require.True(t, ses.isCached("cached")) + + advanceOptimizerStatsVersion(service, key) + require.False(t, ses.isCached("cached")) + require.Zero(t, stmt.freed, "an executing wrapper may still borrow this AST") + + require.Nil(t, ses.getCachedPlan("cached")) + require.Equal(t, 1, stmt.freed) +} + +var optimizerStatsVersionsCurrentSink bool + +func BenchmarkOptimizerStatsVersionsCurrent(b *testing.B) { + const ( + service = "optimizer-stats-version-benchmark" + accountID = uint32(7) + ) + InitServerLevelVars(service) + b.Cleanup(func() { serverVarsMap.Delete(service) }) + + for _, tc := range []struct { + name string + dependency int + }{ + {name: "no-dependency", dependency: 0}, + {name: "one-table", dependency: 1}, + {name: "four-tables", dependency: 4}, + {name: "sixteen-tables", dependency: 16}, + } { + b.Run(tc.name, func(b *testing.B) { + versions := make(map[optimizerStatsTableKey]uint64, tc.dependency) + for tableID := 1; tableID <= tc.dependency; tableID++ { + versions[optimizerStatsTableKey{ + accountID: accountID, + tableID: uint64(tableID), + }] = 0 + } + b.ReportAllocs() + for b.Loop() { + optimizerStatsVersionsCurrentSink = optimizerStatsVersionsCurrent(service, versions) + } + }) + } +} + func TestSessionRemoveCachedPlanOnlyEvictsTarget(t *testing.T) { ses := &Session{planCache: newPlanCache(2)} first := &trackedStatement{} diff --git a/pkg/frontend/server.go b/pkg/frontend/server.go index 0ce712a820427..155c6304efa49 100644 --- a/pkg/frontend/server.go +++ b/pkg/frontend/server.go @@ -553,6 +553,16 @@ func nextConnectionID() uint32 { var serverVarsMap sync.Map +const ( + optimizerStatsPublisherStripes = 64 + optimizerStatsVersionEntries = 64 * 1024 +) + +type optimizerStatsTableKey struct { + accountID uint32 + tableID uint64 +} + func init() { InitServerLevelVars("") } @@ -567,10 +577,25 @@ func getServerLevelVars(service string) *ServerLevelVariables { } func InitServerLevelVars(service string) { - serverVarsMap.LoadOrStore(service, &ServerLevelVariables{}) + vars := &ServerLevelVariables{ + optimizerStatsVersions: make(map[optimizerStatsTableKey]uint64), + } + for i := range vars.optimizerStatsPublish { + vars.optimizerStatsPublish[i] = make(chan struct{}, 1) + } + serverVarsMap.LoadOrStore(service, vars) getServerLevelVars(service) } +func getOptimizerStatsVars(service string) *ServerLevelVariables { + vars := getServerLevelVars(service) + if vars == nil { + InitServerLevelVars(service) + vars = getServerLevelVars(service) + } + return vars +} + func getSessionAlloc(service string) Allocator { return getServerLevelVars(service).sessionAlloc.Load().(Allocator) } @@ -633,6 +658,89 @@ func getPu(service string) *config.ParameterUnit { return pu } +func currentOptimizerStatsClock(service string) uint64 { + vars := getOptimizerStatsVars(service) + vars.optimizerStatsMu.RLock() + defer vars.optimizerStatsMu.RUnlock() + return vars.optimizerStatsClock +} + +func currentOptimizerStatsVersion(service string, key optimizerStatsTableKey) uint64 { + vars := getOptimizerStatsVars(service) + vars.optimizerStatsMu.RLock() + defer vars.optimizerStatsMu.RUnlock() + return currentOptimizerStatsVersionLocked(vars, key) +} + +func advanceOptimizerStatsVersion(service string, key optimizerStatsTableKey) uint64 { + vars := getOptimizerStatsVars(service) + vars.optimizerStatsMu.Lock() + defer vars.optimizerStatsMu.Unlock() + return advanceOptimizerStatsVersionLocked(vars, key, optimizerStatsVersionEntries) +} + +func currentOptimizerStatsVersionLocked(vars *ServerLevelVariables, key optimizerStatsTableKey) uint64 { + if version, ok := vars.optimizerStatsVersions[key]; ok { + return version + } + return vars.optimizerStatsReset +} + +func advanceOptimizerStatsVersionLocked( + vars *ServerLevelVariables, + key optimizerStatsTableKey, + maxEntries int, +) uint64 { + if _, exists := vars.optimizerStatsVersions[key]; !exists && len(vars.optimizerStatsVersions) >= maxEntries { + // Explicit ANALYZE of many short-lived tables must not grow process + // metadata forever. A rare compaction advances the missing-key token, + // making every older cache label conservatively stale before reuse. + vars.optimizerStatsClock++ + vars.optimizerStatsReset = vars.optimizerStatsClock + clear(vars.optimizerStatsVersions) + } + vars.optimizerStatsClock++ + vars.optimizerStatsVersions[key] = vars.optimizerStatsClock + return vars.optimizerStatsClock +} + +func optimizerStatsVersionsCurrent( + service string, + versions map[optimizerStatsTableKey]uint64, +) bool { + if len(versions) == 0 { + return true + } + vars := getOptimizerStatsVars(service) + vars.optimizerStatsMu.RLock() + defer vars.optimizerStatsMu.RUnlock() + for key, version := range versions { + if currentOptimizerStatsVersionLocked(vars, key) != version { + return false + } + } + return true +} + +func optimizerStatsPublisherStripe(key optimizerStatsTableKey) int { + mixed := key.tableID ^ uint64(key.accountID)*0x9e3779b97f4a7c15 + return int(mixed % optimizerStatsPublisherStripes) +} + +func acquireOptimizerStatsPublisher( + ctx context.Context, + service string, + key optimizerStatsTableKey, +) (func(), error) { + admission := getOptimizerStatsVars(service).optimizerStatsPublish[optimizerStatsPublisherStripe(key)] + select { + case admission <- struct{}{}: + return func() { <-admission }, nil + case <-ctx.Done(): + return nil, context.Cause(ctx) + } +} + func setAicm(service string, aicm *defines.AutoIncrCacheManager) { getServerLevelVars(service).Aicm.Store(aicm) } diff --git a/pkg/frontend/session.go b/pkg/frontend/session.go index 49bce8765a3d5..e3be559b5a111 100644 --- a/pkg/frontend/session.go +++ b/pkg/frontend/session.go @@ -45,6 +45,7 @@ import ( "github.com/matrixorigin/matrixone/pkg/logutil" "github.com/matrixorigin/matrixone/pkg/pb/plan" "github.com/matrixorigin/matrixone/pkg/pb/query" + pbstats "github.com/matrixorigin/matrixone/pkg/pb/statsinfo" "github.com/matrixorigin/matrixone/pkg/pb/status" "github.com/matrixorigin/matrixone/pkg/pb/timestamp" "github.com/matrixorigin/matrixone/pkg/perfcounter" @@ -295,8 +296,10 @@ type Session struct { planCache *planCache - statsCache *plan2.StatsCache - seqCurValues map[uint64]string + statsCacheMu sync.Mutex + statsCache *plan2.StatsCache + statsCacheVersions map[uint64]optimizerStatsCacheTag + seqCurValues map[uint64]string /* CORNER CASE: @@ -808,9 +811,85 @@ func (ses *Session) GetProc() *process.Process { } func (ses *Session) GetStatsCache() *plan2.StatsCache { + ses.statsCacheMu.Lock() + defer ses.statsCacheMu.Unlock() return ses.statsCache } +func (ses *Session) optimizerStatsKey(tableID uint64) optimizerStatsTableKey { + return optimizerStatsTableKey{ + accountID: ses.GetAccountId(), + tableID: tableID, + } +} + +type optimizerStatsCacheTag struct { + key optimizerStatsTableKey + version uint64 +} + +func (ses *Session) getStatsCacheWithVersion(key optimizerStatsTableKey) (*plan2.StatsCache, uint64) { + ses.statsCacheMu.Lock() + defer ses.statsCacheMu.Unlock() + ses.initStatsCacheLocked() + version := currentOptimizerStatsVersion(ses.GetService(), key) + wrapper := ses.statsCache.Get(key.tableID) + tag, tagged := ses.statsCacheVersions[key.tableID] + if !wrapper.Exists() { + delete(ses.statsCacheVersions, key.tableID) + } else if !tagged && version == 0 { + // Accept caches created before version tracking only in the initial + // generation. Once any publication has happened, an untagged entry is + // conservatively stale. + ses.statsCacheVersions[key.tableID] = optimizerStatsCacheTag{key: key, version: version} + } else if tag.key != key || tag.version != version { + ses.statsCache.Delete(key.tableID) + delete(ses.statsCacheVersions, key.tableID) + } + return ses.statsCache, version +} + +func (ses *Session) cacheStatsIfCurrent( + key optimizerStatsTableKey, + version uint64, + stats *pbstats.StatsInfo, +) bool { + ses.statsCacheMu.Lock() + defer ses.statsCacheMu.Unlock() + if currentOptimizerStatsVersion(ses.GetService(), key) != version { + return false + } + ses.initStatsCacheLocked() + if ses.statsCache.SetAndReportReset(key.tableID, stats) { + clear(ses.statsCacheVersions) + } + ses.statsCacheVersions[key.tableID] = optimizerStatsCacheTag{key: key, version: version} + return true +} + +func (ses *Session) cachePublishedStats( + key optimizerStatsTableKey, + version uint64, + stats *pbstats.StatsInfo, +) { + ses.statsCacheMu.Lock() + defer ses.statsCacheMu.Unlock() + ses.initStatsCacheLocked() + if ses.statsCache.SetAndReportReset(key.tableID, stats) { + clear(ses.statsCacheVersions) + } + ses.statsCacheVersions[key.tableID] = optimizerStatsCacheTag{key: key, version: version} +} + +func (ses *Session) initStatsCacheLocked() { + if ses.statsCache == nil { + ses.statsCache = plan2.NewStatsCache() + } + if ses.statsCacheVersions == nil { + ses.statsCacheVersions = make(map[uint64]optimizerStatsCacheTag) + } +} + func (ses *Session) GetSessionStart() time.Time { ses.mu.Lock() defer ses.mu.Unlock() @@ -1172,7 +1251,6 @@ func NewSession( var txnOp TxnOperator var err error txnHandler := InitTxnHandler(service, getPu(service).StorageEngine, connCtx, txnOp) - ses := &Session{ feSessionImpl: feSessionImpl{ pool: mp, @@ -1195,8 +1273,9 @@ func NewSession( startedAt: time.Now(), connType: ConnTypeUnset, - timestampMap: map[TS]time.Time{}, - statsCache: plan2.NewStatsCache(), + timestampMap: map[TS]time.Time{}, + statsCache: plan2.NewStatsCache(), + statsCacheVersions: make(map[uint64]optimizerStatsCacheTag), } atomic.StoreInt32(&ses.sqlModeNoAutoValueOnZero, -1) @@ -1439,8 +1518,21 @@ func (ses *Session) IsBackgroundSession() bool { } func (ses *Session) cachePlan(sql string, stmts []tree.Statement, plans []*plan.Plan, versions ...int64) { - ses.cachePlanWithSnapshots( - sql, stmts, plans, make([]timestamp.Timestamp, len(plans)), versions...) + ses.cachePlanWithSnapshotsAndStatsVersions( + sql, stmts, plans, make([]timestamp.Timestamp, len(plans)), + make([]map[optimizerStatsTableKey]uint64, len(plans)), versions...) +} + +func (ses *Session) cachePlanWithStatsVersions( + sql string, + stmts []tree.Statement, + plans []*plan.Plan, + statsVersions map[optimizerStatsTableKey]uint64, + versions ...int64, +) { + ses.cachePlanWithSnapshotsAndStatsVersions( + sql, stmts, plans, make([]timestamp.Timestamp, len(plans)), + planStatsVersionsFromAggregate(len(plans), statsVersions), versions...) } func (ses *Session) cachePlanWithSnapshots( @@ -1449,10 +1541,30 @@ func (ses *Session) cachePlanWithSnapshots( plans []*plan.Plan, planSnapshotTS []timestamp.Timestamp, versions ...int64, +) { + ses.cachePlanWithSnapshotsAndStatsVersions( + sql, stmts, plans, planSnapshotTS, + make([]map[optimizerStatsTableKey]uint64, len(plans)), versions...) +} + +func (ses *Session) cachePlanWithSnapshotsAndStatsVersions( + sql string, + stmts []tree.Statement, + plans []*plan.Plan, + planSnapshotTS []timestamp.Timestamp, + planStatsVersions []map[optimizerStatsTableKey]uint64, + versions ...int64, ) { if len(sql) == 0 { return } + statsVersions, versionsConsistent := aggregatePlanStatsVersions(planStatsVersions) + if !versionsConsistent || !optimizerStatsVersionsCurrent(ses.GetService(), statsVersions) { + // The plan crossed a statistics publication boundary while compiling. + // It may execute, but must not enter the cache with stale dependencies. + freeStmts(stmts) + return + } ses.mu.Lock() defer ses.mu.Unlock() if ses.planCache == nil { @@ -1463,7 +1575,8 @@ func (ses *Session) cachePlanWithSnapshots( if len(versions) > 0 { protocolVersion = versions[0] } - ses.planCache.cacheWithPlanSnapshots(sql, stmts, plans, planSnapshotTS, protocolVersion) + ses.planCache.cacheWithPlanSnapshotsAndStatsVersions( + sql, stmts, plans, planSnapshotTS, planStatsVersions, protocolVersion) } func (ses *Session) getCachedPlan(sql string) *cachedPlan { @@ -1476,7 +1589,8 @@ func (ses *Session) getCachedPlan(sql string) *cachedPlan { return nil } cached := ses.planCache.get(sql) - if cached != nil && cached.protocolVersion != currentProtocolVersion(ses.proc) { + if cached != nil && (cached.protocolVersion != currentProtocolVersion(ses.proc) || + !optimizerStatsVersionsCurrent(ses.GetService(), cached.statsVersions)) { ses.planCache.remove(sql) return nil } @@ -1492,7 +1606,18 @@ func (ses *Session) isCached(sql string) bool { if ses.planCache == nil { return false } - return ses.planCache.isCached(sql) + if !ses.planCache.isCached(sql) { + return false + } + cached := ses.planCache.cachePool[sql].Value.(*cachedPlan) + if cached.protocolVersion != currentProtocolVersion(ses.proc) || + !optimizerStatsVersionsCurrent(ses.GetService(), cached.statsVersions) { + // isCached is also queried while wrappers still borrow the cached AST at + // the end of execution. Report staleness without releasing that owner; + // the next getCachedPlan lookup removes it after all borrowers are gone. + return false + } + return true } func (ses *Session) removeCachedPlan(sql string) { @@ -1512,6 +1637,7 @@ func (ses *Session) updateCachedPlanGeneration( expectedPlan *plan.Plan, newPlan *plan.Plan, planSnapshotTS timestamp.Timestamp, + statsVersions map[optimizerStatsTableKey]uint64, ) bool { if len(sql) == 0 { return false @@ -1522,7 +1648,7 @@ func (ses *Session) updateCachedPlanGeneration( return false } return ses.planCache.updatePlanGeneration( - sql, index, expectedPlan, newPlan, planSnapshotTS) + sql, index, expectedPlan, newPlan, planSnapshotTS, statsVersions) } func (ses *Session) invalidateCachedPlanGeneration( diff --git a/pkg/frontend/types.go b/pkg/frontend/types.go index 59a4aa3f6acaf..41af120003a7f 100644 --- a/pkg/frontend/types.go +++ b/pkg/frontend/types.go @@ -2063,4 +2063,10 @@ type ServerLevelVariables struct { Aicm atomic.Value moServerStarted atomic.Bool sessionAlloc atomic.Value + + optimizerStatsMu sync.RWMutex + optimizerStatsClock uint64 + optimizerStatsReset uint64 + optimizerStatsVersions map[optimizerStatsTableKey]uint64 + optimizerStatsPublish [optimizerStatsPublisherStripes]chan struct{} } diff --git a/pkg/sql/plan/stats.go b/pkg/sql/plan/stats.go index 9f679bf01426e..08974c66e0c53 100644 --- a/pkg/sql/plan/stats.go +++ b/pkg/sql/plan/stats.go @@ -149,17 +149,38 @@ func (sc *StatsCache) Get(tableID uint64) StatsInfoWrapper { // Set caches the stats result for the table. func (sc *StatsCache) Set(tableID uint64, stats *pb.StatsInfo) { + sc.set(tableID, stats) +} + +// SetAndReportReset caches stats and reports whether the bounded cache evicted +// all prior entries. Callers that keep side metadata can reset it in lockstep. +func (sc *StatsCache) SetAndReportReset(tableID uint64, stats *pb.StatsInfo) bool { + return sc.set(tableID, stats) +} + +func (sc *StatsCache) set(tableID uint64, stats *pb.StatsInfo) bool { if sc == nil { - return + return false } + reset := false if len(sc.cache) > statsCacheMaxSize { sc.cache = make(map[uint64]StatsInfoWrapper, statsCacheInitSize) logutil.Infof("statscache entries more than %v in long session, release memory", statsCacheMaxSize) + reset = true } sc.cache[tableID] = StatsInfoWrapper{ stats: stats, lastVisit: time.Now().Unix(), } + return reset +} + +// Delete removes one table without disturbing unrelated session statistics. +func (sc *StatsCache) Delete(tableID uint64) { + if sc == nil { + return + } + delete(sc.cache, tableID) } func NewStatsInfo() *pb.StatsInfo { diff --git a/pkg/sql/plan/stats_test.go b/pkg/sql/plan/stats_test.go index 9a74cd62cf693..b3cce5ca3fa76 100644 --- a/pkg/sql/plan/stats_test.go +++ b/pkg/sql/plan/stats_test.go @@ -345,6 +345,19 @@ func newStatsTestBuilderWithNDV(colName string, ndv float64) *QueryBuilder { return builder } +func TestStatsCacheReportsWholeCacheReset(t *testing.T) { + statsCache := NewStatsCache() + stats := NewStatsInfo() + for tableID := uint64(0); tableID <= statsCacheMaxSize; tableID++ { + require.False(t, statsCache.SetAndReportReset(tableID, stats)) + } + require.True(t, statsCache.SetAndReportReset(statsCacheMaxSize+1, stats)) + removed := statsCache.Get(0) + retained := statsCache.Get(statsCacheMaxSize + 1) + require.False(t, removed.Exists()) + require.True(t, retained.Exists()) +} + type statsCacheCompilerContext struct { *MockCompilerContext statsCache *StatsCache diff --git a/pkg/vm/engine/disttae/engine.go b/pkg/vm/engine/disttae/engine.go index 07198d6d6f3bb..d7dd008cb0ae8 100644 --- a/pkg/vm/engine/disttae/engine.go +++ b/pkg/vm/engine/disttae/engine.go @@ -58,6 +58,7 @@ import ( ) var _ engine.Engine = new(Engine) +var _ engine.StatsRefresher = new(Engine) const ( workspaceRSSCacheFamilyEvictTimeout = 10 * time.Second @@ -1278,6 +1279,20 @@ func (e *Engine) Stats(ctx context.Context, key pb.StatsInfoKey, sync bool) *pb. return e.globalStats.Get(ctx, key, sync) } +// RefreshTableStats synchronously replaces the local optimizer statistics for +// key. The cache swap is the publication boundary observed by later plans. +func (e *Engine) RefreshTableStats(ctx context.Context, key pb.StatsInfoKey) (*pb.StatsInfo, error) { + return refreshTableStats(ctx, key, e.globalStats) +} + +type optimizerStatsStore interface { + refreshStatsWithMode(context.Context, pb.StatsInfoKey, string) (*pb.StatsInfo, error) +} + +func refreshTableStats(ctx context.Context, key pb.StatsInfoKey, store optimizerStatsStore) (*pb.StatsInfo, error) { + return store.refreshStatsWithMode(ctx, key, "auto") +} + // GetGlobalStats returns the GlobalStats instance func (e *Engine) GetGlobalStats() *GlobalStats { return e.globalStats diff --git a/pkg/vm/engine/disttae/engine_stats_test.go b/pkg/vm/engine/disttae/engine_stats_test.go new file mode 100644 index 0000000000000..0231213777bea --- /dev/null +++ b/pkg/vm/engine/disttae/engine_stats_test.go @@ -0,0 +1,336 @@ +// Copyright 2026 Matrix Origin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package disttae + +import ( + "context" + "errors" + "sync" + "testing" + "time" + + "github.com/matrixorigin/matrixone/pkg/objectio" + pb "github.com/matrixorigin/matrixone/pkg/pb/statsinfo" + "github.com/matrixorigin/matrixone/pkg/util/fault" + "github.com/stretchr/testify/require" +) + +type optimizerStatsStoreStub struct { + stats *pb.StatsInfo + refreshErr error + key pb.StatsInfoKey + mode string +} + +func (s *optimizerStatsStoreStub) refreshStatsWithMode( + _ context.Context, + key pb.StatsInfoKey, + mode string, +) (*pb.StatsInfo, error) { + s.key = key + s.mode = mode + return s.stats, s.refreshErr +} + +func TestRefreshTableStatsDefinesPublicationBoundary(t *testing.T) { + key := pb.StatsInfoKey{TableID: 42, DbName: "db", TableName: "events"} + t.Run("success", func(t *testing.T) { + fresh := &pb.StatsInfo{TableCnt: 1_000_000} + store := &optimizerStatsStoreStub{stats: fresh} + + got, err := refreshTableStats(context.Background(), key, store) + require.NoError(t, err) + require.Same(t, fresh, got) + require.Equal(t, key, store.key) + require.Equal(t, "auto", store.mode) + }) + + t.Run("refresh failure is not published", func(t *testing.T) { + wantErr := errors.New("refresh failed") + store := &optimizerStatsStoreStub{refreshErr: wantErr} + + got, err := refreshTableStats(context.Background(), key, store) + require.ErrorIs(t, err, wantErr) + require.Nil(t, got) + }) +} + +func TestOptimizerStatsRefreshAdmissionIsTableScopedAndCancelable(t *testing.T) { + gs := &GlobalStats{} + gs.initStatsRefreshAdmission() + key := pb.StatsInfoKey{AccId: 1, TableID: 42} + + release, err := gs.acquireStatsRefresh(context.Background(), key) + require.NoError(t, err) + defer release() + + otherKey := key + otherKey.TableID++ + releaseOther, err := gs.acquireStatsRefresh(context.Background(), otherKey) + require.NoError(t, err, "unrelated tables must not queue behind this table") + releaseOther() + + canceled, cancel := context.WithCancel(context.Background()) + cancel() + releaseCanceled, err := gs.acquireStatsRefresh(canceled, key) + require.ErrorIs(t, err, context.Canceled) + require.Nil(t, releaseCanceled) +} + +func TestOptimizerStatsRefreshAdmissionRejectsStoppedOwner(t *testing.T) { + ownerCtx, stopOwner := context.WithCancel(context.Background()) + stopOwner() + gs := &GlobalStats{ctx: ownerCtx} + gs.initStatsRefreshAdmission() + + release, err := gs.acquireStatsRefresh( + context.Background(), pb.StatsInfoKey{AccId: 1, TableID: 42}) + require.ErrorIs(t, err, context.Canceled) + require.Nil(t, release, + "a stopped GlobalStats owner must not transfer a refresh admission token") +} + +func TestStatsRefreshContextPreservesRequestAndObservesOwnerStop(t *testing.T) { + type requestValueKey struct{} + ownerCtx, stopOwner := context.WithCancel(context.Background()) + requestCtx := context.WithValue(context.Background(), requestValueKey{}, "request-value") + gs := &GlobalStats{ctx: ownerCtx} + + refreshCtx, stopRefresh, err := gs.newStatsRefreshContext(requestCtx) + require.NoError(t, err) + t.Cleanup(stopRefresh) + require.Equal(t, "request-value", refreshCtx.Value(requestValueKey{})) + + stopOwner() + select { + case <-refreshCtx.Done(): + require.ErrorIs(t, context.Cause(refreshCtx), context.Canceled) + case <-time.After(time.Second): + t.Fatal("owner shutdown did not cancel downstream refresh work") + } +} + +func TestStatsRefreshContextClosesOwnerWatcherRegistrationRace(t *testing.T) { + ownerCtx := newDelayedLifecycleContext() + ownerCtx.cancelOnRegister = true + gs := &GlobalStats{ctx: ownerCtx} + + refreshCtx, stopRefresh, err := gs.newStatsRefreshContext(context.Background()) + require.ErrorIs(t, err, context.Canceled) + require.Nil(t, refreshCtx) + require.Nil(t, stopRefresh) +} + +func TestCoordinateStatsUpdateCancellationReleasesUpdateGeneration(t *testing.T) { + gs := &GlobalStats{} + gs.initStatsRefreshAdmission() + gs.updatingMu.updating = make(map[pb.StatsInfoKey]*updateRecord) + key := pb.StatsInfoKey{AccId: 1, TableID: 42} + + release, err := gs.acquireStatsRefresh(context.Background(), key) + require.NoError(t, err) + defer release() + + canceled, cancel := context.WithCancel(context.Background()) + cancel() + generation := gs.currentOrCreateUpdateRecord(key) + gs.coordinateStatsUpdateJob(statsUpdateJob{ + wrapKey: pb.StatsInfoKeyWithContext{Ctx: canceled, Key: key}, + expectedRecord: generation, + }) + + gs.updatingMu.Lock() + record := gs.updatingMu.updating[key] + gs.updatingMu.Unlock() + require.NotNil(t, record) + require.False(t, record.inProgress, + "cancellation while waiting for refresh admission must close the update generation") +} + +func TestCompleteAutomaticStatsRefreshKeepsMetadataInsideAdmission(t *testing.T) { + gs := &GlobalStats{} + gs.initStatsRefreshAdmission() + gs.updatingMu.updating = make(map[pb.StatsInfoKey]*updateRecord) + gs.mu.statsInfoMap = make(map[pb.StatsInfoKey]*pb.StatsInfo) + gs.mu.cond = sync.NewCond(&gs.mu) + key := pb.StatsInfoKey{AccId: 1, TableID: 42} + + generation := &updateRecord{inProgress: true} + gs.updatingMu.updating[key] = generation + oldRelease, err := gs.acquireStatsRefresh(context.Background(), key) + require.NoError(t, err) + + newerDone := make(chan error, 1) + go func() { + newRelease, acquireErr := gs.acquireStatsRefresh(context.Background(), key) + if acquireErr != nil { + newerDone <- acquireErr + return + } + gs.markAutomaticUpdateComplete(key, generation, true, 100, 0.5) + newRelease() + newerDone <- nil + }() + + // Waiting here after releasing forces the newer refresh to commit between + // release and any code that might incorrectly update the old baseline late. + var newerErr error + gs.completeAutomaticStatsRefresh( + key, generation, &pb.StatsInfo{}, true, 1, 1.0, func() { + oldRelease() + newerErr = <-newerDone + }) + require.NoError(t, newerErr) + + gs.updatingMu.Lock() + record := *gs.updatingMu.updating[key] + gs.updatingMu.Unlock() + require.False(t, record.inProgress) + require.Equal(t, int64(100), record.baseObjectCount) + require.Equal(t, 0.5, record.samplingRatio) +} + +func TestCoordinateStatsUpdateSubscribeFailurePreservesLastPublishedStats(t *testing.T) { + key := pb.StatsInfoKey{ + AccId: 1, DatabaseID: 10, TableID: 42, DbName: "db", TableName: "events", + } + newStats := func() *pb.StatsInfo { + return &pb.StatsInfo{TableCnt: 1_000_000} + } + newGlobalStats := func() *GlobalStats { + gs := &GlobalStats{engine: &Engine{}} + gs.initStatsRefreshAdmission() + gs.updatingMu.updating = make(map[pb.StatsInfoKey]*updateRecord) + gs.mu.statsInfoMap = make(map[pb.StatsInfoKey]*pb.StatsInfo) + gs.mu.cond = sync.NewCond(&gs.mu) + return gs + } + + fault.Enable() + t.Cleanup(func() { fault.Disable() }) + removeFault, err := objectio.InjectLogging( + objectio.FJ_CNSubscribeTableFail, key.DbName, key.TableName, 0, true, + ) + require.NoError(t, err) + t.Cleanup(removeFault) + + t.Run("retain last successful publication", func(t *testing.T) { + gs := newGlobalStats() + lastGood := newStats() + gs.mu.statsInfoMap[key] = lastGood + + generation := gs.currentOrCreateUpdateRecord(key) + gs.coordinateStatsUpdateJob(statsUpdateJob{ + wrapKey: pb.StatsInfoKeyWithContext{Ctx: context.Background(), Key: key}, + expectedRecord: generation, + }) + + gs.mu.Lock() + got, exists := gs.mu.statsInfoMap[key] + gs.mu.Unlock() + require.True(t, exists) + require.Same(t, lastGood, got, + "a failed automatic refresh must not erase the last successful publication") + }) + + t.Run("complete first failed generation", func(t *testing.T) { + gs := newGlobalStats() + + generation := gs.currentOrCreateUpdateRecord(key) + gs.coordinateStatsUpdateJob(statsUpdateJob{ + wrapKey: pb.StatsInfoKeyWithContext{Ctx: context.Background(), Key: key}, + expectedRecord: generation, + }) + + gs.mu.Lock() + got, exists := gs.mu.statsInfoMap[key] + gs.mu.Unlock() + require.True(t, exists, + "the first failed automatic generation must still wake synchronous waiters") + require.Nil(t, got) + }) +} + +func TestExplicitStatsRefreshSubscribeFailureDoesNotRetainGeneration(t *testing.T) { + key := pb.StatsInfoKey{ + AccId: 1, DatabaseID: 10, TableID: 42, DbName: "db", TableName: "events", + } + gs := &GlobalStats{engine: &Engine{}} + gs.initStatsRefreshAdmission() + gs.updatingMu.updating = make(map[pb.StatsInfoKey]*updateRecord) + + fault.Enable() + t.Cleanup(func() { fault.Disable() }) + removeFault, err := objectio.InjectLogging( + objectio.FJ_CNSubscribeTableFail, key.DbName, key.TableName, 0, true, + ) + require.NoError(t, err) + t.Cleanup(removeFault) + + stats, err := gs.refreshStatsWithMode(context.Background(), key, "auto") + require.Error(t, err) + require.Nil(t, stats) + + gs.updatingMu.Lock() + _, retained := gs.updatingMu.updating[key] + gs.updatingMu.Unlock() + require.False(t, retained, + "a failed subscription has no cleanup owner and must not create a generation") +} + +func TestInitialStatsGetSubscribeFailureDoesNotQueueOwnerlessGeneration(t *testing.T) { + key := pb.StatsInfoKey{ + AccId: 1, DatabaseID: 10, TableID: 42, DbName: "db", TableName: "events", + } + gs := &GlobalStats{ + engine: &Engine{}, + updateC: make(chan statsUpdateJob, 1), + queueWatcher: newQueueWatcher(), + } + gs.updatingMu.updating = make(map[pb.StatsInfoKey]*updateRecord) + gs.mu.statsInfoMap = make(map[pb.StatsInfoKey]*pb.StatsInfo) + gs.mu.cond = sync.NewCond(&gs.mu) + + fault.Enable() + t.Cleanup(func() { fault.Disable() }) + removeFault, err := objectio.InjectLogging( + objectio.FJ_CNSubscribeTableFail, key.DbName, key.TableName, 0, true, + ) + require.NoError(t, err) + t.Cleanup(removeFault) + + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + result := make(chan *pb.StatsInfo, 1) + go func() { result <- gs.Get(ctx, key, true) }() + + select { + case got := <-result: + require.Nil(t, got) + case <-gs.updateC: + cancel() + <-result + t.Fatal("a failed initial subscription queued work without a cleanup owner") + case <-time.After(time.Second): + cancel() + t.Fatal("stats get did not terminate after subscription failure") + } + + gs.updatingMu.Lock() + _, retained := gs.updatingMu.updating[key] + gs.updatingMu.Unlock() + require.False(t, retained) + require.Empty(t, gs.updateC) +} diff --git a/pkg/vm/engine/disttae/stats.go b/pkg/vm/engine/disttae/stats.go index 11162cd984ee1..a47c9619bc4ed 100644 --- a/pkg/vm/engine/disttae/stats.go +++ b/pkg/vm/engine/disttae/stats.go @@ -17,6 +17,7 @@ package disttae import ( "context" "encoding/binary" + "errors" "math" "runtime" "sync" @@ -54,7 +55,7 @@ import ( // logtailConsumer (1个 goroutine) // │ // │ 判断入队条件（第一层）： -// │ - keyExists(): key 必须已存在 +// │ - cache/generation 原子检查：key 必须已存在 // │ - CkpLocation: checkpoint 时触发 // │ - MetaEntry: object 元数据变更时触发 // │ @@ -65,10 +66,10 @@ import ( // spawnUpdateWorkers (16-27个 goroutine) // │ // │ 判断执行条件（第二层）： 便于统一 debounce force/normal update request -// │ - shouldExecuteUpdate(): 检查 inProgress 和 MinUpdateInterval (15s) +// │ - startAutomaticUpdate(): 检查 generation、inProgress 和 MinUpdateInterval (15s) // │ // ▼ -// coordinateStatsUpdate() +// coordinateStatsUpdateJob() // │ // ├─→ 订阅表获取 PartitionState // ├─→ 从 CatalogCache 获取 TableDef @@ -181,6 +182,8 @@ type GlobalStatsConfig struct { LogtailUpdateStatsThreshold int } +const optimizerStatsRefreshStripes = 64 + type GlobalStatsOption func(s *GlobalStats) // WithUpdateWorkerFactor set the update worker factor. @@ -199,6 +202,11 @@ func WithApproxObjectNumUpdater(f func() int64) GlobalStatsOption { // updateRecord records the update status of a key. type updateRecord struct { + // queued is the number of registered enqueue attempts that have not reached + // worker admission or rolled back. It includes a forced sender blocked on a + // full queue. Together with inProgress it is the durable predicate that + // proves a synchronous waiter still has a producer. + queued int // inProgress indicates if the stats of a table is being updated. inProgress bool // lastUpdate is the time of the stats last updated. @@ -211,6 +219,17 @@ type updateRecord struct { samplingRatio float64 } +// statsUpdateJob carries the scheduling generation observed by the producer. +// Pointer identity prevents a queued worker from publishing into a newer table +// generation after RemoveTid deleted the old record. +type statsUpdateJob struct { + wrapKey pb.StatsInfoKeyWithContext + expectedRecord *updateRecord + // registered means enqueueStatsUpdateForRecord accounted this job in + // expectedRecord. Direct test helpers leave it false. + registered bool +} + type GlobalStats struct { ctx context.Context @@ -222,7 +241,7 @@ type GlobalStats struct { // TODO(volgariver6): add metrics of the chan length. tailC chan *logtail.TableLogtail - updateC chan pb.StatsInfoKeyWithContext + updateC chan statsUpdateJob // queueWatcher keeps the table id and its enqueue time. // and watch the queue item in the queue. @@ -233,6 +252,11 @@ type GlobalStats struct { updating map[pb.StatsInfoKey]*updateRecord } + // Explicit ANALYZE refreshes and automatic logtail refreshes share this + // bounded admission layer. The same table is calculated and published in + // order; unrelated tables normally remain parallel. + refreshAdmission [optimizerStatsRefreshStripes]chan struct{} + // statsInfoMap is the global stats info in engine which // contains all subscribed tables stats info. mu struct { @@ -263,6 +287,14 @@ type GlobalStats struct { // beforeSubscribeTable is for test only. beforeSubscribeTable func(pb.StatsInfoKey) + + // beforeStatsWait is for deterministic wait-protocol tests only. It runs + // with gs.mu held immediately before cond.Wait atomically releases it. + beforeStatsWait func(pb.StatsInfoKey, *updateRecord) + + // afterAutomaticUpdateStarted is for deterministic producer-cancellation + // tests only. It runs after worker admission and before refresh admission. + afterAutomaticUpdateStarted func(pb.StatsInfoKey, *updateRecord) } func NewGlobalStats( @@ -272,13 +304,17 @@ func NewGlobalStats( ctx: ctx, engine: e, tailC: make(chan *logtail.TableLogtail, 10000), - updateC: make(chan pb.StatsInfoKeyWithContext, 3000), + updateC: make(chan statsUpdateJob, 3000), KeyRouter: keyRouter, queueWatcher: newQueueWatcher(), } s.updatingMu.updating = make(map[pb.StatsInfoKey]*updateRecord) s.mu.statsInfoMap = make(map[pb.StatsInfoKey]*pb.StatsInfo) s.mu.cond = sync.NewCond(&s.mu) + // One lifecycle callback wakes every current waiter when update workers + // stop. Register it once per GlobalStats rather than once per cache miss. + context.AfterFunc(ctx, s.notifyStatsWaiters) + s.initStatsRefreshAdmission() for _, opt := range opts { opt(s) } @@ -316,26 +352,112 @@ func NewGlobalStats( return s } -// keyExists returns true only if key already exists in the map. -func (gs *GlobalStats) keyExists(key pb.StatsInfoKey) bool { - gs.mu.Lock() - defer gs.mu.Unlock() - _, ok := gs.mu.statsInfoMap[key] - return ok +func (gs *GlobalStats) initStatsRefreshAdmission() { + for i := range gs.refreshAdmission { + gs.refreshAdmission[i] = make(chan struct{}, 1) + } +} + +func optimizerStatsRefreshStripe(key pb.StatsInfoKey) int { + mixed := key.TableID ^ uint64(key.AccId)*0x9e3779b97f4a7c15 + return int(mixed % optimizerStatsRefreshStripes) } -// RemoveTid removes all statsInfoMap entries for the given table ID. +func (gs *GlobalStats) acquireStatsRefresh( + ctx context.Context, + key pb.StatsInfoKey, +) (func(), error) { + if cause := gs.statsRefreshCancellationCause(ctx); cause != nil { + return nil, cause + } + admission := gs.refreshAdmission[optimizerStatsRefreshStripe(key)] + select { + case admission <- struct{}{}: + // Cancellation can race the select and make both cases ready. Recheck + // the authoritative caller and owner contexts before transferring the + // admission token to the caller. + if cause := gs.statsRefreshCancellationCause(ctx); cause != nil { + <-admission + return nil, cause + } + return func() { <-admission }, nil + case <-ctx.Done(): + return nil, context.Cause(ctx) + case <-gs.lifecycleDone(): + return nil, gs.statsRefreshCancellationCause(ctx) + } +} + +// statsRefreshCancellationCause treats the request context and the GlobalStats +// owner lifecycle as durable predicates. Async callbacks may wake blocked work, +// but they are never the source of truth for admission or publication. +func (gs *GlobalStats) statsRefreshCancellationCause(ctx context.Context) error { + if cause := context.Cause(ctx); cause != nil { + return cause + } + if gs.ctx != nil { + return context.Cause(gs.ctx) + } + return nil +} + +// newStatsRefreshContext preserves request values and deadlines while linking +// downstream subscription and object I/O to the GlobalStats owner lifecycle. +// The returned context delivers cancellation; admission and publication still +// re-read statsRefreshCancellationCause as their durable predicate. +func (gs *GlobalStats) newStatsRefreshContext( + ctx context.Context, +) (context.Context, func(), error) { + if cause := gs.statsRefreshCancellationCause(ctx); cause != nil { + return nil, nil, cause + } + if gs.ctx == nil { + return ctx, func() {}, nil + } + refreshCtx, cancelRefresh := context.WithCancelCause(ctx) + stopOwnerWatch := context.AfterFunc(gs.ctx, func() { + cause := context.Cause(gs.ctx) + if cause == nil { + cause = context.Canceled + } + cancelRefresh(cause) + }) + stop := func() { + stopOwnerWatch() + cancelRefresh(nil) + } + // Close the owner check/register race without depending on callback + // dispatch. No downstream operation is admitted on the error path. + if cause := gs.statsRefreshCancellationCause(ctx); cause != nil { + stop() + return nil, nil, cause + } + return refreshCtx, stop, nil +} + +// RemoveTid removes every GlobalStats entry owned by the given table ID. // Called from cleanMemoryTableWithTable (1+ hour after unsubscribe/drop) -// to prevent unbounded map growth. Safe because no queries target a -// table that has been unsubscribed for over an hour. +// to prevent both published statistics and refresh-scheduling metadata from +// growing for the process lifetime. Safe because no queries target a table +// that has been unsubscribed for over an hour. func (gs *GlobalStats) RemoveTid(tableID uint64) { + // Keep the established gs.mu -> updatingMu lock order used by + // broadcastStats. Table cleanup is the common lifetime boundary for both + // maps, so a removed table cannot retain one owner after losing the other. gs.mu.Lock() defer gs.mu.Unlock() + gs.updatingMu.Lock() + defer gs.updatingMu.Unlock() for key := range gs.mu.statsInfoMap { if key.TableID == tableID { delete(gs.mu.statsInfoMap, key) } } + for key := range gs.updatingMu.updating { + if key.TableID == tableID { + delete(gs.updatingMu.updating, key) + } + } if gs.mu.cond != nil { gs.mu.cond.Broadcast() } @@ -346,7 +468,69 @@ func (gs *GlobalStats) PrefetchTableMeta(ctx context.Context, key pb.StatsInfoKe Ctx: ctx, Key: key, } - return gs.enqueueStatsUpdate(wrapkey, false) + generation, ok := gs.currentOrCreateSubscribedUpdateRecord(key) + if !ok { + return false + } + return gs.enqueueStatsUpdateForRecord(wrapkey, false, generation) +} + +// currentOrCreateUpdateRecord returns the table-lifetime token that every +// queued or explicit refresh must carry. In particular, the first refresh must +// not use nil as an "expected absence" token: RemoveTid can make absence true +// again, allowing old queued work to cross the cleanup boundary and recreate +// table-owned state. Production callers must already hold or have proved a +// cleanup owner; use currentOrCreateSubscribedUpdateRecord for request-driven +// work and shouldEnqueueExistingStatsUpdateGeneration for logtail work. +func (gs *GlobalStats) currentOrCreateUpdateRecord(key pb.StatsInfoKey) *updateRecord { + gs.updatingMu.Lock() + defer gs.updatingMu.Unlock() + rec := gs.updatingMu.updating[key] + if rec == nil { + rec = &updateRecord{} + gs.updatingMu.updating[key] = rec + } + return rec +} + +// currentOrCreateSubscribedUpdateRecord captures the scheduling generation +// only while the subscription that owns its cleanup is still current. Explicit +// refreshes call this after toSubscribeTable succeeds, so a failed subscription +// cannot leave metadata that no unsubscribe path can reclaim. Holding the +// subscription read lock across record creation also prevents cleanup from +// falling between subscription validation and token capture. +func (gs *GlobalStats) currentOrCreateSubscribedUpdateRecord( + key pb.StatsInfoKey, +) (*updateRecord, bool) { + gs.engine.pClient.subscribed.rw.RLock() + defer gs.engine.pClient.subscribed.rw.RUnlock() + + ent, ok := gs.engine.pClient.subscribed.m[key.TableID] + if !ok || ent == nil || ent.dbID != key.DatabaseID || ent.state != Subscribed { + return nil, false + } + return gs.currentOrCreateUpdateRecord(key), true +} + +// currentOrCreateExactSubscribedUpdateRecord is the first-read variant. It +// validates the exact subscription generation before mutating scheduling +// state, so an old read cannot create even idle metadata for a replacement +// lifetime. +func (gs *GlobalStats) currentOrCreateExactSubscribedUpdateRecord( + key pb.StatsInfoKey, + expectedEnt *subEntry, +) (*updateRecord, bool) { + if expectedEnt == nil { + return nil, false + } + gs.engine.pClient.subscribed.rw.RLock() + defer gs.engine.pClient.subscribed.rw.RUnlock() + + ent, ok := gs.engine.pClient.subscribed.m[key.TableID] + if !ok || ent != expectedEnt || ent.dbID != key.DatabaseID || ent.state != Subscribed { + return nil, false + } + return gs.currentOrCreateUpdateRecord(key), true } func (gs *GlobalStats) subscribedEntry(key pb.StatsInfoKey) *subEntry { @@ -422,11 +606,23 @@ func (gs *GlobalStats) Get(ctx context.Context, key pb.StatsInfoKey, sync bool) key.DatabaseID, key.DbName) - if err == nil && ps.ApproxDataObjectsNum() == 0 { + if err != nil { + // A failed initial subscription has no table-lifetime cleanup owner. + // Retrying through updateC would create a scheduling generation (and + // potentially a nil cache sentinel) that RemoveTid can never reclaim. + return nil + } + if ps.ApproxDataObjectsNum() == 0 { return nil } subscribedEnt := gs.subscribedEntry(key) + if subscribedEnt == nil { + // Cleanup crossed the successful subscribe return before this read could + // capture the exact lifetime. Do not let nil mean "accept any later + // subscription" when the synchronous producer is created below. + return nil + } var remoteInfo *pb.StatsInfo if _, ok = ctx.Value(perfcounter.CalcTableStatsKey{}).(bool); ok { stats := statistic.StatsInfoFromContext(ctx) @@ -460,57 +656,71 @@ func (gs *GlobalStats) Get(ctx context.Context, key pb.StatsInfoKey, sync bool) } } - if sync { - stopWake := context.AfterFunc(ctx, func() { - gs.mu.Lock() - gs.mu.cond.Broadcast() - gs.mu.Unlock() - }) - defer stopWake() - } - + // Another producer may have published while subscription or remote lookup + // was in progress. Preserve this recheck for both synchronous and non-blocking + // callers. For a synchronous caller, an existing nil sentinel still admits a + // background retry, as before. gs.mu.Lock() - defer gs.mu.Unlock() - - // Recheck local cache after lock reacquired, another goroutine may have updated it. - info, ok = gs.mu.statsInfoMap[key] - if ok && info != nil { + info = gs.mu.statsInfoMap[key] + gs.mu.Unlock() + if info != nil { return info } + if !sync { + return nil + } - ok = false - if sync { - for !ok { - if ctx.Err() != nil { - return nil - } + // Capture the producer generation only while the exact subscription + // observed above still owns cleanup. A replacement subscription must not + // silently retarget this read to a new table lifetime. + generation, generationOwned := + gs.currentOrCreateExactSubscribedUpdateRecord(key, subscribedEnt) + if !generationOwned { + return nil + } - func() { - // A forced update can block while the channel is full. A worker - // draining it may need gs.mu, so unlock before enqueueing. - gs.mu.Unlock() - defer gs.mu.Lock() - // If the trigger condition is not satisfied, the stats will not be updated - // for long time. So we trigger the update here to get the stats info as soon - // as possible. - gs.enqueueStatsUpdate(wrapkey, true) - }() - - info, ok = gs.mu.statsInfoMap[key] - if ok { - break - } - if ctx.Err() != nil { - return nil - } + // A forced enqueue is the ownership transfer to a producer. If admission + // is canceled, no waiter may depend on work that was never accepted. + if !gs.enqueueStatsUpdateForRecord(wrapkey, true, generation) { + return nil + } + return gs.waitForStatsUpdate(ctx, key, generation) +} - // Wait until stats info of the key is updated. - gs.mu.cond.Wait() +// waitForStatsUpdate waits on durable state, not on a Broadcast edge. The +// cache predicate and exact queued/running producer predicate are both checked +// while gs.mu is held; RemoveTid uses the same gs.mu -> updatingMu order. +// Cleanup or producer exhaustion therefore either changes the predicate before +// this check, or broadcasts after cond.Wait has atomically registered the +// waiter and released gs.mu. +func (gs *GlobalStats) waitForStatsUpdate( + ctx context.Context, + key pb.StatsInfoKey, + generation *updateRecord, +) *pb.StatsInfo { + stopWake := context.AfterFunc(ctx, gs.notifyStatsWaiters) + defer stopWake() - info, ok = gs.mu.statsInfoMap[key] + gs.mu.Lock() + defer gs.mu.Unlock() + for { + if info, complete := gs.mu.statsInfoMap[key]; complete { + return info + } + if ctx.Err() != nil { + return nil } + if gs.ctx != nil && gs.ctx.Err() != nil { + return nil + } + if !gs.statsUpdateProducerActive(key, generation) { + return nil + } + if gs.beforeStatsWait != nil { + gs.beforeStatsWait(key, generation) + } + gs.mu.cond.Wait() } - return info } func (gs *GlobalStats) enqueue(tail *logtail.TableLogtail) { @@ -541,43 +751,104 @@ func (gs *GlobalStats) spawnUpdateWorkers(ctx context.Context, num int) { case <-ctx.Done(): return - case key := <-gs.updateC: + case job := <-gs.updateC: // after dequeue from the chan, remove the table ID from the queue watcher. - gs.queueWatcher.del(key.Key.TableID) + gs.queueWatcher.del(job.wrapKey.Key.TableID) v2.StatsTriggerConsumeCounter.Add(1) - gs.coordinateStatsUpdate(key) + gs.coordinateStatsUpdateJob(job) } } }() } } -func (gs *GlobalStats) enqueueStatsUpdate(key pb.StatsInfoKeyWithContext, force bool) bool { +func (gs *GlobalStats) enqueueStatsUpdateForRecord( + key pb.StatsInfoKeyWithContext, + force bool, + expectedRecord *updateRecord, +) bool { + if expectedRecord == nil { + return false + } defer func() { v2.StatsTriggerQueueSizeGauge.Set(float64(len(gs.updateC))) }() + gs.registerStatsUpdateJob(expectedRecord) + job := statsUpdateJob{ + wrapKey: key, + expectedRecord: expectedRecord, + registered: true, + } if force { select { - case gs.updateC <- key: + case gs.updateC <- job: gs.queueWatcher.add(key.Key.TableID) v2.StatsTriggerForcedCounter.Add(1) return true case <-key.Ctx.Done(): + if gs.unregisterStatsUpdateJob(key.Key, expectedRecord) { + gs.notifyStatsWaiters() + } + return false + case <-gs.lifecycleDone(): + if gs.unregisterStatsUpdateJob(key.Key, expectedRecord) { + gs.notifyStatsWaiters() + } return false } } select { - case gs.updateC <- key: + case gs.updateC <- job: gs.queueWatcher.add(key.Key.TableID) v2.StatsTriggerUnforcedCounter.Add(1) return true default: + if gs.unregisterStatsUpdateJob(key.Key, expectedRecord) { + gs.notifyStatsWaiters() + } return false } } +func (gs *GlobalStats) registerStatsUpdateJob(generation *updateRecord) { + gs.updatingMu.Lock() + generation.queued++ + gs.updatingMu.Unlock() +} + +// unregisterStatsUpdateJob rolls back a job that never reached worker +// admission. It returns true only when the current generation has no remaining +// queued or running producer and waiters must re-evaluate their predicate. +func (gs *GlobalStats) unregisterStatsUpdateJob( + key pb.StatsInfoKey, + generation *updateRecord, +) bool { + gs.updatingMu.Lock() + defer gs.updatingMu.Unlock() + if generation.queued > 0 { + generation.queued-- + } + current, ok := gs.updatingMu.updating[key] + return ok && current == generation && generation.queued == 0 && !generation.inProgress +} + +func (gs *GlobalStats) notifyStatsWaiters() { + gs.mu.Lock() + if gs.mu.cond != nil { + gs.mu.cond.Broadcast() + } + gs.mu.Unlock() +} + +func (gs *GlobalStats) lifecycleDone() <-chan struct{} { + if gs.ctx == nil { + return nil + } + return gs.ctx.Done() +} + func (gs *GlobalStats) processLogtail(ctx context.Context, tail *logtail.TableLogtail) { key := pb.StatsInfoKey{ AccId: tail.Table.AccId, @@ -598,11 +869,14 @@ func (gs *GlobalStats) processLogtail(ctx context.Context, tail *logtail.TableLo } if len(tail.CkpLocation) > 0 || metaChanges > 0 { - if gs.keyExists(key) && gs.shouldEnqueueUpdate(key, metaChanges, len(tail.CkpLocation) > 0) { - gs.enqueueStatsUpdate(pb.StatsInfoKeyWithContext{ + record, ok := gs.shouldEnqueueExistingStatsUpdateGeneration( + key, metaChanges, len(tail.CkpLocation) > 0, + ) + if ok { + gs.enqueueStatsUpdateForRecord(pb.StatsInfoKeyWithContext{ Ctx: ctx, Key: key, - }, false) + }, false, record) } } } @@ -617,17 +891,52 @@ func (gs *GlobalStats) processLogtail(ctx context.Context, tail *logtail.TableLo // - Accumulated change rate >= 5%, OR // - Time since last update > 30min func (gs *GlobalStats) shouldEnqueueUpdate(key pb.StatsInfoKey, metaChanges int, hasCheckpoint bool) bool { + _, ok := gs.shouldEnqueueUpdateGeneration(key, metaChanges, hasCheckpoint) + return ok +} +func (gs *GlobalStats) shouldEnqueueUpdateGeneration( + key pb.StatsInfoKey, + metaChanges int, + hasCheckpoint bool, +) (*updateRecord, bool) { + gs.updatingMu.Lock() + defer gs.updatingMu.Unlock() + return gs.shouldEnqueueUpdateGenerationLocked(key, metaChanges, hasCheckpoint) +} + +// shouldEnqueueExistingStatsUpdateGeneration links the logtail producer to the +// same table-lifetime boundary as RemoveTid. The cache-existence check and +// scheduling-record capture are atomic in the established gs.mu -> updatingMu +// lock order, so cleanup cannot fall between them. +func (gs *GlobalStats) shouldEnqueueExistingStatsUpdateGeneration( + key pb.StatsInfoKey, + metaChanges int, + hasCheckpoint bool, +) (*updateRecord, bool) { + gs.mu.Lock() + defer gs.mu.Unlock() + if _, ok := gs.mu.statsInfoMap[key]; !ok { + return nil, false + } gs.updatingMu.Lock() defer gs.updatingMu.Unlock() + return gs.shouldEnqueueUpdateGenerationLocked(key, metaChanges, hasCheckpoint) +} +func (gs *GlobalStats) shouldEnqueueUpdateGenerationLocked( + key pb.StatsInfoKey, + metaChanges int, + hasCheckpoint bool, +) (*updateRecord, bool) { rec, ok := gs.updatingMu.updating[key] if !ok { // First time: create record and enqueue - gs.updatingMu.updating[key] = &updateRecord{ + rec = &updateRecord{ pendingChanges: metaChanges, } - return true + gs.updatingMu.updating[key] = rec + return rec, true } // Accumulate pending changes @@ -635,7 +944,7 @@ func (gs *GlobalStats) shouldEnqueueUpdate(key pb.StatsInfoKey, metaChanges int, // Small table: enqueue on any change if rec.baseObjectCount < LargeTableThreshold { - return metaChanges > 0 || hasCheckpoint + return rec, metaChanges > 0 || hasCheckpoint } // Large table: check two conditions (enqueue if either is true) @@ -643,52 +952,126 @@ func (gs *GlobalStats) shouldEnqueueUpdate(key pb.StatsInfoKey, metaChanges int, if rec.baseObjectCount > 0 { changeRate := float64(rec.pendingChanges) / float64(rec.baseObjectCount) if changeRate >= LargeTableChangeRateThreshold { - return true + return rec, true } } // Condition 2: Time since last update > 30min if time.Since(rec.lastUpdate) > LargeTableMaxUpdateInterval { - return true + return rec, true } - return false + return rec, false } -// shouldUpdate implements a debounce mechanism to prevent excessive stats updates. +func (gs *GlobalStats) startAutomaticUpdate( + key pb.StatsInfoKey, + expectedRecord *updateRecord, +) (*updateRecord, bool) { + if expectedRecord == nil { + return nil, false + } + gs.updatingMu.Lock() + defer gs.updatingMu.Unlock() + return gs.startAutomaticUpdateLocked(key, expectedRecord) +} -// shouldExecuteUpdate implements a debounce mechanism to prevent excessive stats updates. -// Only checks inProgress and MinUpdateInterval. -// Change rate is NOT checked here (already checked in shouldEnqueueUpdate). -func (gs *GlobalStats) shouldExecuteUpdate(key pb.StatsInfoKey) bool { +func (gs *GlobalStats) startAutomaticUpdateJob( + job statsUpdateJob, +) (*updateRecord, bool, bool) { + if job.expectedRecord == nil { + return nil, false, false + } gs.updatingMu.Lock() defer gs.updatingMu.Unlock() + if job.registered && job.expectedRecord.queued > 0 { + job.expectedRecord.queued-- + } + generation, started := + gs.startAutomaticUpdateLocked(job.wrapKey.Key, job.expectedRecord) + if started { + return generation, true, false + } + current, ok := gs.updatingMu.updating[job.wrapKey.Key] + noProducer := ok && current == job.expectedRecord && + job.expectedRecord.queued == 0 && !job.expectedRecord.inProgress + return nil, false, noProducer +} + +func (gs *GlobalStats) startAutomaticUpdateLocked( + key pb.StatsInfoKey, + expectedRecord *updateRecord, +) (*updateRecord, bool) { rec, ok := gs.updatingMu.updating[key] - if !ok { - gs.updatingMu.updating[key] = &updateRecord{ - inProgress: true, - } - return true + if !ok || rec != expectedRecord { + return nil, false } if rec.inProgress { - return false + return nil, false } if time.Since(rec.lastUpdate) > MinUpdateInterval { rec.inProgress = true - return true + return rec, true } - return false + return nil, false } -func (gs *GlobalStats) markUpdateComplete(key pb.StatsInfoKey, updated bool, actualObjectCount int64, samplingRatio float64) { +func (gs *GlobalStats) statsUpdateGenerationActive(key pb.StatsInfoKey, generation *updateRecord) bool { gs.updatingMu.Lock() defer gs.updatingMu.Unlock() rec, ok := gs.updatingMu.updating[key] - if !ok { - // set new record for RefreshWithMode - rec = &updateRecord{} - gs.updatingMu.updating[key] = rec + return ok && rec == generation +} + +func (gs *GlobalStats) statsUpdateProducerActive(key pb.StatsInfoKey, generation *updateRecord) bool { + gs.updatingMu.Lock() + defer gs.updatingMu.Unlock() + rec, ok := gs.updatingMu.updating[key] + return ok && rec == generation && (rec.queued > 0 || rec.inProgress) +} + +// markExplicitUpdateComplete advances the refresh baseline without stealing +// the in-progress bit from an automatic refresh that was admitted before the +// explicit refresh acquired the shared table stripe. +func (gs *GlobalStats) markExplicitUpdateComplete( + key pb.StatsInfoKey, + generation *updateRecord, + actualObjectCount int64, + samplingRatio float64, +) { + gs.updatingMu.Lock() + defer gs.updatingMu.Unlock() + rec, ok := gs.updatingMu.updating[key] + if !ok || rec != generation { + return } + rec.lastUpdate = time.Now() + rec.baseObjectCount = actualObjectCount + rec.pendingChanges = 0 + rec.samplingRatio = samplingRatio +} + +// markAutomaticUpdateComplete closes a generation opened by +// shouldExecuteUpdate. Unlike the explicit RefreshWithMode completion path, it +// must not recreate scheduling metadata that table-lifetime cleanup removed +// while an old worker was still unwinding. +func (gs *GlobalStats) markAutomaticUpdateComplete( + key pb.StatsInfoKey, + generation *updateRecord, + updated bool, + actualObjectCount int64, + samplingRatio float64, +) { + gs.updatingMu.Lock() + defer gs.updatingMu.Unlock() + rec, ok := gs.updatingMu.updating[key] + if !ok || rec != generation { + return + } + completeUpdateRecord(rec, updated, actualObjectCount, samplingRatio) +} + +func completeUpdateRecord(rec *updateRecord, updated bool, actualObjectCount int64, samplingRatio float64) { rec.inProgress = false // only if the stats is updated, set the update time and reset baseline. if updated { @@ -872,33 +1255,100 @@ func (gs *GlobalStats) broadcastStats(key pb.StatsInfoKey) { }) } -func (gs *GlobalStats) coordinateStatsUpdate(wrapKey pb.StatsInfoKeyWithContext) { +// completeAutomaticStatsCacheUpdate is the only automatic-refresh transition +// for statsInfoMap. A successful generation replaces the published value. A +// failed generation never destroys the last successful value; it installs a +// nil sentinel only when no generation has completed before, so synchronous +// first-read waiters can terminate without treating failure as publication. +func (gs *GlobalStats) completeAutomaticStatsCacheUpdate( + key pb.StatsInfoKey, + generation *updateRecord, + stats *pb.StatsInfo, + updated bool, +) bool { + gs.mu.Lock() + defer gs.mu.Unlock() + // GlobalStats shutdown is an authoritative failed-publication predicate. + // In particular, do not install a first-generation nil sentinel here: the + // lifecycle watcher already wakes waiters, and shutdown must leave the cache + // and its scheduling baseline unchanged. + if gs.ctx != nil && context.Cause(gs.ctx) != nil { + if gs.mu.cond != nil { + gs.mu.cond.Broadcast() + } + return false + } + // The update record is also the automatic generation's table-lifetime + // token. RemoveTid deletes it under the same gs.mu -> updatingMu order; an + // old worker that completes afterward must not resurrect either cache. + if !gs.statsUpdateGenerationActive(key, generation) { + if gs.mu.cond != nil { + gs.mu.cond.Broadcast() + } + return false + } + if updated { + gs.mu.statsInfoMap[key] = stats + gs.broadcastStats(key) + } else if _, ok := gs.mu.statsInfoMap[key]; !ok { + gs.mu.statsInfoMap[key] = nil + } + if gs.mu.cond != nil { + gs.mu.cond.Broadcast() + } + return updated +} + +func (gs *GlobalStats) coordinateStatsUpdateJob(job statsUpdateJob) { + wrapKey := job.wrapKey statser := statistic.StatsInfoFromContext(wrapKey.Ctx) crs := new(perfcounter.CounterSet) - if !gs.shouldExecuteUpdate(wrapKey.Key) { + generation, ok, noProducer := gs.startAutomaticUpdateJob(job) + if !ok { + if noProducer { + gs.notifyStatsWaiters() + } return } + if gs.afterAutomaticUpdateStarted != nil { + gs.afterAutomaticUpdateStarted(wrapKey.Key, generation) + } // updated is used to mark that the stats info is updated. var updated bool var actualObjectCount int64 var samplingRatio float64 + var stats *pb.StatsInfo + release, err := gs.acquireStatsRefresh(wrapKey.Ctx, wrapKey.Key) + if err != nil { + // Worker admission opened this generation before refresh admission. Close it + // even when cancellation prevents this worker from acquiring the stripe. + gs.markAutomaticUpdateComplete(wrapKey.Key, generation, false, 0, 0) + gs.notifyStatsWaiters() + return + } + if !gs.statsUpdateGenerationActive(wrapKey.Key, generation) { + // Table cleanup removed this queued generation while it waited for + // admission. Avoid re-subscribing and doing object work for stale state. + release() + return + } defer func() { - gs.markUpdateComplete(wrapKey.Key, updated, actualObjectCount, samplingRatio) + gs.completeAutomaticStatsRefresh( + wrapKey.Key, generation, stats, updated, + actualObjectCount, samplingRatio, release) }() - - broadcastWithoutUpdate := func() { - gs.mu.Lock() - defer gs.mu.Unlock() - gs.mu.statsInfoMap[wrapKey.Key] = nil - gs.mu.cond.Broadcast() + refreshCtx, stopRefresh, err := gs.newStatsRefreshContext(wrapKey.Ctx) + if err != nil { + return } + defer stopRefresh() // Get the latest partition state of the table. //Notice that for snapshot read, subscribing the table maybe failed since the invalid table id, //We should handle this case in next PR if needed. ps, err := gs.engine.pClient.toSubscribeTable( - wrapKey.Ctx, + refreshCtx, uint64(wrapKey.Key.AccId), wrapKey.Key.TableID, wrapKey.Key.TableName, @@ -910,12 +1360,11 @@ func (gs *GlobalStats) coordinateStatsUpdate(wrapKey pb.StatsInfoKeyWithContext) wrapKey.Key.TableID, wrapKey.Key.TableName, err) - broadcastWithoutUpdate() return } - stats := plan2.NewStatsInfo() + stats = plan2.NewStatsInfo() - newCtx := perfcounter.AttachS3RequestKey(wrapKey.Ctx, crs) + newCtx := perfcounter.AttachS3RequestKey(refreshCtx, crs) updated, samplingRatio = gs.executeStatsUpdate(newCtx, ps, wrapKey.Key, stats) // Get actual object count for baseline update @@ -932,37 +1381,82 @@ func (gs *GlobalStats) coordinateStatsUpdate(wrapKey pb.StatsInfoKeyWithContext) DeleteMul: crs.FileService.S3.DeleteMulti.Load(), }) - gs.mu.Lock() - defer gs.mu.Unlock() - if updated { - gs.mu.statsInfoMap[wrapKey.Key] = stats - gs.broadcastStats(wrapKey.Key) - } else if _, ok := gs.mu.statsInfoMap[wrapKey.Key]; !ok { - gs.mu.statsInfoMap[wrapKey.Key] = nil - } +} - // Notify all the waiters to read the new stats info. - gs.mu.cond.Broadcast() +// completeAutomaticStatsRefresh is the sole terminal owner for an admitted +// automatic refresh. Cache publication decides whether the result committed; +// only that committed result may advance scheduling metadata, and both happen +// before the table-scoped admission token is released. +func (gs *GlobalStats) completeAutomaticStatsRefresh( + key pb.StatsInfoKey, + generation *updateRecord, + stats *pb.StatsInfo, + calculated bool, + actualObjectCount int64, + samplingRatio float64, + release func(), +) { + committed := gs.completeAutomaticStatsCacheUpdate( + key, generation, stats, calculated) + gs.markAutomaticUpdateComplete( + key, generation, committed, actualObjectCount, samplingRatio) + release() } // RefreshWithMode triggers a stats refresh with the specified sampling mode func (gs *GlobalStats) RefreshWithMode(ctx context.Context, key pb.StatsInfoKey, samplingMode string) error { + _, err := gs.refreshStatsWithMode(ctx, key, samplingMode) + return err +} + +// refreshStatsWithMode returns the exact statistics object published while +// same-table refresh admission is still held. Callers that define a synchronous +// publication boundary must use this result instead of re-reading the map after +// admission has been released. +func (gs *GlobalStats) refreshStatsWithMode( + ctx context.Context, + key pb.StatsInfoKey, + samplingMode string, +) (*pb.StatsInfo, error) { + release, err := gs.acquireStatsRefresh(ctx, key) + if err != nil { + return nil, err + } + defer release() + refreshCtx, stopRefresh, err := gs.newStatsRefreshContext(ctx) + if err != nil { + return nil, err + } + defer stopRefresh() + // Get partition state ps, err := gs.engine.pClient.toSubscribeTable( - ctx, + refreshCtx, uint64(key.AccId), key.TableID, key.TableName, key.DatabaseID, key.DbName) if err != nil { - return moerr.NewInternalErrorNoCtxf("failed to subscribe table: %v", err) + if cause := gs.statsRefreshCancellationCause(ctx); cause != nil { + return nil, cause + } + return nil, moerr.NewInternalErrorNoCtxf("failed to subscribe table: %v", err) } // Get table definition table := gs.engine.GetLatestCatalogCache().GetTableById(key.AccId, key.DatabaseID, key.TableID) if table == nil || table.TableDef == nil { - return moerr.NewInternalErrorNoCtx("table not found") + return nil, moerr.NewInternalErrorNoCtx("table not found") + } + + // The subscription owns eventual RemoveTid cleanup. Capture the refresh + // generation only after subscription and catalog resolution succeed, and + // only while that exact subscription lifetime is still current. + generation, ok := gs.currentOrCreateSubscribedUpdateRecord(key) + if !ok { + return nil, moerr.NewInternalErrorNoCtxf( + "table statistics refresh crossed subscription boundary for table %d", key.TableID) } // Create stats info @@ -987,21 +1481,65 @@ func (gs *GlobalStats) RefreshWithMode(ctx context.Context, key pb.StatsInfoKey, } // Execute stats update - samplingRatio, err := CollectAndCalculateStats(ctx, req, gs.concurrentExecutor) + samplingRatio, err := CollectAndCalculateStats(refreshCtx, req, gs.concurrentExecutor) if err != nil { - return moerr.NewInternalErrorNoCtxf("failed to update stats: %v", err) + if cause := gs.statsRefreshCancellationCause(ctx); cause != nil { + return nil, cause + } + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + return nil, err + } + return nil, moerr.NewInternalErrorNoCtxf("failed to update stats: %v", err) + } + if cause := gs.statsRefreshCancellationCause(ctx); cause != nil { + return nil, cause + } + + if !gs.publishStatsForGeneration(key, generation, stats) { + if cause := gs.statsRefreshCancellationCause(ctx); cause != nil { + return nil, cause + } + return nil, moerr.NewInternalErrorNoCtxf( + "table statistics refresh crossed cleanup boundary for table %d", key.TableID) } + // Record the baseline only if this exact table lifetime is still current. + // Preserve a concurrently admitted automatic refresh's in-progress bit. + gs.markExplicitUpdateComplete( + key, generation, stats.AccurateObjectNumber, samplingRatio) - // Update cache + return stats, nil +} + +// publishStatsForGeneration replaces the cache only while the exact table +// lifetime captured by the refresh remains current. The gs.mu -> updatingMu +// order matches RemoveTid, making validation and publication atomic with +// respect to cleanup. A late explicit refresh therefore cannot resurrect an +// unsubscribed table or publish into a replacement generation. +func (gs *GlobalStats) publishStatsForGeneration( + key pb.StatsInfoKey, + generation *updateRecord, + stats *pb.StatsInfo, +) bool { + if generation == nil || stats == nil { + return false + } gs.mu.Lock() defer gs.mu.Unlock() + if gs.ctx != nil && context.Cause(gs.ctx) != nil { + return false + } + if !gs.statsUpdateGenerationActive(key, generation) { + if gs.mu.cond != nil { + gs.mu.cond.Broadcast() + } + return false + } gs.mu.statsInfoMap[key] = stats gs.broadcastStats(key) - gs.mu.cond.Broadcast() - // Record sampling ratio in updateRecord - gs.markUpdateComplete(key, true, stats.AccurateObjectNumber, samplingRatio) - - return nil + if gs.mu.cond != nil { + gs.mu.cond.Broadcast() + } + return true } func (gs *GlobalStats) executeStatsUpdate(ctx context.Context, ps *logtailreplay.PartitionState, key pb.StatsInfoKey, stats *pb.StatsInfo) (bool, float64) { @@ -1040,6 +1578,9 @@ func (gs *GlobalStats) executeStatsUpdate(ctx context.Context, ps *logtailreplay logutil.Errorf("failed to init stats info for table %v, err: %v", key, err) return false, 0 } + if context.Cause(ctx) != nil { + return false, 0 + } v2.StatsUpdateDurationHistogram.Observe(time.Since(start).Seconds()) v2.StatsUpdateBlockCounter.Add(float64(stats.BlockNumber)) return true, samplingRatio @@ -1174,7 +1715,7 @@ func collectTableStats( var sampledRowCount float64 var sampledObjectCount int64 - onObjFn := func(obj objectio.ObjectEntry) error { + onObjFn := func(objCtx context.Context, obj objectio.ObjectEntry) error { objName := obj.ObjectShortName() // ===== Phase 1: Get exact values from ObjectStats (no IO) ===== @@ -1205,7 +1746,7 @@ func collectTableStats( // Sampled object: read ObjectMeta (requires IO) location := obj.Location() - objMeta, err := objectio.FastLoadObjectMeta(ctx, &location, false, fs) + objMeta, err := objectio.FastLoadObjectMeta(objCtx, &location, false, fs) if err != nil { return err } @@ -1344,6 +1885,7 @@ func collectTableStats( } if err := ForeachVisibleObjects( + ctx, req.partitionState, req.ts, onObjFn, @@ -1398,6 +1940,16 @@ func CollectAndCalculateStats(ctx context.Context, req *updateStatsRequest, exec defer func() { v2.TxnStatementUpdateStatsDurationHistogram.Observe(time.Since(start).Seconds()) }() + // The zero-object fast path below does not enter ForeachVisibleObjects. Keep + // executor shutdown as a failed refresh by checking its authoritative + // lifecycle before any successful early return. + if executor != nil { + if lifecycle := executor.LifecycleContext(); lifecycle != nil { + if cause := context.Cause(lifecycle); cause != nil { + return 0, cause + } + } + } lenCols := len(req.tableDef.Cols) - 1 /* row-id */ info := plan2.NewTableStatsInfo(lenCols) if req.approxObjectNum == 0 { diff --git a/pkg/vm/engine/disttae/stats_test.go b/pkg/vm/engine/disttae/stats_test.go index 75b118d14fdf4..daed5f4c76fb7 100644 --- a/pkg/vm/engine/disttae/stats_test.go +++ b/pkg/vm/engine/disttae/stats_test.go @@ -231,11 +231,16 @@ func TestGlobalStats_ShouldUpdate(t *testing.T) { DatabaseID: 100, TableID: 101, } - assert.True(t, gs.shouldExecuteUpdate(k1)) - assert.False(t, gs.shouldExecuteUpdate(k1)) - gs.markUpdateComplete(k1, true, 1, 1.0) + generation := gs.currentOrCreateUpdateRecord(k1) + _, started := gs.startAutomaticUpdate(k1, generation) + assert.True(t, started) + _, started = gs.startAutomaticUpdate(k1, generation) + assert.False(t, started) + gs.markAutomaticUpdateComplete( + k1, generation, true, 1, 1.0) time.Sleep(MinUpdateInterval) - assert.True(t, gs.shouldExecuteUpdate(k1)) + _, started = gs.startAutomaticUpdate(k1, generation) + assert.True(t, started) }) t.Run("parallel", func(t *testing.T) { @@ -256,11 +261,13 @@ func TestGlobalStats_ShouldUpdate(t *testing.T) { var wg sync.WaitGroup updateFn := func() { defer wg.Done() - if !gs.shouldExecuteUpdate(k1) { + generation := gs.currentOrCreateUpdateRecord(k1) + if _, started := gs.startAutomaticUpdate(k1, generation); !started { return } count.Add(1) - gs.markUpdateComplete(k1, true, 2, 1.0) + gs.markAutomaticUpdateComplete( + k1, generation, true, 2, 1.0) } for i := 0; i < 20; i++ { wg.Add(1) @@ -1267,7 +1274,8 @@ func TestSamplingForceAtLeastOneObject(t *testing.T) { // by calling shouldEnqueueUpdate once and then markUpdateComplete to set the baseObjectCount func initTableForTest(gs *GlobalStats, key statsinfo.StatsInfoKey, baseObjectCount int64) { gs.shouldEnqueueUpdate(key, 0, false) - gs.markUpdateComplete(key, true, baseObjectCount, 1.0) + gs.markExplicitUpdateComplete( + key, gs.currentOrCreateUpdateRecord(key), baseObjectCount, 1.0) } // TestGlobalStats_ShouldEnqueue tests the shouldEnqueue logic for large table throttling @@ -1732,21 +1740,135 @@ func TestRemoveTid(t *testing.T) { gs.mu.statsInfoMap[k2] = nil // simulate failed update gs.mu.statsInfoMap[k3] = plan2.NewStatsInfo() gs.mu.Unlock() + generation := gs.currentOrCreateUpdateRecord(k1) + gs.currentOrCreateUpdateRecord(k2) + gs.currentOrCreateUpdateRecord(k3) + gs.markExplicitUpdateComplete(k1, generation, 1, 1) + gs.markAutomaticUpdateComplete( + k2, gs.currentOrCreateUpdateRecord(k2), false, 0, 0) + gs.markExplicitUpdateComplete( + k3, gs.currentOrCreateUpdateRecord(k3), 2, 1) // Remove table 1001 entries gs.RemoveTid(1001) + // A worker admitted before cleanup may finish afterward. Its stale + // publication or completion must not recreate table-owned state. + queuedAfterCleanup, enqueueAfterCleanup := + gs.shouldEnqueueExistingStatsUpdateGeneration(k1, 1, false) + assert.False(t, enqueueAfterCleanup) + assert.Nil(t, queuedAfterCleanup) + gs.completeAutomaticStatsCacheUpdate(k1, generation, plan2.NewStatsInfo(), true) + gs.completeAutomaticStatsRefresh( + k1, generation, plan2.NewStatsInfo(), true, 3, 1, func() {}) gs.mu.Lock() - defer gs.mu.Unlock() _, ok1 := gs.mu.statsInfoMap[k1] _, ok2 := gs.mu.statsInfoMap[k2] _, ok3 := gs.mu.statsInfoMap[k3] + gs.mu.Unlock() assert.False(t, ok1, "k1 should be removed") assert.False(t, ok2, "k2 should be removed") assert.True(t, ok3, "k3 should not be removed") + + gs.updatingMu.Lock() + _, updating1 := gs.updatingMu.updating[k1] + _, updating2 := gs.updatingMu.updating[k2] + _, updating3 := gs.updatingMu.updating[k3] + gs.updatingMu.Unlock() + assert.False(t, updating1, "k1 scheduling metadata should be removed") + assert.False(t, updating2, "k2 scheduling metadata should be removed") + assert.True(t, updating3, "unrelated scheduling metadata should remain") + + // Reuse of the same table key creates a distinct generation. Neither + // an old queued job nor its late callbacks may publish into it. + replacement := &updateRecord{inProgress: true, pendingChanges: 7} + gs.updatingMu.Lock() + gs.updatingMu.updating[k1] = replacement + gs.updatingMu.Unlock() + gs.completeAutomaticStatsCacheUpdate(k1, generation, plan2.NewStatsInfo(), true) + gs.completeAutomaticStatsRefresh( + k1, generation, plan2.NewStatsInfo(), true, 4, 0.5, func() {}) + _, oldGenerationStarted := gs.startAutomaticUpdate(k1, generation) + assert.False(t, oldGenerationStarted, "an old queued generation should be rejected") + + gs.mu.Lock() + _, oldStatsPublished := gs.mu.statsInfoMap[k1] + gs.mu.Unlock() + assert.False(t, oldStatsPublished, "an old generation should not publish into its replacement") + gs.updatingMu.Lock() + current := gs.updatingMu.updating[k1] + gs.updatingMu.Unlock() + require.Same(t, replacement, current) + assert.True(t, current.inProgress) + assert.Equal(t, 7, current.pendingChanges) }) }) + t.Run("first_queued_and_explicit_refreshes_cannot_cross_cleanup_generation", func(t *testing.T) { + gs := &GlobalStats{ + updateC: make(chan statsUpdateJob, 1), + queueWatcher: newQueueWatcher(), + } + gs.updatingMu.updating = make(map[statsinfo.StatsInfoKey]*updateRecord) + gs.mu.statsInfoMap = make(map[statsinfo.StatsInfoKey]*statsinfo.StatsInfo) + gs.mu.cond = sync.NewCond(&gs.mu) + + key := statsinfo.StatsInfoKey{DatabaseID: 100, TableID: 1001, TableName: "t1"} + generation := gs.currentOrCreateUpdateRecord(key) + require.True(t, gs.enqueueStatsUpdateForRecord(statsinfo.StatsInfoKeyWithContext{ + Ctx: context.Background(), + Key: key, + }, false, generation)) + job := <-gs.updateC + require.NotNil(t, job.expectedRecord, + "the first queued refresh must own a concrete lifetime token") + + gs.RemoveTid(key.TableID) + gs.coordinateStatsUpdateJob(job) + gs.markAutomaticUpdateComplete(key, job.expectedRecord, true, 1, 1) + assert.False(t, gs.publishStatsForGeneration( + key, job.expectedRecord, plan2.NewStatsInfo())) + + gs.mu.Lock() + _, cached := gs.mu.statsInfoMap[key] + gs.mu.Unlock() + gs.updatingMu.Lock() + _, scheduled := gs.updatingMu.updating[key] + gs.updatingMu.Unlock() + assert.False(t, cached, "old work must not recreate the statistics cache") + assert.False(t, scheduled, "old work must not recreate scheduling metadata") + + replacement := gs.currentOrCreateUpdateRecord(key) + assert.False(t, gs.publishStatsForGeneration( + key, job.expectedRecord, plan2.NewStatsInfo()), + "old explicit work must not publish into a replacement lifetime") + fresh := plan2.NewStatsInfo() + fresh.TableCnt = 42 + require.True(t, gs.publishStatsForGeneration(key, replacement, fresh)) + gs.mu.Lock() + assert.Same(t, fresh, gs.mu.statsInfoMap[key]) + gs.mu.Unlock() + }) + + t.Run("explicit_completion_preserves_admitted_automatic_refresh", func(t *testing.T) { + gs := &GlobalStats{} + gs.updatingMu.updating = make(map[statsinfo.StatsInfoKey]*updateRecord) + key := statsinfo.StatsInfoKey{DatabaseID: 100, TableID: 1001, TableName: "t1"} + generation := &updateRecord{inProgress: true, pendingChanges: 7} + gs.updatingMu.updating[key] = generation + + gs.markExplicitUpdateComplete(key, generation, 42, 0.5) + + gs.updatingMu.Lock() + got := *gs.updatingMu.updating[key] + gs.updatingMu.Unlock() + assert.True(t, got.inProgress, + "explicit completion must not reopen admission for another automatic refresh") + assert.Equal(t, int64(42), got.baseObjectCount) + assert.Zero(t, got.pendingChanges) + assert.Equal(t, 0.5, got.samplingRatio) + }) + t.Run("remove_nonexistent_table", func(t *testing.T) { runTest(t, func(ctx context.Context, e *Engine) { gs := e.globalStats @@ -1766,46 +1888,42 @@ func TestRemoveTid(t *testing.T) { }) }) - t.Run("remove_wakes_waiting_goroutines", func(t *testing.T) { - runTest(t, func(ctx context.Context, e *Engine) { - gs := e.globalStats - - // Set up: goroutine will cond.Wait() on a key, RemoveTid should broadcast - targetKey := statsinfo.StatsInfoKey{TableID: 42, DatabaseID: 1} - - woken := make(chan bool, 1) - gs.mu.Lock() - go func() { - gs.mu.Lock() - defer gs.mu.Unlock() - // Block on cond.Wait() like production GlobalStats.Get does - for { - if _, ok := gs.mu.statsInfoMap[targetKey]; ok { - break - } - // cond.Wait releases the lock and waits for Broadcast - gs.mu.cond.Wait() - // After Broadcast, check if our condition changed - break - } - woken <- true - }() - gs.mu.Unlock() - - // Small sleep to let goroutine enter cond.Wait() - time.Sleep(50 * time.Millisecond) - - // RemoveTid broadcasts to cond, which should wake the waiting goroutine - gs.RemoveTid(999) +} - select { - case <-woken: - // ok — goroutine was woken by Broadcast - case <-time.After(2 * time.Second): - t.Fatal("RemoveTid did not wake goroutine blocked on cond.Wait()") - } - }) - }) +func TestStatsPublicationRejectsStoppedOwnerLifecycle(t *testing.T) { + ownerCtx, stopOwner := context.WithCancel(context.Background()) + stopOwner() + key := statsinfo.StatsInfoKey{AccId: 1, DatabaseID: 10, TableID: 42} + lastGood := plan2.NewStatsInfo() + lastGood.TableCnt = 7 + generation := &updateRecord{ + inProgress: true, + baseObjectCount: 7, + samplingRatio: 0.25, + } + gs := &GlobalStats{ctx: ownerCtx} + gs.mu.statsInfoMap = map[statsinfo.StatsInfoKey]*statsinfo.StatsInfo{key: lastGood} + gs.mu.cond = sync.NewCond(&gs.mu) + gs.updatingMu.updating = map[statsinfo.StatsInfoKey]*updateRecord{key: generation} + + fresh := plan2.NewStatsInfo() + fresh.TableCnt = 42 + releases := 0 + gs.completeAutomaticStatsRefresh( + key, generation, fresh, true, 42, 1, func() { releases++ }) + require.False(t, gs.publishStatsForGeneration(key, generation, fresh)) + require.Equal(t, 1, releases) + + gs.mu.Lock() + require.Same(t, lastGood, gs.mu.statsInfoMap[key], + "shutdown must preserve the last successfully published statistics") + gs.mu.Unlock() + gs.updatingMu.Lock() + require.False(t, generation.inProgress) + require.Equal(t, int64(7), generation.baseObjectCount, + "failed publication must not advance the object-count baseline") + require.Equal(t, 0.25, generation.samplingRatio) + gs.updatingMu.Unlock() } func TestGlobalStatsGetDoesNotHoldMuWhileSubscribing(t *testing.T) { @@ -1816,6 +1934,17 @@ func TestGlobalStatsGetDoesNotHoldMuWhileSubscribing(t *testing.T) { e.pClient.eng = e e.pClient.subscribed.eng = e + partition := e.GetOrCreateLatestPart(ctx, 0, dbID, tblID) + state, commit := partition.MutateState() + objectID := objectio.NewObjectid() + objectStats := objectio.NewObjectStatsWithObjectID( + &objectID, false, false, false) + require.NoError(t, objectio.SetObjectStatsSize(objectStats, 1)) + require.NoError(t, state.HandleObjectEntry(ctx, nil, objectio.ObjectEntry{ + ObjectStats: *objectStats, + CreateTime: types.BuildTS(1, 0), + }, false)) + commit() ent := &subEntry{dbID: dbID, state: Subscribed} ent.lastTs.Store(time.Now().UnixNano()) @@ -1854,10 +1983,11 @@ func TestGlobalStatsGetDoesNotHoldMuWhileSubscribing(t *testing.T) { getCtx, cancel := context.WithTimeout(ctx, 2*time.Second) defer cancel() - getDone := make(chan struct{}) + getDone := make(chan *statsinfo.StatsInfo, 1) + getCompleted := make(chan struct{}) go func() { - defer close(getDone) - _ = gs.Get(getCtx, key, false) + getDone <- gs.Get(getCtx, key, false) + close(getCompleted) }() require.Eventually(t, func() bool { @@ -1869,16 +1999,19 @@ func TestGlobalStatsGetDoesNotHoldMuWhileSubscribing(t *testing.T) { } }, time.Second, 10*time.Millisecond, "GlobalStats.Get did not reach subscribe path") + published := plan2.NewStatsInfo() + published.TableCnt = 42 muAcquired := make(chan struct{}) go func() { gs.mu.Lock() + gs.mu.statsInfoMap[key] = published gs.mu.Unlock() close(muAcquired) }() require.Eventually(t, func() bool { select { - case <-getDone: + case <-getCompleted: return false default: } @@ -1894,13 +2027,58 @@ func TestGlobalStatsGetDoesNotHoldMuWhileSubscribing(t *testing.T) { e.pClient.subscribed.rw.Unlock() select { - case <-getDone: + case result := <-getDone: + require.Same(t, published, result, + "a non-blocking Get must recheck publication after subscription") case <-time.After(time.Second): t.Fatal("GlobalStats.Get did not return after subscribe lock released") } }) } +func newSynchronousStatsGetHarness( + t *testing.T, + ctx context.Context, + e *Engine, + key statsinfo.StatsInfoKey, +) (*GlobalStats, *subEntry) { + t.Helper() + partition := e.GetOrCreateLatestPart(ctx, uint64(key.AccId), key.DatabaseID, key.TableID) + state, commit := partition.MutateState() + objectID := objectio.NewObjectid() + objectStats := objectio.NewObjectStatsWithObjectID( + &objectID, false, false, false) + require.NoError(t, objectio.SetObjectStatsSize(objectStats, 1)) + require.NoError(t, state.HandleObjectEntry(ctx, nil, objectio.ObjectEntry{ + ObjectStats: *objectStats, + CreateTime: types.BuildTS(1, 0), + }, false)) + commit() + + e.pClient.eng = e + e.pClient.subscribed.eng = e + ent := &subEntry{dbID: key.DatabaseID, state: Subscribed} + ent.lastTs.Store(time.Now().UnixNano()) + e.pClient.subscribed.rw.Lock() + if e.pClient.subscribed.m == nil { + e.pClient.subscribed.m = make(map[uint64]*subEntry) + } + e.pClient.subscribed.m[key.TableID] = ent + e.pClient.subscribed.rw.Unlock() + + gs := &GlobalStats{ + ctx: ctx, + engine: e, + updateC: make(chan statsUpdateJob, 1), + queueWatcher: newQueueWatcher(), + } + gs.updatingMu.updating = make(map[statsinfo.StatsInfoKey]*updateRecord) + gs.mu.statsInfoMap = make(map[statsinfo.StatsInfoKey]*statsinfo.StatsInfo) + gs.mu.cond = sync.NewCond(&gs.mu) + gs.initStatsRefreshAdmission() + return gs, ent +} + func TestGlobalStatsGetReturnsWhenContextCanceledWhileWaiting(t *testing.T) { runTest(t, func(ctx context.Context, e *Engine) { const dbID uint64 = 100 @@ -1913,41 +2091,15 @@ func TestGlobalStatsGetReturnsWhenContextCanceledWhileWaiting(t *testing.T) { DbName: "d", } - partition := e.GetOrCreateLatestPart(ctx, 0, dbID, tblID) - state, commit := partition.MutateState() - objectID := objectio.NewObjectid() - objectStats := objectio.NewObjectStatsWithObjectID( - &objectID, false, false, false) - require.NoError(t, objectio.SetObjectStatsSize(objectStats, 1)) - require.NoError(t, state.HandleObjectEntry(ctx, nil, objectio.ObjectEntry{ - ObjectStats: *objectStats, - CreateTime: types.BuildTS(1, 0), - }, false)) - commit() - - e.pClient.eng = e - e.pClient.subscribed.eng = e - e.pClient.subscribed.rw.Lock() - if e.pClient.subscribed.m == nil { - e.pClient.subscribed.m = make(map[uint64]*subEntry) - } - e.pClient.subscribed.m[tblID] = &subEntry{ - dbID: dbID, - state: Subscribed, - } - e.pClient.subscribed.rw.Unlock() - // This isolated GlobalStats has no update worker. Receiving its forced // request below proves Get reached the synchronous update path, after // which no data-path goroutine can broadcast the condition variable. - gs := &GlobalStats{ - ctx: ctx, - engine: e, - updateC: make(chan statsinfo.StatsInfoKeyWithContext, 1), - queueWatcher: newQueueWatcher(), + gs, _ := newSynchronousStatsGetHarness(t, ctx, e, key) + waitEntered := make(chan struct{}) + var waitOnce sync.Once + gs.beforeStatsWait = func(statsinfo.StatsInfoKey, *updateRecord) { + waitOnce.Do(func() { close(waitEntered) }) } - gs.mu.statsInfoMap = make(map[statsinfo.StatsInfoKey]*statsinfo.StatsInfo) - gs.mu.cond = sync.NewCond(&gs.mu) getCtx, cancel := context.WithCancel(ctx) defer cancel() @@ -1956,11 +2108,17 @@ func TestGlobalStatsGetReturnsWhenContextCanceledWhileWaiting(t *testing.T) { result <- gs.Get(getCtx, key, true) }() + var job statsUpdateJob select { - case <-gs.updateC: + case job = <-gs.updateC: case <-time.After(time.Second): t.Fatal("GlobalStats.Get did not enqueue the synchronous update") } + select { + case <-waitEntered: + case <-time.After(time.Second): + t.Fatal("GlobalStats.Get did not register its condition wait") + } cancel() select { @@ -1969,29 +2127,350 @@ func TestGlobalStatsGetReturnsWhenContextCanceledWhileWaiting(t *testing.T) { case <-time.After(time.Second): t.Fatal("GlobalStats.Get did not return after context cancellation") } + gs.unregisterStatsUpdateJob(key, job.expectedRecord) gs.queueWatcher.del(tblID) }) } +func TestGlobalStatsGetReturnsPublishedStatsFromAcceptedProducer(t *testing.T) { + runTest(t, func(ctx context.Context, e *Engine) { + key := statsinfo.StatsInfoKey{ + DatabaseID: 100, + TableID: 10006, + TableName: "t", + DbName: "d", + } + gs, _ := newSynchronousStatsGetHarness(t, ctx, e, key) + result := make(chan *statsinfo.StatsInfo, 1) + go func() { result <- gs.Get(ctx, key, true) }() + + var job statsUpdateJob + select { + case job = <-gs.updateC: + case <-time.After(time.Second): + t.Fatal("synchronous read did not enqueue its producer") + } + generation, started, noProducer := gs.startAutomaticUpdateJob(job) + require.True(t, started) + require.False(t, noProducer) + published := plan2.NewStatsInfo() + published.TableCnt = 42 + gs.completeAutomaticStatsCacheUpdate(key, generation, published, true) + gs.markAutomaticUpdateComplete(key, generation, true, 1, 1) + + select { + case info := <-result: + require.Same(t, published, info) + case <-time.After(time.Second): + t.Fatal("synchronous read did not observe accepted producer publication") + } + gs.queueWatcher.del(key.TableID) + }) +} + +func TestGlobalStatsGetReturnsWhenCleanupPrecedesWait(t *testing.T) { + runTest(t, func(ctx context.Context, e *Engine) { + key := statsinfo.StatsInfoKey{ + DatabaseID: 100, + TableID: 10002, + TableName: "t", + DbName: "d", + } + gs, _ := newSynchronousStatsGetHarness(t, ctx, e, key) + + // Hold the watcher after updateC accepts the job. This is an observable + // barrier between producer ownership transfer and wait registration. + gs.queueWatcher.Lock() + watcherLocked := true + defer func() { + if watcherLocked { + gs.queueWatcher.Unlock() + } + }() + + waitEntered := make(chan struct{}) + var waitOnce sync.Once + gs.beforeStatsWait = func(statsinfo.StatsInfoKey, *updateRecord) { + waitOnce.Do(func() { close(waitEntered) }) + } + result := make(chan *statsinfo.StatsInfo, 1) + go func() { result <- gs.Get(ctx, key, true) }() + + var staleJob statsUpdateJob + select { + case staleJob = <-gs.updateC: + case <-time.After(time.Second): + t.Fatal("GlobalStats.Get did not transfer the refresh job") + } + + // Cleanup broadcasts before Get is able to enter cond.Wait. Processing + // the stale job produces no second notification. + gs.RemoveTid(key.TableID) + gs.coordinateStatsUpdateJob(staleJob) + watcherLocked = false + gs.queueWatcher.Unlock() + + select { + case info := <-result: + require.Nil(t, info) + case <-time.After(time.Second): + t.Fatal("GlobalStats.Get lost the cleanup wake before cond.Wait") + } + select { + case <-waitEntered: + t.Fatal("GlobalStats.Get waited after its producer generation was removed") + default: + } + gs.queueWatcher.del(key.TableID) + }) +} + +func TestGlobalStatsGetReturnsWhenCleanupFollowsWait(t *testing.T) { + runTest(t, func(ctx context.Context, e *Engine) { + key := statsinfo.StatsInfoKey{ + DatabaseID: 100, + TableID: 10003, + TableName: "t", + DbName: "d", + } + gs, _ := newSynchronousStatsGetHarness(t, ctx, e, key) + waitEntered := make(chan struct{}) + var waitOnce sync.Once + gs.beforeStatsWait = func(statsinfo.StatsInfoKey, *updateRecord) { + waitOnce.Do(func() { close(waitEntered) }) + } + result := make(chan *statsinfo.StatsInfo, 1) + go func() { result <- gs.Get(ctx, key, true) }() + + var staleJob statsUpdateJob + select { + case staleJob = <-gs.updateC: + case <-time.After(time.Second): + t.Fatal("GlobalStats.Get did not transfer the refresh job") + } + select { + case <-waitEntered: + case <-time.After(time.Second): + t.Fatal("GlobalStats.Get did not reach cond.Wait") + } + + gs.RemoveTid(key.TableID) + select { + case info := <-result: + require.Nil(t, info) + case <-time.After(time.Second): + t.Fatal("GlobalStats.Get was not released by cleanup after cond.Wait") + } + gs.coordinateStatsUpdateJob(staleJob) + gs.queueWatcher.del(key.TableID) + }) +} + +func TestGlobalStatsGetDoesNotOutliveCanceledSharedProducer(t *testing.T) { + runTest(t, func(ctx context.Context, e *Engine) { + key := statsinfo.StatsInfoKey{ + DatabaseID: 100, + TableID: 10004, + TableName: "t", + DbName: "d", + } + gs, _ := newSynchronousStatsGetHarness(t, ctx, e, key) + + // Occupy the table stripe so the first producer can be canceled after + // worker admission but before it starts object work. + releaseStripe, err := gs.acquireStatsRefresh(context.Background(), key) + require.NoError(t, err) + stripeHeld := true + defer func() { + if stripeHeld { + releaseStripe() + } + }() + + producerStarted := make(chan struct{}) + var producerOnce sync.Once + gs.afterAutomaticUpdateStarted = func(statsinfo.StatsInfoKey, *updateRecord) { + producerOnce.Do(func() { close(producerStarted) }) + } + + producerCtx, cancelProducer := context.WithCancel(ctx) + defer cancelProducer() + producerResult := make(chan *statsinfo.StatsInfo, 1) + go func() { producerResult <- gs.Get(producerCtx, key, true) }() + var producerJob statsUpdateJob + select { + case producerJob = <-gs.updateC: + case <-time.After(time.Second): + t.Fatal("first synchronous read did not enqueue its producer") + } + producerDone := make(chan struct{}) + go func() { + gs.coordinateStatsUpdateJob(producerJob) + close(producerDone) + }() + select { + case <-producerStarted: + case <-time.After(time.Second): + t.Fatal("first producer did not reach worker admission") + } + + sharedResult := make(chan *statsinfo.StatsInfo, 1) + go func() { sharedResult <- gs.Get(ctx, key, true) }() + var sharedJob statsUpdateJob + select { + case sharedJob = <-gs.updateC: + case <-time.After(time.Second): + t.Fatal("second synchronous read did not enqueue its shared job") + } + // This job coalesces behind the in-progress producer. Its waiter must + // still be released if that producer is canceled before publication. + gs.coordinateStatsUpdateJob(sharedJob) + cancelProducer() + + select { + case <-producerDone: + case <-time.After(time.Second): + t.Fatal("canceled producer did not leave refresh admission") + } + select { + case info := <-producerResult: + require.Nil(t, info) + case <-time.After(time.Second): + t.Fatal("producer caller did not observe cancellation") + } + select { + case info := <-sharedResult: + require.Nil(t, info) + case <-time.After(time.Second): + t.Fatal("shared waiter outlived its canceled producer") + } + + stripeHeld = false + releaseStripe() + gs.queueWatcher.del(key.TableID) + }) +} + +func TestGlobalStatsGetReturnsWhenStatsWorkersStop(t *testing.T) { + runTest(t, func(ctx context.Context, e *Engine) { + key := statsinfo.StatsInfoKey{ + DatabaseID: 100, + TableID: 10005, + TableName: "t", + DbName: "d", + } + gs, _ := newSynchronousStatsGetHarness(t, ctx, e, key) + workerCtx, stopWorkers := context.WithCancel(ctx) + gs.ctx = workerCtx + context.AfterFunc(workerCtx, gs.notifyStatsWaiters) + waitEntered := make(chan struct{}) + var waitOnce sync.Once + gs.beforeStatsWait = func(statsinfo.StatsInfoKey, *updateRecord) { + waitOnce.Do(func() { close(waitEntered) }) + } + + result := make(chan *statsinfo.StatsInfo, 1) + go func() { result <- gs.Get(ctx, key, true) }() + var abandonedJob statsUpdateJob + select { + case abandonedJob = <-gs.updateC: + case <-time.After(time.Second): + t.Fatal("synchronous read did not enqueue before worker shutdown") + } + select { + case <-waitEntered: + case <-time.After(time.Second): + t.Fatal("synchronous read did not wait for its queued producer") + } + + stopWorkers() + select { + case info := <-result: + require.Nil(t, info) + case <-time.After(time.Second): + t.Fatal("synchronous read outlived the statistics worker lifecycle") + } + gs.unregisterStatsUpdateJob(key, abandonedJob.expectedRecord) + gs.queueWatcher.del(key.TableID) + }) +} + +func TestStatsUpdateGenerationRequiresLiveSubscriptionOwner(t *testing.T) { + key := statsinfo.StatsInfoKey{DatabaseID: 10, TableID: 42} + e := &Engine{} + gs := &GlobalStats{ + engine: e, + updateC: make(chan statsUpdateJob, 1), + queueWatcher: newQueueWatcher(), + } + gs.updatingMu.updating = make(map[statsinfo.StatsInfoKey]*updateRecord) + gs.mu.statsInfoMap = make(map[statsinfo.StatsInfoKey]*statsinfo.StatsInfo) + gs.mu.cond = sync.NewCond(&gs.mu) + e.pClient.subscribed.rw.Lock() + e.pClient.subscribed.m = make(map[uint64]*subEntry) + e.pClient.subscribed.rw.Unlock() + + require.False(t, gs.PrefetchTableMeta(context.Background(), key)) + gs.updatingMu.Lock() + _, retained := gs.updatingMu.updating[key] + gs.updatingMu.Unlock() + require.False(t, retained, + "prefetch without a subscription cleanup owner must not create a generation") + + oldEnt := &subEntry{dbID: key.DatabaseID, state: Subscribed} + e.pClient.subscribed.rw.Lock() + e.pClient.subscribed.m[key.TableID] = oldEnt + e.pClient.subscribed.rw.Unlock() + oldGeneration, ok := gs.currentOrCreateExactSubscribedUpdateRecord(key, oldEnt) + require.True(t, ok) + require.True(t, gs.PrefetchTableMeta(context.Background(), key)) + job := <-gs.updateC + require.Same(t, oldGeneration, job.expectedRecord) + gs.queueWatcher.del(key.TableID) + + gs.RemoveTid(key.TableID) + gs.coordinateStatsUpdateJob(job) + gs.updatingMu.Lock() + require.Zero(t, oldGeneration.queued) + gs.updatingMu.Unlock() + newEnt := &subEntry{dbID: key.DatabaseID, state: Subscribed} + e.pClient.subscribed.rw.Lock() + e.pClient.subscribed.m[key.TableID] = newEnt + e.pClient.subscribed.rw.Unlock() + _, ok = gs.currentOrCreateExactSubscribedUpdateRecord(key, oldEnt) + require.False(t, ok, + "work captured from an old subscription must not target its replacement") + gs.updatingMu.Lock() + _, retained = gs.updatingMu.updating[key] + gs.updatingMu.Unlock() + require.False(t, retained, + "rejecting an old subscription must not create idle replacement metadata") + newGeneration, ok := gs.currentOrCreateExactSubscribedUpdateRecord(key, newEnt) + require.True(t, ok) + require.NotSame(t, oldGeneration, newGeneration) +} + func TestEnqueueStatsUpdateForceReturnsWhenContextCanceled(t *testing.T) { gs := &GlobalStats{ - updateC: make(chan statsinfo.StatsInfoKeyWithContext, 1), + updateC: make(chan statsUpdateJob, 1), queueWatcher: newQueueWatcher(), } + gs.updatingMu.updating = make(map[statsinfo.StatsInfoKey]*updateRecord) queued := statsinfo.StatsInfoKeyWithContext{ Ctx: context.Background(), Key: statsinfo.StatsInfoKey{TableID: 1}, } - gs.updateC <- queued + gs.updateC <- statsUpdateJob{wrapKey: queued} ctx, cancel := context.WithCancel(context.Background()) cancel() - accepted := gs.enqueueStatsUpdate(statsinfo.StatsInfoKeyWithContext{ + generation := gs.currentOrCreateUpdateRecord(statsinfo.StatsInfoKey{TableID: 2}) + accepted := gs.enqueueStatsUpdateForRecord(statsinfo.StatsInfoKeyWithContext{ Ctx: ctx, Key: statsinfo.StatsInfoKey{TableID: 2}, - }, true) + }, true, generation) require.False(t, accepted) - require.Equal(t, queued, <-gs.updateC) + require.Equal(t, queued, (<-gs.updateC).wrapKey) } func TestCacheRemoteInfoIfSubscribedBroadcastsWaiters(t *testing.T) { diff --git a/pkg/vm/engine/disttae/txn_table.go b/pkg/vm/engine/disttae/txn_table.go index 597dbf53a25cf..33c02f0c530ac 100644 --- a/pkg/vm/engine/disttae/txn_table.go +++ b/pkg/vm/engine/disttae/txn_table.go @@ -363,36 +363,118 @@ func (tbl *txnTable) Size(ctx context.Context, columnName string) (uint64, error } func ForeachVisibleObjects( + ctx context.Context, state *logtailreplay.PartitionState, ts types.TS, - fn func(obj objectio.ObjectEntry) error, + fn func(context.Context, objectio.ObjectEntry) error, executor ConcurrentExecutor, visitTombstone bool, ) (err error) { + if cause := context.Cause(ctx); cause != nil { + return cause + } + var executorLifecycle context.Context + if executor != nil { + executorLifecycle = executor.LifecycleContext() + if executorLifecycle != nil { + if cause := context.Cause(executorLifecycle); cause != nil { + return cause + } + } + } iter, err := state.NewObjectsIter(ts, true, visitTombstone) if err != nil { return err } defer iter.Close() - var wg sync.WaitGroup + + taskCtx, cancelTasks := context.WithCancelCause(ctx) + defer cancelTasks(nil) + if executorLifecycle != nil { + stopLifecycleWatch := context.AfterFunc(executorLifecycle, func() { + cause := context.Cause(executorLifecycle) + if cause == nil { + cause = context.Canceled + } + cancelTasks(cause) + }) + defer stopLifecycleWatch() + // Close the check/register race. AfterFunc is only a cancellation + // delivery mechanism; its callback runs asynchronously and is not the + // authoritative lifecycle predicate. + if cause := context.Cause(executorLifecycle); cause != nil { + cancelTasks(cause) + return cause + } + } + var ( + wg sync.WaitGroup + firstErrOnce sync.Once + firstErr error + ) + completeTask := func(taskErr error) { + if taskErr != nil { + firstErrOnce.Do(func() { + firstErr = taskErr + // Stop sibling object I/O and prevent further admission. Already + // admitted work is still joined below before its accumulator dies. + cancelTasks(taskErr) + }) + } + wg.Done() + } + for iter.Next() { + if cause := context.Cause(taskCtx); cause != nil { + err = cause + break + } entry := iter.Entry() if executor != nil { wg.Add(1) - executor.AppendTask(func() error { - defer wg.Done() - return fn(entry) - }) + appendErr := executor.AppendTask( + taskCtx, + func() error { return fn(taskCtx, entry) }, + completeTask, + ) + if appendErr != nil { + // Ownership was not transferred to the executor. + completeTask(appendErr) + err = appendErr + break + } } else { - if err = fn(entry); err != nil { + if err = fn(taskCtx, entry); err != nil { + cancelTasks(err) break } } } if executor != nil { wg.Wait() + if firstErr != nil { + return firstErr + } + // Executor shutdown is a failed traversal even when a running callback + // ignores taskCtx and happens to return nil. Without this check, the + // caller can publish a partial accumulator after the executor lifecycle + // has already canceled the work group. + if cause := context.Cause(taskCtx); cause != nil { + return cause + } + // The lifecycle callback can be scheduled but not yet run. Read the + // executor-owned predicate directly before declaring the joined group a + // success. + if executorLifecycle != nil { + if cause := context.Cause(executorLifecycle); cause != nil { + return cause + } + } } - return + if err != nil { + return err + } + return context.Cause(ctx) } // not accurate! only used by stats @@ -430,10 +512,10 @@ func (tbl *txnTable) MaxAndMinValues(ctx context.Context) ([][2]any, []uint8, er return nil, nil, err } var updateMu sync.Mutex - onObjFn := func(obj objectio.ObjectEntry) error { + onObjFn := func(objCtx context.Context, obj objectio.ObjectEntry) error { var err error location := obj.Location() - if objMeta, err = objectio.FastLoadObjectMeta(ctx, &location, false, fs); err != nil { + if objMeta, err = objectio.FastLoadObjectMeta(objCtx, &location, false, fs); err != nil { return err } updateMu.Lock() @@ -460,6 +542,7 @@ func (tbl *txnTable) MaxAndMinValues(ctx context.Context) ([][2]any, []uint8, er } if err = ForeachVisibleObjects( + ctx, part, types.TimestampToTS(tbl.db.op.SnapshotTS()), onObjFn, @@ -516,7 +599,7 @@ func (tbl *txnTable) GetColumMetadataScanInfo(ctx context.Context, name string, } infoList := make([]*plan.MetadataScanInfo, 0, state.ApproxDataObjectsNum()) var updateMu sync.Mutex - onObjFn := func(obj objectio.ObjectEntry) error { + onObjFn := func(objCtx context.Context, obj objectio.ObjectEntry) error { createTs, err := obj.CreateTime.Marshal() if err != nil { return err @@ -548,7 +631,7 @@ func (tbl *txnTable) GetColumMetadataScanInfo(ctx context.Context, name string, return nil } - objMeta, err := objectio.FastLoadObjectMeta(ctx, &location, false, fs) + objMeta, err := objectio.FastLoadObjectMeta(objCtx, &location, false, fs) if err != nil { return err } @@ -578,6 +661,7 @@ func (tbl *txnTable) GetColumMetadataScanInfo(ctx context.Context, name string, } if err = ForeachVisibleObjects( + ctx, state, types.TimestampToTS(tbl.db.op.SnapshotTS()), onObjFn, @@ -3351,7 +3435,7 @@ func (tbl *txnTable) GetNonAppendableObjectStats(ctx context.Context) ([]objecti sortKeyPos, _ := tbl.getSortKeyPosAndSortKeyIsPK() objStats := make([]objectio.ObjectStats, 0, tbl.ApproxObjectsNum(ctx)) - err = ForeachVisibleObjects(state, snapshot, func(obj objectio.ObjectEntry) error { + err = ForeachVisibleObjects(ctx, state, snapshot, func(_ context.Context, obj objectio.ObjectEntry) error { if obj.GetAppendable() { return nil } diff --git a/pkg/vm/engine/disttae/util.go b/pkg/vm/engine/disttae/util.go index f678d64647da4..d014be34dec1f 100644 --- a/pkg/vm/engine/disttae/util.go +++ b/pkg/vm/engine/disttae/util.go @@ -18,9 +18,10 @@ import ( "bytes" "context" "encoding/hex" - "errors" "fmt" "reflect" + "sync" + "sync/atomic" "go.uber.org/zap" @@ -627,10 +628,16 @@ type concurrentTask func() error // ConcurrentExecutor is an interface that runs tasks concurrently. type ConcurrentExecutor interface { - // AppendTask append the concurrent task to the exuecutor. - AppendTask(concurrentTask) + // AppendTask admits a task or returns without taking ownership. Once + // admitted, the executor calls complete exactly once with either the task + // result or its own lifecycle cancellation. + AppendTask(context.Context, concurrentTask, func(error)) error // Run starts receive task to execute. Run(context.Context) + // LifecycleContext is canceled when the executor stops. A task group uses + // it to cancel work that a worker already owns, while queued work is rejected + // through the completion callback above. + LifecycleContext() context.Context // GetConcurrency returns the concurrency of this executor. GetConcurrency() int } @@ -639,40 +646,145 @@ type concurrentExecutor struct { // concurrency is the concurrency to run the tasks at the same time. concurrency int // task contains all the tasks needed to run. - tasks chan concurrentTask + tasks chan queuedConcurrentTask + + runOnce sync.Once + stopOnce sync.Once + stopCh chan struct{} + stopCause atomic.Pointer[concurrentExecutorStop] + workers sync.WaitGroup + + // submitMu closes the race between a producer admitting work and shutdown + // draining the queue. Shutdown signals stopCh before taking the write lock, + // so a producer blocked on a full queue can always leave promptly. + submitMu sync.RWMutex + stopped bool + lifecycle context.Context +} + +type queuedConcurrentTask struct { + run concurrentTask + complete func(error) +} + +type concurrentExecutorStop struct { + err error } func newConcurrentExecutor(concurrency int) ConcurrentExecutor { return &concurrentExecutor{ concurrency: concurrency, - tasks: make(chan concurrentTask, 2048), + tasks: make(chan queuedConcurrentTask, 2048), + stopCh: make(chan struct{}), } } // AppendTask implements the ConcurrentExecutor interface. -func (e *concurrentExecutor) AppendTask(t concurrentTask) { - e.tasks <- t +func (e *concurrentExecutor) AppendTask( + ctx context.Context, + t concurrentTask, + complete func(error), +) error { + e.submitMu.RLock() + defer e.submitMu.RUnlock() + if e.stopped { + return e.stoppedError() + } + + // Prefer rejection once shutdown is already observable. The second select + // still covers shutdown or caller cancellation racing queue admission. + select { + case <-e.stopCh: + return e.stoppedError() + default: + } + select { + case e.tasks <- queuedConcurrentTask{run: t, complete: complete}: + return nil + case <-ctx.Done(): + return context.Cause(ctx) + case <-e.stopCh: + return e.stoppedError() + } } // Run implements the ConcurrentExecutor interface. func (e *concurrentExecutor) Run(ctx context.Context) { - for i := 0; i < e.concurrency; i++ { - go func() { - for { - select { - case <-ctx.Done(): - return - - case t := <-e.tasks: - if err := t(); err != nil { - if !errors.Is(err, context.Canceled) { - logutil.Errorf("failed to execute task: %v", err) - } - } - } + e.runOnce.Do(func() { + e.submitMu.Lock() + e.lifecycle = ctx + e.submitMu.Unlock() + e.workers.Add(e.concurrency) + for i := 0; i < e.concurrency; i++ { + go e.runWorker() + } + go e.stopWhenDone(ctx) + }) +} + +func (e *concurrentExecutor) LifecycleContext() context.Context { + e.submitMu.RLock() + defer e.submitMu.RUnlock() + return e.lifecycle +} + +func (e *concurrentExecutor) runWorker() { + defer e.workers.Done() + for { + // Make shutdown deterministic after the current task. Without this + // priority check, a ready queue and a closed stop channel could keep + // selecting more data work while shutdown waits. + select { + case <-e.stopCh: + return + default: + } + select { + case <-e.stopCh: + return + case task := <-e.tasks: + err := task.run() + if task.complete != nil { + task.complete(err) + } + } + } +} + +func (e *concurrentExecutor) stopWhenDone(ctx context.Context) { + <-ctx.Done() + cause := context.Cause(ctx) + if cause == nil { + cause = context.Canceled + } + e.stopOnce.Do(func() { + e.stopCause.Store(&concurrentExecutorStop{err: cause}) + close(e.stopCh) + }) + + // Wait for every producer that could still admit a task, then prevent all + // future admission before workers stop and the remaining queue is rejected. + e.submitMu.Lock() + e.stopped = true + e.submitMu.Unlock() + e.workers.Wait() + for { + select { + case task := <-e.tasks: + if task.complete != nil { + task.complete(cause) } - }() + default: + return + } + } +} + +func (e *concurrentExecutor) stoppedError() error { + if stopped := e.stopCause.Load(); stopped != nil && stopped.err != nil { + return stopped.err } + return context.Canceled } // GetConcurrency implements the ConcurrentExecutor interface. diff --git a/pkg/vm/engine/disttae/util_test.go b/pkg/vm/engine/disttae/util_test.go index 35183131add4c..15d936058201d 100644 --- a/pkg/vm/engine/disttae/util_test.go +++ b/pkg/vm/engine/disttae/util_test.go @@ -16,10 +16,12 @@ package disttae import ( "context" + "errors" "io" "math/rand" "slices" "sync" + "sync/atomic" "testing" "time" @@ -661,24 +663,452 @@ func TestConcurrentExecutor_Run(t *testing.T) { require.Equal(t, 3, ex.GetConcurrency()) var wg sync.WaitGroup - wg.Add(1) - ex.AppendTask(func() error { - defer wg.Done() + submit := func(task concurrentTask) { + wg.Add(1) + err := ex.AppendTask(ctx, task, func(error) { + wg.Done() + }) + if err != nil { + wg.Done() + } + require.NoError(t, err) + } + submit(func() error { return nil }) + submit(func() error { return context.Canceled }) + submit(func() error { return io.EOF }) + wg.Wait() +} + +func TestConcurrentExecutorRejectsQueuedTasksOnShutdown(t *testing.T) { + ex := newConcurrentExecutor(1) + executorCtx, stopExecutor := context.WithCancel(context.Background()) + defer stopExecutor() + ex.Run(executorCtx) + + firstStarted := make(chan struct{}) + releaseFirst := make(chan struct{}) + require.NoError(t, ex.AppendTask(context.Background(), func() error { + close(firstStarted) + <-releaseFirst return nil - }) + }, func(err error) { + if err != nil { + t.Errorf("running task was unexpectedly rejected: %v", err) + } + })) + select { + case <-firstStarted: + case <-time.After(time.Second): + close(releaseFirst) + t.Fatal("executor did not start its admitted task") + } - wg.Add(1) - ex.AppendTask(func() error { - defer wg.Done() - return context.Canceled - }) + secondRan := atomic.Bool{} + secondRejected := make(chan error, 1) + require.NoError(t, ex.AppendTask(context.Background(), func() error { + secondRan.Store(true) + return nil + }, func(err error) { + secondRejected <- err + })) + + stopExecutor() + <-ex.(*concurrentExecutor).stopCh + close(releaseFirst) + select { + case err := <-secondRejected: + require.ErrorIs(t, err, context.Canceled) + case <-time.After(time.Second): + t.Fatal("executor shutdown abandoned an admitted task") + } + require.False(t, secondRan.Load()) - wg.Add(1) - ex.AppendTask(func() error { - defer wg.Done() - return io.EOF - }) - wg.Wait() + err := ex.AppendTask(context.Background(), func() error { return nil }, nil) + require.ErrorIs(t, err, context.Canceled) +} + +func visibleObjectStateForExecutorTest(t *testing.T, count int) *logtailreplay.PartitionState { + t.Helper() + ctx := context.Background() + state := logtailreplay.NewPartitionState("", true, 42, false) + state.UpdateDuration(types.TS{}, types.MaxTs()) + for i := 0; i < count; i++ { + oid := types.NewObjectid() + stats := objectio.NewObjectStatsWithObjectID(&oid, false, false, false) + require.NoError(t, objectio.SetObjectStatsSize(stats, 1)) + require.NoError(t, state.HandleObjectEntry(ctx, nil, objectio.ObjectEntry{ + ObjectStats: *stats, + CreateTime: types.BuildTS(int64(i+1), 0), + }, false)) + } + return state +} + +// delayedLifecycleContext models a conforming lifecycle whose Done predicate is +// observable before its AfterFunc notification is dispatched. This keeps the +// cancellation race deterministic without sleeps or scheduler assumptions. +type delayedLifecycleContext struct { + done chan struct{} + + mu sync.Mutex + err error + callbacks []*delayedLifecycleCallback + // cancelOnRegister closes Done while context.AfterFunc is registering its + // callback, but before that callback is dispatched. + cancelOnRegister bool +} + +type delayedLifecycleCallback struct { + active bool + fn func() +} + +func newDelayedLifecycleContext() *delayedLifecycleContext { + return &delayedLifecycleContext{done: make(chan struct{})} +} + +func (c *delayedLifecycleContext) Deadline() (time.Time, bool) { return time.Time{}, false } +func (c *delayedLifecycleContext) Done() <-chan struct{} { return c.done } +func (c *delayedLifecycleContext) Value(any) any { return nil } + +func (c *delayedLifecycleContext) Err() error { + c.mu.Lock() + defer c.mu.Unlock() + return c.err +} + +func (c *delayedLifecycleContext) AfterFunc(fn func()) func() bool { + c.mu.Lock() + callback := &delayedLifecycleCallback{active: true, fn: fn} + c.callbacks = append(c.callbacks, callback) + if c.cancelOnRegister && c.err == nil { + c.err = context.Canceled + close(c.done) + } + c.mu.Unlock() + return func() bool { + c.mu.Lock() + defer c.mu.Unlock() + if !callback.active { + return false + } + callback.active = false + return true + } +} + +func (c *delayedLifecycleContext) cancelWithoutNotification() { + c.mu.Lock() + defer c.mu.Unlock() + if c.err != nil { + return + } + c.err = context.Canceled + close(c.done) +} + +type inlineLifecycleExecutor struct { + lifecycle context.Context +} + +func (e *inlineLifecycleExecutor) AppendTask( + ctx context.Context, + task concurrentTask, + complete func(error), +) error { + if cause := context.Cause(ctx); cause != nil { + return cause + } + err := task() + if complete != nil { + complete(err) + } + return nil +} + +func (*inlineLifecycleExecutor) Run(context.Context) {} + +func (e *inlineLifecycleExecutor) LifecycleContext() context.Context { + return e.lifecycle +} + +func (*inlineLifecycleExecutor) GetConcurrency() int { return 1 } + +func TestForeachVisibleObjectsPropagatesConcurrentTaskError(t *testing.T) { + state := visibleObjectStateForExecutorTest(t, 2) + ex := newConcurrentExecutor(2) + executorCtx, stopExecutor := context.WithCancel(context.Background()) + defer stopExecutor() + ex.Run(executorCtx) + + wantErr := errors.New("object metadata unavailable") + allStarted := make(chan struct{}) + var releaseOnce sync.Once + t.Cleanup(func() { releaseOnce.Do(func() { close(allStarted) }) }) + var calls atomic.Int32 + result := make(chan error, 1) + go func() { + result <- ForeachVisibleObjects( + context.Background(), state, types.MaxTs(), + func(_ context.Context, _ objectio.ObjectEntry) error { + call := calls.Add(1) + if call == 2 { + releaseOnce.Do(func() { close(allStarted) }) + } + <-allStarted + if call == 1 { + return wantErr + } + return nil + }, + ex, + false, + ) + }() + var err error + select { + case err = <-result: + case <-time.After(time.Second): + t.Fatal("concurrent visible-object traversal did not join its tasks") + } + require.ErrorIs(t, err, wantErr) + require.Equal(t, int32(2), calls.Load(), "one successful and one failed task must both be joined") +} + +func TestForeachVisibleObjectsHonorsCancellationBeforeAdmission(t *testing.T) { + state := visibleObjectStateForExecutorTest(t, 1) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + cancel() + + var called atomic.Bool + err := ForeachVisibleObjects( + ctx, state, types.MaxTs(), + func(context.Context, objectio.ObjectEntry) error { + called.Store(true) + return nil + }, + newConcurrentExecutor(1), + false, + ) + require.ErrorIs(t, err, context.Canceled) + require.False(t, called.Load()) +} + +func TestForeachVisibleObjectsCancelsInFlightTask(t *testing.T) { + state := visibleObjectStateForExecutorTest(t, 1) + ex := newConcurrentExecutor(1) + executorCtx, stopExecutor := context.WithCancel(context.Background()) + defer stopExecutor() + ex.Run(executorCtx) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + taskStarted := make(chan struct{}) + result := make(chan error, 1) + go func() { + result <- ForeachVisibleObjects( + ctx, state, types.MaxTs(), + func(taskCtx context.Context, _ objectio.ObjectEntry) error { + close(taskStarted) + <-taskCtx.Done() + return context.Cause(taskCtx) + }, + ex, + false, + ) + }() + + select { + case <-taskStarted: + case <-time.After(time.Second): + t.Fatal("executor did not start the admitted visible-object task") + } + cancel() + select { + case err := <-result: + require.ErrorIs(t, err, context.Canceled) + case <-time.After(time.Second): + t.Fatal("in-flight object task ignored caller cancellation") + } +} + +func TestForeachVisibleObjectsCancelsInFlightTaskOnExecutorShutdown(t *testing.T) { + state := visibleObjectStateForExecutorTest(t, 1) + ex := newConcurrentExecutor(1) + executorCtx, stopExecutor := context.WithCancel(context.Background()) + defer stopExecutor() + ex.Run(executorCtx) + + taskStarted := make(chan struct{}) + result := make(chan error, 1) + go func() { + result <- ForeachVisibleObjects( + context.Background(), state, types.MaxTs(), + func(taskCtx context.Context, _ objectio.ObjectEntry) error { + close(taskStarted) + <-taskCtx.Done() + return context.Cause(taskCtx) + }, + ex, + false, + ) + }() + + select { + case <-taskStarted: + case <-time.After(time.Second): + t.Fatal("executor did not start the admitted visible-object task") + } + stopExecutor() + select { + case err := <-result: + require.ErrorIs(t, err, context.Canceled) + case <-time.After(time.Second): + t.Fatal("in-flight object task outlived executor shutdown") + } +} + +func TestForeachVisibleObjectsRejectsExecutorShutdownWhenTaskReturnsNil(t *testing.T) { + state := visibleObjectStateForExecutorTest(t, 1) + ex := newConcurrentExecutor(1) + executorCtx, stopExecutor := context.WithCancel(context.Background()) + t.Cleanup(stopExecutor) + ex.Run(executorCtx) + + taskStarted := make(chan struct{}) + result := make(chan error, 1) + go func() { + result <- ForeachVisibleObjects( + context.Background(), state, types.MaxTs(), + func(taskCtx context.Context, _ objectio.ObjectEntry) error { + close(taskStarted) + <-taskCtx.Done() + // Model a callback that observes shutdown only as a release + // signal and fails to propagate the context error itself. + return nil + }, + ex, + false, + ) + }() + + select { + case <-taskStarted: + case <-time.After(time.Second): + t.Fatal("executor did not start the admitted visible-object task") + } + stopExecutor() + select { + case err := <-result: + require.ErrorIs(t, err, context.Canceled, + "executor shutdown must remain visible even when a callback returns nil") + case <-time.After(time.Second): + t.Fatal("visible-object traversal did not join the shutdown task") + } +} + +func TestForeachVisibleObjectsReadsLifecyclePredicateAfterJoin(t *testing.T) { + state := visibleObjectStateForExecutorTest(t, 1) + lifecycle := newDelayedLifecycleContext() + ex := &inlineLifecycleExecutor{lifecycle: lifecycle} + + err := ForeachVisibleObjects( + context.Background(), state, types.MaxTs(), + func(context.Context, objectio.ObjectEntry) error { + lifecycle.cancelWithoutNotification() + return nil + }, + ex, + false, + ) + require.ErrorIs(t, err, context.Canceled, + "executor shutdown is a failed traversal even before async notification dispatch") +} + +func TestForeachVisibleObjectsClosesLifecycleWatcherRegistrationRace(t *testing.T) { + state := visibleObjectStateForExecutorTest(t, 1) + lifecycle := newDelayedLifecycleContext() + lifecycle.cancelOnRegister = true + ex := &inlineLifecycleExecutor{lifecycle: lifecycle} + var called atomic.Bool + + err := ForeachVisibleObjects( + context.Background(), state, types.MaxTs(), + func(context.Context, objectio.ObjectEntry) error { + called.Store(true) + return nil + }, + ex, + false, + ) + require.ErrorIs(t, err, context.Canceled) + require.False(t, called.Load(), + "work must not be admitted after lifecycle cancellation becomes observable") +} + +func TestCollectAndCalculateStatsRejectsStoppedExecutorOnZeroObjectFastPath(t *testing.T) { + lifecycle, cancel := context.WithCancel(context.Background()) + cancel() + ex := &inlineLifecycleExecutor{lifecycle: lifecycle} + req := &updateStatsRequest{ + statsInfo: plan2.NewStatsInfo(), + tableDef: &plan.TableDef{Cols: []*plan.ColDef{{Name: "__mo_rowid"}}}, + approxObjectNum: 0, + } + + _, err := CollectAndCalculateStats(context.Background(), req, ex) + require.ErrorIs(t, err, context.Canceled, + "the zero-object fast path must not bypass executor lifecycle failure") +} + +func TestCollectAndCalculateStatsDoesNotApplyFailedObjectScan(t *testing.T) { + ctx := context.Background() + state := logtailreplay.NewPartitionState("", true, 42, false) + state.UpdateDuration(types.TS{}, types.MaxTs()) + location := objectio.NewRandomLocation(1, 128) + stats := objectio.NewObjectStats() + objectio.SetObjectStatsLocation(stats, location) + require.NoError(t, objectio.SetObjectStatsSize(stats, 1)) + require.NoError(t, state.HandleObjectEntry(ctx, nil, objectio.ObjectEntry{ + ObjectStats: *stats, + CreateTime: types.BuildTS(1, 0), + }, false)) + require.Equal(t, 1, state.ApproxDataObjectsNum()) + + fs, err := fileservice.NewMemoryFS( + defines.SharedFileServiceName, + fileservice.DisabledCacheConfig, + nil, + ) + require.NoError(t, err) + published := plan2.NewStatsInfo() + req := &updateStatsRequest{ + statsInfo: published, + tableDef: &plan.TableDef{ + Name: "events", + Cols: []*plan.ColDef{ + {Name: "event_id", Seqnum: 0, Typ: plan.Type{Id: int32(types.T_int64)}}, + {Name: "__mo_rowid", Seqnum: 1, Typ: plan.Type{Id: int32(types.T_Rowid)}}, + }, + }, + partitionState: state, + fs: fs, + ts: types.MaxTs(), + approxObjectNum: 1, + samplingMode: "full", + } + ex := newConcurrentExecutor(1) + executorCtx, stopExecutor := context.WithCancel(context.Background()) + defer stopExecutor() + ex.Run(executorCtx) + + _, err = CollectAndCalculateStats(ctx, req, ex) + require.Error(t, err) + require.Zero(t, published.TableCnt) + require.Zero(t, published.AccurateObjectNumber) + require.Empty(t, published.NdvMap) + require.Empty(t, published.ShuffleRangeMap) } func TestShrinkBatchWithRowids(t *testing.T) { diff --git a/pkg/vm/engine/types.go b/pkg/vm/engine/types.go index 10ec9f4a75dcf..9586a069f8591 100644 --- a/pkg/vm/engine/types.go +++ b/pkg/vm/engine/types.go @@ -1327,6 +1327,14 @@ type Engine interface { LatestLogtailAppliedTime() timestamp.Timestamp } +// StatsRefresher is an optional engine capability for statements that define +// a synchronous statistics-publication boundary, such as ANALYZE TABLE. +// Implementations must not return until Stats() can observe the returned +// statistics on the local engine instance. +type StatsRefresher interface { + RefreshTableStats(ctx context.Context, key pb.StatsInfoKey) (*pb.StatsInfo, error) +} + type VectorPool interface { PutBatch(bat *batch.Batch) GetVector(typ types.Type) *vector.Vector diff --git a/test/distributed/cases/analyze/analyze_stmt.result b/test/distributed/cases/analyze/analyze_stmt.result index d1017a5365976..d410c1feddcce 100644 --- a/test/distributed/cases/analyze/analyze_stmt.result +++ b/test/distributed/cases/analyze/analyze_stmt.result @@ -42,6 +42,13 @@ approx_count_distinct(select) approx_count_distinct(a-b) approx_count_distinct(t select 'AFTER_EXPANDED_QUOTED'; AFTER_EXPANDED_QUOTED AFTER_EXPANDED_QUOTED +create view v_analyze as select a, b from t_analyze_01; +analyze table v_analyze(a); +approx_count_distinct(a) +2 +select 'AFTER_VIEW_ANALYZE'; +AFTER_VIEW_ANALYZE +AFTER_VIEW_ANALYZE create database `select-db`; create table `select-db`.`tick``table`(`a-b` int); insert into `select-db`.`tick``table` values (1),(1),(2); @@ -139,6 +146,20 @@ AFTER_TXN_MULTI AFTER_TXN_MULTI rollback; begin; +create table txn_created_analyze(a int); +insert into txn_created_analyze values (1), (2); +analyze table txn_created_analyze(a); +approx_count_distinct(a) +2 +rollback; +drop table if exists txn_created_analyze; +begin; +insert into t_analyze_01 values (3, 30); +analyze table t_analyze_01(a); +approx_count_distinct(a) +3 +rollback; +begin; check table t_analyze_01; not supported: CHECK TABLE is not supported in MatrixOne rollback; @@ -146,6 +167,7 @@ begin; show profile; not supported: SHOW PROFILE is not supported in MatrixOne rollback; +drop view v_analyze; drop table t_analyze_01; drop table t_analyze_02; drop table quoted_cols; diff --git a/test/distributed/cases/analyze/analyze_stmt.sql b/test/distributed/cases/analyze/analyze_stmt.sql index 291777a59de94..9ac29b2c4f055 100644 --- a/test/distributed/cases/analyze/analyze_stmt.sql +++ b/test/distributed/cases/analyze/analyze_stmt.sql @@ -41,6 +41,12 @@ select 'AFTER_EXPLICIT_QUOTED'; analyze table quoted_cols; select 'AFTER_EXPANDED_QUOTED'; +-- views retain the legacy derived-query result and must not be subscribed as +-- physical optimizer-statistics tables +create view v_analyze as select a, b from t_analyze_01; +analyze table v_analyze(a); +select 'AFTER_VIEW_ANALYZE'; + -- quoted database, table, and column identifiers create database `select-db`; create table `select-db`.`tick``table`(`a-b` int); @@ -111,6 +117,22 @@ analyze table t_analyze_01, t_analyze_02; select 'AFTER_TXN_MULTI'; rollback; +-- A transaction-local table is visible to the derived ANALYZE query but not +-- to the process-global optimizer statistics refresher. +begin; +create table txn_created_analyze(a int); +insert into txn_created_analyze values (1), (2); +analyze table txn_created_analyze(a); +rollback; +drop table if exists txn_created_analyze; + +-- The result includes uncommitted workspace rows, while global optimizer +-- statistics remain at the committed visibility boundary. +begin; +insert into t_analyze_01 values (3, 30); +analyze table t_analyze_01(a); +rollback; + begin; check table t_analyze_01; rollback; @@ -120,6 +142,7 @@ show profile; rollback; -- cleanup +drop view v_analyze; drop table t_analyze_01; drop table t_analyze_02; drop table quoted_cols;