generate requests with lua script

.gitignore

@@ -1,2 +1,3 @@
-obj/*
+*.o
+*.a
 wrk

Makefile

@@ -6,33 +6,53 @@ TARGET  := $(shell uname -s | tr [A-Z] [a-z] 2>/dev/null || echo unknown)
 ifeq ($(TARGET), sunos)
 	CFLAGS += -D_PTHREADS -D_POSIX_C_SOURCE=200112L
 	LIBS   += -lsocket
+else ifeq ($(TARGET), darwin)
+	LDFLAGS += -pagezero_size 10000 -image_base 100000000
+else
+	LDFLAGS += -Wl,-E
 endif
 
-SRC  := wrk.c net.c ssl.c aprintf.c stats.c units.c ae.c zmalloc.c http_parser.c tinymt64.c
+SRC  := wrk.c net.c ssl.c aprintf.c stats.c script.c units.c \
+		ae.c zmalloc.c http_parser.c tinymt64.c
 BIN  := wrk
 
 ODIR := obj
 OBJ  := $(patsubst %.c,$(ODIR)/%.o,$(SRC))
 
+LDIR     = deps/luajit/src
+CFLAGS  += -I $(LDIR)
+LDFLAGS += -L $(LDIR) -lluajit
+
 all: $(BIN)
 
 clean:
 	$(RM) $(BIN) obj/*
 
-$(BIN): $(OBJ)
+$(BIN): $(OBJ) $(ODIR)/bytecode.o
-	$(CC) $(LDFLAGS) -o $@ $^ $(LIBS)
+	@echo LINK $(BIN)
+	@$(CC) $(LDFLAGS) -o $@ $^ $(LIBS)
 
 $(OBJ): config.h Makefile | $(ODIR)
 
-$(ODIR):
+$(ODIR): $(LDIR)/libluajit.a
-	@mkdir $@
+	@mkdir -p $@
+
+$(ODIR)/bytecode.o: scripts/wrk.lua
+	@echo LUAJIT $<
+	@$(SHELL) -c 'cd $(LDIR) && luajit -b $(PWD)/$< $(PWD)/$@'
 
 $(ODIR)/%.o : %.c
-	$(CC) $(CFLAGS) -c -o $@ $<
+	@echo CC $<
+	@$(CC) $(CFLAGS) -c -o $@ $<
+
+$(LDIR)/libluajit.a:
+	@echo Building LuaJit...
+	@$(MAKE) -C $(LDIR) BUILDMODE=static
 
 .PHONY: all clean
 .SUFFIXES:
-.SUFFIXES: .c .o
+.SUFFIXES: .c .o .lua
 
 vpath %.c   src
 vpath %.h   src
+vpath %.lua scripts

NOTICE

@@ -81,6 +81,32 @@ DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
 OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
 USE OR OTHER DEALINGS IN THE SOFTWARE.
 
+=========================================================================
+==  LuaJIT Notice                                                      ==
+=========================================================================
+
+LuaJIT -- a Just-In-Time Compiler for Lua. http://luajit.org/
+
+Copyright (C) 2005-2013 Mike Pall. All rights reserved.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.
+
 =========================================================================
 ==  Tiny Mersenne Twister (TinyMT) Notice                              ==
 =========================================================================

README

@@ -22,6 +22,42 @@ Basic Usage
   Requests/sec: 748868.53
   Transfer/sec:    606.33MB
 
+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()
+    request = function()
+
+    wrk = {
+      scheme  = "http",
+      host    = "localhost",
+      port    = nil,
+      method  = "GET",
+      path    = "/",
+      headers = {},
+      body    = nil
+    }
+
+    function wrk.format(method, path, headers, body)
+
+      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
+
+  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.
+
 Benchmarking Tips
 
   The machine running wrk must have a sufficient number of ephemeral ports
@@ -32,6 +68,6 @@ Benchmarking Tips
 Acknowledgements
 
   wrk contains code from a number of open source projects including the
-  'ae' event loop from redis, the nginx/joyent/node.js 'http-parser' and
+  'ae' event loop from redis, the nginx/joyent/node.js 'http-parser',
-  the Tiny Mersenne Twister PRNG. Please consult the NOTICE file for
+  Mike Pall's LuaJIT, and the Tiny Mersenne Twister PRNG. Please consult
-  licensing details.
+  the NOTICE file for licensing details.

scripts/wrk.lua

@@ -0,0 +1,39 @@
+local wrk = {
+   scheme  = "http",
+   host    = "localhost",
+   port    = nil,
+   method  = "GET",
+   path    = "/",
+   headers = {},
+   body    = nil
+}
+
+function wrk.format(method, path, headers, body)
+   local host    = wrk.host
+   local method  = method  or wrk.method
+   local path    = path    or wrk.path
+   local headers = headers or wrk.headers
+   local body    = body    or wrk.body
+   local s       = {}
+
+   headers["Host"]           = port and (host .. ":" .. port) or host
+   headers["Content-Length"] = body and string.len(body)
+
+   s[1] = string.format("%s %s HTTP/1.1", method, path)
+   for name, value in pairs(headers) do
+      s[#s+1] = string.format("%s: %s", name, value)
+   end
+
+   s[#s+1] = ""
+   s[#s+1] = body or ""
+
+   return table.concat(s, "\r\n")
+end
+
+function wrk.init()    req = wrk.format() end
+function wrk.request() return req         end
+
+init    = wrk.init
+request = wrk.request
+
+return wrk

src/main.h

@@ -45,7 +45,6 @@ static int request_complete(http_parser *);
 static uint64_t time_us();
 
 static char *extract_url_part(char *, struct http_parser_url *, enum http_parser_url_fields);
-static char *format_request(char *, char *, char *, char **);
 
 static int parse_args(struct config *, char **, char **, int, char **);
 static void print_stats_header();

src/script.c

@@ -0,0 +1,64 @@
+// Copyright (C) 2013 - Will Glozer.  All rights reserved.
+
+#include <string.h>
+#include "script.h"
+
+lua_State *script_create(char *scheme, char *host, int port, char *path) {
+    lua_State *L = luaL_newstate();
+    luaL_openlibs(L);
+    luaL_dostring(L, "wrk = require \"wrk\"");
+
+    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");
+    lua_pop(L, 1);
+
+    return L;
+}
+
+void script_headers(lua_State *L, char **headers) {
+    lua_getglobal(L, "wrk");
+    lua_getfield(L, 1, "headers");
+    for (char **h = headers; *h; h++) {
+        char *p = strchr(*h, ':');
+        if (p && p[1] == ' ') {
+            lua_pushlstring(L, *h, p - *h);
+            lua_pushstring(L, p + 2);
+            lua_settable(L, 2);
+        }
+    }
+    lua_pop(L, 2);
+}
+
+void script_init(lua_State *L, char *script) {
+    if (script && luaL_dofile(L, script)) {
+        const char *cause = lua_tostring(L, -1);
+        fprintf(stderr, "script %s failed: %s", script, cause);
+    }
+
+    lua_getglobal(L, "init");
+    lua_call(L, 0, 0);
+}
+
+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);
+    lua_pop(L, 1);
+}
+
+bool script_is_static(lua_State *L) {
+    lua_getglobal(L, "wrk");
+    lua_getfield(L, 1, "request");
+    lua_getglobal(L, "request");
+    bool is_static = lua_equal(L, 2, 3);
+    lua_pop(L, 3);
+    return is_static;
+}

src/script.h

@@ -0,0 +1,15 @@
+#ifndef SCRIPT_H
+#define SCRIPT_H
+
+#include <stdbool.h>
+#include <lua.h>
+#include <lualib.h>
+#include <lauxlib.h>
+
+lua_State *script_create(char *, char *, int, char *);
+void script_headers(lua_State *, char **);
+void script_init(lua_State *, char *);
+void script_request(lua_State *, char **, size_t *);
+bool script_is_static(lua_State *);
+
+#endif /* SCRIPT_H */

