semantic/src/GitmonClient.hs

{-# LANGUAGE RecordWildCards, BangPatterns, DeriveGeneric #-}
module GitmonClient where

import Data.Aeson
import Data.Aeson.Types
import Data.ByteString.Lazy (toStrict)
import Data.Text (pack, unpack, toLower)
import qualified Data.Yaml as Y
import GHC.Generics
import Git.Libgit2
import Network.Socket
import Network.Socket.ByteString (sendAll)
import Prelude
import Prologue hiding (toStrict)
import System.Clock
import System.Directory (getCurrentDirectory)
import System.Environment
import System.IO

data ProcIO = ProcIO {
    read_bytes :: Integer
  , write_bytes :: Integer
} deriving (Show, Generic)

instance FromJSON ProcIO


data ProcessStats =
    ProcessBeforeStats { gitDir :: String
                       , program :: String
                       , realIP :: Maybe String
                       , repoID :: Maybe String
                       , repoName :: Maybe String
                       , userID :: Maybe String
                       , via :: String }
  | ProcessAfterStats { cpu :: Integer
                      , diskReadBytes :: Integer
                      , diskWriteBytes :: Integer
                      , resultCode :: Integer } deriving (Generic, Show)

instance ToJSON ProcessStats where
  toJSON = genericToJSON defaultOptions { fieldLabelModifier = camelTo2 '_' }



data GitmonCommand = Update
                   | Finish
                   | Schedule deriving (Generic, Show)

instance ToJSON GitmonCommand where
  toJSON = genericToJSON defaultOptions { constructorTagModifier = unpack . toLower . pack }


data GitmonMsg = GitmonMsg { command :: GitmonCommand, stats :: ProcessStats } deriving (Show)

instance ToJSON GitmonMsg where
  toJSON GitmonMsg{..} = object [
    "command" .= command,
    "data" .= stats
    ]


gitmonSocketAddr :: String
gitmonSocketAddr = "/tmp/gitstats.sock"

procFileAddr :: String
procFileAddr = "/proc/self/io"

clock :: Clock
clock = Realtime

processJSON :: GitmonCommand -> ProcessStats -> ByteString
processJSON command stats = toStrict . encode $ GitmonMsg command stats

type ProcInfo = Either Y.ParseException (Maybe ProcIO)

reportGitmon :: String -> Arguments -> ReaderT LgRepo IO a -> ReaderT LgRepo IO a
reportGitmon program Arguments{..} gitCommand = do
  soc <- liftIO $ socket AF_UNIX Stream defaultProtocol
  safeIO $ connect soc (SockAddrUnix gitmonSocketAddr)

  safeIO $ sendAll soc (processJSON Update ProcessBeforeStats { gitDir = gitDir, via = "semantic-diff", program = program, realIP = realIP, repoID = repoID, repoName = repoName, userID = userID })

  (startTime, beforeProcIOContents) <- liftIO collectStats

  !result <- gitCommand

  (afterTime, afterProcIOContents) <- liftIO collectStats

  let (cpuTime, diskReadBytes', diskWriteBytes', resultCode') = procStats startTime afterTime beforeProcIOContents afterProcIOContents

  safeIO $ sendAll soc (processJSON Finish ProcessAfterStats { cpu = cpuTime, diskReadBytes = diskReadBytes', diskWriteBytes = diskWriteBytes', resultCode = resultCode' })

  safeIO $ close soc

  return result

  where safeIO :: MonadIO m => IO () -> m ()
        safeIO command = liftIO $ command `catch` noop

        noop :: IOException -> IO ()
        noop _ = return ()

        collectStats :: IO (TimeSpec, ProcInfo)
        collectStats = do
          time <- getTime clock
          procIOContents <- Y.decodeFileEither procFileAddr :: IO ProcInfo
          return (time, procIOContents)

        procStats :: TimeSpec -> TimeSpec -> ProcInfo -> ProcInfo -> ( Integer, Integer, Integer, Integer )
        procStats beforeTime afterTime beforeProcIOContents afterProcIOContents = ( cpuTime, diskReadBytes, diskWriteBytes, resultCode )
          where
            cpuTime = toNanoSecs afterTime - toNanoSecs beforeTime
            beforeDiskReadBytes = either (const 0) (maybe 0 read_bytes) beforeProcIOContents
            afterDiskReadBytes = either (const 0) (maybe 0 read_bytes) afterProcIOContents
            beforeDiskWriteBytes = either (const 0) (maybe 0 write_bytes) beforeProcIOContents
            afterDiskWriteBytes = either (const 0) (maybe 0 write_bytes) afterProcIOContents
            diskReadBytes = afterDiskReadBytes - beforeDiskReadBytes
            diskWriteBytes = afterDiskWriteBytes - beforeDiskWriteBytes
            resultCode = 0