Merge pull request #338 from rchincha/gc

GC: add support for GC policies

oci/casext/gc.go

@@ -25,6 +25,9 @@ import (
 	"golang.org/x/net/context"
 )
 
+// GCPolicy is a policy function that returns 'true' if a blob can be GC'ed
+type GCPolicy func(ctx context.Context, digest digest.Digest) (bool, error)
+
 // GC will perform a mark-and-sweep garbage collection of the OCI image
 // referenced by the given CAS engine. The root set is taken to be the set of
 // references stored in the image, and all blobs not reachable by following a
@@ -35,7 +38,11 @@ import (
 // functions. In other words, it assumes it is the only user of the image that
 // is making modifications. Things will not go well if this assumption is
 // challenged.
-func (e Engine) GC(ctx context.Context) error {
+//
+// Furthermore, GC policies (zero or more) can also be specified which given a
+// blob's digest can indicate whether that blob needs to garbage collected. The
+// blob is skipped for garbage collection if a policy returns false.
+func (e Engine) GC(ctx context.Context, policies ...GCPolicy) error {
 	// Generate the root set of descriptors.
 	var root []ispec.Descriptor
 
@@ -74,12 +81,26 @@ func (e Engine) GC(ctx context.Context) error {
 	}
 
 	n := 0
+sweep:
 	for _, digest := range blobs {
 		if _, ok := black[digest]; ok {
 			// Digest is in the black set.
 			continue
 		}
-		log.Infof("garbage collecting blob: %s", digest)
+
+		for i, policy := range policies {
+			ok, err := policy(ctx, digest)
+			if err != nil {
+				return errors.Wrapf(err, "invoking policy %d failed", i)
+			}
+
+			if !ok {
+				// skip this blob for GC
+				log.Debugf("skipping garbage collection of blob %s because of policy %d", digest, i)
+				continue sweep
+			}
+		}
+		log.Debugf("garbage collecting blob: %s", digest)
 
 		if err := e.DeleteBlob(ctx, digest); err != nil {
 			return errors.Wrapf(err, "remove unmarked blob %s", digest)

oci/casext/gc_test.go

@@ -3,12 +3,14 @@ package casext
 import (
 	"bytes"
 	"encoding/json"
+	"fmt"
 	"io/ioutil"
 	"os"
 	"path/filepath"
 	"strings"
 	"testing"
 
+	"github.com/opencontainers/go-digest"
 	imeta "github.com/opencontainers/image-spec/specs-go"
 	ispec "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/opencontainers/umoci/oci/cas/dir"
@@ -175,3 +177,143 @@ func TestGCWithNonEmptyIndex(t *testing.T) {
 		t.Fatalf("expected single-entry blob list after GC")
 	}
 }
+
+func gcOkFunc(t *testing.T, expectedDigest digest.Digest, unexpectedDigest digest.Digest) GCPolicy {
+	return func(ctx context.Context, digest digest.Digest) (bool, error) {
+		if digest == "" || digest == unexpectedDigest {
+			t.Errorf("got incorrect digest to gc policy callback: unexpected %v", digest)
+		}
+		if digest != expectedDigest {
+			t.Errorf("got incorrect digest to gc policy callback: expected %v, got %v", expectedDigest, digest)
+		}
+		return true, nil
+	}
+}
+
+func gcSkipFunc(t *testing.T, expectedDigest digest.Digest) GCPolicy {
+	return func(ctx context.Context, digest digest.Digest) (bool, error) {
+		if digest != expectedDigest {
+			t.Errorf("got incorrect digest to gc policy callback: expected %v, got %v", expectedDigest, digest)
+		}
+		return false, nil
+	}
+}
+
+func errFunc(ctx context.Context, digest digest.Digest) (bool, error) {
+	return false, fmt.Errorf("err policy")
+}
+
+func TestGCWithPolicy(t *testing.T) {
+	ctx := context.Background()
+
+	root, err := ioutil.TempDir("", "umoci-TestEngineReference")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.RemoveAll(root)
+
+	image := filepath.Join(root, "image")
+	if err := dir.Create(image); err != nil {
+		t.Fatalf("unexpected error creating image: %+v", err)
+	}
+
+	engine, err := dir.Open(image)
+	if err != nil {
+		t.Fatalf("unexpected error opening image: %+v", err)
+	}
+	engineExt := NewEngine(engine)
+	defer engine.Close()
+
+	// build a orphan blob that should be GC'ed
+	content := "this is a orphan blob"
+	br := strings.NewReader(content)
+	odigest, size, err := engine.PutBlob(ctx, br)
+	if err != nil {
+		t.Fatalf("error writing blob: %+v", err)
+	}
+	if size != int64(len(content)) {
+		t.Fatalf("partially written blob")
+	}
+
+	// build a blob, manifest, index that will survive GC
+	content = "this is a test blob"
+	br = strings.NewReader(content)
+	digest, size, err := engine.PutBlob(ctx, br)
+	if err != nil {
+		t.Fatalf("error writing blob: %+v", err)
+	}
+	if size != int64(len(content)) {
+		t.Fatalf("partially written blob")
+	}
+
+	digest, size, err = engineExt.PutBlobJSON(ctx,
+		ispec.Manifest{
+			Versioned: imeta.Versioned{
+				SchemaVersion: 2,
+			},
+			Config: ispec.Descriptor{
+				MediaType: ispec.MediaTypeImageLayer,
+				Digest:    digest,
+				Size:      size,
+			},
+			Layers: []ispec.Descriptor{
+				{
+					MediaType: ispec.MediaTypeImageLayer,
+					Digest:    digest,
+					Size:      size,
+				},
+			},
+		})
+	if err != nil {
+		t.Fatalf("error writing blob: %+v", err)
+	}
+
+	idx := ispec.Index{
+		Versioned: imeta.Versioned{
+			SchemaVersion: 2,
+		},
+		Manifests: []ispec.Descriptor{
+			{
+				MediaType: ispec.MediaTypeImageManifest,
+				Digest:    digest,
+				Size:      size,
+			},
+		},
+	}
+	if err := engine.PutIndex(ctx, idx); err != nil {
+		t.Fatalf("error writing index: %+v", err)
+	}
+
+	err = engineExt.GC(ctx, errFunc)
+	// expect this to fail
+	if err == nil {
+		t.Fatalf("GC failed: %+v", err)
+	}
+
+	err = engineExt.GC(ctx, gcSkipFunc(t, odigest))
+	// expect this to succeed but not perform GC
+	if err != nil {
+		t.Fatalf("GC failed: %+v", err)
+	}
+	b, err := engine.ListBlobs(ctx)
+	if err != nil {
+		t.Fatalf("unable to list blobs: %+v", err)
+	}
+	if len(b) != 3 {
+		t.Fatalf("expected all entries in blob list after skip GC policy")
+	}
+
+	err = engineExt.GC(ctx, gcOkFunc(t, odigest, digest))
+	// expect this to succeed
+	if err != nil {
+		t.Fatalf("GC failed: %+v", err)
+	}
+
+	b, err = engine.ListBlobs(ctx)
+	if err != nil {
+		t.Fatalf("unable to list blobs: %+v", err)
+	}
+	if len(b) != 2 {
+		t.Fatalf("expected blob list with two entries after GC")
+	}
+}