Enforce lease when building SOCI index

Previously, in our CLI create command, we would create zTOCs, push them
to the content store, then label the index to refer to the zTOCs.
This created a problem where a user using the containerd content store
might have their zTOCs deleted by the containerd garbage collector
before it could be labeled. While a lease was used, it was only done
during the latter step, so it did not prevent zTOCs from being deleted
before being labeled.

This commit fixes this by using a lease on the entire process via a new
command, BuildAndWriteIndex, which is called by "soci create".

Signed-off-by: David Son <[email protected]>

Author: sondavidb

cmd/soci/commands/create.go

@@ -122,12 +122,7 @@ var CreateCommand = cli.Command{
 				return err
 			}
 
-			sociIndexWithMetadata, err := builder.Build(ctx, srcImg)
-			if err != nil {
-				return err
-			}
-
-			err = soci.WriteSociIndex(ctx, sociIndexWithMetadata, blobStore, builder.ArtifactsDb)
+			_, err = builder.Build(ctx, srcImg)
 			if err != nil {
 				return err
 			}

integration/create_test.go

@@ -19,6 +19,7 @@ package integration
 import (
 	"bytes"
 	"encoding/json"
+	"fmt"
 	"path/filepath"
 	"testing"
 
@@ -202,3 +203,32 @@ func TestSociCreate(t *testing.T) {
 		})
 	}
 }
+
+// TestSociCreateGarbageCollection ensures that SOCI index artifacts are not garbage collected after creation.
+func TestSociCreateGarbageCollection(t *testing.T) {
+	image := rabbitmqImage
+	smallImage := alpineImage
+	sh, done := newSnapshotterBaseShell(t)
+	defer done()
+
+	extraContainerdConfig := `
+[plugins."io.containerd.gc.v1.scheduler"]
+	deletion_threshold = 1`
+
+	rebootContainerd(t, sh, getContainerdConfigToml(t, false, extraContainerdConfig), getSnapshotterConfigToml(t, false, tcpMetricsConfig, GetContentStoreConfigToml(store.WithType(config.ContainerdContentStoreType))))
+
+	imgInfo := dockerhub(image)
+	sh.X("nerdctl", "pull", "-q", imgInfo.ref)
+	indexDigest := buildIndex(sh, imgInfo, withMinLayerSize(0), withContentStoreType(config.ContainerdContentStoreType))
+
+	// Pull and remove a small image to automatically trigger GC.
+	smallImageInfo := dockerhub(smallImage)
+	sh.X("nerdctl", "pull", "-q", smallImageInfo.ref)
+	sh.X("nerdctl", "rmi", smallImageInfo.ref)
+
+	sh.X(append(imagePullCmd, "--soci-index-digest", indexDigest, imgInfo.ref)...)
+	curlOutput := string(sh.O("curl", tcpMetricsAddress+metricsPath))
+	if err := checkOverlayFallbackCount(curlOutput, 0); err != nil {
+		t.Fatal(fmt.Errorf("resources unexpectedly garbage collected: %v", err))
+	}
+}

soci/soci_index.go

@@ -39,7 +39,6 @@ import (
 	"github.com/opencontainers/go-digest"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 
-	orascontent "oras.land/oras-go/v2/content"
 	"oras.land/oras-go/v2/errdef"
 )
 
