"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 }