From 93348e28147b085b1003f88b62e16422593976ab Mon Sep 17 00:00:00 2001
From: Will <will@glozer.net>
Date: Sat, 7 Feb 2015 10:35:21 +0900
Subject: [PATCH] add CHANGES file and simplify script init

---
 CHANGES              |  4 ++++
 README               | 13 +++++++------
 scripts/pipeline.lua |  2 --
 src/script.c         | 40 ++++++++++++++++++++++------------------
 src/wrk.lua          | 17 ++++++++---------
 5 files changed, 41 insertions(+), 35 deletions(-)
 create mode 100644 CHANGES

diff --git a/CHANGES b/CHANGES
new file mode 100644
index 0000000..d2b4de6
--- /dev/null
+++ b/CHANGES
@@ -0,0 +1,4 @@
+master
+
+ * The wrk global variable is the only global defined by default.
+ * wrk.init() calls the global init(), remove calls to wrk.init().
diff --git a/README b/README
index f4c865f..76bda73 100644
--- a/README
+++ b/README
@@ -50,14 +50,15 @@ Scripting
       wrk.format returns a HTTP request string containing the passed
       parameters merged with values from the wrk table.
 
-    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 following globals are optional, and if defined must be functions:
+
+    global init     -- called when the thread is initialized
+    global request  -- returning the HTTP message for each request
+    global response -- called with HTTP response data
+    global done     -- 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 "--"
-  and scripts that override init() but not request() must call wrk.init()
+  script which must be separated from wrk arguments with "--".
 
   The done() function receives a table containing result data, and two
   statistics objects representing the sampled per-request latency and
diff --git a/scripts/pipeline.lua b/scripts/pipeline.lua
index 4bc0285..269bc92 100644
--- a/scripts/pipeline.lua
+++ b/scripts/pipeline.lua
@@ -1,8 +1,6 @@
 -- example script demonstrating HTTP pipelining
 
 init = function(args)
-   wrk.init(args)
-
    local r = {}
    r[1] = wrk.format(nil, "/?foo")
    r[2] = wrk.format(nil, "/?bar")
diff --git a/src/script.c b/src/script.c
index ad032a7..7561e3c 100644
--- a/src/script.c
+++ b/src/script.c
@@ -65,22 +65,30 @@ void script_init(lua_State *L, char *script, int argc, char **argv) {
         fprintf(stderr, "%s: %s\n", script, cause);
     }
 
-    lua_getglobal(L, "init");
+    lua_getglobal(L, "wrk");
+    lua_getfield(L, -1, "init");
     lua_newtable(L);
     for (int i = 0; i < argc; i++) {
         lua_pushstring(L, argv[i]);
-        lua_rawseti(L, 2, i);
+        lua_rawseti(L, -2, i);
     }
     lua_call(L, 1, 0);
+    lua_pop(L, 1);
 }
 
 void script_request(lua_State *L, char **buf, size_t *len) {
+    int pop = 1;
     lua_getglobal(L, "request");
+    if (!lua_isfunction(L, -1)) {
+        lua_getglobal(L, "wrk");
+        lua_getfield(L, -1, "request");
+        pop += 2;
+    }
     lua_call(L, 0, 1);
-    const char *str = lua_tolstring(L, 1, len);
+    const char *str = lua_tolstring(L, -1, len);
     *buf = realloc(*buf, *len);
     memcpy(*buf, str, *len);
-    lua_pop(L, 1);
+    lua_pop(L, pop);
 }
 
 void script_response(lua_State *L, int status, buffer *headers, buffer *body) {
@@ -101,27 +109,23 @@ void script_response(lua_State *L, int status, buffer *headers, buffer *body) {
     buffer_reset(body);
 }
 
+bool script_is_function(lua_State *L, char *name) {
+    lua_getglobal(L, name);
+    bool is_function = lua_isfunction(L, -1);
+    lua_pop(L, 1);
+    return is_function;
+}
+
 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;
+    return !script_is_function(L, "request");
 }
 
 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;
+    return script_is_function(L, "response");
 }
 
 bool script_has_done(lua_State *L) {
-    lua_getglobal(L, "done");
-    bool defined = lua_type(L, 1) == LUA_TFUNCTION;
-    lua_pop(L, 1);
-    return defined;
+    return script_is_function(L, "done");
 }
 
 void script_header_done(lua_State *L, luaL_Buffer *buffer) {
diff --git a/src/wrk.lua b/src/wrk.lua
index 1bfa571..17dc6d7 100644
--- a/src/wrk.lua
+++ b/src/wrk.lua
@@ -42,16 +42,15 @@ function wrk.init(args)
 
       wrk.headers["Host"] = host
    end
-   req = wrk.format()
-end
 
-function wrk.request()
-   return req
-end
+   if type(init) == "function" then
+      init(args)
+   end
 
-init     = wrk.init
-request  = wrk.request
-response = nil
-done     = nil
+   local req = wrk.format()
+   wrk.request = function()
+      return req
+   end
+end
 
 return wrk