133 lines
5.3 KiB
Python
133 lines
5.3 KiB
Python
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, 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)
|
|
|
|
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.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.network.membership.update_self(roles=list(roles))
|
|
self.logger.info(f"Roles initialized: {list(roles.keys())}. ")
|
|
|
|
def get_role_instances(self):
|
|
return self.network.self_node.get_role_instances()
|
|
|
|
def start_all(self):
|
|
for role_name, role in self.get_role_instances().items():
|
|
self.logger.info(f"Starting role {role_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...")
|
|
|
|
new_roles = self._validate_roles(new_roles)
|
|
|
|
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]
|
|
|
|
if not roles_to_remove and not roles_to_add:
|
|
self.logger.info("Roles already up to date")
|
|
return
|
|
|
|
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)
|
|
|
|
final_roles = {
|
|
name: current_roles[name] if name in current_roles else added[name]
|
|
for name in new_roles
|
|
}
|
|
|
|
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}")
|
|
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())}")
|
|
|
|
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}
|