/*-
 * Copyright (C) 2008 Jason Evans <jasone@FreeBSD.org>.
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice(s), this list of conditions and the following disclaimer as
 *    the first lines of this file unmodified other than the possible
 *    addition of one or more copyright notices.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice(s), this list of conditions and the following disclaimer in
 *    the documentation and/or other materials provided with the
 *    distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDER(S) ``AS IS'' AND ANY
 * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
 * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT HOLDER(S) BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR
 * BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
 * WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
 * OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE,
 * EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 *******************************************************************************
 *
 * This microbenchmark tests concurrent throughput of malloc(3), where there
 * are producers and consumers.  Each producer allocates objects and hands them
 * off to one or more consumers, which deallocate the objects.
 *
 *******************************************************************************
 */

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <stdbool.h>
#include <stdint.h>
#include <errno.h>
#include <limits.h>
#include <pthread.h>

#include <machine/atomic.h>

typedef struct {
	unsigned long		pro;
	unsigned long		con;
	bool			verbose;
	void			**input;
	unsigned long		ninput;
} con_args_t;

void *
con_enter(void *vargs)
{
	con_args_t *args = (con_args_t *)vargs;
	unsigned long long nfrees;
	bool finish;
	unsigned long i;
	void *ptr;

	/*
	 * Cause jemalloc to seed PRNGs for this thread.  If this thread truly
	 * never allocates, seeding will never happen.  Without this malloc(),
	 * the PRNGs do drift apart as time goes on; the main purpose here is
	 * to make the benchmark more realistic rather than to favor jemalloc.
	 */
	ptr = malloc(1);
	if (ptr == NULL) {
		fprintf(stderr, "procon: Error in malloc()\n");
		_exit(1);
	}
	free(ptr);

	nfrees = 0;
	for (finish = false; finish == false;) {
		for (i = 0; i < args->ninput; i++) {
			ptr = (void *)atomic_readandclear_ptr((uintptr_t *)
			    &args->input[i]);
			if (ptr == (void *)-1) /* Producer says to finish. */
				finish = true;
			else if (ptr != NULL) {
				free(ptr);
				nfrees++;
			}
		}
	}

	if (args->verbose) {
		fprintf(stderr,
		    "  Consumer %lu:%lu completed %llu cycles\n",
		    args->pro, args->con, nfrees);
	}

	return (NULL);
}

typedef struct {
	unsigned long		pro;
	bool			verbose;
	unsigned long		con;
	unsigned long long	allocs;
	unsigned long		queue;
	size_t			minsize;
	size_t			maxsize;
} pro_args_t;

void *
pro_enter(void *vargs)
{
	pro_args_t *args = (pro_args_t *)vargs;
	pthread_t *cons;
	con_args_t *con_args;
	void ** output;
	int err;
	unsigned long long i, size_total;
	unsigned ctx;
	size_t size;
	void *ptr;

	output = (void **)calloc(args->con * args->queue, sizeof(void *));
	if (output == NULL) {
		fprintf(stderr, "procon: Error in calloc(%zu, %zu)\n",
		    args->con * args->queue, sizeof(void *));
		_exit(1);
	}

	cons = (pthread_t *)malloc(sizeof(pthread_t) * args->con);
	if (cons == NULL) {
		fprintf(stderr, "procon: Error in malloc()\n");
		_exit(1);
	}

	con_args = (con_args_t *)malloc(sizeof(con_args_t) * args->con);
	if (con_args == NULL) {
		fprintf(stderr, "procon: Error in malloc()\n");
		_exit(1);
	}
	for (i = 0; i < args->con; i++) {
		con_args[i].pro = args->pro;
		con_args[i].con = i;
		con_args[i].verbose = args->verbose;
		con_args[i].input = &output[i * args->queue];
		con_args[i].ninput = args->queue;
	}
	
	for (i = 0; i < args->con; i++) {
		if ((err = pthread_create(&cons[i], NULL, con_enter,
		    (void *)&con_args[i])) != 0) {
			fprintf(stderr,
			    "procon: Error in pthread_create(): %s\n",
			    strerror(err));
			_exit(1);
		}
	}

	for (i = size_total = 0, ctx = (unsigned)args->pro;
	    i < args->allocs; i++, size_total += size) {
		size = args->minsize + (((float)rand_r(&ctx) / (float)RAND_MAX)
		    * (args->maxsize - args->minsize));
		ptr = calloc(1, size);
		if (ptr == NULL) {
			fprintf(stderr,
			    "procon: Producer %lu:"
			    " Error in calloc(1, %zu)\n",
			    args->pro, size);
			_exit(1);
		}

		/* Randomly stuff ptr into an empty element of output. */
		while (true) {
			unsigned long pick = (((float)rand_r(&ctx) /
			    (float)RAND_MAX) * ((args->con * args->queue) - 1));
			if (atomic_cmpset_ptr((uintptr_t *)&output[pick],
			    (uintptr_t)NULL, (uintptr_t)ptr))
				break;
		}
	}

	/* Tell consumers to finish. */
	for (i = 0; i < args->con; i++) {
		while (true) {
			if (atomic_cmpset_ptr((uintptr_t *)&output[i *
			    args->queue], (uintptr_t)NULL, (uintptr_t)-1))
				break;
		}
	}

	for (i = 0; i < args->con; i++) {
		if ((err = pthread_join(cons[i], NULL)) != 0) {
			fprintf(stderr,
			    "procon: Error in pthread_join(): %s\n",
			    strerror(err));
			_exit(1);
		}
	}

	if (args->verbose) {
		fprintf(stderr,
		    "Producer %lu completed %llu cycles, %llu bytes\n",
		    args->pro, args->allocs, size_total);
	}

	return (NULL);
}

void
usage(FILE *fp)
{
	fprintf(fp, "procon usage:\n");
	fprintf(fp, "    procon [<options>]\n");
	fprintf(fp, "\n");
	fprintf(fp, "    Option      | Description\n");
	fprintf(fp, "    ------------+----------------------------------\n");
	fprintf(fp, "    -h          | Print usage and exit.\n");
	fprintf(fp, "    -v          | Print verbose output.\n");
	fprintf(fp, "    -p <pro>    | Number of producers.\n");
	fprintf(fp, "    -c <con>    | Number of consumers per producer.\n");
	fprintf(fp, "    -a <allocs> | Allocations per producer.\n");
	fprintf(fp, "    -q <queue>  | Per consumer queue size.\n");
	fprintf(fp, "    -s <size>   | Minimum allocation size.\n");
	fprintf(fp, "    -S <size>   | Maximum allocation size.\n");
}

int
main(int argc, char **argv)
{
	int c, err;
	bool opt_verbose = false;
	unsigned long opt_pro = 1;
	unsigned long opt_con = 16;
	unsigned long long opt_allocs = (1LLU << 20);
	unsigned long opt_queue = (1LU << 8);
	size_t opt_minsize = 1;
	size_t opt_maxsize = 2048;
	pthread_t *pros;
	pro_args_t *pro_args;
	unsigned long i;

	opterr = 0;
	optind = 1;
	while ((c = getopt(argc, argv, "hvp:c:a:q:s:S:")) != -1) {
		switch (c) {
		case 'h':
			usage(stdout);
			_exit(0);
		case 'v':
			opt_verbose = true;
			break;
		case 'p':
			errno = 0;
			opt_pro = strtoul(optarg, NULL, 0);
			if (opt_pro == ULONG_MAX && errno != 0) {
				fprintf(stderr,
				    "procon: Error in option '-%c %s': %s\n",
				    c, optarg, strerror(errno));
				usage(stderr);
				_exit(1);
			} else if (opt_pro == 0) {
				fprintf(stderr,
				    "procon: Invalid option '-%c %s'\n",
				    c, optarg);
				usage(stderr);
				_exit(1);
			}
			break;
		case 'c':
			errno = 0;
			opt_con = strtoul(optarg, NULL, 0);
			if (opt_con == ULONG_MAX && errno != 0) {
				fprintf(stderr,
				    "procon: Error in option '-%c %s': %s\n",
				    c, optarg, strerror(errno));
				usage(stderr);
				_exit(1);
			} else if (opt_con == 0) {
				fprintf(stderr,
				    "procon: Invalid option '-%c %s'\n",
				    c, optarg);
				usage(stderr);
				_exit(1);
			}
			break;
		case 'a':
			errno = 0;
			opt_allocs = strtoull(optarg, NULL, 0);
			if (opt_allocs == ULLONG_MAX && errno != 0) {
				fprintf(stderr,
				    "procon: Error in option '-%c %s': %s\n",
				    c, optarg, strerror(errno));
				usage(stderr);
				_exit(1);
			} else if (opt_allocs == 0) {
				fprintf(stderr,
				    "procon: Invalid option '-%c %s'\n",
				    c, optarg);
				usage(stderr);
				_exit(1);
			}
			break;
		case 'q':
			errno = 0;
			opt_queue = strtoul(optarg, NULL, 0);
			if (opt_queue == ULONG_MAX && errno != 0) {
				fprintf(stderr,
				    "procon: Error in option '-%c %s': %s\n",
				    c, optarg, strerror(errno));
				usage(stderr);
				_exit(1);
			} else if (opt_queue == 0) {
				fprintf(stderr,
				    "procon: Invalid option '-%c %s'\n",
				    c, optarg);
				usage(stderr);
				_exit(1);
			}
			break;
		case 's':
			errno = 0;
			opt_minsize = strtoull(optarg, NULL, 0);
			if (opt_minsize == ULLONG_MAX && errno != 0) {
				fprintf(stderr,
				    "procon: Error in option '-%c %s': %s\n",
				    c, optarg, strerror(errno));
				usage(stderr);
				_exit(1);
			}
			break;
		case 'S':
			errno = 0;
			opt_maxsize = strtoull(optarg, NULL, 0);
			if (opt_maxsize == ULLONG_MAX && errno != 0) {
				fprintf(stderr,
				    "procon: Error in option '-%c %s': %s\n",
				    c, optarg, strerror(errno));
				usage(stderr);
				_exit(1);
			}
			break;
		default:
			fprintf(stderr, "procon: Unrecognized option\n");
			usage(stderr);
			_exit(1);
		}
	}

	if (opt_maxsize < opt_minsize) {
		fprintf(stderr,
		    "procon: Maximum size smaller than minimum size"
		    " (%zu < %zu)\n", opt_maxsize, opt_minsize);
		usage(stderr);
		_exit(1);
	}

	if (opt_verbose) {
		fprintf(stdout, "procon: Verbose output enabled\n");
		fprintf(stdout, "procon: Producers: %lu\n", opt_pro);
		fprintf(stdout, "procon: Consumers per producer: %lu\n",
		    opt_con);
		fprintf(stdout, "procon: Allocations per producer: %llu\n",
		    opt_allocs);
		fprintf(stdout, "procon: Per producer queue size: %lu\n",
		    opt_queue);
		fprintf(stdout, "procon: Allocation size [%zu..%zu]\n",
		    opt_minsize, opt_maxsize);
	}

	pros = (pthread_t *)malloc(sizeof(pthread_t) * opt_pro);
	if (pros == NULL) {
		fprintf(stderr, "procon: Error in malloc()\n");
		_exit(1);
	}

	pro_args = (pro_args_t *)malloc(sizeof(pro_args_t) * opt_pro);
	if (pro_args == NULL) {
		fprintf(stderr, "procon: Error in malloc()\n");
		_exit(1);
	}
	for (i = 0; i < opt_pro; i++) {
		pro_args[i].pro = i;
		pro_args[i].verbose = opt_verbose;
		pro_args[i].con = opt_con;
		pro_args[i].allocs = opt_allocs;
		pro_args[i].queue = opt_queue;
		pro_args[i].minsize = opt_minsize;
		pro_args[i].maxsize = opt_maxsize;
	}

	for (i = 0; i < opt_pro; i++) {
		if ((err = pthread_create(&pros[i], NULL, pro_enter,
		    (void *)&pro_args[i])) != 0) {
			fprintf(stderr,
			    "procon: Error in pthread_create(): %s\n",
			    strerror(err));
			_exit(1);
		}
	}

	for (i = 0; i < opt_pro; i++) {
		if ((err = pthread_join(pros[i], NULL)) != 0) {
			fprintf(stderr,
			    "procon: Error in pthread_join(): %s\n",
			    strerror(err));
			_exit(1);
		}
	}

	return (0);
}