member-console/internal/testkit/dbtest/wait.go

// SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-Commercial
// SPDX-FileCopyrightText: 2025-2026 Christian Galo

package dbtest

import (
	"database/sql"
	"fmt"
	"testing"
	"time"
)

// lockWaitDeadline bounds how long AwaitLockWait polls, well over any wait a
// test provokes.
const lockWaitDeadline = 10 * time.Second

// Backend selects the server backend a test waits on. Exactly one of PID,
// BlockedBy and Query is set.
type Backend struct {
	PID       int    // pg_backend_pid() of the connection that is expected to wait
	BlockedBy int    // pg_backend_pid() of the connection whose lock some backend waits on
	Query     string // text the waiting backend's current statement contains, within this database
	// Waiters is how many selected backends must wait at once; zero means
	// one. Two requests queued on one row lock are Waiters 2 with Query;
	// BlockedBy counts only the first of them, because the second waits on
	// the first's tuple lock.
	Waiters int
	// Abort is called between polls. A non-nil error fails the test with its
	// text, so a test whose waiting call can return without ever waiting
	// reports that instead of polling to the deadline.
	Abort func() error
}

// waiting is the statement that counts the backends b selects that wait on a
// lock, and its argument.
func (b Backend) waiting() (string, any) {
	switch {
	case b.BlockedBy != 0:
		return `SELECT count(*) FROM pg_stat_activity WHERE $1 = ANY(pg_blocking_pids(pid))`, b.BlockedBy
	case b.Query != "":
		return `SELECT count(*) FROM pg_stat_activity
		        WHERE datname = current_database() AND pid <> pg_backend_pid()
		          AND wait_event_type = 'Lock' AND position($1 in query) > 0`, b.Query
	default:
		return `SELECT count(*) FROM pg_stat_activity WHERE pid = $1 AND wait_event_type = 'Lock'`, b.PID
	}
}

// AwaitLockWait polls pg_stat_activity until b waits on a lock (b.Waiters
// backends at once, when it is above one), so a test can assert what the
// waiting backends hold without sleeping for a fixed time. It fails the test
// when b has not waited after ten seconds or when b.Abort reports an error.
func AwaitLockWait(tb testing.TB, d *sql.DB, b Backend) {
	tb.Helper()
	query, arg := b.waiting()
	deadline := time.Now().Add(lockWaitDeadline)
	for {
		var waiting int
		if err := d.QueryRow(query, arg).Scan(&waiting); err != nil {
			tb.Fatalf("read pg_stat_activity: %v", err)
		}
		if waiting >= b.waiters() {
			return
		}
		if b.Abort != nil {
			if err := b.Abort(); err != nil {
				tb.Fatalf("%v", err)
			}
		}
		if time.Now().After(deadline) {
			tb.Fatal(b.neverWaited())
		}
		time.Sleep(5 * time.Millisecond)
	}
}

// waiters is how many selected backends AwaitLockWait waits for.
func (b Backend) waiters() int {
	return max(b.Waiters, 1)
}

// neverWaited is the failure message for a selection that not enough
// backends matched before the deadline.
func (b Backend) neverWaited() string {
	switch {
	case b.Waiters > 1 && b.BlockedBy != 0:
		return fmt.Sprintf("fewer than %d backends waited on a lock held by backend %d", b.Waiters, b.BlockedBy)
	case b.Waiters > 1:
		return fmt.Sprintf("fewer than %d backends running a statement containing %q waited on a lock", b.Waiters, b.Query)
	case b.BlockedBy != 0:
		return fmt.Sprintf("no backend waited on a lock held by backend %d", b.BlockedBy)
	case b.Query != "":
		return fmt.Sprintf("no backend running a statement containing %q waited on a lock", b.Query)
	default:
		return fmt.Sprintf("backend %d never waited on a lock", b.PID)
	}
}