Refactor

db/pg/leader.go

@@ -4,80 +4,160 @@ import (
 	"context"
 	"database/sql"
 	"sync"
+	"time"
+
+	"heckel.io/ntfy/v2/log"
 )
 
 const (
+	tagLeader = "leader"
+
 	tryAdvisoryLockQuery = `SELECT pg_try_advisory_lock($1)`
 	advisoryUnlockQuery  = `SELECT pg_advisory_unlock($1)`
+
+	defaultRenewInterval = 5 * time.Second
+	leaderMissedRenewals = 3
+	leaderHoldoffFactor  = 2
 )
 
-// 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.
+// Leader implements singleton-job leader election via a Postgres advisory lock held on a
+// pinned connection. The lock auto-releases when the holding connection dies, so a crashed
+// leader is replaced without manual fencing; distinct keys elect independently. The Leader
+// renews its lease on its own loop; callers only ask IsLeader and eventually Close.
+//
+// Holding the lock is not the same as believing to be the leader: IsLeader also requires a
+// recent renewal (lease duration) and a completed hold-off after winning the lock. The
+// hold-off outlasts the lease duration by construction, so on failover the old belief always
+// expires before the new one begins: a short no-leader gap, never two leaders. Defaults:
+// renew every 5s, lease duration 15s, hold-off 30s -> up to ~35s without a leader.
 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
+	db            *sql.DB
+	key           int64
+	renewInterval time.Duration
+	conn          *sql.Conn          // holds the advisory lock while this process is leader
+	acquiredAt    time.Time          // When the lock was won (this tenure), for the hold-off
+	renewedAt     time.Time          // Last successful renewal, for the lease duration; zero = lock not held
+	cancel        context.CancelFunc // Stops the renew loop and aborts its in-flight query on Close
+	closeOnce     sync.Once
+	wg            sync.WaitGroup
+	mu            sync.Mutex // Protects conn, acquiredAt and renewedAt
 }
 
-// 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}
+// NewLeader creates a Leader competing for the lock identified by key and starts its renew
+// loop. renewInterval is for tests; pass 0 for the default.
+func NewLeader(db *sql.DB, key int64, renewInterval time.Duration) *Leader {
+	if renewInterval <= 0 {
+		renewInterval = defaultRenewInterval
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	l := &Leader{
+		db:            db,
+		key:           key,
+		renewInterval: renewInterval,
+		cancel:        cancel,
+	}
+	l.wg.Add(1)
+	go l.runAcquireOrRenewLoop(ctx)
+	return l
 }
 
-// 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 {
+// IsLeader reports whether this process should act as the leader: lock held, lease renewed
+// recently, hold-off elapsed (see the Leader doc comment).
+func (l *Leader) IsLeader() 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
+	defer l.mu.Unlock()
+	leaseDuration := leaderMissedRenewals * l.renewInterval
+	holdoff := leaderHoldoffFactor * leaseDuration
+	return time.Since(l.renewedAt) < leaseDuration && time.Since(l.acquiredAt) >= holdoff
+}
+
+// Close stops competing for leadership and releases the lock. Idempotent.
+func (l *Leader) Close() {
+	l.closeOnce.Do(func() {
+		l.cancel() // Also aborts an in-flight renewal query
+		l.wg.Wait()
+		if l.IsLeader() {
+			log.Tag(tagLeader).Info("Lost leadership: closed (lock key %d)", l.key)
 		}
-		l.Release() // Connection died; the lock is already gone, re-acquire below
+		l.release()
+	})
+}
+
+// runAcquireOrRenewLoop acquires or renews the lock every renewInterval until ctx is canceled
+func (l *Leader) runAcquireOrRenewLoop(ctx context.Context) {
+	defer l.wg.Done()
+	ticker := time.NewTicker(l.renewInterval)
+	defer ticker.Stop()
+	wasLeader := false
+	for {
+		attemptCtx, cancel := context.WithTimeout(ctx, l.renewInterval)
+		l.tryAcquireOrRenew(attemptCtx)
+		cancel()
+		if isLeader := l.IsLeader(); isLeader != wasLeader {
+			wasLeader = isLeader
+			if isLeader {
+				log.Tag(tagLeader).Info("Became leader (lock key %d)", l.key)
+			} else {
+				log.Tag(tagLeader).Info("Lost leadership (lock key %d)", l.key)
+			}
+		}
+		select {
+		case <-ticker.C:
+		case <-ctx.Done():
+			return
+		}
+	}
+}
+
+// tryAcquireOrRenew renews the lock on a healthy leader (a cheap ping) or retries acquiring
+// it on a follower, on a pinned connection.
+func (l *Leader) tryAcquireOrRenew(ctx context.Context) {
+	l.mu.Lock()
+	conn := l.conn
+	l.mu.Unlock()
+	if conn != nil {
+		if conn.PingContext(ctx) == nil {
+			// Still holding the lock, connection healthy: renew the lease
+			l.mu.Lock()
+			l.renewedAt = time.Now()
+			l.mu.Unlock()
+			log.Tag(tagLeader).Trace("Renewed leader lease (lock key %d)", l.key)
+			return
+		}
+		log.Tag(tagLeader).Debug("Leader lock connection died, lock lost (lock key %d)", l.key)
+		l.release() // Connection died; the lock is already gone, re-acquire below
 	}
 	newConn, err := l.db.Conn(ctx)
 	if err != nil {
-		return false
+		log.Tag(tagLeader).Debug("Cannot get connection to compete for leader lock (lock key %d): %s", l.key, err.Error())
+		return
 	}
 	var acquired bool
 	if err := newConn.QueryRowContext(ctx, tryAdvisoryLockQuery, l.key).Scan(&acquired); err != nil || !acquired {
 		newConn.Close()
-		return false
+		log.Tag(tagLeader).Trace("Leader lock held elsewhere (lock key %d)", l.key)
+		return
 	}
+	log.Tag(tagLeader).Debug("Acquired leader lock (lock key %d); leadership after the hold-off", l.key)
 	l.mu.Lock()
 	l.conn = newConn
-	l.held = true
+	l.acquiredAt = time.Now()
+	l.renewedAt = l.acquiredAt
 	l.mu.Unlock()
-	return true
 }
 
-// Release unlocks the advisory lock and returns the pinned connection to the pool.
-func (l *Leader) Release() {
+// 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.renewedAt = time.Time{} // Zero revokes belief; without it, IsLeader would linger a lease duration
 	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.
+		// Unlock explicitly: sql.Conn.Close() returns the connection to the pool, so the
+		// session-scoped lock would otherwise stay held
 		conn.ExecContext(context.Background(), advisoryUnlockQuery, l.key)
 		conn.Close()
+		log.Tag(tagLeader).Debug("Released leader lock (lock key %d)", l.key)
 	}
 }
-
-// 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
-}