WIP: Clsuter support (cross node delivery, leader election)

db/pg/leader.go

@@ -0,0 +1,83 @@
+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
+}

db/pg/leader_test.go

@@ -0,0 +1,84 @@
+package pg_test
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"heckel.io/ntfy/v2/db/pg"
+	dbtest "heckel.io/ntfy/v2/db/test"
+)
+
+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))
+	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()
+	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.
+	schemaDSN := dbtest.CreateTestPostgresSchema(t)
+	hostA, err := pg.Open(schemaDSN)
+	require.Nil(t, err)
+	defer hostA.DB.Close()
+	hostB, err := pg.Open(schemaDSN)
+	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))
+	// 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() {
+			return
+		}
+		time.Sleep(20 * time.Millisecond)
+	}
+	t.Fatal("timed out waiting for condition")
+}
+
+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
+}

db/pg/pg.go

@@ -17,6 +17,7 @@ import (
 // ntfy key is defined here, following the "ntfy"+2586+letter scheme
 const (
 	SchemaLockKey = int64(0x6e7466792586a) // Schema setup serialization (transaction-scoped, see db/schema)
+	LeaderLockKey = int64(0x6e7466792586b) // Cluster singleton-job leader (session-scoped, held for process lifetime)
 )
 
 // Open opens a PostgreSQL connection pool for a primary database. It pings the database

db/schema/schema_test.go

@@ -14,7 +14,9 @@ import (
 	_ "github.com/mattn/go-sqlite3"
 )
 
-const testCreateQuery = `CREATE TABLE IF NOT EXISTS things (id TEXT PRIMARY KEY, name TEXT NOT NULL)`
+const (
+	testCreateQuery = `CREATE TABLE IF NOT EXISTS things (id TEXT PRIMARY KEY, name TEXT NOT NULL)`
+)
 
 func testCreate(tx *sql.Tx) error {
 	_, err := tx.Exec(testCreateQuery)

db/test/test.go

@@ -17,7 +17,7 @@ const testPoolMaxConns = "2"
 // CreateTestPostgresSchema creates a temporary PostgreSQL schema and returns the DSN pointing to it.
 // It registers a cleanup function to drop the schema when the test finishes.
 // If NTFY_TEST_DATABASE_URL is not set, the test is skipped.
-func CreateTestPostgresSchema(t *testing.T) string {
+func CreateTestPostgresSchema(t testing.TB) string {
 	t.Helper()
 	dsn := os.Getenv("NTFY_TEST_DATABASE_URL")
 	if dsn == "" {
@@ -51,7 +51,7 @@ func CreateTestPostgresSchema(t *testing.T) string {
 // CreateTestPostgres creates a temporary PostgreSQL schema and returns an open *db.DB connection to it.
 // It registers cleanup functions to close the DB and drop the schema when the test finishes.
 // If NTFY_TEST_DATABASE_URL is not set, the test is skipped.
-func CreateTestPostgres(t *testing.T) *db.DB {
+func CreateTestPostgres(t testing.TB) *db.DB {
 	t.Helper()
 	schemaDSN := CreateTestPostgresSchema(t)
 	testHost, err := pg.Open(schemaDSN)