!context store lifecycle in DistributedActorContext; internal only (#976)

* !context store lifecycle in DistributedActorContext; internal only

* cleanup

* wip

* formatting

* trying CI workarounds

* workarounds

* formatting

Author: ktoso

Sources/ActorSingletonPlugin/ActorSingletonManager.swift

@@ -100,7 +100,7 @@ extension ActorSingletonManager {
 
 extension ActorID {
     static func _singletonManager(name: String, on node: UniqueNode) -> ActorID {
-        ActorID(local: node, path: ._singletonManager(name: name), incarnation: .wellKnown)
+        ._make(local: node, path: ._singletonManager(name: name), incarnation: .wellKnown)
     }
 }
 

Sources/ActorSingletonPlugin/ActorSingletonProxy.swift

@@ -239,7 +239,7 @@ extension ActorSingletonProxy {
 
 extension ActorID {
     static func _singletonProxy(name: String, remote node: UniqueNode) -> ActorID {
-        .init(remote: node, path: ._singletonProxy(name: name), incarnation: .wellKnown)
+        ._make(remote: node, path: ._singletonProxy(name: name), incarnation: .wellKnown)
     }
 }
 

Sources/DistributedActors/ActorAddress.swift renamed to Sources/DistributedActors/ActorID.swift

@@ -99,6 +99,9 @@ extension ClusterSystem {
         /// - SeeAlso: `ActorTags` for a detailed discussion of some frequently used tags.
         public var tags: ActorTags
 
+        /// Internal "actor context" which is used as storage for additional cluster actor features, such as watching.
+        internal var context: DistributedActorContext
+
         /// Underlying path representation, not attached to a specific Actor instance.
         // FIXME(distributed): make optional
         public var path: ActorPath {
@@ -123,30 +126,59 @@ extension ClusterSystem {
         /// Uniquely identifies the specific "incarnation" of this actor.
         public let incarnation: ActorIncarnation
 
-        /// :nodoc:
-        public init(local node: UniqueNode, path: ActorPath?, incarnation: ActorIncarnation) {
+        // TODO(distributed): remove this workaround; only exists for old ActorSingletonManager
+        public static func _make(local node: UniqueNode, path: ActorPath?, incarnation: ActorIncarnation) -> Self {
+            .init(local: node, path: path, incarnation: incarnation)
+        }
+
+        // TODO(distributed): remove this initializer, as it is only for Behavior actors
+        init(local node: UniqueNode, path: ActorPath?, incarnation: ActorIncarnation) {
+            self.context = .init(lifecycle: nil)
             self._location = .local(node)
             self.tags = ActorTags()
             self.incarnation = incarnation
             if let path {
                 self.tags[ActorTags.path] = path
             }
+            traceLog_DeathWatch("Made ID: \(self)")
         }
 
-        /// :nodoc:
-        public init(remote node: UniqueNode, path: ActorPath?, incarnation: ActorIncarnation) {
+        // TODO(distributed): remove this workaround; only exists for old ActorSingletonManager
+        public static func _make(remote node: UniqueNode, path: ActorPath?, incarnation: ActorIncarnation) -> Self {
+            .init(remote: node, path: path, incarnation: incarnation)
+        }
+
+        // TODO(distributed): remove this initializer, as it is only for Behavior actors
+        init(remote node: UniqueNode, path: ActorPath?, incarnation: ActorIncarnation) {
+            self.context = .init(lifecycle: nil)
             self._location = .remote(node)
             self.incarnation = incarnation
             self.tags = ActorTags()
             if let path {
                 self.tags[ActorTags.path] = path
             }
+            traceLog_DeathWatch("Made ID: \(self)")
         }
 
-        /// :nodoc:
-        public init<Act>(local node: UniqueNode, type: Act.Type, incarnation: ActorIncarnation)
+        public init<Act>(remote node: UniqueNode, type: Act.Type, incarnation: ActorIncarnation)
             where Act: DistributedActor, Act.ActorSystem == ClusterSystem
         {
+            self.context = .init(lifecycle: nil)
+            self._location = .remote(node)
+            self.incarnation = incarnation
+            self.tags = ActorTags()
+            // TODO: avoid mangling names on every spawn?
+            if let mangledName = _mangledTypeName(type) {
+                self.tags[ActorTags.type] = .init(mangledName: mangledName)
+            }
+            traceLog_DeathWatch("Made ID: \(self)")
+        }
+
+        init<Act>(local node: UniqueNode, type: Act.Type, incarnation: ActorIncarnation,
+                  context: DistributedActorContext)
+            where Act: DistributedActor, Act.ActorSystem == ClusterSystem
+        {
+            self.context = context
             self._location = .local(node)
             self.tags = ActorTags()
             self.incarnation = incarnation
@@ -155,18 +187,36 @@ extension ClusterSystem {
             if let mangledName = _mangledTypeName(type) {
                 self.tags[ActorTags.type] = .init(mangledName: mangledName)
             }
+            traceLog_DeathWatch("Made ID: \(self)")
         }
 
-        /// :nodoc:
-        public init(remote node: UniqueNode, type: (some DistributedActor).Type, incarnation: ActorIncarnation) {
+        init<Act>(remote node: UniqueNode, type: Act.Type, incarnation: ActorIncarnation,
+                  context: DistributedActorContext)
+            where Act: DistributedActor, Act.ActorSystem == ClusterSystem
+        {
+            self.context = context
             self._location = .remote(node)
             self.incarnation = incarnation
             self.tags = ActorTags()
             // TODO: avoid mangling names on every spawn?
             if let mangledName = _mangledTypeName(type) {
                 self.tags[ActorTags.type] = .init(mangledName: mangledName)
             }
+            traceLog_DeathWatch("Made ID: \(self)")
         }
+
+        internal var withoutContext: Self {
+            var copy = self
+            copy.context = .init(lifecycle: nil)
+            return copy
+        }
+    }
+}
+
+extension DistributedActor where ActorSystem == ClusterSystem {
+    /// INTERNAL: Provides the actor context for use within this actor.
+    internal var context: DistributedActorContext {
+        self.id.context
     }
 }
 
@@ -237,7 +287,6 @@ extension ActorID {
 
 extension ActorID {
     /// :nodoc:
-    @inlinable
     public var _isLocal: Bool {
         switch self._location {
         case .local: return true
@@ -246,20 +295,17 @@ extension ActorID {
     }
 
     /// :nodoc:
-    @inlinable
     public var _isRemote: Bool {
         !self._isLocal
     }
 
     /// :nodoc:
-    @inlinable
     public var _asRemote: Self {
         let remote = Self(remote: self.uniqueNode, path: self.path, incarnation: self.incarnation)
         return remote
     }
 
     /// :nodoc:
-    @inlinable
     public var _asLocal: Self {
         let local = Self(local: self.uniqueNode, path: self.path, incarnation: self.incarnation)
         return local

Sources/DistributedActors/ClusterSystem.swift

@@ -55,9 +55,6 @@ public class ClusterSystem: DistributedActorSystem, @unchecked Sendable {
     // TODO: collapse it with the other initialization lock; the other one is not needed now I think?
     private let initLock = Lock()
 
-    internal let lifecycleWatchLock = Lock()
-    internal var _lifecycleWatches: [ActorID: LifecycleWatchContainer] = [:]
-
     private var _associationTombstoneCleanupTask: RepeatedTask?
 
     private let dispatcher: InternalMessageDispatcher
@@ -907,17 +904,27 @@ extension ClusterSystem {
         where Act: DistributedActor
     {
         let props = _Props.forSpawn // task-local read for any properties this actor should have
-        let address = try! self._reserveName(type: Act.self, props: props)
+        var id = try! self._reserveName(type: Act.self, props: props)
+
+        let lifecycleContainer: LifecycleWatchContainer?
+        if Act.self is (any(LifecycleWatch).Type) {
+            lifecycleContainer = LifecycleWatchContainer(watcherID: id.withoutContext, actorSystem: self)
+        } else {
+            lifecycleContainer = nil
+        }
+        traceLog_DeathWatch("Make LifecycleWatchContainer for \(id):::: \(lifecycleContainer)")
+
+        id.context = .init(lifecycle: lifecycleContainer)
 
         self.log.warning("Assign identity", metadata: [
             "actor/type": "\(actorType)",
-            "actor/id": "\(address)",
-            "actor/id/uniqueNode": "\(address.uniqueNode)",
+            "actor/id": "\(id)",
+            "actor/id/uniqueNode": "\(id.uniqueNode)",
         ])
 
         return self.namingLock.withLock {
-            self._reservedNames.insert(address)
-            return address
+            self._reservedNames.insert(id)
+            return id
         }
     }
 
@@ -931,12 +938,12 @@ extension ClusterSystem {
         defer { self.namingLock.unlock() }
         precondition(self._reservedNames.remove(actor.id) != nil, "Attempted to ready an identity that was not reserved: \(actor.id)")
 
-        if let watcher = actor as? any LifecycleWatch {
-            func doMakeLifecycleWatch<Watcher: LifecycleWatch & DistributedActor>(watcher: Watcher) {
-                _ = self._makeLifecycleWatch(watcher: watcher)
-            }
-            _openExistential(watcher, do: doMakeLifecycleWatch)
-        }
+//        if let watcher = actor as? any LifecycleWatch {
+//            func doMakeLifecycleWatch<Watcher: LifecycleWatch & DistributedActor>(watcher: Watcher) {
+//                _ = LifecycleWatchContainer(watcher)
+//            }
+//            _openExistential(watcher, do: doMakeLifecycleWatch)
+//        }
 
         let behavior = InvocationBehavior.behavior(instance: Weak(actor))
         let ref = self._spawnDistributedActor(behavior, identifiedBy: actor.id)
@@ -953,11 +960,12 @@ extension ClusterSystem {
                 ref._sendSystemMessage(.stop, file: #file, line: #line)
             }
         }
-        self.lifecycleWatchLock.withLockVoid {
-            if let watch = self._lifecycleWatches.removeValue(forKey: id) {
-                watch.notifyWatchersWeDied()
-            }
-        }
+        id.context.terminate()
+//        self.lifecycleWatchLock.withLockVoid {
+//            if let watch = self._lifecycleWatches.removeValue(forKey: id) {
+//                watch.notifyWatchersWeDied()
+//            }
+//        }
         self.namingLock.withLockVoid {
             self._managedRefs.removeValue(forKey: id) // TODO: should not be necessary in the future
             _ = self._managedDistributedActors.removeActor(identifiedBy: id)

Sources/DistributedActors/DistributedActorContext.swift

@@ -0,0 +1,38 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift Distributed Actors open source project
+//
+// 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
+// See CONTRIBUTORS.md for the list of Swift Distributed Actors project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+import Distributed
+
+/// Internal context object used by the actor system to support per-actor state, such as necessary to implement lifecycle watch etc.
+///
+/// Access must be carefully managed to only be performed from the actor itself, once the context has been initialized.
+/// And also it is allowed to modify this state once the system receives it in the `resignID` (as the context is carried inside the ID),
+/// as at that point in time it can no longer be used–by the now deallocated–actor itself.
+public final class DistributedActorContext {
+    let lifecycle: LifecycleWatchContainer?
+
+    init(lifecycle: LifecycleWatchContainer?) {
+        self.lifecycle = lifecycle
+        traceLog_DeathWatch("Create context; Lifecycle: \(lifecycle)")
+    }
+
+    /// Invoked by the actor system when the owning actor is terminating, so we can clean up all stored data
+    func terminate() {
+        guard let lifecycle = self.lifecycle else {
+            return
+        }
+        traceLog_DeathWatch("Terminate: \(lifecycle.watcherID)")
+        lifecycle.clear()
+    }
+}

Sources/DistributedActors/InvocationBehavior.swift

@@ -43,8 +43,8 @@ enum InvocationBehavior {
                 }
 
                 // `InvocationMessage`s are handled in `UserMessageHandler`
-//                 await context.system.receiveInvocation(actor: instance, message: message)
-                return .same
+                // old impl: await context.system.receiveInvocation(actor: instance, message: message)
+                return fatalErrorBacktrace("We don't invoke distributed actors via the behavior runtime anymore! ")
             }.receiveSignal { _, signal in
 
                 // We received a signal, but our target actor instance was already released;
@@ -55,8 +55,13 @@ enum InvocationBehavior {
 
                 if let terminated = signal as? _Signals.Terminated {
                     if let watcher = instance as? (any LifecycleWatch) {
-                        let watch = watcher.actorSystem._getLifecycleWatch(watcher: watcher)
-                        watch?.receiveTerminated(terminated)
+                        Task {
+                            await instance.whenLocal { __secretlyKnownToBeLocalK in
+                                let __secretlyKnownToBeLocal: any LifecycleWatch = __secretlyKnownToBeLocalK as! any LifecycleWatch
+                                await __secretlyKnownToBeLocal._receiveActorTerminated(id: terminated.id)
+                            }
+                        }
+
                         return .same
                     }
                 }