From b90d822dfa10f19fec5813f25f477013248a1db4 Mon Sep 17 00:00:00 2001 From: Ken McGrady Date: Tue, 12 Mar 2024 09:52:44 -0700 Subject: [PATCH] Revert "Expire session storage cache on an async timer (#8083)" (#8281) ## Describe your changes This reverts commit 44227adc2fda53e7ccb00a7099b23902e6fb1604. --- **Contribution License Agreement** By submitting this pull request you agree that all contributions to this project are made under the Apache 2.0 license. --- .../runtime/caching/cache_resource_api.py | 4 +- .../in_memory_cache_storage_wrapper.py | 5 ++- .../runtime/memory_session_storage.py | 5 ++- lib/streamlit/util.py | 39 +------------------ lib/tests/streamlit/util_test.py | 25 ------------ 5 files changed, 9 insertions(+), 69 deletions(-) diff --git a/lib/streamlit/runtime/caching/cache_resource_api.py b/lib/streamlit/runtime/caching/cache_resource_api.py index acb89221beb0..7abacf6c1155 100644 --- a/lib/streamlit/runtime/caching/cache_resource_api.py +++ b/lib/streamlit/runtime/caching/cache_resource_api.py @@ -22,6 +22,7 @@ from datetime import timedelta from typing import Any, Callable, Final, TypeVar, cast, overload +from cachetools import TTLCache from typing_extensions import TypeAlias import streamlit as st @@ -47,7 +48,6 @@ from streamlit.runtime.metrics_util import gather_metrics from streamlit.runtime.scriptrunner.script_run_context import get_script_run_ctx from streamlit.runtime.stats import CacheStat, CacheStatsProvider, group_stats -from streamlit.util import TimedCleanupCache _LOGGER: Final = get_logger(__name__) @@ -472,7 +472,7 @@ def __init__( super().__init__() self.key = key self.display_name = display_name - self._mem_cache: TimedCleanupCache[str, MultiCacheResults] = TimedCleanupCache( + self._mem_cache: TTLCache[str, MultiCacheResults] = TTLCache( maxsize=max_entries, ttl=ttl_seconds, timer=cache_utils.TTLCACHE_TIMER ) self._mem_cache_lock = threading.Lock() diff --git a/lib/streamlit/runtime/caching/storage/in_memory_cache_storage_wrapper.py b/lib/streamlit/runtime/caching/storage/in_memory_cache_storage_wrapper.py index dc750c394676..bb0d19e78294 100644 --- a/lib/streamlit/runtime/caching/storage/in_memory_cache_storage_wrapper.py +++ b/lib/streamlit/runtime/caching/storage/in_memory_cache_storage_wrapper.py @@ -16,6 +16,8 @@ import math import threading +from cachetools import TTLCache + from streamlit.logger import get_logger from streamlit.runtime.caching import cache_utils from streamlit.runtime.caching.storage.cache_storage_protocol import ( @@ -24,7 +26,6 @@ CacheStorageKeyNotFoundError, ) from streamlit.runtime.stats import CacheStat -from streamlit.util import TimedCleanupCache _LOGGER = get_logger(__name__) @@ -61,7 +62,7 @@ def __init__(self, persist_storage: CacheStorage, context: CacheStorageContext): self.function_display_name = context.function_display_name self._ttl_seconds = context.ttl_seconds self._max_entries = context.max_entries - self._mem_cache: TimedCleanupCache[str, bytes] = TimedCleanupCache( + self._mem_cache: TTLCache[str, bytes] = TTLCache( maxsize=self.max_entries, ttl=self.ttl_seconds, timer=cache_utils.TTLCACHE_TIMER, diff --git a/lib/streamlit/runtime/memory_session_storage.py b/lib/streamlit/runtime/memory_session_storage.py index 1dcb91082ace..3f31e85b09fe 100644 --- a/lib/streamlit/runtime/memory_session_storage.py +++ b/lib/streamlit/runtime/memory_session_storage.py @@ -16,8 +16,9 @@ from typing import MutableMapping +from cachetools import TTLCache + from streamlit.runtime.session_manager import SessionInfo, SessionStorage -from streamlit.util import TimedCleanupCache class MemorySessionStorage(SessionStorage): @@ -56,7 +57,7 @@ def __init__( inaccessible and will be removed eventually. """ - self._cache: MutableMapping[str, SessionInfo] = TimedCleanupCache( + self._cache: MutableMapping[str, SessionInfo] = TTLCache( maxsize=maxsize, ttl=ttl_seconds ) diff --git a/lib/streamlit/util.py b/lib/streamlit/util.py index 982f71a4b925..445e1c710206 100644 --- a/lib/streamlit/util.py +++ b/lib/streamlit/util.py @@ -16,16 +16,13 @@ from __future__ import annotations -import asyncio import dataclasses import functools import hashlib import os import subprocess import sys -from typing import Any, Callable, Final, Generic, Iterable, Mapping, TypeVar - -from cachetools import TTLCache +from typing import Any, Callable, Final, Iterable, Mapping, TypeVar from streamlit import env_util @@ -202,37 +199,3 @@ def extract_key_query_params( ] for item in sublist } - - -K = TypeVar("K") -V = TypeVar("V") - - -class TimedCleanupCache(TTLCache, Generic[K, V]): - """A TTLCache that asynchronously expires its entries.""" - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self._task: asyncio.Task[Any] | None = None - - def __setitem__(self, key: K, value: V) -> None: - # Set an expiration task to run periodically - # Can't be created in init because that only runs once and - # the event loop might not exist yet. - if self._task is None: - try: - self._task = asyncio.create_task(expire_cache(self)) - except RuntimeError: - # Just continue if the event loop isn't started yet. - pass - super().__setitem__(key, value) - - def __del__(self): - if self._task is not None: - self._task.cancel() - - -async def expire_cache(cache: TTLCache) -> None: - while True: - await asyncio.sleep(30) - cache.expire() diff --git a/lib/tests/streamlit/util_test.py b/lib/tests/streamlit/util_test.py index 628ea3893149..fd26c6fd0182 100644 --- a/lib/tests/streamlit/util_test.py +++ b/lib/tests/streamlit/util_test.py @@ -12,8 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import asyncio -import gc import random import unittest from typing import Dict, List, Set @@ -187,26 +185,3 @@ def test_calc_md5_can_handle_bytes_and_strings(self): util.calc_md5("eventually bytes"), util.calc_md5("eventually bytes".encode("utf-8")), ) - - def test_timed_cleanup_cache_gc(self): - """Test that the TimedCleanupCache does not leave behind tasks when - the cache is not externally reachable""" - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) - - async def create_cache(): - cache = util.TimedCleanupCache(maxsize=2, ttl=10) - cache["foo"] = "bar" - - # expire_cache and create_cache - assert len(asyncio.all_tasks()) > 1 - - asyncio.run(create_cache()) - - gc.collect() - - async def check(): - # Only has this function running - assert len(asyncio.all_tasks()) == 1 - - asyncio.run(check())