progress: fix sync issue in controller.

A panic was encountered where writes to progress status hit a nil
writer. While not easy to reproduce, this commit fixes a likely culprit.
The use of atomics in the controller was not safe against a race such as
the following:
1. One goroutine calls Start, incrementing count but not yet setting
   writer.
2. A second goroutine calls Start, increments count again but sees that
   count >1 and thus doesn't try to set writer.
3. The second goroutine then calls Status, assuming the writer has been
   set, but the first goroutine still hasn't set the writer yet, causing
   a nil pointer exception.

This commit fixes that issue by just using a mutex instead of atomics.
It also adds a nil check for the writer just to be safe against panics
due to unknown issues in the future as missing a status update is much
better than buildkitd crashing.

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

util/progress/controller/controller.go

@@ -2,7 +2,7 @@ package controller
 
 import (
 	"context"
-	"sync/atomic"
+	"sync"
 	"time"
 
 	"github.com/moby/buildkit/client"
@@ -15,6 +15,7 @@ type Controller struct {
 	count   int64
 	started *time.Time
 	writer  progress.Writer
+	mu      sync.Mutex
 
 	Digest        digest.Digest
 	Name          string
@@ -25,7 +26,10 @@ type Controller struct {
 var _ progress.Controller = &Controller{}
 
 func (c *Controller) Start(ctx context.Context) (context.Context, func(error)) {
-	if atomic.AddInt64(&c.count, 1) == 1 {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.count++
+	if c.count == 1 {
 		if c.started == nil {
 			now := time.Now()
 			c.started = &now
@@ -42,7 +46,10 @@ func (c *Controller) Start(ctx context.Context) (context.Context, func(error)) {
 		}
 	}
 	return progress.WithProgress(ctx, c.writer), func(err error) {
-		if atomic.AddInt64(&c.count, -1) == 0 {
+		c.mu.Lock()
+		defer c.mu.Unlock()
+		c.count--
+		if c.count == 0 {
 			now := time.Now()
 			var errString string
 			if err != nil {
@@ -59,18 +66,22 @@ func (c *Controller) Start(ctx context.Context) (context.Context, func(error)) {
 				})
 			}
 			c.writer.Close()
+			c.started = nil
 		}
 	}
 }
 
 func (c *Controller) Status(id string, action string) func() {
 	start := time.Now()
+	if c.writer != nil {
 		c.writer.Write(id, progress.Status{
 			Action:  action,
 			Started: &start,
 		})
+	}
 	return func() {
 		complete := time.Now()
+		if c.writer != nil {
 			c.writer.Write(id, progress.Status{
 				Action:    action,
 				Started:   &start,
@@ -78,3 +89,4 @@ func (c *Controller) Status(id string, action string) func() {
 			})
 		}
 	}
+}