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
-}

db/pg/leader_internal_test.go

@@ -0,0 +1,41 @@
+package pg
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+// The lease logic is pure time arithmetic, so it is unit-tested here without a database; the
+// external leader tests cover the loop end to end.
+
+func TestLeader_Lease_HoldoffMeansNoLeaderRatherThanTwo(t *testing.T) {
+	// Freshly acquired lock: belief must wait out the hold-off
+	l := &Leader{renewInterval: 20 * time.Second} // Lease duration 1m, hold-off 2m
+	l.acquiredAt = time.Now()
+	l.renewedAt = l.acquiredAt
+	require.False(t, l.IsLeader())
+	// Once the hold-off has passed (and verification is fresh), belief begins
+	l.acquiredAt = time.Now().Add(-3 * time.Minute)
+	l.renewedAt = time.Now()
+	require.True(t, l.IsLeader())
+}
+
+func TestLeader_Lease_ExpiredLeaseRevokesLeadership(t *testing.T) {
+	// A leader that cannot renew its lease (wedged process, long GC pause) must stop
+	// believing once the lease expires, even though the lock may still be held
+	l := &Leader{renewInterval: 20 * time.Second} // Lease duration 1m, hold-off 2m
+	l.acquiredAt = time.Now().Add(-time.Hour)
+	l.renewedAt = time.Now().Add(-2 * time.Minute) // Lease expired
+	require.False(t, l.IsLeader())
+	l.renewedAt = time.Now() // Fresh renewal restores belief
+	require.True(t, l.IsLeader())
+}
+
+func TestLeader_Lease_ReleasedIsNeverLeader(t *testing.T) {
+	// release() zeroes renewedAt, which fails the lease check no matter how old the tenure
+	l := &Leader{renewInterval: 20 * time.Second} // Lease duration 1m, hold-off 2m
+	l.acquiredAt = time.Now().Add(-time.Hour)
+	require.False(t, l.IsLeader())
+}

db/pg/leader_test.go

