{-# 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
(Int -> DiskStats -> ShowS)
-> (DiskStats -> String)
-> ([DiskStats] -> ShowS)
-> Show DiskStats
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> DiskStats -> ShowS
showsPrec :: Int -> DiskStats -> ShowS
$cshow :: DiskStats -> String
show :: DiskStats -> String
$cshowList :: [DiskStats] -> ShowS
showList :: [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 DiskStats -> Getting Text DiskStats Text -> Text
forall s a. s -> Getting a s a -> a
^. Getting Text DiskStats Text
Lens' DiskStats Text
diskFilesystem))
  let dropHeader :: String -> [Text]
dropHeader = Int -> [Text] -> [Text]
forall a. Int -> [a] -> [a]
drop Int
1 ([Text] -> [Text]) -> (String -> [Text]) -> String -> [Text]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> [Text]
T.lines (Text -> [Text]) -> (String -> Text) -> String -> [Text]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Text
T.strip (Text -> Text) -> (String -> Text) -> String -> Text
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    -> [DiskStats] -> IO [DiskStats]
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ([DiskStats] -> IO [DiskStats]) -> [DiskStats] -> IO [DiskStats]
forall a b. (a -> b) -> a -> b
$ (DiskStats -> Bool) -> [DiskStats] -> [DiskStats]
forall a. (a -> Bool) -> [a] -> [a]
filter DiskStats -> Bool
diskOnly ([DiskStats] -> [DiskStats])
-> ([Text] -> [DiskStats]) -> [Text] -> [DiskStats]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Text -> Maybe DiskStats) -> [Text] -> [DiskStats]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe Text -> Maybe DiskStats
mkDiskStats ([Text] -> [DiskStats]) -> [Text] -> [DiskStats]
forall a b. (a -> b) -> a -> b
$ String -> [Text]
dropHeader String
rawLines
    ExitFailure Int
ec -> do
      Logger
logger Severity
ErrorS (Text -> IO ()) -> Text -> IO ()
forall a b. (a -> b) -> a -> b
$ Text
"getDiskStats exited with error code " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
T.pack (Int -> String
forall a. Show a => a -> String
show Int
ec) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
": " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
T.pack String
errors
      [DiskStats] -> IO [DiskStats]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [DiskStats]
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 (Text -> Double -> Double -> DiskStats)
-> Maybe Text -> Maybe (Double -> Double -> DiskStats)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Text -> Maybe Text
forall a. a -> Maybe a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Text
fs
                                        Maybe (Double -> Double -> DiskStats)
