fix(cluster): 修复身份事务阻塞与 Worker 容量抖动

将身份协调和安全事件心跳改为 PostgreSQL 单 Leader 执行，并从独立 Worker 进程中移除身份运行时，避免多副本重复写同一状态。为安全事件事务增加锁等待、空闲事务超时及独立回滚上下文；Worker 需连续丢失六次心跳后才判定失效，降低跨节点抖动导致的容量反复扩缩。验证：Go 全量测试、go vet、Race 聚焦测试、PostgreSQL 18 Leader/安全事件/Worker 分配集成测试及迁移测试通过。
2026-07-29 23:04:41 +08:00
parent 3886048e0f
commit 91451b3c86
10 changed files with 497 additions and 76 deletions
@@ -7,9 +7,11 @@ import (
 	"time"

 	"github.com/easyai/easyai-ai-gateway/apps/api/internal/identity"
+	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
 )

 const identityPairingReconcileInterval = 5 * time.Second
+const identityPairingLeadershipTimeout = 5 * time.Second

 type identityPairingReconciliationStore interface {
 	PendingIdentityPairingExchanges(context.Context) ([]identity.PairingExchange, error)
@@ -21,9 +23,9 @@ type identityPairingRestorer interface {
 	RestoreCompletedSecurityEvents(context.Context, string) error
 }

-func (s *Server) reconcileCanonicalIdentityPairing() {
+func (s *Server) reconcileCanonicalIdentityPairing(ctx context.Context) {
 	reconcileCanonicalIdentityPairing(
-		s.ctx,
+		ctx,
 		s.store,
 		s.identityPairing,
 		&s.identityRestoredPairings,
@@ -34,15 +36,57 @@ func (s *Server) reconcileCanonicalIdentityPairing() {
 }

 func (s *Server) runIdentityPairingCoordinator() {
-	ticker := time.NewTicker(identityPairingReconcileInterval)
-	defer ticker.Stop()
+	retry := time.NewTicker(identityPairingReconcileInterval)
+	defer retry.Stop()
 	for {
+		acquireCtx, cancel := context.WithTimeout(s.ctx, identityPairingLeadershipTimeout)
+		leadership, acquired, err := s.store.TryAcquireIdentityCoordinatorLeadership(acquireCtx)
+		cancel()
+		if err != nil && s.logger != nil {
+			s.logger.Warn("identity pairing coordinator leader election failed",
+				"error_category", "identity_coordinator_leader_election_failed",
+				"error", err,
+			)
+		}
+		if acquired {
+			if s.logger != nil {
+				s.logger.Info("identity pairing coordinator leadership acquired")
+			}
+			if !s.runIdentityPairingLeader(leadership) {
+				return
+			}
+		}
 		select {
 		case <-s.ctx.Done():
 			return
+		case <-retry.C:
+		}
+	}
+}
+
+func (s *Server) runIdentityPairingLeader(leadership store.Leadership) bool {
+	defer leadership.Release()
+	ticker := time.NewTicker(identityPairingReconcileInterval)
+	defer ticker.Stop()
+	for {
+		cycleCtx, cancel := context.WithTimeout(s.ctx, identityPairingLeadershipTimeout)
+		if err := leadership.KeepAlive(cycleCtx); err != nil {
+			cancel()
+			if s.logger != nil {
+				s.logger.Warn("identity pairing coordinator leadership lost",
+					"error_category", "identity_coordinator_leadership_lost",
+					"error", err,
+				)
+			}
+			return true
+		}
+		s.reconcileIdentityRuntime(cycleCtx)
+		s.reconcileCanonicalIdentityPairing(cycleCtx)
+		cancel()
+		select {
+		case <-s.ctx.Done():
+			return false
 		case <-ticker.C:
-			s.reconcileIdentityRuntime()
-			s.reconcileCanonicalIdentityPairing()
 		}
 	}
 }
@@ -105,14 +149,15 @@ func reconcileCanonicalIdentityPairing(
 	}
 }

-func (s *Server) reconcileIdentityRuntime() {
+func (s *Server) reconcileIdentityRuntime(ctx context.Context) {
 	if s.identityRuntime == nil || !s.identityRuntime.ReconciliationRequired() {
 		return
 	}
-	ctx, cancel := context.WithTimeout(s.ctx, 5*time.Second)
-	defer cancel()
 	if err := s.identityRuntime.ReconcileActive(ctx); err != nil && s.logger != nil {
-		s.logger.Warn("identity runtime reconciliation failed and will retry", "error_category", "identity_runtime_reconciliation_failed")
+		s.logger.Warn("identity runtime reconciliation failed and will retry",
+			"error_category", "identity_runtime_reconciliation_failed",
+			"error", err,
+		)
 	}
 }

@@ -81,43 +81,8 @@ func NewServerWithContext(ctx context.Context, cfg config.Config, db *store.Stor
 	}
 	server.auth.ServerMainInternalKey = cfg.ServerMainInternalKey
 	server.auth.ServerMainInternalSecret = cfg.ServerMainInternalSecret
-	secretStore, err := identitySecretStore(cfg)
-	if err != nil {
-		panic("invalid identity SecretStore: " + err.Error())
-	}
-	go ssfreceiver.RunSecurityEventRetirementWorker(ctx, db)
-	go ssfreceiver.RunIdentitySecretCleanupWorker(ctx, db, secretStore)
-	runtimeBuilder := identityruntime.NewRuntimeBuilder(ctx, db, secretStore, identityruntime.RuntimeBuilderConfig{
-		AppEnv: cfg.AppEnv, JWKSCacheTTL: 5 * time.Minute,
-		HeartbeatInterval: time.Duration(cfg.IdentitySecurityEventHeartbeatIntervalSeconds) * time.Second,
-		StaleAfter:        time.Duration(cfg.IdentitySecurityEventStaleAfterSeconds) * time.Second,
-		ClockSkew:         time.Duration(cfg.IdentitySecurityEventClockSkewSeconds) * time.Second,
-	}, securityEventMetrics)
-	server.identityRuntime = identityruntime.NewManager(db, runtimeBuilder)
-	if err := server.identityRuntime.LoadActive(ctx); err != nil && logger != nil {
-		logger.Error("load active identity runtime failed; local management login remains available", "error_category", "identity_runtime_load_failed")
-	}
-	server.identityPairing = identity.NewPairingService(db, secretStore, func(baseURL string) (identity.OnboardingRemote, error) {
-		return identity.NewOnboardingClient(baseURL, nil, cfg.AppEnv)
-	}, server.identityRuntime, cfg.AppEnv)
-	server.reconcileCanonicalIdentityPairing()
-	go server.runIdentityPairingCoordinator()
-	server.auth.OIDCVerifierProvider = func() *auth.OIDCVerifier {
-		if runtime := server.identityRuntime.Current(); runtime != nil {
-			return runtime.Verifier
-		}
-		return nil
-	}
-	server.auth.OIDCSessionResolverProvider = func(ctx context.Context, sessionID string) (*auth.User, error) {
-		runtime := server.identityRuntime.Current()
-		if runtime == nil || runtime.Sessions == nil {
-			return nil, auth.ErrUnauthorized
-		}
-		user, resolveErr := runtime.Sessions.Resolve(ctx, sessionID)
-		return user, oidcSessionRequestError(resolveErr)
-	}
-	server.auth.LegacyJWTEnabledProvider = func() bool {
-		return server.identityRuntime.LegacyJWTEnabled()
+	if cfg.RunsPublicHTTP() {
+		server.initializeIdentityRuntime(ctx, db, securityEventMetrics)
 	}
 	server.auth.LocalAPIKeyVerifier = db.VerifyLocalAPIKey
 	server.runner.StartAdmissionNotifier(ctx)
@@ -359,6 +324,54 @@ func NewServerWithContext(ctx context.Context, cfg config.Config, db *store.Stor
 	return server.recover(server.cors(server.protectOIDCSessionCookie(mux)))
 }

+func (server *Server) initializeIdentityRuntime(
+	ctx context.Context,
+	db *store.Store,
+	securityEventMetrics *ssfreceiver.Metrics,
+) {
+	secretStore, err := identitySecretStore(server.cfg)
+	if err != nil {
+		panic("invalid identity SecretStore: " + err.Error())
+	}
+	go ssfreceiver.RunSecurityEventRetirementWorker(ctx, db)
+	go ssfreceiver.RunIdentitySecretCleanupWorker(ctx, db, secretStore)
+	runtimeBuilder := identityruntime.NewRuntimeBuilder(ctx, db, secretStore, identityruntime.RuntimeBuilderConfig{
+		AppEnv: server.cfg.AppEnv, JWKSCacheTTL: 5 * time.Minute,
+		HeartbeatInterval: time.Duration(server.cfg.IdentitySecurityEventHeartbeatIntervalSeconds) * time.Second,
+		StaleAfter:        time.Duration(server.cfg.IdentitySecurityEventStaleAfterSeconds) * time.Second,
+		ClockSkew:         time.Duration(server.cfg.IdentitySecurityEventClockSkewSeconds) * time.Second,
+	}, securityEventMetrics)
+	server.identityRuntime = identityruntime.NewManager(db, runtimeBuilder)
+	if err := server.identityRuntime.LoadActive(ctx); err != nil && server.logger != nil {
+		server.logger.Error(
+			"load active identity runtime failed; local management login remains available",
+			"error_category", "identity_runtime_load_failed",
+			"error", err,
+		)
+	}
+	server.identityPairing = identity.NewPairingService(db, secretStore, func(baseURL string) (identity.OnboardingRemote, error) {
+		return identity.NewOnboardingClient(baseURL, nil, server.cfg.AppEnv)
+	}, server.identityRuntime, server.cfg.AppEnv)
+	go server.runIdentityPairingCoordinator()
+	server.auth.OIDCVerifierProvider = func() *auth.OIDCVerifier {
+		if runtime := server.identityRuntime.Current(); runtime != nil {
+			return runtime.Verifier
+		}
+		return nil
+	}
+	server.auth.OIDCSessionResolverProvider = func(ctx context.Context, sessionID string) (*auth.User, error) {
+		runtime := server.identityRuntime.Current()
+		if runtime == nil || runtime.Sessions == nil {
+			return nil, auth.ErrUnauthorized
+		}
+		user, resolveErr := runtime.Sessions.Resolve(ctx, sessionID)
+		return user, oidcSessionRequestError(resolveErr)
+	}
+	server.auth.LegacyJWTEnabledProvider = func() bool {
+		return server.identityRuntime.LegacyJWTEnabled()
+	}
+}
+
 func identitySecretStore(cfg config.Config) (ssfreceiver.SecretStore, error) {
 	switch strings.ToLower(strings.TrimSpace(cfg.IdentitySecretStore)) {
 	case "", "file":
@@ -215,12 +215,13 @@ func (s *Service) loadAsyncWorkerCapacity(ctx context.Context) (store.AsyncWorke
 		return store.AsyncWorkerCapacitySnapshot{}, err
 	}
 	allocation, err := s.store.RegisterWorkerInstance(ctx, store.WorkerRegistrationInput{
-		InstanceID:      s.workerInstanceID,
-		PodUID:          strings.TrimSpace(os.Getenv("POD_UID")),
-		PodName:         strings.TrimSpace(os.Getenv("POD_NAME")),
-		Site:            strings.TrimSpace(os.Getenv("EASYAI_SITE")),
-		Revision:        strings.TrimSpace(os.Getenv("AI_GATEWAY_REVISION")),
-		DesiredCapacity: snapshot.Capacity,
+		InstanceID:          s.workerInstanceID,
+		PodUID:              strings.TrimSpace(os.Getenv("POD_UID")),
+		PodName:             strings.TrimSpace(os.Getenv("POD_NAME")),
+		Site:                strings.TrimSpace(os.Getenv("EASYAI_SITE")),
+		Revision:            strings.TrimSpace(os.Getenv("AI_GATEWAY_REVISION")),
+		DesiredCapacity:     snapshot.Capacity,
+		HeartbeatStaleAfter: time.Duration(s.cfg.AsyncWorkerRefreshIntervalSeconds) * 6 * time.Second,
 	})
 	if err != nil {
 		return store.AsyncWorkerCapacitySnapshot{}, err
@@ -52,6 +52,16 @@ type Service struct {
 	metrics    *Metrics
 }

+type heartbeatLeadershipRepository interface {
+	TryAcquireSecurityEventHeartbeatLeadership(context.Context) (store.Leadership, bool, error)
+}
+
+const (
+	heartbeatLeadershipRetryInterval = 5 * time.Second
+	heartbeatLeadershipKeepAlive     = 5 * time.Second
+	heartbeatExecutionTimeout        = 30 * time.Second
+)
+
 func (s *Service) SetMetrics(metrics *Metrics) { s.metrics = metrics }

 func NewService(repository StateRepository, config ServiceConfig) (*Service, error) {
@@ -108,7 +118,34 @@ func (s *Service) Evaluate(ctx context.Context, identity auth.OIDCSecurityEventI
 }

 func (s *Service) run(ctx context.Context) {
-	s.sendVerification(ctx)
+	elector, distributed := s.repository.(heartbeatLeadershipRepository)
+	if !distributed {
+		s.runHeartbeatLoop(ctx)
+		return
+	}
+	retry := time.NewTicker(heartbeatLeadershipRetryInterval)
+	defer retry.Stop()
+	for {
+		acquireCtx, cancel := context.WithTimeout(ctx, heartbeatLeadershipKeepAlive)
+		leadership, acquired, err := elector.TryAcquireSecurityEventHeartbeatLeadership(acquireCtx)
+		cancel()
+		if err != nil {
+			_ = s.repository.RecordSecurityEventHeartbeatFailure(ctx, s.config.Issuer, s.config.Audience, "leader_election_failed")
+		} else if acquired {
+			if !s.runLeaderHeartbeatLoop(ctx, leadership) {
+				return
+			}
+		}
+		select {
+		case <-ctx.Done():
+			return
+		case <-retry.C:
+		}
+	}
+}
+
+func (s *Service) runHeartbeatLoop(ctx context.Context) {
+	s.sendScheduledVerification(ctx)
 	ticker := time.NewTicker(s.config.HeartbeatInterval)
 	defer ticker.Stop()
 	for {
@@ -116,11 +153,42 @@ func (s *Service) run(ctx context.Context) {
 		case <-ctx.Done():
 			return
 		case <-ticker.C:
-			s.sendVerification(ctx)
+			s.sendScheduledVerification(ctx)
 		}
 	}
 }

+func (s *Service) runLeaderHeartbeatLoop(ctx context.Context, leadership store.Leadership) bool {
+	defer leadership.Release()
+	s.sendScheduledVerification(ctx)
+	heartbeat := time.NewTicker(s.config.HeartbeatInterval)
+	defer heartbeat.Stop()
+	keepalive := time.NewTicker(heartbeatLeadershipKeepAlive)
+	defer keepalive.Stop()
+	for {
+		select {
+		case <-ctx.Done():
+			return false
+		case <-heartbeat.C:
+			s.sendScheduledVerification(ctx)
+		case <-keepalive.C:
+			keepaliveCtx, cancel := context.WithTimeout(ctx, heartbeatLeadershipKeepAlive)
+			err := leadership.KeepAlive(keepaliveCtx)
+			cancel()
+			if err != nil {
+				_ = s.repository.RecordSecurityEventHeartbeatFailure(ctx, s.config.Issuer, s.config.Audience, "leader_keepalive_failed")
+				return true
+			}
+		}
+	}
+}
+
+func (s *Service) sendScheduledVerification(ctx context.Context) {
+	verificationCtx, cancel := context.WithTimeout(ctx, heartbeatExecutionTimeout)
+	defer cancel()
+	s.sendVerification(verificationCtx)
+}
+
 func (s *Service) sendVerification(ctx context.Context) {
 	outcome := "failed"
 	defer func() {
@@ -15,6 +15,8 @@ import (
 type fakeStateRepository struct {
 	mutex        sync.Mutex
 	stateHash    []byte
+	beginCount   int
+	beginSignal  chan struct{}
 	verification chan struct{}
 	evaluation   store.SecurityEventEvaluation
 }
@@ -25,6 +27,13 @@ func (*fakeStateRepository) EnsureSecurityEventStreamState(context.Context, stri
 func (f *fakeStateRepository) BeginSecurityEventVerification(_ context.Context, _, _ string, hash []byte, _ time.Time) error {
 	f.mutex.Lock()
 	f.stateHash = append([]byte(nil), hash...)
+	f.beginCount++
+	if f.beginSignal != nil {
+		select {
+		case f.beginSignal <- struct{}{}:
+		default:
+		}
+	}
 	f.mutex.Unlock()
 	return nil
 }
@@ -38,6 +47,105 @@ func (f *fakeStateRepository) EvaluateOIDCSecurityEvent(context.Context, string,
 	return f.evaluation, nil
 }

+type deniedHeartbeatLeadershipRepository struct {
+	*fakeStateRepository
+}
+
+func (*deniedHeartbeatLeadershipRepository) TryAcquireSecurityEventHeartbeatLeadership(context.Context) (store.Leadership, bool, error) {
+	return nil, false, nil
+}
+
+type fakeHeartbeatLeadership struct {
+	released chan struct{}
+	once     sync.Once
+}
+
+func (*fakeHeartbeatLeadership) KeepAlive(context.Context) error { return nil }
+func (leadership *fakeHeartbeatLeadership) Release() {
+	leadership.once.Do(func() { close(leadership.released) })
+}
+
+type grantedHeartbeatLeadershipRepository struct {
+	*fakeStateRepository
+	leadership *fakeHeartbeatLeadership
+}
+
+func (repository *grantedHeartbeatLeadershipRepository) TryAcquireSecurityEventHeartbeatLeadership(context.Context) (store.Leadership, bool, error) {
+	return repository.leadership, true, nil
+}
+
+func TestServiceFollowerDoesNotRunVerificationHeartbeat(t *testing.T) {
+	repository := &deniedHeartbeatLeadershipRepository{fakeStateRepository: &fakeStateRepository{}}
+	service, err := NewService(repository, ServiceConfig{
+		Issuer: "https://auth.test.example/ssf", Audience: "urn:easyai:ssf:receiver:app",
+		StreamID:           "00000000-0000-4000-8000-000000000001",
+		ManagementTokenURL: "https://auth.test.example/token",
+		ManagementClientID: "gateway-ssf", ManagementSecret: "management-secret",
+		HeartbeatInterval: time.Hour, StaleAfter: 2 * time.Hour,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	done := make(chan struct{})
+	go func() {
+		service.Run(ctx)
+		close(done)
+	}()
+	time.Sleep(100 * time.Millisecond)
+	cancel()
+	<-done
+	repository.mutex.Lock()
+	defer repository.mutex.Unlock()
+	if repository.beginCount != 0 {
+		t.Fatalf("follower persisted %d verification heartbeats", repository.beginCount)
+	}
+}
+
+func TestServiceLeaderRunsHeartbeatAndReleasesLeadership(t *testing.T) {
+	repository := &grantedHeartbeatLeadershipRepository{
+		fakeStateRepository: &fakeStateRepository{beginSignal: make(chan struct{}, 1)},
+		leadership:          &fakeHeartbeatLeadership{released: make(chan struct{})},
+	}
+	tokenServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path == "/token" {
+			w.Header().Set("Content-Type", "application/json")
+			_, _ = w.Write([]byte(`{"access_token":"management-token","token_type":"Bearer","expires_in":300}`))
+			return
+		}
+		w.WriteHeader(http.StatusNoContent)
+	}))
+	defer tokenServer.Close()
+	service, err := NewService(repository, ServiceConfig{
+		Issuer: tokenServer.URL, Audience: "urn:easyai:ssf:receiver:app",
+		StreamID:           "00000000-0000-4000-8000-000000000001",
+		ManagementTokenURL: tokenServer.URL + "/token",
+		ManagementClientID: "gateway-ssf", ManagementSecret: "management-secret",
+		HeartbeatInterval: time.Hour, StaleAfter: 2 * time.Hour, HTTPClient: tokenServer.Client(),
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	done := make(chan struct{})
+	go func() {
+		service.Run(ctx)
+		close(done)
+	}()
+	select {
+	case <-repository.beginSignal:
+	case <-time.After(2 * time.Second):
+		t.Fatal("leader did not persist a verification heartbeat")
+	}
+	cancel()
+	<-done
+	select {
+	case <-repository.leadership.released:
+	default:
+		t.Fatal("leader did not release PostgreSQL leadership")
+	}
+}
+
 func TestServiceRequestsVerificationWithMachineTokenAndPersistsStateFirst(t *testing.T) {
 	repository := &fakeStateRepository{verification: make(chan struct{}, 1)}
 	var server *httptest.Server
@@ -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
 	}