Decouple ztoc creation and compression algorithm

Signed-off-by: Jin Dong <jindon@amazon.com>

cmd/go.sum

@@ -2133,7 +2133,6 @@ k8s.io/utils v0.0.0-20210819203725-bdf08cb9a70a/go.mod h1:jPW/WVKK9YHAvNhRxK0md/
 k8s.io/utils v0.0.0-20210930125809-cb0fa318a74b/go.mod h1:jPW/WVKK9YHAvNhRxK0md/EJ228hCsBRufyofKtW8HA=
 k8s.io/utils v0.0.0-20220728103510-ee6ede2d64ed h1:jAne/RjBTyawwAy0utX5eqigAwz/lQhTmy+Hr/Cpue4=
 k8s.io/utils v0.0.0-20220728103510-ee6ede2d64ed/go.mod h1:jPW/WVKK9YHAvNhRxK0md/EJ228hCsBRufyofKtW8HA=
-oras.land/oras-go/v2 v2.0.0-rc.6/go.mod h1:iVExH1NxrccIxjsiq17L91WCZ4KIw6jVQyCLsZsu1gc=
 oras.land/oras-go/v2 v2.0.0 h1:+LRAz92WF7AvYQsQjPEAIw3Xb2zPPhuydjpi4pIHmc0=
 oras.land/oras-go/v2 v2.0.0/go.mod h1:iVExH1NxrccIxjsiq17L91WCZ4KIw6jVQyCLsZsu1gc=
 rsc.io/binaryregexp v0.2.0/go.mod h1:qTv7/COck+e2FymRvadv62gMdZztPaShugOCi3I+8D8=

go.mod

