Skip to content

Commit 59db543

Browse files
sreyasreya
committed
port ReadModifyUpdate
1 parent f6cf3b7 commit 59db543

File tree

3 files changed

+132
-3
lines changed

3 files changed

+132
-3
lines changed

‎coderd/database/tx.go

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
package database
2+
3+
import (
4+
"database/sql"
5+
6+
"github.com/lib/pq"
7+
"golang.org/x/xerrors"
8+
)
9+
10+
const maxRetries = 5
11+
12+
// ReadModifyUpdate is a helper function to run a db transaction that reads some
13+
// object(s), modifies some of the data, and writes the modified object(s) back
14+
// to the database. It is run in a transaction at RepeatableRead isolation so
15+
// that if another database client also modifies the data we are writing and
16+
// commits, then the transaction is rolled back and restarted.
17+
//
18+
// This is needed because we typically read all object columns, modify some
19+
// subset, and then write all columns. Consider an object with columns A, B and
20+
// initial values A=1, B=1. Two database clients work simultaneously, with one
21+
// client attempting to set A=2, and another attempting to set B=2. They both
22+
// initially read A=1, B=1, and then one writes A=2, B=1, and the other writes
23+
// A=1, B=2. With default PostgreSQL isolation of ReadCommitted, both of these
24+
// transactions would succeed and we end up with either A=2, B=1 or A=1, B=2.
25+
// One or other client gets their transaction wiped out even though the data
26+
// they wanted to change didn't conflict.
27+
//
28+
// If we run at RepeatableRead isolation, then one or other transaction will
29+
// fail. Let's say the transaction that sets A=2 succeeds. Then the first B=2
30+
// transaction fails, but here we retry. The second attempt we read A=2, B=1,
31+
// then write A=2, B=2 as desired, and this succeeds.
32+
func ReadModifyUpdate(db Store, f func(tx Store) error,
33+
) error {
34+
var err error
35+
for retries := 0; retries < maxRetries; retries++ {
36+
err = db.InTx(f, &sql.TxOptions{
37+
Isolation: sql.LevelRepeatableRead,
38+
})
39+
var pqe *pq.Error
40+
if xerrors.As(err, &pqe) {
41+
if pqe.Code == "40001" {
42+
// serialization error, retry
43+
continue
44+
}
45+
}
46+
return err
47+
}
48+
return xerrors.Errorf("too many errors; last error: %w", err)
49+
}

‎coderd/database/tx_test.go

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
package database_test
2+
3+
import (
4+
"database/sql"
5+
"testing"
6+
7+
"github.com/golang/mock/gomock"
8+
"github.com/lib/pq"
9+
"github.com/stretchr/testify/require"
10+
"golang.org/x/xerrors"
11+
12+
"github.com/coder/coder/coderd/database"
13+
"github.com/coder/coder/coderd/database/dbmock"
14+
)
15+
16+
func TestReadModifyUpdate_OK(t *testing.T) {
17+
t.Parallel()
18+
19+
mDB := dbmock.NewMockStore(gomock.NewController(t))
20+
21+
mDB.EXPECT().
22+
InTx(gomock.Any(), &sql.TxOptions{Isolation: sql.LevelRepeatableRead}).
23+
Times(1).
24+
Return(nil)
25+
err := database.ReadModifyUpdate(mDB, func(tx database.Store) error {
26+
return nil
27+
})
28+
require.NoError(t, err)
29+
}
30+
31+
func TestReadModifyUpdate_RetryOK(t *testing.T) {
32+
t.Parallel()
33+
34+
mDB := dbmock.NewMockStore(gomock.NewController(t))
35+
36+
firstUpdate := mDB.EXPECT().
37+
InTx(gomock.Any(), &sql.TxOptions{Isolation: sql.LevelRepeatableRead}).
38+
Times(1).
39+
Return(&pq.Error{Code: pq.ErrorCode("40001")})
40+
mDB.EXPECT().
41+
InTx(gomock.Any(), &sql.TxOptions{Isolation: sql.LevelRepeatableRead}).
42+
After(firstUpdate).
43+
Times(1).
44+
Return(nil)
45+
46+
err := database.ReadModifyUpdate(mDB, func(tx database.Store) error {
47+
return nil
48+
})
49+
require.NoError(t, err)
50+
}
51+
52+
func TestReadModifyUpdate_HardError(t *testing.T) {
53+
t.Parallel()
54+
55+
mDB := dbmock.NewMockStore(gomock.NewController(t))
56+
57+
mDB.EXPECT().
58+
InTx(gomock.Any(), &sql.TxOptions{Isolation: sql.LevelRepeatableRead}).
59+
Times(1).
60+
Return(xerrors.New("a bad thing happened"))
61+
62+
err := database.ReadModifyUpdate(mDB, func(tx database.Store) error {
63+
return nil
64+
})
65+
require.ErrorContains(t, err, "a bad thing happened")
66+
}
67+
68+
func TestReadModifyUpdate_TooManyRetries(t *testing.T) {
69+
t.Parallel()
70+
71+
mDB := dbmock.NewMockStore(gomock.NewController(t))
72+
73+
mDB.EXPECT().
74+
InTx(gomock.Any(), &sql.TxOptions{Isolation: sql.LevelRepeatableRead}).
75+
Times(5).
76+
Return(&pq.Error{Code: pq.ErrorCode("40001")})
77+
err := database.ReadModifyUpdate(mDB, func(tx database.Store) error {
78+
return nil
79+
})
80+
require.ErrorContains(t, err, "too many errors")
81+
}

‎enterprise/coderd/groups.go

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -190,7 +190,7 @@ func (api *API) patchGroup(rw http.ResponseWriter, r *http.Request) {
190190
}
191191
}
192192

193-
err := api.Database.InTx(func(tx database.Store) error {
193+
err := database.ReadModifyUpdate(api.Database, func(tx database.Store) error {
194194
var err error
195195
group, err = tx.GetGroupByID(ctx, group.ID)
196196
if err != nil {
@@ -251,9 +251,8 @@ func (api *API) patchGroup(rw http.ResponseWriter, r *http.Request) {
251251
}
252252
}
253253
return nil
254-
}, &sql.TxOptions{
255-
Isolation: sql.LevelRepeatableRead,
256254
})
255+
257256
if database.IsUniqueViolation(err) {
258257
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
259258
Message: "Cannot add the same user to a group twice!",

0 commit comments

Comments
 (0)