@@ -4,12 +4,13 @@ module TestTransport where
44import Prelude hiding (catch, (>>=), (>>), return, fail)
55import TestAuxiliary (forkTry, runTests, trySome, randomThreadDelay)
66import Control.Concurrent (forkIO, killThread, yield)
7−import Control.Concurrent.MVar (newEmptyMVar, takeMVar, putMVar, readMVar, tryTakeMVar)
8−import Control.Exception (evaluate, throw, throwIO)
9−import Control.Monad (replicateM, replicateM_, when, guard, forM, forM_, unless)
7+import Control.Concurrent.MVar (newEmptyMVar, takeMVar, putMVar, readMVar, tryTakeMVar, modifyMVar_, newMVar)
8+import Control.Exception (evaluate, throw, throwIO, bracket)
9+import Control.Monad (replicateM, replicateM_, when, guard, forM_, unless)
1010import Control.Monad.Error ()
11+import Control.Applicative ((<$>))
1112import Network.Transport
12−import Network.Transport.Internal (tlog)
13+import Network.Transport.Internal (tlog, tryIO, timeoutMaybe)
1314import Network.Transport.Util (spawn)
1415import System.Random (randomIO)
1516import Data.ByteString (ByteString)
@@ -271,30 +272,37 @@ testCloseOneDirection transport numPings = do
271272
272273 mapM_ takeMVar [doneA, doneB]
273274
274−-- | Collect a given number of events and order them by connection ID
275−collect :: EndPoint -> Int -> IO [(ConnectionId, [[ByteString]])]
276−collect endPoint numEvents = go numEvents Map.empty Map.empty
275+-- | Collect events and order them by connection ID
276+collect :: EndPoint -> Maybe Int -> Maybe Int -> IO [(ConnectionId, [[ByteString]])]
277+collect endPoint maxEvents timeout = go maxEvents Map.empty Map.empty
277278 where
278279 -- TODO: for more serious use of this function we'd need to make these arguments strict
279− go 0 open closed = if Map.null open
280− then return . Map.toList . Map.map reverse $ closed
281− else fail "Open connections"
280+ go (Just 0) open closed = finish open closed
282281 go n open closed = do
283− event <- receive endPoint
284− case event of
285− ConnectionOpened cid _ _ ->
286− go (n - 1) (Map.insert cid [] open) closed
287− ConnectionClosed cid ->
288− let list = Map.findWithDefault (error "Invalid ConnectionClosed") cid open in
289− go (n - 1) (Map.delete cid open) (Map.insert cid list closed)
290− Received cid msg ->
291− go (n - 1) (Map.adjust (msg :) cid open) closed
292− ReceivedMulticast _ _ ->
293− fail "Unexpected multicast"
294− ErrorEvent _ ->
295− fail "Unexpected error"
296− EndPointClosed ->
297− fail "Unexpected endpoint closure"
282+ mEvent <- tryIO . timeoutMaybe timeout (userError "timeout") $ receive endPoint
283+ case mEvent of
284+ Left _ -> finish open closed
285+ Right event -> do
286+ let n' = (\x -> x - 1) <$> n
287+ case event of
288+ ConnectionOpened cid _ _ ->
289+ go n' (Map.insert cid [] open) closed
290+ ConnectionClosed cid ->
291+ let list = Map.findWithDefault (error "Invalid ConnectionClosed") cid open in
292+ go n' (Map.delete cid open) (Map.insert cid list closed)
293+ Received cid msg ->
294+ go n' (Map.adjust (msg :) cid open) closed
295+ ReceivedMulticast _ _ ->
296+ fail "Unexpected multicast"
297+ ErrorEvent _ ->
298+ fail "Unexpected error"
299+ EndPointClosed ->
300+ fail "Unexpected endpoint closure"
301+
302+ finish open closed =
303+ if Map.null open
304+ then return . Map.toList . Map.map reverse $ closed
305+ else fail "Open connections"
298306
299307-- | Open connection, close it, then reopen it
300308-- (In the TCP transport this means the socket will be closed, then reopened)
@@ -342,7 +350,7 @@ testCloseReopen transport numPings = do
342350 Right endpoint <- newEndPoint transport
343351 putMVar addrB (address endpoint)
344352
345− eventss <- collect endpoint (numRepeats * (numPings + 2))
353+ eventss <- collect endpoint (Just (numRepeats * (numPings + 2))) Nothing
346354
347355 forM_ (zip [1 .. numRepeats] eventss) $ \(i, (_, events)) -> do
348356 forM_ (zip [1 .. numPings] events) $ \(j, event) -> do
@@ -368,7 +376,7 @@ testParallelConnects transport numPings = do
368376 close conn
369377
370378 forkTry $ do
371− eventss <- collect endpoint (numPings * 4)
379+ eventss <- collect endpoint (Just (numPings * 4)) Nothing
372380 -- Check that no pings got sent to the wrong connection
373381 forM_ eventss $ \(_, [[ping1], [ping2]]) ->
374382 guard (ping1 == ping2)
@@ -519,7 +527,7 @@ testConnectToSelfTwice transport numPings = do
519527 forkTry $ do
520528 tlog $ "reading"
521529
522− [(_, events1), (_, events2)] <- collect endpoint (2 * (numPings + 2))
530+ [(_, events1), (_, events2)] <- collect endpoint (Just (2 * (numPings + 2))) Nothing
523531 True <- return $ events1 == replicate numPings ["pingA"]
524532 True <- return $ events2 == replicate numPings ["pingB"]
525533
@@ -826,31 +834,34 @@ testKill newTransport numThreads = do
826834 Right endpoint1 <- newEndPoint transport1
827835 Right endpoint2 <- newEndPoint transport2
828836
829− Right conn <- connect endpoint1 (address endpoint2) ReliableOrdered defaultConnectHints
830− ConnectionOpened _ _ _ <- receive endpoint2
831−
832837 threads <- replicateM numThreads . forkIO $ do
833− randomThreadDelay 10
834− Right () <- send conn ["ping"]
835− return ()
836−
837− -- Kill half of those threads, and wait on the rest
838− killerDone <- newEmptyMVar
839− forkIO $ do
840− killed <- forM threads $ \tid -> do
841− shouldKill <- randomIO
842− when shouldKill $ randomThreadDelay 10 >> killThread tid
843− return shouldKill
844−
845− -- We should receive at least the pings from the threads that we didn't
846− -- kill (we might get more, depending on when exactly the kill happens)
847− forM_ killed $ \wasKilled -> unless wasKilled $ do
848− Received _ ["ping"] <- receive endpoint2
849− return ()
850−
851− putMVar killerDone ()
838+ randomThreadDelay 100
839+ bracket (connect endpoint1 (address endpoint2) ReliableOrdered defaultConnectHints)
840+ (\(Right conn) -> randomThreadDelay 100 >> close conn)
841+ (\(Right conn) -> randomThreadDelay 100 >> do Right () <- send conn ["ping"] ; return ())
842+
843+ numAlive <- newMVar (0 :: Int)
844+
845+ -- Kill half of those threads
846+ forkIO . forM_ threads $ \tid -> do
847+ shouldKill <- randomIO
848+ if shouldKill
849+ then randomThreadDelay 600 >> killThread tid
850+ else modifyMVar_ numAlive (return . (+ 1))
851+
852+ -- Since it is impossible to predict when the kill exactly happens, we don't
853+ -- know how many connects were opened and how many pings were sent. But we
854+ -- should not have any open connections (if we do, collect will throw an
855+ -- error) and we should have at least the number of pings equal to the number
856+ -- of threads we did *not* kill
857+ eventss <- collect endpoint2 Nothing (Just 1000000)
858+ let actualPings = sum . map (length . snd) $ eventss
859+ expectedPings <- takeMVar numAlive
860+ unless (actualPings >= expectedPings) $
861+ throwIO (userError "Missing pings")
862+
863+-- print (actualPings, expectedPings)
852864
853− takeMVar killerDone
854865
855866-- | Set up conditions with a high likelyhood of "crossing" (for transports
856867-- that multiplex lightweight connections across heavyweight connections)