<commit_before>package metafora

import (
	"math/rand"
	"runtime"
	"sort"
	"sync"
	"time"
)

var (
	// balance calls are randomized and this is the upper bound of the random
	// amount
	balanceJitterMax = 10 * int64(time.Second)

	//FIXME should probably be improved, see usage in Run()
	consumerRetryDelay = 10 * time.Second
)

// Consumer is the core Metafora task runner.
type Consumer struct {
	// Func to create new handlers
	handler HandlerFunc

	// Map of task:Handler
	running map[string]*task

	// Mutex to protect access to running
	runL sync.Mutex

	// WaitGroup for running handlers and consumer goroutines
	hwg sync.WaitGroup

	// WaitGroup so Shutdown() can block on Run() exiting fully
	runwg  sync.WaitGroup
	runwgL sync.Mutex

	bal      Balancer
	balEvery time.Duration
	coord    Coordinator
	stop     chan struct{} // closed by Shutdown to cause Run to exit

	// ticked on each loop of the main loop to enforce sequential interaction
	// with coordinator and balancer
	tick chan int

	watch chan string // channel for watcher to send tasks to main loop

	// Set by command handler, read anywhere via Consumer.frozen()
	freezeL sync.Mutex
	freeze  bool
}

// NewConsumer returns a new consumer and calls Init on the Balancer and Coordinator.
func NewConsumer(coord Coordinator, h HandlerFunc, b Balancer) (*Consumer, error) {
	c := &Consumer{
		running:  make(map[string]*task),
		handler:  h,
		bal:      b,
		balEvery: 15 * time.Minute, //TODO make balance wait configurable
		coord:    coord,
		stop:     make(chan struct{}),
		tick:     make(chan int),
		watch:    make(chan string),
	}

	// initialize balancer with the consumer and a prefixed logger
	b.Init(c)

	if err := coord.Init(&coordinatorContext{c}); err != nil {
		return nil, err
	}
	return c, nil
}

// Run is the core run loop of Metafora. It is responsible for calling into the
// Coordinator to claim work and Balancer to rebalance work.
//
// Run blocks until Shutdown is called or an internal error occurs.
func (c *Consumer) Run() {
	Debug("Starting consumer")

	// Increment run wait group so Shutdown() can block on Run() exiting fully.
	c.runwgL.Lock()
	c.runwg.Add(1)
	c.runwgL.Unlock()
	defer c.runwg.Done()

	// chans for core goroutines to communicate with main loop
	balance := make(chan bool)
	cmdChan := make(chan Command)

	// Balance is called by the main loop when the balance channel is ticked
	go func() {
		randInt := rand.New(rand.NewSource(time.Now().UnixNano())).Int63n
		for {
			select {
			case <-c.stop:
				// Shutdown has been called.
				return
			case <-time.After(c.balEvery + time.Duration(randInt(balanceJitterMax))):
				Info("Balancing")
				select {
				case balance <- true:
					// Ticked balance
				case <-c.stop:
					// Shutdown has been called.
					return
				}
			}
			// Wait for main loop to signal balancing is done
			select {
			case <-c.stop:
				return
			case <-c.tick:
			}
		}
	}()

	// Watch for new tasks in a goroutine
	go c.watcher()

	// Watch for new commands in a goroutine
	go func() {
		defer close(cmdChan)
		for {
			cmd, err := c.coord.Command()
			if err != nil {
				Errorf("Coordinator returned an error during command, waiting and retrying. %v", err)
				select {
				case <-c.stop:
					return
				case <-time.After(consumerRetryDelay):
				}
				continue
			}
			if cmd == nil {
				Debug("Command coordinator exited")
				return
			}
			// Send command to watcher (or shutdown)
			select {
			case <-c.stop:
				return
			case cmdChan <- cmd:
			}
			// Wait for main loop to signal command has been handled
			select {
			case <-c.stop:
				return
			case <-c.tick:
			}
		}
	}()

	// Make sure Run() cleans up on exit (stops coordinator, releases tasks, etc)
	defer c.shutdown()

	// Main Loop ensures events are processed synchronously
	for {
		if c.Frozen() {
			// Only recv commands while frozen
			select {
			case <-c.stop:
				// Shutdown has been called.
				return
			case cmd, ok := <-cmdChan:
				if !ok {
					Debug("Command channel closed. Exiting main loop.")
					return
				}
				Debugf("Received command: %s", cmd)
				c.handleCommand(cmd)
			}
			// Must send tick whenever main loop restarts
			select {
			case <-c.stop:
				return
			case c.tick <- 1:
			}
			continue
		}

		select {
		case <-c.stop:
			// Shutdown has been called.
			return
		case <-balance:
			c.balance()
		case task, ok := <-c.watch:
			if !ok {
				Debug("Watch channel closed. Exiting main loop.")
				return
			}
			if !c.bal.CanClaim(task) {
				Infof("Balancer rejected task %s", task)
				break
			}
			if !c.coord.Claim(task) {
				Infof("Coordinator unable to claim task %s", task)
				break
			}
			c.claimed(task)
		case cmd, ok := <-cmdChan:
			if !ok {
				Debug("Command channel closed. Exiting main loop.")
				return
			}
			c.handleCommand(cmd)
		}
		// Signal that main loop is restarting after handling an event
		select {
		case <-c.stop:
			return
		case c.tick <- 1:
		}
	}
}

func (c *Consumer) watcher() {
	defer close(c.watch)
	Debug("Consumer watching")

	for {
		task, err := c.coord.Watch()
		if err != nil {
			//FIXME add more sophisticated error handling
			Errorf("Coordinator returned an error during watch, waiting and retrying: %v", err)
			select {
			case <-c.stop:
				return
			case <-time.After(consumerRetryDelay):
			}
			continue
		}
		if task == "" {
			Info("Coordinator has closed, no longer watching for tasks.")
			return
		}
		// Send task to watcher (or shutdown)
		select {
		case <-c.stop:
			return
		case c.watch <- task:
		}
		// Wait for main loop to signal task has been handled
		select {
		case <-c.stop:
			return
		case <-c.tick:
		}
	}
}

func (c *Consumer) balance() {
	tasks := c.bal.Balance()
	if len(tasks) > 0 {
		Infof("Balancer releasing: %v", tasks)
	}
	for _, task := range tasks {
		c.stopTask(task)
	}
}

// shutdown is the actual shutdown logic called when Run() exits.
func (c *Consumer) shutdown() {
	// Build list of of currently running tasks
	tasks := c.Tasks()
	Infof("Sending stop signal to %d handler(s)", len(tasks))

	for _, id := range tasks {
		c.stopTask(id.ID())
	}

	Info("Waiting for handlers to exit")
	c.hwg.Wait()

	Debug("Closing Coordinator")
	c.coord.Close()
}

// Shutdown stops the main Run loop, calls Stop on all handlers, and calls
// Close on the Coordinator. Running tasks will be released for other nodes to
// claim.
func (c *Consumer) Shutdown() {
	// acquire the runL lock to make sure we don't race with claimed()'s <-c.stop
	// check
	c.runL.Lock()
	select {
	case <-c.stop:
		// already stopped
	default:
		Debug("Stopping Run loop")
		close(c.stop)
	}
	c.runL.Unlock()

	// Wait for task handlers to exit.
	c.hwg.Wait()

	// Make sure Run() exits, otherwise Shutdown() might exit before
	// coord.Close() is called.
	c.runwgL.Lock()
	c.runwg.Wait()
	c.runwgL.Unlock()
}

// Tasks returns a lexicographically sorted list of running Task IDs.
func (c *Consumer) Tasks() []Task {
	c.runL.Lock()
	defer c.runL.Unlock()

	// Create a sorted list of task IDs
	ids := make([]string, len(c.running))
	i := 0
	for id, _ := range c.running {
		ids[i] = id
		i++
	}
	sort.Strings(ids)

	// Add tasks in lexicographic order
	t := make([]Task, len(ids))
	for i, id := range ids {
		t[i] = c.running[id]
	}
	return t
}

// claimed starts a handler for a claimed task. It is the only method to
// manipulate c.running and closes the task channel when a handler's Run
// method exits.
func (c *Consumer) claimed(taskID string) {
	h := c.handler()

	Debugf("Attempting to start task " + taskID)
	// Associate handler with taskID
	// **This is the only place tasks should be added to c.running**
	c.runL.Lock()
	defer c.runL.Unlock()
	select {
	case <-c.stop:
		// We're closing, don't bother starting this task
		return
	default:
	}
	if _, ok := c.running[taskID]; ok {
		// If a coordinator returns an already claimed task from Watch(), then it's
		// a coordinator (or broker) bug.
		Warnf("Attempted to claim already running task %s", taskID)
		return
	}
	rt := newTask(taskID, h)
	c.running[taskID] = rt

	// This must be done in the runL lock after the stop chan check so Shutdown
	// doesn't close(stop) and start Wait()ing concurrently.
	// See "Note" http://golang.org/pkg/sync/#WaitGroup.Add
	c.hwg.Add(1)

	// Start handler in its own goroutine
	go func() {
		defer c.hwg.Done() // Must be run after task exit and Done/Release called

		// Run the task
		Infof("Task started: %s", taskID)
		done := c.runTask(h.Run, taskID)
		if done {
			Infof("Task exited: %s (marking done)", taskID)
			c.coord.Done(taskID)
		} else {
			Infof("Task exited: %s (releasing)", taskID)
			c.coord.Release(taskID)
		}
	}()
}

// runTask executes a handler's Run method and recovers from panic()s.
func (c *Consumer) runTask(run func(string) bool, task string) bool {
	done := false
	func() {
		defer func() {
			if err := recover(); err != nil {
				stack := make([]byte, 50*1024)
				sz := runtime.Stack(stack, false)
				Errorf("Handler %s panic()'d: %v\n%s", task, err, stack[:sz])
				// panics are considered fatal errors. Make sure the task isn't
				// rescheduled.
				done = true
			}

			// **This is the only place tasks should be removed from c.running**
			c.runL.Lock()
			delete(c.running, task)
			c.runL.Unlock()
		}()
		done = run(task)
	}()
	return done
}

// stopTask asynchronously calls the task handlers' Stop method. While stopTask
// calls don't block, calls to task handler's Stop method are serialized with a
// lock.
func (c *Consumer) stopTask(taskID string) {
	c.runL.Lock()
	task, ok := c.running[taskID]
	c.runL.Unlock()

	if !ok {
		// This can happen if a task completes during Balance() and is not an error.
		Warnf("Tried to release a non-running task: %s", taskID)
		return
	}

	// all handler methods must be wrapped in a recover to prevent a misbehaving
	// handler from crashing the entire consumer
	go func() {
		defer func() {
			if err := recover(); err != nil {
				stack := make([]byte, 50*1024)
				sz := runtime.Stack(stack, false)
				Errorf("Handler %s panic()'d on Stop: %v\n%s", taskID, err, stack[:sz])
			}
		}()

		// Serialize calls to Stop as a convenience to handler implementors.
		task.stop()
	}()
}

// Frozen returns true if Metafora is no longer watching for new tasks or
// rebalancing.
//
// Metafora will remain frozen until receiving an Unfreeze command or it is
// restarted (frozen state is not persisted).
func (c *Consumer) Frozen() bool {
	c.freezeL.Lock()
	r := c.freeze
	c.freezeL.Unlock()
	return r
}

func (c *Consumer) handleCommand(cmd Command) {
	switch cmd.Name() {
	case cmdFreeze:
		if c.Frozen() {
			Info("Ignoring freeze command: already frozen")
			return
		}
		Info("Freezing")
		c.freezeL.Lock()
		c.freeze = true
		c.freezeL.Unlock()
	case cmdUnfreeze:
		if !c.Frozen() {
			Info("Ignoring unfreeze command: not frozen")
			return
		}
		Info("Unfreezing")
		c.freezeL.Lock()
		c.freeze = false
		c.freezeL.Unlock()
	case cmdBalance:
		Info("Balancing due to command")
		c.balance()
		Debug("Finished balancing due to command")
	case cmdStopTask:
		taskI, ok := cmd.Parameters()["task"]
		task, ok2 := taskI.(string)
		if !ok || !ok2 {
			Error("Stop task command didn't contain a valid task")
			return
		}
		Info("Stopping task %s due to command", task)
		c.stopTask(task)
	default:
		Warnf("Discarding unknown command: %s", cmd.Name())
	}
}
<commit_msg>Include task shutdown latency in logs<commit_after>package metafora

import (
	"math/rand"
	"runtime"
	"sort"
	"sync"
	"time"
)

var (
	// balance calls are randomized and this is the upper bound of the random
	// amount
	balanceJitterMax = 10 * int64(time.Second)

	//FIXME should probably be improved, see usage in Run()
	consumerRetryDelay = 10 * time.Second
)

// Consumer is the core Metafora task runner.
type Consumer struct {
	// Func to create new handlers
	handler HandlerFunc

	// Map of task:Handler
	running map[string]*task

	// Mutex to protect access to running
	runL sync.Mutex

	// WaitGroup for running handlers and consumer goroutines
	hwg sync.WaitGroup

	// WaitGroup so Shutdown() can block on Run() exiting fully
	runwg  sync.WaitGroup
	runwgL sync.Mutex

	bal      Balancer
	balEvery time.Duration
	coord    Coordinator
	stop     chan struct{} // closed by Shutdown to cause Run to exit

	// ticked on each loop of the main loop to enforce sequential interaction
	// with coordinator and balancer
	tick chan int

	watch chan string // channel for watcher to send tasks to main loop

	// Set by command handler, read anywhere via Consumer.frozen()
	freezeL sync.Mutex
	freeze  bool
}

