StrictRWLockSpec.hs

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

module Test.LinearLocks.StrictRWLockSpec where

import Control.Concurrent.MVar qualified as MVar
import Control.Concurrent.ReadWriteLock qualified as Conc
import Control.Exception (SomeException, throwIO, try)
import Control.Functor.Linear qualified as L
import Control.Monad (when)
import Control.Monad.IO.Class.Linear qualified as L
import Data.IORef qualified as IORef
import LinearLocks
import LinearLocks.Internal qualified as Internal
import LinearLocks.Internal.StrictRWLock qualified as Internal
import LinearLocks.RWLock.Strict qualified as RWLock
import Prelude.Linear (Ur (..))
import Prelude.Linear qualified as L hiding (IO)
import System.IO.Resource.Linear (RIO)
import Test.Hspec.Expectations.Pretty (errorCall, shouldThrow)
import "tasty-hunit-compat" Test.Tasty.HUnit

-- | 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
-- ...
unit_read_mutex :: IO ()
unit_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 @?= "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 @?= "hello"

unit_write_mutex :: IO ()
unit_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 @?= "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 @?= "world"

  str <- IORef.readIORef rwl.var
  str.unNF @?= "world"

unit_realeases_ioref_in_read_mode :: IO ()
unit_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

unit_realeases_ioref_in_write_mode :: IO ()
unit_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

unit_rolls_back_on_exception :: IO ()
unit_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.unNF @?= "hello"

unit_rolls_back_on_imprecise_exception :: IO ()
unit_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.unNF @?= "hello"

unit_new_evaluates_value_to_normal_form :: IO ()
unit_new_evaluates_value_to_normal_form = do
  RWLock.new @[Int] 0 [1, 2, error "oops", 4]
    `shouldThrow` errorCall "oops"

unit_release_evaluates_value_to_normal_form :: IO ()
unit_release_evaluates_value_to_normal_form = do
  mutex <- RWLock.new @[Int] 0 [1]

  logs <- MVar.newMVar @[String] []
  let logMsg :: String -> RIO ()
      logMsg msg = L.liftSystemIO do
        MVar.modifyMVar_ logs \logs -> pure (logs <> [msg])

  let run =
        lockScope \key -> L.do
          (mg, key) <- RWLock.acquireWrite key mutex
          logMsg "ran 'acquire'"
          mg <- RWLock.write mg [1, 2, error "oops", 4]
          logMsg "ran 'write'"
          RWLock.releaseWrite mg
          logMsg "ran 'release'"
          dropKeyAndReturn key ()

  run `shouldThrow` errorCall "oops"

  -- The exception should be thrown WHILE running `release`.
  -- `write` should NOT throw.
  msgs <- MVar.takeMVar logs
  msgs @?= ["ran 'acquire'", "ran 'write'"]

assertCanRead :: RWLock.RWLock lvl a -> Bool -> IO ()
assertCanRead rwl expected = do
  canRead <- Conc.tryAcquireRead rwl.lock
  canRead @?= 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 @?= expected
  -- Release the lock if it was acquired.
  when canWrite do
    Conc.releaseWrite rwl.lock