fix: memory leak in bidi classes (#770)
* clean unneeded fields after close
* added assertions to tests
---------
Co-authored-by: Anthonios Partheniou <partheniou@google.com>
diff --git a/google/api_core/bidi.py b/google/api_core/bidi.py
index 78d98b9..4e800c8 100644
--- a/google/api_core/bidi.py
+++ b/google/api_core/bidi.py
@@ -306,6 +306,8 @@
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.
@@ -717,6 +719,7 @@
_LOGGER.warning("Background thread did not exit.")
self._thread = None
+ self._on_response = None
@property
def is_active(self):
diff --git a/tests/unit/test_bidi.py b/tests/unit/test_bidi.py
index c196a68..cca9a21 100644
--- a/tests/unit/test_bidi.py
+++ b/tests/unit/test_bidi.py
@@ -296,6 +296,9 @@
# 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)
@@ -623,6 +626,8 @@
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")
@@ -777,6 +782,7 @@
consumer.stop()
assert consumer.is_active is False
+ assert consumer._on_response is None
def test_wake_on_error(self):
should_continue = threading.Event()
@@ -884,6 +890,7 @@
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()
