skip deduping running nested execs (dagger#10264)

* skip deduping running nested exec function calls

De-duping currently running nested execs, such as module function calls, can break when:
1. Client A starts the function call
2. Client B makes the same call, gets deduped with the existing one
3. Client A cancels the call (i.e. ctrl-C's, CI job shuts down for whatever reason, etc.)
4. The call continues running, but now is trying to talk back to a session that doesn't exist anymore when dagger APIs are invoked

Errors vary, but example from the new integ test that repros this:
```
    testctx.go:197: 29  : │ [15.0s] | failed to return error: input: get or init client: initialize client: failed to load module: load moduleSource(refString: ".").asModule: Module!: load base: load: failed to stat local path: failed to stat path: failed to get requester session: session for "b9y05fidme5ahpym64tr8ymq5" not found
    testctx.go:197: 29  : │ [15.0s] |
    testctx.go:197: 29  : │ [15.0s] | original error: input: get or init client: initialize client: failed to load module: load moduleSource(refString: ".").asModule: Module!: load base: load: failed to stat local path: failed to stat path: failed to get requester session: session for "b9y05fidme5ahpym64tr8ymq5" not found
    testctx.go:197: 29  : │ [15.0s] | Stderr:
    testctx.go:197: 29  : │ [15.0s] | time=2025-04-24T20:39:19.608Z level=ERROR msg="failed to emit telemetry" error="failed to upload metrics: failed to send metrics to http://127.0.0.1:38113/v1/metrics: 400 Bad Request"
```

Fixing this by adding support for skipping the dedupe behavior for selected calls.

Signed-off-by: Erik Sipsma <erik@sipsma.dev>

* use optional buildkit edge merging

Signed-off-by: Erik Sipsma <erik@sipsma.dev>

* cleanup zillions of return vals

Signed-off-by: Erik Sipsma <erik@sipsma.dev>

* add extra unit test coverage for skip dedupe

Signed-off-by: Erik Sipsma <erik@sipsma.dev>

* add changelog

Signed-off-by: Erik Sipsma <erik@sipsma.dev>

* de-dupe all dagql calls for now

There were occasional failures in the newly added integ test from
non-nested calls. This turned out to be because buildkit refs were
attempted to be used from the session that closed, which makes sense
since those buildkit refs get released when the session closes.

This is fairly tricky to address while buildkit refs are still being
used in our codebase, so this disables all de-duping on the dagql level
for now. The ability to de-dupe is left in (and still used in other
codepaths that use that code, namely filesync), so it can be re-enabled
in the future.

So far there doesn't seem to be a huge difference in test runtime, so
the performance impact is likely not that great. The dagql cache is
still enabled so we only lose de-duping of exactly overlapping
concurrent calls.

Signed-off-by: Erik Sipsma <erik@sipsma.dev>

---------

Signed-off-by: Erik Sipsma <erik@sipsma.dev>

Author: sipsma

.changes/unreleased/Fixed-20250501-184424.yaml

@@ -0,0 +1,7 @@
+kind: Fixed
+body: Fix "failed to return error" and "failed to emit telemetry" errors when two
+  identical functions execute at the same time and one client cancels the request.
+time: 2025-05-01T18:44:24.032929067-07:00
+custom:
+  Author: sipsma
+  PR: "10264"

core/container_exec.go

@@ -308,6 +308,9 @@ func (container *Container) WithExec(ctx context.Context, opts ContainerExecOpts
 		llb.Platform(platform.Spec()),
 		execMDOpt,
 	}
+	if opts.ExperimentalPrivilegedNesting {
+		marshalOpts = append(marshalOpts, llb.SkipEdgeMerge)
+	}
 	execDef, err := execSt.Root().Marshal(ctx, marshalOpts...)
 	if err != nil {
 		return nil, fmt.Errorf("marshal root: %w", err)

core/integration/cross_session_test.go

@@ -18,6 +18,7 @@ import (
 	"github.com/moby/buildkit/identity"
 	"github.com/stretchr/testify/require"
 	fs "github.com/tonistiigi/fsutil/copy"
+	"golang.org/x/sync/errgroup"
 )
 
 // TODO: add more tests for longer chains of cache hits that then diverge (i.e. constructor + some function cache hit, then diverge)
@@ -1518,3 +1519,71 @@ func (*Test) Fn2(ctx context.Context, secret *dagger.Secret) *dagger.Container {
 		}
 	})
 }
