<commit_before>package firewaller

import (
	"fmt"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/log"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/state/watcher"
	"launchpad.net/juju-core/worker"
	"launchpad.net/tomb"
)

// Firewaller watches the state for ports opened or closed
// and reflects those changes onto the backing environment.
type Firewaller struct {
	tomb            tomb.Tomb
	st              *state.State
	environ         environs.Environ
	environWatcher  *state.EnvironConfigWatcher
	machinesWatcher *state.MachinesWatcher
	machineds       map[int]*machineData
	unitsChange     chan *unitsChange
	unitds          map[string]*unitData
	portsChange     chan *portsChange
	serviceds       map[string]*serviceData
	exposedChange   chan *exposedChange
}

// NewFirewaller returns a new Firewaller.
func NewFirewaller(st *state.State) *Firewaller {
	fw := &Firewaller{
		st:              st,
		environWatcher:  st.WatchEnvironConfig(),
		machinesWatcher: st.WatchMachines(),
		machineds:       make(map[int]*machineData),
		unitsChange:     make(chan *unitsChange),
		unitds:          make(map[string]*unitData),
		portsChange:     make(chan *portsChange),
		serviceds:       make(map[string]*serviceData),
		exposedChange:   make(chan *exposedChange),
	}
	go func() {
		defer fw.tomb.Done()
		fw.tomb.Kill(fw.loop())
	}()
	return fw
}

func (fw *Firewaller) loop() error {
	defer fw.stopWatchers()

	var err error
	fw.environ, err = worker.WaitForEnviron(fw.environWatcher, fw.tomb.Dying())
	if err != nil {
		return err
	}
	for {
		select {
		case <-fw.tomb.Dying():
			return tomb.ErrDying
		case change, ok := <-fw.environWatcher.Changes():
			if !ok {
				return watcher.MustErr(fw.environWatcher)
			}
			if err := fw.environ.SetConfig(change); err != nil {
				log.Printf("firewaller loaded invalid environment configuration: %v", err)
			}
		case change, ok := <-fw.machinesWatcher.Changes():
			if !ok {
				return watcher.MustErr(fw.machinesWatcher)
			}
			for _, id := range change.Dead {
				machined := fw.machineds[id]
				delete(fw.machineds, id)
				if err := machined.Stop(); err != nil {
					log.Printf("machine %d watcher returned error when stopping: %v", id, err)
				}
				log.Debugf("firewaller: stopped watching machine %d", id)
			}
			for _, id := range change.Alive {
				fw.machineds[id] = newMachineData(id, fw)
				log.Debugf("firewaller: started watching machine %d", id)
			}
		case change := <-fw.unitsChange:
			changed := []*unitData{}
			for _, unit := range change.Removed {
				unitd, ok := fw.unitds[unit.Name()]
				if !ok {
					panic("trying to remove unit that was not added")
				}
				delete(fw.unitds, unit.Name())
				delete(unitd.machined.unitds, unit.Name())
				delete(unitd.serviced.unitds, unit.Name())
				changed = append(changed, unitd)
				if err := unitd.Stop(); err != nil {
					log.Printf("unit watcher %q returned error when stopping: %v", unit.Name(), err)
				}
				log.Debugf("firewaller: stopped watching unit %s", unit.Name())
			}
			for _, unit := range change.Added {
				unitd := newUnitData(unit, fw)
				fw.unitds[unit.Name()] = unitd
				machineId, err := unit.AssignedMachineId()
				if err != nil {
					fw.tomb.Kill(err)
				}
				if fw.machineds[machineId] == nil {
					panic("machine of added unit is not watched")
				}
				unitd.machined = fw.machineds[machineId]
				unitd.machined.unitds[unit.Name()] = unitd
				if fw.serviceds[unit.ServiceName()] == nil {
					service, err := fw.st.Service(unit.ServiceName())
					if err != nil {
						return err
					}
					fw.serviceds[unit.ServiceName()] = newServiceData(service, fw)
				}
				unitd.serviced = fw.serviceds[unit.ServiceName()]
				unitd.serviced.unitds[unit.ServiceName()] = unitd
				fw.serviceds[unit.ServiceName()].unitds[unit.Name()] = unitd
				changed = append(changed, unitd)
				log.Debugf("firewaller: started watching unit %s", unit.Name())
			}
			if err := fw.flushUnits(changed); err != nil {
				return fmt.Errorf("cannot change firewall ports: %v", err)
			}
		case change := <-fw.portsChange:
			change.unitd.ports = change.ports
			if err := fw.flushUnits([]*unitData{change.unitd}); err != nil {
				return fmt.Errorf("cannot change firewall ports: %v", err)
			}
		case change := <-fw.exposedChange:
			change.serviced.exposed = change.exposed
			unitds := []*unitData{}
			for _, unitd := range change.serviced.unitds {
				unitds = append(unitds, unitd)
			}
			if err := fw.flushUnits(unitds); err != nil {
				return fmt.Errorf("cannot change firewall ports: %v", err)
			}
		}
	}
	panic("not reached")
}

