more p2p progress meters

Display progress meter on send and receive from remote.

Added a new hGetMetered that can read an exact number of bytes (or
less), updating a meter as it goes.

This commit was sponsored by Andreas on Patreon.

1 files changed, 21 insertions, 9 deletions

diff --git a/Utility/Metered.hs b/Utility/Metered.hs
index aa65efd4d..b80d3ae3f 100644
--- a/Utility/Metered.hs
+++ b/Utility/Metered.hs
@@ -1,6 +1,6 @@
 {- Metered IO and actions
  -
- - Copyright 2012-2106 Joey Hess <id@joeyh.name>
+ - Copyright 2012-2016 Joey Hess <id@joeyh.name>
  -
  - License: BSD-2-clause
  -}
@@ -115,24 +115,24 @@ offsetMeterUpdate base offset = \n -> base (offset `addBytesProcessed` n)
  - meter updates, so use caution.
  -}
 hGetContentsMetered :: Handle -> MeterUpdate -> IO L.ByteString
-hGetContentsMetered h = hGetUntilMetered h (const True)
+hGetContentsMetered h = hGetMetered h Nothing
 
-{- Reads from the Handle, updating the meter after each chunk.
+{- Reads from the Handle, updating the meter after each chunk is read.
+ -
+ - Stops at EOF, or when the requested number of bytes have been read.
+ - Closes the Handle at EOF, but otherwise leaves it open.
  -
  - 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.
- -
- - Stops at EOF, or when keepgoing evaluates to False.
- - Closes the Handle at EOF, but otherwise leaves it open.
  -}
-hGetUntilMetered :: Handle -> (Integer -> Bool) -> MeterUpdate -> IO L.ByteString
-hGetUntilMetered h keepgoing meterupdate = lazyRead zeroBytesProcessed
+hGetMetered :: Handle -> Maybe Integer -> MeterUpdate -> IO L.ByteString
+hGetMetered h wantsize meterupdate = lazyRead zeroBytesProcessed
   where
 	lazyRead sofar = unsafeInterleaveIO $ loop sofar
 
 	loop sofar = do
-		c <- S.hGet h defaultChunkSize
+		c <- S.hGet h (nextchunksize (fromBytesProcessed sofar))
 		if S.null c
 			then do
 				hClose h
@@ -148,6 +148,18 @@ hGetUntilMetered h keepgoing meterupdate = lazyRead zeroBytesProcessed
 						cs <- lazyRead sofar'
 						return $ L.append (L.fromChunks [c]) cs
 					else return $ L.fromChunks [c]
+	
+	keepgoing n = case wantsize of
+		Nothing -> True
+		Just sz -> n < sz
+	
+	nextchunksize n = case wantsize of
+		Nothing -> defaultChunkSize
+		Just sz -> 
+			let togo = sz - n
+			in if togo < toInteger defaultChunkSize
+				then fromIntegral togo
+				else defaultChunkSize
 
 {- Same default chunk size Lazy ByteStrings use. -}
 defaultChunkSize :: Int