Utils.hs 6.69 KB
Newer Older
1
{-# LANGUAGE LambdaCase, OverloadedStrings, DataKinds,
2
  FlexibleInstances, ScopedTypeVariables, TypeOperators #-}
3 4 5

module Argo.Utils where

Valentin Reis's avatar
Valentin Reis committed
6
import           Argo.Args
7 8 9 10 11 12
import           Turtle
import           Prelude                 hiding ( FilePath )
import           System.Console.ANSI
import           System.Console.ANSI.Types      ( Color )
import           System.Posix.Signals
import           System.Process          hiding ( shell )
13 14 15 16
import           Data.Conduit
import           Data.Conduit.Process    hiding ( shell )
import           Data.ByteString               as B
                                         hiding ( empty )
17 18 19 20 21 22 23

import           System.IO                      ( BufferMode(NoBuffering)
                                                , hSetBuffering
                                                )
import           Control.Monad.IO.Unlift        ( MonadIO(..)
                                                , withRunInIO
                                                )
24 25 26 27 28
import           Data.Text.Encoding            as TE
import           Data.Conduit.Combinators      as CC
import           Control.Exception.Base
import           Data.Typeable
import           Data.Text                     as T
29
import qualified System.IO                     as IO
30 31 32

-- | Miscellaneous printing utilities
colorShell :: Color -> Shell () -> Shell ()
33
colorShell color she = setC color *> she *> setC White
34 35 36
  where setC c = liftIO $ setSGR [SetColor Foreground Dull c]

printInfo :: Text -> Shell ()
37
printCommand :: Text -> Shell ()
38 39
printError :: Text -> Shell ()
printWarning :: Text -> Shell ()
40
printSuccess :: Text -> Shell ()
41
printTest :: Text -> Shell ()
42 43
dieRed :: Text -> Shell ()

44
printInfo = printf ("Info: " % s % "\n")
45
printCommand = printf ("Running: " % s % "\n")
46 47 48
printWarning = colorShell Yellow . printf ("Warning: " % s % "\n")
printError = colorShell Red . printf ("Error: " % s % "\n")
printSuccess = colorShell Green . printf ("Success: " % s % "\n")
Valentin Reis's avatar
Valentin Reis committed
49
printTest = colorShell Green . printf ("RUNNING TEST: " % s % "\n")
50 51
dieRed str =
  colorShell Red (printf ("Failure: " % s) str) >> exit (ExitFailure 1)
52

53
myWhich :: FilePath -> Shell FilePath
54 55
myWhich str = which str >>= \case
  (Just p) ->
56
    printInfo (format ("Found " % fp % " at " % fp) str p) >> return p
57 58
  Nothing -> die $ format ("Argo `" % fp % "` not in $PATH.") str

59
myWhichMaybe :: FilePath -> Shell (Maybe FilePath)
60
myWhichMaybe str = which str >>= \case
61 62
  (Just p) ->
    printInfo (format ("Found " % fp % " at " % fp) str p) >> return (Just p)
63 64
  Nothing -> return Nothing

65 66 67 68
sudoRemoveFile :: (Text -> Shell ()) -> Text -> FilePath -> Shell ()
sudoRemoveFile printer desc filePath = do
  foundSocket <- testfile filePath
  when foundSocket $ go False
69
  printInfo $ format ("OK: " % s % " " % fp) desc filePath
70 71 72 73
 where
  go useSudo = do
    printer $ format ("found stale " % s % " at " % fp % ".. ") desc filePath
    shell
Valentin Reis's avatar
Valentin Reis committed
74
        (format ((if useSudo then "sudo " else "") % "rm -rf " % fp) filePath)
75
        Turtle.empty
76 77
      >>= \case
            ExitSuccess -> colorShell Green $ printf " Successfully removed.\n"
78
            ExitFailure _ -> if useSudo
79 80 81 82 83 84 85 86 87
              then printer $ format
                ("Failed to remove stale " % s % ", even with sudo.")
                desc
              else do
                printer $ format
                  ("Failed to remove stale " % s % ". Trying sudo..\n")
                  desc
                go True

88
verboseShell :: Text -> Shell Line -> Shell ExitCode
89
verboseShell command i = printCommand command >> shell command i
90

91
verboseShell' :: Text -> Shell Line -> Shell (ExitCode, Text, Text)
92
verboseShell' command i = printCommand command >> shellStrictWithErr command i
93

94
cleanSocket :: FilePath -> Shell ()
95
cleanSocket = sudoRemoveFile printError "socket"
96
cleanLog :: FilePath -> Shell ()
97
cleanLog = sudoRemoveFile printWarning "log folder"
98 99 100

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

102 103
newtype MonitoringResult = PatternMatched Text deriving (Show, Typeable)
instance Exception MonitoringResult
104 105

data Instrumentation = Instrumentation
Valentin Reis's avatar
Valentin Reis committed
106 107 108 109 110
  CreateProcess
  StdOutLog
  StdErrLog
  (Maybe TestText)
  deriving (Show)
111

112 113 114 115 116
data TracebackScan = WarningTraceback | Clean deriving (Show)

runI
  :: Instrumentation
  -> IO (Either MonitoringResult (ExitCode, TracebackScan, TracebackScan))
117 118
runI (Instrumentation crProc (StdOutLog stdOut) (StdErrLog stdErr) t) = try
  (reroutedDaemon crProc)
119
 where
120
  {-reroutedDaemon :: CreateProcess -> IO (ExitCode, (), ())-}
121
  reroutedDaemon process =
122 123 124 125
    withSinkFileNoBuffering (T.unpack stdOut) $ \outSink ->
      withSinkFile (T.unpack stdErr) $ \errSink -> sourceProcessWithStreams
        process
        mempty
126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154
        (makeBehavior outTest `fuseUpstream` outSink)
        (makeBehavior errTest `fuseUpstream` errSink)
  makeBehavior
    :: Maybe TextBehavior -> ConduitT ByteString ByteString IO TracebackScan
  makeBehavior = \case
    Just ExpectClean       -> warnOnTraceback False
    Just (WaitFor message) -> untilMatch message False
    Nothing                -> awaitForever yield >> return Clean
  warnOnTraceback :: Bool -> ConduitT ByteString ByteString IO TracebackScan
  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

  untilMatch :: Text -> Bool -> ConduitT ByteString ByteString IO TracebackScan
  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
  {-withSinkFileNoBuffering-}
    {-:: (MonadUnliftIO m, MonadIO n)-}
    {-=> IO.FilePath-}
    {--> (ConduitM ByteString o n () -> m a)-}
    {--> m a-}
155 156 157 158
  withSinkFileNoBuffering filepath inner =
    withRunInIO $ \run -> IO.withBinaryFile filepath IO.WriteMode $ \h -> do
      hSetBuffering h NoBuffering
      run $ inner $ sinkHandle h
159 160 161 162 163 164 165
  outTest :: Maybe TextBehavior
  errTest :: Maybe TextBehavior
  (outTest, errTest) = case t of
    Just (TestText (TextBehaviorStdout tOut) (TextBehaviorStderr tErr)) ->
      (Just tOut, Just tErr)
    Nothing -> (Nothing, Nothing)

166

Valentin Reis's avatar
Valentin Reis committed
167
processBehaviorToI :: CreateProcess -> ProcessBehavior -> Maybe Instrumentation
168
processBehaviorToI crProc = \case
169 170
  DontRun               -> Nothing
  JustRun stdOut stdErr -> Just $ Instrumentation crProc stdOut stdErr Nothing
171
  Test t stdOut stdErr  -> Just $ Instrumentation crProc stdOut stdErr (Just t)