mirror of
https://github.com/multipleof4/ntfy.git
synced 2026-10-08 21:05:21 +00:00
84 lines
2.5 KiB
Go
84 lines
2.5 KiB
Go
package pg
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"sync"
|
|
)
|
|
|
|
const (
|
|
tryAdvisoryLockQuery = `SELECT pg_try_advisory_lock($1)`
|
|
advisoryUnlockQuery = `SELECT pg_advisory_unlock($1)`
|
|
)
|
|
|
|
// Leader implements singleton-job leader election via a Postgres advisory lock held on a pinned
|
|
// connection. The lock auto-releases if the holding connection dies, so a crashed leader is
|
|
// replaced without manual fencing or lease bookkeeping. Multiple independent leaderships can
|
|
// coexist by using distinct lock keys.
|
|
type Leader struct {
|
|
db *sql.DB
|
|
key int64
|
|
conn *sql.Conn // holds the advisory lock while this process is leader
|
|
held bool
|
|
mu sync.Mutex // Protects conn and held
|
|
}
|
|
|
|
// NewLeader creates a Leader competing for the advisory lock identified by key. It does not
|
|
// attempt to acquire the lock; call TryAcquire periodically.
|
|
func NewLeader(db *sql.DB, key int64) *Leader {
|
|
return &Leader{db: db, key: key}
|
|
}
|
|
|
|
// TryAcquire attempts to grab (or confirm) the advisory lock on a pinned connection, without
|
|
// blocking. It returns whether this process is the leader after the attempt, i.e. it acquired
|
|
// the lock now or still holds it. It is meant to be called periodically; on a healthy leader it
|
|
// is a cheap ping, on a follower it retries the lock.
|
|
func (l *Leader) TryAcquire(ctx context.Context) bool {
|
|
l.mu.Lock()
|
|
conn, held := l.conn, l.held
|
|
l.mu.Unlock()
|
|
if held {
|
|
if conn != nil && conn.PingContext(ctx) == nil {
|
|
return true // Still leader, connection healthy
|
|
}
|
|
l.Release() // Connection died; the lock is already gone, re-acquire below
|
|
}
|
|
newConn, err := l.db.Conn(ctx)
|
|
if err != nil {
|
|
return false
|
|
}
|
|
var acquired bool
|
|
if err := newConn.QueryRowContext(ctx, tryAdvisoryLockQuery, l.key).Scan(&acquired); err != nil || !acquired {
|
|
newConn.Close()
|
|
return false
|
|
}
|
|
l.mu.Lock()
|
|
l.conn = newConn
|
|
l.held = true
|
|
l.mu.Unlock()
|
|
return true
|
|
}
|
|
|
|
// Release unlocks the advisory lock and returns the pinned connection to the pool.
|
|
func (l *Leader) Release() {
|
|
l.mu.Lock()
|
|
conn := l.conn
|
|
l.conn = nil
|
|
l.held = false
|
|
l.mu.Unlock()
|
|
if conn != nil {
|
|
// Unlock explicitly before returning the connection to the pool: sql.Conn.Close() returns
|
|
// the physical connection to the pool rather than closing it, so the session-scoped lock
|
|
// would otherwise stay held.
|
|
conn.ExecContext(context.Background(), advisoryUnlockQuery, l.key)
|
|
conn.Close()
|
|
}
|
|
}
|
|
|
|
// IsLeader reports whether this process currently holds the advisory lock.
|
|
func (l *Leader) IsLeader() bool {
|
|
l.mu.Lock()
|
|
defer l.mu.Unlock()
|
|
return l.held
|
|
}
|