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

import           Keter.Common
import           Keter.Context
import           Data.Text           (pack)
import           Control.Applicative     ((<$>))
import           Control.Concurrent.MVar
import           Control.Exception
import           Control.Monad.IO.Class (liftIO)
import           Control.Monad.IO.Unlift (withRunInIO)
import           Control.Monad.Logger
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 :: PortPool
        -> KeterM cfg (Either SomeException Port)
getPort :: forall cfg. PortPool -> KeterM cfg (Either SomeException Port)
getPort (PortPool MVar PPState
mstate) =
    forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> IO a) -> IO b) -> m b
withRunInIO forall a b. (a -> b) -> a -> b
$ \forall a. KeterM cfg a -> IO a
rio -> forall a b. MVar a -> (a -> IO (a, b)) -> IO b
modifyMVar MVar PPState
mstate (forall a. KeterM cfg a -> IO a
rio forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall cfg.
PPState -> KeterM cfg (PPState, Either SomeException Port)
loop)
  where
    removePortMsg :: a -> Text
removePortMsg a
p = String -> Text
pack forall a b. (a -> b) -> a -> b
$
        String
"Port in use, removing from port pool: "
        forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show a
p

    loop :: PPState -> KeterM cfg (PPState, Either SomeException Port)
    loop :: forall cfg.
PPState -> KeterM cfg (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 <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall e a. Exception e => IO a -> IO (Either e a)
try forall a b. (a -> b) -> a -> b
$ String -> IO Socket
listenOn forall a b. (a -> b) -> a -> b
$ forall a. Show a => a -> String
show Port
p
                case Either SomeException Socket
res of
                    Left (SomeException
_ :: SomeException) -> do
                        $Port
String
LogLevel
String -> Text
String -> String -> String -> CharPos -> CharPos -> Loc
Text -> KeterM cfg ()
forall a. a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall (m :: * -> *) msg.
(MonadLogger m, ToLogStr msg) =>
Loc -> Text -> LogLevel -> msg -> m ()
monadLoggerLog :: forall (m :: * -> *) msg.
(MonadLogger m, ToLogStr msg) =>
Loc -> Text -> LogLevel -> msg -> m ()
id :: forall a. a -> a
. :: forall b c a. (b -> c) -> (a -> b) -> a -> c
pack :: String -> Text
logInfo forall a b. (a -> b) -> a -> b
$ forall {a}. Show a => a -> Text
removePortMsg Port
p
                        forall cfg.
PPState -> KeterM cfg (PPState, Either SomeException Port)
loop PPState
next
                    Right Socket
socket' -> do
                        Either SomeException ()
res' <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall e a. Exception e => IO a -> IO (Either e a)
try @SomeException forall a b. (a -> b) -> a -> b
$ Socket -> IO ()
close Socket
socket'
                        case Either SomeException ()
res' of
                            Left SomeException
e -> do
                                $Port
String
LogLevel
String -> Text
String -> String -> String -> CharPos -> CharPos -> Loc
Text -> KeterM cfg ()
forall a. a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall (m :: * -> *) msg.
(MonadLogger m, ToLogStr msg) =>
Loc -> Text -> LogLevel -> msg -> m ()
monadLoggerLog :: forall (m :: * -> *) msg.
(MonadLogger m, ToLogStr msg) =>
Loc -> Text -> LogLevel -> msg -> m ()
id :: forall a. a -> a
. :: forall b c a. (b -> c) -> (a -> b) -> a -> c
pack :: String -> Text
logError forall a b. (a -> b) -> a -> b
$ String -> Text
pack forall a b. (a -> b) -> a -> b
$ forall a. Show a => a -> String
show SomeException
e
                                $Port
String
LogLevel
String -> Text
String -> String -> String -> CharPos -> CharPos -> Loc
Text -> KeterM cfg ()
forall a. a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall (m :: * -> *) msg.
(MonadLogger m, ToLogStr msg) =>
Loc -> Text -> LogLevel -> msg -> m ()
monadLoggerLog :: forall (m :: * -> *) msg.
(MonadLogger m, ToLogStr msg) =>
Loc -> Text -> LogLevel -> msg -> m ()
id :: forall a. a -> a
. :: forall b c a. (b -> c) -> (a -> b) -> a -> c
pack :: String -> Text
logInfo forall a b. (a -> b) -> a -> b
$ forall {a}. Show a => a -> Text
removePortMsg Port
p
                                forall cfg.
PPState -> KeterM cfg (PPState, Either SomeException Port)
loop PPState
next
                            Right () -> forall (m :: * -> *) a. Monad m => a -> m a
return (PPState
next, forall a b. b -> Either a b
Right Port
p)
            [] ->
                case [Port] -> [Port]
ppRecycled [] of
                    [] -> forall (m :: * -> *) a. Monad m => a -> m a
return ([Port] -> ([Port] -> [Port]) -> PPState
PPState [] forall a. a -> a
id, forall a b. a -> Either a b
Left forall a b. (a -> b) -> a -> b
$ forall e. Exception e => e -> SomeException
toException KeterException
NoPortsAvailable)
                    [Port]
ps -> forall cfg.
PPState -> KeterM cfg (PPState, Either SomeException Port)
loop forall a b. (a -> b) -> a -> b
$ [Port] -> ([Port] -> [Port]) -> PPState
PPState [Port]
ps forall a. a -> a
id

    listenOn :: String -> IO Socket
listenOn String
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 String -> Maybe String -> IO [AddrInfo]
getAddrInfo (forall a. a -> Maybe a
Just AddrInfo
hints) forall a. Maybe a
Nothing (forall a. a -> Maybe a
Just String
port)
        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
                 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 =
    forall a. MVar a -> (a -> IO a) -> IO ()
modifyMVar_ MVar PPState
mstate forall a b. (a -> b) -> a -> b
$ \(PPState [Port]
avail [Port] -> [Port]
recycled) -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ [Port] -> ([Port] -> [Port]) -> PPState
PPState [Port]
avail forall a b. (a -> b) -> a -> b
$ [Port] -> [Port]
recycled forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Port
pforall a. a -> [a] -> [a]
:)

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