make commitMigration interuption safe

Fixed inversion of control issue, so the tree is recorded
in streamLogFile finalizer.

Sponsored-by: Leon Schuermann on Patreon

Git/Tree.hs

@@ -23,7 +23,8 @@ module Git.Tree (
 	graftTree',
 	withMkTreeHandle,
 	MkTreeHandle,
-	mkTree,
+	sendMkTree,
+	finishMkTree,
 	treeMode,
 ) where
 
@@ -92,39 +93,37 @@ recordTree :: Tree -> Repo -> IO Sha
 recordTree t repo = withMkTreeHandle repo $ \h -> recordTree' h t
 
 recordTree' :: MkTreeHandle -> Tree -> IO Sha
-recordTree' h (Tree l) = mkTree' h =<< mapM (recordSubTree h) l
+recordTree' h (Tree l) = mkTree h =<< mapM (recordSubTree h) l
 
 {- Note that the returned RecordedSubTree does not have its [TreeContent]
  - list populated. This is a memory optimisation, since the list is not
  - used. -}
 recordSubTree :: MkTreeHandle -> TreeContent -> IO TreeContent
 recordSubTree h (NewSubTree d l) = do
-	sha <- mkTree' h =<< mapM (recordSubTree h) l
+	sha <- mkTree h =<< mapM (recordSubTree h) l
 	return (RecordedSubTree d sha [])
 recordSubTree _ alreadyrecorded = return alreadyrecorded
 
-{- Note that this creates a single tree. It cannot create a recursive tree
- - with subtrees in a single call. -}
-mkTree
-	:: (MonadIO m, MonadCatch m)
-	=> MkTreeHandle
-	-> ((FileMode -> ObjectType -> Sha -> TopFilePath -> m ()) -> m ())
-	-> m Sha
-mkTree (MkTreeHandle cp) a = CoProcess.query cp send receive
-  where
-	send h = do
-		a $ \fm ot s f -> liftIO $ hPutStr h (mkTreeOutput fm ot s f)
-		-- NUL to signal end of tree to --batch
-		liftIO $ hPutStr h "\NUL"
-	receive h = liftIO $ getSha "mktree" (S8.hGetLine h)
+sendMkTree :: MkTreeHandle -> FileMode -> ObjectType -> Sha -> TopFilePath -> IO ()
+sendMkTree (MkTreeHandle cp) fm ot s f =
+	CoProcess.send cp $ \h -> 
+		hPutStr h (mkTreeOutput fm ot s f)
 
-mkTree' :: MkTreeHandle -> [TreeContent] -> IO Sha
-mkTree' h l = mkTree h $ \send -> 
+finishMkTree :: MkTreeHandle -> IO Sha
+finishMkTree (MkTreeHandle cp) = do
+	CoProcess.send cp $ \h ->
+		-- NUL to signal end of tree to --batch
+		hPutStr h "\NUL"
+	getSha "mktree" (CoProcess.receive cp S8.hGetLine)
+
+mkTree :: MkTreeHandle -> [TreeContent] -> IO Sha
+mkTree h l = do
 	forM_ l $ \case
-		TreeBlob f fm s -> send fm BlobObject s f
-		RecordedSubTree f s _ -> send treeMode TreeObject s f
+		TreeBlob f fm s -> sendMkTree h fm BlobObject s f
+		RecordedSubTree f s _ -> sendMkTree h treeMode TreeObject s f
 		NewSubTree _ _ -> error "recordSubTree internal error; unexpected NewSubTree"
-		TreeCommit f fm s -> send fm CommitObject s f
+		TreeCommit f fm s -> sendMkTree h fm CommitObject s f
+	finishMkTree h
 
 treeMode :: FileMode
 treeMode = 0o040000
