Fix a race condition in RWVar

src-control/Control/Concurrent/Class/MonadSTM/RWVar.hs

@@ -75,47 +75,85 @@ withReadAccess rwvar k =
       k
 
 {-# SPECIALISE unsafeAcquireWriteAccess :: RWVar IO a -> IO a #-}
-unsafeAcquireWriteAccess :: MonadSTM m => RWVar m a -> m a
-unsafeAcquireWriteAccess rw@(RWVar !var) = do
-    may <- atomically $ readTVar var >>= \case
-      Reading n x
-        | n == 0 -> do
-            writeTVar var Writing
-            pure (Just x)
-        | otherwise -> do
-            writeTVar var (WaitingToWrite n x)
-            pure Nothing
-      WaitingToWrite n x
-        | n == 0 -> do
-            writeTVar var Writing
-            pure (Just x)
-        | otherwise -> retry
-      -- If we see this, someone else has already acquired a write lock, so we retry.
-      Writing -> retry
-    -- When the previous STM transaction returns Nothing, it signals that we set
-    -- the RWState to WaitingToWrite. All that remains is to try acquiring write
-    -- access again.
+-- | Acquire write access. This function assumes that it runs in a masked
+-- context, and that is properly paired with an 'unsafeReleaseWriteAccess'!
+--
+-- If multiple threads try to acquire write access concurrently, then they will
+-- race for access. However, if a thread has set RWState to WaitingToWrite, then
+-- it is guaranteed that the same thread will acquire write access when all
+-- readers have finished. That is, other writes can not "jump the queue". When
+-- the writer finishes, then all other waiting threads will race for write
+-- access again.
+--
+-- TODO: unsafeReleaseWriteAccess will set RWState to Reading 0. In case we have
+-- readers *and* writers waiting for a writer to finish, once the writer is
+-- finished there will be a race. In this race, readers and writers are just as
+-- likely to acquire access first. However, if we wanted to make RWVar even more
+-- biased towards writers, then we could ensure that all waiting writers get
+-- access before the readers get a chance. This would probably require us to
+-- change RWState to represent the case where writers are waiting for a writer
+-- to finish.
+unsafeAcquireWriteAccess :: (MonadSTM m, MonadCatch m) => RWVar m a -> m a
+unsafeAcquireWriteAccess (RWVar !var) =
+    -- trySetWriting is interruptible, but it is fine if it is interrupted
+    -- because the RWState can not be changed before the interruption.
     --
-    -- If multiple threads try to acquire write access concurrently, then they
-    -- will race for access. Even if a thread has set RWState to WaitingToWrite,
-    -- there is no guarantee that the same thread will acquire write access
-    -- first when all readers have finished. However, if the writer has
-    -- finished, then all other waiting threads will try to get write again
-    -- until they succeed.
+    -- trySetWriting might update the RWState. There are interruptible
+    -- operations in the body of the bracketOnError (in waitToWrite), so async
+    -- exceptions can be delivered there. If an async exception happens because
+    -- of an interrupt, we undo the RWState change using undoWaitingToWrite.
+    --
+    -- Note that if waitToWrite is interrupted, that it is impossible for the
+    -- RWState to have changed from WaitingToWrite to either Reading or Writing.
+    -- Therefore, undoWaitingToWrite can assume that it will find WaitingToWrite
+    -- in the lock.
+    bracketOnError trySetWriting undoWaitingToWrite $
+      -- When Nothing is returned, it means that we set the RWState to
+      -- WaitingToWrite, and so we wait to acquire the final write access.
+      --
+      -- When Just is returned, we already have write access.
+      maybe waitToWrite pure
+  where
+    -- Try to acquire a write lock immediately, or otherwise set the internal
+    -- state to WaitingToWrite as soon as possible.
+    --
+    -- Note: this is interruptible
+    trySetWriting = atomically $ readTVar var >>= \case
+        Reading n x
+          | n == 0 -> do
+              writeTVar var Writing
+              pure (Just x)
+          | otherwise -> do
+              writeTVar var (WaitingToWrite n x)
+              pure Nothing
+        -- The following two branches are interruptible
+        WaitingToWrite _n _x -> retry
+        Writing -> retry
+
+    -- Note: this is uninterruptible
+    undoWaitingToWrite Nothing  =  atomically $ readTVar var >>= \case
+        Reading _n _x -> error "undoWaitingToWrite: found Reading but expected WaitingToWrite"
+        WaitingToWrite n x -> writeTVar var (Reading n x)
+        Writing -> error "undoWaitingToWrite: found Writing but expected WaitingToWrite"
+    undoWaitingToWrite (Just _) = error "undoWaitingToWrite: found Just but expected Nothing"
+
+    -- Wait for the number of readers to go to 0, and then finally acquire write
+    -- access.
     --
-    -- TODO: unsafeReleaseWriteAccess will set RWState to Reading 0. In case we
-    -- have readers *and* writers waiting for a writer to finish, once the
-    -- writer is finished there will be a race. In this race, readers and
-    -- writers are just as likely to acquire access first. However, if we wanted
-    -- to make RWVar even more biased towards writers, then we could ensure that
-    -- all waiting writers get access before the readers get a chance. This
-    -- would probably require us to change RWState to represent the case where
-    -- writers are waiting for a writer to finish.
-    case may of
-      Nothing -> unsafeAcquireWriteAccess rw
-      Just x  -> pure x
+    -- Note: this is interruptible
+    waitToWrite = atomically $ readTVar var >>= \case
+        Reading _n _x -> error "waitToWrite: found Reading but expected WaitingToWrite"
+        WaitingToWrite n x
+          | n == 0 -> do
+              writeTVar var Writing
+              pure x
+          -- This branch is interruptible
+          | otherwise -> retry
+        Writing -> error "waitToWrite: found Reading but expected Writing"
 
 {-# SPECIALISE unsafeReleaseWriteAccess :: RWVar IO a -> a -> STM IO () #-}
+-- | Release write access. This function assumes that it runs in a masked
+-- context, and that is properly paired with an 'unsafeAcquireWriteAccess'!
 unsafeReleaseWriteAccess :: MonadSTM m => RWVar m a -> a -> STM m ()
 unsafeReleaseWriteAccess (RWVar !var) x = do
     readTVar var >>= \case

test-control/Test/Control/Concurrent/Class/MonadSTM/RWVar.hs

@@ -3,7 +3,8 @@ module Test.Control.Concurrent.Class.MonadSTM.RWVar (tests) where
 import qualified Control.Concurrent.Class.MonadSTM.RWVar as RW
 import           Control.Monad.Class.MonadAsync
 import           Control.Monad.Class.MonadSay (MonadSay (say))
-import           Control.Monad.Class.MonadTimer
+import           Control.Monad.Class.MonadTest (MonadTest (exploreRaces))
+import           Control.Monad.Class.MonadThrow
 import           Control.Monad.IOSim
 import           Test.QuickCheck
 import           Test.Tasty
@@ -17,44 +18,67 @@ tests = testGroup "Test.Control.Concurrent.Class.MonadSTM.RWVar" [
 data Action a = Read a | Incr a
   deriving stock (Show, Eq, Read)
 
-instance Arbitrary (Action (Small Word, Small Word)) where
-  arbitrary = do
-    n <- oneof [ arbitrary, Small <$> choose (1, 5) ]
-    m <- oneof [ arbitrary, Small <$> choose (1, 5) ]
-    elements [ Read (n, m), Incr (n, m) ]
+instance Arbitrary a => Arbitrary (Action a) where
+  arbitrary = oneof [
+        Read <$> arbitrary
+      , Incr <$> arbitrary
+      ]
+  shrink (Read x) = [Read x' | x' <- shrink x]
+  shrink (Incr x) = [Incr x' | x' <- shrink x]
 
 -- | Performing reads and increments on an @'RWVar' Int@ in parallel should not
--- lead to races. We observe this by looking at the trace and seeing that the
--- value inside the 'RWVar' increases monotonically.
-prop_noRace :: [Action (Small Word, Small Word)] -> Property
-prop_noRace as = exploreSimTrace id prop (\_ st -> sayMonotonic 0 st)
+-- lead to races. We let IOSimPOR check that there are no deadlocks. We also
+-- look at the trace to see if the value inside the 'RWVar' increases
+-- monotonically.
+prop_noRace ::
+     Action ()
+  -> Action ()
+  -> Action ()
+  -> Property
+prop_noRace a1 a2 a3 = exploreSimTrace modifyOpts prop $ \_ tr ->
+    case traceResult False tr of
+      Left e -> counterexample (show e) (property False)
+      _      -> propSayMonotonic tr
   where
+    modifyOpts = id
+
     prop :: IOSim s ()
     prop = do
+      exploreRaces
       var <- RW.new (0 :: Int)
-      forConcurrently_ as $ \case
-        Read (Small n, Small m) -> do
-          threadDelay (fromIntegral n)
-          RW.withReadAccess var $ \x -> do
-            threadDelay (fromIntegral m)
-            say (show (Read x))
-        Incr (Small n, Small m) -> do
-          threadDelay (fromIntegral n)
-          RW.withWriteAccess_ var $ \x -> do
-            threadDelay (fromIntegral m)
-            let x' = x + 1
-            say (show (Incr x'))
-            pure x'
-
-    sayMonotonic ::  Int -> SimTrace () -> Property
-    sayMonotonic _ (Nil MainReturn{}) = property True
-    sayMonotonic prev (Cons se st') = (case se of
-        SimPOREvent{seType} -> case seType of
-          EventSay s -> case read s of
-            Read x -> counterexample "Unexpected Read result"
-                    $ prev     === x .&&. sayMonotonic x st'
-            Incr x -> counterexample "Unexpected Write result"
-                    $ prev + 1 === x .&&. sayMonotonic x st'
-          _ -> sayMonotonic prev st'
-        _ -> property False)
-    sayMonotonic _ _ = property False
+      let g = \case
+            Read () ->
+              RW.withReadAccess var $ \x -> do
+                say (show (Read x))
+            Incr () ->
+              RW.withWriteAccess_ var $ \x -> do
+                let x' = x + 1
+                say (show (Incr x'))
+                pure x'
+      t1 <- async (g a1)
+      t2 <- async (g a2)
+      t3 <- async (g a3)
+      async (cancel t1) >>= wait
+      (_ :: Either AsyncCancelled ()) <- try (wait t1)
+      (_ :: Either AsyncCancelled ()) <- try (wait t2)
+      (_ :: Either AsyncCancelled ()) <- try (wait t3)
+      pure ()
+
+propSayMonotonic :: SimTrace () -> Property
+propSayMonotonic simTrace = propResultsMonotonic (actionResults simTrace)
+
+propResultsMonotonic :: [Action Int] -> Property
+propResultsMonotonic as =
+    counterexample
+      ("Action results are not monotonic: " ++ show as)
+      (resultsMonotonic as)
+
+resultsMonotonic :: [Action Int] -> Bool
+resultsMonotonic = go 0
+  where
+    go _ []               = True
+    go prev (Read x : as) = prev == x     && go x as
+    go prev (Incr x : as) = prev + 1 == x && go x as
+
+actionResults :: SimTrace () -> [Action Int]
+actionResults = map read . selectTraceEventsSay