WIP: Check flakiness against exceptions in lsp-test teardown

ghcide/src/Development/IDE/LSP/LanguageServer.hs

@@ -40,7 +40,8 @@ import           Control.Concurrent.Extra              (newBarrier,
                                                         signalBarrier,
                                                         waitBarrier)
 import           Control.Monad.IO.Unlift               (MonadUnliftIO)
-import           Control.Monad.Trans.Cont              (evalContT)
+import           Control.Monad.Trans.Cont              (ContT (..), evalContT)
+import           Data.Foldable                         (traverse_)
 import           Development.IDE.Core.IdeConfiguration
 import           Development.IDE.Core.Service          (shutdown)
 import           Development.IDE.Core.Shake            hiding (Log)
@@ -54,6 +55,7 @@ import           Language.LSP.Server                   (LanguageContextEnv,
                                                         LspServerLog,
                                                         type (<~>))
 import           System.Timeout                        (timeout)
+
 data Log
   = LogRegisteringIdeConfig !IdeConfiguration
   | LogReactorThreadException !SomeException
@@ -291,9 +293,8 @@ handleInit lifecycleCtx env (TRequestMessage _ _ m params) = otTracedHandler "In
     ideMVar <- newEmptyMVar
 
     let
-      handleServerExceptionOrShutDown me = do
+      loggedTeardown me = do
         -- shutdown shake
-        tryReadMVar ideMVar >>= mapM_ shutdown
         case me of
           Left e -> do
             lifetimeConfirm "due to exception in reactor thread"
@@ -326,8 +327,12 @@ handleInit lifecycleCtx env (TRequestMessage _ _ m params) = otTracedHandler "In
                 $ \(e :: SomeException) -> do
                     exceptionInHandler e
                     k $ TResponseError (InR ErrorCodes_InternalError) (T.pack $ show e) Nothing
-    _ <- flip forkFinally handleServerExceptionOrShutDown $ do
-      runWithWorkerThreads (cmapWithPrio LogSession recorder) dbLoc $ \withHieDb' threadQueue' -> do
+    _ <- flip forkFinally loggedTeardown $ do
+      -- Need to be careful about when the shutdown occurs, it needs to be shut
+      -- down after the session loader and restarting threads, and before the
+      -- hiedb connections are closed.
+      let shutdownSession = tryReadMVar ideMVar >>= traverse_ shutdown
+      runWithWorkerThreads (cmapWithPrio LogSession recorder) dbLoc shutdownSession $ \withHieDb' threadQueue' -> do
         ide <- ctxGetIdeState lifecycleCtx env root withHieDb' threadQueue'
         putMVar ideMVar ide
         -- Keep this after putMVar ideMVar ide; otherwise shutdown during
@@ -349,12 +354,16 @@ handleInit lifecycleCtx env (TRequestMessage _ _ m params) = otTracedHandler "In
 -- | runWithWorkerThreads
 -- create several threads to run the session, db and session loader
 -- see Note [Serializing runs in separate thread]
-runWithWorkerThreads :: Recorder (WithPriority Session.Log) -> FilePath -> (WithHieDb -> ThreadQueue -> IO ()) -> IO ()
-runWithWorkerThreads recorder dbLoc f = evalContT $ do
-            (WithHieDbShield hiedb, threadQueue) <- runWithDb recorder dbLoc
-            sessionRestartTQueue <- withWorkerQueueSimple (cmapWithPrio Session.LogSessionWorkerThread recorder) "RestartTQueue"
-            sessionLoaderTQueue <- withWorkerQueueSimple (cmapWithPrio Session.LogSessionWorkerThread recorder) "SessionLoaderTQueue"
-            liftIO $ f hiedb (ThreadQueue threadQueue sessionRestartTQueue sessionLoaderTQueue)
+runWithWorkerThreads :: Recorder (WithPriority Session.Log) -> FilePath -> IO () -> (WithHieDb -> ThreadQueue -> IO ()) -> IO ()
+runWithWorkerThreads recorder dbLoc shutdownSession f = evalContT $ do
+  (WithHieDbShield hiedb, threadQueue) <- runWithDb recorder dbLoc
+  -- The shake session needs to be shut down prior to the hiedb connections
+  -- being cleaned up, otherwise shake could be referencing dead connections.
+  -- This is passed in via the callsites.
+  ContT $ \action -> action () `finally` shutdownSession
+  sessionRestartTQueue <- withWorkerQueueSimple (cmapWithPrio Session.LogSessionWorkerThread recorder) "RestartTQueue"
+  sessionLoaderTQueue <- withWorkerQueueSimple (cmapWithPrio Session.LogSessionWorkerThread recorder) "SessionLoaderTQueue"
+  liftIO $ f hiedb (ThreadQueue threadQueue sessionRestartTQueue sessionLoaderTQueue)
 
 -- | Runs the action until it ends or until the given MVar is put.
 --   It is important, that the thread that puts the 'MVar' is not dropped before it puts the 'MVar' i.e. it should

