{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

-- | The Unix domain socket server backing the POSIX semaphore implementation.
--
-- The server manages a shared token pool 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 a token, so each token (obtained by
-- @waitOnSemaphore@) has its own connection.
module System.Semaphore.Internal.Posix.Server
  ( serverLoop
  , pattern CmdWait, pattern CmdTryWait, pattern CmdRelease
  , pattern RspOk, pattern RspFail
  ) where

-- base
import Control.Concurrent
  ( ThreadId, forkIOWithUnmask, killThread )
import Control.Concurrent.MVar
  ( MVar, newEmptyMVar, newMVar, putMVar
  , readMVar, takeMVar, tryTakeMVar )
import Control.Exception
  ( IOException, SomeException
  , catch, finally, mask, mask_
  , onException, throw, try, uninterruptibleMask_
  )
import Control.Monad
  ( forM_, forever, void, when )
import Data.Word ( Word8 )
import Foreign.C.Error ( Errno(Errno), eCONNABORTED )
import GHC.IO.Exception ( ioe_errno )
import System.IO.Error ( isFullError )

-- stm
import Control.Concurrent.STM
  ( TVar, atomically, newTVarIO, readTVar, readTVarIO
  , modifyTVar', writeTVar, retry )

-- unix
import System.Posix.IO ( closeFd )
import System.Posix.Types ( Fd )

import System.Semaphore.Internal.DomainSocket
  ( pollAcceptSocket, AcceptResult(..)
  , fdReadByte, fdWriteByte
  , fdShutdown )

---------------------------------------
-- Protocol byte constants

pattern CmdWait, CmdTryWait, CmdRelease :: Word8
pattern $mCmdWait :: forall {r}. Word8 -> ((# #) -> r) -> ((# #) -> r) -> r
$bCmdWait :: Word8
CmdWait    = 0x2D -- '-'
pattern $mCmdTryWait :: forall {r}. Word8 -> ((# #) -> r) -> ((# #) -> r) -> r
$bCmdTryWait :: Word8
CmdTryWait = 0x3F -- '?'
pattern $mCmdRelease :: forall {r}. Word8 -> ((# #) -> r) -> ((# #) -> r) -> r
$bCmdRelease :: Word8
CmdRelease = 0x2B -- '+'

pattern RspOk, RspFail :: Word8
pattern $mRspOk :: forall {r}. Word8 -> ((# #) -> r) -> ((# #) -> r) -> r
$bRspOk :: Word8
RspOk   = 0x2E -- '.'
pattern $mRspFail :: forall {r}. Word8 -> ((# #) -> r) -> ((# #) -> r) -> r
$bRspFail :: Word8
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
  { Child -> MVar ThreadId
childThread :: !(MVar ThreadId)
  , Child -> MVar Fd
childFdLock :: !(MVar Fd)
  }

-- | Run the server accept loop on a listening socket until cancellation is
-- signalled via the cancel pipe.  All accepted connections are served on
-- child threads, which are cleaned up before this returns.
serverLoop :: TVar Int -- ^ shared token pool
           -> Fd        -- ^ listening socket
           -> Fd        -- ^ read end of the cancel pipe
           -> IO ()
serverLoop :: TVar Int -> Fd -> Fd -> IO ()
serverLoop TVar Int
pool Fd
listenFd Fd
cancelFd = do
    TVar [Child]
children <- [Child] -> IO (TVar [Child])
forall a. a -> IO (TVar a)
newTVarIO ([] :: [Child])
    TVar [Child] -> IO ()
loop TVar [Child]
children IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO a
`finally` TVar [Child] -> IO ()
killChildren TVar [Child]
children
  where
    loop :: TVar [Child] -> IO ()
loop TVar [Child]
children = do
      Bool
continueLoop <- IO Bool -> IO Bool
forall a. IO a -> IO a
mask_ (IO Bool -> IO Bool) -> IO Bool -> IO Bool
forall a b. (a -> b) -> a -> b
$ do
        AcceptResult
r <- IO AcceptResult
acceptWithRetry
        case AcceptResult
r of
          AcceptResult
AcceptCancelled     -> Bool -> IO Bool
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
          AcceptedFd Fd
clientFd -> do
            TVar [Child] -> Fd -> IO ()
forkServeChild TVar [Child]
children Fd
clientFd
            Bool -> IO Bool
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
      Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
continueLoop (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar [Child] -> IO ()
loop TVar [Child]
children

    -- pollAcceptSocket only returns if the accept succeeded, or cancellation
    -- was signalled via the cancel pipe.
    acceptWithRetry :: IO AcceptResult
    acceptWithRetry :: IO AcceptResult
acceptWithRetry = Fd -> Fd -> IO AcceptResult
pollAcceptSocket Fd
listenFd Fd
cancelFd IO AcceptResult
-> (IOException -> IO AcceptResult) -> IO AcceptResult
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` IOException -> IO AcceptResult
handleIOError

    -- Retry accept on transient errors.
    handleIOError :: IOException -> IO AcceptResult
    handleIOError :: IOException -> IO AcceptResult
handleIOError IOException
e
      -- 'isFullError' catches ResourceExhausted (EMFILE/ENFILE/ENOBUFS/ENOMEM and more that accept doesn't produce but are harmless to retry).
      | IOException -> Bool
isFullError IOException
e                             = IO AcceptResult
acceptWithRetry
      -- ECONNABORTED is also transient but categorised as OtherError, we additionaly match that.
      | Just CInt
err <- IOException -> Maybe CInt
ioe_errno IOException
e
      , CInt -> Errno
Errno CInt
err Errno -> Errno -> Bool
forall a. Eq a => a -> a -> Bool
== Errno
eCONNABORTED                 = IO AcceptResult
acceptWithRetry
      -- EINTR is absorbed in hs_poll_accept.
      -- everything else (EBADF, EINVAL, ENOTSOCK, ...) is rethrown.
      | Bool
otherwise                                 = IOException -> IO AcceptResult
forall a e. Exception e => e -> a
throw IOException
e

    forkServeChild :: TVar [Child] -> Fd -> IO ()
forkServeChild TVar [Child]
children Fd
clientFd = do
      MVar Fd
fdLock <- Fd -> IO (MVar Fd)
forall a. a -> IO (MVar a)
newMVar Fd
clientFd
      MVar ThreadId
tidVar <- IO (MVar ThreadId)
forall a. IO (MVar a)
newEmptyMVar
      let child :: Child
child = MVar ThreadId -> MVar Fd -> Child
Child MVar ThreadId
tidVar MVar Fd
fdLock
      STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar [Child] -> ([Child] -> [Child]) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar [Child]
children (Child
child Child -> [Child] -> [Child]
forall a. a -> [a] -> [a]
:)
      ThreadId
childTid <- ((forall a. IO a -> IO a) -> IO ()) -> IO ThreadId
forkIOWithUnmask (((forall a. IO a -> IO a) -> IO ()) -> IO ThreadId)
-> ((forall a. IO a -> IO a) -> IO ()) -> IO ThreadId
forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
unmask ->
        (forall a. IO a -> IO a)
-> TVar Int -> TVar [Child] -> Fd -> Child -> IO ()
serve IO a -> IO a
forall a. IO a -> IO a
unmask TVar Int
pool TVar [Child]
children Fd
clientFd Child
child
      MVar ThreadId -> ThreadId -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar ThreadId
tidVar ThreadId
childTid

    -- Interrupt all children blocked on a read from the FD, then kill them.
    killChildren :: TVar [Child] -> IO ()
    killChildren :: TVar [Child] -> IO ()
killChildren TVar [Child]
children = do
      [Child]
kids <- TVar [Child] -> IO [Child]
forall a. TVar a -> IO a
readTVarIO TVar [Child]
children
      [Child] -> (Child -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Child]
kids ((Child -> IO ()) -> IO ()) -> (Child -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Child
child -> do
        -- If the child is in a read(), interrupt it
        -- by calling 'fdShutdown'.
        Maybe Fd
mb <- MVar Fd -> IO (Maybe Fd)
forall a. MVar a -> IO (Maybe a)
tryTakeMVar (Child -> MVar Fd
childFdLock Child
child)
        case Maybe Fd
mb of
          Just Fd
cfd -> do
            IO (Either IOException ()) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Either IOException ()) -> IO ())
-> IO (Either IOException ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ forall e a. Exception e => IO a -> IO (Either e a)
try @IOException (IO () -> IO (Either IOException ()))
-> IO () -> IO (Either IOException ())
forall a b. (a -> b) -> a -> b
$ Fd -> IO ()
fdShutdown Fd
cfd
            MVar Fd -> Fd -> IO ()
forall a. MVar a -> a -> IO ()
putMVar (Child -> MVar Fd
childFdLock Child
child) Fd
cfd
          Maybe Fd
Nothing  -> () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
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.
      [Child] -> (Child -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Child]
kids ((Child -> IO ()) -> IO ()) -> (Child -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Child
child -> do
        ThreadId
tid <- MVar ThreadId -> IO ThreadId
forall a. MVar a -> IO a
readMVar (Child -> MVar ThreadId
childThread Child
child)
        ThreadId -> IO ()
killThread ThreadId
tid

-- | Per-connection server loop.
serve :: (forall a. IO a -> IO a)
      -> TVar Int -> TVar [Child] -> Fd -> Child
      -> IO ()
serve :: (forall a. IO a -> IO a)
-> TVar Int -> TVar [Child] -> Fd -> Child -> IO ()
serve forall a. IO a -> IO a
restore TVar Int
pool TVar [Child]
children Fd
fd (Child MVar ThreadId
_ MVar Fd
fdLock) = do
    TVar Int
myCount <- Int -> IO (TVar Int)
forall a. a -> IO (TVar a)
newTVarIO (Int
0 :: Int)
    let loop :: IO b
loop = IO () -> IO b
forall (f :: * -> *) a b. Applicative f => f a -> f b
forever (IO () -> IO b) -> IO () -> IO b
forall a b. (a -> b) -> a -> b
$ ((forall a. IO a -> IO a) -> IO ()) -> IO ()
forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
mask (((forall a. IO a -> IO a) -> IO ()) -> IO ())
-> ((forall a. IO a -> IO a) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
restoreInner -> do
            -- fdReadByte is a safe-FFI read(2), interrupted by
            -- fdShutdown from killChildren, not by throwTo.
            Word8
msg <- HasCallStack => Fd -> IO Word8
Fd -> IO Word8
fdReadByte Fd
fd
            case Word8
msg of
              Word8
CmdWait -> do
                -- Block until a token is available.
                -- restoreInner keeps retry interruptible under mask.
                IO () -> IO ()
forall a. IO a -> IO a
restoreInner (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
                    Int
n <- TVar Int -> STM Int
forall a. TVar a -> STM a
readTVar TVar Int
pool
                    Bool -> STM () -> STM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0) STM ()
forall a. STM a
retry
                    TVar Int -> Int -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar Int
pool (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
                    TVar Int -> (Int -> Int) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar Int
myCount (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                HasCallStack => Fd -> Word8 -> IO ()
Fd -> Word8 -> IO ()
fdWriteByte Fd
fd Word8
RspOk

              Word8
CmdRelease -> do
                Bool
ok <- STM Bool -> IO Bool
forall a. STM a -> IO a
atomically (STM Bool -> IO Bool) -> STM Bool -> IO Bool
forall a b. (a -> b) -> a -> b
$ do
                    Int
mc <- TVar Int -> STM Int
forall a. TVar a -> STM a
readTVar TVar Int
myCount
                    if Int
mc Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0
                      then do
                        TVar Int -> (Int -> Int) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar Int
pool (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                        TVar Int -> (Int -> Int) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar Int
myCount (Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
1)
                        Bool -> STM Bool
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
                      else Bool -> STM Bool
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
                HasCallStack => Fd -> Word8 -> IO ()
Fd -> Word8 -> IO ()
fdWriteByte Fd
fd (if Bool
ok then Word8
RspOk else Word8
RspFail)

              Word8
CmdTryWait -> do
                Bool
acquired <- STM Bool -> IO Bool
forall a. STM a -> IO a
atomically (STM Bool -> IO Bool) -> STM Bool -> IO Bool
forall a b. (a -> b) -> a -> b
$ do
                    Int
n <- TVar Int -> STM Int
forall a. TVar a -> STM a
readTVar TVar Int
pool
                    if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0
                      then do
                        TVar Int -> Int -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar Int
pool (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
                        TVar Int -> (Int -> Int) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar Int
myCount (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                        Bool -> STM Bool
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
                      else Bool -> STM Bool
forall a. a -> STM a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
                if Bool
acquired
                  then HasCallStack => Fd -> Word8 -> IO ()
Fd -> Word8 -> IO ()
fdWriteByte Fd
fd Word8
RspOk
                  else HasCallStack => Fd -> Word8 -> IO ()
Fd -> Word8 -> IO ()
fdWriteByte Fd
fd Word8
RspFail

              -- Unknown command: reply so the client doesn't hang in read().
              Word8
_ -> HasCallStack => Fd -> Word8 -> IO ()
Fd -> Word8 -> IO ()
fdWriteByte Fd
fd Word8
RspFail

        cleanup :: IO ()
cleanup = do
          -- Return tokens to the pool and remove from children list.
          STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
            Int
n <- TVar Int -> STM Int
forall a. TVar a -> STM a
readTVar TVar Int
myCount
            Bool -> STM () -> STM ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (STM () -> STM ()) -> STM () -> STM ()
forall a b. (a -> b) -> a -> b
$ TVar Int -> (Int -> Int) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar Int
pool (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n)
            TVar [Child] -> ([Child] -> [Child]) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar [Child]
children ((Child -> Bool) -> [Child] -> [Child]
forall a. (a -> Bool) -> [a] -> [a]
filter (\Child
c -> Child -> MVar Fd
childFdLock Child
c MVar Fd -> MVar Fd -> Bool
forall a. Eq a => a -> a -> Bool
/= MVar Fd
fdLock))
          -- Take fd ownership and close.
          -- prevents killChildren from double closing fd
          IO Fd -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Fd -> IO ()) -> IO Fd -> IO ()
forall a b. (a -> b) -> a -> b
$ MVar Fd -> IO Fd
forall a. MVar a -> IO a
takeMVar MVar Fd
fdLock
          IO (Either IOException ()) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Either IOException ()) -> IO ())
-> IO (Either IOException ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ forall e a. Exception e => IO a -> IO (Either e a)
try @IOException (IO () -> IO (Either IOException ()))
-> IO () -> IO (Either IOException ())
forall a b. (a -> b) -> a -> b
$ Fd -> IO ()
closeFd Fd
fd

    -- restore so thread can be killed in between loop iterations
    -- Catch IOException (EOF/disconnect) silently.
    (IO () -> IO ()
forall a. IO a -> IO a
restore IO ()
forall {b}. IO b
loop IO () -> (IOException -> IO ()) -> IO ()
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` \(IOException
_ :: IOException) -> () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ())
      IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO a
`finally` IO ()
cleanup