fix(initialize): ensure update-engine is restarted after group/server

changes

initialize/config.go

@@ -20,11 +20,9 @@ type CloudConfigFile interface {
 }
 
 // CloudConfigUnit represents a CoreOS specific configuration option that can generate
-// an associated system.Unit to be created/enabled appropriately
+// associated system.Units to be created/enabled appropriately
 type CloudConfigUnit interface {
-	// Unit should either return (*system.Unit, error), or (nil, nil) if nothing
-	// needs to be done for this configuration option.
-	Unit(root string) (*system.Unit, error)
+	Units(root string) ([]system.Unit, error)
 }
 
 // CloudConfig encapsulates the entire cloud-config configuration file and maps directly to YAML
@@ -215,13 +213,11 @@ func Apply(cfg CloudConfig, env *Environment) error {
 	}
 
 	for _, ccu := range []CloudConfigUnit{cfg.Coreos.Etcd, cfg.Coreos.Fleet, cfg.Coreos.Update} {
-		u, err := ccu.Unit(env.Root())
+		u, err := ccu.Units(env.Root())
 		if err != nil {
 			return err
 		}
-		if u != nil {
-			cfg.Coreos.Units = append(cfg.Coreos.Units, *u)
-		}
+		cfg.Coreos.Units = append(cfg.Coreos.Units, u...)
 	}
 
 	for _, file := range cfg.WriteFiles {

initialize/etcd.go

@@ -28,9 +28,9 @@ func (ee EtcdEnvironment) String() (out string) {
 	return
 }
 
-// Unit creates a Unit file drop-in for etcd, using any configured
+// Units creates a Unit file drop-in for etcd, using any configured
 // options and adding a default MachineID if unset.
-func (ee EtcdEnvironment) Unit(root string) (*system.Unit, error) {
+func (ee EtcdEnvironment) Units(root string) ([]system.Unit, error) {
 	if ee == nil {
 		return nil, nil
 	}
@@ -45,10 +45,11 @@ func (ee EtcdEnvironment) Unit(root string) (*system.Unit, error) {
 		}
 	}
 
-	return &system.Unit{
+	etcd := system.Unit{
 		Name:    "etcd.service",
 		Runtime: true,
 		DropIn:  true,
 		Content: ee.String(),
-	}, nil
+	}
+	return []system.Unit{etcd}, nil
 }

initialize/etcd_test.go

@@ -59,7 +59,7 @@ Environment="ETCD_PEER_BIND_ADDR=127.0.0.1:7002"
 }
 
 func TestEtcdEnvironmentWrittenToDisk(t *testing.T) {
-	ec := EtcdEnvironment{
+	ee := EtcdEnvironment{
 		"name":           "node001",
 		"discovery":      "http://disco.example.com/foobar",
 		"peer-bind-addr": "127.0.0.1:7002",
@@ -70,17 +70,18 @@ func TestEtcdEnvironmentWrittenToDisk(t *testing.T) {
 	}
 	defer os.RemoveAll(dir)
 
-	u, err := ec.Unit(dir)
+	uu, err := ee.Units(dir)
 	if err != nil {
 		t.Fatalf("Generating etcd unit failed: %v", err)
 	}
-	if u == nil {
-		t.Fatalf("Returned nil etcd unit unexpectedly")
+	if len(uu) != 1 {
+		t.Fatalf("Expected 1 unit to be returned, got %d", len(uu))
 	}
+	u := uu[0]
 
-	dst := system.UnitDestination(u, dir)
+	dst := system.UnitDestination(&u, dir)
 	os.Stderr.WriteString("writing to " + dir + "\n")
-	if err := system.PlaceUnit(u, dst); err != nil {
+	if err := system.PlaceUnit(&u, dst); err != nil {
 		t.Fatalf("Writing of EtcdEnvironment failed: %v", err)
 	}
 
@@ -124,17 +125,18 @@ func TestEtcdEnvironmentWrittenToDiskDefaultToMachineID(t *testing.T) {
 		t.Fatalf("Failed writing out /etc/machine-id: %v", err)
 	}
 
-	u, err := ee.Unit(dir)
+	uu, err := ee.Units(dir)
 	if err != nil {
 		t.Fatalf("Generating etcd unit failed: %v", err)
 	}
-	if u == nil {
-		t.Fatalf("Returned nil etcd unit unexpectedly")
+	if len(uu) == 0 {
+		t.Fatalf("Returned empty etcd units unexpectedly")
 	}
+	u := uu[0]
 
-	dst := system.UnitDestination(u, dir)
+	dst := system.UnitDestination(&u, dir)
 	os.Stderr.WriteString("writing to " + dir + "\n")
-	if err := system.PlaceUnit(u, dst); err != nil {
+	if err := system.PlaceUnit(&u, dst); err != nil {
 		t.Fatalf("Writing of EtcdEnvironment failed: %v", err)
 	}
 
@@ -159,8 +161,8 @@ func TestEtcdEnvironmentWhenNil(t *testing.T) {
 	if ee != nil {
 		t.Fatalf("EtcdEnvironment is not nil")
 	}
-	u, err := ee.Unit("")
-	if u != nil || err != nil {
-		t.Fatalf("Unit returned a non-nil value for nil input")
+	uu, err := ee.Units("")
+	if len(uu) != 0 || err != nil {
+		t.Fatalf("Units returned value for nil input")
 	}
 }

initialize/fleet.go

@@ -19,16 +19,17 @@ func (fe FleetEnvironment) String() (out string) {
 	return
 }
 
-// Unit generates a Unit file drop-in for fleet, if any fleet options were
+// Units generates a Unit file drop-in for fleet, if any fleet options were
 // configured in cloud-config
-func (fe FleetEnvironment) Unit(root string) (*system.Unit, error) {
+func (fe FleetEnvironment) Units(root string) ([]system.Unit, error) {
 	if len(fe) < 1 {
 		return nil, nil
 	}
-	return &system.Unit{
+	fleet := system.Unit{
 		Name:    "fleet.service",
 		Runtime: true,
 		DropIn:  true,
 		Content: fe.String(),
-	}, nil
+	}
+	return []system.Unit{fleet}, nil
 }

initialize/fleet_test.go

@@ -19,20 +19,21 @@ Environment="FLEET_PUBLIC_IP=12.34.56.78"
 
 func TestFleetUnit(t *testing.T) {
 	cfg := make(FleetEnvironment, 0)
-	u, err := cfg.Unit("/")
-	if u != nil {
+	uu, err := cfg.Units("/")
+	if len(uu) != 0 {
 		t.Errorf("unexpectedly generated unit with empty FleetEnvironment")
 	}
 
 	cfg["public-ip"] = "12.34.56.78"
 
-	u, err = cfg.Unit("/")
+	uu, err = cfg.Units("/")
 	if err != nil {
 		t.Errorf("error generating fleet unit: %v", err)
 	}
-	if u == nil {
-		t.Fatalf("unexpectedly got nil unit generating fleet unit!")
+	if len(uu) != 1 {
+		t.Fatalf("expected 1 unit generated, got %d", len(uu))
 	}
+	u := uu[0]
 	if !u.Runtime {
 		t.Errorf("bad Runtime for generated fleet unit!")
 	}

initialize/update.go

@@ -126,24 +126,39 @@ func (uc UpdateConfig) File(root string) (*system.File, error) {
 	}, nil
 }
 
-// GetUnit generates a locksmith system.Unit, if reboot-strategy was set in
-// cloud-config, for the cloud-init initializer to act on appropriately
-func (uc UpdateConfig) Unit(root string) (*system.Unit, error) {
-	strategy, ok := uc["reboot-strategy"]
-	if !ok {
-		return nil, nil
+// Units generates units for the cloud-init initializer to act on:
+// - a locksmith system.Unit, if "reboot-strategy" was set in cloud-config
+// - an update_engine system.Unit, if "group" was set in cloud-config
+func (uc UpdateConfig) Units(root string) ([]system.Unit, error) {
+	var units []system.Unit
+	if strategy, ok := uc["reboot-strategy"]; ok {
+		ls := &system.Unit{
+			Name:    locksmithUnit,
+			Command: "restart",
+			Mask:    false,
+		}
+
+		if strategy == "off" {
+			ls.Command = "stop"
+			ls.Mask = true
+		}
+		units = append(units, *ls)
 	}
 
-	u := &system.Unit{
-		Name:    locksmithUnit,
-		Command: "restart",
-		Mask:    false,
+	rue := false
+	if _, ok := uc["group"]; ok {
+		rue = true
+	}
+	if _, ok := uc["server"]; ok {
+		rue = true
+	}
+	if rue {
+		ue := system.Unit{
+			Name:    "update-engine",
+			Command: "restart",
+		}
+		units = append(units, ue)
 	}
 
-	if strategy == "off" {
-		u.Command = "stop"
-		u.Mask = true
-	}
-
-	return u, nil
+	return units, nil
 }

initialize/update_test.go

@@ -38,12 +38,12 @@ func TestEmptyUpdateConfig(t *testing.T) {
 	if f != nil {
 		t.Errorf("getting file from empty UpdateConfig should have returned nil, got %v", f)
 	}
-	u, err := uc.Unit("")
+	uu, err := uc.Units("")
 	if err != nil {
 		t.Error("unexpected error getting unit from empty UpdateConfig")
 	}
-	if u != nil {
-		t.Errorf("getting unit from empty UpdateConfig should have returned nil, got %v", u)
+	if len(uu) != 0 {
+		t.Errorf("getting unit from empty UpdateConfig should have returned zero units, got %d", len(uu))
 	}
 }
 
@@ -106,6 +106,21 @@ SERVER=http://foo.com`
 			t.Errorf("File has incorrect contents, got %v, want %v", got, want)
 		}
 	}
+
+	uu, err := u.Units(dir)
+	if err != nil {
+		t.Errorf("unexpected error getting units from UpdateConfig: %v", err)
+	} else if len(uu) != 1 {
+		t.Errorf("unexpected number of files returned from UpdateConfig: want 1, got %d", len(uu))
+	} else {
+		unit := uu[0]
+		if unit.Name != "update-engine" {
+			t.Errorf("bad name for generated unit: want update-engine, got %s", unit.Name)
+		}
+		if unit.Command != "restart" {
+			t.Errorf("bad command for generated unit: want restart, got %s", unit.Command)
+		}
+	}
 }
 
 func TestRebootStrategies(t *testing.T) {
@@ -145,12 +160,13 @@ func TestRebootStrategies(t *testing.T) {
 				t.Errorf("couldn't find expected line %v for reboot-strategy=%v", s.line)
 			}
 		}
-		u, err := uc.Unit(dir)
+		uu, err := uc.Units(dir)
 		if err != nil {
 			t.Errorf("failed to generate unit for reboot-strategy=%v!", s.name)
-		} else if u == nil {
-			t.Errorf("generated empty unit for reboot-strategy=%v", s.name)
+		} else if len(uu) != 1 {
+			t.Errorf("unexpected number of units for reboot-strategy=%v: %d", s.name, len(uu))
 		} else {
+			u := uu[0]
 			if u.Name != locksmithUnit {
 				t.Errorf("unit generated for reboot strategy=%v had bad name: %v", s.name, u.Name)
 			}