"Fossies" - the Fresh Open Source Software Archive

Member "wrk-4.2.0/src/script.c" (7 Feb 2021, 16607 Bytes) of package /linux/www/wrk-4.2.0.tar.gz:


As a special service "Fossies" has tried to format the requested source page into HTML format using (guessed) C and C++ source code syntax highlighting (style: standard) with prefixed line numbers and code folding option. Alternatively you can here view or download the uninterpreted source code file. For more information about "script.c" see the Fossies "Dox" file reference documentation and the last Fossies "Diffs" side-by-side code changes report: 4.0.2_vs_4.1.0.

    1 // Copyright (C) 2013 - Will Glozer.  All rights reserved.
    2 
    3 #include <stdlib.h>
    4 #include <string.h>
    5 #include "script.h"
    6 #include "http_parser.h"
    7 #include "zmalloc.h"
    8 
    9 typedef struct {
   10     char *name;
   11     int   type;
   12     void *value;
   13 } table_field;
   14 
   15 static int script_addr_tostring(lua_State *);
   16 static int script_addr_gc(lua_State *);
   17 static int script_stats_call(lua_State *);
   18 static int script_stats_len(lua_State *);
   19 static int script_stats_index(lua_State *);
   20 static int script_thread_index(lua_State *);
   21 static int script_thread_newindex(lua_State *);
   22 static int script_wrk_lookup(lua_State *);
   23 static int script_wrk_connect(lua_State *);
   24 
   25 static void set_fields(lua_State *, int, const table_field *);
   26 static void set_field(lua_State *, int, char *, int);
   27 static int push_url_part(lua_State *, char *, struct http_parser_url *, enum http_parser_url_fields);
   28 
   29 static const struct luaL_Reg addrlib[] = {
   30     { "__tostring", script_addr_tostring   },
   31     { "__gc"    ,   script_addr_gc         },
   32     { NULL,         NULL                   }
   33 };
   34 
   35 static const struct luaL_Reg statslib[] = {
   36     { "__call",     script_stats_call      },
   37     { "__index",    script_stats_index     },
   38     { "__len",      script_stats_len       },
   39     { NULL,         NULL                   }
   40 };
   41 
   42 static const struct luaL_Reg threadlib[] = {
   43     { "__index",    script_thread_index    },
   44     { "__newindex", script_thread_newindex },
   45     { NULL,         NULL                   }
   46 };
   47 
   48 lua_State *script_create(char *file, char *url, char **headers) {
   49     lua_State *L = luaL_newstate();
   50     luaL_openlibs(L);
   51     (void) luaL_dostring(L, "wrk = require \"wrk\"");
   52 
   53     luaL_newmetatable(L, "wrk.addr");
   54     luaL_register(L, NULL, addrlib);
   55     luaL_newmetatable(L, "wrk.stats");
   56     luaL_register(L, NULL, statslib);
   57     luaL_newmetatable(L, "wrk.thread");
   58     luaL_register(L, NULL, threadlib);
   59 
   60     struct http_parser_url parts = {};
   61     script_parse_url(url, &parts);
   62     char *path = "/";
   63 
   64     if (parts.field_set & (1 << UF_PATH)) {
   65         path = &url[parts.field_data[UF_PATH].off];
   66     }
   67 
   68     const table_field fields[] = {
   69         { "lookup",  LUA_TFUNCTION, script_wrk_lookup  },
   70         { "connect", LUA_TFUNCTION, script_wrk_connect },
   71         { "path",    LUA_TSTRING,   path               },
   72         { NULL,      0,             NULL               },
   73     };
   74 
   75     lua_getglobal(L, "wrk");
   76 
   77     set_field(L, 4, "scheme", push_url_part(L, url, &parts, UF_SCHEMA));
   78     set_field(L, 4, "host",   push_url_part(L, url, &parts, UF_HOST));
   79     set_field(L, 4, "port",   push_url_part(L, url, &parts, UF_PORT));
   80     set_fields(L, 4, fields);
   81 
   82     lua_getfield(L, 4, "headers");
   83     for (char **h = headers; *h; h++) {
   84         char *p = strchr(*h, ':');
   85         if (p && p[1] == ' ') {
   86             lua_pushlstring(L, *h, p - *h);
   87             lua_pushstring(L, p + 2);
   88             lua_settable(L, 5);
   89         }
   90     }
   91     lua_pop(L, 5);
   92 
   93     if (file && luaL_dofile(L, file)) {
   94         const char *cause = lua_tostring(L, -1);
   95         fprintf(stderr, "%s: %s\n", file, cause);
   96     }
   97 
   98     return L;
   99 }
  100 
  101 bool script_resolve(lua_State *L, char *host, char *service) {
  102     lua_getglobal(L, "wrk");
  103 
  104     lua_getfield(L, -1, "resolve");
  105     lua_pushstring(L, host);
  106     lua_pushstring(L, service);
  107     lua_call(L, 2, 0);
  108 
  109     lua_getfield(L, -1, "addrs");
  110     size_t count = lua_objlen(L, -1);
  111     lua_pop(L, 2);
  112     return count > 0;
  113 }
  114 
  115 void script_push_thread(lua_State *L, thread *t) {
  116     thread **ptr = (thread **) lua_newuserdata(L, sizeof(thread **));
  117     *ptr = t;
  118     luaL_getmetatable(L, "wrk.thread");
  119     lua_setmetatable(L, -2);
  120 }
  121 
  122 void script_init(lua_State *L, thread *t, int argc, char **argv) {
  123     lua_getglobal(t->L, "wrk");
  124 
  125     script_push_thread(t->L, t);
  126     lua_setfield(t->L, -2, "thread");
  127 
  128     lua_getglobal(L, "wrk");
  129     lua_getfield(L, -1, "setup");
  130     script_push_thread(L, t);
  131     lua_call(L, 1, 0);
  132     lua_pop(L, 1);
  133 
  134     lua_getfield(t->L, -1, "init");
  135     lua_newtable(t->L);
  136     for (int i = 0; i < argc; i++) {
  137         lua_pushstring(t->L, argv[i]);
  138         lua_rawseti(t->L, -2, i);
  139     }
  140     lua_call(t->L, 1, 0);
  141     lua_pop(t->L, 1);
  142 }
  143 
  144 uint64_t script_delay(lua_State *L) {
  145     lua_getglobal(L, "delay");
  146     lua_call(L, 0, 1);
  147     uint64_t delay = lua_tonumber(L, -1);
  148     lua_pop(L, 1);
  149     return delay;
  150 }
  151 
  152 void script_request(lua_State *L, char **buf, size_t *len) {
  153     int pop = 1;
  154     lua_getglobal(L, "request");
  155     if (!lua_isfunction(L, -1)) {
  156         lua_getglobal(L, "wrk");
  157         lua_getfield(L, -1, "request");
  158         pop += 2;
  159     }
  160     lua_call(L, 0, 1);
  161     const char *str = lua_tolstring(L, -1, len);
  162     *buf = realloc(*buf, *len);
  163     memcpy(*buf, str, *len);
  164     lua_pop(L, pop);
  165 }
  166 
  167 void script_response(lua_State *L, int status, buffer *headers, buffer *body) {
  168     lua_getglobal(L, "response");
  169     lua_pushinteger(L, status);
  170     lua_newtable(L);
  171 
  172     for (char *c = headers->buffer; c < headers->cursor; ) {
  173         c = buffer_pushlstring(L, c);
  174         c = buffer_pushlstring(L, c);
  175         lua_rawset(L, -3);
  176     }
  177 
  178     lua_pushlstring(L, body->buffer, body->cursor - body->buffer);
  179     lua_call(L, 3, 0);
  180 
  181     buffer_reset(headers);
  182     buffer_reset(body);
  183 }
  184 
  185 bool script_is_function(lua_State *L, char *name) {
  186     lua_getglobal(L, name);
  187     bool is_function = lua_isfunction(L, -1);
  188     lua_pop(L, 1);
  189     return is_function;
  190 }
  191 
  192 bool script_is_static(lua_State *L) {
  193     return !script_is_function(L, "request");
  194 }
  195 
  196 bool script_want_response(lua_State *L) {
  197     return script_is_function(L, "response");
  198 }
  199 
  200 bool script_has_delay(lua_State *L) {
  201     return script_is_function(L, "delay");
  202 }
  203 
  204 bool script_has_done(lua_State *L) {
  205     return script_is_function(L, "done");
  206 }
  207 
  208 void script_header_done(lua_State *L, luaL_Buffer *buffer) {
  209     luaL_pushresult(buffer);
  210 }
  211 
  212 void script_summary(lua_State *L, uint64_t duration, uint64_t requests, uint64_t bytes) {
  213     const table_field fields[] = {
  214         { "duration", LUA_TNUMBER, &duration },
  215         { "requests", LUA_TNUMBER, &requests },
  216         { "bytes",    LUA_TNUMBER, &bytes    },
  217         { NULL,       0,           NULL      },
  218     };
  219     lua_newtable(L);
  220     set_fields(L, 1, fields);
  221 }
  222 
  223 void script_errors(lua_State *L, errors *errors) {
  224     uint64_t e[] = {
  225         errors->connect,
  226         errors->read,
  227         errors->write,
  228         errors->status,
  229         errors->timeout
  230     };
  231     const table_field fields[] = {
  232         { "connect", LUA_TNUMBER, &e[0] },
  233         { "read",    LUA_TNUMBER, &e[1] },
  234         { "write",   LUA_TNUMBER, &e[2] },
  235         { "status",  LUA_TNUMBER, &e[3] },
  236         { "timeout", LUA_TNUMBER, &e[4] },
  237         { NULL,      0,           NULL  },
  238     };
  239     lua_newtable(L);
  240     set_fields(L, 2, fields);
  241     lua_setfield(L, 1, "errors");
  242 }
  243 
  244 void script_push_stats(lua_State *L, stats *s) {
  245     stats **ptr = (stats **) lua_newuserdata(L, sizeof(stats **));
  246     *ptr = s;
  247     luaL_getmetatable(L, "wrk.stats");
  248     lua_setmetatable(L, -2);
  249 }
  250 
  251 void script_done(lua_State *L, stats *latency, stats *requests) {
  252     lua_getglobal(L, "done");
  253     lua_pushvalue(L, 1);
  254 
  255     script_push_stats(L, latency);
  256     script_push_stats(L, requests);
  257 
  258     lua_call(L, 3, 0);
  259     lua_pop(L, 1);
  260 }
  261 
  262 static int verify_request(http_parser *parser) {
  263     size_t *count = parser->data;
  264     (*count)++;
  265     return 0;
  266 }
  267 
  268 size_t script_verify_request(lua_State *L) {
  269     http_parser_settings settings = {
  270         .on_message_complete = verify_request
  271     };
  272     http_parser parser;
  273     char *request = NULL;
  274     size_t len, count = 0;
  275 
  276     script_request(L, &request, &len);
  277     http_parser_init(&parser, HTTP_REQUEST);
  278     parser.data = &count;
  279 
  280     size_t parsed = http_parser_execute(&parser, &settings, request, len);
  281 
  282     if (parsed != len || count == 0) {
  283         enum http_errno err = HTTP_PARSER_ERRNO(&parser);
  284         const char *desc = http_errno_description(err);
  285         const char *msg = err != HPE_OK ? desc : "incomplete request";
  286         int line = 1, column = 1;
  287 
  288         for (char *c = request; c < request + parsed; c++) {
  289             column++;
  290             if (*c == '\n') {
  291                 column = 1;
  292                 line++;
  293             }
  294         }
  295 
  296         fprintf(stderr, "%s at %d:%d\n", msg, line, column);
  297         exit(1);
  298     }
  299 
  300     return count;
  301 }
  302 
  303 static struct addrinfo *checkaddr(lua_State *L) {
  304     struct addrinfo *addr = luaL_checkudata(L, -1, "wrk.addr");
  305     luaL_argcheck(L, addr != NULL, 1, "`addr' expected");
  306     return addr;
  307 }
  308 
  309 void script_addr_copy(struct addrinfo *src, struct addrinfo *dst) {
  310     *dst = *src;
  311     dst->ai_addr = zmalloc(src->ai_addrlen);
  312     memcpy(dst->ai_addr, src->ai_addr, src->ai_addrlen);
  313 }
  314 
  315 struct addrinfo *script_addr_clone(lua_State *L, struct addrinfo *addr) {
  316     struct addrinfo *udata = lua_newuserdata(L, sizeof(*udata));
  317     luaL_getmetatable(L, "wrk.addr");
  318     lua_setmetatable(L, -2);
  319     script_addr_copy(addr, udata);
  320     return udata;
  321 }
  322 
  323 static int script_addr_tostring(lua_State *L) {
  324     struct addrinfo *addr = checkaddr(L);
  325     char host[NI_MAXHOST];
  326     char service[NI_MAXSERV];
  327 
  328     int flags = NI_NUMERICHOST | NI_NUMERICSERV;
  329     int rc = getnameinfo(addr->ai_addr, addr->ai_addrlen, host, NI_MAXHOST, service, NI_MAXSERV, flags);
  330     if (rc != 0) {
  331         const char *msg = gai_strerror(rc);
  332         return luaL_error(L, "addr tostring failed %s", msg);
  333     }
  334 
  335     lua_pushfstring(L, "%s:%s", host, service);
  336     return 1;
  337 }
  338 
  339 static int script_addr_gc(lua_State *L) {
  340     struct addrinfo *addr = checkaddr(L);
  341     zfree(addr->ai_addr);
  342     return 0;
  343 }
  344 
  345 static stats *checkstats(lua_State *L) {
  346     stats **s = luaL_checkudata(L, 1, "wrk.stats");
  347     luaL_argcheck(L, s != NULL, 1, "`stats' expected");
  348     return *s;
  349 }
  350 
  351 static int script_stats_percentile(lua_State *L) {
  352     stats *s = checkstats(L);
  353     lua_Number p = luaL_checknumber(L, 2);
  354     lua_pushnumber(L, stats_percentile(s, p));
  355     return 1;
  356 }
  357 
  358 static int script_stats_call(lua_State *L) {
  359     stats *s = checkstats(L);
  360     uint64_t index = lua_tonumber(L, 2);
  361     uint64_t count;
  362     lua_pushnumber(L, stats_value_at(s, index - 1, &count));
  363     lua_pushnumber(L, count);
  364     return 2;
  365 }
  366 
  367 static int script_stats_index(lua_State *L) {
  368     stats *s = checkstats(L);
  369     const char *method = lua_tostring(L, 2);
  370     if (!strcmp("min",   method)) lua_pushnumber(L, s->min);
  371     if (!strcmp("max",   method)) lua_pushnumber(L, s->max);
  372     if (!strcmp("mean",  method)) lua_pushnumber(L, stats_mean(s));
  373     if (!strcmp("stdev", method)) lua_pushnumber(L, stats_stdev(s, stats_mean(s)));
  374     if (!strcmp("percentile", method)) {
  375         lua_pushcfunction(L, script_stats_percentile);
  376     }
  377     return 1;
  378 }
  379 
  380 static int script_stats_len(lua_State *L) {
  381     stats *s = checkstats(L);
  382     lua_pushinteger(L, stats_popcount(s));
  383     return 1;
  384 }
  385 
  386 static thread *checkthread(lua_State *L) {
  387     thread **t = luaL_checkudata(L, 1, "wrk.thread");
  388     luaL_argcheck(L, t != NULL, 1, "`thread' expected");
  389     return *t;
  390 }
  391 
  392 static int script_thread_get(lua_State *L) {
  393     thread *t = checkthread(L);
  394     const char *key = lua_tostring(L, -1);
  395     lua_getglobal(t->L, key);
  396     script_copy_value(t->L, L, -1);
  397     lua_pop(t->L, 1);
  398     return 1;
  399 }
  400 
  401 static int script_thread_set(lua_State *L) {
  402     thread *t = checkthread(L);
  403     const char *name = lua_tostring(L, -2);
  404     script_copy_value(L, t->L, -1);
  405     lua_setglobal(t->L, name);
  406     return 0;
  407 }
  408 
  409 static int script_thread_stop(lua_State *L) {
  410     thread *t = checkthread(L);
  411     aeStop(t->loop);
  412     return 0;
  413 }
  414 
  415 static int script_thread_index(lua_State *L) {
  416     thread *t = checkthread(L);
  417     const char *key = lua_tostring(L, 2);
  418     if (!strcmp("get",  key)) lua_pushcfunction(L, script_thread_get);
  419     if (!strcmp("set",  key)) lua_pushcfunction(L, script_thread_set);
  420     if (!strcmp("stop", key)) lua_pushcfunction(L, script_thread_stop);
  421     if (!strcmp("addr", key)) script_addr_clone(L, t->addr);
  422     return 1;
  423 }
  424 
  425 static int script_thread_newindex(lua_State *L) {
  426     thread *t = checkthread(L);
  427     const char *key = lua_tostring(L, -2);
  428     if (!strcmp("addr", key)) {
  429         struct addrinfo *addr = checkaddr(L);
  430         if (t->addr) zfree(t->addr->ai_addr);
  431         t->addr = zrealloc(t->addr, sizeof(*addr));
  432         script_addr_copy(addr, t->addr);
  433     } else {
  434         luaL_error(L, "cannot set '%s' on thread", luaL_typename(L, -1));
  435     }
  436     return 0;
  437 }
  438 
  439 static int script_wrk_lookup(lua_State *L) {
  440     struct addrinfo *addrs;
  441     struct addrinfo hints = {
  442         .ai_family   = AF_UNSPEC,
  443         .ai_socktype = SOCK_STREAM
  444     };
  445     int rc, index = 1;
  446 
  447     const char *host    = lua_tostring(L, -2);
  448     const char *service = lua_tostring(L, -1);
  449 
  450     if ((rc = getaddrinfo(host, service, &hints, &addrs)) != 0) {
  451         const char *msg = gai_strerror(rc);
  452         fprintf(stderr, "unable to resolve %s:%s %s\n", host, service, msg);
  453         exit(1);
  454     }
  455 
  456     lua_newtable(L);
  457     for (struct addrinfo *addr = addrs; addr != NULL; addr = addr->ai_next) {
  458         script_addr_clone(L, addr);
  459         lua_rawseti(L, -2, index++);
  460     }
  461 
  462     freeaddrinfo(addrs);
  463     return 1;
  464 }
  465 
  466 static int script_wrk_connect(lua_State *L) {
  467     struct addrinfo *addr = checkaddr(L);
  468     int fd, connected = 0;
  469     if ((fd = socket(addr->ai_family, addr->ai_socktype, addr->ai_protocol)) != -1) {
  470         connected = connect(fd, addr->ai_addr, addr->ai_addrlen) == 0;
  471         close(fd);
  472     }
  473     lua_pushboolean(L, connected);
  474     return 1;
  475 }
  476 
  477 void script_copy_value(lua_State *src, lua_State *dst, int index) {
  478     switch (lua_type(src, index)) {
  479         case LUA_TBOOLEAN:
  480             lua_pushboolean(dst, lua_toboolean(src, index));
  481             break;
  482         case LUA_TNIL:
  483             lua_pushnil(dst);
  484             break;
  485         case LUA_TNUMBER:
  486             lua_pushnumber(dst, lua_tonumber(src, index));
  487             break;
  488         case LUA_TSTRING:
  489             lua_pushstring(dst, lua_tostring(src, index));
  490             break;
  491         case LUA_TTABLE:
  492             lua_newtable(dst);
  493             lua_pushnil(src);
  494             while (lua_next(src, index - 1)) {
  495                 script_copy_value(src, dst, -2);
  496                 script_copy_value(src, dst, -1);
  497                 lua_settable(dst, -3);
  498                 lua_pop(src, 1);
  499             }
  500             lua_pop(src, 1);
  501             break;
  502         default:
  503             luaL_error(src, "cannot transfer '%s' to thread", luaL_typename(src, index));
  504     }
  505 }
  506 
  507 int script_parse_url(char *url, struct http_parser_url *parts) {
  508     if (!http_parser_parse_url(url, strlen(url), 0, parts)) {
  509         if (!(parts->field_set & (1 << UF_SCHEMA))) return 0;
  510         if (!(parts->field_set & (1 << UF_HOST)))   return 0;
  511         return 1;
  512     }
  513     return 0;
  514 }
  515 
  516 static int push_url_part(lua_State *L, char *url, struct http_parser_url *parts, enum http_parser_url_fields field) {
  517     int type = parts->field_set & (1 << field) ? LUA_TSTRING : LUA_TNIL;
  518     uint16_t off, len;
  519     switch (type) {
  520         case LUA_TSTRING:
  521             off = parts->field_data[field].off;
  522             len = parts->field_data[field].len;
  523             lua_pushlstring(L, &url[off], len);
  524             break;
  525         case LUA_TNIL:
  526             lua_pushnil(L);
  527     }
  528     return type;
  529 }
  530 
  531 static void set_field(lua_State *L, int index, char *field, int type) {
  532     (void) type;
  533     lua_setfield(L, index, field);
  534 }
  535 
  536 static void set_fields(lua_State *L, int index, const table_field *fields) {
  537     for (int i = 0; fields[i].name; i++) {
  538         table_field f = fields[i];
  539         switch (f.value == NULL ? LUA_TNIL : f.type) {
  540             case LUA_TFUNCTION:
  541                 lua_pushcfunction(L, (lua_CFunction) f.value);
  542                 break;
  543             case LUA_TNUMBER:
  544                 lua_pushinteger(L, *((lua_Integer *) f.value));
  545                 break;
  546             case LUA_TSTRING:
  547                 lua_pushstring(L, (const char *) f.value);
  548                 break;
  549             case LUA_TNIL:
  550                 lua_pushnil(L);
  551                 break;
  552         }
  553         lua_setfield(L, index, f.name);
  554     }
  555 }
  556 
  557 void buffer_append(buffer *b, const char *data, size_t len) {
  558     size_t used = b->cursor - b->buffer;
  559     while (used + len + 1 >= b->length) {
  560         b->length += 1024;
  561         b->buffer  = realloc(b->buffer, b->length);
  562         b->cursor  = b->buffer + used;
  563     }
  564     memcpy(b->cursor, data, len);
  565     b->cursor += len;
  566 }
  567 
  568 void buffer_reset(buffer *b) {
  569     b->cursor = b->buffer;
  570 }
  571 
  572 char *buffer_pushlstring(lua_State *L, char *start) {
  573     char *end = strchr(start, 0);
  574     lua_pushlstring(L, start, end - start);
  575     return end + 1;
  576 }