recv-0.1.3: test/BufferPoolSpec.hs

{-# 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)
import Foreign.ForeignPtr (withForeignPtr)
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)

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 = do
    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

    describe "tryWithBufferPool" $ do
        it "returns the filled prefix when the filler fills" $ do
            pool <- newBufferPool 2048 16384
            tryWithBufferPool pool (blitBuffer (B.take 10 wantData))
                `shouldReturn` Just (B.take 10 wantData)

        it "returns an empty ByteString when the filler consumes nothing" $ do
            -- 'receiveNoWait' reports EOF this way, so it must not be 'Nothing'.
            pool <- newBufferPool 2048 16384
            tryWithBufferPool pool (\_ _ -> return 0) `shouldReturn` Just B.empty

        it "returns Nothing when the filler declines" $ do
            pool <- newBufferPool 2048 16384
            tryWithBufferPool pool (\_ _ -> return (-1)) `shouldReturn` Nothing

        it "keeps the leftover buffer when the filler declines" $ do
            pool <- newBufferPool 2048 16384
            -- Consume 10000 of the 16384 bytes, leaving 6384 in the pool.
            _ <- tryWithBufferPool pool $ \_ _ -> return 10000
            tryWithBufferPool pool (\_ _ -> return (-1)) `shouldReturn` Nothing
            -- The declined call must neither consume, drop nor reallocate the
            -- leftover: the next filler is offered exactly those 6384 bytes.
            offered <- offeredSize pool
            offered `shouldBe` 6384

        it "does not corrupt buffered data across a decline" $ do
            pool <- newBufferPool 2048 16384
            _ <- tryWithBufferPool pool $ \_ _ -> return 0
            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
offeredSize pool = do
    ref <- newIORef 0
    _ <- tryWithBufferPool pool $ \_ size -> writeIORef ref size >> return 0
    readIORef ref

-- 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