在工作中經常用wrk對接口進行簡單的壓測，最近工作中測試接口需要對參數進行簽名校驗，借這個機會，打算仔細研究下wrk；
wrk命令選項
wrk命令選項如下圖所示：
image.png
例子：
./wrk -c 1 -d 10s -t 1 -s test.lua http://10.221.84.140:8080/test/anchor/info
其中-s選項指定lua腳本文件，下面舉一個腳本的例子：
package.path="/Users/allan/Softwares/wrk/?.lua;;"
local sha1=require("sha1")
local timestamp=os.time()*1000
local token=sha1("name=allan".."timestamp"..timestamp.."age=56".."4dfp*&ddddd4445")
wrk.method = 'POST'
wrk.body = "timestamp="..timestamp.."&name=allan&age=56"
wrk.headers["App-Key"] = "android"
--wrk.headers["Timestamp"]=timestamp
wrk.headers["Authorization"] = token
wrk.headers["Content-Type"] = "application/x-www-form-urlencoded"
function response(status, headers, body)
print(body)
end
可以看到lua腳本中可以調用第三方庫，動態設置參數，而且wrk基于epoll,性能強悍;
wrk源碼
wrk是開源的，其源碼地址為https://github.com/wg/wrk, 采用C語言實現；
wrk定義了全局變量wrk, 提供了如下函數供擴展:
- setup(thread):
初始化時調用,在init方法調用之前,thread有get、set、stop方法，有addr屬性，例如：
function setup(thread)
thread:set("id", counter)
table.insert(threads, thread)
counter = counter + 1
end
- init(args):先調用setup,再調用init,只調用一次
- request()：發送請求之前調用，例如
request = function()
path = paths[counter]
counter = counter + 1
if counter > #paths then
counter = 0
end
return wrk.format(nil, path)
end
- response(status, headers, body): 接收到請求響應之后執行，例如
function response()
if counter == 100 then
wrk.thread:stop()
end
counter = counter + 1
end
- done(summary, latency, requests): 所有線程執行完畢，最后調用
done = function(summary, latency, requests)
io.write("------------------------------\n")
for _, p in pairs({ 50, 90, 99, 99.999 }) do
n = latency:percentile(p)
io.write(string.format("%g%%,%d\n", p, n))
end
end
wrk源文件
wrk實現很簡潔，主要的源文件包括：
- wrk.c:
啟動入口，包含main方法 - script.c：
通過c調用luajit - ae.c:
網絡多路復用層的實現，包括epoll、evport、kqueue和select；FreeBSD和Apple默認使用kqueue,Linux使用epoll,Sun使用evport
下面具體看看wrk是如何實現的：
wrk有幾個重要的數據結構，包括thread和connection:
typedef struct {
pthread_t thread;//操作系統線程對象
aeEventLoop *loop;//持有epoll對象
struct addrinfo *addr;//連接地址信息
uint64_t connections;//連接數
uint64_t complete;//完成請求數
uint64_t requests;//發送請求數
uint64_t bytes;//發送字節數
uint64_t start;
lua_State *L;//lua句柄
errors errors;
struct connection *cs;//連接對象
} thread;
typedef struct connection {
thread *thread; //所屬線程
http_parser parser;
enum {
FIELD, VALUE
} state;
int fd;
SSL *ssl;
bool delayed;
uint64_t start;
char *request;
size_t length;
size_t written;
uint64_t pending;
buffer headers;
buffer body;
char buf[RECVBUF];
} connection;
wrk初始化邏輯：
- 通過parse_args函數解析命令行參數；
- 創建lua_State，檢查是否可以調用lua,是否可以連接到測試地址；
- 根據命令行參數中指定的線程數，初始化線程對象，為每個線程創建lua_State,線程的執行函數定義在thread_main方法：
void *thread_main(void *arg) {
thread *thread = arg;
char *request = NULL;
size_t length = 0;
if (!cfg.dynamic) {//如果定義了request函數，則在每次發送請求前執行request函數；否則執行wrk.request函數
script_request(thread->L, &request, &length);
}
thread->cs = zcalloc(thread->connections * sizeof(connection));
connection *c = thread->cs;
for (uint64_t i = 0; i < thread->connections; i++, c++) {
//創建connection對象，每個線程要創建的connection對象個數，由命令行參數-c,-t決定，例如-c 100 -t 4,則意味著每個線程創建25個連接
c->thread = thread;
c->ssl = cfg.ctx ? SSL_new(cfg.ctx) : NULL;
c->request = request;
c->length = length;
c->delayed = cfg.delay;
connect_socket(thread, c);
}
aeEventLoop *loop = thread->loop;
aeCreateTimeEvent(loop, RECORD_INTERVAL_MS, record_rate, thread, NULL);//每100ms觸發一次record_rate函數
thread->start = time_us();
aeMain(loop);
aeDeleteEventLoop(loop);
zfree(thread->cs);
return NULL;
}
//wrk.init中定義的，req是字符串，和我們抓包看到的http請求內容一樣
local req = wrk.format()
wrk.request = function()
return req
end
static void socket_connected(aeEventLoop *loop, int fd, void *data, int mask) {
connection *c = data;
switch (sock.connect(c, cfg.host)) {
case OK: break;
case ERROR: goto error;
case RETRY: return;
}
http_parser_init(&c->parser, HTTP_RESPONSE);
c->written = 0;
//注冊事件，當有數據可讀時，觸發socket_readable函數，當可以寫入數據時，觸發socket_writeable函數
aeCreateFileEvent(c->thread->loop, fd, AE_READABLE, socket_readable, c);
aeCreateFileEvent(c->thread->loop, fd, AE_WRITABLE, socket_writeable, c);
return;
error:
c->thread->errors.connect++;
reconnect_socket(c->thread, c);
}
