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
|
||||
data Connection = Connection {
|
||||
read :: IO ByteString,
|
||||
unread :: ByteString -> IO (),
|
||||
readExactly :: Int -> IO ByteString,
|
||||
write :: ByteString -> IO (),
|
||||
flush :: 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
|
||||
-- ^ Make connection form handle
|
||||
fromHandle handle = do
|
||||
restRef <- newIORef mempty
|
||||
return Connection
|
||||
{ read = do
|
||||
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 <>)
|
||||
{ readExactly = ByteString.hGet handle
|
||||
, write = ByteString.hPut handle
|
||||
, flush = hFlush handle
|
||||
, close = hClose handle
|
||||
|
|
|
@ -117,8 +117,8 @@ readMessage :: Connection -> IO Response
|
|||
-- ^ read response from a connection
|
||||
readMessage conn = readResp where
|
||||
readResp = do
|
||||
len <- fromEnum . decodeSize <$> Connection.readExactly conn 4
|
||||
runGet getReply <$> Connection.readExactly conn len
|
||||
len <- fromEnum . decodeSize . L.fromStrict <$> Connection.readExactly conn 4
|
||||
runGet getReply . L.fromStrict <$> Connection.readExactly conn len
|
||||
decodeSize = subtract 4 . runGet getInt32
|
||||
|
||||
type FullCollection = Text
|
||||
|
|
Loading…
Reference in a new issue