fix: bind py4j callback server to a dynamic port to avoid 25334 collision

PythonCallback started the py4j callback server with no port argument, so
py4j bound the hardcoded default port 25334. Concurrent or repeated runs on
the same host using a lambda-based Check then failed with
"OSError: [Errno 98] Address already in use (127.0.0.1:25334)".

Start the callback server on port 0 and reset the JVM-side callback client to
the actually-bound port (the documented py4j pattern), fixing both the bind
collision and the "Error while obtaining a new communication channel" failure
seen when only port=0 was set.

Closes #86, #19, #7, #72, #156, #173, #198

Author: linguado-app

pydeequ/scala_utils.py

@@ -1,7 +1,37 @@
 # -*- coding: utf-8 -*-
 """A collection of utility functions and classes for manipulating with scala objects anc classes through py4j
 """
-from py4j.java_gateway import JavaObject
+from py4j.java_gateway import CallbackServerParameters, JavaObject
+
+
+def _start_dynamic_callback_server(gateway):
+    """Start the py4j callback server on a dynamically-allocated port and point
+    the JVM-side callback client back at the port that was actually bound.
+
+    Historically PyDeequ called ``gateway.start_callback_server()`` with no
+    arguments, which makes py4j bind the hardcoded default port (25334). When
+    two pyspark applications, or two runs in the same host, used a lambda-based
+    ``Check`` at the same time, the second one failed with
+    ``OSError: [Errno 98] Address already in use (127.0.0.1:25334)``
+    (issues #86, #7, #72, #156, #173, #198).
+
+    Binding ``port=0`` lets the OS pick a free port and removes the collision.
+    However, just passing ``CallbackServerParameters(port=0)`` is not enough:
+    the JVM still believes the callback client lives on the previously
+    configured port, so assertions fail with
+    ``Error while obtaining a new communication channel`` (issue #19). The
+    documented fix is to read back the real listening port and reset the
+    JVM-side callback client to it. This mirrors py4j's own
+    ``ResetCallbackClientTest`` (py4j/tests/java_callback_test.py).
+    """
+    gateway.start_callback_server(CallbackServerParameters(port=0))
+    callback_server = gateway.get_callback_server()
+    listening_port = callback_server.get_listening_port()
+    gateway.java_gateway_server.resetCallbackClient(
+        gateway.java_gateway_server.getCallbackClient().getAddress(),
+        listening_port,
+    )
+    return listening_port
 
 
 class PythonCallback:
@@ -12,17 +42,20 @@ def __init__(self, gateway):
         # P4j will return false if the callback server is already started
         # https://github.com/bartdag/py4j/blob/master/py4j-python/src/py4j/java_gateway.py
         callback_server = self.gateway.get_callback_server()
-        # TODO clean
         if callback_server is None:
-            self.gateway.start_callback_server()
-            print("Python Callback server started!")  # TODO Logging
+            # No callback server yet: start one on a dynamic port to avoid the
+            # default-port (25334) collision and reset the JVM callback client.
+            listening_port = _start_dynamic_callback_server(self.gateway)
+            print(f"Python Callback server started on port {listening_port}!")  # TODO Logging
         elif callback_server.is_shutdown:
+            # The previous callback server was shut down (e.g. a prior Spark
+            # session was stopped). Tear it down fully and re-derive a fresh
+            # dynamic port instead of reverting to the default port.
             callback_server.close()
-            self.gateway.restart_callback_server()
-            # Have you tried turning it off and on again?
-            # TODO why do we need to restart this every time?
-            # TODO Will this break during chained function calls?
-            print("PythonCallback server restarted!")
+            self.gateway.shutdown_callback_server()
+            self.gateway._callback_server = None
+            listening_port = _start_dynamic_callback_server(self.gateway)
+            print(f"PythonCallback server restarted on port {listening_port}!")
 
 
 class ScalaFunction1(PythonCallback):

tests/test_checks.py

@@ -452,6 +452,60 @@ def test_isUnique(self):
         self.assertEqual(self.isUnique("b", "All rows are unique"), [Row(constraint_status="Success")])
         self.assertEqual(self.isUnique("email", "All rows are unique"), [Row(constraint_status="Success")])
 
+    def test_lambda_check_uses_dynamic_callback_port(self):
+        """Regression test for the py4j callback-server default-port (25334) collision.
+
+        Closes #86, #19, #7, #72, #156, #173, #198.
+
+        The old PythonCallback called ``gateway.start_callback_server()`` with no
+        port argument, so py4j bound the hardcoded default port 25334. If that port
+        was already taken (a concurrent/repeated run on the same host), a lambda
+        assertion failed with ``OSError: [Errno 98] Address already in use``.
+
+        This test occupies 25334 *before* the callback server is started, then runs
+        a lambda-assertion Check. With the fix it binds a free dynamic port and
+        succeeds; with the old code it would raise the bind error. It also asserts
+        the callback server is not listening on the legacy default port.
+        """
+        import socket
+
+        gateway = self.spark.sparkContext._gateway
+        # Ensure no callback server is lingering from earlier tests so that this
+        # Check is the one that (re)starts it. shutdown is a no-op if not started.
+        gateway.shutdown_callback_server()
+        gateway._callback_server = None
+
+        # Squat on the legacy default port so the old code path would collide.
+        blocker = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        blocker.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 0)
+        squatting_default_port = False
+        try:
+            blocker.bind(("127.0.0.1", 25334))
+            blocker.listen(1)
+            squatting_default_port = True
+        except OSError:
+            # If 25334 cannot be bound here (rare), the test still validates the
+            # success path; it just cannot prove the collision specifically.
+            blocker.close()
+
+        try:
+            # A lambda assertion forces a py4j callback server start.
+            result = self.hasSize(lambda x: x == 3.0)
+            self.assertEqual(result, [Row(constraint_status="Success")])
+
+            callback_server = gateway.get_callback_server()
+            self.assertIsNotNone(callback_server)
+            listening_port = callback_server.get_listening_port()
+            if squatting_default_port:
+                self.assertNotEqual(
+                    listening_port,
+                    25334,
+                    "Callback server must not use the hardcoded default port 25334",
+                )
+        finally:
+            if squatting_default_port:
+                blocker.close()
+
     def test_fail_isUnique(self):
         self.assertEqual(self.isUnique("d"), [Row(constraint_status="Failure")])
         self.assertEqual(self.isUnique("f", "All rows are unique"), [Row(constraint_status="Failure")])