solver: allow finalizing history record traces

api/services/control/control.proto

@@ -220,6 +220,7 @@ message UpdateBuildHistoryRequest {
 	string Ref = 1;
 	bool Pinned = 2;
 	bool Delete = 3;
+	bool Finalize = 4;
 }
 
 message UpdateBuildHistoryResponse {}

control/control.go

@@ -283,18 +283,23 @@ func (c *Controller) ListenBuildHistory(req *controlapi.BuildHistoryRequest, srv
 }
 
 func (c *Controller) UpdateBuildHistory(ctx context.Context, req *controlapi.UpdateBuildHistoryRequest) (*controlapi.UpdateBuildHistoryResponse, error) {
-	if !req.Delete {
-		err := c.history.UpdateRef(ctx, req.Ref, func(r *controlapi.BuildHistoryRecord) error {
-			if req.Pinned == r.Pinned {
-				return nil
-			}
-			r.Pinned = req.Pinned
-			return nil
-		})
+	if req.Delete {
+		err := c.history.Delete(ctx, req.Ref)
 		return &controlapi.UpdateBuildHistoryResponse{}, err
 	}
 
-	err := c.history.Delete(ctx, req.Ref)
+	if req.Finalize {
+		err := c.history.Finalize(ctx, req.Ref)
+		return &controlapi.UpdateBuildHistoryResponse{}, err
+	}
+
+	err := c.history.UpdateRef(ctx, req.Ref, func(r *controlapi.BuildHistoryRecord) error {
+		if req.Pinned == r.Pinned {
+			return nil
+		}
+		r.Pinned = req.Pinned
+		return nil
+	})
 	return &controlapi.UpdateBuildHistoryResponse{}, err
 }
 

frontend/dockerfile/dockerfile_test.go

@@ -194,6 +194,7 @@ var allTests = integration.TestFuncs(
 	testSourcePolicyWithNamedContext,
 	testInvalidJSONCommands,
 	testHistoryError,
+	testHistoryFinalizeTrace,
 )
 
 // Tests that depend on the `security.*` entitlements
@@ -7521,6 +7522,66 @@ COPY notexist /foo
 	}
 }
 
