environs/state: simplify watcher implementation

state/watcher.go

 package state
 
 import (
         "launchpad.net/goyaml"
         "launchpad.net/juju-core/log"
         "launchpad.net/juju-core/state/watcher"
-        "launchpad.net/tomb"
 )
 
 // contentWatcher holds behaviour common to all ContentWatcher clients in
 // the state package.
 type contentWatcher struct {
         st      *State
-        tomb    tomb.Tomb
         path    string
         updated bool
+        err     error
+        handler contentHandler
+        stop    <-chan struct{}
 }
 
-func newContentWatcher(st *State, path string) contentWatcher {
-        return contentWatcher{st: st, path: path}
+func newContentWatcher(st *State, path string, stop <-chan struct{}) contentWatcher {
+        return contentWatcher{
+                st:   st,
+                path: path,
+                stop: stop,
+                err:  watcher.ErrNotStopped,
+        }
 }
 
 // contentHandler must be implemented by watchers that intend to make use
 // of contentWatcher.
 type contentHandler interface {
+        // update is called whenever the content changes.
         update(watcher.ContentChange) error
+
+        // wait reads values from the change channel until it is closed.
+        wait()
+
+        // done cleans up and closes the change channel.
         done()
 }
 
-// loop handles the common tasks of receiving changes from a watcher.ContentWatcher,
+// watcher handles the common tasks of receiving changes from a watcher.ContentWatcher,
 // and dispatching them to the contentHandler's update method.
-func (w *contentWatcher) loop(handler contentHandler) {
-        defer w.tomb.Done()
-        defer handler.done()
-        cw := watcher.NewContentWatcher(w.st.zk, w.path)
-        defer watcher.Stop(cw, &w.tomb)
-        for {
-                select {
-                case <-w.tomb.Dying():
-                        return
-                case ch, ok := <-cw.Changes():
-                        if !ok {
-                                w.tomb.Kill(watcher.MustErr(cw))
-                                return
-                        }
+func (w *contentWatcher) runWatcher(handler contentHandler) {
+        w.handler = handler
+        cw := watcher.NewContentWatcher(w.st.zk, w.path, w.stop)
+        go func() {
+                defer handler.done()
+                for ch := range cw.Changes() {
                         if err := handler.update(ch); err != nil {
-                                w.tomb.Kill(err)
+                                w.err = err
                                 return
                         }
                         w.updated = true
                 }
-        }
+                w.err = cw.Wait()
+        }()
 }
 
-// Stop stops the watcher and returns any errors encountered while watching.
-func (w *contentWatcher) Stop() error {
-        w.tomb.Kill(nil)
-        return w.tomb.Wait()
-}
-
-// Err returns any error encountered while stopping the watcher, or
-// tome.ErrStillAlive if the watcher is still running.
-func (w *contentWatcher) Err() error {
-        return w.tomb.Err()
+// Wait reads from the Changes channel until it's closed,
+// then returns any error encountered while the watcher
+// was running.
+func (w *contentWatcher) Wait() error {
+        w.handler.wait()
+        return w.err
 }
 
 // ConfigWatcher observes changes to any configuration node.
 type ConfigWatcher struct {
         contentWatcher
         changeChan chan *ConfigNode
 }
 
 // newConfigWatcher creates and starts a new config watcher for
-// the given path.
-func newConfigWatcher(st *State, path string) *ConfigWatcher {
+// the given path. After stop is closed, the watcher will terminate and
+// close the Changes channel.
+func newConfigWatcher(st *State, path string, stop <-chan struct{}) *ConfigWatcher {
         w := &ConfigWatcher{
-                contentWatcher: newContentWatcher(st, path),
+                contentWatcher: newContentWatcher(st, path, stop),
                 changeChan:     make(chan *ConfigNode),
         }
-        go w.loop(w)
+        w.runWatcher(w)
         return w
 }
 
 // Changes returns a channel that will receive the new
 // *ConfigNode when a change is detected. Note that multiple
 // changes may be observed as a single event in the channel.
 // The first event on the channel holds the initial state
 // as returned by Service.Config.
 func (w *ConfigWatcher) Changes() <-chan *ConfigNode {
         return w.changeChan
 }
 
 func (w *ConfigWatcher) update(change watcher.ContentChange) error {
         // A non-existent node is treated as an empty node.
         configNode, err := parseConfigNode(w.st.zk, w.path, change.Content)
         if err != nil {
                 return err
         }
-        select {
-        case <-w.tomb.Dying():
-                return tomb.ErrDying
-        case w.changeChan <- configNode:
+        w.changeChan <- configNode
+        return nil
+}
+
+func (w *ConfigWatcher) wait() {
+        for _ = range w.changeChan {
         }
-        return nil
 }
 
 func (w *ConfigWatcher) done() {
         close(w.changeChan)
 }
 
 // FlagWatcher observes whether a given flag is on or off.
 type FlagWatcher struct {
         contentWatcher
         changeChan chan bool
         exists     bool
 }
 
 // newFlagWatcher creates and starts a new flag watcher for
 // the given path.
-func newFlagWatcher(st *State, path string) *FlagWatcher {
+func newFlagWatcher(st *State, path string, stop <-chan struct{}) *FlagWatcher {
         w := &FlagWatcher{
-                contentWatcher: newContentWatcher(st, path),
+                contentWatcher: newContentWatcher(st, path, stop),
                 changeChan:     make(chan bool),
         }
-        go w.loop(w)
+        w.runWatcher(w)
         return w
 }
 
 // Changes returns a channel that will receive true when a
 // flag is set and false if it is cleared. Note that multiple
 // changes may be observed as a single event in the channel.
 // The first event on the channel holds the initial state.
 func (w *FlagWatcher) Changes() <-chan bool {
         return w.changeChan
 }
 
 func (w *FlagWatcher) update(change watcher.ContentChange) error {
         if w.updated && change.Exists == w.exists {
                 return nil
         }
-        select {
-        case <-w.tomb.Dying():
-                return tomb.ErrDying
-        case w.changeChan <- change.Exists:
-                w.exists = change.Exists
+        w.changeChan <- change.Exists
+        w.exists = change.Exists
+        return nil
+}
+
+func (w *FlagWatcher) wait() {
+        for _ = range w.changeChan {
         }
-        return nil
 }
 
 func (w *FlagWatcher) done() {
         close(w.changeChan)
 }
 
 // NeedsUpgradeWatcher observes changes to a unit's upgrade flag.
 type NeedsUpgradeWatcher struct {
         contentWatcher
         changeChan chan NeedsUpgrade
 }
 
 // newNeedsUpgradeWatcher creates and starts a new resolved flag node
 // watcher for the given path.
-func newNeedsUpgradeWatcher(st *State, path string) *NeedsUpgradeWatcher {
+func newNeedsUpgradeWatcher(st *State, path string, stop <-chan struct{}) *NeedsUpgradeWatcher {
         w := &NeedsUpgradeWatcher{
-                contentWatcher: newContentWatcher(st, path),
+                contentWatcher: newContentWatcher(st, path, stop),
                 changeChan:     make(chan NeedsUpgrade),
         }
-        go w.loop(w)
+        w.runWatcher(w)
         return w
 }
 
 // Changes returns a channel that will receive notifications
 // about upgrades for the unit. Note that multiple changes
 // may be observed as a single event in the channel.
 // The first event on the channel holds the initial
 // state as returned by Unit.NeedsUpgrade.
 func (w *NeedsUpgradeWatcher) Changes() <-chan NeedsUpgrade {
         return w.changeChan
 }
 
 func (w *NeedsUpgradeWatcher) update(change watcher.ContentChange) error {
         var needsUpgrade NeedsUpgrade
         if change.Exists {
                 needsUpgrade.Upgrade = true
                 var setting needsUpgradeNode
                 if err := goyaml.Unmarshal([]byte(change.Content), &setting); err != nil {
                         return err
                 }
                 needsUpgrade.Force = setting.Force
         }
-        select {
-        case <-w.tomb.Dying():
-                return tomb.ErrDying
-        case w.changeChan <- needsUpgrade:
+        w.changeChan <- needsUpgrade
+        return nil
+}
+
+func (w *NeedsUpgradeWatcher) wait() {
+        for _ = range w.changeChan {
         }
-        return nil
 }
 
 func (w *NeedsUpgradeWatcher) done() {
         close(w.changeChan)
 }
 
 // ResolvedWatcher observes changes to a unit's resolved
 // mode. See SetResolved for details.
 type ResolvedWatcher struct {
         contentWatcher
         changeChan chan ResolvedMode
 }
 
 // newResolvedWatcher returns a new ResolvedWatcher watching path.
-func newResolvedWatcher(st *State, path string) *ResolvedWatcher {
+func newResolvedWatcher(st *State, path string, stop <-chan struct{}) *ResolvedWatcher {
         w := &ResolvedWatcher{
-                contentWatcher: newContentWatcher(st, path),
+                contentWatcher: newContentWatcher(st, path, stop),
                 changeChan:     make(chan ResolvedMode),
         }
-        go w.loop(w)
+        w.runWatcher(w)
         return w
 }
 
 // Changes returns a channel that will receive the new
 // resolved mode when a change is detected. Note that multiple
 // changes may be observed as a single event in the channel.
 // The first event on the channel holds the initial
 // state as returned by Unit.Resolved.
 func (w *ResolvedWatcher) Changes() <-chan ResolvedMode {
         return w.changeChan
 }
 
 func (w *ResolvedWatcher) update(change watcher.ContentChange) error {
         mode := ResolvedNone
         if change.Exists {
                 var err error
                 mode, err = parseResolvedMode(change.Content)
                 if err != nil {
                         return err
                 }
         }
-        select {
-        case <-w.tomb.Dying():
-                return tomb.ErrDying
-        case w.changeChan <- mode:
+        w.changeChan <- mode
+        return nil
+}
+
+func (w *ResolvedWatcher) wait() {
+        for _ = range w.changeChan {
         }
-        return nil
 }
 
 func (w *ResolvedWatcher) done() {
         close(w.changeChan)
 }
 
 // PortsWatcher observes changes to a unit's open ports.
 // See OpenPort for details.
 type PortsWatcher struct {
         contentWatcher
         changeChan chan []Port
 }
 
 // newPortsWatcher creates and starts a new ports node
 // watcher for the given path.
-func newPortsWatcher(st *State, path string) *PortsWatcher {
+func newPortsWatcher(st *State, path string, stop <-chan struct{}) *PortsWatcher {
         w := &PortsWatcher{
-                contentWatcher: newContentWatcher(st, path),
+                contentWatcher: newContentWatcher(st, path, stop),
                 changeChan:     make(chan []Port),
         }
-        go w.loop(w)
+        w.runWatcher(w)
         return w
 }
 
 // Changes returns a channel that will receive the actual
 // open ports when a change is detected. Note that multiple
 // changes may be observed as a single event in the channel.
 // The first event on the channel holds the initial
 // state as returned by Unit.OpenPorts.
 func (w *PortsWatcher) Changes() <-chan []Port {
         return w.changeChan
 }
 
 func (w *PortsWatcher) update(change watcher.ContentChange) error {
         var ports openPortsNode
         if err := goyaml.Unmarshal([]byte(change.Content), &ports); err != nil {
                 return err
         }
-        select {
-        case <-w.tomb.Dying():
-                return tomb.ErrDying
-        case w.changeChan <- ports.Open:
+        w.changeChan <- ports.Open
+        return nil
+}
+
+func (w *PortsWatcher) wait() {
+        for _ = range w.changeChan {
         }
-        return nil
 }
 
 func (w *PortsWatcher) done() {
         close(w.changeChan)
 }
 
 // MachinesWatcher notifies about machines being added or removed
 // from the environment.
 type MachinesWatcher struct {
         contentWatcher
         changeChan       chan *MachinesChange
         watcher          *watcher.ContentWatcher
         knownMachineKeys []string
 }
 
 // MachinesChange contains information about
 // machines that have been added or deleted.
 type MachinesChange struct {
         Added   []*Machine
         Removed []*Machine
 }
 
 // newMachinesWatcher creates and starts a new watcher for changes to
 // the set of machines known to the topology.
-func newMachinesWatcher(st *State) *MachinesWatcher {
+func newMachinesWatcher(st *State, stop <-chan struct{}) *MachinesWatcher {
         w := &MachinesWatcher{
-                contentWatcher: newContentWatcher(st, zkTopologyPath),
+                contentWatcher: newContentWatcher(st, zkTopologyPath, stop),
                 changeChan:     make(chan *MachinesChange),
         }
-        go w.loop(w)
+        w.runWatcher(w)
         return w
 }
 
 // Changes returns a channel that will receive changes when machines are
 // added or deleted.  The Added field in the first event on the channel
 // holds the initial state as returned by State.AllMachines.
 func (w *MachinesWatcher) Changes() <-chan *MachinesChange {
         return w.changeChan
 }
 
 func (w *MachinesWatcher) update(change watcher.ContentChange) error {
         topology, err := parseTopology(change.Content)
         if err != nil {
                 return err
         }
         currentMachineKeys := topology.MachineKeys()
         added := diff(currentMachineKeys, w.knownMachineKeys)
         removed := diff(w.knownMachineKeys, currentMachineKeys)
         w.knownMachineKeys = currentMachineKeys
         if w.updated && len(added) == 0 && len(removed) == 0 {
                 return nil
         }
         mc := &MachinesChange{}
         for _, m := range added {
                 mc.Added = append(mc.Added, &Machine{w.st, m})
         }
         for _, m := range removed {
                 mc.Removed = append(mc.Removed, &Machine{w.st, m})
         }
-        select {
-        case <-w.tomb.Dying():
-                return tomb.ErrDying
-        case w.changeChan <- mc:
+        w.changeChan <- mc
+        return nil
+}
+
+func (w *MachinesWatcher) wait() {
+        for _ = range w.changeChan {
         }
-        return nil
 }
 
 func (w *MachinesWatcher) done() {
         close(w.changeChan)
 }
 
 type MachineUnitsWatcher struct {
         contentWatcher
         machine       *Machine
         changeChan    chan *MachineUnitsChange
         knownUnitKeys []string
         knownUnits    map[string]*Unit
 }
 
 type MachineUnitsChange struct {
         Added   []*Unit
         Removed []*Unit
 }
 
 // newMachineUnitsWatcher creates and starts a new machine units watcher.
-func newMachineUnitsWatcher(m *Machine) *MachineUnitsWatcher {
+func newMachineUnitsWatcher(m *Machine, stop <-chan struct{}) *MachineUnitsWatcher {
         w := &MachineUnitsWatcher{
-                contentWatcher: newContentWatcher(m.st, zkTopologyPath),
+                contentWatcher: newContentWatcher(m.st, zkTopologyPath, stop),
                 machine:        m,
                 changeChan:     make(chan *MachineUnitsChange),
                 knownUnits:     make(map[string]*Unit),
         }
-        go w.loop(w)
+        w.runWatcher(w)
         return w
 }
 
 // Changes returns a channel that will receive changes when
 // units are assigned or unassigned from a machine.
 // The Added field in the first event on the channel holds the initial
 // state as returned by machine.Units.
 func (w *MachineUnitsWatcher) Changes() <-chan *MachineUnitsChange {
         return w.changeChan
 }
@@ skipping 21 matching lines @@
         }
         for _, ukey := range added {
                 unit, err := w.st.unitFromKey(topology, ukey)
                 if err != nil {
                         log.Printf("inconsistent topology: %v", err)
                         continue
                 }
                 w.knownUnits[ukey] = unit
                 uc.Added = append(uc.Added, unit)
         }
-        select {
-        case <-w.tomb.Dying():
-                return tomb.ErrDying
-        case w.changeChan <- uc:
+        w.changeChan <- uc
+        return nil
+}
+
+func (w *MachineUnitsWatcher) wait() {
+        for _ = range w.changeChan {
         }
-        return nil
 }
 
 func (w *MachineUnitsWatcher) done() {
         close(w.changeChan)
 }
 
 // ServicesWatcher observes the addition and removal of services.
 type ServicesWatcher struct {
         contentWatcher
         knownServices    map[string]*Service
         knownServiceKeys []string
         changeChan       chan *ServicesChange
 }
 
 // ServicesChange holds services that were added or removed
 // from the environment.
 type ServicesChange struct {
         Added   []*Service
         Removed []*Service
 }
 
 // newServicesWatcher returns a new ServicesWatcher.
-func newServicesWatcher(st *State) *ServicesWatcher {
+func newServicesWatcher(st *State, stop <-chan struct{}) *ServicesWatcher {
         w := &ServicesWatcher{
-                contentWatcher: newContentWatcher(st, zkTopologyPath),
+                contentWatcher: newContentWatcher(st, zkTopologyPath, stop),
                 knownServices:  make(map[string]*Service),
                 changeChan:     make(chan *ServicesChange),
         }
-        go w.loop(w)
+        w.runWatcher(w)
         return w
 }
 
 // Changes returns a channel that will receive a notification when services
 // are added to or removed from the state. The Added field in
 // the first event on the channel holds the initial state as would be·
 // returned by Service.AllServices.
 func (w *ServicesWatcher) Changes() <-chan *ServicesChange {
         return w.changeChan
 }
@@ skipping 23 matching lines @@
                         continue
                 }
                 service, err := w.st.Service(serviceName)
                 if err != nil {
                         log.Printf("can't read service %q: %v", serviceName, err)
                         continue
                 }
                 w.knownServices[serviceKey] = service
                 servicesChange.Added = append(servicesChange.Added, service)
         }
-        select {
-        case <-w.tomb.Dying():
-                return tomb.ErrDying
-        case w.changeChan <- servicesChange:
+        w.changeChan <- servicesChange
+        return nil
+}
+
+func (w *ServicesWatcher) wait() {
+        for _ = range w.changeChan {
         }
-        return nil
 }
 
 func (w *ServicesWatcher) done() {
         close(w.changeChan)
 }
 
 // ServiceUnitsWatcher observes the addition and removal
 // of units to and from a service.
 type ServiceUnitsWatcher struct {
         contentWatcher
         serviceKey    string
         knownUnits    map[string]*Unit
         knownUnitKeys []string
         changeChan    chan *ServiceUnitsChange
 }
 
 // ServiceUnitsChange contains information about
 // units that have been added to or removed from
 // services.
 type ServiceUnitsChange struct {
         Added   []*Unit
         Removed []*Unit
 }
 
 // newServiceUnitsWatcher creates and starts a new watcher
 // for service unit changes.
-func newServiceUnitsWatcher(service *Service) *ServiceUnitsWatcher {
+func newServiceUnitsWatcher(service *Service, stop <-chan struct{}) *ServiceUnitsWatcher {
         w := &ServiceUnitsWatcher{
-                contentWatcher: newContentWatcher(service.st, zkTopologyPath),
+                contentWatcher: newContentWatcher(service.st, zkTopologyPath, stop),
                 serviceKey:     service.key,
                 knownUnits:     make(map[string]*Unit),
                 changeChan:     make(chan *ServiceUnitsChange),
         }
-        go w.loop(w)
+        w.runWatcher(w)
         return w
 }
 
 // Changes returns a channel that will receive changes when units
 // are added to or removed from the service. The Added field in
 // the first event on the channel holds the initial state as returned
 // by Service.AllUnits.
 func (w *ServiceUnitsWatcher) Changes() <-chan *ServiceUnitsChange {
         return w.changeChan
 }
@@ skipping 21 matching lines @@
         }
         for _, unitKey := range added {
                 unit, err := w.st.unitFromKey(topology, unitKey)
                 if err != nil {
                         log.Printf("can't read unit %q: %v", unitKey, err)
                         continue
                 }
                 w.knownUnits[unitKey] = unit
                 serviceUnitsChange.Added = append(serviceUnitsChange.Added, unit)
         }
-        select {
-        case <-w.tomb.Dying():
-                return tomb.ErrDying
-        case w.changeChan <- serviceUnitsChange:
+        w.changeChan <- serviceUnitsChange
+        return nil
+}
+
+func (w *ServiceUnitsWatcher) wait() {
+        for _ = range w.changeChan {
         }
-        return nil
 }
 
 func (w *ServiceUnitsWatcher) done() {
         close(w.changeChan)
 }
 
 // ServiceRelationsWatcher notifies of changes to a service's relations.
 type ServiceRelationsWatcher struct {
         contentWatcher
         changeChan chan RelationsChange
         service    *Service
         current    map[string]*Relation
 }
 
 type RelationsChange struct {
         Added, Removed []*Relation
 }
 
 // newServiceRelationsWatcher creates and starts a new service relations watcher.
-func newServiceRelationsWatcher(s *Service) *ServiceRelationsWatcher {
+func newServiceRelationsWatcher(s *Service, stop <-chan struct{}) *ServiceRelationsWatcher {
         w := &ServiceRelationsWatcher{
-                contentWatcher: newContentWatcher(s.st, zkTopologyPath),
+                contentWatcher: newContentWatcher(s.st, zkTopologyPath, stop),
                 changeChan:     make(chan RelationsChange),
                 service:        s,
                 current:        make(map[string]*Relation),
         }
-        go w.loop(w)
+        w.runWatcher(w)
         return w
 }
 
 // Changes returns a channel that will receive changes when
 // the service enters and leaves relations.
 // The Added field in the first event on the channel holds the initial
 // state, corresponding to that returned by service.Relations.
 func (w *ServiceRelationsWatcher) Changes() <-chan RelationsChange {
         return w.changeChan
 }
@@ skipping 18 matching lines @@
                 }
         }
         for key, rel := range w.current {
                 if latest[key] == nil {
                         ch.Removed = append(ch.Removed, rel)
                 }
         }
         if w.updated && len(ch.Added) == 0 && len(ch.Removed) == 0 {
                 return nil
         }
-        select {
-        case <-w.tomb.Dying():
-                return tomb.ErrDying
-        case w.changeChan <- ch:
-                w.current = latest
+        w.changeChan <- ch
+        w.current = latest
+        return nil
+}
+
+func (w *ServiceRelationsWatcher) wait() {
+        for _ = range w.changeChan {
         }
-        return nil
 }
 
 func (w *ServiceRelationsWatcher) done() {
         close(w.changeChan)
 }