pool: a settled archive intent keeps a replicated claim's row

docs/cluster.md

@@ -187,7 +187,10 @@ replication rounds write it, so the claim path is unchanged.
   answers 404. A tombstone is never adopted. The next exec or file call wakes
   the claim from the store as usual. An owner with a live lease answers for
   itself, and a node that tries to take its row is refused with a 404.
-  A draining node never adopts. Adoption checks no tenant `max_claims` or node
+  A draining node never adopts, nor does a fenced node or one that lifted its
+  fence less than 8 seconds ago (one renewal interval plus one registry query),
+  so after a registry outage every healthy owner renews before a peer can take
+  its rows. Adoption checks no tenant `max_claims` or node
   `max_claims`: the claim already exists, and a refusal could strand it.
 - **Fencing.** Take, consume and delete are compare-and-sets on the row's
   epoch, and so is the write of a replica round and of the archive of a claim
@@ -217,7 +220,8 @@ replication rounds write it, so the claim path is unchanged.
   restart that finds the intent and no row keeps the claim while that
   checkpoint is still in the store, since the row never committed, and drops it
   once the checkpoint is gone, since only a node that consumed the row deletes
-  it.
+  it. A replicated claim's row stays: one the archive moved goes back to the
+  replica, so the row is never absent while the claim is held.
   A put whose outcome is unknown leaves the intent in place, and a release,
   the next archive and the hourly sweep settle it by the same rule before
   they act on the claim.
@@ -359,9 +363,11 @@ checkpoint store and keeps each claim's registry row on the latest snapshot:
   the node, and the node destroys nothing. Claims of pools that do not
   replicate stay served. A node starts fenced and lifts the fence after its
   first renewal that is followed by a confirmation. The confirmation runs beside the renew loop, which keeps
-  renewing: for each replicated claim the node takes the row back for itself in
-  one compare-and-set, which the row lock serializes against adopters, and
-  retries once on a stale epoch. A claim whose row another node took, or took
+  renewing: the node takes back every live row it holds in one statement, a
+  compare-and-set on each row's epoch that the row locks serialize against
+  adopters, and records the new epochs in one journal write; a claim whose row
+  that statement did not move takes it in its own compare-and-set and retries
+  once on a stale epoch. A claim whose row another node took, or took
   and finished, is forgotten at once and its VM removed in the background; a
   claim whose row is still this node's but keeps refusing the write leaves the
   node fenced and is retried at the next renewal. The lift replaces the fenced

sandboxd/archives/archives.go

@@ -37,6 +37,8 @@ type Registry interface {
 	Replace(ctx context.Context, rec Record, expect int64) (int64, error)
 	// Take moves the live row at epoch expect to owner at a new epoch and returns it; another node's row moves only once its owner's lease expired.
 	Take(ctx context.Context, id string, expect int64, owner string) (int64, error)
+	// TakeOwned moves each live row owner holds at its epoch in expect to a new epoch in one call and returns the new epochs of the rows it moved.
+	TakeOwned(ctx context.Context, owner string, expect map[string]int64) (map[string]int64, error)
 	// Consume turns the live row at epoch expect into a tombstone at a new epoch and returns it.
 	Consume(ctx context.Context, id string, expect int64) (int64, error)
 	// Delete removes the row only while it is at epoch.

sandboxd/archives/archivestest/archivestest.go

@@ -83,6 +83,20 @@ func (m *Mem) Take(_ context.Context, id string, expect int64, owner string) (in
 	})
 }
 
