chore(asr-worker; datashare-python): refactor caching to prevent premature evictions

datashare-python/datashare_python/utils.py

@@ -10,6 +10,7 @@
 import threading
 import time
 from collections.abc import (
+    AsyncGenerator,
     AsyncIterable,
     Awaitable,
     Callable,
@@ -766,59 +767,108 @@ def __init__(
         eviction_callback: Callable[[Any, Any], None] | None = None,
     ) -> None:
         self._eviction_callback = eviction_callback
-        self._cache = LRU(cache_size, self._eviction_callback)
+        self._cache = LRU(cache_size, self._on_evict)
         self._sentinel = object()
+        # Leased resources are tracked by id, evicting a leased resource defers the
+        # eviction callback until its last lease is released
+        self._lock = threading.Lock()
+        self._n_leases: dict[int, int] = dict()
+        self._pending_evictions: dict[int, tuple[Any, Any]] = dict()
 
     def __enter__(self) -> Self:
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb) -> None:  # noqa: ANN001
-        if self._eviction_callback is not None:
-            for k, v in self._cache.items():
-                self._eviction_callback(k, v)
+        for k, v in self._cache.items():
+            self._on_evict(k, v)
 
+    @contextlib.contextmanager
     def get_or_cache_resource(
         self, key: str, default_factory: Callable[[], Any]
-    ) -> Any:
-        value = self._cache.get(key, self._sentinel)
-        if value is not self._sentinel:
-            return value
-        # Get the value first to be sure to return the right one in case of concurrent
-        # access.
-        #
-        # Additionally, the factory can be long to complete (that's one of the reason
-        # we want to cache its results). Because we don't lock the factory call,
-        # together with the cache update, in case of concurrent access,
-        # quicker factories will complete first and will stay longer in the cache.
-        # This is fine, the alternative is to lock factory + cache update to avoid
-        # concurrent access but this will mean waiting for the first factory call to
-        # complete, this can potentially imply longer waits than no getting the value
-        # from the cache
-        value = default_factory()
-        self._cache[key] = value
-        return value
+    ) -> Generator[Any, None, None]:
+        """Get or cache a resource, leasing it for the duration of the context.
+
+        A resource evicted while leased stays usable, its eviction callback only runs
+        once all its leases have been released.
+
+        :param key: resource cache key
+        :param default_factory: resource factory, called on cache miss
+        :return: the cached resource
+        """
+        value = self._get_and_acquire(key)
+        if value is self._sentinel:
+            value = default_factory()
+            self._cache_and_acquire(key, value)
+        try:
+            yield value
+        finally:
+            self._release(value)
 
+    @contextlib.asynccontextmanager
     async def async_get_or_cache_resource(
         self, key: str, default_factory: Callable[[], Awaitable[Any]]
-    ) -> Any:
-        value = self._cache.get(key, self._sentinel)
-        if value is not self._sentinel:
-            return value
-        # Get the value first to be sure to return the right one in case of concurrent
-        # access.
-        #
-        # Additionally, the factory can be long to complete (that's one of the reason
-        # we want to cache its results). Because we don't lock the factory call,
-        # together with the cache update, in case of concurrent access,
-        # quicker factories will complete first and will stay longer in the cache.
-        # This is fine, the alternative is to lock factory + cache update to avoid
-        # concurrent access but this will mean waiting for the first factory call to
-        # complete, this can potentially imply longer waits than no getting the value
-        # from the cache
-        value = await default_factory()
-        self._cache[key] = value
+    ) -> AsyncGenerator[Any, None]:
+        """Async version of get_or_cache_resource.
+
+        :param key: resource cache key
+        :param default_factory: async resource factory, awaited on cache miss
+        :return: the cached resource
+        """
+        value = self._get_and_acquire(key)
+        if value is self._sentinel:
+            value = await default_factory()
+            self._cache_and_acquire(key, value)
+        try:
+            yield value
+        finally:
+            self._release(value)
+
+    def _get_and_acquire(self, key: str) -> Any:
+        with self._lock:
+            value = self._cache.get(key, self._sentinel)
+            if value is not self._sentinel:
+                self._acquire(value)
         return value
 
