<commit_before>package bootstrap

import (
	"fmt"
	"io/ioutil"
	"os"

	"html/template"

	"github.com/tomogoma/authms/api"
	"github.com/tomogoma/authms/config"
	"github.com/tomogoma/authms/db"
	"github.com/tomogoma/authms/facebook"
	"github.com/tomogoma/authms/logging"
	"github.com/tomogoma/authms/model"
	"github.com/tomogoma/authms/sms/africas_talking"
	"github.com/tomogoma/authms/sms/messagebird"
	"github.com/tomogoma/authms/sms/twilio"
	"github.com/tomogoma/authms/smtp"
	"github.com/tomogoma/crdb"
	token "github.com/tomogoma/jwt"
)

func InstantiateRoach(lg logging.Logger, conf crdb.Config) *db.Roach {
	var opts []db.Option
	if dsn := conf.FormatDSN(); dsn != "" {
		opts = append(opts, db.WithDSN(dsn))
	}
	if dbn := conf.DBName; dbn != "" {
		opts = append(opts, db.WithDBName(dbn))
	}
	rdb := db.NewRoach(opts...)
	err := rdb.InitDBIfNot()
	logging.LogWarnOnError(lg, err, "Initiate Cockroach DB connection")
	return rdb
}

func InstantiateJWTHandler(lg logging.Logger, tknKyF string) *token.Handler {
	JWTKey, err := ioutil.ReadFile(tknKyF)
	logging.LogFatalOnError(lg, err, "Read JWT key file")
	jwter, err := token.NewHandler(JWTKey)
	logging.LogFatalOnError(lg, err, "Instantiate JWT handler")
	return jwter
}

func InstantiateFacebook(conf config.Facebook) (*facebook.FacebookOAuth, error) {
	if conf.ID < 1 {
		return nil, nil
	}
	fbSecret, err := readFile(conf.SecretFilePath)
	if err != nil {
		return nil, fmt.Errorf("read facebook secret from file: %v", err)
	}
	fb, err := facebook.New(conf.ID, fbSecret)
	if err != nil {
		return nil, err
	}
	return fb, nil
}

func InstantiateSMSer(conf config.SMS) (model.SMSer, error) {
	if conf.ActiveAPI == "" {
		return nil, nil
	}
	var s model.SMSer
	switch conf.ActiveAPI {
	case config.SMSAPIAfricasTalking:
		apiKey, err := readFile(conf.AfricasTalking.APIKeyFile)
		if err != nil {
			return nil, fmt.Errorf("read africa's talking API key: %v", err)
		}
		s, err = africas_talking.NewSMSCl(conf.AfricasTalking.UserName, apiKey)
		if err != nil {
			return nil, fmt.Errorf("new africasTalking client: %v", err)
		}
	case config.SMSAPITwilio:
		tkn, err := readFile(conf.Twilio.TokenKeyFile)
		if err != nil {
			return nil, fmt.Errorf("read twilio token: %v", err)
		}
		s, err = twilio.NewSMSCl(conf.Twilio.ID, tkn, conf.Twilio.SenderPhone)
		if err != nil {
			return nil, fmt.Errorf("new twilio client: %v", err)
		}
	case config.SMSAPIMessageBird:
		apiKey, err := readFile(conf.MessageBird.APIKeyFile)
		if err != nil {
			return nil, fmt.Errorf("read messageBird API key: %v", err)
		}
		s, err = messagebird.NewClient(conf.MessageBird.AccountName, apiKey)
		if err != nil {
			return nil, fmt.Errorf("new messageBird client: %v", err)
		}
	default:
		return nil, fmt.Errorf("invalid API selected can be %s or %s",
			config.SMSAPIAfricasTalking, config.SMSAPITwilio)
	}
	var testMessage string
	host := hostname()
	testMessage = fmt.Sprintf("The SMS API is being used on %s", host)
	if err := s.SMS(conf.TestNumber, testMessage); err != nil {
		return s, fmt.Errorf("test SMS: %v", err)
	}
	return s, nil
}

func InstantiateSMTP(rdb *db.Roach, lg logging.Logger, conf config.SMTP) *smtp.Mailer {

	emailCl, err := smtp.New(rdb)
	logging.LogFatalOnError(lg, err, "Instantiate email API")

	err = emailCl.Configured()
	if err == nil {
		return emailCl
	}

	if !emailCl.IsNotFoundError(err) {
		logging.LogWarnOnError(lg, err, "Check email API configured")
		return emailCl
	}

	pass, err := readFile(conf.PasswordFile)
	logging.LogWarnOnError(lg, err, "Read SMTP password file")

	host := hostname()
	err = emailCl.SetConfig(
		smtp.Config{
			ServerAddress: conf.ServerAddress,
			TLSPort:       conf.TLSPort,
			SSLPort:       conf.SSLPort,
			Username:      conf.Username,
			FromEmail:     conf.FromEmail,
			Password:      pass,
		},
		model.SendMail{
			ToEmails: []string{conf.TestEmail},
			Subject:  "Authentication Micro-Service Started on " + host,
			Body:     template.HTML("The authentication micro-service is being used on " + host),
		},
	)
	logging.LogWarnOnError(lg, err, "Set default SMTP config")

	return emailCl
}

func Instantiate(confFile string, lg logging.Logger) (config.General, *model.Authentication, *api.Guard, *db.Roach, model.JWTEr, model.SMSer, *smtp.Mailer) {

	conf, err := config.ReadFile(confFile)
	logging.LogFatalOnError(lg, err, "Read config file")

	rdb := InstantiateRoach(lg, conf.Database)
	tg := InstantiateJWTHandler(lg, conf.Token.TokenKeyFile)

	var authOpts []model.Option
	fb, err := InstantiateFacebook(conf.Authentication.Facebook)
	logging.LogWarnOnError(lg, err, "Set up OAuth options")
	if fb != nil {
		authOpts = append(authOpts, model.WithFacebookCl(fb))
	}

	sms, err := InstantiateSMSer(conf.SMS)
	logging.LogWarnOnError(lg, err, "Instantiate SMS API")
	if sms != nil {
		authOpts = append(authOpts, model.WithSMSCl(sms))
	}

	emailCl := InstantiateSMTP(rdb, lg, conf.SMTP)
	authOpts = append(authOpts, model.WithEmailCl(emailCl))

	authOpts = append(
		authOpts,
		model.WithAppName(conf.Service.AppName),
		model.WithWebAppURL(conf.Service.WebAppURL),
		model.WithDevLockedToUser(conf.Authentication.LockDevsToUsers),
		model.WithSelfRegAllowed(conf.Authentication.AllowSelfReg),
	)

	a, err := model.NewAuthentication(rdb, tg, authOpts...)
	logging.LogFatalOnError(lg, err, "Instantiate Auth Model")

	g, err := api.NewGuard(rdb, api.WithMasterKey(conf.Service.MasterAPIKey))
	logging.LogFatalOnError(lg, err, "Instantate API access guard")

	return *conf, a, g, rdb, tg, sms, emailCl
}

func hostname() string {
	hostName, err := os.Hostname()
	if err != nil {
		return "an unknown host"
	}
	return hostName
}

func readFile(path string) (string, error) {
	contentB, err := ioutil.ReadFile(path)
	if err != nil {
		return "", fmt.Errorf("read: %v", err)
	}
	return string(contentB), nil
}
<commit_msg>don't send test sms if test number not provided<commit_after>package bootstrap

import (
	"fmt"
	"io/ioutil"
	"os"

	"html/template"

	"github.com/tomogoma/authms/api"
	"github.com/tomogoma/authms/config"
	"github.com/tomogoma/authms/db"
	"github.com/tomogoma/authms/facebook"
	"github.com/tomogoma/authms/logging"
	"github.com/tomogoma/authms/model"
	"github.com/tomogoma/authms/sms/africas_talking"
	"github.com/tomogoma/authms/sms/messagebird"
	"github.com/tomogoma/authms/sms/twilio"
	"github.com/tomogoma/authms/smtp"
	"github.com/tomogoma/crdb"
	token "github.com/tomogoma/jwt"
)

func InstantiateRoach(lg logging.Logger, conf crdb.Config) *db.Roach {
	var opts []db.Option
	if dsn := conf.FormatDSN(); dsn != "" {
		opts = append(opts, db.WithDSN(dsn))
	}
	if dbn := conf.DBName; dbn != "" {
		opts = append(opts, db.WithDBName(dbn))
	}
	rdb := db.NewRoach(opts...)
	err := rdb.InitDBIfNot()
	logging.LogWarnOnError(lg, err, "Initiate Cockroach DB connection")
	return rdb
}

func InstantiateJWTHandler(lg logging.Logger, tknKyF string) *token.Handler {
	JWTKey, err := ioutil.ReadFile(tknKyF)
	logging.LogFatalOnError(lg, err, "Read JWT key file")
	jwter, err := token.NewHandler(JWTKey)
	logging.LogFatalOnError(lg, err, "Instantiate JWT handler")
	return jwter
}

func InstantiateFacebook(conf config.Facebook) (*facebook.FacebookOAuth, error) {
	if conf.ID < 1 {
		return nil, nil
	}
	fbSecret, err := readFile(conf.SecretFilePath)
	if err != nil {
		return nil, fmt.Errorf("read facebook secret from file: %v", err)
	}
	fb, err := facebook.New(conf.ID, fbSecret)
	if err != nil {
		return nil, err
	}
	return fb, nil
}

func InstantiateSMSer(conf config.SMS) (model.SMSer, error) {
	if conf.ActiveAPI == "" {
		return nil, nil
	}
	var s model.SMSer
	switch conf.ActiveAPI {
	case config.SMSAPIAfricasTalking:
		apiKey, err := readFile(conf.AfricasTalking.APIKeyFile)
		if err != nil {
			return nil, fmt.Errorf("read africa's talking API key: %v", err)
		}
		s, err = africas_talking.NewSMSCl(conf.AfricasTalking.UserName, apiKey)
		if err != nil {
			return nil, fmt.Errorf("new africasTalking client: %v", err)
		}
	case config.SMSAPITwilio:
		tkn, err := readFile(conf.Twilio.TokenKeyFile)
		if err != nil {
			return nil, fmt.Errorf("read twilio token: %v", err)
		}
		s, err = twilio.NewSMSCl(conf.Twilio.ID, tkn, conf.Twilio.SenderPhone)
		if err != nil {
			return nil, fmt.Errorf("new twilio client: %v", err)
		}
	case config.SMSAPIMessageBird:
		apiKey, err := readFile(conf.MessageBird.APIKeyFile)
		if err != nil {
			return nil, fmt.Errorf("read messageBird API key: %v", err)
		}
		s, err = messagebird.NewClient(conf.MessageBird.AccountName, apiKey)
		if err != nil {
			return nil, fmt.Errorf("new messageBird client: %v", err)
		}
	default:
		return nil, fmt.Errorf("invalid API selected can be %s or %s",
			config.SMSAPIAfricasTalking, config.SMSAPITwilio)
	}
	if conf.TestNumber != "" {
		var testMessage string
		host := hostname()
		testMessage = fmt.Sprintf("The SMS API is being used on %s", host)
		if err := s.SMS(conf.TestNumber, testMessage); err != nil {
			return s, fmt.Errorf("test SMS: %v", err)
		}
	}
	return s, nil
}

func InstantiateSMTP(rdb *db.Roach, lg logging.Logger, conf config.SMTP) *smtp.Mailer {

	emailCl, err := smtp.New(rdb)
	logging.LogFatalOnError(lg, err, "Instantiate email API")

	err = emailCl.Configured()
	if err == nil {
		return emailCl
	}

	if !emailCl.IsNotFoundError(err) {
		logging.LogWarnOnError(lg, err, "Check email API configured")
		return emailCl
	}

	pass, err := readFile(conf.PasswordFile)
	logging.LogWarnOnError(lg, err, "Read SMTP password file")

	host := hostname()
	err = emailCl.SetConfig(
		smtp.Config{
			ServerAddress: conf.ServerAddress,
			TLSPort:       conf.TLSPort,
			SSLPort:       conf.SSLPort,
			Username:      conf.Username,
			FromEmail:     conf.FromEmail,
			Password:      pass,
		},
		model.SendMail{
			ToEmails: []string{conf.TestEmail},
			Subject:  "Authentication Micro-Service Started on " + host,
			Body:     template.HTML("The authentication micro-service is being used on " + host),
		},
	)
	logging.LogWarnOnError(lg, err, "Set default SMTP config")

	return emailCl
}

func Instantiate(confFile string, lg logging.Logger) (config.General, *model.Authentication, *api.Guard, *db.Roach, model.JWTEr, model.SMSer, *smtp.Mailer) {

	conf, err := config.ReadFile(confFile)
	logging.LogFatalOnError(lg, err, "Read config file")

	rdb := InstantiateRoach(lg, conf.Database)
	tg := InstantiateJWTHandler(lg, conf.Token.TokenKeyFile)

	var authOpts []model.Option
	fb, err := InstantiateFacebook(conf.Authentication.Facebook)
	logging.LogWarnOnError(lg, err, "Set up OAuth options")
	if fb != nil {
		authOpts = append(authOpts, model.WithFacebookCl(fb))
	}

	sms, err := InstantiateSMSer(conf.SMS)
	logging.LogWarnOnError(lg, err, "Instantiate SMS API")
	if sms != nil {
		authOpts = append(authOpts, model.WithSMSCl(sms))
	}

	emailCl := InstantiateSMTP(rdb, lg, conf.SMTP)
	authOpts = append(authOpts, model.WithEmailCl(emailCl))

	authOpts = append(
		authOpts,
		model.WithAppName(conf.Service.AppName),
		model.WithWebAppURL(conf.Service.WebAppURL),
		model.WithDevLockedToUser(conf.Authentication.LockDevsToUsers),
		model.WithSelfRegAllowed(conf.Authentication.AllowSelfReg),
	)

	a, err := model.NewAuthentication(rdb, tg, authOpts...)
	logging.LogFatalOnError(lg, err, "Instantiate Auth Model")

	g, err := api.NewGuard(rdb, api.WithMasterKey(conf.Service.MasterAPIKey))
	logging.LogFatalOnError(lg, err, "Instantate API access guard")

	return *conf, a, g, rdb, tg, sms, emailCl
}

