{-# LANGUAGE ScopedTypeVariables #-}

-- | Database scan for one gauge snapshot.
module Arbiter.Otel.Gauges.Scan
  ( scanSnapshot
  ) where

import Arbiter.Core.Health qualified as Health
import Arbiter.Core.Job.Schema (SchemaName, TableName)
import Arbiter.Core.MonadArbiter (MonadArbiter, withDbTransaction)
import Arbiter.Core.Operations
  ( QueueOverview (..)
  , QueueStats (..)
  , getAllQueueStats
  , listConcurrencyPolicies
  , listRateLimitPolicies
  , setLocalStatementTimeout
  )
import Data.Map.Strict qualified as Map
import Data.Text (Text)
import Data.Time (NominalDiffTime)

import Arbiter.Otel.Gauges.Cache (Snapshot (..))

-- | Read all database-global values used by gauge instruments.
scanSnapshot
  :: forall m
   . (MonadArbiter m)
  => NominalDiffTime
  -> SchemaName
  -> [(TableName, [Text])]
  -> m Snapshot
scanSnapshot :: forall (m :: * -> *).
MonadArbiter m =>
NominalDiffTime
-> SchemaName -> [(SchemaName, [SchemaName])] -> m Snapshot
scanSnapshot NominalDiffTime
statementTimeout SchemaName
schema [(SchemaName, [SchemaName])]
queueKinds = do
  overviews <- (QueueOverview -> QueueOverview)
-> [QueueOverview] -> [QueueOverview]
forall a b. (a -> b) -> [a] -> [b]
map QueueOverview -> QueueOverview
zeroFilled ([QueueOverview] -> [QueueOverview])
-> m [QueueOverview] -> m [QueueOverview]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m [QueueOverview] -> m [QueueOverview]
forall a. m a -> m a
boundedRead (SchemaName -> [(SchemaName, [SchemaName])] -> m [QueueOverview]
forall (m :: * -> *).
MonadArbiter m =>
SchemaName -> [(SchemaName, [SchemaName])] -> m [QueueOverview]
getAllQueueStats SchemaName
schema [(SchemaName, [SchemaName])]
queueKinds)
  (dbHealth, tableHealth) <- boundedRead (Health.getPgHealth schema queueTables)
  concurrencyPolicies <- boundedRead (listConcurrencyPolicies schema)
  rateLimitPolicies <- boundedRead (listRateLimitPolicies schema [])
  pure
    Snapshot
      { queues = overviews
      , db = dbHealth
      , tables = tableHealth
      , concurrency = concurrencyPolicies
      , rateLimits = rateLimitPolicies
      }
  where
    boundedRead :: forall a. m a -> m a
    boundedRead :: forall a. m a -> m a
boundedRead m a
query =
      m a -> m a
forall a. m a -> m a
forall (m :: * -> *) a. MonadArbiter m => m a -> m a
withDbTransaction (NominalDiffTime -> m ()
forall (m :: * -> *). MonadArbiter m => NominalDiffTime -> m ()
setLocalStatementTimeout NominalDiffTime
statementTimeout m () -> m a -> m a
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> m a
query)
    queueTables :: [SchemaName]
queueTables = ((SchemaName, [SchemaName]) -> SchemaName)
-> [(SchemaName, [SchemaName])] -> [SchemaName]
forall a b. (a -> b) -> [a] -> [b]
map (SchemaName, [SchemaName]) -> SchemaName
forall a b. (a, b) -> a
fst [(SchemaName, [SchemaName])]
queueKinds
    declaredKinds :: Map SchemaName [SchemaName]
declaredKinds = [(SchemaName, [SchemaName])] -> Map SchemaName [SchemaName]
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(SchemaName, [SchemaName])]
queueKinds
    zeroFilled :: QueueOverview -> QueueOverview
zeroFilled QueueOverview
overview =
      let zeros :: Map SchemaName Int64
zeros = [(SchemaName, Int64)] -> Map SchemaName Int64
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(SchemaName
kind, Int64
0) | SchemaName
kind <- [SchemaName]
-> SchemaName -> Map SchemaName [SchemaName] -> [SchemaName]
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault [] (QueueOverview -> SchemaName
overviewQueue QueueOverview
overview) Map SchemaName [SchemaName]
declaredKinds]
          stats :: QueueStats
stats = QueueOverview -> QueueStats
overviewStats QueueOverview
overview
       in QueueOverview
overview {overviewStats = stats {kindCounts = Map.union (kindCounts stats) zeros}}