Revert "[SPARK-45167][CONNECT][PYTHON][FOLLOW-UP] Use lighter threading Rlock, and use the existing eventually util function"

HyukjinKwon · HyukjinKwon · commit 522af6971350 · 2023-10-05T09:39:36.000+09:00
This reverts commit 2a9dd2b.
diff --git a/python/pyspark/sql/connect/client/reattach.py b/python/pyspark/sql/connect/client/reattach.py
@@ -18,11 +18,12 @@
 
 check_dependencies(__name__)
 
-from threading import RLock
 import warnings
 import uuid
 from collections.abc import Generator
 from typing import Optional, Dict, Any, Iterator, Iterable, Tuple, Callable, cast, Type, ClassVar
+from multiprocessing import RLock
+from multiprocessing.synchronize import RLock as RLockBase
 from multiprocessing.pool import ThreadPool
 import os
 
@@ -55,7 +56,7 @@ class ExecutePlanResponseReattachableIterator(Generator):
     """
 
     # Lock to manage the pool
-    _lock: ClassVar[RLock] = RLock()
+    _lock: ClassVar[RLockBase] = RLock()
     _release_thread_pool: Optional[ThreadPool] = ThreadPool(os.cpu_count() if os.cpu_count() else 8)
 
     @classmethod
diff --git a/python/pyspark/sql/tests/connect/client/test_client.py b/python/pyspark/sql/tests/connect/client/test_client.py
@@ -25,7 +25,6 @@
 from pyspark.sql.connect.client import SparkConnectClient, ChannelBuilder
 import pyspark.sql.connect.proto as proto
 from pyspark.testing.connectutils import should_test_connect, connect_requirement_message
-from pyspark.testing.utils import eventually
 
 from pyspark.sql.connect.client.core import Retrying
 from pyspark.sql.connect.client.reattach import (
@@ -153,6 +152,20 @@ def _stub_with(self, execute=None, attach=None):
             attach_ops=ResponseGenerator(attach) if attach is not None else None,
         )
 
+    def assertEventually(self, callable, timeout_ms=1000):
+        """Helper method that will continuously evaluate the callable to not raise an
+        exception."""
+        import time
+
+        limit = time.monotonic_ns() + timeout_ms * 1000 * 1000
+        while time.monotonic_ns() < limit:
+            try:
+                callable()
+                break
+            except Exception:
+                time.sleep(0.1)
+        callable()
+
     def test_basic_flow(self):
         stub = self._stub_with([self.response, self.finished])
         ite = ExecutePlanResponseReattachableIterator(self.request, stub, self.policy, [])
@@ -165,7 +178,7 @@ def check_all():
             self.assertEqual(1, stub.release_calls)
             self.assertEqual(1, stub.execute_calls)
 
-        eventually(timeout=1, catch_assertions=True)(check_all)()
+        self.assertEventually(check_all, timeout_ms=1000)
 
     def test_fail_during_execute(self):
         def fatal():
@@ -183,7 +196,7 @@ def check():
             self.assertEqual(1, stub.release_until_calls)
             self.assertEqual(1, stub.execute_calls)
 
-        eventually(timeout=1, catch_assertions=True)(check)()
+        self.assertEventually(check, timeout_ms=1000)
 
     def test_fail_and_retry_during_execute(self):
         def non_fatal():
@@ -202,7 +215,7 @@ def check():
             self.assertEqual(3, stub.release_until_calls)
             self.assertEqual(1, stub.execute_calls)
 
-        eventually(timeout=1, catch_assertions=True)(check)()
+        self.assertEventually(check, timeout_ms=1000)
 
     def test_fail_and_retry_during_reattach(self):
         count = 0
@@ -228,7 +241,7 @@ def check():
             self.assertEqual(1, stub.release_calls)
             self.assertEqual(1, stub.execute_calls)
 
-        eventually(timeout=1, catch_assertions=True)(check)()
+        self.assertEventually(check, timeout_ms=1000)
 
 
 class TestException(grpc.RpcError, grpc.Call):
