Utils.hs 8.99 KB
Newer Older
1
{-# language LambdaCase #-}
Valentin Reis's avatar
Valentin Reis committed
2
{-# language TupleSections #-}
3 4 5 6 7
{-# language OverloadedStrings #-}
{-# language DataKinds #-}
{-# language FlexibleInstances #-}
{-# language ScopedTypeVariables #-}
{-# language TypeOperators #-}
Valentin Reis's avatar
Valentin Reis committed
8

Valentin Reis's avatar
Valentin Reis committed
9 10 11 12 13 14 15 16
{-|
Module      : Argo.Utils
Description : Argo stack library
Copyright   : (c) Valentin Reis, 2018
License     : MIT
Maintainer  : fre@freux.fr
-}

17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32
module Argo.Utils
  ( printInfo
  , printWarning
  , printSuccess
  , printError
  , printTest
  , MonitoringResult(..)
  , Instrumentation(..)
  , TracebackScan(..)
  , processBehaviorToI
  , kbInstallHandler
  , runI
  , cleanSocket
  , myWhich
  )
where
Valentin Reis's avatar
Valentin Reis committed
33

Valentin Reis's avatar
Valentin Reis committed
34
import           Protolude
35
import           Argo.Types
Valentin Reis's avatar
Valentin Reis committed
36
import           System.Exit                    ( ExitCode(..) )
37 38 39
import           Data.ByteString               as B
                                         hiding ( empty )
import           Data.Conduit
Valentin Reis's avatar
Valentin Reis committed
40
import           Data.Conduit.Combinators      as CC
Valentin Reis's avatar
Valentin Reis committed
41 42
import           System.Console.ANSI
import           System.Console.ANSI.Types      ( Color )
43 44 45 46
import           System.Posix.Signals           ( installHandler
                                                , keyboardSignal
                                                , Handler(..)
                                                )
Valentin Reis's avatar
Valentin Reis committed
47
import           System.Directory
Valentin Reis's avatar
Valentin Reis committed
48

49 50 51
import           Control.Exception.Base         ( Exception
                                                , try
                                                , throw
Valentin Reis's avatar
Valentin Reis committed
52 53 54 55 56
                                                )
import           Control.Monad.IO.Unlift        ( MonadIO(..)
                                                , withRunInIO
                                                )
import           Data.Text                     as T
57 58 59 60 61 62 63 64 65 66 67
                                                ( unpack
                                                , Text
                                                )
import           Data.Text.Encoding            as TE
                                                ( decodeUtf8
                                                , encodeUtf8
                                                )
import           Data.Typeable                  ( Typeable )
import           System.IO                      ( BufferMode(NoBuffering)
                                                , hSetBuffering
                                                )
Valentin Reis's avatar
Valentin Reis committed
68
import qualified System.IO                     as IO
Valentin Reis's avatar
Valentin Reis committed
69 70
import           Data.Conduit.Process.Typed
import           Data.Conduit.Binary           as CB
Valentin Reis's avatar
Valentin Reis committed
71 72

-- | Miscellaneous printing utilities
Valentin Reis's avatar
Valentin Reis committed
73
colorShell :: Color -> IO () -> IO ()
Valentin Reis's avatar
Valentin Reis committed
74 75 76
colorShell color she = setC color *> she *> setC White
  where setC c = liftIO $ setSGR [SetColor Foreground Dull c]

Valentin Reis's avatar
Valentin Reis committed
77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94
printInfo :: Text -> IO ()
printError :: Text -> IO ()
printWarning :: Text -> IO ()
printSuccess :: Text -> IO ()
printTest :: Text -> IO ()

printInfo s = putText $ "info: " <> s <> "\n"
printWarning s = colorShell Yellow $ putText $ "warning: " <> s <> "\n"
printError s = colorShell Red $ putText $ "error: " <> s <> "\n"
printSuccess s = colorShell Green $ putText $ "success: " <> s <> "\n"
printTest s = colorShell Green $ putText $ "running stack: " <> s <> "\n"

myWhich :: Text -> IO Text
myWhich str =
  (toS <$> readProcessStdout_ (shell $ toS $ "which " <> str)) >>= \case
    "" -> die $ "Argo `" <> str <> "` not in $PATH."
    p  -> printInfo ("Found " <> str <> " at " <> p) $> p

Valentin Reis's avatar
Valentin Reis committed
95
sudoRemoveFile :: (Text -> IO ()) -> Text -> Text -> IO ()
Valentin Reis's avatar
Valentin Reis committed
96
sudoRemoveFile printer desc filePath = do
Valentin Reis's avatar
Valentin Reis committed
97
  foundSocket <- doesFileExist $ toS filePath
Valentin Reis's avatar
Valentin Reis committed
98
  when foundSocket $ go False
Valentin Reis's avatar
Valentin Reis committed
99
  printInfo $ "OK: " <> desc <> " " <> filePath
Valentin Reis's avatar
Valentin Reis committed
100 101
 where
  go useSudo = do
Valentin Reis's avatar
Valentin Reis committed
102 103 104 105 106
    printer $ "found stale " <> desc <> " at " <> filePath <> ".. "
    runProcess
        ( shell
        $ toS ((if useSudo then "sudo " else "") <> "rm -rf " <> filePath)
        )
Valentin Reis's avatar
Valentin Reis committed
107
      >>= \case
Valentin Reis's avatar
Valentin Reis committed
108 109
            ExitSuccess ->
              colorShell Green $ putText " Successfully removed.\n"
Valentin Reis's avatar
Valentin Reis committed
110
            ExitFailure _ -> if useSudo
Valentin Reis's avatar
Valentin Reis committed
111 112
              then printer
                ("Failed to remove stale " <> desc <> ", even with sudo.")
Valentin Reis's avatar
Valentin Reis committed
113
              else do
Valentin Reis's avatar
Valentin Reis committed
114 115
                printer
                  ("Failed to remove stale " <> desc <> ". Trying sudo..\n")
Valentin Reis's avatar
Valentin Reis committed
116 117
                go True

Valentin Reis's avatar
Valentin Reis committed
118 119
cleanSocket :: Text -> IO ()
cleanSocket = sudoRemoveFile putText "socket"
Valentin Reis's avatar
Valentin Reis committed
120 121 122 123 124 125 126 127

kbInstallHandler :: IO () -> IO Handler
kbInstallHandler h = installHandler keyboardSignal (Catch h) Nothing

newtype MonitoringResult = PatternMatched Text deriving (Show, Typeable)
instance Exception MonitoringResult

data Instrumentation = Instrumentation
Valentin Reis's avatar
Valentin Reis committed
128
  (ProcessConfig () () ())
Valentin Reis's avatar
Valentin Reis committed
129 130 131 132 133 134 135
  StdOutLog
  StdErrLog
  (Maybe TestText)
  deriving (Show)

data TracebackScan = WarningTraceback | Clean deriving (Show)

Valentin Reis's avatar
Valentin Reis committed
136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151
withAsyncConduitsOnProcess
  :: Process () (ConduitT () ByteString IO ()) (ConduitT () ByteString IO ())
  -> ConduitT ByteString Void IO ()
  -> Maybe TextBehavior
  -> ConduitT ByteString Void IO ()
  -> Maybe TextBehavior
  -> (Async TracebackScan -> Async TracebackScan -> IO b)
  -> IO b
withAsyncConduitsOnProcess p outSink outTest errSink errTest = withAsyncs
  (doFilter outTest (getStdout p) outSink)
  (doFilter errTest (getStderr p) errSink)

withAsyncs :: IO a -> IO a1 -> (Async a -> Async a1 -> IO b) -> IO b
withAsyncs io1 io2 f = withAsync io1 $ \a1 -> withAsync io2 $ \a2 -> f a1 a2

doFilter
Valentin Reis's avatar
Valentin Reis committed
152 153 154 155 156
  :: (MonadIO m)
  => Maybe TextBehavior
  -> ConduitT () ByteString m ()
  -> ConduitT ByteString Void m ()
  -> m TracebackScan
Valentin Reis's avatar
Valentin Reis committed
157 158 159 160
doFilter behavior source sink =
  runConduit $ source .| CB.lines .| makeBehavior behavior `fuseUpstream` sink

makeBehavior
Valentin Reis's avatar
Valentin Reis committed
161 162 163
  :: (MonadIO m)
  => Maybe TextBehavior
  -> ConduitT ByteString ByteString m TracebackScan
Valentin Reis's avatar
Valentin Reis committed
164 165 166 167 168
makeBehavior = \case
  Just ExpectClean       -> warnOnTraceback False
  Just (WaitFor message) -> untilMatch message False
  Nothing                -> awaitForever yield $> Clean

Valentin Reis's avatar
Valentin Reis committed
169 170
warnOnTraceback
  :: (MonadIO m) => Bool -> ConduitT ByteString ByteString m TracebackScan
Valentin Reis's avatar
Valentin Reis committed
171 172 173 174 175
warnOnTraceback sawTraceback = await >>= \case
  Just b | B.isInfixOf "Traceback" b -> yield b >> warnOnTraceback True
         | otherwise                 -> yield b >> warnOnTraceback sawTraceback
  Nothing -> if sawTraceback then return WarningTraceback else return Clean

Valentin Reis's avatar
Valentin Reis committed
176 177 178 179 180
untilMatch
  :: (MonadIO m)
  => Text
  -> Bool
  -> ConduitT ByteString ByteString m TracebackScan
Valentin Reis's avatar
Valentin Reis committed
181 182 183 184 185 186 187 188 189 190 191
untilMatch msg sawTraceback = await >>= \case
  Just b
    | B.isInfixOf "Traceback" b
    -> untilMatch msg True >> yield b >> untilMatch msg True
    | B.isInfixOf (TE.encodeUtf8 msg) b && not sawTraceback
    -> throw (PatternMatched $ TE.decodeUtf8 b)
    | otherwise
    -> yield b >> untilMatch msg sawTraceback
  Nothing -> return Clean

configureConduits
Valentin Reis's avatar
Valentin Reis committed
192 193
  :: WorkingDirectory
  -> ProcessConfig () () ()
Valentin Reis's avatar
Valentin Reis committed
194 195 196 197 198
  -> ProcessConfig
       ()
       (ConduitM () ByteString IO ())
       (ConduitM () ByteString IO ())
configureConduits (WorkingDirectory wd) p =
Valentin Reis's avatar
Valentin Reis committed
199 200 201 202
  setStdout createSource
    $ setStderr createSource
    $ setStdin closed
    $ setWorkingDir (toS wd) p
Valentin Reis's avatar
Valentin Reis committed
203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219

withConduitSinks
  :: Text
  -> Text
  -> (ConduitT ByteString o IO () -> ConduitT ByteString o1 IO () -> IO b)
  -> IO b
withConduitSinks outName errName f =
  withSinkFileNoBuffering (T.unpack outName) $ \outSink ->
    withSinkFileNoBuffering (T.unpack errName) $ \errSink -> f outSink errSink

withSinkFileNoBuffering
  :: FilePath -> (ConduitT ByteString o IO () -> IO b) -> IO b
withSinkFileNoBuffering filepath inner =
  withRunInIO $ \run -> IO.withBinaryFile filepath IO.WriteMode $ \h -> do
    hSetBuffering h NoBuffering
    run $ inner $ sinkHandle h

Valentin Reis's avatar
Valentin Reis committed
220
runI
Valentin Reis's avatar
Valentin Reis committed
221 222
  :: WorkingDirectory
  -> Instrumentation
Valentin Reis's avatar
Valentin Reis committed
223
  -> IO (Either MonitoringResult (ExitCode, TracebackScan, TracebackScan))
Valentin Reis's avatar
Valentin Reis committed
224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241
runI workDir@(WorkingDirectory wd) (Instrumentation crProc (StdOutLog stdOut) (StdErrLog stdErr) t)
  = try
    $ withConduitSinks (wd <> "/" <> stdOut) (wd <> "/" <> stdErr)
    $ \outSink errSink ->
        withProcess (configureConduits workDir crProc) $ \p ->
          withAsyncConduitsOnProcess p
                                     outSink
                                     outTest
                                     errSink
                                     errTest
                                     waitEither
            >>= \case
                  Left  Clean -> (, Clean, Clean) <$> waitExitCode p
                  Right Clean -> (, Clean, Clean) <$> waitExitCode p
                  Right WarningTraceback ->
                    (, Clean, WarningTraceback) <$> waitExitCode p
                  Left WarningTraceback ->
                    (, WarningTraceback, Clean) <$> waitExitCode p
Valentin Reis's avatar
Valentin Reis committed
242 243 244 245 246 247 248 249
 where
  outTest :: Maybe TextBehavior
  errTest :: Maybe TextBehavior
  (outTest, errTest) = case t of
    Just (TestText (TextBehaviorStdout tOut) (TextBehaviorStderr tErr)) ->
      (Just tOut, Just tErr)
    Nothing -> (Nothing, Nothing)

Valentin Reis's avatar
Valentin Reis committed
250 251
processBehaviorToI
  :: ProcessConfig () () () -> ProcessBehavior -> Maybe Instrumentation
Valentin Reis's avatar
Valentin Reis committed
252 253 254 255
processBehaviorToI crProc = \case
  DontRun               -> Nothing
  JustRun stdOut stdErr -> Just $ Instrumentation crProc stdOut stdErr Nothing
  Test t stdOut stdErr  -> Just $ Instrumentation crProc stdOut stdErr (Just t)