ims/util/db/pgsql_lock.go

package db

import (
	"fmt"
	"hash/fnv"
	"ims/util/backoff"

	"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.
	pgsqlTryAdvisoryXactLock = "SELECT pg_try_advisory_xact_lock(?)"
)

// pgsqlLock represents a PostgreSQL transaction-level advisory pgsqlLock.
type pgsqlLock struct {
	tx      *gorm.DB
	lockKey string
}

func (p *pgsqlLock) 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", ErrExecSQL, sqlstr)
}

func (p *pgsqlLock) acquire() error {
	key := generateLockKey(p.lockKey)
	ok, err := p.execute(pgsqlTryAdvisoryXactLock, key)
	if err != nil || !ok {
		return fmt.Errorf("%w for key %s: %v", 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) (release func() error, err error) {
	l := &pgsqlLock{tx: tx, lockKey: lockKey}

	if opts == nil {
		err = l.acquire()
	} else {
		err = backoff.Retry(
			func() error { return l.acquire() },
			func(o *backoff.Options) { *o = *opts },
		)
	}

	release = func() error {
		return nil
	}

	return
}