From 6aaa9cb458d8368194276ed107417109538134de Mon Sep 17 00:00:00 2001
From: Christoph Hagen <github@christophhagen.de>
Date: Wed, 6 Dec 2023 09:13:59 +0100
Subject: [PATCH] Simplify async scheduler

---
 Sources/App/EventLoopScheduler.swift | 29 +++-------------------------
 Sources/App/configure.swift          | 13 ++++++++++++-
 Sources/Run/main.swift               |  5 ++++-
 3 files changed, 19 insertions(+), 28 deletions(-)

diff --git a/Sources/App/EventLoopScheduler.swift b/Sources/App/EventLoopScheduler.swift
index 3997c05..b21c496 100644
--- a/Sources/App/EventLoopScheduler.swift
+++ b/Sources/App/EventLoopScheduler.swift
@@ -3,32 +3,9 @@ import Clairvoyant
 import Vapor
 import NIOCore
 
-final class EventLoopScheduler {
+extension MultiThreadedEventLoopGroup: AsyncScheduler {
 
-    private let backgroundGroup: EventLoopGroup
-
-    init(numberOfThreads: Int = 2) {
-        backgroundGroup = MultiThreadedEventLoopGroup(numberOfThreads: numberOfThreads)
-    }
-
-    func next() -> EventLoop {
-        backgroundGroup.next()
-    }
-
-    func provider() -> NIOEventLoopGroupProvider {
-        return .shared(backgroundGroup)
-    }
-
-    func shutdown() {
-        backgroundGroup.shutdownGracefully { _ in
-
-        }
-    }
-}
-
-extension EventLoopScheduler: AsyncScheduler {
-
-    func schedule(asyncJob: @escaping @Sendable () async throws -> Void) {
-        _ = backgroundGroup.any().makeFutureWithTask(asyncJob)
+    public func schedule(asyncJob: @escaping @Sendable () async throws -> Void) {
+        _ = any().makeFutureWithTask(asyncJob)
     }
 }
diff --git a/Sources/App/configure.swift b/Sources/App/configure.swift
index b0eb373..e621a52 100755
--- a/Sources/App/configure.swift
+++ b/Sources/App/configure.swift
@@ -7,7 +7,7 @@ import ClairvoyantBinaryCodable
 private var provider: VaporMetricProvider!
 private var serverStatus: Metric<ServerStatus>!
 
-private let asyncScheduler = EventLoopScheduler()
+private let asyncScheduler = MultiThreadedEventLoopGroup(numberOfThreads: 2)
 
 private var server: CapServer!
 
@@ -67,6 +67,17 @@ public func configure(_ app: Application) throws {
     print("[\(df.string(from: Date()))] Server started (\(server.capCount) caps)")
 }
 
+public func shutdown() {
+    print("[\(df.string(from: Date()))] Server shutdown")
+    asyncScheduler.schedule {
+        do {
+            try await asyncScheduler.shutdownGracefully()
+        } catch {
+            print("Failed to shut down MultiThreadedEventLoopGroup: \(error)")
+        }
+    }
+}
+
 func log(_ message: String) {
     guard let observer = MetricObserver.standard else {
         print(message)
diff --git a/Sources/Run/main.swift b/Sources/Run/main.swift
index e715f81..9163d1e 100755
--- a/Sources/Run/main.swift
+++ b/Sources/Run/main.swift
@@ -4,6 +4,9 @@ import Vapor
 var env = Environment.production
 try LoggingSystem.bootstrap(from: &env)
 let app = Application(env)
-defer { app.shutdown() }
+defer {
+    shutdown()
+    app.shutdown()
+}
 try configure(app)
 try app.run()