Non-copyable ColumnWriter

Author: ahoppen

Package.swift

@@ -54,7 +54,7 @@ let package = Package(
                 .product(name: "NIOFoundationCompat", package: "swift-nio"),
                 .product(name: "ServiceLifecycle", package: "swift-service-lifecycle"),
             ],
-            swiftSettings: swiftSettings
+            swiftSettings: swiftSettings + [.enableExperimentalFeature("Lifetimes")]
         ),
         .target(
             name: "_ConnectionPoolModule",

Sources/PostgresNIO/Connection/PostgresConnection+CopyFrom.swift

@@ -98,27 +98,42 @@ public struct PostgresCopyFromWriter: Sendable {
     }
 }
 
-
+// PostgresBinaryCopyFromWriter relies on non-Escapable types, which were only introduced in Swift 6.2
+#if compiler(>=6.2)
 /// Handle to send binary data for a `COPY ... FROM STDIN` query to the backend.
 ///
 /// It takes care of serializing `PostgresEncodable` column types into the binary format that Postgres expects.
 public struct PostgresBinaryCopyFromWriter: ~Copyable {
     /// Handle to serialize columns into a row that is being written by `PostgresBinaryCopyFromWriter`.
-    public struct ColumnWriter: ~Copyable {
-        /// The `PostgresBinaryCopyFromWriter` that is gathering the serialized data.
-        ///
-        /// We need to model this as `UnsafeMutablePointer` because we can't express in the Swift type system that
-        /// `ColumnWriter` never exceeds the lifetime of `PostgresBinaryCopyFromWriter`.
+    public struct ColumnWriter: ~Escapable, ~Copyable {
+        /// Pointer to the `PostgresBinaryCopyFromWriter` that is gathering the serialized data.
         @usableFromInline
         let underlying: UnsafeMutablePointer<PostgresBinaryCopyFromWriter>
 
         /// The number of columns that have been written by this `ColumnWriter`.
         @usableFromInline
         var columns: UInt16 = 0
 
+        /// - Warning: Do not call directly, call `withColumnWriter` instead
         @usableFromInline
-        init(underlying: UnsafeMutablePointer<PostgresBinaryCopyFromWriter>) {
-            self.underlying = underlying
+        init(_underlying: UnsafeMutablePointer<PostgresBinaryCopyFromWriter>) {
+            self.underlying = _underlying
+        }
+
+        @usableFromInline
+        static func withColumnWriter<T>(
+            writingTo underlying: inout PostgresBinaryCopyFromWriter,
+            body: (inout ColumnWriter) throws -> T
+        ) rethrows -> T {
+            return try withUnsafeMutablePointer(to: &underlying) { pointerToUnderlying in
+                // We can guarantee that `ColumWriter` never outlives `underlying` because `ColumnWriter` is
+                // `~Escapable` and thus cannot escape the context of the closure to `withUnsafeMutablePointer`.
+                // To model this without resorting to unsafe pointers, we would need to be able to declare an `inout`
+                // reference to `PostgresBinaryCopyFromWriter` as a member of `ColumnWriter`, which isn't possible at
+                // the moment (https://github.com/swiftlang/swift/issues/85832).
+                var columnWriter = ColumnWriter(_underlying: pointerToUnderlying)
+                return try body(&columnWriter)
+            }
         }
 
         /// Serialize a single column to a row.
@@ -128,16 +143,19 @@ public struct PostgresBinaryCopyFromWriter: ~Copyable {
         ///   be called with an `Int32`. Serializing an integer of a different width will cause a deserialization
         ///   failure in the backend.
         @inlinable
+        #if compiler(<6.3)
+        @_lifetime(&self)
+        #endif
         public mutating func writeColumn(_ column: (some PostgresEncodable)?) throws {
             columns += 1
             try invokeWriteColumn(on: underlying, column)
         }
 
-        // Needed to work around https://github.com/swiftlang/swift/issues/83309, copying the implementation into 
+        // Needed to work around https://github.com/swiftlang/swift/issues/83309, copying the implementation into
         // `writeColumn` causes an assertion failure when thread sanitizer is enabled.
-        @inlinable 
+        @inlinable
         func invokeWriteColumn(
-            on writer: UnsafeMutablePointer<PostgresBinaryCopyFromWriter>, 
+            on writer: UnsafeMutablePointer<PostgresBinaryCopyFromWriter>,
             _ column: (some PostgresEncodable)?
         ) throws {
             try writer.pointee.writeColumn(column)
@@ -169,17 +187,8 @@ public struct PostgresBinaryCopyFromWriter: ~Copyable {
         let columnIndex = buffer.writerIndex
         buffer.writeInteger(UInt16(0))
 
-        let columns = try withUnsafeMutablePointer(to: &self) { pointerToSelf in
-            // Important: We need to ensure that `pointerToSelf` (and thus `ColumnWriter`) does not exceed the lifetime
-            // of `self` because it is holding an unsafe reference to it.
-            //
-            // We achieve this because `ColumnWriter` is non-Copyable and thus the client can't store a copy to it.
-            // Furthermore, `columnWriter` is destroyed before the end of `withUnsafeMutablePointer`, which holds `self`
-            // alive.
-            var columnWriter = ColumnWriter(underlying: pointerToSelf)
-
+        let columns = try ColumnWriter.withColumnWriter(writingTo: &self) { columnWriter in
             try body(&columnWriter)
-
             return columnWriter.columns
         }
 
@@ -212,6 +221,7 @@ public struct PostgresBinaryCopyFromWriter: ~Copyable {
         buffer.clear()
     }
 }
+#endif
 
 /// Specifies the format in which data is transferred to the backend in a COPY operation.
 ///
@@ -289,6 +299,7 @@ private func buildCopyFromQuery(
 }
 
 extension PostgresConnection {
+    #if compiler(>=6.2)
     /// Copy data into a table using a `COPY <table name> FROM STDIN` query, transferring data in a binary format.
     ///
     /// - Parameters:
@@ -331,6 +342,7 @@ extension PostgresConnection {
             try await binaryWriter.flush()
         }
     }
+    #endif
 
     /// Copy data into a table using a `COPY <table name> FROM STDIN` query.
     ///

Tests/PostgresNIOTests/New/PostgresConnectionTests.swift

@@ -936,6 +936,7 @@ import Synchronization
         }
     }
 
+    #if compiler(>=6.2) // copyFromBinary is only available in Swift 6.2+
     @Test func testCopyFromBinary() async throws {
         try await self.withAsyncTestingChannel { connection, channel in
             try await withThrowingTaskGroup(of: Void.self) { taskGroup async throws -> Void in
@@ -997,6 +998,7 @@ import Synchronization
             }
         }
     }
+    #endif
 
     func withAsyncTestingChannel(_ body: (PostgresConnection, NIOAsyncTestingChannel) async throws -> ()) async throws {
         let eventLoop = NIOAsyncTestingEventLoop()