{-# LANGUAGE RecordWildCards     #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell     #-}
-- | Manages a pool of available ports and allocates them.
module Keter.PortPool
    ( -- * Types
      PortPool
      -- * Actions
    , getPort
    , releasePort
      -- * Initialize
    , start
    ) where

import           Keter.Common
import           Control.Applicative     ((<$>))
import           Control.Concurrent.MVar
import           Control.Exception
import           Keter.Config
import           Network.Socket
import           Prelude                 hiding (log)

data PPState = PPState
    { PPState -> [Port]
ppAvail    :: ![Port]
    , PPState -> [Port] -> [Port]
ppRecycled :: !([Port] -> [Port])
    }

newtype PortPool = PortPool (MVar PPState)

-- | Gets an unassigned port number.
getPort :: (LogMessage -> IO ())
        -> PortPool
        -> IO (Either SomeException Port)
getPort :: (LogMessage -> IO ()) -> PortPool -> IO (Either SomeException Port)
getPort LogMessage -> IO ()
log (PortPool MVar PPState
mstate) =
    MVar PPState
-> (PPState -> IO (PPState, Either SomeException Port))
-> IO (Either SomeException Port)
forall a b. MVar a -> (a -> IO (a, b)) -> IO b
modifyMVar MVar PPState
mstate PPState -> IO (PPState, Either SomeException Port)
loop
  where
    loop :: PPState -> IO (PPState, Either SomeException Port)
    loop :: PPState -> IO (PPState, Either SomeException Port)
loop PPState {[Port]
[Port] -> [Port]
ppRecycled :: [Port] -> [Port]
ppAvail :: [Port]
ppRecycled :: PPState -> [Port] -> [Port]
ppAvail :: PPState -> [Port]
..} =
        case [Port]
ppAvail of
            Port
p:[Port]
ps -> do
                let next :: PPState
next = [Port] -> ([Port] -> [Port]) -> PPState
PPState [Port]
ps [Port] -> [Port]
ppRecycled
                Either SomeException Socket
res <- IO Socket -> IO (Either SomeException Socket)
forall e a. Exception e => IO a -> IO (Either e a)
try (IO Socket -> IO (Either SomeException Socket))
-> IO Socket -> IO (Either SomeException Socket)
forall a b. (a -> b) -> a -> b
$ ServiceName -> IO Socket
listenOn (ServiceName -> IO Socket) -> ServiceName -> IO Socket
forall a b. (a -> b) -> a -> b
$ Port -> ServiceName
forall a. Show a => a -> ServiceName
show Port
p
                case Either SomeException Socket
res of
                    Left (SomeException
_ :: SomeException) -> do
                        LogMessage -> IO ()
log (LogMessage -> IO ()) -> LogMessage -> IO ()
forall a b. (a -> b) -> a -> b
$ Port -> LogMessage
RemovingPort Port
p
                        PPState -> IO (PPState, Either SomeException Port)
loop PPState
next
                    Right Socket
socket' -> do
                        Either SomeException ()
res' <- IO () -> IO (Either SomeException ())
forall e a. Exception e => IO a -> IO (Either e a)
try (IO () -> IO (Either SomeException ()))
-> IO () -> IO (Either SomeException ())
forall a b. (a -> b) -> a -> b
$ Socket -> IO ()
close Socket
socket'
                        case Either SomeException ()
res' of
                            Left SomeException
e -> do
                                ServiceName
ServiceName -> Text
Text -> SomeException -> LogMessage
(LogMessage -> IO ())
-> (SomeException -> LogMessage) -> SomeException -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
pack :: ServiceName -> Text
. :: forall b c a. (b -> c) -> (a -> b) -> a -> c
$logEx LogMessage -> IO ()
log SomeException
e
                                LogMessage -> IO ()
log (LogMessage -> IO ()) -> LogMessage -> IO ()
forall a b. (a -> b) -> a -> b
$ Port -> LogMessage
RemovingPort Port
p
                                PPState -> IO (PPState, Either SomeException Port)
loop PPState
next
                            Right () -> (PPState, Either SomeException Port)
-> IO (PPState, Either SomeException Port)
forall (m :: * -> *) a. Monad m => a -> m a
return (PPState
next, Port -> Either SomeException Port
forall a b. b -> Either a b
Right Port
p)
            [] ->
                case [Port] -> [Port]
ppRecycled [] of
                    [] -> (PPState, Either SomeException Port)
-> IO (PPState, Either SomeException Port)
forall (m :: * -> *) a. Monad m => a -> m a
return ([Port] -> ([Port] -> [Port]) -> PPState
PPState [] [Port] -> [Port]
forall a. a -> a
id, SomeException -> Either SomeException Port
forall a b. a -> Either a b
Left (SomeException -> Either SomeException Port)
-> SomeException -> Either SomeException Port
forall a b. (a -> b) -> a -> b
$ KeterException -> SomeException
forall e. Exception e => e -> SomeException
toException KeterException
NoPortsAvailable)
                    [Port]
ps -> PPState -> IO (PPState, Either SomeException Port)
loop (PPState -> IO (PPState, Either SomeException Port))
-> PPState -> IO (PPState, Either SomeException Port)
forall a b. (a -> b) -> a -> b
$ [Port] -> ([Port] -> [Port]) -> PPState
PPState [Port]
ps [Port] -> [Port]
forall a. a -> a
id

    listenOn :: ServiceName -> IO Socket
listenOn ServiceName
port = do
        let hints :: AddrInfo
hints = AddrInfo
defaultHints {
                addrFlags :: [AddrInfoFlag]
addrFlags = [AddrInfoFlag
AI_PASSIVE]
              , addrSocketType :: SocketType
addrSocketType = SocketType
Stream
              }
        AddrInfo
addr:[AddrInfo]
_ <- Maybe AddrInfo
-> Maybe ServiceName -> Maybe ServiceName -> IO [AddrInfo]
getAddrInfo (AddrInfo -> Maybe AddrInfo
forall a. a -> Maybe a
Just AddrInfo
hints) Maybe ServiceName
forall a. Maybe a
Nothing (ServiceName -> Maybe ServiceName
forall a. a -> Maybe a
Just ServiceName
port)
        IO Socket
-> (Socket -> IO ()) -> (Socket -> IO Socket) -> IO Socket
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracketOnError
             (Family -> SocketType -> ProtocolNumber -> IO Socket
socket (AddrInfo -> Family
addrFamily AddrInfo
addr) (AddrInfo -> SocketType
addrSocketType AddrInfo
addr) (AddrInfo -> ProtocolNumber
addrProtocol AddrInfo
addr))
             (Socket -> IO ()
close)
             (\Socket
sock -> do
                 Socket -> SocketOption -> Port -> IO ()
setSocketOption Socket
sock SocketOption
ReuseAddr Port
1
                 Socket -> SockAddr -> IO ()
bind Socket
sock (AddrInfo -> SockAddr
addrAddress AddrInfo
addr)
                 Socket -> Port -> IO ()
listen Socket
sock Port
maxListenQueue
                 Socket -> IO Socket
forall (m :: * -> *) a. Monad m => a -> m a
return Socket
sock
             )

-- | Return a port to the recycled collection of the pool.  Note that recycling
-- puts the new ports at the end of the queue (FIFO), so that if an application
-- holds onto the port longer than expected, there should be no issues.
releasePort :: PortPool -> Port -> IO ()
releasePort :: PortPool -> Port -> IO ()
releasePort (PortPool MVar PPState
mstate) Port
p =
    MVar PPState -> (PPState -> IO PPState) -> IO ()
forall a. MVar a -> (a -> IO a) -> IO ()
modifyMVar_ MVar PPState
mstate ((PPState -> IO PPState) -> IO ())
-> (PPState -> IO PPState) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(PPState [Port]
avail [Port] -> [Port]
recycled) -> PPState -> IO PPState
forall (m :: * -> *) a. Monad m => a -> m a
return (PPState -> IO PPState) -> PPState -> IO PPState
forall a b. (a -> b) -> a -> b
$ [Port] -> ([Port] -> [Port]) -> PPState
PPState [Port]
avail (([Port] -> [Port]) -> PPState) -> ([Port] -> [Port]) -> PPState
forall a b. (a -> b) -> a -> b
$ [Port] -> [Port]
recycled ([Port] -> [Port]) -> ([Port] -> [Port]) -> [Port] -> [Port]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Port
pPort -> [Port] -> [Port]
forall a. a -> [a] -> [a]
:)

start :: PortSettings -> IO PortPool
start :: PortSettings -> IO PortPool
start PortSettings{[Port]
portRange :: PortSettings -> [Port]
portRange :: [Port]
..} =
    MVar PPState -> PortPool
PortPool (MVar PPState -> PortPool) -> IO (MVar PPState) -> IO PortPool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> PPState -> IO (MVar PPState)
forall a. a -> IO (MVar a)
newMVar PPState
freshState
  where
    freshState :: PPState
freshState = [Port] -> ([Port] -> [Port]) -> PPState
PPState [Port]
portRange [Port] -> [Port]
forall a. a -> a
id