Add ability to attach a drone set to a host. This will restrict the host
from running on any other drones.

Signed-off-by: James Ren <[email protected]>

--- autotest/cli/host.py        2010-03-24 11:22:06.000000000 -0700
+++ autotest/cli/host.py        2010-03-24 11:22:06.000000000 -0700
@@ -85,6 +85,52 @@
         return self.hosts
 
 
+class has_drone_set(object):
+    """Mixin for handling drone sets"""
+    _static_data = None
+
+    def _parse_drone_set(self, options):
+        self.drone_set = options.drone_set
+        if self.drone_set:
+            self.data['drone_set'] = self.drone_set
+        self.remove_drone_set = getattr(options, 'remove_drone_set', None)
+
+
+    def _add_drone_set_parser_option(self, remove_option):
+        self.parser.add_option('-d', '--drone-set',
+                               help='Specify the drone set for this host')
+        if remove_option:
+            self.parser.add_option('-x', '--remove-drone-set',
+                                   help=('Remove the associated drone set. '
+                                         'Cannot be used in conjunction with '
+                                         '-d'),
+                                   action='store_true')
+
+
+    def _check_drone_set(self, check_required, check_remove):
+        if not self._static_data:
+            self._static_data = self.execute_rpc('get_static_data')
+
+        drone_sets = self._static_data['drone_sets']
+        drone_sets_str = ', '.join('"%s"' % d for d in drone_sets)
+
+        if (check_required and self._static_data['drone_set_required']
+                and not self.drone_set):
+            raise self.invalid_syntax(
+                    'Drone set is required. Choices are: %s' % drone_sets_str)
+        if self.drone_set and self.drone_set not in drone_sets:
+            raise self.invalid_syntax(
+                    'Drone set must be one of: %s' % drone_sets_str)
+        if check_remove and self.remove_drone_set:
+            if self.drone_set:
+                raise self.invalid_syntax('Cannot specify both -x and -d')
+            if self._static_data['drone_set_required']:
+                raise self.invalid_syntax(
+                        'Drone set is required, cannot remove')
+            self.data['drone_set'] = None
+            self.messages.append('Drone set removed')
+
+
 class host_help(host):
     """Just here to get the atest logic working.
     Usage is set by its parent"""
@@ -304,8 +350,9 @@
                                                 'status'])
 
 
-class host_mod(host):
+class host_mod(host, has_drone_set):
     """atest host mod --lock|--unlock|--protection
+    --drone-set <name>|--remove-drone-set
     --mlist <file>|<hosts>"""
     usage_action = 'mod'
 
@@ -326,6 +373,7 @@
                                      ', '.join('"%s"' % p
                                                for p in self.protections)),
                                choices=self.protections)
+        self._add_drone_set_parser_option(remove_option=True)
 
 
     def parse(self):
@@ -338,8 +386,12 @@
             self.data['protection'] = options.protection
             self.messages.append('Protection set to "%s"' % options.protection)
 
+        self._parse_drone_set(options)
+        self._check_drone_set(check_required=False, check_remove=True)
+
         if len(self.data) == 0:
             self.invalid_syntax('No modification requested')
+
         return (options, leftover)
 
 
@@ -365,11 +417,12 @@
 
 
 
-class host_create(host):
+class host_create(host, has_drone_set):
     """atest host create [--lock|--unlock --platform <arch>
     --labels <labels>|--blist <label_file>
     --acls <acls>|--alist <acl_file>
     --protection <protection_type>
+    --drone-set <name>
     --mlist <mach_file>] <hosts>"""
     usage_action = 'create'
 
@@ -403,6 +456,7 @@
                                      ', '.join('"%s"' % p
                                                for p in self.protections)),
                                choices=self.protections)
+        self._add_drone_set_parser_option(remove_option=False)
 
 
     def parse(self):
@@ -422,6 +476,10 @@
         self.platform = getattr(options, 'platform', None)
         if options.protection:
             self.data['protection'] = options.protection
+
+        self._parse_drone_set(options)
+        self._check_drone_set(check_required=True, check_remove=False)
+
         return (options, leftover)
 
 
--- autotest/cli/host_unittest.py       2010-03-24 11:22:06.000000000 -0700
+++ autotest/cli/host_unittest.py       2010-03-24 11:22:06.000000000 -0700
@@ -1198,7 +1198,15 @@
                                    'testjob'])
 
 
