feat: wasm_memory pool

runtime/include/module.h

@@ -10,6 +10,7 @@
 #include "awsm_abi.h"
 #include "http.h"
 #include "panic.h"
+#include "pool.h"
 #include "types.h"
 
 #define MODULE_DEFAULT_REQUEST_RESPONSE_SIZE (PAGE_SIZE)
@@ -17,6 +18,8 @@
 #define MODULE_MAX_NAME_LENGTH 32
 #define MODULE_MAX_PATH_LENGTH 256
 
+extern thread_local int worker_thread_idx;
+
 /*
  * Defines the listen backlog, the queue length for completely established socketeds waiting to be accepted
  * If this value is greater than the value in /proc/sys/net/core/somaxconn (typically 128), then it is silently
@@ -32,12 +35,18 @@
   "MODULE_MAX_PENDING_CLIENT_REQUESTS likely exceeds the value in /proc/sys/net/core/somaxconn and thus may be silently truncated";
 #endif
 
+/* TODO: Dynamically size based on number of threads */
+#define MAX_WORKER_THREADS 64
+
+struct module_pools {
+	struct pool memory[MAX_WORKER_THREADS];
+};
+
 struct module {
 	/* Metadata from JSON Config */
 	char                   name[MODULE_MAX_NAME_LENGTH];
 	char                   path[MODULE_MAX_PATH_LENGTH];
 	uint32_t               stack_size; /* a specification? */
-	uint64_t               max_memory; /* perhaps a specification of the module. (max 4GB) */
 	uint32_t               relative_deadline_us;
 	int                    port;
 	struct admissions_info admissions_info;
@@ -55,6 +64,8 @@ struct module {
 
 	_Atomic uint32_t   reference_count; /* ref count how many instances exist here. */
 	struct wasm_table *indirect_table;
+
+	struct module_pools pools;
 };
 
 /*************************
@@ -141,12 +152,42 @@ module_release(struct module *module)
 	return;
 }
 
+static inline struct wasm_memory *
+module_allocate_linear_memory(struct module *module)
+{
+	assert(module != NULL);
+
+	char *error_message = NULL;
+
+	size_t initial = (size_t)module->abi.starting_pages * WASM_PAGE_SIZE;
+	size_t max     = (size_t)module->abi.max_pages * WASM_PAGE_SIZE;
+
+	assert(initial <= (size_t)UINT32_MAX + 1);
+	assert(max <= (size_t)UINT32_MAX + 1);
+
+	struct wasm_memory *linear_memory = (struct wasm_memory *)pool_remove_nolock(
+	  &module->pools.memory[worker_thread_idx]);
+	if (linear_memory == NULL) {
+		linear_memory = wasm_memory_allocate(initial, max);
+		if (unlikely(linear_memory == NULL)) return NULL;
+	}
+
+	return linear_memory;
+}
+
+static inline void
+module_free_linear_memory(struct module *module, struct wasm_memory *memory)
+{
+	wasm_memory_wipe(memory);
+	wasm_memory_reinit(memory, module->abi.starting_pages * WASM_PAGE_SIZE);
+	pool_add_nolock(&module->pools.memory[worker_thread_idx], memory);
+}
+
 /********************************
  * Public Methods from module.c *
  *******************************/
 
-void module_free(struct module *module);
-struct module *
-    module_new(char *mod_name, char *mod_path, uint32_t stack_sz, uint32_t max_heap, uint32_t relative_deadline_us,
-               int port, int req_sz, int resp_sz, int admissions_percentile, uint32_t expected_execution_us);
-int module_new_from_json(char *filename);
+void           module_free(struct module *module);
+struct module *module_new(char *mod_name, char *mod_path, uint32_t stack_sz, uint32_t relative_deadline_us, int port,
+                          int req_sz, int resp_sz, int admissions_percentile, uint32_t expected_execution_us);
+int            module_new_from_json(char *filename);

runtime/include/pool.h

@@ -0,0 +1,119 @@
+#pragma once
+
+#include <assert.h>
+#include <stdbool.h>
+#include <stdlib.h>
+
+#include "generic_thread.h"
+#include "lock.h"
+#include "ps_list.h"
+#include "wasm_memory.h"
+
+struct pool {
+	bool                use_lock;
+	lock_t              lock;
+	struct ps_list_head list;
+};
+
+static inline bool
+pool_is_empty(struct pool *self)
+{
+	assert(self != NULL);
+
+	return ps_list_head_empty(&self->list);
+}
+
+static inline void
+pool_init(struct pool *self, bool use_lock)
+{
+	ps_list_head_init(&self->list);
+	self->use_lock = use_lock;
+	if (use_lock) LOCK_INIT(&self->lock);
+}
+
+static inline void
+pool_deinit(struct pool *self)
+{
+	if (pool_is_empty(self)) return;
+
+	struct wasm_memory *iterator = NULL;
+	struct wasm_memory *buffer   = NULL;
+
+	ps_list_foreach_del_d(&self->list, iterator, buffer)
+	{
+		ps_list_rem_d(iterator);
+		wasm_memory_free(iterator);
+	}
+}
+
+static inline struct wasm_memory *
+pool_remove_nolock(struct pool *self)
+{
+	assert(self != NULL);
+	assert(!self->use_lock || LOCK_IS_LOCKED(&self->lock));
+
+	struct wasm_memory *obj = NULL;
+
+	if (pool_is_empty(self)) return obj;
+
+	obj = ps_list_head_first_d(&self->list, struct wasm_memory);
+	assert(obj);
+	ps_list_rem_d(obj);
+
+	return obj;
+}
+
+static inline struct wasm_memory *
+pool_remove(struct pool *self)
+{
+	assert(self != NULL);
+	assert(self->use_lock);
+
+	struct wasm_memory *obj = NULL;
+
+	if (pool_is_empty(self)) return obj;
+
+	LOCK_LOCK(&self->lock);
+	if (pool_is_empty(self)) {
+		LOCK_UNLOCK(&self->lock);
+		return obj;
+	}
+
+	obj = ps_list_head_first_d(&self->list, struct wasm_memory);
+	assert(obj);
+	ps_list_rem_d(obj);
+	LOCK_UNLOCK(&self->lock);
+	return obj;
+}
+
+static inline int
+pool_add_nolock(struct pool *self, struct wasm_memory *obj)
+{
+	assert(self != NULL);
+	assert(obj != NULL);
+	assert(!self->use_lock || LOCK_IS_LOCKED(&self->lock));
+
+	ps_list_head_add_d(&self->list, obj);
+	return 0;
+}
+
+static inline int
+pool_add(struct pool *self, struct wasm_memory *obj)
+{
+	assert(self != NULL);
+	assert(obj != NULL);
+	assert(self->use_lock);
+
+	LOCK_LOCK(&self->lock);
+	ps_list_head_add_d(&self->list, obj);
+	LOCK_UNLOCK(&self->lock);
+	return 0;
+}
+
+static inline void
+pool_free(struct pool *self)
+{
+	while (!pool_is_empty(self)) free(pool_remove(self));
+
+	free(self);
+}

runtime/include/sandbox_functions.h

@@ -36,7 +36,9 @@ sandbox_close_http(struct sandbox *sandbox)
 static inline void
 sandbox_free_linear_memory(struct sandbox *sandbox)
 {
-	wasm_memory_free(sandbox->memory);
+	assert(sandbox != NULL);
+	assert(sandbox->memory != NULL);
+	module_free_linear_memory(sandbox->module, sandbox->memory);
 	sandbox->memory = NULL;
 }
 

runtime/include/wasm_memory.h

@@ -7,16 +7,18 @@
 #include <string.h>
 #include <sys/mman.h>
 
+#include "ps_list.h"
 #include "types.h" /* PAGE_SIZE */
 #include "wasm_types.h"
 
 #define WASM_MEMORY_MAX (size_t) UINT32_MAX + 1
 
 struct wasm_memory {
-	size_t  size;     /* Initial Size in bytes */
-	size_t  capacity; /* Size backed by actual pages */
-	size_t  max;      /* Soft cap in bytes. Defaults to 4GB */
-	uint8_t data[];
+	struct ps_list list;     /* Linked List Node used for object pool */
+	size_t         size;     /* Initial Size in bytes */
+	size_t         capacity; /* Size backed by actual pages */
+	size_t         max;      /* Soft cap in bytes. Defaults to 4GB */
+	uint8_t        data[];
 };
 
 static inline struct wasm_memory *
@@ -48,6 +50,7 @@ wasm_memory_allocate(size_t initial, size_t max)
 		return NULL;
 	}
 
