awsloadtest.py

awsloadtest.py

#!/usr/bin/python

from multiprocessing import Pool
from subprocess import call, Popen
from tempfile import NamedTemporaryFile
import time

from boto import ec2

INSTANCE_NAME = 'vegeta-loadtest'
INSTANCE_COUNT = 10
AMI_ID = "ami-6cba8329"

c = ec2.connect_to_region('us-west-1')


def run_instances(count):
    r = c.run_instances(
        image_id=AMI_ID,
        key_name="oremj",
        security_groups=["wide-open"],
        instance_type="m3.large",
        max_count=count,
        min_count=count,
    )

    for i in r.instances:
        while i.update() != "running":
            time.sleep(5)

    c.create_tags([i.id for i in r.instances],
                  {'Name': INSTANCE_NAME})


def get_hosts():
    instances = c.get_only_instances(filters={'tag:Name': INSTANCE_NAME})

    if len(instances) < INSTANCE_COUNT:
        run_instances(INSTANCE_COUNT - len(instances))
        instances = c.get_only_instances(filters={'tag:Name': INSTANCE_NAME})

    return [i.public_dns_name for i in instances]


def copy_vegeta(host):
    if call(['ssh', host, '[[ -f /tmp/vegeta ]]']) == 0:
        return

    p = Popen(['scp', './vegeta', host + ":/tmp/."])
    p.communicate()
    if p.returncode != 0:
        raise Exception("Bad returncode: " + host)


def run_loadtest(args):
    host, outfile = args
    p = Popen(['ssh', host,
               'echo GET https://localhost/ '
               '| /tmp/vegeta attack -duration=2s'],
              stdout=outfile)
    p.communicate()


def print_report(files):
    Popen(['./vegeta', 'report', '-input', ','.join(files)]).communicate()  # noqa


hosts = get_hosts()

tmp_files = [NamedTemporaryFile() for t in range(len(hosts))]
tmpfile_names = [t.name for t in tmp_files]
tmpfile_fds = [t.file.fileno() for t in tmp_files]

process_pool = Pool(processes=10)
process_pool.map(copy_vegeta, hosts)

process_pool.map(run_loadtest, zip(hosts, tmpfile_fds))

print_report(tmpfile_names)