@@ -0,0 +1,393 @@
1+{-# LANGUAGE CPP #-}
2+{-# LANGUAGE BangPatterns #-}
3+{-# LANGUAGE FlexibleContexts #-}
4+{-# LANGUAGE MagicHash #-}
5+{-# LANGUAGE RecordWildCards #-}
6+{-# LANGUAGE UnboxedTuples #-}
7+
8+#include "inline.hs"
9+
10+-- |
11+-- Module : Streamly.Network.Socket
12+-- Copyright : (c) 2018 Harendra Kumar
13+--
14+-- License : BSD3
15+-- Maintainer : harendra.kumar@gmail.com
16+-- Stability : experimental
17+-- Portability : GHC
18+--
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".
22+--
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
28+-- framing behavior.
29+--
30+-- > import qualified Streamly.Network.Socket as SK
31+--
32+
33+module Streamly.Network.Socket
34+ (
35+ -- ** Listen for Connections
36+ ServerSpec(..)
37+ , recvConnectionsWith
38+ , recvConnectionsOn
39+
40+ -- ** Read a stream from a connection
41+ , read
42+ -- , readUtf8
43+ -- , readLines
44+ -- , readFrames
45+ -- , readByChunks
46+
47+ -- -- * Array Read
48+ -- , readArrayUpto
49+ -- , readArrayOf
50+
51+ -- , readArraysUpto
52+ -- , readArraysOf
53+ , readArrays
54+
55+ -- ** Write a stream to a connection
56+ , write
57+ -- , writeUtf8
58+ -- , writeUtf8ByLines
59+ -- , writeByFrames
60+ , writeByChunks
61+
62+ -- -- * Array Write
63+ , writeArray
64+ , writeArrays
65+ )
66+where
67+
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)
78+
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)
84+
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
88+
89+import Network.Socket hiding (listen)
90+import qualified Network.Socket as Net
91+import Streamly (MonadAsync)
92+
93+-------------------------------------------------------------------------------
94+-- Listen
95+-------------------------------------------------------------------------------
96+
97+-- | Specify the configuration of a server.
98+data ServerSpec = ServerSpec
99+ {
100+ serverAddressFamily :: !Family
101+ , serverAddress :: !SockAddr
102+ , serverSockOpts :: ![(SocketOption, Int)]
103+ }
104+
105+initListener :: Int -> ServerSpec -> IO Socket
106+initListener tcpListenQ ServerSpec{..} =
107+ withSocketsDo $ do
108+ sock <- socket serverAddressFamily Stream defaultProtocol
109+ mapM_ (\(opt, val) -> setSocketOption sock opt val) serverSockOpts
110+ bind sock serverAddress
111+ Net.listen sock tcpListenQ
112+ return sock
113+
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
118+-- backlog.
119+{-# INLINE recvConnectionsWith #-}
120+recvConnectionsWith :: MonadAsync m => Int -> ServerSpec -> SerialT m (Socket, SockAddr)
121+recvConnectionsWith tcpListenQ opts = S.unfoldrM step Nothing
122+ where
123+ step Nothing = do
124+ listener <- liftIO $ initListener tcpListenQ opts
125+ r <- liftIO $ accept listener
126+ -- XXX error handling
127+ return $ Just (r, Just listener)
128+
129+ step (Just listener) = do
130+ r <- liftIO $ accept listener
131+ -- XXX error handling
132+ return $ Just (r, Just listener)
133+
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)]
141+ }
142+
143+-------------------------------------------------------------------------------
144+-- Array IO (Input)
145+-------------------------------------------------------------------------------
146+
147+{-# INLINABLE readArrayUptoWith #-}
148+readArrayUptoWith
149+ :: (h -> Ptr Word8 -> Int -> IO Int)
150+ -> Int
151+ -> h
152+ -> IO (Array Word8)
153+readArrayUptoWith f size h = do
154+ ptr <- mallocPlainForeignPtrBytes size
155+ -- ptr <- mallocPlainForeignPtrAlignedBytes size (alignment (undefined :: Word8))
156+ withForeignPtr ptr $ \p -> do
157+ n <- f h p size
158+ let v = Array
159+ { aStart = ptr
160+ , aEnd = p `plusPtr` n
161+ , aBound = p `plusPtr` size
162+ }
163+ -- XXX shrink only if the diff is significant
164+ -- A.shrinkToFit v
165+ return v
166+
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
174+
175+-------------------------------------------------------------------------------
176+-- Array IO (output)
177+-------------------------------------------------------------------------------
178+
179+waitWhen0 :: Int -> Socket -> IO ()
180+waitWhen0 0 s = when rtsSupportsBoundThreads $
181+ withFdSocket s $ \fd -> threadWaitWrite $ fromIntegral fd
182+waitWhen0 _ _ = return ()
183+
184+sendAll :: Socket -> Ptr Word8 -> Int -> IO ()
185+sendAll _ _ len | len <= 0 = return ()
186+sendAll s p len = do
187+ sent <- sendBuf s p len
188+ waitWhen0 sent s
189+ -- assert (sent <= len)
190+ when (sent >= 0) $ sendAll s (p `plusPtr` sent) (len - sent)
191+
192+{-# INLINABLE writeArrayWith #-}
193+writeArrayWith :: Storable a
194+ => (h -> Ptr Word8 -> Int -> IO ())
195+ -> h
196+ -> Array a
197+ -> IO ()
198+writeArrayWith _ _ arr | A.length arr == 0 = return ()
199+writeArrayWith f h Array{..} = withForeignPtr aStart $ \p ->
200+ f h (castPtr p) aLen
201+ where
202+ aLen =
203+ let p = unsafeForeignPtrToPtr aStart
204+ in aEnd `minusPtr` p
205+
206+-- | Write an Array to a file handle.
207+--
208+-- @since 0.7.0
209+{-# INLINABLE writeArray #-}
210+writeArray :: Storable a => Socket -> Array a -> IO ()
211+writeArray = writeArrayWith sendAll
212+
213+-------------------------------------------------------------------------------
214+-- Stream of Arrays IO
215+-------------------------------------------------------------------------------
216+
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
222+ where
223+ -- XXX use cons/nil instead
224+ go = mkStream $ \_ yld _ stp -> do
225+ arr <- liftIO $ f size h
226+ if A.length arr == 0
227+ then stp
228+ else yld arr go
229+
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'
233+-- on the 'Handle'.
234+{-# INLINABLE readArraysUpto #-}
235+readArraysUpto :: (IsStream t, MonadIO m)
236+ => Int -> Socket -> t m (Array Word8)
237+readArraysUpto = readArraysUptoWith readArrayUpto
238+
239+-- XXX read 'Array a' instead of Word8
240+--
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'
244+-- on the 'Handle'.
245+--
246+-- @since 0.7.0
247+{-# INLINE readArrays #-}
248+readArrays :: (IsStream t, MonadIO m) => Socket -> t m (Array Word8)
249+readArrays = readArraysUpto A.defaultChunkSize
250+
251+-------------------------------------------------------------------------------
252+-- Read File to Stream
253+-------------------------------------------------------------------------------
254+
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.
258+
259+{-
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.
263+--
264+{-# INLINE readByChunksUpto #-}
265+readByChunksUpto :: (IsStream t, MonadIO m) => Int -> Handle -> t m Word8
266+readByChunksUpto chunkSize h = A.flattenArrays $ readArraysUpto chunkSize h
267+-}
268+
269+-- TODO
270+-- read :: (IsStream t, MonadIO m, Storable a) => Handle -> t m a
271+--
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.
275+--
276+-- @since 0.7.0
277+{-# INLINE read #-}
278+read :: (IsStream t, MonadIO m) => Socket -> t m Word8
279+read = A.flattenArrays . readArrays
280+
281+-------------------------------------------------------------------------------
282+-- Writing
283+-------------------------------------------------------------------------------
284+
285+-- | Write a stream of arrays to a handle.
286+--
287+-- @since 0.7.0
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
291+
292+-- GHC buffer size dEFAULT_FD_BUFFER_SIZE=8192 bytes.
293+--
294+-- XXX test this
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
298+-- writes as well.
299+
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
302+-- input elements.
303+--
304+-- @since 0.7.0
305+{-# INLINE writeByChunks #-}
306+writeByChunks :: MonadIO m => Int -> Socket -> SerialT m Word8 -> m ()
307+writeByChunks n h m = writeArrays h $ A.arraysOf n m
308+
309+-- > write = 'writeByChunks' A.defaultChunkSize
310+--
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.
314+--
315+-- @since 0.7.0
316+{-# INLINE write #-}
317+write :: MonadIO m => Socket -> SerialT m Word8 -> m ()
318+write = writeByChunks A.defaultChunkSize
319+
320+{-
321+{-# INLINE write #-}
322+write :: (MonadIO m, Storable a) => Handle -> SerialT m a -> m ()
323+write = toHandleWith A.defaultChunkSize
324+-}
325+
326+-------------------------------------------------------------------------------
327+-- IO with encoding/decoding Unicode characters
328+-------------------------------------------------------------------------------
329+
330+{-
331+-- |
332+-- > readUtf8 = decodeUtf8 . read
333+--
334+-- Read a UTF8 encoded stream of unicode characters from a file handle.
335+--
336+-- @since 0.7.0
337+{-# INLINE readUtf8 #-}
338+readUtf8 :: (IsStream t, MonadIO m) => Handle -> t m Char
339+readUtf8 = decodeUtf8 . read
340+
341+-- |
342+-- > writeUtf8 h s = write h $ encodeUtf8 s
343+--
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.
346+--
347+-- @since 0.7.0
348+{-# INLINE writeUtf8 #-}
349+writeUtf8 :: MonadIO m => Handle -> SerialT m Char -> m ()
350+writeUtf8 h s = write h $ encodeUtf8 s
351+
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'
356+-- option.
357+--
358+-- @since 0.7.0
359+{-# INLINE writeUtf8ByLines #-}
360+writeUtf8ByLines :: (IsStream t, MonadIO m) => Handle -> t m Char -> m ()
361+writeUtf8ByLines = undefined
362+
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.
365+--
366+-- @since 0.7.0
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
370+
371+-------------------------------------------------------------------------------
372+-- Framing on a sequence
373+-------------------------------------------------------------------------------
374+
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
377+-- frame.
378+--
379+-- @since 0.7.0
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
384+
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@.
387+--
388+-- @since 0.7.0
389+{-# INLINE writeByFrames #-}
390+writeByFrames :: (IsStream t, MonadIO m, Storable a)
391+ => Array a -> Handle -> t m a -> m ()
392+writeByFrames = undefined
393+-}