<commit_before>package main

import (
	"fmt"
	"gnd.la/log"
	"math"
	"os"
	"os/exec"
	"reflect"
	"strings"
	"sync"
	"syscall"
	"time"
)

var (
	// avoids starting multiple services
	// at the same time, to enforce resource limits
	startLock sync.Mutex
)

type State uint8

const (
	StateStopped State = iota
	StateStopping
	StateStarted
	StateStarting
	StateBackoff
	StateFailed
)

func (s State) isRunState() bool {
	return s == StateStarted || s == StateStarting
}

func (s State) canStop() bool {
	return s.isRunState() || s == StateBackoff
}

const (
	minTime    = time.Second
	maxRetries = 10
)

type Service struct {
	mu         sync.Mutex // protects access to the fields
	st         sync.Mutex // prevents start/stop from running concurrently
	Config     *Config
	Cmd        *exec.Cmd
	State      State
	Started    time.Time
	Restarts   int
	Err        error
	stopCh     chan error
	errCh      chan error
	retries    int
	startTimer *time.Timer
	nextStart  time.Time
}

func newService(cfg *Config) *Service {
	return &Service{Config: cfg, stopCh: make(chan error)}
}

func (s *Service) Name() string {
	return s.Config.ServiceName()
}

func (s *Service) sendErr(ch *chan<- error, err error) {
	if err != nil && s.State != StateStopping {
		s.errorf("%v", err)
	}
	s.Err = err
	if ch != nil && *ch != nil {
		select {
		case *ch <- err:
			*ch = nil
		default:
		}
	}
	if s.errCh != nil {
		select {
		case s.errCh <- err:
		default:
		}
	}
}

func (s *Service) Start() error {
	s.st.Lock()
	defer s.st.Unlock()
	if s.State.isRunState() {
		return nil
	}
	if err := s.Config.Log.Open(); err != nil {
		return err
	}
	if err := s.startService(); err != nil {
		return err
	}
	if err := s.startWatchdog(); err != nil {
		s.errorf("error starting watchdog %s: %s", s.Name(), err)
	}
	return nil
}

func (s *Service) startIn(d time.Duration) {
	s.stopTimer()
	s.startTimer = time.AfterFunc(d, func() {
		s.Run(nil)
	})
	s.nextStart = time.Now().Add(d)
}

func (s *Service) started(ch *chan<- error) {
	s.State = StateStarted
	s.stopTimer()
	s.sendErr(ch, nil)
	s.retries = 0
}

func (s *Service) startFailed(ch *chan<- error, err error) {
	s.sendErr(ch, err)
	if s.retries < maxRetries-1 {
		s.State = StateBackoff
		duration := time.Second * time.Duration(math.Pow(2, float64(s.retries)))
		s.startIn(duration)
		s.retries++
		s.infof("will retry in %s", duration)
	} else {
		s.State = StateFailed
		s.Cmd = nil
		s.errorf("maximum retries reached")
	}
}

func (s *Service) untilNextRestart() time.Duration {
	return s.nextStart.Sub(time.Now())
}

func (s *Service) Run(ch chan<- error) {
	s.mu.Lock()
	s.State = StateStarted
	s.mu.Unlock()
	for {
		s.mu.Lock()
		if s.State != StateStarted {
			s.mu.Unlock()
			break
		}
		cmd, err := s.Config.Cmd()
		if err != nil {
			s.State = StateFailed
			s.errorf("could not initialize: %s", err)
			s.sendErr(&ch, err)
			s.mu.Unlock()
			break
		}
		if s.Config.Log != nil {
			cmd.Stdout = s.Config.Log.Stdout
			cmd.Stderr = s.Config.Log.Stderr
		}
		s.Cmd = cmd
		s.Started = time.Now()
		s.infof("starting")
		startLock.Lock()
		limits, err := SetLimits(s.Config)
		if err != nil {
			s.errorf("error setting service limits: %s", err)
		}
		serr := s.Cmd.Start()
		if err := RestoreLimits(limits); err != nil {
			s.errorf("error restoring limits: %s", err)
		}
		startLock.Unlock()
		if serr != nil {
			s.errorf("failed to start: %s", serr)
			s.startFailed(&ch, serr)
			s.mu.Unlock()
			break
		}
		s.mu.Unlock()
		timer := time.AfterFunc(time.Duration(float64(minTime)*1.1), func() {
			s.mu.Lock()
			// Clear any potentially stored errors
			s.started(&ch)
			s.mu.Unlock()
			s.infof("started")
		})
		err = s.Cmd.Wait()
		timer.Stop()
		s.mu.Lock()
		if s.State != StateStarted {
			s.Cmd = nil
			s.sendErr(&ch, err)
			s.stopCh <- nil
			s.mu.Unlock()
			break
		}
		if since := time.Since(s.Started); since < minTime {
			// Consider failure
			s.startFailed(&ch, fmt.Errorf("exited too fast (%s)", since))
			s.mu.Unlock()
			break
		}
		s.Restarts++
		s.mu.Unlock()
		if err != nil {
			s.infof("exited with error %s - restarting", err)
		} else {
			s.infof("exited without error - restarting")
		}
	}
}

func (s *Service) Stop() error {
	s.st.Lock()
	defer s.st.Unlock()
	s.stopWatchdog()
	if err := s.stopService(); err != nil {
		s.startWatchdog()
		return err
	}
	return nil
}

func (s *Service) startWatchdog() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.Config.Watchdog != nil {
		interval := s.Config.WatchdogInterval
		if interval < 0 {
			interval = defaultWatchdogInterval
		}
		return s.Config.Watchdog.Start(s, interval)
	}
	return nil
}

func (s *Service) stopWatchdog() {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.Config.Watchdog != nil {
		s.Config.Watchdog.Stop()
	}
}

func (s *Service) stopTimerLocking() {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.stopTimer()
}

func (s *Service) stopTimer() {
	if s.startTimer != nil {
		s.startTimer.Stop()
		s.startTimer = nil
		s.nextStart = time.Time{}
	}
}

func (s *Service) startService() error {
	ch := make(chan error)
	s.stopTimerLocking()
	go s.Run(ch)
	err := <-ch
	close(ch)
	return err
}

func (s *Service) stopService() error {
	s.stopTimer()
	s.mu.Lock()
	if !s.State.isRunState() {
		if s.State.canStop() {
			s.infof("stopped")
		}
		s.State = StateStopped
		s.mu.Unlock()
		return nil
	}
	prevState := s.State
	s.State = StateStopping
	s.infof("stopping")
	p := s.Cmd.Process
	s.mu.Unlock()
	if s != nil {
		stopped := isStoppedErr(p.Signal(os.Signal(syscall.SIGTERM)))
		if !stopped {
			select {
			case <-s.stopCh:
				stopped = true
			case <-time.After(10 * time.Second):
				stopped = isStoppedErr(p.Kill())
			}
			if !stopped {
				select {
				case <-s.stopCh:
				case <-time.After(2 * time.Second):
					// sending signal 0 checks that the process is
					// alive and we're allowed to send the signal
					// without actually sending anything
					if isStoppedErr(p.Signal(syscall.Signal(0))) {
						break
					}
					s.mu.Lock()
					s.State = prevState
					s.mu.Unlock()
					err := fmt.Errorf("could not stop, probably stuck")
					s.errorf("%v", err)
					return err
				}
			}
		}
	}
	s.mu.Lock()
	s.State = StateStopped
	s.Restarts = 0
	s.mu.Unlock()
	if s.Config.Log != nil {
		s.Config.Log.Close()
	}
	s.infof("stopped")
	return nil
}

func (s *Service) log(level log.LLevel, prefix string, format string, args ...interface{}) {
	msg := fmt.Sprintf(format, args...)
	log.Logf(level, "[%s] %s", s.Name(), msg)
	if s.Config.Log != nil {
		s.Config.Log.WriteString(prefix, msg)
	}
}

func (s *Service) errorf(format string, args ...interface{}) {
	s.log(log.LError, "error", format, args...)
}

func (s *Service) infof(format string, args ...interface{}) {
	s.log(log.LInfo, "info", format, args...)
}

func (s *Service) debugf(format string, args ...interface{}) {
	s.log(log.LDebug, "debug", format, args...)
}

func (s *Service) updateConfig(cfg *Config) {
	if reflect.DeepEqual(s.Config, cfg) {
		// there were changes to the file which don't affect the conf
		return
	}
	log.Debugf("changed service %s's configuration", s.Name())
	start := false
	if s.State == StateStarted {
		start = s.Stop() == nil
	}
	s.Config = cfg
	if start {
		s.Start()
	}
}

func isStoppedErr(err error) bool {
	return err != nil && strings.Contains(err.Error(), "process already finished")
}
<commit_msg>Refactor code to allow releasing locks with defer<commit_after>package main

import (
	"fmt"
	"gnd.la/log"
	"math"
	"os"
	"os/exec"
	"reflect"
	"strings"
	"sync"
	"syscall"
	"time"
)

var (
	// avoids starting multiple services
	// at the same time, to enforce resource limits
	startLock sync.Mutex
)

type State uint8

const (
	StateStopped State = iota
	StateStopping
	StateStarted
	StateStarting
	StateBackoff
	StateFailed
)

func (s State) isRunState() bool {
	return s == StateStarted || s == StateStarting
}

func (s State) canStop() bool {
	return s.isRunState() || s == StateBackoff
}

const (
	minTime    = time.Second
	maxRetries = 10
)

type Service struct {
	mu         sync.Mutex // protects access to the fields
	st         sync.Mutex // prevents start/stop from running concurrently
	Config     *Config
	Cmd        *exec.Cmd
	State      State
	Started    time.Time
	Restarts   int
	Err        error
	stopCh     chan error
	errCh      chan error
	retries    int
	startTimer *time.Timer
	nextStart  time.Time
}

func newService(cfg *Config) *Service {
	return &Service{Config: cfg, stopCh: make(chan error)}
}

func (s *Service) Name() string {
	return s.Config.ServiceName()
}

func (s *Service) sendErr(ch *chan<- error, err error) {
	if err != nil && s.State != StateStopping {
		s.errorf("%v", err)
	}
	s.Err = err
	if ch != nil && *ch != nil {
		select {
		case *ch <- err:
			*ch = nil
		default:
		}
	}
	if s.errCh != nil {
		select {
		case s.errCh <- err:
		default:
		}
	}
}

func (s *Service) Start() error {
	s.st.Lock()
	defer s.st.Unlock()
	if s.State.isRunState() {
		return nil
	}
	if err := s.Config.Log.Open(); err != nil {
		return err
	}
	if err := s.startService(); err != nil {
		return err
	}
	if err := s.startWatchdog(); err != nil {
		s.errorf("error starting watchdog %s: %s", s.Name(), err)
	}
	return nil
}

func (s *Service) startIn(d time.Duration) {
	s.stopTimer()
	s.startTimer = time.AfterFunc(d, func() {
		s.Run(nil)
	})
	s.nextStart = time.Now().Add(d)
}

func (s *Service) started(ch *chan<- error) {
	s.State = StateStarted
	s.stopTimer()
	s.sendErr(ch, nil)
	s.retries = 0
}

func (s *Service) startFailed(ch *chan<- error, err error) {
	s.sendErr(ch, err)
	if s.retries < maxRetries-1 {
		s.State = StateBackoff
		duration := time.Second * time.Duration(math.Pow(2, float64(s.retries)))
		s.startIn(duration)
		s.retries++
		s.infof("will retry in %s", duration)
	} else {
		s.State = StateFailed
		s.Cmd = nil
		s.errorf("maximum retries reached")
	}
}

func (s *Service) untilNextRestart() time.Duration {
	return s.nextStart.Sub(time.Now())
}

func (s *Service) Run(ch chan<- error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.State = StateStarted
	if s.State != StateStarted {
		s.sendErr(&ch, nil)
		return
	}
	cmd, err := s.Config.Cmd()
	if err != nil {
		s.State = StateFailed
		s.errorf("could not initialize: %s", err)
		s.sendErr(&ch, err)
		return
	}
	if s.Config.Log != nil {
		cmd.Stdout = s.Config.Log.Stdout
		cmd.Stderr = s.Config.Log.Stderr
	}
	s.Cmd = cmd
	s.Started = time.Now()
	s.infof("starting")
	startLock.Lock()
	limits, err := SetLimits(s.Config)
	if err != nil {
		s.errorf("error setting service limits: %s", err)
	}
	startLock.Unlock()
	serr := s.Cmd.Start()
	if err := RestoreLimits(limits); err != nil {
		s.errorf("error restoring limits: %s", err)
	}
	if serr != nil {
		s.errorf("failed to start: %s", serr)
		s.startFailed(&ch, serr)
		return
	}
	var timer *time.Timer
	timer = time.AfterFunc(minTime, func() {
		s.mu.Lock()
		defer s.mu.Unlock()
		if timer == nil {
			return
		}
		timer = nil
		// Clear any potentially stored errors
		s.started(&ch)
		s.infof("started")
	})

	go func() {
		err := s.Cmd.Wait()
		s.mu.Lock()
		defer s.mu.Unlock()
		if timer != nil {
			// Consider failure, no mintime has passed
			timer.Stop()
			timer = nil
			since := time.Since(s.Started)
			s.startFailed(&ch, fmt.Errorf("exited too fast (%s)", since))
			return
		}
		if s.State != StateStarted {
			s.Cmd = nil
			s.sendErr(&ch, err)
			s.stopCh <- nil
			return
		}
		s.Restarts++
		if err != nil {
			s.infof("exited with error %s - restarting", err)
		} else {
			s.infof("exited without error - restarting")
		}
		// Spawn a goroutine so this function ends and
		// the lock is released before Run() is executed
		// again.
		go s.Run(ch)
	}()
}