@@ -1,7 +1,6 @@
 package pg_test
 
 import (
-	"context"
 	"testing"
 	"time"
 
@@ -10,33 +9,33 @@ import (
 	dbtest "heckel.io/ntfy/v2/db/test"
 )
 
+const testRenewInterval = 20 * time.Millisecond // Lease duration 60ms, hold-off 120ms
+
 func TestLeader_AcquireAndFailover(t *testing.T) {
 	testDB := dbtest.CreateTestPostgres(t) // skips if NTFY_TEST_DATABASE_URL is unset
 	const key = int64(42)
-	ctx := context.Background()
-	l1 := pg.NewLeader(testDB.Primary(), key)
-	l2 := pg.NewLeader(testDB.Primary(), key)
-	defer l1.Release()
-	defer l2.Release()
-	// First to try wins; the second stays follower
-	require.True(t, l1.TryAcquire(ctx))
-	require.False(t, l2.TryAcquire(ctx))
+	l1 := pg.NewLeader(testDB.Primary(), key, testRenewInterval)
+	defer l1.Close()
+	// Belief follows the hold-off, it is never instant
+	require.False(t, l1.IsLeader())
+	waitForLeader(t, l1)
+	// A competitor never becomes leader while the leader lives
+	l2 := pg.NewLeader(testDB.Primary(), key, testRenewInterval)
+	defer l2.Close()
+	time.Sleep(300 * time.Millisecond) // Several verification rounds
+	require.False(t, l2.IsLeader())
 	require.True(t, l1.IsLeader())
-	require.False(t, l2.IsLeader())
-	// Repeated TryAcquire on the leader is a no-op ping and reports leadership
-	require.True(t, l1.TryAcquire(ctx))
-	// Release -> the follower can take over
-	l1.Release()
+	// Close -> the follower takes over
+	l1.Close()
+	require.False(t, l1.IsLeader())
+	waitForLeader(t, l2)
 	require.False(t, l1.IsLeader())
-	require.True(t, l2.TryAcquire(ctx))
-	require.True(t, l2.IsLeader())
 }
 
 func TestLeader_ConnectionLossFailover(t *testing.T) {
-	// A leader that dies without calling Release (crash, network loss) must not wedge the
-	// cluster: the advisory lock is session-scoped, so Postgres releases it when the pinned
-	// connection dies, and a follower can take over. Simulated by terminating the lock-holding
-	// backend server-side.
+	// A crashed leader must not wedge the cluster: Postgres releases the session-scoped lock
+	// when the pinned connection dies (simulated by terminating the backend), and someone
+	// re-acquires. Either node may win; the invariant is one leader eventually, never two.
 	schemaDSN := dbtest.CreateTestPostgresSchema(t)
 	hostA, err := pg.Open(schemaDSN)
 	require.Nil(t, err)
@@ -45,40 +44,47 @@ func TestLeader_ConnectionLossFailover(t *testing.T) {
 	require.Nil(t, err)
 	defer hostB.DB.Close()
 	const key = int64(43)
-	ctx := context.Background()
-	l1 := pg.NewLeader(hostA.DB, key)
-	l2 := pg.NewLeader(hostB.DB, key)
-	defer l1.Release()
-	defer l2.Release()
-	require.True(t, l1.TryAcquire(ctx))
-	require.False(t, l2.TryAcquire(ctx))
+	l1 := pg.NewLeader(hostA.DB, key, testRenewInterval)
+	defer l1.Close()
+	waitForLeader(t, l1)
+	l2 := pg.NewLeader(hostB.DB, key, testRenewInterval)
+	defer l2.Close()
 	// Kill the backend holding the lock (advisory lock keys map to classid/objid)
 	_, err = hostB.DB.Exec(`SELECT pg_terminate_backend(pid) FROM pg_locks WHERE locktype = 'advisory' AND objid = $1 AND granted`, key)
 	require.Nil(t, err)
-	// The lock is auto-released; the follower takes over
-	waitForCond(t, func() bool { return l2.TryAcquire(ctx) })
-	// The old leader discovers its dead connection on the next attempt and stays follower
-	require.False(t, l1.TryAcquire(ctx))
-}
-
-func waitForCond(t *testing.T, f func() bool) {
-	t.Helper()
-	for i := 0; i < 100; i++ {
-		if f() {
+	// Eventually exactly one leader again, and never two along the way
+	deadline := time.Now().Add(5 * time.Second)
+	for time.Now().Before(deadline) {
+		leader1, leader2 := l1.IsLeader(), l2.IsLeader()
+		require.False(t, leader1 && leader2, "two leaders at once")
+		if leader1 != leader2 {
 			return
 		}
-		time.Sleep(20 * time.Millisecond)
+		time.Sleep(10 * time.Millisecond)
 	}
-	t.Fatal("timed out waiting for condition")
+	t.Fatal("no leader re-emerged after connection loss")
 }
 
 func TestLeader_DistinctKeysAreIndependent(t *testing.T) {
 	testDB := dbtest.CreateTestPostgres(t)
-	ctx := context.Background()
-	l1 := pg.NewLeader(testDB.Primary(), 1)
-	l2 := pg.NewLeader(testDB.Primary(), 2)
-	defer l1.Release()
-	defer l2.Release()
-	require.True(t, l1.TryAcquire(ctx))
-	require.True(t, l2.TryAcquire(ctx)) // Different keys do not compete
+	l1 := pg.NewLeader(testDB.Primary(), 1, testRenewInterval)
+	defer l1.Close()
+	l2 := pg.NewLeader(testDB.Primary(), 2, testRenewInterval)
+	defer l2.Close()
+	// Different keys do not compete: both become effective leaders
+	waitForLeader(t, l1)
+	waitForLeader(t, l2)
+}
+
+// waitForLeader waits until the node believes it is the leader, or fails the test
+func waitForLeader(t *testing.T, l *pg.Leader) {
+	t.Helper()
+	deadline := time.Now().Add(5 * time.Second)
+	for time.Now().Before(deadline) {
+		if l.IsLeader() {
+			return
+		}
+		time.Sleep(10 * time.Millisecond)
+	}
+	t.Fatal("node never became effective leader")
 }