git.iko.soy/distributed-process
on 87bb129 help sign in
distributed-process/src/Control/Distributed/Process/Internal/ThreadPool.hs
78 lines 2.9 kB history raw
1-- | An implementation of a pool of threads
2--
3{-# LANGUAGE RecursiveDo #-}
4module Control.Distributed.Process.Internal.ThreadPool
5 ( newThreadPool
6 , submitTask
7 , lookupWorker
8 , ThreadPool
9 ) where
11import Control.Exception
12import Control.Monad
13import Data.IORef
14import qualified Data.Map as Map
17-- | A pool of worker threads that execute tasks.
18--
19-- Each worker thread is named with a key @k@. Tasks are submitted to a
20-- specific worker using its key. While the worker is busy the tasks are queued.
21-- When there are no more queued tasks the worker ceases to exist.
22--
23-- The next time a task is submitted the worker will be respawned.
24--
25newtype ThreadPool k w = ThreadPool (IORef (Map.Map k (Maybe (IO ()), w)))
27-- Each worker has an entry in the map with a closure that contains all
28-- queued actions fot it.
29--
30-- No entry in the map is kept for defunct workers.
32-- | Creates a pool with no workers.
33newThreadPool :: IO (ThreadPool k w)
34newThreadPool = fmap ThreadPool $ newIORef Map.empty
36-- | @submitTask pool fork k task@ submits a task for the worker @k@.
37--
38-- If worker @k@ is busy, then the task is queued until the worker is available.
39--
40-- If worker @k@ does not exist, then the given @fork@ operation is used to
41-- spawn the worker. @fork@ returns whatever information is deemed useful for
42-- later retrieval via 'lookupWorker'.
43--
44submitTask :: Ord k
45 => ThreadPool k w
46 -> (IO () -> IO w)
47 -> k -> IO () -> IO ()
48submitTask (ThreadPool mapRef) fork k task = mdo
49 m' <- join $ atomicModifyIORef mapRef $ \m ->
50 case Map.lookup k m of
51 -- There is no worker for this key, create one.
52 Nothing -> ( m'
53 , do w <- fork $ flip onException terminateWorker $ do
54 task
55 continue
56 return $ Map.insert k (Nothing, w) m
57 )
58 -- Queue an action for the existing worker.
59 Just (mp, w) ->
60 (m', return $ Map.insert k (Just $ maybe task (>> task) mp, w) m)
61 return ()
62 where
63 continue = join $ atomicModifyIORef mapRef $ \m ->
64 case Map.lookup k m of
65 -- Execute the next batch of queued actions.
66 Just (Just p, w) -> (Map.insert k (Nothing, w) m, p >> continue)
67 -- There are no more queued actions. Terminate the worker.
68 Just (Nothing, w) -> (Map.delete k m, return ())
69 -- The worker key was removed already (?)
70 Nothing -> (m, return ())
71 -- Remove the worker key regardless of whether there are more queued
72 -- actions.
73 terminateWorker = atomicModifyIORef mapRef $ \m -> (Map.delete k m, ())
75-- | Looks up a worker with the given key.
76lookupWorker :: Ord k => ThreadPool k w -> k -> IO (Maybe w)
77lookupWorker (ThreadPool mapRef) k =
78 atomicModifyIORef mapRef $ \m -> (m, fmap snd $ Map.lookup k m)