From 0286f554eaae4454b085f092b43777d6c5f5d674 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ra=C3=BAl=20Cumplido?= Date: Tue, 28 Jul 2026 16:19:57 +0200 Subject: [PATCH 1/4] GH-50684: [Python][FlightRPC] Break the reference cycle between the C++ FlightServerBase and the Python object to avoid leaking server --- python/pyarrow/_flight.pyx | 2 + python/pyarrow/includes/libarrow_flight.pxd | 1 + python/pyarrow/src/arrow/python/flight.cc | 6 +++ python/pyarrow/src/arrow/python/flight.h | 3 ++ python/pyarrow/tests/test_flight.py | 43 +++++++++++++++++++++ 5 files changed, 55 insertions(+) diff --git a/python/pyarrow/_flight.pyx b/python/pyarrow/_flight.pyx index 82cbf87f516d..07a001dfdd76 100644 --- a/python/pyarrow/_flight.pyx +++ b/python/pyarrow/_flight.pyx @@ -2887,6 +2887,7 @@ cdef class _FlightServerFinalizer(_Weakrefable): status = server.Wait() check_flight_status(status) finally: + server.ReleasePythonServerRef() self.server.reset() @@ -3234,6 +3235,7 @@ cdef class FlightServerBase(_Weakrefable): raise ValueError("shutdown() on uninitialized FlightServerBase") with nogil: check_flight_status(self.server.get().Shutdown()) + self.server.get().ReleasePythonServerRef() def wait(self): """Block until server is terminated with shutdown.""" diff --git a/python/pyarrow/includes/libarrow_flight.pxd b/python/pyarrow/includes/libarrow_flight.pxd index a89137b8455e..a88cfa66644e 100644 --- a/python/pyarrow/includes/libarrow_flight.pxd +++ b/python/pyarrow/includes/libarrow_flight.pxd @@ -536,6 +536,7 @@ cdef extern from "arrow/python/flight.h" namespace "arrow::py::flight" nogil: CStatus ServeWithSignals() except * CStatus Shutdown() CStatus Wait() + void ReleasePythonServerRef() cdef cppclass PyServerAuthHandler\ " arrow::py::flight::PyServerAuthHandler"(CServerAuthHandler): diff --git a/python/pyarrow/src/arrow/python/flight.cc b/python/pyarrow/src/arrow/python/flight.cc index 5ef8a1dd6b05..db868e88269c 100644 --- a/python/pyarrow/src/arrow/python/flight.cc +++ b/python/pyarrow/src/arrow/python/flight.cc @@ -86,6 +86,12 @@ PyFlightServer::PyFlightServer(PyObject* server, const PyFlightServerVtable& vta server_.reset(server); } +void PyFlightServer::ReleasePythonServerRef() { + // Resets OwnedRefNoGIL to break the reference cycle between the C++ FlightServerBase + // and the Python object. + server_.reset(); +} + Status PyFlightServer::ListFlights( const arrow::flight::ServerCallContext& context, const arrow::flight::Criteria* criteria, diff --git a/python/pyarrow/src/arrow/python/flight.h b/python/pyarrow/src/arrow/python/flight.h index 8a1f4c750aac..283e2c26ccfe 100644 --- a/python/pyarrow/src/arrow/python/flight.h +++ b/python/pyarrow/src/arrow/python/flight.h @@ -171,6 +171,9 @@ class ARROW_PYFLIGHT_EXPORT PyFlightServer : public arrow::flight::FlightServerB Status ListActions(const arrow::flight::ServerCallContext& context, std::vector* actions) override; + // Breaks the reference cycle between the C++ FlightServerBase and the Python object. + void ReleasePythonServerRef(); + private: OwnedRefNoGIL server_; PyFlightServerVtable vtable_; diff --git a/python/pyarrow/tests/test_flight.py b/python/pyarrow/tests/test_flight.py index a66be6186d66..60e9c9a7c52c 100644 --- a/python/pyarrow/tests/test_flight.py +++ b/python/pyarrow/tests/test_flight.py @@ -1128,6 +1128,49 @@ def test_flight_server_location_argument(): assert isinstance(server, FlightServerBase) +# The following tests are for GH-50684, which is a memory leak +# in FlightServerBase.__del__(). +def test_flight_server_is_freed(): + # Calling server.shutdown manually should free the server object. + import weakref + import gc + server = FlightServerBase('grpc://localhost:0') + server.shutdown() + server.wait() + ref = weakref.ref(server) + del server + gc.collect() + assert ref() is None + + +def test_flight_server_is_freed_on_exit(): + # Using FlightServerBase as a context manager should free + # the server object on exit. + import weakref + import gc + with FlightServerBase('grpc://localhost:0') as server: + ref = weakref.ref(server) + del server + gc.collect() + assert ref() is None + + +@pytest.mark.xfail( + reason="GH-50684: FlightServerBase is not freed on delete without shutdown" +) +def test_flight_server_is_freed_without_shutdown(): + # GH-50684: Not calling server.shutdown() currently leaks the server object. + # This test is expected to fail until the issue is fixed but is included + # for completeness and further discussion. + import weakref + import gc + server = FlightServerBase('grpc://localhost:0') + ref = weakref.ref(server) + del server + gc.collect() + assert ref() is None + + def test_server_exit_reraises_exception(): with pytest.raises(ValueError): with FlightServerBase(): From 6c7c46a5f05894690dbeefda11f95d23f9dd2955 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ra=C3=BAl=20Cumplido?= Date: Tue, 28 Jul 2026 16:33:06 +0200 Subject: [PATCH 2/4] Some minor improvements --- python/pyarrow/tests/test_flight.py | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) diff --git a/python/pyarrow/tests/test_flight.py b/python/pyarrow/tests/test_flight.py index 60e9c9a7c52c..596e6ead7623 100644 --- a/python/pyarrow/tests/test_flight.py +++ b/python/pyarrow/tests/test_flight.py @@ -17,7 +17,10 @@ import ast import base64 +from datetime import datetime +import gc import itertools +import json import os import pathlib import signal @@ -28,8 +31,7 @@ import threading import time import traceback -import json -from datetime import datetime +import weakref try: import numpy as np @@ -1128,12 +1130,10 @@ def test_flight_server_location_argument(): assert isinstance(server, FlightServerBase) -# The following tests are for GH-50684, which is a memory leak -# in FlightServerBase.__del__(). +# The following tests are for GH-50684, which was a memory leak +# in FlightServerBase. def test_flight_server_is_freed(): # Calling server.shutdown manually should free the server object. - import weakref - import gc server = FlightServerBase('grpc://localhost:0') server.shutdown() server.wait() @@ -1146,8 +1146,6 @@ def test_flight_server_is_freed(): def test_flight_server_is_freed_on_exit(): # Using FlightServerBase as a context manager should free # the server object on exit. - import weakref - import gc with FlightServerBase('grpc://localhost:0') as server: ref = weakref.ref(server) del server @@ -1162,8 +1160,6 @@ def test_flight_server_is_freed_without_shutdown(): # GH-50684: Not calling server.shutdown() currently leaks the server object. # This test is expected to fail until the issue is fixed but is included # for completeness and further discussion. - import weakref - import gc server = FlightServerBase('grpc://localhost:0') ref = weakref.ref(server) del server From decdc0a5ba6714ab002e406d598e41c862642680 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ra=C3=BAl=20Cumplido?= Date: Wed, 29 Jul 2026 15:27:11 +0200 Subject: [PATCH 3/4] Remove unnecessary server.wait() on test --- python/pyarrow/tests/test_flight.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/pyarrow/tests/test_flight.py b/python/pyarrow/tests/test_flight.py index 596e6ead7623..e1c404d81e72 100644 --- a/python/pyarrow/tests/test_flight.py +++ b/python/pyarrow/tests/test_flight.py @@ -1136,7 +1136,6 @@ def test_flight_server_is_freed(): # Calling server.shutdown manually should free the server object. server = FlightServerBase('grpc://localhost:0') server.shutdown() - server.wait() ref = weakref.ref(server) del server gc.collect() From 5b600c2c3841780a64af436d97760c604c87afbd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ra=C3=BAl=20Cumplido?= Date: Thu, 30 Jul 2026 10:36:29 +0200 Subject: [PATCH 4/4] Remove unnecessary server.Wait after server.Shutdown --- python/pyarrow/_flight.pyx | 2 -- 1 file changed, 2 deletions(-) diff --git a/python/pyarrow/_flight.pyx b/python/pyarrow/_flight.pyx index 07a001dfdd76..f8118fbf61c3 100644 --- a/python/pyarrow/_flight.pyx +++ b/python/pyarrow/_flight.pyx @@ -2883,8 +2883,6 @@ cdef class _FlightServerFinalizer(_Weakrefable): try: with nogil: status = server.Shutdown() - if status.ok(): - status = server.Wait() check_flight_status(status) finally: server.ReleasePythonServerRef()