driver.ErrBadConnでのリトライ

データベースの再起動などで切断されたコネクションをコネクションプールから引き当ててエラーになる問題について、SetConnMaxLifetime()を設定して、定期的にコネクションをリフレッシュするしかないと思っていたのですが、こちらの記事でdriver.ErrBadConnのときにリトライしてくれることを知りました。

たしかにドキュメントには

ErrBadConn should be returned by a driver to signal to the database/sql package that a driver.Conn is in a bad state (such as the server having earlier closed the connection) and the database/sql package should retry on a new connection.

と書いてあり、database/sqlのコードを読むとfunc (db *DB) retry()でリトライ処理を行っています。

// go/src/database/sql/sql.go

// maxBadConnRetries is the number of maximum retries if the driver returns
// driver.ErrBadConn to signal a broken connection before forcing a new
// connection to be opened.
const maxBadConnRetries = 2

func (db *DB) retry(fn func(strategy connReuseStrategy) error) error {
    for i := int64(0); i < maxBadConnRetries; i++ {
        err := fn(cachedOrNewConn)
        // retry if err is driver.ErrBadConn
        if err == nil || !errors.Is(err, driver.ErrBadConn) {
            return err
        }
    }

    return fn(alwaysNewConn)
}

pgxのコードではSafeToRetry()がtrueを返すときにErrBadConnを返していました。特定のエラーや、エラー発生時にデータ送信がなかった場合などにリトライが許可されるようです。

cf. https://github.com/search?q=repo%3Ajackc%2Fpgx%20SafeToRetry&type=code

// pgx/stdlib/sql.go

    if err != nil {
        if pgconn.SafeToRetry(err) {
            return nil, driver.ErrBadConn
        }
        return nil, err
    }

以下のコードで動作を確認してみました。

package main

import (
    "database/sql"
    "fmt"
    "net/url"
    "time"

    _ "github.com/jackc/pgx/v5/stdlib"
    "github.com/mattn/go-tty"
)

func main() {
    url := &url.URL{
        Scheme: "postgres",
        User:   url.UserPassword("postgres", "xxx"),
        Host:   "xxx.ap-northeast-1.rds.amazonaws.com:5432",
        Path:   "postgres",
    }

    db, err := sql.Open("pgx", url.String())

    if err != nil {
        panic(err)
    }

    defer db.Close()
    db.SetConnMaxLifetime(0)
    db.SetConnMaxIdleTime(0)
    db.SetMaxIdleConns(1)
    db.SetMaxOpenConns(1)

    tty, err := tty.Open()

    if err != nil {
        panic(err)
    }

    defer tty.Close()

    for {
        // キー入力を待つ
        tty.ReadRune()

        var n int
        err = db.QueryRow("select 1").Scan(&n)

        if err != nil {
            fmt.Println(err)
            continue
        }

        fmt.Printf("select 1 => %d\n", n)
    }
}

// database/sql/sql.go

func (db *DB) retry(fn func(strategy connReuseStrategy) error) error {
    for i := int64(0); i < maxBadConnRetries; i++ {
        err := fn(cachedOrNewConn)
        // retry if err is driver.ErrBadConn
        if err == nil || !errors.Is(err, driver.ErrBadConn) {
            return err
        }
        // リトライ時の出力を追加
        fmt.Printf("[INFO] retried with error: %s\n", err)
    }

    return fn(alwaysNewConn)
}

キー入力でselect 1を実行しながら途中でデータベースを再起動してみると、リトライされていることが確認できました。

select 1 => 1
select 1 => 1
select 1 => 1
# ここでデータベースを再起動
[INFO] retried with error: driver: bad connection
select 1 => 1

target_session_attrs

こちらは別のブログ記事で知ったのですが、go-sql-driver/mysqlにはrejectReadOnlyというパラメーターがあり、Auroraがフェイルオーバーした際に降格したreaderノードに書き込みを行う問題を回避できるようになっていました。

pgxでも同様の機能がないか調べたところtarget_session_attrsというパラメーターで接続するノードの種別を指定できるようになっていました。

