nodecontainer.diff

diff --git a/provision/swarm/docker.go b/provision/swarm/docker.go
index 8d7b4eb..64b4bcf 100644
--- a/provision/swarm/docker.go
+++ b/provision/swarm/docker.go
@@ -601,7 +601,7 @@ func serviceSpecForNodeContainer(name, pool string) (*swarm.ServiceSpec, error)
 	return service, nil
 }
 
-func upsertService(spec *swarm.ServiceSpec, client *docker.Client) (error, bool) {
+func upsertService(spec *swarm.ServiceSpec, client *docker.Client, placementOnly bool) (error, bool) {
 	currService, err := client.InspectService(spec.Name)
 	if err != nil {
 		if _, ok := err.(*docker.NoSuchService); !ok {
@@ -614,6 +614,10 @@ func upsertService(spec *swarm.ServiceSpec, client *docker.Client) (error, bool)
 		}
 		return nil, true
 	}
+	if placementOnly {
+		currService.Spec.TaskTemplate.Placement = spec.TaskTemplate.Placement
+		spec = &currService.Spec
+	}
 	opts := docker.UpdateServiceOptions{
 		ServiceSpec: *spec,
 		Version:     currService.Version.Index,
diff --git a/provision/swarm/provisioner.go b/provision/swarm/provisioner.go
index 912706a..92c2516 100644
--- a/provision/swarm/provisioner.go
+++ b/provision/swarm/provisioner.go
@@ -877,66 +877,36 @@ func (p *swarmProvisioner) UpgradeNodeContainer(name string, pool string, writer
 		}
 		return err
 	}
-	poolsToRun := []string{pool}
+	poolsToRun := []string{pool, ""}
 	if pool == "" {
 		poolMap, errLoad := nodecontainer.LoadNodeContainersForPools(name)
 		if errLoad != nil {
 			return errors.WithStack(errLoad)
 		}
-		poolsToRun = make([]string, len(poolMap))
-		i := 0
+		poolsToRun = make([]string, 0, len(poolMap))
 		for k, v := range poolMap {
 			if !v.Valid() {
 				continue
 			}
-			if k == "" {
-				poolsToRun[len(poolMap)-1] = k
-				continue
-			}
-			poolsToRun[i] = k
-			i++
+			poolsToRun = append(poolsToRun, k)
 		}
 	}
 	var allErrors []error
-	created := false
 	for _, v := range poolsToRun {
 		serviceSpec, errUpsert := serviceSpecForNodeContainer(name, v)
 		if errUpsert != nil {
 			errUpsert = errors.Wrapf(errUpsert, "[node containers] failed retrieve service spec for node container %q [%s]", name, v)
 			allErrors = append(allErrors, errUpsert)
 		}
+		placementOnly := v == "" && pool != ""
 		log.Debugf("[node containers] upserting service %q for node container %s [%s]", serviceSpec.Name, name, v)
 		fmt.Fprintf(writer, "upserting service %q for node container %q [%s]\n", serviceSpec.Name, name, v)
-		errUpsert, created = upsertService(serviceSpec, client)
+		errUpsert, _ = upsertService(serviceSpec, client, placementOnly)
 		if errUpsert != nil {
 			errUpsert = errors.Wrapf(errUpsert, "[node containers] failed upsert service %q for node container %q [%s]", serviceSpec.Name, name, v)
 			allErrors = append(allErrors, errUpsert)
 		}
 	}
-	if pool != "" && created {
-		serviceName := nodeContainerServiceName(name, "")
-		baseSpec, err := client.InspectService(serviceName)
-		if err != nil {
-			if _, ok := err.(*docker.NoSuchService); ok {
-				return nil
-			}
-			err = errors.Wrapf(err, "[node containers] failed inspect base service %q for node container %q", serviceName, name)
-			allErrors = append(allErrors, err)
-		}
-		newBaseSpec, err := serviceSpecForNodeContainer(name, "")
-		if err != nil {
-			err = errors.Wrapf(err, "[node containers] failed retrieve base service spec %q for node container %q", serviceName, name)
-			allErrors = append(allErrors, err)
-		}
-		baseSpec.Spec.TaskTemplate.Placement = newBaseSpec.TaskTemplate.Placement
-		log.Debugf("[node containers] updating base service %q for node container %s constraints", serviceName, name)
-		fmt.Fprintf(writer, "updating base service %q for node container %s constraints\n", serviceName, name)
-		err, _ = upsertService(&baseSpec.Spec, client)
-		if err != nil {
-			err = errors.Wrapf(err, "[node containers] failed update base service %q for node container %q", serviceName, name)
-			allErrors = append(allErrors, err)
-		}
-	}
 	if len(allErrors) == 0 {
 		return nil
 	}
@@ -961,15 +931,9 @@ func (p *swarmProvisioner) ensureNodeContainersCreated() error {
 		return errors.WithStack(err)
 	}
 	for _, n := range names {
-		poolMap, err := nodecontainer.LoadNodeContainersForPools(n)
+		err = p.UpgradeNodeContainer(n, "", ioutil.Discard)
 		if err != nil {
-			return errors.WithStack(err)
-		}
-		for pool := range poolMap {
-			err = p.UpgradeNodeContainer(n, pool, ioutil.Discard)
-			if err != nil {
-				return err
-			}
+			return err
 		}
 	}
 	return nil