func hostname() string {
	hostName, err := os.Hostname()
	if err != nil {
		return "an unknown host"
	}
	return hostName
}

func readFile(path string) (string, error) {
	contentB, err := ioutil.ReadFile(path)
	if err != nil {
		return "", fmt.Errorf("read: %v", err)
	}
	return string(contentB), nil
}
<|endoftext|>


<commit_before>package redisocket

import (
	"errors"
	"fmt"
	"log"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/garyburd/redigo/redis"
	"github.com/gorilla/websocket"
)

type User interface {
	Trigger(event string, p *Payload) (err error)
	Close()
}

type Payload struct {
	Data           []byte
	PrepareMessage *websocket.PreparedMessage
	IsPrepare      bool
	Event          string
}

type ReceiveMsg struct {
	Event        string
	EventHandler EventHandler
	Sub          bool
	ResponseMsg  []byte
}

type WebsocketOptional struct {
	WriteWait      time.Duration
	PongWait       time.Duration
	PingPeriod     time.Duration
	MaxMessageSize int64
	Upgrader       websocket.Upgrader
}

var (
	DefaultWebsocketOptional = WebsocketOptional{
		WriteWait:      10 * time.Second,
		PongWait:       60 * time.Second,
		PingPeriod:     (60 * time.Second * 9) / 10,
		MaxMessageSize: 512,
		Upgrader: websocket.Upgrader{
			CheckOrigin: func(r *http.Request) bool { return true },
		},
	}
)

var APPCLOSE = errors.New("APP_CLOSE")

type EventHandler func(event string, payload *Payload) error

type ReceiveMsgHandler func([]byte) (*ReceiveMsg, error)

func NewSender(m *redis.Pool) (e *Sender) {

	return &Sender{
		redisManager: m,
	}
}

type Sender struct {
	redisManager *redis.Pool
}

type BatchData struct {
	Event string
	Data  []byte
}

func (s *Sender) GetChannels(channelPrefix string, appKey string, pattern string) (channels []string, err error) {
	keyPrefix := fmt.Sprintf("%s%s@channels:", channelPrefix, appKey)
	conn := s.redisManager.Get()
	defer conn.Close()
	tmp, err := redis.Strings(conn.Do("keys", keyPrefix+pattern))
	channels = make([]string, len(tmp))
	for _, v := range tmp {
		channel := strings.Replace(v, keyPrefix, "", -1)
		if channel == "" {
			continue
		}
		channels = append(channels, channel)
	}

	return
}
func (s *Sender) GetOnlineByChannel(channelPrefix string, appKey string, channel string) (online []string, err error) {
	memberKey := fmt.Sprintf("%s%s@channels:%s", channelPrefix, appKey, channel)
	conn := s.redisManager.Get()
	defer conn.Close()
	online, err = redis.Strings(conn.Do("smembers", memberKey))
	return
}
func (s *Sender) GetOnline(channelPrefix string, appKey string) (online []string, err error) {
	memberKey := fmt.Sprintf("%s%s@online", channelPrefix, appKey)
	conn := s.redisManager.Get()
	defer conn.Close()
	online, err = redis.Strings(conn.Do("smembers", memberKey))
	return
}

func (s *Sender) PushBatch(channelPrefix, appKey string, data []BatchData) {
	conn := s.redisManager.Get()
	defer conn.Close()
	for _, d := range data {
		conn.Do("PUBLISH", channelPrefix+appKey+"@"+d.Event, d.Data)
	}
	return
}

func (s *Sender) Push(channelPrefix, appKey string, event string, data []byte) (val int, err error) {
	conn := s.redisManager.Get()
	defer conn.Close()
	val, err = redis.Int(conn.Do("PUBLISH", channelPrefix+appKey+"@"+event, data))
	return
}

//NewApp It's create a Hub
func NewHub(m *redis.Pool, debug bool) (e *Hub) {

	l := log.New(os.Stdout, "[redisocket.v2]", log.Lshortfile|log.Ldate|log.Lmicroseconds)
	pool := &Pool{

		users:     make(map[*Client]bool),
		broadcast: make(chan *eventPayload),
		join:      make(chan *Client),
		leave:     make(chan *Client),
		rpool:     m,
	}
	go pool.Run()
	return &Hub{

		Config:       DefaultWebsocketOptional,
		redisManager: m,
		psc:          &redis.PubSubConn{m.Get()},
		Pool:         pool,
		debug:        debug,
		log:          l,
	}

}
func (e *Hub) Upgrade(w http.ResponseWriter, r *http.Request, responseHeader http.Header, uid string, prefix string) (c *Client, err error) {
	ws, err := e.Config.Upgrader.Upgrade(w, r, responseHeader)
	if err != nil {
		return
	}
	c = &Client{
		prefix:  prefix,
		uid:     uid,
		ws:      ws,
		send:    make(chan *Payload, 32),
		RWMutex: new(sync.RWMutex),
		hub:     e,
		events:  make(map[string]EventHandler),
	}
	e.Join(c)
	return
}

type Hub struct {
	ChannelPrefix string
	Config        WebsocketOptional
	psc           *redis.PubSubConn
	redisManager  *redis.Pool
	*Pool
	debug bool
	log   *log.Logger
}

func (a *Hub) Ping() (err error) {
	_, err = a.redisManager.Get().Do("PING")
	if err != nil {
		return
	}
	return
}
func (a *Hub) logger(format string, v ...interface{}) {
	if a.debug {
		a.log.Printf(format, v...)
	}
}
func (a *Hub) CountOnlineUsers() (i int) {
	return len(a.Pool.users)
}
func (a *Hub) listenRedis() <-chan error {

	errChan := make(chan error, 1)
	go func() {
		for {
			switch v := a.psc.Receive().(type) {
			case redis.PMessage:

				//過濾掉前綴
				channel := strings.Replace(v.Channel, a.ChannelPrefix, "", -1)
				//過濾掉@ 之前的字
				sch := strings.SplitN(channel, "@", 2)
				if len(sch) != 2 {
					continue
				}

				//過濾掉星號
				channel = strings.Replace(sch[1], "*", "", -1)
				pMsg, err := websocket.NewPreparedMessage(websocket.TextMessage, v.Data)
				if err != nil {
					continue
				}
				p := &Payload{
					PrepareMessage: pMsg,
					IsPrepare:      true,
				}
				a.Broadcast(channel, p)

			case error:
				errChan <- v

				break
			}
		}
	}()
	return errChan
}

func (a *Hub) Listen(channelPrefix string) error {
	a.Pool.channelPrefix = channelPrefix
	a.ChannelPrefix = channelPrefix
	a.psc.PSubscribe(channelPrefix + "*")
	redisErr := a.listenRedis()
	select {
	case e := <-redisErr:
		return e
	}
}
func (a *Hub) Close() {
	return

}
<commit_msg>update vendor<commit_after>package redisocket

import (
	"errors"
	"fmt"
	"log"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/garyburd/redigo/redis"
	"github.com/gorilla/websocket"
)

type User interface {
	Trigger(event string, p *Payload) (err error)
	Close()
}

type Payload struct {
	Data           []byte
	PrepareMessage *websocket.PreparedMessage
	IsPrepare      bool
	Event          string
}

type ReceiveMsg struct {
	Event        string
	EventHandler EventHandler
	Sub          bool
	ResponseMsg  []byte
}

type WebsocketOptional struct {
	WriteWait      time.Duration
	PongWait       time.Duration
	PingPeriod     time.Duration
	MaxMessageSize int64
	Upgrader       websocket.Upgrader
}

var (
	DefaultWebsocketOptional = WebsocketOptional{
		WriteWait:      10 * time.Second,
		PongWait:       60 * time.Second,
		PingPeriod:     (60 * time.Second * 9) / 10,
		MaxMessageSize: 512,
		Upgrader: websocket.Upgrader{
			CheckOrigin: func(r *http.Request) bool { return true },
		},
	}
)

var APPCLOSE = errors.New("APP_CLOSE")

type EventHandler func(event string, payload *Payload) error

type ReceiveMsgHandler func([]byte) (*ReceiveMsg, error)

func NewSender(m *redis.Pool) (e *Sender) {

	return &Sender{
		redisManager: m,
	}
}

type Sender struct {
	redisManager *redis.Pool
}

type BatchData struct {
	Event string
	Data  []byte
}

func (s *Sender) GetChannels(channelPrefix string, appKey string, pattern string) (channels []string, err error) {
	keyPrefix := fmt.Sprintf("%s%s@channels:", channelPrefix, appKey)
	conn := s.redisManager.Get()
	defer conn.Close()
	tmp, err := redis.Strings(conn.Do("keys", keyPrefix+pattern))
	channels = make([]string, 0)
	for _, v := range tmp {
		channel := strings.Replace(v, keyPrefix, "", -1)
		if channel == "" {
			continue
		}
		channels = append(channels, channel)
	}

	return
}
func (s *Sender) GetOnlineByChannel(channelPrefix string, appKey string, channel string) (online []string, err error) {
	memberKey := fmt.Sprintf("%s%s@channels:%s", channelPrefix, appKey, channel)
	conn := s.redisManager.Get()
	defer conn.Close()
	online, err = redis.Strings(conn.Do("smembers", memberKey))
	return
}
func (s *Sender) GetOnline(channelPrefix string, appKey string) (online []string, err error) {
	memberKey := fmt.Sprintf("%s%s@online", channelPrefix, appKey)
	conn := s.redisManager.Get()
	defer conn.Close()
	online, err = redis.Strings(conn.Do("smembers", memberKey))
	return
}

func (s *Sender) PushBatch(channelPrefix, appKey string, data []BatchData) {
	conn := s.redisManager.Get()
	defer conn.Close()
	for _, d := range data {
		conn.Do("PUBLISH", channelPrefix+appKey+"@"+d.Event, d.Data)
	}
	return
}

func (s *Sender) Push(channelPrefix, appKey string, event string, data []byte) (val int, err error) {
	conn := s.redisManager.Get()
	defer conn.Close()
	val, err = redis.Int(conn.Do("PUBLISH", channelPrefix+appKey+"@"+event, data))
	return
}

//NewApp It's create a Hub
func NewHub(m *redis.Pool, debug bool) (e *Hub) {

	l := log.New(os.Stdout, "[redisocket.v2]", log.Lshortfile|log.Ldate|log.Lmicroseconds)
	pool := &Pool{

		users:     make(map[*Client]bool),
		broadcast: make(chan *eventPayload),
		join:      make(chan *Client),
		leave:     make(chan *Client),
		rpool:     m,
	}
	go pool.Run()
	return &Hub{

		Config:       DefaultWebsocketOptional,
		redisManager: m,
		psc:          &redis.PubSubConn{m.Get()},
		Pool:         pool,
		debug:        debug,
		log:          l,
	}

}
func (e *Hub) Upgrade(w http.ResponseWriter, r *http.Request, responseHeader http.Header, uid string, prefix string) (c *Client, err error) {
	ws, err := e.Config.Upgrader.Upgrade(w, r, responseHeader)
	if err != nil {
		return
	}
	c = &Client{
		prefix:  prefix,
		uid:     uid,
		ws:      ws,
		send:    make(chan *Payload, 32),
		RWMutex: new(sync.RWMutex),
		hub:     e,
		events:  make(map[string]EventHandler),
	}
	e.Join(c)
	return
}

type Hub struct {
	ChannelPrefix string
	Config        WebsocketOptional
	psc           *redis.PubSubConn
	redisManager  *redis.Pool
	*Pool
	debug bool
	log   *log.Logger
}

func (a *Hub) Ping() (err error) {
	_, err = a.redisManager.Get().Do("PING")
	if err != nil {
		return
	}
	return
}
func (a *Hub) logger(format string, v ...interface{}) {
	if a.debug {
		a.log.Printf(format, v...)
	}
}
func (a *Hub) CountOnlineUsers() (i int) {
	return len(a.Pool.users)
}
func (a *Hub) listenRedis() <-chan error {

	errChan := make(chan error, 1)
	go func() {
		for {
			switch v := a.psc.Receive().(type) {
			case redis.PMessage:

				//過濾掉前綴
				channel := strings.Replace(v.Channel, a.ChannelPrefix, "", -1)
				//過濾掉@ 之前的字
				sch := strings.SplitN(channel, "@", 2)
				if len(sch) != 2 {
					continue
				}

				//過濾掉星號
				channel = strings.Replace(sch[1], "*", "", -1)
				pMsg, err := websocket.NewPreparedMessage(websocket.TextMessage, v.Data)
				if err != nil {
					continue
				}
				p := &Payload{
					PrepareMessage: pMsg,
					IsPrepare:      true,
				}
				a.Broadcast(channel, p)

			case error:
				errChan <- v

				break
			}
		}
	}()
	return errChan
}

func (a *Hub) Listen(channelPrefix string) error {
	a.Pool.channelPrefix = channelPrefix
	a.ChannelPrefix = channelPrefix
	a.psc.PSubscribe(channelPrefix + "*")
	redisErr := a.listenRedis()
	select {
	case e := <-redisErr:
		return e
	}
}
func (a *Hub) Close() {
	return

}
<|endoftext|>


<commit_before>// Copyright 2019 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package bhyve

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

	"github.com/google/syzkaller/pkg/config"
	"github.com/google/syzkaller/pkg/log"
	"github.com/google/syzkaller/pkg/osutil"
	"github.com/google/syzkaller/vm/vmimpl"
)

func init() {
	vmimpl.Register("bhyve", ctor, true)
}

type Config struct {
	Bridge  string `json:"bridge"`  // name of network bridge device
	Count   int    `json:"count"`   // number of VMs to use
	HostIP  string `json:"hostip"`  // VM host IP address
	Mem     string `json:"mem"`     // amount of VM memory
	Dataset string `json:"dataset"` // ZFS dataset containing VM image
}

type Pool struct {
	env *vmimpl.Env
	cfg *Config
}

type instance struct {
	cfg      *Config
	snapshot string
	tapdev   string
	image    string
	debug    bool
	os       string
	sshkey   string
	sshuser  string
	sshhost  string
	merger   *vmimpl.OutputMerger
	vmName   string
	bhyve    *exec.Cmd
}

