Implement readExactly using hGet
This commit is contained in:
parent
b3effd4439
commit
19e631c9f4
2 changed files with 4 additions and 37 deletions
|
@ -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
|
||||||
|
|
|
@ -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
|
||||||
|
|
Loading…
Reference in a new issue