diff --git a/README.md b/README.md index fdb5375..5d8fb5b 100644 --- a/README.md +++ b/README.md @@ -57,6 +57,9 @@ Say we `accept` a connection in one thread, and want to send the file descriptor to another for handling. ```zig +// Spawn a thread with a queue of 16 entries. When this function returns, the +// the thread is idle and waiting to receive tasks via msgRing +const thread = main_rt.spawnThread(16); const target_task = try main_rt.getTask(); target_task.* { .userdata = &foo, @@ -65,11 +68,12 @@ target_task.* { .req = .{ .userfd = fd }, }; -// Send target_task from the main_rt thread to the thread_rt Ring. The +// Send target_task from the main_rt thread to the thread Ring. The // thread_rt Ring will then // process the task as a completion, ie -// Worker.onCompletion will be called with // this task. That thread can then -// schedule a recv, a write, etc on the file // descriptor it just received. -_ = try main_rt.msgRing(thread_rt, target_task, .{}); +// Worker.onCompletion will be called with this task. That thread can then +// schedule a recv, a write, etc on the file descriptor it just received. Or do +// arbitrary work +_ = try main_rt.msgRing(&thread.ring, target_task, .{}); ``` ### Multiple Rings on the same thread @@ -84,6 +88,7 @@ _ = try rt2.poll(fd, .{ .cb = onCompletion, .msg = @intFromEnum(Msg.rt1_has_completions)} ); + ``` ## Example diff --git a/src/Uring.zig b/src/Uring.zig index 62d7354..9f46cf5 100644 --- a/src/Uring.zig +++ b/src/Uring.zig @@ -346,19 +346,21 @@ pub fn reapCompletions(self: *Uring, rt: *io.Ring) anyerror!void { .userfd, .userptr => unreachable, }; - try task.callback(rt, task.*); - - if (cqe.flags & msg_ring_received_cqe != 0) { - // This message was received from another ring. We don't decrement inflight for this. - // But we do need to set the task as free because we will add it to our free list - rt.free_q.push(task); - } else if (cqe.flags & linux.IORING_CQE_F_MORE == 0) { - // If the cqe doesn't have IORING_CQE_F_MORE set, then this task is complete and free to - // be rescheduled - task.state = .complete; - self.in_flight.remove(task); - rt.free_q.push(task); + defer { + if (cqe.flags & msg_ring_received_cqe != 0) { + // This message was received from another ring. We don't decrement inflight for this. + // But we do need to set the task as free because we will add it to our free list + rt.free_q.push(task); + } else if (cqe.flags & linux.IORING_CQE_F_MORE == 0) { + // If the cqe doesn't have IORING_CQE_F_MORE set, then this task is complete and free to + // be rescheduled + task.state = .complete; + self.in_flight.remove(task); + rt.free_q.push(task); + } } + + try task.callback(rt, task.*); } } diff --git a/src/main.zig b/src/main.zig index 276131e..42e379b 100644 --- a/src/main.zig +++ b/src/main.zig @@ -394,6 +394,70 @@ pub const Ring = struct { self.submission_q.push(task); return task; } + + /// Spawns a thread with a Ring instance. The thread will be idle and waiting to receive work + /// via msgRing when this function returns. Call kill on the returned thread to signal it to + /// shutdown. + pub fn spawnThread(self: *Ring, entries: u16) !*Thread { + const thread = try self.gpa.create(Thread); + errdefer self.gpa.destroy(thread); + + var wg: std.Thread.WaitGroup = .{}; + wg.start(); + thread.thread = try std.Thread.spawn(.{}, Thread.run, .{ thread, self, &wg, entries }); + wg.wait(); + + return thread; + } +}; + +pub const Thread = struct { + thread: std.Thread, + ring: io.Ring = undefined, + + pub const Msg = enum { + kill, + }; + + pub fn run(self: *Thread, parent: *io.Ring, wg: *std.Thread.WaitGroup, entries: u16) !void { + self.ring = try parent.initChild(entries); + wg.finish(); + + defer self.ring.deinit(); + + // Run forever, because we may not start with a task. Inter-thread messaging means we could + // receive work at any time + self.ring.run(.forever) catch |err| { + switch (err) { + error.ThreadKilled => return, + else => return err, + } + }; + } + + /// Kill sends a message to the thread telling it to exit. Callers of this thread can safely + /// join and deinit the Thread in the Context callback + pub fn kill(self: *Thread, rt: *io.Ring, ctx: Context) Allocator.Error!*io.Task { + const target_task = try rt.getTask(); + target_task.* = .{ + .userdata = self, + .msg = @intFromEnum(Thread.Msg.kill), + .callback = Thread.onCompletion, + .result = .noop, + }; + + return rt.msgRing(&self.ring, target_task, ctx); + } + + pub fn join(self: Thread) void { + self.thread.join(); + } + + fn onCompletion(_: *io.Ring, task: Task) anyerror!void { + switch (task.msgToEnum(Thread.Msg)) { + .kill => return error.ThreadKilled, + } + } }; pub const Op = enum { @@ -715,3 +779,58 @@ test "runtime: msgRing" { try std.testing.expect(foo.rt1); try std.testing.expect(foo.rt2); } + +test "runtime: spawnThread" { + const gpa = std.testing.allocator; + var rt = try io.Ring.init(gpa, 16); + defer rt.deinit(); + + const thread = try rt.spawnThread(4); + + const Foo2 = struct { + kill: bool = false, + did_work: bool = false, + + gpa: Allocator, + thread: *Thread, + + const Msg = enum { kill, work }; + + fn callback(_: *io.Ring, task: io.Task) anyerror!void { + const self = task.userdataCast(@This()); + const msg = task.msgToEnum(Msg); + switch (msg) { + .kill => { + self.kill = true; + self.thread.join(); + self.gpa.destroy(self.thread); + }, + .work => self.did_work = true, + } + } + }; + + var foo: Foo2 = .{ .thread = thread, .gpa = gpa }; + + // Send work to the thread + const target_task = try rt.getTask(); + target_task.* = .{ + .userdata = &foo, + .callback = Foo2.callback, + .msg = @intFromEnum(Foo2.Msg.work), + .result = .{ .usermsg = 0 }, + }; + + _ = try rt.msgRing(&thread.ring, target_task, .{}); + + try rt.run(.until_done); + _ = try thread.kill(&rt, .{ + .ptr = &foo, + .cb = Foo2.callback, + .msg = @intFromEnum(Foo2.Msg.kill), + }); + try rt.run(.until_done); + + try std.testing.expect(foo.did_work); + try std.testing.expect(foo.kill); +}