packages feed

recv (empty) → 0.0.0

raw patch · 9 files changed

+449/−0 lines, 9 filesdep +basedep +bytestringdep +hspec

Dependencies added: base, bytestring, hspec, network

Files

+ LICENSE view
@@ -0,0 +1,29 @@+Copyright (c) 2022, Internet Initiative Japan Inc.+All rights reserved.++Redistribution and use in source and binary forms, with or without+modification, are permitted provided that the following conditions+are met:++  * Redistributions of source code must retain the above copyright+    notice, this list of conditions and the following disclaimer.+  * Redistributions in binary form must reproduce the above copyright+    notice, this list of conditions and the following disclaimer in+    the documentation and/or other materials provided with the+    distribution.+  * Neither the name of the copyright holders nor the names of its+    contributors may be used to endorse or promote products derived+    from this software without specific prior written permission.++THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS+"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT+LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS+FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE+COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,+INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,+BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;+LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT+LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN+ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE+POSSIBILITY OF SUCH DAMAGE.
+ Network/Socket/BufferPool.hs view
@@ -0,0 +1,47 @@+-- | This module provides efficient receiving functions from the network.+--   'Network.Socket.ByteString.recv' uses 'createAndTrim'+--   which behaves as follows:+--+--   * Allocates a buffer whose size is decided from the+--     first argument.+--   * Receives data with the buffer.+--   * Allocates another buffer whose size fits the received data.+--   * Copies the data from the first buffer to the second buffer.+--+--  On 64bit machines, the global lock is taken for the allocation of+--  a byte string whose length is larger than or equal to 3272 bytes.+--  So, for instance, if 4,096 is specified to 'recv' and the size of+--  received data is 3,300, the global lock is taken twice with the copy+--  overhead.+--+--  The efficient receiving functions provided here use a buffer pool.+--  A large buffer is allocated at the beginning and it is divided into+--  a used one and a leftover when receiving.+--  The latter is kept in the buffer pooll and will be used next time.+--  When the buffer gets small+--  and usefless, a new large buffer is allocated.+module Network.Socket.BufferPool (+  -- * Recv+    Recv+  , receive+  , BufferPool+  , newBufferPool+  , withBufferPool+  -- * RecvBuf+  , RecvBuf+  , receiveBuf+  -- * RecvN+  , RecvN+  , makeReceiveN+  , makePlainReceiveN+  -- * Types+  , Buffer+  , BufSize+  -- * Utilities+  , mallocBS+  , copy+  ) where++import Network.Socket.BufferPool.Buffer+import Network.Socket.BufferPool.Recv+import Network.Socket.BufferPool.Types
+ Network/Socket/BufferPool/Buffer.hs view
@@ -0,0 +1,64 @@+module Network.Socket.BufferPool.Buffer (+    newBufferPool+  , withBufferPool+  , mallocBS+  , copy+  ) where++import qualified Data.ByteString as BS+import Data.ByteString.Internal (ByteString(..), memcpy)+import Data.ByteString.Unsafe (unsafeTake, unsafeDrop)+import Data.IORef (newIORef, readIORef, writeIORef)+import Foreign.ForeignPtr+import Foreign.Marshal.Alloc (mallocBytes, finalizerFree)+import Foreign.Ptr (castPtr, plusPtr)++import Network.Socket.BufferPool.Types++----------------------------------------------------------------++-- | Creating a buffer pool.+--   The first argument is the lower limit.+--   When the size of the buffer in the poll is lower than this limit,+--   the buffer is thrown awany (and is eventually freed).+--   Then a new buffer is allocated.+--   The second argument is the size for the new allocation.+newBufferPool :: Int -> Int -> IO BufferPool+newBufferPool l h = BufferPool l h <$> newIORef BS.empty++----------------------------------------------------------------++-- | Using a buffer pool.+--   The second argument is a function which returns+--   how many bytes are filled in the buffer.+--   The buffer in the buffer pool is automatically managed.+withBufferPool :: BufferPool -> (Buffer -> BufSize -> IO Int) -> IO ByteString+withBufferPool (BufferPool l h ref) f = do+    buf0 <- readIORef ref+    buf  <- if BS.length buf0 >= l then return buf0+                                   else mallocBS h+    consumed <- withForeignBuffer buf f+    writeIORef ref $ unsafeDrop consumed buf+    return $ unsafeTake consumed buf++withForeignBuffer :: ByteString -> (Buffer -> BufSize -> IO Int) -> IO Int+withForeignBuffer (PS ps s l) f = withForeignPtr ps $ \p -> f (castPtr p `plusPtr` s) l+{-# INLINE withForeignBuffer #-}++----------------------------------------------------------------++-- | Allocating a byte string.+mallocBS :: Int -> IO ByteString+mallocBS size = do+    ptr <- mallocBytes size+    fptr <- newForeignPtr finalizerFree ptr+    return $ PS fptr 0 size+{-# INLINE mallocBS #-}++-- | Copying the bytestring to the buffer.+--   This function returns the point where the next copy should start.+copy :: Buffer -> ByteString -> IO Buffer+copy ptr (PS fp o l) = withForeignPtr fp $ \p -> do+    memcpy ptr (p `plusPtr` o) (fromIntegral l)+    return $ ptr `plusPtr` l+{-# INLINE copy #-}
+ Network/Socket/BufferPool/Recv.hs view
@@ -0,0 +1,165 @@+{-# LANGUAGE ForeignFunctionInterface, OverloadedStrings #-}+{-# LANGUAGE CPP #-}++module Network.Socket.BufferPool.Recv (+    receive+  , receiveBuf+  , makeReceiveN+  , makePlainReceiveN+  ) where++import qualified Data.ByteString as BS+import Data.ByteString.Internal (ByteString(..))+import Data.IORef+import Foreign.C.Error (eAGAIN, getErrno, throwErrno)+import Foreign.C.Types+import Foreign.ForeignPtr (withForeignPtr)+import Foreign.Ptr (Ptr, castPtr, plusPtr)+import GHC.Conc (threadWaitRead)+import Network.Socket (Socket, withFdSocket)+import System.Posix.Types (Fd(..))++#ifdef mingw32_HOST_OS+import GHC.IO.FD (FD(..), readRawBufferPtr)+import Network.Socket.BufferPool.Windows+#endif++import Network.Socket.BufferPool.Types+import Network.Socket.BufferPool.Buffer++----------------------------------------------------------------++-- | The receiving function with a buffer pool.+--   The buffer pool is automatically managed.+receive :: Socket -> BufferPool -> Recv+receive sock pool = withBufferPool pool $ \ptr size -> do+#if MIN_VERSION_network(3,1,0)+  withFdSocket sock $ \fd -> do+#elif MIN_VERSION_network(3,0,0)+    fd <- fdSocket sock+#else+    let fd = fdSocket sock+#endif+    let size' = fromIntegral size+    fromIntegral <$> tryRecv fd ptr size'++----------------------------------------------------------------++-- | The receiving function with a buffer.+--   This tries to fill the buffer.+--   This returns when the buffer is filled or reaches EOF.+receiveBuf :: Socket -> RecvBuf+receiveBuf sock buf0 siz0 = do+#if MIN_VERSION_network(3,1,0)+  withFdSocket sock $ \fd -> do+#elif MIN_VERSION_network(3,0,0)+    fd <- fdSocket sock+#else+    let fd = fdSocket sock+#endif+    loop fd buf0 siz0+  where+    loop _  _   0   = return True+    loop fd buf siz = do+        n <- fromIntegral <$> tryRecv fd buf (fromIntegral siz)+        -- fixme: what should we do in the case of n == 0+        if n == 0 then+            return False+          else+            loop fd (buf `plusPtr` n) (siz - n)++----------------------------------------------------------------++tryRecv :: CInt -> Buffer -> CSize -> IO CInt+tryRecv sock ptr size = go+  where+    go = do+#ifdef mingw32_HOST_OS+      bytes <- windowsThreadBlockHack $ fromIntegral <$> readRawBufferPtr "tryRecv" (FD sock 1) (castPtr ptr) 0 size+#else+      bytes <- c_recv sock (castPtr ptr) size 0+#endif+      if bytes == -1 then do+          errno <- getErrno+          if errno == eAGAIN then do+              threadWaitRead (Fd sock)+              go+            else+              throwErrno "tryRecv"+         else+          return bytes++----------------------------------------------------------------++-- | This function returns a receiving function+--   based on two receiving functions.+--   The returned function receives exactly N bytes.+--   The first argument is an initial received data.+--   After consuming the initial data, the two functions is used.+--   When N is less than equal to 4096, the buffer pool is used.+--   Otherwise, a new buffer is allocated.+--   In this case, the global lock is taken.+makeReceiveN :: ByteString -> Recv -> RecvBuf -> IO RecvN+makeReceiveN bs0 recv recvBuf = do+    ref <- newIORef bs0+    return $ receiveN ref recv recvBuf++-- | This function returns a receiving function with two receiving+--   functions is created internally.+--   The second argument is the lower limit of the buffer pool.+--   The third argument is the size of the allocated buffer in the pool.+--   The fourth argument is an initial received data.+--   The returned function behaves as described in 'makeReceiveN'.+makePlainReceiveN :: Socket -> Int -> Int -> ByteString -> IO RecvN+makePlainReceiveN s l h bs0 = do+    ref <- newIORef bs0+    pool <- newBufferPool l h+    return $ receiveN ref (receive s pool) (receiveBuf s)++-- | The receiving function which receives exactly N bytes+--   (the fourth argument).+receiveN :: IORef ByteString -> Recv -> RecvBuf -> RecvN+receiveN ref recv recvBuf size = do+    cached <- readIORef ref+    (bs, leftover) <- tryRecvN cached size recv recvBuf+    writeIORef ref leftover+    return bs++----------------------------------------------------------------++tryRecvN :: ByteString -> Int -> IO ByteString -> RecvBuf -> IO (ByteString, ByteString)+tryRecvN init0 siz0 recv recvBuf+  | siz0 <= len0 = return $ BS.splitAt siz0 init0+  -- fixme: hard coding 4096+  | siz0 <= 4096 = recvWithPool [init0] (siz0 - len0)+  | otherwise    = recvWithNewBuf+  where+    len0 = BS.length init0+    recvWithPool bss siz = do+        bs <- recv+        let len = BS.length bs+        if len == 0 then+            return ("", "")+          else if len >= siz then do+            let (consume, leftover) = BS.splitAt siz bs+                ret = BS.concat $ reverse (consume : bss)+            return (ret, leftover)+          else do+            let bss' = bs : bss+                siz' = siz - len+            recvWithPool bss' siz'+    recvWithNewBuf = do+      bs@(PS fptr _ _) <- mallocBS siz0+      withForeignPtr fptr $ \ptr -> do+          ptr' <- copy ptr init0+          full <- recvBuf ptr' (siz0 - len0)+          if full then+              return (bs, "")+            else+              return ("", "") -- fixme++#ifndef mingw32_HOST_OS+-- fixme: the type of the return value+foreign import ccall unsafe "recv"+    c_recv :: CInt -> Ptr CChar -> CSize -> CInt -> IO CInt+#endif
+ Network/Socket/BufferPool/Types.hs view
@@ -0,0 +1,32 @@+module Network.Socket.BufferPool.Types where++import Data.ByteString (ByteString)+import Data.IORef+import Data.Word (Word8)+import Foreign.Ptr (Ptr)++-- | Type for buffer.+type Buffer = Ptr Word8++-- | Type for buffer size.+type BufSize = Int++-- | Type for read buffer pool.+data BufferPool = BufferPool {+    minBufSize :: Int -- ^ If the buffer is larger than or equal to this size,+                      --   the buffer is used.+                      --   Otherwise, a new buffer is allocated.+                      --   The thrown buffer is eventually freed.+  , maxBufSize :: Int+  , poolBuffer :: IORef ByteString+  }++-- | Type for the receiving function with a buffer pool.+type Recv = IO ByteString++-- | Type for the receiving function with a buffer.+--   The result boolean indicates whether or not the buffer is fully filled.+type RecvBuf = Buffer -> BufSize -> IO Bool++-- | Type for the receiving function which receives N bytes.+type RecvN = Int -> IO ByteString
+ Network/Socket/BufferPool/Windows.hs view
@@ -0,0 +1,25 @@+{-# LANGUAGE CPP #-}+module Network.Socket.BufferPool.Windows+  ( windowsThreadBlockHack+  ) where++#ifdef mingw32_HOST_OS+import Control.Concurrent.MVar+import Control.Concurrent+import qualified Control.Exception+import Control.Monad++-- | Allow main socket listening thread to be interrupted on Windows platform+windowsThreadBlockHack :: IO a -> IO a+windowsThreadBlockHack act = do+    var <- newEmptyMVar :: IO (MVar (Either Control.Exception.SomeException a))+    -- Catch and rethrow even async exceptions, so don't bother with UnliftIO+    void . forkIO $ Control.Exception.try act >>= putMVar var+    res <- takeMVar var+    case res of+      Left  e -> Control.Exception.throwIO e+      Right r -> return r+#else+windowsThreadBlockHack :: IO a -> IO a+windowsThreadBlockHack = id+#endif
+ recv.cabal view
@@ -0,0 +1,45 @@+Name:                recv+Version:             0.0.0+Synopsis:            Efficient netowrk recv+License:             BSD3+License-file:        LICENSE+Author:              Kazu Yamamoto+Maintainer:          kazu@iij.ad.jp+Homepage:            http://github.com/yesodweb/wai+Category:            Network+Build-Type:          Simple+Cabal-Version:       >= 1.10+Stability:           Stable+description:         Network recv based on buffer pools++Library+  Build-Depends:     base                      >= 4.12       && < 5+                   , bytestring                >= 0.9.1.4+                   , network                   >= 3.1.0+  Exposed-modules:   Network.Socket.BufferPool+  Other-modules:     Network.Socket.BufferPool.Buffer+                     Network.Socket.BufferPool.Recv+                     Network.Socket.BufferPool.Types+                     Network.Socket.BufferPool.Windows+  if impl(ghc >= 8)+      Default-Extensions:  Strict StrictData+  Ghc-Options:       -Wall+  Default-Language:  Haskell2010++Test-Suite spec+  Main-Is:           Spec.hs+  Other-modules:     BufferPoolSpec+                     Network.Socket.BufferPool+                     Network.Socket.BufferPool.Buffer+                     Network.Socket.BufferPool.Recv+                     Network.Socket.BufferPool.Types+                     Network.Socket.BufferPool.Windows+  Hs-Source-Dirs:    test, .+  Type:              exitcode-stdio-1.0+  Build-Depends:     base                      >= 4.12       && < 5+                   , bytestring                >= 0.9.1.4+                   , network                   >= 3.1.0+                   , hspec+  Ghc-Options:       -Wall+  Default-Language:  Haskell2010+  Build-Tool-Depends: hspec-discover:hspec-discover
+ test/BufferPoolSpec.hs view
@@ -0,0 +1,41 @@+module BufferPoolSpec where++import qualified Data.ByteString as B+import qualified Data.ByteString.Internal as B (ByteString(PS))+import Foreign.ForeignPtr (withForeignPtr)+import Foreign.Marshal.Utils (copyBytes)+import Foreign.Ptr (plusPtr)++import Network.Socket.BufferPool+import Test.Hspec (Spec, hspec, shouldBe, describe, it)++main :: IO ()+main = hspec spec++-- Two ByteStrings each big enough to fill a buffer (16K).+wantData, otherData :: B.ByteString+wantData  = B.replicate 16384 0xac+otherData = B.replicate 16384 0x77++spec :: Spec+spec = describe "withBufferPool" $ do+    it "does not clobber buffers" $ do+        pool <- newBufferPool 2048 16384+        -- 'pool' contains B.empty; prime it to contain a real buffer.+        _ <- withBufferPool pool $ \_ _ -> return 0+        -- 'pool' contains a 16K buffer; fill it with \xac and keep the result.+        got <- withBufferPool pool $ blitBuffer wantData+        got `shouldBe` wantData+        -- 'pool' should now be empty and reallocate, rather than clobber the+        -- previous buffer.+        _ <- withBufferPool pool $ blitBuffer otherData+        got `shouldBe` wantData++-- Fill the Buffer with the contents of the ByteString and return the number of+-- bytes written.  To be used with 'withBufferPool'.+blitBuffer :: B.ByteString -> Buffer -> BufSize -> IO Int+blitBuffer (B.PS fp off len) dst len' = withForeignPtr fp $ \ptr -> do+    let src = ptr `plusPtr` off+        n = min len len'+    copyBytes dst src n+    return n
+ test/Spec.hs view
@@ -0,0 +1,1 @@+{-# OPTIONS_GHC -F -pgmF hspec-discover #-}