eval: parallelize environment loading

These changes add support for loading environments in parallel.

Parallelization is breadth-first, then depth-first. For example, given
the following environments:

```yaml
imports: # env a
  - b
  - c

--
imports; # env b
  - e
--
imports: # env c
  - f
  - g
```

Environments `b` and `c` would be loaded in parallel, then environment
`e` would be loaded, then environments `f` and `g` would be loaded in
parallel. This simplifies the detection of cyclic imports and the
collection of diagnostics.

This improves performance for scenarios that are dominated by
environment load time (e.g. import graphs with high degrees of fanout).

Local benchmark results:

goos: darwin
goarch: arm64
pkg: github.com/pulumi/esc/eval
cpu: Apple M1 Max
BenchmarkEval-10           	    162	  6249736 ns/op	4636941 B/op	  23697 allocs/op
BenchmarkEval-10           	    192	  6229592 ns/op	4637483 B/op	  23697 allocs/op
BenchmarkEval-10           	    192	  6213844 ns/op	4638117 B/op	  23699 allocs/op
BenchmarkEval-10           	    192	  6215693 ns/op	4637189 B/op	  23696 allocs/op
BenchmarkEval-10           	    192	  6286186 ns/op	4637032 B/op	  23696 allocs/op
BenchmarkEval-10           	    192	  6250083 ns/op	4637796 B/op	  23698 allocs/op
BenchmarkEval-10           	    194	  6201700 ns/op	4637262 B/op	  23697 allocs/op
BenchmarkEval-10           	    192	  6256509 ns/op	4637151 B/op	  23697 allocs/op
BenchmarkEval-10           	    193	  6220107 ns/op	4638638 B/op	  23699 allocs/op
BenchmarkEval-10           	    192	  6196454 ns/op	4636411 B/op	  23696 allocs/op
BenchmarkEvalOpen-10       	      9	124156394 ns/op	4632673 B/op	  23703 allocs/op
BenchmarkEvalOpen-10       	      9	123760278 ns/op	4636146 B/op	  23713 allocs/op
BenchmarkEvalOpen-10       	      9	123941329 ns/op	4640512 B/op	  23718 allocs/op
BenchmarkEvalOpen-10       	      9	122636315 ns/op	4637026 B/op	  23710 allocs/op
BenchmarkEvalOpen-10       	      9	123189880 ns/op	4636938 B/op	  23702 allocs/op
BenchmarkEvalOpen-10       	      9	122790926 ns/op	4635157 B/op	  23710 allocs/op
BenchmarkEvalOpen-10       	      9	123945481 ns/op	4637705 B/op	  23711 allocs/op
BenchmarkEvalOpen-10       	      9	123235093 ns/op	4640275 B/op	  23710 allocs/op
BenchmarkEvalOpen-10       	      9	122647329 ns/op	4639136 B/op	  23708 allocs/op
BenchmarkEvalOpen-10       	      9	123545866 ns/op	4638529 B/op	  23715 allocs/op
BenchmarkEvalEnvLoad-10    	     60	 17289035 ns/op	4639187 B/op	  23721 allocs/op
BenchmarkEvalEnvLoad-10    	     72	 17035170 ns/op	4638442 B/op	  23721 allocs/op
BenchmarkEvalEnvLoad-10    	     72	 17176440 ns/op	4639614 B/op	  23722 allocs/op
BenchmarkEvalEnvLoad-10    	     72	 17151064 ns/op	4639322 B/op	  23722 allocs/op
BenchmarkEvalEnvLoad-10    	     70	 17086652 ns/op	4638750 B/op	  23719 allocs/op
BenchmarkEvalEnvLoad-10    	     70	 16979492 ns/op	4638801 B/op	  23721 allocs/op
BenchmarkEvalEnvLoad-10    	     70	 17061474 ns/op	4638384 B/op	  23720 allocs/op
BenchmarkEvalEnvLoad-10    	     70	 17121283 ns/op	4640587 B/op	  23723 allocs/op
BenchmarkEvalEnvLoad-10    	     69	 17093001 ns/op	4639506 B/op	  23722 allocs/op
BenchmarkEvalEnvLoad-10    	     70	 17046285 ns/op	4639028 B/op	  23720 allocs/op
BenchmarkEvalAll-10        	      8	135173646 ns/op	4635106 B/op	  23730 allocs/op
BenchmarkEvalAll-10        	      8	133903672 ns/op	4638396 B/op	  23735 allocs/op
BenchmarkEvalAll-10        	      8	133961172 ns/op	4640463 B/op	  23735 allocs/op
BenchmarkEvalAll-10        	      8	134953359 ns/op	4639089 B/op	  23724 allocs/op
BenchmarkEvalAll-10        	      8	134443724 ns/op	4639118 B/op	  23739 allocs/op
BenchmarkEvalAll-10        	      8	134042062 ns/op	4638356 B/op	  23733 allocs/op
BenchmarkEvalAll-10        	      8	135336984 ns/op	4642280 B/op	  23739 allocs/op
BenchmarkEvalAll-10        	      8	135161682 ns/op	4638414 B/op	  23730 allocs/op
BenchmarkEvalAll-10        	      8	133572880 ns/op	4639524 B/op	  23727 allocs/op
BenchmarkEvalAll-10        	      8	137661594 ns/op	4642510 B/op	  23736 allocs/op

