feat: Request-ID plugin add snowflake algorithm (#4559)

Author: dickens7

apisix/plugins/request-id.lua

@@ -14,25 +14,58 @@
 -- See the License for the specific language governing permissions and
 -- limitations under the License.
 --
-local core          = require("apisix.core")
-local plugin_name   = "request-id"
-local ngx           = ngx
-local uuid          = require("resty.jit-uuid")
+
+local ngx = ngx
+local bit = require("bit")
+local core = require("apisix.core")
+local snowflake = require("snowflake")
+local uuid = require("resty.jit-uuid")
+local process = require("ngx.process")
+local timers = require("apisix.timers")
+local tostring = tostring
+local math_pow = math.pow
+local math_ceil = math.ceil
+local math_floor = math.floor
+
+local plugin_name = "request-id"
+
+local data_machine = nil
+local snowflake_inited = nil
+
+local attr = nil
 
 local schema = {
     type = "object",
     properties = {
         header_name = {type = "string", default = "X-Request-Id"},
-        include_in_response = {type = "boolean", default = true}
+        include_in_response = {type = "boolean", default = true},
+        algorithm = {type = "string", enum = {"uuid", "snowflake"}, default = "uuid"}
     }
 }
 
+local attr_schema = {
+    type = "object",
+    properties = {
+        snowflake = {
+            type = "object",
+            properties = {
+                enable = {type = "boolean", default = false},
+                snowflake_epoc = {type = "integer", minimum = 1, default = 1609459200000},
+                data_machine_bits = {type = "integer", minimum = 1, maximum = 31, default = 12},
+                sequence_bits = {type = "integer", minimum = 1, default = 10},
+                delta_offset = {type = "integer", default = 1, enum = {1, 10, 100, 1000}},
+                data_machine_ttl = {type = "integer", minimum = 1, default = 30},
+                data_machine_interval = {type = "integer", minimum = 1, default = 10}
+            }
+        }
+    }
+}
 
 local _M = {
     version = 0.1,
     priority = 11010,
     name = plugin_name,
-    schema = schema,
+    schema = schema
 }
 
 
@@ -41,9 +74,144 @@ function _M.check_schema(conf)
 end
 
 
+-- Generates the current process data machine
+local function gen_data_machine(max_number)
+    if data_machine == nil then
+        local etcd_cli, prefix = core.etcd.new()
+        local prefix = prefix .. "/plugins/request-id/snowflake/"
+        local uuid = uuid.generate_v4()
+        local id = 1
+        ::continue::
+        while (id <= max_number) do
+            local res, err = etcd_cli:grant(attr.snowflake.data_machine_ttl)
+            if err then
+                id = id + 1
+                core.log.error("Etcd grant failure, err: ".. err)
+                goto continue
+            end
+
+            local _, err1 = etcd_cli:setnx(prefix .. tostring(id), uuid)
+            local res2, err2 = etcd_cli:get(prefix .. tostring(id))
+
+            if err1 or err2 or res2.body.kvs[1].value ~= uuid then
+                core.log.notice("data_machine " .. id .. " is not available")
+                id = id + 1
+            else
+                data_machine = id
+
+                local _, err3 =
+                    etcd_cli:set(
+                    prefix .. tostring(id),
+                    uuid,
+                    {
+                        prev_kv = true,
+                        lease = res.body.ID
+                    }
+                )
+
+                if err3 then
+                    id = id + 1
+                    etcd_cli:delete(prefix .. tostring(id))
+                    core.log.error("set data_machine " .. id .. " lease error: " .. err3)
+                    goto continue
+                end
+
+                local lease_id = res.body.ID
+                local start_at = ngx.time()
+                local handler = function()
+                    local now = ngx.time()
+                    if now - start_at < attr.snowflake.data_machine_interval then
+                        return
+                    end
+
+                    local _, err4 = etcd_cli:keepalive(lease_id)
+                    if err4 then
+                        snowflake_inited = nil
+                        data_machine = nil
+                        core.log.error("snowflake data_machine: " .. id .." lease faild.")
+                    end
+                    start_at = now
+                    core.log.info("snowflake data_machine: " .. id .." lease success.")
+                end
+
+                timers.register_timer("plugin#request-id", handler)
+                core.log.info(
+                    "timer created to lease snowflake algorithm data_machine, interval: ",
+                    attr.snowflake.data_machine_interval)
+                core.log.notice("lease snowflake data_machine: " .. id)
+                break
+            end
+        end
+
+        if data_machine == nil then
+            core.log.error("No data_machine is not available")
+            return nil
+        end
+    end
+    return data_machine
+end
+
+
+-- Split 'Data Machine' into 'Worker ID' and 'datacenter ID'
+local function split_data_machine(data_machine, node_id_bits, datacenter_id_bits)
+    local num = bit.tobit(data_machine)
+    local worker_id = bit.band(num, math_pow(2, node_id_bits) - 1)
+    num = bit.rshift(num, node_id_bits)
+    local datacenter_id = bit.band(num, math_pow(2, datacenter_id_bits) - 1)
+    return worker_id, datacenter_id
+end
+
+
+-- Initialize the snowflake algorithm
+local function snowflake_init()
+    if snowflake_inited == nil then
+        local max_number = math_pow(2, (attr.snowflake.data_machine_bits))
+        local datacenter_id_bits = math_floor(attr.snowflake.data_machine_bits / 2)
+        local node_id_bits = math_ceil(attr.snowflake.data_machine_bits / 2)
+        data_machine = gen_data_machine(max_number)
+        if data_machine == nil then
+            return ""
+        end
+
+        local worker_id, datacenter_id = split_data_machine(data_machine,
+            node_id_bits, datacenter_id_bits)
+
+        core.log.info("snowflake init datacenter_id: " ..
+            datacenter_id .. " worker_id: " .. worker_id)
+        snowflake.init(
+            datacenter_id,
+            worker_id,
+            attr.snowflake.snowflake_epoc,
+            node_id_bits,
+            datacenter_id_bits,
+            attr.snowflake.sequence_bits,
+            attr.delta_offset
+        )
+        snowflake_inited = true
+    end
+end
+
+
+-- generate snowflake id
+local function next_id()
+    if snowflake_inited == nil then
+        snowflake_init()
+    end
+    return snowflake:next_id()
+end
+
+
+local function get_request_id(algorithm)
+    if algorithm == "uuid" then
+        return uuid()
+    end
+    return next_id()
+end
+
+
 function _M.rewrite(conf, ctx)
     local headers = ngx.req.get_headers()
-    local uuid_val = uuid()
+    local uuid_val = get_request_id(conf.algorithm)
     if not headers[conf.header_name] then
         core.request.set_header(ctx, conf.header_name, uuid_val)
     end
@@ -53,7 +221,6 @@ function _M.rewrite(conf, ctx)
     end
 end
 
-
 function _M.header_filter(conf, ctx)
     if not conf.include_in_response then
         return
@@ -65,4 +232,25 @@ function _M.header_filter(conf, ctx)
     end
 end
 
+function _M.init()
+    local local_conf = core.config.local_conf()
+    attr = core.table.try_read_attr(local_conf, "plugin_attr", plugin_name)
+    local ok, err = core.schema.check(attr_schema, attr)
+    if not ok then
+        core.log.error("failed to check the plugin_attr[", plugin_name, "]", ": ", err)
+        return
+    end
+    if attr.snowflake.enable then
+        if process.type() == "worker" then
+            ngx.timer.at(0, snowflake_init)
+        end
+    end
+end
+
+function _M.destroy()
+    if snowflake_inited then
+        timers.unregister_timer("plugin#request-id")
+    end
+end
+
 return _M

conf/config-default.yaml

@@ -349,3 +349,12 @@ plugin_attr:
     report_ttl: 3600     # live time for server info in etcd (unit: second)
   dubbo-proxy:
     upstream_multiplex_count: 32
+  request-id:
+    snowflake:
+      enable: false
+      snowflake_epoc: 1609459200000   # the starting timestamp is expressed in milliseconds
+      data_machine_bits: 12           # data machine bit, maximum 31, because Lua cannot do bit operations greater than 31
+      sequence_bits: 10               # each machine generates a maximum of (1 << sequence_bits) serial numbers per millisecond
+      data_machine_ttl: 30            # live time for data_machine in etcd (unit: second)
+      data_machine_interval: 10       # lease renewal interval in etcd (unit: second)
+

docs/en/latest/plugins/request-id.md

@@ -39,7 +39,8 @@ API request. The plugin will not add a request id if the `header_name` is alread
 | Name                | Type    | Requirement | Default        | Valid | Description                                                    |
 | ------------------- | ------- | ----------- | -------------- | ----- | -------------------------------------------------------------- |
 | header_name         | string  | optional    | "X-Request-Id" |       | Request ID header name                                         |
-| include_in_response | boolean | optional    | true          |       | Option to include the unique request ID in the response header |
+| include_in_response | boolean | optional    | true           |       | Option to include the unique request ID in the response header |
+| algorithm           | string  | optional    | "uuid"         | ["uuid", "snowflake"] | ID generation algorithm |
 
 ## How To Enable
 
@@ -72,6 +73,60 @@ X-Request-Id: fe32076a-d0a5-49a6-a361-6c244c1df956
 ......
 ```
 
+### Use the snowflake algorithm to generate an ID
+
+> supports using the Snowflake algorithm to generate ID.
+> read the documentation first before deciding to use snowflake. Because once the configuration information is enabled, you can not arbitrarily adjust the configuration information. Failure to do so may result in duplicate ID being generated.
+
+The Snowflake algorithm is not enabled by default and needs to be configured in 'conf/config.yaml'.
+
+```yaml
+plugin_attr:
+  request-id:
+    snowflake:
+      enable: true
+      snowflake_epoc: 1609459200000
+      data_machine_bits: 12
+      sequence_bits: 10
+      data_machine_ttl: 30
+      data_machine_interval: 10
+```
+
+#### Configuration parameters
+
+| Name                | Type    | Requirement   | Default        |  Valid  | Description                    |
+| ------------------- | ------- | ------------- | -------------- | ------- | ------------------------------ |
+| enable                     | boolean  | optional   | false          |  | When set it to true, enable the snowflake algorithm.  |
+| snowflake_epoc             | integer  | optional   | 1609459200000  |  | Start timestamp (in milliseconds)       |
+| data_machine_bits          | integer  | optional   | 12             |  | Maximum number of supported machines (processes) `1 << data_machine_bits` |
+| sequence_bits              | integer  | optional   | 10             |  | Maximum number of generated ID per millisecond per node `1 << sequence_bits` |
+| data_machine_ttl           | integer  | optional   | 30             |  | Valid time of registration of 'data_machine' in 'etcd' (unit: seconds) |
+| data_machine_interval      | integer  | optional   | 10             |  | Time between 'data_machine' renewal in 'etcd' (unit: seconds) |
+
+- `snowflake_epoc` default start time is  `2021-01-01T00:00:00Z`, and it can support `69 year` approximately to `2090-09-0715:47:35Z` according to the default configuration
+- `data_machine_bits` corresponds to the set of workIDs and datacEnteridd in the snowflake definition. The plug-in aslocates a unique ID to each process. Maximum number of supported processes is `pow(2, data_machine_bits)`. The default number of `12 bits` is up to `4096`.
+- `sequence_bits` defaults to `10 bits` and each process generates up to `1024` ID per second
+
+#### example
+
+> Snowflake supports flexible configuration to meet a wide variety of needs
+
+- Snowflake original configuration
+
+> - Start time 2014-10-20 T15:00:00.000z, accurate to milliseconds. It can last about 69 years
+> - supports up to `1024` processes
+> - Up to `4096` ID per second per process
+
+```yaml
+plugin_attr:
+  request-id:
+    snowflake:
+      enable: true
+      snowflake_epoc: 1413817200000
+      data_machine_bits: 10
+      sequence_bits: 12
+```
+
 ## Disable Plugin
 
 Remove the corresponding json configuration in the plugin configuration to disable the `request-id`.