var ipRegex = regexp.MustCompile(`bound to (([0-9]+\.){3}[0-9]+) `)
var tapRegex = regexp.MustCompile(`^tap[0-9]+`)

func ctor(env *vmimpl.Env) (vmimpl.Pool, error) {
	cfg := &Config{
		Count: 1,
		Mem:   "512M",
	}
	if err := config.LoadData(env.Config, cfg); err != nil {
		return nil, fmt.Errorf("failed to parse bhyve vm config: %v", err)
	}
	if cfg.Count < 1 || cfg.Count > 128 {
		return nil, fmt.Errorf("invalid config param count: %v, want [1-128]", cfg.Count)
	}
	if env.Debug && cfg.Count > 1 {
		log.Logf(0, "limiting number of VMs from %v to 1 in debug mode", cfg.Count)
		cfg.Count = 1
	}
	pool := &Pool{
		cfg: cfg,
		env: env,
	}
	return pool, nil
}

func (pool *Pool) Count() int {
	return pool.cfg.Count
}

func (pool *Pool) Create(workdir string, index int) (vmimpl.Instance, error) {
	inst := &instance{
		cfg:     pool.cfg,
		debug:   pool.env.Debug,
		os:      pool.env.OS,
		sshkey:  pool.env.SSHKey,
		sshuser: pool.env.SSHUser,
		vmName:  fmt.Sprintf("syzkaller-%v-%v", pool.env.Name, index),
	}

	dataset := inst.cfg.Dataset
	mountpoint, err := osutil.RunCmd(time.Minute, "", "zfs", "get", "-H", "-o", "value", "mountpoint", dataset)
	if err != nil {
		return nil, err
	}

	snapshot := fmt.Sprintf("%v@bhyve-%v", dataset, inst.vmName)
	clone := fmt.Sprintf("%v/bhyve-%v", dataset, inst.vmName)

	prefix := strings.TrimSuffix(string(mountpoint), "\n") + "/"
	image := strings.TrimPrefix(pool.env.Image, prefix)
	if image == pool.env.Image {
		return nil, fmt.Errorf("image file %v not contained in dataset %v", image, prefix)
	}
	inst.image = prefix + fmt.Sprintf("bhyve-%v", inst.vmName) + "/" + image

	// Stop the instance from a previous run in case it's still running.
	osutil.RunCmd(time.Minute, "", "bhyvectl", "--destroy", fmt.Sprintf("--vm=%v", inst.vmName))
	// Destroy a lingering snapshot and clone.
	osutil.RunCmd(time.Minute, "", "zfs", "destroy", "-R", snapshot)

	// Create a snapshot of the data set containing the VM image.  bhyve will
	// use a clone of the snapshot, which gets recreated every time the VM
	// is restarted.  This is all to work around bhyve's current lack of an
	// image snapshot facility.
	if _, err := osutil.RunCmd(time.Minute, "", "zfs", "snapshot", snapshot); err != nil {
		inst.Close()
		return nil, err
	}
	inst.snapshot = snapshot
	if _, err := osutil.RunCmd(time.Minute, "", "zfs", "clone", snapshot, clone); err != nil {
		inst.Close()
		return nil, err
	}

	tapdev, err := osutil.RunCmd(time.Minute, "", "ifconfig", "tap", "create")
	if err != nil {
		inst.Close()
		return nil, err
	}
	inst.tapdev = tapRegex.FindString(string(tapdev))
	if _, err := osutil.RunCmd(time.Minute, "", "ifconfig", inst.cfg.Bridge, "addm", inst.tapdev); err != nil {
		inst.Close()
		return nil, err
	}

	if err := inst.Boot(); err != nil {
		inst.Close()
		return nil, err
	}

	return inst, nil
}

func (inst *instance) Boot() error {
	loaderArgs := []string{
		"-c", "stdio",
		"-m", inst.cfg.Mem,
		"-d", inst.image,
		"-e", "autoboot_delay=0",
		inst.vmName,
	}

	// Stop the instance from the previous run in case it's still running.
	osutil.RunCmd(time.Minute, "", "bhyvectl", "--destroy", fmt.Sprintf("--vm=%v", inst.vmName))

	_, err := osutil.RunCmd(time.Minute, "", "bhyveload", loaderArgs...)
	if err != nil {
		return err
	}

	bhyveArgs := []string{
		"-H", "-A", "-P",
		"-c", "1",
		"-m", inst.cfg.Mem,
		"-s", "0:0,hostbridge",
		"-s", "1:0,lpc",
		"-s", fmt.Sprintf("2:0,virtio-net,%v", inst.tapdev),
		"-s", fmt.Sprintf("3:0,virtio-blk,%v", inst.image),
		"-l", "com1,stdio",
		inst.vmName,
	}

	outr, outw, err := osutil.LongPipe()
	if err != nil {
		return err
	}

	bhyve := osutil.Command("bhyve", bhyveArgs...)
	bhyve.Stdout = outw
	bhyve.Stderr = outw
	if err := bhyve.Start(); err != nil {
		outr.Close()
		outw.Close()
		return err
	}
	outw.Close()
	outw = nil
	inst.bhyve = bhyve

	var tee io.Writer
	if inst.debug {
		tee = os.Stdout
	}
	inst.merger = vmimpl.NewOutputMerger(tee)
	inst.merger.Add("console", outr)
	outr = nil

	var bootOutput []byte
	bootOutputStop := make(chan bool)
	ipch := make(chan string, 1)
	go func() {
		gotip := false
		for {
			select {
			case out := <-inst.merger.Output:
				bootOutput = append(bootOutput, out...)
			case <-bootOutputStop:
				close(bootOutputStop)
				return
			}
			if gotip {
				continue
			}
			if ip := parseIP(bootOutput); ip != "" {
				ipch <- ip
				gotip = true
			}
		}
	}()

	select {
	case ip := <-ipch:
		inst.sshhost = ip
	case <-inst.merger.Err:
		bootOutputStop <- true
		<-bootOutputStop
		return vmimpl.BootError{Title: "bhyve exited", Output: bootOutput}
	case <-time.After(10 * time.Minute):
		bootOutputStop <- true
		<-bootOutputStop
		return vmimpl.BootError{Title: "no IP found", Output: bootOutput}
	}

	if err := vmimpl.WaitForSSH(inst.debug, 10*time.Minute, inst.sshhost,
		inst.sshkey, inst.sshuser, inst.os, 22, nil); err != nil {
		bootOutputStop <- true
		<-bootOutputStop
		return vmimpl.MakeBootError(err, bootOutput)
	}
	bootOutputStop <- true
	return nil
}

func (inst *instance) Close() {
	if inst.bhyve != nil {
		inst.bhyve.Process.Kill()
		inst.bhyve.Wait()
		inst.bhyve = nil
	}
	if inst.snapshot != "" {
		osutil.RunCmd(time.Minute, "", "zfs", "destroy", "-R", inst.snapshot)
		inst.snapshot = ""
	}
	if inst.tapdev != "" {
		osutil.RunCmd(time.Minute, "", "ifconfig", inst.tapdev, "destroy")
		inst.tapdev = ""
	}
}

func (inst *instance) Forward(port int) (string, error) {
	return fmt.Sprintf("%v:%v", inst.cfg.HostIP, port), nil
}

func (inst *instance) Copy(hostSrc string) (string, error) {
	vmDst := filepath.Join("/root", filepath.Base(hostSrc))
	args := append(vmimpl.SCPArgs(inst.debug, inst.sshkey, 22),
		hostSrc, inst.sshuser+"@"+inst.sshhost+":"+vmDst)
	if inst.debug {
		log.Logf(0, "running command: scp %#v", args)
	}
	_, err := osutil.RunCmd(10*time.Minute, "", "scp", args...)
	if err != nil {
		return "", err
	}
	return vmDst, nil
}

func (inst *instance) Run(timeout time.Duration, stop <-chan bool, command string) (
	<-chan []byte, <-chan error, error) {
	rpipe, wpipe, err := osutil.LongPipe()
	if err != nil {
		return nil, nil, err
	}
	inst.merger.Add("ssh", rpipe)

	args := append(vmimpl.SSHArgs(inst.debug, inst.sshkey, 22),
		inst.sshuser+"@"+inst.sshhost, command)
	if inst.debug {
		log.Logf(0, "running command: ssh %#v", args)
	}
	cmd := osutil.Command("ssh", args...)
	cmd.Stdout = wpipe
	cmd.Stderr = wpipe
	if err := cmd.Start(); err != nil {
		wpipe.Close()
		return nil, nil, err
	}
	wpipe.Close()
	errc := make(chan error, 1)
	signal := func(err error) {
		select {
		case errc <- err:
		default:
		}
	}

	go func() {
		select {
		case <-time.After(timeout):
			signal(vmimpl.ErrTimeout)
		case <-stop:
			signal(vmimpl.ErrTimeout)
		case err := <-inst.merger.Err:
			cmd.Process.Kill()
			if cmdErr := cmd.Wait(); cmdErr == nil {
				// If the command exited successfully, we got EOF error from merger.
				// But in this case no error has happened and the EOF is expected.
				err = nil
			}
			signal(err)
			return
		}
		cmd.Process.Kill()
		cmd.Wait()
	}()
	return inst.merger.Output, errc, nil
}

func (inst *instance) Diagnose() ([]byte, bool) {
	return nil, false
}

func parseIP(output []byte) string {
	matches := ipRegex.FindSubmatch(output)
	if len(matches) < 2 {
		return ""
	}
	return string(matches[1])
}
<commit_msg>vm/bhyve: ensure the VM is destroyed after closing<commit_after>// Copyright 2019 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package bhyve

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

	"github.com/google/syzkaller/pkg/config"
	"github.com/google/syzkaller/pkg/log"
	"github.com/google/syzkaller/pkg/osutil"
	"github.com/google/syzkaller/vm/vmimpl"
)

func init() {
	vmimpl.Register("bhyve", ctor, true)
}

type Config struct {
	Bridge  string `json:"bridge"`  // name of network bridge device
	Count   int    `json:"count"`   // number of VMs to use
	HostIP  string `json:"hostip"`  // VM host IP address
	Mem     string `json:"mem"`     // amount of VM memory
	Dataset string `json:"dataset"` // ZFS dataset containing VM image
}

type Pool struct {
	env *vmimpl.Env
	cfg *Config
}

type instance struct {
	cfg      *Config
	snapshot string
	tapdev   string
	image    string
	debug    bool
	os       string
	sshkey   string
	sshuser  string
	sshhost  string
	merger   *vmimpl.OutputMerger
	vmName   string
	bhyve    *exec.Cmd
}

var ipRegex = regexp.MustCompile(`bound to (([0-9]+\.){3}[0-9]+) `)
var tapRegex = regexp.MustCompile(`^tap[0-9]+`)

func ctor(env *vmimpl.Env) (vmimpl.Pool, error) {
	cfg := &Config{
		Count: 1,
		Mem:   "512M",
	}
	if err := config.LoadData(env.Config, cfg); err != nil {
		return nil, fmt.Errorf("failed to parse bhyve vm config: %v", err)
	}
	if cfg.Count < 1 || cfg.Count > 128 {
		return nil, fmt.Errorf("invalid config param count: %v, want [1-128]", cfg.Count)
	}
	if env.Debug && cfg.Count > 1 {
		log.Logf(0, "limiting number of VMs from %v to 1 in debug mode", cfg.Count)
		cfg.Count = 1
	}
	pool := &Pool{
		cfg: cfg,
		env: env,
	}
	return pool, nil
}

func (pool *Pool) Count() int {
	return pool.cfg.Count
}

func (pool *Pool) Create(workdir string, index int) (vmimpl.Instance, error) {
	inst := &instance{
		cfg:     pool.cfg,
		debug:   pool.env.Debug,
		os:      pool.env.OS,
		sshkey:  pool.env.SSHKey,
		sshuser: pool.env.SSHUser,
		vmName:  fmt.Sprintf("syzkaller-%v-%v", pool.env.Name, index),
	}

	dataset := inst.cfg.Dataset
	mountpoint, err := osutil.RunCmd(time.Minute, "", "zfs", "get", "-H", "-o", "value", "mountpoint", dataset)
	if err != nil {
		return nil, err
	}

	snapshot := fmt.Sprintf("%v@bhyve-%v", dataset, inst.vmName)
	clone := fmt.Sprintf("%v/bhyve-%v", dataset, inst.vmName)

	prefix := strings.TrimSuffix(string(mountpoint), "\n") + "/"
	image := strings.TrimPrefix(pool.env.Image, prefix)
	if image == pool.env.Image {
		return nil, fmt.Errorf("image file %v not contained in dataset %v", image, prefix)
	}
	inst.image = prefix + fmt.Sprintf("bhyve-%v", inst.vmName) + "/" + image

	// Stop the instance from a previous run in case it's still running.
	osutil.RunCmd(time.Minute, "", "bhyvectl", "--destroy", fmt.Sprintf("--vm=%v", inst.vmName))
	// Destroy a lingering snapshot and clone.
	osutil.RunCmd(time.Minute, "", "zfs", "destroy", "-R", snapshot)

	// Create a snapshot of the data set containing the VM image.  bhyve will
	// use a clone of the snapshot, which gets recreated every time the VM
	// is restarted.  This is all to work around bhyve's current lack of an
	// image snapshot facility.
	if _, err := osutil.RunCmd(time.Minute, "", "zfs", "snapshot", snapshot); err != nil {
		inst.Close()
		return nil, err
	}
	inst.snapshot = snapshot
	if _, err := osutil.RunCmd(time.Minute, "", "zfs", "clone", snapshot, clone); err != nil {
		inst.Close()
		return nil, err
	}

	tapdev, err := osutil.RunCmd(time.Minute, "", "ifconfig", "tap", "create")
	if err != nil {
		inst.Close()
		return nil, err
	}
	inst.tapdev = tapRegex.FindString(string(tapdev))
	if _, err := osutil.RunCmd(time.Minute, "", "ifconfig", inst.cfg.Bridge, "addm", inst.tapdev); err != nil {
		inst.Close()
		return nil, err
	}

	if err := inst.Boot(); err != nil {
		inst.Close()
		return nil, err
	}

	return inst, nil
}