+    def _cache_and_acquire(self, key: str, value: Any) -> None:
+        # The factory can be long to complete (that's one of the reason we want to
+        # cache its results). Because we don't lock the factory call, together with
+        # the cache update, in case of concurrent access, quicker factories will
+        # complete first and will stay longer in the cache. This is fine, the
+        # alternative is to lock factory + cache update to avoid concurrent access but
+        # this will mean waiting for the first factory call to complete, this can
+        # potentially imply longer waits than no getting the value from the cache
+        with self._lock:
+            self._acquire(value)
+        # Can evict other resources, must be called without holding the lock
+        self._cache[key] = value
+
+    def _acquire(self, value: Any) -> None:
+        value_id = id(value)
+        self._n_leases[value_id] = self._n_leases.get(value_id, 0) + 1
+
+    def _release(self, value: Any) -> None:
+        value_id = id(value)
+        with self._lock:
+            self._n_leases[value_id] -= 1
+            if self._n_leases[value_id]:
+                return
+            del self._n_leases[value_id]
+            pending = self._pending_evictions.pop(value_id, None)
+        if pending is not None:
+            self._evict(*pending)
+
+    def _on_evict(self, key: Any, value: Any) -> None:
+        with self._lock:
+            if id(value) in self._n_leases:
+                self._pending_evictions[id(value)] = (key, value)
+                return
+        self._evict(key, value)
+
+    def _evict(self, key: Any, value: Any) -> None:
+        if self._eviction_callback is not None:
+            self._eviction_callback(key, value)
+
 
 def close_cm_callback(key: str, value: Any) -> None:  # noqa: ARG001
     if hasattr(value, "__exit__"):

datashare-python/tests/test_utils.py

@@ -2,6 +2,7 @@
 import fcntl
 import json
 import os
+import threading
 import uuid
 from concurrent.futures import ProcessPoolExecutor
 from datetime import timedelta
@@ -385,9 +386,9 @@ def factory() -> str:
         return "cached"
 
     # When
-    from_cache = shared.get_or_cache_resource(new_key, factory)
-    # Then
-    assert from_cache == "cached"
+    with shared.get_or_cache_resource(new_key, factory) as from_cache:
+        # Then
+        assert from_cache == "cached"
 
 
 def test_get_or_cache_resource_existing_key() -> None:
@@ -399,12 +400,31 @@ def factory() -> str:
         return uuid.uuid4().hex
 
     # When
-    first = shared.get_or_cache_resource(existing_key, factory)
-    second = shared.get_or_cache_resource(existing_key, factory)
+    with shared.get_or_cache_resource(existing_key, factory) as first:
+        pass
+    with shared.get_or_cache_resource(existing_key, factory) as second:
+        pass
     # Then
     assert second == first
 
 
+def test_get_or_cache_resource_released_resource_stays_cached() -> None:
+    # Given
+    mock = MagicMock()
+    eviction_callback = mock.evict
+    factory = MagicMock(side_effect=object)
+    shared = SharedResources(cache_size=1, eviction_callback=eviction_callback)
+    # When
+    with shared.get_or_cache_resource("k", factory) as first:
+        pass
+    with shared.get_or_cache_resource("k", factory) as second:
+        pass
+    # Then
+    assert second is first
+    factory.assert_called_once()
+    eviction_callback.assert_not_called()
+
+
 def test_get_or_cache_eviction_callback_is_called_on_exit() -> None:
     # Given
     mock = MagicMock()
@@ -413,12 +433,131 @@ def test_get_or_cache_eviction_callback_is_called_on_exit() -> None:
     factory = lambda: "value"  # noqa: E731
     shared = SharedResources(eviction_callback=eviction_callback)
     # When
-    with shared:
-        shared.get_or_cache_resource(key, factory)
+    with shared, shared.get_or_cache_resource(key, factory):
+        pass
     # Then
     eviction_callback.assert_called_once_with(key, "value")
 
 
