Handled safe channel closure

Author: MetsysEht

pkg/distribution/oci/remote/remote.go

@@ -743,8 +743,8 @@ func Write(ref reference.Reference, img oci.Image, w io.Writer, opts ...Option)
 		safeWriter = &syncWriter{w: w}
 	}
 
-	var completed int64
 	for _, layer := range layers {
+		var completed int64
 		digest, err := layer.Digest()
 		if err != nil {
 			return fmt.Errorf("getting layer digest: %w", err)
@@ -790,6 +790,8 @@ func Write(ref reference.Reference, img oci.Image, w io.Writer, opts ...Option)
 						Complete: completed,
 						Total:    size,
 					}
+					closeProgress(progressChan)
+					closeReporter(pr)
 				}
 				continue
 			}
@@ -808,6 +810,7 @@ func Write(ref reference.Reference, img oci.Image, w io.Writer, opts ...Option)
 			cw.Close()
 			rc.Close()
 			closeProgress(progressChan)
+			closeReporter(pr)
 			return fmt.Errorf("writing layer: %w", err)
 		}
 
@@ -816,6 +819,7 @@ func Write(ref reference.Reference, img oci.Image, w io.Writer, opts ...Option)
 			rc.Close()
 			if !errdefs.IsAlreadyExists(err) && !strings.Contains(err.Error(), "already exists") {
 				closeProgress(progressChan)
+				closeReporter(pr)
 				return fmt.Errorf("committing layer: %w", err)
 			}
 			// If it already exists, we still want to update progress
@@ -837,6 +841,7 @@ func Write(ref reference.Reference, img oci.Image, w io.Writer, opts ...Option)
 			}
 		}
 		closeProgress(progressChan)
+		closeReporter(pr)
 		cw.Close()
 		rc.Close()
 	}
@@ -933,6 +938,15 @@ func closeProgress(ch chan<- oci.Update) {
 	}
 }
 
+// closeReporter safely closes the progress reporter if not nil
+func closeReporter(pr *progress.Reporter) {
+	if pr != nil {
+		if waitErr := pr.Wait(); waitErr != nil {
+			fmt.Printf("reporter finished with non-fatal error: %v\n", waitErr)
+		}
+	}
+}
+
 // Ensure remoteImage is cleaned up properly
 func (i *remoteImage) Close() error {
 	// The local content store doesn't expose its root path, so cleanup

pkg/distribution/registry/client.go

@@ -263,15 +263,12 @@ func (t *Target) Write(ctx context.Context, model types.ModelArtifact, progressW
 		}
 		imageSize += size
 	}
-	//pr := progress.NewProgressReporter(progressWriter, progress.PushMsg, imageSize, nil)
-	//defer pr.Wait()
 
 	// Set up authentication options
 	authOpts := []remote.Option{
 		remote.WithContext(ctx),
 		remote.WithTransport(t.transport),
 		remote.WithUserAgent(t.userAgent),
-		//remote.WithProgress(pr.Updates()),
 		remote.WithPlainHTTP(t.plainHTTP),
 	}