Author: pgavlin

eval/eval.go

@@ -113,7 +113,14 @@ func evalEnvironment(
 		return nil, nil
 	}
 
-	ec := newEvalContext(ctx, validating, name, env, decrypter, providers, envs, map[string]*imported{}, execContext, showSecrets)
+	loader := newLoader(ctx, envs)
+	ec := newEvalContext(ctx, validating, name, env, decrypter, providers, loader, map[string]*imported{}, execContext, showSecrets)
+
+	diags := ec.load()
+	if diags.HasErrors() {
+		return nil, diags
+	}
+
 	v, diags := ec.evaluate()
 
 	s := schema.Never().Schema()
@@ -139,22 +146,23 @@ func evalEnvironment(
 }
 
 type imported struct {
-	evaluating bool
-	value      *value
+	loading bool
+	ctx     *evalContext
+	value   *value
 }
 
 // An evalContext carries the state necessary to evaluate an environment.
 type evalContext struct {
-	ctx          context.Context      // the cancellation context for evaluation
-	validating   bool                 // true if we are only checking the environment
-	showSecrets  bool                 // true if secrets should be decrypted during validation
-	name         string               // the name of the environment
-	env          *ast.EnvironmentDecl // the root of the environment AST
-	decrypter    Decrypter            // the decrypter to use for the environment
-	providers    ProviderLoader       // the provider loader to use
-	environments EnvironmentLoader    // the environment loader to use
-	imports      map[string]*imported // the shared set of imported environments
-	execContext  *esc.ExecContext     // evaluation context used for interpolation
+	ctx         context.Context      // the cancellation context for evaluation
+	validating  bool                 // true if we are only checking the environment
+	showSecrets bool                 // true if secrets should be decrypted during validation
+	name        string               // the name of the environment
+	env         *ast.EnvironmentDecl // the root of the environment AST
+	decrypter   Decrypter            // the decrypter to use for the environment
+	providers   ProviderLoader       // the provider loader to use
+	loader      *loader              // the environment loader
+	imports     map[string]*imported // the shared set of imported environments
+	execContext *esc.ExecContext     // evaluation context used for interpolation
 
 	myContext *value // evaluated context to be used to interpolate properties
 	myImports *value // directly-imported environments
@@ -171,22 +179,22 @@ func newEvalContext(
 	env *ast.EnvironmentDecl,
 	decrypter Decrypter,
 	providers ProviderLoader,
-	environments EnvironmentLoader,
+	loader *loader,
 	imports map[string]*imported,
 	execContext *esc.ExecContext,
 	showSecrets bool,
 ) *evalContext {
 	return &evalContext{
-		ctx:          ctx,
-		validating:   validating,
-		showSecrets:  showSecrets,
-		name:         name,
-		env:          env,
-		decrypter:    decrypter,
-		providers:    providers,
-		environments: environments,
-		imports:      imports,
-		execContext:  execContext.CopyForEnv(name),
+		ctx:         ctx,
+		validating:  validating,
+		showSecrets: showSecrets,
+		name:        name,
+		env:         env,
+		decrypter:   decrypter,
+		providers:   providers,
+		loader:      loader,
+		imports:     imports,
+		execContext: execContext.CopyForEnv(name),
 	}
 }
 
@@ -349,6 +357,55 @@ func (e *evalContext) isReserveTopLevelKey(k string) bool {
 	}
 }
 
