Win32.hs

{-# LANGUAGE TypeApplications #-}

module System.Semaphore.Internal.Win32
  ( ClientSemaphore(..), ServerSemaphore(..)
  , SemaphoreToken(..)
  , create_sem, open_sem_raw
  , waitOnSemaphore, tryWaitOnSemaphore
  , releaseSemaphoreToken
  , destroyClientSemaphore, destroyServerSemaphore
  , getSemaphoreValue
  , getTimeSeed
  ) where

import Control.Concurrent
  ( forkIO, killThread )
import Control.Concurrent.MVar
  ( MVar, newEmptyMVar, newMVar, putMVar, takeMVar, tryReadMVar, tryTakeMVar )
import Control.Exception ( IOException, uninterruptibleMask_ )
import Control.Monad
import System.IO.Error ( isDoesNotExistError )

import qualified Control.Monad.Catch as MC

import qualified System.Win32.Event     as Win32
  ( createEvent, setEvent
  , waitForSingleObject, waitForMultipleObjects
  , wAIT_OBJECT_0 )
import qualified System.Win32.File      as Win32
  ( closeHandle )
import qualified System.Win32.Process   as Win32
  ( iNFINITE )
import qualified System.Win32.Semaphore as Win32
  ( Semaphore(..), sEMAPHORE_ALL_ACCESS
  , createSemaphore, openSemaphore, releaseSemaphore )
import qualified System.Win32.Time      as Win32
  ( FILETIME(..), getSystemTimeAsFileTime )

import System.Semaphore.Internal.Common

-- | A client-side semaphore (semaphore for a jobclient).
data ClientSemaphore =
  ClientSemaphore
    { clientSemaphoreName :: !SemaphoreName
    , semaphore           :: !(MVar Win32.Semaphore)
      -- ^ MVar is emptied when the handle is closed to prevent double close.
    }

-- | A held semaphore token.
newtype SemaphoreToken = SemaphoreToken
  { tokenSemaphore :: Win32.Semaphore
  }

-- | A server-side semaphore: a 'ClientSemaphore' with additional
-- book-keeping to keep track of resource usage by all clients.
--
-- Use 'serverClientSemaphore' to retrieve the client-side semaphore.
newtype ServerSemaphore = ServerSemaphore
  { serverClientSemaphore :: ClientSemaphore
  }

create_sem :: SemaphoreName -> Int -> IO (Either SemaphoreError ServerSemaphore)
create_sem sem_nm init_toks = MC.mask_ $ do
  let toks = fromIntegral init_toks
  mb_sem <- MC.try @_ @IOException $
    Win32.createSemaphore Nothing toks 0x7FFFFFFF (Just (semaphoreIdentifier sem_nm))
  case mb_sem of
    Left err ->
      return $ Left $ SemaphoreOtherError err
    Right (sem, True) -> do
      -- Semaphore already existed; createSemaphore still returned a valid
      -- handle to it, which we must close to avoid leaking it.
      void $ MC.try @_ @IOException $ Win32.closeHandle (Win32.semaphoreHandle sem)
      return $ Left $ SemaphoreAlreadyExists (semaphoreIdentifier sem_nm)
    Right (sem, False) -> do
      semVar <- newMVar sem
      return $ Right $
        ServerSemaphore
          { serverClientSemaphore =
              ClientSemaphore
                { clientSemaphoreName = sem_nm
                , semaphore     = semVar }
          }

open_sem_raw :: SemaphoreName -> IO (Either SemaphoreError ClientSemaphore)
open_sem_raw nm = MC.mask_ $ do
  mb_res <- MC.try @_ @IOException $
    Win32.openSemaphore Win32.sEMAPHORE_ALL_ACCESS True (semaphoreIdentifier nm)
  case mb_res of
    Left  err
      | isDoesNotExistError err
      -> return $ Left $ SemaphoreDoesNotExist (semaphoreIdentifier nm)
      | otherwise
      -> return $ Left $ SemaphoreOtherError err
    Right sem -> do
      semVar <- newMVar sem
      return $ Right $
        ClientSemaphore
          { clientSemaphoreName = nm
          , semaphore     = semVar }

-- | Read the kernel handle from a 'ClientSemaphore'.  Throws if the
-- client has been destroyed.
readClientHandle :: ClientSemaphore -> String -> IO Win32.Semaphore
readClientHandle client opName = do
  mb <- tryReadMVar (semaphore client)
  case mb of
    Just sem -> return sem
    Nothing  -> ioError $ userError $
      "semaphore-compat." ++ opName ++ ": semaphore has been destroyed"

-- | Acquire a token from the semaphore, blocking until one is available.
--
-- This operation is interruptible. If interrupted after the token was acquired,
-- the token is automatically released back to the pool.
--
-- Returns a 'SemaphoreToken' that must be released with 'releaseSemaphoreToken'.
waitOnSemaphore :: ClientSemaphore -> IO SemaphoreToken
waitOnSemaphore client = do
  resultVar <- newEmptyMVar
  MC.mask_ $ do
    sem <- readClientHandle client "waitOnSemaphore"
    cancelEvent <- Win32.createEvent Nothing True False ""
    workerTid <- (forkIO $ MC.mask_ $ do
      res <- MC.try @_ @MC.SomeException $ do
        wait_res <-
          Win32.waitForMultipleObjects
            [ Win32.semaphoreHandle sem
            , cancelEvent ]
            False -- WaitAny
            Win32.iNFINITE
        if wait_res == Win32.wAIT_OBJECT_0
          then return (SemaphoreToken sem)
          else ioError (userError "semaphore wait cancelled")
      putMVar resultVar res
      ) `MC.onException` Win32.closeHandle cancelEvent
    let cleanup =
          -- Signal the cancel event and wait for the worker to exit.
          (do ok <- MC.try @_ @IOException $ Win32.setEvent cancelEvent
              case ok of
                Right () ->
                  -- Cancel event signaled. Worker's WaitForMultipleObjects
                  -- returns immediately. Safe to wait uninterruptibly.
                  uninterruptibleMask_ $ killThread workerTid
                Left _ ->
                  -- setEvent failed. We can't really do anything.
                  -- The worker is in a safe ffi call so killThread
                  -- will just block forever.
                  return ()
          ) `MC.finally` (do
              -- Recover any transiently acquired token.  If the worker
              -- acquired the semaphore before seeing the cancel event,
              -- release it back to the pool.
              mbRes <- tryTakeMVar resultVar
              case mbRes of
                Just (Right tok) ->
                  void $ MC.try @_ @MC.SomeException $ releaseSemaphoreToken tok
                _ -> return ()
              Win32.closeHandle cancelEvent)
    res <- takeMVar resultVar `MC.onException` cleanup
    Win32.closeHandle cancelEvent
    case res of
      Right tok -> return tok
      Left e    -> MC.throwM e

-- | Try to acquire a token from the semaphore without blocking.
--
-- Returns @Just token@ if a token was available, @Nothing@ otherwise.
tryWaitOnSemaphore :: ClientSemaphore -> IO (Maybe SemaphoreToken)
tryWaitOnSemaphore client =
  MC.mask_ $ do
    sem <- readClientHandle client "tryWaitOnSemaphore"
    wait_res <- Win32.waitForSingleObject (Win32.semaphoreHandle sem) 0
    if wait_res == Win32.wAIT_OBJECT_0
      then return (Just (SemaphoreToken sem))
      else return Nothing

-- | Release a semaphore token, returning it to the pool.
releaseSemaphoreToken :: SemaphoreToken -> IO ()
releaseSemaphoreToken (SemaphoreToken sem) =
  MC.mask_ $ do
    void $ Win32.releaseSemaphore sem 1

-- | Destroy a client-side semaphore.
--
-- Idempotent: subsequent calls after the first are no-ops.
destroyClientSemaphore :: ClientSemaphore -> IO ()
destroyClientSemaphore client = uninterruptibleMask_ $ do
  mb <- tryTakeMVar (semaphore client)
  case mb of
    Nothing  -> return ()  -- already destroyed
    Just sem -> Win32.closeHandle (Win32.semaphoreHandle sem)

-- | Destroy a server-side semaphore.
--
-- Idempotent: subsequent calls after the first are no-ops.
destroyServerSemaphore :: ServerSemaphore -> IO ()
destroyServerSemaphore server =
  destroyClientSemaphore (serverClientSemaphore server)

-- | 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 =
  MC.mask_ $ do
    sem <- readClientHandle (serverClientSemaphore server) "getSemaphoreValue"
    wait_res <- Win32.waitForSingleObject (Win32.semaphoreHandle sem) 0
    if wait_res == Win32.wAIT_OBJECT_0
    then
      (+1) . fromIntegral <$> Win32.releaseSemaphore sem 1
    else
      return 0

getTimeSeed :: IO Int
getTimeSeed = do
  Win32.FILETIME t <- Win32.getSystemTimeAsFileTime
  return (fromIntegral t)