func (s *Service) Stop() error {
	s.st.Lock()
	defer s.st.Unlock()
	s.stopWatchdog()
	if err := s.stopService(); err != nil {
		s.startWatchdog()
		return err
	}
	return nil
}

func (s *Service) startWatchdog() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.Config.Watchdog != nil {
		interval := s.Config.WatchdogInterval
		if interval < 0 {
			interval = defaultWatchdogInterval
		}
		return s.Config.Watchdog.Start(s, interval)
	}
	return nil
}

func (s *Service) stopWatchdog() {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.Config.Watchdog != nil {
		s.Config.Watchdog.Stop()
	}
}

func (s *Service) stopTimerLocking() {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.stopTimer()
}

func (s *Service) stopTimer() {
	if s.startTimer != nil {
		s.startTimer.Stop()
		s.startTimer = nil
		s.nextStart = time.Time{}
	}
}

func (s *Service) startService() error {
	ch := make(chan error)
	s.stopTimerLocking()
	s.Run(ch)
	err := <-ch
	close(ch)
	return err
}

func (s *Service) stopService() error {
	s.stopTimer()
	s.mu.Lock()
	if !s.State.isRunState() {
		if s.State.canStop() {
			s.infof("stopped")
		}
		s.State = StateStopped
		s.mu.Unlock()
		return nil
	}
	prevState := s.State
	s.State = StateStopping
	s.infof("stopping")
	p := s.Cmd.Process
	s.mu.Unlock()
	if s != nil {
		stopped := isStoppedErr(p.Signal(os.Signal(syscall.SIGTERM)))
		if !stopped {
			select {
			case <-s.stopCh:
				stopped = true
			case <-time.After(10 * time.Second):
				stopped = isStoppedErr(p.Kill())
			}
			if !stopped {
				select {
				case <-s.stopCh:
				case <-time.After(2 * time.Second):
					// sending signal 0 checks that the process is
					// alive and we're allowed to send the signal
					// without actually sending anything
					if isStoppedErr(p.Signal(syscall.Signal(0))) {
						break
					}
					s.mu.Lock()
					s.State = prevState
					s.mu.Unlock()
					err := fmt.Errorf("could not stop, probably stuck")
					s.errorf("%v", err)
					return err
				}
			}
		}
	}
	s.mu.Lock()
	s.State = StateStopped
	s.Restarts = 0
	s.mu.Unlock()
	if s.Config.Log != nil {
		s.Config.Log.Close()
	}
	s.infof("stopped")
	return nil
}

func (s *Service) log(level log.LLevel, prefix string, format string, args ...interface{}) {
	msg := fmt.Sprintf(format, args...)
	log.Logf(level, "[%s] %s", s.Name(), msg)
	if s.Config.Log != nil {
		s.Config.Log.WriteString(prefix, msg)
	}
}

func (s *Service) errorf(format string, args ...interface{}) {
	s.log(log.LError, "error", format, args...)
}

func (s *Service) infof(format string, args ...interface{}) {
	s.log(log.LInfo, "info", format, args...)
}

func (s *Service) debugf(format string, args ...interface{}) {
	s.log(log.LDebug, "debug", format, args...)
}

func (s *Service) updateConfig(cfg *Config) {
	if reflect.DeepEqual(s.Config, cfg) {
		// there were changes to the file which don't affect the conf
		return
	}
	log.Debugf("changed service %s's configuration", s.Name())
	start := false
	if s.State == StateStarted {
		start = s.Stop() == nil
	}
	s.Config = cfg
	if start {
		s.Start()
	}
}

func isStoppedErr(err error) bool {
	return err != nil && strings.Contains(err.Error(), "process already finished")
}
<|endoftext|>


<commit_before>package pooly

import (
	"github.com/cactus/go-statsd-client/statsd"
	"log"
	"runtime"
	"sync"
	"time"
)

// Computer describes the interface responsible of computing the resulting score of a host.
// It takes the initial score and returns one computed using a predefined function (e.g exp, log ...)
type Computer interface {
	Compute(float64) float64
}

// Selecter describes the interface responsible of selecting a host among the ones registered in the service.
type Selecter interface {
	Select(map[string]*Host) *Host
}

// ServiceConfig defines the service configuration options.
type ServiceConfig struct {
	PoolConfig

	// Number of connections to prespawn on hosts additions (DefaultPrespawnConns by default).
	PrespawnConns uint

	// Maximum attempts to get a connection from the service before giving up (DefaultMaxAttempts by default).
	MaxAttempts uint

	// Deadline after which pools are forced closed (see Pool.ForceClose) (DefaultCloseDeadline by default).
	CloseDeadline time.Duration

	// Defines the time interval taken into account when scores are computed (DefaultDecayDuration by default).
	// Scores are calculated using a weighted average over the course of this duration (recent feedbacks get higher weight).
	DecayDuration time.Duration

	// Time interval between two successive hosts scores computations.
	// Each score is calculated and cached for this duration (DefaultMemoizeScoreDuration by default).
	MemoizeScoreDuration time.Duration

	// Optional score calculator (none by default).
	ScoreCalculator Computer

	// Multi-armed bandit strategy used for host selection (RoundRobin by default).
	// The tradeoff faced by the service at each GetConn is between "exploitation" (choose hosts having the highest score)
	// and "exploration" (find about the expected score of other hosts).
	// The key here is to find the right balance between "exploration" and "exploitation" of hosts given their respective score.
	// Some strategies will favor fairness while others will prefer to pick hosts based on how well they perform.
	BanditStrategy Selecter

	// Address and port of a statsd server to collect and aggregate pooly service metrics (none by default).
	StatsdAddr string
}

// Service manages several hosts, every one of them having a connection pool (see Pool).
// It computes periodically hosts scores and learns about the best alternatives according to the BanditStrategy option.
// Hosts are added or removed from the service via Add and Remove respectively.
// The application calls the GetConn method to get a connection and releases it through the Conn.Release interface.
// When one is done with the pool, Close will cleanup all the service resources.
type Service struct {
	*ServiceConfig

	sync.RWMutex
	name    string
	hosts   map[string]*Host
	decay   *time.Ticker
	memoize *time.Ticker
	add, rm chan string
	stop    chan struct{}
	stats   statsd.Statter
}

// NewService creates a new service given a unique name.
// If no configuration is specified (nil), defaults values are used.
func NewService(name string, c *ServiceConfig) *Service {
	var err error

	if c == nil {
		c = new(ServiceConfig)
	}
	if c.PrespawnConns == 0 {
		c.PrespawnConns = DefaultPrespawnConns
	}
	if c.MaxAttempts == 0 {
		c.MaxAttempts = DefaultMaxAttempts
	}
	if c.CloseDeadline == 0 {
		c.CloseDeadline = DefaultCloseDeadline
	}
	if c.DecayDuration == 0 {
		c.DecayDuration = DefaultDecayDuration
	}
	if c.MemoizeScoreDuration == 0 {
		c.MemoizeScoreDuration = DefaultMemoizeScoreDuration
	}
	if c.BanditStrategy == nil {
		c.BanditStrategy = NewRoundRobin()
	}

	s := &Service{
		ServiceConfig: c,
		name:          name,
		hosts:         make(map[string]*Host),
		add:           make(chan string),
		rm:            make(chan string),
		stop:          make(chan struct{}),
	}
	if _, ok := s.BanditStrategy.(*RoundRobin); !ok {
		s.decay = time.NewTicker(c.DecayDuration / seriesNum)
		s.memoize = time.NewTicker(c.MemoizeScoreDuration)
	}
	if c.StatsdAddr != "" {
		s.stats, err = statsd.New(c.StatsdAddr, "service."+name)
		log.Println("pooly:", err)
	}
	if s.stats == nil {
		s.stats, _ = statsd.NewNoop()
	} else {
		runtime.SetFinalizer(s.stats, func(s statsd.Statter) { s.Close() })
		s.stats.Gauge("conns.count", 0, sampleRate)
		s.stats.Gauge("hosts.count", 0, sampleRate)
		s.stats.Gauge("hosts.score.last", 0, sampleRate)

		go s.monitor()
	}

	go s.serve()
	return s
}

func (s *Service) monitor() {
	t := time.NewTicker(1 * time.Second)
	for {
		select {
		case <-t.C:
			var n int64
			for _, c := range s.Status() {
				n += int64(c)
			}
			s.stats.Gauge("conns.count", n, sampleRate)
		case <-s.stop:
			t.Stop()
			return
		}
	}
}

func (s *Service) serve() {
	var decay, memoize <-chan time.Time

	if s.decay != nil {
		decay = s.decay.C
	}
	if s.memoize != nil {
		memoize = s.memoize.C
	}
	for {
		select {
		case a := <-s.add:
			s.newHost(a)
		case a := <-s.rm:
			s.deleteHost(a)
		case <-decay:
			for _, h := range s.hosts {
				h.decay()
			}
		case <-memoize:
			// XXX lock to prevent selecting hosts during scores computation
			//s.Lock()
			for _, h := range s.hosts {
				h.computeScore(s.ScoreCalculator)
			}
			//c.Unlock()
		case <-s.stop:
			for a := range s.hosts {
				s.deleteHost(a)
			}
			if s.decay != nil {
				s.decay.Stop()
			}
			if s.memoize != nil {
				s.memoize.Stop()
			}
			return
		}
	}
}

func (s *Service) newHost(a string) {
	s.Lock()
	if h := s.hosts[a]; h != nil {
		s.Unlock()
		return
	}

	p := NewPool(a, &s.PoolConfig)
	p.setStats(s.stats)

	p.New(s.PrespawnConns)
	s.hosts[a] = &Host{
		pool:       p,
		timeSeries: make([]serie, 1, seriesNum),
		score:      -1,
		stats:      s.stats,
	}
	s.Unlock()
	s.stats.GaugeDelta("hosts.count", 1, sampleRate)
}

func (s *Service) deleteHost(a string) {
	s.Lock()
	h := s.hosts[a]
	delete(s.hosts, a)
	s.Unlock()

	if h == nil {
		return
	}
	go func() {
		time.AfterFunc(s.CloseDeadline, func() {
			h.pool.ForceClose()
		})
		h.pool.Close()
	}()
	s.stats.GaugeDelta("hosts.count", -1, sampleRate)
}

// Name returns the name of the service.
func (s *Service) Name() string {
	return s.name
}

// Add adds a given host to the service.
// The effect of such operation may not be reflected immediately.
func (s *Service) Add(address string) {
	s.add <- address
}

// Remove removes a given host from the service.
// The effect of such operation may not be reflected immediately.
func (s *Service) Remove(address string) {
	s.rm <- address
}

// GetConn returns a connection from the service.
// The host serving the connection is chosen according to the BanditStrategy policy in place.
func (s *Service) GetConn() (*Conn, error) {
	var attempts uint

	start := time.Now()
again:
	s.RLock()
	if len(s.hosts) == 0 {
		s.RUnlock()
		if attempts < s.MaxAttempts {
			attempts++
			goto again
		}
		return nil, ErrNoHostAvailable
	}
	h := s.BanditStrategy.Select(s.hosts)
	s.RUnlock()

	c, err := h.pool.Get()
	if err != nil {
		// Pool is closed or timed out, demote the host and start over
		s.stats.Inc("conns.get.fails", 1, sampleRate)
		h.rate(HostDown)
		if attempts < s.MaxAttempts {
			attempts++
			goto again
		}
		return nil, err
	}

	// Send statsd metrics
	end := time.Now()
	dt := int64(end.Sub(start).Seconds() * 1000)
	s.stats.Timing("conns.get.delay", dt, sampleRate)
	s.stats.Inc("conns.get.count", 1, sampleRate)
	if _, ok := s.BanditStrategy.(*RoundRobin); !ok {
		// TODO hosts.score.avg
		p := int64(h.Score() * 100)
		s.stats.Gauge("hosts.score.last", p, sampleRate)
	}

	c.setTime(end)
	c.setHost(h)
	return c, nil
}