func (inst *instance) Boot() error {
	loaderArgs := []string{
		"-c", "stdio",
		"-m", inst.cfg.Mem,
		"-d", inst.image,
		"-e", "autoboot_delay=0",
		inst.vmName,
	}

	// Stop the instance from the previous run in case it's still running.
	osutil.RunCmd(time.Minute, "", "bhyvectl", "--destroy", fmt.Sprintf("--vm=%v", inst.vmName))

	_, err := osutil.RunCmd(time.Minute, "", "bhyveload", loaderArgs...)
	if err != nil {
		return err
	}

	bhyveArgs := []string{
		"-H", "-A", "-P",
		"-c", "1",
		"-m", inst.cfg.Mem,
		"-s", "0:0,hostbridge",
		"-s", "1:0,lpc",
		"-s", fmt.Sprintf("2:0,virtio-net,%v", inst.tapdev),
		"-s", fmt.Sprintf("3:0,virtio-blk,%v", inst.image),
		"-l", "com1,stdio",
		inst.vmName,
	}

	outr, outw, err := osutil.LongPipe()
	if err != nil {
		return err
	}

	bhyve := osutil.Command("bhyve", bhyveArgs...)
	bhyve.Stdout = outw
	bhyve.Stderr = outw
	if err := bhyve.Start(); err != nil {
		outr.Close()
		outw.Close()
		return err
	}
	outw.Close()
	outw = nil
	inst.bhyve = bhyve

	var tee io.Writer
	if inst.debug {
		tee = os.Stdout
	}
	inst.merger = vmimpl.NewOutputMerger(tee)
	inst.merger.Add("console", outr)
	outr = nil

	var bootOutput []byte
	bootOutputStop := make(chan bool)
	ipch := make(chan string, 1)
	go func() {
		gotip := false
		for {
			select {
			case out := <-inst.merger.Output:
				bootOutput = append(bootOutput, out...)
			case <-bootOutputStop:
				close(bootOutputStop)
				return
			}
			if gotip {
				continue
			}
			if ip := parseIP(bootOutput); ip != "" {
				ipch <- ip
				gotip = true
			}
		}
	}()

	select {
	case ip := <-ipch:
		inst.sshhost = ip
	case <-inst.merger.Err:
		bootOutputStop <- true
		<-bootOutputStop
		return vmimpl.BootError{Title: "bhyve exited", Output: bootOutput}
	case <-time.After(10 * time.Minute):
		bootOutputStop <- true
		<-bootOutputStop
		return vmimpl.BootError{Title: "no IP found", Output: bootOutput}
	}

	if err := vmimpl.WaitForSSH(inst.debug, 10*time.Minute, inst.sshhost,
		inst.sshkey, inst.sshuser, inst.os, 22, nil); err != nil {
		bootOutputStop <- true
		<-bootOutputStop
		return vmimpl.MakeBootError(err, bootOutput)
	}
	bootOutputStop <- true
	return nil
}

func (inst *instance) Close() {
	if inst.bhyve != nil {
		inst.bhyve.Process.Kill()
		inst.bhyve.Wait()
		osutil.RunCmd(time.Minute, "", "bhyvectl", fmt.Sprintf("--vm=%v", inst.vmName), "--destroy")
		inst.bhyve = nil
	}
	if inst.snapshot != "" {
		osutil.RunCmd(time.Minute, "", "zfs", "destroy", "-R", inst.snapshot)
		inst.snapshot = ""
	}
	if inst.tapdev != "" {
		osutil.RunCmd(time.Minute, "", "ifconfig", inst.tapdev, "destroy")
		inst.tapdev = ""
	}
}

func (inst *instance) Forward(port int) (string, error) {
	return fmt.Sprintf("%v:%v", inst.cfg.HostIP, port), nil
}

func (inst *instance) Copy(hostSrc string) (string, error) {
	vmDst := filepath.Join("/root", filepath.Base(hostSrc))
	args := append(vmimpl.SCPArgs(inst.debug, inst.sshkey, 22),
		hostSrc, inst.sshuser+"@"+inst.sshhost+":"+vmDst)
	if inst.debug {
		log.Logf(0, "running command: scp %#v", args)
	}
	_, err := osutil.RunCmd(10*time.Minute, "", "scp", args...)
	if err != nil {
		return "", err
	}
	return vmDst, nil
}

func (inst *instance) Run(timeout time.Duration, stop <-chan bool, command string) (
	<-chan []byte, <-chan error, error) {
	rpipe, wpipe, err := osutil.LongPipe()
	if err != nil {
		return nil, nil, err
	}
	inst.merger.Add("ssh", rpipe)

	args := append(vmimpl.SSHArgs(inst.debug, inst.sshkey, 22),
		inst.sshuser+"@"+inst.sshhost, command)
	if inst.debug {
		log.Logf(0, "running command: ssh %#v", args)
	}
	cmd := osutil.Command("ssh", args...)
	cmd.Stdout = wpipe
	cmd.Stderr = wpipe
	if err := cmd.Start(); err != nil {
		wpipe.Close()
		return nil, nil, err
	}
	wpipe.Close()
	errc := make(chan error, 1)
	signal := func(err error) {
		select {
		case errc <- err:
		default:
		}
	}

	go func() {
		select {
		case <-time.After(timeout):
			signal(vmimpl.ErrTimeout)
		case <-stop:
			signal(vmimpl.ErrTimeout)
		case err := <-inst.merger.Err:
			cmd.Process.Kill()
			if cmdErr := cmd.Wait(); cmdErr == nil {
				// If the command exited successfully, we got EOF error from merger.
				// But in this case no error has happened and the EOF is expected.
				err = nil
			}
			signal(err)
			return
		}
		cmd.Process.Kill()
		cmd.Wait()
	}()
	return inst.merger.Output, errc, nil
}

func (inst *instance) Diagnose() ([]byte, bool) {
	return nil, false
}