-> Maybe Double -> Maybe (Double -> DiskStats)
forall a b. Maybe (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> String -> Maybe Double
forall a. Read a => String -> Maybe a
readMaybe (Text -> String
T.unpack Text
used)
                                        Maybe (Double -> DiskStats) -> Maybe Double -> Maybe DiskStats
forall a b. Maybe (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> String -> Maybe Double
forall a. Read a => String -> Maybe a
readMaybe (Text -> String
T.unpack Text
free)
#endif
     [Text]
_ -> Maybe DiskStats
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
_dskMetricUsed :: DiskMetric -> Gauge
_dskMetricFree :: DiskMetric -> Gauge
_dskMetricUsed :: Gauge
_dskMetricFree :: Gauge
..} DiskStats
d Bool
_ = do
  Double -> Gauge -> IO ()
P.set (DiskStats
d DiskStats -> Getting Double DiskStats Double -> Double
forall s a. s -> Getting a s a -> a
^. Getting Double DiskStats Double
Lens' DiskStats Double
diskUsed) Gauge
_dskMetricUsed
  Double -> Gauge -> IO ()
P.set (DiskStats
d DiskStats -> Getting Double DiskStats Double -> Double
forall s a. s -> Getting a s a -> a
^. Getting Double DiskStats Double
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
  [DiskStats] -> (DiskStats -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [DiskStats]
diskStats ((DiskStats -> IO ()) -> IO ()) -> (DiskStats -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \DiskStats
d -> do
    let key :: Text
key = DiskStats
d DiskStats -> Getting Text DiskStats Text -> Text
forall s a. s -> Getting a s a -> a
^. Getting Text DiskStats Text
Lens' DiskStats Text
diskFilesystem
    case Text -> DiskUsageMetrics -> Maybe DiskMetric
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Text
key DiskUsageMetrics
dmetrics of
      Maybe DiskMetric
Nothing -> () -> IO ()
forall a. a -> IO a
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 DiskStats -> Getting Text DiskStats Text -> Text
forall s a. s -> Getting a s a -> a
^. Getting Text DiskStats Text
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 (Gauge -> Gauge -> DiskMetric)
-> RegistryT m Gauge -> RegistryT m (Gauge -> DiskMetric)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> Labels -> RegistryT m Gauge
forall (m :: * -> *).
MonadIO m =>
Name -> Labels -> RegistryT m Gauge
P.registerGauge Name
"disk_used_bytes_blocks" Labels
finalLabels
                       RegistryT m (Gauge -> DiskMetric)
-> RegistryT m Gauge -> RegistryT m DiskMetric
forall a b. RegistryT m (a -> b) -> RegistryT m a -> RegistryT m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Name -> Labels -> RegistryT m Gauge
forall (m :: * -> *).
MonadIO m =>
Name -> Labels -> RegistryT m Gauge
P.registerGauge Name
"disk_free_bytes_blocks" Labels
finalLabels
  IO () -> RegistryT m ()
forall a. IO a -> RegistryT m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> RegistryT m ()) -> IO () -> RegistryT m ()
forall a b. (a -> b) -> a -> b
$ DiskMetric -> DiskStats -> Bool -> IO ()
updateDiskUsageMetric DiskMetric
metric DiskStats
d Bool
False
  DiskUsageMetrics -> RegistryT m DiskUsageMetrics
forall a. a -> RegistryT m a
forall (m :: * -> *) a. Monad m => a -> m a
return (DiskUsageMetrics -> RegistryT m DiskUsageMetrics)
-> DiskUsageMetrics -> RegistryT m DiskUsageMetrics
forall a b. (a -> b) -> a -> b
$! Text -> DiskMetric -> DiskUsageMetrics -> DiskUsageMetrics
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Text
fs DiskMetric
metric (DiskUsageMetrics -> DiskUsageMetrics)
-> DiskUsageMetrics -> DiskUsageMetrics
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 <- IO [DiskStats]
-> RidleyT (RegistryT (KatipContextT IO)) [DiskStats]
forall a. IO a -> RidleyT (RegistryT (KatipContextT IO)) a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Logger -> IO [DiskStats]
getDiskStats Logger
logger)
  DiskUsageMetrics
metrics   <- RegistryT (KatipContextT IO) DiskUsageMetrics
-> RidleyT (RegistryT (KatipContextT IO)) DiskUsageMetrics
forall (m :: * -> *) a. Monad m => m a -> RidleyT m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (RegistryT (KatipContextT IO) DiskUsageMetrics
 -> RidleyT (RegistryT (KatipContextT IO)) DiskUsageMetrics)
-> RegistryT (KatipContextT IO) DiskUsageMetrics
-> RidleyT (RegistryT (KatipContextT IO)) DiskUsageMetrics
forall a b. (a -> b) -> a -> b
$ (DiskUsageMetrics
 -> DiskStats -> RegistryT (KatipContextT IO) DiskUsageMetrics)
-> DiskUsageMetrics
-> [DiskStats]
-> RegistryT (KatipContextT IO) DiskUsageMetrics
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (Labels
-> DiskUsageMetrics
-> DiskStats
-> RegistryT (KatipContextT IO) DiskUsageMetrics
forall (m :: * -> *).
MonadIO m =>
Labels
-> DiskUsageMetrics -> DiskStats -> RegistryT m DiskUsageMetrics
mkDiskGauge (RidleyOptions
opts RidleyOptions -> Getting Labels RidleyOptions Labels -> Labels
forall s a. s -> Getting a s a -> a
^. (PrometheusOptions -> Const Labels PrometheusOptions)
-> RidleyOptions -> Const Labels RidleyOptions
Lens' RidleyOptions PrometheusOptions
prometheusOptions ((PrometheusOptions -> Const Labels PrometheusOptions)
 -> RidleyOptions -> Const Labels RidleyOptions)
-> ((Labels -> Const Labels Labels)
    -> PrometheusOptions -> Const Labels PrometheusOptions)
-> Getting Labels RidleyOptions Labels
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Labels -> Const Labels Labels)
-> PrometheusOptions -> Const Labels PrometheusOptions
Lens' PrometheusOptions Labels
labels)) DiskUsageMetrics
forall k a. Map k a
M.empty [DiskStats]
diskStats
  RidleyMetricHandler -> Ridley RidleyMetricHandler
forall a. a -> RidleyT (RegistryT (KatipContextT IO)) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (RidleyMetricHandler -> Ridley RidleyMetricHandler)
-> RidleyMetricHandler -> Ridley RidleyMetricHandler
forall a b. (a -> b) -> a -> b
$ Text
-> DiskUsageMetrics
-> (DiskUsageMetrics -> Bool -> IO ())
-> Bool
-> RidleyMetricHandler
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