-class host_mod_unittest(cli_mock.cli_unittest):
+class unittest_with_drone_set(cli_mock.cli_unittest):
+    def setUp(self):
+        super(unittest_with_drone_set, self).setUp()
+        self.mock_rpcs([(('get_static_data'), {}, True,
+                        {'drone_set_required': False,
+                         'drone_sets': ('drone_set',)})])
+
+
+class host_mod_unittest(unittest_with_drone_set):
     def test_execute_lock_one_host(self):
         self.run_cmd(argv=['atest', 'host', 'mod',
                            '--lock', 'host0', '--ignore_site_file'],
@@ -1264,7 +1272,7 @@
 
 
 
-class host_create_unittest(cli_mock.cli_unittest):
+class host_create_unittest(unittest_with_drone_set):
     def test_execute_create_muliple_hosts_all_options(self):
         self.run_cmd(argv=['atest', 'host', 'create', '--lock',
                            '-b', 'label0', '--acls', 'acl0', 'host0', 'host1',
--- autotest/frontend/afe/admin.py      2010-03-24 11:22:06.000000000 -0700
+++ autotest/frontend/afe/admin.py      2010-03-24 11:22:06.000000000 -0700
@@ -135,6 +135,9 @@
 
 admin.site.register(models.AclGroup, AclGroupAdmin)
 
+admin.site.register(models.Drone)
+admin.site.register(models.DroneSet)
+
 
 if settings.FULL_ADMIN:
     class JobAdmin(SiteAdmin):
--- autotest/frontend/afe/doctests/001_rpc_test.txt     2010-03-24 
11:22:06.000000000 -0700
+++ autotest/frontend/afe/doctests/001_rpc_test.txt     2010-03-24 
11:22:06.000000000 -0700
@@ -87,12 +87,15 @@
 
 # all the add*, modify*, delete*, and get* methods work the same way
 # hosts...
->>> rpc_interface.add_host(hostname='ipaj1', locked=True)
+>>> drone_set = models.DroneSet.objects.create(name='drone_set')
+>>> rpc_interface.add_host(hostname='ipaj1', locked=True, 
drone_set='drone_set')
 1
 >>> data = rpc_interface.get_hosts()
 
 # delete the lock_time field, since that can't be reliably checked
 >>> del data[0]['lock_time']
+>>> if drone_set:
+...     drone_set = drone_set.name
 >>> data == [{'id': 1,
 ...           'hostname': 'ipaj1',
 ...           'locked': 1,
@@ -106,7 +109,8 @@
 ...           'invalid': 0,
 ...           'protection': 'No protection',
 ...           'locked_by': 'debug_user',
-...           'dirty': True}]
+...           'dirty': True,
+...           'drone_set': 'drone_set'}]
 True
 >>> rpc_interface.modify_host('ipaj1', status='Hello')
 Traceback (most recent call last):
@@ -216,9 +220,9 @@
 # ###################################
 
 # first, create some hosts and labels to play around with
->>> rpc_interface.add_host(hostname='host1')
+>>> rpc_interface.add_host(hostname='host1', drone_set='drone_set')
 2
->>> rpc_interface.add_host(hostname='host2')
+>>> rpc_interface.add_host(hostname='host2', drone_set='drone_set')
 3
 >>> rpc_interface.add_label(name='label1')
 2
@@ -404,9 +408,9 @@
 ...                        test_category='Test',
 ...                        test_class='Kernel', path=test_control_path)
 2
->>> rpc_interface.add_host(hostname='my_label_host1')
+>>> rpc_interface.add_host(hostname='my_label_host1', drone_set='drone_set')
 4
->>> rpc_interface.add_host(hostname='my_label_host2')
+>>> rpc_interface.add_host(hostname='my_label_host2', drone_set='drone_set')
 5
 >>> rpc_interface.label_add_hosts(id='my_label', hosts=['my_label_host1', 
 >>> 'my_label_host2'])
 
@@ -554,7 +558,8 @@
 ...           'protection': 'No protection',
 ...           'locked_by': None,
 ...           'lock_time': None,
-...           'dirty': True},
+...           'dirty': True,
+...           'drone_set': drone_set},
 ...  'id': 1,
 ...  'job': job, # full job info here
 ...  'meta_host': None,
@@ -721,10 +726,10 @@
 >>> data = rpc_interface.get_atomic_groups()
 >>> assert data[0]['max_number_of_machines'] == 8, data
 
->>> unused = rpc_interface.add_host(hostname='ahost1')
->>> unused = rpc_interface.add_host(hostname='ahost2')
->>> unused = rpc_interface.add_host(hostname='ah3-blue')
->>> unused = rpc_interface.add_host(hostname='ah4-blue')
+>>> unused = rpc_interface.add_host(hostname='ahost1', drone_set='drone_set')
+>>> unused = rpc_interface.add_host(hostname='ahost2', drone_set='drone_set')
+>>> unused = rpc_interface.add_host(hostname='ah3-blue', drone_set='drone_set')
+>>> unused = rpc_interface.add_host(hostname='ah4-blue', drone_set='drone_set')
 >>> two_id = rpc_interface.add_label(name='two-label')
 >>> rpc_interface.label_add_hosts(two_id, ['ahost1', 'ahost2',
 ...                                        'ah3-blue', 'ah4-blue'])
--- autotest/frontend/afe/frontend_test_utils.py        2010-03-24 
11:22:06.000000000 -0700
+++ autotest/frontend/afe/frontend_test_utils.py        2010-03-24 
11:22:06.000000000 -0700
@@ -11,7 +11,13 @@
         acl_group = models.AclGroup.objects.create(name='my_acl')
         acl_group.users.add(models.User.current_user())
 
