Fix Graceful shutdown

Author: Andrea Scuderi

Sources/BreezeLambdaAPI/BreezeLambdaAPI.swift

@@ -57,5 +57,6 @@ public actor BreezeLambdaAPI<T: BreezeCodable>: Service {
     public func run() async throws {
         logger.info("Starting BreezeLambdaAPI...")
         try await serviceGroup.run()
+        logger.info("BreezeLambdaAPI is stopped successfully")
     }
 }

Sources/BreezeLambdaAPI/BreezeLambdaService.swift

@@ -46,25 +46,40 @@ actor BreezeLambdaService<T: BreezeCodable>: Service {
 
     func run() async throws {
         let dbManager = await dynamoDBService.dbManager()
-        try await withGracefulShutdownHandler {
+        let breezeApi = BreezeLambdaHandler<T>(dbManager: dbManager, operation: self.operation)
+        self.breezeApi = breezeApi
+        logger.info("Starting BreezeLambdaService...")
+        let runtime = LambdaRuntime(body: handler)
+        try await runTaskWithCancellationOnGracefulShutdown {
             do {
-                let breezeApi = BreezeLambdaHandler<T>(dbManager: dbManager, operation: operation)
-                self.breezeApi = breezeApi
-                logger.info("Starting BreezeLambdaService...")
-                logger.info("Starting BreezeLambdaService...")
-                let runtime = LambdaRuntime(body: handler)
                 try await runtime.run()
             } catch {
-                logger.error("\(error.localizedDescription)")
+                self.logger.error("\(error.localizedDescription)")
                 throw error
             }
         } onGracefulShutdown: {
-            Task {
-                self.logger.info("Gracefully stoping BreezeLambdaService ...")
-                try await self.dynamoDBService.gracefulShutdown()
-                self.logger.info("BreezeLambdaService stopped.")
-                exit(EXIT_SUCCESS)
+            self.logger.info("Gracefully stoping BreezeLambdaService ...")
+            try await self.dynamoDBService.gracefulShutdown()
+            self.logger.info("BreezeLambdaService is stopped.")
+        }
+    }
+    
+    private func runTaskWithCancellationOnGracefulShutdown(
+        operation: @escaping @Sendable () async throws -> Void,
+        onGracefulShutdown: () async throws -> Void
+    ) async throws {
+        let (cancelOrGracefulShutdown, cancelOrGracefulShutdownContinuation) = AsyncStream<Void>.makeStream()
+        let task = Task {
+            try await withTaskCancellationOrGracefulShutdownHandler {
+                try await operation()
+            } onCancelOrGracefulShutdown: {
+                cancelOrGracefulShutdownContinuation.yield()
+                cancelOrGracefulShutdownContinuation.finish()
             }
         }
+        for await _ in cancelOrGracefulShutdown {
+            try await onGracefulShutdown()
+            task.cancel()
+        }
     }
 }

Tests/BreezeLambdaAPITests/BreezeLambdaAPIServiceTests.swift

@@ -35,9 +35,6 @@ struct BreezeLambdaAPIServiceTests {
 
     let logger = Logger(label: "BreezeHTTPClientServiceTests")
 
-    /*
-    This test is commented out because it requires to manage the process exit(EXIT_SUCCESS).
-     
     @Test
     func test_breezeLambdaAPIService_whenValidEnvironment() async throws {
         try await testGracefulShutdown { gracefulShutdownTestTrigger in
@@ -62,12 +59,11 @@ struct BreezeLambdaAPIServiceTests {
                 }
                 for await _ in gracefulStream {
                     logger.info("Graceful shutdown stream received")
-                    await group.cancelAll()
+                    group.cancelAll()
                 }
             }
         }
     }
-    */
 
     @Test
     func test_breezeLambdaAPIService_whenInvalidEnvironment() async throws {