Skip to content

Commit a79a1e1

Browse files
ucirelloucirello
authored
Use row-level locking to handle database operations (#118)
* pglock: remove unnecessary transaction to execute heartbeats * pglock: remove unnecessary transaction to execute releases * pglock: remove unnecessary transaction to execute lock acquisitions
1 parent 1b6d0d1 commit a79a1e1

File tree

3 files changed

+62
-163
lines changed

3 files changed

+62
-163
lines changed

‎client.go

+52-60
Original file line numberDiff line numberDiff line change
@@ -223,51 +223,51 @@ func (c *Client) storeAcquire(ctx context.Context, l *Lock) error {
223223
return typedError(err, "cannot run query to read record version number")
224224
}
225225

226-
tx, err := c.db.BeginTx(ctx, &sql.TxOptions{Isolation: sql.LevelSerializable})
227-
if err != nil {
228-
return typedError(err, "cannot create transaction for lock acquisition")
229-
}
230226
c.log.Debug("storeAcquire in: %v %v %v %v", l.name, rvn, l.data, l.recordVersionNumber)
231227
defer func() {
232228
c.log.Debug("storeAcquire out: %v %v %v %v", l.name, rvn, l.data, l.recordVersionNumber)
233229
}()
234-
_, err = tx.ExecContext(ctx, `
230+
rowLockInfo := c.db.QueryRowContext(ctx, `
235231
INSERT INTO `+c.tableName+`
236232
("name", "record_version_number", "data", "owner")
237233
VALUES
238234
($1, $2, $3, $6)
239235
ON CONFLICT ("name") DO UPDATE
240236
SET
241-
"record_version_number" = $2,
237+
"record_version_number" = CASE
238+
WHEN COALESCE(`+c.tableName+`."record_version_number" = $4, TRUE) THEN $2
239+
ELSE `+c.tableName+`."record_version_number"
240+
END,
242241
"data" = CASE
243-
WHEN $5 THEN $3
242+
WHEN COALESCE(`+c.tableName+`."record_version_number" = $4, TRUE) THEN
243+
CASE
244+
WHEN $5 THEN $3
245+
ELSE `+c.tableName+`."data"
246+
END
244247
ELSE `+c.tableName+`."data"
245248
END,
246-
"owner" = $6
247-
WHERE
248-
`+c.tableName+`."record_version_number" IS NULL
249-
OR `+c.tableName+`."record_version_number" = $4
249+
"owner" = CASE
250+
WHEN COALESCE(`+c.tableName+`."record_version_number" = $4, TRUE) THEN $6
251+
ELSE `+c.tableName+`."owner"
252+
END
253+
RETURNING
254+
"record_version_number", "data", "owner"
250255
`, l.name, rvn, l.data, l.recordVersionNumber, l.replaceData, c.owner)
251-
if err != nil {
252-
return typedError(err, "cannot run query to acquire lock")
253-
}
254-
rowLockInfo := tx.QueryRowContext(ctx, `SELECT "record_version_number", "data", "owner" FROM `+c.tableName+` WHERE name = $1 FOR UPDATE`, l.name)
255-
var actualRVN int64
256-
var data []byte
257-
var actualOwner string
258-
if err := rowLockInfo.Scan(&actualRVN, &data, &actualOwner); err != nil {
256+
var (
257+
actualRVN int64
258+
actualData []byte
259+
actualOwner string
260+
)
261+
if err := rowLockInfo.Scan(&actualRVN, &actualData, &actualOwner); err != nil && !errors.Is(err, sql.ErrNoRows) {
259262
return typedError(err, "cannot load information for lock acquisition")
260263
}
261264
l.owner = actualOwner
262265
if actualRVN != rvn {
263266
l.recordVersionNumber = actualRVN
264267
return ErrNotAcquired
265268
}
266-
if err := tx.Commit(); err != nil {
267-
return typedError(err, "cannot commit lock acquisition")
268-
}
269269
l.recordVersionNumber = rvn
270-
l.data = data
270+
l.data = actualData
271271
return nil
272272
}
273273

@@ -312,21 +312,34 @@ func (c *Client) storeRelease(ctx context.Context, l *Lock) error {
312312
defer l.mu.Unlock()
313313
ctx, cancel := context.WithTimeout(ctx, l.leaseDuration)
314314
defer cancel()
315-
tx, err := c.db.BeginTx(ctx, &sql.TxOptions{Isolation: sql.LevelSerializable})
316-
if err != nil {
317-
return typedError(err, "cannot create transaction for lock acquisition")
318-
}
319-
result, err := tx.ExecContext(ctx, `
320-
UPDATE
321-
`+c.tableName+`
322-
SET
323-
"record_version_number" = NULL
324-
WHERE
325-
"name" = $1
326-
AND "record_version_number" = $2
327-
`, l.name, l.recordVersionNumber)
328-
if err != nil {
329-
return typedError(err, "cannot run query to release lock")
315+
var result sql.Result
316+
switch l.keepOnRelease {
317+
case true:
318+
res, err := c.db.ExecContext(ctx, `
319+
UPDATE
320+
`+c.tableName+`
321+
SET
322+
"record_version_number" = NULL
323+
WHERE
324+
"name" = $1
325+
AND "record_version_number" = $2
326+
`, l.name, l.recordVersionNumber)
327+
if err != nil {
328+
return typedError(err, "cannot run query to release lock (keep)")
329+
}
330+
result = res
331+
case false:
332+
res, err := c.db.ExecContext(ctx, `
333+
DELETE FROM
334+
`+c.tableName+`
335+
WHERE
336+
"name" = $1
337+
AND "record_version_number" = $2
338+
`, l.name, l.recordVersionNumber)
339+
if err != nil {
340+
return typedError(err, "cannot run query to delete lock (delete)")
341+
}
342+
result = res
330343
}
331344
affected, err := result.RowsAffected()
332345
if err != nil {
@@ -335,20 +348,6 @@ func (c *Client) storeRelease(ctx context.Context, l *Lock) error {
335348
l.isReleased = true
336349
return ErrLockAlreadyReleased
337350
}
338-
if !l.keepOnRelease {
339-
_, err := tx.ExecContext(ctx, `
340-
DELETE FROM
341-
`+c.tableName+`
342-
WHERE
343-
"name" = $1
344-
AND "record_version_number" IS NULL`, l.name)
345-
if err != nil {
346-
return typedError(err, "cannot run query to delete lock")
347-
}
348-
}
349-
if err := tx.Commit(); err != nil {
350-
return typedError(err, "cannot commit lock release")
351-
}
352351
l.isReleased = true
353352
l.heartbeatCancel()
354353
return nil
@@ -392,11 +391,7 @@ func (c *Client) storeHeartbeat(ctx context.Context, l *Lock) error {
392391
if err != nil {
393392
return typedError(err, "cannot run query to read record version number")
394393
}
395-
tx, err := c.db.BeginTx(ctx, &sql.TxOptions{Isolation: sql.LevelSerializable})
396-
if err != nil {
397-
return typedError(err, "cannot create transaction for lock acquisition")
398-
}
399-
result, err := tx.ExecContext(ctx, `
394+
result, err := c.db.ExecContext(ctx, `
400395
UPDATE
401396
`+c.tableName+`
402397
SET
@@ -414,9 +409,6 @@ func (c *Client) storeHeartbeat(ctx context.Context, l *Lock) error {
414409
} else if affected == 0 {
415410
return ErrLockAlreadyReleased
416411
}
417-
if err := tx.Commit(); err != nil {
418-
return typedError(err, "cannot commit lock heartbeat")
419-
}
420412
l.recordVersionNumber = rvn
421413
return nil
422414
}

‎client_internal_test.go

+8-101
Original file line numberDiff line numberDiff line change
@@ -120,15 +120,6 @@ func TestDBErrorHandling(t *testing.T) {
120120
}
121121
}
122122
t.Run("acquire", func(t *testing.T) {
123-
t.Run("bad tx", func(t *testing.T) {
124-
client, mock, _ := setup()
125-
badTx := errors.New("transaction begin error")
126-
mock.ExpectQuery(`SELECT nextval\('locks_rvn'\)`).WillReturnRows(sqlmock.NewRows([]string{"nextval"}).AddRow(1))
127-
mock.ExpectBegin().WillReturnError(badTx)
128-
if _, err := client.Acquire("bad-tx"); !errors.Is(err, badTx) {
129-
t.Errorf("expected tx error missing: %v", err)
130-
}
131-
})
132123
t.Run("bad rvn", func(t *testing.T) {
133124
client, mock, _ := setup()
134125
badRVN := errors.New("cannot load next RVN")
@@ -141,103 +132,32 @@ func TestDBErrorHandling(t *testing.T) {
141132
client, mock, _ := setup()
142133
badInsert := errors.New("cannot insert")
143134
mock.ExpectQuery(`SELECT nextval\('locks_rvn'\)`).WillReturnRows(sqlmock.NewRows([]string{"nextval"}).AddRow(1))
144-
mock.ExpectBegin()
145-
mock.ExpectExec(`INSERT INTO locks (.+)`).WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg()).WillReturnError(badInsert)
135+
mock.ExpectQuery(`INSERT INTO locks (.+)`).WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg()).WillReturnError(badInsert)
146136
if _, err := client.Acquire("bad-insert"); !errors.Is(err, badInsert) {
147137
t.Errorf("expected RVN error missing: %v", err)
148138
}
149139
})
150-
t.Run("bad RVN confirmation", func(t *testing.T) {
151-
client, mock, _ := setup()
152-
badRVN := errors.New("cannot confirm RVN")
153-
mock.ExpectQuery(`SELECT nextval\('locks_rvn'\)`).WillReturnRows(sqlmock.NewRows([]string{"nextval"}).AddRow(1))
154-
mock.ExpectBegin()
155-
mock.ExpectExec(`INSERT INTO locks (.+)`).WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg()).WillReturnResult(sqlmock.NewResult(0, 1))
156-
mock.ExpectQuery(`SELECT "record_version_number", "data", "owner" FROM locks WHERE name = (.+)`).WithArgs(sqlmock.AnyArg()).WillReturnError(badRVN)
157-
if _, err := client.Acquire("bad-insert"); !errors.Is(err, badRVN) {
158-
t.Errorf("expected RVN confirmation error missing: %v", err)
159-
}
160-
})
161-
t.Run("bad commit", func(t *testing.T) {
162-
client, mock, _ := setup()
163-
badCommit := errors.New("cannot confirm RVN")
164-
mock.ExpectQuery(`SELECT nextval\('locks_rvn'\)`).WillReturnRows(sqlmock.NewRows([]string{"nextval"}).AddRow(1))
165-
mock.ExpectBegin()
166-
mock.ExpectExec(`INSERT INTO locks (.+)`).WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg()).WillReturnResult(sqlmock.NewResult(0, 1))
167-
mock.ExpectQuery(`SELECT "record_version_number", "data", "owner" FROM locks WHERE name = (.+)`).
168-
WithArgs(sqlmock.AnyArg()).
169-
WillReturnRows(
170-
sqlmock.NewRows([]string{
171-
"record_version_number",
172-
"data",
173-
"owner",
174-
}).AddRow(1, []byte{}, "owner"),
175-
)
176-
mock.ExpectCommit().WillReturnError(badCommit)
177-
if _, err := client.Acquire("bad-insert"); !errors.Is(err, badCommit) {
178-
t.Errorf("expected commit error missing: %v", err)
179-
}
180-
})
181140
})
182141
t.Run("release", func(t *testing.T) {
183-
t.Run("bad tx", func(t *testing.T) {
184-
client, mock, fakeLock := setup()
185-
badTx := errors.New("transaction begin error")
186-
mock.ExpectBegin().WillReturnError(badTx)
187-
if err := client.Release(fakeLock); !errors.Is(err, badTx) {
188-
t.Errorf("expected tx error missing: %v", err)
189-
}
190-
})
191142
t.Run("bad update", func(t *testing.T) {
192143
client, mock, fakeLock := setup()
193-
badUpdate := errors.New("cannot update")
194-
mock.ExpectBegin()
195-
mock.ExpectExec(`UPDATE locks (.+)`).WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg()).WillReturnError(badUpdate)
196-
if err := client.Release(fakeLock); !errors.Is(err, badUpdate) {
197-
t.Errorf("expected update error missing: %v", err)
198-
}
199-
})
200-
t.Run("bad update result", func(t *testing.T) {
201-
client, mock, fakeLock := setup()
202-
badUpdateResult := errors.New("cannot grab update result")
203-
mock.ExpectBegin()
204-
mock.ExpectExec(`UPDATE locks (.+)`).WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg()).WillReturnResult(sqlmock.NewErrorResult(badUpdateResult))
205-
if err := client.Release(fakeLock); !errors.Is(err, badUpdateResult) {
206-
t.Errorf("expected update result error missing: %v", err)
144+
badDelete := errors.New("cannot delete lock entry")
145+
fakeLock.keepOnRelease = true
146+
mock.ExpectExec(`UPDATE locks (.+)`).WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg()).WillReturnError(badDelete)
147+
if err := client.Release(fakeLock); !errors.Is(err, badDelete) {
148+
t.Errorf("expected delete error missing: %v", err)
207149
}
208150
})
209151
t.Run("bad delete", func(t *testing.T) {
210152
client, mock, fakeLock := setup()
211153
badDelete := errors.New("cannot delete lock entry")
212-
mock.ExpectBegin()
213-
mock.ExpectExec(`UPDATE locks (.+)`).WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg()).WillReturnResult(sqlmock.NewResult(0, 1))
214-
mock.ExpectExec(`DELETE FROM locks (.+)`).WithArgs(sqlmock.AnyArg()).WillReturnError(badDelete)
154+
mock.ExpectExec(`DELETE FROM locks (.+)`).WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg()).WillReturnError(badDelete)
215155
if err := client.Release(fakeLock); !errors.Is(err, badDelete) {
216156
t.Errorf("expected delete error missing: %v", err)
217157
}
218158
})
219-
t.Run("bad commit", func(t *testing.T) {
220-
client, mock, fakeLock := setup()
221-
badCommit := errors.New("cannot commit release")
222-
mock.ExpectBegin()
223-
mock.ExpectExec(`UPDATE locks (.+)`).WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg()).WillReturnResult(sqlmock.NewResult(0, 1))
224-
mock.ExpectExec(`DELETE FROM locks (.+)`).WithArgs(sqlmock.AnyArg()).WillReturnResult(sqlmock.NewResult(0, 1))
225-
mock.ExpectCommit().WillReturnError(badCommit)
226-
if err := client.Release(fakeLock); !errors.Is(err, badCommit) {
227-
t.Errorf("expected commit error missing: %v", err)
228-
}
229-
})
230159
})
231160
t.Run("heartbeat", func(t *testing.T) {
232-
t.Run("bad tx", func(t *testing.T) {
233-
client, mock, fakeLock := setup()
234-
badTx := errors.New("transaction begin error")
235-
mock.ExpectQuery(`SELECT nextval\('locks_rvn'\)`).WillReturnRows(sqlmock.NewRows([]string{"nextval"}).AddRow(1))
236-
mock.ExpectBegin().WillReturnError(badTx)
237-
if err := client.SendHeartbeat(context.Background(), fakeLock); !errors.Is(err, badTx) {
238-
t.Errorf("expected tx error missing: %v", err)
239-
}
240-
})
241161
t.Run("bad rvn", func(t *testing.T) {
242162
client, mock, fakeLock := setup()
243163
badRVN := errors.New("cannot load next RVN")
@@ -246,11 +166,10 @@ func TestDBErrorHandling(t *testing.T) {
246166
t.Errorf("expected RVN error missing: %v", err)
247167
}
248168
})
249-
t.Run("bad insert", func(t *testing.T) {
169+
t.Run("bad update", func(t *testing.T) {
250170
client, mock, fakeLock := setup()
251171
badUpdate := errors.New("cannot insert")
252172
mock.ExpectQuery(`SELECT nextval\('locks_rvn'\)`).WillReturnRows(sqlmock.NewRows([]string{"nextval"}).AddRow(1))
253-
mock.ExpectBegin()
254173
mock.ExpectExec(`UPDATE locks (.+)`).WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg()).WillReturnError(badUpdate)
255174
if err := client.SendHeartbeat(context.Background(), fakeLock); !errors.Is(err, badUpdate) {
256175
t.Errorf("expected RVN error missing: %v", err)
@@ -260,23 +179,11 @@ func TestDBErrorHandling(t *testing.T) {
260179
client, mock, fakeLock := setup()
261180
badRVN := errors.New("cannot confirm RVN")
262181
mock.ExpectQuery(`SELECT nextval\('locks_rvn'\)`).WillReturnRows(sqlmock.NewRows([]string{"nextval"}).AddRow(1))
263-
mock.ExpectBegin()
264182
mock.ExpectExec(`UPDATE locks (.+)`).WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg()).WillReturnResult(sqlmock.NewErrorResult(badRVN))
265183
if err := client.SendHeartbeat(context.Background(), fakeLock); !errors.Is(err, badRVN) {
266184
t.Errorf("expected RVN confirmation error missing: %v", err)
267185
}
268186
})
269-
t.Run("bad commit", func(t *testing.T) {
270-
client, mock, fakeLock := setup()
271-
badCommit := errors.New("cannot confirm RVN")
272-
mock.ExpectQuery(`SELECT nextval\('locks_rvn'\)`).WillReturnRows(sqlmock.NewRows([]string{"nextval"}).AddRow(1))
273-
mock.ExpectBegin()
274-
mock.ExpectExec(`UPDATE locks (.+)`).WithArgs(sqlmock.AnyArg(), sqlmock.AnyArg(), sqlmock.AnyArg()).WillReturnResult(sqlmock.NewResult(0, 1))
275-
mock.ExpectCommit().WillReturnError(badCommit)
276-
if err := client.SendHeartbeat(context.Background(), fakeLock); !errors.Is(err, badCommit) {
277-
t.Errorf("expected commit error missing: %v", err)
278-
}
279-
})
280187
})
281188

282189
t.Run("GetAllLocks", func(t *testing.T) {

‎client_test.go

+2-2
Original file line numberDiff line numberDiff line change
@@ -350,14 +350,14 @@ func TestKeepOnRelease(t *testing.T) {
350350
expected := []byte("42")
351351
l, err := c.Acquire(name, pglock.KeepOnRelease(), pglock.WithData(expected))
352352
if err != nil {
353-
t.Fatal("unexpected error while acquiring lock:", err)
353+
t.Fatal("unexpected error while acquiring lock (take 1):", err)
354354
}
355355
t.Log("lock acquired")
356356
l.Close()
357357

358358
l2, err := c.Acquire(name)
359359
if err != nil {
360-
t.Fatal("unexpected error while acquiring lock:", err)
360+
t.Fatal("unexpected error while acquiring lock (take 2):", err)
361361
}
362362
defer l2.Close()
363363
t.Log("lock reacquired")

0 commit comments

Comments
 (0)