fix(cluster): 修复身份事务阻塞与 Worker 容量抖动
将身份协调和安全事件心跳改为 PostgreSQL 单 Leader 执行,并从独立 Worker 进程中移除身份运行时,避免多副本重复写同一状态。 为安全事件事务增加锁等待、空闲事务超时及独立回滚上下文;Worker 需连续丢失六次心跳后才判定失效,降低跨节点抖动导致的容量反复扩缩。 验证:Go 全量测试、go vet、Race 聚焦测试、PostgreSQL 18 Leader/安全事件/Worker 分配集成测试及迁移测试通过。
This commit is contained in:
@@ -488,8 +488,8 @@ func TestWorkerCapacityAllocationAndFailover(t *testing.T) {
|
||||
}
|
||||
if _, err := db.pool.Exec(ctx, `
|
||||
UPDATE gateway_worker_instances
|
||||
SET heartbeat_at = now() - interval '16 seconds'
|
||||
WHERE instance_id = $1`, secondID); err != nil {
|
||||
SET heartbeat_at = now() - $2::interval
|
||||
WHERE instance_id = $1`, secondID, (workerHeartbeatStaleAfter + time.Second).String()); err != nil {
|
||||
t.Fatalf("expire second heartbeat: %v", err)
|
||||
}
|
||||
first, err = db.RegisterWorkerInstance(ctx, WorkerRegistrationInput{InstanceID: firstID, DesiredCapacity: 5})
|
||||
|
||||
@@ -0,0 +1,93 @@
|
||||
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"
|
||||
)
|
||||
|
||||
// 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) 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
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestPostgresLeadershipIsExclusiveAndRecoverable(t *testing.T) {
|
||||
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
|
||||
if databaseURL == "" {
|
||||
t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run PostgreSQL leadership integration tests")
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
|
||||
defer cancel()
|
||||
first, err := Connect(ctx, databaseURL)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer first.Close()
|
||||
second, err := Connect(ctx, databaseURL)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer second.Close()
|
||||
|
||||
leader, acquired, err := first.TryAcquireSecurityEventHeartbeatLeadership(ctx)
|
||||
if err != nil || !acquired {
|
||||
t.Fatalf("first leadership acquired=%v err=%v", acquired, err)
|
||||
}
|
||||
defer leader.Release()
|
||||
if err := leader.KeepAlive(ctx); err != nil {
|
||||
t.Fatalf("leadership keepalive: %v", err)
|
||||
}
|
||||
blocked, acquired, err := second.TryAcquireSecurityEventHeartbeatLeadership(ctx)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if acquired {
|
||||
blocked.Release()
|
||||
t.Fatal("second store acquired leadership while first store held it")
|
||||
}
|
||||
|
||||
leader.Release()
|
||||
replacement, acquired, err := second.TryAcquireSecurityEventHeartbeatLeadership(ctx)
|
||||
if err != nil || !acquired {
|
||||
t.Fatalf("replacement leadership acquired=%v err=%v", acquired, err)
|
||||
}
|
||||
replacement.Release()
|
||||
}
|
||||
@@ -10,6 +10,42 @@ import (
|
||||
"github.com/jackc/pgx/v5"
|
||||
)
|
||||
|
||||
const (
|
||||
securityEventTransactionLockTimeout = 5 * time.Second
|
||||
securityEventTransactionIdleTimeout = 15 * time.Second
|
||||
transactionRollbackTimeout = 5 * time.Second
|
||||
)
|
||||
|
||||
func (s *Store) beginSecurityEventTransaction(ctx context.Context) (pgx.Tx, error) {
|
||||
tx, err := s.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
settings := []struct {
|
||||
name string
|
||||
value string
|
||||
}{
|
||||
{name: "idle_in_transaction_session_timeout", value: securityEventTransactionIdleTimeout.String()},
|
||||
{name: "lock_timeout", value: securityEventTransactionLockTimeout.String()},
|
||||
}
|
||||
for _, setting := range settings {
|
||||
if _, err := tx.Exec(ctx, `SELECT set_config($1, $2, true)`, setting.name, setting.value); err != nil {
|
||||
rollbackSecurityEventTransaction(tx)
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
return tx, nil
|
||||
}
|
||||
|
||||
func rollbackSecurityEventTransaction(tx pgx.Tx) {
|
||||
if tx == nil {
|
||||
return
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), transactionRollbackTimeout)
|
||||
defer cancel()
|
||||
_ = tx.Rollback(ctx)
|
||||
}
|
||||
|
||||
type ApplySessionRevokedInput struct {
|
||||
Issuer string
|
||||
Audience string
|
||||
@@ -41,11 +77,11 @@ func (s *Store) ApplySessionRevoked(ctx context.Context, input ApplySessionRevok
|
||||
if input.SubjectType == "" {
|
||||
input.SubjectType = "principal"
|
||||
}
|
||||
tx, err := s.pool.Begin(ctx)
|
||||
tx, err := s.beginSecurityEventTransaction(ctx)
|
||||
if err != nil {
|
||||
return ApplySecurityEventResult{}, err
|
||||
}
|
||||
defer func() { _ = tx.Rollback(ctx) }()
|
||||
defer rollbackSecurityEventTransaction(tx)
|
||||
subjectHash := shortSecurityEventHash(input.Subject)
|
||||
tag, err := tx.Exec(ctx, `
|
||||
INSERT INTO gateway_security_event_receipts (
|
||||
@@ -146,11 +182,11 @@ WHERE session.gateway_user_id = gateway_user.id
|
||||
}
|
||||
|
||||
func (s *Store) EnsureSecurityEventStreamState(ctx context.Context, issuer, audience, streamID string) error {
|
||||
tx, err := s.pool.Begin(ctx)
|
||||
tx, err := s.beginSecurityEventTransaction(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer func() { _ = tx.Rollback(ctx) }()
|
||||
defer rollbackSecurityEventTransaction(tx)
|
||||
if _, err := tx.Exec(ctx, `
|
||||
INSERT INTO gateway_security_event_stream_state(issuer,audience,stream_id,mode)
|
||||
VALUES($1,$2,$3::uuid,'bootstrap')
|
||||
@@ -184,11 +220,11 @@ func (s *Store) BeginSecurityEventVerification(ctx context.Context, issuer, audi
|
||||
if len(stateHash) != sha256.Size {
|
||||
return errors.New("verification state hash must be SHA-256")
|
||||
}
|
||||
tx, err := s.pool.Begin(ctx)
|
||||
tx, err := s.beginSecurityEventTransaction(ctx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer func() { _ = tx.Rollback(ctx) }()
|
||||
defer rollbackSecurityEventTransaction(tx)
|
||||
tag, err := tx.Exec(ctx, `
|
||||
UPDATE gateway_security_event_stream_state
|
||||
SET pending_state_hash=$3,pending_state_created_at=$4,updated_at=now()
|
||||
@@ -212,11 +248,11 @@ ON CONFLICT(issuer,audience,state_hash) DO NOTHING`, issuer, audience, stateHash
|
||||
}
|
||||
|
||||
func (s *Store) ConfirmSecurityEventVerification(ctx context.Context, issuer, audience, streamID, jti string, stateHash []byte, now time.Time) (bool, error) {
|
||||
tx, err := s.pool.Begin(ctx)
|
||||
tx, err := s.beginSecurityEventTransaction(ctx)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
defer func() { _ = tx.Rollback(ctx) }()
|
||||
defer rollbackSecurityEventTransaction(tx)
|
||||
tag, err := tx.Exec(ctx, `
|
||||
INSERT INTO gateway_security_event_receipts(issuer,jti,audience,event_type,subject_hash)
|
||||
VALUES($1,$2::uuid,$3,$4,$5)
|
||||
@@ -294,11 +330,11 @@ ON CONFLICT(issuer,jti) DO NOTHING`, issuer, jti, audience, eventType, shortSecu
|
||||
// status change and immediately moves OIDC traffic to introspection. Push is
|
||||
// trusted again only after the normal consecutive-verification recovery path.
|
||||
func (s *Store) ApplySecurityEventStreamUpdated(ctx context.Context, issuer, audience, jti, streamID, status, reason string, now time.Time) (bool, error) {
|
||||
tx, err := s.pool.Begin(ctx)
|
||||
tx, err := s.beginSecurityEventTransaction(ctx)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
defer func() { _ = tx.Rollback(ctx) }()
|
||||
defer rollbackSecurityEventTransaction(tx)
|
||||
tag, err := tx.Exec(ctx, `
|
||||
INSERT INTO gateway_security_event_receipts(issuer,jti,audience,event_type,subject_hash)
|
||||
VALUES($1,$2::uuid,$3,$4,$5)
|
||||
|
||||
@@ -9,15 +9,16 @@ import (
|
||||
"github.com/jackc/pgx/v5"
|
||||
)
|
||||
|
||||
const workerHeartbeatStaleAfter = 15 * time.Second
|
||||
const workerHeartbeatStaleAfter = 30 * time.Second
|
||||
|
||||
type WorkerRegistrationInput struct {
|
||||
InstanceID string
|
||||
PodUID string
|
||||
PodName string
|
||||
Site string
|
||||
Revision string
|
||||
DesiredCapacity int
|
||||
InstanceID string
|
||||
PodUID string
|
||||
PodName string
|
||||
Site string
|
||||
Revision string
|
||||
DesiredCapacity int
|
||||
HeartbeatStaleAfter time.Duration
|
||||
}
|
||||
|
||||
type WorkerAllocation struct {
|
||||
@@ -36,6 +37,10 @@ func (s *Store) RegisterWorkerInstance(ctx context.Context, input WorkerRegistra
|
||||
if input.DesiredCapacity < 0 {
|
||||
return WorkerAllocation{}, errors.New("worker desired capacity cannot be negative")
|
||||
}
|
||||
staleAfter := input.HeartbeatStaleAfter
|
||||
if staleAfter < workerHeartbeatStaleAfter {
|
||||
staleAfter = workerHeartbeatStaleAfter
|
||||
}
|
||||
tx, err := s.pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return WorkerAllocation{}, err
|
||||
@@ -73,7 +78,7 @@ UPDATE gateway_worker_instances
|
||||
SET allocated_capacity = 0,
|
||||
updated_at = now()
|
||||
WHERE status <> 'active'
|
||||
OR heartbeat_at <= now() - $1::interval`, workerHeartbeatStaleAfter.String()); err != nil {
|
||||
OR heartbeat_at <= now() - $1::interval`, staleAfter.String()); err != nil {
|
||||
return WorkerAllocation{}, err
|
||||
}
|
||||
|
||||
@@ -83,7 +88,7 @@ FROM gateway_worker_instances
|
||||
WHERE status = 'active'
|
||||
AND heartbeat_at > now() - $1::interval
|
||||
ORDER BY instance_id ASC
|
||||
FOR UPDATE`, workerHeartbeatStaleAfter.String())
|
||||
FOR UPDATE`, staleAfter.String())
|
||||
if err != nil {
|
||||
return WorkerAllocation{}, err
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user