solver: allow finalizing history record traces

Client can send a finalize update to build record that
will complete saving the traces and block until the record
has been updated. If no request is sent then the traces will be
sent after a 3 second timeout as before.

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

Author: tonistiigi

api/services/control/control.pb.go

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

api/services/control/control.proto

@@ -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
 }
 

control/control.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

frontend/dockerfile/dockerfile_test.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/history.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 {