Add atomic tag initialization to cache counters

Author: Gonghan-Princess

diskcache/core.py

@@ -1015,12 +1015,15 @@ def add(self, key, value, expire=None, read=False, tag=None, retry=False):
 
             return True
 
-    def incr(self, key, delta=1, default=0, retry=False):
+    def incr(self, key, delta=1, default=0, retry=False, tag=None):
         """Increment value by delta for item with key.
 
         If key is missing and default is None then raise KeyError. Else if key
         is missing and default is not None then use default for value.
 
+        The tag is applied only when creating a missing or expired item.
+        Existing items retain their tag and expiration time.
+
         Operation is atomic. All concurrent increment operations will be
         counted individually.
 
@@ -1035,6 +1038,7 @@ def incr(self, key, delta=1, default=0, retry=False):
         :param int delta: amount to increment (default 1)
         :param int default: value if key is missing (default 0)
         :param bool retry: retry if database timeout occurs (default False)
+        :param str tag: tag for a newly created item (default None)
         :return: new value for item
         :raises KeyError: if key is not found and default is None
         :raises Timeout: if database timeout occurs
@@ -1055,9 +1059,7 @@ def incr(self, key, delta=1, default=0, retry=False):
                     raise KeyError(key)
 
                 value = default + delta
-                columns = (None, None) + self._disk.store(
-                    value, False, key=key
-                )
+                columns = (None, tag) + self._disk.store(value, False, key=key)
                 self._row_insert(db_key, raw, now, columns)
                 self._cull(now, sql, cleanup)
                 return value
@@ -1069,9 +1071,7 @@ def incr(self, key, delta=1, default=0, retry=False):
                     raise KeyError(key)
 
                 value = default + delta
-                columns = (None, None) + self._disk.store(
-                    value, False, key=key
-                )
+                columns = (None, tag) + self._disk.store(value, False, key=key)
                 self._row_update(rowid, now, columns)
                 self._cull(now, sql, cleanup)
                 cleanup(filename)
@@ -1090,12 +1090,15 @@ def incr(self, key, delta=1, default=0, retry=False):
 
             return value
 
-    def decr(self, key, delta=1, default=0, retry=False):
+    def decr(self, key, delta=1, default=0, retry=False, tag=None):
         """Decrement value by delta for item with key.
 
         If key is missing and default is None then raise KeyError. Else if key
         is missing and default is not None then use default for value.
 
+        The tag is applied only when creating a missing or expired item.
+        Existing items retain their tag and expiration time.
+
         Operation is atomic. All concurrent decrement operations will be
         counted individually.
 
@@ -1113,12 +1116,13 @@ def decr(self, key, delta=1, default=0, retry=False):
         :param int delta: amount to decrement (default 1)
         :param int default: value if key is missing (default 0)
         :param bool retry: retry if database timeout occurs (default False)
+        :param str tag: tag for a newly created item (default None)
         :return: new value for item
         :raises KeyError: if key is not found and default is None
         :raises Timeout: if database timeout occurs
 
         """
-        return self.incr(key, -delta, default, retry)
+        return self.incr(key, -delta, default, retry, tag=tag)
 
     def get(
         self,

diskcache/fanout.py

@@ -189,12 +189,15 @@ def add(self, key, value, expire=None, read=False, tag=None, retry=False):
         except Timeout:
             return False
 
-    def incr(self, key, delta=1, default=0, retry=False):
+    def incr(self, key, delta=1, default=0, retry=False, tag=None):
         """Increment value by delta for item with key.
 
         If key is missing and default is None then raise KeyError. Else if key
         is missing and default is not None then use default for value.
 
+        The tag is applied only when creating a missing or expired item.
+        Existing items retain their tag and expiration time.
+
         Operation is atomic. All concurrent increment operations will be
         counted individually.
 
@@ -209,23 +212,27 @@ def incr(self, key, delta=1, default=0, retry=False):
         :param int delta: amount to increment (default 1)
         :param int default: value if key is missing (default 0)
         :param bool retry: retry if database timeout occurs (default False)
+        :param str tag: tag for a newly created item (default None)
         :return: new value for item on success else None
         :raises KeyError: if key is not found and default is None
 
         """
         index = self._hash(key) % self._count
         shard = self._shards[index]
         try:
-            return shard.incr(key, delta, default, retry)
+            return shard.incr(key, delta, default, retry, tag=tag)
         except Timeout:
             return None
 
-    def decr(self, key, delta=1, default=0, retry=False):
+    def decr(self, key, delta=1, default=0, retry=False, tag=None):
         """Decrement value by delta for item with key.
 
         If key is missing and default is None then raise KeyError. Else if key
         is missing and default is not None then use default for value.
 
