Skip to content

Commit

Permalink
Add collecting stat for EC2
Browse files Browse the repository at this point in the history
  • Loading branch information
asmorodskyi committed May 30, 2024
1 parent 2c40fab commit dec26df
Show file tree
Hide file tree
Showing 4 changed files with 71 additions and 9 deletions.
27 changes: 27 additions & 0 deletions ocw/lib/dump_state.py
Original file line number Diff line number Diff line change
@@ -1,14 +1,22 @@
import os
import logging
import traceback
from webui.PCWConfig import PCWConfig
from ocw.lib.azure import Azure
from ocw.lib.ec2 import EC2
from ocw.enums import ProviderChoice
from ocw.lib.influx import Influx

logger = logging.getLogger(__name__)


def dump_state():
if os.getenv("INFLUX_TOKEN") is None:
logger.warning("INFLUX_TOKEN is not set, dumping state is not possible")
return
if not PCWConfig.has("influxdb/url"):
logger.warning("pcw.ini missing influxdb configuration, dumping state is not possible")
return
for namespace in PCWConfig.get_namespaces_for("influxdb"):
try:
providers = PCWConfig.get_providers_for("influxdb", namespace)
Expand Down Expand Up @@ -38,6 +46,25 @@ def dump_state():
namespace,
Azure(namespace).get_img_versions_count,
)
if ProviderChoice.EC2 in providers:
Influx().dump_resource(
ProviderChoice.EC2.value,
Influx.VMS_QUANTITY,
namespace,
EC2(namespace).count_all_instances
)
Influx().dump_resource(
ProviderChoice.EC2.value,
Influx.IMAGES_QUANTITY,
namespace,
EC2(namespace).count_all_images
)
Influx().dump_resource(
ProviderChoice.EC2.value,
Influx.VOLUMES_QUANTITY,
namespace,
EC2(namespace).count_all_volumes
)
except Exception:
logger.exception(
"[%s] Dump state failed!: \n %s", namespace, traceback.format_exc()
Expand Down
21 changes: 21 additions & 0 deletions ocw/lib/ec2.py
Original file line number Diff line number Diff line change
Expand Up @@ -116,6 +116,13 @@ def volume_protected(self, volume: dict) -> bool:
def list_instances(self, region: str) -> list:
return list(self.ec2_resource(region).instances.all())

def count_all_instances(self) -> int:
instance_quantity = 0
for region in self.all_regions:
instances = self.list_instances(region=region)
instance_quantity += len(instances)
return instance_quantity

def get_all_regions(self) -> list:
regions_resp = self.ec2_client(EC2.default_region).describe_regions()
regions = [region['RegionName'] for region in regions_resp['Regions']]
Expand Down Expand Up @@ -324,6 +331,20 @@ def report_cleanup_results(self, vpc_errors: list, vpc_notify: list, vpc_locked:
if len(vpc_locked) > 0:
send_mail('VPC deletion locked by running VMs', '\n'.join(vpc_locked))

def count_all_images(self) -> int:
all_images_cnt = 0
for region in self.all_regions:
response = self.ec2_client(region).describe_images(Owners=['self'])
all_images_cnt += len(response['Images'])
return all_images_cnt

def count_all_volumes(self) -> int:
all_volumes_cnt = 0
for region in self.all_regions:
response = self.ec2_client(region).describe_volumes()
all_volumes_cnt += len(response['Volumes'])
return all_volumes_cnt

def cleanup_images(self, valid_period_days: float) -> None:
self.log_dbg('Call cleanup_images')
for region in self.all_regions:
Expand Down
1 change: 1 addition & 0 deletions ocw/lib/influx.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ class Influx:
VMS_QUANTITY: str = "vms_quantity"
IMAGES_QUANTITY: str = "images_quantity"
DISK_QUANTITY: str = "disk_quantity"
VOLUMES_QUANTITY: str = "volumes_quanity"
IMAGE_VERSION_QUANTITY: str = "img_version_quantity"
NAMESPACE_TAG: str = "namespace"

Expand Down
31 changes: 22 additions & 9 deletions tests/test_ec2.py
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,22 @@ def mock_local_get_feature_property(feature: str, property: str, namespace: str
monkeypatch.setattr('smtplib.SMTP', lambda arg1, arg2: MockedSMTP())
return EC2('fake')

class MockedInstances:
is_empty = False

def __init__(self, subnet=True) -> None:
self.subnet = subnet

def all(self):
if self.subnet:
if MockedInstances.is_empty:
return []
else:
return ['subnet']
else:
return 5



class MockedEC2Client():
response = {}
Expand Down Expand Up @@ -127,6 +143,9 @@ class MockedEC2Client():
}
]}

def __init__(self, region):
self.instances = MockedInstances(subnet=False)

ec2_snapshots = {snapshotid_to_delete: 'snapshot', snapshotid_i_have_ami: 'snapshot'}

def describe_images(self, *args, **kwargs):
Expand Down Expand Up @@ -185,15 +204,6 @@ def sendmail(self, sender_email, receiver_email, mimetext):
MockedSMTP.mimetext = mimetext


class MockedInstances:
is_empty = False

def all(self):
if MockedInstances.is_empty:
return []
else:
return ['subnet']


class MockedInterface:
delete_called = False
Expand Down Expand Up @@ -507,3 +517,6 @@ def mocked_get_boolean(config_path, field=None):
ec2_patch.cleanup_all()

assert called_stack == ['cleanup_images', 'cleanup_snapshots', 'cleanup_volumes', 'cleanup_vpcs']

def test_list_all_instances(ec2_patch):
assert ec2_patch.count_all_instances() == 5

0 comments on commit dec26df

Please sign in to comment.