1 files changed, 32 insertions, 25 deletions

diff --git a/src/Network/QueryResponse/TCP.hs b/src/Network/QueryResponse/TCP.hs
index 83ae367f..7efe6966 100644
--- a/src/Network/QueryResponse/TCP.hs
+++ b/src/Network/QueryResponse/TCP.hs
@@ -45,77 +45,84 @@ data TCPCache st = TCPCache
     , 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.
+data SessionProtocol x y = SessionProtocol
+    { streamGoodbye :: Handle -> IO ()        -- ^ "Goodbye" protocol upon termination.
+    , streamDecode  :: Handle -> IO (Maybe x) -- ^ Parse inbound messages.
+    , streamEncode  :: y -> IO ByteString     -- ^ Serialize outbound messages.
     }
 
-acquireConnection :: MVar (Maybe (Either a (x, SockAddr)))
-                           -> TCPCache st
-                           -> StreamTransform st x y
-                           -> SockAddr
+data StreamHandshake addr x y = StreamHandshake
+    { streamHello :: addr -> Handle -> IO (SessionProtocol x y) -- ^ "Hello" protocol upon fresh connection.
+    , streamAddr  :: addr -> SockAddr
+    }
+
+acquireConnection :: MVar (Maybe (Either a (x, addr)))
+                           -> TCPCache (SessionProtocol x y)
+                           -> StreamHandshake addr x y
+                           -> addr
                            -> IO (Maybe (y -> IO ()))
 acquireConnection mvar tcpcache stream addr = do
     cache <- atomically $ readTVar (lru tcpcache)
-    case MM.lookup' (TCPAddress addr) cache of
+    case MM.lookup' (TCPAddress $ streamAddr stream addr) cache of
         Nothing -> do
             proto <- getProtocolNumber "tcp"
-            sock <- socket (socketFamily addr) Stream proto
-            connect sock addr `catchIOError` (\e -> close sock)
+            sock <- socket (socketFamily $ streamAddr stream addr) Stream proto
+            connect sock (streamAddr stream addr) `catchIOError` (\e -> close sock)
             h <- socketToHandle sock ReadWriteMode
-            st <- streamHello stream h
+            st <- streamHello stream addr h
             t <- getPOSIXTime
             mh <- newMVar h
             rthread <- forkIO $ fix $ \loop -> do
-                x <- streamDecode stream st h
+                x <- streamDecode 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)
+                        atomically $ modifyTVar' (lru tcpcache)
+                                   $ MM.delete (TCPAddress $ streamAddr stream addr)
                         hClose h
-            labelThread rthread ("tcp:"++show addr)
+            let showAddr a = show (streamAddr stream a)
+            labelThread rthread ("tcp:"++showAddr addr)
             let v = TCPSession
                     { tcpHandle = mh
                     , tcpState  = st
                     , tcpThread = rthread
                     }
-            let (retires,cache') = MM.takeView (tcpMax tcpcache) $ MM.insert' (TCPAddress addr) v t cache
+            let (retires,cache') = MM.takeView (tcpMax tcpcache)
+                                    $ MM.insert' (TCPAddress $ streamAddr stream 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
+                streamGoodbye st h
                 hClose h
             atomically $ writeTVar (lru tcpcache) cache'
 
             return $ Just $ \y -> do
-                                bs <- streamEncode stream st y
+                                bs <- streamEncode 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
+                cache' = MM.insert' (TCPAddress $ streamAddr stream addr) v t cache
             atomically $ writeTVar (lru tcpcache) cache'
             return $ Just $ \y -> do
-                                bs <- streamEncode stream st y
+                                bs <- streamEncode st y
                                 withMVar mh (`hPut` bs)
 
-closeAll :: TCPCache st -> StreamTransform st x y -> IO ()
+closeAll :: TCPCache (SessionProtocol x y) -> StreamHandshake addr 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
+        streamGoodbye st h
         hClose h
 
 tcpTransport :: Int -- ^ maximum number of TCP links to maintain.
-                -> StreamTransform st x y
-                -> IO (TransportA err SockAddr x y)
+                -> StreamHandshake addr x y
+                -> IO (TransportA err addr x y)
 tcpTransport maxcon stream = do
     msgvar <- newEmptyMVar
     tcpcache <- atomically $ (`TCPCache` maxcon) <$> newTVar (MM.empty)