<commit_before>package plugin

import (
	log "github.com/Sirupsen/logrus"
	"github.com/davecgh/go-spew/spew"
	. "github.com/talbright/go-curator"
	"github.com/talbright/go-zookeeper/zk"

	"fmt"
	"path"
	"sync"
)

type WorkLeader struct {
	curator       *Curator
	client        *Client
	workPath      string
	stopChn       chan struct{}
	eventChn      chan Event
	workWatch     *ChildWatch
	supervisor    *WorkSupervisor
	workerTracker map[string]*Znode //otherwise we need to hold a reference to the discovery plugin...
	mutex         *sync.RWMutex
	leader        bool
}

func (p *WorkLeader) Name() string {
	return "WorkLeader"
}

func (p *WorkLeader) Accepts(eventType EventType) bool {
	mask := DiscoveryEventInactive | DiscoveryEventActive | LeaderEventElected | LeaderEventCandidate | LeaderEventResigned | ConnectionEvent
	return mask&eventType != 0
}

func (p *WorkLeader) Notify(event Event) {
	p.eventChn <- event
}

func (p *WorkLeader) OnLoad(curator *Curator) {
	p.workPath = path.Join(curator.Settings.RootPath, "work")
	p.client = curator.Client
	p.curator = curator
	p.startWorkLeader()
}

func (p *WorkLeader) OnUnload() {
	close(p.stopChn)
}

// AddWork adds work to the master list, and is independent of leadership
// status. It is a convenience method that simply writes znodes to the
// appropriate path in zk.
func (p *WorkLeader) AddWork(node *Znode) (err error) {
	workPath := path.Join(p.workPath, node.Name)
	_, err = p.client.Create(workPath, node.Data, zk.FlagPersistent, zk.WorldACLPermAll)
	return
}

// RemoveWork removes work from the master list, and is independent of
// leadership status. It is a convenience method that simply removes znodes
// from the appropriate path in zk.
func (p *WorkLeader) RemoveWork(node *Znode) (err error) {
	workPath := path.Join(p.workPath, node.Name)
	var version int32
	if node.Stat != nil {
		version = node.Stat.Version
	} else {
		version = -1
	}
	err = p.client.Delete(workPath, version)
	return
}

func (p *WorkLeader) IsLeader() bool {
	p.mutex.RLock()
	defer p.mutex.RUnlock()
	return p.leader == true
}

func (p *WorkLeader) setLeader(leader bool) {
	p.mutex.Lock()
	defer p.mutex.Unlock()
	p.leader = leader
}

func (p *WorkLeader) startWorkLeader() {
	p.eventChn = make(chan Event, 10)
	p.stopChn = make(chan struct{})
	p.workerTracker = make(map[string]*Znode)
	p.mutex = &sync.RWMutex{}
	go p.loop()
}

func (p *WorkLeader) loop() {
	workWatchChn := make(chan Event)
	var err error
	entry := p.curator.LogEntry("work_leader")

	for {
		select {
		case event := <-p.eventChn:
			if event.Type == LeaderEventElected {
				entry.Info("work leader: elected")
				if !p.IsLeader() {
					p.becomeLeader()
					if workWatchChn, err = p.workWatch.WatchChildren(); err != nil {
						panic(err)
					}
				} else {
					entry.Warn("work leader: already elected")
				}
			} else if event.Type == LeaderEventCandidate {
				entry.Info("work leader: candidate")
			} else if event.Type == LeaderEventResigned {
				entry.Info("work leader: resigned")
				if p.IsLeader() {
					p.resignLeader()
					workWatchChn = make(chan Event)
				}
				//Continuously track workers. An alternative would be to hold a reference to
				//the Discovery plugin, but that hasn't been flushed out yet. Curator holds
				//references to all plugins, however there are concurrency issues to consider.
			} else if event.Type == DiscoveryEventActive {
				entry.WithField("path", event.Data["path"]).Info("work leader discovery: active")
				p.addWorker(event.Data["path"].(string))
			} else if event.Type == DiscoveryEventInactive {
				entry.WithField("path", event.Data["path"]).Info("work leader discovery: inactive")
				p.removeWorker(event.Data["path"].(string))
			}
		case event := <-workWatchChn:
			p.processWorkEvents(event)
		case <-p.stopChn:
			return
		}
	}
}

