fix: improve query planning time

router-tests/go.mod

@@ -23,11 +23,11 @@ require (
 	github.com/twmb/franz-go v1.16.1
 	github.com/twmb/franz-go/pkg/kadm v1.11.0
 	github.com/wundergraph/astjson v0.0.0-20250106123708-be463c97e083
-	github.com/wundergraph/cosmo/demo v0.0.0-20250912064154-106e871ee32e
+	github.com/wundergraph/cosmo/demo v0.0.0-20251029114720-2b84bc4e4e77
 	github.com/wundergraph/cosmo/demo/pkg/subgraphs/projects v0.0.0-20250715110703-10f2e5f9c79e
-	github.com/wundergraph/cosmo/router v0.0.0-20250912064154-106e871ee32e
+	github.com/wundergraph/cosmo/router v0.0.0-20251029114720-2b84bc4e4e77
 	github.com/wundergraph/cosmo/router-plugin v0.0.0-20250808194725-de123ba1c65e
-	github.com/wundergraph/graphql-go-tools/v2 v2.0.0-rc.232
+	github.com/wundergraph/graphql-go-tools/v2 v2.0.0-rc.235
 	go.opentelemetry.io/otel v1.36.0
 	go.opentelemetry.io/otel/sdk v1.36.0
 	go.opentelemetry.io/otel/sdk/metric v1.36.0

router-tests/go.sum

@@ -354,8 +354,8 @@ github.com/wundergraph/astjson v0.0.0-20250106123708-be463c97e083 h1:8/D7f8gKxTB
 github.com/wundergraph/astjson v0.0.0-20250106123708-be463c97e083/go.mod h1:eOTL6acwctsN4F3b7YE+eE2t8zcJ/doLm9sZzsxxxrE=
 github.com/wundergraph/consul/sdk v0.0.0-20250204115147-ed842a8fd301 h1:EzfKHQoTjFDDcgaECCCR2aTePqMu9QBmPbyhqIYOhV0=
 github.com/wundergraph/consul/sdk v0.0.0-20250204115147-ed842a8fd301/go.mod h1:wxI0Nak5dI5RvJuzGyiEK4nZj0O9X+Aw6U0tC1wPKq0=
-github.com/wundergraph/graphql-go-tools/v2 v2.0.0-rc.232 h1:G04FDSXlEaQZS9cBrKlP8djbzquoQElB7w7i/d4sAHg=
-github.com/wundergraph/graphql-go-tools/v2 v2.0.0-rc.232/go.mod h1:ErOQH1ki2+SZB8JjpTyGVnoBpg5picIyjvuWQJP4abg=
+github.com/wundergraph/graphql-go-tools/v2 v2.0.0-rc.235 h1:GiqY9zm5OR6SElIohw/rzfqejm3R1HjlSbMhbtZ4zWM=
+github.com/wundergraph/graphql-go-tools/v2 v2.0.0-rc.235/go.mod h1:ErOQH1ki2+SZB8JjpTyGVnoBpg5picIyjvuWQJP4abg=
 github.com/xrash/smetrics v0.0.0-20250705151800-55b8f293f342 h1:FnBeRrxr7OU4VvAzt5X7s6266i6cSVkkFPS0TuXWbIg=
 github.com/xrash/smetrics v0.0.0-20250705151800-55b8f293f342/go.mod h1:Ohn+xnUBiLI6FVj/9LpzZWtj1/D6lUovWYBkxHVV3aM=
 github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4=

router/core/plan_generator.go

@@ -74,28 +74,14 @@ func NewPlanner(planConfiguration *plan.Configuration, definition *ast.Document,
 	}, nil
 }
 