-        self.hosts = [models.Host.objects.create(hostname=hostname)
+        drone_set = models.DroneSet.objects.create(name='drone_set')
+        if models.Host.default_drone_set_name:
+            models.DroneSet.objects.create(
+                    name=models.Host.default_drone_set_name)
+
+        self.hosts = [models.Host.objects.create(hostname=hostname,
+                                                 drone_set=drone_set)
                       for hostname in
                       ('host1', 'host2', 'host3', 'host4', 'host5', 'host6',
                        'host7', 'host8', 'host9')]
--- autotest/frontend/afe/model_logic.py        2010-03-24 11:22:06.000000000 
-0700
+++ autotest/frontend/afe/model_logic.py        2010-03-24 11:22:06.000000000 
-0700
@@ -629,7 +629,10 @@
             if data.get(name) is not None:
                 continue
             if obj.default is not dbmodels.fields.NOT_PROVIDED:
-                new_data[name] = obj.default
+                if callable(obj.default):
+                    new_data[name] = obj.default()
+                else:
+                    new_data[name] = obj.default
             elif (isinstance(obj, dbmodels.CharField) or
                   isinstance(obj, dbmodels.TextField)):
                 new_data[name] = ''
--- autotest/frontend/afe/models.py     2010-03-24 11:22:06.000000000 -0700
+++ autotest/frontend/afe/models.py     2010-03-24 11:22:06.000000000 -0700
@@ -179,6 +179,71 @@
         return unicode(self.login)
 
 
+class Drone(dbmodels.Model, model_logic.ModelExtensions):
+    """A scheduler drone
+
+    hostname: the drone's hostname
+    """
+    hostname = dbmodels.CharField(max_length=255, unique=True)
+
+    name_field = 'hostname'
+    objects = model_logic.ExtendedManager()
+
+
+    def save(self, *args, **kwargs):
+        if not User.current_user().is_superuser():
+            raise Exception('Only superusers may edit drones')
+        super(Drone, self).save(*args, **kwargs)
+
+
+    def delete(self):
+        if not User.current_user().is_superuser():
+            raise Exception('Only superusers may delete drones')
+        super(Drone, self).delete()
+
+
+    class Meta:
+        db_table = 'afe_drones'
+
+    def __unicode__(self):
+        return unicode(self.hostname)
+
+
+class DroneSet(dbmodels.Model, model_logic.ModelExtensions):
+    """A set of scheduler drones
+
+    These will be used by the scheduler to decide what drones a host is allowed
+    to run on.
+
+    name: the drone set's name
+    drones: the drones that are part of this set
+    """
+    name = dbmodels.CharField(max_length=255, unique=True)
+    drones = dbmodels.ManyToManyField(Drone, db_table='afe_drone_sets_drones')
+
+    name_field = 'name'
+    objects = model_logic.ExtendedManager()
+
+
+    def save(self, *args, **kwargs):
+        if not User.current_user().is_superuser():
+            raise Exception('Only superusers may edit drone sets')
+        super(DroneSet, self).save(*args, **kwargs)
+
+
+    def delete(self):
+        if not User.current_user().is_superuser():
+            raise Exception('Only superusers may delete drone sets')
+        super(DroneSet, self).delete()
+
+
+    class Meta:
+        db_table = 'afe_drone_sets'
+
+    def __unicode__(self):
+        return unicode(self.name)
+
+
 class Host(model_logic.ModelWithInvalid, dbmodels.Model,
            model_logic.ModelWithAttributes):
     """\
@@ -196,11 +261,17 @@
     locked_by: user that locked the host, or null if the host is unlocked
     lock_time: DateTime at which the host was locked
     dirty: true if the host has been used without being rebooted
+    drone_set: the set of drones this host is allowed to run on. A NULL value
+               here means that the host can run on any drone
     """
     Status = enum.Enum('Verifying', 'Running', 'Ready', 'Repairing',
                        'Repair Failed', 'Dead', 'Cleaning', 'Pending',
                        string_values=True)
     Protection = host_protections.Protection
+    drone_set_required = global_config.global_config.get_config_value(
+            'SCHEDULER', 'drone_set_required', type=bool)
+    default_drone_set_name = (global_config.global_config.get_config_value(
+            'SCHEDULER', 'default_drone_set_name', default=None))
 
     hostname = dbmodels.CharField(max_length=255, unique=True)
     labels = dbmodels.ManyToManyField(Label, blank=True,
@@ -220,6 +291,10 @@
     lock_time = dbmodels.DateTimeField(null=True, blank=True, editable=False)
     dirty = dbmodels.BooleanField(default=True, editable=settings.FULL_ADMIN)
 
+    drone_set = dbmodels.ForeignKey(DroneSet,
+                                    null=(not drone_set_required),
+                                    blank=(not drone_set_required))
+
     name_field = 'hostname'
     objects = model_logic.ModelWithInvalidManager()
     valid_objects = model_logic.ValidObjectsManager()
@@ -230,11 +305,36 @@
         self._record_attributes(['status'])
 
 
+    @classmethod
+    def verify_default_drone_set_exists(cls):
+        """Verifies that the specified default drone set exists
+
+        Raises an Exception if a default drone set is specified, but does not
+        exist in the database
+        """
+        if Host.default_drone_set_name:
+            try:
+                DroneSet.objects.get(
+                        name=Host.default_drone_set_name)
+            except DroneSet.DoesNotExist:
+                raise Exception(
+                        'One-time host default drone set %s does not exist'
+                        % Host.default_drone_set_name)
+
+
+    @classmethod
+    def default_drone_set(cls):
+        if cls.default_drone_set_name:
+            return DroneSet.objects.get(name=cls.default_drone_set_name)
+        return None
+
+
     @staticmethod
     def create_one_time_host(hostname):
+        drone_set = Host.default_drone_set()
         query = Host.objects.filter(hostname=hostname)
         if query.count() == 0:
-            host = Host(hostname=hostname, invalid=True)
+            host = Host(hostname=hostname, invalid=True, drone_set=drone_set)
             host.do_validate()
         else:
             host = query[0]
@@ -244,6 +344,7 @@
                         'Select it rather than entering it as a one time '
                         'host.' % hostname
                     })
+            host.drone_set = drone_set
         host.protection = host_protections.Protection.DO_NOT_REPAIR
         host.locked = False
         host.save()
