From fb3783319769f2301536e6a6787570e8d0db140e Mon Sep 17 00:00:00 2001
From: Yiling-J <njjyl723@gmail.com>
Date: Sat, 11 Nov 2023 21:50:15 +0800
Subject: [PATCH] add close method

---
 README.md           |  6 ++++++
 tests/test_cache.py |  8 ++++++++
 theine/theine.py    | 26 +++++++++++++++++++++++---
 3 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index 20afd6d..de37c87 100644
--- a/README.md
+++ b/README.md
@@ -89,6 +89,12 @@ cache.set("key", {"foo": "bar"}, timedelta(seconds=100))
 
 # delete from cache
 cache.delete("key")
+
+# close cache, stop timing wheel thread
+cache.close()
+
+# clear cache
+cache.clear()
 ```
 
 ## Decorator
diff --git a/tests/test_cache.py b/tests/test_cache.py
index 31cea0e..053bee0 100644
--- a/tests/test_cache.py
+++ b/tests/test_cache.py
@@ -129,3 +129,11 @@ def test_ttl_high_workload(policy):
             break
     assert len(cache.key_gen.kh) == 0
     assert len(cache.key_gen.hk) == 0
+
+
+def test_close_cache(policy):
+    for _ in range(10):
+        cache = Cache(policy, 500)
+        cache.set("foo", "bar", timedelta(seconds=60))
+        cache.close()
+        assert cache._maintainer.is_alive() is False
diff --git a/theine/theine.py b/theine/theine.py
index 4aec668..eec2744 100644
--- a/theine/theine.py
+++ b/theine/theine.py
@@ -6,8 +6,19 @@
 from datetime import timedelta
 from functools import _make_key, update_wrapper
 from threading import Event, Thread
-from typing import (Any, Callable, Dict, Hashable, List, Optional, Tuple, Type,
-                    TypeVar, Union, cast)
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Hashable,
+    List,
+    Optional,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+    cast,
+)
 
 from theine_core import ClockProCore, LruCore, TlfuCore
 from typing_extensions import ParamSpec, Protocol
@@ -263,6 +274,7 @@ def __init__(self, policy: str, size: int):
             self._cache = [sentinel] * (2 * size + 500)
             setattr(self, "set", self._set_clockpro)
         self.key_gen = KeyGen()
+        self._closed = False
         self._maintainer = Thread(target=self.maintenance, daemon=True)
         self._maintainer.start()
 
@@ -411,10 +423,18 @@ def maintenance(self):
         """
         Remove expired keys.
         """
-        while True:
+        while not self._closed:
             self.core.advance(self._cache, sentinel, self.key_gen.kh, self.key_gen.hk)
             time.sleep(0.5)
 
     def clear(self):
         self.core.clear()
         self._cache = [sentinel] * len(self._cache)
+
+    def close(self):
+        self._closed = True
+        self._maintainer.join()
+
+    def __del__(self):
+        self.clear()
+        self.close()