wangbo
7786692d32
将 Worker 发现、路由画像、容量与执行传输抽象为平台无关接口，新增 Kubernetes 和静态容量适配器，并以 shadow 模式接入生产配置。 实现网络与容量评分、路由防抖、池队列、同步 Worker 租约、一次性执行令牌，以及提交状态不明时禁止重复分配的安全语义。 新增 0105 兼容迁移、管理接口、指标、OpenAPI 和回归测试。已执行全量 Go 测试、go vet、OpenAPI、迁移安全、Compose 与 Kustomize 验证。
103 lines
3.0 KiB
Go
103 lines
3.0 KiB
Go
package store
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
)
|
|
|
|
const leadershipReleaseTimeout = 5 * time.Second
|
|
|
|
const (
|
|
identityCoordinatorLeadershipKey = "easyai-gateway-identity-coordinator"
|
|
securityEventHeartbeatLeadershipKey = "easyai-gateway-security-event-heartbeat"
|
|
capacityControllerLeadershipKey = "easyai-gateway-capacity-controller"
|
|
)
|
|
|
|
// Leadership holds a PostgreSQL session advisory lock. Callers must keep the
|
|
// lease alive and release it when leadership is no longer needed.
|
|
type Leadership interface {
|
|
KeepAlive(context.Context) error
|
|
Release()
|
|
}
|
|
|
|
type postgresLeadership struct {
|
|
conn *pgxpool.Conn
|
|
key string
|
|
once sync.Once
|
|
}
|
|
|
|
func (s *Store) TryAcquireIdentityCoordinatorLeadership(ctx context.Context) (Leadership, bool, error) {
|
|
return s.tryAcquireLeadership(ctx, identityCoordinatorLeadershipKey)
|
|
}
|
|
|
|
func (s *Store) TryAcquireSecurityEventHeartbeatLeadership(ctx context.Context) (Leadership, bool, error) {
|
|
return s.tryAcquireLeadership(ctx, securityEventHeartbeatLeadershipKey)
|
|
}
|
|
|
|
func (s *Store) TryAcquireCapacityControllerLeadership(ctx context.Context) (Leadership, bool, error) {
|
|
return s.tryAcquireLeadership(ctx, capacityControllerLeadershipKey)
|
|
}
|
|
|
|
func (s *Store) TryAcquireRouteProbeLeadership(ctx context.Context, poolID, routeProfileKey string) (Leadership, bool, error) {
|
|
return s.tryAcquireLeadership(ctx, "easyai-gateway-route-probe:"+poolID+":"+routeProfileKey)
|
|
}
|
|
|
|
func (s *Store) tryAcquireLeadership(ctx context.Context, key string) (Leadership, bool, error) {
|
|
conn, err := s.pool.Acquire(ctx)
|
|
if err != nil {
|
|
return nil, false, err
|
|
}
|
|
var acquired bool
|
|
if err := conn.QueryRow(ctx,
|
|
`SELECT pg_try_advisory_lock(hashtextextended($1, 0))`,
|
|
key,
|
|
).Scan(&acquired); err != nil {
|
|
conn.Release()
|
|
return nil, false, err
|
|
}
|
|
if !acquired {
|
|
conn.Release()
|
|
return nil, false, nil
|
|
}
|
|
return &postgresLeadership{conn: conn, key: key}, true, nil
|
|
}
|
|
|
|
func (leadership *postgresLeadership) KeepAlive(ctx context.Context) error {
|
|
if leadership == nil || leadership.conn == nil {
|
|
return errors.New("leadership lease is closed")
|
|
}
|
|
var alive int
|
|
return leadership.conn.QueryRow(ctx, `SELECT 1`).Scan(&alive)
|
|
}
|
|
|
|
func (leadership *postgresLeadership) Release() {
|
|
if leadership == nil {
|
|
return
|
|
}
|
|
leadership.once.Do(func() {
|
|
unlockCtx, unlockCancel := context.WithTimeout(context.Background(), leadershipReleaseTimeout)
|
|
var unlocked bool
|
|
err := leadership.conn.QueryRow(unlockCtx,
|
|
`SELECT pg_advisory_unlock(hashtextextended($1, 0))`,
|
|
leadership.key,
|
|
).Scan(&unlocked)
|
|
unlockCancel()
|
|
if err != nil || !unlocked {
|
|
// A session-scoped advisory lock must never return to the pool when
|
|
// unlock is uncertain. Closing the connection lets PostgreSQL
|
|
// release the lock and forces the pool to replace the session.
|
|
conn := leadership.conn.Hijack()
|
|
closeCtx, closeCancel := context.WithTimeout(context.Background(), leadershipReleaseTimeout)
|
|
_ = conn.Close(closeCtx)
|
|
closeCancel()
|
|
} else {
|
|
leadership.conn.Release()
|
|
}
|
|
leadership.conn = nil
|
|
})
|
|
}
|