Merge lp:~axwalk/juju-core/provider-unit-assignment-policy into lp:~go-bot/juju-core/trunk

[Andrew Wilkins]

Introduce InstanceDistributor policy

This change introduces an InstanceDistributor
policy, which will be used to implement automatic
Availability Zone management for AWS and OpenStack.

When an attempt is made to assign a unit to a clean
machine, we gather the instances that have already
been provisioned, and provide those and the unit's
distribution group to the provider. The provider
may select as many or as few of the candidates as it
wishes to allow; the assignment policy will then
choose one of them for assigning to. If no candidates
are returned, the unit will be assigned to a new
machine.

There are currently no implementations of this policy;
they will be forthcoming when manual AZ support lands
in the amazon/openstack providers.

https://codereview.appspot.com/99660047/

Introduce InstanceDistributor policy

This change introduces an InstanceDistributor
policy, which will be used to implement automatic
Availability Zone management for AWS and OpenStack.

When an attempt is made to assign a unit to a clean
machine, we gather the instances that have already
been provisioned, and provide those and the unit's
distribution group to the provider. The provider
may select as many or as few of the candidates as it
wishes to allow; the assignment policy will then
choose one of them for assigning to. If no candidates
are returned, the unit will be assigned to a new
machine.

There are currently no implementations of this policy;
they will be forthcoming when manual AZ support lands
in the amazon/openstack providers.

https://codereview.appspot.com/99660047/

[Andrew Wilkins (axwalk)]

Reviewers: mp+221481_code.launchpad.net,

Message:
Please take a look.

Description:
Introduce InstanceDistributor policy

This change introduces an InstanceDistributor
policy, which will be used to implement automatic
Availability Zone management for AWS and OpenStack.

When an attempt is made to assign a unit to a clean
machine, we gather the instances that have already
been provisioned, and provide those and the unit's
distribution group to the provider. The provider
may select as many or as few of the candidates as it
wishes to allow; the assignment policy will then
choose one of them for assigning to. If no candidates
are returned, the unit will be assigned to a new
machine.

There are currently no implementations of this policy;
they will be forthcoming when manual AZ support lands
in the amazon/openstack providers.

https://code.launchpad.net/~axwalk/juju-core/provider-unit-assignment-policy/+merge/221481

(do not edit description out of merge proposal)

Please review this at https://codereview.appspot.com/99660047/

Affected files (+379, -40 lines):
   A [revision details]
   M environs/statepolicy.go
   M state/apiserver/provisioner/provisioner.go
   M state/conn_test.go
   A state/distribution.go
   A state/distribution_test.go
   M state/policy.go
   M state/unit.go

[William Reade (fwereade)]

LGTM with a bit of mindful rearrangement.

https://codereview.appspot.com/99660047/diff/1/state/distribution.go
File state/distribution.go (right):

https://codereview.appspot.com/99660047/diff/1/state/distribution.go#newcode42
state/distribution.go:42: distributionGroup, err :=
service.ServiceInstances()
fwiw, this doesn't feel *quite* right because we don't actually ever
need/use the *Service directly -- the AllUnits could be figured out from
the unit name alone, so grabbing the service doc is just an unnecessary
db hit (I think?).

But it's a balance: using the methods we already have for
familarity/clarity's sake, vs writing more primitives that operate more
on ids and less on instantiated objects and *might* eventually lead us
in a leaner/cleaner direction. I have a slight preference for the latter
but I'll leave it to your judgment.