+func (m *Mem) TakeOwned(_ context.Context, owner string, expect map[string]int64) (map[string]int64, error) {
+	taken := map[string]int64{}
+	for id, epoch := range expect {
+		next, err := m.advance(id, epoch, func(rec *archives.Record) bool { return rec.Owner == owner })
+		switch {
+		case err == nil:
+			taken[id] = next
+		case !errors.Is(err, archives.ErrLost):
+			return nil, err
+		}
+	}
+	return taken, nil
+}
+
 func (m *Mem) Consume(_ context.Context, id string, expect int64) (int64, error) {
 	return m.advance(id, expect, func(rec *archives.Record) bool {
 		rec.Consumed = true
@@ -206,6 +220,7 @@ func Run(t *testing.T, open Opener) {
 	t.Run("put fences other owners", func(t *testing.T) { testPutFences(t, open(t)) })
 	t.Run("take is a compare-and-set", func(t *testing.T) { testTake(t, open(t)) })
 	t.Run("one of concurrent takes wins", func(t *testing.T) { testTakeRace(t, open(t)) })
+	t.Run("take owned moves only the owner's live rows at their epochs", func(t *testing.T) { testTakeOwned(t, open(t)) })
 	t.Run("consume leaves a tombstone", func(t *testing.T) { testConsume(t, open(t)) })
 	t.Run("delete only at its epoch", func(t *testing.T) { testDelete(t, open(t)) })
 	t.Run("epochs are never reused", func(t *testing.T) { testEpochsNeverRepeat(t, open(t)) })
@@ -306,6 +321,37 @@ func testTakeRace(t *testing.T, r archives.Registry) {
 	}
 }
 
+func testTakeOwned(t *testing.T, r archives.Registry) {
+	ctx := t.Context()
+	mine := mustPut(t, r, record("sb_mine", "a", time.Time{}))
+	stale := mustPut(t, r, record("sb_moved", "a", time.Time{}))
+	moved := mustPut(t, r, record("sb_moved", "a", time.Time{}))
+	other := mustPut(t, r, record("sb_other", "b", time.Time{}))
+	spent, err := r.Consume(ctx, "sb_spent", mustPut(t, r, record("sb_spent", "a", time.Time{})))
+	if err != nil {
+		t.Fatalf("consume: %v", err)
+	}
+	taken, err := r.TakeOwned(ctx, "a", map[string]int64{"sb_mine": mine, "sb_moved": stale, "sb_other": other, "sb_spent": spent, "sb_none": 1})
+	if err != nil || len(taken) != 1 || taken["sb_mine"] <= spent {
+		t.Fatalf("take owned: %v %v, want only sb_mine at a new epoch", taken, err)
+	}
+	put := record("sb_mine", "a", time.Time{})
+	if got, _, _ := r.Get(ctx, "sb_mine"); got.Owner != "a" || got.Epoch != taken["sb_mine"] || got.Consumed || got.Checkpoint != put.Checkpoint || string(got.Claim) != string(put.Claim) {
+		t.Errorf("row after take owned: %+v, want a@%d with the row", got, taken["sb_mine"])
+	}
+	for id, want := range map[string]int64{"sb_moved": moved, "sb_other": other, "sb_spent": spent} {
+		if got, _, _ := r.Get(ctx, id); got.Epoch != want {
+			t.Errorf("%s moved to epoch %d, want it left at %d", id, got.Epoch, want)
+		}
+	}
+	if _, ok, _ := r.Get(ctx, "sb_none"); ok {
+		t.Error("take owned inserted a row")
+	}
+	if taken, err = r.TakeOwned(ctx, "a", nil); err != nil || len(taken) != 0 {
+		t.Errorf("take owned of no ids: %v %v, want none", taken, err)
+	}
+}
+
 func testConsume(t *testing.T, r archives.Registry) {
 	ctx := t.Context()
 	first := mustPut(t, r, record("sb_1", "a", time.Time{}))

sandboxd/archives/pg/pg.go

@@ -78,6 +78,34 @@ func (r *Registry) Take(ctx context.Context, id string, expect int64, owner stri
 		RETURNING epoch`, r.cell, id, expect, owner)
 }
 
+func (r *Registry) TakeOwned(ctx context.Context, owner string, expect map[string]int64) (map[string]int64, error) {
+	ids, epochs := make([]string, 0, len(expect)), make([]int64, 0, len(expect))
+	for id, epoch := range expect {
+		ids, epochs = append(ids, id), append(epochs, epoch)
+	}
+	qctx, cancel := context.WithTimeout(ctx, metastore.QueryTimeout)
+	defer cancel()
+	rows, err := r.pool.Query(qctx, `UPDATE sandboxd_archives a SET epoch = nextval('sandboxd_archives_epoch'), updated_at = now()
+		FROM unnest($2::text[], $3::bigint[]) AS t(id, epoch)
+		WHERE a.cell = $1 AND a.id = t.id AND a.epoch = t.epoch AND a.owner = $4 AND NOT a.consumed
+		RETURNING a.id, a.epoch`, r.cell, ids, epochs, owner)
+	if err != nil {
+		return nil, unavailable(err)
+	}
+	taken := make(map[string]int64, len(expect))
+	var (
+		id    string
+		epoch int64
+	)
+	if _, err = pgx.ForEachRow(rows, []any{&id, &epoch}, func() error {
+		taken[id] = epoch
+		return nil
+	}); err != nil {
+		return nil, unavailable(err)
+	}
+	return taken, nil
+}
+
 func (r *Registry) Consume(ctx context.Context, id string, expect int64) (int64, error) {
 	return r.advance(ctx, `UPDATE sandboxd_archives SET consumed = true, epoch = nextval('sandboxd_archives_epoch'), updated_at = now()
 		WHERE cell = $1 AND id = $2 AND epoch = $3 AND NOT consumed

sandboxd/archives/pg/pg_test.go

@@ -30,6 +30,9 @@ func TestCellsKeepTheirOwnRows(t *testing.T) {
 	if _, err = c2.Take(ctx, "sb_1", epoch, "n2"); !errors.Is(err, archives.ErrLost) {
 		t.Errorf("c2 takes c1's row: %v, want ErrLost", err)
 	}
+	if taken, takeErr := c2.TakeOwned(ctx, "n1", map[string]int64{"sb_1": epoch}); takeErr != nil || len(taken) != 0 {
+		t.Errorf("c2 takes c1's row as its owner: %v %v", taken, takeErr)
+	}
 	if _, err = c2.Consume(ctx, "sb_1", epoch); !errors.Is(err, archives.ErrLost) {
 		t.Errorf("c2 consumes c1's row: %v, want ErrLost", err)
 	}

sandboxd/pool/archive.go

@@ -153,16 +153,22 @@ func (m *Manager) archive(ctx context.Context, sb *types.Sandbox) error {
 		m.mu.Lock()
 		var rb claimSnapshot
 		// A Release that landed meanwhile already deleted the claim; re-setting it would resurrect it
-		if m.claimed[sb.ID] == sb {
+		live := m.claimed[sb.ID] == sb
+		if live {
 			sb.ArchiveCk, sb.ArchiveEpoch, sb.ArchivePending, sb.ReplicaCk = "", 0, prevPending, replica
+			if replica != "" {
+				sb.ArchiveEpoch = epoch
+			}
 			sb.HibernateSnap, sb.VsockSocket, sb.VMName = snap, sock, vmName
 			sb.Deadline = prevDeadline
 			rb = m.store.set(sb)
 		}
 		m.mu.Unlock()
 		sb.Transition.Unlock()
 		m.store.recommit(ctx, rb)
-		m.abandonArchive(ctx, sb, epoch, ck.ID)
+		if !live || replica == "" {
+			m.abandonArchive(ctx, sb, epoch, ck.ID)
+		}
 		return fmt.Errorf("archive %s: persist claims: %w", sb.ID, saveErr)
 	}
 	m.counters.archives.Add(1)

sandboxd/pool/failover.go

@@ -32,12 +32,15 @@ func (m *Manager) adoptArchived(ctx context.Context, id, token string) (bool, er
 	m.mu.Lock()
 	draining := m.draining
 	_, held := m.authed(id, token)
+	fenced := m.adoptFenced()
 	m.mu.Unlock()
 	switch {
 	case draining:
 		return false, nil
 	case held:
 		return true, nil
+	case fenced:
+		return false, nil
 	}
 	ctx = context.WithoutCancel(ctx)
 	rec, ok, err := m.archives.Get(ctx, id)
@@ -557,7 +560,7 @@ func (m *Manager) fenceArchived(ctx context.Context, claims map[string]*types.Sa
 				sb.ReplicaPending = ""
 			}
 		case sb.ArchivePending != "" || ok && !rec.Consumed:
-			moved, err := m.settleIntent(ctx, id, sb.ArchivePending, rec, ok)
+			moved, err := m.settleIntent(ctx, sb, sb.ArchivePending, rec, ok)
 			switch {
 			case err != nil:
 				logger.Warnf(ctx, "archive intent of %s left for the next restart: %v", id, err)
@@ -580,20 +583,26 @@ func (m *Manager) fenceArchived(ctx context.Context, claims map[string]*types.Sa
 	return true
 }
 
-func (m *Manager) settleIntent(ctx context.Context, id, pending string, rec archives.Record, ok bool) (bool, error) {
+func (m *Manager) settleIntent(ctx context.Context, sb *types.Sandbox, pending string, rec archives.Record, ok bool) (moved bool, err error) {
+	m.mu.Lock()
+	replica := sb.ReplicaCk
+	m.mu.Unlock()
+	if ok && rec.Owner == m.nodeID && replica != "" {
+		return false, m.restoreReplicaRow(ctx, sb, replica, rec)
+	}
 	if ok && rec.Owner == m.nodeID {
-		err := m.archives.Delete(ctx, id, rec.Epoch)
+		err = m.archives.Delete(ctx, sb.ID, rec.Epoch)
 		if err == nil {
 			return false, nil
 		}
 		if !errors.Is(err, archives.ErrLost) {
 			return false, err
 		}
-		if rec, ok, err = m.archives.Get(ctx, id); err != nil {
+		if rec, ok, err = m.archives.Get(ctx, sb.ID); err != nil {
 			return false, err
 		}
 		if ok && rec.Owner == m.nodeID {
-			return false, fmt.Errorf("archive row %s moved under its settle", id)
+			return false, fmt.Errorf("archive row %s moved under its settle", sb.ID)
 		}
 	}
 	switch {
@@ -602,13 +611,30 @@ func (m *Manager) settleIntent(ctx context.Context, id, pending string, rec arch
 	case pending == "":
 		return false, nil
 	}
-	_, err := m.ckpts.ReadMeta(ctx, pending)
+	_, err = m.ckpts.ReadMeta(ctx, pending)
 	if errors.Is(err, store.ErrNotFound) {
 		return true, nil
 	}
 	return false, err
 }
 
+func (m *Manager) restoreReplicaRow(ctx context.Context, sb *types.Sandbox, replica string, rec archives.Record) error {
+	m.mu.Lock()
+	sb.ArchiveEpoch = rec.Epoch
+	m.mu.Unlock()
+	if rec.Checkpoint == replica || rec.Consumed {
+		return nil
+	}
+	epoch, err := m.replaceReplica(ctx, sb, replica)
+	if err != nil {
+		return err
+	}
+	m.mu.Lock()
+	sb.ArchiveEpoch = epoch
+	m.mu.Unlock()
+	return nil
+}
+
 func (m *Manager) resolveIntent(ctx context.Context, sb *types.Sandbox) error {
 	if m.archives == nil {
 		return nil
@@ -626,7 +652,7 @@ func (m *Manager) resolveIntent(ctx context.Context, sb *types.Sandbox) error {
 	if err != nil {
 		return err
 	}
-	moved, err := m.settleIntent(ctx, sb.ID, pending, rec, ok)
+	moved, err := m.settleIntent(ctx, sb, pending, rec, ok)
 	switch {
 	case err != nil:
 		return err

sandboxd/pool/lease.go

@@ -20,6 +20,7 @@ import (
 type fenceState struct {
 	ch     chan struct{}
 	fenced bool
+	lifted time.Time
 }
 
 // CloseOnFence closes c when fenced closes; stop ends the watch.
@@ -72,6 +73,11 @@ func (m *Manager) isFenced() bool {
 	return m.fence.Load().fenced
 }
 
+func (m *Manager) adoptFenced() bool {
+	st := m.fence.Load()
+	return st.fenced || time.Since(st.lifted) < adoptGrace
+}
+
 func (m *Manager) trip(ctx context.Context) {
 	for {
 		old := m.fence.Load()
@@ -107,7 +113,15 @@ func (m *Manager) unfence(ctx context.Context, start time.Time) {
 			return
 		}
 	}
+	taken, err := m.takeOwnRows(ctx, held, rows)
+	if err != nil {
+		logger.Warnf(ctx, "replicated claims stay fenced: %v", err)
+		return
+	}
 	for id, sb := range held {
+		if _, ok := taken[id]; ok {
+			continue
+		}
 		keep, err := m.confirmOwnership(ctx, sb, rows[id])
 		if err != nil {
 			logger.Warnf(ctx, "replicated claims stay fenced: %v", err)
@@ -129,7 +143,7 @@ func (m *Manager) unfence(ctx context.Context, start time.Time) {
 
 // lift replaces the fenced state this renewal saw, and trips again when the window closed before the lift landed.
 func (m *Manager) lift(ctx context.Context, fenced *fenceState, start time.Time) {
-	if !m.fence.CompareAndSwap(fenced, &fenceState{ch: make(chan struct{})}) {
+	if !m.fence.CompareAndSwap(fenced, &fenceState{ch: make(chan struct{}), lifted: time.Now()}) {
 		return
 	}
 	if time.Since(start) > leaseTTL-leaseMargin {
@@ -140,6 +154,36 @@ func (m *Manager) lift(ctx context.Context, fenced *fenceState, start time.Time)
 	log.WithFunc("pool.unfence").Info(ctx, "node lease renewed; replicated claims are served again")
 }
 
+func (m *Manager) takeOwnRows(ctx context.Context, held map[string]*types.Sandbox, rows map[string]archives.Record) (map[string]int64, error) {
+	expect := make(map[string]int64, len(rows))
+	for id, rec := range rows {
+		if rec.Owner == m.nodeID && !rec.Consumed {
+			expect[id] = rec.Epoch
+		}
+	}
+	if len(expect) == 0 {
+		return nil, nil
+	}
+	taken, err := m.archives.TakeOwned(ctx, m.nodeID, expect)
+	if err != nil {
+		return nil, err
+	}
+	live := make([]*types.Sandbox, 0, len(taken))
+	m.mu.Lock()
+	for id, epoch := range taken {
+		if sb := held[id]; m.claimed[id] == sb {
+			sb.ArchiveEpoch = max(sb.ArchiveEpoch, epoch)
+			live = append(live, sb)
+		}
+	}
+	js := m.store.set(live...)
+	m.mu.Unlock()
+	if err = m.store.commit(js); err != nil {
+		m.store.recommit(ctx, js)
+	}
+	return taken, nil
+}
+
 // confirmOwnership takes sb's row for this node in one compare-and-set; false means another node holds the claim.
 func (m *Manager) confirmOwnership(ctx context.Context, sb *types.Sandbox, rec archives.Record) (bool, error) {
 	for attempt := range 3 {