Mercurial

comparison third_party/wrk/src/script.c @ 178:94705b5986b3

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
[ThirdParty] Added WRK and luajit for load testing.
author MrJuneJune <me@mrjunejune.com>
date Thu, 22 Jan 2026 20:10:30 -0800
parents
children
comparison
equal deleted inserted replaced
177:24fe8ff94056 178:94705b5986b3
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 }