diff --git a/ChangeLog.md b/ChangeLog.md --- a/ChangeLog.md +++ b/ChangeLog.md @@ -1,5 +1,11 @@ # ChangeLog for "recv" +## v0.1.3 + +* Adding `receiveWith`, a `receive` the caller can give up on through an + `STM` action: nothing is closed and no thread is interrupted to make it + return. Requires network 3.2.10. + ## v0.1.2 * Providing `receiveNoWait` and `tryWithBufferPool`. diff --git a/Network/Socket/BufferPool.hs b/Network/Socket/BufferPool.hs --- a/Network/Socket/BufferPool.hs +++ b/Network/Socket/BufferPool.hs @@ -25,6 +25,7 @@ Recv, receive, receiveNoWait, + receiveWith, BufferPool, newBufferPool, withBufferPool, diff --git a/Network/Socket/BufferPool/Buffer.hs b/Network/Socket/BufferPool/Buffer.hs --- a/Network/Socket/BufferPool/Buffer.hs +++ b/Network/Socket/BufferPool/Buffer.hs @@ -2,6 +2,7 @@ newBufferPool, withBufferPool, tryWithBufferPool, + eitherWithBufferPool, mallocBS, copy, ) where @@ -34,6 +35,8 @@ -- how many bytes are filled in the buffer. -- This function should return non negative 'Int'. -- The buffer in the buffer pool is automatically managed. +-- +-- /Caveats of 'BufferPool' apply./ withBufferPool :: BufferPool -> (Buffer -> BufSize -> IO Int) -> IO ByteString withBufferPool pool@(BufferPool _ _ ref) f = do (buf, consumed) <- applyBufferPool pool f @@ -43,6 +46,8 @@ -- | L ike 'withBufferPool' for fillers that can decline to fill: -- a negative return value from the filler leaves the pool untouched -- and produces 'Nothing'. +-- +-- /Caveats of 'BufferPool' apply./ tryWithBufferPool :: BufferPool -> (Buffer -> BufSize -> IO Int) -> IO (Maybe ByteString) tryWithBufferPool pool@(BufferPool _ _ ref) f = do @@ -55,8 +60,26 @@ writeIORef ref $ BS.drop consumed buf return $ Just $ BS.take consumed buf +-- | Like 'withBufferPool' for fillers which can give up: 'Left' leaves +-- the pool untouched. +-- +-- /Caveats of 'BufferPool' apply./ +eitherWithBufferPool + :: BufferPool + -> (Buffer -> BufSize -> IO (Either a Int)) + -> IO (Either a ByteString) +eitherWithBufferPool pool@(BufferPool _ _ ref) f = do + (buf, r) <- applyBufferPool pool f + case r of + Left a -> do + writeIORef ref buf + return $ Left a + Right consumed -> do + writeIORef ref $ BS.drop consumed buf + return $ Right $ BS.take consumed buf + applyBufferPool - :: BufferPool -> (Buffer -> BufSize -> IO Int) -> IO (ByteString, Int) + :: BufferPool -> (Buffer -> BufSize -> IO r) -> IO (ByteString, r) applyBufferPool (BufferPool l h ref) f = do buf0 <- readIORef ref buf <- @@ -66,7 +89,7 @@ consumed <- withForeignBuffer buf f return (buf, consumed) -withForeignBuffer :: ByteString -> (Buffer -> BufSize -> IO Int) -> IO Int +withForeignBuffer :: ByteString -> (Buffer -> BufSize -> IO r) -> IO r withForeignBuffer (PS ps s l) f = withForeignPtr ps $ \p -> f (castPtr p `plusPtr` s) l {-# INLINE withForeignBuffer #-} diff --git a/Network/Socket/BufferPool/Recv.hs b/Network/Socket/BufferPool/Recv.hs --- a/Network/Socket/BufferPool/Recv.hs +++ b/Network/Socket/BufferPool/Recv.hs @@ -3,13 +3,15 @@ module Network.Socket.BufferPool.Recv ( receive, receiveNoWait, + receiveWith, makeRecvN, ) where import qualified Data.ByteString as BS import Data.ByteString.Internal (ByteString (..), unsafeCreate) import Data.IORef -import Network.Socket (Socket, recvBuf, recvBufNoWait) +import Control.Concurrent.STM (STM) +import Network.Socket (Socket, recvBuf, recvBufNoWait, recvBufWith) import Network.Socket.BufferPool.Buffer import Network.Socket.BufferPool.Types @@ -18,12 +20,16 @@ -- | The receiving function with a buffer pool. -- The buffer pool is automatically managed. +-- +-- /Caveats of 'BufferPool' apply./ receive :: Socket -> BufferPool -> Recv receive sock pool = withBufferPool pool $ \ptr size -> recvBuf sock ptr size -- | Like 'receive' but never blocks and never involves the IO manager: -- 'Nothing' means no data was available (or an error occurred, which a -- subsequent blocking 'receive' will report properly). @Just \"\"@ is EOF. +-- +-- /Caveats of 'BufferPool' apply./ receiveNoWait :: Socket -> BufferPool -> IO (Maybe ByteString) receiveNoWait sock pool = tryWithBufferPool pool $ \ptr size -> -- Both EAGAIN and real errors map to a negative result, deferring @@ -31,6 +37,24 @@ -- intact. fromIntegral <$> recvBufNoWait sock ptr size +-- | Like 'receive' but the caller can give up on it: whichever settles +-- first decides, 'Left' when the 'STM' action fires, 'Right' when data +-- arrives. Data already waiting is taken without blocking at all and +-- so wins over a caller which has given up, as a completed operation +-- does in 'Network.Socket.withReadable'; the give-up is read there +-- and nowhere else. +-- +-- Nothing is closed and no thread is interrupted to make this return. +-- +-- /Caveats of 'BufferPool' apply./ +receiveWith :: Socket -> BufferPool -> STM a -> IO (Either a ByteString) +receiveWith sock pool giveUp = do + early <- receiveNoWait sock pool + case early of + Just bs -> return $ Right bs + Nothing -> eitherWithBufferPool pool $ \ptr size -> + recvBufWith sock ptr size giveUp + ---------------------------------------------------------------- -- | This function returns a receiving function @@ -49,6 +73,8 @@ -- ("abc","") -- >>> tryRecvN "a" 3 =<< _iorefRecv ["b"] -- ("ab","") +-- +-- /The resulting 'RecvN' is designed to be used in one thread only./ makeRecvN :: ByteString -> Recv -> IO RecvN makeRecvN bs0 recv = do ref <- newIORef bs0 diff --git a/Network/Socket/BufferPool/Types.hs b/Network/Socket/BufferPool/Types.hs --- a/Network/Socket/BufferPool/Types.hs +++ b/Network/Socket/BufferPool/Types.hs @@ -12,6 +12,10 @@ type BufSize = Int -- | Type for read buffer pool. +-- +-- /Caveats: designed to be used by one thread only./ +-- /The buffer is not behind a lock, so using a 'BufferPool'/ +-- /in more than one thread might create race conditions./ data BufferPool = BufferPool { minBufSize :: Int -- ^ If the buffer is larger than or equal to this size, diff --git a/recv.cabal b/recv.cabal --- a/recv.cabal +++ b/recv.cabal @@ -1,6 +1,6 @@ cabal-version: >=1.10 name: recv -version: 0.1.2 +version: 0.1.3 license: BSD3 license-file: LICENSE maintainer: kazu@iij.ad.jp @@ -25,7 +25,8 @@ build-depends: base >=4.12 && <5, bytestring >=0.9.1.4, - network >=3.2.9 + network >=3.2.10, + stm >=2.5 && <2.6 if impl(ghc >=8) default-extensions: Strict StrictData @@ -43,10 +44,11 @@ Network.Socket.BufferPool.Types default-language: Haskell2010 - ghc-options: -Wall + ghc-options: -Wall -threaded build-depends: base >=4.12 && <5, bytestring >=0.9.1.4, - network >=3.2.9, + network >=3.2.10, + stm >=2.5 && <2.6, hspec, recv diff --git a/test/BufferPoolSpec.hs b/test/BufferPoolSpec.hs --- a/test/BufferPoolSpec.hs +++ b/test/BufferPoolSpec.hs @@ -1,5 +1,10 @@ +{-# LANGUAGE OverloadedStrings #-} module BufferPoolSpec where +import Control.Concurrent (forkIO, threadDelay) +import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar) +import Control.Concurrent.STM +import Control.Monad (void) import qualified Data.ByteString as B import qualified Data.ByteString.Internal as B (ByteString (PS)) import Data.IORef (newIORef, readIORef, writeIORef) @@ -7,6 +12,8 @@ import Foreign.Marshal.Utils (copyBytes) import Foreign.Ptr (plusPtr) +import Network.Socket +import qualified Network.Socket.ByteString as NSB import Network.Socket.BufferPool import Test.Hspec (Spec, describe, hspec, it, shouldBe, shouldReturn) @@ -64,6 +71,60 @@ tryWithBufferPool pool (\_ _ -> return (-1)) `shouldReturn` Nothing tryWithBufferPool pool (blitBuffer wantData) `shouldReturn` Just wantData + + describe "receiveWith" $ do + it "gives up while blocked, without closing the socket" $ do + withPair $ \_ srv pool -> do + stop <- newTVarIO False + done <- newEmptyMVar + void $ forkIO $ do + r <- receiveWith srv pool $ readTVar stop >>= check + putMVar done r + threadDelay 100000 + atomically $ writeTVar stop True + takeMVar done `shouldReturn` Left () + -- The socket is the caller's, and still theirs to read. + void $ getSocketName srv + + -- Data already there wins over a give-up, as it does in + -- withReadable: answering Left would leave it unread. + it "takes what is there even once the caller has given up" $ do + withPair $ \cli srv pool -> do + void $ NSB.send cli "unread" + threadDelay 100000 + gone <- newTVarIO True + receiveWith srv pool (readTVar gone >>= check) + `shouldReturn` Right "unread" + + it "reads nothing once the caller has given up and none is there" $ do + withPair $ \_cli srv pool -> do + gone <- newTVarIO True + receiveWith srv pool (readTVar gone >>= check) + `shouldReturn` Left () + + it "takes what is already there without waiting" $ do + withPair $ \cli srv pool -> do + void $ NSB.send cli "hello" + threadDelay 100000 + receiveWith srv pool (retry :: STM ()) `shouldReturn` Right "hello" + +-- A connected pair on the loopback, and a pool to receive with. +withPair :: (Socket -> Socket -> BufferPool -> IO a) -> IO a +withPair action = do + lsock <- socket AF_INET Stream defaultProtocol + setSocketOption lsock ReuseAddr 1 + bind lsock $ SockAddrInet 0 $ tupleToHostAddress (127, 0, 0, 1) + listen lsock 1 + port <- socketPort lsock + cli <- socket AF_INET Stream defaultProtocol + connect cli $ SockAddrInet port $ tupleToHostAddress (127, 0, 0, 1) + (srv, _) <- accept lsock + close lsock + pool <- newBufferPool 2048 16384 + r <- action cli srv pool + close cli + close srv + return r -- The 'BufSize' the pool offers to the next filler, leaving the pool as it was. offeredSize :: BufferPool -> IO Int