fix: reduce memory usage when abnormal weights are given in chash

Unlike the planned, I choose to do the gcd in APISIX because in this way we don't need to take care of dynamically resize. Signed-off-by: spacewander <spacewanderlzx@gmail.com>
apache · May 23, 2022 · 4b89dcf · 4b89dcf
1 parent adf7042
commit 4b89dcf
Show file tree

Hide file tree

Showing 5 changed files with 183 additions and 9 deletions.
diff --git a/apisix/balancer/chash.lua b/apisix/balancer/chash.lua
@@ -71,11 +71,27 @@ function _M.new(up_nodes, upstream)
 
     local nodes_count = 0
     local safe_limit = 0
+    local gcd = 0
     local servers, nodes = {}, {}
+
+    for serv, weight in pairs(up_nodes) do
+        if gcd == 0 then
+            gcd = weight
+        else
+            gcd = core.math.gcd(gcd, weight)
+        end
+    end
+
+    if gcd == 0 then
+        -- all nodes' weight are 0
+        gcd = 1
+    end
+
     for serv, weight in pairs(up_nodes) do
         local id = str_gsub(serv, ":", str_null)
 
         nodes_count = nodes_count + 1
+        weight = weight / gcd
         safe_limit = safe_limit + weight
         servers[id] = serv
         nodes[id] = weight

diff --git a/apisix/core.lua b/apisix/core.lua
@@ -53,4 +53,5 @@ return {
     resolver    = require("apisix.core.resolver"),
     os          = require("apisix.core.os"),
     pubsub      = require("apisix.core.pubsub"),
+    math        = require("apisix.core.math"),
 }
diff --git a/apisix/core/dns/client.lua b/apisix/core/dns/client.lua
@@ -24,6 +24,7 @@ local config_local = require("apisix.core.config_local")
 local log = require("apisix.core.log")
 local json = require("apisix.core.json")
 local table = require("apisix.core.table")
+local gcd = require("apisix.core.math").gcd
 local insert_tab = table.insert
 local math_random = math.random
 local package_loaded = package.loaded
@@ -38,15 +39,6 @@ local _M = {
 }
 
 
-local function gcd(a, b)
-    if b == 0 then
-        return a
-    end
-
-    return gcd(b, a % b)
-end
-
-
 local function resolve_srv(client, answers)
     if #answers == 0 then
         return nil, "empty SRV record"

diff --git a/apisix/core/math.lua b/apisix/core/math.lua
@@ -0,0 +1,41 @@
+--
+-- Licensed to the Apache Software Foundation (ASF) under one or more
+-- contributor license agreements.  See the NOTICE file distributed with
+-- this work for additional information regarding copyright ownership.
+-- The ASF licenses this file to You under the Apache License, Version 2.0
+-- (the "License"); you may not use this file except in compliance with
+-- the License.  You may obtain a copy of the License at
+--
+--     http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing, software
+-- distributed under the License is distributed on an "AS IS" BASIS,
+-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+-- See the License for the specific language governing permissions and
+-- limitations under the License.
+--
+
+--- Common library about math
+--
+-- @module core.math
+local _M = {}
+
+
+---
+-- Calculate the GCD of two numbers
+--
+-- @function core.math.gcd
+-- @tparam number a
+-- @tparam number b
+-- @treturn number the GCD of a and b
+local function gcd(a, b)
+    if b == 0 then
+        return a
+    end
+
+    return gcd(b, a % b)
+end
+_M.gcd = gcd
+
+
+return _M
diff --git a/t/node/chash-balance.t b/t/node/chash-balance.t
@@ -556,3 +556,127 @@ passed
 GET /t
 --- response_body
 200
+
+
+
+=== TEST 15: set routes with very big weights
+--- config
+    location /t {
+        content_by_lua_block {
+            local t = require("lib.test_admin").test
+            local code, body = t('/apisix/admin/routes/1',
+                 ngx.HTTP_PUT,
+                 [[{
+                    "uri": "/server_port",
+                    "upstream": {
+                        "key": "arg_device_id",
+                        "type": "chash",
+                        "nodes": {
+                            "127.0.0.1:1980": 1000000000,
+                            "127.0.0.1:1981": 2000000000,
+                            "127.0.0.1:1982": 1000000000
+                        }
+                    }
+                }]]
+                )
+
+            if code >= 300 then
+                ngx.status = code
+            end
+            ngx.say(body)
+        }
+    }
+--- request
+GET /t
+--- response_body
+passed
+--- no_error_log
+[error]
+
+
+
+=== TEST 16: hit
+--- config
+    location /t {
+        content_by_lua_block {
+            local http = require "resty.http"
+            local uri = "http://127.0.0.1:" .. ngx.var.server_port
+                        .. "/server_port?device_id=1"
+
+            local httpc = http.new()
+            local res, err = httpc:request_uri(uri, {method = "GET"})
+            if not res then
+                ngx.say(err)
+                return
+            end
+
+            -- a `size too large` error will be thrown if we don't reduce the weight
+            ngx.say(res.status)
+        }
+    }
+--- request
+GET /t
+--- response_body
+200
+
+
+
+=== TEST 17: set routes with very big weights, some nodes have zero weight
+--- config
+    location /t {
+        content_by_lua_block {
+            local t = require("lib.test_admin").test
+            local code, body = t('/apisix/admin/routes/1',
+                 ngx.HTTP_PUT,
+                 [[{
+                    "uri": "/server_port",
+                    "upstream": {
+                        "key": "arg_device_id",
+                        "type": "chash",
+                        "nodes": {
+                            "127.0.0.1:1980": 1000000000,
+                            "127.0.0.1:1981": 0,
+                            "127.0.0.1:1982": 4000000000
+                        }
+                    }
+                }]]
+                )
+
+            if code >= 300 then
+                ngx.status = code
+            end
+            ngx.say(body)
+        }
+    }
+--- request
+GET /t
+--- response_body
+passed
+--- no_error_log
+[error]
+
+
+
+=== TEST 18: hit
+--- config
+    location /t {
+        content_by_lua_block {
+            local http = require "resty.http"
+            local uri = "http://127.0.0.1:" .. ngx.var.server_port
+                        .. "/server_port?device_id=1"
+
+            local httpc = http.new()
+            local res, err = httpc:request_uri(uri, {method = "GET"})
+            if not res then
+                ngx.say(err)
+                return
+            end
+
+            -- a `size too large` error will be thrown if we don't reduce the weight
+            ngx.say(res.status)
+        }
+    }
+--- request
+GET /t
+--- response_body
+200