Skip to content

Commit 3dfb0ff

Browse files
prinkova.prinkov
prinkov
and
a.prinkov
authored
used uber atomic bool instead standard in lock/unlock db (#580)
* used uber atomic bool instead standard in lock/unlock db * fixing local lock optimization implementation * fixing mysql local lock optimization implementation * fixing cockroachdb local lock optimization implementation * CAS wrapper to automatically restore * added test for cas with restore * added atomic lock to sqlite & mongo * added atomic lock to sqlite & mongo * fix/refactor test cas restore function * disable lock for mongo * mongo local lock before acquiring the db lock Co-authored-by: a.prinkov <[email protected]>
1 parent 3b3c1b6 commit 3dfb0ff

File tree

19 files changed

+347
-299
lines changed

19 files changed

+347
-299
lines changed

‎database/cassandra/cassandra.go

+6-4
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package cassandra
33
import (
44
"errors"
55
"fmt"
6+
"go.uber.org/atomic"
67
"io"
78
"io/ioutil"
89
nurl "net/url"
@@ -45,7 +46,7 @@ type Config struct {
4546

4647
type Cassandra struct {
4748
session *gocql.Session
48-
isLocked bool
49+
isLocked atomic.Bool
4950

5051
// Open and WithInstance need to guarantee that config is never nil
5152
config *Config
@@ -182,15 +183,16 @@ func (c *Cassandra) Close() error {
182183
}
183184

184185
func (c *Cassandra) Lock() error {
185-
if c.isLocked {
186+
if !c.isLocked.CAS(false, true) {
186187
return database.ErrLocked
187188
}
188-
c.isLocked = true
189189
return nil
190190
}
191191

192192
func (c *Cassandra) Unlock() error {
193-
c.isLocked = false
193+
if !c.isLocked.CAS(true, false) {
194+
return database.ErrNotLocked
195+
}
194196
return nil
195197
}
196198

‎database/clickhouse/clickhouse.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -285,7 +285,7 @@ func (ch *ClickHouse) Lock() error {
285285
}
286286
func (ch *ClickHouse) Unlock() error {
287287
if !ch.isLocked.CAS(true, false) {
288-
return database.ErrLocked
288+
return database.ErrNotLocked
289289
}
290290

291291
return nil

‎database/cockroachdb/cockroachdb.go

+49-52
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import (
44
"context"
55
"database/sql"
66
"fmt"
7+
"go.uber.org/atomic"
78
"io"
89
"io/ioutil"
910
nurl "net/url"
@@ -46,7 +47,7 @@ type Config struct {
4647

4748
type CockroachDb struct {
4849
db *sql.DB
49-
isLocked bool
50+
isLocked atomic.Bool
5051

5152
// Open and WithInstance need to guarantee that config is never nil
5253
config *Config
@@ -152,71 +153,67 @@ func (c *CockroachDb) Close() error {
152153
// Locking is done manually with a separate lock table. Implementing advisory locks in CRDB is being discussed
153154
// See: https://github.com/cockroachdb/cockroach/issues/13546
154155
func (c *CockroachDb) Lock() error {
155-
err := crdb.ExecuteTx(context.Background(), c.db, nil, func(tx *sql.Tx) (err error) {
156-
aid, err := database.GenerateAdvisoryLockId(c.config.DatabaseName)
157-
if err != nil {
158-
return err
159-
}
160-
161-
query := "SELECT * FROM " + c.config.LockTable + " WHERE lock_id = $1"
162-
rows, err := tx.Query(query, aid)
163-
if err != nil {
164-
return database.Error{OrigErr: err, Err: "failed to fetch migration lock", Query: []byte(query)}
165-
}
166-
defer func() {
167-
if errClose := rows.Close(); errClose != nil {
168-
err = multierror.Append(err, errClose)
156+
return database.CasRestoreOnErr(&c.isLocked, false, true, database.ErrLocked, func() (err error) {
157+
return crdb.ExecuteTx(context.Background(), c.db, nil, func(tx *sql.Tx) (err error) {
158+
aid, err := database.GenerateAdvisoryLockId(c.config.DatabaseName)
159+
if err != nil {
160+
return err
169161
}
170-
}()
171162

172-
// If row exists at all, lock is present
173-
locked := rows.Next()
174-
if locked && !c.config.ForceLock {
175-
return database.ErrLocked
176-
}
163+
query := "SELECT * FROM " + c.config.LockTable + " WHERE lock_id = $1"
164+
rows, err := tx.Query(query, aid)
165+
if err != nil {
166+
return database.Error{OrigErr: err, Err: "failed to fetch migration lock", Query: []byte(query)}
167+
}
168+
defer func() {
169+
if errClose := rows.Close(); errClose != nil {
170+
err = multierror.Append(err, errClose)
171+
}
172+
}()
173+
174+
// If row exists at all, lock is present
175+
locked := rows.Next()
176+
if locked && !c.config.ForceLock {
177+
return database.ErrLocked
178+
}
177179

178-
query = "INSERT INTO " + c.config.LockTable + " (lock_id) VALUES ($1)"
179-
if _, err := tx.Exec(query, aid); err != nil {
180-
return database.Error{OrigErr: err, Err: "failed to set migration lock", Query: []byte(query)}
181-
}
180+
query = "INSERT INTO " + c.config.LockTable + " (lock_id) VALUES ($1)"
181+
if _, err := tx.Exec(query, aid); err != nil {
182+
return database.Error{OrigErr: err, Err: "failed to set migration lock", Query: []byte(query)}
183+
}
182184

183-
return nil
185+
return nil
186+
})
184187
})
185-
186-
if err != nil {
187-
return err
188-
} else {
189-
c.isLocked = true
190-
return nil
191-
}
192188
}
193189

194190
// Locking is done manually with a separate lock table. Implementing advisory locks in CRDB is being discussed
195191
// See: https://github.com/cockroachdb/cockroach/issues/13546
196192
func (c *CockroachDb) Unlock() error {
197-
aid, err := database.GenerateAdvisoryLockId(c.config.DatabaseName)
198-
if err != nil {
199-
return err
200-
}
193+
return database.CasRestoreOnErr(&c.isLocked, true, false, database.ErrNotLocked, func() (err error) {
194+
aid, err := database.GenerateAdvisoryLockId(c.config.DatabaseName)
195+
if err != nil {
196+
return err
197+
}
201198

202-
// In the event of an implementation (non-migration) error, it is possible for the lock to not be released. Until
203-
// a better locking mechanism is added, a manual purging of the lock table may be required in such circumstances
204-
query := "DELETE FROM " + c.config.LockTable + " WHERE lock_id = $1"
205-
if _, err := c.db.Exec(query, aid); err != nil {
206-
if e, ok := err.(*pq.Error); ok {
207-
// 42P01 is "UndefinedTableError" in CockroachDB
208-
// https://github.com/cockroachdb/cockroach/blob/master/pkg/sql/pgwire/pgerror/codes.go
209-
if e.Code == "42P01" {
210-
// On drops, the lock table is fully removed; This is fine, and is a valid "unlocked" state for the schema
211-
c.isLocked = false
212-
return nil
199+
// In the event of an implementation (non-migration) error, it is possible for the lock to not be released. Until
200+
// a better locking mechanism is added, a manual purging of the lock table may be required in such circumstances
201+
query := "DELETE FROM " + c.config.LockTable + " WHERE lock_id = $1"
202+
if _, err := c.db.Exec(query, aid); err != nil {
203+
if e, ok := err.(*pq.Error); ok {
204+
// 42P01 is "UndefinedTableError" in CockroachDB
205+
// https://github.com/cockroachdb/cockroach/blob/master/pkg/sql/pgwire/pgerror/codes.go
206+
if e.Code == "42P01" {
207+
// On drops, the lock table is fully removed; This is fine, and is a valid "unlocked" state for the schema
208+
return nil
209+
}
213210
}
211+
212+
return database.Error{OrigErr: err, Err: "failed to release migration lock", Query: []byte(query)}
214213
}
215-
return database.Error{OrigErr: err, Err: "failed to release migration lock", Query: []byte(query)}
216-
}
217214

218-
c.isLocked = false
219-
return nil
215+
return nil
216+
})
220217
}
221218

222219
func (c *CockroachDb) Run(migration io.Reader) error {

‎database/firebird/firebird.go

+6-4
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import (
1010
"github.com/golang-migrate/migrate/v4/database"
1111
"github.com/hashicorp/go-multierror"
1212
_ "github.com/nakagami/firebirdsql"
13+
"go.uber.org/atomic"
1314
"io"
1415
"io/ioutil"
1516
nurl "net/url"
@@ -36,7 +37,7 @@ type Firebird struct {
3637
// Locking and unlocking need to use the same connection
3738
conn *sql.Conn
3839
db *sql.DB
39-
isLocked bool
40+
isLocked atomic.Bool
4041

4142
// Open and WithInstance need to guarantee that config is never nil
4243
config *Config
@@ -106,15 +107,16 @@ func (f *Firebird) Close() error {
106107
}
107108

108109
func (f *Firebird) Lock() error {
109-
if f.isLocked {
110+
if !f.isLocked.CAS(false, true) {
110111
return database.ErrLocked
111112
}
112-
f.isLocked = true
113113
return nil
114114
}
115115

116116
func (f *Firebird) Unlock() error {
117-
f.isLocked = false
117+
if !f.isLocked.CAS(true, false) {
118+
return database.ErrNotLocked
119+
}
118120
return nil
119121
}
120122

‎database/mongodb/mongodb.go

+52-45
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import (
1010
"go.mongodb.org/mongo-driver/mongo"
1111
"go.mongodb.org/mongo-driver/mongo/options"
1212
"go.mongodb.org/mongo-driver/x/mongo/driver/connstring"
13+
"go.uber.org/atomic"
1314
"io"
1415
"io/ioutil"
1516
"net/url"
@@ -40,9 +41,10 @@ var (
4041
)
4142

4243
type Mongo struct {
43-
client *mongo.Client
44-
db *mongo.Database
45-
config *Config
44+
client *mongo.Client
45+
db *mongo.Database
46+
config *Config
47+
isLocked atomic.Bool
4648
}
4749

4850
type Locking struct {
@@ -327,55 +329,60 @@ func (m *Mongo) ensureVersionTable() (err error) {
327329
// Utilizes advisory locking on the config.LockingCollection collection
328330
// This uses a unique index on the `locking_key` field.
329331
func (m *Mongo) Lock() error {
330-
if !m.config.Locking.Enabled {
331-
return nil
332-
}
333-
pid := os.Getpid()
334-
hostname, err := os.Hostname()
335-
if err != nil {
336-
hostname = fmt.Sprintf("Could not determine hostname. Error: %s", err.Error())
337-
}
332+
return database.CasRestoreOnErr(&m.isLocked, false, true, database.ErrLocked, func() error {
333+
if !m.config.Locking.Enabled {
334+
return nil
335+
}
338336

339-
newLockObj := lockObj{
340-
Key: lockKeyUniqueValue,
341-
Pid: pid,
342-
Hostname: hostname,
343-
CreatedAt: time.Now(),
344-
}
345-
operation := func() error {
346-
timeout, cancelFunc := context.WithTimeout(context.Background(), contextWaitTimeout)
347-
_, err := m.db.Collection(m.config.Locking.CollectionName).InsertOne(timeout, newLockObj)
348-
defer cancelFunc()
349-
return err
350-
}
351-
exponentialBackOff := backoff.NewExponentialBackOff()
352-
duration := time.Duration(m.config.Locking.Timeout) * time.Second
353-
exponentialBackOff.MaxElapsedTime = duration
354-
exponentialBackOff.MaxInterval = time.Duration(m.config.Locking.Interval) * time.Second
337+
pid := os.Getpid()
338+
hostname, err := os.Hostname()
339+
if err != nil {
340+
hostname = fmt.Sprintf("Could not determine hostname. Error: %s", err.Error())
341+
}
355342

356-
err = backoff.Retry(operation, exponentialBackOff)
357-
if err != nil {
358-
return database.ErrLocked
359-
}
343+
newLockObj := lockObj{
344+
Key: lockKeyUniqueValue,
345+
Pid: pid,
346+
Hostname: hostname,
347+
CreatedAt: time.Now(),
348+
}
349+
operation := func() error {
350+
timeout, cancelFunc := context.WithTimeout(context.Background(), contextWaitTimeout)
351+
_, err := m.db.Collection(m.config.Locking.CollectionName).InsertOne(timeout, newLockObj)
352+
defer cancelFunc()
353+
return err
354+
}
355+
exponentialBackOff := backoff.NewExponentialBackOff()
356+
duration := time.Duration(m.config.Locking.Timeout) * time.Second
357+
exponentialBackOff.MaxElapsedTime = duration
358+
exponentialBackOff.MaxInterval = time.Duration(m.config.Locking.Interval) * time.Second
360359

361-
return nil
360+
err = backoff.Retry(operation, exponentialBackOff)
361+
if err != nil {
362+
return database.ErrLocked
363+
}
362364

365+
return nil
366+
})
363367
}
368+
364369
func (m *Mongo) Unlock() error {
365-
if !m.config.Locking.Enabled {
366-
return nil
367-
}
370+
return database.CasRestoreOnErr(&m.isLocked, true, false, database.ErrNotLocked, func() error {
371+
if !m.config.Locking.Enabled {
372+
return nil
373+
}
368374

369-
filter := findFilter{
370-
Key: lockKeyUniqueValue,
371-
}
375+
filter := findFilter{
376+
Key: lockKeyUniqueValue,
377+
}
372378

373-
ctx, cancel := context.WithTimeout(context.Background(), contextWaitTimeout)
374-
_, err := m.db.Collection(m.config.Locking.CollectionName).DeleteMany(ctx, filter)
375-
defer cancel()
379+
ctx, cancel := context.WithTimeout(context.Background(), contextWaitTimeout)
380+
_, err := m.db.Collection(m.config.Locking.CollectionName).DeleteMany(ctx, filter)
381+
defer cancel()
376382

377-
if err != nil {
378-
return err
379-
}
380-
return nil
383+
if err != nil {
384+
return err
385+
}
386+
return nil
387+
})
381388
}

‎database/mongodb/mongodb_test.go

+1-12
Original file line numberDiff line numberDiff line change
@@ -221,18 +221,7 @@ func TestLockWorks(t *testing.T) {
221221
t.Fatal(err)
222222
}
223223

224-
// disable locking, validate wer can lock twice
225-
mc.config.Locking.Enabled = false
226-
err = mc.Lock()
227-
if err != nil {
228-
t.Fatal(err)
229-
}
230-
err = mc.Lock()
231-
if err != nil {
232-
t.Fatal(err)
233-
}
234-
235-
// re-enable locking,
224+
// enable locking,
236225
//try to hit a lock conflict
237226
mc.config.Locking.Enabled = true
238227
mc.config.Locking.Timeout = 1

0 commit comments

Comments
 (0)