WIP: submission protocol

Author: lhoward

Sources/IORing/Submission.swift

@@ -21,24 +21,53 @@ internal import CIOURing
 import Glibc
 import SystemPackage
 
-class Submission<T: Sendable>: CustomStringConvertible, @unchecked Sendable {
+// MARK: - Submission Protocol
+
+protocol Submission: CustomStringConvertible, Sendable, Hashable {
+  associatedtype Result: Sendable
+
+  var core: SubmissionCore<Result> { get }
+
+  func onCompletion(cqe: io_uring_cqe)
+  func onCancel(cqe: io_uring_cqe)
+}
+
+extension Submission {
+  var ring: IORing { core.ring }
+  var fd: FileDescriptorRepresentable { core.fd }
+  var opcode: IORingOperation { core.opcode }
+
+  nonisolated var description: String {
+    "(\(type(of: self)))(fd: \(fd.fileDescriptor), opcode: \(opcode), handler: \(String(describing: core.handler)))"
+  }
+
+  func cancel(ring: isolated IORing) throws {
+    try core.cancel(ring: ring, onCancel: onCancel)
+  }
+
+  func throwingErrno(
+    cqe: io_uring_cqe,
+    _ body: @escaping @Sendable (_: io_uring_cqe) throws -> Result
+  ) throws -> Result {
+    try core.throwingErrno(cqe: cqe, body)
+  }
+}
+
+// MARK: - Submission Core
+
+final class SubmissionCore<T: Sendable>: @unchecked Sendable {
   // reference to owner which owns ring
   let ring: IORing
   /// user-supplied callback to transform a completion queue entry to a result
-  fileprivate let handler: @Sendable (io_uring_cqe) throws -> T
+  let handler: @Sendable (io_uring_cqe) throws -> T
   /// file descriptor, stored so that it is not closed before the completion handler is run
-  fileprivate let fd: FileDescriptorRepresentable
-
+  let fd: FileDescriptorRepresentable
   /// opcode, useful for debugging
-  fileprivate let opcode: IORingOperation
+  let opcode: IORingOperation
   /// assigned submission queue entry for this object
   private let sqe: UnsafeMutablePointer<io_uring_sqe>
   private var cancellationToken: UnsafeMutableRawPointer?
 
-  nonisolated var description: String {
-    "(\(type(of: self)))(fd: \(fd.fileDescriptor), opcode: \(opcode), handler: \(String(describing: handler)))"
-  }
-
   private func prepare(
     _ opcode: IORingOperation,
     sqe: UnsafeMutablePointer<io_uring_sqe>,
@@ -86,24 +115,24 @@ class Submission<T: Sendable>: CustomStringConvertible, @unchecked Sendable {
   /// because actors are reentrant, `setBlock()` must be called immediately after
   /// the io_uring assigned a SQE (or, at least before any suspension point)
   /// FIXME: `swift_allocObject()` here appears to be a potential performance issue
-  private func setBlock() {
+  private func setBlock(onCompletion: @escaping (io_uring_cqe) -> ()) {
     cancellationToken = io_uring_sqe_set_block(sqe) { cqe in
       let cqe = cqe.pointee
-      self.onCompletion(cqe: cqe)
+      onCompletion(cqe)
     }
   }
 
-  func cancel(ring: isolated IORing) throws {
+  func cancel(ring: isolated IORing, onCancel: @escaping (io_uring_cqe) -> ()) throws {
     do {
       precondition(cancellationToken != nil)
       let sqe = try ring.getSqe()
       io_uring_prep_cancel(sqe, cancellationToken, AsyncCancelFlags.userData.rawValue)
       _ = io_uring_sqe_set_block(sqe) { cqe in
-        self.onCancel(cqe: cqe.pointee)
+        onCancel(cqe.pointee)
       }
       try ring.submit()
     } catch {
-      IORing.shared.logger.debug("failed to cancel submission \(self)")
+      IORing.shared.logger.debug("failed to cancel submission")
       throw error
     }
   }
@@ -120,7 +149,8 @@ class Submission<T: Sendable>: CustomStringConvertible, @unchecked Sendable {
     moreFlags: UInt32 = 0,
     bufferIndexOrGroup: UInt16 = 0,
     socketAddress: sockaddr_storage? = nil,
-    handler: @escaping @Sendable (io_uring_cqe) throws -> T
+    handler: @escaping @Sendable (io_uring_cqe) throws -> T,
+    onCompletion: @escaping (io_uring_cqe) -> ()
   ) throws {
     sqe = try ring.getSqe()
     self.ring = ring
@@ -140,12 +170,9 @@ class Submission<T: Sendable>: CustomStringConvertible, @unchecked Sendable {
         try setSocketAddress(sqe: sqe, socketAddress: socketAddress)
       }
     }
-    setBlock()
+    setBlock(onCompletion: onCompletion)
   }
 
-  func onCompletion(cqe: io_uring_cqe) { fatalError("must be implemented by concrete class") }
-  func onCancel(cqe: io_uring_cqe) {}
-
   func throwingErrno(
     cqe: io_uring_cqe,
     _ body: @escaping @Sendable (_: io_uring_cqe) throws -> T
@@ -155,7 +182,7 @@ class Submission<T: Sendable>: CustomStringConvertible, @unchecked Sendable {
       if error != .brokenPipe {
         IORing.shared.logger
           .debug(
-            "\(type(of: self)) completion fileDescriptor: \(fd) opcode: \(opcode) error: \(Errno(rawValue: -cqe.res))"
+            "completion fileDescriptor: \(fd) opcode: \(opcode) error: \(Errno(rawValue: -cqe.res))"
           )
       }
       throw error
@@ -164,11 +191,19 @@ class Submission<T: Sendable>: CustomStringConvertible, @unchecked Sendable {
   }
 }
 
-final class SingleshotSubmission<T: Sendable>: Submission<T>, @unchecked Sendable {
+// MARK: - SingleshotSubmission
+
+struct SingleshotSubmission<T: Sendable>: Submission, @unchecked Sendable {
+  typealias Result = T
+
+  let core: SubmissionCore<T>
   weak var group: SubmissionGroup<T>?
 
-  private typealias Continuation = UnsafeContinuation<T, Error>
-  private var continuation: Continuation!
+  private final class ContinuationHolder: @unchecked Sendable {
+    var continuation: UnsafeContinuation<T, Error>?
+  }
+
+  private let holder: ContinuationHolder
 
   init(
     ring: isolated IORing,
@@ -186,7 +221,9 @@ final class SingleshotSubmission<T: Sendable>: Submission<T>, @unchecked Sendabl
     handler: @escaping @Sendable (io_uring_cqe) throws -> T
   ) async throws {
     self.group = group
-    try super.init(
+    holder = ContinuationHolder()
+
+    core = try SubmissionCore(
       ring: ring,
       opcode,
       fd: fd,
@@ -198,8 +235,17 @@ final class SingleshotSubmission<T: Sendable>: Submission<T>, @unchecked Sendabl
       moreFlags: moreFlags,
       bufferIndexOrGroup: bufferIndex,
       socketAddress: socketAddress,
-      handler: handler
+      handler: handler,
+      onCompletion: { [holder] cqe in
+        guard let continuation = holder.continuation else { return }
+        do {
+          try continuation.resume(returning: handler(cqe))
+        } catch {
+          continuation.resume(throwing: error)
+        }
+      }
     )
+
     if let group {
       await group.enqueue(submission: self)
     }
@@ -209,7 +255,7 @@ final class SingleshotSubmission<T: Sendable>: Submission<T>, @unchecked Sendabl
     try await withTaskCancellationHandler(operation: {
       try await withUnsafeThrowingContinuation { continuation in
         // guaranteed to run immediately
-        self.continuation = continuation
+        holder.continuation = continuation
         if group != nil {
           Task { await ready() }
         } else {
@@ -226,13 +272,16 @@ final class SingleshotSubmission<T: Sendable>: Submission<T>, @unchecked Sendabl
     try await _submit(ring: ring)
   }
 
-  override func onCompletion(cqe: io_uring_cqe) {
+  func onCompletion(cqe: io_uring_cqe) {
+    guard let continuation = holder.continuation else { return }
     do {
-      try continuation.resume(returning: throwingErrno(cqe: cqe, handler))
+      try continuation.resume(returning: throwingErrno(cqe: cqe, core.handler))
     } catch {
       continuation.resume(throwing: error)
     }
   }
+
+  func onCancel(cqe: io_uring_cqe) {}
 }
 
 struct BufferCount: FileDescriptorRepresentable {
@@ -243,7 +292,13 @@ struct BufferCount: FileDescriptorRepresentable {
   }
 }
 
-final class BufferSubmission<U>: Submission<()>, @unchecked Sendable {
+// MARK: - BufferSubmission
+
+final class BufferSubmission<U>: Submission, @unchecked Sendable {
+  typealias Result = ()
+
+  let core: SubmissionCore<()>
+
   nonisolated var count: Int {
     Int(fd.fileDescriptor)
   }
@@ -253,7 +308,8 @@ final class BufferSubmission<U>: Submission<()>, @unchecked Sendable {
   let buffer: UnsafeMutablePointer<U>
   let deallocate: Bool
 
-  override func onCompletion(cqe: io_uring_cqe) {}
+  func onCompletion(cqe: io_uring_cqe) {}
+  func onCancel(cqe: io_uring_cqe) {}
 
   private func _submit(ring: isolated IORing) throws {
     try ring.submit()
@@ -287,16 +343,18 @@ final class BufferSubmission<U>: Submission<()>, @unchecked Sendable {
     self.deallocate = deallocate
     self.buffer = buffer
 
-    try super.init(
+    core = try SubmissionCore(
       ring: ring,
       .provide_buffers,
       fd: BufferCount(count: count),
       address: buffer,
       length: UInt32(size),
       offset: offset,
       flags: flags,
-      bufferIndexOrGroup: bufferGroup
-    ) { _ in }
+      bufferIndexOrGroup: bufferGroup,
+      handler: { _ in },
+      onCompletion: { _ in }
+    )
   }
 
   convenience init(
@@ -382,7 +440,13 @@ final class BufferSubmission<U>: Submission<()>, @unchecked Sendable {
   }
 }
 
-final class MultishotSubmission<T: Sendable>: Submission<T>, @unchecked Sendable {
+// MARK: - MultishotSubmission
+
+final class MultishotSubmission<T: Sendable>: Submission, @unchecked Sendable {
+  typealias Result = T
+
+  let core: SubmissionCore<T>
+
   // Shared holder ensures continuation is accessible across resubmissions
   private final class _StreamHolder: Sendable {
     let stream: AsyncThrowingStream<T, Error>
@@ -434,7 +498,7 @@ final class MultishotSubmission<T: Sendable>: Submission<T>, @unchecked Sendable
     self.socketAddress = socketAddress
     self.holder = holder
 
-    try super.init(
+    core = try SubmissionCore(
       ring: ring,
       opcode,
       fd: fd,
@@ -446,7 +510,15 @@ final class MultishotSubmission<T: Sendable>: Submission<T>, @unchecked Sendable
       moreFlags: moreFlags,
       bufferIndexOrGroup: bufferIndexOrGroup,
       socketAddress: socketAddress,
-      handler: handler
+      handler: handler,
+      onCompletion: { [holder] cqe in
+        do {
+          let result = try handler(cqe)
+          holder.continuation.yield(result)
+        } catch {
+          holder.continuation.finish(throwing: error)
+        }
+      }
     )
   }
 
@@ -464,11 +536,11 @@ final class MultishotSubmission<T: Sendable>: Submission<T>, @unchecked Sendable
       bufferIndexOrGroup: submission.bufferIndexOrGroup,
       socketAddress: submission.socketAddress,
       holder: submission.holder,
-      handler: submission.handler
+      handler: submission.core.handler
     )
   }
 
-  override convenience init(
+  convenience init(
     ring: isolated IORing,
     _ opcode: IORingOperation,
     fd: FileDescriptorRepresentable,
@@ -522,22 +594,24 @@ final class MultishotSubmission<T: Sendable>: Submission<T>, @unchecked Sendable
     }
   }
 
-  override func onCompletion(cqe: io_uring_cqe) {
+  func onCompletion(cqe: io_uring_cqe) {
     do {
-      let result = try throwingErrno(cqe: cqe, handler)
+      let result = try throwingErrno(cqe: cqe, core.handler)
       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
         Task { await resubmit(ring: ring) }
       }
       if Task.isCancelled {
-        Task { try await cancel(ring: ring) }
+        Task { try? await cancel(ring: ring) }
       }
     } catch {
       holder.continuation.finish(throwing: error)
     }
   }
+
+  func onCancel(cqe: io_uring_cqe) {}
 }
 
 enum IORingOperation: UInt32 {
@@ -609,13 +683,33 @@ struct AsyncCancelFlags: OptionSet {
   static let op = AsyncCancelFlags(rawValue: 1 << 5)
 }
 
-extension Submission: Equatable {
-  nonisolated static func == (lhs: Submission, rhs: Submission) -> Bool {
+// MARK: - Equatable and Hashable conformances
+
+extension SingleshotSubmission: Equatable {
+  nonisolated static func == (lhs: SingleshotSubmission, rhs: SingleshotSubmission) -> Bool {
+    lhs.core === rhs.core
+  }
+
+  nonisolated func hash(into hasher: inout Hasher) {
+    ObjectIdentifier(core).hash(into: &hasher)
+  }
+}
+
+extension BufferSubmission: Equatable {
+  nonisolated static func == (lhs: BufferSubmission, rhs: BufferSubmission) -> Bool {
     lhs === rhs
   }
+
+  nonisolated func hash(into hasher: inout Hasher) {
+    ObjectIdentifier(self).hash(into: &hasher)
+  }
 }
 
-extension Submission: Hashable {
+extension MultishotSubmission: Equatable {
+  nonisolated static func == (lhs: MultishotSubmission, rhs: MultishotSubmission) -> Bool {
+    lhs === rhs
+  }
+
   nonisolated func hash(into hasher: inout Hasher) {
     ObjectIdentifier(self).hash(into: &hasher)
   }