git.iko.soy/distributed-process
on master help sign in
Add transport and endpoint closing.
4d91083Nicolas Wu
examples/DemoProcess.hs+3−1
compare with master
unifiedsplit
@@ -3,6 +3,7 @@ module DemoProcess where
3import Control.Distributed.Process
3import Control.Distributed.Process
4import Control.Concurrent (threadDelay)
4import Control.Concurrent (threadDelay)
5
5
6+import Network.Transport (closeTransport)
6import Network.Transport.TCP (mkTransport, TCPConfig (..))
7import Network.Transport.TCP (mkTransport, TCPConfig (..))
7
8
8import Debug.Trace
9import Debug.Trace
@@ -16,8 +17,9 @@ demo1 = do
16 lpid1 <- spawnLocal (logger 1)
17 lpid1 <- spawnLocal (logger 1)
17 spawnLocal (chatty "jim1" lpid1)
18 spawnLocal (chatty "jim1" lpid1)
18 spawnLocal (chatty "bob1" lpid1)
19 spawnLocal (chatty "bob1" lpid1)
19− liftIO $ threadDelay 500000
20 return ()
20 return ()
21+ threadDelay 600000
22+ closeTransport trans
21
23
22chatty :: String -> ProcessId -> Process ()
24chatty :: String -> ProcessId -> Process ()
23chatty name target = do
25chatty name target = do
examples/DemoTransport.hs+22−6
@@ -30,7 +30,7 @@ demo0 = do
30 mapM_ (\n -> send sourceEnd [BS.pack ("hello " ++ show n)]) [1 .. 10]
30 mapM_ (\n -> send sourceEnd [BS.pack ("hello " ++ show n)]) [1 .. 10]
31 threadDelay 100000
31 threadDelay 100000
32
32
33− closeTargetEnd targetEnd
33+ closeTransport trans
34
34
35-- | Check endpoint serialization and deserialization.
35-- | Check endpoint serialization and deserialization.
36demo1 :: IO ()
36demo1 :: IO ()
@@ -52,6 +52,7 @@ demo1 = do
52 send sourceEnd [BS.pack "hello 1"]
52 send sourceEnd [BS.pack "hello 1"]
53 send sourceEnd' [BS.pack "hello 2"]
53 send sourceEnd' [BS.pack "hello 2"]
54 threadDelay 100000
54 threadDelay 100000
55+ closeTransport trans
55
56
56-- | Check that messages can be sent before receive is set up.
57-- | Check that messages can be sent before receive is set up.
57demo2 :: IO ()
58demo2 :: IO ()
@@ -66,7 +67,9 @@ demo2 = do
66 threadDelay 100000
67 threadDelay 100000
67
68
68 forkIO $ logServer "logServer" targetEnd
69 forkIO $ logServer "logServer" targetEnd
69− return ()
70+ threadDelay 100000
71+
72+ closeTransport trans
70
73
71-- | Check that two different transports can be created.
74-- | Check that two different transports can be created.
72demo3 :: IO ()
75demo3 :: IO ()
@@ -77,8 +80,8 @@ demo3 = do
77 (sourceAddr1, targetEnd1) <- newConnection trans1
80 (sourceAddr1, targetEnd1) <- newConnection trans1
78 (sourceAddr2, targetEnd2) <- newConnection trans2
81 (sourceAddr2, targetEnd2) <- newConnection trans2
79
82
80− forkIO $ logServer "logServer1" targetEnd1
83+ threadId1 <- forkIO $ logServer "logServer1" targetEnd1
81− forkIO $ logServer "logServer2" targetEnd2
84+ threadId2 <- forkIO $ logServer "logServer2" targetEnd2
82
85
83 sourceEnd1 <- connect sourceAddr1
86 sourceEnd1 <- connect sourceAddr1
84 sourceEnd2 <- connect sourceAddr2
87 sourceEnd2 <- connect sourceAddr2
@@ -86,6 +89,13 @@ demo3 = do
86 send sourceEnd1 [BS.pack "hello1"]
89 send sourceEnd1 [BS.pack "hello1"]
87 send sourceEnd2 [BS.pack "hello2"]
90 send sourceEnd2 [BS.pack "hello2"]
88
91
92+ threadDelay 100000
93+
94+ killThread threadId1
95+ killThread threadId2
96+ closeTransport trans1
97+ closeTransport trans2
98+
89-- | Check that two different connections on the same transport can be created.
99-- | Check that two different connections on the same transport can be created.
90demo4 :: IO ()
100demo4 :: IO ()
91demo4 = do
101demo4 = do
@@ -95,8 +105,8 @@ demo4 = do
95 (sourceAddr1, targetEnd1) <- newConnection trans
105 (sourceAddr1, targetEnd1) <- newConnection trans
96 (sourceAddr2, targetEnd2) <- newConnection trans
106 (sourceAddr2, targetEnd2) <- newConnection trans
97
107
98− forkIO $ logServer "logServer1" targetEnd1
108+ threadId1 <- forkIO $ logServer "logServer1" targetEnd1
99− forkIO $ logServer "logServer2" targetEnd2
109+ threadId2 <- forkIO $ logServer "logServer2" targetEnd2
100
110
101 sourceEnd1 <- connect sourceAddr1
111 sourceEnd1 <- connect sourceAddr1
102 sourceEnd2 <- connect sourceAddr2
112 sourceEnd2 <- connect sourceAddr2
@@ -104,6 +114,12 @@ demo4 = do
104 send sourceEnd1 [BS.pack "hello1"]
114 send sourceEnd1 [BS.pack "hello1"]
105 send sourceEnd2 [BS.pack "hello2"]
115 send sourceEnd2 [BS.pack "hello2"]
106
116
117+ threadDelay 100000
118+
119+ killThread threadId1
120+ killThread threadId2
121+ closeTransport trans
122+
107logServer :: String -> TargetEnd -> IO ()
123logServer :: String -> TargetEnd -> IO ()
108logServer name targetEnd = forever $ do
124logServer name targetEnd = forever $ do
109 x <- receive targetEnd
125 x <- receive targetEnd
network-transport/src/Network/Transport.hs+1
@@ -33,6 +33,7 @@ data Transport = Transport
33 { newConnectionWith :: Hints -> IO (SourceAddr, TargetEnd)
33 { newConnectionWith :: Hints -> IO (SourceAddr, TargetEnd)
34 , newMulticastWith :: Hints -> IO (MulticastSourceEnd, MulticastTargetAddr)
34 , newMulticastWith :: Hints -> IO (MulticastSourceEnd, MulticastTargetAddr)
35 , deserialize :: ByteString -> Maybe SourceAddr
35 , deserialize :: ByteString -> Maybe SourceAddr
36+ , closeTransport :: IO ()
36 }
37 }
37
38
38-- | This is a convenience function that creates a new connection on a transport
39-- | This is a convenience function that creates a new connection on a transport
network-transport/src/Network/Transport/TCP.hs+68−33
@@ -5,10 +5,10 @@ module Network.Transport.TCP
5
5
6import Network.Transport
6import Network.Transport
7
7
8−import Control.Concurrent (forkIO)
8+import Control.Concurrent (forkIO, ThreadId, killThread)
9import Control.Concurrent.Chan
9import Control.Concurrent.Chan
10import Control.Concurrent.MVar
10import Control.Concurrent.MVar
11−import Control.Monad (forever)
11+import Control.Monad (forever, forM_)
12import Data.ByteString.Lazy.Char8 (ByteString)
12import Data.ByteString.Lazy.Char8 (ByteString)
13import Data.IntMap (IntMap)
13import Data.IntMap (IntMap)
14import Data.Int
14import Data.Int
@@ -25,8 +25,10 @@ import qualified Data.IntMap as IntMap
25import qualified Network.Socket as N
25import qualified Network.Socket as N
26import qualified Network.Socket.ByteString.Lazy as NBS
26import qualified Network.Socket.ByteString.Lazy as NBS
27
27
28+import Debug.Trace
29+
28type ChanId = Int
30type ChanId = Int
29−type Chans = MVar (ChanId, IntMap (Chan ByteString, MVar [Socket]))
31+type Chans = MVar (ChanId, IntMap (Chan ByteString, [(ThreadId, Socket)]))
30
32
31-- | This deals with several different configuration properties:
33-- | This deals with several different configuration properties:
32-- * Buffer size, specified in Hints
34-- * Buffer size, specified in Hints
@@ -42,7 +44,7 @@ data TCPConfig = TCPConfig Hints HostName ServiceName
42-- Messages are all queued using an unbounded Chan.
44-- Messages are all queued using an unbounded Chan.
43mkTransport :: TCPConfig -> IO Transport
45mkTransport :: TCPConfig -> IO Transport
44mkTransport (TCPConfig _hints host service) = withSocketsDo $ do
46mkTransport (TCPConfig _hints host service) = withSocketsDo $ do
45− channels <- newMVar (0, IntMap.empty)
47+ chans <- newMVar (0, IntMap.empty)
46 serverAddrs <- getAddrInfo
48 serverAddrs <- getAddrInfo
47 (Just (N.defaultHints { addrFlags = [AI_PASSIVE] }))
49 (Just (N.defaultHints { addrFlags = [AI_PASSIVE] }))
48 Nothing
50 Nothing
@@ -51,21 +53,31 @@ mkTransport (TCPConfig _hints host service) = withSocketsDo $ do
51 [] -> error "mkTransport: getAddrInfo returned []"
53 [] -> error "mkTransport: getAddrInfo returned []"
52 as -> head as
54 as -> head as
53 sock <- socket (addrFamily serverAddr) Stream defaultProtocol
55 sock <- socket (addrFamily serverAddr) Stream defaultProtocol
56+ setSocketOption sock ReuseAddr 1
54 bindSocket sock (addrAddress serverAddr)
57 bindSocket sock (addrAddress serverAddr)
55 listen sock 5
58 listen sock 5
56− _ <- forkIO $ procConnections channels sock
59+ threadId <- forkIO $ procConnections chans sock
57
60
58 return Transport
61 return Transport
59 { newConnectionWith = \_ -> do
62 { newConnectionWith = \_ -> do
60− (chanId, chanMap) <- takeMVar channels
63+ (chanId, chanMap) <- takeMVar chans
61 chan <- newChan
64 chan <- newChan
62− socks <- newMVar [sock]
63− putMVar channels (chanId + 1, IntMap.insert chanId (chan, socks) chanMap)
65+ putMVar chans (chanId + 1, IntMap.insert chanId (chan, []) chanMap)
64− return (mkSourceAddr host service chanId, mkTargetEnd chan socks)
66+ return (mkSourceAddr host service chanId, mkTargetEnd chans chanId chan)
65 , newMulticastWith = error "newMulticastWith: not defined"
67 , newMulticastWith = error "newMulticastWith: not defined"
66 , deserialize = \bs ->
68 , deserialize = \bs ->
67 let (host, service, chanId) = B.decode bs in
69 let (host, service, chanId) = B.decode bs in
68 Just $ mkSourceAddr host service chanId
70 Just $ mkSourceAddr host service chanId
71+ , closeTransport = do
72+ -- Kill the transport channel process
73+ killThread threadId
74+ sClose sock
75+ -- Kill all target end processes
76+ (chanId, chanMap) <- takeMVar chans
77+ forM_ (IntMap.elems chanMap) (\(chan, socks) ->
78+ forM_ socks (\(threadId', sock') -> do
79+ killThread threadId'
80+ sClose sock'))
69 }
81 }
70
82
71 where
83 where
@@ -93,13 +105,22 @@ mkTransport (TCPConfig _hints host service) = withSocketsDo $ do
93 , closeSourceEnd = sClose sock
105 , closeSourceEnd = sClose sock
94 }
106 }
95
107
96− mkTargetEnd :: Chan ByteString -> MVar [Socket] -> TargetEnd
108+ mkTargetEnd :: Chans -> ChanId -> Chan ByteString -> TargetEnd
97− mkTargetEnd chan socks = TargetEnd
109+ mkTargetEnd chans chanId chan = TargetEnd
98 { -- for now we will implement this as a Chan
110 { -- for now we will implement this as a Chan
99 receive = do
111 receive = do
100 bs <- readChan chan
112 bs <- readChan chan
101 return [bs]
113 return [bs]
102− , closeTargetEnd = takeMVar socks >>= mapM_ sClose
114+ , closeTargetEnd = do
115+ (chanId', chanMap) <- takeMVar chans
116+ case IntMap.lookup chanId chanMap of
117+ Nothing -> putMVar chans (chanId', chanMap)
118+ Just (_, socks) -> do
119+ forM_ socks $ \(threadId, sock) -> do
120+ killThread threadId
121+ sClose sock
122+ let chanMap' = IntMap.delete chanId chanMap
123+ putMVar chans (chanId', chanMap')
103 }
124 }
104
125
105 procConnections :: Chans -> Socket -> IO ()
126 procConnections :: Chans -> Socket -> IO ()
@@ -108,36 +129,50 @@ mkTransport (TCPConfig _hints host service) = withSocketsDo $ do
108 -- decode the first message to find the correct chanId
129 -- decode the first message to find the correct chanId
109 bs <- recvExact clientSock 8
130 bs <- recvExact clientSock 8
110 let chanId = fromIntegral (B.decode bs :: Int64)
131 let chanId = fromIntegral (B.decode bs :: Int64)
111− (_, chanMap) <- readMVar chans
132+ (chanId', chanMap) <- takeMVar chans
112 case IntMap.lookup chanId chanMap of
133 case IntMap.lookup chanId chanMap of
134+ Nothing -> do
135+ putMVar chans (chanId', chanMap)
113− Nothing -> error "procConnections: cannot find chanId"
136+ error "procConnections: cannot find chanId"
114 Just (chan, socks) -> do
137 Just (chan, socks) -> do
115− modifyMVar_ socks (return . (clientSock :))
116− forkIO $ procMessages chan clientSock
138+ threadId <- forkIO $ procMessages chans chanId chan clientSock
139+ let chanMap' = IntMap.insert chanId (chan, (threadId, clientSock):socks) chanMap
140+ putMVar chans (chanId', chanMap')
117
141
118 -- This function first extracts a header of type Int64, which determines
142 -- This function first extracts a header of type Int64, which determines
119 -- the size of the ByteString that follows. The ByteString is then
143 -- the size of the ByteString that follows. The ByteString is then
120 -- extracted from the socket, and then written to the Chan only when
144 -- extracted from the socket, and then written to the Chan only when
121− -- complete.
145+ -- complete. A length of -1 indicates that the socket has finished
146+ -- communicating.
122− procMessages :: Chan ByteString -> Socket -> IO ()
147+ procMessages :: Chans -> ChanId -> Chan ByteString -> Socket -> IO ()
123− procMessages chan sock = forever $ do
148+ procMessages chans chanId chan sock = do
124 sizeBS <- recvExact sock 8
149 sizeBS <- recvExact sock 8
125 let size = B.decode sizeBS
150 let size = B.decode sizeBS
151+ if size == -1
152+ then do
153+ (chanId', chanMap) <- takeMVar chans
154+ case IntMap.lookup chanId chanMap of
155+ Nothing -> do
156+ putMVar chans (chanId', chanMap)
157+ error "procMessages: chanId not found."
158+ Just (chan, socks) -> do
159+ let socks' = filter ((/= sock) . snd) socks
160+ let chanMap' = IntMap.insert chanId (chan, socks') chanMap
161+ putMVar chans (chanId', chanMap')
162+ sClose sock
163+ else do
126− bs <- recvExact sock size
164+ bs <- recvExact sock size
127− writeChan chan bs
165+ writeChan chan bs
166+ procMessages chans chanId chan sock
128
167
129−-- The result of `recvExact sock n` is a `ByteString` whose length is exactly
168+-- | The result of `recvExact sock n` is a `ByteString` of length `n`, received
130−-- `n`. No more bytes than necessary are read from the socket.
169+-- from `sock`. No more bytes than necessary are read from the socket.
131-- NB: This uses Network.Socket.ByteString.recv, which may *discard*
170-- NB: This uses Network.Socket.ByteString.recv, which may *discard*
132-- superfluous input depending on the socket type.
171-- superfluous input depending on the socket type.
133recvExact :: Socket -> Int64 -> IO ByteString
172recvExact :: Socket -> Int64 -> IO ByteString
134−recvExact sock n = do
173+recvExact sock n = go BS.empty sock n
174+ where
175+ go bs sock 0 = return bs
176+ go bs sock n = do
135− bs <- NBS.recv sock n
177+ bs' <- NBS.recv sock n
136− let remainder = n - BS.length bs
178+ go (BS.append bs bs') sock (n - BS.length bs')
137− if remainder > 0
138− then do
139− bs' <- recvExact sock remainder
140− return (BS.append bs bs')
141− else
142− return bs
143−