// NewConsumer returns a new consumer and calls Init on the Balancer and Coordinator.
func NewConsumer(coord Coordinator, h HandlerFunc, b Balancer) (*Consumer, error) {
	c := &Consumer{
		running:  make(map[string]*task),
		handler:  h,
		bal:      b,
		balEvery: 15 * time.Minute, //TODO make balance wait configurable
		coord:    coord,
		stop:     make(chan struct{}),
		tick:     make(chan int),
		watch:    make(chan string),
	}

	// initialize balancer with the consumer and a prefixed logger
	b.Init(c)

	if err := coord.Init(&coordinatorContext{c}); err != nil {
		return nil, err
	}
	return c, nil
}

// Run is the core run loop of Metafora. It is responsible for calling into the
// Coordinator to claim work and Balancer to rebalance work.
//
// Run blocks until Shutdown is called or an internal error occurs.
func (c *Consumer) Run() {
	Debug("Starting consumer")

	// Increment run wait group so Shutdown() can block on Run() exiting fully.
	c.runwgL.Lock()
	c.runwg.Add(1)
	c.runwgL.Unlock()
	defer c.runwg.Done()

	// chans for core goroutines to communicate with main loop
	balance := make(chan bool)
	cmdChan := make(chan Command)

	// Balance is called by the main loop when the balance channel is ticked
	go func() {
		randInt := rand.New(rand.NewSource(time.Now().UnixNano())).Int63n
		for {
			select {
			case <-c.stop:
				// Shutdown has been called.
				return
			case <-time.After(c.balEvery + time.Duration(randInt(balanceJitterMax))):
				Info("Balancing")
				select {
				case balance <- true:
					// Ticked balance
				case <-c.stop:
					// Shutdown has been called.
					return
				}
			}
			// Wait for main loop to signal balancing is done
			select {
			case <-c.stop:
				return
			case <-c.tick:
			}
		}
	}()

	// Watch for new tasks in a goroutine
	go c.watcher()

	// Watch for new commands in a goroutine
	go func() {
		defer close(cmdChan)
		for {
			cmd, err := c.coord.Command()
			if err != nil {
				Errorf("Coordinator returned an error during command, waiting and retrying. %v", err)
				select {
				case <-c.stop:
					return
				case <-time.After(consumerRetryDelay):
				}
				continue
			}
			if cmd == nil {
				Debug("Command coordinator exited")
				return
			}
			// Send command to watcher (or shutdown)
			select {
			case <-c.stop:
				return
			case cmdChan <- cmd:
			}
			// Wait for main loop to signal command has been handled
			select {
			case <-c.stop:
				return
			case <-c.tick:
			}
		}
	}()

	// Make sure Run() cleans up on exit (stops coordinator, releases tasks, etc)
	defer c.shutdown()

	// Main Loop ensures events are processed synchronously
	for {
		if c.Frozen() {
			// Only recv commands while frozen
			select {
			case <-c.stop:
				// Shutdown has been called.
				return
			case cmd, ok := <-cmdChan:
				if !ok {
					Debug("Command channel closed. Exiting main loop.")
					return
				}
				Debugf("Received command: %s", cmd)
				c.handleCommand(cmd)
			}
			// Must send tick whenever main loop restarts
			select {
			case <-c.stop:
				return
			case c.tick <- 1:
			}
			continue
		}

		select {
		case <-c.stop:
			// Shutdown has been called.
			return
		case <-balance:
			c.balance()
		case task, ok := <-c.watch:
			if !ok {
				Debug("Watch channel closed. Exiting main loop.")
				return
			}
			if !c.bal.CanClaim(task) {
				Infof("Balancer rejected task %s", task)
				break
			}
			if !c.coord.Claim(task) {
				Infof("Coordinator unable to claim task %s", task)
				break
			}
			c.claimed(task)
		case cmd, ok := <-cmdChan:
			if !ok {
				Debug("Command channel closed. Exiting main loop.")
				return
			}
			c.handleCommand(cmd)
		}
		// Signal that main loop is restarting after handling an event
		select {
		case <-c.stop:
			return
		case c.tick <- 1:
		}
	}
}

func (c *Consumer) watcher() {
	defer close(c.watch)
	Debug("Consumer watching")

	for {
		task, err := c.coord.Watch()
		if err != nil {
			//FIXME add more sophisticated error handling
			Errorf("Coordinator returned an error during watch, waiting and retrying: %v", err)
			select {
			case <-c.stop:
				return
			case <-time.After(consumerRetryDelay):
			}
			continue
		}
		if task == "" {
			Info("Coordinator has closed, no longer watching for tasks.")
			return
		}
		// Send task to watcher (or shutdown)
		select {
		case <-c.stop:
			return
		case c.watch <- task:
		}
		// Wait for main loop to signal task has been handled
		select {
		case <-c.stop:
			return
		case <-c.tick:
		}
	}
}

func (c *Consumer) balance() {
	tasks := c.bal.Balance()
	if len(tasks) > 0 {
		Infof("Balancer releasing: %v", tasks)
	}
	for _, task := range tasks {
		c.stopTask(task)
	}
}

// shutdown is the actual shutdown logic called when Run() exits.
func (c *Consumer) shutdown() {
	// Build list of of currently running tasks
	tasks := c.Tasks()
	Infof("Sending stop signal to %d handler(s)", len(tasks))

	for _, id := range tasks {
		c.stopTask(id.ID())
	}

	Info("Waiting for handlers to exit")
	c.hwg.Wait()

	Debug("Closing Coordinator")
	c.coord.Close()
}

// Shutdown stops the main Run loop, calls Stop on all handlers, and calls
// Close on the Coordinator. Running tasks will be released for other nodes to
// claim.
func (c *Consumer) Shutdown() {
	// acquire the runL lock to make sure we don't race with claimed()'s <-c.stop
	// check
	c.runL.Lock()
	select {
	case <-c.stop:
		// already stopped
	default:
		Debug("Stopping Run loop")
		close(c.stop)
	}
	c.runL.Unlock()

	// Wait for task handlers to exit.
	c.hwg.Wait()

	// Make sure Run() exits, otherwise Shutdown() might exit before
	// coord.Close() is called.
	c.runwgL.Lock()
	c.runwg.Wait()
	c.runwgL.Unlock()
}

// Tasks returns a lexicographically sorted list of running Task IDs.
func (c *Consumer) Tasks() []Task {
	c.runL.Lock()
	defer c.runL.Unlock()

	// Create a sorted list of task IDs
	ids := make([]string, len(c.running))
	i := 0
	for id, _ := range c.running {
		ids[i] = id
		i++
	}
	sort.Strings(ids)

	// Add tasks in lexicographic order
	t := make([]Task, len(ids))
	for i, id := range ids {
		t[i] = c.running[id]
	}
	return t
}

// claimed starts a handler for a claimed task. It is the only method to
// manipulate c.running and closes the task channel when a handler's Run
// method exits.
func (c *Consumer) claimed(taskID string) {
	h := c.handler()

	Debugf("Attempting to start task " + taskID)
	// Associate handler with taskID
	// **This is the only place tasks should be added to c.running**
	c.runL.Lock()
	defer c.runL.Unlock()
	select {
	case <-c.stop:
		// We're closing, don't bother starting this task
		return
	default:
	}
	if _, ok := c.running[taskID]; ok {
		// If a coordinator returns an already claimed task from Watch(), then it's
		// a coordinator (or broker) bug.
		Warnf("Attempted to claim already running task %s", taskID)
		return
	}
	rt := newTask(taskID, h)
	c.running[taskID] = rt

	// This must be done in the runL lock after the stop chan check so Shutdown
	// doesn't close(stop) and start Wait()ing concurrently.
	// See "Note" http://golang.org/pkg/sync/#WaitGroup.Add
	c.hwg.Add(1)

	// Start handler in its own goroutine
	go func() {
		defer c.hwg.Done() // Must be run after task exit and Done/Release called

		// Run the task
		Infof("Task %q started", taskID)
		done := c.runTask(h.Run, taskID)
		var status string
		if done {
			status = "done"
			c.coord.Done(taskID)
		} else {
			status = "released"
			c.coord.Release(taskID)
		}

		stopped := rt.Stopped()
		if stopped.IsZero() {
			// Task exited on its own
			Infof("Task %q exited (%s)", taskID, status)
		} else {
			// Task exited due to Stop() being called
			Infof("Task %q exited (%s) after %s", taskID, status, time.Now().Sub(stopped))
		}
	}()
}

// runTask executes a handler's Run method and recovers from panic()s.
func (c *Consumer) runTask(run func(string) bool, task string) bool {
	done := false
	func() {
		defer func() {
			if err := recover(); err != nil {
				stack := make([]byte, 50*1024)
				sz := runtime.Stack(stack, false)
				Errorf("Handler %s panic()'d: %v\n%s", task, err, stack[:sz])
				// panics are considered fatal errors. Make sure the task isn't
				// rescheduled.
				done = true
			}

			// **This is the only place tasks should be removed from c.running**
			c.runL.Lock()
			delete(c.running, task)
			c.runL.Unlock()
		}()
		done = run(task)
	}()
	return done
}

// stopTask asynchronously calls the task handlers' Stop method. While stopTask
// calls don't block, calls to task handler's Stop method are serialized with a
// lock.
func (c *Consumer) stopTask(taskID string) {
	c.runL.Lock()
	task, ok := c.running[taskID]
	c.runL.Unlock()

	if !ok {
		// This can happen if a task completes during Balance() and is not an error.
		Warnf("Tried to release a non-running task: %s", taskID)
		return
	}

	// all handler methods must be wrapped in a recover to prevent a misbehaving
	// handler from crashing the entire consumer
	go func() {
		defer func() {
			if err := recover(); err != nil {
				stack := make([]byte, 50*1024)
				sz := runtime.Stack(stack, false)
				Errorf("Handler %s panic()'d on Stop: %v\n%s", taskID, err, stack[:sz])
			}
		}()

		// Serialize calls to Stop as a convenience to handler implementors.
		task.stop()
	}()
}

// Frozen returns true if Metafora is no longer watching for new tasks or
// rebalancing.
//
// Metafora will remain frozen until receiving an Unfreeze command or it is
// restarted (frozen state is not persisted).
func (c *Consumer) Frozen() bool {
	c.freezeL.Lock()
	r := c.freeze
	c.freezeL.Unlock()
	return r
}

func (c *Consumer) handleCommand(cmd Command) {
	switch cmd.Name() {
	case cmdFreeze:
		if c.Frozen() {
			Info("Ignoring freeze command: already frozen")
			return
		}
		Info("Freezing")
		c.freezeL.Lock()
		c.freeze = true
		c.freezeL.Unlock()
	case cmdUnfreeze:
		if !c.Frozen() {
			Info("Ignoring unfreeze command: not frozen")
			return
		}
		Info("Unfreezing")
		c.freezeL.Lock()
		c.freeze = false
		c.freezeL.Unlock()
	case cmdBalance:
		Info("Balancing due to command")
		c.balance()
		Debug("Finished balancing due to command")
	case cmdStopTask:
		taskI, ok := cmd.Parameters()["task"]
		task, ok2 := taskI.(string)
		if !ok || !ok2 {
			Error("Stop task command didn't contain a valid task")
			return
		}
		Info("Stopping task %s due to command", task)
		c.stopTask(task)
	default:
		Warnf("Discarding unknown command: %s", cmd.Name())
	}
}
<|endoftext|>


<commit_before>package chord

import (
	"bytes"
	"encoding/binary"
	"log"
	"math/big"
	"math/rand"
	"time"
)

// Initializes a local vnode
func (vn *localVnode) init(idx int) error {
	// Generate an ID
	vn.genId(uint16(idx))

	// Initialize all state
	vn.successors = make([]*Vnode, vn.ring.config.NumSuccessors)
	vn.finger = make([]*Vnode, vn.ring.config.HashBits)
	return nil
}

// Schedules the Vnode to do regular maintenence
func (vn *localVnode) schedule() {
	// Setup our stabilize timer
	vn.timer = time.AfterFunc(randStabilize(vn.ring.config), vn.stabilize)
}

// Generates an ID for the node
func (vn *localVnode) genId(idx uint16) {
	// Use the hash funciton
	conf := vn.ring.config
	hash := conf.HashFunc()
	hash.Write([]byte(conf.Hostname))
	binary.Write(hash, binary.BigEndian, idx)

	// Use the hash as the ID
	vn.Id = hash.Sum(nil)
}

// Called to periodically stabilize the vnode
func (vn *localVnode) stabilize() {
	// Clear the timer
	vn.timer = nil

	// Check for shutdown
	if vn.ring.shutdown {
		return
	}

	// Setup the next stabilize timer
	defer vn.schedule()

	// Check for new successor
	if err := vn.checkNewSuccessor(); err != nil {
		log.Printf("[ERR] Error checking for new successor: %s", err)
	}

	// Notify the successor
	if err := vn.notifySuccessor(); err != nil {
		log.Printf("[ERR] Error notifying successor: %s", err)
	}

	// Finger table fix up
	if err := vn.fixFingerTable(); err != nil {
		log.Printf("[ERR] Error fixing finger table: %s", err)
	}

	// Check the predecessor
	if err := vn.checkPredecessor(); err != nil {
		log.Printf("[ERR] Error checking predecessor: %s", err)
	}

	// Set the last stabilized time
	vn.stabilized = time.Now()
}

// Generates a random stabilization time
func randStabilize(conf *Config) time.Duration {
	min := conf.StabilizeMin
	max := conf.StabilizeMax
	r := rand.Float64()
	return time.Duration((r * float64(max-min)) + float64(min))
}

