Fix pandas-dev#587 Add a best effort cache option for list_libraries

list_libraries is a frequently called operation that can get
quite slow with a high amount of libraries. I noticed 4-5s+
for 7-8k libraries. Also the number of libraries is fairly
constant, so not much point in refetching the entire list every
time.

This adds a new collection in mongo which acts as a global cache
for all queries. The meta_db can be used for more accounting and
metadata storage in future, and this collection can be used for
caching other stuff as well.

It's best effort as if you don't have permissions or the data
does not exist in the cache it will just use the list_libraries.
The idea is to not require admin for reads and not expect the user
to have write access to a collection on writes.

Author: shashank88

arctic/_cache.py

@@ -0,0 +1,67 @@
+import logging
+from datetime import datetime, timedelta
+
+from pymongo.errors import OperationFailure
+
+from ._config import CACHE_COLL, CACHE_DB
+
+logger = logging.getLogger(__name__)
+
+
+class Cache:
+    def __init__(self, client, cache_expiry=3600, cache_db=CACHE_DB, cache_col=CACHE_COLL):
+        self._client = client
+        self._cachedb = client[cache_db]
+        self._cachecol = None
+        try:
+            if cache_col not in self._cachedb.collection_names():
+                self._cachedb.create_collection(cache_col).create_index("date", expireAfterSeconds=cache_expiry)
+        except OperationFailure as op:
+            logging.debug("This is fine if you are not admin. The collection should already be created for you: %s", op)
+
+        self._cachecol = self._cachedb[cache_col]
+
+    def get(self, key, newer_than_secs=-1):
+        """
+
+        :param key: Key for the dataset. eg. list_libraries.
+        :param newer_than_secs: -1 to indicate use cache if available. Used to indicate what level of staleness
+        in seconds is tolerable.
+        :return: None unless if there is non stale data present in the cache.
+        """
+        try:
+            if not self._cachecol:
+                # Collection not created or no permissions to read from it.
+                return None
+            coll_data = self._cachecol.find_one({"type": key})
+            # Check that there is data in cache and it's not stale.
+            if coll_data and (
+                    newer_than_secs == -1 or
+                    datetime.utcnow() < coll_data['date'] + timedelta(seconds=newer_than_secs)
+            ):
+                return coll_data['data']
+        except OperationFailure as op:
+            logging.warning("Could not read from cache due to: %s. Ask your admin to give read permissions on %s:%s",
+                            op, CACHE_DB, CACHE_COLL)
+
+        return None
+
+    def set(self, key, data):
+        try:
+            self._cachecol.update_one(
+                {"type": key},
+                {"$set": {"type": key, "date": datetime.utcnow(), "data": data}},
+                upsert=True
+            )
+        except OperationFailure as op:
+            logging.debug("This operation is to be run with admin permissions. Should be fine: %s", op)
+
+    def append(self, key, append_data):
+        # Not upserting here as this is meant to be use when there is already data for this key.
+        try:
+            self._cachecol.update_one(
+                {'type': key},
+                {'$push': {'data': append_data}}
+            )
+        except OperationFailure as op:
+            logging.debug("Admin is required to append to the cache: %s", op)

arctic/_config.py

@@ -92,3 +92,11 @@ class FwPointersCfg(Enum):
 # ---------------------------
 # Configures the size of the workers pools used for async arctic requests
 ARCTIC_ASYNC_NWORKERS = os.environ.get('ARCTIC_ASYNC_NWORKERS', 4)
+
+
+# -------------------------------
+# Flag used for indicating caching levels. For now just for list_libraries.
+# -------------------------------
+ENABLE_CACHE = not bool(os.environ.get('DISABLE_CACHE'))
+CACHE_COLL = 'cache'
+CACHE_DB = 'meta_db'

arctic/arctic.py

@@ -7,6 +7,8 @@
 from pymongo.errors import OperationFailure, AutoReconnect
 from six import string_types
 
+from ._cache import Cache
+from ._config import ENABLE_CACHE
 from ._util import indent
 from .auth import authenticate, get_auth
 from .chunkstore import chunkstore
