fix(worker): 防止事务泄漏并恢复滞留队列
为 PostgreSQL 连接增加可配置的事务空闲与锁等待超时,并在请求取消或 Worker 退出后使用独立有界上下文回滚事务。\n\n恢复过期任务时按批次使用 SKIP LOCKED,周期重建缺失或终态 River Job;锁等待超时只在确认尚未提交上游时安全重排,避免重复执行和重复结算。\n\n验证:完整 Go 测试、go vet、生产 Kustomize 渲染、gofmt 与 git diff --check 均通过。
This commit is contained in:
@@ -8,6 +8,7 @@ import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
"unicode"
|
||||
@@ -29,6 +30,12 @@ const (
|
||||
postgresPoolWarmLimit = 64
|
||||
)
|
||||
|
||||
type PostgresPoolOptions struct {
|
||||
MaxConns int
|
||||
IdleInTransactionTimeout time.Duration
|
||||
LockTimeout time.Duration
|
||||
}
|
||||
|
||||
func defaultAPIKeyScopes() []string {
|
||||
return []string{"chat", "embedding", "rerank", "image", "image_vectorize", "video", "video_enhance", "music", "audio", "voice_clone"}
|
||||
}
|
||||
@@ -81,7 +88,11 @@ func Connect(ctx context.Context, databaseURL string) (*Store, error) {
|
||||
}
|
||||
|
||||
func ConnectWithMaxConns(ctx context.Context, databaseURL string, maxConns int) (*Store, error) {
|
||||
config, err := postgresPoolConfig(databaseURL, maxConns)
|
||||
return ConnectWithPoolOptions(ctx, databaseURL, PostgresPoolOptions{MaxConns: maxConns})
|
||||
}
|
||||
|
||||
func ConnectWithPoolOptions(ctx context.Context, databaseURL string, options PostgresPoolOptions) (*Store, error) {
|
||||
config, err := postgresPoolConfigWithOptions(databaseURL, options)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -97,19 +108,37 @@ func ConnectWithMaxConns(ctx context.Context, databaseURL string, maxConns int)
|
||||
}
|
||||
|
||||
func postgresPoolConfig(databaseURL string, maxConns ...int) (*pgxpool.Config, error) {
|
||||
options := PostgresPoolOptions{}
|
||||
if len(maxConns) > 0 {
|
||||
options.MaxConns = maxConns[0]
|
||||
}
|
||||
return postgresPoolConfigWithOptions(databaseURL, options)
|
||||
}
|
||||
|
||||
func postgresPoolConfigWithOptions(databaseURL string, options PostgresPoolOptions) (*pgxpool.Config, error) {
|
||||
config, err := pgxpool.ParseConfig(databaseURL)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
config.ConnConfig.ConnectTimeout = postgresConnectTimeout
|
||||
config.ConnConfig.RuntimeParams["application_name"] = postgresApplicationName
|
||||
if len(maxConns) > 0 && maxConns[0] > 0 {
|
||||
config.MaxConns = int32(maxConns[0])
|
||||
config.MinIdleConns = int32(min(maxConns[0], postgresPoolWarmLimit))
|
||||
if options.IdleInTransactionTimeout > 0 {
|
||||
config.ConnConfig.RuntimeParams["idle_in_transaction_session_timeout"] = postgresDuration(options.IdleInTransactionTimeout)
|
||||
}
|
||||
if options.LockTimeout > 0 {
|
||||
config.ConnConfig.RuntimeParams["lock_timeout"] = postgresDuration(options.LockTimeout)
|
||||
}
|
||||
if options.MaxConns > 0 {
|
||||
config.MaxConns = int32(options.MaxConns)
|
||||
config.MinIdleConns = int32(min(options.MaxConns, postgresPoolWarmLimit))
|
||||
}
|
||||
return config, nil
|
||||
}
|
||||
|
||||
func postgresDuration(value time.Duration) string {
|
||||
return strconv.FormatInt(max(value.Milliseconds(), 1), 10) + "ms"
|
||||
}
|
||||
|
||||
func warmPostgresPool(ctx context.Context, pool *pgxpool.Pool, count int) error {
|
||||
if count <= 0 {
|
||||
return pool.Ping(ctx)
|
||||
@@ -152,6 +181,11 @@ func IsPostgresUnavailable(err error) bool {
|
||||
return pgconn.SafeToRetry(err)
|
||||
}
|
||||
|
||||
func IsPostgresLockTimeout(err error) bool {
|
||||
var postgresError *pgconn.PgError
|
||||
return errors.As(err, &postgresError) && postgresError.Code == "55P03"
|
||||
}
|
||||
|
||||
func (s *Store) Close() {
|
||||
s.pool.Close()
|
||||
}
|
||||
@@ -938,7 +972,7 @@ func (s *Store) DeletePlatform(ctx context.Context, id string) error {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer tx.Rollback(ctx)
|
||||
defer rollbackTransaction(tx)
|
||||
|
||||
rows, err := tx.Query(ctx, `SELECT id::text FROM platform_models WHERE platform_id = $1::uuid`, id)
|
||||
if err != nil {
|
||||
@@ -1638,7 +1672,7 @@ func (s *Store) DeleteAPIKey(ctx context.Context, apiKeyID string, user *auth.Us
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer tx.Rollback(ctx)
|
||||
defer rollbackTransaction(tx)
|
||||
|
||||
result, err := tx.Exec(ctx, `
|
||||
UPDATE gateway_api_keys
|
||||
@@ -1792,7 +1826,7 @@ func (s *Store) RegisterLocalUser(ctx context.Context, input LocalRegisterInput)
|
||||
if err != nil {
|
||||
return GatewayUser{}, err
|
||||
}
|
||||
defer tx.Rollback(ctx)
|
||||
defer rollbackTransaction(tx)
|
||||
|
||||
var tenantID string
|
||||
userGroupID := ""
|
||||
@@ -2065,7 +2099,7 @@ func (s *Store) CreateTaskIdempotent(ctx context.Context, input CreateTaskInput,
|
||||
if err != nil {
|
||||
return CreateTaskResult{}, err
|
||||
}
|
||||
defer tx.Rollback(ctx)
|
||||
defer rollbackTransaction(tx)
|
||||
|
||||
task, err := scanGatewayTask(tx.QueryRow(ctx, `
|
||||
INSERT INTO gateway_tasks (
|
||||
|
||||
Reference in New Issue
Block a user