154 lines
5.7 KiB
Haskell
154 lines
5.7 KiB
Haskell
{- | Pipelining is sending multiple requests over a socket and receiving the responses later, in the same order. This is faster than sending one request, waiting for the response, then sending the next request, and so on. This implementation returns a /promise (future)/ response for each request that when invoked waits for the response if not already arrived. Multiple threads can send on the same pipe (and get promises back); the pipe will pipeline each thread's request right away without waiting. -}
|
|
|
|
{-# LANGUAGE DoRec, RecordWildCards, MultiParamTypeClasses, FlexibleContexts #-}
|
|
|
|
module Control.Pipeline (
|
|
-- * Pipe
|
|
Pipe, newPipe, send, call,
|
|
-- * Util
|
|
Size,
|
|
Length(..),
|
|
Resource(..),
|
|
Flush(..),
|
|
Stream(..), getN
|
|
) where
|
|
|
|
import Prelude hiding (length)
|
|
import Control.Applicative ((<$>))
|
|
import Control.Monad (forever)
|
|
import Control.Exception (assert)
|
|
import System.IO.Error (try, mkIOError, eofErrorType)
|
|
import System.IO (Handle, hFlush, hClose, hIsClosed)
|
|
import qualified Data.ByteString as S
|
|
import qualified Data.ByteString.Lazy as L
|
|
import Data.Monoid (Monoid(..))
|
|
import Control.Concurrent (ThreadId, forkIO, killThread)
|
|
import Control.Concurrent.MVar
|
|
import Control.Concurrent.Chan
|
|
|
|
-- * Length
|
|
|
|
type Size = Int
|
|
|
|
class Length list where
|
|
length :: list -> Size
|
|
|
|
instance Length S.ByteString where
|
|
length = S.length
|
|
|
|
instance Length L.ByteString where
|
|
length = fromEnum . L.length
|
|
|
|
-- * Resource
|
|
|
|
class Resource m r where
|
|
close :: r -> m ()
|
|
-- ^ Close resource
|
|
isClosed :: r -> m Bool
|
|
-- ^ Is resource closed
|
|
|
|
instance Resource IO Handle where
|
|
close = hClose
|
|
isClosed = hIsClosed
|
|
|
|
-- * Flush
|
|
|
|
class Flush handle where
|
|
flush :: handle -> IO ()
|
|
-- ^ Flush written bytes to destination
|
|
|
|
instance Flush Handle where
|
|
flush = hFlush
|
|
|
|
-- * Stream
|
|
|
|
class (Length bytes, Monoid bytes, Flush handle) => Stream handle bytes where
|
|
put :: handle -> bytes -> IO ()
|
|
-- ^ Write bytes to handle
|
|
get :: handle -> Int -> IO bytes
|
|
-- ^ Read up to N bytes from handle; if EOF return empty bytes, otherwise block until at least 1 byte is available
|
|
|
|
getN :: (Stream h b) => h -> Int -> IO b
|
|
-- ^ Read N bytes from hande, blocking until all N bytes are read. If EOF is reached before N bytes then throw EOF exception.
|
|
getN h n = assert (n >= 0) $ do
|
|
bytes <- get h n
|
|
let x = length bytes
|
|
if x >= n then return bytes
|
|
else if x == 0 then ioError (mkIOError eofErrorType "Control.Pipeline" Nothing Nothing)
|
|
else mappend bytes <$> getN h (n - x)
|
|
|
|
instance Stream Handle S.ByteString where
|
|
put = S.hPut
|
|
get = S.hGet
|
|
|
|
instance Stream Handle L.ByteString where
|
|
put = L.hPut
|
|
get = L.hGet
|
|
|
|
-- * Pipe
|
|
|
|
-- | Thread-safe and pipelined socket
|
|
data Pipe handle bytes = Pipe {
|
|
encodeSize :: Size -> bytes,
|
|
decodeSize :: bytes -> Size,
|
|
vHandle :: MVar handle, -- ^ Mutex on handle, so only one thread at a time can write to it
|
|
responseQueue :: Chan (MVar (Either IOError bytes)), -- ^ Queue of threads waiting for responses. Every time a response arrive we pop the next thread and give it the response.
|
|
listenThread :: ThreadId
|
|
}
|
|
|
|
-- | Create new Pipe with given encodeInt, decodeInt, and handle. You should 'close' pipe when finished, which will also close handle. If pipe is not closed but eventually garbage collected, it will be closed along with handle.
|
|
newPipe :: (Stream h b, Resource IO h) =>
|
|
(Size -> b) -- ^ Convert Size to bytes of fixed length. Every Int must translate to same number of bytes.
|
|
-> (b -> Size) -- ^ Convert bytes of fixed length to Size. Must be exact inverse of encodeSize.
|
|
-> h -- ^ Underlying socket (handle) this pipe will read/write from
|
|
-> IO (Pipe h b)
|
|
newPipe encodeSize decodeSize handle = do
|
|
vHandle <- newMVar handle
|
|
responseQueue <- newChan
|
|
rec
|
|
let pipe = Pipe{..}
|
|
listenThread <- forkIO (listen pipe)
|
|
addMVarFinalizer vHandle $ do
|
|
killThread listenThread
|
|
close handle
|
|
return pipe
|
|
|
|
instance (Resource IO h) => Resource IO (Pipe h b) where
|
|
-- | Close pipe and underlying socket (handle)
|
|
close Pipe{..} = do
|
|
killThread listenThread
|
|
close =<< readMVar vHandle
|
|
isClosed Pipe{..} = isClosed =<< readMVar vHandle
|
|
|
|
listen :: (Stream h b) => Pipe h b -> IO ()
|
|
-- ^ Listen for responses and supply them to waiting threads in order
|
|
listen Pipe{..} = do
|
|
let n = length (encodeSize 0)
|
|
h <- readMVar vHandle
|
|
forever $ do
|
|
e <- try $ do
|
|
len <- decodeSize <$> getN h n
|
|
getN h len
|
|
var <- readChan responseQueue
|
|
putMVar var e
|
|
|
|
send :: (Stream h b) => Pipe h b -> [b] -> IO ()
|
|
-- ^ Send messages all together to destination (no messages will be interleaved between them). None of the messages can induce a response, i.e. the destination must not reply to any of these messages (otherwise future 'call's will get these responses instead of their own).
|
|
-- Each message is preceeded by its length when written to socket.
|
|
send Pipe{..} messages = withMVar vHandle $ \h -> do
|
|
mapM_ (write encodeSize h) messages
|
|
flush h
|
|
|
|
call :: (Stream h b) => Pipe h b -> [b] -> IO (IO b)
|
|
-- ^ Send messages all together to destination (no messages will be interleaved between them), and return /promise/ of response from one message only. One and only one message in the list must induce a response, i.e. the destination must reply to exactly one message only (otherwise promises will have the wrong responses in them).
|
|
-- Each message is preceeded by its length when written to socket. Likewise, the response must be preceeded by its length.
|
|
call Pipe{..} messages = withMVar vHandle $ \h -> do
|
|
mapM_ (write encodeSize h) messages
|
|
flush h
|
|
var <- newEmptyMVar
|
|
writeChan responseQueue var
|
|
return (either ioError return =<< readMVar var) -- return promise
|
|
|
|
write :: (Stream h b, Monoid b, Length b) => (Size -> b) -> h -> b -> IO ()
|
|
write encodeSize h bytes = put h (mappend lenBytes bytes) where lenBytes = encodeSize (length bytes)
|