hGetUntilMetered

1 files changed, 21 insertions, 11 deletions

diff --git a/Utility/Metered.hs b/Utility/Metered.hs
index 4618aecfe..447eab2e8 100644
--- a/Utility/Metered.hs
+++ b/Utility/Metered.hs
@@ -99,15 +99,23 @@ offsetMeterUpdate base offset = \n -> base (offset `addBytesProcessed` n)
 {- This is like L.hGetContents, but after each chunk is read, a meter
  - is updated based on the size of the chunk.
  -
+ - All the usual caveats about using unsafeInterleaveIO apply to the
+ - meter updates, so use caution.
+ -}
+hGetContentsMetered :: Handle -> MeterUpdate -> IO L.ByteString
+hGetContentsMetered h = hGetUntilMetered h (const True)
+
+{- Reads from the Handle, updating the meter after each chunk.
+ -
  - Note that the meter update is run in unsafeInterleaveIO, which means that
  - it can be run at any time. It's even possible for updates to run out
  - of order, as different parts of the ByteString are consumed.
  -
- - All the usual caveats about using unsafeInterleaveIO apply to the
- - meter updates, so use caution.
+ - Stops at EOF, or when keepgoing evaluates to False.
+ - Closes the Handle at EOF, but otherwise leaves it open.
  -}
-hGetContentsMetered :: Handle -> MeterUpdate -> IO L.ByteString
-hGetContentsMetered h meterupdate = lazyRead zeroBytesProcessed
+hGetUntilMetered :: Handle -> (Integer -> Bool) -> MeterUpdate -> IO L.ByteString
+hGetUntilMetered h keepgoing meterupdate = lazyRead zeroBytesProcessed
   where
 	lazyRead sofar = unsafeInterleaveIO $ loop sofar
 
@@ -118,14 +126,16 @@ hGetContentsMetered h meterupdate = lazyRead zeroBytesProcessed
 				hClose h
 				return $ L.empty
 			else do
-				let sofar' = addBytesProcessed sofar $
-					S.length c
+				let sofar' = addBytesProcessed sofar (S.length c)
 				meterupdate sofar'
-				{- unsafeInterleaveIO causes this to be
-				 - deferred until the data is read from the
-				 - ByteString. -}
-				cs <- lazyRead sofar'
-				return $ L.append (L.fromChunks [c]) cs
+				if keepgoing (fromBytesProcessed sofar')
+					then do
+						{- unsafeInterleaveIO causes this to be
+						 - deferred until the data is read from the
+						 - ByteString. -}
+						cs <- lazyRead sofar'
+						return $ L.append (L.fromChunks [c]) cs
+					else return $ L.fromChunks [c]
 
 {- Same default chunk size Lazy ByteStrings use. -}
 defaultChunkSize :: Int