Fix bug for ZeroDivisionError

it can happen when it has one node whose weight is 0. https://github.com/fluent/fluentd/blob/2502cb8b57ece25e7bf83aa29d4ded48acab63b3/lib/fluent/plugin/out_forward/load_balancer.rb#L93 Signed-off-by: Yuta Iwama <[email protected]>
fluent · May 12, 2020 · b8f3ee7 · b8f3ee7
1 parent 55d11dd
commit b8f3ee7
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 1 deletion.
diff --git a/lib/fluent/plugin/out_forward/load_balancer.rb b/lib/fluent/plugin/out_forward/load_balancer.rb
@@ -56,7 +56,7 @@ def select_healthy_node
       end
 
       def rebuild_weight_array(nodes)
-        standby_nodes, regular_nodes = nodes.partition {|n|
+        standby_nodes, regular_nodes = nodes.select { |e| e.weight > 0 }.partition {|n|
           n.standby?
         }
 

diff --git a/test/plugin/out_forward/test_load_balancer.rb b/test/plugin/out_forward/test_load_balancer.rb
@@ -49,12 +49,58 @@ class LoadBalancerTest < Test::Unit::TestCase
       end
     end
 
+    test 'call like round robin without weight=0 node' do
+      lb = Fluent::Plugin::ForwardOutput::LoadBalancer.new($log)
+      n1 = flexmock('node', :'standby?' => false, :'available?' => true, weight: 1)
+      n2 = flexmock('node', :'standby?' => false, :'available?' => true, weight: 1)
+      n3 = flexmock('node', :'standby?' => false, :'available?' => true, weight: 0)
+
+      lb.rebuild_weight_array([n1, n2, n3])
+
+      lb.select_healthy_node do |node|
+        # to handle random choice
+        if node == n1
+          lb.select_healthy_node do |node|
+            assert_equal(node, n2)
+          end
+
+          lb.select_healthy_node do |node|
+            assert_equal(node, n1)
+          end
+
+          lb.select_healthy_node do |node|
+            assert_equal(node, n2)
+          end
+        else
+          lb.select_healthy_node do |node|
+            assert_equal(node, n1)
+          end
+
+          lb.select_healthy_node do |node|
+            assert_equal(node, n2)
+          end
+
+          lb.select_healthy_node do |node|
+            assert_equal(node, n1)
+          end
+        end
+      end
+    end
+
     test 'raise an error if all node are unavialble' do
       lb = Fluent::Plugin::ForwardOutput::LoadBalancer.new($log)
       lb.rebuild_weight_array([flexmock('node', :'standby?' => false, :'available?' => false, weight: 1)])
       assert_raise(Fluent::Plugin::ForwardOutput::NoNodesAvailable) do
         lb.select_healthy_node
       end
     end
+
+    test 'it regards weight=0 node as unavialble' do
+      lb = Fluent::Plugin::ForwardOutput::LoadBalancer.new($log)
+      lb.rebuild_weight_array([flexmock('node', :'standby?' => false, :'available?' => true, weight: 0)])
+      assert_raise(Fluent::Plugin::ForwardOutput::NoNodesAvailable) do
+        lb.select_healthy_node
+      end
+    end
   end
 end