[aws][fix] Fix parallel collection (#2071)

Co-authored-by: Matthias Veit <matthias_veit@yahoo.de>
Co-authored-by: Matthias Veit <aquamatthias@users.noreply.github.com>

Author: 1101-1

fixlib/fixlib/baseresources.py

@@ -148,16 +148,18 @@ def __str__(self) -> str:
 
     # load balancers
     RequestCount = "request"  # _count will be added to the end because of the unit
-    ActiveConnection = "active_connection"
+    ActiveConnectionCount = "active_connection"  # _count will be added to the end because of the unit
+    ALBActiveConnectionCount = "alb_active_connection"  # _count will be added to the end because of the unit
     ConnectionAttemptCount = "connection_attempt"  # _count will be added to the end because of the unit
     ConnectionEstablishedCount = "connection_established"  # _count will be added to the end because of the unit
     StatusCode2XX = "status_code_2xx"
     StatusCode4XX = "status_code_4xx"
     StatusCode5XX = "status_code_5xx"
     Latency = "latency"
+    TargetResponseTime = "target_response_time"
     ProcessedBytes = "processed"  # _bytes will be added to the end because of the unit
     HealthyHostCount = "healthy_host"  # _count will be added to the end because of the unit
-    UnhealthyHostCount = "unhealthy_host"  # _count will be added to the end because of the unit
+    UnHealthyHostCount = "unhealthy_host"  # _count will be added to the end because of the unit
     HealthyStateRouting = "healthy_state_routing"
     UnhealthyStateRouting = "unhealthy_state_routing"
     HealthyStateDNS = "healthy_state_dns"

plugins/aws/fix_plugin_aws/collector.py

@@ -1,8 +1,10 @@
+from collections import defaultdict
 import logging
-from attrs import define
 from concurrent.futures import Future, ThreadPoolExecutor
-from typing import List, Type, Optional, ClassVar, Union
-from datetime import datetime, timezone
+from datetime import datetime, timedelta, timezone
+from typing import List, Type, Optional, ClassVar, Union, cast
+
+from attrs import define
 
 from fix_plugin_aws.aws_client import AwsClient
 from fix_plugin_aws.configuration import AwsConfig
@@ -46,16 +48,14 @@
     waf,
 )
 from fix_plugin_aws.resource.base import AwsAccount, AwsApiSpec, AwsRegion, AwsResource, GraphBuilder
-
 from fixlib.baseresources import Cloud, EdgeType, BaseOrganizationalRoot, BaseOrganizationalUnit
 from fixlib.core.actions import CoreFeedback, ErrorAccumulator
 from fixlib.core.progress import ProgressDone, ProgressTree
 from fixlib.graph import Graph, BySearchCriteria, ByNodeId
+from fixlib.json import value_in_path
 from fixlib.proc import set_thread_name
 from fixlib.threading import ExecutorQueue, GatherFutures
 from fixlib.types import Json
-from fixlib.json import value_in_path
-
 from .utils import global_region_by_partition
 
 log = logging.getLogger("fix.plugins.aws")
@@ -193,6 +193,7 @@ def get_last_run() -> Optional[datetime]:
                 self.cloud,
                 self.account,
                 self.global_region,
+                {r.id: r for r in self.regions},
                 self.client,
                 shared_queue,
                 self.core_feedback,
@@ -218,6 +219,16 @@ def get_last_run() -> Optional[datetime]:
                 self.collect_resources(regional_resources, global_builder.for_region(region))
             shared_queue.wait_for_submitted_work()
 
+            if global_builder.config.collect_usage_metrics:
+                try:
+                    log.info(f"[Aws:{self.account.id}] Collect usage metrics.")
+                    self.collect_usage_metrics(global_builder)
+                except Exception as e:
+                    log.warning(
+                        f"Failed to collect usage metrics on account {self.account.id} in region {global_builder.region.id}: {e}"
+                    )
+            shared_queue.wait_for_submitted_work()
+
             # connect nodes
             log.info(f"[Aws:{self.account.id}] Connect resources and create edges.")
             for node, data in list(self.graph.nodes(data=True)):
