feat(jobs/runners): worktree prune two-phase + workspace prune queries/repo methods

This commit is contained in:
2026-05-06 07:55:56 +08:00
parent 0111135e78
commit c50621eeca
8 changed files with 638 additions and 12 deletions
+4 -3
View File
@@ -83,9 +83,10 @@ type Worktree struct {
Path string
Status WorktreeStatus
ActiveHolder string // "user:<uid>" / "session:<sid>" / 空
AcquiredAt *time.Time
LastUsedAt time.Time
CreatedAt time.Time
AcquiredAt *time.Time
LastUsedAt time.Time
PruneWarningAt *time.Time
CreatedAt time.Time
}
// Credential 是 git_credentials 表的领域投影。Secret 仅在创建/更新时有值,
+32
View File
@@ -49,3 +49,35 @@ RETURNING *;
-- name: DeleteWorktree :exec
DELETE FROM workspace_worktrees WHERE id = $1;
-- name: ListWorktreesNeedingPruneWarning :many
SELECT * FROM workspace_worktrees
WHERE status = 'idle'
AND prune_warning_at IS NULL
AND last_used_at < $1
ORDER BY last_used_at ASC;
-- name: ListWorktreesNeedingPrune :many
SELECT * FROM workspace_worktrees
WHERE status = 'idle'
AND last_used_at < $1
ORDER BY last_used_at ASC;
-- name: MarkPruneWarning :exec
UPDATE workspace_worktrees
SET prune_warning_at = now()
WHERE id = $1 AND prune_warning_at IS NULL;
-- name: TryStartPrune :one
UPDATE workspace_worktrees
SET status = 'pruning'
WHERE id = $1 AND status = 'idle'
RETURNING *;
-- name: FinishPruneSuccess :exec
DELETE FROM workspace_worktrees WHERE id = $1;
-- name: FinishPruneRollback :exec
UPDATE workspace_worktrees
SET status = 'idle'
WHERE id = $1 AND status = 'pruning';
+74 -9
View File
@@ -50,6 +50,12 @@ type Repository interface {
GetWorktreeByID(ctx context.Context, id uuid.UUID) (*Worktree, error)
ListWorktreesByWorkspace(ctx context.Context, wsID uuid.UUID) ([]*Worktree, error)
DeleteWorktree(ctx context.Context, id uuid.UUID) error
ListWorktreesNeedingPruneWarning(ctx context.Context, lastUsedBefore time.Time) ([]*Worktree, error)
ListWorktreesNeedingPrune(ctx context.Context, lastUsedBefore time.Time) ([]*Worktree, error)
MarkPruneWarning(ctx context.Context, id uuid.UUID) error
TryStartPrune(ctx context.Context, id uuid.UUID) (*Worktree, error)
FinishPruneSuccess(ctx context.Context, id uuid.UUID) error
FinishPruneRollback(ctx context.Context, id uuid.UUID) error
// Credential
UpsertCredential(ctx context.Context, c *Credential) (*Credential, error)
@@ -240,15 +246,16 @@ func (r *pgRepo) ResetStuckSyncStatuses(ctx context.Context) error {
func rowToWorktree(r workspacesqlc.WorkspaceWorktree) *Worktree {
return &Worktree{
ID: fromPgUUID(r.ID),
WorkspaceID: fromPgUUID(r.WorkspaceID),
Branch: r.Branch,
Path: r.Path,
Status: WorktreeStatus(r.Status),
ActiveHolder: fromPtrString(r.ActiveHolder),
AcquiredAt: fromPgTimePtr(r.AcquiredAt),
LastUsedAt: r.LastUsedAt.Time,
CreatedAt: r.CreatedAt.Time,
ID: fromPgUUID(r.ID),
WorkspaceID: fromPgUUID(r.WorkspaceID),
Branch: r.Branch,
Path: r.Path,
Status: WorktreeStatus(r.Status),
ActiveHolder: fromPtrString(r.ActiveHolder),
AcquiredAt: fromPgTimePtr(r.AcquiredAt),
LastUsedAt: r.LastUsedAt.Time,
PruneWarningAt: fromPgTimePtr(r.PruneWarningAt),
CreatedAt: r.CreatedAt.Time,
}
}
@@ -283,6 +290,64 @@ func (r *pgRepo) DeleteWorktree(ctx context.Context, id uuid.UUID) error {
return nil
}
func (r *pgRepo) ListWorktreesNeedingPruneWarning(ctx context.Context, lastUsedBefore time.Time) ([]*Worktree, error) {
ts := pgtype.Timestamptz{Time: lastUsedBefore, Valid: true}
rows, err := r.q.ListWorktreesNeedingPruneWarning(ctx, ts)
if err != nil {
return nil, errs.Wrap(err, errs.CodeInternal, "list worktrees needing prune warning")
}
out := make([]*Worktree, 0, len(rows))
for _, row := range rows {
out = append(out, rowToWorktree(row))
}
return out, nil
}
func (r *pgRepo) ListWorktreesNeedingPrune(ctx context.Context, lastUsedBefore time.Time) ([]*Worktree, error) {
ts := pgtype.Timestamptz{Time: lastUsedBefore, Valid: true}
rows, err := r.q.ListWorktreesNeedingPrune(ctx, ts)
if err != nil {
return nil, errs.Wrap(err, errs.CodeInternal, "list worktrees needing prune")
}
out := make([]*Worktree, 0, len(rows))
for _, row := range rows {
out = append(out, rowToWorktree(row))
}
return out, nil
}
func (r *pgRepo) MarkPruneWarning(ctx context.Context, id uuid.UUID) error {
if err := r.q.MarkPruneWarning(ctx, toPgUUID(id)); err != nil {
return errs.Wrap(err, errs.CodeInternal, "mark prune warning")
}
return nil
}
func (r *pgRepo) TryStartPrune(ctx context.Context, id uuid.UUID) (*Worktree, error) {
row, err := r.q.TryStartPrune(ctx, toPgUUID(id))
if errors.Is(err, pgx.ErrNoRows) {
return nil, errs.New(errs.CodeWorktreeAlreadyActive, "worktree no longer idle")
}
if err != nil {
return nil, errs.Wrap(err, errs.CodeInternal, "try start prune")
}
return rowToWorktree(row), nil
}
func (r *pgRepo) FinishPruneSuccess(ctx context.Context, id uuid.UUID) error {
if err := r.q.FinishPruneSuccess(ctx, toPgUUID(id)); err != nil {
return errs.Wrap(err, errs.CodeInternal, "finish prune success")
}
return nil
}
func (r *pgRepo) FinishPruneRollback(ctx context.Context, id uuid.UUID) error {
if err := r.q.FinishPruneRollback(ctx, toPgUUID(id)); err != nil {
return errs.Wrap(err, errs.CodeInternal, "finish prune rollback")
}
return nil
}
// ===== Credential =====
func rowToCredential(r workspacesqlc.GitCredential) *Credential {
+81
View File
@@ -13,6 +13,7 @@ import (
tcpg "github.com/testcontainers/testcontainers-go/modules/postgres"
"github.com/yan1h/agent-coding-workflow/internal/infra/db"
"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
)
// setupRepo 起一个临时 postgres 容器、跑迁移、注入一条 user/project 行做外键,
@@ -186,3 +187,83 @@ func TestRepo_SetWorktreeActive_ClearsPruneWarning(t *testing.T) {
require.NoError(t, pool.QueryRow(ctx, `SELECT prune_warning_at FROM workspace_worktrees WHERE id = $1`, wtID).Scan(&pwAt))
require.Nil(t, pwAt, "SetWorktreeIdle should clear prune_warning_at")
}
// TestRepo_TryStartPrune_CASLossReturnsAlreadyActive 验证:对一个已变为 active 的
// worktree 调用 TryStartPrune 会返回 CodeWorktreeAlreadyActive。
func TestRepo_TryStartPrune_CASLossReturnsAlreadyActive(t *testing.T) {
t.Parallel()
repo, pool, pid := setupRepo(t)
ctx := context.Background()
wsID := uuid.New()
_, err := repo.CreateWorkspace(ctx, &Workspace{
ID: wsID, ProjectID: pid, Slug: "ws-cas", Name: "CAS",
GitRemoteURL: "https://x", DefaultBranch: "main", MainPath: "/x", SyncStatus: SyncStatusIdle,
})
require.NoError(t, err)
wtID := uuid.New()
require.NoError(t, repo.InTx(ctx, func(tx Tx) error {
_, e := tx.InsertWorktree(ctx, &Worktree{
ID: wtID, WorkspaceID: wsID, Branch: "cas-branch",
Path: "/x/.worktrees/cas", Status: WorktreeStatusIdle,
})
return e
}))
// 绕过 service 直接把 status 改成 active,模拟 CAS 竞争丢失。
_, err = pool.Exec(ctx, `UPDATE workspace_worktrees SET status = 'active' WHERE id = $1`, wtID)
require.NoError(t, err)
_, err = repo.TryStartPrune(ctx, wtID)
require.Error(t, err)
appErr, ok := errs.As(err)
require.True(t, ok, "expected AppError")
require.Equal(t, errs.CodeWorktreeAlreadyActive, appErr.Code)
}
// TestRepo_ListWorktreesNeedingPruneWarning_FiltersExisting 验证:已有
// prune_warning_at 的 worktree 不会被返回,仅返回尚未标记的。
func TestRepo_ListWorktreesNeedingPruneWarning_FiltersExisting(t *testing.T) {
t.Parallel()
repo, pool, pid := setupRepo(t)
ctx := context.Background()
wsID := uuid.New()
_, err := repo.CreateWorkspace(ctx, &Workspace{
ID: wsID, ProjectID: pid, Slug: "ws-pwf", Name: "PWF",
GitRemoteURL: "https://x", DefaultBranch: "main", MainPath: "/x", SyncStatus: SyncStatusIdle,
})
require.NoError(t, err)
cutoff := time.Now().Add(-24 * time.Hour)
// worktree A:没有 prune_warning_at,last_used_at 在 cutoff 之前 → 应返回。
wtA := uuid.New()
require.NoError(t, repo.InTx(ctx, func(tx Tx) error {
_, e := tx.InsertWorktree(ctx, &Worktree{
ID: wtA, WorkspaceID: wsID, Branch: "prune-a",
Path: "/x/.worktrees/a", Status: WorktreeStatusIdle,
})
return e
}))
_, err = pool.Exec(ctx, `UPDATE workspace_worktrees SET last_used_at = $1 WHERE id = $2`, cutoff.Add(-time.Hour), wtA)
require.NoError(t, err)
// worktree B:已经有 prune_warning_at,last_used_at 同样在 cutoff 前 → 不应返回。
wtB := uuid.New()
require.NoError(t, repo.InTx(ctx, func(tx Tx) error {
_, e := tx.InsertWorktree(ctx, &Worktree{
ID: wtB, WorkspaceID: wsID, Branch: "prune-b",
Path: "/x/.worktrees/b", Status: WorktreeStatusIdle,
})
return e
}))
_, err = pool.Exec(ctx,
`UPDATE workspace_worktrees SET last_used_at = $1, prune_warning_at = now() WHERE id = $2`,
cutoff.Add(-time.Hour), wtB)
require.NoError(t, err)
rows, err := repo.ListWorktreesNeedingPruneWarning(ctx, cutoff)
require.NoError(t, err)
require.Len(t, rows, 1)
require.Equal(t, wtA, rows[0].ID)
}
+133
View File
@@ -63,6 +63,26 @@ func (q *Queries) DeleteWorktree(ctx context.Context, id pgtype.UUID) error {
return err
}
const finishPruneRollback = `-- name: FinishPruneRollback :exec
UPDATE workspace_worktrees
SET status = 'idle'
WHERE id = $1 AND status = 'pruning'
`
func (q *Queries) FinishPruneRollback(ctx context.Context, id pgtype.UUID) error {
_, err := q.db.Exec(ctx, finishPruneRollback, id)
return err
}
const finishPruneSuccess = `-- name: FinishPruneSuccess :exec
DELETE FROM workspace_worktrees WHERE id = $1
`
func (q *Queries) FinishPruneSuccess(ctx context.Context, id pgtype.UUID) error {
_, err := q.db.Exec(ctx, finishPruneSuccess, id)
return err
}
const getWorktreeByBranchForUpdate = `-- name: GetWorktreeByBranchForUpdate :one
SELECT id, workspace_id, branch, path, status, active_holder, acquired_at, last_used_at, created_at, prune_warning_at FROM workspace_worktrees
WHERE workspace_id = $1 AND branch = $2
@@ -174,6 +194,94 @@ func (q *Queries) ListWorktreesByWorkspace(ctx context.Context, workspaceID pgty
return items, nil
}
const listWorktreesNeedingPrune = `-- name: ListWorktreesNeedingPrune :many
SELECT id, workspace_id, branch, path, status, active_holder, acquired_at, last_used_at, created_at, prune_warning_at FROM workspace_worktrees
WHERE status = 'idle'
AND last_used_at < $1
ORDER BY last_used_at ASC
`
func (q *Queries) ListWorktreesNeedingPrune(ctx context.Context, lastUsedAt pgtype.Timestamptz) ([]WorkspaceWorktree, error) {
rows, err := q.db.Query(ctx, listWorktreesNeedingPrune, lastUsedAt)
if err != nil {
return nil, err
}
defer rows.Close()
var items []WorkspaceWorktree
for rows.Next() {
var i WorkspaceWorktree
if err := rows.Scan(
&i.ID,
&i.WorkspaceID,
&i.Branch,
&i.Path,
&i.Status,
&i.ActiveHolder,
&i.AcquiredAt,
&i.LastUsedAt,
&i.CreatedAt,
&i.PruneWarningAt,
); err != nil {
return nil, err
}
items = append(items, i)
}
if err := rows.Err(); err != nil {
return nil, err
}
return items, nil
}
const listWorktreesNeedingPruneWarning = `-- name: ListWorktreesNeedingPruneWarning :many
SELECT id, workspace_id, branch, path, status, active_holder, acquired_at, last_used_at, created_at, prune_warning_at FROM workspace_worktrees
WHERE status = 'idle'
AND prune_warning_at IS NULL
AND last_used_at < $1
ORDER BY last_used_at ASC
`
func (q *Queries) ListWorktreesNeedingPruneWarning(ctx context.Context, lastUsedAt pgtype.Timestamptz) ([]WorkspaceWorktree, error) {
rows, err := q.db.Query(ctx, listWorktreesNeedingPruneWarning, lastUsedAt)
if err != nil {
return nil, err
}
defer rows.Close()
var items []WorkspaceWorktree
for rows.Next() {
var i WorkspaceWorktree
if err := rows.Scan(
&i.ID,
&i.WorkspaceID,
&i.Branch,
&i.Path,
&i.Status,
&i.ActiveHolder,
&i.AcquiredAt,
&i.LastUsedAt,
&i.CreatedAt,
&i.PruneWarningAt,
); err != nil {
return nil, err
}
items = append(items, i)
}
if err := rows.Err(); err != nil {
return nil, err
}
return items, nil
}
const markPruneWarning = `-- name: MarkPruneWarning :exec
UPDATE workspace_worktrees
SET prune_warning_at = now()
WHERE id = $1 AND prune_warning_at IS NULL
`
func (q *Queries) MarkPruneWarning(ctx context.Context, id pgtype.UUID) error {
_, err := q.db.Exec(ctx, markPruneWarning, id)
return err
}
const setWorktreeActive = `-- name: SetWorktreeActive :one
UPDATE workspace_worktrees
SET status = 'active',
@@ -261,3 +369,28 @@ func (q *Queries) SetWorktreePruning(ctx context.Context, id pgtype.UUID) (Works
)
return i, err
}
const tryStartPrune = `-- name: TryStartPrune :one
UPDATE workspace_worktrees
SET status = 'pruning'
WHERE id = $1 AND status = 'idle'
RETURNING id, workspace_id, branch, path, status, active_holder, acquired_at, last_used_at, created_at, prune_warning_at
`
func (q *Queries) TryStartPrune(ctx context.Context, id pgtype.UUID) (WorkspaceWorktree, error) {
row := q.db.QueryRow(ctx, tryStartPrune, 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,
&i.PruneWarningAt,
)
return i, err
}
@@ -86,6 +86,18 @@ func (f *fakeRepo) ListWorktreesByWorkspace(_ context.Context, _ uuid.UUID) ([]*
return nil, nil
}
func (f *fakeRepo) DeleteWorktree(_ context.Context, _ uuid.UUID) error { return nil }
func (f *fakeRepo) ListWorktreesNeedingPruneWarning(_ context.Context, _ time.Time) ([]*Worktree, error) {
return nil, nil
}
func (f *fakeRepo) ListWorktreesNeedingPrune(_ context.Context, _ time.Time) ([]*Worktree, error) {
return nil, nil
}
func (f *fakeRepo) MarkPruneWarning(_ context.Context, _ uuid.UUID) error { return nil }
func (f *fakeRepo) TryStartPrune(_ context.Context, _ uuid.UUID) (*Worktree, error) {
return nil, nil
}
func (f *fakeRepo) FinishPruneSuccess(_ context.Context, _ uuid.UUID) error { return nil }
func (f *fakeRepo) FinishPruneRollback(_ context.Context, _ uuid.UUID) error { return nil }
func (f *fakeRepo) UpsertCredential(_ context.Context, c *Credential) (*Credential, error) {
cp := *c
f.creds[c.WorkspaceID] = &cp