From d1d653317623d2656ca393c34961795ac0b12481 Mon Sep 17 00:00:00 2001
From: ken zhou <nanamikon@gmail.com>
Date: Fri, 15 Jan 2021 10:45:46 +0800
Subject: [PATCH] feat: support hash on vars_combinations (#3278)

related #3261
---
 apisix/admin/upstreams.lua |  4 ++
 apisix/balancer/chash.lua  |  6 +++
 apisix/schema_def.lua      |  6 +++
 doc/admin-api.md           |  1 +
 t/node/chash-hashon.t      | 91 ++++++++++++++++++++++++++++++++++++++
 5 files changed, 108 insertions(+)

diff --git a/apisix/admin/upstreams.lua b/apisix/admin/upstreams.lua
index bee4d734..312bf77f 100644
--- a/apisix/admin/upstreams.lua
+++ b/apisix/admin/upstreams.lua
@@ -45,6 +45,10 @@ local function get_chash_key_schema(hash_on)
         return nil, nil
     end
 
+    if hash_on == "vars_combinations" then
+        return core.schema.upstream_hash_vars_combinations_schema
+    end
+
     return nil, "invalid hash_on type " .. hash_on
 end
 
diff --git a/apisix/balancer/chash.lua b/apisix/balancer/chash.lua
index 7c955266..df1568a5 100644
--- a/apisix/balancer/chash.lua
+++ b/apisix/balancer/chash.lua
@@ -38,6 +38,12 @@ local function fetch_chash_hash_key(ctx, upstream)
         chash_key = ctx.var["http_" .. key]
     elseif hash_on == "cookie" then
         chash_key = ctx.var["cookie_" .. key]
+    elseif hash_on == "vars_combinations" then
+        local err
+        chash_key, err = core.utils.resolve_var(key, ctx.var);
+        if err then
+            core.log.error("could not resolve vars in ", key, " error: ", err)
+        end
     end
 
     if not chash_key then
diff --git a/apisix/schema_def.lua b/apisix/schema_def.lua
index 1d945662..7fee31b2 100644
--- a/apisix/schema_def.lua
+++ b/apisix/schema_def.lua
@@ -334,6 +334,7 @@ local upstream_schema = {
               "header",
               "cookie",
               "consumer",
+              "vars_combinations",
             },
         },
         key = {
@@ -393,6 +394,11 @@ _M.upstream_hash_header_schema = {
     pattern = [[^[a-zA-Z0-9-_]+$]]
 }
 
+-- validates string only
+_M.upstream_hash_vars_combinations_schema = {
+    type = "string"
+}
+
 
 _M.route = {
     type = "object",
diff --git a/doc/admin-api.md b/doc/admin-api.md
index 6c7deb5e..c542e952 100644
--- a/doc/admin-api.md
+++ b/doc/admin-api.md
@@ -534,6 +534,7 @@ In addition to the basic complex equalization algorithm selection, APISIX's Upst
 1. when it is `header`, the `key` is required. It is equal to "http_`key`".
 1. when it is `cookie`, the `key` is required. It is equal to "cookie_`key`".
 1. when it is `consumer`, the `key` is optional. The key is the `consumer_name` set by authentication plugin.
+1. when it is `vars_combinations`, the `key` is required. The `key` can be any [Nginx builtin variables](http://nginx.org/en/docs/varindex.html) combinations, such as `$request_uri$remote_addr`.
 1. If there is no value for either `hash_on` or `key`, `remote_addr` will be used as key.
 
 Config Example:
diff --git a/t/node/chash-hashon.t b/t/node/chash-hashon.t
index e83fe4cb..ae81bc44 100644
--- a/t/node/chash-hashon.t
+++ b/t/node/chash-hashon.t
@@ -625,3 +625,94 @@ chash_key: "chash_val_3"
 chash_key: "chash_val_4"
 chash_key: "chash_val_5"
 chash_key: "chash_val_6"
+
+
+
+=== TEST 13: set route(two upstream nodes, type chash), hash_on vars_combinations
+--- 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": "$http_custom_header-$http_custom_header_second",
+                        "type": "chash",
+                        "hash_on": "vars_combinations",
+                        "nodes": {
+                            "127.0.0.1:1980": 1,
+                            "127.0.0.1:1981": 1
+                        }
+                    }
+                }]]
+                )
+
+            if code >= 300 then
+                ngx.status = code
+            end
+            ngx.say(body)
+        }
+    }
+--- request
+GET /t
+--- response_body
+passed
+--- no_error_log
+[error]
+
+
+
+=== TEST 14: hit routes, hash_on custom header combinations
+--- 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"
+
+            local request_headers = {}
+            request_headers["custom_header"] = "custom-one"
+            request_headers["custom_header_second"] = "custom-two"
+
+            local ports_count = {}
+            for i = 1, 4 do
+                local httpc = http.new()
+                local res, err = httpc:request_uri(uri, {method = "GET", headers = request_headers})
+                if not res then
+                    ngx.say(err)
+                    return
+                end
+                ports_count[res.body] = (ports_count[res.body] or 0) + 1
+            end
+
+            local ports_arr = {}
+            for port, count in pairs(ports_count) do
+                table.insert(ports_arr, {port = port, count = count})
+            end
+
+            local function cmd(a, b)
+                return a.port > b.port
+            end
+            table.sort(ports_arr, cmd)
+
+            ngx.say(require("toolkit.json").encode(ports_arr))
+            ngx.exit(200)
+        }
+    }
+--- request
+GET /t
+--- response_body
+[{"count":4,"port":"1980"}]
+--- grep_error_log eval
+qr/hash_on: vars_combinations|chash_key: "custom-one-custom-two"/
+--- grep_error_log_out
+hash_on: vars_combinations
+chash_key: "custom-one-custom-two"
+hash_on: vars_combinations
+chash_key: "custom-one-custom-two"
+hash_on: vars_combinations
+chash_key: "custom-one-custom-two"
+hash_on: vars_combinations
+chash_key: "custom-one-custom-two"