// Checks for a new successor
func (vn *localVnode) checkNewSuccessor() error {
	// Ask our successor for it's predecessor
	succ := vn.successors[0]
	maybe_suc, err := vn.ring.transport.GetPredecessor(succ)
	if err != nil {
		return err
	}

	// Check if we should replace our successor
	if between(vn.Id, succ.Id, maybe_suc.Id) {
		vn.successors[0] = maybe_suc
	}
	return nil
}

// Notifies our successor of us, updates successor list
func (vn *localVnode) notifySuccessor() error {
	// Notify successor
	succ := vn.successors[0]
	succ_list, err := vn.ring.transport.Notify(succ, &vn.Vnode)
	if err != nil {
		return err
	}

	// Trim the successors list if too long
	max_succ := vn.ring.config.NumSuccessors
	if len(succ_list) > max_succ-1 {
		succ_list = succ_list[:max_succ-1]
	}

	// Update local successors list
	for idx, s := range succ_list {
		vn.successors[idx+1] = s
	}
	return nil
}

// RpcNotified is invoked when a Vnode gets notified
func (vn *localVnode) RpcNotified(maybe_pred *Vnode) ([]*Vnode, error) {
	// Check if we should update our predecessor
	if vn.predecessor == nil || between(vn.predecessor.Id, vn.Id, maybe_pred.Id) {
		vn.predecessor = maybe_pred
	}

	// Return our successors list
	return vn.successors, nil
}

// Fixes up the finger table
func (vn *localVnode) fixFingerTable() error {
	// Increment to the index to repair
	hb := vn.ring.config.HashBits
	if vn.last_finger+1 == hb {
		vn.last_finger = 0
	} else {
		vn.last_finger++
	}

	// Determine the offset
	offset := powerOffset(vn.Id, vn.last_finger, hb)

	// Find the successor
	node, err := vn.findSuccessor(offset)
	if err != nil {
		return err
	}

	// Update the finger table
	vn.finger[vn.last_finger] = node

	// Try to skip as many finger entries as possible
	for {
		offset := powerOffset(vn.Id, vn.last_finger+1, hb)

		// While the node is the successor, update the finger entries
		if bytes.Compare(node.Id, offset) == 1 {
			vn.finger[vn.last_finger+1] = node
			vn.last_finger++
		} else {
			break
		}
	}
	return nil
}

// Checks the health of our predecessor
func (vn *localVnode) checkPredecessor() error {
	// Check predecessor
	if vn.predecessor != nil {
		res, err := vn.ring.transport.Ping(vn.predecessor)
		if err != nil {
			return nil
		}

		// Predecessor is dead
		if !res {
			vn.predecessor = nil
		}
	}
	return nil
}

func (vn *localVnode) findSuccessor(key []byte) (*Vnode, error) {

}

// Checks if a key is STRICTLY between two ID's exclusively
func between(id1, id2, key []byte) bool {
	if bytes.Compare(id1, key) != -1 {
		return false
	}
	if bytes.Compare(id2, key) != 1 {
		return false
	}
	return true
}

// Computes the offset by (n + 2^exp) % (2^mod)
func powerOffset(id []byte, exp int, mod int) []byte {
	// Copy the existing slice
	off := make([]byte, len(id))
	copy(off, id)

	// Convert the ID to a bigint
	idInt := big.Int{}
	idInt.SetBytes(id)

	// Get the offset
	two := big.NewInt(2)
	offset := big.Int{}
	offset.Exp(two, big.NewInt(int64(exp)), nil)

	// Sum
	sum := big.Int{}
	sum.Add(&idInt, &offset)

	// Get the ceiling
	ceil := big.Int{}
	ceil.Exp(two, big.NewInt(int64(mod)), nil)

	// Apply the mod
	idInt.Mod(&sum, &ceil)

	// Add together
	return idInt.Bytes()
}
<commit_msg>Handle the modulous math<commit_after>package chord

import (
	"bytes"
	"encoding/binary"
	"log"
	"math/big"
	"math/rand"
	"time"
)

// Initializes a local vnode
func (vn *localVnode) init(idx int) error {
	// Generate an ID
	vn.genId(uint16(idx))

	// Initialize all state
	vn.successors = make([]*Vnode, vn.ring.config.NumSuccessors)
	vn.finger = make([]*Vnode, vn.ring.config.HashBits)
	return nil
}

// Schedules the Vnode to do regular maintenence
func (vn *localVnode) schedule() {
	// Setup our stabilize timer
	vn.timer = time.AfterFunc(randStabilize(vn.ring.config), vn.stabilize)
}

// Generates an ID for the node
func (vn *localVnode) genId(idx uint16) {
	// Use the hash funciton
	conf := vn.ring.config
	hash := conf.HashFunc()
	hash.Write([]byte(conf.Hostname))
	binary.Write(hash, binary.BigEndian, idx)

	// Use the hash as the ID
	vn.Id = hash.Sum(nil)
}

// Called to periodically stabilize the vnode
func (vn *localVnode) stabilize() {
	// Clear the timer
	vn.timer = nil

	// Check for shutdown
	if vn.ring.shutdown {
		return
	}

	// Setup the next stabilize timer
	defer vn.schedule()

	// Check for new successor
	if err := vn.checkNewSuccessor(); err != nil {
		log.Printf("[ERR] Error checking for new successor: %s", err)
	}

	// Notify the successor
	if err := vn.notifySuccessor(); err != nil {
		log.Printf("[ERR] Error notifying successor: %s", err)
	}

	// Finger table fix up
	if err := vn.fixFingerTable(); err != nil {
		log.Printf("[ERR] Error fixing finger table: %s", err)
	}

	// Check the predecessor
	if err := vn.checkPredecessor(); err != nil {
		log.Printf("[ERR] Error checking predecessor: %s", err)
	}

	// Set the last stabilized time
	vn.stabilized = time.Now()
}

// Generates a random stabilization time
func randStabilize(conf *Config) time.Duration {
	min := conf.StabilizeMin
	max := conf.StabilizeMax
	r := rand.Float64()
	return time.Duration((r * float64(max-min)) + float64(min))
}

// Checks for a new successor
func (vn *localVnode) checkNewSuccessor() error {
	// Ask our successor for it's predecessor
	succ := vn.successors[0]
	maybe_suc, err := vn.ring.transport.GetPredecessor(succ)
	if err != nil {
		return err
	}

	// Check if we should replace our successor
	if between(vn.Id, succ.Id, maybe_suc.Id) {
		vn.successors[0] = maybe_suc
	}
	return nil
}

// Notifies our successor of us, updates successor list
func (vn *localVnode) notifySuccessor() error {
	// Notify successor
	succ := vn.successors[0]
	succ_list, err := vn.ring.transport.Notify(succ, &vn.Vnode)
	if err != nil {
		return err
	}

	// Trim the successors list if too long
	max_succ := vn.ring.config.NumSuccessors
	if len(succ_list) > max_succ-1 {
		succ_list = succ_list[:max_succ-1]
	}

	// Update local successors list
	for idx, s := range succ_list {
		vn.successors[idx+1] = s
	}
	return nil
}

// RpcNotified is invoked when a Vnode gets notified
func (vn *localVnode) RpcNotified(maybe_pred *Vnode) ([]*Vnode, error) {
	// Check if we should update our predecessor
	if vn.predecessor == nil || between(vn.predecessor.Id, vn.Id, maybe_pred.Id) {
		vn.predecessor = maybe_pred
	}

	// Return our successors list
	return vn.successors, nil
}

// Fixes up the finger table
func (vn *localVnode) fixFingerTable() error {
	// Increment to the index to repair
	hb := vn.ring.config.HashBits
	if vn.last_finger+1 == hb {
		vn.last_finger = 0
	} else {
		vn.last_finger++
	}

	// Determine the offset
	offset := powerOffset(vn.Id, vn.last_finger, hb)

	// Find the successor
	node, err := vn.findSuccessor(offset)
	if err != nil {
		return err
	}

	// Update the finger table
	vn.finger[vn.last_finger] = node

	// Try to skip as many finger entries as possible
	for {
		next := (vn.last_finger + 1) % hb
		offset := powerOffset(vn.Id, next, hb)

		// While the node is the successor, update the finger entries
		if bytes.Compare(node.Id, offset) == 1 {
			vn.finger[next] = node
			vn.last_finger = next
		} else {
			break
		}
	}
	return nil
}

// Checks the health of our predecessor
func (vn *localVnode) checkPredecessor() error {
	// Check predecessor
	if vn.predecessor != nil {
		res, err := vn.ring.transport.Ping(vn.predecessor)
		if err != nil {
			return nil
		}

		// Predecessor is dead
		if !res {
			vn.predecessor = nil
		}
	}
	return nil
}

func (vn *localVnode) findSuccessor(key []byte) (*Vnode, error) {

}

// Checks if a key is STRICTLY between two ID's exclusively
func between(id1, id2, key []byte) bool {
	if bytes.Compare(id1, key) != -1 {
		return false
	}
	if bytes.Compare(id2, key) != 1 {
		return false
	}
	return true
}

// Computes the offset by (n + 2^exp) % (2^mod)
func powerOffset(id []byte, exp int, mod int) []byte {
	// Copy the existing slice
	off := make([]byte, len(id))
	copy(off, id)

	// Convert the ID to a bigint
	idInt := big.Int{}
	idInt.SetBytes(id)

	// Get the offset
	two := big.NewInt(2)
	offset := big.Int{}
	offset.Exp(two, big.NewInt(int64(exp)), nil)

	// Sum
	sum := big.Int{}
	sum.Add(&idInt, &offset)

	// Get the ceiling
	ceil := big.Int{}
	ceil.Exp(two, big.NewInt(int64(mod)), nil)

	// Apply the mod
	idInt.Mod(&sum, &ceil)

	// Add together
	return idInt.Bytes()
}
<|endoftext|>


<commit_before>// Copyright 2012 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This example demonstrates a priority queue built using the heap interface.
package heap_test

import (
	"container/heap"
	"fmt"
)

// An Item is something we manage in a priority queue.
type Item struct {
	value    string // The value of the item; arbitrary.
	priority int    // The priority of the item in the queue.
	// The index is needed by changePriority and is maintained by the heap.Interface methods.
	index int // The index of the item in the heap.
}

// A PriorityQueue implements heap.Interface and holds Items.
type PriorityQueue []*Item

func (pq PriorityQueue) Len() int { return len(pq) }

func (pq PriorityQueue) Less(i, j int) bool {
	// We want Pop to give us the highest, not lowest, priority so we use greater than here.
	return pq[i].priority > pq[j].priority
}

func (pq PriorityQueue) Swap(i, j int) {
	pq[i], pq[j] = pq[j], pq[i]
	pq[i].index = i
	pq[j].index = j
}

func (pq *PriorityQueue) Push(x interface{}) {
	n := len(*pq)
	item := x.(*Item)
	item.index = n
	*pq = append(*pq, item)
}

func (pq *PriorityQueue) Pop() interface{} {
	old := *pq
	n := len(old)
	item := old[n-1]
	item.index = -1 // for safety
	*pq = old[0 : n-1]
	return item
}

// update modifies the priority and value of an Item in the queue.
func (pq *PriorityQueue) update(item *Item, value string, priority int) {
	heap.Remove(pq, item.index)
	item.value = value
	item.priority = priority
	heap.Push(pq, item)
}

// This example inserts some items into a PriorityQueue, manipulates an item,
// and then removes the items in priority order.
func Example_priorityQueue() {
	// Some items and their priorities.
	items := map[string]int{
		"banana": 3, "apple": 2, "pear": 4,
	}

	// Create a priority queue and put the items in it.
	pq := &PriorityQueue{}
	heap.Init(pq)
	for value, priority := range items {
		item := &Item{
			value:    value,
			priority: priority,
		}
		heap.Push(pq, item)
	}

	// Insert a new item and then modify its priority.
	item := &Item{
		value:    "orange",
		priority: 1,
	}
	heap.Push(pq, item)
	pq.update(item, item.value, 5)

	// Take the items out; they arrive in decreasing priority order.
	for pq.Len() > 0 {
		item := heap.Pop(pq).(*Item)
		fmt.Printf("%.2d:%s ", item.priority, item.value)
	}
	// Output:
	// 05:orange 04:pear 03:banana 02:apple
}
<commit_msg>container/heap: fix comment typo in example test<commit_after>// Copyright 2012 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This example demonstrates a priority queue built using the heap interface.
package heap_test

import (
	"container/heap"
	"fmt"
)

// An Item is something we manage in a priority queue.
type Item struct {
	value    string // The value of the item; arbitrary.
	priority int    // The priority of the item in the queue.
	// The index is needed by update and is maintained by the heap.Interface methods.
	index int // The index of the item in the heap.
}

// A PriorityQueue implements heap.Interface and holds Items.
type PriorityQueue []*Item

func (pq PriorityQueue) Len() int { return len(pq) }

func (pq PriorityQueue) Less(i, j int) bool {
	// We want Pop to give us the highest, not lowest, priority so we use greater than here.
	return pq[i].priority > pq[j].priority
}

func (pq PriorityQueue) Swap(i, j int) {
	pq[i], pq[j] = pq[j], pq[i]
	pq[i].index = i
	pq[j].index = j
}

func (pq *PriorityQueue) Push(x interface{}) {
	n := len(*pq)
	item := x.(*Item)
	item.index = n
	*pq = append(*pq, item)
}

func (pq *PriorityQueue) Pop() interface{} {
	old := *pq
	n := len(old)
	item := old[n-1]
	item.index = -1 // for safety
	*pq = old[0 : n-1]
	return item
}

