Add ability to associate drone sets with jobs. This restricts a job to
running on a specified set of drones.

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

--- autotest/frontend/afe/admin.py      2010-04-16 17:35:33.000000000 -0700
+++ autotest/frontend/afe/admin.py      2010-04-16 17:35:33.000000000 -0700
@@ -136,6 +136,14 @@
 admin.site.register(models.AclGroup, AclGroupAdmin)
 
 
+class DroneSetAdmin(SiteAdmin):
+    filter_horizontal = ('drones',)
+
+admin.site.register(models.DroneSet, DroneSetAdmin)
+
+admin.site.register(models.Drone)
+
+
 if settings.FULL_ADMIN:
     class JobAdmin(SiteAdmin):
         list_display = ('id', 'owner', 'name', 'control_type')
--- autotest/frontend/afe/doctests/001_rpc_test.txt     2010-04-16 
17:35:33.000000000 -0700
+++ autotest/frontend/afe/doctests/001_rpc_test.txt     2010-04-16 
17:35:33.000000000 -0700
@@ -30,6 +30,10 @@
 >>> from autotest_lib.client.common_lib import logging_manager
 >>> logging_manager.logger.setLevel(100)
 
+>>> drone_set = models.DroneSet.default_drone_set_name()
+>>> if drone_set:
+...     _ = models.DroneSet.objects.create(name=drone_set)
+
 # basic interface test
 ######################
 
@@ -184,6 +188,7 @@
 ...           'access_level': 1,
 ...           'reboot_before': 'If dirty',
 ...           'reboot_after': 'Always',
+...           'drone_set': None,
 ...           'show_experimental': False}]
 True
 >>> rpc_interface.delete_user('showard')
@@ -538,7 +543,8 @@
 ...         'email_list': '',
 ...         'reboot_before': 'If dirty',
 ...         'reboot_after': 'Always',
-...         'parse_failed_repair': True}
+...         'parse_failed_repair': True,
+...         'drone_set': drone_set}
 True
 
 # get_host_queue_entries returns a lot of data, so let's only check a couple
--- autotest/frontend/afe/frontend_test_utils.py        2010-04-16 
17:35:33.000000000 -0700
+++ autotest/frontend/afe/frontend_test_utils.py        2010-04-16 
17:35:33.000000000 -0700
@@ -8,6 +8,10 @@
 class FrontendTestMixin(object):
     def _fill_in_test_data(self):
         """Populate the test database with some hosts and labels."""
+        if models.DroneSet.drone_sets_enabled():
+            models.DroneSet.objects.create(
+                    name=models.DroneSet.default_drone_set_name())
+
         acl_group = models.AclGroup.objects.create(name='my_acl')
         acl_group.users.add(models.User.current_user())
 
@@ -72,7 +76,8 @@
 
 
     def _create_job(self, hosts=[], metahosts=[], priority=0, active=False,
-                    synchronous=False, atomic_group=None, hostless=False):
+                    synchronous=False, atomic_group=None, hostless=False,
+                    drone_set=None):
         """
         Create a job row in the test database.
 
@@ -93,6 +98,10 @@
 
         @returns A Django frontend.afe.models.Job instance.
         """
+        if not drone_set:
+            drone_set = (models.DroneSet.default_drone_set_name()
+                         and models.DroneSet.get_default())
+
         assert not (atomic_group and active)  # TODO(gps): support this
         synch_count = synchronous and 2 or 1
         created_on = datetime.datetime(2008, 1, 1)
@@ -102,7 +111,8 @@
         job = models.Job.objects.create(
             name='test', owner='autotest_system', priority=priority,
             synch_count=synch_count, created_on=created_on,
-            reboot_before=model_attributes.RebootBefore.NEVER)
+            reboot_before=model_attributes.RebootBefore.NEVER,
+            drone_set=drone_set)
         for host_id in hosts:
             models.HostQueueEntry.objects.create(job=job, host_id=host_id,
                                                  status=status,
@@ -126,11 +136,12 @@
 
 
     def _create_job_simple(self, hosts, use_metahost=False,
-                          priority=0, active=False):
+                          priority=0, active=False, drone_set=None):
         """An alternative interface to _create_job"""
         args = {'hosts' : [], 'metahosts' : []}
         if use_metahost:
             args['metahosts'] = hosts
         else:
             args['hosts'] = hosts
