{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnliftedFFITypes #-}

-- | A reentrant mutex, so the thread that owns a connection's active response
-- can issue another request (which cancels the first) without deadlocking.
--
-- The owner is recorded by its thread /number/, not its 'ThreadId'.  Holding
-- a 'ThreadId' would keep the owner's thread object reachable from the
-- connection, so an owner that blocks forever (on an 'MVar' nobody else can
-- fill, say) could never be found unreachable by the garbage collector: it
-- would not receive 'Control.Exception.BlockedIndefinitelyOnMVar', never
-- release the mutex, and every other user of the connection would hang with
-- it.  hpgsql uses a weak 'ThreadId' for the same reason; a number also works
-- after the collector has declared the thread dead but resurrected it to run
-- its exception handlers, which may re-enter the mutex (a rollback, say).
--
-- Adapted from @Hpgsql.Locking@ in hpgsql
-- (<https://github.com/mzabani/hpgsql>), Copyright Marcelo Zabani (c) 2026,
-- BSD-3-Clause.
module Hmssql.Connection.Locking
  ( Mutex
  , newMutex
  , withMutex
  , mutexOwner
  ) where

import Control.Concurrent (ThreadId, myThreadId)
import Control.Concurrent.STM
import Control.Exception (bracket_)
import Data.Functor ((<&>))
import Data.Word (Word64)

#if MIN_VERSION_base(4,19,0)
import GHC.Conc.Sync (fromThreadId)
#else
import Foreign.C.Types (CULLong (..))
import GHC.Conc.Sync (ThreadId (..))
import GHC.Exts (ThreadId#)
#endif

newtype Mutex = Mutex (TVar Owner)

-- | Strict, so that no thunk referring to the owner's 'ThreadId' is stored.
data Owner
  = Unowned
  | Owner !Word64
    -- ^ The owner's thread number.  Nested entries don't count themselves:
    -- only the outermost one takes and releases the mutex.

newMutex :: IO Mutex
newMutex :: IO Mutex
newMutex = TVar Owner -> Mutex
Mutex (TVar Owner -> Mutex) -> IO (TVar Owner) -> IO Mutex
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Owner -> IO (TVar Owner)
forall a. a -> IO (TVar a)
newTVarIO Owner
Unowned

-- | Run the action holding the mutex, waiting for any other owner to leave.
-- Waiting is interruptible; the mutex is released however the action ends.
--
-- Reentry takes a fast path: when this thread holds the mutex already, the
-- action just runs.  Nothing needs undoing afterwards, since the entry that
-- took the mutex releases it, and nothing can change the owner under us:
-- only the owner writes the variable while it is owned.  This saves two
-- STM transactions and an exception frame on every nested use (every token
-- read inside 'Hmssql.Query.query', say).
withMutex :: Mutex -> IO a -> IO a
withMutex :: forall a. Mutex -> IO a -> IO a
withMutex (Mutex TVar Owner
var) IO a
act = do
  !me <- ThreadId -> Word64
threadNumber (ThreadId -> Word64) -> IO ThreadId -> IO Word64
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO ThreadId
myThreadId
  readTVarIO var >>= \case
    Owner Word64
owner | Word64
owner Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
== Word64
me -> IO a
act
    Owner
_ -> IO () -> IO () -> IO a -> IO a
forall a b c. IO a -> IO b -> IO c -> IO c
bracket_ (Word64 -> IO ()
acquire Word64
me) (Word64 -> IO ()
release Word64
me) IO a
act
  where
    -- Another thread's (this one's, it was just seen not to be).
    acquire :: Word64 -> IO ()
acquire Word64
me = STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar Owner -> STM Owner
forall a. TVar a -> STM a
readTVar TVar Owner
var STM Owner -> (Owner -> STM ()) -> STM ()
forall a b. STM a -> (a -> STM b) -> STM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
      Owner
Unowned -> TVar Owner -> Owner -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar Owner
var (Word64 -> Owner
Owner Word64
me)
      Owner Word64
_ -> STM ()
forall a. STM a
retry
    release :: Word64 -> IO ()
release Word64
me = STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar Owner -> STM Owner
forall a. TVar a -> STM a
readTVar TVar Owner
var STM Owner -> (Owner -> STM ()) -> STM ()
forall a b. STM a -> (a -> STM b) -> STM b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
      Owner Word64
owner | Word64
owner Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
== Word64
me -> TVar Owner -> Owner -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar Owner
var Owner
Unowned
      Owner
_ -> () -> STM ()
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- | The number of the thread holding the mutex (as in 'show' of its
-- 'ThreadId').  For diagnostics and tests.
mutexOwner :: Mutex -> IO (Maybe Word64)
mutexOwner :: Mutex -> IO (Maybe Word64)
mutexOwner (Mutex TVar Owner
var) = TVar Owner -> IO Owner
forall a. TVar a -> IO a
readTVarIO TVar Owner
var IO Owner -> (Owner -> Maybe Word64) -> IO (Maybe Word64)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \case
  Owner
Unowned -> Maybe Word64
forall a. Maybe a
Nothing
  Owner Word64
owner -> Word64 -> Maybe Word64
forall a. a -> Maybe a
Just Word64
owner

threadNumber :: ThreadId -> Word64
#if MIN_VERSION_base(4,19,0)
threadNumber :: ThreadId -> Word64
threadNumber = ThreadId -> Word64
fromThreadId
#else
threadNumber (ThreadId t) = fromIntegral (rtsGetThreadId t)

foreign import ccall unsafe "rts_getThreadId"
  rtsGetThreadId :: ThreadId# -> CULLong
#endif