/* Copyright (c) 2024 Barcelona Supercomputing Center (BSC)
 * SPDX-License-Identifier: MIT
 * Author: Rodrigo Arias Mallo <rodrigo.arias@bsc.es> */

/* This is just a small tool to exercise the memory which attempts to
 * stress the virtual memory, in a crude attempt to reproduce the hangs
 * that we were observing while booting NixOS. */

/* Changelog:
 *   v0.0.1 (2024-07-10): Start version with "chain" and "fill" tests.
 */

#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>

#define MAX_SIZE (1024L * 1024L)

struct block {
	struct block *next;
	size_t size;
	uint32_t data[];
};

struct chain {
	struct block *front;
	struct block *tail;
	long maxsize;
	long nbytes;
	long nblocks;
};

static int
allocate(struct chain *chain)
{
	/* Constraint the number of elements based on the maxsize */
	long maxn = chain->maxsize / sizeof(uint32_t);
	long n = (long) rand() % maxn;

	size_t size = sizeof(struct block) + n * sizeof(uint32_t);

	printf("allocating...\n");

	struct block *b = malloc(size);

	/* No mem */
	if (b == NULL)
		return -1;

	b->size = size;
	b->next = NULL;

	/* Populate the block with some data */
	printf("filling...\n");
	for (long i = 0; i < n; i++)
		b->data[i] = rand();

	/* Add it to the chain */
	if (chain->tail)
		chain->tail->next = b;

	chain->tail = b;

	/* And to the front if it is the first */
	if (!chain->front)
		chain->front = b;

	chain->nblocks++;
	chain->nbytes += size;

	return 0;
}

static int
deallocate(struct chain *chain)
{
	/* May run out of blocks */
	if (!chain->front)
		return -1;

	struct block *b = chain->front;

	chain->front = b->next;

	/* Last block */
	if (chain->tail == b)
		chain->tail = NULL;

	chain->nblocks--;
	chain->nbytes -= b->size;

	printf("deallocating...\n");
	free(b);

	return 0;
}

static void
do_chain(int argc, char *argv[])
{
	struct chain chain = {0};

	/* Default 1 MiB */
	chain.maxsize = 1024L * 1024L;

	if (argc > 0)
		chain.maxsize = atol(argv[0]);

	printf("mode chain: maxsize=%ldK\n", chain.maxsize / 1024);

	srand(123);

	for (long iter = 0; ; iter++) {
		int p = rand() % 100;
		int is_alloc = (p > 10);
		int ret = 0;
		char c;
		if (is_alloc) {
			if (allocate(&chain) == 0)
				c = 'A';
			else
				c = '-';
		} else {
			if (deallocate(&chain) == 0)
				c = 'D';
			else
				c = '-';
		}

		printf("iter=%ld nblocks=%ld allocated=%ldK (%c)\n",
				iter, chain.nblocks, chain.nbytes / 1024,
				c);
	}
}

static void
do_fill(int argc, char *argv[])
{
	/* Default: 256 MiB */
	long nbytes = 256L * 1024L * 1024L;

	if (argc > 0)
		nbytes = atol(argv[0]);

	long n = nbytes / sizeof(int);

	printf("mode fill: nbytes=%ldM, n=%ld\n",
			nbytes / (1024L * 1024L), n);

	int *buf = malloc(nbytes);

	if (!buf) {
		perror("malloc failed");
		exit(1);
	}

	for (long i = 0; i < n; i++) {
		buf[i] = i;
		if ((i % (1024L * 1024L)) == 0)
			printf("written=%ldK, addr=%p OK\n",
					i * sizeof(int) / 1024L,
					&buf[i]);
	}

	free(buf);

	printf("fill test OK\n");
}

static void
usage(void)
{
	printf(
"Usage: memtool <command> [<options>...]\n"
"\n"
"Available commands:\n"
"  chain [<maxsize>]\n"
"        Creates a chain of blocks of random size, each up to maxsize\n"
"        or 1MiB if not given. Blocks are freed with 10% probability\n"
"        starting from the oldest.\n"
"\n"
"  fill [<size>]\n"
"        Allocates a vector of the given size (or 256 MiB if not given)\n"
"        and initializes it with a increasing value per element.\n"
"\n");

	exit(1);
}

int main(int argc, char *argv[])
{
	printf("memtool v0.0.1 - Rodrigo Arias Mallo <rodrigo.arias@bsc.es>\n");

	if (argc < 2)
		usage();

	/* Skip program name */
	argc--; argv++;

	const char *mode = argv[0];

	/* Skip mode */
	argc--; argv++;

	if (strcmp(mode, "chain") == 0)
		do_chain(argc, argv);
	else if (strcmp(mode, "fill") == 0)
		do_fill(argc, argv);
	else
		usage();

	return 0;
}