-- SPDX-FileCopyrightText: 2026 Alexandra de Wit
-- SPDX-License-Identifier: MIT

-- | Local TTL retention with shared bounds and per-store recency eviction.
module Ecluse.Core.Server.Cache.Backend.Local (
    newPooledRetention,
    LocalPool,
    newLocalPool,
) where

import Data.Cache (Cache)
import Data.Cache qualified as Cache
import Data.HashSet qualified as HashSet
import Data.Map.Strict qualified as Map
import Data.Time (NominalDiffTime)
import System.Clock (TimeSpec, fromNanoSecs)

import Ecluse.Core.Server.Cache.Backend (CacheOccupancy (..), Recency (..), RetentionOperations (..))
import Ecluse.Core.Server.Cache.Backend.Local.Internal
import Ecluse.Core.Server.Cache.Types (StoreBudget (..))

data Weighted v = Weighted
    { forall v. Weighted v -> v
wValue :: v
    , forall v. Weighted v -> Int
wWeight :: Int
    , forall v. Weighted v -> Integer
wStamp :: Integer
    , forall v. Weighted v -> TimeSpec
wExpires :: TimeSpec
    }

data LocalStore k v = LocalStore
    { forall k v. LocalStore k v -> Cache k (Weighted v)
lsStore :: Cache k (Weighted v)
    , forall k v. LocalStore k v -> LocalPool
lsPool :: LocalPool
    , forall k v. LocalStore k v -> StoreBudget
lsFloor :: StoreBudget
    , forall k v. LocalStore k v -> v -> Int
lsWeigh :: v -> Int
    , forall k v. LocalStore k v -> TVar Integer
lsClock :: TVar Integer
    , forall k v. LocalStore k v -> TVar (Map Integer k)
lsRecency :: TVar (Map Integer k)
    , forall k v. LocalStore k v -> TimeSpec
lsTTL :: TimeSpec
    , forall k v. LocalStore k v -> TVar CacheOccupancy
lsOccupancy :: TVar CacheOccupancy
    , forall k v. LocalStore k v -> TVar (Map TimeSpec (HashSet k))
lsExpiry :: TVar (Map TimeSpec (HashSet k))
    , forall k v. LocalStore k v -> TVar (CacheOccupancy -> IO ())
lsRecord :: TVar (CacheOccupancy -> IO ())
    }

-- | Share aggregate capacity, evicting only this store's entries above its floor.
newPooledRetention :: (Hashable k) => LocalPool -> NominalDiffTime -> StoreBudget -> (v -> Int) -> IO (RetentionOperations k v)
newPooledRetention :: forall k v.
Hashable k =>
LocalPool
-> NominalDiffTime
-> StoreBudget
-> (v -> Int)
-> IO (RetentionOperations k v)
newPooledRetention LocalPool
pool NominalDiffTime
ttl StoreBudget
floorBudget v -> Int
weigh = do
    store <- Maybe TimeSpec -> IO (Cache k (Weighted v))
forall k v. Maybe TimeSpec -> IO (Cache k v)
Cache.newCache Maybe TimeSpec
forall a. Maybe a
Nothing
    clock <- newTVarIO 0
    recency <- newTVarIO Map.empty
    occupancy <- newTVarIO (CacheOccupancy 0 0)
    expiry <- newTVarIO Map.empty
    record <- newTVarIO (const pass)
    let storeState = Cache k (Weighted v)
-> LocalPool
-> StoreBudget
-> (v -> Int)
-> TVar Integer
-> TVar (Map Integer k)
-> TimeSpec
-> TVar CacheOccupancy
-> TVar (Map TimeSpec (HashSet k))
-> TVar (CacheOccupancy -> IO ())
-> LocalStore k v
forall k v.
Cache k (Weighted v)
-> LocalPool
-> StoreBudget
-> (v -> Int)
-> TVar Integer
-> TVar (Map Integer k)
-> TimeSpec
-> TVar CacheOccupancy
-> TVar (Map TimeSpec (HashSet k))
-> TVar (CacheOccupancy -> IO ())
-> LocalStore k v
LocalStore Cache k (Weighted v)
store LocalPool
pool StoreBudget
floorBudget v -> Int
weigh TVar Integer
clock TVar (Map Integer k)
recency (NominalDiffTime -> TimeSpec
toTimeSpec NominalDiffTime
ttl) TVar CacheOccupancy
occupancy TVar (Map TimeSpec (HashSet k))
expiry TVar (CacheOccupancy -> IO ())
record
    registerStore pool (fmap fst . Map.lookupMin <$> readTVar expiry) (purgeExpired storeState) ((,) <$> readTVar occupancy <*> readTVar record)
    pure RetentionOperations{roLookup = lookupStore storeState, roInsert = insertBounded storeState}

