Repository: streamlit/streamlit

Issue:
Better control of cached object lifecycles
# Problem

Cached objects have the lifetime of the server itself, which is oftem more of a  memory leak than a desired feaure. In particular, I have observed several interrelated problems.

1. In most cases, **cached objects shouldn't persist past the current session** ([example](https://discuss.streamlit.io/t/session-specific-caching/271?u=adrien_treuille)).
2. We've been asked for months for an option so that **cache entries can expire after a certain time** ([example](https://discuss.streamlit.io/t/session-specific-caching/271/2?u=adrien_treuille)).
3. When a cached object represents an OS resource, like an open database connnectionn ([example](https://discuss.streamlit.io/t/on-refresh-database-app-errors-out/344/3?u=adrien_treuille)), **we need to have a way of [finalizing that obect](https://www.geeksforgeeks.org/destructors-in-python/)**, e.g. closing that connection.
4. We've also been asked to **add numerical limits to the cache** (I can't find an example right now.)

# Solution

## MVP

The simplest solution would be to add the following keyword options to 
`st.cache`:

| Option                  | Default | Meaning
|-------------------------|---------|--------
| `global`                | `False` | Whether the object is cached across sessions.
| `ttl`                   | `None`  | The number of seconds to keep the object cached
| `max_entries`           | `None`  | The entries to keep _for that function_


### Reference Implementation

A starting point for implementing these ideas is [this gist](https://gist.github.com/treuille/f988f78c4610c78322d089eb77f74598). Note that the difference between these ideas and this gist is that we should **let cached items expire when they're no longer needed!**

### Why is `max_entries` None?

A basic goal of Streamlit is that the default options should "do the right thing." It's not clear what the "right thing" is in this case but I think there are two possible failure modes:

1. If `max_entries` is a particular number, then the user might experience sudden, discontinuous drop in performance as the cache starts evicting entries 
2. But if `max_entries` is `None`, then the user might see a slow degradation in performance due to a continuous increase in the number of cached items.

I think that (1) is more mysterious than (2) and thus less desirable.

## Next Step: Explicit Finalizing

The next past the MVP would be to allow a function to be called on cache eviction:

| Option                  | Default | Meaning
|-------------------------|---------|--------
| `finalize_func`         | `None`  | Function called on the object on cached eviction

as follows:

```python
from db import Connection # <- making this example up

# close() is a method of Connnection
@st.cache(finalize_func=Connection.close) 
def get_db_connection():
    return Connectin(...)
```

## Next Step Part Deux: Fancy Finalizing

Sometimes the user might want to write their own finalizer after the `st.cache`, which looks prettier, which we could enable as follows:

```python
from db import Connection # <- making this example up

# close() is a method of Connnection
@st.cache(finalize_func=Connection.close) 
def get_db_connection():
    return Connectin(...)

@get_db_connection.finalizer
def finalize_db_connection(conn):
    conn.close()
```

which would be equivalent to the previous code snippet.



Discussion:

Resolving patch:
diff --git a/lib/streamlit/caching.py b/lib/streamlit/caching.py
--- a/lib/streamlit/caching.py
+++ b/lib/streamlit/caching.py
@@ -19,15 +19,18 @@
 import contextlib
 import hashlib
 import inspect
+import math
 import os
 import pickle
 import shutil
 import struct
 import textwrap
 import threading
-import types
+import time
 from collections import namedtuple
-from typing import Any, Dict, List
+from typing import Any, Dict
+
+from cachetools import LRUCache, TTLCache
 
 import streamlit as st
 from streamlit.util import functools_wraps
@@ -53,6 +56,12 @@
 
 LOGGER = get_logger(__name__)
 
+# The timer function we use with TTLCache.
+TTLCACHE_TIMER = time.monotonic
+
+# A list of all mem-caches we've created.
+_all_mem_caches = []
+
 
 class CacheError(Exception):
     pass
@@ -71,10 +80,6 @@ def __init__(self, cached_value):
 DiskCacheEntry = namedtuple("DiskCacheEntry", ["value"])
 
 
-# The in memory cache.
-_mem_cache = {}  # type: Dict[str, CacheEntry]
-
-
 # A thread-local counter that's incremented when we enter @st.cache
 # and decremented when we exit.
 class ThreadLocalCacheInfo(threading.local):
@@ -201,9 +206,37 @@ def _get_mutated_output_error_message():
     return message
 
 
-def _read_from_mem_cache(key, allow_output_mutation, hash_funcs):
-    if key in _mem_cache:
-        entry = _mem_cache[key]
+def _create_mem_cache(max_entries, ttl):
+    """Create an in-memory cache object with the given parameters."""
+    if max_entries is None and ttl is None:
+        # If we have no max_entries or TTL, we can just use a regular dict.
+        mem_cache = {}
+    else:
+        if max_entries is None:
+            max_entries = math.inf
+        elif not isinstance(max_entries, int):
+            raise Exception("`max_entries` must be an int or None")
+
+        if not isinstance(ttl, (float, int)) and ttl is not None:
+            raise Exception("`ttl` must be a float or None")
+
+        # If ttl is none, just create an LRUCache. (TTLCache is simply an
+        # LRUCache that adds a ttl option.)
+        if ttl is None:
+            mem_cache = LRUCache(maxsize=max_entries)
+        else:
+            mem_cache = TTLCache(maxsize=max_entries, ttl=ttl, timer=TTLCACHE_TIMER)
+
+    # Stick the new cache in our global list
+    global _all_mem_caches
+    _all_mem_caches.append(mem_cache)
+
+    return mem_cache
+
+
+def _read_from_mem_cache(mem_cache, key, allow_output_mutation, hash_funcs):
+    if key in mem_cache:
+        entry = mem_cache[key]
 
         if (
             allow_output_mutation
@@ -219,13 +252,13 @@ def _read_from_mem_cache(key, allow_output_mutation, hash_funcs):
         raise CacheKeyNotFoundError("Key not found in mem cache")
 
 
-def _write_to_mem_cache(key, value, allow_output_mutation, hash_funcs):
+def _write_to_mem_cache(mem_cache, key, value, allow_output_mutation, hash_funcs):
     if allow_output_mutation:
         hash = None
     else:
         hash = get_hash(value, hash_funcs=hash_funcs)
 
-    _mem_cache[key] = CacheEntry(value=value, hash=hash)
+    mem_cache[key] = CacheEntry(value=value, hash=hash)
 
 
 def _read_from_disk_cache(key):
@@ -263,9 +296,7 @@ def _write_to_disk_cache(key, value):
         raise CacheError("Unable to write to cache: %s" % e)
 
 
-def _read_from_cache(
-    key, persisted, allow_output_mutation, func_or_code, hash_funcs=None
-):
+def _read_from_cache(mem_cache, key, persisted, allow_output_mutation, hash_funcs=None):
     """
     Read the value from the cache. Our goal is to read from memory
     if possible. If the data was mutated (hash changed), we show a
@@ -273,20 +304,24 @@ def _read_from_cache(
     or rerun the code.
     """
     try:
-        return _read_from_mem_cache(key, allow_output_mutation, hash_funcs)
+        return _read_from_mem_cache(mem_cache, key, allow_output_mutation, hash_funcs)
     except CachedObjectWasMutatedError as e:
         st.warning(_get_mutated_output_error_message())
         return e.cached_value
     except CacheKeyNotFoundError as e:
         if persisted:
             value = _read_from_disk_cache(key)
-            _write_to_mem_cache(key, value, allow_output_mutation, hash_funcs)
+            _write_to_mem_cache(
+                mem_cache, key, value, allow_output_mutation, hash_funcs
+            )
             return value
         raise e
 
 
-def _write_to_cache(key, value, persist, allow_output_mutation, hash_funcs=None):
-    _write_to_mem_cache(key, value, allow_output_mutation, hash_funcs)
+def _write_to_cache(
+    mem_cache, key, value, persist, allow_output_mutation, hash_funcs=None
+):
+    _write_to_mem_cache(mem_cache, key, value, allow_output_mutation, hash_funcs)
     if persist:
         _write_to_disk_cache(key, value)
 
@@ -299,6 +334,8 @@ def cache(
     suppress_st_warning=False,
     hash_funcs=None,
     ignore_hash=False,
+    max_entries=None,
+    ttl=None,
 ):
     """Function decorator to memoize function executions.
 
@@ -332,6 +369,15 @@ def cache(
         check to see if its type matches a key in this dict and, if so, will use the provided
         function to generate a hash for it. See below for an example of how this can be used.
 
+    max_entries : int or None
+        The maximum number of entries to keep in the cache, or None
+        for an unbounded cache. (When a new entry is added to a full cache,
+        the oldest cached entry will be removed.) The default is None.
+
+    ttl : float or None
+        The maximum number of seconds to keep an entry in the cache, or
+        None if cache entries should not expire. The default is None.
+
     ignore_hash : boolean
         DEPRECATED. Please use allow_output_mutation instead.
         This argument will be fully removed after 2020-03-16.
@@ -395,8 +441,13 @@ def cache(
             show_spinner=show_spinner,
             suppress_st_warning=suppress_st_warning,
             hash_funcs=hash_funcs,
+            max_entries=max_entries,
+            ttl=ttl,
         )
 
+    # Create the function's in-memory cache.
+    mem_cache = _create_mem_cache(max_entries, ttl)
+
     @functools_wraps(func)
     def wrapped_func(*args, **kwargs):
         """This function wrapper will only call the underlying function in
@@ -414,7 +465,7 @@ def wrapped_func(*args, **kwargs):
         else:
             message = "Running %s(...)." % name
 
-        def get_or_set_cache():
+        def get_or_create_cached_value():
             hasher = hashlib.new("md5")
 
             args_hasher = CodeHasher("md5", hasher, hash_funcs)
@@ -430,7 +481,11 @@ def get_or_set_cache():
 
             try:
                 return_value = _read_from_cache(
-                    key, persist, allow_output_mutation, func, hash_funcs
+                    mem_cache=mem_cache,
+                    key=key,
+                    persisted=persist,
+                    allow_output_mutation=allow_output_mutation,
+                    hash_funcs=hash_funcs,
                 )
                 LOGGER.debug("Cache hit: %s", func)
             except CacheKeyNotFoundError:
@@ -444,6 +499,7 @@ def get_or_set_cache():
                         return_value = func(*args, **kwargs)
 
                 _write_to_cache(
+                    mem_cache=mem_cache,
                     key=key,
                     value=return_value,
                     persist=persist,
@@ -455,9 +511,9 @@ def get_or_set_cache():
 
         if show_spinner:
             with st.spinner(message):
-                return get_or_set_cache()
+                return get_or_create_cached_value()
         else:
-            return get_or_set_cache()
+            return get_or_create_cached_value()
 
     # Make this a well-behaved decorator by preserving important function
     # attributes.
@@ -500,6 +556,7 @@ class Cache(Dict[Any, Any]):
     def __init__(self, persist=False, allow_output_mutation=False):
         self._persist = persist
         self._allow_output_mutation = allow_output_mutation
+        self._mem_cache = _create_mem_cache(None, None)
 
         dict.__init__(self)
 
@@ -554,7 +611,10 @@ def has_changes(self) -> bool:
 
         try:
             value, _ = _read_from_cache(
-                key, self._persist, self._allow_output_mutation, code, None
+                mem_cache=self._mem_cache,
+                key=key,
+                persisted=self._persist,
+                allow_output_mutation=self._allow_output_mutation,
             )
             self.update(value)
         except CacheKeyNotFoundError:
@@ -562,12 +622,17 @@ def has_changes(self) -> bool:
                 # If we don't hash the results, we don't need to use exec and just return True.
                 # This way line numbers will be correct.
                 _write_to_cache(
-                    key=key, value=self, persist=False, allow_output_mutation=True
+                    mem_cache=self._mem_cache,
+                    key=key,
+                    value=self,
+                    persist=False,
+                    allow_output_mutation=True,
                 )
                 return True
 
             exec(code, caller_frame.f_globals, caller_frame.f_locals)
             _write_to_cache(
+                mem_cache=self._mem_cache,
                 key=key,
                 value=self,
                 persist=self._persist,
@@ -621,5 +686,7 @@ def _clear_disk_cache():
 
 
 def _clear_mem_cache():
-    global _mem_cache
-    _mem_cache = {}
+    global _all_mem_caches
+    # Copy _all_mem_caches to guard against threading errors
+    for mem_cache in list(_all_mem_caches):
+        mem_cache.clear()


Test patch:
diff --git a/lib/tests/streamlit/caching_test.py b/lib/tests/streamlit/caching_test.py
--- a/lib/tests/streamlit/caching_test.py
+++ b/lib/tests/streamlit/caching_test.py
@@ -209,6 +209,146 @@ def thread_test():
         # The other thread should not have modified the main thread
         self.assertEqual(1, get_counter())
 
+    def test_max_size(self):
+        """The oldest object should be evicted when maxsize is reached."""
+        # Create 2 cached functions to test that they don't interfere
+        # with each other.
+
+        foo_vals = []
+
+        @st.cache(max_entries=2)
+        def foo(x):
+            foo_vals.append(x)
+            return x
+
+        bar_vals = []
+
+        @st.cache(max_entries=3)
+        def bar(x):
+            bar_vals.append(x)
+            return x
+
+        self.assertEqual([], foo_vals)
+        self.assertEqual([], bar_vals)
+
+        # Stick two items in both caches. foo will be filled.
+        foo(0), foo(1)
+        bar(0), bar(1)
+        self.assertEqual([0, 1], foo_vals)
+        self.assertEqual([0, 1], bar_vals)
+
+        # 0, 1 are already cached, so called_values shouldn't change.
+        foo(0), foo(1)
+        bar(0), bar(1)
+        self.assertEqual([0, 1], foo_vals)
+        self.assertEqual([0, 1], bar_vals)
+
+        # Add a new item to the cache.
+        # foo: 0 should be evicted; 1 and 2 should be present.
+        # bar: 0, 1, 2 present.
+        foo(2)
+        bar(2)
+
+        # foo(0) again should cause 0 to be added again, since it was
+        # previously evicted. Nothing will have been evicted from bar.
+        foo(1), foo(0)
+        bar(1), bar(0)
+        self.assertEqual([0, 1, 2, 0], foo_vals)
+        self.assertEqual([0, 1, 2], bar_vals)
+
+    # Reduce the huge amount of logspam we get from hashing/caching
+    @patch("streamlit.hashing.LOGGER.debug")
+    @patch("streamlit.caching.LOGGER.debug")
+    def test_no_max_size(self, _1, _2):
+        """If max_size is None, the cache is unbounded."""
+        called_values = []
+
+        @st.cache(max_entries=None)
+        def f(x):
+            called_values.append(x)
+            return x
+
+        # Stick a bunch of items in the cache.
+        for ii in range(256):
+            f(ii)
+
+        # Clear called_values, and test that accessing the same bunch of
+        # items doesn't result in f() being called.
+        called_values = []
+        for ii in range(256):
+            f(ii)
+        self.assertEqual([], called_values)
+
+    @patch("streamlit.caching.TTLCACHE_TIMER")
+    def test_ttl(self, timer_patch):
+        """Entries should expire after the given ttl."""
+        # Create 2 cached functions to test that they don't interfere
+        # with each other.
+        foo_vals = []
+
+        @st.cache(ttl=1)
+        def foo(x):
+            foo_vals.append(x)
+            return x
+
+        bar_vals = []
+
+        @st.cache(ttl=5)
+        def bar(x):
+            bar_vals.append(x)
+            return x
+
+        # Store a value at time 0
+        timer_patch.return_value = 0
+        foo(0)
+        bar(0)
+        self.assertEqual([0], foo_vals)
+        self.assertEqual([0], bar_vals)
+
+        # Advance our timer, but not enough to expire our value.
+        timer_patch.return_value = 0.5
+        foo(0)
+        bar(0)
+        self.assertEqual([0], foo_vals)
+        self.assertEqual([0], bar_vals)
+
+        # Advance our timer enough to expire foo, but not bar.
+        timer_patch.return_value = 1.5
+        foo(0)
+        bar(0)
+        self.assertEqual([0, 0], foo_vals)
+        self.assertEqual([0], bar_vals)
+
+    def test_clear_cache(self):
+        """Clear cache should do its thing."""
+        foo_vals = []
+
+        @st.cache
+        def foo(x):
+            foo_vals.append(x)
+            return x
+
+        bar_vals = []
+
+        @st.cache
+        def bar(x):
+            bar_vals.append(x)
+            return x
+
+        foo(0), foo(1), foo(2)
+        bar(0), bar(1), bar(2)
+        self.assertEqual([0, 1, 2], foo_vals)
+        self.assertEqual([0, 1, 2], bar_vals)
+
+        # Clear the cache and access our original values again. They
+        # should be recomputed.
+        caching.clear_cache()
+
+        foo(0), foo(1), foo(2)
+        bar(0), bar(1), bar(2)
+        self.assertEqual([0, 1, 2, 0, 1, 2], foo_vals)
+        self.assertEqual([0, 1, 2, 0, 1, 2], bar_vals)
+
 
 # Temporarily turn off these tests since there's no Cache object in __init__
 # right now.
diff --git a/lib/tests/streamlit/help_test.py b/lib/tests/streamlit/help_test.py
--- a/lib/tests/streamlit/help_test.py
+++ b/lib/tests/streamlit/help_test.py
@@ -108,8 +108,15 @@ def test_st_cache(self):
             self.assertEqual(
                 ds.signature,
                 (
-                    "(func=None, persist=False, "
-                    "allow_output_mutation=False, show_spinner=True, suppress_st_warning=False, hash_funcs=None, ignore_hash=False)"
+                    "(func=None, "
+                    "persist=False, "
+                    "allow_output_mutation=False, "
+                    "show_spinner=True, "
+                    "suppress_st_warning=False, "
+                    "hash_funcs=None, "
+                    "ignore_hash=False, "
+                    "max_entries=None, "
+                    "ttl=None)"
                 ),
             )
             self.assertTrue(ds.doc_string.startswith("Function decorator to"))