<commit_before>package scheduler

import (
	"errors"
	"fmt"
	// "strings"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/intelsdi-x/gomit"

	"github.com/intelsdi-x/pulse/core"
	"github.com/intelsdi-x/pulse/core/ctypes"
	"github.com/intelsdi-x/pulse/core/perror"
	"github.com/intelsdi-x/pulse/core/scheduler_event"
	"github.com/intelsdi-x/pulse/pkg/schedule"
	"github.com/intelsdi-x/pulse/scheduler/wmap"
)

var (
	HandlerRegistrationName = "scheduler"

	ErrMetricManagerNotSet = errors.New("MetricManager is not set.")
	ErrSchedulerNotStarted = errors.New("Scheduler is not started.")
)

type schedulerState int

const (
	schedulerStopped schedulerState = iota
	schedulerStarted
)

// ManagesMetric is implemented by control
// On startup a scheduler will be created and passed a reference to control
type managesMetrics interface {
	collectsMetrics
	publishesMetrics
	processesMetrics
	managesPluginContentTypes
	ValidateDeps([]core.Metric, []core.SubscribedPlugin) []perror.PulseError
	SubscribeDeps(string, []core.Metric, []core.Plugin) []perror.PulseError
	UnsubscribeDeps(string, []core.Metric, []core.Plugin) []perror.PulseError
}

// ManagesPluginContentTypes is an interface to a plugin manager that can tell us what content accept and returns are supported.
type managesPluginContentTypes interface {
	GetPluginContentTypes(n string, t core.PluginType, v int) ([]string, []string, error)
}

type collectsMetrics interface {
	CollectMetrics([]core.Metric, time.Time) ([]core.Metric, []error)
}

type publishesMetrics interface {
	PublishMetrics(contentType string, content []byte, pluginName string, pluginVersion int, config map[string]ctypes.ConfigValue) []error
}

type processesMetrics interface {
	ProcessMetrics(contentType string, content []byte, pluginName string, pluginVersion int, config map[string]ctypes.ConfigValue) (string, []byte, []error)
}

type scheduler struct {
	workManager     *workManager
	metricManager   managesMetrics
	tasks           *taskCollection
	state           schedulerState
	logger          *log.Entry
	eventManager    *gomit.EventController
	taskWatcherColl *taskWatcherCollection
}

type managesWork interface {
	Work(job) job
}

// New returns an instance of the scheduler
// The MetricManager must be set before the scheduler can be started.
// The MetricManager must be started before it can be used.
func New(opts ...workManagerOption) *scheduler {
	s := &scheduler{
		tasks: newTaskCollection(),
		logger: log.WithFields(log.Fields{
			"_module": "scheduler",
		}),
		eventManager:    gomit.NewEventController(),
		taskWatcherColl: newTaskWatcherCollection(),
	}

	// we are setting the size of the queue and number of workers for
	// collect, process and publish consistently for now
	s.workManager = newWorkManager(opts...)
	s.workManager.Start()
	s.eventManager.RegisterHandler(HandlerRegistrationName, s)

	return s
}

type taskErrors struct {
	errs []perror.PulseError
}

func (t *taskErrors) Errors() []perror.PulseError {
	return t.errs
}

func (s *scheduler) Name() string {
	return "scheduler"
}

func (s *scheduler) RegisterEventHandler(name string, h gomit.Handler) error {
	return s.eventManager.RegisterHandler(name, h)
}

// CreateTask creates and returns task
func (s *scheduler) CreateTask(sch schedule.Schedule, wfMap *wmap.WorkflowMap, startOnCreate bool, opts ...core.TaskOption) (core.Task, core.TaskErrors) {
	logger := s.logger.WithField("_block", "create-task")
	// Create a container for task errors
	te := &taskErrors{
		errs: make([]perror.PulseError, 0),
	}

	// Return error if we are not started.
	if s.state != schedulerStarted {
		te.errs = append(te.errs, perror.New(ErrSchedulerNotStarted))
		f := buildErrorsLog(te.Errors(), logger)
		f.Error("scheduler not started")
		return nil, te
	}

	// Ensure the schedule is valid at this point and time.
	if err := sch.Validate(); err != nil {
		te.errs = append(te.errs, perror.New(err))
		f := buildErrorsLog(te.Errors(), logger)
		f.Error("schedule passed not valid")
		return nil, te
	}

	// Generate a workflow from the workflow map
	wf, err := wmapToWorkflow(wfMap)
	if err != nil {
		te.errs = append(te.errs, perror.New(err))
		f := buildErrorsLog(te.Errors(), logger)
		f.Error(ErrSchedulerNotStarted.Error())
		return nil, te
	}

	// Bind plugin content type selections in workflow
	err = wf.BindPluginContentTypes(s.metricManager)

	// validate plugins and metrics
	mts, plugins := s.gatherPlugins(wf)
	errs := s.metricManager.ValidateDeps(mts, plugins)
	if len(errs) > 0 {
		te.errs = append(te.errs, errs...)
		return nil, te
	}

	// Create the task object
	task := newTask(sch, wf, s.workManager, s.metricManager, s.eventManager, opts...)

	// Add task to taskCollection
	if err := s.tasks.add(task); err != nil {
		te.errs = append(te.errs, perror.New(err))
		f := buildErrorsLog(te.Errors(), logger)
		f.Error("errors during task creation")
		return nil, te
	}

	logger.WithFields(log.Fields{
		"task-id":    task.ID(),
		"task-state": task.State(),
	}).Info("task created")

	event := &scheduler_event.TaskCreatedEvent{
		TaskID:        task.id,
		StartOnCreate: startOnCreate,
	}
	defer s.eventManager.Emit(event)

	if startOnCreate {
		logger.WithFields(log.Fields{
			"task-id": task.ID(),
		}).Info("starting task on creation")

		cps := make([]core.Plugin, len(plugins))
		for i, plugin := range plugins {
			cps[i] = plugin
		}
		errs := s.metricManager.SubscribeDeps(task.ID(), mts, cps)
		if len(errs) > 0 {
			te.errs = append(te.errs, errs...)
			return nil, te
		}

		task.Spin()
	}

	return task, te
}

// RemoveTask given a tasks id.  The task must be stopped.
// Can return errors ErrTaskNotFound and ErrTaskNotStopped.
func (s *scheduler) RemoveTask(id string) error {
	t, err := s.getTask(id)
	if err != nil {
		log.WithFields(log.Fields{
			"_module": "scheduler",
			"block":   "RemoveTask",
			"task id": id,
		}).Error(ErrTaskNotFound)
		return err
	}
	event := &scheduler_event.TaskDeletedEvent{
		TaskID: t.id,
	}
	defer s.eventManager.Emit(event)
	return s.tasks.remove(t)
}

// GetTasks returns a copy of the tasks in a map where the task id is the key
func (s *scheduler) GetTasks() map[string]core.Task {
	tasks := make(map[string]core.Task)
	for id, t := range s.tasks.Table() {
		tasks[id] = t
	}
	return tasks
}

// GetTask provided the task id a task is returned
func (s *scheduler) GetTask(id string) (core.Task, error) {
	t, err := s.getTask(id)
	if err != nil {
		return nil, err // We do this to send back an explicit nil on the interface
	}
	return t, nil
}

// StartTask provided a task id a task is started
func (s *scheduler) StartTask(id string) []perror.PulseError {
	t, err := s.getTask(id)
	if err != nil {
		return []perror.PulseError{
			perror.New(err),
		}
	}

	mts, plugins := s.gatherPlugins(t.workflow)
	cps := make([]core.Plugin, len(plugins))
	for i, plugin := range plugins {
		cps[i] = plugin
	}
	errs := s.metricManager.SubscribeDeps(id, mts, cps)
	if len(errs) > 0 {
		return errs
	}

	event := new(scheduler_event.TaskStartedEvent)
	event.TaskID = t.ID()
	defer s.eventManager.Emit(event)
	t.Spin()
	s.logger.WithFields(log.Fields{
		"_block":     "start-task",
		"task-id":    t.ID(),
		"task-state": t.State(),
	}).Info("task started")
	return nil
}

// StopTask provided a task id a task is stopped
func (s *scheduler) StopTask(id string) []perror.PulseError {
	t, err := s.getTask(id)
	if err != nil {
		s.logger.WithFields(log.Fields{
			"_block":  "stop-task",
			"_error":  err.Error(),
			"task-id": id,
		}).Warning("error on stopping of task")
		return []perror.PulseError{
			perror.New(err),
		}
	}

	mts, plugins := s.gatherPlugins(t.workflow)
	cps := make([]core.Plugin, len(plugins))
	for i, plugin := range plugins {
		cps[i] = plugin
	}
	errs := s.metricManager.UnsubscribeDeps(t.ID(), mts, cps)
	if len(errs) > 0 {
		return errs
	}

	event := new(scheduler_event.TaskStoppedEvent)
	event.TaskID = t.ID()
	defer s.eventManager.Emit(event)
	t.Stop()
	s.logger.WithFields(log.Fields{
		"_block":     "stop-task",
		"task-id":    t.ID(),
		"task-state": t.State(),
	}).Info("task stopped")
	return nil
}

// Start starts the scheduler
func (s *scheduler) Start() error {
	if s.metricManager == nil {
		s.logger.WithFields(log.Fields{
			"_block": "start-scheduler",
			"_error": ErrMetricManagerNotSet.Error(),
		}).Error("error on scheduler start")
		return ErrMetricManagerNotSet
	}
	s.state = schedulerStarted
	s.logger.WithFields(log.Fields{
		"_block": "start-scheduler",
	}).Info("scheduler started")
	return nil
}

func (s *scheduler) Stop() {
	s.state = schedulerStopped
	// stop all tasks that are not already stopped
	for _, t := range s.tasks.table {
		// Kill ensure another task can't turn it back on while we are shutting down
		t.Kill()
	}
	s.logger.WithFields(log.Fields{
		"_block": "stop-scheduler",
	}).Info("scheduler stopped")
}

// Set metricManager for scheduler
func (s *scheduler) SetMetricManager(mm managesMetrics) {
	s.metricManager = mm
	s.logger.WithFields(log.Fields{
		"_block": "set-metric-manager",
	}).Debug("metric manager linked")
}

//
func (s *scheduler) WatchTask(id string, tw core.TaskWatcherHandler) (core.TaskWatcherCloser, error) {
	task, err := s.getTask(id)
	if err != nil {
		return nil, err
	}
	return s.taskWatcherColl.add(task.ID(), tw)
}

// Central handling for all async events in scheduler
func (s *scheduler) HandleGomitEvent(e gomit.Event) {

	switch v := e.Body.(type) {
	case *scheduler_event.MetricCollectedEvent:
		log.WithFields(log.Fields{
			"_module":         "scheduler-events",
			"_block":          "handle-events",
			"event-namespace": e.Namespace(),
			"task-id":         v.TaskID,
			"metric-count":    len(v.Metrics),
		}).Debug("event received")
		s.taskWatcherColl.handleMetricCollected(v.TaskID, v.Metrics)
	case *scheduler_event.MetricCollectionFailedEvent:
		log.WithFields(log.Fields{
			"_module":         "scheduler-events",
			"_block":          "handle-events",
			"event-namespace": e.Namespace(),
			"task-id":         v.TaskID,
			"errors-count":    v.Errors,
		}).Debug("event received")
	case *scheduler_event.TaskStartedEvent:
		log.WithFields(log.Fields{
			"_module":         "scheduler-events",
			"_block":          "handle-events",
			"event-namespace": e.Namespace(),
			"task-id":         v.TaskID,
		}).Debug("event received")
		s.taskWatcherColl.handleTaskStarted(v.TaskID)
	case *scheduler_event.TaskStoppedEvent:
		log.WithFields(log.Fields{
			"_module":         "scheduler-events",
			"_block":          "handle-events",
			"event-namespace": e.Namespace(),
			"task-id":         v.TaskID,
		}).Debug("event received")
		s.taskWatcherColl.handleTaskStopped(v.TaskID)
	case *scheduler_event.TaskDisabledEvent:
		log.WithFields(log.Fields{
			"_module":         "scheduler-events",
			"_block":          "handle-events",
			"event-namespace": e.Namespace(),
			"task-id":         v.TaskID,
			"disabled-reason": v.Why,
		}).Debug("event received")
		s.taskWatcherColl.handleTaskDisabled(v.TaskID, v.Why)
	default:
		log.WithFields(log.Fields{
			"_module":         "scheduler-events",
			"_block":          "handle-events",
			"event-namespace": e.Namespace(),
		}).Debug("event received")
	}
}