insertBounded :: (Hashable k) => LocalStore k v -> (CacheOccupancy -> IO ()) -> IO () -> k -> v -> IO ()
insertBounded :: forall k v.
Hashable k =>
LocalStore k v
-> (CacheOccupancy -> IO ()) -> IO () -> k -> v -> IO ()
insertBounded LocalStore k v
storeState CacheOccupancy -> IO ()
record IO ()
refused k
key v
value
    | Bool -> Bool
not (LocalPool -> Bool
poolEnabled LocalPool
pool) = IO ()
refused
    | Bool -> Bool
not (LocalPool -> Int -> Bool
poolAcceptsWeight LocalPool
pool Int
weight) = IO ()
refused
    | Bool
otherwise = do
        retained <- LocalPool -> (TimeSpec -> STM Bool) -> IO Bool
forall a. LocalPool -> (TimeSpec -> STM a) -> IO a
runPool LocalPool
pool ((TimeSpec -> STM Bool) -> IO Bool)
-> (TimeSpec -> STM Bool) -> IO Bool
forall a b. (a -> b) -> a -> b
$ \TimeSpec
now -> do
            TVar (CacheOccupancy -> IO ())
-> (CacheOccupancy -> IO ()) -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (LocalStore k v -> TVar (CacheOccupancy -> IO ())
forall k v. LocalStore k v -> TVar (CacheOccupancy -> IO ())
lsRecord LocalStore k v
storeState) CacheOccupancy -> IO ()
record
            LocalStore k v -> k -> STM ()
forall k v. Hashable k => LocalStore k v -> k -> STM ()
deleteStored LocalStore k v
storeState k
key
            fits <- LocalStore k v -> Int -> STM Bool
forall k v. Hashable k => LocalStore k v -> Int -> STM Bool
evictToBudget LocalStore k v
storeState Int
weight
            when fits (insertStored storeState now key value weight)
            pure fits
        unless retained refused
  where
    pool :: LocalPool
pool = LocalStore k v -> LocalPool
forall k v. LocalStore k v -> LocalPool
lsPool LocalStore k v
storeState
    weight :: Int
weight = LocalStore k v -> v -> Int
forall k v. LocalStore k v -> v -> Int
lsWeigh LocalStore k v
storeState v
value

insertStored :: (Hashable k) => LocalStore k v -> TimeSpec -> k -> v -> Int -> STM ()
insertStored :: forall k v.
Hashable k =>
LocalStore k v -> TimeSpec -> k -> v -> Int -> STM ()
insertStored LocalStore k v
storeState TimeSpec
now k
key v
value Int
weight = do
    stamp <- LocalStore k v -> STM Integer
forall k v. LocalStore k v -> STM Integer
nextStamp LocalStore k v
storeState
    let expires = TimeSpec
now TimeSpec -> TimeSpec -> TimeSpec
forall a. Num a => a -> a -> a
+ LocalStore k v -> TimeSpec
forall k v. LocalStore k v -> TimeSpec
lsTTL LocalStore k v
storeState
        weighted = v -> Int -> Integer -> TimeSpec -> Weighted v
forall v. v -> Int -> Integer -> TimeSpec -> Weighted v
Weighted v
value Int
weight Integer
stamp TimeSpec
expires
    Cache.insertSTM key weighted (lsStore storeState) Nothing
    modifyTVar' (lsExpiry storeState) (Map.insertWith HashSet.union expires (HashSet.singleton key))
    modifyTVar' (lsRecency storeState) (Map.insert stamp key)
    adjustOccupancy storeState 1 weight

