diff --git a/README b/README index 6ceb90e..7b83596 100644 --- a/README +++ b/README @@ -7,7 +7,7 @@ wrk - a HTTP benchmarking tool This "scripted" branch of wrk includes LuaJIT and a Lua script may be used to perform minor alterations to the default HTTP request or even even generate a completely new HTTP request each time. The script may - also perform custom reporting at the end of a run. + also inspect each response and perform custom reporting. Basic Usage @@ -31,9 +31,10 @@ Scripting wrk's public Lua API is: - init = function(args) - request = function() - done = function(summary, latency, requests) + init = function(args) + request = function() + response = function(status, headers, body) + done = function(summary, latency, requests) wrk = { scheme = "http", @@ -50,9 +51,10 @@ Scripting wrk.format returns a HTTP request string containing the passed parameters merged with values from the wrk table. - global init -- function to be called when the thread is initialized - global request -- function returning the HTTP message for each request - global done -- optional function to be called with results of run + global init -- function called when the thread is initialized + global request -- function returning the HTTP message for each request + global response -- optional function called with HTTP response data + global done -- optional function called with results of run The init() function receives any extra command line arguments for the script. Script arguments must be separated from wrk arguments with "--" @@ -68,6 +70,7 @@ Scripting latency.mean -- average value seen latency.stdev -- standard deviation latency:percentile(99.0) -- 99th percentile value + latency[i] -- raw sample value summary = { duration = N, -- run duration in microseconds @@ -93,7 +96,8 @@ Benchmarking Tips a body, will have no performance impact. If multiple HTTP requests are necessary they should be pre-generated and returned via a quick lookup in the request() call. Per-request actions, particularly building a new HTTP - request, will necessarily reduce the amount of load that can be generated. + request, and use of response() will necessarily reduce the amount of load + that can be generated. Acknowledgements diff --git a/scripts/wrk.lua b/scripts/wrk.lua index 4301037..03cbced 100644 --- a/scripts/wrk.lua +++ b/scripts/wrk.lua @@ -33,8 +33,9 @@ end function wrk.init(args) req = wrk.format() end function wrk.request() return req end -init = wrk.init -request = wrk.request -done = nil +init = wrk.init +request = wrk.request +response = nil +done = nil return wrk diff --git a/src/main.h b/src/main.h index d9c5a17..a1e1cdb 100644 --- a/src/main.h +++ b/src/main.h @@ -40,7 +40,11 @@ static int check_timeouts(aeEventLoop *, long long, void *); static void socket_connected(aeEventLoop *, int, void *, int); static void socket_writeable(aeEventLoop *, int, void *, int); static void socket_readable(aeEventLoop *, int, void *, int); -static int request_complete(http_parser *); + +static int response_complete(http_parser *); +static int header_field(http_parser *, const char *, size_t); +static int header_value(http_parser *, const char *, size_t); +static int response_body(http_parser *, const char *, size_t); static uint64_t time_us(); diff --git a/src/script.c b/src/script.c index f7f2a41..9f31cf2 100644 --- a/src/script.c +++ b/src/script.c @@ -1,5 +1,6 @@ // Copyright (C) 2013 - Will Glozer. All rights reserved. +#include #include #include "script.h" @@ -75,11 +76,28 @@ void script_init(lua_State *L, char *script, int argc, char **argv) { void script_request(lua_State *L, char **buf, size_t *len) { lua_getglobal(L, "request"); lua_call(L, 0, 1); - *buf = (char *) lua_tostring(L, 1); - *len = (size_t) lua_strlen(L, 1); + *buf = (char *) lua_tolstring(L, 1, len); lua_pop(L, 1); } +void script_response(lua_State *L, int status, buffer *headers, buffer *body) { + lua_getglobal(L, "response"); + lua_pushinteger(L, status); + lua_newtable(L); + + for (char *c = headers->buffer; c < headers->cursor; ) { + c = buffer_pushlstring(L, c); + c = buffer_pushlstring(L, c); + lua_rawset(L, -3); + } + + lua_pushlstring(L, body->buffer, body->cursor - body->buffer); + lua_call(L, 3, 0); + + buffer_reset(headers); + buffer_reset(body); +} + bool script_is_static(lua_State *L) { lua_getglobal(L, "wrk"); lua_getfield(L, 1, "request"); @@ -89,11 +107,22 @@ bool script_is_static(lua_State *L) { return is_static; } +bool script_want_response(lua_State *L) { + lua_getglobal(L, "response"); + bool defined = lua_type(L, 1) == LUA_TFUNCTION; + lua_pop(L, 1); + return defined; +} + bool script_has_done(lua_State *L) { lua_getglobal(L, "done"); - bool has_done = lua_type(L, 1) == LUA_TFUNCTION; + bool defined = lua_type(L, 1) == LUA_TFUNCTION; lua_pop(L, 1); - return has_done; + return defined; +} + +void script_header_done(lua_State *L, luaL_Buffer *buffer) { + luaL_pushresult(buffer); } void script_summary(lua_State *L, uint64_t duration, uint64_t requests, uint64_t bytes) { @@ -145,7 +174,7 @@ void script_done(lua_State *L, stats *latency, stats *requests) { lua_setmetatable(L, 5); lua_call(L, 3, 0); - lua_pop(L, 5); + lua_pop(L, 1); } static stats *checkstats(lua_State *L) { @@ -202,3 +231,24 @@ static void set_fields(lua_State *L, int index, const table_field *fields) { lua_setfield(L, index, f.name); } } + +void buffer_append(buffer *b, const char *data, size_t len) { + size_t used = b->cursor - b->buffer; + while (used + len + 1 >= b->length) { + b->length += 1024; + b->buffer = realloc(b->buffer, b->length); + b->cursor = b->buffer + used; + } + memcpy(b->cursor, data, len); + b->cursor += len; +} + +void buffer_reset(buffer *b) { + b->cursor = b->buffer; +} + +char *buffer_pushlstring(lua_State *L, char *start) { + char *end = strchr(start, 0); + lua_pushlstring(L, start, end - start); + return end + 1; +} diff --git a/src/script.h b/src/script.h index 8a29b80..72fdaf9 100644 --- a/src/script.h +++ b/src/script.h @@ -7,16 +7,28 @@ #include #include "stats.h" +typedef struct { + char *buffer; + size_t length; + char *cursor; +} buffer; + lua_State *script_create(char *, char *, char *, char *); void script_headers(lua_State *, char **); void script_init(lua_State *, char *, int, char **); void script_done(lua_State *, stats *, stats *); void script_request(lua_State *, char **, size_t *); +void script_response(lua_State *, int, buffer *, buffer *); bool script_is_static(lua_State *); +bool script_want_response(lua_State *L); bool script_has_done(lua_State *L); void script_summary(lua_State *, uint64_t, uint64_t, uint64_t); void script_errors(lua_State *, errors *); +void buffer_append(buffer *, const char *, size_t); +void buffer_reset(buffer *); +char *buffer_pushlstring(lua_State *, char *); + #endif /* SCRIPT_H */ diff --git a/src/wrk.c b/src/wrk.c index a2fc0cc..37cc3d9 100644 --- a/src/wrk.c +++ b/src/wrk.c @@ -10,6 +10,7 @@ static struct config { uint64_t duration; uint64_t timeout; bool latency; + bool dynamic; char *script; SSL_CTX *ctx; } cfg; @@ -27,8 +28,8 @@ static struct sock sock = { .write = sock_write }; -static const struct http_parser_settings parser_settings = { - .on_message_complete = request_complete +static struct http_parser_settings parser_settings = { + .on_message_complete = response_complete }; static volatile sig_atomic_t stop = 0; @@ -130,7 +131,6 @@ int main(int argc, char **argv) { thread *threads = zcalloc(cfg.threads * sizeof(thread)); uint64_t connections = cfg.connections / cfg.threads; uint64_t stop_at = time_us() + (cfg.duration * 1000000); - lua_State *L = NULL; for (uint64_t i = 0; i < cfg.threads; i++) { thread *t = &threads[i]; @@ -140,7 +140,15 @@ int main(int argc, char **argv) { t->L = script_create(schema, host, port, path); script_headers(t->L, headers); script_init(t->L, cfg.script, argc - optind, &argv[optind]); - if (L == NULL) L = t->L; + + if (i == 0) { + cfg.dynamic = !script_is_static(t->L); + if (script_want_response(t->L)) { + parser_settings.on_header_field = header_field; + parser_settings.on_header_value = header_value; + parser_settings.on_body = response_body; + } + } if (pthread_create(&t->thread, NULL, &thread_main, t)) { char *msg = strerror(errno); @@ -204,6 +212,7 @@ int main(int argc, char **argv) { printf("Requests/sec: %9.2Lf\n", req_per_s); printf("Transfer/sec: %10sB\n", format_binary(bytes_per_s)); + lua_State *L = threads[0].L; if (script_has_done(L)) { script_summary(L, runtime_us, complete, bytes); script_errors(L, &errors); @@ -225,9 +234,8 @@ void *thread_main(void *arg) { char *request = NULL; size_t length = 0; - if (script_is_static(thread->L)) { + if (!cfg.dynamic) { script_request(thread->L, &request, &length); - thread->L = NULL; } connection *c = thread->cs; @@ -340,20 +348,53 @@ static int sample_rate(aeEventLoop *loop, long long id, void *data) { return thread->interval; } -static int request_complete(http_parser *parser) { +static int header_field(http_parser *parser, const char *at, size_t len) { + connection *c = parser->data; + if (c->state == VALUE) { + *c->headers.cursor++ = '\0'; + c->state = FIELD; + } + buffer_append(&c->headers, at, len); + return 0; +} + +static int header_value(http_parser *parser, const char *at, size_t len) { + connection *c = parser->data; + if (c->state == FIELD) { + *c->headers.cursor++ = '\0'; + c->state = VALUE; + } + buffer_append(&c->headers, at, len); + return 0; +} + +static int response_body(http_parser *parser, const char *at, size_t len) { + connection *c = parser->data; + buffer_append(&c->body, at, len); + return 0; +} + +static int response_complete(http_parser *parser) { connection *c = parser->data; thread *thread = c->thread; uint64_t now = time_us(); + int status = parser->status_code; thread->complete++; thread->requests++; stats_record(thread->latency, now - c->start); - if (parser->status_code > 399) { + if (status > 399) { thread->errors.status++; } + if (c->headers.buffer) { + *c->headers.cursor++ = '\0'; + script_response(thread->L, status, &c->headers, &c->body); + c->state = FIELD; + } + if (now >= thread->stop_at) { aeStop(thread->loop); goto done; @@ -420,7 +461,7 @@ static void socket_writeable(aeEventLoop *loop, int fd, void *data, int mask) { connection *c = data; thread *thread = c->thread; - if (!c->written && thread->L) { + if (!c->written && cfg.dynamic) { script_request(thread->L, &c->request, &c->length); } diff --git a/src/wrk.h b/src/wrk.h index 2f6e427..db28dca 100644 --- a/src/wrk.h +++ b/src/wrk.h @@ -44,12 +44,17 @@ typedef struct { typedef struct connection { thread *thread; http_parser parser; + enum { + FIELD, VALUE + } state; int fd; SSL *ssl; uint64_t start; char *request; size_t length; size_t written; + buffer headers; + buffer body; char buf[RECVBUF]; } connection;