semaphore-compat-2.0.0: src/System/Semaphore/Internal/Posix.hs
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
module System.Semaphore.Internal.Posix
( ClientSemaphore(..), ServerSemaphore(..)
, SemaphoreToken(..)
, create_sem, open_sem_raw
, waitOnSemaphore, tryWaitOnSemaphore
, releaseSemaphoreToken
, destroyClientSemaphore, destroyServerSemaphore
, getSemaphoreValue
, getTimeSeed
) where
-- base
import Control.Concurrent
( ThreadId, forkIOWithUnmask, killThread )
import Control.Concurrent.MVar
( MVar, mkWeakMVar, newEmptyMVar, newMVar, putMVar
, readMVar, takeMVar, tryTakeMVar )
import Control.Exception ( IOException )
import Control.Monad
import Data.Word ( Word8 )
import Data.Bits ( xor )
import Foreign.C.Error ( Errno(Errno), eCONNABORTED )
import GHC.Clock ( getMonotonicTimeNSec )
import GHC.IO.Exception ( ioe_errno )
import GHC.Stack ( HasCallStack )
import System.IO.Error ( isFullError )
-- exceptions
import qualified Control.Monad.Catch as MC
-- stm
import Control.Concurrent.STM
( TVar, atomically, newTVarIO, readTVar, readTVarIO
, modifyTVar', writeTVar, retry )
-- directory
import System.Directory ( doesPathExist )
-- unix
import System.Posix.IO ( closeFd, createPipe )
import System.Posix.Files ( removeLink )
import System.Posix.Types ( Fd )
import System.Posix.Process ( getProcessID )
import System.Semaphore.Internal.Common
import System.Semaphore.Internal.DomainSocket
( connectDomainSocket, listenDomainSocket
, pollAcceptSocket, AcceptResult(..)
, fdReadByte, fdWriteByte
, fdShutdown )
-- | A semaphore identity (name + socket path).
-- Each operation that needs a connection opens one internally.
data ClientSemaphore =
ClientSemaphore
{ clientSemaphoreName :: !SemaphoreName
, semSocketPath :: !FilePath
}
-- | A held semaphore token, bound to one acquired resource.
--
-- If all references to the 'SemaphoreToken' are dropped without being
-- released, a finalizer closes the underlying connection and the server
-- returns the token to the pool. Use 'releaseSemaphoreToken' or
-- 'System.Semaphore.withSemaphoreToken' for prompt release rather than
-- relying on GC timing.
--
-- The fd is held in an internal 'MVar' so 'releaseSemaphoreToken' takes
-- ownership atomically: a second (erroneous) release is a safe no-op.
newtype SemaphoreToken = SemaphoreToken
{ tokenFdLock :: MVar Fd
}
-- | A server-side semaphore (owns the server thread, listen socket, and token pool).
data ServerSemaphore = ServerSemaphore
{ serverClientSemaphore :: !ClientSemaphore
, serverThreadId :: !ThreadId
, serverPool :: !(TVar Int)
, serverState :: !(MVar ServerState) -- ^ MVar is emptied when resources/fds are freed to prevent double close
}
data ServerState = ServerState
{ serverListenFd :: !Fd
, serverCancelFd :: !Fd
-- ^ Write end of the cancel pipe. Writing a byte signals the
-- server loop to exit its poll-accept.
}
create_sem :: SemaphoreName -> Int -> IO (Either SemaphoreError ServerSemaphore)
create_sem sem_nm init_toks = do
mb_res <- MC.try @_ @IOException $ MC.mask_ $ do
socketPath <- getSemaphoreSocketPath sem_nm
listenFd <- listenDomainSocket socketPath
-- ^^ creates the socket file, unlink it too on failure.
let cleanupListen = do
void $ MC.try @_ @IOException $ closeFd listenFd
void $ MC.try @_ @IOException $ removeLink socketPath
flip MC.onException cleanupListen $ do
pool <- newTVarIO init_toks
(cancelRd, cancelWr) <- createPipe
tid <- forkIOWithUnmask $ \unmask ->
unmask (serverLoop pool listenFd cancelRd)
`MC.finally` closeFd cancelRd
stateVar <- newMVar ServerState
{ serverListenFd = listenFd
, serverCancelFd = cancelWr
}
return ServerSemaphore
{ serverClientSemaphore = ClientSemaphore { clientSemaphoreName = sem_nm
, semSocketPath = socketPath }
, serverThreadId = tid
, serverPool = pool
, serverState = stateVar
}
return $ case mb_res of
Left err -> Left $ SemaphoreOtherError err
Right sem -> Right sem
open_sem_raw :: SemaphoreName -> IO (Either SemaphoreError ClientSemaphore)
open_sem_raw nm = do
mb_res <- MC.try @_ @IOException $ do
socketPath <- getSemaphoreSocketPath nm
exists <- doesPathExist socketPath
return (socketPath, exists)
return $ case mb_res of
Left err -> Left $ SemaphoreOtherError err
Right (_, False) -> Left $ SemaphoreDoesNotExist (semaphoreIdentifier nm)
Right (socketPath, _) -> Right $
ClientSemaphore
{ clientSemaphoreName = nm
, semSocketPath = socketPath
}
-- | Acquire a token from the semaphore, blocking until one is available.
--
-- This operation is interruptible: it can be cancelled by
-- 'Control.Concurrent.throwTo', 'Control.Concurrent.killThread', etc. If
-- interrupted, any transiently acquired token is automatically returned to the
-- pool.
--
-- For prompt and predictable release of resources, callers should use
-- 'System.Semaphore.withSemaphoreToken' or `releaseSemaphoreToken'.
waitOnSemaphore :: HasCallStack => ClientSemaphore -> IO SemaphoreToken
waitOnSemaphore sem = do
resultVar <- newEmptyMVar
-- Mask exceptions until we get to the interruptible takeMVar
MC.mask_ $ do
fd <- connectDomainSocket (semSocketPath sem)
-- The read() runs in a forked thread
workerTid <- forkIOWithUnmask $ \_ -> do
res <- MC.try @_ @MC.SomeException $ do
fdWriteByte fd CmdWait
fdReadByte fd
putMVar resultVar res
-- uninterruptibleMask_: killThread is interruptible, and a
-- second async between killThread and closeFd would leak fd.
let cleanup = MC.uninterruptibleMask_ $ do
-- shutdown(SHUT_RDWR) causes the worker's read() to return EOF immediately
void $ MC.try @_ @IOException $ fdShutdown fd
-- Wait for the worker to exit before closing the fd.
-- this prevents double close bugs
killThread workerTid
-- Finally close the fd
void $ MC.try @_ @IOException $ closeFd fd
-- We achieve interruptiblity by relying on the interruptiblity of takeMVar
-- The worker thread is blocked on read(), fdShutdown in cleanup interrupts
-- it if we are interrupt by an async exception here.
res <- takeMVar resultVar `MC.onException` cleanup
case res of
Right resp
| resp == RspOk -> mkToken fd
| otherwise -> do void $ MC.try @_ @IOException $ closeFd fd
fail $ "semaphore-compat: unexpected response in waitOnSemaphore: " ++ show resp
Left e -> do void $ MC.try @_ @IOException $ closeFd fd
MC.throwM e
-- | Try to acquire a token from the semaphore without blocking.
--
-- Returns @Just token@ if a token was available, @Nothing@ otherwise.
--
-- Not interruptible, but this shouldn't block for long as the server is
-- supposed to respond immediately.
tryWaitOnSemaphore :: HasCallStack => ClientSemaphore -> IO (Maybe SemaphoreToken)
tryWaitOnSemaphore sem =
MC.mask_ $ do
fd <- connectDomainSocket (semSocketPath sem)
resp <- flip MC.onException (closeFd fd) $ do
fdWriteByte fd CmdTryWait
fdReadByte fd
case resp of
RspOk -> Just <$> mkToken fd
_ -> do void $ MC.try @_ @IOException $ closeFd fd
return Nothing
mkToken :: Fd -> IO SemaphoreToken
mkToken fd = do
fdVar <- newMVar fd
_ <- mkWeakMVar fdVar $ do
mb <- tryTakeMVar fdVar
case mb of
Nothing -> return () -- already released
-- Closing the fd triggers the server's disconnect handling,
-- which returns the token to the pool.
Just fd' -> void $ MC.try @_ @IOException $ closeFd fd'
return (SemaphoreToken fdVar)
-- | Release a semaphore token, returning it to the pool.
--
-- Sends a release command on the token's connection, then closes it.
-- Idempotent: a second call on the same token is a safe no-op.
--
-- Not interruptible; only returns when the release has succeeded.
releaseSemaphoreToken :: HasCallStack => SemaphoreToken -> IO ()
releaseSemaphoreToken (SemaphoreToken fdVar) =
MC.mask_ $ do
mb <- tryTakeMVar fdVar
case mb of
Nothing -> return () -- already released
Just fd -> do
resp <- (do fdWriteByte fd CmdRelease
fdReadByte fd
) `MC.finally` (void $ MC.try @_ @IOException $ closeFd fd)
case resp of
RspOk -> return ()
-- myCount <= 0 on the server means the token is effectively
-- already released; stay idempotent.
RspFail -> return ()
_ -> fail "semaphore-compat: unexpected response in releaseSemaphoreToken"
-- | Destroy a client-side semaphore.
--
-- On POSIX this is a no-op: 'ClientSemaphore' holds no live connection.
destroyClientSemaphore :: ClientSemaphore -> IO ()
destroyClientSemaphore _ = return ()
-- | Destroy a server-side semaphore.
--
-- Idempotent. Subsequent calls after the first one are no-ops
-- Not interruptible. Only returns when the server and all resources
-- have been cleaned up
destroyServerSemaphore :: ServerSemaphore -> IO ()
destroyServerSemaphore server = MC.uninterruptibleMask_ $ do
-- we justify the uninterruptibleMask_ here by analysis of the server, which must exit
-- in a bounded amount of time once it receives the cancel signal on serverCancelFd
-- It has a bounded number of child threads it needs to cleanup before returning,
-- but otherwise it should exit promptly.
--
-- Without uninterruptibleMask_, we potentially leak serverListenFd and path if an
-- exception arrives when we are in `killThread`.
mbState <- tryTakeMVar (serverState server)
case mbState of
Nothing -> return () -- already destroyed
Just ServerState{..} -> do
let path = semSocketPath $ serverClientSemaphore server
-- Signal the server loop to exit pollAcceptSocket, then wait for it.
void $ MC.try @_ @IOException $ fdWriteByte serverCancelFd 0
void $ MC.try @_ @IOException $ closeFd serverCancelFd
killThread (serverThreadId server)
void $ MC.try @_ @IOException $ closeFd serverListenFd
void $ MC.try @_ @IOException $ removeLink path
-- | Query the current semaphore value (how many tokens it has available).
--
-- This is mainly for debugging use, as it is easy to introduce race conditions
-- when nontrivial program logic depends on the value returned by this function.
getSemaphoreValue :: ServerSemaphore -> IO Int
getSemaphoreValue server = readTVarIO (serverPool server)
getTimeSeed :: IO Int
getTimeSeed = do
ns <- getMonotonicTimeNSec
pid <- getProcessID
return $ fromIntegral ns `xor` fromIntegral pid
---------------------------------------
-- Server (Unix domain socket)
--
-- The server manages a shared token pool (TVar Int) and accepts multiple
-- client connections, each served on its own thread.
--
-- Protocol (SOCK_STREAM, one byte per command):
--
-- "-" Wait (blocking acquire). Decrements semaphore; replies ".".
-- "?" Try-wait. Decrements if positive and replies "."; otherwise replies "!".
-- "+" Release. Increments pool; replies ".". Rejected with "!" if
-- myCount <= 0 (client has not acquired any tokens on this connection).
--
-- Unrecognised bytes are rejected with RspFail.
--
-- Per-connection token tracking: the server counts tokens held by each
-- connection (myCount). On disconnect (EOF / ResourceVanished) held
-- tokens are returned to the pool, so a crashing client cannot leak
-- tokens.
--
-- Connections are tied to SemaphoreToken, so each token (obtained by waitOnSemaphore)
-- has its own connection
--
---------------------------------------
-- Protocol byte constants
pattern CmdWait, CmdTryWait, CmdRelease :: Word8
pattern CmdWait = 0x2D -- '-'
pattern CmdTryWait = 0x3F -- '?'
pattern CmdRelease = 0x2B -- '+'
pattern RspOk, RspFail :: Word8
pattern RspOk = 0x2E -- '.'
pattern RspFail = 0x21 -- '!'
-- | Children of the server, to be cleaned up
-- childFdLock is full when the 'Fd' is still known to be valid,
-- and subject to being shutdown or closed
data Child = Child
{ childThread :: !(MVar ThreadId)
, childFdLock :: !(MVar Fd)
}
serverLoop :: TVar Int -> Fd -> Fd -> IO ()
serverLoop pool listenFd cancelFd = do
children <- newTVarIO ([] :: [Child])
loop children `MC.finally` killChildren children
where
loop children = do
continueLoop <- MC.mask_ $ do
r <- acceptWithRetry
case r of
AcceptCancelled -> return False
AcceptedFd clientFd -> do
forkServeChild children clientFd
return True
when continueLoop $ loop children
-- pollAcceptSocket only returns if the accept succeeded, or cancellation
-- was signalled via the cancel pipe.
acceptWithRetry :: IO AcceptResult
acceptWithRetry = pollAcceptSocket listenFd cancelFd `MC.catch` handleIOError
-- Retry accept on transient errors.
handleIOError :: IOException -> IO AcceptResult
handleIOError e
-- 'isFullError' catches ResourceExhausted (EMFILE/ENFILE/ENOBUFS/ENOMEM and more that accept doesn't produce but are harmless to retry).
| isFullError e = acceptWithRetry
-- ECONNABORTED is also transient but categorised as OtherError, we additionaly match that.
| Just err <- ioe_errno e
, Errno err == eCONNABORTED = acceptWithRetry
-- EINTR is absorbed in hs_poll_accept.
-- everything else (EBADF, EINVAL, ENOTSOCK, ...) is rethrown.
| otherwise = MC.throwM e
forkServeChild children clientFd = do
fdLock <- newMVar clientFd
tidVar <- newEmptyMVar
let child = Child tidVar fdLock
atomically $ modifyTVar' children (child :)
childTid <- forkIOWithUnmask $ \unmask ->
serve unmask pool children clientFd child
putMVar tidVar childTid
-- Interrupt all children blocked on a read from the FD, then kill them.
killChildren :: TVar [Child] -> IO ()
killChildren children = do
kids <- readTVarIO children
forM_ kids $ \child -> do
-- If the child is in a read(), interrupt it
-- by calling 'fdShutdown'.
mb <- tryTakeMVar (childFdLock child)
case mb of
Just cfd -> do
void $ MC.try @_ @IOException $ fdShutdown cfd
putMVar (childFdLock child) cfd
Nothing -> return () -- serve thread already closing
-- No children blocked on read: terminate them all.
-- childThread was filled by the parent before mask exit; readMVar of
-- a definitely-full MVar is non-interruptible.
forM_ kids $ \child -> do
tid <- readMVar (childThread child)
killThread tid
-- | Per-connection server loop.
serve :: (forall a. IO a -> IO a)
-> TVar Int -> TVar [Child] -> Fd -> Child
-> IO ()
serve restore pool children fd (Child _ fdLock) = do
myCount <- newTVarIO (0 :: Int)
let loop = forever $ MC.mask $ \restoreInner -> do
-- fdReadByte is a safe-FFI read(2), interrupted by
-- fdShutdown from killChildren, not by throwTo.
msg <- fdReadByte fd
case msg of
CmdWait -> do
-- Block until a token is available.
-- restoreInner keeps retry interruptible under mask.
restoreInner $ atomically $ do
n <- readTVar pool
when (n <= 0) retry
writeTVar pool (n - 1)
modifyTVar' myCount (+ 1)
fdWriteByte fd RspOk
CmdRelease -> do
ok <- atomically $ do
mc <- readTVar myCount
if mc > 0
then do
modifyTVar' pool (+ 1)
modifyTVar' myCount (subtract 1)
return True
else return False
fdWriteByte fd (if ok then RspOk else RspFail)
CmdTryWait -> do
acquired <- atomically $ do
n <- readTVar pool
if n > 0
then do
writeTVar pool (n - 1)
modifyTVar' myCount (+ 1)
return True
else return False
if acquired
then fdWriteByte fd RspOk
else fdWriteByte fd RspFail
-- Unknown command: reply so the client doesn't hang in read().
_ -> fdWriteByte fd RspFail
cleanup = do
-- Return tokens to the pool and remove from children list.
atomically $ do
n <- readTVar myCount
when (n > 0) $ modifyTVar' pool (+ n)
modifyTVar' children (filter (\c -> childFdLock c /= fdLock))
-- Take fd ownership and close.
-- prevents killChildren from double closing fd
void $ takeMVar fdLock
void $ MC.try @_ @IOException $ closeFd fd
-- restore so thread can be killed in between loop iterations
-- Catch IOException (EOF/disconnect) silently.
(restore loop `MC.catch` \(_ :: IOException) -> return ())
`MC.finally` cleanup