{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TemplateHaskell #-}
module System.Metrics.Prometheus.Ridley.Metrics.DiskUsage (
  newDiskUsageMetrics
  ) where

import           Control.Monad
import           Control.Monad.IO.Class
import           Control.Monad.Trans.Class
import           Data.Maybe
import           Katip
import           Lens.Micro
import           Lens.Micro.TH
import           System.Exit
import           System.Metrics.Prometheus.Ridley.Types
import           System.Metrics.Prometheus.Ridley.Types.Internal
import           System.Process
import           System.Remote.Monitoring.Prometheus (labels)
import           Text.Read hiding (lift)
import qualified Data.Map.Strict as M
import qualified Data.Text as T
import qualified System.Metrics.Prometheus.Metric.Gauge as P
import qualified System.Metrics.Prometheus.MetricId as P
import qualified System.Metrics.Prometheus.RegistryT as P


--------------------------------------------------------------------------------
data DiskStats = DiskStats {
    DiskStats -> Text
_diskFilesystem :: T.Text
  , DiskStats -> Double
_diskUsed       :: Double
  , DiskStats -> Double
_diskFree       :: Double
  } deriving Int -> DiskStats -> ShowS
[DiskStats] -> ShowS
DiskStats -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [DiskStats] -> ShowS
$cshowList :: [DiskStats] -> ShowS
show :: DiskStats -> String
$cshow :: DiskStats -> String
showsPrec :: Int -> DiskStats -> ShowS
$cshowsPrec :: Int -> DiskStats -> ShowS
Show

makeLenses ''DiskStats

--------------------------------------------------------------------------------
data DiskMetric = DiskMetric {
    DiskMetric -> Gauge
_dskMetricUsed :: P.Gauge
  , DiskMetric -> Gauge
_dskMetricFree :: P.Gauge
  }

--------------------------------------------------------------------------------
type DiskUsageMetrics = M.Map T.Text DiskMetric

--------------------------------------------------------------------------------
getDiskStats :: Logger -> IO [DiskStats]
getDiskStats :: Logger -> IO [DiskStats]
getDiskStats Logger
logger = do
  let diskOnly :: DiskStats -> Bool
diskOnly = (\DiskStats
d -> Text
"/dev" Text -> Text -> Bool
`T.isInfixOf` (DiskStats
d forall s a. s -> Getting a s a -> a
^. Lens' DiskStats Text
diskFilesystem))
  let dropHeader :: String -> [Text]
dropHeader = forall a. Int -> [a] -> [a]
drop Int
1 forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> [Text]
T.lines forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Text
T.strip forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Text
T.pack
  (ExitCode
exitCode, String
rawLines, String
errors) <- String -> [String] -> String -> IO (ExitCode, String, String)
readProcessWithExitCode String
"df" [] []
  case ExitCode
exitCode of
    ExitCode
ExitSuccess    -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. (a -> Bool) -> [a] -> [a]
filter DiskStats -> Bool
diskOnly forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe Text -> Maybe DiskStats
mkDiskStats forall a b. (a -> b) -> a -> b
$ String -> [Text]
dropHeader String
rawLines
    ExitFailure Int
ec -> do
      Logger
logger Severity
ErrorS forall a b. (a -> b) -> a -> b
$ Text
"getDiskStats exited with error code " forall a. Semigroup a => a -> a -> a
<> String -> Text
T.pack (forall a. Show a => a -> String
show Int
ec) forall a. Semigroup a => a -> a -> a
<> Text
": " forall a. Semigroup a => a -> a -> a
<> String -> Text
T.pack String
errors
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Monoid a => a
mempty
  where
    mkDiskStats :: T.Text -> Maybe DiskStats
    mkDiskStats :: Text -> Maybe DiskStats
mkDiskStats Text
rawLine = case Text -> [Text]
T.words Text
rawLine of
#ifdef darwin_HOST_OS
     [fs,_, used,free,_,_,_,_,_] -> DiskStats <$> pure fs
                                              <*> readMaybe (T.unpack used)
                                              <*> readMaybe (T.unpack free)
#else
     -- On Linux, `df` shows less things by default, example
     -- Filesystem     1K-blocks     Used Available Use% Mounted on
     -- /dev/xvda1      52416860 27408532  25008328  53% /
     [Text
fs,Text
_, Text
used,Text
free,Text
_,Text
_] -> Text -> Double -> Double -> DiskStats
DiskStats forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (f :: * -> *) a. Applicative f => a -> f a
pure Text
fs
                                        forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a. Read a => String -> Maybe a
readMaybe (Text -> String
T.unpack Text
used)
                                        forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a. Read a => String -> Maybe a
readMaybe (Text -> String
T.unpack Text
free)
#endif
     [Text]
_ -> forall a. Maybe a
Nothing

--------------------------------------------------------------------------------
-- | As this is a gauge, it makes no sense flushing it.
updateDiskUsageMetric :: DiskMetric -> DiskStats -> Bool -> IO ()
updateDiskUsageMetric :: DiskMetric -> DiskStats -> Bool -> IO ()
updateDiskUsageMetric DiskMetric{Gauge
_dskMetricFree :: Gauge
_dskMetricUsed :: Gauge
_dskMetricFree :: DiskMetric -> Gauge
_dskMetricUsed :: DiskMetric -> Gauge
..} DiskStats
d Bool
_ = do
  Double -> Gauge -> IO ()
P.set (DiskStats
d forall s a. s -> Getting a s a -> a
^. Lens' DiskStats Double
diskUsed) Gauge
_dskMetricUsed
  Double -> Gauge -> IO ()
P.set (DiskStats
d forall s a. s -> Getting a s a -> a
^. Lens' DiskStats Double
diskFree) Gauge
_dskMetricFree

--------------------------------------------------------------------------------
updateDiskUsageMetrics :: Logger -> DiskUsageMetrics -> Bool -> IO ()
updateDiskUsageMetrics :: Logger -> DiskUsageMetrics -> Bool -> IO ()
updateDiskUsageMetrics Logger
logger DiskUsageMetrics
dmetrics Bool
flush = do
  [DiskStats]
diskStats <- Logger -> IO [DiskStats]
getDiskStats Logger
logger
  forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [DiskStats]
diskStats forall a b. (a -> b) -> a -> b
$ \DiskStats
d -> do
    let key :: Text
key = DiskStats
d forall s a. s -> Getting a s a -> a
^. Lens' DiskStats Text
diskFilesystem
    case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Text
key DiskUsageMetrics
dmetrics of
      Maybe DiskMetric
Nothing -> forall (m :: * -> *) a. Monad m => a -> m a
return ()
      Just DiskMetric
m  -> DiskMetric -> DiskStats -> Bool -> IO ()
updateDiskUsageMetric DiskMetric
m DiskStats
d Bool
flush

--------------------------------------------------------------------------------
mkDiskGauge :: MonadIO m => P.Labels -> DiskUsageMetrics -> DiskStats -> P.RegistryT m DiskUsageMetrics
mkDiskGauge :: forall (m :: * -> *).
MonadIO m =>
Labels
-> DiskUsageMetrics -> DiskStats -> RegistryT m DiskUsageMetrics
mkDiskGauge Labels
currentLabels DiskUsageMetrics
dmap DiskStats
d = do
  let fs :: Text
fs = DiskStats
d forall s a. s -> Getting a s a -> a
^. Lens' DiskStats Text
diskFilesystem
  let finalLabels :: Labels
finalLabels = Text -> Text -> Labels -> Labels
P.addLabel Text
"filesystem" Text
fs Labels
currentLabels
  DiskMetric
metric <- Gauge -> Gauge -> DiskMetric
DiskMetric forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *).
MonadIO m =>
Name -> Labels -> RegistryT m Gauge
P.registerGauge Name
"disk_used_bytes_blocks" Labels
finalLabels
                       forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (m :: * -> *).
