在工作中经常用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);
}