You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
ims/util/db/pgsql/lock.go

63 lines
1.5 KiB

package pgsql
import (
"fmt"
"hash/fnv"
"ims/util/backoff"
"ims/util/db"
"gorm.io/gorm"
)
// SQL statements for PostgreSQL advisory locks.
// https://www.postgresql.org/docs/16/functions-admin.html#FUNCTIONS-ADVISORY-LOCKS
const (
// pg_try_advisory_xact_lock ( key bigint ) → boolean.
sqlTryAdvisoryXactLock = "SELECT pg_try_advisory_xact_lock(?)"
)
// lock represents a PostgreSQL transaction-level advisory lock.
type lock struct {
tx *gorm.DB
lockKey string
}
func (p *lock) execute(sqlstr string, args ...any) (bool, error) {
var result bool
if err := p.tx.Raw(sqlstr, args...).Scan(&result).Error; err == nil && result {
return true, nil
}
return false, fmt.Errorf("%w: %s", db.ErrExecSQL, sqlstr)
}
func (p *lock) acquire() error {
key := GenerateLockKey(p.lockKey)
ok, err := p.execute(sqlTryAdvisoryXactLock, key)
if err != nil || !ok {
return fmt.Errorf("%w for key %s: %v", db.ErrAcquireLock, p.lockKey, err)
}
return nil
}
func GenerateLockKey(s string) int64 {
hasher := fnv.New64a()
hasher.Write([]byte(s))
hash := hasher.Sum64()
return int64(hash)
}
// AcquireXact acquires a PostgreSQL transaction-level advisory lock.
// The caller is responsible for ensuring that a transaction is active,
// and that the lock is released after use.
func AcquireXact(tx *gorm.DB, lockKey string, opts *backoff.Options) error {
l := &lock{tx: tx, lockKey: lockKey}
if opts == nil {
return l.acquire()
}
return backoff.Retry(
func() error { return l.acquire() },
func(o *backoff.Options) { *o = *opts },
)
}