state: Make review draft updates atomic

internal/spice/state/review_draft.go

@@ -2,76 +2,97 @@ package state
 
 import (
 	"context"
+	"encoding/json/jsontext"
+	json "encoding/json/v2"
 	"errors"
 	"fmt"
+	"maps"
 	"path"
 	"slices"
 
+	"go.abhg.dev/gs/internal/jsonmut"
 	"go.abhg.dev/gs/internal/review"
 	"go.abhg.dev/gs/internal/spice/state/storage"
 )
 
 const _reviewDraftsDir = "comments"
 
 type reviewDraftState struct {
-	NextID review.DraftID      `json:"nextID"`
-	Drafts []storedReviewDraft `json:"comments"`
+	LastID review.DraftID                       `json:"lastID"`
+	Drafts map[review.DraftID]storedReviewDraft `json:"drafts"`
 }
 
 type storedReviewDraft struct {
-	ID       review.DraftID `json:"id"`
-	File     string         `json:"file"`
-	Line     int            `json:"line"`
-	Body     string         `json:"body"`
-	ThreadID string         `json:"threadID,omitempty"`
+	File     string `json:"file"`
+	Line     int    `json:"line"`
+	Body     string `json:"body"`
+	ThreadID string `json:"threadID,omitempty"`
 }
 
