How to implement the equivalent of Go's select statement for Haskell STM channels?

Question

The Go language has a select statement that can be used to poll multiple channels and carry out a particular action depending on which channel is non-empty first.

E.g.

select {
  case a := <- chanA:
    foo(a)
  case b := <- chanB:
    baz(b)
  case c := <- chanC:
    bar(c)
}

This will wait until either chanA, chanB or chanC is non-empty, then if for instance chanB is non-empty, it will read from chanB and store the result in b, then call baz(b). A default: clause can also be added, meaning the select statement will not wait on the channels and instead will do whatever the default clause is if all the channels are empty.

What would be the best way to implement something like this for STM TChans in Haskell? It could be done naively by an if-else chain: checking if each chan isEmptyChan, and if it's not empty then reading from it and calling the appropriate function, or else calling retry if all of the channels are empty. I was wondering if there would be a more elegant/idiomatic way to do this?

Note that Go's select statement can also include send statements in its cases, and will only complete a send statement if its channel is empty. It would be great if that functionality could be duplicated too, although I'm not sure whether there would be an elegant way to do so.

Only slightly related but something I just noticed and I'm not sure where to post it: there's a typo on the Control.Monad.STM page in the description for retry:

"The implementation may block the thread until one of the TVars that it has read from has been udpated."

Shu

You might want to look at race from Control.Concurrent.Async. – Folger

It's worth noting that go doesn't perform the first action available, but any one available, selected randomly. It specifically will not starve channels just because they're defined later or unlucky in the selection path. – Flatter

This is completely unlike Go's select. Channels in Go are bounded, unlike TChan (making them actually useful) and select can be used with send operations. – Junction

Answer 1

Starvation avoiding

foreverK :: (a -> m a) -> a -> m ()
foreverK loop = go
 where go = loop >=> go

-- Existential, not really required, but feels more like the Go version
data ChanAct = Action (TChan a) (a -> STM ())

perform :: STM ()
perform (Action c a) = readTChan c >>= a

foreverSelectE :: [ChanAct] -> STM ()
foreverSelectE = foreverSelect . map perform

foreverSelect :: [STM ()] -> STM ()
foreverSelect = foreverK $ \xs -> first xs >> return (rotate1 xs)

-- Should only be defined for non-empty sequences, but return () is an okay default.
-- Will NOT block the thread, but might do nothing.
first :: [STM ()] -> STM ()
first = foldr orElse (return ())

-- Should only be defined for non-empty sequences, really.
-- Also, using a list with O(1) viewL and snoc could be better.
rotate1 :: [a] -> [a]
rotate1 []    = []
rotate1 (h:t) = t ++ [h]

example = foreverSelectE
    [ Action chanA foo
    , Action charB baz
    , Action chanC bar
    ]

To avoid forever, you could instead have a mkSelect :: [STM ()] -> STM (STM ()) that "hides" a TVar [STM ()] and rotates it each it is used like:

example1 :: STM ()
example1 = do
    select <- mkSelect [actions] -- Just set-up
    stuff1
    select -- does one of the actions
    stuff2
    select -- does one of the actions

main = OpenGL.idleCallback $= atomically example1

Extending that technique you could have a select that reported if it performed an action or which action it performed or even looped until all the actions would block, etc.

Enjoin

Answer 2

You can implement select semantics (both for read and write) using orElse (note: it is specific to ghc.) For example:

forever $ atomically $
  writeTChan chan1 "hello" `orElse` writeTChan chan2 "world" `orElse` ...

The idea is that when one action retries (e.g. you are writing to chan, but it is full; or you are reading from chan, but it is empty), the second action is performed. The default statement is just a return () as the last action in the chain.

Add: As @Dustin noted, go selects random branch for good reason. Probably the easiest solution is to shuffle actions on each iteration, it should be ok in most cases. Replicating go semantics properly (shuffle active branches only) is a bit harder. Probably manually inspecting isEmptyChan for all branches is the way to go.

Th

Won't this have problems with starvation though, as Dustin commented above? – Ekaterinburg

@Ekaterinburg I believe you're right: (assuming Yuras' example were using readTChan) readTChan chanN would only ever be read from while all chans < N are empty. So the possibility of starvation is actually even worse than you might expect at first. – Czechoslovak

Thank you! I didn't realise orElse could be chained like that. I also hadn't considered the possibility of starvation. – Shu