@@ -699,6 +800,7 @@
         and filling in the rest of the necessary information.
         """
         AclGroup.check_for_acl_violation_hosts(hosts)
+        cls._check_for_conflicting_drone_sets(hosts)
 
         user = User.current_user()
         if options.get('reboot_before') is None:
@@ -731,6 +833,30 @@
         return job
 
 
+    @classmethod
+    def _check_for_conflicting_drone_sets(cls, hosts):
+        host_objs = set()
+        for host in hosts:
+            if isinstance(host, Host):
+                host_objs.add(host)
+            elif isinstance(host, Label):
+                host_objs.update(host.host_set.filter(
+                        aclgroup__users=User.current_user()))
+            else:
+                raise TypeError('Unexpected type %s' % host.__class__)
+
+        if host_objs:
+            drone_sets = set(host.drone_set for host in host_objs)
+            if len(drone_sets) != 1:
+                str_hosts = set(str(host) for host in hosts)
+                str_drone_sets = (str(drone_set) for drone_set in drone_sets)
+                raise model_logic.ValidationError(
+                        {'hosts': ('Hosts/Labels %s have specified more than '
+                                   'one set of drones: %s'
+                                   % (', '.join(str_hosts),
+                                      ', '.join(str_drone_sets)))})
+
+
     def queue(self, hosts, atomic_group=None, is_template=False):
         """Enqueue a job on the given hosts."""
         if not hosts:
--- autotest/frontend/afe/models_test.py        2010-03-24 11:22:06.000000000 
-0700
+++ autotest/frontend/afe/models_test.py        2010-03-24 11:22:06.000000000 
-0700
@@ -60,7 +60,8 @@
         host.status = models.Host.Status.RUNNING
         host.save()
 
-        host2 = models.Host.add_object(hostname='othost')
+        drone_set = None or (host.drone_set and host.drone_set.name)
+        host2 = models.Host.add_object(hostname='othost', drone_set=drone_set)
         self.assertEquals(host2.id, host.id)
         self.assertEquals(host2.status, models.Host.Status.RUNNING)
 
@@ -213,5 +214,46 @@
         self.assertEqual(0, models.Job.objects.all().count())
 
 
+class JobTest(unittest.TestCase,
+              frontend_test_utils.FrontendTestMixin):
+    def setUp(self):
+        self._frontend_common_setup()
+
+
+    def tearDown(self):
+        self._frontend_common_teardown()
+
+
+    def _test_conflicting_drone_sets_setup(self):
+        drone_set_1 = models.DroneSet.objects.create(name='drone_set_1')
+        drone_set_2 = models.DroneSet.objects.create(name='drone_set_2')
+
+        for x in (0,1):
+            self.hosts[x].drone_set = drone_set_1
+        for x in (2,3):
+            self.hosts[x].drone_set = drone_set_2
+
+
+    def test_conflicting_drone_sets_no_errors(self):
+        self._test_conflicting_drone_sets_setup()
+        models.Job._check_for_conflicting_drone_sets(self.hosts[0:2])
+        models.Job._check_for_conflicting_drone_sets(self.hosts[2:4])
+        models.Job._check_for_conflicting_drone_sets(self.hosts[4:6])
+
+
+    def test_conflicting_drone_sets_different_sets(self):
+        self._test_conflicting_drone_sets_setup()
+        self.assertRaises(
+                model_logic.ValidationError,
+                models.Job._check_for_conflicting_drone_sets, self.hosts[1:3])
+
+
+    def test_conflicting_drone_sets_one_without(self):
+        self._test_conflicting_drone_sets_setup()
+        self.assertRaises(
+                model_logic.ValidationError,
+                models.Job._check_for_conflicting_drone_sets, self.hosts[3:5])
+
+
 if __name__ == '__main__':
     unittest.main()
--- autotest/frontend/afe/resources.py  2010-03-24 11:22:06.000000000 -0700
+++ autotest/frontend/afe/resources.py  2010-03-24 11:22:06.000000000 -0700
@@ -300,8 +300,10 @@
         cls._check_for_required_fields(input_dict, ('hostname',))
         # include locked here, rather than waiting for update(), to avoid race
         # conditions
-        host = models.Host.add_object(hostname=input_dict['hostname'],
-                                      locked=input_dict.get('locked', False))
+        host = models.Host.add_object(
+                hostname=input_dict['hostname'],
+                locked=input_dict.get('locked', False),
+                drone_set=input_dict.get('drone_set', None))
         return host
 
     def update(self, input_dict):
--- autotest/frontend/afe/resources_test.py     2010-03-24 11:22:06.000000000 
-0700
+++ autotest/frontend/afe/resources_test.py     2010-03-24 11:22:06.000000000 
-0700
@@ -315,7 +315,8 @@
     def test_post(self):
         data = {'hostname': 'newhost',
                 'platform': {'href': self.URI_PREFIX + '/labels/myplatform'},
-                'protection_level': 'Do not verify'}
+                'protection_level': 'Do not verify',
+                'drone_set': 'drone_set'}
         response = self.request('post', 'hosts', data=data)
         self.assertEquals(response, self.URI_PREFIX + '/hosts/newhost')
 
--- autotest/frontend/afe/rpc_interface.py      2010-03-24 11:22:06.000000000 
-0700
+++ autotest/frontend/afe/rpc_interface.py      2010-03-24 11:22:06.000000000 
-0700
@@ -107,9 +107,11 @@
 
 # hosts
 
-def add_host(hostname, status=None, locked=None, protection=None):
+def add_host(hostname, status=None, locked=None,
+             protection=None, drone_set=None):
     return models.Host.add_object(hostname=hostname, status=status,
-                                  locked=locked, protection=protection).id
+                                  locked=locked, protection=protection,
+                                  drone_set=drone_set).id
 
 
 def modify_host(id, **data):
@@ -823,6 +825,10 @@
     result['reboot_before_options'] = model_attributes.RebootBefore.names
     result['reboot_after_options'] = model_attributes.RebootAfter.names
     result['motd'] = rpc_utils.get_motd()
+    result['drone_set_required'] = models.Host.drone_set_required
+    result['drone_sets'] = [drone_set.name for drone_set
+                            in models.DroneSet.objects.all()]
+    result['default_drone_set_name'] = models.Host.default_drone_set_name
 
     result['status_dictionary'] = {"Aborted": "Aborted",
                                    "Verifying": "Verifying Host",
--- autotest/frontend/afe/rpc_interface_unittest.py     2010-03-24 
11:22:06.000000000 -0700
+++ autotest/frontend/afe/rpc_interface_unittest.py     2010-03-24 
11:22:06.000000000 -0700
@@ -32,7 +32,7 @@
                           name=None)
         # violate uniqueness constraint
         self.assertRaises(model_logic.ValidationError, rpc_interface.add_host,
-                          hostname='host1')
+                          hostname='host1', drone_set='drone_set')
 
 
     def test_multiple_platforms(self):
--- autotest/frontend/client/src/autotest/afe/HostSelector.java 2010-03-24 
11:22:06.000000000 -0700
+++ autotest/frontend/client/src/autotest/afe/HostSelector.java 2010-03-24 
11:22:06.000000000 -0700
@@ -1,5 +1,6 @@
 package autotest.afe;
 
+import autotest.common.StaticDataRepository;
 import autotest.common.Utils;
 import autotest.common.table.ArrayDataSource;
 import autotest.common.table.SelectionManager;
@@ -19,6 +20,7 @@
 import com.google.gwt.json.client.JSONNumber;
 import com.google.gwt.json.client.JSONObject;
 import com.google.gwt.json.client.JSONString;
+import com.google.gwt.json.client.JSONValue;
 import com.google.gwt.user.client.ui.HasText;
 import com.google.gwt.user.client.ui.HasValue;
 import com.google.gwt.user.client.ui.Widget;
@@ -52,6 +54,7 @@
     public interface Display {
         public HasText getHostnameField();
         public HasValue<Boolean> getAllowOneTimeHostsField();
+        public HasText getAllowOneTimeHostsLabel();
         public HasClickHandlers getAddByHostnameButton();
         public SimplifiedList getLabelList();
         public HasText getLabelNumberField();
@@ -140,6 +143,15 @@
         display.addTables(availableDecorator, selectedDecorator);
         
         populateLabels(display.getLabelList());
+        
+        JSONValue oneTimeHostDefaultDroneSetName = 
StaticDataRepository.getRepository().getData(
+                "default_drone_set_name");
+        if (oneTimeHostDefaultDroneSetName.isNull() == null) {
+            String origText = display.getAllowOneTimeHostsLabel().getText();
+            String newText = origText + " (default drone set: " +
+                    Utils.jsonToString(oneTimeHostDefaultDroneSetName) + ")";
+            display.getAllowOneTimeHostsLabel().setText(newText);
+        }
     }
     
     @Override
--- autotest/frontend/client/src/autotest/afe/HostSelectorDisplay.java  
2010-03-24 11:22:06.000000000 -0700
+++ autotest/frontend/client/src/autotest/afe/HostSelectorDisplay.java  
2010-03-24 11:22:06.000000000 -0700
@@ -95,6 +95,11 @@
     public HasValue<Boolean> getAllowOneTimeHostsField() {
         return allowOneTimeHostsBox;
     }
+    
+    @Override
+    public HasText getAllowOneTimeHostsLabel() {
+        return allowOneTimeHostsBox;
+    }
 
     @Override
     public HasClickHandlers getAddByLabelButton() {
--- autotest/frontend/frontend_unittest.py      2010-03-24 11:22:06.000000000 
-0700
+++ autotest/frontend/frontend_unittest.py      2010-03-24 11:22:06.000000000 
-0700
@@ -3,18 +3,18 @@
 import unittest, os
 import common
 from autotest_lib.frontend import setup_django_environment
-from autotest_lib.frontend import setup_test_environment
-from autotest_lib.frontend.afe import test, readonly_connection
+from autotest_lib.frontend.afe import frontend_test_utils, test
 
 _APP_DIR = os.path.join(os.path.dirname(__file__), 'afe')
 
-class FrontendTest(unittest.TestCase):
+class FrontendTest(unittest.TestCase,
+                   frontend_test_utils.FrontendTestMixin):
     def setUp(self):
-        setup_test_environment.set_up()
+        self._frontend_common_setup(fill_data=False)
 
 
     def tearDown(self):
-        setup_test_environment.tear_down()
+        self._frontend_common_teardown()
 
 
     def test_all(self):
--- /dev/null   2009-12-17 12:29:38.000000000 -0800
+++ autotest/frontend/migrations/052_drone_management.py        2010-03-24 
11:22:06.000000000 -0700
@@ -0,0 +1,60 @@
+UP_SQL = """
+CREATE TABLE afe_drones (
+  id INT AUTO_INCREMENT NOT NULL PRIMARY KEY,
+  hostname VARCHAR(255) NOT NULL
+) ENGINE=InnoDB;
+
+ALTER TABLE afe_drones
+ADD CONSTRAINT afe_drones_unique
+UNIQUE KEY (hostname);
+
+
+CREATE TABLE afe_drone_sets (
+  id INT AUTO_INCREMENT NOT NULL PRIMARY KEY,
+  name VARCHAR(255) NOT NULL
+) ENGINE=InnoDB;
+
+ALTER TABLE afe_drone_sets
+ADD CONSTRAINT afe_drone_sets_unique
+UNIQUE KEY (name);
+
+
+CREATE TABLE afe_drone_sets_drones (
+  id INT AUTO_INCREMENT NOT NULL PRIMARY KEY,
+  droneset_id INT NOT NULL,
+  drone_id INT NOT NULL
+) ENGINE=InnoDB;
+
+ALTER TABLE afe_drone_sets_drones
+ADD CONSTRAINT afe_drone_sets_drones_droneset_ibfk
+FOREIGN KEY (droneset_id) REFERENCES afe_drone_sets (id);
+
+ALTER TABLE afe_drone_sets_drones
+ADD CONSTRAINT afe_drone_sets_drones_drone_ibfk
+FOREIGN KEY (drone_id) REFERENCES afe_drones (id);
+
+ALTER TABLE afe_drone_sets_drones
+ADD CONSTRAINT afe_drone_sets_drones_unique
+UNIQUE KEY (droneset_id, drone_id);
+
+
+ALTER TABLE afe_hosts
+ADD COLUMN drone_set_id INT;
+
+ALTER TABLE afe_hosts
+ADD CONSTRAINT afe_hosts_drone_set_ibfk
+FOREIGN KEY (drone_set_id) REFERENCES afe_drone_sets (id);
+"""
+
+
+DOWN_SQL = """
+ALTER TABLE afe_hosts
+DROP FOREIGN KEY afe_hosts_drone_set_ibfk;
+
+ALTER TABLE afe_hosts
+DROP COLUMN drone_set_id;
+
+DROP TABLE IF EXISTS afe_drone_sets_drones;
+DROP TABLE IF EXISTS afe_drone_sets;
+DROP TABLE IF EXISTS afe_drones;
+"""
--- autotest/global_config.ini  2010-03-24 10:07:59.000000000 -0700
+++ autotest/global_config.ini  2010-03-24 11:22:06.000000000 -0700
@@ -75,6 +75,11 @@
 gc_stats_interval_mins: 360
 # set nonzero to enable periodic reverification of all dead hosts
 reverify_period_minutes: 0
+drone_set_required: False
+# if drone_set_required is set to True, you MUST specify a
+# default_drone_set_name, which will be used for one-time hosts and hostless
+# jobs
+default_drone_set_name:
 
 [HOSTS]
 wait_up_processes:
--- autotest/scheduler/drone_manager.py 2010-03-24 11:14:47.000000000 -0700
+++ autotest/scheduler/drone_manager.py 2010-03-24 11:22:06.000000000 -0700
@@ -412,14 +412,19 @@
         return sum(drone.active_processes for drone in self.get_drones())
 
 
-    def max_runnable_processes(self, username):
+    def max_runnable_processes(self, username, drone_hostnames_allowed):
         """
         Return the maximum number of processes that can be run (in a single
         execution) given the current load on drones.
         @param username: login of user to run a process.  may be None.
