fix: fix progress bar stuck when error occurred during pull/push (#65)

Signed-off-by: chlins <[email protected]>

Author: chlins

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()

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
 		}
 	}

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
 		}
 	}