1 files changed, 161 insertions, 0 deletions

diff --git a/dht/src/Control/Concurrent/PingMachine.hs b/dht/src/Control/Concurrent/PingMachine.hs
new file mode 100644
index 00000000..a8f10e83
--- /dev/null
+++ b/dht/src/Control/Concurrent/PingMachine.hs
@@ -0,0 +1,161 @@
+{-# LANGUAGE CPP #-}
+{-# LANGUAGE TupleSections #-}
+module Control.Concurrent.PingMachine where
+
+import Control.Monad
+import Data.Function
+#ifdef THREAD_DEBUG
+import Control.Concurrent.Lifted.Instrument
+#else
+import Control.Concurrent.Lifted
+import GHC.Conc                  (labelThread)
+#endif
+import Control.Concurrent.STM
+
+import Control.Concurrent.Delay
+
+type Miliseconds = Int
+type TimeOut = Miliseconds
+type PingInterval = Miliseconds
+
+-- | Events that occur as a result of the 'PingMachine' watchdog.
+--
+-- Use 'pingWait' to wait for one of these to occur.
+data PingEvent
+    = PingIdle    -- ^ You should send a ping if you observe this event.
+    | PingTimeOut -- ^ You should give up on the connection in case of this event.
+
+data PingMachine = PingMachine
+    { pingFlag          :: TVar Bool
+    , pingInterruptible :: InterruptibleDelay
+    , pingEvent         :: TMVar PingEvent
+    , pingStarted       :: TMVar Bool
+    }
+
+-- | Fork a thread to monitor a connection for a ping timeout.
+--
+-- If 'pingBump' is not invoked after a idle is signaled, a timeout event will
+-- occur.  When that happens, even if the caller chooses to ignore this event,
+-- the watchdog thread will be terminated and no more ping events will be
+-- signaled.
+--
+-- An idle connection will be signaled by:
+--
+--   (1) 'pingFlag' is set 'True'
+--
+--   (2) 'pingWait' returns 'PingIdle'
+--
+-- Either may be tested to determine whether a ping should be sent but
+-- 'pingFlag' is difficult to use properly because it is up to the caller to
+-- remember that the ping is already in progress.
+forkPingMachine
+    :: String
+    -> PingInterval -- ^ Milliseconds of idle before a ping is considered necessary.
+    -> TimeOut      -- ^ Milliseconds after 'PingIdle' before we signal 'PingTimeOut'.
+    -> IO PingMachine
+forkPingMachine label idle timeout = do
+    d <- interruptibleDelay
+    flag <- atomically $ newTVar False
+    canceled <- atomically $ newTVar False
+    event <- atomically newEmptyTMVar
+    started <- atomically $ newEmptyTMVar
+    when (idle/=0) $ void . forkIO $ do
+        myThreadId >>= flip labelThread ("Ping." ++ label) -- ("ping.watchdog")
+        (>>=) (atomically (readTMVar started)) $ flip when $ do
+        fix $ \loop -> do
+            atomically $ writeTVar flag False
+            fin <- startDelay d (1000*idle)
+            (>>=) (atomically (readTMVar started)) $ flip when $ do
+            if (not fin) then loop
+            else do
+            -- Idle event
+            atomically $ do
+                tryTakeTMVar event
+                putTMVar event PingIdle
+                writeTVar flag True
+            fin <- startDelay d (1000*timeout)
+            (>>=) (atomically (readTMVar started)) $ flip when $ do
+            me <- myThreadId
+            if (not fin) then loop
+            else do
+            -- Timeout event
+            atomically $ do
+                tryTakeTMVar event
+                writeTVar flag False
+                putTMVar event PingTimeOut
+    return PingMachine
+        { pingFlag          = flag
+        , pingInterruptible = d
+        , pingEvent         = event
+        , pingStarted       = started
+        }
+
+-- | like 'forkPingMachine'  but the timeout and idle parameters can be changed dynamically
+--  Unlike 'forkPingMachine', 'forkPingMachineDynamic' always launches a thread
+--  regardless of idle value.
+forkPingMachineDynamic
+    :: String
+    -> TVar PingInterval -- ^ Milliseconds of idle before a ping is considered necessary.
+    -> TVar TimeOut      -- ^ Milliseconds after 'PingIdle' before we signal 'PingTimeOut'.
+    -> IO PingMachine
+forkPingMachineDynamic label idleV timeoutV = do
+    d <- interruptibleDelay
+    flag <- atomically $ newTVar False
+    canceled <- atomically $ newTVar False
+    event <- atomically newEmptyTMVar
+    started <- atomically $ newEmptyTMVar
+    void . forkIO $ do
+        myThreadId >>= flip labelThread ("Ping." ++ label) -- ("ping.watchdog")
+        (>>=) (atomically (readTMVar started)) $ flip when $ do
+        fix $ \loop -> do
+            atomically $ writeTVar flag False
+            (idle,timeout) <- atomically $ (,) <$> readTVar idleV <*> readTVar timeoutV
+            fin <- startDelay d (1000*idle)
+            (>>=) (atomically (readTMVar started)) $ flip when $ do
+            if (not fin) then loop
+            else do
+            -- Idle event
+            atomically $ do
+                tryTakeTMVar event
+                putTMVar event PingIdle
+                writeTVar flag True
+            fin <- startDelay d (1000*timeout)
+            (>>=) (atomically (readTMVar started)) $ flip when $ do
+            me <- myThreadId
+            if (not fin) then loop
+            else do
+            -- Timeout event
+            atomically $ do
+                tryTakeTMVar event
+                writeTVar flag False
+                putTMVar event PingTimeOut
+    return PingMachine
+        { pingFlag          = flag
+        , pingInterruptible = d
+        , pingEvent         = event
+        , pingStarted       = started
+        }
+
+-- | Terminate the watchdog thread.  Call this upon connection close.
+--
+-- You should ensure no threads are waiting on 'pingWait' because there is no
+-- 'PingEvent' signaling termination.
+pingCancel :: PingMachine -> IO ()
+pingCancel me = do
+    atomically $ do tryTakeTMVar (pingStarted me)
+                    putTMVar (pingStarted me) False
+    interruptDelay (pingInterruptible me)
+
+-- | Reset the ping timer.  Call this regularly to prevent 'PingTimeOut'.
+pingBump :: PingMachine -> IO ()
+pingBump me = do
+    atomically $ do
+            b <- tryReadTMVar (pingStarted me)
+            when (b/=Just False) $ do
+                tryTakeTMVar (pingStarted me)
+                putTMVar (pingStarted me) True
+    interruptDelay (pingInterruptible me)
+
+-- | Retries until a 'PingEvent' occurs.
+pingWait :: PingMachine -> STM PingEvent
+pingWait me = takeTMVar (pingEvent me)