Skip to content

Commit

Permalink
fix: memory leak in bidi classes (#770)
Browse files Browse the repository at this point in the history
* clean unneeded fields after close

* added assertions to tests

---------

Co-authored-by: Anthonios Partheniou <[email protected]>
  • Loading branch information
daniel-sanche and parthea authored Jan 22, 2025
1 parent fb1c3a9 commit c1b8afa
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 0 deletions.
3 changes: 3 additions & 0 deletions google/api_core/bidi.py
Original file line number Diff line number Diff line change
Expand Up @@ -306,6 +306,8 @@ def close(self):
self._request_queue.put(None)
self.call.cancel()
self._request_generator = None
self._initial_request = None
self._callbacks = []
# Don't set self.call to None. Keep it around so that send/recv can
# raise the error.

Expand Down Expand Up @@ -717,6 +719,7 @@ def stop(self):
_LOGGER.warning("Background thread did not exit.")

self._thread = None
self._on_response = None

@property
def is_active(self):
Expand Down
7 changes: 7 additions & 0 deletions tests/unit/test_bidi.py
Original file line number Diff line number Diff line change
Expand Up @@ -296,6 +296,9 @@ def test_close(self):
# ensure the request queue was signaled to stop.
assert bidi_rpc.pending_requests == 1
assert bidi_rpc._request_queue.get() is None
# ensure request and callbacks are cleaned up
assert bidi_rpc._initial_request is None
assert not bidi_rpc._callbacks

def test_close_no_rpc(self):
bidi_rpc = bidi.BidiRpc(None)
Expand Down Expand Up @@ -623,6 +626,8 @@ def cancel_side_effect():
assert bidi_rpc.pending_requests == 1
assert bidi_rpc._request_queue.get() is None
assert bidi_rpc._finalized
assert bidi_rpc._initial_request is None
assert not bidi_rpc._callbacks

def test_reopen_failure_on_rpc_restart(self):
error1 = ValueError("1")
Expand Down Expand Up @@ -777,6 +782,7 @@ def on_response(response):
consumer.stop()

assert consumer.is_active is False
assert consumer._on_response is None

def test_wake_on_error(self):
should_continue = threading.Event()
Expand Down Expand Up @@ -884,6 +890,7 @@ def close_side_effect():

consumer.stop()
assert consumer.is_active is False
assert consumer._on_response is None

# calling stop twice should not result in an error.
consumer.stop()

0 comments on commit c1b8afa

Please sign in to comment.