diff --git a/src/api-service/__app__/onefuzzlib/workers/pools.py b/src/api-service/__app__/onefuzzlib/workers/pools.py index 409b9f8331..dd1f83a157 100644 --- a/src/api-service/__app__/onefuzzlib/workers/pools.py +++ b/src/api-service/__app__/onefuzzlib/workers/pools.py @@ -135,8 +135,7 @@ def get_pool_queue(self) -> str: def init(self) -> None: create_queue(self.get_pool_queue(), StorageType.corpus) ShrinkQueue(self.pool_id).create() - self.state = PoolState.running - self.save() + self.set_state(PoolState.running) def schedule_workset(self, work_set: WorkSet) -> bool: # Don't schedule work for pools that can't and won't do work. @@ -183,15 +182,10 @@ def search_states(cls, *, states: Optional[List[PoolState]] = None) -> List["Poo return cls.search(query=query) def set_shutdown(self, now: bool) -> None: - if self.state in [PoolState.halt, PoolState.shutdown]: - return - if now: - self.state = PoolState.halt + self.set_state(PoolState.halt) else: - self.state = PoolState.shutdown - - self.save() + self.set_state(PoolState.shutdown) def shutdown(self) -> None: """shutdown allows nodes to finish current work then delete""" @@ -202,8 +196,6 @@ def shutdown(self) -> None: nodes = Node.search(query={"pool_name": [self.name]}) if not scalesets and not nodes: logging.info("pool stopped, deleting: %s", self.name) - - self.state = PoolState.halt self.delete() return @@ -227,7 +219,6 @@ def halt(self) -> None: delete_queue(self.get_pool_queue(), StorageType.corpus) ShrinkQueue(self.pool_id).delete() logging.info("pool stopped, deleting: %s", self.name) - self.state = PoolState.halt self.delete() return @@ -239,6 +230,17 @@ def halt(self) -> None: self.save() + def set_state(self, state: PoolState) -> None: + if self.state == state: + return + + # scalesets should never leave the `halt` state + if self.state == PoolState.halt: + return + + self.state = state + self.save() + @classmethod def key_fields(cls) -> Tuple[str, str]: return ("name", "pool_id") diff --git a/src/api-service/__app__/onefuzzlib/workers/scalesets.py b/src/api-service/__app__/onefuzzlib/workers/scalesets.py index 4aa20a5751..2e1921bcef 100644 --- a/src/api-service/__app__/onefuzzlib/workers/scalesets.py +++ b/src/api-service/__app__/onefuzzlib/workers/scalesets.py @@ -635,22 +635,11 @@ def reimage_nodes(self, nodes: List[Node]) -> None: node.delete() def set_shutdown(self, now: bool) -> None: - if self.state in [ScalesetState.halt, ScalesetState.shutdown]: - return - - logging.info( - SCALESET_LOG_PREFIX + "scaleset set_shutdown: scaleset_id:%s now:%s", - self.scaleset_id, - now, - ) - if now: self.set_state(ScalesetState.halt) else: self.set_state(ScalesetState.shutdown) - self.save() - def shutdown(self) -> None: size = get_vmss_size(self.scaleset_id) if size is None: @@ -807,6 +796,10 @@ def set_state(self, state: ScalesetState) -> None: if self.state == state: return + # scalesets should never leave the `halt` state + if self.state == ScalesetState.halt: + return + self.state = state self.save() if self.state == ScalesetState.resize: