{-# LANGUAGE MagicHash #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
module Z.IO.Network.TCP (
TCPClientConfig(..)
, defaultTCPClientConfig
, initTCPClient
, TCPServerConfig(..)
, defaultTCPServerConfig
, startTCPServer
) where
import Control.Concurrent.MVar
import Control.Monad
import Control.Monad.IO.Class
import Data.Primitive.PrimArray
import Foreign.Ptr
import GHC.Ptr
import Z.IO.Buffered
import Z.IO.Exception
import Z.IO.Network.SocketAddr
import Z.IO.Resource
import Z.IO.UV.FFI
import Z.IO.UV.Manager
import Data.Coerce
data TCPClientConfig = TCPClientConfig
{ TCPClientConfig -> Maybe SocketAddr
tcpClientAddr :: Maybe SocketAddr
, TCPClientConfig -> SocketAddr
tcpRemoteAddr :: SocketAddr
, TCPClientConfig -> Bool
tcpNoDelay :: Bool
}
defaultTCPClientConfig :: TCPClientConfig
defaultTCPClientConfig :: TCPClientConfig
defaultTCPClientConfig = Maybe SocketAddr -> SocketAddr -> Bool -> TCPClientConfig
TCPClientConfig Maybe SocketAddr
forall a. Maybe a
Nothing (PortNumber -> InetAddr -> SocketAddr
SocketAddrInet PortNumber
8888 InetAddr
inetLoopback) Bool
True
initTCPClient :: HasCallStack => TCPClientConfig -> Resource UVStream
initTCPClient :: TCPClientConfig -> Resource UVStream
initTCPClient TCPClientConfig{Bool
Maybe SocketAddr
SocketAddr
tcpNoDelay :: Bool
tcpRemoteAddr :: SocketAddr
tcpClientAddr :: Maybe SocketAddr
tcpNoDelay :: TCPClientConfig -> Bool
tcpRemoteAddr :: TCPClientConfig -> SocketAddr
tcpClientAddr :: TCPClientConfig -> Maybe SocketAddr
..} = do
UVManager
uvm <- IO UVManager -> Resource UVManager
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO UVManager
getUVManager
UVStream
client <- HasCallStack => UVManager -> Resource UVStream
UVManager -> Resource UVStream
initTCPStream UVManager
uvm
let hdl :: Ptr UVHandle
hdl = UVStream -> Ptr UVHandle
uvsHandle UVStream
client
IO () -> Resource ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> Resource ())
-> ((Ptr SocketAddr -> IO ()) -> IO ())
-> (Ptr SocketAddr -> IO ())
-> Resource ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SocketAddr -> (Ptr SocketAddr -> IO ()) -> IO ()
forall a. SocketAddr -> (Ptr SocketAddr -> IO a) -> IO a
withSocketAddr SocketAddr
tcpRemoteAddr ((Ptr SocketAddr -> IO ()) -> Resource ())
-> (Ptr SocketAddr -> IO ()) -> Resource ()
forall a b. (a -> b) -> a -> b
$ \ Ptr SocketAddr
targetPtr -> do
Maybe SocketAddr -> (SocketAddr -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ Maybe SocketAddr
tcpClientAddr ((SocketAddr -> IO ()) -> IO ()) -> (SocketAddr -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ SocketAddr
tcpClientAddr' ->
SocketAddr -> (Ptr SocketAddr -> IO ()) -> IO ()
forall a. SocketAddr -> (Ptr SocketAddr -> IO a) -> IO a
withSocketAddr SocketAddr
tcpClientAddr' ((Ptr SocketAddr -> IO ()) -> IO ())
-> (Ptr SocketAddr -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ Ptr SocketAddr
localPtr ->
IO CInt -> IO ()
forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (Ptr UVHandle -> Ptr SocketAddr -> CUInt -> IO CInt
uv_tcp_bind Ptr UVHandle
hdl Ptr SocketAddr
localPtr CUInt
0)
Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
tcpNoDelay (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ IO CInt -> IO ()
forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (Ptr UVHandle -> CInt -> IO CInt
uv_tcp_nodelay Ptr UVHandle
hdl CInt
1)
IO Int -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Int -> IO ())
-> ((Ptr UVLoop -> IO UVSlotUnSafe) -> IO Int)
-> (Ptr UVLoop -> IO UVSlotUnSafe)
-> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HasCallStack =>
UVManager -> (Ptr UVLoop -> IO UVSlotUnSafe) -> IO Int
UVManager -> (Ptr UVLoop -> IO UVSlotUnSafe) -> IO Int
withUVRequest UVManager
uvm ((Ptr UVLoop -> IO UVSlotUnSafe) -> IO ())
-> (Ptr UVLoop -> IO UVSlotUnSafe) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ Ptr UVLoop
_ -> Ptr UVHandle -> Ptr SocketAddr -> IO UVSlotUnSafe
hs_uv_tcp_connect Ptr UVHandle
hdl Ptr SocketAddr
targetPtr
UVStream -> Resource UVStream
forall (m :: * -> *) a. Monad m => a -> m a
return UVStream
client
data TCPServerConfig = TCPServerConfig
{ TCPServerConfig -> SocketAddr
tcpListenAddr :: SocketAddr
, TCPServerConfig -> Int
tcpListenBacklog :: Int
, TCPServerConfig -> Bool
tcpServerWorkerNoDelay :: Bool
, TCPServerConfig -> UVStream -> IO ()
tcpServerWorker :: UVStream -> IO ()
}
defaultTCPServerConfig :: TCPServerConfig
defaultTCPServerConfig :: TCPServerConfig
defaultTCPServerConfig = SocketAddr -> Int -> Bool -> (UVStream -> IO ()) -> TCPServerConfig
TCPServerConfig
(PortNumber -> InetAddr -> SocketAddr
SocketAddrInet PortNumber
8888 InetAddr
inetAny)
Int
256
Bool
True
(\ UVStream
uvs -> UVStream -> Ptr Word8 -> Int -> IO ()
forall o.
(Output o, HasCallStack) =>
o -> Ptr Word8 -> Int -> IO ()
writeOutput UVStream
uvs (Addr# -> Ptr Word8
forall a. Addr# -> Ptr a
Ptr Addr#
"hello world"#) Int
11)
startTCPServer :: HasCallStack => TCPServerConfig -> IO ()
startTCPServer :: TCPServerConfig -> IO ()
startTCPServer TCPServerConfig{Bool
Int
SocketAddr
UVStream -> IO ()
tcpServerWorker :: UVStream -> IO ()
tcpServerWorkerNoDelay :: Bool
tcpListenBacklog :: Int
tcpListenAddr :: SocketAddr
tcpServerWorker :: TCPServerConfig -> UVStream -> IO ()
tcpServerWorkerNoDelay :: TCPServerConfig -> Bool
tcpListenBacklog :: TCPServerConfig -> Int
tcpListenAddr :: TCPServerConfig -> SocketAddr
..} = do
let backLog :: Int
backLog = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
tcpListenBacklog Int
128
UVManager
serverUVManager <- IO UVManager
getUVManager
Resource UVStream -> (UVStream -> IO ()) -> IO ()
forall (m :: * -> *) a b.
(MonadMask m, MonadIO m, HasCallStack) =>
Resource a -> (a -> m b) -> m b
withResource (HasCallStack => UVManager -> Resource UVStream
UVManager -> Resource UVStream
initTCPStream UVManager
serverUVManager) ((UVStream -> IO ()) -> IO ()) -> (UVStream -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ (UVStream Ptr UVHandle
serverHandle Int
serverSlot UVManager
_ IORef Bool
_) ->
IO (Ptr UVHandle)
-> (Ptr UVHandle -> IO ()) -> (Ptr UVHandle -> IO ()) -> IO ()
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket
(IO (Ptr UVHandle) -> IO (Ptr UVHandle)
forall a. HasCallStack => IO (Ptr a) -> IO (Ptr a)
throwOOMIfNull (IO (Ptr UVHandle) -> IO (Ptr UVHandle))
-> IO (Ptr UVHandle) -> IO (Ptr UVHandle)
forall a b. (a -> b) -> a -> b
$ Ptr UVHandle -> IO (Ptr UVHandle)
hs_uv_accept_check_alloc Ptr UVHandle
serverHandle)
Ptr UVHandle -> IO ()
hs_uv_accept_check_close ((Ptr UVHandle -> IO ()) -> IO ())
-> (Ptr UVHandle -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$
\ Ptr UVHandle
check -> do
IO CInt -> IO ()
forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (IO CInt -> IO ()) -> IO CInt -> IO ()
forall a b. (a -> b) -> a -> b
$ Ptr UVHandle -> IO CInt
hs_uv_accept_check_init Ptr UVHandle
check
SocketAddr -> (Ptr SocketAddr -> IO ()) -> IO ()
forall a. SocketAddr -> (Ptr SocketAddr -> IO a) -> IO a
withSocketAddr SocketAddr
tcpListenAddr ((Ptr SocketAddr -> IO ()) -> IO ())
-> (Ptr SocketAddr -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ Ptr SocketAddr
addrPtr -> do
MVar Int
m <- UVManager -> Int -> IO (MVar Int)
getBlockMVar UVManager
serverUVManager Int
serverSlot
MutablePrimArray RealWorld Int32
acceptBuf <- Int -> IO (MutablePrimArray (PrimState IO) Int32)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPinnedPrimArray Int
backLog
let acceptBufPtr :: Ptr Word8
acceptBufPtr = Ptr Int32 -> Ptr Word8
coerce (MutablePrimArray RealWorld Int32 -> Ptr Int32
forall s a. MutablePrimArray s a -> Ptr a
mutablePrimArrayContents MutablePrimArray RealWorld Int32
acceptBuf :: Ptr UVFD)
UVManager -> IO () -> IO ()
forall a. HasCallStack => UVManager -> IO a -> IO a
withUVManager' UVManager
serverUVManager (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
UVManager -> Int -> Ptr Word8 -> Int -> IO ()
pokeBufferTable UVManager
serverUVManager Int
serverSlot Ptr Word8
acceptBufPtr (Int
backLogInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
IO CInt -> IO ()
forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (Ptr UVHandle -> Ptr SocketAddr -> CUInt -> IO CInt
uv_tcp_bind Ptr UVHandle
serverHandle Ptr SocketAddr
addrPtr CUInt
0)
IO CInt -> IO ()
forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (Ptr UVHandle -> CInt -> IO CInt
hs_uv_listen Ptr UVHandle
serverHandle (CInt -> CInt -> CInt
forall a. Ord a => a -> a -> a
max CInt
4 (Int -> CInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
backLog)))
IO () -> IO ()
forall (f :: * -> *) a b. Applicative f => f a -> f b
forever (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
!Int
acceptCountDown <- MVar Int -> IO Int
forall a. MVar a -> IO a
takeMVar MVar Int
m
PrimArray Int32
acceptBufCopy <- UVManager -> IO (PrimArray Int32) -> IO (PrimArray Int32)
forall a. HasCallStack => UVManager -> IO a -> IO a
withUVManager' UVManager
serverUVManager (IO (PrimArray Int32) -> IO (PrimArray Int32))
-> IO (PrimArray Int32) -> IO (PrimArray Int32)
forall a b. (a -> b) -> a -> b
$ do
Maybe Int
_ <- MVar Int -> IO (Maybe Int)
forall a. MVar a -> IO (Maybe a)
tryTakeMVar MVar Int
m
UVManager -> Int -> Ptr Word8 -> Int -> IO ()
pokeBufferTable UVManager
serverUVManager Int
serverSlot Ptr Word8
acceptBufPtr (Int
backLogInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
acceptCountDown Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== -Int
1) (Ptr UVHandle -> IO ()
hs_uv_listen_resume Ptr UVHandle
serverHandle)
let acceptCount :: Int
acceptCount = Int
backLog Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
acceptCountDown
MutablePrimArray RealWorld Int32
acceptBuf' <- Int -> IO (MutablePrimArray (PrimState IO) Int32)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
acceptCount
MutablePrimArray (PrimState IO) Int32
-> Int
-> MutablePrimArray (PrimState IO) Int32
-> Int
-> Int
-> IO ()
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> MutablePrimArray (PrimState m) a -> Int -> Int -> m ()
copyMutablePrimArray MutablePrimArray RealWorld Int32
MutablePrimArray (PrimState IO) Int32
acceptBuf' Int
0 MutablePrimArray RealWorld Int32
MutablePrimArray (PrimState IO) Int32
acceptBuf (Int
acceptCountDownInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
acceptCount
MutablePrimArray (PrimState IO) Int32 -> IO (PrimArray Int32)
forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray RealWorld Int32
MutablePrimArray (PrimState IO) Int32
acceptBuf'
[Int] -> (Int -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0..PrimArray Int32 -> Int
forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray Int32
acceptBufCopyInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1] ((Int -> IO ()) -> IO ()) -> (Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ Int
i -> do
let fd :: Int32
fd = PrimArray Int32 -> Int -> Int32
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Int32
acceptBufCopy Int
i
if Int32
fd Int32 -> Int32 -> Bool
forall a. Ord a => a -> a -> Bool
< Int32
0
then IO Int32 -> IO ()
forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (Int32 -> IO Int32
forall (m :: * -> *) a. Monad m => a -> m a
return Int32
fd)
else IO ThreadId -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO ThreadId -> IO ()) -> (IO () -> IO ThreadId) -> IO () -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO () -> IO ThreadId
forkBa (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
UVManager
uvm <- IO UVManager
getUVManager
Resource UVStream -> (UVStream -> IO ()) -> IO ()
forall (m :: * -> *) a b.
(MonadMask m, MonadIO m, HasCallStack) =>
Resource a -> (a -> m b) -> m b
withResource (HasCallStack =>
(Ptr UVLoop -> Ptr UVHandle -> IO ())
-> UVManager -> Resource UVStream
(Ptr UVLoop -> Ptr UVHandle -> IO ())
-> UVManager -> Resource UVStream
initUVStream (\ Ptr UVLoop
loop Ptr UVHandle
hdl -> do
IO CInt -> IO ()
forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (Ptr UVLoop -> Ptr UVHandle -> IO CInt
uv_tcp_init Ptr UVLoop
loop Ptr UVHandle
hdl)
IO CInt -> IO ()
forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (Ptr UVHandle -> Int32 -> IO CInt
hs_uv_tcp_open Ptr UVHandle
hdl Int32
fd)) UVManager
uvm) ((UVStream -> IO ()) -> IO ()) -> (UVStream -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ UVStream
uvs -> do
Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
tcpServerWorkerNoDelay (IO () -> IO ()) -> (IO CInt -> IO ()) -> IO CInt -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO CInt -> IO ()
forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (IO CInt -> IO ()) -> IO CInt -> IO ()
forall a b. (a -> b) -> a -> b
$
Ptr UVHandle -> CInt -> IO CInt
uv_tcp_nodelay (UVStream -> Ptr UVHandle
uvsHandle UVStream
uvs) CInt
1
UVStream -> IO ()
tcpServerWorker UVStream
uvs
initTCPStream :: HasCallStack => UVManager -> Resource UVStream
initTCPStream :: UVManager -> Resource UVStream
initTCPStream = HasCallStack =>
(Ptr UVLoop -> Ptr UVHandle -> IO ())
-> UVManager -> Resource UVStream
(Ptr UVLoop -> Ptr UVHandle -> IO ())
-> UVManager -> Resource UVStream
initUVStream (\ Ptr UVLoop
loop Ptr UVHandle
hdl ->
IO CInt -> IO ()
forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (Ptr UVLoop -> Ptr UVHandle -> IO CInt
uv_tcp_init Ptr UVLoop
loop Ptr UVHandle
hdl))