test_workqueue.c 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453
  1. /* Copyright (c) 2001-2004, Roger Dingledine.
  2. * Copyright (c) 2004-2006, Roger Dingledine, Nick Mathewson.
  3. * Copyright (c) 2007-2016, 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 workqueue_reply_t
  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 workqueue_reply_t
  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 workqueue_reply_t
  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 workqueue_reply_t
  108. workqueue_shutdown_error(void *state, void *work)
  109. {
  110. (void)state;
  111. (void)work;
  112. return WQ_RPL_REPLY;
  113. }
  114. static void *
  115. new_state(void *arg)
  116. {
  117. state_t *st;
  118. (void)arg;
  119. st = tor_malloc(sizeof(*st));
  120. /* Every thread gets its own keys. not a problem for benchmarking */
  121. st->rsa = crypto_pk_new();
  122. if (crypto_pk_generate_key_with_bits(st->rsa, 1024) < 0) {
  123. crypto_pk_free(st->rsa);
  124. tor_free(st);
  125. return NULL;
  126. }
  127. curve25519_secret_key_generate(&st->ecdh, 0);
  128. st->magic = 13371337;
  129. return st;
  130. }
  131. static void
  132. free_state(void *arg)
  133. {
  134. state_t *st = arg;
  135. crypto_pk_free(st->rsa);
  136. tor_free(st);
  137. }
  138. static tor_weak_rng_t weak_rng;
  139. static int n_sent = 0;
  140. static int rsa_sent = 0;
  141. static int ecdh_sent = 0;
  142. static int n_received = 0;
  143. static int no_shutdown = 0;
  144. #ifdef TRACK_RESPONSES
  145. bitarray_t *received;
  146. #endif
  147. static void
  148. handle_reply(void *arg)
  149. {
  150. #ifdef TRACK_RESPONSES
  151. rsa_work_t *rw = arg; /* Naughty cast, but only looking at serial. */
  152. tor_assert(! bitarray_is_set(received, rw->serial));
  153. bitarray_set(received,rw->serial);
  154. #endif
  155. tor_free(arg);
  156. ++n_received;
  157. }
  158. /* This should never get called. */
  159. static void
  160. handle_reply_shutdown(void *arg)
  161. {
  162. (void)arg;
  163. no_shutdown = 1;
  164. }
  165. static workqueue_entry_t *
  166. add_work(threadpool_t *tp)
  167. {
  168. int add_rsa =
  169. opt_ratio_rsa == 0 ||
  170. tor_weak_random_range(&weak_rng, opt_ratio_rsa) == 0;
  171. if (add_rsa) {
  172. rsa_work_t *w = tor_malloc_zero(sizeof(*w));
  173. w->serial = n_sent++;
  174. crypto_rand((char*)w->msg, 20);
  175. w->msglen = 20;
  176. ++rsa_sent;
  177. return threadpool_queue_work(tp, workqueue_do_rsa, handle_reply, w);
  178. } else {
  179. ecdh_work_t *w = tor_malloc_zero(sizeof(*w));
  180. w->serial = n_sent++;
  181. /* Not strictly right, but this is just for benchmarks. */
  182. crypto_rand((char*)w->u.pk.public_key, 32);
  183. ++ecdh_sent;
  184. return threadpool_queue_work(tp, workqueue_do_ecdh, handle_reply, w);
  185. }
  186. }
  187. static int n_failed_cancel = 0;
  188. static int n_successful_cancel = 0;
  189. static int
  190. add_n_work_items(threadpool_t *tp, int n)
  191. {
  192. int n_queued = 0;
  193. int n_try_cancel = 0, i;
  194. workqueue_entry_t **to_cancel;
  195. workqueue_entry_t *ent;
  196. to_cancel = tor_malloc(sizeof(workqueue_entry_t*) * opt_n_cancel);
  197. while (n_queued++ < n) {
  198. ent = add_work(tp);
  199. if (! ent) {
  200. puts("Z");
  201. tor_event_base_loopexit(tor_libevent_get_base(), NULL);
  202. return -1;
  203. }
  204. if (n_try_cancel < opt_n_cancel &&
  205. tor_weak_random_range(&weak_rng, n) < opt_n_cancel) {
  206. to_cancel[n_try_cancel++] = ent;
  207. }
  208. }
  209. for (i = 0; i < n_try_cancel; ++i) {
  210. void *work = workqueue_entry_cancel(to_cancel[i]);
  211. if (! work) {
  212. n_failed_cancel++;
  213. } else {
  214. n_successful_cancel++;
  215. tor_free(work);
  216. }
  217. }
  218. tor_free(to_cancel);
  219. return 0;
  220. }
  221. static int shutting_down = 0;
  222. static void
  223. replysock_readable_cb(tor_socket_t sock, short what, void *arg)
  224. {
  225. threadpool_t *tp = arg;
  226. replyqueue_t *rq = threadpool_get_replyqueue(tp);
  227. int old_r = n_received;
  228. (void) sock;
  229. (void) what;
  230. replyqueue_process(rq);
  231. if (old_r == n_received)
  232. return;
  233. if (opt_verbose) {
  234. printf("%d / %d", n_received, n_sent);
  235. if (opt_n_cancel)
  236. printf(" (%d cancelled, %d uncancellable)",
  237. n_successful_cancel, n_failed_cancel);
  238. puts("");
  239. }
  240. #ifdef TRACK_RESPONSES
  241. tor_mutex_acquire(&bitmap_mutex);
  242. for (i = 0; i < opt_n_items; ++i) {
  243. if (bitarray_is_set(received, i))
  244. putc('o', stdout);
  245. else if (bitarray_is_set(handled, i))
  246. putc('!', stdout);
  247. else
  248. putc('.', stdout);
  249. }
  250. puts("");
  251. tor_mutex_release(&bitmap_mutex);
  252. #endif
  253. if (n_sent - (n_received+n_successful_cancel) < opt_n_lowwater) {
  254. int n_to_send = n_received + opt_n_inflight - n_sent;
  255. if (n_to_send > opt_n_items - n_sent)
  256. n_to_send = opt_n_items - n_sent;
  257. add_n_work_items(tp, n_to_send);
  258. }
  259. if (shutting_down == 0 &&
  260. n_received+n_successful_cancel == n_sent &&
  261. n_sent >= opt_n_items) {
  262. shutting_down = 1;
  263. threadpool_queue_update(tp, NULL,
  264. workqueue_do_shutdown, NULL, NULL);
  265. // Anything we add after starting the shutdown must not be executed.
  266. threadpool_queue_work(tp, workqueue_shutdown_error,
  267. handle_reply_shutdown, NULL);
  268. {
  269. struct timeval limit = { 2, 0 };
  270. tor_event_base_loopexit(tor_libevent_get_base(), &limit);
  271. }
  272. }
  273. }
  274. static void
  275. help(void)
  276. {
  277. puts(
  278. "Options:\n"
  279. " -h Display this information\n"
  280. " -v Be verbose\n"
  281. " -N <items> Run this many items of work\n"
  282. " -T <threads> Use this many threads\n"
  283. " -I <inflight> Have no more than this many requests queued at once\n"
  284. " -L <lowwater> Add items whenever fewer than this many are pending\n"
  285. " -C <cancel> Try to cancel N items of every batch that we add\n"
  286. " -R <ratio> Make one out of this many items be a slow (RSA) one\n"
  287. " --no-{eventfd2,eventfd,pipe2,pipe,socketpair}\n"
  288. " Disable one of the alert_socket backends.");
  289. }
  290. int
  291. main(int argc, char **argv)
  292. {
  293. replyqueue_t *rq;
  294. threadpool_t *tp;
  295. int i;
  296. tor_libevent_cfg evcfg;
  297. struct event *ev;
  298. uint32_t as_flags = 0;
  299. for (i = 1; i < argc; ++i) {
  300. if (!strcmp(argv[i], "-v")) {
  301. opt_verbose = 1;
  302. } else if (!strcmp(argv[i], "-T") && i+1<argc) {
  303. opt_n_threads = atoi(argv[++i]);
  304. } else if (!strcmp(argv[i], "-N") && i+1<argc) {
  305. opt_n_items = atoi(argv[++i]);
  306. } else if (!strcmp(argv[i], "-I") && i+1<argc) {
  307. opt_n_inflight = atoi(argv[++i]);
  308. } else if (!strcmp(argv[i], "-L") && i+1<argc) {
  309. opt_n_lowwater = atoi(argv[++i]);
  310. } else if (!strcmp(argv[i], "-R") && i+1<argc) {
  311. opt_ratio_rsa = atoi(argv[++i]);
  312. } else if (!strcmp(argv[i], "-C") && i+1<argc) {
  313. opt_n_cancel = atoi(argv[++i]);
  314. } else if (!strcmp(argv[i], "--no-eventfd2")) {
  315. as_flags |= ASOCKS_NOEVENTFD2;
  316. } else if (!strcmp(argv[i], "--no-eventfd")) {
  317. as_flags |= ASOCKS_NOEVENTFD;
  318. } else if (!strcmp(argv[i], "--no-pipe2")) {
  319. as_flags |= ASOCKS_NOPIPE2;
  320. } else if (!strcmp(argv[i], "--no-pipe")) {
  321. as_flags |= ASOCKS_NOPIPE;
  322. } else if (!strcmp(argv[i], "--no-socketpair")) {
  323. as_flags |= ASOCKS_NOSOCKETPAIR;
  324. } else if (!strcmp(argv[i], "-h")) {
  325. help();
  326. return 0;
  327. } else {
  328. help();
  329. return 1;
  330. }
  331. }
  332. if (opt_n_threads < 1 ||
  333. opt_n_items < 1 || opt_n_inflight < 1 || opt_n_lowwater < 0 ||
  334. opt_n_cancel > opt_n_inflight || opt_n_inflight > MAX_INFLIGHT ||
  335. opt_ratio_rsa < 0) {
  336. help();
  337. return 1;
  338. }
  339. if (opt_n_inflight > opt_n_items) {
  340. opt_n_inflight = opt_n_items;
  341. }
  342. init_logging(1);
  343. network_init();
  344. if (crypto_global_init(1, NULL, NULL) < 0) {
  345. printf("Couldn't initialize crypto subsystem; exiting.\n");
  346. return 1;
  347. }
  348. if (crypto_seed_rng() < 0) {
  349. printf("Couldn't seed RNG; exiting.\n");
  350. return 1;
  351. }
  352. rq = replyqueue_new(as_flags);
  353. tor_assert(rq);
  354. tp = threadpool_new(opt_n_threads,
  355. rq, new_state, free_state, NULL);
  356. tor_assert(tp);
  357. crypto_seed_weak_rng(&weak_rng);
  358. memset(&evcfg, 0, sizeof(evcfg));
  359. tor_libevent_initialize(&evcfg);
  360. ev = tor_event_new(tor_libevent_get_base(),
  361. replyqueue_get_socket(rq), EV_READ|EV_PERSIST,
  362. replysock_readable_cb, tp);
  363. event_add(ev, NULL);
  364. #ifdef TRACK_RESPONSES
  365. handled = bitarray_init_zero(opt_n_items);
  366. received = bitarray_init_zero(opt_n_items);
  367. tor_mutex_init(&bitmap_mutex);
  368. handled_len = opt_n_items;
  369. #endif
  370. for (i = 0; i < opt_n_inflight; ++i) {
  371. if (! add_work(tp)) {
  372. puts("Couldn't add work.");
  373. return 1;
  374. }
  375. }
  376. {
  377. struct timeval limit = { 180, 0 };
  378. tor_event_base_loopexit(tor_libevent_get_base(), &limit);
  379. }
  380. event_base_loop(tor_libevent_get_base(), 0);
  381. if (n_sent != opt_n_items || n_received+n_successful_cancel != n_sent) {
  382. printf("%d vs %d\n", n_sent, opt_n_items);
  383. printf("%d+%d vs %d\n", n_received, n_successful_cancel, n_sent);
  384. puts("FAIL");
  385. return 1;
  386. } else if (no_shutdown) {
  387. puts("Accepted work after shutdown\n");
  388. puts("FAIL");
  389. } else {
  390. puts("OK");
  391. return 0;
  392. }
  393. }