-// AddReviewDraft assigns a branch-local ID and saves a review draft.
+// AddReviewDraft atomically assigns a branch-local ID and saves a draft.
 func (s *Store) AddReviewDraft(
 	ctx context.Context,
 	branch string,
 	draft review.Draft,
 ) (review.Draft, error) {
-	state, err := s.loadReviewDraftState(ctx, branch)
+	stored, err := json.Marshal(storeReviewDraft(draft))
 	if err != nil {
-		return review.Draft{}, err
-	}
-	if state == nil {
-		state = &reviewDraftState{NextID: 1}
+		return review.Draft{}, fmt.Errorf("encode review draft: %w", err)
 	}
-
-	draft.ID = state.NextID
-	state.NextID++
-	state.Drafts = append(state.Drafts, storeReviewDraft(draft))
-	if err := s.saveReviewDraftState(ctx, branch, state); err != nil {
-		return review.Draft{}, err
+	id, err := storage.MutateJSON(
+		ctx,
+		s.db,
+		storage.JSONMutationRequest{
+			Key:       reviewDraftsJSON(branch),
+			IfMissing: jsontext.Value(`{}`),
+			Message:   fmt.Sprintf("%v: add review draft", branch),
+		},
+		jsonmut.InsertAutoIncrement(
+			"/lastID",
+			"/drafts",
+			jsontext.Value(stored),
+		),
+	)
+	if err != nil {
+		return review.Draft{}, fmt.Errorf("add review draft: %w", err)
 	}
+	draft.ID = review.DraftID(id)
 	return draft, nil
 }
 
-// UpdateReviewDraftBody replaces the body of one branch-local draft.
+// UpdateReviewDraftBody atomically replaces one draft's body.
 func (s *Store) UpdateReviewDraftBody(
 	ctx context.Context,
 	branch string,
 	id review.DraftID,
 	body string,
 ) error {
-	state, err := s.loadReviewDraftState(ctx, branch)
+	bodyJSON, err := json.Marshal(body)
 	if err != nil {
-		return err
+		return fmt.Errorf("encode review draft body: %w", err)
 	}
-	if state == nil {
+	err = storage.UpdateJSON(
+		ctx,
+		s.db,
+		storage.JSONMutationRequest{
+			Key:       reviewDraftsJSON(branch),
+			IfMissing: jsontext.Value(`{}`),
+			Message:   fmt.Sprintf("%v: update review draft", branch),
+		},
+		jsonmut.Replace(
+			jsontext.Pointer("/drafts").
+				AppendToken(id.String()).
+				AppendToken("body"),
+			jsontext.Value(bodyJSON),
+		),
+	)
+	if errors.Is(err, jsonmut.ErrNotExist) {
 		return fmt.Errorf("draft comment %d not found", id)
 	}
-
-	idx := slices.IndexFunc(state.Drafts, func(draft storedReviewDraft) bool {
-		return draft.ID == id
-	})
-	if idx < 0 {
-		return fmt.Errorf("draft comment %d not found", id)
+	if err != nil {
+		return fmt.Errorf("update review draft: %w", err)
 	}
-	state.Drafts[idx].Body = body
-	return s.saveReviewDraftState(ctx, branch, state)
+	return nil
 }
 
 // LoadReviewDrafts retrieves the unpublished review comments for branch.
@@ -85,19 +106,21 @@ func (s *Store) LoadReviewDrafts(
 		return nil, err
 	}
 
-	drafts := make([]review.Draft, len(state.Drafts))
-	for i, stored := range state.Drafts {
+	ids := slices.Sorted(maps.Keys(state.Drafts))
+	drafts := make([]review.Draft, len(ids))
+	for i, id := range ids {
+		stored := state.Drafts[id]
 		if stored.ThreadID != "" {
 			drafts[i] = review.Draft{
-				ID:      stored.ID,
+				ID:      id,
 				Body:    stored.Body,
 				ReplyTo: stored.ThreadID,
 			}
 			continue
 		}
 
 		drafts[i] = review.Draft{
-			ID:   stored.ID,
+			ID:   id,
 			Body: stored.Body,
 			Anchor: review.Anchor{
 				Path:      stored.File,
@@ -135,25 +158,8 @@ func (s *Store) loadReviewDraftState(
 	return &state, nil
 }
 
-func (s *Store) saveReviewDraftState(
-	ctx context.Context,
-	branch string,
-	state *reviewDraftState,
-) error {
-	if err := s.db.Set(
-		ctx,
-		reviewDraftsJSON(branch),
-		state,
-		fmt.Sprintf("%v: save review drafts", branch),
-	); err != nil {
-		return fmt.Errorf("set review drafts: %w", err)
-	}
-	return nil
-}
-
 func storeReviewDraft(draft review.Draft) storedReviewDraft {
 	stored := storedReviewDraft{
-		ID:   draft.ID,
 		Body: draft.Body,
 	}
 	if draft.ReplyTo != "" {

internal/spice/state/review_draft_concurrency_test.go

@@ -0,0 +1,217 @@
+package state_test
+
+import (
+	"context"
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.abhg.dev/gs/internal/git"
+	"go.abhg.dev/gs/internal/review"
+	"go.abhg.dev/gs/internal/silog/silogtest"
+	"go.abhg.dev/gs/internal/spice/state"
+	"go.abhg.dev/gs/internal/spice/state/storage"
+)
+
+func TestReviewDraftsConcurrentAdd(t *testing.T) {
+	ctx := t.Context()
+	stores, repo := newConcurrentReviewDraftStores(t)
+
+	added := make([]review.Draft, len(stores))
+	// Hold both initial compare-and-swap attempts until they have derived
+	// mutations from the same storage commit.
+	repo.pauseNextRefUpdates(2)
+	errs := runConcurrently(
+		func() (err error) {
+			added[0], err = stores[0].AddReviewDraft(
+				ctx,
+				"feat",
+				review.Draft{
+					ID:   0,
+					Body: "First",
+					Anchor: review.Anchor{
+						Path:      "first.go",
+						StartLine: 1,
+						EndLine:   1,
+					},
+				},
+			)
+			return err
+		},
+		func() (err error) {
+			added[1], err = stores[1].AddReviewDraft(
+				ctx,
+				"feat",
+				review.Draft{
+					ID:   0,
+					Body: "Second",
+					Anchor: review.Anchor{
+						Path:      "second.go",
+						StartLine: 2,
+						EndLine:   2,
+					},
+				},
+			)
+			return err
+		},
+	)
+	require.NoError(t, errs[0])
+	require.NoError(t, errs[1])
+
+	drafts, err := stores[0].LoadReviewDrafts(ctx, "feat")
+	require.NoError(t, err)
+	require.NotNil(t, drafts)
+	assert.ElementsMatch(
+		t,
+		[]review.DraftID{1, 2},
+		[]review.DraftID{added[0].ID, added[1].ID},
+	)
+	assert.ElementsMatch(t, added, drafts)
+}
+
+func TestReviewDraftsConcurrentEdit(t *testing.T) {
+	ctx := t.Context()
+	stores, repo := newConcurrentReviewDraftStores(t)
+	_, err := stores[0].AddReviewDraft(
+		ctx,
+		"feat",
+		review.Draft{
+			ID:   0,
+			Body: "First",
+			Anchor: review.Anchor{
+				Path:      "first.go",
+				StartLine: 1,
+				EndLine:   1,
+			},
+		},
+	)
+	require.NoError(t, err)
+	_, err = stores[0].AddReviewDraft(
+		ctx,
+		"feat",
+		review.Draft{
+			ID:   0,
+			Body: "Second",
+			Anchor: review.Anchor{
+				Path:      "second.go",
+				StartLine: 2,
+				EndLine:   2,
+			},
+		},
+	)
+	require.NoError(t, err)
+
+	// Hold both initial compare-and-swap attempts until they have derived
+	// mutations from the same storage commit.
+	repo.pauseNextRefUpdates(2)
+	errs := runConcurrently(
+		func() error {
+			return stores[0].UpdateReviewDraftBody(
+				ctx, "feat", 1, "First edited",
+			)
+		},
+		func() error {
+			return stores[1].UpdateReviewDraftBody(
+				ctx, "feat", 2, "Second edited",
+			)
+		},
+	)
+	require.NoError(t, errs[0])
+	require.NoError(t, errs[1])
+
+	drafts, err := stores[0].LoadReviewDrafts(ctx, "feat")
+	require.NoError(t, err)
+	require.NotNil(t, drafts)
+	assert.Equal(t, "First edited", drafts[0].Body)
+	assert.Equal(t, "Second edited", drafts[1].Body)
+}
+
+func newConcurrentReviewDraftStores(
+	t *testing.T,
+) ([2]*state.Store, *pausingGitRepository) {
+	t.Helper()
+	ctx := t.Context()
+	repo, _, err := git.Init(ctx, t.TempDir(), git.InitOptions{
+		Log: silogtest.New(t),
+	})
+	require.NoError(t, err)
+
+	pausingRepo := &pausingGitRepository{GitRepository: repo}
+	newDB := func() *storage.DB {
+		return storage.NewDB(storage.NewGitBackend(storage.GitConfig{
+			Repo:        pausingRepo,
+			Ref:         "refs/data",
+			AuthorName:  "Test Author",
+			AuthorEmail: "test@example.com",
+			Log:         silogtest.New(t),
+		}))
+	}
+	dbs := [2]*storage.DB{newDB(), newDB()}
+	_, err = state.InitStore(ctx, state.InitStoreRequest{
+		DB:    dbs[0],
+		Trunk: "main",
+	})
+	require.NoError(t, err)
+
+	var stores [2]*state.Store
+	for i, db := range dbs {
+		stores[i], err = state.OpenStore(ctx, db, silogtest.New(t))
+		require.NoError(t, err)
+	}
+	return stores, pausingRepo
+}
+
+func runConcurrently(operations ...func() error) []error {
+	var wg sync.WaitGroup
+	errs := make([]error, len(operations))
+	for i, operation := range operations {
+		wg.Go(func() {
+			errs[i] = operation()
+		})
+	}
+	wg.Wait()
+	return errs
+}
+
+type pausingGitRepository struct {
+	storage.GitRepository
+
+	mu      sync.Mutex
+	waitFor int
+	ready   int
+	release chan struct{}
+}
+
+func (r *pausingGitRepository) pauseNextRefUpdates(count int) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.waitFor = count
+	r.ready = 0
+	r.release = make(chan struct{})
+}
+
+func (r *pausingGitRepository) SetRef(
+	ctx context.Context,
+	req git.SetRefRequest,
+) error {
+	r.mu.Lock()
+	wait := r.ready < r.waitFor
+	if wait {
+		r.ready++
+		if r.ready == r.waitFor {
+			close(r.release)
+		}
+	}
+	release := r.release
+	r.mu.Unlock()
+
+	if wait {
+		select {
+		case <-release:
+		case <-ctx.Done():
+			return ctx.Err()
+		}
+	}
+	return r.GitRepository.SetRef(ctx, req)
+}