git.iko.soy/streamly
on 7beefce help sign in
streamly/src/Asyncly/Threads.hs
449 lines 17 kB history raw
1{-# LANGUAGE ConstraintKinds #-}
2{-# LANGUAGE FlexibleContexts #-}
3{-# LANGUAGE RecordWildCards #-}
5-- |
6-- Module : Asyncly.Threads
7-- Copyright : (c) 2017 Harendra Kumar
8--
9-- License : MIT-style
10-- Maintainer : harendra.kumar@gmail.com
11-- Stability : experimental
12-- Portability : GHC
13--
14module Asyncly.Threads
15 ( -- ChildEvent(..)
16 MonadAsync
17 , Context(..)
18 , initContext
19 --, runAsyncTask
20 --, makeCont
21 , threadCtl
22 -- , Location(..)
23 -- , getLocation
24 -- , setLocation
26 -- , waitForChildren
27 --, handleResult
28 )
29where
31import Control.Concurrent (ThreadId, forkIO, killThread,
32 myThreadId)
33import Control.Concurrent.STM (TChan, atomically, newTChan,
34 readTChan, tryReadTChan,
35 writeTChan)
36import Control.Exception (fromException,
37 SomeException (..))
38import qualified Control.Exception.Lifted as EL
39import Control.Monad (when)
40import Control.Monad.Catch (MonadThrow, throwM)
41import Control.Monad.State (MonadIO (..),
42 StateT (..), runStateT,
43 get, gets, modify)
44import Control.Monad.Trans.Class (MonadTrans (lift))
45import Control.Monad.Trans.Control (MonadBaseControl, liftBaseWith)
46import Data.Atomics (atomicModifyIORefCAS)
47import Data.IORef (IORef, modifyIORef, newIORef,
48 readIORef, writeIORef)
49import Data.List (delete)
50import Data.Maybe (isJust)
51import GHC.Prim (Any)
52import Unsafe.Coerce (unsafeCoerce)
54import Control.Monad.Trans.Recorder (Paused (..), Recording)
56type MonadAsync m = (MonadIO m, MonadBaseControl IO m, MonadThrow m)
58------------------------------------------------------------------------------
59-- Parent child thread communication types
60------------------------------------------------------------------------------
62data ChildEvent = ChildDone ThreadId (Maybe SomeException)
64------------------------------------------------------------------------------
65-- State threaded around the monad
66------------------------------------------------------------------------------
68data Context = Context
69 {
70 ---------------------------------------------------------------------------
71 -- Execution state
72 ---------------------------------------------------------------------------
74 -- When we suspend we save the logs in this IORef and exit.
75 logsRef :: Maybe (IORef [Recording])
76 -- XXX this functionality can be in a separate layer?
77 , location :: Location
79 ---------------------------------------------------------------------------
80 -- Thread creation, communication and cleanup
81 ---------------------------------------------------------------------------
82 -- When a new thread is created the parent records it in the
83 -- 'pendingThreads' field and 'threadCredit' is decremented. When a child
84 -- thread is done it sends a done event to the parent on an unbounded
85 -- channel and goes away. Before starting a new thread the parent always
86 -- processes the unbounded channel to clear up the pending zombies. This
87 -- strategy ensures that the zombie queue will never grow more than the
88 -- number of running threads. Before exiting, the parent thread waits on
89 -- the channel until all its children are cleaned up.
90 ---------------------------------------------------------------------------
92 -- XXX setup a cleanup computation to run rather than passing all these
93 -- params.
95 , childChannel :: TChan ChildEvent
96 -- ^ A channel for the immediate children to communicate to us when they
97 -- die. Each thread has its own dedicated channel for its children
99 -- We always track the child threads, otherwise the programmer needs to
100 -- worry about if some threads may remain hanging because they are stuck in
101 -- an infinite loop or a forever blocked IO. Also, it is not clear whether
102 -- there is a significant gain by disabling tracking. Instead of using
103 -- threadIds we can also keep a count instead, that will allow us to wait
104 -- for all children to drain but we cannot kill them.
105 --
106 -- We need these as IORefs since we need any changes to these to be
107 -- reflected back in the calling code that waits for the children.
109 , pendingThreads :: IORef [ThreadId]
110 -- ^ Active immediate child threads spawned by this thread, modified only
111 -- by this thread, therefore atomic modification is not needed.
113 -- By default this limit is shared by the current tree of threads. But at
114 -- any point the 'threads' primitive can be used to start a new limit for
115 -- the computation under that primitive.
116 --
117 -- Shared across many threads, therefore atomic access is required.
118 -- This is incremented by the child thread itelf before it exits.
119 --
120 , threadCredit :: IORef Int
121 -- ^ How many more threads are allowed to be created?
122 }
124initContext
125 :: TChan ChildEvent
126 -> IORef [ThreadId]
127 -> IORef Int
128 -> Maybe (IORef [Recording])
129 -> Context
130initContext childChan pending credit lref =
131 Context { logsRef = lref
132 , location = Worker
133 , childChannel = childChan
134 , pendingThreads = pending
135 , threadCredit = credit }
137------------------------------------------------------------------------------
138-- Where is the computation running?
139------------------------------------------------------------------------------
141-- RemoteNode is used in the Alternative instance. We stop the computation when
142-- the location is RemoteNode, we do not execute any alternatives.
143--
144data Location = Worker | WaitingParent | RemoteNode
145 deriving (Eq, Show)
147getLocation :: Monad m => StateT Context m Location
148getLocation = gets location
150setLocation :: Monad m => Location -> StateT Context m ()
151setLocation loc = modify $ \Context { .. } -> Context { location = loc, .. }
153{-
154------------------------------------------------------------------------------
155-- Pick up from where we left in the previous thread
156------------------------------------------------------------------------------
158{-# SPECIALIZE continue :: [Any] -> a -> StateT Context IO (Maybe a) #-}
159continue :: Monad m => [Any] -> a -> StateT Context m (Maybe a)
160continue [] x = do
161 return (Just x)
163continue [f] x = do
164 modify $ \ctx -> ctx {continuations = []}
165 (unsafeCoerce f) x
167continue (f:fs) x = do
168 modify $ \ctx -> ctx {continuations = fs}
169 y <- (unsafeCoerce f) x
170 case y of
171 Nothing -> return y
172 Just z -> continue fs z
174{-# SPECIALIZE runContext :: Context -> StateT Context IO (Maybe a) -> StateT Context IO () #-}
175runContext :: MonadAsync m => Context -> StateT Context m (Maybe a) -> StateT Context m ()
176runContext ctx action = do
177 _ <- lift $ runStateT m ctx
178 return ()
180 where
182 m = do
183 let conts = continuations ctx
184 x <- action
185 case x of
186 Just y -> continue conts y >> return ()
187 Nothing -> return ()
189------------------------------------------------------------------------------
190-- Thread Management (creation, reaping and killing)
191------------------------------------------------------------------------------
193-- XXX We are using unbounded channels so this will not block on writing to
194-- pchan. We can use bounded channels to throttle the creation of threads based
195-- on consumption rate.
196processOneEvent :: MonadIO m
197 => ChildEvent
198 -> [ThreadId]
199 -> m ([ThreadId], Maybe SomeException)
200processOneEvent (ChildDone tid e) pending = do
201 when (isJust e) $ liftIO $ mapM_ killThread pending
202 return (delete tid pending, Nothing)
204drainChildren :: MonadIO m
205 => TChan ChildEvent
206 -> [ThreadId]
207 -> m ([ThreadId], Maybe SomeException)
208drainChildren cchan pending =
209 case pending of
210 [] -> return (pending, Nothing)
211 _ -> do
212 ev <- liftIO $ atomically $ readTChan cchan
213 (p, e) <- processOneEvent ev pending
214 maybe (drainChildren cchan p) (const $ return (p, e)) e
216waitForChildren :: MonadIO m => Context -> m (Maybe SomeException)
217waitForChildren ctx = do
218 let pendingRef = pendingThreads ctx
219 pending <- liftIO $ readIORef pendingRef
220 (p, e) <- drainChildren (childChannel ctx) pending
221 liftIO $ writeIORef pendingRef p
222 return e
224tryReclaimZombies :: (MonadIO m, MonadThrow m) => Context -> m ()
225tryReclaimZombies ctx = do
226 let cchan = childChannel ctx
227 pendingRef = pendingThreads ctx
229 pending <- liftIO $ readIORef pendingRef
230 case pending of
231 [] -> return ()
232 _ -> do
233 mev <- liftIO $ atomically $ tryReadTChan cchan
234 case mev of
235 Nothing -> return ()
236 Just ev -> do
237 (p, e) <- processOneEvent ev pending
238 liftIO $ writeIORef pendingRef p
239 maybe (return ()) throwM e
240 tryReclaimZombies ctx
242{-# NOINLINE waitForOneEvent #-}
243waitForOneEvent :: (MonadIO m, MonadThrow m) => Context -> m ()
244waitForOneEvent ctx = do
245 -- XXX assert pending must have at least one element
246 -- assert that the tid is found in our list
247 let cchan = childChannel ctx
248 pendingRef = pendingThreads ctx
250 ev <- liftIO $ atomically $ readTChan cchan
251 pending <- liftIO $ readIORef pendingRef
252 (p, e) <- processOneEvent ev pending
253 liftIO $ writeIORef pendingRef p
254 maybe (return ()) throwM e
256-- XXX this is not a real semaphore as it does not really block on wait,
257-- instead it returns whether the value is zero or non-zero.
258--
259{-# INLINE waitQSemB #-}
260waitQSemB :: IORef Int -> IO Bool
261waitQSemB sem = atomicModifyIORefCAS sem $ \n ->
262 if n > 0
263 then (n - 1, True)
264 else (n, False)
266signalQSemB :: IORef Int -> IO ()
267signalQSemB sem = atomicModifyIORefCAS sem $ \n -> (n + 1, ())
269-- Allocation of threads
270--
271-- global thread limit
272-- thread fan-out i.e. per thread children limit
273-- min per thread allocation to avoid starvation
274--
275-- dynamic adjustment based on the cost, speed of consumption, cpu utilization
276-- etc. We need to adjust the limits based on cost, throughput and latencies.
277--
278-- The event producer thread must put the work on a work-queue and the child
279-- threads can pick it up from there. But if there is just one consumer then it
280-- may not make sense to have a separate producer unless the producing cost is
281-- high.
282--
284{-# INLINE handleResult #-}
285handleResult :: MonadIO m
286 => Context -> Either SomeException a -> m (Maybe SomeException)
287handleResult ctx res =
288 case res of
289 Left e -> do
290 case fromException e of
291 Just (Paused recording) ->
292 maybe (passException e)
293 (handleRecording recording) (logsRef ctx)
294 Nothing -> passException e
295 Right _ -> waitForChildren ctx
297 where
299 handleRecording recording ref = do
300 liftIO $ atomicModifyIORefCAS ref $ \logs -> (recording : logs, ())
301 waitForChildren ctx
303 passException e = do
304 liftIO $ readIORef (pendingThreads ctx) >>= mapM_ killThread
305 return (Just e)
307forkFinally1 :: MonadAsync m
308 => Context
309 -> StateT Context m (Maybe a)
310 -> (Either SomeException () -> IO ())
311 -> StateT Context m ThreadId
312forkFinally1 ctx action preExit =
313 EL.mask $ \restore ->
314 liftBaseWith $ \runInIO -> forkIO $ do
315 _ <- runInIO $ EL.try (restore (runContext ctx action))
316 >>= liftIO . preExit
317 -- XXX restore state here
318 return ()
320-- | Run a given context in a new thread.
321--
322forkContext :: (MonadBaseControl IO m, MonadIO m, MonadThrow m)
323 => StateT Context m (Maybe a) -> StateT Context m ()
324forkContext action = do
325 parentCtx <- get
326 childCtx <- childContext parentCtx
327 tid <- forkFinally1 childCtx action
328 (beforeExit childCtx (childChannel parentCtx))
329 updatePendingThreads parentCtx tid
331 where
333 updatePendingThreads :: (MonadIO m, MonadThrow m)
334 => Context -> ThreadId -> m ()
335 updatePendingThreads ctx tid = do
336 -- update the new thread before reclaiming zombies so that if it exited
337 -- already reclaim finds it in the list and does not panic.
338 liftIO $ modifyIORef (pendingThreads ctx) $ (\ts -> tid:ts)
339 tryReclaimZombies ctx
341 childContext ctx = do
342 pendingRef <- liftIO $ newIORef []
343 chan <- liftIO $ atomically newTChan
344 -- shares the threadCredit of the parent by default
345 return $ ctx
346 { pendingThreads = pendingRef
347 , childChannel = chan
348 }
350 beforeExit ctx pchan res = do
351 r <- handleResult ctx res
352 signalQSemB (threadCredit ctx)
353 tid <- myThreadId
354 liftIO $ atomically $ writeTChan pchan (ChildDone tid r)
356-- | Decide whether to resume the context in the same thread or a new thread
357--
358canFork :: Context -> IO Bool
359canFork context = do
360 gotCredit <- waitQSemB (threadCredit context)
361 case gotCredit of
362 False -> do
363 pending <- readIORef $ pendingThreads context
364 case pending of
365 [] -> return False
366 _ -> do
367 -- XXX If we have unreclaimable child threads e.g.
368 -- infinite loop, this is going to deadlock us. We need
369 -- special handling for those cases. Add those to
370 -- unreclaimable list? And always execute them in an
371 -- async thread, cannot use sync for those.
372 --
373 waitForOneEvent context
374 canFork context
375 True -> return True
377-- Housekeeping, invoked after spawning of all child tasks is done and the
378-- parent task needs to terminate. Either the task is fully done or we handed
379-- it over to another thread, in any case the current thread is done.
381spawningParentDone :: MonadIO m => StateT Context m (Maybe a)
382spawningParentDone = do
383 loc <- getLocation
384 when (loc /= RemoteNode) $ setLocation WaitingParent
385 return Nothing
387-- | The task may be run in the same thread or in a new thread depending on the
388-- forceAsync parameter and the current thread quota.
389--
390{-# SPECIALIZE runAsyncTask :: Bool -> StateT Context IO (Maybe a)
391 -> StateT Context IO (Maybe a) #-}
392{-# INLINE runAsyncTask #-}
393runAsyncTask :: MonadAsync m
394 => Bool
395 -> StateT Context m (Maybe a)
396 -> StateT Context m (Maybe a)
397runAsyncTask forceAsync action = do
398 if forceAsync
399 then forkContext action
400 else do
401 context <- get
402 can <- liftIO $ canFork context
403 case can of
404 False -> runContext context action
405 True -> forkContext action
406 return Nothing
408-- XXX After setting the callback we should wait for the callback otherwise we
409-- will just go away. We should use the same finalization mechanism that we use
410-- in fork.
411--
412-- | Makes an asyncly action from a callback setter function; can be used to
413-- convert existing callback style code into asyncly style code. The first
414-- parameter is a callback setter function. When we are called back,
415-- 'makeAsync' behaves as if it was an async computation that just returned a
416-- value of type 'a'. After the computation is done the result of the action
417-- in second parameter is returned to the callback.
418--
419makeCont :: MonadAsync m => ((a -> m ()) -> m ()) -> StateT Context m (Maybe a)
420makeCont cbsetter = do
421 -- XXX should we force fork a thread or keep running in the context of the
422 -- callback?
423 ctx <- get
424 lift $ cbsetter $ \a -> do
425 let s = runAsyncTask False (return (Just a))
426 _ <- runStateT s ctx
427 return ()
428 spawningParentDone
430-}
431------------------------------------------------------------------------------
432-- Controlling thread quota
433------------------------------------------------------------------------------
435-- XXX Should n be Word32 instead?
436-- | Runs a computation under a given thread limit. A limit of 0 means all new
437-- tasks start synchronously in the current thread unless overridden by
438-- 'async'.
439threadCtl :: MonadAsync m
440 => Int
441 -> StateT Context m a
442 -> StateT Context m a
443threadCtl n action = do
444 oldCr <- gets threadCredit
445 newCr <- liftIO $ newIORef n
446 modify $ \s -> s { threadCredit = newCr }
447 r <- action
448 modify $ \s -> s { threadCredit = oldCr }
449 return r