fix flakes related to context cancellation while establishing postgres connections

hugodutka · hugodutka · commit cf0daf48b26b · 2025-06-05T12:58:32.000Z
diff --git a/cli/clitest/clitest.go b/cli/clitest/clitest.go
@@ -168,6 +168,11 @@ func StartWithAssert(t *testing.T, inv *serpent.Invocation, assertCallback func(
 		switch {
 		case errors.Is(err, context.Canceled):
 			return
+		case err != nil && strings.Contains(err.Error(), "driver: bad connection"):
+			// When a test exits it cancels the context. If a transaction is open
+			// and a query is being executed, instead of a context.Canceled error
+			// we get a "driver: bad connection" error.
+			return
 		default:
 			assert.NoError(t, err)
 		}
diff --git a/cli/server.go b/cli/server.go
@@ -2359,6 +2359,10 @@ func ConnectToPostgres(ctx context.Context, logger slog.Logger, driver string, d
 	if err != nil {
 		return nil, xerrors.Errorf("get postgres version: %w", err)
 	}
+	defer version.Close()
+	if version.Err() != nil {
+		return nil, xerrors.Errorf("version select: %w", version.Err())
+	}
 	if !version.Next() {
 		return nil, xerrors.Errorf("no rows returned for version select")
 	}
@@ -2367,7 +2371,6 @@ func ConnectToPostgres(ctx context.Context, logger slog.Logger, driver string, d
 	if err != nil {
 		return nil, xerrors.Errorf("scan version: %w", err)
 	}
-	_ = version.Close()
 
 	if versionNum < 130000 {
 		return nil, xerrors.Errorf("PostgreSQL version must be v13.0.0 or higher! Got: %d", versionNum)
diff --git a/coderd/database/pubsub/pubsub.go b/coderd/database/pubsub/pubsub.go
@@ -552,15 +552,14 @@ func (p *PGPubsub) startListener(ctx context.Context, connectURL string) error {
 			sentErrCh = true
 		}),
 	}
-	select {
-	case err = <-errCh:
-		if err != nil {
-			_ = p.pgListener.Close()
-			return xerrors.Errorf("create pq listener: %w", err)
-		}
-	case <-ctx.Done():
+	// We don't respect context cancellation here. There's a bug in the pq library
+	// where if you close the listener before or while the connection is being
+	// established, the connection will be established anyway, and will not be
+	// closed.
+	// https://github.com/lib/pq/issues/1192
+	if err := <-errCh; err != nil {
 		_ = p.pgListener.Close()
-		return ctx.Err()
+		return xerrors.Errorf("create pq listener: %w", err)
 	}
 	return nil
 }

Original file line number	Diff line number	Diff line change
`@@ -2359,6 +2359,10 @@ func ConnectToPostgres(ctx context.Context, logger slog.Logger, driver string, d`
`2359`	`2359`	`if err != nil {`
`2360`	`2360`	`return nil, xerrors.Errorf("get postgres version: %w", err)`
`2361`	`2361`	`}`
	`2362`	`+ defer version.Close()`
	`2363`	`+ if version.Err() != nil {`
	`2364`	`+ return nil, xerrors.Errorf("version select: %w", version.Err())`
	`2365`	`+ }`
`2362`	`2366`	`if !version.Next() {`
`2363`	`2367`	`return nil, xerrors.Errorf("no rows returned for version select")`
`2364`	`2368`	`}`
`@@ -2367,7 +2371,6 @@ func ConnectToPostgres(ctx context.Context, logger slog.Logger, driver string, d`
`2367`	`2371`	`if err != nil {`
`2368`	`2372`	`return nil, xerrors.Errorf("scan version: %w", err)`
`2369`	`2373`	`}`
`2370`		`- _ = version.Close()`
`2371`	`2374`
`2372`	`2375`	`if versionNum < 130000 {`
`2373`	`2376`	`return nil, xerrors.Errorf("PostgreSQL version must be v13.0.0 or higher! Got: %d", versionNum)`