feat(workspace): WorktreeService with acquire/release state machine

This commit is contained in:
2026-05-02 08:03:30 +08:00
parent bd91880898
commit d76ed67666
6 changed files with 325 additions and 0 deletions
+4
View File
@@ -7,6 +7,10 @@ RETURNING *;
-- name: GetWorktreeByID :one
SELECT * FROM workspace_worktrees WHERE id = $1;
-- name: GetWorktreeByIDForUpdate :one
SELECT * FROM workspace_worktrees WHERE id = $1
FOR UPDATE;
-- name: GetWorktreeByBranchForUpdate :one
SELECT * FROM workspace_worktrees
WHERE workspace_id = $1 AND branch = $2
+14
View File
@@ -23,6 +23,9 @@ import (
// 状态转移方法,不直接操作 pgx.Tx。
type Tx interface {
GetWorktreeForUpdate(ctx context.Context, wsID uuid.UUID, branch string) (*Worktree, error)
// GetWorktreeByIDForUpdate 在事务内对 id 行加 FOR UPDATE — Acquire/Release 必须用它,
// GetWorktreeByID 不带行锁,并发时会 TOCTOU。
GetWorktreeByIDForUpdate(ctx context.Context, id uuid.UUID) (*Worktree, error)
GetWorktreeByID(ctx context.Context, id uuid.UUID) (*Worktree, error)
InsertWorktree(ctx context.Context, w *Worktree) (*Worktree, error)
SetWorktreeActive(ctx context.Context, id uuid.UUID, holder string) (*Worktree, error)
@@ -355,6 +358,17 @@ func (t *pgTx) GetWorktreeByID(ctx context.Context, id uuid.UUID) (*Worktree, er
return rowToWorktree(row), nil
}
func (t *pgTx) GetWorktreeByIDForUpdate(ctx context.Context, id uuid.UUID) (*Worktree, error) {
row, err := t.q.GetWorktreeByIDForUpdate(ctx, toPgUUID(id))
if errors.Is(err, pgx.ErrNoRows) {
return nil, errs.New(errs.CodeNotFound, "worktree not found")
}
if err != nil {
return nil, errs.Wrap(err, errs.CodeInternal, "select for update")
}
return rowToWorktree(row), nil
}
func (t *pgTx) InsertWorktree(ctx context.Context, w *Worktree) (*Worktree, error) {
row, err := t.q.CreateWorktree(ctx, workspacesqlc.CreateWorktreeParams{
ID: toPgUUID(w.ID),
+22
View File
@@ -111,6 +111,28 @@ func (q *Queries) GetWorktreeByID(ctx context.Context, id pgtype.UUID) (Workspac
return i, err
}
const getWorktreeByIDForUpdate = `-- name: GetWorktreeByIDForUpdate :one
SELECT id, workspace_id, branch, path, status, active_holder, acquired_at, last_used_at, created_at FROM workspace_worktrees WHERE id = $1
FOR UPDATE
`
func (q *Queries) GetWorktreeByIDForUpdate(ctx context.Context, id pgtype.UUID) (WorkspaceWorktree, error) {
row := q.db.QueryRow(ctx, getWorktreeByIDForUpdate, id)
var i WorkspaceWorktree
err := row.Scan(
&i.ID,
&i.WorkspaceID,
&i.Branch,
&i.Path,
&i.Status,
&i.ActiveHolder,
&i.AcquiredAt,
&i.LastUsedAt,
&i.CreatedAt,
)
return i, err
}
const listWorktreesByWorkspace = `-- name: ListWorktreesByWorkspace :many
SELECT id, workspace_id, branch, path, status, active_holder, acquired_at, last_used_at, created_at FROM workspace_worktrees
WHERE workspace_id = $1
@@ -118,6 +118,13 @@ func (t *fakeTx) GetWorktreeByID(_ context.Context, id uuid.UUID) (*Worktree, er
}
return nil, &errsNotFound{}
}
func (t *fakeTx) GetWorktreeByIDForUpdate(_ context.Context, id uuid.UUID) (*Worktree, error) {
if w, ok := t.f.wts[id]; ok {
cp := *w
return &cp, nil
}
return nil, &errsNotFound{}
}
func (t *fakeTx) InsertWorktree(_ context.Context, w *Worktree) (*Worktree, error) {
cp := *w
t.f.wts[w.ID] = &cp
+195
View File
@@ -0,0 +1,195 @@
package workspace
import (
"context"
"errors"
"time"
"github.com/google/uuid"
"github.com/yan1h/agent-coding-workflow/internal/audit"
"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
"github.com/yan1h/agent-coding-workflow/internal/infra/git"
)
type worktreeService struct {
repo Repository
rec audit.Recorder
gitr git.Runner
pa ProjectAccess
dataDir string
}
// NewWorktreeService 构造 WorktreeService。
func NewWorktreeService(repo Repository, rec audit.Recorder, gitr git.Runner, pa ProjectAccess, dataDir string) WorktreeService {
return &worktreeService{repo: repo, rec: rec, gitr: gitr, pa: pa, dataDir: dataDir}
}
func (s *worktreeService) Create(ctx context.Context, c Caller, wsID uuid.UUID, in CreateWorktreeInput) (*Worktree, error) {
ws, err := s.repo.GetWorkspaceByID(ctx, wsID)
if err != nil {
return nil, err
}
if _, canWrite, err := s.pa.ResolveByID(ctx, c.UserID, c.IsAdmin, ws.ProjectID); err != nil {
return nil, err
} else if !canWrite {
return nil, errs.New(errs.CodeForbidden, "no write access")
}
if err := ValidateBranch(in.Branch); err != nil {
return nil, errs.Wrap(err, errs.CodeInvalidInput, "invalid branch")
}
base := in.BaseBranch
if base == "" {
base = ws.DefaultBranch
}
if err := ValidateBranch(base); err != nil {
return nil, errs.Wrap(err, errs.CodeInvalidInput, "invalid base branch")
}
wtPath := WorktreePath(s.dataDir, wsID, in.Branch)
if err := s.gitr.WorktreeAdd(ctx, ws.MainPath, in.Branch, wtPath); err != nil {
if git.IsConflict(err) {
return nil, errs.Wrap(err, errs.CodeWorktreeBranchConflict, "branch already used")
}
return nil, errs.Wrap(err, errs.CodeGitCmdFailed, "git worktree add")
}
wt := &Worktree{
ID: uuid.New(), WorkspaceID: wsID, Branch: in.Branch, Path: wtPath,
Status: WorktreeStatusIdle,
}
var saved *Worktree
err = s.repo.InTx(ctx, func(tx Tx) error {
w, e := tx.InsertWorktree(ctx, wt)
if e != nil {
return e
}
saved = w
return nil
})
if err != nil {
// 回滚磁盘
_ = s.gitr.WorktreeRemove(ctx, ws.MainPath, wtPath)
return nil, err
}
s.audit(ctx, &c.UserID, "worktree.create", "worktree", saved.ID.String(), map[string]any{"branch": in.Branch})
return saved, nil
}
func (s *worktreeService) List(ctx context.Context, c Caller, wsID uuid.UUID) ([]*Worktree, error) {
ws, err := s.repo.GetWorkspaceByID(ctx, wsID)
if err != nil {
return nil, err
}
if canRead, _, err := s.pa.ResolveByID(ctx, c.UserID, c.IsAdmin, ws.ProjectID); err != nil {
return nil, err
} else if !canRead {
return nil, errs.New(errs.CodeForbidden, "no read access")
}
return s.repo.ListWorktreesByWorkspace(ctx, wsID)
}
func (s *worktreeService) Delete(ctx context.Context, c Caller, wtID uuid.UUID) error {
wt, err := s.repo.GetWorktreeByID(ctx, wtID)
if err != nil {
return err
}
ws, err := s.repo.GetWorkspaceByID(ctx, wt.WorkspaceID)
if err != nil {
return err
}
if _, canWrite, err := s.pa.ResolveByID(ctx, c.UserID, c.IsAdmin, ws.ProjectID); err != nil {
return err
} else if !canWrite {
return errs.New(errs.CodeForbidden, "no write access")
}
if wt.Status == WorktreeStatusActive {
return errs.New(errs.CodeWorktreeAlreadyActive, "release before delete")
}
// 先 set pruning,再调 git rm,然后 DELETE row。中途失败 row 留 pruning。
err = s.repo.InTx(ctx, func(tx Tx) error {
_, e := tx.SetWorktreePruning(ctx, wtID)
return e
})
if err != nil {
return err
}
if rmErr := s.gitr.WorktreeRemove(ctx, ws.MainPath, wt.Path); rmErr != nil {
return errs.Wrap(rmErr, errs.CodeGitCmdFailed, "git worktree remove")
}
if err := s.repo.DeleteWorktree(ctx, wtID); err != nil {
return err
}
s.audit(ctx, &c.UserID, "worktree.delete", "worktree", wtID.String(), nil)
return nil
}
func (s *worktreeService) Acquire(ctx context.Context, c Caller, wtID uuid.UUID) (*Worktree, error) {
holder := c.Holder()
var out *Worktree
err := s.repo.InTx(ctx, func(tx Tx) error {
wt, err := tx.GetWorktreeByIDForUpdate(ctx, wtID)
if err != nil {
return err
}
switch wt.Status {
case WorktreeStatusPruning:
return errs.New(errs.CodeNotFound, "worktree being pruned")
case WorktreeStatusActive:
if wt.ActiveHolder == holder {
out = wt // idempotent
return nil
}
return errs.New(errs.CodeWorktreeAlreadyActive, "held by another holder")
case WorktreeStatusIdle:
updated, err := tx.SetWorktreeActive(ctx, wtID, holder)
if err != nil {
return err
}
out = updated
return nil
}
return errs.New(errs.CodeInternal, "unknown worktree status")
})
if err != nil {
return nil, err
}
s.audit(ctx, &c.UserID, "worktree.acquire", "worktree", wtID.String(), map[string]any{"holder": holder})
return out, nil
}
func (s *worktreeService) Release(ctx context.Context, c Caller, wtID uuid.UUID) (*Worktree, error) {
holder := c.Holder()
var out *Worktree
err := s.repo.InTx(ctx, func(tx Tx) error {
wt, err := tx.GetWorktreeByIDForUpdate(ctx, wtID)
if err != nil {
return err
}
if wt.Status != WorktreeStatusActive {
return errs.New(errs.CodeWorktreeNotHeldByCaller, "worktree not active")
}
if wt.ActiveHolder != holder && !c.IsAdmin {
return errs.New(errs.CodeWorktreeNotHeldByCaller, "not held by caller")
}
updated, err := tx.SetWorktreeIdle(ctx, wtID)
if err != nil {
return err
}
out = updated
return nil
})
if err != nil {
return nil, err
}
s.audit(ctx, &c.UserID, "worktree.release", "worktree", wtID.String(), map[string]any{"holder": holder})
return out, nil
}
func (s *worktreeService) audit(ctx context.Context, uid *uuid.UUID, action, ttype, tid string, meta map[string]any) {
_ = s.rec.Record(ctx, audit.Entry{
UserID: uid, Action: action, TargetType: ttype, TargetID: tid, Metadata: meta,
})
}
// 防止未使用 import lint 移除
var _ = errors.Is
var _ = time.Now
@@ -0,0 +1,83 @@
package workspace
import (
"context"
"testing"
"github.com/google/uuid"
"github.com/stretchr/testify/require"
)
func newWtSvc(t *testing.T) (*worktreeService, *fakeRepo, *fakePA, *fakeGit, uuid.UUID) {
t.Helper()
pa := &fakePA{pid: uuid.New(), canRead: true, canWrite: true}
gitr := &fakeGit{}
repo := newFakeRepo()
svc := NewWorktreeService(repo, noopAudit{}, gitr, pa, "/data").(*worktreeService)
wsID := uuid.New()
repo.wss[wsID] = &Workspace{ID: wsID, ProjectID: pa.pid, MainPath: "/data/main", DefaultBranch: "main"}
return svc, repo, pa, gitr, wsID
}
func TestWorktreeService_CreateAndList(t *testing.T) {
t.Parallel()
svc, _, _, _, wsID := newWtSvc(t)
wt, err := svc.Create(context.Background(), Caller{UserID: uuid.New()}, wsID, CreateWorktreeInput{Branch: "feat-x"})
require.NoError(t, err)
require.Equal(t, "feat-x", wt.Branch)
require.Equal(t, WorktreeStatusIdle, wt.Status)
}
func TestWorktreeService_AcquireIdempotent(t *testing.T) {
t.Parallel()
svc, _, _, _, wsID := newWtSvc(t)
caller := Caller{UserID: uuid.New()}
wt, err := svc.Create(context.Background(), caller, wsID, CreateWorktreeInput{Branch: "feat-y"})
require.NoError(t, err)
a1, err := svc.Acquire(context.Background(), caller, wt.ID)
require.NoError(t, err)
require.Equal(t, WorktreeStatusActive, a1.Status)
a2, err := svc.Acquire(context.Background(), caller, wt.ID) // idempotent
require.NoError(t, err)
require.Equal(t, a1.ID, a2.ID)
}
func TestWorktreeService_AcquireConflict(t *testing.T) {
t.Parallel()
svc, _, _, _, wsID := newWtSvc(t)
a := Caller{UserID: uuid.New()}
b := Caller{UserID: uuid.New()}
wt, err := svc.Create(context.Background(), a, wsID, CreateWorktreeInput{Branch: "feat-z"})
require.NoError(t, err)
_, err = svc.Acquire(context.Background(), a, wt.ID)
require.NoError(t, err)
_, err = svc.Acquire(context.Background(), b, wt.ID)
require.Error(t, err)
}
func TestWorktreeService_Release(t *testing.T) {
t.Parallel()
svc, _, _, _, wsID := newWtSvc(t)
caller := Caller{UserID: uuid.New()}
wt, err := svc.Create(context.Background(), caller, wsID, CreateWorktreeInput{Branch: "feat-r"})
require.NoError(t, err)
_, err = svc.Acquire(context.Background(), caller, wt.ID)
require.NoError(t, err)
_, err = svc.Release(context.Background(), caller, wt.ID)
require.NoError(t, err)
got, _ := svc.repo.GetWorktreeByID(context.Background(), wt.ID)
require.Equal(t, WorktreeStatusIdle, got.Status)
}
func TestWorktreeService_Delete_BlockedWhenActive(t *testing.T) {
t.Parallel()
svc, _, _, _, wsID := newWtSvc(t)
caller := Caller{UserID: uuid.New()}
wt, err := svc.Create(context.Background(), caller, wsID, CreateWorktreeInput{Branch: "feat-d"})
require.NoError(t, err)
_, err = svc.Acquire(context.Background(), caller, wt.ID)
require.NoError(t, err)
err = svc.Delete(context.Background(), caller, wt.ID)
require.Error(t, err)
}