diff options
Diffstat (limited to 'users/Profpatsch/lorri-wait-for-eval/LorriWaitForEval.hs')
-rw-r--r-- | users/Profpatsch/lorri-wait-for-eval/LorriWaitForEval.hs | 173 |
1 files changed, 173 insertions, 0 deletions
diff --git a/users/Profpatsch/lorri-wait-for-eval/LorriWaitForEval.hs b/users/Profpatsch/lorri-wait-for-eval/LorriWaitForEval.hs new file mode 100644 index 000000000000..a1a45864019c --- /dev/null +++ b/users/Profpatsch/lorri-wait-for-eval/LorriWaitForEval.hs @@ -0,0 +1,173 @@ +{-# LANGUAGE DerivingStrategies #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE NumericUnderscores #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE QuasiQuotes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# OPTIONS_GHC -Wall #-} + +module Main where + +import Conduit +import Conduit qualified as Cond +import Control.Concurrent +import Control.Concurrent.Async qualified as Async +import Control.Monad +import Data.Aeson.BetterErrors qualified as Json +import Data.Bifunctor +import Data.Conduit.Binary qualified as Conduit.Binary +import Data.Conduit.Combinators qualified as Cond +import Data.Conduit.Process +import Data.Error +import Data.Function +import Data.Functor +import Data.Text.IO (hPutStrLn) +import MyPrelude +import System.Directory qualified as Dir +import System.Environment qualified as Env +import System.Exit qualified as Exit +import System.FilePath (takeDirectory) +import System.FilePath.Posix qualified as FilePath +import System.IO (stderr) +import System.Posix qualified as Posix +import Prelude hiding (log) + +data LorriEvent = LorriEvent + { nixFile :: Text, + eventType :: LorriEventType + } + deriving stock (Show) + +data LorriEventType + = Completed + | Started + | EvalFailure + deriving stock (Show) + +main :: IO () +main = do + argv <- Env.getArgs <&> nonEmpty + + dir <- Dir.getCurrentDirectory + shellNix <- + findShellNix dir >>= \case + Nothing -> Exit.die [fmt|could not find any shell.nix in or above the directory {dir}|] + Just s -> pure s + getEventChan :: MVar (Chan LorriEvent) <- newEmptyMVar + Async.race_ + ( do + sendEventChan :: Chan LorriEvent <- newChan + (exitCode, ()) <- + sourceProcessWithConsumer + (proc "lorri" ["internal", "stream-events"]) + $ + -- first, we want to send a message over the chan that the process is running (for timeout) + liftIO (putMVar getEventChan sendEventChan) + *> Conduit.Binary.lines + .| Cond.mapC + ( \jsonBytes -> + (jsonBytes :: ByteString) + & Json.parseStrict + ( Json.key + "Completed" + ( do + nixFile <- Json.key "nix_file" Json.asText + pure LorriEvent {nixFile, eventType = Completed} + ) + Json.<|> Json.key + "Started" + ( do + nixFile <- Json.key "nix_file" Json.asText + pure LorriEvent {nixFile, eventType = Started} + ) + Json.<|> Json.key + "Failure" + ( do + nixFile <- Json.key "nix_file" Json.asText + pure LorriEvent {nixFile, eventType = EvalFailure} + ) + ) + & first Json.displayError' + & first (map newError) + & first (smushErrors [fmt|Cannot parse line returned by lorri: {jsonBytes & bytesToTextUtf8Lenient}|]) + & unwrapError + ) + .| (Cond.mapM_ (\ev -> writeChan sendEventChan ev)) + + log [fmt|lorri internal stream-events exited {show exitCode}|] + ) + ( do + let waitMs ms = threadDelay (ms * 1000) + + -- log [fmt|Waiting for lorri event for {shellNix}|] + + eventChan <- takeMVar getEventChan + + let isOurEvent ev = FilePath.normalise (ev & nixFile & textToString) == FilePath.normalise shellNix + + let handleEvent ev = + case ev & eventType of + Started -> + log [fmt|waiting for lorri build to finish|] + Completed -> do + log [fmt|build completed|] + exec (inDirenvDir (takeDirectory shellNix) <$> argv) + EvalFailure -> do + log [fmt|evaluation failed! for path {ev & nixFile}|] + Exit.exitWith (Exit.ExitFailure 111) + + -- wait for 100ms for the first message from lorri, + -- or else assume lorri is not building the project yet + Async.race + (waitMs 100) + ( do + -- find the first event that we can use + let go = do + ev <- readChan eventChan + if isOurEvent ev then pure ev else go + go + ) + >>= \case + Left () -> do + log [fmt|No event received from lorri, assuming this is the first evaluation|] + exec argv + Right ch -> handleEvent ch + + runConduit $ + repeatMC (readChan eventChan) + .| filterC isOurEvent + .| mapM_C handleEvent + ) + where + inDirenvDir dir' argv' = ("direnv" :| ["exec", dir']) <> argv' + exec = \case + Just (exe :| args') -> Posix.executeFile exe True args' Nothing + Nothing -> Exit.exitSuccess + +log :: Text -> IO () +log msg = hPutStrLn stderr [fmt|lorri-wait-for-eval: {msg}|] + +-- | Searches from the current directory upwards, until it finds the `shell.nix`. +findShellNix :: FilePath -> IO (Maybe FilePath) +findShellNix curDir = do + let go :: (FilePath -> IO (Maybe FilePath)) + go dir = do + let file = dir FilePath.</> "shell.nix" + Dir.doesFileExist file >>= \case + True -> pure (Just file) + False -> do + let parent = FilePath.takeDirectory dir + if parent == dir + then pure Nothing + else go parent + go (FilePath.normalise curDir) + +smushErrors :: Foldable t => Text -> t Error -> Error +smushErrors msg errs = + errs + -- hrm, pretty printing and then creating a new error is kinda shady + & foldMap (\err -> "\n- " <> prettyError err) + & newError + & errorContext msg |