https://codereview.appspot.com/99660047/diff/1/state/distribution.go#newcode54
state/distribution.go:54: func (service *Service) ServiceInstances()
([]instance.Id, error) {
Suggestion: ServiceInstances(*state.Service). I retain a mild discomfort
over spreading the methods on a particular type across files, but just
making the receiver a regular param pretty much eliminates that feeling
of ick. I'm not sure it's *entirely* rational, but...

Anyway, same applies to Unit.distribute above. Doesn't feel like it's
quite in the right place.

https://codereview.appspot.com/99660047/diff/1/state/policy.go
File state/policy.go (right):

https://codereview.appspot.com/99660047/diff/1/state/policy.go#newcode204
state/policy.go:204: // TODO(axw) move this comment
where to? ;p

https://codereview.appspot.com/99660047/

[Andrew Wilkins (axwalk)]

Please take a look.

https://codereview.appspot.com/99660047/diff/1/state/distribution.go
File state/distribution.go (right):

https://codereview.appspot.com/99660047/diff/1/state/distribution.go#newcode42
state/distribution.go:42: distributionGroup, err :=
service.ServiceInstances()
On 2014/05/30 08:51:12, fwereade wrote:
> fwiw, this doesn't feel *quite* right because we don't actually ever
need/use
> the *Service directly -- the AllUnits could be figured out from the
unit name
> alone, so grabbing the service doc is just an unnecessary db hit (I
think?).

Indeed it is unnecessary. I've changed ServiceInstances to take a
service name, and extracted the logic of Service.AllUnits out to a free
function which takes a service name.

> But it's a balance: using the methods we already have for
familarity/clarity's
> sake, vs writing more primitives that operate more on ids and less on
> instantiated objects and *might* eventually lead us in a
leaner/cleaner
> direction. I have a slight preference for the latter but I'll leave it
to your
> judgment.

https://codereview.appspot.com/99660047/diff/1/state/distribution.go#newcode54
state/distribution.go:54: func (service *Service) ServiceInstances()
([]instance.Id, error) {
On 2014/05/30 08:51:12, fwereade wrote:
> Suggestion: ServiceInstances(*state.Service). I retain a mild
discomfort over
> spreading the methods on a particular type across files, but just
making the
> receiver a regular param pretty much eliminates that feeling of ick.
I'm not
> sure it's *entirely* rational, but...

> Anyway, same applies to Unit.distribute above. Doesn't feel like it's
quite in
> the right place.

Changed both to free functions.

https://codereview.appspot.com/99660047/diff/1/state/policy.go
File state/policy.go (right):

https://codereview.appspot.com/99660047/diff/1/state/policy.go#newcode204
state/policy.go:204: // TODO(axw) move this comment
On 2014/05/30 08:51:12, fwereade wrote:
> where to? ;p

Yeah, it's probably best just left here.

https://codereview.appspot.com/99660047/

=== modified file 'environs/statepolicy.go'
--- environs/statepolicy.go	2014-04-24 12:33:19 +0000
+++ environs/statepolicy.go	2014-05-30 12:33:30 +0000
@@ -4,6 +4,8 @@
 package environs
 
 import (
+	"github.com/juju/errors"
+
 	"launchpad.net/juju-core/constraints"
 	"launchpad.net/juju-core/environs/config"
 	"launchpad.net/juju-core/state"
@@ -44,3 +46,14 @@
 	}
 	return env.ConstraintsValidator()
 }
+
+func (environStatePolicy) InstanceDistributor(cfg *config.Config) (state.InstanceDistributor, error) {
+	env, err := New(cfg)
+	if err != nil {
+		return nil, err
+	}
+	if p, ok := env.(state.InstanceDistributor); ok {
+		return p, nil
+	}
+	return nil, errors.NotImplementedf("InstanceDistributor")
+}

=== modified file 'state/apiserver/provisioner/provisioner.go'
--- state/apiserver/provisioner/provisioner.go	2014-05-13 12:57:53 +0000
+++ state/apiserver/provisioner/provisioner.go	2014-05-30 12:33:30 +0000
@@ -416,33 +416,12 @@
 		if !unit.IsPrincipal() {
 			continue
 		}
-		service, err := unit.Service()
-		if err != nil {
-			return nil, err
-		}
-		allUnits, err := service.AllUnits()
-		if err != nil {
-			return nil, err
-		}
-		for _, unit := range allUnits {
-			machineId, err := unit.AssignedMachineId()
-			if state.IsNotAssigned(err) {
-				continue
-			} else if err != nil {
-				return nil, err
-			}
-			machine, err := st.Machine(machineId)
-			if err != nil {
-				return nil, err
-			}
-			instanceId, err := machine.InstanceId()
-			if err == nil {
-				instanceIdSet.Add(string(instanceId))
-			} else if state.IsNotProvisionedError(err) {
-				continue
-			} else {
-				return nil, err
-			}
+		instanceIds, err := state.ServiceInstances(st, unit.ServiceName())
+		if err != nil {
+			return nil, err
+		}
+		for _, instanceId := range instanceIds {
+			instanceIdSet.Add(string(instanceId))
 		}
 	}
 	instanceIds := make([]instance.Id, instanceIdSet.Size())

=== modified file 'state/conn_test.go'
--- state/conn_test.go	2014-05-20 04:27:02 +0000
+++ state/conn_test.go	2014-05-30 12:33:30 +0000
@@ -105,6 +105,7 @@
 	getConfigValidator      func(string) (state.ConfigValidator, error)
 	getEnvironCapability    func(*config.Config) (state.EnvironCapability, error)
 	getConstraintsValidator func(*config.Config) (constraints.Validator, error)
+	getInstanceDistributor  func(*config.Config) (state.InstanceDistributor, error)
 }
 
 func (p *mockPolicy) Prechecker(cfg *config.Config) (state.Prechecker, error) {
@@ -134,3 +135,10 @@
 	}
 	return nil, errors.NewNotImplemented(nil, "ConstraintsValidator")
 }
