Skip to content

Commit

Permalink
feat: Request-ID plugin add snowflake algorithm (#4559)
Browse files Browse the repository at this point in the history
  • Loading branch information
@dickens7
dickens7 authored Aug 9, 2021
1 parent a7c040f commit e127cc7
Show file tree
Hide file tree
Showing 6 changed files with 585 additions and 12 deletions.
204 changes: 196 additions & 8 deletions apisix/plugins/request-id.lua
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
}


Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
9 changes: 9 additions & 0 deletions conf/config-default.yaml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)

57 changes: 56 additions & 1 deletion docs/en/latest/plugins/request-id.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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`.
Expand Down
Loading

0 comments on commit e127cc7

Please sign in to comment.