From 363db3da0b7c74f7f331a4eda2a5a107a5b90b38 Mon Sep 17 00:00:00 2001
From: CrazyMax <github@crazymax.dev>
Date: Tue, 8 Aug 2023 08:38:52 +0200
Subject: [PATCH] build: sublogger to show docker load progress output

Signed-off-by: CrazyMax <crazy-max@users.noreply.github.com>
---
 build/build.go              | 10 ++---
 util/dockerutil/client.go   |  6 ++-
 util/dockerutil/progress.go | 90 +++++++++++++++++++++++++------------
 3 files changed, 68 insertions(+), 38 deletions(-)

diff --git a/build/build.go b/build/build.go
index b2705f2d23c..8ac57bebd30 100644
--- a/build/build.go
+++ b/build/build.go
@@ -552,7 +552,9 @@ func toSolveOpt(ctx context.Context, node builder.Node, multiDriver bool, opt Op
 						return nil, nil, err
 					}
 					defers = append(defers, cancel)
-					opt.Exports[i].Output = wrapWriteCloser(w)
+					opt.Exports[i].Output = func(_ map[string]string) (io.WriteCloser, error) {
+						return w, nil
+					}
 				}
 			} else if !nodeDriver.Features(ctx)[driver.DockerExporter] {
 				return nil, nil, notSupported(nodeDriver, driver.DockerExporter)
@@ -1607,12 +1609,6 @@ func handleLowercaseDockerfile(dir, p string) string {
 	return p
 }
 
-func wrapWriteCloser(wc io.WriteCloser) func(map[string]string) (io.WriteCloser, error) {
-	return func(map[string]string) (io.WriteCloser, error) {
-		return wc, nil
-	}
-}
-
 var nodeIdentifierMu sync.Mutex
 
 func tryNodeIdentifier(configDir string) (out string) {
diff --git a/util/dockerutil/client.go b/util/dockerutil/client.go
index 764c42eedce..831545b0a6f 100644
--- a/util/dockerutil/client.go
+++ b/util/dockerutil/client.go
@@ -60,8 +60,10 @@ func (c *Client) LoadImage(ctx context.Context, name string, status progress.Wri
 				return
 			}
 
-			prog := progress.WithPrefix(status, "", false)
-			if err := fromReader(prog, "importing to docker", resp.Body); err != nil {
+			status = progress.ResetTime(status)
+			if err := progress.Wrap("importing to docker", status.Write, func(l progress.SubLogger) error {
+				return fromReader(l, resp.Body)
+			}); err != nil {
 				handleErr(err)
 			}
 		},
diff --git a/util/dockerutil/progress.go b/util/dockerutil/progress.go
index d54be24e83f..cb64e452e8a 100644
--- a/util/dockerutil/progress.go
+++ b/util/dockerutil/progress.go
@@ -1,47 +1,79 @@
 package dockerutil
 
 import (
-	"errors"
+	"encoding/json"
 	"io"
 	"time"
 
 	"github.com/docker/buildx/util/progress"
-	"github.com/docker/cli/cli/streams"
 	"github.com/docker/docker/pkg/jsonmessage"
 	"github.com/moby/buildkit/client"
-	"github.com/moby/buildkit/identity"
-	"github.com/opencontainers/go-digest"
 )
 
-func fromReader(w progress.Writer, name string, rc io.ReadCloser) error {
-	dgst := digest.FromBytes([]byte(identity.NewID()))
-	tm := time.Now()
+const minProgressDelta = 0.50 // %
 
-	vtx := client.Vertex{
-		Digest:  dgst,
-		Name:    name,
-		Started: &tm,
-	}
+func fromReader(l progress.SubLogger, rc io.ReadCloser) error {
+	started := map[string]client.VertexStatus{}
 
-	w.Write(&client.SolveStatus{
-		Vertexes: []*client.Vertex{&vtx},
-	})
+	defer func() {
+		for _, st := range started {
+			st := st
+			if st.Completed == nil {
+				now := time.Now()
+				st.Completed = &now
+				l.SetStatus(&st)
+			}
+		}
+	}()
 
-	err := jsonmessage.DisplayJSONMessagesToStream(rc, streams.NewOut(io.Discard), nil)
-	if err != nil {
-		if jerr, ok := err.(*jsonmessage.JSONError); ok {
-			err = errors.New(jerr.Message)
+	dec := json.NewDecoder(rc)
+	var parsedErr error
+	var jm jsonmessage.JSONMessage
+	for {
+		if err := dec.Decode(&jm); err != nil {
+			if parsedErr != nil {
+				return parsedErr
+			}
+			if err == io.EOF {
+				break
+			}
+			return err
+		}
+		if jm.Error != nil {
+			parsedErr = jm.Error
+		}
+		if jm.ID == "" {
+			continue
+		}
+		id := "loading layer " + jm.ID
+		st, ok := started[id]
+		if !ok {
+			if jm.Progress != nil {
+				now := time.Now()
+				st = client.VertexStatus{
+					ID:      id,
+					Started: &now,
+				}
+			} else {
+				continue
+			}
 		}
+		st.Timestamp = time.Now()
+		var progressDelta float64
+		if jm.Progress != nil && jm.Progress.Total > 0 {
+			progressDelta = float64(jm.Progress.Current-st.Current) / float64(jm.Progress.Total)
+		}
+		if jm.Progress != nil && jm.Status == "Loading layer" && (st.Current == 0 || progressDelta >= minProgressDelta) {
+			st.Current = jm.Progress.Current
+			st.Total = jm.Progress.Total
+		}
+		if jm.Error != nil {
+			now := time.Now()
+			st.Completed = &now
+		}
+		started[id] = st
+		l.SetStatus(&st)
 	}
 
-	tm2 := time.Now()
-	vtx2 := vtx
-	vtx2.Completed = &tm2
-	if err != nil {
-		vtx2.Error = err.Error()
-	}
-	w.Write(&client.SolveStatus{
-		Vertexes: []*client.Vertex{&vtx2},
-	})
-	return err
+	return nil
 }