@@ -107,6 +109,7 @@ def __init__(self, mongo_host, app_name=APPLICATION_NAME, allow_secondary=False,
         self._lock = threading.RLock()
         self._pid = os.getpid()
         self._pymongo_kwargs = kwargs
+        self._cache = None
 
         if isinstance(mongo_host, string_types):
             self._given_instance = False
@@ -118,6 +121,7 @@ def __init__(self, mongo_host, app_name=APPLICATION_NAME, allow_secondary=False,
             mongo_host.server_info()
             self.mongo_host = ",".join(["{}:{}".format(x[0], x[1]) for x in mongo_host.nodes])
             self._adminDB = self._conn.admin
+            self._cache = Cache(self._conn)
 
     @property
     @mongo_retry
@@ -143,6 +147,7 @@ def _conn(self):
                                                   serverSelectionTimeoutMS=self._server_selection_timeout,
                                                   **self._pymongo_kwargs)
                 self._adminDB = self.__conn.admin
+                self._cache = Cache(self.__conn)
 
                 # Authenticate against admin for the user
                 auth = get_auth(self.mongo_host, self._application_name, 'admin')
@@ -183,13 +188,16 @@ def __getstate__(self):
     def __setstate__(self, state):
         return Arctic.__init__(self, **state)
 
-    @mongo_retry
-    def list_libraries(self):
+    def list_libraries(self, newer_than_secs=-1):
         """
         Returns
         -------
         list of Arctic library names
         """
+        return self._list_libraries_cached(newer_than_secs) if ENABLE_CACHE else self._list_libraries()
+
+    @mongo_retry
+    def _list_libraries(self):
         libs = []
         for db in self._conn.list_database_names():
             if db.startswith(self.DB_PREFIX + '_'):
@@ -202,6 +210,25 @@ def list_libraries(self):
                         libs.append(coll[:-1 * len(self.METADATA_COLL) - 1])
         return libs
 
+    # Better to be pessimistic here and not retry.
+    def _list_libraries_cached(self, newer_than_secs=-1):
+        """
+        Returns
+        -------
+        List of Arctic library names from a cached collection (global per mongo cluster) in mongo.
+        Long term list_libraries should have a use_cached argument.
+        """
+        _ = self._conn  # Ensures the connection exists and cache is initialized with it.
+        cache_data = self._cache.get('list_libraries', newer_than_secs)
+        if cache_data:
+            logger.debug('Library names are in cache.')
+            return cache_data
+
+        return self._list_libraries()
+
+    def reload_cache(self):
+        self._cache.set('list_libraries', self._list_libraries())
+
     def library_exists(self, library):
         """
         Check whether a given library exists.
@@ -260,6 +287,8 @@ def initialize_library(self, library, lib_type=VERSION_STORE, **kwargs):
         if not lib.get_quota():
             lib.set_quota(10 * 1024 * 1024 * 1024)
 
+        self._cache.append('list_libraries', library)
+
     @mongo_retry
     def delete_library(self, library):
         """
@@ -550,9 +579,9 @@ def to_gigabytes(bytes_):
         count = stats['totals']['count']
         if size >= self.quota:
             raise QuotaExceededException("Mongo Quota Exceeded: %s %.3f / %.0f GB used" % (
-                                         '.'.join([self.database_name, self.library]),
-                                         to_gigabytes(size),
-                                         to_gigabytes(self.quota)))
+                '.'.join([self.database_name, self.library]),
+                to_gigabytes(size),
+                to_gigabytes(self.quota)))
 
         # Quota not exceeded, print an informational message and return
         try:

tests/integration/test_arctic.py

@@ -233,3 +233,62 @@ def test_library_exists_no_auth(arctic):
         assert arctic.library_exists('test')
         assert AB.return_value.get_library_type.called
         assert not arctic.library_exists('nonexistentlib')
+
+
+def test_list_libraries_cached(arctic):
+    libs = ['test1', 'test2']
+    for lib in libs:
+        arctic.initialize_library(lib)
+
+    # Should use uncached data as reload_cache is required for initializing for the first time.
+    assert sorted(libs) == sorted(arctic.list_libraries())
+
+    # Check that the cached collection does not exist right now.
+    cache = arctic._cache
+    assert not cache.get('list_libraries')
+
+    # Should call list_libraries if cache is empty.
+    with patch('arctic.arctic.Arctic._list_libraries', return_value=libs) as uncached_list_libraries:
+        assert arctic._list_libraries_cached() == libs
+        uncached_list_libraries.assert_called()
+
+    # Reload cache and check that it has data
+    arctic.reload_cache()
+    assert sorted(cache.get('list_libraries')) == sorted(libs)
+
+    # Should fetch it from cache the second time.
+    with patch('arctic.arctic.Arctic._list_libraries', return_value=libs) as uncached_list_libraries:
+        assert arctic._list_libraries_cached() == libs
+        uncached_list_libraries.assert_not_called()
+
+
+def test_initialize_library_adds_to_cache(arctic):
+    libs = ['test1', 'test2']
+
+    for lib in libs:
+        arctic.initialize_library(lib)
+
+    arctic.reload_cache()
+    assert arctic._list_libraries_cached() == arctic._list_libraries()
+
+    # Add another lib
+    arctic.initialize_library('test3')
+
+    assert sorted(arctic._cache.get('list_libraries')) == ['test1', 'test2', 'test3']
+
+
+def test_cache_does_not_return_stale_data(arctic):
+    libs = ['test1', 'test2']
+
+    for lib in libs:
+        arctic.initialize_library(lib)
+
+    arctic.reload_cache()
+    assert arctic._list_libraries_cached() == arctic._list_libraries()
+
+    time.sleep(0.2)
+
+    # Should call uncached list_libraries if the data is stale according to caller.
+    with patch('arctic.arctic.Arctic._list_libraries', return_value=libs) as uncached_list_libraries:
+        assert arctic._list_libraries_cached(newer_than_secs=0.1) == libs
+        uncached_list_libraries.assert_called()