mirror of
https://github.com/mongodb/mongo.git
synced 2024-12-01 09:32:32 +01:00
dcaf6432b7
- Add capability to generate report.json in powertest.py - Add capability to generate and read a configuration YML file in powertest.py
355 lines
14 KiB
Python
Executable File
355 lines
14 KiB
Python
Executable File
#!/usr/bin/env python
|
|
|
|
"""AWS EC2 instance launcher and controller."""
|
|
|
|
from __future__ import print_function
|
|
|
|
import collections
|
|
import datetime
|
|
import optparse
|
|
import sys
|
|
import time
|
|
|
|
import boto3
|
|
import botocore
|
|
import yaml
|
|
|
|
_MODES = ("status", "create", "start", "stop", "force-stop", "reboot", "terminate")
|
|
|
|
|
|
class AwsEc2(object):
|
|
"""Class to support controlling AWS EC2 istances."""
|
|
|
|
InstanceStatus = collections.namedtuple("InstanceStatus", [
|
|
"instance_id",
|
|
"image_id",
|
|
"instance_type",
|
|
"state",
|
|
"private_ip_address",
|
|
"public_ip_address",
|
|
"private_dns_name",
|
|
"public_dns_name",
|
|
"tags"
|
|
])
|
|
|
|
def __init__(self):
|
|
try:
|
|
self.connection = boto3.resource("ec2")
|
|
except botocore.exceptions.BotoCoreError:
|
|
print("Please configure your Boto environment variables or files, see"
|
|
" http://boto3.readthedocs.io/en/latest/guide/configuration.html"
|
|
" for the variable names, file names and precedence order.")
|
|
raise
|
|
|
|
@staticmethod
|
|
def wait_for_state(instance, state, wait_time_secs=0, show_progress=False):
|
|
"""Wait up to 'wait_time_secs' for instance to be in 'state'.
|
|
Return 0 if 'state' reached, 1 otherwise."""
|
|
if show_progress:
|
|
print("Waiting for instance {} to reach '{}' state".format(instance, state),
|
|
end="",
|
|
file=sys.stdout)
|
|
reached_state = False
|
|
end_time = time.time() + wait_time_secs
|
|
while True:
|
|
if show_progress:
|
|
print(".", end="", file=sys.stdout)
|
|
sys.stdout.flush()
|
|
try:
|
|
client_error = ""
|
|
time_left = end_time - time.time()
|
|
instance.load()
|
|
if instance.state["Name"] == state:
|
|
reached_state = True
|
|
break
|
|
if time_left <= 0:
|
|
break
|
|
except botocore.exceptions.ClientError as err:
|
|
# A ClientError exception can sometimes be generated, due to RequestLimitExceeded,
|
|
# so we ignore it and retry until we time out.
|
|
client_error = " {}".format(err.message)
|
|
|
|
wait_interval_secs = 15 if time_left > 15 else time_left
|
|
time.sleep(wait_interval_secs)
|
|
if show_progress:
|
|
if reached_state:
|
|
print(" Instance {}!".format(instance.state["Name"]), file=sys.stdout)
|
|
else:
|
|
print(" Instance in state '{}', failed to reach state '{}'{}!".format(
|
|
instance.state["Name"], state, client_error), file=sys.stdout)
|
|
sys.stdout.flush()
|
|
return 0 if reached_state else 1
|
|
|
|
def control_instance(self, mode, image_id, wait_time_secs=0, show_progress=False):
|
|
"""Controls an AMI instance. Returns 0 & status information, if successful."""
|
|
if mode not in _MODES:
|
|
raise ValueError(
|
|
"Invalid mode '{}' specified, choose from {}.".format(mode, _MODES))
|
|
|
|
sys.stdout.flush()
|
|
instance = self.connection.Instance(image_id)
|
|
try:
|
|
if mode == "start":
|
|
state = "running"
|
|
instance.start()
|
|
elif mode == "stop":
|
|
state = "stopped"
|
|
instance.stop()
|
|
elif mode == "force-stop":
|
|
state = "stopped"
|
|
instance.stop(Force=True)
|
|
elif mode == "terminate":
|
|
state = "terminated"
|
|
instance.terminate()
|
|
elif mode == "reboot":
|
|
state = "running"
|
|
instance.reboot()
|
|
else:
|
|
state = None
|
|
wait_time_secs = 0
|
|
except botocore.exceptions.ClientError as err:
|
|
return 1, err.message
|
|
|
|
ret = 0
|
|
if wait_time_secs > 0:
|
|
ret = self.wait_for_state(
|
|
instance=instance,
|
|
state=state,
|
|
wait_time_secs=wait_time_secs,
|
|
show_progress=show_progress)
|
|
try:
|
|
# Always provide status after executing command.
|
|
status = self.InstanceStatus(
|
|
getattr(instance, "instance_id", None),
|
|
getattr(instance, "image_id", None),
|
|
getattr(instance, "instance_type", None),
|
|
getattr(instance, "state", None),
|
|
getattr(instance, "private_ip_address", None),
|
|
getattr(instance, "public_ip_address", None),
|
|
getattr(instance, "private_dns_name", None),
|
|
getattr(instance, "public_dns_name", None),
|
|
getattr(instance, "tags", None))
|
|
except botocore.exceptions.ClientError as err:
|
|
return 1, err.message
|
|
|
|
return ret, status
|
|
|
|
def tag_instance(self, image_id, tags):
|
|
"""Tags an AMI instance. """
|
|
if tags:
|
|
# It's possible that ClientError code InvalidInstanceID.NotFound could be returned,
|
|
# even if the 'image_id' exists. We will retry up to 5 times, with increasing wait,
|
|
# if this occurs.
|
|
# http://docs.aws.amazon.com/AWSEC2/latest/APIReference/query-api-troubleshooting.html
|
|
for i in range(5):
|
|
try:
|
|
instance = self.connection.Instance(image_id)
|
|
break
|
|
except botocore.exceptions.ClientError as err:
|
|
if err.response["Error"]["Code"] != "InvalidInstanceID.NotFound":
|
|
raise
|
|
time.sleep(i + 1)
|
|
instance.create_tags(Tags=tags)
|
|
|
|
def launch_instance(self,
|
|
ami,
|
|
instance_type,
|
|
block_devices=None,
|
|
key_name=None,
|
|
security_groups=None,
|
|
tags=None,
|
|
wait_time_secs=0,
|
|
show_progress=False,
|
|
**kwargs):
|
|
"""Launches and tags an AMI instance.
|
|
|
|
Returns the tuple (0, status_information), if successful."""
|
|
|
|
bdms = []
|
|
if block_devices is None:
|
|
block_devices = {}
|
|
for block_device in block_devices:
|
|
bdm = {}
|
|
bdm["DeviceName"] = block_device
|
|
bdm["Ebs"] = {"DeleteOnTermination": True, "VolumeSize": block_devices[block_device]}
|
|
bdms.append(bdm)
|
|
if bdms:
|
|
kwargs["BlockDeviceMappings"] = bdms
|
|
if security_groups:
|
|
kwargs["SecurityGroups"] = security_groups
|
|
if key_name:
|
|
kwargs["KeyName"] = key_name
|
|
|
|
try:
|
|
instances = self.connection.create_instances(
|
|
ImageId=ami,
|
|
InstanceType=instance_type,
|
|
MaxCount=1,
|
|
MinCount=1,
|
|
**kwargs)
|
|
except (botocore.exceptions.ClientError, botocore.exceptions.ParamValidationError) as err:
|
|
return 1, err.message
|
|
|
|
instance = instances[0]
|
|
if wait_time_secs > 0:
|
|
self.wait_for_state(
|
|
instance=instance,
|
|
state="running",
|
|
wait_time_secs=wait_time_secs,
|
|
show_progress=show_progress)
|
|
|
|
self.tag_instance(instance.instance_id, tags)
|
|
|
|
return self.control_instance("status", instance.instance_id)
|
|
|
|
|
|
def main():
|
|
"""Main program."""
|
|
|
|
required_create_options = ["ami", "key_name"]
|
|
|
|
parser = optparse.OptionParser(description=__doc__)
|
|
control_options = optparse.OptionGroup(parser, "Control options")
|
|
create_options = optparse.OptionGroup(parser, "Create options")
|
|
|
|
parser.add_option("--mode",
|
|
dest="mode",
|
|
choices=_MODES,
|
|
default="status",
|
|
help="Operations to perform on an EC2 instance, choose one of"
|
|
" '{}', defaults to '%default'.".format(", ".join(_MODES)))
|
|
|
|
control_options.add_option("--imageId",
|
|
dest="image_id",
|
|
default=None,
|
|
help="EC2 image_id to perform operation on [REQUIRED for control].")
|
|
|
|
control_options.add_option("--waitTimeSecs",
|
|
dest="wait_time_secs",
|
|
type=int,
|
|
default=5 * 60,
|
|
help="Time to wait for EC2 instance to reach it's new state,"
|
|
" defaults to '%default'.")
|
|
|
|
create_options.add_option("--ami",
|
|
dest="ami",
|
|
default=None,
|
|
help="EC2 AMI to launch [REQUIRED for create].")
|
|
|
|
create_options.add_option("--blockDevice",
|
|
dest="block_devices",
|
|
metavar="DEVICE-NAME DEVICE-SIZE-GB",
|
|
action="append",
|
|
default=[],
|
|
nargs=2,
|
|
help="EBS device name and volume size in GiB."
|
|
" More than one device can be attached, by specifying"
|
|
" this option more than once."
|
|
" The device will be deleted on termination of the instance.")
|
|
|
|
create_options.add_option("--instanceType",
|
|
dest="instance_type",
|
|
default="t1.micro",
|
|
help="EC2 instance type to launch, defaults to '%default'.")
|
|
|
|
create_options.add_option("--keyName",
|
|
dest="key_name",
|
|
default=None,
|
|
help="EC2 key name [REQUIRED for create].")
|
|
|
|
create_options.add_option("--securityGroup",
|
|
dest="security_groups",
|
|
action="append",
|
|
default=[],
|
|
help="EC2 security group. More than one security group can be added,"
|
|
" by specifying this option more than once.")
|
|
|
|
create_options.add_option("--tagExpireHours",
|
|
dest="tag_expire_hours",
|
|
type=int,
|
|
default=2,
|
|
help="EC2 tag expire time in hours, defaults to '%default'.")
|
|
|
|
create_options.add_option("--tagName",
|
|
dest="tag_name",
|
|
default="",
|
|
help="EC2 tag and instance name.")
|
|
|
|
create_options.add_option("--tagOwner",
|
|
dest="tag_owner",
|
|
default="",
|
|
help="EC2 tag owner.")
|
|
|
|
create_options.add_option("--extraArgs",
|
|
dest="extra_args",
|
|
metavar="{key1: value1, key2: value2, ..., keyN: valueN}",
|
|
default=None,
|
|
help="EC2 create instance keyword args. The argument is specified as"
|
|
" bracketed YAML - i.e. JSON with support for single quoted"
|
|
" and unquoted keys. Example, '{DryRun: True}'")
|
|
|
|
parser.add_option_group(control_options)
|
|
parser.add_option_group(create_options)
|
|
|
|
(options, _) = parser.parse_args()
|
|
|
|
aws_ec2 = AwsEc2()
|
|
|
|
if options.mode == "create":
|
|
for option in required_create_options:
|
|
if not getattr(options, option, None):
|
|
parser.print_help()
|
|
parser.error("Missing required create option")
|
|
|
|
block_devices = {}
|
|
for (device_name, device_size) in options.block_devices:
|
|
try:
|
|
device_size = int(device_size)
|
|
except ValueError:
|
|
parser.print_help()
|
|
parser.error("Block size must be an integer")
|
|
block_devices[device_name] = device_size
|
|
|
|
# The 'expire-on' key is a UTC time.
|
|
expire_dt = datetime.datetime.utcnow() + datetime.timedelta(hours=options.tag_expire_hours)
|
|
tags = [{"Key": "expire-on", "Value": expire_dt.strftime("%Y-%m-%d %H:%M:%S")},
|
|
{"Key": "Name", "Value": options.tag_name},
|
|
{"Key": "owner", "Value": options.tag_owner}]
|
|
|
|
my_kwargs = {}
|
|
if options.extra_args is not None:
|
|
my_kwargs = yaml.safe_load(options.extra_args)
|
|
|
|
(ret_code, instance_status) = aws_ec2.launch_instance(
|
|
ami=options.ami,
|
|
instance_type=options.instance_type,
|
|
block_devices=block_devices,
|
|
key_name=options.key_name,
|
|
security_groups=options.security_groups,
|
|
tags=tags,
|
|
wait_time_secs=options.wait_time_secs,
|
|
show_progress=True,
|
|
**my_kwargs)
|
|
else:
|
|
if not getattr(options, "image_id", None):
|
|
parser.print_help()
|
|
parser.error("Missing required control option")
|
|
|
|
(ret_code, instance_status) = aws_ec2.control_instance(
|
|
mode=options.mode,
|
|
image_id=options.image_id,
|
|
wait_time_secs=options.wait_time_secs,
|
|
show_progress=True)
|
|
|
|
print("Return code: {}, Instance status:".format(ret_code))
|
|
if ret_code:
|
|
print(instance_status)
|
|
else:
|
|
for field in instance_status._fields:
|
|
print("\t{}: {}".format(field, getattr(instance_status, field)))
|
|
|
|
sys.exit(ret_code)
|
|
|
|
if __name__ == "__main__":
|
|
main()
|