+
+func (p *mockPolicy) InstanceDistributor(cfg *config.Config) (state.InstanceDistributor, error) {
+	if p.getInstanceDistributor != nil {
+		return p.getInstanceDistributor(cfg)
+	}
+	return nil, errors.NewNotImplemented(nil, "InstanceDistributor")
+}

=== added file 'state/distribution.go'
--- state/distribution.go	1970-01-01 00:00:00 +0000
+++ state/distribution.go	2014-05-30 12:33:30 +0000
@@ -0,0 +1,77 @@
+// Copyright 2014 Canonical Ltd.
+// Licensed under the AGPLv3, see LICENCE file for details.
+
+package state
+
+import (
+	"fmt"
+
+	"github.com/juju/errors"
+
+	"launchpad.net/juju-core/instance"
+)
+
+// distributeuUnit takes a unit and set of clean, possibly empty, instances
+// and asks the InstanceDistributor policy (if any) which ones are suitable
+// for assigning the unit to. If there is no InstanceDistributor, or the
+// distribution group is empty, then all of the candidates will be returned.
+func distributeUnit(u *Unit, candidates []instance.Id) ([]instance.Id, error) {
+	if len(candidates) == 0 {
+		return nil, nil
+	}
+	if u.st.policy == nil {
+		return candidates, nil
+	}
+	cfg, err := u.st.EnvironConfig()
+	if err != nil {
+		return nil, err
+	}
+	distributor, err := u.st.policy.InstanceDistributor(cfg)
+	if errors.IsNotImplemented(err) {
+		return candidates, nil
+	} else if err != nil {
+		return nil, err
+	}
+	if distributor == nil {
+		return nil, fmt.Errorf("policy returned nil instance distributor without an error")
+	}
+	distributionGroup, err := ServiceInstances(u.st, u.doc.Service)
+	if err != nil {
+		return nil, err
+	}
+	if len(distributionGroup) == 0 {
+		return candidates, nil
+	}
+	return distributor.DistributeInstances(candidates, distributionGroup)
+}
+
+// ServiceInstances returns the instance IDs of provisioned
+// machines that are assigned units of the specified service.
+func ServiceInstances(st *State, service string) ([]instance.Id, error) {
+	units, err := allUnits(st, service)
+	if err != nil {
+		return nil, err
+	}
+	instanceIds := make([]instance.Id, 0, len(units))
+	for _, unit := range units {
+		machineId, err := unit.AssignedMachineId()
+		if IsNotAssigned(err) {
+			continue
+		} else if err != nil {
+			return nil, err
+		}
+		machine, err := st.Machine(machineId)
+		if err != nil {
+			return nil, err
+		}
+		instanceId, err := machine.InstanceId()
+		if err == nil {
+			instanceIds = append(instanceIds, instanceId)
+		} else if IsNotProvisionedError(err) {
+			continue
+		} else {
+			return nil, err
+		}
+	}
+	return instanceIds, nil
+}

