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

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

package dbtest

import (
	"context"
	"database/sql"
	"errors"
	"strings"
	"testing"
	"time"
)

// holdAndWait holds an advisory lock on one transaction and parks a second
// connection on it, returning both backends' pids and the statement the
// second one runs. The caller rolls the holder back to release the waiter.
func holdAndWait(t *testing.T) (holderPID, waiterPID int, stmt string, release func()) {
	t.Helper()
	d := Open(t)
	ctx := context.Background()
	key := time.Now().UnixNano()
	stmt = "SELECT pg_advisory_xact_lock($1) /* dbtest wait test */"

	holder, err := d.BeginTx(ctx, nil)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = holder.Rollback() })
	if err := holder.QueryRowContext(ctx, `SELECT pg_backend_pid()`).Scan(&holderPID); err != nil {
		t.Fatal(err)
	}
	if _, err := holder.ExecContext(ctx, stmt, key); err != nil {
		t.Fatal(err)
	}

	waiter, err := d.BeginTx(ctx, nil)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = waiter.Rollback() })
	if err := waiter.QueryRowContext(ctx, `SELECT pg_backend_pid()`).Scan(&waiterPID); err != nil {
		t.Fatal(err)
	}
	done := make(chan error, 1)
	go func() {
		_, err := waiter.ExecContext(ctx, stmt, key)
		done <- err
	}()
	return holderPID, waiterPID, stmt, func() {
		if err := holder.Rollback(); err != nil {
			t.Fatal(err)
		}
		if err := <-done; err != nil {
			t.Fatal(err)
		}
	}
}

func TestAwaitLockWaitSelectsTheWaitingBackendThreeWays(t *testing.T) {
	holderPID, waiterPID, _, release := holdAndWait(t)
	defer release()
	d := Open(t)

	for name, b := range map[string]Backend{
		"by pid":        {PID: waiterPID},
		"by the holder": {BlockedBy: holderPID},
		"by statement":  {Query: "dbtest wait test"},
	} {
		t.Run(name, func(t *testing.T) { AwaitLockWait(t, d, b) })
	}
}

func TestAwaitLockWaitFailsWithTheAbortProbesText(t *testing.T) {
	d := Open(t)
	r := &recorder{TB: t}
	func() {
		defer func() {
			if p := recover(); p != nil {
				if _, ok := p.(stopped); !ok {
					panic(p)
				}
			}
		}()
		AwaitLockWait(r, d, Backend{PID: -1, Abort: func() error { return errors.New("the call returned without waiting") }})
	}()
	if !strings.Contains(r.failure, "the call returned without waiting") {
		t.Errorf("Fatalf = %q, want the probe's text", r.failure)
	}
}

// TestAwaitLockWaitCountsTheWaiters: with Waiters 2 the wait does not end
// while one backend waits on a held lock; it ends once a second one queues
// behind it. The probe starts the second waiter on its first call, so a wait
// that ended on one waiter never calls it.
func TestAwaitLockWaitCountsTheWaiters(t *testing.T) {
	d := Open(t)
	ctx := context.Background()
	key := time.Now().UnixNano()
	const stmt = "SELECT pg_advisory_xact_lock($1) /* dbtest waiters test */"
	// The holder is first, so a failing test releases it before it waits for
	// the waiters' transactions.
	var open []*sql.Tx
	t.Cleanup(func() {
		for _, tx := range open {
			_ = tx.Rollback()
		}
	})
	begin := func() *sql.Tx {
		tx, err := d.BeginTx(ctx, nil)
		if err != nil {
			t.Fatal(err)
		}
		open = append(open, tx)
		return tx
	}
	holder := begin()
	if _, err := holder.ExecContext(ctx, stmt, key); err != nil {
		t.Fatal(err)
	}
	done := make(chan error, 2)
	wait := func(tx *sql.Tx) {
		_, err := tx.ExecContext(ctx, stmt, key)
		if rerr := tx.Rollback(); err == nil {
			err = rerr
		}
		done <- err
	}
	go wait(begin())
	AwaitLockWait(t, d, Backend{Query: "dbtest waiters test"})

	second := begin()
	probes := 0
	AwaitLockWait(t, d, Backend{Query: "dbtest waiters test", Waiters: 2, Abort: func() error {
		if probes == 0 {
			go wait(second)
		}
		probes++
		return nil
	}})
	if probes == 0 {
		t.Fatal("the wait for two waiters ended while one waited")
	}
	if err := holder.Rollback(); err != nil {
		t.Fatal(err)
	}
	for range 2 {
		if err := <-done; err != nil {
			t.Fatal(err)
		}
	}
}