+        The tag is applied only when creating a missing or expired item.
+        Existing items retain their tag and expiration time.
+
         Operation is atomic. All concurrent decrement operations will be
         counted individually.
 
@@ -243,14 +250,15 @@ def decr(self, key, delta=1, default=0, retry=False):
         :param int delta: amount to decrement (default 1)
         :param int default: value if key is missing (default 0)
         :param bool retry: retry if database timeout occurs (default False)
+        :param str tag: tag for a newly created item (default None)
         :return: new value for item on success else None
         :raises KeyError: if key is not found and default is None
 
         """
         index = self._hash(key) % self._count
         shard = self._shards[index]
         try:
-            return shard.decr(key, delta, default, retry)
+            return shard.decr(key, delta, default, retry, tag=tag)
         except Timeout:
             return None
 

docs/tutorial.rst

@@ -173,6 +173,20 @@ decrementing a missing key will raise a :exc:`KeyError`.
         ...
     KeyError: 'carol'
 
+Use `tag` to associate a newly created counter with a group for eviction. The
+tag is applied atomically when the key is missing or expired. Incrementing or
+decrementing an existing item preserves its tag and expiration time, even if
+a different `tag` is supplied. :class:`FanoutCache` supports the same option.
+
+    >>> cache.incr('requests', tag='counters')
+    1
+    >>> cache.decr('requests', tag='another-group')
+    0
+    >>> cache.get('requests', tag=True)
+    (0, 'counters')
+    >>> cache.evict('counters')
+    1
+
 Increment and decrement operations are atomic and assume the value may be
 stored in a SQLite integer column. SQLite supports 64-bit signed integers.
 

tests/test_counter_tags.py

@@ -0,0 +1,83 @@
+"""Test atomic initialization of tagged counters."""
+
+from concurrent.futures import ThreadPoolExecutor
+
+import pytest
+
+import diskcache as dc
+
+
+@pytest.fixture(params=[dc.Cache, dc.FanoutCache])
+def cache(request, tmp_path):
+    with request.param(str(tmp_path)) as cache:
+        yield cache
+
+
+@pytest.mark.parametrize('method,delta', [('incr', 3), ('decr', -3)])
+@pytest.mark.parametrize('tag', [None, '', 'counter'])
+@pytest.mark.parametrize('expired', [False, True])
+def test_counter_tag_creation(cache, method, delta, tag, expired):
+    if expired:
+        cache.set('count', 100, expire=-1, tag='old')
+
+    assert (
+        getattr(cache, method)('count', 3, default=10, tag=tag) == 10 + delta
+    )
+    assert cache.get('count', expire_time=True, tag=True) == (
+        10 + delta,
+        None,
+        tag,
+    )
+    if tag is not None:
+        assert cache.evict(tag) == 1
+        assert 'count' not in cache
+
+
+@pytest.mark.parametrize('method,delta', [('incr', 3), ('decr', -3)])
+@pytest.mark.parametrize('tag', [None, '', 'replacement'])
+def test_counter_tag_preserves_live_metadata(cache, method, delta, tag):
+    cache.set('count', 10, expire=60, tag='original')
+    _, expiry, _ = cache.get('count', expire_time=True, tag=True)
+
+    assert getattr(cache, method)('count', 3, tag=tag) == 10 + delta
+    assert cache.get('count', expire_time=True, tag=True) == (
+        10 + delta,
+        expiry,
+        'original',
+    )
+
+
+@pytest.mark.parametrize('method', ['incr', 'decr'])
+@pytest.mark.parametrize('expired', [False, True])
+def test_counter_tag_missing_default_none(cache, method, expired):
+    if expired:
+        cache.set('count', 100, expire=-1, tag='old')
+
+    with pytest.raises(KeyError):
+        getattr(cache, method)('count', default=None, tag='counter')
+    assert 'count' not in cache
+
+
+@pytest.mark.parametrize('method,delta', [('incr', 3), ('decr', -3)])
+def test_counter_original_positional_arguments(cache, method, delta):
+    assert getattr(cache, method)('count', 3, 10, True) == 10 + delta
+    assert cache.get('count', tag=True) == (10 + delta, None)
+
+
+def test_counter_tag_concurrent_initialization(cache):
+    def increment():
+        try:
+            return [
+                cache.incr('count', retry=True, tag='counter')
+                for _ in range(20)
+            ]
+        finally:
+            cache.close()
+
+    with ThreadPoolExecutor(max_workers=4) as executor:
+        futures = [executor.submit(increment) for _ in range(4)]
+        results = [value for future in futures for value in future.result(30)]
+
+    assert sorted(results) == list(range(1, 81))
+    assert cache.get('count', tag=True) == (80, 'counter')
+    assert cache.evict('counter') == 1