Repository: aio-libs/async-lru

Issue:
Error when the cache is cleared before the done callback is called
I'm sometimes getting a KeyError on self.__tasks.remove(task) on line 165 in the __init__.py. I think it's because the cache_clear() method is called before the task finishes, which clears the __tasks set and then the task can no longer be found there.  The code in the _task_done_callback method should be probably wrapped in the try/except block, or something like that. 

Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/aio-libs/async-lru/pull/506?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=aio-libs) Report
> Merging [#506](https://app.codecov.io/gh/aio-libs/async-lru/pull/506?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=aio-libs) (1cb3d48) into [master](https://app.codecov.io/gh/aio-libs/async-lru/commit/b3d180deec606fd0cec8cafab7ba1cf1a41b2470?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=aio-libs) (b3d180d) will **increase** coverage by `0.13%`.
> The diff coverage is `100.00%`.

```diff
@@            Coverage Diff             @@
##           master     #506      +/-   ##
==========================================
+ Coverage   96.43%   96.56%   +0.13%     
==========================================
  Files          12       12              
  Lines         730      758      +28     
  Branches       83       87       +4     
==========================================
+ Hits          704      732      +28     
  Misses         23       23              
  Partials        3        3              
```

| Flag | Coverage Δ | |
|---|---|---|
| unit | `96.43% <100.00%> (+0.13%)` | :arrow_up: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=aio-libs#carryforward-flags-in-the-pull-request-comment) to find out more.

| [Impacted Files](https://app.codecov.io/gh/aio-libs/async-lru/pull/506?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=aio-libs) | Coverage Δ | |
|---|---|---|
| [async\_lru/\_\_init\_\_.py](https://app.codecov.io/gh/aio-libs/async-lru/pull/506?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=aio-libs#diff-YXN5bmNfbHJ1L19faW5pdF9fLnB5) | `88.78% <100.00%> (+0.15%)` | :arrow_up: |
| [tests/test\_cache\_clear.py](https://app.codecov.io/gh/aio-libs/async-lru/pull/506?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=aio-libs#diff-dGVzdHMvdGVzdF9jYWNoZV9jbGVhci5weQ==) | `100.00% <100.00%> (ø)` | |

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=aio-libs)


Patch:
diff --git a/async_lru/__init__.py b/async_lru/__init__.py
index 4e5620f..0bdffc7 100644
--- a/async_lru/__init__.py
+++ b/async_lru/__init__.py
@@ -119,6 +119,10 @@ class _LRUCacheWrapper(Generic[_R]):
     def cache_clear(self) -> None:
         self.__hits = 0
         self.__misses = 0
+
+        for c in self.__cache.values():
+            if c.later_call:
+                c.later_call.cancel()
         self.__cache.clear()
         self.__tasks.clear()
 
@@ -162,7 +166,7 @@ class _LRUCacheWrapper(Generic[_R]):
     def _task_done_callback(
         self, fut: "asyncio.Future[_R]", key: Hashable, task: "asyncio.Task[_R]"
     ) -> None:
-        self.__tasks.remove(task)
+        self.__tasks.discard(task)
 
         cache_item = self.__cache.get(key)
         if self.__ttl is not None and cache_item is not None:
diff --git a/setup.cfg b/setup.cfg
index 6e94731..e7b500e 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -53,7 +53,8 @@ install_requires =
 [flake8]
 exclude = .git,.env,__pycache__,.eggs
 max-line-length = 88
-ignore = N801,N802,N803,E252,W503,E133,E203
+extend-select = B950
+ignore = N801,N802,N803,E252,W503,E133,E203,E501
 
 [coverage:run]
 branch = True


Test patch:
diff --git a/tests/test_basic.py b/tests/test_basic.py
index 1994047..267668e 100644
--- a/tests/test_basic.py
+++ b/tests/test_basic.py
@@ -143,9 +143,9 @@ async def test_alru_cache_dict_not_shared(check_lru: Callable[..., None]) -> Non
         coro1._LRUCacheWrapper__cache[1].fut.result()  # type: ignore[attr-defined]
         == coro2._LRUCacheWrapper__cache[1].fut.result()  # type: ignore[attr-defined]
     )
-    assert coro1._LRUCacheWrapper__cache != coro2._LRUCacheWrapper__cache  # type: ignore[attr-defined]  # noqa: E501
-    assert coro1._LRUCacheWrapper__cache.keys() == coro2._LRUCacheWrapper__cache.keys()  # type: ignore[attr-defined]  # noqa: E501
-    assert coro1._LRUCacheWrapper__cache is not coro2._LRUCacheWrapper__cache  # type: ignore[attr-defined]  # noqa: E501
+    assert coro1._LRUCacheWrapper__cache != coro2._LRUCacheWrapper__cache  # type: ignore[attr-defined]
+    assert coro1._LRUCacheWrapper__cache.keys() == coro2._LRUCacheWrapper__cache.keys()  # type: ignore[attr-defined]
+    assert coro1._LRUCacheWrapper__cache is not coro2._LRUCacheWrapper__cache  # type: ignore[attr-defined]
 
 
 async def test_alru_cache_parameters() -> None:
diff --git a/tests/test_cache_clear.py b/tests/test_cache_clear.py
index f7e955c..5593dc8 100644
--- a/tests/test_cache_clear.py
+++ b/tests/test_cache_clear.py
@@ -18,3 +18,38 @@ async def test_cache_clear(check_lru: Callable[..., None]) -> None:
     coro.cache_clear()
 
     check_lru(coro, hits=0, misses=0, cache=0, tasks=0)
+
+
+async def test_cache_clear_pending_task() -> None:
+    @alru_cache()
+    async def coro() -> str:
+        await asyncio.sleep(0.5)
+        return "foo"
+
+    t = asyncio.create_task(coro())
+    await asyncio.sleep(0)
+    assert len(coro._LRUCacheWrapper__tasks) == 1  # type: ignore[attr-defined]
+    inner_task = next(iter(coro._LRUCacheWrapper__tasks))  # type: ignore[attr-defined]
+    assert not inner_task.done()
+
+    coro.cache_clear()
+    await inner_task
+
+    assert await t == "foo"
+    assert inner_task.done()
+
+
+async def test_cache_clear_ttl_callback(check_lru: Callable[..., None]) -> None:
+    @alru_cache(ttl=0.5)
+    async def coro() -> str:
+        return "foo"
+
+    await coro()
+    assert len(coro._LRUCacheWrapper__cache) == 1  # type: ignore[attr-defined]
+    cache_item = next(iter(coro._LRUCacheWrapper__cache.values()))  # type: ignore[attr-defined]
+    assert not cache_item.later_call.cancelled()
+
+    coro.cache_clear()
+
+    assert cache_item.later_call.cancelled()
+    await asyncio.sleep(0.5)