evictToBudget :: (Hashable k) => LocalStore k v -> Int -> STM Bool
evictToBudget :: forall k v. Hashable k => LocalStore k v -> Int -> STM Bool
evictToBudget LocalStore k v
storeState Int
incoming = do
    fits <- LocalPool -> Int -> STM Bool
poolFits (LocalStore k v -> LocalPool
forall k v. LocalStore k v -> LocalPool
lsPool LocalStore k v
storeState) Int
incoming
    if fits
        then pure True
        else do
            recency <- readTVar (lsRecency storeState)
            case Map.lookupMin recency of
                Maybe (Integer, k)
Nothing -> Bool -> STM Bool
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
                Just (Integer
_, k
key) -> do
                    held <- Bool
-> k
-> Cache k (Weighted v)
-> TimeSpec
-> STM (Maybe (Weighted v))
forall k v.
(Eq k, Hashable k) =>
Bool -> k -> Cache k v -> TimeSpec -> STM (Maybe v)
Cache.lookupSTM Bool
False k
key (LocalStore k v -> Cache k (Weighted v)
forall k v. LocalStore k v -> Cache k (Weighted v)
lsStore LocalStore k v
storeState) (Integer -> TimeSpec
fromNanoSecs Integer
0)
                    occupancy <- readTVar (lsOccupancy storeState)
                    case held of
                        Just Weighted v
weighted | CacheOccupancy -> Weighted v -> Bool
forall {v}. CacheOccupancy -> Weighted v -> Bool
aboveFloor CacheOccupancy
occupancy Weighted v
weighted -> do
                            LocalStore k v -> k -> STM ()
forall k v. Hashable k => LocalStore k v -> k -> STM ()
deleteStored LocalStore k v
storeState k
key
                            LocalStore k v -> Int -> STM Bool
forall k v. Hashable k => LocalStore k v -> Int -> STM Bool
evictToBudget LocalStore k v
storeState Int
incoming
                        Maybe (Weighted v)
_ -> Bool -> STM Bool
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
  where
    aboveFloor :: CacheOccupancy -> Weighted v -> Bool
aboveFloor CacheOccupancy
occupancy Weighted v
weighted =
        CacheOccupancy -> Int
occEntries CacheOccupancy
occupancy Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (StoreBudget -> Int
sbMinEntries (LocalStore k v -> StoreBudget
forall k v. LocalStore k v -> StoreBudget
lsFloor LocalStore k v
storeState))
            Bool -> Bool -> Bool
&& CacheOccupancy -> Int
occBytes CacheOccupancy
occupancy Int -> Int -> Int
forall a. Num a => a -> a -> a
- Weighted v -> Int
forall v. Weighted v -> Int
wWeight Weighted v
weighted Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (StoreBudget -> Int
sbMinBytes (LocalStore k v -> StoreBudget
forall k v. LocalStore k v -> StoreBudget
lsFloor LocalStore k v
storeState))

deleteStored :: (Hashable k) => LocalStore k v -> k -> STM ()
deleteStored :: forall k v. Hashable k => LocalStore k v -> k -> STM ()
deleteStored LocalStore k v
storeState k
key = do
    held <- Bool
-> k
-> Cache k (Weighted v)
-> TimeSpec
-> STM (Maybe (Weighted v))
forall k v.
(Eq k, Hashable k) =>
Bool -> k -> Cache k v -> TimeSpec -> STM (Maybe v)
Cache.lookupSTM Bool
False k
key (LocalStore k v -> Cache k (Weighted v)
forall k v. LocalStore k v -> Cache k (Weighted v)
lsStore LocalStore k v
storeState) (Integer -> TimeSpec
fromNanoSecs Integer
0)
    for_ held $ \Weighted v
weighted -> do
        k -> Cache k (Weighted v) -> STM ()
forall k v. (Eq k, Hashable k) => k -> Cache k v -> STM ()
Cache.deleteSTM k
key (LocalStore k v -> Cache k (Weighted v)
forall k v. LocalStore k v -> Cache k (Weighted v)
lsStore LocalStore k v
storeState)
        TVar (Map TimeSpec (HashSet k))