+        @param drone_hostnames_allowed: list of drones that can be used. May be
+                                        None
         """
         usable_drone_wrappers = [wrapper for wrapper in self._drone_queue
-                                 if wrapper.drone.usable_by(username)]
+                                 if wrapper.drone.usable_by(username) and
+                                 (drone_hostnames_allowed is None or
+                                          wrapper.drone.hostname in
+                                                  drone_hostnames_allowed)]
         if not usable_drone_wrappers:
             # all drones disabled or inaccessible
             return 0
@@ -437,7 +442,8 @@
         return drone_to_use
 
 
-    def _choose_drone_for_execution(self, num_processes, username):
+    def _choose_drone_for_execution(self, num_processes, username,
+                                    drone_hostnames_allowed):
         # cycle through drones is order of increasing used capacity until
         # we find one that can handle these processes
         checked_drones = []
@@ -448,12 +454,19 @@
             checked_drones.append(drone)
             if not drone.usable_by(username):
                 continue
+
+            drone_allowed = (drone_hostnames_allowed is None or
+                             drone.hostname in drone_hostnames_allowed)
+            if not drone_allowed:
+                continue
+
             usable_drones.append(drone)
+
             if drone.active_processes + num_processes <= drone.max_processes:
                 drone_to_use = drone
                 break
 
-        if not drone_to_use:
+        if not drone_to_use and usable_drones:
             drone_summary = ','.join('%s %s/%s' % (drone.hostname,
                                                    drone.active_processes,
                                                    drone.max_processes)
@@ -478,7 +491,7 @@
 
     def execute_command(self, command, working_directory, pidfile_name,
                         num_processes, log_file=None, paired_with_pidfile=None,
-                        username=None):
+                        username=None, drone_hostnames_allowed=None):
         """
         Execute the given command, taken as an argv list.
 
