More performant pendingRoutes (IntMap -> Array).

3 files changed, 23 insertions, 22 deletions

diff --git a/OnionRouter.hs b/OnionRouter.hs
index e3c32fbe..b11a0bf6 100644
--- a/OnionRouter.hs
+++ b/OnionRouter.hs
@@ -72,7 +72,7 @@ data OnionRouter = OnionRouter
       -- | Each of the 12 routes has a version number here that is set larger
       -- than the 'routeVersion' set in 'routeMap' when the route should be
       -- discarded and replaced with a fresh one.
-    , pendingRoutes   :: IntMap (TVar Int)
+    , pendingRoutes   :: TArray Int Int
       -- | Debug prints are written to this channel which is then flushed to
       -- 'routeLogger'.
     , routeLog        :: TChan String
@@ -134,7 +134,7 @@ gotTimeout rr = rr
     { timeoutCount = succ $ timeoutCount rr
     }
 
-data RouteEvent = BuildRoute RouteId
+newtype RouteEvent = BuildRoute RouteId
 
 newOnionRouter :: (String -> IO ()) -> IO OnionRouter
 newOnionRouter perror = do
@@ -148,10 +148,10 @@ newOnionRouter perror = do
         tn <- newTVar IntMap.empty
         ti <- newTVar HashMap.empty
         tc <- newTVar 0
-        vs <- sequence $ replicate 12 (newTVar 0)
+        pr <- newArray (0,11) 0
         rlog <- newTChan
         return OnionRouter