// Status returns every host addresses managed by the service along with
// the number of connections handled by their respective pool thus far.
func (s *Service) Status() map[string]int32 {
	s.RLock()
	m := make(map[string]int32, len(s.hosts))
	for a, h := range s.hosts {
		m[a] = h.pool.ActiveConns()
	}
	s.RUnlock()
	return m
}

// Close closes the service, thus destroying all hosts and their respective pool.
// After a call to Close, the service can not be used again.
func (s *Service) Close() {
	close(s.stop)
}
<commit_msg>Move hosts metric to monitor<commit_after>package pooly

import (
	"github.com/cactus/go-statsd-client/statsd"
	"log"
	"runtime"
	"sync"
	"time"
)

// Computer describes the interface responsible of computing the resulting score of a host.
// It takes the initial score and returns one computed using a predefined function (e.g exp, log ...)
type Computer interface {
	Compute(float64) float64
}

// Selecter describes the interface responsible of selecting a host among the ones registered in the service.
type Selecter interface {
	Select(map[string]*Host) *Host
}

// ServiceConfig defines the service configuration options.
type ServiceConfig struct {
	PoolConfig

	// Number of connections to prespawn on hosts additions (DefaultPrespawnConns by default).
	PrespawnConns uint

	// Maximum attempts to get a connection from the service before giving up (DefaultMaxAttempts by default).
	MaxAttempts uint

	// Deadline after which pools are forced closed (see Pool.ForceClose) (DefaultCloseDeadline by default).
	CloseDeadline time.Duration

	// Defines the time interval taken into account when scores are computed (DefaultDecayDuration by default).
	// Scores are calculated using a weighted average over the course of this duration (recent feedbacks get higher weight).
	DecayDuration time.Duration

	// Time interval between two successive hosts scores computations.
	// Each score is calculated and cached for this duration (DefaultMemoizeScoreDuration by default).
	MemoizeScoreDuration time.Duration

	// Optional score calculator (none by default).
	ScoreCalculator Computer

	// Multi-armed bandit strategy used for host selection (RoundRobin by default).
	// The tradeoff faced by the service at each GetConn is between "exploitation" (choose hosts having the highest score)
	// and "exploration" (find about the expected score of other hosts).
	// The key here is to find the right balance between "exploration" and "exploitation" of hosts given their respective score.
	// Some strategies will favor fairness while others will prefer to pick hosts based on how well they perform.
	BanditStrategy Selecter

	// Address and port of a statsd server to collect and aggregate pooly service metrics (none by default).
	StatsdAddr string
}

// Service manages several hosts, every one of them having a connection pool (see Pool).
// It computes periodically hosts scores and learns about the best alternatives according to the BanditStrategy option.
// Hosts are added or removed from the service via Add and Remove respectively.
// The application calls the GetConn method to get a connection and releases it through the Conn.Release interface.
// When one is done with the pool, Close will cleanup all the service resources.
type Service struct {
	*ServiceConfig

	sync.RWMutex
	name    string
	hosts   map[string]*Host
	decay   *time.Ticker
	memoize *time.Ticker
	add, rm chan string
	stop    chan struct{}
	stats   statsd.Statter
}

// NewService creates a new service given a unique name.
// If no configuration is specified (nil), defaults values are used.
func NewService(name string, c *ServiceConfig) *Service {
	var err error

	if c == nil {
		c = new(ServiceConfig)
	}
	if c.PrespawnConns == 0 {
		c.PrespawnConns = DefaultPrespawnConns
	}
	if c.MaxAttempts == 0 {
		c.MaxAttempts = DefaultMaxAttempts
	}
	if c.CloseDeadline == 0 {
		c.CloseDeadline = DefaultCloseDeadline
	}
	if c.DecayDuration == 0 {
		c.DecayDuration = DefaultDecayDuration
	}
	if c.MemoizeScoreDuration == 0 {
		c.MemoizeScoreDuration = DefaultMemoizeScoreDuration
	}
	if c.BanditStrategy == nil {
		c.BanditStrategy = NewRoundRobin()
	}

	s := &Service{
		ServiceConfig: c,
		name:          name,
		hosts:         make(map[string]*Host),
		add:           make(chan string),
		rm:            make(chan string),
		stop:          make(chan struct{}),
	}
	if _, ok := s.BanditStrategy.(*RoundRobin); !ok {
		s.decay = time.NewTicker(c.DecayDuration / seriesNum)
		s.memoize = time.NewTicker(c.MemoizeScoreDuration)
	}
	if c.StatsdAddr != "" {
		s.stats, err = statsd.New(c.StatsdAddr, "service."+name)
		log.Println("pooly:", err)
	}
	if s.stats == nil {
		s.stats, _ = statsd.NewNoop()
	} else {
		runtime.SetFinalizer(s.stats, func(s statsd.Statter) { s.Close() })
		s.stats.Gauge("conns.count", 0, sampleRate)
		s.stats.Gauge("hosts.count", 0, sampleRate)
		s.stats.Gauge("hosts.score.last", 0, sampleRate)

		go s.monitor()
	}

	go s.serve()
	return s
}

func (s *Service) monitor() {
	t := time.NewTicker(1 * time.Second)
	for {
		select {
		case <-t.C:
			status := s.Status()
			n := int64(len(status))
			s.stats.Gauge("hosts.count", n, sampleRate)
			n = 0
			for _, c := range status {
				n += int64(c)
			}
			s.stats.Gauge("conns.count", n, sampleRate)

		case <-s.stop:
			t.Stop()
			return
		}
	}
}

func (s *Service) serve() {
	var decay, memoize <-chan time.Time

	if s.decay != nil {
		decay = s.decay.C
	}
	if s.memoize != nil {
		memoize = s.memoize.C
	}
	for {
		select {
		case a := <-s.add:
			s.newHost(a)
		case a := <-s.rm:
			s.deleteHost(a)
		case <-decay:
			for _, h := range s.hosts {
				h.decay()
			}
		case <-memoize:
			// XXX lock to prevent selecting hosts during scores computation
			//s.Lock()
			for _, h := range s.hosts {
				h.computeScore(s.ScoreCalculator)
			}
			//c.Unlock()
		case <-s.stop:
			for a := range s.hosts {
				s.deleteHost(a)
			}
			if s.decay != nil {
				s.decay.Stop()
			}
			if s.memoize != nil {
				s.memoize.Stop()
			}
			return
		}
	}
}

func (s *Service) newHost(a string) {
	s.Lock()
	if h := s.hosts[a]; h != nil {
		s.Unlock()
		return
	}

	p := NewPool(a, &s.PoolConfig)
	p.setStats(s.stats)

	p.New(s.PrespawnConns)
	s.hosts[a] = &Host{
		pool:       p,
		timeSeries: make([]serie, 1, seriesNum),
		score:      -1,
		stats:      s.stats,
	}
	s.Unlock()
}

func (s *Service) deleteHost(a string) {
	s.Lock()
	h := s.hosts[a]
	delete(s.hosts, a)
	s.Unlock()

	if h == nil {
		return
	}
	go func() {
		time.AfterFunc(s.CloseDeadline, func() {
			h.pool.ForceClose()
		})
		h.pool.Close()
	}()
}

// Name returns the name of the service.
func (s *Service) Name() string {
	return s.name
}

// Add adds a given host to the service.
// The effect of such operation may not be reflected immediately.
func (s *Service) Add(address string) {
	s.add <- address
}

// Remove removes a given host from the service.
// The effect of such operation may not be reflected immediately.
func (s *Service) Remove(address string) {
	s.rm <- address
}

// GetConn returns a connection from the service.
// The host serving the connection is chosen according to the BanditStrategy policy in place.
func (s *Service) GetConn() (*Conn, error) {
	var attempts uint

	start := time.Now()
again:
	s.RLock()
	if len(s.hosts) == 0 {
		s.RUnlock()
		if attempts < s.MaxAttempts {
			attempts++
			goto again
		}
		return nil, ErrNoHostAvailable
	}
	h := s.BanditStrategy.Select(s.hosts)
	s.RUnlock()

	c, err := h.pool.Get()
	if err != nil {
		// Pool is closed or timed out, demote the host and start over
		s.stats.Inc("conns.get.fails", 1, sampleRate)
		h.rate(HostDown)
		if attempts < s.MaxAttempts {
			attempts++
			goto again
		}
		return nil, err
	}

	// Send statsd metrics
	end := time.Now()
	dt := int64(end.Sub(start).Seconds() * 1000)
	s.stats.Timing("conns.get.delay", dt, sampleRate)
	s.stats.Inc("conns.get.count", 1, sampleRate)
	if _, ok := s.BanditStrategy.(*RoundRobin); !ok {
		// TODO hosts.score.avg
		p := int64(h.Score() * 100)
		s.stats.Gauge("hosts.score.last", p, sampleRate)
	}

	c.setTime(end)
	c.setHost(h)
	return c, nil
}

// Status returns every host addresses managed by the service along with
// the number of connections handled by their respective pool thus far.
func (s *Service) Status() map[string]int32 {
	s.RLock()
	m := make(map[string]int32, len(s.hosts))
	for a, h := range s.hosts {
		m[a] = h.pool.ActiveConns()
	}
	s.RUnlock()
	return m
}

// Close closes the service, thus destroying all hosts and their respective pool.
// After a call to Close, the service can not be used again.
func (s *Service) Close() {
	close(s.stop)
}
<|endoftext|>


<commit_before>// Copyright 2017 ibelie, Chen Jie, Joungtao. All rights reserved.
// Use of this source code is governed by The MIT License
// that can be found in the LICENSE file.

package tygo

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"path"
	"sort"
	"strconv"
	"strings"

	"io/ioutil"
)

var (
	PY_WRITER  io.Writer
	PY_TYPES   map[string]bool
	PY_OBJECTS map[string]*Object
)

const PY_HEADER = `#-*- coding: utf-8 -*-
# Generated by ibelie-rpc.  DO NOT EDIT!
`

func Python(dir string, name string, types []Type) {
	var buffer bytes.Buffer
	buffer.Write([]byte(PY_HEADER))
	buffer.Write([]byte(`
import typy
`))

	PY_WRITER = &buffer
	PY_TYPES = make(map[string]bool)
	var codes []string
	for _, t := range types {
		codes = append(codes, t.Python())
	}
	PY_TYPES = nil
	PY_WRITER = nil

	buffer.Write([]byte(strings.Join(codes, "")))
	ioutil.WriteFile(path.Join(dir, name+".py"), buffer.Bytes(), 0666)
}

func (t *Enum) Python() string {
	if ok, exist := PY_TYPES[t.Name]; exist && ok {
		return ""
	}

	PY_TYPES[t.Name] = true
	var enums []string
	for _, name := range t.Sorted() {
		enums = append(enums, fmt.Sprintf(`
	%s = %d, "%s"`, name, t.Values[name], name))
	}

	return fmt.Sprintf(`
class %s(typy.Enum):%s
`, t.Name, strings.Join(enums, ""))
}

func (t *Object) Python() string {
	if ok, exist := PY_TYPES[t.Name]; exist && ok {
		return ""
	}

	PY_TYPES[t.Name] = true
	var fields []string
	var sequences []string
	for _, f := range t.VisibleFields() {
		sequences = append(sequences, fmt.Sprintf("'%s'", f.Name))
		fields = append(fields, fmt.Sprintf(`
	%s = %s`, f.Name, strings.Replace(f.Python(), "typy.", "typy.pb.", 1)))
	}

	var sequence string
	if sequences == nil {
		return ""
	} else if len(sequences) > 1 {
		sequence = fmt.Sprintf(`
	____propertySequence__ = %s`, strings.Join(sequences, ", "))
	}

	parent := "typy.Object"
	if t.HasParent() {
		parent = t.Parent.Name
	}

	return fmt.Sprintf(`
class %s(%s):%s%s
`, t.Name, parent, sequence, strings.Join(fields, ""))
}

func (t UnknownType) Python() string {
	return ""
}

func (t SimpleType) Python() string {
	switch t {
	case SimpleType_INT32:
		fallthrough
	case SimpleType_INT64:
		fallthrough
	case SimpleType_UINT32:
		fallthrough
	case SimpleType_UINT64:
		return "typy.Integer"
	case SimpleType_FLOAT32:
		fallthrough
	case SimpleType_FLOAT64:
		return "typy.Float"
	case SimpleType_BYTES:
		return "typy.Bytes"
	case SimpleType_STRING:
		return "typy.String"
	case SimpleType_BOOL:
		return "typy.Boolean"
	default:
		log.Fatalf("[Tygo][SimpleType] Unexpect enum value for Python: %d", t)
		return "Unknown"
	}
}

func (t *EnumType) Python() string {
	PY_WRITER.Write([]byte(t.Enum.Python()))
	return fmt.Sprintf("typy.Enum(%s)", t.Name)
}

func (t *InstanceType) Python() string {
	if object, ok := TS_OBJECTS[t.Name]; ok {
		PY_WRITER.Write([]byte(object.Python()))
		return fmt.Sprintf("typy.Instance(%s)", t.Name)
	} else {
		return fmt.Sprintf("typy.Python(%s)", t.Name)
	}
}

func (t *FixedPointType) Python() string {
	return fmt.Sprintf("typy.FixedPoint(%d, %d)", t.Precision, t.Floor)
}

func (t *ListType) Python() string {
	return fmt.Sprintf("typy.List(%s)", t.E.Python())
}

func (t *DictType) Python() string {
	return fmt.Sprintf("typy.Dict(%s, %s)", t.K.Python(), t.V.Python())
}

func (t *VariantType) Python() string {
	var variants []string
	for _, v := range t.Ts {
		if inst, ok := v.(*InstanceType); ok {
			if object, ok := TS_OBJECTS[inst.Name]; ok {
				PY_WRITER.Write([]byte(object.Python()))
				variants = append(variants, inst.Name)
				continue
			}
		}
		variants = append(variants, v.Python())
	}
	return fmt.Sprintf("typy.Instance(%s)", strings.Join(variants, ", "))
}

const (
	TYPYD_ENUM       = 0
	TYPYD_INT32      = 1
	TYPYD_INT64      = 2
	TYPYD_UINT32     = 3
	TYPYD_UINT64     = 4
	TYPYD_FIXEDPOINT = 5
	TYPYD_DOUBLE     = 6
	TYPYD_FLOAT      = 7
	TYPYD_BOOL       = 8
	TYPYD_BYTES      = 9
	TYPYD_STRING     = 10
	TYPYD_OBJECT     = 11
	TYPYD_VARIANT    = 12
	TYPYD_LIST       = 13
	TYPYD_DICT       = 14
	TYPYD_PYTHON     = 15
	MAX_TYPYD_TYPE   = 16
)

func Typyd(dir string, name string, types []Type) {
	var buffer bytes.Buffer
	buffer.Write([]byte(PY_HEADER))
	buffer.Write([]byte(`
from typy import _typyd
`))

	PY_WRITER = &buffer
	PY_OBJECTS = ObjectMap(types)
	PY_TYPES = make(map[string]bool)
	var codes []string
	for _, t := range types {
		codes = append(codes, t.Typyd())
	}
	PY_TYPES = nil
	PY_WRITER = nil
	PY_OBJECTS = nil

	buffer.Write([]byte(strings.Join(codes, "")))
	ioutil.WriteFile(path.Join(dir, name+".py"), buffer.Bytes(), 0666)
}

func (t *Enum) Typyd() string {
	if ok, exist := PY_TYPES[t.Name]; exist && ok {
		return ""
	} else {
		PY_TYPES[t.Name] = true
		return fmt.Sprintf(`
%s = _typyd.Enum('%s')`, t.Name, t.Name)
	}
}

func (t *Object) Typyd() string {
	if ok, exist := PY_TYPES[t.Name]; exist && ok {
		return ""
	}

	PY_TYPES[t.Name] = true
	var fields []string
	for i, f := range t.AllFields(PY_OBJECTS, true) {
		wiretype := f.WireType()
		fields = append(fields, fmt.Sprintf(`
	('%s', %d, %d, %d, %s),`, f.Name, _MAKE_TAG(i+1, wiretype), TAG_SIZE(i+1), wiretype, f.Typyd()))
	}

	if len(fields) <= 0 {
		return ""
	}

	return fmt.Sprintf(`
%s = _typyd.Object('%s', (%s
))`, t.Name, t.Name, strings.Join(fields, ""))
}

func (t UnknownType) Typyd() string {
	return ""
}

func (t SimpleType) Typyd() string {
	switch t {
	case SimpleType_INT32:
		return strconv.Itoa(TYPYD_INT32)
	case SimpleType_INT64:
		return strconv.Itoa(TYPYD_INT64)
	case SimpleType_UINT32:
		return strconv.Itoa(TYPYD_UINT32)
	case SimpleType_UINT64:
		return strconv.Itoa(TYPYD_UINT64)
	case SimpleType_FLOAT32:
		return strconv.Itoa(TYPYD_FLOAT)
	case SimpleType_FLOAT64:
		return strconv.Itoa(TYPYD_DOUBLE)
	case SimpleType_BYTES:
		return strconv.Itoa(TYPYD_BYTES)
	case SimpleType_STRING:
		return strconv.Itoa(TYPYD_STRING)
	case SimpleType_BOOL:
		return strconv.Itoa(TYPYD_BOOL)
	default:
		log.Fatalf("[Tygo][SimpleType] Unexpect enum value for Typyd: %d", t)
		return "unknown"
	}
}

func (t *EnumType) Typyd() string {
	PY_WRITER.Write([]byte(t.Enum.Typyd()))
	return fmt.Sprintf("%d, %s", TYPYD_ENUM, t.Name)
}

func (t *InstanceType) Typyd() string {
	if _, ok := TS_OBJECTS[t.Name]; ok {
		return t.Name
	} else {
		return "Type"
	}
}

func (t *FixedPointType) Typyd() string {
	identifier := t.Identifier()
	if ok, exist := PY_TYPES[identifier]; !exist || !ok {
		PY_WRITER.Write([]byte(fmt.Sprintf(`
%s = _typyd.FixedPoint(%d, %d)`, identifier, t.Floor, t.Precision)))
		PY_TYPES[identifier] = true
	}
	return fmt.Sprintf("%d, %s", TYPYD_FIXEDPOINT, identifier)
}

func (t *ListType) Typyd() string {
	identifier := t.Identifier()
	if ok, exist := PY_TYPES[identifier]; !exist || !ok {
		PY_WRITER.Write([]byte(fmt.Sprintf(`
%s = _typyd.List('%s', (%d, %s))`, identifier, identifier, t.E.WireType(), t.E.Typyd())))
		PY_TYPES[identifier] = true
	}
	return fmt.Sprintf("%d, %s", TYPYD_LIST, identifier)
}

func (t *DictType) Typyd() string {
	identifier := t.Identifier()
	if ok, exist := PY_TYPES[identifier]; !exist || !ok {
		PY_WRITER.Write([]byte(fmt.Sprintf(`
%s = _typyd.Dict('%s', (%d, %s), (%d, %s))`, identifier, identifier,
			t.K.WireType(), t.K.Typyd(), t.V.WireType(), t.V.Typyd())))
		PY_TYPES[identifier] = true
	}
	return fmt.Sprintf("%d, %s", TYPYD_DICT, identifier)
}

func (t *VariantType) Typyd() string {
	identifier := t.Identifier()
	if ok, exist := PY_TYPES[identifier]; !exist || !ok {
		properties := make(map[string]Type)
		for _, typ := range t.Ts {
			switch v := typ.(type) {
			case SimpleType:
				switch v {
				case SimpleType_INT32:
					fallthrough
				case SimpleType_INT64:
					fallthrough
				case SimpleType_UINT32:
					fallthrough
				case SimpleType_UINT64:
					properties["Integer"] = v
				case SimpleType_FLOAT32:
					properties["Float"] = v
				case SimpleType_FLOAT64:
					properties["Double"] = v
				case SimpleType_BYTES:
					properties["Bytes"] = v
				case SimpleType_STRING:
					properties["String"] = v
				case SimpleType_BOOL:
					properties["Boolean"] = v
				default:
					log.Fatalf("[Tygo][VariantType] Unexpect enum value for SimpleType: %d", v)
				}
			case *EnumType:
				properties["Enum"] = v
			case *InstanceType:
				properties[v.Name] = v
			case *FixedPointType:
				properties["Dict"] = v
			case *ListType:
				properties["Dict"] = v
			case *DictType:
				properties["Dict"] = v
			default:
				log.Fatalf("[Tygo][VariantType] Unexpect type for Typyd: %v", v)
			}
		}

		var sortedProperties []string
		for n, _ := range properties {
			sortedProperties = append(sortedProperties, n)
		}
		sort.Strings(sortedProperties)

		var fields []string
		for i, n := range sortedProperties {
			wiretype := properties[n].WireType()
			fields = append(fields, fmt.Sprintf(`
	('%s', %d, %d, %d, %s),`, n, _MAKE_TAG(i+1, wiretype), TAG_SIZE(i+1), wiretype, properties[n].Typyd()))
		}

		PY_WRITER.Write([]byte(fmt.Sprintf(`
%s = _typyd.Variant('%s', (%s
))`, identifier, identifier, strings.Join(fields, ""))))
		PY_TYPES[identifier] = true
	}
	return fmt.Sprintf("%d, %s", TYPYD_VARIANT, identifier)
}
<commit_msg>change interface<commit_after>// Copyright 2017 ibelie, Chen Jie, Joungtao. All rights reserved.
// Use of this source code is governed by The MIT License
// that can be found in the LICENSE file.

package tygo

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"sort"
	"strconv"
	"strings"
)

var (
	PY_WRITER  io.Writer
	PY_TYPES   map[string]bool
	PY_OBJECTS map[string]*Object
)

func Python(types []Type) []byte {
	var buffer bytes.Buffer
	buffer.Write([]byte(`
import typy
`))

	PY_WRITER = &buffer
	PY_TYPES = make(map[string]bool)
	var codes []string
	for _, t := range types {
		codes = append(codes, t.Python())
	}
	PY_TYPES = nil
	PY_WRITER = nil

	buffer.Write([]byte(strings.Join(codes, "")))
	return buffer.Bytes()
}

func (t *Enum) Python() string {
	if ok, exist := PY_TYPES[t.Name]; exist && ok {
		return ""
	}

	PY_TYPES[t.Name] = true
	var enums []string
	for _, name := range t.Sorted() {
		enums = append(enums, fmt.Sprintf(`
	%s = %d, "%s"`, name, t.Values[name], name))
	}

	return fmt.Sprintf(`
class %s(typy.Enum):%s
`, t.Name, strings.Join(enums, ""))
}

func (t *Object) Python() string {
	if ok, exist := PY_TYPES[t.Name]; exist && ok {
		return ""
	}

	PY_TYPES[t.Name] = true
	var fields []string
	var sequences []string
	for _, f := range t.VisibleFields() {
		sequences = append(sequences, fmt.Sprintf("'%s'", f.Name))
		fields = append(fields, fmt.Sprintf(`
	%s = %s`, f.Name, strings.Replace(f.Python(), "typy.", "typy.pb.", 1)))
	}

	var sequence string
	if sequences == nil {
		return ""
	} else if len(sequences) > 1 {
		sequence = fmt.Sprintf(`
	____propertySequence__ = %s`, strings.Join(sequences, ", "))
	}

	parent := "typy.Object"
	if t.HasParent() {
		parent = t.Parent.Name
	}

	return fmt.Sprintf(`
class %s(%s):%s%s
`, t.Name, parent, sequence, strings.Join(fields, ""))
}

func (t UnknownType) Python() string {
	return ""
}

func (t SimpleType) Python() string {
	switch t {
	case SimpleType_INT32:
		fallthrough
	case SimpleType_INT64:
		fallthrough
	case SimpleType_UINT32:
		fallthrough
	case SimpleType_UINT64:
		return "typy.Integer"
	case SimpleType_FLOAT32:
		fallthrough
	case SimpleType_FLOAT64:
		return "typy.Float"
	case SimpleType_BYTES:
		return "typy.Bytes"
	case SimpleType_STRING:
		return "typy.String"
	case SimpleType_BOOL:
		return "typy.Boolean"
	default:
		log.Fatalf("[Tygo][SimpleType] Unexpect enum value for Python: %d", t)
		return "Unknown"
	}
}

func (t *EnumType) Python() string {
	PY_WRITER.Write([]byte(t.Enum.Python()))
	return fmt.Sprintf("typy.Enum(%s)", t.Name)
}

func (t *InstanceType) Python() string {
	if object, ok := TS_OBJECTS[t.Name]; ok {
		PY_WRITER.Write([]byte(object.Python()))
		return fmt.Sprintf("typy.Instance(%s)", t.Name)
	} else {
		return fmt.Sprintf("typy.Python(%s)", t.Name)
	}
}

func (t *FixedPointType) Python() string {
	return fmt.Sprintf("typy.FixedPoint(%d, %d)", t.Precision, t.Floor)
}

func (t *ListType) Python() string {
	return fmt.Sprintf("typy.List(%s)", t.E.Python())
}

func (t *DictType) Python() string {
	return fmt.Sprintf("typy.Dict(%s, %s)", t.K.Python(), t.V.Python())
}

func (t *VariantType) Python() string {
	var variants []string
	for _, v := range t.Ts {
		if inst, ok := v.(*InstanceType); ok {
			if object, ok := TS_OBJECTS[inst.Name]; ok {
				PY_WRITER.Write([]byte(object.Python()))
				variants = append(variants, inst.Name)
				continue
			}
		}
		variants = append(variants, v.Python())
	}
	return fmt.Sprintf("typy.Instance(%s)", strings.Join(variants, ", "))
}

