Add Async Storage shutdown

Sources/Development/routes.swift

@@ -268,6 +268,7 @@ public func routes(_ app: Application) throws {
         return [cred1]
     }
 
+    @Sendable
     func opaqueRouteTester(_ req: Request) async throws -> some AsyncResponseEncodable {
         "Hello World"
     }

Sources/Vapor/Application.swift

@@ -120,6 +120,7 @@ public final class Application: Sendable {
         _ eventLoopGroupProvider: EventLoopGroupProvider = .singleton
     ) {
         self.init(environment, eventLoopGroupProvider, async: false)
+        self.asyncCommands.use(self.servers.command, as: "serve", isDefault: true)
         DotEnvFile.load(for: environment, on: .shared(self.eventLoopGroup), fileio: self.fileio, logger: self.logger)
     }
 
@@ -155,12 +156,12 @@ public final class Application: Sendable {
         self.servers.use(.http)
         self.clients.initialize()
         self.clients.use(.http)
-        self.asyncCommands.use(self.servers.command, as: "serve", isDefault: true)
         self.asyncCommands.use(RoutesCommand(), as: "routes")
     }
 
     public static func make(_ environment: Environment = .development, _ eventLoopGroupProvider: EventLoopGroupProvider = .singleton) async throws -> Application {
         let app = Application(environment, eventLoopGroupProvider, async: true)
+        await app.asyncCommands.use(app.servers.asyncCommand, as: "serve", isDefault: true)
         await DotEnvFile.load(for: app.environment, fileio: app.fileio, logger: app.logger)
         return app
     }
@@ -268,8 +269,11 @@ public final class Application: Sendable {
 
         self.logger.trace("Shutting down providers")
         self.lifecycle.handlers.reversed().forEach { $0.shutdown(self) }
-
-        triggerShutdown()
+        self.lifecycle.handlers = []
+
+        self.logger.trace("Clearing Application storage")
+        self.storage.shutdown()
+        self.storage.clear()
 
         switch self.eventLoopGroupProvider {
         case .shared:
@@ -295,8 +299,11 @@ public final class Application: Sendable {
         for handler in self.lifecycle.handlers.reversed()  {
             await handler.shutdownAsync(self)
         }
-
-        triggerShutdown()
+        self.lifecycle.handlers = []
+
+        self.logger.trace("Clearing Application storage")
+        await self.storage.asyncShutdown()
+        self.storage.clear()
 
         switch self.eventLoopGroupProvider {
         case .shared:
@@ -313,14 +320,6 @@ public final class Application: Sendable {
         self._didShutdown.withLockedValue { $0 = true }
         self.logger.trace("Application shutdown complete")
     }
-
-    private func triggerShutdown() {
-        self.lifecycle.handlers = []
-
-        self.logger.trace("Clearing Application storage")
-        self.storage.shutdown()
-        self.storage.clear()
-    }
 
     deinit {
         self.logger.trace("Application deinitialized, goodbye!")

Sources/Vapor/Commands/ServeCommand.swift

@@ -103,6 +103,7 @@ public final class ServeCommand: AsyncCommand, Sendable {
         self.box.withLockedValue { $0 = box }
     }
 
+    @available(*, noasync, message: "Use the async asyncShutdown() method instead.")
     func shutdown() {
         var box = self.box.withLockedValue { $0 }
         box.didShutdown = true
@@ -115,6 +116,18 @@ public final class ServeCommand: AsyncCommand, Sendable {
         self.box.withLockedValue { $0 = box }
     }
 
+    func asyncShutdown() async {
+        var box = self.box.withLockedValue { $0 }
+        box.didShutdown = true
+        box.running?.stop()
+        if let server = box.server {
+            await server.shutdown()
+        }
+        box.signalSources.forEach { $0.cancel() } // clear refs
+        box.signalSources = []
+        self.box.withLockedValue { $0 = box }
+    }
+
     deinit {
         assert(self.box.withLockedValue({ $0.didShutdown }), "ServeCommand did not shutdown before deinit")
     }

Sources/Vapor/Server/Application+Servers.swift

@@ -51,6 +51,7 @@ extension Application {
             self.storage.makeServer.withLockedValue { $0 = .init(factory: makeServer) }
         }
 
+        @available(*, noasync, renamed: "asyncCommand", message: "Use the async property instead.")
         public var command: ServeCommand {
             if let existing = self.application.storage.get(CommandKey.self) {
                 return existing
@@ -62,6 +63,20 @@ extension Application {
                 return new
             }
         }
+
+        public var asyncCommand: ServeCommand {
+            get async {
+                if let existing = self.application.storage.get(CommandKey.self) {
+                    return existing
+                } else {
+                    let new = ServeCommand()
+                    await self.application.storage.setWithAsyncShutdown(CommandKey.self, to: new) {
+                        await $0.asyncShutdown()
+                    }
+                    return new
+                }
+            }
+        }
 
         let application: Application
 

Sources/Vapor/Utilities/Storage.swift

@@ -11,13 +11,21 @@ public struct Storage: Sendable {
     struct Value<T: Sendable>: AnyStorageValue {
         var value: T
         var onShutdown: (@Sendable (T) throws -> ())?
+        var onAsyncShutdown: (@Sendable (T) async throws -> ())?
         func shutdown(logger: Logger) {
             do {
                 try self.onShutdown?(self.value)
             } catch {
                 logger.warning("Could not shutdown \(T.self): \(error)")
             }
         }
+        func asyncShutdown(logger: Logger) async {
+            do {
+                try await self.onAsyncShutdown?(self.value)
+            } catch {
+                logger.warning("Could not shutdown \(T.self): \(error)")
+            }
+        }
     }
 
     /// The logger provided to shutdown closures.
@@ -94,20 +102,49 @@ public struct Storage: Sendable {
             existing.shutdown(logger: self.logger)
         }
     }
+
+    /// Set or remove a value for a given key, optionally providing an async shutdown closure for the value.
+    ///
+    /// If a key that has a shutdown closure is removed by this method, the closure **is** invoked.
+    public mutating func setWithAsyncShutdown<Key>(
+        _ key: Key.Type,
+        to value: Key.Value?,
+        onShutdown: (@Sendable (Key.Value) async throws -> ())? = nil
+    ) async
+        where Key: StorageKey
+    {
+        let key = ObjectIdentifier(Key.self)
+        if let value = value {
+            self.storage[key] = Value(value: value, onShutdown: nil, onAsyncShutdown: onShutdown)
+        } else if let existing = self.storage[key] {
+            self.storage[key] = nil
+            await existing.asyncShutdown(logger: self.logger)
+        }
+    }
 
     /// For every key in the container having a shutdown closure, invoke the closure. Designed to
     /// be invoked during an explicit app shutdown process or in a reference type's `deinit`.
+    @available(*, noasync, message: "Use the async asyncShutdown() method instead.")
     public func shutdown() {
         self.storage.values.forEach {
             $0.shutdown(logger: self.logger)
         }
     }
+
+    /// For every key in the container having a shutdown closure, invoke the closure. Designed to
+    /// be invoked during an explicit app shutdown process or in a reference type's `deinit`.
+    public func asyncShutdown() async {
+        for value in self.storage.values {
+            await value.asyncShutdown(logger: self.logger)
+        }
+    }
 }
 
 /// ``Storage`` uses this protocol internally to generically invoke shutdown closures for arbitrarily-
 /// typed key values.
 protocol AnyStorageValue: Sendable {
     func shutdown(logger: Logger)
+    func asyncShutdown(logger: Logger) async
 }
 
 /// A key used to store values in a ``Storage`` must conform to this protocol.