From b31831aa17b6208c69d01d9cf8962504ac7d359f Mon Sep 17 00:00:00 2001 From: Jens Neuse Date: Thu, 11 Jun 2026 11:15:19 +0200 Subject: [PATCH 1/5] perf(resolve): pre-size buffers and slab-allocate parallel fetch results Port of the fork's hot-path pre-sizing (818aaa2): single-slab result allocation in resolveParallel, HTTP request buffer pre-sizing, singleflight key buffer reuse, and resolvable field-path buffer pre-sizing. Verified at parity on the flags-off A/B benchmark gate. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../datasource/httpclient/nethttpclient.go | 16 ++++++++++++++++ .../resolve/inbound_request_singleflight.go | 19 +++++++++++++++++++ v2/pkg/engine/resolve/loader.go | 5 ++++- v2/pkg/engine/resolve/resolvable.go | 15 +++++++++++---- 4 files changed, 50 insertions(+), 5 deletions(-) diff --git a/v2/pkg/engine/datasource/httpclient/nethttpclient.go b/v2/pkg/engine/datasource/httpclient/nethttpclient.go index 02025437c2..04a6ef578c 100644 --- a/v2/pkg/engine/datasource/httpclient/nethttpclient.go +++ b/v2/pkg/engine/datasource/httpclient/nethttpclient.go @@ -251,6 +251,22 @@ func makeHTTPRequest(client *http.Client, ctx context.Context, baseHeaders http. // It'll know best the lifecycle of the buffer // Using an arena here just increased overall memory usage out := buffer(ctx) + // Pre-size the buffer from the advertised Content-Length to avoid the repeated + // grow-and-copy reallocations bytes.Buffer.ReadFrom would otherwise perform while + // reading the response body. This reduces total allocation volume and GC pressure. + // Capped to avoid over-allocating on a large or forged Content-Length; the buffer + // still grows by doubling beyond the cap. For compressed responses Content-Length is + // the compressed size, so this remains a safe lower-bound hint. + if cl := response.ContentLength; cl > 0 { + const maxPreAlloc = 1 << 21 // 2 MiB + want := int(cl) + if want > maxPreAlloc { + want = maxPreAlloc + } + if want > out.Cap() { + out.Grow(want) + } + } _, err = out.ReadFrom(respReader) if err != nil { return nil, err diff --git a/v2/pkg/engine/resolve/inbound_request_singleflight.go b/v2/pkg/engine/resolve/inbound_request_singleflight.go index 5affdb4d36..21742c7339 100644 --- a/v2/pkg/engine/resolve/inbound_request_singleflight.go +++ b/v2/pkg/engine/resolve/inbound_request_singleflight.go @@ -90,6 +90,25 @@ func (r *InboundRequestSingleFlight) GetOrCreate(ctx *Context, response *GraphQL shard := r.shardFor(key) + // Fast path: a leader is already inflight for this key -> become a follower without + // allocating a new InflightRequest. In a deduplication burst most callers are followers, + // so this avoids one heap allocation per follower. The LoadOrStore below still resolves + // the leader-election race for the case where Load misses but another goroutine stores + // concurrently before our LoadOrStore. + if existing, ok := shard.m.Load(key); ok { + request := existing.(*InflightRequest) + request.AddFollower() + select { + case <-request.Done: + if request.Err != nil { + return nil, request.Err + } + return request, nil + case <-ctx.ctx.Done(): + return nil, ctx.ctx.Err() + } + } + request := &InflightRequest{ Done: make(chan struct{}), ID: key, diff --git a/v2/pkg/engine/resolve/loader.go b/v2/pkg/engine/resolve/loader.go index 6e6c82ef0d..daddafde9a 100644 --- a/v2/pkg/engine/resolve/loader.go +++ b/v2/pkg/engine/resolve/loader.go @@ -230,6 +230,9 @@ func (l *Loader) resolveParallel(nodes []*FetchTreeNode) error { return nil } results := make([]*result, len(nodes)) + // Allocate the result structs as a single slab instead of len(nodes) individual heap + // allocations; the pointers into it stay valid for the function lifetime. + resultStore := make([]result, len(nodes)) defer func() { for i := range results { // no-op if tools == nil @@ -239,7 +242,7 @@ func (l *Loader) resolveParallel(nodes []*FetchTreeNode) error { itemsItems := make([][]*astjson.Value, len(nodes)) g, ctx := errgroup.WithContext(l.ctx.ctx) for i := range nodes { - results[i] = &result{} + results[i] = &resultStore[i] itemsItems[i] = l.selectItemsForPath(nodes[i].Item.FetchPath) f := nodes[i].Item.Fetch item := nodes[i].Item diff --git a/v2/pkg/engine/resolve/resolvable.go b/v2/pkg/engine/resolve/resolvable.go index 13695e1e4b..79ed6475ce 100644 --- a/v2/pkg/engine/resolve/resolvable.go +++ b/v2/pkg/engine/resolve/resolvable.go @@ -7,7 +7,6 @@ import ( "fmt" "io" "strconv" - "strings" "github.com/cespare/xxhash/v2" "github.com/pkg/errors" @@ -57,6 +56,9 @@ type Resolvable struct { typeNames [][]byte marshalBuf []byte + // fieldPathBuf is a reusable scratch buffer for currentFieldPath to avoid a + // per-call []string + strings.Join allocation on the resolve hot path. + fieldPathBuf []byte enclosingTypeNames []string @@ -1044,13 +1046,18 @@ func (r *Resolvable) recordObjectTypeStats(obj *Object, typeName []byte) { // Helper to build JSON path (field names only, no array indices) func (r *Resolvable) currentFieldPath() string { - var parts []string + buf := r.fieldPathBuf[:0] for _, elem := range r.path { if elem.Name != "" { - parts = append(parts, elem.Name) + if len(buf) > 0 { + buf = append(buf, '.') + } + buf = append(buf, elem.Name...) } } - return strings.Join(parts, ".") + r.fieldPathBuf = buf + // string(buf) copies; required because the result is used as a map key. + return string(buf) } func (r *Resolvable) walkNull() bool { From 7c355083ba8816916016040acfc6d44da3cb5633 Mon Sep 17 00:00:00 2001 From: Jens Neuse Date: Thu, 11 Jun 2026 11:43:47 +0200 Subject: [PATCH 2/5] feat(postprocess): dominance-based schedule-tree scheduler (WithBuildScheduleTree) Ports the improve-query-order scheduler with two load-bearing deviations, both found by adversarial review against real planner output: 1. validateParallelLeaves trusts FetchID dependency edges ONLY. The upstream path-containment heuristics false-positive on the canonical federation plan (a fetch attaching under an ancestor-provided subtree is independent of a sibling merging into a prefix path; the planner emits explicit edges when a dependency is real) and panicked every affected request. 2. buildScheduleTreeProcessor never panics: any scheduler/validator error falls back to the legacy orderSequenceByDependencies+createParallelNodes pipeline, planning exactly as with the option off. Includes the property/scenario suites, a 13-fetch real-plan schedule pin, and the ADR. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/adr/0001-query-plan-scheduler.md | 517 ++++++++++++ .../query_order_baseline_test.go | 345 ++++++++ .../engine/postprocess/build_schedule_tree.go | 797 ++++++++++++++++++ .../build_schedule_tree_bench_plan_test.go | 113 +++ .../build_schedule_tree_property_test.go | 163 ++++ .../build_schedule_tree_scenarios_test.go | 307 +++++++ v2/pkg/engine/postprocess/postprocess.go | 48 +- 7 files changed, 2274 insertions(+), 16 deletions(-) create mode 100644 docs/adr/0001-query-plan-scheduler.md create mode 100644 v2/pkg/engine/datasource/graphql_datasource/query_order_baseline_test.go create mode 100644 v2/pkg/engine/postprocess/build_schedule_tree.go create mode 100644 v2/pkg/engine/postprocess/build_schedule_tree_bench_plan_test.go create mode 100644 v2/pkg/engine/postprocess/build_schedule_tree_property_test.go create mode 100644 v2/pkg/engine/postprocess/build_schedule_tree_scenarios_test.go diff --git a/docs/adr/0001-query-plan-scheduler.md b/docs/adr/0001-query-plan-scheduler.md new file mode 100644 index 0000000000..e41b8f88c0 --- /dev/null +++ b/docs/adr/0001-query-plan-scheduler.md @@ -0,0 +1,517 @@ +# ADR 0001: Query Plan Scheduler + +Status: Accepted. +Date: 2026-05-09. + +## Context + +The query planner emits fetches as a flat sequence and then runs postprocessors that sort by dependencies and group eligible siblings into parallel nodes. +That design is simple, but it leaves latency on the table when independent dependency chains coexist in the same operation. +The motivating shape is a query that fetches `me` from one subgraph, fetches `organisations` from another subgraph, and then fetches `currentPractice` through an entity hop that depends only on `me`. +The old postprocessor emits `Sequence(Parallel(user, organisation), practice)`. +That forces `practice` to wait for `organisation`, even though `organisation` does not provide any dependency required by `practice`. +The better tree is `Parallel(Sequence(user, practice), organisation)`. +That tree has the critical path `max(user + practice, organisation)` instead of `max(user, organisation) + practice`. + +The old executor could not safely run that better tree. +Its parallel implementation assumed every child of a `Parallel` node was a leaf fetch. +It preselected merge targets, ran network loads concurrently, and then merged results serially after all loads completed. +If a `Parallel` child became a nested `Sequence`, the loader would have to recurse inside a goroutine. +That would make selection, input rendering, merge, JSON arena writes, and tainted-object map writes happen concurrently. +Those operations are not thread safe. + +The scheduler therefore had two coupled requirements. +First, the executor needed to support nested sequence and parallel fetch trees without racing shared response state. +Second, the postprocessor needed a dependency-aware tree builder that could emit nested plans while preserving existing default behavior until the new path is explicitly enabled. + +## Decision + +We keep the legacy postprocessors wired by default. +The new scheduler is enabled only with `WithBuildScheduleTree()`. +When enabled, it replaces the `orderSequenceByDependencies` and `createParallelNodes` pair after missing nested dependencies are added and after concrete fetch types are created. +The old processors remain in the codebase and continue to support the existing options. + +The new scheduler builds two candidate trees from the same fetch DAG. +It builds a component-aware level scheduler and a component-aware eager-inline scheduler. +It then picks the eager-inline tree only when a symbolic dominance check proves that it can never be slower than the level tree for any non-negative duration vector. +If the dominance check fails, the scheduler returns the level tree. +This is deliberately conservative. +It guarantees no regression against the level scheduler while still taking the stronger nested-chain shape when the proof applies. + +The executor now supports nested parallel plans with a three-phase leaf protocol. +Prepare selects merge targets and renders input while holding the merge mutex when nested concurrency is possible. +Load performs the network round trip without the merge mutex. +Merge writes the result into the response while holding the merge mutex when nested concurrency is possible. +For flat plans, the existing fast path remains in use and does not acquire the mutex. + +## Algorithms + +### scheduleLevel + +The level scheduler receives a set of fetch nodes and the dependency DAG. +It first partitions the set into weakly connected components. +Independent components are scheduled recursively and wrapped in a `Parallel` node sorted by minimum reachable fetch ID. +Inside a single component, it finds all roots whose parents are not present in the current node set. +Those roots form the current level. +The remaining nodes are scheduled recursively after stripping the root level from the component. +A single root is emitted as a leaf. +Multiple roots are emitted as `Parallel(root_1, root_2, ...)`. +When there is a recursive remainder, the result is `Sequence(current_level, remainder)`. +If a non-empty component has no root, the DAG is cyclic and scheduling fails with a clear error. + +This algorithm matches the old level-barrier intuition but adds component awareness. +It is optimal under uniform fetch durations because every node is scheduled at its earliest possible topological level. +It can be suboptimal under skewed durations because it synchronizes all chains at every level boundary. + +### scheduleSP + +The eager-inline scheduler also begins with weakly connected component partitioning. +That step prevents unrelated components from being serialized together by an outer batch. +For a single component, the scheduler processes ready roots in batches. +Each processed node emits itself, then recursively consumes descendants that are uniquely ready within that branch. +Children with multiple pending parents are recorded in an unhandled map rather than emitted immediately. +When sibling branches are merged, the remaining-parent sets are intersected. +If the intersection becomes empty, the child becomes ready at the outer level. + +This intersection rule is the critical recurrence. +If `A` and `B` both feed `X`, processing `A` reports `X` pending on `B`, and processing `B` reports `X` pending on `A`. +The intersection is empty, so `X` is ready after the parallel batch containing `A` and `B`. +This preserves joins while allowing independent chains such as `A -> C` and `B -> D` to run as `Parallel(Sequence(A, C), Sequence(B, D))`. + +The eager-inline tree is often better for skewed chains that join later. +It is not universally better. +On non-series-parallel shapes with a leaf side branch and a shared join, eager inlining can place the shared join behind unrelated work. +That is why the driver does not pick it based on uniform makespan alone. + +### Dominance Theorem + +For a fetch tree `T`, define `paths(T)` recursively. +For a single fetch, `paths(Single(n)) = {{n}}`. +For a sequence, `paths(Sequence(c1, ..., ck))` is the cross product of child paths with set union. +For a parallel node, `paths(Parallel(c1, ..., ck))` is the union of the child path sets. +For any non-negative duration vector `d`, `makespan(T, d)` is the maximum sum of durations over any path in `paths(T)`. + +Tree `A` dominates tree `B` if and only if every path in `A` is set-contained in some path in `B`. +If that condition holds, choose the heaviest path in `A` under any duration vector. +By containment, there is a path in `B` containing every node from that path. +Because durations are non-negative, the `B` path has weight at least the `A` path. +Therefore `makespan(A, d) <= makespan(B, d)` for all `d`. +For the converse, suppose some path in `A` is not contained in any path in `B`. +Assign duration one to nodes on that `A` path and zero to all other nodes. +Every `B` path misses at least one node from the `A` path, so every `B` path is strictly lighter. +That contradicts dominance for all duration vectors. + +The driver uses this theorem directly. +It computes `scheduleSP` and `scheduleLevel`. +It validates both. +It returns `scheduleSP` only when `dominates(scheduleSP, scheduleLevel)` is true. +Otherwise it returns `scheduleLevel`. + +### Validation + +The validator walks the produced tree and enforces dependency order. +For a sequence, each child is validated with the fetch IDs provided by all previous children. +For a parallel node, every branch is validated with the same incoming dependency set. +The validator flattens every leaf under each parallel branch and checks all cross-branch leaf pairs. +It rejects fetch-ID dependencies across branches. +It also rejects response-path containment across branches because that indicates one branch may need data produced by the other. +Provided paths are computed by fetch kind using `SingleFetch`, `EntityFetch`, or `BatchEntityFetch` postprocessing merge paths. +The scheduler and validator use the `Fetch.Dependencies()` interface rather than assuming a concrete single fetch type. + +### Executor Split + +The loader now computes `useMergeMu` once at `LoadGraphQLResponseData` entry. +The predicate is true when any `Parallel` node has at least one non-`Single` child. +Flat plans keep the old path. +The old path preselects items, runs leaf loads concurrently, and merges serially after the errgroup completes. +Nested plans recurse through the fetch tree inside errgroup branches. +For those plans, each leaf runs prepare, load, and merge. +Prepare and merge call `maybeLock`, which acquires `mergeMu` only when `useMergeMu` is true. +Load does not hold the mutex. + +This keeps network I/O concurrent while serializing access to `resolvable.data`, the JSON arena, and tainted-object tracking. +It also avoids adding mutex overhead to the only tree shape produced by the default legacy postprocessors. + +## Consequences + +The new option can produce nested fetch trees that reduce latency for independent chains. +Default behavior is unchanged, so existing golden query-plan fixtures remain stable. +The executor is more capable, but flat plans still use the existing fast path. +The validator may surface previously hidden dependency omissions. +That is intentional. +If a fetch needs data from another branch, the planner should declare that dependency rather than relying on a level barrier to hide the missing edge. + +The scheduler performs more work than the legacy pair because it builds two trees and enumerates path sets. +The expected federation DAG sizes are small enough for this to be negligible. +The implementation keeps deterministic ordering by sorting branches by minimum reachable fetch ID. + +## Alternatives Considered + +We considered replacing the legacy scheduler outright. +That would force fixture churn and make executor changes harder to review. +We rejected that in favor of opt-in migration. + +We considered choosing the smaller uniform-duration makespan. +That is unsound under skewed durations. +A tree can be better under uniform weights while worse under realistic latency distributions. +The path-set dominance theorem gives a stronger and exact condition. + +We considered a cost-based scheduler using observed subgraph durations. +That could close more gaps, but the planner does not currently have reliable runtime duration inputs. +Adding that feedback loop is a separate product and operational decision. + +We considered changing the runtime model to support per-parent dependencies inside a parallel block. +That would address non-series-parallel representation limits. +It would also be a larger executor contract change, so it is outside this decision. + +## Acknowledged Residuals + +Gap A is the non-series-parallel representation limit. +Some DAGs cannot be represented as a `Sequence` and `Parallel` tree that matches the critical path for every duration vector. +The scheduler falls back to the level tree when SP does not dominate, but that is still a representation-level compromise. + +Gap B is incomparable trees under skew. +When SP and Level each win under different duration assignments, the hybrid picks Level. +That preserves the no-regression guarantee but can leave a runtime win on the table. +Closing this gap requires runtime statistics or a richer execution model. + +## Worked Examples + +Each example below shows the same five pieces: +the dependency graph drawn in ASCII, +the two candidate trees produced by `scheduleSP` and `scheduleLevel`, +the root-to-leaf path sets used by the dominance check, +the dominance verdict and which tree the hybrid returns, +and a numerical makespan comparison under at least one skewed duration vector that explains why the choice matters at runtime. + +### Example 1 — the baseline (independent components) + +Federation pattern: a query selects `me { firstName lastName currentPractice { id } }` plus `organisations(...)`. +`me` lives in the user subgraph. +`currentPractice` is an entity hop into a third subgraph that depends on `me.id`. +`organisations` is in a fourth subgraph that has nothing to do with the user. + +DAG: + +``` +user ───▶ practice +organisation (no edges) +``` + +Candidate trees: + +``` +scheduleSP : Parallel(Sequence(user, practice), organisation) +scheduleLevel : Parallel(Sequence(user, practice), organisation) +``` + +Both schedulers see two weakly connected components — `{user, practice}` and `{organisation}` — and emit a `Parallel` over the components. +The trees are identical. + +Path sets: + +``` +SP paths : {user, practice}, {organisation} +Level paths : {user, practice}, {organisation} +``` + +Every SP path is contained in some Level path (each path is contained in itself). +`dominates(SP, Level)` returns `true`. +The hybrid returns `scheduleSP`, which is the same shape as `scheduleLevel` here. + +Makespan with `user = 100ms`, `practice = 50ms`, `organisation = 200ms`: + +``` +runtime makespan = max(100 + 50, 200) = 200ms +old plan makespan = max(100, 200) + 50 = 250ms +``` + +Why this matters: the legacy postprocessor would have emitted `Sequence(Parallel(user, organisation), practice)` and forced `practice` to wait for `organisation` even though `practice` does not need `organisation`'s data. +The new scheduler reads the dependency graph instead of the topological levels and lets the unrelated branches run end-to-end in parallel. + +### Example 2 — two chains joining at a third subgraph (SP dominates) + +Federation pattern: two independent entity chains feed a third subgraph. +For instance, `User → Address` and `Order → Discount` both contribute fields that a downstream `Recommendation` resolver needs. + +DAG: + +``` +A ────▶ C ─┐ + ├─▶ E +B ────▶ D ─┘ +``` + +Edges: `A→C`, `B→D`, `C→E`, `D→E`. + +Candidate trees: + +``` +scheduleSP : Sequence( + Parallel( + Sequence(A, C), + Sequence(B, D) + ), + E + ) + +scheduleLevel : Sequence( + Parallel(A, B), + Parallel(C, D), + E + ) +``` + +`scheduleSP` recognises that `A → C` and `B → D` are independent chains that meet at `E`. +It runs each chain end-to-end in its own parallel branch. +`scheduleLevel` strips the topological layers — first `{A, B}`, then `{C, D}`, then `E` — and synchronises every layer. + +Path sets: + +``` +SP paths : {A, C, E}, {B, D, E} +Level paths : {A, C, E}, {A, D, E}, {B, C, E}, {B, D, E} +``` + +Every SP path appears in the Level set as well. +`dominates(SP, Level)` returns `true`. +The hybrid returns the SP tree. + +Why dominance is the right rule here: under uniform durations both trees finish in the same time (three "ticks"), but their behaviour diverges as soon as the chain durations become skewed. +With `A = 1ms`, `B = 100ms`, `C = 100ms`, `D = 1ms`, and `E = 10ms`: + +``` +SP makespan : max(A+C, B+D) + E = max(101, 101) + 10 = 111ms +Level makespan : max(A, B) + max(C, D) + E = 100 + 100 + 10 = 210ms +``` + +The SP tree carries the long `A → C` arm and the long `B → D` arm in parallel. +The Level tree synchronises at every layer and ends up paying the maximum of each layer twice — once for the slow root and once for the slow middle node. +The dominance check guaranteed at plan time that this was safe, so the hybrid picked the win. + +### Example 3 — asymmetric chain merge with a leaf side-branch (Level wins, SP regresses) + +Federation pattern: a root entity `A` has a chain branch `A → B → D` and a sibling branch `A → C` where `C` both feeds the join `D` and produces an extra leaf `E`. + +DAG: + +``` + ┌───▶ B ───┐ + │ ├─▶ D +A ──────┼───▶ C ───┘ + │ + └───▶ C ─▶ E (same C, drawn twice for clarity) +``` + +Edges: `A→B`, `A→C`, `B→D`, `C→D`, `C→E`. + +Candidate trees: + +``` +scheduleSP : Sequence( + A, + Parallel( + B, + Sequence(C, E) + ), + D + ) + +scheduleLevel : Sequence( + A, + Parallel(B, C), + Parallel(D, E) + ) +``` + +`scheduleSP` eagerly inlines `C → E` into `C`'s branch. +That binds `E` to the parallel block, which means `D` cannot start until both `B` and `C → E` have finished. +`scheduleLevel` keeps `D` and `E` at the same topological level after `{B, C}` and lets them run in parallel. + +Path sets: + +``` +SP paths : {A, B, D}, {A, C, E, D} +Level paths : {A, B, D}, {A, B, E}, {A, C, D}, {A, C, E} +``` + +The SP path `{A, C, E, D}` has four nodes. +Every Level path has exactly three. +There is no Level path that contains all four SP nodes. +`dominates(SP, Level)` returns `false`. + +The dominance-only driver therefore returns `scheduleLevel`. + +Makespan under uniform durations, weight `1` for every node: + +``` +SP makespan : 1 + max(1, 2) + 1 = 4 ticks +Level makespan : 1 + max(1, 1) + max(1, 1) = 3 ticks +``` + +The hybrid avoided a one-tick regression on uniform durations. +Under skewed durations the SP tree can lose by even more, because the chain `C → E` keeps growing while `D` sits idle behind it. + +This is the canonical motivation for the dominance check: `scheduleSP` is not always better. +Picking it on a "uniform-tie" or "smaller uniform makespan" rule would silently regress this class of plans. +The path-set test catches the regression at plan time. + +### Example 4 — independent leaf alongside a chain (component-awareness) + +Federation pattern: a root `A` produces three children — two of them (`B` and `C`) form a chain that joins inside `C`, the third (`D`) is an unrelated leaf hanging off `A`. + +DAG: + +``` + ┌──▶ B ──▶ C +A ──────┼──▶ C (multi-parent join via B) + │ + └──▶ D +``` + +Edges: `A→B`, `A→C`, `A→D`, `B→C`. + +Candidate trees: + +``` +scheduleSP : Sequence( + A, + Parallel(B, D), + C + ) + +scheduleLevel : Sequence( + A, + Parallel( + Sequence(B, C), + D + ) + ) +``` + +`scheduleSP` batches everything ready after `A` into a single parallel block (`B` and `D`), then runs `C` once `B` is done. +`scheduleLevel` notices that after stripping `A`, the rest of the DAG splits into two components — `{B, C}` and `{D}` — and runs each as its own parallel arm. + +Path sets: + +``` +SP paths : {A, B, C}, {A, D, C} +Level paths : {A, B, C}, {A, D} +``` + +The SP path `{A, D, C}` is not contained in any Level path. +Level has `{A, D}` (no `C`) and `{A, B, C}` (no `D`). +`dominates(SP, Level)` returns `false`. + +The hybrid returns `scheduleLevel`. + +Why this matters under skew: with `A = 1ms`, `B = 1ms`, `C = 100ms`, `D = 50ms`, + +``` +SP makespan : 1 + max(1, 50) + 100 = 151ms +Level makespan : 1 + max(1 + 100, 50) = 102ms +``` + +The SP tree drags `C` behind whichever parallel sibling is slowest, so a slow but unrelated `D` blocks the join. +The Level tree sees that `D` does not feed `C` and lets the `B → C` chain run end-to-end alongside `D`. + +This case was discovered during plan review. +It is the reason `scheduleSP` does its own weakly-connected-component pass at every recursion level — without that step, `scheduleSP` would have produced `Sequence(Parallel(A, B, D), C)` and lost even more time. +With the component pass, `scheduleSP` produces the shape above, which still loses to `scheduleLevel`, and the dominance check correctly falls back. + +### Example 5 — incomparable trees (residual gap) + +Federation pattern: a non-series-parallel DAG. +There is no `Sequence`/`Parallel` tree that is universally optimal; the best choice depends on the actual durations. + +DAG: + +``` + ┌──▶ D +A ──────┤ + └──▶ E ◀──┐ + │ +B ─▶ C ─▶ D │ (E has parents A and B; D has parents A and C) + └────▶ E ─┘ +``` + +Edges: `A→D`, `A→E`, `B→C`, `B→E`, `C→D`. + +Candidate trees: + +``` +scheduleSP : Sequence( + Parallel(A, Sequence(B, C)), + Parallel(D, E) + ) + +scheduleLevel : Sequence( + Parallel(A, B), + Parallel( + Sequence(C, D), + E + ) + ) +``` + +Path sets: + +``` +SP paths : {A, D}, {A, E}, {B, C, D}, {B, C, E} +Level paths : {A, C, D}, {A, E}, {B, C, D}, {B, E} +``` + +SP dominance check: +the SP path `{B, C, E}` is not contained in any Level path (Level has `{B, C, D}` without `E` and `{B, E}` without `C`). +`dominates(SP, Level)` returns `false`. + +Level dominance is also false (the dominance-only driver does not check this, but for the analysis the symmetric case fails too: +Level path `{A, C, D}` is not contained in any SP path). + +Neither tree dominates. +The hybrid returns `scheduleLevel`. + +Why "incomparable" is real and not an artefact of the algorithm: +under one duration vector SP is faster, under another Level is faster. +Concrete example with `A=25, B=227, C=647, D=3, E=5`: + +``` +SP makespan : max(A, B+C) + max(D, E) = max(25, 874) + max(3, 5) = 874 + 5 = 879ms +Level makespan : max(A, B) + max(C+D, E) = max(25, 227) + max(650, 5) = 227 + 650 = 877ms +``` + +Level is slightly better. + +But under `A=100, B=1, C=1, D=1, E=1`: + +``` +SP makespan : max(100, 2) + max(1, 1) = 100 + 1 = 101ms +Level makespan : max(100, 1) + max(1+1, 1) = 100 + 2 = 102ms +``` + +SP is slightly better. + +The dominance theorem proves that no single SP tree on this DAG is universally optimal, so at plan time we have no way to choose correctly without runtime statistics. +The hybrid takes the conservative option — `scheduleLevel`, the same shape today's algorithm would produce — and accepts that some duration distributions leave a small win on the table. +This is the "incomparable trees" residual documented in Gap B above. + +### Reading the decision table + +The four examples cover the four interesting outcomes: + +``` +Example 1 : trees identical → hybrid returns SP (same shape) +Example 2 : SP path-set-dominates Level → hybrid returns SP (skew win) +Example 3 : SP does not dominate, Level wins on uniform too → hybrid returns Level (regression caught) +Example 4 : SP does not dominate, regression only under skew → hybrid returns Level (regression caught) +Example 5 : incomparable, neither tree universally better → hybrid returns Level (conservative) +``` + +Production federation queries are dominated by Example 1 and Example 2 patterns — a single chain plus an unrelated subgraph, or two chains joining at a downstream entity. +Examples 3, 4, and 5 are uncommon but real, and the dominance check guarantees the hybrid never regresses against the legacy level-based shape on any of them. + +## References + +Valdes, Tarjan, and Lawler, 1982, The Recognition of Series Parallel Digraphs. +GraphQL federation query planning shapes: Sequence, Parallel, Fetch, Flatten, and entity fetch dependencies. +Local implementation plan: `docs/superpowers/specs/2026-05-08-improve-query-order-plan.md`. diff --git a/v2/pkg/engine/datasource/graphql_datasource/query_order_baseline_test.go b/v2/pkg/engine/datasource/graphql_datasource/query_order_baseline_test.go new file mode 100644 index 0000000000..aa7f4d13b1 --- /dev/null +++ b/v2/pkg/engine/datasource/graphql_datasource/query_order_baseline_test.go @@ -0,0 +1,345 @@ +package graphql_datasource + +import ( + "strings" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/wundergraph/graphql-go-tools/v2/pkg/astnormalization" + "github.com/wundergraph/graphql-go-tools/v2/pkg/asttransform" + "github.com/wundergraph/graphql-go-tools/v2/pkg/astvalidation" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/plan" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/postprocess" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" + "github.com/wundergraph/graphql-go-tools/v2/pkg/internal/unsafeparser" + "github.com/wundergraph/graphql-go-tools/v2/pkg/operationreport" +) + +func TestQueryOrderBaseline(t *testing.T) { + definition := ` + type Query { + me: User + organisations(ids: [ID!]!): [Organisation!]! + } + + type User { + id: ID! + firstName: String! + lastName: String! + currentPractice: Practice + } + + type Organisation { + id: ID! + name: String! + shortCode: String! + } + + type Practice { + id: ID! + } + ` + + operation := ` + query Baseline($a: [ID!]!) { + me { + firstName + lastName + currentPractice { + id + } + } + organisations(ids: $a) { + name + shortCode + id + } + } + ` + + userSubgraphSDL := ` + type Query { + me: User + } + + type User @key(fields: "id") { + id: ID! + firstName: String! + lastName: String! + } + ` + + organisationSubgraphSDL := ` + type Query { + organisations(ids: [ID!]!): [Organisation!]! + } + + type Organisation { + id: ID! + name: String! + shortCode: String! + } + ` + + practiceSubgraphSDL := ` + type User @key(fields: "id") { + id: ID! @external + currentPractice: Practice + } + + type Practice { + id: ID! + } + ` + + config := plan.Configuration{ + DataSources: []plan.DataSource{ + mustDataSourceConfiguration( + t, + "user-subgraph", + &plan.DataSourceMetadata{ + RootNodes: []plan.TypeField{ + { + TypeName: "Query", + FieldNames: []string{"me"}, + }, + { + TypeName: "User", + FieldNames: []string{"id", "firstName", "lastName"}, + }, + }, + FederationMetaData: plan.FederationMetaData{ + Keys: plan.FederationFieldConfigurations{ + { + TypeName: "User", + SelectionSet: "id", + }, + }, + }, + }, + mustCustomConfiguration(t, ConfigurationInput{ + Fetch: &FetchConfiguration{ + URL: "http://user-subgraph", + }, + SchemaConfiguration: mustSchema(t, + &FederationConfiguration{ + Enabled: true, + ServiceSDL: userSubgraphSDL, + }, + userSubgraphSDL, + ), + }), + ), + mustDataSourceConfiguration( + t, + "organisation-subgraph", + &plan.DataSourceMetadata{ + RootNodes: []plan.TypeField{ + { + TypeName: "Query", + FieldNames: []string{"organisations"}, + }, + }, + ChildNodes: []plan.TypeField{ + { + TypeName: "Organisation", + FieldNames: []string{"id", "name", "shortCode"}, + }, + }, + }, + mustCustomConfiguration(t, ConfigurationInput{ + Fetch: &FetchConfiguration{ + URL: "http://organisation-subgraph", + }, + SchemaConfiguration: mustSchema(t, + &FederationConfiguration{ + Enabled: true, + ServiceSDL: organisationSubgraphSDL, + }, + organisationSubgraphSDL, + ), + }), + ), + mustDataSourceConfiguration( + t, + "practice-subgraph", + &plan.DataSourceMetadata{ + RootNodes: []plan.TypeField{ + { + TypeName: "User", + FieldNames: []string{"id", "currentPractice"}, + }, + }, + ChildNodes: []plan.TypeField{ + { + TypeName: "Practice", + FieldNames: []string{"id"}, + }, + }, + FederationMetaData: plan.FederationMetaData{ + Keys: plan.FederationFieldConfigurations{ + { + TypeName: "User", + SelectionSet: "id", + }, + }, + }, + }, + mustCustomConfiguration(t, ConfigurationInput{ + Fetch: &FetchConfiguration{ + URL: "http://practice-subgraph", + }, + SchemaConfiguration: mustSchema(t, + &FederationConfiguration{ + Enabled: true, + ServiceSDL: practiceSubgraphSDL, + }, + practiceSubgraphSDL, + ), + }), + ), + }, + DisableResolveFieldPositions: true, + Fields: plan.FieldConfigurations{ + { + TypeName: "Query", + FieldName: "organisations", + Arguments: plan.ArgumentsConfigurations{ + { + Name: "ids", + SourceType: plan.FieldArgumentSource, + }, + }, + }, + }, + Debug: plan.DebugConfiguration{}, + } + + def := unsafeparser.ParseGraphqlDocumentString(definition) + op := unsafeparser.ParseGraphqlDocumentString(operation) + + err := asttransform.MergeDefinitionWithBaseSchema(&def) + require.NoError(t, err) + + norm := astnormalization.NewWithOpts( + astnormalization.WithExtractVariables(), + astnormalization.WithInlineFragmentSpreads(), + astnormalization.WithRemoveFragmentDefinitions(), + astnormalization.WithRemoveUnusedVariables(), + ) + + var report operationreport.Report + norm.NormalizeOperation(&op, &def, &report) + require.False(t, report.HasErrors(), report.Error()) + + valid := astvalidation.DefaultOperationValidator() + valid.Validate(&op, &def, &report) + require.False(t, report.HasErrors(), report.Error()) + + p, err := plan.NewPlanner(config) + require.NoError(t, err) + + actualPlan := p.Plan(&op, &def, "Baseline", &report, plan.IncludeQueryPlanInResponse()) + require.False(t, report.HasErrors(), report.Error()) + postprocess.NewProcessor().Process(actualPlan) + + responsePlan, ok := actualPlan.(*plan.SynchronousResponsePlan) + require.True(t, ok) + require.NotNil(t, responsePlan.Response) + require.NotNil(t, responsePlan.Response.Fetches) + + queryPlan := responsePlan.Response.Fetches.QueryPlan() + require.NotNil(t, queryPlan) + + assertQueryOrderBaselineShape(t, queryPlan) + + actual := queryPlan.PrettyPrint() + require.Equal(t, strings.TrimSpace(expectedQueryOrderBaselinePlan), strings.TrimSpace(actual)) +} + +func assertQueryOrderBaselineShape(t *testing.T, queryPlan *resolve.FetchTreeQueryPlanNode) { + t.Helper() + + require.Equal(t, resolve.FetchTreeNodeKindSequence, queryPlan.Kind) + require.Len(t, queryPlan.Children, 2) + + parallel := queryPlan.Children[0] + require.Equal(t, resolve.FetchTreeNodeKindParallel, parallel.Kind) + require.Len(t, parallel.Children, 2) + + fetchByService := make(map[string]*resolve.FetchTreeQueryPlan, len(parallel.Children)) + for _, child := range parallel.Children { + require.Equal(t, resolve.FetchTreeNodeKindSingle, child.Kind) + require.NotNil(t, child.Fetch) + require.Equal(t, "Single", child.Fetch.Kind) + fetchByService[child.Fetch.SubgraphName] = child.Fetch + } + require.Contains(t, fetchByService, "user-subgraph") + require.Contains(t, fetchByService, "organisation-subgraph") + + practice := queryPlan.Children[1] + require.Equal(t, resolve.FetchTreeNodeKindSingle, practice.Kind) + require.NotNil(t, practice.Fetch) + require.Equal(t, "Entity", practice.Fetch.Kind) + require.Equal(t, "practice-subgraph", practice.Fetch.SubgraphName) + + require.Equal(t, + []int{fetchByService["user-subgraph"].FetchID}, + practice.Fetch.DependsOnFetchIDs, + "practice entity fetch must depend only on the user-subgraph fetch", + ) + + require.Len(t, practice.Fetch.Representations, 1) + rep := practice.Fetch.Representations[0] + require.Equal(t, resolve.RepresentationKindKey, rep.Kind) + require.Equal(t, "User", rep.TypeName) + require.Contains(t, rep.Fragment, "__typename") + require.Contains(t, rep.Fragment, "id") +} + +const expectedQueryOrderBaselinePlan = ` +QueryPlan { + Sequence { + Parallel { + Fetch(service: "user-subgraph") { + { + me { + firstName + lastName + __typename + id + } + } + } + Fetch(service: "organisation-subgraph") { + { + organisations(ids: $a){ + name + shortCode + id + } + } + } + } + Fetch(service: "practice-subgraph") { + { + fragment Key on User { + __typename + id + } + } => + { + _entities(representations: $representations){ + ... on User { + __typename + currentPractice { + id + } + } + } + } + } + } +} +` diff --git a/v2/pkg/engine/postprocess/build_schedule_tree.go b/v2/pkg/engine/postprocess/build_schedule_tree.go new file mode 100644 index 0000000000..83e1a0de5f --- /dev/null +++ b/v2/pkg/engine/postprocess/build_schedule_tree.go @@ -0,0 +1,797 @@ +package postprocess + +import ( + "fmt" + "slices" + "strings" + + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" +) + +type buildScheduleTreeProcessor struct { + disable bool +} + +func (b *buildScheduleTreeProcessor) ProcessFetchTree(root *resolve.FetchTreeNode) { + if b.disable || root == nil || root.Kind != resolve.FetchTreeNodeKindSequence { + return + } + // DEVIATION from upstream (improve-query-order): upstream panics on any + // scheduler/validator error, which turns a plan-time edge case into a 500 for + // every affected operation. Fall back to the legacy wave pipeline instead — + // the same processors that would have run with the flag off — so a failed + // schedule degrades to today's shipping behavior, never a panic. + if err := b.buildSchedule(root); err != nil { + (&orderSequenceByDependencies{}).ProcessFetchTree(root) + (&createParallelNodes{}).ProcessFetchTree(root) + } +} + +func (b *buildScheduleTreeProcessor) buildSchedule(root *resolve.FetchTreeNode) error { + dag, err := newFetchDAG(root.ChildNodes) + if err != nil { + return err + } + tree, err := buildScheduleTree(root.ChildNodes, dag) + if err != nil { + return err + } + if err := validateSchedule(tree, dag); err != nil { + return err + } + if tree == nil { + root.ChildNodes = nil + return nil + } + *root = *tree + return nil +} + +type fetchDAG struct { + nodes map[int]*resolve.FetchTreeNode + parents map[int]map[int]struct{} + children map[int]map[int]struct{} +} + +func newFetchDAG(nodes []*resolve.FetchTreeNode) (*fetchDAG, error) { + dag := &fetchDAG{ + nodes: make(map[int]*resolve.FetchTreeNode, len(nodes)), + parents: make(map[int]map[int]struct{}, len(nodes)), + children: make(map[int]map[int]struct{}, len(nodes)), + } + for _, node := range nodes { + if node == nil || node.Item == nil || node.Item.Fetch == nil { + continue + } + id := node.Item.Fetch.Dependencies().FetchID + if _, exists := dag.nodes[id]; exists { + return nil, fmt.Errorf("duplicate fetch id %d", id) + } + dag.nodes[id] = node + dag.parents[id] = make(map[int]struct{}) + dag.children[id] = make(map[int]struct{}) + } + for id, node := range dag.nodes { + for _, dep := range node.Item.Fetch.Dependencies().DependsOnFetchIDs { + if _, exists := dag.nodes[dep]; !exists { + continue + } + dag.parents[id][dep] = struct{}{} + dag.children[dep][id] = struct{}{} + } + } + return dag, nil +} + +func (d *fetchDAG) sortedIDs() []int { + ids := make([]int, 0, len(d.nodes)) + for id := range d.nodes { + ids = append(ids, id) + } + slices.Sort(ids) + return ids +} + +func (d *fetchDAG) hasCycle() bool { + _, err := scheduleLevel(d.sortedIDs(), d) + return err != nil +} + +func buildScheduleTree(roots []*resolve.FetchTreeNode, dag *fetchDAG) (*resolve.FetchTreeNode, error) { + ids := make([]int, 0, len(roots)) + for _, root := range roots { + if root == nil || root.Item == nil || root.Item.Fetch == nil { + continue + } + ids = append(ids, root.Item.Fetch.Dependencies().FetchID) + } + slices.Sort(ids) + treeLevel, err := scheduleLevel(ids, dag) + if err != nil { + return nil, err + } + treeSP, err := scheduleSP(ids, dag) + if err != nil { + return nil, err + } + if err := validateSchedule(treeLevel, dag); err != nil { + return nil, err + } + if err := validateSchedule(treeSP, dag); err != nil { + return nil, err + } + if dominates(treeSP, treeLevel) { + return treeSP, nil + } + return treeLevel, nil +} + +func scheduleLevel(nodes []int, dag *fetchDAG) (*resolve.FetchTreeNode, error) { + nodes = sortedUnique(nodes) + if len(nodes) == 0 { + return nil, nil + } + if len(nodes) == 1 { + return dag.nodes[nodes[0]], nil + } + components := weaklyConnectedComponents(nodes, dag) + if len(components) > 1 { + branches := make([]*resolve.FetchTreeNode, 0, len(components)) + for _, component := range components { + child, err := scheduleLevel(component, dag) + if err != nil { + return nil, err + } + if child != nil { + branches = append(branches, child) + } + } + sortBranches(branches) + return parallelOf(branches...), nil + } + allowed := idSet(nodes) + roots := make([]int, 0, len(nodes)) + for _, id := range nodes { + hasParent := false + for parent := range dag.parents[id] { + if _, ok := allowed[parent]; ok { + hasParent = true + break + } + } + if !hasParent { + roots = append(roots, id) + } + } + if len(roots) == 0 { + return nil, fmt.Errorf("cycle detected in fetch dependency graph") + } + rootNodes := make([]*resolve.FetchTreeNode, 0, len(roots)) + rootSet := idSet(roots) + rest := make([]int, 0, len(nodes)-len(roots)) + for _, id := range roots { + rootNodes = append(rootNodes, dag.nodes[id]) + } + for _, id := range nodes { + if _, ok := rootSet[id]; !ok { + rest = append(rest, id) + } + } + rootTree := parallelOf(rootNodes...) + restTree, err := scheduleLevel(rest, dag) + if err != nil { + return nil, err + } + if restTree == nil { + return rootTree, nil + } + return sequenceOf(rootTree, restTree), nil +} + +func scheduleSP(nodes []int, dag *fetchDAG) (*resolve.FetchTreeNode, error) { + nodes = sortedUnique(nodes) + if len(nodes) == 0 { + return nil, nil + } + if len(nodes) == 1 { + return dag.nodes[nodes[0]], nil + } + components := weaklyConnectedComponents(nodes, dag) + if len(components) > 1 { + branches := make([]*resolve.FetchTreeNode, 0, len(components)) + for _, component := range components { + child, err := scheduleSP(component, dag) + if err != nil { + return nil, err + } + if child != nil { + branches = append(branches, child) + } + } + sortBranches(branches) + return parallelOf(branches...), nil + } + return scheduleSPInline(nodes, dag) +} + +type processingState struct { + ready []int + unhandled map[int]map[int]struct{} +} + +func scheduleSPInline(nodes []int, dag *fetchDAG) (*resolve.FetchTreeNode, error) { + allowed := idSet(nodes) + roots := make([]int, 0, len(nodes)) + for _, id := range nodes { + hasParent := false + for parent := range dag.parents[id] { + if _, ok := allowed[parent]; ok { + hasParent = true + break + } + } + if !hasParent { + roots = append(roots, id) + } + } + if len(roots) == 0 { + return nil, fmt.Errorf("cycle detected in fetch dependency graph") + } + state := processingState{ + ready: roots, + unhandled: map[int]map[int]struct{}{}, + } + sequence := make([]*resolve.FetchTreeNode, 0, len(nodes)) + parallelFirst := len(roots) > 1 + processed := make(map[int]struct{}, len(nodes)) + for len(state.ready) != 0 { + batch, nextState, err := processBatch(state, parallelFirst, dag, allowed, processed, map[int]struct{}{}) + if err != nil { + return nil, err + } + if batch != nil { + sequence = append(sequence, batch) + } + state = nextState + parallelFirst = true + } + if len(processed) != len(nodes) || len(state.unhandled) != 0 { + return nil, fmt.Errorf("cycle detected in fetch dependency graph") + } + return sequenceOf(sequence...), nil +} + +func processBatch(state processingState, parallelFirst bool, dag *fetchDAG, allowed map[int]struct{}, processed map[int]struct{}, branchDone map[int]struct{}) (*resolve.FetchTreeNode, processingState, error) { + ready := sortedUnique(state.ready) + branches := make([]*resolve.FetchTreeNode, 0, len(ready)) + merged := processingState{ + unhandled: clonePending(state.unhandled), + } + for _, id := range ready { + if _, ok := processed[id]; ok { + continue + } + subtree, subState, err := processNode(id, dag, allowed, processed, cloneSet(branchDone)) + if err != nil { + return nil, processingState{}, err + } + if subtree != nil { + branches = append(branches, subtree) + } + merged = mergeStates(merged, subState, dag) + for processedID := range processed { + delete(merged.unhandled, processedID) + } + } + processedReady := idSet(ready) + merged.ready = removeIDs(sortedUnique(merged.ready), processedReady, processed) + if len(branches) == 0 { + return nil, merged, nil + } + sortBranches(branches) + if len(branches) == 1 { + return branches[0], merged, nil + } + if parallelFirst { + return parallelOf(branches...), merged, nil + } + return sequenceOf(branches...), merged, nil +} + +func processNode(id int, dag *fetchDAG, allowed map[int]struct{}, processed map[int]struct{}, branchDone map[int]struct{}) (*resolve.FetchTreeNode, processingState, error) { + if _, ok := processed[id]; ok { + return nil, processingState{unhandled: map[int]map[int]struct{}{}}, nil + } + processed[id] = struct{}{} + branchDone[id] = struct{}{} + childState := processingState{ + unhandled: map[int]map[int]struct{}{}, + } + for _, child := range sortedSet(dag.children[id]) { + if _, ok := allowed[child]; !ok { + continue + } + parents := filterSet(dag.parents[child], allowed) + if len(parents) == 1 { + childState.ready = append(childState.ready, child) + continue + } + pending := cloneSet(parents) + for done := range branchDone { + delete(pending, done) + } + if len(pending) == 0 { + childState.ready = append(childState.ready, child) + continue + } + childState.unhandled[child] = pending + } + if len(childState.ready) == 0 { + return dag.nodes[id], childState, nil + } + sequence := []*resolve.FetchTreeNode{dag.nodes[id]} + state := childState + for len(state.ready) != 0 { + batch, nextState, err := processBatch(state, true, dag, allowed, processed, branchDone) + if err != nil { + return nil, processingState{}, err + } + if batch != nil { + sequence = append(sequence, batch) + } + state = nextState + } + return sequenceOf(sequence...), state, nil +} + +func mergeStates(a, b processingState, dag *fetchDAG) processingState { + merged := processingState{ + ready: append([]int{}, a.ready...), + unhandled: map[int]map[int]struct{}{}, + } + merged.ready = append(merged.ready, b.ready...) + keys := make(map[int]struct{}, len(a.unhandled)+len(b.unhandled)) + for key := range a.unhandled { + keys[key] = struct{}{} + } + for key := range b.unhandled { + keys[key] = struct{}{} + } + for key := range keys { + left, leftOK := a.unhandled[key] + if !leftOK { + left = dag.parents[key] + } + right, rightOK := b.unhandled[key] + if !rightOK { + right = dag.parents[key] + } + pending := intersectSets(left, right) + if len(pending) == 0 { + merged.ready = append(merged.ready, key) + continue + } + merged.unhandled[key] = pending + } + merged.ready = sortedUnique(merged.ready) + for _, id := range merged.ready { + delete(merged.unhandled, id) + } + return merged +} + +func dominates(treeA, treeB *resolve.FetchTreeNode) bool { + pathsA := enumeratePaths(treeA) + pathsB := enumeratePaths(treeB) + for _, pathA := range pathsA { + found := false + for _, pathB := range pathsB { + if setContainsAll(pathB, pathA) { + found = true + break + } + } + if !found { + return false + } + } + return true +} + +func enumeratePaths(node *resolve.FetchTreeNode) []map[int]struct{} { + if node == nil { + return []map[int]struct{}{{}} + } + switch node.Kind { + case resolve.FetchTreeNodeKindSingle: + return []map[int]struct{}{{ + node.Item.Fetch.Dependencies().FetchID: {}, + }} + case resolve.FetchTreeNodeKindParallel: + var paths []map[int]struct{} + for _, child := range node.ChildNodes { + paths = append(paths, enumeratePaths(child)...) + } + return paths + case resolve.FetchTreeNodeKindSequence: + paths := []map[int]struct{}{{}} + for _, child := range node.ChildNodes { + childPaths := enumeratePaths(child) + next := make([]map[int]struct{}, 0, len(paths)*len(childPaths)) + for _, existing := range paths { + for _, childPath := range childPaths { + combined := cloneSet(existing) + for id := range childPath { + combined[id] = struct{}{} + } + next = append(next, combined) + } + } + paths = next + } + return paths + default: + return nil + } +} + +func validateSchedule(root *resolve.FetchTreeNode, dag *fetchDAG) error { + _, err := validateScheduleNode(root, dag, map[int]struct{}{}) + return err +} + +func validateScheduleNode(node *resolve.FetchTreeNode, dag *fetchDAG, before map[int]struct{}) ([]*resolve.FetchTreeNode, error) { + if node == nil { + return nil, nil + } + switch node.Kind { + case resolve.FetchTreeNodeKindSingle: + id := node.Item.Fetch.Dependencies().FetchID + for _, dep := range node.Item.Fetch.Dependencies().DependsOnFetchIDs { + if _, known := dag.nodes[dep]; !known { + continue + } + if _, ok := before[dep]; !ok { + return nil, fmt.Errorf("fetch %d depends on fetch %d before it is available", id, dep) + } + } + return []*resolve.FetchTreeNode{node}, nil + case resolve.FetchTreeNodeKindSequence: + available := cloneSet(before) + leaves := make([]*resolve.FetchTreeNode, 0) + for _, child := range node.ChildNodes { + childLeaves, err := validateScheduleNode(child, dag, available) + if err != nil { + return nil, err + } + for _, leaf := range childLeaves { + available[leaf.Item.Fetch.Dependencies().FetchID] = struct{}{} + leaves = append(leaves, leaf) + } + } + return leaves, nil + case resolve.FetchTreeNodeKindParallel: + branchLeaves := make([][]*resolve.FetchTreeNode, len(node.ChildNodes)) + leaves := make([]*resolve.FetchTreeNode, 0) + for i, child := range node.ChildNodes { + childLeaves, err := validateScheduleNode(child, dag, before) + if err != nil { + return nil, err + } + branchLeaves[i] = childLeaves + leaves = append(leaves, childLeaves...) + } + for i := range branchLeaves { + for j := i + 1; j < len(branchLeaves); j++ { + for _, a := range branchLeaves[i] { + for _, b := range branchLeaves[j] { + if err := validateParallelLeaves(a, b); err != nil { + return nil, err + } + } + } + } + } + return leaves, nil + default: + return nil, nil + } +} + +// validateParallelLeaves enforces that two leaves scheduled in parallel do not +// depend on each other via FetchID edges. +// +// DEVIATION from upstream (improve-query-order): upstream additionally rejected +// (a) a leaf whose response path is a strict-prefix extension of a sibling's +// provided path and (b) two siblings providing the same path. Both heuristics +// FALSE-POSITIVE on real cosmo planner output, where the engine's actual +// correctness contract is complete DependsOnFetchIDs — the exact invariant +// createParallelNodes (legacy waves) and resolveParallel already rely on. +// Canonical counter-example (the hive gateways-benchmark plan): fetch 3 +// (accounts, users.@.reviews.@.product.reviews.@.author, deps [1]) attaches to +// review/author objects provided ENTIRELY by fetch 1; sibling fetch 2 (products, +// users.@.reviews.@.product, deps [1]) merges product fields into a PREFIX of +// 3's path without providing 3's attach points. The planner proves the contract +// in the same plan: fetch 9 (inventory, users.@.reviews.@.product, deps [1, 2]) +// DOES list fetch 2 because it actually reads fetch-2-provided data. Same-path +// siblings (heuristic b) are ordinary federation: multiple subgraphs +// contributing disjoint fields to one entity set (fetches 6 and 11 at +// topProducts). The FetchID edges are authoritative; the path heuristics are +// strictly weaker and are dropped. +func validateParallelLeaves(a, b *resolve.FetchTreeNode) error { + aID := a.Item.Fetch.Dependencies().FetchID + bID := b.Item.Fetch.Dependencies().FetchID + if slices.Contains(a.Item.Fetch.Dependencies().DependsOnFetchIDs, bID) { + return fmt.Errorf("fetch %d depends on parallel fetch %d", aID, bID) + } + if slices.Contains(b.Item.Fetch.Dependencies().DependsOnFetchIDs, aID) { + return fmt.Errorf("fetch %d depends on parallel fetch %d", bID, aID) + } + return nil +} + +func providedPath(node *resolve.FetchTreeNode) []string { + base := responsePath(node) + var merge []string + switch fetch := node.Item.Fetch.(type) { + case *resolve.SingleFetch: + merge = fetch.PostProcessing.MergePath + case *resolve.EntityFetch: + merge = fetch.PostProcessing.MergePath + case *resolve.BatchEntityFetch: + merge = fetch.PostProcessing.MergePath + } + out := make([]string, 0, len(base)+len(merge)) + out = append(out, base...) + out = append(out, merge...) + return out +} + +func responsePath(node *resolve.FetchTreeNode) []string { + if len(node.Item.ResponsePathElements) != 0 { + return append([]string{}, node.Item.ResponsePathElements...) + } + if node.Item.ResponsePath == "" { + return nil + } + return strings.Split(node.Item.ResponsePath, ".") +} + +func pathStrictPrefix(prefix, path []string) bool { + if len(prefix) >= len(path) { + return false + } + for i := range prefix { + if prefix[i] != path[i] { + return false + } + } + return true +} + +func weaklyConnectedComponents(nodes []int, dag *fetchDAG) [][]int { + allowed := idSet(nodes) + seen := map[int]struct{}{} + components := make([][]int, 0) + for _, id := range nodes { + if _, ok := seen[id]; ok { + continue + } + queue := []int{id} + seen[id] = struct{}{} + component := make([]int, 0) + for len(queue) != 0 { + current := queue[0] + queue = queue[1:] + component = append(component, current) + for neighbor := range dag.parents[current] { + if _, ok := allowed[neighbor]; !ok { + continue + } + if _, ok := seen[neighbor]; ok { + continue + } + seen[neighbor] = struct{}{} + queue = append(queue, neighbor) + } + for neighbor := range dag.children[current] { + if _, ok := allowed[neighbor]; !ok { + continue + } + if _, ok := seen[neighbor]; ok { + continue + } + seen[neighbor] = struct{}{} + queue = append(queue, neighbor) + } + } + slices.Sort(component) + components = append(components, component) + } + slices.SortFunc(components, func(a, b []int) int { + return a[0] - b[0] + }) + return components +} + +func uniformMakespan(node *resolve.FetchTreeNode) int { + durations := map[int]int{} + for _, path := range enumeratePaths(node) { + for id := range path { + durations[id] = 1 + } + } + return weightedMakespan(node, durations) +} + +func weightedMakespan(node *resolve.FetchTreeNode, durations map[int]int) int { + max := 0 + for _, path := range enumeratePaths(node) { + sum := 0 + for id := range path { + sum += durations[id] + } + if sum > max { + max = sum + } + } + return max +} + +func sequenceOf(children ...*resolve.FetchTreeNode) *resolve.FetchTreeNode { + children = compactNodes(children) + children = flattenKind(children, resolve.FetchTreeNodeKindSequence) + if len(children) == 0 { + return nil + } + if len(children) == 1 { + return children[0] + } + return resolve.Sequence(children...) +} + +func parallelOf(children ...*resolve.FetchTreeNode) *resolve.FetchTreeNode { + children = compactNodes(children) + children = flattenKind(children, resolve.FetchTreeNodeKindParallel) + if len(children) == 0 { + return nil + } + if len(children) == 1 { + return children[0] + } + sortBranches(children) + return resolve.Parallel(children...) +} + +func flattenKind(nodes []*resolve.FetchTreeNode, kind resolve.FetchTreeNodeKind) []*resolve.FetchTreeNode { + out := nodes[:0] + for _, node := range nodes { + if node != nil && node.Kind == kind { + out = append(out, node.ChildNodes...) + continue + } + out = append(out, node) + } + return out +} + +func compactNodes(nodes []*resolve.FetchTreeNode) []*resolve.FetchTreeNode { + out := nodes[:0] + for _, node := range nodes { + if node != nil { + out = append(out, node) + } + } + return out +} + +func sortBranches(nodes []*resolve.FetchTreeNode) { + slices.SortFunc(nodes, func(a, b *resolve.FetchTreeNode) int { + return minReachableFetchID(a) - minReachableFetchID(b) + }) +} + +func minReachableFetchID(node *resolve.FetchTreeNode) int { + if node == nil { + return 0 + } + if node.Kind == resolve.FetchTreeNodeKindSingle { + return node.Item.Fetch.Dependencies().FetchID + } + min := int(^uint(0) >> 1) + for _, child := range node.ChildNodes { + childMin := minReachableFetchID(child) + if childMin < min { + min = childMin + } + } + return min +} + +func sortedUnique(ids []int) []int { + if len(ids) == 0 { + return nil + } + out := append([]int{}, ids...) + slices.Sort(out) + return slices.Compact(out) +} + +func sortedSet(set map[int]struct{}) []int { + out := make([]int, 0, len(set)) + for id := range set { + out = append(out, id) + } + slices.Sort(out) + return out +} + +func idSet(ids []int) map[int]struct{} { + set := make(map[int]struct{}, len(ids)) + for _, id := range ids { + set[id] = struct{}{} + } + return set +} + +func clonePending(in map[int]map[int]struct{}) map[int]map[int]struct{} { + out := make(map[int]map[int]struct{}, len(in)) + for id, pending := range in { + out[id] = cloneSet(pending) + } + return out +} + +func cloneSet(in map[int]struct{}) map[int]struct{} { + out := make(map[int]struct{}, len(in)) + for id := range in { + out[id] = struct{}{} + } + return out +} + +func filterSet(in map[int]struct{}, allowed map[int]struct{}) map[int]struct{} { + out := make(map[int]struct{}, len(in)) + for id := range in { + if _, ok := allowed[id]; ok { + out[id] = struct{}{} + } + } + return out +} + +func intersectSets(a, b map[int]struct{}) map[int]struct{} { + out := make(map[int]struct{}) + for id := range a { + if _, ok := b[id]; ok { + out[id] = struct{}{} + } + } + return out +} + +func setContainsAll(set, subset map[int]struct{}) bool { + for id := range subset { + if _, ok := set[id]; !ok { + return false + } + } + return true +} + +func removeIDs(ids []int, sets ...map[int]struct{}) []int { + out := ids[:0] + for _, id := range ids { + remove := false + for _, set := range sets { + if _, ok := set[id]; ok { + remove = true + break + } + } + if !remove { + out = append(out, id) + } + } + return out +} diff --git a/v2/pkg/engine/postprocess/build_schedule_tree_bench_plan_test.go b/v2/pkg/engine/postprocess/build_schedule_tree_bench_plan_test.go new file mode 100644 index 0000000000..620fa7d467 --- /dev/null +++ b/v2/pkg/engine/postprocess/build_schedule_tree_bench_plan_test.go @@ -0,0 +1,113 @@ +package postprocess + +import ( + "testing" + + "github.com/stretchr/testify/require" + + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" +) + +func sfDeps(id int, responsePath string, deps ...int) *resolve.FetchTreeNode { + return resolve.SingleWithPath(&resolve.SingleFetch{ + FetchDependencies: resolve.FetchDependencies{FetchID: id, DependsOnFetchIDs: deps}, + }, responsePath) +} + +// TestBuildScheduleTreeRealGatewaysBenchmarkPlan pins the schedule for the +// REAL hive gateways-benchmark plan (13 fetches, 4 legacy waves), extracted +// verbatim from cosmo's plan-generator on the bench supergraph. Upstream's +// path-containment validator heuristics rejected this plan (fetch 3 nests +// under fetch 2's provided path without depending on it; fetches 6 and 11 +// provide the same path) and the upstream processor PANICKED at plan time — +// a 500 for every benchmark request. With the fork's FetchID-edge-only +// validator the scheduler produces the eager-inline SP schedule below: the +// users chain and the topProducts chain are fully decoupled (no cross-branch +// barriers), which is exactly the skew win the scheduler exists to capture. +func TestBuildScheduleTreeRealGatewaysBenchmarkPlan(t *testing.T) { + build := func() *resolve.FetchTreeNode { + return resolve.Sequence( + sfDeps(0, ""), + sfDeps(5, ""), + sfDeps(1, "users", 0), + sfDeps(6, "topProducts", 5), + sfDeps(11, "topProducts", 5), + sfDeps(2, "users.@.reviews.@.product", 1), + sfDeps(3, "users.@.reviews.@.product.reviews.@.author", 1), + sfDeps(4, "users.@.reviews.@.product.reviews.@.author.reviews.@.product", 1), + sfDeps(7, "topProducts.@.reviews.@.author", 6), + sfDeps(8, "topProducts.@.reviews.@.author.reviews.@.product", 6), + sfDeps(9, "users.@.reviews.@.product", 1, 2), + sfDeps(10, "users.@.reviews.@.product.reviews.@.author.reviews.@.product", 1, 4), + sfDeps(12, "topProducts.@.reviews.@.author.reviews.@.product", 6, 8), + ) + } + + root := build() + dag, err := newFetchDAG(root.ChildNodes) + require.NoError(t, err) + tree, err := buildScheduleTree(root.ChildNodes, dag) + require.NoError(t, err) + require.NoError(t, validateSchedule(tree, dag)) + + expected := par( + seq( + sfDeps(0, ""), + sfDeps(1, "users", 0), + par( + seq( + sfDeps(2, "users.@.reviews.@.product", 1), + sfDeps(9, "users.@.reviews.@.product", 1, 2), + ), + sfDeps(3, "users.@.reviews.@.product.reviews.@.author", 1), + seq( + sfDeps(4, "users.@.reviews.@.product.reviews.@.author.reviews.@.product", 1), + sfDeps(10, "users.@.reviews.@.product.reviews.@.author.reviews.@.product", 1, 4), + ), + ), + ), + seq( + sfDeps(5, ""), + par( + seq( + sfDeps(6, "topProducts", 5), + par( + sfDeps(7, "topProducts.@.reviews.@.author", 6), + seq( + sfDeps(8, "topProducts.@.reviews.@.author.reviews.@.product", 6), + sfDeps(12, "topProducts.@.reviews.@.author.reviews.@.product", 6, 8), + ), + ), + ), + sfDeps(11, "topProducts", 5), + ), + ), + ) + require.Equal(t, expected, tree) + + // End to end through the processor: must take the schedule path (the + // schedule differs from the legacy wave tree), not the error fallback. + processed := build() + (&buildScheduleTreeProcessor{}).ProcessFetchTree(processed) + require.Equal(t, expected, processed) +} + +// TestBuildScheduleTreeProcessorFallsBackOnError pins the fork's no-panic +// contract: on any scheduler/validator error (duplicate FetchIDs here) the +// processor degrades to the LEGACY wave pipeline — byte-identical planning to +// the flag-off path — instead of panicking like upstream. +func TestBuildScheduleTreeProcessorFallsBackOnError(t *testing.T) { + build := func() *resolve.FetchTreeNode { + return seq(sf(7), sf(7, 7)) + } + + legacy := build() + (&orderSequenceByDependencies{}).ProcessFetchTree(legacy) + (&createParallelNodes{}).ProcessFetchTree(legacy) + + scheduled := build() + require.NotPanics(t, func() { + (&buildScheduleTreeProcessor{}).ProcessFetchTree(scheduled) + }) + require.Equal(t, legacy, scheduled) +} diff --git a/v2/pkg/engine/postprocess/build_schedule_tree_property_test.go b/v2/pkg/engine/postprocess/build_schedule_tree_property_test.go new file mode 100644 index 0000000000..63b3cb3e25 --- /dev/null +++ b/v2/pkg/engine/postprocess/build_schedule_tree_property_test.go @@ -0,0 +1,163 @@ +package postprocess + +import ( + "math" + "math/rand" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve" +) + +func TestBuildScheduleTreePropertiesExhaustive(t *testing.T) { + for n := 0; n <= 6; n++ { + edgeCount := n * (n - 1) / 2 + for mask := 0; mask < 1< 0 && dag.hasCycle() { + require.EqualError(t, spErr, "cycle detected in fetch dependency graph") + require.EqualError(t, levelErr, "cycle detected in fetch dependency graph") + require.EqualError(t, hybridErr, "cycle detected in fetch dependency graph") + return + } + require.NoErrorf(t, spErr, "input=%v", dependencyList(input)) + require.NoErrorf(t, levelErr, "input=%v", dependencyList(input)) + require.NoErrorf(t, hybridErr, "input=%v", dependencyList(input)) + require.NoError(t, validateSchedule(sp, dag)) + require.NoError(t, validateSchedule(level, dag)) + require.NoError(t, validateSchedule(hybrid, dag)) + + if dominates(sp, level) { + require.Equal(t, sp, hybrid) + } else { + require.Equal(t, level, hybrid) + } + require.LessOrEqual(t, uniformMakespan(hybrid), uniformMakespan(level)) + for _, durations := range profiles { + hybridMakespan := weightedMakespan(hybrid, durations) + levelMakespan := weightedMakespan(level, durations) + require.LessOrEqual(t, hybridMakespan, levelMakespan) + if dominates(sp, level) { + require.LessOrEqual(t, weightedMakespan(sp, durations), levelMakespan) + } + _ = !dominates(sp, level) && weightedMakespan(sp, durations) < levelMakespan + } +} + +func dependencyList(input []*resolve.FetchTreeNode) []resolve.FetchDependencies { + out := make([]resolve.FetchDependencies, 0, len(input)) + for _, node := range input { + out = append(out, *node.Item.Fetch.Dependencies()) + } + return out +} + +func dagFromMask(n, mask int) []*resolve.FetchTreeNode { + deps := make([][]int, n) + bit := 0 + for from := range n { + for to := from + 1; to < n; to++ { + if mask&(1< Date: Thu, 11 Jun 2026 16:07:09 +0200 Subject: [PATCH 3/5] feat(resolve): 3-phase prepare/load/merge protocol for nested schedule trees MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ports the improve-query-order loader protocol with three deviations, each pinned by a regression test: 1. resolveParallelWithCtx routes flat Parallels through the locked nested path whenever useMergeMu is set — upstream's bare shape check lets a flat Parallel inside a nested tree merge lock-free against locked sibling merges (verified data race; TestLoaderNestedTreeFlatParallelTakesLockedPath fails under -race if the gate is reverted, verified empirically). 2. shouldSkipErroredDependencyLocked is gated on useMergeMu so flat-plan error-path bytes stay identical to the legacy executor. 3. The pooled batch tools are registered for Put BEFORE the prepare-error return; upstream leaks the pooled arena on every failed batch prepare. Also documents the nested executor's completion-order sink-append property and its tree-order staging follow-up. Sink audit re-run against master: the three order-bearing sinks (resolvable.errors, ctx.subgraphErrors, resolvable.subgraphExtensions) are written only from merge-phase functions; resolvable.go's growth since v2.4.1 added no new writers; executeSourceLoad remains arena-free. Co-Authored-By: Claude Opus 4.8 (1M context) --- v2/pkg/engine/resolve/loader.go | 528 +++++++++++++++++- .../resolve/loader_nested_parallel_test.go | 472 ++++++++++++++++ 2 files changed, 971 insertions(+), 29 deletions(-) create mode 100644 v2/pkg/engine/resolve/loader_nested_parallel_test.go diff --git a/v2/pkg/engine/resolve/loader.go b/v2/pkg/engine/resolve/loader.go index daddafde9a..fb8da45260 100644 --- a/v2/pkg/engine/resolve/loader.go +++ b/v2/pkg/engine/resolve/loader.go @@ -177,6 +177,10 @@ type Loader struct { validateRequiredExternalFields bool taintedObjs taintedObjects + mergeMu sync.Mutex + useMergeMu bool + + erroredFetchIDs map[int]struct{} // jsonArena is the arena for JSON allocation, supplied by the Resolver. // Not thread safe — only use from the main goroutine. @@ -199,6 +203,8 @@ func (l *Loader) Free() { l.ctx = nil l.resolvable = nil l.taintedObjs = nil + l.erroredFetchIDs = nil + l.useMergeMu = false } func (l *Loader) LoadGraphQLResponseData(ctx *Context, response *GraphQLResponse, resolvable *Resolvable) (err error) { @@ -206,29 +212,52 @@ func (l *Loader) LoadGraphQLResponseData(ctx *Context, response *GraphQLResponse l.ctx = ctx l.info = response.Info l.taintedObjs = make(taintedObjects) + l.erroredFetchIDs = nil + // useMergeMu must be set BEFORE the dataflow branch: resolveDataflow's internal + // fallbacks call resolveFetchNode, which must take the locked nested path for + // nested (schedule-tree) plans. + l.useMergeMu = fetchTreeHasNestedParallel(response.Fetches) return l.resolveFetchNode(response.Fetches) } func (l *Loader) resolveFetchNode(node *FetchTreeNode) error { + return l.resolveFetchNodeWithCtx(l.ctx.ctx, node) +} + +func (l *Loader) resolveFetchNodeWithCtx(ctx context.Context, node *FetchTreeNode) error { if node == nil { return nil } switch node.Kind { case FetchTreeNodeKindSingle: - return l.resolveSingle(node.Item) + return l.resolveSingleWithCtx(ctx, node.Item) case FetchTreeNodeKindSequence: - return l.resolveSerial(node.ChildNodes) + return l.resolveSerialWithCtx(ctx, node.ChildNodes) case FetchTreeNodeKindParallel: - return l.resolveParallel(node.ChildNodes) + return l.resolveParallelWithCtx(ctx, node.ChildNodes) default: return nil } } func (l *Loader) resolveParallel(nodes []*FetchTreeNode) error { + return l.resolveParallelWithCtx(l.ctx.ctx, nodes) +} + +func (l *Loader) resolveParallelWithCtx(ctx context.Context, nodes []*FetchTreeNode) error { if len(nodes) == 0 { return nil } + // DEVIATION from upstream (improve-query-order): the gate also routes flat + // Parallel nodes through the locked nested path whenever useMergeMu is set. + // Upstream checks only !allChildrenAreSingle(nodes), which lets a flat Parallel + // INSIDE a nested tree run this lock-free fast path and merge concurrently with + // sibling branches merging under mergeMu — a verified data race. + // With the gate, the lock-free fast path runs only when the whole tree is flat + // (useMergeMu false), where single-threaded merge order is guaranteed. + if l.useMergeMu || !allChildrenAreSingle(nodes) { + return l.resolveParallelNested(ctx, nodes) + } results := make([]*result, len(nodes)) // Allocate the result structs as a single slab instead of len(nodes) individual heap // allocations; the pointers into it stay valid for the function lifetime. @@ -240,7 +269,7 @@ func (l *Loader) resolveParallel(nodes []*FetchTreeNode) error { } }() itemsItems := make([][]*astjson.Value, len(nodes)) - g, ctx := errgroup.WithContext(l.ctx.ctx) + g, ctx := errgroup.WithContext(ctx) for i := range nodes { results[i] = &resultStore[i] itemsItems[i] = l.selectItemsForPath(nodes[i].Item.FetchPath) @@ -277,8 +306,12 @@ func (l *Loader) resolveParallel(nodes []*FetchTreeNode) error { } func (l *Loader) resolveSerial(nodes []*FetchTreeNode) error { + return l.resolveSerialWithCtx(l.ctx.ctx, nodes) +} + +func (l *Loader) resolveSerialWithCtx(ctx context.Context, nodes []*FetchTreeNode) error { for i := range nodes { - err := l.resolveFetchNode(nodes[i]) + err := l.resolveFetchNodeWithCtx(ctx, nodes[i]) if err != nil { return errors.WithStack(err) } @@ -287,43 +320,480 @@ func (l *Loader) resolveSerial(nodes []*FetchTreeNode) error { } func (l *Loader) resolveSingle(item *FetchItem) error { + return l.resolveSingleWithCtx(l.ctx.ctx, item) +} + +func (l *Loader) resolveSingleWithCtx(ctx context.Context, item *FetchItem) error { if item == nil { return nil } - items := l.selectItemsForPath(item.FetchPath) + prepared, err := l.preparePhase(item) + // DEVIATION from upstream (codex P1 on this port): register the tools Put + // BEFORE the error return. prepareBatchEntityFetch acquires res.tools and can + // still fail afterwards (header/item/separator/footer render, + // SetInputUndefinedVariables, validatePreFetch); upstream returns on err before + // its defer is registered, leaking the pooled arena on every such request. + // Put(nil) is a no-op, so registering for any non-nil prepared is safe for all + // fetch kinds. + if prepared != nil { + defer batchEntityToolPool.Put(prepared.res.tools) + } + if err != nil { + return errors.WithStack(err) + } + if prepared == nil { + return nil + } + if err := l.loadPhase(ctx, prepared); err != nil { + return errors.WithStack(err) + } + return l.mergePhase(prepared) +} + +type preparedFetch struct { + item *FetchItem + items []*astjson.Value + res *result + source DataSource + input []byte + trace *DataSourceLoadTrace + skipLoad bool + batchFetch bool +} + +// resolveParallelNested executes nested (schedule-tree) Parallel branches +// concurrently; each leaf's mergePhase runs under mergeMu at COMPLETION time. +// KNOWN PROPERTY (upstream design, documented for ENGINE_ENABLE_SCHEDULE_TREE): +// order-bearing sinks (resolvable.errors, subgraphExtensions, subgraphErrors) +// are appended in completion order across branches, so the ERROR-path byte +// order is timing-dependent — unlike the legacy flat Parallel, which merges by +// child index after its barrier. This holds for ANY cross-branch parallel +// (e.g. two independent chains), not just same-path siblings; the success-path +// data is unaffected (serialized in query order by the resolvable). The named +// follow-up is to stage per-leaf sink appends and flush in tree order, exactly +// as resolveDataflow's swap-capture staging already does (loader_dataflow.go +// invariant 2); not done here to keep the port surgical. +func (l *Loader) resolveParallelNested(ctx context.Context, nodes []*FetchTreeNode) error { + g, ctx := errgroup.WithContext(ctx) + for i := range nodes { + node := nodes[i] + g.Go(func() error { + return l.resolveFetchNodeWithCtx(ctx, node) + }) + } + if err := g.Wait(); err != nil { + return errors.WithStack(err) + } + return nil +} - switch f := item.Fetch.(type) { +func (l *Loader) preparePhase(item *FetchItem) (*preparedFetch, error) { + unlock := l.maybeLock() + defer unlock() + + // DEVIATION from upstream (improve-query-order): the skip-on-errored-dependency + // check is gated on useMergeMu. Upstream runs it unconditionally, which observably + // changes flat-plan error-path bytes vs the legacy executor (dependents of an + // errored fetch would be skipped instead of loaded). Flags-off byte-identity is a + // hard constraint of this fork; nested (schedule-tree) plans keep upstream's + // skip-dependents behavior. + if l.useMergeMu && l.shouldSkipErroredDependencyLocked(item) { + return nil, nil + } + + items := l.selectItemsForPath(item.FetchPath) + res := &result{} + prepared := &preparedFetch{ + item: item, + items: items, + res: res, + } + switch fetch := item.Fetch.(type) { case *SingleFetch: - res := &result{} - err := l.loadSingleFetch(l.ctx.ctx, f, item, items, res) - if err != nil { - return err + err := l.prepareSingleFetch(item, fetch, items, res, prepared) + return prepared, err + case *EntityFetch: + err := l.prepareEntityFetch(item, fetch, items, res, prepared) + return prepared, err + case *BatchEntityFetch: + prepared.batchFetch = true + err := l.prepareBatchEntityFetch(item, fetch, items, res, prepared) + return prepared, err + default: + return nil, nil + } +} + +func (l *Loader) loadPhase(ctx context.Context, prepared *preparedFetch) error { + if prepared.skipLoad { + return nil + } + l.executeSourceLoad(ctx, prepared.item, prepared.source, prepared.input, prepared.res, prepared.trace) + if prepared.res.err != nil { + l.recordErroredFetchID(prepared.item) + } + return nil +} + +func (l *Loader) mergePhase(prepared *preparedFetch) error { + unlock := l.maybeLock() + defer unlock() + + res := prepared.res + var err error + if res.nestedMergeItems != nil { + for j := range res.nestedMergeItems { + err = l.mergeResult(prepared.item, res.nestedMergeItems[j], prepared.items[j:j+1]) + l.callOnFinished(res.nestedMergeItems[j]) + if err != nil { + return errors.WithStack(err) + } + } + return nil + } + err = l.mergeResult(prepared.item, res, prepared.items) + l.callOnFinished(res) + return err +} + +func (l *Loader) prepareSingleFetch(fetchItem *FetchItem, fetch *SingleFetch, items []*astjson.Value, res *result, prepared *preparedFetch) error { + res.init(fetch.PostProcessing, fetch.Info) + buf := bytes.NewBuffer(nil) + + inputData := l.itemsData(items) + if l.ctx.TracingOptions.Enable { + fetch.Trace = &DataSourceLoadTrace{} + if !l.ctx.TracingOptions.ExcludeRawInputData && inputData != nil { + fetch.Trace.RawInputData, _ = l.compactJSON(inputData.MarshalTo(nil)) } - err = l.mergeResult(item, res, items) - l.callOnFinished(res) + } + + // When we don't have parent data it makes no sense to proceed with next fetches in a sequence + // Right now, it is the case only for the introspection - because introspection uses + // only single fetches. + // Having null means that the previous fetch returned null as data + if len(items) == 1 && items[0].Type() == astjson.TypeNull { + res.fetchSkipped = true + prepared.skipLoad = true + if l.ctx.TracingOptions.Enable { + fetch.Trace.LoadSkipped = true + } + return nil + } + + err := fetch.InputTemplate.Render(l.ctx, inputData, buf) + if err != nil { + res.out = l.renderErrorsInvalidInput(fetchItem) + prepared.skipLoad = true + return nil + } + fetchInput := buf.Bytes() + allowed, err := l.validatePreFetch(fetchInput, fetch.Info, res) + if err != nil { return err - case *BatchEntityFetch: - res := &result{} - defer batchEntityToolPool.Put(res.tools) - err := l.loadBatchEntityFetch(l.ctx.ctx, item, f, items, res) - if err != nil { - return errors.WithStack(err) + } + if !allowed { + prepared.skipLoad = true + return nil + } + prepared.source = fetch.DataSource + prepared.input = fetchInput + prepared.trace = fetch.Trace + return nil +} + +func (l *Loader) prepareEntityFetch(fetchItem *FetchItem, fetch *EntityFetch, items []*astjson.Value, res *result, prepared *preparedFetch) error { + res.init(fetch.PostProcessing, fetch.Info) + input := l.itemsData(items) + if l.ctx.TracingOptions.Enable { + fetch.Trace = &DataSourceLoadTrace{} + if !l.ctx.TracingOptions.ExcludeRawInputData && input != nil { + fetch.Trace.RawInputData, _ = l.compactJSON(input.MarshalTo(nil)) + } + } + + preparedInput := bytes.NewBuffer(nil) + item := bytes.NewBuffer(nil) + + var undefinedVariables []string + + err := fetch.Input.Header.RenderAndCollectUndefinedVariables(l.ctx, nil, preparedInput, &undefinedVariables) + if err != nil { + return errors.WithStack(err) + } + + err = fetch.Input.Item.Render(l.ctx, input, item) + if err != nil { + if fetch.Input.SkipErrItem { + // skip fetch on render item error + if l.ctx.TracingOptions.Enable { + fetch.Trace.LoadSkipped = true + } + res.fetchSkipped = true + prepared.skipLoad = true + return nil } - err = l.mergeResult(item, res, items) - l.callOnFinished(res) + return errors.WithStack(err) + } + renderedItem := item.Bytes() + if bytes.Equal(renderedItem, null) { + // skip fetch if item is null + res.fetchSkipped = true + if l.ctx.TracingOptions.Enable { + fetch.Trace.LoadSkipped = true + } else { + prepared.skipLoad = true + return nil + } + } + if bytes.Equal(renderedItem, emptyObject) { + // skip fetch if item is empty + res.fetchSkipped = true + if l.ctx.TracingOptions.Enable { + fetch.Trace.LoadSkipped = true + } else { + prepared.skipLoad = true + return nil + } + } + _, _ = item.WriteTo(preparedInput) + err = fetch.Input.Footer.RenderAndCollectUndefinedVariables(l.ctx, nil, preparedInput, &undefinedVariables) + if err != nil { + return errors.WithStack(err) + } + + err = SetInputUndefinedVariables(preparedInput, undefinedVariables) + if err != nil { + return errors.WithStack(err) + } + fetchInput := preparedInput.Bytes() + + if l.ctx.TracingOptions.Enable && res.fetchSkipped { + l.setTracingInput(fetchItem, fetchInput, fetch.Trace) + prepared.skipLoad = true + return nil + } + + allowed, err := l.validatePreFetch(fetchInput, fetch.Info, res) + if err != nil { return err - case *EntityFetch: - res := &result{} - err := l.loadEntityFetch(l.ctx.ctx, item, f, items, res) - if err != nil { - return errors.WithStack(err) + } + if !allowed { + prepared.skipLoad = true + return nil + } + prepared.source = fetch.DataSource + prepared.input = fetchInput + prepared.trace = fetch.Trace + return nil +} + +func (l *Loader) prepareBatchEntityFetch(fetchItem *FetchItem, fetch *BatchEntityFetch, items []*astjson.Value, res *result, prepared *preparedFetch) error { + res.init(fetch.PostProcessing, fetch.Info) + + if l.ctx.TracingOptions.Enable { + fetch.Trace = &DataSourceLoadTrace{} + if !l.ctx.TracingOptions.ExcludeRawInputData && len(items) != 0 { + data := l.itemsData(items) + if data != nil { + fetch.Trace.RawInputData, _ = l.compactJSON(data.MarshalTo(nil)) + } + } + } + + res.tools = batchEntityToolPool.Get(len(items)) + preparedInput := arena.NewArenaBuffer(res.tools.a) + itemInput := arena.NewArenaBuffer(res.tools.a) + batchStats := arena.AllocateSlice[[]*astjson.Value](res.tools.a, 0, len(items)) + defer func() { + // we need to clear the batchStats slice to avoid memory corruption + // once the outer func returns, we must not keep pointers to items on the arena + for i := range batchStats { + // nolint:ineffassign + batchStats[i] = nil } - err = l.mergeResult(item, res, items) - l.callOnFinished(res) + // nolint:ineffassign + batchStats = nil + }() + + // I tried using arena here, but it only worsened the situation + var undefinedVariables []string + + err := fetch.Input.Header.RenderAndCollectUndefinedVariables(l.ctx, nil, preparedInput, &undefinedVariables) + if err != nil { + return errors.WithStack(err) + } + batchItemIndex := 0 + addSeparator := false + +WithNextItem: + for i, item := range items { + for j := range fetch.Input.Items { + itemInput.Reset() + err = fetch.Input.Items[j].Render(l.ctx, item, itemInput) + if err != nil { + if fetch.Input.SkipErrItems { + err = nil // nolint:ineffassign + continue + } + if l.ctx.TracingOptions.Enable { + fetch.Trace.LoadSkipped = true + } + return errors.WithStack(err) + } + if fetch.Input.SkipNullItems && itemInput.Len() == 4 && bytes.Equal(itemInput.Bytes(), null) { + continue + } + if fetch.Input.SkipEmptyObjectItems && itemInput.Len() == 2 && bytes.Equal(itemInput.Bytes(), emptyObject) { + continue + } + + res.tools.keyGen.Reset() + _, _ = res.tools.keyGen.Write(itemInput.Bytes()) + itemHash := res.tools.keyGen.Sum64() + if existingIndex, ok := res.tools.batchHashToIndex[itemHash]; ok { + batchStats[existingIndex] = arena.SliceAppend(res.tools.a, batchStats[existingIndex], items[i]) + continue WithNextItem + } else { + if addSeparator { + err = fetch.Input.Separator.Render(l.ctx, nil, preparedInput) + if err != nil { + return errors.WithStack(err) + } + } + _, _ = itemInput.WriteTo(preparedInput) + // new unique representation + res.tools.batchHashToIndex[itemHash] = batchItemIndex + // create a new targets bucket for this unique index + batchStats = arena.SliceAppend(res.tools.a, batchStats, []*astjson.Value{items[i]}) + batchItemIndex++ + addSeparator = true + } + } + } + + if len(batchStats) == 0 { + // all items were skipped - discard fetch + res.fetchSkipped = true + if l.ctx.TracingOptions.Enable { + fetch.Trace.LoadSkipped = true + } else { + prepared.skipLoad = true + return nil + } + } + + err = fetch.Input.Footer.RenderAndCollectUndefinedVariables(l.ctx, nil, preparedInput, &undefinedVariables) + if err != nil { + return errors.WithStack(err) + } + + err = SetInputUndefinedVariables(preparedInput, undefinedVariables) + if err != nil { + return errors.WithStack(err) + } + + fetchInput := preparedInput.Bytes() + // it's important to copy the *astjson.Value's off the arena to avoid memory corruption + res.batchStats = make([][]*astjson.Value, len(batchStats)) + for i := range batchStats { + res.batchStats[i] = make([]*astjson.Value, len(batchStats[i])) + copy(res.batchStats[i], batchStats[i]) + } + + if l.ctx.TracingOptions.Enable && res.fetchSkipped { + l.setTracingInput(fetchItem, fetchInput, fetch.Trace) + prepared.skipLoad = true + return nil + } + + allowed, err := l.validatePreFetch(fetchInput, fetch.Info, res) + if err != nil { return err - default: + } + if !allowed { + prepared.skipLoad = true return nil } + + prepared.source = fetch.DataSource + prepared.input = fetchInput + prepared.trace = fetch.Trace + return nil +} + +func (l *Loader) maybeLock() func() { + if !l.useMergeMu { + return func() {} + } + l.mergeMu.Lock() + return l.mergeMu.Unlock +} + +func (l *Loader) shouldSkipErroredDependencyLocked(item *FetchItem) bool { + if item == nil || item.Fetch == nil || len(l.erroredFetchIDs) == 0 { + return false + } + dependencies := item.Fetch.Dependencies() + if dependencies == nil { + return false + } + for _, dependencyID := range dependencies.DependsOnFetchIDs { + if _, ok := l.erroredFetchIDs[dependencyID]; ok { + l.recordErroredFetchIDLocked(item) + return true + } + } + return false +} + +func (l *Loader) recordErroredFetchID(item *FetchItem) { + unlock := l.maybeLock() + defer unlock() + + l.recordErroredFetchIDLocked(item) +} + +func (l *Loader) recordErroredFetchIDLocked(item *FetchItem) { + if item == nil || item.Fetch == nil { + return + } + dependencies := item.Fetch.Dependencies() + if dependencies == nil { + return + } + if l.erroredFetchIDs == nil { + l.erroredFetchIDs = make(map[int]struct{}) + } + l.erroredFetchIDs[dependencies.FetchID] = struct{}{} +} + +func allChildrenAreSingle(nodes []*FetchTreeNode) bool { + for _, node := range nodes { + if node == nil || node.Kind != FetchTreeNodeKindSingle { + return false + } + } + return true +} + +func fetchTreeHasNestedParallel(node *FetchTreeNode) bool { + if node == nil { + return false + } + if node.Kind == FetchTreeNodeKindParallel { + for _, child := range node.ChildNodes { + if child != nil && child.Kind != FetchTreeNodeKindSingle { + return true + } + } + } + for _, child := range node.ChildNodes { + if fetchTreeHasNestedParallel(child) { + return true + } + } + return false } func (l *Loader) callOnFinished(res *result) { diff --git a/v2/pkg/engine/resolve/loader_nested_parallel_test.go b/v2/pkg/engine/resolve/loader_nested_parallel_test.go new file mode 100644 index 0000000000..9cf12adfd0 --- /dev/null +++ b/v2/pkg/engine/resolve/loader_nested_parallel_test.go @@ -0,0 +1,472 @@ +package resolve + +import ( + "bytes" + "context" + "errors" + "net/http" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/wundergraph/graphql-go-tools/v2/pkg/ast" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/httpclient" + "github.com/wundergraph/graphql-go-tools/v2/pkg/fastjsonext" +) + +type controlledLoaderDataSource struct { + response []byte + err error + + waitFor <-chan struct{} + waitForCancel bool + + startedOnce sync.Once + doneOnce sync.Once + started chan struct{} + done chan struct{} + + mu sync.Mutex + inputs []string + + cancelled atomic.Bool + loadCalls atomic.Int64 +} + +func newControlledLoaderDataSource(response string) *controlledLoaderDataSource { + return &controlledLoaderDataSource{ + response: []byte(response), + started: make(chan struct{}), + done: make(chan struct{}), + } +} + +func (d *controlledLoaderDataSource) Load(ctx context.Context, _ http.Header, input []byte) ([]byte, error) { + d.loadCalls.Add(1) + + d.mu.Lock() + d.inputs = append(d.inputs, string(input)) + d.mu.Unlock() + + d.startedOnce.Do(func() { + close(d.started) + }) + defer d.doneOnce.Do(func() { + close(d.done) + }) + + if d.waitForCancel { + select { + case <-ctx.Done(): + d.cancelled.Store(true) + return nil, ctx.Err() + case <-time.After(5 * time.Second): + return nil, errors.New("timed out waiting for cancellation") + } + } + + if d.waitFor != nil { + select { + case <-d.waitFor: + case <-ctx.Done(): + d.cancelled.Store(true) + return nil, ctx.Err() + case <-time.After(5 * time.Second): + return nil, errors.New("timed out waiting for load release") + } + } + + if d.err != nil { + return nil, d.err + } + return d.response, nil +} + +func (d *controlledLoaderDataSource) LoadWithFiles(ctx context.Context, headers http.Header, input []byte, _ []*httpclient.FileUpload) ([]byte, error) { + return d.Load(ctx, headers, input) +} + +func (d *controlledLoaderDataSource) requireInputs(t *testing.T, expected ...string) { + t.Helper() + + d.mu.Lock() + defer d.mu.Unlock() + + require.Equal(t, expected, d.inputs) +} + +func TestLoaderNestedParallelCorrectness(t *testing.T) { + a := newControlledLoaderDataSource(`{"data":{"a":"A"}}`) + b := newControlledLoaderDataSource(`{"data":{"b":"B"}}`) + c := newControlledLoaderDataSource(`{"data":{"c":"C"}}`) + + response := &GraphQLResponse{ + Info: &GraphQLResponseInfo{OperationType: ast.OperationTypeQuery}, + Fetches: Parallel( + Sequence( + Single(nestedParallelSingleFetch(a, `{"fetch":"A"}`)), + Single(nestedParallelSingleFetchWithTemplate(b, nestedParallelInputForFields("a"))), + ), + Single(nestedParallelSingleFetch(c, `{"fetch":"C"}`)), + ), + Data: nestedParallelData("a", "b", "c"), + } + + ctx := NewContext(context.Background()) + ctx.ExecutionOptions.DisableSubgraphRequestDeduplication = true + resolver := newResolver(context.Background()) + buf := &bytes.Buffer{} + + _, err := resolver.ResolveGraphQLResponse(ctx, response, nil, buf) + require.NoError(t, err) + require.Equal(t, `{"data":{"a":"A","b":"B","c":"C"}}`, buf.String()) + b.requireInputs(t, `{"a":"A"}`) +} + +func TestLoaderNestedParallelRaceMutexDiscipline(t *testing.T) { + a := newControlledLoaderDataSource(`{"data":{"a":"A"}}`) + b := newControlledLoaderDataSource(`{"data":{"b":"B"}}`) + c := newControlledLoaderDataSource(`{"data":{"c":"C"}}`) + d := newControlledLoaderDataSource(`{"data":{"d":"D"}}`) + a.waitFor = d.done + + response := &GraphQLResponse{ + Info: &GraphQLResponseInfo{OperationType: ast.OperationTypeQuery}, + Fetches: Parallel( + Sequence( + Single(nestedParallelSingleFetch(a, `{"fetch":"A"}`)), + Single(nestedParallelSingleFetchWithTemplate(b, nestedParallelInputForFields("a"))), + ), + Sequence( + Single(nestedParallelSingleFetch(c, `{"fetch":"C"}`)), + Single(nestedParallelSingleFetchWithTemplate(d, nestedParallelInputForFields("c"))), + ), + ), + Data: nestedParallelData("a", "b", "c", "d"), + } + + ctx := NewContext(context.Background()) + ctx.ExecutionOptions.DisableSubgraphRequestDeduplication = true + resolver := newResolver(context.Background()) + buf := &bytes.Buffer{} + + _, err := resolver.ResolveGraphQLResponse(ctx, response, nil, buf) + require.NoError(t, err) + require.Equal(t, `{"data":{"a":"A","b":"B","c":"C","d":"D"}}`, buf.String()) + require.Eventually(t, func() bool { + select { + case <-a.started: + return true + default: + return false + } + }, time.Second, time.Millisecond) + require.Eventually(t, func() bool { + select { + case <-d.done: + return true + default: + return false + } + }, time.Second, time.Millisecond) + b.requireInputs(t, `{"a":"A"}`) + d.requireInputs(t, `{"c":"C"}`) +} + +func TestLoaderFlatParallelKeepsFastPathWithoutMergeMutex(t *testing.T) { + a := newControlledLoaderDataSource(`{"data":{"a":"A"}}`) + b := newControlledLoaderDataSource(`{"data":{"b":"B"}}`) + response := &GraphQLResponse{ + Info: &GraphQLResponseInfo{OperationType: ast.OperationTypeQuery}, + Fetches: Parallel( + Single(nestedParallelSingleFetch(a, `{"fetch":"A"}`)), + Single(nestedParallelSingleFetch(b, `{"fetch":"B"}`)), + ), + Data: nestedParallelData("a", "b"), + } + + ctx := NewContext(context.Background()) + ctx.ExecutionOptions.DisableSubgraphRequestDeduplication = true + resolvable := NewResolvable(nil, ResolvableOptions{}) + loader := &Loader{} + + require.NoError(t, resolvable.Init(ctx, nil, ast.OperationTypeQuery)) + require.NoError(t, loader.LoadGraphQLResponseData(ctx, response, resolvable)) + require.False(t, loader.useMergeMu) + require.Equal(t, `{"data":{"a":"A","b":"B"}}`, fastjsonext.PrintGraphQLResponse(resolvable.data, resolvable.errors)) +} + +// TestLoaderNestedTreeFlatParallelTakesLockedPath is the regression test for the +// fork's routing-gate deviation in resolveParallelWithCtx: a FLAT Parallel (all-Single +// children) nested INSIDE a larger tree must take the locked nested path, not the +// lock-free fast path. Upstream's bare !allChildrenAreSingle gate lets branch 1's +// inner Parallel(B,C) merge lock-free while branch 2 merges E under mergeMu — a data +// race on the shared arena. The waitFor wiring forces the two branches' merges to +// overlap; the test must fail under -race if the gate is reverted to upstream's check. +func TestLoaderNestedTreeFlatParallelTakesLockedPath(t *testing.T) { + a := newControlledLoaderDataSource(`{"data":{"a":"A"}}`) + b := newControlledLoaderDataSource(`{"data":{"b":"B"}}`) + c := newControlledLoaderDataSource(`{"data":{"c":"C"}}`) + d := newControlledLoaderDataSource(`{"data":{"d":"D"}}`) + e := newControlledLoaderDataSource(`{"data":{"e":"E"}}`) + // Branch 2 holds D until branch 1's inner flat parallel is in flight, + // then E's start releases C — so B/C merges overlap E's merge. + d.waitFor = b.started + c.waitFor = e.started + + response := &GraphQLResponse{ + Info: &GraphQLResponseInfo{OperationType: ast.OperationTypeQuery}, + Fetches: Parallel( + Sequence( + Single(nestedParallelSingleFetch(a, `{"fetch":"A"}`)), + Parallel( + Single(nestedParallelSingleFetch(b, `{"fetch":"B"}`)), + Single(nestedParallelSingleFetch(c, `{"fetch":"C"}`)), + ), + ), + Sequence( + Single(nestedParallelSingleFetch(d, `{"fetch":"D"}`)), + Single(nestedParallelSingleFetch(e, `{"fetch":"E"}`)), + ), + ), + Data: nestedParallelData("a", "b", "c", "d", "e"), + } + + ctx := NewContext(context.Background()) + ctx.ExecutionOptions.DisableSubgraphRequestDeduplication = true + resolver := newResolver(context.Background()) + buf := &bytes.Buffer{} + + _, err := resolver.ResolveGraphQLResponse(ctx, response, nil, buf) + require.NoError(t, err) + require.Equal(t, `{"data":{"a":"A","b":"B","c":"C","d":"D","e":"E"}}`, buf.String()) + require.Equal(t, int64(1), b.loadCalls.Load()) + require.Equal(t, int64(1), c.loadCalls.Load()) + require.Equal(t, int64(1), e.loadCalls.Load()) +} + +func TestFetchTreeHasNestedParallel(t *testing.T) { + largeFlatChildren := make([]*FetchTreeNode, 100) + for i := range largeFlatChildren { + largeFlatChildren[i] = Single(&SingleFetch{}) + } + + tests := []struct { + name string + node *FetchTreeNode + want bool + }{ + {name: "nil", node: nil, want: false}, + {name: "single", node: Single(&SingleFetch{}), want: false}, + {name: "sequence of singles", node: Sequence(Single(&SingleFetch{}), Single(&SingleFetch{})), want: false}, + {name: "flat parallel of singles", node: Parallel(Single(&SingleFetch{}), Single(&SingleFetch{})), want: false}, + {name: "parallel with sequence child", node: Parallel(Sequence(Single(&SingleFetch{}), Single(&SingleFetch{})), Single(&SingleFetch{})), want: true}, + {name: "deeply buried nested parallel", node: Sequence(Sequence(Sequence(Parallel(Single(&SingleFetch{}), Sequence(Single(&SingleFetch{})))))), want: true}, + {name: "large flat parallel", node: Parallel(largeFlatChildren...), want: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + require.Equal(t, tt.want, fetchTreeHasNestedParallel(tt.node)) + }) + } +} + +func TestResolveParallelNestedDoesNotCancelSiblings(t *testing.T) { + boom := errors.New("boom") + failing := newControlledLoaderDataSource(``) + failing.err = boom + sibling := newControlledLoaderDataSource(`{"data":{"sibling":"ok"}}`) + sibling.waitFor = failing.done + + response := &GraphQLResponse{ + Info: &GraphQLResponseInfo{OperationType: ast.OperationTypeQuery}, + Fetches: Parallel( + Sequence(SingleWithPath(nestedParallelSingleFetchWithInfo(failing, `{"fetch":"failing"}`, "Failing", "Failing"), "query.failing")), + Sequence(SingleWithPath(nestedParallelSingleFetchWithInfo(sibling, `{"fetch":"sibling"}`, "Sibling", "Sibling"), "query.sibling")), + ), + Data: nestedParallelNullableData("failing", "sibling"), + } + + ctx := NewContext(context.Background()) + ctx.ExecutionOptions.DisableSubgraphRequestDeduplication = true + resolver := newResolver(context.Background()) + buf := &bytes.Buffer{} + + _, err := resolver.ResolveGraphQLResponse(ctx, response, nil, buf) + require.NoError(t, err) + require.Equal(t, `{"errors":[{"message":"Failed to fetch from Subgraph 'Failing' at Path 'query.failing'."}],"data":{"failing":null,"sibling":"ok"}}`, buf.String()) + require.False(t, sibling.cancelled.Load()) + require.Equal(t, int64(1), sibling.loadCalls.Load()) +} + +func TestResolveParallelNestedSequenceStopsAfterFailedDependency(t *testing.T) { + boom := errors.New("boom") + a := newControlledLoaderDataSource(``) + a.err = boom + b := newControlledLoaderDataSource(`{"data":{"b":"should not run"}}`) + c := newControlledLoaderDataSource(`{"data":{"c":"ok"}}`) + aFetch := nestedParallelSingleFetchWithInfo(a, `{"fetch":"a"}`, "A", "A") + aFetch.FetchDependencies.FetchID = 0 + bFetch := nestedParallelSingleFetchWithInfoAndTemplate(b, nestedParallelInputForFields("a"), "B", "B") + bFetch.FetchDependencies.FetchID = 1 + bFetch.FetchDependencies.DependsOnFetchIDs = []int{0} + bFetch.PostProcessing.SelectResponseErrorsPath = []string{"errors"} + + response := &GraphQLResponse{ + Info: &GraphQLResponseInfo{OperationType: ast.OperationTypeQuery}, + Fetches: Parallel( + Sequence( + SingleWithPath(aFetch, "query.a"), + SingleWithPath(bFetch, "query.b"), + ), + SingleWithPath(nestedParallelSingleFetchWithInfo(c, `{"fetch":"c"}`, "C", "C"), "query.c"), + ), + Data: nestedParallelNullableData("a", "b", "c"), + } + + ctx := NewContext(context.Background()) + ctx.ExecutionOptions.DisableSubgraphRequestDeduplication = true + resolver := newResolver(context.Background()) + buf := &bytes.Buffer{} + + _, err := resolver.ResolveGraphQLResponse(ctx, response, nil, buf) + require.NoError(t, err) + require.Equal(t, int64(0), b.loadCalls.Load()) + require.Equal(t, int64(1), c.loadCalls.Load()) + require.Equal(t, `{"errors":[{"message":"Failed to fetch from Subgraph 'A' at Path 'query.a'."}],"data":{"a":null,"b":null,"c":"ok"}}`, buf.String()) +} + +func TestResolveSerialStopsTransitivelyOnFailedDependency(t *testing.T) { + boom := errors.New("boom") + a := newControlledLoaderDataSource(``) + a.err = boom + b := newControlledLoaderDataSource(`{"data":{"b":"should not run"}}`) + c := newControlledLoaderDataSource(`{"data":{"c":"should not run"}}`) + + aFetch := nestedParallelSingleFetchWithInfo(a, `{"fetch":"a"}`, "A", "A") + aFetch.FetchDependencies.FetchID = 0 + bFetch := nestedParallelSingleFetchWithInfoAndTemplate(b, nestedParallelInputForFields("a"), "B", "B") + bFetch.FetchDependencies.FetchID = 1 + bFetch.FetchDependencies.DependsOnFetchIDs = []int{0} + cFetch := nestedParallelSingleFetchWithInfoAndTemplate(c, nestedParallelInputForFields("b"), "C", "C") + cFetch.FetchDependencies.FetchID = 2 + cFetch.FetchDependencies.DependsOnFetchIDs = []int{1} + + response := &GraphQLResponse{ + Info: &GraphQLResponseInfo{OperationType: ast.OperationTypeQuery}, + // Nested shape (Parallel with a non-Single child) so the loader takes the + // locked nested path: skip-on-errored-dependency is gated on useMergeMu in + // this fork and deliberately does NOT fire on flat plans (flat plans keep + // the legacy load-everything error semantics for flags-off byte-identity). + Fetches: Sequence( + SingleWithPath(aFetch, "query.a"), + Parallel( + Sequence( + SingleWithPath(bFetch, "query.b"), + SingleWithPath(cFetch, "query.c"), + ), + ), + ), + Data: nestedParallelNullableData("a", "b", "c"), + } + + ctx := NewContext(context.Background()) + ctx.ExecutionOptions.DisableSubgraphRequestDeduplication = true + resolver := newResolver(context.Background()) + buf := &bytes.Buffer{} + + _, err := resolver.ResolveGraphQLResponse(ctx, response, nil, buf) + require.NoError(t, err) + require.Equal(t, int64(0), b.loadCalls.Load()) + require.Equal(t, int64(0), c.loadCalls.Load()) + require.Equal(t, `{"errors":[{"message":"Failed to fetch from Subgraph 'A' at Path 'query.a'."}],"data":{"a":null,"b":null,"c":null}}`, buf.String()) +} + +func nestedParallelSingleFetch(ds DataSource, input string) *SingleFetch { + return nestedParallelSingleFetchWithTemplate(ds, InputTemplate{ + Segments: []TemplateSegment{{ + SegmentType: StaticSegmentType, + Data: []byte(input), + }}, + }) +} + +func nestedParallelSingleFetchWithInfo(ds DataSource, input, id, name string) *SingleFetch { + return nestedParallelSingleFetchWithInfoAndTemplate(ds, InputTemplate{ + Segments: []TemplateSegment{{ + SegmentType: StaticSegmentType, + Data: []byte(input), + }}, + }, id, name) +} + +func nestedParallelSingleFetchWithInfoAndTemplate(ds DataSource, input InputTemplate, id, name string) *SingleFetch { + fetch := nestedParallelSingleFetchWithTemplate(ds, input) + fetch.Info = &FetchInfo{ + DataSourceID: id, + DataSourceName: name, + } + return fetch +} + +func nestedParallelSingleFetchWithTemplate(ds DataSource, input InputTemplate) *SingleFetch { + return &SingleFetch{ + InputTemplate: input, + FetchConfiguration: FetchConfiguration{ + DataSource: ds, + PostProcessing: PostProcessingConfiguration{ + SelectResponseDataPath: []string{"data"}, + }, + }, + } +} + +func nestedParallelInputForFields(fields ...string) InputTemplate { + object := &Object{Fields: make([]*Field, 0, len(fields))} + for _, field := range fields { + object.Fields = append(object.Fields, &Field{ + Name: []byte(field), + Value: &String{ + Path: []string{field}, + }, + }) + } + return InputTemplate{ + Segments: []TemplateSegment{{ + SegmentType: VariableSegmentType, + VariableKind: ResolvableObjectVariableKind, + Renderer: NewGraphQLVariableResolveRenderer(object), + }}, + } +} + +func nestedParallelData(fields ...string) *Object { + data := &Object{Fields: make([]*Field, 0, len(fields))} + for _, field := range fields { + data.Fields = append(data.Fields, &Field{ + Name: []byte(field), + Value: &String{ + Path: []string{field}, + }, + }) + } + return data +} + +func nestedParallelNullableData(fields ...string) *Object { + data := &Object{Fields: make([]*Field, 0, len(fields))} + for _, field := range fields { + data.Fields = append(data.Fields, &Field{ + Name: []byte(field), + Value: &String{ + Path: []string{field}, + Nullable: true, + }, + }) + } + return data +} From 6c19ba03f59707518d673505213bbd32dd7a1fac Mon Sep 17 00:00:00 2001 From: Jens Neuse Date: Thu, 11 Jun 2026 16:14:20 +0200 Subject: [PATCH 4/5] =?UTF-8?q?feat(resolve):=20dataflow=20executor=20?= =?UTF-8?q?=E2=80=94=20per-FetchID=20gating=20with=20coordinator-owned=20a?= =?UTF-8?q?rena?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces the per-wave barrier with per-FetchID dependency gating for flat query plans (default off; ResolverOptions.EnableDataflowExecution). Each fetch's network load starts when its OWN dependencies have merged, which collapses wall-clock latency to the dependency-graph critical path under skewed subgraph latency (measured -37% on the gateways-benchmark workload). Correctness by construction, with nine falsifiable invariants documented in the file header: every arena access runs on the coordinator goroutine (workers execute only the arena-free source load); the three audited error sinks are swap-captured per leaf and replayed in leaf order, making error/extension output byte-identical to the wave executor; dispatch pops the global leaf-index minimum so pre-fetch hook order matches wave spawn order whenever the DAG permits; nested (schedule-tree) plans are structurally rejected and fall back to the wave executor. Both dispatch-order falsification checks verified on this base: FIFO pop and FetchID-order pop each make TestDataflowPreFetchHookOrder fail. Co-Authored-By: Claude Opus 4.8 (1M context) --- v2/pkg/engine/resolve/loader.go | 7 + v2/pkg/engine/resolve/loader_dataflow.go | 486 +++++++++++ v2/pkg/engine/resolve/loader_dataflow_test.go | 767 ++++++++++++++++++ v2/pkg/engine/resolve/resolve.go | 7 + 4 files changed, 1267 insertions(+) create mode 100644 v2/pkg/engine/resolve/loader_dataflow.go create mode 100644 v2/pkg/engine/resolve/loader_dataflow_test.go diff --git a/v2/pkg/engine/resolve/loader.go b/v2/pkg/engine/resolve/loader.go index fb8da45260..76c5e7777e 100644 --- a/v2/pkg/engine/resolve/loader.go +++ b/v2/pkg/engine/resolve/loader.go @@ -176,6 +176,10 @@ type Loader struct { validateRequiredExternalFields bool + // enableDataflow runs query fetch DAGs through resolveDataflow (per-FetchID + // dependency gating) instead of the per-wave-barrier resolveSerial/resolveParallel. + enableDataflow bool + taintedObjs taintedObjects mergeMu sync.Mutex useMergeMu bool @@ -217,6 +221,9 @@ func (l *Loader) LoadGraphQLResponseData(ctx *Context, response *GraphQLResponse // fallbacks call resolveFetchNode, which must take the locked nested path for // nested (schedule-tree) plans. l.useMergeMu = fetchTreeHasNestedParallel(response.Fetches) + if l.enableDataflow && l.dataflowEligibleOperation() { + return l.resolveDataflow(response.Fetches) + } return l.resolveFetchNode(response.Fetches) } diff --git a/v2/pkg/engine/resolve/loader_dataflow.go b/v2/pkg/engine/resolve/loader_dataflow.go new file mode 100644 index 0000000000..85219e9413 --- /dev/null +++ b/v2/pkg/engine/resolve/loader_dataflow.go @@ -0,0 +1,486 @@ +package resolve + +import ( + "context" + stderrors "errors" + "maps" + "slices" + + "github.com/pkg/errors" + "github.com/wundergraph/astjson" + "github.com/wundergraph/graphql-go-tools/v2/pkg/ast" +) + +// resolveDataflow executes the fetch DAG WITHOUT per-wave barriers. Each fetch's +// network load starts as soon as its OWN dependencies (DependsOnFetchIDs) have +// merged, rather than waiting for the slowest fetch of the previous Parallel wave +// (which resolveSerial/resolveParallel do via g.Wait()). Under skewed subgraph +// latency this collapses the wall-clock toward the true critical path. +// +// INVARIANTS (each falsifiable; codex is explicitly tasked to attack them): +// +// 1. Single-coordinator arena ownership: EVERY arena read (selectItemsForPath, +// itemsData, input-template rendering, batch dedup) and EVERY arena +// write (response parse, merge) runs on this coordinator goroutine, via +// preparePhase (prepare) and mergeResult (merge). +// The spawned worker closure contains ONLY +// executeSourceLoad, which is arena-free (no l.jsonArena / l.resolvable / +// selectItemsForPath / itemsData / MustParse references — greppable). +// No mutex is needed; there is nothing concurrent to lock against. +// +// 2. Deterministic error ordering via swap-capture staging: the three audited +// order-bearing sinks — l.resolvable.errors (all appends inside mergeErrors, +// addApolloRouterCompatibilityError, renderErrorsFailedDeps/FailedToFetch/ +// StatusFallback, renderAuthorizationRejectedErrors, +// renderRateLimitRejectedErrors, all reachable only from mergeResult / +// mergeErrors), l.ctx.subgraphErrors +// (ctx.appendSubgraphErrors, same reachability), and +// l.resolvable.subgraphExtensions (mergeResult) — are +// nil-swapped around each merge and replayed in ascending LEAF order after +// the drain. Both swap targets are nil-gated lazy-init +// (Resolvable.ensureErrorsInitialized, Context.appendSubgraphErrors), so the +// swap is transparent to the merge code: ZERO diff to the loader.go +// error paths. If a fourth order-bearing sink exists, this +// staging silently misses it — that is the falsification target. +// +// 3. Leaf order == wave merge order: leaves are indexed by their position in +// collectDataflowLeaves' depth-first flatten, which equals the wave +// executor's merge order (resolveSerial walks Sequence children in order; +// resolveParallel merges by node index after g.Wait()). Flushing staged +// sinks and OnFinished hooks in ascending leaf index therefore reproduces +// the wave executor's error/extension array order byte-for-byte. +// Flags written during merge (taintedObjs, skipValueCompletion) stay +// UN-staged: they are order-free, and dependents' selectItemsForPath must +// observe parent taint (DAG order guarantees parents merged first). +// +// 4. Tools lifetime: batch input bytes live on the pooled res.tools +// arena, so tools are collected into toolsToPut AT PREPARE TIME (abort-safe) +// and Put ONLY after the drain loop exits — on every path including +// prepare-fatal and merge-fatal drains. Put(nil) is a no-op. +// +// 5. Exactly-once completion: ch is buffered to n and every spawned worker +// sends exactly once, unconditionally; the coordinator never returns before +// inflight == 0. No goroutine can leak or block, even after cancel(). +// cancel() propagates to in-flight HTTP via the derived ctx. A worker panic +// crashes the process exactly like the wave executor's errgroup — no +// recover(), by parity. +// +// 6. Fatal selection: prepare/merge errors record per-leaf fatals and cancel(); +// dispatching stops, in-flight loads drain, staged sinks still flush (hook/ +// telemetry consistency), and the LOWEST-leaf-index fatal is returned — +// deterministic, though not guaranteed bitwise wave-identical when multiple +// fetches fail fatally (the wave executor short-circuits dispatch instead; +// its own errgroup error selection was already completion-order +// nondeterministic). The flush may therefore fire OnFinished for staged +// leaves the wave executor would never have merged after ITS first merge +// error — a deterministic superset, telemetry-only (codex P2, accepted). +// Acceptable: a fatal LoadGraphQLResponseData error discards the resolvable. +// +// 7. OnFinished hooks fire during the flush, in leaf order — each hook sees +// ctx.subgraphErrors in exactly the post-leaf-i state the wave executor +// would show it. Wall-clock timing is later than the wave executor's +// per-merge hooks; cosmo's hooks consume the already-captured responseInfo, +// so this is timing-only. +// +// 8. Complete-DependsOnFetchIDs: taint visibility for concurrently dispatched +// fetches relies on every fetch listing ALL fetches whose data it reads — +// the same invariant resolveParallel/createParallelNodes already require. +// +// 9. Pre-fetch hook call order (Authorizer/RateLimiter via validatePreFetch, +// which runs at PREPARE time and whose implementations can render +// accumulated state into response extensions): every dispatch pops the +// GLOBAL leaf-index minimum among ready fetches (not FIFO — an inline +// skipLoad completion may enqueue a lower-leaf dependent behind an +// already-queued sibling), so hook call order equals the wave executor's +// spawn order whenever the DAG permits it. Residual contract: when a +// WORKER completion makes a lower-leaf fetch ready after a higher-leaf +// fetch was already dispatched, hook calls interleave earlier than the +// wave BARRIER would allow — the wave executor already calls these hooks +// concurrently (unordered) WITHIN a wave, so implementations must already +// be order-tolerant; dataflow extends that tolerance requirement across +// waves. Order-SENSITIVE extension renderers are not byte-stable under +// ENGINE_ENABLE_DATAFLOW. +// +// validatePreFetch (authorization + rate limiting) runs inside the prepare +// functions on the coordinator, exactly as in the nested 3-phase protocol +// (upstream holds it under mergeMu there). Rate-limit I/O therefore serializes +// on the coordinator; if that ever matters, splitting it back into the worker +// is the named lever — do not do it speculatively. +// +// Eligible only for queries whose fetch tree is the flat federation shape +// (see collectDataflowLeaves) with UNIQUE FetchIDs forming a DAG. Mutations, +// subscriptions, non-unique FetchIDs, cyclic deps, nested (schedule-tree) +// plans, or any unexpected node kind fall back to the wave executor. +// +// STATUS: experimental, default-off (ResolverOptions.EnableDataflowExecution / +// ENGINE_ENABLE_DATAFLOW). Recovers up to ~37% wall-clock under skewed +// subgraph latency. +func (l *Loader) resolveDataflow(root *FetchTreeNode) error { + leaves, ok := collectDataflowLeaves(root) + if !ok { + return l.resolveFetchNode(root) + } + n := len(leaves) + switch n { + case 0: + return nil + case 1: + return l.resolveSingle(leaves[0].Item) + } + + byID := make(map[int]*FetchTreeNode, n) + leafIndexByID := make(map[int]int, n) + for i, lf := range leaves { + id := lf.Item.Fetch.Dependencies().FetchID + if _, dup := byID[id]; dup { + // Non-unique FetchIDs mean ordering is expressed by the tree STRUCTURE, + // not by the FetchID dependency edges (e.g. hand-built plans, or any plan + // where the planner did not assign distinct FetchIDs). The dataflow + // scheduler keys solely on FetchID deps, so it cannot honor structural + // ordering — fall back to the wave executor. Real planner output always + // has unique FetchIDs with complete deps (createParallelNodes groups BY + // those deps), so production plans take the dataflow path. + return l.resolveFetchNode(root) + } + byID[id] = lf + leafIndexByID[id] = i + } + // Count only dependencies that actually exist in this leaf set, and record the + // reverse edges so a completed fetch can unblock its dependents. + remaining := make(map[int]int, n) + dependents := make(map[int][]int, n) + for _, lf := range leaves { + id := lf.Item.Fetch.Dependencies().FetchID + cnt := 0 + for _, dep := range lf.Item.Fetch.Dependencies().DependsOnFetchIDs { + if _, exists := byID[dep]; exists { + cnt++ + dependents[dep] = append(dependents[dep], id) + } + } + remaining[id] = cnt + } + + // Schedulability pre-check (Kahn's algorithm on a copy of the in-degrees). If the + // in-set dependency graph is not a DAG, a cycle leaves some fetches permanently + // unschedulable; without this guard the coordinator loop would exit with those + // fetches never dispatched and return a silently-incomplete response (codex P1). + // Real planner output is always a DAG, so this falls back rather than executes. + { + indeg := make(map[int]int, n) + maps.Copy(indeg, remaining) + queue := make([]int, 0, n) + for id, d := range indeg { + if d == 0 { + queue = append(queue, id) + } + } + scheduled := 0 + for len(queue) > 0 { + id := queue[0] + queue = queue[1:] + scheduled++ + for _, dep := range dependents[id] { + indeg[dep]-- + if indeg[dep] == 0 { + queue = append(queue, dep) + } + } + } + if scheduled != n { + // Not a DAG (cycle / unschedulable). Fall back before merging anything. + return l.resolveFetchNode(root) + } + } + + // mergeStage holds one leaf's captured error sinks until the leaf-order flush + // (invariants 2 and 3). + type mergeStage struct { + errors *astjson.Value + extensions []*astjson.Object + subgraphErrs map[string]error + res *result + merged bool + } + stages := make([]mergeStage, n) + + type completion struct { + id int + idx int + } + + ctx, cancel := context.WithCancel(l.ctx.ctx) + defer cancel() + ch := make(chan completion, n) + inflight := 0 + toolsToPut := make([]*batchEntityTools, 0, n) + preparedByIdx := make([]*preparedFetch, n) + + fatalByIdx := make([]error, n) + hasFatal := false + recordFatal := func(idx int, err error) { + if fatalByIdx[idx] == nil { + fatalByIdx[idx] = err + } + hasFatal = true + cancel() + } + + // stagedMerge runs the wave executor's merge dispatch for one leaf on the + // coordinator, with the three audited sinks nil-swapped and captured + // (invariant 2). callOnFinished is deliberately NOT called here — it moves to + // the leaf-order flush (invariant 7). + stagedMerge := func(idx int, p *preparedFetch) error { + savedErrors := l.resolvable.errors + savedExtensions := l.resolvable.subgraphExtensions + savedSubgraphErrors := l.ctx.subgraphErrors + l.resolvable.errors = nil + l.resolvable.subgraphExtensions = nil + l.ctx.subgraphErrors = nil + + var err error + switch { + case p.res.nestedMergeItems != nil: + // Vestigial in this fork — nestedMergeItems is never assigned — but kept + // for parity with resolveParallel's merge dispatch (surgical-changes rule). + for j := range p.res.nestedMergeItems { + if err = l.mergeResult(p.item, p.res.nestedMergeItems[j], p.items[j:j+1]); err != nil { + break + } + } + default: + err = l.mergeResult(p.item, p.res, p.items) + } + + stages[idx] = mergeStage{ + errors: l.resolvable.errors, + extensions: l.resolvable.subgraphExtensions, + subgraphErrs: l.ctx.subgraphErrors, + res: p.res, + merged: true, + } + l.resolvable.errors = savedErrors + l.resolvable.subgraphExtensions = savedExtensions + l.ctx.subgraphErrors = savedSubgraphErrors + return err + } + + // Dispatch ordering is LEAF order (= tree order = the wave executor's spawn + // order), NOT FetchID order: validatePreFetch calls user-supplied + // Authorizer/RateLimiter hooks at prepare time, and those hooks can render + // order-sensitive accumulated state into response extensions (codex P1 on + // this hardening). Leaf-ordered seeding and unblock batches reproduce the + // wave executor's call order whenever the DAG permits. + byLeafIndex := func(a, b int) int { return leafIndexByID[a] - leafIndexByID[b] } + var ready []int + unblock := func(id int) { + for _, dep := range dependents[id] { + remaining[dep]-- + if remaining[dep] == 0 { + ready = append(ready, dep) + } + } + } + + // dispatchOne prepares one fetch on the coordinator (all arena reads, + // invariant 1) and either completes it inline (skip paths) or hands ONLY the + // arena-free load to a worker goroutine. + dispatchOne := func(id int) { + node := byID[id] + idx := leafIndexByID[id] + p, err := l.preparePhase(node.Item) + if p != nil && p.res.tools != nil { + // Collect at prepare time so abort paths still Put (invariant 4). + toolsToPut = append(toolsToPut, p.res.tools) + } + if err != nil { + recordFatal(idx, err) + return + } + if p == nil { + // Unknown fetch kind: the wave executor's resolveSingle default case + // performs no load and no merge. Just unblock dependents. + unblock(id) + return + } + if p.skipLoad { + // Skip paths still merge (fetchSkipped / rendered-error / denial state + // lives on res), inline on the coordinator — no goroutine, no recursion. + if mErr := stagedMerge(idx, p); mErr != nil { + recordFatal(idx, mErr) + return + } + unblock(id) + return + } + preparedByIdx[idx] = p + inflight++ + go func() { + // Worker: arena-free by invariant 1. executeSourceLoad stores failures + // in res.err (never a Go return); merge renders them deterministically. + l.executeSourceLoad(ctx, p.item, p.source, p.input, p.res, p.trace) + ch <- completion{id: id, idx: idx} + }() + } + + // Seed every dependency-free fetch. + for id, r := range remaining { + if r == 0 { + ready = append(ready, id) + } + } + + for { + for len(ready) > 0 && !hasFatal { + // Pop the GLOBAL leaf-index minimum, not FIFO: an inline skipLoad + // completion can enqueue a lower-leaf dependent behind an already-queued + // higher-leaf sibling, and FIFO would call its pre-fetch hooks out of + // wave order even though the DAG permits wave order (codex P1 round 2). + // n is small; a sort per pop is cheaper than a heap at this size. + slices.SortFunc(ready, byLeafIndex) + id := ready[0] + ready = ready[1:] + dispatchOne(id) + } + if inflight == 0 { + break + } + c := <-ch + inflight-- + if hasFatal { + continue // draining after a fatal + } + if err := stagedMerge(c.idx, preparedByIdx[c.idx]); err != nil { + recordFatal(c.idx, err) + continue + } + unblock(c.id) + } + + // Flush staged sinks + OnFinished hooks in ascending leaf index — the wave + // executor's merge order (invariants 3 and 7). Runs even after a fatal, for + // hook/telemetry consistency; the resolvable is discarded on fatal anyway. + for i := range stages { + st := &stages[i] + if !st.merged { + continue + } + if st.errors != nil && len(st.errors.GetArray()) > 0 { + l.resolvable.ensureErrorsInitialized() + l.resolvable.errors.AppendArrayItems(l.jsonArena, st.errors) + } + if len(st.extensions) > 0 { + l.resolvable.subgraphExtensions = append(l.resolvable.subgraphExtensions, st.extensions...) + } + if len(st.subgraphErrs) > 0 { + if l.ctx.subgraphErrors == nil { + l.ctx.subgraphErrors = make(map[string]error, len(st.subgraphErrs)) + } + // Sorted for determinism (one key per fetch in practice). When the key is + // new, the captured value is assigned directly — exact wave parity, since + // the capture started from nil exactly like a fresh map entry. When the + // key repeats across leaves, errors.Join nests the captured chain one + // level deeper than the wave executor's sequential appends; the flattened + // Error() string and errors.Is/As behavior are identical. + for _, k := range slices.Sorted(maps.Keys(st.subgraphErrs)) { + if existing, exists := l.ctx.subgraphErrors[k]; exists { + l.ctx.subgraphErrors[k] = stderrors.Join(existing, st.subgraphErrs[k]) + } else { + l.ctx.subgraphErrors[k] = st.subgraphErrs[k] + } + } + } + // The hook sees ctx.subgraphErrors in exactly the post-leaf-i state the + // wave executor would show it (invariant 7). + if st.res.nestedMergeItems != nil { + for j := range st.res.nestedMergeItems { + l.callOnFinished(st.res.nestedMergeItems[j]) + } + } else { + l.callOnFinished(st.res) + } + } + + for _, t := range toolsToPut { + batchEntityToolPool.Put(t) + } + if hasFatal { + for i := range fatalByIdx { + if fatalByIdx[i] != nil { + // Lowest-leaf-index fatal: deterministic regardless of completion order + // (invariant 6). + return errors.WithStack(fatalByIdx[i]) + } + } + } + return nil +} + +// collectDataflowLeaves returns the Single leaves of a FLAT fetch tree: nil, a +// bare Single, or a Sequence whose children are Single or Parallel-of-Single — +// exactly the shape createParallelNodes emits. Anything else, in particular the +// NESTED Parallel(Sequence(...)) trees built by the schedule-tree processor +// (WithBuildScheduleTree), reports ok=false and resolveDataflow falls back to the +// wave executor, which handles nested trees under mergeMu. This structural guard +// is what makes ENGINE_ENABLE_DATAFLOW + ENGINE_ENABLE_SCHEDULE_TREE safe to +// combine: the dataflow scheduler keys solely on FetchID dependency edges and +// would otherwise ignore the schedule tree's structural ordering. +func collectDataflowLeaves(node *FetchTreeNode) ([]*FetchTreeNode, bool) { + if node == nil { + return nil, true + } + switch node.Kind { + case FetchTreeNodeKindSingle: + return singleDataflowLeaf(node) + case FetchTreeNodeKindSequence: + out := make([]*FetchTreeNode, 0, len(node.ChildNodes)) + for _, child := range node.ChildNodes { + if child == nil { + return nil, false + } + switch child.Kind { + case FetchTreeNodeKindSingle: + leaf, ok := singleDataflowLeaf(child) + if !ok { + return nil, false + } + out = append(out, leaf...) + case FetchTreeNodeKindParallel: + for _, pc := range child.ChildNodes { + if pc == nil || pc.Kind != FetchTreeNodeKindSingle { + return nil, false + } + leaf, ok := singleDataflowLeaf(pc) + if !ok { + return nil, false + } + out = append(out, leaf...) + } + default: + return nil, false + } + } + return out, true + default: + return nil, false + } +} + +func singleDataflowLeaf(node *FetchTreeNode) ([]*FetchTreeNode, bool) { + if node.Item == nil || node.Item.Fetch == nil { + return nil, false + } + return []*FetchTreeNode{node}, true +} + +// dataflowEligibleOperation reports whether the current operation may use the +// dataflow executor. Mutations are serialized by side effect (ordering not +// captured in DependsOnFetchIDs) and subscriptions resolve differently, so only +// queries are eligible. response.Info is authoritative when present; the +// context value (which defaults to Query when unset) is the fallback. +func (l *Loader) dataflowEligibleOperation() bool { + if l.info != nil { + return l.info.OperationType == ast.OperationTypeQuery + } + return GetOperationTypeFromContext(l.ctx.ctx) == ast.OperationTypeQuery +} diff --git a/v2/pkg/engine/resolve/loader_dataflow_test.go b/v2/pkg/engine/resolve/loader_dataflow_test.go new file mode 100644 index 0000000000..f766971f4d --- /dev/null +++ b/v2/pkg/engine/resolve/loader_dataflow_test.go @@ -0,0 +1,767 @@ +package resolve + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "math/rand" + "net/http" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/require" + "go.uber.org/goleak" + + "github.com/wundergraph/graphql-go-tools/v2/pkg/ast" + "github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/httpclient" +) + +// The tests in this file verify the dataflow executor against the wave executor +// by running the IDENTICAL plan through both and asserting full-output equality +// (rule: assert.Equal on complete bytes, never Contains). They use the direct +// Loader harness (no Resolver background goroutines) so goleak can verify the +// dataflow coordinator leaks nothing. +// +// Batch entity fetches are exercised through the dataflow path end-to-end +// by the router-level byte-identity gate (perf-analysis/q_byteid.sh runs the +// full federated query corpus with ENGINE_ENABLE_DATAFLOW=true); the unit tests +// here pin the dataflow-specific machinery: scheduling, coordinator-owned +// prepare, swap-capture error staging, leaf-order flush, fallbacks, +// cancellation, and goroutine hygiene. + +type delayDataSource struct { + response []byte + delay time.Duration + loadErr error + + mu sync.Mutex + inputs []string + + loadCalls int64 +} + +func newDelayDataSource(response string, delay time.Duration) *delayDataSource { + return &delayDataSource{response: []byte(response), delay: delay} +} + +func (d *delayDataSource) Load(ctx context.Context, _ http.Header, input []byte) ([]byte, error) { + d.mu.Lock() + d.loadCalls++ + d.inputs = append(d.inputs, string(input)) + d.mu.Unlock() + if d.delay > 0 { + select { + case <-time.After(d.delay): + case <-ctx.Done(): + return nil, ctx.Err() + } + } + if d.loadErr != nil { + return nil, d.loadErr + } + return d.response, nil +} + +func (d *delayDataSource) LoadWithFiles(ctx context.Context, headers http.Header, input []byte, _ []*httpclient.FileUpload) ([]byte, error) { + return d.Load(ctx, headers, input) +} + +func (d *delayDataSource) requireInputs(t *testing.T, expected ...string) { + t.Helper() + d.mu.Lock() + defer d.mu.Unlock() + require.Equal(t, expected, d.inputs) +} + +func (d *delayDataSource) calls() int64 { + d.mu.Lock() + defer d.mu.Unlock() + return d.loadCalls +} + +// dataflowFetch builds a SingleFetch with a static input, FetchInfo (data source +// name feeds error messages) and explicit FetchID/dependency edges. +func dataflowFetch(ds DataSource, id int, deps []int, name, staticInput string) *SingleFetch { + f := nestedParallelSingleFetchWithInfo(ds, staticInput, name, name) + f.FetchDependencies.FetchID = id + f.FetchDependencies.DependsOnFetchIDs = deps + return f +} + +// dataflowFetchReading builds a SingleFetch whose input template renders the +// given fields from the parent-merged response data (arena reads at prepare). +func dataflowFetchReading(ds DataSource, id int, deps []int, name string, fields ...string) *SingleFetch { + f := nestedParallelSingleFetchWithInfoAndTemplate(ds, nestedParallelInputForFields(fields...), name, name) + f.FetchDependencies.FetchID = id + f.FetchDependencies.DependsOnFetchIDs = deps + return f +} + +// runDataflowScenario executes the response plan through the direct Loader +// harness and returns the SERIALIZED response (the real query-order renderer, +// resolvable.Resolve — raw arena insertion order is merge-order-dependent by +// design and is not part of the byte-identity contract) plus the load error. +func runDataflowScenario(t *testing.T, reqCtx context.Context, enableDataflow bool, opType ast.OperationType, response *GraphQLResponse, configure func(*Loader, *Context)) (string, error) { + t.Helper() + ctx := NewContext(reqCtx) + ctx.ExecutionOptions.DisableSubgraphRequestDeduplication = true + resolvable := NewResolvable(nil, ResolvableOptions{}) + loader := &Loader{ + enableDataflow: enableDataflow, + propagateSubgraphErrors: true, + propagateSubgraphStatusCodes: true, + subgraphErrorPropagationMode: SubgraphErrorPropagationModeWrapped, + // pass-through mode filters subgraph error fields by allowlist; mirror the + // production default of allowing "message" + allowedSubgraphErrorFields: map[string]struct{}{"message": {}}, + } + if configure != nil { + configure(loader, ctx) + } + require.NoError(t, resolvable.Init(ctx, nil, opType)) + err := loader.LoadGraphQLResponseData(ctx, response, resolvable) + buf := &bytes.Buffer{} + require.NoError(t, resolvable.Resolve(ctx.ctx, response.Data, response.Fetches, buf)) + return buf.String(), err +} + +// runBothExecutors builds a fresh plan per executor (fresh datasources, fresh +// channels) and returns both outputs; both runs must succeed. +func runBothExecutors(t *testing.T, build func() *GraphQLResponse, configure func(*Loader, *Context)) (wave string, dataflow string) { + t.Helper() + wave, waveErr := runDataflowScenario(t, context.Background(), false, ast.OperationTypeQuery, build(), configure) + require.NoError(t, waveErr) + dataflow, dataflowErr := runDataflowScenario(t, context.Background(), true, ast.OperationTypeQuery, build(), configure) + require.NoError(t, dataflowErr) + return wave, dataflow +} + +func TestDataflowByteIdenticalToWave(t *testing.T) { + t.Run("two waves", func(t *testing.T) { + build := func() *GraphQLResponse { + r := newDelayDataSource(`{"data":{"r":"R"}}`, 0) + b := newDelayDataSource(`{"data":{"b":"B"}}`, 4*time.Millisecond) + c := newDelayDataSource(`{"data":{"c":"C"}}`, time.Millisecond) + return &GraphQLResponse{ + Info: &GraphQLResponseInfo{OperationType: ast.OperationTypeQuery}, + Fetches: Sequence( + Single(dataflowFetch(r, 0, nil, "R", `{"fetch":"r"}`)), + Parallel( + Single(dataflowFetch(b, 1, []int{0}, "B", `{"fetch":"b"}`)), + Single(dataflowFetch(c, 2, []int{0}, "C", `{"fetch":"c"}`)), + ), + ), + Data: nestedParallelData("r", "b", "c"), + } + } + wave, dataflow := runBothExecutors(t, build, nil) + require.Equal(t, `{"data":{"r":"R","b":"B","c":"C"}}`, wave) + require.Equal(t, wave, dataflow) + }) + + t.Run("diamond reads merged data", func(t *testing.T) { + var dWave, dDataflow *delayDataSource + build := func() *delayDataSource { + return newDelayDataSource(`{"data":{"d":"D"}}`, 0) + } + buildResponse := func(d *delayDataSource) *GraphQLResponse { + r := newDelayDataSource(`{"data":{"r":"R"}}`, 0) + b := newDelayDataSource(`{"data":{"b":"B"}}`, 3*time.Millisecond) + c := newDelayDataSource(`{"data":{"c":"C"}}`, time.Millisecond) + return &GraphQLResponse{ + Info: &GraphQLResponseInfo{OperationType: ast.OperationTypeQuery}, + Fetches: Sequence( + Single(dataflowFetch(r, 0, nil, "R", `{"fetch":"r"}`)), + Parallel( + Single(dataflowFetchReading(b, 1, []int{0}, "B", "r")), + Single(dataflowFetchReading(c, 2, []int{0}, "C", "r")), + ), + Single(dataflowFetchReading(d, 3, []int{1, 2}, "D", "b", "c")), + ), + Data: nestedParallelData("r", "b", "c", "d"), + } + } + dWave = build() + wave, waveErr := runDataflowScenario(t, context.Background(), false, ast.OperationTypeQuery, buildResponse(dWave), nil) + require.NoError(t, waveErr) + dDataflow = build() + dataflow, dataflowErr := runDataflowScenario(t, context.Background(), true, ast.OperationTypeQuery, buildResponse(dDataflow), nil) + require.NoError(t, dataflowErr) + require.Equal(t, `{"data":{"r":"R","b":"B","c":"C","d":"D"}}`, wave) + require.Equal(t, wave, dataflow) + // d's input renders from BOTH parents' merged data — identical in both modes. + dWave.requireInputs(t, `{"b":"B","c":"C"}`) + dDataflow.requireInputs(t, `{"b":"B","c":"C"}`) + }) + + t.Run("fan-out 8", func(t *testing.T) { + delays := []time.Duration{5, 0, 3, 1, 4, 2, 0, 1} + build := func() *GraphQLResponse { + r := newDelayDataSource(`{"data":{"r":"R"}}`, 0) + children := make([]*FetchTreeNode, 8) + fields := []string{"r"} + for i := range 8 { + name := fmt.Sprintf("f%d", i) + ds := newDelayDataSource(fmt.Sprintf(`{"data":{"%s":"V%d"}}`, name, i), delays[i]*time.Millisecond) + children[i] = Single(dataflowFetch(ds, i+1, []int{0}, name, fmt.Sprintf(`{"fetch":"%s"}`, name))) + fields = append(fields, name) + } + return &GraphQLResponse{ + Info: &GraphQLResponseInfo{OperationType: ast.OperationTypeQuery}, + Fetches: Sequence( + append([]*FetchTreeNode{Single(dataflowFetch(r, 0, nil, "R", `{"fetch":"r"}`))}, + Parallel(children...))..., + ), + Data: nestedParallelData(fields...), + } + } + wave, dataflow := runBothExecutors(t, build, nil) + require.Equal(t, `{"data":{"r":"R","f0":"V0","f1":"V1","f2":"V2","f3":"V3","f4":"V4","f5":"V5","f6":"V6","f7":"V7"}}`, wave) + require.Equal(t, wave, dataflow) + }) +} + +func TestDataflowRaceStress(t *testing.T) { + // 12-fetch DAG with seeded random per-fetch delays. Run repeatedly so the + // dispatch/merge interleavings vary; -race is the gate, full-byte equality + // the assertion. + rng := rand.New(rand.NewSource(0x5eed)) + for iter := range 30 { + d := func() time.Duration { return time.Duration(rng.Intn(4)) * time.Millisecond } + build := func() *GraphQLResponse { + r := newDelayDataSource(`{"data":{"r":"R"}}`, d()) + mk := func(name string, id int, deps []int, readFields ...string) *FetchTreeNode { + ds := newDelayDataSource(fmt.Sprintf(`{"data":{"%s":"%s"}}`, name, name), d()) + if len(readFields) > 0 { + return Single(dataflowFetchReading(ds, id, deps, name, readFields...)) + } + return Single(dataflowFetch(ds, id, deps, name, fmt.Sprintf(`{"fetch":"%s"}`, name))) + } + return &GraphQLResponse{ + Info: &GraphQLResponseInfo{OperationType: ast.OperationTypeQuery}, + Fetches: Sequence( + Single(dataflowFetch(r, 0, nil, "R", `{"fetch":"r"}`)), + Parallel( + mk("a1", 1, []int{0}, "r"), + mk("a2", 2, []int{0}, "r"), + mk("a3", 3, []int{0}, "r"), + ), + Parallel( + mk("b1", 4, []int{1}, "a1"), + mk("b2", 5, []int{1}), + mk("b3", 6, []int{2}, "a2"), + mk("b4", 7, []int{2}), + mk("b5", 8, []int{3}, "a3"), + mk("b6", 9, []int{3}), + ), + Parallel( + mk("c1", 10, []int{4, 5}, "b1", "b2"), + mk("c2", 11, []int{8, 9}, "b5", "b6"), + ), + ), + Data: nestedParallelData("r", "a1", "a2", "a3", "b1", "b2", "b3", "b4", "b5", "b6", "c1", "c2"), + } + } + wave, dataflow := runBothExecutors(t, build, nil) + require.Equal(t, `{"data":{"r":"R","a1":"a1","a2":"a2","a3":"a3","b1":"b1","b2":"b2","b3":"b3","b4":"b4","b5":"b5","b6":"b6","c1":"c1","c2":"c2"}}`, wave, "iteration %d", iter) + require.Equal(t, wave, dataflow, "iteration %d", iter) + } +} + +// TestDataflowPrepareMergeOverlap reconstructs the ORIGINAL arena race the +// hardening eliminated: sibling A merges into the root object while sibling B's +// input would render from the same object. With coordinator-owned prepare, B's +// input is rendered at dispatch time on the coordinator — before either sibling +// completes — so no arena access can overlap a merge. The gate is -race +// cleanliness plus full-byte equality. +func TestDataflowPrepareMergeOverlap(t *testing.T) { + for range 20 { + var bWave, bDataflow *delayDataSource + buildWith := func(b *delayDataSource) func() *GraphQLResponse { + return func() *GraphQLResponse { + p := newDelayDataSource(`{"data":{"p":"P"}}`, 0) + a := newDelayDataSource(`{"data":{"a":"A"}}`, 0) + return &GraphQLResponse{ + Info: &GraphQLResponseInfo{OperationType: ast.OperationTypeQuery}, + Fetches: Sequence( + Single(dataflowFetch(p, 0, nil, "P", `{"fetch":"p"}`)), + Parallel( + Single(dataflowFetchReading(a, 1, []int{0}, "A", "p")), + Single(dataflowFetchReading(b, 2, []int{0}, "B", "p")), + ), + ), + Data: nestedParallelData("p", "a", "b"), + } + } + } + bWave = newDelayDataSource(`{"data":{"b":"B"}}`, 4*time.Millisecond) + wave, waveErr := runDataflowScenario(t, context.Background(), false, ast.OperationTypeQuery, buildWith(bWave)(), nil) + require.NoError(t, waveErr) + bDataflow = newDelayDataSource(`{"data":{"b":"B"}}`, 4*time.Millisecond) + dataflow, dataflowErr := runDataflowScenario(t, context.Background(), true, ast.OperationTypeQuery, buildWith(bDataflow)(), nil) + require.NoError(t, dataflowErr) + require.Equal(t, `{"data":{"p":"P","a":"A","b":"B"}}`, wave) + require.Equal(t, wave, dataflow) + bWave.requireInputs(t, `{"p":"P"}`) + bDataflow.requireInputs(t, `{"p":"P"}`) + } +} + +// TestDataflowErrorOrderDeterminism inverts completion order against plan order +// (first leaf slowest) and asserts the errors array is byte-identical to the +// wave executor — the property the swap-capture staging plus leaf-order flush +// exists to guarantee. +func TestDataflowErrorOrderDeterminism(t *testing.T) { + build := func() *GraphQLResponse { + p := newDelayDataSource(`{"data":{"p":"P"}}`, 0) + mkErr := func(name string, id int, delay time.Duration) *FetchTreeNode { + ds := newDelayDataSource(fmt.Sprintf(`{"errors":[{"message":"%s exploded"}],"data":{"%s":null}}`, name, name), delay) + f := dataflowFetch(ds, id, []int{0}, name, fmt.Sprintf(`{"fetch":"%s"}`, name)) + f.PostProcessing.SelectResponseErrorsPath = []string{"errors"} + return SingleWithPath(f, "query."+name) + } + return &GraphQLResponse{ + Info: &GraphQLResponseInfo{OperationType: ast.OperationTypeQuery}, + Fetches: Sequence( + Single(dataflowFetch(p, 0, nil, "P", `{"fetch":"p"}`)), + Parallel( + mkErr("e1", 1, 6*time.Millisecond), + mkErr("e2", 2, 3*time.Millisecond), + mkErr("e3", 3, 0), + ), + ), + Data: nestedParallelNullableData("p", "e1", "e2", "e3"), + } + } + + for _, mode := range []struct { + name string + mode SubgraphErrorPropagationMode + }{ + {name: "wrap", mode: SubgraphErrorPropagationModeWrapped}, + {name: "pass-through", mode: SubgraphErrorPropagationModePassThrough}, + } { + t.Run(mode.name, func(t *testing.T) { + configure := func(l *Loader, _ *Context) { + l.subgraphErrorPropagationMode = mode.mode + } + wave, dataflow := runBothExecutors(t, build, configure) + // Completion order is e3, e2, e1 — the errors array must still be in + // PLAN order e1, e2, e3, exactly as the wave executor emits it. + require.Equal(t, wave, dataflow) + switch mode.mode { + case SubgraphErrorPropagationModeWrapped: + require.Equal(t, `{"errors":[{"message":"Failed to fetch from Subgraph 'e1' at Path 'query.e1'.","extensions":{"errors":[{"message":"e1 exploded"}]}},{"message":"Failed to fetch from Subgraph 'e2' at Path 'query.e2'.","extensions":{"errors":[{"message":"e2 exploded"}]}},{"message":"Failed to fetch from Subgraph 'e3' at Path 'query.e3'.","extensions":{"errors":[{"message":"e3 exploded"}]}}],"data":{"p":"P","e1":null,"e2":null,"e3":null}}`, dataflow) + case SubgraphErrorPropagationModePassThrough: + require.Equal(t, `{"errors":[{"message":"e1 exploded"},{"message":"e2 exploded"},{"message":"e3 exploded"}],"data":{"p":"P","e1":null,"e2":null,"e3":null}}`, dataflow) + } + }) + } +} + +type stubRateLimiter struct { + denyName string + errName string +} + +func (s *stubRateLimiter) RateLimitPreFetch(_ *Context, info *FetchInfo, _ json.RawMessage) (*RateLimitDeny, error) { + if info != nil && info.DataSourceName == s.errName { + return nil, errors.New("rate limiter exploded") + } + if info != nil && info.DataSourceName == s.denyName { + return &RateLimitDeny{Reason: "test denied"}, nil + } + return nil, nil +} + +func (s *stubRateLimiter) RenderResponseExtension(_ *Context, _ io.Writer) error { + return nil +} + +// TestDataflowRateLimitDenyAbortPath: a mid-DAG fetch is denied at prepare +// (skipLoad), its rendered denial error goes through the staged merge, and its +// dependent still dispatches — byte-identical to the wave executor. +func TestDataflowRateLimitDenyAbortPath(t *testing.T) { + var dWave, dDataflow *delayDataSource + buildWith := func(d *delayDataSource) *GraphQLResponse { + p := newDelayDataSource(`{"data":{"p":"P"}}`, 0) + m := newDelayDataSource(`{"data":{"m":"M"}}`, 0) + mFetch := dataflowFetch(m, 1, []int{0}, "M", `{"fetch":"m"}`) + return &GraphQLResponse{ + Info: &GraphQLResponseInfo{OperationType: ast.OperationTypeQuery}, + Fetches: Sequence( + Single(dataflowFetch(p, 0, nil, "P", `{"fetch":"p"}`)), + SingleWithPath(mFetch, "query.m"), + // static input: the dependent must still dispatch and LOAD after its + // parent was denied (denial nulls m, but d does not read it) + Single(dataflowFetch(d, 2, []int{1}, "D", `{"fetch":"d"}`)), + ), + Data: nestedParallelNullableData("p", "m", "d"), + } + } + configure := func(_ *Loader, ctx *Context) { + ctx.RateLimitOptions = RateLimitOptions{Enable: true} + ctx.SetRateLimiter(&stubRateLimiter{denyName: "M"}) + } + + dWave = newDelayDataSource(`{"data":{"d":"D"}}`, 0) + wave, waveErr := runDataflowScenario(t, context.Background(), false, ast.OperationTypeQuery, buildWith(dWave), configure) + require.NoError(t, waveErr) + dDataflow = newDelayDataSource(`{"data":{"d":"D"}}`, 0) + dataflow, dataflowErr := runDataflowScenario(t, context.Background(), true, ast.OperationTypeQuery, buildWith(dDataflow), configure) + require.NoError(t, dataflowErr) + + require.Equal(t, wave, dataflow) + require.Equal(t, `{"errors":[{"message":"Rate limit exceeded for Subgraph 'M' at Path 'query.m', Reason: test denied."}],"data":{"p":"P","m":null,"d":"D"}}`, wave) + // The denied fetch loads nothing; the dependent still runs in both modes. + require.Equal(t, int64(1), dWave.calls()) + require.Equal(t, int64(1), dDataflow.calls()) + dWave.requireInputs(t, `{"fetch":"d"}`) + dDataflow.requireInputs(t, `{"fetch":"d"}`) +} + +type recordingRateLimiter struct { + mu sync.Mutex + seen []string +} + +func (r *recordingRateLimiter) RateLimitPreFetch(_ *Context, info *FetchInfo, _ json.RawMessage) (*RateLimitDeny, error) { + r.mu.Lock() + defer r.mu.Unlock() + if info != nil { + r.seen = append(r.seen, info.DataSourceName) + } + return nil, nil +} + +func (r *recordingRateLimiter) RenderResponseExtension(_ *Context, _ io.Writer) error { + return nil +} + +func (r *recordingRateLimiter) order() []string { + r.mu.Lock() + defer r.mu.Unlock() + return r.seen +} + +// denyThenRecordRateLimiter records every pre-fetch call (via inner) and denies +// the named data source. +type denyThenRecordRateLimiter struct { + deny string + inner *recordingRateLimiter +} + +func (d *denyThenRecordRateLimiter) RateLimitPreFetch(ctx *Context, info *FetchInfo, input json.RawMessage) (*RateLimitDeny, error) { + _, _ = d.inner.RateLimitPreFetch(ctx, info, input) + if info != nil && info.DataSourceName == d.deny { + return &RateLimitDeny{Reason: "test denied"}, nil + } + return nil, nil +} + +func (d *denyThenRecordRateLimiter) RenderResponseExtension(_ *Context, _ io.Writer) error { + return nil +} + +// TestDataflowPreFetchHookOrder is the regression test for the codex P1 on the +// hardening: pre-fetch hooks (RateLimitPreFetch/AuthorizePreFetch) fire at +// PREPARE time and can render order-sensitive accumulated state into response +// extensions, so dataflow must dispatch in LEAF order (= the wave executor's +// spawn order), not FetchID order. FetchIDs here are deliberately INVERTED +// against tree order — FetchID-ordered seeding would call the hook for B first. +func TestDataflowPreFetchHookOrder(t *testing.T) { + t.Run("serial sequence, inverted fetch IDs", func(t *testing.T) { + run := func(enableDataflow bool) []string { + a := newDelayDataSource(`{"data":{"a":"A"}}`, 0) + b := newDelayDataSource(`{"data":{"b":"B"}}`, 0) + response := &GraphQLResponse{ + Info: &GraphQLResponseInfo{OperationType: ast.OperationTypeQuery}, + Fetches: Sequence( + Single(dataflowFetch(a, 2, nil, "A", `{"fetch":"a"}`)), + Single(dataflowFetch(b, 1, nil, "B", `{"fetch":"b"}`)), + ), + Data: nestedParallelData("a", "b"), + } + limiter := &recordingRateLimiter{} + configure := func(_ *Loader, ctx *Context) { + ctx.RateLimitOptions = RateLimitOptions{Enable: true} + ctx.SetRateLimiter(limiter) + } + out, err := runDataflowScenario(t, context.Background(), enableDataflow, ast.OperationTypeQuery, response, configure) + require.NoError(t, err) + require.Equal(t, `{"data":{"a":"A","b":"B"}}`, out) + return limiter.order() + } + require.Equal(t, []string{"A", "B"}, run(false)) + require.Equal(t, []string{"A", "B"}, run(true)) + }) + + t.Run("inline skip unblocks ahead of queued siblings", func(t *testing.T) { + // codex P1 round 2: A is rate-limit denied (inline skipLoad completion), + // which makes C (leaf 1, depends on A) ready while B (leaf 2) is already + // queued. FIFO dispatch would call hooks as A,B,C; the wave executor's + // serial order is A,C,B — global leaf-minimum popping must reproduce it. + run := func(enableDataflow bool) []string { + a := newDelayDataSource(`{"data":{"a":"never"}}`, 0) + c := newDelayDataSource(`{"data":{"c":"C"}}`, 0) + b := newDelayDataSource(`{"data":{"b":"B"}}`, 0) + aFetch := dataflowFetch(a, 0, nil, "A", `{"fetch":"a"}`) + response := &GraphQLResponse{ + Info: &GraphQLResponseInfo{OperationType: ast.OperationTypeQuery}, + Fetches: Sequence( + SingleWithPath(aFetch, "query.a"), + Single(dataflowFetch(c, 2, []int{0}, "C", `{"fetch":"c"}`)), + Single(dataflowFetch(b, 1, nil, "B", `{"fetch":"b"}`)), + ), + Data: nestedParallelNullableData("a", "c", "b"), + } + limiter := &recordingRateLimiter{} + deny := &denyThenRecordRateLimiter{deny: "A", inner: limiter} + configure := func(_ *Loader, ctx *Context) { + ctx.RateLimitOptions = RateLimitOptions{Enable: true} + ctx.SetRateLimiter(deny) + } + out, err := runDataflowScenario(t, context.Background(), enableDataflow, ast.OperationTypeQuery, response, configure) + require.NoError(t, err) + require.Equal(t, `{"errors":[{"message":"Rate limit exceeded for Subgraph 'A' at Path 'query.a', Reason: test denied."}],"data":{"a":null,"c":"C","b":"B"}}`, out) + return limiter.order() + } + require.Equal(t, []string{"A", "C", "B"}, run(false)) + require.Equal(t, []string{"A", "C", "B"}, run(true)) + }) + + t.Run("parallel wave, inverted fetch IDs", func(t *testing.T) { + // The wave executor calls hooks CONCURRENTLY within a Parallel wave + // (unordered), so only the dataflow order is asserted: deterministic + // leaf order, root first. + r := newDelayDataSource(`{"data":{"r":"R"}}`, 0) + x := newDelayDataSource(`{"data":{"x":"X"}}`, 0) + y := newDelayDataSource(`{"data":{"y":"Y"}}`, 0) + response := &GraphQLResponse{ + Info: &GraphQLResponseInfo{OperationType: ast.OperationTypeQuery}, + Fetches: Sequence( + Single(dataflowFetch(r, 0, nil, "R", `{"fetch":"r"}`)), + Parallel( + Single(dataflowFetch(x, 5, []int{0}, "X", `{"fetch":"x"}`)), + Single(dataflowFetch(y, 3, []int{0}, "Y", `{"fetch":"y"}`)), + ), + ), + Data: nestedParallelData("r", "x", "y"), + } + limiter := &recordingRateLimiter{} + configure := func(_ *Loader, ctx *Context) { + ctx.RateLimitOptions = RateLimitOptions{Enable: true} + ctx.SetRateLimiter(limiter) + } + out, err := runDataflowScenario(t, context.Background(), true, ast.OperationTypeQuery, response, configure) + require.NoError(t, err) + require.Equal(t, `{"data":{"r":"R","x":"X","y":"Y"}}`, out) + require.Equal(t, []string{"R", "X", "Y"}, limiter.order()) + }) +} + +func TestDataflowFallbacks(t *testing.T) { + t.Run("duplicate fetch IDs", func(t *testing.T) { + build := func() *GraphQLResponse { + a := newDelayDataSource(`{"data":{"a":"A"}}`, 0) + b := newDelayDataSource(`{"data":{"b":"B"}}`, 0) + return &GraphQLResponse{ + Info: &GraphQLResponseInfo{OperationType: ast.OperationTypeQuery}, + Fetches: Sequence( + Single(dataflowFetch(a, 7, nil, "A", `{"fetch":"a"}`)), + Single(dataflowFetch(b, 7, nil, "B", `{"fetch":"b"}`)), + ), + Data: nestedParallelData("a", "b"), + } + } + wave, dataflow := runBothExecutors(t, build, nil) + require.Equal(t, `{"data":{"a":"A","b":"B"}}`, wave) + require.Equal(t, wave, dataflow) + }) + + t.Run("dependency cycle", func(t *testing.T) { + build := func() *GraphQLResponse { + a := newDelayDataSource(`{"data":{"a":"A"}}`, 0) + b := newDelayDataSource(`{"data":{"b":"B"}}`, 0) + return &GraphQLResponse{ + Info: &GraphQLResponseInfo{OperationType: ast.OperationTypeQuery}, + Fetches: Sequence( + Single(dataflowFetch(a, 1, []int{2}, "A", `{"fetch":"a"}`)), + Single(dataflowFetch(b, 2, []int{1}, "B", `{"fetch":"b"}`)), + ), + Data: nestedParallelData("a", "b"), + } + } + wave, dataflow := runBothExecutors(t, build, nil) + require.Equal(t, `{"data":{"a":"A","b":"B"}}`, wave) + require.Equal(t, wave, dataflow) + }) + + t.Run("mutation", func(t *testing.T) { + build := func() *GraphQLResponse { + a := newDelayDataSource(`{"data":{"a":"A"}}`, 0) + b := newDelayDataSource(`{"data":{"b":"B"}}`, 0) + return &GraphQLResponse{ + Info: &GraphQLResponseInfo{OperationType: ast.OperationTypeMutation}, + Fetches: Sequence( + Single(dataflowFetch(a, 0, nil, "A", `{"fetch":"a"}`)), + Single(dataflowFetch(b, 1, []int{0}, "B", `{"fetch":"b"}`)), + ), + Data: nestedParallelData("a", "b"), + } + } + wave, waveErr := runDataflowScenario(t, context.Background(), false, ast.OperationTypeMutation, build(), nil) + require.NoError(t, waveErr) + dataflow, dataflowErr := runDataflowScenario(t, context.Background(), true, ast.OperationTypeMutation, build(), nil) + require.NoError(t, dataflowErr) + require.Equal(t, `{"data":{"a":"A","b":"B"}}`, wave) + require.Equal(t, wave, dataflow) + }) +} + +// TestDataflowCancellation cancels the request context while a load is blocked +// on it; both executors must return promptly with identical output and fully +// drain their goroutines. +func TestDataflowCancellation(t *testing.T) { + defer goleak.VerifyNone(t) + run := func(enableDataflow bool) string { + blocked := newControlledLoaderDataSource(`{"data":{"s":"never"}}`) + blocked.waitForCancel = true + fast := newDelayDataSource(`{"data":{"f":"F"}}`, 0) + sFetch := nestedParallelSingleFetchWithInfo(blocked, `{"fetch":"s"}`, "S", "S") + sFetch.FetchDependencies.FetchID = 0 + fFetch := dataflowFetch(fast, 1, nil, "F", `{"fetch":"f"}`) + response := &GraphQLResponse{ + Info: &GraphQLResponseInfo{OperationType: ast.OperationTypeQuery}, + Fetches: Sequence( + Parallel( + SingleWithPath(sFetch, "query.s"), + Single(fFetch), + ), + ), + Data: nestedParallelNullableData("s", "f"), + } + reqCtx, cancel := context.WithCancel(context.Background()) + timer := time.AfterFunc(30*time.Millisecond, cancel) + defer timer.Stop() + defer cancel() + start := time.Now() + out, err := runDataflowScenario(t, reqCtx, enableDataflow, ast.OperationTypeQuery, response, nil) + require.NoError(t, err) + require.Less(t, time.Since(start), 2*time.Second) + require.True(t, blocked.cancelled.Load()) + return out + } + wave := run(false) + dataflow := run(true) + require.Equal(t, `{"errors":[{"message":"Failed to fetch from Subgraph 'S' at Path 'query.s'."}],"data":{"s":null,"f":"F"}}`, wave) + require.Equal(t, wave, dataflow) +} + +// TestDataflowNoGoroutineLeaks: a prepare-time fatal (rate limiter error) fires +// while a sibling load is in flight; the coordinator must cancel, drain the +// in-flight worker, flush, and return the fatal — leaking nothing. +func TestDataflowNoGoroutineLeaks(t *testing.T) { + defer goleak.VerifyNone(t) + build := func() *GraphQLResponse { + p := newDelayDataSource(`{"data":{"p":"P"}}`, 0) + slow := newDelayDataSource(`{"data":{"s":"S"}}`, 50*time.Millisecond) + boom := newDelayDataSource(`{"data":{"x":"X"}}`, 0) + return &GraphQLResponse{ + Info: &GraphQLResponseInfo{OperationType: ast.OperationTypeQuery}, + Fetches: Sequence( + Single(dataflowFetch(p, 0, nil, "P", `{"fetch":"p"}`)), + Parallel( + Single(dataflowFetch(slow, 1, []int{0}, "S", `{"fetch":"s"}`)), + Single(dataflowFetch(boom, 2, []int{0}, "X", `{"fetch":"x"}`)), + ), + ), + Data: nestedParallelNullableData("p", "s", "x"), + } + } + configure := func(_ *Loader, ctx *Context) { + ctx.RateLimitOptions = RateLimitOptions{Enable: true} + ctx.SetRateLimiter(&stubRateLimiter{errName: "X"}) + } + _, waveErr := runDataflowScenario(t, context.Background(), false, ast.OperationTypeQuery, build(), configure) + require.EqualError(t, waveErr, "rate limiter exploded") + _, dataflowErr := runDataflowScenario(t, context.Background(), true, ast.OperationTypeQuery, build(), configure) + require.EqualError(t, dataflowErr, "rate limiter exploded") +} + +// TestCollectDataflowLeavesAcceptsOnlyFlatTrees pins the structural guard that +// makes ENGINE_ENABLE_DATAFLOW safe to combine with schedule-tree plans: only the +// flat createParallelNodes shape (Sequence of Single / Parallel-of-Single) is +// eligible; any nested tree must report ok=false so resolveDataflow falls back to +// the (mergeMu-protected) wave executor. +func TestCollectDataflowLeavesAcceptsOnlyFlatTrees(t *testing.T) { + single := func() *FetchTreeNode { return Single(&SingleFetch{}) } + tests := []struct { + name string + node *FetchTreeNode + wantOK bool + wantCount int + }{ + {name: "nil", node: nil, wantOK: true, wantCount: 0}, + {name: "bare single", node: single(), wantOK: true, wantCount: 1}, + {name: "flat sequence", node: Sequence(single(), single()), wantOK: true, wantCount: 2}, + {name: "sequence with parallel of singles", node: Sequence(single(), Parallel(single(), single())), wantOK: true, wantCount: 3}, + // root Parallel is not createParallelNodes output (root is always a Sequence) + {name: "root parallel", node: Parallel(single(), single()), wantOK: false, wantCount: 0}, + {name: "parallel containing sequence", node: Sequence(single(), Parallel(Sequence(single(), single()))), wantOK: false, wantCount: 0}, + {name: "sequence containing sequence", node: Sequence(Sequence(single())), wantOK: false, wantCount: 0}, + {name: "single without fetch", node: &FetchTreeNode{Kind: FetchTreeNodeKindSingle, Item: &FetchItem{}}, wantOK: false, wantCount: 0}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + leaves, ok := collectDataflowLeaves(tt.node) + require.Equal(t, tt.wantOK, ok) + require.Equal(t, tt.wantCount, len(leaves)) + }) + } +} + +// TestDataflowFallsBackOnNestedScheduleTreePlan asserts the flag-interaction +// contract: with EnableDataflowExecution on, a nested (schedule-tree style) plan +// is structurally rejected by collectDataflowLeaves and runs through the wave +// executor, producing bytes identical to the dataflow-off run. +func TestDataflowFallsBackOnNestedScheduleTreePlan(t *testing.T) { + run := func(t *testing.T, enableDataflow bool) string { + t.Helper() + a := newControlledLoaderDataSource(`{"data":{"a":"A"}}`) + b := newControlledLoaderDataSource(`{"data":{"b":"B"}}`) + c := newControlledLoaderDataSource(`{"data":{"c":"C"}}`) + aFetch := nestedParallelSingleFetch(a, `{"fetch":"A"}`) + aFetch.FetchDependencies.FetchID = 0 + bFetch := nestedParallelSingleFetch(b, `{"fetch":"B"}`) + bFetch.FetchDependencies.FetchID = 1 + bFetch.FetchDependencies.DependsOnFetchIDs = []int{0} + cFetch := nestedParallelSingleFetch(c, `{"fetch":"C"}`) + cFetch.FetchDependencies.FetchID = 2 + cFetch.FetchDependencies.DependsOnFetchIDs = []int{1} + + response := &GraphQLResponse{ + Info: &GraphQLResponseInfo{OperationType: ast.OperationTypeQuery}, + Fetches: Sequence( + Single(aFetch), + Parallel( + Sequence(Single(bFetch), Single(cFetch)), + ), + ), + Data: nestedParallelData("a", "b", "c"), + } + + out, err := runDataflowScenario(t, context.Background(), enableDataflow, ast.OperationTypeQuery, response, nil) + require.NoError(t, err) + return out + } + + withDataflow := run(t, true) + withoutDataflow := run(t, false) + require.Equal(t, `{"data":{"a":"A","b":"B","c":"C"}}`, withDataflow) + require.Equal(t, withoutDataflow, withDataflow) +} diff --git a/v2/pkg/engine/resolve/resolve.go b/v2/pkg/engine/resolve/resolve.go index a45f3487ec..6e95be7273 100644 --- a/v2/pkg/engine/resolve/resolve.go +++ b/v2/pkg/engine/resolve/resolve.go @@ -180,6 +180,12 @@ type ResolverOptions struct { // This helps keep the Heap size more maintainable if you regularly perform large queries. MaxRecyclableParserSize int + // EnableDataflowExecution runs query fetch DAGs through the dataflow executor + // (per-FetchID dependency gating) instead of the per-wave-barrier executor. + // Reduces wall-clock latency under skewed subgraph latency; byte-identical + // responses. Off by default. + EnableDataflowExecution bool + // ResolvableOptions are configuration options for the Resolvable struct ResolvableOptions ResolvableOptions @@ -325,6 +331,7 @@ func newTools(options ResolverOptions, allowedExtensionFields map[string]struct{ apolloRouterCompatibilitySubrequestHTTPError: options.ApolloRouterCompatibilitySubrequestHTTPError, propagateFetchReasons: options.PropagateFetchReasons, validateRequiredExternalFields: options.ValidateRequiredExternalFields, + enableDataflow: options.EnableDataflowExecution, singleFlight: sf, jsonArena: a, }, From f75f280f48ddc314fd8166b7087c5c3a8e539b8f Mon Sep 17 00:00:00 2001 From: Jens Neuse Date: Thu, 11 Jun 2026 18:06:33 +0200 Subject: [PATCH 5/5] chore: satisfy gci and modernize linters Co-Authored-By: Claude Opus 4.8 (1M context) --- v2/pkg/engine/datasource/httpclient/nethttpclient.go | 5 +---- v2/pkg/engine/resolve/loader.go | 7 +------ v2/pkg/engine/resolve/loader_dataflow.go | 2 ++ 3 files changed, 4 insertions(+), 10 deletions(-) diff --git a/v2/pkg/engine/datasource/httpclient/nethttpclient.go b/v2/pkg/engine/datasource/httpclient/nethttpclient.go index 04a6ef578c..a8a0dd4850 100644 --- a/v2/pkg/engine/datasource/httpclient/nethttpclient.go +++ b/v2/pkg/engine/datasource/httpclient/nethttpclient.go @@ -259,10 +259,7 @@ func makeHTTPRequest(client *http.Client, ctx context.Context, baseHeaders http. // the compressed size, so this remains a safe lower-bound hint. if cl := response.ContentLength; cl > 0 { const maxPreAlloc = 1 << 21 // 2 MiB - want := int(cl) - if want > maxPreAlloc { - want = maxPreAlloc - } + want := min(int(cl), maxPreAlloc) if want > out.Cap() { out.Grow(want) } diff --git a/v2/pkg/engine/resolve/loader.go b/v2/pkg/engine/resolve/loader.go index 76c5e7777e..59a49d0bf9 100644 --- a/v2/pkg/engine/resolve/loader.go +++ b/v2/pkg/engine/resolve/loader.go @@ -795,12 +795,7 @@ func fetchTreeHasNestedParallel(node *FetchTreeNode) bool { } } } - for _, child := range node.ChildNodes { - if fetchTreeHasNestedParallel(child) { - return true - } - } - return false + return slices.ContainsFunc(node.ChildNodes, fetchTreeHasNestedParallel) } func (l *Loader) callOnFinished(res *result) { diff --git a/v2/pkg/engine/resolve/loader_dataflow.go b/v2/pkg/engine/resolve/loader_dataflow.go index 85219e9413..01cd484baf 100644 --- a/v2/pkg/engine/resolve/loader_dataflow.go +++ b/v2/pkg/engine/resolve/loader_dataflow.go @@ -7,7 +7,9 @@ import ( "slices" "github.com/pkg/errors" + "github.com/wundergraph/astjson" + "github.com/wundergraph/graphql-go-tools/v2/pkg/ast" )