From 076dc6bdba5f58fa0a3b5ebb714e9ab6153948eb Mon Sep 17 00:00:00 2001 From: aman Date: Tue, 1 Sep 2026 19:26:56 +0530 Subject: [PATCH 1/2] fix(db): run advisory lock and unlock on the same session TryLock asked for the lock through the pool, so the lock landed on a different connection than the pinned one Unlock later used and the release never worked. Run the lock query on the pinned connection, keep the release working when the caller's context is canceled, and report when the session does not hold the lock. Co-Authored-By: Claude Fable 5 --- pkg/db/db.go | 17 +++++- pkg/db/db_test.go | 147 ++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 161 insertions(+), 3 deletions(-) diff --git a/pkg/db/db.go b/pkg/db/db.go index 732f12609..e3be53a2d 100644 --- a/pkg/db/db.go +++ b/pkg/db/db.go @@ -16,7 +16,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 { @@ -86,10 +87,17 @@ type Lock struct { // Unlock uses postgres advisory locks to release a lock on a given id func (l Lock) Unlock(ctx context.Context) error { + // 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) + var errs []error - _, err := l.conn.ExecContext(ctx, "SELECT pg_advisory_unlock($1)", l.ID) + var released bool + err := l.conn.GetContext(ctx, &released, "SELECT pg_advisory_unlock($1)", l.ID) if err != nil { errs = append(errs, err) + } else if !released { + errs = append(errs, fmt.Errorf("advisory lock %d: %w", l.ID, ErrLockNotHeld)) } err = l.conn.Close() @@ -117,7 +125,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 { diff --git a/pkg/db/db_test.go b/pkg/db/db_test.go index 2f3a11128..75b5ab06b 100644 --- a/pkg/db/db_test.go +++ b/pkg/db/db_test.go @@ -5,6 +5,8 @@ import ( "database/sql" "database/sql/driver" "errors" + "io" + "sync" "testing" "github.com/jmoiron/sqlx" @@ -139,3 +141,148 @@ 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 { 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 + } + c.rec.add(c.id, query, args) + 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 +} + +type queryRecorder struct { + mu sync.Mutex + queries []recordedQuery +} + +func (r *queryRecorder) add(connID int, query string, args []driver.NamedValue) { + r.mu.Lock() + defer r.mu.Unlock() + q := recordedQuery{connID: connID, query: query} + 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...) +} + +// 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) + }) +} From 9323df4e950aa8c90685fff4e5fa984429c1c68e Mon Sep 17 00:00:00 2001 From: aman Date: Wed, 2 Sep 2026 11:15:00 +0530 Subject: [PATCH 2/2] fix(db): time-bound the advisory unlock and discard failed sessions WithoutCancel also removed the caller's deadline, so a stalled release could block forever. Cap the release with the client's query timeout, and when it fails drop the physical connection instead of pooling it, so postgres frees the lock as soon as the session ends. Co-Authored-By: Claude Fable 5 --- pkg/db/db.go | 41 +++++++++++++---------- pkg/db/db_test.go | 85 ++++++++++++++++++++++++++++++++++++++++++----- 2 files changed, 100 insertions(+), 26 deletions(-) diff --git a/pkg/db/db.go b/pkg/db/db.go index e3be53a2d..eebd571cc 100644 --- a/pkg/db/db.go +++ b/pkg/db/db.go @@ -3,6 +3,7 @@ package db import ( "context" "database/sql" + "database/sql/driver" "errors" "fmt" "math" @@ -81,8 +82,9 @@ 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 @@ -90,25 +92,29 @@ func (l Lock) Unlock(ctx context.Context) error { // 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() + } - var errs []error var released bool - err := l.conn.GetContext(ctx, &released, "SELECT pg_advisory_unlock($1)", l.ID) - if err != nil { - errs = append(errs, err) - } else if !released { - errs = append(errs, fmt.Errorf("advisory lock %d: %w", l.ID, ErrLockNotHeld)) + 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 } - err = l.conn.Close() - if err != nil { - errs = append(errs, err) + var errs []error + if !released { + errs = append(errs, fmt.Errorf("advisory lock %d: %w", l.ID, ErrLockNotHeld)) } - - if len(errs) > 0 { - return errors.Join(errs...) + 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 @@ -145,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 } diff --git a/pkg/db/db_test.go b/pkg/db/db_test.go index 75b5ab06b..a9b69ad03 100644 --- a/pkg/db/db_test.go +++ b/pkg/db/db_test.go @@ -8,6 +8,7 @@ import ( "io" "sync" "testing" + "time" "github.com/jmoiron/sqlx" "github.com/stretchr/testify/assert" @@ -154,14 +155,18 @@ type lockConn struct { } func (c *lockConn) Prepare(string) (driver.Stmt, error) { return nil, errors.New("not implemented") } -func (c *lockConn) Close() error { return nil } +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 } - c.rec.add(c.id, query, args) + 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 } @@ -183,20 +188,23 @@ func (r *boolRows) Next(dest []driver.Value) error { } type recordedQuery struct { - connID int - query string - arg driver.Value + connID int + query string + arg driver.Value + hasDeadline bool } type queryRecorder struct { - mu sync.Mutex - queries []recordedQuery + mu sync.Mutex + queries []recordedQuery + closes []int + failNext error } -func (r *queryRecorder) add(connID int, query string, args []driver.NamedValue) { +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} + q := recordedQuery{connID: connID, query: query, hasDeadline: hasDeadline} if len(args) > 0 { q.arg = args[0].Value } @@ -209,6 +217,32 @@ func (r *queryRecorder) all() []recordedQuery { 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 { @@ -285,4 +319,37 @@ func TestUnlock(t *testing.T) { 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) + }) }