MonadIO m =>
Name -> Labels -> RegistryT m Gauge
P.registerGauge Name
"disk_free_bytes_blocks" Labels
finalLabels
  forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ DiskMetric -> DiskStats -> Bool -> IO ()
updateDiskUsageMetric DiskMetric
metric DiskStats
d Bool
False
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Text
fs DiskMetric
metric forall a b. (a -> b) -> a -> b
$! DiskUsageMetrics
dmap

-- | Creates a new 'RidleyMetricHandler' to monitor disk usage.
newDiskUsageMetrics :: Ridley RidleyMetricHandler
newDiskUsageMetrics :: Ridley RidleyMetricHandler
newDiskUsageMetrics = do
  Logger
logger <- Ridley Logger
ioLogger
  RidleyOptions
opts <- Ridley RidleyOptions
getRidleyOptions
  [DiskStats]
diskStats <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Logger -> IO [DiskStats]
getDiskStats Logger
logger)
  DiskUsageMetrics
metrics   <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (forall (m :: * -> *).
MonadIO m =>
Labels
-> DiskUsageMetrics -> DiskStats -> RegistryT m DiskUsageMetrics
mkDiskGauge (RidleyOptions
opts forall s a. s -> Getting a s a -> a
^. Lens' RidleyOptions PrometheusOptions
prometheusOptions forall b c a. (b -> c) -> (a -> b) -> a -> c
. Lens' PrometheusOptions Labels
labels)) forall k a. Map k a
M.empty [DiskStats]
diskStats
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall c.
HasCallStack =>
Text -> c -> (c -> Bool -> IO ()) -> Bool -> RidleyMetricHandler
mkRidleyMetricHandler Text
"ridley-disk-usage" DiskUsageMetrics
metrics (Logger -> DiskUsageMetrics -> Bool -> IO ()
updateDiskUsageMetrics Logger
logger) Bool
False