@@ -78,14 +78,14 @@ import Control.Monad.IO.Class (MonadIO(..))
7878import Control.Monad.Reader.Class (MonadReader(..))
7979import Control.Monad.Trans.Class (MonadTrans (lift))
8080import Control.Monad.Trans.Control (MonadBaseControl, liftBaseWith)
81−import Data.Atomics (atomicModifyIORefCAS,
82− atomicModifyIORefCAS_)
81+import Data.Atomics (casIORef, readForCAS, peekTicket
82+ ,atomicModifyIORefCAS_)
8383import Data.Concurrent.Queue.MichaelScott (LinkedQueue, newQ, pushL,
8484 tryPopR, nullQ)
8585import Data.Functor (void)
8686import Data.IORef (IORef, modifyIORef, newIORef,
87− readIORef)
88−import Data.Maybe (isNothing)
87+ readIORef, atomicModifyIORef)
88+import Data.Maybe (isNothing, fromJust)
8989import Data.Semigroup (Semigroup(..))
9090import Data.Set (Set)
9191import qualified Data.Set as S
@@ -158,6 +158,25 @@ data SVar m a =
158158 , svarStyle :: SVarStyle
159159 }
160160
161+-- Slightly faster version of CAS. Gained some improvement by avoiding the use
162+-- of "evaluate" because we know we do not have exceptions in fn.
163+{-# INLINE atomicModifyIORefCAS #-}
164+atomicModifyIORefCAS :: IORef a -> (a -> (a,b)) -> IO b
165+atomicModifyIORefCAS ref fn = do
166+ tkt <- readForCAS ref
167+ loop tkt retries
168+
169+ where
170+
171+ retries = 30 :: Int
172+ loop _ 0 = atomicModifyIORef ref fn
173+ loop old tries = do
174+ let (new, result) = fn $ peekTicket old
175+ (success, tkt) <- casIORef ref old new
176+ if success
177+ then return result
178+ else loop tkt (tries - 1)
179+
161180------------------------------------------------------------------------------
162181-- The stream type
163182------------------------------------------------------------------------------
@@ -376,14 +395,16 @@ delThread sv tid =
376395
377396-- If present then delete else add. This takes care of out of order add and
378397-- delete i.e. a delete arriving before we even added a thread.
398+-- This occurs when the forked thread is done even before the 'addThread' right
399+-- after the fork is executed.
379400{-# INLINE modifyThread #-}
380401modifyThread :: MonadIO m => SVar m a -> ThreadId -> m ()
381402modifyThread sv tid = do
382403 changed <- liftIO $ atomicModifyIORefCAS (runningThreads sv) $ \old ->
383404 if (S.member tid old)
384− then let new = (S.delete tid old) in (new, null new)
385− else let new = (S.insert tid old) in (new, null old)
386− if changed
405+ then let new = (S.delete tid old) in (new, new)
406+ else let new = (S.insert tid old) in (new, old)
407+ if null changed
387408 then liftIO $ void $ tryPutMVar (doorBell sv) True
388409 else return ()
389410
@@ -631,13 +652,6 @@ toStreamVar sv m = do
631652-- TBD Note 4: for pure work (when we are not in the IO monad) we can divide it
632653-- into just the number of CPUs.
633654
634−{-# NOINLINE withNewSVar2 #-}
635−withNewSVar2 :: MonadParallel m
636− => SVarStyle -> Stream m a -> Stream m a -> Stream m a
637−withNewSVar2 style m1 m2 = Stream $ \_ stp sng yld -> do
638− sv <- newStreamVar2 style m1 m2
639− (runStream (fromStreamVar sv)) Nothing stp sng yld
640−
641655-- | Join two computations on the currently running 'SVar' queue for concurrent
642656-- execution. When we are using parallel composition, an SVar is passed around
643657-- as a state variable. We try to schedule a new parallel computation on the
@@ -671,7 +685,13 @@ joinStreamVar2 style m1 m2 = Stream $ \svr stp sng yld ->
671685 case svr of
672686 Just sv | svarStyle sv == style ->
673687 liftIO ((enqueue sv) m2) >> (runStream m1) svr stp sng yld
674− _ -> (runStream (withNewSVar2 style m1 m2)) Nothing stp sng yld
688+ _ -> do
689+ sv <- newStreamVar1 style (concurrently m1 m2)
690+ (runStream (fromStreamVar sv)) Nothing stp sng yld
691+ where
692+ concurrently ma mb = Stream $ \svr stp sng yld -> do
693+ liftIO $ (enqueue (fromJust svr)) mb
694+ (runStream ma) svr stp sng yld
675695
676696{-# INLINE joinStreamVarPar #-}
677697joinStreamVarPar :: MonadParallel m