Swift-NIO + WebSocket-Kit: Proper Setup/Cleanup in a Mac App

Question

Context

I'm developing a Mac app. In this app, I want to run a websocket server. To do this, I'm using Swift NIO and Websocket-Kit. My full setup is below.

Question

All of the documentation for Websocket-Kit and SwiftNIO is geared towards a creating a single server-side process that starts up when you launch it from the command line and then runs infinitely.

In my app, I must be able to start the websocket server and then shut it down and restart it on demand, without re-launching my application. The code below does that, but I would like confirmation of two things:

1. In the test() function, I send some text to all connected clients. I am unsure if this is thread-safe and correct. Can I store the WebSocket instances as I'm doing here and message them from the main thread of my application?

2. Am I shutting down the websocket server correctly? The result of the call to serverBootstrap(group:)[...].bind(host:port:).wait() creates a Channel and then waits infinitely. When I call shutdownGracefully() on the associated EventLoopGroup, is that server cleaned up correctly? (I can confirm that port 5759 is free again after this shutdown, so I'm guessing everything is cleaned up?)

Thanks for the input; it's tough to find examples of using SwiftNIO and Websocket-Kit inside an application.

Code

import Foundation
import NIO
import NIOHTTP1
import NIOWebSocket
import WebSocketKit


@objc class WebsocketServer: NSObject
{
    private var queue: DispatchQueue?
    private var eventLoopGroup: MultiThreadedEventLoopGroup?
    private var websocketClients: [WebSocket] = []
    
    
    @objc func startServer()
    {
        queue = DispatchQueue.init(label: "socketServer")
        queue?.async
        {
            let upgradePipelineHandler: (Channel, HTTPRequestHead) -> EventLoopFuture<Void> = { channel, req in
                
                WebSocket.server(on: channel) { ws in
                    ws.send("You have connected to WebSocket")
                    
                    DispatchQueue.main.async {
                        self.websocketClients.append(ws)
                        print("websocketClients after connection: \(self.websocketClients)")
                    }
                
                    ws.onText { ws, string in
                        print("received")
                        ws.send(string.trimmingCharacters(in: .whitespacesAndNewlines).reversed())
                    }
                
                    ws.onBinary { ws, buffer in
                        print(buffer)
                    }
                
                    ws.onClose.whenSuccess { value in
                        print("onClose")
                        
                        DispatchQueue.main.async
                        {
                            self.websocketClients.removeAll { (socketToTest) -> Bool in
                                return socketToTest === ws
                            }
                            
                            print("websocketClients after close: \(self.websocketClients)")
                        }
                    }
                }
            }

            self.eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 2)
            let port: Int = 5759

            let promise = self.eventLoopGroup!.next().makePromise(of: String.self)
            
            let server = try? ServerBootstrap(group: self.eventLoopGroup!)
                
                // Specify backlog and enable SO_REUSEADDR for the server itself
                .serverChannelOption(ChannelOptions.backlog, value: 256)
                .serverChannelOption(ChannelOptions.socketOption(.so_reuseaddr), value: 1)
                
                .childChannelInitializer { channel in
              
                let webSocket = NIOWebSocketServerUpgrader(
                    shouldUpgrade: { channel, req in
                        return channel.eventLoop.makeSucceededFuture([:])
                    },
                    upgradePipelineHandler: upgradePipelineHandler
                )
              
                return channel.pipeline.configureHTTPServerPipeline(
                    withServerUpgrade: (
                        upgraders: [webSocket],
                        completionHandler: { ctx in
                            // complete
                        })
                )
            }.bind(host: "0.0.0.0", port: port).wait()                  

            _ = try! promise.futureResult.wait()
        }
    }
    
    
    
    ///
    ///  Send a message to connected clients, then shut down the server.
    ///
    @objc func test()
    {
        self.websocketClients.forEach { (ws) in
            ws.eventLoop.execute {
                ws.send("This is a message being sent to all websockets.")
            }
        }
        
        stopServer()
    }
    
    
    
    @objc func stopServer()
    {
        self.websocketClients.forEach { (ws) in
            try? ws.eventLoop.submit { () -> Void in
                print("closing websocket: \(ws)")
                _ = ws.close()
            }.wait()                        // Block until complete so we don't shut down the eventLoop before all clients get closed.
        }
        
        eventLoopGroup?.shutdownGracefully(queue: .main, { (error: Error?) in

            print("Eventloop shutdown now complete.")
            self.eventLoopGroup = nil
            self.queue = nil
        })
    }
}

Answer 1

In the test() function, I send some text to all connected clients. I am unsure if this is thread-safe and correct. Can I store the WebSocket instances as I'm doing here and message them from the main thread of my application?

Exactly as you're doing here, yes, that should be safe. ws.eventLoop.execute will execute that block on the event loop thread belonging to that WebSocket connection. This will be safe.

When I call shutdownGracefully() on the associated EventLoopGroup, is that server cleaned up correctly? (I can confirm that port 5759 is free again after this shutdown, so I'm guessing everything is cleaned up?)

Yes. shutdownGracefully forces all connections and listening sockets closed.