// update modifies the priority and value of an Item in the queue.
func (pq *PriorityQueue) update(item *Item, value string, priority int) {
	heap.Remove(pq, item.index)
	item.value = value
	item.priority = priority
	heap.Push(pq, item)
}

// This example inserts some items into a PriorityQueue, manipulates an item,
// and then removes the items in priority order.
func Example_priorityQueue() {
	// Some items and their priorities.
	items := map[string]int{
		"banana": 3, "apple": 2, "pear": 4,
	}

	// Create a priority queue and put the items in it.
	pq := &PriorityQueue{}
	heap.Init(pq)
	for value, priority := range items {
		item := &Item{
			value:    value,
			priority: priority,
		}
		heap.Push(pq, item)
	}

	// Insert a new item and then modify its priority.
	item := &Item{
		value:    "orange",
		priority: 1,
	}
	heap.Push(pq, item)
	pq.update(item, item.value, 5)

	// Take the items out; they arrive in decreasing priority order.
	for pq.Len() > 0 {
		item := heap.Pop(pq).(*Item)
		fmt.Printf("%.2d:%s ", item.priority, item.value)
	}
	// Output:
	// 05:orange 04:pear 03:banana 02:apple
}
<|endoftext|>


<commit_before>package tbeer

import (
	//sqlite3 "code.google.com/p/go-sqlite/go1/sqlite3"
	"crypto/rand"
	"database/sql"
	"encoding/binary"
	"fmt"
	"strings"
)

type randTable struct {
	stmt *sql.Stmt
	ids  []int64
}

func (t *randTable) put(args ...interface{}) {
	res, err := t.stmt.Exec(args...)
	if err != nil {
		fmt.Println(err)
	} else {
		id, _ := res.LastInsertId()
		t.ids = append(t.ids, id)
	}
}

func (t *randTable) randId() int64 {
	var r int
	binary.Read(rand.Reader, binary.LittleEndian, &r)
	return t.ids[r%len(t.ids)]
}

type randContext struct {
	db *sql.DB
	tx *sql.Tx
}

func (rc *randContext) table(q string) *randTable {
	fmt.Println(q)
	stmt, err := rc.db.Prepare(q)
	if err != nil {
		fmt.Println(err)
		return nil
	} else {
		return &randTable{rc.tx.Stmt(stmt), make([]int64, 0, 1000)}
	}
}

func newRandContext() *randContext {
	rc := new(randContext)
	rc.db = GlobalDB
	rc.tx, _ = rc.db.Begin()
	return rc
}

// return a slice that contains n copies of str
// BUG: isn't there a more elegant way of doing this in Go?
func repeat(str string, n int) []string {
	ret := make([]string, n)
	for i := 0; i < n; i++ {
		ret[i] = str
	}
	return ret
}

// generate an INSERT INTO query
func iq(table string, columns []string) string {
	return "INSERT INTO " + table + " (" +
		strings.Join(columns, ",") + ") VALUES (" +
		strings.Join(repeat("?", len(columns)), ",") + ")"
}

func PopulateRandom() {
	rc := newRandContext()

	fmt.Println("populating random data")

	rc.db.Begin()

	user := rc.table(iq("user", []string{"alias", "login", "email"}))
	place := rc.table(iq("place", []string{"name", "lat", "long"}))
	part := rc.table(iq("participant", []string{"ownerid", "alias", "description"}))
	meeting := rc.table(iq("meeting", []string{"ownerid", "name"}))
	mpart := rc.table(iq("meeting_participant", []string{"meetingid", "participantid"}))
	mplace := rc.table(iq("meeting_place", []string{"meetingid", "placeid"}))

	for i := 0; i < 10; i++ {
		user.put("yo", "login", "test@mail.com")
	}

	for i := 0; i < 10; i++ {
		place.put("place", 59.95+(float64(i)/100.0), 10.75+(float64(i)/100.0))
	}

	for i := 0; i < 10; i++ {
		part.put(user.randId(), "participant alias", "participant description")
	}

	for i := 0; i < 10; i++ {
		meeting.put(user.randId(), "meeting name")
	}

	for i := 0; i < 20; i++ {
		mpart.put(meeting.randId(), part.randId())
	}

	for i := 0; i < 20; i++ {
		mplace.put(meeting.randId(), place.randId())
	}

	fmt.Println("committing...")
	err := rc.tx.Commit()
	if err != nil {
		fmt.Println(err)
	} else {
		fmt.Println("... done")
	}
}
<commit_msg>Fix the random id generator<commit_after>package tbeer

import (
	//sqlite3 "code.google.com/p/go-sqlite/go1/sqlite3"
	"crypto/rand"
	"database/sql"
	"encoding/binary"
	"fmt"
	"strings"
)

type randTable struct {
	stmt *sql.Stmt
	ids  []int64
}

func (t *randTable) put(args ...interface{}) {
	res, err := t.stmt.Exec(args...)
	if err != nil {
		fmt.Println(err)
	} else {
		id, _ := res.LastInsertId()
		t.ids = append(t.ids, id)
	}
}

func (t *randTable) randId() int64 {
	var r int32
	binary.Read(rand.Reader, binary.LittleEndian, &r)
	if r < 0 {
		r = -r
	}
	return t.ids[int(r)%len(t.ids)]
}

type randContext struct {
	db *sql.DB
	tx *sql.Tx
}

func (rc *randContext) table(q string) *randTable {
	fmt.Println(q)
	stmt, err := rc.db.Prepare(q)
	if err != nil {
		fmt.Println(err)
		return nil
	} else {
		return &randTable{rc.tx.Stmt(stmt), make([]int64, 0, 1000)}
	}
}

func newRandContext() *randContext {
	rc := new(randContext)
	rc.db = GlobalDB
	rc.tx, _ = rc.db.Begin()
	return rc
}

// return a slice that contains n copies of str
// BUG: isn't there a more elegant way of doing this in Go?
func repeat(str string, n int) []string {
	ret := make([]string, n)
	for i := 0; i < n; i++ {
		ret[i] = str
	}
	return ret
}

// generate an INSERT INTO query
func iq(table string, columns []string) string {
	return "INSERT INTO " + table + " (" +
		strings.Join(columns, ",") + ") VALUES (" +
		strings.Join(repeat("?", len(columns)), ",") + ")"
}

func PopulateRandom() {
	rc := newRandContext()

	fmt.Println("populating random data")

	rc.db.Begin()

	user := rc.table(iq("user", []string{"alias", "login", "email"}))
	place := rc.table(iq("place", []string{"name", "lat", "long"}))
	part := rc.table(iq("participant", []string{"ownerid", "alias", "description"}))
	meeting := rc.table(iq("meeting", []string{"ownerid", "name"}))
	mpart := rc.table(iq("meeting_participant", []string{"meetingid", "participantid"}))
	mplace := rc.table(iq("meeting_place", []string{"meetingid", "placeid"}))

	for i := 0; i < 10; i++ {
		user.put("yo", "login", "test@mail.com")
	}

	for i := 0; i < 10; i++ {
		place.put("place", 59.95+(float64(i)/100.0), 10.75+(float64(i)/100.0))
	}

	for i := 0; i < 10; i++ {
		part.put(user.randId(), "participant alias", "participant description")
	}

	for i := 0; i < 10; i++ {
		meeting.put(user.randId(), "meeting name")
	}

	for i := 0; i < 20; i++ {
		mpart.put(meeting.randId(), part.randId())
	}

	for i := 0; i < 20; i++ {
		mplace.put(meeting.randId(), place.randId())
	}

	fmt.Println("committing...")
	err := rc.tx.Commit()
	if err != nil {
		fmt.Println(err)
	} else {
		fmt.Println("... done")
	}
}
<|endoftext|>


<commit_before>package transloadit

import (
	"encoding/json"
	"errors"
	"fmt"
	"github.com/mitchellh/go-homedir"
	"gopkg.in/fsnotify.v1"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strings"
	"time"
)

var sanitizeRe = regexp.MustCompile(`[^a-zA-Z0-9\-\_\.]`)

type WatchOptions struct {
	// Directory to watch files in (only if Watch is true)
	Input string
	// Directoy to put result files in
	Output string
	// Watch the input directory or just compile files in input directory
	Watch bool
	// Template id to convert files with
	TemplateId string
	// Optional notify url for each assembly.
	NotifyUrl string
	// Instead of using templates you can define steps
	Steps map[string]map[string]interface{}
	// If true the original files will be copied in the output directoy with `-original_0_` prefix.
	// If false input files will be deleted.
	Preserve bool
	// Read template from JSON file. See exmaples for information about the format.
	TemplateFile string
	// If true file which are already in the directory won't be compiled.
	DontProcessDir bool
}

type Watcher struct {
	client *Client
	// Original options passed to Client.Watch.
	// If Input or Output are relative to the home directory
	// they will be extended (~/input -> /home/user/input)
	Options *WatchOptions
	stopped bool
	// Listen for errors
	Error chan error
	// Listen for completed assemblies
	Done chan *AssemblyInfo
	// Listen for file changes (only if Watch == true)
	Change       chan string
	stop         chan bool
	recentWrites map[string]time.Time
	blacklist    map[string]bool
}

// Watch a directory for changes and convert all changes files and download the result.
// It will create a new assembly for each file.
// If the directory already contains some they are all converted.
// See WatchOptions for possible configuration.
func (client *Client) Watch(options *WatchOptions) *Watcher {
	options.Input = expandPath(options.Input)
	options.Output = expandPath(options.Output)

	if options.TemplateFile != "" {
		options.TemplateFile = expandPath(options.TemplateFile)
		options.Steps = readJson(options.TemplateFile)
	}

	if _, err := os.Stat(options.Input); os.IsNotExist(err) {
		panic(fmt.Errorf("Input directory does not exist: %s", options.Input))
	}

	if _, err := os.Stat(options.Output); os.IsNotExist(err) {
		panic(fmt.Errorf("Output directory does not exist: %s", options.Output))
	}

	if options.Input == options.Output {
		panic(fmt.Errorf("Input and output directory are both: %s", options.Output))
	}

	watcher := &Watcher{
		client:       client,
		Options:      options,
		Error:        make(chan error, 1),
		Done:         make(chan *AssemblyInfo),
		Change:       make(chan string),
		stop:         make(chan bool, 1),
		recentWrites: make(map[string]time.Time),
		blacklist:    make(map[string]bool),
	}

	watcher.start()

	return watcher
}

func (watcher *Watcher) start() {
	if !watcher.Options.DontProcessDir {
		watcher.processDir()
	}

	if watcher.Options.Watch {
		go watcher.startWatcher()
	}
}

// Stop the watcher.
func (watcher *Watcher) Stop() {
	if watcher.stopped {
		return
	}

	watcher.stopped = true

	close(watcher.Done)
	close(watcher.Error)
	close(watcher.Change)
	close(watcher.stop)
}

func (watcher *Watcher) processDir() {
	files, err := ioutil.ReadDir(watcher.Options.Input)
	if err != nil {
		watcher.error(err)
		return
	}

	input := watcher.Options.Input

	for _, file := range files {
		if !file.IsDir() {
			go watcher.processFile(path.Join(input, file.Name()))
		}
	}
}

func (watcher *Watcher) processFile(name string) {
	file, err := os.Open(name)
	if err != nil {
		watcher.error(err)
		return
	}

	// Add file to blacklist
	watcher.blacklist[name] = true

	assembly := watcher.client.CreateAssembly()

	if watcher.Options.TemplateId != "" {
		assembly.TemplateId = watcher.Options.TemplateId
	}

	if watcher.Options.NotifyUrl != "" {
		assembly.NotifyUrl = watcher.Options.NotifyUrl
	}

	for name, step := range watcher.Options.Steps {
		assembly.AddStep(name, step)
	}

	assembly.Blocking = true

	assembly.AddReader("file", path.Base(name), file)

	info, err := assembly.Upload()
	if err != nil {
		watcher.error(err)
		return
	}

	if info.Error != "" {
		watcher.error(errors.New(info.Error))
		return
	}

	for stepName, results := range info.Results {
		for index, result := range results {
			go func() {
				watcher.downloadResult(stepName, index, result)
				watcher.handleOriginalFile(name)
				delete(watcher.blacklist, name)
				watcher.Done <- info
			}()
		}
	}
}

func (watcher *Watcher) downloadResult(stepName string, index int, result *FileInfo) {
	fileName := sanitizeRe.ReplaceAllString(fmt.Sprintf("%s_%d_%s", stepName, index, result.Name), "-")

	resp, err := http.Get(result.Url)
	if err != nil {
		watcher.error(err)
		return
	}

	defer resp.Body.Close()

	out, err := os.Create(path.Join(watcher.Options.Output, fileName))
	if err != nil {
		watcher.error(err)
		return
	}

	defer out.Close()

	io.Copy(out, resp.Body)
}

func (watcher *Watcher) startWatcher() {
	fsWatcher, err := fsnotify.NewWatcher()
	if err != nil {
		watcher.error(err)
		return
	}

	defer fsWatcher.Close()

	if err = fsWatcher.Add(watcher.Options.Input); err != nil {
		watcher.error(err)
		return
	}

	for {
		select {
		case _, more := <-watcher.stop:
			if !more {
				return
			}
		case err := <-fsWatcher.Errors:
			watcher.error(err)
		case evt := <-fsWatcher.Events:
			// Ignore the event if the file is currently processed
			if _, ok := watcher.blacklist[evt.Name]; ok == true {
				continue
			}
			if evt.Op&fsnotify.Create == fsnotify.Create || evt.Op&fsnotify.Write == fsnotify.Write {
				watcher.recentWrites[evt.Name] = time.Now()
			}
		case <-time.Tick(1 * time.Second):
			now := time.Now()

			for name, lastEvent := range watcher.recentWrites {
				diff := now.Sub(lastEvent)
				if diff > (time.Millisecond * 500) {
					delete(watcher.recentWrites, name)
					watcher.Change <- name
					go watcher.processFile(name)
				}
			}
		}
	}
}