func parseIP(output []byte) string {
	matches := ipRegex.FindSubmatch(output)
	if len(matches) < 2 {
		return ""
	}
	return string(matches[1])
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Kubernetes Authors.

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

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

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

package cmd

import (
	"fmt"
	"io"
	"sort"
	"strings"

	"github.com/spf13/cobra"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/kubernetes/pkg/kubectl/cmd/templates"
	cmdutil "k8s.io/kubernetes/pkg/kubectl/cmd/util"
	"k8s.io/kubernetes/pkg/kubectl/genericclioptions"
	"k8s.io/kubernetes/pkg/printers"
)

var (
	apiresourcesExample = templates.Examples(`
		# Print the supported API Resources
		kubectl api-resources

		# Print the supported API Resources with more information
		kubectl api-resources -o wide

		# Print the supported namespaced resources
		kubectl api-resources --namespaced=true

		# Print the supported non-namespaced resources
		kubectl api-resources --namespaced=false

		# Print the supported API Resources with specific APIGroup
		kubectl api-resources --api-group=extensions`)
)

// ApiResourcesOptions is the start of the data required to perform the operation.  As new fields are added, add them here instead of
// referencing the cmd.Flags()
type ApiResourcesOptions struct {
	Output     string
	APIGroup   string
	Namespaced bool
	Verbs      []string
	NoHeaders  bool
	Cached     bool

	genericclioptions.IOStreams
}

// groupResource contains the APIGroup and APIResource
type groupResource struct {
	APIGroup    string
	APIResource metav1.APIResource
}

func NewAPIResourceOptions(ioStreams genericclioptions.IOStreams) *ApiResourcesOptions {
	return &ApiResourcesOptions{
		IOStreams:  ioStreams,
		Namespaced: true,
	}
}

func NewCmdApiResources(f cmdutil.Factory, ioStreams genericclioptions.IOStreams) *cobra.Command {
	o := NewAPIResourceOptions(ioStreams)

	cmd := &cobra.Command{
		Use:     "api-resources",
		Short:   "Print the supported API resources on the server",
		Long:    "Print the supported API resources on the server",
		Example: apiresourcesExample,
		Run: func(cmd *cobra.Command, args []string) {
			cmdutil.CheckErr(o.Validate(cmd))
			cmdutil.CheckErr(o.RunApiResources(cmd, f))
		},
	}

	cmd.Flags().BoolVar(&o.NoHeaders, "no-headers", o.NoHeaders, "When using the default or custom-column output format, don't print headers (default print headers).")
	cmd.Flags().StringVarP(&o.Output, "output", "o", o.Output, "Output format. One of: wide|name.")

	cmd.Flags().StringVar(&o.APIGroup, "api-group", o.APIGroup, "Limit to resources in the specified API group.")
	cmd.Flags().BoolVar(&o.Namespaced, "namespaced", o.Namespaced, "If false, non-namespaced resources will be returned, otherwise returning namespaced resources by default.")
	cmd.Flags().StringSliceVar(&o.Verbs, "verbs", o.Verbs, "Limit to resources that support the specified verbs.")
	cmd.Flags().BoolVar(&o.Cached, "cached", o.Cached, "Use the cached list of resources if available.")
	return cmd
}

func (o *ApiResourcesOptions) Validate(cmd *cobra.Command) error {
	supportedOutputTypes := sets.NewString("", "wide", "name")
	if !supportedOutputTypes.Has(o.Output) {
		return fmt.Errorf("--output %v is not available", o.Output)
	}
	return nil
}

func (o *ApiResourcesOptions) RunApiResources(cmd *cobra.Command, f cmdutil.Factory) error {
	w := printers.GetNewTabWriter(o.Out)
	defer w.Flush()

	discoveryclient, err := f.ToDiscoveryClient()
	if err != nil {
		return err
	}

	if !o.Cached {
		// Always request fresh data from the server
		discoveryclient.Invalidate()
	}

	lists, err := discoveryclient.ServerPreferredResources()
	if err != nil {
		return err
	}

	resources := []groupResource{}

	groupChanged := cmd.Flags().Changed("api-group")
	nsChanged := cmd.Flags().Changed("namespaced")

	for _, list := range lists {
		if len(list.APIResources) == 0 {
			continue
		}
		gv, err := schema.ParseGroupVersion(list.GroupVersion)
		if err != nil {
			continue
		}
		for _, resource := range list.APIResources {
			if len(resource.Verbs) == 0 {
				continue
			}
			// filter apiGroup
			if groupChanged && o.APIGroup != gv.Group {
				continue
			}
			// filter namespaced
			if nsChanged && o.Namespaced != resource.Namespaced {
				continue
			}
			// filter to resources that support the specified verbs
			if len(o.Verbs) > 0 && !sets.NewString(resource.Verbs...).HasAll(o.Verbs...) {
				continue
			}
			resources = append(resources, groupResource{
				APIGroup:    gv.Group,
				APIResource: resource,
			})
		}
	}

	if o.NoHeaders == false && o.Output != "name" {
		if err = printContextHeaders(w, o.Output); err != nil {
			return err
		}
	}

	sort.Stable(sortableGroupResource(resources))
	for _, r := range resources {
		switch o.Output {
		case "name":
			name := r.APIResource.Name
			if len(r.APIGroup) > 0 {
				name += "." + r.APIGroup
			}
			if _, err := fmt.Fprintf(w, "%s\n", name); err != nil {
				return err
			}
		case "wide":
			if _, err := fmt.Fprintf(w, "%s\t%s\t%s\t%v\t%s\t%v\n",
				r.APIResource.Name,
				strings.Join(r.APIResource.ShortNames, ","),
				r.APIGroup,
				r.APIResource.Namespaced,
				r.APIResource.Kind,
				r.APIResource.Verbs); err != nil {
				return err
			}
		case "":
			if _, err := fmt.Fprintf(w, "%s\t%s\t%s\t%v\t%s\n",
				r.APIResource.Name,
				strings.Join(r.APIResource.ShortNames, ","),
				r.APIGroup,
				r.APIResource.Namespaced,
				r.APIResource.Kind); err != nil {
				return err
			}
		}
	}
	return nil
}

func printContextHeaders(out io.Writer, output string) error {
	columnNames := []string{"NAME", "SHORTNAMES", "APIGROUP", "NAMESPACED", "KIND"}
	if output == "wide" {
		columnNames = append(columnNames, "VERBS")
	}
	_, err := fmt.Fprintf(out, "%s\n", strings.Join(columnNames, "\t"))
	return err
}

type sortableGroupResource []groupResource

func (s sortableGroupResource) Len() int      { return len(s) }
func (s sortableGroupResource) Swap(i, j int) { s[i], s[j] = s[j], s[i] }
func (s sortableGroupResource) Less(i, j int) bool {
	ret := strings.Compare(s[i].APIGroup, s[j].APIGroup)
	if ret > 0 {
		return false
	} else if ret == 0 {
		return strings.Compare(s[i].APIResource.Name, s[j].APIResource.Name) < 0
	}
	return true
}
<commit_msg>UPSTREAM: 73035: make api-resource discovery errors non-fatal<commit_after>/*
Copyright 2018 The Kubernetes Authors.

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

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

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

package cmd

import (
	"fmt"
	"io"
	"sort"
	"strings"

	"github.com/spf13/cobra"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime/schema"
	"k8s.io/apimachinery/pkg/util/errors"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/kubernetes/pkg/kubectl/cmd/templates"
	cmdutil "k8s.io/kubernetes/pkg/kubectl/cmd/util"
	"k8s.io/kubernetes/pkg/kubectl/genericclioptions"
	"k8s.io/kubernetes/pkg/printers"
)

var (
	apiresourcesExample = templates.Examples(`
		# Print the supported API Resources
		kubectl api-resources

		# Print the supported API Resources with more information
		kubectl api-resources -o wide

		# Print the supported namespaced resources
		kubectl api-resources --namespaced=true

		# Print the supported non-namespaced resources
		kubectl api-resources --namespaced=false

		# Print the supported API Resources with specific APIGroup
		kubectl api-resources --api-group=extensions`)
)

// ApiResourcesOptions is the start of the data required to perform the operation.  As new fields are added, add them here instead of
// referencing the cmd.Flags()
type ApiResourcesOptions struct {
	Output     string
	APIGroup   string
	Namespaced bool
	Verbs      []string
	NoHeaders  bool
	Cached     bool

	genericclioptions.IOStreams
}

// groupResource contains the APIGroup and APIResource
type groupResource struct {
	APIGroup    string
	APIResource metav1.APIResource
}

func NewAPIResourceOptions(ioStreams genericclioptions.IOStreams) *ApiResourcesOptions {
	return &ApiResourcesOptions{
		IOStreams:  ioStreams,
		Namespaced: true,
	}
}

func NewCmdApiResources(f cmdutil.Factory, ioStreams genericclioptions.IOStreams) *cobra.Command {
	o := NewAPIResourceOptions(ioStreams)

	cmd := &cobra.Command{
		Use:     "api-resources",
		Short:   "Print the supported API resources on the server",
		Long:    "Print the supported API resources on the server",
		Example: apiresourcesExample,
		Run: func(cmd *cobra.Command, args []string) {
			cmdutil.CheckErr(o.Validate(cmd))
			cmdutil.CheckErr(o.RunApiResources(cmd, f))
		},
	}

	cmd.Flags().BoolVar(&o.NoHeaders, "no-headers", o.NoHeaders, "When using the default or custom-column output format, don't print headers (default print headers).")
	cmd.Flags().StringVarP(&o.Output, "output", "o", o.Output, "Output format. One of: wide|name.")

	cmd.Flags().StringVar(&o.APIGroup, "api-group", o.APIGroup, "Limit to resources in the specified API group.")
	cmd.Flags().BoolVar(&o.Namespaced, "namespaced", o.Namespaced, "If false, non-namespaced resources will be returned, otherwise returning namespaced resources by default.")
	cmd.Flags().StringSliceVar(&o.Verbs, "verbs", o.Verbs, "Limit to resources that support the specified verbs.")
	cmd.Flags().BoolVar(&o.Cached, "cached", o.Cached, "Use the cached list of resources if available.")
	return cmd
}

func (o *ApiResourcesOptions) Validate(cmd *cobra.Command) error {
	supportedOutputTypes := sets.NewString("", "wide", "name")
	if !supportedOutputTypes.Has(o.Output) {
		return fmt.Errorf("--output %v is not available", o.Output)
	}
	return nil
}

func (o *ApiResourcesOptions) RunApiResources(cmd *cobra.Command, f cmdutil.Factory) error {
	w := printers.GetNewTabWriter(o.Out)
	defer w.Flush()

	discoveryclient, err := f.ToDiscoveryClient()
	if err != nil {
		return err
	}

	if !o.Cached {
		// Always request fresh data from the server
		discoveryclient.Invalidate()
	}

	errs := []error{}
	lists, err := discoveryclient.ServerPreferredResources()
	if err != nil {
		errs = append(errs, err)
	}

	resources := []groupResource{}

	groupChanged := cmd.Flags().Changed("api-group")
	nsChanged := cmd.Flags().Changed("namespaced")

	for _, list := range lists {
		if len(list.APIResources) == 0 {
			continue
		}
		gv, err := schema.ParseGroupVersion(list.GroupVersion)
		if err != nil {
			continue
		}
		for _, resource := range list.APIResources {
			if len(resource.Verbs) == 0 {
				continue
			}
			// filter apiGroup
			if groupChanged && o.APIGroup != gv.Group {
				continue
			}
			// filter namespaced
			if nsChanged && o.Namespaced != resource.Namespaced {
				continue
			}
			// filter to resources that support the specified verbs
			if len(o.Verbs) > 0 && !sets.NewString(resource.Verbs...).HasAll(o.Verbs...) {
				continue
			}
			resources = append(resources, groupResource{
				APIGroup:    gv.Group,
				APIResource: resource,
			})
		}
	}

	if o.NoHeaders == false && o.Output != "name" {
		if err = printContextHeaders(w, o.Output); err != nil {
			return err
		}
	}

	sort.Stable(sortableGroupResource(resources))
	for _, r := range resources {
		switch o.Output {
		case "name":
			name := r.APIResource.Name
			if len(r.APIGroup) > 0 {
				name += "." + r.APIGroup
			}
			if _, err := fmt.Fprintf(w, "%s\n", name); err != nil {
				errs = append(errs, err)
			}
		case "wide":
			if _, err := fmt.Fprintf(w, "%s\t%s\t%s\t%v\t%s\t%v\n",
				r.APIResource.Name,
				strings.Join(r.APIResource.ShortNames, ","),
				r.APIGroup,
				r.APIResource.Namespaced,
				r.APIResource.Kind,
				r.APIResource.Verbs); err != nil {
				errs = append(errs, err)
			}
		case "":
			if _, err := fmt.Fprintf(w, "%s\t%s\t%s\t%v\t%s\n",
				r.APIResource.Name,
				strings.Join(r.APIResource.ShortNames, ","),
				r.APIGroup,
				r.APIResource.Namespaced,
				r.APIResource.Kind); err != nil {
				errs = append(errs, err)
			}
		}
	}

	if len(errs) > 0 {
		return errors.NewAggregate(errs)
	}
	return nil
}

func printContextHeaders(out io.Writer, output string) error {
	columnNames := []string{"NAME", "SHORTNAMES", "APIGROUP", "NAMESPACED", "KIND"}
	if output == "wide" {
		columnNames = append(columnNames, "VERBS")
	}
	_, err := fmt.Fprintf(out, "%s\n", strings.Join(columnNames, "\t"))
	return err
}

type sortableGroupResource []groupResource

func (s sortableGroupResource) Len() int      { return len(s) }
func (s sortableGroupResource) Swap(i, j int) { s[i], s[j] = s[j], s[i] }
func (s sortableGroupResource) Less(i, j int) bool {
	ret := strings.Compare(s[i].APIGroup, s[j].APIGroup)
	if ret > 0 {
		return false
	} else if ret == 0 {
		return strings.Compare(s[i].APIResource.Name, s[j].APIResource.Name) < 0
	}
	return true
}
<|endoftext|>


<commit_before>package simulator

import (
	"bytes"
	"math"
	"math/rand"
	"sort"
	"strconv"

	"github.com/golang/glog"
	"github.com/octo47/tsgen/generator"
)

// Simulator keeps state of simulation.
type Simulator struct {
	conf        Configuration
	currentTime uint64
	machines    []*Machine
	tags        tagsDef
}

// Configuration keeps simulation configuration
type Configuration struct {
	// total machines count
	Machines int
	// group machines by clusters.
	// machine in every cluster will drain unique tags from cluster private set of possible tags
	Clusters int
	// number of tags that will be assigned to every machine
	// value will be unique to machine
	GlobalTags int
	// unique tags are tags assigned per machine
	ClusterTags int
	// minimun tags
	MinimumTags int
	// Total metrics
	MetricsTotal int
	// Set of base metrics that reported by all machines
	BaseMetrics int
	// Maximum metrics per machine
	MaxMetrics int
	// Metrics per namespace
	MetricsPerNamespace int
	// Tags per metric
	TagsPerMetric int
	// TagValuesPerMetricTag configures how deiverse values of metric tags.
	TagValuesPerMetricTag int
	// StartSplay configures how far machines start time could drift away from startTime
	StartSplay int
}

func NewConfiguration(machines int, metrics int, clusters int) Configuration {
	return Configuration{
		Machines:              machines,
		Clusters:              clusters,
		GlobalTags:            8,
		ClusterTags:           clusters/48 + 16, // suppose 48 machines average per service
		MinimumTags:           4,
		MetricsTotal:          int(math.Log10(float64(machines)) * 300),
		BaseMetrics:           metrics/10 + 1,
		MaxMetrics:            metrics,
		MetricsPerNamespace:   metrics/100 + 50,
		TagsPerMetric:         2,
		TagValuesPerMetricTag: 64,
		StartSplay:            300,
	}
}

type clusterDef struct {
	cID     int
	tags    tagsDef
	metrics []metricDef
}

func NewSimulator(rnd *rand.Rand, conf Configuration, startTime uint64) *Simulator {
	clusters := generateClusters(rnd, conf)
	if glog.V(1) {
		glog.Info("Generated ", len(clusters), " clusters")
	}
	clusterGen := rand.NewZipf(rnd, 1.2, 1.1, uint64(conf.Clusters-1))
	globalTags := NewTagsDef(rnd, "global", "gv", conf.GlobalTags, conf.GlobalTags)
	metrics := generateMetrics(rnd, conf)
	if glog.V(1) {
		glog.Info("Generated ", len(metrics), " metrics")
	}
	machines := make([]*Machine, conf.Machines)
	for machine := 0; machine < conf.Machines; machine++ {
		var tags Tags
		machineName := genName("machine-", machine)
		clusterIdx := clusterGen.Uint64()
		tags = append(tags, globalTags.selectTags(conf.MinimumTags)...)
		tags = append(tags, clusters[clusterIdx].tags.selectTags(conf.MinimumTags)...)
		tags = append(tags, Tag{Name: "machine", Value: machineName})
		tags = append(tags, Tag{Name: "rack", Value: strconv.Itoa(machine / 48)})
		metricsCount := rnd.Intn(conf.MaxMetrics - conf.BaseMetrics)
		machines[machine] = NewMachine(machineName, tags,
			startTime+uint64(rnd.Intn(conf.StartSplay)), metricsCount+conf.BaseMetrics)
		if glog.V(1) {
			glog.Info("Machine ", machineName, " cluster=", clusterIdx,
				" tags=", tags, " startTime=", machines[machine].lastTs)
		}
		for metric := 0; metric < conf.BaseMetrics; metric++ {
			gen, name := metrics[metric].genMaker()
			selectedTags := metrics[metric].tags.selectTags(conf.TagsPerMetric)
			machines[machine].AddTimeseriesWithTags(
				metrics[metric].namespace,
				name,
				selectedTags,
				gen,
				metrics[metric].period)
		}
		if glog.V(1) {
			glog.Info("Machine ", machineName, " has ", conf.BaseMetrics, " base metrics")
		}
		metricsSelected := make(map[int]bool)
		for i := 0; i < metricsCount; i++ {
			metricsSelected[rnd.Intn(conf.MaxMetrics-conf.BaseMetrics)+conf.BaseMetrics] = true
		}
		for metric := range metricsSelected {
			gen, name := metrics[metric].genMaker()
			selectedTags := metrics[metric].tags.selectTags(conf.TagsPerMetric)
			glog.Info("Machine ", machineName, " adding metric ", metrics[metric])
			machines[machine].AddTimeseriesWithTags(
				metrics[metric].namespace,
				name,
				selectedTags,
				gen,
				metrics[metric].period)
		}
		if glog.V(1) {
			glog.Info("Machine ", machineName, " has ", len(metricsSelected), " unique metrics")
		}
	}
	return &Simulator{conf, startTime, machines, globalTags}
}

// Run simulator for specified runFor time uints (whatever you use timestamp for, usually seconds)
// Callback will get tagged points usable to be sent to monitoring systems
func (s *Simulator) Run(shard int, shardCount int, runTo int64, cb func(points *[]TaggedPoints)) {
	s.currentTime = uint64(runTo)
	for i := shard; i < len(s.machines); i += shardCount {
		tick := s.machines[i].Tick(s.currentTime)
		if len(*tick) > 0 {
			cb(tick)
		}
	}
}

type tagsDef struct {
	tags              Tags
	tagDistribution   *rand.Zipf
	countDistribution *rand.Zipf
}

type metricDef struct {
	namespace string
	period    uint64
	tags      tagsDef
	genMaker  func() (generator.Generator, string)
}

func NewTagsDef(rnd *rand.Rand, namePrefix, valuePrefix string, maxTags int, maxCount int) tagsDef {
	tags := generateTags(rnd, namePrefix, valuePrefix, maxTags)
	if len(tags) == 0 {
		panic("Tags is zero for " + namePrefix)
	}
	return tagsDef{
		tags:              tags,
		tagDistribution:   rand.NewZipf(rnd, 1.3, 1.1, uint64(maxTags-1)),
		countDistribution: rand.NewZipf(rnd, 1.2, 1.1, uint64(maxCount-1)),
	}
}
func (td *tagsDef) selectTags(minimum int) Tags {
	if len(td.tags) == 0 {
		return nil
	}
	count := int(td.countDistribution.Uint64())
	if count < minimum {
		count = minimum
	}
	tags := make(Tags, 0, count)
	toGo := count
	for toGo > 0 {
		for idx := len(tags); idx < count; idx++ {
			tagIdx := int(td.tagDistribution.Uint64())
			tags = append(tags, td.tags[tagIdx])
		}
		sort.Sort(&tags)
		deduplicateTags(&tags)
		toGo -= len(tags)
	}
	return tags
}

func generateClusters(rnd *rand.Rand, conf Configuration) []clusterDef {
	clusters := make([]clusterDef, conf.Clusters)
	for cID := 0; cID < conf.Clusters; cID++ {
		clusters[cID].cID = cID
		clusters[cID].tags = NewTagsDef(rnd,
			"svc"+strconv.Itoa(cID)+"-", "v", conf.ClusterTags, conf.ClusterTags)
		if glog.V(2) {
			glog.Info("Cluster ", cID, " tags=", clusters[cID].tags.tags)
		}
	}
	return clusters
}

func genName(prefix string, id int) string {
	var buf bytes.Buffer
	buf.WriteString(prefix)
	buf.WriteString(strconv.Itoa(id))
	return buf.String()
}

func genOrCache(cache *map[uint64]string, prefix string, key uint64) string {
	name, ok := (*cache)[key]
	if !ok {
		name = genName(prefix, int(key))
		(*cache)[key] = name
	}
	return name
}

func generateTags(rnd *rand.Rand, namePrefix, valuePrefix string, numTags int) Tags {
	if numTags == 0 {
		panic("numTags shouldn't be zero for " + namePrefix)
	}
	tagsCache := make(map[uint64]string)
	valuesCache := make(map[uint64]string)
	tagsGen := rand.NewZipf(rnd, 1.3, 1.1, uint64(numTags-1))
	valuesGen := rand.NewZipf(rnd, 1.3, 1.1, uint64(numTags-1))
	tags := make(Tags, 0, numTags)
	tagsToGo := numTags
	for tagsToGo > 0 {
		for tagIdx := len(tags); tagIdx < numTags; tagIdx++ {
			tagName := genOrCache(&tagsCache, namePrefix, tagsGen.Uint64())
			value := genOrCache(&valuesCache, valuePrefix, valuesGen.Uint64())
			tags = append(tags, Tag{tagName, value})
		}
		sort.Sort(&tags)
		deduplicateTags(&tags)
		tagsToGo = numTags - len(tags)
	}
	return tags
}

func deduplicateTags(tags *Tags) {
	var newTags Tags
	if len(*tags) < 2 {
		return
	}
	prev := 0
	newTags = append(newTags, (*tags)[prev])
	for i := 1; i < len(*tags); i++ {
		if (*tags)[i].Value != (*tags)[prev].Value {
			prev = i
			newTags = append(newTags, (*tags)[prev])
		}
	}
	*tags = newTags
}

func generateMetrics(rnd *rand.Rand, conf Configuration) []metricDef {
	count := conf.MaxMetrics
	periodDF := rand.NewZipf(rnd, 1.3, 1.2, 19)
	namespaces := count/conf.MetricsPerNamespace + 1
	metrics := make([]metricDef, count)
	for i := 0; i < count; i++ {
		genNum := rnd.Intn(10) // make randomwalk more frequent
		scale := float64(rnd.Intn(1e6)) * 0.001
		ns := rnd.Intn(namespaces)
		metrics[i].namespace = genName("ns", ns)
		metrics[i].period = uint64(15 * (int(periodDF.Uint64()) + 1))
		tagsCount := rnd.Intn(conf.TagsPerMetric * 2)
		if tagsCount > conf.TagsPerMetric/2 {
			metrics[i].tags = NewTagsDef(rnd, "mtag", "mv",
				conf.TagsPerMetric, tagsCount/2)
		}
		metrics[i].genMaker = func() (generator.Generator, string) {
			genNum := genNum
			scale := scale
			var gen generator.Generator
			var metricPrefix string
			switch genNum {
			case 4:
				gen = generator.NewIncreasingGenerator(rnd, 0.8, 0.01, 0.1, 0.0, 100.0)
				metricPrefix = "metricI"
			case 1:
				gen = generator.NewSpikesGenerator(rnd, 100.0)
				metricPrefix = "metricS"
			case 3:
				gen = generator.NewCyclicGenerator(rnd, 0.0, 100.0)
				metricPrefix = "metricC"
			default:
				gen = generator.NewRandomWalkGenerator(rnd, 1.0, 0.0, 100)
				metricPrefix = "metricR"
			}
			return generator.NewScalingGenerator(gen, scale), genName(metricPrefix, i)
		}
	}
	return metrics
}
<commit_msg>making less logs<commit_after>package simulator

import (
	"bytes"
	"math"
	"math/rand"
	"sort"
	"strconv"

	"github.com/golang/glog"
	"github.com/octo47/tsgen/generator"
)

// Simulator keeps state of simulation.
type Simulator struct {
	conf        Configuration
	currentTime uint64
	machines    []*Machine
	tags        tagsDef
}

// Configuration keeps simulation configuration
type Configuration struct {
	// total machines count
	Machines int
	// group machines by clusters.
	// machine in every cluster will drain unique tags from cluster private set of possible tags
	Clusters int
	// number of tags that will be assigned to every machine
	// value will be unique to machine
	GlobalTags int
	// unique tags are tags assigned per machine
	ClusterTags int
	// minimun tags
	MinimumTags int
	// Total metrics
	MetricsTotal int
	// Set of base metrics that reported by all machines
	BaseMetrics int
	// Maximum metrics per machine
	MaxMetrics int
	// Metrics per namespace
	MetricsPerNamespace int
	// Tags per metric
	TagsPerMetric int
	// TagValuesPerMetricTag configures how deiverse values of metric tags.
	TagValuesPerMetricTag int
	// StartSplay configures how far machines start time could drift away from startTime
	StartSplay int
}

func NewConfiguration(machines int, metrics int, clusters int) Configuration {
	return Configuration{
		Machines:              machines,
		Clusters:              clusters,
		GlobalTags:            8,
		ClusterTags:           clusters/48 + 16, // suppose 48 machines average per service
		MinimumTags:           4,
		MetricsTotal:          int(math.Log10(float64(machines)) * 300),
		BaseMetrics:           metrics/10 + 1,
		MaxMetrics:            metrics,
		MetricsPerNamespace:   metrics/100 + 50,
		TagsPerMetric:         2,
		TagValuesPerMetricTag: 64,
		StartSplay:            300,
	}
}

type clusterDef struct {
	cID     int
	tags    tagsDef
	metrics []metricDef
}

func NewSimulator(rnd *rand.Rand, conf Configuration, startTime uint64) *Simulator {
	clusters := generateClusters(rnd, conf)
	if glog.V(1) {
		glog.Info("Generated ", len(clusters), " clusters")
	}
	clusterGen := rand.NewZipf(rnd, 1.2, 1.1, uint64(conf.Clusters-1))
	globalTags := NewTagsDef(rnd, "global", "gv", conf.GlobalTags, conf.GlobalTags)
	metrics := generateMetrics(rnd, conf)
	if glog.V(1) {
		glog.Info("Generated ", len(metrics), " metrics")
	}
	machines := make([]*Machine, conf.Machines)
	for machine := 0; machine < conf.Machines; machine++ {
		var tags Tags
		machineName := genName("machine-", machine)
		clusterIdx := clusterGen.Uint64()
		tags = append(tags, globalTags.selectTags(conf.MinimumTags)...)
		tags = append(tags, clusters[clusterIdx].tags.selectTags(conf.MinimumTags)...)
		tags = append(tags, Tag{Name: "machine", Value: machineName})
		tags = append(tags, Tag{Name: "rack", Value: strconv.Itoa(machine / 48)})
		metricsCount := rnd.Intn(conf.MaxMetrics - conf.BaseMetrics)
		machines[machine] = NewMachine(machineName, tags,
			startTime+uint64(rnd.Intn(conf.StartSplay)), metricsCount+conf.BaseMetrics)
		if glog.V(1) {
			glog.Info("Machine ", machineName, " cluster=", clusterIdx,
				" tags=", tags, " startTime=", machines[machine].lastTs)
		}
		for metric := 0; metric < conf.BaseMetrics; metric++ {
			gen, name := metrics[metric].genMaker()
			selectedTags := metrics[metric].tags.selectTags(conf.TagsPerMetric)
			machines[machine].AddTimeseriesWithTags(
				metrics[metric].namespace,
				name,
				selectedTags,
				gen,
				metrics[metric].period)
		}
		if glog.V(1) {
			glog.Info("Machine ", machineName, " has ", conf.BaseMetrics, " base metrics")
		}
		metricsSelected := make(map[int]bool)
		for i := 0; i < metricsCount; i++ {
			metricsSelected[rnd.Intn(conf.MaxMetrics-conf.BaseMetrics)+conf.BaseMetrics] = true
		}
		for metric := range metricsSelected {
			gen, name := metrics[metric].genMaker()
			selectedTags := metrics[metric].tags.selectTags(conf.TagsPerMetric)
			if glog.V(2) {
				glog.Info("Machine ", machineName, " adding metric ", metrics[metric])
			}
			machines[machine].AddTimeseriesWithTags(
				metrics[metric].namespace,
				name,
				selectedTags,
				gen,
				metrics[metric].period)
		}
		if glog.V(1) {
			glog.Info("Machine ", machineName, " has ", len(metricsSelected), " unique metrics")
		}
	}
	return &Simulator{conf, startTime, machines, globalTags}
}

// Run simulator for specified runFor time uints (whatever you use timestamp for, usually seconds)
// Callback will get tagged points usable to be sent to monitoring systems
func (s *Simulator) Run(shard int, shardCount int, runTo int64, cb func(points *[]TaggedPoints)) {
	s.currentTime = uint64(runTo)
	for i := shard; i < len(s.machines); i += shardCount {
		tick := s.machines[i].Tick(s.currentTime)
		if len(*tick) > 0 {
			cb(tick)
		}
	}
}

type tagsDef struct {
	tags              Tags
	tagDistribution   *rand.Zipf
	countDistribution *rand.Zipf
}

type metricDef struct {
	namespace string
	period    uint64
	tags      tagsDef
	genMaker  func() (generator.Generator, string)
}

func NewTagsDef(rnd *rand.Rand, namePrefix, valuePrefix string, maxTags int, maxCount int) tagsDef {
	tags := generateTags(rnd, namePrefix, valuePrefix, maxTags)
	if len(tags) == 0 {
		panic("Tags is zero for " + namePrefix)
	}
	return tagsDef{
		tags:              tags,
		tagDistribution:   rand.NewZipf(rnd, 1.3, 1.1, uint64(maxTags-1)),
		countDistribution: rand.NewZipf(rnd, 1.2, 1.1, uint64(maxCount-1)),
	}
}
func (td *tagsDef) selectTags(minimum int) Tags {
	if len(td.tags) == 0 {
		return nil
	}
	count := int(td.countDistribution.Uint64())
	if count < minimum {
		count = minimum
	}
	tags := make(Tags, 0, count)
	toGo := count
	for toGo > 0 {
		for idx := len(tags); idx < count; idx++ {
			tagIdx := int(td.tagDistribution.Uint64())
			tags = append(tags, td.tags[tagIdx])
		}
		sort.Sort(&tags)
		deduplicateTags(&tags)
		toGo -= len(tags)
	}
	return tags
}

func generateClusters(rnd *rand.Rand, conf Configuration) []clusterDef {
	clusters := make([]clusterDef, conf.Clusters)
	for cID := 0; cID < conf.Clusters; cID++ {
		clusters[cID].cID = cID
		clusters[cID].tags = NewTagsDef(rnd,
			"svc"+strconv.Itoa(cID)+"-", "v", conf.ClusterTags, conf.ClusterTags)
		if glog.V(2) {
			glog.Info("Cluster ", cID, " tags=", clusters[cID].tags.tags)
		}
	}
	return clusters
}

func genName(prefix string, id int) string {
	var buf bytes.Buffer
	buf.WriteString(prefix)
	buf.WriteString(strconv.Itoa(id))
	return buf.String()
}

func genOrCache(cache *map[uint64]string, prefix string, key uint64) string {
	name, ok := (*cache)[key]
	if !ok {
		name = genName(prefix, int(key))
		(*cache)[key] = name
	}
	return name
}

func generateTags(rnd *rand.Rand, namePrefix, valuePrefix string, numTags int) Tags {
	if numTags == 0 {
		panic("numTags shouldn't be zero for " + namePrefix)
	}
	tagsCache := make(map[uint64]string)
	valuesCache := make(map[uint64]string)
	tagsGen := rand.NewZipf(rnd, 1.3, 1.1, uint64(numTags-1))
	valuesGen := rand.NewZipf(rnd, 1.3, 1.1, uint64(numTags-1))
	tags := make(Tags, 0, numTags)
	tagsToGo := numTags
	for tagsToGo > 0 {
		for tagIdx := len(tags); tagIdx < numTags; tagIdx++ {
			tagName := genOrCache(&tagsCache, namePrefix, tagsGen.Uint64())
			value := genOrCache(&valuesCache, valuePrefix, valuesGen.Uint64())
			tags = append(tags, Tag{tagName, value})
		}
		sort.Sort(&tags)
		deduplicateTags(&tags)
		tagsToGo = numTags - len(tags)
	}
	return tags
}

func deduplicateTags(tags *Tags) {
	var newTags Tags
	if len(*tags) < 2 {
		return
	}
	prev := 0
	newTags = append(newTags, (*tags)[prev])
	for i := 1; i < len(*tags); i++ {
		if (*tags)[i].Value != (*tags)[prev].Value {
			prev = i
			newTags = append(newTags, (*tags)[prev])
		}
	}
	*tags = newTags
}

func generateMetrics(rnd *rand.Rand, conf Configuration) []metricDef {
	count := conf.MaxMetrics
	periodDF := rand.NewZipf(rnd, 1.3, 1.2, 19)
	namespaces := count/conf.MetricsPerNamespace + 1
	metrics := make([]metricDef, count)
	for i := 0; i < count; i++ {
		genNum := rnd.Intn(10) // make randomwalk more frequent
		scale := float64(rnd.Intn(1e6)) * 0.001
		ns := rnd.Intn(namespaces)
		metrics[i].namespace = genName("ns", ns)
		metrics[i].period = uint64(15 * (int(periodDF.Uint64()) + 1))
		tagsCount := rnd.Intn(conf.TagsPerMetric * 2)
		if tagsCount > conf.TagsPerMetric/2 {
			metrics[i].tags = NewTagsDef(rnd, "mtag", "mv",
				conf.TagsPerMetric, tagsCount/2)
		}
		metrics[i].genMaker = func() (generator.Generator, string) {
			genNum := genNum
			scale := scale
			var gen generator.Generator
			var metricPrefix string
			switch genNum {
			case 4:
				gen = generator.NewIncreasingGenerator(rnd, 0.8, 0.01, 0.1, 0.0, 100.0)
				metricPrefix = "metricI"
			case 1:
				gen = generator.NewSpikesGenerator(rnd, 100.0)
				metricPrefix = "metricS"
			case 3:
				gen = generator.NewCyclicGenerator(rnd, 0.0, 100.0)
				metricPrefix = "metricC"
			default:
				gen = generator.NewRandomWalkGenerator(rnd, 1.0, 0.0, 100)
				metricPrefix = "metricR"
			}
			return generator.NewScalingGenerator(gen, scale), genName(metricPrefix, i)
		}
	}
	return metrics
}
<|endoftext|>


<commit_before>package cmd

// Version defines the current Pop version.
const Version = "development"
<commit_msg>Bump version<commit_after>package cmd

// Version defines the current Pop version.
const Version = "v4.10.0"
<|endoftext|>


<commit_before>// Copyright (c) 2017 Cisco and/or its affiliates.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at:
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package etcd

import (
	"fmt"
	"time"

	"github.com/ligato/cn-infra/core"
	"github.com/ligato/cn-infra/datasync/resync"
	"github.com/ligato/cn-infra/db/keyval"
	"github.com/ligato/cn-infra/db/keyval/kvproto"
	"github.com/ligato/cn-infra/flavors/local"
	"github.com/ligato/cn-infra/health/statuscheck"
	"github.com/ligato/cn-infra/utils/safeclose"
)

const (
	// healthCheckProbeKey is a key used to probe Etcd state
	healthCheckProbeKey = "/probe-etcd-connection"
	// ETCD reconnect interval
	defaultReconnectInterval time.Duration = 2000000000 // 2 seconds
)

// Plugin implements etcd plugin.
type Plugin struct {
	Deps
	// Plugin is disabled if there is no config file available
	disabled bool
	// Set if connected to ETCD db
	connected bool
	// ETCD connection encapsulation
	connection *BytesConnectionEtcd
	// Read/Write proto modelled data
	protoWrapper *kvproto.ProtoWrapper

	// plugin config
	config *Config

	// List of callback functions, used in case ETCD is not connected immediately. All plugins using
	// ETCD as dependency add their own function if cluster is not reachable. After connection, all
	// functions are executed.
	onConnection []func() error

	autoCompactDone chan struct{}
	lastConnErr     error
}

// Deps lists dependencies of the etcd plugin.
// If injected, etcd plugin will use StatusCheck to signal the connection status.
type Deps struct {
	local.PluginInfraDeps
	Resync *resync.Plugin
}

// Init retrieves ETCD configuration and establishes a new connection
// with the etcd data store.
// If the configuration file doesn't exist or cannot be read, the returned error
// will be of os.PathError type. An untyped error is returned in case the file
// doesn't contain a valid YAML configuration.
// The function may also return error if TLS connection is selected and the
// CA or client certificate is not accessible(os.PathError)/valid(untyped).
// Check clientv3.New from coreos/etcd for possible errors returned in case
// the connection cannot be established.
func (plugin *Plugin) Init() (err error) {
	// Read ETCD configuration file. Returns error if does not exists.
	plugin.config, err = plugin.getEtcdConfig()
	if err != nil || plugin.disabled {
		return err
	}
	// Transforms .yaml config to ETCD client configuration
	etcdClientCfg, err := ConfigToClient(plugin.config)
	if err != nil {
		return err
	}
	// Uses config file to establish connection with the database
	plugin.connection, err = NewEtcdConnectionWithBytes(*etcdClientCfg, plugin.Log)
	// Register for providing status reports (polling mode).
	if plugin.StatusCheck != nil {
		plugin.StatusCheck.Register(plugin.PluginName, plugin.statusCheckProbe)
	} else {
		plugin.Log.Warnf("Unable to start status check for etcd")
	}
	if err != nil && plugin.config.AllowDelayedStart {
		// If the connection cannot be established during init, keep trying in another goroutine (if allowed) and
		// end the init
		go plugin.etcdReconnectionLoop(etcdClientCfg)
		return nil
	} else if err != nil {
		// If delayed start is not allowed, return error
		return fmt.Errorf("error connecting to ETCD: %v", err)
	}

	// If successful, configure and return
	plugin.configureConnection()

	// Mark plugin as connected at this point
	plugin.connected = true

	return nil
}

// Close shutdowns the connection.
func (plugin *Plugin) Close() error {
	err := safeclose.Close(plugin.autoCompactDone)
	return err
}

// NewBroker creates new instance of prefixed broker that provides API with arguments of type proto.Message.
func (plugin *Plugin) NewBroker(keyPrefix string) keyval.ProtoBroker {
	return plugin.protoWrapper.NewBroker(keyPrefix)
}

// NewWatcher creates new instance of prefixed broker that provides API with arguments of type proto.Message.
func (plugin *Plugin) NewWatcher(keyPrefix string) keyval.ProtoWatcher {
	return plugin.protoWrapper.NewWatcher(keyPrefix)
}

// Disabled returns *true* if the plugin is not in use due to missing
// etcd configuration.
func (plugin *Plugin) Disabled() (disabled bool) {
	return plugin.disabled
}

// OnConnect executes callback if plugin is connected, or gathers functions from all plugin with ETCD as dependency
func (plugin *Plugin) OnConnect(callback func() error) {
	if plugin.connected {
		if err := callback(); err != nil {
			plugin.Log.Error(err)
		}
	} else {
		plugin.onConnection = append(plugin.onConnection, callback)
	}
}

// GetPluginName returns name of the plugin
func (plugin *Plugin) GetPluginName() core.PluginName {
	return plugin.PluginName
}

// PutIfNotExists puts given key-value pair into etcd if there is no value set for the key. If the put was successful
// succeeded is true. If the key already exists succeeded is false and the value for the key is untouched.
func (plugin *Plugin) PutIfNotExists(key string, value []byte) (succeeded bool, err error) {
	if plugin.connection != nil {
		return plugin.connection.PutIfNotExists(key, value)
	}
	return false, fmt.Errorf("connection is not established")
}

// Compact compatcs the ETCD database to the specific revision
func (plugin *Plugin) Compact(rev ...int64) (toRev int64, err error) {
	if plugin.connection != nil {
		return plugin.connection.Compact(rev...)
	}
	return 0, fmt.Errorf("connection is not established")
}

// Method starts loop which attempt to connect to the ETCD. If successful, send signal callback with resync,
// which will be started when datasync confirms successful registration
func (plugin *Plugin) etcdReconnectionLoop(clientCfg *ClientConfig) {
	var err error
	// Set reconnect interval
	interval := plugin.config.ReconnectInterval
	if interval == 0 {
		interval = defaultReconnectInterval
	}
	plugin.Log.Infof("ETCD server %s not reachable in init phase. Agent will continue to try to connect every %d second(s)",
		plugin.config.Endpoints, interval)
	for {
		time.Sleep(interval)

		plugin.Log.Infof("Connecting to ETCD %v ...", plugin.config.Endpoints)
		plugin.connection, err = NewEtcdConnectionWithBytes(*clientCfg, plugin.Log)
		if err != nil {
			continue
		}
		plugin.Log.Infof("ETCD server %s connected", plugin.config.Endpoints)
		// Configure connection and set as connected
		plugin.configureConnection()
		plugin.connected = true
		// Execute callback functions (if any)
		for _, callback := range plugin.onConnection {
			if err := callback(); err != nil {
				plugin.Log.Error(err)
			}
		}
		// Call resync if any callback was executed. Otherwise there is nothing to resync
		if plugin.Resync != nil && len(plugin.onConnection) > 0 {
			plugin.Resync.DoResync()
		}
		plugin.Log.Debugf("Etcd reconnection loop ended")

		return
	}
}

// If ETCD is connected, complete all other procedures
func (plugin *Plugin) configureConnection() {
	if plugin.config.AutoCompact > 0 {
		if plugin.config.AutoCompact < time.Duration(time.Minute*60) {
			plugin.Log.Warnf("Auto compact option for ETCD is set to less than 60 minutes!")
		}
		plugin.startPeriodicAutoCompact(plugin.config.AutoCompact)
	}
	plugin.protoWrapper = kvproto.NewProtoWrapperWithSerializer(plugin.connection, &keyval.SerializerJSON{})
}

// ETCD status check probe function
func (plugin *Plugin) statusCheckProbe() (statuscheck.PluginState, error) {
	if plugin.connection == nil {
		plugin.connected = false
		return statuscheck.Error, fmt.Errorf("no ETCD connection available")
	}
	if _, _, _, err := plugin.connection.GetValue(healthCheckProbeKey); err != nil {
		plugin.lastConnErr = err
		plugin.connected = false
		return statuscheck.Error, err
	}
	if plugin.config.ReconnectResync && plugin.lastConnErr != nil {
		if plugin.Resync != nil {
			plugin.Resync.DoResync()
			plugin.lastConnErr = nil
		} else {
			plugin.Log.Warn("Expected resync after ETCD reconnect could not start beacuse of missing Resync plugin")
		}
	}
	plugin.connected = true
	return statuscheck.OK, nil
}

func (plugin *Plugin) getEtcdConfig() (*Config, error) {
	var etcdCfg Config
	found, err := plugin.PluginConfig.GetValue(&etcdCfg)
	if err != nil {
		return &etcdCfg, err
	}
	if !found {
		plugin.Log.Info("ETCD config not found, skip loading this plugin")
		plugin.disabled = true
		return &etcdCfg, nil
	}
	return &etcdCfg, nil
}

func (plugin *Plugin) startPeriodicAutoCompact(period time.Duration) {
	plugin.autoCompactDone = make(chan struct{})
	go func() {
		plugin.Log.Infof("Starting periodic auto compacting every %v", period)
		for {
			select {
			case <-time.After(period):
				plugin.Log.Debugf("Executing auto compact")
				if toRev, err := plugin.connection.Compact(); err != nil {
					plugin.Log.Errorf("Periodic auto compacting failed: %v", err)
				} else {
					plugin.Log.Infof("Auto compacting finished (to revision %v)", toRev)
				}
			case <-plugin.autoCompactDone:
				return
			}
		}
	}()
}
<commit_msg>address comments<commit_after>// Copyright (c) 2017 Cisco and/or its affiliates.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at:
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package etcd

import (
	"fmt"
	"time"

	"github.com/ligato/cn-infra/core"
	"github.com/ligato/cn-infra/datasync/resync"
	"github.com/ligato/cn-infra/db/keyval"
	"github.com/ligato/cn-infra/db/keyval/kvproto"
	"github.com/ligato/cn-infra/flavors/local"
	"github.com/ligato/cn-infra/health/statuscheck"
	"github.com/ligato/cn-infra/utils/safeclose"
)

const (
	// healthCheckProbeKey is a key used to probe Etcd state
	healthCheckProbeKey = "/probe-etcd-connection"
	// ETCD reconnect interval
	defaultReconnectInterval = 2 * time.Second // 2 seconds
)

// Plugin implements etcd plugin.
type Plugin struct {
	Deps
	// Plugin is disabled if there is no config file available
	disabled bool
	// Set if connected to ETCD db
	connected bool
	// ETCD connection encapsulation
	connection *BytesConnectionEtcd
	// Read/Write proto modelled data
	protoWrapper *kvproto.ProtoWrapper

	// plugin config
	config *Config

	// List of callback functions, used in case ETCD is not connected immediately. All plugins using
	// ETCD as dependency add their own function if cluster is not reachable. After connection, all
	// functions are executed.
	onConnection []func() error

	autoCompactDone chan struct{}
	lastConnErr     error
}

// Deps lists dependencies of the etcd plugin.
// If injected, etcd plugin will use StatusCheck to signal the connection status.
type Deps struct {
	local.PluginInfraDeps
	Resync *resync.Plugin
}

// Init retrieves ETCD configuration and establishes a new connection
// with the etcd data store.
// If the configuration file doesn't exist or cannot be read, the returned error
// will be of os.PathError type. An untyped error is returned in case the file
// doesn't contain a valid YAML configuration.
// The function may also return error if TLS connection is selected and the
// CA or client certificate is not accessible(os.PathError)/valid(untyped).
// Check clientv3.New from coreos/etcd for possible errors returned in case
// the connection cannot be established.
func (plugin *Plugin) Init() (err error) {
	// Read ETCD configuration file. Returns error if does not exists.
	plugin.config, err = plugin.getEtcdConfig()
	if err != nil || plugin.disabled {
		return err
	}
	// Transforms .yaml config to ETCD client configuration
	etcdClientCfg, err := ConfigToClient(plugin.config)
	if err != nil {
		return err
	}
	// Uses config file to establish connection with the database
	plugin.connection, err = NewEtcdConnectionWithBytes(*etcdClientCfg, plugin.Log)
	// Register for providing status reports (polling mode).
	if plugin.StatusCheck != nil {
		plugin.StatusCheck.Register(plugin.PluginName, plugin.statusCheckProbe)
	} else {
		plugin.Log.Warnf("Unable to start status check for etcd")
	}
	if err != nil && plugin.config.AllowDelayedStart {
		// If the connection cannot be established during init, keep trying in another goroutine (if allowed) and
		// end the init
		go plugin.etcdReconnectionLoop(etcdClientCfg)
		return nil
	} else if err != nil {
		// If delayed start is not allowed, return error
		return fmt.Errorf("error connecting to ETCD: %v", err)
	}

	// If successful, configure and return
	plugin.configureConnection()

	// Mark plugin as connected at this point
	plugin.connected = true

	return nil
}

// Close shutdowns the connection.
func (plugin *Plugin) Close() error {
	err := safeclose.Close(plugin.autoCompactDone)
	return err
}

// NewBroker creates new instance of prefixed broker that provides API with arguments of type proto.Message.
func (plugin *Plugin) NewBroker(keyPrefix string) keyval.ProtoBroker {
	return plugin.protoWrapper.NewBroker(keyPrefix)
}

// NewWatcher creates new instance of prefixed broker that provides API with arguments of type proto.Message.
func (plugin *Plugin) NewWatcher(keyPrefix string) keyval.ProtoWatcher {
	return plugin.protoWrapper.NewWatcher(keyPrefix)
}

// Disabled returns *true* if the plugin is not in use due to missing
// etcd configuration.
func (plugin *Plugin) Disabled() (disabled bool) {
	return plugin.disabled
}

// OnConnect executes callback if plugin is connected, or gathers functions from all plugin with ETCD as dependency
func (plugin *Plugin) OnConnect(callback func() error) {
	if plugin.connected {
		if err := callback(); err != nil {
			plugin.Log.Error(err)
		}
	} else {
		plugin.onConnection = append(plugin.onConnection, callback)
	}
}

// GetPluginName returns name of the plugin
func (plugin *Plugin) GetPluginName() core.PluginName {
	return plugin.PluginName
}

// PutIfNotExists puts given key-value pair into etcd if there is no value set for the key. If the put was successful
// succeeded is true. If the key already exists succeeded is false and the value for the key is untouched.
func (plugin *Plugin) PutIfNotExists(key string, value []byte) (succeeded bool, err error) {
	if plugin.connection != nil {
		return plugin.connection.PutIfNotExists(key, value)
	}
	return false, fmt.Errorf("connection is not established")
}

// Compact compatcs the ETCD database to the specific revision
func (plugin *Plugin) Compact(rev ...int64) (toRev int64, err error) {
	if plugin.connection != nil {
		return plugin.connection.Compact(rev...)
	}
	return 0, fmt.Errorf("connection is not established")
}

// Method starts loop which attempt to connect to the ETCD. If successful, send signal callback with resync,
// which will be started when datasync confirms successful registration
func (plugin *Plugin) etcdReconnectionLoop(clientCfg *ClientConfig) {
	var err error
	// Set reconnect interval
	interval := plugin.config.ReconnectInterval
	if interval == 0 {
		interval = defaultReconnectInterval
	}
	plugin.Log.Infof("ETCD server %s not reachable in init phase. Agent will continue to try to connect every %d second(s)",
		plugin.config.Endpoints, interval)
	for {
		time.Sleep(interval)

		plugin.Log.Infof("Connecting to ETCD %v ...", plugin.config.Endpoints)
		plugin.connection, err = NewEtcdConnectionWithBytes(*clientCfg, plugin.Log)
		if err != nil {
			continue
		}
		plugin.Log.Infof("ETCD server %s connected", plugin.config.Endpoints)
		// Configure connection and set as connected
		plugin.configureConnection()
		plugin.connected = true
		// Execute callback functions (if any)
		for _, callback := range plugin.onConnection {
			if err := callback(); err != nil {
				plugin.Log.Error(err)
			}
		}
		// Call resync if any callback was executed. Otherwise there is nothing to resync
		if plugin.Resync != nil && len(plugin.onConnection) > 0 {
			plugin.Resync.DoResync()
		}
		plugin.Log.Debugf("Etcd reconnection loop ended")

		return
	}
}

// If ETCD is connected, complete all other procedures
func (plugin *Plugin) configureConnection() {
	if plugin.config.AutoCompact > 0 {
		if plugin.config.AutoCompact < time.Duration(time.Minute*60) {
			plugin.Log.Warnf("Auto compact option for ETCD is set to less than 60 minutes!")
		}
		plugin.startPeriodicAutoCompact(plugin.config.AutoCompact)
	}
	plugin.protoWrapper = kvproto.NewProtoWrapperWithSerializer(plugin.connection, &keyval.SerializerJSON{})
}

// ETCD status check probe function
func (plugin *Plugin) statusCheckProbe() (statuscheck.PluginState, error) {
	if plugin.connection == nil {
		plugin.connected = false
		return statuscheck.Error, fmt.Errorf("no ETCD connection available")
	}
	if _, _, _, err := plugin.connection.GetValue(healthCheckProbeKey); err != nil {
		plugin.lastConnErr = err
		plugin.connected = false
		return statuscheck.Error, err
	}
	if plugin.config.ReconnectResync && plugin.lastConnErr != nil {
		if plugin.Resync != nil {
			plugin.Resync.DoResync()
			plugin.lastConnErr = nil
		} else {
			plugin.Log.Warn("Expected resync after ETCD reconnect could not start beacuse of missing Resync plugin")
		}
	}
	plugin.connected = true
	return statuscheck.OK, nil
}

func (plugin *Plugin) getEtcdConfig() (*Config, error) {
	var etcdCfg Config
	found, err := plugin.PluginConfig.GetValue(&etcdCfg)
	if err != nil {
		return nil, err
	}
	if !found {
		plugin.Log.Info("ETCD config not found, skip loading this plugin")
		plugin.disabled = true
		return &etcdCfg, nil
	}
	return &etcdCfg, nil
}

func (plugin *Plugin) startPeriodicAutoCompact(period time.Duration) {
	plugin.autoCompactDone = make(chan struct{})
	go func() {
		plugin.Log.Infof("Starting periodic auto compacting every %v", period)
		for {
			select {
			case <-time.After(period):
				plugin.Log.Debugf("Executing auto compact")
				if toRev, err := plugin.connection.Compact(); err != nil {
					plugin.Log.Errorf("Periodic auto compacting failed: %v", err)
				} else {
					plugin.Log.Infof("Auto compacting finished (to revision %v)", toRev)
				}
			case <-plugin.autoCompactDone:
				return
			}
		}
	}()
}
<|endoftext|>


<commit_before>package main

import (
	//"github.com/bjeanes/go-lifx"

	"flag"
	"strconv"
	"strings"
	"time"

	log "github.com/cihub/seelog"
	"github.com/lucasb-eyer/go-colorful"
	"github.com/stamp/go-lifx/client"
	"github.com/stampzilla/stampzilla-go/nodes/basenode"
	"github.com/stampzilla/stampzilla-go/protocol"
)

// GLOBAL VARS
var node *protocol.Node
var state *State
var serverSendChannel chan interface{}
var serverRecvChannel chan protocol.Command
var lifxClient *client.Client

func main() { // {{{
	// Parse all commandline arguments, host and port parameters are added in the basenode init function
	flag.Parse()

	//Get a config with the correct parameters
	config := basenode.NewConfig()

	//Activate the config
	basenode.SetConfig(config)

	node = protocol.NewNode("lifx")

	//Start communication with the server
	connection := basenode.Connect()

	// This worker keeps track on our connection state, if we are connected or not
	go monitorState(node, connection)

	// This worker recives all incomming commands
	go serverRecv(connection)

	// Create new node description
	state = NewState()
	node.SetState(state)

	lifxClient = client.New()
	go discoverWorker(lifxClient, connection)

	select {}

} /*}}}*/

func discoverWorker(client *client.Client, connection basenode.Connection) {

	go monitorLampCollection(client.Lights, connection)

	for {
		log.Info("Try to discover bulbs:")
		for _ = range client.Discover() {
		}

		<-time.After(60 * time.Second)
	}
}

// WORKER that monitors the list of lamps
func monitorLampCollection(lights *client.LightCollection, connection basenode.Connection) {
	for s := range lights.Lights.Changes {
		switch s.Event {
		case client.LampAdded:
			log.Warnf("Added: %s (%s)", s.Lamp.Id(), s.Lamp.Label())

			state.AddDevice(s.Lamp)

			node.AddElement(&protocol.Element{
				Type: protocol.ElementTypeToggle,
				Name: s.Lamp.Label(),
				Command: &protocol.Command{
					Cmd:  "set",
					Args: []string{s.Lamp.Id()},
				},
				Feedback: "Devices[2].State",
			})
			node.AddElement(&protocol.Element{
				Type: protocol.ElementTypeColorPicker,
				Name: s.Lamp.Label(),
				Command: &protocol.Command{
					Cmd:  "color",
					Args: []string{s.Lamp.Id()},
				},
				Feedback: "Devices[4].State",
			})
			//log.Infof("Collection: %#v", lights.Lights)
			connection.Send(node)

		case client.LampUpdated:
			log.Warnf("Changed: %s (%s)", s.Lamp.Id(), s.Lamp.Label())
		case client.LampRemoved:
			log.Warnf("Removed: %s (%s)", s.Lamp.Id(), s.Lamp.Label())
		default:
			log.Infof("Received unknown event: %d", s.Event)
		}
	}
}

// WORKER that monitors the current connection state
func monitorState(node *protocol.Node, connection basenode.Connection) {
	for s := range connection.State() {
		switch s {
		case basenode.ConnectionStateConnected:
			connection.Send(node)
		case basenode.ConnectionStateDisconnected:
		}
	}
}

// WORKER that recives all incoming commands
func serverRecv(connection basenode.Connection) {
	for d := range connection.Receive() {
		processCommand(d)
	}
}

// This is called on each incoming command
func processCommand(cmd protocol.Command) {
	log.Info("Incoming command from server:", cmd)

	if len(cmd.Args) < 1 {
		return
	}

	lamp, err := lifxClient.Lights.GetById(cmd.Args[0])
	if err != nil {
		log.Error(err)
		return
	}

	switch cmd.Cmd {
	case "set":
		lamp.TurnOn()
		if len(cmd.Args) > 1 {
			if strings.Index(cmd.Args[1], "#") != 0 {
				cmd.Args[1] = "#" + cmd.Args[1]
			}

			c, err := colorful.Hex(cmd.Args[1])
			if err != nil {
				log.Error("Failed to decode color: ", err)
				return
			}

			h, s, v := c.Hsv()

			if len(cmd.Args) > 2 {
				duration, err := strconv.Atoi(cmd.Args[2])
				if err == nil {
					if len(cmd.Args) > 3 {
						kelvin, err := strconv.Atoi(cmd.Args[3])
						if err == nil {
							lamp.SetState(h, s, v, uint32(duration), uint32(kelvin))
							return
						}
					}
					lamp.SetState(h, s, v, uint32(duration), 6500)
					return
				}
				log.Error("Failed to decode duration: ", err)
			}

			lamp.SetState(h, s, v, 1000, 6500)
		}
	case "on":
		lamp.TurnOn()
	case "off":
		lamp.TurnOff()
	case "color":
		c, err := colorful.Hex(cmd.Params[0])
		if err != nil {
			log.Error(err)
			return
		}

		h, s, v := c.Hsv()

		lamp.SetState(h, s, v, 1000, 6500)
	}
}
<commit_msg>Added error handling when creating a lifx instance<commit_after>package main

import (
	//"github.com/bjeanes/go-lifx"

	"flag"
	"strconv"
	"strings"
	"time"

	log "github.com/cihub/seelog"
	"github.com/lucasb-eyer/go-colorful"
	"github.com/stamp/go-lifx/client"
	"github.com/stampzilla/stampzilla-go/nodes/basenode"
	"github.com/stampzilla/stampzilla-go/protocol"
)

// GLOBAL VARS
var node *protocol.Node
var state *State
var serverSendChannel chan interface{}
var serverRecvChannel chan protocol.Command
var lifxClient *client.Client

func main() { // {{{
	// Parse all commandline arguments, host and port parameters are added in the basenode init function
	flag.Parse()

	//Get a config with the correct parameters
	config := basenode.NewConfig()

	//Activate the config
	basenode.SetConfig(config)

	node = protocol.NewNode("lifx")

	//Start communication with the server
	connection := basenode.Connect()

	// This worker keeps track on our connection state, if we are connected or not
	go monitorState(node, connection)

	// This worker recives all incomming commands
	go serverRecv(connection)

	// Create new node description
	state = NewState()
	node.SetState(state)

	var err error
	lifxClient, err = client.New()
	if err != nil {
		log.Error(err)
		return
	}
	go discoverWorker(lifxClient, connection)

	select {}

} /*}}}*/

func discoverWorker(client *client.Client, connection basenode.Connection) {

	go monitorLampCollection(client.Lights, connection)

	for {
		log.Info("Try to discover bulbs:")
		for _ = range client.Discover() {
		}

		<-time.After(60 * time.Second)
	}
}

// WORKER that monitors the list of lamps
func monitorLampCollection(lights *client.LightCollection, connection basenode.Connection) {
	for s := range lights.Lights.Changes {
		switch s.Event {
		case client.LampAdded:
			log.Warnf("Added: %s (%s)", s.Lamp.Id(), s.Lamp.Label())

			state.AddDevice(s.Lamp)

			node.AddElement(&protocol.Element{
				Type: protocol.ElementTypeToggle,
				Name: s.Lamp.Label(),
				Command: &protocol.Command{
					Cmd:  "set",
					Args: []string{s.Lamp.Id()},
				},
				Feedback: "Devices[2].State",
			})
			node.AddElement(&protocol.Element{
				Type: protocol.ElementTypeColorPicker,
				Name: s.Lamp.Label(),
				Command: &protocol.Command{
					Cmd:  "color",
					Args: []string{s.Lamp.Id()},
				},
				Feedback: "Devices[4].State",
			})
			//log.Infof("Collection: %#v", lights.Lights)
			connection.Send(node)

		case client.LampUpdated:
			log.Warnf("Changed: %s (%s)", s.Lamp.Id(), s.Lamp.Label())
		case client.LampRemoved:
			log.Warnf("Removed: %s (%s)", s.Lamp.Id(), s.Lamp.Label())
		default:
			log.Infof("Received unknown event: %d", s.Event)
		}
	}
}

// WORKER that monitors the current connection state
func monitorState(node *protocol.Node, connection basenode.Connection) {
	for s := range connection.State() {
		switch s {
		case basenode.ConnectionStateConnected:
			connection.Send(node)
		case basenode.ConnectionStateDisconnected:
		}
	}
}

// WORKER that recives all incoming commands
func serverRecv(connection basenode.Connection) {
	for d := range connection.Receive() {
		processCommand(d)
	}
}

// This is called on each incoming command
func processCommand(cmd protocol.Command) {
	log.Info("Incoming command from server:", cmd)

	if len(cmd.Args) < 1 {
		return
	}

	lamp, err := lifxClient.Lights.GetById(cmd.Args[0])
	if err != nil {
		log.Error(err)
		return
	}

	switch cmd.Cmd {
	case "set":
		lamp.TurnOn()
		if len(cmd.Args) > 1 {
			if strings.Index(cmd.Args[1], "#") != 0 {
				cmd.Args[1] = "#" + cmd.Args[1]
			}

			c, err := colorful.Hex(cmd.Args[1])
			if err != nil {
				log.Error("Failed to decode color: ", err)
				return
			}

			h, s, v := c.Hsv()

			if len(cmd.Args) > 2 {
				duration, err := strconv.Atoi(cmd.Args[2])
				if err == nil {
					if len(cmd.Args) > 3 {
						kelvin, err := strconv.Atoi(cmd.Args[3])
						if err == nil {
							lamp.SetState(h, s, v, uint32(duration), uint32(kelvin))
							return
						}
					}
					lamp.SetState(h, s, v, uint32(duration), 6500)
					return
				}
				log.Error("Failed to decode duration: ", err)
			}

			lamp.SetState(h, s, v, 1000, 6500)
		}
	case "on":
		lamp.TurnOn()
	case "off":
		lamp.TurnOff()
	case "color":
		c, err := colorful.Hex(cmd.Params[0])
		if err != nil {
			log.Error(err)
			return
		}

		h, s, v := c.Hsv()

		lamp.SetState(h, s, v, 1000, 6500)
	}
}
<|endoftext|>