/*-
 * Copyright (c) 2005 Robert N. M. Watson
 * 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, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, 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 AUTHOR AND CONTRIBUTORS ``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 AUTHOR OR CONTRIBUTORS 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.
 *
 * $FreeBSD$
 */

#include <sys/param.h>
#include <sys/sysctl.h>

#include <err.h>
#include <errno.h>
#include <memstat.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

/*
 * libmemstat(3) allows us to maintain some of our own data with each memory
 * type.  Take advantage of this to store history information between runs so
 * that we can calculate a delta.
 */
#define	memstat_get_lastnumallocs(p)	memstat_get_caller_uint64(p, 0)
#define	memstat_set_lastnumallocs(p, v)	memstat_set_caller_uint64(p, 0, v)

#define	memstat_get_lastnumfrees(p)	memstat_get_caller_uint64(p, 1)
#define	memstat_set_lastnumfrees(p, v)	memstat_set_caller_uint64(p, 1, v)

#define	memstat_get_lastmemalloced(p)	memstat_get_caller_uint64(p, 2)
#define	memstat_set_lastmemalloced(p, v) memstat_set_caller_uint64(p, 2, v)

#define	memstat_get_lastmemfreed(p)	memstat_get_caller_uint64(p, 3)
#define	memstat_set_lastmemfreed(p, v)	memstat_set_caller_uint64(p, 3, v)

#define	memstat_get_activity(p)		memstat_get_caller_uint64(p, 4)
#define	memstat_set_activity(p, v)	memstat_set_caller_uint64(p, 4, v)

#define	memstat_get_score(p)		memstat_get_caller_uint64(p, 5)
#define	memstat_set_score(p, v)		memstat_set_caller_uint64(p, 5, v)

#define	memstat_get_lastcount(p)	memstat_get_caller_uint64(p, 6)
#define	memstat_set_lastcount(p, v)	memstat_set_caller_uint64(p, 6, v)

#define	memstat_get_lastbytes(p)	memstat_get_caller_uint64(p, 7)
#define	memstat_set_lastbytes(p, v)	memstat_set_caller_uint64(p, 7, v)

#define	memstat_get_lastfailures(p)	memstat_get_caller_uint64(p, 8)
#define	memstat_set_lastfailures(p, v)	memstat_set_caller_uint64(p, 8, v)

#define	DISPLAY_ALLOCATOR	0x00000001

/*
 * Measurements over the lifetime of all memory types.
 */
static uint64_t	lifetime_bytes_allocated;
static uint64_t	lifetime_bytes_freed;
static uint64_t	lifetime_allocs_count;
static uint64_t	lifetime_frees_count;
static uint64_t	lifetime_fails_count;

/*
 * Measurements since the last update.
 */
static int64_t	update_delta_bytes;
static int64_t	update_delta_count;

static uint64_t	new_bytes_allocated;
static uint64_t	new_bytes_freed;
static uint64_t	new_allocs_count;
static uint64_t	new_frees_count;
static uint64_t	new_fails_count;

/*
 * Current condition.
 */
static uint64_t	current_bytes_allocated;
static uint64_t	current_allocs_count;

/*
 * Given a memory type list that has been updated from the live kernel,
 * update our own derived fields relating to liveliness and inter-snapshot
 * changes.
 */
static void
update(struct memory_type_list *mtlp, int (*func)(struct memory_type_list *,
    int), const char *func_name)
{
	struct memory_type *mtp;
	int64_t delta_allocs, delta_bytes, delta_count, delta_fails;
	int64_t delta_frees;
	uint64_t delta_memalloced, delta_memfreed;
	uint64_t activity, score;

	uint64_t old_lifetime_bytes_allocated, old_lifetime_bytes_freed;
	uint64_t old_lifetime_allocs_count, old_lifetime_frees_count;
	uint64_t old_bytes_allocated, old_allocs_count;
	uint64_t old_lifetime_fails_count;

	old_lifetime_bytes_allocated = lifetime_bytes_allocated;
	old_lifetime_bytes_freed = lifetime_bytes_freed;
	old_lifetime_allocs_count = lifetime_allocs_count;
	old_lifetime_frees_count = lifetime_frees_count;
	old_lifetime_fails_count = lifetime_fails_count;

	old_bytes_allocated = current_bytes_allocated;
	old_allocs_count = current_allocs_count;

	update_delta_bytes = 0;
	update_delta_count = 0;

	new_bytes_allocated = 0;
	new_bytes_freed = 0;
	new_allocs_count = 0;
	new_frees_count = 0;
	new_fails_count = 0;

	current_bytes_allocated = 0;
	current_allocs_count = 0;

	for (mtp = memstat_mtl_first(mtlp); mtp != NULL;
	    mtp = memstat_mtl_next(mtp)) {
		memstat_set_lastnumallocs(mtp, memstat_get_numallocs(mtp));
		memstat_set_lastnumfrees(mtp, memstat_get_numfrees(mtp));
		memstat_set_lastmemalloced(mtp, memstat_get_memalloced(mtp));
		memstat_set_lastmemfreed(mtp, memstat_get_memfreed(mtp));
		memstat_set_lastcount(mtp, memstat_get_count(mtp));
		memstat_set_lastbytes(mtp, memstat_get_bytes(mtp));
		memstat_set_lastfailures(mtp, memstat_get_failures(mtp));
	}

	/*
	 * Update/add all types.
	 */
	if (func(mtlp, 0) < 0) {
		warnx("update: %s: %s", func_name,
		    memstat_strerror(memstat_mtl_geterror(mtlp)));
		exit(-1);
	}

	/*
	 * Calculate and save the score for each type: the score is the
	 * number of allocations and frees that has occurred since the last
	 * pass, so reflects the level of activity of the type.
	 *
	 * We also calculate and maintain a general 'activity' counter, which
	 * allows us to keep a type higher on the scoreboard if it was active
	 * in the last few passes, even if it's not active in the current
	 * pass.
	 *
	 * Update global counters and measurements for the summary bar above
	 * the memory type list output.
	 */
	for (mtp = memstat_mtl_first(mtlp); mtp != NULL;
	    mtp = memstat_mtl_next(mtp)) {
		delta_allocs = memstat_get_numallocs(mtp) -
		    memstat_get_lastnumallocs(mtp);
		delta_frees = memstat_get_numfrees(mtp) -
		    memstat_get_lastnumfrees(mtp);
		delta_fails = memstat_get_failures(mtp) -
		    memstat_get_lastfailures(mtp);
		delta_count = memstat_get_count(mtp) -
		    memstat_get_lastcount(mtp);
		delta_bytes = memstat_get_bytes(mtp) -
		    memstat_get_lastbytes(mtp);
		delta_memalloced = memstat_get_memalloced(mtp) -
		    memstat_get_lastmemalloced(mtp);
		delta_memfreed = memstat_get_memfreed(mtp) -
		    memstat_get_lastmemfreed(mtp);
		activity = memstat_get_activity(mtp);

		lifetime_bytes_allocated += delta_memalloced;
		lifetime_bytes_freed += delta_memfreed;
		lifetime_allocs_count += delta_allocs;
		lifetime_frees_count += delta_frees;
		lifetime_fails_count += delta_fails;

		update_delta_bytes += delta_bytes;
		update_delta_count += delta_count;

		new_bytes_allocated += delta_memalloced;
		new_bytes_freed += delta_memfreed;
		new_allocs_count += delta_allocs;
		new_frees_count += delta_frees;
		new_fails_count += delta_fails;

		current_bytes_allocated += memstat_get_bytes(mtp);
		current_allocs_count += memstat_get_count(mtp);

		score = delta_allocs + delta_frees;
		if (score != 0) {
			if (activity < 5)
				activity++;
		} else {
			if (activity > 0)
				activity--;
		}

		memstat_set_activity(mtp, activity);
		memstat_set_score(mtp, score);
	}
}

/*
 * In order to sort by "activity", we need a sort score.  We generate the
 * score by considering how many malloc and free operations have taken place.
 * We may well want to include size of allocation in the sort order as well.
 */
static int
memory_type_activity_score(const struct memory_type *mtp)
{
	uint64_t score;

	score = memstat_get_score(mtp);
	if (score != 0)
		return (score + 5);

	if (memstat_get_activity(mtp));

	return (memstat_get_score(mtp));
}

/*
 * qsort comparison routine, comparing the activity level of two types.
 */
static int
mtp_compare(const void *a, const void *b)
{
	const struct memory_type **mtpap, **mtpbp;
	const struct memory_type *mtpa, *mtpb;
	int a_score, b_score;

	mtpap = (const struct memory_type **)a;
	mtpbp = (const struct memory_type **)b;

	mtpa = *mtpap;
	mtpb = *mtpbp;

	a_score = memory_type_activity_score(mtpa);
	b_score = memory_type_activity_score(mtpb);

	return (b_score - a_score);
}

static void
display(struct memory_type_list *mtlp, int flags)
{
	struct memory_type *mtp, **mtp_array;
	long delta_alloc, delta_free, delta_bytes;
	int count, i, j, shown;

	count = 0;
	for (mtp = memstat_mtl_first(mtlp); mtp != NULL;
	    mtp = memstat_mtl_next(mtp))
		count++;

	mtp_array = malloc(sizeof(mtp) * count);

	i = 0;
	for (mtp = memstat_mtl_first(mtlp); mtp != NULL;
	    mtp = memstat_mtl_next(mtp)) {
		mtp_array[i] = mtp;
		i++;
	}

	qsort(mtp_array, count, sizeof(mtp), mtp_compare);

	/* XXXRW: In the ncurses world order, this would be nicer. */
	for (i = 0; i < 25; i++)
		printf("\n");

	printf("Lifetime bytes allocated:  %12llu\n", lifetime_bytes_allocated);
	printf("Lifetime bytes freed:      %12llu\n", lifetime_bytes_freed);

	printf("Lifetime allocations:      %12llu\n", lifetime_allocs_count);
	printf("Lifetime frees:            %12llu\n", lifetime_frees_count);
	printf("Lifetime failures:         %12llu\n", lifetime_fails_count);

	printf("Delta bytes:               %12lld\n", update_delta_bytes);
	printf("Delta count:               %12lld\n", update_delta_count);

	printf("New bytes allocated:       %12lld\n", new_bytes_allocated);
	printf("New bytes freed:           %12lld\n", new_bytes_freed);
	printf("New count allocates:       %12lld\n", new_allocs_count);
	printf("New count frees:           %12lld\n", new_frees_count);
	printf("New count fails:           %12lld\n", new_fails_count);

	printf("Current bytes allocated:   %12llu\n", current_bytes_allocated);
	printf("Current allocs count:      %12llu\n", current_allocs_count);

	printf("\n");

	printf("%12s ", "Type");

	if (flags & DISPLAY_ALLOCATOR)
		printf("%8s ", "Alloc");

	printf("%8s %8s %8s %8s %10s %10s\n", "#Alloc'd", "#Free'd", "DAlloc",
	    "DBytes", "TAlloc", "TBytes");

	shown = 0;
	for (j = 0; j < count; j++) {
		mtp = mtp_array[j];
		if (memstat_get_numallocs(mtp) == 0)
			continue;
#if 0
		if (memory_type_activity_score(mtp) == 0)
			continue;
#endif

		printf("%12s ", memstat_get_name(mtp));

		if (flags & DISPLAY_ALLOCATOR) {
			switch(memstat_get_allocator(mtp)) {
			case ALLOCATOR_MALLOC:
				printf("%8s ", "malloc");
				break;
			case ALLOCATOR_UMA:
				printf("%8s ", "uma");
				break;
			default:
				printf("%8s ", "unknown");
			}
		}

		delta_alloc = (long)(memstat_get_numallocs(mtp) -
		    memstat_get_lastnumallocs(mtp));
		delta_free = (long)(memstat_get_numfrees(mtp) -
		    memstat_get_lastnumfrees(mtp));
		delta_bytes = (long)
		    (memstat_get_memalloced(mtp) -
		    memstat_get_memfreed(mtp)) -
		    (memstat_get_lastmemalloced(mtp) -
		    memstat_get_lastmemfreed(mtp));
		printf("%8ld %8ld %8ld %8ld ", delta_alloc, delta_free,
		    delta_alloc - delta_free, delta_bytes);

		printf("%10llu %10llu ",
		    (unsigned long long)(memstat_get_count(mtp)),
		    (unsigned long long)(memstat_get_bytes(mtp)));

		printf("\n");
		shown++;
		if (shown > 20)
			break;
	}

	free(mtp_array);
}

static void
usage(void)
{

	fprintf(stderr, "memtop [-amz]\n");
	exit(-1);
}

int
main(int argc, char *argv[])
{
	int (*func)(struct memory_type_list *, int) = NULL;
	struct memory_type_list *mtlp;
	const char *func_name = "";
	int flags;

	if (argc != 2)
		usage();

	flags = 0;
	if (strcmp(argv[1], "-a") == 0) {
		func = memstat_sysctl_all;
		func_name = "memstat_sysctl_all";
		flags = DISPLAY_ALLOCATOR;
	} else if (strcmp(argv[1], "-m") == 0) {
		func = memstat_sysctl_malloc;
		func_name = "memstat_sysctl_malloc";
	} else if (strcmp(argv[1], "-z") == 0) {
		func = memstat_sysctl_uma;
		func_name = "memstat_sysctl_uma";
	} else
		usage();

	mtlp = memstat_mtl_alloc();
	if (mtlp == NULL)
		err(-1, "memstat_mtl_init");
	update(mtlp, func, func_name);
	while (1) {
		sleep(1);
		update(mtlp, func, func_name);
		display(mtlp, flags);
	}

	memstat_mtl_free(mtlp);

	return (0);
}