@@ -33,6 +33,7 @@ import (
33
33
"github.com/cockroachdb/cockroach/pkg/util/retry"
34
34
"github.com/cockroachdb/cockroach/pkg/util/syncutil"
35
35
"github.com/cockroachdb/cockroach/pkg/util/timeutil"
36
+ "github.com/pkg/errors"
36
37
)
37
38
38
39
const (
@@ -46,28 +47,35 @@ type bankClient struct {
46
47
count uint64
47
48
}
48
49
49
- func (client * bankClient ) transferMoney (numAccounts , maxTransfer int ) error {
50
+ func (client * bankClient ) transferMoney (ctx context. Context , numAccounts , maxTransfer int ) error {
50
51
from := rand .Intn (numAccounts )
51
52
to := rand .Intn (numAccounts - 1 )
52
53
if from == to {
53
54
to = numAccounts - 1
54
55
}
55
56
amount := rand .Intn (maxTransfer )
56
57
57
- const update = `
58
+ tBegin := timeutil .Now ()
59
+
60
+ // If this statement gets stuck, the test harness will get stuck. Run with a
61
+ // statement timeout, which unfortunately precludes the use of prepared
62
+ // statements.
63
+ q := fmt .Sprintf (`
64
+ SET statement_timeout = '31s';
58
65
UPDATE bank.accounts
59
- SET balance = CASE id WHEN $1 THEN balance-$3 WHEN $2 THEN balance+$3 END
60
- WHERE id IN ($1, $2) AND (SELECT balance >= $3 FROM bank.accounts WHERE id = $1)
61
- `
66
+ SET balance = CASE id WHEN %[1]d THEN balance-%[3]d WHEN %[2]d THEN balance+%[3]d END
67
+ WHERE id IN (%[1]d, %[2]d) AND (SELECT balance >= %[3]d FROM bank.accounts WHERE id = %[1]d);
68
+ ` , from , to , amount )
69
+
62
70
client .RLock ()
63
71
defer client .RUnlock ()
64
- _ , err := client .db .Exec ( update , from , to , amount )
72
+ _ , err := client .db .ExecContext ( ctx , q )
65
73
if err == nil {
66
74
// Do all increments under the read lock so that grabbing a write lock in
67
75
// startChaosMonkey below guarantees no more increments could be incoming.
68
76
atomic .AddUint64 (& client .count , 1 )
69
77
}
70
- return err
78
+ return errors . Wrapf ( err , "after %.1fs" , timeutil . Since ( tBegin ). Seconds ())
71
79
}
72
80
73
81
type bankState struct {
@@ -111,12 +119,12 @@ func (s *bankState) initBank(ctx context.Context, t *test, c *cluster) {
111
119
db := c .Conn (ctx , 1 )
112
120
defer db .Close ()
113
121
114
- if _ , err := db .Exec ( `CREATE DATABASE IF NOT EXISTS bank` ); err != nil {
122
+ if _ , err := db .ExecContext ( ctx , `CREATE DATABASE IF NOT EXISTS bank` ); err != nil {
115
123
t .Fatal (err )
116
124
}
117
125
118
126
// Delete table created by a prior instance of a test.
119
- if _ , err := db .Exec ( `DROP TABLE IF EXISTS bank.accounts` ); err != nil {
127
+ if _ , err := db .ExecContext ( ctx , `DROP TABLE IF EXISTS bank.accounts` ); err != nil {
120
128
t .Fatal (err )
121
129
}
122
130
@@ -125,7 +133,7 @@ CREATE TABLE bank.accounts (
125
133
id INT PRIMARY KEY,
126
134
balance INT NOT NULL
127
135
)`
128
- if _ , err := db .Exec ( schema ); err != nil {
136
+ if _ , err := db .ExecContext ( ctx , schema ); err != nil {
129
137
t .Fatal (err )
130
138
}
131
139
@@ -139,7 +147,7 @@ CREATE TABLE bank.accounts (
139
147
values = append (values , i )
140
148
}
141
149
stmt := `INSERT INTO bank.accounts (id, balance) VALUES ` + placeholders .String ()
142
- if _ , err := db .Exec ( stmt , values ... ); err != nil {
150
+ if _ , err := db .ExecContext ( ctx , stmt , values ... ); err != nil {
143
151
t .Fatal (err )
144
152
}
145
153
}
@@ -151,7 +159,7 @@ func (s *bankState) transferMoney(
151
159
defer c .l .Printf ("client %d shutting down\n " , idx )
152
160
client := & s .clients [idx - 1 ]
153
161
for ! s .done (ctx ) {
154
- if err := client .transferMoney (numAccounts , maxTransfer ); err != nil {
162
+ if err := client .transferMoney (ctx , numAccounts , maxTransfer ); err != nil {
155
163
// Ignore some errors.
156
164
if ! pgerror .IsSQLRetryableError (err ) {
157
165
// Report the err and terminate.
@@ -180,7 +188,7 @@ func (s *bankState) verifyAccounts(ctx context.Context, t *test) {
180
188
// chaos monkey.
181
189
client .RLock ()
182
190
defer client .RUnlock ()
183
- err := client .db .QueryRow ( "SELECT count(*), sum(balance) FROM bank.accounts" ).Scan (& accounts , & sum )
191
+ err := client .db .QueryRowContext ( ctx , "SELECT count(*), sum(balance) FROM bank.accounts" ).Scan (& accounts , & sum )
184
192
if err != nil && ! pgerror .IsSQLRetryableError (err ) {
185
193
t .Fatal (err )
186
194
}
@@ -311,7 +319,7 @@ func (s *bankState) startSplitMonkey(ctx context.Context, d time.Duration, c *cl
311
319
zipF := accountDistribution (r )
312
320
key := zipF .Uint64 ()
313
321
c .l .Printf ("round %d: splitting key %v\n " , curRound , key )
314
- _ , err := client .db .Exec ( fmt .Sprintf (
322
+ _ , err := client .db .ExecContext ( ctx , fmt .Sprintf (
315
323
`SET experimental_force_split_at = true; ALTER TABLE bank.accounts SPLIT AT VALUES (%d)` , key ))
316
324
if err != nil && ! (pgerror .IsSQLRetryableError (err ) || isExpectedRelocateError (err )) {
317
325
s .errChan <- err
@@ -333,7 +341,7 @@ func (s *bankState) startSplitMonkey(ctx context.Context, d time.Duration, c *cl
333
341
c .l .Printf ("round %d: relocating key %d to nodes %s\n " ,
334
342
curRound , key , nodes [1 :])
335
343
336
- _ , err := client .db .Exec ( relocateQuery )
344
+ _ , err := client .db .ExecContext ( ctx , relocateQuery )
337
345
if err != nil && ! (pgerror .IsSQLRetryableError (err ) || isExpectedRelocateError (err )) {
338
346
s .errChan <- err
339
347
}
0 commit comments