Stuff

colony/core/roles/manager.py

@@ -1,19 +1,32 @@
 import threading
 
 from colony.core.context import RoleContext
+from colony.core.roles.policy import DefaultRolePolicy
 from colony.core.state_manager import StateManager
+from colony.role_registry import RolePriority
 from colony.core.logger.logging import get_logger
 
 
 class RoleManager:
-    def __init__(self, network, role_registry, state_manager: StateManager):
+    def __init__(self, network, role_registry, state_manager: StateManager, role_policy=None):
         self.network = network
         self.role_registry = role_registry
         self.state_manager = state_manager
+        self.role_policy = role_policy or DefaultRolePolicy()
 
         self.logger = get_logger(self.__class__.__name__)
         self._lock = threading.RLock()
 
+    def ensure_general_roles(self):
+        general = sorted(
+            name for name, stats in self.role_registry.roles.items()
+            if stats.priority == RolePriority.GENERAL
+        )
+        current = self.state_manager.get(StateManager.Data.Roles)
+        wanted = list(dict.fromkeys(general + current))
+        if wanted != current:
+            self.state_manager.set(StateManager.Data.Roles, wanted)
+
     def initialize_roles(self):
         self.logger.info("Initializing roles...")
         context = RoleContext(self.network, self)
@@ -31,6 +44,7 @@ class RoleManager:
 
             self.network.self_node.roles = list(roles)
             self.network.self_node.set_role_instances(roles)
+            self.network.membership.update_self(roles=list(roles))
         self.logger.info(f"Roles initialized: {list(roles.keys())}. ")
 
     def get_role_instances(self):
@@ -80,6 +94,7 @@ class RoleManager:
             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.network.membership.update_self(roles=list(final_roles))
 
             for role_name in roles_to_remove:
                 self.logger.debug(f"Stopping removed role {role_name}")
@@ -90,3 +105,28 @@ class RoleManager:
                 role.start()
 
         self.logger.info(f"Roles updated to: {list(final_roles.keys())}")
+
+    def request_role(self, role_name):
+        membership = self.network.membership
+
+        with self._lock:
+            if self.role_registry.get_role_class_by_name(role_name) is None:
+                return {"accepted": False, "reason": "unknown role"}
+
+            held = list(self.get_role_instances())
+            if role_name in held:
+                return {"accepted": True, "reason": "already held"}
+
+            if membership.is_declined(role_name):
+                return {"accepted": False, "reason": "recently declined"}
+
+            scores = membership.self_record().scores
+            accepted, reason = self.role_policy.decide(role_name, held, scores, self.role_registry.roles)
+            if not accepted:
+                self.logger.info(f"Declining role {role_name}: {reason}")
+                membership.decline(role_name, self.role_policy.decline_cooldown)
+                return {"accepted": False, "reason": reason}
+
+            self.set_roles(held + [role_name])
+            membership.clear_decline(role_name)
+            return {"accepted": True, "reason": reason}