static int
add_n_work_items(threadpool_t *tp, int n)
{
int n_queued = 0;
int n_try_cancel = 0, i;
workqueue_entry_t **to_cancel;
workqueue_entry_t *ent;
to_cancel = tor_malloc(sizeof(workqueue_entry_t*) * opt_n_cancel);
while (n_queued++ < n) {
ent = add_work(tp);
if (! ent) {
tor_event_base_loopexit(tor_libevent_get_base(), NULL);
return -1;
}
if (n_try_cancel < opt_n_cancel &&
tor_weak_random_range(&weak_rng, n) < opt_n_cancel) {
to_cancel[n_try_cancel++] = ent;
}
}
for (i = 0; i < n_try_cancel; ++i) {
void *work = workqueue_entry_cancel(to_cancel[i]);
if (! work) {
n_failed_cancel++;
} else {
n_successful_cancel++;
tor_free(work);
}
}
tor_free(to_cancel);
return 0;
}
static void
replysock_readable_cb(tor_socket_t sock, short what, void *arg)
{
threadpool_t *tp = arg;
replyqueue_t *rq = threadpool_get_replyqueue(tp);
int old_r = n_received;
(void) sock;
(void) what;
replyqueue_process(rq);
if (old_r == n_received)
return;
if (opt_verbose) {
printf("%d / %d", n_received, n_sent);
if (opt_n_cancel)
printf(" (%d cancelled, %d uncancellable)",
n_successful_cancel, n_failed_cancel);
puts("");
}
#ifdef TRACK_RESPONSES
tor_mutex_acquire(&bitmap_mutex);
for (i = 0; i < opt_n_items; ++i) {
if (bitarray_is_set(received, i))
putc('o', stdout);
else if (bitarray_is_set(handled, i))
putc('!', stdout);
else
putc('.', stdout);
}
puts("");
tor_mutex_release(&bitmap_mutex);
#endif
if (n_sent - (n_received+n_successful_cancel) < opt_n_lowwater) {
int n_to_send = n_received + opt_n_inflight - n_sent;
if (n_to_send > opt_n_items - n_sent)
n_to_send = opt_n_items - n_sent;
add_n_work_items(tp, n_to_send);
}
if (shutting_down == 0 &&
n_received+n_successful_cancel == n_sent &&
n_sent >= opt_n_items) {
shutting_down = 1;
threadpool_queue_update(tp, NULL,
workqueue_do_shutdown, NULL, NULL);
// Anything we add after starting the shutdown must not be executed.
threadpool_queue_work(tp, workqueue_shutdown_error,
handle_reply_shutdown, NULL);
{
struct timeval limit = { 2, 0 };
tor_event_base_loopexit(tor_libevent_get_base(), &limit);
}
}
}
int
main(int argc, char **argv)
{
replyqueue_t *rq;
threadpool_t *tp;
int i;
tor_libevent_cfg evcfg;
struct event *ev;
uint32_t as_flags = 0;
for (i = 1; i < argc; ++i) {
if (!strcmp(argv[i], "-v")) {
opt_verbose = 1;
} else if (!strcmp(argv[i], "-T") && i+1<argc) {
opt_n_threads = atoi(argv[++i]);
} else if (!strcmp(argv[i], "-N") && i+1<argc) {
opt_n_items = atoi(argv[++i]);
} else if (!strcmp(argv[i], "-I") && i+1<argc) {
opt_n_inflight = atoi(argv[++i]);
} else if (!strcmp(argv[i], "-L") && i+1<argc) {
opt_n_lowwater = atoi(argv[++i]);
} else if (!strcmp(argv[i], "-R") && i+1<argc) {
opt_ratio_rsa = atoi(argv[++i]);
} else if (!strcmp(argv[i], "-C") && i+1<argc) {
opt_n_cancel = atoi(argv[++i]);
} else if (!strcmp(argv[i], "--no-eventfd2")) {
as_flags |= ASOCKS_NOEVENTFD2;
} else if (!strcmp(argv[i], "--no-eventfd")) {
as_flags |= ASOCKS_NOEVENTFD;
} else if (!strcmp(argv[i], "--no-pipe2")) {
as_flags |= ASOCKS_NOPIPE2;
} else if (!strcmp(argv[i], "--no-pipe")) {
as_flags |= ASOCKS_NOPIPE;
} else if (!strcmp(argv[i], "--no-socketpair")) {
as_flags |= ASOCKS_NOSOCKETPAIR;
} else if (!strcmp(argv[i], "-h")) {
help();
return 0;
} else {
help();
return 1;
}
}
if (opt_n_threads < 1 ||
opt_n_items < 1 || opt_n_inflight < 1 || opt_n_lowwater < 0 ||
opt_n_cancel > opt_n_inflight || opt_n_inflight > MAX_INFLIGHT ||
opt_ratio_rsa < 0) {
help();
return 1;
}
init_logging(1);
crypto_global_init(1, NULL, NULL);
crypto_seed_rng(1);
rq = replyqueue_new(as_flags);
tor_assert(rq);
tp = threadpool_new(opt_n_threads,
rq, new_state, free_state, NULL);
tor_assert(tp);
crypto_seed_weak_rng(&weak_rng);
memset(&evcfg, 0, sizeof(evcfg));
tor_libevent_initialize(&evcfg);
ev = tor_event_new(tor_libevent_get_base(),
replyqueue_get_socket(rq), EV_READ|EV_PERSIST,
replysock_readable_cb, tp);
event_add(ev, NULL);
#ifdef TRACK_RESPONSES
handled = bitarray_init_zero(opt_n_items);
received = bitarray_init_zero(opt_n_items);
tor_mutex_init(&bitmap_mutex);
handled_len = opt_n_items;
#endif
for (i = 0; i < opt_n_inflight; ++i) {
if (! add_work(tp)) {
puts("Couldn't add work.");
return 1;
}
}
{
struct timeval limit = { 30, 0 };
tor_event_base_loopexit(tor_libevent_get_base(), &limit);
}
event_base_loop(tor_libevent_get_base(), 0);
if (n_sent != opt_n_items || n_received+n_successful_cancel != n_sent) {
printf("%d vs %d\n", n_sent, opt_n_items);
printf("%d+%d vs %d\n", n_received, n_successful_cancel, n_sent);
puts("FAIL");
return 1;
} else {
puts("OK");
return 0;
}
}
