diff --git a/python/pyarrow/_flight.pyx b/python/pyarrow/_flight.pyx index 82cbf87f516d..f8118fbf61c3 100644 --- a/python/pyarrow/_flight.pyx +++ b/python/pyarrow/_flight.pyx @@ -2883,10 +2883,9 @@ cdef class _FlightServerFinalizer(_Weakrefable): try: with nogil: status = server.Shutdown() - if status.ok(): - status = server.Wait() check_flight_status(status) finally: + server.ReleasePythonServerRef() self.server.reset() @@ -3234,6 +3233,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..e1c404d81e72 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,6 +1130,42 @@ def test_flight_server_location_argument(): assert isinstance(server, FlightServerBase) +# 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. + server = FlightServerBase('grpc://localhost:0') + server.shutdown() + 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. + 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. + 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():