{-# LANGUAGE ScopedTypeVariables #-}

-- | A bit like 'Fence', but not thread safe and optimised for avoiding taking the fence
module General.Thread(
    withThreadsBoth,
    withThreadSlave,
    allocateThread,
    Thread, newThreadFinally, stopThreads
    ) where

import General.Cleanup
import Data.Hashable
import Control.Concurrent.Extra
import Control.Exception
import General.Extra
import Control.Monad.Extra


data Thread = Thread ThreadId (Barrier ())

instance Eq Thread where
    Thread ThreadId
a Barrier ()
_ == :: Thread -> Thread -> Bool
== Thread ThreadId
b Barrier ()
_ = ThreadId
a forall a. Eq a => a -> a -> Bool
== ThreadId
b

instance Hashable Thread where
    hashWithSalt :: Int -> Thread -> Int
hashWithSalt Int
salt (Thread ThreadId
a Barrier ()
_) = forall a. Hashable a => Int -> a -> Int
hashWithSalt Int
salt ThreadId
a

-- | The inner thread is unmasked even if you started masked.
newThreadFinally :: IO a -> (Thread -> Either SomeException a -> IO ()) -> IO Thread
newThreadFinally :: forall a.
IO a -> (Thread -> Either SomeException a -> IO ()) -> IO Thread
newThreadFinally IO a
act Thread -> Either SomeException a -> IO ()
cleanup = do
    Barrier ()
bar <- forall a. IO (Barrier a)
newBarrier
    ThreadId
t <- forall a. IO a -> IO a
mask_ forall a b. (a -> b) -> a -> b
$ ((forall a. IO a -> IO a) -> IO ()) -> IO ThreadId
forkIOWithUnmask forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
unmask -> forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a b. IO a -> IO b -> IO a
finally (forall a. Partial => Barrier a -> a -> IO ()
signalBarrier Barrier ()
bar ()) forall a b. (a -> b) -> a -> b
$ do
        Either SomeException a
res <- forall e a. Exception e => IO a -> IO (Either e a)
try forall a b. (a -> b) -> a -> b
$ forall a. IO a -> IO a
unmask IO a
act
        ThreadId
me <- IO ThreadId
myThreadId
        Thread -> Either SomeException a -> IO ()
cleanup (ThreadId -> Barrier () -> Thread
Thread ThreadId
me Barrier ()
bar) Either SomeException a
res
    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ ThreadId -> Barrier () -> Thread
Thread ThreadId
t Barrier ()
bar


stopThreads :: [Thread] -> IO ()
stopThreads :: [Thread] -> IO ()
stopThreads [Thread]
threads = do
    -- if a thread is in a masked action, killing it may take some time, so kill them in parallel
    [Barrier ()]
bars <- forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence [do IO () -> IO ThreadId
forkIO forall a b. (a -> b) -> a -> b
$ ThreadId -> IO ()
killThread ThreadId
t; forall (f :: * -> *) a. Applicative f => a -> f a
pure Barrier ()
bar | Thread ThreadId
t Barrier ()
bar <- [Thread]
threads]
    forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ forall a. Barrier a -> IO a
waitBarrier [Barrier ()]
bars


-- Run both actions. If either throws an exception, both threads
-- are killed and an exception reraised.
-- Not called much, so simplicity over performance (2 threads).
withThreadsBoth :: IO a -> IO b -> IO (a, b)
withThreadsBoth :: forall a b. IO a -> IO b -> IO (a, b)
withThreadsBoth IO a
act1 IO b
act2 = do
    Barrier (Either SomeException a)
bar1 <- forall a. IO (Barrier a)
newBarrier
    Barrier (Either SomeException b)
bar2 <- forall a. IO (Barrier a)
newBarrier
    ThreadId
parent <- IO ThreadId
myThreadId
    Var Bool
ignore <- forall a. a -> IO (Var a)
newVar Bool
False
    forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
mask forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
unmask -> do
        ThreadId
t1 <- ((forall a. IO a -> IO a) -> IO ()) -> IO ThreadId
forkIOWithUnmask forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
unmask -> do
            Either SomeException a
res1 :: Either SomeException a <- forall e a. Exception e => IO a -> IO (Either e a)
try forall a b. (a -> b) -> a -> b
$ forall a. IO a -> IO a
unmask IO a
act1
            forall (m :: * -> *). Monad m => m Bool -> m () -> m ()
unlessM (forall a. Var a -> IO a
readVar Var Bool
ignore) forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a b.
Applicative m =>
Either a b -> (a -> m ()) -> m ()
whenLeft Either SomeException a
res1 forall a b. (a -> b) -> a -> b
$ forall e. Exception e => ThreadId -> e -> IO ()
throwTo ThreadId
parent
            forall a. Partial => Barrier a -> a -> IO ()
signalBarrier Barrier (Either SomeException a)
bar1 Either SomeException a
res1
        ThreadId
t2 <- ((forall a. IO a -> IO a) -> IO ()) -> IO ThreadId
forkIOWithUnmask forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
unmask -> do
            Either SomeException b
res2 :: Either SomeException b <- forall e a. Exception e => IO a -> IO (Either e a)
try forall a b. (a -> b) -> a -> b
$ forall a. IO a -> IO a
unmask IO b
act2
            forall (m :: * -> *). Monad m => m Bool -> m () -> m ()
unlessM (forall a. Var a -> IO a
readVar Var Bool
ignore) forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a b.
Applicative m =>
Either a b -> (a -> m ()) -> m ()
whenLeft Either SomeException b
res2 forall a b. (a -> b) -> a -> b
$ forall e. Exception e => ThreadId -> e -> IO ()
throwTo ThreadId
parent
            forall a. Partial => Barrier a -> a -> IO ()
signalBarrier Barrier (Either SomeException b)
bar2 Either SomeException b
res2
        Either SomeException (a, b)
res :: Either SomeException (a,b) <- forall e a. Exception e => IO a -> IO (Either e a)
try forall a b. (a -> b) -> a -> b
$ forall a. IO a -> IO a
unmask forall a b. (a -> b) -> a -> b
$ do
            Right a
v1 <- forall a. Barrier a -> IO a
waitBarrier Barrier (Either SomeException a)
bar1
            Right b
v2 <- forall a. Barrier a -> IO a
waitBarrier Barrier (Either SomeException b)
bar2
            forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
v1,b
v2)
        forall a. Var a -> a -> IO ()
