{-# LANGUAGE Safe #-}
{-# LANGUAGE BangPatterns #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Control.Concurrent.QSem
-- Copyright   :  (c) The University of Glasgow 2001
-- License     :  BSD-style (see the file libraries/base/LICENSE)
-- 
-- Maintainer  :  libraries@haskell.org
-- Stability   :  experimental
-- Portability :  non-portable (concurrency)
--
-- Simple quantity semaphores.
--
-----------------------------------------------------------------------------

module Control.Concurrent.QSem
        ( -- * Simple Quantity Semaphores
          QSem,         -- abstract
          newQSem,      -- :: Int  -> IO QSem
          waitQSem,     -- :: QSem -> IO ()
          signalQSem    -- :: QSem -> IO ()
        ) where

import Control.Concurrent.MVar ( MVar, newEmptyMVar, takeMVar, tryTakeMVar
                          , putMVar, newMVar, tryPutMVar)
import Control.Exception
import Data.Maybe

-- | 'QSem' is a quantity semaphore in which the resource is acquired
-- and released in units of one. It provides guaranteed FIFO ordering
-- for satisfying blocked `waitQSem` calls.
--
-- The pattern
--
-- >   bracket_ waitQSem signalQSem (...)
--
-- is safe; it never loses a unit of the resource.
--
newtype QSem = QSem (MVar (Int, [MVar ()], [MVar ()]))

-- The semaphore state (i, xs, ys):
--
--   i is the current resource value
--
--   (xs,ys) is the queue of blocked threads, where the queue is
--           given by xs ++ reverse ys.  We can enqueue new blocked threads
--           by consing onto ys, and dequeue by removing from the head of xs.
--
-- A blocked thread is represented by an empty (MVar ()).  To unblock
-- the thread, we put () into the MVar.
--
-- A thread can dequeue itself by also putting () into the MVar, which
-- it must do if it receives an exception while blocked in waitQSem.
-- This means that when unblocking a thread in signalQSem we must
-- first check whether the MVar is already full; the MVar lock on the
-- semaphore itself resolves race conditions between signalQSem and a
-- thread attempting to dequeue itself.

-- |Build a new 'QSem' with a supplied initial quantity.
--  The initial quantity must be at least 0.
newQSem :: Int -> IO QSem
newQSem :: Int -> IO QSem
newQSem Int
initial
  | Int
initial Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = String -> IO QSem
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"newQSem: Initial quantity must be non-negative"
  | Bool
otherwise   = do
      MVar (Int, [MVar ()], [MVar ()])
sem <- (Int, [MVar ()], [MVar ()])
-> IO (MVar (Int, [MVar ()], [MVar ()]))
forall a. a -> IO (MVar a)
newMVar (Int
initial, [], [])
      QSem -> IO QSem
forall (m :: * -> *) a. Monad m => a -> m a
return (MVar (Int, [MVar ()], [MVar ()]) -> QSem
QSem MVar (Int, [MVar ()], [MVar ()])
sem)

-- |Wait for a unit to become available
waitQSem :: QSem -> IO ()
waitQSem :: QSem -> IO ()
waitQSem (QSem MVar (Int, [MVar ()], [MVar ()])
m) =
  IO () -> IO ()
forall a. IO a -> IO a
mask_ (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
    (Int
i,[MVar ()]
b1,[MVar ()]
b2) <- MVar (Int, [MVar ()], [MVar ()]) -> IO (Int, [MVar ()], [MVar ()])
forall a. MVar a -> IO a
takeMVar MVar (Int, [MVar ()], [MVar ()])
m
    if Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
       then do
         MVar ()
b <- IO (MVar ())
forall a. IO (MVar a)
newEmptyMVar
         MVar (Int, [MVar ()], [MVar ()])
-> (Int, [MVar ()], [MVar ()]) -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar (Int, [MVar ()], [MVar ()])
m (Int
i, [MVar ()]
b1, MVar ()
bMVar () -> [MVar ()] -> [MVar ()]
forall a. a -> [a] -> [a]
:[MVar ()]
b2)
         MVar () -> IO ()
wait MVar ()
b
       else do
         let !z :: Int
z = Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1
         MVar (Int, [MVar ()], [MVar ()])
-> (Int, [MVar ()], [MVar ()]) -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar (Int, [MVar ()], [MVar ()])
m (Int
z, [MVar ()]
b1, [MVar ()]
b2)
         () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
  where
    wait :: MVar () -> IO ()
wait MVar ()
b = MVar () -> IO ()
forall a. MVar a -> IO a
takeMVar MVar ()
b IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO a
`onException` do
                (IO () -> IO ()
forall a. IO a -> IO a
uninterruptibleMask_ (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do -- Note [signal uninterruptible]
                   (Int
i,[MVar ()]
b1,[MVar ()]
b2) <- MVar (Int, [MVar ()], [MVar ()]) -> IO (Int, [MVar ()], [MVar ()])
forall a. MVar a -> IO a
takeMVar MVar (Int, [MVar ()], [MVar ()])
m
                   Maybe ()
r <- MVar () -> IO (Maybe ())
forall a. MVar a -> IO (Maybe a)
tryTakeMVar MVar ()
b
                   (Int, [MVar ()], [MVar ()])
r' <- if Maybe () -> Bool
forall a. Maybe a -> Bool
isJust Maybe ()
r
                            then (Int, [MVar ()], [MVar ()]) -> IO (Int, [MVar ()], [MVar ()])
signal (Int
i,[MVar ()]
b1,[MVar ()]
b2)
                            else do MVar () -> () -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar ()
b (); (Int, [MVar ()], [MVar ()]) -> IO (Int, [MVar ()], [MVar ()])
forall (m :: * -> *) a. Monad m => a -> m a
return (Int
i,[MVar ()]
b1,[MVar ()]
b2)
                   MVar (Int, [MVar ()], [MVar ()])
-> (Int, [MVar ()], [MVar ()]) -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar (Int, [MVar ()], [MVar ()])
m (Int, [MVar ()], [MVar ()])
r')

-- |Signal that a unit of the 'QSem' is available
signalQSem :: QSem -> IO ()
signalQSem :: QSem -> IO ()
signalQSem (QSem MVar (Int, [MVar ()], [MVar ()])
m) =
  IO () -> IO ()
forall a. IO a -> IO a
uninterruptibleMask_ (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do -- Note [signal uninterruptible]
    (Int, [MVar ()], [MVar ()])
r <- MVar (Int, [MVar ()], [MVar ()]) -> IO (Int, [MVar ()], [MVar ()])
forall a. MVar a -> IO a
takeMVar MVar (Int, [MVar ()], [MVar ()])
m
    (Int, [MVar ()], [MVar ()])
r' <- (Int, [MVar ()], [MVar ()]) -> IO (Int, [MVar ()], [MVar ()])
signal (Int, [MVar ()], [MVar ()])
r
    MVar (Int, [MVar ()], [MVar ()])
-> (Int, [MVar ()], [MVar ()]) -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar (Int, [MVar ()], [MVar ()])
m (Int, [MVar ()], [MVar ()])
r'

-- Note [signal uninterruptible]
--
--   If we have
--
--      bracket waitQSem signalQSem (...)
--
--   and an exception arrives at the signalQSem, then we must not lose
--   the resource.  The signalQSem is masked by bracket, but taking
--   the MVar might block, and so it would be interruptible.  Hence we
--   need an uninterruptibleMask here.
--
--   This isn't ideal: during high contention, some threads won't be
--   interruptible.  The QSemSTM implementation has better behaviour
--   here, but it performs much worse than this one in some
--   benchmarks.

signal :: (Int,[MVar ()],[MVar ()]) -> IO (Int,[MVar ()],[MVar ()])
signal :: (Int, [MVar ()], [MVar ()]) -> IO (Int, [MVar ()], [MVar ()])
signal (Int
i,[MVar ()]
a1,[MVar ()]
a2) =
 if Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
   then [MVar ()] -> [MVar ()] -> IO (Int, [MVar ()], [MVar ()])
forall {a}.
Num a =>
[MVar ()] -> [MVar ()] -> IO (a, [MVar ()], [MVar ()])
loop [MVar ()]
a1 [MVar ()]
a2
   else let !z :: Int
z = Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1 in (Int, [MVar ()], [MVar ()]) -> IO (Int, [MVar ()], [MVar ()])
forall (m :: * -> *) a. Monad m => a -> m a
return (Int
z, [MVar ()]
a1, [MVar ()]
a2)
 where
   loop :: [MVar ()] -> [MVar ()] -> IO (a, [MVar ()], [MVar ()])
loop [] [] = (a, [MVar ()], [MVar ()]) -> IO (a, [MVar ()], [MVar ()])
forall (m :: * -> *) a. Monad m => a -> m a
return (a
1, [], [])
   loop [] [MVar ()]
b2 = [MVar ()] -> [MVar ()] -> IO (a, [MVar ()], [MVar ()])
loop ([MVar ()] -> [MVar ()]
forall a. [a] -> [a]
reverse [MVar ()]
b2) []
   loop (MVar ()
b:[MVar ()]
bs) [MVar ()]
b2 = do
     Bool
r <- MVar () -> () -> IO Bool
forall a. MVar a -> a -> IO Bool
tryPutMVar MVar ()
b ()
     if Bool
r then (a, [MVar ()], [MVar ()]) -> IO (a, [MVar ()], [MVar ()])
forall (m :: * -> *) a. Monad m => a -> m a
return (a
0, [MVar ()]
bs, [MVar ()]
b2)
          else [MVar ()] -> [MVar ()] -> IO (a, [MVar ()], [MVar ()])
loop [MVar ()]
bs [MVar ()]
b2