4-- Module : Streamly.Internal.Network.Socket
5-- Copyright : (c) 2018 Composewell Technologies
8-- Maintainer : streamly@composewell.com
9-- Stability : experimental
12module Streamly.Internal.Network.Socket
19 -- * Accept connections
27 -- * Read from connection
42 -- * Write to connection
57 -- reading/writing datagrams
61 , readChunksWithBufferOf
63 , writeChunksWithBufferOf
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.Internal.Data.Unboxed (Unbox)
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)
83import Network.Socket (fdSocket)
85import Prelude hiding (read)
87import qualified Network.Socket as Net
89import Streamly.Internal.Data.Array.Type (Array(..))
90import Streamly.Internal.Data.Stream.Chunked (lpackArraysChunksOf)
91import Streamly.Internal.Data.Fold (Fold)
92import Streamly.Internal.Data.Stream.StreamD.Type (Stream)
93import Streamly.Internal.Data.Unfold.Type (Unfold(..))
94-- import Streamly.String (encodeUtf8, decodeUtf8, foldLines)
95import Streamly.Internal.System.IO (defaultChunkSize)
97import qualified Streamly.Data.Array as A (reader, length, writeN)
98import qualified Streamly.Data.Fold as FL
99import qualified Streamly.Internal.Data.Array.Type as A
100 (unsafeFreeze, asPtrUnsafe, byteLength, writeNUnsafe, chunksOf)
101import qualified Streamly.Internal.Data.Array.Mut as MArray
102 (MutArray(..), newPinnedBytes, asPtrUnsafe)
103import qualified Streamly.Internal.Data.Stream.StreamD as S
104import qualified Streamly.Internal.Data.Stream.StreamD.Type as D
105 (Stream(..), Step(..))
106import qualified Streamly.Data.Unfold as UF
107import qualified Streamly.Internal.Data.Unfold as UF (first, map)
108import qualified Streamly.Internal.Data.Stream.StreamK.Type as K (mkStream)
112-- >>> import Streamly.Internal.System.IO (defaultChunkSize)
113-- >>> import qualified Streamly.Internal.Network.Socket as Socket
115-- | @'forSocketM' action socket@ runs the monadic computation @action@ passing
116-- the socket handle to it. The handle will be closed on exit from
117-- 'forSocketM', whether by normal termination or by raising an exception. If
118-- closing the handle raises an exception, then this exception will be raised
119-- by 'forSocketM' rather than any exception raised by 'action'.
121{-# INLINE forSocketM #-}
122forSocketM :: (MonadMask m, MonadIO m) => (Socket -> m ()) -> Socket -> m ()
123forSocketM f sk = finally (f sk) (liftIO (Net.close sk))
125-- | Like 'forSocketM' but runs a streaming computation instead of a monadic
128-- /Inhibits stream fusion/
131{-# INLINE withSocket #-}
132withSocket :: (MonadIO m, MonadCatch m) =>
133 Socket -> (Socket -> Stream m a) -> Stream m a
134withSocket sk f = S.finallyIO (Net.close sk) (f sk)
136-------------------------------------------------------------------------------
138-------------------------------------------------------------------------------
140-- XXX Protocol specific socket options should be separated from socket level
143-- | Specify the socket protocol details.
144data SockSpec = SockSpec
146 sockFamily :: !Family
147 , sockType :: !SocketType
148 , sockProto :: !ProtocolNumber
149 , sockOpts :: ![(SocketOption, Int)]
152initListener :: Int -> SockSpec -> SockAddr -> IO Socket
153initListener listenQLen SockSpec{..} addr =
155 sock <- socket sockFamily sockType sockProto
156 use sock `onException` Net.close sock
162 mapM_ (\(opt, val) -> setSocketOption sock opt val) sockOpts
164 Net.listen sock listenQLen
166{-# INLINE listenTuples #-}
167listenTuples :: MonadIO m
168 => Unfold m (Int, SockSpec, SockAddr) (Socket, SockAddr)
169listenTuples = Unfold step inject
171 inject (listenQLen, spec, addr) =
172 liftIO $ initListener listenQLen spec addr
175 r <- liftIO (Net.accept listener `onException` Net.close listener)
176 return $ D.Yield r listener
178-- | Unfold a three tuple @(listenQLen, spec, addr)@ into a stream of connected
179-- protocol sockets corresponding to incoming connections. @listenQLen@ is the
180-- maximum number of pending connections in the backlog. @spec@ is the socket
181-- protocol and options specification and @addr@ is the protocol address where
182-- the server listens for incoming connections.
184{-# INLINE acceptor #-}
185acceptor :: MonadIO m => Unfold m (Int, SockSpec, SockAddr) Socket
186acceptor = UF.map fst listenTuples
188{-# INLINE connectCommon #-}
189connectCommon :: SockSpec -> Maybe SockAddr -> SockAddr -> IO Socket
190connectCommon SockSpec{..} local remote = withSocketsDo $ do
191 sock <- socket sockFamily sockType sockProto
192 use sock `onException` Net.close sock
198 mapM_ (\(opt, val) -> setSocketOption sock opt val) sockOpts
199 forM_ local (bind sock)
200 Net.connect sock remote
202-- | Connect to a remote host using the given socket specification and remote
203-- address. Returns a connected socket or throws an exception.
207{-# INLINE connect #-}
208connect :: SockSpec -> SockAddr -> IO Socket
209connect spec = connectCommon spec Nothing
211-- | Connect to a remote host using the given socket specification, a local
212-- address to bind to and a remote address to connect to. Returns a connected
213-- socket or throws an exception.
217{-# INLINE connectFrom #-}
218connectFrom :: SockSpec -> SockAddr -> SockAddr -> IO Socket
219connectFrom spec local = connectCommon spec (Just local)
221-------------------------------------------------------------------------------
223-------------------------------------------------------------------------------
225{-# INLINE recvConnectionTuplesWith #-}
226recvConnectionTuplesWith :: MonadIO m
227 => Int -> SockSpec -> SockAddr -> Stream m (Socket, SockAddr)
228recvConnectionTuplesWith tcpListenQ spec addr = S.unfoldrM step Nothing
231 listener <- liftIO $ initListener tcpListenQ spec addr
232 r <- liftIO (Net.accept listener `onException` Net.close listener)
233 return $ Just (r, Just listener)
235 step (Just listener) = do
236 r <- liftIO (Net.accept listener `onException` Net.close listener)
237 return $ Just (r, Just listener)
239-- | Start a TCP stream server that listens for connections on the supplied
240-- server address specification (address family, local interface IP address and
241-- port). The server generates a stream of connected sockets. The first
242-- argument is the maximum number of pending connections in the backlog.
246accept :: MonadIO m => Int -> SockSpec -> SockAddr -> Stream m Socket
247accept tcpListenQ spec addr =
248 fst <$> recvConnectionTuplesWith tcpListenQ spec addr
250-------------------------------------------------------------------------------
252-------------------------------------------------------------------------------
254{-# INLINABLE readArrayUptoWith #-}
256 :: (h -> Ptr Word8 -> Int -> IO Int)
260readArrayUptoWith f size h = do
261 arr <- MArray.newPinnedBytes size
262 -- ptr <- mallocPlainForeignPtrAlignedBytes size (alignment (undefined :: Word8))
263 MArray.asPtrUnsafe arr $ \p -> do
265 let v = A.unsafeFreeze
266 $ arr { MArray.arrEnd = n, MArray.arrBound = size }
268 -- XXX shrink only if the diff is significant
272-- | Read a byte array from a file handle up to a maximum of the requested
273-- size. If no data is available on the handle it blocks until some data
274-- becomes available. If data is available then it immediately returns that
275-- data without blocking.
277{-# INLINABLE getChunk #-}
278getChunk :: Int -> Socket -> IO (Array Word8)
279getChunk = readArrayUptoWith recvBuf
281-------------------------------------------------------------------------------
283-------------------------------------------------------------------------------
285waitWhen0 :: Int -> Socket -> IO ()
286waitWhen0 0 s = when rtsSupportsBoundThreads $
287#if MIN_VERSION_network(3,1,0)
288 withFdSocket s $ \fd -> threadWaitWrite $ fromIntegral fd
289#elif MIN_VERSION_network(3,0,0)
290 fdSocket s >>= threadWaitWrite . fromIntegral
292 let fd = fdSocket s in threadWaitWrite $ fromIntegral fd
294waitWhen0 _ _ = return ()
296sendAll :: Socket -> Ptr Word8 -> Int -> IO ()
297sendAll _ _ len | len <= 0 = return ()
299 sent <- sendBuf s p len
301 -- assert (sent <= len)
302 when (sent >= 0) $ sendAll s (p `plusPtr` sent) (len - sent)
304{-# INLINABLE writeArrayWith #-}
305writeArrayWith :: Unbox a
306 => (h -> Ptr Word8 -> Int -> IO ())
310writeArrayWith _ _ arr | A.length arr == 0 = return ()
311writeArrayWith f h arr = A.asPtrUnsafe arr $ \ptr -> f h (castPtr ptr) aLen
315 aLen = A.byteLength arr
317-- | Write an Array to a socket.
319{-# INLINABLE putChunk #-}
320putChunk :: Unbox a => Socket -> Array a -> IO ()
321putChunk = writeArrayWith sendAll
323-------------------------------------------------------------------------------
324-- Stream of Arrays IO
325-------------------------------------------------------------------------------
327{-# INLINE _readChunksUptoWith #-}
328_readChunksUptoWith :: (MonadIO m)
329 => (Int -> h -> IO (Array Word8))
330 -> Int -> h -> Stream m (Array Word8)
331_readChunksUptoWith f size h = S.fromStreamK go
333 -- XXX use cons/nil instead
334 go = K.mkStream $ \_ yld _ stp -> do
335 arr <- liftIO $ f size h
340-- | @readChunksWith bufsize socket@ reads a stream of arrays from @socket@.
341-- The maximum size of a single array is limited to @bufsize@.
344{-# INLINE_NORMAL readChunksWith #-}
345readChunksWith :: MonadIO m => Int -> Socket -> Stream m (Array Word8)
346-- readChunksWith = _readChunksUptoWith readChunk
347readChunksWith size h = D.Stream step ()
349 {-# INLINE_LATE step #-}
351 arr <- liftIO $ getChunk size h
357-- | Read a stream of byte arrays from a socket. The maximum size of a single
358-- array is limited to @defaultChunkSize@.
360-- >>> readChunks = Socket.readChunksWith defaultChunkSize
363{-# INLINE readChunks #-}
364readChunks :: MonadIO m => Socket -> Stream m (Array Word8)
365readChunks = readChunksWith defaultChunkSize
367-- | Unfold the tuple @(bufsize, socket)@ into a stream of 'Word8' arrays.
368-- Read requests to the socket are performed using a buffer of size @bufsize@.
369-- The size of an array in the resulting stream is always less than or equal to
372{-# INLINE_NORMAL chunkReaderWith #-}
373chunkReaderWith :: MonadIO m => Unfold m (Int, Socket) (Array Word8)
374chunkReaderWith = Unfold step return
376 {-# INLINE_LATE step #-}
378 arr <- liftIO $ getChunk size h
382 _ -> D.Yield arr (size, h)
384-- | Same as 'chunkReaderWith'
386{-# DEPRECATED readChunksWithBufferOf "Please use 'chunkReaderWith' instead" #-}
387{-# INLINE_NORMAL readChunksWithBufferOf #-}
388readChunksWithBufferOf :: MonadIO m => Unfold m (Int, Socket) (Array Word8)
389readChunksWithBufferOf = chunkReaderWith
391-- | Unfolds a socket into a stream of 'Word8' arrays. Requests to the socket
392-- are performed using a buffer of size
393-- 'Streamly.Internal.Data.Array.Type.defaultChunkSize'. The
394-- size of arrays in the resulting stream are therefore less than or equal to
395-- 'Streamly.Internal.Data.Array.Type.defaultChunkSize'.
397{-# INLINE chunkReader #-}
398chunkReader :: MonadIO m => Unfold m Socket (Array Word8)
399chunkReader = UF.first defaultChunkSize chunkReaderWith
401-------------------------------------------------------------------------------
402-- Read File to Stream
403-------------------------------------------------------------------------------
405{-# INLINE concatChunks #-}
406concatChunks :: (Monad m, Unbox a) => Stream m (Array a) -> Stream m a
407concatChunks = S.unfoldMany A.reader
409-- | Generate a byte stream from a socket using a buffer of the given size.
412{-# INLINE readWith #-}
413readWith :: MonadIO m => Int -> Socket -> Stream m Word8
414readWith size = concatChunks . readChunksWith size
416-- | Generate a byte stream from a socket.
418-- >>> read = Socket.readWith defaultChunkSize
422read :: MonadIO m => Socket -> Stream m Word8
423read = readWith defaultChunkSize
425-- | Unfolds the tuple @(bufsize, socket)@ into a byte stream, read requests
426-- to the socket are performed using buffers of @bufsize@.
428{-# INLINE readerWith #-}
429readerWith :: MonadIO m => Unfold m (Int, Socket) Word8
430readerWith = UF.many A.reader chunkReaderWith
432-- | Same as 'readWith'
434{-# DEPRECATED readWithBufferOf "Please use 'readerWith' instead" #-}
435{-# INLINE readWithBufferOf #-}
436readWithBufferOf :: MonadIO m => Unfold m (Int, Socket) Word8
437readWithBufferOf = readerWith
439-- | Unfolds a 'Socket' into a byte stream. IO requests to the socket are
440-- performed in sizes of
441-- 'Streamly.Internal.Data.Array.Type.defaultChunkSize'.
444reader :: MonadIO m => Unfold m Socket Word8
445reader = UF.first defaultChunkSize readerWith
447-------------------------------------------------------------------------------
449-------------------------------------------------------------------------------
451-- | Write a stream of arrays to a handle.
453{-# INLINE putChunks #-}
454putChunks :: (MonadIO m, Unbox a)
455 => Socket -> Stream m (Array a) -> m ()
456putChunks h = S.fold (FL.drainMapM (liftIO . putChunk h))
458-- | Write a stream of arrays to a socket. Each array in the stream is written
459-- to the socket as a separate IO request.
461{-# INLINE writeChunks #-}
462writeChunks :: (MonadIO m, Unbox a) => Socket -> Fold m (Array a) ()
463writeChunks h = FL.drainMapM (liftIO . putChunk h)
465-- | @writeChunksWith bufsize socket@ writes a stream of arrays to
466-- @socket@ after coalescing the adjacent arrays in chunks of @bufsize@.
467-- Multiple arrays are coalesed as long as the total size remains below the
468-- specified size. It never splits an array, if a single array is bigger than
469-- the specified size it emitted as it is.
471{-# INLINE writeChunksWith #-}
472writeChunksWith :: (MonadIO m, Unbox a)
473 => Int -> Socket -> Fold m (Array a) ()
474writeChunksWith n h = lpackArraysChunksOf n (writeChunks h)
476-- | Same as 'writeChunksWith'
478{-# DEPRECATED writeChunksWithBufferOf "Please use 'writeChunksWith' instead" #-}
479{-# INLINE writeChunksWithBufferOf #-}
480writeChunksWithBufferOf :: (MonadIO m, Unbox a)
481 => Int -> Socket -> Fold m (Array a) ()
482writeChunksWithBufferOf = writeChunksWith
484-- GHC buffer size dEFAULT_FD_BUFFER_SIZE=8192 bytes.
487-- Note that if you use a chunk size less than 8K (GHC's default buffer
488-- size) then you are advised to use 'NOBuffering' mode on the 'Handle' in case you
489-- do not want buffering to occur at GHC level as well. Same thing applies to
492-- | Like 'write' but provides control over the write buffer. Output will
493-- be written to the IO device as soon as we collect the specified number of
496{-# INLINE putBytesWith #-}
497putBytesWith :: MonadIO m => Int -> Socket -> Stream m Word8 -> m ()
498putBytesWith n h m = putChunks h $ A.chunksOf n m
500-- | Write a byte stream to a socket. Accumulates the input in chunks of
501-- specified number of bytes before writing.
503{-# INLINE writeWith #-}
504writeWith :: MonadIO m => Int -> Socket -> Fold m Word8 ()
505writeWith n h = FL.groupsOf n (A.writeNUnsafe n) (writeChunks h)
507-- | Same as 'writeWith'
509{-# DEPRECATED writeWithBufferOf "Please use 'writeWith' instead" #-}
510{-# INLINE writeWithBufferOf #-}
511writeWithBufferOf :: MonadIO m => Int -> Socket -> Fold m Word8 ()
512writeWithBufferOf = writeWith
514-- | Write a stream of 'Maybe' values. Keep buffering the 'Just' values in an
515-- array. Write the array to the 'Handle' as soon as a 'Nothing' is encountered
516-- or the buffer size exceeds the specified limit.
519{-# INLINE writeMaybesWith #-}
520writeMaybesWith :: (MonadIO m )
521 => Int -> Socket -> Fold m (Maybe Word8) ()
523 let writeNJusts = FL.lmap fromJust $ A.writeN n
524 writeOnNothing = FL.takeEndBy_ isNothing writeNJusts
525 in FL.many writeOnNothing (writeChunks h)
527-- | Write a byte stream to a file handle. Combines the bytes in chunks of size
528-- up to 'defaultChunkSize' before writing. Note that the write behavior
529-- depends on the 'IOMode' and the current seek position of the handle.
531{-# INLINE putBytes #-}
532putBytes :: MonadIO m => Socket -> Stream m Word8 -> m ()
533putBytes = putBytesWith defaultChunkSize
535-- | Write a byte stream to a socket. Accumulates the input in chunks of
536-- up to 'defaultChunkSize' bytes before writing.
538-- >>> write = Socket.writeWith defaultChunkSize
541write :: MonadIO m => Socket -> Fold m Word8 ()
542write = writeWith defaultChunkSize
keysesc to close
- j k
- move down, up
- enter
- open
- esc
- back, or up a level
- 1 2 3 4
- readme, files, commits, refs
- /
- filter
- y
- copy the clone command
- c
- compare the chosen commit with another (or shift-click both)
- ?
- show or hide this list