More refactoring, and working timeout support on ManagedIORing

Catfish-Man · Catfish-Man · commit 7ea32aefce1f · 2024-12-11T23:00:50.000Z
diff --git a/Sources/System/AsyncFileDescriptor.swift b/Sources/System/AsyncFileDescriptor.swift
@@ -21,8 +21,7 @@ public struct AsyncFileDescriptor: ~Copyable {
         let cstr = path.withCString {
             return $0  // bad
         }
-        let res = await ring.submitAndWait(
-            IORequest(
+        let res = try await ring.submit(request: IORequest(
                 opening: cstr, 
                 in: directory, 
                 into: fileSlot, 
@@ -46,7 +45,7 @@ public struct AsyncFileDescriptor: ~Copyable {
 
     @inlinable @inline(__always)
     public consuming func close(isolation actor: isolated (any Actor)? = #isolation) async throws {
-        let res = await ring.submitAndWait(IORequest(closing: fileSlot))
+        let res = try await ring.submit(request: IORequest(closing: fileSlot))
         if res.result < 0 {
             throw Errno(rawValue: -res.result)
         }
@@ -59,8 +58,7 @@ public struct AsyncFileDescriptor: ~Copyable {
         atAbsoluteOffset offset: UInt64 = UInt64.max,
         isolation actor: isolated (any Actor)? = #isolation
     ) async throws -> UInt32 {
-        let res = await ring.submitAndWait(
-            IORequest(
+        let res = try await ring.submit(request: IORequest(
                 reading: fileSlot, 
                 into: buffer, 
                 at: offset
@@ -78,8 +76,7 @@ public struct AsyncFileDescriptor: ~Copyable {
         atAbsoluteOffset offset: UInt64 = UInt64.max,
         isolation actor: isolated (any Actor)? = #isolation
     ) async throws -> UInt32 {
-        let res = await ring.submitAndWait(
-            IORequest(
+        let res = try await ring.submit(request: IORequest(
                 reading: fileSlot,
                 into: buffer,
                 at: offset
diff --git a/Sources/System/IORing.swift b/Sources/System/IORing.swift
@@ -138,54 +138,77 @@ internal func _blockingGetSubmissionEntry(ring: inout SQRing, submissionQueueEnt
     io_uring_sqe
 > {
     while true {
-        if let entry = _getSubmissionEntry(ring: &ring, submissionQueueEntries: submissionQueueEntries) {
+        if let entry = _getSubmissionEntry(
+            ring: &ring, 
+            submissionQueueEntries: submissionQueueEntries
+        ) {
             return entry
         }
         // TODO: actually block here instead of spinning
     }
 
 }
 
-internal func _submitRequests(ring: inout SQRing, ringDescriptor: Int32) {
-    let flushedEvents = _flushQueue(ring: &ring)
-
-    // Ring always needs enter right now;
+//TODO: omitting signal mask for now
+//Tell the kernel that we've submitted requests and/or are waiting for completions
+internal func _enter(
+    ringDescriptor: Int32, 
+    numEvents: UInt32,
+    minCompletions: UInt32,
+    flags: UInt32
+) throws -> Int32 {
+  // Ring always needs enter right now;
     // TODO: support SQPOLL here
     while true {
-        let ret = io_uring_enter(ringDescriptor, flushedEvents, 0, 0, nil)
+        let ret = io_uring_enter(ringDescriptor, numEvents, minCompletions, flags, nil)
         // error handling:
         //     EAGAIN / EINTR (try again),
         //     EBADF / EBADFD / EOPNOTSUPP / ENXIO
         //     (failure in ring lifetime management, fatal),
         //     EINVAL (bad constant flag?, fatal),
         //     EFAULT (bad address for argument from library, fatal)
         if ret == -EAGAIN || ret == -EINTR {
+            //TODO: should we wait a bit on AGAIN?
             continue
         } else if ret < 0 {
             fatalError(
                 "fatal error in submitting requests: " + Errno(rawValue: -ret).debugDescription
             )
         } else {
-            break
+            return ret
         }
     }
 }
 
+internal func _submitRequests(ring: inout SQRing, ringDescriptor: Int32) throws {
+    let flushedEvents = _flushQueue(ring: &ring)
+    _ = try _enter(ringDescriptor: ringDescriptor, numEvents: flushedEvents, minCompletions: 0, flags: 0)
+}
+
+internal func _getUnconsumedSubmissionCount(ring: inout SQRing) -> UInt32 {
+    return ring.userTail - ring.kernelHead.pointee.load(ordering: .acquiring)
+}
+
+internal func _getUnconsumedCompletionCount(ring: inout CQRing) -> UInt32 {
+    return ring.kernelTail.pointee.load(ordering: .acquiring) - ring.kernelHead.pointee.load(ordering: .acquiring)
+}
+
+//TODO: pretty sure this is supposed to do more than it does
 internal func _flushQueue(ring: inout SQRing) -> UInt32 {
     ring.kernelTail.pointee.store(
-        ring.userTail, ordering: .relaxed
+        ring.userTail, ordering: .releasing
     )
-    return ring.userTail - ring.kernelHead.pointee.load(ordering: .relaxed)
+    return _getUnconsumedSubmissionCount(ring: &ring)
 }
 
 @usableFromInline @inline(__always)
 internal func _getSubmissionEntry(ring: inout SQRing, submissionQueueEntries: UnsafeMutableBufferPointer<io_uring_sqe>) -> UnsafeMutablePointer<
     io_uring_sqe
 >? {
-    let next = ring.userTail + 1
+    let next = ring.userTail &+ 1 //this is expected to wrap
 
     // FEAT: smp load when SQPOLL in use (not in MVP)
-    let kernelHead = ring.kernelHead.pointee.load(ordering: .relaxed)
+    let kernelHead = ring.kernelHead.pointee.load(ordering: .acquiring)
 
     // FEAT: 128-bit event support (not in MVP)
     if next - kernelHead <= ring.array.count {
@@ -383,18 +406,45 @@ public struct IORing: @unchecked Sendable, ~Copyable {
 
     func _tryConsumeCompletion(ring: inout CQRing) -> IOCompletion? {
         let tail = ring.kernelTail.pointee.load(ordering: .acquiring)
-        let head = ring.kernelHead.pointee.load(ordering: .relaxed)
+        let head = ring.kernelHead.pointee.load(ordering: .acquiring)
 
         if tail != head {
             // 32 byte copy - oh well
             let res = ring.cqes[Int(head & ring.ringMask)]
-            ring.kernelHead.pointee.store(head + 1, ordering: .relaxed)
+            ring.kernelHead.pointee.store(head &+ 1, ordering: .releasing)
             return IOCompletion(rawValue: res)
         }
 
         return nil
     }
 
+    internal func handleRegistrationResult(_ result: Int32) throws {
+        //TODO: error handling
+    }
+
+    public mutating func registerEventFD(ring: inout IORing, _ descriptor: FileDescriptor) throws {
+        var rawfd = descriptor.rawValue
+        let result = withUnsafePointer(to: &rawfd) { fdptr in
+            return io_uring_register(
+                ring.ringDescriptor, 
+                IORING_REGISTER_EVENTFD,
+                UnsafeMutableRawPointer(mutating: fdptr), 
+                1
+            )
+        }
+        try handleRegistrationResult(result)
+    }
+
+    public mutating func unregisterEventFD(ring: inout IORing) throws {
+        let result = io_uring_register(
+            ring.ringDescriptor, 
+            IORING_UNREGISTER_EVENTFD,
+            nil, 
+            0
+        )
+        try handleRegistrationResult(result)
+    }
+
     public mutating func registerFiles(count: UInt32) {
         guard self.registeredFiles == nil else { fatalError() }
         let fileBuf = UnsafeMutableBufferPointer<UInt32>.allocate(capacity: Int(count))
@@ -445,9 +495,9 @@ public struct IORing: @unchecked Sendable, ~Copyable {
         fatalError("failed to unregister buffers: TODO")
     }
 
-    public func submitRequests() {
-        submissionMutex.withLock { ring in
-            _submitRequests(ring: &ring, ringDescriptor: ringDescriptor)
+    public func submitRequests() throws {
+        try submissionMutex.withLock { ring in
+            try _submitRequests(ring: &ring, ringDescriptor: ringDescriptor)
         }
     }
 
diff --git a/Sources/System/IORingError.swift b/Sources/System/IORingError.swift
@@ -1,3 +1,6 @@
-enum IORingError: Error {
+//TODO: make this not an enum
+public enum IORingError: Error, Equatable {
     case missingRequiredFeatures
+    case operationCanceled
+    case unknown
 }
diff --git a/Sources/System/ManagedIORing.swift b/Sources/System/ManagedIORing.swift
@@ -1,3 +1,16 @@
+fileprivate func handleCompletionError(
+    _ result: Int32, 
+    for continuation: UnsafeContinuation<IOCompletion, any Error>) {
+    var error: IORingError = .unknown
+    switch result {
+        case -(_ECANCELED):
+            error = .operationCanceled
+        default:
+            error = .unknown
+    }
+    continuation.resume(throwing: error)
+}
+
 final public class ManagedIORing: @unchecked Sendable {
     var internalRing: IORing
 
@@ -11,25 +24,64 @@ final public class ManagedIORing: @unchecked Sendable {
     private func startWaiter() {
         Task.detached {
             while !Task.isCancelled {
+                //TODO: should timeout handling be sunk into IORing?
                 let cqe = self.internalRing.blockingConsumeCompletion()
 
+                if cqe.userData == 0 {
+                    continue
+                }
                 let cont = unsafeBitCast(
-                    cqe.userData, to: UnsafeContinuation<IOCompletion, Never>.self)
-                cont.resume(returning: cqe)
+                cqe.userData, to: UnsafeContinuation<IOCompletion, any Error>.self)
+                
+                if cqe.result < 0 {
+                    var err = system_strerror(cqe.result * -1)
+                    let len = system_strlen(err!)
+                    err!.withMemoryRebound(to: UInt8.self, capacity: len) {
+                        let errStr = String(decoding: UnsafeBufferPointer(start: $0, count: len), as: UTF8.self)
+                        print("\(errStr)")
+                    }
+                    handleCompletionError(cqe.result, for: cont)
+                } else {
+                    cont.resume(returning: cqe)
+                }
             }
         }
     }
 
-    public func submitAndWait(_ request: __owned IORequest, isolation actor: isolated (any Actor)? = #isolation) async -> IOCompletion {
+    public func submit(
+        request: __owned IORequest,
+        timeout: Duration? = nil,
+        isolation actor: isolated (any Actor)? = #isolation
+    ) async throws -> IOCompletion {
         var consumeOnceWorkaround: IORequest? = request
-        return await withUnsafeContinuation { cont in
-            return internalRing.submissionMutex.withLock { ring in
-                let request = consumeOnceWorkaround.take()!
-                let entry = _blockingGetSubmissionEntry(
-                    ring: &ring, submissionQueueEntries: internalRing.submissionQueueEntries)
-                entry.pointee = request.makeRawRequest().rawValue
-                entry.pointee.user_data = unsafeBitCast(cont, to: UInt64.self)
-                _submitRequests(ring: &ring, ringDescriptor: internalRing.ringDescriptor)
+        return try await withUnsafeThrowingContinuation { cont in
+            do {
+                try internalRing.submissionMutex.withLock { ring in
+                    let request = consumeOnceWorkaround.take()!
+                    let entry = _blockingGetSubmissionEntry(
+                        ring: &ring, submissionQueueEntries: internalRing.submissionQueueEntries)
+                    entry.pointee = request.makeRawRequest().rawValue
+                    entry.pointee.user_data = unsafeBitCast(cont, to: UInt64.self)
+                    if let timeout {
+                        //TODO: if IORING_FEAT_MIN_TIMEOUT is supported we can do this more efficiently
+                        let timeoutEntry = _blockingGetSubmissionEntry(
+                            ring: &ring, 
+                            submissionQueueEntries: internalRing.submissionQueueEntries
+                        )
+                        try RawIORequest.withTimeoutRequest(
+                            linkedTo: entry,
+                            in: timeoutEntry,
+                            duration: timeout, 
+                            flags: .relativeTime
+                        ) {
+                            try _submitRequests(ring: &ring, ringDescriptor: internalRing.ringDescriptor)
+                        }
+                    } else {
+                        try _submitRequests(ring: &ring, ringDescriptor: internalRing.ringDescriptor)
+                    }
+                }
+            } catch (let e) {
+                cont.resume(throwing: e)
             }
         }
 
diff --git a/Sources/System/RawIORequest.swift b/Sources/System/RawIORequest.swift
@@ -24,6 +24,8 @@ extension RawIORequest {
         case sendMessage = 9
         case receiveMessage = 10
         // ...
+        case link_timeout = 15
+        // ...
         case openAt = 18
         case close = 19
         case filesUpdate = 20
@@ -103,7 +105,7 @@ extension RawIORequest {
         // poll32_events
         // sync_range_flags
         // msg_flags
-        // timeout_flags
+        case timeoutFlags(TimeOutFlags)
         // accept_flags
         // cancel_flags
         case openFlags(FileDescriptor.OpenOptions)
@@ -132,12 +134,48 @@ extension RawIORequest {
         // append to end of the file
         public static let append = ReadWriteFlags(rawValue: 1 << 4)
     }
+
+    public struct TimeOutFlags: OptionSet, Hashable, Codable {
+        public var rawValue: UInt32
+
+        public init(rawValue: UInt32) {
+            self.rawValue = rawValue
+        }
+
+        public static let relativeTime: RawIORequest.TimeOutFlags = TimeOutFlags(rawValue: 0)
+        public static let absoluteTime: RawIORequest.TimeOutFlags = TimeOutFlags(rawValue: 1 << 0)
+    }
 }
 
 extension RawIORequest {
     static func nop() -> RawIORequest {
-        var req = RawIORequest()
+        var req: RawIORequest = RawIORequest()
         req.operation = .nop
         return req
     }
+
+    //TODO: typed errors
+    static func withTimeoutRequest<R>(
+        linkedTo opEntry: UnsafeMutablePointer<io_uring_sqe>,
+        in timeoutEntry: UnsafeMutablePointer<io_uring_sqe>,
+        duration: Duration, 
+        flags: TimeOutFlags, 
+        work: () throws -> R) rethrows -> R {
+
+        opEntry.pointee.flags |= Flags.linkRequest.rawValue
+        opEntry.pointee.off = 1
+        var ts = __kernel_timespec(
+            tv_sec: duration.components.seconds, 
+            tv_nsec: duration.components.attoseconds / 1_000_000_000 
+        )
+        return try withUnsafePointer(to: &ts) { tsPtr in
+            var req: RawIORequest = RawIORequest()
+            req.operation = .link_timeout
+            req.rawValue.timeout_flags = flags.rawValue
+            req.rawValue.len = 1
+            req.rawValue.addr = UInt64(UInt(bitPattern: tsPtr))
+            timeoutEntry.pointee = req.rawValue
+            return try work()
+        }
+    }
 }
diff --git a/Tests/SystemTests/IORingTests.swift b/Tests/SystemTests/IORingTests.swift
@@ -14,7 +14,7 @@ final class IORingTests: XCTestCase {
     func testNop() throws {
         var ring = try IORing(queueDepth: 32)
         ring.writeRequest(IORequest())
-        ring.submitRequests()
+        try ring.submitRequests()
         let completion = ring.blockingConsumeCompletion()
         XCTAssertEqual(completion.result, 0)
     }
diff --git a/Tests/SystemTests/ManagedIORingTests.swift b/Tests/SystemTests/ManagedIORingTests.swift

Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,6 @@`
`1`		`-enum IORingError: Error {`
	`1`	`+//TODO: make this not an enum`
	`2`	`+public enum IORingError: Error, Equatable {`
`2`	`3`	`case missingRequiredFeatures`
	`4`	`+ case operationCanceled`
	`5`	`+ case unknown`
`3`	`6`	`}`
Original file line number	Diff line number	Diff line change
`@@ -14,7 +14,7 @@ final class IORingTests: XCTestCase {`
`14`	`14`	`func testNop() throws {`
`15`	`15`	`var ring = try IORing(queueDepth: 32)`
`16`	`16`	`ring.writeRequest(IORequest())`
`17`		`- ring.submitRequests()`
	`17`	`+ try ring.submitRequests()`
`18`	`18`	`let completion = ring.blockingConsumeCompletion()`
`19`	`19`	`XCTAssertEqual(completion.result, 0)`
`20`	`20`	`}`