@@ -281,6 +292,40 @@ def work_done(_: Future[None]) -> None:
         when_done.add_done_callback(work_done)
         return when_done
 
+    def collect_usage_metrics(self, builder: GraphBuilder) -> None:
+        metrics_queries = defaultdict(list)
+        two_hours = timedelta(hours=2)
+        thirty_minutes = timedelta(minutes=30)
+        lookup_map = {}
+        for resource in builder.graph.nodes:
+            if not isinstance(resource, AwsResource):
+                continue
+            # region can be overridden in the query: s3 is global, but need to be queried per region
+            if region := cast(AwsRegion, resource.region()):
+                lookup_map[resource.id] = resource
+                resource_queries: List[cloudwatch.AwsCloudwatchQuery] = resource.collect_usage_metrics(builder)
+                for query in resource_queries:
+                    query_region = query.region or region
+                    start = query.start_delta or builder.metrics_delta
+                    if query.period and query.period < thirty_minutes:
+                        start = min(start, two_hours)
+                    metrics_queries[(query_region, start)].append(query)
+        for (region, start), queries in metrics_queries.items():
+
+            def collect_and_set_metrics(
+                start: timedelta, region: AwsRegion, queries: List[cloudwatch.AwsCloudwatchQuery]
+            ) -> None:
+                start_at = builder.created_at - start
+                try:
+                    result = cloudwatch.AwsCloudwatchMetricData.query_for_multiple(
+                        builder.for_region(region), start_at, builder.created_at, queries
+                    )
+                    cloudwatch.update_resource_metrics(lookup_map, result)
+                except Exception as e:
+                    log.warning(f"Error occurred in region {region}: {e}")
+
+            builder.submit_work("cloudwatch", collect_and_set_metrics, start, region, queries)
+
     # TODO: move into separate AwsAccountSettings
     def update_account(self) -> None:
         log.info(f"Collecting AWS IAM Account Summary in account {self.account.dname}")

plugins/aws/fix_plugin_aws/resource/apigateway.py

@@ -513,64 +513,56 @@ def called_mutator_apis(cls) -> List[AwsApiSpec]:
         ]
 
     @classmethod
-    def collect(cls: Type[AwsResource], json: List[Json], builder: GraphBuilder) -> List[AwsResource]:
-        instances: List[AwsResource] = []
+    def collect(cls: Type[AwsResource], json: List[Json], builder: GraphBuilder) -> None:
+        def add_instance(api_instance: AwsResource) -> None:
+            for deployment in builder.client.list(service_name, "get-deployments", "items", restApiId=api_instance.id):
+                if deploy_instance := AwsApiGatewayDeployment.from_api(deployment, builder):
+                    deploy_instance.set_arn(
+                        builder=builder,
+                        account="",
+                        resource=f"/restapis/{api_instance.id}/deployments/{deploy_instance.id}",
+                    )
+                    deploy_instance.api_link = api_instance.id
+                    builder.add_node(deploy_instance, deployment)
+                    builder.add_edge(api_instance, EdgeType.default, node=deploy_instance)
+                    for stage in builder.client.list(
+                        service_name,
+                        "get-stages",
+                        "item",
+                        restApiId=api_instance.id,
+                        deploymentId=deploy_instance.id,
+                    ):
+                        stage["syntheticId"] = f'{api_instance.id}_{stage["stageName"]}'  # create unique id
+                        if stage_instance := AwsApiGatewayStage.from_api(stage, builder):
+                            stage_instance.api_link = api_instance.id
+                            builder.add_node(stage_instance, stage)
+                            # reference kinds for this edge are maintained in AwsApiGatewayDeployment.reference_kinds # noqa: E501
+                            builder.add_edge(deploy_instance, EdgeType.default, node=stage_instance)
+            for authorizer in builder.client.list(service_name, "get-authorizers", "items", restApiId=api_instance.id):
+                if auth_instance := AwsApiGatewayAuthorizer.from_api(authorizer, builder):
+                    auth_instance.api_link = api_instance.id
+                    builder.add_node(auth_instance, authorizer)
+                    builder.add_edge(api_instance, EdgeType.default, node=auth_instance)
+            for resource in builder.client.list(service_name, "get-resources", "items", restApiId=api_instance.id):
+                if resource_instance := AwsApiGatewayResource.from_api(resource, builder):
+                    resource_instance.api_link = api_instance.id
+                    if resource_instance.resource_methods:
+                        for method in resource_instance.resource_methods:
+                            mapped = bend(AwsApiGatewayMethod.mapping, resource["resourceMethods"][method])
+                            if gm := parse_json(mapped, AwsApiGatewayMethod, builder):
+                                resource_instance.resource_methods[method] = gm
+                    builder.add_node(resource_instance, resource)
+                    builder.add_edge(api_instance, EdgeType.default, node=resource_instance)
+
         for js in json:
             if api_instance := cls.from_api(js, builder):
                 api_instance.set_arn(
                     builder=builder,
                     account="",
                     resource=f"/restapis/{api_instance.id}",
                 )
