Make default `context.temp_allocator` thread safe when using `package thread`

4 files changed, 41 insertions, 8 deletions

diff --git a/core/mem/allocators.odin b/core/mem/allocators.odin
index f9bffb2cd..5d8260de0 100644
--- a/core/mem/allocators.odin
+++ b/core/mem/allocators.odin
@@ -108,6 +108,18 @@ scratch_allocator_init :: proc(scratch: ^Scratch_Allocator, data: []byte, backup
 	scratch.backup_allocator = backup_allocator;
 }
 
+scratch_allocator_destroy :: proc(using scratch: ^Scratch_Allocator) {
+	if scratch == nil {
+		return;
+	}
+	for ptr in leaked_allocations {
+		free(ptr, backup_allocator);
+	}
+	delete(leaked_allocations);
+	delete(data, backup_allocator);
+	scratch^ = {};
+}
+
 scratch_allocator_proc :: proc(allocator_data: rawptr, mode: Allocator_Mode,
                                size, alignment: int,
                                old_memory: rawptr, old_size: int, flags: u64 = 0, loc := #caller_location) -> rawptr {
diff --git a/core/runtime/core.odin b/core/runtime/core.odin
index e70f2a52f..d09dfd72b 100644
--- a/core/runtime/core.odin
+++ b/core/runtime/core.odin
@@ -253,7 +253,10 @@ Context :: struct {
 	derived:    any, // May be used for derived data types
 }
 
-global_scratch_allocator_data: mem.Scratch_Allocator;
+@thread_local global_scratch_allocator_data: mem.Scratch_Allocator;
+global_scratch_allocator_proc :: mem.scratch_allocator_proc;
+global_scratch_allocator_init :: mem.scratch_allocator_init;
+global_scratch_allocator_destroy :: mem.scratch_allocator_destroy;
 
 
 
@@ -392,7 +395,7 @@ __init_context :: proc "contextless" (c: ^Context) {
 	c.allocator.procedure = os.heap_allocator_proc;
 	c.allocator.data = nil;
 
-	c.temp_allocator.procedure = mem.scratch_allocator_proc;
+	c.temp_allocator.procedure = global_scratch_allocator_proc;
 	c.temp_allocator.data = &global_scratch_allocator_data;
 
 	c.thread_id = os.current_thread_id(); // NOTE(bill): This is "contextless" so it is okay to call
@@ -408,7 +411,7 @@ __init_context :: proc "contextless" (c: ^Context) {
 
 @builtin
 init_global_temporary_allocator :: proc(data: []byte, backup_allocator := context.allocator) {
-	mem.scratch_allocator_init(&global_scratch_allocator_data, data, backup_allocator);
+	global_scratch_allocator_init(&global_scratch_allocator_data, data, backup_allocator);
 }
 
 default_assertion_failure_proc :: proc(prefix, message: string, loc: Source_Code_Location) {
diff --git a/core/thread/thread_unix.odin b/core/thread/thread_unix.odin
index 52b8a81f4..741e52566 100644
--- a/core/thread/thread_unix.odin
+++ b/core/thread/thread_unix.odin
@@ -1,8 +1,9 @@
 // +build linux, darwin
 package thread;
 
-import "core:sys/unix"
+import "core:runtime"
 import "core:sync"
+import "core:sys/unix"
 
 // NOTE(tetra): Aligned here because of core/unix/pthread_linux.odin/pthread_t.
 // Also see core/sys/darwin/mach_darwin.odin/semaphore_t.
@@ -62,6 +63,13 @@ create :: proc(procedure: Thread_Proc, priority := Thread_Priority.Normal) -> ^T
 		context = c;
 
 		t.procedure(t);
+
+		if !t.use_init_context {
+			if context.temp_allocator.data == &runtime.global_scratch_allocator_data {
+				runtime.global_scratch_allocator_destroy(auto_cast context.temp_allocator.data);
+			}
+		}
+
 		sync.atomic_store(&t.done, true, .Sequentially_Consistent);
 		return nil;
 	}
diff --git a/core/thread/thread_windows.odin b/core/thread/thread_windows.odin
index 5b956940f..79ebe48d6 100644
--- a/core/thread/thread_windows.odin
+++ b/core/thread/thread_windows.odin
@@ -1,5 +1,6 @@
 package thread
 
+import "core:runtime"
 import "core:sync"
 import "core:sys/win32"
 
@@ -34,6 +35,13 @@ create :: proc(procedure: Thread_Proc, priority := Thread_Priority.Normal) -> ^T
 		context = c;
 
 		t.procedure(t);
+
+		if !t.use_init_context {
+			if context.temp_allocator.data == &runtime.global_scratch_allocator_data {
+				runtime.global_scratch_allocator_destroy(auto_cast context.temp_allocator.data);
+			}
+		}
+
 		sync.atomic_store(&t.done, true, .Sequentially_Consistent);
 		return 0;
 	}
@@ -69,9 +77,11 @@ is_done :: proc(using thread: ^Thread) -> bool {
 }
 
 join :: proc(using thread: ^Thread) {
-	win32.wait_for_single_object(win32_thread, win32.INFINITE);
-	win32.close_handle(win32_thread);
-	win32_thread = win32.INVALID_HANDLE;
+	if win32_thread != win32.INVALID_HANDLE {
+		win32.wait_for_single_object(win32_thread, win32.INFINITE);
+		win32.close_handle(win32_thread);
+		win32_thread = win32.INVALID_HANDLE;
+	}
 }
 
 destroy :: proc(thread: ^Thread) {
@@ -81,4 +91,4 @@ destroy :: proc(thread: ^Thread) {
 
 terminate :: proc(using thread : ^Thread, exit_code : u32) {
 	win32.terminate_thread(win32_thread, exit_code);
-}
\ No newline at end of file
+}