fix: failed to migrate
This commit is contained in:
+269
-123
@@ -13,11 +13,12 @@ import (
|
||||
"net"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5/internal/iobufpool"
|
||||
"github.com/jackc/pgx/v5/internal/nbconn"
|
||||
"github.com/jackc/pgx/v5/internal/pgio"
|
||||
"github.com/jackc/pgx/v5/pgconn/internal/bgreader"
|
||||
"github.com/jackc/pgx/v5/pgconn/internal/ctxwatch"
|
||||
"github.com/jackc/pgx/v5/pgproto3"
|
||||
)
|
||||
@@ -65,17 +66,24 @@ type NotificationHandler func(*PgConn, *Notification)
|
||||
|
||||
// PgConn is a low-level PostgreSQL connection handle. It is not safe for concurrent usage.
|
||||
type PgConn struct {
|
||||
conn nbconn.Conn // the non-blocking wrapper for the underlying TCP or unix domain socket connection
|
||||
conn net.Conn
|
||||
pid uint32 // backend pid
|
||||
secretKey uint32 // key to use to send a cancel query message to the server
|
||||
parameterStatuses map[string]string // parameters that have been reported by the server
|
||||
txStatus byte
|
||||
frontend *pgproto3.Frontend
|
||||
bgReader *bgreader.BGReader
|
||||
slowWriteTimer *time.Timer
|
||||
|
||||
config *Config
|
||||
|
||||
status byte // One of connStatus* constants
|
||||
|
||||
bufferingReceive bool
|
||||
bufferingReceiveMux sync.Mutex
|
||||
bufferingReceiveMsg pgproto3.BackendMessage
|
||||
bufferingReceiveErr error
|
||||
|
||||
peekedMsg pgproto3.BackendMessage
|
||||
|
||||
// Reusable / preallocated resources
|
||||
@@ -89,7 +97,7 @@ type PgConn struct {
|
||||
}
|
||||
|
||||
// Connect establishes a connection to a PostgreSQL server using the environment and connString (in URL or DSN format)
|
||||
// to provide configuration. See documentation for ParseConfig for details. ctx can be used to cancel a connect attempt.
|
||||
// to provide configuration. See documentation for [ParseConfig] for details. ctx can be used to cancel a connect attempt.
|
||||
func Connect(ctx context.Context, connString string) (*PgConn, error) {
|
||||
config, err := ParseConfig(connString)
|
||||
if err != nil {
|
||||
@@ -100,7 +108,7 @@ func Connect(ctx context.Context, connString string) (*PgConn, error) {
|
||||
}
|
||||
|
||||
// Connect establishes a connection to a PostgreSQL server using the environment and connString (in URL or DSN format)
|
||||
// and ParseConfigOptions to provide additional configuration. See documentation for ParseConfig for details. ctx can be
|
||||
// and ParseConfigOptions to provide additional configuration. See documentation for [ParseConfig] for details. ctx can be
|
||||
// used to cancel a connect attempt.
|
||||
func ConnectWithOptions(ctx context.Context, connString string, parseConfigOptions ParseConfigOptions) (*PgConn, error) {
|
||||
config, err := ParseConfigWithOptions(connString, parseConfigOptions)
|
||||
@@ -112,7 +120,7 @@ func ConnectWithOptions(ctx context.Context, connString string, parseConfigOptio
|
||||
}
|
||||
|
||||
// Connect establishes a connection to a PostgreSQL server using config. config must have been constructed with
|
||||
// ParseConfig. ctx can be used to cancel a connect attempt.
|
||||
// [ParseConfig]. ctx can be used to cancel a connect attempt.
|
||||
//
|
||||
// If config.Fallbacks are present they will sequentially be tried in case of error establishing network connection. An
|
||||
// authentication error will terminate the chain of attempts (like libpq:
|
||||
@@ -146,12 +154,15 @@ func ConnectConfig(octx context.Context, config *Config) (pgConn *PgConn, err er
|
||||
|
||||
foundBestServer := false
|
||||
var fallbackConfig *FallbackConfig
|
||||
for _, fc := range fallbackConfigs {
|
||||
for i, fc := range fallbackConfigs {
|
||||
// ConnectTimeout restricts the whole connection process.
|
||||
if config.ConnectTimeout != 0 {
|
||||
var cancel context.CancelFunc
|
||||
ctx, cancel = context.WithTimeout(octx, config.ConnectTimeout)
|
||||
defer cancel()
|
||||
// create new context first time or when previous host was different
|
||||
if i == 0 || (fallbackConfigs[i].Host != fallbackConfigs[i-1].Host) {
|
||||
var cancel context.CancelFunc
|
||||
ctx, cancel = context.WithTimeout(octx, config.ConnectTimeout)
|
||||
defer cancel()
|
||||
}
|
||||
} else {
|
||||
ctx = octx
|
||||
}
|
||||
@@ -166,7 +177,7 @@ func ConnectConfig(octx context.Context, config *Config) (pgConn *PgConn, err er
|
||||
const ERRCODE_INVALID_CATALOG_NAME = "3D000" // db does not exist
|
||||
const ERRCODE_INSUFFICIENT_PRIVILEGE = "42501" // missing connect privilege
|
||||
if pgerr.Code == ERRCODE_INVALID_PASSWORD ||
|
||||
pgerr.Code == ERRCODE_INVALID_AUTHORIZATION_SPECIFICATION ||
|
||||
pgerr.Code == ERRCODE_INVALID_AUTHORIZATION_SPECIFICATION && fc.TLSConfig != nil ||
|
||||
pgerr.Code == ERRCODE_INVALID_CATALOG_NAME ||
|
||||
pgerr.Code == ERRCODE_INSUFFICIENT_PRIVILEGE {
|
||||
break
|
||||
@@ -203,6 +214,8 @@ func ConnectConfig(octx context.Context, config *Config) (pgConn *PgConn, err er
|
||||
func expandWithIPs(ctx context.Context, lookupFn LookupFunc, fallbacks []*FallbackConfig) ([]*FallbackConfig, error) {
|
||||
var configs []*FallbackConfig
|
||||
|
||||
var lookupErrors []error
|
||||
|
||||
for _, fb := range fallbacks {
|
||||
// skip resolve for unix sockets
|
||||
if isAbsolutePath(fb.Host) {
|
||||
@@ -217,7 +230,8 @@ func expandWithIPs(ctx context.Context, lookupFn LookupFunc, fallbacks []*Fallba
|
||||
|
||||
ips, err := lookupFn(ctx, fb.Host)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
lookupErrors = append(lookupErrors, err)
|
||||
continue
|
||||
}
|
||||
|
||||
for _, ip := range ips {
|
||||
@@ -242,11 +256,18 @@ func expandWithIPs(ctx context.Context, lookupFn LookupFunc, fallbacks []*Fallba
|
||||
}
|
||||
}
|
||||
|
||||
// See https://github.com/jackc/pgx/issues/1464. When Go 1.20 can be used in pgx consider using errors.Join so all
|
||||
// errors are reported.
|
||||
if len(configs) == 0 && len(lookupErrors) > 0 {
|
||||
return nil, lookupErrors[0]
|
||||
}
|
||||
|
||||
return configs, nil
|
||||
}
|
||||
|
||||
func connect(ctx context.Context, config *Config, fallbackConfig *FallbackConfig,
|
||||
ignoreNotPreferredErr bool) (*PgConn, error) {
|
||||
ignoreNotPreferredErr bool,
|
||||
) (*PgConn, error) {
|
||||
pgConn := new(PgConn)
|
||||
pgConn.config = config
|
||||
pgConn.cleanupDone = make(chan struct{})
|
||||
@@ -257,14 +278,13 @@ func connect(ctx context.Context, config *Config, fallbackConfig *FallbackConfig
|
||||
if err != nil {
|
||||
return nil, &connectError{config: config, msg: "dial error", err: normalizeTimeoutError(ctx, err)}
|
||||
}
|
||||
nbNetConn := nbconn.NewNetConn(netConn, false)
|
||||
|
||||
pgConn.conn = nbNetConn
|
||||
pgConn.contextWatcher = newContextWatcher(nbNetConn)
|
||||
pgConn.conn = netConn
|
||||
pgConn.contextWatcher = newContextWatcher(netConn)
|
||||
pgConn.contextWatcher.Watch(ctx)
|
||||
|
||||
if fallbackConfig.TLSConfig != nil {
|
||||
nbTLSConn, err := startTLS(nbNetConn, fallbackConfig.TLSConfig)
|
||||
nbTLSConn, err := startTLS(netConn, fallbackConfig.TLSConfig)
|
||||
pgConn.contextWatcher.Unwatch() // Always unwatch `netConn` after TLS.
|
||||
if err != nil {
|
||||
netConn.Close()
|
||||
@@ -280,7 +300,10 @@ func connect(ctx context.Context, config *Config, fallbackConfig *FallbackConfig
|
||||
|
||||
pgConn.parameterStatuses = make(map[string]string)
|
||||
pgConn.status = connStatusConnecting
|
||||
pgConn.frontend = config.BuildFrontend(pgConn.conn, pgConn.conn)
|
||||
pgConn.bgReader = bgreader.New(pgConn.conn)
|
||||
pgConn.slowWriteTimer = time.AfterFunc(time.Duration(math.MaxInt64), pgConn.bgReader.Start)
|
||||
pgConn.slowWriteTimer.Stop()
|
||||
pgConn.frontend = config.BuildFrontend(pgConn.bgReader, pgConn.conn)
|
||||
|
||||
startupMsg := pgproto3.StartupMessage{
|
||||
ProtocolVersion: pgproto3.ProtocolVersionNumber,
|
||||
@@ -298,9 +321,9 @@ func connect(ctx context.Context, config *Config, fallbackConfig *FallbackConfig
|
||||
}
|
||||
|
||||
pgConn.frontend.Send(&startupMsg)
|
||||
if err := pgConn.frontend.Flush(); err != nil {
|
||||
if err := pgConn.flushWithPotentialWriteReadDeadlock(); err != nil {
|
||||
pgConn.conn.Close()
|
||||
return nil, &connectError{config: config, msg: "failed to write startup message", err: err}
|
||||
return nil, &connectError{config: config, msg: "failed to write startup message", err: normalizeTimeoutError(ctx, err)}
|
||||
}
|
||||
|
||||
for {
|
||||
@@ -383,7 +406,7 @@ func newContextWatcher(conn net.Conn) *ctxwatch.ContextWatcher {
|
||||
)
|
||||
}
|
||||
|
||||
func startTLS(conn *nbconn.NetConn, tlsConfig *tls.Config) (*nbconn.TLSConn, error) {
|
||||
func startTLS(conn net.Conn, tlsConfig *tls.Config) (net.Conn, error) {
|
||||
err := binary.Write(conn, binary.BigEndian, []int32{8, 80877103})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -398,17 +421,12 @@ func startTLS(conn *nbconn.NetConn, tlsConfig *tls.Config) (*nbconn.TLSConn, err
|
||||
return nil, errors.New("server refused TLS connection")
|
||||
}
|
||||
|
||||
tlsConn, err := nbconn.TLSClient(conn, tlsConfig)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return tlsConn, nil
|
||||
return tls.Client(conn, tlsConfig), nil
|
||||
}
|
||||
|
||||
func (pgConn *PgConn) txPasswordMessage(password string) (err error) {
|
||||
pgConn.frontend.Send(&pgproto3.PasswordMessage{Password: password})
|
||||
return pgConn.frontend.Flush()
|
||||
return pgConn.flushWithPotentialWriteReadDeadlock()
|
||||
}
|
||||
|
||||
func hexMD5(s string) string {
|
||||
@@ -417,6 +435,24 @@ func hexMD5(s string) string {
|
||||
return hex.EncodeToString(hash.Sum(nil))
|
||||
}
|
||||
|
||||
func (pgConn *PgConn) signalMessage() chan struct{} {
|
||||
if pgConn.bufferingReceive {
|
||||
panic("BUG: signalMessage when already in progress")
|
||||
}
|
||||
|
||||
pgConn.bufferingReceive = true
|
||||
pgConn.bufferingReceiveMux.Lock()
|
||||
|
||||
ch := make(chan struct{})
|
||||
go func() {
|
||||
pgConn.bufferingReceiveMsg, pgConn.bufferingReceiveErr = pgConn.frontend.Receive()
|
||||
pgConn.bufferingReceiveMux.Unlock()
|
||||
close(ch)
|
||||
}()
|
||||
|
||||
return ch
|
||||
}
|
||||
|
||||
// ReceiveMessage receives one wire protocol message from the PostgreSQL server. It must only be used when the
|
||||
// connection is not busy. e.g. It is an error to call ReceiveMessage while reading the result of a query. The messages
|
||||
// are still handled by the core pgconn message handling system so receiving a NotificationResponse will still trigger
|
||||
@@ -445,7 +481,8 @@ func (pgConn *PgConn) ReceiveMessage(ctx context.Context) (pgproto3.BackendMessa
|
||||
err = &pgconnError{
|
||||
msg: "receive message failed",
|
||||
err: normalizeTimeoutError(ctx, err),
|
||||
safeToRetry: true}
|
||||
safeToRetry: true,
|
||||
}
|
||||
}
|
||||
return msg, err
|
||||
}
|
||||
@@ -456,13 +493,25 @@ func (pgConn *PgConn) peekMessage() (pgproto3.BackendMessage, error) {
|
||||
return pgConn.peekedMsg, nil
|
||||
}
|
||||
|
||||
msg, err := pgConn.frontend.Receive()
|
||||
var msg pgproto3.BackendMessage
|
||||
var err error
|
||||
if pgConn.bufferingReceive {
|
||||
pgConn.bufferingReceiveMux.Lock()
|
||||
msg = pgConn.bufferingReceiveMsg
|
||||
err = pgConn.bufferingReceiveErr
|
||||
pgConn.bufferingReceiveMux.Unlock()
|
||||
pgConn.bufferingReceive = false
|
||||
|
||||
// If a timeout error happened in the background try the read again.
|
||||
var netErr net.Error
|
||||
if errors.As(err, &netErr) && netErr.Timeout() {
|
||||
msg, err = pgConn.frontend.Receive()
|
||||
}
|
||||
} else {
|
||||
msg, err = pgConn.frontend.Receive()
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
if errors.Is(err, nbconn.ErrWouldBlock) {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Close on anything other than timeout error - everything else is fatal
|
||||
var netErr net.Error
|
||||
isNetErr := errors.As(err, &netErr)
|
||||
@@ -510,7 +559,8 @@ func (pgConn *PgConn) receiveMessage() (pgproto3.BackendMessage, error) {
|
||||
return msg, nil
|
||||
}
|
||||
|
||||
// Conn returns the underlying net.Conn. This rarely necessary.
|
||||
// Conn returns the underlying net.Conn. This rarely necessary. If the connection will be directly used for reading or
|
||||
// writing then SyncConn should usually be called before Conn.
|
||||
func (pgConn *PgConn) Conn() net.Conn {
|
||||
return pgConn.conn
|
||||
}
|
||||
@@ -573,7 +623,7 @@ func (pgConn *PgConn) Close(ctx context.Context) error {
|
||||
//
|
||||
// See https://github.com/jackc/pgx/issues/637
|
||||
pgConn.frontend.Send(&pgproto3.Terminate{})
|
||||
pgConn.frontend.Flush()
|
||||
pgConn.flushWithPotentialWriteReadDeadlock()
|
||||
|
||||
return pgConn.conn.Close()
|
||||
}
|
||||
@@ -600,7 +650,7 @@ func (pgConn *PgConn) asyncClose() {
|
||||
pgConn.conn.SetDeadline(deadline)
|
||||
|
||||
pgConn.frontend.Send(&pgproto3.Terminate{})
|
||||
pgConn.frontend.Flush()
|
||||
pgConn.flushWithPotentialWriteReadDeadlock()
|
||||
}()
|
||||
}
|
||||
|
||||
@@ -775,7 +825,7 @@ func (pgConn *PgConn) Prepare(ctx context.Context, name, sql string, paramOIDs [
|
||||
pgConn.frontend.SendParse(&pgproto3.Parse{Name: name, Query: sql, ParameterOIDs: paramOIDs})
|
||||
pgConn.frontend.SendDescribe(&pgproto3.Describe{ObjectType: 'S', Name: name})
|
||||
pgConn.frontend.SendSync(&pgproto3.Sync{})
|
||||
err := pgConn.frontend.Flush()
|
||||
err := pgConn.flushWithPotentialWriteReadDeadlock()
|
||||
if err != nil {
|
||||
pgConn.asyncClose()
|
||||
return nil, err
|
||||
@@ -848,9 +898,28 @@ func (pgConn *PgConn) CancelRequest(ctx context.Context) error {
|
||||
// the connection config. This is important in high availability configurations where fallback connections may be
|
||||
// specified or DNS may be used to load balance.
|
||||
serverAddr := pgConn.conn.RemoteAddr()
|
||||
cancelConn, err := pgConn.config.DialFunc(ctx, serverAddr.Network(), serverAddr.String())
|
||||
var serverNetwork string
|
||||
var serverAddress string
|
||||
if serverAddr.Network() == "unix" {
|
||||
// for unix sockets, RemoteAddr() calls getpeername() which returns the name the
|
||||
// server passed to bind(). For Postgres, this is always a relative path "./.s.PGSQL.5432"
|
||||
// so connecting to it will fail. Fall back to the config's value
|
||||
serverNetwork, serverAddress = NetworkAddress(pgConn.config.Host, pgConn.config.Port)
|
||||
} else {
|
||||
serverNetwork, serverAddress = serverAddr.Network(), serverAddr.String()
|
||||
}
|
||||
cancelConn, err := pgConn.config.DialFunc(ctx, serverNetwork, serverAddress)
|
||||
if err != nil {
|
||||
return err
|
||||
// In case of unix sockets, RemoteAddr() returns only the file part of the path. If the
|
||||
// first connect failed, try the config.
|
||||
if serverAddr.Network() != "unix" {
|
||||
return err
|
||||
}
|
||||
serverNetwork, serverAddr := NetworkAddress(pgConn.config.Host, pgConn.config.Port)
|
||||
cancelConn, err = pgConn.config.DialFunc(ctx, serverNetwork, serverAddr)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
defer cancelConn.Close()
|
||||
|
||||
@@ -868,17 +937,11 @@ func (pgConn *PgConn) CancelRequest(ctx context.Context) error {
|
||||
binary.BigEndian.PutUint32(buf[4:8], 80877102)
|
||||
binary.BigEndian.PutUint32(buf[8:12], uint32(pgConn.pid))
|
||||
binary.BigEndian.PutUint32(buf[12:16], uint32(pgConn.secretKey))
|
||||
// Postgres will process the request and close the connection
|
||||
// so when don't need to read the reply
|
||||
// https://www.postgresql.org/docs/current/protocol-flow.html#id-1.10.6.7.10
|
||||
_, err = cancelConn.Write(buf)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
_, err = cancelConn.Read(buf)
|
||||
if err != io.EOF {
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
return err
|
||||
}
|
||||
|
||||
// WaitForNotification waits for a LISTON/NOTIFY message to be received. It returns an error if a notification was not
|
||||
@@ -944,7 +1007,7 @@ func (pgConn *PgConn) Exec(ctx context.Context, sql string) *MultiResultReader {
|
||||
}
|
||||
|
||||
pgConn.frontend.SendQuery(&pgproto3.Query{String: sql})
|
||||
err := pgConn.frontend.Flush()
|
||||
err := pgConn.flushWithPotentialWriteReadDeadlock()
|
||||
if err != nil {
|
||||
pgConn.asyncClose()
|
||||
pgConn.contextWatcher.Unwatch()
|
||||
@@ -1055,7 +1118,7 @@ func (pgConn *PgConn) execExtendedSuffix(result *ResultReader) {
|
||||
pgConn.frontend.SendExecute(&pgproto3.Execute{})
|
||||
pgConn.frontend.SendSync(&pgproto3.Sync{})
|
||||
|
||||
err := pgConn.frontend.Flush()
|
||||
err := pgConn.flushWithPotentialWriteReadDeadlock()
|
||||
if err != nil {
|
||||
pgConn.asyncClose()
|
||||
result.concludeCommand(CommandTag{}, err)
|
||||
@@ -1088,7 +1151,7 @@ func (pgConn *PgConn) CopyTo(ctx context.Context, w io.Writer, sql string) (Comm
|
||||
// Send copy to command
|
||||
pgConn.frontend.SendQuery(&pgproto3.Query{String: sql})
|
||||
|
||||
err := pgConn.frontend.Flush()
|
||||
err := pgConn.flushWithPotentialWriteReadDeadlock()
|
||||
if err != nil {
|
||||
pgConn.asyncClose()
|
||||
pgConn.unlock()
|
||||
@@ -1144,85 +1207,91 @@ func (pgConn *PgConn) CopyFrom(ctx context.Context, r io.Reader, sql string) (Co
|
||||
defer pgConn.contextWatcher.Unwatch()
|
||||
}
|
||||
|
||||
// Send copy to command
|
||||
// Send copy from query
|
||||
pgConn.frontend.SendQuery(&pgproto3.Query{String: sql})
|
||||
err := pgConn.frontend.Flush()
|
||||
err := pgConn.flushWithPotentialWriteReadDeadlock()
|
||||
if err != nil {
|
||||
pgConn.asyncClose()
|
||||
return CommandTag{}, err
|
||||
}
|
||||
|
||||
err = pgConn.conn.SetReadDeadline(nbconn.NonBlockingDeadline)
|
||||
if err != nil {
|
||||
pgConn.asyncClose()
|
||||
return CommandTag{}, err
|
||||
}
|
||||
nonblocking := true
|
||||
defer func() {
|
||||
if nonblocking {
|
||||
pgConn.conn.SetReadDeadline(time.Time{})
|
||||
// Send copy data
|
||||
abortCopyChan := make(chan struct{})
|
||||
copyErrChan := make(chan error, 1)
|
||||
signalMessageChan := pgConn.signalMessage()
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(1)
|
||||
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
buf := iobufpool.Get(65536)
|
||||
defer iobufpool.Put(buf)
|
||||
(*buf)[0] = 'd'
|
||||
|
||||
for {
|
||||
n, readErr := r.Read((*buf)[5:cap(*buf)])
|
||||
if n > 0 {
|
||||
*buf = (*buf)[0 : n+5]
|
||||
pgio.SetInt32((*buf)[1:], int32(n+4))
|
||||
|
||||
writeErr := pgConn.frontend.SendUnbufferedEncodedCopyData(*buf)
|
||||
if writeErr != nil {
|
||||
// Write errors are always fatal, but we can't use asyncClose because we are in a different goroutine. Not
|
||||
// setting pgConn.status or closing pgConn.cleanupDone for the same reason.
|
||||
pgConn.conn.Close()
|
||||
|
||||
copyErrChan <- writeErr
|
||||
return
|
||||
}
|
||||
}
|
||||
if readErr != nil {
|
||||
copyErrChan <- readErr
|
||||
return
|
||||
}
|
||||
|
||||
select {
|
||||
case <-abortCopyChan:
|
||||
return
|
||||
default:
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
buf := iobufpool.Get(65536)
|
||||
defer iobufpool.Put(buf)
|
||||
buf[0] = 'd'
|
||||
|
||||
var readErr, pgErr error
|
||||
for pgErr == nil {
|
||||
// Read chunk from r.
|
||||
var n int
|
||||
n, readErr = r.Read(buf[5:cap(buf)])
|
||||
|
||||
// Send chunk to PostgreSQL.
|
||||
if n > 0 {
|
||||
buf = buf[0 : n+5]
|
||||
pgio.SetInt32(buf[1:], int32(n+4))
|
||||
|
||||
writeErr := pgConn.frontend.SendUnbufferedEncodedCopyData(buf)
|
||||
if writeErr != nil {
|
||||
pgConn.asyncClose()
|
||||
return CommandTag{}, err
|
||||
}
|
||||
}
|
||||
|
||||
// Abort loop if there was a read error.
|
||||
if readErr != nil {
|
||||
break
|
||||
}
|
||||
|
||||
// Read messages until error or none available.
|
||||
for pgErr == nil {
|
||||
msg, err := pgConn.receiveMessage()
|
||||
if err != nil {
|
||||
if errors.Is(err, nbconn.ErrWouldBlock) {
|
||||
break
|
||||
}
|
||||
pgConn.asyncClose()
|
||||
var pgErr error
|
||||
var copyErr error
|
||||
for copyErr == nil && pgErr == nil {
|
||||
select {
|
||||
case copyErr = <-copyErrChan:
|
||||
case <-signalMessageChan:
|
||||
// If pgConn.receiveMessage encounters an error it will call pgConn.asyncClose. But that is a race condition with
|
||||
// the goroutine. So instead check pgConn.bufferingReceiveErr which will have been set by the signalMessage. If an
|
||||
// error is found then forcibly close the connection without sending the Terminate message.
|
||||
if err := pgConn.bufferingReceiveErr; err != nil {
|
||||
pgConn.status = connStatusClosed
|
||||
pgConn.conn.Close()
|
||||
close(pgConn.cleanupDone)
|
||||
return CommandTag{}, normalizeTimeoutError(ctx, err)
|
||||
}
|
||||
msg, _ := pgConn.receiveMessage()
|
||||
|
||||
switch msg := msg.(type) {
|
||||
case *pgproto3.ErrorResponse:
|
||||
pgErr = ErrorResponseToPgError(msg)
|
||||
break
|
||||
default:
|
||||
signalMessageChan = pgConn.signalMessage()
|
||||
}
|
||||
}
|
||||
}
|
||||
close(abortCopyChan)
|
||||
// Make sure io goroutine finishes before writing.
|
||||
wg.Wait()
|
||||
|
||||
err = pgConn.conn.SetReadDeadline(time.Time{})
|
||||
if err != nil {
|
||||
pgConn.asyncClose()
|
||||
return CommandTag{}, err
|
||||
}
|
||||
nonblocking = false
|
||||
|
||||
if readErr == io.EOF || pgErr != nil {
|
||||
if copyErr == io.EOF || pgErr != nil {
|
||||
pgConn.frontend.Send(&pgproto3.CopyDone{})
|
||||
} else {
|
||||
pgConn.frontend.Send(&pgproto3.CopyFail{Message: readErr.Error()})
|
||||
pgConn.frontend.Send(&pgproto3.CopyFail{Message: copyErr.Error()})
|
||||
}
|
||||
err = pgConn.frontend.Flush()
|
||||
err = pgConn.flushWithPotentialWriteReadDeadlock()
|
||||
if err != nil {
|
||||
pgConn.asyncClose()
|
||||
return CommandTag{}, err
|
||||
@@ -1274,7 +1343,6 @@ func (mrr *MultiResultReader) ReadAll() ([]*Result, error) {
|
||||
|
||||
func (mrr *MultiResultReader) receiveMessage() (pgproto3.BackendMessage, error) {
|
||||
msg, err := mrr.pgConn.receiveMessage()
|
||||
|
||||
if err != nil {
|
||||
mrr.pgConn.contextWatcher.Unwatch()
|
||||
mrr.err = normalizeTimeoutError(mrr.ctx, err)
|
||||
@@ -1417,7 +1485,8 @@ func (rr *ResultReader) NextRow() bool {
|
||||
}
|
||||
|
||||
// FieldDescriptions returns the field descriptions for the current result set. The returned slice is only valid until
|
||||
// the ResultReader is closed.
|
||||
// the ResultReader is closed. It may return nil (for example, if the query did not return a result set or an error was
|
||||
// encountered.)
|
||||
func (rr *ResultReader) FieldDescriptions() []FieldDescription {
|
||||
return rr.fieldDescriptions
|
||||
}
|
||||
@@ -1583,6 +1652,8 @@ func (pgConn *PgConn) ExecBatch(ctx context.Context, batch *Batch) *MultiResultR
|
||||
|
||||
batch.buf = (&pgproto3.Sync{}).Encode(batch.buf)
|
||||
|
||||
pgConn.enterPotentialWriteReadDeadlock()
|
||||
defer pgConn.exitPotentialWriteReadDeadlock()
|
||||
_, err := pgConn.conn.Write(batch.buf)
|
||||
if err != nil {
|
||||
multiResult.closed = true
|
||||
@@ -1611,29 +1682,99 @@ func (pgConn *PgConn) EscapeString(s string) (string, error) {
|
||||
return strings.Replace(s, "'", "''", -1), nil
|
||||
}
|
||||
|
||||
// CheckConn checks the underlying connection without writing any bytes. This is currently implemented by reading and
|
||||
// buffering until the read would block or an error occurs. This can be used to check if the server has closed the
|
||||
// connection. If this is done immediately before sending a query it reduces the chances a query will be sent that fails
|
||||
// CheckConn checks the underlying connection without writing any bytes. This is currently implemented by doing a read
|
||||
// with a very short deadline. This can be useful because a TCP connection can be broken such that a write will appear
|
||||
// to succeed even though it will never actually reach the server. Reading immediately before a write will detect this
|
||||
// condition. If this is done immediately before sending a query it reduces the chances a query will be sent that fails
|
||||
// without the client knowing whether the server received it or not.
|
||||
//
|
||||
// Deprecated: CheckConn is deprecated in favor of Ping. CheckConn cannot detect all types of broken connections where
|
||||
// the write would still appear to succeed. Prefer Ping unless on a high latency connection.
|
||||
func (pgConn *PgConn) CheckConn() error {
|
||||
err := pgConn.conn.BufferReadUntilBlock()
|
||||
if err != nil && !errors.Is(err, nbconn.ErrWouldBlock) {
|
||||
return err
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 1*time.Millisecond)
|
||||
defer cancel()
|
||||
|
||||
_, err := pgConn.ReceiveMessage(ctx)
|
||||
if err != nil {
|
||||
if !Timeout(err) {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Ping pings the server. This can be useful because a TCP connection can be broken such that a write will appear to
|
||||
// succeed even though it will never actually reach the server. Pinging immediately before sending a query reduces the
|
||||
// chances a query will be sent that fails without the client knowing whether the server received it or not.
|
||||
func (pgConn *PgConn) Ping(ctx context.Context) error {
|
||||
return pgConn.Exec(ctx, "-- ping").Close()
|
||||
}
|
||||
|
||||
// makeCommandTag makes a CommandTag. It does not retain a reference to buf or buf's underlying memory.
|
||||
func (pgConn *PgConn) makeCommandTag(buf []byte) CommandTag {
|
||||
return CommandTag{s: string(buf)}
|
||||
}
|
||||
|
||||
// enterPotentialWriteReadDeadlock must be called before a write that could deadlock if the server is simultaneously
|
||||
// blocked writing to us.
|
||||
func (pgConn *PgConn) enterPotentialWriteReadDeadlock() {
|
||||
// The time to wait is somewhat arbitrary. A Write should only take as long as the syscall and memcpy to the OS
|
||||
// outbound network buffer unless the buffer is full (which potentially is a block). It needs to be long enough for
|
||||
// the normal case, but short enough not to kill performance if a block occurs.
|
||||
//
|
||||
// In addition, on Windows the default timer resolution is 15.6ms. So setting the timer to less than that is
|
||||
// ineffective.
|
||||
if pgConn.slowWriteTimer.Reset(15 * time.Millisecond) {
|
||||
panic("BUG: slow write timer already active")
|
||||
}
|
||||
}
|
||||
|
||||
// exitPotentialWriteReadDeadlock must be called after a call to enterPotentialWriteReadDeadlock.
|
||||
func (pgConn *PgConn) exitPotentialWriteReadDeadlock() {
|
||||
// The state of the timer is not relevant upon exiting the potential slow write. It may both
|
||||
// fire (due to a slow write), or not fire (due to a fast write).
|
||||
_ = pgConn.slowWriteTimer.Stop()
|
||||
pgConn.bgReader.Stop()
|
||||
}
|
||||
|
||||
func (pgConn *PgConn) flushWithPotentialWriteReadDeadlock() error {
|
||||
pgConn.enterPotentialWriteReadDeadlock()
|
||||
defer pgConn.exitPotentialWriteReadDeadlock()
|
||||
err := pgConn.frontend.Flush()
|
||||
return err
|
||||
}
|
||||
|
||||
// SyncConn prepares the underlying net.Conn for direct use. PgConn may internally buffer reads or use goroutines for
|
||||
// background IO. This means that any direct use of the underlying net.Conn may be corrupted if a read is already
|
||||
// buffered or a read is in progress. SyncConn drains read buffers and stops background IO. In some cases this may
|
||||
// require sending a ping to the server. ctx can be used to cancel this operation. This should be called before any
|
||||
// operation that will use the underlying net.Conn directly. e.g. Before Conn() or Hijack().
|
||||
//
|
||||
// This should not be confused with the PostgreSQL protocol Sync message.
|
||||
func (pgConn *PgConn) SyncConn(ctx context.Context) error {
|
||||
for i := 0; i < 10; i++ {
|
||||
if pgConn.bgReader.Status() == bgreader.StatusStopped && pgConn.frontend.ReadBufferLen() == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
err := pgConn.Ping(ctx)
|
||||
if err != nil {
|
||||
return fmt.Errorf("SyncConn: Ping failed while syncing conn: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
// This should never happen. Only way I can imagine this occuring is if the server is constantly sending data such as
|
||||
// LISTEN/NOTIFY or log notifications such that we never can get an empty buffer.
|
||||
return errors.New("SyncConn: conn never synchronized")
|
||||
}
|
||||
|
||||
// HijackedConn is the result of hijacking a connection.
|
||||
//
|
||||
// Due to the necessary exposure of internal implementation details, it is not covered by the semantic versioning
|
||||
// compatibility.
|
||||
type HijackedConn struct {
|
||||
Conn nbconn.Conn // the non-blocking wrapper of the underlying TCP or unix domain socket connection
|
||||
Conn net.Conn
|
||||
PID uint32 // backend pid
|
||||
SecretKey uint32 // key to use to send a cancel query message to the server
|
||||
ParameterStatuses map[string]string // parameters that have been reported by the server
|
||||
@@ -1642,9 +1783,9 @@ type HijackedConn struct {
|
||||
Config *Config
|
||||
}
|
||||
|
||||
// Hijack extracts the internal connection data. pgConn must be in an idle state. pgConn is unusable after hijacking.
|
||||
// Hijacking is typically only useful when using pgconn to establish a connection, but taking complete control of the
|
||||
// raw connection after that (e.g. a load balancer or proxy).
|
||||
// Hijack extracts the internal connection data. pgConn must be in an idle state. SyncConn should be called immediately
|
||||
// before Hijack. pgConn is unusable after hijacking. Hijacking is typically only useful when using pgconn to establish
|
||||
// a connection, but taking complete control of the raw connection after that (e.g. a load balancer or proxy).
|
||||
//
|
||||
// Due to the necessary exposure of internal implementation details, it is not covered by the semantic versioning
|
||||
// compatibility.
|
||||
@@ -1668,6 +1809,8 @@ func (pgConn *PgConn) Hijack() (*HijackedConn, error) {
|
||||
// Construct created a PgConn from an already established connection to a PostgreSQL server. This is the inverse of
|
||||
// PgConn.Hijack. The connection must be in an idle state.
|
||||
//
|
||||
// hc.Frontend is replaced by a new pgproto3.Frontend built by hc.Config.BuildFrontend.
|
||||
//
|
||||
// Due to the necessary exposure of internal implementation details, it is not covered by the semantic versioning
|
||||
// compatibility.
|
||||
func Construct(hc *HijackedConn) (*PgConn, error) {
|
||||
@@ -1686,6 +1829,10 @@ func Construct(hc *HijackedConn) (*PgConn, error) {
|
||||
}
|
||||
|
||||
pgConn.contextWatcher = newContextWatcher(pgConn.conn)
|
||||
pgConn.bgReader = bgreader.New(pgConn.conn)
|
||||
pgConn.slowWriteTimer = time.AfterFunc(time.Duration(math.MaxInt64), pgConn.bgReader.Start)
|
||||
pgConn.slowWriteTimer.Stop()
|
||||
pgConn.frontend = hc.Config.BuildFrontend(pgConn.bgReader, pgConn.conn)
|
||||
|
||||
return pgConn, nil
|
||||
}
|
||||
@@ -1808,7 +1955,7 @@ func (p *Pipeline) Flush() error {
|
||||
return errors.New("pipeline closed")
|
||||
}
|
||||
|
||||
err := p.conn.frontend.Flush()
|
||||
err := p.conn.flushWithPotentialWriteReadDeadlock()
|
||||
if err != nil {
|
||||
err = normalizeTimeoutError(p.ctx, err)
|
||||
|
||||
@@ -1887,7 +2034,6 @@ func (p *Pipeline) GetResults() (results any, err error) {
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
func (p *Pipeline) getResultsPrepare() (*StatementDescription, error) {
|
||||
|
||||
Reference in New Issue
Block a user