init

1 files changed, 520 insertions, 0 deletions

diff --git a/app/Sensor.hs b/app/Sensor.hs
new file mode 100644
index 0000000..2342807
--- /dev/null
+++ b/app/Sensor.hs
@@ -0,0 +1,520 @@
+{-# LANGUAGE BlockArguments #-}
+{-# LANGUAGE FunctionalDependencies #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE OverloadedRecordDot #-}
+{-# LANGUAGE QuasiQuotes #-}
+{-# LANGUAGE StrictData #-}
+{-# LANGUAGE ViewPatterns #-}
+{-# OPTIONS_GHC -fno-warn-name-shadowing #-}
+{-# OPTIONS_GHC -fno-warn-type-defaults #-}
+{-# OPTIONS_GHC -fno-warn-unused-top-binds #-}
+
+module Sensor
+  ( SensorM,
+    runSensorM,
+    poll,
+    Sensor (..),
+    cpu,
+    io,
+    net,
+    mem,
+    bat,
+    Sensor.snd,
+    disk,
+    date,
+    time,
+    wmName,
+    wmWorkspaces,
+  )
+where
+
+import Control.Monad
+import Control.Monad.Reader
+import Data.Char
+import Data.Dynamic
+import Data.Functor.WithIndex
+import Data.List
+import Data.Map qualified as M
+import Data.Maybe
+import Data.Set qualified as S
+import Data.Time.Clock
+import Data.Time.Format
+import Data.Time.LocalTime
+import Graphics.X11 qualified as X
+import Graphics.X11.Xlib.Extras qualified as X
+import Pretty qualified as P
+import Process.Shell (sh)
+import Safe
+import System.FilePath
+import System.IO.Error
+import System.Linux.Inotify qualified as I
+import System.Posix.StatVFS
+import Text.Printf
+import UnliftIO.Concurrent
+import UnliftIO.Exception
+import UnliftIO.Memoize
+import UnliftIO.STM
+import Witherable (ifilter)
+
+type SensorM a =
+  ReaderT (MVar (M.Map String (Memoized Dynamic))) IO a
+
+runSensorM :: SensorM a -> IO a
+runSensorM m = do
+  g <- newMVar M.empty
+  runReaderT m g
+
+poll :: (Typeable a) => Sensor a -> SensorM (TVar a)
+poll (Sensor id vM') = do
+  gM <- ask
+  g <- takeMVar gM
+  vM <- case M.lookup id g of
+    Just vM -> do
+      putMVar gM g
+      pure (fromJust . fromDynamic <$> vM)
+    Nothing -> do
+      -- liftIO (printf "[%s] start\n" id)
+      vM <- vM'
+      putMVar gM (M.insert id (toDyn <$> vM) g)
+      pure vM
+  runMemoized vM
+
+data Sensor a = Sensor String (SensorM (Memoized (TVar a)))
+
+data CpuStat = CpuStat
+  { total :: Int,
+    idle :: Int
+  }
+  deriving (Show, Eq)
+
+cpuStat :: Sensor CpuStat
+cpuStat = pollFile "cpuStat" aggregate "/proc/stat" (5 * 10 ^ 5)
+  where
+    aggregate fp = liftIO do
+      stat <- readFile fp
+      case filter ("cpu " `isPrefixOf`) (lines stat) of
+        [] -> error "/proc/stat: no cpu line"
+        (cpu : _) ->
+          case map read (drop 1 (words cpu)) of
+            xs@(_ : _ : _ : idle : _) ->
+              pure $ CpuStat (sum xs) idle
+            _ -> error "/proc/stat: unexpected cpu line"
+
+cpu' :: Sensor [CpuStat]
+cpu' = histogram "cpu'" [] (\b a -> take 7 (a : b)) cpuStat
+
+data Cpu = Cpu {unCpu :: [Float]}
+
+instance P.Pretty Cpu where
+  pretty = P.diagram 3 . unCpu
+
+cpu :: Sensor Cpu
+cpu = transform "cpu" f cpu'
+  where
+    f xs = do
+      let xs' =
+            zipWith
+              ( \next prev ->
+                  let total = next.total - prev.total
+                      idle = next.idle - prev.idle
+                   in fromIntegral (total - idle) / fromIntegral total
+              )
+              (fromMaybe [] (initMay xs))
+              (fromMaybe [] (tailMay xs))
+      pure (Cpu xs')
+
+data IoStat = IoStat {unIoStat :: Int} deriving (Show, Eq)
+
+ioStat :: Sensor IoStat
+ioStat = pollFile "ioStat" aggregate "/sys/block/nvme0n1/stat" (5 * 10 ^ 5)
+  where
+    aggregate fp = liftIO do
+      stat <- readFile fp
+      case words stat of
+        [_, _, read', _, _, _, write', _, _, _, _, _, _, _, _, _, _] ->
+          pure (IoStat ((read read') + (read write')))
+        _ -> error "/sys/block/nvme0n1/stat: malformed"
+
+io' :: Sensor [IoStat]
+io' = histogram "io'" [] (\b a -> take 7 (a : b)) ioStat
+
+data Io = Io {unIo :: [Float]} deriving (Show, Eq)
+
+instance P.Pretty Io where
+  pretty = P.diagram 3 . unIo
+
+io :: Sensor Io
+io = transform "io" f io'
+  where
+    f (map unIoStat -> xs) = do
+      let xs' =
+            zipWith
+              (-)
+              (fromMaybe [] (initMay xs))
+              (fromMaybe [] (tailMay xs))
+          x' = max 1 (fromMaybe 0 (maximumMay xs'))
+      pure . Io $
+        map (\x -> fromIntegral x / fromIntegral x') xs'
+
+data NetStat = NetStat {unNetStat :: Int} deriving (Show, Eq)
+
+netStat :: Sensor NetStat
+netStat =
+  pollFile
+    "netStat"
+    aggregate
+    "/sys/class/net/wlan0/statistics/rx_bytes"
+    (5 * 10 ^ 5)
+  where
+    aggregate fp = liftIO do
+      NetStat
+        <$> ( (+)
+                <$> (read <$> readFile fp)
+                <*> (read <$> readFile (takeDirectory fp </> "tx_bytes"))
+            )
+
+net' :: Sensor [NetStat]
+net' = histogram "net'" [] (\b a -> take 7 (a : b)) netStat
+
+data Net = Net {unNet :: [Float]} deriving (Show, Eq)
+
+instance P.Pretty Net where
+  pretty = P.diagram 3 . unNet
+
+net :: Sensor Net
+net = transform "net" f net'
+  where
+    f (map unNetStat -> xs) = do
+      let xs' =
+            zipWith
+              (-)
+              (fromMaybe [] (initMay xs))
+              (fromMaybe [] (tailMay xs))
+          x' = max 1 (fromMaybe 0 (maximumMay xs'))
+      pure $ Net $ map (\x -> fromIntegral x / fromIntegral x') xs'
+
+data MemStat = MemStat {unMemStat :: Float} deriving (Show, Eq)
+
+memStat :: Sensor MemStat
+memStat = pollFile "memStat" aggregate "/proc/meminfo" (5 * 10 ^ 5)
+  where
+    aggregate fp = liftIO do
+      meminfo <- readFile fp
+      case foldl
+        ( \(total, avail) xs ->
+            case xs of
+              ["MemTotal:", v, "kB"] -> (Just (read v), avail)
+              ["MemAvailable:", v, "kB"] -> (total, Just (read v))
+              _ -> (total, avail)
+        )
+        (Nothing, Nothing)
+        (map words (lines meminfo)) of
+        (Just total, Just avail) -> pure $ MemStat (1 - avail / total)
+        (Nothing, _) -> error (printf "%s: MemTotal missing" fp)
+        (_, Nothing) -> error (printf "%s: MemAvail missing" fp)
+
+mem' :: Sensor [MemStat]
+mem' = histogram "mem'" [] (\b a -> take 7 (a : b)) memStat
+
+data Mem = Mem {unMem :: [Float]} deriving (Show, Eq)
+
+instance P.Pretty Mem where
+  pretty = P.diagram 3 . unMem
+
+mem :: Sensor Mem
+mem = transform "mem" (\(map unMemStat -> xs) -> pure (Mem xs)) mem'
+
+data BatStat = BatStat {unBatStat :: Float} deriving (Show, Eq)
+
+batStat :: Sensor BatStat
+batStat = pollFile "batStat" aggregate "/sys/class/power_supply/macsmc-battery/charge_now" (5 * 10 ^ 5)
+  where
+    aggregate fp = liftIO do
+      BatStat
+        <$> ( (\now full -> fromIntegral now / fromIntegral full)
+                <$> (read <$> readFile fp)
+                <*> (read <$> readFile (takeDirectory fp </> "charge_full"))
+            )
+
+bat' :: Sensor [BatStat]
+bat' = histogram "bat'" [] (\b a -> take 7 (a : b)) batStat
+
+data Bat = Bat {unBat :: [Float]} deriving (Show, Eq)
+
+instance P.Pretty Bat where
+  pretty = P.diagram 1 . unBat
+
+bat :: Sensor Bat
+bat = transform "bat" (\(map unBatStat -> xs) -> pure (Bat xs)) bat'
+
+data Snd = Snd {unSnd :: Float} deriving (Show, Eq)
+
+instance P.Pretty Snd where
+  pretty = P.diagram 1 . (: []) . unSnd
+
+snd :: Sensor Snd
+snd = Sensor "snd" $ memoizeMVar do
+  stateT <- newTVarIO (Snd 1)
+  void $ forkIO $ forever do
+    x <- liftIO aggregate
+    atomically do writeTVar stateT (Snd x)
+    threadDelay (2 * 10 ^ 6)
+  pure stateT
+  where
+    aggregate = do
+      (/ 153) . read <$> [sh|pamixer --get-volume|]
+
+data Disk = Disk {unDisk :: [Float]} deriving (Show, Eq)
+
+instance P.Pretty Disk where
+  pretty = P.diagram 1 . unDisk
+
+disk :: Sensor Disk
+disk = Sensor "disk" $ memoizeMVar do
+  stateT <- newTVarIO (Disk [])
+  void $ forkIO $ forever do
+    x <- liftIO aggregate
+    atomically do
+      readTVar stateT >>= \(Disk xs) ->
+        writeTVar stateT (Disk (x : xs))
+    threadDelay (2 * 10 ^ 6)
+  pure stateT
+  where
+    aggregate = do
+      stat <- statVFS "/"
+      pure $
+        (fromIntegral stat.statVFS_bfree)
+          / fromIntegral (stat.statVFS_bfree + stat.statVFS_bavail)
+
+currentTime :: Sensor UTCTime
+currentTime = Sensor "currentTime" $ memoizeMVar do
+  stateT <- newTVarIO =<< liftIO getCurrentTime
+  void $ forkIO $ forever do
+    x <- liftIO getCurrentTime
+    atomically . writeTVar stateT =<< liftIO getCurrentTime
+    threadDelay (((10 ^ 12) - (fromEnum (utctDayTime x) `mod` (10 ^ 12))) `div` (10 ^ 6))
+  pure stateT
+
+data Date = Date {unDate :: String} deriving (Show, Eq)
+
+instance P.Pretty Date where
+  pretty = P.color P.White . P.pretty . unDate
+
+date :: Sensor Date
+date = transform "date" f currentTime
+  where
+    f x = do
+      z <- liftIO getCurrentTimeZone
+      pure . Date . formatTime defaultTimeLocale "%b %e" $
+        utcToLocalTime z x
+
+data Time = Time {unTime :: String} deriving (Show, Eq)
+
+instance P.Pretty Time where
+  pretty = P.color P.White . P.pretty . unTime
+
+time :: Sensor Time
+time = transform "time" f currentTime
+  where
+    f x = do
+      z <- liftIO getCurrentTimeZone
+      pure . Time . formatTime defaultTimeLocale "%R" $
+        utcToLocalTime z x
+
+data WmName = WmName {unWmName :: String} deriving (Show, Eq)
+
+instance P.Pretty WmName where
+  pretty = P.color P.White . P.pretty . unWmName
+
+wmName' :: Sensor WmName
+wmName' = watchXPropertyChanges "wmName'" as aggregate
+  where
+    as = ["_NET_ACTIVE_WINDOW"]
+    aggregate dpy = liftIO do
+      let root = X.rootWindowOfScreen (X.defaultScreenOfDisplay dpy)
+      netActiveWindow <- X.internAtom dpy "_NET_ACTIVE_WINDOW" False
+      X.getWindowProperty32 dpy netActiveWindow root >>= \case
+        Just (win : _) -> do
+          netWmName <- X.internAtom dpy "_NET_WM_NAME" False
+          fmap (WmName . strip . head) . X.wcTextPropertyToTextList dpy
+            =<< do
+              let loop 0 = do
+                    error "failed to read _NET_WM_NAME"
+                  loop n = do
+                    catch
+                      (X.getTextProperty dpy (fromIntegral win) netWmName)
+                      ( \(e :: IOError) -> do
+                          if isUserError e
+                            then threadDelay 1000 >> loop (n - 1)
+                            else throwIO e
+                      )
+              loop 16
+        _ -> pure (WmName "")
+    strip = reverse . dropWhile isSpace . reverse . dropWhile isSpace
+
+wmName :: Sensor String
+wmName = transform "wmName" (\(WmName x) -> pure x) wmName'
+
+data WmWorkspaces = WmWorkspaces {unWmWorkspaces :: [Workspace]}
+  deriving (Show, Eq)
+
+data Workspace
+  = Active String
+  | Inactive String
+  deriving (Show, Eq)
+
+instance P.Pretty WmWorkspaces where
+  pretty =
+    P.Col . intersperse (P.pretty " ") . map P.pretty . unWmWorkspaces
+
+instance P.Pretty Workspace where
+  pretty (Active s) = P.color P.White (P.pretty s)
+  pretty (Inactive s) = P.colorDull P.White (P.pretty s)
+
+wmWorkspaces :: Sensor WmWorkspaces
+wmWorkspaces = watchXPropertyChanges "wmWorkspaces" as aggregate
+  where
+    as =
+      [ "_NET_CLIENT_LIST",
+        "_NET_CURRENT_DESKTOP",
+        "_NET_DESKTOP_NAMES",
+        "_NET_WM_DESKTOP"
+      ]
+    aggregate dpy = liftIO do
+      let root = X.rootWindowOfScreen (X.defaultScreenOfDisplay dpy)
+      netClientList <- liftIO $ X.internAtom dpy "_NET_CLIENT_LIST" False
+      netCurrentDesktop <- liftIO $ X.internAtom dpy "_NET_CURRENT_DESKTOP" False
+      netDesktopNames <- liftIO $ X.internAtom dpy "_NET_DESKTOP_NAMES" False
+      netWmDesktop <- liftIO $ X.internAtom dpy "_NET_WM_DESKTOP" False
+      currentDesktop <- fmap (fromIntegral . head) <$> X.getWindowProperty32 dpy netCurrentDesktop root
+      occupiedDesktops <-
+        fmap (S.unions . catMaybes)
+          . mapM
+            ( \win -> do
+                fmap (S.singleton . fromIntegral . head)
+                  <$> X.getWindowProperty32 dpy netWmDesktop win
+            )
+          . map fromIntegral
+          . fromMaybe []
+          =<< X.getWindowProperty32 dpy netClientList root
+      fmap
+        ( WmWorkspaces
+            . ifilter (\i _ -> Just i == currentDesktop || i `S.member` occupiedDesktops)
+            . imap
+              ( \i ->
+                  if Just i == currentDesktop
+                    then Active
+                    else Inactive
+              )
+        )
+        . X.wcTextPropertyToTextList dpy
+        =<< X.getTextProperty dpy root netDesktopNames
+
+watchFile :: String -> (FilePath -> SensorM a) -> FilePath -> Sensor a
+watchFile id aggregate fp = Sensor id $ memoizeMVar do
+  i <- liftIO do
+    i <- I.init
+    void (I.addWatch i fp I.in_CLOSE)
+    pure i
+  stateT <- newTVarIO =<< aggregate fp
+  void $ forkIO $ forever do
+    void (liftIO (I.getEvent i))
+    -- liftIO (printf "[%s] aggregate\n" id)
+    x <- aggregate fp
+    atomically (writeTVar stateT x)
+  pure stateT
+
+pollFile :: String -> (FilePath -> SensorM a) -> FilePath -> Int -> Sensor a
+pollFile id aggregate fp delay = Sensor id $ memoizeMVar do
+  stateT <- newTVarIO =<< aggregate fp
+  void $ forkIO $ forever do
+    -- liftIO (printf "[%s] aggregate\n" id)
+    x <- aggregate fp
+    threadDelay delay
+    atomically (writeTVar stateT x)
+  pure stateT
+
+watchXPropertyChanges ::
+  String ->
+  [String] ->
+  (X.Display -> SensorM a) ->
+  Sensor a
+watchXPropertyChanges id as' aggregate = Sensor id $ memoizeMVar do
+  (dpy, as) <- liftIO do
+    dpy <- X.openDisplay ""
+    let root = X.rootWindowOfScreen (X.defaultScreenOfDisplay dpy)
+    X.selectInput dpy root X.propertyChangeMask
+    as <- S.fromList <$> mapM (\a -> X.internAtom dpy a False) as'
+    pure (dpy, as)
+  stateT <- newTVarIO =<< aggregate dpy
+  void $ forkIO $ forever $ do
+    as' <- liftIO $ X.allocaXEvent $ \ev -> do
+      X.nextEvent dpy ev
+      e <- X.getEvent ev
+      let loop es = do
+            isTimeout <- X.waitForEvent dpy 0
+            if isTimeout
+              then pure (S.fromList (map (X.ev_atom) es))
+              else do
+                X.nextEvent dpy ev
+                e <- X.getEvent ev
+                loop (e : es)
+      loop [e]
+    unless (S.null (as `S.intersection` as')) do
+      atomically . writeTVar stateT =<< aggregate dpy
+  pure stateT
+
+histogram ::
+  (Eq a, Typeable a) =>
+  String ->
+  b ->
+  (b -> a -> b) ->
+  Sensor a ->
+  Sensor b
+histogram id def acc sensor = transform' id def f sensor
+  where
+    f b a = pure (acc b a)
+
+transform' ::
+  (Eq a, Typeable a) =>
+  String ->
+  b ->
+  (b -> a -> SensorM b) ->
+  Sensor a ->
+  Sensor b
+transform' id def fM sensor = Sensor id $ memoizeMVar do
+  xT <- poll sensor
+  x <- readTVarIO xT
+  x'T <- newTVarIO x
+  stateT <- newTVarIO =<< fM def x
+  void $ forkIO $ forever $ do
+    (def, x) <- atomically do
+      x' <- readTVar x'T
+      x <- readTVar xT
+      checkSTM (x /= x')
+      def <- readTVar stateT
+      writeTVar x'T x
+      pure (def, x)
+    atomically . writeTVar stateT =<< fM def x
+  pure stateT
+
+transform ::
+  (Eq a, Typeable a) =>
+  String ->
+  (a -> SensorM b) ->
+  Sensor a ->
+  Sensor b
+transform id fM sensor = Sensor id $ memoizeMVar do
+  xT <- poll sensor
+  x <- readTVarIO xT
+  x'T <- newTVarIO x
+  stateT <- newTVarIO =<< fM x
+  void $ forkIO $ forever $ do
+    x <- atomically do
+      x' <- readTVar x'T
+      x <- readTVar xT
+      checkSTM (x /= x')
+      writeTVar x'T x
+      pure x
+    atomically . writeTVar stateT =<< fM x
+  pure stateT