-        return self._create_job(priority=priority, active=active, **args)
+        return self._create_job(priority=priority, active=active,
+                                drone_set=drone_set, **args)
--- autotest/frontend/afe/models.py     2010-04-15 14:59:52.000000000 -0700
+++ autotest/frontend/afe/models.py     2010-04-16 17:35:33.000000000 -0700
@@ -116,6 +116,119 @@
         return unicode(self.name)
 
 
+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 job is allowed
+    to run on.
+
+    name: the drone set's name
+    drones: the drones that are part of the set
+    """
+    DRONE_SETS_ENABLED = global_config.global_config.get_config_value(
+            'SCHEDULER', 'drone_sets_enabled', type=bool, default=False)
+    DEFAULT_DRONE_SET_NAME = global_config.global_config.get_config_value(
+            'SCHEDULER', 'default_drone_set_name', default=None)
+
+    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()
+
+
+    @classmethod
+    def drone_sets_enabled(cls):
+        return cls.DRONE_SETS_ENABLED
+
+
+    @classmethod
+    def default_drone_set_name(cls):
+        return cls.DEFAULT_DRONE_SET_NAME
+
+
+    @classmethod
+    def get_default(cls):
+        return cls.smart_get(cls.DEFAULT_DRONE_SET_NAME)
+
+
+    @classmethod
+    def resolve_name(cls, drone_set_name):
+        """
+        Returns the name of one of these, if not None, in order of preference:
+        1) the drone set given,
+        2) the current user's default drone set, or
+        3) the global default drone set
+
+        or returns None if drone sets are disabled
+        """
+        if not cls.drone_sets_enabled():
+            return None
+
+        user = User.current_user()
+        user_drone_set_name = user.drone_set and user.drone_set.name
+
+        return drone_set_name or user_drone_set_name or cls.get_default().name
+
+
+    def get_drone_hostnames(self):
+        """
+        Gets the hostnames of all drones in this drone set
+        """
+        return set(self.drones.all().values_list('hostname', flat=True))
+
+
+    class Meta:
+        db_table = 'afe_drone_sets'
+
+    def __unicode__(self):
+        return unicode(self.name)
+
+
 class User(dbmodels.Model, model_logic.ModelExtensions):
     """\
     Required:
@@ -140,6 +253,7 @@
     reboot_after = dbmodels.SmallIntegerField(
         choices=model_attributes.RebootAfter.choices(), blank=True,
         default=DEFAULT_REBOOT_AFTER)
+    drone_set = dbmodels.ForeignKey(DroneSet, null=True, blank=True)
     show_experimental = dbmodels.BooleanField(default=False)
 
     name_field = 'login'
@@ -644,6 +758,7 @@
     reboot_after: Never, If all tests passed, or Always
     parse_failed_repair: if True, a failed repair launched by this job will 
have
     its results parsed as part of the job.
