mstate: added unit status

mstate/unit.go

 package mstate
 
 import (
         "errors"
         "fmt"
         "labix.org/v2/mgo"
         "labix.org/v2/mgo/txn"
+        "launchpad.net/juju-core/mstate/presence"
         "launchpad.net/juju-core/trivial"
+        "time"
 )
 
 // ResolvedMode describes the way state transition errors·
 // are resolved.·
 type ResolvedMode int
 
 const (
         ResolvedNone ResolvedMode = iota
         ResolvedRetryHooks
         ResolvedNoHooks
         nResolvedModes
 )
 
 // AssignmentPolicy controls what machine a unit will be assigned to.
 type AssignmentPolicy string
 
 const (
         // AssignLocal indicates that all service units should be assigned·
         // to machine 0.
         AssignLocal AssignmentPolicy = "local"
         // AssignUnused indicates that every service unit should be assigned
         // to a dedicated machine, and that new machines should be launched
         // if required.
         AssignUnused AssignmentPolicy = "unused"
 )
 
+// UnitStatus represents the status of the unit agent.
+type UnitStatus string
+
+const (
+        UnitPending   UnitStatus = "pending"   // Agent hasn't started
+        UnitInstalled UnitStatus = "installed" // Agent has run the installed hook
+        UnitStarted   UnitStatus = "started"   // Agent is running properly
+        UnitStopped   UnitStatus = "stopped"   // Agent has stopped running on request
+        UnitError     UnitStatus = "error"     // Agent is waiting in an error state
+        UnitDown      UnitStatus = "down"      // Agent is down or not communicating
+)
+
 // NeedsUpgrade describes if a unit needs an
 // upgrade and if this is forced.
 type NeedsUpgrade struct {
         Upgrade bool
         Force   bool
 }
 
 // Port identifies a network port number for a particular protocol.
 type Port struct {
         Protocol string `yaml:"proto"`
@@ skipping 41 matching lines @@
 // String returns the unit as string.
 func (u *Unit) String() string {
         return u.doc.Name
 }
 
 // Name returns the unit name.
 func (u *Unit) Name() string {
         return u.doc.Name
 }
 
+// globalKey returns the global database key for the unit.
+func (u *Unit) globalKey() string {
+        return "u#" + u.doc.Name
+}
+
 // Life returns whether the unit is Alive, Dying or Dead.
 func (u *Unit) Life() Life {
         return u.doc.Life
 }
 
 // Kill sets the unit lifecycle to Dying if it is Alive.
 // It does nothing otherwise.
 func (u *Unit) Kill() error {
         err := ensureLife(u.st, u.st.units, u.doc.Name, Dying, "unit")
         if err != nil {
@@ skipping 42 matching lines @@
 }
 
 func (u *Unit) Refresh() error {
         err := u.st.units.FindId(u.doc.Name).One(&u.doc)
         if err != nil {
                 return fmt.Errorf("cannot refresh unit %q: %v", u, err)
         }
         return nil
 }
 
+// Status returns the status of the unit's agent.
+func (u *Unit) Status() (s UnitStatus, info string, err error) {
+        config, err := u.Config()
+        if err != nil {
+                return "", "", fmt.Errorf("cannot read status of unit %q: %v", u, err)
+        }
+        raw, found := config.Get("status")
+        if !found {
+                return UnitPending, "", nil
+        }
+        s = UnitStatus(raw.(string))
+        switch s {
+        case UnitError:
+                // We always expect an info if status is 'error'.
+                raw, found = config.Get("status-info")
+                if !found {
+                        panic("no status-info found for unit error")
+                }
+                return s, raw.(string), nil
+        case UnitStopped:
+                return UnitStopped, "", nil
+        }
+        alive := u.AgentAlive()
+        if !alive {
+                s = UnitDown
+        }
+        return s, "", nil
+}
+
+// SetStatus sets the status of the unit.
+func (u *Unit) SetStatus(status UnitStatus, info string) error {
+        if status == UnitPending {
+                panic("unit status must not be set to pending")
+        }
+        config, err := u.Config()
+        if err != nil {
+                return err
+        }
+        config.Set("status", status)
+        config.Set("status-info", info)
+        _, err = config.Write()
+        if err != nil {
+                return fmt.Errorf("cannot set status of unit %q: %v", u, err)
+        }
+        return nil
+}
+
+// AgentAlive returns whether the respective remote agent is alive.
+func (u *Unit) AgentAlive() bool {
+        return u.st.presencew.Alive(u.globalKey())
+}
+
+// WaitAgentAlive blocks until the respective agent is alive.
+func (u *Unit) WaitAgentAlive(timeout time.Duration) error {
+        ch := make(chan presence.Change)
+        u.st.presencew.Add(u.globalKey(), ch)
+        defer u.st.presencew.Remove(u.globalKey(), ch)
+        // Initial check.
+        select {
+        case change := <-ch:
+                if change.Alive {
+                        return nil
+                }
+        case <-time.After(timeout):
+                return fmt.Errorf("waiting for agent of unit %q: still not alive after timeout", u)
+        }
+        // Hasn't been alive, so now wait for change.
+        select {
+        case change := <-ch:
+                if change.Alive {
+                        return nil
+                }
+                panic(fmt.Sprintf("presence reported dead status twice in a row for unit %q", u))
+        case <-time.After(timeout):
+                return fmt.Errorf("waiting for agent of unit %q: still not alive after timeout", u)
+        }
+        panic("unreachable")
+}
+
+// SetAgentAlive signals that the agent for unit u is alive.·
+// It returns the started pinger.
+func (u *Unit) SetAgentAlive() (*presence.Pinger, error) {
+        p := presence.NewPinger(u.st.presence, u.globalKey())
+        err := p.Start()
+        if err != nil {
+                return nil, err
+        }
+        return p, nil
+}
+
 // AssignedMachineId returns the id of the assigned machine.
 func (u *Unit) AssignedMachineId() (id int, err error) {
         defer trivial.ErrorContextf(&err, "cannot get machine id of unit %q", u)
         if u.IsPrincipal() {
                 if u.doc.MachineId == nil {
                         return 0, errors.New("unit not assigned to machine")
                 }
                 return *u.doc.MachineId, nil
         }
         pudoc := unitDoc{}
@@ skipping 184 matching lines @@
 func (u *Unit) ClearNeedsUpgrade() error {
         change := D{{"$set", D{{"needsupgrade", nil}}}}
         sel := D{{"_id", u.doc.Name}}
         err := u.st.units.Update(sel, change)
         if err != nil {
                 return fmt.Errorf("upgrade notification for unit %q cannot be reset: %v", u, err)
         }
         u.doc.NeedsUpgrade = nil
         return nil
 }
+
+// Config returns the configuration node for the unit.
+func (u *Unit) Config() (config *ConfigNode, err error) {
+        config, err = readConfigNode(u.st, u.globalKey())
+        if err != nil {
+                return nil, fmt.Errorf("cannot get configuration of unit %q: %v", u, err)
+        }
+        return config, nil
+}