solver: add jobCtx support to metadata resolver

Fixes issue where HTTP result can be released after metadata resolve
and rest of the build might pull new (potentially different) data in again.

Signed-off-by: Tonis Tiigi <[email protected]>

Author: tonistiigi

client/client_test.go

@@ -245,6 +245,7 @@ var allTests = []func(t *testing.T, sb integration.Sandbox){
 	testGitResolveSourceMetadata,
 	testHTTPResolveSourceMetadata,
 	testHTTPPruneAfterCacheKey,
+	testHTTPPruneAfterResolveMeta,
 }
 
 func TestIntegration(t *testing.T) {
@@ -12208,6 +12209,67 @@ func testHTTPPruneAfterCacheKey(t *testing.T, sb integration.Sandbox) {
 	<-done
 }
 
+// testHTTPPruneAfterResolveMeta ensures that pruning after ResolveSourceMetadata
+// doesn't pull in new data for same build. Once URL has been resolved once for a specific
+// build, the data should be considered immutable and remote changes don't affect ongoing build.
+func testHTTPPruneAfterResolveMeta(t *testing.T, sb integration.Sandbox) {
+	ctx := sb.Context()
+	c, err := New(ctx, sb.Address())
+	require.NoError(t, err)
+	defer c.Close()
+
+	resp := &httpserver.Response{
+		Etag:    identity.NewID(),
+		Content: []byte("content1"),
+	}
+	server := httpserver.NewTestServer(map[string]*httpserver.Response{
+		"/foo": resp,
+	})
+	defer server.Close()
+
+	dest := t.TempDir()
+
+	_, err = c.Build(ctx, SolveOpt{
+		Exports: []ExportEntry{
+			{
+				Type:      ExporterLocal,
+				OutputDir: dest,
+			},
+		},
+	}, "test", func(ctx context.Context, gc gateway.Client) (*gateway.Result, error) {
+		id := server.URL + "/foo"
+		md, err := gc.ResolveSourceMetadata(ctx, &pb.SourceOp{
+			Identifier: id,
+		}, sourceresolver.Opt{})
+		if err != nil {
+			return nil, err
+		}
+		require.NotNil(t, md.HTTP)
+
+		// prune all
+		err = c.Prune(ctx, nil)
+		require.NoError(t, err)
+
+		resp.Content = []byte("content2") // etag is same so should hit cache if record not pruned
+
+		st := llb.Scratch().File(llb.Copy(llb.HTTP(id), "foo", "bar"))
+		def, err := st.Marshal(sb.Context())
+		if err != nil {
+			return nil, err
+		}
+		return gc.Solve(ctx, gateway.SolveRequest{
+			Definition: def.ToPB(),
+		})
+	}, nil)
+	require.NoError(t, err)
+
+	dt, err := os.ReadFile(filepath.Join(dest, "bar"))
+	require.NoError(t, err)
+	require.Equal(t, "content1", string(dt))
+
+	checkAllReleasable(t, c, sb, false)
+}
+
 func runInDir(dir string, cmds ...string) error {
 	for _, args := range cmds {
 		var cmd *exec.Cmd

solver/jobs.go

@@ -28,7 +28,7 @@ type ResolveOpFunc func(Vertex, Builder) (Op, error)
 
 type Builder interface {
 	Build(ctx context.Context, e Edge) (CachedResultWithProvenance, error)
-	InContext(ctx context.Context, f func(ctx context.Context, g session.Group) error) error
+	InContext(ctx context.Context, f func(ctx context.Context, jobCtx JobContext) error) error
 	EachValue(ctx context.Context, key string, fn func(any) error) error
 }
 
@@ -299,7 +299,7 @@ func (sb *subBuilder) Build(ctx context.Context, e Edge) (CachedResultWithProven
 	return &withProvenance{CachedResult: res}, nil
 }
 
-func (sb *subBuilder) InContext(ctx context.Context, f func(context.Context, session.Group) error) error {
+func (sb *subBuilder) InContext(ctx context.Context, f func(context.Context, JobContext) error) error {
 	ctx = progress.WithProgress(ctx, sb.mpw)
 	if sb.mspan.Span != nil {
 		ctx = trace.ContextWithSpan(ctx, sb.mspan)
@@ -328,6 +328,7 @@ type Job struct {
 	id            string
 	startedTime   time.Time
 	completedTime time.Time
+	releasers     []func() error
 
 	progressCloser func(error)
 	SessionID      string
@@ -812,6 +813,13 @@ func (j *Job) Discard() error {
 		st.mu.Unlock()
 	}
 
+	for _, r := range j.releasers {
+		if err := r(); err != nil {
+			bklog.G(context.TODO()).WithError(err).Error("failed to cleanup job resources")
+		}
+	}
+	j.releasers = nil
+
 	go func() {
 		// don't clear job right away. there might still be a status request coming to read progress
 		time.Sleep(10 * time.Second)
@@ -839,8 +847,19 @@ func (j *Job) UniqueID() string {
 	return j.uniqueID
 }
 
-func (j *Job) InContext(ctx context.Context, f func(context.Context, session.Group) error) error {
-	return f(progress.WithProgress(ctx, j.pw), session.NewGroup(j.SessionID))
+func (j *Job) InContext(ctx context.Context, f func(context.Context, JobContext) error) error {
+	return f(progress.WithProgress(ctx, j.pw), j)
+}
+
+func (j *Job) Session() session.Group {
+	return session.NewGroup(j.SessionID)
+}
+
+func (j *Job) Cleanup(fn func() error) error {
+	j.mu.Lock()
+	j.releasers = append(j.releasers, fn)
+	j.mu.Unlock()
+	return nil
 }
 
 func (j *Job) SetValue(key string, v any) {

solver/llbsolver/bridge.go

@@ -49,7 +49,7 @@ type llbBridge struct {
 }
 
 func (b *llbBridge) Warn(ctx context.Context, dgst digest.Digest, msg string, opts frontend.WarnOpts) error {
-	return b.builder.InContext(ctx, func(ctx context.Context, g session.Group) error {
+	return b.builder.InContext(ctx, func(ctx context.Context, _ solver.JobContext) error {
 		pw, ok, _ := progress.NewFromContext(ctx, progress.WithMetadata("vertex", dgst))
 		if !ok {
 			return nil
@@ -111,11 +111,15 @@ func (b *llbBridge) loadResult(ctx context.Context, def *pb.Definition, cacheImp
 			func(cmID string, im gw.CacheOptionsEntry) {
 				cm = newLazyCacheManager(cmID, func() (solver.CacheManager, error) {
 					var cmNew solver.CacheManager
-					if err := inBuilderContext(context.TODO(), b.builder, "importing cache manifest from "+cmID, "", func(ctx context.Context, g session.Group) error {
+					if err := inBuilderContext(context.TODO(), b.builder, "importing cache manifest from "+cmID, "", func(ctx context.Context, jobCtx solver.JobContext) error {
 						resolveCI, ok := b.resolveCacheImporterFuncs[im.Type]
 						if !ok {
 							return errors.Errorf("unknown cache importer: %s", im.Type)
 						}
+						var g session.Group
+						if jobCtx != nil {
+							g = jobCtx.Session()
+						}
 						ci, desc, err := resolveCI(ctx, g, im.Attrs)
 						if err != nil {
 							return errors.Wrapf(err, "failed to configure %v cache importer", im.Type)
@@ -373,8 +377,8 @@ func (b *llbBridge) ResolveSourceMetadata(ctx context.Context, op *pb.SourceOp,
 	// policy is evaluated, so we can remove it from the options
 	opt.SourcePolicies = nil
 
-	err = inBuilderContext(ctx, b.builder, opt.LogName, id, func(ctx context.Context, g session.Group) error {
-		resp, err = w.ResolveSourceMetadata(ctx, op, opt, b.sm, dummyJobContext{g: g})
+	err = inBuilderContext(ctx, b.builder, opt.LogName, id, func(ctx context.Context, jobCtx solver.JobContext) error {
+		resp, err = w.ResolveSourceMetadata(ctx, op, opt, b.sm, jobCtx)
 		return err
 	})
 	if err != nil {
@@ -383,18 +387,6 @@ func (b *llbBridge) ResolveSourceMetadata(ctx context.Context, op *pb.SourceOp,
 	return resp, nil
 }
 
-type dummyJobContext struct {
-	g session.Group
-}
-
-func (d dummyJobContext) Session() session.Group {
-	return d.g
-}
-
-func (d dummyJobContext) Cleanup(fn func() error) error {
-	return errors.Errorf("cleanup not implemented for %T", d)
-}
-
 type lazyCacheManager struct {
 	id   string
 	main solver.CacheManager

solver/llbsolver/solver.go

@@ -776,7 +776,7 @@ func runCacheExporters(ctx context.Context, exporters []RemoteCacheExporter, j *
 		i, exp := i, exp
 		eg.Go(func() (err error) {
 			id := fmt.Sprint(j.SessionID, "-cache-", i)
-			err = inBuilderContext(ctx, j, exp.Name(), id, func(ctx context.Context, _ session.Group) error {
+			err = inBuilderContext(ctx, j, exp.Name(), id, func(ctx context.Context, _ solver.JobContext) error {
 				prepareDone := progress.OneOff(ctx, "preparing build cache for export")
 				if err := result.EachRef(cached, inp, func(res solver.CachedResult, ref cache.ImmutableRef) error {
 					ctx := withDescHandlerCacheOpts(ctx, ref)
@@ -852,7 +852,7 @@ func (s *Solver) runExporters(ctx context.Context, exporters []exporter.Exporter
 		i, exp := i, exp
 		eg.Go(func() error {
 			id := fmt.Sprint(job.SessionID, "-export-", i)
-			return inBuilderContext(ctx, job, exp.Name(), id, func(ctx context.Context, _ session.Group) error {
+			return inBuilderContext(ctx, job, exp.Name(), id, func(ctx context.Context, _ solver.JobContext) error {
 				span, ctx := tracing.StartSpan(ctx, exp.Name())
 				defer span.End()
 
@@ -884,7 +884,7 @@ func (s *Solver) runExporters(ctx context.Context, exporters []exporter.Exporter
 	}
 
 	if len(exporters) == 0 && len(warnings) > 0 {
-		err := inBuilderContext(ctx, job, "Verifying build result", identity.NewID(), func(ctx context.Context, _ session.Group) error {
+		err := inBuilderContext(ctx, job, "Verifying build result", identity.NewID(), func(ctx context.Context, _ solver.JobContext) error {
 			pw, _, _ := progress.NewFromContext(ctx)
 			for _, w := range warnings {
 				pw.Write(identity.NewID(), w)
@@ -1152,19 +1152,19 @@ func allWorkers(wc *worker.Controller) func(func(w worker.Worker) error) error {
 	}
 }
 
-func inBuilderContext(ctx context.Context, b solver.Builder, name, id string, f func(ctx context.Context, g session.Group) error) error {
+func inBuilderContext(ctx context.Context, b solver.Builder, name, id string, f func(ctx context.Context, jobCtx solver.JobContext) error) error {
 	if id == "" {
 		id = name
 	}
 	v := client.Vertex{
 		Digest: digest.FromBytes([]byte(id)),
 		Name:   name,
 	}
-	return b.InContext(ctx, func(ctx context.Context, g session.Group) error {
+	return b.InContext(ctx, func(ctx context.Context, jobCtx solver.JobContext) error {
 		pw, _, ctx := progress.NewFromContext(ctx, progress.WithMetadata("vertex", v.Digest))
 		notifyCompleted := notifyStarted(ctx, &v)
 		defer pw.Close()
-		err := f(ctx, g)
+		err := f(ctx, jobCtx)
 		notifyCompleted(err)
 		return err
 	})

source/http/source.go

@@ -386,7 +386,10 @@ func (hs *httpSourceHandler) resolveMetadata(ctx context.Context, jobCtx solver.
 		return ref.Release(context.TODO())
 	}
 	if jobCtx != nil {
-		jobCtx.Cleanup(cleanup)
+		if err := jobCtx.Cleanup(cleanup); err != nil {
+			_ = cleanup()
+			return nil, err
+		}
 	} else {
 		cleanup()
 	}