Agent-System-Backend/internal/runmgr/admission.go

package runmgr

import (
	"context"
	"errors"
	"log"
	"sync"

	"agent-system-backend/internal/store"
)

// 协作存储 U3 切片 2：执行许可门（R8）。预算经 budgetFn 动态读取（设置
// runtime.execBudgetGlobal / runtime.execBudgetPerUser，默认全零=不限，本地档
// 行为与历史一致）。许可的共享权威在 store.admission_leases（PG 咨询锁下
// 原子清点）；本门只负责有界排队（接纳数）与让出/重取（父等待子代理或
// HITL 时让出执行许可、保留会话所有权——避免额度占尽时父子互等）。

// ErrAdmissionQueueFull 待执行队列已满（接纳被拒；HTTP 层映射 429）。
var ErrAdmissionQueueFull = errors.New("runmgr: execution queue full")

// admissionGate 每进程一个；释放时广播唤醒全部等待者（无公平队列，Golang
// cond 语义；公平分配留 U7 压测后按实测加）。
type admissionGate struct {
	mu        sync.Mutex
	broadcast chan struct{} // 释放时 close+重建
	waiters   int
	budgetFn  func() store.ExecutionBudget
	queueMax  int
}

func newAdmissionGate(budgetFn func() store.ExecutionBudget, queueMax int) *admissionGate {
	if budgetFn == nil {
		budgetFn = func() store.ExecutionBudget { return store.ExecutionBudget{} }
	}
	if queueMax <= 0 {
		queueMax = 200 // 方案容量合同的默认队列上限（待实测调整）
	}
	return &admissionGate{broadcast: make(chan struct{}), budgetFn: budgetFn, queueMax: queueMax}
}

// budget 加锁读取当前预算（与 SetExecutionBudgetFn 并发安全）。
func (g *admissionGate) budget() store.ExecutionBudget {
	if g == nil {
		return store.ExecutionBudget{} // 直接构造的 Manager（测试）：预算关闭
	}
	g.mu.Lock()
	defer g.mu.Unlock()
	return g.budgetFn()
}

// QueueDepth 返回当前排队数（诊断/验收）。
func (g *admissionGate) QueueDepth() int {
	if g == nil {
		return 0
	}
	g.mu.Lock()
	defer g.mu.Unlock()
	return g.waiters
}

// QueueFull 非阻塞探测：排队是否已满（API 边缘在落库前拒绝，429 且不写数据）。
func (g *admissionGate) QueueFull() bool {
	budget := g.budget()
	if budget.Off() {
		return false
	}
	g.mu.Lock()
	defer g.mu.Unlock()
	return g.waiters >= g.queueMax
}

// acquire 阻塞获取执行许可（ctx 取消则放弃排队）。预算关闭时零开销直通。
func (g *admissionGate) acquire(ctx context.Context, st *store.Store, runID, owner string, accountGuard ...bool) error {
	if g == nil {
		return nil
	}
	budget := g.budget()
	if budget.Off() {
		return ctx.Err()
	}
	for {
		var err error
		if len(accountGuard) > 0 && accountGuard[0] {
			allowImplicit := len(accountGuard) > 1 && accountGuard[1]
			err = st.TryAcquireExecutionSlotAccountChecked(owner, runID, budget, allowImplicit)
		} else {
			err = st.TryAcquireExecutionSlot(runID, owner, budget)
		}
		if err == nil {
			return nil
		} else if !errors.Is(err, store.ErrBudgetExceeded) {
			return err // 存储故障：不放行（fail-closed，宁可拒绝执行）
		}
		g.mu.Lock()
		if g.waiters >= g.queueMax {
			g.mu.Unlock()
			return ErrAdmissionQueueFull
		}
		g.waiters++
		ch := g.broadcast
		g.mu.Unlock()
		select {
		case <-ch:
			g.mu.Lock()
			g.waiters--
			g.mu.Unlock()
		case <-ctx.Done():
			g.mu.Lock()
			g.waiters--
			g.mu.Unlock()
			return ctx.Err()
		}
	}
}

// release 释放执行许可并广播（幂等；预算关闭时零开销）。
func (g *admissionGate) release(st *store.Store, runID string) {
	if g == nil || g.budget().Off() {
		return
	}
	if err := st.ReleaseExecutionSlot(runID); err != nil {
		// 释放失败只影响账本整洁（下次复活即修正），不中断运行。
		log.Printf("[runmgr] admission release %s: %v", runID, err)
	}
	g.mu.Lock()
	close(g.broadcast)
	g.broadcast = make(chan struct{})
	g.mu.Unlock()
}