6dedc5973c
pgx 不识别 DSN 中 libpq 的 keepalives* 参数，会作为服务端运行参数发送导致连接被拒；改为在建连前消费并应用到 TCP 拨号器。 未显式配置时为每条连接默认设置 idle_in_transaction_session_timeout=5min，避免卡住的事务长期持锁；DSN 显式配置时以 DSN 为准。 建连钩子先复制共享的 RuntimeParams，避免多条连接之间累积修改。 Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
87 lines
2.9 KiB
Go
87 lines
2.9 KiB
Go
package database
|
||
|
||
import (
|
||
"context"
|
||
"fmt"
|
||
"maps"
|
||
"net"
|
||
"strconv"
|
||
"strings"
|
||
"time"
|
||
|
||
"github.com/jackc/pgx/v5"
|
||
)
|
||
|
||
// defaultIdleInTransactionSessionTimeout 限制事务内等待下一条命令的时间，
|
||
// 避免卡住的 idle in transaction 会话长期持有锁。DSN 显式配置时以 DSN 为准。
|
||
const defaultIdleInTransactionSessionTimeout = "5min"
|
||
|
||
var libpqKeepaliveParams = []string{"keepalives", "keepalives_idle", "keepalives_interval", "keepalives_count"}
|
||
|
||
// configurePostgresConnection 在每次建立物理连接前调整 pgx 连接配置。
|
||
func configurePostgresConnection(_ context.Context, config *pgx.ConnConfig) error {
|
||
// stdlib 传入的是浅拷贝，RuntimeParams 与连接器共享，修改前必须复制，否则会跨连接累积。
|
||
config.RuntimeParams = maps.Clone(config.RuntimeParams)
|
||
if config.RuntimeParams == nil {
|
||
config.RuntimeParams = make(map[string]string)
|
||
}
|
||
if err := applyLibpqKeepalives(config); err != nil {
|
||
return err
|
||
}
|
||
if _, exists := config.RuntimeParams["idle_in_transaction_session_timeout"]; !exists {
|
||
// PostgreSQL 按顺序处理 options，DSN 原有 options 放在后面，同名设置以 DSN 为准。
|
||
options := "-c idle_in_transaction_session_timeout=" + defaultIdleInTransactionSessionTimeout
|
||
if existing := strings.TrimSpace(config.RuntimeParams["options"]); existing != "" {
|
||
options += " " + existing
|
||
}
|
||
config.RuntimeParams["options"] = options
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// applyLibpqKeepalives 消费 libpq 的 keepalives* 参数并应用到 TCP 拨号器。
|
||
// pgx 不识别这些参数，会把它们当作服务端运行参数发送，导致 PostgreSQL 拒绝连接。
|
||
// 未配置任何 keepalives* 参数时保持 pgx 默认拨号器（Go 默认 keepalive）不变。
|
||
func applyLibpqKeepalives(config *pgx.ConnConfig) error {
|
||
values := make(map[string]int)
|
||
for _, key := range libpqKeepaliveParams {
|
||
raw, ok := config.RuntimeParams[key]
|
||
if !ok {
|
||
continue
|
||
}
|
||
delete(config.RuntimeParams, key)
|
||
value, err := strconv.Atoi(strings.TrimSpace(raw))
|
||
if err != nil || value < 0 || (key == "keepalives" && value > 1) {
|
||
return fmt.Errorf("invalid PostgreSQL connection parameter %s=%q", key, raw)
|
||
}
|
||
values[key] = value
|
||
}
|
||
if len(values) == 0 {
|
||
return nil
|
||
}
|
||
dialer := &net.Dialer{Timeout: config.ConnectTimeout}
|
||
if enabled, ok := values["keepalives"]; ok && enabled == 0 {
|
||
dialer.KeepAlive = -1
|
||
} else {
|
||
// 未配置的项使用 Go 默认值（0）；libpq 的 0 表示操作系统默认值，对应 Go 的 -1。
|
||
setting := func(key string) int {
|
||
value, ok := values[key]
|
||
if !ok {
|
||
return 0
|
||
}
|
||
if value == 0 {
|
||
return -1
|
||
}
|
||
return value
|
||
}
|
||
dialer.KeepAliveConfig = net.KeepAliveConfig{
|
||
Enable: true,
|
||
Idle: time.Duration(setting("keepalives_idle")) * time.Second,
|
||
Interval: time.Duration(setting("keepalives_interval")) * time.Second,
|
||
Count: setting("keepalives_count"),
|
||
}
|
||
}
|
||
config.DialFunc = dialer.DialContext
|
||
return nil
|
||
}
|