+	ps_list_init_d(self);
 	self->size     = initial;
 	self->capacity = initial;
 	self->max      = max;
@@ -67,6 +70,13 @@ wasm_memory_wipe(struct wasm_memory *self)
 	memset(self->data, 0, self->size);
 }
 
+static inline void
+wasm_memory_reinit(struct wasm_memory *self, size_t initial)
+{
+	wasm_memory_wipe(self);
+	self->size = initial;
+}
+
 static inline int
 wasm_memory_expand(struct wasm_memory *self, size_t size_to_expand)
 {

runtime/src/module.c

@@ -135,7 +135,6 @@ module_free(struct module *module)
  * @param name
  * @param path
  * @param stack_size
- * @param max_memory
  * @param relative_deadline_us
  * @param port
  * @param request_size
@@ -143,8 +142,8 @@ module_free(struct module *module)
  */
 
 struct module *
-module_new(char *name, char *path, uint32_t stack_size, uint32_t max_memory, uint32_t relative_deadline_us, int port,
-           int request_size, int response_size, int admissions_percentile, uint32_t expected_execution_us)
+module_new(char *name, char *path, uint32_t stack_size, uint32_t relative_deadline_us, int port, int request_size,
+           int response_size, int admissions_percentile, uint32_t expected_execution_us)
 {
 	int rc = 0;
 
@@ -165,7 +164,6 @@ module_new(char *name, char *path, uint32_t stack_size, uint32_t max_memory, uin
 
 	module->stack_size = ((uint32_t)(round_up_to_page(stack_size == 0 ? WASM_STACK_SIZE : stack_size)));
 	debuglog("Stack Size: %u", module->stack_size);
-	module->max_memory        = max_memory == 0 ? ((uint64_t)WASM_PAGE_SIZE * WASM_MEMORY_PAGES_MAX) : max_memory;
 	module->socket_descriptor = -1;
 	module->port              = port;
 
@@ -212,6 +210,8 @@ module_new(char *name, char *path, uint32_t stack_size, uint32_t max_memory, uin
 	module_initialize_table(module);
 	current_wasm_module_instance.table = NULL;
 
+	for (int i = 0; i < MAX_WORKER_THREADS; i++) { pool_init(&module->pools.memory[i], false); }
+
 	/* Start listening for requests */
 	rc = module_listen(module);
 	if (rc < 0) goto err_listen;
@@ -421,7 +421,7 @@ module_new_from_json(char *file_name)
 #endif
 
 		/* Allocate a module based on the values from the JSON */
-		struct module *module = module_new(module_name, module_path, 0, 0, relative_deadline_us, port,
+		struct module *module = module_new(module_name, module_path, 0, relative_deadline_us, port,
 		                                   request_size, response_size, admissions_percentile,
 		                                   expected_execution_us);
 		if (module == NULL) goto module_new_err;

runtime/src/sandbox.c

@@ -5,6 +5,7 @@
 #include "current_sandbox.h"
 #include "debuglog.h"
 #include "panic.h"
+#include "pool.h"
 #include "runtime.h"
 #include "sandbox_functions.h"
 #include "sandbox_set_as_error.h"
@@ -34,16 +35,7 @@ static inline int
 sandbox_allocate_linear_memory(struct sandbox *sandbox)
 {
 	assert(sandbox != NULL);
-
-	char *error_message = NULL;
-
-	size_t initial = (size_t)sandbox->module->abi.starting_pages * WASM_PAGE_SIZE;
-	size_t max     = (size_t)sandbox->module->abi.max_pages * WASM_PAGE_SIZE;
-
-	assert(initial <= (size_t)UINT32_MAX + 1);
-	assert(max <= (size_t)UINT32_MAX + 1);
-
-	sandbox->memory = wasm_memory_allocate(initial, max);
+	sandbox->memory = module_allocate_linear_memory(sandbox->module);
 	if (unlikely(sandbox->memory == NULL)) return -1;
 
 	return 0;