func (watcher *Watcher) handleOriginalFile(name string) {
	var err error
	if watcher.Options.Preserve {
		_, file := path.Split(name)
		err = os.Rename(name, watcher.Options.Output+"/-original_0_"+basename(file))
	} else {
		err = os.Remove(name)
	}

	if err != nil {
		watcher.error(err)
	}
}

func (watcher *Watcher) error(err error) {
	watcher.Error <- err
}

func basename(name string) string {
	i := strings.LastIndex(name, string(os.PathSeparator))
	return name[i+1:]
}

func expandPath(str string) string {
	expanded, err := homedir.Expand(str)
	if err != nil {
		panic(err)
	}

	expanded, err = filepath.Abs(expanded)
	if err != nil {
		panic(err)
	}

	return expanded
}

func readJson(file string) map[string]map[string]interface{} {
	content, err := ioutil.ReadFile(file)
	if err != nil {
		panic(fmt.Errorf("Error reading template file: %s", err))
	}

	steps := make(map[string]map[string]interface{})

	err = json.Unmarshal(content, &steps)
	if err != nil {
		panic(fmt.Errorf("Error parsing template file: %s", err))
	}

	return steps
}
<commit_msg>watch: avoid deadlock with empty input directory<commit_after>package transloadit

import (
	"encoding/json"
	"errors"
	"fmt"
	"github.com/mitchellh/go-homedir"
	"gopkg.in/fsnotify.v1"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"strings"
	"sync"
	"time"
)

var sanitizeRe = regexp.MustCompile(`[^a-zA-Z0-9\-\_\.]`)

type WatchOptions struct {
	// Directory to watch files in (only if Watch is true)
	Input string
	// Directoy to put result files in
	Output string
	// Watch the input directory or just compile files in input directory
	Watch bool
	// Template id to convert files with
	TemplateId string
	// Optional notify url for each assembly.
	NotifyUrl string
	// Instead of using templates you can define steps
	Steps map[string]map[string]interface{}
	// If true the original files will be copied in the output directoy with `-original_0_` prefix.
	// If false input files will be deleted.
	Preserve bool
	// Read template from JSON file. See exmaples for information about the format.
	TemplateFile string
	// If true file which are already in the directory won't be compiled.
	DontProcessDir bool
}

type Watcher struct {
	client *Client
	// Original options passed to Client.Watch.
	// If Input or Output are relative to the home directory
	// they will be extended (~/input -> /home/user/input)
	Options *WatchOptions
	stopped bool
	// Listen for errors
	Error chan error
	// Listen for completed assemblies
	Done chan *AssemblyInfo
	// Listen for file changes (only if Watch == true)
	Change       chan string
	stop         chan bool
	recentWrites map[string]time.Time
	blacklist    map[string]bool
}

// Watch a directory for changes and convert all changes files and download the result.
// It will create a new assembly for each file.
// If the directory already contains some they are all converted.
// See WatchOptions for possible configuration.
func (client *Client) Watch(options *WatchOptions) *Watcher {
	options.Input = expandPath(options.Input)
	options.Output = expandPath(options.Output)

	if options.TemplateFile != "" {
		options.TemplateFile = expandPath(options.TemplateFile)
		options.Steps = readJson(options.TemplateFile)
	}

	if _, err := os.Stat(options.Input); os.IsNotExist(err) {
		panic(fmt.Errorf("Input directory does not exist: %s", options.Input))
	}

	if _, err := os.Stat(options.Output); os.IsNotExist(err) {
		panic(fmt.Errorf("Output directory does not exist: %s", options.Output))
	}

	if options.Input == options.Output {
		panic(fmt.Errorf("Input and output directory are both: %s", options.Output))
	}

	watcher := &Watcher{
		client:       client,
		Options:      options,
		Error:        make(chan error, 1),
		Done:         make(chan *AssemblyInfo),
		Change:       make(chan string),
		stop:         make(chan bool, 1),
		recentWrites: make(map[string]time.Time),
		blacklist:    make(map[string]bool),
	}

	watcher.start()

	return watcher
}

func (watcher *Watcher) start() {
	if watcher.Options.Watch {
		go watcher.startWatcher()
	}

	if !watcher.Options.DontProcessDir {
		go watcher.processDir()
	}

	// If we have nothing to do (neither processing the input directory nor
	// watching it for changes) stop everything again
	if !watcher.Options.Watch && watcher.Options.DontProcessDir {
		watcher.Stop()
	}
}

// Stop the watcher.
func (watcher *Watcher) Stop() {
	if watcher.stopped {
		return
	}

	watcher.stopped = true

	close(watcher.Done)
	close(watcher.Error)
	close(watcher.Change)
	close(watcher.stop)
}

func (watcher *Watcher) processDir() {
	files, err := ioutil.ReadDir(watcher.Options.Input)
	if err != nil {
		watcher.error(err)
		return
	}

	input := watcher.Options.Input
	var wg sync.WaitGroup

	for _, file := range files {
		if !file.IsDir() {
			wg.Add(1)
			go func(file os.FileInfo) {
				watcher.processFile(path.Join(input, file.Name()))
				wg.Done()
			}(file)
		}
	}

	wg.Wait()

	// If watching is not enabled, stop and close all channels to avoid a
	// deadlock since we are done with everything.
	if !watcher.Options.Watch {
		watcher.Stop()
	}
}

func (watcher *Watcher) processFile(name string) {
	file, err := os.Open(name)
	if err != nil {
		watcher.error(err)
		return
	}

	// Add file to blacklist
	watcher.blacklist[name] = true

	assembly := watcher.client.CreateAssembly()

	if watcher.Options.TemplateId != "" {
		assembly.TemplateId = watcher.Options.TemplateId
	}

	if watcher.Options.NotifyUrl != "" {
		assembly.NotifyUrl = watcher.Options.NotifyUrl
	}

	for name, step := range watcher.Options.Steps {
		assembly.AddStep(name, step)
	}

	assembly.Blocking = true

	assembly.AddReader("file", path.Base(name), file)

	info, err := assembly.Upload()
	if err != nil {
		watcher.error(err)
		return
	}

	if info.Error != "" {
		watcher.error(errors.New(info.Error))
		return
	}

	var wg sync.WaitGroup

	for stepName, results := range info.Results {
		for index, result := range results {
			wg.Add(1)
			go func() {
				watcher.downloadResult(stepName, index, result)
				watcher.handleOriginalFile(name)
				delete(watcher.blacklist, name)
				watcher.Done <- info
				wg.Done()
			}()
		}
	}

	wg.Wait()

	if !watcher.Options.Watch {
		watcher.Stop()
	}
}

func (watcher *Watcher) downloadResult(stepName string, index int, result *FileInfo) {
	fileName := sanitizeRe.ReplaceAllString(fmt.Sprintf("%s_%d_%s", stepName, index, result.Name), "-")

	resp, err := http.Get(result.Url)
	if err != nil {
		watcher.error(err)
		return
	}

	defer resp.Body.Close()

	out, err := os.Create(path.Join(watcher.Options.Output, fileName))
	if err != nil {
		watcher.error(err)
		return
	}

	defer out.Close()

	io.Copy(out, resp.Body)
}

func (watcher *Watcher) startWatcher() {
	fsWatcher, err := fsnotify.NewWatcher()
	if err != nil {
		watcher.error(err)
		return
	}

	defer fsWatcher.Close()

	if err = fsWatcher.Add(watcher.Options.Input); err != nil {
		watcher.error(err)
		return
	}

	for {
		select {
		case _, more := <-watcher.stop:
			if !more {
				return
			}
		case err := <-fsWatcher.Errors:
			watcher.error(err)
		case evt := <-fsWatcher.Events:
			// Ignore the event if the file is currently processed
			if _, ok := watcher.blacklist[evt.Name]; ok == true {
				continue
			}
			if evt.Op&fsnotify.Create == fsnotify.Create || evt.Op&fsnotify.Write == fsnotify.Write {
				watcher.recentWrites[evt.Name] = time.Now()
			}
		case <-time.Tick(1 * time.Second):
			now := time.Now()

			for name, lastEvent := range watcher.recentWrites {
				diff := now.Sub(lastEvent)
				if diff > (time.Millisecond * 500) {
					delete(watcher.recentWrites, name)
					watcher.Change <- name
					go watcher.processFile(name)
				}
			}
		}
	}
}

func (watcher *Watcher) handleOriginalFile(name string) {
	var err error
	if watcher.Options.Preserve {
		_, file := path.Split(name)
		err = os.Rename(name, watcher.Options.Output+"/-original_0_"+basename(file))
	} else {
		err = os.Remove(name)
	}

	if err != nil {
		watcher.error(err)
	}
}

func (watcher *Watcher) error(err error) {
	watcher.Error <- err
}

func basename(name string) string {
	i := strings.LastIndex(name, string(os.PathSeparator))
	return name[i+1:]
}

func expandPath(str string) string {
	expanded, err := homedir.Expand(str)
	if err != nil {
		panic(err)
	}

	expanded, err = filepath.Abs(expanded)
	if err != nil {
		panic(err)
	}

	return expanded
}

func readJson(file string) map[string]map[string]interface{} {
	content, err := ioutil.ReadFile(file)
	if err != nil {
		panic(fmt.Errorf("Error reading template file: %s", err))
	}

	steps := make(map[string]map[string]interface{})

	err = json.Unmarshal(content, &steps)
	if err != nil {
		panic(fmt.Errorf("Error parsing template file: %s", err))
	}

	return steps
}
<|endoftext|>


<commit_before>// Copyright 2014-2015 Chadev. 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 (
	"fmt"
	"strings"

	"github.com/danryan/hal"
)

var whoisHandler = hear(`who\s?is ([A-Za-z0-9\-_\{\}\[\]\\]+)`, "who is", "Find out who a user is", func(res *hal.Response) error {
	name := res.Match[1]
	key := strings.ToUpper(name)
	val, err := res.Robot.Store.Get(key)
	if err != nil {
		res.Send(fmt.Sprintf("%s is no one to me\n", name))
		return err
	}
	return res.Send(fmt.Sprintf("%s is %s", name, string(val)))
})

var isHandler = hear(`([^(who|remember)].+) is (.+)`, "(name) is (role)", "Tell Ash who the user is", func(res *hal.Response) error {
	name := res.Match[1]
	key := strings.ToUpper(name)
	role := res.Match[2]

	storedRoles, err := res.Robot.Store.Get(key)
	roleToStore := role
	if len(storedRoles) > 0 {
		roleToStore = roleToStore + ", " + string(storedRoles)
	}

	err = res.Robot.Store.Set(key, []byte(roleToStore))
	if err != nil {
		res.Send("There's something wrong")
		return err
	}
	return res.Send(fmt.Sprintf("Got it. %s is %s\n", name, role))
})

var whoamHandler = hear(`who am (?i)I`, "who am I", "Ask ash who you are", func(res *hal.Response) error {
	name := res.UserName()
	key := strings.ToUpper(name)
	val, err := res.Robot.Store.Get(key)
	if err != nil {
		return res.Send(fmt.Sprintf("%s?  I have no memory of who you are.", name))
	}

	return res.Send(fmt.Sprintf("%s, you are %s", name, string(val)))
})
<commit_msg>clarify the "who is (role)" usage in help<commit_after>// Copyright 2014-2015 Chadev. 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 (
	"fmt"
	"strings"

	"github.com/danryan/hal"
)

var whoisHandler = hear(`who\s?is ([A-Za-z0-9\-_\{\}\[\]\\]+)`, "who is (name)", "Find out who a user is", func(res *hal.Response) error {
	name := res.Match[1]
	key := strings.ToUpper(name)
	val, err := res.Robot.Store.Get(key)
	if err != nil {
		res.Send(fmt.Sprintf("%s is no one to me\n", name))
		return err
	}
	return res.Send(fmt.Sprintf("%s is %s", name, string(val)))
})

var isHandler = hear(`([^(who|remember)].+) is (.+)`, "(name) is (role)", "Tell Ash who the user is", func(res *hal.Response) error {
	name := res.Match[1]
	key := strings.ToUpper(name)
	role := res.Match[2]

	storedRoles, err := res.Robot.Store.Get(key)
	roleToStore := role
	if len(storedRoles) > 0 {
		roleToStore = roleToStore + ", " + string(storedRoles)
	}

	err = res.Robot.Store.Set(key, []byte(roleToStore))
	if err != nil {
		res.Send("There's something wrong")
		return err
	}
	return res.Send(fmt.Sprintf("Got it. %s is %s\n", name, role))
})

var whoamHandler = hear(`who am (?i)I`, "who am I", "Ask ash who you are", func(res *hal.Response) error {
	name := res.UserName()
	key := strings.ToUpper(name)
	val, err := res.Robot.Store.Get(key)
	if err != nil {
		return res.Send(fmt.Sprintf("%s?  I have no memory of who you are.", name))
	}

	return res.Send(fmt.Sprintf("%s, you are %s", name, string(val)))
})
<|endoftext|>


<commit_before>package main

import (
	"crypto/tls"
	"flag"
	"fmt"
	"github.com/gorilla/websocket"
	"github.com/pkg/errors"
	"github.com/pterodactyl/wings/config"
	"github.com/pterodactyl/wings/server"
	"github.com/pterodactyl/wings/sftp"
	"github.com/remeh/sizedwaitgroup"
	"go.uber.org/zap"
	"net/http"
	"os"
)

