wrk 使用

---
--- 环境信息
--- Lua 5.1.5
--- LuaJIT 2.0.5
---

local bit = require "bit"

local zlib = require "zlib"

local data = '%s-%s'

local mac_data_file = "file_1.txt"
local imei_data_file = "file_2.txt"

---
--- 从文本文件载入数据到 table
--- 参数：
---     tb          载入数据的 table
---     file_name   数据文件路径
---
local function load_data(tb, file_name)
    local tpl = "%s"
    for e in io.lines(file_name) do
        table.insert(tb, tpl.format(e))
    end
end

---
--- 从文本文件载入数据到 table
--- 参数：
---     tb          载入数据的 table
---     file_name   数据文件路径
---
local function private_enc(raw, key)
    local compressed_data = zlib.compress(raw)
    local compressed_bytes = { string.byte(compressed_data, 1, -1) }
    local key_bytes = { string.byte(key, 1, -1) }

    local encryped_data = {}
    local key_len = #key_bytes

    for i in pairs(compressed_bytes) do
        local a = bit.bxor(compressed_bytes[i], key_bytes[(i - 1) % key_len + 1])
        table.insert(encryped_data, string.char(a))
    end

    return table.concat(encryped_data, "")
end

local mac_table = {}
local imei_table = {}

load_data(mac_table, mac_data_file)
load_data(imei_table, imei_data_file)

local mac_table_len = #mac_table
local imei_table_len = #imei_table

local counter = 1
local all_threads = {}

---
--- wrk 就会在测试线程已经初始化但还没有启动的时候调用该方法。
--- wrk 会为每一个测试线程调用一次 setup 方法，并传入代表测试线程的对象 thread 作为参数。
--- setup 方法中可操作该 thread 对象，获取信息、存储信息、甚至关闭该线程。
---
function setup(thread)
    thread:set("id", counter)
    counter = counter + 1
    table.insert(all_threads, thread)
end

---
--- 每个线程仅调用 1 次，args 用于获取命令行中传入的参数
---
function init(args)
    -- ok_count、all_count 会存储在线程变量中
    -- wrk.thread:set("all_count", 0)
    -- wrk.thread:set("ok_count", 0)
    ok_count = 0
    all_count = 0
end

---
--- 请求
---
function request()
    local mac_idx = math.random(mac_table_len)
    local imei_idx = math.random(imei_table_len)

    local mac = mac_table[mac_idx]
    local imei = imei_table[imei_idx]

    all_count=all_count + 1

    local raw_data = string.format(data, mac, imei)

    local key = "key"
    local post_data = private_enc(raw_data, key)
    local post_data_len = string.len(post_data)
    
    local header = {
        ["Content-Type"] = "application/json",
        ["Content-Length"] = post_data_len
    }

    return wrk.format("POST", "/r?v=5.2&e=1&c=1", header, post_data)
end

function response(status, headers, body)
        
    if 200 == status then
        ok_count = ok_count + 1
    end
end

function done(summary, latency, requests)
    local msg = "made %d requests and got %d responses ok"
    print(msg:format(all_count, ok_count))
end