Search code examples
haskellselectchannelstm

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


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."


Solution

  • 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.