{-# 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 Control.Concurrent.MVar
import Control.Exception
import Control.Monad.IO.Class (liftIO)
import Control.Monad.IO.Unlift (withRunInIO)
import Control.Monad.Logger
import Data.Text (pack)
import Keter.Common
import Keter.Config
import Keter.Context
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 a. KeterM cfg a -> IO a)
 -> IO (Either SomeException Port))
-> KeterM cfg (Either SomeException Port)
forall b.
((forall a. KeterM cfg a -> IO a) -> IO b) -> KeterM cfg b
forall (m :: * -> *) b.
MonadUnliftIO m =>
((forall a. m a -> IO a) -> IO b) -> m b
withRunInIO (((forall a. KeterM cfg a -> IO a)
  -> IO (Either SomeException Port))
 -> KeterM cfg (Either SomeException Port))
-> ((forall a. KeterM cfg a -> IO a)
    -> IO (Either SomeException Port))
-> KeterM cfg (Either SomeException Port)
forall a b. (a -> b) -> a -> b
$ \forall a. KeterM cfg a -> IO a
rio -> 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 (KeterM cfg (PPState, Either SomeException Port)
-> IO (PPState, Either SomeException Port)
forall a. KeterM cfg a -> IO a
rio (KeterM cfg (PPState, Either SomeException Port)
 -> IO (PPState, Either SomeException Port))
-> (PPState -> KeterM cfg (PPState, Either SomeException Port))
-> PPState
-> IO (PPState, Either SomeException Port)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PPState -> KeterM cfg (PPState, Either SomeException Port)
forall cfg.
PPState -> KeterM cfg (PPState, Either SomeException Port)
loop)
  where
    removePortMsg :: a -> Text
removePortMsg a
p = String -> Text
pack (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$
        String
"Port in use, removing from port pool: "
        String -> String -> String
forall a. [a] -> [a] -> [a]
++ a -> String
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]
ppAvail :: PPState -> [Port]
ppRecycled :: PPState -> [Port] -> [Port]
ppAvail :: [Port]
ppRecycled :: [Port] -> [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 (Either SomeException Socket)
-> KeterM cfg (Either SomeException Socket)
forall a. IO a -> KeterM cfg a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Either SomeException Socket)
 -> KeterM cfg (Either SomeException Socket))
-> IO (Either SomeException Socket)
-> KeterM cfg (Either SomeException Socket)
forall a b. (a -> b) -> a -> b
$ 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
$ String -> IO Socket
listenOn (String -> IO Socket) -> String -> IO Socket
forall a b. (a -> b) -> a -> b
$ Port -> String
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 -> Text
Loc -> Text -> LogLevel -> Text -> KeterM cfg ()
(Text -> KeterM cfg ()) -> (Text -> Text) -> Text -> KeterM cfg ()
forall a. a -> a
forall msg.
ToLogStr msg =>
Loc -> Text -> LogLevel -> msg -> KeterM cfg ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall (m :: * -> *) msg.
(MonadLogger m, ToLogStr msg) =>
Loc -> Text -> LogLevel -> msg -> m ()
pack :: String -> Text
. :: forall b c a. (b -> c) -> (a -> b) -> a -> c
id :: forall a. a -> a
monadLoggerLog :: forall (m :: * -> *) msg.
(MonadLogger m, ToLogStr msg) =>
Loc -> Text -> LogLevel -> msg -> m ()
logInfo (Text -> KeterM cfg ()) -> Text -> KeterM cfg ()
forall a b. (a -> b) -> a -> b
$ Port -> Text
forall {a}. Show a => a -> Text
removePortMsg Port
p
                        PPState -> KeterM cfg (PPState, Either SomeException Port)
forall cfg.
PPState -> KeterM cfg (PPState, Either SomeException Port)
loop PPState
next
                    Right Socket
socket' -> do
                        Either SomeException ()
res' <- IO (Either SomeException ())
-> KeterM cfg (Either SomeException ())
forall a. IO a -> KeterM cfg a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Either SomeException ())
 -> KeterM cfg (Either SomeException ()))
