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

static int
allocate(void)
{
	long n = (long) rand() % MAX_SIZE;
	/* Constraint the size */
	n = n % MAX_SIZE;

	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, nbytes %.1fM (%c)\n",
				iter, nblocks, (double) nbytes / (1024. * 1024.),
				c);
	}
}

int main(int argc, char *argv[])
{
	torture();
	return 0;
}