diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,30 @@ +Copyright (c) 2014, EURL Tweag + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided + with the distribution. + + * Neither the name of Alexander Vershilov nor the names of other + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/benchmarks/Channels.hs b/benchmarks/Channels.hs new file mode 100644 --- /dev/null +++ b/benchmarks/Channels.hs @@ -0,0 +1,42 @@ +-- | Like Latency, but creating lots of channels +import System.Environment +import Control.Monad +import Control.Applicative +import Control.Distributed.Process +import Control.Distributed.Process.Node +import Network.Transport.ZMQ (createTransport, defaultZMQParameters) +import Data.Binary (encode, decode) +import Data.ByteString.Char8 (pack) +import qualified Data.ByteString.Lazy as BSL + +pingServer :: Process () +pingServer = forever $ do + them <- expect + sendChan them () + -- TODO: should this be automatic? + reconnectPort them + +pingClient :: Int -> ProcessId -> Process () +pingClient n them = do + replicateM_ n $ do + (sc, rc) <- newChan :: Process (SendPort (), ReceivePort ()) + send them sc + receiveChan rc + liftIO . putStrLn $ "Did " ++ show n ++ " pings" + +initialProcess :: String -> Process () +initialProcess "SERVER" = do + us <- getSelfPid + liftIO $ BSL.writeFile "pingServer.pid" (encode us) + pingServer +initialProcess "CLIENT" = do + n <- liftIO $ getLine + them <- liftIO $ decode <$> BSL.readFile "pingServer.pid" + pingClient (read n) them + +main :: IO () +main = do + [role, host] <- getArgs + Right transport <- createTransport defaultZMQParameters (pack host) + node <- newLocalNode transport initRemoteTable + runProcess node $ initialProcess role diff --git a/benchmarks/Latency.hs b/benchmarks/Latency.hs new file mode 100644 --- /dev/null +++ b/benchmarks/Latency.hs @@ -0,0 +1,37 @@ +import System.Environment +import Control.Monad +import Control.Applicative +import Control.Distributed.Process +import Control.Distributed.Process.Node +import Network.Transport.ZMQ (createTransport, defaultZMQParameters) +import Data.Binary (encode, decode) +import Data.ByteString.Char8 (pack) +import qualified Data.ByteString.Lazy as BSL + +pingServer :: Process () +pingServer = forever $ do + them <- expect + send them () + +pingClient :: Int -> ProcessId -> Process () +pingClient n them = do + us <- getSelfPid + replicateM_ n $ send them us >> (expect :: Process ()) + liftIO . putStrLn $ "Did " ++ show n ++ " pings" + +initialProcess :: String -> Process () +initialProcess "SERVER" = do + us <- getSelfPid + liftIO $ BSL.writeFile "pingServer.pid" (encode us) + pingServer +initialProcess "CLIENT" = do + n <- liftIO $ getLine + them <- liftIO $ decode <$> BSL.readFile "pingServer.pid" + pingClient (read n) them + +main :: IO () +main = do + [role, host] <- getArgs + Right transport <- createTransport defaultZMQParameters (pack host) + node <- newLocalNode transport initRemoteTable + runProcess node $ initialProcess role diff --git a/benchmarks/Throughput.hs b/benchmarks/Throughput.hs new file mode 100644 --- /dev/null +++ b/benchmarks/Throughput.hs @@ -0,0 +1,71 @@ +{-# LANGUAGE DeriveDataTypeable #-} +import System.Environment +import Control.Monad +import Control.Applicative +import Control.Distributed.Process +import Control.Distributed.Process.Node +import Network.Transport.ZMQ (createTransport, defaultZMQParameters) +import Data.Binary +import Data.ByteString.Char8 ( pack ) +import qualified Data.ByteString.Lazy as BSL +import Data.Typeable + +data SizedList a = SizedList { size :: Int , elems :: [a] } + deriving (Typeable) + +instance Binary a => Binary (SizedList a) where + put (SizedList sz xs) = put sz >> mapM_ put xs + get = do + sz <- get + xs <- getMany sz + return (SizedList sz xs) + +-- Copied from Data.Binary +getMany :: Binary a => Int -> Get [a] +getMany = go [] + where + go xs 0 = return $! reverse xs + go xs i = do x <- get + x `seq` go (x:xs) (i-1) +{-# INLINE getMany #-} + +nats :: Int -> SizedList Int +nats = \n -> SizedList n (aux n) + where + aux 0 = [] + aux n = n : aux (n - 1) + +counter :: Process () +counter = go 0 + where + go :: Int -> Process () + go !n = + receiveWait + [ match $ \xs -> go (n + size (xs :: SizedList Int)) + , match $ \them -> send them n >> go 0 + ] + +count :: (Int, Int) -> ProcessId -> Process () +count (packets, sz) them = do + us <- getSelfPid + replicateM_ packets $ send them (nats sz) + send them us + n' <- expect + liftIO $ print (packets * sz, n' == packets * sz) + +initialProcess :: String -> Process () +initialProcess "SERVER" = do + us <- getSelfPid + liftIO $ BSL.writeFile "counter.pid" (encode us) + counter +initialProcess "CLIENT" = do + n <- liftIO getLine + them <- liftIO $ decode <$> BSL.readFile "counter.pid" + count (read n) them + +main :: IO () +main = do + [role, host] <- getArgs + Right transport <- createTransport defaultZMQParameters (pack host) + node <- newLocalNode transport initRemoteTable + runProcess node $ initialProcess role diff --git a/network-transport-zeromq.cabal b/network-transport-zeromq.cabal new file mode 100644 --- /dev/null +++ b/network-transport-zeromq.cabal @@ -0,0 +1,184 @@ +name: network-transport-zeromq +version: 0.1.0.0 +synopsis: ZeroMQ backend for network-transport +license: BSD3 +license-file: LICENSE +author: Alexander Vershilov +maintainer: Alexander Vershilov +copyright: (c) 2014, EURL Tweag +category: Network +build-type: Simple +cabal-version: >=1.10 + +Source-repository head + Type: git + Location: https://github.com/tweag/network-transport-zeromq + +flag benchmarks + description: build benchmarks + default: False + +flag unsafe-constructs + description: use faster primitives, but it will introduce undefined behaviour on some usecases + default: False + +library + build-depends: base >=4.6 && < 4.8, + binary >= 0.6, + network-transport >= 0.3, + zeromq4-haskell >= 0.2, + async >= 2.0, + stm >= 2.4, + stm-chans >= 0.3, + containers >= 0.5, + bytestring >= 0.10, + transformers >= 0.3, + semigroups >= 0.12, + exceptions >= 0.3, + void >= 0.6, + random >= 1.0 + exposed-modules: Network.Transport.ZMQ + Network.Transport.ZMQ.Types + other-modules: System.ZMQ4.Utils + hs-source-dirs: src + ghc-options: -Wall + if flag(unsafe-constructs) + cpp-options: -DUNSAFE_SEND + default-extensions: DeriveGeneric + DeriveDataTypeable + OverloadedStrings + LambdaCase + ScopedTypeVariables + StandaloneDeriving + default-language: Haskell2010 + + +Test-Suite test-zeromq + type: exitcode-stdio-1.0 + main-is: TestZMQ.hs + build-depends: base >= 4.4 && < 5, + network-transport >= 0.2 && < 0.4, + network-transport-zeromq, + zeromq4-haskell >= 0.2, + network-transport-tests >= 0.1.0.1 && < 0.2 + ghc-options: -Wall -threaded -rtsopts -with-rtsopts=-N -fno-warn-unused-do-bind + hs-source-dirs: tests + default-language: Haskell2010 + +Test-Suite test-api + type: exitcode-stdio-1.0 + main-is: TestAPI.hs + build-depends: base >= 4.4 && < 5, + network-transport >= 0.2 && < 0.4, + network-transport-zeromq, + zeromq4-haskell >= 0.2 + ghc-options: -Wall -threaded -rtsopts -with-rtsopts=-N -fno-warn-unused-do-bind + hs-source-dirs: tests + default-language: Haskell2010 + + + +Test-Suite test-ch-core + Type: exitcode-stdio-1.0 + Main-Is: test-ch.hs + CPP-Options: -DTEST_SUITE_MODULE=Control.Distributed.Process.Tests.CH + Build-Depends: base >= 4.4 && < 5, + network-transport-zeromq, + distributed-process-tests >= 0.4, + network >= 2.3 && < 2.5, + network-transport >= 0.3 && < 0.4, + test-framework >= 0.6 && < 0.9, + containers, + stm, + stm-chans, + bytestring + default-extensions: CPP + ghc-options: -Wall -threaded -rtsopts -with-rtsopts=-N -fno-warn-unused-do-bind + HS-Source-Dirs: tests + default-language: Haskell2010 + + +Test-Suite test-ch-closure + Type: exitcode-stdio-1.0 + Main-Is: test-ch.hs + CPP-Options: -DTEST_SUITE_MODULE=Control.Distributed.Process.Tests.Closure + Build-Depends: base >= 4.4 && < 5, + network-transport-zeromq, + distributed-process-tests >= 0.4, + network >= 2.3 && < 2.5, + network-transport >= 0.3 && < 0.4, + test-framework >= 0.6 && < 0.9, + containers, + stm, + stm-chans, + bytestring + default-extensions: CPP + ghc-options: -Wall -threaded -rtsopts -with-rtsopts=-N -fno-warn-unused-do-bind -caf-all -auto-all + HS-Source-Dirs: tests + default-language: Haskell2010 + +Test-Suite test-ch-stat + Type: exitcode-stdio-1.0 + Main-Is: test-ch.hs + CPP-Options: -DTEST_SUITE_MODULE=Control.Distributed.Process.Tests.Stats + Build-Depends: base >= 4.4 && < 5, + network-transport-zeromq, + distributed-process-tests >= 0.4, + network >= 2.3 && < 2.5, + network-transport >= 0.3 && < 0.4, + test-framework >= 0.6 && < 0.9, + containers, + stm, + stm-chans, + bytestring + default-extensions: CPP + ghc-options: -Wall -threaded -rtsopts -with-rtsopts=-N -fno-warn-unused-do-bind + HS-Source-Dirs: tests + default-language: Haskell2010 + + + +--- benchmark files +executable bench-dp-latency + main-is: Latency.hs + build-depends: base >= 4.4 && < 5, + network-transport-zeromq, + bytestring, + binary, + distributed-process + hs-source-dirs: benchmarks + if flag(benchmarks) + buildable: True + else + buildable: False + default-language: Haskell2010 + +Executable bench-dp-throughput + if flag(benchmarks) + Build-Depends: base >= 4.4 && < 5, + distributed-process, + network-transport-zeromq, + bytestring >= 0.9 && < 0.11, + binary >= 0.5 && < 0.8 + buildable: True + else + buildable: False + main-is: Throughput.hs + hs-source-dirs: benchmarks + ghc-options: -Wall + default-extensions: BangPatterns + +Executable bench-dp-channels + if flag(benchmarks) + Build-Depends: base >= 4.4 && < 5, + distributed-process, + network-transport-zeromq, + bytestring >= 0.9 && < 0.11, + binary >= 0.5 && < 0.8 + buildable: True + else + buildable: False + main-is: Channels.hs + hs-source-dirs: benchmarks + ghc-options: -Wall + default-extensions: BangPatterns diff --git a/src/Network/Transport/ZMQ.hs b/src/Network/Transport/ZMQ.hs new file mode 100644 --- /dev/null +++ b/src/Network/Transport/ZMQ.hs @@ -0,0 +1,1182 @@ +-- | +-- Module: Network.Transport.ZMQ +-- Copyright: (C) 2014, EURL Tweag +-- License: BSD-3 +{-# LANGUAGE DeriveGeneric, StandaloneDeriving, OverloadedStrings, DeriveDataTypeable, CPP #-} +module Network.Transport.ZMQ + ( -- * Main API + createTransport -- :: ZMQParameters -> ByteString -> IO (Either ZMQError Transport) + , ZMQParameters(..) + , ZMQAuthType(..) + , defaultZMQParameters -- :: ZMQParameters + -- * Internals + -- $internals + , createTransportEx + , breakConnectionEndPoint -- :: ZMQTransport -> EndPointAddress -> EndPointAddress -> IO () + , breakConnection -- :: ZMQTransport -> EndPointAddress -> EndPointAddress -> IO () + , unsafeConfigurePush + -- * Design + -- $design + -- ** Multicast + -- $multicast + ) where + +import Network.Transport.ZMQ.Types + +import Control.Applicative +import Control.Concurrent + ( yield + , threadDelay + ) +import qualified Control.Concurrent.Async as Async +import Control.Concurrent.MVar +import Control.Concurrent.STM +import Control.Concurrent.STM.TMChan +import Control.Monad + ( void + , forever + , unless + , join + , forM_ + , foldM + , when + , (<=<) + ) +import Control.Exception + ( AsyncException ) +import Control.Monad.Catch + ( finally + , try + , throwM + , Exception + , SomeException + , fromException + , mask + , mask_ + , uninterruptibleMask_ + , onException + ) +import Data.Binary +import Data.ByteString (ByteString) +import qualified Data.ByteString as B +import qualified Data.ByteString.Lazy as BL +import qualified Data.ByteString.Char8 as B8 +import Data.Foldable ( traverse_ ) +import qualified Data.Foldable as Foldable +import Data.IORef + ( newIORef + , modifyIORef + , readIORef + , writeIORef + ) +import Data.List.NonEmpty +import qualified Data.Map.Strict as Map +-- import Data.Set (Set) +import qualified Data.Set as Set +import qualified Data.Traversable as Traversable +import Data.Typeable +import Data.Void +import GHC.Generics + ( Generic ) + +import Network.Transport +import Network.Transport.ZMQ.Types +import System.IO + ( fixIO ) +import System.ZMQ4 + ( Context ) +import qualified System.ZMQ4 as ZMQ +import qualified System.ZMQ4.Utils as ZMQ + +import Text.Printf + +-------------------------------------------------------------------------------- +--- Internal datatypes -- +-------------------------------------------------------------------------------- +-- $design +-- +-- In the zeromq backend we are using following address scheme: +-- +-- @ +-- tcp:\/\/host:port\/ +-- | | | +-- | +---+---------------- can be configured by user, one host, +-- | port pair per distributed process +-- | instance +-- +---------------------------- In feature it will be possible to add another schemas. +-- @ +-- +-- Transport specifies host that will be used, and port will be +-- automatically generated by the transport for each endpoint. +-- +-- * Connections Reliability +-- +-- Currently only reliable connections are supportred. In case if +-- Unreliable type was passed reliable connection will be created. +-- This may be changed in future versions. +-- +-- Network-transport-zeromq maintains one thread for each endpoint that is +-- used to read incomming requests, all sends and connection requests are +-- handled from the user thread and mostly asynchronous. +-- +-- Each endpoint obtains one pull socket and push socket for each remote +-- end point, all lightweight threads abrigged into one heavyweight +-- connection. +-- +-- Virually connections looks like the following plot: +-- +-- @ +-- +------------------+ +--------------+ +-- | Endpoint | | Endpoint | +-- | | | | +-- | +------+ +------+ | +-- | | pull |<-----------------------| push | | +-- + +------+ +------+ | +-- | | +--------------------+ | | +-- + +------+/~~~~ connection 1 ~~~~\\+------+ | +-- | | push |~~~~~ connection 2 ~~~~~| pull | | +-- | +------+\\ /+------+ | +-- | | +--------------------+ | | +-- +------------------+ +--------------+ +-- @ +-- +-- Physically 0mq may choose better representations and mapping on +-- a real connections. +-- +-- ZeroMQ connection takes care of reliability thus for correct lifeness, +-- it stores messages so in case of connecdtion break in may be restarted +-- with no message loss. So heartbeating procedure should be introduced on +-- the top on network-transport, and 'breakConnection' should be used to notify +-- connection death. However if High Water Mark will be reached connection +-- to endpoint considered failed and connection break procedure started +-- automatically. +-- + +-- Naming conventions. +-- api* -- functions that can be called by user +-- localEndPoint -- internal functions that are used in endpoints +-- remoteEndPoint -- internal functions that are used in endpoints +-- +-- Transport +-- |--- Valid transport itself +-- | |--- newEndpoint +-- | |--- close +-- |--- Closed +-- +-- LocaEndPoint +-- RemoteEndPoint +-- Connection + +-- Connection closing procedure. +-- Once endpoint is gracefully closed by user or by transport close +-- RemoteEndPoint goes into 'RemoteEndPointClosing' state and sends +-- a message to remote endpoint if link is known as alive, in +-- RemoteEndPointClosing all messages are ignored and EndPointClosed +-- is returned. Uppon a EndPointClose delivery endpoint replies with +-- EndPointClose message, and starts cleanup procedure (marking link +-- as not alive). When endpoint receives EndPointCloseOk message it +-- moved RemoteEndPoint to close state and removes all data structures. +-- Together with marking a endpoint as closed asynchronous timeout is +-- set, and if EndPointCloseOk is not delivered withing that timeperiod +-- EndPoint is closed. +-- +-- endpoint-1 endpoint-2 +-- 1. mark as closing +-- [RemoteEndPoint:Closing] ----endPointClose-> [remoteEndPoint:Valid] +-- 2. mark as closed +-- <-endPointCloseOk- [remoteEndPoint:Closed] +-- 2. mark as closed +-- [RemoteEndPoint:Closed] 4. cleanup remote end point +-- +-- EndPoint can be closed in a two ways: normally and abnormally (in case +-- of exception or invariant vionation on a remove side). If EndPoint is +-- closed normally all opened connections will recevice ConnectionClosed +-- events, otherwise. +-- +-- Multicast group schema: +-- +-- +--------------+ +-- | EndPoint |<----------------+ +-- +--------------+ | events +-- | | +-- +------------------+ | +-- | | +-- +------------------------------+ +-- | Multicast Group (Server) | +-- +--+ +--+ +--+ +---+ +-- | Rep | | Pub | | Sub | +-- +-----+ +-----+ +-----+ +-- ^ | ^ +-- request | | | +-- +------------+ +--------+ messages +-- | | +-- | +-----------+ +-- | | messages +-- +-----+ +-----+ +-- | Req | | Sub | +-- +---+ +--+ +--+ +-- | Multicast Group | +-- +---------------------+ +-- +-- Resolve Group automatically subscribes to the group to receive events, +-- however unless 'subscribe' will be called endpoint will not reveive any +-- message. This is done in order to keep ability to notify receiver about +-- multcast group close. +-- +-- Current solution incorrectly keeps track of incomming connections, this +-- means that we can't really guarantee that close message is delivered to +-- every recipient and it disconnected correctly. +-- + + +-- $internals +-- Internal function provides additional 0mq specific set of configuration +-- options, this functionality should be used with caution as it may break +-- required socket properties. + +-- | Messages +data ZMQMessage + = MessageConnect !EndPointAddress -- ^ Connection greeting + | MessageInitConnection !EndPointAddress !ConnectionId !Reliability + | MessageInitConnectionOk !EndPointAddress !ConnectionId !ConnectionId + | MessageCloseConnection !ConnectionId + | MessageData !ConnectionId + | MessageEndPointClose !EndPointAddress !Bool + | MessageEndPointCloseOk !EndPointAddress + deriving (Generic) + +{- +instance Binary ZMQMessage where + put (MessageConnect ep) = putWord64le 0 >> put ep + put (MessageInitConnection ep cid rel) = putWord64le 2 >> put ep >> put cid >> put rel + put (MessageInitConnectionOk ep cid rid) = putWord64le 3 >> put ep >> put cid >> put rid + put (MessageCloseConnection cid) = putWord64le 4 >> put cid + put (MessageEndPointClose ep) = putWord64le 5 >> put ep + put (MessageData cid) = putWord64le cid + get = do x <- getWord64be + case x of + 0 -> MessageConnect <$> get + 2 -> MessageInitConnection <$> get <*> get <*> get + 3 -> MessageInitConnectionOk <$> get <*> get <*> get + 4 -> MessageCloseConnection <$> get + 5 -> MessageEndPointClose <$> get + 6 -> MessageData <$> pure x + +reservedConnectionId = 7 +-} + +instance Binary ZMQMessage + +data ZMQError = InvariantViolation String + | IncorrectState String + | ConnectionFailed + deriving (Typeable, Show) + +instance Exception ZMQError + +-- | Create 0MQ based transport. +createTransport :: ZMQParameters + -> ByteString -- ^ Transport address (IP or hostname) + -> IO (Either (TransportError Void) Transport) +createTransport z b = fmap (fmap snd) (createTransportEx z b) + +-- | Create 'Transport' and export internal transport state that may be +-- used in API. +createTransportEx :: ZMQParameters -- ^ Transport features. + -> ByteString -- ^ Host name or IP address + -> IO (Either (TransportError Void) (ZMQTransport, Transport)) +createTransportEx params host = do + ctx <- ZMQ.context + mtid <- case authorizationType params of + ZMQAuthPlain user pass -> Just <$> ZMQ.authManager ctx user pass + _ -> return Nothing + transport <- ZMQTransport + <$> pure addr + <*> newMVar (TransportValid $ ValidTransportState ctx Map.empty mtid) + return $ Right (transport, Transport + { newEndPoint = apiNewEndPoint params transport + , closeTransport = apiTransportClose transport + }) + where + addr = B.concat ["tcp://", host] + +-- Synchronous +apiTransportClose :: ZMQTransport -> IO () +apiTransportClose transport = mask_ $ do + old <- swapMVar (_transportState transport) TransportClosed + case old of + TransportClosed -> return () + TransportValid (ValidTransportState ctx m mtid) -> do + case mtid of + Nothing -> return () + Just tid -> Async.cancel tid + forM_ (Map.elems m) $ apiCloseEndPoint transport + ZMQ.term ctx + +apiNewEndPoint :: ZMQParameters -> ZMQTransport -> IO (Either (TransportError NewEndPointErrorCode) EndPoint) +apiNewEndPoint params transport = do +-- printf "[transport] endpoint create\n" + elep <- modifyMVar (_transportState transport) $ \case + TransportClosed -> return (TransportClosed, Left $ TransportError NewEndPointFailed "Transport is closed.") + v@(TransportValid i@(ValidTransportState ctx _ _)) -> do + eEndPoint <- endPointCreate params ctx (B8.unpack addr) + case eEndPoint of + Right (_port, ep, chan) -> return + ( TransportValid i{_transportEndPoints = Map.insert (localEndPointAddress ep) ep (_transportEndPoints i)} + , Right (ep, ctx, chan)) + Left _ -> return (v, Left $ TransportError NewEndPointFailed "Failed to create new endpoint.") + case elep of + Right (ep,ctx, chOut) -> + return $ Right $ EndPoint + { receive = atomically $ do + mx <- readTMChan chOut + case mx of + Nothing -> error "channel is closed" + Just x -> return x + , address = localEndPointAddress ep + , connect = apiConnect params ctx ep + , closeEndPoint = apiCloseEndPoint transport ep + , newMulticastGroup = apiNewMulticastGroup params transport ep + , resolveMulticastGroup = apiResolveMulticastGroup transport ep + } + Left x -> return $ Left x + where + addr = transportAddress transport + +-- | Asynchronous operation, shutdown of the remote end point may take a while +apiCloseEndPoint :: ZMQTransport + -> LocalEndPoint + -> IO () +apiCloseEndPoint transport lep = mask_ $ do +-- printf "[%s][go] close endpoint\n" +-- (B8.unpack $ endPointAddressToByteString $ _localEndPointAddress lep) + old <- readMVar (localEndPointState lep) + case old of + LocalEndPointValid (ValidLocalEndPoint x _ _ threadId _ _) -> do + -- close channel, no events will be received + atomically $ do + writeTMChan x EndPointClosed + closeTMChan x + Async.cancel threadId + void $ Async.waitCatch threadId + LocalEndPointClosed -> return () + modifyMVar_ (_transportState transport) $ \case + TransportClosed -> return TransportClosed + TransportValid v -> return $ TransportValid + v{_transportEndPoints = Map.delete (localEndPointAddress lep) (_transportEndPoints v)} + +endPointCreate :: ZMQParameters + -> Context + -> String + -> IO (Either (TransportError NewEndPointErrorCode) (Int,LocalEndPoint, TMChan Event)) +endPointCreate params ctx addr = do + em <- try $ do + pull <- ZMQ.socket ctx ZMQ.Pull + case authorizationType params of + ZMQNoAuth -> return () + ZMQAuthPlain{} -> do + ZMQ.setPlainServer True pull + ZMQ.setSendHighWM (ZMQ.restrict (highWaterMark params)) pull + port <- ZMQ.bindFromRangeRandom pull addr (minPort params) (maxPort params) (maxTries params) + return (port, pull) + case em of + Right (port, pull) -> (do + chOut <- newTMChanIO + lep <- LocalEndPoint <$> pure (EndPointAddress $ B8.pack (addr ++ ":" ++ show port)) + <*> newEmptyMVar + <*> pure port + opened <- newIORef True + mask $ \restore -> do + thread <- Async.async $ (restore (receiver pull lep chOut)) + `finally` finalizeEndPoint lep port pull + putMVar (localEndPointState lep) $ LocalEndPointValid + (ValidLocalEndPoint chOut (Counter 0 Map.empty) Map.empty thread opened Map.empty) + return $ Right (port, lep, chOut)) + `onException` (ZMQ.close pull) + Left (_e::SomeException) -> do + return $ Left $ TransportError NewEndPointInsufficientResources "no free sockets" + where + + finalizer pull ourEp = forever $ do + (cmd:_) <- ZMQ.receiveMulti pull + case decode' cmd of + MessageEndPointCloseOk theirAddress -> getRemoteEndPoint ourEp theirAddress >>= \case + Nothing -> return () + Just rep -> do + state <- swapMVar (remoteEndPointState rep) RemoteEndPointClosed + closeRemoteEndPoint ourEp rep state + _ -> return () -- XXX: send exception + + receiver :: ZMQ.Socket ZMQ.Pull + -> LocalEndPoint + -> TMChan Event + -> IO () + receiver pull ourEp chan = forever $ do + (cmd:msgs) <- ZMQ.receiveMulti pull + case decode' cmd of + MessageData idx -> atomically $ writeTMChan chan (Received idx msgs) + MessageConnect theirAddress -> do +-- printf "[%s] message connect from %s\n" +-- (B8.unpack $ endPointAddressToByteString ourAddr) +-- (B8.unpack $ endPointAddressToByteString theirAddress) + void $ createOrGetRemoteEndPoint params ctx ourEp theirAddress + MessageInitConnection theirAddress theirId rel -> do +-- printf "[%s] message init connection from %s\n" +-- (B8.unpack $ endPointAddressToByteString ourAddr) +-- (B8.unpack $ endPointAddressToByteString theirAddress) + join $ do + modifyMVar (localEndPointState ourEp) $ \case + LocalEndPointValid v -> + case theirAddress `Map.lookup` r of + Nothing -> return (LocalEndPointValid v, throwM $ InvariantViolation "Remote endpoint should exist.") + Just rep -> modifyMVar (remoteEndPointState rep) $ \case + RemoteEndPointFailed -> throwM $ InvariantViolation "RemoteEndPoint should be valid." + RemoteEndPointClosed -> throwM $ InvariantViolation "RemoteEndPoint should be valid." + RemoteEndPointClosing{} -> throwM $ InvariantViolation "RemoteEndPoint should be valid." + w@RemoteEndPointValid{} -> do + conn <- ZMQConnection <$> pure ourEp + <*> pure rep + <*> pure rel + <*> newMVar (ZMQConnectionValid $ ValidZMQConnection Nothing i) + <*> newEmptyMVar + w' <- register (succ i) w + return ( w' + , ( LocalEndPointValid v{ _localEndPointConnections = Counter (succ i) (Map.insert (succ i) conn m) } + , return ()) + ) + z@(RemoteEndPointPending w) -> do + conn <- ZMQConnection <$> pure ourEp + <*> pure rep + <*> pure rel + <*> newMVar (ZMQConnectionValid $ ValidZMQConnection Nothing i) + <*> newEmptyMVar + modifyIORef w (\xs -> (register (succ i)) : xs) + return ( z + , ( LocalEndPointValid v{ _localEndPointConnections = Counter (succ i) (Map.insert (succ i) conn m) } + , return ()) + ) + where + r = _localEndPointRemotes v + (Counter i m) = _localEndPointConnections v + _ -> throwM $ InvariantViolation "RemoteEndPoint should be valid." + where + register i RemoteEndPointFailed = do + atomically $ do + writeTMChan chan (ConnectionOpened i rel theirAddress) + writeTMChan chan (ConnectionClosed i) + return RemoteEndPointFailed + register i RemoteEndPointClosed = do + atomically $ do + writeTMChan chan (ConnectionOpened i rel theirAddress) + writeTMChan chan (ConnectionClosed i) + return RemoteEndPointClosed + register i (RemoteEndPointClosing x) = do + atomically $ do + writeTMChan chan (ConnectionOpened i rel theirAddress) + writeTMChan chan (ConnectionClosed i) + return $ RemoteEndPointClosing x + register _ RemoteEndPointPending{} = throwM $ InvariantViolation "RemoteEndPoint can't be pending." + register i (RemoteEndPointValid v@(ValidRemoteEndPoint sock _ s _)) = do + ZMQ.send sock [] $ encode' (MessageInitConnectionOk ourAddr theirId i) + atomically $ writeTMChan chan (ConnectionOpened i rel theirAddress) + return $ RemoteEndPointValid + v{_remoteEndPointIncommingConnections = Set.insert i s} + MessageCloseConnection idx -> join $ do +-- printf "[%s] message close connection: %i\n" +-- (B8.unpack $ endPointAddressToByteString ourAddr) +-- idx + modifyMVar (localEndPointState ourEp) $ \case + LocalEndPointValid v -> + case idx `Map.lookup` m of + Nothing -> return (LocalEndPointValid v, return ()) + Just conn -> do + old <- swapMVar (connectionState conn) ZMQConnectionFailed + return ( LocalEndPointValid v{ _localEndPointConnections = Counter i (idx `Map.delete` m)} + , case old of + ZMQConnectionFailed -> return () + ZMQConnectionInit -> return () -- throwM InvariantViolation + ZMQConnectionClosed -> return () + ZMQConnectionValid (ValidZMQConnection _ _) -> do + atomically $ writeTMChan chan (ConnectionClosed idx) + connectionCleanup (connectionRemoteEndPoint conn) idx) + where + (Counter i m) = _localEndPointConnections v + LocalEndPointClosed -> return (LocalEndPointClosed, return ()) + MessageInitConnectionOk theirAddress ourId theirId -> do +-- printf "[%s] message init connection ok: %i -> %i\n" +-- (B8.unpack $ endPointAddressToByteString ourAddr) +-- ourId +-- theirId + join $ withMVar (localEndPointState ourEp) $ \case + LocalEndPointValid v -> + case theirAddress `Map.lookup` r of + Nothing -> return (return ()) -- XXX: send message to the host + Just rep -> modifyMVar (remoteEndPointState rep) $ \case + RemoteEndPointFailed -> return (RemoteEndPointFailed, return ()) + RemoteEndPointClosed -> throwM $ InvariantViolation "RemoteEndPoint should be valid or failed." + RemoteEndPointClosing{} -> throwM $ InvariantViolation "RemoteEndPoint should be valid or failed." + t@(RemoteEndPointValid (ValidRemoteEndPoint sock (Counter x m) s z)) -> do + case ourId `Map.lookup` m of + Nothing -> return (t, return ()) -- XXX: send message to the hostv + Just c -> mask_ $ do + return (RemoteEndPointValid (ValidRemoteEndPoint sock (Counter x (ourId `Map.delete` m)) s (z+1)) + , do + modifyMVar_ (connectionState c) $ \case + ZMQConnectionFailed -> return ZMQConnectionFailed + ZMQConnectionInit -> return $ ZMQConnectionValid (ValidZMQConnection (Just sock) theirId) + ZMQConnectionClosed -> do + ZMQ.send sock [] $ encode' (MessageCloseConnection theirId) + return ZMQConnectionClosed + ZMQConnectionValid _ -> throwM $ InvariantViolation "RemoteEndPoint should be closed" + void $ tryPutMVar (connectionReady c) () + ) + RemoteEndPointPending p -> return (RemoteEndPointPending p, undefined) + where + r = _localEndPointRemotes v + LocalEndPointClosed -> return $ return () +-- printf "[%s] message init connection ok [ok]\n" +-- (B8.unpack $ endPointAddressToByteString ourAddr) + MessageEndPointClose theirAddress True -> getRemoteEndPoint ourEp theirAddress >>= \case + Nothing -> return () + Just rep -> do + onValidRemote rep $ \v -> + ZMQ.send (_remoteEndPointChan v) [] (encode' $ MessageEndPointCloseOk $ localEndPointAddress ourEp) + remoteEndPointClose True ourEp rep + MessageEndPointClose theirAddress False -> getRemoteEndPoint ourEp theirAddress >>= \case + Nothing -> return () + Just rep -> do + mst <- cleanupRemoteEndPoint ourEp rep Nothing + case mst of + Nothing -> return () + Just st -> do + onValidEndPoint ourEp $ \v -> atomically $ writeTMChan (_localEndPointChan v) $ + ErrorEvent $ TransportError (EventConnectionLost theirAddress) "Exception on remote side" + closeRemoteEndPoint ourEp rep st + MessageEndPointCloseOk theirAddress -> getRemoteEndPoint ourEp theirAddress >>= \case + Nothing -> return () + Just rep -> do + state <- swapMVar (remoteEndPointState rep) RemoteEndPointClosed + closeRemoteEndPoint ourEp rep state + where + ourAddr = localEndPointAddress ourEp + finalizeEndPoint ourEp port pull = do +-- printf "[%s] finalize-end-point\n" +-- (B8.unpack $ endPointAddressToByteString $ _localEndPointAddress ourEp) + join $ withMVar (localEndPointState ourEp) $ \case + LocalEndPointClosed -> afterP () + LocalEndPointValid v -> do + writeIORef (_localEndPointOpened v) False + return $ do + tid <- Async.async $ finalizer pull ourEp + void $ Async.mapConcurrently (remoteEndPointClose False ourEp) + $ _localEndPointRemotes v + Async.cancel tid + ZMQ.unbind pull (addr ++ ":" ++ show port) + ZMQ.close pull + void $ swapMVar (localEndPointState ourEp) LocalEndPointClosed + +apiSend :: ZMQConnection -> [ByteString] -> IO (Either (TransportError SendErrorCode) ()) +#ifdef UNSAFE_SEND +apiSend c@(ZMQConnection l e _ s _) b = mask_ $ join $ withMVar s $ \case + ZMQConnectionInit -> return $ yield >> apiSend c b + ZMQConnectionClosed -> afterP $ Left $ TransportError SendClosed "Connection is closed" + ZMQConnectionFailed -> afterP $ Left $ TransportError SendFailed "Connection is failed" + ZMQConnectionValid (ValidZMQConnection (Just ch) idx) -> do + o <- readIORef (remoteEndPointOpened e) + if o + then do + evs <- ZMQ.events ch + if ZMQ.Out `elem` evs + then do ZMQ.sendMulti ch $ encode' (MessageData idx) :| b + afterP $ Right () + else return $ do + mz <- cleanupRemoteEndPoint l e Nothing + case mz of + Nothing -> return () + Just z -> do + onValidEndPoint l $ \v -> atomically $ writeTMChan (_localEndPointChan v) $ + ErrorEvent $ TransportError (EventConnectionLost (remoteEndPointAddress e)) "Exception on remote side" + closeRemoteEndPoint l e z + return $ Left $ TransportError SendFailed "Connection broken." + else afterP $ Left $ TransportError SendFailed "Connection broken." + _ -> afterP $ Left $ TransportError SendFailed "Incorrect channel." +#else +apiSend (ZMQConnection l e _ s _) b = mask_ $ do + result <- trySome inner + case result of + Left ex -> do + cleanup + return $ Left $ TransportError SendFailed (show ex) + Right x -> return x + where + inner = join $ withMVar (remoteEndPointState e) $ \x -> case x of + RemoteEndPointFailed -> afterP $ Left $ TransportError SendFailed "Remote end point is failed." + RemoteEndPointClosed -> afterP $ Left $ TransportError SendFailed "Remote end point is closed." + RemoteEndPointClosing{} -> afterP $ Left $ TransportError SendFailed "Remote end point is closing." + RemoteEndPointPending{} -> return $ yield >> inner + RemoteEndPointValid v -> withMVar s $ \case + ZMQConnectionInit -> return $ yield >> inner -- readMVar (connectionReady c) >> inner + ZMQConnectionClosed -> afterP $ Left $ TransportError SendClosed "Connection is closed" + ZMQConnectionFailed -> afterP $ Left $ TransportError SendFailed "Connection is failed" + ZMQConnectionValid (ValidZMQConnection _ idx) -> do + evs <- ZMQ.events (_remoteEndPointChan v) + if ZMQ.Out `elem` evs + then do ZMQ.sendMulti (_remoteEndPointChan v) $ encode' (MessageData idx) :| b + afterP $ Right () + else return $ do + mz <- cleanupRemoteEndPoint l e Nothing + case mz of + Nothing -> return () + Just z -> do + onValidEndPoint l $ \w -> atomically $ writeTMChan (_localEndPointChan w) $ + ErrorEvent $ TransportError (EventConnectionLost (remoteEndPointAddress e)) "Exception on remote side" + closeRemoteEndPoint l e z + return $ Left $ TransportError SendFailed "Connection broken." + cleanup = do + void $ cleanupRemoteEndPoint l e + (Just $ \v -> ZMQ.send (_remoteEndPointChan v) [] $ encode' (MessageEndPointClose (localEndPointAddress l) False)) + onValidEndPoint l $ \v -> atomically $ do + writeTMChan (_localEndPointChan v) $ ErrorEvent $ TransportError + (EventConnectionLost (remoteEndPointAddress e)) "Exception on send." +#endif + + +-- 'apiClose' function is asynchronous, as connection may not exists by the +-- time of the calling to this function. In this case function just marks +-- connection as closed, so all subsequent calls from the user side will +-- "think" that the connection is closed, and remote side will be contified +-- only after connection will be up. +apiClose :: ZMQConnection -> IO () +apiClose (ZMQConnection _ e _ s _) = uninterruptibleMask_ $ join $ do + modifyMVar s $ \case + ZMQConnectionInit -> return (ZMQConnectionClosed, return ()) + ZMQConnectionClosed -> return (ZMQConnectionClosed, return ()) + ZMQConnectionFailed -> return (ZMQConnectionClosed, return ()) + ZMQConnectionValid (ValidZMQConnection _ idx) -> do + return (ZMQConnectionClosed, do + modifyMVar_ (remoteEndPointState e) $ \case + v@RemoteEndPointClosed -> return v + v@RemoteEndPointClosing{} -> return v + v@RemoteEndPointFailed -> return v + v@RemoteEndPointValid{} -> notify idx v + v@(RemoteEndPointPending p) -> modifyIORef p (\xs -> notify idx : xs) >> return v + ) + where + notify _ RemoteEndPointFailed = return RemoteEndPointFailed + notify _ (RemoteEndPointClosing x) = return $ RemoteEndPointClosing x + notify _ RemoteEndPointClosed = return RemoteEndPointClosed + notify _ RemoteEndPointPending{} = throwM $ InvariantViolation "RemoteEndPoint can't be pending." + notify idx w@(RemoteEndPointValid (ValidRemoteEndPoint sock _ _ _)) = do + ZMQ.send sock [] $ encode' (MessageCloseConnection idx) + return w + +apiConnect :: ZMQParameters + -> Context + -> LocalEndPoint + -> EndPointAddress + -> Reliability + -> ConnectHints + -> IO (Either (TransportError ConnectErrorCode) Connection) +apiConnect params ctx ourEp theirAddr reliability _hints = do +-- printf "[%s] apiConnect to %s\n" +-- (B8.unpack $ endPointAddressToByteString $ _localEndPointAddress ourEp) +-- (B8.unpack $ endPointAddressToByteString theirAddr) + eRep <- createOrGetRemoteEndPoint params ctx ourEp theirAddr + case eRep of + Left{} -> return $ Left $ TransportError ConnectFailed "LocalEndPoint is closed." + Right rep -> do + conn <- ZMQConnection <$> pure ourEp + <*> pure rep + <*> pure reliability + <*> newMVar ZMQConnectionInit + <*> newEmptyMVar + let apiConn = Connection + { send = apiSend conn + , close = apiClose conn + } + join $ modifyMVar (remoteEndPointState rep) $ \w -> case w of + RemoteEndPointClosed -> do + return ( RemoteEndPointClosed + , return $ Left $ TransportError ConnectFailed "Transport is closed.") + RemoteEndPointClosing x -> do + return ( RemoteEndPointClosing x + , return $ Left $ TransportError ConnectFailed "RemoteEndPoint closed.") + RemoteEndPointValid _ -> do + s' <- go conn w + return (s', waitReady conn apiConn) + RemoteEndPointPending z -> do + modifyIORef z (\zs -> go conn : zs) + return ( RemoteEndPointPending z, waitReady conn apiConn) + RemoteEndPointFailed -> + return ( RemoteEndPointFailed + , return $ Left $ TransportError ConnectFailed "RemoteEndPoint failed.") + where + waitReady conn apiConn = join $ withMVar (connectionState conn) $ \case + ZMQConnectionInit{} -> return $ yield {-readMVar (connectionReady conn)-} >> waitReady conn apiConn + ZMQConnectionValid{} -> afterP $ Right apiConn + ZMQConnectionFailed{} -> afterP $ Left $ TransportError ConnectFailed "Connection failed." + ZMQConnectionClosed{} -> throwM $ InvariantViolation "Connection closed." + ourAddr = localEndPointAddress ourEp + go conn s = inner conn s + inner _ RemoteEndPointClosed = return RemoteEndPointClosed + inner _ RemoteEndPointPending{} = throwM $ InvariantViolation "Connection pending." + inner _ (RemoteEndPointClosing x) = return $ RemoteEndPointClosing x + inner _ RemoteEndPointFailed = return RemoteEndPointFailed + inner conn (RemoteEndPointValid (ValidRemoteEndPoint sock (Counter i m) s z)) = do + ZMQ.send sock [] $ encode' (MessageInitConnection ourAddr i' reliability) + return $ RemoteEndPointValid (ValidRemoteEndPoint sock (Counter i' (Map.insert i' conn m)) s z) + where i' = succ i + + +getRemoteEndPoint :: LocalEndPoint -> EndPointAddress -> IO (Maybe RemoteEndPoint) +getRemoteEndPoint ourEp theirAddr = do + withMVar (localEndPointState ourEp) $ \case + LocalEndPointValid v -> return $ theirAddr `Map.lookup` (_localEndPointRemotes v) + LocalEndPointClosed -> return Nothing + +createOrGetRemoteEndPoint :: ZMQParameters + -> Context + -> LocalEndPoint + -> EndPointAddress + -> IO (Either ZMQError RemoteEndPoint) +createOrGetRemoteEndPoint params ctx ourEp theirAddr = join $ do +-- printf "[%s] apiConnect to %s\n" +-- saddr +-- (B8.unpack $ endPointAddressToByteString theirAddr) + modifyMVar (localEndPointState ourEp) $ \case + LocalEndPointValid v@(ValidLocalEndPoint _ _ m _ o _) -> do + opened <- readIORef o + if opened + then do + case theirAddr `Map.lookup` m of + Nothing -> create v m + Just rep -> do + withMVar (remoteEndPointState rep) $ \case + RemoteEndPointFailed -> create v m + _ -> return (LocalEndPointValid v, return $ Right rep) + else return (LocalEndPointValid v, return $ Left $ IncorrectState "EndPointClosing") + LocalEndPointClosed -> + return ( LocalEndPointClosed + , return $ Left $ IncorrectState "EndPoint is closed" + ) + where + create v m = do +-- printf "[%s] apiConnect: remoteEndPoint not found, creating %s\n" +-- (B8.unpack $ endPointAddressToByteString $ _localEndPointAddress ourEp) +-- (B8.unpack $ endPointAddressToByteString theirAddr) + push <- ZMQ.socket ctx ZMQ.Push + case authorizationType params of + ZMQNoAuth -> return () + ZMQAuthPlain p u -> do + ZMQ.setPlainPassword (ZMQ.restrict p) push + ZMQ.setPlainUserName (ZMQ.restrict u) push + state <- newMVar . RemoteEndPointPending =<< newIORef [] + opened <- newIORef False + let rep = RemoteEndPoint theirAddr state opened + return ( LocalEndPointValid v{ _localEndPointRemotes = Map.insert theirAddr rep m} + , initialize push rep >> return (Right rep)) + ourAddr = localEndPointAddress ourEp + initialize push rep = do + lock <- newEmptyMVar + x <- Async.async $ do + takeMVar lock + ZMQ.connect push (B8.unpack $ endPointAddressToByteString theirAddr) + monitor <- ZMQ.monitor [ZMQ.ConnectedEvent] ctx push + putMVar lock () + r <- Async.race (threadDelay 100000) (monitor True) + void $ monitor False + case r of + Left _ -> do + _ <- swapMVar (remoteEndPointState rep) RemoteEndPointFailed + Async.cancel x + ZMQ.closeZeroLinger push + Right _ -> do + ZMQ.send push [] $ encode' $ MessageConnect ourAddr + let v = ValidRemoteEndPoint push (Counter 0 Map.empty) Set.empty 0 + modifyMVar_ (remoteEndPointState rep) $ \case + RemoteEndPointPending p -> do + z <- foldM (\y f -> f y) (RemoteEndPointValid v) . Prelude.reverse =<< readIORef p + modifyIORef (remoteEndPointOpened rep) (const True) + return z + RemoteEndPointValid _ -> throwM $ InvariantViolation "RemoteEndPoint valid." + RemoteEndPointClosed -> return RemoteEndPointClosed + RemoteEndPointClosing j -> return (RemoteEndPointClosing j) + RemoteEndPointFailed -> return RemoteEndPointFailed + +-- | Close all endpoint connections, return previous state in case +-- if it was alive, for further cleanup actions. +cleanupRemoteEndPoint :: LocalEndPoint + -> RemoteEndPoint + -> (Maybe (ValidRemoteEndPoint -> IO ())) + -> IO (Maybe RemoteEndPointState) +cleanupRemoteEndPoint lep rep actions = modifyMVar (localEndPointState lep) $ \case + LocalEndPointValid v -> do + modifyIORef (remoteEndPointOpened rep) (const False) + oldState <- swapMVar (remoteEndPointState rep) newState + case oldState of + RemoteEndPointValid w -> do + let (Counter _ cn) = _remoteEndPointPendingConnections w + traverse_ (\c -> void $ swapMVar (connectionState c) ZMQConnectionFailed) cn + cn' <- foldM + (\(Counter i' cn') idx -> case idx `Map.lookup` cn of + Nothing -> return (Counter i' cn') + Just c -> do + void $ swapMVar (connectionState c) ZMQConnectionFailed + return $ Counter i' (Map.delete idx cn') + ) + (_localEndPointConnections v) + (Set.toList (_remoteEndPointIncommingConnections w)) + case actions of + Nothing -> return () + Just f -> f w + return ( LocalEndPointValid v { _localEndPointConnections=cn' } + , Just oldState) + _ -> return (LocalEndPointValid v, Nothing) + c -> return (c, Nothing) + where + newState = RemoteEndPointFailed + +-- | Close, all network connections. +closeRemoteEndPoint :: LocalEndPoint -> RemoteEndPoint -> RemoteEndPointState -> IO () +closeRemoteEndPoint lep rep state = step1 >> step2 state + where + step1 = modifyMVar_ (localEndPointState lep) $ \case + LocalEndPointValid v -> return $ + LocalEndPointValid v{_localEndPointRemotes=Map.delete + (remoteEndPointAddress rep) + (_localEndPointRemotes v)} + c -> return c + step2 (RemoteEndPointValid v) = do + ZMQ.closeZeroLinger (_remoteEndPointChan v) + step2 (RemoteEndPointClosing (ClosingRemoteEndPoint sock rd)) = do + _ <- readMVar rd + ZMQ.closeZeroLinger sock + step2 _ = return () + + +remoteEndPointClose :: Bool -> LocalEndPoint -> RemoteEndPoint -> IO () +remoteEndPointClose silent lep rep = do +-- printf "[???] remoteEndPointClose %s\n" +-- (B8.unpack $ endPointAddressToByteString $ lAddr) +-- (B8.unpack $ endPointAddressToByteString $ remoteEndPointAddress rep) + modifyIORef (remoteEndPointOpened rep) (const False) + join $ modifyMVar (remoteEndPointState rep) $ \o -> case o of + RemoteEndPointFailed -> return (o, return ()) + RemoteEndPointClosed -> return (o, return ()) + RemoteEndPointClosing (ClosingRemoteEndPoint _ l) -> return (o, void $ readMVar l) + RemoteEndPointPending _ -> closing undefined o -- XXX: store socket, or delay + RemoteEndPointValid v -> closing (_remoteEndPointChan v) o + where + closing sock old = do + lock <- newEmptyMVar + return (RemoteEndPointClosing (ClosingRemoteEndPoint sock lock), go lock old) + go lock old@(RemoteEndPointValid (ValidRemoteEndPoint c _ s i)) = do + -- close all connections + void $ cleanupRemoteEndPoint lep rep Nothing + withMVar (localEndPointState lep) $ \case + LocalEndPointClosed -> return () + LocalEndPointValid v -> do + -- notify about all connections close (?) do we really want it? + traverse_ (atomically . writeTMChan (_localEndPointChan v) . ConnectionClosed) (Set.toList s) + -- if we have outgoing connections, then we have connection error + when (i > 0) $ atomically + $ writeTMChan (_localEndPointChan v) + $ ErrorEvent $ TransportError (EventConnectionLost (remoteEndPointAddress rep)) "Remote end point closed." + -- notify other side about closing connection + unless silent $ do + ZMQ.send c [] (encode' (MessageEndPointClose (localEndPointAddress lep) True)) + yield + void $ Async.race (readMVar lock) (threadDelay 1000000) + void $ tryPutMVar lock () + closeRemoteEndPoint lep rep old + return () + go _ _ = return () + +connectionCleanup :: RemoteEndPoint -> ConnectionId -> IO () +connectionCleanup rep cid = modifyMVar_ (remoteEndPointState rep) $ \case + RemoteEndPointValid v -> return $ + RemoteEndPointValid v{_remoteEndPointIncommingConnections = Set.delete cid (_remoteEndPointIncommingConnections v)} + c -> return c + +encode' :: Binary a => a -> ByteString +encode' = B.concat . BL.toChunks . encode + +decode' :: Binary a => ByteString -> a +decode' s = decode . BL.fromChunks $ [s] + +onValidEndPoint :: LocalEndPoint -> (ValidLocalEndPoint -> IO ()) -> IO () +onValidEndPoint lep f = withMVar (localEndPointState lep) $ \case + LocalEndPointValid v -> f v + _ -> return () + +onValidRemote :: RemoteEndPoint -> (ValidRemoteEndPoint -> IO ()) -> IO () +onValidRemote rep f = withMVar (remoteEndPointState rep) $ \case + RemoteEndPointValid v -> f v + _ -> return () + +afterP :: a -> IO (IO a) +afterP = return . return + +trySome :: IO a -> IO (Either SomeException a) +trySome f = try f >>= \case + Left e -> case (fromException e) of + Just m -> throwM (m::AsyncException) + Nothing -> return $ Left e + Right x -> return $ Right x + +-- | Break endpoint connection, all endpoints that will be affected. +breakConnection :: ZMQTransport + -> EndPointAddress + -> EndPointAddress + -> IO () +breakConnection zmqt from to = Foldable.sequence_ <=< withMVar (_transportState zmqt) $ \case + TransportValid v -> Traversable.forM (_transportEndPoints v) $ \x -> + withMVar (localEndPointState x) $ \case + LocalEndPointValid w -> return $ Foldable.sequence_ $ flip Map.mapWithKey (_localEndPointRemotes w) $ \key rep -> + if onDeadHost key + then do + mz <- cleanupRemoteEndPoint x rep Nothing + flip traverse_ mz $ \z -> do + atomically $ writeTMChan (_localEndPointChan w) $ + ErrorEvent $ TransportError (EventConnectionLost key) "Manual connection break" + closeRemoteEndPoint x rep z + else return () + LocalEndPointClosed -> afterP () + TransportClosed -> return Map.empty + where + dh = B8.init . fst $ B8.spanEnd (/=':') (endPointAddressToByteString to) + onDeadHost = B8.isPrefixOf dh . endPointAddressToByteString + +-- | Break connection between two endpoints, other endpoints on the same +-- remove will not be affected. +breakConnectionEndPoint :: ZMQTransport + -> EndPointAddress + -> EndPointAddress + -> IO () +breakConnectionEndPoint zmqt from to = one from to >> one to from + where + one f t = join $ withMVar (_transportState zmqt) $ \case + TransportValid v -> case f `Map.lookup` _transportEndPoints v of + Nothing -> afterP () + Just x -> withMVar (localEndPointState x) $ \case + LocalEndPointValid w -> case t `Map.lookup` _localEndPointRemotes w of + Nothing -> afterP () + Just y -> return $ do + mz <- cleanupRemoteEndPoint x y Nothing + case mz of + Nothing -> return () + Just z -> do + onValidEndPoint x $ \j -> atomically $ writeTMChan (_localEndPointChan j) $ + ErrorEvent $ TransportError (EventConnectionLost to) "Exception on remote side" + closeRemoteEndPoint x y z + LocalEndPointClosed -> afterP () + TransportClosed -> afterP () + +-- | Configure socket after creation +unsafeConfigurePush :: ZMQTransport + -> EndPointAddress + -> EndPointAddress + -> (ZMQ.Socket ZMQ.Push -> IO ()) + -> IO () +unsafeConfigurePush zmqt from to f = withMVar (_transportState zmqt) $ \case + TransportValid v -> case from `Map.lookup` _transportEndPoints v of + Nothing -> return () + Just x -> withMVar (localEndPointState x) $ \case + LocalEndPointValid w -> case to `Map.lookup` _localEndPointRemotes w of + Nothing -> return () + Just y -> onValidRemote y $ f . _remoteEndPointChan + LocalEndPointClosed -> return () + TransportClosed -> return () + +apiNewMulticastGroup :: ZMQParameters -> ZMQTransport -> LocalEndPoint -> IO ( Either (TransportError NewMulticastGroupErrorCode) MulticastGroup) +apiNewMulticastGroup params zmq lep = withMVar (_transportState zmq) $ \case + TransportClosed -> return $ Left $ TransportError NewMulticastGroupFailed "Transport is closed." + TransportValid vt -> modifyMVar (localEndPointState lep) $ \case + LocalEndPointClosed -> return (LocalEndPointClosed, Left $ TransportError NewMulticastGroupFailed "Transport is closed.") + LocalEndPointValid vl -> mask_ $ do + (pub, portPub, rep, portRep, wrkThread) <- mkPublisher vt + + let addr = MulticastAddress $ + B8.concat [transportAddress zmq,":",B8.pack (show portPub), ":",B8.pack (show portRep)] + subAddr = extractPubAddress addr + repAddr = extractRepAddress addr + +-- printf "new multicastaddress: %s\n" (show addr) +-- printf "[%s] subscribe address: %s\n" (show addr) (B8.unpack subAddr) +-- printf "[%s] send address: %s\n" (show addr) (B8.unpack repAddr) + + -- subscriber api + sub <- ZMQ.socket (_transportContext vt) ZMQ.Sub + ZMQ.connect sub (B8.unpack subAddr) + ZMQ.subscribe sub "" + + subscribed <- newIORef False + (subThread, mgroup) <- fixIO $ \ ~(tid,_) -> do + v <- newMVar (MulticastGroupValid (ValidMulticastGroup subscribed)) + let mgroup = ZMQMulticastGroup v + (apiDeleteMulticastGroupLocal v lep addr rep repAddr pub sub subAddr (Just tid) wrkThread) + tid' <- Async.async $ do + repeatWhile $ do + (ctrl: msg) <- ZMQ.receiveMulti sub + if B8.null ctrl + then do opened <- readIORef subscribed + when opened $ atomically $ writeTMChan (_localEndPointChan vl) + $ ReceivedMulticast addr msg + return True + else return False + return (tid', mgroup) + return ( LocalEndPointValid vl + , Right $ MulticastGroup + { multicastAddress = addr + , deleteMulticastGroup = apiDeleteMulticastGroupLocal (multicastGroupState mgroup) lep addr rep repAddr pub sub subAddr (Just subThread) wrkThread + , maxMsgSize = Nothing + , multicastSend = apiMulticastSendLocal mgroup pub + , multicastSubscribe = apiMulticastSubscribe mgroup + , multicastUnsubscribe = apiMulticastUnsubscribe mgroup + , multicastClose = apiMulticastClose + } + ) + where + mkPublisher vt = do + pub <- ZMQ.socket (_transportContext vt) ZMQ.Pub + portPub <- ZMQ.bindFromRangeRandom pub (B8.unpack $ transportAddress zmq) (minPort params) (maxPort params) (maxTries params) + rep <- ZMQ.socket (_transportContext vt) ZMQ.Rep + portRep <- ZMQ.bindFromRangeRandom rep (B8.unpack $ transportAddress zmq) (minPort params) (maxPort params) (maxTries params) + wrkThread <- Async.async $ forever $ do + msg <- ZMQ.receiveMulti rep + ZMQ.sendMulti pub $ "" :| msg + ZMQ.send rep [] "OK" + return (pub,portPub,rep,portRep, wrkThread) + +apiResolveMulticastGroup :: ZMQTransport -> LocalEndPoint -> MulticastAddress -> IO (Either (TransportError ResolveMulticastGroupErrorCode) MulticastGroup) +apiResolveMulticastGroup zmq lep addr = withMVar (_transportState zmq) $ \case + TransportClosed -> return $ Left $ TransportError ResolveMulticastGroupFailed "Transport is closed." + TransportValid vt -> modifyMVar (localEndPointState lep) $ \case + LocalEndPointClosed -> return (LocalEndPointClosed, Left $ TransportError ResolveMulticastGroupFailed "Transport is closed.") + LocalEndPointValid vl -> mask_ $ do + -- socket allocation + req <- ZMQ.socket (_transportContext vt) ZMQ.Req + ZMQ.connect req (B8.unpack reqAddr) + sub <- ZMQ.socket (_transportContext vt) ZMQ.Sub + ZMQ.connect sub (B8.unpack subAddr) + ZMQ.subscribe sub "" + + subscribed <- newIORef False + (subThread, mgroup) <- fixIO $ \ ~(tid,_) -> do + v <- newMVar (MulticastGroupValid (ValidMulticastGroup subscribed)) + let mgroup = ZMQMulticastGroup v + (apiDeleteMulticastGroupRemote v lep addr req reqAddr sub subAddr (Just tid)) + tid' <- Async.async $ repeatWhile $ do + (ctrl: msg) <- ZMQ.receiveMulti sub + if B8.null ctrl + then do opened <- readIORef subscribed + when opened $ atomically $ writeTMChan (_localEndPointChan vl) + $ ReceivedMulticast addr msg + return True + else do apiDeleteMulticastGroupRemote v lep addr req reqAddr sub subAddr Nothing + return False + return (tid', mgroup) + + return ( LocalEndPointValid + vl{_localEndPointMulticastGroups = Map.insert addr mgroup (_localEndPointMulticastGroups vl)} + , Right $ MulticastGroup + { multicastAddress = addr + , deleteMulticastGroup = + apiDeleteMulticastGroupRemote (multicastGroupState mgroup) lep addr req reqAddr sub subAddr (Just subThread) + , maxMsgSize = Nothing + , multicastSend = apiMulticastSendRemote mgroup req + , multicastSubscribe = apiMulticastSubscribe mgroup + , multicastUnsubscribe = apiMulticastUnsubscribe mgroup + , multicastClose = apiMulticastClose + } + ) + where + reqAddr = extractRepAddress addr + subAddr = extractPubAddress addr + + +apiDeleteMulticastGroupRemote :: MVar MulticastGroupState -> LocalEndPoint -> MulticastAddress + -> ZMQ.Socket ZMQ.Req -> ByteString -> ZMQ.Socket ZMQ.Sub -> ByteString -> Maybe (Async.Async ()) + -> IO () +apiDeleteMulticastGroupRemote mstate lep addr req reqAddr sub subAddr mtid = mask_ $ do + Foldable.traverse_ (\tid -> Async.cancel tid >> void (Async.waitCatch tid)) mtid + modifyMVar_ mstate $ \case + MulticastGroupClosed -> return MulticastGroupClosed + MulticastGroupValid v -> do + modifyIORef (_multicastGroupSubscribed v) (const False) +-- ZMQ.disconnect req (B8.unpack reqAddr) + ZMQ.close req + ZMQ.unsubscribe sub "" +-- ZMQ.disconnect sub (B8.unpack subAddr) + ZMQ.close sub + return MulticastGroupClosed + modifyMVar_ (localEndPointState lep) $ \case + LocalEndPointClosed -> return LocalEndPointClosed + LocalEndPointValid v -> return $ LocalEndPointValid + v{_localEndPointMulticastGroups = Map.delete addr (_localEndPointMulticastGroups v)} + +apiDeleteMulticastGroupLocal :: MVar MulticastGroupState -> LocalEndPoint -> MulticastAddress + -> ZMQ.Socket ZMQ.Rep -> ByteString -> ZMQ.Socket ZMQ.Pub -> ZMQ.Socket ZMQ.Sub -> ByteString -> Maybe (Async.Async ()) + -> Async.Async () + -> IO () +apiDeleteMulticastGroupLocal mstate lep addr rep repAddr pub sub pubAddr mtid wrk = mask_ $ do + Foldable.traverse_ (\tid -> Async.cancel tid >> void (Async.waitCatch tid)) mtid + void $ Async.cancel wrk >> Async.waitCatch wrk + modifyMVar_ mstate $ \case + MulticastGroupClosed -> return MulticastGroupClosed + MulticastGroupValid v -> do + modifyIORef (_multicastGroupSubscribed v) (const False) + ZMQ.sendMulti pub $ "C" :| [] + threadDelay 50000 + ZMQ.unbind rep (B8.unpack repAddr) + ZMQ.close rep + ZMQ.unsubscribe sub "" +-- ZMQ.disconnect sub (B8.unpack pubAddr) + ZMQ.close sub + ZMQ.unbind pub (B8.unpack pubAddr) + ZMQ.close pub + return MulticastGroupClosed + modifyMVar_ (localEndPointState lep) $ \case + LocalEndPointClosed -> return LocalEndPointClosed + LocalEndPointValid v -> return $ LocalEndPointValid + v{_localEndPointMulticastGroups = Map.delete addr (_localEndPointMulticastGroups v)} + +apiMulticastSendLocal :: (ZMQ.Sender a) => ZMQMulticastGroup -> ZMQ.Socket a -> [ByteString] -> IO () +apiMulticastSendLocal mgroup sock msg = withMVar (multicastGroupState mgroup) $ \case + MulticastGroupClosed -> return () + MulticastGroupValid _ -> do + ZMQ.sendMulti sock $ "" :| msg + +apiMulticastSendRemote :: (ZMQ.Sender a, ZMQ.Receiver a) => ZMQMulticastGroup -> ZMQ.Socket a -> [ByteString] -> IO () +apiMulticastSendRemote _ _ [] = return () +apiMulticastSendRemote mgroup sock (m:msg) = withMVar (multicastGroupState mgroup) $ \case + MulticastGroupClosed -> return () + MulticastGroupValid _ -> do + ZMQ.sendMulti sock $ m :| msg + "OK" <- ZMQ.receive sock -- XXX: timeout + return () + +apiMulticastSubscribe :: ZMQMulticastGroup -> IO () +apiMulticastSubscribe mgroup = withMVar (multicastGroupState mgroup) $ \case + MulticastGroupClosed -> return () + MulticastGroupValid v -> do + modifyIORef (_multicastGroupSubscribed v) (const True) + +apiMulticastUnsubscribe :: ZMQMulticastGroup -> IO () +apiMulticastUnsubscribe mgroup = withMVar (multicastGroupState mgroup) $ \case + MulticastGroupClosed -> return () + MulticastGroupValid v -> modifyIORef (_multicastGroupSubscribed v) (const False) + +apiMulticastClose :: IO () +apiMulticastClose = return () + +extractRepAddress :: MulticastAddress -> ByteString +extractRepAddress (MulticastAddress bs) = B8.concat [a,":",p] + where + (x,p) = B8.spanEnd (/=':') bs + a = B8.init . fst . B8.spanEnd (/=':') . B8.init $ x +extractPubAddress :: MulticastAddress -> ByteString +extractPubAddress (MulticastAddress bs) = B8.init . fst $ B8.spanEnd (/=':') bs + +repeatWhile :: Monad m => (m Bool) -> m () +repeatWhile f = f >>= \x -> if x then repeatWhile f else return () diff --git a/src/Network/Transport/ZMQ/Types.hs b/src/Network/Transport/ZMQ/Types.hs new file mode 100644 --- /dev/null +++ b/src/Network/Transport/ZMQ/Types.hs @@ -0,0 +1,222 @@ +-- | +-- Module: Network.Transport.ZMQ.Types +-- Copyright: (C) 2014, EURL Tweag +-- Licese: BSD-3 +-- +{-# LANGUAGE LambdaCase #-} +module Network.Transport.ZMQ.Types + ( ZMQParameters(..) + , ZMQAuthType(..) + , defaultZMQParameters + -- * Internal types + , ZMQTransport(..) + , TransportState(..) + , ValidTransportState(..) + -- ** RemoteEndPoint + , RemoteEndPoint(..) + , RemoteEndPointState(..) + , ValidRemoteEndPoint(..) + , ClosingRemoteEndPoint(..) + -- ** LocalEndPoint + , LocalEndPoint(..) + , LocalEndPointState(..) + , ValidLocalEndPoint(..) + -- ** ZeroMQ connection + , ZMQConnection(..) + , ZMQConnectionState(..) + , ValidZMQConnection(..) + -- ** ZeroMQ multicast + , ZMQMulticastGroup(..) + , MulticastGroupState(..) + , ValidMulticastGroup(..) + -- * Internal data structures + , Counter(..) + , nextElement + , nextElement' + , nextElementM + , nextElementM' + ) where + +import Control.Concurrent.Async +import Control.Concurrent.MVar +import Control.Concurrent.STM.TMChan +import Data.Word +import Data.ByteString +import Data.IORef +import Data.Map.Strict (Map) +import qualified Data.Map.Strict as M +import Data.Set + ( Set + ) +import System.ZMQ4 + ( Socket + , Push + ) + +import Network.Transport + +import qualified System.ZMQ4 as ZMQ + + +-- | Parameters for ZeroMQ connection +data ZMQParameters = ZMQParameters + { highWaterMark :: Word64 -- uint64_t + , authorizationType :: ZMQAuthType + , minPort :: Int + , maxPort :: Int + , maxTries :: Int + } +-- High Watermark + +defaultZMQParameters :: ZMQParameters +defaultZMQParameters = ZMQParameters + { highWaterMark = 0 + , authorizationType = ZMQNoAuth + , minPort = 40000 + , maxPort = 60000 + , maxTries = 10000 + } + +data ZMQAuthType + = ZMQNoAuth + | ZMQAuthPlain + { zmqAuthPlainPassword :: ByteString + , zmqAutnPlainUserName :: ByteString + } + +type TransportAddress = ByteString + +-- | Transport data type. +data ZMQTransport = ZMQTransport + { transportAddress :: !TransportAddress + -- ^ Transport address (used as identifier). + , _transportState :: !(MVar TransportState) + -- ^ Internal state. + } + +-- | Transport state. +data TransportState + = TransportValid !ValidTransportState -- ^ Transport is in active state. + | TransportClosed -- ^ Transport is closed. + +-- | Transport state. +data ValidTransportState = ValidTransportState + { _transportContext :: !ZMQ.Context + , _transportEndPoints :: !(Map EndPointAddress LocalEndPoint) + , _transportAuth :: !(Maybe (Async ())) + } + +data LocalEndPoint = LocalEndPoint + { localEndPointAddress :: !EndPointAddress + , localEndPointState :: !(MVar LocalEndPointState) + , localEndPointPort :: !(Int) + } + +data LocalEndPointState + = LocalEndPointValid !ValidLocalEndPoint + | LocalEndPointClosed + +data ValidLocalEndPoint = ValidLocalEndPoint + { _localEndPointChan :: !(TMChan Event) + -- ^ channel for n-t - user communication + , _localEndPointConnections :: !(Counter ConnectionId ZMQConnection) + -- ^ list of incomming connections + , _localEndPointRemotes :: !(Map EndPointAddress RemoteEndPoint) + -- ^ list of remote end points + , _localEndPointThread :: !(Async ()) + -- ^ thread id + , _localEndPointOpened :: !(IORef Bool) + -- ^ is remote endpoint opened + , _localEndPointMulticastGroups :: !(Map MulticastAddress ZMQMulticastGroup) + -- ^ list of multicast nodes + } + +data ZMQConnection = ZMQConnection + { connectionLocalEndPoint :: !LocalEndPoint + , connectionRemoteEndPoint :: !RemoteEndPoint + , connectionReliability :: !Reliability + , connectionState :: !(MVar ZMQConnectionState) + , connectionReady :: !(MVar ()) + } + +data ZMQConnectionState + = ZMQConnectionInit + | ZMQConnectionValid !ValidZMQConnection + | ZMQConnectionClosed + | ZMQConnectionFailed + +data ValidZMQConnection = ValidZMQConnection + { _connectionSocket :: !(Maybe (ZMQ.Socket ZMQ.Push)) + , _connectionId :: !Word64 + } + +data RemoteEndPoint = RemoteEndPoint + { remoteEndPointAddress :: !EndPointAddress + , remoteEndPointState :: !(MVar RemoteEndPointState) + , remoteEndPointOpened :: !(IORef Bool) + } + +data RemoteEndPointState + = RemoteEndPointValid ValidRemoteEndPoint + | RemoteEndPointClosed + | RemoteEndPointFailed + | RemoteEndPointPending (IORef [RemoteEndPointState -> IO RemoteEndPointState]) + | RemoteEndPointClosing ClosingRemoteEndPoint + +data ValidRemoteEndPoint = ValidRemoteEndPoint + { _remoteEndPointChan :: !(Socket Push) + , _remoteEndPointPendingConnections :: !(Counter ConnectionId ZMQConnection) + , _remoteEndPointIncommingConnections :: !(Set ConnectionId) + , _remoteEndPointOutgoingCount :: !Int + } + +data ClosingRemoteEndPoint = ClosingRemoteEndPoint + { _remoteEndPointClosingSocket :: !(Socket Push) + , _remoteEndPointDone :: !(MVar ()) + } + +data ZMQMulticastGroup = ZMQMulticastGroup + { multicastGroupState :: MVar MulticastGroupState + , multicastGroupClose :: IO () + } + +data MulticastGroupState + = MulticastGroupValid ValidMulticastGroup + | MulticastGroupClosed + +data ValidMulticastGroup = ValidMulticastGroup + { _multicastGroupSubscribed :: IORef Bool + } + +data Counter a b = Counter { counterNext :: !a + , counterValue :: !(Map a b) + } + +nextElement :: (Enum a, Ord a) => (b -> IO Bool) -> b -> Counter a b -> IO (Counter a b, (a, b)) +nextElement t e c = nextElement' t (const e) c + +nextElement' :: (Enum a, Ord a) => (b -> IO Bool) -> (a -> b) -> Counter a b -> IO (Counter a b, (a,b)) +nextElement' t e c = nextElementM t (return . e) c + +nextElementM :: (Enum a, Ord a) => (b -> IO Bool) -> (a -> IO b) -> Counter a b -> IO (Counter a b, (a,b)) +nextElementM t me (Counter n m) = + case n' `M.lookup` m of + Nothing -> mv >>= \v' -> return (Counter n' (M.insert n' v' m), (n', v')) + Just v -> t v >>= \case + True -> mv >>= \v' -> return (Counter n' (M.insert n' v' m), (n', v')) + False -> nextElementM t me (Counter n' m) + where + n' = succ n + mv = me n' + +nextElementM' :: (Enum a, Ord a) => (b -> IO Bool) -> (a -> IO (b,c)) -> Counter a b -> IO (Counter a b, (a,c)) +nextElementM' t me (Counter n m) = + case n' `M.lookup` m of + Nothing -> mv >>= \(v',r) -> return (Counter n' (M.insert n' v' m), (n', r)) + Just v -> t v >>= \case + True -> mv >>= \(v',r) -> return (Counter n' (M.insert n' v' m), (n', r)) + False -> nextElementM' t me (Counter n' m) + where + n' = succ n + mv = me n' + diff --git a/src/System/ZMQ4/Utils.hs b/src/System/ZMQ4/Utils.hs new file mode 100644 --- /dev/null +++ b/src/System/ZMQ4/Utils.hs @@ -0,0 +1,87 @@ +{-# LANGUAGE OverloadedStrings #-} +module System.ZMQ4.Utils + ( bindFromRangeRandom + , bindFromRangeRandomM + , authManager + , closeZeroLinger + ) + where + +import Control.Monad ( forever ) +import Control.Monad.Catch +import Control.Concurrent.Async +import Data.ByteString +import Data.List.NonEmpty +import System.Random ( randomRIO ) +import Text.Printf + + +import System.ZMQ4.Monadic + ( ZMQ + ) +import qualified System.ZMQ4.Monadic as M +import System.ZMQ4 + ( Context + , errno + ) +import qualified System.ZMQ4 as P + +-- | Bind socket to the random port in a given range. +bindFromRangeRandomM :: M.Socket z t + -> String -- ^ Address + -> Int -- ^ Min port + -> Int -- ^ Max port + -> Int -- ^ Max tries + -> ZMQ z Int +bindFromRangeRandomM sock addr mI mA tr = go tr + where + go 0 = error "!" -- XXX: throw correct error + go x = do + port <- M.liftIO $ randomRIO (mI,mA) + result <- try $ M.bind sock (printf "%s:%i" addr port) + case result of + Left e + | errno e == -1 -> go (x - 1) + | otherwise -> throwM e + Right () -> return port + +bindFromRangeRandom :: P.Socket t + -> String -- ^ Address + -> Int -- ^ Min port + -> Int -- ^ Max port + -> Int -- ^ Max tries + -> IO Int +bindFromRangeRandom sock addr mI mA tr = go tr + where + go 0 = error "!" -- XXX: throw correct error + go x = do + port <- randomRIO (mI,mA) + result <- try $ P.bind sock (printf "%s:%i" addr port) + case result of + Left e + | errno e == -1 -> go (x - 1) + | otherwise -> throwM e + Right () -> return port + +-- | One possible password authentification +authManager :: P.Context -> ByteString -> ByteString -> IO (Async ()) +authManager ctx user pass = do + req <- P.socket ctx P.Rep + P.bind req "inproc://zeromq.zap.01" + async $ forever $ do + ("1.0":requestId:domain:ipAddress:identity:mech:xs) <- P.receiveMulti req + case mech of + "PLAIN" -> case xs of + (pass':user':_) + | user == user' && pass == pass' -> do + P.sendMulti req $ "1.0" :| [requestId, "200", "OK", "", ""] + | otherwise -> P.sendMulti req $ "1.0" :| [requestId, "400", "Credentials are not implemented", "", ""] + _ -> P.sendMulti req $ "1.0" :| [requestId, "500", "Method not implemented", "", ""] + _ -> P.sendMulti req $ "1.0" :| [requestId, "500", "Method not implemented", "", ""] + + +-- | Close socket immideately. +closeZeroLinger :: P.Socket a -> IO () +closeZeroLinger sock = do + P.setLinger (P.restrict 0) sock + P.close sock diff --git a/tests/TestAPI.hs b/tests/TestAPI.hs new file mode 100644 --- /dev/null +++ b/tests/TestAPI.hs @@ -0,0 +1,97 @@ +{-# LANGUAGE OverloadedStrings #-} +module Main where + +import Control.Concurrent +import Control.Exception +import Control.Monad ( (<=<), replicateM ) + +import Network.Transport +import Network.Transport.ZMQ +import System.Exit + +main :: IO () +main = finish <=< trySome $ do + putStr "simple: " + Right (zmq, transport) <- createTransportEx defaultZMQParameters "127.0.0.1" + Right ep1 <- newEndPoint transport + Right ep2 <- newEndPoint transport + Right c1 <- connect ep1 (address ep2) ReliableOrdered defaultConnectHints + Right c2 <- connect ep2 (address ep1) ReliableOrdered defaultConnectHints + Right _ <- send c1 ["123"] + Right _ <- send c2 ["321"] + close c1 + close c2 + [ConnectionOpened 1 ReliableOrdered _, Received 1 ["321"], ConnectionClosed 1] <- replicateM 3 $ receive ep1 + [ConnectionOpened 1 ReliableOrdered _, Received 1 ["123"], ConnectionClosed 1] <- replicateM 3 $ receive ep2 + putStrLn "OK" + + putStr "connection break: " + Right ep3 <- newEndPoint transport + + Right c21 <- connect ep1 (address ep2) ReliableOrdered defaultConnectHints + Right c22 <- connect ep2 (address ep1) ReliableOrdered defaultConnectHints + Right c23 <- connect ep3 (address ep1) ReliableOrdered defaultConnectHints + + ConnectionOpened 2 ReliableOrdered _ <- receive ep2 + ConnectionOpened 2 ReliableOrdered _ <- receive ep1 + ConnectionOpened 3 ReliableOrdered _ <- receive ep1 + + breakConnection zmq (address ep1) (address ep2) + + ErrorEvent (TransportError (EventConnectionLost _ ) _) <- receive $ ep1 + ErrorEvent (TransportError (EventConnectionLost _ ) _) <- receive $ ep2 + + Left (TransportError SendFailed _) <- send c21 ["test"] + Left (TransportError SendFailed _) <- send c22 ["test"] + + Left (TransportError SendFailed _) <- send c23 ["test"] + ErrorEvent (TransportError (EventConnectionLost _) _ ) <- receive ep1 + Right c24 <- connect ep3 (address ep1) ReliableOrdered defaultConnectHints + Right () <- send c24 ["final"] + ConnectionOpened 4 ReliableOrdered _ <- receive ep1 + Received 4 ["final"] <- receive ep1 + putStrLn "OK" + multicast transport + where + multicast transport = do + Right ep1 <- newEndPoint transport + Right ep2 <- newEndPoint transport + putStr "multicast: " + Right g1 <- newMulticastGroup ep1 + multicastSubscribe g1 + threadDelay 1000000 + multicastSend g1 ["test"] + ReceivedMulticast _ ["test"] <- receive ep1 + Right g2 <- resolveMulticastGroup ep2 (multicastAddress g1) + multicastSubscribe g2 + threadDelay 100000 + multicastSend g2 ["test-2"] + ReceivedMulticast _ ["test-2"] <- receive ep2 + ReceivedMulticast _ ["test-2"] <- receive ep1 + putStrLn "OK" + + putStr "Auth:" + Right tr2 <- createTransport + defaultZMQParameters {authorizationType=ZMQAuthPlain "user" "password"} + "127.0.0.1" + Right ep3 <- newEndPoint tr2 + Right ep4 <- newEndPoint tr2 + Right c3 <- connect ep3 (address ep4) ReliableOrdered defaultConnectHints + Right _ <- send c3 ["4456"] + [ConnectionOpened 1 ReliableOrdered _, Received 1 ["4456"]] <- replicateM 2 $ receive ep4 + Right c4 <- connect ep3 (address ep4) ReliableOrdered defaultConnectHints + Right _ <- send c4 ["5567"] + [ConnectionOpened 2 ReliableOrdered _, Received 2 ["5567"]] <- replicateM 2 $ receive ep4 + putStrLn "OK" + + putStr "Connect to non existing host: " + Left (TransportError ConnectFailed _) <- connect ep3 (EndPointAddress "tcp://128.0.0.1:7689") ReliableOrdered defaultConnectHints + putStrLn "OK" + finish (Left e) = print e >> exitWith (ExitFailure 1) + finish Right{} = exitWith ExitSuccess + +trySome :: IO a -> IO (Either SomeException a) +trySome = try + +forkTry :: IO () -> IO ThreadId +forkTry = forkIO diff --git a/tests/TestZMQ.hs b/tests/TestZMQ.hs new file mode 100644 --- /dev/null +++ b/tests/TestZMQ.hs @@ -0,0 +1,38 @@ +{-# LANGUAGE OverloadedStrings #-} +module Main where + +import Control.Applicative ((<$>)) + +import Network.Transport +import Network.Transport.ZMQ +import Network.Transport.Tests +import Network.Transport.Tests.Auxiliary (runTests) + +main :: IO () +main = testTransport' (either (Left . show) (Right) <$> createTransport defaultZMQParameters "127.0.0.1") + +testTransport' :: IO (Either String Transport) -> IO () +testTransport' newTransport = do + Right transport <- newTransport + runTests + [ ("PingPong", testPingPong transport numPings) + , ("EndPoints", testEndPoints transport numPings) + , ("Connections", testConnections transport numPings) + , ("CloseOneConnection", testCloseOneConnection transport numPings) + , ("CloseOneDirection", testCloseOneDirection transport numPings) + , ("CloseReopen", testCloseReopen transport numPings) + , ("ParallelConnects", testParallelConnects transport 10) + , ("SendAfterClose", testSendAfterClose transport 100) + , ("Crossing", testCrossing transport 10) + , ("CloseTwice", testCloseTwice transport 100) + , ("ConnectToSelf", testConnectToSelf transport numPings) + , ("ConnectToSelfTwice", testConnectToSelfTwice transport numPings) + , ("CloseSelf", testCloseSelf newTransport) + , ("CloseEndPoint", testCloseEndPoint transport numPings) + , ("CloseTransport", testCloseTransport newTransport) + , ("ExceptionOnReceive", testExceptionOnReceive newTransport) + , ("SendException", testSendException newTransport) + , ("Kill", testKill newTransport 1000) + ] + where + numPings = 10000 :: Int diff --git a/tests/test-ch.hs b/tests/test-ch.hs new file mode 100644 --- /dev/null +++ b/tests/test-ch.hs @@ -0,0 +1,34 @@ +{-# LANGUAGE OverloadedStrings #-} +-- | +-- Copyright: (c) 2014, EURL Tweag +-- License: BSD3 +-- +-- Runner for given CH test suite. The test suite to run is given by the +-- expansion of the @TEST_SUITE_MODULE@ macro. + +module Main where + +import TEST_SUITE_MODULE (tests) + +import Network.Transport.Test (TestTransport(..)) +import Network.Transport.ZMQ + ( createTransportEx + , defaultZMQParameters + , breakConnection + ) +import Control.Concurrent.MVar +import Control.Concurrent.STM +import Control.Concurrent.STM.TMChan +import qualified Data.Map as Map +import qualified Data.ByteString.Char8 as B +import Network.Transport +import Network.Transport.ZMQ.Types +import Test.Framework (defaultMain) + +main :: IO () +main = do + Right (zmqt, transport) <- createTransportEx defaultZMQParameters "127.0.0.1" + defaultMain =<< tests TestTransport + { testTransport = transport + , testBreakConnection = breakConnection zmqt + }