RWLockSpec.hs

{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE QualifiedDo #-}
{-# LANGUAGE NoFieldSelectors #-}

module Test.LinearLocks.RWLockSpec where

import Control.Concurrent.ReadWriteLock qualified as Conc
import Control.Exception (SomeException, throwIO, try)
import Control.Functor.Linear qualified as L
import Control.Monad (void, when)
import Control.Monad.IO.Class.Linear qualified as L
import Data.IORef qualified as IORef
import LinearLocks
import LinearLocks.Internal.RWLock qualified as Internal
import LinearLocks.RWLock qualified as RWLock
import Prelude.Linear (Ur (..))
import Prelude.Linear qualified as L hiding (IO)
import Test.Syd

-- | Doctests
--
-- >>> :{
-- >>> unit_mutexes_cannot_be_locked_in_wrong_order :: IO ()
-- >>> unit_mutexes_cannot_be_locked_in_wrong_order = do
-- >>>   m1 <- RWLock.new 2 "hello"
-- >>>   m2 <- RWLock.new 4 "world"
-- >>>   lockScope \key -> L.do
-- >>>     (g2, key) <- RWLock.acquireRead key m2
-- >>>     (g1, key) <- RWLock.acquireRead key m1
-- >>>     RWLock.releaseRead g1
-- >>>     RWLock.releaseRead g2
-- >>>     dropKeyAndReturn key ()
-- >>> :}
-- ...
-- ... • Cannot satisfy: 5 <= 2
-- ... • In a stmt of a 'do' block: (g1, key) <- RWLock.acquireRead key m1
-- ...
spec :: Spec
spec = describe "RWLock" do
  it "read mutex" do
    rwl <- RWLock.new 0 "hello"
    -- Read in "read mode"
    str <- lockScope \key -> L.do
      (guard, key) <- RWLock.acquireRead key rwl
      (Ur str, guard) <- RWLock.read guard
      RWLock.releaseRead guard
      dropKeyAndReturn key str
    str `shouldBe` "hello"

    -- Read in "write mode"
    str <- lockScope \key -> L.do
      (guard, key) <- RWLock.acquireWrite key rwl
      (Ur str, guard) <- RWLock.read guard
      RWLock.releaseWrite guard
      dropKeyAndReturn key str
    str `shouldBe` "hello"

  it "write mutex" do
    rwl <- RWLock.new 0 "hello"

    -- Write in "write mode"
    lockScope \key -> L.do
      (guard, key) <- RWLock.acquireWrite key rwl
      guard <- RWLock.write guard "world"
      RWLock.releaseWrite guard
      dropKeyAndReturn key ()

    -- Read in "read mode"
    str <- lockScope \key -> L.do
      (guard, key) <- RWLock.acquireRead key rwl
      (Ur str, guard) <- RWLock.read guard
      RWLock.releaseRead guard
      dropKeyAndReturn key str
    str `shouldBe` "world"

    -- Read in "write mode"
    str <- lockScope \key -> L.do
      (guard, key) <- RWLock.acquireWrite key rwl
      (Ur str, guard) <- RWLock.read guard
      RWLock.releaseWrite guard
      dropKeyAndReturn key str
    str `shouldBe` "world"

    str <- IORef.readIORef rwl.var
    str `shouldBe` "world"

  it "realeases ioref in read mode" do
    rwl <- RWLock.new 0 "hello"
    lockScope \key -> L.do
      (mg, key) <- RWLock.acquireRead key rwl

      -- If the lock was acquired in "read mode",
      -- we shouldn't be able to acquire it again in "write mode",
      -- but we should be able to acquire it in "read mode".
      L.liftSystemIO do
        assertCanRead rwl True
        assertCanWrite rwl False

      RWLock.releaseRead mg

      --  The lock was released, we should be able to acquire it in both "read mode" and "write mode".
      L.liftSystemIO do
        assertCanRead rwl True
        assertCanWrite rwl True

      dropKeyAndReturn key ()

    --  The lock was released, we should be able to acquire it in both "read mode" and "write mode".
    assertCanRead rwl True
    assertCanWrite rwl True

  it "realeases ioref in write mode" do
    rwl <- RWLock.new 0 "hello"
    lockScope \key -> L.do
      (mg, key) <- RWLock.acquireWrite key rwl

      -- If the lock was acquired in "write mode",
      -- we shouldn't be able to acquire it again in "write mode" or "read mode".
      L.liftSystemIO do
        assertCanRead rwl False
        assertCanWrite rwl False

      RWLock.releaseWrite mg

      --  The lock was released, we should be able to acquire it in both "read mode" and "write mode".
      L.liftSystemIO do
        assertCanRead rwl True
        assertCanWrite rwl True

      dropKeyAndReturn key ()

    --  The lock was released, we should be able to acquire it in both "read mode" and "write mode".
    assertCanRead rwl True
    assertCanWrite rwl True

  it "rolls back on exception" do
    rwl <- RWLock.new 0 "hello"
    Left _ <- try @SomeException $ lockScope \key -> L.do
      (mg, key) <- RWLock.acquireWrite key rwl
      mg <- RWLock.write mg "world"
      L.liftSystemIO L.$ throwIO (userError "oops")
      RWLock.releaseWrite mg
      dropKeyAndReturn key ()

    -- The IORef should have been released, and the original value should have been put back into the IORef.
    assertCanRead rwl True
    assertCanWrite rwl True
    mbResult <- IORef.readIORef rwl.var
    mbResult `shouldBe` "hello"

  it "rolls back on imprecise exception" do
    rwl <- RWLock.new 0 "hello"
    Left _ <- try @SomeException $ lockScope \key -> L.do
      (mg, key) <- RWLock.acquireWrite key rwl
      mg <- RWLock.write mg "world"
      error "err"
      RWLock.releaseWrite mg
      dropKeyAndReturn key ()

    -- The IORef should have been released, and the original value should have been put back into the IORef.
    assertCanRead rwl True
    assertCanWrite rwl True
    mbResult <- IORef.readIORef rwl.var
    mbResult `shouldBe` "hello"

  it "new doesn't evaluate value to normal form" do
    -- This should not throw, the "error" thunk should not be evaluated
    void $ RWLock.new @[Int] 0 [1, 2, error "oops", 4]

  it "release doesn't evaluate value to normal form" do
    mutex <- RWLock.new @[Int] 0 [1]

    lockScope \key -> L.do
      (mg, key) <- RWLock.acquireWrite key mutex
      -- This should not throw, the "error" thunk should not be evaluated
      mg <- RWLock.write mg [1, 2, error "oops", 4]
      -- This should not throw
      RWLock.releaseWrite mg
      dropKeyAndReturn key ()

assertCanRead :: RWLock.RWLock lvl a -> Bool -> IO ()
assertCanRead rwl expected = do
  canRead <- Conc.tryAcquireRead rwl.lock
  canRead `shouldBe` expected
  -- Release the lock if it was acquired.
  when canRead do
    Conc.releaseRead rwl.lock

assertCanWrite :: RWLock.RWLock lvl a -> Bool -> IO ()
assertCanWrite rwl expected = do
  canWrite <- Conc.tryAcquireWrite rwl.lock
  canWrite `shouldBe` expected
  -- Release the lock if it was acquired.
  when canWrite do
    Conc.releaseWrite rwl.lock