// flushUnits opens and closes ports for the passed unit data.
func (fw *Firewaller) flushUnits(unitds []*unitData) error {
	machineds := map[int]*machineData{}
	for _, unitd := range unitds {
		machineds[unitd.machined.id] = unitd.machined
	}
	for _, machined := range machineds {
		if err := fw.flushMachine(machined); err != nil {
			return err
		}
	}
	return nil
}

// flushMachine opens and closes ports for the passed machine.
func (fw *Firewaller) flushMachine(machined *machineData) error {
	// Gather ports to open and close.
	ports := map[state.Port]bool{}
	for _, unitd := range machined.unitds {
		if unitd.serviced.exposed {
			for _, port := range unitd.ports {
				ports[port] = true
			}
		}
	}
	want := []state.Port{}
	for port := range ports {
		want = append(want, port)
	}
	toOpen := diff(want, machined.ports)
	toClose := diff(machined.ports, want)
	machined.ports = want

	// If there's nothing to do, do nothing.
	// This is important because when a machine is first created,
	// it will have no instance id but also no open ports -
	// InstanceId will fail but we don't care.
	if len(toOpen) == 0 && len(toClose) == 0 {
		return nil
	}
	// Open and close the ports.
	m, err := machined.machine()
	if _, ok := err.(*state.NotFoundError); ok {
		return nil
	}
	if err != nil {
		return err
	}
	instanceId, err := m.InstanceId()
	if err != nil {
		return err
	}
	instances, err := fw.environ.Instances([]string{instanceId})
	if err != nil {
		return err
	}
	if len(toOpen) > 0 {
		err = instances[0].OpenPorts(machined.id, toOpen)
		if err != nil {
			// TODO(mue) Add local retry logic.
			return err
		}
		state.SortPorts(toOpen)
		log.Printf("firewaller: opened ports %v on machine %d", toOpen, machined.id)
	}
	if len(toClose) > 0 {
		err = instances[0].ClosePorts(machined.id, toClose)
		if err != nil {
			// TODO(mue) Add local retry logic.
			return err
		}
		state.SortPorts(toClose)
		log.Printf("firewaller: closed ports %v on machine %d", toClose, machined.id)
	}
	return nil
}

// stopWatchers stops all the firewaller's watchers.
func (fw *Firewaller) stopWatchers() {
	watcher.Stop(fw.environWatcher, &fw.tomb)
	watcher.Stop(fw.machinesWatcher, &fw.tomb)
	for _, unitd := range fw.unitds {
		watcher.Stop(unitd, &fw.tomb)
	}
	for _, serviced := range fw.serviceds {
		watcher.Stop(serviced, &fw.tomb)
	}
	for _, machined := range fw.machineds {
		watcher.Stop(machined, &fw.tomb)
	}
}

func (fw *Firewaller) String() string {
	return "firewaller"
}

// Dying returns a channel that signals a Firewaller exit.
func (fw *Firewaller) Dying() <-chan struct{} {
	return fw.tomb.Dying()
}

// Err returns the reason why the firewaller has stopped or tomb.ErrStillAlive
// when it is still alive.
func (fw *Firewaller) Err() (reason error) {
	return fw.tomb.Err()
}

// Wait waits for the Firewaller to exit.
func (fw *Firewaller) Wait() error {
	return fw.tomb.Wait()
}