=== added file 'state/distribution_test.go'
--- state/distribution_test.go	1970-01-01 00:00:00 +0000
+++ state/distribution_test.go	2014-05-30 12:33:30 +0000
@@ -0,0 +1,185 @@
+// Copyright 2014 Canonical Ltd.
+// Licensed under the AGPLv3, see LICENCE file for details.
+
+package state_test
+
+import (
+	"fmt"
+
+	"github.com/juju/errors"
+	jc "github.com/juju/testing/checkers"
+	gc "launchpad.net/gocheck"
+
+	"launchpad.net/juju-core/environs/config"
+	"launchpad.net/juju-core/instance"
+	"launchpad.net/juju-core/state"
+)
+
+type InstanceDistributorSuite struct {
+	ConnSuite
+	distributor mockInstanceDistributor
+	wordpress   *state.Service
+	machines    []*state.Machine
+}
+
+var _ = gc.Suite(&InstanceDistributorSuite{})
+
+type mockInstanceDistributor struct {
+	candidates        []instance.Id
+	distributionGroup []instance.Id
+	result            []instance.Id
+	err               error
+}
+
+func (p *mockInstanceDistributor) DistributeInstances(candidates, distributionGroup []instance.Id) ([]instance.Id, error) {
+	p.candidates = candidates
+	p.distributionGroup = distributionGroup
+	result := p.result
+	if result == nil {
+		result = candidates
+	}
+	return result, p.err
+}
+
+func (s *InstanceDistributorSuite) SetUpTest(c *gc.C) {
+	s.ConnSuite.SetUpTest(c)
+	s.distributor = mockInstanceDistributor{}
+	s.policy.getInstanceDistributor = func(*config.Config) (state.InstanceDistributor, error) {
+		return &s.distributor, nil
+	}
+	s.wordpress = s.AddTestingServiceWithNetworks(
+		c,
+		"wordpress",
+		s.AddTestingCharm(c, "wordpress"),
+		[]string{"net1", "net2"},
+		[]string{"net3", "net4"},
+	)
+	s.machines = make([]*state.Machine, 3)
+	for i := range s.machines {
+		var err error
+		s.machines[i], err = s.State.AddOneMachine(state.MachineTemplate{
+			Series: "quantal",
+			Jobs:   []state.MachineJob{state.JobHostUnits},
+		})
+		c.Assert(err, gc.IsNil)
+	}
+}
+
+func (s *InstanceDistributorSuite) setupScenario(c *gc.C) {
+	// Assign a unit so we have a non-empty distribution group, and
+	// provision all instances so we have candidates.
+	unit, err := s.wordpress.AddUnit()
+	c.Assert(err, gc.IsNil)
+	err = unit.AssignToMachine(s.machines[0])
+	c.Assert(err, gc.IsNil)
+	for i, m := range s.machines {
+		instId := instance.Id(fmt.Sprintf("i-blah-%d", i))
+		err = m.SetProvisioned(instId, "fake-nonce", nil)
+		c.Assert(err, gc.IsNil)
+	}
+}
+
+func (s *InstanceDistributorSuite) TestDistributeInstances(c *gc.C) {
+	s.setupScenario(c)
+	unit, err := s.wordpress.AddUnit()
+	c.Assert(err, gc.IsNil)
+	_, err = unit.AssignToCleanMachine()
+	c.Assert(err, gc.IsNil)
+	c.Assert(s.distributor.candidates, jc.SameContents, []instance.Id{"i-blah-1", "i-blah-2"})
+	c.Assert(s.distributor.distributionGroup, jc.SameContents, []instance.Id{"i-blah-0"})
+	s.distributor.result = []instance.Id{}
+	_, err = unit.AssignToCleanMachine()
+	c.Assert(err, gc.ErrorMatches, eligibleMachinesInUse)
+}
+
+func (s *InstanceDistributorSuite) TestDistributeInstancesInvalidInstances(c *gc.C) {
+	s.setupScenario(c)
+	unit, err := s.wordpress.AddUnit()
+	c.Assert(err, gc.IsNil)
+	s.distributor.result = []instance.Id{"notthere"}
+	_, err = unit.AssignToCleanMachine()
+	c.Assert(err, gc.ErrorMatches, `cannot assign unit "wordpress/1" to clean machine: invalid instance returned: notthere`)
+}
+
+func (s *InstanceDistributorSuite) TestDistributeInstancesNoEmptyMachines(c *gc.C) {
+	for i := range s.machines {
+		// Assign a unit so we have a non-empty distribution group.
+		unit, err := s.wordpress.AddUnit()
+		c.Assert(err, gc.IsNil)
+		m, err := unit.AssignToCleanMachine()
+		c.Assert(err, gc.IsNil)
+		instId := instance.Id(fmt.Sprintf("i-blah-%d", i))
+		err = m.SetProvisioned(instId, "fake-nonce", nil)
+		c.Assert(err, gc.IsNil)
+	}
+
+	// InstanceDistributor is not called if there are no empty instances.
+	s.distributor.err = fmt.Errorf("no assignment for you")
+	unit, err := s.wordpress.AddUnit()
+	c.Assert(err, gc.IsNil)
+	_, err = unit.AssignToCleanMachine()
+	c.Assert(err, gc.ErrorMatches, eligibleMachinesInUse)
+}
+
+func (s *InstanceDistributorSuite) TestDistributeInstancesErrors(c *gc.C) {
+	s.setupScenario(c)
+	unit, err := s.wordpress.AddUnit()
+	c.Assert(err, gc.IsNil)
+
+	// Ensure that assignment fails when DistributeInstances returns an error.
+	s.distributor.err = fmt.Errorf("no assignment for you")
+	_, err = unit.AssignToCleanMachine()
+	c.Assert(err, gc.ErrorMatches, ".*no assignment for you")
+	_, err = unit.AssignToCleanEmptyMachine()
+	c.Assert(err, gc.ErrorMatches, ".*no assignment for you")
+	// If the policy's InstanceDistributor method fails, that will be returned first.
+	s.policy.getInstanceDistributor = func(*config.Config) (state.InstanceDistributor, error) {
+		return nil, fmt.Errorf("incapable of InstanceDistributor")
+	}
+	_, err = unit.AssignToCleanMachine()
+	c.Assert(err, gc.ErrorMatches, ".*incapable of InstanceDistributor")
+}
+
+func (s *InstanceDistributorSuite) TestDistributeInstancesEmptyDistributionGroup(c *gc.C) {
+	s.distributor.err = fmt.Errorf("no assignment for you")
+
+	// InstanceDistributor is not called if the distribution group is empty.
+	unit0, err := s.wordpress.AddUnit()
+	c.Assert(err, gc.IsNil)
+	_, err = unit0.AssignToCleanMachine()
+	c.Assert(err, gc.IsNil)
+
+	// Distribution group is still empty, because the machine assigned to has
+	// not been provisioned.
+	unit1, err := s.wordpress.AddUnit()
+	c.Assert(err, gc.IsNil)
+	_, err = unit1.AssignToCleanMachine()
+	c.Assert(err, gc.IsNil)
+}
+
+func (s *InstanceDistributorSuite) TestInstanceDistributorUnimplemented(c *gc.C) {
+	s.setupScenario(c)
+	var distributorErr error
+	s.policy.getInstanceDistributor = func(*config.Config) (state.InstanceDistributor, error) {
+		return nil, distributorErr
+	}
+	unit, err := s.wordpress.AddUnit()
+	c.Assert(err, gc.IsNil)
+	_, err = unit.AssignToCleanMachine()
+	c.Assert(err, gc.ErrorMatches, `cannot assign unit "wordpress/1" to clean machine: policy returned nil instance distributor without an error`)
+	distributorErr = errors.NotImplementedf("InstanceDistributor")
+	_, err = unit.AssignToCleanMachine()
+	c.Assert(err, gc.IsNil)
+}
+
+func (s *InstanceDistributorSuite) TestDistributeInstancesNoPolicy(c *gc.C) {
+	s.policy.getInstanceDistributor = func(*config.Config) (state.InstanceDistributor, error) {
+		c.Errorf("should not have been invoked")
+		return nil, nil
+	}
+	state.SetPolicy(s.State, nil)
+	unit, err := s.wordpress.AddUnit()
+	c.Assert(err, gc.IsNil)
+	_, err = unit.AssignToCleanMachine()
+	c.Assert(err, gc.IsNil)
+}

