/* 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. */

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

#define MAX_SIZE (1024L * 1024L)

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

struct block *front = NULL;
struct block *tail = NULL;
long nblocks = 0;
long nbytes = 0;
long maxsize = MAX_SIZE;

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

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

	struct block *b = malloc(size);

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

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

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

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

	tail = b;

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

	nblocks++;
	nbytes += size;

	return 0;
}

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

	struct block *b = front;

	front = b->next;

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

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

	free(b);

	return 0;
}

static void
torture(void)
{
	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() == 0)
				c = 'A';
			else
				c = '-';
		} else {
			if (deallocate() == 0)
				c = 'D';
			else
				c = '-';
		}

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

int main(int argc, char *argv[])
{
	if (argc > 1)
		maxsize = atol(argv[1]);

	printf("memtool v1.0.0 maxsize=%ldK\n", maxsize / 1024);

	torture();

	return 0;
}