Skip to content

Commit c66548a

Browse files
committed
sql: add overload of crdb_internal.force_retry that counts retries
Add an overload of `crdb_internal.force_retry` that continues to throw retryable errors until the number of retries matches the argument. This overload is a bit easier to work with than the overload that uses an interval. Future commits will make use of this overload for testing. To power this overload, we add the current autoRetryCounter to planner before every statement execution, which means `crdb_internal.force_retry` is no longer DistSQL-safe. Informs: #145377 Release note: None
1 parent 1cfffc1 commit c66548a

7 files changed

Lines changed: 96 additions & 5 deletions

File tree

pkg/sql/conn_executor_exec.go

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1079,6 +1079,10 @@ func (ex *connExecutor) execStmtInOpenState(
10791079
}()
10801080
}
10811081

1082+
if ex.state.mu.autoRetryReason != nil {
1083+
p.autoRetryCounter = int(ex.state.mu.autoRetryCounter)
1084+
}
1085+
10821086
if ex.executorType != executorTypeInternal &&
10831087
ex.state.mu.txn.IsoLevel() == isolation.ReadCommitted &&
10841088
!ex.implicitTxn() {
@@ -2127,6 +2131,10 @@ func (ex *connExecutor) execStmtInOpenStateWithPausablePortal(
21272131
}()
21282132
}
21292133

2134+
if ex.state.mu.autoRetryReason != nil {
2135+
p.autoRetryCounter = int(ex.state.mu.autoRetryCounter)
2136+
}
2137+
21302138
if ex.executorType != executorTypeInternal &&
21312139
ex.state.mu.txn.IsoLevel() == isolation.ReadCommitted &&
21322140
!ex.implicitTxn() {

pkg/sql/faketreeeval/evalctx.go

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -572,6 +572,11 @@ func (ep *DummyEvalPlanner) ClearQueryPlanCache() {}
572572
// ClearTableStatsCache is part of the eval.Planner interface.
573573
func (ep *DummyEvalPlanner) ClearTableStatsCache() {}
574574

575+
// RetryCounter is part of the eval.Planner interface.
576+
func (ep *DummyEvalPlanner) RetryCounter() int {
577+
return 0
578+
}
579+
575580
// DummyPrivilegedAccessor implements the tree.PrivilegedAccessor interface by returning errors.
576581
type DummyPrivilegedAccessor struct{}
577582

pkg/sql/logictest/testdata/logic_test/txn_retry

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,3 +45,44 @@ SELECT cluster_logical_timestamp(); INSERT INTO test_retry VALUES (1);
4545

4646
statement ok
4747
COMMIT
48+
49+
statement ok
50+
RESET kv_transaction_buffered_writes_enabled
51+
52+
# Test the overload of crdb_internal.force_retry that counts retries.
53+
subtest force_retry
54+
55+
statement ok
56+
SET tracing = on
57+
58+
query I
59+
SELECT crdb_internal.force_retry(0)
60+
----
61+
0
62+
63+
statement ok
64+
SET tracing = off
65+
66+
query T nosort
67+
SELECT message FROM [SHOW TRACE FOR SESSION] WHERE message LIKE '%executing after % retries%'
68+
----
69+
executing after 0 retries, last retry reason: <nil>
70+
71+
statement ok
72+
SET tracing = on
73+
74+
query I
75+
SELECT crdb_internal.force_retry(3)
76+
----
77+
0
78+
79+
statement ok
80+
SET tracing = off
81+
82+
query T nosort
83+
SELECT message FROM [SHOW TRACE FOR SESSION] WHERE message LIKE '%executing after % retries%'
84+
----
85+
executing after 0 retries, last retry reason: <nil>
86+
executing after 1 retries, last retry reason: TransactionRetryWithProtoRefreshError: forced by crdb_internal.force_retry()
87+
executing after 2 retries, last retry reason: TransactionRetryWithProtoRefreshError: forced by crdb_internal.force_retry()
88+
executing after 3 retries, last retry reason: TransactionRetryWithProtoRefreshError: forced by crdb_internal.force_retry()

pkg/sql/planner.go

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -289,6 +289,10 @@ type planner struct {
289289

290290
// datumAlloc is used when decoding datums and running subqueries.
291291
datumAlloc *tree.DatumAlloc
292+
293+
// This is a copy of txnState.mu.autoRetryCounter from when we started the
294+
// statement.
295+
autoRetryCounter int
292296
}
293297

294298
// hasFlowForPausablePortal returns true if the planner is for re-executing a
@@ -949,6 +953,7 @@ func (p *planner) resetPlanner(
949953
p.skipDescriptorCache = false
950954
p.typeResolutionDbID = descpb.InvalidID
951955
p.pausablePortal = nil
956+
p.autoRetryCounter = 0
952957
}
953958

954959
// GetReplicationStreamManager returns a ReplicationStreamManager.
@@ -1052,3 +1057,8 @@ func (p *planner) StartHistoryRetentionJob(
10521057
func (p *planner) ExtendHistoryRetention(ctx context.Context, jobID jobspb.JobID) error {
10531058
return ExtendHistoryRetention(ctx, p.EvalContext(), p.InternalSQLTxn(), jobID)
10541059
}
1060+
1061+
// RetryCounter is part of the eval.Planner interface.
1062+
func (p *planner) RetryCounter() int {
1063+
return p.autoRetryCounter
1064+
}

pkg/sql/sem/builtins/builtins.go

Lines changed: 28 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6061,15 +6061,15 @@ SELECT
60616061
},
60626062
),
60636063

