Skip to content

fix(bigquery): Fix bigquery socket leak - #17953

Merged
shuoweil merged 8 commits into
mainfrom
fix/bigquery-socket-leak-whitebox
Jul 31, 2026
Merged

fix(bigquery): Fix bigquery socket leak#17953
shuoweil merged 8 commits into
mainfrom
fix/bigquery-socket-leak-whitebox

Conversation

@chalmerlowe

@chalmerlowe chalmerlowe commented Jul 30, 2026

Copy link
Copy Markdown
Contributor

This Pull Request addresses flakiness in the socket leak system test and improves resource cleanup in the Database API (DB-API) connection.

Problem

  1. The test test_dbapi_connection_does_not_leak_sockets was flaky in Continuous Integration (CI) environments because it relied on the psutil library to count Operating System sockets. Socket pooling, delayed garbage collection, and background noise (other processes in multiprocess systems creating sockets in parallelized systems) make counting Operating System sockets non-deterministic.
  2. Closing a database connection could leave underlying resources hanging if cursors were not closed in the correct order or if clients were not fully initialized.

Solution

  1. Deterministic Testing: Rewrote the test to use Whitebox Object Tracking. Instead of counting Operating System sockets, the test intercepts the creation and closure of requests.Session and grpc.Channel objects using instance-level monkey patching. This ensures that every session or channel opened by the connection is explicitly accounted for and closed, independent of Operating System quirks.
  2. Improved Cleanup: Updated Connection.close() in connection.py to close all created cursors before closing the clients themselves. This ensures that cursors release their references to query data and transports first. Added defensive checks to prevent errors if clients are None.

Note

  • This test no longer uses psutil or artificial sleeps (that were used to try and wait long enough for garbage collection to kick in), making it faster and fully deterministic.
  • Uses Python sets to track objects thus avoiding false positives due to double closures.
  • Instance-level patching is used to isolate tracking to only the objects created within the test block.

Closes #17956

@chalmerlowe chalmerlowe self-assigned this Jul 30, 2026

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request refactors the test_dbapi_connection_does_not_leak_sockets system test to use a more deterministic approach by patching and tracking HTTP sessions and gRPC channels instead of relying on flaky psutil socket counts. The review feedback suggests using sets instead of lists to track these objects, which prevents potential flakiness from duplicate close() calls and simplifies both the assertions and the cleanup logic.

Comment thread packages/google-cloud-bigquery/tests/system/test_client.py Outdated
Comment thread packages/google-cloud-bigquery/tests/system/test_client.py Outdated
@chalmerlowe
chalmerlowe marked this pull request as ready for review July 31, 2026 08:11
@chalmerlowe
chalmerlowe requested review from a team as code owners July 31, 2026 08:11
@chalmerlowe
chalmerlowe requested review from shuoweil and removed request for a team July 31, 2026 08:11
@chalmerlowe

Copy link
Copy Markdown
Contributor Author

Related to work here.

@parthea parthea assigned shuoweil and unassigned chalmerlowe Jul 31, 2026
@parthea

parthea commented Jul 31, 2026

Copy link
Copy Markdown
Contributor

Googlers see b/540939659

@shuoweil
shuoweil merged commit 8c26b38 into main Jul 31, 2026
43 checks passed
@shuoweil
shuoweil deleted the fix/bigquery-socket-leak-whitebox branch July 31, 2026 18:20
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants