test_workqueue.c 9.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416
  1. /* Copyright (c) 2001-2004, Roger Dingledine.
  2. * Copyright (c) 2004-2006, Roger Dingledine, Nick Mathewson.
  3. * Copyright (c) 2007-2015, The Tor Project, Inc. */
  4. /* See LICENSE for licensing information */
  5. #include "or.h"
  6. #include "compat_threads.h"
  7. #include "onion.h"
  8. #include "workqueue.h"
  9. #include "crypto.h"
  10. #include "crypto_curve25519.h"
  11. #include "compat_libevent.h"
  12. #include <stdio.h>
  13. #ifdef HAVE_EVENT2_EVENT_H
  14. #include <event2/event.h>
  15. #else
  16. #include <event.h>
  17. #endif
  18. #define MAX_INFLIGHT (1<<16)
  19. static int opt_verbose = 0;
  20. static int opt_n_threads = 8;
  21. static int opt_n_items = 10000;
  22. static int opt_n_inflight = 1000;
  23. static int opt_n_lowwater = 250;
  24. static int opt_n_cancel = 0;
  25. static int opt_ratio_rsa = 5;
  26. #ifdef TRACK_RESPONSES
  27. tor_mutex_t bitmap_mutex;
  28. int handled_len;
  29. bitarray_t *handled;
  30. #endif
  31. typedef struct state_s {
  32. int magic;
  33. int n_handled;
  34. crypto_pk_t *rsa;
  35. curve25519_secret_key_t ecdh;
  36. int is_shutdown;
  37. } state_t;
  38. typedef struct rsa_work_s {
  39. int serial;
  40. uint8_t msg[128];
  41. uint8_t msglen;
  42. } rsa_work_t;
  43. typedef struct ecdh_work_s {
  44. int serial;
  45. union {
  46. curve25519_public_key_t pk;
  47. uint8_t msg[32];
  48. } u;
  49. } ecdh_work_t;
  50. static void
  51. mark_handled(int serial)
  52. {
  53. #ifdef TRACK_RESPONSES
  54. tor_mutex_acquire(&bitmap_mutex);
  55. tor_assert(serial < handled_len);
  56. tor_assert(! bitarray_is_set(handled, serial));
  57. bitarray_set(handled, serial);
  58. tor_mutex_release(&bitmap_mutex);
  59. #else
  60. (void)serial;
  61. #endif
  62. }
  63. static int
  64. workqueue_do_rsa(void *state, void *work)
  65. {
  66. rsa_work_t *rw = work;
  67. state_t *st = state;
  68. crypto_pk_t *rsa = st->rsa;
  69. uint8_t sig[256];
  70. int len;
  71. tor_assert(st->magic == 13371337);
  72. len = crypto_pk_private_sign(rsa, (char*)sig, 256,
  73. (char*)rw->msg, rw->msglen);
  74. if (len < 0) {
  75. rw->msglen = 0;
  76. return WQ_RPL_ERROR;
  77. }
  78. memset(rw->msg, 0, sizeof(rw->msg));
  79. rw->msglen = len;
  80. memcpy(rw->msg, sig, len);
  81. ++st->n_handled;
  82. mark_handled(rw->serial);
  83. return WQ_RPL_REPLY;
  84. }
  85. static int
  86. workqueue_do_shutdown(void *state, void *work)
  87. {
  88. (void)state;
  89. (void)work;
  90. crypto_pk_free(((state_t*)state)->rsa);
  91. tor_free(state);
  92. return WQ_RPL_SHUTDOWN;
  93. }
  94. static int
  95. workqueue_do_ecdh(void *state, void *work)
  96. {
  97. ecdh_work_t *ew = work;
  98. uint8_t output[CURVE25519_OUTPUT_LEN];
  99. state_t *st = state;
  100. tor_assert(st->magic == 13371337);
  101. curve25519_handshake(output, &st->ecdh, &ew->u.pk);
  102. memcpy(ew->u.msg, output, CURVE25519_OUTPUT_LEN);
  103. ++st->n_handled;
  104. mark_handled(ew->serial);
  105. return WQ_RPL_REPLY;
  106. }
  107. static void *
  108. new_state(void *arg)
  109. {
  110. state_t *st;
  111. (void)arg;
  112. st = tor_malloc(sizeof(*st));
  113. /* Every thread gets its own keys. not a problem for benchmarking */
  114. st->rsa = crypto_pk_new();
  115. if (crypto_pk_generate_key_with_bits(st->rsa, 1024) < 0) {
  116. crypto_pk_free(st->rsa);
  117. tor_free(st);
  118. return NULL;
  119. }
  120. curve25519_secret_key_generate(&st->ecdh, 0);
  121. st->magic = 13371337;
  122. return st;
  123. }
  124. static void
  125. free_state(void *arg)
  126. {
  127. state_t *st = arg;
  128. crypto_pk_free(st->rsa);
  129. tor_free(st);
  130. }
  131. static tor_weak_rng_t weak_rng;
  132. static int n_sent = 0;
  133. static int rsa_sent = 0;
  134. static int ecdh_sent = 0;
  135. static int n_received = 0;
  136. #ifdef TRACK_RESPONSES
  137. bitarray_t *received;
  138. #endif
  139. static void
  140. handle_reply(void *arg)
  141. {
  142. #ifdef TRACK_RESPONSES
  143. rsa_work_t *rw = arg; /* Naughty cast, but only looking at serial. */
  144. tor_assert(! bitarray_is_set(received, rw->serial));
  145. bitarray_set(received,rw->serial);
  146. #endif
  147. tor_free(arg);
  148. ++n_received;
  149. }
  150. static workqueue_entry_t *
  151. add_work(threadpool_t *tp)
  152. {
  153. int add_rsa =
  154. opt_ratio_rsa == 0 ||
  155. tor_weak_random_range(&weak_rng, opt_ratio_rsa) == 0;
  156. if (add_rsa) {
  157. rsa_work_t *w = tor_malloc_zero(sizeof(*w));
  158. w->serial = n_sent++;
  159. crypto_rand((char*)w->msg, 20);
  160. w->msglen = 20;
  161. ++rsa_sent;
  162. return threadpool_queue_work(tp, workqueue_do_rsa, handle_reply, w);
  163. } else {
  164. ecdh_work_t *w = tor_malloc_zero(sizeof(*w));
  165. w->serial = n_sent++;
  166. /* Not strictly right, but this is just for benchmarks. */
  167. crypto_rand((char*)w->u.pk.public_key, 32);
  168. ++ecdh_sent;
  169. return threadpool_queue_work(tp, workqueue_do_ecdh, handle_reply, w);
  170. }
  171. }
  172. static int n_failed_cancel = 0;
  173. static int n_successful_cancel = 0;
  174. static int
  175. add_n_work_items(threadpool_t *tp, int n)
  176. {
  177. int n_queued = 0;
  178. int n_try_cancel = 0, i;
  179. workqueue_entry_t **to_cancel;
  180. workqueue_entry_t *ent;
  181. to_cancel = tor_malloc(sizeof(workqueue_entry_t*) * opt_n_cancel);
  182. while (n_queued++ < n) {
  183. ent = add_work(tp);
  184. if (! ent) {
  185. puts("Z");
  186. tor_event_base_loopexit(tor_libevent_get_base(), NULL);
  187. return -1;
  188. }
  189. if (n_try_cancel < opt_n_cancel &&
  190. tor_weak_random_range(&weak_rng, n) < opt_n_cancel) {
  191. to_cancel[n_try_cancel++] = ent;
  192. }
  193. }
  194. for (i = 0; i < n_try_cancel; ++i) {
  195. void *work = workqueue_entry_cancel(to_cancel[i]);
  196. if (! work) {
  197. n_failed_cancel++;
  198. } else {
  199. n_successful_cancel++;
  200. tor_free(work);
  201. }
  202. }
  203. tor_free(to_cancel);
  204. return 0;
  205. }
  206. static int shutting_down = 0;
  207. static void
  208. replysock_readable_cb(tor_socket_t sock, short what, void *arg)
  209. {
  210. threadpool_t *tp = arg;
  211. replyqueue_t *rq = threadpool_get_replyqueue(tp);
  212. int old_r = n_received;
  213. (void) sock;
  214. (void) what;
  215. replyqueue_process(rq);
  216. if (old_r == n_received)
  217. return;
  218. if (opt_verbose) {
  219. printf("%d / %d", n_received, n_sent);
  220. if (opt_n_cancel)
  221. printf(" (%d cancelled, %d uncancellable)",
  222. n_successful_cancel, n_failed_cancel);
  223. puts("");
  224. }
  225. #ifdef TRACK_RESPONSES
  226. tor_mutex_acquire(&bitmap_mutex);
  227. for (i = 0; i < opt_n_items; ++i) {
  228. if (bitarray_is_set(received, i))
  229. putc('o', stdout);
  230. else if (bitarray_is_set(handled, i))
  231. putc('!', stdout);
  232. else
  233. putc('.', stdout);
  234. }
  235. puts("");
  236. tor_mutex_release(&bitmap_mutex);
  237. #endif
  238. if (n_sent - (n_received+n_successful_cancel) < opt_n_lowwater) {
  239. int n_to_send = n_received + opt_n_inflight - n_sent;
  240. if (n_to_send > opt_n_items - n_sent)
  241. n_to_send = opt_n_items - n_sent;
  242. add_n_work_items(tp, n_to_send);
  243. }
  244. if (shutting_down == 0 &&
  245. n_received+n_successful_cancel == n_sent &&
  246. n_sent >= opt_n_items) {
  247. shutting_down = 1;
  248. threadpool_queue_update(tp, NULL,
  249. workqueue_do_shutdown, NULL, NULL);
  250. {
  251. struct timeval limit = { 2, 0 };
  252. tor_event_base_loopexit(tor_libevent_get_base(), &limit);
  253. }
  254. }
  255. }
  256. static void
  257. help(void)
  258. {
  259. puts(
  260. "Options:\n"
  261. " -N <items> Run this many items of work\n"
  262. " -T <threads> Use this many threads\n"
  263. " -I <inflight> Have no more than this many requests queued at once\n"
  264. " -L <lowwater> Add items whenever fewer than this many are pending\n"
  265. " -C <cancel> Try to cancel N items of every batch that we add\n"
  266. " -R <ratio> Make one out of this many items be a slow (RSA) one\n"
  267. " --no-{eventfd2,eventfd,pipe2,pipe,socketpair}\n"
  268. " Disable one of the alert_socket backends.");
  269. }
  270. int
  271. main(int argc, char **argv)
  272. {
  273. replyqueue_t *rq;
  274. threadpool_t *tp;
  275. int i;
  276. tor_libevent_cfg evcfg;
  277. struct event *ev;
  278. uint32_t as_flags = 0;
  279. for (i = 1; i < argc; ++i) {
  280. if (!strcmp(argv[i], "-v")) {
  281. opt_verbose = 1;
  282. } else if (!strcmp(argv[i], "-T") && i+1<argc) {
  283. opt_n_threads = atoi(argv[++i]);
  284. } else if (!strcmp(argv[i], "-N") && i+1<argc) {
  285. opt_n_items = atoi(argv[++i]);
  286. } else if (!strcmp(argv[i], "-I") && i+1<argc) {
  287. opt_n_inflight = atoi(argv[++i]);
  288. } else if (!strcmp(argv[i], "-L") && i+1<argc) {
  289. opt_n_lowwater = atoi(argv[++i]);
  290. } else if (!strcmp(argv[i], "-R") && i+1<argc) {
  291. opt_ratio_rsa = atoi(argv[++i]);
  292. } else if (!strcmp(argv[i], "-C") && i+1<argc) {
  293. opt_n_cancel = atoi(argv[++i]);
  294. } else if (!strcmp(argv[i], "--no-eventfd2")) {
  295. as_flags |= ASOCKS_NOEVENTFD2;
  296. } else if (!strcmp(argv[i], "--no-eventfd")) {
  297. as_flags |= ASOCKS_NOEVENTFD;
  298. } else if (!strcmp(argv[i], "--no-pipe2")) {
  299. as_flags |= ASOCKS_NOPIPE2;
  300. } else if (!strcmp(argv[i], "--no-pipe")) {
  301. as_flags |= ASOCKS_NOPIPE;
  302. } else if (!strcmp(argv[i], "--no-socketpair")) {
  303. as_flags |= ASOCKS_NOSOCKETPAIR;
  304. } else if (!strcmp(argv[i], "-h")) {
  305. help();
  306. return 0;
  307. } else {
  308. help();
  309. return 1;
  310. }
  311. }
  312. if (opt_n_threads < 1 ||
  313. opt_n_items < 1 || opt_n_inflight < 1 || opt_n_lowwater < 0 ||
  314. opt_n_cancel > opt_n_inflight || opt_n_inflight > MAX_INFLIGHT ||
  315. opt_ratio_rsa < 0) {
  316. help();
  317. return 1;
  318. }
  319. init_logging(1);
  320. crypto_global_init(1, NULL, NULL);
  321. crypto_seed_rng();
  322. rq = replyqueue_new(as_flags);
  323. tor_assert(rq);
  324. tp = threadpool_new(opt_n_threads,
  325. rq, new_state, free_state, NULL);
  326. tor_assert(tp);
  327. crypto_seed_weak_rng(&weak_rng);
  328. memset(&evcfg, 0, sizeof(evcfg));
  329. tor_libevent_initialize(&evcfg);
  330. ev = tor_event_new(tor_libevent_get_base(),
  331. replyqueue_get_socket(rq), EV_READ|EV_PERSIST,
  332. replysock_readable_cb, tp);
  333. event_add(ev, NULL);
  334. #ifdef TRACK_RESPONSES
  335. handled = bitarray_init_zero(opt_n_items);
  336. received = bitarray_init_zero(opt_n_items);
  337. tor_mutex_init(&bitmap_mutex);
  338. handled_len = opt_n_items;
  339. #endif
  340. for (i = 0; i < opt_n_inflight; ++i) {
  341. if (! add_work(tp)) {
  342. puts("Couldn't add work.");
  343. return 1;
  344. }
  345. }
  346. {
  347. struct timeval limit = { 180, 0 };
  348. tor_event_base_loopexit(tor_libevent_get_base(), &limit);
  349. }
  350. event_base_loop(tor_libevent_get_base(), 0);
  351. if (n_sent != opt_n_items || n_received+n_successful_cancel != n_sent) {
  352. printf("%d vs %d\n", n_sent, opt_n_items);
  353. printf("%d+%d vs %d\n", n_received, n_successful_cancel, n_sent);
  354. puts("FAIL");
  355. return 1;
  356. } else {
  357. puts("OK");
  358. return 0;
  359. }
  360. }