/
tx.go
108 lines (91 loc) · 2.63 KB
/
tx.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
package crdb
import (
"context"
"errors"
"fmt"
"github.com/jackc/pgconn"
"github.com/jackc/pgx/v4"
"github.com/prometheus/client_golang/prometheus"
"github.com/rs/zerolog/log"
)
const (
crdbRetryErrCode = "40001"
errUnableToRetry = "failed to retry conflicted transaction: %w"
errReachedMaxRetry = "maximum retries reached"
)
var retryHistogram = prometheus.NewHistogram(prometheus.HistogramOpts{
Name: "crdb_client_retries",
Help: "cockroachdb client-side retry distribution",
Buckets: []float64{0, 1, 2, 5, 10, 20, 50},
})
func init() {
prometheus.MustRegister(retryHistogram)
}
// conn is satisfied by both pgx.conn and pgxpool.Pool.
type conn interface {
Begin(context.Context) (pgx.Tx, error)
BeginTx(context.Context, pgx.TxOptions) (pgx.Tx, error)
}
type transactionFn func(tx pgx.Tx) error
type executeTxRetryFunc func(context.Context, conn, pgx.TxOptions, transactionFn) error
func executeWithMaxRetries(max int) executeTxRetryFunc {
return func(ctx context.Context, conn conn, txOptions pgx.TxOptions, fn transactionFn) (err error) {
return execute(ctx, conn, txOptions, fn, max)
}
}
// adapted from https://github.com/cockroachdb/cockroach-go
func execute(ctx context.Context, conn conn, txOptions pgx.TxOptions, fn transactionFn, maxRetries int) (err error) {
var tx pgx.Tx
tx, err = conn.BeginTx(ctx, txOptions)
if err != nil {
return err
}
defer func() {
if err == nil {
_ = tx.Commit(ctx)
return
}
_ = tx.Rollback(ctx)
}()
if _, err = tx.Exec(ctx, "SAVEPOINT cockroach_restart"); err != nil {
return
}
var i int
defer func() {
retryHistogram.Observe(float64(i))
}()
releasedFn := func(tx pgx.Tx) error {
if err := fn(tx); err != nil {
return err
}
// RELEASE acts like COMMIT in CockroachDB. We use it since it gives us an
// opportunity to react to retryable errors, whereas tx.Commit() doesn't.
// RELEASE SAVEPOINT itself can fail, in which case the entire
// transaction needs to be retried
if _, err := tx.Exec(ctx, "RELEASE SAVEPOINT cockroach_restart"); err != nil {
return err
}
return nil
}
for i = 0; i < maxRetries; i++ {
if err = releasedFn(tx); err != nil {
if !retriable(ctx, err) {
return err
}
if _, retryErr := tx.Exec(ctx, "ROLLBACK TO SAVEPOINT cockroach_restart"); retryErr != nil {
return fmt.Errorf(errUnableToRetry, err)
}
continue
}
return nil
}
return errors.New(errReachedMaxRetry)
}
func retriable(ctx context.Context, err error) bool {
var pgerr *pgconn.PgError
if !errors.As(err, &pgerr) {
log.Ctx(ctx).Error().Err(err).Msg("error not retriable")
return false
}
return pgerr.SQLState() == crdbRetryErrCode
}