From 01b24c7090af9a8901ec3968a38b4666d53010ef Mon Sep 17 00:00:00 2001
From: Jeff Becker <jeff@i2p.rocks>
Date: Tue, 3 Dec 2019 12:03:19 -0500
Subject: [PATCH] limit connections

---
 llarp/link/i_link_manager.hpp |  3 +++
 llarp/link/link_manager.cpp   | 17 +++++++++++++++++
 llarp/link/link_manager.hpp   |  3 +++
 llarp/link/server.hpp         |  8 ++++++++
 llarp/path/path_context.hpp   |  7 +++++--
 llarp/router/router.cpp       | 18 +++++++++++++++---
 6 files changed, 51 insertions(+), 5 deletions(-)

diff --git a/llarp/link/i_link_manager.hpp b/llarp/link/i_link_manager.hpp
index d42c9e98d..88f315e05 100644
--- a/llarp/link/i_link_manager.hpp
+++ b/llarp/link/i_link_manager.hpp
@@ -67,6 +67,9 @@ namespace llarp
     virtual size_t
     NumberOfConnectedClients() const = 0;
 
+    virtual size_t
+    NumberOfPendingConnections() const = 0;
+
     virtual bool
     GetRandomConnectedRouter(RouterContact &router) const = 0;
 
diff --git a/llarp/link/link_manager.cpp b/llarp/link/link_manager.cpp
index 09aee0d4d..6d0062df5 100644
--- a/llarp/link/link_manager.cpp
+++ b/llarp/link/link_manager.cpp
@@ -241,6 +241,23 @@ namespace llarp
     return connectedClients.size();
   }
 
+  size_t
+  LinkManager::NumberOfPendingConnections() const
+  {
+    size_t pending = 0;
+    for(const auto &link : inboundLinks)
+    {
+      pending += link->NumberOfPendingSessions();
+    }
+
+    for(const auto &link : outboundLinks)
+    {
+      pending += link->NumberOfPendingSessions();
+    }
+
+    return pending;
+  }
+
   bool
   LinkManager::GetRandomConnectedRouter(RouterContact &router) const
   {
diff --git a/llarp/link/link_manager.hpp b/llarp/link/link_manager.hpp
index 103ef5869..0d3f3666d 100644
--- a/llarp/link/link_manager.hpp
+++ b/llarp/link/link_manager.hpp
@@ -66,6 +66,9 @@ namespace llarp
     size_t
     NumberOfConnectedClients() const override;
 
+    size_t
+    NumberOfPendingConnections() const override;
+
     bool
     GetRandomConnectedRouter(RouterContact &router) const override;
 
diff --git a/llarp/link/server.hpp b/llarp/link/server.hpp
index 86b15ef1b..e024ee90a 100644
--- a/llarp/link/server.hpp
+++ b/llarp/link/server.hpp
@@ -205,6 +205,14 @@ namespace llarp
     // void
     // RemovePending(ILinkSession* s) LOCKS_EXCLUDED(m_PendingMutex);
 
+    /// count the number of sessions that are yet to be fully connected
+    size_t
+    NumberOfPendingSessions() const
+    {
+      ACQUIRE_LOCK(Lock_t lock, m_PendingMutex);
+      return m_Pending.size();
+    }
+
    private:
     static void
     on_timer_tick(void* user, uint64_t orig, uint64_t left)
diff --git a/llarp/path/path_context.hpp b/llarp/path/path_context.hpp
index b06f5728e..76a28bc35 100644
--- a/llarp/path/path_context.hpp
+++ b/llarp/path/path_context.hpp
@@ -12,6 +12,7 @@
 #include <util/types.hpp>
 
 #include <memory>
+#include <unordered_map>
 
 namespace llarp
 {
@@ -104,7 +105,8 @@ namespace llarp
       void
       RemovePathSet(PathSet_ptr set);
 
-      using TransitHopsMap_t = std::multimap< PathID_t, TransitHop_ptr >;
+      using TransitHopsMap_t =
+          std::unordered_multimap< PathID_t, TransitHop_ptr, PathID_t::Hash >;
 
       struct SyncTransitMap_t
       {
@@ -125,7 +127,8 @@ namespace llarp
       };
 
       // maps path id -> pathset owner of path
-      using OwnedPathsMap_t = std::map< PathID_t, Path_ptr >;
+      using OwnedPathsMap_t =
+          std::unordered_map< PathID_t, Path_ptr, PathID_t::Hash >;
 
       struct SyncOwnedPathsMap_t
       {
diff --git a/llarp/router/router.cpp b/llarp/router/router.cpp
index d2441e95a..074dfcf4d 100644
--- a/llarp/router/router.cpp
+++ b/llarp/router/router.cpp
@@ -831,8 +831,12 @@ namespace llarp
 
     _linkManager.CheckPersistingSessions(now);
 
-    const size_t connected = NumberOfConnectedRouters();
-    const size_t N         = nodedb()->num_loaded();
+    size_t connected = NumberOfConnectedRouters();
+    if(not isSvcNode)
+    {
+      connected += _linkManager.NumberOfPendingConnections();
+    }
+    const size_t N = nodedb()->num_loaded();
     if(N < llarp::path::default_len)
     {
       LogInfo("We need at least ", llarp::path::default_len,
@@ -1231,8 +1235,16 @@ namespace llarp
   }
 
   void
-  Router::ConnectToRandomRouters(int want)
+  Router::ConnectToRandomRouters(int _want)
   {
+    const size_t want = _want;
+    auto connected    = NumberOfConnectedRouters();
+    if(not IsServiceNode())
+    {
+      connected += _linkManager.NumberOfPendingConnections();
+    }
+    if(connected >= want)
+      return;
     _outboundSessionMaker.ConnectToRandomRouters(want, Now());
   }