func (p *WorkLeader) becomeLeader() {
	entry := p.curator.LogEntry("work_leader")
	p.setLeader(true)
	p.supervisor = NewWorkSupervisor(p.client, p.workPath)
	p.supervisor.Logger = p.curator.Logger()
	p.supervisor.LogComponent = fmt.Sprintf("%s.%s", p.curator.Settings.LogComponent, "work_supervisor")
	p.supervisor.Load()
	entry.WithField("spew", p.workerTracker).Debug("loading work onto supervisor")
	for _, v := range p.workerTracker {
		if err := p.supervisor.AddWorker(v); err != nil {
			entry.WithError(err).WithField("worker", spew.Sprintf("%#v", v)).Error("unable to add worker")
		}
	}<commit_msg>Don't process initial event to avoid duplicates<commit_after>package plugin

import (
	log "github.com/Sirupsen/logrus"
	"github.com/davecgh/go-spew/spew"
	. "github.com/talbright/go-curator"
	"github.com/talbright/go-zookeeper/zk"

	"fmt"
	"path"
	"sync"
)

type WorkLeader struct {
	curator       *Curator
	client        *Client
	workPath      string
	stopChn       chan struct{}
	eventChn      chan Event
	workWatch     *ChildWatch
	supervisor    *WorkSupervisor
	workerTracker map[string]*Znode //otherwise we need to hold a reference to the discovery plugin...
	mutex         *sync.RWMutex
	leader        bool
}

func (p *WorkLeader) Name() string {
	return "WorkLeader"
}

func (p *WorkLeader) Accepts(eventType EventType) bool {
	mask := DiscoveryEventInactive | DiscoveryEventActive | LeaderEventElected | LeaderEventCandidate | LeaderEventResigned | ConnectionEvent
	return mask&eventType != 0
}

func (p *WorkLeader) Notify(event Event) {
	p.eventChn <- event
}

func (p *WorkLeader) OnLoad(curator *Curator) {
	p.workPath = path.Join(curator.Settings.RootPath, "work")
	p.client = curator.Client
	p.curator = curator
	p.startWorkLeader()
}

func (p *WorkLeader) OnUnload() {
	close(p.stopChn)
}

// AddWork adds work to the master list, and is independent of leadership
// status. It is a convenience method that simply writes znodes to the
// appropriate path in zk.
func (p *WorkLeader) AddWork(node *Znode) (err error) {
	workPath := path.Join(p.workPath, node.Name)
	_, err = p.client.Create(workPath, node.Data, zk.FlagPersistent, zk.WorldACLPermAll)
	return
}

// RemoveWork removes work from the master list, and is independent of
// leadership status. It is a convenience method that simply removes znodes
// from the appropriate path in zk.
func (p *WorkLeader) RemoveWork(node *Znode) (err error) {
	workPath := path.Join(p.workPath, node.Name)
	var version int32
	if node.Stat != nil {
		version = node.Stat.Version
	} else {
		version = -1
	}
	err = p.client.Delete(workPath, version)
	return
}

func (p *WorkLeader) IsLeader() bool {
	p.mutex.RLock()
	defer p.mutex.RUnlock()
	return p.leader == true
}

func (p *WorkLeader) setLeader(leader bool) {
	p.mutex.Lock()
	defer p.mutex.Unlock()
	p.leader = leader
}

func (p *WorkLeader) startWorkLeader() {
	p.eventChn = make(chan Event, 10)
	p.stopChn = make(chan struct{})
	p.workerTracker = make(map[string]*Znode)
	p.mutex = &sync.RWMutex{}
	go p.loop()
}

func (p *WorkLeader) loop() {
	workWatchChn := make(chan Event)
	var err error
	entry := p.curator.LogEntry("work_leader")

	for {
		select {
		case event := <-p.eventChn:
			if event.Type == LeaderEventElected {
				entry.Info("work leader: elected")
				if !p.IsLeader() {
					p.becomeLeader()
					if workWatchChn, err = p.workWatch.WatchChildren(); err != nil {
						panic(err)
					}
				} else {
					entry.Warn("work leader: already elected")
				}
			} else if event.Type == LeaderEventCandidate {
				entry.Info("work leader: candidate")
			} else if event.Type == LeaderEventResigned {
				entry.Info("work leader: resigned")
				if p.IsLeader() {
					p.resignLeader()
					workWatchChn = make(chan Event)
				}
				//Continuously track workers. An alternative would be to hold a reference to
				//the Discovery plugin, but that hasn't been flushed out yet. Curator holds
				//references to all plugins, however there are concurrency issues to consider.
			} else if event.Type == DiscoveryEventActive {
				entry.WithField("path", event.Data["path"]).Info("work leader discovery: active")
				p.addWorker(event.Data["path"].(string))
			} else if event.Type == DiscoveryEventInactive {
				entry.WithField("path", event.Data["path"]).Info("work leader discovery: inactive")
				p.removeWorker(event.Data["path"].(string))
			}
		case event := <-workWatchChn:
			// Don't processs this event: ChildrenWatchLoadedEvent
			if event.Type != ChildrenWatchLoadedEvent {
				p.processWorkEvents(event)
			}
		case <-p.stopChn:
			return
		}
	}
}

func (p *WorkLeader) becomeLeader() {
	entry := p.curator.LogEntry("work_leader")
	p.setLeader(true)
	p.supervisor = NewWorkSupervisor(p.client, p.workPath)
	p.supervisor.Logger = p.curator.Logger()
	p.supervisor.LogComponent = fmt.Sprintf("%s.%s", p.curator.Settings.LogComponent, "work_supervisor")
	p.supervisor.Load()
	entry.WithField("spew", p.workerTracker).Debug("loading work onto supervisor")
	for _, v := range p.workerTracker {
		if err := p.supervisor.AddWorker(v); err != nil {
			entry.WithError(err).WithField("worker", spew.Sprintf("%#v", v)).Error("unable to add worker")
		}
	}
{'fraction_non_alphanumeric': 0.1290608714935137, 'fraction_numerical': 0.001330524448386739, 'mean_word_length': 3.8132337246531485, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 3, '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': '7182057', 'n_tokens_mistral': 3326, 'n_tokens_neox': 3053, 'n_words': 788}