Commit e06ab1b0107e8a6a1720703a6df0f61f535b5e5a | zig

author	Luuk de Gram <luuk@degram.dev> 2023-06-23 19:14:55+02:00
committer	Luuk de Gram <luuk@degram.dev> 2023-06-26 20:00:58+02:00
log	e06ab1b0107e8a6a1720703a6df0f61f535b5e5a
tree	822263fe88040667a7482175c703f7e40dad8e86
parent	622b7c47468bc4508f4cfe840e0f8c51b54949dd
signature	Commit is signed but in an unrecognized format.

std: implement `detach` for WASI-threads

When a thread is detached from the main thread, we automatically
cleanup any allocated memory. For this we first reset the stack-pointer
to the original stack-pointer of the main-thread so we can safely clear
the memory which also contains the thread's stack.

2 files changed, 64 insertions(+), 27 deletions(-)

lib
- std
  - Thread.zig
  - Thread
    Futex.zig

lib/std/Thread.zig+62-25

...	...	@@ -757,6 +757,8 @@ const WasiThreadImpl = struct {
757	757	/// The allocator used to allocate the thread's memory,
758	758	/// which is also used during `join` to ensure clean-up.
759	759	allocator: std.mem.Allocator,
	760	/// The current state of the thread.
	761	state: State = State.init(.running),
760	762	};
761	763
762	764	/// A meta-data structure used to bootstrap a thread
...	...	@@ -775,8 +777,15 @@ const WasiThreadImpl = struct {
775	777	/// function upon thread spawn. The above mentioned pointer will be passed
776	778	/// to this function pointer as its argument.
777	779	call_back: *const fn (usize) void,
	780	/// When a thread is in `detached` state, we must free all of its memory
	781	/// upon thread completion. However, as this is done while still within
	782	/// the thread, we must first jump back to the main thread's stack or else
	783	/// we end up freeing the stack that we're currently using.
	784	original_stack_pointer: [*]u8,
778	785	};
779	786
	787	const State = Atomic(enum(u8) { running, completed, detached });
	788
