Use forkFinally in ResultShare to be sure we always catch exceptions & set MVar

src/Control/Concurrent/ResultShare.hs

@@ -42,6 +42,7 @@ import Prelude
 
 import Control.Concurrent
 import Control.Concurrent.STM.TVar
+import Control.Exception
 import Control.Monad
 import Control.Monad.IO.Class
 import Control.Monad.STM
@@ -51,7 +52,7 @@ import Data.HashMap.Strict (HashMap)
 import qualified Data.HashMap.Strict as M
 
 data ResultShare k v a = ResultShare
-    { _rsMap    :: TVar (HashMap k (MVar v))
+    { _rsMap    :: TVar (HashMap k (MVar (Either SomeException v)))
     , _rsAction :: k -> a -> IO v
     }
 
@@ -70,7 +71,11 @@ stateTVar var f = do
     return a
 
 runShared
-    :: (MonadIO m, Eq k, Hashable k) => ResultShare k v a -> k -> a -> m v
+    :: (MonadIO m, Eq k, Hashable k)
+    => ResultShare k v a
+    -> k
+    -> a
+    -> m (Either SomeException v)
 runShared (ResultShare tvar action) key param = liftIO $ do
     (mvar, new) <- do
         existing <- M.lookup key <$> readTVarIO tvar
@@ -82,8 +87,7 @@ runShared (ResultShare tvar action) key param = liftIO $ do
                     case M.lookup key m of
                         Just v' -> ((v', False), m)
                         Nothing -> ((v , True) , M.insert key v m)
-    when new $ void $ forkIO $ do
-        result <- action key param
+    when new $ void $ forkFinally (action key param) $ \ result -> do
         atomically $ modifyTVar' tvar $ M.delete key
         putMVar mvar result
     readMVar mvar