// Stop stops the Firewaller and returns any error encountered while stopping.
func (fw *Firewaller) Stop() error {
	fw.tomb.Kill(nil)
	return fw.tomb.Wait()
}

// unitsChange contains the changed units for one specific machine. 
type unitsChange struct {
	machined *machineData
	*state.MachineUnitsChange
}

// machineData holds machine details and watches units added or removed.
type machineData struct {
	tomb       tomb.Tomb
	firewaller *Firewaller
	id         int
	unitds     map[string]*unitData
	ports      []state.Port
}

// newMachineData returns a new data value for tracking details of the
// machine, and starts watching the machine for units added or removed.
func newMachineData(id int, fw *Firewaller) *machineData {
	// BUG(niemeyer): The firewaller must watch *ALL* units, not just principals.
	md := &machineData{
		firewaller: fw,
		id:         id,
		unitds:     make(map[string]*unitData),
		ports:      make([]state.Port, 0),
	}
	go md.watchLoop()
	return md
}

func (md *machineData) machine() (*state.Machine, error) {
	return md.fw.st.Machine(md.id)
}

// watchLoop watches the machine for units added or removed.
func (md *machineData) watchLoop() {
	defer md.tomb.Done()
	m, err := md.machine()
	if  != nil {
	}
	watcher = machine.WatchPrincipalUnits(),
	defer md.watcher.Stop()
	for {
		select {
		case <-md.tomb.Dying():
			return
		case change, ok := <-md.watcher.Changes():
			if !ok {
				md.firewaller.tomb.Kill(watcher.MustErr(md.watcher))
				return
			}
			select {
			case md.firewaller.unitsChange <- &unitsChange{md, change}:
			case <-md.tomb.Dying():
				return
			}
		}
	}
}

// stopWatch stops the machine watching.
func (md *machineData) Stop() error {
	md.tomb.Kill(nil)
	return md.tomb.Wait()
}
<commit_msg>Ported firewaller to new watcher.<commit_after>package firewaller

import (
	"fmt"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/log"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/state/watcher"
	"launchpad.net/juju-core/worker"
	"launchpad.net/tomb"
)

// Firewaller watches the state for ports opened or closed
// and reflects those changes onto the backing environment.
type Firewaller struct {
	tomb            tomb.Tomb
	st              *state.State
	environ         environs.Environ
	environWatcher  *state.EnvironConfigWatcher
	machinesWatcher *state.MachinesWatcher
	machineds       map[int]*machineData
	unitsChange     chan *unitsChange
	unitds          map[string]*unitData
	portsChange     chan *portsChange
	serviceds       map[string]*serviceData
	exposedChange   chan *exposedChange
}

// NewFirewaller returns a new Firewaller.
func NewFirewaller(st *state.State) *Firewaller {
	fw := &Firewaller{
		st:              st,
		environWatcher:  st.WatchEnvironConfig(),
		machinesWatcher: st.WatchMachines(),
		machineds:       make(map[int]*machineData),
		unitsChange:     make(chan *unitsChange),
		unitds:          make(map[string]*unitData),
		portsChange:     make(chan *portsChange),
		serviceds:       make(map[string]*serviceData),
		exposedChange:   make(chan *exposedChange),
	}
	go func() {
		defer fw.tomb.Done()
		fw.tomb.Kill(fw.loop())
	}()
	return fw
}