780	789	fn getCurrentId() Id {
781	790	return tls_thread_id;
782	791	}
...	...	@@ -786,7 +795,11 @@ const WasiThreadImpl = struct {
786	795	}
787	796
788	797	fn detach(self: Impl) void {
789		_ = self;
	798	switch (self.thread.state.swap(.detached, .SeqCst)) {
	799	.running => {},
	800	.completed => self.join(),
	801	.detached => unreachable,
	802	}
790	803	}
791	804
792	805	fn join(self: Impl) void {
...	...	@@ -836,7 +849,7 @@ const WasiThreadImpl = struct {
836	849	const Wrapper = struct {
837	850	args: @TypeOf(args),
838	851	fn entry(ptr: usize) void {
839		const w = @intToPtr(*@This(), ptr);
	852	const w: *@This() = @ptrFromInt(ptr);
840	853	@call(.auto, f, w.args);
841	854	}
842	855	};
...	...	@@ -854,7 +867,7 @@ const WasiThreadImpl = struct {
854	867	// start with atleast a single page, which is used as a guard to prevent
855	868	// other threads clobbering our new thread.
856	869	// Unfortunately, WebAssembly has no notion of read-only segments, so this
857		// is only a temporary measure until the entire page is "run over".
	870	// is only a best effort.
858	871	var bytes: usize = std.wasm.page_size;
859	872
860	873	bytes = std.mem.alignForward(usize, bytes, 16); // align stack to 16 bytes
...	...	@@ -880,16 +893,17 @@ const WasiThreadImpl = struct {
880	893	// Allocate the amount of memory required for all meta data.
881	894	const allocated_memory = try config.allocator.?.alloc(u8, map_bytes);
882	895
883		const wrapper = @ptrCast(*Wrapper, @alignCast(@alignOf(Wrapper), &allocated_memory[wrapper_offset]));
	896	const wrapper: *Wrapper = @ptrCast(@alignCast(&allocated_memory[wrapper_offset]));
884	897	wrapper.* = .{ .args = args };
885	898
886		const instance = @ptrCast(*Instance, @alignCast(@alignOf(Instance), &allocated_memory[instance_offset]));
	899	const instance: *Instance = @ptrCast(@alignCast(&allocated_memory[instance_offset]));
887	900	instance.* = .{
888	901	.thread = .{ .memory = allocated_memory, .allocator = config.allocator.? },
889	902	.tls_offset = tls_offset,
890	903	.stack_offset = stack_offset,
891		.raw_ptr = @ptrToInt(wrapper),
	904	.raw_ptr = @intFromPtr(wrapper),
892	905	.call_back = &Wrapper.entry,
	906	.original_stack_pointer = __get_stack_pointer(),
893	907	};
894	908
895	909	const tid = spawnWasiThread(instance);
...	...	@@ -903,32 +917,46 @@ const WasiThreadImpl = struct {
903	917	return .{ .thread = &instance.thread };
904	918	}
905	919
906		/// Bootstrap procedure, called by the HOST environment after thread creation.
	920	/// Bootstrap procedure, called by the host environment after thread creation.
907	921	export fn wasi_thread_start(tid: i32, arg: *Instance) void {
908	922	__set_stack_pointer(arg.thread.memory.ptr + arg.stack_offset);
909	923	__wasm_init_tls(arg.thread.memory.ptr + arg.tls_offset);
910		WasiThreadImpl.tls_thread_id = @intCast(u32, tid);
	924	@atomicStore(u32, &WasiThreadImpl.tls_thread_id, @intCast(tid), .SeqCst);
911	925
912	926	// Finished bootstrapping, call user's procedure.
913	927	arg.call_back(arg.raw_ptr);
914	928
915		// Thread finished. Reset Thread ID and wake up the main thread if needed.
916		// We use inline assembly here as we must ensure not to use the stack.
917		asm volatile (
918		\\ local.get %[ptr]
919		\\ i32.const 0
920		\\ i32.atomic.store 0
921		:
922		: [ptr] "r" (&arg.thread.tid.value),
923		);
924		asm volatile (
925		\\ local.get %[ptr]
926		\\ i32.const 1 # waiters
927		\\ memory.atomic.notify 0
928		\\ drop # no need to know the waiters
929		:
930		: [ptr] "r" (&arg.thread.tid.value),
931		);
	929	switch (arg.thread.state.swap(.completed, .SeqCst)) {
	930	.running => {
	931	// reset the Thread ID
	932	asm volatile (
	933	\\ local.get %[ptr]
	934	\\ i32.const 0
	935	\\ i32.atomic.store 0
	936	:
	937	: [ptr] "r" (&arg.thread.tid.value),
	938	);
	939
	940	// Wake the main thread listening to this thread
	941	asm volatile (
	942	\\ local.get %[ptr]
	943	\\ i32.const 1 # waiters
	944	\\ memory.atomic.notify 0
	945	\\ drop # no need to know the waiters
	946	:
	947	: [ptr] "r" (&arg.thread.tid.value),
	948	);
	949	},
	950	.completed => unreachable,
	951	.detached => {
	952	// restore the original stack pointer so we can free the memory
	953	// without having to worry about freeing the stack
	954	__set_stack_pointer(arg.original_stack_pointer);
	955	// Ensure a copy so we don't free the allocator reference itself
	956	var allocator = arg.thread.allocator;
	957	allocator.free(arg.thread.memory);
	958	},
	959	}
932	960	}
933	961
934	962	/// Asks the host to create a new thread for us.
...	...	@@ -980,6 +1008,15 @@ const WasiThreadImpl = struct {
980	1008	: [ptr] "r" (addr),
981	1009	);
982	1010	}
	1011
	1012	/// Returns the current value of the stack pointer
	1013	inline fn __get_stack_pointer() [*]u8 {
	1014	return asm (
	1015	\\ global.get __stack_pointer
	1016	\\ local.set %[stack_ptr]
	1017	: [stack_ptr] "=r" (-> [*]u8),
	1018	);
	1019	}
983	1020	};
984	1021
985	1022	const LinuxThreadImpl = struct {

lib/std/Thread/Futex.zig+2-2

...	...	@@ -453,7 +453,7 @@ const WasmImpl = struct {
453	453	if (!comptime std.Target.wasm.featureSetHas(builtin.target.cpu.features, .atomics)) {
454	454	@compileError("WASI target missing cpu feature 'atomics'");
455	455	}
456		const to: i64 = if (timeout) \|to\| @intCast(i64, to) else -1;
	456	const to: i64 = if (timeout) \|to\| @intCast(to) else -1;
457	457	const result = asm (
458	458	\\local.get %[ptr]
459	459	\\local.get %[expected]
...	...	@@ -462,7 +462,7 @@ const WasmImpl = struct {
462	462	\\local.set %[ret]
463	463	: [ret] "=r" (-> u32),
464	464	: [ptr] "r" (&ptr.value),
465		[expected] "r" (@bitCast(i32, expect)),
	465	[expected] "r" (@as(i32, @bitCast(expect))),
466	466	[timeout] "r" (to),
467	467	);
468	468	switch (result) {