-            { pendingRoutes = IntMap.fromList $ zip [0..11] vs
+            { pendingRoutes = pr
             , onionDRG = drg
             , pendingQueries = pq
             , routeMap = rm
@@ -170,17 +170,19 @@ forkRouteBuilder or getnodes = do
         me <- myThreadId
         labelThread me "OnionRouter"
         forever $ do
-            let checkRebuild rid want_build stm = flip orElse stm $ do
-                    wanted_ver <- readTVar want_build
+            let checkRebuild :: Int -> Int -> STM RouteEvent
+                checkRebuild rid wanted_ver = do
                     current_ver <- fmap routeVersion <$> readArray (routeMap or) rid
                     writeTChan (routeLog or) $ "ONION checkRebuild "++show(current_ver,wanted_ver)
                     check $ maybe True (< wanted_ver) current_ver
                     return $ BuildRoute $ RouteId rid
-            io <- atomically $
+            io <- atomically $ {-# SCC "forkRouteBuilder.log" #-}
                    (readTChan (routeLog or) >>= return . routeLogger or)
-                `orElse`
-                   (IntMap.foldrWithKey checkRebuild retry (pendingRoutes or)
-                        >>= return . handleEvent getnodes or { routeThread = me })
+                `orElse` {-# SCC "forkRouteBuilder.checkRebuild" #-}
+                   (let stms = map (\rid -> checkRebuild rid =<< readArray (pendingRoutes or) rid)
+                                   [0..11]
+                        in do event <- foldr1 orElse stms
+                              return $ handleEvent getnodes or { routeThread = me } event)
             io
     return or { routeThread = tid }
 
@@ -240,7 +242,9 @@ selectTrampolines or = do
             routeLogger or $ unwords
                 ( "ONION Discarding insecure trampolines:" : (map (show . nodeAddr) ns) )
             myThreadId >>= flip labelThread ("OnionRouter.selectTrampolines.sleep")
-            threadDelay 1000000
+            case ns of
+                [_,_,_] -> threadDelay 1000000 -- wait 1 second if we failed the distinct3by predicate.
+                _       -> threadDelay 5000000 -- wait 5 seconds if insufficient nodes.
             myThreadId >>= flip labelThread ("OnionRouter.selectTrampolines")
             selectTrampolines or
         Right ns -> do
@@ -347,7 +351,7 @@ handleEvent getnodes or e@(BuildRoute (RouteId rid)) = do
                                              (freshRoute now r)
                                              rid
                                  v <- routeVersion . fromJust <$> readArray (routeMap or) rid
-                                 writeTVar (pendingRoutes or IntMap.! rid) v
+                                 writeArray (pendingRoutes or) rid v
                                  )
                        mb
     case mb of
@@ -390,7 +394,7 @@ hookQueries or t8 tmethods = TransactionMethods
     { dispatchRegister = \mvar od d -> {-# SCC "hookQ.dispatchRegister" #-} do --  :: MVar x -> d -> STM (tid, d)
         let ni = onionNodeInfo od
             rid@(RouteId ridn) = fromMaybe (routeId (nodeId ni)) $ onionRouteSpec od
-        wanted <- {-# SCC "hookQ.wanted" #-} (readTVar (pendingRoutes or IntMap.! ridn))
+        wanted <- {-# SCC "hookQ.wanted" #-} (readArray (pendingRoutes or) ridn)
         mr <- {-# SCC "hookQ.mr_action" #-} (readArray (routeMap or) ridn)
         -- Block query until a route is ready.
         check $ fromMaybe False $ do
@@ -431,7 +435,7 @@ hookQueries or t8 tmethods = TransactionMethods
             mrr <- readArray (routeMap or) rid
             forM_ mrr $ \rr -> do
             when (routeVersion rr == pendingVersion pq) $ do
-            let expireRoute = modifyTVar' (pendingRoutes or IntMap.! rid) expire
+            let expireRoute = modifyArray (pendingRoutes or) expire rid
                 expire ver | ver <= succ (pendingVersion pq) = succ (pendingVersion pq)
                            | otherwise                       = ver
             modifyArray (routeMap or) (fmap gotTimeout) rid
diff --git a/examples/dhtd.hs b/examples/dhtd.hs
index 8e9e7692..7ed77d3a 100644
--- a/examples/dhtd.hs
+++ b/examples/dhtd.hs
@@ -935,7 +935,7 @@ clientSession s@Session{..} sock cnum h = do
             ts <- readTVar $ trampolineNodes onionRouter
             tcnt <- readTVar $ trampolineCount onionRouter
             icnt <- HashMap.size <$> readTVar (trampolineIds onionRouter)
-            rs <- mapM readTVar (pendingRoutes onionRouter)
+            rs <- getAssocs (pendingRoutes onionRouter)
             pqs <- readTVar (pendingQueries onionRouter)
             let showRecord :: Int -> Int -> [String]
                 showRecord n wanted_ver
@@ -945,7 +945,7 @@ clientSession s@Session{..} sock cnum h = do
                                      then show routeVersion
                                      else show routeVersion ++ "(pending)" ]
                     | otherwise = [show n, "error!","","",""]
-                r = map (uncurry showRecord) $ IntMap.toAscList rs
+                r = map (uncurry showRecord) rs
             return $ do
                 hPutClientChunk h $ unlines [ "trampolines: " ++ show (IntMap.size ts,tcnt,icnt)
                                             , "pending: " ++ show (W64.size pqs) ]
diff --git a/src/Network/Tox/Onion/Transport.hs b/src/Network/Tox/Onion/Transport.hs
index e2e136ae..7ed9702a 100644
--- a/src/Network/Tox/Onion/Transport.hs
+++ b/src/Network/Tox/Onion/Transport.hs
@@ -250,12 +250,9 @@ encodeOnionAddr crypto getRoute (msg,OnionDestination _ ni (Just rid)) = do
             req <- wrapForRoute crypto msg ni route
             return ( runPut $ putRequest req
                    , nodeAddr $ routeNodeA route)
-        mapM' f x = do
-            let _ = x :: Maybe OnionRoute
-            -- dput XMisc $ "ONION encode sending to  " ++ show ni
-            -- dput XMisc $ "ONION encode getRoute -> " ++ show (fmap (\y -> map ($ y) [routeNodeA,routeNodeB,routeNodeC]) x)
-            mapM f x         -- ONION encode getRoute -> Nothing
-    getRoute ni rid >>= mapM' go
+    m <- {-# SCC "encodeOnionAddr.getRoute" #-} getRoute ni rid
+    x <- {-# SCC "encodeOnionAddr.wrapForRoute" #-} mapM go m
+    return x
 
 
 forwardOnions :: TransportCrypto -> UDPTransport -> UDPTransport