cf. https://github.com/jackc/pgx/blob/70f7cad2226dc12406b105f8bb5be9c62780aaf7/pgconn/config.go#L402-L417

   switch tsa := settings["target_session_attrs"]; tsa {
    case "read-write":
        config.ValidateConnect = ValidateConnectTargetSessionAttrsReadWrite
    case "read-only":
        config.ValidateConnect = ValidateConnectTargetSessionAttrsReadOnly
    case "primary":
        config.ValidateConnect = ValidateConnectTargetSessionAttrsPrimary
    case "standby":
        config.ValidateConnect = ValidateConnectTargetSessionAttrsStandby
    case "prefer-standby":
        config.ValidateConnect = ValidateConnectTargetSessionAttrsPreferStandby
    case "any":
        // do nothing
    default:
        return nil, &ParseConfigError{ConnString: connString, msg: fmt.Sprintf("unknown target_session_attrs value: %v", tsa)}
    }

libpqにある機能ですがpgxも独自に実装しているようです。

// ValidateConnectTargetSessionAttrsReadWrite is a ValidateConnectFunc that implements libpq compatible
// target_session_attrs=read-write.
func ValidateConnectTargetSessionAttrsReadWrite(ctx context.Context, pgConn *PgConn) error {

以下のコードで動作を確認してみました。

package main

import (
    "database/sql"
    "fmt"
    "net/url"
    "time"

    _ "github.com/jackc/pgx/v5/stdlib"
)

func main() {
    params := url.Values{}
    // params.Add("target_session_attrs", "read-write")

    url := &url.URL{
        Scheme:   "postgres",
        User:     url.UserPassword("postgres", "xxx"),
        Host:     "xxx.ap-northeast-1.rds.amazonaws.com:5432",
        Path:     "postgres",
        RawQuery: params.Encode(),
    }

    db, err := sql.Open("pgx", url.String())

    if err != nil {
        panic(err)
    }

    defer db.Close()
    db.SetConnMaxLifetime(0)
    db.SetConnMaxIdleTime(0)
    db.SetMaxIdleConns(1)
    db.SetMaxOpenConns(1)

    for {
        time.Sleep(1 * time.Second)

        r, err := db.Exec("insert into test values ($1)", time.Now().String())

        if err != nil {
            fmt.Println(err)
            continue
        }

        n, _ := r.RowsAffected()
        fmt.Printf("RowsAffected: %d\n", n)
    }
}

target_session_attrsを設定しないコードを動かしてフェイルオーバーを行うと、切り替え後にERROR: cannot execute INSERT in a read-only transactionが発生しつづけてしまいます。

RowsAffected: 1
RowsAffected: 1
RowsAffected: 1
unexpected EOF
failed to connect to `user=postgres database=postgres`:
    xxx.xxx.xxx.xxx:5432 (xxx.ap-northeast-1.rds.amazonaws.com): dial error: dial tcp xxx.xxx.xxx.xxx:5432: connect: connection refused
    xxx.xxx.xxx.xxx:5432 (xxx.ap-northeast-1.rds.amazonaws.com): dial error: dial tcp xxx.xxx.xxx.xxx:5432: connect: connection refused
...
ERROR: cannot execute INSERT in a read-only transaction (SQLSTATE 25006)
ERROR: cannot execute INSERT in a read-only transaction (SQLSTATE 25006)
ERROR: cannot execute INSERT in a read-only transaction (SQLSTATE 25006)
...

target_session_attrs=read-writeを設定した場合には、切り替え後に検証が行われ書き込み可能なコネクションに接続することを確認できました。

RowsAffected: 1
RowsAffected: 1
RowsAffected: 1
unexpected EOF
failed to connect to `user=postgres database=postgres`:
    xxx.xxx.xxx.xxx:5432 (xxx.ap-northeast-1.rds.amazonaws.com): dial error: dial tcp xxx.xxx.xxx.xxx:5432: connect: connection refused
    xxx.xxx.xxx.xxx:5432 (xxx.ap-northeast-1.rds.amazonaws.com): dial error: dial tcp xxx.xxx.xxx.xxx:5432: connect: connection refused
...
failed to connect to `user=postgres database=postgres`:
    xxx.xxx.xxx.xxx:5432 (xxx.ap-northeast-1.rds.amazonaws.com): ValidateConnect failed: read only connection
    xxx.xxx.xxx.xxx:5432 (xxx.ap-northeast-1.rds.amazonaws.com): ValidateConnect failed: read only connection
RowsAffected: 1
RowsAffected: 1
RowsAffected: 1
...