add optional response() script function

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
 

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

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();
 

src/script.c

@@ -1,5 +1,6 @@
 // Copyright (C) 2013 - Will Glozer.  All rights reserved.
 
+#include <stdlib.h>
 #include <string.h>
 #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;
+}

src/script.h

@@ -7,16 +7,28 @@
 #include <lauxlib.h>
 #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 */

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);
     }
 

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;