@@ -233,7 +232,7 @@ adjustTree adjusttreeitem addtreeitems resolveaddconflict removefiles r repo =
 		(l, cleanup) <- liftIO $ lsTreeWithObjects LsTree.LsTreeRecursive r repo
 		(l', _, _) <- go h False [] 1 inTopTree l
 		l'' <- adjustlist h 0 inTopTree (const True) l'
-		sha <- liftIO $ mkTree' h l''
+		sha <- liftIO $ mkTree h l''
 		void $ liftIO cleanup
 		return sha
   where
@@ -344,7 +343,7 @@ graftTree' subtree graftloc basetree repo hdl = go basetree subdirs graftdirs
 							return $ RecordedSubTree tloc tsha'' []
 					_ -> graftin (topmostgraphdir:restgraphdirs)
 				return (newshas ++ rest)
-		mkTree' hdl t'
+		mkTree hdl t'
 	go _ _ [] = return subtree
 	go _ [] _ = return subtree
 	

Logs/Migrate.hs

@@ -66,16 +66,31 @@ commitMigration = do
 	logf <- fromRawFilePath <$> fromRepo gitAnnexMigrateLog
 	lckf <- fromRepo gitAnnexMigrateLock
 	nv <- liftIO $ newTVarIO (0 :: Integer)
-	newtv <- liftIO newEmptyTMVarIO
 	g <- Annex.gitRepo
-	oldt <- withMkTreeHandle g $ \oldh ->
+	withMkTreeHandle g $ \oldh ->
 		withMkTreeHandle g $ \newh ->
-			mkTree oldh $ \oldsend -> do
-				newt <- mkTree newh $ \newsend ->
-					streamLogFile logf lckf noop $
-						processor nv oldsend newsend
-				liftIO $ atomically $ writeTMVar newtv newt
-	newt <- liftIO $ atomically $ takeTMVar newtv
+			streamLogFile logf lckf 
+				(finalizer nv oldh newh g)
+				(processor nv oldh newh)
+  where
+	processor nv oldh newh s = case words s of
+		(old:new:[]) -> do
+			fn <- liftIO $ atomically $ do
+				n <- readTVar nv
+				let !n' = succ n
+				writeTVar nv n'
+				return (asTopFilePath (encodeBS (show n')))
+			let rec h r = liftIO $ sendMkTree h
+				(fromTreeItemType TreeFile)
+				BlobObject
+				(Git.Ref (encodeBS r))
+				fn
+			rec oldh old
+			rec newh new
+		_ -> error "migrate.log parse error"
+	finalizer nv oldh newh g = do
+		oldt <- liftIO $ finishMkTree oldh
+		newt <- liftIO $ finishMkTree newh
 		n <- liftIO $ atomically $ readTVar nv
 		when (n > 0) $ do
 			treesha <- liftIO $ flip recordTree g $ Tree
@@ -84,19 +99,3 @@ commitMigration = do
 				]
 			Annex.Branch.rememberTreeish treesha
 				(asTopFilePath migrationTreeGraftPoint)
-  where
-	processor nv oldsend newsend s = case words s of
-		(old:new:[]) -> do
-			fn <- liftIO $ atomically $ do
-				n <- readTVar nv
-				let !n' = succ n
-				writeTVar nv n'
-				return (asTopFilePath (encodeBS (show n')))
-			let rec f r = f
-				(fromTreeItemType TreeFile)
-				BlobObject
-				(Git.Ref (encodeBS r))
-				fn
-			void $ rec oldsend old
-			void $ rec newsend new
-		_ -> error "migrate.log parse error"

Utility/CoProcess.hs

@@ -1,7 +1,7 @@
 {- Interface for running a shell command as a coprocess,
  - sending it queries and getting back results.
  -
- - Copyright 2012-2013 Joey Hess <id@joeyh.name>
+ - Copyright 2012-2023 Joey Hess <id@joeyh.name>
  -
  - License: BSD-2-clause
  -}
@@ -14,6 +14,8 @@ module Utility.CoProcess (
 	start,
 	stop,
 	query,
+	send,
+	receive,
 ) where
 
 import Common
@@ -67,11 +69,11 @@ stop ch = do
  - receive actions are assumed to mean communication with the process
  - failed, and the failed action is re-run with a new process. -}
 query :: (MonadIO m, MonadCatch m) => CoProcessHandle -> (Handle -> m a) -> (Handle -> m b) -> m b
-query ch send receive = do
+query ch sender receiver = do
 	s <- liftIO $ readMVar ch
-	restartable s (send $ coProcessTo s) $ const $
+	restartable s (sender $ coProcessTo s) $ const $
 		restartable s (liftIO $ hFlush $ coProcessTo s) $ const $
-			restartable s (receive $ coProcessFrom s)
+			restartable s (receiver $ coProcessFrom s)
 				return
   where
 	restartable s a cont
@@ -87,4 +89,15 @@ query ch send receive = do
 		s' <- liftIO $ start' $ (coProcessSpec s)
 			{ coProcessNumRestarts = coProcessNumRestarts (coProcessSpec s) - 1 }
 		liftIO $ putMVar ch s'
-		query ch send receive
+		query ch sender receiver
+
+send :: MonadIO m => CoProcessHandle -> (Handle -> m a) -> m a
+send ch a = do
+	s <- liftIO $ readMVar ch
+	a (coProcessTo s)
+
+receive :: MonadIO m => CoProcessHandle -> (Handle -> m a) -> m a
+receive ch a = do
+	s <- liftIO $ readMVar ch
+	liftIO $ hFlush $ coProcessTo s
+	a (coProcessFrom s)