@@ -273,14 +272,14 @@ func WithArtifactsDb(db *ArtifactsDb) BuildOption {
 // IndexBuilder creates soci indices.
 type IndexBuilder struct {
 	contentStore content.Store
-	blobStore    orascontent.Storage
+	blobStore    store.Store
 	ArtifactsDb  *ArtifactsDb
 	config       *buildConfig
 	ztocBuilder  *ztoc.Builder
 }
 
 // NewIndexBuilder returns an `IndexBuilder` that is used to create soci indices.
-func NewIndexBuilder(contentStore content.Store, blobStore orascontent.Storage, artifactsDb *ArtifactsDb, opts ...BuildOption) (*IndexBuilder, error) {
+func NewIndexBuilder(contentStore content.Store, blobStore store.Store, artifactsDb *ArtifactsDb, opts ...BuildOption) (*IndexBuilder, error) {
 	defaultPlatform := platforms.DefaultSpec()
 	config := &buildConfig{
 		spanSize:            defaultSpanSize,
@@ -305,8 +304,35 @@ func NewIndexBuilder(contentStore content.Store, blobStore orascontent.Storage,
 	}, nil
 }
 
-// Build builds a soci index for `img` and return the index with metadata.
+// Build builds a soci index for `img` and pushes it with its corresponding zTOCs to the blob store.
+// Returns the SOCI index and its metadata.
 func (b *IndexBuilder) Build(ctx context.Context, img images.Image) (*IndexWithMetadata, error) {
+	// batch will prevent content from being garbage collected in the middle of the following operations
+	ctx, done, err := b.blobStore.BatchOpen(ctx)
+	if err != nil {
+		return nil, err
+	}
+	defer done(ctx)
+
+	// Create and push zTOCs to blob store
+	index, err := b.build(ctx, img)
+	if err != nil {
+		return nil, err
+	}
+
+	// Label zTOCs and push SOCI index
+	err = b.writeSociIndex(ctx, index)
+	if err != nil {
+		return nil, err
+	}
+
+	return index, nil
+}
+
+// build attempts to create a zTOC in each layer and pushes the zTOC to the blob store.
+// It then creates the SOCI index and returns it with some metadata.
+// This should be done within a Batch and followed by writeSociIndex() to prevent garbage collection.
+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)
 	imgManifestDesc, err := GetImageManifestDescriptor(ctx, b.contentStore, img.Target, platforms.OnlyStrict(b.config.platform))
@@ -357,7 +383,6 @@ func (b *IndexBuilder) Build(ctx context.Context, img images.Image) (*IndexWithM
 		for _, err := range errs {
 			errWrap = fmt.Errorf("%w; %v", errWrap, err)
 		}
-
 		return nil, errWrap
 	}
 
@@ -393,6 +418,7 @@ func (b *IndexBuilder) Build(ctx context.Context, img images.Image) (*IndexWithM
 
 // 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.
+// This should be done within a Batch and followed by Label calls to prevent garbage collection.
 func (b *IndexBuilder) buildSociLayer(ctx context.Context, desc ocispec.Descriptor) (*ocispec.Descriptor, error) {
 	if !images.IsLayerType(desc.MediaType) {
 		return nil, errNotLayerType
@@ -541,15 +567,9 @@ func GetImageManifestDescriptor(ctx context.Context, cs content.Store, imageTarg
 	return nil, nil
 }
 
-// WriteSociIndex writes the SociIndex manifest to oras `store`.
-func WriteSociIndex(ctx context.Context, indexWithMetadata *IndexWithMetadata, contentStore store.Store, artifactsDb *ArtifactsDb) error {
-	// batch will prevent content from being garbage collected in the middle of the following operations
-	ctx, batchDone, err := contentStore.BatchOpen(ctx)
-	if err != nil {
-		return err
-	}
-	defer batchDone(ctx)
-
+// writeSociIndex writes the SociIndex manifest to the blob store.
+// This should be done within a Batch to prevent garbage collection.
+func (b *IndexBuilder) writeSociIndex(ctx context.Context, indexWithMetadata *IndexWithMetadata) error {
 	manifest, err := MarshalIndex(indexWithMetadata.Index)
 	if err != nil {
 		return err
@@ -559,7 +579,7 @@ func WriteSociIndex(ctx context.Context, indexWithMetadata *IndexWithMetadata, c
 	// empty config objct in the store as well. We will need to push this to the
 	// registry later.
 	if indexWithMetadata.Index.MediaType == ocispec.MediaTypeImageManifest {
-		err = contentStore.Push(ctx, defaultConfigDescriptor, bytes.NewReader(defaultConfigContent))
+		err = b.blobStore.Push(ctx, defaultConfigDescriptor, bytes.NewReader(defaultConfigContent))
 		if err != nil && !errors.Is(err, errdef.ErrAlreadyExists) {
 			return fmt.Errorf("error creating OCI 1.0 empty config: %w", err)
 		}
@@ -572,25 +592,25 @@ func WriteSociIndex(ctx context.Context, indexWithMetadata *IndexWithMetadata, c
 		Size:   size,
 	}
 
-	err = contentStore.Push(ctx, desc, bytes.NewReader(manifest))
+	err = b.blobStore.Push(ctx, desc, bytes.NewReader(manifest))
 	if err != nil && !errors.Is(err, errdef.ErrAlreadyExists) {
 		return fmt.Errorf("cannot write SOCI index to local store: %w", err)
 	}
 
 	log.G(ctx).WithField("digest", dgst.String()).Debugf("soci index has been written")
 
-	err = store.LabelGCRoot(ctx, contentStore, desc)
+	err = store.LabelGCRoot(ctx, b.blobStore, desc)
 	if err != nil {
 		return fmt.Errorf("cannot apply garbage collection label to index %s: %w", desc.Digest.String(), err)
 	}
-	err = store.LabelGCRefContent(ctx, contentStore, desc, "config", defaultConfigDescriptor.Digest.String())
+	err = store.LabelGCRefContent(ctx, b.blobStore, desc, "config", defaultConfigDescriptor.Digest.String())
 	if err != nil {
 		return fmt.Errorf("cannot apply garbage collection label to index %s referencing default config: %w", desc.Digest.String(), err)
 	}
 
 	var allErr error
 	for i, blob := range indexWithMetadata.Index.Blobs {
-		err = store.LabelGCRefContent(ctx, contentStore, desc, "ztoc."+strconv.Itoa(i), blob.Digest.String())
+		err = store.LabelGCRefContent(ctx, b.blobStore, desc, "ztoc."+strconv.Itoa(i), blob.Digest.String())
 		if err != nil {
 			errors.Join(allErr, err)
 		}
@@ -617,7 +637,7 @@ func WriteSociIndex(ctx context.Context, indexWithMetadata *IndexWithMetadata, c
 		MediaType:      indexWithMetadata.Index.MediaType,
 		CreatedAt:      indexWithMetadata.CreatedAt,
 	}
-	return artifactsDb.WriteArtifactEntry(entry)
+	return b.ArtifactsDb.WriteArtifactEntry(entry)
 }
 
 func (b *IndexBuilder) maybeAddDisableXattrAnnotation(ztocDesc *ocispec.Descriptor, ztoc *ztoc.Ztoc) {

soci/soci_index_test.go

@@ -28,7 +28,6 @@ import (
 	"github.com/google/go-cmp/cmp"
 	"github.com/opencontainers/go-digest"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
-	"oras.land/oras-go/v2/content/memory"
 )
 
 func TestSkipBuildingZtoc(t *testing.T) {
@@ -136,7 +135,7 @@ func TestBuildSociIndexNotLayer(t *testing.T) {
 	spanSize := int64(65535)
 	ctx := context.Background()
 	cs := newFakeContentStore()
-	blobStore := memory.New()
+	blobStore := NewOrasMemoryStore()
 
 	artifactsDb, err := newTestableDb()
 	if err != nil {
@@ -209,7 +208,7 @@ func TestBuildSociIndexWithLimits(t *testing.T) {
 				Size:      tc.layerSize,
 			}
 			spanSize := int64(65535)
-			blobStore := memory.New()
+			blobStore := NewOrasMemoryStore()
 			artifactsDb, err := newTestableDb()
 			if err != nil {
 				t.Fatalf("can't create a test db")
@@ -305,7 +304,7 @@ func TestDisableXattrs(t *testing.T) {
 			}
 
 			cs := newFakeContentStore()
-			blobStore := memory.New()
+			blobStore := NewOrasMemoryStore()
 			artifactsDb, err := newTestableDb()
 			if err != nil {
 				t.Fatalf("can't create a test db")

soci/store/store.go

@@ -141,7 +141,7 @@ func GetContentStorePath(contentStoreType ContentStoreType) (string, error) {
 
 type CleanupFunc func(context.Context) error
 
-func nopCleanup(context.Context) error { return nil }
+func NopCleanup(context.Context) error { return nil }
 
 func NewContentStore(ctx context.Context, opts ...Option) (context.Context, Store, error) {
 	storeConfig := NewStoreConfig(opts...)
@@ -185,7 +185,7 @@ func (s *SociStore) Delete(_ context.Context, _ digest.Digest) error {
 
 // BatchOpen is a no-op for sociStore; it does not support batching operations.
 func (s *SociStore) BatchOpen(ctx context.Context) (context.Context, CleanupFunc, error) {
-	return ctx, nopCleanup, nil
+	return ctx, NopCleanup, nil
 }
 
 type ContainerdStore struct {
@@ -336,7 +336,7 @@ func (s *ContainerdStore) Delete(ctx context.Context, dgst digest.Digest) error
 func (s *ContainerdStore) BatchOpen(ctx context.Context) (context.Context, CleanupFunc, error) {
 	ctx, leaseDone, err := s.client.WithLease(ctx)
 	if err != nil {
-		return ctx, nopCleanup, fmt.Errorf("unable to open batch: %w", err)
+		return ctx, NopCleanup, fmt.Errorf("unable to open batch: %w", err)
 	}
 	return ctx, leaseDone, nil
 }

soci/util_test.go

@@ -20,16 +20,52 @@ import (
 	"context"
 	"io"
 
+	"github.com/awslabs/soci-snapshotter/soci/store"
 	"github.com/containerd/containerd/content"
 	"github.com/opencontainers/go-digest"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
+	"oras.land/oras-go/v2/content/memory"
 )
 
 func parseDigest(digestString string) digest.Digest {
 	dgst, _ := digest.Parse(digestString)
 	return dgst
 }
 
+type OrasMemoryStore struct {
+	s *memory.Store
+}
+
+func (*OrasMemoryStore) BatchOpen(ctx context.Context) (context.Context, store.CleanupFunc, error) {
+	return ctx, store.NopCleanup, nil
+}
+
+func (m *OrasMemoryStore) Exists(ctx context.Context, target ocispec.Descriptor) (bool, error) {
+	return m.s.Exists(ctx, target)
+}
+
+func (m *OrasMemoryStore) Fetch(ctx context.Context, target ocispec.Descriptor) (io.ReadCloser, error) {
+	return m.s.Fetch(ctx, target)
+}
+
+func (m *OrasMemoryStore) Push(ctx context.Context, expected ocispec.Descriptor, reader io.Reader) error {
+	return m.s.Push(ctx, expected, reader)
+}
+
+func (m *OrasMemoryStore) Label(ctx context.Context, target ocispec.Descriptor, label string, value string) error {
+	return nil
+}
+
+func (m *OrasMemoryStore) Delete(ctx context.Context, dgst digest.Digest) error {
+	return nil
+}
+
+func NewOrasMemoryStore() *OrasMemoryStore {
+	return &OrasMemoryStore{
+		s: memory.New(),
+	}
+}
+
 type fakeContentStore struct {
 }