TCP-based Network.QueryResponse.Transport.

1 files changed, 128 insertions, 0 deletions

diff --git a/src/Network/QueryResponse/TCP.hs b/src/Network/QueryResponse/TCP.hs
new file mode 100644
index 00000000..83ae367f
--- /dev/null
+++ b/src/Network/QueryResponse/TCP.hs
@@ -0,0 +1,128 @@
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE CPP #-}
+module Network.QueryResponse.TCP where
+
+#ifdef THREAD_DEBUG
+import Control.Concurrent.Lifted.Instrument
+#else
+import Control.Concurrent.Lifted
+import GHC.Conc                  (labelThread)
+#endif
+
+import Control.Concurrent.STM
+import Control.Monad
+import Data.ByteString (ByteString,hPut)
+import Data.Function
+import Data.Hashable
+import Data.Time.Clock.POSIX
+import Data.Word
+import Network.BSD
+import Network.Socket
+import System.IO
+import System.IO.Error
+
+import Connection.Tcp (socketFamily)
+import qualified Data.MinMaxPSQ as MM
+import Network.QueryResponse
+
+data TCPSession st = TCPSession
+    { tcpHandle :: MVar Handle
+    , tcpState  :: st
+    , tcpThread :: ThreadId
+    }
+
+newtype TCPAddress = TCPAddress SockAddr
+ deriving (Eq,Ord)
+
+instance Hashable TCPAddress where
+    hashWithSalt salt (TCPAddress x) = case x of
+        SockAddrInet port addr   -> hashWithSalt salt (fromIntegral port :: Word16,addr)
+        SockAddrInet6 port b c d -> hashWithSalt salt (fromIntegral port :: Word16,b,c,d)
+        _                        -> 0
+
+data TCPCache st = TCPCache
+    { lru    :: TVar (MM.MinMaxPSQ' TCPAddress POSIXTime (TCPSession st))
+    , tcpMax :: Int
+    }
+
+data StreamTransform st x y = StreamTransform
+    { streamHello :: Handle -> IO st               -- ^ "Hello" protocol upon fresh connection.
+    , streamGoodbye :: st -> Handle -> IO ()       -- ^ "Goodbye" protocol upon termination.
+    , streamDecode :: st -> Handle -> IO (Maybe x) -- ^ Parse inbound messages.
+    , streamEncode :: st -> y -> IO ByteString     -- ^ Serialize outbound messages.
+    }
+
+acquireConnection :: MVar (Maybe (Either a (x, SockAddr)))
+                           -> TCPCache st
+                           -> StreamTransform st x y
+                           -> SockAddr
+                           -> IO (Maybe (y -> IO ()))
+acquireConnection mvar tcpcache stream addr = do
+    cache <- atomically $ readTVar (lru tcpcache)
+    case MM.lookup' (TCPAddress addr) cache of
+        Nothing -> do
+            proto <- getProtocolNumber "tcp"
+            sock <- socket (socketFamily addr) Stream proto
+            connect sock addr `catchIOError` (\e -> close sock)
+            h <- socketToHandle sock ReadWriteMode
+            st <- streamHello stream h
+            t <- getPOSIXTime
+            mh <- newMVar h
+            rthread <- forkIO $ fix $ \loop -> do
+                x <- streamDecode stream st h
+                putMVar mvar $ fmap (\u -> Right (u, addr)) x
+                case x of
+                    Just _ -> loop
+                    Nothing -> do
+                        atomically $ modifyTVar' (lru tcpcache) $ MM.delete (TCPAddress addr)
+                        hClose h
+            labelThread rthread ("tcp:"++show addr)
+            let v = TCPSession
+                    { tcpHandle = mh
+                    , tcpState  = st
+                    , tcpThread = rthread
+                    }
+            let (retires,cache') = MM.takeView (tcpMax tcpcache) $ MM.insert' (TCPAddress addr) v t cache
+            forM_ retires $ \(MM.Binding (TCPAddress k) r _) -> void $ forkIO $ do
+                myThreadId >>= flip labelThread ("tcp-close:"++show k)
+                killThread (tcpThread r)
+                h <- takeMVar (tcpHandle r)
+                streamGoodbye stream st h
+                hClose h
+            atomically $ writeTVar (lru tcpcache) cache'
+
+            return $ Just $ \y -> do
+                                bs <- streamEncode stream st y
+                                withMVar mh (`hPut` bs)
+        Just (tm,v) -> do
+            t <- getPOSIXTime
+            let TCPSession { tcpHandle = mh, tcpState = st } = v
+                cache' = MM.insert' (TCPAddress addr) v t cache
+            atomically $ writeTVar (lru tcpcache) cache'
+            return $ Just $ \y -> do
+                                bs <- streamEncode stream st y
+                                withMVar mh (`hPut` bs)
+
+closeAll :: TCPCache st -> StreamTransform st x y -> IO ()
+closeAll tcpcache stream = do
+    cache <- atomically $ readTVar (lru tcpcache)
+    forM_ (MM.toList cache) $ \(MM.Binding (TCPAddress addr) r tm) -> do
+        let st = tcpState r
+        killThread (tcpThread r)
+        h <- takeMVar $ tcpHandle r
+        streamGoodbye stream st h
+        hClose h
+
+tcpTransport :: Int -- ^ maximum number of TCP links to maintain.
+                -> StreamTransform st x y
+                -> IO (TransportA err SockAddr x y)
+tcpTransport maxcon stream = do
+    msgvar <- newEmptyMVar
+    tcpcache <- atomically $ (`TCPCache` maxcon) <$> newTVar (MM.empty)
+    return Transport
+        { awaitMessage   = (takeMVar msgvar >>=)
+        , sendMessage    = \addr y -> do
+                                msock <- acquireConnection msgvar tcpcache stream addr
+                                mapM_ ($ y) msock
+        , closeTransport = closeAll tcpcache stream
+        }