From c2d556f64d20941343c24879e7c4f2aafdbf8c74 Mon Sep 17 00:00:00 2001
From: ortur <orturb@yandex.ru>
Date: Mon, 29 Apr 2024 00:43:56 +0300
Subject: [PATCH] Feat: Implemented ip-hash

---
 configs/config.yaml      |  3 +-
 src/lb/tcp/selectors.cpp | 72 ++++++++++++++++++++++++++++++++++++++--
 src/lb/tcp/selectors.hpp | 20 +++++++++--
 3 files changed, 89 insertions(+), 6 deletions(-)

diff --git a/configs/config.yaml b/configs/config.yaml
index 1fb1b21..8f54400 100644
--- a/configs/config.yaml
+++ b/configs/config.yaml
@@ -13,8 +13,7 @@ acceptor:
   ip_version: 4
 
 load_balancing:
-  algorithm: consistent_hash
-  replicas: 5
+  algorithm: ip_hash
   endpoints:
     - ip: "127.0.0.1"
       port: 8001
diff --git a/src/lb/tcp/selectors.cpp b/src/lb/tcp/selectors.cpp
index 77258cc..fd892c4 100644
--- a/src/lb/tcp/selectors.cpp
+++ b/src/lb/tcp/selectors.cpp
@@ -86,7 +86,8 @@ SelectorPtr DetectSelector(const YAML::Node& node)
     static const std::unordered_map<std::string, SelectorType> selector_switch = {
         {"round_robin", SelectorType::ROUND_ROBIN},
         {"weighted_round_robin", SelectorType::WEIGHTED_ROUND_ROBIN},
-        {"consistent_hash", SelectorType::CONSISTENT_HASH}
+        {"consistent_hash", SelectorType::CONSISTENT_HASH},
+        {"ip_hash", SelectorType::IP_HASH}
     };
 
     if (!balancing_node["algorithm"].IsDefined()) {
@@ -112,6 +113,11 @@ SelectorPtr DetectSelector(const YAML::Node& node)
             return wrr;
         }
         break;
+        case SelectorType::IP_HASH: {
+            SelectorPtr ip = std::make_shared<IpHashSelector>();
+            ip->Configure(balancing_node);
+            return ip;
+        } break;
         case SelectorType::CONSISTENT_HASH: {
             if (!balancing_node["replicas"].as<std::size_t>()) {
                 EXCEPTION("Missed replicas field!");
@@ -120,7 +126,8 @@ SelectorPtr DetectSelector(const YAML::Node& node)
             SelectorPtr ch = std::make_shared<ConsistentHashSelector>(replicas_num);
             ch->Configure(balancing_node);
             return ch;
-        }
+        } break;
+
         default: {
             STACKTRACE("Selector {} is not implemented", name);
         }
@@ -290,6 +297,67 @@ void WeightedRoundRobinSelector::AdvanceCounter()
     }
 }
 
+// ============================ IpHashSelector ============================
+
+void IpHashSelector::Configure(const YAML::Node& balancing_node)
+{
+    if (!balancing_node["endpoints"].IsDefined()) {
+        STACKTRACE("Ip-hash endpoints node is missed");
+    }
+    const YAML::Node& ep_node = balancing_node["endpoints"];
+    if (!ep_node.IsSequence()) {
+        EXCEPTION("endpoints node must be a sequence");
+    }
+    backends_.reserve(ep_node.size());
+    for (const YAML::Node& ep : ep_node) {
+        if (ep["url"].IsDefined()) {
+            backends_.emplace_back(ep["url"].as<std::string>());
+            continue;
+        }
+
+        if (!ep["ip"].IsDefined()) {
+            STACKTRACE("{} missed {} field", ep, "ip");
+        }
+        if (!ep["port"].IsDefined()) {
+            STACKTRACE("{} missed {} field", ep, "port");
+        }
+
+        backends_.emplace_back(ep["ip"].as<std::string>(), ep["port"].as<int>());
+    }
+
+    for (const auto& backend : backends_) {
+        DEBUG("\t{}", backend);
+    }
+}
+
+Backend IpHashSelector::SelectBackend(const boost::asio::ip::tcp::endpoint& client_address)
+{
+    boost::mutex::scoped_lock lock(mutex_);
+    static auto ComputeHash = [&](const boost::asio::ip::tcp::endpoint& ep) -> std::size_t {
+        return std::hash<std::string>(ep.address().to_string().size()) * 37 +
+               ep.port() * 37 * 37;
+    };
+
+    Backend result = backends_[ComputeHash(client_address) % backends_.size()];
+    // DEBUG("Ip hash selected: {}", result);
+    // DEBUG("Hash: {}", ComputeHash(client_address));
+    return result;
+}
+
+void IpHashSelector::ExcludeBackend(const Backend& backend)
+{
+    boost::mutex::scoped_lock lock(mutex_);
+    backends_.erase(std::remove(backends_.begin(), backends_.end(), backend), backends_.end());
+    if (backends_.empty()) {
+        EXCEPTION("All backends are excluded!");
+    }
+}
+
+SelectorType IpHashSelector::Type() const
+{
+    return SelectorType::IP_HASH;
+}
+
 // ============================ ConsistentHashSelector ============================
 
 BackendCHTraits::HashType BackendCHTraits::GetHash(const Backend& backend)
diff --git a/src/lb/tcp/selectors.hpp b/src/lb/tcp/selectors.hpp
index 97936fa..f7c6aee 100644
--- a/src/lb/tcp/selectors.hpp
+++ b/src/lb/tcp/selectors.hpp
@@ -43,7 +43,8 @@ std::ostream& operator<<(std::ostream& out, const Backend& backend);
 enum class SelectorType {
     ROUND_ROBIN=0,
     WEIGHTED_ROUND_ROBIN,
-    CONSISTENT_HASH
+    IP_HASH,
+    CONSISTENT_HASH,
 };
 
 struct ISelector {
@@ -101,6 +102,20 @@ class WeightedRoundRobinSelector final : public ISelector {
     std::size_t counter_ = 0;
 };
 
+class IpHashSelector final : public ISelector {
+public:
+    void Configure(const YAML::Node& config) override;
+
+    Backend SelectBackend(const boost::asio::ip::tcp::endpoint& client_socket) override;
+
+    void ExcludeBackend(const Backend& backend) override;
+
+    SelectorType Type() const override;
+private:
+    boost::mutex mutex_;
+    std::vector<Backend> backends_;
+};
+
 
 struct BackendCHTraits {
     using HashType = std::size_t;
@@ -119,7 +134,7 @@ class ConsistentHashSelector final : public ISelector {
 
     void Configure(const YAML::Node& config) override;
 
-    Backend SelectBackend(const boost::asio::ip::tcp::endpoint& client_socket) override;
+    Backend SelectBackend(const boost::asio::ip::tcp::endpoint& client_address) override;
 
     void ExcludeBackend(const Backend& backend) override;
 
@@ -130,4 +145,5 @@ class ConsistentHashSelector final : public ISelector {
     pumba::ConsistentHashingRouter<Backend, BackendCHTraits> ring_; // guarded by mutex
 };
 
+
 } // namespace lb::tcp
\ No newline at end of file