=== modified file 'state/policy.go'
--- state/policy.go	2014-05-13 04:30:48 +0000
+++ state/policy.go	2014-05-30 12:33:30 +0000
@@ -10,6 +10,7 @@
 
 	"launchpad.net/juju-core/constraints"
 	"launchpad.net/juju-core/environs/config"
+	"launchpad.net/juju-core/instance"
 )
 
 // Policy is an interface provided to State that may
@@ -22,21 +23,24 @@
 // be ignored. Any other error will cause an error
 // in the use of the policy.
 type Policy interface {
-	// Prechecker takes a *config.Config and returns
-	// a (possibly nil) Prechecker or an error.
+	// Prechecker takes a *config.Config and returns a Prechecker or an error.
 	Prechecker(*config.Config) (Prechecker, error)
 
-	// ConfigValidator takes a provider type name and returns
-	// a (possibly nil) ConfigValidator or an error.
+	// ConfigValidator takes a provider type name and returns a ConfigValidator
+	// or an error.
 	ConfigValidator(providerType string) (ConfigValidator, error)
 
-	// EnvironCapability takes a *config.Config and returns
-	// a (possibly nil) EnvironCapability or an error.
+	// EnvironCapability takes a *config.Config and returns an EnvironCapability
+	// or an error.
 	EnvironCapability(*config.Config) (EnvironCapability, error)
 
-	// ConstraintsValidator takes a *config.Config and returns
-	// a (possibly nil) constraints.Validator or an error.
+	// ConstraintsValidator takes a *config.Config and returns a
+	// constraints.Validator or an error.
 	ConstraintsValidator(*config.Config) (constraints.Validator, error)
+
+	// InstanceDistributor takes a *config.Config and returns an
+	// InstanceDistributor or an error.
+	InstanceDistributor(*config.Config) (InstanceDistributor, error)
 }
 
 // Prechecker is a policy interface that is provided to State
