ConnectWithEndlessRetry

1 files changed, 100 insertions, 7 deletions

diff --git a/Presence/Server.hs b/Presence/Server.hs
index 5e54ff27..1141f0b0 100644
--- a/Presence/Server.hs
+++ b/Presence/Server.hs
@@ -60,14 +60,16 @@ import Network.BSD
     ( getProtocolNumber
     )   
 import Debug.Trace
-import Data.Time.Clock (UTCTime,getCurrentTime)
+import Data.Time.Clock (UTCTime,getCurrentTime,diffUTCTime)
 import Data.Time.Format (formatTime)
 import System.Locale (defaultTimeLocale)
 
 todo = error "unimplemented"
 
-type TimeOut = Int      -- ^ miliseconds
-type PingInterval = Int -- ^ miliseconds
+type Microseconds = Int
+type Miliseconds = Int
+type TimeOut = Miliseconds
+type PingInterval = Miliseconds
 
 -- | This object is passed with the 'Listen' and 'Connect'
 -- instructions in order to control the behavior of the 
@@ -128,6 +130,10 @@ data ServerInstruction conkey
     -- ^ listen for incomming connections
     | Connect SockAddr  (ConnectionParameters conkey)
     -- ^ connect to addresses
+    | ConnectWithEndlessRetry SockAddr
+                              (ConnectionParameters conkey)
+                              Miliseconds
+    -- ^ keep retrying the connection
     | Ignore PortNumber
     -- ^ stop listening on specified port
     | Send conkey ByteString
@@ -178,6 +184,7 @@ data Server a
           , serverReleaseKey :: ReleaseKey
           , conmap :: TVar (Map a (TMVar (STM (IO ())), ConnectionState))
           , listenmap :: TVar (Map PortNumber (ThreadId,Socket))
+          , retrymap :: TVar (Map SockAddr (TVar Bool,InterruptableDelay))
           }
 
 -- | Construct a 'Server' object.  Use 'Control.Monad.Trans.Resource.ResourceT'
@@ -210,11 +217,13 @@ server = do
         tchan    <- newTChan
         conmap   <- newTVar Map.empty
         listenmap<- newTVar Map.empty
+        retrymap <- newTVar Map.empty
         return Server { serverCommand = cmds
                       , serverEvent = tchan
                       , serverReleaseKey = key
                       , conmap = conmap
                       , listenmap = listenmap
+                      , retrymap = retrymap
                       }
     liftIO $ do
         forkIO $ fix $ \loop -> do
@@ -230,6 +239,10 @@ server = do
   closeAll server = liftIO $ do
     listening <- atomically . readTVar $ listenmap server
     mapM_ killListener   (Map.elems listening)
+    let stopRetry (v,d) = do atomically $ writeTVar v False
+                             interruptDelay d
+    retriers <- atomically . readTVar $ retrymap server
+    mapM_ stopRetry (Map.elems retriers)
     cons <- atomically . readTVar $ conmap server
     atomically $ mapM_ (connClose . snd) (Map.elems cons)
     atomically $ mapM_ (connWait . snd) (Map.elems cons)
@@ -278,7 +291,7 @@ server = do
   doit server (Connect addr params) = liftIO $ do
     void . forkIO $ do
         proto <- getProtocolNumber "tcp"
-        bracketOnError
+        sock <- bracketOnError
             (socket (socketFamily addr) Stream proto)
             (\sock -> do -- only done if there's an error
                          -- Weird hack: puting the would-be peer address
@@ -288,12 +301,66 @@ server = do
                          atomically 
                             $ writeTChan (serverEvent server) 
                                 $ (conkey,ConnectFailure addr))
-            $ \sock -> do
-        connect sock addr
+            $ \sock -> do connect sock addr
+                          return sock
         me <- getSocketName sock
         conkey <- makeConnKey params (sock,me)
         h <- socketToHandle sock ReadWriteMode
         newConnection server params conkey h Out
+        return ()
+
+  doit server (ConnectWithEndlessRetry addr params interval) = liftIO $ do
+    proto <- getProtocolNumber "tcp"
+    void . forkIO $ do
+        resultVar <- atomically newEmptyTMVar
+        timer <- interruptableDelay
+        (retryVar,action) <- atomically $ do
+            let noop = return ()
+            map <- readTVar (retrymap server)
+            let mb = Map.lookup addr map
+            action <-
+                maybe (return noop)
+                      (\(v,d) -> do writeTVar v False
+                                    return $ interruptDelay d)
+                      mb
+            v <- newTVar True
+            writeTVar (retrymap server) (Map.insert addr (v,timer) map)
+            return (v,action)
+        action
+        fix $ \retryLoop -> do
+        utc <- getCurrentTime
+        forkIO $ do
+            sock <- bracketOnError
+                (socket (socketFamily addr) Stream proto)
+                (\sock -> do -- only done if there's an error
+                             -- Weird hack: puting the would-be peer address
+                             -- instead of local socketName
+                             conkey <- makeConnKey params (sock,addr) -- XXX: ?
+                             sClose sock
+                             atomically $ do
+                                writeTChan (serverEvent server) 
+                                    $ (conkey,ConnectFailure addr)
+                                retry <- readTVar retryVar
+                                putTMVar resultVar retry)
+                $ \sock -> do connect sock addr
+                              return sock
+            me <- getSocketName sock
+            conkey <- makeConnKey params (sock,me)
+            h <- socketToHandle sock ReadWriteMode
+            threads <- newConnection server params conkey h Out
+            atomically $ do threadsWait threads
+                            retry <- readTVar retryVar
+                            putTMVar resultVar retry
+        retry <- atomically $ takeTMVar resultVar
+        fin_utc <- getCurrentTime
+        when retry $ do
+        let elapsed = 1000.0 * (fin_utc `diffUTCTime` utc)
+            expected = fromIntegral interval
+        when (retry && elapsed < expected) $ do
+            warn $ "Waiting to retry " <> bshow addr
+            startDelay timer (round $ 1000 * (expected-elapsed))
+        warn $ "retry " <> bshow (retry,addr)
+        when retry $ retryLoop
 
 
 -- INTERNAL ----------------------------------------------------------
@@ -397,7 +464,7 @@ newConnection server params conkey h inout = do
               mb
         warn $ "quit-gots: " <> bshow (conkey,inout)
 #endif
-    return ()
+    return new
  where
 
   announce e = writeTChan (serverEvent server) e
@@ -754,3 +821,29 @@ pingWait me = do
     return e
 
 
+data InterruptableDelay = InterruptableDelay 
+        { delayThread :: TMVar ThreadId
+        }
+
+interruptableDelay :: IO InterruptableDelay
+interruptableDelay = do
+    fmap InterruptableDelay
+       $ atomically newEmptyTMVar
+startDelay :: InterruptableDelay -> Microseconds -> IO ()
+startDelay d interval = do
+    thread <- myThreadId
+    atomically $ putTMVar (delayThread d) thread
+    handle (\(ErrorCall _)-> warn $ "delay interrupted" ) $ do 
+        threadDelay interval
+    void . atomically $ takeTMVar (delayThread d)
+
+interruptDelay :: InterruptableDelay -> IO ()
+interruptDelay d = do
+    mthread <- atomically $ do
+        tryTakeTMVar (delayThread d)
+    flip (maybe $ return ()) mthread $ \thread -> do
+    throwTo thread (ErrorCall "Interrupted delay")
+    atomically $ do
+        putTMVar (delayThread d) thread
+
+