Skip to content

Commit c1b8afa

Browse files
daniel-sancheparthea
andauthoredJan 22, 2025··
fix: memory leak in bidi classes (#770)
* clean unneeded fields after close * added assertions to tests --------- Co-authored-by: Anthonios Partheniou <partheniou@google.com>
1 parent fb1c3a9 commit c1b8afa

File tree

2 files changed

+10
-0
lines changed

2 files changed

+10
-0
lines changed
 

Diff for: ‎google/api_core/bidi.py

+3
Original file line numberDiff line numberDiff line change
@@ -306,6 +306,8 @@ def close(self):
306306
self._request_queue.put(None)
307307
self.call.cancel()
308308
self._request_generator = None
309+
self._initial_request = None
310+
self._callbacks = []
309311
# Don't set self.call to None. Keep it around so that send/recv can
310312
# raise the error.
311313

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

719721
self._thread = None
722+
self._on_response = None
720723

721724
@property
722725
def is_active(self):

Diff for: ‎tests/unit/test_bidi.py

+7
Original file line numberDiff line numberDiff line change
@@ -296,6 +296,9 @@ def test_close(self):
296296
# ensure the request queue was signaled to stop.
297297
assert bidi_rpc.pending_requests == 1
298298
assert bidi_rpc._request_queue.get() is None
299+
# ensure request and callbacks are cleaned up
300+
assert bidi_rpc._initial_request is None
301+
assert not bidi_rpc._callbacks
299302

300303
def test_close_no_rpc(self):
301304
bidi_rpc = bidi.BidiRpc(None)
@@ -623,6 +626,8 @@ def cancel_side_effect():
623626
assert bidi_rpc.pending_requests == 1
624627
assert bidi_rpc._request_queue.get() is None
625628
assert bidi_rpc._finalized
629+
assert bidi_rpc._initial_request is None
630+
assert not bidi_rpc._callbacks
626631

627632
def test_reopen_failure_on_rpc_restart(self):
628633
error1 = ValueError("1")
@@ -777,6 +782,7 @@ def on_response(response):
777782
consumer.stop()
778783

779784
assert consumer.is_active is False
785+
assert consumer._on_response is None
780786

781787
def test_wake_on_error(self):
782788
should_continue = threading.Event()
@@ -884,6 +890,7 @@ def close_side_effect():
884890

885891
consumer.stop()
886892
assert consumer.is_active is False
893+
assert consumer._on_response is None
887894

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

0 commit comments

Comments
 (0)
Please sign in to comment.