@@ -187,3 +191,21 @@
 	}
 	return capability.SupportsUnitPlacement()
 }
+
+// InstanceDistributor is a policy interface that is provided
+// to State to perform distribution of units across instances
+// for high availability.
+type InstanceDistributor interface {
+	// DistributeInstance takes a set of clean, empty
+	// instances, and a distribution group, and returns
+	// the subset of candidates which the policy will
+	// allow entry into the distribution group.
+	//
+	// The AssignClean and AssignCleanEmpty unit
+	// assignment policies will attempt to assign a
+	// unit to each of the resulting instances until
+	// one is successful. If no instances can be assigned
+	// to (e.g. because of concurrent deployments), then
+	// a new machine will be allocated.
+	DistributeInstances(candidates, distributionGroup []instance.Id) ([]instance.Id, error)
+}

=== modified file 'state/service.go'
--- state/service.go	2014-05-26 21:16:46 +0000
+++ state/service.go	2014-05-30 12:33:30 +0000
@@ -723,13 +723,17 @@
 
 // AllUnits returns all units of the service.
 func (s *Service) AllUnits() (units []*Unit, err error) {
+	return allUnits(s.st, s.doc.Name)
+}
+
+func allUnits(st *State, service string) (units []*Unit, err error) {
 	docs := []unitDoc{}
-	err = s.st.units.Find(bson.D{{"service", s.doc.Name}}).All(&docs)
+	err = st.units.Find(bson.D{{"service", service}}).All(&docs)
 	if err != nil {
-		return nil, fmt.Errorf("cannot get all units from service %q: %v", s, err)
+		return nil, fmt.Errorf("cannot get all units from service %q: %v", service, err)
 	}
 	for i := range docs {
-		units = append(units, newUnit(s.st, &docs[i]))
+		units = append(units, newUnit(st, &docs[i]))
 	}
 	return units, nil
 }

