set pool_id on node creation (#1049)

src/api-service/__app__/agent_registration/__init__.py

@@ -102,7 +102,8 @@ def post(req: func.HttpRequest) -> func.HttpResponse:
         node.delete()
 
     node = Node.create(
-        pool_name=registration_request.pool_name,
+        pool_id=pool.pool_id,
+        pool_name=pool.name,
         machine_id=registration_request.machine_id,
         scaleset_id=registration_request.scaleset_id,
         version=registration_request.version,

src/api-service/__app__/onefuzzlib/workers/nodes.py

@@ -48,6 +48,7 @@ class Node(BASE_NODE, ORMMixin):
     def create(
         cls,
         *,
+        pool_id: UUID,
         pool_name: PoolName,
         machine_id: UUID,
         scaleset_id: Optional[UUID],
@@ -55,6 +56,7 @@ class Node(BASE_NODE, ORMMixin):
         new: bool = False,
     ) -> "Node":
         node = cls(
+            pool_id=pool_id,
             pool_name=pool_name,
             machine_id=machine_id,
             scaleset_id=scaleset_id,
@@ -78,11 +80,14 @@ class Node(BASE_NODE, ORMMixin):
     def search_states(
         cls,
         *,
+        pool_id: Optional[UUID] = None,
         scaleset_id: Optional[UUID] = None,
         states: Optional[List[NodeState]] = None,
         pool_name: Optional[PoolName] = None,
     ) -> List["Node"]:
         query: QueryFilter = {}
+        if pool_id:
+            query["pool_id"] = [pool_id]
         if scaleset_id:
             query["scaleset_id"] = [scaleset_id]
         if states:
@@ -95,6 +100,7 @@ class Node(BASE_NODE, ORMMixin):
     def search_outdated(
         cls,
         *,
+        pool_id: Optional[UUID] = None,
         scaleset_id: Optional[UUID] = None,
         states: Optional[List[NodeState]] = None,
         pool_name: Optional[PoolName] = None,
@@ -102,6 +108,8 @@ class Node(BASE_NODE, ORMMixin):
         num_results: Optional[int] = None,
     ) -> List["Node"]:
         query: QueryFilter = {}
+        if pool_id:
+            query["pool_id"] = [pool_id]
         if scaleset_id:
             query["scaleset_id"] = [scaleset_id]
         if states:

src/api-service/__app__/onefuzzlib/workers/scalesets.py

@@ -291,6 +291,8 @@ class Scaleset(BASE_SCALESET, ORMMixin):
 
     # result = 'did I modify the scaleset in azure'
     def cleanup_nodes(self) -> bool:
+        from .pools import Pool
+
         logging.info(
             SCALESET_LOG_PREFIX + "cleaning up nodes. scaleset_id:%s", self.scaleset_id
         )
@@ -302,6 +304,16 @@ class Scaleset(BASE_SCALESET, ORMMixin):
             self.halt()
             return True
 
+        pool = Pool.get_by_name(self.pool_name)
+        if isinstance(pool, Error):
+            logging.error(
+                "unable to find pool during cleanup: %s - %s",
+                self.scaleset_id,
+                pool,
+            )
+            self.set_failed(pool)
+            return True
+
         Node.reimage_long_lived_nodes(self.scaleset_id)
 
         to_reimage = []
@@ -346,6 +358,7 @@ class Scaleset(BASE_SCALESET, ORMMixin):
             # Note, using `new=True` makes it such that if a node already has
             # checked in, this won't overwrite it.
             Node.create(
+                pool_id=pool.pool_id,
                 pool_name=self.pool_name,
                 machine_id=machine_id,
                 scaleset_id=self.scaleset_id,