+// PlanOperation creates a query plan from an operation file in a pretty-printed text or JSON format
 func (pl *Planner) PlanOperation(operationFilePath string, outputFormat PlanOutputFormat) (string, error) {
-	operation, err := pl.parseOperation(operationFilePath)
-	if err != nil {
-		return "", &PlannerOperationValidationError{err: err}
-	}
-
-	operationName := findOperationName(operation)
-	if operationName == nil {
-		return "", &PlannerOperationValidationError{err: errors.New("operation name not found")}
-	}
-
-	err = pl.normalizeOperation(operation, operationName)
-	if err != nil {
-		return "", &PlannerOperationValidationError{err: err}
-	}
-
-	err = pl.validateOperation(operation)
+	operation, err := pl.ParseAndPrepareOperation(operationFilePath)
 	if err != nil {
-		return "", &PlannerOperationValidationError{err: err}
+		return "", err
 	}
 
-	rawPlan, err := pl.planOperation(operation)
+	rawPlan, err := pl.PlanPreparedOperation(operation)
 	if err != nil {
 		return "", fmt.Errorf("failed to plan operation: %w", err)
 	}
@@ -111,31 +97,35 @@ func (pl *Planner) PlanOperation(operationFilePath string, outputFormat PlanOutp
 		return string(marshal), nil
 	}
 
-	return "", fmt.Errorf("invalid type specified: %q", outputFormat)
+	return "", fmt.Errorf("invalid outputFormat specified: %q", outputFormat)
 }
 