var configPath = "config.yml"
var debug = false

// Entrypoint for the Wings application. Configures the logger and checks any
// flags that were passed through in the boot arguments.
func main() {
	flag.StringVar(&configPath, "config", "config.yml", "set the location for the configuration file")
	flag.BoolVar(&debug, "debug", false, "pass in order to run wings in debug mode")

	flag.Parse()

	c, err := config.ReadConfiguration(configPath)
	if err != nil {
		panic(err)
		return
	}

	if debug {
		c.Debug = true
	}

	printLogo()
	if err := configureLogging(c.Debug); err != nil {
		panic(err)
	}

	zap.S().Infof("using configuration from path: %s", configPath)
	if c.Debug {
		zap.S().Debugw("running in debug mode")
		zap.S().Infow("certificate checking is disabled")

		http.DefaultTransport.(*http.Transport).TLSClientConfig = &tls.Config{
			InsecureSkipVerify: true,
		}
	}

	config.Set(c)
	config.SetDebugViaFlag(debug)

	zap.S().Infof("checking for pterodactyl system user \"%s\"", c.System.User)
	if su, err := c.EnsurePterodactylUser(); err != nil {
		zap.S().Panicw("failed to create pterodactyl system user", zap.Error(err))
		return
	} else {
		zap.S().Infow("configured system user", zap.String("username", su.Username), zap.String("uid", su.Uid), zap.String("gid", su.Gid))
	}

	zap.S().Infow("beginnning file permission setting on server data directories")
	if err := c.EnsureFilePermissions(); err != nil {
		zap.S().Errorw("failed to properly chown data directories", zap.Error(err))
	} else {
		zap.S().Infow("finished ensuring file permissions")
	}

	if err := server.LoadDirectory("data/servers",&c.System); err != nil {
		zap.S().Fatalw("failed to load server configurations", zap.Error(err))
		return
	}

	if err := ConfigureDockerEnvironment(&c.Docker); err != nil {
		zap.S().Fatalw("failed to configure docker environment", zap.Error(errors.WithStack(err)))
		os.Exit(1)
	}

	if err := c.WriteToDisk(); err != nil {
		zap.S().Errorw("failed to save configuration to disk", zap.Error(errors.WithStack(err)))
	}

	// Just for some nice log output.
	for _, s := range server.GetServers().All() {
		zap.S().Infow("loaded configuration for server", zap.String("server", s.Uuid))
	}

	// Create a new WaitGroup that limits us to 4 servers being bootstrapped at a time
	// on Wings. This allows us to ensure the environment exists, write configurations,
	// and reboot processes without causing a slow-down due to sequential booting.
	wg := sizedwaitgroup.New(4)

	for _, serv := range server.GetServers().All() {
		wg.Add()

		go func(s *server.Server) {
			defer wg.Done()

			// Create a server environment if none exists currently. This allows us to recover from Docker
			// being reinstalled on the host system for example.
			zap.S().Infow("ensuring envrionment exists", zap.String("server", s.Uuid))
			if err := s.Environment.Create(); err != nil {
				zap.S().Errorw("failed to create an environment for server", zap.String("server", s.Uuid), zap.Error(err))
			}

			if r, err := s.Environment.IsRunning(); err != nil {
				zap.S().Errorw("error checking server environment status", zap.String("server", s.Uuid), zap.Error(err))
			} else if r {
				// If the server is currently running on Docker, mark the process as being in that state.
				// We never want to stop an instance that is currently running external from Wings since
				// that is a good way of keeping things running even if Wings gets in a very corrupted state.
				zap.S().Infow("detected server is running, re-attaching to process", zap.String("server", s.Uuid))
				if err := s.Sync(); err != nil {
					zap.S().Errorw("failed to sync server state, cannot mark as running", zap.String("server", s.Uuid), zap.Error(errors.WithStack(err)))
				} else {
					s.SetState(server.ProcessRunningState)

					// If we cannot attach to the environment go ahead and mark the processs as being offline.
					if err := s.Environment.Attach(); err != nil {
						zap.S().Warnw("error attaching to server environment", zap.String("server", s.Uuid), zap.Error(err))
						s.SetState(server.ProcessOfflineState)
					}
				}
			} else if !r {
				// If the server is not in a running state right now but according to the configuration it
				// should be, we want to go ahead and restart the instance.
				if s.State == server.ProcessRunningState || s.State == server.ProcessStartingState {
					zap.S().Infow(
						"server state does not match last recorded state in configuration, starting instance now",
						zap.String("server", s.Uuid),
					)

					if err := s.Environment.Start(); err != nil {
						zap.S().Warnw(
							"failed to put server instance back in running state",
							zap.String("server", s.Uuid),
							zap.Error(errors.WithStack(err)),
						)
					}
				} else {
					if s.State == "" {
						// Addresses potentially invalid data in the stored file that can cause Wings to lose
						// track of what the actual server state is.
						s.SetState(server.ProcessOfflineState)
					}
				}
			}
		}(serv)
	}

	// Wait until all of the servers are ready to go before we fire up the HTTP server.
	wg.Wait()

	// If the SFTP subsystem should be started, do so now.
	if c.System.Sftp.UseInternalSystem {
		sftp.Initialize(c)
	}

	r := &Router{
		token: c.AuthenticationToken,
		upgrader: websocket.Upgrader{
			// Ensure that the websocket request is originating from the Panel itself,
			// and not some other location.
			CheckOrigin: func(r *http.Request) bool {
				return r.Header.Get("Origin") == c.PanelLocation
			},
		},
	}

	router := r.ConfigureRouter()
	zap.S().Infow("configuring webserver", zap.Bool("ssl", c.Api.Ssl.Enabled), zap.String("host", c.Api.Host), zap.Int("port", c.Api.Port))

	addr := fmt.Sprintf("%s:%d", c.Api.Host, c.Api.Port)
	if c.Api.Ssl.Enabled {
		if err := http.ListenAndServeTLS(addr, c.Api.Ssl.CertificateFile, c.Api.Ssl.KeyFile, router); err != nil {
			zap.S().Fatalw("failed to configure HTTPS server", zap.Error(err))
		}
	} else {
		if err := http.ListenAndServe(addr, router); err != nil {
			zap.S().Fatalw("failed to configure HTTP server", zap.Error(err))
		}
	}
}

// Configures the global logger for Zap so that we can call it from any location
// in the code without having to pass around a logger instance.
func configureLogging(debug bool) error {
	cfg := zap.NewProductionConfig()
	if debug {
		cfg = zap.NewDevelopmentConfig()
	}

	cfg.Encoding = "console"
	cfg.OutputPaths = []string{
		"stdout",
	}

	logger, err := cfg.Build()
	if err != nil {
		return err
	}

	zap.ReplaceGlobals(logger)

	return nil
}

// Prints the wings logo, nothing special here!
func printLogo() {
	fmt.Println()
	fmt.Println(`                     ____`)
	fmt.Println(`__ Pterodactyl _____/___/_______ _______ ______`)
	fmt.Println(`\_____\    \/\/    /   /       /  __   /   ___/`)
	fmt.Println(`   \___\          /   /   /   /  /_/  /___   /`)
	fmt.Println(`        \___/\___/___/___/___/___    /______/`)
	fmt.Println(`                            /_______/ v` + Version)
	fmt.Println()
}
<commit_msg>Show stack on error better<commit_after>package main

import (
	"crypto/tls"
	"flag"
	"fmt"
	"github.com/gorilla/websocket"
	"github.com/pkg/errors"
	"github.com/pterodactyl/wings/config"
	"github.com/pterodactyl/wings/server"
	"github.com/pterodactyl/wings/sftp"
	"github.com/remeh/sizedwaitgroup"
	"go.uber.org/zap"
	"net/http"
	"os"
)

var configPath = "config.yml"
var debug = false

// Entrypoint for the Wings application. Configures the logger and checks any
// flags that were passed through in the boot arguments.
func main() {
	flag.StringVar(&configPath, "config", "config.yml", "set the location for the configuration file")
	flag.BoolVar(&debug, "debug", false, "pass in order to run wings in debug mode")

	flag.Parse()

	c, err := config.ReadConfiguration(configPath)
	if err != nil {
		panic(err)
		return
	}

	if debug {
		c.Debug = true
	}

	printLogo()
	if err := configureLogging(c.Debug); err != nil {
		panic(err)
	}

	zap.S().Infof("using configuration from path: %s", configPath)
	if c.Debug {
		zap.S().Debugw("running in debug mode")
		zap.S().Infow("certificate checking is disabled")

		http.DefaultTransport.(*http.Transport).TLSClientConfig = &tls.Config{
			InsecureSkipVerify: true,
		}
	}

	config.Set(c)
	config.SetDebugViaFlag(debug)

	zap.S().Infof("checking for pterodactyl system user \"%s\"", c.System.User)
	if su, err := c.EnsurePterodactylUser(); err != nil {
		zap.S().Panicw("failed to create pterodactyl system user", zap.Error(err))
		return
	} else {
		zap.S().Infow("configured system user", zap.String("username", su.Username), zap.String("uid", su.Uid), zap.String("gid", su.Gid))
	}

	zap.S().Infow("beginnning file permission setting on server data directories")
	if err := c.EnsureFilePermissions(); err != nil {
		zap.S().Errorw("failed to properly chown data directories", zap.Error(err))
	} else {
		zap.S().Infow("finished ensuring file permissions")
	}

	if err := server.LoadDirectory("data/servers", &c.System); err != nil {
		zap.S().Fatalw("failed to load server configurations", zap.Error(errors.WithStack(err)))
		return
	}

	if err := ConfigureDockerEnvironment(&c.Docker); err != nil {
		zap.S().Fatalw("failed to configure docker environment", zap.Error(errors.WithStack(err)))
		os.Exit(1)
	}

	if err := c.WriteToDisk(); err != nil {
		zap.S().Errorw("failed to save configuration to disk", zap.Error(errors.WithStack(err)))
	}

	// Just for some nice log output.
	for _, s := range server.GetServers().All() {
		zap.S().Infow("loaded configuration for server", zap.String("server", s.Uuid))
	}

	// Create a new WaitGroup that limits us to 4 servers being bootstrapped at a time
	// on Wings. This allows us to ensure the environment exists, write configurations,
	// and reboot processes without causing a slow-down due to sequential booting.
	wg := sizedwaitgroup.New(4)

	for _, serv := range server.GetServers().All() {
		wg.Add()

		go func(s *server.Server) {
			defer wg.Done()

			// Create a server environment if none exists currently. This allows us to recover from Docker
			// being reinstalled on the host system for example.
			zap.S().Infow("ensuring envrionment exists", zap.String("server", s.Uuid))
			if err := s.Environment.Create(); err != nil {
				zap.S().Errorw("failed to create an environment for server", zap.String("server", s.Uuid), zap.Error(err))
			}

			if r, err := s.Environment.IsRunning(); err != nil {
				zap.S().Errorw("error checking server environment status", zap.String("server", s.Uuid), zap.Error(err))
			} else if r {
				// If the server is currently running on Docker, mark the process as being in that state.
				// We never want to stop an instance that is currently running external from Wings since
				// that is a good way of keeping things running even if Wings gets in a very corrupted state.
				zap.S().Infow("detected server is running, re-attaching to process", zap.String("server", s.Uuid))
				if err := s.Sync(); err != nil {
					zap.S().Errorw("failed to sync server state, cannot mark as running", zap.String("server", s.Uuid), zap.Error(errors.WithStack(err)))
				} else {
					s.SetState(server.ProcessRunningState)

					// If we cannot attach to the environment go ahead and mark the processs as being offline.
					if err := s.Environment.Attach(); err != nil {
						zap.S().Warnw("error attaching to server environment", zap.String("server", s.Uuid), zap.Error(err))
						s.SetState(server.ProcessOfflineState)
					}
				}
			} else if !r {
				// If the server is not in a running state right now but according to the configuration it
				// should be, we want to go ahead and restart the instance.
				if s.State == server.ProcessRunningState || s.State == server.ProcessStartingState {
					zap.S().Infow(
						"server state does not match last recorded state in configuration, starting instance now",
						zap.String("server", s.Uuid),
					)

					if err := s.Environment.Start(); err != nil {
						zap.S().Warnw(
							"failed to put server instance back in running state",
							zap.String("server", s.Uuid),
							zap.Error(errors.WithStack(err)),
						)
					}
				} else {
					if s.State == "" {
						// Addresses potentially invalid data in the stored file that can cause Wings to lose
						// track of what the actual server state is.
						s.SetState(server.ProcessOfflineState)
					}
				}
			}
		}(serv)
	}

	// Wait until all of the servers are ready to go before we fire up the HTTP server.
	wg.Wait()

	// If the SFTP subsystem should be started, do so now.
	if c.System.Sftp.UseInternalSystem {
		sftp.Initialize(c)
	}

	r := &Router{
		token: c.AuthenticationToken,
		upgrader: websocket.Upgrader{
			// Ensure that the websocket request is originating from the Panel itself,
			// and not some other location.
			CheckOrigin: func(r *http.Request) bool {
				return r.Header.Get("Origin") == c.PanelLocation
			},
		},
	}

	router := r.ConfigureRouter()
	zap.S().Infow("configuring webserver", zap.Bool("ssl", c.Api.Ssl.Enabled), zap.String("host", c.Api.Host), zap.Int("port", c.Api.Port))

	addr := fmt.Sprintf("%s:%d", c.Api.Host, c.Api.Port)
	if c.Api.Ssl.Enabled {
		if err := http.ListenAndServeTLS(addr, c.Api.Ssl.CertificateFile, c.Api.Ssl.KeyFile, router); err != nil {
			zap.S().Fatalw("failed to configure HTTPS server", zap.Error(err))
		}
	} else {
		if err := http.ListenAndServe(addr, router); err != nil {
			zap.S().Fatalw("failed to configure HTTP server", zap.Error(err))
		}
	}
}

