More performant pendingRoutes (IntMap -> Array).

1 files changed, 18 insertions, 14 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