[GitHub] [apisix] zuiyangqingzhou commented on a change in pull request #4559: feat: Request-ID plugin add snowflake algorithm

[GitBox <gi...@apache.org>]

zuiyangqingzhou commented on a change in pull request #4559:
URL: https://github.com/apache/apisix/pull/4559#discussion_r671948175



##########
File path: apisix/plugins/request-id.lua
##########
@@ -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
+

Review comment:
       agree




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: notifications-unsubscribe@apisix.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org