Clean up the leftover connection for finished threads in pinned thread mode (#471)

HyukjinKwon · web-flow · commit 799ac4eca035 · 2022-03-14T10:24:51.000+09:00
## What is the problem? Correctly, there is resource leak when using the pinned thread mode (see also apache/spark#24898). For example, if you repeat the codes below multiple times to create Py4J connections in multiple threads, ```python # PySpark application import threading def print_prop(): # Py4J connection is used under the hood. print(spark.sparkContext.getLocalProperty("a")) threading.Thread(target=print_prop).start() ``` the number of leftover connections grows: ```python spark._jvm._gateway_client.deque deque([<py4j.clientserver.ClientServerConnection object at 0x7fdc60170940>, <py4j.clientserver.ClientServerConnection object at 0x7fdca011e760>, <py4j.clientserver.ClientServerConnection object at 0x7fdcb01acdc0>, <py4j.clientserver.ClientServerConnection object at 0x7fdc60170100>, <py4j.clientserver.ClientServerConnection object at 0x7fdcb0232d30>]) ``` In the environment where multiple threads are used without a pool, it easily causes "Too many files open" due to the lack of file descriptors (as they are all occupied by unclosed sockets). ## How do you fix? This PR adds another variable to thread local that cleans up the connection right before the thread is finished. We need it as a separate thread local because `connection` is NOT cleaned because the reference is being held at `JavaClient.deque`. See also 50fe45e for more details.
diff --git a/py4j-python/src/py4j/clientserver.py b/py4j-python/src/py4j/clientserver.py
@@ -9,7 +9,7 @@
 
 from __future__ import unicode_literals, absolute_import
 
-from collections import deque
+from collections import deque, Callable
 import logging
 import socket
 from threading import local, Thread
@@ -244,7 +244,10 @@ def set_thread_connection(self, connection):
         :param connection: The ClientServerConnection to associate with the
             current thread.
         """
-        self.thread_connection.connection = weakref.ref(connection)
+        conn = weakref.ref(connection)
+        self.thread_connection._cleaner = (
+            ThreadLocalConnectionFinalizer(conn, self.deque))
+        self.thread_connection.connection = conn
 
     def shutdown_gateway(self):
         try:
@@ -300,6 +303,39 @@ def _create_connection_guard(self, connection):
         return ClientServerConnectionGuard(self, connection)
 
 
+class ThreadLocalConnectionFinalizer(object):
+    """Cleans :class:`ClientServerConnection` held by a thread local by
+    closing it properly and removing it from the :class:`JavaClient`
+    deque. Right before the Python thread is terminated, this
+    instance will be garbage-collected, which triggers a call
+    to __del__  that contains the cleanup logic.
+    """
+    def __init__(self, connection, dequeue):
+        assert (
+            isinstance(connection, Callable) and
+            connection() is not None and
+            isinstance(connection(), ClientServerConnection))
+        self.connection = connection
+        self.deque = dequeue
+
+    def __del__(self):
+        """Removes the connection associated with the current thread
+        from the deque.
+
+        Expected to be called when the thread that started the
+        connection is garbage-collected.
+        """
+        conn = self.connection()
+        if conn is not None:
+            try:
+                # This dequeue is thread-safe, and shared across other
+                # threads.
+                self.deque.remove(conn)
+            except ValueError:
+                # Should never reach this point
+                pass
+
+
 class ClientServerConnectionGuard(GatewayConnectionGuard):
     """Connection guard that does nothing on exit because there is no need to
     close or give back a connection.
@@ -603,6 +639,12 @@ def _get_params(self, input):
             temp = smart_decode(input.readline())[:-1]
         return params
 
+    def __del__(self):
+        # In case new connection is set via
+        # `JavaClient.set_thread_connection`, this connection will be
+        # garbage-collected with closing the underlying socket properly.
+        self.close()
+
 
 class ClientServer(JavaGateway):
     """Subclass of JavaGateway that implements a different threading model: a
diff --git a/py4j-python/src/py4j/tests/client_server_test.py b/py4j-python/src/py4j/tests/client_server_test.py
@@ -14,6 +14,7 @@
 from py4j.tests.java_callback_test import IHelloImpl, IHelloFailingImpl
 from py4j.tests.java_gateway_test import (
     PY4J_JAVA_PATH, check_connection, sleep, WaitOperator)
+from py4j.tests.memory_leak_test import python_gc
 from py4j.tests.py4j_callback_recursive_example import (
     PythonPing, HelloState)
 
@@ -147,6 +148,44 @@ def testSendObjects(self):
             client_server.shutdown()
             self.assertEquals(1000, hello.calls)
 
+    def testCleanConnections(self):
+        """This test intentionally create multiple connections in multiple
+        threads so each connection is in a thread local of each thread.
+        After that, it verifies that if the connection is cleaned and closed
+        properly without a resource leak.
+        """
+        with clientserver_example_app_process():
+            client_server = ClientServer(
+                JavaParameters(), PythonParameters())
+            connections = client_server._gateway_client.deque
+            conditions = []
+
+            def assert_connection():
+                # Creates a connection.
+                client_server.jvm.System.currentTimeMillis()
+                # Should at least create one connection.
+                conditions.append(0 < len(connections))
+
+            threads = [
+                threading.Thread(target=assert_connection),
+                threading.Thread(target=assert_connection),
+                threading.Thread(target=assert_connection),
+            ]
+            for t in threads:
+                t.start()
+            for t in threads:
+                t.join()
+
+            # Here we explicitly call garbage collection to clean
+            # `ClientServerConnection`s by
+            # `JavaClient.ThreadLocalConnectionCleaner`
+            python_gc()
+
+            # Should have zero connections left.
+            self.assertEqual(len(client_server._gateway_client.deque), 0)
+            self.assertTrue(all(connections))
+            client_server.shutdown()
+
 
 class RetryTest(unittest.TestCase):
 
