tarball: pass imageToTags (#1563)

thepwagner · web-flow · commit 3624968304fd · 2023-02-09T11:53:35.000-05:00
* tarball: godocs typo

* tarball: memoize imageToTags
diff --git a/pkg/v1/tarball/write.go b/pkg/v1/tarball/write.go
@@ -69,7 +69,7 @@ func Write(ref name.Reference, img v1.Image, w io.Writer, opts ...WriteOption) e
 	return MultiRefWrite(map[name.Reference]v1.Image{ref: img}, w, opts...)
 }
 
-// MultiWrite writes the contents of each image to the provided reader, in the compressed format.
+// MultiWrite writes the contents of each image to the provided writer, in the compressed format.
 // The contents are written in the following format:
 // One manifest.json file at the top level containing information about several images.
 // One file for each layer, named after the layer's SHA.
@@ -82,7 +82,7 @@ func MultiWrite(tagToImage map[name.Tag]v1.Image, w io.Writer, opts ...WriteOpti
 	return MultiRefWrite(refToImage, w, opts...)
 }
 
-// MultiRefWrite writes the contents of each image to the provided reader, in the compressed format.
+// MultiRefWrite writes the contents of each image to the provided writer, in the compressed format.
 // The contents are written in the following format:
 // One manifest.json file at the top level containing information about several images.
 // One file for each layer, named after the layer's SHA.
@@ -98,12 +98,13 @@ func MultiRefWrite(refToImage map[name.Reference]v1.Image, w io.Writer, opts ...
 		}
 	}
 
-	size, mBytes, err := getSizeAndManifest(refToImage)
+	imageToTags := dedupRefToImage(refToImage)
+	size, mBytes, err := getSizeAndManifest(imageToTags)
 	if err != nil {
 		return sendUpdateReturn(o, err)
 	}
 
-	return writeImagesToTar(refToImage, mBytes, size, w, o)
+	return writeImagesToTar(imageToTags, mBytes, size, w, o)
 }
 
 // sendUpdateReturn return the passed in error message, also sending on update channel, if it exists
@@ -125,11 +126,10 @@ func sendProgressWriterReturn(pw *progressWriter, err error) error {
 }
 
 // writeImagesToTar writes the images to the tarball
-func writeImagesToTar(refToImage map[name.Reference]v1.Image, m []byte, size int64, w io.Writer, o *writeOptions) (err error) {
+func writeImagesToTar(imageToTags map[v1.Image][]string, m []byte, size int64, w io.Writer, o *writeOptions) (err error) {
 	if w == nil {
 		return sendUpdateReturn(o, errors.New("must pass valid writer"))
 	}
-	imageToTags := dedupRefToImage(refToImage)
 
 	tw := w
 	var pw *progressWriter
@@ -219,9 +219,7 @@ func writeImagesToTar(refToImage map[name.Reference]v1.Image, m []byte, size int
 }
 
 // calculateManifest calculates the manifest and optionally the size of the tar file
-func calculateManifest(refToImage map[name.Reference]v1.Image) (m Manifest, err error) {
-	imageToTags := dedupRefToImage(refToImage)
-
+func calculateManifest(imageToTags map[v1.Image][]string) (m Manifest, err error) {
 	if len(imageToTags) == 0 {
 		return nil, errors.New("set of images is empty")
 	}
@@ -290,12 +288,13 @@ func calculateManifest(refToImage map[name.Reference]v1.Image) (m Manifest, err
 
 // CalculateSize calculates the expected complete size of the output tar file
 func CalculateSize(refToImage map[name.Reference]v1.Image) (size int64, err error) {
-	size, _, err = getSizeAndManifest(refToImage)
+	imageToTags := dedupRefToImage(refToImage)
+	size, _, err = getSizeAndManifest(imageToTags)
 	return size, err
 }
 
-func getSizeAndManifest(refToImage map[name.Reference]v1.Image) (int64, []byte, error) {
-	m, err := calculateManifest(refToImage)
+func getSizeAndManifest(imageToTags map[v1.Image][]string) (int64, []byte, error) {
+	m, err := calculateManifest(imageToTags)
 	if err != nil {
 		return 0, nil, fmt.Errorf("unable to calculate manifest: %w", err)
 	}
@@ -304,17 +303,15 @@ func getSizeAndManifest(refToImage map[name.Reference]v1.Image) (int64, []byte,
 		return 0, nil, fmt.Errorf("could not marshall manifest to bytes: %w", err)
 	}
 
-	size, err := calculateTarballSize(refToImage, mBytes)
+	size, err := calculateTarballSize(imageToTags, mBytes)
 	if err != nil {
 		return 0, nil, fmt.Errorf("error calculating tarball size: %w", err)
 	}
 	return size, mBytes, nil
 }
 
 // calculateTarballSize calculates the size of the tar file
-func calculateTarballSize(refToImage map[name.Reference]v1.Image, mBytes []byte) (size int64, err error) {
-	imageToTags := dedupRefToImage(refToImage)
-
+func calculateTarballSize(imageToTags map[v1.Image][]string, mBytes []byte) (size int64, err error) {
 	seenLayerDigests := make(map[string]struct{})
 	for img, name := range imageToTags {
 		manifest, err := img.Manifest()
@@ -386,7 +383,8 @@ func writeTarEntry(tf *tar.Writer, path string, r io.Reader, size int64) error {
 // ComputeManifest get the manifest.json that will be written to the tarball
 // for multiple references
 func ComputeManifest(refToImage map[name.Reference]v1.Image) (Manifest, error) {
-	return calculateManifest(refToImage)
+	imageToTags := dedupRefToImage(refToImage)
+	return calculateManifest(imageToTags)
 }
 
 // WriteOption a function option to pass to Write()
