Store the machine ID in the payload doc.

nammn · May 19, 2016 · bd4b5c5 · bd4b5c5
1 parent 642a288
commit bd4b5c5
Show file tree

Hide file tree

Showing 12 changed files with 151 additions and 179 deletions.
diff --git a/component/all/payload.go b/component/all/payload.go
@@ -205,8 +205,8 @@ func (payloads) registerState() {
 		return payloadstate.NewUnitPayloads(persist, unit, machine), nil
 	}
 
-	newEnvPayloads := func(db state.Persistence, st state.PayloadsEnvPersistence) (state.EnvPayloads, error) {
-		envPersist := persistence.NewEnvPersistence(db, st)
+	newEnvPayloads := func(db state.Persistence) (state.EnvPayloads, error) {
+		envPersist := persistence.NewEnvPersistence(db)
 		envPayloads := payloadstate.EnvPayloads{
 			Persist: envPersist,
 		}

diff --git a/payload/persistence/base_test.go b/payload/persistence/base_test.go
@@ -18,9 +18,10 @@ import (
 type BaseSuite struct {
 	testing.BaseSuite
 
-	Stub  *gitjujutesting.Stub
-	State *fakeStatePersistence
-	Unit  string
+	Stub    *gitjujutesting.Stub
+	State   *fakeStatePersistence
+	Unit    string
+	Machine string
 }
 
 func (s *BaseSuite) SetUpTest(c *gc.C) {
@@ -29,6 +30,7 @@ func (s *BaseSuite) SetUpTest(c *gc.C) {
 	s.Stub = &gitjujutesting.Stub{}
 	s.State = &fakeStatePersistence{Stub: s.Stub}
 	s.Unit = "a-unit/0"
+	s.Machine = "0"
 }
 
 type PayloadDoc payloadDoc
@@ -37,19 +39,20 @@ func (doc PayloadDoc) convert() *payloadDoc {
 	return (*payloadDoc)(&doc)
 }
 
-func (s *BaseSuite) NewDoc(id string, pl payload.Payload) *payloadDoc {
+func (s *BaseSuite) NewDoc(id string, pl payload.FullPayloadInfo) *payloadDoc {
 	return &payloadDoc{
-		DocID:   "payload#" + s.Unit + "#" + pl.Name,
-		UnitID:  s.Unit,
-		Name:    pl.Name,
-		StateID: id,
-		Type:    pl.Type,
-		RawID:   pl.ID,
-		State:   pl.Status,
+		DocID:     "payload#" + s.Unit + "#" + pl.Name,
+		UnitID:    s.Unit,
+		Name:      pl.Name,
+		MachineID: pl.Machine,
+		StateID:   id,
+		Type:      pl.Type,
+		RawID:     pl.ID,
+		State:     pl.Status,
 	}
 }
 
-func (s *BaseSuite) SetDoc(id string, pl payload.Payload) *payloadDoc {
+func (s *BaseSuite) SetDoc(id string, pl payload.FullPayloadInfo) *payloadDoc {
 	payloadDoc := s.NewDoc(id, pl)
 	s.State.SetDocs(payloadDoc)
 	return payloadDoc
@@ -68,28 +71,31 @@ func (s *BaseSuite) SetUnit(id string) {
 	s.Unit = id
 }
 
-func (s *BaseSuite) NewPayloads(pType string, ids ...string) []payload.Payload {
-	var payloads []payload.Payload
+func (s *BaseSuite) NewPayloads(pType string, ids ...string) []payload.FullPayloadInfo {
+	var payloads []payload.FullPayloadInfo
 	for _, id := range ids {
 		pl := s.NewPayload(pType, id)
 		payloads = append(payloads, pl)
 	}
 	return payloads
 }
 
-func (s *BaseSuite) NewPayload(pType string, id string) payload.Payload {
+func (s *BaseSuite) NewPayload(pType string, id string) payload.FullPayloadInfo {
 	name, pluginID := payload.ParseID(id)
 	if pluginID == "" {
 		pluginID = fmt.Sprintf("%s-%s", name, utils.MustNewUUID())
 	}
 
-	return payload.Payload{
-		PayloadClass: charm.PayloadClass{
-			Name: name,
-			Type: pType,
+	return payload.FullPayloadInfo{
+		Payload: payload.Payload{
+			PayloadClass: charm.PayloadClass{
+				Name: name,
+				Type: pType,
+			},
+			ID:     pluginID,
+			Status: "running",
+			Unit:   s.Unit,
 		},
-		ID:     pluginID,
-		Status: "running",
-		Unit:   s.Unit,
+		Machine: s.Machine,
 	}
 }
diff --git a/payload/persistence/env.go b/payload/persistence/env.go
@@ -9,25 +9,16 @@ import (
 	"github.com/juju/juju/payload"
 )
 
-// EnvPersistenceEntities provides all the information needed to produce
-// a new EnvPersistence value.
-type EnvPersistenceEntities interface {
-	// AssignedMachineID the machine to which the identfies unit is assigned.
-	AssignedMachineID(unitName string) (string, error)
-}
-
 // EnvPersistence provides the persistence functionality for the
 // Juju environment as a whole.
 type EnvPersistence struct {
 	db *Persistence
-	st EnvPersistenceEntities
 }
 
 // NewEnvPersistence wraps the "db" in a new EnvPersistence.
-func NewEnvPersistence(db PersistenceBase, st EnvPersistenceEntities) *EnvPersistence {
+func NewEnvPersistence(db PersistenceBase) *EnvPersistence {
 	return &EnvPersistence{
 		db: NewPersistence(db, ""),
-		st: st,
 	}
 }
 
@@ -40,21 +31,10 @@ func (ep *EnvPersistence) ListAll() ([]payload.FullPayloadInfo, error) {
 		return nil, errors.Trace(err)
 	}
 
-	unitMachines := make(map[string]string)
 	var fullPayloads []payload.FullPayloadInfo
 	for _, doc := range docs {
-		machineID, ok := unitMachines[doc.UnitID]
-		if !ok {
-			machineID, err = ep.st.AssignedMachineID(doc.UnitID)
-			if err != nil {
-				return nil, errors.Trace(err)
-			}
-			unitMachines[doc.UnitID] = machineID
-		}
-		fullPayloads = append(fullPayloads, payload.FullPayloadInfo{
-			Payload: doc.payload(doc.UnitID),
-			Machine: machineID,
-		})
+		p := doc.payload()
+		fullPayloads = append(fullPayloads, p)
 	}
 	return fullPayloads, nil
 }
diff --git a/payload/persistence/env_test.go b/payload/persistence/env_test.go
@@ -56,37 +56,32 @@ func (s *envPersistenceSuite) TestListAllOkay(c *gc.C) {
 	p2 := s.newPayload("eggs")
 	s.base.setPayloads(p1, p2)
 
-	persist := NewEnvPersistence(s.base, s.base)
+	persist := NewEnvPersistence(s.base)
 
 	payloads, err := persist.ListAll()
 	c.Assert(err, jc.ErrorIsNil)
 
 	checkPayloads(c, payloads, p1, p2)
-	s.Stub.CheckCallNames(c,
-		"All",
-		"AssignedMachineID",
-	)
+	s.Stub.CheckCallNames(c, "All")
 }
 
 func (s *envPersistenceSuite) TestListAllEmpty(c *gc.C) {
 	s.base.setUnits("0")
 	s.base.setUnits("1", "a-service/0", "a-service/1")
-	persist := NewEnvPersistence(s.base, s.base)
+	persist := NewEnvPersistence(s.base)
 
 	payloads, err := persist.ListAll()
 	c.Assert(err, jc.ErrorIsNil)
 
 	c.Check(payloads, gc.HasLen, 0)
-	s.Stub.CheckCallNames(c,
-		"All",
-	)
+	s.Stub.CheckCallNames(c, "All")
 }
 
 func (s *envPersistenceSuite) TestListAllFailed(c *gc.C) {
 	failure := errors.Errorf("<failed!>")
 	s.Stub.SetErrors(failure)
 
-	persist := NewEnvPersistence(s.base, s.base)
+	persist := NewEnvPersistence(s.base)
 
 	_, err := persist.ListAll()
 
@@ -138,7 +133,7 @@ func (s *stubEnvPersistenceBase) setPayloads(payloads ...payload.FullPayloadInfo
 	for _, pl := range payloads {
 		s.setUnits(pl.Machine, pl.Unit)
 
-		doc := newPayloadDoc(pl.Unit, "0", pl.Payload)
+		doc := newPayloadDoc("0", pl)
 		s.SetDocs(doc)
 	}
 }
@@ -151,12 +146,3 @@ func (s *stubEnvPersistenceBase) setUnits(machine string, units ...string) {
 		s.unitMachines[unit] = machine
 	}
 }
-
-func (s *stubEnvPersistenceBase) AssignedMachineID(unit string) (string, error) {
-	s.stub.AddCall("AssignedMachineID", unit)
-	if err := s.stub.NextErr(); err != nil {
-		return "", errors.Trace(err)
-	}
-
-	return s.unitMachines[unit], nil
-}
diff --git a/payload/persistence/mongo.go b/payload/persistence/mongo.go
@@ -20,12 +20,13 @@ const (
 
 // TODO(ericsnow) Move the methods under their own type (payloadcollection?).
 
-func (pp Persistence) extractPayload(id string, payloadDocs map[string]payloadDoc) (*payload.Payload, bool) {
+func (pp Persistence) extractPayload(id string, payloadDocs map[string]payloadDoc) (*payload.FullPayloadInfo, bool) {
 	doc, ok := payloadDocs[id]
 	if !ok {
 		return nil, false
 	}
-	p := doc.payload(pp.unit)
+	p := doc.payload()
+	p.Unit = pp.unit
 	return &p, true
 }
 
@@ -48,7 +49,7 @@ func payloadID(unit, name string) string {
 	return fmt.Sprintf("payload#%s#%s", unit, name)
 }
 
-func (pp Persistence) newInsertPayloadOps(id string, p payload.Payload) []txn.Op {
+func (pp Persistence) newInsertPayloadOps(id string, p payload.FullPayloadInfo) []txn.Op {
 	// We must also ensure that there isn't any collision on the
 	// state-provided ID. However, that isn't something we can do in
 	// a transaction.
@@ -100,6 +101,8 @@ type payloadDoc struct {
 	UnitID string `bson:"unitid"`
 	Name   string `bson:"name"`
 
+	MachineID string `bson:"machine-id"`
+
 	// StateID is the unique ID that State gave this payload for this unit.
 	StateID string `bson:"state-id"`
 
@@ -116,15 +119,18 @@ type payloadDoc struct {
 	RawID string `bson:"rawid"`
 }
 
-func (d payloadDoc) payload(unit string) payload.Payload {
+func (d payloadDoc) payload() payload.FullPayloadInfo {
 	labels := make([]string, len(d.Labels))
 	copy(labels, d.Labels)
-	p := payload.Payload{
-		PayloadClass: d.definition(),
-		ID:           d.RawID,
-		Status:       d.State,
-		Labels:       labels,
-		Unit:         unit,
+	p := payload.FullPayloadInfo{
+		Payload: payload.Payload{
+			PayloadClass: d.definition(),
+			ID:           d.RawID,
+			Status:       d.State,
+			Labels:       labels,
+			Unit:         d.UnitID,
+		},
+		Machine: d.MachineID,
 	}
 	return p
 }
@@ -147,12 +153,13 @@ func (d payloadDoc) match(name, rawID string) bool {
 	return true
 }
 
-func (pp Persistence) newPayloadDoc(stID string, p payload.Payload) *payloadDoc {
-	return newPayloadDoc(pp.unit, stID, p)
+func (pp Persistence) newPayloadDoc(stID string, p payload.FullPayloadInfo) *payloadDoc {
+	p.Unit = pp.unit
+	return newPayloadDoc(stID, p)
 }
 
-func newPayloadDoc(unit, stID string, p payload.Payload) *payloadDoc {
-	id := payloadID(unit, p.Name)
+func newPayloadDoc(stID string, p payload.FullPayloadInfo) *payloadDoc {
+	id := payloadID(p.Unit, p.Name)
 
 	definition := p.PayloadClass
 
@@ -161,9 +168,11 @@ func newPayloadDoc(unit, stID string, p payload.Payload) *payloadDoc {
 
 	return &payloadDoc{
 		DocID:  id,
-		UnitID: unit,
+		UnitID: p.Unit,
 		Name:   definition.Name,
 
+		MachineID: p.Machine,
+
 		StateID: stID,
 
 		Type: definition.Type,

diff --git a/payload/persistence/unit.go b/payload/persistence/unit.go
@@ -53,8 +53,9 @@ func NewPersistence(db PersistenceBase, unit string) *Persistence {
 // Track adds records for the payload to persistence. If the payload
 // is already there then false gets returned (true if inserted).
 // Existing records are not checked for consistency.
-func (pp Persistence) Track(id string, pl payload.Payload) error {
-	logger.Tracef("insertng %q - %#v", id, pl)
+func (pp Persistence) Track(id string, pl payload.FullPayloadInfo) error {
+	pl.Unit = pp.unit
+	logger.Tracef("inserting %q - %#v", id, pl)
 
 	docs, err := pp.payloads([]string{id})
 	if err != nil {
@@ -121,15 +122,15 @@ func (pp Persistence) SetStatus(id, status string) error {
 // List builds the list of payloads found in persistence which match
 // the provided IDs. The lists of IDs with missing records is also
 // returned.
-func (pp Persistence) List(ids ...string) ([]payload.Payload, []string, error) {
+func (pp Persistence) List(ids ...string) ([]payload.FullPayloadInfo, []string, error) {
 	// TODO(ericsnow) Ensure that the unit is Alive?
 
 	docs, err := pp.payloads(ids)
 	if err != nil {
 		return nil, nil, errors.Trace(err)
 	}
 
-	var results []payload.Payload
+	var results []payload.FullPayloadInfo
 	var missing []string
 	for _, id := range ids {
 		p, ok := pp.extractPayload(id, docs)
@@ -144,15 +145,15 @@ func (pp Persistence) List(ids ...string) ([]payload.Payload, []string, error) {
 
 // ListAll builds the list of all payloads found in persistence.
 // Inconsistent records result in errors.NotValid.
-func (pp Persistence) ListAll() ([]payload.Payload, error) {
+func (pp Persistence) ListAll() ([]payload.FullPayloadInfo, error) {
 	// TODO(ericsnow) Ensure that the unit is Alive?
 
 	docs, err := pp.allPayloads()
 	if err != nil {
 		return nil, errors.Trace(err)
 	}
 
-	var results []payload.Payload
+	var results []payload.FullPayloadInfo
 	for id := range docs {
 		p, _ := pp.extractPayload(id, docs)
 		results = append(results, *p)