+func testHistoryFinalizeTrace(t *testing.T, sb integration.Sandbox) {
+	integration.SkipOnPlatform(t, "windows")
+	workers.CheckFeatureCompat(t, sb, workers.FeatureDirectPush)
+	ctx := sb.Context()
+
+	c, err := client.New(ctx, sb.Address())
+	require.NoError(t, err)
+	defer c.Close()
+
+	f := getFrontend(t, sb)
+
+	dockerfile := []byte(`
+FROM scratch
+COPY Dockerfile /foo
+`)
+	dir := integration.Tmpdir(
+		t,
+		fstest.CreateFile("Dockerfile", dockerfile, 0600),
+	)
+
+	ref := identity.NewID()
+
+	_, err = f.Solve(sb.Context(), c, client.SolveOpt{
+		Ref: ref,
+		LocalMounts: map[string]fsutil.FS{
+			dockerui.DefaultLocalNameDockerfile: dir,
+			dockerui.DefaultLocalNameContext:    dir,
+		},
+	}, nil)
+	require.NoError(t, err)
+
+	_, err = c.ControlClient().UpdateBuildHistory(sb.Context(), &controlapi.UpdateBuildHistoryRequest{
+		Ref:      ref,
+		Finalize: true,
+	})
+	require.NoError(t, err)
+
+	cl, err := c.ControlClient().ListenBuildHistory(sb.Context(), &controlapi.BuildHistoryRequest{
+		EarlyExit: true,
+		Ref:       ref,
+	})
+	require.NoError(t, err)
+
+	got := false
+	for {
+		resp, err := cl.Recv()
+		if err == io.EOF {
+			require.Equal(t, true, got)
+			break
+		}
+		require.NoError(t, err)
+		got = true
+
+		trace := resp.Record.Trace
+		require.NotEmpty(t, trace)
+
+		require.NotEmpty(t, trace.Digest)
+	}
+}
+
 func runShell(dir string, cmds ...string) error {
 	for _, args := range cmds {
 		var cmd *exec.Cmd

solver/llbsolver/history.go

@@ -52,12 +52,20 @@ type HistoryQueue struct {
 	opt           HistoryQueueOpt
 	ps            *pubsub[*controlapi.BuildHistoryEvent]
 	active        map[string]*controlapi.BuildHistoryRecord
+	finalizers    map[string]*finalizer
 	refs          map[string]int
 	deleted       map[string]struct{}
 	hContentStore *containerdsnapshot.Store
 	hLeaseManager *leaseutil.Manager
 }
 
+// finalizer controls completion of saving traces for a
+// record and making it immutable
+type finalizer struct {
+	trigger func()
+	done    chan struct{}
+}
+
 type StatusImportResult struct {
 	Descriptor        ocispecs.Descriptor
 	NumCachedSteps    int
@@ -77,9 +85,10 @@ func NewHistoryQueue(opt HistoryQueueOpt) (*HistoryQueue, error) {
 		ps: &pubsub[*controlapi.BuildHistoryEvent]{
 			m: map[*channel[*controlapi.BuildHistoryEvent]]struct{}{},
 		},
-		active:  map[string]*controlapi.BuildHistoryRecord{},
-		refs:    map[string]int{},
-		deleted: map[string]struct{}{},
+		active:     map[string]*controlapi.BuildHistoryRecord{},
+		refs:       map[string]int{},
+		deleted:    map[string]struct{}{},
+		finalizers: map[string]*finalizer{},
 	}
 
 	ns := h.opt.ContentStore.Namespace()
@@ -570,6 +579,40 @@ func (h *HistoryQueue) update(ctx context.Context, rec controlapi.BuildHistoryRe
 	})
 }
 
+func (h *HistoryQueue) AcquireFinalizer(ref string) (<-chan struct{}, func()) {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+	trigger := make(chan struct{})
+	f := &finalizer{
+		trigger: sync.OnceFunc(func() {
+			close(trigger)
+		}),
+		done: make(chan struct{}),
+	}
+	h.finalizers[ref] = f
+	go func() {
+		<-f.done
+		h.mu.Lock()
+		delete(h.finalizers, ref)
+		h.mu.Unlock()
+	}()
+	return trigger, sync.OnceFunc(func() {
+		close(f.done)
+	})
+}
+
+func (h *HistoryQueue) Finalize(ctx context.Context, ref string) error {
+	h.mu.Lock()
+	f, ok := h.finalizers[ref]
+	h.mu.Unlock()
+	if !ok {
+		return nil
+	}
+	f.trigger()
+	<-f.done
+	return nil
+}
+
 func (h *HistoryQueue) Update(ctx context.Context, e *controlapi.BuildHistoryEvent) error {
 	h.init()
 	h.mu.Lock()

solver/llbsolver/solver.go

@@ -385,6 +385,9 @@ func (s *Solver) recordBuildHistory(ctx context.Context, id string, req frontend
 			releasers = append(releasers, release)
 			rec.Error = status
 		}
+
+		ready, done := s.history.AcquireFinalizer(rec.Ref)
+
 		if err1 := s.history.Update(ctx, &controlapi.BuildHistoryEvent{
 			Type:   controlapi.BuildHistoryEventType_COMPLETE,
 			Record: rec,
@@ -396,10 +399,17 @@ func (s *Solver) recordBuildHistory(ctx context.Context, id string, req frontend
 
 		if stopTrace == nil {
 			bklog.G(ctx).Warn("no trace recorder found, skipping")
+			done()
 			return err
 		}
 		go func() {
-			time.Sleep(3 * time.Second)
+			defer done()
+
+			// if there is no finalizer request then stop tracing after 3 seconds
+			select {
+			case <-time.After(3 * time.Second):
+			case <-ready:
+			}
 			spans := stopTrace()
 
 			if len(spans) == 0 {