Agent-System-Backend/internal/store/lease_pg_test.go

package store

import (
	"errors"
	"testing"
	"time"
)

// 协作存储 U5（KTD4）：会话执行租约与 epoch fencing。
// 双档（PG 由 CI 真库 job 覆盖咨询锁与跨实例语义；SQLite 验证条件逻辑）。

func TestExecutionClaimAndContention(t *testing.T) {
	st := openTestStore(t)
	// 首次占用：epoch=1。
	e1, err := st.TryClaimExecution("cnv_1", "instA", "run_1", 60)
	if err != nil || e1 != 1 {
		t.Fatalf("claim A: epoch=%d err=%v", e1, err)
	}
	// 他实例争抢：租约未到期 → ErrExecutionBusy。
	_, err = st.TryClaimExecution("cnv_1", "instB", "run_2", 60)
	if !errors.Is(err, ErrExecutionBusy) {
		t.Fatalf("B should be busy: %v", err)
	}
	// 同实例重占（新回合）：允许，epoch 递增。
	e2, err := st.TryClaimExecution("cnv_1", "instA", "run_3", 60)
	if err != nil || e2 != 2 {
		t.Fatalf("re-claim A: epoch=%d err=%v", e2, err)
	}
	// 续租 fencing：正确 (instance,epoch) 成功；旧 epoch 失败。
	if !st.RenewExecution("cnv_1", "instA", e2, 60) {
		t.Fatal("renew with current epoch should succeed")
	}
	if st.RenewExecution("cnv_1", "instA", e1, 60) {
		t.Fatal("renew with stale epoch must fail")
	}
	// 主动释放后他实例可占，epoch 继续递增（单调 fencing）。
	st.ReleaseExecution("cnv_1", "instA", e2)
	e3, err := st.TryClaimExecution("cnv_1", "instB", "run_4", 60)
	if err != nil || e3 != 3 {
		t.Fatalf("takeover B: epoch=%d err=%v", e3, err)
	}
	// fence 校验：B 当前 OK；A 的旧代次失败。
	if !st.ExecutionFenceOK("cnv_1", "instB", e3) {
		t.Fatal("B fence should be ok")
	}
	if st.ExecutionFenceOK("cnv_1", "instA", e2) {
		t.Fatal("A stale fence must fail")
	}
}

func TestExecutionLeaseExpiryTakeover(t *testing.T) {
	st := openTestStore(t)
	// 1 秒租约到期后他实例接管。
	if _, err := st.TryClaimExecution("cnv_2", "instA", "run_1", 1); err != nil {
		t.Fatal(err)
	}
	time.Sleep(1100 * time.Millisecond)
	e2, err := st.TryClaimExecution("cnv_2", "instB", "run_2", 60)
	if err != nil || e2 != 2 {
		t.Fatalf("expired takeover: epoch=%d err=%v", e2, err)
	}
	// 旧 owner 续租被拒（fencing）。
	if st.RenewExecution("cnv_2", "instA", 1, 60) {
		t.Fatal("stale renew must fail after takeover")
	}
}

// 命令收件箱（KTD5）：幂等投递、原子认领、并发认领不重复。
func TestRunCommandsClaim(t *testing.T) {
	st := openTestStore(t)
	if err := st.EnqueueCommand("cmd_1", "cancel", "run_9", "cnv_9", 0, map[string]string{"owner": "o"}); err != nil {
		t.Fatal(err)
	}
	// 幂等：重复投递不新增。
	if err := st.EnqueueCommand("cmd_1", "cancel", "run_9", "cnv_9", 0, nil); err != nil {
		t.Fatal(err)
	}
	got, err := st.ClaimPendingCommands("instA", 8)
	if err != nil || len(got) != 1 || got[0].CommandID != "cmd_1" || got[0].Kind != "cancel" {
		t.Fatalf("claim: %v err=%v", got, err)
	}
	state, consumer, err := st.CommandState("cmd_1")
	if err != nil || state != "consumed" || consumer != "instA" {
		t.Fatalf("state=%s consumer=%s err=%v", state, consumer, err)
	}
	// 已消费不再返回。
	got, err = st.ClaimPendingCommands("instB", 8)
	if err != nil || len(got) != 0 {
		t.Fatalf("consumed must not re-deliver: %v err=%v", got, err)
	}
}

func TestRunCommandsConcurrentSingleClaim(t *testing.T) {
	st := openTestStore(t)
	for i := 0; i < 8; i++ {
		if err := st.EnqueueCommand(newID("cmd"), "cancel", newID("run"), "cnv_c", 0, nil); err != nil {
			t.Fatal(err)
		}
	}
	var mu = make(chan int, 2)
	total := 0
	done := make(chan struct{})
	for w := 0; w < 2; w++ {
		go func(w int) {
			got, err := st.ClaimPendingCommands("inst"+string(rune('A'+w)), 8)
			if err == nil {
				mu <- len(got)
			}
			done <- struct{}{}
		}(w)
	}
	<-done
	<-done
	close(mu)
	for n := range mu {
		total += n
	}
	if total != 8 {
		t.Fatalf("并发认领应恰好 8 条（不重复不丢失）：total=%d", total)
	}
}