writeVar Var Bool
ignore Bool
True
        ThreadId -> IO ()
killThread ThreadId
t1
        IO () -> IO ThreadId
forkIO forall a b. (a -> b) -> a -> b
$ ThreadId -> IO ()
killThread ThreadId
t2
        forall a. Barrier a -> IO a
waitBarrier Barrier (Either SomeException a)
bar1
        forall a. Barrier a -> IO a
waitBarrier Barrier (Either SomeException b)
bar2
        forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either forall e a. Exception e => e -> IO a
throwIO forall (f :: * -> *) a. Applicative f => a -> f a
pure Either SomeException (a, b)
res


-- | Run an action in a separate thread.
--   After the first action terminates, the thread will be killed.
--   If the action raises an exception it will be rethrown on the parent thread.
withThreadSlave :: IO () -> IO a -> IO a
withThreadSlave :: forall a. IO () -> IO a -> IO a
withThreadSlave IO ()
slave IO a
act = forall a. (Cleanup -> IO a) -> IO a
withCleanup forall a b. (a -> b) -> a -> b
$ \Cleanup
cleanup -> do
    Cleanup -> IO () -> IO ()
allocateThread Cleanup
cleanup IO ()
slave
    IO a
act


-- | Run the given action in a separate thread.
--   On cleanup, the thread will be killed before continuing.
--   If the action raises an exception it will be rethrown on the parent thread.
allocateThread :: Cleanup -> IO () -> IO ()
allocateThread :: Cleanup -> IO () -> IO ()
allocateThread Cleanup
cleanup IO ()
act = do
    Barrier ()
bar <- forall a. IO (Barrier a)
newBarrier
    ThreadId
parent <- IO ThreadId
myThreadId
    Var Bool
ignore <- forall a. a -> IO (Var a)
newVar Bool
False
    forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ forall a. Cleanup -> IO a -> (a -> IO ()) -> IO a
allocate Cleanup
cleanup
        (forall a. IO a -> IO a
mask_ forall a b. (a -> b) -> a -> b
$ ((forall a. IO a -> IO a) -> IO ()) -> IO ThreadId
forkIOWithUnmask forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
unmask -> do
            Either SomeException ()
res :: Either SomeException () <- forall e a. Exception e => IO a -> IO (Either e a)
try forall a b. (a -> b) -> a -> b
$ forall a. IO a -> IO a
unmask IO ()
act
            forall (m :: * -> *). Monad m => m Bool -> m () -> m ()
unlessM (forall a. Var a -> IO a
readVar Var Bool
ignore) forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a b.
Applicative m =>
Either a b -> (a -> m ()) -> m ()
whenLeft Either SomeException ()
res forall a b. (a -> b) -> a -> b
$ forall e. Exception e => ThreadId -> e -> IO ()
throwTo ThreadId
parent
            forall a. Partial => Barrier a -> a -> IO ()
signalBarrier Barrier ()
bar ()
        )
        (\ThreadId
t -> do forall a. Var a -> a -> IO ()
writeVar Var Bool
ignore Bool
True; ThreadId -> IO ()
killThread ThreadId
t; forall a. Barrier a -> IO a
waitBarrier Barrier ()
bar)