RWLock.hs

{-# OPTIONS_GHC -Wall #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  RWLock
-- Copyright   :  (c) Masahiro Sakai 2023
-- License     :  BSD-3-Clause
--
-- Simple implement of various variants of RWLocks.
--
-- References:
--
-- * <https://en.wikipedia.org/wiki/Readers%E2%80%93writers_problem>
--
-- * <https://en.wikipedia.org/wiki/Readers%E2%80%93writer_lock>
--
-- * <https://hackage.haskell.org/package/concurrent-extra>
--
-- * <https://pypi.org/project/readerwriterlock/>
--
-----------------------------------------------------------------------------

module RWLock
  (
  -- * Basic typees
    IsRWLock (..)
  , RWLockState (..)

  -- * MVar-based implementation
  , RWLockReadMVar
  , RWLockWriteMVar
  , RWLockFairMVar

  -- * STM-based implementation
  , RWLockReadSTM
  , RWLockWriteSTM
  , RWLockFairSTM
  ) where

import Control.Concurrent.MVar
import Control.Concurrent.STM
import Control.Exception
import Control.Monad

-- ------------------------------------------------------------------------

class IsRWLock rw where
  newRWLock :: IO rw
  acquireRead :: rw -> IO ()
  releaseRead :: rw -> IO ()
  acquireWrite :: rw -> IO ()
  releaseWrite :: rw -> IO ()
  withRead :: rw -> IO b -> IO b
  withWrite :: rw -> IO b -> IO b
  withRead rw = bracket_ (acquireRead rw) (releaseRead rw)
  withWrite rw = bracket_ (acquireWrite rw) (releaseWrite rw)

data RWLockState
  = Unlocked
  | RLocked !Int
  | WLocked
  deriving (Eq)

-- ------------------------------------------------------------------------

-- | MVar-based RWLock with reader priority (aka first readers-writers problem)
data RWLockReadMVar = RWLockReadMVar !(MVar ()) !(MVar Int)

instance IsRWLock RWLockReadMVar where
  newRWLock = do
    resource <- newMVar ()
    rmutex <- newMVar 0
    return (RWLockReadMVar resource rmutex)

  acquireRead (RWLockReadMVar resource rmutex) = do
    modifyMVarMasked_ rmutex $ \i -> do
      -- takeMVar can be blocked thus interruptible
      when (i == 0) $ takeMVar resource
      return $! i+1

  releaseRead (RWLockReadMVar resource rmutex) = do
    modifyMVarMasked_ rmutex $ \i -> do
      -- resource should be empty, thus putMVar is uninterruptible
      when (i == 1) $ putMVar resource ()
      return $! i-1

  acquireWrite (RWLockReadMVar resource _rmutex) = takeMVar resource

  releaseWrite (RWLockReadMVar resource _rmutex) = putMVar resource ()

-- ------------------------------------------------------------------------

-- | MVar-based RWLock with writer priority (aka second readers-writers problem)
data RWLockWriteMVar = RWLockWriteMVar !(MVar Int) !(MVar Int) !(MVar ()) !(MVar ())

instance IsRWLock RWLockWriteMVar where
  newRWLock = do
    rmutex <- newMVar 0
    wmutex <- newMVar 0
    readTry <- newMVar ()
    resource <- newMVar ()
    return (RWLockWriteMVar rmutex wmutex readTry resource)

  acquireRead (RWLockWriteMVar rmutex _wmutex readTry resource) = do
    withMVar readTry $ \_ -> do
      modifyMVarMasked_ rmutex $ \i -> do
        -- takeMVar can be blocked thus interruptible
        when (i == 0) $ takeMVar resource
        return $! i+1

  releaseRead (RWLockWriteMVar rmutex _wmutex _readTry resource) = do
    modifyMVarMasked_ rmutex $ \i -> do
      -- resource should be empty, thus putMVar is uninterruptible
      when (i == 1) $ putMVar resource ()
      return $! i-1

  acquireWrite (RWLockWriteMVar _rmutex wmutex readTry resource) = do
    modifyMVarMasked_ wmutex $ \i -> do
      if i == 0 then do
        takeMVar readTry
        takeMVar resource `onException` putMVar readTry ()
      else do
        takeMVar resource
      return $! i+1

  releaseWrite (RWLockWriteMVar _rmutex wmutex readTry resource) = do
    modifyMVarMasked_ wmutex $ \i -> do
      -- resource should be empty, thus putMVar is uninterruptible
      putMVar resource ()
      -- readTry should be empty, thus putMVar is uninterruptible
      when (i == 1) $ putMVar readTry ()
      return $! i-1

-- ------------------------------------------------------------------------

-- | MVar-based RWLock with fair priority (aka third readers-writers problem)
data RWLockFairMVar = RWLockFairMVar !RWLockReadMVar !(MVar ())

instance IsRWLock RWLockFairMVar where
  newRWLock = do
    base <- newRWLock
    serviceQueue <- newMVar ()  -- assume that this MVar is fair
    return (RWLockFairMVar base serviceQueue)

  acquireRead (RWLockFairMVar base serviceQueue) = do
    withMVar serviceQueue $ \_ -> do
      acquireRead base

  releaseRead (RWLockFairMVar base _serviceQueue) = do
    releaseRead base

  acquireWrite (RWLockFairMVar base serviceQueue) = do
    withMVar serviceQueue $ \_ -> acquireWrite base

  releaseWrite (RWLockFairMVar base _serviceQueue) = do
    releaseWrite base

-- ------------------------------------------------------------------------

-- | STM-based RWLock with reader priority (aka first readers-writers problem)
newtype RWLockReadSTM = RWLockReadSTM (TVar RWLockState)

instance IsRWLock RWLockReadSTM where
  newRWLock = do
    tv <- newTVarIO Unlocked
    return $ RWLockReadSTM tv

  acquireRead (RWLockReadSTM tv) = atomically $ do
    st <- readTVar tv
    case st of
      Unlocked -> writeTVar tv (RLocked 0)
      RLocked n -> writeTVar tv (RLocked (n+1))
      WLocked -> retry

  releaseRead (RWLockReadSTM tv) = atomically $ do
    st <- readTVar tv
    case st of
      RLocked 1 -> writeTVar tv Unlocked
      RLocked n -> writeTVar tv (RLocked (n-1))
      _ -> undefined

  acquireWrite (RWLockReadSTM tv) = atomically $ do
    st <- readTVar tv
    guard $ st == Unlocked
    writeTVar tv WLocked

  releaseWrite (RWLockReadSTM tv) = atomically $ do
    st <- readTVar tv
    case st of
      WLocked -> writeTVar tv Unlocked
      _ -> undefined

-- ------------------------------------------------------------------------

-- | STM-based RWLock with writer priority (aka second readers-writers problem)
data RWLockWriteSTM = RWLockWriteSTM (TVar RWLockState) (TVar Int)

instance IsRWLock RWLockWriteSTM where
  newRWLock = do
    tv <- newTVarIO Unlocked
    writers <- newTVarIO 0
    return $ RWLockWriteSTM tv writers

  acquireRead (RWLockWriteSTM tv writers) = atomically $ do
    m <- readTVar writers
    guard $ m == 0
    st <- readTVar tv
    case st of
      Unlocked -> writeTVar tv (RLocked 0)
      RLocked n -> writeTVar tv (RLocked (n+1))
      WLocked -> retry

  releaseRead (RWLockWriteSTM tv _writers) = atomically $ do
    st <- readTVar tv
    case st of
      RLocked 1 -> writeTVar tv Unlocked
      RLocked n -> writeTVar tv (RLocked (n-1))
      _ -> undefined

  acquireWrite (RWLockWriteSTM tv writers) = mask_ $ do
    atomically (modifyTVar writers (+1))
    let body = do
          st <- readTVar tv
          guard $ st == Unlocked
          writeTVar tv WLocked
    atomically body `onException` atomically (modifyTVar writers (subtract 1))

  releaseWrite (RWLockWriteSTM tv writers) = atomically $ do
    modifyTVar writers (subtract 1)
    st <- readTVar tv
    case st of
      WLocked -> writeTVar tv Unlocked
      _ -> undefined

-- ------------------------------------------------------------------------

-- | STM-based RWLock with fair priority (aka third readers-writers problem)
--
-- It uses MVar for guaranteeing fairness.
data RWLockFairSTM = RWLockFairSTM !RWLockReadSTM !(MVar ())

instance IsRWLock RWLockFairSTM where
  newRWLock = do
    base <- newRWLock
    serviceQueue <- newMVar ()  -- assume that this MVar is fair
    return (RWLockFairSTM base serviceQueue)

  acquireRead (RWLockFairSTM base serviceQueue) = do
    withMVar serviceQueue $ \_ -> do
      acquireRead base

  releaseRead (RWLockFairSTM base _serviceQueue) = do
    releaseRead base

  acquireWrite (RWLockFairSTM base serviceQueue) = do
    withMVar serviceQueue $ \_ -> do
      acquireWrite base

  releaseWrite (RWLockFairSTM base _serviceQueue) = do
    releaseWrite base

-- ------------------------------------------------------------------------

TestRWLock.hs

module Main where

import Control.Concurrent
import Control.Concurrent.Async
import Control.Monad
import Data.Proxy
import Data.IORef
import RWLock

import Test.Tasty
import Test.Tasty.HUnit

main :: IO ()
main = defaultMain tests

tests :: TestTree
tests = testGroup "Tests" [unitTests]

unitTests :: TestTree
unitTests = testGroup "Unit tests"
  [ testCase "RWLockReadMVar reader do not starve" $ do
      b <- readerStarvationScenario (Proxy :: Proxy RWLockReadMVar)
      b @?= True
  , testCase "RWLockReadMVar writer starves" $ do
      b <- writerStarvationScenario (Proxy :: Proxy RWLockReadMVar)
      b @?= False

   -- , testCase "RWLockReadSTM reader do not starve" $ do
   --     b <- readerStarvationScenario (Proxy :: Proxy RWLockReadSTM)
   --     b @?= True
  , testCase "RWLockReadSTM writer starves" $ do
      b <- writerStarvationScenario (Proxy :: Proxy RWLockReadSTM)
      b @?= False

  , testCase "RWLockWriteMVar reader starves" $ do
      b <- readerStarvationScenario (Proxy :: Proxy RWLockWriteMVar)
      b @?= False
  , testCase "RWLockWriteMVar writer do not starve" $ do
      b <- writerStarvationScenario (Proxy :: Proxy RWLockWriteMVar)
      b @?= True

  , testCase "RWLockWriteSTM reader starves" $ do
      b <- readerStarvationScenario (Proxy :: Proxy RWLockWriteSTM)
      b @?= False
  , testCase "RWLockWriteSTM writer do not starve" $ do
      b <- writerStarvationScenario (Proxy :: Proxy RWLockWriteSTM)
      b @?= True

  , testCase "RWLockFairMVar reader do not starve" $ do
      b <- readerStarvationScenario (Proxy :: Proxy RWLockFairMVar)
      b @?= True
  , testCase "RWLockFairMVar writer do not starve" $ do
      b <- writerStarvationScenario (Proxy :: Proxy RWLockFairMVar)
      b @?= True

  , testCase "RWLockFairSTM reader do not starve" $ do
      b <- readerStarvationScenario (Proxy :: Proxy RWLockFairSTM)
      b @?= True
  , testCase "RWLockFairSTM writer do not starve" $ do
      b <- writerStarvationScenario (Proxy :: Proxy RWLockFairSTM)
      b @?= True
  ]

readerStarvationScenario :: forall lock. IsRWLock lock => Proxy lock -> IO Bool
readerStarvationScenario _ = do
  (lock :: lock) <- newRWLock
  ref <- newIORef False
  let w1 = forever $ withWrite lock $ threadDelay (120*1000)
      w2 = threadDelay (40*1000) >> w1
      r1 = threadDelay (80*1000) >> withRead lock (writeIORef ref True)
  withAsync w1 $ \_ ->
    withAsync w2 $ \_ ->
      withAsync r1 $ \_ -> do
        threadDelay (400*1000)
  readIORef ref

writerStarvationScenario :: forall lock. IsRWLock lock => Proxy lock -> IO Bool
writerStarvationScenario _ = do
  (lock :: lock) <- newRWLock
  ref <- newIORef False
  let r1 = forever $ withRead lock $ threadDelay (40*1000)
      r2 = threadDelay (20*1000) >> r1
      w1 = threadDelay (20*1000) >> withWrite lock (writeIORef ref True)
  withAsync r1 $ \_ ->
    withAsync r2 $ \_ ->
      withAsync w1 $ \_ -> do
        threadDelay (400*1000)
  readIORef ref