progress: fix progress writer pause and unpause to prevent panics

This changes the progress printer's pause and unpause implementation to
be reentrant to prevent race conditions and it also allows the status
updates to be buffered when the display is paused.

The previous implementation mixed the pause implementation with the
finish implementation and could cause a send on closed channel panic
because it could close the status channel before it had finished being
used. Now, the status channel is not closed.

When the display is enabled, the status channel will be forwarded to an
internal channel that is used to display the updates. When the display
is paused, the status channel will have the statuses buffered in memory
to be sent when the progress display is resumed.

The `Unpause` method has also been renamed to `Resume`.

Signed-off-by: Jonathan A. Sternberg <[email protected]>

Author: jsternberg

monitor/monitor.go

@@ -113,10 +113,8 @@ func (m *Monitor) Close() error {
 
 // RunMonitor provides an interactive session for running and managing containers via specified IO.
 func RunMonitor(ctx context.Context, invokeConfig *build.InvokeConfig, rCtx *build.ResultHandle, stdin io.ReadCloser, stdout, stderr io.WriteCloser, progress *progress.Printer) error {
-	if err := progress.Pause(); err != nil {
-		return err
-	}
-	defer progress.Unpause()
+	progress.Pause()
+	defer progress.Resume()
 
 	defer stdin.Close()
 
@@ -475,10 +473,10 @@ func printError(err error, printer *progress.Printer) error {
 	if err == nil {
 		return nil
 	}
-	if err := printer.Pause(); err != nil {
-		return err
-	}
-	defer printer.Unpause()
+
+	printer.Pause()
+	defer printer.Resume()
+
 	for _, s := range errdefs.Sources(err) {
 		s.Print(os.Stderr)
 	}

util/progress/printer.go

@@ -16,12 +16,24 @@ import (
 	"go.opentelemetry.io/otel/metric"
 )
 
+type printerState int
+
+const (
+	printerStateDone printerState = iota
+	printerStateRunning
+	printerStatePaused
+)
+
 type Printer struct {
-	status chan *client.SolveStatus
+	out  console.File
+	mode progressui.DisplayMode
+	opt  *printerOpts
+
+	status    chan *client.SolveStatus
+	interrupt chan interruptRequest
+	state     printerState
 
-	ready     chan struct{}
 	done      chan struct{}
-	paused    chan struct{}
 	closeOnce sync.Once
 
 	err          error
@@ -54,13 +66,23 @@ func (p *Printer) IsDone() bool {
 }
 
 func (p *Printer) Pause() error {
-	p.paused = make(chan struct{})
-	return p.Wait()
+	done := make(chan struct{})
+	p.interrupt <- interruptRequest{
+		desiredState: printerStatePaused,
+		done:         done,
+	}
+
+	// Need to wait for a response to confirm we have control
+	// of the console output.
+	<-done
+	return nil
 }
 
-func (p *Printer) Unpause() {
-	close(p.paused)
-	<-p.ready
+func (p *Printer) Resume() {
+	p.interrupt <- interruptRequest{
+		desiredState: printerStateRunning,
+	}
+	// Do not care about waiting for a response.
 }
 
 func (p *Printer) Write(s *client.SolveStatus) {
@@ -115,42 +137,114 @@ func NewPrinter(ctx context.Context, out console.File, mode progressui.DisplayMo
 	}
 
 	pw := &Printer{
-		ready:   make(chan struct{}),
-		metrics: opt.mw,
+		out:       out,
+		mode:      mode,
+		opt:       opt,
+		status:    make(chan *client.SolveStatus),
+		interrupt: make(chan interruptRequest),
+		state:     printerStateRunning,
+		done:      make(chan struct{}),
 	}
+	go pw.run(ctx, d)
+
+	return pw, nil
+}
+
+func (p *Printer) run(ctx context.Context, d progressui.Display) {
+	defer close(p.done)
+	defer close(p.interrupt)
+
+	var ss []*client.SolveStatus
+	for p.state != printerStateDone {
+		switch p.state {
+		case printerStatePaused:
+			ss, p.err = p.bufferDisplay(ctx, ss)
+		case printerStateRunning:
+			var warnings []client.VertexWarning
+			warnings, ss, p.err = p.updateDisplay(ctx, d, ss)
+			p.warnings = append(p.warnings, warnings...)
+
+			d, _ = p.newDisplay()
+		}
+	}
+
+	if p.opt.onclose != nil {
+		p.opt.onclose()
+	}
+}
+
+func (p *Printer) newDisplay() (progressui.Display, error) {
+	return progressui.NewDisplay(p.out, p.mode, p.opt.displayOpts...)
+}
+
+func (p *Printer) updateDisplay(ctx context.Context, d progressui.Display, ss []*client.SolveStatus) ([]client.VertexWarning, []*client.SolveStatus, error) {
+	p.logMu.Lock()
+	p.logSourceMap = map[digest.Digest]any{}
+	p.logMu.Unlock()
+
+	resumeLogs := logutil.Pause(logrus.StandardLogger())
+	defer resumeLogs()
+
+	interruptCh := make(chan interruptRequest, 1)
+	ingress := make(chan *client.SolveStatus)
+
 	go func() {
+		defer close(ingress)
+		defer close(interruptCh)
+
+		for _, s := range ss {
+			ingress <- s
+		}
+
 		for {
-			pw.status = make(chan *client.SolveStatus)
-			pw.done = make(chan struct{})
-			pw.closeOnce = sync.Once{}
-
-			pw.logMu.Lock()
-			pw.logSourceMap = map[digest.Digest]any{}
-			pw.logMu.Unlock()
-
-			resumeLogs := logutil.Pause(logrus.StandardLogger())
-			close(pw.ready)
-			// not using shared context to not disrupt display but let is finish reporting errors
-			pw.warnings, pw.err = d.UpdateFrom(ctx, pw.status)
-			resumeLogs()
-			close(pw.done)
-
-			if opt.onclose != nil {
-				opt.onclose()
-			}
-			if pw.paused == nil {
-				break
+			select {
+			case s, ok := <-p.status:
+				if !ok {
+					return
+				}
+				ingress <- s
+			case req := <-p.interrupt:
+				interruptCh <- req
+				return
+			case <-ctx.Done():
+				return
 			}
+		}
+	}()
+
+	warnings, err := d.UpdateFrom(context.Background(), ingress)
+	if err == nil {
+		err = context.Cause(ctx)
+	}
 
-			pw.ready = make(chan struct{})
-			<-pw.paused
-			pw.paused = nil
+	interrupt := <-interruptCh
+	p.state = interrupt.desiredState
+	interrupt.close()
+	return warnings, nil, err
+}
 
-			d, _ = progressui.NewDisplay(out, mode, opt.displayOpts...)
+// bufferDisplay will buffer display updates from the status channel into a
+// slice.
+//
+// This method returns if either status gets closed or if an interrupt is received.
+func (p *Printer) bufferDisplay(ctx context.Context, ss []*client.SolveStatus) ([]*client.SolveStatus, error) {
+	for {
+		select {
+		case s, ok := <-p.status:
+			if !ok {
+				p.state = printerStateDone
+				return ss, nil
+			}
+			ss = append(ss, s)
+		case req := <-p.interrupt:
+			p.state = req.desiredState
+			req.close()
+			return ss, nil
+		case <-ctx.Done():
+			p.state = printerStateDone
+			return nil, context.Cause(ctx)
 		}
-	}()
-	<-pw.ready
-	return pw, nil
+	}
 }
 
 func (p *Printer) WriteBuildRef(target string, ref string) {
@@ -221,3 +315,14 @@ func dedupWarnings(inp []client.VertexWarning) []client.VertexWarning {
 	}
 	return res
 }
+
+type interruptRequest struct {
+	desiredState printerState
+	done         chan<- struct{}
+}
+
+func (req *interruptRequest) close() {
+	if req.done != nil {
+		close(req.done)
+	}
+}