{-# LANGUAGE RankNTypes #-}

module Network.Wai.Handler.Warp.Counter (
    Counter,
    newCounter,
    HasDecreased (..),
    waitForZero,
    forkConnection,
    waitForDecreased,
    getCount,
    getCountSTM,
    threadManager,
) where

import Control.Concurrent.STM
import System.ThreadManager (
    ThreadManager,
    countManaged,
    forkManagedBy,
    newThreadManager,
    waitUntilAllGone,
 )

-- | The connections of a server: how many there are, and which threads
--   are serving them.
--
--   Both come from the 'ThreadManager', which counts a connection from
--   when the accept loop keeps a place for it -- before the thread that
--   will serve it has been scheduled, which is too late for a graceful
--   shutdown to notice it.
newtype Counter = Counter ThreadManager

newCounter :: IO Counter
newCounter :: IO Counter
newCounter = ThreadManager -> Counter
Counter (ThreadManager -> Counter) -> IO ThreadManager -> IO Counter
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO ThreadManager
newThreadManager

-- | The manager owning the connections, for supervising them.
threadManager :: Counter -> ThreadManager
threadManager :: Counter -> ThreadManager
threadManager (Counter ThreadManager
tm) = ThreadManager
tm

waitForZero :: Counter -> IO ()
waitForZero :: Counter -> IO ()
waitForZero (Counter ThreadManager
tm) = ThreadManager -> IO ()
waitUntilAllGone ThreadManager
tm

data HasDecreased = HasDecreased | NoConnections
    deriving (HasDecreased -> HasDecreased -> Bool
(HasDecreased -> HasDecreased -> Bool)
-> (HasDecreased -> HasDecreased -> Bool) -> Eq HasDecreased
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: HasDecreased -> HasDecreased -> Bool
== :: HasDecreased -> HasDecreased -> Bool
$c/= :: HasDecreased -> HasDecreased -> Bool
/= :: HasDecreased -> HasDecreased -> Bool
Eq, Int -> HasDecreased -> ShowS
[HasDecreased] -> ShowS
HasDecreased -> String
(Int -> HasDecreased -> ShowS)
-> (HasDecreased -> String)
-> ([HasDecreased] -> ShowS)
-> Show HasDecreased
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> HasDecreased -> ShowS
showsPrec :: Int -> HasDecreased -> ShowS
$cshow :: HasDecreased -> String
show :: HasDecreased -> String
$cshowList :: [HasDecreased] -> ShowS
showList :: [HasDecreased] -> ShowS
Show)

waitForDecreased :: Counter -> IO HasDecreased
waitForDecreased :: Counter -> IO HasDecreased
waitForDecreased (Counter ThreadManager
tm) = do
    Int
n0 <- STM Int -> IO Int
forall a. STM a -> IO a
atomically (STM Int -> IO Int) -> STM Int -> IO Int
forall a b. (a -> b) -> a -> b
$ ThreadManager -> STM Int
countManaged ThreadManager
tm
    if Int
n0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0
        then HasDecreased -> IO HasDecreased
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure HasDecreased
NoConnections
        else STM HasDecreased -> IO HasDecreased
forall a. STM a -> IO a
atomically (STM HasDecreased -> IO HasDecreased)
-> STM HasDecreased -> IO HasDecreased
forall a b. (a -> b) -> a -> b
$ do
            Int
n <- ThreadManager -> STM Int
countManaged ThreadManager
tm
            Bool -> STM ()
check (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n0)
            HasDecreased -> STM HasDecreased
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure HasDecreased
HasDecreased

-- | Serving a connection on a thread made by the given fork.  The
--   connection is counted from here rather than from when that thread
--   starts: the accept loop does not wait for it, and a connection which
--   has been accepted must not look to 'waitForZero' like one that is
--   gone.
forkConnection
    :: Counter
    -> (((forall x. IO x -> IO x) -> IO ()) -> IO ())
    -- ^ The fork from the settings
    -> ((forall x. IO x -> IO x) -> IO ())
    -- ^ Serving the connection
    -> IO ()
forkConnection :: Counter
-> (((forall x. IO x -> IO x) -> IO ()) -> IO ())
-> ((forall x. IO x -> IO x) -> IO ())
-> IO ()
forkConnection (Counter ThreadManager
tm) ((forall x. IO x -> IO x) -> IO ()) -> IO ()
fork (forall x. IO x -> IO x) -> IO ()
io = ThreadManager
-> (((forall x. IO x -> IO x) -> IO ()) -> IO ())
-> String
-> ((forall x. IO x -> IO x) -> IO ())
-> IO ()
forall a.
ThreadManager
-> (((forall x. IO x -> IO x) -> IO ()) -> IO a)
-> String
-> ((forall x. IO x -> IO x) -> IO ())
-> IO a
forkManagedBy ThreadManager
tm ((forall x. IO x -> IO x) -> IO ()) -> IO ()
fork String
"Warp just forked" (forall x. IO x -> IO x) -> IO ()
io

-- | Get the current count of open connections.
--
-- @since 3.4.11
getCount :: Counter -> IO Int
getCount :: Counter -> IO Int
getCount = STM Int -> IO Int
forall a. STM a -> IO a
atomically (STM Int -> IO Int) -> (Counter -> STM Int) -> Counter -> IO Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Counter -> STM Int
getCountSTM

-- | Get the current count in an 'STM' transaction.
--
-- @since 3.4.13
getCountSTM :: Counter -> STM Int
getCountSTM :: Counter -> STM Int
getCountSTM (Counter ThreadManager
tm) = ThreadManager -> STM Int
countManaged ThreadManager
tm