Mercurial
diff third_party/wrk/src/script.c @ 178:94705b5986b3
[ThirdParty] Added WRK and luajit for load testing.
| author | MrJuneJune <me@mrjunejune.com> |
|---|---|
| date | Thu, 22 Jan 2026 20:10:30 -0800 |
| parents | |
| children |
line wrap: on
line diff
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/third_party/wrk/src/script.c Thu Jan 22 20:10:30 2026 -0800 @@ -0,0 +1,576 @@ +// Copyright (C) 2013 - Will Glozer. All rights reserved. + +#include <stdlib.h> +#include <string.h> +#include "script.h" +#include "http_parser.h" +#include "zmalloc.h" + +typedef struct { + char *name; + int type; + void *value; +} table_field; + +static int script_addr_tostring(lua_State *); +static int script_addr_gc(lua_State *); +static int script_stats_call(lua_State *); +static int script_stats_len(lua_State *); +static int script_stats_index(lua_State *); +static int script_thread_index(lua_State *); +static int script_thread_newindex(lua_State *); +static int script_wrk_lookup(lua_State *); +static int script_wrk_connect(lua_State *); + +static void set_fields(lua_State *, int, const table_field *); +static void set_field(lua_State *, int, char *, int); +static int push_url_part(lua_State *, char *, struct http_parser_url *, enum http_parser_url_fields); + +static const struct luaL_Reg addrlib[] = { + { "__tostring", script_addr_tostring }, + { "__gc" , script_addr_gc }, + { NULL, NULL } +}; + +static const struct luaL_Reg statslib[] = { + { "__call", script_stats_call }, + { "__index", script_stats_index }, + { "__len", script_stats_len }, + { NULL, NULL } +}; + +static const struct luaL_Reg threadlib[] = { + { "__index", script_thread_index }, + { "__newindex", script_thread_newindex }, + { NULL, NULL } +}; + +lua_State *script_create(char *file, char *url, char **headers) { + lua_State *L = luaL_newstate(); + luaL_openlibs(L); + (void) luaL_dostring(L, "wrk = require \"wrk\""); + + luaL_newmetatable(L, "wrk.addr"); + luaL_register(L, NULL, addrlib); + luaL_newmetatable(L, "wrk.stats"); + luaL_register(L, NULL, statslib); + luaL_newmetatable(L, "wrk.thread"); + luaL_register(L, NULL, threadlib); + + struct http_parser_url parts = {}; + script_parse_url(url, &parts); + char *path = "/"; + + if (parts.field_set & (1 << UF_PATH)) { + path = &url[parts.field_data[UF_PATH].off]; + } + + const table_field fields[] = { + { "lookup", LUA_TFUNCTION, script_wrk_lookup }, + { "connect", LUA_TFUNCTION, script_wrk_connect }, + { "path", LUA_TSTRING, path }, + { NULL, 0, NULL }, + }; + + lua_getglobal(L, "wrk"); + + set_field(L, 4, "scheme", push_url_part(L, url, &parts, UF_SCHEMA)); + set_field(L, 4, "host", push_url_part(L, url, &parts, UF_HOST)); + set_field(L, 4, "port", push_url_part(L, url, &parts, UF_PORT)); + set_fields(L, 4, fields); + + lua_getfield(L, 4, "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, 5); + } + } + lua_pop(L, 5); + + if (file && luaL_dofile(L, file)) { + const char *cause = lua_tostring(L, -1); + fprintf(stderr, "%s: %s\n", file, cause); + } + + return L; +} + +bool script_resolve(lua_State *L, char *host, char *service) { + lua_getglobal(L, "wrk"); + + lua_getfield(L, -1, "resolve"); + lua_pushstring(L, host); + lua_pushstring(L, service); + lua_call(L, 2, 0); + + lua_getfield(L, -1, "addrs"); + size_t count = lua_objlen(L, -1); + lua_pop(L, 2); + return count > 0; +} + +void script_push_thread(lua_State *L, thread *t) { + thread **ptr = (thread **) lua_newuserdata(L, sizeof(thread **)); + *ptr = t; + luaL_getmetatable(L, "wrk.thread"); + lua_setmetatable(L, -2); +} + +void script_init(lua_State *L, thread *t, int argc, char **argv) { + lua_getglobal(t->L, "wrk"); + + script_push_thread(t->L, t); + lua_setfield(t->L, -2, "thread"); + + lua_getglobal(L, "wrk"); + lua_getfield(L, -1, "setup"); + script_push_thread(L, t); + lua_call(L, 1, 0); + lua_pop(L, 1); + + lua_getfield(t->L, -1, "init"); + lua_newtable(t->L); + for (int i = 0; i < argc; i++) { + lua_pushstring(t->L, argv[i]); + lua_rawseti(t->L, -2, i); + } + lua_call(t->L, 1, 0); + lua_pop(t->L, 1); +} + +uint64_t script_delay(lua_State *L) { + lua_getglobal(L, "delay"); + lua_call(L, 0, 1); + uint64_t delay = lua_tonumber(L, -1); + lua_pop(L, 1); + return delay; +} + +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); + *buf = realloc(*buf, *len); + memcpy(*buf, str, *len); + lua_pop(L, pop); +} + +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_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) { + return !script_is_function(L, "request"); +} + +bool script_want_response(lua_State *L) { + return script_is_function(L, "response"); +} + +bool script_has_delay(lua_State *L) { + return script_is_function(L, "delay"); +} + +bool script_has_done(lua_State *L) { + return script_is_function(L, "done"); +} + +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) { + 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_push_stats(lua_State *L, stats *s) { + stats **ptr = (stats **) lua_newuserdata(L, sizeof(stats **)); + *ptr = s; + luaL_getmetatable(L, "wrk.stats"); + lua_setmetatable(L, -2); +} + +void script_done(lua_State *L, stats *latency, stats *requests) { + lua_getglobal(L, "done"); + lua_pushvalue(L, 1); + + script_push_stats(L, latency); + script_push_stats(L, requests); + + lua_call(L, 3, 0); + lua_pop(L, 1); +} + +static int verify_request(http_parser *parser) { + size_t *count = parser->data; + (*count)++; + return 0; +} + +size_t script_verify_request(lua_State *L) { + http_parser_settings settings = { + .on_message_complete = verify_request + }; + http_parser parser; + char *request = NULL; + size_t len, count = 0; + + script_request(L, &request, &len); + http_parser_init(&parser, HTTP_REQUEST); + parser.data = &count; + + size_t parsed = http_parser_execute(&parser, &settings, request, len); + + if (parsed != len || count == 0) { + enum http_errno err = HTTP_PARSER_ERRNO(&parser); + const char *desc = http_errno_description(err); + const char *msg = err != HPE_OK ? desc : "incomplete request"; + int line = 1, column = 1; + + for (char *c = request; c < request + parsed; c++) { + column++; + if (*c == '\n') { + column = 1; + line++; + } + } + + fprintf(stderr, "%s at %d:%d\n", msg, line, column); + exit(1); + } + + return count; +} + +static struct addrinfo *checkaddr(lua_State *L) { + struct addrinfo *addr = luaL_checkudata(L, -1, "wrk.addr"); + luaL_argcheck(L, addr != NULL, 1, "`addr' expected"); + return addr; +} + +void script_addr_copy(struct addrinfo *src, struct addrinfo *dst) { + *dst = *src; + dst->ai_addr = zmalloc(src->ai_addrlen); + memcpy(dst->ai_addr, src->ai_addr, src->ai_addrlen); +} + +struct addrinfo *script_addr_clone(lua_State *L, struct addrinfo *addr) { + struct addrinfo *udata = lua_newuserdata(L, sizeof(*udata)); + luaL_getmetatable(L, "wrk.addr"); + lua_setmetatable(L, -2); + script_addr_copy(addr, udata); + return udata; +} + +static int script_addr_tostring(lua_State *L) { + struct addrinfo *addr = checkaddr(L); + char host[NI_MAXHOST]; + char service[NI_MAXSERV]; + + int flags = NI_NUMERICHOST | NI_NUMERICSERV; + int rc = getnameinfo(addr->ai_addr, addr->ai_addrlen, host, NI_MAXHOST, service, NI_MAXSERV, flags); + if (rc != 0) { + const char *msg = gai_strerror(rc); + return luaL_error(L, "addr tostring failed %s", msg); + } + + lua_pushfstring(L, "%s:%s", host, service); + return 1; +} + +static int script_addr_gc(lua_State *L) { + struct addrinfo *addr = checkaddr(L); + zfree(addr->ai_addr); + return 0; +} + +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_call(lua_State *L) { + stats *s = checkstats(L); + uint64_t index = lua_tonumber(L, 2); + uint64_t count; + lua_pushnumber(L, stats_value_at(s, index - 1, &count)); + lua_pushnumber(L, count); + return 2; +} + +static int script_stats_index(lua_State *L) { + stats *s = checkstats(L); + 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, stats_popcount(s)); + return 1; +} + +static thread *checkthread(lua_State *L) { + thread **t = luaL_checkudata(L, 1, "wrk.thread"); + luaL_argcheck(L, t != NULL, 1, "`thread' expected"); + return *t; +} + +static int script_thread_get(lua_State *L) { + thread *t = checkthread(L); + const char *key = lua_tostring(L, -1); + lua_getglobal(t->L, key); + script_copy_value(t->L, L, -1); + lua_pop(t->L, 1); + return 1; +} + +static int script_thread_set(lua_State *L) { + thread *t = checkthread(L); + const char *name = lua_tostring(L, -2); + script_copy_value(L, t->L, -1); + lua_setglobal(t->L, name); + return 0; +} + +static int script_thread_stop(lua_State *L) { + thread *t = checkthread(L); + aeStop(t->loop); + return 0; +} + +static int script_thread_index(lua_State *L) { + thread *t = checkthread(L); + const char *key = lua_tostring(L, 2); + if (!strcmp("get", key)) lua_pushcfunction(L, script_thread_get); + if (!strcmp("set", key)) lua_pushcfunction(L, script_thread_set); + if (!strcmp("stop", key)) lua_pushcfunction(L, script_thread_stop); + if (!strcmp("addr", key)) script_addr_clone(L, t->addr); + return 1; +} + +static int script_thread_newindex(lua_State *L) { + thread *t = checkthread(L); + const char *key = lua_tostring(L, -2); + if (!strcmp("addr", key)) { + struct addrinfo *addr = checkaddr(L); + if (t->addr) zfree(t->addr->ai_addr); + t->addr = zrealloc(t->addr, sizeof(*addr)); + script_addr_copy(addr, t->addr); + } else { + luaL_error(L, "cannot set '%s' on thread", luaL_typename(L, -1)); + } + return 0; +} + +static int script_wrk_lookup(lua_State *L) { + struct addrinfo *addrs; + struct addrinfo hints = { + .ai_family = AF_UNSPEC, + .ai_socktype = SOCK_STREAM + }; + int rc, index = 1; + + const char *host = lua_tostring(L, -2); + const char *service = lua_tostring(L, -1); + + if ((rc = getaddrinfo(host, service, &hints, &addrs)) != 0) { + const char *msg = gai_strerror(rc); + fprintf(stderr, "unable to resolve %s:%s %s\n", host, service, msg); + exit(1); + } + + lua_newtable(L); + for (struct addrinfo *addr = addrs; addr != NULL; addr = addr->ai_next) { + script_addr_clone(L, addr); + lua_rawseti(L, -2, index++); + } + + freeaddrinfo(addrs); + return 1; +} + +static int script_wrk_connect(lua_State *L) { + struct addrinfo *addr = checkaddr(L); + int fd, connected = 0; + if ((fd = socket(addr->ai_family, addr->ai_socktype, addr->ai_protocol)) != -1) { + connected = connect(fd, addr->ai_addr, addr->ai_addrlen) == 0; + close(fd); + } + lua_pushboolean(L, connected); + return 1; +} + +void script_copy_value(lua_State *src, lua_State *dst, int index) { + switch (lua_type(src, index)) { + case LUA_TBOOLEAN: + lua_pushboolean(dst, lua_toboolean(src, index)); + break; + case LUA_TNIL: + lua_pushnil(dst); + break; + case LUA_TNUMBER: + lua_pushnumber(dst, lua_tonumber(src, index)); + break; + case LUA_TSTRING: + lua_pushstring(dst, lua_tostring(src, index)); + break; + case LUA_TTABLE: + lua_newtable(dst); + lua_pushnil(src); + while (lua_next(src, index - 1)) { + script_copy_value(src, dst, -2); + script_copy_value(src, dst, -1); + lua_settable(dst, -3); + lua_pop(src, 1); + } + lua_pop(src, 1); + break; + default: + luaL_error(src, "cannot transfer '%s' to thread", luaL_typename(src, index)); + } +} + +int script_parse_url(char *url, struct http_parser_url *parts) { + if (!http_parser_parse_url(url, strlen(url), 0, parts)) { + if (!(parts->field_set & (1 << UF_SCHEMA))) return 0; + if (!(parts->field_set & (1 << UF_HOST))) return 0; + return 1; + } + return 0; +} + +static int push_url_part(lua_State *L, char *url, struct http_parser_url *parts, enum http_parser_url_fields field) { + int type = parts->field_set & (1 << field) ? LUA_TSTRING : LUA_TNIL; + uint16_t off, len; + switch (type) { + case LUA_TSTRING: + off = parts->field_data[field].off; + len = parts->field_data[field].len; + lua_pushlstring(L, &url[off], len); + break; + case LUA_TNIL: + lua_pushnil(L); + } + return type; +} + +static void set_field(lua_State *L, int index, char *field, int type) { + (void) type; + lua_setfield(L, index, field); +} + +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_TFUNCTION: + lua_pushcfunction(L, (lua_CFunction) f.value); + break; + 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); + } +} + +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; +}