Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Flag to disable EC2 API usage #312

Merged
merged 1 commit into from
Aug 2, 2023
+186 −45
Merged

Flag to disable EC2 API usage #312

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions config/src/main/java/ai/asserts/aws/config/ScrapeConfig.java
View file
Open in desktop
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
Open in desktop
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
Open in desktop
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
Open in desktop
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
Loading