call optional done() script function

Makefile

@@ -8,7 +8,11 @@ ifeq ($(TARGET), sunos)
 	LIBS   += -lsocket
 else ifeq ($(TARGET), darwin)
 	LDFLAGS += -pagezero_size 10000 -image_base 100000000
-else
+else ifeq ($(TARGET), linux)
+	LIBS    += -ldl
+	LDFLAGS += -Wl,-E
+else ifeq ($(TARGET), freebsd)
+	CFLAGS  += -D_DECLARE_C99_LDBL_MATH
 	LDFLAGS += -Wl,-E
 endif
 
@@ -20,8 +24,9 @@ ODIR := obj
 OBJ  := $(patsubst %.c,$(ODIR)/%.o,$(SRC))
 
 LDIR     = deps/luajit/src
+LIBS    := -lluajit $(LIBS)
 CFLAGS  += -I $(LDIR)
-LDFLAGS += -L $(LDIR) -lluajit
+LDFLAGS += -L $(LDIR)
 
 all: $(BIN)
 
@@ -39,7 +44,7 @@ $(ODIR): $(LDIR)/libluajit.a
 
 $(ODIR)/bytecode.o: scripts/wrk.lua
 	@echo LUAJIT $<
-	@$(SHELL) -c 'cd $(LDIR) && luajit -b $(PWD)/$< $(PWD)/$@'
+	@$(SHELL) -c 'cd $(LDIR) && ./luajit -b $(PWD)/$< $(PWD)/$@'
 
 $(ODIR)/%.o : %.c
 	@echo CC $<

README

@@ -4,6 +4,11 @@ wrk - a HTTP benchmarking tool
   load when run on a single multi-core CPU. It combines a multithreaded
   design with scalable event notification systems such as epoll and kqueue.
 
+  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.
+
 Basic Usage
 
   wrk -t12 -c400 -d30s http://127.0.0.1:8080/index.html
@@ -24,16 +29,11 @@ Basic Usage
 
 Scripting
 
-  The "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. Per-request
-  actions, particularly building a new HTTP request, will necessarily
-  reduce the amount of load that can be generated.
-
   wrk's public Lua API is:
 
-    init    = function()
+    init    = function(args)
     request = function()
