WIP: use AsyncThrowingStream instead of AsyncThrowingChannel

Author: lhoward

Sources/IORing/IORing.swift

@@ -391,7 +391,7 @@ public final class IORing: CustomStringConvertible {
     moreFlags: UInt32 = 0,
     bufferIndexOrGroup: UInt16 = 0,
     handler: @escaping @Sendable (io_uring_cqe) throws -> T
-  ) throws -> AsyncThrowingChannel<T, Error> {
+  ) throws -> AsyncThrowingStream<T, Error> {
     try MultishotSubmission(
       ring: self,
       opcode,
@@ -616,7 +616,7 @@ private extension IORing {
     fd: FileDescriptorRepresentable,
     count: Int,
     link: Bool = false
-  ) throws -> AsyncThrowingChannel<[UInt8], Error> {
+  ) throws -> AsyncThrowingStream<[UInt8], Error> {
     var buffer = [UInt8]._unsafelyInitialized(count: count)
     return try prepareAndSubmitMultishot(
       .recv,
@@ -654,7 +654,7 @@ private extension IORing {
     count: Int,
     capacity: Int,
     flags: UInt32 = 0
-  ) async throws -> AsyncThrowingChannel<Message, Error> {
+  ) async throws -> AsyncThrowingStream<Message, Error> {
     // FIXME: combine message holder buffer registration with multishot registration to avoid extra system call
     let holder = try await MessageHolder(ring: self, size: count, count: capacity)
     return try await holder.withUnsafeMutablePointer { @IORingActor pointer in
@@ -669,7 +669,7 @@ private extension IORing {
         bufferIndexOrGroup: holder.bufferGroup
       ) { [holder] cqe in
         // because the default for multishots is to resubmit when IORING_CQE_F_MORE is unset,
-        // we don't need to deallocate the buffer here. FIXME: do this when channel closes.
+        // we don't need to deallocate the buffer here. FIXME: do this when stream finishes.
         // we know that handlers are always executed in an @IORingActor actor's execution context
         // so it's safe to access the holder's buffer. But we should make this more explicit
         // by making the callback take an async function.
@@ -715,7 +715,7 @@ private extension IORing {
   func io_uring_op_multishot_accept(
     fd: FileDescriptorRepresentable,
     flags: UInt32 = 0
-  ) throws -> AsyncThrowingChannel<FileDescriptorRepresentable, Error> {
+  ) throws -> AsyncThrowingStream<FileDescriptorRepresentable, Error> {
     try prepareAndSubmitMultishot(
       .accept,
       fd: fd,

Sources/IORing/Submission.swift

@@ -371,7 +371,22 @@ final class BufferSubmission<U>: Submission<()> {
 }
 
 final class MultishotSubmission<T: Sendable>: Submission<T> {
-  private let channel = AsyncThrowingChannel<T, Error>()
+  // Shared holder ensures continuation is accessible across resubmissions
+  private final class StreamHolder: Sendable {
+    let stream: AsyncThrowingStream<T, Error>
+    let continuation: AsyncThrowingStream<T, Error>.Continuation
+
+    init() {
+      var cont: AsyncThrowingStream<T, Error>.Continuation!
+      let str = AsyncThrowingStream<T, Error> { continuation in
+        cont = continuation
+      }
+      stream = str
+      continuation = cont
+    }
+  }
+
+  private let holder: StreamHolder
 
   // state for resubmission
   private let address: UnsafeRawPointer?
@@ -383,7 +398,8 @@ final class MultishotSubmission<T: Sendable>: Submission<T> {
   private let bufferIndexOrGroup: UInt16
   private let socketAddress: sockaddr_storage?
 
-  override init(
+  // Private init that allows specifying a holder (for resubmission)
+  private init(
     ring: IORing,
     _ opcode: IORingOperation,
     fd: FileDescriptorRepresentable,
@@ -395,8 +411,10 @@ final class MultishotSubmission<T: Sendable>: Submission<T> {
     moreFlags: UInt32 = 0,
     bufferIndexOrGroup: UInt16 = 0,
     socketAddress: sockaddr_storage? = nil,
+    holder: StreamHolder?,
     handler: @escaping @Sendable (io_uring_cqe) throws -> T
   ) throws {
+    self.holder = holder ?? StreamHolder()
     self.address = address
     self.length = length
     self.offset = offset
@@ -422,44 +440,73 @@ final class MultishotSubmission<T: Sendable>: Submission<T> {
     )
   }
 
-  private convenience init(_ submission: MultishotSubmission) throws {
+  // Public convenience init (creates new holder/stream)
+  override convenience init(
+    ring: IORing,
+    _ opcode: IORingOperation,
+    fd: FileDescriptorRepresentable,
+    address: UnsafeRawPointer? = nil,
+    length: CUnsignedInt = 0,
+    offset: IORing.Offset = 0,
+    flags: IORing.SqeFlags = IORing.SqeFlags(),
+    ioprio: UInt16 = 0,
+    moreFlags: UInt32 = 0,
+    bufferIndexOrGroup: UInt16 = 0,
+    socketAddress: sockaddr_storage? = nil,
+    handler: @escaping @Sendable (io_uring_cqe) throws -> T
+  ) throws {
     try self.init(
-      ring: submission.ring,
-      submission.opcode,
-      fd: submission.fd,
-      address: submission.address,
-      length: submission.length,
-      offset: submission.offset,
-      flags: submission.flags,
-      ioprio: submission.ioprio,
-      moreFlags: submission.moreFlags,
-      bufferIndexOrGroup: submission.bufferIndexOrGroup,
-      socketAddress: submission.socketAddress,
-      handler: submission.handler
+      ring: ring,
+      opcode,
+      fd: fd,
+      address: address,
+      length: length,
+      offset: offset,
+      flags: flags,
+      ioprio: ioprio,
+      moreFlags: moreFlags,
+      bufferIndexOrGroup: bufferIndexOrGroup,
+      socketAddress: socketAddress,
+      holder: nil,
+      handler: handler
     )
   }
 
-  func submit() throws -> AsyncThrowingChannel<T, Error> {
+  func submit() throws -> AsyncThrowingStream<T, Error> {
     try ring.submit()
-    return channel
+    return holder.stream
   }
 
   private func resubmit() {
     do {
-      // this will allocate a new SQE with the same channel, fd, opcode and handler
-      let resubmission = try MultishotSubmission(self)
+      // Create new SQE with same holder (shared stream/continuation)
+      let resubmission = try MultishotSubmission(
+        ring: ring,
+        opcode,
+        fd: fd,
+        address: address,
+        length: length,
+        offset: offset,
+        flags: flags,
+        ioprio: ioprio,
+        moreFlags: moreFlags,
+        bufferIndexOrGroup: bufferIndexOrGroup,
+        socketAddress: socketAddress,
+        holder: holder, // Share the same holder!
+        handler: handler
+      )
       IORing.shared.logger.debug("resubmitting multishot submission \(resubmission)")
       _ = try resubmission.submit()
     } catch {
       IORing.shared.logger.debug("resubmitting multishot submission failed: \(error)")
-      channel.fail(error)
+      holder.continuation.finish(throwing: error)
     }
   }
 
   override func onCompletion(cqe: io_uring_cqe) async {
     do {
       let result = try throwingErrno(cqe: cqe, handler)
-      await channel.send(result)
+      holder.continuation.yield(result) // No suspension point!
       if cqe.flags & IORING_CQE_F_MORE == 0 {
         // if IORING_CQE_F_MORE is not set, we need to issue a new request
         // try to do this implictily
@@ -469,7 +516,7 @@ final class MultishotSubmission<T: Sendable>: Submission<T> {
         try await cancel()
       }
     } catch {
-      channel.fail(error)
+      holder.continuation.finish(throwing: error)
     }
   }
 }