Fix race condition in IORef backend

src/Data/Vault/ST/backends/IORef.hs

@@ -1,3 +1,4 @@
+import Control.Concurrent.MVar
 import Data.IORef
 import System.IO.Unsafe (unsafePerformIO)
 
@@ -7,19 +8,22 @@
 {-----------------------------------------------------------------------------
     Locker
 ------------------------------------------------------------------------------}
-data Key s a  = Key    !Unique (IORef (Maybe a))
+data Key s a  = Key    !Unique (IORef (Maybe a)) (MVar ())
 data Locker s = Locker !Unique (IO ())
 
 #if IsStrict
-lock (Key u ref) x = x `seq` (Locker u $ writeIORef ref $ Just x)
+lock (Key u ref _) x = x `seq` (Locker u $ writeIORef ref $ Just x)
 #else
-lock (Key u ref) x = Locker u $ writeIORef ref $ Just x
+lock (Key u ref _) x = Locker u $ writeIORef ref $ Just x
 #endif
 
-unlock (Key k ref) (Locker k' m)
+unlock (Key k ref lock) (Locker k' m)
     | k == k' = unsafePerformIO $ do
+        takeMVar lock
         m
-        readIORef ref     -- FIXME: race condition!
+        val <- readIORef ref
+        putMVar lock ()
+        pure val
     | otherwise = Nothing
 
 {-----------------------------------------------------------------------------
@@ -28,13 +32,13 @@
 -- implemented as a collection of lockers
 newtype Vault s = Vault (Map Unique (Locker s))
 
-newKey = unsafeIOToST $ Key <$> newUnique <*> newIORef Nothing
+newKey = unsafeIOToST $ Key <$> newUnique <*> newIORef Nothing <*> newMVar ()
 
-lookup key@(Key k _)   (Vault m) = unlock key =<< Map.lookup k m
+lookup key@(Key k _ _)   (Vault m) = unlock key =<< Map.lookup k m
 
-insert key@(Key k _) x (Vault m) = Vault $ Map.insert k (lock key x) m
+insert key@(Key k _ _) x (Vault m) = Vault $ Map.insert k (lock key x) m
 
-adjust f key@(Key k _) (Vault m) = Vault $ Map.update f' k m
+adjust f key@(Key k _ _) (Vault m) = Vault $ Map.update f' k m
     where f' = fmap (lock key . f) . unlock key
 
-delete (Key k _) (Vault m)  = Vault $ Map.delete k m
+delete (Key k _ _) (Vault m)  = Vault $ Map.delete k m