Skip to content

Commit

Permalink
Merge pull request #312 from asserts/radha/sc-16243-disable-ec2-apis
Browse files Browse the repository at this point in the history 
Flag to disable EC2 API usage
  • Loading branch information
@jradhakrishnan
jradhakrishnan authored Aug 2, 2023
2 parents 9c36eb1 + f85b90f commit 58e7532
Show file tree
Hide file tree
Showing 5 changed files with 186 additions and 45 deletions.
3 changes: 3 additions & 0 deletions config/src/main/java/ai/asserts/aws/config/ScrapeConfig.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,9 @@ public class ScrapeConfig {
@Builder.Default
private boolean discoverOnlySubnetTasks = false;

@Builder.Default
private boolean fetchEC2Metadata = false;

@Builder.Default
private Map<String, SubnetDetails> primaryExporterByAccount = new TreeMap<>();

Expand Down
34 changes: 23 additions & 11 deletions src/main/java/ai/asserts/aws/exporter/EC2ToEBSVolumeExporter.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,9 +4,9 @@
*/
package ai.asserts.aws.exporter;

import ai.asserts.aws.AWSApiCallRateLimiter;
import ai.asserts.aws.AWSClientProvider;
import ai.asserts.aws.CollectionBuilderTask;
import ai.asserts.aws.AWSApiCallRateLimiter;
import ai.asserts.aws.ScrapeConfigProvider;
import ai.asserts.aws.TagUtil;
import ai.asserts.aws.TaskExecutorUtil;
Expand All @@ -32,6 +32,7 @@
import software.amazon.awssdk.services.resourcegroupstaggingapi.model.Tag;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
Expand Down Expand Up @@ -69,7 +70,8 @@ public class EC2ToEBSVolumeExporter extends Collector implements MetricProvider,

public EC2ToEBSVolumeExporter(AccountProvider accountProvider,
AWSClientProvider awsClientProvider, MetricSampleBuilder metricSampleBuilder,
CollectorRegistry collectorRegistry, AWSApiCallRateLimiter rateLimiter, TagUtil tagUtil,
CollectorRegistry collectorRegistry, AWSApiCallRateLimiter rateLimiter,
TagUtil tagUtil,
ECSServiceDiscoveryExporter ecsServiceDiscoveryExporter,
TaskExecutorUtil taskExecutorUtil, ScrapeConfigProvider scrapeConfigProvider) {
this.accountProvider = accountProvider;
Expand Down Expand Up @@ -108,12 +110,13 @@ public List<Sample> call() {
}
}));
volumeFutures.add(
taskExecutorUtil.executeTenantTask(awsAccount.getTenant(), new CollectionBuilderTask<ResourceRelation>() {
@Override
public List<ResourceRelation> call() {
return buildResourceRelations(awsAccount, region);
}
}));
taskExecutorUtil.executeTenantTask(awsAccount.getTenant(),
new CollectionBuilderTask<ResourceRelation>() {
@Override
public List<ResourceRelation> call() {
return buildResourceRelations(awsAccount, region);
}
}));
}));
taskExecutorUtil.awaitAll(futures, allSamples::addAll);
taskExecutorUtil.awaitAll(volumeFutures, newAttachedVolumes::addAll);
Expand All @@ -126,9 +129,15 @@ public List<ResourceRelation> call() {
}

