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,20 @@ 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
+
+
 return _M

conf/config-default.yaml

@@ -316,3 +316,13 @@ 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
+      delta_offset: 1                 # timestamp delta offset (unit: milliseconds)
+      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/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,6 +72,99 @@ X-Request-Id: fe32076a-d0a5-49a6-a361-6c244c1df956
 ......
 ```
 
+### 使用 snowflake 算法
+
+> 支持使用 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             |  | 数据机器位`datacenterId` + `workerId`（1 << node_id_bits） |
+| sequence_bits              | integer  | 可选 | 10             |  | 每个节点每毫秒内最多产生ID数量 （1 << sequence_bits） |
+| delta_offset               | integer  | 可选 | 1              |  | 时间戳增量偏移（单位： 毫秒） |
+| 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
+- delta_offset 时间戳增量偏移 (单位: 毫秒) [`delta_offset=1` 每毫秒, `delta_offset=10` 每10毫秒, `delta_offset=100` 每100毫秒, `delta_offset=1000` 每1秒`]
+
+#### 配置示例
+
+> 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
+```
+
+- [sonyflake](https://github.com/sony/sonyflake)
+
+> - 39 bit 为时间戳，精确到 `10ms` 大约可以使用174年。
+> - 16 bit 做为机器号， 可同时支持 `65536` 个进程运行。
+> - 8 bit 做为序列号，每10毫最大生成256个，1秒最多生成25600个。
+
+```yaml
+plugin_attr:
+  request-id:
+    snowflake:
+      enable: true
+      snowflake_epoc: 1413817200000
+      data_machine_bits: 16
+      sequence_bits: 8
+      delta_offset: 10
+```
+
+- [baidu UidGenerator](https://github.com/baidu/uid-generator)
+
+不支持 Double RingBuffer
+
+> - delta seconds (28 bits)
+> - 当前时间，相对于时间基点"2016-05-20"的增量值，单位：秒，最多可支持约8.7年
+> - worker id (22 bits) 机器id，最多可支持约420w次机器启动。内置实现为在启动时由数据库分配，默认分配策略为用后即弃，后续可提供复用策略。
+> - sequence (13 bits) 每秒下的并发序列，13 bits可支持每秒8192个并发。
+
+```yaml
+plugin_attr:
+  request-id:
+    snowflake:
+      enable: true
+      snowflake_epoc: 1463644800000
+      data_machine_bits: 22
+      sequence_bits: 13
+      delta_offset: 1000
+```
+
 ## 禁用插件
 
 在路由 `plugins` 配置块中删除 `request-id 配置，即可禁用该插件，无需重启 APISIX。

rockspec/apisix-master-0.rockspec

@@ -67,6 +67,7 @@ dependencies = {
     "lua-resty-consul = 0.3-2",
     "penlight = 1.9.2-1",
     "ext-plugin-proto = 0.1.1",
+    "api7-snowflake = 2.0-1",
 }
 
 build = {