@@ -496,6 +509,9 @@
                 same drone as the previous process.
         @param username (optional): login of the user responsible for this
                 process.
+        @param drone_hostnames_allowed (optional): hostnames of the drones that
+                                                   this command is allowed to
+                                                   execute on
         """
         abs_working_directory = self.absolute_path(working_directory)
         if not log_file:
@@ -508,7 +524,12 @@
         if paired_with_pidfile:
             drone = self._get_drone_for_pidfile_id(paired_with_pidfile)
         else:
-            drone = self._choose_drone_for_execution(num_processes, username)
+            drone = self._choose_drone_for_execution(num_processes, username,
+                                                     drone_hostnames_allowed)
+        if not drone:
+            raise DroneManagerError('command %s failed; no drones available'
+                                    % command)
+
         logging.info("command = %s" % command)
         logging.info('log file = %s:%s' % (drone.hostname, log_file))
         self._write_attached_files(working_directory, drone)
--- autotest/scheduler/drone_manager_unittest.py        2010-03-24 
11:14:47.000000000 -0700
+++ autotest/scheduler/drone_manager_unittest.py        2010-03-24 
11:22:06.000000000 -0700
@@ -12,6 +12,7 @@
                  allowed_users=None):
         super(MockDrone, self).__init__()
         self.name = name
+        self.hostname = name
         self.active_processes = active_processes
         self.max_processes = max_processes
         self.allowed_users = allowed_users
@@ -100,7 +101,7 @@
                                                   max_processes))
 
         return self.manager._choose_drone_for_execution(requested_processes,
-                                                        self._USERNAME)
+                                                        self._USERNAME, None)
 
 
     def test_choose_drone_for_execution(self):
@@ -136,9 +137,10 @@
                                               allowed_users=[self._USERNAME]))
 
         self.assertEquals(2,
-                          self.manager.max_runnable_processes(self._USERNAME))
+                          self.manager.max_runnable_processes(self._USERNAME,
+                                                              None))
         drone = self.manager._choose_drone_for_execution(
-                1, username=self._USERNAME)
+                1, username=self._USERNAME, drone_hostnames_allowed=None)
         self.assertEquals(drone.name, 2)
 
 
@@ -152,12 +154,69 @@
                                               allowed_users=[self._USERNAME]))
 
         self.assertEquals(0,
-                          self.manager.max_runnable_processes(self._USERNAME))
+                          self.manager.max_runnable_processes(self._USERNAME,
+                                                              None))
         drone = self.manager._choose_drone_for_execution(
-                1, username=self._USERNAME)
+                1, username=self._USERNAME, drone_hostnames_allowed=None)
         self.assertEquals(drone.name, 2)
 
 
+    def _setup_test_drone_restrictions(self, active_processes=0):
+        self.manager._enqueue_drone(MockDrone(
+                1, active_processes=active_processes, max_processes=10))
+        self.manager._enqueue_drone(MockDrone(
+                2, active_processes=active_processes, max_processes=5))
+        self.manager._enqueue_drone(MockDrone(
+                3, active_processes=active_processes, max_processes=2))
+
+
+    def test_drone_restrictions_allow_any(self):
+        self._setup_test_drone_restrictions()
+        self.assertEquals(10,
+                          self.manager.max_runnable_processes(self._USERNAME,
+                                                              None))
+        drone = self.manager._choose_drone_for_execution(
+                1, username=self._USERNAME, drone_hostnames_allowed=None)
+        self.assertEqual(drone.name, 1)
+
+
+    def test_drone_restrictions_under_capacity(self):
+        self._setup_test_drone_restrictions()
+        drone_hostnames_allowed = (2, 3)
+        self.assertEquals(
+                5, self.manager.max_runnable_processes(self._USERNAME,
+                                                       
drone_hostnames_allowed))
+        drone = self.manager._choose_drone_for_execution(
+                1, username=self._USERNAME,
+                drone_hostnames_allowed=drone_hostnames_allowed)
+
+        self.assertEqual(drone.name, 2)
+
+
+    def test_drone_restrictions_over_capacity(self):
+        self._setup_test_drone_restrictions(active_processes=6)
+        drone_hostnames_allowed = (2, 3)
+        self.assertEquals(
+                0, self.manager.max_runnable_processes(self._USERNAME,
+                                                       
drone_hostnames_allowed))
+        drone = self.manager._choose_drone_for_execution(
+                7, username=self._USERNAME,
+                drone_hostnames_allowed=drone_hostnames_allowed)
+        self.assertEqual(drone.name, 2)
+
+
+    def test_drone_restrictions_allow_none(self):
+        self._setup_test_drone_restrictions()
+        drone_hostnames_allowed = ()
+        self.assertEquals(
+                0, self.manager.max_runnable_processes(self._USERNAME,
+                                                       
drone_hostnames_allowed))
+        drone = self.manager._choose_drone_for_execution(
+                1, username=self._USERNAME,
+                drone_hostnames_allowed=drone_hostnames_allowed)
+        self.assertEqual(drone, None)
+
+
     def test_initialize(self):
         results_hostname = 'results_repo'
         results_install_dir = '/results/install'
--- autotest/scheduler/monitor_db.py    2010-03-24 11:22:07.000000000 -0700
+++ autotest/scheduler/monitor_db.py    2010-03-24 11:22:07.000000000 -0700
@@ -74,7 +74,13 @@
         'get_metahost_schedulers', lambda : ())
 
 
+def _sanity_check():
+    """Make sure the configs are consistent before starting the scheduler"""
+    models.Host.verify_default_drone_set_exists()
+
+
 def main():
+    _sanity_check()
     try:
         try:
             main_without_exception_handling()
@@ -1138,7 +1144,8 @@
             return False
         # total process throttling
         max_runnable_processes = _drone_manager.max_runnable_processes(
-                agent.task.owner_username)
+                agent.task.owner_username,
+                agent.task.get_drone_hostnames_allowed())
         if agent.task.num_processes > max_runnable_processes:
             return False
         # if a single agent exceeds the per-cycle throttling, still allow it to
@@ -1244,7 +1251,7 @@
 
     def run(self, command, working_directory, num_processes, nice_level=None,
             log_file=None, pidfile_name=None, paired_with_pidfile=None,
-            username=None):
+            username=None, drone_hostnames_allowed=None):
         assert command is not None
         if nice_level is not None:
             command = ['nice', '-n', str(nice_level)] + command
@@ -1252,7 +1259,8 @@
         self.pidfile_id = _drone_manager.execute_command(
             command, working_directory, pidfile_name=pidfile_name,
             num_processes=num_processes, log_file=log_file,
-            paired_with_pidfile=paired_with_pidfile, username=username)
+            paired_with_pidfile=paired_with_pidfile, username=username,
+            drone_hostnames_allowed=drone_hostnames_allowed)
 
 
     def attach_to_existing_process(self, execution_path,
@@ -1665,7 +1673,55 @@
                 nice_level=AUTOSERV_NICE_LEVEL, log_file=self._log_file(),
                 pidfile_name=self._pidfile_name(),
                 paired_with_pidfile=self._paired_with_monitor().pidfile_id,
-                username=self.owner_username)
+                username=self.owner_username,
+                drone_hostnames_allowed=self.get_drone_hostnames_allowed())
+
+
+    def get_drone_hostnames_allowed(self):
+        if self.host_ids:
+            hosts = models.Host.objects.filter(id__in=self.host_ids)
+        else:
+            # Figure it out from queue entries if self.host_ids is not
+            # explicitly set
+            hosts = models.Host.objects.filter(
+                    hostqueueentry__id__in=self.queue_entry_ids)
+
+        if hosts:
+            return self._get_drone_hostnames(hosts)
+
+        # No hosts associated (i.e., hostless job); get the default drone set
+        drone_set = models.Host.default_drone_set()
+        if drone_set:
+            return set(
+                    drone_set.drones.all().values_list('hostname', flat=True))
+        return None
+
+
+    def _get_drone_hostnames(self, hosts):
+        """
+        Finds the drones associated with a list of hosts.
+
+        Raises a SchedulerError if more than one drone set is specified on the
+        hosts.
+        """
+        drone_sets = hosts.values_list('drone_set', flat=True).distinct()
+        if len(drone_sets) != 1:
+            str_hosts = (str(host) for host in hosts)
+            str_drone_sets = (
+                    str(drone_set) for drone_set in
+                    models.DroneSet.objects.filter(
+                            id__in=drone_sets).values_list('name', flat=True))
+            raise SchedulerError('Hosts %s have specified more than one set of 
'
+                                 'drones: %s'
+                                 % (','.join(str_hosts),
+                                    ','.join(str_drone_sets)))
+
+        drone_set_id = drone_sets[0]
+        if drone_set_id:
+            drone_set = models.DroneSet.objects.get(id=drone_set_id)
+            return set(
+                    drone_set.drones.all().values_list('hostname', flat=True))
+        return None
 
 
     def register_necessary_pidfiles(self):
--- autotest/scheduler/monitor_db_functional_test.py    2010-03-24 
11:22:07.000000000 -0700
+++ autotest/scheduler/monitor_db_functional_test.py    2010-03-24 
11:22:07.000000000 -0700
@@ -185,7 +185,7 @@
                    for pidfile_id in self.nonfinished_pidfile_ids())
 
 
-    def max_runnable_processes(self, username):
+    def max_runnable_processes(self, username, drone_hostnames_allowed):
         return self.process_capacity - self.total_running_processes()
 
 
@@ -232,7 +232,7 @@
 
     def execute_command(self, command, working_directory, pidfile_name,
                         num_processes, log_file=None, paired_with_pidfile=None,
-                        username=None):
+                        username=None, drone_hostnames_allowed=None):
         logging.debug('Executing %s in %s', command, working_directory)
         pidfile_id = self._DummyPidfileId(working_directory, pidfile_name)
         if pidfile_id.key() in self._pidfile_index:
--- autotest/scheduler/monitor_db_unittest.py   2010-03-24 11:22:07.000000000 
-0700
+++ autotest/scheduler/monitor_db_unittest.py   2010-03-24 11:22:07.000000000 
-0700
@@ -20,6 +20,9 @@
     num_processes = 1
     owner_username = 'my_user'
 
+    def get_drone_hostnames_allowed(self):
+        return None
+
 
 class DummyAgent(object):
     started = False
@@ -730,7 +733,7 @@
         scheduler_config.config.max_processes_started_per_cycle = (
             self._MAX_STARTED)
 
-        def fake_max_runnable_processes(fake_self, username):
+        def fake_max_runnable_processes(fake_self, username, 
drone_hostnames_allowed):
             running = sum(agent.task.num_processes
                           for agent in self._agents
                           if agent.started and not agent.is_done())
@@ -1358,5 +1361,56 @@
         self.assertEqual(expected_command_line, command_line)
 
 
+class AgentTaskTest(unittest.TestCase,
+                    frontend_test_utils.FrontendTestMixin):
+    def setUp(self):
+        self._frontend_common_setup()
+
+
+    def tearDown(self):
+        self._frontend_common_teardown()
+
+
+    def _setup_drones(self):
+        drones = []
+        for x in xrange(4):
+            drones.append(models.Drone.objects.create(hostname=str(x)))
+
+        drone_set_1 = models.DroneSet.objects.create(name='1')
+        drone_set_1.drones.add(*drones[0:2])
+        drone_set_2 = models.DroneSet.objects.create(name='2')
+        drone_set_2.drones.add(*drones[2:4])
+
+        for host in self.hosts[0:2]:
+            host.drone_set = drone_set_1
+            host.save()
+        self.hosts[2].drone_set = drone_set_2
+        self.hosts[2].save()
+
+        return drones
+
+
+    def test_get_drone_hostnames_allowed_no_associated_sets(self):
+        self._setup_drones()
+        task = monitor_db.AgentTask()
+        task.host_ids = (4,5)
+        self.assertEqual(set(), task.get_drone_hostnames_allowed())
+
+
+    def test_get_drone_hostnames_allowed_success(self):
+        drones = self._setup_drones()
+        task = monitor_db.AgentTask()
+        task.host_ids = (1,2)
+        self.assertEqual(set(('0','1')), task.get_drone_hostnames_allowed())
+
+
+    def test_get_drone_hostnames_allowed_different_drones(self):
+        self._setup_drones()
+        task = monitor_db.AgentTask()
+        task.host_ids = (2,3)
+        self.assertRaises(monitor_db.SchedulerError,
+                          task.get_drone_hostnames_allowed)
+
+
 if __name__ == '__main__':
     unittest.main()
--- autotest/scheduler/scheduler_models.py      2010-03-24 11:22:07.000000000 
-0700
+++ autotest/scheduler/scheduler_models.py      2010-03-24 11:22:07.000000000 
-0700
@@ -357,7 +357,8 @@
 class Host(DBObject):
     _table_name = 'afe_hosts'
     _fields = ('id', 'hostname', 'locked', 'synch_id', 'status',
-               'invalid', 'protection', 'locked_by_id', 'lock_time', 'dirty')
+               'invalid', 'protection', 'locked_by_id', 'lock_time', 'dirty',
+               'drone_set_id')
 
 
     def set_status(self,status):
_______________________________________________
Autotest mailing list
[email protected]
http://test.kernel.org/cgi-bin/mailman/listinfo/autotest

Reply via email to