From 174af51dde9678a803839baee32d28d4bdcf6dcf Mon Sep 17 00:00:00 2001
From: chlins <chlins.zhang@gmail.com>
Date: Mon, 10 Feb 2025 17:47:38 +0800
Subject: [PATCH] fix: fix progress bar stuck when error occurred during
 pull/push

Signed-off-by: chlins <chlins.zhang@gmail.com>
---
 pkg/backend/progress.go | 11 +++++++++++
 pkg/backend/pull.go     |  2 ++
 pkg/backend/push.go     |  2 ++
 3 files changed, 15 insertions(+)

diff --git a/pkg/backend/progress.go b/pkg/backend/progress.go
index 19d8a630..c29f07fd 100644
--- a/pkg/backend/progress.go
+++ b/pkg/backend/progress.go
@@ -105,6 +105,17 @@ func (p *ProgressBar) PrintMessage(prompt string, desc ocispec.Descriptor, messa
 	p.bars[desc.Digest.String()] = bar
 }
 
+// Abort aborts the progress bar by the given desc.
+func (p *ProgressBar) Abort(desc ocispec.Descriptor) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	// if the bar already exists, abort it.
+	bar, ok := p.bars[desc.Digest.String()]
+	if ok {
+		bar.Abort(false)
+	}
+}
+
 // Wait waits for the progress bar to finish.
 func (p *ProgressBar) Wait() {
 	p.mpb.Wait()
diff --git a/pkg/backend/pull.go b/pkg/backend/pull.go
index 84b5b9ef..81384bad 100644
--- a/pkg/backend/pull.go
+++ b/pkg/backend/pull.go
@@ -163,6 +163,7 @@ func pullIfNotExist(ctx context.Context, pb *ProgressBar, prompt string, src *re
 	if desc.MediaType == ocispec.MediaTypeImageManifest {
 		body, err := io.ReadAll(pb.Add(prompt, desc, content))
 		if err != nil {
+			pb.Abort(desc)
 			return fmt.Errorf("failed to read the manifest: %w", err)
 		}
 
@@ -171,6 +172,7 @@ func pullIfNotExist(ctx context.Context, pb *ProgressBar, prompt string, src *re
 		}
 	} else {
 		if _, _, err := dst.PushBlob(ctx, repo, pb.Add(prompt, desc, content)); err != nil {
+			pb.Abort(desc)
 			return err
 		}
 	}
diff --git a/pkg/backend/push.go b/pkg/backend/push.go
index 4f0c350b..72806e8f 100644
--- a/pkg/backend/push.go
+++ b/pkg/backend/push.go
@@ -142,6 +142,7 @@ func pushIfNotExist(ctx context.Context, pb *ProgressBar, prompt string, src sto
 		}
 
 		if err := dst.Manifests().Push(ctx, desc, pb.Add(prompt, desc, bytes.NewReader(manifestRaw))); err != nil {
+			pb.Abort(desc)
 			return err
 		}
 
@@ -159,6 +160,7 @@ func pushIfNotExist(ctx context.Context, pb *ProgressBar, prompt string, src sto
 		defer content.Close()
 
 		if err := dst.Blobs().Push(ctx, desc, pb.Add(prompt, desc, content)); err != nil {
+			pb.Abort(desc)
 			return err
 		}
 	}