From b68c7d9c0fd13df235788bc37c3df67e7595a092 Mon Sep 17 00:00:00 2001
From: Joey Hess <joey@kitenet.net>
Date: Mon, 10 Nov 2014 12:06:29 -0400
Subject: merge from git-annex

---
 Utility/Metered.hs | 34 ++++++++++++++++++++++------------
 1 file changed, 22 insertions(+), 12 deletions(-)

(limited to 'Utility/Metered.hs')

diff --git a/Utility/Metered.hs b/Utility/Metered.hs
index 4618aec..f27eee2 100644
--- a/Utility/Metered.hs
+++ b/Utility/Metered.hs
@@ -98,34 +98,44 @@ 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
 
 	loop sofar = do
-		c <- S.hGetSome h defaultChunkSize
+		c <- S.hGet h defaultChunkSize
 		if S.null c
 			then do
 				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
-- 
cgit v1.2.3