Merge pull request #2156 from binwiederhier/replica-least-busy

Send reads to the least busy replica

db/db.go

@@ -17,8 +17,8 @@ const (
 )
 
 // DB wraps a primary *sql.DB and optional read replicas. All standard query/exec methods
-// delegate to the primary. The ReadOnly() method returns a *sql.DB from a healthy replica
-// (round-robin), falling back to the primary if no replicas are configured or all are unhealthy.
+// delegate to the primary. The ReadOnly() method returns a *sql.DB from the least busy healthy
+// replica, falling back to the primary if none is configured or healthy.
 type DB struct {
 	primary  *Host
 	replicas []*Host
@@ -92,38 +92,52 @@ func (d *DB) Primary() *sql.DB {
 func (d *DB) Stats() []*PoolStats {
 	stats := []*PoolStats{
 		{
-			Role:    PoolRolePrimary,
-			Addr:    d.primary.Addr,
-			Healthy: d.primary.healthy.Load(),
-			Stats:   d.primary.DB.Stats(),
+			Role:         PoolRolePrimary,
+			Addr:         d.primary.Addr,
+			Healthy:      d.primary.healthy.Load(),
+			Stats:        d.primary.DB.Stats(),
+			PingDuration: time.Duration(d.primary.ping.Load()),
 		},
 	}
 	for _, r := range d.replicas {
 		stats = append(stats, &PoolStats{
-			Role:    PoolRoleReplica,
-			Addr:    r.Addr,
-			Healthy: r.healthy.Load(),
-			Stats:   r.DB.Stats(),
+			Role:         PoolRoleReplica,
+			Addr:         r.Addr,
+			Healthy:      r.healthy.Load(),
+			Stats:        r.DB.Stats(),
+			PingDuration: time.Duration(r.ping.Load()),
 		})
 	}
 	return stats
 }
 
-// ReadOnly returns a *sql.DB suitable for read-only queries. It round-robins across healthy
-// replicas. If all replicas are unhealthy or none are configured, the primary is returned.
+// ReadOnly returns a *sql.DB suitable for read-only queries: the healthy replica with the fewest
+// connections in use, round-robin among equally busy ones. A replica that stalls piles up
+// connections and stops getting new reads. If every replica's pool is full, the read waits for a
+// connection rather than moving to the primary, which takes every write. If no replica is healthy
+// or none is configured, the primary is returned.
 func (d *DB) ReadOnly() *sql.DB {
 	if len(d.replicas) == 0 {
 		return d.primary.DB
 	}
 	n := len(d.replicas)
 	start := int(d.counter.Add(1) - 1)
+	var best *Host
+	bestInUse := 0
 	for i := 0; i < n; i++ {
 		r := d.replicas[(start+i)%n]
-		if r.healthy.Load() {
-			return r.DB
+		if !r.healthy.Load() {
+			continue
+		}
+		inUse := r.DB.Stats().InUse
+		if best == nil || inUse < bestInUse {
+			best, bestInUse = r, inUse
 		}
 	}
-	return d.primary.DB
+	if best == nil {
+		return d.primary.DB
+	}
+	return best.DB
 }
 
 // MarkReplicasHealthyForTest immediately marks all configured replicas as healthy, bypassing the
@@ -175,7 +189,9 @@ func (d *DB) checkHealth(ctx context.Context) {
 func checkHost(ctx context.Context, h *Host, kind string) {
 	wasHealthy := h.healthy.Load()
 	pingCtx, cancel := context.WithTimeout(ctx, healthCheckTimeout)
+	started := time.Now()
 	err := h.DB.PingContext(pingCtx)
+	h.ping.Store(int64(time.Since(started)))
 	cancel()
 	if err != nil {
 		h.healthy.Store(false)

db/db_test.go

@@ -93,3 +93,80 @@ func TestDB_HealthCheck(t *testing.T) {
 		return !d.Stats()[0].Healthy
 	}, 15*time.Second, 50*time.Millisecond)
 }
+
+func TestDB_ReadOnly_PrefersLeastBusyReplica(t *testing.T) {
+	// A replica that stalls piles up connections in use; new reads go to the other one
+	primary := openSQLite(t, "primary.db", 5)
+	busy := openSQLite(t, "busy.db", 5)
+	idle := openSQLite(t, "idle.db", 5)
+	d := db.New(&db.Host{DB: primary}, []*db.Host{{DB: busy}, {DB: idle}})
+	defer d.Close()
+	d.MarkReplicasHealthyForTest()
+
+	// Hold two connections on the busy replica
+	for i := 0; i < 2; i++ {
+		conn, err := busy.Conn(context.Background())
+		require.Nil(t, err)
+		defer conn.Close()
+	}
+	for i := 0; i < 10; i++ {
+		require.Same(t, idle, d.ReadOnly())
+	}
+}
+
+func TestDB_ReadOnly_RoundRobinWhenEquallyBusy(t *testing.T) {
+	primary := openSQLite(t, "primary.db", 5)
+	r1 := openSQLite(t, "r1.db", 5)
+	r2 := openSQLite(t, "r2.db", 5)
+	d := db.New(&db.Host{DB: primary}, []*db.Host{{DB: r1}, {DB: r2}})
+	defer d.Close()
+	d.MarkReplicasHealthyForTest()
+
+	seen := map[*sql.DB]int{}
+	for i := 0; i < 10; i++ {
+		seen[d.ReadOnly()]++
+	}
+	require.Equal(t, 5, seen[r1])
+	require.Equal(t, 5, seen[r2])
+}
+
+func TestDB_ReadOnly_ReplicaWhenReplicaPoolsFull(t *testing.T) {
+	// Full pools wait for a connection on the least busy replica; the primary, which takes every
+	// write, does not take the reads of stalled replicas
+	primary := openSQLite(t, "primary.db", 5)
+	r1 := openSQLite(t, "r1.db", 1)
+	r2 := openSQLite(t, "r2.db", 2)
+	d := db.New(&db.Host{DB: primary}, []*db.Host{{DB: r1}, {DB: r2}})
+	defer d.Close()
+	d.MarkReplicasHealthyForTest()
+
+	// Fill both pools, r1 with one connection, r2 with two
+	held := []*sql.DB{r1, r2, r2}
+	for _, replica := range held {
+		conn, err := replica.Conn(context.Background())
+		require.Nil(t, err)
+		defer conn.Close()
+	}
+	for i := 0; i < 10; i++ {
+		require.Same(t, r1, d.ReadOnly())
+	}
+}
+
+func TestDB_Stats_PingDuration(t *testing.T) {
+	// The health check records how long each database took to answer, for ntfy_db_ping_seconds
+	primary := openSQLite(t, "primary.db", 5)
+	replica := openSQLite(t, "replica.db", 5)
+	d := db.New(&db.Host{DB: primary}, []*db.Host{{DB: replica}})
+	defer d.Close()
+	require.Eventually(t, func() bool {
+		stats := d.Stats()
+		return stats[0].PingDuration > 0 && stats[1].PingDuration > 0
+	}, 15*time.Second, 50*time.Millisecond)
+}
+
+func openSQLite(t *testing.T, name string, maxOpen int) *sql.DB {
+	conn, err := sql.Open("sqlite3", filepath.Join(t.TempDir(), name))
+	require.Nil(t, err)
+	conn.SetMaxOpenConns(maxOpen)
+	return conn
+}

db/types.go

@@ -3,6 +3,7 @@ package db
 import (
 	"database/sql"
 	"sync/atomic"
+	"time"
 )
 
 // Beginner is an interface for types that can begin a database transaction.
@@ -22,6 +23,7 @@ type Host struct {
 	Addr    string // "host:port"
 	DB      *sql.DB
 	healthy atomic.Bool
+	ping    atomic.Int64 // Duration of the last health check ping, in nanoseconds
 }
 
 // Roles of a connection pool, see PoolStats
@@ -36,4 +38,8 @@ type PoolStats struct {
 	Addr    string // "host:port"
 	Healthy bool   // The last health check (a ping) succeeded
 	Stats   sql.DBStats
+
+	// PingDuration is how long the last health check ping took, including any wait for a free
+	// connection, so a stalled database shows even when its pool is idle; 0 before the first check
+	PingDuration time.Duration
 }

docs/config.md

@@ -437,8 +437,9 @@ You can also set this via the environment variable `NTFY_DATABASE_URL` or the co
 
 To offload read-heavy queries from the primary database, you can optionally configure one or more read replicas
 using the `database-replica-urls` option. When configured, non-critical read-only queries (e.g. fetching messages, checking access permissions, etc)
-are distributed across the replicas using round-robin, while all writes and correctness-critical reads continue to go
-to the primary. If a replica becomes unhealthy, ntfy automatically falls back to the primary until the replica recovers.
+go to the replica with the fewest connections in use (round-robin among equally busy ones), while all writes and
+correctness-critical reads continue to go to the primary. A replica that slows down piles up connections and stops
+getting new reads until it catches up. If every replica becomes unhealthy, ntfy falls back to the primary.
 You can also set this via the environment variable `NTFY_DATABASE_REPLICA_URLS` (comma-separated) or the command line
 flag `--database-replica-urls`.
 
@@ -2580,7 +2581,7 @@ variable before running the `ntfy` command (e.g. `export NTFY_LISTEN_HTTP=:80`).
 | `cluster-role`                | `NTFY_CLUSTER_ROLE`                | `active` or `passive`                               | `active`          | Experimental, not ready for production. A `passive` node never leads the cluster's singleton jobs and never closes its subscribers to send them to another node; other active nodes do not count it as a place for theirs. Use it for a node that is not behind the load balancer (check `/v1/ready`, see [Health checks](#health-checks)). While a passive node is cut off from the database, its clients can miss messages published on other nodes, so give it test clients only. |
 | `cluster-drain-delay`         | `NTFY_CLUSTER_DRAIN_DELAY`         | *duration*                                          | `15s`             | Experimental, not ready for production. When [draining on shutdown](#draining-on-shutdown), how long a node withdraws: it is not ready but still serves everything, so the load balancer stops sending it requests before the node turns read-only. |
 | `cluster-drain-duration`      | `NTFY_CLUSTER_DRAIN_DURATION`      | *duration*                                          | `1m`              | Experimental, not ready for production. When [draining on shutdown](#draining-on-shutdown), how long an active node takes to close its subscriber connections, so they reconnect to the other nodes gradually; `0` closes them all at once. The service manager's stop timeout must be longer than the delay plus this. |
-| `database-replica-urls`                    | `NTFY_DATABASE_REPLICA_URLS`                    | *list of strings (connection URLs)*                 | -                 | PostgreSQL read replica connection strings. Non-critical read-only queries are distributed across replicas (round-robin) with automatic fallback to primary. Requires `database-url`.                                                   |
+| `database-replica-urls`                    | `NTFY_DATABASE_REPLICA_URLS`                    | *list of strings (connection URLs)*                 | -                 | PostgreSQL read replica connection strings. Non-critical read-only queries go to the least busy replica, with automatic fallback to primary.   Requires `database-url`.                                                   |
 | `cache-file`                               | `NTFY_CACHE_FILE`                               | *filename*                                          | -                 | If set, messages are cached in a local SQLite database instead of only in-memory. This allows for service restarts without losing messages in support of the since= parameter. See [message cache](#message-cache).                     |
 | `cache-duration`                           | `NTFY_CACHE_DURATION`                           | *duration*                                          | 12h               | Duration for which messages will be buffered before they are deleted. This is required to support the `since=...` and `poll=1` parameter. Set this to `0` to disable the cache entirely.                                                |
 | `cache-startup-queries`                    | `NTFY_CACHE_STARTUP_QUERIES`                    | *string (SQL queries)*                              | -                 | SQL queries to run during database startup; this is useful for tuning and [enabling WAL mode](#message-cache)                                                                                                                           |

docs/releases.md

@@ -2358,6 +2358,7 @@ and the [ntfy Android app](https://github.com/binwiederhier/ntfy-android/release
 * Do not send TCP keep-alive probes on the client listeners (`listen-http`, `listen-https`): Go sends one every 15 seconds on an idle connection by default, which, without a proxy in front of ntfy, woke phones holding a stream three times as often as ntfy's own keepalive event ([#2152](https://github.com/binwiederhier/ntfy/pull/2152))
 * `ntfy cluster status` lists the node itself as the first row of the table (marked "this node", and "leader" if it leads) instead of a line above it ([#2153](https://github.com/binwiederhier/ntfy/pull/2153))
 * The cluster and admin socket options lost their `experimental-` prefix, e.g. `cluster-listen` instead of `experimental-cluster-listen` (and `NTFY_CLUSTER_LISTEN`); the old names still work ([#2154](https://github.com/binwiederhier/ntfy/pull/2154))
+* Send read-only queries to the least busy read replica instead of round-robin, so a replica that stalls stops getting new reads; the new `ntfy_db_ping_seconds` metric shows how long each database takes to answer its health check ([#2156](https://github.com/binwiederhier/ntfy/pull/2156))
 
 ### ntfy iOS app v1.8.0 (UNRELEASED)
 

metrics/metrics.go

@@ -6,6 +6,7 @@ package metrics
 import (
 	"database/sql"
 	"sync/atomic"
+	"time"
 
 	"github.com/prometheus/client_golang/prometheus"
 )
@@ -173,6 +174,8 @@ type DBPool struct {
 	Addr  string // "host:port"
 	Up    bool   // The last health check of the pool's database succeeded
 	Stats sql.DBStats
+
+	PingDuration time.Duration // How long the last health check ping took
 }
 
 // SetDBPoolsFunc sets the function that the ntfy_db_* metrics read their values from on
@@ -205,6 +208,9 @@ func newDBPoolCollector() *dbPoolCollector {
 				}
 				return 0
 			}),
+			newDBPoolMetric("ntfy_db_ping_seconds", "How long the last health check ping took (every 30s), including any wait for a free connection, per pool", prometheus.GaugeValue, func(pool *DBPool) float64 {
+				return pool.PingDuration.Seconds()
+			}),
 			newDBPoolMetric("ntfy_db_pool_connections_max", "Maximum open database connections (pool_max_conns), per pool", prometheus.GaugeValue, func(pool *DBPool) float64 {
 				return float64(pool.Stats.MaxOpenConnections)
 			}),

metrics/metrics_test.go

@@ -30,6 +30,7 @@ var expectedMetricNames = []string{
 	"ntfy_cluster_peers",
 	"ntfy_cluster_queue_dropped_total",
 	"ntfy_cluster_send_errors_total",
+	"ntfy_db_ping_seconds",
 	"ntfy_db_pool_connections_idle",
 	"ntfy_db_pool_connections_in_use",
 	"ntfy_db_pool_connections_max",
@@ -101,9 +102,10 @@ func TestDBPoolMetrics(t *testing.T) {
 	SetDBPoolsFunc(func() []*DBPool {
 		return []*DBPool{
 			{
-				Role: db.PoolRolePrimary,
-				Addr: "db:5432",
-				Up:   true,
+				Role:         db.PoolRolePrimary,
+				Addr:         "db:5432",
+				Up:           true,
+				PingDuration: 250 * time.Millisecond,
 				Stats: sql.DBStats{
 					MaxOpenConnections: 30,
 					OpenConnections:    12,
@@ -127,6 +129,10 @@ func TestDBPoolMetrics(t *testing.T) {
 	})
 	defer SetDBPoolsFunc(nil)
 	expected := fmt.Sprintf(`
+# HELP ntfy_db_ping_seconds How long the last health check ping took (every 30s), including any wait for a free connection, per pool
+# TYPE ntfy_db_ping_seconds gauge
+ntfy_db_ping_seconds{host="db:5432",role="%[1]s"} 0.25
+ntfy_db_ping_seconds{host="replica:5432",role="%[2]s"} 0
 # HELP ntfy_db_pool_connections_in_use Database connections currently in use, per pool
 # TYPE ntfy_db_pool_connections_in_use gauge
 ntfy_db_pool_connections_in_use{host="db:5432",role="%[1]s"} 5
@@ -149,7 +155,7 @@ ntfy_db_up{host="db:5432",role="%[1]s"} 1
 ntfy_db_up{host="replica:5432",role="%[2]s"} 0
 `, db.PoolRolePrimary, db.PoolRoleReplica)
 	require.Nil(t, testutil.GatherAndCompare(prometheus.DefaultGatherer, strings.NewReader(expected),
-		"ntfy_db_pool_connections_in_use", "ntfy_db_pool_connections_max", "ntfy_db_pool_waits_total", "ntfy_db_pool_wait_seconds_total", "ntfy_db_up"))
+		"ntfy_db_ping_seconds", "ntfy_db_pool_connections_in_use", "ntfy_db_pool_connections_max", "ntfy_db_pool_waits_total", "ntfy_db_pool_wait_seconds_total", "ntfy_db_up"))
 	open, err := testutil.GatherAndCount(prometheus.DefaultGatherer, "ntfy_db_pool_connections_open", "ntfy_db_pool_connections_idle")
 	require.Nil(t, err)
 	require.Equal(t, 4, open)

server/server.go

@@ -402,10 +402,11 @@ func (s *Server) dbPoolStats() []*metrics.DBPool {
 	pools := make([]*metrics.DBPool, 0)
 	for _, p := range s.db.Stats() {
 		pools = append(pools, &metrics.DBPool{
-			Role:  p.Role,
-			Addr:  p.Addr,
-			Up:    p.Healthy,
-			Stats: p.Stats,
+			Role:         p.Role,
+			Addr:         p.Addr,
+			Up:           p.Healthy,
+			Stats:        p.Stats,
+			PingDuration: p.PingDuration,
 		})
 	}
 	return pools