Improve ctrl-c and sig INT catching in main UCM loop (#2734)

unison-cli/package.yaml

@@ -42,6 +42,9 @@ dependencies:
 
 library:
   source-dirs: src
+  when:
+    - condition: '!os(windows)'
+      dependencies: unix
 
 tests:
   tests:
@@ -65,9 +68,6 @@ executables:
       - template-haskell
       - temporary
       - unison-cli
-    when:
-      - condition: '!os(windows)'
-        dependencies: unix
 
   transcripts:
     source-dirs: transcripts

unison-cli/src/Compat.hs

@@ -0,0 +1,56 @@
+{-# LANGUAGE CPP #-}
+
+module Compat where
+
+import Control.Concurrent (mkWeakThreadId, myThreadId)
+import Control.Exception (AsyncException (UserInterrupt), throwTo)
+import System.Mem.Weak (deRefWeak)
+import Unison.Prelude
+import qualified UnliftIO
+
+#if defined(mingw32_HOST_OS)
+import qualified GHC.ConsoleHandler as WinSig
+#else
+import qualified System.Posix.Signals as Sig
+#endif
+
+-- | Constructs a default interrupt handler which builds an interrupt handler which throws a
+-- UserInterrupt exception to the thread in which the setup was initially called.
+defaultInterruptHandler :: IO (IO ())
+defaultInterruptHandler = do
+  main_thread <- myThreadId
+  wtid <- mkWeakThreadId main_thread
+  let interrupt = do
+        r <- deRefWeak wtid
+        case r of
+          Nothing -> return ()
+          Just t -> throwTo t UserInterrupt
+  pure interrupt
+
+-- | Replaces any existing interrupt handlers with the provided IO action while the provided
+-- action is running, restoring any existing handlers afterwards.
+withInterruptHandler :: IO () -> IO a -> IO a
+withInterruptHandler handler action = do
+  UnliftIO.bracket
+    installNewHandlers
+    restoreOldHandlers
+    (\_ -> action)
+  where
+    -- Installs the new handler and returns an action to restore the old handlers.
+    installNewHandlers :: IO (IO ())
+    installNewHandlers = do
+#if defined(mingw32_HOST_OS)
+      let sig_handler WinSig.ControlC = handler
+          sig_handler WinSig.Break    = handler
+          sig_handler _               = return ()
+      oldHandler <- WinSig.installHandler (WinSig.Catch sig_handler)
+      pure (void $ WinSig.installHandler oldHandler)
+#else
+      oldQuitHandler <- Sig.installHandler Sig.sigQUIT  (Sig.Catch handler) Nothing
+      oldInterruptHandler <- Sig.installHandler Sig.sigINT   (Sig.Catch handler) Nothing
+      pure do
+        void $ Sig.installHandler Sig.sigQUIT oldQuitHandler Nothing
+        void $ Sig.installHandler Sig.sigINT oldInterruptHandler Nothing
+#endif
+    restoreOldHandlers :: IO () -> IO ()
+    restoreOldHandlers restore = restore

unison-cli/src/Unison/CommandLine/Main.hs

@@ -10,7 +10,7 @@ module Unison.CommandLine.Main
 import Unison.Prelude
 
 import Control.Concurrent.STM (atomically)
-import Control.Exception (finally, catch, AsyncException(UserInterrupt), asyncExceptionFromException)
+import Control.Exception (finally, catch)
 import Control.Monad.State (runStateT)
 import Data.Configurator.Types (Config)
 import Data.IORef
@@ -45,11 +45,12 @@ import qualified Unison.Util.TQueue as Q
 import qualified Unison.CommandLine.Welcome as Welcome
 import Control.Lens (view)
 import Control.Error (rightMay)
-import UnliftIO (catchSyncOrAsync, throwIO, withException)
+import qualified UnliftIO
 import System.IO (hPutStrLn, stderr)
 import Unison.Codebase.Editor.Output (Output)
 import qualified Unison.Codebase.Editor.HandleInput.LoopState as LoopState
 import qualified Unison.Codebase.Editor.HandleInput as HandleInput
+import Compat (withInterruptHandler)
 
 getUserInput
   :: forall m v a
@@ -157,10 +158,7 @@ main dir welcome initialPath (config, cancelConfig) initialInputs runtime codeba
               (putPrettyNonempty o)
               (putPrettyLnUnpaged o))
 
-    let interruptHandler :: SomeException -> IO (Either Event Input)
-        interruptHandler (asyncExceptionFromException -> Just UserInterrupt) = awaitInput
-        interruptHandler e = hPutStrLn stderr ("Exception: " <> show e) *> throwIO e
-        cleanup = do
+    let cleanup = do
           Runtime.terminate runtime
           cancelConfig
           cancelFileSystemWatch
@@ -168,8 +166,7 @@ main dir welcome initialPath (config, cancelConfig) initialInputs runtime codeba
         awaitInput :: IO (Either Event Input)
         awaitInput = do
           -- use up buffered input before consulting external events
-          i <- readIORef initialInputsRef
-          (case i of
+          readIORef initialInputsRef >>= \case
             h:t -> writeIORef initialInputsRef t >> pure h
             [] ->
               -- Race the user input and file watch.
@@ -180,30 +177,57 @@ main dir welcome initialPath (config, cancelConfig) initialInputs runtime codeba
                   e
                 x      -> do
                   writeIORef pageOutput True
-                  pure x) `catchSyncOrAsync` interruptHandler
+                  pure x
+    (onInterrupt, waitForInterrupt) <- buildInterruptHandler
+    withInterruptHandler onInterrupt $ do
+      let loop :: LoopState.LoopState IO Symbol -> IO ()
+          loop state = do
+            writeIORef pathRef (view LoopState.currentPath state)
+            let free = runStateT (runMaybeT HandleInput.loop) state
+            let handleCommand = HandleCommand.commandLine config awaitInput
+                                         (writeIORef rootRef)
+                                         runtime
+                                         notify
+                                         (\o -> let (p, args) = notifyNumbered o in
+                                          putPrettyNonempty p $> args)
+                                         loadSourceFile
+                                         codebase
+                                         serverBaseUrl
+                                         (const Random.getSystemDRG)
+                                         free
+            UnliftIO.race waitForInterrupt (try handleCommand) >>= \case
+              -- SIGINT
+              Left () -> do hPutStrLn stderr "\nAborted."
+                            loop state
+              -- Exception during command execution
+              Right (Left e) -> do printException e
+                                   loop state
+              -- Success
+              Right (Right (o, state')) -> do
+                case o of
+                  Nothing -> pure ()
+                  Just () -> do
+                    writeIORef numberedArgsRef (LoopState._numberedArgs state')
+                    loop state'
 
-    let loop :: LoopState.LoopState IO Symbol -> IO ()
-        loop state = do
-          writeIORef pathRef (view LoopState.currentPath state)
-          let free = runStateT (runMaybeT HandleInput.loop) state
-          (o, state') <- HandleCommand.commandLine config awaitInput
-                                       (writeIORef rootRef)
-                                       runtime
-                                       notify
-                                       (\o -> let (p, args) = notifyNumbered o in
-                                        putPrettyNonempty p $> args)
-                                       loadSourceFile
-                                       codebase
-                                       serverBaseUrl
-                                       (const Random.getSystemDRG)
-                                       free
-          case o of
-            Nothing -> pure ()
-            Just () -> do
-              writeIORef numberedArgsRef (LoopState._numberedArgs state')
-              loop state'
-    -- Run the main program loop, always run cleanup,
-    -- If an exception occurred, print it before exiting.
-    (loop (LoopState.loopState0 root initialPath)
-      `withException` \e -> hPutStrLn stderr ("Exception: " <> show (e :: SomeException)))
-      `finally` cleanup
+      -- Run the main program loop, always run cleanup,
+      -- If an exception occurred, print it before exiting.
+      loop (LoopState.loopState0 root initialPath)
+        `finally` cleanup
+  where
+    printException :: SomeException -> IO ()
+    printException e = hPutStrLn stderr ("Encountered Exception: " <> show (e :: SomeException))
+
+
+-- | Installs a posix interrupt handler for catching SIGINT.
+-- This replaces GHC's default sigint handler which throws a UserInterrupt async exception
+-- and kills the entire process.
+--
+-- Returns an IO action which blocks until a ctrl-c is detected. It may be used multiple
+-- times.
+buildInterruptHandler :: IO (IO (), IO ())
+buildInterruptHandler = do
+    ctrlCMarker <- UnliftIO.newEmptyMVar
+    let onInterrupt = void $ UnliftIO.tryPutMVar ctrlCMarker ()
+    let waitForInterrupt =  UnliftIO.takeMVar ctrlCMarker
+    pure $ (onInterrupt,waitForInterrupt)

unison-cli/unison-cli.cabal

@@ -23,6 +23,7 @@ flag optimized
 
 library
   exposed-modules:
+      Compat
       Unison.Codebase.Editor.AuthorInfo
       Unison.Codebase.Editor.Command
       Unison.Codebase.Editor.HandleCommand
@@ -109,6 +110,9 @@ library
     , unliftio
   if flag(optimized)
     ghc-options: -O2 -funbox-strict-fields
+  if !os(windows)
+    build-depends:
+        unix
   default-language: Haskell2010
 
 executable integration-tests
@@ -253,7 +257,6 @@ executable unison
   main-is: Main.hs
   other-modules:
       ArgParse
-      Compat
       System.Path
       Version
       Paths_unison_cli
@@ -320,9 +323,6 @@ executable unison
     , unliftio
   if flag(optimized)
     ghc-options: -O2 -funbox-strict-fields
-  if !os(windows)
-    build-depends:
-        unix
   default-language: Haskell2010
 
 test-suite tests

unison-cli/unison/Compat.hs

@@ -1,34 +0,0 @@
-{- ORMOLU_DISABLE -} -- Remove this when the file is ready to be auto-formatted
-{-# LANGUAGE CPP #-}
-module Compat where
-
-import Control.Concurrent (mkWeakThreadId, myThreadId)
-import Control.Exception (AsyncException (UserInterrupt), throwTo)
-import System.Mem.Weak (deRefWeak)
-
-#if defined(mingw32_HOST_OS)
-import qualified GHC.ConsoleHandler as WinSig
-#else
-import qualified System.Posix.Signals as Sig
-#endif
-
-installSignalHandlers :: IO ()
-installSignalHandlers = do
-  main_thread <- myThreadId
-  wtid <- mkWeakThreadId main_thread
-  let interrupt = do
-        r <- deRefWeak wtid
-        case r of
-          Nothing -> return ()
-          Just t  -> throwTo t UserInterrupt
-
-#if defined(mingw32_HOST_OS)
-  let sig_handler WinSig.ControlC = interrupt
-      sig_handler WinSig.Break    = interrupt
-      sig_handler _               = return ()
-  _ <- WinSig.installHandler (WinSig.Catch sig_handler)
-#else
-  _ <- Sig.installHandler Sig.sigQUIT  (Sig.Catch interrupt) Nothing
-  _ <- Sig.installHandler Sig.sigINT   (Sig.Catch interrupt) Nothing
-#endif
-  return ()

unison-cli/unison/Main.hs

@@ -50,7 +50,7 @@ import Unison.Symbol (Symbol)
 import qualified Unison.Util.Pretty as P
 import qualified Version
 import UnliftIO.Directory ( getHomeDirectory )
-import Compat ( installSignalHandlers )
+import Compat ( defaultInterruptHandler, withInterruptHandler )
 import ArgParse
     ( UsageRenderer,
       GlobalOptions(GlobalOptions, codebasePathOption),
@@ -68,10 +68,11 @@ import Unison.CommandLine.Welcome (CodebaseInitStatus(..))
 
 main :: IO ()
 main = do
+ interruptHandler <- defaultInterruptHandler
+ withInterruptHandler interruptHandler $ do
   progName <- getProgName
   -- hSetBuffering stdout NoBuffering -- cool
 
-  void installSignalHandlers
   (renderUsageInfo, globalOptions, command) <- parseCLIArgs progName Version.gitDescribeWithDate
   let GlobalOptions{codebasePathOption=mCodePathOption} = globalOptions
   let mcodepath = fmap codebasePathOptionToPath mCodePathOption