Implement receiveChanTimeout

This closes #47.

Author: edsko

distributed-process/src/Control/Distributed/Process.hs

@@ -19,13 +19,15 @@ module Control.Distributed.Process
     -- * Basic messaging
   , send 
   , expect
+  , expectTimeout
     -- * Channels
   , ReceivePort
   , SendPort
   , sendPortId
   , newChan
   , sendChan
   , receiveChan
+  , receiveChanTimeout
   , mergePortsBiased
   , mergePortsRR
     -- * Advanced messaging
@@ -91,7 +93,6 @@ module Control.Distributed.Process
   , bracket_
   , finally
     -- * Auxiliary API
-  , expectTimeout
   , spawnAsync
   , spawnSupervised
   , spawnLink
@@ -215,6 +216,7 @@ import Control.Distributed.Process.Internal.Primitives
   , finally
     -- Auxiliary API
   , expectTimeout
+  , receiveChanTimeout
   , spawnAsync
     -- Reconnecting
   , reconnect

distributed-process/src/Control/Distributed/Process/Internal/Primitives.hs

@@ -54,6 +54,7 @@ module Control.Distributed.Process.Internal.Primitives
   , finally
     -- * Auxiliary API
   , expectTimeout
+  , receiveChanTimeout
   , spawnAsync
   , linkNode
   , linkPort
@@ -74,6 +75,7 @@ import Data.Binary (decode)
 import Data.Time.Clock (getCurrentTime)
 import Data.Time.Format (formatTime)
 import System.Locale (defaultTimeLocale)
+import System.Timeout (timeout)
 import Control.Monad.Reader (ask)
 import Control.Monad.IO.Class (MonadIO, liftIO)
 import Control.Applicative ((<$>))
@@ -188,7 +190,7 @@ newChan = do
              )
   where
     finalizer :: StrictMVar LocalProcessState -> LocalSendPortId -> IO ()
-    finalizer processState lcid = modifyMVar_ processState $ 
+    finalizer st lcid = modifyMVar_ st $ 
       return . (typedChannelWithId lcid ^= Nothing)
 
 -- | Send a message on a typed channel
@@ -204,18 +206,29 @@ sendChan (SendPort cid) msg = do
 -- | Wait for a message on a typed channel
 receiveChan :: Serializable a => ReceivePort a -> Process a
 receiveChan = liftIO . atomically . receiveSTM 
-  where
-    receiveSTM :: ReceivePort a -> STM a
-    receiveSTM (ReceivePortSingle c) = 
-      readTQueue c
-    receiveSTM (ReceivePortBiased ps) =
-      foldr1 orElse (map receiveSTM ps)
-    receiveSTM (ReceivePortRR psVar) = do
-      ps <- readTVar psVar
-      a  <- foldr1 orElse (map receiveSTM ps)
-      writeTVar psVar (rotate ps)
-      return a
 
+-- | Like 'receiveChan' but with a timeout. If the timeout is 0, do a 
+-- non-blocking check for a message.
+receiveChanTimeout :: Serializable a => Int -> ReceivePort a -> Process (Maybe a)
+receiveChanTimeout 0 ch = liftIO . atomically $ 
+  (Just <$> receiveSTM ch) `orElse` return Nothing
+receiveChanTimeout n ch = liftIO . timeout n . atomically $ 
+  receiveSTM ch
+
+-- | Receive a message from a typed channel as an STM transaction.
+-- 
+-- The transaction retries when no message is available.
+receiveSTM :: ReceivePort a -> STM a
+receiveSTM (ReceivePortSingle c) = 
+    readTQueue c
+receiveSTM (ReceivePortBiased ps) =
+    foldr1 orElse (map receiveSTM ps)
+receiveSTM (ReceivePortRR psVar) = do
+    ps <- readTVar psVar
+    a  <- foldr1 orElse (map receiveSTM ps)
+    writeTVar psVar (rotate ps)
+    return a
+  where
     rotate :: [a] -> [a]
     rotate []     = []
     rotate (x:xs) = xs ++ [x]
@@ -459,7 +472,7 @@ finally a sequel = bracket_ (return ()) sequel a
 
 -- | Like 'expect' but with a timeout
 expectTimeout :: forall a. Serializable a => Int -> Process (Maybe a)
-expectTimeout timeout = receiveTimeout timeout [match return] 
+expectTimeout n = receiveTimeout n [match return] 
 
 -- | Asynchronous version of 'spawn'
 -- 

distributed-process/tests/TestCH.hs

@@ -695,6 +695,47 @@ testMatchAny transport = do
 
   takeMVar clientDone
 
+-- Test 'receiveChanTimeout'
+testReceiveChanTimeout :: NT.Transport -> IO ()
+testReceiveChanTimeout transport = do
+  done <- newEmptyMVar
+  sendPort <- newEmptyMVar
+  
+  forkIO $ do
+    localNode <- newLocalNode transport initRemoteTable
+    runProcess localNode $ do
+      -- Create a typed channel
+      (sp, rp) <- newChan :: Process (SendPort Bool, ReceivePort Bool)
+      liftIO $ putMVar sendPort sp
+
+      -- Wait for a message with a delay. No message arrives, we should get Nothing after 1 second 
+      Nothing <- receiveChanTimeout 1000000 rp
+
+      -- Wait for a message with a delay again. Now a message arrives after 0.5 seconds
+      Just True <- receiveChanTimeout 1000000 rp
+
+      -- Wait for a message with zero timeout: non-blocking check. No message is available, we get Nothing
+      Nothing <- receiveChanTimeout 0 rp
+
+      -- Again, but now there is a message available
+      liftIO $ threadDelay 1000000
+      Just False <- receiveChanTimeout 0 rp 
+
+      liftIO $ putMVar done ()
+
+  forkIO $ do
+    localNode <- newLocalNode transport initRemoteTable
+    runProcess localNode $ do
+      sp <- liftIO $ readMVar sendPort
+
+      liftIO $ threadDelay 1500000
+      sendChan sp True
+
+      liftIO $ threadDelay 500000
+      sendChan sp False
+
+  takeMVar done
+
 main :: IO ()
 main = do
   Right (transport, transportInternals) <- createTransportExposeInternals "127.0.0.1" "8080" defaultTCPParameters
@@ -739,4 +780,6 @@ main = do
     , ("MonitorChannel",               testMonitorChannel             transport)
       -- Reconnect
     , ("Reconnect",                    testReconnect                  transport transportInternals)
+      -- ReceiveChanTimeout
+    , ("ReceiveChanTimeout",           testReceiveChanTimeout         transport)
     ]