feat(users/Profpatsch): add lorri-wait-for-eval r/4265

A small exec wrapper which will query the lorri daemon for the last
few events, and if it sees a build running for the current
project (searching upwards for shell.nix), it will wait for the build
to finish before executing the command (in the new direnv
environment).

TODO: should patch lorri so that it can provide this information in a
better digestive format; right now it might have a later evaluation
running, so it’s hard to know which completion to wait for …

Change-Id: I8fa4a10484830a731fe3ec58f2694498f46a496c
Reviewed-on: https://cl.tvl.fyi/c/depot/+/5903
Tested-by: BuildkiteCI
Reviewed-by: Profpatsch <mail@profpatsch.de>

3 files changed, 217 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..05c5eb9f2bb6
--- /dev/null
+++ b/users/Profpatsch/lorri-wait-for-eval/LorriWaitForEval.hs
@@ -0,0 +1,189 @@
+{-# 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 qualified Conduit as Cond
+import Control.Concurrent
+import qualified Control.Concurrent.Async as Async
+import Control.Monad
+import qualified Data.Aeson.BetterErrors as Json
+import Data.Bifunctor
+import Data.ByteString (ByteString)
+import qualified Data.Conduit.Binary as Conduit.Binary
+import qualified Data.Conduit.Combinators as Cond
+import Data.Conduit.Process
+import Data.Error
+import Data.Function
+import Data.Functor
+import Data.List.NonEmpty (NonEmpty ((:|)), nonEmpty)
+import Data.Text (Text)
+import qualified Data.Text as Text
+import qualified Data.Text.Encoding
+import qualified Data.Text.Encoding.Error
+import Data.Text.IO (hPutStrLn)
+import PyF
+import qualified System.Directory as Dir
+import qualified System.Environment as Env
+import qualified System.Exit as Exit
+import System.FilePath (takeDirectory)
+import qualified System.FilePath.Posix as FilePath
+import System.IO (stderr)
+import qualified System.Posix as Posix
+import Prelude hiding (log)
+
+data LorriEvent = LorriEvent
+  { nixFile :: Text,
+    eventType :: LorriEventType
+  }
+  deriving stock (Show)
+
+data ChanToken a
+  = -- | so we can see that the lorri thread has been initialized
+    NoEventYet
+  | ChanEvent a
+
+data LorriEventType
+  = Completed
+  | Started
+  | Failure
+  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 = Failure}
+                              )
+                        )
+                      & 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)
+                Failure -> 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 -> pure Nothing
+  go curDir
+
+textToString :: Text -> String
+textToString = Text.unpack
+
+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
+
+-- | decode a Text from a ByteString that is assumed to be UTF-8,
+-- replace non-UTF-8 characters with the replacment char U+FFFD.
+bytesToTextUtf8Lenient :: Data.ByteString.ByteString -> Data.Text.Text
+bytesToTextUtf8Lenient =
+  Data.Text.Encoding.decodeUtf8With Data.Text.Encoding.Error.lenientDecode
diff --git a/users/Profpatsch/lorri-wait-for-eval/default.nix b/users/Profpatsch/lorri-wait-for-eval/default.nix
new file mode 100644
index 000000000000..af8135839a52
--- /dev/null
+++ b/users/Profpatsch/lorri-wait-for-eval/default.nix
@@ -0,0 +1,19 @@
+{ depot, pkgs, lib, ... }:
+
+let
+  lorri-wait-for-eval = pkgs.writers.writeHaskell "lorri-wait-for-eval"
+    {
+      libraries = [
+        pkgs.haskellPackages.async
+        pkgs.haskellPackages.aeson-better-errors
+        pkgs.haskellPackages.conduit-extra
+        pkgs.haskellPackages.error
+        pkgs.haskellPackages.PyF
+        pkgs.haskellPackages.unliftio
+      ];
+      ghcArgs = [ "-threaded" ];
+
+    } ./LorriWaitForEval.hs;
+
+in
+lorri-wait-for-eval
diff --git a/users/Profpatsch/nix-home/default.nix b/users/Profpatsch/nix-home/default.nix
index 3f0b7c9c39c5..ee154c549a6b 100644
--- a/users/Profpatsch/nix-home/default.nix
+++ b/users/Profpatsch/nix-home/default.nix
@@ -150,6 +150,15 @@ let
               name = "scripts/ytextr";
               path = depot.users.Profpatsch.ytextr;
             }
+            {
+              name = "scripts/lorri-wait-for-eval";
+              path = depot.users.Profpatsch.lorri-wait-for-eval;
+            }
+            {
+              name = "scripts/lw";
+              path = depot.users.Profpatsch.lorri-wait-for-eval;
+            }
+
           ]
           ++
           (lib.pipe depot.users.Profpatsch.aliases [