content: {-
    Copyright (c) Meta Platforms, Inc. and affiliates.
    All rights reserved.
  
    This source code is licensed under the BSD-style license found in the
    LICENSE file in the root directory of this source tree.
  -}
  
  module Glean.RTS.Foreign.LookupCache
   ( LookupCache, new, clear, withCache, ReplacementPolicy(..)
   , Stats, StatValues, Stat(..)
   , isCounter, getStat, newStats, readStatsAndResetCounters
   , countFailuresAsMisses
   )
  where
  
  import Control.Exception
  import Data.List
  import qualified Data.Vector.Storable as V
  import qualified Data.Vector.Storable.Mutable as VM
  import Data.Word
  import Foreign.C.String
  import Foreign.C.Types
  import Foreign.ForeignPtr
  import Foreign.Ptr
  
  import Util.FFI
  
  import Glean.FFI
  import Glean.RTS.Foreign.Lookup
  
  -- NOTE: This must be kept in sync with the ReplacementPolicy in rts/cache.h
  data ReplacementPolicy
    = LRU  -- evict least-recently-used entries
    | FIFO -- evict the oldest entries. Faster with many concurrent clients
    deriving(Eq,Ord,Enum,Bounded,Show)
  
  newtype LookupCache = LookupCache (ForeignPtr LookupCache)
  
  instance Object LookupCache where
    wrap = LookupCache
    unwrap (LookupCache p) = p
    destroy = glean_lookupcache_free
  
  -- | Pair the 'LookupCache' with a base 'Lookup' for looking up facts. It is
  -- the user's responsibility to ensure that the 'Lookup's are morally the same
  -- throughout the cache's lifetime (such as different instances of the same
  -- database).
  withCache
    :: CanLookup base
    => base
    -> LookupCache
    -> ReplacementPolicy
    -> (Lookup -> IO a)
    -> IO a
  withCache base cache replacement_policy f =
    withLookup base $ \base_lookup ->
    with cache $ \cache_ptr ->
    bracket
      (invoke $ glean_lookupcache_anchor_new
        base_lookup
        cache_ptr
        (fromEnum replacement_policy))
      glean_lookupcache_anchor_free
      (\p -> f (Lookup p ("anchor:" <> lookupName base)))
  
  -- | The 'Stat' object can be shared between different 'LookupCache's which will
  -- accumulate their statistics into it (cf. 'new').
  newtype Stats = Stats (ForeignPtr Stats)
  
  instance Object Stats where
    wrap = Stats
    unwrap (Stats p) = p
    destroy = glean_lookupcache_stats_free
  
  
  -- | An approximate snapshot of the values in a 'Stats' object
  newtype StatValues = StatValues (V.Vector Word64)
  
  instance Semigroup StatValues where
    StatValues old <> StatValues new = StatValues $
      V.izipWith (\i o n -> if isCounter (toEnum i) then o+n else n) old new
  
  instance Monoid StatValues where
    mempty = StatValues $ V.replicate sTAT_COUNT 0
  
  instance Show StatValues where
    show statValues = intercalate ", " $ map (f statValues) $ enumFrom minBound
      where
        f statValues stat =
          show stat <> ": " <> show (getStat statValues stat)
  
  -- NOTE: This must be kept in sync with the stats in rts/cache.h
  data Stat
    = IdByKey_hits
    | IdByKey_misses
    | IdByKey_failures
    | IdByKey_deletes
    | TypeById_hits
    | TypeById_misses
    | TypeById_failures
    | FactById_hits
    | FactById_misses
    | FactById_failures
    | FactById_deletes
  
      -- slightly ugly names because we want them to be nice on ODS
    | Fact_bytes
    | Fact_count
    deriving(Eq,Ord,Enum,Bounded,Show)
  
  sTAT_COUNT :: Int
  sTAT_COUNT = fromEnum (maxBound :: Stat) + 1
  
  lAST_COUNTER :: Stat
  lAST_COUNTER = FactById_deletes
  
  -- | Counters are values which we want to bump in ODS; the others we want to
  -- set.
  isCounter :: Stat -> Bool
  isCounter x = x <= lAST_COUNTER
  
  getStat :: StatValues -> Stat -> Word64
  getStat (StatValues as) c = as V.! fromEnum c
  
  newStats :: IO Stats
  newStats = construct $ invoke glean_lookupcache_stats_new
  
  -- | Obtain an approximate snapshot of the values in 'Stats' and reset all
  -- counters to 0 (but not sums like total cache size).
  readStatsAndResetCounters :: Stats -> IO StatValues
  readStatsAndResetCounters stats = do
    v <- VM.new sTAT_COUNT
    with stats $ \pstats ->
      VM.unsafeWith v $ \pdata ->
      glean_lookupcache_stats_read_and_reset_counters
        pstats
        pdata
        (fromIntegral $ VM.length v)
    StatValues <$> V.unsafeFreeze v
  
  -- | Replace miss counts with failure counts. This is useful when the cache is
  -- being used without a base DB, e.g. in SendAndRebaseQueue. In that case there
  -- will always be zero cache misses, and we're interested in failures.
  countFailuresAsMisses :: StatValues -> StatValues
  countFailuresAsMisses stats@(StatValues vec) = StatValues $
    vec V.//
      [ (fromEnum IdByKey_misses, getStat stats IdByKey_failures)
      , (fromEnum TypeById_misses, getStat stats TypeById_failures)
      , (fromEnum FactById_misses, getStat stats FactById_failures)
      ]
  
  new :: Word64 -> Int -> Stats -> IO LookupCache
  new capacity shards stats =
    with stats $ construct . invoke . glean_lookupcache_new
      (fromIntegral capacity)
      (fromIntegral shards)
  
  clear :: LookupCache -> IO ()
  clear cache = with cache $ invoke . glean_lookupcache_clear
  
  foreign import ccall unsafe glean_lookupcache_stats_new
    :: Ptr (Ptr Stats) -> IO CString
  foreign import ccall unsafe "&glean_lookupcache_stats_free"
    glean_lookupcache_stats_free :: Destroy Stats
  foreign import ccall safe glean_lookupcache_stats_read_and_reset_counters
    :: Ptr Stats -> Ptr Word64 -> CSize -> IO ()
  
  
  foreign import ccall unsafe glean_lookupcache_new
    :: CSize
    -> CSize
    -> Ptr Stats
    -> Ptr (Ptr LookupCache)
    -> IO CString
  foreign import ccall unsafe "&glean_lookupcache_free" glean_lookupcache_free
    :: Destroy LookupCache
  foreign import ccall safe glean_lookupcache_clear
    :: Ptr LookupCache -> IO ()
  
  foreign import ccall unsafe glean_lookupcache_anchor_new
    :: Ptr Lookup
    -> Ptr LookupCache
    -> Int
    -> Ptr (Ptr Lookup)
    -> IO CString
  foreign import ccall unsafe glean_lookupcache_anchor_free
    :: Ptr Lookup -> IO ()
  