From 42b8b5fbcbe02b02878f8f6e2b98aafc713204be Mon Sep 17 00:00:00 2001 From: Tom Smeding Date: Wed, 1 Apr 2026 13:59:42 +0200 Subject: Cache, drop chronos --- src/Cache.hs | 122 +++++++++++++++++++---------------------------------------- 1 file changed, 39 insertions(+), 83 deletions(-) (limited to 'src/Cache.hs') 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 ZNC - - --- value, previous (more recently accessed), next (less recently accessed) -data LRUNode a = LRUNode a (TVar (Maybe (LRUNode a))) (TVar (Maybe (LRUNode a))) - --- head (most recent), last (least recent) -data LRUChain a = LRUChain (TVar (Maybe (LRUNode a))) (TVar (Maybe (LRUNode a))) - -lruchainNew :: IO (LRUChain a) -lruchainNew = LRUChain <$> newTVarIO Nothing <*> newTVarIO Nothing - -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 +import System.Clock qualified as Clock +import System.Random + + +data Cache k v = Cache Int (IORef (Map k (v, Clock.TimeSpec))) (IORef StdGen) + +cacheNew :: Int -> IO (Cache k v) +cacheNew maxsize = Cache maxsize <$> newIORef Map.empty <*> (newIORef =<< initStdGen) + +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 + +cacheLookup :: Ord k => Cache k v -> k -> IO (Maybe v) +cacheLookup (Cache _ ref _) key = fmap fst . Map.lookup key <$> readIORef ref + +-- 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' -- cgit v1.3