-> (Map TimeSpec (HashSet k) -> Map TimeSpec (HashSet k)) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' (LocalStore k v -> TVar (Map TimeSpec (HashSet k))
forall k v. LocalStore k v -> TVar (Map TimeSpec (HashSet k))
lsExpiry LocalStore k v
storeState) ((HashSet k -> Maybe (HashSet k))
-> TimeSpec -> Map TimeSpec (HashSet k) -> Map TimeSpec (HashSet k)
forall k a. Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a
Map.update HashSet k -> Maybe (HashSet k)
dropKey (Weighted v -> TimeSpec
forall v. Weighted v -> TimeSpec
wExpires Weighted v
weighted))
        TVar (Map Integer k) -> (Map Integer k -> Map Integer k) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' (LocalStore k v -> TVar (Map Integer k)
forall k v. LocalStore k v -> TVar (Map Integer k)
lsRecency LocalStore k v
storeState) (Integer -> Map Integer k -> Map Integer k
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete (Weighted v -> Integer
forall v. Weighted v -> Integer
wStamp Weighted v
weighted))
        LocalStore k v -> Int -> Int -> STM ()
forall k v. LocalStore k v -> Int -> Int -> STM ()
adjustOccupancy LocalStore k v
storeState (-Int
1) (Int -> Int
forall a. Num a => a -> a
negate (Weighted v -> Int
forall v. Weighted v -> Int
wWeight Weighted v
weighted))
  where
    dropKey :: HashSet k -> Maybe (HashSet k)
dropKey HashSet k
bucket =
        let remaining :: HashSet k
remaining = k -> HashSet k -> HashSet k
forall a. (Eq a, Hashable a) => a -> HashSet a -> HashSet a
HashSet.delete k
key HashSet k
bucket
         in if HashSet k -> Bool
forall a. HashSet a -> Bool
HashSet.null HashSet k
remaining then Maybe (HashSet k)
forall a. Maybe a
Nothing else HashSet k -> Maybe (HashSet k)
forall a. a -> Maybe a
Just HashSet k
remaining

adjustOccupancy :: LocalStore k v -> Int -> Int -> STM ()
adjustOccupancy :: forall k v. LocalStore k v -> Int -> Int -> STM ()
adjustOccupancy LocalStore k v
storeState Int
entries Int
bytes = do
    TVar CacheOccupancy -> (CacheOccupancy -> CacheOccupancy) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' (LocalStore k v -> TVar CacheOccupancy
forall k v. LocalStore k v -> TVar CacheOccupancy
lsOccupancy LocalStore k v
storeState) ((CacheOccupancy -> CacheOccupancy) -> STM ())
-> (CacheOccupancy -> CacheOccupancy) -> STM ()
forall a b. (a -> b) -> a -> b
$ \CacheOccupancy
occupancy ->
        Int -> Int -> CacheOccupancy
CacheOccupancy (CacheOccupancy -> Int
occEntries CacheOccupancy
occupancy Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
entries) (CacheOccupancy -> Int
occBytes CacheOccupancy
occupancy Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
bytes)
    LocalPool -> Int -> Int -> STM ()
adjustPool (LocalStore k v -> LocalPool
forall k v. LocalStore k v -> LocalPool
lsPool LocalStore k v
storeState) Int
entries Int
bytes

purgeExpired :: (Hashable k) => LocalStore k v -> TimeSpec -> STM ()
purgeExpired :: forall k v. Hashable k => LocalStore k v -> TimeSpec -> STM ()
purgeExpired LocalStore k v
storeState TimeSpec
now = do
    expiry <- TVar (Map TimeSpec (HashSet k)) -> STM (Map TimeSpec (HashSet k))
