Factored the cache and retry pattern into module SendMessage

3 files changed, 241 insertions, 75 deletions

diff --git a/Presence/SendMessage.hs b/Presence/SendMessage.hs
new file mode 100644
index 00000000..6e8ea2b9
--- /dev/null
+++ b/Presence/SendMessage.hs
@@ -0,0 +1,208 @@
+{-# LANGUAGE CPP #-}
+{-# LANGUAGE TypeFamilies #-}
+module SendMessage
+    ( sendMessage
+    , CommandCache(..)
+    , newOutgoingConnections
+    , OutgoingConnections
+    ) where
+
+import Control.Monad
+import Control.Monad.Fix
+import Control.Monad.IO.Class
+import Control.Monad.Trans.Maybe
+import Control.Concurrent.STM
+import Control.Concurrent (forkIO)
+import Control.Concurrent.Async (async,waitSTM)
+import Control.Exception as E
+    ( bracketOnError 
+    , finally 
+    )
+import GHC.Conc 
+    ( threadStatus
+    , ThreadStatus(..)
+    , ThreadId
+    )
+import Data.IORef
+import qualified Data.Map as Map
+import           Data.Map as Map (Map)
+import Data.XML.Types as XML (Event)
+import System.IO
+    ( BufferMode(..)
+    , IOMode(..)
+    , hSetBuffering
+    )
+import Network.BSD 
+    ( getProtocolNumber
+    )
+import Network.Socket 
+    ( connect
+    , socketToHandle
+    , sClose
+    , Socket(..)
+    , socket
+    , SocketType(..)
+    )
+
+import XMPPTypes (Peer(..),discardPort,peerAddr,withPort,socketFamily)
+import SocketLike
+import ServerC (packetSink)
+import ControlMaybe (handleIO,handleIO_)
+import Data.Conduit (Sink,Source)
+import qualified Data.ByteString as S (ByteString)
+import XMLToByteStrings
+
+type ByteStringSink = Sink S.ByteString IO ()
+
+-- |Strict version of 'modifyIORef'
+modifyIORef' :: IORef a -> (a -> a) -> IO ()
+modifyIORef' ref f = do
+    x <- readIORef ref
+    let x' = f x
+    x' `seq` writeIORef ref x'
+
+
+class CommandCache cache where
+    type CacheableCommand cache
+
+    emptyCache :: cache
+
+    updateCache :: CacheableCommand cache -> cache -> cache
+
+
+data OutgoingConnections cache = 
+        OutgoingConnections (TVar (Map Peer (TChan (CacheableCommand cache), ThreadId)))
+                            (OutBoundXML RestrictedSocket cache (CacheableCommand cache))                        
+
+
+newOutgoingConnections ::
+  ( RestrictedSocket
+       -> cache
+       -> TChan (CacheableCommand cache)
+       -> (Maybe (CacheableCommand cache) -> IO ())
+       -> Source IO [XML.Event] )
+
+  -> STM (OutgoingConnections cache)
+
+newOutgoingConnections interpretCommands = do
+    remotes <- newTVar (Map.empty)
+    return (OutgoingConnections remotes interpretCommands)
+    
+
+
+sendMessage
+  :: CommandCache a =>
+     OutgoingConnections a -> CacheableCommand a -> Peer -> IO ()
+sendMessage (OutgoingConnections cons interpretCommands) msg peer0 = do
+    let peer = discardPort peer0
+    found <- atomically $ do
+        consmap <- readTVar cons
+        return (Map.lookup peer consmap)
+    let newEntry = do
+          chan <- atomically newTChan
+          t <- forkIO $ connect_to_server chan peer interpretCommands -- (newCache,updateCache,toPeer)
+          -- L.putStrLn $ "remote-map new: " <++> showPeer peer
+          return (True,(chan,t))
+    (is_new,entry) <- maybe newEntry
+          ( \(chan,t) -> do 
+                        st <- threadStatus t
+                        let running = do
+                                -- L.putStrLn $ "remote-map, thread running: " <++> showPeer peer
+                                return (False,(chan,t))
+                            died = do
+                                -- L.putStrLn $ "remote-map, thread died("<++>bshow st<++>"): " <++> showPeer peer
+                                newEntry 
+                        case st of 
+                          ThreadRunning   -> running
+                          ThreadBlocked _ -> running
+                          ThreadDied      -> died
+                          ThreadFinished  -> died
+                        )
+          found
+    -- L.putStrLn $ "sendMessage ->"<++>showPeer peer<++>": "<++>bshow msg
+    atomically $ writeTChan (fst entry) msg
+    when is_new . atomically $
+        readTVar cons >>= writeTVar cons . Map.insert peer entry
+
+
+connect_to_server chan peer toPeer = (>> return ()) . runMaybeT $ do
+    let port = 5269 :: Int
+    -- We'll cache Presence notifications until the socket
+    -- is ready.
+    cached <- liftIO $ newIORef emptyCache
+
+    let cacheCmd msg cached = modifyIORef' cached (updateCache msg)
+
+    fix $ \sendmsgs -> do
+        connected <- liftIO . async $ connect' (peerAddr peer) port
+
+        sock <- MaybeT . fix $ \loop -> do
+            e <- atomically $ orElse
+                    (fmap Right $ waitSTM connected)
+                    (fmap Left  $ readTChan chan)
+            case e of
+                Left cmd   -> cacheCmd cmd cached >> loop
+                Right sock -> return sock
+
+        retry <- do
+            (cache,snk) <- liftIO $ do
+                h <- socketToHandle sock ReadWriteMode
+                hSetBuffering h NoBuffering
+                cache <- readIORef $ cached
+                -- hint garbage collector: we're done with this...
+                writeIORef cached emptyCache 
+                return (cache,packetSink h)
+            MaybeT $ handleOutgoingToPeer toPeer (restrictSocket sock) cache chan snk
+
+        liftIO $ cacheCmd retry cached
+        -- liftIO $ putStrLn $ "retrying " ++ show retry
+        sendmsgs
+
+connect' :: SockAddr -> Int -> IO (Maybe Socket)
+connect' addr port = do
+    proto <- getProtocolNumber "tcp"
+    let getport (SockAddrInet port _) = port
+        getport (SockAddrInet6 port _ _ _) = port
+    let doException e = do
+            -- L.putStrLn $  "\nFailed to reach "<++> showPeer (RemotePeer addr) <++> " on port "<++>bshow port<++>": " <++> bshow e
+            return Nothing
+    handleIO doException
+         $ tryToConnect proto (addr `withPort` port)
+  where
+  tryToConnect proto addr =
+    bracketOnError
+        (socket (socketFamily addr) Stream proto)
+        (sClose )  -- only done if there's an error
+        (\sock -> do
+          connect sock addr
+          return (Just sock) -- socketToHandle sock ReadWriteMode
+        )
+
+
+type OutBoundXML sock cache msg =
+         sock
+         -> cache
+         -> TChan msg
+         -> (Maybe msg -> IO ())
+         -> Source IO [XML.Event]
+
+handleOutgoingToPeer
+  :: SocketLike sock =>
+     OutBoundXML sock cache msg
+     -> sock
+     -> cache 
+     -> TChan msg
+     -> ByteStringSink
+     -> IO (Maybe msg)
+handleOutgoingToPeer toPeer sock cache chan snk = do
+    p <- getPeerName sock
+    -- L.putStrLn $ "(>P) connected " <++> showPeer (RemotePeer p)
+    failed <- newIORef Nothing
+    let failure cmd = do
+            writeIORef failed cmd
+            -- putStrLn $ "Failed: " ++ show cmd
+    finally (
+        handleIO_ (return ()) $ toPeer sock cache chan failure `xmlToByteStrings` snk
+        ) $ return () -- logging L.putStrLn $ "(>P) disconnected " <++> showPeer (RemotePeer p)
+    readIORef failed
+
diff --git a/Presence/XMPP.hs b/Presence/XMPP.hs
index a469c08e..339cc6f2 100644
--- a/Presence/XMPP.hs
+++ b/Presence/XMPP.hs
@@ -1,6 +1,7 @@
 {-# LANGUAGE OverloadedStrings   #-}
 {-# LANGUAGE FlexibleContexts    #-}
 {-# LANGUAGE ViewPatterns        #-}
+{-# LANGUAGE TypeFamilies        #-}
 module XMPP 
     ( module XMPPTypes
     , listenForXmppClients
@@ -9,6 +10,10 @@ module XMPP
     , seekRemotePeers
     , quitListening
     , OutBoundMessage(..)
+    , OutgoingConnections
+    , CachedMessages
+    , toPeer
+    , newOutgoingConnections
     , sendMessage
     )  where
 
@@ -17,14 +22,13 @@ import XMPPTypes
 import ByteStringOperators
 import ControlMaybe
 import XMLToByteStrings
-
+import SendMessage
 
 import Data.Maybe (catMaybes)
 import Data.HList
 import Network.Socket 
     ( Family
     , connect
-    , socketToHandle
     , sClose
     , Socket(..)
     , socket
@@ -33,11 +37,6 @@ import Network.Socket
 import Network.BSD 
     ( getProtocolNumber
     )
-import System.IO
-    ( BufferMode(..)
-    , IOMode(..)
-    , hSetBuffering
-    )
 import Control.Concurrent.STM
 import Data.Conduit
 import Data.ByteString (ByteString)
@@ -45,7 +44,6 @@ import qualified Data.ByteString.Lazy.Char8 as L
     ( putStrLn
     , fromChunks
     )
-import Control.Concurrent (forkIO)
 import Control.Concurrent.Async
 import Control.Exception as E
     ( handle
@@ -72,11 +70,6 @@ import           Data.Set as Set (Set,(\\))
 import qualified Data.Set as Set
 import qualified Data.Map as Map
 import           Data.Map as Map (Map)
-import GHC.Conc 
-    ( threadStatus
-    , ThreadStatus(..)
-    , ThreadId
-    )
 
 textToByteString x = L.fromChunks [S.encodeUtf8 x]
 
@@ -914,65 +907,28 @@ data CachedMessages = CachedMessages
     , probes :: Map JID (Set (Bool,JID)) -- False means solicitation rather than probe
     , approvals :: Map JID (Set (Bool,JID) ) -- False means rejection rather than approval
     }
-newCache = CachedMessages Map.empty Map.empty Map.empty
-
-connect_to_server chan peer = (>> return ()) . runMaybeT $ do
-    let port = 5269 :: Int
-    -- We'll cache Presence notifications until the socket
-    -- is ready.
-    cached <- liftIO $ newIORef newCache
-
-    let mmInsert val key mm = Map.alter f key mm
-            where 
-                f Nothing    = Just $ Set.singleton val
-                f (Just set) = Just $ Set.insert val set
-    let cacheCmd (OutBoundPresence (Presence jid Offline)) cached = do
-                    cache <- readIORef cached
-                    writeIORef cached (cache { presences=Map.delete jid . presences $ cache })
-        cacheCmd (OutBoundPresence p@(Presence jid st)) cached = do
-                    cache <- readIORef cached
-                    writeIORef cached (cache { presences=Map.insert jid st . presences $ cache })
-        cacheCmd (PresenceProbe from to) cached = do
-                    cache <- readIORef cached
-                    let probes' = mmInsert  (True,from) to $ probes cache
-                    writeIORef cached (cache { probes=probes' })
-        cacheCmd (Solicitation from to) cached = do
-                    cache <- readIORef cached
-                    let probes' = mmInsert (False,from) to $ probes cache
-                    writeIORef cached (cache { probes=probes' })
-        cacheCmd (Approval from to) cached = do
-                    cache <- readIORef cached
-                    let approvals' = mmInsert (True,from) to $ approvals cache
-                    writeIORef cached (cache { approvals=approvals' })
-        cacheCmd (Rejection from to) cached = do
-                    cache <- readIORef cached
-                    let approvals' = mmInsert (False,from) to $ approvals cache
-                    writeIORef cached (cache { approvals=approvals' })
-
-    fix $ \sendmsgs -> do
-        connected <- liftIO . async $ connect' (peerAddr peer) port
-
-        sock <- MaybeT . fix $ \loop -> do
-            e <- atomically $ orElse
-                    (fmap Right $ waitSTM connected)
-                    (fmap Left  $ readTChan chan)
-            case e of
-                Left cmd   -> cacheCmd cmd cached >> loop
-                Right sock -> return sock
-
-        retry <- do
-            (cache,snk) <- liftIO $ do
-                h <- socketToHandle sock ReadWriteMode
-                hSetBuffering h NoBuffering
-                cache <- readIORef $ cached
-                -- hint garbage collector: we're done with this...
-                writeIORef cached newCache
-                return (cache,packetSink h)
-            MaybeT $ handleOutgoingToPeer (restrictSocket sock) cache chan snk
-
-        liftIO $ cacheCmd retry cached
-        liftIO $ putStrLn $ "retrying " ++ show retry
-        sendmsgs
+
+instance CommandCache CachedMessages where
+    type CacheableCommand CachedMessages = OutBoundMessage
+    emptyCache = CachedMessages Map.empty Map.empty Map.empty
+
+    updateCache (OutBoundPresence (Presence jid Offline)) cache = 
+                cache { presences=Map.delete jid . presences $ cache }
+    updateCache (OutBoundPresence p@(Presence jid st)) cache =
+                cache { presences=Map.insert jid st . presences $ cache }
+    updateCache (PresenceProbe from to) cache =
+                cache { probes = mmInsert  (True,from) to $ probes cache }
+    updateCache (Solicitation from to) cache =
+                cache { probes= mmInsert (False,from) to $ probes cache }
+    updateCache (Approval from to) cache =
+                cache { approvals= mmInsert (True,from) to $ approvals cache }
+    updateCache (Rejection from to) cache =
+                cache { approvals= mmInsert (False,from) to $ approvals cache }
+
+mmInsert val key mm = Map.alter f key mm
+ where 
+    f Nothing    = Just $ Set.singleton val
+    f (Just set) = Just $ Set.insert val set
 
 
 greetPeer =
@@ -1139,6 +1095,7 @@ connect' addr port = do
         )
 
 
+{-
 sendMessage  :: TVar (Map Peer (TChan OutBoundMessage, ThreadId)) -> OutBoundMessage -> Peer -> IO ()
 sendMessage cons msg peer0 = do
     let peer = discardPort peer0
@@ -1170,11 +1127,12 @@ sendMessage cons msg peer0 = do
     atomically $ writeTChan (fst entry) msg
     when is_new . atomically $
         readTVar cons >>= writeTVar cons . Map.insert peer entry
+-}
 
 
 
 seekRemotePeers :: JabberPeerSession config =>
-     XMPPPeerClass config -> TChan Presence -> TVar (Map Peer (TChan OutBoundMessage, ThreadId)) -> IO b0
+     XMPPPeerClass config -> TChan Presence -> OutgoingConnections CachedMessages -> IO b0
 seekRemotePeers config chan server_connections = do
     fix $ \loop -> do
         event <- atomically $ readTChan chan
diff --git a/Presence/main.hs b/Presence/main.hs
index 317dc71e..3d5ac046 100644
--- a/Presence/main.hs
+++ b/Presence/main.hs
@@ -89,7 +89,7 @@ data PresenceState = PresenceState
     , remoteUsers :: TVar (Map Peer (RefCount,TVar (MultiMap JabberUser (JabberResource,JabberShow))))
 
     -- outGoingConnections  - a set of channels that may be used to send messages to remote peers.
-    , outGoingConnections :: TVar (Map Peer (TChan OutBoundMessage, ThreadId)) 
+    , outGoingConnections :: OutgoingConnections CachedMessages
     }
 
 
@@ -108,7 +108,7 @@ newPresenceState hostname = atomically $ do
     locals_greedy <- newEmptyTMVar
     rchan <- newEmptyTMVar
     remotes <- newTVar (Map.empty)
-    server_connections <- newServerConnections
+    server_connections <- newOutgoingConnections toPeer
     return $ PresenceState hostname tty us subs locals_greedy rchan remotes server_connections