use live reposizes in balanced preferred content

Annex/RepoSize.hs

@@ -9,6 +9,7 @@
 
 module Annex.RepoSize (
 	getRepoSizes,
+	getLiveRepoSizes,
 ) where
 
 import Annex.Common
@@ -29,7 +30,11 @@ import Control.Concurrent.Async
 import qualified Data.Map.Strict as M
 import qualified Data.Set as S
 
-{- Gets the repo size map. Cached for speed. -}
+{- Gets the repo size map. Cached for speed.
+ -
+ - Note that this is the size of all repositories as of the first time it
+ - was called. It does not update while git-annex is running.
+ -}
 getRepoSizes :: Bool -> Annex (M.Map UUID RepoSize)
 getRepoSizes quiet = do
 	rsv <- Annex.getRead Annex.reposizes
@@ -39,14 +44,29 @@ getRepoSizes quiet = do
 			return sizemap
 		Nothing -> calcRepoSizes quiet rsv
 
+{- Like getRepoSizes, but with live updates. -}
+getLiveRepoSizes :: Bool -> Annex (M.Map UUID RepoSize)
+getLiveRepoSizes quiet = do
+	h <- Db.getRepoSizeHandle
+	liftIO (Db.estimateLiveRepoSizes h) >>= \case
+		Just (m, annexbranchsha) -> return m
+		Nothing -> do
+			-- Db.estimateLiveRepoSizes needs the
+			-- reposizes to be calculated first.
+			m <- getRepoSizes quiet
+			liftIO (Db.estimateLiveRepoSizes h) >>= \case
+				Just (m', annexbranchsha) -> return m'
+				Nothing -> return m
+
 {- Fills an empty Annex.reposizes MVar with current information
  - from the git-annex branch, supplimented with journalled but
  - not yet committed information.
  -}
 calcRepoSizes :: Bool -> MVar (Maybe (M.Map UUID RepoSize)) -> Annex (M.Map UUID RepoSize)
-calcRepoSizes quiet rsv = bracket setup cleanup $ \h -> go h `onException` failed
+calcRepoSizes quiet rsv = go `onException` failed
   where
-	go h = do
+	go = do
+		h <- Db.getRepoSizeHandle
 		(oldsizemap, moldbranchsha) <- liftIO $ Db.getRepoSizes h
 		!sizemap <- case moldbranchsha of
 			Nothing -> calculatefromscratch h
@@ -70,10 +90,6 @@ calcRepoSizes quiet rsv = bracket setup cleanup $ \h -> go h `onException` faile
 		liftIO $ Db.setRepoSizes h sizemap branchsha
 		calcJournalledRepoSizes sizemap branchsha
 
-	setup = Db.getRepoSizeHandle
-
-	cleanup _ = return ()
-
 	failed = do
 		liftIO $ putMVar rsv (Just M.empty)
 		return M.empty

Command/MaxSize.hs

@@ -12,7 +12,6 @@ module Command.MaxSize where
 import Command
 import qualified Remote
 import Annex.RepoSize
-import Types.RepoSize
 import Logs.MaxSize
 import Logs.Trust
 import Utility.DataUnits

Database/RepoSize.hs

@@ -25,7 +25,7 @@ module Database.RepoSize (
 	closeDb,
 	getRepoSizes,
 	setRepoSizes,
-	getLiveRepoSizes,
+	estimateLiveRepoSizes,
 	startingLiveSizeChange,
 	successfullyFinishedLiveSizeChange,
 	removeStaleLiveSizeChange,
@@ -291,10 +291,9 @@ getRecentChange u k = do
  - This is only expensive when there are a lot of live changes happening at
  - the same time.
  -}
-getLiveRepoSizes :: RepoSizeHandle -> IO (Maybe (M.Map UUID RepoSize, Sha))
+estimateLiveRepoSizes :: RepoSizeHandle -> IO (Maybe (M.Map UUID RepoSize, Sha))
-getLiveRepoSizes (RepoSizeHandle (Just h)) = H.queryDb h $ do
+estimateLiveRepoSizes (RepoSizeHandle (Just h)) = H.queryDb h $ do
 	getAnnexBranchCommit >>= \case
-		Nothing -> return Nothing
 		Just annexbranchsha -> do
 			sizechanges <- getSizeChanges
 			livechanges <- getLiveSizeChanges
@@ -302,6 +301,7 @@ getLiveRepoSizes (RepoSizeHandle (Just h)) = H.queryDb h $ do
 			m <- M.fromList <$> forM reposizes
 				(go sizechanges livechanges)
 			return (Just (m, annexbranchsha))
+		Nothing -> return Nothing
   where
 	go
 		:: M.Map UUID FileSize
@@ -350,4 +350,4 @@ getLiveRepoSizes (RepoSizeHandle (Just h)) = H.queryDb h $ do
 		filter (\(sc', cid') -> cid /= cid' && sc' == AddingKey)
 			(fromMaybe [] $ M.lookup k livechangesbykey)
 	competinglivechanges _ _ AddingKey _ = []
-getLiveRepoSizes (RepoSizeHandle Nothing) = return Nothing
+estimateLiveRepoSizes (RepoSizeHandle Nothing) = return Nothing

Limit.hs

@@ -599,9 +599,9 @@ limitFullyBalanced :: Maybe UUID -> Annex GroupMap -> MkLimit Annex
 limitFullyBalanced = limitFullyBalanced' "fullybalanced"
 
 limitFullyBalanced' :: String -> Maybe UUID -> Annex GroupMap -> MkLimit Annex
-limitFullyBalanced' = limitFullyBalanced'' $ \lu n key candidates -> do
+limitFullyBalanced' = limitFullyBalanced'' $ \n key candidates -> do
 	maxsizes <- getMaxSizes
-	sizemap <- getRepoSizes False
+	sizemap <- getLiveRepoSizes False
 	threshhold <- annexFullyBalancedThreshhold <$> Annex.getGitConfig
 	let toofull u =
 		case (M.lookup u maxsizes, M.lookup u sizemap) of
@@ -633,7 +633,7 @@ repoHasSpace keysize inrepo (RepoSize reposize) (MaxSize maxsize)
 		reposize + keysize <= maxsize
 
 limitFullyBalanced''
-	:: (LiveUpdate -> Int -> Key -> S.Set UUID -> Annex (S.Set UUID))
+	:: (Int -> Key -> S.Set UUID -> Annex (S.Set UUID))
 	-> String
 	-> Maybe UUID
 	-> Annex GroupMap
@@ -651,7 +651,7 @@ limitFullyBalanced'' filtercandidates termname mu getgroupmap want =
 		getgroupmap (toGroup s) n want
 
 limitFullyBalanced'''
-	:: (LiveUpdate -> Int -> Key -> S.Set UUID -> Annex (S.Set UUID))
+	:: (Int -> Key -> S.Set UUID -> Annex (S.Set UUID))
 	-> String
 	-> Maybe UUID
 	-> Annex GroupMap
@@ -663,8 +663,7 @@ limitFullyBalanced''' filtercandidates termname mu getgroupmap g n want = Right
 		gm <- getgroupmap
 		let groupmembers = fromMaybe S.empty $
 			M.lookup g (uuidsByGroup gm)
-		-- TODO locking for liveupdate
+		candidates <- filtercandidates n key groupmembers
-		candidates <- filtercandidates lu n key groupmembers
 		let wanted = if S.null candidates
 			then False
 			else case (mu, M.lookup g (balancedPickerByGroup gm)) of
@@ -690,9 +689,9 @@ limitFullySizeBalanced :: Maybe UUID -> Annex GroupMap -> MkLimit Annex
 limitFullySizeBalanced = limitFullySizeBalanced' "fullysizebalanced"
 
 limitFullySizeBalanced' :: String -> Maybe UUID -> Annex GroupMap -> MkLimit Annex
-limitFullySizeBalanced' = limitFullyBalanced'' $ \lu n key candidates -> do
+limitFullySizeBalanced' = limitFullyBalanced'' $ \n key candidates -> do
 	maxsizes <- getMaxSizes
-	sizemap <- getRepoSizes False
+	sizemap <- getLiveRepoSizes False
 	filterCandidatesFullySizeBalanced maxsizes sizemap n key candidates
 
 filterCandidatesFullySizeBalanced