Merge pull request #638 from jonboulle/638

Scope unit state by machine

registry/unit_state.go

@@ -11,57 +11,109 @@ import (
 )
 
 const (
+	// Legacy namespace for unit states
 	statePrefix = "/state/"
+	// Namespace for unit states stored per-machine
+	statesPrefix = "/states/"
 )
 
-// Get the current UnitState of the provided Job's Unit
+// legacyUnitStatePath returns the path where UnitState objects were formerly
+// reported before being moved to a machine-specific namespace
+// https://github.com/coreos/fleet/issues/638
+func (r *EtcdRegistry) legacyUnitStatePath(jobName string) string {
+	return path.Join(r.keyPrefix, statePrefix, jobName)
+}
+
+// unitStatesNamespace generates a keypath of a namespace containing all
+// UnitState objects for a particular job
+func (r *EtcdRegistry) unitStatesNamespace(jobName string) string {
+	return path.Join(r.keyPrefix, statesPrefix, jobName)
+}
+
+// unitStatePath generates a keypath where the UnitState object for a given
+// machine ID + jobName combination is stored
+func (r *EtcdRegistry) unitStatePath(machID, jobName string) string {
+	return path.Join(r.unitStatesNamespace(jobName), machID)
+}
+
+// getUnitState retrieves the current UnitState of the provided Job's Unit
 func (r *EtcdRegistry) getUnitState(jobName string) *unit.UnitState {
+	// TODO(jonboulle): deal with multiple UnitStates
+	legacyKey := r.legacyUnitStatePath(jobName)
 	req := etcd.Get{
-		Key:       path.Join(r.keyPrefix, statePrefix, jobName),
+		Key:       legacyKey,
 		Recursive: true,
 	}
 	resp, err := r.etcd.Do(&req)
 
-	// Assume the error was KeyNotFound and return an empty data structure
 	if err != nil {
+		if !isKeyNotFound(err) {
+			log.Errorf("Error retrieving UnitState(%s): %v", jobName, err)
+		}
 		return nil
 	}
 
 	var usm unitStateModel
-	//TODO: Handle the error generated by unmarshal
-	unmarshal(resp.Node.Value, &usm)
+	if err := unmarshal(resp.Node.Value, &usm); err != nil {
+		log.Errorf("Error unmarshalling UnitState(%s): %v", jobName, err)
+		return nil
+	}
 
 	return modelToUnitState(&usm)
 }
 