const (
	TYPYD_ENUM       = 0
	TYPYD_INT32      = 1
	TYPYD_INT64      = 2
	TYPYD_UINT32     = 3
	TYPYD_UINT64     = 4
	TYPYD_FIXEDPOINT = 5
	TYPYD_DOUBLE     = 6
	TYPYD_FLOAT      = 7
	TYPYD_BOOL       = 8
	TYPYD_BYTES      = 9
	TYPYD_STRING     = 10
	TYPYD_OBJECT     = 11
	TYPYD_VARIANT    = 12
	TYPYD_LIST       = 13
	TYPYD_DICT       = 14
	TYPYD_PYTHON     = 15
	MAX_TYPYD_TYPE   = 16
)

func Typyd(types []Type) []byte {
	var buffer bytes.Buffer
	buffer.Write([]byte(`
from typy import _typyd
`))

	PY_WRITER = &buffer
	PY_OBJECTS = ObjectMap(types)
	PY_TYPES = make(map[string]bool)
	var codes []string
	for _, t := range types {
		codes = append(codes, t.Typyd())
	}
	PY_TYPES = nil
	PY_WRITER = nil
	PY_OBJECTS = nil

	buffer.Write([]byte(strings.Join(codes, "")))
	return buffer.Bytes()
}

func (t *Enum) Typyd() string {
	if ok, exist := PY_TYPES[t.Name]; exist && ok {
		return ""
	} else {
		PY_TYPES[t.Name] = true
		return fmt.Sprintf(`
%s = _typyd.Enum('%s')`, t.Name, t.Name)
	}
}

func (t *Object) Typyd() string {
	if ok, exist := PY_TYPES[t.Name]; exist && ok {
		return ""
	}

	PY_TYPES[t.Name] = true
	var fields []string
	for i, f := range t.AllFields(PY_OBJECTS, true) {
		wiretype := f.WireType()
		fields = append(fields, fmt.Sprintf(`
	('%s', %d, %d, %d, %s),`, f.Name, _MAKE_TAG(i+1, wiretype), TAG_SIZE(i+1), wiretype, f.Typyd()))
	}

	if len(fields) <= 0 {
		return ""
	}

	return fmt.Sprintf(`
%s = _typyd.Object('%s', (%s
))`, t.Name, t.Name, strings.Join(fields, ""))
}

func (t UnknownType) Typyd() string {
	return ""
}

func (t SimpleType) Typyd() string {
	switch t {
	case SimpleType_INT32:
		return strconv.Itoa(TYPYD_INT32)
	case SimpleType_INT64:
		return strconv.Itoa(TYPYD_INT64)
	case SimpleType_UINT32:
		return strconv.Itoa(TYPYD_UINT32)
	case SimpleType_UINT64:
		return strconv.Itoa(TYPYD_UINT64)
	case SimpleType_FLOAT32:
		return strconv.Itoa(TYPYD_FLOAT)
	case SimpleType_FLOAT64:
		return strconv.Itoa(TYPYD_DOUBLE)
	case SimpleType_BYTES:
		return strconv.Itoa(TYPYD_BYTES)
	case SimpleType_STRING:
		return strconv.Itoa(TYPYD_STRING)
	case SimpleType_BOOL:
		return strconv.Itoa(TYPYD_BOOL)
	default:
		log.Fatalf("[Tygo][SimpleType] Unexpect enum value for Typyd: %d", t)
		return "unknown"
	}
}

func (t *EnumType) Typyd() string {
	PY_WRITER.Write([]byte(t.Enum.Typyd()))
	return fmt.Sprintf("%d, %s", TYPYD_ENUM, t.Name)
}

func (t *InstanceType) Typyd() string {
	if _, ok := TS_OBJECTS[t.Name]; ok {
		return t.Name
	} else {
		return "Type"
	}
}

func (t *FixedPointType) Typyd() string {
	identifier := t.Identifier()
	if ok, exist := PY_TYPES[identifier]; !exist || !ok {
		PY_WRITER.Write([]byte(fmt.Sprintf(`
%s = _typyd.FixedPoint(%d, %d)`, identifier, t.Floor, t.Precision)))
		PY_TYPES[identifier] = true
	}
	return fmt.Sprintf("%d, %s", TYPYD_FIXEDPOINT, identifier)
}

func (t *ListType) Typyd() string {
	identifier := t.Identifier()
	if ok, exist := PY_TYPES[identifier]; !exist || !ok {
		PY_WRITER.Write([]byte(fmt.Sprintf(`
%s = _typyd.List('%s', (%d, %s))`, identifier, identifier, t.E.WireType(), t.E.Typyd())))
		PY_TYPES[identifier] = true
	}
	return fmt.Sprintf("%d, %s", TYPYD_LIST, identifier)
}

func (t *DictType) Typyd() string {
	identifier := t.Identifier()
	if ok, exist := PY_TYPES[identifier]; !exist || !ok {
		PY_WRITER.Write([]byte(fmt.Sprintf(`
%s = _typyd.Dict('%s', (%d, %s), (%d, %s))`, identifier, identifier,
			t.K.WireType(), t.K.Typyd(), t.V.WireType(), t.V.Typyd())))
		PY_TYPES[identifier] = true
	}
	return fmt.Sprintf("%d, %s", TYPYD_DICT, identifier)
}

func (t *VariantType) Typyd() string {
	identifier := t.Identifier()
	if ok, exist := PY_TYPES[identifier]; !exist || !ok {
		properties := make(map[string]Type)
		for _, typ := range t.Ts {
			switch v := typ.(type) {
			case SimpleType:
				switch v {
				case SimpleType_INT32:
					fallthrough
				case SimpleType_INT64:
					fallthrough
				case SimpleType_UINT32:
					fallthrough
				case SimpleType_UINT64:
					properties["Integer"] = v
				case SimpleType_FLOAT32:
					properties["Float"] = v
				case SimpleType_FLOAT64:
					properties["Double"] = v
				case SimpleType_BYTES:
					properties["Bytes"] = v
				case SimpleType_STRING:
					properties["String"] = v
				case SimpleType_BOOL:
					properties["Boolean"] = v
				default:
					log.Fatalf("[Tygo][VariantType] Unexpect enum value for SimpleType: %d", v)
				}
			case *EnumType:
				properties["Enum"] = v
			case *InstanceType:
				properties[v.Name] = v
			case *FixedPointType:
				properties["Dict"] = v
			case *ListType:
				properties["Dict"] = v
			case *DictType:
				properties["Dict"] = v
			default:
				log.Fatalf("[Tygo][VariantType] Unexpect type for Typyd: %v", v)
			}
		}

		var sortedProperties []string
		for n, _ := range properties {
			sortedProperties = append(sortedProperties, n)
		}
		sort.Strings(sortedProperties)

		var fields []string
		for i, n := range sortedProperties {
			wiretype := properties[n].WireType()
			fields = append(fields, fmt.Sprintf(`
	('%s', %d, %d, %d, %s),`, n, _MAKE_TAG(i+1, wiretype), TAG_SIZE(i+1), wiretype, properties[n].Typyd()))
		}

		PY_WRITER.Write([]byte(fmt.Sprintf(`
%s = _typyd.Variant('%s', (%s
))`, identifier, identifier, strings.Join(fields, ""))))
		PY_TYPES[identifier] = true
	}
	return fmt.Sprintf("%d, %s", TYPYD_VARIANT, identifier)
}
<|endoftext|>


<commit_before>package nds

import (
	"encoding/binary"
	"errors"
	"math/rand"
	"reflect"
	"time"

	"appengine/datastore"
	"appengine/memcache"
)

const (
	// memcachePrefix is the namespace memcache uses to store entities.
	memcachePrefix = "NDS0:"

	// memcacheLockTime is the maximum length of time a memcache lock will be
	// held for. 32 seconds is choosen as 30 seconds is the maximum amount of
	// time an underlying datastore call will retry even if the API reports a
	// success to the user.
	memcacheLockTime = 32 * time.Second
)

var (
	typeOfPropertyLoadSaver = reflect.TypeOf(
		(*datastore.PropertyLoadSaver)(nil)).Elem()
	typeOfPropertyList = reflect.TypeOf(datastore.PropertyList(nil))

	memcacheAddMulti            = memcache.AddMulti
	memcacheCompareAndSwapMulti = memcache.CompareAndSwapMulti
	memcacheGetMulti            = memcache.GetMulti
	memcacheDeleteMulti         = memcache.DeleteMulti
	memcacheSetMulti            = memcache.SetMulti

	datastoreDeleteMulti = datastore.DeleteMulti
	datastoreGetMulti    = datastore.GetMulti
	datastorePutMulti    = datastore.PutMulti
)

const (
	noneItem uint32 = iota
	entityItem
	lockItem
)

func itemLock() []byte {
	b := make([]byte, 4)
	binary.LittleEndian.PutUint32(b, rand.Uint32())
	return b
}

func checkArgs(key *datastore.Key, val interface{}) error {
	if key == nil {
		return errors.New("nds: key is nil")
	}

	if val == nil {
		return errors.New("nds: val is nil")
	}

	v := reflect.ValueOf(val)
	if v.Type() == typeOfPropertyList {
		return errors.New("nds: PropertyList not supported")
	}

	switch v.Kind() {
	case reflect.Ptr:
		return nil
	default:
		return errors.New("nds: val must be a slice or pointer")
	}

}

func checkMultiArgs(keys []*datastore.Key, v reflect.Value) error {
	if v.Kind() != reflect.Slice {
		return errors.New("nds: vals is not a slice")
	}

	if len(keys) != v.Len() {
		return errors.New("nds: keys and vals slices have different length")
	}

	if v.Type() == typeOfPropertyList {
		return errors.New("nds: PropertyList not supported")
	}

	elemType := v.Type().Elem()
	if reflect.PtrTo(elemType).Implements(typeOfPropertyLoadSaver) {
		return nil
	}

	switch elemType.Kind() {
	case reflect.Struct, reflect.Interface:
		return nil
	case reflect.Ptr:
		elemType = elemType.Elem()
		if elemType.Kind() == reflect.Struct {
			return nil
		}
	}
	return errors.New("nds: unsupported vals type")
}

func createMemcacheKey(key *datastore.Key) string {
	return memcachePrefix + key.Encode()
}

// SaveStruct saves src to a datastore.PropertyList. src must be a struct
// pointer.
func SaveStruct(src interface{}, pl *datastore.PropertyList) error {
	c, err := make(chan datastore.Property), make(chan error)
	go func() {
		err <- datastore.SaveStruct(src, c)
	}()
	for p := range c {
		*pl = append(*pl, p)
	}
	return <-err
}

// LoadStruct loads a datastore.PropertyList into dst. src must be a struct
// pointer.
func LoadStruct(dst interface{}, pl datastore.PropertyList) error {
	c := make(chan datastore.Property)
	go func() {
		for _, p := range pl {
			c <- p
		}
		close(c)
	}()
	return datastore.LoadStruct(dst, c)
}

func propertyLoadSaverToPropertyList(
	pls datastore.PropertyLoadSaver, pl *datastore.PropertyList) error {
	c, err := make(chan datastore.Property), make(chan error)
	go func() {
		err <- pls.Save(c)
	}()
	for p := range c {
		*pl = append(*pl, p)
	}
	return <-err
}

func propertyListToPropertyLoadSaver(
	pl datastore.PropertyList, pls datastore.PropertyLoadSaver) error {

	c := make(chan datastore.Property)
	go func() {
		for _, p := range pl {
			c <- p
		}
		close(c)
	}()

	return pls.Load(c)
}
<commit_msg>Bumped memcache namespace up.<commit_after>package nds

import (
	"encoding/binary"
	"errors"
	"math/rand"
	"reflect"
	"time"

	"appengine/datastore"
	"appengine/memcache"
)

const (
	// memcachePrefix is the namespace memcache uses to store entities.
	memcachePrefix = "NDS1:"

	// memcacheLockTime is the maximum length of time a memcache lock will be
	// held for. 32 seconds is choosen as 30 seconds is the maximum amount of
	// time an underlying datastore call will retry even if the API reports a
	// success to the user.
	memcacheLockTime = 32 * time.Second
)

var (
	typeOfPropertyLoadSaver = reflect.TypeOf(
		(*datastore.PropertyLoadSaver)(nil)).Elem()
	typeOfPropertyList = reflect.TypeOf(datastore.PropertyList(nil))

	memcacheAddMulti            = memcache.AddMulti
	memcacheCompareAndSwapMulti = memcache.CompareAndSwapMulti
	memcacheGetMulti            = memcache.GetMulti
	memcacheDeleteMulti         = memcache.DeleteMulti
	memcacheSetMulti            = memcache.SetMulti

	datastoreDeleteMulti = datastore.DeleteMulti
	datastoreGetMulti    = datastore.GetMulti
	datastorePutMulti    = datastore.PutMulti
)

const (
	noneItem uint32 = iota
	entityItem
	lockItem
)