forall a. TVar a -> STM a
readTVar (LocalStore k v -> TVar (Map TimeSpec (HashSet k))
forall k v. LocalStore k v -> TVar (Map TimeSpec (HashSet k))
lsExpiry LocalStore k v
storeState)
    case Map.lookupMin expiry of
        Just (TimeSpec
deadline, HashSet k
bucket) | TimeSpec
deadline TimeSpec -> TimeSpec -> Bool
forall a. Ord a => a -> a -> Bool
< TimeSpec
now -> do
            (k -> STM ()) -> HashSet k -> STM ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (LocalStore k v -> k -> STM ()
forall k v. Hashable k => LocalStore k v -> k -> STM ()
deleteStored LocalStore k v
storeState) HashSet k
bucket
            LocalStore k v -> TimeSpec -> STM ()
forall k v. Hashable k => LocalStore k v -> TimeSpec -> STM ()
purgeExpired LocalStore k v
storeState TimeSpec
now
        Maybe (TimeSpec, HashSet k)
_ -> STM ()
forall (f :: * -> *). Applicative f => f ()
pass

nextStamp :: LocalStore k v -> STM Integer
nextStamp :: forall k v. LocalStore k v -> STM Integer
nextStamp LocalStore k v
storeState = do
    stamp <- (Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
1) (Integer -> Integer) -> STM Integer -> STM Integer
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar Integer -> STM Integer
forall a. TVar a -> STM a
readTVar (LocalStore k v -> TVar Integer
forall k v. LocalStore k v -> TVar Integer
lsClock LocalStore k v
storeState)
    writeTVar (lsClock storeState) stamp
    pure stamp

lookupStore :: (Hashable k) => LocalStore k v -> (CacheOccupancy -> IO ()) -> Recency -> k -> IO (Maybe v)
lookupStore :: forall k v.
Hashable k =>
LocalStore k v
-> (CacheOccupancy -> IO ()) -> Recency -> k -> IO (Maybe v)
lookupStore LocalStore k v
storeState CacheOccupancy -> IO ()
record Recency
recency k
key = LocalPool -> (TimeSpec -> STM (Maybe v)) -> IO (Maybe v)
forall a. LocalPool -> (TimeSpec -> STM a) -> IO a
readPool (LocalStore k v -> LocalPool
forall k v. LocalStore k v -> LocalPool
lsPool LocalStore k v
storeState) ((TimeSpec -> STM (Maybe v)) -> IO (Maybe v))
-> (TimeSpec -> STM (Maybe v)) -> IO (Maybe v)
forall a b. (a -> b) -> a -> b
$ \TimeSpec
now -> do
    TVar (CacheOccupancy -> IO ())
-> (CacheOccupancy -> IO ()) -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (LocalStore k v -> TVar (CacheOccupancy -> IO ())
forall k v. LocalStore k v -> TVar (CacheOccupancy -> IO ())
lsRecord LocalStore k v
storeState) CacheOccupancy -> IO ()
record
    held <- Bool
-> k
-> Cache k (Weighted v)
-> TimeSpec
-> STM (Maybe (Weighted v))
forall k v.
(Eq k, Hashable k) =>
Bool -> k -> Cache k v -> TimeSpec -> STM (Maybe v)
Cache.lookupSTM Bool
False k
key (LocalStore k v -> Cache k (Weighted v)
forall k v. LocalStore k v -> Cache k (Weighted v)
lsStore LocalStore k v
storeState) TimeSpec
now
    for_ held $ \Weighted v
weighted -> case Recency
recency of
        Recency
PreserveRecency -> STM ()
forall (f :: * -> *). Applicative f => f ()
pass
        Recency
RefreshRecency -> do
            stamp <- LocalStore k v -> STM Integer
forall k v. LocalStore k v -> STM Integer
nextStamp LocalStore k v
storeState
            modifyTVar' (lsRecency storeState) (Map.insert stamp key . Map.delete (wStamp weighted))
            Cache.insertSTM key weighted{wStamp = stamp} (lsStore storeState) Nothing
    pure (wValue <$> held)

toTimeSpec :: NominalDiffTime -> TimeSpec
toTimeSpec :: NominalDiffTime -> TimeSpec
toTimeSpec NominalDiffTime
ttl = Integer -> TimeSpec
fromNanoSecs (Integer -> Integer -> Integer
forall a. Ord a => a -> a -> a
max Integer
0 (Double -> Integer
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
round (NominalDiffTime -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac NominalDiffTime
ttl Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
1e9 :: Double) :: Integer))