private List<ResourceRelation> buildResourceRelations(AWSAccount awsAccount, String region) {
ScrapeConfig scrapeConfig = scrapeConfigProvider.getScrapeConfig(awsAccount.getTenant());
if (!scrapeConfig.isFetchEC2Metadata()) {
log.info("Skipping EC2 Metadata fetch");
return Collections.emptyList();
}

Set<ResourceRelation> newAttachedVolumes = new HashSet<>();
String accountId = awsAccount.getAccountId();
Ec2Client ec2Client = awsClientProvider.getEc2Client(region, awsAccount);
Set<ResourceRelation> newAttachedVolumes = new HashSet<>();
try {
AtomicReference<String> nextToken = new AtomicReference<>();
do {
Expand Down Expand Up @@ -179,9 +188,12 @@ private List<ResourceRelation> buildResourceRelations(AWSAccount awsAccount, Str

private List<Sample> buildEC2InstanceMetrics(String region, AWSAccount awsAccount) {
List<Sample> samples = new ArrayList<>();

Ec2Client ec2Client = awsClientProvider.getEc2Client(region, awsAccount);
ScrapeConfig scrapeConfig = scrapeConfigProvider.getScrapeConfig(awsAccount.getTenant());
if (!scrapeConfig.isFetchEC2Metadata()) {
log.info("Skipping EC2 Metadata fetch");
return samples;
}
Ec2Client ec2Client = awsClientProvider.getEc2Client(region, awsAccount);
String accountId = awsAccount.getAccountId();
SortedMap<String, String> telemetryLabels = new TreeMap<>();
telemetryLabels.put(SCRAPE_REGION_LABEL, region);
Expand Down
44 changes: 27 additions & 17 deletions src/main/java/ai/asserts/aws/exporter/ECSTaskUtil.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@

import ai.asserts.aws.AWSApiCallRateLimiter;
import ai.asserts.aws.AWSClientProvider;
import ai.asserts.aws.ScrapeConfigProvider;
import ai.asserts.aws.TagUtil;
import ai.asserts.aws.TaskExecutorUtil;
import ai.asserts.aws.account.AWSAccount;
Expand Down Expand Up @@ -60,6 +61,8 @@ public class ECSTaskUtil {
private final AWSApiCallRateLimiter rateLimiter;
private final TagUtil tagUtil;
private final TaskExecutorUtil taskExecutorUtil;

private final ScrapeConfigProvider scrapeConfigProvider;
private final String envName;

public final Cache<String, TaskDefinition> taskDefsByARN = CacheBuilder.newBuilder()
Expand All @@ -77,12 +80,14 @@ public class ECSTaskUtil {

public ECSTaskUtil(AWSClientProvider awsClientProvider, ResourceMapper resourceMapper,
AWSApiCallRateLimiter rateLimiter,
TagUtil tagUtil, TaskExecutorUtil taskExecutorUtil) {
TagUtil tagUtil, TaskExecutorUtil taskExecutorUtil,
ScrapeConfigProvider scrapeConfigProvider) {
this.awsClientProvider = awsClientProvider;
this.resourceMapper = resourceMapper;
this.rateLimiter = rateLimiter;
this.tagUtil = tagUtil;
this.taskExecutorUtil = taskExecutorUtil;
this.scrapeConfigProvider = scrapeConfigProvider;
// If the exporter's environment name is marked, use this for ECS metrics
envName = getInstallEnvName();
}
Expand Down Expand Up @@ -322,21 +327,26 @@ private SubnetDetails getSubnetDetails(Task task, Resource taskResource) {
}

private String getVpcId(Resource taskResource, AtomicReference<String> subnetId) {
AtomicReference<String> id = new AtomicReference<>("");
Ec2Client ec2Client = awsClientProvider.getEc2Client(taskResource.getRegion(),
AWSAccount.builder()
.accountId(taskResource.getAccount())
.build());
DescribeSubnetsResponse r = rateLimiter.doWithRateLimit("EC2Client/describeSubnets",
ImmutableSortedMap.of(
SCRAPE_ACCOUNT_ID_LABEL, taskResource.getAccount(),
SCRAPE_REGION_LABEL, taskResource.getRegion(),
SCRAPE_OPERATION_LABEL, "EC2Client/describeSubnets"
),
() -> ec2Client.describeSubnets(DescribeSubnetsRequest.builder()
.subnetIds(subnetId.get())
.build()));
r.subnets().stream().findFirst().ifPresent(subnet -> id.set(subnet.vpcId()));
return id.get();
ScrapeConfig scrapeConfig = scrapeConfigProvider.getScrapeConfig(taskResource.getTenant());
if (scrapeConfig.isFetchEC2Metadata()) {
AtomicReference<String> id = new AtomicReference<>("");
Ec2Client ec2Client = awsClientProvider.getEc2Client(taskResource.getRegion(),
AWSAccount.builder()
.accountId(taskResource.getAccount())
.build());
DescribeSubnetsResponse r = rateLimiter.doWithRateLimit("EC2Client/describeSubnets",
ImmutableSortedMap.of(
SCRAPE_ACCOUNT_ID_LABEL, taskResource.getAccount(),
SCRAPE_REGION_LABEL, taskResource.getRegion(),
SCRAPE_OPERATION_LABEL, "EC2Client/describeSubnets"
),
() -> ec2Client.describeSubnets(DescribeSubnetsRequest.builder()
.subnetIds(subnetId.get())
.build()));
r.subnets().stream().findFirst().ifPresent(subnet -> id.set(subnet.vpcId()));
return id.get();
} else {
return "";
}
}
}
23 changes: 18 additions & 5 deletions src/test/java/ai/asserts/aws/exporter/EC2ToEBSVolumeExporterTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,14 +4,14 @@
*/
package ai.asserts.aws.exporter;

import ai.asserts.aws.AWSApiCallRateLimiter;
import ai.asserts.aws.AWSClientProvider;
import ai.asserts.aws.ScrapeConfigProvider;
import ai.asserts.aws.TagUtil;
import ai.asserts.aws.TaskExecutorUtil;
import ai.asserts.aws.TestTaskThreadPool;
import ai.asserts.aws.account.AccountProvider;
import ai.asserts.aws.account.AWSAccount;
import ai.asserts.aws.AWSApiCallRateLimiter;
import ai.asserts.aws.TagUtil;
import ai.asserts.aws.account.AccountProvider;
import ai.asserts.aws.config.ScrapeConfig;
import ai.asserts.aws.resource.Resource;
import ai.asserts.aws.resource.ResourceRelation;
Expand Down Expand Up @@ -103,8 +103,8 @@ public void afterPropertiesSet() throws Exception {
public void updateCollect() {
expect(accountProvider.getAccounts()).andReturn(ImmutableSet.of(account));
expect(awsClientProvider.getEc2Client("region", account)).andReturn(ec2Client).anyTimes();
expect(scrapeConfigProvider.getScrapeConfig("acme")).andReturn(scrapeConfig);

expect(scrapeConfigProvider.getScrapeConfig("acme")).andReturn(scrapeConfig).anyTimes();
expect(scrapeConfig.isFetchEC2Metadata()).andReturn(true).anyTimes();
DescribeInstancesRequest request = DescribeInstancesRequest.builder()
.filters(Filter.builder()
.name("vpc-id")
Expand Down Expand Up @@ -232,4 +232,17 @@ public void updateCollect() {
assertEquals(ImmutableList.of(metricFamilySamples), testClass.collect());
verifyAll();
}

@Test
public void updateCollect_skipEC2MetaFetch() {
expect(accountProvider.getAccounts()).andReturn(ImmutableSet.of(account));
expect(scrapeConfigProvider.getScrapeConfig("acme")).andReturn(scrapeConfig).anyTimes();
expect(scrapeConfig.isFetchEC2Metadata()).andReturn(false).anyTimes();

replayAll();
testClass.update();
assertEquals(ImmutableSet.of(), testClass.getAttachedVolumes());
assertEquals(ImmutableList.of(), testClass.collect());
verifyAll();
}
}
Loading

0 comments on commit 58e7532

Please sign in to comment.