+
+func (ModuleSuite) TestCrossSessionDedupeOfNestedExec(ctx context.Context, t *testctx.T) {
+	callMod := func(c *dagger.Client) error {
+		_, err := goGitBase(t, c).
+			WithWorkdir("/work").
+			WithEnvVariable("CACHEBUSTER", identity.NewID()).
+			With(daggerExec("init", "--source=.", "--name=test", "--sdk=go")).
+			WithNewFile("main.go", `package main
+
+import (
+	"context"
+	"fmt"
+	"strconv"
+	"time"
+)
+
+type Test struct {}
+
+func (Test) Fn(ctx context.Context) error {
+	ctr, err := dag.Container().
+		From("alpine:3.20").
+		WithExec([]string{"sh", "-c", "echo "+strconv.Itoa(int(time.Now().UnixNano()))+"> /foo.txt"}).
+		Sync(ctx)
+	if err != nil {
+		return err
+	}
+
+	fmt.Println("sleeping", time.Now().UnixNano())
+	time.Sleep(20 * time.Second)
+	fmt.Println("awoken", time.Now().UnixNano())
+
+	ctr, err = ctr.WithExec([]string{"true"}).Sync(ctx)
+	return err
+}
+	`,
+			).
+			With(daggerCall("fn")).
+			Sync(ctx)
+		return err
+	}
+
+	c1 := connect(ctx, t)
+	c2 := connect(ctx, t)
+
+	var eg errgroup.Group
+
+	eg.Go(func() error {
+		time.Sleep(10 * time.Second)
+		t.Log("closing c1")
+		c1.Close()
+		t.Log("closed c1")
+		return nil
+	})
+
+	eg.Go(func() error {
+		callMod(c1)
+		t.Log("c1 call complete")
+		return nil
+	})
+
+	eg.Go(func() error {
+		time.Sleep(5 * time.Second)
+		return callMod(c2)
+	})
+
+	err := eg.Wait()
+	require.NoError(t, err)
+}

core/postcall.go

