diff --git a/.gitignore b/.gitignore new file mode 100644 --- /dev/null +++ b/.gitignore @@ -0,0 +1,1 @@ +dist diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,19 @@ +Copyright (c) 2016 Charles Strahan + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/bench/Main.hs b/bench/Main.hs new file mode 100644 --- /dev/null +++ b/bench/Main.hs @@ -0,0 +1,135 @@ +{-# LANGUAGE BangPatterns #-} + +module Main where + +import Control.Concurrent hiding (Chan) +import Control.Concurrent.GoChan +import Control.Monad +import Criterion.Main +import System.Random + +for xs f = map f xs + +main :: IO () +main = + defaultMain $ + (for + [0, 2, 5] + (\size -> + bgroup + ("buffer size " ++ show size) + [ bench "create" $ whnfIO (chanMake size :: IO (Chan Int)) + , bench "send & recv" $ whnfIO (sendRecv size) + , bench "select send & recv with 1 case" $ + whnfIO (select1 size) + , bench "select send & recv with 2 case" $ + whnfIO (select2 size) + , bench "select send & recv with 3 case" $ + whnfIO (select3 size)])) ++ + [bench "no-threads send & recv" $ whnfIO sendRecvBuf] + +{-# INLINE iters #-} + +iters = 15 + +sendRecv :: Int -> IO () +sendRecv !size = do + c <- chanMake size + forkIO $ + do sendN c 0 + chanClose c + drain c + where + drain !ch = do + mn <- chanRecv ch + case mn of + Msg n -> do + drain ch + _ -> return () + sendN !ch !n = do + chanSend ch n + when (n < iters) (sendN ch (n + 1)) + +sendRecvBuf :: IO () +sendRecvBuf = do + c <- chanMake 1 + sendRecN c + where + sendRecN !ch = do + chanSend ch 42 + void $ chanRecv ch + +select1 + :: Int -> IO () +select1 !size = do + c1 <- chanMake size + forkIO $ ping c1 0 + pong c1 0 + where + ping !c1 !n = do + if (n < iters) + then do + chanSelect [Send c1 n (return ())] Nothing + ping c1 (n + 1) + else return () + pong !c1 !n = do + if (n < iters) + then do + chanSelect [Recv c1 (const (return ()))] Nothing + pong c1 (n + 1) + else return () + +select2 :: Int -> IO () +select2 !size = do + --lock <- newEmptyMVar + c1 <- chanMake size + c2 <- chanMake size + forkIO $ ping c1 c2 0 + pong c1 c2 0 + where + ping :: Chan Int -> Chan Int -> Int -> IO () + ping !c1 !c2 !n = do + if (n < iters) + then do + chanSelect [Send c1 n (return ()), Send c2 n (return ())] Nothing + ping c1 c2 (n + 1) + else return () + pong :: Chan Int -> Chan Int -> Int -> IO () + pong !c1 !c2 !n = do + if (n < iters) + then do + chanSelect + [Recv c1 (const (return ())), Recv c2 (const (return ()))] + Nothing + pong c1 c2 (n + 1) + else return () + +select3 :: Int -> IO () +select3 !size = do + --lock <- newEmptyMVar + c1 <- chanMake size + c2 <- chanMake size + c3 <- chanMake size + forkIO $ ping c1 c2 c3 0 + pong c1 c2 c3 0 + where + ping !c1 !c2 !c3 !n = do + if (n < iters) + then do + chanSelect + [ Send c1 n (return ()) + , Send c2 n (return ()) + , Send c3 n (return ())] + Nothing + ping c1 c2 c3 (n + 1) + else return () + pong !c1 !c2 !c3 !n = do + if (n < iters) + then do + chanSelect + [ Recv c1 (const (return ())) + , Recv c2 (const (return ())) + , Recv c3 (const (return ()))] + Nothing + pong c1 c2 c3 (n + 1) + else return () diff --git a/fuzz/Main.hs b/fuzz/Main.hs new file mode 100644 --- /dev/null +++ b/fuzz/Main.hs @@ -0,0 +1,174 @@ +{-# LANGUAGE LambdaCase #-} + +module Main where + +import Control.Concurrent hiding (Chan) +import Control.Concurrent.GoChan +import Control.Monad +import Data.IORef +import Test.Hspec + +main :: IO () +main = + hspec $ + do describe "with buffer size 0" $ + do it "chanMake doesn't blow up" $ do void (chanMake 0 :: IO (Chan Int)) + it "send & recv doesn't blow up" $ do sendRecv 0 1 10 sendN drain + it "send & recv/select doesn't blow up" $ + do sendRecv 0 1 10 sendN drainSelect + it "send/select & recv doesn't blow up" $ + do sendRecv 0 1 10 sendNSelect drain + it "send/select & recv/select doesn't blow up" $ + do sendRecv 0 1 10 sendNSelect drainSelect + it "multi-case select doesn't blow up" $ do multiTest 0 + describe "with buffer size 1" $ + do it "chanMake doesn't blow up" $ void (chanMake 1 :: IO (Chan Int)) + it "send & recv doesn't blow up" $ do sendRecv 1 1 10 sendN drain + it "send & recv/select doesn't blow up" $ + do sendRecv 1 1 10 sendN drainSelect + it "send/select & recv doesn't blow up" $ + do sendRecv 1 1 10 sendNSelect drain + it "send/select & recv/select doesn't blow up" $ + do sendRecv 1 1 10 sendNSelect drainSelect + it "multi-case select doesn't blow up" $ do multiTest 1 + describe "with buffer size 2" $ + do it "chanMake doesn't blow up" $ void (chanMake 2 :: IO (Chan Int)) + it "send & recv doesn't blow up" $ do sendRecv 2 1 10 sendN drain + it "send & recv/select doesn't blow up" $ + do sendRecv 2 1 10 sendN drainSelect + it "send/select & recv doesn't blow up" $ + do sendRecv 2 1 10 sendNSelect drain + it "send/select & recv/select doesn't blow up" $ + do sendRecv 2 1 10 sendNSelect drainSelect + it "multi-case select doesn't blow up" $ do multiTest 2 + describe "with buffer size 3" $ + do it "chanMake doesn't blow up" $ do void (chanMake 3 :: IO (Chan Int)) + it "send & recv doesn't blow up" $ do sendRecv 3 1 10 sendN drain + it "send & recv/select doesn't blow up" $ + do sendRecv 3 1 10 sendN drainSelect + it "send/select & recv doesn't blow up" $ + do sendRecv 3 1 10 sendNSelect drain + it "send/select & recv/select doesn't blow up" $ + do sendRecv 3 1 10 sendNSelect drainSelect + it "multi-case select doesn't blow up" $ do multiTest 3 + +type Sender = Chan Int -> Int -> Int -> IO () + +type Drainer = Chan Int -> (Int -> IO ()) -> IO () -> IO () + +drain :: Drainer +drain ch recvAct closeAct = do + mn <- chanRecv ch + case mn of + Msg n -> do + recvAct n + drain ch recvAct closeAct + _ -> closeAct + +drainSelect :: Drainer +drainSelect ch recvAct closeAct = do + chanSelect + [ Recv + ch + (\case + Msg n -> do + recvAct n + drainSelect ch recvAct closeAct + _ -> closeAct)] + Nothing + +sendN :: Sender +sendN ch low hi = do + chanSend ch low + when (low < hi) (sendN ch (low + 1) hi) + +sendNSelect :: Sender +sendNSelect ch low hi = do + chanSelect [Send ch low (return ())] Nothing + when (low < hi) (sendNSelect ch (low + 1) hi) + +sendRecv :: Int -> Int -> Int -> Sender -> Drainer -> Expectation +sendRecv size low hi sender drainer = do + lock <- newEmptyMVar + c <- chanMake size + totalRef <- newIORef 0 + forkIO $ + do sender c low hi + chanClose c + drainer + c + (\n -> + modifyIORef' totalRef (+ n)) + (when (size > 0) (putMVar lock ())) + readIORef totalRef + -- when the channel is un-buffered, draining should act as synchronization; + -- only lock when the buffer size is greater than 0. + when + (size > 0) + (takeMVar lock) + total <- readIORef totalRef + total `shouldBe` sum [low .. hi] + +multiTest :: Int -> Expectation +multiTest size = do + lock1 <- newEmptyMVar + lock2 <- newEmptyMVar + c1 <- chanMake size + c2 <- chanMake size + c1sentRef <- newIORef 0 + c1recvdRef <- newIORef 0 + c2sentRef <- newIORef 0 + c2recvdRef <- newIORef 0 + forkIO $ ping2 c1sentRef c2sentRef c1 c2 0 (putMVar lock2 ()) + pong2 c1recvdRef c2recvdRef c1 c2 0 (putMVar lock1 ()) + takeMVar lock1 + takeMVar lock2 + c1sent <- readIORef c1sentRef + c1recvd <- readIORef c1recvdRef + c2sent <- readIORef c2sentRef + c2recvd <- readIORef c2recvdRef + -- each channel should recv as often as it is sent on. + (c1sent, c2sent) `shouldBe` + (c1recvd, c2recvd) + +ping2 + :: IORef Int + -> IORef Int + -> Chan Int + -> Chan Int + -> Int + -> IO () + -> IO () +ping2 ref1 ref2 c1 c2 n doneAct = do + if (n < 20) + then do + chanSelect + [ Send c1 n (void (modifyIORef' ref1 (+ 1))) + , Send c2 n (void (modifyIORef' ref2 (+ 1)))] + Nothing + ping2 ref1 ref2 c1 c2 (n + 1) doneAct + else doneAct + +pong2 + :: IORef Int + -> IORef Int + -> Chan Int + -> Chan Int + -> Int + -> IO () + -> IO () +pong2 ref1 ref2 c1 c2 n doneAct = do + if (n < 20) + then do + chanSelect + [ Recv + c1 + (\case + Msg n -> modifyIORef' ref1 (+ 1)) + , Recv + c2 + (\case + Msg n -> modifyIORef' ref2 (+ 1))] + Nothing + pong2 ref1 ref2 c1 c2 (n + 1) doneAct + else doneAct diff --git a/gochan.cabal b/gochan.cabal new file mode 100644 --- /dev/null +++ b/gochan.cabal @@ -0,0 +1,78 @@ +name: gochan +version: 0.0.1 +cabal-version: >=1.10 +build-type: Simple +license: MIT +license-file: LICENSE +copyright: Copyright (C) 2016 Charles Strahan +maintainer: Charles Strahan +homepage: http://github.com/cstrahan/gochan +bug-reports: http://github.com/cstrahan/gochan/issues +synopsis: Go-style channels +description: This library provides bounded channels similar to those popularized + by the Go programming language. + . + THIS LIBRARY IS STILL ALPHA AND SUBJECT TO CHANGE. + . + Breaking changes during the 0.0.X series won't result in a major + version bump (I hope to release a stable 1.0.0 soon). + . + Please give gochan a try and file an issue or pull request if you + discover any ways to improve this library. +category: Concurrency +author: Charles Strahan +extra-source-files: + .gitignore + +source-repository head + type: git + location: git://github.com/cstrahan/gochan.git + +library + exposed-modules: + Control.Concurrent.GoChan + build-depends: + base >=4 && <5, + array -any, + random -any, + vector -any, + vector-algorithms -any, + primitive -any, + ghc-prim -any + default-language: Haskell2010 + hs-source-dirs: src + ghc-options: -O2 + +executable bench + main-is: Main.hs + build-depends: + base >=4 && <5, + criterion -any, + gochan -any, + random -any + default-language: Haskell98 + hs-source-dirs: bench + ghc-options: -threaded -rtsopts -with-rtsopts=-N -O2 + +executable weight + main-is: Main.hs + build-depends: + base >=4 && <5, + weigh -any, + gochan -any + default-language: Haskell98 + hs-source-dirs: weight + ghc-options: -threaded -rtsopts -with-rtsopts=-N -O2 + +test-suite fuzz + type: exitcode-stdio-1.0 + main-is: Main.hs + build-depends: + base >=4 && <5, + hspec -any, + hspec-core -any, + hspec-core -any, + gochan -any + default-language: Haskell2010 + hs-source-dirs: fuzz + ghc-options: -Wall diff --git a/src/Control/Concurrent/GoChan.hs b/src/Control/Concurrent/GoChan.hs new file mode 100644 --- /dev/null +++ b/src/Control/Concurrent/GoChan.hs @@ -0,0 +1,793 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE ExistentialQuantification #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE Rank2Types #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} + +-- | This module provides bounded channels similar to those seen in the Go +-- programming language. +module Control.Concurrent.GoChan + ( -- * Types + Chan + ,Result(..) + ,Case(..) + -- * Functions + ,chanMake + ,chanClose + ,chanRecv + ,chanTryRecv + ,chanSend + ,chanTrySend + ,chanSelect) + where + +import Control.Concurrent.MVar +import Control.Monad +import Control.Monad.Primitive (RealWorld) +import Data.Array.IO +import Data.IORef +import Data.List (intercalate) +import Data.Maybe (fromJust, isJust, isNothing) +import qualified Data.Vector as V +import qualified Data.Vector.Algorithms.Heap as VAH +import qualified Data.Vector.Generic as VG +import qualified Data.Vector.Generic.Mutable as VGM +import qualified Data.Vector.Mutable as VM +import Data.Word +import GHC.Prim (Any) +import System.IO.Unsafe +import System.Random +import Unsafe.Coerce + +-- | The core data type. A @'Chan' α@ provides a conduit through which messages +-- of type @α@ can be sent and received. +data Chan a = Chan + { _qcount :: {-# UNPACK #-} !(IORef Int) + , _qsize :: {-# UNPACK #-} !Int + , _buf :: {-# UNPACK #-} !(IOArray Int a) + , _sendx :: {-# UNPACK #-} !(IORef Int) + , _recvx :: {-# UNPACK #-} !(IORef Int) + , _sendq :: {-# UNPACK #-} !SuspQ + , _recvq :: {-# UNPACK #-} !SuspQ + , _lock :: {-# UNPACK #-} !(MVar ()) + , _closed :: {-# UNPACK #-} !(IORef Bool) + , _id :: {-# UNPACK #-} !Word64 + } + +data SuspQ = SuspQ + { _first :: {-# UNPACK #-} !(IORef (Maybe SomeSuspend)) + , _last :: {-# UNPACK #-} !(IORef (Maybe SomeSuspend)) + } + +data SomeSuspend = + forall a. SomeSuspend (Suspend a) + +data Suspend a = forall b. Suspend + { _selectDone :: !(Maybe (IORef Bool)) + , _case :: !(Maybe (Case b)) + , _next :: !(IORef (Maybe SomeSuspend)) + , _prev :: !(IORef (Maybe SomeSuspend)) + , _elem :: !(Maybe (IORef a)) + , _chan :: !(Chan a) + , _park :: !(MVar (Maybe (Suspend a))) -- park, and when unparked, the awoken sleeper is given + , _sid :: !Word64 + } + +-- | The `Result` record represents the result of waiting to recieve a message +-- on a given channel. +data Result a + = Msg a + | Closed + +-- | When used with 'chanSelect`, a @'Case' α@ provides a means of waiting for a +-- 'Send' or 'Recv' to complete on a given channel, subsequently invoking the +-- given callback. +-- +-- [@'Recv' (chan :: 'Chan' β) (act :: Result β -> IO α)@] Wait to receive a +-- @msg@ on @chan@, invoking @act result@ when selected. +-- +-- [@'Send' (chan :: 'Chan' β) (msg :: β) (act :: IO α)@] Wait to send @msg@ on +-- @chan@, invoking @act@ when selected. +data Case a + = forall b. Recv (Chan b) + (Result b -> IO a) + | forall b. Send (Chan b) + b + (IO a) + +{-# INLINE caseChanId #-} + +caseChanId :: Case a -> Word64 +caseChanId (Recv chan _) = _id chan +caseChanId (Send chan _ _) = _id chan + +{-# NOINLINE currIdRef #-} + +currIdRef :: IORef Word64 +currIdRef = unsafePerformIO (newIORef 0) + +{-# NOINLINE currSIdRef #-} + +currSIdRef :: IORef Word64 +currSIdRef = unsafePerformIO (newIORef 0) + +shuffleVector + :: (VGM.MVector v e) + => v RealWorld e -> IO () +shuffleVector !xs = do + let !size = VGM.length xs + forM_ [1 .. size - 1] $ + \i -> do + j <- randomRIO (0, i) + vi <- VGM.read xs i + vj <- VGM.read xs j + VGM.write xs j vi + VGM.write xs i vj + +-- | When no default action is given, blocks until one of the cases can run, +-- then it executes that case. It chooses one at random if multiple are ready. +-- +-- If given a default action, and no cases can run, immediately executes the +-- default action. +chanSelect + :: [Case a] -- ^ The list of cases. + -> Maybe (IO a) -- ^ When 'Nothing', wait synchronously to select a + -- case; when @'Just' act@, run @act@ as a default + -- instead of blocking. + -> IO a -- ^ The value resulting from the selected 'Case' + -- (or default action, if given). +chanSelect cases mdefault = do + -- randomize the poll order to enforce fairness. + !pollOrder <- + -- for reasons unknown to me, we get better performance when we don't + -- explicitly lift out the common fromList expression. + do vec <- V.thaw (V.fromList cases) + shuffleVector vec + V.unsafeFreeze vec + let !ncases = VG.length pollOrder + -- we need to aquire locks in a consistent order so we don't deadlock. + !lockOrder <- + do vec <- V.thaw (V.fromList cases) + VAH.sortBy + (\cas1 cas2 -> + caseChanId cas1 `compare` caseChanId cas2) + vec + V.unsafeFreeze vec + selLock lockOrder + -- PASS 1 - attempt to dequeue a Suspend from one of the channels. + let pass1 !n = do + if n /= ncases + then case pollOrder VG.! n of + Recv chan act -> do + ms <- dequeue (_sendq chan) + case ms of + Just (SomeSuspend s) -> do + elemRef <- newIORef undefined + recv chan (unsafeCoerceSuspend s) (Just elemRef) (selUnlock lockOrder) + val <- readIORef elemRef + act (Msg val) + _ -> do + !qcount <- readIORef (_qcount chan) + if qcount > 0 + then do + !recvx <- readIORef (_recvx chan) + val <- readArray (_buf chan) recvx + let !recvx' = + let !x = recvx + 1 + in if x == _qsize chan + then 0 + else x + writeIORef (_recvx chan) $! recvx' + writeIORef (_qcount chan) (qcount - 1) + selUnlock lockOrder + act (Msg val) + else do + !isClosed <- readIORef (_closed chan) + if isClosed + then do + selUnlock lockOrder + act Closed + else do + pass1 (n + 1) + Send chan val act -> do + !isClosed <- readIORef (_closed chan) + if isClosed + then do + selUnlock lockOrder + fail "send on closed channel" + else do + ms <- dequeue (_recvq chan) + case ms of + Just (SomeSuspend s) -> do + send chan (unsafeCoerceSuspend s) val (selUnlock lockOrder) + act + _ -> do + !qcount <- readIORef (_qcount chan) + if qcount < _qsize chan + then do + !sendx <- readIORef (_sendx chan) + writeArray (_buf chan) sendx val + let !sendx' = + let !x = sendx + 1 + in if x == _qsize chan + then 0 + else x + writeIORef (_sendx chan) sendx' + writeIORef (_qcount chan) (qcount + 1) + selUnlock lockOrder + act + else do + pass1 (n + 1) + else case mdefault of + Just def -> do + selUnlock lockOrder + def + _ -> do + -- PASS 2 - enqueue a Suspend for each case. + -- + -- the shared 'park' MVark is used wake up this + -- thread. + park <- newEmptyMVar + -- the shared 'selectDone' IORef allows the other + -- threads to know who won the race to wake us up + -- (via CAS). + selectDone <- newIORef False + ss <- + V.generateM + ncases + (\n -> do + next <- newIORef Nothing + prev <- newIORef Nothing + id <- + atomicModifyIORef' + currSIdRef + (\currId -> + (currId + 1, currId)) + case lockOrder V.! n of + cas@(Send chan val _) -> do + elemRef <- newIORef (unsafeCoerce val) + let !s = + SomeSuspend + (Suspend + (Just selectDone) + (Just (unsafeCoerceCase cas)) + next + prev + (Just elemRef) + (unsafeCoerceChan chan) + park + id) + enqueue (_sendq chan) s + return s + cas@(Recv chan _) -> do + elemRef <- newIORef undefined + let !s = + SomeSuspend + (Suspend + (Just selectDone) + (Just (unsafeCoerceCase cas)) + next + prev + (Just elemRef) + (unsafeCoerceChan chan) + park + id) + enqueue (_recvq chan) s + return s) + selUnlock lockOrder + ms <- takeMVar park + selLock lockOrder + -- PASS 3 - dequeue each Suspend we enqueued + -- earlier. + let pass3 !n = do + case ss VG.! n of + someS@(SomeSuspend s) -> + case s of + (Suspend _ cas _ _ _ _ _ _) -> + case cas of + Just (Send _ _ _) -> do + -- + dequeueSuspend + (_sendq (_chan s)) + someS + Just (Recv _ _) -> do + dequeueSuspend (_recvq (_chan s)) someS + when ((n + 1) /= ncases) (pass3 (n + 1)) + pass3 0 + case ms of + Just s -> do + case s of + (Suspend _ cas _ _ _ _ _ _) -> + case cas of + Just (Send chan _ act) -> do + selUnlock lockOrder + unsafeCoerceSendAction act + Just (Recv chan act) -> do + !val <- readIORef (fromJust (_elem s)) + selUnlock lockOrder + unsafeCoerceRecvAction act (Msg val) + _ -> do + -- channel closed, restart loop to figure + -- out which one. + pass1 + 0 + pass1 0 + +{-# INLINE unsafeCoerceSendAction #-} + +unsafeCoerceSendAction :: IO a -> IO b +unsafeCoerceSendAction = unsafeCoerce + +{-# INLINE unsafeCoerceRecvAction #-} + +unsafeCoerceRecvAction :: (Result b -> IO a) -> (Result d -> IO c) +unsafeCoerceRecvAction = unsafeCoerce + +{-# INLINE unsafeCoerceSuspend #-} + +unsafeCoerceSuspend :: Suspend a -> Suspend b +unsafeCoerceSuspend = unsafeCoerce + +{-# INLINE unsafeCoerceChan #-} + +unsafeCoerceChan :: Chan a -> Chan b +unsafeCoerceChan = unsafeCoerce + +{-# INLINE unsafeCoerceCase #-} + +unsafeCoerceCase :: Case a -> Case b +unsafeCoerceCase = unsafeCoerce + +{-# INLINE lockCase #-} + +lockCase :: Case a -> IO () +lockCase (Recv chan _) = takeMVar (_lock chan) +lockCase (Send chan _ _) = takeMVar (_lock chan) + +{-# INLINE unlockCase #-} + +unlockCase :: Case a -> IO () +unlockCase (Recv chan _) = putMVar (_lock chan) () +unlockCase (Send chan _ _) = putMVar (_lock chan) () + +selLock + :: (VG.Vector v e, e ~ Case a) + => v (Case a) -> IO () +selLock !vec = do + go 0 maxBound + where + len = VG.length vec + go n prevId = do + let !cas = vec VG.! n + if n == len - 1 + then do + lockCase cas + else do + when (caseChanId cas /= prevId) $ do lockCase cas + go (n + 1) (caseChanId cas) + +selUnlock + :: (VG.Vector v e, e ~ Case a) + => v (Case a) -> IO () +selUnlock !vec = do + go (len - 1) maxBound + where + len = VG.length vec + go n prevId = do + let !cas = vec VG.! n + if n == 0 + then do + unlockCase cas + else do + when (caseChanId cas /= prevId) $ do (unlockCase cas) + go (n - 1) (caseChanId cas) + +-- | Make a channel with the given buffer size. +chanMake + :: Int -> IO (Chan a) +chanMake !size = do + ary <- newArray_ (0, size - 1) + qcount <- newIORef 0 + sendx <- newIORef 0 + recvx <- newIORef 0 + sendq_first <- newIORef Nothing + sendq_last <- newIORef Nothing + recvq_first <- newIORef Nothing + recvq_last <- newIORef Nothing + lock <- newMVar () + closed <- newIORef False + id <- + atomicModifyIORef' + currIdRef + (\currId -> + (currId + 1, currId)) + return + Chan + { _qcount = qcount + , _qsize = size + , _buf = ary + , _sendx = sendx + , _recvx = recvx + , _sendq = SuspQ sendq_first sendq_last + , _recvq = SuspQ recvq_first recvq_last + , _lock = lock + , _closed = closed + , _id = id + } + +-- | Wait to successfully send a message on a channel. +-- +-- Throws an exception when sending on a closed channel. +chanSend + :: Chan a -> a -> IO () +chanSend !chan !val = void $ chanSendInternal chan val True + +-- | Attempt to send a message on a channel. The message will be sent iff the +-- channel has spare space in its buffer or another thread is waiting to recieve +-- a message on this channel. +-- +-- Returns 'True' iff the message was sent. +-- +-- Throws an exception when sending on a closed channel. +chanTrySend + :: Chan a -> a -> IO Bool +chanTrySend !chan !val = chanSendInternal chan val False + +-- returns True if a message was sent; otherwise returns False +chanSendInternal :: Chan a -> a -> Bool -> IO Bool +chanSendInternal !chan !val !block = do + !isClosed <- readIORef (_closed chan) + !recvq_first <- readIORef (_first (_recvq chan)) + !qcount <- readIORef (_qcount chan) + -- Fast path: check for failed non-blocking operation without acquiring the lock. + if not block && not isClosed && ((_qsize chan == 0 && isJust recvq_first) || (_qsize chan > 0 && qcount == _qsize chan)) + then return False + else do + takeMVar (_lock chan) + !isClosed <- readIORef (_closed chan) + if isClosed + then do + putMVar (_lock chan) () + fail "send on closed channel" + else do + ms <- dequeue (_recvq chan) + case ms of + Just (SomeSuspend s) -> do + send chan (unsafeCoerceSuspend s) val (putMVar (_lock chan) ()) + return True + Nothing -> do + !qcount <- readIORef (_qcount chan) + if qcount < _qsize chan + then do + !sendx <- readIORef (_sendx chan) + writeArray (_buf chan) sendx val + writeIORef (_sendx chan) $! (sendx + 1) + let !sendx' = sendx + 1 + if sendx' == _qsize chan + then writeIORef (_sendx chan) 0 + else writeIORef (_sendx chan) $! sendx' + writeIORef (_qcount chan) (qcount + 1) + putMVar (_lock chan) () + return True + else if not block + then do + putMVar (_lock chan) () + return False + else do + next <- newIORef Nothing + prev <- newIORef Nothing + elem <- newIORef val + park <- newEmptyMVar -- we're about to park + id <- + atomicModifyIORef' + currSIdRef + (\currId -> + (currId + 1, currId)) + let !s = (SomeSuspend (Suspend Nothing Nothing next prev (Just elem) chan park id)) + enqueue (_sendq chan) s + putMVar (_lock chan) () + ms' <- takeMVar park + case ms' of + Nothing -> do + !isClosed <- readIORef (_closed chan) + unless isClosed (fail "chansend: spurious wakeup") + fail "send on closed channel" + _ -> return True + +send :: Chan a -> Suspend a -> a -> IO () -> IO () +send !chan !s !val !unlock = do + case _elem s of + Just elemRef -> do + writeIORef elemRef val + _ -> do + return () + unlock + putMVar (_park s) (Just s) -- unpark + +-- | Close a channel. +chanClose + :: Chan a -> IO () +chanClose !chan = do + takeMVar (_lock chan) + !isClosed <- readIORef (_closed chan) + when isClosed $ + do putMVar (_lock chan) () + fail "close of closed channel" + writeIORef (_closed chan) True + ss <- releaseReaders [] chan + ss <- releaseWriters ss chan + putMVar (_lock chan) () + wakeSuspends ss + where + releaseReaders ss chan = do + ms <- dequeue (_recvq chan) + case ms of + Nothing -> return ss + Just s -> releaseReaders (s : ss) chan + releaseWriters ss chan = do + ms <- dequeue (_sendq chan) + case ms of + Nothing -> return ss + Just s -> releaseReaders (s : ss) chan + wakeSuspends ss = + forM_ + ss + (\(SomeSuspend s) -> + putMVar (_park s) Nothing) + +-- data type used internally to represent the result of waiting on a channel +data RecvResult + = RecvWouldBlock + | RecvGotMessage + | RecvClosed + +-- | Attempt to receive a message on a channel. A message will be recieved iff +-- the channel has a message in its buffer or another thread is waiting to +-- send a message on this channel. +chanTryRecv + :: Chan a -> IO (Maybe (Result a)) +chanTryRecv !chan = do + ref <- newIORef undefined + chanRecvInternal chan (Just ref) False >>= + \case + RecvWouldBlock -> return Nothing + RecvClosed -> return (Just Closed) + RecvGotMessage-> Just <$> Msg <$> readIORef ref + +-- | Wait to receive a message on a channel. +chanRecv + :: Chan a -> IO (Result a) +chanRecv !chan = do + ref <- newIORef undefined + chanRecvInternal chan (Just ref) True >>= + \case + RecvWouldBlock -> fail "the impossible happened" + RecvClosed -> return Closed + RecvGotMessage -> Msg <$> readIORef ref + +chanRecvInternal :: Chan a -> Maybe (IORef a) -> Bool -> IO RecvResult +chanRecvInternal !chan !melemRef !block = do + -- Fast path: check for failed non-blocking operation without acquiring the lock. + -- WARNING: the order of these reads is important. + !sendq_first <- readIORef (_first (_sendq chan)) + !qcount <- atomicReadIORef (_qcount chan) + !isClosed <- atomicReadIORef (_closed chan) + if not block && ((_qsize chan == 0 && isNothing sendq_first) || (_qsize chan > 0 && qcount == _qsize chan)) && not isClosed + then return RecvWouldBlock + else do + takeMVar (_lock chan) + !isClosed <- readIORef (_closed chan) + !qcount <- readIORef (_qcount chan) + if isClosed && qcount == 0 + then do + putMVar (_lock chan) () + return RecvClosed + else do + ms <- dequeue (_sendq chan) + case ms of + Just (SomeSuspend s) -> do + recv chan (unsafeCoerceSuspend s) melemRef (putMVar (_lock chan) ()) + return RecvGotMessage + _ -> + if qcount > 0 + then do + !recvx <- readIORef (_recvx chan) + val <- readArray (_buf chan) recvx + case melemRef of + Just elemRef -> writeIORef elemRef val + _ -> return () + let !recvx' = + let !x = recvx + 1 + in if x == _qsize chan + then 0 + else x + writeIORef (_recvx chan) $! recvx' + modifyIORef' (_qcount chan) (subtract 1) + putMVar (_lock chan) () + return RecvGotMessage + else if not block + then do + putMVar (_lock chan) () + return RecvWouldBlock + else do + next <- newIORef Nothing + prev <- newIORef Nothing + park <- newEmptyMVar -- we're about to park + id <- + atomicModifyIORef' + currSIdRef + (\currId -> + (currId + 1, currId)) + let !s = SomeSuspend (Suspend Nothing Nothing next prev melemRef chan park id) + enqueue (_recvq chan) s + putMVar (_lock chan) () + ms' <- takeMVar park -- park + if isJust ms' + then return RecvGotMessage + else return RecvClosed + +recv :: Chan a -> Suspend a -> Maybe (IORef a) -> IO () -> IO () +recv !chan !s !melemRef !unlock = do + if _qsize chan == 0 + then case melemRef of + Just elemRef -> do + !val <- readIORef (fromJust (_elem s)) + writeIORef elemRef val + _ -> return () + else do + !recvx <- readIORef (_recvx chan) + val <- readArray (_buf chan) recvx + case melemRef of + Just elemRef -> writeIORef elemRef val + _ -> return () + !val' <- readIORef (fromJust (_elem s)) + writeArray (_buf chan) recvx val' + let !recvx' = + let !x = recvx + 1 + in if x == _qsize chan + then 0 + else x + writeIORef (_recvx chan) $! recvx' + writeIORef (_sendx chan) $! recvx' + unlock + putMVar (_park s) (Just s) -- unpark + +-- enqueue a Supsend. +enqueue + :: SuspQ -> SomeSuspend -> IO () +enqueue !q someS@(SomeSuspend s) = do + writeIORef (_next s) Nothing + mx <- readIORef . _last $ q + case mx of + Nothing -> do + writeIORef (_prev s) Nothing + writeIORef (_first q) (Just someS) + writeIORef (_last q) (Just someS) + Just someX@(SomeSuspend x) -> do + writeIORef (_prev s) (Just someX) + writeIORef (_next x) (Just someS) + writeIORef (_last q) (Just someS) + +-- dequeue each Suspend until one is found that can be resumed; +-- if we dequeue one that participates in a select, and it is already +-- flagged as selected, continue dequeuing. +dequeue + :: SuspQ -> IO (Maybe SomeSuspend) +dequeue !q = do + !ms <- readIORef (_first q) + case ms of + Nothing -> return Nothing + Just someS@(SomeSuspend s) -> do + !my <- readIORef (_next s) + case my of + Nothing -> do + writeIORef (_first q) Nothing + writeIORef (_last q) Nothing + Just someY@(SomeSuspend y) -> do + writeIORef (_prev y) Nothing + writeIORef (_first q) (Just someY) + writeIORef (_next s) Nothing + case _selectDone s of + Nothing -> return (Just someS) + Just doneRef -> do + done <- readIORef doneRef + if not done + then do + -- attempt to set the "selectdone" flag and return + -- the Suspend. + -- if someone beat us to it, try dequeuing again. + oldDone <- + atomicModifyIORef' + doneRef + (\oldDone -> + (True, oldDone)) + if oldDone + then do + -- we did *not* win the race; try again + dequeue q + else do + -- we won! done. + return + (Just someS) + else do + -- we did *not* win the race; try again + dequeue q + +{-# INLINE atomicReadIORef #-} + +atomicReadIORef :: IORef a -> IO a +atomicReadIORef !ref = + atomicModifyIORef' + ref + (\oldVal -> + (oldVal, oldVal)) + +-- TODO: Consider (carefully) using pointer equality instead of maintaining a +-- unique ID. +eqSuspend + :: Suspend a -> Suspend b -> Bool +eqSuspend !s1 !s2 = _sid s1 == _sid s2 + +dequeueSuspend :: SuspQ -> SomeSuspend -> IO () +dequeueSuspend !q someS@(SomeSuspend s) = do + !mx <- readIORef (_prev s) + !my <- readIORef (_next s) + case mx of + Just someX@(SomeSuspend x) -> + case my of + Just someY@(SomeSuspend y) -> do + writeIORef (_next x) (Just someY) + writeIORef (_prev y) (Just someX) + writeIORef (_next s) Nothing + writeIORef (_prev s) Nothing + _ -> do + writeIORef (_next x) Nothing + writeIORef (_last q) (Just someX) + writeIORef (_prev s) Nothing + _ -> + case my of + Just someY@(SomeSuspend y) -> do + writeIORef (_prev y) Nothing + writeIORef (_first q) (Just someY) + writeIORef (_next s) Nothing + _ -> do + !mfirst <- readIORef (_first q) + case mfirst of + Just someFirst@(SomeSuspend first) -> + when (first `eqSuspend` s) $ + do writeIORef (_first q) Nothing + writeIORef (_last q) Nothing + _ -> return () + +-------------------------------------------------------------------------------- +-- misc. debug utils +waitqToList + :: SuspQ -> IO [SomeSuspend] +waitqToList q = do + !ms <- readIORef (_first q) + case ms of + Just s -> sleeperChain s + _ -> return [] + +sleeperChain :: SomeSuspend -> IO [SomeSuspend] +sleeperChain someS@(SomeSuspend s) = do + !mnext <- readIORef (_next s) + case mnext of + Just next -> do + ss <- sleeperChain next + return (someS : ss) + _ -> return [someS] + +printSuspQ :: SuspQ -> IO () +printSuspQ q = do + ss <- waitqToList q + let !chain = + intercalate + "->" + (map + (\(SomeSuspend s) -> + "(SID: " ++ show (_sid s) ++ ", CID: " ++ show (_id (_chan s)) ++ ")") + ss) + putStrLn $ "WAITQ: " ++ chain diff --git a/weight/Main.hs b/weight/Main.hs new file mode 100644 --- /dev/null +++ b/weight/Main.hs @@ -0,0 +1,79 @@ +module Main where + +import Control.Concurrent hiding (Chan) +import Control.Concurrent.GoChan +import Control.Monad +import System.IO.Unsafe +import Weigh + +main :: IO () +main = + mainWith + (do io "creating a channel" chanConstruction () + io "unbuffered send & recv" unbufferedSendRecv chanUnbuffered + io "unbuffered select send & recv" unselectSendRecv chanUnbuffered + io "unbuffered recv" simpleRecv chanUnbufferedAlreadySent + io "unbuffered select recv" unbufferedSendRecv chanUnbufferedAlreadySent + io "buffered send" simpleSend chanBufferedEmpty + io "buffered recv" simpleRecv chanBufferedAlreadySent + io "buffered select send" selectSend chanBufferedEmpty + io "buffered select recv" selectRecv chanBufferedAlreadySent) + +chanConstruction :: () -> IO () +chanConstruction = const (void (chanMake 0 :: IO (Chan Int))) + +unbufferedSendRecv :: Chan Int -> IO () +unbufferedSendRecv c = do + forkIO $ chanSend c 0 + chanRecv c + return () + +unselectSendRecv :: Chan Int -> IO () +unselectSendRecv c = do + forkIO $ chanSelect [Send c 0 (return ())] Nothing + chanSelect [Recv c (const (return ()))] Nothing + return () + +simpleSend :: Chan Int -> IO () +simpleSend c = chanSend c 0 + +simpleRecv :: Chan Int -> IO () +simpleRecv c = void (chanRecv c) + +selectSend :: Chan Int -> IO () +selectSend c = chanSelect [Send c 0 (return ())] Nothing + +selectRecv :: Chan Int -> IO () +selectRecv c = chanSelect [Recv c (const (return ()))] Nothing + +-------------------------------------------------------------------------------- +-- We don't want these to show up in allocations, so carefully unsafePerformIO +-- here. The threading is scary, but I can't think of any other way to test +-- this, given the concurrent nature of this library. +{-# NOINLINE chanUnbuffered #-} + +chanUnbuffered :: Chan Int +chanUnbuffered = unsafePerformIO (chanMake 0) + +{-# NOINLINE chanUnbufferedAlreadySent #-} + +chanUnbufferedAlreadySent :: Chan Int +chanUnbufferedAlreadySent = + unsafePerformIO $ + do c <- chanMake 0 + forkIO $ chanSend c 0 + return c + +{-# NOINLINE chanBufferedEmpty #-} + +chanBufferedEmpty :: Chan Int +chanBufferedEmpty = unsafePerformIO (chanMake 1) + +{-# NOINLINE chanBufferedAlreadySent #-} + +chanBufferedAlreadySent :: Chan Int +chanBufferedAlreadySent = + unsafePerformIO $ + do c <- chanMake 1 + chanSend c 0 + return c