-// Persist the changes in a provided Machine's Job
+// SaveUnitState persists the given UnitState to the Registry
 func (r *EtcdRegistry) SaveUnitState(jobName string, unitState *unit.UnitState) {
 	usm := unitStateToModel(unitState)
 	if usm == nil {
 		log.Errorf("Unable to save nil UnitState model")
 		return
 	}
 
-	//TODO: Handle the error generated by marshal
-	json, _ := marshal(usm)
+	json, err := marshal(usm)
+	if err != nil {
+		log.Errorf("Error marshalling UnitState: %v", err)
+		return
+	}
 
+	legacyKey := r.legacyUnitStatePath(jobName)
 	req := etcd.Set{
-		Key:   path.Join(r.keyPrefix, statePrefix, jobName),
+		Key:   legacyKey,
+		Value: json,
+	}
+	r.etcd.Do(&req)
+
+	newKey := r.unitStatePath(unitState.MachineID, jobName)
+	req = etcd.Set{
+		Key:   newKey,
 		Value: json,
 	}
 	r.etcd.Do(&req)
 }
 
 // Delete the state from the Registry for the given Job's Unit
 func (r *EtcdRegistry) RemoveUnitState(jobName string) error {
+	// TODO(jonboulle): consider https://github.com/coreos/fleet/issues/465
+	legacyKey := r.legacyUnitStatePath(jobName)
 	req := etcd.Delete{
-		Key: path.Join(r.keyPrefix, statePrefix, jobName),
+		Key: legacyKey,
 	}
 	_, err := r.etcd.Do(&req)
-	if isKeyNotFound(err) {
-		err = nil
+	if err != nil && !isKeyNotFound(err) {
+		return err
+	}
+
+	// TODO(jonboulle): deal properly with multiple states
+	newKey := r.unitStatesNamespace(jobName)
+	req = etcd.Delete{
+		Key:       newKey,
+		Recursive: true,
+	}
+	_, err = r.etcd.Do(&req)
+	if err != nil && !isKeyNotFound(err) {
+		return err
 	}
-	return err
+	return nil
 }
 
 type unitStateModel struct {

registry/unit_state_test.go

@@ -0,0 +1,91 @@
+package registry
+
+import (
+	"reflect"
+	"testing"
+
+	"github.com/coreos/fleet/etcd"
+	"github.com/coreos/fleet/unit"
+)
+
+type action struct {
+	key string
+	val string
+	rec bool
+}
+
+type testEtcdClient struct {
+	sets    []action
+	deletes []action
+}
+
+func (t *testEtcdClient) Do(req etcd.Action) (*etcd.Result, error) {
+	if s, ok := req.(*etcd.Set); ok {
+		t.sets = append(t.sets, action{key: s.Key, val: s.Value})
+	} else if d, ok := req.(*etcd.Delete); ok {
+		t.deletes = append(t.deletes, action{key: d.Key, rec: d.Recursive})
+	}
+	return nil, nil
+}
+func (t *testEtcdClient) Wait(req etcd.Action, ch <-chan bool) (*etcd.Result, error) {
+	return nil, nil
+}
+
+func TestUnitStatePaths(t *testing.T) {
+	r := &EtcdRegistry{nil, "/fleet/"}
+	j := "foo.service"
+	want := "/fleet/state/foo.service"
+	got := r.legacyUnitStatePath(j)
+	if got != want {
+		t.Errorf("bad unit state path: got %v, want %v", got, want)
+	}
+	m := "abcdefghij"
+	want = "/fleet/states/foo.service/abcdefghij"
+	got = r.unitStatePath(m, j)
+	if got != want {
+		t.Errorf("bad unit state path: got %v, want %v", got, want)
+	}
+}
+
+func TestSaveUnitState(t *testing.T) {
+	e := &testEtcdClient{}
+	r := &EtcdRegistry{e, "/fleet/"}
+	j := "foo.service"
+	mID := "mymachine"
+	us := unit.NewUnitState("abc", "def", "ghi", mID)
+
+	r.SaveUnitState(j, us)
+
+	json := `{"loadState":"abc","activeState":"def","subState":"ghi","machineState":{"ID":"mymachine","PublicIP":"","Metadata":null,"Version":"","TotalResources":{"Cores":0,"Memory":0,"Disk":0},"FreeResources":{"Cores":0,"Memory":0,"Disk":0},"LoadedUnits":0}}`
+	p1 := "/fleet/state/foo.service"
+	p2 := "/fleet/states/foo.service/mymachine"
+	want := []action{
+		action{key: p1, val: json},
+		action{key: p2, val: json},
+	}
+	got := e.sets
+	if !reflect.DeepEqual(got, want) {
+		t.Errorf("bad result from SaveUnitState: \ngot\n%#v\nwant\n%#v", got, want)
+	}
+	if e.deletes != nil {
+		t.Errorf("unexpected deletes during SaveUnitState: %#v", e.deletes)
+	}
+}
+
+func TestRemoveUnitState(t *testing.T) {
+	e := &testEtcdClient{}
+	r := &EtcdRegistry{e, "/fleet/"}
+	j := "foo.service"
+	r.RemoveUnitState(j)
+	want := []action{
+		action{key: "/fleet/state/foo.service", rec: false},
+		action{key: "/fleet/states/foo.service", rec: true},
+	}
+	got := e.deletes
+	if !reflect.DeepEqual(got, want) {
+		t.Errorf("bad result from RemoveUnitState: \ngot\n%#v\nwant\n%#v", got, want)
+	}
+	if e.sets != nil {
+		t.Errorf("unexpected sets during RemoveUnitState: %#v", e.sets)
+	}
+}