{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TemplateHaskell #-}
module System.Metrics.Prometheus.Ridley.Metrics.DiskUsage (
    getDiskStats
  , mkDiskGauge
  , diskUsageMetrics
  ) where

import           Control.Monad
import           Control.Monad.IO.Class
import qualified Data.Map.Strict as M
import           Data.Maybe
import qualified Data.Text as T
import           Lens.Micro
import           Lens.Micro.TH
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
import           System.Metrics.Prometheus.Ridley.Types
import           System.Process
import           Text.Read


--------------------------------------------------------------------------------
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
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 :: IO [DiskStats]
getDiskStats :: IO [DiskStats]
getDiskStats = 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 :: [a] -> [a]
dropHeader = Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop Int
1
  [Text]
rawLines <- [Text] -> [Text]
forall a. [a] -> [a]
dropHeader ([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 (String -> [Text]) -> IO String -> IO [Text]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> [String] -> String -> IO String
readProcess String
"df" [] []
  [DiskStats] -> IO [DiskStats]
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
$ [Text]
rawLines
  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 (f :: * -> *) a. Applicative f => a -> f a
pure Text
fs
                                        Maybe (Double -> Double -> DiskStats)
-> Maybe Double -> Maybe (Double -> DiskStats)
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 (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
_dskMetricFree :: Gauge
_dskMetricUsed :: Gauge
_dskMetricFree :: DiskMetric -> Gauge
_dskMetricUsed :: DiskMetric -> 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 :: DiskUsageMetrics -> Bool -> IO ()
updateDiskUsageMetrics :: DiskUsageMetrics -> Bool -> IO ()
updateDiskUsageMetrics DiskUsageMetrics
dmetrics Bool
flush = do
  [DiskStats]
diskStats <- IO [DiskStats]
getDiskStats
  [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 (m :: * -> *) a. Monad m => a -> m a
return ()
      Just DiskMetric
m  -> DiskMetric -> DiskStats -> Bool -> IO ()
updateDiskUsageMetric DiskMetric
m DiskStats
d Bool
flush

--------------------------------------------------------------------------------
diskUsageMetrics :: DiskUsageMetrics -> RidleyMetricHandler
diskUsageMetrics :: DiskUsageMetrics -> RidleyMetricHandler
diskUsageMetrics DiskUsageMetrics
g = Text
-> DiskUsageMetrics
-> (DiskUsageMetrics -> Bool -> IO ())
-> Bool
-> RidleyMetricHandler
forall c.
HasCallStack =>
Text -> c -> (c -> Bool -> IO ()) -> Bool -> RidleyMetricHandler
mkRidleyMetricHandler Text
"ridley-disk-usage" DiskUsageMetrics
g DiskUsageMetrics -> Bool -> IO ()
updateDiskUsageMetrics Bool
False

--------------------------------------------------------------------------------
mkDiskGauge :: MonadIO m => P.Labels -> DiskUsageMetrics -> DiskStats -> P.RegistryT m DiskUsageMetrics
mkDiskGauge :: 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 (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 (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 (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