func (s *scheduler) getTask(id string) (*task, error) {
	task := s.tasks.Get(id)
	if task == nil {
		return nil, fmt.Errorf("No task found with id '%v'", id)
	}
	return task, nil
}

func (s *scheduler) gatherPlugins(wf *schedulerWorkflow) ([]core.Metric, []core.SubscribedPlugin) {
	var (
		mts     []core.Metric
		plugins []core.SubscribedPlugin
	)

	for _, m := range wf.metrics {
		mts = append(mts, &metric{
			namespace: m.Namespace(),
			version:   m.Version(),
			config:    wf.configTree.Get(m.Namespace()),
		})
	}
	s.walkWorkflow(wf.processNodes, wf.publishNodes, &plugins)

	return mts, plugins
}

func (s *scheduler) walkWorkflow(prnodes []*processNode, pbnodes []*publishNode, plugins *[]core.SubscribedPlugin) {
	for _, pr := range prnodes {
		*plugins = append(*plugins, pr)
		s.walkWorkflow(pr.ProcessNodes, pr.PublishNodes, plugins)
	}
	for _, pb := range pbnodes {
		*plugins = append(*plugins, pb)
	}
}

func buildErrorsLog(errs []perror.PulseError, logger *log.Entry) *log.Entry {
	for i, e := range errs {
		logger = logger.WithField(fmt.Sprintf("%s[%d]", "error", i), e.Error())
	}
	return logger
}
<commit_msg>Update function name to reflect that it returns both metrics and plugins<commit_after>package scheduler

import (
	"errors"
	"fmt"
	// "strings"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/intelsdi-x/gomit"

	"github.com/intelsdi-x/pulse/core"
	"github.com/intelsdi-x/pulse/core/ctypes"
	"github.com/intelsdi-x/pulse/core/perror"
	"github.com/intelsdi-x/pulse/core/scheduler_event"
	"github.com/intelsdi-x/pulse/pkg/schedule"
	"github.com/intelsdi-x/pulse/scheduler/wmap"
)

var (
	HandlerRegistrationName = "scheduler"

	ErrMetricManagerNotSet = errors.New("MetricManager is not set.")
	ErrSchedulerNotStarted = errors.New("Scheduler is not started.")
)

type schedulerState int

const (
	schedulerStopped schedulerState = iota
	schedulerStarted
)

// ManagesMetric is implemented by control
// On startup a scheduler will be created and passed a reference to control
type managesMetrics interface {
	collectsMetrics
	publishesMetrics
	processesMetrics
	managesPluginContentTypes
	ValidateDeps([]core.Metric, []core.SubscribedPlugin) []perror.PulseError
	SubscribeDeps(string, []core.Metric, []core.Plugin) []perror.PulseError
	UnsubscribeDeps(string, []core.Metric, []core.Plugin) []perror.PulseError
}

// ManagesPluginContentTypes is an interface to a plugin manager that can tell us what content accept and returns are supported.
type managesPluginContentTypes interface {
	GetPluginContentTypes(n string, t core.PluginType, v int) ([]string, []string, error)
}

type collectsMetrics interface {
	CollectMetrics([]core.Metric, time.Time) ([]core.Metric, []error)
}

type publishesMetrics interface {
	PublishMetrics(contentType string, content []byte, pluginName string, pluginVersion int, config map[string]ctypes.ConfigValue) []error
}

type processesMetrics interface {
	ProcessMetrics(contentType string, content []byte, pluginName string, pluginVersion int, config map[string]ctypes.ConfigValue) (string, []byte, []error)
}

type scheduler struct {
	workManager     *workManager
	metricManager   managesMetrics
	tasks           *taskCollection
	state           schedulerState
	logger          *log.Entry
	eventManager    *gomit.EventController
	taskWatcherColl *taskWatcherCollection
}

type managesWork interface {
	Work(job) job
}

// New returns an instance of the scheduler
// The MetricManager must be set before the scheduler can be started.
// The MetricManager must be started before it can be used.
func New(opts ...workManagerOption) *scheduler {
	s := &scheduler{
		tasks: newTaskCollection(),
		logger: log.WithFields(log.Fields{
			"_module": "scheduler",
		}),
		eventManager:    gomit.NewEventController(),
		taskWatcherColl: newTaskWatcherCollection(),
	}

	// we are setting the size of the queue and number of workers for
	// collect, process and publish consistently for now
	s.workManager = newWorkManager(opts...)
	s.workManager.Start()
	s.eventManager.RegisterHandler(HandlerRegistrationName, s)

	return s
}

type taskErrors struct {
	errs []perror.PulseError
}

func (t *taskErrors) Errors() []perror.PulseError {
	return t.errs
}

func (s *scheduler) Name() string {
	return "scheduler"
}

func (s *scheduler) RegisterEventHandler(name string, h gomit.Handler) error {
	return s.eventManager.RegisterHandler(name, h)
}

// CreateTask creates and returns task
func (s *scheduler) CreateTask(sch schedule.Schedule, wfMap *wmap.WorkflowMap, startOnCreate bool, opts ...core.TaskOption) (core.Task, core.TaskErrors) {
	logger := s.logger.WithField("_block", "create-task")
	// Create a container for task errors
	te := &taskErrors{
		errs: make([]perror.PulseError, 0),
	}

	// Return error if we are not started.
	if s.state != schedulerStarted {
		te.errs = append(te.errs, perror.New(ErrSchedulerNotStarted))
		f := buildErrorsLog(te.Errors(), logger)
		f.Error("scheduler not started")
		return nil, te
	}

	// Ensure the schedule is valid at this point and time.
	if err := sch.Validate(); err != nil {
		te.errs = append(te.errs, perror.New(err))
		f := buildErrorsLog(te.Errors(), logger)
		f.Error("schedule passed not valid")
		return nil, te
	}

	// Generate a workflow from the workflow map
	wf, err := wmapToWorkflow(wfMap)
	if err != nil {
		te.errs = append(te.errs, perror.New(err))
		f := buildErrorsLog(te.Errors(), logger)
		f.Error(ErrSchedulerNotStarted.Error())
		return nil, te
	}

	// Bind plugin content type selections in workflow
	err = wf.BindPluginContentTypes(s.metricManager)

	// validate plugins and metrics
	mts, plugins := s.gatherMetricsAndPlugins(wf)
	errs := s.metricManager.ValidateDeps(mts, plugins)
	if len(errs) > 0 {
		te.errs = append(te.errs, errs...)
		return nil, te
	}

	// Create the task object
	task := newTask(sch, wf, s.workManager, s.metricManager, s.eventManager, opts...)

	// Add task to taskCollection
	if err := s.tasks.add(task); err != nil {
		te.errs = append(te.errs, perror.New(err))
		f := buildErrorsLog(te.Errors(), logger)
		f.Error("errors during task creation")
		return nil, te
	}

	logger.WithFields(log.Fields{
		"task-id":    task.ID(),
		"task-state": task.State(),
	}).Info("task created")

	event := &scheduler_event.TaskCreatedEvent{
		TaskID:        task.id,
		StartOnCreate: startOnCreate,
	}
	defer s.eventManager.Emit(event)

	if startOnCreate {
		logger.WithFields(log.Fields{
			"task-id": task.ID(),
		}).Info("starting task on creation")

		cps := make([]core.Plugin, len(plugins))
		for i, plugin := range plugins {
			cps[i] = plugin
		}
		errs := s.metricManager.SubscribeDeps(task.ID(), mts, cps)
		if len(errs) > 0 {
			te.errs = append(te.errs, errs...)
			return nil, te
		}

		task.Spin()
	}

	return task, te
}

// RemoveTask given a tasks id.  The task must be stopped.
// Can return errors ErrTaskNotFound and ErrTaskNotStopped.
func (s *scheduler) RemoveTask(id string) error {
	t, err := s.getTask(id)
	if err != nil {
		log.WithFields(log.Fields{
			"_module": "scheduler",
			"block":   "RemoveTask",
			"task id": id,
		}).Error(ErrTaskNotFound)
		return err
	}
	event := &scheduler_event.TaskDeletedEvent{
		TaskID: t.id,
	}
	defer s.eventManager.Emit(event)
	return s.tasks.remove(t)
}

// GetTasks returns a copy of the tasks in a map where the task id is the key
func (s *scheduler) GetTasks() map[string]core.Task {
	tasks := make(map[string]core.Task)
	for id, t := range s.tasks.Table() {
		tasks[id] = t
	}
	return tasks
}

// GetTask provided the task id a task is returned
func (s *scheduler) GetTask(id string) (core.Task, error) {
	t, err := s.getTask(id)
	if err != nil {
		return nil, err // We do this to send back an explicit nil on the interface
	}
	return t, nil
}

// StartTask provided a task id a task is started
func (s *scheduler) StartTask(id string) []perror.PulseError {
	t, err := s.getTask(id)
	if err != nil {
		return []perror.PulseError{
			perror.New(err),
		}
	}

	mts, plugins := s.gatherMetricsAndPlugins(t.workflow)
	cps := make([]core.Plugin, len(plugins))
	for i, plugin := range plugins {
		cps[i] = plugin
	}
	errs := s.metricManager.SubscribeDeps(id, mts, cps)
	if len(errs) > 0 {
		return errs
	}

	event := new(scheduler_event.TaskStartedEvent)
	event.TaskID = t.ID()
	defer s.eventManager.Emit(event)
	t.Spin()
	s.logger.WithFields(log.Fields{
		"_block":     "start-task",
		"task-id":    t.ID(),
		"task-state": t.State(),
	}).Info("task started")
	return nil
}

// StopTask provided a task id a task is stopped
func (s *scheduler) StopTask(id string) []perror.PulseError {
	t, err := s.getTask(id)
	if err != nil {
		s.logger.WithFields(log.Fields{
			"_block":  "stop-task",
			"_error":  err.Error(),
			"task-id": id,
		}).Warning("error on stopping of task")
		return []perror.PulseError{
			perror.New(err),
		}
	}

	mts, plugins := s.gatherMetricsAndPlugins(t.workflow)
	cps := make([]core.Plugin, len(plugins))
	for i, plugin := range plugins {
		cps[i] = plugin
	}
	errs := s.metricManager.UnsubscribeDeps(t.ID(), mts, cps)
	if len(errs) > 0 {
		return errs
	}

	event := new(scheduler_event.TaskStoppedEvent)
	event.TaskID = t.ID()
	defer s.eventManager.Emit(event)
	t.Stop()
	s.logger.WithFields(log.Fields{
		"_block":     "stop-task",
		"task-id":    t.ID(),
		"task-state": t.State(),
	}).Info("task stopped")
	return nil
}

// Start starts the scheduler
func (s *scheduler) Start() error {
	if s.metricManager == nil {
		s.logger.WithFields(log.Fields{
			"_block": "start-scheduler",
			"_error": ErrMetricManagerNotSet.Error(),
		}).Error("error on scheduler start")
		return ErrMetricManagerNotSet
	}
	s.state = schedulerStarted
	s.logger.WithFields(log.Fields{
		"_block": "start-scheduler",
	}).Info("scheduler started")
	return nil
}

func (s *scheduler) Stop() {
	s.state = schedulerStopped
	// stop all tasks that are not already stopped
	for _, t := range s.tasks.table {
		// Kill ensure another task can't turn it back on while we are shutting down
		t.Kill()
	}
	s.logger.WithFields(log.Fields{
		"_block": "stop-scheduler",
	}).Info("scheduler stopped")
}

// Set metricManager for scheduler
func (s *scheduler) SetMetricManager(mm managesMetrics) {
	s.metricManager = mm
	s.logger.WithFields(log.Fields{
		"_block": "set-metric-manager",
	}).Debug("metric manager linked")
}

//
func (s *scheduler) WatchTask(id string, tw core.TaskWatcherHandler) (core.TaskWatcherCloser, error) {
	task, err := s.getTask(id)
	if err != nil {
		return nil, err
	}
	return s.taskWatcherColl.add(task.ID(), tw)
}