+    drone_set: The set of drones to run this job on
     """
     DEFAULT_TIMEOUT = global_config.global_config.get_config_value(
         'AUTOTEST_WEB', 'job_timeout_default', default=240)
@@ -682,6 +797,7 @@
     parse_failed_repair = dbmodels.BooleanField(
         default=DEFAULT_PARSE_FAILED_REPAIR)
     max_runtime_hrs = dbmodels.IntegerField(default=DEFAULT_MAX_RUNTIME_HRS)
+    drone_set = dbmodels.ForeignKey(DroneSet, null=True, blank=True)
 
 
     # custom manager
@@ -706,6 +822,8 @@
         if options.get('reboot_after') is None:
             options['reboot_after'] = user.get_reboot_after_display()
 
+        drone_set = DroneSet.resolve_name(options.get('drone_set'))
+
         job = cls.add_object(
             owner=owner,
             name=options['name'],
@@ -720,7 +838,8 @@
             reboot_before=options.get('reboot_before'),
             reboot_after=options.get('reboot_after'),
             parse_failed_repair=options.get('parse_failed_repair'),
-            created_on=datetime.now())
+            created_on=datetime.now(),
+            drone_set=drone_set)
 
         job.dependency_labels = options['dependencies']
 
@@ -995,7 +1114,7 @@
     host = dbmodels.ForeignKey(Host, blank=False, null=False)
     task = dbmodels.CharField(max_length=64, choices=Task.choices(),
                               blank=False, null=False)
-    requested_by = dbmodels.ForeignKey(User, blank=True, null=True)
+    requested_by = dbmodels.ForeignKey(User)
     time_requested = dbmodels.DateTimeField(auto_now_add=True, blank=False,
                                             null=False)
     is_active = dbmodels.BooleanField(default=False, blank=False, null=False)
--- autotest/frontend/afe/models_test.py        2010-04-16 17:35:33.000000000 
-0700
+++ autotest/frontend/afe/models_test.py        2010-04-16 17:35:33.000000000 
-0700
@@ -77,7 +77,8 @@
 
     def _create_task(self):
         return models.SpecialTask.objects.create(
-                host=self.hosts[0], task=models.SpecialTask.Task.VERIFY)
+                host=self.hosts[0], task=models.SpecialTask.Task.VERIFY,
+                requested_by=models.User.current_user())
 
 
     def test_execution_path(self):
--- autotest/frontend/afe/resources.py  2010-04-16 16:51:49.000000000 -0700
+++ autotest/frontend/afe/resources.py  2010-04-16 17:35:33.000000000 -0700
@@ -643,9 +643,11 @@
         rep = super(Job, self).full_representation()
         queue_entries = QueueEntryCollection(self._request)
         queue_entries.set_query_parameters(job=self.instance.id)
+        drone_set = self.instance.drone_set and self.instance.drone_set.name
         rep.update({'email_list': self.instance.email_list,
                     'parse_failed_repair':
                         bool(self.instance.parse_failed_repair),
+                    'drone_set': drone_set,
                     'execution_info':
                         ExecutionInfo.execution_info_from_job(self.instance),
                     'queue_entries': queue_entries.link(),
@@ -686,6 +688,7 @@
                 reboot_before=execution_info.get('cleanup_before_job'),
                 reboot_after=execution_info.get('cleanup_after_job'),
                 parse_failed_repair=input_dict.get('parse_failed_repair', 
None),
+                drone_set=input_dict.get('drone_set', None),
                 keyvals=input_dict.get('keyvals', None))
 
         host_objects, metahost_label_objects, atomic_group = [], [], None
--- autotest/frontend/afe/resources_test.py     2010-04-16 16:51:49.000000000 
-0700
+++ autotest/frontend/afe/resources_test.py     2010-04-16 17:35:33.000000000 
-0700
@@ -350,6 +350,7 @@
                 'execution_info': {'control_file': self.CONTROL_FILE_CONTENTS,
                                    'is_server': True},
                 'owner': owner,
+                'drone_set': models.DroneSet.default_drone_set_name(),
                 'queue_entries':
                 [{'host': {'href': self.URI_PREFIX + '/hosts/host1'}},
                  {'host': {'href': self.URI_PREFIX + '/hosts/host2'}}]}
--- autotest/frontend/afe/rpc_interface.py      2010-04-16 17:35:33.000000000 
-0700
+++ autotest/frontend/afe/rpc_interface.py      2010-04-16 17:35:33.000000000 
-0700
@@ -403,7 +403,7 @@
                timeout=None, max_runtime_hrs=None, run_verify=True,
                email_list='', dependencies=(), reboot_before=None,
                reboot_after=None, parse_failed_repair=None, hostless=False,
-               keyvals=None):
+               keyvals=None, drone_set=None):
     """\
     Create and enqueue a job.
 
@@ -432,6 +432,7 @@
     one host will be chosen from that label to run the job on.
     @param one_time_hosts List of hosts not in the database to run the job on.
     @param atomic_group_name The name of an atomic group to schedule the job 
on.
+    @param drone_set The name of the drone set to run this test on.
 
 
     @returns The created Job id number.
@@ -529,7 +530,8 @@
                    reboot_before=reboot_before,
                    reboot_after=reboot_after,
                    parse_failed_repair=parse_failed_repair,
