allow Retriever action to update the progress meter

Needed for eg, Remote.External.

Generally, any Retriever that stores content in a file is responsible for
updating the meter, while ones that procude a lazy bytestring cannot update
the meter, so are not asked to.

3 files changed, 31 insertions, 15 deletions

diff --git a/Remote/Directory.hs b/Remote/Directory.hs
index 9f2775965..37942a295 100644
--- a/Remote/Directory.hs
+++ b/Remote/Directory.hs
@@ -137,8 +137,8 @@ store d chunkconfig k b p = liftIO $ do
 
 retrieve :: FilePath -> ChunkConfig -> Preparer Retriever
 retrieve d (LegacyChunks _) = Legacy.retrieve locations d
-retrieve d _ = simplyPrepare $ byteRetriever $ 
-	\k -> liftIO $ L.readFile =<< getLocation d k
+retrieve d _ = simplyPrepare $ byteRetriever $ \k ->
+	liftIO $ L.readFile =<< getLocation d k
 
 retrieveCheap :: FilePath -> ChunkConfig -> Key -> FilePath -> Annex Bool
 -- no cheap retrieval possible for chunks
diff --git a/Remote/Helper/Chunked.hs b/Remote/Helper/Chunked.hs
index ccdd35271..102ced8f4 100644
--- a/Remote/Helper/Chunked.hs
+++ b/Remote/Helper/Chunked.hs
@@ -221,7 +221,7 @@ retrieveChunks
 	-> Key
 	-> FilePath
 	-> MeterUpdate
-	-> (Handle -> MeterUpdate -> L.ByteString -> IO ())
+	-> (Handle -> Maybe MeterUpdate -> L.ByteString -> IO ())
 	-> Annex Bool
 retrieveChunks retriever u chunkconfig encryptor basek dest basep sink
 	| noChunks chunkconfig =
@@ -245,18 +245,18 @@ retrieveChunks retriever u chunkconfig encryptor basek dest basep sink
 	firstavail _ [] = return False
 	firstavail currsize ([]:ls) = firstavail currsize ls
 	firstavail currsize ((k:ks):ls) = do
-		v <- tryNonAsyncAnnex $ retriever (encryptor k)
+		let offset = resumeOffset currsize k
+		let p = maybe basep
+			(offsetMeterUpdate basep . toBytesProcessed)
+			offset
+		v <- tryNonAsyncAnnex $ retriever (encryptor k) p
 		case v of
 			Left e
 				| null ls -> giveup e
 				| otherwise -> firstavail currsize ls
 			Right content -> do
-				let offset = resumeOffset currsize k
-				let p = maybe basep
-					(offsetMeterUpdate basep . toBytesProcessed)
-					offset
 				bracketIO (maybe opennew openresume offset) hClose $ \h -> do
-					withBytes content $ liftIO . sink h p
+					tosink h p content
 					let sz = toBytesProcessed $
 						fromMaybe 0 $ keyChunkSize k
 					getrest p h sz sz ks
@@ -264,13 +264,11 @@ retrieveChunks retriever u chunkconfig encryptor basek dest basep sink
 	getrest _ _ _ _ [] = return True
 	getrest p h sz bytesprocessed (k:ks) = do
 		let p' = offsetMeterUpdate p bytesprocessed
-		content <- retriever (encryptor k)
-		withBytes content $ liftIO . sink h p'
+		tosink h p' =<< retriever (encryptor k) p'
 		getrest p h sz (addBytesProcessed bytesprocessed sz) ks
 
 	getunchunked = bracketIO opennew hClose $ \h -> do
-		content <- retriever (encryptor basek)
-		withBytes content $ liftIO . sink h basep
+		tosink h basep =<< retriever (encryptor basek) basep
 		return True
 
 	opennew = openBinaryFile dest WriteMode
@@ -282,6 +280,19 @@ retrieveChunks retriever u chunkconfig encryptor basek dest basep sink
 		hSeek h AbsoluteSeek startpoint
 		return h
 
+ 	{- Progress meter updating is a bit tricky: If the Retriever
+	 - populates a file, it is responsible for updating progress
+	 - as the file is being retrieved. 
+	 -
+	 - However, if the Retriever generates a lazy ByteString,
+	 - it is not responsible for updating progress (often it cannot).
+	 - Instead, the sink is passed a meter to update  as it consumes
+	 - the ByteString. -}
+	tosink h p (ByteContent b) = liftIO $
+		sink h (Just p) b
+	tosink h _ (FileContent f) = liftIO $
+		sink h Nothing =<< L.readFile f
+
 {- Can resume when the chunk's offset is at or before the end of
  - the dest file. -}
 resumeOffset :: Maybe Integer -> Key -> Maybe Integer
diff --git a/Remote/Helper/ChunkedEncryptable.hs b/Remote/Helper/ChunkedEncryptable.hs
index b851ecd94..024a53309 100644
--- a/Remote/Helper/ChunkedEncryptable.hs
+++ b/Remote/Helper/ChunkedEncryptable.hs
@@ -14,6 +14,7 @@ module Remote.Helper.ChunkedEncryptable (
 	Storer,
 	Retriever,
 	simplyPrepare,
+	ContentSource,
 	checkPrepare,
 	fileStorer,
 	byteStorer,
@@ -36,6 +37,8 @@ import Remote.Helper.Encryptable as X
 import Annex.Content
 import Annex.Exception
 
+import qualified Data.ByteString.Lazy as L
+
 simplyPrepare :: helper -> Preparer helper
 simplyPrepare helper _ a = a $ Just helper
 
@@ -101,8 +104,10 @@ chunkedEncryptableRemote c preparestorer prepareretriever baser = encr
 			retrieveChunks retriever (uuid baser) chunkconfig
 				enck k dest p' sink
 		go Nothing = return False
-		sink h p' b = do
-			let write = meteredWrite p' h
+		sink h mp b = do
+			let write = case mp of
+				Just p' -> meteredWrite p' h
+				Nothing -> L.hPut h
 			case enc of
 				Nothing -> write b
 				Just (cipher, _) ->