fix(db): run advisory lock and unlock on the same session

pkg/db/db.go

@@ -3,6 +3,7 @@ package db
 import (
 	"context"
 	"database/sql"
+	"database/sql/driver"
 	"errors"
 	"fmt"
 	"math"
@@ -16,7 +17,8 @@ import (
 )
 
 var (
-	ErrLockBusy = errors.New("lock busy")
+	ErrLockBusy    = errors.New("lock busy")
+	ErrLockNotHeld = errors.New("lock not held by this session")
 )
 
 type Client struct {
@@ -80,27 +82,39 @@ func (c Client) WithTxn(ctx context.Context, txnOptions sql.TxOptions, txFunc fu
 }
 
 type Lock struct {
-	ID   int64
-	conn *sqlx.Conn
+	ID      int64
+	conn    *sqlx.Conn
+	timeout time.Duration
 }
 
 // Unlock uses postgres advisory locks to release a lock on a given id
 func (l Lock) Unlock(ctx context.Context) error {
-	var errs []error
-	_, err := l.conn.ExecContext(ctx, "SELECT pg_advisory_unlock($1)", l.ID)
-	if err != nil {
-		errs = append(errs, err)
+	// the release must run even if the caller's context is already canceled,
+	// otherwise the lock stays held until the pool closes this connection
+	ctx = context.WithoutCancel(ctx)
+	if l.timeout > 0 {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(ctx, l.timeout)
+		defer cancel()
 	}
 
-	err = l.conn.Close()
-	if err != nil {
-		errs = append(errs, err)
+	var released bool
+	if err := l.conn.GetContext(ctx, &released, "SELECT pg_advisory_unlock($1)", l.ID); err != nil {
+		// the session may still hold the lock, so drop the physical
+		// connection instead of returning it to the pool: postgres releases
+		// the lock as soon as the session ends
+		_ = l.conn.Raw(func(any) error { return driver.ErrBadConn })
+		return err
 	}
 
-	if len(errs) > 0 {
-		return errors.Join(errs...)
+	var errs []error
+	if !released {
+		errs = append(errs, fmt.Errorf("advisory lock %d: %w", l.ID, ErrLockNotHeld))
+	}
+	if err := l.conn.Close(); err != nil {
+		errs = append(errs, err)
 	}
-	return nil
+	return errors.Join(errs...)
 }
 
 // TryLock uses postgres advisory locks to acquire a lock on a given id
@@ -117,7 +131,10 @@ func (c Client) TryLock(ctx context.Context, id string) (*Lock, error) {
 	intHash := int64(hash % uint64(math.MaxInt64)) // Reduce hash to fit within int64 range
 	query := "SELECT pg_try_advisory_lock($1)"
 	var acquired bool
-	if err := c.GetContext(ctx, &acquired, query, intHash); err != nil {
+	// the lock query must run on the pinned connection: an advisory lock
+	// belongs to the session that acquired it, and only that session can
+	// release it
+	if err := newConn.GetContext(ctx, &acquired, query, intHash); err != nil {
 		var errs []error
 		errs = append(errs, err)
 		if connErr := newConn.Close(); connErr != nil {
@@ -134,8 +151,9 @@ func (c Client) TryLock(ctx context.Context, id string) (*Lock, error) {
 	}
 
 	lock := &Lock{
-		ID:   intHash,
-		conn: newConn,
+		ID:      intHash,
+		conn:    newConn,
+		timeout: c.queryTimeOut,
 	}
 	return lock, nil
 }

pkg/db/db_test.go

@@ -5,7 +5,10 @@ import (
 	"database/sql"
 	"database/sql/driver"
 	"errors"
+	"io"
+	"sync"
 	"testing"
+	"time"
 
 	"github.com/jmoiron/sqlx"
 	"github.com/stretchr/testify/assert"
@@ -139,3 +142,214 @@ func TestWithTxn(t *testing.T) {
 		})
 	})
 }
+
+// lockConn is a fake connection for the advisory lock tests. It answers the
+// try-lock query with a fixed result and records every query together with
+// the id of the connection that ran it, so tests can check that lock and
+// unlock happen on the same session. It refuses to run queries on a canceled
+// context, like a real driver.
+type lockConn struct {
+	id       int
+	rec      *queryRecorder
+	acquired bool
+}
+
+func (c *lockConn) Prepare(string) (driver.Stmt, error) { return nil, errors.New("not implemented") }
+func (c *lockConn) Close() error                        { c.rec.closed(c.id); return nil }
+func (c *lockConn) Begin() (driver.Tx, error)           { return nil, errors.New("not implemented") }
+
+func (c *lockConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+	if err := c.rec.takeFailure(); err != nil {
+		return nil, err
+	}
+	_, hasDeadline := ctx.Deadline()
+	c.rec.add(c.id, query, args, hasDeadline)
+	return &boolRows{value: c.acquired}, nil
+}
+
+// boolRows is a result set with a single row holding a single boolean column.
+type boolRows struct {
+	value bool
+	done  bool
+}
+
+func (r *boolRows) Columns() []string { return []string{"acquired"} }
+func (r *boolRows) Close() error      { return nil }
+func (r *boolRows) Next(dest []driver.Value) error {
+	if r.done {
+		return io.EOF
+	}
+	r.done = true
+	dest[0] = r.value
+	return nil
+}
+
+type recordedQuery struct {
+	connID      int
+	query       string
+	arg         driver.Value
+	hasDeadline bool
+}
+
+type queryRecorder struct {
+	mu       sync.Mutex
+	queries  []recordedQuery
+	closes   []int
+	failNext error
+}
+
+func (r *queryRecorder) add(connID int, query string, args []driver.NamedValue, hasDeadline bool) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	q := recordedQuery{connID: connID, query: query, hasDeadline: hasDeadline}
+	if len(args) > 0 {
+		q.arg = args[0].Value
+	}
+	r.queries = append(r.queries, q)
+}
+
+func (r *queryRecorder) all() []recordedQuery {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	return append([]recordedQuery(nil), r.queries...)
+}
+
+func (r *queryRecorder) closed(connID int) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.closes = append(r.closes, connID)
+}
+
+func (r *queryRecorder) closedConns() []int {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	return append([]int(nil), r.closes...)
+}
+
+func (r *queryRecorder) failNextQuery(err error) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.failNext = err
+}
+
+func (r *queryRecorder) takeFailure() error {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	err := r.failNext
+	r.failNext = nil
+	return err
+}
+
+// lockConnector hands out a fresh numbered connection on every Connect call,
+// the way a real pool dials new sessions.
+type lockConnector struct {
+	rec      *queryRecorder
+	acquired bool
+	mu       sync.Mutex
+	next     int
+}
+
+func (f *lockConnector) Connect(context.Context) (driver.Conn, error) {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	f.next++
+	return &lockConn{id: f.next, rec: f.rec, acquired: f.acquired}, nil
+}
+func (f *lockConnector) Driver() driver.Driver { return nil }
+
+func newLockClient(t *testing.T, acquired bool) (Client, *queryRecorder) {
+	t.Helper()
+	rec := &queryRecorder{}
+	client := Client{DB: sqlx.NewDb(sql.OpenDB(&lockConnector{rec: rec, acquired: acquired}), "postgres")}
+	t.Cleanup(func() { _ = client.Close() })
+	return client, rec
+}
+
+func TestTryLock(t *testing.T) {
+	t.Run("acquires and releases on the same connection", func(t *testing.T) {
+		client, rec := newLockClient(t, true)
+
+		lock, err := client.TryLock(context.Background(), "some-job")
+		require.NoError(t, err)
+		require.NotNil(t, lock)
+		require.NoError(t, lock.Unlock(context.Background()))
+
+		queries := rec.all()
+		require.Len(t, queries, 2)
+		assert.Contains(t, queries[0].query, "pg_try_advisory_lock")
+		assert.Contains(t, queries[1].query, "pg_advisory_unlock")
+		assert.Equal(t, queries[0].connID, queries[1].connID)
+		assert.Equal(t, queries[0].arg, queries[1].arg)
+	})
+
+	t.Run("returns ErrLockBusy when the lock is already held", func(t *testing.T) {
+		client, _ := newLockClient(t, false)
+
+		lock, err := client.TryLock(context.Background(), "some-job")
+		assert.ErrorIs(t, err, ErrLockBusy)
+		assert.Nil(t, lock)
+	})
+}
+
+func TestUnlock(t *testing.T) {
+	t.Run("releases even when the caller's context is canceled", func(t *testing.T) {
+		client, rec := newLockClient(t, true)
+
+		lock, err := client.TryLock(context.Background(), "some-job")
+		require.NoError(t, err)
+
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+		require.NoError(t, lock.Unlock(ctx))
+
+		queries := rec.all()
+		require.Len(t, queries, 2)
+		assert.Contains(t, queries[1].query, "pg_advisory_unlock")
+	})
+
+	t.Run("reports when the session does not hold the lock", func(t *testing.T) {
+		client, _ := newLockClient(t, false)
+
+		conn, err := client.Connx(context.Background())
+		require.NoError(t, err)
+		lock := &Lock{ID: 42, conn: conn}
+
+		assert.ErrorIs(t, lock.Unlock(context.Background()), ErrLockNotHeld)
+	})
+
+	t.Run("discards the connection when the release fails", func(t *testing.T) {
+		client, rec := newLockClient(t, true)
+
+		lock, err := client.TryLock(context.Background(), "some-job")
+		require.NoError(t, err)
+
+		queryErr := errors.New("network down")
+		rec.failNextQuery(queryErr)
+		assert.ErrorIs(t, lock.Unlock(context.Background()), queryErr)
+
+		queries := rec.all()
+		require.NotEmpty(t, queries)
+		assert.Contains(t, rec.closedConns(), queries[0].connID)
+	})
+
+	t.Run("applies the client's query timeout to the release", func(t *testing.T) {
+		rec := &queryRecorder{}
+		client := Client{
+			DB:           sqlx.NewDb(sql.OpenDB(&lockConnector{rec: rec, acquired: true}), "postgres"),
+			queryTimeOut: time.Second,
+		}
+		t.Cleanup(func() { _ = client.Close() })
+
+		lock, err := client.TryLock(context.Background(), "some-job")
+		require.NoError(t, err)
+		require.NoError(t, lock.Unlock(context.Background()))
+
+		queries := rec.all()
+		require.Len(t, queries, 2)
+		assert.Contains(t, queries[1].query, "pg_advisory_unlock")
+		assert.True(t, queries[1].hasDeadline)
+	})
+}