Cleaning up threads in Haskell

When you press Ctrl+C in ghci or when your code gets reloaded in ghcid, only the main thread is killed (via a UserInterrupt exception). All other spawned threads continue to run. Here's how to clean them up:

Since the addition of listThreads in GHC 9.6.1, you can label your threads with labelThread then list all, filter, and kill them later. This won't work for threads spawned by dependencies since they're probably not labeled. Perhaps a labeling convention could emerge to support this use case, or GHC could apply descriptive labels to all threads to enable user code to disambiguate between user threads and runtime threads (e.g. the IO manager). One hacky idea that doesn't need labels is to sort the list of threads, drop the first 3 (which probably correspond to runtime threads, dangerous assumption though), and kill the rest.

Prior to GHC 9.6.1, it was not possible to list threads, so it was up to you the programmer to clean up. Here's an idea for cleaning up at least some threads (although it suffers from the same shortcoming as above in that it can't catch threads spawned by dependencies):

• Keep track of all currently-running threads in an IORef (Set (Async a))
• When an exception occurs in any thread, hand it to the main thread with putMVar exceptionVar exception
• Have the main thread wait for an exception then kill all currently-running threads

Implementation:

main = do
  safeAsync $ runApp

  dieOnException

-- NOINLINE ensures that there is only one MVar
{-# NOINLINE errorVar #-}
errorVar = unsafePerformIO $ newEmptyMVar :: MVar SomeException

-- Set of currently-running threads
{-# NOINLINE asyncsRef #-}
asyncsRef = unsafePerformIO $ newIORef Set.empty :: IORef (Set (Async a))

-- Waits for an exception the kills all threads
dieOnException = (readMVar errorVar >>= print) `finally` (do
  mapM_ cancel =<< readIORef asyncsRef
  putStrLn "exiting entire Haskell process"
  exitFailure)

-- A version of `async` that additionally stores itself in the currently-running threads.
-- It also notifies the main thread when there's an exception.
safeAsync a = do
  asyncValue <- async a
  forkIO $ do
    (wait asyncValue `finally` (atomicModifyIORef' asyncsRef (Set.delete asyncValue))) `catches`
      [ Handler $ \(e :: AsyncCancelled) -> return ()
      , Handler $ \(e :: SomeException) -> putMVar errorVar e
      ]
  atomicModifyIORef' asyncsRef (Set.insert asyncValue)
  return asyncValue