6064-
// If force_retry is called during the specified interval from the beginning
6065-
// of the transaction it returns a retryable error. If not, 0 is returned
6066-
// instead of an error.
6067-
// The second version allows one to create an error intended for a transaction
6068-
// different than the current statement's transaction.
60696064
"crdb_internal.force_retry": makeBuiltin(
60706065
tree.FunctionProperties{
60716066
Category: builtinconstants.CategorySystemInfo,
60726067
},
6068+
// This overload takes an interval parameter. If force_retry is called
6069+
// within this interval from the beginning of the transaction, it returns a
6070+
// retryable error. If force_retry is called after this interval, it returns
6071+
// 0. This allows the transaction to eventually succeed after the interval
6072+
// has passed, assuming it is being retried.
60736073
tree.Overload{
60746074
Types: tree.ParamTypes{{Name: "val", Typ: types.Interval}},
60756075
ReturnType: tree.FixedReturnType(types.Int),
@@ -6085,6 +6085,29 @@ SELECT
60856085
Info: "This function is used only by CockroachDB's developers for testing purposes.",
60866086
Volatility: volatility.Volatile,
60876087
},
6088+
// This overload takes an integer parameter. If the statement or transaction
6089+
// has already been retried < the parameter number of times, force_retry
6090+
// will return a retryable error. If the statement or transaction has
6091+
// already been retried >= the parameter number of times, force_retry will
6092+
// return 0. This allows precise control of the number of times the
6093+
// statement or transaction is retied.
6094+
tree.Overload{
6095+
Types: tree.ParamTypes{{Name: "val", Typ: types.Int}},
6096+
ReturnType: tree.FixedReturnType(types.Int),
6097+
Fn: func(ctx context.Context, evalCtx *eval.Context, args tree.Datums) (tree.Datum, error) {
6098+
retries := int64(evalCtx.Planner.RetryCounter())
6099+
maxRetries := int64(tree.MustBeDInt(args[0]))
6100+
if retries < maxRetries {
6101+
return nil, evalCtx.Txn.GenerateForcedRetryableErr(
6102+
ctx, "forced by crdb_internal.force_retry()",
6103+
)
6104+
}
6105+
return tree.DZero, nil
6106+
},
6107+
Info: "This function is used only by CockroachDB's developers for testing purposes.",
6108+
Volatility: volatility.Volatile,
6109+
DistsqlBlocklist: true, // applicable only on the gateway
6110+
},
60886111
),
60896112

60906113
// Fetches the corresponding lease_holder for the request key.

pkg/sql/sem/builtins/fixed_oids.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2662,6 +2662,7 @@ var builtinOidsArray = []string{
26622662
2699: `jsonb_path_match(target: jsonb, path: jsonpath) -> bool`,
26632663
2700: `jsonb_path_match(target: jsonb, path: jsonpath, vars: jsonb) -> bool`,
26642664
2701: `jsonb_path_match(target: jsonb, path: jsonpath, vars: jsonb, silent: bool) -> bool`,
2665+
2702: `crdb_internal.force_retry(val: int) -> int`,
26652666
}
26662667

26672668
var builtinOidsBySignature map[string]oid.Oid

pkg/sql/sem/eval/deps.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -448,6 +448,9 @@ type Planner interface {
448448

449449
// ClearTableStatsCache removes all entries from the node's table stats cache.
450450
ClearTableStatsCache()
451+
452+
// RetryCounter is the number of times this statement has been retried.
453+
RetryCounter() int
451454
}
452455

453456
// InternalRows is an iterator interface that's exposed by the internal

0 commit comments

Comments
 (0)