{-
  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.Write.Stats (
  Stats, new, Bump, bump, tick,
  mutatorLatency, mutatorInput, mutatorThroughput, mutatorDedupedThroughput,
  mutatorDupThroughput,
  renameThroughput, commitThroughput,
  lookupCacheStats
) where

import Control.Monad
import Data.ByteString (ByteString)
import qualified Data.ByteString.Char8 as Char8
import Data.Char (toLower)
import Data.Functor.Contravariant
import Data.IORef
import Data.Word
import System.Clock (TimeSpec)

import qualified Glean.RTS.Foreign.LookupCache as LookupCache

import ServiceData.GlobalStats
import ServiceData.Types
import Util.Log

import Glean.Util.Metric
import Glean.Util.Throttle

data Stat a = Stat
  { statNew :: !a
  , statAll :: !a
  }

data Stats = Stats
  { statsLogger :: Execute
  , statsMutatorLatency :: !(IORef (Stat Tick))
  , statsMutatorInput :: !(IORef (Stat Tick))
  , statsMutatorThroughput :: !(IORef (Stat Tick))
  , statsMutatorDedupedThroughput :: !(IORef (Stat Tick))
  , statsMutatorDupThroughput :: !(IORef (Stat Tick))
  , statsRenameThroughput :: !(IORef (Stat Tick))
  , statsCommitThroughput :: !(IORef (Stat Tick))
  , statsLookupCache :: !(IORef (Stat LookupCache.StatValues))
  }

data Ref = forall a. Monoid a => Ref (Stats -> IORef (Stat a))

refs :: [Ref]
refs =
  [ Ref statsMutatorLatency
  , Ref statsMutatorInput
  , Ref statsMutatorThroughput
  , Ref statsMutatorDedupedThroughput
  , Ref statsMutatorDupThroughput
  , Ref statsRenameThroughput
  , Ref statsCommitThroughput
  , Ref statsLookupCache
  ]

bumper :: Monoid a => (Stats -> IORef (Stat a)) -> Bump a
bumper f = Bump $ \stats x -> atomicModifyIORef' (f stats) $
  \(Stat new all) -> (Stat (new <> x) (all <> x), ())

-- | A metric that can be bumped using 'bump'.  Metrics of the same
-- type can be bumped together using the `Monoid` instance.
newtype Bump a = Bump { runBump :: Stats -> a -> IO () }

instance Semigroup (Bump a) where
  Bump p <> Bump q = Bump $ \stats a -> p stats a >> q stats a

instance Monoid (Bump a) where
  mempty = Bump $ const $ const $ return ()

instance Contravariant Bump where
  contramap f (Bump g) = Bump $ \stats a -> g stats (f a)

-- | The latency between a write request being received, and it being completed.
mutatorLatency :: Bump Tick
mutatorLatency =
  bumper statsMutatorLatency <>
  bumpServiceCounterLatency "glean.db.write.latency_ms"

-- | The throughput of writes in bytes per unit time
mutatorThroughput :: Bump Tick
mutatorThroughput =
  bumper statsMutatorThroughput <>
  contramap tickValue (bumpServiceCounter "glean.db.write.throughput")

mutatorInput :: Bump Tick
mutatorInput =
  bumper statsMutatorInput

-- | The throughput of writes in bytes per unit time
mutatorDedupedThroughput :: Bump Tick
mutatorDedupedThroughput =
  bumper statsMutatorDedupedThroughput <>
  contramap tickValue (bumpServiceCounter "glean.db.write.deduped.bytes")

-- | The throughput of writes in bytes per unit time
mutatorDupThroughput :: Bump Tick
mutatorDupThroughput =
  bumper statsMutatorDupThroughput <>
  contramap tickValue (bumpServiceCounter "glean.db.write.undeduped.bytes") <>
  contramap tickMillis (bumpServiceCounter "glean.db.write.undeduped.ms")

-- | The throughput of renames in bytes per unit time
renameThroughput :: Bump Tick
renameThroughput =
  bumper statsRenameThroughput <>
  contramap tickValue (bumpServiceCounter "glean.db.write.rename.bytes") <>
  contramap tickMillis (bumpServiceCounter "glean.db.write.rename.ms")

-- | The throughput of commits in bytes per unit time
commitThroughput :: Bump Tick
commitThroughput =
  bumper statsCommitThroughput <>
  contramap tickValue (bumpServiceCounter "glean.db.write.commit.bytes") <>
  contramap tickMillis (bumpServiceCounter "glean.db.write.commit.ms")

lookupCacheStats :: Bump LookupCache.StatValues
lookupCacheStats = mconcat $ bumper statsLookupCache
  : [contramap (`LookupCache.getStat` c) bump | (c,bump) <- lookupCacheCounters]

lookupCacheCounters :: [(LookupCache.Stat, Bump Word64)]
lookupCacheCounters =
  [(c, bump c $ Char8.pack $ "glean.db.write.cache." ++ map toLower (show c))
    | c <- [minBound .. maxBound]]
  where
    bump c
      | LookupCache.isCounter c = bumpServiceCounter
      | otherwise = setServiceCounter

-- | Bump an fb303 counter representing latency
bumpServiceCounterLatency :: ByteString -> Bump Tick
bumpServiceCounterLatency counter = Bump $ \_ tick ->
  let ms = fromIntegral (tickLatencyNanoSecs tick `div` 1000000) in
  addStatValueType counter ms Avg

-- | Bump an fb303 counter
bumpServiceCounter :: ByteString -> Bump Word64
bumpServiceCounter counter = Bump $ \_ value ->
  addStatValueType counter (fromIntegral value) Sum

-- | Set an fb303 counter
setServiceCounter :: ByteString -> Bump Word64
setServiceCounter counter = Bump $ \_ value ->
  void $ setCounter counter (fromIntegral value)

allStats :: [(String, Stats -> IO String)]
allStats =
  [("mut_lat", counter showLatency statsMutatorLatency)
  ,("mut_thp", counter showThroughput statsMutatorThroughput)
  ,("ded_thp", counter showThroughput statsMutatorDedupedThroughput)
  ,("dup_thp", counter showThroughput statsMutatorDupThroughput)
  ,("rnm_thp", counter showThroughput statsRenameThroughput)
  ,("cmt_thp", counter showThroughput statsCommitThroughput)
  ,("ibk_mis", counter
    (showMissRate LookupCache.IdByKey_hits LookupCache.IdByKey_misses)
    statsLookupCache)
  ,("tbi_mis", counter
    (showMissRate LookupCache.TypeById_hits LookupCache.TypeById_misses)
    statsLookupCache)
  ,("fbi_mis", counter
    (showMissRate LookupCache.FactById_hits LookupCache.FactById_misses)
    statsLookupCache)
  ,("lch_mem", value
    (LookupCache.Fact_bytes `using` showMemory)
    statsLookupCache)
  ,("lch_cnt", value
    (LookupCache.Fact_count `using` showCount)
    statsLookupCache)
  ]
  where
    counter :: Monoid a
            => (a -> String) -> (Stats -> IORef (Stat a)) -> Stats -> IO String
    counter shw get stats = do
      stat <- readIORef (get stats)
      return $ concat [shw $ statNew stat, " [", shw $ statAll stat, "]"]

    showMissRate h m stats = pad 5 $ case hits+misses of
      0 -> "-"
      n ->
        let !rate = (misses * 1000) `div` n
        in concat [show (rate `div` 10), ".", show (rate `mod` 10), "%"]
      where
        hits = LookupCache.getStat stats h
        misses = LookupCache.getStat stats m

    value :: (a -> String) -> (Stats -> IORef (Stat a)) -> Stats -> IO String
    value shw get stats = shw . statNew <$> readIORef (get stats)

    using c f stats = f (LookupCache.getStat stats c)

class Initial a where
  initial :: IO a -> IO Stats

instance Initial Stats where
  initial = id

instance (Monoid a, Initial b) => Initial (IORef (Stat a) -> b) where
  initial p = initial $ do
    f <- p
    f <$> newIORef (Stat mempty mempty)

new :: TimeSpec -> IO Stats
new period = initial $ Stats <$> atMostEvery period

-- | Bump a metric
bump :: Stats -> Bump a -> a -> IO ()
bump stats b a = do
  runBump b stats a
  statsLogger stats $ do
    s <- fmap unwords $ forM allStats $ \(tag, f) -> do
      s <- f stats
      return $ tag ++ ": " ++ s
    logInfo s
    forM_ refs $ \(Ref f) -> atomicModifyIORef' (f stats) $
      \stat -> (stat { statNew = mempty }, ())

tick :: Stats -> Bump Tick -> Word64 -> IO a -> IO a
tick stats b val io = do
  p <- beginTick val
  x <- io
  t <- endTick p
  bump stats b t
  return x