-                instances.append(api_instance)
                 builder.add_node(api_instance, js)
-                for deployment in builder.client.list(
-                    service_name, "get-deployments", "items", restApiId=api_instance.id
-                ):
-                    if deploy_instance := AwsApiGatewayDeployment.from_api(deployment, builder):
-                        deploy_instance.set_arn(
-                            builder=builder,
-                            account="",
-                            resource=f"/restapis/{api_instance.id}/deployments/{deploy_instance.id}",
-                        )
-                        deploy_instance.api_link = api_instance.id
-                        instances.append(deploy_instance)
-                        builder.add_node(deploy_instance, deployment)
-                        builder.add_edge(api_instance, EdgeType.default, node=deploy_instance)
-                        for stage in builder.client.list(
-                            service_name,
-                            "get-stages",
-                            "item",
-                            restApiId=api_instance.id,
-                            deploymentId=deploy_instance.id,
-                        ):
-                            stage["syntheticId"] = f'{api_instance.id}_{stage["stageName"]}'  # create unique id
-                            if stage_instance := AwsApiGatewayStage.from_api(stage, builder):
-                                stage_instance.api_link = api_instance.id
-                                instances.append(stage_instance)
-                                builder.add_node(stage_instance, stage)
-                                # reference kinds for this edge are maintained in AwsApiGatewayDeployment.reference_kinds # noqa: E501
-                                builder.add_edge(deploy_instance, EdgeType.default, node=stage_instance)
-                for authorizer in builder.client.list(
-                    service_name, "get-authorizers", "items", restApiId=api_instance.id
-                ):
-                    if auth_instance := AwsApiGatewayAuthorizer.from_api(authorizer, builder):
-                        auth_instance.api_link = api_instance.id
-                        instances.append(auth_instance)
-                        builder.add_node(auth_instance, authorizer)
-                        builder.add_edge(api_instance, EdgeType.default, node=auth_instance)
-                for resource in builder.client.list(service_name, "get-resources", "items", restApiId=api_instance.id):
-                    if resource_instance := AwsApiGatewayResource.from_api(resource, builder):
-                        resource_instance.api_link = api_instance.id
-                        instances.append(resource_instance)
-                        if resource_instance.resource_methods:
-                            for method in resource_instance.resource_methods:
-                                mapped = bend(AwsApiGatewayMethod.mapping, resource["resourceMethods"][method])
-                                if gm := parse_json(mapped, AwsApiGatewayMethod, builder):
-                                    resource_instance.resource_methods[method] = gm
-                        builder.add_node(resource_instance, resource)
-                        builder.add_edge(api_instance, EdgeType.default, node=resource_instance)
-        return instances
+                builder.submit_work(service_name, add_instance, api_instance)
 
     def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
         if self.api_endpoint_configuration:

plugins/aws/fix_plugin_aws/resource/athena.py

@@ -1,6 +1,7 @@
 from typing import ClassVar, Dict, Optional, List, Any
 from typing import Type
 
+
 from attrs import define
 
 from fix_plugin_aws.aws_client import AwsClient
@@ -133,27 +134,18 @@ def called_mutator_apis(cls) -> List[AwsApiSpec]:
         ]
 
     @classmethod
-    def collect(cls: Type[AwsResource], json: List[Json], builder: GraphBuilder) -> List[AwsResource]:
-        workgroups: List[AwsResource] = []
-
-        def fetch_workgroup(name: str) -> Optional[AwsAthenaWorkGroup]:
+    def collect(cls: Type[AwsResource], json: List[Json], builder: GraphBuilder) -> None:
+        def fetch_workgroup(name: str) -> None:
             result = builder.client.get(
                 aws_service=service_name, action="get-work-group", result_name="WorkGroup", WorkGroup=name
             )
-            if result is None:
-                return None
-
-            if workgroup := AwsAthenaWorkGroup.from_api(result, builder):
+            if result and (workgroup := AwsAthenaWorkGroup.from_api(result, builder)):
                 workgroup.set_arn(
                     builder=builder,
                     resource=f"workgroup/{workgroup.name}",
                 )
-                workgroups.append(workgroup)
                 builder.add_node(workgroup, result)
                 builder.submit_work(service_name, add_tags, workgroup)
-                return workgroup
-            else:
-                return None
 
         def add_tags(data_catalog: AwsAthenaWorkGroup) -> None:
             tags = builder.client.list(
@@ -167,8 +159,7 @@ def add_tags(data_catalog: AwsAthenaWorkGroup) -> None:
 
         for js in json:
             if (name := js.get("Name")) is not None and isinstance(name, str):
-                fetch_workgroup(name)
-        return workgroups
+                builder.submit_work(service_name, fetch_workgroup, name)
 
     # noinspection PyUnboundLocalVariable
     def connect_in_graph(self, builder: GraphBuilder, source: Json) -> None:
@@ -252,25 +243,18 @@ def called_mutator_apis(cls) -> List[AwsApiSpec]:
         ]
 
     @classmethod
-    def collect(cls: Type[AwsResource], json: List[Json], builder: GraphBuilder) -> List[AwsResource]:
-        catalogs: List[AwsResource] = []
-
-        def fetch_data_catalog(data_catalog_name: str) -> Optional[AwsAthenaDataCatalog]:
+    def collect(cls: Type[AwsResource], json: List[Json], builder: GraphBuilder) -> None:
+        def fetch_data_catalog(data_catalog_name: str) -> None:
             result = builder.client.get(
                 aws_service=service_name,
                 action="get-data-catalog",
                 result_name="DataCatalog",
                 Name=data_catalog_name,
             )
-            if result is None:
-                return None
-            if catalog := AwsAthenaDataCatalog.from_api(result, builder):
+            if result and (catalog := AwsAthenaDataCatalog.from_api(result, builder)):
                 catalog.set_arn(builder=builder, resource=f"datacatalog/{catalog.name}")
-                catalogs.append(catalog)
                 builder.add_node(catalog, result)
                 builder.submit_work(service_name, add_tags, catalog)
-                return catalog
-            return None
 
         def add_tags(data_catalog: AwsAthenaDataCatalog) -> None:
             tags = builder.client.list(
@@ -285,8 +269,7 @@ def add_tags(data_catalog: AwsAthenaDataCatalog) -> None:
         for js in json:
             # we filter out the default data catalog as it is not possible to do much with it
             if (name := js.get("CatalogName")) is not None and isinstance(name, str) and name != "AwsDataCatalog":
-                fetch_data_catalog(name)
-        return catalogs
+                builder.submit_work(service_name, fetch_data_catalog, name)
 
     def update_resource_tag(self, client: AwsClient, key: str, value: str) -> bool:
         client.call(