Multiplexer: Implement channel management

qrpc.cabal

@@ -94,5 +94,6 @@ test-suite qrpc-test
   main-is: Spec.hs
   other-modules:
     Network.RpcSpec
+    Network.Rpc.MultiplexerSpec
   hs-source-dirs:
     test

src/Network/Rpc.hs

@@ -2,7 +2,7 @@ module Network.Rpc where
 
 import Control.Concurrent (forkFinally)
 import Control.Concurrent.Async (Async, async, link, withAsync)
-import Control.Exception (SomeException, bracket, bracketOnError, bracketOnError)
+import Control.Exception (SomeException, bracket, bracketOnError, bracketOnError, interruptible)
 import Control.Monad ((>=>), when, forever)
 import Control.Monad.State (State, execState)
 import qualified Control.Monad.State as State
@@ -12,14 +12,13 @@ import qualified Data.ByteString.Lazy as BSL
 import Data.Hashable (Hashable)
 import qualified Data.HashMap.Strict as HM
 import Data.Maybe (isNothing)
-import Language.Haskell.TH
+import Language.Haskell.TH hiding (interruptible)
 import Language.Haskell.TH.Syntax
 import Network.Rpc.Multiplexer
 import Network.Rpc.Connection
 import qualified Network.Socket as Socket
 import Prelude
 import GHC.Generics
-import GHC.IO (unsafeUnmask)
 import System.Posix.Files (getFileStatus, isSocket)
 
 
@@ -261,11 +260,11 @@ emptyClientState = ClientState {
 }
 
 clientSend :: RpcProtocol p => Client p -> ProtocolRequest p -> IO ()
-clientSend client req = channelSend_ client.channel (encode req) []
+clientSend client req = channelSend_ client.channel [] (encode req)
 clientRequestBlocking :: forall p. RpcProtocol p => Client p -> ProtocolRequest p -> IO (ProtocolResponse p)
 clientRequestBlocking client req = do
   resultMVar <- newEmptyMVar
-  channelSend client.channel (encode req) [] $ \msgId ->
+  channelSend client.channel [] (encode req) $ \msgId ->
     modifyMVar_ client.stateMVar $
       \state -> pure state{callbacks = HM.insert msgId (requestCompletedCallback resultMVar msgId) state.callbacks}
   -- Block on resultMVar until the request completes
@@ -305,7 +304,7 @@ serverHandleChannelMessage protocolImpl channel msgId headers msg = case decodeO
     serverHandleChannelRequest :: ProtocolRequest p -> IO ()
     serverHandleChannelRequest req = handleMessage @p protocolImpl req >>= maybe (pure ()) serverSendResponse
     serverSendResponse :: ProtocolResponse p -> IO ()
-    serverSendResponse response = channelSend_ channel (encode wrappedResponse) []
+    serverSendResponse response = channelSend_ channel [] (encode wrappedResponse)
       where
         wrappedResponse :: ProtocolResponseWrapper p
         wrappedResponse = (msgId, response)
