Stuff

colony/core/roles/manager.py

@@ -12,20 +12,25 @@ class RoleManager:
         self.state_manager = state_manager
 
         self.logger = get_logger(self.__class__.__name__)
+        self._lock = threading.RLock()
 
     def initialize_roles(self):
         self.logger.info("Initializing roles...")
         context = RoleContext(self.network, self)
 
-        roles = {}
-        for role_name in self.state_manager.get(StateManager.Data.Roles):
-            role_cls = self.role_registry.get_role_class_by_name(role_name)
-            self.logger.debug(f"Creating instance of role {role_name}. ")
-            role_instance = role_cls(context=context)
-            roles[role_name] = role_instance
+        with self._lock:
+            roles = {}
+            for role_name in self.state_manager.get(StateManager.Data.Roles):
+                role_cls = self.role_registry.get_role_class_by_name(role_name)
+                if role_cls is None:
+                    self.logger.warning(f"Skipping unknown role {role_name}")
+                    continue
 
-        self.network.self_node.roles = [role for role in roles.keys()]
-        self.network.self_node.set_role_instances(roles)
+                self.logger.debug(f"Creating instance of role {role_name}. ")
+                roles[role_name] = role_cls(context=context)
+
+            self.network.self_node.roles = list(roles)
+            self.network.self_node.set_role_instances(roles)
         self.logger.info(f"Roles initialized: {list(roles.keys())}. ")
 
     def get_role_instances(self):
@@ -34,39 +39,54 @@ class RoleManager:
     def start_all(self):
         for role_name, role in self.get_role_instances().items():
             self.logger.info(f"Starting role {role_name}")
-            t = threading.Thread(target=role.start, daemon=True, name=role_name)
-            t.start()
-            self.logger.debug(f"Thread started for role {role_name}: {t.name}")
+            role.start()
+
+    def _validate_roles(self, new_roles):
+        if not isinstance(new_roles, list):
+            raise ValueError("Roles must be a list of role names")
+
+        for role_name in new_roles:
+            if self.role_registry.get_role_class_by_name(role_name) is None:
+                raise ValueError(f"Unknown role {role_name!r}")
+
+        return list(dict.fromkeys(new_roles))
 
     def set_roles(self, new_roles):
         self.logger.info("Updating roles...")
 
-        current_roles = self.get_role_instances()
-        roles_to_remove = set(current_roles) - set(new_roles)
-        roles_to_add = set(new_roles) - set(current_roles)
+        new_roles = self._validate_roles(new_roles)
 
-        for role_name in roles_to_remove:
-            self.logger.debug(f"Stopping and removing role {role_name}")
-            current_roles[role_name].stop()
-            del current_roles[role_name]
+        with self._lock:
+            current_roles = self.get_role_instances()
+            roles_to_remove = [name for name in current_roles if name not in new_roles]
+            roles_to_add = [name for name in new_roles if name not in current_roles]
 
-        context = RoleContext(self.network, self)
-        for role_name in roles_to_add:
-            self.logger.debug(f"Creating and adding new role {role_name}")
-            role_cls = self.role_registry.get_role_class_by_name(role_name)
-            current_roles[role_name] = role_cls(context=context)
+            if not roles_to_remove and not roles_to_add:
+                self.logger.info("Roles already up to date")
+                return
 
-        self.state_manager.set(StateManager.Data.Roles, list(current_roles.keys()))
+            context = RoleContext(self.network, self)
+            added = {}
+            for role_name in roles_to_add:
+                self.logger.debug(f"Creating new role {role_name}")
+                role_cls = self.role_registry.get_role_class_by_name(role_name)
+                added[role_name] = role_cls(context=context)
 
-        self.network.self_node.roles = list(current_roles.keys())
-        self.network.self_node.set_role_instances(current_roles)
+            final_roles = {
+                name: current_roles[name] if name in current_roles else added[name]
+                for name in new_roles
+            }
 
-        for role_name in roles_to_add:
-            role = current_roles[role_name]
-            self.logger.info(f"Starting role {role_name}")
-            t = threading.Thread(target=role.start, daemon=True, name=role_name)
-            t.start()
-            self.logger.debug(f"Thread started for role {role_name}: {t.name}")
+            self.state_manager.set(StateManager.Data.Roles, list(final_roles.keys()))
+            self.network.self_node.set_role_instances(final_roles)
+            self.network.self_node.roles = list(final_roles.keys())
 
-        self.logger.info(f"Roles updated to: {list(current_roles.keys())}")
+            for role_name in roles_to_remove:
+                self.logger.debug(f"Stopping removed role {role_name}")
+                current_roles[role_name].stop()
 
+            for role_name, role in added.items():
+                self.logger.info(f"Starting role {role_name}")
+                role.start()
+
+        self.logger.info(f"Roles updated to: {list(final_roles.keys())}")