From 9b93f751b3cd0fb84e27ad1f938da8121ce1e83b Mon Sep 17 00:00:00 2001
From: qyuzh <40571736+qyuzh@users.noreply.github.com>
Date: Sat, 28 Oct 2023 04:52:59 +0800
Subject: [PATCH] fix: expiration bug (#137)

---
 src/db.rs | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/src/db.rs b/src/db.rs
index a02d105..b98d675 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -177,8 +177,6 @@ impl Db {
                 .map(|expiration| expiration > when)
                 .unwrap_or(true);
 
-            // Track the expiration.
-            state.expirations.insert((when, key.clone()));
             when
         });
 
@@ -197,10 +195,17 @@ impl Db {
         if let Some(prev) = prev {
             if let Some(when) = prev.expires_at {
                 // clear expiration
-                state.expirations.remove(&(when, key));
+                state.expirations.remove(&(when, key.clone()));
             }
         }
 
+        // Track the expiration. If we insert before remove that will cause bug
+        // when current `(when, key)` equals prev `(when, key)`. Remove then insert
+        // can avoid this.
+        if let Some(when) = expires_at {
+            state.expirations.insert((when, key));
+        }
+
         // Release the mutex before notifying the background task. This helps
         // reduce contention by avoiding the background task waking up only to
         // be unable to acquire the mutex due to this function still holding it.