diff --git a/buildman/manager/enterprise.py b/buildman/manager/enterprise.py index d7fdea39a..c56830a1c 100644 --- a/buildman/manager/enterprise.py +++ b/buildman/manager/enterprise.py @@ -5,7 +5,7 @@ from buildman.component.basecomponent import BaseComponent from buildman.component.buildcomponent import BuildComponent from buildman.manager.basemanager import BaseManager -from trollius import From, Return, coroutine, async +from trollius import From, Return, coroutine REGISTRATION_REALM = 'registration' logger = logging.getLogger(__name__) @@ -31,8 +31,11 @@ class EnterpriseManager(BaseManager): def __init__(self, *args, **kwargs): self.ready_components = set() + self.all_components = set() self.shutting_down = False + super(EnterpriseManager, self).__init__(*args, **kwargs) + def initialize(self, manager_config): # Add a component which is used by build workers for dynamic registration. Unlike # production, build workers in enterprise are long-lived and register dynamically. @@ -47,7 +50,8 @@ class EnterpriseManager(BaseManager): """ Adds a new build component for an Enterprise Registry. """ # Generate a new unique realm ID for the build worker. realm = str(uuid.uuid4()) - self.register_component(realm, BuildComponent, token="") + new_component = self.register_component(realm, BuildComponent, token="") + self.all_components.add(new_component) return realm @coroutine @@ -74,8 +78,9 @@ class EnterpriseManager(BaseManager): self.job_complete_callback(build_job, job_status) def build_component_disposed(self, build_component, timed_out): + self.all_components.remove(build_component) if build_component in self.ready_components: self.ready_components.remove(build_component) def num_workers(self): - return len(self.build_components) + return len(self.all_components)