// Central handling for all async events in scheduler
func (s *scheduler) HandleGomitEvent(e gomit.Event) {

	switch v := e.Body.(type) {
	case *scheduler_event.MetricCollectedEvent:
		log.WithFields(log.Fields{
			"_module":         "scheduler-events",
			"_block":          "handle-events",
			"event-namespace": e.Namespace(),
			"task-id":         v.TaskID,
			"metric-count":    len(v.Metrics),
		}).Debug("event received")
		s.taskWatcherColl.handleMetricCollected(v.TaskID, v.Metrics)
	case *scheduler_event.MetricCollectionFailedEvent:
		log.WithFields(log.Fields{
			"_module":         "scheduler-events",
			"_block":          "handle-events",
			"event-namespace": e.Namespace(),
			"task-id":         v.TaskID,
			"errors-count":    v.Errors,
		}).Debug("event received")
	case *scheduler_event.TaskStartedEvent:
		log.WithFields(log.Fields{
			"_module":         "scheduler-events",
			"_block":          "handle-events",
			"event-namespace": e.Namespace(),
			"task-id":         v.TaskID,
		}).Debug("event received")
		s.taskWatcherColl.handleTaskStarted(v.TaskID)
	case *scheduler_event.TaskStoppedEvent:
		log.WithFields(log.Fields{
			"_module":         "scheduler-events",
			"_block":          "handle-events",
			"event-namespace": e.Namespace(),
			"task-id":         v.TaskID,
		}).Debug("event received")
		s.taskWatcherColl.handleTaskStopped(v.TaskID)
	case *scheduler_event.TaskDisabledEvent:
		log.WithFields(log.Fields{
			"_module":         "scheduler-events",
			"_block":          "handle-events",
			"event-namespace": e.Namespace(),
			"task-id":         v.TaskID,
			"disabled-reason": v.Why,
		}).Debug("event received")
		s.taskWatcherColl.handleTaskDisabled(v.TaskID, v.Why)
	default:
		log.WithFields(log.Fields{
			"_module":         "scheduler-events",
			"_block":          "handle-events",
			"event-namespace": e.Namespace(),
		}).Debug("event received")
	}
}

func (s *scheduler) getTask(id string) (*task, error) {
	task := s.tasks.Get(id)
	if task == nil {
		return nil, fmt.Errorf("No task found with id '%v'", id)
	}
	return task, nil
}

func (s *scheduler) gatherMetricsAndPlugins(wf *schedulerWorkflow) ([]core.Metric, []core.SubscribedPlugin) {
	var (
		mts     []core.Metric
		plugins []core.SubscribedPlugin
	)

	for _, m := range wf.metrics {
		mts = append(mts, &metric{
			namespace: m.Namespace(),
			version:   m.Version(),
			config:    wf.configTree.Get(m.Namespace()),
		})
	}
	s.walkWorkflow(wf.processNodes, wf.publishNodes, &plugins)

	return mts, plugins
}

func (s *scheduler) walkWorkflow(prnodes []*processNode, pbnodes []*publishNode, plugins *[]core.SubscribedPlugin) {
	for _, pr := range prnodes {
		*plugins = append(*plugins, pr)
		s.walkWorkflow(pr.ProcessNodes, pr.PublishNodes, plugins)
	}
	for _, pb := range pbnodes {
		*plugins = append(*plugins, pb)
	}
}

func buildErrorsLog(errs []perror.PulseError, logger *log.Entry) *log.Entry {
	for i, e := range errs {
		logger = logger.WithField(fmt.Sprintf("%s[%d]", "error", i), e.Error())
	}
	return logger
}
<|endoftext|>


<commit_before>package carbon

import (
	"net"
	"testing"
	"time"

	"github.com/k0kubun/pp"
)

type udpTestCase struct {
	*testing.T
	receiver *UdpReceiver
	conn     net.Conn
	rcvChan  chan *Message
}

func newUdpTestCase(t *testing.T) *udpTestCase {
	test := &udpTestCase{
		T: t,
	}

	addr, err := net.ResolveUDPAddr("udp", "localhost:1818")
	if err != nil {
		t.Fatal(err)
	}

	test.rcvChan = make(chan *Message, 128)
	test.receiver = NewUdpReceiver(test.rcvChan)
	// defer receiver.Stop()

	if err = test.receiver.Listen(addr); err != nil {
		t.Fatal(err)
	}

	test.conn, err = net.Dial("udp", addr.String())
	if err != nil {
		t.Fatal(err)
	}
	// defer conn.Close()
	return test
}

func (test *udpTestCase) Finish() {
	if test.conn != nil {
		test.conn.Close()
		test.conn = nil
	}
	if test.receiver != nil {
		test.receiver.Stop()
		test.receiver = nil
	}
}

func (test *udpTestCase) Send(text string) {
	if _, err := test.conn.Write([]byte(text)); err != nil {
		test.Fatal(err)
	}
}

func TestUdpReceiver1(t *testing.T) {

	test := newUdpTestCase(t)
	defer test.Finish()

	test.Send("hello.world 42.15 1422698155\n")

	time.Sleep(10 * time.Millisecond)

	for i := 0; i < 1; i++ {
		select {
		case msg := <-test.rcvChan:
			pp.Println(msg)
		default:
			t.Fatalf("Message #%d not received", i)
		}
	}

}
<commit_msg>receiver simple tests<commit_after>package carbon

import (
	"net"
	"testing"
	"time"
)

type udpTestCase struct {
	*testing.T
	receiver *UdpReceiver
	conn     net.Conn
	rcvChan  chan *Message
}

func newUdpTestCase(t *testing.T) *udpTestCase {
	test := &udpTestCase{
		T: t,
	}

	addr, err := net.ResolveUDPAddr("udp", "localhost:1818")
	if err != nil {
		t.Fatal(err)
	}

	test.rcvChan = make(chan *Message, 128)
	test.receiver = NewUdpReceiver(test.rcvChan)
	// defer receiver.Stop()

	if err = test.receiver.Listen(addr); err != nil {
		t.Fatal(err)
	}

	test.conn, err = net.Dial("udp", addr.String())
	if err != nil {
		t.Fatal(err)
	}
	// defer conn.Close()
	return test
}

func (test *udpTestCase) Finish() {
	if test.conn != nil {
		test.conn.Close()
		test.conn = nil
	}
	if test.receiver != nil {
		test.receiver.Stop()
		test.receiver = nil
	}
}

func (test *udpTestCase) Send(text string) {
	if _, err := test.conn.Write([]byte(text)); err != nil {
		test.Fatal(err)
	}
}

func (test *udpTestCase) Eq(a *Message, b *Message) {
	if a.Name != b.Name || a.Value != b.Value || a.Timestamp != b.Timestamp {
		test.Fatalf("%#v != %#v", a, b)
	}
}

func TestUdpReceiver1(t *testing.T) {
	test := newUdpTestCase(t)
	defer test.Finish()

	test.Send("hello.world 42.15 1422698155\n")

	time.Sleep(10 * time.Millisecond)

	select {
	case msg := <-test.rcvChan:
		test.Eq(msg, &Message{
			Name:      "hello.world",
			Value:     42.15,
			Timestamp: 1422698155,
		})
	default:
		t.Fatalf("Message #0 not received")
	}
}

func TestUdpReceiver2(t *testing.T) {
	test := newUdpTestCase(t)
	defer test.Finish()

	test.Send("hello.world 42.15 1422698155\nmetric.name -72.11 1422698155\n")

	time.Sleep(10 * time.Millisecond)

	select {
	case msg := <-test.rcvChan:
		test.Eq(msg, &Message{
			Name:      "hello.world",
			Value:     42.15,
			Timestamp: 1422698155,
		})
	default:
		t.Fatalf("Message #0 not received")
	}

	select {
	case msg := <-test.rcvChan:
		test.Eq(msg, &Message{
			Name:      "metric.name",
			Value:     -72.11,
			Timestamp: 1422698155,
		})
	default:
		t.Fatalf("Message #1 not received")
	}
}
<|endoftext|>


<commit_before>package wikifier

type tocBlock struct {
	*parserBlock
}

func newTocBlock(name string, b *parserBlock) block {
	return &tocBlock{b}
}

func (toc *tocBlock) html(page *Page, el element) {
	el.setTag("ul")
	el.addHTML(HTML("Contents"))
	// add each top-level section
	for _, child := range page.main.blockContent() {
		if sec, ok := child.(*secBlock); ok {
			tocAdd(sec, el)
		}
	}
}

func tocAdd(sec *secBlock, addTo element) {

	// create an item for this section
	var subList element
	if !sec.isIntro {
		li := addTo.createChild("li", "")
		a := li.createChild("a", "link-internal")
		a.setAttr("href", "#"+sec.headingID)
		a.addText(sec.title)
		addTo = li
	} else {
		subList = addTo
	}

	// create a sub-list for each section underneath
	for _, child := range sec.blockContent() {
		if secChild, ok := child.(*secBlock); ok {
			if subList == nil {
				subList = addTo.createChild("ul", "")
			}
			tocAdd(secChild, subList)
		}
	}
}
<commit_msg>add formatted title to TOC<commit_after>package wikifier

type tocBlock struct {
	*parserBlock
}

func newTocBlock(name string, b *parserBlock) block {
	return &tocBlock{b}
}

func (toc *tocBlock) html(page *Page, el element) {
	el.setTag("ul")
	el.addHTML(HTML("Contents"))
	// add each top-level section
	for _, child := range page.main.blockContent() {
		if sec, ok := child.(*secBlock); ok {
			tocAdd(sec, el)
		}
	}
}

