fix: fix load image failed

Signed-off-by: ningmingxiao <ning.mingxiao@zte.com.cn>

Author: ningmingxiao

pkg/imgutil/load/load.go

@@ -84,32 +84,13 @@ func FromArchive(ctx context.Context, client *containerd.Client, options types.I
 	storeOpts = append(storeOpts, transferimage.WithUnpack(platUnpack, options.GOptions.Snapshotter))
 	storeOpts = append(storeOpts, transferimage.WithDigestRef("import", true, true))
 
-	var loadedImages []images.Image
-	pf, done := transferutil.ProgressHandler(ctx, options.Stdout)
-
+	pf, done, loadedImages := transferutil.ProgressHandlerLoadImage(ctx, client, beforeSet, options)
 	err = client.Transfer(ctx,
 		tarchive.NewImageImportStream(options.Stdin, ""),
 		transferimage.NewStore("", storeOpts...),
-		transfer.WithProgress(func(p transfer.Progress) {
-			if p.Event == "saved" {
-				if img, err := imageService.Get(ctx, p.Name); err == nil {
-					if !beforeSet[img.Name] {
-						loadedImages = append(loadedImages, img)
-					}
-				}
-			}
-			pf(p)
-		}),
+		transfer.WithProgress(pf),
 	)
-
 	done()
-
-	if !options.Quiet {
-		for _, img := range loadedImages {
-			fmt.Fprintf(options.Stdout, "Loaded image: %s\n", img.Name)
-		}
-	}
-
 	return loadedImages, err
 }
 

pkg/transferutil/progress.go

@@ -25,8 +25,11 @@ import (
 
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 
+	containerd "github.com/containerd/containerd/v2/client"
+	"github.com/containerd/containerd/v2/core/images"
 	"github.com/containerd/containerd/v2/core/transfer"
 	"github.com/containerd/containerd/v2/pkg/progress"
+	"github.com/containerd/nerdctl/v2/pkg/api/types"
 )
 
 // From https://github.com/containerd/containerd/blob/v2.2.0-rc.0/cmd/ctr/commands/image/pull.go#L240-L473
@@ -156,6 +159,125 @@ func ProgressHandler(ctx context.Context, out io.Writer) (transfer.ProgressFunc,
 	return progressFn, done
 }
 
+func ProgressHandlerLoadImage(ctx context.Context, client *containerd.Client, beforeSet map[string]bool, options types.ImageLoadOptions) (transfer.ProgressFunc, func(), []images.Image) {
+	ctx, cancel := context.WithCancel(ctx)
+	var (
+		fw            = progress.NewWriter(options.Stdout)
+		start         = time.Now()
+		statuses      = map[string]*progressNode{}
+		roots         = []*progressNode{}
+		pc            = make(chan transfer.Progress, 5)
+		status        string
+		closeC        = make(chan struct{})
+		loadedImages  []images.Image
+		imagesDisplay []string
+	)
+	progressFn := func(p transfer.Progress) {
+		select {
+		case pc <- p:
+		case <-ctx.Done():
+		}
+	}
+
+	done := func() {
+		cancel()
+		<-closeC
+		if !options.Quiet {
+			for _, img := range imagesDisplay {
+				fmt.Fprintf(options.Stdout, "Loaded image: %s\n", img)
+			}
+		}
+	}
+
+	go func() {
+		defer close(closeC)
+		for {
+			select {
+			case p := <-pc:
+				if p.Name == "" {
+					status = p.Event
+					continue
+				}
+				if p.Event == "saved" {
+					if img, err := client.ImageService().Get(ctx, p.Name); err == nil {
+						if !beforeSet[img.Name] {
+							loadedImages = append(loadedImages, img)
+						}
+						imagesDisplay = append(imagesDisplay, img.Name)
+					}
+				}
+				if node, ok := statuses[p.Name]; !ok {
+					node = &progressNode{
+						Progress: p,
+						root:     true,
+					}
+					if len(p.Parents) == 0 {
+						roots = append(roots, node)
+					} else {
+						var parents []string
+						for _, parent := range p.Parents {
+							pStatus, ok := statuses[parent]
+							if ok {
+								parents = append(parents, parent)
+								pStatus.children = append(pStatus.children, node)
+								node.root = false
+							}
+						}
+						node.Progress.Parents = parents
+						if node.root {
+							roots = append(roots, node)
+						}
+					}
+					statuses[p.Name] = node
+				} else {
+					if len(node.Progress.Parents) != len(p.Parents) {
+						var parents []string
+						var removeRoot bool
+						for _, parent := range p.Parents {
+							pStatus, ok := statuses[parent]
+							if ok {
+								parents = append(parents, parent)
+								var found bool
+								for _, child := range pStatus.children {
+									if child.Progress.Name == p.Name {
+										found = true
+										break
+									}
+								}
+								if !found {
+									pStatus.children = append(pStatus.children, node)
+								}
+								if node.root {
+									removeRoot = true
+								}
+								node.root = false
+							}
+						}
+						p.Parents = parents
+						// Check if needs to remove from root
+						if removeRoot {
+							for i := range roots {
+								if roots[i] == node {
+									roots = append(roots[:i], roots[i+1:]...)
+									break
+								}
+							}
+						}
+					}
+					node.Progress = p
+				}
+
+				displayHierarchy(fw, status, roots, start)
+				fw.Flush()
+
+			case <-ctx.Done():
+				return
+			}
+		}
+	}()
+	return progressFn, done, loadedImages
+}
+
 func displayHierarchy(w io.Writer, status string, roots []*progressNode, start time.Time) {
 	total := displayNode(w, "", roots)
 	for _, r := range roots {