Set AWSAccount in context instead of just tenant name

grafana · Aug 22, 2023 · fce9503 · fce9503
1 parent c6b81ed
commit fce9503
Show file tree

Hide file tree

Showing 35 changed files with 116 additions and 72 deletions.
diff --git a/src/main/java/ai/asserts/aws/MetricNameUtil.java b/src/main/java/ai/asserts/aws/MetricNameUtil.java
@@ -19,6 +19,8 @@ public class MetricNameUtil {
     public static final String SCRAPE_LATENCY_METRIC = "aws_exporter_milliseconds";
     public static final String ASSERTS_ERROR_TYPE = "asserts_error_type";
     public static final String TENANT = "tenant";
+    public static final String ENV = "asserts_env";
+    public static final String SITE = "asserts_site";
     public static final String SCRAPE_ERROR_COUNT_METRIC = "aws_exporter_error_total";
     public static final String SCRAPE_OPERATION_LABEL = "operation";
     public static final String SCRAPE_REGION_LABEL = "region";

diff --git a/src/main/java/ai/asserts/aws/TaskExecutorUtil.java b/src/main/java/ai/asserts/aws/TaskExecutorUtil.java
@@ -4,6 +4,7 @@
  */
 package ai.asserts.aws;
 
+import ai.asserts.aws.account.AWSAccount;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.beans.factory.annotation.Qualifier;
 import org.springframework.stereotype.Component;
@@ -21,15 +22,15 @@ public class TaskExecutorUtil {
     private final TaskThreadPool taskThreadPool;
 
     private final AWSApiCallRateLimiter rateLimiter;
-    private static final ThreadLocal<String> tenantName = new ThreadLocal<>();
+    private static final ThreadLocal<AWSAccount> tenantName = new ThreadLocal<>();
 
 
     public TaskExecutorUtil(@Qualifier("aws-api-calls-thread-pool") TaskThreadPool taskThreadPool, AWSApiCallRateLimiter rateLimiter) {
         this.taskThreadPool = taskThreadPool;
         this.rateLimiter = rateLimiter;
     }
 
-    public <T> Future<T> executeTenantTask(String tenant, TenantTask<T> task) {
+    public <T> Future<T> executeTenantTask(AWSAccount tenant, TenantTask<T> task) {
         return taskThreadPool.getExecutorService().submit(() -> {
             tenantName.set(tenant);
             try {
@@ -53,7 +54,7 @@ public <K> void awaitAll(List<Future<K>> futures, Consumer<K> consumer) {
         });
     }
 
-    public String getTenant() {
+    public AWSAccount getAccountDetails() {
         return tenantName.get();
     }
 }
diff --git a/src/main/java/ai/asserts/aws/cloudwatch/alarms/AlarmFetcher.java b/src/main/java/ai/asserts/aws/cloudwatch/alarms/AlarmFetcher.java
@@ -103,7 +103,7 @@ public void update() {
                 continue;
             }
             accountRegion.getRegions().forEach(region ->
-                    futures.add(taskExecutorUtil.executeTenantTask(accountRegion.getTenant(),
+                    futures.add(taskExecutorUtil.executeTenantTask(accountRegion,
                             new CollectionBuilderTask<Sample>() {
                                 @Override
                                 public List<Sample> call() {

diff --git a/src/main/java/ai/asserts/aws/cloudwatch/query/MetricQueryProvider.java b/src/main/java/ai/asserts/aws/cloudwatch/query/MetricQueryProvider.java
@@ -102,7 +102,7 @@ Map<String, Map<String, Map<Integer, List<MetricQuery>>>> getQueriesInternal() {
             log.info("Will discover metrics and build metric queries for tenant {}, account {}",
                     accountRegion.getTenant(), accountRegion.getAccountId());
             accountRegion.getRegions().forEach(region -> futures.add(
-                    taskExecutorUtil.executeTenantTask(accountRegion.getTenant(), new SimpleTenantTask<Void>() {
+                    taskExecutorUtil.executeTenantTask(accountRegion, new SimpleTenantTask<Void>() {
                         @Override
                         public Void call() {
                             buildQueries(scrapeConfig, region, accountRegion, account, queriesByAccount);

diff --git a/src/main/java/ai/asserts/aws/exporter/ApiGatewayToLambdaBuilder.java b/src/main/java/ai/asserts/aws/exporter/ApiGatewayToLambdaBuilder.java
@@ -96,7 +96,7 @@ public void update() {
         List<Future<List<Sample>>> futures = new ArrayList<>();
         for (AWSAccount accountRegion : accountProvider.getAccounts()) {
             accountRegion.getRegions().forEach(region ->
-                    futures.add(taskExecutorUtil.executeTenantTask(accountRegion.getTenant(),
+                    futures.add(taskExecutorUtil.executeTenantTask(accountRegion,
                             new CollectionBuilderTask<Sample>() {
                                 @Override
                                 public List<Sample> call() {
@@ -182,18 +182,19 @@ private void captureIntegrations(ApiGatewayClient client, Set<ResourceRelation>
                     }
                     Matcher matcher = LAMBDA_URI_PATTERN.matcher(uri);
                     if (matcher.matches()) {
+                        String tenant = taskExecutorUtil.getAccountDetails().getTenant();
                         ResourceRelation resourceRelation = ResourceRelation.builder()
-                                .tenant(taskExecutorUtil.getTenant())
+                                .tenant(tenant)
                                 .from(Resource.builder()
-                                        .tenant(taskExecutorUtil.getTenant())
+                                        .tenant(tenant)
                                         .type(ApiGateway)
                                         .name(restApi.name())
                                         .id(restApi.id())
                                         .region(region)
                                         .account(accountId)
                                         .build())
                                 .to(Resource.builder()
-                                        .tenant(taskExecutorUtil.getTenant())
+                                        .tenant(tenant)
                                         .type(LambdaFunction)
                                         .name(matcher.group(4))
                                         .region(matcher.group(2))

diff --git a/src/main/java/ai/asserts/aws/exporter/DynamoDBExporter.java b/src/main/java/ai/asserts/aws/exporter/DynamoDBExporter.java
@@ -83,7 +83,7 @@ public void update() {
         List<Future<List<Sample>>> futures = new ArrayList<>();
         accountProvider.getAccounts().forEach(account -> account.getRegions().forEach(region ->
                 futures.add(
-                        taskExecutorUtil.executeTenantTask(account.getTenant(), new CollectionBuilderTask<Sample>() {
+                        taskExecutorUtil.executeTenantTask(account, new CollectionBuilderTask<Sample>() {
                             @Override
                             public List<Sample> call() {
                                 return buildSamples(region, account, allSamples);

diff --git a/src/main/java/ai/asserts/aws/exporter/EC2ToEBSVolumeExporter.java b/src/main/java/ai/asserts/aws/exporter/EC2ToEBSVolumeExporter.java
@@ -103,14 +103,14 @@ public void update() {
         List<Future<List<Sample>>> futures = new ArrayList<>();
         List<Future<List<ResourceRelation>>> volumeFutures = new ArrayList<>();
         accountProvider.getAccounts().forEach(awsAccount -> awsAccount.getRegions().forEach(region -> {
-            futures.add(taskExecutorUtil.executeTenantTask(awsAccount.getTenant(), new CollectionBuilderTask<Sample>() {
+            futures.add(taskExecutorUtil.executeTenantTask(awsAccount, new CollectionBuilderTask<Sample>() {
                 @Override
                 public List<Sample> call() {
                     return buildEC2InstanceMetrics(region, awsAccount);
                 }
             }));
             volumeFutures.add(
-                    taskExecutorUtil.executeTenantTask(awsAccount.getTenant(),
+                    taskExecutorUtil.executeTenantTask(awsAccount,
                             new CollectionBuilderTask<ResourceRelation>() {
                                 @Override
                                 public List<ResourceRelation> call() {
@@ -156,17 +156,18 @@ private List<ResourceRelation> buildResourceRelations(AWSAccount awsAccount, Str
                     newAttachedVolumes.addAll(resp.volumes().stream()
                             .flatMap(volume -> volume.attachments().stream())
                             .map(volumeAttachment -> {
+                                String tenant = taskExecutorUtil.getAccountDetails().getTenant();
                                 Resource ec2Instance = Resource.builder()
-                                        .tenant(taskExecutorUtil.getTenant())
+                                        .tenant(tenant)
                                         .account(accountId)
                                         .region(region)
                                         .type(EC2Instance)
                                         .name(volumeAttachment.instanceId())
                                         .build();
                                 return ResourceRelation.builder()
-                                        .tenant(taskExecutorUtil.getTenant())
+                                        .tenant(tenant)
                                         .from(Resource.builder()
-                                                .tenant(taskExecutorUtil.getTenant())
+                                                .tenant(tenant)
                                                 .account(accountId)
                                                 .region(region)
                                                 .type(EBSVolume)

diff --git a/src/main/java/ai/asserts/aws/exporter/ECSTaskProvider.java b/src/main/java/ai/asserts/aws/exporter/ECSTaskProvider.java
@@ -147,7 +147,7 @@ public void run() {
         for (AWSAccount account : accountProvider.getAccounts()) {
             ScrapeConfig scrapeConfig = scrapeConfigProvider.getScrapeConfig(account.getTenant());
             for (String region : account.getRegions()) {
-                taskExecutorUtil.executeTenantTask(account.getTenant(), new SimpleTenantTask<Void>() {
+                taskExecutorUtil.executeTenantTask(account, new SimpleTenantTask<Void>() {
                     @Override
                     public Void call() {
                         Map<Resource, List<Resource>> clusterWiseNewTasks = new HashMap<>();

diff --git a/src/main/java/ai/asserts/aws/exporter/ECSTaskUtil.java b/src/main/java/ai/asserts/aws/exporter/ECSTaskUtil.java
@@ -236,9 +236,10 @@ private LabelsBuilder getLabelsBuilder(AWSAccount account, Resource cluster, Opt
 
         LabelsBuilder labelsBuilder;
         taskSubnetMap.computeIfAbsent(taskResource.getName(), k -> getSubnetDetails(task, taskResource));
+        String tenant = taskExecutorUtil.getAccountDetails().getTenant();
         if (service.isPresent()) {
             labelsBuilder = Labels.builder()
-                    .tenant(taskExecutorUtil.getTenant())
+                    .tenant(tenant)
                     .workload(service.get())
                     .taskId(taskResource.getName())
                     .pod(service.get() + "-" + taskResource.getName())
@@ -254,7 +255,7 @@ private LabelsBuilder getLabelsBuilder(AWSAccount account, Resource cluster, Opt
                     .metricsPath("/metrics");
         } else {
             labelsBuilder = Labels.builder()
-                    .tenant(taskExecutorUtil.getTenant())
+                    .tenant(tenant)
                     .workload(taskDefResource.getName())
                     .taskId(taskResource.getName())
                     .pod(taskDefResource.getName() + "-" + taskResource.getName())

diff --git a/src/main/java/ai/asserts/aws/exporter/EMRExporter.java b/src/main/java/ai/asserts/aws/exporter/EMRExporter.java
@@ -71,7 +71,7 @@ public void update() {
         List<Future<List<Sample>>> futures = new ArrayList<>();
         accountProvider.getAccounts().forEach(account -> account.getRegions().forEach(region ->
                 futures.add(
-                        taskExecutorUtil.executeTenantTask(account.getTenant(), new CollectionBuilderTask<Sample>() {
+                        taskExecutorUtil.executeTenantTask(account, new CollectionBuilderTask<Sample>() {
                             @Override
                             public List<Sample> call() {
                                 return buildSamples(region, account);

diff --git a/src/main/java/ai/asserts/aws/exporter/KinesisAnalyticsExporter.java b/src/main/java/ai/asserts/aws/exporter/KinesisAnalyticsExporter.java
@@ -77,7 +77,7 @@ public void update() {
         List<Future<List<Sample>>> futures = new ArrayList<>();
         accountProvider.getAccounts().forEach(account -> account.getRegions().forEach(region ->
                 futures.add(
-                        taskExecutorUtil.executeTenantTask(account.getTenant(), new CollectionBuilderTask<Sample>() {
+                        taskExecutorUtil.executeTenantTask(account, new CollectionBuilderTask<Sample>() {
                             @Override
                             public List<Sample> call() {
                                 return buildSamples(region, account);

diff --git a/src/main/java/ai/asserts/aws/exporter/KinesisFirehoseExporter.java b/src/main/java/ai/asserts/aws/exporter/KinesisFirehoseExporter.java
@@ -90,7 +90,7 @@ public void update() {
         List<Future<List<Sample>>> futures = new ArrayList<>();
         accountProvider.getAccounts().forEach(account -> account.getRegions().forEach(region ->
                 futures.add(
-                        taskExecutorUtil.executeTenantTask(account.getTenant(), new CollectionBuilderTask<Sample>() {
+                        taskExecutorUtil.executeTenantTask(account, new CollectionBuilderTask<Sample>() {
                             @Override
                             public List<Sample> call() {
                                 return buildMetricSamples(region, account);

diff --git a/src/main/java/ai/asserts/aws/exporter/KinesisStreamExporter.java b/src/main/java/ai/asserts/aws/exporter/KinesisStreamExporter.java
@@ -82,7 +82,7 @@ public void update() {
         List<Sample> allSamples = new ArrayList<>();
         List<Future<List<Sample>>> futures = new ArrayList<>();
         accountProvider.getAccounts().forEach(account -> account.getRegions().forEach(region ->
-                futures.add(taskExecutorUtil.executeTenantTask(account.getTenant(), new CollectionBuilderTask<Sample>() {
+                futures.add(taskExecutorUtil.executeTenantTask(account, new CollectionBuilderTask<Sample>() {
                     @Override
                     public List<Sample> call()  {
                         return buildSamples(region, account);

diff --git a/src/main/java/ai/asserts/aws/exporter/LBToASGRelationBuilder.java b/src/main/java/ai/asserts/aws/exporter/LBToASGRelationBuilder.java
@@ -97,7 +97,7 @@ public void updateRouting() {
         List<Future<List<Sample>>> futures = new ArrayList<>();
         for (AWSAccount accountRegion : accountProvider.getAccounts()) {
             accountRegion.getRegions().forEach(region ->
-                    futures.add(taskExecutorUtil.executeTenantTask(accountRegion.getTenant(),
+                    futures.add(taskExecutorUtil.executeTenantTask(accountRegion,
                             new CollectionBuilderTask<Sample>() {
                                 @Override
                                 public List<Sample> call() {

diff --git a/src/main/java/ai/asserts/aws/exporter/LBToECSRoutingBuilder.java b/src/main/java/ai/asserts/aws/exporter/LBToECSRoutingBuilder.java
@@ -69,7 +69,7 @@ public void run() {
         Set<ResourceRelation> newRouting = new HashSet<>();
         List<Future<Set<ResourceRelation>>> futures = new ArrayList<>();
         accountProvider.getAccounts().forEach(awsAccount -> awsAccount.getRegions().forEach(region ->
-                futures.add(taskExecutorUtil.executeTenantTask(awsAccount.getTenant(),
+                futures.add(taskExecutorUtil.executeTenantTask(awsAccount,
                         new SimpleTenantTask<Set<ResourceRelation>>() {
                             @Override
                             public Set<ResourceRelation> call() {

diff --git a/src/main/java/ai/asserts/aws/exporter/LBToLambdaRoutingBuilder.java b/src/main/java/ai/asserts/aws/exporter/LBToLambdaRoutingBuilder.java
@@ -55,7 +55,7 @@ public Set<ResourceRelation> getRoutings() {
         List<Future<Pair<Set<ResourceRelation>, Set<Resource>>>> futures = new ArrayList<>();
         for (AWSAccount accountRegion : accountProvider.getAccounts()) {
             accountRegion.getRegions().forEach(region ->
-                    futures.add(taskExecutorUtil.executeTenantTask(accountRegion.getTenant(),
+                    futures.add(taskExecutorUtil.executeTenantTask(accountRegion,
                             new SimpleTenantTask<Pair<Set<ResourceRelation>, Set<Resource>>>() {
                                 @Override
                                 public Pair<Set<ResourceRelation>, Set<Resource>> call() {

diff --git a/src/main/java/ai/asserts/aws/exporter/LambdaCapacityExporter.java b/src/main/java/ai/asserts/aws/exporter/LambdaCapacityExporter.java
@@ -115,7 +115,7 @@ private List<MetricFamilySamples> getMetrics() {
                 Map<String, Map<String, LambdaFunction>> byRegion = byAccountByRegion.getOrDefault(account,
                         Collections.emptyMap());
                 byRegion.forEach((region, functions) -> futures.add(taskExecutorUtil.executeTenantTask(
-                        accountRegion.getTenant(), new SimpleTenantTask<Map<String, List<Sample>>>() {
+                        accountRegion, new SimpleTenantTask<Map<String, List<Sample>>>() {
                             @Override
                             public Map<String, List<Sample>> call() {
                                 return buildSamples(region, accountRegion, account, accountLimitMetric,

diff --git a/src/main/java/ai/asserts/aws/exporter/LambdaEventSourceExporter.java b/src/main/java/ai/asserts/aws/exporter/LambdaEventSourceExporter.java
@@ -95,7 +95,7 @@ private List<MetricFamilySamples> getMappings() {
             ScrapeConfig scrapeConfig = scrapeConfigProvider.getScrapeConfig(accountRegion.getTenant());
             Optional<NamespaceConfig> lambdaConfig = scrapeConfig.getLambdaConfig();
             lambdaConfig.ifPresent(namespaceConfig -> accountRegion.getRegions().forEach(region ->
-                    futures.add(taskExecutorUtil.executeTenantTask(accountRegion.getTenant(),
+                    futures.add(taskExecutorUtil.executeTenantTask(accountRegion,
                             new SimpleTenantTask<Map<String, List<Sample>>>() {
                                 @Override
                                 public Map<String, List<Sample>> call() {

diff --git a/src/main/java/ai/asserts/aws/exporter/LambdaInvokeConfigExporter.java b/src/main/java/ai/asserts/aws/exporter/LambdaInvokeConfigExporter.java
@@ -108,7 +108,7 @@ List<MetricFamilySamples> getInvokeConfigs() {
                 Map<String, Map<String, LambdaFunction>> byRegion =
                         new ConcurrentHashMap<>(byAccountByRegion.get(account));
                 byRegion.forEach((region, byARN) ->
-                        futures.add(taskExecutorUtil.executeTenantTask(accountRegion.getTenant(),
+                        futures.add(taskExecutorUtil.executeTenantTask(accountRegion,
                                 new CollectionBuilderTask<Sample>() {
                                     @Override
                                     public List<Sample> call() {

diff --git a/src/main/java/ai/asserts/aws/exporter/LoadBalancerExporter.java b/src/main/java/ai/asserts/aws/exporter/LoadBalancerExporter.java
@@ -84,7 +84,7 @@ public void update() {
         List<Future<Pair<List<Sample>, List<Sample>>>> futures = new ArrayList<>();
         accountProvider.getAccounts().forEach(awsAccount -> awsAccount.getRegions().forEach(region -> {
             ScrapeConfig scrapeConfig = scrapeConfigProvider.getScrapeConfig(awsAccount.getTenant());
-            futures.add(taskExecutorUtil.executeTenantTask(awsAccount.getTenant(),
+            futures.add(taskExecutorUtil.executeTenantTask(awsAccount,
                     new SimpleTenantTask<Pair<List<Sample>, List<Sample>>>() {
                         @Override
                         public Pair<List<Sample>, List<Sample>> call() {

diff --git a/src/main/java/ai/asserts/aws/exporter/MetricSampleBuilder.java b/src/main/java/ai/asserts/aws/exporter/MetricSampleBuilder.java
@@ -6,6 +6,7 @@
 
 import ai.asserts.aws.MetricNameUtil;
 import ai.asserts.aws.TaskExecutorUtil;
+import ai.asserts.aws.account.AWSAccount;
 import ai.asserts.aws.cloudwatch.query.MetricQuery;
 import io.prometheus.client.Collector.MetricFamilySamples;
 import io.prometheus.client.Collector.MetricFamilySamples.Sample;
@@ -20,8 +21,12 @@
 import java.util.Optional;
 import java.util.TreeMap;
 
+import static ai.asserts.aws.MetricNameUtil.ENV;
+import static ai.asserts.aws.MetricNameUtil.SCRAPE_REGION_LABEL;
+import static ai.asserts.aws.MetricNameUtil.SITE;
 import static ai.asserts.aws.MetricNameUtil.TENANT;
 import static io.prometheus.client.Collector.Type.GAUGE;
+import static org.springframework.util.StringUtils.hasLength;
 
 @Component
 @Slf4j
@@ -38,7 +43,13 @@ public List<Sample> buildSamples(String account, String region, MetricQuery metr
         String metricName = metricNameUtil.exportedMetricName(metricQuery.getMetric(), metricQuery.getMetricStat());
         if (metricDataResult.timestamps().size() > 0) {
             Map<String, String> labels = labelBuilder.buildLabels(account, region, metricQuery);
-            labels.putIfAbsent(TENANT, taskExecutorUtil.getTenant());
+            labels.putIfAbsent(TENANT, taskExecutorUtil.getAccountDetails().getTenant());
+            if (hasLength(taskExecutorUtil.getAccountDetails().getName())) {
+                labels.putIfAbsent(ENV, taskExecutorUtil.getAccountDetails().getName());
+            } else {
+                labels.putIfAbsent(ENV, account);
+            }
+            labels.putIfAbsent(SITE, region);
             labels.entrySet().removeIf(entry -> entry.getValue() == null);
             for (int i = 0; i < metricDataResult.timestamps().size(); i++) {
                 Sample sample = new Sample(
@@ -55,7 +66,16 @@ public List<Sample> buildSamples(String account, String region, MetricQuery metr
     public Optional<Sample> buildSingleSample(String metricName, Map<String, String> labels,
                                               Double metric) {
         labels = new TreeMap<>(labels);
-        labels.putIfAbsent(TENANT, taskExecutorUtil.getTenant());
+        AWSAccount accountDetails = taskExecutorUtil.getAccountDetails();
+        labels.putIfAbsent(TENANT, accountDetails.getTenant());
+        if (hasLength(accountDetails.getName())) {
+            labels.putIfAbsent(ENV, accountDetails.getName());
+        } else {
+            labels.putIfAbsent(ENV, accountDetails.getAccountId());
+        }
+        if (labels.containsKey(SCRAPE_REGION_LABEL)) {
+            labels.putIfAbsent(SITE, labels.get(SCRAPE_REGION_LABEL));
+        }
         labels.entrySet().removeIf(entry -> entry.getValue() == null);
         return Optional.of(new Sample(
                 metricName,