1
1
mirror of https://github.com/github/semantic.git synced 2024-12-22 22:31:36 +03:00
semantic/src/GitmonClient.hs

205 lines
8.6 KiB
Haskell
Raw Normal View History

{-# LANGUAGE RecordWildCards, BangPatterns, DeriveGeneric, RankNTypes #-}
2017-03-02 21:27:16 +03:00
{-# OPTIONS_GHC -fno-warn-unused-do-bind #-}
2017-02-15 04:42:35 +03:00
module GitmonClient where
import Data.Aeson
import Data.Aeson.Types
2017-02-24 01:03:26 +03:00
import Data.ByteString.Lazy (toStrict)
2017-03-03 00:16:38 +03:00
import Data.Text (pack, unpack, toLower, isInfixOf)
2017-03-01 02:27:33 +03:00
import qualified Data.Yaml as Y
2017-02-24 01:03:26 +03:00
import GHC.Generics
2017-02-22 22:15:58 +03:00
import Git.Libgit2
2017-03-03 00:16:38 +03:00
import Network.Socket hiding (recv)
import Network.Socket.ByteString (sendAll, recv)
2017-02-24 01:03:26 +03:00
import Prelude
2017-03-03 00:16:38 +03:00
import Prologue hiding (toStrict, error)
import System.Clock
2017-03-01 02:27:33 +03:00
import System.Directory (getCurrentDirectory)
import System.Environment
import System.Timeout
2017-02-15 04:42:35 +03:00
2017-03-03 00:35:51 +03:00
data ProcIO = ProcIO { read_bytes :: Integer
, write_bytes :: Integer } deriving (Show, Generic)
instance FromJSON ProcIO
2017-03-03 00:35:51 +03:00
data ProcessData = ProcessUpdateData { gitDir :: String
, program :: String
, realIP :: Maybe String
, repoName :: Maybe String
, userID :: Maybe String
, via :: String }
| ProcessScheduleData
| ProcessFinishData { cpu :: Integer
, diskReadBytes :: Integer
, diskWriteBytes :: Integer
, resultCode :: Integer } deriving (Generic, Show)
2017-02-23 00:34:28 +03:00
2017-03-01 02:46:29 +03:00
instance ToJSON ProcessData where
2017-02-15 04:42:35 +03:00
toJSON = genericToJSON defaultOptions { fieldLabelModifier = camelTo2 '_' }
2017-02-24 01:04:24 +03:00
2017-02-15 04:42:35 +03:00
data GitmonCommand = Update
| Finish
| Schedule deriving (Generic, Show)
2017-02-15 04:42:35 +03:00
instance ToJSON GitmonCommand where
toJSON = genericToJSON defaultOptions { constructorTagModifier = unpack . toLower . pack }
2017-02-15 04:42:35 +03:00
2017-02-24 01:04:24 +03:00
2017-03-03 00:35:51 +03:00
data GitmonMsg = GitmonMsg { command :: GitmonCommand
, processData :: ProcessData } deriving (Show)
2017-02-15 04:42:35 +03:00
instance ToJSON GitmonMsg where
toJSON GitmonMsg{..} = object ["command" .= command, "data" .= processData]
2017-03-02 01:06:58 +03:00
2017-03-03 01:05:50 +03:00
type ProcInfo = Either Y.ParseException (Maybe ProcIO)
newtype SocketFactory = SocketFactory { withSocket :: forall a. (Socket -> IO a) -> IO a }
2017-03-07 00:51:38 +03:00
reportGitmon''' :: String -> ReaderT LgRepo IO a -> ReaderT LgRepo IO a
reportGitmon''' = reportGitmon'' SocketFactory { withSocket = withGitmonSocket }
reportGitmon'' :: SocketFactory -> String -> ReaderT LgRepo IO a -> ReaderT LgRepo IO a
reportGitmon'' SocketFactory{..} program gitCommand = do
2017-03-07 00:51:38 +03:00
(gitDir, realIP, repoName, userID) <- liftIO loadEnvVars
(startTime, beforeProcIOContents) <- liftIO collectStats
maybeCommand <- safeIO . timeout gitmonTimeout . withSocket $ \s -> do
sendAll s (processJSON Update (ProcessUpdateData gitDir program realIP repoName userID "semantic-diff"))
sendAll s (processJSON Schedule ProcessScheduleData)
recv s 1024
!r <- case join maybeCommand of
Just command | "fail" `isInfixOf` decodeUtf8 command -> error . unpack $ "Received '" <> decodeUtf8 command <> "' from Gitmon"
_ -> gitCommand
(afterTime, afterProcIOContents) <- liftIO collectStats
let (cpuTime, diskReadBytes, diskWriteBytes, resultCode) = procStats startTime afterTime beforeProcIOContents afterProcIOContents
safeIO . withSocket $ \s ->
sendAll s (processJSON Finish (ProcessFinishData cpuTime diskReadBytes diskWriteBytes resultCode))
pure r
where
collectStats :: IO (TimeSpec, ProcInfo)
collectStats = do
time <- getTime clock
procIOContents <- Y.decodeFileEither procFileAddr :: IO ProcInfo
pure (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
loadEnvVars :: IO (String, Maybe String, Maybe String, Maybe String)
loadEnvVars = do
pwd <- getCurrentDirectory
gitDir <- fromMaybe pwd <$> lookupEnv "GIT_DIR"
realIP <- lookupEnv "GIT_SOCKSTAT_VAR_real_ip"
repoName <- lookupEnv "GIT_SOCKSTAT_VAR_repo_name"
userID <- lookupEnv "GIT_SOCKSTAT_VAR_user_id"
pure (gitDir, realIP, repoName, userID)
withGitmonSocket :: (Socket -> IO c) -> IO c
withGitmonSocket = bracket connectSocket close
2017-03-07 00:51:38 +03:00
where
connectSocket = do
s <- socket AF_UNIX Stream defaultProtocol
connect s (SockAddrUnix gitmonSocketAddr)
pure s
2017-03-03 01:05:50 +03:00
reportGitmon :: String -> ReaderT LgRepo IO a -> ReaderT LgRepo IO a
reportGitmon program gitCommand = do
2017-03-02 21:27:16 +03:00
maybeSoc <- safeIO $ socket AF_UNIX Stream defaultProtocol
case maybeSoc of
Nothing -> gitCommand
2017-03-02 22:56:33 +03:00
Just soc -> do
safeIO $ connect soc (SockAddrUnix gitmonSocketAddr)
result <- reportGitmon' soc program gitCommand
safeIO $ close soc
pure result
`catchError` (\e -> do
safeIO $ close soc
throwIO e)
2017-03-02 01:06:58 +03:00
reportGitmon' :: Socket -> String -> ReaderT LgRepo IO a -> ReaderT LgRepo IO a
reportGitmon' soc program gitCommand = do
2017-03-06 22:09:24 +03:00
(gitDir, realIP, repoName, userID) <- liftIO loadEnvVars
safeIO $ sendAll soc (processJSON Update (ProcessUpdateData gitDir program realIP repoName userID "semantic-diff"))
2017-03-01 02:46:29 +03:00
safeIO $ sendAll soc (processJSON Schedule ProcessScheduleData)
2017-03-03 00:59:13 +03:00
shouldContinue error $ do
(startTime, beforeProcIOContents) <- liftIO collectStats
!result <- gitCommand
(afterTime, afterProcIOContents) <- liftIO collectStats
2017-03-03 01:05:50 +03:00
let (cpuTime, diskReadBytes, diskWriteBytes, resultCode) = procStats startTime afterTime beforeProcIOContents afterProcIOContents
safeIO $ sendAll soc (processJSON Finish (ProcessFinishData cpuTime diskReadBytes diskWriteBytes resultCode))
pure result
2017-02-22 23:54:43 +03:00
2017-03-02 01:06:58 +03:00
where collectStats :: IO (TimeSpec, ProcInfo)
collectStats = do
time <- getTime clock
procIOContents <- Y.decodeFileEither procFileAddr :: IO ProcInfo
pure (time, procIOContents)
2017-03-03 00:59:13 +03:00
shouldContinue :: MonadIO m => (String -> m b) -> m b -> m b
shouldContinue err action = do
2017-03-03 00:35:51 +03:00
maybeCommand <- safeIO $ timeout gitmonTimeout (safeIO $ recv soc 1024)
case (join . join) maybeCommand of
2017-03-03 01:01:13 +03:00
Just command | "fail" `isInfixOf` decodeUtf8 command -> err . unpack $ "Received '" <> decodeUtf8 command <> "' from Gitmon"
_ -> action
procStats :: TimeSpec -> TimeSpec -> ProcInfo -> ProcInfo -> ( Integer, Integer, Integer, Integer )
procStats beforeTime afterTime beforeProcIOContents afterProcIOContents = ( cpuTime, diskReadBytes, diskWriteBytes, resultCode )
2017-02-24 03:56:02 +03:00
where
cpuTime = toNanoSecs afterTime - toNanoSecs beforeTime
2017-02-24 03:56:02 +03:00
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
2017-03-06 22:09:24 +03:00
loadEnvVars :: IO (String, Maybe String, Maybe String, Maybe String)
2017-03-01 02:30:42 +03:00
loadEnvVars = do
pwd <- getCurrentDirectory
gitDir <- fromMaybe pwd <$> lookupEnv "GIT_DIR"
realIP <- lookupEnv "GIT_SOCKSTAT_VAR_real_ip"
repoName <- lookupEnv "GIT_SOCKSTAT_VAR_repo_name"
userID <- lookupEnv "GIT_SOCKSTAT_VAR_user_id"
2017-03-06 22:09:24 +03:00
pure (gitDir, realIP, repoName, userID)
2017-03-03 00:35:51 +03:00
-- Timeout in nanoseconds to wait before giving up on Gitmon response to schedule.
gitmonTimeout :: Int
gitmonTimeout = 1 * 1000 * 1000
gitmonSocketAddr :: String
gitmonSocketAddr = "/tmp/gitstats.sock"
procFileAddr :: String
procFileAddr = "/proc/self/io"
clock :: Clock
clock = Realtime
processJSON :: GitmonCommand -> ProcessData -> ByteString
processJSON command processData = (toStrict . encode $ GitmonMsg command processData) <> "\n"
safeIO :: MonadIO m => IO a -> m (Maybe a)
safeIO command = liftIO $ (Just <$> command) `catch` noop
noop :: IOException -> IO (Maybe a)
noop _ = pure Nothing