@@ -0,0 +1,393 @@
2+{-# LANGUAGE BangPatterns #-}
3+{-# LANGUAGE FlexibleContexts #-}
4+{-# LANGUAGE MagicHash #-}
5+{-# LANGUAGE RecordWildCards #-}
6+{-# LANGUAGE UnboxedTuples #-}
11+-- Module : Streamly.Network.Socket
12+-- Copyright : (c) 2018 Harendra Kumar
15+-- Maintainer : harendra.kumar@gmail.com
16+-- Stability : experimental
19+-- Read and write streams and arrays to and from network sockets. Socket IO
20+-- APIs are quite similar to "Streamly.Mem.Array" read write APIs and almost
21+-- identical to the sequential streaming APIs in "Streamly.FileSystem.File".
23+-- Read IO requests to the socket are performed in chunks of 32KiB, this is
24+-- referred to as @defaultChunkSize@ in the documentation. One IO request may
25+-- or may not read the full chunk. Unless specified otherwise in the API,
26+-- writes are collected into chunks of @defaultChunkSize@ before they are
27+-- written to the socket. APIs are provided to control the chunking and
30+-- > import qualified Streamly.Network.Socket as SK
33+module Streamly.Network.Socket
35+ -- ** Listen for Connections
37+ , recvConnectionsWith
40+ -- ** Read a stream from a connection
55+ -- ** Write a stream to a connection
58+ -- , writeUtf8ByLines
68+import Control.Concurrent (threadWaitWrite, rtsSupportsBoundThreads)
69+import Control.Monad.IO.Class (MonadIO(..))
70+import Control.Monad (when)
71+import Data.Word (Word8)
72+import Foreign.ForeignPtr (withForeignPtr)
73+import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
74+import Foreign.Ptr (minusPtr, plusPtr, Ptr, castPtr)
75+import Foreign.Storable (Storable(..))
76+import GHC.ForeignPtr (mallocPlainForeignPtrBytes)
77+import Prelude hiding (read)
79+import Streamly.Mem.Array.Types (Array(..))
80+import Streamly.Streams.Serial (SerialT)
81+import Streamly.Streams.StreamK.Type (IsStream, mkStream)
82+-- import Streamly.Fold (Fold)
83+-- import Streamly.String (encodeUtf8, decodeUtf8, foldLines)
85+import qualified Streamly.Mem.Array as A
86+import qualified Streamly.Mem.Array.Types as A hiding (flattenArrays)
87+import qualified Streamly.Prelude as S
89+import Network.Socket hiding (listen)
90+import qualified Network.Socket as Net
91+import Streamly (MonadAsync)
93+-------------------------------------------------------------------------------
95+-------------------------------------------------------------------------------
97+-- | Specify the configuration of a server.
98+data ServerSpec = ServerSpec
100+ serverAddressFamily :: !Family
101+ , serverAddress :: !SockAddr
102+ , serverSockOpts :: ![(SocketOption, Int)]
105+initListener :: Int -> ServerSpec -> IO Socket
106+initListener tcpListenQ ServerSpec{..} =
108+ sock <- socket serverAddressFamily Stream defaultProtocol
109+ mapM_ (\(opt, val) -> setSocketOption sock opt val) serverSockOpts
110+ bind sock serverAddress
111+ Net.listen sock tcpListenQ
114+-- | Start a TCP stream server. The server listens for connections on the
115+-- supplied address (address family, local interface IP address and port) and
116+-- generates a stream of connected sockets and the endpoint they are connected
117+-- to. The first argument is the maximum number of pending connections in the
119+{-# INLINE recvConnectionsWith #-}
120+recvConnectionsWith :: MonadAsync m => Int -> ServerSpec -> SerialT m (Socket, SockAddr)
121+recvConnectionsWith tcpListenQ opts = S.unfoldrM step Nothing
124+ listener <- liftIO $ initListener tcpListenQ opts
125+ r <- liftIO $ accept listener
126+ -- XXX error handling
127+ return $ Just (r, Just listener)
129+ step (Just listener) = do
130+ r <- liftIO $ accept listener
131+ -- XXX error handling
132+ return $ Just (r, Just listener)
134+{-# INLINE recvConnectionsOn #-}
135+recvConnectionsOn :: MonadAsync m => PortNumber -> SerialT m (Socket, SockAddr)
136+recvConnectionsOn port =
137+ recvConnectionsWith maxListenQueue ServerSpec
138+ { serverAddressFamily = AF_INET
139+ , serverAddress = SockAddrInet port 0
140+ , serverSockOpts = [(NoDelay,1), (ReuseAddr,1)]
143+-------------------------------------------------------------------------------
145+-------------------------------------------------------------------------------
147+{-# INLINABLE readArrayUptoWith #-}
149+ :: (h -> Ptr Word8 -> Int -> IO Int)
153+readArrayUptoWith f size h = do
154+ ptr <- mallocPlainForeignPtrBytes size
155+ -- ptr <- mallocPlainForeignPtrAlignedBytes size (alignment (undefined :: Word8))
156+ withForeignPtr ptr $ \p -> do
160+ , aEnd = p `plusPtr` n
161+ , aBound = p `plusPtr` size
163+ -- XXX shrink only if the diff is significant
167+-- | Read a 'ByteArray' from a file handle. If no data is available on the
168+-- handle it blocks until some data becomes available. If data is available
169+-- then it immediately returns that data without blocking. It reads a maximum
170+-- of up to the size requested.
171+{-# INLINABLE readArrayUpto #-}
172+readArrayUpto :: Int -> Socket -> IO (Array Word8)
173+readArrayUpto = readArrayUptoWith recvBuf
175+-------------------------------------------------------------------------------
177+-------------------------------------------------------------------------------
179+waitWhen0 :: Int -> Socket -> IO ()
180+waitWhen0 0 s = when rtsSupportsBoundThreads $
181+ withFdSocket s $ \fd -> threadWaitWrite $ fromIntegral fd
182+waitWhen0 _ _ = return ()
184+sendAll :: Socket -> Ptr Word8 -> Int -> IO ()
185+sendAll _ _ len | len <= 0 = return ()
187+ sent <- sendBuf s p len
189+ -- assert (sent <= len)
190+ when (sent >= 0) $ sendAll s (p `plusPtr` sent) (len - sent)
192+{-# INLINABLE writeArrayWith #-}
193+writeArrayWith :: Storable a
194+ => (h -> Ptr Word8 -> Int -> IO ())
198+writeArrayWith _ _ arr | A.length arr == 0 = return ()
199+writeArrayWith f h Array{..} = withForeignPtr aStart $ \p ->
200+ f h (castPtr p) aLen
203+ let p = unsafeForeignPtrToPtr aStart
204+ in aEnd `minusPtr` p
206+-- | Write an Array to a file handle.
209+{-# INLINABLE writeArray #-}
210+writeArray :: Storable a => Socket -> Array a -> IO ()
211+writeArray = writeArrayWith sendAll
213+-------------------------------------------------------------------------------
214+-- Stream of Arrays IO
215+-------------------------------------------------------------------------------
217+{-# INLINABLE readArraysUptoWith #-}
218+readArraysUptoWith :: (IsStream t, MonadIO m)
219+ => (Int -> h -> IO (Array Word8))
220+ -> Int -> h -> t m (Array Word8)
221+readArraysUptoWith f size h = go
223+ -- XXX use cons/nil instead
224+ go = mkStream $ \_ yld _ stp -> do
225+ arr <- liftIO $ f size h
226+ if A.length arr == 0
230+-- | @readArraysUpto size h@ reads a stream of arrays from file handle @h@.
231+-- The maximum size of a single array is limited to @size@.
232+-- 'fromHandleArraysUpto' ignores the prevailing 'TextEncoding' and 'NewlineMode'
234+{-# INLINABLE readArraysUpto #-}
235+readArraysUpto :: (IsStream t, MonadIO m)
236+ => Int -> Socket -> t m (Array Word8)
237+readArraysUpto = readArraysUptoWith readArrayUpto
239+-- XXX read 'Array a' instead of Word8
241+-- | @readArrays h@ reads a stream of arrays from file handle @h@.
242+-- The maximum size of a single array is limited to @defaultChunkSize@.
243+-- 'readArrays' ignores the prevailing 'TextEncoding' and 'NewlineMode'
247+{-# INLINE readArrays #-}
248+readArrays :: (IsStream t, MonadIO m) => Socket -> t m (Array Word8)
249+readArrays = readArraysUpto A.defaultChunkSize
251+-------------------------------------------------------------------------------
252+-- Read File to Stream
253+-------------------------------------------------------------------------------
255+-- TODO for concurrent streams implement readahead IO. We can send multiple
256+-- read requests at the same time. For serial case we can use async IO. We can
257+-- also control the read throughput in mbps or IOPS.
260+-- | @readByChunksUpto chunkSize handle@ reads a byte stream from a file
261+-- handle, reads are performed in chunks of up to @chunkSize@. The stream ends
262+-- as soon as EOF is encountered.
264+{-# INLINE readByChunksUpto #-}
265+readByChunksUpto :: (IsStream t, MonadIO m) => Int -> Handle -> t m Word8
266+readByChunksUpto chunkSize h = A.flattenArrays $ readArraysUpto chunkSize h
270+-- read :: (IsStream t, MonadIO m, Storable a) => Handle -> t m a
272+-- > read = 'readByChunks' A.defaultChunkSize
273+-- | Generate a stream of elements of the given type from a file 'Handle'. The
274+-- stream ends when EOF is encountered.
278+read :: (IsStream t, MonadIO m) => Socket -> t m Word8
279+read = A.flattenArrays . readArrays
281+-------------------------------------------------------------------------------
283+-------------------------------------------------------------------------------
285+-- | Write a stream of arrays to a handle.
288+{-# INLINE writeArrays #-}
289+writeArrays :: (MonadIO m, Storable a) => Socket -> SerialT m (Array a) -> m ()
290+writeArrays h m = S.mapM_ (liftIO . writeArray h) m
292+-- GHC buffer size dEFAULT_FD_BUFFER_SIZE=8192 bytes.
295+-- Note that if you use a chunk size less than 8K (GHC's default buffer
296+-- size) then you are advised to use 'NOBuffering' mode on the 'Handle' in case you
297+-- do not want buffering to occur at GHC level as well. Same thing applies to
300+-- | Like 'write' but provides control over the write buffer. Output will
301+-- be written to the IO device as soon as we collect the specified number of
305+{-# INLINE writeByChunks #-}
306+writeByChunks :: MonadIO m => Int -> Socket -> SerialT m Word8 -> m ()
307+writeByChunks n h m = writeArrays h $ A.arraysOf n m
309+-- > write = 'writeByChunks' A.defaultChunkSize
311+-- | Write a byte stream to a file handle. Combines the bytes in chunks of size
312+-- up to 'A.defaultChunkSize' before writing. Note that the write behavior
313+-- depends on the 'IOMode' and the current seek position of the handle.
317+write :: MonadIO m => Socket -> SerialT m Word8 -> m ()
318+write = writeByChunks A.defaultChunkSize
322+write :: (MonadIO m, Storable a) => Handle -> SerialT m a -> m ()
323+write = toHandleWith A.defaultChunkSize
326+-------------------------------------------------------------------------------
327+-- IO with encoding/decoding Unicode characters
328+-------------------------------------------------------------------------------
332+-- > readUtf8 = decodeUtf8 . read
334+-- Read a UTF8 encoded stream of unicode characters from a file handle.
337+{-# INLINE readUtf8 #-}
338+readUtf8 :: (IsStream t, MonadIO m) => Handle -> t m Char
339+readUtf8 = decodeUtf8 . read
342+-- > writeUtf8 h s = write h $ encodeUtf8 s
344+-- Encode a stream of unicode characters to UTF8 and write it to the given file
345+-- handle. Default block buffering applies to the writes.
348+{-# INLINE writeUtf8 #-}
349+writeUtf8 :: MonadIO m => Handle -> SerialT m Char -> m ()
350+writeUtf8 h s = write h $ encodeUtf8 s
352+-- | Write a stream of unicode characters after encoding to UTF-8 in chunks
353+-- separated by a linefeed character @'\n'@. If the size of the buffer exceeds
354+-- @defaultChunkSize@ and a linefeed is not yet found, the buffer is written
355+-- anyway. This is similar to writing to a 'Handle' with the 'LineBuffering'
359+{-# INLINE writeUtf8ByLines #-}
360+writeUtf8ByLines :: (IsStream t, MonadIO m) => Handle -> t m Char -> m ()
361+writeUtf8ByLines = undefined
363+-- | Read UTF-8 lines from a file handle and apply the specified fold to each
364+-- line. This is similar to reading a 'Handle' with the 'LineBuffering' option.
367+{-# INLINE readLines #-}
368+readLines :: (IsStream t, MonadIO m) => Handle -> Fold m Char b -> t m b
369+readLines h f = foldLines (readUtf8 h) f
371+-------------------------------------------------------------------------------
372+-- Framing on a sequence
373+-------------------------------------------------------------------------------
375+-- | Read a stream from a file handle and split it into frames delimited by
376+-- the specified sequence of elements. The supplied fold is applied on each
380+{-# INLINE readFrames #-}
381+readFrames :: (IsStream t, MonadIO m, Storable a)
382+ => Array a -> Handle -> Fold m a b -> t m b
383+readFrames = undefined -- foldFrames . read
385+-- | Write a stream to the given file handle buffering up to frames separated
386+-- by the given sequence or up to a maximum of @defaultChunkSize@.
389+{-# INLINE writeByFrames #-}
390+writeByFrames :: (IsStream t, MonadIO m, Storable a)
391+ => Array a -> Handle -> t m a -> m ()
392+writeByFrames = undefined