+    done    = function(summary, latency, requests)
 
     wrk = {
       scheme  = "http",
@@ -50,13 +50,37 @@ 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 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
 
-  A user script that only changes the HTTP method, path, adds headers or
-  a body, will have no performance impact. If multiple HTTP requests are
-  necessary they should be generated in the call to init() and returned
-  via a quick lookup in the request() call.
+  The init() function receives any extra command line arguments for the
+  script. Script arguments must be separated from wrk arguments with "--"
+  and scripts that override init() but not request() must call wrk.init()
+
+  The done() function receives a table containing result data, and two
+  statistics objects representing the sampled per-request latency and
+  per-thread request rate. Duration and latency are microsecond values
+  and rate is measured in requests per second.
+
+    latency.min              -- minimum value seen
+    latency.max              -- maximum value seen
+    latency.mean             -- average value seen
+    latency.stdev            -- standard deviation
+    latency:percentile(99.0) -- 99th percentile value
+
+    summary = {
+      duration = N,  -- run duration in microseconds
+      requests = N,  -- total completed requests
+      bytes    = N,  -- total bytes received
+      errors   = {
+        connect = N, -- total socket connection errors
+        read    = N, -- total socket read errors
+        write   = N, -- total socket write errors
+        status  = N, -- total HTTP status codes > 399
+        timeout = N  -- total request timeouts
+      }
+    }
 
 Benchmarking Tips
 
@@ -65,6 +89,12 @@ Benchmarking Tips
   initial connection burst the server's listen(2) backlog should be greater
   than the number of concurrent connections being tested.
 
+  A user script that only changes the HTTP method, path, adds headers or
+  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.
+
 Acknowledgements
 
   wrk contains code from a number of open source projects including the

scripts/wrk.lua

@@ -30,10 +30,11 @@ function wrk.format(method, path, headers, body)
    return table.concat(s, "\r\n")
 end
 
-function wrk.init()    req = wrk.format() end
+function wrk.init(args) req = wrk.format() end
 function wrk.request()  return req         end
 
 init    = wrk.init
 request = wrk.request
+done    = nil
 
 return wrk

src/script.c

@@ -3,20 +3,41 @@
 #include <string.h>
 #include "script.h"
 
-lua_State *script_create(char *scheme, char *host, int port, char *path) {
+typedef struct {
+    char *name;
+    int   type;
+    void *value;
+} table_field;
+
+static int script_stats_len(lua_State *);
+static int script_stats_get(lua_State *);
+static void set_fields(lua_State *, int index, const table_field *);
+
+static const struct luaL_reg statslib[] = {
+    { "__index", script_stats_get },
+    { "__len",   script_stats_len },
+    { NULL,      NULL             }
+};
+
+lua_State *script_create(char *scheme, char *host, char *port, char *path) {
     lua_State *L = luaL_newstate();
     luaL_openlibs(L);
     luaL_dostring(L, "wrk = require \"wrk\"");
 
+    luaL_newmetatable(L, "wrk.stats");
+    luaL_register(L, NULL, statslib);
+    lua_pop(L, 1);
+
+    const table_field fields[] = {
+        { "scheme", LUA_TSTRING, scheme },
+        { "host",   LUA_TSTRING, host   },
+        { "port",   LUA_TSTRING, port   },
+        { "path",   LUA_TSTRING, path   },
+        { NULL,     0,           NULL   },
+    };
+
     lua_getglobal(L, "wrk");
-    lua_pushstring(L, scheme);
-    lua_pushstring(L, host);
-    lua_pushinteger(L, port);
-    lua_pushstring(L, path);
-    lua_setfield(L, 1, "path");
-    lua_setfield(L, 1, "port");
-    lua_setfield(L, 1, "host");
-    lua_setfield(L, 1, "scheme");
+    set_fields(L, 1, fields);
     lua_pop(L, 1);
 
     return L;
@@ -36,14 +57,19 @@ void script_headers(lua_State *L, char **headers) {
     lua_pop(L, 2);
 }
 
-void script_init(lua_State *L, char *script) {
+void script_init(lua_State *L, char *script, int argc, char **argv) {
     if (script && luaL_dofile(L, script)) {
         const char *cause = lua_tostring(L, -1);
-        fprintf(stderr, "script %s failed: %s", script, cause);
+        fprintf(stderr, "%s: %s\n", script, cause);
     }
 
     lua_getglobal(L, "init");
-    lua_call(L, 0, 0);
+    lua_newtable(L);
+    for (int i = 0; i < argc; i++) {
+        lua_pushstring(L, argv[i]);
+        lua_rawseti(L, 2, i);
+    }
+    lua_call(L, 1, 0);
 }
 
 void script_request(lua_State *L, char **buf, size_t *len) {
@@ -62,3 +88,117 @@ bool script_is_static(lua_State *L) {
     lua_pop(L, 3);
     return is_static;
 }
+
+bool script_has_done(lua_State *L) {
+    lua_getglobal(L, "done");
+    bool has_done = lua_type(L, 1) == LUA_TFUNCTION;
+    lua_pop(L, 1);
+    return has_done;
+}
+
+void script_summary(lua_State *L, uint64_t duration, uint64_t requests, uint64_t bytes) {
+    const table_field fields[] = {
+        { "duration", LUA_TNUMBER, &duration },
+        { "requests", LUA_TNUMBER, &requests },
+        { "bytes",    LUA_TNUMBER, &bytes    },
+        { NULL,       0,           NULL      },
+    };
+    lua_newtable(L);
+    set_fields(L, 1, fields);
+}
+
+void script_errors(lua_State *L, errors *errors) {
+    uint64_t e[] = {
+        errors->connect,
+        errors->read,
+        errors->write,
+        errors->status,
+        errors->timeout
+    };
+    const table_field fields[] = {
+        { "connect", LUA_TNUMBER, &e[0] },
+        { "read",    LUA_TNUMBER, &e[1] },
+        { "write",   LUA_TNUMBER, &e[2] },
+        { "status",  LUA_TNUMBER, &e[3] },
+        { "timeout", LUA_TNUMBER, &e[4] },
+        { NULL,      0,           NULL  },
+    };
+    lua_newtable(L);
+    set_fields(L, 2, fields);
+    lua_setfield(L, 1, "errors");
+}
+
+void script_done(lua_State *L, stats *latency, stats *requests) {
+    stats **s;
+
+    lua_getglobal(L, "done");
+    lua_pushvalue(L, 1);
+
+    s = (stats **) lua_newuserdata(L, sizeof(stats **));
+    *s = latency;
+    luaL_getmetatable(L, "wrk.stats");
+    lua_setmetatable(L, 4);
+
+    s = (stats **) lua_newuserdata(L, sizeof(stats **));
+    *s = requests;
+    luaL_getmetatable(L, "wrk.stats");
+    lua_setmetatable(L, 5);
+
+    lua_call(L, 3, 0);
+    lua_pop(L, 5);
+}
+
+static stats *checkstats(lua_State *L) {
+    stats **s = luaL_checkudata(L, 1, "wrk.stats");
+    luaL_argcheck(L, s != NULL, 1, "`stats' expected");
+    return *s;
+}
+
+static int script_stats_percentile(lua_State *L) {
+    stats *s = checkstats(L);
+    lua_Number p = luaL_checknumber(L, 2);
+    lua_pushnumber(L, stats_percentile(s, p));
+    return 1;
+}
+
+static int script_stats_get(lua_State *L) {
+    stats *s = checkstats(L);
+    if (lua_isnumber(L, 2)) {
+        int index = luaL_checkint(L, 2);
+        lua_pushnumber(L, s->data[index - 1]);
+    } else if (lua_isstring(L, 2)) {
+        const char *method = lua_tostring(L, 2);
+        if (!strcmp("min",   method)) lua_pushnumber(L, s->min);
+        if (!strcmp("max",   method)) lua_pushnumber(L, s->max);
+        if (!strcmp("mean",  method)) lua_pushnumber(L, stats_mean(s));
+        if (!strcmp("stdev", method)) lua_pushnumber(L, stats_stdev(s, stats_mean(s)));
+        if (!strcmp("percentile", method)) {
+            lua_pushcfunction(L, script_stats_percentile);
+        }
+    }
+    return 1;
+}
+
+static int script_stats_len(lua_State *L) {
+    stats *s = checkstats(L);
+    lua_pushinteger(L, s->limit);
+    return 1;
+}
+
+static void set_fields(lua_State *L, int index, const table_field *fields) {
+    for (int i = 0; fields[i].name; i++) {
+        table_field f = fields[i];
+        switch (f.value == NULL ? LUA_TNIL : f.type) {
+            case LUA_TNUMBER:
+                lua_pushinteger(L, *((lua_Integer *) f.value));
+                break;
+            case LUA_TSTRING:
+                lua_pushstring(L, (const char *) f.value);
+                break;
+            case LUA_TNIL:
+                lua_pushnil(L);
+                break;
+        }
+        lua_setfield(L, index, f.name);
+    }
+}

src/script.h

@@ -5,11 +5,18 @@
 #include <lua.h>
 #include <lualib.h>
 #include <lauxlib.h>
+#include "stats.h"
 
-lua_State *script_create(char *, char *, int, char *);
+lua_State *script_create(char *, char *, char *, char *);
 void script_headers(lua_State *, char **);
-void script_init(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 *);
+
 bool script_is_static(lua_State *);
+bool script_has_done(lua_State *L);
+void script_summary(lua_State *, uint64_t, uint64_t, uint64_t);
+void script_errors(lua_State *, errors *);
 
 #endif /* SCRIPT_H */

src/stats.c

@@ -46,7 +46,10 @@ static int stats_compare(const void *a, const void *b) {
 
 long double stats_summarize(stats *stats) {
     qsort(stats->data, stats->limit, sizeof(uint64_t), &stats_compare);
+    return stats_mean(stats);
+}
 
+long double stats_mean(stats *stats) {
     if (stats->limit == 0) return 0.0;
 
     uint64_t sum = 0;

src/stats.h

@@ -7,6 +7,14 @@
 #define MAX(X, Y) ((X) > (Y) ? (X) : (Y))
 #define MIN(X, Y) ((X) < (Y) ? (X) : (Y))
 
+typedef struct {
+    uint32_t connect;
+    uint32_t read;
+    uint32_t write;
+    uint32_t status;
+    uint32_t timeout;
+} errors;
+
 typedef struct {
     uint64_t samples;
     uint64_t index;
@@ -24,6 +32,7 @@ void stats_rewind(stats *);
 void stats_record(stats *, uint64_t);
 
 long double stats_summarize(stats *);
+long double stats_mean(stats *);
 long double stats_stdev(stats *stats, long double);
 long double stats_within_stdev(stats *, long double, long double, uint64_t);
 uint64_t stats_percentile(stats *, long double);

src/wrk.c

@@ -130,15 +130,17 @@ 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];
         t->connections = connections;
         t->stop_at     = stop_at;
 
-        t->L = script_create(schema, host, (int) strtol(port, NULL, 10), path);
+        t->L = script_create(schema, host, port, path);
         script_headers(t->L, headers);
-        script_init(t->L, cfg.script);
+        script_init(t->L, cfg.script, argc - optind, &argv[optind]);
+        if (L == NULL) L = t->L;
 
         if (pthread_create(&t->thread, NULL, &thread_main, t)) {
             char *msg = strerror(errno);
@@ -202,6 +204,12 @@ 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));
 
+    if (script_has_done(L)) {
+        script_summary(L, runtime_us, complete, bytes);
+        script_errors(L, &errors);
+        script_done(L, statistics.latency, statistics.requests);
+    }
+
     return 0;
 }
 

src/wrk.h

@@ -22,14 +22,6 @@
 #define CALIBRATE_DELAY_MS  500
 #define TIMEOUT_INTERVAL_MS 2000
 
-typedef struct {
-    uint32_t connect;
-    uint32_t read;
-    uint32_t write;
-    uint32_t status;
-    uint32_t timeout;
-} errors;
-
 typedef struct {
     pthread_t thread;
     aeEventLoop *loop;