fix(storage): defer blob deletes for cache-unknown blobs until dedupe rebuild completes (#4195)

With dedupe enabled on cloud storage, duplicate blobs are stored as
zero-size placeholders whose content lives in a single original blob,
resolvable only through the dedupe cache. When the cache is lost or
incomplete (e.g. an ephemeral local cache database with cloud storage),
deleteBlob's cache lookup misses, the move-content-to-next-duplicate
logic is skipped entirely, and the storage delete destroys the only
content-bearing copy. Every duplicate of that digest in other
repositories is then permanently broken (416 RequestedRangeNotSatisfiable
on read), which is the data loss reported in issue #2625.

The startup dedupe/restore walk (RunDedupeBlobs) already rebuilds the
cache from storage, but GC/retention can run before or while it does,
which is exactly the corruption window.

Fix: track completion of the dedupe/restore walk on the ImageStore
(reusing the DedupeTaskGenerator completion callback, renamed to
OnRunComplete and now fired for both dedupe and restore runs) and
defer deletes of content-bearing blobs that the cache cannot account
for until the walk has completed. Such deletes
return the new ErrDedupeRebuildInProgress; GC/retention retries on a
later interval, and blobs known to the cache (or zero-size
placeholders) are unaffected. Local storage is exempt: it dedupes via
hardlinks, so deleting one path never destroys shared content.

The gate is armed only when RunDedupeBlobs is invoked (the server
always does at startup), so stores that never schedule the walk keep
existing behavior. Generator submission order in StartBackgroundTasks
is deliberately unchanged: the walk being scheduled a few statements
after the GC generator leaves only a sub-millisecond arming window,
whereas reordering the submissions shifts the GC generator's first
pass by a scheduler slot and changes long-standing GC timing behavior.

Fixes #2625

Signed-off-by: Johnathan Aretos <johnathan@squareup.com>

Author: johnathan-sq

errors/errors.go

@@ -130,6 +130,7 @@ var (
 	ErrTimeout                          = errors.New("operation timeout")
 	ErrNotImplemented                   = errors.New("not implemented")
 	ErrDedupeRebuild                    = errors.New("couldn't rebuild dedupe index")
+	ErrDedupeRebuildInProgress          = errors.New("dedupe cache rebuild in progress, blob delete deferred")
 	ErrMissingAuthHeader                = errors.New("required authorization header is missing")
 	ErrUserAPIKeyNotFound               = errors.New("user info for given API key hash not found")
 	ErrUserSessionNotFound              = errors.New("user session for given ID not found")

pkg/storage/common/common.go

@@ -918,9 +918,9 @@ type restoreRunState struct {
 	// completed successfully. It is incremented when a task is generated and decremented
 	// only by onDone, which a task calls on success (see dedupeTask.DoWork). A failed task
 	// is therefore never decremented, so the count intentionally never reaches zero again
-	// for this run, preventing checkCompletion from firing OnRestoreComplete after a failure.
+	// for this run, preventing checkCompletion from firing OnRunComplete after a failure.
 	pendingTaskCount atomic.Int64
-	// completeOnce ensures OnRestoreComplete is called at most once for this run.
+	// completeOnce ensures OnRunComplete is called at most once for this run.
 	completeOnce sync.Once
 	// done is set to true once all tasks for this run have been generated. It is read
 	// both by IsDone() on the scheduler goroutine and by checkCompletion() from onDone
@@ -943,10 +943,12 @@ type DedupeTaskGenerator struct {
 	lastDigests []godigest.Digest
 	repos       []string // list of repos on which we run dedupe
 	Log         zlog.Logger
-	// OnRestoreComplete is called exactly once after ALL restore tasks have executed
-	// successfully. Used to write the restore-complete marker so the next startup with
-	// dedupe=false can skip the expensive per-digest scan.
-	OnRestoreComplete func()
+	// OnRunComplete is called exactly once after all tasks of a run have
+	// succeeded, whether the run dedupes blobs (Dedupe=true) or restores them
+	// (Dedupe=false). Restore runs use it to write the restore-complete marker;
+	// both kinds use it to signal the image store that blob deletes may proceed
+	// (see ImageStore.deleteBlob).
+	OnRunComplete func()
 	// run holds the completion-tracking state for the current run. onDone closures
 	// capture the *restoreRunState pointer directly and operate on it independently
 	// of any later Reset(). It is an atomic.Pointer because checkCompletion() compares
@@ -975,7 +977,7 @@ func (gen *DedupeTaskGenerator) getRun() *restoreRunState {
 
 // All restore tasks for this run have been generated AND all of them have completed successfully.
 func (gen *DedupeTaskGenerator) checkCompletion(run *restoreRunState) {
-	if gen.OnRestoreComplete != nil &&
+	if gen.OnRunComplete != nil &&
 		run.done.Load() &&
 		run.pendingTaskCount.Load() == 0 {
 		// Dispatch asynchronously so the executor goroutine that triggered completion
@@ -985,11 +987,11 @@ func (gen *DedupeTaskGenerator) checkCompletion(run *restoreRunState) {
 				defer func() {
 					if r := recover(); r != nil {
 						gen.Log.Error().Interface("panic", r).Str("component", "dedupe").
-							Msg("panic in OnRestoreComplete")
+							Msg("panic in OnRunComplete")
 					}
 				}()
 
-				gen.OnRestoreComplete()
+				gen.OnRunComplete()
 			}()
 		})
 	}
@@ -1044,10 +1046,10 @@ func (gen *DedupeTaskGenerator) Next() (scheduler.Task, error) {
 	// mark digest as processed before running its task
 	gen.lastDigests = append(gen.lastDigests, gen.digest)
 
-	// For restore passes, track each task so the marker is only written after all succeed.
+	// Track each task so completion fires only after all of them succeed.
 	var onDone func()
 
-	if !gen.Dedupe && gen.OnRestoreComplete != nil {
+	if gen.OnRunComplete != nil {
 		run.pendingTaskCount.Add(1)
 
 		onDone = func() {
@@ -1074,10 +1076,10 @@ func (gen *DedupeTaskGenerator) Reset() {
 
 	// Only start a fresh run if the current one has no in-flight tasks (or completion
 	// isn't tracked at all). If tasks are still executing, keep the same run state so
-	// their onDone callbacks can still drive checkCompletion to fire OnRestoreComplete
+	// their onDone callbacks can still drive checkCompletion to fire OnRunComplete
 	// once they finish; replacing gen.run here would otherwise make that run's
 	// completion unobservable forever.
-	if gen.OnRestoreComplete == nil || (run.done.Load() && run.pendingTaskCount.Load() == 0) {
+	if gen.OnRunComplete == nil || (run.done.Load() && run.pendingTaskCount.Load() == 0) {
 		gen.run.Store(&restoreRunState{})
 	}
 
@@ -1095,7 +1097,7 @@ type dedupeTask struct {
 	duplicateBlobs []string
 	dedupe         bool
 	log            zlog.Logger
-	// onDone is called when this restore task succeeds. nil for dedupe (non-restore) tasks.
+	// onDone is called when this task succeeds; nil when completion isn't tracked.
 	onDone func()
 }
 

pkg/storage/common/common_test.go

@@ -624,7 +624,7 @@ func TestDedupeGeneratorErrors(t *testing.T) {
 func TestDedupeTaskGeneratorRestoreComplete(t *testing.T) {
 	testLog := log.NewTestLogger()
 
-	Convey("OnRestoreComplete fires once after all restore tasks finish successfully", t, func(c C) {
+	Convey("OnRunComplete fires once after all restore tasks finish successfully", t, func(c C) {
 		digest := godigest.FromString("blob1")
 		duplicateBlobs := []string{"/repo/blob1-a", "/repo/blob1-b"}
 
@@ -662,7 +662,7 @@ func TestDedupeTaskGeneratorRestoreComplete(t *testing.T) {
 			ImgStore: imgStore,
 			Dedupe:   false,
 			Log:      testLog,
-			OnRestoreComplete: func() {
+			OnRunComplete: func() {
 				callCountMutex.Lock()
 				callCount++
 				callCountMutex.Unlock()
@@ -687,11 +687,11 @@ func TestDedupeTaskGeneratorRestoreComplete(t *testing.T) {
 		So(task, ShouldBeNil)
 		So(generator.IsDone(), ShouldBeTrue)
 
-		// OnRestoreComplete is dispatched asynchronously
+		// OnRunComplete is dispatched asynchronously
 		select {
 		case <-done:
 		case <-time.After(5 * time.Second):
-			t.Fatal("timed out waiting for OnRestoreComplete")
+			t.Fatal("timed out waiting for OnRunComplete")
 		}
 
 		callCountMutex.Lock()
@@ -703,6 +703,76 @@ func TestDedupeTaskGeneratorRestoreComplete(t *testing.T) {
 		So(generator.IsDone(), ShouldBeFalse)
 	})
 
+	Convey("OnRunComplete fires once after all dedupe tasks finish successfully", t, func(c C) {
+		digest := godigest.FromString("blob-dedupe")
+		duplicateBlobs := []string{"/repo/blob-dedupe-a", "/repo/blob-dedupe-b"}
+
+		getNextCalls := 0
+
+		imgStore := &mocks.MockedImageStore{
+			GetRepositoriesFn: func() ([]string, error) {
+				return []string{"repo1"}, nil
+			},
+			GetNextDigestWithBlobPathsFn: func(repos []string, lastDigests []godigest.Digest) (
+				godigest.Digest, []string, error,
+			) {
+				getNextCalls++
+				if getNextCalls == 1 {
+					return digest, duplicateBlobs, nil
+				}
+
+				return "", nil, nil
+			},
+			RunDedupeForDigestFn: func(ctx context.Context, digest godigest.Digest, dedupe bool,
+				duplicateBlobs []string,
+			) error {
+				return nil
+			},
+		}
+
+		var (
+			callCountMutex sync.Mutex
+			callCount      int
+		)
+
+		done := make(chan struct{})
+
+		generator := &common.DedupeTaskGenerator{
+			ImgStore: imgStore,
+			Dedupe:   true,
+			Log:      testLog,
+			OnRunComplete: func() {
+				callCountMutex.Lock()
+				callCount++
+				callCountMutex.Unlock()
+				close(done)
+			},
+		}
+
+		task, err := generator.Next()
+		So(err, ShouldBeNil)
+		So(task, ShouldNotBeNil)
+		So(generator.IsDone(), ShouldBeFalse)
+
+		err = task.DoWork(context.Background())
+		So(err, ShouldBeNil)
+
+		task, err = generator.Next()
+		So(err, ShouldBeNil)
+		So(task, ShouldBeNil)
+		So(generator.IsDone(), ShouldBeTrue)
+
+		select {
+		case <-done:
+		case <-time.After(5 * time.Second):
+			t.Fatal("timed out waiting for OnRunComplete")
+		}
+
+		callCountMutex.Lock()
+		So(callCount, ShouldEqual, 1)
+		callCountMutex.Unlock()
+	})
+
 	Convey("Reset keeps the same run while a restore task is in-flight", t, func(c C) {
 		digest := godigest.FromString("blob2")
 		duplicateBlobs := []string{"/repo/blob2-a"}
@@ -741,7 +811,7 @@ func TestDedupeTaskGeneratorRestoreComplete(t *testing.T) {
 			ImgStore: imgStore,
 			Dedupe:   false,
 			Log:      testLog,
-			OnRestoreComplete: func() {
+			OnRunComplete: func() {
 				callCountMutex.Lock()
 				callCount++
 				callCountMutex.Unlock()
@@ -760,26 +830,26 @@ func TestDedupeTaskGeneratorRestoreComplete(t *testing.T) {
 		So(err, ShouldBeNil)
 		So(generator.IsDone(), ShouldBeTrue)
 
-		// Reset() must keep the same run state, since OnRestoreComplete hasn't fired yet
+		// Reset() must keep the same run state, since OnRunComplete hasn't fired yet
 		generator.Reset()
 		So(generator.IsDone(), ShouldBeTrue)
 
-		// finishing the in-flight task now drives checkCompletion to fire OnRestoreComplete
+		// finishing the in-flight task now drives checkCompletion to fire OnRunComplete
 		err = task.DoWork(context.Background())
 		So(err, ShouldBeNil)
 
 		select {
 		case <-done:
 		case <-time.After(5 * time.Second):
-			t.Fatal("timed out waiting for OnRestoreComplete")
+			t.Fatal("timed out waiting for OnRunComplete")
 		}
 
 		callCountMutex.Lock()
 		So(callCount, ShouldEqual, 1)
 		callCountMutex.Unlock()
 	})
 
-	Convey("checkCompletion recovers if OnRestoreComplete panics", t, func(c C) {
+	Convey("checkCompletion recovers if OnRunComplete panics", t, func(c C) {
 		digest := godigest.FromString("blob3")
 		duplicateBlobs := []string{"/repo/blob3-a"}
 
@@ -816,7 +886,7 @@ func TestDedupeTaskGeneratorRestoreComplete(t *testing.T) {
 			ImgStore: imgStore,
 			Dedupe:   false,
 			Log:      panicLog,
-			OnRestoreComplete: func() {
+			OnRunComplete: func() {
 				defer close(done)
 				panic("boom")
 			},
@@ -836,18 +906,18 @@ func TestDedupeTaskGeneratorRestoreComplete(t *testing.T) {
 		select {
 		case <-done:
 		case <-time.After(5 * time.Second):
-			t.Fatal("timed out waiting for OnRestoreComplete")
+			t.Fatal("timed out waiting for OnRunComplete")
 		}
 
 		for range 100 {
-			if strings.Contains(logBuf.String(), "panic in OnRestoreComplete") {
+			if strings.Contains(logBuf.String(), "panic in OnRunComplete") {
 				break
 			}
 
 			time.Sleep(10 * time.Millisecond)
 		}
 
-		So(logBuf.String(), ShouldContainSubstring, "panic in OnRestoreComplete")
+		So(logBuf.String(), ShouldContainSubstring, "panic in OnRunComplete")
 	})
 
 	Convey("getRun is safe when the first call races across goroutines", t, func(c C) {

pkg/storage/imagestore/imagestore.go

@@ -13,6 +13,7 @@ import (
 	"slices"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 	"unicode/utf8"
 
@@ -54,6 +55,9 @@ type ImageStore struct {
 	linter      common.Lint
 	commit      bool
 	compat      []compat.MediaCompatibility
+	// dedupeRebuildDone is set once RunDedupeBlobs has walked all blobs, i.e. the
+	// cache accounts for every pre-existing blob; see deleteBlob.
+	dedupeRebuildDone atomic.Bool
 }
 
 func (is *ImageStore) Name() string {
@@ -95,6 +99,9 @@ func NewImageStore(rootDir string, cacheDir string, dedupe, commit bool, log zlo
 		events:      recorder,
 	}
 
+	// Deletes are only gated while a dedupe/restore walk is pending; see RunDedupeBlobs.
+	imgStore.dedupeRebuildDone.Store(true)
+
 	return imgStore
 }
 
@@ -1983,7 +1990,7 @@ func (is *ImageStore) CleanupRepo(repo string, blobs []godigest.Digest, removeRe
 func (is *ImageStore) deleteBlob(repo string, digest godigest.Digest) error {
 	blobPath := is.BlobPath(repo, digest)
 
-	_, err := is.storeDriver.Stat(blobPath)
+	binfo, err := is.storeDriver.Stat(blobPath)
 	if err != nil {
 		var pathNotFoundErr driver.PathNotFoundError
 		if errors.As(err, &pathNotFoundErr) {
@@ -2009,6 +2016,16 @@ func (is *ImageStore) deleteBlob(repo string, digest godigest.Digest) error {
 			return err
 		}
 
+		// Cache miss: with dedupe on remote storage this blob may be the only
+		// content copy backing zero-size duplicates elsewhere.
+		// Defer until the startup walk has rebuilt the cache; GC retries later.
+		if dstRecord == "" && binfo.Size() > 0 && !is.dedupeRebuildDone.Load() {
+			is.log.Warn().Str("digest", digest.String()).Str("blobPath", blobPath).Str("component", "dedupe").
+				Msg("no cache record for content blob while dedupe rebuild is still running, deferring delete")
+
+			return zerr.ErrDedupeRebuildInProgress
+		}
+
 		// remove cache entry and move blob contents to the next candidate if there is any
 		if ok := is.cache.HasBlob(digest, blobPath); ok {
 			if err := is.cache.DeleteBlob(digest, blobPath); err != nil {
@@ -2056,6 +2073,15 @@ func (is *ImageStore) deleteBlob(repo string, digest godigest.Digest) error {
 		}
 	}
 
+	// No cache (dedupe off): leftover placeholders are refilled by the restore
+	// walk; until it completes this blob may be their only content copy.
+	if fmt.Sprintf("%v", is.cache) == fmt.Sprintf("%v", nil) && binfo.Size() > 0 && !is.dedupeRebuildDone.Load() {
+		is.log.Warn().Str("digest", digest.String()).Str("blobPath", blobPath).Str("component", "dedupe").
+			Msg("content blob delete requested before dedupe restore walk finished, deferring delete")
+
+		return zerr.ErrDedupeRebuildInProgress
+	}
+
 	if err := is.storeDriver.Delete(blobPath); err != nil {
 		var pathNotFoundErr driver.PathNotFoundError
 		if errors.As(err, &pathNotFoundErr) {
@@ -2442,6 +2468,12 @@ func (is *ImageStore) RunDedupeForDigest(ctx context.Context, digest godigest.Di
 func (is *ImageStore) RunDedupeBlobs(interval time.Duration, sch *scheduler.Scheduler) {
 	markerPath := path.Join(is.rootDir, storageConstants.DedupeRestoreCompleteMarker)
 
+	// Gate deletes of cache-unknown blobs until the walk completes (see deleteBlob).
+	// Local storage dedupes via hardlinks, so deletes there never destroy shared content.
+	if is.storeDriver.Name() != storageConstants.LocalStorageDriverName {
+		is.dedupeRebuildDone.Store(false)
+	}
+
 	if is.dedupe {
 		// Dedupe is active: remove the restore-complete marker so that a future dedupe→false
 		// transition knows it must run restore again.
@@ -2470,6 +2502,9 @@ func (is *ImageStore) RunDedupeBlobs(interval time.Duration, sch *scheduler.Sche
 				is.log.Info().Str("component", "dedupe").
 					Msg("restore-complete marker present, skipping dedupe restore scan")
 
+				// storage holds no deduped blobs, so deletes are safe without a walk
+				is.dedupeRebuildDone.Store(true)
+
 				return
 			}
 
@@ -2490,16 +2525,21 @@ func (is *ImageStore) RunDedupeBlobs(interval time.Duration, sch *scheduler.Sche
 		Log:      is.log,
 	}
 
-	if !is.dedupe {
-		generator.OnRestoreComplete = func() {
-			if _, err := is.storeDriver.WriteFile(markerPath,
-				[]byte(storageConstants.DedupeRestoreMarkerComplete)); err != nil {
-				is.log.Error().Err(err).Str("component", "dedupe").
-					Msg("failed to write restore-complete marker")
-			} else {
-				is.log.Info().Str("component", "dedupe").
-					Msg("restore-complete marker written; future startups will skip the restore scan")
-			}
+	generator.OnRunComplete = func() {
+		// walk finished: deferred blob deletes may proceed (see deleteBlob)
+		is.dedupeRebuildDone.Store(true)
+
+		if is.dedupe {
+			return
+		}
+
+		if _, err := is.storeDriver.WriteFile(markerPath,
+			[]byte(storageConstants.DedupeRestoreMarkerComplete)); err != nil {
+			is.log.Error().Err(err).Str("component", "dedupe").
+				Msg("failed to write restore-complete marker")
+		} else {
+			is.log.Info().Str("component", "dedupe").
+				Msg("restore-complete marker written; future startups will skip the restore scan")
 		}
 	}