Merge pull request #68 from loads/feat/48

add more attack node tags (issue #48)

loadsbroker/aws.py

@@ -25,6 +25,7 @@
 import time
 from collections import defaultdict, namedtuple
 from datetime import datetime, timedelta
+from typing import Dict, Optional
 
 from boto.ec2 import connect_to_region
 from tornado import gen
@@ -52,6 +53,13 @@
 AWS_AMI_IDS = {k: {} for k in AWS_REGIONS}
 
 
+# How long after est. run times to trigger the reaper
+REAPER_DELTA = timedelta(hours=5)
+# Force the reaper for run times less than
+REAPER_FORCE = timedelta(hours=24)
+REAPER_STATE = 'ThirdState'
+
+
 def populate_ami_ids(aws_access_key_id=None, aws_secret_access_key=None,
                      port=None, owner_id="595879546273", use_filters=True):
     """Populate all the AMI ID's with the latest CoreOS stable info.
@@ -355,7 +363,7 @@ def __init__(self, broker_id, access_key=None, secret_key=None,
         self.security = security
         self.user_data = user_data
         self._instances = defaultdict(list)
-        self._tag_filters = {"tag:Name": "loads-%s" % self.broker_id,
+        self._tag_filters = {"tag:Name": "loads-%s*" % self.broker_id,
                              "tag:Project": "loads"}
         self._conns = {}
         self._recovered = {}
@@ -493,7 +501,7 @@ def _locate_existing_instances(self, count, inst_type, region):
 
         for inst in region_instances:
             if available_instance(inst) and inst_type == inst.instance_type:
-                    instances.append(inst)
+                instances.append(inst)
             else:
                 remaining.append(inst)
 
@@ -518,12 +526,15 @@ async def _allocate_instances(self, conn, count, inst_type, region):
         return reservations.instances
 
     async def request_instances(self,
-                                run_id,
-                                uuid,
+                                run_id: str,
+                                uuid: str,
                                 count=1,
                                 inst_type="t1.micro",
                                 region="us-west-2",
-                                allocate_missing=True):
+                                allocate_missing=True,
+                                plan: Optional[str] = None,
+                                owner: Optional[str] = None,
+                                run_max_time: Optional[int] = None):
         """Allocate a collection of instances.
 
         :param run_id: Run ID for these instances
@@ -535,6 +546,10 @@ async def request_instances(self,
             If there's insufficient existing instances for this uuid,
             whether existing or new instances should be allocated to the
             collection.
+        :param plan: Name of the instances' plan
+        :param owner: Owner name of the instances
+        :param run_max_time: Maximum expected run-time of instances in
+            seconds
         :returns: Collection of allocated instances
         :rtype: :class:`EC2Collection`
 
@@ -563,26 +578,33 @@ async def request_instances(self,
         if num > 0:
             new_instances = await self._allocate_instances(
                 conn, num, inst_type, region)
-            logger.debug("Allocated instances: %s", new_instances)
+            logger.debug("Allocated instances%s: %s",
+                         " (Owner: %s)" % owner if owner else "",
+                         new_instances)
             instances.extend(new_instances)
 
         # Tag all the instances
         if self.use_filters:
-            # Sometimes, we can get instance data back before the AWS API fully
-            # recognizes it, so we wait as needed.
+            tags = {
+                "Name": "loads-{}{}".format(self.broker_id,
+                                            "-" + plan if plan else ""),
+                "Project": "loads",
+                "RunId": run_id,
+                "Uuid": uuid,
+            }
+            if owner:
+                tags["Owner"] = owner
+            if run_max_time is not None:
+                self._tag_for_reaping(tags, run_max_time)
+
+            # Sometimes, we can get instance data back before the AWS
+            # API fully recognizes it, so we wait as needed.
             async def tag_instance(instance):
                 retries = 0
                 while True:
                     try:
                         await self._run_in_executor(
-                            conn.create_tags, [instance.id],
-                            {
-                                "Name": "loads-%s" % self.broker_id,
-                                "Project": "loads",
-                                "RunId": run_id,
-                                "Uuid": uuid
-                            }
-                        )
+                            conn.create_tags, [instance.id], tags)
                         break
                     except:
                         if retries > 5:
@@ -592,6 +614,23 @@ async def tag_instance(instance):
             await gen.multi([tag_instance(x) for x in instances])
         return EC2Collection(run_id, uuid, conn, instances, self._loop)
 
+    def _tag_for_reaping(self,
+                         tags: Dict[str, str],
+                         run_max_time: int) -> None:
+        """Tag an instance for the mozilla-services reaper
+
+        Set instances to stop after run_max_time + REAPER_DELTA
+
+        """
+        now = datetime.utcnow()
+        reap = now + timedelta(seconds=run_max_time) + REAPER_DELTA
+        tags['REAPER'] = "{}|{:{dfmt}}".format(
+            REAPER_STATE, reap, dfmt="%Y-%m-%d %I:%M%p UTC")
+        if reap < now + REAPER_FORCE:
+            # the reaper ignores instances younger than REAPER_FORCE
+            # unless forced w/ REAP_ME
+            tags['REAP_ME'] = ""
+
     async def release_instances(self, collection):
         """Return a collection of instances to the pool.
 

loadsbroker/broker.py

@@ -98,10 +98,12 @@ class RunHelpers:
 
 
 class Broker:
-    def __init__(self, io_loop, sqluri, ssh_key,
+    def __init__(self, name, io_loop, sqluri, ssh_key,
                  heka_options, influx_options, aws_port=None,
                  aws_owner_id="595879546273", aws_use_filters=True,
                  aws_access_key=None, aws_secret_key=None, initial_db=None):
+        self.name = name
+        logger.debug("loads-broker (%s)", self.name)
 
         self.loop = io_loop
         self._base_env = BASE_ENV.copy()
@@ -133,7 +135,7 @@ def __init__(self, io_loop, sqluri, ssh_key,
                 raise ImportError('You need to install the influx lib')
             self.influx = InfluxDBClient(**influx_args)
 
-        self.pool = aws.EC2Pool("1234", user_data=user_data,
+        self.pool = aws.EC2Pool(self.name, user_data=user_data,
                                 io_loop=self.loop, port=aws_port,
                                 owner_id=aws_owner_id,
                                 use_filters=aws_use_filters,
@@ -228,7 +230,7 @@ def run_plan(self, strategy_id, create_db=True, **kwargs):
 
         log_threadid("Running strategy: %s" % strategy_id)
         uuid = kwargs.pop('run_uuid', None)
-        creator = kwargs.pop('creator', None)
+        owner = kwargs.pop('owner', None)
 
         # now we can start a new run
         try:
@@ -240,7 +242,7 @@ def run_plan(self, strategy_id, create_db=True, **kwargs):
                 plan_uuid=strategy_id,
                 run_uuid=uuid,
                 additional_env=kwargs,
-                creator=creator)
+                owner=owner)
         except NoResultFound as e:
             raise LoadsException(str(e))
 
@@ -335,7 +337,7 @@ def _get_state(self):
 
     @classmethod
     def new_run(cls, run_helpers, db_session, pool, io_loop, plan_uuid,
-                run_uuid=None, additional_env=None, creator=None):
+                run_uuid=None, additional_env=None, owner=None):
         """Create a new run manager for the given strategy name
 
         This creates a new run for this strategy and initializes it.
@@ -354,7 +356,7 @@ def new_run(cls, run_helpers, db_session, pool, io_loop, plan_uuid,
         """
         # Create the run for this manager
         logger.debug('Starting a new run manager')
-        run = Run.new_run(db_session, plan_uuid, creator)
+        run = Run.new_run(db_session, plan_uuid, owner)
         if run_uuid:
             run.uuid = run_uuid
         db_session.add(run)
@@ -391,10 +393,15 @@ async def _get_steps(self):
         logger.debug('Getting steps & collections')
         steps = self.run.plan.steps
         collections = await gen.multi(
-            [self._pool.request_instances(self.run.uuid, s.uuid,
-                                          count=s.instance_count,
-                                          inst_type=s.instance_type,
-                                          region=s.instance_region)
+            [self._pool.request_instances(
+                self.run.uuid,
+                s.uuid,
+                count=s.instance_count,
+                inst_type=s.instance_type,
+                region=s.instance_region,
+                plan=self.run.plan.name,
+                owner=self.run.owner,
+                run_max_time=s.run_delay + s.run_max_time)
              for s in steps])
 
         try:

loadsbroker/client/cmd_delete.py

@@ -9,4 +9,5 @@ def __call__(self, args):
         url = '/run/' + args.run_id + '?purge=1'
         return self.session.delete(self.root + url).json()
 
+
 cmd = Delete

loadsbroker/client/cmd_run.py

@@ -24,4 +24,5 @@ def __call__(self, args):
                               data=json.dumps(options), headers=headers)
         return r.json()
 
+
 cmd = Run

loadsbroker/client/cmd_status.py

@@ -12,4 +12,5 @@ def __call__(self, args):
         res = self.session.get(url)
         return res.json()
 
+
 cmd = Status

loadsbroker/db.py

@@ -342,7 +342,7 @@ class Run(Base):
     """
     state = Column(Integer, default=INITIALIZING)
 
-    creator = Column(String, nullable=True)
+    owner = Column(String, nullable=True)
 
     created_at = Column(DateTime, default=datetime.datetime.utcnow,
                         doc="When the Run was created.")
@@ -358,7 +358,7 @@ class Run(Base):
     plan_id = Column(Integer, ForeignKey("plan.id"))
 
     @classmethod
-    def new_run(cls, session, plan_uuid, creator=None):
+    def new_run(cls, session, plan_uuid, owner=None):
         """Create a new run with appropriate running container set
         linkage for a given strategy"""
         plan = Plan.load_with_steps(session, plan_uuid)
@@ -367,7 +367,7 @@ def new_run(cls, session, plan_uuid, creator=None):
 
         run = cls()
         run.plan = plan
-        run.creator = creator
+        run.owner = owner
 
         # Setup step records for each step in this plan
         run.step_records = [StepRecord.from_step(step) for step in plan.steps]
@@ -383,7 +383,7 @@ def json(self, fields=None):
                 'completed_at': self._datetostr(self.completed_at),
                 'started_at': self._datetostr(self.started_at),
                 'plan_id': self.plan_id, 'plan_name': self.plan.name,
-                'creator': self.creator}
+                'owner': self.owner}
 
 
 run_table = Run.__table__

loadsbroker/extensions.py

@@ -359,7 +359,6 @@ async def start(self,
 
         # Upload heka config to all the instances
         def upload_files(inst):
-            hostname = ""
             hostname = "%s%s" % (
                 series_name,
                 inst.instance.ip_address.replace('.', '_')

loadsbroker/main.py

@@ -17,6 +17,8 @@ def _parse(sysargs=None):
         sysargs = sys.argv[1:]
 
     parser = argparse.ArgumentParser(description='Runs a Loads broker.')
+    parser.add_argument('--name', help="Name of this broker instance",
+                        type=str, default="1234")
     parser.add_argument('-p', '--port', help='HTTP Port', type=int,
                         default=8080)
     parser.add_argument('--debug', help='Debug Info.', action='store_true',
@@ -86,7 +88,7 @@ def main(sysargs=None):
                                        args.influx_user, args.influx_password,
                                        args.influx_secure)
 
-    application.broker = Broker(loop, args.database, args.ssh_key,
+    application.broker = Broker(args.name, loop, args.database, args.ssh_key,
                                 heka_options,
                                 influx_options,
                                 aws_port=args.aws_port,

loadsbroker/tests/fakedocker.py

@@ -106,5 +106,6 @@ def main():
     except KeyboardInterrupt:
         pass
 
+
 if __name__ == "__main__":
     main()

loadsbroker/tests/fakeinflux.py

@@ -81,5 +81,6 @@ def main():
     except KeyboardInterrupt:
         pass
 
+
 if __name__ == "__main__":
     main()

loadsbroker/tests/test_api.py

@@ -38,7 +38,7 @@ def get_app(self):
     def _createBroker(self):
         from loadsbroker.broker import Broker
         from mock import Mock
-        return Broker(self.io_loop, self.db_uri, None,
+        return Broker("1234", self.io_loop, self.db_uri, None,
                       Mock(spec=HekaOptions),
                       Mock(spec=InfluxOptions),
                       aws_use_filters=False, initial_db=None,

loadsbroker/tests/test_units/test_aws.py

@@ -1,4 +1,6 @@
 import unittest
+from datetime import datetime, timedelta
+
 from tornado.testing import AsyncTestCase, gen_test
 from moto import mock_ec2
 import boto
@@ -264,6 +266,60 @@ async def test_allocates_instances_for_collection(self):
                                             region=region)
         self.assertEqual(len(coll.instances), 5)
 
+    @gen_test
+    async def test_tags(self):
+        region = "us-west-2"
+        conn = boto.ec2.connect_to_region(region)
+        reservation = conn.run_instances('ami-1234abcd')
+        instance = reservation.instances[0]
+        conn.create_image(instance.id, "CoreOS stable")
+
+        broker_id = "br12"
+        owner = "otto.push"
+        plan = "Loadtest"
+
+        pool = self._callFUT(broker_id)
+        pool.use_filters = True
+        await pool.ready
+
+        coll = await pool.request_instances("run_12", "12423", 5,
+                                            inst_type="m1.small",
+                                            region=region,
+                                            plan=plan,
+                                            owner=owner)
+        ids = {ec2instance.instance.id for ec2instance in coll.instances}
+
+        tagged = 0
+        for reservation in conn.get_all_instances():
+            for instance in reservation.instances:
+                if instance.id not in ids:
+                    continue
+                self.assertEqual(instance.tags['Owner'], owner)
+                self.assertEqual(instance.tags['Name'],
+                                 "loads-{}-{}".format(broker_id, plan))
+                tagged += 1
+        self.assertEqual(tagged, len(ids))
+
+    @gen_test
+    async def test_reaper_tags(self):
+        pool = self._callFUT("br12")
+        await pool.ready
+
+        run_max_time = 64800
+        now = datetime.utcnow().replace(second=0, microsecond=0)
+        reap_time = now + timedelta(hours=5, seconds=run_max_time)
+
+        tags = {}
+        pool._tag_for_reaping(tags, run_max_time)
+        self.assertIn('REAPER', tags)
+        # reap time < 24 hours so force reaping
+        self.assertIn('REAP_ME', tags)
+
+        state, ts = tags['REAPER'].split('|')
+        self.assertEqual(state, 'ThirdState')
+        dt = datetime.strptime(ts, '%Y-%m-%d %I:%M%p %Z')
+        self.assertTrue(reap_time <= dt <= reap_time + timedelta(minutes=1))
+
     @gen_test
     async def test_allocates_recovered_for_collection(self):
         region = "us-west-2"