moved ping logic into ConnectionState

1 files changed, 42 insertions, 20 deletions

diff --git a/Presence/Server.hs b/Presence/Server.hs
index c7162265..530b6669 100644
--- a/Presence/Server.hs
+++ b/Presence/Server.hs
@@ -311,7 +311,17 @@ killListener (thread,sock) = do sClose sock
 
 newConnection server params conkey h inout = do
     hSetBuffering h NoBuffering
-    new <- connectionThreads h
+    let (forward,idle_ms,timeout_ms) =
+            case (inout,duplex params) of
+                (Out,True) -> ( const $ return ()
+                              , 0
+                              , 0 )
+                _ -> ( announce . (conkey,) . Got
+                     , pingInterval params
+                     , timeout params )
+
+    new <- do pinglogic <- pingMachine idle_ms timeout_ms
+              connectionThreads h pinglogic
     started <- atomically $ newEmptyTMVar
     kontvar <- atomically newEmptyTMVar
     forkIO $ do
@@ -319,7 +329,6 @@ newConnection server params conkey h inout = do
         kont <- atomically getkont
         kont
     
-    pinglogic <- pingMachine (pingInterval params) (timeout params)
     atomically $ do
         current <- fmap (Map.lookup conkey) $ readTVar (conmap server)
         case current of
@@ -335,25 +344,21 @@ newConnection server params conkey h inout = do
             announce e
             putTMVar kontvar $ return $ do
                 atomically $ putTMVar started ()
-                handleEOF conkey kontvar newCon pinglogic
+                handleEOF conkey kontvar newCon
           Just what@(mvar,_) -> do
             putTMVar kontvar $ return $ return ()
             putTMVar mvar $ do
-                kont <- updateConMap conkey new what pinglogic
+                kont <- updateConMap conkey new what
                 putTMVar started ()
                 return kont
     forkIO $ do -- inout==In || duplex params then forkIO $ do
         -- warn $ "waiting read thread: " <> bshow (conkey,inout)
-        let forward =
-                case (inout,duplex params) of
-                    (Out,True) -> const $ return ()
-                    _ -> announce . (conkey,) . Got
         atomically $ takeTMVar started
-        pingBump pinglogic -- start the ping timer
+        -- pingBump pinglogic -- start the ping timer
         fix $ \loop -> do
         -- warn $ "read thread: " <> bshow (conkey,inout)
         mb <- threadsRead new
-        pingBump pinglogic
+        -- pingBump pinglogic
         -- warn $ "got: " <> bshow (mb,(conkey,inout))
         maybe (return ())
               (atomically . forward >=> const loop)
@@ -363,17 +368,18 @@ newConnection server params conkey h inout = do
 
   announce e = writeTChan (serverEvent server) e
 
-  handleEOF conkey mvar newCon pingTimer  = do
+  handleEOF conkey mvar newCon = do
     action <- atomically . foldr1 orElse $
         [ takeTMVar mvar >>= id -- passed continuation
         , connWait newCon     >> return eof
-        , pingWait pingTimer >>= return . sendPing
+        , connWaitPing newCon >>= return . sendPing
+        -- , pingWait pingTimer >>= return . sendPing
         ]
     action :: IO ()
    where
     eof = do
             -- warn $ "EOF " <>bshow conkey
-            pingCancel pingTimer -- force cleanup of ping thread
+            connCancelPing newCon
             atomically $ do connFlush newCon
                             announce (conkey,EOF)
                             modifyTVar' (conmap server)
@@ -384,10 +390,10 @@ newConnection server params conkey h inout = do
                               eof
     sendPing PingIdle = do
         atomically . announce $ (conkey,RequiresPing)
-        handleEOF conkey mvar newCon pingTimer
+        handleEOF conkey mvar newCon
 
 
-  updateConMap conkey new (mvar,replaced) pingTimer = do
+  updateConMap conkey new (mvar,replaced) = do
     new' <-
         if duplex params then do
             announce (conkey,EOF)
@@ -410,7 +416,7 @@ newConnection server params conkey h inout = do
                                 In  -> ReadOnlyConnection new
                                 Out -> WriteOnlyConnection new
     modifyTVar' (conmap server) $ Map.insert conkey (mvar,new')
-    return $ handleEOF conkey mvar new' pingTimer
+    return $ handleEOF conkey mvar new'
 
 acceptLoop server params sock = handle (acceptException server params sock) $ do
     con <- accept sock
@@ -445,12 +451,13 @@ data ConnectionThreads = ConnectionThreads
     { threadsWriter :: TMVar (Maybe ByteString)
     , threadsChannel :: TChan ByteString
     , threadsWait :: STM () -- ^ waits for a 'ConnectionThreads' object to close
+    , threadsPing :: PingMachine
     }
 
 -- | This spawns the reader and writer threads and returns a newly
 -- constructed 'ConnectionThreads' object.
-connectionThreads :: Handle -> IO ConnectionThreads
-connectionThreads h = do
+connectionThreads :: Handle -> PingMachine -> IO ConnectionThreads
+connectionThreads h pinglogic = do
 
     (donew,outs) <- atomically $ liftM2 (,) newEmptyTMVar newEmptyTMVar
     writerThread <- forkIO . fix $ \loop -> do
@@ -471,9 +478,12 @@ connectionThreads h = do
                           let _ = fmap ioeGetErrorType e -- type hint
                           atomically $ do putTMVar outs Nothing -- quit writer
                                           putTMVar doner ()
-        handle (finished . Just) . fix $ \loop -> do
+        handle (finished . Just) $ do
+            pingBump pinglogic -- start the ping timer
+            fix $ \loop -> do
             packet <- getPacket h
             atomically $ writeTChan incomming packet
+            pingBump pinglogic
             isEof <- liftIO $ hIsEOF h
             if isEof then finished Nothing else loop
 
@@ -482,7 +492,8 @@ connectionThreads h = do
                   return ()
     return ConnectionThreads { threadsWriter = outs
                              , threadsChannel = incomming
-                             , threadsWait = wait }
+                             , threadsWait = wait
+                             , threadsPing = pinglogic }
 
 
 -- | 'threadsWrite' writes the given 'ByteString' to the
@@ -555,6 +566,17 @@ connClose c = mapConn True threadsClose c
 connWait :: ConnectionState -> STM ()
 connWait c = mapConn False threadsWait c
 
+connPingTimer c =
+    case c of
+        SaneConnection rw     -> threadsPing rw
+        ReadOnlyConnection r  -> threadsPing r
+        WriteOnlyConnection w -> threadsPing w -- should be disabled.
+        ConnectionPair r w    -> threadsPing r
+
+connCancelPing c = pingCancel (connPingTimer c)
+connWaitPing c = pingWait (connPingTimer c)
+
+
 connFlush c =
     case c of
         SaneConnection rw     -> waitChan rw