-func (pl *Planner) PlanParsedOperation(operation *ast.Document) (*resolve.FetchTreeQueryPlanNode, error) {
-	operationName := findOperationName(operation)
-	if operationName == nil {
-		return nil, errors.New("operation name not found")
+// ParseAndPrepareOperation parses, normalizes and validates the operation
+func (pl *Planner) ParseAndPrepareOperation(operationFilePath string) (*ast.Document, error) {
+	operation, err := pl.parseOperation(operationFilePath)
+	if err != nil {
+		return nil, &PlannerOperationValidationError{err: err}
 	}
 
-	err := pl.normalizeOperation(operation, operationName)
-	if err != nil {
-		return nil, fmt.Errorf("failed to normalize operation: %w", err)
+	return pl.PrepareOperation(operation)
+}
+
+// PrepareOperation normalizes and validates the operation
+func (pl *Planner) PrepareOperation(operation *ast.Document) (*ast.Document, error) {
+	operationName := findOperationName(operation)
+	if operationName == nil {
+		return nil, &PlannerOperationValidationError{err: errors.New("operation name not found")}
 	}
 
-	err = pl.validateOperation(operation)
-	if err != nil {
+	if err := pl.normalizeOperation(operation, operationName); err != nil {
 		return nil, &PlannerOperationValidationError{err: err}
 	}
 
-	rawPlan, err := pl.planOperation(operation)
-	if err != nil {
-		return nil, fmt.Errorf("failed to plan operation: %w", err)
+	if err := pl.validateOperation(operation); err != nil {
+		return nil, &PlannerOperationValidationError{err: err}
 	}
 
-	return rawPlan, nil
+	return operation, nil
 }
 
 func (pl *Planner) normalizeOperation(operation *ast.Document, operationName []byte) (err error) {
@@ -169,7 +159,8 @@ func (pl *Planner) normalizeOperation(operation *ast.Document, operationName []b
 	return nil
 }
 
-func (pl *Planner) planOperation(operation *ast.Document) (planNode *resolve.FetchTreeQueryPlanNode, err error) {
+// PlanPreparedOperation creates a query plan from a normalized and validated operation
+func (pl *Planner) PlanPreparedOperation(operation *ast.Document) (planNode *resolve.FetchTreeQueryPlanNode, err error) {
 	defer func() {
 		if r := recover(); r != nil {
 			err = fmt.Errorf("panic during plan generation: %v", r)
@@ -358,8 +349,8 @@ func (pg *PlanGenerator) loadConfiguration(routerConfig *nodev1.RouterConfig, lo
 	// we need to merge the base schema, it contains the __schema and __type queries
 	// these are not usually part of a regular GraphQL schema
 	// the engine needs to have them defined, otherwise it cannot resolve such fields
-	err = asttransform.MergeDefinitionWithBaseSchema(&definition)
-	if err != nil {
+
+	if err := asttransform.MergeDefinitionWithBaseSchema(&definition); err != nil {
 		return fmt.Errorf("failed to merge graphql schema with base schema: %w", err)
 	}
 
@@ -410,5 +401,6 @@ func findOperationName(operation *ast.Document) (operationName []byte) {
 			return operation.OperationDefinitionNameBytes(operation.RootNodes[i].Ref)
 		}
 	}
+	// TODO: assign static operation name if we have single anonymous operation
 	return nil
 }

router/core/plan_generator_test.go

@@ -29,7 +29,7 @@ func TestPlanOperationPanic(t *testing.T) {
 	}
 
 	assert.NotPanics(t, func() {
-		_, err = planner.planOperation(invalidOperation)
+		_, err = planner.PlanPreparedOperation(invalidOperation)
 		assert.Error(t, err)
 	})
 }

router/go.mod

@@ -31,7 +31,7 @@ require (
 	github.com/tidwall/gjson v1.18.0
 	github.com/tidwall/sjson v1.2.5
 	github.com/twmb/franz-go v1.16.1
-	github.com/wundergraph/graphql-go-tools/v2 v2.0.0-rc.232
+	github.com/wundergraph/graphql-go-tools/v2 v2.0.0-rc.235
 	// Do not upgrade, it renames attributes we rely on
 	go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.58.0
 	go.opentelemetry.io/contrib/propagators/b3 v1.23.0
@@ -194,4 +194,4 @@ replace (
 // Remember you can use Go workspaces to avoid using replace directives in multiple go.mod files
 // Use what is best for your personal workflow. See CONTRIBUTING.md for more information
 
-// replace github.com/wundergraph/graphql-go-tools/v2 => ../../graphql-go-tools/v2
+//replace github.com/wundergraph/graphql-go-tools/v2 => ../../graphql-go-tools/v2

router/go.sum

@@ -322,8 +322,8 @@ github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/
 github.com/wk8/go-ordered-map/v2 v2.1.8/go.mod h1:5nJHM5DyteebpVlHnWMV0rPz6Zp7+xBAnxjb1X5vnTw=
 github.com/wundergraph/astjson v0.0.0-20250106123708-be463c97e083 h1:8/D7f8gKxTBjW+SZK4mhxTTBVpxcqeBgWF1Rfmltbfk=
 github.com/wundergraph/astjson v0.0.0-20250106123708-be463c97e083/go.mod h1:eOTL6acwctsN4F3b7YE+eE2t8zcJ/doLm9sZzsxxxrE=
-github.com/wundergraph/graphql-go-tools/v2 v2.0.0-rc.232 h1:G04FDSXlEaQZS9cBrKlP8djbzquoQElB7w7i/d4sAHg=
-github.com/wundergraph/graphql-go-tools/v2 v2.0.0-rc.232/go.mod h1:ErOQH1ki2+SZB8JjpTyGVnoBpg5picIyjvuWQJP4abg=
+github.com/wundergraph/graphql-go-tools/v2 v2.0.0-rc.235 h1:GiqY9zm5OR6SElIohw/rzfqejm3R1HjlSbMhbtZ4zWM=
+github.com/wundergraph/graphql-go-tools/v2 v2.0.0-rc.235/go.mod h1:ErOQH1ki2+SZB8JjpTyGVnoBpg5picIyjvuWQJP4abg=
 github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4=
 github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4=
 github.com/yuin/gopher-lua v1.1.1 h1:kYKnWBjvbNP4XLT3+bPEwAXJx262OhaHDWDVOPjL46M=

router/internal/planningbenchmark/.gitignore

@@ -0,0 +1,4 @@
+*.json
+*.txt
+*.out
+*.test

router/internal/planningbenchmark/README.md

@@ -0,0 +1,27 @@
+# Planning Benchmark
+
+Allows to benchmark query planning performance of a single graphql query
+
+## Prerequisites
+
+- Compose subgraphs to get execution config.json
+- Create `benchmark_config.json` file in the `router/internal/planningbenchmark` directory with the following content:
+
+```json
+{
+  "executionConfigPath": "<path-to-composed-config.json>",
+  "operationPath": "<path-to-graphql-query-file>.graphql",
+}
+```
+
+## Running the Benchmark
+
+Run `BenchmarkPlanning` benchmark from benchmark_test.go
+
+## Running benchmark using taskfile
+
+- Modify step name in `router/internal/planningbenchmark/Taskfile.yml` if needed
+- Run `task bench-cpu` from `router/internal/planningbenchmark` directory
+- Profile will be generated in file `<step>_cpu.out`
+
+You could use different profiles or combine profile. See Taskfile.yml for more details.

router/internal/planningbenchmark/Taskfile.yaml

@@ -0,0 +1,26 @@
+version: '3'
+
+vars:
+  step: bench_planning_01
+
+tasks:
+  bench:
+    go test -benchtime 60s -bench=. -benchmem -memprofile {{.step}}_mem.out -cpuprofile {{.step}}_cpu.out > {{.step}}_benchmark.txt
+
+  bench-cpu:
+    go test -benchtime 60s -bench=. -cpuprofile {{.step}}_cpu.out
+
+  bench-mem:
+    go test -benchtime 60s -bench=. -memprofile {{.step}}_mem.out
+
+  profile-cpu-web:
+    go tool pprof -http=localhost:8080 {{.step}}_cpu.out
+
+  profile-cpu:
+    go tool pprof {{.step}}_cpu.out
+
+  profile-mem-web:
+    go tool pprof -http=localhost:8080 {{.step}}_mem.out
+
+  profile-mem:
+    go tool pprof {{.step}}_mem.out

router/internal/planningbenchmark/benchmark_config.json.example

@@ -0,0 +1,4 @@
+{
+  "executionConfigPath": "",
+  "operationPath": ""
+}

router/internal/planningbenchmark/benchmark_test.go

@@ -0,0 +1,76 @@
+package planningbenchmark
+
+import (
+	"encoding/json"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"go.uber.org/zap"
+
+	"github.com/wundergraph/cosmo/router/core"
+)
+
+type BenchmarkConfig struct {
+	ExecutionConfigPath string `json:"executionConfigPath"`
+	OperationPath       string `json:"operationPath"`
+}
+
+func TestPlanning(t *testing.T) {
+	cfgContent, err := os.ReadFile("benchmark_config.json")
+	if err != nil {
+		t.Skipf("unable to read benchmark_config.json: %v", err)
+	}
+
+	var cfg BenchmarkConfig
+	require.NoError(t, json.Unmarshal(cfgContent, &cfg))
+
+	logger := zap.NewNop()
+
+	pg, err := core.NewPlanGenerator(cfg.ExecutionConfigPath, logger, 0)
+	require.NoError(t, err)
+
+	pl, err := pg.GetPlanner()
+	require.NoError(t, err)
+
+	opDoc, err := pl.ParseAndPrepareOperation(cfg.OperationPath)
+	require.NoError(t, err)
+
+	start := time.Now()
+	_, err = pl.PlanPreparedOperation(opDoc)
+	require.NoError(t, err)
+	t.Logf("Planning completed in %v", time.Since(start))
+}
+
+func BenchmarkPlanning(b *testing.B) {
+	cfgContent, err := os.ReadFile("benchmark_config.json")
+	if err != nil {
+		b.Skipf("unable to read benchmark_config.json: %v", err)
+	}
+
+	var cfg BenchmarkConfig
+	require.NoError(b, json.Unmarshal(cfgContent, &cfg))
+
+	logger := zap.NewNop()
+
+	pg, err := core.NewPlanGenerator(cfg.ExecutionConfigPath, logger, 0)
+	require.NoError(b, err)
+
+	pl, err := pg.GetPlanner()
+	require.NoError(b, err)
+
+	b.ReportAllocs()
+	b.ResetTimer()
+
+	for b.Loop() {
+		b.StopTimer()
+		opDoc, err := pl.ParseAndPrepareOperation(cfg.OperationPath)
+		require.NoError(b, err)
+		b.SetBytes(int64(len(opDoc.Input.RawBytes)))
+		b.StartTimer()
+
+		_, err = pl.PlanPreparedOperation(opDoc)
+		require.NoError(b, err)
+	}
+}