test_workqueue.c 11 KB

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