Implement readExactly using hGet

This commit is contained in:
Victor Denisov 2016-04-27 00:10:12 -07:00
parent b3effd4439
commit 19e631c9f4
2 changed files with 4 additions and 37 deletions

View file

@ -23,49 +23,16 @@ import System.IO.Error (mkIOError, eofErrorType)
-- --
-- `read` should return `ByteString.null` on EOF -- `read` should return `ByteString.null` on EOF
data Connection = Connection { data Connection = Connection {
read :: IO ByteString, readExactly :: Int -> IO ByteString,
unread :: ByteString -> IO (),
write :: ByteString -> IO (), write :: ByteString -> IO (),
flush :: IO (), flush :: IO (),
close :: IO ()} close :: IO ()}
readExactly :: Connection -> Int -> IO Lazy.ByteString
-- ^ Read specified number of bytes
--
-- If EOF is reached before N bytes then raise EOF exception.
readExactly conn count = go mempty count
where
go acc n = do
-- read until get enough bytes
chunk <- read conn
when (ByteString.null chunk) $
ioError eof
let len = ByteString.length chunk
if len >= n
then do
let (res, rest) = ByteString.splitAt n chunk
unless (ByteString.null rest) $
unread conn rest
return (acc <> Lazy.ByteString.fromStrict res)
else go (acc <> Lazy.ByteString.fromStrict chunk) (n - len)
eof = mkIOError eofErrorType "Database.MongoDB.Internal.Connection"
Nothing Nothing
fromHandle :: Handle -> IO Connection fromHandle :: Handle -> IO Connection
-- ^ Make connection form handle -- ^ Make connection form handle
fromHandle handle = do fromHandle handle = do
restRef <- newIORef mempty
return Connection return Connection
{ read = do { readExactly = ByteString.hGet handle
rest <- readIORef restRef
writeIORef restRef mempty
if ByteString.null rest
-- 32k corresponds to the default chunk size
-- used in bytestring package
then ByteString.hGetSome handle (32 * 1024)
else return rest
, unread = \rest ->
modifyIORef restRef (rest <>)
, write = ByteString.hPut handle , write = ByteString.hPut handle
, flush = hFlush handle , flush = hFlush handle
, close = hClose handle , close = hClose handle

View file

@ -117,8 +117,8 @@ readMessage :: Connection -> IO Response
-- ^ read response from a connection -- ^ read response from a connection
readMessage conn = readResp where readMessage conn = readResp where
readResp = do readResp = do
len <- fromEnum . decodeSize <$> Connection.readExactly conn 4 len <- fromEnum . decodeSize . L.fromStrict <$> Connection.readExactly conn 4
runGet getReply <$> Connection.readExactly conn len runGet getReply . L.fromStrict <$> Connection.readExactly conn len
decodeSize = subtract 4 . runGet getInt32 decodeSize = subtract 4 . runGet getInt32
type FullCollection = Text type FullCollection = Text