Add offset writes and file size support to NonBlockingFileIO (#1408)

* Add offset writes and file size support to NonBlockingFileIO

Motivation:

NonBlockingFileIO currently only supports writes from 0 or appending,
if opened with O_APPEND. To facilitate random read/write file access
the API needs to evolve.

Modifications:

- Add toOffset function parameter to write for offset file write access
- Add changeFileSize function for growing or shrinking the file
- Add readFileSize function for asynchronously reading the file size

Result:

- Improve API for random access scenarios
- Provide method for reading and writing file size

* Conform to code review

* Formatting

* Update Sources/NIO/NonBlockingFileIO.swift

Co-Authored-By: Johannes Weiss <johannesweiss@apple.com>

* Update Sources/NIO/NonBlockingFileIO.swift

Co-Authored-By: Johannes Weiss <johannesweiss@apple.com>

* Modify offset handling to support partial writes.
Clean up stray comment.

* Fix precondition check within offset pwrite handling

* Improve guard for readability

* White space cleanup

* Single spacing issue

Co-authored-by: Johannes Weiss <johannesweiss@apple.com>

Sources/NIO/NonBlockingFileIO.swift

@@ -349,6 +349,44 @@ public struct NonBlockingFileIO {
         }
     }
 
+    /// Changes the file size of `fileHandle` to `size`.
+    ///
+    /// If `size` is smaller than the current file size, the remaining bytes will be truncated and are lost. If `size`
+    /// is larger than the current file size, the gap will be filled with zero bytes.
+    ///
+    /// - parameters:
+    ///   - fileHandle: The `NIOFileHandle` to write to.
+    ///   - size: The new file size in bytes to write.
+    ///   - eventLoop: The `EventLoop` to create the returned `EventLoopFuture` from.
+    /// - returns: An `EventLoopFuture` which is fulfilled if the write was successful or fails on error.
+    public func changeFileSize(fileHandle: NIOFileHandle,
+                               size: Int64,
+                               eventLoop: EventLoop) -> EventLoopFuture<()> {
+        return self.threadPool.runIfActive(eventLoop: eventLoop) {
+            try fileHandle.withUnsafeFileDescriptor { descriptor -> Void in
+                try Posix.ftruncate(descriptor: descriptor, size: off_t(size))
+            }
+        }
+    }
+
+    /// Returns the length of the file associated with `fileHandle`.
+    ///
+    /// - parameters:
+    ///   - fileHandle: The `NIOFileHandle` to read from.
+    ///   - eventLoop: The `EventLoop` to create the returned `EventLoopFuture` from.
+    /// - returns: An `EventLoopFuture` which is fulfilled if the write was successful or fails on error.
+    public func readFileSize(fileHandle: NIOFileHandle,
+                             eventLoop: EventLoop) -> EventLoopFuture<Int64> {
+        return self.threadPool.runIfActive(eventLoop: eventLoop) {
+            return try fileHandle.withUnsafeFileDescriptor { descriptor in
+                let curr = try Posix.lseek(descriptor: descriptor, offset: 0, whence: SEEK_CUR)
+                let eof = try Posix.lseek(descriptor: descriptor, offset: 0, whence: SEEK_END)
+                try Posix.lseek(descriptor: descriptor, offset: curr, whence: SEEK_SET)
+                return Int64(eof)
+            }
+        }
+    }
+
     /// Write `buffer` to `fileHandle` in `NonBlockingFileIO`'s private thread pool which is separate from any `EventLoop` thread.
     ///
     /// - parameters:
@@ -359,7 +397,29 @@ public struct NonBlockingFileIO {
     public func write(fileHandle: NIOFileHandle,
                       buffer: ByteBuffer,
                       eventLoop: EventLoop) -> EventLoopFuture<()> {
-        var byteCount = buffer.readableBytes
+        return self.write0(fileHandle: fileHandle, toOffset: nil, buffer: buffer, eventLoop: eventLoop)
+    }
+
+    /// Write `buffer` starting from `toOffset` to `fileHandle` in `NonBlockingFileIO`'s private thread pool which is separate from any `EventLoop` thread.
+    ///
+    /// - parameters:
+    ///   - fileHandle: The `NIOFileHandle` to write to.
+    ///   - toOffset: The file offset to write to.
+    ///   - buffer: The `ByteBuffer` to write.
+    ///   - eventLoop: The `EventLoop` to create the returned `EventLoopFuture` from.
+    /// - returns: An `EventLoopFuture` which is fulfilled if the write was successful or fails on error.
+    public func write(fileHandle: NIOFileHandle,
+                      toOffset: Int64,
+                      buffer: ByteBuffer,
+                      eventLoop: EventLoop) -> EventLoopFuture<()> {
+        return self.write0(fileHandle: fileHandle, toOffset: toOffset, buffer: buffer, eventLoop: eventLoop)
+    }
+
+    private func write0(fileHandle: NIOFileHandle,
+                        toOffset: Int64?,
+                        buffer: ByteBuffer,
+                        eventLoop: EventLoop) -> EventLoopFuture<()> {
+        let byteCount = buffer.readableBytes
 
         guard byteCount > 0 else {
             return eventLoop.makeSucceededFuture(())
@@ -367,13 +427,22 @@ public struct NonBlockingFileIO {
 
         return self.threadPool.runIfActive(eventLoop: eventLoop) {
             var buf = buffer
-            while byteCount > 0 {
+
+            var offsetAccumulator: Int = 0
+            repeat {
                 let n = try buf.readWithUnsafeReadableBytes { ptr in
-                    precondition(ptr.count == byteCount)
-                    let res = try fileHandle.withUnsafeFileDescriptor { descriptor in
-                        try Posix.write(descriptor: descriptor,
-                                        pointer: ptr.baseAddress!,
-                                        size: byteCount)
+                    precondition(ptr.count == byteCount - offsetAccumulator)
+                    let res: IOResult<ssize_t> = try fileHandle.withUnsafeFileDescriptor { descriptor in
+                        if let toOffset = toOffset {
+                            return try Posix.pwrite(descriptor: descriptor,
+                                                    pointer: ptr.baseAddress!,
+                                                    size: byteCount - offsetAccumulator,
+                                                    offset: off_t(toOffset + Int64(offsetAccumulator)))
+                        } else {
+                            return try Posix.write(descriptor: descriptor,
+                                                   pointer: ptr.baseAddress!,
+                                                   size: byteCount - offsetAccumulator)
+                        }
                     }
                     switch res {
                     case .processed(let n):
@@ -383,9 +452,8 @@ public struct NonBlockingFileIO {
                         throw Error.descriptorSetToNonBlocking
                     }
                 }
-
-                byteCount -= n
-            }
+                offsetAccumulator += n
+            } while offsetAccumulator < byteCount
         }
     }
 

Sources/NIO/System.swift

@@ -57,7 +57,9 @@ private let sysAccept = accept
 private let sysConnect = connect
 private let sysOpen: (UnsafePointer<CChar>, CInt) -> CInt = open
 private let sysOpenWithMode: (UnsafePointer<CChar>, CInt, mode_t) -> CInt = open
+private let sysFtruncate = ftruncate
 private let sysWrite = write
+private let sysPwrite = pwrite
 private let sysRead = read
 private let sysPread = pread
 private let sysLseek = lseek
@@ -353,6 +355,14 @@ internal enum Posix {
         }
     }
 
+    @inline(never)
+    @discardableResult
+    public static func ftruncate(descriptor: CInt, size: off_t) throws -> CInt {
+        return try wrapSyscall {
+            sysFtruncate(descriptor, size)
+        }
+    }
+    
     @inline(never)
     public static func write(descriptor: CInt, pointer: UnsafeRawPointer, size: Int) throws -> IOResult<Int> {
         return try wrapSyscallMayBlock {
@@ -360,6 +370,13 @@ internal enum Posix {
         }
     }
 
+    @inline(never)
+    public static func pwrite(descriptor: CInt, pointer: UnsafeRawPointer, size: Int, offset: off_t) throws -> IOResult<Int> {
+        return try wrapSyscallMayBlock {
+            sysPwrite(descriptor, pointer, size, offset)
+        }
+    }
+
     @inline(never)
     public static func writev(descriptor: CInt, iovecs: UnsafeBufferPointer<IOVector>) throws -> IOResult<Int> {
         return try wrapSyscallMayBlock {

Tests/NIOTests/NonBlockingFileIOTest+XCTest.swift

@@ -46,8 +46,13 @@ extension NonBlockingFileIOTest {
                 ("testFileRegionReadFromPipeFails", testFileRegionReadFromPipeFails),
                 ("testReadFromNonBlockingPipeFails", testReadFromNonBlockingPipeFails),
                 ("testSeekPointerIsSetToFront", testSeekPointerIsSetToFront),
+                ("testReadingFileSize", testReadingFileSize),
+                ("testChangeFileSizeShrink", testChangeFileSizeShrink),
+                ("testChangeFileSizeGrow", testChangeFileSizeGrow),
                 ("testWriting", testWriting),
                 ("testWriteMultipleTimes", testWriteMultipleTimes),
+                ("testWritingWithOffset", testWritingWithOffset),
+                ("testWritingBeyondEOF", testWritingBeyondEOF),
                 ("testFileOpenWorks", testFileOpenWorks),
                 ("testFileOpenWorksWithEmptyFile", testFileOpenWorksWithEmptyFile),
                 ("testFileOpenFails", testFileOpenFails),

Tests/NIOTests/NonBlockingFileIOTest.swift

@@ -453,6 +453,45 @@ class NonBlockingFileIOTest: XCTestCase {
         XCTAssertEqual(2, numCalls)
     }
 
+    func testReadingFileSize() throws {
+        try withTemporaryFile(content: "0123456789") { (fileHandle, _) -> Void in
+            let size = try self.fileIO.readFileSize(fileHandle: fileHandle,
+                                                    eventLoop: eventLoop).wait()
+            XCTAssertEqual(size, 10)
+        }
+    }
+
+    func testChangeFileSizeShrink() throws {
+        try withTemporaryFile(content: "0123456789") { (fileHandle, _) -> Void in
+            try self.fileIO.changeFileSize(fileHandle: fileHandle,
+                                           size: 1,
+                                           eventLoop: eventLoop).wait()
+            let fileRegion = try FileRegion(fileHandle: fileHandle)
+            var buf = try self.fileIO.read(fileRegion: fileRegion,
+                                           allocator: allocator,
+                                           eventLoop: eventLoop).wait()
+            XCTAssertEqual("0", buf.readString(length: buf.readableBytes))
+        }
+    }
+
+    func testChangeFileSizeGrow() throws {
+        try withTemporaryFile(content: "0123456789") { (fileHandle, _) -> Void in
+            try self.fileIO.changeFileSize(fileHandle: fileHandle,
+                                           size: 100,
+                                           eventLoop: eventLoop).wait()
+            let fileRegion = try FileRegion(fileHandle: fileHandle)
+            var buf = try self.fileIO.read(fileRegion: fileRegion,
+                                           allocator: allocator,
+                                           eventLoop: eventLoop).wait()
+            let zeros = (1...90).map { _ in UInt8(0) }
+            guard let bytes = buf.readBytes(length: buf.readableBytes)?.suffix(from: 10) else {
+                XCTFail("readBytes(length:) should not be nil")
+                return
+            }
+            XCTAssertEqual(zeros, Array(bytes))
+        }
+    }
+
     func testWriting() throws {
         var buffer = allocator.buffer(capacity: 3)
         buffer.writeStaticString("123")
@@ -499,6 +538,53 @@ class NonBlockingFileIOTest: XCTestCase {
         }
     }
 
+    func testWritingWithOffset() throws {
+        var buffer = allocator.buffer(capacity: 3)
+        buffer.writeStaticString("123")
+
+        try withTemporaryFile(content: "hello") { (fileHandle, _) -> Void in
+            try self.fileIO.write(fileHandle: fileHandle,
+                                  toOffset: 1,
+                                  buffer: buffer,
+                                  eventLoop: eventLoop).wait()
+            let offset = try fileHandle.withUnsafeFileDescriptor {
+                try Posix.lseek(descriptor: $0, offset: 0, whence: SEEK_SET)
+            }
+            XCTAssertEqual(offset, 0)
+
+            var readBuffer = try self.fileIO.read(fileHandle: fileHandle,
+                                                  byteCount: 5,
+                                                  allocator: self.allocator,
+                                                  eventLoop: self.eventLoop).wait()
+            XCTAssertEqual(5, readBuffer.readableBytes)
+            XCTAssertEqual("h123o", readBuffer.readString(length: readBuffer.readableBytes))
+        }
+    }
+
+    // This is undefined behavior and may cause different
+    // results on other platforms. Please add #if:s according
+    // to platform requirements.
+    func testWritingBeyondEOF() throws {
+        var buffer = allocator.buffer(capacity: 3)
+        buffer.writeStaticString("123")
+
+        try withTemporaryFile(content: "hello") { (fileHandle, _) -> Void in
+            try self.fileIO.write(fileHandle: fileHandle,
+                                  toOffset: 6,
+                                  buffer: buffer,
+                                  eventLoop: eventLoop).wait()
+
+            let fileRegion = try FileRegion(fileHandle: fileHandle)
+            var buf = try self.fileIO.read(fileRegion: fileRegion,
+                                           allocator: self.allocator,
+                                           eventLoop: self.eventLoop).wait()
+            XCTAssertEqual(9, buf.readableBytes)
+            XCTAssertEqual("hello", buf.readString(length: 5))
+            XCTAssertEqual([ UInt8(0) ], buf.readBytes(length: 1))
+            XCTAssertEqual("123", buf.readString(length: buf.readableBytes))
+        }
+    }
+
     func testFileOpenWorks() throws {
         let content = "123"
         try withTemporaryFile(content: content) { (fileHandle, path) -> Void in