feat: Request-ID plugin add snowflake algorithm

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"},
+                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,142 @@ 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
+                core.log.error("Etcd grant failure, err: ".. err)
+            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 +219,6 @@ function _M.rewrite(conf, ctx)
     end
 end
 
-
 function _M.header_filter(conf, ctx)
     if not conf.include_in_response then
         return
@@ -65,4 +230,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

@@ -316,3 +316,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

@@ -72,6 +72,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 cannot 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  | required   | false          |  | When set it to true, enable the snowflake algorithm.  |
+| snowflake_epoc             | integer  | required   | 1609459200000  |  | Start timestamp (in milliseconds)       |
+| data_machine_bits          | integer  | deprecated | 12             |  | Maximum number of supported machines (processes) `1 << data_machine_bits` |
+| sequence_bits              | integer  | deprecated | 10             |  | Maximum number of generated ID per millisecond per node `1 << sequence_bits` |
+| data_machine_ttl           | integer  | deprecated | 30             |  | Valid time of registration of 'data_machine' in 'etcd' (unit: seconds) |
+| data_machine_interval      | integer  | deprecated | 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`.

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

@@ -37,8 +37,9 @@ title: request-id
 
 | 名称                | 类型    | 必选项   | 默认值         | 有效值 | 描述                           |
 | ------------------- | ------- | -------- | -------------- | ------ | ------------------------------ |
-| header_name         | string  | 可选 | "X-Request-Id" |        | Request ID header name         |
-| include_in_response | boolean | 可选 | false          |        | 是否需要在返回头中包含该唯一ID |
+| header_name         | string  | 可选 | "X-Request-Id" |                       | Request ID header name         |
+| include_in_response | boolean | 可选 | false          |                       | 是否需要在返回头中包含该唯一ID |
+| algorithm           | string  | 可选 | "uuid"         | ["uuid", "snowflake"] | ID 生成算法 |
 
 ## 如何启用
 
@@ -71,9 +72,63 @@ X-Request-Id: fe32076a-d0a5-49a6-a361-6c244c1df956
 ......
 ```
 
+### 使用 snowflake 算法生成ID
+
+> 支持使用 snowflake 算法来生成ID。
+> 在决定使用snowflake时，请优先阅读一下文档。因为一旦启用配置信息则不可随意调整配置信息。否则可能会导致生成重复ID。
+
+snowflake 算法默认是不启用的，需要在 `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
+```
+
+#### 配置参数
+
+| 名称                | 类型    | 必选项   | 默认值         | 有效值 | 描述                           |
+| ------------------- | ------- | -------- | -------------- | ------ | ------------------------------ |
+| enable                     | boolean  | 可选 | false          |  | 当设置为true时， 启用snowflake算法。      |
+| snowflake_epoc             | integer  | 可选 | 1609459200000  |  | 起始时间戳（单位： 毫秒）       |
+| data_machine_bits          | integer  | 可选 | 12             |  | 最多支持机器（进程）数量 `1 << data_machine_bits` |
+| sequence_bits              | integer  | 可选 | 10             |  | 每个节点每毫秒内最多产生ID数量 `1 << sequence_bits` |
+| data_machine_ttl           | integer  | 可选 | 30             |  | `etcd` 中 `data_machine` 注册有效时间（单位： 秒）|
+| data_machine_interval      | integer  | 可选 | 10             |  | `etcd` 中 `data_machine` 续约间隔时间（单位： 秒）|
+
+- snowflake_epoc 默认起始时间为 `2021-01-01T00:00:00Z`, 按默认配置可以支持 `69年` 大约可以使用到 `2090-09-07 15:47:35Z`
+- data_machine_bits 对应的是 snowflake 定义中的 WorkerID 和 DatacenterIDd的集合，插件会为每一个进程分配一个唯一ID，最大支持进程数为 `pow(2, data_machine_bits)`。默认占 `12 bits` 最多支持 `4096` 个进程。
+- sequence_bits 默认占 `10 bits`, 每个进程每秒最多生成 `1024` 个ID
+
+#### 配置示例
+
+> snowflake 支持灵活配置来满足各式各样的需求
+
+- snowflake 原版配置
+
+> - 起始时间 2014-10-20T15:00:00.000Z， 精确到毫秒为单位。大约可以使用 `69年`
+> - 最多支持 `1024` 个进程
+> - 每个进程每秒最多产生 `4096` 个ID
+
+```yaml
+plugin_attr:
+  request-id:
+    snowflake:
+      enable: true
+      snowflake_epoc: 1413817200000
+      data_machine_bits: 10
+      sequence_bits: 12
+```
+
 ## 禁用插件
 
-在路由 `plugins` 配置块中删除 `request-id 配置，即可禁用该插件，无需重启 APISIX。
+在路由 `plugins` 配置块中删除 `request-id 配置，reload 即可禁用该插件，无需重启 APISIX。
 
 ```shell
 curl http://127.0.0.1:9080/apisix/admin/routes/5 -H 'X-API-KEY: edd1c9f034335f136f87ad84b625c8f1' -X PUT -d '