@@ -14,6 +14,7 @@ require (
 	github.com/hanwen/go-fuse/v2 v2.1.1-0.20210825171523-3ab5d95a30ae
 	github.com/hashicorp/go-multierror v1.1.1
 	github.com/hashicorp/go-retryablehttp v0.7.2
+	github.com/klauspost/compress v1.15.14
 	github.com/moby/sys/mountinfo v0.6.2
 	github.com/montanaflynn/stats v0.7.0
 	github.com/opencontainers/go-digest v1.0.0
@@ -64,7 +65,6 @@ require (
 	github.com/imdario/mergo v0.3.12 // indirect
 	github.com/josharian/intern v1.0.0 // indirect
 	github.com/json-iterator/go v1.1.12 // indirect
-	github.com/klauspost/compress v1.15.14 // indirect
 	github.com/mailru/easyjson v0.7.6 // indirect
 	github.com/matttproud/golang_protobuf_extensions v1.0.4 // indirect
 	github.com/moby/locker v1.0.1 // indirect

soci/soci_index.go

@@ -40,23 +40,23 @@ import (
 )
 
 const (
-	// artifactType of index SOCI index
+	// SociIndexArtifactType is the artifactType of index SOCI index
 	SociIndexArtifactType = "application/vnd.amazon.soci.index.v1+json"
-	// mediaType of ztoc
+	// SociLayerMediaType is the mediaType of ztoc
 	SociLayerMediaType = "application/octet-stream"
-	// index annotation for image layer media type
+	// IndexAnnotationImageLayerMediaType is the index annotation for image layer media type
 	IndexAnnotationImageLayerMediaType = "com.amazon.soci.image-layer-mediaType"
-	// index annotation for image layer digest
+	// IndexAnnotationImageLayerDigest is the index annotation for image layer digest
 	IndexAnnotationImageLayerDigest = "com.amazon.soci.image-layer-digest"
-	// index annotation for build tool identifier
+	// IndexAnnotationBuildToolIdentifier is the index annotation for build tool identifier
 	IndexAnnotationBuildToolIdentifier = "com.amazon.soci.build-tool-identifier"
-	// media type for OCI Artifact manifest
+	// OCIArtifactManifestMediaType is the media type for OCI Artifact manifest
 	OCIArtifactManifestMediaType = "application/vnd.oci.artifact.manifest.v1+json"
-	// default span size (4MiB)
+	// defaultSpanSize (4MiB)
 	defaultSpanSize = int64(1 << 22)
-	// min layer size (10MiB)
+	// minLayerSize to build a ztoc for a layer (10MiB)
 	minLayerSize = 10 << 20
-	// default build tool identier
+	// defaultBuildToolIdentifier is the default build tool identier
 	defaultBuildToolIdentifier = "AWS SOCI CLI v0.1"
 )
 
@@ -65,7 +65,7 @@ var (
 	errUnsupportedLayerFormat = errors.New("unsupported layer format")
 )
 
-// Index represents an ORAS/OCI Artifact Manifest
+// Index represents an ORAS/OCI Artifact Manifest.
 type Index struct {
 	// The media type of the manifest
 	MediaType string `json:"mediaType"`
@@ -82,18 +82,21 @@ type Index struct {
 	Annotations map[string]string `json:"annotations,omitempty"`
 }
 
+// IndexWithMetadata has a soci `Index` and its metadata.
 type IndexWithMetadata struct {
 	Index       *Index
 	Platform    *ocispec.Platform
 	ImageDigest digest.Digest
 	CreatedAt   time.Time
 }
 
+// IndexDescriptorInfo has a soci index descriptor and its creation time info.
 type IndexDescriptorInfo struct {
 	ocispec.Descriptor
 	CreatedAt time.Time
 }
 
+// GetIndexDescriptorCollection returns all `IndexDescriptorInfo` of the given image and platforms.
 func GetIndexDescriptorCollection(ctx context.Context, cs content.Store, img images.Image, ps []ocispec.Platform) ([]IndexDescriptorInfo, error) {
 	descriptors := []IndexDescriptorInfo{}
 	var entries []ArtifactEntry
@@ -136,42 +139,50 @@ type buildConfig struct {
 	platform            ocispec.Platform
 }
 
+// BuildOption specifies a config change to build soci indices.
 type BuildOption func(c *buildConfig) error
 
+// WithSpanSize specifies span size.
 func WithSpanSize(spanSize int64) BuildOption {
 	return func(c *buildConfig) error {
 		c.spanSize = spanSize
 		return nil
 	}
 }
 
+// WithMinLayerSize specifies min layer size to build a ztoc for a layer.
 func WithMinLayerSize(minLayerSize int64) BuildOption {
 	return func(c *buildConfig) error {
 		c.minLayerSize = minLayerSize
 		return nil
 	}
 }
 
+// WithBuildToolIdentifier specifies the build tool label value.
 func WithBuildToolIdentifier(tool string) BuildOption {
 	return func(c *buildConfig) error {
 		c.buildToolIdentifier = tool
 		return nil
 	}
 }
 
+// WithPlatform specifies platform used to build soci indices.
 func WithPlatform(platform ocispec.Platform) BuildOption {
 	return func(c *buildConfig) error {
 		c.platform = platform
 		return nil
 	}
 }
 
+// IndexBuilder creates soci indices.
 type IndexBuilder struct {
 	contentStore content.Store
 	blobStore    orascontent.Storage
 	config       *buildConfig
+	ztocBuilder  *ztoc.Builder
 }
 
+// NewIndexBuilder returns an `IndexBuilder` that is used to create soci indices.
 func NewIndexBuilder(contentStore content.Store, blobStore orascontent.Storage, opts ...BuildOption) (*IndexBuilder, error) {
 	defaultPlatform := platforms.DefaultSpec()
 	config := &buildConfig{
@@ -191,9 +202,11 @@ func NewIndexBuilder(contentStore content.Store, blobStore orascontent.Storage,
 		contentStore: contentStore,
 		blobStore:    blobStore,
 		config:       config,
+		ztocBuilder:  ztoc.NewBuilder(),
 	}, nil
 }
 
+// Build builds a soci index for `img` and return the index with metadata.
 func (b *IndexBuilder) Build(ctx context.Context, img images.Image) (*IndexWithMetadata, error) {
 	// we get manifest descriptor before calling images.Manifest, since after calling
 	// images.Manifest, images.Children will error out when reading the manifest blob (this happens on containerd side)
@@ -214,7 +227,7 @@ func (b *IndexBuilder) Build(ctx context.Context, img images.Image) (*IndexWithM
 		eg.Go(func() error {
 			desc, err := b.buildSociLayer(ctx, l)
 			if err != nil {
-				return fmt.Errorf("could not build zTOC for %s: %w", l.Digest.String(), err)
+				return fmt.Errorf("could not build zTOC for layer %s: %w", l.Digest.String(), err)
 			}
 			sociLayersDesc[i] = desc
 			return nil
@@ -250,38 +263,8 @@ func (b *IndexBuilder) Build(ctx context.Context, img images.Image) (*IndexWithM
 	}, nil
 }
 
-// Returns a new index.
-func NewIndex(blobs []ocispec.Descriptor, subject *ocispec.Descriptor, annotations map[string]string) *Index {
-	return &Index{
-		Blobs:        blobs,
-		ArtifactType: SociIndexArtifactType,
-		Annotations:  annotations,
-		Subject:      subject,
-		MediaType:    OCIArtifactManifestMediaType,
-	}
-}
-
-// Returns a new index from a Reader.
-func NewIndexFromReader(reader io.Reader) (*Index, error) {
-	index := new(Index)
-	if err := json.NewDecoder(reader).Decode(index); err != nil {
-		return nil, fmt.Errorf("unable to decode reader into index: %v", err)
-	}
-	return index, nil
-}
-
-func skipBuildingZtoc(desc ocispec.Descriptor, cfg *buildConfig) bool {
-	if cfg == nil {
-		return false
-	}
-	// avoid the file access if the layer size is below threshold
-	if desc.Size < cfg.minLayerSize {
-		return true
-	}
-	return false
-}
-
-// buildSociLayer builds the ztoc for an image layer and returns a Descriptor for the new ztoc.
+// buildSociLayer builds a ztoc for an image layer (`desc`) and returns ztoc descriptor.
+// It may skip building ztoc (e.g., if layer size < `minLayerSize`) and return nil.
 func (b *IndexBuilder) buildSociLayer(ctx context.Context, desc ocispec.Descriptor) (*ocispec.Descriptor, error) {
 	if !images.IsLayerType(desc.MediaType) {
 		return nil, errNotLayerType
@@ -291,10 +274,13 @@ func (b *IndexBuilder) buildSociLayer(ctx context.Context, desc ocispec.Descript
 		fmt.Printf("layer %s -> ztoc skipped\n", desc.Digest)
 		return nil, nil
 	}
+
 	compression, err := images.DiffCompression(ctx, desc.MediaType)
 	if err != nil {
 		return nil, fmt.Errorf("could not determine layer compression: %w", err)
 	}
+	// TODO(djdongjin): when we support new compression algorithms (e.g., zstd), we need to
+	// change the check here.
 	if compression != "gzip" {
 		return nil, fmt.Errorf("layer %s (%s) must be compressed by gzip, but got %q: %w",
 			desc.Digest, desc.MediaType, compression, errUnsupportedLayerFormat)
@@ -320,7 +306,7 @@ func (b *IndexBuilder) buildSociLayer(ctx context.Context, desc ocispec.Descript
 		return nil, errors.New("the size of the temp file doesn't match that of the layer")
 	}
 
-	toc, err := ztoc.BuildZtoc(tmpFile.Name(), b.config.spanSize, b.config.buildToolIdentifier)
+	toc, err := b.ztocBuilder.BuildZtoc(tmpFile.Name(), b.config.spanSize, b.config.buildToolIdentifier, ztoc.WithCompression(compression))
 	if err != nil {
 		return nil, err
 	}
@@ -355,13 +341,46 @@ func (b *IndexBuilder) buildSociLayer(ctx context.Context, desc ocispec.Descript
 
 	ztocDesc.MediaType = SociLayerMediaType
 	ztocDesc.Annotations = map[string]string{
+		// TODO(djdongjin): when we support new compression algorithms (e.g., zstd), we need to
+		// update the annotation value here (e.g., ocispec.MediaTypeImageLayerZstd).
 		IndexAnnotationImageLayerMediaType: ocispec.MediaTypeImageLayerGzip,
 		IndexAnnotationImageLayerDigest:    desc.Digest.String(),
 	}
 	return &ztocDesc, err
 }
 
-// getImageManifestDescriptor gets the descriptor of image manifest
+// NewIndex returns a new index.
+func NewIndex(blobs []ocispec.Descriptor, subject *ocispec.Descriptor, annotations map[string]string) *Index {
+	return &Index{
+		Blobs:        blobs,
+		ArtifactType: SociIndexArtifactType,
+		Annotations:  annotations,
+		Subject:      subject,
+		MediaType:    OCIArtifactManifestMediaType,
+	}
+}
+
+// NewIndexFromReader returns a new index from a Reader.
+func NewIndexFromReader(reader io.Reader) (*Index, error) {
+	index := new(Index)
+	if err := json.NewDecoder(reader).Decode(index); err != nil {
+		return nil, fmt.Errorf("unable to decode reader into index: %v", err)
+	}
+	return index, nil
+}
+
+func skipBuildingZtoc(desc ocispec.Descriptor, cfg *buildConfig) bool {
+	if cfg == nil {
+		return false
+	}
+	// avoid the file access if the layer size is below threshold
+	if desc.Size < cfg.minLayerSize {
+		return true
+	}
+	return false
+}
+
+// GetImageManifestDescriptor gets the descriptor of image manifest.
 func GetImageManifestDescriptor(ctx context.Context, cs content.Store, imageTarget ocispec.Descriptor, platform platforms.MatchComparer) (*ocispec.Descriptor, error) {
 	if images.IsIndexType(imageTarget.MediaType) {
 		manifests, err := images.Children(ctx, cs, imageTarget)
@@ -384,7 +403,7 @@ func GetImageManifestDescriptor(ctx context.Context, cs content.Store, imageTarg
 	return nil, nil
 }
 
-// WriteSociIndex writes the SociIndex manifest
+// WriteSociIndex writes the SociIndex manifest to oras `store`.
 func WriteSociIndex(ctx context.Context, indexWithMetadata *IndexWithMetadata, store orascontent.Storage) error {
 	manifest, err := json.Marshal(indexWithMetadata.Index)
 	if err != nil {

util/testutil/tar.go

@@ -36,12 +36,15 @@ package testutil
 
 import (
 	"archive/tar"
+	"bytes"
 	"compress/gzip"
 	"fmt"
 	"io"
 	"os"
 	"strings"
 	"time"
+
+	"github.com/klauspost/compress/zstd"
 )
 
 // TarEntry is an entry of tar.
@@ -66,7 +69,7 @@ func WithPrefix(prefix string) BuildTarOption {
 	}
 }
 
-// BuildTar builds a tar blob
+// BuildTar builds a tar blob.
 func BuildTar(ents []TarEntry, opts ...BuildTarOption) io.Reader {
 	var bo BuildTarOptions
 	for _, o := range opts {
@@ -90,6 +93,7 @@ func BuildTar(ents []TarEntry, opts ...BuildTarOption) io.Reader {
 	return pr
 }
 
+// BuildTarGz builds a tar blob with gzip compression.
 func BuildTarGz(ents []TarEntry, compressionLevel int, opts ...BuildTarOption) io.Reader {
 	var bo BuildTarOptions
 	for _, o := range opts {
@@ -122,6 +126,58 @@ func BuildTarGz(ents []TarEntry, compressionLevel int, opts ...BuildTarOption) i
 	return pr
 }
 
+// BuildTarZstd builds a tar blob with zstd compression.
+func BuildTarZstd(ents []TarEntry, compressionLevel int, opts ...BuildTarOption) io.Reader {
+	var bo BuildTarOptions
+	for _, o := range opts {
+		o(&bo)
+	}
+	pr, pw := io.Pipe()
+	go func() {
+		zw, err := zstd.NewWriter(pw, zstd.WithEncoderLevel(zstd.EncoderLevelFromZstd(compressionLevel)))
+		if err != nil {
+			pw.CloseWithError(err)
+			return
+		}
+		tw := tar.NewWriter(zw)
+		for _, ent := range ents {
+			if err := ent.AppendTar(tw, bo); err != nil {
+				pw.CloseWithError(err)
+				return
+			}
+		}
+		if err := tw.Close(); err != nil {
+			pw.CloseWithError(err)
+			return
+		}
+		if err := zw.Close(); err != nil {
+			pw.CloseWithError(err)
+			return
+		}
+		pw.Close()
+	}()
+	return pr
+}
+
+// WriteTarToTempFile writes the tar blob to a temp file and returns the file
+// and tar data.
+//
+// The caller need to clean up the temp file if it's not used.
+func WriteTarToTempFile(tarReader io.Reader) (*os.File, []byte, error) {
+	tarFile, err := os.CreateTemp("", "tmp.*")
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to create temp file: %v", err)
+	}
+	tarBuf := new(bytes.Buffer)
+	w := io.MultiWriter(tarFile, tarBuf)
+	_, err = io.Copy(w, tarReader)
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to write tar file: %v", err)
+	}
+	tarData := tarBuf.Bytes()
+	return tarFile, tarData, nil
+}
+
 type tarEntryFunc func(*tar.Writer, BuildTarOptions) error
 
 func (f tarEntryFunc) AppendTar(tw *tar.Writer, opts BuildTarOptions) error { return f(tw, opts) }

ztoc/testutils.go

@@ -145,7 +145,7 @@ func tempTarGz(inputDir string, targzName string) (*string, error) {
 	return &outputFileName, nil
 }
 
-// buildZtocReader creates the tar gz file for tar entries.
+// BuildZtocReader creates the tar gz file for tar entries.
 // It returns ztoc and io.SectionReader of the file.
 func BuildZtocReader(ents []testutil.TarEntry, compressionLevel int, spanSize int64, opts ...testutil.BuildTarOption) (*Ztoc, *io.SectionReader, error) {
 	// build tar gz file
@@ -165,7 +165,8 @@ func BuildZtocReader(ents []testutil.TarEntry, compressionLevel int, spanSize in
 	}
 	tarData := tarBuf.Bytes()
 	sr := io.NewSectionReader(bytes.NewReader(tarData), 0, int64(len(tarData)))
-	ztoc, err := BuildZtoc(tarFile.Name(), spanSize, "test")
+
+	ztoc, err := NewBuilder().BuildZtoc(tarFile.Name(), spanSize, "test")
 	if err != nil {
 		return nil, nil, fmt.Errorf("failed to build sample ztoc: %v", err)
 	}