{-# LANGUAGE DeriveAnyClass #-}

-- | Shared state for gauge instruments and the refresh loop.
module Arbiter.Otel.Gauges.Cache
  ( Snapshot (..)
  , Cached (..)
  , Export (..)
  , live
  , lastScan
  , retire
  , GaugeCache (..)
  , Baseline
  , SeriesKey
  , newGaugeCache
  , publishSnapshot
  , setReachable
  , retireCache
  , riseSince
  ) where

import Arbiter.Core.Concurrency.Stats qualified as Conc (ConcurrencyPolicyView)
import Arbiter.Core.Health qualified as Health
import Arbiter.Core.Operations (QueueOverview)
import Arbiter.Core.RateLimit.Stats qualified as RL (RateLimitPolicyView)
import Control.Concurrent.STM (STM, TVar, modifyTVar', newTVarIO, writeTVar)
import Data.Aeson (FromJSON, ToJSON)
import Data.HashMap.Strict (HashMap)
import Data.HashMap.Strict qualified as HM
import Data.IORef (IORef, newIORef)
import Data.Text (Text)
import GHC.Generics (Generic)

-- | Values from one database scan.
data Snapshot = Snapshot
  { Snapshot -> [QueueOverview]
queues :: [QueueOverview]
  , Snapshot -> Maybe PgDbHealth
db :: Maybe Health.PgDbHealth
  , Snapshot -> [PgTableHealth]
tables :: [Health.PgTableHealth]
  , Snapshot -> [ConcurrencyPolicyView]
concurrency :: [Conc.ConcurrencyPolicyView]
  , Snapshot -> [RateLimitPolicyView]
rateLimits :: [RL.RateLimitPolicyView]
  }
  deriving stock ((forall x. Snapshot -> Rep Snapshot x)
-> (forall x. Rep Snapshot x -> Snapshot) -> Generic Snapshot
forall x. Rep Snapshot x -> Snapshot
forall x. Snapshot -> Rep Snapshot x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. Snapshot -> Rep Snapshot x
from :: forall x. Snapshot -> Rep Snapshot x
$cto :: forall x. Rep Snapshot x -> Snapshot
to :: forall x. Rep Snapshot x -> Snapshot
Generic)
  deriving anyclass (Maybe Snapshot
Value -> Parser [Snapshot]
Value -> Parser Snapshot
(Value -> Parser Snapshot)
-> (Value -> Parser [Snapshot])
-> Maybe Snapshot
-> FromJSON Snapshot
forall a.
(Value -> Parser a)
-> (Value -> Parser [a]) -> Maybe a -> FromJSON a
$cparseJSON :: Value -> Parser Snapshot
parseJSON :: Value -> Parser Snapshot
$cparseJSONList :: Value -> Parser [Snapshot]
parseJSONList :: Value -> Parser [Snapshot]
$comittedField :: Maybe Snapshot
omittedField :: Maybe Snapshot
FromJSON, [Snapshot] -> Value
[Snapshot] -> Encoding
Snapshot -> Bool
Snapshot -> Value
Snapshot -> Encoding
(Snapshot -> Value)
-> (Snapshot -> Encoding)
-> ([Snapshot] -> Value)
-> ([Snapshot] -> Encoding)
-> (Snapshot -> Bool)
-> ToJSON Snapshot
forall a.
(a -> Value)
-> (a -> Encoding)
-> ([a] -> Value)
-> ([a] -> Encoding)
-> (a -> Bool)
-> ToJSON a
$ctoJSON :: Snapshot -> Value
toJSON :: Snapshot -> Value
$ctoEncoding :: Snapshot -> Encoding
toEncoding :: Snapshot -> Encoding
$ctoJSONList :: [Snapshot] -> Value
toJSONList :: [Snapshot] -> Value
$ctoEncodingList :: [Snapshot] -> Encoding
toEncodingList :: [Snapshot] -> Encoding
$comitField :: Snapshot -> Bool
omitField :: Snapshot -> Bool
ToJSON)

-- | A snapshot and the monotonic time at which its scan started.
data Cached = Cached
  { Cached -> Double
takenAt :: Double
  , Cached -> Snapshot
reading :: Snapshot
  }

-- | What the instruments export. 'Idle' keeps the last reading's scan time for the
-- staleness series. Before the first scan it has none.
data Export = Live Cached | Idle (Maybe Double)

-- | The scan behind an export, if it has one.
live :: Export -> Maybe Cached
live :: Export -> Maybe Cached
live = \case
  Live Cached
cached -> Cached -> Maybe Cached
forall a. a -> Maybe a
Just Cached
cached
  Idle Maybe Double
_ -> Maybe Cached
forall a. Maybe a
Nothing

-- | When the export was last scanned, if it ever was.
lastScan :: Export -> Maybe Double
lastScan :: Export -> Maybe Double
lastScan = \case
  Live Cached
cached -> Double -> Maybe Double
forall a. a -> Maybe a
Just (Cached -> Double
takenAt Cached
cached)
  Idle Maybe Double
lastAt -> Maybe Double
lastAt

-- | Stop exporting readings, keeping when the last one was scanned.
retire :: Export -> Export
retire :: Export -> Export
retire = Maybe Double -> Export
Idle (Maybe Double -> Export)
-> (Export -> Maybe Double) -> Export -> Export
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Export -> Maybe Double
lastScan

-- | One counter series: its instrument and attributes.
type SeriesKey = (Text, [(Text, Text)])

-- | The scan a counter series was last counted from, and the total it stood at.
data Baseline = Baseline
  { Baseline -> Double
countedFrom :: !Double
  , Baseline -> Double
countedTotal :: !Double
  }

-- | Mutable gauge state and its registration time.
data GaugeCache = GaugeCache
  { GaugeCache -> TVar Export
export :: TVar Export
  , GaugeCache -> TVar (Maybe Bool)
databaseReachable :: TVar (Maybe Bool)
  , GaugeCache -> IORef (HashMap SeriesKey Baseline)
counterBaselines :: IORef (HashMap SeriesKey Baseline)
  , GaugeCache -> Double
registeredAt :: Double
  }

-- | Create an empty gauge cache for a registration starting at @now@.
newGaugeCache :: Double -> IO GaugeCache
newGaugeCache :: Double -> IO GaugeCache
newGaugeCache Double
now =
  TVar Export
-> TVar (Maybe Bool)
-> IORef (HashMap SeriesKey Baseline)
-> Double
-> GaugeCache
GaugeCache
    (TVar Export
 -> TVar (Maybe Bool)
 -> IORef (HashMap SeriesKey Baseline)
 -> Double
 -> GaugeCache)
-> IO (TVar Export)
-> IO
     (TVar (Maybe Bool)
      -> IORef (HashMap SeriesKey Baseline) -> Double -> GaugeCache)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Export -> IO (TVar Export)
forall a. a -> IO (TVar a)
newTVarIO (Maybe Double -> Export
Idle Maybe Double
forall a. Maybe a
Nothing)
    IO
  (TVar (Maybe Bool)
   -> IORef (HashMap SeriesKey Baseline) -> Double -> GaugeCache)
-> IO (TVar (Maybe Bool))
-> IO (IORef (HashMap SeriesKey Baseline) -> Double -> GaugeCache)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe Bool -> IO (TVar (Maybe Bool))
forall a. a -> IO (TVar a)
newTVarIO Maybe Bool
forall a. Maybe a
Nothing
    IO (IORef (HashMap SeriesKey Baseline) -> Double -> GaugeCache)
-> IO (IORef (HashMap SeriesKey Baseline))
-> IO (Double -> GaugeCache)
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> HashMap SeriesKey Baseline
-> IO (IORef (HashMap SeriesKey Baseline))
forall a. a -> IO (IORef a)
newIORef HashMap SeriesKey Baseline
forall k v. HashMap k v
HM.empty
    IO (Double -> GaugeCache) -> IO Double -> IO GaugeCache
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Double -> IO Double
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Double
now

-- | Publish a snapshot to the observable instruments.
publishSnapshot :: GaugeCache -> Cached -> STM ()
publishSnapshot :: GaugeCache -> Cached -> STM ()
publishSnapshot GaugeCache
cache = TVar Export -> Export -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (GaugeCache -> TVar Export
export GaugeCache
cache) (Export -> STM ()) -> (Cached -> Export) -> Cached -> STM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Cached -> Export
Live

-- | Update the result of the last database operation.
setReachable :: GaugeCache -> Bool -> STM ()
setReachable :: GaugeCache -> Bool -> STM ()
setReachable GaugeCache
cache = TVar (Maybe Bool) -> Maybe Bool -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (GaugeCache -> TVar (Maybe Bool)
databaseReachable GaugeCache
cache) (Maybe Bool -> STM ()) -> (Bool -> Maybe Bool) -> Bool -> STM ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> Maybe Bool
forall a. a -> Maybe a
Just

-- | Stop exporting the cached snapshot.
retireCache :: GaugeCache -> STM ()
retireCache :: GaugeCache -> STM ()
retireCache GaugeCache
cache = TVar Export -> (Export -> Export) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' (GaugeCache -> TVar Export
export GaugeCache
cache) Export -> Export
retire

-- | What a total scanned at @scannedAt@ adds to its series. The first reading and an
-- already counted reading add nothing. A reset counter adds the whole total. Any other
-- reading adds the difference.
riseSince
  :: SeriesKey
  -> Double
  -> Double
  -> HashMap SeriesKey Baseline
  -> (HashMap SeriesKey Baseline, Double)
riseSince :: SeriesKey
-> Double
-> Double
-> HashMap SeriesKey Baseline
-> (HashMap SeriesKey Baseline, Double)
riseSince SeriesKey
key Double
scannedAt Double
total HashMap SeriesKey Baseline
seen = case SeriesKey -> HashMap SeriesKey Baseline -> Maybe Baseline
forall k v. Hashable k => k -> HashMap k v -> Maybe v
HM.lookup SeriesKey
key HashMap SeriesKey Baseline
seen of
  Just Baseline
base | Baseline -> Double
countedFrom Baseline
base Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
>= Double
scannedAt -> (HashMap SeriesKey Baseline
seen, Double
0)
  Just Baseline
base -> (HashMap SeriesKey Baseline
counted, if Double
total Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
< Baseline -> Double
countedTotal Baseline
base then Double
total else Double
total Double -> Double -> Double
forall a. Num a => a -> a -> a
- Baseline -> Double
countedTotal Baseline
base)
  Maybe Baseline
Nothing -> (HashMap SeriesKey Baseline
counted, Double
0)
  where
    counted :: HashMap SeriesKey Baseline
counted = SeriesKey
-> Baseline
-> HashMap SeriesKey Baseline
-> HashMap SeriesKey Baseline
forall k v. Hashable k => k -> v -> HashMap k v -> HashMap k v
HM.insert SeriesKey
key (Double -> Double -> Baseline
Baseline Double
scannedAt Double
total) HashMap SeriesKey Baseline
seen