Respect context cancellation when backing off (#437)

Author: ryanslade

pkg/db/db.go

@@ -42,10 +42,13 @@ func (db *RDB) ExecContext(ctx context.Context, query string, args ...interface{
 
 		pqErr := &pq.Error{}
 		if errors.As(err, &pqErr) && pqErr.Code == lockNotAvailableErrorCode {
-			<-time.After(b.Duration())
-		} else {
-			return nil, err
+			if err := sleepCtx(ctx, b.Duration()); err != nil {
+				return nil, err
+			}
+			continue
 		}
+
+		return nil, err
 	}
 }
 
@@ -70,13 +73,25 @@ func (db *RDB) WithRetryableTransaction(ctx context.Context, f func(context.Cont
 
 		pqErr := &pq.Error{}
 		if errors.As(err, &pqErr) && pqErr.Code == lockNotAvailableErrorCode {
-			<-time.After(b.Duration())
-		} else {
-			return err
+			if err := sleepCtx(ctx, b.Duration()); err != nil {
+				return err
+			}
+			continue
 		}
+
+		return err
 	}
 }
 
 func (db *RDB) Close() error {
 	return db.DB.Close()
 }
+
+func sleepCtx(ctx context.Context, d time.Duration) error {
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-time.After(d):
+		return nil
+	}
+}

pkg/db/db_test.go

@@ -9,9 +9,9 @@ import (
 	"testing"
 	"time"
 
-	"github.com/xataio/pgroll/internal/testutils"
-
 	"github.com/stretchr/testify/require"
+
+	"github.com/xataio/pgroll/internal/testutils"
 	"github.com/xataio/pgroll/pkg/db"
 )
 
@@ -37,6 +37,30 @@ func TestExecContext(t *testing.T) {
 	})
 }
 
+func TestExecContextWhenContextCancelled(t *testing.T) {
+	t.Parallel()
+
+	testutils.WithConnectionToContainer(t, func(conn *sql.DB, connStr string) {
+		ctx := context.Background()
+		ctx, cancel := context.WithCancel(ctx)
+
+		// create a table on which an exclusive lock is held for 2 seconds
+		setupTableLock(t, connStr, 2*time.Second)
+
+		// set the lock timeout to 100ms
+		ensureLockTimeout(t, conn, 100)
+
+		// execute a query that should retry until the lock is released
+		rdb := &db.RDB{DB: conn}
+
+		// Cancel the context before the lock times out
+		go time.AfterFunc(500*time.Millisecond, cancel)
+
+		_, err := rdb.ExecContext(ctx, "INSERT INTO test(id) VALUES (1)")
+		require.Errorf(t, err, "context canceled")
+	})
+}
+
 func TestWithRetryableTransaction(t *testing.T) {
 	t.Parallel()
 
@@ -58,6 +82,32 @@ func TestWithRetryableTransaction(t *testing.T) {
 	})
 }
 
+func TestWithRetryableTransactionWhenContextCancelled(t *testing.T) {
+	t.Parallel()
+
+	testutils.WithConnectionToContainer(t, func(conn *sql.DB, connStr string) {
+		ctx := context.Background()
+		ctx, cancel := context.WithCancel(ctx)
+
+		// create a table on which an exclusive lock is held for 2 seconds
+		setupTableLock(t, connStr, 2*time.Second)
+
+		// set the lock timeout to 100ms
+		ensureLockTimeout(t, conn, 100)
+
+		// run a transaction that should retry until the lock is released
+		rdb := &db.RDB{DB: conn}
+
+		// Cancel the context before the lock times out
+		go time.AfterFunc(500*time.Millisecond, cancel)
+
+		err := rdb.WithRetryableTransaction(ctx, func(ctx context.Context, tx *sql.Tx) error {
+			return tx.QueryRowContext(ctx, "SELECT 1 FROM test").Err()
+		})
+		require.Errorf(t, err, "context canceled")
+	})
+}
+
 // setupTableLock:
 // * connects to the database
 // * creates a table in the database