From 453ccd4596b9b0ae3731433eaba281385924eb08 Mon Sep 17 00:00:00 2001
From: Max Wash <mbswash@gmail.com>
Date: Thu, 2 Feb 2023 21:16:17 +0000
Subject: [PATCH] sandbox: multi-threaded kmalloc() stress-test

---
 sandbox/Makefile           |  1 +
 sandbox/base/memory_test.c | 60 ++++++++++++++++++++++++++++++++++----
 2 files changed, 56 insertions(+), 5 deletions(-)

diff --git a/sandbox/Makefile b/sandbox/Makefile
index 81a5206..71f1d85 100644
--- a/sandbox/Makefile
+++ b/sandbox/Makefile
@@ -13,6 +13,7 @@ OBJ := $(addprefix $(BUILD_DIR)/,$(SRC:.c=.o))
 DEPS := $(OBJ:.o=.d)
 
 CFLAGS := $(INCLUDE_DIRS) -g
+LDFLAGS := -pthread
 
 all: $(BUILD_DIR)/$(EXEC_NAME)
 	@for prog in $(SANDBOX_DIR_LIST); do \
diff --git a/sandbox/base/memory_test.c b/sandbox/base/memory_test.c
index 6b528ef..619f30d 100644
--- a/sandbox/base/memory_test.c
+++ b/sandbox/base/memory_test.c
@@ -1,9 +1,10 @@
-#include "socks/queue.h"
+#include <socks/queue.h>
 #include <stdint.h>
 #include <stdio.h>
 #include <stddef.h>
 #include <stdlib.h>
 #include <inttypes.h>
+#include <pthread.h>
 #include <time.h>
 #include <assert.h>
 #include <sys/mman.h>
@@ -12,6 +13,8 @@
 #include <socks/memblock.h>
 #include <socks/vm.h>
 
+#define NR_THREADS 8
+
 /* we're working with 512MiB of simulated system RAM */
 #define MEMORY_SIZE_MB 512
 
@@ -51,17 +54,51 @@ static void print_free_pages(vm_zone_t *z)
 	printf(" * %s:\n", z->z_info.zd_name);
 
 	for (int i = VM_PAGE_MIN_ORDER; i <= VM_PAGE_MAX_ORDER; i++) {
-		if (queue_length(&z->z_free_pages[i]) == 0) {
+		if (queue_empty(&z->z_free_pages[i])) {
 			continue;
 		}
 
 		char size_str[64];
 		data_size_to_string(vm_page_order_to_bytes(i), size_str, sizeof size_str);
 
-		printf("  - %u pages with size %s (order-%u)\n", queue_length(&z->z_free_pages[i]), size_str, i);
+		printf("  - %zu pages with size %s (order-%u)\n", queue_length(&z->z_free_pages[i]), size_str, i);
 	}
 }
 
+static void *kmalloc_test_thread(void *p)
+{
+	size_t thread_id = (size_t)p;
+	struct timespec ts = { .tv_sec = 1 };
+	void *allocated[4096];
+	while (1) {
+		int op = rand() % 2;
+		if (op == 1) {
+			for (int i = 0; i < 4096; i++) {
+				if (!allocated[i]) {
+					unsigned int size = (rand() % 4095) + 1;
+					allocated[i] = kmalloc(size, 0);
+					printf("thread %zu: allocated %u bytes at %p\n", thread_id, size, allocated[i]);
+					assert(allocated[i]);
+					break;
+				}
+			}
+		} else {
+			for (int i = 4095; i >= 0; i--) {
+				if (allocated[i]) {
+					kfree(allocated[i]);
+					printf("thread %zu: freed %p\n", thread_id, allocated[i]);
+					allocated[i] = NULL;
+					break;
+				}
+			}
+		}
+
+		//nanosleep(&ts, NULL);
+	}
+
+	return NULL;
+}
+
 static void print_all_pages(void)
 {
 	for (phys_addr_t i = 0; i < UINTPTR_MAX; ) {
@@ -242,8 +279,21 @@ int memory_test(void)
 	}
 
 	void *p = kmalloc(32, 0);
-	printf("allocated 32 bytes at %p\n", p);
-	
+	printf("kmalloc'd 32 bytes at %p\n", p);
+	kfree(p);
+	printf("kfree'd 32 bytes at %p\n", p);
+	p = kmalloc(32, 0);
+	printf("kmalloc'd 32 bytes at %p\n", p);
+
+	pthread_t threads[NR_THREADS];
+	for (size_t i = 0; i < NR_THREADS; i++) {
+		pthread_create(&threads[i], NULL, kmalloc_test_thread, (void *)i);
+	}
+
+	for (size_t i = 0; i < NR_THREADS; i++) {
+		pthread_join(threads[i], NULL);
+	}
+
 	munmap(system_memory, MB_TO_BYTES(MEMORY_SIZE_MB));
 	return 0;
 }