Merge pull request #158 from vapor/error-handling

Add error handlers to TCPSocketSink

Author: tanner0101

Sources/TCP/Streams/TCPSocketSink.swift

@@ -7,6 +7,9 @@ private let maxExcessSignalCount: Int = 2
 
 /// Data stream wrapper for a dispatch socket.
 public final class TCPSocketSink: Async.InputStream {
+    /// Able to handle errors that are thrown to the Sink
+    public typealias ErrorHandler = (TCPSocketSink, Error) -> ()
+    
     /// See InputStream.Input
     public typealias Input = ByteBuffer
 
@@ -34,15 +37,19 @@ public final class TCPSocketSink: Async.InputStream {
     /// The current number of signals received while downstream was not ready
     /// since it was last ready
     private var excessSignalCount: Int
+    
+    /// This closure will be called with an error thrown from upstream
+    private let onError: ErrorHandler
 
     /// Creates a new `SocketSink`
-    internal init(socket: TCPSocket, on worker: Worker) {
+    internal init(socket: TCPSocket, on worker: Worker, onError: @escaping ErrorHandler) {
         self.socket = socket
         self.eventLoop = worker.eventLoop
         self.inputBuffer = nil
         self.isClosed = false
         self.sourceIsSuspended = true
         self.excessSignalCount = 0
+        self.onError = onError
         let writeSource = self.eventLoop.onWritable(descriptor: socket.descriptor, writeSourceSignal)
         self.writeSource = writeSource
     }
@@ -64,8 +71,7 @@ public final class TCPSocketSink: Async.InputStream {
         case .close:
             close()
         case .error(let e):
-            close()
-            fatalError("\(e)")
+            onError(self, e)
         }
     }
 
@@ -163,8 +169,15 @@ public final class TCPSocketSink: Async.InputStream {
 
 extension TCPSocket {
     /// Creates a data stream for this socket on the supplied event loop.
+    public func sink(on eventLoop: Worker, onError: @escaping TCPSocketSink.ErrorHandler) -> TCPSocketSink {
+        return .init(socket: self, on: eventLoop, onError: onError)
+    }
+    
+    /// Creates a data stream for this socket on the supplied event loop.
+    @available(*, deprecated)
     public func sink(on eventLoop: Worker) -> TCPSocketSink {
-        return .init(socket: self, on: eventLoop)
+        return .init(socket: self, on: eventLoop) { _, error in
+            fatalError("Uncaught error in TCPSocketSink: \(error).")
+        }
     }
 }
-

Sources/TCP/Streams/TCPSocketStream.swift

@@ -16,9 +16,9 @@ public final class TCPSocketStream: ByteStream {
     internal let sink: TCPSocketSink
 
     /// Internal stream init. Use socket convenience method.
-    internal init(socket: TCPSocket, bufferSize: Int, on worker: Worker) {
+    internal init(socket: TCPSocket, bufferSize: Int, on worker: Worker, onError: @escaping TCPSocketSink.ErrorHandler) {
         self.source = socket.source(on: worker, bufferSize: bufferSize)
-        self.sink = socket.sink(on: worker)
+        self.sink = socket.sink(on: worker, onError: onError)
     }
 
     /// See `InputStream.input(_:)`
@@ -34,7 +34,15 @@ public final class TCPSocketStream: ByteStream {
 
 extension TCPSocket {
     /// Create a `TCPSocketStream` for this socket.
+    public func stream(bufferSize: Int = 4096, on worker: Worker, onError: @escaping TCPSocketSink.ErrorHandler) -> TCPSocketStream {
+        return TCPSocketStream(socket: self, bufferSize: bufferSize, on: worker, onError: onError)
+    }
+    
+    /// Create a `TCPSocketStream` for this socket.
+    @available(*, deprecated)
     public func stream(bufferSize: Int = 4096, on worker: Worker) -> TCPSocketStream {
-        return TCPSocketStream(socket: self, bufferSize: bufferSize, on: worker)
+        return TCPSocketStream(socket: self, bufferSize: bufferSize, on: worker) { _, error in
+            fatalError("Uncaught error in TCPSocketStream: \(error).")
+        }
     }
 }

Tests/TCPTests/SocketsTests.swift

@@ -25,7 +25,10 @@ class SocketsTests: XCTestCase {
             /// set up the server stream
             serverStream.drain { client in
                 let clientSource = client.socket.source(on: workerLoop)
-                let clientSink = client.socket.sink(on: workerLoop)
+                let clientSink = client.socket.sink(on: workerLoop) { sink, error in
+                    XCTFail("\(error)")
+                    sink.close()
+                }
                 clientSource.output(to: clientSink)
             }.catch { err in
                 XCTFail("\(err)")