src/wrk.c

@@ -10,16 +10,10 @@ static struct config {
     uint64_t duration;
     uint64_t timeout;
     bool     latency;
+    char    *script;
     SSL_CTX *ctx;
 } cfg;
 
-static struct {
-    char *method;
-    char *body;
-    size_t size;
-    char *buf;
-} req;
-
 static struct {
     stats *latency;
     stats *requests;
@@ -50,9 +44,8 @@ static void usage() {
            "    -d, --duration    <T>  Duration of test           \n"
            "    -t, --threads     <N>  Number of threads to use   \n"
            "                                                      \n"
+           "    -s, --script      <S>  Load Lua script file       \n"
            "    -H, --header      <H>  Add header to request      \n"
-           "    -M, --method      <M>  HTTP method                \n"
-           "        --body        <B>  Request body               \n"
            "        --latency          Print latency statistics   \n"
            "        --timeout     <T>  Socket/request timeout     \n"
            "    -v, --version          Print version details      \n"
@@ -129,8 +122,6 @@ int main(int argc, char **argv) {
     signal(SIGPIPE, SIG_IGN);
     signal(SIGINT,  SIG_IGN);
     cfg.addr = *addr;
-    req.buf  = format_request(host, port, path, headers);
-    req.size = strlen(req.buf);
 
     pthread_mutex_init(&statistics.mutex, NULL);
     statistics.latency  = stats_alloc(SAMPLES);
@@ -145,6 +136,10 @@ int main(int argc, char **argv) {
         t->connections = connections;
         t->stop_at     = stop_at;
 
+        t->L = script_create(schema, host, (int) strtol(port, NULL, 10), path);
+        script_headers(t->L, headers);
+        script_init(t->L, cfg.script);
+
         if (pthread_create(&t->thread, NULL, &thread_main, t)) {
             char *msg = strerror(errno);
             fprintf(stderr, "unable to create thread %"PRIu64" %s\n", i, msg);
@@ -219,11 +214,21 @@ void *thread_main(void *arg) {
     tinymt64_init(&thread->rand, time_us());
     thread->latency = stats_alloc(100000);
 
+    char *request = NULL;
+    size_t length = 0;
+
+    if (script_is_static(thread->L)) {
+        script_request(thread->L, &request, &length);
+        thread->L = NULL;
+    }
+
     connection *c = thread->cs;
 
     for (uint64_t i = 0; i < thread->connections; i++, c++) {
         c->thread = thread;
         c->ssl     = cfg.ctx ? SSL_new(cfg.ctx) : NULL;
+        c->request = request;
+        c->length  = length;
         connect_socket(thread, c);
     }
 
@@ -405,10 +410,17 @@ static void socket_connected(aeEventLoop *loop, int fd, void *data, int mask) {
 
 static void socket_writeable(aeEventLoop *loop, int fd, void *data, int mask) {
     connection *c = data;
-    size_t len = req.size - c->written;
+    thread *thread = c->thread;
+
+    if (!c->written && thread->L) {
+        script_request(thread->L, &c->request, &c->length);
+    }
+
+    char  *buf = c->request + c->written;
+    size_t len = c->length  - c->written;
     size_t n;
 
-    switch (sock.write(c, req.buf + c->written, len, &n)) {
+    switch (sock.write(c, buf, len, &n)) {
         case OK:    break;
         case ERROR: goto error;
         case RETRY: return;
@@ -417,7 +429,7 @@ static void socket_writeable(aeEventLoop *loop, int fd, void *data, int mask) {
     if (!c->written) c->start = time_us();
 
     c->written += n;
-    if (c->written == req.size) {
+    if (c->written == c->length) {
         c->written = 0;
         aeDeleteFileEvent(loop, fd, AE_WRITABLE);
     }
@@ -425,8 +437,8 @@ static void socket_writeable(aeEventLoop *loop, int fd, void *data, int mask) {
     return;
 
   error:
-    c->thread->errors.write++;
+    thread->errors.write++;
-    reconnect_socket(c->thread, c);
+    reconnect_socket(thread, c);
 }
 
 
@@ -469,42 +481,12 @@ static char *extract_url_part(char *url, struct http_parser_url *parser_url, enu
     return part;
 }
 
-static char *format_request(char *host, char *port, char *path, char **headers) {
-    char *buf  = NULL;
-    char *head = NULL;
-
-    for (char **h = headers; *h != NULL; h++) {
-        aprintf(&head, "%s\r\n", *h);
-        if (!strncasecmp(*h, "Host:", 5)) {
-            host = NULL;
-            port = NULL;
-        }
-    }
-
-    if (req.body) {
-        size_t len = strlen(req.body);
-        aprintf(&head, "Content-Length: %zd\r\n", len);
-    }
-
-    aprintf(&buf, "%s %s HTTP/1.1\r\n", req.method, path);
-    if (host) aprintf(&buf, "Host: %s", host);
-    if (port) aprintf(&buf, ":%s", port);
-    if (host) aprintf(&buf, "\r\n");
-
-    aprintf(&buf, "%s\r\n", head ? head : "");
-    aprintf(&buf, "%s", req.body ? req.body : "");
-
-    free(head);
-    return buf;
-}
-
 static struct option longopts[] = {
     { "connections", required_argument, NULL, 'c' },
     { "duration",    required_argument, NULL, 'd' },
     { "threads",     required_argument, NULL, 't' },
+    { "script",      required_argument, NULL, 's' },
     { "header",      required_argument, NULL, 'H' },
-    { "method",      required_argument, NULL, 'M' },
-    { "body",        required_argument, NULL, 'B' },
     { "latency",     no_argument,       NULL, 'L' },
     { "timeout",     required_argument, NULL, 'T' },
     { "help",        no_argument,       NULL, 'h' },
@@ -520,9 +502,8 @@ static int parse_args(struct config *cfg, char **url, char **headers, int argc,
     cfg->connections = 10;
     cfg->duration    = 10;
     cfg->timeout     = SOCKET_TIMEOUT_MS;
-    req.method       = "GET";
 
-    while ((c = getopt_long(argc, argv, "t:c:d:H:M:B:T:Lrv?", longopts, NULL)) != -1) {
+    while ((c = getopt_long(argc, argv, "t:c:d:s:H:T:Lrv?", longopts, NULL)) != -1) {
         switch (c) {
             case 't':
                 if (scan_metric(optarg, &cfg->threads)) return -1;
@@ -533,15 +514,12 @@ static int parse_args(struct config *cfg, char **url, char **headers, int argc,
             case 'd':
                 if (scan_time(optarg, &cfg->duration)) return -1;
                 break;
+            case 's':
+                cfg->script = optarg;
+                break;
             case 'H':
                 *header++ = optarg;
                 break;
-            case 'M':
-                req.method = optarg;
-                break;
-            case 'B':
-                req.body = optarg;
-                break;
             case 'L':
                 cfg->latency = true;
                 break;

src/wrk.h

@@ -11,6 +11,7 @@
 
 #include "stats.h"
 #include "ae.h"
+#include "script.h"
 #include "http_parser.h"
 
 #define VERSION  "2.2.2"
@@ -43,6 +44,7 @@ typedef struct {
     uint64_t missed;
     stats *latency;
     tinymt64_t rand;
+    lua_State *L;
     errors errors;
     struct connection *cs;
 } thread;
@@ -53,6 +55,8 @@ typedef struct connection {
     int fd;
     SSL *ssl;
     uint64_t start;
+    char *request;
+    size_t length;
     size_t written;
     char buf[RECVBUF];
 } connection;