ExternalVendorCode/onefuzz
1
0
Fork 0
mirror of https://github.com/microsoft/onefuzz.git synced 2025-06-16 20:08:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

allow pools & scalesets set to shutdown to halt (#1104)

Browse Source 
Currently, if a pool or scaleset is set to `shutdown`, it cannot be set to `halt`.

While moving from `halt` to `shutdown` would cause issues, moving from `shutdown` to `halt` is fine.
This commit is contained in:
bmc-msft
2021-07-23 09:14:47 -04:00
committed by GitHub
parent 5be9c4dcee
commit 55366e751a
2 changed files with 18 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
src/api-service/__app__/onefuzzlib/workers/pools.py
View File

@ -135,8 +135,7 @@ class Pool(BASE_POOL, ORMMixin):
def init(self) -> None: def init(self) -> None:
create_queue(self.get_pool_queue(), StorageType.corpus) create_queue(self.get_pool_queue(), StorageType.corpus)
ShrinkQueue(self.pool_id).create() ShrinkQueue(self.pool_id).create()
self.state = PoolState.running self.set_state(PoolState.running)
self.save()
def schedule_workset(self, work_set: WorkSet) -> bool: def schedule_workset(self, work_set: WorkSet) -> bool:
# Don't schedule work for pools that can't and won't do work. # Don't schedule work for pools that can't and won't do work.
@ -183,15 +182,10 @@ class Pool(BASE_POOL, ORMMixin):
return cls.search(query=query) return cls.search(query=query)
def set_shutdown(self, now: bool) -> None: def set_shutdown(self, now: bool) -> None:
if self.state in [PoolState.halt, PoolState.shutdown]:
return
if now: if now:
self.state = PoolState.halt self.set_state(PoolState.halt)
else: else:
self.state = PoolState.shutdown self.set_state(PoolState.shutdown)
self.save()
def shutdown(self) -> None: def shutdown(self) -> None:
"""shutdown allows nodes to finish current work then delete""" """shutdown allows nodes to finish current work then delete"""
@ -202,8 +196,6 @@ class Pool(BASE_POOL, ORMMixin):
nodes = Node.search(query={"pool_name": [self.name]}) nodes = Node.search(query={"pool_name": [self.name]})
if not scalesets and not nodes: if not scalesets and not nodes:
logging.info("pool stopped, deleting: %s", self.name) logging.info("pool stopped, deleting: %s", self.name)
self.state = PoolState.halt
self.delete() self.delete()
return return
@ -227,7 +219,6 @@ class Pool(BASE_POOL, ORMMixin):
delete_queue(self.get_pool_queue(), StorageType.corpus) delete_queue(self.get_pool_queue(), StorageType.corpus)
ShrinkQueue(self.pool_id).delete() ShrinkQueue(self.pool_id).delete()
logging.info("pool stopped, deleting: %s", self.name) logging.info("pool stopped, deleting: %s", self.name)
self.state = PoolState.halt
self.delete() self.delete()
return return
@ -239,6 +230,17 @@ class Pool(BASE_POOL, ORMMixin):
self.save() 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 @classmethod
def key_fields(cls) -> Tuple[str, str]: def key_fields(cls) -> Tuple[str, str]:
return ("name", "pool_id") return ("name", "pool_id")

15
src/api-service/__app__/onefuzzlib/workers/scalesets.py
View File

@ -635,22 +635,11 @@ class Scaleset(BASE_SCALESET, ORMMixin):
node.delete() node.delete()
def set_shutdown(self, now: bool) -> None: 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: if now:
self.set_state(ScalesetState.halt) self.set_state(ScalesetState.halt)
else: else:
self.set_state(ScalesetState.shutdown) self.set_state(ScalesetState.shutdown)
self.save()
def shutdown(self) -> None: def shutdown(self) -> None:
size = get_vmss_size(self.scaleset_id) size = get_vmss_size(self.scaleset_id)
if size is None: if size is None:
@ -807,6 +796,10 @@ class Scaleset(BASE_SCALESET, ORMMixin):
if self.state == state: if self.state == state:
return return
# scalesets should never leave the `halt` state
if self.state == ScalesetState.halt:
return
self.state = state self.state = state
self.save() self.save()
if self.state == ScalesetState.resize: if self.state == ScalesetState.resize: