git.iko.soy/streamly
on 5d096df help sign in
streamly/src/Streamly/Internal/Network/Socket.hs
545 lines 19 kB history raw
1#include "inline.hs"
3-- |
4-- Module : Streamly.Internal.Network.Socket
5-- Copyright : (c) 2018 Composewell Technologies
6--
7-- License : BSD3
8-- Maintainer : streamly@composewell.com
9-- Stability : experimental
10-- Portability : GHC
11--
12module Streamly.Internal.Network.Socket
13 (
14 SockSpec (..)
15 -- * Use a socket
16 , forSocketM
17 , withSocket
19 -- * Accept connections
20 , accept
21 , acceptor
23 -- * Connect
24 , connect
25 , connectFrom
27 -- * Read from connection
28 , getChunk
30 -- ** Streams
31 , read
32 , readWith
33 , readChunks
34 , readChunksWith
36 -- ** Unfolds
37 , reader
38 , readerWith
39 , chunkReader
40 , chunkReaderWith
42 -- * Write to connection
43 , putChunk
45 -- ** Folds
46 , write
47 , writeWith
48 , writeChunks
49 , writeChunksWith
50 , writeMaybesWith
52 -- ** Stream writes
53 , putChunks
54 , putBytesWith
55 , putBytes
57 -- reading/writing datagrams
59 -- * Deprecated
60 , readWithBufferOf
61 , readChunksWithBufferOf
62 , writeWithBufferOf
63 , writeChunksWithBufferOf
64 )
65where
67import Control.Concurrent (threadWaitWrite, rtsSupportsBoundThreads)
68import Control.Exception (onException)
69import Control.Monad (forM_, when)
70import Control.Monad.Catch (MonadCatch, finally, MonadMask)
71import Control.Monad.IO.Class (MonadIO(..))
72import Data.Maybe (isNothing, fromJust)
73import Data.Word (Word8)
74import Foreign.Ptr (plusPtr, Ptr, castPtr)
75import Streamly.Data.MutByteArray (Unbox)
76import Network.Socket
77 (Socket, SocketOption(..), Family(..), SockAddr(..),
78 ProtocolNumber, withSocketsDo, SocketType(..), socket, bind,
79 setSocketOption, sendBuf, recvBuf)
80#if MIN_VERSION_network(3,1,0)
81import Network.Socket (withFdSocket)
82#else
83import Network.Socket (fdSocket)
84#endif
85import Prelude hiding (read)
87import qualified Network.Socket as Net
89import Streamly.Internal.Data.Array (Array(..))
90import Streamly.Data.Fold (Fold)
91import Streamly.Data.Stream (Stream)
92import Streamly.Internal.Data.Unfold (Unfold(..))
93-- import Streamly.String (encodeUtf8, decodeUtf8, foldLines)
94import Streamly.Internal.System.IO (defaultChunkSize)
96import qualified Streamly.Data.Array as A
97import qualified Streamly.Data.Fold as FL
98import qualified Streamly.Data.Stream as S
99import qualified Streamly.Data.Unfold as UF
100import qualified Streamly.Internal.Data.Array as A
101 ( unsafeFreeze, unsafePinnedAsPtr, byteLength, pinnedChunksOf,
102 pinnedCreateOf, unsafePinnedCreateOf, lCompactGE )
103import qualified Streamly.Internal.Data.MutArray as MArray
104 (MutArray(..), unsafePinnedAsPtr, pinnedEmptyOf)
105import qualified Streamly.Internal.Data.Stream as S (fromStreamK, Stream(..), Step(..))
106import qualified Streamly.Internal.Data.StreamK as K (mkStream)
108-- $setup
109-- >>> :m
110-- >>> import Streamly.Internal.System.IO (defaultChunkSize)
111-- >>> import qualified Streamly.Internal.Network.Socket as Socket
113-- | @'forSocketM' action socket@ runs the monadic computation @action@ passing
114-- the socket handle to it. The handle will be closed on exit from
115-- 'forSocketM', whether by normal termination or by raising an exception. If
116-- closing the handle raises an exception, then this exception will be raised
117-- by 'forSocketM' rather than any exception raised by 'action'.
118--
119{-# INLINE forSocketM #-}
120forSocketM :: (MonadMask m, MonadIO m) => (Socket -> m ()) -> Socket -> m ()
121forSocketM f sk = finally (f sk) (liftIO (Net.close sk))
123-- | Like 'forSocketM' but runs a streaming computation instead of a monadic
124-- computation.
125--
126-- /Inhibits stream fusion/
127--
128-- /Internal/
129{-# INLINE withSocket #-}
130withSocket :: (MonadIO m, MonadCatch m) =>
131 Socket -> (Socket -> Stream m a) -> Stream m a
132withSocket sk f = S.finallyIO (Net.close sk) (f sk)
134-------------------------------------------------------------------------------
135-- Accept (Unfolds)
136-------------------------------------------------------------------------------
138-- XXX Protocol specific socket options should be separated from socket level
139-- options.
140--
141-- NOTE: the socket config is specified as a record and not by composing
142-- functions because all the fields are mandatory except the sockOpts field.
144-- | Specify the socket protocol details.
145data SockSpec = SockSpec
146 {
147 sockFamily :: !Family
148 , sockType :: !SocketType
149 , sockProto :: !ProtocolNumber
150 , sockOpts :: ![(SocketOption, Int)]
151 }
153initListener :: Int -> SockSpec -> SockAddr -> IO Socket
154initListener listenQLen SockSpec{..} addr =
155 withSocketsDo $ do
156 sock <- socket sockFamily sockType sockProto
157 use sock `onException` Net.close sock
158 return sock
160 where
162 use sock = do
163 mapM_ (\(opt, val) -> setSocketOption sock opt val) sockOpts
164 bind sock addr
165 Net.listen sock listenQLen
167{-# INLINE listenTuples #-}
168listenTuples :: MonadIO m
169 => Unfold m (Int, SockSpec, SockAddr) (Socket, SockAddr)
170listenTuples = Unfold step inject
171 where
172 inject (listenQLen, spec, addr) =
173 liftIO $ initListener listenQLen spec addr
175 step listener = do
176 r <- liftIO (Net.accept listener `onException` Net.close listener)
177 return $ S.Yield r listener
179-- | Unfold a three tuple @(listenQLen, spec, addr)@ into a stream of connected
180-- protocol sockets corresponding to incoming connections. @listenQLen@ is the
181-- maximum number of pending connections in the backlog. @spec@ is the socket
182-- protocol and options specification and @addr@ is the protocol address where
183-- the server listens for incoming connections.
184--
185{-# INLINE acceptor #-}
186acceptor :: MonadIO m => Unfold m (Int, SockSpec, SockAddr) Socket
187acceptor = fmap fst listenTuples
189{-# INLINE connectCommon #-}
190connectCommon :: SockSpec -> Maybe SockAddr -> SockAddr -> IO Socket
191connectCommon SockSpec{..} local remote = withSocketsDo $ do
192 sock <- socket sockFamily sockType sockProto
193 use sock `onException` Net.close sock
194 return sock
196 where
198 use sock = do
199 mapM_ (\(opt, val) -> setSocketOption sock opt val) sockOpts
200 forM_ local (bind sock)
201 Net.connect sock remote
203-- | Connect to a remote host using the given socket specification and remote
204-- address. Returns a connected socket or throws an exception.
205--
206-- /Pre-release/
207--
208{-# INLINE connect #-}
209connect :: SockSpec -> SockAddr -> IO Socket
210connect spec = connectCommon spec Nothing
212-- | Connect to a remote host using the given socket specification, a local
213-- address to bind to and a remote address to connect to. Returns a connected
214-- socket or throws an exception.
215--
216-- /Pre-release/
217--
218{-# INLINE connectFrom #-}
219connectFrom :: SockSpec -> SockAddr -> SockAddr -> IO Socket
220connectFrom spec local = connectCommon spec (Just local)
222-------------------------------------------------------------------------------
223-- Listen (Streams)
224-------------------------------------------------------------------------------
226{-# INLINE recvConnectionTuplesWith #-}
227recvConnectionTuplesWith :: MonadIO m
228 => Int -> SockSpec -> SockAddr -> Stream m (Socket, SockAddr)
229recvConnectionTuplesWith tcpListenQ spec addr = S.unfoldrM step Nothing
230 where
231 step Nothing = do
232 listener <- liftIO $ initListener tcpListenQ spec addr
233 r <- liftIO (Net.accept listener `onException` Net.close listener)
234 return $ Just (r, Just listener)
236 step (Just listener) = do
237 r <- liftIO (Net.accept listener `onException` Net.close listener)
238 return $ Just (r, Just listener)
240-- | Start a TCP stream server that listens for connections on the supplied
241-- server address specification (address family, local interface IP address and
242-- port). The server generates a stream of connected sockets. The first
243-- argument is the maximum number of pending connections in the backlog.
244--
245-- /Pre-release/
246{-# INLINE accept #-}
247accept :: MonadIO m => Int -> SockSpec -> SockAddr -> Stream m Socket
248accept tcpListenQ spec addr =
249 fst <$> recvConnectionTuplesWith tcpListenQ spec addr
251-------------------------------------------------------------------------------
252-- Array IO (Input)
253-------------------------------------------------------------------------------
255-- XXX add an API that compacts the arrays to an exact size.
257{-# INLINABLE readArrayUptoWith #-}
258readArrayUptoWith
259 :: (h -> Ptr Word8 -> Int -> IO Int)
260 -> Int
261 -> h
262 -> IO (Array Word8)
263readArrayUptoWith f size h = do
264 arr <- MArray.pinnedEmptyOf size
265 -- ptr <- mallocPlainForeignPtrAlignedBytes size (alignment (undefined :: Word8))
266 MArray.unsafePinnedAsPtr arr $ \p -> do
267 n <- f h p size
268 let v = A.unsafeFreeze
269 $ arr { MArray.arrEnd = n, MArray.arrBound = size }
271 -- XXX shrink only if the diff is significant
272 -- A.shrinkToFit v
273 return v
275-- | Read a byte array from a file handle up to a maximum of the requested
276-- size. If no data is available on the handle it blocks until some data
277-- becomes available. If data is available then it immediately returns that
278-- data without blocking.
279--
280{-# INLINABLE getChunk #-}
281getChunk :: Int -> Socket -> IO (Array Word8)
282getChunk = readArrayUptoWith recvBuf
284-------------------------------------------------------------------------------
285-- Array IO (output)
286-------------------------------------------------------------------------------
288waitWhen0 :: Int -> Socket -> IO ()
289waitWhen0 0 s = when rtsSupportsBoundThreads $
290#if MIN_VERSION_network(3,1,0)
291 withFdSocket s $ \fd -> threadWaitWrite $ fromIntegral fd
292#elif MIN_VERSION_network(3,0,0)
293 fdSocket s >>= threadWaitWrite . fromIntegral
294#else
295 let fd = fdSocket s in threadWaitWrite $ fromIntegral fd
296#endif
297waitWhen0 _ _ = return ()
299sendAll :: Socket -> Ptr Word8 -> Int -> IO ()
300sendAll _ _ len | len <= 0 = return ()
301sendAll s p len = do
302 sent <- sendBuf s p len
303 waitWhen0 sent s
304 -- assert (sent <= len)
305 when (sent >= 0) $ sendAll s (p `plusPtr` sent) (len - sent)
307{-# INLINABLE writeArrayWith #-}
308writeArrayWith :: Unbox a
309 => (h -> Ptr Word8 -> Int -> IO ())
310 -> h
311 -> Array a
312 -> IO ()
313writeArrayWith _ _ arr | A.length arr == 0 = return ()
314writeArrayWith f h arr = A.unsafePinnedAsPtr arr $ \ptr -> f h (castPtr ptr) aLen
316 where
318 aLen = A.byteLength arr
320-- | Write an Array to a socket.
321--
322{-# INLINABLE putChunk #-}
323putChunk :: Unbox a => Socket -> Array a -> IO ()
324putChunk = writeArrayWith sendAll
326-------------------------------------------------------------------------------
327-- Stream of Arrays IO
328-------------------------------------------------------------------------------
330{-# INLINE _readChunksUptoWith #-}
331_readChunksUptoWith :: (MonadIO m)
332 => (Int -> h -> IO (Array Word8))
333 -> Int -> h -> Stream m (Array Word8)
334_readChunksUptoWith f size h = S.fromStreamK go
335 where
336 -- XXX use cons/nil instead
337 go = K.mkStream $ \_ yld _ stp -> do
338 arr <- liftIO $ f size h
339 if A.length arr == 0
340 then stp
341 else yld arr go
343-- | @readChunksWith bufsize socket@ reads a stream of arrays from @socket@.
344-- The maximum size of a single array is limited to @bufsize@.
345--
346-- /Pre-release/
347{-# INLINE_NORMAL readChunksWith #-}
348readChunksWith :: MonadIO m => Int -> Socket -> Stream m (Array Word8)
349-- readChunksWith = _readChunksUptoWith readChunk
350readChunksWith size h = S.Stream step ()
351 where
352 {-# INLINE_LATE step #-}
353 step _ _ = do
354 arr <- liftIO $ getChunk size h
355 return $
356 case A.length arr of
357 0 -> S.Stop
358 _ -> S.Yield arr ()
360-- | Read a stream of byte arrays from a socket. The maximum size of a single
361-- array is limited to @defaultChunkSize@.
362--
363-- >>> readChunks = Socket.readChunksWith defaultChunkSize
364--
365-- /Pre-release/
366{-# INLINE readChunks #-}
367readChunks :: MonadIO m => Socket -> Stream m (Array Word8)
368readChunks = readChunksWith defaultChunkSize
370-- | Unfold the tuple @(bufsize, socket)@ into a stream of 'Word8' arrays.
371-- Read requests to the socket are performed using a buffer of size @bufsize@.
372-- The size of an array in the resulting stream is always less than or equal to
373-- @bufsize@.
374--
375{-# INLINE_NORMAL chunkReaderWith #-}
376chunkReaderWith :: MonadIO m => Unfold m (Int, Socket) (Array Word8)
377chunkReaderWith = Unfold step return
378 where
379 {-# INLINE_LATE step #-}
380 step (size, h) = do
381 arr <- liftIO $ getChunk size h
382 return $
383 case A.length arr of
384 0 -> S.Stop
385 _ -> S.Yield arr (size, h)
387-- | Same as 'chunkReaderWith'
388--
389{-# DEPRECATED readChunksWithBufferOf "Please use 'chunkReaderWith' instead" #-}
390{-# INLINE_NORMAL readChunksWithBufferOf #-}
391readChunksWithBufferOf :: MonadIO m => Unfold m (Int, Socket) (Array Word8)
392readChunksWithBufferOf = chunkReaderWith
394-- | Unfolds a socket into a stream of 'Word8' arrays. Requests to the socket
395-- are performed using a buffer of size
396-- 'Streamly.Internal.Data.Array.Type.defaultChunkSize'. The
397-- size of arrays in the resulting stream are therefore less than or equal to
398-- 'Streamly.Internal.Data.Array.Type.defaultChunkSize'.
399--
400{-# INLINE chunkReader #-}
401chunkReader :: MonadIO m => Unfold m Socket (Array Word8)
402chunkReader = UF.first defaultChunkSize chunkReaderWith
404-------------------------------------------------------------------------------
405-- Read File to Stream
406-------------------------------------------------------------------------------
408{-# INLINE concatChunks #-}
409concatChunks :: (Monad m, Unbox a) => Stream m (Array a) -> Stream m a
410concatChunks = S.unfoldMany A.reader
412-- | Generate a byte stream from a socket using a buffer of the given size.
413--
414-- /Pre-release/
415{-# INLINE readWith #-}
416readWith :: MonadIO m => Int -> Socket -> Stream m Word8
417readWith size = concatChunks . readChunksWith size
419-- | Generate a byte stream from a socket.
420--
421-- >>> read = Socket.readWith defaultChunkSize
422--
423-- /Pre-release/
424{-# INLINE read #-}
425read :: MonadIO m => Socket -> Stream m Word8
426read = readWith defaultChunkSize
428-- | Unfolds the tuple @(bufsize, socket)@ into a byte stream, read requests
429-- to the socket are performed using buffers of @bufsize@.
430--
431{-# INLINE readerWith #-}
432readerWith :: MonadIO m => Unfold m (Int, Socket) Word8
433readerWith = UF.many A.reader chunkReaderWith
435-- | Same as 'readWith'
436--
437{-# DEPRECATED readWithBufferOf "Please use 'readerWith' instead" #-}
438{-# INLINE readWithBufferOf #-}
439readWithBufferOf :: MonadIO m => Unfold m (Int, Socket) Word8
440readWithBufferOf = readerWith
442-- | Unfolds a 'Socket' into a byte stream. IO requests to the socket are
443-- performed in sizes of
444-- 'Streamly.Internal.Data.Array.Type.defaultChunkSize'.
445--
446{-# INLINE reader #-}
447reader :: MonadIO m => Unfold m Socket Word8
448reader = UF.first defaultChunkSize readerWith
450-------------------------------------------------------------------------------
451-- Writing
452-------------------------------------------------------------------------------
454-- | Write a stream of arrays to a handle.
455--
456{-# INLINE putChunks #-}
457putChunks :: (MonadIO m, Unbox a)
458 => Socket -> Stream m (Array a) -> m ()
459putChunks h = S.fold (FL.drainMapM (liftIO . putChunk h))
461-- | Write a stream of arrays to a socket. Each array in the stream is written
462-- to the socket as a separate IO request.
463--
464{-# INLINE writeChunks #-}
465writeChunks :: (MonadIO m, Unbox a) => Socket -> Fold m (Array a) ()
466writeChunks h = FL.drainMapM (liftIO . putChunk h)
468-- | @writeChunksWith bufsize socket@ writes a stream of arrays to
469-- @socket@ after coalescing the adjacent arrays in chunks of @bufsize@.
470-- Multiple arrays are coalesed as long as the total size remains below the
471-- specified size. It never splits an array, if a single array is bigger than
472-- the specified size it emitted as it is.
473--
474{-# INLINE writeChunksWith #-}
475writeChunksWith :: (MonadIO m, Unbox a)
476 => Int -> Socket -> Fold m (Array a) ()
477writeChunksWith n h = A.lCompactGE n (writeChunks h)
479-- | Same as 'writeChunksWith'
480--
481{-# DEPRECATED writeChunksWithBufferOf "Please use 'writeChunksWith' instead" #-}
482{-# INLINE writeChunksWithBufferOf #-}
483writeChunksWithBufferOf :: (MonadIO m, Unbox a)
484 => Int -> Socket -> Fold m (Array a) ()
485writeChunksWithBufferOf = writeChunksWith
487-- GHC buffer size dEFAULT_FD_BUFFER_SIZE=8192 bytes.
488--
489-- XXX test this
490-- Note that if you use a chunk size less than 8K (GHC's default buffer
491-- size) then you are advised to use 'NOBuffering' mode on the 'Handle' in case you
492-- do not want buffering to occur at GHC level as well. Same thing applies to
493-- writes as well.
495-- | Like 'write' but provides control over the write buffer. Output will
496-- be written to the IO device as soon as we collect the specified number of
497-- input elements.
498--
499{-# INLINE putBytesWith #-}
500putBytesWith :: MonadIO m => Int -> Socket -> Stream m Word8 -> m ()
501putBytesWith n h m = putChunks h $ A.pinnedChunksOf n m
503-- | Write a byte stream to a socket. Accumulates the input in chunks of
504-- specified number of bytes before writing.
505--
506{-# INLINE writeWith #-}
507writeWith :: MonadIO m => Int -> Socket -> Fold m Word8 ()
508writeWith n h = FL.groupsOf n (A.unsafePinnedCreateOf n) (writeChunks h)
510-- | Same as 'writeWith'
511--
512{-# DEPRECATED writeWithBufferOf "Please use 'writeWith' instead" #-}
513{-# INLINE writeWithBufferOf #-}
514writeWithBufferOf :: MonadIO m => Int -> Socket -> Fold m Word8 ()
515writeWithBufferOf = writeWith
517-- | Write a stream of 'Maybe' values. Keep buffering the 'Just' values in an
518-- array. Write the array to the 'Handle' as soon as a 'Nothing' is encountered
519-- or the buffer size exceeds the specified limit.
520--
521-- /Pre-release/
522{-# INLINE writeMaybesWith #-}
523writeMaybesWith :: (MonadIO m )
524 => Int -> Socket -> Fold m (Maybe Word8) ()
525writeMaybesWith n h =
526 let writeNJusts = FL.lmap fromJust $ A.pinnedCreateOf n
527 writeOnNothing = FL.takeEndBy_ isNothing writeNJusts
528 in FL.many writeOnNothing (writeChunks h)
530-- | Write a byte stream to a file handle. Combines the bytes in chunks of size
531-- up to 'defaultChunkSize' before writing. Note that the write behavior
532-- depends on the 'IOMode' and the current seek position of the handle.
533--
534{-# INLINE putBytes #-}
535putBytes :: MonadIO m => Socket -> Stream m Word8 -> m ()
536putBytes = putBytesWith defaultChunkSize
538-- | Write a byte stream to a socket. Accumulates the input in chunks of
539-- up to 'defaultChunkSize' bytes before writing.
540--
541-- >>> write = Socket.writeWith defaultChunkSize
542--
543{-# INLINE write #-}
544write :: MonadIO m => Socket -> Fold m Word8 ()
545write = writeWith defaultChunkSize