diff --git a/src/GitmonClient.hs b/src/GitmonClient.hs index 84535fc40..29f077b93 100644 --- a/src/GitmonClient.hs +++ b/src/GitmonClient.hs @@ -59,30 +59,29 @@ type ProcInfo = Either Y.ParseException (Maybe ProcIO) newtype SocketFactory = SocketFactory { withSocket :: forall a. (Socket -> IO a) -> IO a } -reportGitmon''' :: String -> ReaderT LgRepo IO a -> ReaderT LgRepo IO a -reportGitmon''' = reportGitmon'' SocketFactory { withSocket = withGitmonSocket } +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 +reportGitmon' :: SocketFactory -> String -> ReaderT LgRepo IO a -> ReaderT LgRepo IO a +reportGitmon' SocketFactory{..} program gitCommand = do (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) + 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" + !result <- 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)) + safeIO . withSocket . flip sendAll $ processJSON Finish (ProcessFinishData cpuTime diskReadBytes diskWriteBytes resultCode) - pure r + pure result where collectStats :: IO (TimeSpec, ProcInfo) @@ -120,67 +119,6 @@ withGitmonSocket = bracket connectSocket close connect s (SockAddrUnix gitmonSocketAddr) pure s -reportGitmon :: String -> ReaderT LgRepo IO a -> ReaderT LgRepo IO a -reportGitmon program gitCommand = do - maybeSoc <- safeIO $ socket AF_UNIX Stream defaultProtocol - case maybeSoc of - Nothing -> gitCommand - 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) - -reportGitmon' :: Socket -> String -> ReaderT LgRepo IO a -> ReaderT LgRepo IO a -reportGitmon' soc program gitCommand = do - (gitDir, realIP, repoName, userID) <- liftIO loadEnvVars - safeIO $ sendAll soc (processJSON Update (ProcessUpdateData gitDir program realIP repoName userID "semantic-diff")) - safeIO $ sendAll soc (processJSON Schedule ProcessScheduleData) - shouldContinue error $ do - (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 (ProcessFinishData cpuTime diskReadBytes diskWriteBytes resultCode)) - pure result - - where collectStats :: IO (TimeSpec, ProcInfo) - collectStats = do - time <- getTime clock - procIOContents <- Y.decodeFileEither procFileAddr :: IO ProcInfo - pure (time, procIOContents) - - shouldContinue :: MonadIO m => (String -> m b) -> m b -> m b - shouldContinue err action = do - maybeCommand <- safeIO $ timeout gitmonTimeout (safeIO $ recv soc 1024) - case (join . join) maybeCommand of - 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 ) - 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) - -- Timeout in nanoseconds to wait before giving up on Gitmon response to schedule. gitmonTimeout :: Int gitmonTimeout = 1 * 1000 * 1000 diff --git a/test/GitmonClientSpec.hs b/test/GitmonClientSpec.hs index 1f7cef950..ea83ed122 100644 --- a/test/GitmonClientSpec.hs +++ b/test/GitmonClientSpec.hs @@ -23,11 +23,11 @@ spec :: Spec spec = parallel $ do describe "gitmon" $ do let wd = "test/fixtures/git/examples/all-languages.git" - it "receives commands in order" . withSocketPair $ \(client, server) -> + it "receives commands in order" . withSocketPair $ \(_, server, socketFactory) -> withRepository lgFactory wd $ do liftIO $ sendAll server "continue" object <- parseObjOid (pack "dfac8fd681b0749af137aebf3203e77a06fbafc2") - commit <- reportGitmon' client "cat-file" $ lookupCommit object + commit <- reportGitmon' socketFactory "cat-file" $ lookupCommit object info <- liftIO $ recv server 1024 let [update, schedule, finish] = infoToCommands info @@ -37,7 +37,7 @@ spec = parallel $ do liftIO $ shouldBe schedule (Just "schedule") liftIO $ shouldBe finish (Just "finish") - it "receives update command with correct data" . withSocketPair $ \(client, server) -> + it "receives update command with correct data" . withSocketPair $ \(_, server, socketFactory) -> withRepository lgFactory wd $ do liftIO $ setEnv "GIT_DIR" wd liftIO $ setEnv "GIT_SOCKSTAT_VAR_real_ip" "127.0.0.1" @@ -46,7 +46,7 @@ spec = parallel $ do liftIO $ sendAll server "continue" object <- parseObjOid (pack "dfac8fd681b0749af137aebf3203e77a06fbafc2") - commit <- reportGitmon' client "cat-file" $ lookupCommit object + commit <- reportGitmon' socketFactory "cat-file" $ lookupCommit object info <- liftIO $ recv server 1024 let [updateData, _, finishData] = infoToData info @@ -64,32 +64,34 @@ spec = parallel $ do liftIO $ shouldSatisfy (either (const (-1)) diskWriteBytes finishData) (>= 0) liftIO $ shouldSatisfy (either (const (-1)) resultCode finishData) (>= 0) - it "returns the correct git result if the socket is unavailable" . withSocketPair $ \(client, server) -> + it "returns the correct git result if the socket is unavailable" . withSocketPair $ \(client, server, socketFactory) -> withRepository lgFactory wd $ do liftIO $ close client object <- parseObjOid (pack "dfac8fd681b0749af137aebf3203e77a06fbafc2") - commit <- reportGitmon' client "cat-file" $ lookupCommit object + commit <- reportGitmon' socketFactory "cat-file" $ lookupCommit object info <- liftIO $ recv server 1024 liftIO $ shouldBe (commitOid commit) object liftIO $ shouldBe "" info - it "throws if schedule response is fail" . withSocketPair $ \(client, server) -> + it "throws if schedule response is fail" . withSocketPair $ \(_, server, socketFactory) -> withRepository lgFactory wd $ do repo <- getRepository liftIO $ sendAll server "fail too busy" object <- parseObjOid (pack "dfac8fd681b0749af137aebf3203e77a06fbafc2") - liftIO $ shouldThrow (runReaderT (reportGitmon' client "cat-file" (lookupCommit object)) repo) anyErrorCall + liftIO $ shouldThrow (runReaderT (reportGitmon' socketFactory "cat-file" (lookupCommit object)) repo) anyErrorCall -withSocketPair :: ((Socket, Socket) -> IO c) -> IO c -withSocketPair = - bracket - (socketPair AF_UNIX Stream defaultProtocol) - (\(client, server) -> do - close client - close server) +withSocketPair :: ((Socket, Socket, SocketFactory) -> IO c) -> IO c +withSocketPair = bracket create release + where + create = do + (client, server) <- socketPair AF_UNIX Stream defaultProtocol + pure (client, server, SocketFactory (\f -> f client)) + release (client, server, _) = do + close client + close server infoToCommands :: ByteString -> [Maybe Text] infoToCommands input = command' . toObject <$> Prelude.take 3 (split '\n' input)