fix(sandboxd): retain checkpoints through failed claim writes

sandboxd/pool/archive.go

@@ -122,10 +122,7 @@ func (m *Manager) archive(ctx context.Context, sb *types.Sandbox) error {
 		intent := m.store.set(sb)
 		m.mu.Unlock()
 		if err = m.store.commit(intent); err != nil {
-			m.mu.Lock()
-			sb.ArchivePending = ""
-			rb := m.store.set(sb)
-			m.mu.Unlock()
+			rb, _ := m.setIfLive(sb, func() { sb.ArchivePending = "" })
 			sb.Transition.Unlock()
 			m.store.recommit(ctx, rb)
 			m.deleteOrphanArchiveCk(ctx, ck.ID)
@@ -280,7 +277,7 @@ func (m *Manager) commitWake(ctx context.Context, sb *types.Sandbox, vmName, soc
 }
 
 func (m *Manager) deleteOrphanArchiveCk(ctx context.Context, ckID string) {
-	if err := m.deleteArchiveCk(ctx, ckID); err != nil {
+	if _, err := m.deleteArchiveCk(ctx, ckID); err != nil {
 		log.WithFunc("pool.deleteOrphanArchiveCk").Warnf(ctx, "delete orphaned archive ck %s: %v", ckID, err)
 	}
 }
@@ -351,17 +348,28 @@ func (m *Manager) clearArchiveCk(ckID string) error {
 	return err
 }
 
-func (m *Manager) deleteArchiveCk(ctx context.Context, ckID string) error {
+func (m *Manager) deleteArchiveCk(ctx context.Context, ckID string) (bool, error) {
 	if err := m.markArchiveCk(ckID); err != nil {
-		return fmt.Errorf("track: %w", err)
+		return false, fmt.Errorf("track: %w", err)
 	}
-	if err := m.deleteCkLocked(ctx, ckID); err != nil {
-		return err
+	return m.finishArchiveCkDelete(ctx, ckID)
+}
+
+func (m *Manager) finishArchiveCkDelete(ctx context.Context, ckID string) (bool, error) {
+	unlock := m.lockRec(ckID)
+	if m.store.referencesCheckpoint(ckID) {
+		unlock()
+		return false, nil
+	}
+	err := m.ckpts.Delete(ctx, ckID)
+	unlock()
+	if err != nil {
+		return false, err
 	}
 	if err := m.clearArchiveCk(ckID); err != nil {
-		return fmt.Errorf("clear: %w", err)
+		return false, fmt.Errorf("clear: %w", err)
 	}
-	return nil
+	return true, nil
 }
 
 func (m *Manager) retryArchiveDeletes(ctx context.Context) {

sandboxd/pool/archive_test.go

@@ -522,7 +522,7 @@ func TestArchiveWakeRequiresDeleteMarker(t *testing.T) {
 	if !archived || !ckExists(t, m, ck) {
 		t.Fatal("failed wake did not retain the archived claim and checkpoint")
 	}
-	claims, err := m.store.load()
+	claims, err := newClaimStore(m.dataDir, false).load()
 	if err != nil || claims[sb.ID] == nil || claims[sb.ID].ArchiveCk != ck {
 		t.Fatalf("failed wake changed the durable archive claim: %+v, %v", claims[sb.ID], err)
 	}
@@ -680,11 +680,14 @@ func TestArchiveDeleteRetryRechecksWakeRollback(t *testing.T) {
 		m.retryArchiveDeletes(t.Context())
 		close(retried)
 	}()
-	waitFor(t, func() bool {
-		m.recLocksMu.Lock()
-		defer m.recLocksMu.Unlock()
-		return m.recRefs[ck] >= 2
-	})
+	select {
+	case <-retried:
+	case <-time.After(3 * time.Second):
+		t.Fatal("retry did not defer the checkpoint still referenced by the journal")
+	}
+	if !ckExists(t, m, ck) {
+		t.Fatal("retry deleted the checkpoint before wake persistence completed")
+	}
 	breakStore(t, m)
 	m.store.writeMu.Unlock()
 	locked = false

sandboxd/pool/checkpoint.go

@@ -318,18 +318,8 @@ func (m *Manager) vetoIfHealPending(ckptID string) {
 func (m *Manager) pinnedArchiveCks() map[string]struct{} {
 	m.mu.Lock()
 	defer m.mu.Unlock()
-	pinned := maps.Clone(m.pendingCks)
-	for _, sb := range m.claimed {
-		if sb.ArchiveCk != "" {
-			pinned[sb.ArchiveCk] = struct{}{}
-		}
-		if sb.ReplicaCk != "" {
-			pinned[sb.ReplicaCk] = struct{}{}
-		}
-		if sb.ReplicaPending != "" {
-			pinned[sb.ReplicaPending] = struct{}{}
-		}
-	}
+	pinned := m.store.checkpointReferences()
+	maps.Copy(pinned, m.pendingCks)
 	return pinned
 }
 

sandboxd/pool/checkpoint_bench_test.go

@@ -0,0 +1,67 @@
+package pool
+
+import (
+	"fmt"
+	"path/filepath"
+	"testing"
+	"time"
+
+	"github.com/cocoonstack/sandbox/sandboxd/config"
+	"github.com/cocoonstack/sandbox/sandboxd/types"
+)
+
+func BenchmarkClaimCommitWithCheckpoints(b *testing.B) {
+	for _, checkpoints := range []bool{false, true} {
+		b.Run(fmt.Sprintf("claims=5000/checkpoints=%t", checkpoints), func(b *testing.B) {
+			claims := checkpointBenchClaims(checkpoints)
+			s := newClaimStore(b.TempDir(), false)
+			if err := s.save(claims); err != nil {
+				b.Fatal(err)
+			}
+			one := claims["sb_0000000000000000"]
+			b.ReportAllocs()
+			for b.Loop() {
+				one.Deadline = one.Deadline.Add(time.Second)
+				if err := s.commit(s.set(one)); err != nil {
+					b.Fatal(err)
+				}
+			}
+			loaded, err := newClaimStore(filepath.Dir(s.path), false).load()
+			if err != nil || len(loaded) != 5000 || !loaded[one.ID].Deadline.Equal(one.Deadline) || loaded[one.ID].ReplicaCk != one.ReplicaCk {
+				b.Fatalf("commit output: rows=%d err=%v", len(loaded), err)
+			}
+		})
+	}
+}
+
+func BenchmarkPinnedArchiveCheckpoints(b *testing.B) {
+	for _, checkpoints := range []bool{false, true} {
+		b.Run(fmt.Sprintf("claims=5000/checkpoints=%t", checkpoints), func(b *testing.B) {
+			m := newTestManager(b, newFakeEngine(), config.PoolSpec{PoolKey: testKey})
+			m.claimed = checkpointBenchClaims(checkpoints)
+			if err := m.store.save(m.claimed); err != nil {
+				b.Fatal(err)
+			}
+			want := 0
+			if checkpoints {
+				want = 5000
+			}
+			b.ReportAllocs()
+			for b.Loop() {
+				if got := len(m.pinnedArchiveCks()); got != want {
+					b.Fatalf("pinned=%d want=%d", got, want)
+				}
+			}
+		})
+	}
+}
+
+func checkpointBenchClaims(checkpoints bool) map[string]*types.Sandbox {
+	claims := benchClaims(5000)
+	if checkpoints {
+		for id, sb := range claims {
+			sb.ReplicaCk = "ck_" + id[3:]
+		}
+	}
+	return claims
+}

sandboxd/pool/checkpoint_cleanup_test.go

@@ -0,0 +1,223 @@
+package pool
+
+import (
+	"errors"
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/cocoonstack/sandbox/sandboxd/config"
+	"github.com/cocoonstack/sandbox/sandboxd/types"
+)
+
+func TestCheckpointCleanupWaitsForTheClaimWrite(t *testing.T) {
+	replicated := func(t *testing.T, m *Manager, _ *failoverCluster, sb *types.Sandbox) string {
+		replicateNow(t, m, sb)
+		return sb.ReplicaCk
+	}
+	keepsTombstone := func(t *testing.T, m *Manager, c *failoverCluster, sb *types.Sandbox) {
+		rec, ok, err := c.reg.Get(t.Context(), sb.ID)
+		if err != nil || !ok || !rec.Consumed {
+			t.Fatalf("retirement lost its tombstone: %+v, %t, %v", rec, ok, err)
+		}
+		m.reconcileOwnRow(t.Context(), rec, time.Now())
+		if _, ok, err = c.reg.Get(t.Context(), sb.ID); err != nil || !ok {
+			t.Fatalf("sweep erased an unacknowledged retirement: %t, %v", ok, err)
+		}
+	}
+	for _, tc := range []struct {
+		name  string
+		setup func(t *testing.T, m *Manager, c *failoverCluster, sb *types.Sandbox) string
+		act   func(t *testing.T, m *Manager, c *failoverCluster, sb *types.Sandbox)
+	}{
+		{
+			name:  "replica retirement",
+			setup: replicated,
+			act: func(t *testing.T, m *Manager, c *failoverCluster, sb *types.Sandbox) {
+				sb.Transition.Lock()
+				err := m.dropReplica(t.Context(), sb)
+				sb.Transition.Unlock()
+				if err == nil || sb.HasReplica() {
+					t.Fatalf("dropReplica: replica=%t, %v; want a memory clear and the write error", sb.HasReplica(), err)
+				}
+				keepsTombstone(t, m, c, sb)
+			},
+		},
+		{
+			name:  "retirement of a released claim",
+			setup: replicated,
+			act: func(t *testing.T, m *Manager, c *failoverCluster, sb *types.Sandbox) {
+				m.retireReplica(t.Context(), sb)
+				m.mu.Lock()
+				delete(m.claimed, sb.ID)
+				m.store.del(sb.ID)
+				m.mu.Unlock()
+				keepsTombstone(t, m, c, sb)
+			},
+		},
+		{
+			name: "unpublished replica intent",
+			setup: func(t *testing.T, m *Manager, c *failoverCluster, sb *types.Sandbox) string {
+				m.archives = &hookRegistry{Registry: c.reg, failPut: true}
+				if err := m.replicate(t.Context(), sb); !errors.Is(err, ErrArchiveStoreDown) {
+					t.Fatal(err)
+				}
+				m.archives = c.reg
+				return sb.ReplicaPending
+			},
+			act: func(t *testing.T, m *Manager, _ *failoverCluster, sb *types.Sandbox) {
+				m.retireReplica(t.Context(), sb)
+				if sb.ReplicaPending != "" {
+					t.Fatal("retirement kept the intent in memory")
+				}
+			},
+		},
+		{
+			name: "settled replica intent",
+			setup: func(t *testing.T, m *Manager, c *failoverCluster, sb *types.Sandbox) string {
+				replicateNow(t, m, sb)
+				previous := sb.ReplicaCk
+				m.archives = &hookRegistry{Registry: c.reg, losePutReply: true}
+				if err := m.replicate(t.Context(), sb); !errors.Is(err, ErrArchiveStoreDown) {
+					t.Fatal(err)
+				}
+				m.archives = c.reg
+				return previous
+			},
+			act: func(t *testing.T, m *Manager, _ *failoverCluster, sb *types.Sandbox) {
+				pending := sb.ReplicaPending
+				sb.Transition.Lock()
+				err := m.settleReplicaIntent(t.Context(), sb)
+				sb.Transition.Unlock()
+				if err == nil || sb.ReplicaPending != "" || sb.ReplicaCk != pending {
+					t.Fatalf("settle: %+v, %v; want the intent settled in memory and the write error", sb, err)
+				}
+			},
+		},
+	} {
+		for _, restart := range []bool{false, true} {
+			t.Run(fmt.Sprintf("%s/restart=%t", tc.name, restart), func(t *testing.T) {
+				c := newReplicaCluster(t)
+				m := c.node(t, "a", t.TempDir(), keepAnHour)
+				sb := mustClaim(t, m, testKey)
+				ck := tc.setup(t, m, c, sb)
+				restoreWrites := failClaimWrites(t, m.store)
+				tc.act(t, m, c, sb)
+				m.retryArchiveDeletes(t.Context())
+				if !ckExists(t, m, ck) || !m.deleteMarked(ck) {
+					t.Fatal("cleanup dropped the checkpoint or its delete marker before the claim write landed")
+				}
+				restoreWrites()
+				if restart {
+					m = c.restart(t, m)
+					if err := m.Reconcile(t.Context()); err != nil {
+						t.Fatal(err)
+					}
+					var ok bool
+					if sb, ok = m.byID(sb.ID); !ok {
+						t.Fatal("restart dropped the claim")
+					}
+					m.retireReplica(t.Context(), sb)
+				} else if err := m.store.commit(m.store.mark()); err != nil {
+					t.Fatal(err)
+				}
+				other := mustClaim(t, m, testKey)
+				m.mu.Lock()
+				other.LeaseSeconds++
+				m.store.set(other)
+				m.mu.Unlock()
+				m.retryArchiveDeletes(t.Context())
+				if ckExists(t, m, ck) || m.deleteMarked(ck) {
+					t.Fatal("cleanup did not finish once the claim write landed")
+				}
+				for _, live := range checkpointRefs(sb) {
+					if live != "" && !ckExists(t, m, live) {
+						t.Errorf("cleanup deleted live checkpoint %s", live)
+					}
+				}
+				m.sweepOrphanArchives(t.Context())
+				if rec, ok, _ := c.reg.Get(t.Context(), sb.ID); ok && rec.Consumed {
+					t.Error("sweep kept the tombstone after the claim write landed")
+				}
+			})
+		}
+	}
+}
+
+func TestFailoverDisabledExitKeepsReplicaCheckpoint(t *testing.T) {
+	states := []struct {
+		name  string
+		setup func(t *testing.T, m *Manager, c *failoverCluster, sb *types.Sandbox) string
+	}{
+		{"published replica", func(t *testing.T, m *Manager, _ *failoverCluster, sb *types.Sandbox) string {
+			replicateNow(t, m, sb)
+			return sb.ReplicaCk
+		}},
+		{"pending replica", func(t *testing.T, m *Manager, c *failoverCluster, sb *types.Sandbox) string {
+			m.archives = &hookRegistry{Registry: c.reg, failPut: true}
+			if err := m.replicate(t.Context(), sb); !errors.Is(err, ErrArchiveStoreDown) {
+				t.Fatal(err)
+			}
+			return sb.ReplicaPending
+		}},
+		{"archive intent", func(t *testing.T, m *Manager, c *failoverCluster, sb *types.Sandbox) string {
+			if err := m.Hibernate(t.Context(), sb.ID, Cred{Token: sb.Token}); err != nil {
+				t.Fatal(err)
+			}
+			m.archives = &hookRegistry{Registry: c.reg, failPut: true}
+			if err := m.archive(t.Context(), sb); !errors.Is(err, ErrArchiveStoreDown) {
+				t.Fatal(err)
+			}
+			return sb.ArchivePending
+		}},
+	}
+	exits := []struct {
+		name string
+		run  func(t *testing.T, m *Manager, sb *types.Sandbox)
+	}{
+		{"release", func(t *testing.T, m *Manager, sb *types.Sandbox) {
+			if err := m.Release(t.Context(), sb.ID, Cred{Token: sb.Token}); err != nil {
+				t.Fatal(err)
+			}
+		}},
+		{"expire", func(t *testing.T, m *Manager, sb *types.Sandbox) {
+			m.mu.Lock()
+			sb.Deadline = time.Now().Add(-time.Minute)
+			m.mu.Unlock()
+			m.reapOnce(t.Context())
+			waitFor(t, func() bool { return m.Counters().Reaps == 1 })
+		}},
+	}
+	for _, state := range states {
+		for _, exit := range exits {
+			t.Run(state.name+"/"+exit.name, func(t *testing.T) {
+				c := newReplicaCluster(t)
+				old := c.node(t, "a", t.TempDir(), keepAnHour)
+				sb := mustClaim(t, old, testKey)
+				ck := state.setup(t, old, c, sb)
+				cfg := &config.Config{
+					APIToken: "root", DataDir: old.dataDir, CheckpointDir: c.ckDir,
+					Pools: []config.PoolSpec{{PoolKey: testKey}},
+				}
+				m, err := NewManager(t.Context(), cfg, old.eng, testSecrets(t), testBackends(t, cfg))
+				if err != nil {
+					t.Fatal(err)
+				}
+				if err = m.Reconcile(t.Context()); err != nil {
+					t.Fatal(err)
+				}
+				restored, ok := m.byID(sb.ID)
+				if !ok || checkpointRefs(restored) != checkpointRefs(sb) {
+					t.Fatalf("startup without a registry changed the claim: %+v", restored)
+				}
+				if !ckExists(t, m, ck) || m.deleteMarked(ck) {
+					t.Fatalf("startup without a registry queued or deleted checkpoint %s", ck)
+				}
+				exit.run(t, m, restored)
+				if _, held := m.byID(sb.ID); held || !ckExists(t, m, ck) {
+					t.Fatalf("%s without a registry: held=%t checkpoint kept=%t", exit.name, held, ckExists(t, m, ck))
+				}
+			})
+		}
+	}
+}