git-annex-10.20261113: Types/WorkerPool.hs

{- Worker thread pool.
 -
 - Copyright 2019-2026 Joey Hess <id@joeyh.name>
 -
 - Licensed under the GNU AGPL version 3 or higher.
 -}

{-# LANGUAGE DeriveFunctor #-}

module Types.WorkerPool where

import Types.Direction

import Control.Concurrent
import Control.Concurrent.Async
import qualified Data.Set as S
import Data.Unique
import Control.Monad

-- | Pool of worker threads. 
data WorkerPool t = WorkerPool
	{ usedStages :: UsedStages
	, workerList :: [Worker t]
	, spareVals :: [t]
	-- ^ Normally there is one value for each IdleWorker,
	-- but there can temporarily be fewer values, when a thread is
	-- changing between stages.
	} 

instance Show (WorkerPool t) where
	show p = unwords
		[ "WorkerPool"
		, show (usedStages p)
		, show (workerList p)
		, show (length (spareVals p))
		]

-- | A worker can either be idle or running an Async action.
-- And it is used for some stage.
data Worker t
	= IdleWorker WorkerStage
	| ActiveWorker (Async t) WorkerStage
	deriving (Functor)

instance Show (Worker t) where
	show (IdleWorker s) = "IdleWorker " ++ show s
	show (ActiveWorker _ s) = "ActiveWorker " ++ show s

data WorkerStage
	= StartStage
	-- ^ All threads start in this stage, and then transition away from
	-- it to the initialStage when they begin doing work. This should
	-- never be included in UsedStages, because transition from some
	-- other stage back to this one could result in a deadlock.
	| PerformStage
	-- ^ Running a CommandPerform action.
	| CleanupStage
	-- ^ Running a CommandCleanup action.
	| TransferStage Direction
	-- ^ Transferring content to or from a remote.
	| VerifyStage
	-- ^ Verifying content, eg by calculating a checksum.
	deriving (Show, Eq, Ord)

-- | Set of stages that make sense to be used while performing an action,
-- and the stage to use initially.
-- 
-- Transitions between these stages will block as needed until there's a
-- free Worker in the pool for the new stage.
-- 
-- Actions that indicate they are in some other stage won't change the
-- stage, and so there will be no blocking before starting them.
data UsedStages = UsedStages
	{ initialStage :: WorkerStage
	, stageSet :: S.Set WorkerStage
	}
	deriving (Show)

memberStage :: WorkerStage -> UsedStages -> Bool
memberStage s u = S.member s (stageSet u)

-- | The default is to use only the CommandPerform and CommandCleanup
-- stages. Since cleanup actions often don't contend much with
-- perform actions, this prevents blocking starting the next perform action
-- on finishing the previous cleanup action.
commandStages :: UsedStages
commandStages = UsedStages
	{ initialStage = PerformStage
	, stageSet = S.fromList [PerformStage, CleanupStage]
	}

-- | This is mostly useful for downloads, not for uploads. A download
-- is often bottlenecked on the network or another disk than the one
-- containing the repository. When verification is not done incrementally,
-- it bottlenecks on the disk containing the repository or on the CPU.
-- So it makes sense to run the download and verify stages separately.
-- 
-- For uploads, there is no separate verify step to this is less likely
-- to be useful than commandStages. However, a separate stage is provided
-- for Uploads. That can be useful when a command downloads from one remote
-- (eg using the network) and uploads to another remote (eg using a disk).
transferStages :: UsedStages
transferStages = UsedStages
	{ initialStage = TransferStage Download
	, stageSet = S.fromList
		[ TransferStage Download
		, TransferStage Upload
		, VerifyStage
		]
	}

workerStage :: Worker t -> WorkerStage
workerStage (IdleWorker s) = s
workerStage (ActiveWorker _ s) = s

workerAsync :: Worker t -> Maybe (Async t)
workerAsync (IdleWorker _) = Nothing
workerAsync (ActiveWorker aid _) = Just aid

-- | Allocates a WorkerPool that has the specified number of workers
-- in it, of each stage.
allocateWorkerPool :: t -> Int -> UsedStages -> WorkerPool t
allocateWorkerPool t n u =
	let stages = workerPoolStages u
	in allocateWorkerPool' u
		(replicate (workerPoolSize n stages) t)
		stages

allocateWorkerPoolM :: IO t -> Int -> UsedStages -> IO (WorkerPool t)
allocateWorkerPoolM mkt n u = do
	let stages = workerPoolStages u
	ts <- replicateM (workerPoolSize n stages) mkt
	return $ allocateWorkerPool' u ts stages

allocateWorkerPool' :: UsedStages -> [t] -> [WorkerStage] -> WorkerPool t
allocateWorkerPool' u ts stages = WorkerPool
	{ usedStages = u
	, workerList = map IdleWorker $
		take (length ts) $ concat $ repeat stages
	, spareVals = ts
	}

workerPoolSize :: Int -> [WorkerStage] -> Int
workerPoolSize n stages = n * length stages

workerPoolStages :: UsedStages -> [WorkerStage]
workerPoolStages u = StartStage : S.toList (stageSet u)

addWorkerPool :: Worker t -> WorkerPool t -> WorkerPool t
addWorkerPool w pool = pool { workerList = w : workerList pool }

-- | Removes a worker from the pool whose Async uses the ThreadId.
--
-- Each Async has its own ThreadId, so this stops once it finds
-- a match.
removeThreadIdWorkerPool :: ThreadId -> WorkerPool t -> Maybe ((Async t, WorkerStage), WorkerPool t)
removeThreadIdWorkerPool tid pool = go [] (workerList pool)
  where
	go _ [] = Nothing
	go c (ActiveWorker a stage : rest)
		| asyncThreadId a == tid =
			let pool' = pool { workerList = (c++rest) }
			in Just ((a, stage), pool')
	go c (v : rest) = go (v:c) rest

deactivateWorker :: WorkerPool t -> Async t -> t -> WorkerPool t
deactivateWorker pool aid t = pool
	{ workerList = go (workerList pool)
	, spareVals = t : spareVals pool
	}
  where
	go [] = []
	go (w@(IdleWorker _) : rest) = w : go rest
	go (w@(ActiveWorker a st) : rest)
		| a == aid = IdleWorker st : rest
		| otherwise = w : go rest

allIdle :: WorkerPool t -> Bool
allIdle pool = all idle (workerList pool)
	-- If this does not hold, a thread must be transitioning between
	-- states, so it's not really idle.
	&& length (spareVals pool) == length (workerList pool)
  where
	idle (IdleWorker _) = True
	idle (ActiveWorker _ _) = False

{- Tags each value in a WorkerPool with a Unique identifier. -}
taggedWorkerPool :: WorkerPool t -> IO (WorkerPool (t, Unique))
taggedWorkerPool wp = do
	wl <- forM (workerList wp) $ \w -> do
		u <- newUnique
		return (fmap (\v -> (v, u)) w)
	sl <- zip (spareVals wp)
		<$> replicateM (length (spareVals wp)) newUnique
	return $ WorkerPool
		{ usedStages = usedStages wp
		, workerList = wl
		, spareVals = sl
		}