git.iko.soy/streamly
on 2ea9fb0 help sign in
streamly/src/Streamly/Internal/Network/Socket.hs
563 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 Control.Monad.Trans.Control (MonadBaseControl)
73import Data.Maybe (isNothing, fromJust)
74import Data.Word (Word8)
75import Foreign.Ptr (plusPtr, Ptr, castPtr)
76import Streamly.Internal.Data.Unboxed (Unboxed)
77import Network.Socket
78 (Socket, SocketOption(..), Family(..), SockAddr(..),
79 ProtocolNumber, withSocketsDo, SocketType(..), socket, bind,
80 setSocketOption, sendBuf, recvBuf)
81#if MIN_VERSION_network(3,1,0)
82import Network.Socket (withFdSocket)
83#else
84import Network.Socket (fdSocket)
85#endif
86import Prelude hiding (read)
88import qualified Network.Socket as Net
90import Streamly.Internal.Data.Array.Unboxed.Type (Array(..))
91import Streamly.Internal.Data.Array.Unboxed.Stream (lpackArraysChunksOf)
92import Streamly.Internal.Data.Fold (Fold)
93import Streamly.Internal.Data.Stream.Type (Stream)
94import Streamly.Internal.Data.Unfold.Type (Unfold(..))
95-- import Streamly.String (encodeUtf8, decodeUtf8, foldLines)
96import Streamly.Internal.System.IO (defaultChunkSize)
98import qualified Streamly.Data.Array.Unboxed as A (read, length, writeN)
99import qualified Streamly.Data.Fold as FL
100import qualified Streamly.Internal.Data.Array.Unboxed.Type as A
101 (unsafeFreeze, asPtrUnsafe, byteLength, writeNUnsafe)
102import qualified Streamly.Internal.Data.Array.Unboxed.Mut as MArray
103 (Array(..), newPinnedBytes, asPtrUnsafe)
104import qualified Streamly.Internal.Data.Array.Unboxed.Stream as AS
105import qualified Streamly.Internal.Data.Stream as S
106import qualified Streamly.Internal.Data.Stream.StreamD.Type as D
107 (Stream(..), Step(..))
108import qualified Streamly.Data.Unfold as UF
109import qualified Streamly.Internal.Data.Unfold as UF (first, map)
110import qualified Streamly.Internal.Data.Stream.StreamK.Type as K (mkStream)
112-- $setup
113-- >>> :m
114-- >>> import Streamly.Internal.System.IO (defaultChunkSize)
115-- >>> import qualified Streamly.Internal.Network.Socket as Socket
117-- | @'forSocketM' action socket@ runs the monadic computation @action@ passing
118-- the socket handle to it. The handle will be closed on exit from
119-- 'forSocketM', whether by normal termination or by raising an exception. If
120-- closing the handle raises an exception, then this exception will be raised
121-- by 'forSocketM' rather than any exception raised by 'action'.
122--
123-- @since 0.8.0
124{-# INLINE forSocketM #-}
125forSocketM :: (MonadMask m, MonadIO m) => (Socket -> m ()) -> Socket -> m ()
126forSocketM f sk = finally (f sk) (liftIO (Net.close sk))
128-- | Like 'forSocketM' but runs a streaming computation instead of a monadic
129-- computation.
130--
131-- /Inhibits stream fusion/
132--
133-- /Internal/
134{-# INLINE withSocket #-}
135withSocket :: (MonadIO m, MonadBaseControl IO m, MonadCatch m) =>
136 Socket -> (Socket -> Stream m a) -> Stream m a
137withSocket sk f = S.finally (liftIO $ Net.close sk) (f sk)
139-------------------------------------------------------------------------------
140-- Accept (Unfolds)
141-------------------------------------------------------------------------------
143-- XXX Protocol specific socket options should be separated from socket level
144-- options.
145--
146-- | Specify the socket protocol details.
147data SockSpec = SockSpec
148 {
149 sockFamily :: !Family
150 , sockType :: !SocketType
151 , sockProto :: !ProtocolNumber
152 , sockOpts :: ![(SocketOption, Int)]
153 }
155initListener :: Int -> SockSpec -> SockAddr -> IO Socket
156initListener listenQLen SockSpec{..} addr =
157 withSocketsDo $ do
158 sock <- socket sockFamily sockType sockProto
159 use sock `onException` Net.close sock
160 return sock
162 where
164 use sock = do
165 mapM_ (\(opt, val) -> setSocketOption sock opt val) sockOpts
166 bind sock addr
167 Net.listen sock listenQLen
169{-# INLINE listenTuples #-}
170listenTuples :: MonadIO m
171 => Unfold m (Int, SockSpec, SockAddr) (Socket, SockAddr)
172listenTuples = Unfold step inject
173 where
174 inject (listenQLen, spec, addr) =
175 liftIO $ initListener listenQLen spec addr
177 step listener = do
178 r <- liftIO (Net.accept listener `onException` Net.close listener)
179 return $ D.Yield r listener
181-- | Unfold a three tuple @(listenQLen, spec, addr)@ into a stream of connected
182-- protocol sockets corresponding to incoming connections. @listenQLen@ is the
183-- maximum number of pending connections in the backlog. @spec@ is the socket
184-- protocol and options specification and @addr@ is the protocol address where
185-- the server listens for incoming connections.
186--
187-- @since 0.9.0
188{-# INLINE acceptor #-}
189acceptor :: MonadIO m => Unfold m (Int, SockSpec, SockAddr) Socket
190acceptor = UF.map fst listenTuples
192{-# INLINE connectCommon #-}
193connectCommon :: SockSpec -> Maybe SockAddr -> SockAddr -> IO Socket
194connectCommon SockSpec{..} local remote = withSocketsDo $ do
195 sock <- socket sockFamily sockType sockProto
196 use sock `onException` Net.close sock
197 return sock
199 where
201 use sock = do
202 mapM_ (\(opt, val) -> setSocketOption sock opt val) sockOpts
203 forM_ local (bind sock)
204 Net.connect sock remote
206-- | Connect to a remote host using the given socket specification and remote
207-- address. Returns a connected socket or throws an exception.
208--
209-- /Pre-release/
210--
211{-# INLINE connect #-}
212connect :: SockSpec -> SockAddr -> IO Socket
213connect spec = connectCommon spec Nothing
215-- | Connect to a remote host using the given socket specification, a local
216-- address to bind to and a remote address to connect to. Returns a connected
217-- socket or throws an exception.
218--
219-- /Pre-release/
220--
221{-# INLINE connectFrom #-}
222connectFrom :: SockSpec -> SockAddr -> SockAddr -> IO Socket
223connectFrom spec local = connectCommon spec (Just local)
225-------------------------------------------------------------------------------
226-- Listen (Streams)
227-------------------------------------------------------------------------------
229{-# INLINE recvConnectionTuplesWith #-}
230recvConnectionTuplesWith :: MonadIO m
231 => Int -> SockSpec -> SockAddr -> Stream m (Socket, SockAddr)
232recvConnectionTuplesWith tcpListenQ spec addr = S.unfoldrM step Nothing
233 where
234 step Nothing = do
235 listener <- liftIO $ initListener tcpListenQ spec addr
236 r <- liftIO (Net.accept listener `onException` Net.close listener)
237 return $ Just (r, Just listener)
239 step (Just listener) = do
240 r <- liftIO (Net.accept listener `onException` Net.close listener)
241 return $ Just (r, Just listener)
243-- | Start a TCP stream server that listens for connections on the supplied
244-- server address specification (address family, local interface IP address and
245-- port). The server generates a stream of connected sockets. The first
246-- argument is the maximum number of pending connections in the backlog.
247--
248-- /Pre-release/
249{-# INLINE accept #-}
250accept :: MonadIO m => Int -> SockSpec -> SockAddr -> Stream m Socket
251accept tcpListenQ spec addr =
252 fst <$> recvConnectionTuplesWith tcpListenQ spec addr
254-------------------------------------------------------------------------------
255-- Array IO (Input)
256-------------------------------------------------------------------------------
258{-# INLINABLE readArrayUptoWith #-}
259readArrayUptoWith
260 :: (h -> Ptr Word8 -> Int -> IO Int)
261 -> Int
262 -> h
263 -> IO (Array Word8)
264readArrayUptoWith f size h = do
265 arr <- MArray.newPinnedBytes size
266 -- ptr <- mallocPlainForeignPtrAlignedBytes size (alignment (undefined :: Word8))
267 MArray.asPtrUnsafe arr $ \p -> do
268 n <- f h p size
269 let v = A.unsafeFreeze
270 $ arr { MArray.arrEnd = n, MArray.arrBound = size }
272 -- XXX shrink only if the diff is significant
273 -- A.shrinkToFit v
274 return v
276-- | Read a byte array from a file handle up to a maximum of the requested
277-- size. If no data is available on the handle it blocks until some data
278-- becomes available. If data is available then it immediately returns that
279-- data without blocking.
280--
281-- @since 0.9.0
282{-# INLINABLE getChunk #-}
283getChunk :: Int -> Socket -> IO (Array Word8)
284getChunk = readArrayUptoWith recvBuf
286-------------------------------------------------------------------------------
287-- Array IO (output)
288-------------------------------------------------------------------------------
290waitWhen0 :: Int -> Socket -> IO ()
291waitWhen0 0 s = when rtsSupportsBoundThreads $
292#if MIN_VERSION_network(3,1,0)
293 withFdSocket s $ \fd -> threadWaitWrite $ fromIntegral fd
294#elif MIN_VERSION_network(3,0,0)
295 fdSocket s >>= threadWaitWrite . fromIntegral
296#else
297 let fd = fdSocket s in threadWaitWrite $ fromIntegral fd
298#endif
299waitWhen0 _ _ = return ()
301sendAll :: Socket -> Ptr Word8 -> Int -> IO ()
302sendAll _ _ len | len <= 0 = return ()
303sendAll s p len = do
304 sent <- sendBuf s p len
305 waitWhen0 sent s
306 -- assert (sent <= len)
307 when (sent >= 0) $ sendAll s (p `plusPtr` sent) (len - sent)
309{-# INLINABLE writeArrayWith #-}
310writeArrayWith :: Unboxed a
311 => (h -> Ptr Word8 -> Int -> IO ())
312 -> h
313 -> Array a
314 -> IO ()
315writeArrayWith _ _ arr | A.length arr == 0 = return ()
316writeArrayWith f h arr = A.asPtrUnsafe arr $ \ptr -> f h (castPtr ptr) aLen
318 where
320 aLen = A.byteLength arr
322-- | Write an Array to a file handle.
323--
324-- @since 0.9.0
325{-# INLINABLE putChunk #-}
326putChunk :: Unboxed a => Socket -> Array a -> IO ()
327putChunk = writeArrayWith sendAll
329-------------------------------------------------------------------------------
330-- Stream of Arrays IO
331-------------------------------------------------------------------------------
333{-# INLINE _readChunksUptoWith #-}
334_readChunksUptoWith :: (MonadIO m)
335 => (Int -> h -> IO (Array Word8))
336 -> Int -> h -> Stream m (Array Word8)
337_readChunksUptoWith f size h = S.fromStreamK go
338 where
339 -- XXX use cons/nil instead
340 go = K.mkStream $ \_ yld _ stp -> do
341 arr <- liftIO $ f size h
342 if A.length arr == 0
343 then stp
344 else yld arr go
346-- | @readChunksWith bufsize socket@ reads a stream of arrays from @socket@.
347-- The maximum size of a single array is limited to @bufsize@.
348--
349-- /Pre-release/
350{-# INLINE_NORMAL readChunksWith #-}
351readChunksWith :: MonadIO m => Int -> Socket -> Stream m (Array Word8)
352-- readChunksWith = _readChunksUptoWith readChunk
353readChunksWith size h = S.fromStreamD (D.Stream step ())
354 where
355 {-# INLINE_LATE step #-}
356 step _ _ = do
357 arr <- liftIO $ getChunk size h
358 return $
359 case A.length arr of
360 0 -> D.Stop
361 _ -> D.Yield arr ()
363-- | Read a stream of byte arrays from a socket. The maximum size of a single
364-- array is limited to @defaultChunkSize@.
365--
366-- >>> readChunks = Socket.readChunksWith defaultChunkSize
367--
368-- /Pre-release/
369{-# INLINE readChunks #-}
370readChunks :: MonadIO m => Socket -> Stream m (Array Word8)
371readChunks = readChunksWith defaultChunkSize
373-- | Unfold the tuple @(bufsize, socket)@ into a stream of 'Word8' arrays.
374-- Read requests to the socket are performed using a buffer of size @bufsize@.
375-- The size of an array in the resulting stream is always less than or equal to
376-- @bufsize@.
377--
378-- @since 0.9.0
379{-# INLINE_NORMAL chunkReaderWith #-}
380chunkReaderWith :: MonadIO m => Unfold m (Int, Socket) (Array Word8)
381chunkReaderWith = Unfold step return
382 where
383 {-# INLINE_LATE step #-}
384 step (size, h) = do
385 arr <- liftIO $ getChunk size h
386 return $
387 case A.length arr of
388 0 -> D.Stop
389 _ -> D.Yield arr (size, h)
391-- | Same as 'chunkReaderWith'
392--
393-- @since 0.7.0
394{-# DEPRECATED readChunksWithBufferOf "Please use 'chunkReaderWith' instead" #-}
395{-# INLINE_NORMAL readChunksWithBufferOf #-}
396readChunksWithBufferOf :: MonadIO m => Unfold m (Int, Socket) (Array Word8)
397readChunksWithBufferOf = chunkReaderWith
399-- | Unfolds a socket into a stream of 'Word8' arrays. Requests to the socket
400-- are performed using a buffer of size
401-- 'Streamly.Internal.Data.Array.Unboxed.Type.defaultChunkSize'. The
402-- size of arrays in the resulting stream are therefore less than or equal to
403-- 'Streamly.Internal.Data.Array.Unboxed.Type.defaultChunkSize'.
404--
405-- @since 0.9.0
406{-# INLINE chunkReader #-}
407chunkReader :: MonadIO m => Unfold m Socket (Array Word8)
408chunkReader = UF.first defaultChunkSize chunkReaderWith
410-------------------------------------------------------------------------------
411-- Read File to Stream
412-------------------------------------------------------------------------------
414-- | Generate a byte stream from a socket using a buffer of the given size.
415--
416-- /Pre-release/
417{-# INLINE readWith #-}
418readWith :: MonadIO m => Int -> Socket -> Stream m Word8
419readWith size = AS.concat . readChunksWith size
421-- | Generate a byte stream from a socket.
422--
423-- >>> read = Socket.readWith defaultChunkSize
424--
425-- /Pre-release/
426{-# INLINE read #-}
427read :: MonadIO m => Socket -> Stream m Word8
428read = readWith defaultChunkSize
430-- | Unfolds the tuple @(bufsize, socket)@ into a byte stream, read requests
431-- to the socket are performed using buffers of @bufsize@.
432--
433-- @since 0.9.0
434{-# INLINE readerWith #-}
435readerWith :: MonadIO m => Unfold m (Int, Socket) Word8
436readerWith = UF.many A.read chunkReaderWith
438-- | Same as 'readWith'
439--
440-- @since 0.7.0
441{-# DEPRECATED readWithBufferOf "Please use 'readerWith' instead" #-}
442{-# INLINE readWithBufferOf #-}
443readWithBufferOf :: MonadIO m => Unfold m (Int, Socket) Word8
444readWithBufferOf = readerWith
446-- | Unfolds a 'Socket' into a byte stream. IO requests to the socket are
447-- performed in sizes of
448-- 'Streamly.Internal.Data.Array.Unboxed.Type.defaultChunkSize'.
449--
450-- @since 0.9.0
451{-# INLINE reader #-}
452reader :: MonadIO m => Unfold m Socket Word8
453reader = UF.first defaultChunkSize readerWith
455-------------------------------------------------------------------------------
456-- Writing
457-------------------------------------------------------------------------------
459-- | Write a stream of arrays to a handle.
460--
461-- @since 0.7.0
462{-# INLINE putChunks #-}
463putChunks :: (MonadIO m, Unboxed a)
464 => Socket -> Stream m (Array a) -> m ()
465putChunks h = S.fold (FL.drainBy (liftIO . putChunk h))
467-- | Write a stream of arrays to a socket. Each array in the stream is written
468-- to the socket as a separate IO request.
469--
470-- @since 0.7.0
471{-# INLINE writeChunks #-}
472writeChunks :: (MonadIO m, Unboxed a) => Socket -> Fold m (Array a) ()
473writeChunks h = FL.drainBy (liftIO . putChunk h)
475-- | @writeChunksWith bufsize socket@ writes a stream of arrays to
476-- @socket@ after coalescing the adjacent arrays in chunks of @bufsize@.
477-- Multiple arrays are coalesed as long as the total size remains below the
478-- specified size. It never splits an array, if a single array is bigger than
479-- the specified size it emitted as it is.
480--
481-- @since 0.9.0
482{-# INLINE writeChunksWith #-}
483writeChunksWith :: (MonadIO m, Unboxed a)
484 => Int -> Socket -> Fold m (Array a) ()
485writeChunksWith n h = lpackArraysChunksOf n (writeChunks h)
487-- | Same as 'writeChunksWith'
488--
489-- @since 0.7.0
490{-# DEPRECATED writeChunksWithBufferOf "Please use 'writeChunksWith' instead" #-}
491{-# INLINE writeChunksWithBufferOf #-}
492writeChunksWithBufferOf :: (MonadIO m, Unboxed a)
493 => Int -> Socket -> Fold m (Array a) ()
494writeChunksWithBufferOf = writeChunksWith
496-- GHC buffer size dEFAULT_FD_BUFFER_SIZE=8192 bytes.
497--
498-- XXX test this
499-- Note that if you use a chunk size less than 8K (GHC's default buffer
500-- size) then you are advised to use 'NOBuffering' mode on the 'Handle' in case you
501-- do not want buffering to occur at GHC level as well. Same thing applies to
502-- writes as well.
504-- | Like 'write' but provides control over the write buffer. Output will
505-- be written to the IO device as soon as we collect the specified number of
506-- input elements.
507--
508-- @since 0.9.0
509{-# INLINE putBytesWith #-}
510putBytesWith :: MonadIO m => Int -> Socket -> Stream m Word8 -> m ()
511putBytesWith n h m = putChunks h $ AS.arraysOf n m
513-- | Write a byte stream to a socket. Accumulates the input in chunks of
514-- specified number of bytes before writing.
515--
516-- @since 0.9.0
517{-# INLINE writeWith #-}
518writeWith :: MonadIO m => Int -> Socket -> Fold m Word8 ()
519writeWith n h = FL.chunksOf n (A.writeNUnsafe n) (writeChunks h)
521-- | Same as 'writeWith'
522--
523-- @since 0.7.0
524{-# DEPRECATED writeWithBufferOf "Please use 'writeWith' instead" #-}
525{-# INLINE writeWithBufferOf #-}
526writeWithBufferOf :: MonadIO m => Int -> Socket -> Fold m Word8 ()
527writeWithBufferOf = writeWith
529-- | Write a stream of 'Maybe' values. Keep buffering the 'Just' values in an
530-- array. Write the array to the 'Handle' as soon as a 'Nothing' is encountered
531-- or the buffer size exceeds the specified limit.
532--
533-- /Pre-release/
534{-# INLINE writeMaybesWith #-}
535writeMaybesWith :: (MonadIO m )
536 => Int -> Socket -> Fold m (Maybe Word8) ()
537writeMaybesWith n h =
538 let writeNJusts = FL.lmap fromJust $ A.writeN n
539 writeOnNothing = FL.takeEndBy_ isNothing writeNJusts
540 in FL.many writeOnNothing (writeChunks h)
542-- > write = 'writeWith' defaultChunkSize
543--
544-- | Write a byte stream to a file handle. Combines the bytes in chunks of size
545-- up to 'defaultChunkSize' before writing. Note that the write behavior
546-- depends on the 'IOMode' and the current seek position of the handle.
547--
548-- @since 0.7.0
549{-# INLINE putBytes #-}
550putBytes :: MonadIO m => Socket -> Stream m Word8 -> m ()
551putBytes = putBytesWith defaultChunkSize
553-- | Write a byte stream to a socket. Accumulates the input in chunks of
554-- up to 'defaultChunkSize' bytes before writing.
555--
556-- @
557-- write = 'writeWith' 'defaultChunkSize'
558-- @
559--
560-- @since 0.7.0
561{-# INLINE write #-}
562write :: MonadIO m => Socket -> Fold m Word8 ()
563write = writeWith defaultChunkSize