Rework GHCi restarting logic

This now more properly detects restarting conditions and handles them.
Should fix the situation where infinite output breaks the ghci process
but it's detected only on the next message received.

1 files changed, 38 insertions, 40 deletions

diff --git a/src/Ghci.hs b/src/Ghci.hs
index 4c6dca1..7f06f99 100644
--- a/src/Ghci.hs
+++ b/src/Ghci.hs
@@ -13,7 +13,7 @@ module Ghci (
   parseSettingsPaste,
 ) where
 
-import Control.Exception (catch, SomeException)
+import Control.Exception (catch, catchNoPropagate, rethrowIO, ExceptionWithContext, SomeException)
 import Control.Monad (replicateM, when)
 import Data.Bifunctor (first)
 import qualified Data.ByteString as BS
@@ -26,7 +26,7 @@ import qualified Data.ByteString.Short as BSS
 import Data.Char (isSpace, toLower)
 import Data.List (nub)
 import Foreign (allocaBytes)
-import System.IO (hFlush, hIsClosed, hGetBufSome, hGetLine, hPutStrLn, stdout, stderr, Handle)
+import System.IO (hFlush, hIsClosed, hGetBufSome, hPutStrLn, stdout, stderr, Handle)
 import System.Process
 import System.Random (getStdRandom, uniformR)
 import System.Timeout (timeout, Timeout)
@@ -64,6 +64,9 @@ parseSettingsPaste = ParseSettings 50000 False
 data Result a = Error String | Ignored | Return a | Exited
   deriving (Show)
 
+data ShouldTerminate = TermNow | StillAlive
+  deriving (Show)
+
 makeGhci :: IO Ghci
 makeGhci = do
   (pipeOut, pipeIn) <- createPipe
@@ -111,9 +114,28 @@ runStmtClever ghci pset line =
     long `startsWith` short = take (length short) long == short
 
 runStmt :: Ghci -> ParseSettings -> String -> IO (Ghci, Result String)
-runStmt ghci pset line = timeouting 2_000_000 (restarting (\g -> runStmt' g pset line)) ghci
+runStmt ghci pset line = do
+  closed <- hIsClosed (ghciStdin ghci)
+  ghci' <- if closed
+             then do
+               putStrLn "ghci: restarting due to closed stdin"
+               hFlush stdout
+               terminateGhci ghci
+               makeGhci
+             else return ghci
+
+  (sterm, res) <- timeouting 2_000_000 (restarting (\g -> runStmt' g pset line)) ghci'
+
+  ghci'' <- case sterm of
+    TermNow -> do
+      terminateGhci ghci'
+      ghci'' <- makeGhci
+      return ghci''
+    StillAlive -> return ghci'
 
-timeouting :: Int -> (Ghci -> IO (Ghci, Result a)) -> Ghci -> IO (Ghci, Result a)
+  return (ghci'', res)
+
+timeouting :: Int -> (Ghci -> IO (ShouldTerminate, Result a)) -> Ghci -> IO (ShouldTerminate, Result a)
 timeouting microseconds f ghci =
   -- TODO: The timeout handling code never actually runs, because the timeout
   -- exception is already handled by the catch-all exception handler in
@@ -121,62 +143,38 @@ timeouting microseconds f ghci =
   timeout microseconds (f ghci) >>= \case
     Nothing -> do putStrLn "ghci: restarting due to timeout"
                   hFlush stdout
-                  terminateGhci ghci
-                  ghci' <- makeGhci
-                  return (ghci', Error "<timeout>")
+                  return (TermNow, Error "<timeout>")
     Just pair -> return pair
 
-restarting :: (Ghci -> IO a) -> Ghci -> IO (Ghci, Result a)
-restarting f ghci = do
-  closed <- hIsClosed (ghciStdin ghci)
-  ghci' <- if closed
-             then do
-               putStrLn "ghci: restarting due to closed stdin"
-               hFlush stdout
-               terminateGhci ghci
-               makeGhci
-             else return ghci
-
-  fmap (\x -> (ghci', Return x)) (f ghci')
-    `catch` (\e -> do let _ = e :: Timeout
-                      putStrLn "ghci: restarting due to timeout (caught in restarting)"
-                      hFlush stdout
-                      terminateGhci ghci
-                      -- putStrLn $ "ghci: terminated"
-                      -- hFlush stdout
-                      ghci'' <- makeGhci
-                      -- putStrLn $ "ghci: new made"
-                      -- hFlush stdout
-                      return (ghci'', Error "<timeout>"))
+restarting :: (Ghci -> IO (ShouldTerminate, a)) -> Ghci -> IO (ShouldTerminate, Result a)
+restarting f ghci =
+  fmap (\(st, x) -> (st, Return x)) (f ghci)
+    -- propagate Timeout exceptions to 'timeouting'
+    `catchNoPropagate` (\e -> rethrowIO (e :: ExceptionWithContext Timeout))
     `catch` (\e -> do let _ = e :: SomeException
                       putStrLn $ "ghci: restarting due to exception: " ++ show e
                       hFlush stdout
-                      terminateGhci ghci'
-                      -- putStrLn $ "ghci: terminated"
-                      -- hFlush stdout
-                      ghci'' <- makeGhci
-                      -- putStrLn $ "ghci: new made"
-                      -- hFlush stdout
-                      return (ghci'', Error "Oops, something went wrong"))
+                      return (TermNow, Error "Oops, something went wrong"))
 
 terminateGhci :: Ghci -> IO ()
 terminateGhci ghci = terminateProcess (ghciProc ghci)
 
-runStmt' :: Ghci -> ParseSettings -> String -> IO String
+runStmt' :: Ghci -> ParseSettings -> String -> IO (ShouldTerminate, String)
 runStmt' ghci pset stmt = do
   tag <- updatePrompt ghci
   ghciPutStrLn (ghciStdin ghci) stmt
   hFlush (ghciStdin ghci)
   let readmax = psMaxOutputLen pset + 8192
   (output, reason) <- hGetUntilUTF8 (ghciStdout ghci) (Just readmax) tag
+  hPutStrLn stderr $ "runStmt': reason = " ++ show reason
   case reason of
     ReachedMaxLen -> do
       terminateGhci ghci  -- because we lost the new prompt
-      return (formatOutput output)  -- don't need to strip tag because we read more than the max output len
-    ReachedTag -> return (formatOutput $ take (length output - length tag) output)
+      return (TermNow, formatOutput output)  -- don't need to strip tag because we read more than the max output len
+    ReachedTag -> return (StillAlive, formatOutput $ take (length output - length tag) output)
     ReachedEOF -> do
       terminateGhci ghci
-      return (formatOutput output)
+      return (TermNow, formatOutput output)
   where
     formatOutput output =
       let output' | psJoinLines pset = replaceNewlines (dropBothSlow isSpace output)