func itemLock() []byte {
	b := make([]byte, 4)
	binary.LittleEndian.PutUint32(b, rand.Uint32())
	return b
}

func checkArgs(key *datastore.Key, val interface{}) error {
	if key == nil {
		return errors.New("nds: key is nil")
	}

	if val == nil {
		return errors.New("nds: val is nil")
	}

	v := reflect.ValueOf(val)
	if v.Type() == typeOfPropertyList {
		return errors.New("nds: PropertyList not supported")
	}

	switch v.Kind() {
	case reflect.Ptr:
		return nil
	default:
		return errors.New("nds: val must be a slice or pointer")
	}

}

func checkMultiArgs(keys []*datastore.Key, v reflect.Value) error {
	if v.Kind() != reflect.Slice {
		return errors.New("nds: vals is not a slice")
	}

	if len(keys) != v.Len() {
		return errors.New("nds: keys and vals slices have different length")
	}

	if v.Type() == typeOfPropertyList {
		return errors.New("nds: PropertyList not supported")
	}

	elemType := v.Type().Elem()
	if reflect.PtrTo(elemType).Implements(typeOfPropertyLoadSaver) {
		return nil
	}

	switch elemType.Kind() {
	case reflect.Struct, reflect.Interface:
		return nil
	case reflect.Ptr:
		elemType = elemType.Elem()
		if elemType.Kind() == reflect.Struct {
			return nil
		}
	}
	return errors.New("nds: unsupported vals type")
}

func createMemcacheKey(key *datastore.Key) string {
	return memcachePrefix + key.Encode()
}

// SaveStruct saves src to a datastore.PropertyList. src must be a struct
// pointer.
func SaveStruct(src interface{}, pl *datastore.PropertyList) error {
	c, err := make(chan datastore.Property), make(chan error)
	go func() {
		err <- datastore.SaveStruct(src, c)
	}()
	for p := range c {
		*pl = append(*pl, p)
	}
	return <-err
}

// LoadStruct loads a datastore.PropertyList into dst. src must be a struct
// pointer.
func LoadStruct(dst interface{}, pl datastore.PropertyList) error {
	c := make(chan datastore.Property)
	go func() {
		for _, p := range pl {
			c <- p
		}
		close(c)
	}()
	return datastore.LoadStruct(dst, c)
}

func propertyLoadSaverToPropertyList(
	pls datastore.PropertyLoadSaver, pl *datastore.PropertyList) error {
	c, err := make(chan datastore.Property), make(chan error)
	go func() {
		err <- pls.Save(c)
	}()
	for p := range c {
		*pl = append(*pl, p)
	}
	return <-err
}

func propertyListToPropertyLoadSaver(
	pl datastore.PropertyList, pls datastore.PropertyLoadSaver) error {

	c := make(chan datastore.Property)
	go func() {
		for _, p := range pl {
			c <- p
		}
		close(c)
	}()

	return pls.Load(c)
}
<|endoftext|>


<commit_before>package bruxism

import (
	"fmt"
	"io"
	"log"
	"regexp"

	"github.com/iopred/discordgo"
)

// The number of guilds supported by one shard.
const numGuildsPerShard = 2400

// DiscordServiceName is the service name for the Discord service.
const DiscordServiceName string = "Discord"

// DiscordMessage is a Message wrapper around discordgo.Message.
type DiscordMessage struct {
	Discord          *Discord
	DiscordgoMessage *discordgo.Message
	MessageType      MessageType
	Nick             *string
	Content          *string
}

// Channel returns the channel id for this message.
func (m *DiscordMessage) Channel() string {
	return m.DiscordgoMessage.ChannelID
}

// UserName returns the user name for this message.
func (m *DiscordMessage) UserName() string {
	me := m.DiscordgoMessage
	if me.Author == nil {
		return ""
	}

	if m.Nick == nil {
		n := m.Discord.NicknameForID(me.Author.ID, me.Author.Username, me.ChannelID)
		m.Nick = &n
	}
	return *m.Nick
}

// UserID returns the user id for this message.
func (m *DiscordMessage) UserID() string {
	if m.DiscordgoMessage.Author == nil {
		return ""
	}

	return m.DiscordgoMessage.Author.ID
}

// UserAvatar returns the avatar url for this message.
func (m *DiscordMessage) UserAvatar() string {
	if m.DiscordgoMessage.Author == nil {
		return ""
	}

	return discordgo.EndpointUserAvatar(m.DiscordgoMessage.Author.ID, m.DiscordgoMessage.Author.Avatar)
}

// Message returns the message content for this message.
func (m *DiscordMessage) Message() string {
	if m.Content == nil {
		c := m.DiscordgoMessage.ContentWithMentionsReplaced()
		c = m.Discord.replaceRoleNames(m.DiscordgoMessage, c)
		c = m.Discord.replaceChannelNames(m.DiscordgoMessage, c)

		m.Content = &c
	}
	return *m.Content
}

// RawMessage returns the raw message content for this message.
func (m *DiscordMessage) RawMessage() string {
	return m.DiscordgoMessage.Content
}

// MessageID returns the message ID for this message.
func (m *DiscordMessage) MessageID() string {
	return m.DiscordgoMessage.ID
}

// Type returns the type of message.
func (m *DiscordMessage) Type() MessageType {
	return m.MessageType
}

// Discord is a Service provider for Discord.
type Discord struct {
	args        []interface{}
	messageChan chan Message

	Shards int

	// The first session, used to send messages (and maintain backwards compatibility).
	Session             *discordgo.Session
	Sessions            []*discordgo.Session
	OwnerUserID         string
	ApplicationClientID string
}

// NewDiscord creates a new discord service.
func NewDiscord(args ...interface{}) *Discord {
	return &Discord{
		args:        args,
		messageChan: make(chan Message, 200),
	}
}

var channelIDRegex = regexp.MustCompile("<#[0-9]*>")

func (d *Discord) replaceChannelNames(message *discordgo.Message, content string) string {
	return channelIDRegex.ReplaceAllStringFunc(content, func(str string) string {
		c, err := d.Channel(str[2 : len(str)-1])
		if err != nil {
			return str
		}

		return "#" + c.Name
	})
}

var roleIDRegex = regexp.MustCompile("<@&[0-9]*>")

func (d *Discord) replaceRoleNames(message *discordgo.Message, content string) string {
	return roleIDRegex.ReplaceAllStringFunc(content, func(str string) string {
		roleID := str[3 : len(str)-1]

		c, err := d.Channel(message.ChannelID)
		if err != nil {
			return str
		}

		g, err := d.Guild(c.GuildID)
		if err != nil {
			return str
		}

		for _, r := range g.Roles {
			if r.ID == roleID {
				return "@" + r.Name
			}
		}

		return str
	})
}

func (d *Discord) onMessageCreate(s *discordgo.Session, message *discordgo.MessageCreate) {
	if message.Content == "" {
		return
	}

	d.messageChan <- &DiscordMessage{
		Discord:          d,
		DiscordgoMessage: message.Message,
		MessageType:      MessageTypeCreate,
	}
}

func (d *Discord) onMessageUpdate(s *discordgo.Session, message *discordgo.MessageUpdate) {
	if message.Content == "" {
		return
	}

	d.messageChan <- &DiscordMessage{
		Discord:          d,
		DiscordgoMessage: message.Message,
		MessageType:      MessageTypeUpdate,
	}
}

func (d *Discord) onMessageDelete(s *discordgo.Session, message *discordgo.MessageDelete) {
	d.messageChan <- &DiscordMessage{
		Discord:          d,
		DiscordgoMessage: message.Message,
		MessageType:      MessageTypeDelete,
	}
}

// Name returns the name of the service.
func (d *Discord) Name() string {
	return DiscordServiceName
}

// Open opens the service and returns a channel which all messages will be sent on.
func (d *Discord) Open() (<-chan Message, error) {
	shards := d.Shards
	if shards < 1 {
		shards = 1
	}

	d.Sessions = make([]*discordgo.Session, shards)

	for i := 0; i < shards; i++ {
		session, err := discordgo.New(d.args...)
		if err != nil {
			return nil, err
		}
		session.ShardCount = shards
		session.ShardID = i
		session.AddHandler(d.onMessageCreate)
		session.AddHandler(d.onMessageUpdate)
		session.AddHandler(d.onMessageDelete)

		d.Sessions[i] = session
	}

	d.Session = d.Sessions[0]

	for i := 0; i < len(d.Sessions); i++ {
		d.Sessions[i].Open()
	}

	return d.messageChan, nil
}

// IsMe returns whether or not a message was sent by the bot.
func (d *Discord) IsMe(message Message) bool {
	if d.Session.State.User == nil {
		return false
	}
	return message.UserID() == d.Session.State.User.ID
}

// SendMessage sends a message.
func (d *Discord) SendMessage(channel, message string) error {
	if channel == "" {
		log.Println("Empty channel could not send message", message)
		return nil
	}

	if _, err := d.Session.ChannelMessageSend(channel, message); err != nil {
		log.Println("Error sending discord message: ", err)
		return err
	}

	return nil
}

// SendAction sends an action.
func (d *Discord) SendAction(channel, message string) error {
	if channel == "" {
		log.Println("Empty channel could not send message", message)
		return nil
	}

	p, err := d.UserChannelPermissions(d.UserID(), channel)
	if err != nil {
		return d.SendMessage(channel, message)
	}

	if p&discordgo.PermissionEmbedLinks == discordgo.PermissionEmbedLinks {
		if _, err := d.Session.ChannelMessageSendEmbed(channel, &discordgo.MessageEmbed{
			Color:       d.UserColor(d.UserID(), channel),
			Description: message,
		}); err != nil {
			return err
		}
		return nil
	}

	return d.SendMessage(channel, message)
}

// DeleteMessage deletes a message.
func (d *Discord) DeleteMessage(channel, messageID string) error {
	return d.Session.ChannelMessageDelete(channel, messageID)
}

// SendFile sends a file.
func (d *Discord) SendFile(channel, name string, r io.Reader) error {
	if _, err := d.Session.ChannelFileSend(channel, name, r); err != nil {
		log.Println("Error sending discord message: ", err)
		return err
	}
	return nil
}

// BanUser bans a user.
func (d *Discord) BanUser(channel, userID string, duration int) error {
	return d.Session.GuildBanCreate(channel, userID, 0)
}

// UnbanUser unbans a user.
func (d *Discord) UnbanUser(channel, userID string) error {
	return d.Session.GuildBanDelete(channel, userID)
}

// UserName returns the bots name.
func (d *Discord) UserName() string {
	if d.Session.State.User == nil {
		return ""
	}
	return d.Session.State.User.Username
}

// UserID returns the bots user id.
func (d *Discord) UserID() string {
	if d.Session.State.User == nil {
		return ""
	}
	return d.Session.State.User.ID
}

// Join accept an invite or return an error.
// If AlreadyJoinedError is return, @me has already accepted that invite.
func (d *Discord) Join(join string) error {
	if i, err := d.Session.Invite(join); err == nil {
		if _, err := d.Guild(i.Guild.ID); err == nil {
			return ErrAlreadyJoined
		}
	}

	if _, err := d.Session.InviteAccept(join); err != nil {
		return err
	}
	return nil
}

// Typing sets that the bot is typing.
func (d *Discord) Typing(channel string) error {
	return d.Session.ChannelTyping(channel)
}

// PrivateMessage will send a private message to a user.
func (d *Discord) PrivateMessage(userID, message string) error {
	c, err := d.Session.UserChannelCreate(userID)
	if err != nil {
		return err
	}
	return d.SendMessage(c.ID, message)
}

// SupportsPrivateMessages returns whether the service supports private messages.
func (d *Discord) SupportsPrivateMessages() bool {
	return true
}

// SupportsMultiline returns whether the service supports multiline messages.
func (d *Discord) SupportsMultiline() bool {
	return true
}

// CommandPrefix returns the command prefix for the service.
func (d *Discord) CommandPrefix() string {
	return fmt.Sprintf("@%s ", d.UserName())
}

// IsBotOwner returns whether or not a message sender was the owner of the bot.
func (d *Discord) IsBotOwner(message Message) bool {
	return message.UserID() == d.OwnerUserID
}

// IsPrivate returns whether or not a message was private.
func (d *Discord) IsPrivate(message Message) bool {
	c, err := d.Channel(message.Channel())
	return err == nil && c.IsPrivate
}

// IsChannelOwner returns whether or not the sender of a message is a moderator.
func (d *Discord) IsChannelOwner(message Message) bool {
	c, err := d.Channel(message.Channel())
	if err != nil {
		return false
	}
	g, err := d.Guild(c.GuildID)
	if err != nil {
		return false
	}
	return g.OwnerID == message.UserID() || d.IsBotOwner(message)
}

