Utils.hs 6.69 KB
Newer Older
Valentin Reis's avatar
Valentin Reis committed
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 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 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171
{-# LANGUAGE LambdaCase, OverloadedStrings, DataKinds,
  FlexibleInstances, ScopedTypeVariables, TypeOperators #-}

module Argo.Utils where

import           Argo.Args
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 )
import           Data.Conduit
import           Data.Conduit.Process    hiding ( shell )
import           Data.ByteString               as B
                                         hiding ( empty )

import           System.IO                      ( BufferMode(NoBuffering)
                                                , hSetBuffering
                                                )
import           Control.Monad.IO.Unlift        ( MonadIO(..)
                                                , withRunInIO
                                                )
import           Data.Text.Encoding            as TE
import           Data.Conduit.Combinators      as CC
import           Control.Exception.Base
import           Data.Typeable
import           Data.Text                     as T
import qualified System.IO                     as IO

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

printInfo :: Text -> Shell ()
printCommand :: Text -> Shell ()
printError :: Text -> Shell ()
printWarning :: Text -> Shell ()
printSuccess :: Text -> Shell ()
printTest :: Text -> Shell ()
dieRed :: Text -> Shell ()

printInfo = printf ("Info: " % s % "\n")
printCommand = printf ("Running: " % s % "\n")
printWarning = colorShell Yellow . printf ("Warning: " % s % "\n")
printError = colorShell Red . printf ("Error: " % s % "\n")
printSuccess = colorShell Green . printf ("Success: " % s % "\n")
printTest = colorShell Green . printf ("RUNNING TEST: " % s % "\n")
dieRed str =
  colorShell Red (printf ("Failure: " % s) str) >> exit (ExitFailure 1)

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

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

sudoRemoveFile :: (Text -> Shell ()) -> Text -> FilePath -> Shell ()
sudoRemoveFile printer desc filePath = do
  foundSocket <- testfile filePath
  when foundSocket $ go False
  printInfo $ format ("OK: " % s % " " % fp) desc filePath
 where
  go useSudo = do
    printer $ format ("found stale " % s % " at " % fp % ".. ") desc filePath
    shell
        (format ((if useSudo then "sudo " else "") % "rm -rf " % fp) filePath)
        Turtle.empty
      >>= \case
            ExitSuccess -> colorShell Green $ printf " Successfully removed.\n"
            ExitFailure _ -> if useSudo
              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

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

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

cleanSocket :: FilePath -> Shell ()
cleanSocket = sudoRemoveFile printError "socket"
cleanLog :: FilePath -> Shell ()
cleanLog = sudoRemoveFile printWarning "log folder"

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

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

data Instrumentation = Instrumentation
  CreateProcess
  StdOutLog
  StdErrLog
  (Maybe TestText)
  deriving (Show)

data TracebackScan = WarningTraceback | Clean deriving (Show)

runI
  :: Instrumentation
  -> IO (Either MonitoringResult (ExitCode, TracebackScan, TracebackScan))
runI (Instrumentation crProc (StdOutLog stdOut) (StdErrLog stdErr) t) = try
  (reroutedDaemon crProc)
 where
  {-reroutedDaemon :: CreateProcess -> IO (ExitCode, (), ())-}
  reroutedDaemon process =
    withSinkFileNoBuffering (T.unpack stdOut) $ \outSink ->
      withSinkFile (T.unpack stdErr) $ \errSink -> sourceProcessWithStreams
        process
        mempty
        (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-}
  withSinkFileNoBuffering filepath inner =
    withRunInIO $ \run -> IO.withBinaryFile filepath IO.WriteMode $ \h -> do
      hSetBuffering h NoBuffering
      run $ inner $ sinkHandle h
  outTest :: Maybe TextBehavior
  errTest :: Maybe TextBehavior
  (outTest, errTest) = case t of
    Just (TestText (TextBehaviorStdout tOut) (TextBehaviorStderr tErr)) ->
      (Just tOut, Just tErr)
    Nothing -> (Nothing, Nothing)


processBehaviorToI :: CreateProcess -> ProcessBehavior -> Maybe Instrumentation
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)