fix: memory leak in bidi classes (#770)

daniel-sanche · parthea · web-flow · commit c1b8afa4e2ab · 2025-01-22T16:46:51.000-05:00
* clean unneeded fields after close

* added assertions to tests

---------

Co-authored-by: Anthonios Partheniou &lt;partheniou@google.com&gt;
diff --git a/google/api_core/bidi.py b/google/api_core/bidi.py
@@ -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.
 
@@ -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):
diff --git a/tests/unit/test_bidi.py b/tests/unit/test_bidi.py
@@ -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)
@@ -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")
@@ -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()
@@ -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()
