Cache, drop chronos

1 files changed, 33 insertions, 77 deletions

diff --git a/src/Cache.hs b/src/Cache.hs
index 53d2185..4694aa0 100644
--- a/src/Cache.hs
+++ b/src/Cache.hs
@@ -1,87 +1,43 @@
 module Cache where
 
-import Control.Concurrent.STM
-import Control.Monad (forM_)
-import Data.ByteString (ByteString)
-import Data.ByteString qualified as BS
+import Control.Monad (when)
+import Data.IORef
 import Data.Map.Strict (Map)
 import Data.Map.Strict qualified as Map
+import System.Clock qualified as Clock
+import System.Random
 
-import ZNC
 
+data Cache k v = Cache Int (IORef (Map k (v, Clock.TimeSpec))) (IORef StdGen)
 
--- value, previous (more recently accessed), next (less recently accessed)
-data LRUNode a = LRUNode a (TVar (Maybe (LRUNode a))) (TVar (Maybe (LRUNode a)))
+cacheNew :: Int -> IO (Cache k v)
+cacheNew maxsize = Cache maxsize <$> newIORef Map.empty <*> (newIORef =<< initStdGen)
 
--- head (most recent), last (least recent)
-data LRUChain a = LRUChain (TVar (Maybe (LRUNode a))) (TVar (Maybe (LRUNode a)))
+cacheAdd :: (Ord k, Show k) => Cache k v -> k -> v -> IO ()
+cacheAdd cache@(Cache maxsize ref _) key val = do
+  now <- Clock.getTime Clock.Monotonic
+  sz <- atomicModifyIORef' ref $ \mp ->
+    case fst (Map.insertLookupWithKey (\_ _ _ -> (val, now)) key (val, now) mp) of
+      Nothing -> (Map.insert key (val, now) mp, Map.size mp + 1)
+      Just _ -> (Map.insert key (val, now) mp, 0)
+  when (sz > maxsize) $
+    cachePrune cache
 
-lruchainNew :: IO (LRUChain a)
-lruchainNew = LRUChain <$> newTVarIO Nothing <*> newTVarIO Nothing
+cacheLookup :: Ord k => Cache k v -> k -> IO (Maybe v)
+cacheLookup (Cache _ ref _) key = fmap fst . Map.lookup key <$> readIORef ref
 
-lruchainAdd :: LRUChain a -> a -> STM (LRUNode a)
-lruchainAdd (LRUChain hd tl) x = do
-  mhdnode <- readTVar hd
-  case mhdnode of
-    Nothing -> do  -- chain empty
-      node <- LRUNode x <$> newTVar Nothing <*> newTVar Nothing
-      writeTVar hd (Just node)
-      writeTVar tl (Just node)
-      return node
-    Just hdnode -> do
-      node <- LRUNode x <$> newTVar Nothing <*> newTVar (Just hdnode)
-      writeTVar hd (Just node)
-      return node
-
-lruchainBump :: LRUChain a -> LRUNode a -> STM ()
-lruchainBump (LRUChain hd _) node@(LRUNode _ prevvar nextvar) = do
-  mprev <- readTVar prevvar
-  case mprev of
-    Nothing -> return ()  -- already most recent
-    Just (LRUNode _ _ prevnextvar) -> do
-      -- remove node from chain where it is now
-      mnext <- readTVar nextvar
-      writeTVar prevnextvar mnext
-      case mnext of
-        Nothing -> return ()
-        Just (LRUNode _ nextprevvar _) -> writeTVar nextprevvar mprev
-      -- add node to chain at head
-      writeTVar nextvar =<< readTVar hd
-      writeTVar hd (Just node)
-
-lruchainEvict :: LRUChain a -> STM (Maybe a)
-lruchainEvict (LRUChain _ tl) = do
-  mtlnode <- readTVar tl
-  case mtlnode of
-    Nothing -> return Nothing
-    Just (LRUNode x prevvar _) -> do
-      mprev <- readTVar prevvar
-      writeTVar tl mprev
-      case mprev of
-        Nothing -> return ()
-        Just (LRUNode _ _ prevnextvar) -> writeTVar prevnextvar Nothing
-      return (Just x)
-
--- maxsize, elements in cache, linked list
-data LRU a = LRU Int (TVar (Map a (LRUNode a))) (LRUChain a)
-
-lruNew :: Ord a => Int -> IO (LRU a)
-lruNew maxsize = LRU maxsize <$> newTVarIO mempty <*> lruchainNew
-
--- | If a value got evicted, it is returned
-lruBump :: Ord a => LRU a -> a -> IO (Maybe a)
-lruBump (LRU maxsize mapvar chain) value = atomically $ do
-  mp <- readTVar mapvar
-  case Map.lookup value mp of
-    Nothing -> do
-      node <- lruchainAdd chain value
-      writeTVar mapvar $! Map.insert value node mp
-      if Map.size mp >= maxsize
-        then do
-          mret <- lruchainEvict chain
-          forM_ mret $ \ret -> writeTVar mapvar $! Map.delete ret mp
-          return mret
-        else return Nothing
-    Just node -> do
-      lruchainBump chain node
-      return Nothing
+-- Uses 2-random LRU (https://danluu.com/2choices-eviction/)
+cachePrune :: Ord k => Cache k v -> IO ()
+cachePrune (Cache maxsize ref genref) = do go =<< atomicModifyIORef' genref splitGen
+  where
+    go gen0 = do
+      (gen', sz') <- atomicModifyIORef' ref $ \mp ->
+        let sz = Map.size mp
+            (idx1, gen1) = uniformR (0, sz - 1) gen0 ; (_, (_, tm1)) = Map.elemAt idx1 mp
+            (idx2, gen2) = uniformR (0, sz - 1) gen1 ; (_, (_, tm2)) = Map.elemAt idx2 mp
+            oldestIdx = if tm1 < tm2 then idx1 else idx2
+        in if sz > maxsize
+             then (Map.deleteAt oldestIdx mp, (gen2, sz - 1))
+             else (mp, (gen0, sz))
+      when (sz' > maxsize) $
+        go gen'