wait on child transfer processes, and invalidate cache

There's still a bug; if the child updates its transfer info file,
then the data from it will superscede the TransferInfo, losing the
info that we should wait on this child.

6 files changed, 40 insertions, 19 deletions

diff --git a/Assistant/DaemonStatus.hs b/Assistant/DaemonStatus.hs
index 40816bb1a..64c441cee 100644
--- a/Assistant/DaemonStatus.hs
+++ b/Assistant/DaemonStatus.hs
@@ -54,8 +54,11 @@ newDaemonStatus = DaemonStatus
 getDaemonStatus :: DaemonStatusHandle -> Annex DaemonStatus
 getDaemonStatus = liftIO . readMVar
 
-modifyDaemonStatus :: DaemonStatusHandle -> (DaemonStatus -> DaemonStatus) -> Annex ()
-modifyDaemonStatus handle a = liftIO $ modifyMVar_ handle (return . a)
+modifyDaemonStatus_ :: DaemonStatusHandle -> (DaemonStatus -> DaemonStatus) -> Annex ()
+modifyDaemonStatus_ handle a = liftIO $ modifyMVar_ handle (return . a)
+
+modifyDaemonStatus :: DaemonStatusHandle -> (DaemonStatus -> (DaemonStatus, b)) -> Annex b
+modifyDaemonStatus handle a = liftIO $ modifyMVar handle (return . a)
 
 {- Load any previous daemon status file, and store it in the MVar for this
  - process to use as its DaemonStatus. Also gets current transfer status. -}
@@ -137,5 +140,15 @@ tenMinutes = 10 * 60
 
 {- Mutates the transfer map. -}
 adjustTransfers :: DaemonStatusHandle -> (TransferMap -> TransferMap) -> Annex ()
-adjustTransfers dstatus a = modifyDaemonStatus dstatus $
+adjustTransfers dstatus a = modifyDaemonStatus_ dstatus $
 	\s -> s { currentTransfers = a (currentTransfers s) }
+
+{- Removes a transfer from the map, and returns its info. -}
+removeTransfer :: DaemonStatusHandle -> Transfer -> Annex (Maybe TransferInfo)
+removeTransfer dstatus t = modifyDaemonStatus dstatus go
+	where
+		go s =
+			let (info, ts) = M.updateLookupWithKey
+				(\_k _v -> Nothing)
+				t (currentTransfers s)
+			in (s { currentTransfers = ts }, info)
diff --git a/Assistant/Threads/SanityChecker.hs b/Assistant/Threads/SanityChecker.hs
index d7b117cd0..c5b99863e 100644
--- a/Assistant/Threads/SanityChecker.hs
+++ b/Assistant/Threads/SanityChecker.hs
@@ -26,7 +26,7 @@ sanityCheckerThread st status transferqueue changechan = forever $ do
 	waitForNextCheck st status
 
 	runThreadState st $
-		modifyDaemonStatus status $ \s -> s
+		modifyDaemonStatus_ status $ \s -> s
 				{ sanityCheckRunning = True }
 
 	now <- getPOSIXTime -- before check started
@@ -34,7 +34,7 @@ sanityCheckerThread st status transferqueue changechan = forever $ do
 		(runThreadState st . warning . show)
 
 	runThreadState st $ do
-		modifyDaemonStatus status $ \s -> s
+		modifyDaemonStatus_ status $ \s -> s
 			{ sanityCheckRunning = False
 			, lastSanityCheck = Just now
 			}
diff --git a/Assistant/Threads/TransferWatcher.hs b/Assistant/Threads/TransferWatcher.hs
index 48c0c79ae..4e468a416 100644
--- a/Assistant/Threads/TransferWatcher.hs
+++ b/Assistant/Threads/TransferWatcher.hs
@@ -13,6 +13,7 @@ import Assistant.DaemonStatus
 import Logs.Transfer
 import Utility.DirWatcher
 import Utility.Types.DirWatcher
+import Annex.BranchState
 
 import Data.Map as M
 
@@ -51,16 +52,27 @@ onErr _ _ msg _ = error msg
 onAdd :: Handler
 onAdd st dstatus file _ = case parseTransferFile file of
 	Nothing -> noop
-	Just t -> do
-		pid <- getProcessID
-		runThreadState st $ go t pid =<< checkTransfer t
+	Just t -> runThreadState st $ go t =<< checkTransfer t
 	where
-		go _ _ Nothing = noop -- transfer already finished
-		go t pid (Just info) = adjustTransfers dstatus $
+		go _ Nothing = noop -- transfer already finished
+		go t (Just info) = adjustTransfers dstatus $
 			M.insertWith' const t info
 
-{- Called when a transfer information file is removed. -}
+{- Called when a transfer information file is removed.
+ -
+ - When the transfer process is a child of this process, wait on it
+ - to avoid zombies.
+ -}
 onDel :: Handler
 onDel st dstatus file _ = case parseTransferFile file of
 	Nothing -> noop
-	Just t -> runThreadState st $ adjustTransfers dstatus $ M.delete t
+	Just t -> maybe noop waitchild
+		=<< runThreadState st (removeTransfer dstatus t)
+	where
+		waitchild info
+			| shouldWait info = case transferPid info of
+				Nothing -> noop
+				Just pid -> do
+					void $ getProcessStatus True False pid
+					runThreadState st invalidateCache
+			| otherwise = noop	
diff --git a/Assistant/Threads/Transferrer.hs b/Assistant/Threads/Transferrer.hs
index 5bc47cfa6..09c0aa036 100644
--- a/Assistant/Threads/Transferrer.hs
+++ b/Assistant/Threads/Transferrer.hs
@@ -44,12 +44,6 @@ shouldTransfer dstatus t = go =<< currentTransfers <$> getDaemonStatus dstatus
 				not <$> inAnnex (transferKey t)
 			| otherwise = return True
 
-{- Waits for any of the transfers in the map to complete. -}
-waitTransfer :: IO ()
-waitTransfer = error "TODO"
---			getProcessStatus True False pid
---			runThreadState st invalidateCache
-
 {- A transfer is run in a separate process, with a *copy* of the Annex
  - state. This is necessary to avoid blocking the rest of the assistant
  - on the transfer completing, and also to allow multiple transfers to run
@@ -81,4 +75,5 @@ runTransfer st dstatus t info
 						M.insertWith' const t info
 							{ startedTime = Just now
 							, transferPid = Just pid
+							, shouldWait = True
 							}
diff --git a/Assistant/Threads/Watcher.hs b/Assistant/Threads/Watcher.hs
index 882aab3a7..9f0eba74e 100644
--- a/Assistant/Threads/Watcher.hs
+++ b/Assistant/Threads/Watcher.hs
@@ -67,7 +67,7 @@ statupScan st dstatus scanner = do
 		showAction "scanning"
 	r <- scanner
 	runThreadState st $
-		modifyDaemonStatus dstatus $ \s -> s { scanComplete = True }
+		modifyDaemonStatus_ dstatus $ \s -> s { scanComplete = True }
 
 	-- Notice any files that were deleted before watching was started.
 	runThreadState st $ do
diff --git a/Assistant/TransferQueue.hs b/Assistant/TransferQueue.hs
index bb65dbae5..5e1fad456 100644
--- a/Assistant/TransferQueue.hs
+++ b/Assistant/TransferQueue.hs
@@ -27,6 +27,7 @@ stubInfo f = TransferInfo
 	, transferRemote = Nothing
 	, bytesComplete = Nothing
 	, associatedFile = f
+	, shouldWait = False
 	}
 
 {- Adds pending transfers to the end of the queue for some of the known