module Thrift.Transport.Framed
( module Thrift.Transport
, FramedTransport
, openFramedTransport
) where
import Thrift.Transport
import Control.Monad (liftM)
import Data.Int (Int32)
import Data.Monoid (mappend, mempty)
import Control.Concurrent.MVar
import qualified Data.Binary as B
import qualified Data.Binary.Builder as BB
import qualified Data.ByteString.Lazy as LBS
data FramedTransport t = FramedTransport {
wrappedTrans :: t,
writeBuffer :: WriteBuffer,
readBuffer :: ReadBuffer
}
openFramedTransport :: Transport t => t -> IO (FramedTransport t)
openFramedTransport trans = do
wbuf <- newWriteBuffer
rbuf <- newReadBuffer
return FramedTransport{ wrappedTrans = trans, writeBuffer = wbuf, readBuffer = rbuf }
instance Transport t => Transport (FramedTransport t) where
tClose = tClose . wrappedTrans
tRead trans n = do
bs <- readBuf (readBuffer trans) n
if LBS.null bs
then
do len <- readFrame trans
if len > 0
then tRead trans n
else return bs
else return bs
tWrite trans = writeBuf (writeBuffer trans)
tFlush trans = do
bs <- flushBuf (writeBuffer trans)
let szBs = B.encode $ (fromIntegral $ LBS.length bs :: Int32)
tWrite (wrappedTrans trans) szBs
tWrite (wrappedTrans trans) bs
tFlush (wrappedTrans trans)
tIsOpen = tIsOpen . wrappedTrans
readFrame :: Transport t => FramedTransport t -> IO Int
readFrame trans = do
szBs <- tRead (wrappedTrans trans) 4
let sz = fromIntegral (B.decode szBs :: Int32)
bs <- tRead (wrappedTrans trans) sz
fillBuf (readBuffer trans) bs
return sz
type WriteBuffer = MVar (BB.Builder)
newWriteBuffer :: IO WriteBuffer
newWriteBuffer = newMVar mempty
writeBuf :: WriteBuffer -> LBS.ByteString -> IO ()
writeBuf w s = modifyMVar_ w $ return . (\builder ->
builder `mappend` (BB.fromLazyByteString s))
flushBuf :: WriteBuffer -> IO (LBS.ByteString)
flushBuf w = BB.toLazyByteString `liftM` swapMVar w mempty
type ReadBuffer = MVar (LBS.ByteString)
newReadBuffer :: IO ReadBuffer
newReadBuffer = newMVar mempty
fillBuf :: ReadBuffer -> LBS.ByteString -> IO ()
fillBuf r s = swapMVar r s >> return ()
readBuf :: ReadBuffer -> Int -> IO (LBS.ByteString)
readBuf r n = modifyMVar r $ return . flipPair . LBS.splitAt (fromIntegral n)
where flipPair (a, b) = (b, a)