performance: compile CUE client side

Restructured the compile logic to happen on the CLI instead of the
BuildKit frontend.

- Avoid uploading the entire workspace to BuildKit on every compilation
- Let the CUE loader scan the files instead of going through the
  BuildKit filesystem gRPC APIs.

Signed-off-by: Andrea Luzzardi <aluzzardi@gmail.com>

client/client.go

@@ -177,11 +177,6 @@ func (c *Client) buildfn(ctx context.Context, st *state.State, env *environment.
 			NoCache: c.cfg.NoCache,
 		})
 
-		lg.Debug().Msg("loading configuration")
-		if err := env.LoadPlan(ctx, s); err != nil {
-			return nil, err
-		}
-
 		// Compute output overlay
 		if fn != nil {
 			if err := fn(ctx, env, s); err != nil {

cmd/dagger/cmd/compute.go

@@ -170,7 +170,25 @@ var computeCmd = &cobra.Command{
 
 		cl := common.NewClient(ctx)
 
-		err := cl.Do(ctx, st, func(ctx context.Context, env *environment.Environment, s solver.Solver) error {
+		v := compiler.NewValue()
+		plan, err := st.CompilePlan(ctx)
+		if err != nil {
+			lg.Fatal().Err(err).Msg("failed to compile plan")
+		}
+		if err := v.FillPath(cue.MakePath(), plan); err != nil {
+			lg.Fatal().Err(err).Msg("failed to compile plan")
+		}
+
+		inputs, err := st.CompileInputs()
+		if err != nil {
+			lg.Fatal().Err(err).Msg("failed to compile inputs")
+		}
+
+		if err := v.FillPath(cue.MakePath(), inputs); err != nil {
+			lg.Fatal().Err(err).Msg("failed to compile inputs")
+		}
+
+		err = cl.Do(ctx, st, func(ctx context.Context, env *environment.Environment, s solver.Solver) error {
 			// check that all inputs are set
 			checkInputs(ctx, env)
 
@@ -178,13 +196,6 @@ var computeCmd = &cobra.Command{
 				return err
 			}
 
-			v := compiler.NewValue()
-			if err := v.FillPath(cue.MakePath(), env.Plan()); err != nil {
-				return err
-			}
-			if err := v.FillPath(cue.MakePath(), env.Input()); err != nil {
-				return err
-			}
 			if err := v.FillPath(cue.MakePath(), env.Computed()); err != nil {
 				return err
 			}

cmd/dagger/cmd/doc.go

@@ -318,7 +318,7 @@ func loadCode(packageName string) (*compiler.Value, error) {
 		stdlib.Path: stdlib.FS,
 	}
 
-	src, err := compiler.Build(sources, packageName)
+	src, err := compiler.Build("/config", sources, packageName)
 	if err != nil {
 		return nil, err
 	}

cmd/dagger/cmd/query.go

@@ -1,15 +1,12 @@
 package cmd
 
 import (
-	"context"
 	"fmt"
 
 	"cuelang.org/go/cue"
 	"go.dagger.io/dagger/cmd/dagger/cmd/common"
 	"go.dagger.io/dagger/cmd/dagger/logger"
 	"go.dagger.io/dagger/compiler"
-	"go.dagger.io/dagger/environment"
-	"go.dagger.io/dagger/solver"
 
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
@@ -45,28 +42,27 @@ var queryCmd = &cobra.Command{
 
 		doneCh := common.TrackWorkspaceCommand(ctx, cmd, workspace, state)
 
-		cl := common.NewClient(ctx)
 		cueVal := compiler.NewValue()
 
-		err := cl.Do(ctx, state, func(ctx context.Context, env *environment.Environment, s solver.Solver) error {
 		if !viper.GetBool("no-plan") {
-				if err := cueVal.FillPath(cue.MakePath(), env.Plan()); err != nil {
+			plan, err := state.CompilePlan(ctx)
-					return err
+			if err != nil {
+				lg.Fatal().Err(err).Msg("failed to compile plan")
+			}
+			if err := cueVal.FillPath(cue.MakePath(), plan); err != nil {
+				lg.Fatal().Err(err).Msg("failed to compile plan")
 			}
 		}
 
 		if !viper.GetBool("no-input") {
-				if err := cueVal.FillPath(cue.MakePath(), env.Input()); err != nil {
+			inputs, err := state.CompileInputs()
-					return err
-				}
-			}
-			return nil
-		})
-
-		<-doneCh
-
 			if err != nil {
-			lg.Fatal().Err(err).Msg("failed to query environment")
+				lg.Fatal().Err(err).Msg("failed to compile inputs")
+			}
+
+			if err := cueVal.FillPath(cue.MakePath(), inputs); err != nil {
+				lg.Fatal().Err(err).Msg("failed to compile inputs")
+			}
 		}
 
 		if !viper.GetBool("no-computed") && state.Computed != "" {
@@ -79,6 +75,8 @@ var queryCmd = &cobra.Command{
 			}
 		}
 
+		<-doneCh
+
 		cueVal = cueVal.LookupPath(cuePath)
 
 		if viper.GetBool("concrete") {
@@ -98,7 +96,7 @@ var queryCmd = &cobra.Command{
 		case "json":
 			fmt.Println(cueVal.JSON().PrettyString())
 		case "yaml":
-			lg.Fatal().Err(err).Msg("yaml format not yet implemented")
+			lg.Fatal().Msg("yaml format not yet implemented")
 		case "text":
 			out, err := cueVal.String()
 			if err != nil {

compiler/build.go

@@ -12,19 +12,16 @@ import (
 )
 
 // Build a cue configuration tree from the files in fs.
-func Build(sources map[string]fs.FS, args ...string) (*Value, error) {
+func Build(src string, overlays map[string]fs.FS, args ...string) (*Value, error) {
 	c := DefaultCompiler
 
 	buildConfig := &cueload.Config{
-		// The CUE overlay needs to be prefixed by a non-conflicting path with the
+		Dir:     src,
-		// local filesystem, otherwise Cue will merge the Overlay with whatever Cue
-		// files it finds locally.
-		Dir:     "/config",
 		Overlay: map[string]cueload.Source{},
 	}
 
 	// Map the source files into the overlay
-	for mnt, f := range sources {
+	for mnt, f := range overlays {
 		f := f
 		mnt := mnt
 		err := fs.WalkDir(f, ".", func(p string, entry fs.DirEntry, err error) error {

environment/environment.go

@@ -3,7 +3,6 @@ package environment
 import (
 	"context"
 	"fmt"
-	"io/fs"
 	"strings"
 	"time"
 
@@ -29,33 +28,38 @@ type Environment struct {
 	// Layer 2: user inputs
 	input *compiler.Value
 
+	// plan + inputs
+	src *compiler.Value
+
 	// Layer 3: computed values
 	computed *compiler.Value
 }
 
 func New(st *state.State) (*Environment, error) {
+	var err error
+
 	e := &Environment{
 		state: st,
-
-		plan:     compiler.NewValue(),
-		input:    compiler.NewValue(),
-		computed: compiler.NewValue(),
 	}
 
-	// Prepare inputs
+	e.plan, err = st.CompilePlan(context.TODO())
-	for key, input := range st.Inputs {
-		v, err := input.Compile(key, st)
 	if err != nil {
 		return nil, err
 	}
-		if key == "" {
+
-			err = e.input.FillPath(cue.MakePath(), v)
+	e.input, err = st.CompileInputs()
-		} else {
-			err = e.input.FillPath(cue.ParsePath(key), v)
-		}
 	if err != nil {
 		return nil, err
 	}
+
+	e.computed = compiler.NewValue()
+
+	e.src = compiler.NewValue()
+	if err := e.src.FillPath(cue.MakePath(), e.plan); err != nil {
+		return nil, err
+	}
+	if err := e.src.FillPath(cue.MakePath(), e.input); err != nil {
+		return nil, err
 	}
 
 	return e, nil
@@ -65,64 +69,10 @@ func (e *Environment) Name() string {
 	return e.state.Name
 }
 
-func (e *Environment) Plan() *compiler.Value {
-	return e.plan
-}
-
-func (e *Environment) Input() *compiler.Value {
-	return e.input
-}
-
 func (e *Environment) Computed() *compiler.Value {
 	return e.computed
 }
 
-// LoadPlan loads the plan
-func (e *Environment) LoadPlan(ctx context.Context, s solver.Solver) error {
-	tr := otel.Tracer("environment")
-	ctx, span := tr.Start(ctx, "environment.LoadPlan")
-	defer span.End()
-
-	// FIXME: universe vendoring
-	// This is already done on `dagger init` and shouldn't be done here too.
-	// However:
-	// 1) As of right now, there's no way to update universe through the
-	// CLI, so we are lazily updating on `dagger up` using the embedded `universe`
-	// 2) For backward compatibility: if the workspace was `dagger
-	// init`-ed before we added support for vendoring universe, it might not
-	// contain a `cue.mod`.
-	if err := e.state.VendorUniverse(ctx); err != nil {
-		return err
-	}
-
-	planSource, err := e.state.Source().Compile("", e.state)
-	if err != nil {
-		return err
-	}
-
-	p := NewPipeline(planSource, s).WithCustomName("[internal] source")
-	// execute updater script
-	if err := p.Run(ctx); err != nil {
-		return err
-	}
-
-	// Build a Cue config by overlaying the source with the stdlib
-	sources := map[string]fs.FS{
-		"/": p.FS(),
-	}
-	args := []string{}
-	if pkg := e.state.Plan.Package; pkg != "" {
-		args = append(args, pkg)
-	}
-	plan, err := compiler.Build(sources, args...)
-	if err != nil {
-		return fmt.Errorf("plan config: %w", compiler.Err(err))
-	}
-	e.plan = plan
-
-	return nil
-}
-
 // Scan all scripts in the environment for references to local directories (do:"local"),
 // and return all referenced directory names.
 // This is used by clients to grant access to local directories when they are referenced
@@ -168,58 +118,33 @@ func (e *Environment) LocalDirs() map[string]string {
 		localdirs(v.Lookup("#up"))
 	}
 
-	// 2. Scan the plan
-	plan, err := e.state.Source().Compile("", e.state)
-	if err != nil {
-		panic(err)
-	}
-	localdirs(plan)
 	return dirs
 }
 
-// prepare initializes the Environment with inputs and plan code
-func (e *Environment) prepare(ctx context.Context) (*compiler.Value, error) {
-	tr := otel.Tracer("environment")
-	_, span := tr.Start(ctx, "environment.Prepare")
-	defer span.End()
-
-	// Reset the computed values
-	e.computed = compiler.NewValue()
-
-	src := compiler.NewValue()
-	if err := src.FillPath(cue.MakePath(), e.plan); err != nil {
-		return nil, err
-	}
-	if err := src.FillPath(cue.MakePath(), e.input); err != nil {
-		return nil, err
-	}
-
-	return src, nil
-}
-
 // Up missing values in environment configuration, and write them to state.
 func (e *Environment) Up(ctx context.Context, s solver.Solver) error {
 	tr := otel.Tracer("environment")
 	ctx, span := tr.Start(ctx, "environment.Up")
 	defer span.End()
 
-	// Set user inputs and plan code
-	src, err := e.prepare(ctx)
-	if err != nil {
-		return err
-	}
-
 	// Orchestrate execution with cueflow
 	flow := cueflow.New(
 		&cueflow.Config{},
-		src.Cue(),
+		e.src.Cue(),
 		newTaskFunc(newPipelineRunner(e.computed, s)),
 	)
 	if err := flow.Run(ctx); err != nil {
 		return err
 	}
 
+	// FIXME: canceling the context makes flow return `nil`
+	// Check explicitly if the context is canceled.
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	default:
 		return nil
+	}
 }
 
 type DownOpts struct{}
@@ -328,20 +253,18 @@ func (e *Environment) ScanInputs(ctx context.Context, mergeUserInputs bool) ([]*
 	src := e.plan
 
 	if mergeUserInputs {
-		// Set user inputs and plan code
+		src = e.src
-		var err error
-		src, err = e.prepare(ctx)
-		if err != nil {
-			return nil, err
-		}
 	}
 
 	return ScanInputs(ctx, src), nil
 }
 
 func (e *Environment) ScanOutputs(ctx context.Context) ([]*compiler.Value, error) {
-	src, err := e.prepare(ctx)
+	src := compiler.NewValue()
-	if err != nil {
+	if err := src.FillPath(cue.MakePath(), e.plan); err != nil {
+		return nil, err
+	}
+	if err := src.FillPath(cue.MakePath(), e.input); err != nil {
 		return nil, err
 	}
 

environment/environment_test.go

@@ -1,26 +0,0 @@
-package environment
-
-import (
-	"testing"
-
-	"github.com/stretchr/testify/require"
-	"go.dagger.io/dagger/state"
-)
-
-func TestLocalDirs(t *testing.T) {
-	st := &state.State{
-		Path: "/tmp/source",
-		Plan: state.Plan{
-			Module: "/tmp/source/plan",
-		},
-	}
-	require.NoError(t, st.SetInput("www.source", state.DirInput("/", []string{}, []string{})))
-
-	environment, err := New(st)
-	require.NoError(t, err)
-
-	localdirs := environment.LocalDirs()
-	require.Len(t, localdirs, 2)
-	require.Contains(t, localdirs, "/")
-	require.Contains(t, localdirs, "/tmp/source/plan")
-}

environment/inputs_scan.go

@@ -4,7 +4,6 @@ import (
 	"context"
 
 	"cuelang.org/go/cue"
-	"github.com/rs/zerolog/log"
 	"go.dagger.io/dagger/compiler"
 )
 
@@ -43,13 +42,11 @@ func isReference(val cue.Value) bool {
 }
 
 func ScanInputs(ctx context.Context, value *compiler.Value) []*compiler.Value {
-	lg := log.Ctx(ctx)
 	inputs := []*compiler.Value{}
 
 	value.Walk(
 		func(val *compiler.Value) bool {
 			if isReference(val.Cue()) {
-				lg.Debug().Str("value.Path", val.Path().String()).Msg("found reference, stop walk")
 				return false
 			}
 
@@ -57,7 +54,6 @@ func ScanInputs(ctx context.Context, value *compiler.Value) []*compiler.Value {
 				return true
 			}
 
-			lg.Debug().Str("value.Path", val.Path().String()).Msg("found input")
 			inputs = append(inputs, val)
 
 			return true
@@ -68,7 +64,6 @@ func ScanInputs(ctx context.Context, value *compiler.Value) []*compiler.Value {
 }
 
 func ScanOutputs(ctx context.Context, value *compiler.Value) []*compiler.Value {
-	lg := log.Ctx(ctx)
 	inputs := []*compiler.Value{}
 
 	value.Walk(
@@ -77,7 +72,6 @@ func ScanOutputs(ctx context.Context, value *compiler.Value) []*compiler.Value {
 				return true
 			}
 
-			lg.Debug().Str("value.Path", val.Path().String()).Msg("found output")
 			inputs = append(inputs, val)
 
 			return true

state/state.go

@@ -3,6 +3,9 @@ package state
 import (
 	"context"
 	"path"
+
+	"cuelang.org/go/cue"
+	"go.dagger.io/dagger/compiler"
 )
 
 // Contents of an environment serialized to a file
@@ -29,13 +32,53 @@ type State struct {
 }
 
 // Cue module containing the environment plan
-func (s *State) Source() Input {
+func (s *State) CompilePlan(ctx context.Context) (*compiler.Value, error) {
 	w := s.Workspace
 	// FIXME: backward compatibility
 	if mod := s.Plan.Module; mod != "" {
 		w = path.Join(w, mod)
 	}
-	return DirInput(w, []string{}, []string{})
+
+	// FIXME: universe vendoring
+	// This is already done on `dagger init` and shouldn't be done here too.
+	// However:
+	// 1) As of right now, there's no way to update universe through the
+	// CLI, so we are lazily updating on `dagger up` using the embedded `universe`
+	// 2) For backward compatibility: if the workspace was `dagger
+	// init`-ed before we added support for vendoring universe, it might not
+	// contain a `cue.mod`.
+	if err := vendorUniverse(ctx, w); err != nil {
+		return nil, err
+	}
+
+	args := []string{}
+	if pkg := s.Plan.Package; pkg != "" {
+		args = append(args, pkg)
+	}
+
+	return compiler.Build(w, nil, args...)
+}
+
+func (s *State) CompileInputs() (*compiler.Value, error) {
+	v := compiler.NewValue()
+
+	// Prepare inputs
+	for key, input := range s.Inputs {
+		i, err := input.Compile(key, s)
+		if err != nil {
+			return nil, err
+		}
+		if key == "" {
+			err = v.FillPath(cue.MakePath(), i)
+		} else {
+			err = v.FillPath(cue.ParsePath(key), i)
+		}
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	return v, nil
 }
 
 // VendorUniverse vendors the latest (built-in) version of the universe into the