func tocAdd(sec *secBlock, addTo element) {

	// create an item for this section
	var subList element
	if !sec.isIntro {
		li := addTo.createChild("li", "")
		a := li.createChild("a", "link-internal")
		a.setAttr("href", "#"+sec.headingID)
		a.addHTML(sec.fmtTitle)
		addTo = li
	} else {
		subList = addTo
	}

	// create a sub-list for each section underneath
	for _, child := range sec.blockContent() {
		if secChild, ok := child.(*secBlock); ok {
			if subList == nil {
				subList = addTo.createChild("ul", "")
			}
			tocAdd(secChild, subList)
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2013 Jari Takkala and Brian Dignan. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"launchpad.net/tomb"
	"log"
	"net"
	"reflect"
	"sort"
	//"strconv"
	"syscall"
	"time"
)

var Protocol = [19]byte{'B', 'i', 't', 'T', 'o', 'r', 'r', 'e', 'n', 't', ' ', 'p', 'r', 'o', 't', 'o', 'c', 'o', 'l'}

// Message ID values
const (
	MsgChoke int = iota
	MsgUnchoke
	MsgInterested
	MsgNotInterested
	MsgHave
	MsgBitfield
	MsgRequest
	MsgPiece
	MsgCancel
	MsgPort
)

// PeerTuple represents a single IP+port pair of a peer
type PeerTuple struct {
	IP   net.IP
	Port uint16
}

type Peer struct {
	conn           *net.TCPConn
	amChoking      bool
	amInterested   bool
	peerChoking    bool
	peerInterested bool
	ourBitfield    []bool
	peerBitfield   []bool
	initiator      bool
	peerID         []byte
	keepalive      <-chan time.Time // channel for sending keepalives
	lastTxKeepalive  time.Time
	lastRxKeepalive  time.Time
	read           chan []byte
	infoHash       []byte
	diskIOChans    diskIOPeerChans
	peerManagerChans peerManagerChans
	contRxChans    ControllerPeerChans
	contTxChans    PeerControllerChans
	stats          PeerStats
	t              tomb.Tomb
}

type PeerStats struct {
	read int
	write int
	errors int
}

type PeerManager struct {
	peers        map[string]*Peer
	infoHash     []byte
	peerChans    peerManagerChans
	serverChans  serverPeerChans
	trackerChans trackerPeerChans
	diskIOChans  diskIOPeerChans
	contChans    ControllerPeerManagerChans
	peerContChans PeerControllerChans
	t            tomb.Tomb
}

type peerManagerChans struct {
	deadPeer chan string
}

type PeerComms struct {
	peerName     string
	chans 		 ControllerPeerChans
}

func NewPeerComms(peerName string, cpc ControllerPeerChans) *PeerComms {
	pc := new(PeerComms)
	pc.peerName = peerName
	pc.chans = cpc
	return pc
}

type PeerInfo struct {
	peerName        string
	isChoked        bool // The peer is connected but choked. Defaults to TRUE (choked)
	availablePieces []bool
	activeRequests  map[int]struct{}
	qtyPiecesNeeded int                   // The quantity of pieces that this peer has that we haven't yet downloaded.
	chans 			ControllerPeerChans
}

type Handshake struct {
	Len      uint8
	Protocol [19]byte
	Reserved [8]uint8
	InfoHash [20]byte
	PeerID   [20]byte
}

func NewPeerInfo(quantityOfPieces int, peerComms PeerComms) *PeerInfo {
	pi := new(PeerInfo)

	pi.peerName = peerComms.peerName
	pi.chans = peerComms.chans

	pi.isChoked = true // By default, a peer starts as being choked by the other side.
	pi.availablePieces = make([]bool, quantityOfPieces)
	pi.activeRequests = make(map[int]struct{})

	return pi
}

// Sent by the peer to controller indicating a 'choke' state change. It either went from unchoked to choked,
// or from choked to unchoked.
type PeerChokeStatus struct {
	peerName   string
	isChoked bool
}

type SortedPeers []*PeerInfo

func (sp SortedPeers) Less(i, j int) bool {
	return sp[i].qtyPiecesNeeded <= sp[j].qtyPiecesNeeded
}

func (sp SortedPeers) Swap(i, j int) {
	tmp := sp[i]
	sp[i] = sp[j]
	sp[j] = tmp
}

func (sp SortedPeers) Len() int {
	return len(sp)
}

func sortedPeersByQtyPiecesNeeded(peers map[string]*PeerInfo) SortedPeers {
	peerInfoSlice := make(SortedPeers, 0)

	for _, peerInfo := range peers {
		peerInfoSlice = append(peerInfoSlice, peerInfo)
	}
	sort.Sort(peerInfoSlice)

	return peerInfoSlice
}

func NewPeerManager(infoHash []byte, diskIOChans diskIOPeerChans, serverChans serverPeerChans, trackerChans trackerPeerChans) *PeerManager {
	pm := new(PeerManager)
	pm.infoHash = infoHash
	pm.diskIOChans = diskIOChans
	pm.serverChans = serverChans
	pm.trackerChans = trackerChans
	pm.peerChans.deadPeer = make(chan string)
	pm.peers = make(map[string]*Peer)
	pm.contChans.newPeer = make(chan PeerComms)
	pm.contChans.deadPeer = make(chan string)
	pm.peerContChans.chokeStatus = make(chan PeerChokeStatus)
	pm.peerContChans.havePiece = make(chan chan HavePiece)
	return pm
}

func ConnectToPeer(peerTuple PeerTuple, connCh chan *net.TCPConn) {
	raddr := net.TCPAddr{peerTuple.IP, int(peerTuple.Port), ""}
	log.Println("Connecting to", raddr)
	conn, err := net.DialTCP("tcp4", nil, &raddr)
	if err != nil {
		if e, ok := err.(*net.OpError); ok {
			if e.Err == syscall.ECONNREFUSED {
				log.Println("ConnectToPeer : Connection Refused:", raddr)
				return
			}
		}
		log.Fatal(err)
	}
	log.Println("ConnectToPeer : Connected:", raddr)
	connCh <- conn
}

func NewPeer(
			infoHash []byte,
			initiator bool,
			diskIOChans diskIOPeerChans,
			contRxChans ControllerPeerChans,
			contTxChans PeerControllerChans) *Peer {
	p := &Peer{
			infoHash: infoHash,
			amChoking: true,
			amInterested: false,
			peerChoking: true,
			peerInterested: false,
			initiator: initiator,
			diskIOChans: diskIOChans,
			contRxChans: contRxChans,
			contTxChans: contTxChans}
	p.read = make(chan []byte)
	return p
}

func constructMessage(id int, payload []byte) (msg []byte, err error) {
	msg = make([]byte, 4)

	// Store the length of payload + id in network byte order
	binary.BigEndian.PutUint32(msg, uint32(len(payload) + 1))
	msg = append(msg, byte(id))
	msg = append(msg, payload...)

	return
}

func verifyHandshake(handshake *Handshake, infoHash []byte) error {
	if int(handshake.Len) != len(Protocol) {
		err := fmt.Sprintf("Unexpected length for pstrlen (wanted %d, got %d)", len(Protocol), int(handshake.Len))
		return errors.New(err)
	}
	if !bytes.Equal(handshake.Protocol[:], Protocol[:]) {
		err := fmt.Sprintf("Protocol mismtach: got %s, expected %s", handshake.Protocol, Protocol)
		return errors.New(err)
	}
	if !bytes.Equal(handshake.InfoHash[:], infoHash) {
		err := fmt.Sprintf("Invalid infoHash: got %x, expected %x", handshake.InfoHash, infoHash)
		return errors.New(err)
	}
	return nil
}

func (p *Peer) decodeMessage(payload []byte) {
	if len(payload) == 0 {
		// keepalive
		return
	}

	switch payload[0] {
	case 0:
		// choke
		break
	case 1:
		// unchoke
		break
	case 2:
		// interested
		break
	case 3:
		// not interested
		break
	case 4:
		// have
		break
	case 5:
		// bitfield
		break
	case 6:
		// request
		break
	case 7:
		// piece
		break
	case 8:
		// cancel
		break
	case 9:
		// port
		break
	}
}

func (p *Peer) Reader() {
	log.Println("Peer : Reader : Started")

	var handshake Handshake
	binary.Read(p.conn, binary.BigEndian, &handshake)
	err := verifyHandshake(&handshake, p.infoHash)
	if err != nil {
		p.conn.Close()
		return
	}
	p.peerID = handshake.PeerID[:]

	for {
		length := make([]byte, 4)
		n, err := io.ReadFull(p.conn, length)
		if err != nil {
			return
		}
		payload := make([]byte, binary.BigEndian.Uint32(length))
		n, err = io.ReadFull(p.conn, payload)
		if err != nil {
			return
		}
		log.Printf("Read %d bytes of %x\n", n, payload)
		p.decodeMessage(payload)
		//p.read <- buf
	}
}

func (p *Peer) sendHandshake() {
	log.Println("Peer : sendHandshake : Started")
	defer log.Println("Peer : sendHandshake : Completed")

	handshake := Handshake{
		Len: uint8(len(Protocol)),
		Protocol: Protocol,
		PeerID: PeerID,
	}
	copy(handshake.InfoHash[:], p.infoHash)
	err := binary.Write(p.conn, binary.BigEndian, &handshake)
	if err != nil {
		// TODO: Handle errors
		log.Fatal(err)
	}
	p.stats.write += int(reflect.TypeOf(handshake).Size())
}

func (p *Peer) Stop() error {
	log.Println("Peer : Stop : Stopping")
	p.t.Kill(nil)
	return p.t.Wait()
}

func (p *Peer) Run() {
	log.Println("Peer : Run : Started")
	defer log.Println("Peer : Run : Completed")

	go p.sendHandshake()
	go p.Reader()

	for {
		select {
		case <-p.keepalive:
		case <-p.read:
			fmt.Println("p.read")
		//case buf := <-p.read:
			//fmt.Println("Read from peer:", buf)
		case <-p.t.Dying():
			p.peerManagerChans.deadPeer <- p.conn.RemoteAddr().String()
			return
		}
	}
}

func (pm *PeerManager) Stop() error {
	log.Println("PeerManager : Stop : Stopping")
	pm.t.Kill(nil)
	return pm.t.Wait()
}

func (pm *PeerManager) Run() {
	log.Println("PeerManager : Run : Started")
	defer pm.t.Done()
	defer log.Println("PeerManager : Run : Completed")

	for {
		select {
		case peer := <-pm.trackerChans.peers:
			peerName := fmt.Sprintf("%s:%d", peer.IP.String(), peer.Port)
			_, ok := pm.peers[peerName]
			if !ok {
				// FIXME Code in this block is duplicated below

				// Create the Controller->Peer chans struct
				contTxChans := *NewControllerPeerChans()

				// Construct the Peer object
				pm.peers[peerName] = NewPeer(
					pm.infoHash,
					true,
					pm.diskIOChans,
					contTxChans,
					pm.peerContChans)

				// Give the controller the channels that it will use to
				// transmit messages to this new peer
				go func() {
					pm.contChans.newPeer <- PeerComms{peerName: peerName, chans: contTxChans}
				}()

				// Have the 'peer' routine create an outbound
				// TCP connection to the remote peer
				go ConnectToPeer(peer, pm.serverChans.conns)
			}
		case conn := <-pm.serverChans.conns:
			_, ok := pm.peers[conn.RemoteAddr().String()]
			if !ok {
				// Create the Controller->Peer chans struct
				contTxChans := *NewControllerPeerChans()

				// Construct the Peer object
				peerName := conn.RemoteAddr().String()
				pm.peers[peerName] = NewPeer(
					pm.infoHash,
					false,
					pm.diskIOChans,
					contTxChans,
					pm.peerContChans)

				// Give the controller the channels that it will use to
				// transmit messages to this new peer
				go func() {
					pm.contChans.newPeer <- PeerComms{peerName: peerName, chans: contTxChans}
				}()
			}
			// Associate the connection with the peer object and start the peer
			pm.peers[conn.RemoteAddr().String()].conn = conn
			go pm.peers[conn.RemoteAddr().String()].Run()
		case peer := <-pm.peerChans.deadPeer:
			log.Printf("PeerManager : Deleting peer %s\n", peer)
			delete(pm.peers, peer)
		case <-pm.t.Dying():
			for _, peer := range pm.peers {
				peer.Stop()
			}
			return
		}
	}
}
<commit_msg>Added peerName to Peer struct to simplify code in situations where it's needed. Added logic for the HaveMessage case in peer.decodeMessage()<commit_after>// Copyright 2013 Jari Takkala and Brian Dignan. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"launchpad.net/tomb"
	"log"
	"net"
	"reflect"
	"sort"
	//"strconv"
	"syscall"
	"time"
)

var Protocol = [19]byte{'B', 'i', 't', 'T', 'o', 'r', 'r', 'e', 'n', 't', ' ', 'p', 'r', 'o', 't', 'o', 'c', 'o', 'l'}

// Message ID values
const (
	MsgChoke int = iota
	MsgUnchoke
	MsgInterested
	MsgNotInterested
	MsgHave
	MsgBitfield
	MsgRequest
	MsgPiece
	MsgCancel
	MsgPort
)

// PeerTuple represents a single IP+port pair of a peer
type PeerTuple struct {
	IP   net.IP
	Port uint16
}

type Peer struct {
	conn           *net.TCPConn
	peerName       string
	amChoking      bool
	amInterested   bool
	peerChoking    bool
	peerInterested bool
	ourBitfield    []bool
	peerBitfield   []bool
	initiator      bool
	peerID         []byte
	keepalive      <-chan time.Time // channel for sending keepalives
	lastTxKeepalive  time.Time
	lastRxKeepalive  time.Time
	read           chan []byte
	infoHash       []byte
	diskIOChans    diskIOPeerChans
	peerManagerChans peerManagerChans
	contRxChans    ControllerPeerChans
	contTxChans    PeerControllerChans
	stats          PeerStats
	t              tomb.Tomb
}

type PeerStats struct {
	read int
	write int
	errors int
}

type PeerManager struct {
	peers        map[string]*Peer
	infoHash     []byte
	peerChans    peerManagerChans
	serverChans  serverPeerChans
	trackerChans trackerPeerChans
	diskIOChans  diskIOPeerChans
	contChans    ControllerPeerManagerChans
	peerContChans PeerControllerChans
	t            tomb.Tomb
}

type peerManagerChans struct {
	deadPeer chan string
}

type PeerComms struct {
	peerName     string
	chans 		 ControllerPeerChans
}

func NewPeerComms(peerName string, cpc ControllerPeerChans) *PeerComms {
	pc := new(PeerComms)
	pc.peerName = peerName
	pc.chans = cpc
	return pc
}

type PeerInfo struct {
	peerName        string
	isChoked        bool // The peer is connected but choked. Defaults to TRUE (choked)
	availablePieces []bool
	activeRequests  map[int]struct{}
	qtyPiecesNeeded int                   // The quantity of pieces that this peer has that we haven't yet downloaded.
	chans 			ControllerPeerChans
}

type Handshake struct {
	Len      uint8
	Protocol [19]byte
	Reserved [8]uint8
	InfoHash [20]byte
	PeerID   [20]byte
}

func NewPeerInfo(quantityOfPieces int, peerComms PeerComms) *PeerInfo {
	pi := new(PeerInfo)

	pi.peerName = peerComms.peerName
	pi.chans = peerComms.chans

	pi.isChoked = true // By default, a peer starts as being choked by the other side.
	pi.availablePieces = make([]bool, quantityOfPieces)
	pi.activeRequests = make(map[int]struct{})

	return pi
}

// Sent by the peer to controller indicating a 'choke' state change. It either went from unchoked to choked,
// or from choked to unchoked.
type PeerChokeStatus struct {
	peerName   string
	isChoked bool
}

type SortedPeers []*PeerInfo

func (sp SortedPeers) Less(i, j int) bool {
	return sp[i].qtyPiecesNeeded <= sp[j].qtyPiecesNeeded
}

func (sp SortedPeers) Swap(i, j int) {
	tmp := sp[i]
	sp[i] = sp[j]
	sp[j] = tmp
}

func (sp SortedPeers) Len() int {
	return len(sp)
}

func sortedPeersByQtyPiecesNeeded(peers map[string]*PeerInfo) SortedPeers {
	peerInfoSlice := make(SortedPeers, 0)

	for _, peerInfo := range peers {
		peerInfoSlice = append(peerInfoSlice, peerInfo)
	}
	sort.Sort(peerInfoSlice)

	return peerInfoSlice
}

func NewPeerManager(infoHash []byte, diskIOChans diskIOPeerChans, serverChans serverPeerChans, trackerChans trackerPeerChans) *PeerManager {
	pm := new(PeerManager)
	pm.infoHash = infoHash
	pm.diskIOChans = diskIOChans
	pm.serverChans = serverChans
	pm.trackerChans = trackerChans
	pm.peerChans.deadPeer = make(chan string)
	pm.peers = make(map[string]*Peer)
	pm.contChans.newPeer = make(chan PeerComms)
	pm.contChans.deadPeer = make(chan string)
	pm.peerContChans.chokeStatus = make(chan PeerChokeStatus)
	pm.peerContChans.havePiece = make(chan chan HavePiece)
	return pm
}

func ConnectToPeer(peerTuple PeerTuple, connCh chan *net.TCPConn) {
	raddr := net.TCPAddr{peerTuple.IP, int(peerTuple.Port), ""}
	log.Println("Connecting to", raddr)
	conn, err := net.DialTCP("tcp4", nil, &raddr)
	if err != nil {
		if e, ok := err.(*net.OpError); ok {
			if e.Err == syscall.ECONNREFUSED {
				log.Println("ConnectToPeer : Connection Refused:", raddr)
				return
			}
		}
		log.Fatal(err)
	}
	log.Println("ConnectToPeer : Connected:", raddr)
	connCh <- conn
}

func NewPeer(
			peerName string,
			infoHash []byte,
			initiator bool,
			diskIOChans diskIOPeerChans,
			contRxChans ControllerPeerChans,
			contTxChans PeerControllerChans) *Peer {
	p := &Peer{
			peerName: peerName,
			infoHash: infoHash,
			amChoking: true,
			amInterested: false,
			peerChoking: true,
			peerInterested: false,
			initiator: initiator,
			diskIOChans: diskIOChans,
			contRxChans: contRxChans,
			contTxChans: contTxChans}
	p.read = make(chan []byte)
	return p
}

func constructMessage(id int, payload []byte) (msg []byte, err error) {
	msg = make([]byte, 4)

	// Store the length of payload + id in network byte order
	binary.BigEndian.PutUint32(msg, uint32(len(payload) + 1))
	msg = append(msg, byte(id))
	msg = append(msg, payload...)

	return
}

func verifyHandshake(handshake *Handshake, infoHash []byte) error {
	if int(handshake.Len) != len(Protocol) {
		err := fmt.Sprintf("Unexpected length for pstrlen (wanted %d, got %d)", len(Protocol), int(handshake.Len))
		return errors.New(err)
	}
	if !bytes.Equal(handshake.Protocol[:], Protocol[:]) {
		err := fmt.Sprintf("Protocol mismtach: got %s, expected %s", handshake.Protocol, Protocol)
		return errors.New(err)
	}
	if !bytes.Equal(handshake.InfoHash[:], infoHash) {
		err := fmt.Sprintf("Invalid infoHash: got %x, expected %x", handshake.InfoHash, infoHash)
		return errors.New(err)
	}
	return nil
}

func (p *Peer) sendHaveToController(pieces []HavePiece) {

}

func (p *Peer) readBytesFromConn(numBytes int) []byte {
	result := make([]byte, numBytes)
	_, err := io.ReadFull(p.conn, result)
	if err != nil {
		log.Fatalf("Encountered an error when attempting to read %d bytes from %s", numBytes, p.peerName)
	}
	return result
}

func (p *Peer) decodeMessage(payload []byte) {
	if len(payload) == 0 {
		// keepalive
		return
	}

	messageID := binary.BigEndian.Uint32(p.readBytesFromConn(1))

	switch messageID {
	case 0:
		// Choke Message
		break
	case 1:
		// Unchoke Message
		break
	case 2:
		// Interested Message
		break
	case 3:
		// Not Interested Message
		break
	case 4:
		// Have Message

		// Determine the piece number
		pieceNum := int(binary.BigEndian.Uint32(p.readBytesFromConn(4)))

		// Update the local peer bitfield
		p.peerBitfield[pieceNum] = true

		// Send a single HavePiece struct to the controller
		have := make([]HavePiece, 1)
		have[1] = HavePiece{pieceNum: pieceNum, peerName: p.peerName}
		break
	case 5:
		// Bitfield Message

		// Record the peer bitfield locally

		// Break the bitfield into a slice of HavePiece structs and send them
		// to the controller
		break
	case 6:
		// Request Message
		break
	case 7:
		// Piece Message
		break
	case 8:
		// Cancel Message
		break
	case 9:
		// Port Message
		break
	}
}

func (p *Peer) Reader() {
	log.Println("Peer : Reader : Started")

	var handshake Handshake
	binary.Read(p.conn, binary.BigEndian, &handshake)
	err := verifyHandshake(&handshake, p.infoHash)
	if err != nil {
		p.conn.Close()
		return
	}
	p.peerID = handshake.PeerID[:]

	for {
		length := make([]byte, 4)
		n, err := io.ReadFull(p.conn, length)
		if err != nil {
			// FIXME Are there any cases where we would not read length?
			return
		}
		payload := make([]byte, binary.BigEndian.Uint32(length))
		n, err = io.ReadFull(p.conn, payload)
		if err != nil {
			// FIXME if this is not a keepalive, we should
			// definitely get a payload
			return
		}
		log.Printf("Read %d bytes of %x\n", n, payload)
		p.decodeMessage(payload)
		//p.read <- buf
	}
}

func (p *Peer) sendHandshake() {
	log.Println("Peer : sendHandshake : Started")
	defer log.Println("Peer : sendHandshake : Completed")

	handshake := Handshake{
		Len: uint8(len(Protocol)),
		Protocol: Protocol,
		PeerID: PeerID,
	}
	copy(handshake.InfoHash[:], p.infoHash)
	err := binary.Write(p.conn, binary.BigEndian, &handshake)
	if err != nil {
		// TODO: Handle errors
		log.Fatal(err)
	}
	p.stats.write += int(reflect.TypeOf(handshake).Size())
}

func (p *Peer) Stop() error {
	log.Println("Peer : Stop : Stopping")
	p.t.Kill(nil)
	return p.t.Wait()
}

func (p *Peer) Run() {
	log.Println("Peer : Run : Started")
	defer log.Println("Peer : Run : Completed")

	go p.sendHandshake()
	go p.Reader()

	for {
		select {
		case <-p.keepalive:
		case <-p.read:
			fmt.Println("p.read")
		//case buf := <-p.read:
			//fmt.Println("Read from peer:", buf)
		case <-p.t.Dying():
			p.peerManagerChans.deadPeer <- p.conn.RemoteAddr().String()
			return
		}
	}
}

func (pm *PeerManager) Stop() error {
	log.Println("PeerManager : Stop : Stopping")
	pm.t.Kill(nil)
	return pm.t.Wait()
}

func (pm *PeerManager) Run() {
	log.Println("PeerManager : Run : Started")
	defer pm.t.Done()
	defer log.Println("PeerManager : Run : Completed")

	for {
		select {
		case peer := <-pm.trackerChans.peers:
			peerName := fmt.Sprintf("%s:%d", peer.IP.String(), peer.Port)
			_, ok := pm.peers[peerName]
			if !ok {
				// FIXME Code in this block is duplicated below

				// Create the Controller->Peer chans struct
				contTxChans := *NewControllerPeerChans()

				// Construct the Peer object
				pm.peers[peerName] = NewPeer(
					peerName,
					pm.infoHash,
					true,
					pm.diskIOChans,
					contTxChans,
					pm.peerContChans)

				// Give the controller the channels that it will use to
				// transmit messages to this new peer
				go func() {
					pm.contChans.newPeer <- PeerComms{peerName: peerName, chans: contTxChans}
				}()

				// Have the 'peer' routine create an outbound
				// TCP connection to the remote peer
				go ConnectToPeer(peer, pm.serverChans.conns)
			}
		case conn := <-pm.serverChans.conns:
			_, ok := pm.peers[conn.RemoteAddr().String()]
			if !ok {
				// Create the Controller->Peer chans struct
				contTxChans := *NewControllerPeerChans()

				// Construct the Peer object
				peerName := conn.RemoteAddr().String()
				pm.peers[peerName] = NewPeer(
					peerName,
					pm.infoHash,
					false,
					pm.diskIOChans,
					contTxChans,
					pm.peerContChans)

				// Give the controller the channels that it will use to
				// transmit messages to this new peer
				go func() {
					pm.contChans.newPeer <- PeerComms{peerName: peerName, chans: contTxChans}
				}()
			}
			// Associate the connection with the peer object and start the peer
			pm.peers[conn.RemoteAddr().String()].conn = conn
			go pm.peers[conn.RemoteAddr().String()].Run()
		case peer := <-pm.peerChans.deadPeer:
			log.Printf("PeerManager : Deleting peer %s\n", peer)
			delete(pm.peers, peer)
		case <-pm.t.Dying():
			for _, peer := range pm.peers {
				peer.Stop()
			}
			return
		}
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package testutil

import (
	"fmt"
	"os"
	"reflect"
	"runtime/pprof"
	"sync"
	"time"

	federation_api "k8s.io/kubernetes/federation/apis/federation/v1beta1"
	"k8s.io/kubernetes/federation/pkg/federation-controller/util"
	"k8s.io/kubernetes/pkg/api"
	api_v1 "k8s.io/kubernetes/pkg/api/v1"
	"k8s.io/kubernetes/pkg/client/testing/core"
	"k8s.io/kubernetes/pkg/runtime"
	"k8s.io/kubernetes/pkg/util/wait"
	"k8s.io/kubernetes/pkg/watch"

	"github.com/golang/glog"
)

// A structure that distributes eventes to multiple watchers.
type WatcherDispatcher struct {
	sync.Mutex
	watchers    []*watch.RaceFreeFakeWatcher
	eventsSoFar []*watch.Event
}

func (wd *WatcherDispatcher) register(watcher *watch.RaceFreeFakeWatcher) {
	wd.Lock()
	defer wd.Unlock()
	wd.watchers = append(wd.watchers, watcher)
	for _, event := range wd.eventsSoFar {
		watcher.Action(event.Type, event.Object)
	}
}

func (wd *WatcherDispatcher) Stop() {
	wd.Lock()
	defer wd.Unlock()
	for _, watcher := range wd.watchers {
		watcher.Stop()
	}
}

func copy(obj runtime.Object) runtime.Object {
	objCopy, err := api.Scheme.DeepCopy(obj)
	if err != nil {
		panic(err)
	}
	return objCopy.(runtime.Object)
}

// Add sends an add event.
func (wd *WatcherDispatcher) Add(obj runtime.Object) {
	wd.Lock()
	defer wd.Unlock()
	wd.eventsSoFar = append(wd.eventsSoFar, &watch.Event{Type: watch.Added, Object: copy(obj)})
	for _, watcher := range wd.watchers {
		if !watcher.IsStopped() {
			watcher.Add(copy(obj))
		}
	}
}

// Modify sends a modify event.
func (wd *WatcherDispatcher) Modify(obj runtime.Object) {
	wd.Lock()
	defer wd.Unlock()
	glog.V(4).Infof("->WatcherDispatcher.Modify(%v)", obj)
	wd.eventsSoFar = append(wd.eventsSoFar, &watch.Event{Type: watch.Modified, Object: copy(obj)})
	for i, watcher := range wd.watchers {
		if !watcher.IsStopped() {
			glog.V(4).Infof("->Watcher(%d).Modify(%v)", i, obj)
			watcher.Modify(copy(obj))
		} else {
			glog.V(4).Infof("->Watcher(%d) is stopped.  Not calling Modify(%v)", i, obj)
		}
	}
}

// Delete sends a delete event.
func (wd *WatcherDispatcher) Delete(lastValue runtime.Object) {
	wd.Lock()
	defer wd.Unlock()
	wd.eventsSoFar = append(wd.eventsSoFar, &watch.Event{Type: watch.Deleted, Object: copy(lastValue)})
	for _, watcher := range wd.watchers {
		if !watcher.IsStopped() {
			watcher.Delete(copy(lastValue))
		}
	}
}

// Error sends an Error event.
func (wd *WatcherDispatcher) Error(errValue runtime.Object) {
	wd.Lock()
	defer wd.Unlock()
	wd.eventsSoFar = append(wd.eventsSoFar, &watch.Event{Type: watch.Error, Object: copy(errValue)})
	for _, watcher := range wd.watchers {
		if !watcher.IsStopped() {
			watcher.Error(copy(errValue))
		}
	}
}

// Action sends an event of the requested type, for table-based testing.
func (wd *WatcherDispatcher) Action(action watch.EventType, obj runtime.Object) {
	wd.Lock()
	defer wd.Unlock()
	wd.eventsSoFar = append(wd.eventsSoFar, &watch.Event{Type: action, Object: copy(obj)})
	for _, watcher := range wd.watchers {
		if !watcher.IsStopped() {
			watcher.Action(action, copy(obj))
		}
	}
}

// RegisterFakeWatch adds a new fake watcher for the specified resource in the given fake client.
// All subsequent requests for a watch on the client will result in returning this fake watcher.
func RegisterFakeWatch(resource string, client *core.Fake) *WatcherDispatcher {
	dispatcher := &WatcherDispatcher{
		watchers:    make([]*watch.RaceFreeFakeWatcher, 0),
		eventsSoFar: make([]*watch.Event, 0),
	}

	client.AddWatchReactor(resource, func(action core.Action) (bool, watch.Interface, error) {
		watcher := watch.NewRaceFreeFake()
		dispatcher.register(watcher)
		return true, watcher, nil
	})
	return dispatcher
}

// RegisterFakeList registers a list response for the specified resource inside the given fake client.
// The passed value will be returned with every list call.
func RegisterFakeList(resource string, client *core.Fake, obj runtime.Object) {
	client.AddReactor("list", resource, func(action core.Action) (bool, runtime.Object, error) {
		return true, obj, nil
	})
}

// RegisterFakeCopyOnCreate registers a reactor in the given fake client that passes
// all created objects to the given watcher and also copies them to a channel for
// in-test inspection.
func RegisterFakeCopyOnCreate(resource string, client *core.Fake, watcher *WatcherDispatcher) chan runtime.Object {
	objChan := make(chan runtime.Object, 100)
	client.AddReactor("create", resource, func(action core.Action) (bool, runtime.Object, error) {
		createAction := action.(core.CreateAction)
		originalObj := createAction.GetObject()
		// Create a copy of the object here to prevent data races while reading the object in go routine.
		obj := copy(originalObj)
		go func() {
			glog.V(4).Infof("Object created. Writing to channel: %v", obj)
			watcher.Add(obj)
			objChan <- obj
		}()
		return true, originalObj, nil
	})
	return objChan
}

// RegisterFakeCopyOnUpdate registers a reactor in the given fake client that passes
// all updated objects to the given watcher and also copies them to a channel for
// in-test inspection.
func RegisterFakeCopyOnUpdate(resource string, client *core.Fake, watcher *WatcherDispatcher) chan runtime.Object {
	objChan := make(chan runtime.Object, 100)
	client.AddReactor("update", resource, func(action core.Action) (bool, runtime.Object, error) {
		updateAction := action.(core.UpdateAction)
		originalObj := updateAction.GetObject()
		// Create a copy of the object here to prevent data races while reading the object in go routine.
		obj := copy(originalObj)
		go func() {
			glog.V(4).Infof("Object updated. Writing to channel: %v", obj)
			watcher.Modify(obj)
			objChan <- obj
		}()
		return true, originalObj, nil
	})
	return objChan
}

// GetObjectFromChan tries to get an api object from the given channel
// within a reasonable time.
func GetObjectFromChan(c chan runtime.Object) runtime.Object {
	select {
	case obj := <-c:
		return obj
	case <-time.After(wait.ForeverTestTimeout):
		pprof.Lookup("goroutine").WriteTo(os.Stderr, 1)
		return nil
	}
}

type CheckingFunction func(runtime.Object) error

// CheckObjectFromChan tries to get an object matching the given check function
// within a reasonable time.
func CheckObjectFromChan(c chan runtime.Object, checkFunction CheckingFunction) error {
	delay := 20 * time.Second
	var lastError error
	for {
		select {
		case obj := <-c:
			if lastError = checkFunction(obj); lastError == nil {
				return nil
			}
			glog.Infof("Check function failed with %v", lastError)
			delay = 5 * time.Second
		case <-time.After(delay):
			pprof.Lookup("goroutine").WriteTo(os.Stderr, 1)
			if lastError == nil {
				return fmt.Errorf("Failed to get an object from channel")
			} else {
				return lastError
			}
		}
	}
}

// CompareObjectMeta returns an error when the given objects are not equivalent.
func CompareObjectMeta(a, b api_v1.ObjectMeta) error {
	if a.Namespace != b.Namespace {
		return fmt.Errorf("Different namespace expected:%s observed:%s", a.Namespace, b.Namespace)
	}
	if a.Name != b.Name {
		return fmt.Errorf("Different name expected:%s observed:%s", a.Namespace, b.Namespace)
	}
	if !reflect.DeepEqual(a.Annotations, b.Annotations) {
		return fmt.Errorf("Annotations are different expected:%v observerd:%v", a.Annotations, b.Annotations)
	}
	if !reflect.DeepEqual(a.Labels, b.Labels) {
		return fmt.Errorf("Annotations are different expected:%v observerd:%v", a.Labels, b.Labels)
	}
	return nil
}

func ToFederatedInformerForTestOnly(informer util.FederatedInformer) util.FederatedInformerForTestOnly {
	inter := informer.(interface{})
	return inter.(util.FederatedInformerForTestOnly)
}

// NewCluster builds a new cluster object.
func NewCluster(name string, readyStatus api_v1.ConditionStatus) *federation_api.Cluster {
	return &federation_api.Cluster{
		ObjectMeta: api_v1.ObjectMeta{
			Name:        name,
			Annotations: map[string]string{},
		},
		Status: federation_api.ClusterStatus{
			Conditions: []federation_api.ClusterCondition{
				{Type: federation_api.ClusterReady, Status: readyStatus},
			},
		},
	}
}

// Ensure a key is in the store before returning (or timeout w/ error)
func WaitForStoreUpdate(store util.FederatedReadOnlyStore, clusterName, key string, timeout time.Duration) error {
	retryInterval := 100 * time.Millisecond
	err := wait.PollImmediate(retryInterval, timeout, func() (bool, error) {
		_, found, err := store.GetByKey(clusterName, key)
		return found, err
	})
	return err
}
<commit_msg>Ensure proper serialization of updates and creates in federation test watcher<commit_after>/*
Copyright 2016 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package testutil

import (
	"fmt"
	"os"
	"reflect"
	"runtime/pprof"
	"sync"
	"time"

	federation_api "k8s.io/kubernetes/federation/apis/federation/v1beta1"
	"k8s.io/kubernetes/federation/pkg/federation-controller/util"
	"k8s.io/kubernetes/pkg/api"
	api_v1 "k8s.io/kubernetes/pkg/api/v1"
	"k8s.io/kubernetes/pkg/client/testing/core"
	"k8s.io/kubernetes/pkg/runtime"
	"k8s.io/kubernetes/pkg/util/wait"
	"k8s.io/kubernetes/pkg/watch"

	"github.com/golang/glog"
)

// A structure that distributes eventes to multiple watchers.
type WatcherDispatcher struct {
	sync.Mutex
	watchers       []*watch.RaceFreeFakeWatcher
	eventsSoFar    []*watch.Event
	orderExecution chan func()
	stopChan       chan struct{}
}

func (wd *WatcherDispatcher) register(watcher *watch.RaceFreeFakeWatcher) {
	wd.Lock()
	defer wd.Unlock()
	wd.watchers = append(wd.watchers, watcher)
	for _, event := range wd.eventsSoFar {
		watcher.Action(event.Type, event.Object)
	}
}

func (wd *WatcherDispatcher) Stop() {
	wd.Lock()
	defer wd.Unlock()
	close(wd.stopChan)
	for _, watcher := range wd.watchers {
		watcher.Stop()
	}
}

func copy(obj runtime.Object) runtime.Object {
	objCopy, err := api.Scheme.DeepCopy(obj)
	if err != nil {
		panic(err)
	}
	return objCopy.(runtime.Object)
}

// Add sends an add event.
func (wd *WatcherDispatcher) Add(obj runtime.Object) {
	wd.Lock()
	defer wd.Unlock()
	wd.eventsSoFar = append(wd.eventsSoFar, &watch.Event{Type: watch.Added, Object: copy(obj)})
	for _, watcher := range wd.watchers {
		if !watcher.IsStopped() {
			watcher.Add(copy(obj))
		}
	}
}

// Modify sends a modify event.
func (wd *WatcherDispatcher) Modify(obj runtime.Object) {
	wd.Lock()
	defer wd.Unlock()
	glog.V(4).Infof("->WatcherDispatcher.Modify(%v)", obj)
	wd.eventsSoFar = append(wd.eventsSoFar, &watch.Event{Type: watch.Modified, Object: copy(obj)})
	for i, watcher := range wd.watchers {
		if !watcher.IsStopped() {
			glog.V(4).Infof("->Watcher(%d).Modify(%v)", i, obj)
			watcher.Modify(copy(obj))
		} else {
			glog.V(4).Infof("->Watcher(%d) is stopped.  Not calling Modify(%v)", i, obj)
		}
	}
}

// Delete sends a delete event.
func (wd *WatcherDispatcher) Delete(lastValue runtime.Object) {
	wd.Lock()
	defer wd.Unlock()
	wd.eventsSoFar = append(wd.eventsSoFar, &watch.Event{Type: watch.Deleted, Object: copy(lastValue)})
	for _, watcher := range wd.watchers {
		if !watcher.IsStopped() {
			watcher.Delete(copy(lastValue))
		}
	}
}

// Error sends an Error event.
func (wd *WatcherDispatcher) Error(errValue runtime.Object) {
	wd.Lock()
	defer wd.Unlock()
	wd.eventsSoFar = append(wd.eventsSoFar, &watch.Event{Type: watch.Error, Object: copy(errValue)})
	for _, watcher := range wd.watchers {
		if !watcher.IsStopped() {
			watcher.Error(copy(errValue))
		}
	}
}

// Action sends an event of the requested type, for table-based testing.
func (wd *WatcherDispatcher) Action(action watch.EventType, obj runtime.Object) {
	wd.Lock()
	defer wd.Unlock()
	wd.eventsSoFar = append(wd.eventsSoFar, &watch.Event{Type: action, Object: copy(obj)})
	for _, watcher := range wd.watchers {
		if !watcher.IsStopped() {
			watcher.Action(action, copy(obj))
		}
	}
}

// RegisterFakeWatch adds a new fake watcher for the specified resource in the given fake client.
// All subsequent requests for a watch on the client will result in returning this fake watcher.
func RegisterFakeWatch(resource string, client *core.Fake) *WatcherDispatcher {
	dispatcher := &WatcherDispatcher{
		watchers:       make([]*watch.RaceFreeFakeWatcher, 0),
		eventsSoFar:    make([]*watch.Event, 0),
		orderExecution: make(chan func()),
		stopChan:       make(chan struct{}),
	}
	go func() {
		for {
			select {
			case fun := <-dispatcher.orderExecution:
				fun()
			case <-dispatcher.stopChan:
				return
			}
		}
	}()

	client.AddWatchReactor(resource, func(action core.Action) (bool, watch.Interface, error) {
		watcher := watch.NewRaceFreeFake()
		dispatcher.register(watcher)
		return true, watcher, nil
	})
	return dispatcher
}

// RegisterFakeList registers a list response for the specified resource inside the given fake client.
// The passed value will be returned with every list call.
func RegisterFakeList(resource string, client *core.Fake, obj runtime.Object) {
	client.AddReactor("list", resource, func(action core.Action) (bool, runtime.Object, error) {
		return true, obj, nil
	})
}

// RegisterFakeCopyOnCreate registers a reactor in the given fake client that passes
// all created objects to the given watcher and also copies them to a channel for
// in-test inspection.
func RegisterFakeCopyOnCreate(resource string, client *core.Fake, watcher *WatcherDispatcher) chan runtime.Object {
	objChan := make(chan runtime.Object, 100)
	client.AddReactor("create", resource, func(action core.Action) (bool, runtime.Object, error) {
		createAction := action.(core.CreateAction)
		originalObj := createAction.GetObject()
		// Create a copy of the object here to prevent data races while reading the object in go routine.
		obj := copy(originalObj)
		watcher.orderExecution <- func() {
			glog.V(4).Infof("Object created. Writing to channel: %v", obj)
			watcher.Add(obj)
			objChan <- obj
		}
		return true, originalObj, nil
	})
	return objChan
}

// RegisterFakeCopyOnUpdate registers a reactor in the given fake client that passes
// all updated objects to the given watcher and also copies them to a channel for
// in-test inspection.
func RegisterFakeCopyOnUpdate(resource string, client *core.Fake, watcher *WatcherDispatcher) chan runtime.Object {
	objChan := make(chan runtime.Object, 100)
	client.AddReactor("update", resource, func(action core.Action) (bool, runtime.Object, error) {
		updateAction := action.(core.UpdateAction)
		originalObj := updateAction.GetObject()
		// Create a copy of the object here to prevent data races while reading the object in go routine.
		obj := copy(originalObj)
		watcher.orderExecution <- func() {
			glog.V(4).Infof("Object updated. Writing to channel: %v", obj)
			watcher.Modify(obj)
			objChan <- obj
		}
		return true, originalObj, nil
	})
	return objChan
}

// GetObjectFromChan tries to get an api object from the given channel
// within a reasonable time.
func GetObjectFromChan(c chan runtime.Object) runtime.Object {
	select {
	case obj := <-c:
		return obj
	case <-time.After(wait.ForeverTestTimeout):
		pprof.Lookup("goroutine").WriteTo(os.Stderr, 1)
		return nil
	}
}

type CheckingFunction func(runtime.Object) error

// CheckObjectFromChan tries to get an object matching the given check function
// within a reasonable time.
func CheckObjectFromChan(c chan runtime.Object, checkFunction CheckingFunction) error {
	delay := 20 * time.Second
	var lastError error
	for {
		select {
		case obj := <-c:
			if lastError = checkFunction(obj); lastError == nil {
				return nil
			}
			glog.Infof("Check function failed with %v", lastError)
			delay = 5 * time.Second
		case <-time.After(delay):
			pprof.Lookup("goroutine").WriteTo(os.Stderr, 1)
			if lastError == nil {
				return fmt.Errorf("Failed to get an object from channel")
			} else {
				return lastError
			}
		}
	}
}

// CompareObjectMeta returns an error when the given objects are not equivalent.
func CompareObjectMeta(a, b api_v1.ObjectMeta) error {
	if a.Namespace != b.Namespace {
		return fmt.Errorf("Different namespace expected:%s observed:%s", a.Namespace, b.Namespace)
	}
	if a.Name != b.Name {
		return fmt.Errorf("Different name expected:%s observed:%s", a.Namespace, b.Namespace)
	}
	if !reflect.DeepEqual(a.Annotations, b.Annotations) {
		return fmt.Errorf("Annotations are different expected:%v observerd:%v", a.Annotations, b.Annotations)
	}
	if !reflect.DeepEqual(a.Labels, b.Labels) {
		return fmt.Errorf("Annotations are different expected:%v observerd:%v", a.Labels, b.Labels)
	}
	return nil
}

func ToFederatedInformerForTestOnly(informer util.FederatedInformer) util.FederatedInformerForTestOnly {
	inter := informer.(interface{})
	return inter.(util.FederatedInformerForTestOnly)
}

// NewCluster builds a new cluster object.
func NewCluster(name string, readyStatus api_v1.ConditionStatus) *federation_api.Cluster {
	return &federation_api.Cluster{
		ObjectMeta: api_v1.ObjectMeta{
			Name:        name,
			Annotations: map[string]string{},
		},
		Status: federation_api.ClusterStatus{
			Conditions: []federation_api.ClusterCondition{
				{Type: federation_api.ClusterReady, Status: readyStatus},
			},
		},
	}
}

// Ensure a key is in the store before returning (or timeout w/ error)
func WaitForStoreUpdate(store util.FederatedReadOnlyStore, clusterName, key string, timeout time.Duration) error {
	retryInterval := 100 * time.Millisecond
	err := wait.PollImmediate(retryInterval, timeout, func() (bool, error) {
		_, found, err := store.GetByKey(clusterName, key)
		return found, err
	})
	return err
}
<|endoftext|>


<commit_before>package assert

import (
	"fmt"
	"strings"
	"testing"
)

// MatchType is used in MockT's CheckPredicates to control whether
// recorded testing.TB operations match predicates.
type MatchType int

const (
	// Success indicates a predicate matches in CheckPredicates
	Success MatchType = iota

	// Failure indicates a predicate does not match in CheckPredicates
	Failure
)

// MockTestOperation represents an invocation of a method on
// testing.TB and its arguments.
type MockTestOperation struct {
	// Name represents the name of an operation on TB. One of: Error,
	// Error, Fatal, Fatalf, Fail, FailNow, Log, Logf, Skip, SkipNow,
	// Skipf.
	Name string

	// Args are the formatted arguments of the operation. For the
	// formatted variants (e.g., Errorf), this is the result of invoke
	// fmt.Sprintf with the format string and arguments. For the
	// non-formatted variables, this is the result of the arguments,
	// converted to strings as in fmt.Sprint, with spaces between each
	// one argument. Finally for the no-arg operations this is an
	// empty string.
	Args string
}

// String provides a string representation of MockTestOperation.
func (op *MockTestOperation) String() string {
	return fmt.Sprintf("%s(%q)", op.Name, op.Args)
}

// Predicate represents a predicate testing a
// MockTestOperation. Typically Predicates are composed from a
// NamePredicate and an ArgPredicate via Match.
type Predicate func(MockTestOperation) MatchType

// NamePredicate compares the given operation name (the name of a
// testing.TB method) to determine if it matches the predicate.
type NamePredicate func(string) MatchType

// OneOfOp returns a NamePredicate requires its op to equal one of the
// given operations.
func OneOfOp(expectedOps ...string) NamePredicate {
	return func(op string) MatchType {
		for _, eop := range expectedOps {
			if op == eop {
				return Success
			}
		}
		return Failure
	}
}

// ExactOp is equivalent to OneOfOp(op).
func ExactOp(op string) NamePredicate { return OneOfOp(op) }

// ErrorOp matches Error or Errorf operations.
func ErrorOp() NamePredicate { return OneOfOp("Error", "Errorf") }

// FatalOp matches Fatal or Fatalf operations.
func FatalOp() NamePredicate { return OneOfOp("Fatal", "Fatalf") }

// Fails matches any testing.TB operation that causes a test failure.
func Fails() NamePredicate {
	return OneOfOp("Error", "Errorf", "Fatal", "Fatalf", "Fail", "FailNow")
}

// ArgsPredicate tests the given operation arguments to determine if
// they match the predicate.
type ArgsPredicate func(string) MatchType

// PrefixedArgs requires the operation arguments to match the given
// prefix.
func PrefixedArgs(prefix string) ArgsPredicate {
	return func(args string) MatchType {
		if strings.HasPrefix(args, prefix) {
			return Success
		}
		return Failure
	}
}

// ArgsContain requires the operation arguments to contain the given
// expression.
func ArgsContain(expr string) ArgsPredicate {
	return func(args string) MatchType {
		if strings.Contains(args, expr) {
			return Success
		}
		return Failure
	}
}

// Match creates a Predicate from the given NamePredicate and
// ArgPredicate. The NamePredicate is tested first. If it fails, the
// entire Predicate fails. Otherwise, the ArgPredicate is tested it's
// result is returned.
func Match(namePred NamePredicate, argPred ArgsPredicate) Predicate {
	return func(op MockTestOperation) MatchType {
		if m := namePred(op.Name); m != Success {
			return m
		}

		return argPred(op.Args)
	}
}

// MockT is an implementation of testing.TB that allows testing of
// code that makes test assertions. Calls to Fatal, Fatalf, FailNow
// and SkipNow are recorded, but do not terminate execution. If the
// code under test depends on termination of execution, unexpected
// results may occur.
type MockT struct {
	testing.TB

	Operations []MockTestOperation

	failed  bool
	skipped bool
	helper  bool
}

func (t *MockT) record(op string, args string) {
	t.Operations = append(t.Operations, MockTestOperation{op, args})
}

func (t *MockT) operationsToString() string {
	s := make([]string, 0, len(t.Operations))
	for _, e := range t.Operations {
		s = append(s, e.String())
	}
	return strings.Join(s, "\n")
}

// Reset resets the MockT's log of operations.
func (t *MockT) Reset() {
	t.Operations = nil
	t.failed = false
	t.skipped = false
}

// CheckHelper fails the test if the Helper method was not invoked on
// this MockT.
func (t *MockT) CheckHelper(realT testing.TB) {
	if !t.helper {
		Tracing(realT).Error("expected Helper invocation, but did not see it")
	}
}

// CheckPredicates evaluates the given predicates against the recorded
// testing.TB operations. Predicates must match operations in order.
func (t *MockT) CheckPredicates(realT testing.TB, predicates ...Predicate) {
	nops := len(t.Operations)
	npreds := len(predicates)

	n := nops
	if npreds < nops {
		n = npreds
	}

	errors := []string{}
	for i := 0; i < n; i++ {
		op := t.Operations[i]
		if predicates[i](op) != Success {
			errors = append(
				errors,
				fmt.Sprintf("operation %d, %s: did not match predicate", i+1, op.String()),
			)
		}
	}

	if nops != npreds {
		Tracing(realT).Errorf(
			"Expected %d operations, but got %d:\n%s",
			npreds,
			nops,
			t.operationsToString(),
		)
	}

	if len(errors) > 0 {
		Tracing(realT).Errorf("Predicate failure(s):\n%s", strings.Join(errors, "\n"))
	}
}

// CheckSuccess fails the test if any operations were logged.
func (t *MockT) CheckSuccess(realT testing.TB) {
	if len(t.Operations) > 0 {
		Tracing(realT).Errorf(
			"expected no testing.TB operations, but saw:\n%s",
			t.operationsToString(),
		)
	}
}

func fmtArgs(args ...interface{}) string {
	// Sprintln (vs Sprint) always adds spaces, but we don't want the
	// trailing newline.
	s := fmt.Sprintln(args...)
	return s[:len(s)-1]
}

// Name returns "MockTest", always.
func (t *MockT) Name() string { return "MockTest" }

// Error records an Error operation.
func (t *MockT) Error(args ...interface{}) {
	t.record("Error", fmtArgs(args...))
}

// Errorf records an Errorf operation.
func (t *MockT) Errorf(format string, args ...interface{}) {
	t.record("Errorf", fmt.Sprintf(format, args...))
}

// Fail records a Fail operation.
func (t *MockT) Fail() {
	t.record("Fail", "")
	t.failed = true
}

// FailNow records a FailNow operation.
func (t *MockT) FailNow() {
	t.record("FailNow", "")
	t.failed = true
}

// Failed reports whether Fail or FailNow was previously invoked.
func (t *MockT) Failed() bool {
	return t.failed
}

// Fatal records a Fatal operation.
func (t *MockT) Fatal(args ...interface{}) {
	t.record("Fatal", fmtArgs(args...))
}

// Fatalf records a Fatalf operation.
func (t *MockT) Fatalf(format string, args ...interface{}) {
	t.record("Fatalf", fmt.Sprintf(format, args...))
}

// Log records a Log operation.
func (t *MockT) Log(args ...interface{}) {
	t.record("Log", fmtArgs(args...))
}

// Logf records a Logf operation.
func (t *MockT) Logf(format string, args ...interface{}) {
	t.record("Logf", fmt.Sprintf(format, args...))
}

// Skip records a Skip operation.
func (t *MockT) Skip(args ...interface{}) {
	t.record("Skip", fmtArgs(args...))
	t.skipped = true
}

// Skipf records a Skipf operation.
func (t *MockT) Skipf(format string, args ...interface{}) {
	t.record("Skipf", fmt.Sprintf(format, args...))
	t.skipped = true
}

// SkipNow records a SkipNow operation.
func (t *MockT) SkipNow() {
	t.record("SkipNow", "")
	t.skipped = true
}

// Skipped reports whether Skip, Skipf, or SkipNow was previously
// invoked.
func (t *MockT) Skipped() bool {
	return t.skipped
}

// Helper records a Helper operation.
func (t *MockT) Helper() {
	t.helper = true
}
<commit_msg>test/assert: MockT matcher for any error<commit_after>package assert

import (
	"fmt"
	"strings"
	"testing"
)

// MatchType is used in MockT's CheckPredicates to control whether
// recorded testing.TB operations match predicates.
type MatchType int

const (
	// Success indicates a predicate matches in CheckPredicates
	Success MatchType = iota

	// Failure indicates a predicate does not match in CheckPredicates
	Failure
)

// MockTestOperation represents an invocation of a method on
// testing.TB and its arguments.
type MockTestOperation struct {
	// Name represents the name of an operation on TB. One of: Error,
	// Error, Fatal, Fatalf, Fail, FailNow, Log, Logf, Skip, SkipNow,
	// Skipf.
	Name string

	// Args are the formatted arguments of the operation. For the
	// formatted variants (e.g., Errorf), this is the result of invoke
	// fmt.Sprintf with the format string and arguments. For the
	// non-formatted variables, this is the result of the arguments,
	// converted to strings as in fmt.Sprint, with spaces between each
	// one argument. Finally for the no-arg operations this is an
	// empty string.
	Args string
}

// String provides a string representation of MockTestOperation.
func (op *MockTestOperation) String() string {
	return fmt.Sprintf("%s(%q)", op.Name, op.Args)
}

// Predicate represents a predicate testing a
// MockTestOperation. Typically Predicates are composed from a
// NamePredicate and an ArgPredicate via Match.
type Predicate func(MockTestOperation) MatchType

// NamePredicate compares the given operation name (the name of a
// testing.TB method) to determine if it matches the predicate.
type NamePredicate func(string) MatchType

// OneOfOp returns a NamePredicate requires its op to equal one of the
// given operations.
func OneOfOp(expectedOps ...string) NamePredicate {
	return func(op string) MatchType {
		for _, eop := range expectedOps {
			if op == eop {
				return Success
			}
		}
		return Failure
	}
}

// ExactOp is equivalent to OneOfOp(op).
func ExactOp(op string) NamePredicate { return OneOfOp(op) }

// ErrorOp matches Error or Errorf operations.
func ErrorOp() NamePredicate { return OneOfOp("Error", "Errorf") }

// FatalOp matches Fatal or Fatalf operations.
func FatalOp() NamePredicate { return OneOfOp("Fatal", "Fatalf") }

// Fails matches any testing.TB operation that causes a test failure.
func Fails() NamePredicate {
	return OneOfOp("Error", "Errorf", "Fatal", "Fatalf", "Fail", "FailNow")
}

// ArgsPredicate tests the given operation arguments to determine if
// they match the predicate.
type ArgsPredicate func(string) MatchType

// PrefixedArgs requires the operation arguments to match the given
// prefix.
func PrefixedArgs(prefix string) ArgsPredicate {
	return func(args string) MatchType {
		if strings.HasPrefix(args, prefix) {
			return Success
		}
		return Failure
	}
}

// ArgsContain requires the operation arguments to contain the given
// expression.
func ArgsContain(expr string) ArgsPredicate {
	return func(args string) MatchType {
		if strings.Contains(args, expr) {
			return Success
		}
		return Failure
	}
}

// Match creates a Predicate from the given NamePredicate and
// ArgPredicate. The NamePredicate is tested first. If it fails, the
// entire Predicate fails. Otherwise, the ArgPredicate is tested it's
// result is returned.
func Match(namePred NamePredicate, argPred ArgsPredicate) Predicate {
	return func(op MockTestOperation) MatchType {
		if m := namePred(op.Name); m != Success {
			return m
		}

		return argPred(op.Args)
	}
}

// MatchWithAnyArgs creates a Predicate from the given NamePredicate without regard
// to the arguments passed.
func MatchWithAnyArgs(namePred NamePredicate) Predicate {
	return func(op MockTestOperation) MatchType {
		return namePred(op.Name)
	}
}

// MockT is an implementation of testing.TB that allows testing of
// code that makes test assertions. Calls to Fatal, Fatalf, FailNow
// and SkipNow are recorded, but do not terminate execution. If the
// code under test depends on termination of execution, unexpected
// results may occur.
type MockT struct {
	testing.TB

	Operations []MockTestOperation

	failed  bool
	skipped bool
	helper  bool
}

func (t *MockT) record(op string, args string) {
	t.Operations = append(t.Operations, MockTestOperation{op, args})
}

func (t *MockT) operationsToString() string {
	s := make([]string, 0, len(t.Operations))
	for _, e := range t.Operations {
		s = append(s, e.String())
	}
	return strings.Join(s, "\n")
}

// Reset resets the MockT's log of operations.
func (t *MockT) Reset() {
	t.Operations = nil
	t.failed = false
	t.skipped = false
}

// CheckHelper fails the test if the Helper method was not invoked on
// this MockT.
func (t *MockT) CheckHelper(realT testing.TB) {
	if !t.helper {
		Tracing(realT).Error("expected Helper invocation, but did not see it")
	}
}

// CheckPredicates evaluates the given predicates against the recorded
// testing.TB operations. Predicates must match operations in order.
func (t *MockT) CheckPredicates(realT testing.TB, predicates ...Predicate) {
	nops := len(t.Operations)
	npreds := len(predicates)

	n := nops
	if npreds < nops {
		n = npreds
	}

	errors := []string{}
	for i := 0; i < n; i++ {
		op := t.Operations[i]
		if predicates[i](op) != Success {
			errors = append(
				errors,
				fmt.Sprintf("operation %d, %s: did not match predicate", i+1, op.String()),
			)
		}
	}

	if nops != npreds {
		Tracing(realT).Errorf(
			"Expected %d operations, but got %d:\n%s",
			npreds,
			nops,
			t.operationsToString(),
		)
	}

	if len(errors) > 0 {
		Tracing(realT).Errorf("Predicate failure(s):\n%s", strings.Join(errors, "\n"))
	}
}

// CheckSuccess fails the test if any operations were logged.
func (t *MockT) CheckSuccess(realT testing.TB) {
	if len(t.Operations) > 0 {
		Tracing(realT).Errorf(
			"expected no testing.TB operations, but saw:\n%s",
			t.operationsToString(),
		)
	}
}

func fmtArgs(args ...interface{}) string {
	// Sprintln (vs Sprint) always adds spaces, but we don't want the
	// trailing newline.
	s := fmt.Sprintln(args...)
	return s[:len(s)-1]
}

// Name returns "MockTest", always.
func (t *MockT) Name() string { return "MockTest" }

// Error records an Error operation.
func (t *MockT) Error(args ...interface{}) {
	t.record("Error", fmtArgs(args...))
}

// Errorf records an Errorf operation.
func (t *MockT) Errorf(format string, args ...interface{}) {
	t.record("Errorf", fmt.Sprintf(format, args...))
}

// Fail records a Fail operation.
func (t *MockT) Fail() {
	t.record("Fail", "")
	t.failed = true
}

// FailNow records a FailNow operation.
func (t *MockT) FailNow() {
	t.record("FailNow", "")
	t.failed = true
}

// Failed reports whether Fail or FailNow was previously invoked.
func (t *MockT) Failed() bool {
	return t.failed
}

// Fatal records a Fatal operation.
func (t *MockT) Fatal(args ...interface{}) {
	t.record("Fatal", fmtArgs(args...))
}

// Fatalf records a Fatalf operation.
func (t *MockT) Fatalf(format string, args ...interface{}) {
	t.record("Fatalf", fmt.Sprintf(format, args...))
}

// Log records a Log operation.
func (t *MockT) Log(args ...interface{}) {
	t.record("Log", fmtArgs(args...))
}

// Logf records a Logf operation.
func (t *MockT) Logf(format string, args ...interface{}) {
	t.record("Logf", fmt.Sprintf(format, args...))
}

// Skip records a Skip operation.
func (t *MockT) Skip(args ...interface{}) {
	t.record("Skip", fmtArgs(args...))
	t.skipped = true
}

// Skipf records a Skipf operation.
func (t *MockT) Skipf(format string, args ...interface{}) {
	t.record("Skipf", fmt.Sprintf(format, args...))
	t.skipped = true
}

// SkipNow records a SkipNow operation.
func (t *MockT) SkipNow() {
	t.record("SkipNow", "")
	t.skipped = true
}

// Skipped reports whether Skip, Skipf, or SkipNow was previously
// invoked.
func (t *MockT) Skipped() bool {
	return t.skipped
}

// Helper records a Helper operation.
func (t *MockT) Helper() {
	t.helper = true
}
<|endoftext|>


<commit_before>package marionette_client

import (
	"time"
	"errors"
	"fmt"
)

type Waiter struct {
	f Finder
	d time.Duration // Milliseconds
}

type Finder interface {
	FindElement(by By, value string) (*WebElement, error)
	FindElements(by By, value string) ([]*WebElement, error)
}

func Wait(f Finder) (*Waiter) {
	return &Waiter{f, time.Duration(1)}
}

func (w *Waiter) For(d time.Duration) *Waiter {
	if d < 0 || d > time.Minute * 10 {
		w.d = time.Duration(time.Second)
		return w
	}

	w.d = d
	return w
}

func (w *Waiter) Until(f func(c Finder) (bool, *WebElement, error)) (bool, *WebElement, error) {
	firstRun := true
	delta := time.Now()
	for time.Since(delta) < w.d || firstRun {
		firstRun = false

		ok, value, err := f(w.f)
		if err != nil {
			_, de := err.(DriverError)
			if de {
				return false, nil, err
			}
		}

		if ok {
			return true, value, err
		}

		time.Sleep(time.Second)
	}

	return false, nil, errors.New("Condition never occurred.")
}
<commit_msg>Removed unused import in Wait.go<commit_after>package marionette_client

import (
	"time"
	"errors"
)

type Waiter struct {
	f Finder
	d time.Duration // Milliseconds
}

type Finder interface {
	FindElement(by By, value string) (*WebElement, error)
	FindElements(by By, value string) ([]*WebElement, error)
}

func Wait(f Finder) (*Waiter) {
	return &Waiter{f, time.Duration(1)}
}

func (w *Waiter) For(d time.Duration) *Waiter {
	if d < 0 || d > time.Minute * 10 {
		w.d = time.Duration(time.Second)
		return w
	}

	w.d = d
	return w
}

func (w *Waiter) Until(f func(c Finder) (bool, *WebElement, error)) (bool, *WebElement, error) {
	firstRun := true
	delta := time.Now()
	for time.Since(delta) < w.d || firstRun {
		firstRun = false

		ok, value, err := f(w.f)
		if err != nil {
			_, de := err.(DriverError)
			if de {
				return false, nil, err
			}
		}

		if ok {
			return true, value, err
		}

		time.Sleep(time.Second)
	}

	return false, nil, errors.New("Condition never occurred.")
}
<|endoftext|>


<commit_before>package dragon

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"sync"
)

type walkFunc func(info fileInfo) error

func concurrentWalk(root string, walkFn walkFunc) error {
	info, err := os.Lstat(root)
	if err != nil {
		return err
	}
	sem := make(chan struct{}, 16)
	return walk(newFileInfo(root, info), walkFn, sem)
}

func walk(info fileInfo, walkFn walkFunc, sem chan struct{}) error {
	walkError := walkFn(info)
	if walkError != nil {
		if info.IsDir() && walkError == filepath.SkipDir {
			return nil
		}
		return walkError
	}

	if !info.IsDir() {
		return nil
	}

	path := info.path
	files, err := ioutil.ReadDir(path)
	if err != nil {
		return err
	}

	wg := &sync.WaitGroup{}
	for _, file := range files {
		f := newFileInfo(filepath.Join(path, file.Name()), file)
		select {
		case sem <- struct{}{}:
			wg.Add(1)
			go func(file fileInfo, wg *sync.WaitGroup) {
				defer wg.Done()
				defer func() { <-sem }()
				walk(file, walkFn, sem)
			}(f, wg)
		default:
			walk(f, walkFn, sem)
		}
	}
	wg.Wait()
	return nil
}
<commit_msg>introduce errgroup<commit_after>package dragon

import (
	"io/ioutil"
	"os"
	"path/filepath"

	"golang.org/x/sync/errgroup"
)

type walkFunc func(info fileInfo) error

func concurrentWalk(root string, walkFn walkFunc) error {
	info, err := os.Lstat(root)
	if err != nil {
		return err
	}
	sem := make(chan struct{}, 16)
	return walk(newFileInfo(root, info), walkFn, sem)
}

func walk(info fileInfo, walkFn walkFunc, sem chan struct{}) error {
	walkError := walkFn(info)
	if walkError != nil {
		if info.IsDir() && walkError == filepath.SkipDir {
			return nil
		}
		return walkError
	}

	if !info.IsDir() {
		return nil
	}

	path := info.path
	files, err := ioutil.ReadDir(path)
	if err != nil {
		return err
	}

	eg := &errgroup.Group{}
	for _, file := range files {
		f := newFileInfo(filepath.Join(path, file.Name()), file)
		select {
		case sem <- struct{}{}:
			eg.Go(func() error {
				defer func() { <-sem }()
				return walk(f, walkFn, sem)
			})
		default:
			if err := walk(f, walkFn, sem); err != nil {
				return err
			}
		}
	}
	return eg.Wait()
}
<|endoftext|>