+def test_get_or_cache_resource_evicted_while_leased_defers_eviction_callback() -> None:
+    # Given
+    mock = MagicMock()
+    eviction_callback = mock.evict
+    shared = SharedResources(cache_size=1, eviction_callback=eviction_callback)
+    # When
+    with shared.get_or_cache_resource("k0", lambda: "leased"):
+        with shared.get_or_cache_resource("k1", lambda: "other"):
+            pass
+        # Then
+        eviction_callback.assert_not_called()
+    eviction_callback.assert_called_once_with("k0", "leased")
+
+
+def test_get_or_cache_resource_eviction_callback_waits_for_last_lease() -> None:
+    # Given
+    mock = MagicMock()
+    eviction_callback = mock.evict
+    shared = SharedResources(cache_size=1, eviction_callback=eviction_callback)
+    # When
+    with shared.get_or_cache_resource("k0", lambda: "leased") as first:
+        with (
+            shared.get_or_cache_resource("k0", lambda: "not_called") as second,
+            shared.get_or_cache_resource("k1", lambda: "other"),
+        ):
+            pass
+        # Then
+        eviction_callback.assert_not_called()
+    assert first == second == "leased"
+    eviction_callback.assert_called_once_with("k0", "leased")
+
+
+def test_get_or_cache_resource_released_resource_is_evicted_immediately() -> None:
+    # Given
+    mock = MagicMock()
+    eviction_callback = mock.evict
+    shared = SharedResources(cache_size=1, eviction_callback=eviction_callback)
+    # When
+    with shared.get_or_cache_resource("k0", lambda: "leased"):
+        pass
+    with shared.get_or_cache_resource("k1", lambda: "other"):
+        pass
+    # Then
+    eviction_callback.assert_called_once_with("k0", "leased")
+
+
+def test_get_or_cache_resource_exit_defers_eviction_callback_of_leased_resource() -> (
+    None
+):
+    # Given
+    mock = MagicMock()
+    eviction_callback = mock.evict
+    shared = SharedResources(cache_size=1, eviction_callback=eviction_callback)
+    lease = shared.get_or_cache_resource("k0", lambda: "leased")
+    # When
+    with shared:
+        lease.__enter__()
+    # Then
+    eviction_callback.assert_not_called()
+    lease.__exit__(None, None, None)
+    eviction_callback.assert_called_once_with("k0", "leased")
+
+
+def test_get_or_cache_resource_is_not_closed_by_concurrent_eviction() -> None:
+    # Given
+    class Resource:
+        def __init__(self) -> None:
+            self.closed = False
+
+    shared = SharedResources(
+        cache_size=1, eviction_callback=lambda _, r: setattr(r, "closed", True)
+    )
+    leased, evicted = threading.Event(), threading.Event()
+    closed_while_leased = []
+
+    def use_resource() -> None:
+        with shared.get_or_cache_resource("k0", Resource) as resource:
+            leased.set()
+            evicted.wait(5)
+            closed_while_leased.append(resource.closed)
+
+    # When
+    thread = threading.Thread(target=use_resource)
+    thread.start()
+    leased.wait(5)
+    with shared.get_or_cache_resource("k1", Resource):
+        pass
+    evicted.set()
+    thread.join()
+    # Then
+    assert closed_while_leased == [False]
+
+
+async def test_async_get_or_cache_resource_evicted_while_leased_defers_eviction() -> (
+    None
+):
+    # Given
+    mock = MagicMock()
+    eviction_callback = mock.evict
+    shared = SharedResources(cache_size=1, eviction_callback=eviction_callback)
+
+    async def factory() -> str:
+        return "leased"
+
+    async def other_factory() -> str:
+        return "other"
+
+    # When
+    async with shared.async_get_or_cache_resource("k0", factory) as first:
+        async with shared.async_get_or_cache_resource("k0", other_factory) as second:
+            pass
+        async with shared.async_get_or_cache_resource("k1", other_factory):
+            pass
+        # Then
+        eviction_callback.assert_not_called()
+    assert first == second == "leased"
+    eviction_callback.assert_called_once_with("k0", "leased")
+
+
 @pytest.mark.parametrize(
     ("wf_context", "act_context", "run_context", "expected_work_dir"),
     [