// Configures the global logger for Zap so that we can call it from any location
// in the code without having to pass around a logger instance.
func configureLogging(debug bool) error {
	cfg := zap.NewProductionConfig()
	if debug {
		cfg = zap.NewDevelopmentConfig()
	}

	cfg.Encoding = "console"
	cfg.OutputPaths = []string{
		"stdout",
	}

	logger, err := cfg.Build()
	if err != nil {
		return err
	}

	zap.ReplaceGlobals(logger)

	return nil
}

// Prints the wings logo, nothing special here!
func printLogo() {
	fmt.Println()
	fmt.Println(`                     ____`)
	fmt.Println(`__ Pterodactyl _____/___/_______ _______ ______`)
	fmt.Println(`\_____\    \/\/    /   /       /  __   /   ___/`)
	fmt.Println(`   \___\          /   /   /   /  /_/  /___   /`)
	fmt.Println(`        \___/\___/___/___/___/___    /______/`)
	fmt.Println(`                            /_______/ v` + Version)
	fmt.Println()
}
<|endoftext|>


<commit_before>// Copyright (c) 2017, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

package main

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/importer"
	"go/parser"
	"go/printer"
	"go/token"
	"go/types"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/mvdan/sh/interp"
	"github.com/mvdan/sh/syntax"
)

var (
	rawPrinter = printer.Config{Mode: printer.RawFormat}

	fastTest = false
)

type reducer struct {
	tdir      string
	logOut    io.Writer
	matchRe   *regexp.Regexp
	shellProg *syntax.File

	fset     *token.FileSet
	origFset *token.FileSet
	pkg      *ast.Package
	files    []*ast.File
	file     *ast.File

	tconf types.Config
	info  *types.Info

	useIdents map[types.Object][]*ast.Ident
	revDefs   map[types.Object]*ast.Ident
	parents   map[ast.Node]ast.Node

	dstBuf *bytes.Buffer

	tmpFiles map[*ast.File]*os.File

	tries     int
	didChange bool

	deleteKeepUnderscore func()
	deleteKeepUnchanged  func()

	tried map[string]bool

	walker
}

var errNoReduction = fmt.Errorf("could not reduce program")

func reduce(dir, match string, logOut io.Writer, shellStr string) error {
	r := &reducer{
		tdir:   dir,
		logOut: logOut,
		tried:  make(map[string]bool, 16),
		dstBuf: bytes.NewBuffer(nil),
	}
	var err error
	if r.tdir, err = ioutil.TempDir("", "goreduce"); err != nil {
		return err
	}
	defer os.RemoveAll(r.tdir)
	if r.matchRe, err = regexp.Compile(match); err != nil {
		return err
	}
	r.fset = token.NewFileSet()
	pkgs, err := parser.ParseDir(r.fset, dir, nil, parser.ParseComments)
	if err != nil {
		return err
	}
	if len(pkgs) != 1 {
		return fmt.Errorf("expected 1 package, got %d", len(pkgs))
	}
	for _, pkg := range pkgs {
		r.pkg = pkg
	}
	switch {
	case shellStr != "":
	case r.pkg.Name == "main":
		shellStr = shellStrRun
	default:
		shellStr = shellStrBuild
	}
	r.shellProg, err = syntax.NewParser().Parse(strings.NewReader(shellStr), "")
	if err != nil {
		return err
	}
	r.origFset = token.NewFileSet()
	parser.ParseDir(r.origFset, dir, nil, 0)
	tfnames := make([]string, 0, len(r.pkg.Files)+1)

	var restoreMain func()
	r.tmpFiles = make(map[*ast.File]*os.File, len(r.pkg.Files))
	for fpath, file := range r.pkg.Files {
		r.files = append(r.files, file)
		tfname := filepath.Join(r.tdir, filepath.Base(fpath))
		f, err := os.Create(tfname)
		if err != nil {
			return err
		}
		if err := rawPrinter.Fprint(f, r.fset, file); err != nil {
			return err
		}
		r.tmpFiles[file] = f
		defer f.Close()
		tfnames = append(tfnames, tfname)
	}
	r.tconf.Importer = importer.Default()
	r.tconf.Error = func(err error) {
		if terr, ok := err.(types.Error); ok && terr.Soft {
			// don't stop type-checking on soft errors
			return
		}
		//panic("types.Check should not error here: " + err.Error())
	}
	// Check that the output matches before we apply any changes
	if !fastTest {
		if err := r.checkRun(); err != nil {
			return err
		}
	}
	r.fillParents()
	if anyChanges := r.reduceLoop(); !anyChanges {
		return errNoReduction
	}
	if restoreMain != nil {
		restoreMain()
	}
	for astFile := range r.tmpFiles {
		astFile.Name.Name = r.pkg.Name
		fname := r.fset.Position(astFile.Pos()).Filename
		f, err := os.Create(fname)
		if err != nil {
			return err
		}
		if err := printer.Fprint(f, r.fset, astFile); err != nil {
			return err
		}
		if err := f.Close(); err != nil {
			return err
		}
	}
	return nil
}

func (r *reducer) logChange(node ast.Node, format string, a ...interface{}) {
	if *verbose {
		pos := r.origFset.Position(node.Pos())
		times := "first try"
		if r.tries != 1 {
			times = fmt.Sprintf("%d tries", r.tries)
		}
		fmt.Fprintf(r.logOut, "%s:%d: %s (%s)\n",
			pos.Filename, pos.Line, fmt.Sprintf(format, a...), times)
	}
	r.tries = 0
}

func (r *reducer) checkRun() error {
	out := r.runCmd()
	if out == nil {
		return fmt.Errorf("expected an error to occur")
	}
	if !r.matchRe.Match(out) {
		return fmt.Errorf("error does not match:\n%s", string(out))
	}
	return nil
}

func (r *reducer) okChangeNoUndo() bool {
	if r.didChange {
		return false
	}
	r.dstBuf.Reset()
	if err := rawPrinter.Fprint(r.dstBuf, r.fset, r.file); err != nil {
		return false
	}
	newSrc := r.dstBuf.String()
	if r.tried[newSrc] {
		return false
	}
	r.tries++
	r.tried[newSrc] = true
	f := r.tmpFiles[r.file]
	if err := f.Truncate(0); err != nil {
		return false
	}
	if _, err := f.Seek(0, io.SeekStart); err != nil {
		return false
	}
	if _, err := f.Write(r.dstBuf.Bytes()); err != nil {
		return false
	}
	if err := r.checkRun(); err != nil {
		return false
	}
	// Reduction worked
	r.didChange = true
	return true
}

func (r *reducer) okChange() bool {
	if r.okChangeNoUndo() {
		r.deleteKeepUnderscore = nil
		r.deleteKeepUnchanged = nil
		return true
	}
	if r.deleteKeepUnderscore != nil {
		r.deleteKeepUnderscore()
		r.deleteKeepUnderscore = nil
		return r.okChange()
	}
	if r.deleteKeepUnchanged != nil {
		r.deleteKeepUnchanged()
		r.deleteKeepUnchanged = nil
	}
	return false
}

func (r *reducer) reduceLoop() (anyChanges bool) {
	r.info = &types.Info{
		Defs: make(map[*ast.Ident]types.Object),
		Uses: make(map[*ast.Ident]types.Object),
	}
	for {
		// Update type info after the AST changes
		r.tconf.Check(r.tdir, r.fset, r.files, r.info)
		r.fillObjs()

		r.didChange = false
		r.walk(r.pkg, r.reduceNode)
		if !r.didChange {
			if *verbose {
				fmt.Fprintf(r.logOut, "gave up after %d final tries\n", r.tries)
			}
			return
		}
		anyChanges = true
	}
}

func (r *reducer) fillObjs() {
	r.revDefs = make(map[types.Object]*ast.Ident, len(r.info.Defs))
	for id, obj := range r.info.Defs {
		if obj == nil {
			continue
		}
		r.revDefs[obj] = id
	}
	r.useIdents = make(map[types.Object][]*ast.Ident, len(r.info.Uses)/2)
	for id, obj := range r.info.Uses {
		if pkg := obj.Pkg(); pkg == nil || pkg.Name() != r.pkg.Name {
			// builtin or declared outside of our pkg
			continue
		}
		r.useIdents[obj] = append(r.useIdents[obj], id)
	}
}

func (r *reducer) fillParents() {
	r.parents = make(map[ast.Node]ast.Node)
	stack := make([]ast.Node, 1, 32)
	ast.Inspect(r.pkg, func(node ast.Node) bool {
		if node == nil {
			stack = stack[:len(stack)-1]
			return true
		}
		r.parents[node] = stack[len(stack)-1]
		stack = append(stack, node)
		return true
	})
}

func findFunc(file *ast.File, name string) *ast.FuncDecl {
	for _, decl := range file.Decls {
		fd, _ := decl.(*ast.FuncDecl)
		if fd != nil && fd.Name.Name == name {
			return fd
		}
	}
	return nil
}

func delFunc(file *ast.File, name string) (undo func()) {
	oldDecls := file.Decls
	for i, decl := range file.Decls {
		fd, _ := decl.(*ast.FuncDecl)
		if fd != nil && fd.Name.Name == name {
			file.Decls = append(file.Decls[:i], file.Decls[i+1:]...)
			return func() { file.Decls = oldDecls }
		}
	}
	return nil
}

func (r *reducer) runCmd() []byte {
	var buf bytes.Buffer
	runner := interp.Runner{
		File:   r.shellProg,
		Dir:    r.tdir,
		Stdout: &buf,
		Stderr: &buf,
	}
	runner.Run()
	return buf.Bytes()
}
<commit_msg>Remove some old, unused code<commit_after>// Copyright (c) 2017, Daniel Martí <mvdan@mvdan.cc>
// See LICENSE for licensing information

package main

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/importer"
	"go/parser"
	"go/printer"
	"go/token"
	"go/types"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/mvdan/sh/interp"
	"github.com/mvdan/sh/syntax"
)

var (
	rawPrinter = printer.Config{Mode: printer.RawFormat}

	fastTest = false
)

type reducer struct {
	tdir      string
	logOut    io.Writer
	matchRe   *regexp.Regexp
	shellProg *syntax.File

	fset     *token.FileSet
	origFset *token.FileSet
	pkg      *ast.Package
	files    []*ast.File
	file     *ast.File

	tconf types.Config
	info  *types.Info

	useIdents map[types.Object][]*ast.Ident
	revDefs   map[types.Object]*ast.Ident
	parents   map[ast.Node]ast.Node

	dstBuf *bytes.Buffer

	tmpFiles map[*ast.File]*os.File

	tries     int
	didChange bool

	deleteKeepUnderscore func()
	deleteKeepUnchanged  func()

	tried map[string]bool

	walker
}

var errNoReduction = fmt.Errorf("could not reduce program")

func reduce(dir, match string, logOut io.Writer, shellStr string) error {
	r := &reducer{
		tdir:   dir,
		logOut: logOut,
		tried:  make(map[string]bool, 16),
		dstBuf: bytes.NewBuffer(nil),
	}
	var err error
	if r.tdir, err = ioutil.TempDir("", "goreduce"); err != nil {
		return err
	}
	defer os.RemoveAll(r.tdir)
	if r.matchRe, err = regexp.Compile(match); err != nil {
		return err
	}
	r.fset = token.NewFileSet()
	pkgs, err := parser.ParseDir(r.fset, dir, nil, parser.ParseComments)
	if err != nil {
		return err
	}
	if len(pkgs) != 1 {
		return fmt.Errorf("expected 1 package, got %d", len(pkgs))
	}
	for _, pkg := range pkgs {
		r.pkg = pkg
	}
	switch {
	case shellStr != "":
	case r.pkg.Name == "main":
		shellStr = shellStrRun
	default:
		shellStr = shellStrBuild
	}
	r.shellProg, err = syntax.NewParser().Parse(strings.NewReader(shellStr), "")
	if err != nil {
		return err
	}
	r.origFset = token.NewFileSet()
	parser.ParseDir(r.origFset, dir, nil, 0)

	var restoreMain func()
	r.tmpFiles = make(map[*ast.File]*os.File, len(r.pkg.Files))
	for fpath, file := range r.pkg.Files {
		r.files = append(r.files, file)
		tfname := filepath.Join(r.tdir, filepath.Base(fpath))
		f, err := os.Create(tfname)
		if err != nil {
			return err
		}
		if err := rawPrinter.Fprint(f, r.fset, file); err != nil {
			return err
		}
		r.tmpFiles[file] = f
		defer f.Close()
	}
	r.tconf.Importer = importer.Default()
	r.tconf.Error = func(err error) {
		if terr, ok := err.(types.Error); ok && terr.Soft {
			// don't stop type-checking on soft errors
			return
		}
		//panic("types.Check should not error here: " + err.Error())
	}
	// Check that the output matches before we apply any changes
	if !fastTest {
		if err := r.checkRun(); err != nil {
			return err
		}
	}
	r.fillParents()
	if anyChanges := r.reduceLoop(); !anyChanges {
		return errNoReduction
	}
	if restoreMain != nil {
		restoreMain()
	}
	for astFile := range r.tmpFiles {
		astFile.Name.Name = r.pkg.Name
		fname := r.fset.Position(astFile.Pos()).Filename
		f, err := os.Create(fname)
		if err != nil {
			return err
		}
		if err := printer.Fprint(f, r.fset, astFile); err != nil {
			return err
		}
		if err := f.Close(); err != nil {
			return err
		}
	}
	return nil
}