func (fw *Firewaller) loop() error {
	defer fw.stopWatchers()

	var err error
	fw.environ, err = worker.WaitForEnviron(fw.environWatcher, fw.tomb.Dying())
	if err != nil {
		return err
	}
	for {
		select {
		case <-fw.tomb.Dying():
			return tomb.ErrDying
		case change, ok := <-fw.environWatcher.Changes():
			if !ok {
				return watcher.MustErr(fw.environWatcher)
			}
			if err := fw.environ.SetConfig(change); err != nil {
				log.Printf("firewaller loaded invalid environment configuration: %v", err)
			}
		case change, ok := <-fw.machinesWatcher.Changes():
			if !ok {
				return watcher.MustErr(fw.machinesWatcher)
			}
			for _, id := range change.Dead {
				machined := fw.machineds[id]
				delete(fw.machineds, id)
				if err := machined.Stop(); err != nil {
					log.Printf("machine %d watcher returned error when stopping: %v", id, err)
				}
				log.Debugf("firewaller: stopped watching machine %d", id)
			}
			for _, id := range change.Alive {
				fw.machineds[id] = newMachineData(id, fw)
				log.Debugf("firewaller: started watching machine %d", id)
			}
		case change := <-fw.unitsChange:
			changed := []*unitData{}
			for _, unit := range change.Removed {
				unitd, ok := fw.unitds[unit.Name()]
				if !ok {
					panic("trying to remove unit that was not added")
				}
				delete(fw.unitds, unit.Name())
				delete(unitd.machined.unitds, unit.Name())
				delete(unitd.serviced.unitds, unit.Name())
				changed = append(changed, unitd)
				if err := unitd.Stop(); err != nil {
					log.Printf("unit watcher %q returned error when stopping: %v", unit.Name(), err)
				}
				log.Debugf("firewaller: stopped watching unit %s", unit.Name())
			}
			for _, unit := range change.Added {
				unitd := newUnitData(unit, fw)
				fw.unitds[unit.Name()] = unitd
				machineId, err := unit.AssignedMachineId()
				if err != nil {
					fw.tomb.Kill(err)
				}
				if fw.machineds[machineId] == nil {
					panic("machine of added unit is not watched")
				}
				unitd.machined = fw.machineds[machineId]
				unitd.machined.unitds[unit.Name()] = unitd
				if fw.serviceds[unit.ServiceName()] == nil {
					service, err := fw.st.Service(unit.ServiceName())
					if err != nil {
						return err
					}
					fw.serviceds[unit.ServiceName()] = newServiceData(service, fw)
				}
				unitd.serviced = fw.serviceds[unit.ServiceName()]
				unitd.serviced.unitds[unit.ServiceName()] = unitd
				fw.serviceds[unit.ServiceName()].unitds[unit.Name()] = unitd
				changed = append(changed, unitd)
				log.Debugf("firewaller: started watching unit %s", unit.Name())
			}
			if err := fw.flushUnits(changed); err != nil {
				return fmt.Errorf("cannot change firewall ports: %v", err)
			}
		case change := <-fw.portsChange:
			change.unitd.ports = change.ports
			if err := fw.flushUnits([]*unitData{change.unitd}); err != nil {
				return fmt.Errorf("cannot change firewall ports: %v", err)
			}
		case change := <-fw.exposedChange:
			change.serviced.exposed = change.exposed
			unitds := []*unitData{}
			for _, unitd := range change.serviced.unitds {
				unitds = append(unitds, unitd)
			}
			if err := fw.flushUnits(unitds); err != nil {
				return fmt.Errorf("cannot change firewall ports: %v", err)
			}
		}
	}
	panic("not reached")
}

// flushUnits opens and closes ports for the passed unit data.
func (fw *Firewaller) flushUnits(unitds []*unitData) error {
	machineds := map[int]*machineData{}
	for _, unitd := range unitds {
		machineds[unitd.machined.id] = unitd.machined
	}
	for _, machined := range machineds {
		if err := fw.flushMachine(machined); err != nil {
			return err
		}
	}
	return nil
}

// flushMachine opens and closes ports for the passed machine.
func (fw *Firewaller) flushMachine(machined *machineData) error {
	// Gather ports to open and close.
	ports := map[state.Port]bool{}
	for _, unitd := range machined.unitds {
		if unitd.serviced.exposed {
			for _, port := range unitd.ports {
				ports[port] = true
			}
		}
	}
	want := []state.Port{}
	for port := range ports {
		want = append(want, port)
	}
	toOpen := diff(want, machined.ports)
	toClose := diff(machined.ports, want)
	machined.ports = want

	// If there's nothing to do, do nothing.
	// This is important because when a machine is first created,
	// it will have no instance id but also no open ports -
	// InstanceId will fail but we don't care.
	if len(toOpen) == 0 && len(toClose) == 0 {
		return nil
	}
	// Open and close the ports.
	m, err := machined.machine()
	if state.IsNotFound(err) {
		return nil
	}
	if err != nil {
		return err
	}
	instanceId, err := m.InstanceId()
	if err != nil {
		return err
	}
	instances, err := fw.environ.Instances([]string{instanceId})
	if err != nil {
		return err
	}
	if len(toOpen) > 0 {
		err = instances[0].OpenPorts(machined.id, toOpen)
		if err != nil {
			// TODO(mue) Add local retry logic.
			return err
		}
		state.SortPorts(toOpen)
		log.Printf("firewaller: opened ports %v on machine %d", toOpen, machined.id)
	}
	if len(toClose) > 0 {
		err = instances[0].ClosePorts(machined.id, toClose)
		if err != nil {
			// TODO(mue) Add local retry logic.
			return err
		}
		state.SortPorts(toClose)
		log.Printf("firewaller: closed ports %v on machine %d", toClose, machined.id)
	}
	return nil
}