-                   keyvals=keyvals)
+                   keyvals=keyvals,
+                   drone_set=drone_set)
     return rpc_utils.create_new_job(owner=owner,
                                     options=options,
                                     host_objects=host_objects,
@@ -656,6 +658,7 @@
     else:
         info['atomic_group_name'] = None
     info['hostless'] = job_info['hostless']
+    info['drone_set'] = job.drone_set and job.drone_set.name
 
     return rpc_utils.prepare_for_serialization(info)
 
@@ -802,6 +805,11 @@
     """
 
     job_fields = models.Job.get_field_dict()
+    default_drone_set_name = models.DroneSet.default_drone_set_name()
+    drone_sets = ([default_drone_set_name] +
+                  sorted(drone_set.name for drone_set in
+                         models.DroneSet.objects.exclude(
+                                 name=default_drone_set_name)))
 
     result = {}
     result['priorities'] = models.Job.Priority.choices()
@@ -824,6 +832,8 @@
     result['reboot_before_options'] = model_attributes.RebootBefore.names
     result['reboot_after_options'] = model_attributes.RebootAfter.names
     result['motd'] = rpc_utils.get_motd()
+    result['drone_sets_enabled'] = models.DroneSet.drone_sets_enabled()
+    result['drone_sets'] = drone_sets
 
     result['status_dictionary'] = {"Aborted": "Aborted",
                                    "Verifying": "Verifying Host",
--- autotest/frontend/afe/rpc_interface_unittest.py     2010-04-16 
17:35:33.000000000 -0700
+++ autotest/frontend/afe/rpc_interface_unittest.py     2010-04-16 
17:35:33.000000000 -0700
@@ -212,13 +212,14 @@
         self.task1 = models.SpecialTask.objects.create(
                 host=host, task=models.SpecialTask.Task.VERIFY,
                 time_started=datetime.datetime(2009, 1, 1), # ran before job 1
-                is_complete=True)
+                is_complete=True, requested_by=models.User.current_user())
         self.task2 = models.SpecialTask.objects.create(
                 host=host, task=models.SpecialTask.Task.VERIFY,
                 queue_entry=entry2, # ran with job 2
-                is_active=True)
+                is_active=True, requested_by=models.User.current_user())
         self.task3 = models.SpecialTask.objects.create(
-                host=host, task=models.SpecialTask.Task.VERIFY) # not yet run
+                host=host, task=models.SpecialTask.Task.VERIFY,
+                requested_by=models.User.current_user()) # not yet run
 
 
     def test_get_special_tasks(self):
--- autotest/frontend/client/src/autotest/afe/AfeUtils.java     2010-04-16 
17:35:33.000000000 -0700
+++ autotest/frontend/client/src/autotest/afe/AfeUtils.java     2010-04-16 
17:35:33.000000000 -0700
@@ -16,6 +16,9 @@
 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.DOM;
+import com.google.gwt.user.client.Element;
+import com.google.gwt.user.client.ui.ListBox;
 
 import java.util.ArrayList;
 import java.util.Collection;
@@ -312,6 +315,28 @@
             chooser.addChoice(Utils.jsonToString(jsonOption));
         }
     }
+    
+    public static void popualateListBox(ListBox box, String staticDataKey) {
+        JSONArray options = staticData.getData(staticDataKey).isArray();
+        for (JSONString jsonOption : new JSONArrayList<JSONString>(options)) {
+            box.addItem(Utils.jsonToString(jsonOption));
+        }
+    }
+    
+    public static void setSelectedItem(ListBox box, String item) {
+        box.setSelectedIndex(0);
+        for (int i = 0; i < box.getItemCount(); i++) {
+            if (box.getItemText(i).equals(item)) {
+                box.setSelectedIndex(i);
+                break;
+            }
+        }
+    }
+    
+    public static void removeElement(String id) {
+        Element element = DOM.getElementById(id);
+        element.getParentElement().removeChild(element);
+    }
 
     public static int parsePositiveIntegerInput(String input, String 
fieldName) {
         final int parsedInt;
--- autotest/frontend/client/src/autotest/afe/CreateJobView.java        
2010-04-16 17:35:33.000000000 -0700
+++ autotest/frontend/client/src/autotest/afe/CreateJobView.java        
2010-04-16 17:35:33.000000000 -0700
@@ -179,6 +179,7 @@
         new CheckBoxPanel<CheckBox>(TEST_COLUMNS);
     private CheckBox runNonProfiledIteration =
         new CheckBox("Run each test without profilers first");
+    private ListBox droneSet = new ListBox();
     protected TextArea controlFile = new TextArea();
     protected DisclosurePanel controlFilePanel = new DisclosurePanel();
     protected ControlTypeSelect controlTypeSelect;
@@ -238,6 +239,9 @@
         if (hostless.getValue()) {
             hostSelector.setEnabled(false);
         }
+        if (cloneObject.get("drone_set").isNull() == null) {
+            AfeUtils.setSelectedItem(droneSet, 
Utils.jsonToString(cloneObject.get("drone_set")));
+        }
 
         controlTypeSelect.setControlType(
                 jobObject.get("control_type").isString().stringValue());
@@ -455,8 +459,9 @@
     @Override
     public void initialize() {
         super.initialize();
+        
         populatePriorities(staticData.getData("priorities").isArray());
-
+        
         BlurHandler kernelBlurHandler = new BlurHandler() {
             public void onBlur(BlurEvent event) {
                 generateControlFile(false);
@@ -617,6 +622,13 @@
         addWidget(createTemplateJobButton, "create_template_job");
         addWidget(resetButton, "create_reset");
         
+        if 
(staticData.getData("drone_sets_enabled").isBoolean().booleanValue()) {
+            AfeUtils.popualateListBox(droneSet, "drone_sets");
+            addWidget(droneSet, "create_drone_set");
+        } else {
+            AfeUtils.removeElement("create_drone_set_wrapper");
+        }
+        
         testSelector.setListener(this);
     }
 
@@ -693,6 +705,11 @@
                 args.put("parse_failed_repair",
                          
JSONBoolean.getInstance(parseFailedRepair.getValue()));
                 args.put("hostless", 
JSONBoolean.getInstance(hostless.getValue()));
+                
+                if 
(staticData.getData("drone_sets_enabled").isBoolean().booleanValue()) {
+                    args.put("drone_set",
+                            new 
JSONString(droneSet.getItemText(droneSet.getSelectedIndex())));
+                }
 
                 HostSelector.HostSelection hosts = 
hostSelector.getSelectedHosts();
                 args.put("hosts", Utils.stringsToJSON(hosts.hosts));
@@ -773,10 +790,17 @@
         String defaultOption = Utils.jsonToString(user.get(name));
         chooser.setDefaultChoice(defaultOption);
     }
-
+    
+    private void selectPreferredDroneSet() {
+        JSONObject user = staticData.getData("current_user").isObject();
+        String preference = Utils.jsonToString(user.get("drone_set"));
+        AfeUtils.setSelectedItem(droneSet, preference);
+    }
+    
     public void onPreferencesChanged() {
         setRebootSelectorDefault(rebootBefore, "reboot_before");
         setRebootSelectorDefault(rebootAfter, "reboot_after");
+        selectPreferredDroneSet();
         testSelector.reset();
     }
 }
--- autotest/frontend/client/src/autotest/afe/JobDetailView.java        
2010-04-16 17:35:33.000000000 -0700
+++ autotest/frontend/client/src/autotest/afe/JobDetailView.java        
2010-04-16 17:35:33.000000000 -0700
@@ -76,6 +76,8 @@
     private Label controlFile = new Label();
     private DisclosurePanel controlFilePanel = new DisclosurePanel("");
     
+    protected StaticDataRepository staticData = 
StaticDataRepository.getRepository();
+    
     public JobDetailView(JobDetailListener listener) {
         this.listener = listener;
     }
@@ -114,6 +116,10 @@
                 showField(jobObject, "synch_count", "view_synch_count");
                 showField(jobObject, "dependencies", "view_dependencies");
                 
+                if 
(staticData.getData("drone_sets_enabled").isBoolean().booleanValue()) {
+                    showField(jobObject, "drone_set", "view_drone_set");
+                }
+                
                 String header = 
Utils.jsonToString(jobObject.get("control_type")) + " control file";
                 controlFilePanel.getHeaderTextAccessor().setText(header);
                 
controlFile.setText(Utils.jsonToString(jobObject.get("control_file")));
@@ -208,6 +214,10 @@
         controlFile.addStyleName("code");
         controlFilePanel.setContent(controlFile);
         addWidget(controlFilePanel, "view_control_file");
+        
+        if 
(!staticData.getData("drone_sets_enabled").isBoolean().booleanValue()) {
+            AfeUtils.removeElement("view_drone_set_wrapper");
+        }
     }
 
     
--- autotest/frontend/client/src/autotest/afe/UserPreferencesView.java  
2010-04-16 17:35:33.000000000 -0700
+++ autotest/frontend/client/src/autotest/afe/UserPreferencesView.java  
2010-04-16 17:35:33.000000000 -0700
@@ -18,6 +18,7 @@
 import com.google.gwt.user.client.ui.CheckBox;
 import com.google.gwt.user.client.ui.FlexTable;
 import com.google.gwt.user.client.ui.HTMLTable;
+import com.google.gwt.user.client.ui.ListBox;
 import com.google.gwt.user.client.ui.Panel;
 import com.google.gwt.user.client.ui.VerticalPanel;
 import com.google.gwt.user.client.ui.Widget;
@@ -35,6 +36,7 @@
     
     private RadioChooser rebootBefore = new RadioChooser();
     private RadioChooser rebootAfter = new RadioChooser();
+    private ListBox droneSet = new ListBox();
     private CheckBox showExperimental = new CheckBox();
     private Button saveButton = new Button("Save preferences");
     private HTMLTable preferencesTable = new FlexTable();
@@ -60,6 +62,11 @@
         addOption("Reboot before", rebootBefore);
         addOption("Reboot after", rebootAfter);
         addOption("Show experimental tests", showExperimental);
+        if 
(staticData.getData("drone_sets_enabled").isBoolean().booleanValue()) {
+            AfeUtils.popualateListBox(droneSet, "drone_sets");
+            addOption("Drone set", droneSet);
+        }
+
         container.add(preferencesTable);
         container.add(saveButton);
         addWidget(container, "user_preferences_table");
@@ -81,6 +88,8 @@
     private void updateValues() {
         rebootBefore.setSelectedChoice(getValue("reboot_before"));
         rebootAfter.setSelectedChoice(getValue("reboot_after"));
+        AfeUtils.setSelectedItem(droneSet, getValue("drone_set"));
+
         
showExperimental.setValue(user.get("show_experimental").isBoolean().booleanValue());
     }
     
@@ -98,6 +107,7 @@
         values.put("id", user.get("id"));
         values.put("reboot_before", new 
JSONString(rebootBefore.getSelectedChoice()));
         values.put("reboot_after", new 
JSONString(rebootAfter.getSelectedChoice()));
+        values.put("drone_set", new 
JSONString(droneSet.getItemText(droneSet.getSelectedIndex())));
         values.put("show_experimental", 
JSONBoolean.getInstance(showExperimental.getValue()));
         proxy.rpcCall("modify_user", values, new JsonRpcCallback() {
             @Override
--- autotest/frontend/client/src/autotest/public/AfeClient.html 2010-04-16 
17:35:33.000000000 -0700
+++ autotest/frontend/client/src/autotest/public/AfeClient.html 2010-04-16 
17:35:33.000000000 -0700
@@ -72,6 +72,12 @@
           <span class="field-name">Reboot options:</span>
           <span id="view_reboot_before"></span> before job,
           <span id="view_reboot_after"></span> after job<br>
+
+          <span id="view_drone_set_wrapper">
+            <span class="field-name">Drone set:</span>
+            <span id="view_drone_set"></span><br>
+          </span>
+
           <span class="field-name">Include failed repair results:</span>
           <span id="view_parse_failed_repair"></span><br>
           <span class="field-name">Dependencies:</span>
@@ -134,6 +140,12 @@
               <td id="create_parse_failed_repair"></td><td></td></tr>
           <tr><td class="field-name">Hostless:</td>
               <td id="create_hostless"></td><td></td></tr>
+
+          <tr id="create_drone_set_wrapper">
+            <td class="field-name">Drone set:</td>
+            <td id="create_drone_set" colspan="2"></td>
+          </tr>
+
           <tr><td class="field-name">Tests:</td>
               <td id="create_tests" colspan="2"></td></tr>
           <tr><td class="field-name">Custom client tests:</td>
--- /dev/null   2009-12-17 12:29:38.000000000 -0800
+++ autotest/frontend/migrations/058_drone_management.py        2010-04-16 
17:35:33.000000000 -0700
@@ -0,0 +1,85 @@
+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_jobs
+ADD COLUMN drone_set_id INT;
+
+ALTER TABLE afe_jobs
+ADD CONSTRAINT afe_jobs_drone_set_ibfk
+FOREIGN KEY (drone_set_id) REFERENCES afe_drone_sets (id);
+
+
+ALTER TABLE afe_users
+ADD COLUMN drone_set_id INT;
+
+ALTER TABLE afe_users
+ADD CONSTRAINT afe_users_drone_set_ibfk
+FOREIGN KEY (drone_set_id) REFERENCES afe_drone_sets (id);
+
+
+UPDATE afe_special_tasks SET requested_by_id = (
+  SELECT id FROM afe_users WHERE login = 'autotest_system')
+WHERE requested_by_id IS NULL;
+
+ALTER TABLE afe_special_tasks
+MODIFY COLUMN requested_by_id INT NOT NULL;
+"""
+
+
+DOWN_SQL = """
+ALTER TABLE afe_special_tasks
+MODIFY COLUMN requested_by_id INT DEFAULT NULL;
+
+ALTER TABLE afe_users
+DROP FOREIGN KEY afe_users_drone_set_ibfk;
+
+ALTER TABLE afe_users
+DROP COLUMN drone_set_id;
+
+ALTER TABLE afe_jobs
+DROP FOREIGN KEY afe_jobs_drone_set_ibfk;
+
+ALTER TABLE afe_jobs
+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-04-16 17:35:33.000000000 -0700
+++ autotest/global_config.ini  2010-04-16 17:35:33.000000000 -0700
@@ -79,6 +79,9 @@
 gc_stats_interval_mins: 360
 # set nonzero to enable periodic reverification of all dead hosts
 reverify_period_minutes: 0
+drone_sets_enabled: False
+# default_drone_set_name: This is required if drone sets are enabled.
+default_drone_set_name:
 
 [HOSTS]
 wait_up_processes:
--- autotest/scheduler/drone_manager.py 2010-04-16 17:35:33.000000000 -0700
+++ autotest/scheduler/drone_manager.py 2010-04-16 17:35:33.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,13 @@
         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 failed; no drones available: %s'
+                                    % 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-04-16 
17:35:33.000000000 -0700
+++ autotest/scheduler/drone_manager_unittest.py        2010-04-16 
17:35:33.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'
@@ -214,7 +273,8 @@
         self.manager.execute_command(command=['test'],
                                      working_directory=self._WORKING_DIRECTORY,
                                      pidfile_name='mypidfile',
-                                     num_processes=1)
+                                     num_processes=1,
+                                     drone_hostnames_allowed=None)
 
         self.assert_(self.mock_drone.was_call_queued(
                 'write_to_file',
--- autotest/scheduler/monitor_db.py    2010-04-16 17:35:33.000000000 -0700
+++ autotest/scheduler/monitor_db.py    2010-04-16 17:35:33.000000000 -0700
@@ -74,6 +74,17 @@
         'get_metahost_schedulers', lambda : ())
 
 
+def _verify_default_drone_set_exists():
+    if (models.DroneSet.drone_sets_enabled() and
+            not models.DroneSet.default_drone_set_name()):
+        raise SchedulerError('Drone sets are enabled, but no default is set')
+
+
+def _sanity_check():
+    """Make sure the configs are consistent before starting the scheduler"""
+    _verify_default_drone_set_exists()
+
+
 def main():
     try:
         try:
@@ -1138,7 +1149,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 +1256,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 +1264,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 +1678,49 @@
                 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 not models.DroneSet.drone_sets_enabled():
+            return None
+
+        hqes = 
models.HostQueueEntry.objects.filter(id__in=self.queue_entry_ids)
+        if not hqes:
+            # Only special tasks could be missing host queue entries
+            assert isinstance(self, SpecialAgentTask)
+            return self._user_or_global_default_drone_set(
+                    self.task, self.task.requested_by)
+
+        job_ids = hqes.values_list('job', flat=True).distinct()
+        assert job_ids.count() == 1, ("AgentTask's queue entries "
+                                      "span multiple jobs")
+
+        job = models.Job.objects.get(id=job_ids[0])
+        drone_set = job.drone_set
+        if not drone_set:
+            return self_user_or_global_default_drone_set(job, job.user())
+
+        return drone_set.get_drone_hostnames()
+
+
+    def _user_or_global_default_drone_set(self, obj_with_owner, user):
+        """
+        Returns the user's default drone set, if present.
+
+        Otherwise, returns the global default drone set.
+        """
+        default_hostnames = models.DroneSet.get_default().get_drone_hostnames()
+        if not user:
+            logging.warn('%s had no owner; using default drone set',
+                         obj_with_owner)
+            return default_hostnames
+        if not user.drone_set:
+            logging.warn('User %s has no default drone set, using global '
+                         'default', user.login)
+            return default_hostnames
+        return user.drone_set.get_drone_hostnames()
 
 
     def register_necessary_pidfiles(self):
--- autotest/scheduler/monitor_db_functional_test.py    2010-04-16 
17:35:33.000000000 -0700
+++ autotest/scheduler/monitor_db_functional_test.py    2010-04-16 
17:35:33.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:
@@ -707,12 +707,16 @@
         queue_entry.save()
 
         # make some dummy SpecialTasks that shouldn't count
-        models.SpecialTask.objects.create(host=queue_entry.host,
-                                          task=models.SpecialTask.Task.VERIFY)
-        models.SpecialTask.objects.create(host=queue_entry.host,
-                                          task=models.SpecialTask.Task.CLEANUP,
-                                          queue_entry=queue_entry,
-                                          is_complete=True)
+        models.SpecialTask.objects.create(
+                host=queue_entry.host,
+                task=models.SpecialTask.Task.VERIFY,
+                requested_by=models.User.current_user())
+        models.SpecialTask.objects.create(
+                host=queue_entry.host,
+                task=models.SpecialTask.Task.CLEANUP,
+                queue_entry=queue_entry,
+                is_complete=True,
+                requested_by=models.User.current_user())
 
         self.assertRaises(monitor_db.SchedulerError, self._initialize_test)
 
--- autotest/scheduler/monitor_db_unittest.py   2010-04-16 17:35:33.000000000 
-0700
+++ autotest/scheduler/monitor_db_unittest.py   2010-04-16 17:35:33.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,8 @@
         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 +1362,65 @@
         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):
+        self.god.stub_function(models.DroneSet, 'drone_sets_enabled')
+        models.DroneSet.drone_sets_enabled.expect_call().and_return(True)
+
+        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])
+        drone_set_3 = models.DroneSet.objects.create(name='3')
+
+        job_1 = self._create_job_simple([self.hosts[0].id],
+                                        drone_set=drone_set_1)
+        job_2 = self._create_job_simple([self.hosts[0].id],
+                                        drone_set=drone_set_2)
+        job_3 = self._create_job_simple([self.hosts[0].id],
+                                        drone_set=drone_set_3)
+
+        hqe_1 = job_1.hostqueueentry_set.all()[0].id
+        hqe_2 = job_2.hostqueueentry_set.all()[0].id
+        hqe_3 = job_3.hostqueueentry_set.all()[0].id
+
+        return (hqe_1, hqe_2, hqe_3), monitor_db.AgentTask()
+
+
+    def test_get_drone_hostnames_allowed_no_associated_sets(self):
+        hqes, task = self._setup_drones()
+        task.queue_entry_ids = (hqes[2],)
+        self.assertEqual(set(), task.get_drone_hostnames_allowed())
+        self.god.check_playback()
+
+
+    def test_get_drone_hostnames_allowed_success(self):
+        hqes, task = self._setup_drones()
+        task.queue_entry_ids = (hqes[0],)
+        self.assertEqual(set(('0','1')), task.get_drone_hostnames_allowed())
+        self.god.check_playback()
+
+
+    def test_get_drone_hostnames_allowed_multiple_jobs(self):
+        hqes, task = self._setup_drones()
+        task.queue_entry_ids = (hqes[0], hqes[1])
+        self.assertRaises(AssertionError,
+                          task.get_drone_hostnames_allowed)
+        self.god.check_playback()
+
+
 if __name__ == '__main__':
     unittest.main()
--- autotest/scheduler/scheduler_models.py      2010-04-16 17:35:33.000000000 
-0700
+++ autotest/scheduler/scheduler_models.py      2010-04-16 17:35:33.000000000 
-0700
@@ -763,7 +763,7 @@
     _fields = ('id', 'owner', 'name', 'priority', 'control_file',
                'control_type', 'created_on', 'synch_count', 'timeout',
                'run_verify', 'email_list', 'reboot_before', 'reboot_after',
-               'parse_failed_repair', 'max_runtime_hrs')
+               'parse_failed_repair', 'max_runtime_hrs', 'drone_set_id')
 
     # This does not need to be a column in the DB.  The delays are likely to
     # be configured short.  If the scheduler is stopped and restarted in
_______________________________________________
Autotest mailing list
[email protected]
http://test.kernel.org/cgi-bin/mailman/listinfo/autotest

Reply via email to