From c4cbbd673180779a8a9564ef3595a9b70754f0a2 Mon Sep 17 00:00:00 2001
From: Stephen Shelton <steve@brewcraft.org>
Date: Mon, 15 Jun 2020 12:53:17 -0600
Subject: [PATCH] RouterHive: store router contexts by routerId instead of
 index

---
 llarp/tooling/router_hive.cpp | 70 ++++++++++++++---------------------
 llarp/tooling/router_hive.hpp | 23 ++++--------
 2 files changed, 34 insertions(+), 59 deletions(-)

diff --git a/llarp/tooling/router_hive.cpp b/llarp/tooling/router_hive.cpp
index 613db5a9a..769f0a77b 100644
--- a/llarp/tooling/router_hive.cpp
+++ b/llarp/tooling/router_hive.cpp
@@ -14,15 +14,19 @@ using namespace std::chrono_literals;
 namespace tooling
 {
   void
-  RouterHive::AddRouter(
-      const std::shared_ptr<llarp::Config>& config, std::vector<llarp_main*>* routers, bool isRelay)
+  RouterHive::AddRouter(const std::shared_ptr<llarp::Config>& config, bool isRelay)
   {
+    auto& container = (isRelay ? relays : clients);
+
     llarp_main* ctx = llarp_main_init_from_config(config->Copy(), isRelay);
     auto result = llarp_main_setup(ctx, isRelay);
     if (result == 0)
     {
-      llarp::Context::Get(ctx)->InjectHive(this);
-      routers->push_back(ctx);
+      auto context = llarp::Context::Get(ctx);
+      auto routerId = llarp::RouterID(context->router->pubkey());
+      context->InjectHive(this);
+      container[routerId] = ctx;
+      std::cout << "Generated router with ID " << routerId << std::endl;
     }
     else
     {
@@ -37,19 +41,21 @@ namespace tooling
   void
   RouterHive::AddRelay(const std::shared_ptr<llarp::Config>& config)
   {
-    AddRouter(config, &relays, true);
+    AddRouter(config, true);
   }
 
   void
   RouterHive::AddClient(const std::shared_ptr<llarp::Config>& config)
   {
-    AddRouter(config, &clients, false);
+    AddRouter(config, false);
   }
 
   void
-  RouterHive::StartRouters(std::vector<llarp_main*>* routers, bool isRelay)
+  RouterHive::StartRouters(bool isRelay)
   {
-    for (llarp_main* ctx : *routers)
+    auto& container = (isRelay ? relays : clients);
+
+    for (auto [routerId, ctx] : container)
     {
       routerMainThreads.emplace_back([=]() {
         llarp_main_run(ctx, llarp_main_runtime_opts{false, false, false, isRelay});
@@ -61,37 +67,37 @@ namespace tooling
   void
   RouterHive::StartRelays()
   {
-    StartRouters(&relays, true);
+    StartRouters(true);
   }
 
   void
   RouterHive::StartClients()
   {
-    StartRouters(&clients, false);
+    StartRouters(false);
   }
 
   void
   RouterHive::StopRouters()
   {
     llarp::LogInfo("Signalling all routers to stop");
-    for (llarp_main* ctx : relays)
+    for (auto [routerId, ctx] : relays)
     {
       llarp_main_signal(ctx, 2 /* SIGINT */);
     }
-    for (llarp_main* ctx : clients)
+    for (auto [routerId, ctx] : clients)
     {
       llarp_main_signal(ctx, 2 /* SIGINT */);
     }
 
     llarp::LogInfo("Waiting on routers to be stopped");
-    for (llarp_main* ctx : relays)
+    for (auto [routerId, ctx] : relays)
     {
       while (llarp_main_is_running(ctx))
       {
         std::this_thread::sleep_for(10ms);
       }
     }
-    for (llarp_main* ctx : clients)
+    for (auto [routerId, ctx] : clients)
     {
       while (llarp_main_is_running(ctx))
       {
@@ -154,28 +160,6 @@ namespace tooling
     LogicCall(ctx->logic, [visit, ctx]() { visit(ctx); });
   }
 
-  void
-  RouterHive::VisitRelay(size_t index, std::function<void(Context_ptr)> visit)
-  {
-    if (index >= relays.size())
-    {
-      visit(nullptr);
-      return;
-    }
-    VisitRouter(relays[index], visit);
-  }
-
-  void
-  RouterHive::VisitClient(size_t index, std::function<void(Context_ptr)> visit)
-  {
-    if (index >= clients.size())
-    {
-      visit(nullptr);
-      return;
-    }
-    VisitRouter(clients[index], visit);
-  }
-
   std::vector<size_t>
   RouterHive::RelayConnectedRelays()
   {
@@ -185,11 +169,11 @@ namespace tooling
 
     size_t i = 0;
     size_t done_count = 0;
-    for (auto relay : relays)
+    for (auto [routerId, ctx] : relays)
     {
-      auto ctx = llarp::Context::Get(relay);
-      LogicCall(ctx->logic, [&, i, ctx]() {
-        size_t count = ctx->router->NumberOfConnectedRouters();
+      auto context = llarp::Context::Get(ctx);
+      LogicCall(context->logic, [&, i, context]() {
+        size_t count = context->router->NumberOfConnectedRouters();
         std::lock_guard<std::mutex> guard{results_lock};
         results[i] = count;
         done_count++;
@@ -220,10 +204,10 @@ namespace tooling
     results.resize(relays.size());
 
     size_t i = 0;
-    for (auto relay : relays)
+    for (auto [routerId, ctx] : relays)
     {
-      auto ctx = llarp::Context::Get(relay);
-      results[i] = ctx->router->rc();
+      auto context = llarp::Context::Get(ctx);
+      results[i] = context->router->rc();
       i++;
     }
     return results;
diff --git a/llarp/tooling/router_hive.hpp b/llarp/tooling/router_hive.hpp
index 9e9982e35..7c547eeea 100644
--- a/llarp/tooling/router_hive.hpp
+++ b/llarp/tooling/router_hive.hpp
@@ -26,26 +26,17 @@ namespace tooling
 
    private:
     void
-    StartRouters(std::vector<llarp_main*>* routers, bool isRelay);
+    StartRouters(bool isRelay);
 
     void
     AddRouter(
         const std::shared_ptr<llarp::Config>& config,
-        std::vector<llarp_main*>* routers,
         bool isRelay);
 
     /// safely visit router
     void
     VisitRouter(llarp_main* router, std::function<void(Context_ptr)> visit);
 
-    /// safely visit relay at index N
-    void
-    VisitRelay(size_t index, std::function<void(Context_ptr)> visit);
-
-    /// safely visit client at index N
-    void
-    VisitClient(size_t index, std::function<void(Context_ptr)> visit);
-
    public:
     RouterHive() = default;
 
@@ -76,18 +67,18 @@ namespace tooling
     void
     ForEachRelay(std::function<void(Context_ptr)> visit)
     {
-      for (size_t idx = 0; idx < relays.size(); ++idx)
+      for (auto [routerId, ctx] : relays)
       {
-        VisitRelay(idx, visit);
+        VisitRouter(ctx, visit);
       }
     }
 
     void
     ForEachClient(std::function<void(Context_ptr)> visit)
     {
-      for (size_t idx = 0; idx < clients.size(); ++idx)
+      for (auto [routerId, ctx] : clients)
       {
-        VisitClient(idx, visit);
+        VisitRouter(ctx, visit);
       }
     }
 
@@ -105,8 +96,8 @@ namespace tooling
     std::vector<llarp::RouterContact>
     GetRelayRCs();
 
-    std::vector<llarp_main*> relays;
-    std::vector<llarp_main*> clients;
+    std::unordered_map<llarp::RouterID, llarp_main*, llarp::RouterID::Hash> relays;
+    std::unordered_map<llarp::RouterID, llarp_main*, llarp::RouterID::Hash> clients;
 
     std::vector<std::thread> routerMainThreads;