// IsModerator returns whether or not the sender of a message is a moderator.
func (d *Discord) IsModerator(message Message) bool {
	p, err := d.UserChannelPermissions(message.UserID(), message.Channel())
	if err == nil {
		if p&discordgo.PermissionAdministrator == discordgo.PermissionAdministrator || p&discordgo.PermissionManageChannels == discordgo.PermissionManageChannels || p&discordgo.PermissionManageServer == discordgo.PermissionManageServer {
			return true
		}
	}

	return d.IsChannelOwner(message)
}

// ChannelCount returns the number of channels the bot is in.
func (d *Discord) ChannelCount() int {
	return len(d.Guilds())
}

// SupportsMessageHistory returns if the service supports message history.
func (d *Discord) SupportsMessageHistory() bool {
	return true
}

// MessageHistory returns the message history for a channel.
func (d *Discord) MessageHistory(channel string) []Message {
	c, err := d.Channel(channel)
	if err != nil {
		return nil
	}

	messages := make([]Message, len(c.Messages))
	for i := 0; i < len(c.Messages); i++ {
		messages[i] = &DiscordMessage{
			Discord:          d,
			DiscordgoMessage: c.Messages[i],
			MessageType:      MessageTypeCreate,
		}
	}

	return messages
}

func (d *Discord) Channel(channelID string) (channel *discordgo.Channel, err error) {
	for _, s := range d.Sessions {
		channel, err = s.State.Channel(channelID)
		if err == nil {
			return channel, nil
		}
	}
	return
}

func (d *Discord) Guild(guildID string) (guild *discordgo.Guild, err error) {
	for _, s := range d.Sessions {
		guild, err = s.State.Guild(guildID)
		if err == nil {
			return guild, nil
		}
	}
	return
}

func (d *Discord) Guilds() []*discordgo.Guild {
	guilds := []*discordgo.Guild{}
	for _, s := range d.Sessions {
		guilds = append(guilds, s.State.Guilds...)
	}
	return guilds
}

func (d *Discord) UserChannelPermissions(userID, channelID string) (apermissions int, err error) {
	for _, s := range d.Sessions {
		apermissions, err = s.State.UserChannelPermissions(userID, channelID)
		if err == nil {
			return apermissions, nil
		}
	}
	return
}

func (d *Discord) UserColor(userID, channelID string) int {
	for _, s := range d.Sessions {
		color := s.State.UserColor(userID, channelID)
		if color != 0 {
			return color
		}
	}
	return 0
}

func (d *Discord) Nickname(message Message) string {
	return d.NicknameForID(message.UserID(), message.UserName(), message.Channel())
}

func (d *Discord) NicknameForID(userID, userName, channelID string) string {
	c, err := d.Channel(channelID)
	if err == nil {
		g, err := d.Guild(c.GuildID)
		if err == nil {
			for _, m := range g.Members {
				if m.User.ID == userID {
					if m.Nick != "" {
						return m.Nick
					}
					break
				}
			}
		}
	}
	return userName
}
<commit_msg>Don't track presences.. This is mainly for Septapus because this kills the Septapus.<commit_after>package bruxism

import (
	"fmt"
	"io"
	"log"
	"regexp"

	"github.com/iopred/discordgo"
)

// The number of guilds supported by one shard.
const numGuildsPerShard = 2400

// DiscordServiceName is the service name for the Discord service.
const DiscordServiceName string = "Discord"

// DiscordMessage is a Message wrapper around discordgo.Message.
type DiscordMessage struct {
	Discord          *Discord
	DiscordgoMessage *discordgo.Message
	MessageType      MessageType
	Nick             *string
	Content          *string
}

// Channel returns the channel id for this message.
func (m *DiscordMessage) Channel() string {
	return m.DiscordgoMessage.ChannelID
}

// UserName returns the user name for this message.
func (m *DiscordMessage) UserName() string {
	me := m.DiscordgoMessage
	if me.Author == nil {
		return ""
	}

	if m.Nick == nil {
		n := m.Discord.NicknameForID(me.Author.ID, me.Author.Username, me.ChannelID)
		m.Nick = &n
	}
	return *m.Nick
}

// UserID returns the user id for this message.
func (m *DiscordMessage) UserID() string {
	if m.DiscordgoMessage.Author == nil {
		return ""
	}

	return m.DiscordgoMessage.Author.ID
}

// UserAvatar returns the avatar url for this message.
func (m *DiscordMessage) UserAvatar() string {
	if m.DiscordgoMessage.Author == nil {
		return ""
	}

	return discordgo.EndpointUserAvatar(m.DiscordgoMessage.Author.ID, m.DiscordgoMessage.Author.Avatar)
}

// Message returns the message content for this message.
func (m *DiscordMessage) Message() string {
	if m.Content == nil {
		c := m.DiscordgoMessage.ContentWithMentionsReplaced()
		c = m.Discord.replaceRoleNames(m.DiscordgoMessage, c)
		c = m.Discord.replaceChannelNames(m.DiscordgoMessage, c)

		m.Content = &c
	}
	return *m.Content
}

// RawMessage returns the raw message content for this message.
func (m *DiscordMessage) RawMessage() string {
	return m.DiscordgoMessage.Content
}

// MessageID returns the message ID for this message.
func (m *DiscordMessage) MessageID() string {
	return m.DiscordgoMessage.ID
}

// Type returns the type of message.
func (m *DiscordMessage) Type() MessageType {
	return m.MessageType
}

// Discord is a Service provider for Discord.
type Discord struct {
	args        []interface{}
	messageChan chan Message

	Shards int

	// The first session, used to send messages (and maintain backwards compatibility).
	Session             *discordgo.Session
	Sessions            []*discordgo.Session
	OwnerUserID         string
	ApplicationClientID string
}

// NewDiscord creates a new discord service.
func NewDiscord(args ...interface{}) *Discord {
	return &Discord{
		args:        args,
		messageChan: make(chan Message, 200),
	}
}

var channelIDRegex = regexp.MustCompile("<#[0-9]*>")

func (d *Discord) replaceChannelNames(message *discordgo.Message, content string) string {
	return channelIDRegex.ReplaceAllStringFunc(content, func(str string) string {
		c, err := d.Channel(str[2 : len(str)-1])
		if err != nil {
			return str
		}

		return "#" + c.Name
	})
}

var roleIDRegex = regexp.MustCompile("<@&[0-9]*>")

func (d *Discord) replaceRoleNames(message *discordgo.Message, content string) string {
	return roleIDRegex.ReplaceAllStringFunc(content, func(str string) string {
		roleID := str[3 : len(str)-1]

		c, err := d.Channel(message.ChannelID)
		if err != nil {
			return str
		}

		g, err := d.Guild(c.GuildID)
		if err != nil {
			return str
		}

		for _, r := range g.Roles {
			if r.ID == roleID {
				return "@" + r.Name
			}
		}

		return str
	})
}

func (d *Discord) onMessageCreate(s *discordgo.Session, message *discordgo.MessageCreate) {
	if message.Content == "" {
		return
	}

	d.messageChan <- &DiscordMessage{
		Discord:          d,
		DiscordgoMessage: message.Message,
		MessageType:      MessageTypeCreate,
	}
}

func (d *Discord) onMessageUpdate(s *discordgo.Session, message *discordgo.MessageUpdate) {
	if message.Content == "" {
		return
	}

	d.messageChan <- &DiscordMessage{
		Discord:          d,
		DiscordgoMessage: message.Message,
		MessageType:      MessageTypeUpdate,
	}
}

func (d *Discord) onMessageDelete(s *discordgo.Session, message *discordgo.MessageDelete) {
	d.messageChan <- &DiscordMessage{
		Discord:          d,
		DiscordgoMessage: message.Message,
		MessageType:      MessageTypeDelete,
	}
}

// Name returns the name of the service.
func (d *Discord) Name() string {
	return DiscordServiceName
}

// Open opens the service and returns a channel which all messages will be sent on.
func (d *Discord) Open() (<-chan Message, error) {
	shards := d.Shards
	if shards < 1 {
		shards = 1
	}

	d.Sessions = make([]*discordgo.Session, shards)

	for i := 0; i < shards; i++ {
		session, err := discordgo.New(d.args...)
		if err != nil {
			return nil, err
		}
		session.ShardCount = shards
		session.ShardID = i
		session.AddHandler(d.onMessageCreate)
		session.AddHandler(d.onMessageUpdate)
		session.AddHandler(d.onMessageDelete)
		session.TrackPresences = False

		d.Sessions[i] = session
	}

	d.Session = d.Sessions[0]

	for i := 0; i < len(d.Sessions); i++ {
		d.Sessions[i].Open()
	}

	return d.messageChan, nil
}

// IsMe returns whether or not a message was sent by the bot.
func (d *Discord) IsMe(message Message) bool {
	if d.Session.State.User == nil {
		return false
	}
	return message.UserID() == d.Session.State.User.ID
}

// SendMessage sends a message.
func (d *Discord) SendMessage(channel, message string) error {
	if channel == "" {
		log.Println("Empty channel could not send message", message)
		return nil
	}

	if _, err := d.Session.ChannelMessageSend(channel, message); err != nil {
		log.Println("Error sending discord message: ", err)
		return err
	}

	return nil
}

// SendAction sends an action.
func (d *Discord) SendAction(channel, message string) error {
	if channel == "" {
		log.Println("Empty channel could not send message", message)
		return nil
	}

	p, err := d.UserChannelPermissions(d.UserID(), channel)
	if err != nil {
		return d.SendMessage(channel, message)
	}

	if p&discordgo.PermissionEmbedLinks == discordgo.PermissionEmbedLinks {
		if _, err := d.Session.ChannelMessageSendEmbed(channel, &discordgo.MessageEmbed{
			Color:       d.UserColor(d.UserID(), channel),
			Description: message,
		}); err != nil {
			return err
		}
		return nil
	}

	return d.SendMessage(channel, message)
}

// DeleteMessage deletes a message.
func (d *Discord) DeleteMessage(channel, messageID string) error {
	return d.Session.ChannelMessageDelete(channel, messageID)
}

// SendFile sends a file.
func (d *Discord) SendFile(channel, name string, r io.Reader) error {
	if _, err := d.Session.ChannelFileSend(channel, name, r); err != nil {
		log.Println("Error sending discord message: ", err)
		return err
	}
	return nil
}

// BanUser bans a user.
func (d *Discord) BanUser(channel, userID string, duration int) error {
	return d.Session.GuildBanCreate(channel, userID, 0)
}

// UnbanUser unbans a user.
func (d *Discord) UnbanUser(channel, userID string) error {
	return d.Session.GuildBanDelete(channel, userID)
}

// UserName returns the bots name.
func (d *Discord) UserName() string {
	if d.Session.State.User == nil {
		return ""
	}
	return d.Session.State.User.Username
}

// UserID returns the bots user id.
func (d *Discord) UserID() string {
	if d.Session.State.User == nil {
		return ""
	}
	return d.Session.State.User.ID
}

// Join accept an invite or return an error.
// If AlreadyJoinedError is return, @me has already accepted that invite.
func (d *Discord) Join(join string) error {
	if i, err := d.Session.Invite(join); err == nil {
		if _, err := d.Guild(i.Guild.ID); err == nil {
			return ErrAlreadyJoined
		}
	}

	if _, err := d.Session.InviteAccept(join); err != nil {
		return err
	}
	return nil
}

// Typing sets that the bot is typing.
func (d *Discord) Typing(channel string) error {
	return d.Session.ChannelTyping(channel)
}

// PrivateMessage will send a private message to a user.
func (d *Discord) PrivateMessage(userID, message string) error {
	c, err := d.Session.UserChannelCreate(userID)
	if err != nil {
		return err
	}
	return d.SendMessage(c.ID, message)
}

// SupportsPrivateMessages returns whether the service supports private messages.
func (d *Discord) SupportsPrivateMessages() bool {
	return true
}

// SupportsMultiline returns whether the service supports multiline messages.
func (d *Discord) SupportsMultiline() bool {
	return true
}

// CommandPrefix returns the command prefix for the service.
func (d *Discord) CommandPrefix() string {
	return fmt.Sprintf("@%s ", d.UserName())
}

// IsBotOwner returns whether or not a message sender was the owner of the bot.
func (d *Discord) IsBotOwner(message Message) bool {
	return message.UserID() == d.OwnerUserID
}

// IsPrivate returns whether or not a message was private.
func (d *Discord) IsPrivate(message Message) bool {
	c, err := d.Channel(message.Channel())
	return err == nil && c.IsPrivate
}

// IsChannelOwner returns whether or not the sender of a message is a moderator.
func (d *Discord) IsChannelOwner(message Message) bool {
	c, err := d.Channel(message.Channel())
	if err != nil {
		return false
	}
	g, err := d.Guild(c.GuildID)
	if err != nil {
		return false
	}
	return g.OwnerID == message.UserID() || d.IsBotOwner(message)
}

