From 422d6c0d5070c22f7e96ae6d18c3b68a558ff570 Mon Sep 17 00:00:00 2001
From: CrazyMax <1951866+crazy-max@users.noreply.github.com>
Date: Tue, 8 Sep 2026 17:19:10 +0200
Subject: [PATCH] build: wait for docker load progress cleanup

Signed-off-by: CrazyMax <1951866+crazy-max@users.noreply.github.com>
---
 build/opt.go                   |  1 +
 util/dockerutil/client.go      |  3 ++
 util/dockerutil/client_test.go | 73 ++++++++++++++++++++++++++++++++++
 3 files changed, 77 insertions(+)
 create mode 100644 util/dockerutil/client_test.go

diff --git a/build/opt.go b/build/opt.go
index fb8b5fe393d2..39262d30d313 100644
--- a/build/opt.go
+++ b/build/opt.go
@@ -501,6 +501,7 @@ func toSolveOpt(ctx context.Context, np *noderesolver.ResolvedNode, multiDriver
 					}
 					defers = append(defers, func(error) {
 						cancel()
+						_ = w.Close()
 					})
 					so.Exports[i].Output = func(_ map[string]string) (io.WriteCloser, error) {
 						return w, nil
diff --git a/util/dockerutil/client.go b/util/dockerutil/client.go
index 68ff9a29c281..421bd911b77d 100644
--- a/util/dockerutil/client.go
+++ b/util/dockerutil/client.go
@@ -110,6 +110,9 @@ func (w *waitingWriter) Write(dt []byte) (int, error) {
 
 func (w *waitingWriter) Close() error {
 	err := w.PipeWriter.Close()
+	w.once.Do(func() {
+		close(w.done)
+	})
 	<-w.done
 	if err == nil {
 		w.mu.Lock()
diff --git a/util/dockerutil/client_test.go b/util/dockerutil/client_test.go
new file mode 100644
index 000000000000..994b9a2339cf
--- /dev/null
+++ b/util/dockerutil/client_test.go
@@ -0,0 +1,73 @@
+package dockerutil
+
+import (
+	"io"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestWaitingWriterCloseWithoutWrite(t *testing.T) {
+	pr, pw := io.Pipe()
+	defer pr.Close()
+
+	started := make(chan struct{})
+	w := &waitingWriter{
+		PipeWriter: pw,
+		f: func() {
+			close(started)
+		},
+		done: make(chan struct{}),
+	}
+
+	require.NoError(t, w.Close())
+	select {
+	case <-started:
+		t.Fatal("loader should not start")
+	default:
+	}
+}
+
+func TestWaitingWriterCloseWaitsForLoader(t *testing.T) {
+	pr, pw := io.Pipe()
+	defer pr.Close()
+
+	started := make(chan struct{})
+	finish := make(chan struct{})
+	done := make(chan struct{})
+	w := &waitingWriter{
+		PipeWriter: pw,
+		f: func() {
+			close(started)
+			<-finish
+			close(done)
+		},
+		done: done,
+	}
+
+	copyDone := make(chan struct{})
+	go func() {
+		_, _ = io.Copy(io.Discard, pr)
+		close(copyDone)
+	}()
+
+	_, err := w.Write([]byte("layer"))
+	require.NoError(t, err)
+	<-started
+
+	closed := make(chan error, 1)
+	go func() {
+		closed <- w.Close()
+	}()
+
+	select {
+	case err := <-closed:
+		require.FailNow(t, "Close returned before loader completed", "err: %v", err)
+	case <-time.After(50 * time.Millisecond):
+	}
+
+	close(finish)
+	require.NoError(t, <-closed)
+	<-copyDone
+}