ghcide/src/Development/IDE/Main.hs

@@ -378,7 +378,7 @@ defaultMain recorder Arguments{..} = withHeapStats (cmapWithPrio LogHeapStats re
         Check argFiles -> do
           let dir = argsProjectRoot
           dbLoc <- getHieDbLoc dir
-          runWithWorkerThreads (cmapWithPrio LogSession recorder) dbLoc $ \hiedb threadQueue -> do
+          runWithWorkerThreads (cmapWithPrio LogSession recorder) dbLoc mempty $ \hiedb threadQueue -> do
             -- GHC produces messages with UTF8 in them, so make sure the terminal doesn't error
             hSetEncoding stdout utf8
             hSetEncoding stderr utf8
@@ -436,7 +436,7 @@ defaultMain recorder Arguments{..} = withHeapStats (cmapWithPrio LogHeapStats re
         Custom (IdeCommand c) -> do
           let root = argsProjectRoot
           dbLoc <- getHieDbLoc root
-          runWithWorkerThreads (cmapWithPrio LogSession recorder) dbLoc $ \hiedb threadQueue -> do
+          runWithWorkerThreads (cmapWithPrio LogSession recorder) dbLoc mempty $ \hiedb threadQueue -> do
             sessionLoader <- loadSessionWithOptions (cmapWithPrio LogSession recorder) argsSessionLoadingOptions "." (tLoaderQueue threadQueue)
             let def_options = argsIdeOptions argsDefaultHlsConfig sessionLoader
                 ideOptions = def_options

hls-test-utils/hls-test-utils.cabal

@@ -50,6 +50,7 @@ library
     , lsp
     , lsp-test                ^>=0.18
     , lsp-types               ^>=2.4
+    , primitive
     , safe-exceptions
     , string-interpolate      >= 0.3.1
     , tasty

hls-test-utils/src/Test/Hls.hs

@@ -82,6 +82,7 @@ import           Control.Lens.Extras                      (is)
 import           Control.Monad                            (guard, unless, void)
 import           Control.Monad.Extra                      (forM)
 import           Control.Monad.IO.Class
+import           Control.Monad.Primitive                  (keepAlive)
 import           Data.Aeson                               (Result (Success),
                                                            Value (Null),
                                                            fromJSON, toJSON)
@@ -830,8 +831,8 @@ wrapClientLogger logger = do
 runSessionWithTestConfig :: Pretty b => TestConfig b -> (FilePath -> Session a) -> IO a
 runSessionWithTestConfig TestConfig{..} session =
     runSessionInVFS testDirLocation $ \root -> shiftRoot root $ do
-    (inR, inW) <- createPipe
-    (outR, outW) <- createPipe
+    pipeIn@(inR, inW) <- createPipe
+    pipeOut@(outR, outW) <- createPipe
     let serverRoot = fromMaybe root testServerRoot
     let clientRoot = fromMaybe root testClientRoot
 
@@ -850,19 +851,21 @@ runSessionWithTestConfig TestConfig{..} session =
     timeoutOverride <- fmap read <$> lookupEnv "LSP_TIMEOUT"
     let sconf' = testConfigSession { lspConfig = hlsConfigToClientConfig testLspConfig, messageTimeout = fromMaybe (messageTimeout defaultConfig) timeoutOverride}
         arguments = testingArgs serverRoot recorderIde plugins
-    server <- async $
-        IDEMain.defaultMain (cmapWithPrio LogIDEMain recorderIde)
-            arguments { argsHandleIn = pure inR , argsHandleOut = pure outW }
-    result <- runSessionWithHandles inW outR sconf' testConfigCaps clientRoot (session root)
-    hClose inW
-    timeout 3 (wait server) >>= \case
-        Just () -> pure ()
-        Nothing -> do
-            logWith testRecorder Info (TestServerExitTimeoutSeconds 3)
-            (t, _) <- duration $ cancel server
-            logWith testRecorder Info (TestServerCancelFinished (showDuration t))
-    logWith testRecorder Info TestRunFinished
-    pure result
+
+    keepAlive (pipeIn, pipeOut) $ do
+      server <- async $
+          IDEMain.defaultMain (cmapWithPrio LogIDEMain recorderIde)
+              arguments { argsHandleIn = pure inR , argsHandleOut = pure outW }
+      result <- runSessionWithHandles inW outR sconf' testConfigCaps clientRoot (session root)
+      hClose inW
+      timeout 3 (wait server) >>= \case
+          Just () -> pure ()
+          Nothing -> do
+              logWith testRecorder Info (TestServerExitTimeoutSeconds 3)
+              (t, _) <- duration $ cancel server
+              logWith testRecorder Info (TestServerCancelFinished (showDuration t))
+      logWith testRecorder Info TestRunFinished
+      pure result
 
     where
         shiftRoot shiftTarget f  =