Add Wait to binaryProcessor

helsaawy · helsaawy · commit ad8b87ba234e · 2022-05-09T17:15:00.000-04:00
Add exported `Wait(ctx context.Context) error` interface that waits on the underlying command (or context cancellation) and returns the error. This fixes a race condition between `.wait()` and `.Err error`: containerd#6914 Signed-off-by: Hamza El-Saawy <hamzaelsaawy@microsoft.com>
diff --git a/diff/stream_unix.go b/diff/stream_unix.go
@@ -89,6 +89,7 @@ func NewBinaryProcessor(ctx context.Context, imt, rmt string, stream StreamProce
 		r:      r,
 		mt:     rmt,
 		stderr: stderr,
+		done:   make(chan struct{}),
 	}
 	go p.wait()
 
@@ -111,6 +112,11 @@ type binaryProcessor struct {
 
 	mu  sync.Mutex
 	err error
+
+	// There is a race condition between waiting on c.cmd.Wait() and setting c.err within
+	// c.wait(), and reading that value from c.Err().
+	// Use done to wait for the returned error to be captured and set.
+	done chan struct{}
 }
 
 func (c *binaryProcessor) Err() error {
@@ -127,6 +133,16 @@ func (c *binaryProcessor) wait() {
 			c.mu.Unlock()
 		}
 	}
+	close(c.done)
+}
+
+func (c *binaryProcessor) Wait(ctx context.Context) error {
+	select {
+	case <-c.done:
+		return c.Err()
+	case <-ctx.Done():
+		return ctx.Err()
+	}
 }
 
 func (c *binaryProcessor) File() *os.File {
diff --git a/diff/stream_windows.go b/diff/stream_windows.go
@@ -98,6 +98,7 @@ func NewBinaryProcessor(ctx context.Context, imt, rmt string, stream StreamProce
 		r:      r,
 		mt:     rmt,
 		stderr: stderr,
+		done:   make(chan struct{}),
 	}
 	go p.wait()
 
@@ -117,6 +118,11 @@ type binaryProcessor struct {
 
 	mu  sync.Mutex
 	err error
+
+	// There is a race condition between waiting on c.cmd.Wait() and setting c.err within
+	// c.wait(), and reading that value from c.Err().
+	// Use done to wait for the returned error to be captured and set.
+	done chan struct{}
 }
 
 func (c *binaryProcessor) Err() error {
@@ -133,6 +139,16 @@ func (c *binaryProcessor) wait() {
 			c.mu.Unlock()
 		}
 	}
+	close(c.done)
+}
+
+func (c *binaryProcessor) Wait(ctx context.Context) error {
+	select {
+	case <-c.done:
+		return c.Err()
+	case <-ctx.Done():
+		return ctx.Err()
+	}
 }
 
 func (c *binaryProcessor) File() *os.File {