// IsModerator returns whether or not the sender of a message is a moderator.
func (d *Discord) IsModerator(message Message) bool {
	p, err := d.UserChannelPermissions(message.UserID(), message.Channel())
	if err == nil {
		if p&discordgo.PermissionAdministrator == discordgo.PermissionAdministrator || p&discordgo.PermissionManageChannels == discordgo.PermissionManageChannels || p&discordgo.PermissionManageServer == discordgo.PermissionManageServer {
			return true
		}
	}

	return d.IsChannelOwner(message)
}

// ChannelCount returns the number of channels the bot is in.
func (d *Discord) ChannelCount() int {
	return len(d.Guilds())
}

// SupportsMessageHistory returns if the service supports message history.
func (d *Discord) SupportsMessageHistory() bool {
	return true
}

// MessageHistory returns the message history for a channel.
func (d *Discord) MessageHistory(channel string) []Message {
	c, err := d.Channel(channel)
	if err != nil {
		return nil
	}

	messages := make([]Message, len(c.Messages))
	for i := 0; i < len(c.Messages); i++ {
		messages[i] = &DiscordMessage{
			Discord:          d,
			DiscordgoMessage: c.Messages[i],
			MessageType:      MessageTypeCreate,
		}
	}

	return messages
}

func (d *Discord) Channel(channelID string) (channel *discordgo.Channel, err error) {
	for _, s := range d.Sessions {
		channel, err = s.State.Channel(channelID)
		if err == nil {
			return channel, nil
		}
	}
	return
}

func (d *Discord) Guild(guildID string) (guild *discordgo.Guild, err error) {
	for _, s := range d.Sessions {
		guild, err = s.State.Guild(guildID)
		if err == nil {
			return guild, nil
		}
	}
	return
}

func (d *Discord) Guilds() []*discordgo.Guild {
	guilds := []*discordgo.Guild{}
	for _, s := range d.Sessions {
		guilds = append(guilds, s.State.Guilds...)
	}
	return guilds
}

func (d *Discord) UserChannelPermissions(userID, channelID string) (apermissions int, err error) {
	for _, s := range d.Sessions {
		apermissions, err = s.State.UserChannelPermissions(userID, channelID)
		if err == nil {
			return apermissions, nil
		}
	}
	return
}

func (d *Discord) UserColor(userID, channelID string) int {
	for _, s := range d.Sessions {
		color := s.State.UserColor(userID, channelID)
		if color != 0 {
			return color
		}
	}
	return 0
}

func (d *Discord) Nickname(message Message) string {
	return d.NicknameForID(message.UserID(), message.UserName(), message.Channel())
}

func (d *Discord) NicknameForID(userID, userName, channelID string) string {
	c, err := d.Channel(channelID)
	if err == nil {
		g, err := d.Guild(c.GuildID)
		if err == nil {
			for _, m := range g.Members {
				if m.User.ID == userID {
					if m.Nick != "" {
						return m.Nick
					}
					break
				}
			}
		}
	}
	return userName
}
<|endoftext|>


<commit_before>package installer

import (
	"bitbucket.org/kardianos/osext"

	"bytes"
	"errors"
	"flag"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"text/template"
)

func GetInitDPath(name string) string {
	return "/etc/init.d/" + name
}

func IsInitDScriptExist(name string) bool {
	if _, err := os.Stat(GetInitDPath(name)); os.IsExist(err) {
		return true
	} else {
		return false
	}
}

func GetInitDScript(folder, command, user string) (ret string, err error) {

	initDTemplate, err := template.New("initd").Parse(templateInitD)
	if nil != err {
		return
	}

	buf := new(bytes.Buffer)
	initDTemplate.Execute(buf, TemplateInitD{Dir: folder, User: user, Command: folder + command, ScriptName: command})
	ret = buf.String()

	return
}

func WriteInitD(name, content string) (rErr error) {
	f, err := os.Create(GetInitDPath(name))
	if nil != err {
		rErr = err
		return
	}
	defer f.Close()
	length, err := f.Write([]byte(content))
	if nil != err {
		rErr = err
		return
	}
	if length != len(content) {
		rErr = errors.New("file could not be written completely")
		return
	}
	if err = MakeInitDExecutable(name); nil != err {
		//@TODO: also need rollback here
		rErr = err
		return
	}
	return
}

func MakeInitDExecutable(name string) (rErr error) {
	cmd := exec.Command("chmod", "+x", GetInitDPath(name))
	if _, execErr := cmd.CombinedOutput(); nil != execErr {
		rErr = execErr
	}
	return
}

func Install(user string) (rErr error) {
	pathToExe, err := osext.Executable()
	folder, command := filepath.Split(pathToExe)
	if nil != err {
		panic(err)
	}
	switch runtime.GOOS {
	case "linux":
	default:
		rErr = errors.New("not yet supported")
		return
	}
	var script string
	if IsInitDScriptExist(command) {
		rErr = errors.New(fmt.Sprintf("Script %s already exists in init.d folder", command))
		return
	}
	if script, err = GetInitDScript(folder, command, user); nil != err {
		rErr = err
		return
	}
	if fileWriteError := WriteInitD(command, script); nil != fileWriteError {
		rErr = fileWriteError
		return
	}
	cmd := exec.Command("update-rc.d", command, "defaults")
	if out, execErr := cmd.CombinedOutput(); nil != execErr {
		//@TODO: rollback , delete init.d script
		rErr = execErr
		return
	} else {
		fmt.Println(string(out))
	}
	return
}

func Register() {
	var install = flag.Bool("install", false, "install this program as service")
	var runAsUser = flag.String("installRunAsUser", "", "which user should the service run as")
	flag.Parse()
	if !*install {
		return
	}
	if *install && "" == *runAsUser {
		fmt.Println("you must specify a user that the service runs as")
		os.Exit(1)
	}
	if err := Install(*runAsUser); err != nil {
		fmt.Printf("installing as service failed: <%v> \n", err)
		os.Exit(1)
	}
	fmt.Println("successfully installed as service")
	os.Exit(0)
}

type TemplateInitD struct {
	Dir, User, Command, ScriptName string
}

const templateInitD = `#!/bin/sh
### BEGIN INIT INFO
# Provides: {{.ScriptName}}
# Required-Start: $remote_fs $syslog
# Required-Stop: $remote_fs $syslog
# Default-Start: 2 3 4 5
# Default-Stop: 0 1 6
# Short-Description: Start daemon at boot time
# Description: Enable service provided by daemon.
### END INIT INFO
dir="{{.Dir}}"
user="{{.User}}"
cmd="{{.Command}}"
name=` + "`basename $0`" + `
pid_file="/var/run/$name.pid"
stdout_log="/var/log/$name.log"
stderr_log="/var/log/$name.err"
get_pid() {
cat "$pid_file"
}
is_running() {
[ -f "$pid_file" ] && ps ` + "`get_pid`" + ` > /dev/null 2>&1
}
case "$1" in
start)
if is_running; then
echo "Already started"
else
echo "Starting $name"
cd "$dir"
sudo -u "$user" $cmd >> "$stdout_log" 2>> "$stderr_log" &
echo $! > "$pid_file"
if ! is_running; then
echo "Unable to start, see $stdout_log and $stderr_log"
exit 1
fi
fi
;;
stop)
if is_running; then
echo -n "Stopping $name.."
kill ` + "`get_pid`" + `
for i in {1..10}
do
if ! is_running; then
break
fi
echo -n "."
sleep 1
done
echo
if is_running; then
echo "Not stopped; may still be shutting down or shutdown may have failed"
exit 1
else
echo "Stopped"
if [ -f "$pid_file" ]; then
rm "$pid_file"
fi
fi
else
echo "Not running"
fi
;;
restart)
$0 stop
if is_running; then
echo "Unable to stop, will not attempt to start"
exit 1
fi
$0 start
;;
status)
if is_running; then
echo "Running"
else
echo "Stopped"
exit 1
fi
;;
*)
echo "Usage: $0 {start|stop|restart|status}"
exit 1
;;
esac
exit 0
`
<commit_msg>updated path to kardianos package<commit_after>package installer

import (
	"github.com/kardianos/osext"

	"bytes"
	"errors"
	"flag"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"text/template"
)

func GetInitDPath(name string) string {
	return "/etc/init.d/" + name
}

func IsInitDScriptExist(name string) bool {
	if _, err := os.Stat(GetInitDPath(name)); os.IsExist(err) {
		return true
	} else {
		return false
	}
}

func GetInitDScript(folder, command, user string) (ret string, err error) {

	initDTemplate, err := template.New("initd").Parse(templateInitD)
	if nil != err {
		return
	}

	buf := new(bytes.Buffer)
	initDTemplate.Execute(buf, TemplateInitD{Dir: folder, User: user, Command: folder + command, ScriptName: command})
	ret = buf.String()

	return
}

func WriteInitD(name, content string) (rErr error) {
	f, err := os.Create(GetInitDPath(name))
	if nil != err {
		rErr = err
		return
	}
	defer f.Close()
	length, err := f.Write([]byte(content))
	if nil != err {
		rErr = err
		return
	}
	if length != len(content) {
		rErr = errors.New("file could not be written completely")
		return
	}
	if err = MakeInitDExecutable(name); nil != err {
		//@TODO: also need rollback here
		rErr = err
		return
	}
	return
}

func MakeInitDExecutable(name string) (rErr error) {
	cmd := exec.Command("chmod", "+x", GetInitDPath(name))
	if _, execErr := cmd.CombinedOutput(); nil != execErr {
		rErr = execErr
	}
	return
}

func Install(user string) (rErr error) {
	pathToExe, err := osext.Executable()
	folder, command := filepath.Split(pathToExe)
	if nil != err {
		panic(err)
	}
	switch runtime.GOOS {
	case "linux":
	default:
		rErr = errors.New("not yet supported")
		return
	}
	var script string
	if IsInitDScriptExist(command) {
		rErr = errors.New(fmt.Sprintf("Script %s already exists in init.d folder", command))
		return
	}
	if script, err = GetInitDScript(folder, command, user); nil != err {
		rErr = err
		return
	}
	if fileWriteError := WriteInitD(command, script); nil != fileWriteError {
		rErr = fileWriteError
		return
	}
	cmd := exec.Command("update-rc.d", command, "defaults")
	if out, execErr := cmd.CombinedOutput(); nil != execErr {
		//@TODO: rollback , delete init.d script
		rErr = execErr
		return
	} else {
		fmt.Println(string(out))
	}
	return
}

func Register() {
	var install = flag.Bool("install", false, "install this program as service")
	var runAsUser = flag.String("installRunAsUser", "", "which user should the service run as")
	flag.Parse()
	if !*install {
		return
	}
	if *install && "" == *runAsUser {
		fmt.Println("you must specify a user that the service runs as")
		os.Exit(1)
	}
	if err := Install(*runAsUser); err != nil {
		fmt.Printf("installing as service failed: <%v> \n", err)
		os.Exit(1)
	}
	fmt.Println("successfully installed as service")
	os.Exit(0)
}

type TemplateInitD struct {
	Dir, User, Command, ScriptName string
}

const templateInitD = `#!/bin/sh
### BEGIN INIT INFO
# Provides: {{.ScriptName}}
# Required-Start: $remote_fs $syslog
# Required-Stop: $remote_fs $syslog
# Default-Start: 2 3 4 5
# Default-Stop: 0 1 6
# Short-Description: Start daemon at boot time
# Description: Enable service provided by daemon.
### END INIT INFO
dir="{{.Dir}}"
user="{{.User}}"
cmd="{{.Command}}"
name=` + "`basename $0`" + `
pid_file="/var/run/$name.pid"
stdout_log="/var/log/$name.log"
stderr_log="/var/log/$name.err"
get_pid() {
cat "$pid_file"
}
is_running() {
[ -f "$pid_file" ] && ps ` + "`get_pid`" + ` > /dev/null 2>&1
}
case "$1" in
start)
if is_running; then
echo "Already started"
else
echo "Starting $name"
cd "$dir"
sudo -u "$user" $cmd >> "$stdout_log" 2>> "$stderr_log" &
echo $! > "$pid_file"
if ! is_running; then
echo "Unable to start, see $stdout_log and $stderr_log"
exit 1
fi
fi
;;
stop)
if is_running; then
echo -n "Stopping $name.."
kill ` + "`get_pid`" + `
for i in {1..10}
do
if ! is_running; then
break
fi
echo -n "."
sleep 1
done
echo
if is_running; then
echo "Not stopped; may still be shutting down or shutdown may have failed"
exit 1
else
echo "Stopped"
if [ -f "$pid_file" ]; then
rm "$pid_file"
fi
fi
else
echo "Not running"
fi
;;
restart)
$0 stop
if is_running; then
echo "Unable to stop, will not attempt to start"
exit 1
fi
$0 start
;;
status)
if is_running; then
echo "Running"
else
echo "Stopped"
exit 1
fi
;;
*)
echo "Usage: $0 {start|stop|restart|status}"
exit 1
;;
esac
exit 0
`
<|endoftext|>