From 16c272cb1fb641544cdf32d9bf1f6f488e0e85e5 Mon Sep 17 00:00:00 2001
From: tom doron <tomer@apple.com>
Date: Wed, 1 Sep 2021 21:52:34 -0700
Subject: [PATCH] shutdown any task that have shutdownIfNotStarted set to true
 when shutting down a never-started-lifecycle

motivation: more reliable shutdown in edge cases where lifecycle never gets started

changes: when shutting down an idle lifecycle, try to shutdown any tasks that have shutdownIfNotStarted set to true
---
 Sources/Lifecycle/Lifecycle.swift             |  8 ++-
 .../ComponentLifecycleTests+XCTest.swift      |  2 +
 .../ComponentLifecycleTests.swift             | 66 +++++++++++++++++--
 3 files changed, 68 insertions(+), 8 deletions(-)

diff --git a/Sources/Lifecycle/Lifecycle.swift b/Sources/Lifecycle/Lifecycle.swift
index e6824b8..00ac1c9 100644
--- a/Sources/Lifecycle/Lifecycle.swift
+++ b/Sources/Lifecycle/Lifecycle.swift
@@ -507,11 +507,17 @@ public class ComponentLifecycle: LifecycleTask {
 
         self.stateLock.lock()
         switch self.state {
-        case .idle:
+        case .idle(let tasks) where tasks.isEmpty:
             self.state = .shutdown(nil)
             self.stateLock.unlock()
             defer { self.shutdownGroup.leave() }
             callback(nil)
+        case .idle(let tasks):
+            self.stateLock.unlock()
+            // attempt to shutdown any registered tasks
+            let stoppable = tasks.filter { $0.shutdownIfNotStarted }
+            setupShutdownListener(.global())
+            self._shutdown(on: .global(), tasks: stoppable, callback: self.shutdownGroup.leave)
         case .shutdown:
             self.stateLock.unlock()
             self.logger.warning("already shutdown")
diff --git a/Tests/LifecycleTests/ComponentLifecycleTests+XCTest.swift b/Tests/LifecycleTests/ComponentLifecycleTests+XCTest.swift
index 9e9a232..30e2390 100644
--- a/Tests/LifecycleTests/ComponentLifecycleTests+XCTest.swift
+++ b/Tests/LifecycleTests/ComponentLifecycleTests+XCTest.swift
@@ -30,6 +30,8 @@ extension ComponentLifecycleTests {
             ("testUserDefinedCallbackQueue", testUserDefinedCallbackQueue),
             ("testShutdownWhileStarting", testShutdownWhileStarting),
             ("testShutdownWhenIdle", testShutdownWhenIdle),
+            ("testShutdownWhenIdleAndNoItems", testShutdownWhenIdleAndNoItems),
+            ("testIfNotStartedWhenIdle", testIfNotStartedWhenIdle),
             ("testShutdownWhenShutdown", testShutdownWhenShutdown),
             ("testShutdownDuringHangingStart", testShutdownDuringHangingStart),
             ("testShutdownErrors", testShutdownErrors),
diff --git a/Tests/LifecycleTests/ComponentLifecycleTests.swift b/Tests/LifecycleTests/ComponentLifecycleTests.swift
index 2500155..e0eec3c 100644
--- a/Tests/LifecycleTests/ComponentLifecycleTests.swift
+++ b/Tests/LifecycleTests/ComponentLifecycleTests.swift
@@ -159,23 +159,75 @@ final class ComponentLifecycleTests: XCTestCase {
 
     func testShutdownWhenIdle() {
         let lifecycle = ComponentLifecycle(label: "test")
-        lifecycle.register(GoodItem())
 
-        let sempahpore1 = DispatchSemaphore(value: 0)
+        let item = GoodItem()
+        lifecycle.register(item)
+
+        let semaphore1 = DispatchSemaphore(value: 0)
         lifecycle.shutdown { errors in
             XCTAssertNil(errors)
-            sempahpore1.signal()
+            semaphore1.signal()
         }
         lifecycle.wait()
-        XCTAssertEqual(.success, sempahpore1.wait(timeout: .now() + 1))
+        XCTAssertEqual(.success, semaphore1.wait(timeout: .now() + 1))
 
-        let sempahpore2 = DispatchSemaphore(value: 0)
+        let semaphore2 = DispatchSemaphore(value: 0)
         lifecycle.shutdown { errors in
             XCTAssertNil(errors)
-            sempahpore2.signal()
+            semaphore2.signal()
         }
         lifecycle.wait()
-        XCTAssertEqual(.success, sempahpore2.wait(timeout: .now() + 1))
+        XCTAssertEqual(.success, semaphore2.wait(timeout: .now() + 1))
+
+        XCTAssertEqual(item.state, .idle, "expected item to be idle")
+    }
+
+    func testShutdownWhenIdleAndNoItems() {
+        let lifecycle = ComponentLifecycle(label: "test")
+
+        let semaphore1 = DispatchSemaphore(value: 0)
+        lifecycle.shutdown { errors in
+            XCTAssertNil(errors)
+            semaphore1.signal()
+        }
+        lifecycle.wait()
+        XCTAssertEqual(.success, semaphore1.wait(timeout: .now() + 1))
+
+        let semaphore2 = DispatchSemaphore(value: 0)
+        lifecycle.shutdown { errors in
+            XCTAssertNil(errors)
+            semaphore2.signal()
+        }
+        lifecycle.wait()
+        XCTAssertEqual(.success, semaphore2.wait(timeout: .now() + 1))
+    }
+
+    func testIfNotStartedWhenIdle() {
+        var shutdown1Called = false
+        var shutdown2Called = false
+        var shutdown3Called = false
+
+        let lifecycle = ComponentLifecycle(label: "test")
+
+        lifecycle.register(label: "shutdown1",
+                           start: .sync {},
+                           shutdown: .sync { shutdown1Called = true },
+                           shutdownIfNotStarted: true)
+
+        lifecycle.register(label: "shutdown2", start: .none, shutdown: .sync {
+            shutdown2Called = true
+        })
+
+        lifecycle.registerShutdown(label: "shutdown3", .sync {
+            shutdown3Called = true
+        })
+
+        lifecycle.shutdown()
+        lifecycle.wait()
+
+        XCTAssertTrue(shutdown1Called, "expected shutdown to be called")
+        XCTAssertTrue(shutdown2Called, "expected shutdown to be called")
+        XCTAssertTrue(shutdown3Called, "expected shutdown to be called")
     }
 
     func testShutdownWhenShutdown() {