+func (e *evalContext) load() syntax.Diagnostics {
+	mine := &imported{loading: true, ctx: e}
+	defer func() { mine.loading = false }()
+	e.imports[e.name] = mine
+
+	loads := make([]*loadedEnvironment, len(e.env.Imports.GetElements()))
+	for i, entry := range e.env.Imports.GetElements() {
+		loads[i] = e.loadImport(entry)
+	}
+
+	for i, entry := range e.env.Imports.GetElements() {
+		l := loads[i]
+		if l == nil {
+			continue
+		}
+		<-l.done
+
+		e.diags.Extend(l.diags...)
+		if l.err != nil {
+			e.errorf(entry.Environment, "%s", l.err.Error())
+			continue
+		}
+
+		imp := newEvalContext(e.ctx, e.validating, l.name, l.env, l.dec, e.providers, e.loader, e.imports, e.execContext, e.showSecrets)
+		diags := imp.load()
+		e.diags.Extend(diags...)
+	}
+
+	return e.diags
+}
+
+func (e *evalContext) loadImport(decl *ast.ImportDecl) *loadedEnvironment {
+	// If the import does not have a name, there's nothing we can do. This can happen for environments
+	// with parse errors.
+	if decl.Environment == nil {
+		return nil
+	}
+	name := decl.Environment.Value
+
+	if imported, ok := e.imports[name]; ok {
+		if imported.loading {
+			e.diags.Extend(syntax.Error(decl.Syntax().Syntax().Range(), fmt.Sprintf("cyclic import of %v", name), decl.Syntax().Syntax().Path()))
+		}
+		return nil
+	}
+
+	return e.loader.load(name)
+}
+
 // evaluate drives the evaluation of the evalContext's environment.
 func (e *evalContext) evaluate() (*value, syntax.Diagnostics) {
 	// Evaluate context. We prepare the context values to later evaluate interpolations.
@@ -393,10 +450,6 @@ func (e *evalContext) evaluateContext() {
 
 // evaluateImports evaluates an environment's imports.
 func (e *evalContext) evaluateImports() {
-	mine := &imported{evaluating: true}
-	defer func() { mine.evaluating = false }()
-	e.imports[e.name] = mine
-
 	myImports := map[string]*value{}
 	for _, entry := range e.env.Imports.GetElements() {
 		e.evaluateImport(myImports, entry)
@@ -432,34 +485,22 @@ func (e *evalContext) evaluateImport(myImports map[string]*value, decl *ast.Impo
 	}
 	name := decl.Environment.Value
 
+	imported, ok := e.imports[name]
+	if !ok {
+		e.diags.Extend(syntax.Error(decl.Syntax().Syntax().Range(), fmt.Sprintf("internal error: missing context for %v", name), decl.Syntax().Syntax().Path()))
+		return
+	}
+
 	merge := true
 	if decl.Meta != nil && decl.Meta.Merge != nil {
 		merge = decl.Meta.Merge.Value
 	}
 
 	var val *value
-	if imported, ok := e.imports[name]; ok {
-		if imported.evaluating {
-			e.diags.Extend(syntax.Error(decl.Syntax().Syntax().Range(), fmt.Sprintf("cyclic import of %v", name), decl.Syntax().Syntax().Path()))
-			return
-		}
+	if imported.value != nil {
 		val = imported.value
 	} else {
-		bytes, dec, err := e.environments.LoadEnvironment(e.ctx, name)
-		if err != nil {
-			e.errorf(decl.Environment, "%s", err.Error())
-			return
-		}
-
-		env, diags, err := LoadYAMLBytes(name, bytes)
-		e.diags.Extend(diags...)
-		if err != nil {
-			e.errorf(decl.Environment, "%s", err.Error())
-			return
-		}
-
-		imp := newEvalContext(e.ctx, e.validating, name, env, dec, e.providers, e.environments, e.imports, e.execContext, e.showSecrets)
-		v, diags := imp.evaluate()
+		v, diags := imported.ctx.evaluate()
 		e.diags.Extend(diags...)
 
 		val = v

eval/eval_test.go

@@ -405,6 +405,8 @@ func benchmarkEval(b *testing.B, openDelay, loadDelay time.Duration) {
 	envs, err := newBenchEnvironments(basePath, loadDelay)
 	require.NoError(b, err)
 
+	b.ResetTimer()
+
 	for i := 0; i < b.N; i++ {
 		execContext, err := esc.NewExecContext(map[string]esc.Value{
 			"pulumi": esc.NewValue(map[string]esc.Value{

eval/loader.go

@@ -0,0 +1,71 @@
+// Copyright 2024, Pulumi Corporation.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package eval
+
+import (
+	"context"
+
+	"github.com/pulumi/esc/ast"
+	"github.com/pulumi/esc/syntax"
+)
+
+type loadedEnvironment struct {
+	done <-chan bool
+
+	name  string
+	env   *ast.EnvironmentDecl
+	dec   Decrypter
+	diags syntax.Diagnostics
+	err   error
+}
+
+type loader struct {
+	ctx          context.Context
+	environments EnvironmentLoader
+	loaded       map[string]*loadedEnvironment
+}
+
+func newLoader(ctx context.Context, environments EnvironmentLoader) *loader {
+	return &loader{
+		ctx:          ctx,
+		environments: environments,
+		loaded:       map[string]*loadedEnvironment{},
+	}
+}
+
+func (l *loader) load(name string) *loadedEnvironment {
+	if loaded, ok := l.loaded[name]; ok {
+		return loaded
+	}
+
+	done := make(chan bool)
+	result := &loadedEnvironment{done: done, name: name}
+	go func() {
+		defer close(done)
+
+		bytes, dec, err := l.environments.LoadEnvironment(l.ctx, name)
+		if err != nil {
+			result.err = err
+			return
+		}
+		result.dec = dec
+
+		result.env, result.diags, result.err = LoadYAMLBytes(name, bytes)
+		return
+	}()
+
+	l.loaded[name] = result
+	return result
+}