File size: 3,883 Bytes
fea99b3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 | package lockr
import (
"context"
"errors"
"fmt"
"log/slog"
"strings"
"github.com/openmeterio/openmeter/openmeter/ent/db"
"github.com/openmeterio/openmeter/pkg/framework/entutils"
"github.com/openmeterio/openmeter/pkg/framework/transaction"
)
type LockerConfig struct {
Logger *slog.Logger
}
func (c *LockerConfig) Validate() error {
if c.Logger == nil {
return fmt.Errorf("logger is required")
}
return nil
}
// Locker is the generic interface for distributed business level locks.
type Locker struct {
cfg *LockerConfig
}
func NewLocker(cfg *LockerConfig) (*Locker, error) {
if err := cfg.Validate(); err != nil {
return nil, fmt.Errorf("invalid locker config: %w", err)
}
return &Locker{
cfg: cfg,
}, nil
}
// ErrLockTimeout is returned when a lock operation times out
var ErrLockTimeout = errors.New("lock operation timed out")
// LockForTX locks the key for the duration of the transaction.
func (l *Locker) LockForTX(ctx context.Context, key Key) error {
l.cfg.Logger.DebugContext(ctx, "locking for tx", "key", key.String(), "hash", key.Hash64())
client, err := l.getTxClient(ctx)
if err != nil {
return err
}
return l.lock(ctx, client, key)
}
func (l *Locker) LockForTXWithScopes(ctx context.Context, scopes ...string) error {
k, err := NewKey(scopes...)
if err != nil {
return err
}
return l.LockForTX(ctx, k)
}
// lock executes the advisory lock query and handles the result set
func (l *Locker) lock(ctx context.Context, client *db.Tx, key Key) error {
rows, err := client.QueryContext(ctx, "SELECT pg_advisory_xact_lock($1)", int64(key.Hash64()))
defer func() {
if rows != nil {
if e := rows.Close(); e != nil {
l.cfg.Logger.WarnContext(ctx, "failed to close result set", "error", e)
}
}
}()
if err != nil {
return checkForTimeout(err)
}
// Consume the result set
for rows.Next() {
// pg_advisory_xact_lock returns void, but we still need to iterate through rows
}
if err := rows.Err(); err != nil {
return checkForTimeout(err)
}
return nil
}
// Note: it would be great to use in-process timeouts with context.WithTimeout
// Unfortunately, due to this https://github.com/jackc/pgx/issues/2100#issuecomment-2395092552 (context cancellation resulting in query cancellation resulting in errored tx states) we rely on the pg timeout which leaves the connection intact
func checkForTimeout(err error) error {
if strings.Contains(err.Error(), pgLockTimeoutErrCode) {
return ErrLockTimeout
}
return err
}
func (l *Locker) getTxClient(ctx context.Context) (*db.Tx, error) {
// If we're not in a transaction this method has to fail
tx, err := entutils.GetDriverFromContext(ctx)
if err != nil {
return nil, fmt.Errorf("lockr only works in a transaction, but driver not found: %w", err)
}
client := db.NewTxClientFromRawConfig(ctx, *tx.GetConfig())
rows, err := client.QueryContext(ctx, "SELECT transaction_timestamp() != statement_timestamp()")
if err != nil {
return nil, fmt.Errorf("failed to check transaction status: %w", err)
}
defer func() {
if rows != nil {
if e := rows.Close(); e != nil {
l.cfg.Logger.WarnContext(ctx, "failed to close result set", "error", e)
}
}
}()
var isInTransaction bool
for rows.Next() {
err = rows.Scan(&isInTransaction)
if err != nil {
return nil, fmt.Errorf("failed to check transaction status: %w", err)
}
}
if err := rows.Err(); err != nil {
return nil, fmt.Errorf("failed to check transaction status: %w", err)
}
if !isInTransaction {
return nil, fmt.Errorf("lockr only works in a postgres transaction")
}
return client, nil
}
type noopTxCreator struct{}
var _ transaction.Creator = (*noopTxCreator)(nil)
func (n *noopTxCreator) Tx(ctx context.Context) (context.Context, transaction.Driver, error) {
return ctx, nil, fmt.Errorf("a transaction should already be accessible from the context")
}
|