add rwlock timeout

Author: Dengsheng-wzh

rock/admin/core/ray_service.py

@@ -1,12 +1,12 @@
-import ray
 import time
 
-from rock.config import RayConfig
-from rock.logger import init_logger
-from rock.utils.rwlock import AsyncRWLock
+import ray
 from apscheduler.schedulers.asyncio import AsyncIOScheduler
 from apscheduler.triggers.interval import IntervalTrigger
 
+from rock.config import RayConfig
+from rock.logger import init_logger
+from rock.utils.rwlock import AsyncRWLock, WriteLockTimeout
 
 logger = init_logger(__name__)
 
@@ -34,7 +34,6 @@ def increment_ray_request_count(self):
 
     def get_ray_rwlock(self):
         return self._ray_rwlock
-    
 
     def _setup_ray_reconnect_scheduler(self):
         self._ray_reconnection_scheduler = AsyncIOScheduler(
@@ -57,19 +56,22 @@ async def _ray_reconnect_with_policy(self):
                 await self._reconnect_ray()
 
     async def _reconnect_ray(self):
-        async with self._ray_rwlock.write_lock():
-            start_time = time.time()
-            logger.info(f"current time {start_time}, Reconnect ray cluster")
-            ray.shutdown()
-            ray.init(
-                address=self._config.address,
-                runtime_env=self._config.runtime_env,
-                namespace=self._config.namespace,
-                resources=self._config.resources,
-            )
-            self._ray_request_count = 0
-            end_time = time.time()
-            self._ray_establish_time = end_time
-            logger.info(
-                f"current time {end_time}, Reconnect ray cluster successfully, duration {end_time - start_time}s"
-            )
+        try:
+            async with self._ray_rwlock.write_lock(timeout=self._config.ray_reconnect_wait_timeout_seconds):
+                start_time = time.time()
+                logger.info(f"current time {start_time}, Reconnect ray cluster")
+                ray.shutdown()
+                ray.init(
+                    address=self._config.address,
+                    runtime_env=self._config.runtime_env,
+                    namespace=self._config.namespace,
+                    resources=self._config.resources,
+                )
+                self._ray_request_count = 0
+                end_time = time.time()
+                self._ray_establish_time = end_time
+                logger.info(
+                    f"current time {end_time}, Reconnect ray cluster successfully, duration {end_time - start_time}s"
+                )
+        except WriteLockTimeout as e:
+            logger.warning("Reconnect ray cluster timeout, skip reconnectting", exc_info=e)

rock/config.py

@@ -22,6 +22,7 @@ class RayConfig:
     ray_reconnect_interval_seconds: int = field(default=60 * 60 * 12)
     ray_reconnect_request_threshold: int = field(default=10 * 1024 * 1024)
     ray_reconnect_check_interval_seconds: int = field(default=60 * 10)
+    ray_reconnect_wait_timeout_seconds: int = field(default=30)
 
 
 @dataclass

rock/utils/rwlock.py

@@ -2,32 +2,37 @@
 from contextlib import asynccontextmanager
 
 
+class WriteLockTimeout(Exception):
+    pass
+
+
 class AsyncRWLock:
     """
     An asynchronous reader-writer lock implementation that allows multiple concurrent readers
     or exclusive writers. This lock ensures that:
-    
+
     - Multiple readers can acquire the lock simultaneously when no writer is active
     - Writers have exclusive access when acquiring the lock (no readers or other writers)
     - Writers are prioritized over new readers when a writer is waiting, preventing reader starvation
     - The lock is awaitable and integrates with Python's asyncio framework
-    
+
     The implementation uses asyncio.Condition to manage the waiting queue and coordinate
     between readers and writers safely in an asynchronous context.
-    
+
     Usage:
         lock = AsyncRWLock()
-        
+
         # For read operations (shared access)
         async with lock.read_lock():
             # Multiple coroutines can enter this block simultaneously
             data = await read_data()
-            
+
         # For write operations (exclusive access)
         async with lock.write_lock():
             # Only one coroutine can enter this block at a time
             await write_data(new_data)
     """
+
     def __init__(self):
         self._readers = 0
         self._writer = False
@@ -56,13 +61,25 @@ async def read_lock(self):
         finally:
             await self.release_read()
 
-    async def acquire_write(self):
+    async def acquire_write(self, timeout=None):
         async with self._cond:
             self._writer_waiting += 1
             try:
-                while self._writer or self._readers > 0:
-                    await self._cond.wait()
+
+                async def wait_for_unlock():
+                    while self._writer or self._readers > 0:
+                        await self._cond.wait()
+
+                try:
+                    if timeout is None:
+                        await wait_for_unlock()
+                    else:
+                        await asyncio.wait_for(wait_for_unlock(), timeout=timeout)
+                except asyncio.TimeoutError:
+                    return False
+
                 self._writer = True
+                return True
             finally:
                 self._writer_waiting -= 1
 
@@ -74,8 +91,10 @@ async def release_write(self):
             self._cond.notify_all()
 
     @asynccontextmanager
-    async def write_lock(self):
-        await self.acquire_write()
+    async def write_lock(self, timeout=None):
+        ready = await self.acquire_write(timeout=timeout)
+        if not ready:
+            raise WriteLockTimeout()
         try:
             yield
         finally:

tests/unit/admin/core/test_ray_reconnect.py

@@ -43,3 +43,27 @@ async def test_reconnect_ray_calls_ray_shutdown_and_init_and_reset_counters(ray_
         assert service._ray_establish_time == old_establish_time + 5
 
         assert service._ray_establish_time != old_establish_time
+
+
+@pytest.mark.asyncio
+async def test_reconnect_ray_skip_when_reader_exists_and_write_lock_timeout(ray_service: RayService):
+    service = ray_service
+
+    service._ray_request_count = 123
+    service._config.ray_reconnect_wait_timeout_seconds = 5
+
+    old_count = service._ray_request_count
+    old_est = service._ray_establish_time
+
+    service._ray_rwlock._readers = 1
+
+    with patch("rock.admin.core.ray_service.ray.shutdown") as mock_shutdown, patch(
+        "rock.admin.core.ray_service.ray.init"
+    ) as mock_init, patch("time.time", return_value=old_est + 5):
+        await service._reconnect_ray()
+
+        mock_shutdown.assert_not_called()
+        mock_init.assert_not_called()
+
+        assert service._ray_request_count == old_count
+        assert service._ray_establish_time == old_est

tests/unit/utils/test_rwlock_timeout.py

@@ -0,0 +1,36 @@
+import asyncio
+
+import pytest
+
+from rock.utils.rwlock import AsyncRWLock, WriteLockTimeout
+
+
+@pytest.mark.asyncio
+async def test_rwlock_write_timeout_then_read_lock_ok():
+    lock = AsyncRWLock()
+
+    lock._readers = 1
+
+    async def writer():
+        with pytest.raises(WriteLockTimeout):
+            async with lock.write_lock(timeout=5):
+                pytest.fail("write_lock should have timed out and not enter this block")
+
+    await writer()
+
+    assert lock._writer is False
+    assert lock._writer_waiting == 0
+
+    lock._readers = 0
+
+    acquired = False
+
+    async def reader():
+        nonlocal acquired
+        async with lock.read_lock():
+            acquired = True
+            await asyncio.sleep(1)
+
+    await reader()
+
+    assert acquired is True, "Read lock should be obtainable after write timeout"