-> IO (Either SomeException ())
-> KeterM cfg (Either SomeException ())
forall a b. (a -> b) -> a -> b
$ forall e a. Exception e => IO a -> IO (Either e a)
try @SomeException (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
                                $Port
String
LogLevel
String -> Text
String -> String -> String -> CharPos -> CharPos -> Loc
Text -> Text
Loc -> Text -> LogLevel -> Text -> KeterM cfg ()
(Text -> KeterM cfg ()) -> (Text -> Text) -> Text -> KeterM cfg ()
forall a. a -> a
forall msg.
ToLogStr msg =>
Loc -> Text -> LogLevel -> msg -> KeterM cfg ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall (m :: * -> *) msg.
(MonadLogger m, ToLogStr msg) =>
Loc -> Text -> LogLevel -> msg -> m ()
pack :: String -> Text
. :: forall b c a. (b -> c) -> (a -> b) -> a -> c
id :: forall a. a -> a
monadLoggerLog :: forall (m :: * -> *) msg.
(MonadLogger m, ToLogStr msg) =>
Loc -> Text -> LogLevel -> msg -> m ()
logError (Text -> KeterM cfg ()) -> Text -> KeterM cfg ()
forall a b. (a -> b) -> a -> b
$ String -> Text
pack (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$ SomeException -> String
forall a. Show a => a -> String
show SomeException
e
                                $Port
String
LogLevel
String -> Text
String -> String -> String -> CharPos -> CharPos -> Loc
Text -> Text
Loc -> Text -> LogLevel -> Text -> KeterM cfg ()
(Text -> KeterM cfg ()) -> (Text -> Text) -> Text -> KeterM cfg ()
forall a. a -> a
forall msg.
ToLogStr msg =>
Loc -> Text -> LogLevel -> msg -> KeterM cfg ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall (m :: * -> *) msg.
(MonadLogger m, ToLogStr msg) =>
Loc -> Text -> LogLevel -> msg -> m ()
pack :: String -> Text
. :: forall b c a. (b -> c) -> (a -> b) -> a -> c
id :: forall a. a -> a
monadLoggerLog :: forall (m :: * -> *) msg.
(MonadLogger m, ToLogStr msg) =>
Loc -> Text -> LogLevel -> msg -> m ()
logInfo (Text -> KeterM cfg ()) -> Text -> KeterM cfg ()
forall a b. (a -> b) -> a -> b
$ Port -> Text
forall {a}. Show a => a -> Text
removePortMsg Port
p
                                PPState -> KeterM cfg (PPState, Either SomeException Port)
forall cfg.
PPState -> KeterM cfg (PPState, Either SomeException Port)
loop PPState
next
                            Right () -> (PPState, Either SomeException Port)
-> KeterM cfg (PPState, Either SomeException Port)
forall a. a -> KeterM cfg a
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)
-> KeterM cfg (PPState, Either SomeException Port)
forall a. a -> KeterM cfg a
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 -> KeterM cfg (PPState, Either SomeException Port)
forall cfg.
PPState -> KeterM cfg (PPState, Either SomeException Port)
loop (PPState -> KeterM cfg (PPState, Either SomeException Port))
-> PPState -> KeterM cfg (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 :: String -> IO Socket
listenOn String
port = do
        let hints :: AddrInfo
hints = AddrInfo
defaultHints {
                addrFlags = [AI_PASSIVE]
              , addrSocketType = Stream
              }
        AddrInfo
addr:[AddrInfo]
_ <- Maybe AddrInfo -> Maybe String -> Maybe String -> IO [AddrInfo]
forall (t :: * -> *).
GetAddrInfo t =>
Maybe AddrInfo -> Maybe String -> Maybe String -> IO (t AddrInfo)
getAddrInfo (AddrInfo -> Maybe AddrInfo
forall a. a -> Maybe a
Just AddrInfo
hints) Maybe String
forall a. Maybe a
Nothing (String -> Maybe String
forall a. a -> Maybe a
Just String
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 a. a -> IO a
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 a. a -> IO a
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 :: [Port]
portRange :: PortSettings -> [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