pool, store: hold a template's record lock only for the re-promote commit

sandboxd/pool/promote_test.go

@@ -2,6 +2,7 @@ package pool
 
 import (
 	"bytes"
+	"context"
 	"errors"
 	"maps"
 	"os"
@@ -12,6 +13,7 @@ import (
 
 	"github.com/cocoonstack/sandbox/sandboxd/config"
 	"github.com/cocoonstack/sandbox/sandboxd/store"
+	"github.com/cocoonstack/sandbox/sandboxd/store/storetest"
 	"github.com/cocoonstack/sandbox/sandboxd/tenants"
 	"github.com/cocoonstack/sandbox/sandboxd/types"
 )
@@ -478,7 +480,7 @@ func TestTemplatesListWhatTheNodeHoldsAcrossARestart(t *testing.T) {
 	if err = os.WriteFile(filepath.Join(staging, store.MetaFile), []byte(`{"id":"`+oldID+`","created_at":"2026-01-01T00:00:00Z"}`), 0o600); err != nil {
 		t.Fatalf("write keyless meta: %v", err)
 	}
-	if _, err = m.tpls.PublishDigested(t.Context(), staging, oldID, ""); err != nil {
+	if _, err = storetest.PublishDigested(t.Context(), m.tpls, staging, oldID); err != nil {
 		t.Fatalf("publish keyless record: %v", err)
 	}
 
@@ -594,6 +596,43 @@ func TestRejectedLabelWritesLeakNoRecLock(t *testing.T) {
 	}
 }
 
+func TestClaimFromATemplateRunsBesideItsRepromote(t *testing.T) {
+	m := newTestManager(t, newFakeEngine())
+	parent := mustClaim(t, m, testKey)
+	key, _, err := m.Promote(t.Context(), parent.ID, Cred{Token: parent.Token}, "tpl:x", "")
+	if err != nil {
+		t.Fatalf("Promote: %v", err)
+	}
+	slow := &slowPrepareStore{Store: m.tpls, entered: make(chan struct{}), release: make(chan struct{})}
+	m.tpls = slow
+	repromoted := make(chan error, 1)
+	go func() {
+		_, _, promoteErr := m.Promote(t.Context(), parent.ID, Cred{Token: parent.Token}, "tpl:x", "")
+		repromoted <- promoteErr
+	}()
+	<-slow.entered
+
+	ctx, cancel := context.WithTimeout(t.Context(), 100*time.Millisecond)
+	defer cancel()
+	claimed := make(chan error, 1)
+	go func() {
+		_, claimErr := claimAny(ctx, m, key, 0)
+		claimed <- claimErr
+	}()
+	select {
+	case err = <-claimed:
+		if err != nil {
+			t.Errorf("claim during a re-promote: %v, want a clone of the committed template", err)
+		}
+	case <-time.After(2 * time.Second):
+		t.Error("a claim with a 100 ms context was still blocked 2 s into a re-promote")
+	}
+	close(slow.release)
+	if err = <-repromoted; err != nil {
+		t.Fatalf("re-promote: %v", err)
+	}
+}
+
 func claimTenant(t *testing.T, m *Manager, tenant string) *types.Sandbox {
 	t.Helper()
 	sb, err := m.ClaimProvision(t.Context(), testKey, ClaimOptions{TTL: time.Hour, Tenant: tenant})
@@ -614,3 +653,15 @@ func stageTemplate(t *testing.T, m *Manager, id string) string {
 	}
 	return staging
 }
+
+type slowPrepareStore struct {
+	store.Store
+	entered chan struct{}
+	release chan struct{}
+}
+
+func (s *slowPrepareStore) PrepareDigested(ctx context.Context, staging, id, lineage string) (string, store.CommitFunc, error) {
+	close(s.entered)
+	<-s.release
+	return s.Store.PrepareDigested(ctx, staging, id, lineage)
+}

sandboxd/pool/template.go

@@ -300,10 +300,6 @@ func (m *Manager) checkTemplateOwner(ctx context.Context, id, tenant string) err
 
 func (m *Manager) commitTemplate(ctx context.Context, staging string, rec templateRecord) (string, error) {
 	id := store.TemplateID(rec.Key.Hash())
-	defer m.lockRec(id)()
-	if ownerErr := m.checkTemplateOwner(ctx, id, rec.Tenant); ownerErr != nil {
-		return "", ownerErr
-	}
 	rec.ID, rec.CreatedAt = id, time.Now()
 	meta, err := json.Marshal(rec)
 	if err != nil {
@@ -312,10 +308,17 @@ func (m *Manager) commitTemplate(ctx context.Context, staging string, rec templa
 	if err = os.WriteFile(filepath.Join(staging, store.MetaFile), meta, 0o600); err != nil {
 		return "", err
 	}
-	digest, err := m.tpls.PublishDigested(ctx, staging, id, "")
+	digest, commit, err := m.tpls.PrepareDigested(ctx, staging, id, "")
 	if err != nil {
 		return "", fmt.Errorf("publish template: %w", err)
 	}
+	defer m.lockRec(id)()
+	if ownerErr := m.checkTemplateOwner(ctx, id, rec.Tenant); ownerErr != nil {
+		return "", ownerErr
+	}
+	if err = commit(ctx); err != nil {
+		return "", fmt.Errorf("publish template: %w", err)
+	}
 	if err = m.tpls.SetLabels(ctx, id, nil); err != nil {
 		log.WithFunc("pool.commitTemplate").Warnf(ctx, "clear the previous promote's labels on %s: %v", id, err)
 	}

sandboxd/store/dir/dir.go

@@ -60,15 +60,12 @@ func (d *Store) Publish(ctx context.Context, staging, id, _ string) error {
 	return d.publish(ctx, staging, id, "")
 }
 
-func (d *Store) PublishDigested(ctx context.Context, staging, id, _ string) (string, error) {
+func (d *Store) PrepareDigested(ctx context.Context, staging, id, _ string) (string, store.CommitFunc, error) {
 	digest, err := hashExport(ctx, filepath.Join(staging, store.ExportDir))
 	if err != nil {
-		return "", fmt.Errorf("digest export: %w", err)
+		return "", nil, fmt.Errorf("digest export: %w", err)
 	}
-	if err := d.publish(ctx, staging, id, digest); err != nil {
-		return "", fmt.Errorf("publish digested record: %w", err)
-	}
-	return digest, nil
+	return digest, func(ctx context.Context) error { return d.publish(ctx, staging, id, digest) }, nil
 }
 
 func (d *Store) Fetch(ctx context.Context, id string) (string, []byte, string, error) {

sandboxd/store/dir/dir_test.go

@@ -134,7 +134,7 @@ func TestPublishDigestedChunkBoundaries(t *testing.T) {
 			if err = os.WriteFile(filepath.Join(staging, store.MetaFile), []byte(metaJSON(tt.name)), 0o600); err != nil {
 				t.Fatalf("write meta: %v", err)
 			}
-			digest, err := st.PublishDigested(t.Context(), staging, id, "")
+			digest, err := storetest.PublishDigested(t.Context(), st, staging, id)
 			if err != nil {
 				t.Fatalf("PublishDigested: %v", err)
 			}
@@ -216,7 +216,7 @@ func TestPublishDigestedFailureAndFreshRetry(t *testing.T) {
 	if err = os.Mkdir(blocker, 0o750); err != nil {
 		t.Fatalf("block digest commit: %v", err)
 	}
-	if digest, publishErr := st.PublishDigested(t.Context(), failed, id, ""); publishErr == nil || digest != "" {
+	if digest, publishErr := storetest.PublishDigested(t.Context(), st, failed, id); publishErr == nil || digest != "" {
 		t.Fatalf("PublishDigested with blocked sidecar = %q, %v, want empty/error", digest, publishErr)
 	}
 	dir, meta, digest, err := st.Fetch(t.Context(), id)
@@ -658,7 +658,7 @@ func mustPublishDigested(t *testing.T, st *Store, id, metaID string) string {
 		t.Fatalf("stage %s: %v", metaID, err)
 	}
 	seedRecord(t, staging, metaID)
-	digest, err := st.PublishDigested(t.Context(), staging, id, "")
+	digest, err := storetest.PublishDigested(t.Context(), st, staging, id)
 	if err != nil {
 		t.Fatalf("publish digested %s: %v", metaID, err)
 	}

sandboxd/store/s3/s3.go

@@ -132,12 +132,15 @@ func (s *Store) Stage(id string) (string, error) {
 }
 
 func (s *Store) Publish(ctx context.Context, staging, id, lineage string) error {
-	_, err := s.publish(ctx, staging, id, lineage, false)
-	return err
+	_, commit, err := s.prepare(ctx, staging, id, lineage, false)
+	if err != nil {
+		return err
+	}
+	return commit(ctx)
 }
 
-func (s *Store) PublishDigested(ctx context.Context, staging, id, lineage string) (string, error) {
-	return s.publish(ctx, staging, id, lineage, true)
+func (s *Store) PrepareDigested(ctx context.Context, staging, id, lineage string) (string, store.CommitFunc, error) {
+	return s.prepare(ctx, staging, id, lineage, true)
 }
 
 func (s *Store) Fetch(ctx context.Context, id string) (string, []byte, string, error) {
@@ -325,22 +328,22 @@ func (s *Store) get(ctx context.Context, key string) (*awss3.GetObjectOutput, er
 	return out, nil
 }
 
-func (s *Store) publish(ctx context.Context, staging, id, lineage string, digested bool) (string, error) {
+func (s *Store) prepare(ctx context.Context, staging, id, lineage string, digested bool) (string, store.CommitFunc, error) {
 	ctx, cancel := context.WithTimeout(ctx, transferBudget)
 	defer cancel()
 	lineage = cmp.Or(lineage, id)
 	metaRaw, err := os.ReadFile(filepath.Join(staging, store.MetaFile)) //nolint:gosec // our own staging dir
 	if err != nil {
-		return "", fmt.Errorf("staging has no %s: %w", store.MetaFile, err)
+		return "", nil, fmt.Errorf("staging has no %s: %w", store.MetaFile, err)
 	}
 	gen := store.ExportGen(metaRaw)
 	files, err := publishFiles(staging)
 	if err != nil {
-		return "", err
+		return "", nil, err
 	}
 	have, err := s.referencedChunks(ctx, id, lineage)
 	if err != nil {
-		return "", err
+		return "", nil, err
 	}
 	man := manifest{Lineage: lineage, Files: make([]manifestFile, len(files))}
 	digestFiles := make([]store.DigestFile, len(files))
@@ -360,37 +363,38 @@ func (s *Store) publish(ctx context.Context, staging, id, lineage string, digest
 		err = fileErr
 	}
 	if err != nil {
-		return "", err
+		return "", nil, err
 	}
 	manifestRaw, err := json.Marshal(man)
 	if err != nil {
-		return "", fmt.Errorf("encode manifest: %w", err)
+		return "", nil, fmt.Errorf("encode manifest: %w", err)
 	}
 	if err = s.putObject(ctx, s.manifestKey(id, gen), manifestRaw, nil); err != nil {
-		return "", err
+		return "", nil, err
 	}
 	if lineage != id {
 		if err = s.putObject(ctx, s.lineageKey(lineage, id), nil, nil); err != nil {
-			return "", err
+			return "", nil, err
 		}
 	}
 	digest := ""
 	var metadata map[string]string
 	if digested {
 		digest, err = store.AssembleDigest(digestFiles)
 		if err != nil {
-			return "", err
+			return "", nil, err
 		}
 		metadata = map[string]string{digestMetadataKey: digest}
 	}
-	if err = s.putObject(ctx, s.key(id, store.MetaFile), metaRaw, metadata); err != nil {
-		return "", err
-	}
-	// Keep old generations: another node may have selected the previous meta.
-	if err := os.RemoveAll(staging); err != nil {
-		return "", err
-	}
-	return digest, nil
+	return digest, func(ctx context.Context) error {
+		ctx, cancel := context.WithTimeout(ctx, transferBudget)
+		defer cancel()
+		if err := s.putObject(ctx, s.key(id, store.MetaFile), metaRaw, metadata); err != nil {
+			return err
+		}
+		// Keep old generations: another node may have selected the previous meta.
+		return os.RemoveAll(staging)
+	}, nil
 }
 
 // populate downloads one cache generation and installs it atomically.

sandboxd/store/s3/s3_test.go

@@ -192,7 +192,7 @@ func TestPublishDigestedRetryAndMetaRequestAccounting(t *testing.T) {
 	fake.setPutFailures(exportKey, 1)
 
 	staging := stageRecord(t, st, id, meta, content)
-	digest, err := st.PublishDigested(t.Context(), staging, id, "")
+	digest, err := storetest.PublishDigested(t.Context(), st, staging, id)
 	if err != nil {
 		t.Fatalf("PublishDigested: %v", err)
 	}
@@ -254,7 +254,7 @@ func TestPublishDigestedFailurePreservesCommittedGeneration(t *testing.T) {
 	st := newTestStore(t, fake)
 	oldMeta := []byte(`{"id":"` + id + `","gen":1}`)
 	oldStaging := stageRecord(t, st, id, oldMeta, "old")
-	oldDigest, err := st.PublishDigested(t.Context(), oldStaging, id, "")
+	oldDigest, err := storetest.PublishDigested(t.Context(), st, oldStaging, id)
 	if err != nil {
 		t.Fatalf("publish old generation: %v", err)
 	}
@@ -266,7 +266,7 @@ func TestPublishDigestedFailurePreservesCommittedGeneration(t *testing.T) {
 	metaKey := "ck/" + id + "/" + store.MetaFile
 	markerPuts := fake.requestCount(http.MethodPut, metaKey)
 	fake.setPutFailures(newExportKey, 100)
-	digest, err := st.PublishDigested(t.Context(), newStaging, id, "")
+	digest, err := storetest.PublishDigested(t.Context(), st, newStaging, id)
 	if err == nil {
 		t.Fatal("PublishDigested succeeded with injected export failure")
 	}
@@ -290,7 +290,7 @@ func TestPublishDigestedFailurePreservesCommittedGeneration(t *testing.T) {
 	}
 
 	fake.setPutFailures(newExportKey, 0)
-	newDigest, err := st.PublishDigested(t.Context(), newStaging, id, "")
+	newDigest, err := storetest.PublishDigested(t.Context(), st, newStaging, id)
 	if err != nil {
 		t.Fatalf("retry PublishDigested: %v", err)
 	}
@@ -534,7 +534,7 @@ func TestPublishDigestedHashesHolesLikeZeros(t *testing.T) {
 	size := 3*store.DigestChunkSize + 7
 	staging, src := stageSparseRecord(t, st, id, []byte(`{"id":"`+id+`"}`), size, map[int64]string{store.DigestChunkSize + 5: "data", size - 1: "z"})
 	sums := denseChunkSums(t, src)
-	digest, err := st.PublishDigested(t.Context(), staging, id, "")
+	digest, err := storetest.PublishDigested(t.Context(), st, staging, id)
 	if err != nil {
 		t.Fatalf("PublishDigested: %v", err)
 	}
@@ -552,7 +552,7 @@ func TestPublishDigestedSplitsChunksAtTheFixedBoundary(t *testing.T) {
 	st := newTestStore(t, &fakeS3{objects: map[string][]byte{}})
 	content := string(bytes.Repeat([]byte{'x'}, int(store.DigestChunkSize))) + "y"
 	staging := stageRecord(t, st, id, []byte(`{"id":"`+id+`"}`), content)
-	digest, err := st.PublishDigested(t.Context(), staging, id, "")
+	digest, err := storetest.PublishDigested(t.Context(), st, staging, id)
 	if err != nil {
 		t.Fatalf("PublishDigested: %v", err)
 	}

sandboxd/store/store.go

@@ -56,14 +56,16 @@ type S3Config struct {
 	UnsignedAcceptEncoding bool   `json:"unsigned_accept_encoding,omitzero"`
 }
 
+type CommitFunc func(ctx context.Context) error
+
 // Store is one record backend.
 type Store interface {
 	// Stage returns a writable staging directory whose Publish is atomic.
 	Stage(id string) (string, error)
 	// Publish atomically replaces the record's previous generation with the staged directory; lineage scopes chunk dedup in a content-addressed backend and "" means the record id.
 	Publish(ctx context.Context, staging, id, lineage string) error
-	// PublishDigested applies Publish semantics and returns the export digest.
-	PublishDigested(ctx context.Context, staging, id, lineage string) (string, error)
+	// PrepareDigested runs a Publish up to the step that makes the generation current, and returns the export digest and that step.
+	PrepareDigested(ctx context.Context, staging, id, lineage string) (string, CommitFunc, error)
 	// Fetch materializes a record's export locally; the caller holds the record lock until the clone ends.
 	Fetch(ctx context.Context, id string) (dir string, meta []byte, digest string, err error)
 	// FetchFile writes one file of a record's export to dst, without fetching the rest of the record; dst is the caller's to delete.