=== modified file 'state/unit.go'
--- state/unit.go	2014-05-27 05:36:00 +0000
+++ state/unit.go	2014-05-30 12:33:30 +0000
@@ -1277,14 +1277,63 @@
 		return nil, err
 	}
 
-	// TODO(rog) Fix so this is more efficient when there are concurrent uses.
-	// Possible solution: pick the highest and the smallest id of all
-	// unused machines, and try to assign to the first one >= a random id in the
-	// middle.
-	iter := query.Batch(1).Prefetch(0).Iter()
-	var mdoc machineDoc
-	for iter.Next(&mdoc) {
-		m := newMachine(u.st, &mdoc)
+	// Find all of the candidate machines, and associated
+	// instances for those that are provisioned. Instances
+	// will be distributed across in preference to
+	// unprovisioned machines.
+	var mdocs []*machineDoc
+	if err := query.All(&mdocs); err != nil {
+		assignContextf(&err, u, context)
+		return nil, err
+	}
+	var unprovisioned []*Machine
+	var instances []instance.Id
+	instanceMachines := make(map[instance.Id]*Machine)
+	for _, mdoc := range mdocs {
+		m := newMachine(u.st, mdoc)
+		instance, err := m.InstanceId()
+		if IsNotProvisionedError(err) {
+			unprovisioned = append(unprovisioned, m)
+		} else if err != nil {
+			assignContextf(&err, u, context)
+			return nil, err
+		} else {
+			instances = append(instances, instance)
+			instanceMachines[instance] = m
+		}
+	}
+
+	// Filter the list of instances that are suitable for
+	// distribution, and then map them back to machines.
+	//
+	// TODO(axw) 2014-05-30 #1324904
+	// Shuffle machines to reduce likelihood of collisions.
+	// The partition of provisioned/unprovisioned machines
+	// must be maintained.
+	if instances, err = distributeUnit(u, instances); err != nil {
+		assignContextf(&err, u, context)
+		return nil, err
+	}
+	machines := make([]*Machine, len(instances), len(instances)+len(unprovisioned))
+	for i, instance := range instances {
+		m, ok := instanceMachines[instance]
+		if !ok {
+			err := fmt.Errorf("invalid instance returned: %v", instance)
+			assignContextf(&err, u, context)
+			return nil, err
+		}
+		machines[i] = m
+	}
+	machines = append(machines, unprovisioned...)
+
+	// TODO(axw) 2014-05-30 #1253704
+	// We should not select a machine that is in the process
+	// of being provisioned. There's no point asserting that
+	// the machine hasn't been provisioned, as there'll still
+	// be a period of time during which the machine may be
+	// provisioned without the fact having yet been recorded
+	// in state.
+	for _, m := range machines {
 		err := u.assignToMachine(m, true)
 		if err == nil {
 			return m, nil
@@ -1294,10 +1343,6 @@
 			return nil, err
 		}
 	}
-	if err := iter.Err(); err != nil {
-		assignContextf(&err, u, context)
-		return nil, err
-	}
 	return nil, noCleanMachines
 }