diff --git a/crdb/common.go b/crdb/common.go index 8039923..7beefa3 100644 --- a/crdb/common.go +++ b/crdb/common.go @@ -62,8 +62,7 @@ func ExecuteInTx(ctx context.Context, tx Tx, fn func() error) (err error) { return err } - // TODO(rafi): make the maxRetryCount configurable. Maybe pass it in the context?) - const maxRetries = 50 + maxRetries := numRetriesFromContext(ctx) retryCount := 0 for { releaseFailed := false @@ -90,7 +89,7 @@ func ExecuteInTx(ctx context.Context, tx Tx, fn func() error) (err error) { } retryCount++ - if retryCount > maxRetries { + if maxRetries > 0 && retryCount > maxRetries { return newMaxRetriesExceededError(err, maxRetries) } } diff --git a/crdb/tx.go b/crdb/tx.go index dffa72d..a957f28 100644 --- a/crdb/tx.go +++ b/crdb/tx.go @@ -101,6 +101,26 @@ func Execute(fn func() error) (err error) { } } +type txConfigKey struct {} + +// WithMaxRetries configures context so that ExecuteTx retries tx specified +// number of times when encountering retryable errors. +// Setting retries to 0 will retry indefinitely. +func WithMaxRetries(ctx context.Context, retries int) context.Context { + return context.WithValue(ctx, txConfigKey{}, retries) +} + +func numRetriesFromContext(ctx context.Context) int { + const defaultRetries = 50 + + if v := ctx.Value(txConfigKey{}); v != nil { + if retries, ok := v.(int); ok && retries >= 0 { + return retries + } + } + return defaultRetries +} + // ExecuteTx runs fn inside a transaction and retries it as needed. On // non-retryable failures, the transaction is aborted and rolled back; on // success, the transaction is committed.