func (r *reducer) logChange(node ast.Node, format string, a ...interface{}) {
	if *verbose {
		pos := r.origFset.Position(node.Pos())
		times := "first try"
		if r.tries != 1 {
			times = fmt.Sprintf("%d tries", r.tries)
		}
		fmt.Fprintf(r.logOut, "%s:%d: %s (%s)\n",
			pos.Filename, pos.Line, fmt.Sprintf(format, a...), times)
	}
	r.tries = 0
}

func (r *reducer) checkRun() error {
	out := r.runCmd()
	if out == nil {
		return fmt.Errorf("expected an error to occur")
	}
	if !r.matchRe.Match(out) {
		return fmt.Errorf("error does not match:\n%s", string(out))
	}
	return nil
}

func (r *reducer) okChangeNoUndo() bool {
	if r.didChange {
		return false
	}
	r.dstBuf.Reset()
	if err := rawPrinter.Fprint(r.dstBuf, r.fset, r.file); err != nil {
		return false
	}
	newSrc := r.dstBuf.String()
	if r.tried[newSrc] {
		return false
	}
	r.tries++
	r.tried[newSrc] = true
	f := r.tmpFiles[r.file]
	if err := f.Truncate(0); err != nil {
		return false
	}
	if _, err := f.Seek(0, io.SeekStart); err != nil {
		return false
	}
	if _, err := f.Write(r.dstBuf.Bytes()); err != nil {
		return false
	}
	if err := r.checkRun(); err != nil {
		return false
	}
	// Reduction worked
	r.didChange = true
	return true
}

func (r *reducer) okChange() bool {
	if r.okChangeNoUndo() {
		r.deleteKeepUnderscore = nil
		r.deleteKeepUnchanged = nil
		return true
	}
	if r.deleteKeepUnderscore != nil {
		r.deleteKeepUnderscore()
		r.deleteKeepUnderscore = nil
		return r.okChange()
	}
	if r.deleteKeepUnchanged != nil {
		r.deleteKeepUnchanged()
		r.deleteKeepUnchanged = nil
	}
	return false
}

func (r *reducer) reduceLoop() (anyChanges bool) {
	r.info = &types.Info{
		Defs: make(map[*ast.Ident]types.Object),
		Uses: make(map[*ast.Ident]types.Object),
	}
	for {
		// Update type info after the AST changes
		r.tconf.Check(r.tdir, r.fset, r.files, r.info)
		r.fillObjs()

		r.didChange = false
		r.walk(r.pkg, r.reduceNode)
		if !r.didChange {
			if *verbose {
				fmt.Fprintf(r.logOut, "gave up after %d final tries\n", r.tries)
			}
			return
		}
		anyChanges = true
	}
}

func (r *reducer) fillObjs() {
	r.revDefs = make(map[types.Object]*ast.Ident, len(r.info.Defs))
	for id, obj := range r.info.Defs {
		if obj == nil {
			continue
		}
		r.revDefs[obj] = id
	}
	r.useIdents = make(map[types.Object][]*ast.Ident, len(r.info.Uses)/2)
	for id, obj := range r.info.Uses {
		if pkg := obj.Pkg(); pkg == nil || pkg.Name() != r.pkg.Name {
			// builtin or declared outside of our pkg
			continue
		}
		r.useIdents[obj] = append(r.useIdents[obj], id)
	}
}

func (r *reducer) fillParents() {
	r.parents = make(map[ast.Node]ast.Node)
	stack := make([]ast.Node, 1, 32)
	ast.Inspect(r.pkg, func(node ast.Node) bool {
		if node == nil {
			stack = stack[:len(stack)-1]
			return true
		}
		r.parents[node] = stack[len(stack)-1]
		stack = append(stack, node)
		return true
	})
}

func (r *reducer) runCmd() []byte {
	var buf bytes.Buffer
	runner := interp.Runner{
		File:   r.shellProg,
		Dir:    r.tdir,
		Stdout: &buf,
		Stderr: &buf,
	}
	runner.Run()
	return buf.Bytes()
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/howeyc/fsnotify"
)

const (
	subSymbol = "{}"
)

var (
	regexString string
	matchAll    = regexp.MustCompile(".*")
)

func init() {
	flag.StringVar(&regexString, "r", "", "The regular expression to match filenames.")
}

func parseRegex() *regexp.Regexp {
	if regexString == "" {
		fmt.Println("Warning: no regex given (-r), so matching all file changes.")
		return matchAll
	}

	r, err := regexp.Compile(regexString)
	if err != nil {
		Fatalln("Bad regular expression provided.\n" + err.Error())
	}
	return r
}

func Fatalln(args ...interface{}) {
	fmt.Println(args...)
	os.Exit(1)
}

func walker(watcher *fsnotify.Watcher) filepath.WalkFunc {
	return func(path string, f os.FileInfo, err error) error {
		if !f.IsDir() {
			return nil
		}
		fmt.Println("Adding watch for path", path)
		if err := watcher.Watch(path); err != nil {
			// TODO: handle this somehow?
			fmt.Printf("Error while watching new path %s: %s\n", path, err)
		}
		return nil
	}
}

func watch(root string, watcher *fsnotify.Watcher, names chan<- string, done chan<- error) {
	if err := watcher.Watch(root); err != nil {
		Fatalln(err)
	}

	for {
		select {
		case e := <-watcher.Event:
			path := strings.TrimPrefix(e.Name, "./")
			names <- path
			if e.IsCreate() {
				if err := filepath.Walk(path, walker(watcher)); err != nil {
					// TODO: handle this somehow?
					fmt.Printf("Error while walking path %s: %s\n", path, err)
				}
			}
			if e.IsDelete() {
				watcher.RemoveWatch(path)
			}
		case err := <-watcher.Error:
			done <- err
			return
		}
	}
}

// runCommand runs the specified command, replacing {} in args with the provided filename. Blocks until the
// command exits.
func runCommand(command []string, path string) {
	replacer := strings.NewReplacer(subSymbol, path)
	args := make([]string, len(command))
	for i, c := range command {
		args[i] = replacer.Replace(c)
	}

	cmd := exec.Command(args[0], args[1:]...)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		fmt.Println("Error running command:", err)
		return
	}
	go io.Copy(os.Stdout, stdout)
	stderr, err := cmd.StderrPipe()
	if err != nil {
		fmt.Println("Error running command:", err)
		return
	}
	go io.Copy(os.Stderr, stderr)

	if err := cmd.Run(); err != nil {
		fmt.Fprintf(os.Stderr, "(error exit: %s)\n", err)
	}
}

// filterMatching passes on messages matching regex.
func filterMatching(in <-chan string, out chan<- string, regex *regexp.Regexp) {
	for name := range in {
		if regex.MatchString(name) {
			out <- name
		}
	}
}

// batchRun receives realtime file notification events and batches them up. It's a bit tricky, but here's what
// it accomplishes:
// * When we initially get a message, wait a bit and batch messages before trying to send anything. This is
//	 because the file events come in quick bursts.
// * Once it's time to send, don't do it until the out channel is unblocked. In the meantime, keep batching.
//   When we've sent off all the batched messages, go back to the beginning.
func batchRun(in <-chan string, out chan<- string, backlog Backlog) {
	for name := range in {
		backlog.Add(name)
		timer := time.NewTimer(500 * time.Millisecond)
	outer:
		for {
			select {
			case name := <-in:
				backlog.Add(name)
			case <-timer.C:
				for {
					select {
					case name := <-in:
						backlog.Add(name)
					case out <- backlog.Next():
						if backlog.RemoveOne() {
							break outer
						}
					}
				}
			}
		}
	}
}

// runEach runs the command on each name that comes through the names channel.
func runEach(names <-chan string, command []string) {
	for name := range names {
		runCommand(command, name)
	}
}

func main() {
	flag.Parse()
	regex := parseRegex()
	command := flag.Args()
	if len(command) == 0 {
		Fatalln("Must give command to execute.")
	}
	substitution := false
	for _, part := range command {
		if strings.Contains(part, subSymbol) {
			substitution = true
			break
		}
	}
	var backlog Backlog
	if substitution {
		backlog = &UniqueFilesBacklog{true, "", make(map[string]struct{})}
	} else {
		backlog = new(UnifiedBacklog)
	}

	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		Fatalln(err)
	}
	defer watcher.Close()

	done := make(chan error)
	rawChanges := make(chan string)
	filtered := make(chan string)
	batched := make(chan string)

	go watch(".", watcher, rawChanges, done)
	go filterMatching(rawChanges, filtered, regex)
	go batchRun(filtered, batched, backlog)
	go runEach(batched, command)

	Fatalln(<-done)
}
<commit_msg>Lower the batching delay from 500ms to 200ms.<commit_after>package main

import (
	"flag"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/howeyc/fsnotify"
)

const (
	subSymbol = "{}"
)

var (
	regexString string
	matchAll    = regexp.MustCompile(".*")
)

func init() {
	flag.StringVar(&regexString, "r", "", "The regular expression to match filenames.")
}

func parseRegex() *regexp.Regexp {
	if regexString == "" {
		fmt.Println("Warning: no regex given (-r), so matching all file changes.")
		return matchAll
	}

	r, err := regexp.Compile(regexString)
	if err != nil {
		Fatalln("Bad regular expression provided.\n" + err.Error())
	}
	return r
}

func Fatalln(args ...interface{}) {
	fmt.Println(args...)
	os.Exit(1)
}

func walker(watcher *fsnotify.Watcher) filepath.WalkFunc {
	return func(path string, f os.FileInfo, err error) error {
		if !f.IsDir() {
			return nil
		}
		fmt.Println("Adding watch for path", path)
		if err := watcher.Watch(path); err != nil {
			// TODO: handle this somehow?
			fmt.Printf("Error while watching new path %s: %s\n", path, err)
		}
		return nil
	}
}

func watch(root string, watcher *fsnotify.Watcher, names chan<- string, done chan<- error) {
	if err := watcher.Watch(root); err != nil {
		Fatalln(err)
	}

	for {
		select {
		case e := <-watcher.Event:
			path := strings.TrimPrefix(e.Name, "./")
			names <- path
			if e.IsCreate() {
				if err := filepath.Walk(path, walker(watcher)); err != nil {
					// TODO: handle this somehow?
					fmt.Printf("Error while walking path %s: %s\n", path, err)
				}
			}
			if e.IsDelete() {
				watcher.RemoveWatch(path)
			}
		case err := <-watcher.Error:
			done <- err
			return
		}
	}
}

// runCommand runs the specified command, replacing {} in args with the provided filename. Blocks until the
// command exits.
func runCommand(command []string, path string) {
	replacer := strings.NewReplacer(subSymbol, path)
	args := make([]string, len(command))
	for i, c := range command {
		args[i] = replacer.Replace(c)
	}

	cmd := exec.Command(args[0], args[1:]...)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		fmt.Println("Error running command:", err)
		return
	}
	go io.Copy(os.Stdout, stdout)
	stderr, err := cmd.StderrPipe()
	if err != nil {
		fmt.Println("Error running command:", err)
		return
	}
	go io.Copy(os.Stderr, stderr)

	if err := cmd.Run(); err != nil {
		fmt.Fprintf(os.Stderr, "(error exit: %s)\n", err)
	}
}

// filterMatching passes on messages matching regex.
func filterMatching(in <-chan string, out chan<- string, regex *regexp.Regexp) {
	for name := range in {
		if regex.MatchString(name) {
			out <- name
		}
	}
}

// batchRun receives realtime file notification events and batches them up. It's a bit tricky, but here's what
// it accomplishes:
// * When we initially get a message, wait a bit and batch messages before trying to send anything. This is
//	 because the file events come in quick bursts.
// * Once it's time to send, don't do it until the out channel is unblocked. In the meantime, keep batching.
//   When we've sent off all the batched messages, go back to the beginning.
func batchRun(in <-chan string, out chan<- string, backlog Backlog) {
	for name := range in {
		backlog.Add(name)
		timer := time.NewTimer(200 * time.Millisecond)
	outer:
		for {
			select {
			case name := <-in:
				backlog.Add(name)
			case <-timer.C:
				for {
					select {
					case name := <-in:
						backlog.Add(name)
					case out <- backlog.Next():
						if backlog.RemoveOne() {
							break outer
						}
					}
				}
			}
		}
	}
}

// runEach runs the command on each name that comes through the names channel.
func runEach(names <-chan string, command []string) {
	for name := range names {
		runCommand(command, name)
	}
}

func main() {
	flag.Parse()
	regex := parseRegex()
	command := flag.Args()
	if len(command) == 0 {
		Fatalln("Must give command to execute.")
	}
	substitution := false
	for _, part := range command {
		if strings.Contains(part, subSymbol) {
			substitution = true
			break
		}
	}
	var backlog Backlog
	if substitution {
		backlog = &UniqueFilesBacklog{true, "", make(map[string]struct{})}
	} else {
		backlog = new(UnifiedBacklog)
	}

	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		Fatalln(err)
	}
	defer watcher.Close()

	done := make(chan error)
	rawChanges := make(chan string)
	filtered := make(chan string)
	batched := make(chan string)

	go watch(".", watcher, rawChanges, done)
	go filterMatching(rawChanges, filtered, regex)
	go batchRun(filtered, batched, backlog)
	go runEach(batched, command)

	Fatalln(<-done)
}
<|endoftext|>