watchTermination should support async 'whenTerminated' (#939)

Resolves #930

Author: yim-lee

Sources/DistributedActors/ClusterSystem.swift

@@ -169,6 +169,8 @@ public class ClusterSystem: DistributedActorSystem, @unchecked Sendable {
         return box.value
     }
 
+    internal var downing: DowningStrategyShell?
+
     // ==== ----------------------------------------------------------------------------------------------------------------
     // MARK: Logging
 
@@ -358,10 +360,14 @@ public class ClusterSystem: DistributedActorSystem, @unchecked Sendable {
         }
 
         // downing strategy (automatic downing)
-        await _Props.$forSpawn.withValue(DowningStrategyShell.props) {
-            if let downingStrategy = self.settings.downingStrategy.make(self.settings) {
-                let downingStrategyShell = await DowningStrategyShell(downingStrategy, system: self)
+        if settings.enabled {
+            await _Props.$forSpawn.withValue(DowningStrategyShell.props) {
+                if let downingStrategy = self.settings.downingStrategy.make(self.settings) {
+                    self.downing = await DowningStrategyShell(downingStrategy, system: self)
+                }
             }
+        } else {
+            self.downing = nil
         }
 
         #if SACT_TESTS_LEAKS
@@ -486,6 +492,7 @@ public class ClusterSystem: DistributedActorSystem, @unchecked Sendable {
             self.userProvider.stopAll()
             self.systemProvider.stopAll()
             self.dispatcher.shutdown()
+            self.downing = nil
 
             self._associationTombstoneCleanupTask?.cancel()
             self._associationTombstoneCleanupTask = nil

Sources/DistributedActors/LifecycleMonitoring/LifecycleWatch.swift

@@ -27,7 +27,7 @@ extension LifecycleWatch {
     @discardableResult
     public func watchTermination<Watchee>(
         of watchee: Watchee,
-        @_inheritActorContext @_implicitSelfCapture whenTerminated: @escaping @Sendable (ID) -> Void,
+        @_inheritActorContext @_implicitSelfCapture whenTerminated: @escaping @Sendable (ID) async -> Void,
         file: String = #file, line: UInt = #line
     ) -> Watchee where Watchee: DistributedActor, Watchee.ActorSystem == ClusterSystem {
         // TODO(distributed): reimplement this as self.id as? _ActorContext which will have the watch things.
@@ -163,7 +163,7 @@ extension LifecycleWatchContainer {
     /// Performed by the sending side of "watch", therefore the `watcher` should equal `context.myself`
     public func termination<Watchee>(
         of watchee: Watchee,
-        @_inheritActorContext @_implicitSelfCapture whenTerminated: @escaping @Sendable (ClusterSystem.ActorID) -> Void,
+        @_inheritActorContext @_implicitSelfCapture whenTerminated: @escaping @Sendable (ClusterSystem.ActorID) async -> Void,
         file: String = #file, line: UInt = #line
     ) where Watchee: DistributedActor, Watchee.ActorSystem == ClusterSystem {
         traceLog_DeathWatch("issue watch: \(watchee) (from \(self.watcherID))")

Tests/DistributedActorsTests/LifecycleWatchTests.swift

@@ -2,7 +2,7 @@
 //
 // This source file is part of the Swift Distributed Actors open source project
 //
-// Copyright (c) 2018-2021 Apple Inc. and the Swift Distributed Actors project authors
+// Copyright (c) 2018-2022 Apple Inc. and the Swift Distributed Actors project authors
 // Licensed under Apache License v2.0
 //
 // See LICENSE.txt for license information
@@ -71,6 +71,22 @@ distributed actor Juliet: LifecycleWatch, CustomStringConvertible {
         }
     }
 
+    distributed func meetWatchAsyncCallback(
+        _ romeo: Romeo,
+        unwatch doUnwatch: Bool
+    ) async throws {
+        @Sendable
+        func asyncTerminated(_ terminatedIdentity: ClusterSystem.ActorID) async {
+            await self.probe.tell("Received terminated: \(terminatedIdentity)")
+        }
+
+        watchTermination(of: romeo, whenTerminated: asyncTerminated)
+
+        if doUnwatch {
+            unwatch(romeo)
+        }
+    }
+
     nonisolated var description: String {
         "\(Self.self)(\(id))"
     }
@@ -99,6 +115,25 @@ final class LifecycleWatchTests: ActorSystemXCTestCase, @unchecked Sendable {
         try pj.expectMessage("Received terminated: /user/Romeo-b")
     }
 
+    func test_watch_shouldTriggerTerminatedWhenWatchedActorDeinits_async() async throws {
+        let pj = self.testKit.makeTestProbe(expecting: String.self)
+        let pr = self.testKit.makeTestProbe(expecting: String.self)
+        let juliet = Juliet(probe: pj, actorSystem: system)
+
+        func meet() async throws {
+            var romeo: Romeo? = Romeo(probe: pr, actorSystem: system)
+
+            try await juliet.meetWatchAsyncCallback(romeo!, unwatch: false)
+            romeo = nil
+        }
+        try await meet()
+
+        try pj.expectMessage("Juliet init")
+        try pr.expectMessage("Romeo init")
+        try pr.expectMessage("Romeo deinit")
+        try pj.expectMessage("Received terminated: /user/Romeo-b")
+    }
+
     func test_watchThenUnwatch_shouldTriggerTerminatedWhenWatchedActorDeinits() async throws {
         let pj = self.testKit.makeTestProbe(expecting: String.self)
         let pr = self.testKit.makeTestProbe(expecting: String.self)