Use async instead of forkIO

saep · saep · commit fc48abb9dd79 · 2018-03-24T15:15:02.000+01:00
diff --git a/library/Neovim/Context.hs b/library/Neovim/Context.hs
@@ -19,7 +19,6 @@ module Neovim.Context (
     FunctionMapEntry,
     mkFunctionMap,
     runNeovim,
-    forkNeovim,
     err,
     errOnInvalidResult,
     restart,
@@ -42,8 +41,7 @@ module Neovim.Context (
 
 import           Neovim.Classes
 import           Neovim.Context.Internal      (FunctionMap, FunctionMapEntry,
-                                               Neovim, forkNeovim,
-                                               mkFunctionMap,
+                                               Neovim, mkFunctionMap,
                                                newUniqueFunctionName, runNeovim)
 import           Neovim.Exceptions            (NeovimException (..))
 
diff --git a/library/Neovim/Context/Internal.hs b/library/Neovim/Context/Internal.hs
@@ -26,7 +26,6 @@ import           Neovim.Plugin.Classes
 import           Neovim.Plugin.IPC            (SomeMessage)
 
 import           Control.Applicative
-import           Control.Concurrent           (MVar, ThreadId, forkIO)
 import           Control.Exception            (ArithException, ArrayException,
                                                ErrorCall, PatternMatchFail)
 import           Control.Monad.Except
@@ -119,21 +118,6 @@ runNeovimInternal f r (Neovim a) =
             (Right <$> f res) `catches` exceptionHandlers
 
 
--- | Fork a neovim thread with the given custom config value and a custom
--- state. The result of the thread is discarded and only the 'ThreadId' is
--- returend immediately.
--- FIXME This function is pretty much unused and mayhave undesired effects,
---       namely that you cannot register autocmds in the forked thread.
-forkNeovim :: NFData a => iEnv -> Neovim iEnv a -> Neovim env ThreadId
-forkNeovim r a = do
-    cfg <- ask'
-    let threadConfig = cfg
-            { pluginSettings = Nothing -- <- slightly problematic
-            , customConfig = r
-            }
-    liftIO . forkIO . void $ runNeovim threadConfig a
-
-
 -- | Create a new unique function name. To prevent possible name clashes, digits
 -- are stripped from the given suffix.
 newUniqueFunctionName :: Neovim env FunctionName
diff --git a/library/Neovim/Debug.hs b/library/Neovim/Debug.hs
@@ -33,14 +33,15 @@ import           Neovim.Main                  (CommandLineOptions (..),
                                                runPluginProvider)
 import           Neovim.RPC.Common            (RPCConfig)
 
-import           Control.Concurrent
-import           Control.Concurrent.STM
 import           Control.Monad
 import qualified Data.Map                     as Map
 import           Foreign.Store
 import           System.IO                    (stdout)
 import           Text.PrettyPrint.ANSI.Leijen hiding ((<$>))
 import qualified Text.PrettyPrint.ANSI.Leijen as Pretty
+import           UnliftIO.Async               (Async, async, cancel)
+import           UnliftIO.STM
+import           UnliftIO.Concurrent          (putMVar, takeMVar)
 
 import           Prelude
 
@@ -68,7 +69,7 @@ debug env a = disableLogger $ do
                 (cfg { Internal.customConfig = env, Internal.pluginSettings = Nothing })
                 a
 
-            mapM_ killThread tids
+            mapM_ cancel tids
             return res
 
         _ ->
@@ -110,7 +111,7 @@ debug' a = debug () a
 --
 develMain
     :: Maybe NeovimConfig
-    -> IO (Either Doc [ThreadId])
+    -> IO (Either Doc [Async ()])
 develMain mcfg = lookupStore 0 >>= \case
     Nothing -> do
         x <- disableLogger $
@@ -126,9 +127,8 @@ develMain mcfg = lookupStore 0 >>= \case
             return $ Left e
 
         Internal.InitSuccess -> do
-            transitionHandlerThread <- forkIO $ do
-                myTid <- myThreadId
-                void $ transitionHandler (myTid:tids) cfg
+            transitionHandlerThread <- async $ do
+                void $ transitionHandler (tids) cfg
             return $ Right (transitionHandlerThread:tids)
 
         Internal.Quit -> do
@@ -139,7 +139,7 @@ develMain mcfg = lookupStore 0 >>= \case
                 Just x ->
                     deleteStore x
 
-            mapM_ killThread tids
+            mapM_ cancel tids
             return . Left $ text "Quit develMain"
 
         _ ->
@@ -155,7 +155,7 @@ quitDevelMain cfg = putMVar (Internal.transitionTo cfg) Internal.Quit
 restartDevelMain
     :: Internal.Config RPCConfig
     -> Maybe NeovimConfig
-    -> IO (Either Doc [ThreadId])
+    -> IO (Either Doc [Async ()])
 restartDevelMain cfg mcfg = do
     quitDevelMain cfg
     develMain mcfg
diff --git a/library/Neovim/Main.hs b/library/Neovim/Main.hs
@@ -33,6 +33,7 @@ import           Data.Monoid
 import           Options.Applicative
 import           System.IO               (stdin, stdout)
 import           System.SetEnv
+import           UnliftIO.Async          (Async, async, cancel)
 
 import           Prelude
 import           System.Environment
@@ -133,7 +134,7 @@ neovim =
 -- 'Internal.Config' with the custom field set to 'RPCConfig'. These information
 -- can be used to properly clean up a session and then do something else.
 -- The transition handler is first called after the plugin provider has started.
-type TransitionHandler a = [ThreadId] -> Internal.Config RPCConfig -> IO a
+type TransitionHandler a = [Async ()] -> Internal.Config RPCConfig -> IO a
 
 
 -- | This main functions can be used to create a custom executable without
@@ -179,11 +180,11 @@ runPluginProvider os mcfg transitionHandler mDyreParams = case (hostPort os, uni
 
         conf <- Internal.newConfig (pure (providerName os)) newRPCConfig
 
-        ehTid <- forkIO $ runEventHandler
+        ehTid <- async $ runEventHandler
                             evHandlerHandle
                             conf { Internal.pluginSettings = Nothing }
 
-        srTid <- forkIO $ runSocketReader sockreaderHandle conf
+        srTid <- async $ runSocketReader sockreaderHandle conf
 
         ghcEnv <- forM ["GHC_PACKAGE_PATH","CABAL_SANDBOX_CONFIG"] $ \var -> do
             val <- lookupEnv var
@@ -216,7 +217,7 @@ finishDyre threads cfg = takeMVar (Internal.transitionTo cfg) >>= \case
 
     Internal.Restart -> do
         debugM logger "Trying to restart nvim-hs"
-        mapM_ killThread threads
+        mapM_ cancel threads
         Dyre.relaunchMaster Nothing
 
     Internal.Failure e ->
diff --git a/library/Neovim/Plugin.hs b/library/Neovim/Plugin.hs
@@ -41,8 +41,6 @@ import qualified Neovim.Plugin.Startup        as Plugin
 import           Neovim.RPC.FunctionCall
 
 import           Control.Applicative
-import           Control.Concurrent           (ThreadId, forkIO)
-import           Control.Concurrent.STM
 import           Control.Monad                (foldM, void)
 import           Control.Monad.Trans.Resource hiding (register)
 import           Data.ByteString              (ByteString)
@@ -56,6 +54,8 @@ import           Data.Traversable             (forM)
 import           System.Log.Logger
 import           Text.PrettyPrint.ANSI.Leijen (Doc)
 import           UnliftIO.Exception           (SomeException, try)
+import           UnliftIO.Async               (Async, async)
+import           UnliftIO.STM
 
 import           Prelude
 
@@ -69,12 +69,12 @@ type StartupConfig = Plugin.StartupConfig NeovimConfig
 
 startPluginThreads :: Internal.Config StartupConfig
                    -> [Neovim StartupConfig NeovimPlugin]
-                   -> IO (Either Doc ([FunctionMapEntry],[ThreadId]))
+                   -> IO (Either Doc ([FunctionMapEntry],[Async ()]))
 startPluginThreads cfg = runNeovimInternal return cfg . foldM go ([], [])
   where
-    go :: ([FunctionMapEntry], [ThreadId])
+    go :: ([FunctionMapEntry], [Async ()])
        -> Neovim StartupConfig NeovimPlugin
-       -> Neovim StartupConfig ([FunctionMapEntry], [ThreadId])
+       -> Neovim StartupConfig ([FunctionMapEntry], [Async ()])
     go (es, tids) iop = do
         NeovimPlugin p <- iop
         (es', tid) <- registerStatefulFunctionality p
@@ -262,17 +262,17 @@ addAutocmd event (opts@AutocmdOptions{..}) f = do
 -- the corresponding 'TQueue's (i.e. communication channels).
 registerStatefulFunctionality
     :: Plugin env
-    -> Neovim anyEnv ([FunctionMapEntry], ThreadId)
+    -> Neovim anyEnv ([FunctionMapEntry], Async ())
 registerStatefulFunctionality (Plugin { environment = env, exports = fs }) = do
-    q <- liftIO newTQueueIO
+    messageQueue <- liftIO newTQueueIO
     route <- liftIO $ newTVarIO Map.empty
 
     cfg <- Internal.ask'
 
     let startupConfig = cfg
             { Internal.customConfig = env
             , Internal.pluginSettings = Just $ Internal.StatefulSettings
-                (registerPlugin (\_ -> return ())) q route
+                (registerPlugin (\_ -> return ())) messageQueue route
             }
     res <- liftIO . runNeovimInternal return startupConfig . forM fs $ \f ->
             registerFunctionality (getDescription f) (getFunction f)
@@ -283,11 +283,11 @@ registerStatefulFunctionality (Plugin { environment = env, exports = fs }) = do
     let pluginThreadConfig = cfg
             { Internal.customConfig = env
             , Internal.pluginSettings = Just $ Internal.StatefulSettings
-                (registerPlugin registerInGlobalFunctionMap) q route
+                (registerPlugin registerInGlobalFunctionMap) messageQueue route
             }
 
-    tid <- liftIO . forkIO . void . runNeovimInternal return pluginThreadConfig $ do
-                listeningThread q route
+    tid <- liftIO . async . void . runNeovim pluginThreadConfig $ do
+                listeningThread messageQueue route
 
     return (map fst es, tid) -- NB: dropping release functions/keys here
 
diff --git a/library/Neovim/RPC/SocketReader.hs b/library/Neovim/RPC/SocketReader.hs
@@ -31,7 +31,6 @@ import           Neovim.RPC.Common
 import           Neovim.RPC.FunctionCall
 
 import           Control.Applicative
-import           Control.Concurrent         (forkIO)
 import           Control.Concurrent.STM
 import           Control.Monad              (void)
 import           Control.Monad.Trans.Class  (lift)
@@ -45,6 +44,8 @@ import           Data.Monoid
 import qualified Data.Serialize             (get)
 import           System.IO                  (Handle)
 import           System.Log.Logger
+import           UnliftIO.Async             (async, race)
+import           UnliftIO.Concurrent        (threadDelay)
 
 import           Prelude
 
@@ -98,6 +99,7 @@ handleResponse i result = do
             atomically' . modifyTVar' answerMap $ Map.delete i
             atomically' $ putTMVar reply result
 
+
 -- | Act upon the received request or notification. The main difference between
 -- the two is that a notification does not generate a reply. The distinction
 -- between those two cases is done via the first paramater which is 'Maybe' the
@@ -106,14 +108,21 @@ handleRequestOrNotification :: Maybe Int64 -> FunctionName -> [Object]
                             -> ConduitT a Void SocketHandler ()
 handleRequestOrNotification mi m params = do
     cfg <- lift Internal.ask'
-    void . liftIO . forkIO $ handle cfg
+    void . liftIO . async $ race logTimeout (handle cfg)
+    return ()
 
   where
     lookupFunction
         :: TMVar Internal.FunctionMap
         -> STM (Maybe (FunctionalityDescription, Internal.FunctionType))
     lookupFunction funMap = Map.lookup m <$> readTMVar funMap
 
+    logTimeout :: IO ()
+    logTimeout = do
+        let seconds = 1000 * 1000
+        threadDelay (10 * seconds)
+        debugM logger $ "Cancelled another action before it was finished"
+
     handle :: Internal.Config RPCConfig -> IO ()
     handle rpc = atomically (lookupFunction (Internal.globalFunctionMap rpc)) >>= \case
 
diff --git a/library/Neovim/Test.hs b/library/Neovim/Test.hs
@@ -21,8 +21,6 @@ import           Neovim.RPC.Common            (newRPCConfig, RPCConfig)
 import           Neovim.RPC.EventHandler      (runEventHandler)
 import           Neovim.RPC.SocketReader      (runSocketReader)
 
-import           Control.Concurrent
-import           Control.Concurrent.STM       (atomically, putTMVar)
 import           Control.Monad.Reader         (runReaderT)
 import           Control.Monad.Trans.Resource (runResourceT)
 import           GHC.IO.Exception             (ioe_filename)
@@ -32,6 +30,9 @@ import           System.IO                    (Handle)
 import           System.Process
 import           Text.PrettyPrint.ANSI.Leijen (red, text, putDoc, (<$$>))
 import           UnliftIO.Exception
+import           UnliftIO.STM                 (atomically, putTMVar)
+import           UnliftIO.Async               (async, cancel)
+import           UnliftIO.Concurrent          (threadDelay)
 
 
 -- | Type synonym for 'Word'.
@@ -55,7 +56,7 @@ testWithEmbeddedNeovim file timeout r (Internal.Neovim a) =
     runTest `catch` catchIfNvimIsNotOnPath
   where
     runTest = do
-        (_, _, ph, cfg) <- startEmbeddedNvim file timeout
+        (_, _, ph, cfg, cleanUp) <- startEmbeddedNvim file timeout
 
         let testCfg = Internal.retypeConfig r cfg
 
@@ -65,14 +66,16 @@ testWithEmbeddedNeovim file timeout r (Internal.Neovim a) =
         -- result of the operation since neovim cannot send a result if it
         -- has quit.
         let Internal.Neovim q = vim_command "qa!"
-        void . forkIO . void $ runReaderT (runResourceT q) testCfg
+        testRunner <- async . void $ runReaderT (runResourceT q) testCfg
 
         waitForProcess ph >>= \case
             ExitFailure i ->
                 fail $ "Neovim returned with an exit status of: " ++ show i
 
             ExitSuccess ->
                 return ()
+        cancel testRunner
+        cleanUp
 
 
 catchIfNvimIsNotOnPath :: IOException -> IO ()
@@ -87,7 +90,7 @@ catchIfNvimIsNotOnPath e = case ioe_filename e of
 startEmbeddedNvim
     :: Maybe FilePath
     -> Seconds
-    -> IO (Handle, Handle, ProcessHandle, Internal.Config RPCConfig)
+    -> IO (Handle, Handle, ProcessHandle, Internal.Config RPCConfig, IO ())
 startEmbeddedNvim file (Seconds timeout) = do
     args <- case file of
                 Nothing ->
@@ -107,21 +110,23 @@ startEmbeddedNvim file (Seconds timeout) = do
 
     cfg <- Internal.newConfig (pure Nothing) newRPCConfig
 
-    void . forkIO $ runSocketReader
+    socketReader <- async . void $ runSocketReader
                     hout
                     (cfg { Internal.pluginSettings = Nothing })
 
-    void . forkIO $ runEventHandler
+    eventHandler <- async . void $ runEventHandler
                     hin
                     (cfg { Internal.pluginSettings = Nothing })
 
     atomically $ putTMVar
                     (Internal.globalFunctionMap cfg)
                     (Internal.mkFunctionMap [])
 
-    void . forkIO $ do
+    timeoutAsync <- async . void $ do
         threadDelay $ (fromIntegral timeout) * 1000 * 1000
         getProcessExitCode ph >>= maybe (terminateProcess ph) (\_ -> return ())
 
-    return (hin, hout, ph, cfg)
+    let cleanUp = mapM_ cancel [socketReader, eventHandler, timeoutAsync]
+
+    return (hin, hout, ph, cfg, cleanUp)
 
