[fix] Fixed a bug in LoadBalancer where a target can be missed out wh…

…en doing a provision()
flomesh-io · Aug 6, 2024 · 7f72a5a · 7f72a5a
1 parent b55dc5b
commit 7f72a5a
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 3 deletions.
diff --git a/src/api/algo.cpp b/src/api/algo.cpp
@@ -665,6 +665,12 @@ LoadBalancer::Options::Options(pjs::Object *options) {
     .check_nullable();
 }
 
+LoadBalancer::~LoadBalancer() {
+  for (auto &p : m_pools) {
+    p->lb = nullptr;
+  }
+}
+
 void LoadBalancer::provision(pjs::Context &ctx, pjs::Array *targets) {
   if (targets) {
     std::map<pjs::Value, Pool*> new_targets;
@@ -683,8 +689,9 @@ void LoadBalancer::provision(pjs::Context &ctx, pjs::Array *targets) {
         auto it = m_targets.find(key);
         if (it != m_targets.end()) {
           p = it->second;
+          m_queue.remove(p);
         } else {
-          p = new Pool(key, target);
+          p = new Pool(this, key, target);
         }
 
         new_targets[key] = p;
@@ -847,6 +854,11 @@ LoadBalancer::Resource::~Resource() {
 }
 
 void LoadBalancer::Resource::free() {
+  if (auto lb = m_pool->lb) {
+    if (lb->m_options.algorithm == Algorithm::LEAST_LOAD) {
+      lb->decrease_load(m_pool);
+    }
+  }
   if (m_load > 0) {
     m_load--;
     if (auto r = back()) {

diff --git a/src/api/algo.hpp b/src/api/algo.hpp
@@ -301,9 +301,10 @@ class LoadBalancer : public pjs::ObjectTemplate<LoadBalancer> {
 
   class Pool : public pjs::RefCount<Pool>, public List<Pool>::Item {
   public:
-    Pool(const pjs::Value &k, const pjs::Value &t)
-      : key(k), target(t) {}
+    Pool(LoadBalancer *l, const pjs::Value &k, const pjs::Value &t)
+      : lb(l), key(k), target(t) {}
 
+    LoadBalancer* lb;
     pjs::Value key;
     pjs::Value target;
     int capacity = 0;
@@ -380,6 +381,8 @@ class LoadBalancer : public pjs::ObjectTemplate<LoadBalancer> {
   LoadBalancer(const Options &options)
     : m_options(options) {}
 
+  ~LoadBalancer();
+
   Options m_options;
   std::map<pjs::Value, Pool*> m_targets;
   std::vector<pjs::Ref<Pool>> m_pools;