kvdb/postgres: fix context cancellation

docs/release-notes/release-notes-0.14.2.md

@@ -45,6 +45,8 @@
 * [Clarify log message about not running within
   systemd](https://github.com/lightningnetwork/lnd/pull/6096)
 
+* [Fix Postgres context cancellation](https://github.com/lightningnetwork/lnd/pull/6108)
+
 ## RPC Server
 
 * [ChanStatusFlags is now

kvdb/postgres/fixture.go

@@ -10,6 +10,7 @@ import (
 	"encoding/hex"
 	"fmt"
 	"strings"
+	"time"
 
 	"github.com/btcsuite/btcwallet/walletdb"
 	embeddedpostgres "github.com/fergusstrange/embedded-postgres"
@@ -82,6 +83,7 @@ func NewFixture(dbName string) (*fixture, error) {
 		context.Background(),
 		&Config{
 			Dsn:     dsn,
+			Timeout: time.Minute,
 		},
 		prefix,
 	)

kvdb/postgres/readwrite_bucket.go

@@ -75,10 +75,12 @@ func (b *readWriteBucket) Get(key []byte) []byte {
 	}
 
 	var value *[]byte
-	err := b.tx.QueryRow(
+	row, cancel := b.tx.QueryRow(
 		"SELECT value FROM "+b.table+" WHERE "+parentSelector(b.id)+
 			" AND key=$1", key,
-	).Scan(&value)
+	)
+	defer cancel()
+	err := row.Scan(&value)
 
 	switch {
 	case err == sql.ErrNoRows:
@@ -106,10 +108,12 @@ func (b *readWriteBucket) NestedReadWriteBucket(
 	}
 
 	var id int64
-	err := b.tx.QueryRow(
+	row, cancel := b.tx.QueryRow(
 		"SELECT id FROM "+b.table+" WHERE "+parentSelector(b.id)+
 			" AND key=$1 AND value IS NULL", key,
-	).Scan(&id)
+	)
+	defer cancel()
+	err := row.Scan(&id)
 
 	switch {
 	case err == sql.ErrNoRows:
@@ -139,10 +143,12 @@ func (b *readWriteBucket) CreateBucket(key []byte) (
 		value *[]byte
 		id    int64
 	)
-	err := b.tx.QueryRow(
+	row, cancel := b.tx.QueryRow(
 		"SELECT id,value FROM "+b.table+" WHERE "+parentSelector(b.id)+
 			" AND key=$1", key,
-	).Scan(&id, &value)
+	)
+	defer cancel()
+	err := row.Scan(&id, &value)
 
 	switch {
 	case err == sql.ErrNoRows:
@@ -159,10 +165,12 @@ func (b *readWriteBucket) CreateBucket(key []byte) (
 
 	// Bucket does not yet exist, so create it. Postgres will generate a
 	// bucket id for the new bucket.
-	err = b.tx.QueryRow(
+	row, cancel = b.tx.QueryRow(
 		"INSERT INTO "+b.table+" (parent_id, key) "+
 			"VALUES($1, $2) RETURNING id", b.id, key,
-	).Scan(&id)
+	)
+	defer cancel()
+	err = row.Scan(&id)
 	if err != nil {
 		return nil, err
 	}
@@ -187,19 +195,23 @@ func (b *readWriteBucket) CreateBucketIfNotExists(key []byte) (
 		value *[]byte
 		id    int64
 	)
-	err := b.tx.QueryRow(
+	row, cancel := b.tx.QueryRow(
 		"SELECT id,value FROM "+b.table+" WHERE "+parentSelector(b.id)+
 			" AND key=$1", key,
-	).Scan(&id, &value)
+	)
+	defer cancel()
+	err := row.Scan(&id, &value)
 
 	switch {
 	// Bucket does not yet exist, so create it now. Postgres will generate a
 	// bucket id for the new bucket.
 	case err == sql.ErrNoRows:
-		err = b.tx.QueryRow(
+		row, cancel := b.tx.QueryRow(
 			"INSERT INTO "+b.table+" (parent_id, key) "+
-				"VALUES($1, $2) RETURNING id", b.id, key).
+				"VALUES($1, $2) RETURNING id", b.id, key,
-			Scan(&id)
+		)
+		defer cancel()
+		err := row.Scan(&id)
 		if err != nil {
 			return nil, err
 		}
@@ -315,10 +327,12 @@ func (b *readWriteBucket) Delete(key []byte) error {
 
 	// Check to see if a bucket with this key exists.
 	var dummy int
-	err := b.tx.QueryRow(
+	row, cancel := b.tx.QueryRow(
 		"SELECT 1 FROM "+b.table+" WHERE "+parentSelector(b.id)+
 			" AND key=$1 AND value IS NULL", key,
-	).Scan(&dummy)
+	)
+	defer cancel()
+	err := row.Scan(&dummy)
 	switch {
 	// No bucket exists, proceed to deletion of the key.
 	case err == sql.ErrNoRows:
@@ -395,11 +409,13 @@ func (b *readWriteBucket) Sequence() uint64 {
 	}
 
 	var seq int64
-	err := b.tx.QueryRow(
+	row, cancel := b.tx.QueryRow(
 		"SELECT sequence FROM "+b.table+" WHERE id=$1 "+
 			"AND sequence IS NOT NULL",
 		b.id,
-	).Scan(&seq)
+	)
+	defer cancel()
+	err := row.Scan(&seq)
 
 	switch {
 	case err == sql.ErrNoRows:

kvdb/postgres/readwrite_cursor.go

@@ -31,11 +31,13 @@ func (c *readWriteCursor) First() ([]byte, []byte) {
 		key   []byte
 		value []byte
 	)
-	err := c.bucket.tx.QueryRow(
+	row, cancel := c.bucket.tx.QueryRow(
-		"SELECT key, value FROM "+c.bucket.table+" WHERE "+
+		"SELECT key, value FROM " + c.bucket.table + " WHERE " +
-			parentSelector(c.bucket.id)+
+			parentSelector(c.bucket.id) +
 			" ORDER BY key LIMIT 1",
-	).Scan(&key, &value)
+	)
+	defer cancel()
+	err := row.Scan(&key, &value)
 
 	switch {
 	case err == sql.ErrNoRows:
@@ -59,11 +61,13 @@ func (c *readWriteCursor) Last() ([]byte, []byte) {
 		key   []byte
 		value []byte
 	)
-	err := c.bucket.tx.QueryRow(
+	row, cancel := c.bucket.tx.QueryRow(
-		"SELECT key, value FROM "+c.bucket.table+" WHERE "+
+		"SELECT key, value FROM " + c.bucket.table + " WHERE " +
-			parentSelector(c.bucket.id)+
+			parentSelector(c.bucket.id) +
 			" ORDER BY key DESC LIMIT 1",
-	).Scan(&key, &value)
+	)
+	defer cancel()
+	err := row.Scan(&key, &value)
 
 	switch {
 	case err == sql.ErrNoRows:
@@ -87,12 +91,14 @@ func (c *readWriteCursor) Next() ([]byte, []byte) {
 		key   []byte
 		value []byte
 	)
-	err := c.bucket.tx.QueryRow(
+	row, cancel := c.bucket.tx.QueryRow(
 		"SELECT key, value FROM "+c.bucket.table+" WHERE "+
 			parentSelector(c.bucket.id)+
 			" AND key>$1 ORDER BY key LIMIT 1",
 		c.currKey,
-	).Scan(&key, &value)
+	)
+	defer cancel()
+	err := row.Scan(&key, &value)
 
 	switch {
 	case err == sql.ErrNoRows:
@@ -116,12 +122,14 @@ func (c *readWriteCursor) Prev() ([]byte, []byte) {
 		key   []byte
 		value []byte
 	)
-	err := c.bucket.tx.QueryRow(
+	row, cancel := c.bucket.tx.QueryRow(
 		"SELECT key, value FROM "+c.bucket.table+" WHERE "+
 			parentSelector(c.bucket.id)+
 			" AND key<$1 ORDER BY key DESC LIMIT 1",
 		c.currKey,
-	).Scan(&key, &value)
+	)
+	defer cancel()
+	err := row.Scan(&key, &value)
 
 	switch {
 	case err == sql.ErrNoRows:
@@ -152,12 +160,14 @@ func (c *readWriteCursor) Seek(seek []byte) ([]byte, []byte) {
 		key   []byte
 		value []byte
 	)
-	err := c.bucket.tx.QueryRow(
+	row, cancel := c.bucket.tx.QueryRow(
 		"SELECT key, value FROM "+c.bucket.table+" WHERE "+
 			parentSelector(c.bucket.id)+
 			" AND key>=$1 ORDER BY key LIMIT 1",
 		seek,
-	).Scan(&key, &value)
+	)
+	defer cancel()
+	err := row.Scan(&key, &value)
 
 	switch {
 	case err == sql.ErrNoRows:
@@ -180,12 +190,14 @@ func (c *readWriteCursor) Seek(seek []byte) ([]byte, []byte) {
 func (c *readWriteCursor) Delete() error {
 	// Get first record at or after cursor.
 	var key []byte
-	err := c.bucket.tx.QueryRow(
+	row, cancel := c.bucket.tx.QueryRow(
 		"SELECT key FROM "+c.bucket.table+" WHERE "+
 			parentSelector(c.bucket.id)+
 			" AND key>=$1 ORDER BY key LIMIT 1",
 		c.currKey,
-	).Scan(&key)
+	)
+	defer cancel()
+	err := row.Scan(&key)
 
 	switch {
 	case err == sql.ErrNoRows:

kvdb/postgres/readwrite_tx.go

@@ -158,21 +158,11 @@ func (tx *readWriteTx) OnCommit(cb func()) {
 }
 
 // QueryRow executes a QueryRow call with a timeout context.
-func (tx *readWriteTx) QueryRow(query string, args ...interface{}) *sql.Row {
+func (tx *readWriteTx) QueryRow(query string, args ...interface{}) (*sql.Row,
-	ctx, cancel := tx.db.getTimeoutCtx()
+	func()) {
-	defer cancel()
-
-	return tx.tx.QueryRowContext(ctx, query, args...)
-}
-
-// Query executes a Query call with a timeout context.
-func (tx *readWriteTx) Query(query string, args ...interface{}) (*sql.Rows,
-	error) {
 
 	ctx, cancel := tx.db.getTimeoutCtx()
-	defer cancel()
+	return tx.tx.QueryRowContext(ctx, query, args...), cancel
-
-	return tx.tx.QueryContext(ctx, query, args...)
 }
 
 // Exec executes a Exec call with a timeout context.