@@ -143,7 +143,7 @@ func ResourceTransferPostCall(
 			// The longer term fix for this type of issue is to have more dagql awareness of edges between
 			// cache results such that a function call return value result inherently results in any referenced
 			// secrets also staying in cache.
-			_, err = destDag.Cache.GetOrInitializeWithCallbacks(ctx, secret.inst.ID().Digest(),
+			_, err = destDag.Cache.GetOrInitializeWithCallbacks(ctx, secret.inst.ID().Digest(), true,
 				func(ctx context.Context) (*dagql.CacheValWithCallbacks, error) {
 					return &dagql.CacheValWithCallbacks{
 						Value:    secret.inst,

dagql/objects.go

@@ -359,30 +359,36 @@ func NoopDone(res Typed, cached bool, rerr error) {}
 
 // Select calls the field on the instance specified by the selector
 func (r Instance[T]) Select(ctx context.Context, s *Server, sel Selector) (Typed, *call.ID, error) {
-	inputArgs, newID, doNotCache, err := r.preselect(ctx, s, sel)
+	preselectResult, err := r.preselect(ctx, s, sel)
 	if err != nil {
 		return nil, nil, err
 	}
-	return r.call(ctx, s, newID, inputArgs, doNotCache)
+	return r.call(ctx, s, preselectResult.newID, preselectResult.inputArgs, preselectResult.doNotCache)
 }
 
 func (r Instance[T]) ReturnType(ctx context.Context, s *Server, sel Selector) (Typed, *call.ID, error) {
-	_, newID, _, err := r.preselect(ctx, s, sel)
+	preselectResult, err := r.preselect(ctx, s, sel)
 	if err != nil {
 		return nil, nil, err
 	}
-	returnType, err := r.returnType(newID)
+	returnType, err := r.returnType(preselectResult.newID)
 	if err != nil {
 		return nil, nil, err
 	}
-	return returnType, newID, nil
+	return returnType, preselectResult.newID, nil
+}
+
+type preselectResult struct {
+	inputArgs  map[string]Input
+	newID      *call.ID
+	doNotCache bool
 }
 
-func (r Instance[T]) preselect(ctx context.Context, s *Server, sel Selector) (map[string]Input, *call.ID, bool, error) {
+func (r Instance[T]) preselect(ctx context.Context, s *Server, sel Selector) (*preselectResult, error) {
 	view := sel.View
 	field, ok := r.Class.Field(sel.Field, view)
 	if !ok {
-		return nil, nil, false, fmt.Errorf("Select: %s has no such field: %q", r.Class.TypeName(), sel.Field)
+		return nil, fmt.Errorf("Select: %s has no such field: %q", r.Class.TypeName(), sel.Field)
 	}
 	if field.Spec.ViewFilter == nil {
 		// fields in the global view shouldn't attach the current view to the
@@ -417,7 +423,7 @@ func (r Instance[T]) preselect(ctx context.Context, s *Server, sel Selector) (ma
 
 		case argSpec.Type.Type().NonNull:
 			// error out if the arg is missing but required
-			return nil, nil, false, fmt.Errorf("missing required argument: %q", argSpec.Name)
+			return nil, fmt.Errorf("missing required argument: %q", argSpec.Name)
 		}
 	}
 	// TODO: it's better DX if it matches schema order
@@ -450,7 +456,7 @@ func (r Instance[T]) preselect(ctx context.Context, s *Server, sel Selector) (ma
 			Digest: origDgst,
 		})
 		if err != nil {
-			return nil, nil, false, fmt.Errorf("failed to compute cache key for %s.%s: %w", r.Type().Name(), sel.Field, err)
+			return nil, fmt.Errorf("failed to compute cache key for %s.%s: %w", r.Type().Name(), sel.Field, err)
 		}
 
 		if len(cacheCfg.UpdatedArgs) > 0 {
@@ -488,7 +494,11 @@ func (r Instance[T]) preselect(ctx context.Context, s *Server, sel Selector) (ma
 		}
 	}
 
-	return inputArgs, newID, doNotCache, nil
+	return &preselectResult{
+		inputArgs:  inputArgs,
+		newID:      newID,
+		doNotCache: doNotCache,
+	}, nil
 }
 
 // Call calls the field on the instance specified by the ID.
@@ -554,7 +564,7 @@ func (r Instance[T]) call(
 			return s.telemetry(ctx, r, newID)
 		}))
 	}
-	res, err := s.Cache.GetOrInitializeWithCallbacks(ctx, callCacheKey, func(ctx context.Context) (*CacheValWithCallbacks, error) {
+	res, err := s.Cache.GetOrInitializeWithCallbacks(ctx, callCacheKey, true, func(ctx context.Context) (*CacheValWithCallbacks, error) {
 		valWithCallbacks, err := r.Class.Call(ctx, r, newID.Field(), View(newID.View()), inputArgs)
 		if err != nil {
 			return nil, err

dagql/session_cache.go

@@ -76,7 +76,7 @@ func (c *SessionCache) GetOrInitialize(
 	fn func(context.Context) (CacheValueType, error),
 	opts ...CacheCallOpt,
 ) (CacheResult, error) {
-	return c.GetOrInitializeWithCallbacks(ctx, key, func(ctx context.Context) (*CacheValWithCallbacks, error) {
+	return c.GetOrInitializeWithCallbacks(ctx, key, false, func(ctx context.Context) (*CacheValWithCallbacks, error) {
 		val, err := fn(ctx)
 		if err != nil {
 			return nil, err
@@ -88,6 +88,7 @@ func (c *SessionCache) GetOrInitialize(
 func (c *SessionCache) GetOrInitializeWithCallbacks(
 	ctx context.Context,
 	key CacheKeyType,
+	skipDedupe bool,
 	fn func(context.Context) (*CacheValWithCallbacks, error),
 	opts ...CacheCallOpt,
 ) (res CacheResult, err error) {
@@ -114,7 +115,7 @@ func (c *SessionCache) GetOrInitializeWithCallbacks(
 		ctx = telemetryCtx
 	}
 
-	res, err = c.cache.GetOrInitializeWithCallbacks(ctx, key, fn)
+	res, err = c.cache.GetOrInitializeWithCallbacks(ctx, key, skipDedupe, fn)
 	if err != nil {
 		return nil, err
 	}

engine/buildkit/ref.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"os"
 	"path"
+	"runtime/debug"
 	"strconv"
 	"strings"
 
@@ -240,7 +241,9 @@ func (r *ref) Result(ctx context.Context) (bksolver.CachedResult, error) {
 	res, err := r.resultProxy.Result(ctx)
 	if err != nil {
 		// writing log w/ %+v so that we can see stack traces embedded in err by buildkit's usage of pkg/errors
-		bklog.G(ctx).Errorf("ref evaluate error: %+v", err)
+		bklog.G(ctx).
+			WithField("caller stack", string(debug.Stack())).
+			Errorf("ref evaluate error: %+v", err)
 		err = includeBuildkitContextCancelledLine(err)
 		return nil, WrapError(ctx, err, r.c)
 	}

engine/cache/cache.go

@@ -26,6 +26,7 @@ type Cache[K comparable, V any] interface {
 	GetOrInitializeWithCallbacks(
 		context.Context,
 		K,
+		bool,
 		func(context.Context) (*ValueWithCallbacks[V], error),
 	) (Result[K, V], error)
 
@@ -60,8 +61,9 @@ func NewCache[K comparable, V any]() Cache[K, V] {
 }
 
 type cache[K comparable, V any] struct {
-	mu    sync.Mutex
-	calls map[K]*result[K, V]
+	mu             sync.Mutex
+	ongoingCalls   map[K]*result[K, V]
+	completedCalls map[K]*result[K, V]
 }
 
 var _ Cache[int, int] = &cache[int, int]{}
@@ -107,7 +109,7 @@ func (c *cache[K, V]) Size() int {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
-	return len(c.calls)
+	return len(c.ongoingCalls) + len(c.completedCalls)
 }
 
 func (c *cache[K, V]) GetOrInitializeValue(
@@ -125,7 +127,7 @@ func (c *cache[K, V]) GetOrInitialize(
 	key K,
 	fn func(context.Context) (V, error),
 ) (Result[K, V], error) {
-	return c.GetOrInitializeWithCallbacks(ctx, key, func(ctx context.Context) (*ValueWithCallbacks[V], error) {
+	return c.GetOrInitializeWithCallbacks(ctx, key, false, func(ctx context.Context) (*ValueWithCallbacks[V], error) {
 		val, err := fn(ctx)
 		if err != nil {
 			return nil, err
@@ -137,6 +139,7 @@ func (c *cache[K, V]) GetOrInitialize(
 func (c *cache[K, V]) GetOrInitializeWithCallbacks(
 	ctx context.Context,
 	key K,
+	skipDedupe bool,
 	fn func(context.Context) (*ValueWithCallbacks[V], error),
 ) (Result[K, V], error) {
 	var zeroKey K
@@ -160,15 +163,29 @@ func (c *cache[K, V]) GetOrInitializeWithCallbacks(
 	}
 
 	c.mu.Lock()
-	if c.calls == nil {
-		c.calls = make(map[K]*result[K, V])
+	if c.ongoingCalls == nil {
+		c.ongoingCalls = make(map[K]*result[K, V])
+	}
+	if c.completedCalls == nil {
+		c.completedCalls = make(map[K]*result[K, V])
 	}
 
-	if res, ok := c.calls[key]; ok {
-		// already an ongoing call
-		res.waiters++
+	if res, ok := c.completedCalls[key]; ok {
+		res.refCount++
 		c.mu.Unlock()
-		return c.wait(ctx, key, res)
+		return &perCallResult[K, V]{
+			result:   res,
+			hitCache: true,
+		}, nil
+	}
+
+	if !skipDedupe {
+		if res, ok := c.ongoingCalls[key]; ok {
+			// already an ongoing call
+			res.waiters++
+			c.mu.Unlock()
+			return c.wait(ctx, key, res)
+		}
 	}
 
 	// make a new call with ctx that's only canceled when all caller contexts are canceled
@@ -183,7 +200,11 @@ func (c *cache[K, V]) GetOrInitializeWithCallbacks(
 		cancel:  cancel,
 		waiters: 1,
 	}
-	c.calls[key] = res
+
+	if !skipDedupe {
+		c.ongoingCalls[key] = res
+	}
+
 	go func() {
 		defer close(res.waitCh)
 		valWithCallbacks, err := fn(callCtx)
@@ -237,6 +258,13 @@ func (c *cache[K, V]) wait(ctx context.Context, key K, res *result[K, V]) (*perC
 	}
 
 	if err == nil {
+		delete(c.ongoingCalls, key)
+		if existingRes, ok := c.completedCalls[key]; ok {
+			res = existingRes
+		} else {
+			c.completedCalls[key] = res
+		}
+
 		res.refCount++
 		return &perCallResult[K, V]{
 			result:   res,
@@ -246,7 +274,8 @@ func (c *cache[K, V]) wait(ctx context.Context, key K, res *result[K, V]) (*perC
 
 	if res.refCount == 0 {
 		// error happened and no refs left, delete it now
-		delete(c.calls, key)
+		delete(c.ongoingCalls, key)
+		delete(c.completedCalls, key)
 	}
 	return nil, err
 }
@@ -270,7 +299,8 @@ func (res *result[K, V]) Release(ctx context.Context) error {
 	var onRelease OnReleaseFunc
 	if res.refCount == 0 && res.waiters == 0 {
 		// no refs left and no one waiting on it, delete from cache
-		delete(res.cache.calls, res.key)
+		delete(res.cache.ongoingCalls, res.key)
+		delete(res.cache.completedCalls, res.key)
 		onRelease = res.onRelease
 	}
 	res.cache.mu.Unlock()