2 files changed, 130 insertions, 79 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
+
+
diff --git a/xmppServer.hs b/xmppServer.hs
index 4fbb775b..987e7dbc 100644
--- a/xmppServer.hs
+++ b/xmppServer.hs
@@ -27,14 +27,12 @@ import Data.Maybe (catMaybes,fromJust)
 import Data.Monoid ( (<>) )
 import Data.Text (Text)
 import qualified Data.Text as Text (pack)
-import qualified Data.Map as Map
 
 import qualified Control.Concurrent.STM.UpdateStream as Slotted
 import ControlMaybe
 import Nesting
 import EventUtil
 import Server
-import Data.Time.Clock (UTCTime,getCurrentTime)
 
 addrToText :: SockAddr -> Text
 addrToText (addr@(SockAddrInet _ _)) = Text.pack $ stripColon (show addr)
@@ -288,66 +286,49 @@ forkConnection k pingflag src snk stanzas = do
         wlog $ "end reader fork: " ++ show k
     return output
 
+data ConnectionKey
+       = PeerKey { callBackAddress :: SockAddr }
+       | ClientKey { localAddress :: SockAddr }
+ deriving (Show, Ord, Eq)
+
+{-
+data Peer = Peer
+    { peerWanted :: TVar Bool -- ^ False when this peer is on a you-call-me basis
+    , peerState :: TVar PeerState
+    }
+data PeerState
+    = PeerPendingConnect UTCTime
+    | PeerPendingAccept UTCTime
+    | PeerConnected (TChan Stanza)
+-}
+
+peerKey (sock,addr) = do
+    peer <- 
+      sIsConnected sock >>= \c -> 
+        if c then getPeerName sock -- addr is normally socketName
+             else return addr      -- Weird hack: addr is would-be peer name
+    return $ PeerKey (peer `withPort` fromIntegral peerport)
+
+clientKey (sock,addr) = return $ ClientKey addr
+
 monitor sv params = do
     chan <- return $ serverEvent sv
     stanzas <- atomically newTChan
-    peersVar <- atomically $ newTVar Map.empty
-    let doConnect utc k = do
-                  peers <- readTVar peersVar
-                  let mb = Map.lookup k peers
-                  maybe (do false <- newTVar False
-                            pending <- newTVar (PeerPendingConnect utc)
-                            let v = Peer { peerWanted = false
-                                         , peerState  = pending }
-                            writeTVar (peersVar) $ Map.insert k v peers)
-                        (\peer -> 
-                            writeTVar (peerState peer)
-                                $ PeerPendingConnect utc)
-                        mb
     fix $ \loop -> do
     action <- atomically $ foldr1 orElse
       [ readTChan chan >>= \(k,e) -> return $ do 
         case e of
             Connection pingflag conread conwrite -> do
-                wlog $ tomsg k "Connection"
-                let (xsrc,xsnk) = xmlStream conread conwrite
-                outs <- forkConnection k pingflag xsrc xsnk stanzas
-                atomically $ do
-                  peers <- readTVar peersVar
-                  let mb = Map.lookup k peers
-                  maybe (do false <- newTVar True -- False -- TODO: should be False
-                            connected <- newTVar (PeerConnected outs)
-                            let v = Peer { peerWanted = false
-                                         , peerState  = connected }
-                            writeTVar (peersVar) $ Map.insert k v peers)
-                          (\peer -> do
-                              writeTVar (peerWanted peer) True -- TODO REMOVE
-                              writeTVar (peerState peer)
-                                  $ PeerConnected outs)
-                          mb
-                return ()
+                 wlog $ tomsg k "Connection"
+                 let (xsrc,xsnk) = xmlStream conread conwrite
+                 forkConnection k pingflag xsrc xsnk stanzas
+                 return ()
             ConnectFailure addr -> do
                 wlog $ tomsg k "ConnectFailure"
-                action <- atomically $ do
-                  peers <- readTVar peersVar
-                  let mb = Map.lookup k peers
-                  maybe (return $ return ())
-                        (\peer -> do
-                            wanted <- readTVar (peerWanted peer)
-                            if wanted then return $ do
-                                            utc <- getCurrentTime
-                                            control sv (Connect addr params)
-                                            wlog $ tomsg k "Retry"
-                                            atomically $ doConnect utc k
-                                      else return $ return ())
-                        mb
-                action
             EOF        -> wlog $ tomsg k "EOF"
             HalfConnection In -> do
                 wlog $ tomsg k "ReadOnly"
-                utc <- getCurrentTime
                 control sv (Connect (callBackAddress k) params)
-                atomically $ doConnect utc k
             HalfConnection Out -> wlog $ tomsg k "WriteOnly"
             RequiresPing -> wlog $ tomsg k "RequiresPing"
             _ -> return ()
@@ -363,29 +344,6 @@ monitor sv params = do
         where
             _ = str :: String
 
-data ConnectionKey
-       = PeerKey { callBackAddress :: SockAddr }
-       | ClientKey { localAddress :: SockAddr }
- deriving (Show, Ord, Eq)
-
-data Peer = Peer
-    { peerWanted :: TVar Bool -- ^ False when this peer is on a you-call-me basis
-    , peerState :: TVar PeerState
-    }
-data PeerState
-    = PeerPendingConnect UTCTime
-    | PeerPendingAccept UTCTime
-    | PeerConnected (TChan Stanza)
-
-peerKey (sock,addr) = do
-    peer <- 
-      sIsConnected sock >>= \c -> 
-        if c then getPeerName sock -- addr is normally socketName
-             else return addr      -- Weird hack: addr is would-be peer name
-    return $ PeerKey (peer `withPort` fromIntegral peerport)
-
-clientKey (sock,addr) = return $ ClientKey addr
-
 
 peerport = 5269
 clientport = 5222
@@ -404,7 +362,7 @@ main = runResourceT $ do
                 (Just testaddr0)
                 (Just "5269")
     putStrLn $ "Connecting to "++show testaddr
-    control sv (Connect testaddr peer_params)
+    control sv (ConnectWithEndlessRetry testaddr peer_params 2000)
     forkIO $ monitor sv peer_params
     control sv (Listen peerport peer_params)
     -- control sv (Listen clientport client_params)