// stopWatchers stops all the firewaller's watchers.
func (fw *Firewaller) stopWatchers() {
	watcher.Stop(fw.environWatcher, &fw.tomb)
	watcher.Stop(fw.machinesWatcher, &fw.tomb)
	for _, unitd := range fw.unitds {
		watcher.Stop(unitd, &fw.tomb)
	}
	for _, serviced := range fw.serviceds {
		watcher.Stop(serviced, &fw.tomb)
	}
	for _, machined := range fw.machineds {
		watcher.Stop(machined, &fw.tomb)
	}
}

func (fw *Firewaller) String() string {
	return "firewaller"
}

// Dying returns a channel that signals a Firewaller exit.
func (fw *Firewaller) Dying() <-chan struct{} {
	return fw.tomb.Dying()
}

// Err returns the reason why the firewaller has stopped or tomb.ErrStillAlive
// when it is still alive.
func (fw *Firewaller) Err() (reason error) {
	return fw.tomb.Err()
}

// Wait waits for the Firewaller to exit.
func (fw *Firewaller) Wait() error {
	return fw.tomb.Wait()
}

// Stop stops the Firewaller and returns any error encountered while stopping.
func (fw *Firewaller) Stop() error {
	fw.tomb.Kill(nil)
	return fw.tomb.Wait()
}

// unitsChange contains the changed units for one specific machine. 
type unitsChange struct {
	machined *machineData
	*state.MachineUnitsChange
}

// machineData holds machine details and watches units added or removed.
type machineData struct {
	tomb   tomb.Tomb
	fw     *Firewaller
	id     int
	unitds map[string]*unitData
	ports  []state.Port
}

// newMachineData returns a new data value for tracking details of the
// machine, and starts watching the machine for units added or removed.
func newMachineData(id int, fw *Firewaller) *machineData {
	md := &machineData{
		fw:     fw,
		id:     id,
		unitds: make(map[string]*unitData),
		ports:  make([]state.Port, 0),
	}
	go md.watchLoop()
	return md
}

func (md *machineData) machine() (*state.Machine, error) {
	return md.fw.st.Machine(md.id)
}

// watchLoop watches the machine for units added or removed.
func (md *machineData) watchLoop() {
	defer md.tomb.Done()
	m, err := md.machine()
	if state.IsNotFound(err) {
		return
	}
	if err != nil {
		md.fw.tomb.Killf("firewaller: cannot watch machine units: %v", err)
		return
	}
	// BUG(niemeyer): The firewaller must watch all units, not just principals.
	w := m.WatchPrincipalUnits()
	defer w.Stop()
	for {
		select {
		case <-md.tomb.Dying():
			return
		case change, ok := <-w.Changes():
			if !ok {
				md.fw.tomb.Kill(watcher.MustErr(w))
				return
			}
			select {
			case md.fw.unitsChange <- &unitsChange{md, change}:
			case <-md.tomb.Dying():
				return
			}
		}
	}
}

// stopWatch stops the machine watching.
func (md *machineData) Stop() error {
	md.tomb.Kill(nil)
	return md.tomb.Wait()
}

{'fraction_non_alphanumeric': 0.13258656901287091, 'fraction_numerical': 0.000753944746620712, 'mean_word_length': 3.331700489853044, 'pattern_counts': {'":': 0, '<': 25, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2624285', 'n_tokens_mistral': 7256, 'n_tokens_neox': 6547, 'n_words': 1709}