@@ -337,11 +336,7 @@ withClient :: forall p a b. (IsConnection a, RpcProtocol p) => a -> (Client p ->
 withClient x = bracket (newClient x) clientClose
 
 newClient :: forall p a. (IsConnection a, RpcProtocol p) => a -> IO (Client p)
-newClient x = do
-  clientMVar <- newEmptyMVar
-  -- 'runMultiplexerProtcol' needs to be interruptible (so it can terminate when it is closed), so 'unsafeUnmask' is used to ensure that this function also works when used in 'bracket'
-  link =<< async (unsafeUnmask (runMultiplexerProtocol (newChannelClient >=> putMVar clientMVar) (toSocketConnection x)))
-  takeMVar clientMVar
+newClient x = newChannelClient =<< newMultiplexer (toSocketConnection x)
 
 
 newChannelClient :: RpcProtocol p => Channel -> IO (Client p)
@@ -400,7 +395,7 @@ listenOnBoundSocket protocolImpl sock = do
       Socket.gracefulClose conn 2000
 
 runServerHandler :: forall p a. (RpcProtocol p, HasProtocolImpl p, IsConnection a) => ProtocolImpl p -> a -> IO ()
-runServerHandler protocolImpl = runMultiplexerProtocol (registerChannelServerHandler @p protocolImpl) . toSocketConnection
+runServerHandler protocolImpl = runMultiplexer (registerChannelServerHandler @p protocolImpl) . toSocketConnection
 
 
 -- ** Test implementation
@@ -483,10 +478,10 @@ buildTupleType fields = buildTupleType' =<< fields
     go t (f:fs) = go (AppT t f) fs
 
 buildFunctionType :: Q [Type] -> Q Type -> Q Type
-buildFunctionType argTypes pureType = go =<< argTypes
+buildFunctionType argTypes returnType = go =<< argTypes
   where
     go :: [Type] -> Q Type
-    go [] = pureType
+    go [] = returnType
     go (t:ts) = pure t `funT` go ts
 
 defaultBangType  :: Q Type -> Q BangType

src/Network/Rpc/Connection.hs

@@ -3,12 +3,11 @@ module Network.Rpc.Connection where
 import Control.Concurrent (threadDelay)
 import Control.Concurrent.Async (Async, async, cancel, link, waitCatch, withAsync)
 import Control.Concurrent.MVar
-import Control.Exception (Exception(..), SomeException, bracketOnError, finally, throwIO, bracketOnError, onException)
+import Control.Exception (Exception(..), SomeException, bracketOnError, interruptible, finally, throwIO, bracketOnError, onException)
 import Control.Monad ((>=>), unless, forM_)
 import qualified Data.ByteString as BS
 import qualified Data.ByteString.Lazy as BSL
 import Data.List (intercalate)
-import GHC.IO (unsafeUnmask)
 import qualified Network.Socket as Socket
 import qualified Network.Socket.ByteString as Socket
 import qualified Network.Socket.ByteString.Lazy as SocketL
@@ -37,6 +36,8 @@ newtype ConnectionFailed = ConnectionFailed [(Socket.AddrInfo, SomeException)]
 instance Exception ConnectionFailed where
   displayException (ConnectionFailed attemts) = "Connection attempts failed:\n" <> intercalate "\n" (map (\(addr, err) -> show (Socket.addrAddress addr) <> ": " <> displayException err) attemts)
 
+-- | Open a TCP connection to target host and port. Will start multiple connection attempts (i.e. retry quickly and then try other addresses) but only return the first successful connection.
+-- Throws a 'ConnectionFailed' on failure, which contains the exceptions from all failed connection attempts.
 connectTCP :: Socket.HostName -> Socket.ServiceName -> IO Socket.Socket
 connectTCP host port = do
   -- 'getAddrInfo' either pures a non-empty list or throws an exception
@@ -74,7 +75,7 @@ connectTCP host port = do
 
   -- The 'raceConnections'-async is 'link'ed to this thread, so 'readMVar' is interrupted when all connection attempts fail
   sock <-
-    (withAsync (unsafeUnmask raceConnections) (link >=> const (readMVar sockMVar))
+    (withAsync (interruptible raceConnections) (link >=> const (readMVar sockMVar))
       `finally` (mapM_ (cancel . snd) =<< readMVar connectTasksMVar))
         `onException` (mapM_ Socket.close =<< tryTakeMVar sockMVar)
     -- As soon as we have an open connection, stop spawning more connections

test/Network/Rpc/MultiplexerSpec.hs

+module Network.Rpc.MultiplexerSpec where
+
+import Control.Concurrent.Async (concurrently_)
+import Control.Concurrent.MVar
+import Control.Exception (bracket, mask_)
+import Prelude
+import Network.Rpc.Multiplexer
+import Network.Rpc.Connection
+import Test.Hspec
+
+spec :: Spec
+spec = describe "runMultiplexerProtocol" $ parallel $ do
+  it "can be closed from the channelSetupHook" $ do
+    (x, _) <- newDummySocketPair
+    runMultiplexer channelClose x
+
+  it "fails when run in masked state" $ do
+    (x, _) <- newDummySocketPair
+    mask_ $ runMultiplexer channelClose x `shouldThrow` anyException
+
+  it "closes when the remote is closed" $ do
+    (x, y) <- newDummySocketPair
+    concurrently_
+      (runMultiplexer (const (pure ())) x)
+      (runMultiplexer channelClose y)
+
+  it "it can send and receive simple messages" $ do
+    withEchoServer $ \channel -> do
+      recvMVar <- newEmptyMVar
+      channelSetHandler channel $ simpleMessageHandler $ \_ _ -> putMVar recvMVar
+      channelSend_ channel [] "foobar"
+      takeMVar recvMVar `shouldReturn` "foobar"
+      channelSend_ channel [] "test"
+      takeMVar recvMVar `shouldReturn` "test"
+
+
+withEchoServer :: (Channel -> IO a) -> IO a
+withEchoServer fn = bracket setup close (\(channel, _) -> fn channel)
+  where
+    setup :: IO (Channel, Channel)
+    setup = do
+      (x, y) <- newDummySocketPair
+      echoChannel <- newMultiplexer y
+      configureEchoHandler echoChannel
+      mainChannel <- newMultiplexer x
+      pure (mainChannel, echoChannel)
+    close :: (Channel, Channel) -> IO ()
+    close (x, y) = channelClose x >> channelClose y
+    configureEchoHandler :: Channel -> IO ()
+    configureEchoHandler channel = channelSetHandler channel (echoHandler channel)
+    echoHandler :: Channel -> ChannelMessageHandler
+    echoHandler channel = simpleMessageHandler $ \_msgId headers msg -> do
+      mapM_ echoHeaderHandler headers
+      channelSend_ channel [] msg
+    echoHeaderHandler :: MessageHeaderResult -> IO ()
+    echoHeaderHandler (CreateChannelHeaderResult channel) = configureEchoHandler channel