<commit_before>// Package resources provides unfancy resources embedding with Go.
package resources

import (
	"bytes"
	"fmt"
	"go/format"
	"io"
	"log"
	"os"
	"strings"
	"text/template"
)

// File mimicks the os.File and http.File interface.
type File interface {
	io.Reader
	Stat() (os.FileInfo, error)
}

// New creates a new Package.
func New() *Package {
	return &Package{
		Config: Config{
			Pkg:     "resources",
			Var:     "FS",
			Declare: true,
		},
		Files: make(map[string]File),
	}
}

// Config defines some details about the output Go file.
type Config struct {
	Pkg     string // Package name
	Var     string // Variable name to assign the file system to.
	Tag     string // Build tag, leave empty for no tag.
	Declare bool   // Dictates whatever there should be a defintion Variable
	Format  bool   // Whether gofmt should be applied to output.
}

// Package describes...
type Package struct {
	Config
	Files map[string]File
}

// Add a file to the package at the give path.
func (p *Package) Add(path string, file File) error {
	p.Files[path] = file
	return nil
}

// AddFile is a helper function that adds the files from the path into the package under the path file.
func (p *Package) AddFile(path string, file string) error {
	f, err := os.Open(file)
	if err != nil {
		return err
	}
	return p.Add(path, f)
}

// Build the package
func (p *Package) Build(out io.Writer) error {
	return pkg.Execute(out, p)
}

// Write the build to a file, you don't need to call Build.
func (p *Package) Write(path string) error {
	f, err := os.Create(path)
	if err != nil {
		return err
	}
	defer func() {
		err := f.Close()
		if err != nil {
			log.Panicf("Failed to close file: %s", err)
		}
	}()

	if !p.Format {
		return p.Build(f)
	}

	buf := &bytes.Buffer{}
	if e := p.Build(buf); e != nil {
		return e
	}

	fmted, e := format.Source(buf.Bytes())
	if e != nil {
		return e
	}
	_, e = f.Write(fmted)
	return e
}

var (
	// Template
	pkg *template.Template

	// BlockWidth allows to adjust the number of bytes per line in the generated file
	BlockWidth = 12
)

func reader(input io.Reader, indent int) (string, error) {
	var (
		buff      bytes.Buffer
		err       error
		curblock  = 0
		linebreak = "\n" + strings.Repeat("\t", indent)
	)

	b := make([]byte, BlockWidth)

	for n, e := input.Read(b); e == nil; n, e = input.Read(b) {
		for i := 0; i < n; i++ {
			_, e = fmt.Fprintf(&buff, "0x%02x,", b[i])
			if e != nil {
				err = e
				break
			}
			curblock++
			if curblock < BlockWidth {
				buff.WriteRune(' ')
				continue
			}
			buff.WriteString(linebreak)
			curblock = 0
		}
	}

	return buff.String(), err
}

func init() {
	pkg = template.Must(template.New("file").Funcs(template.FuncMap{"reader": reader}).Parse(`File{
				data: []byte{
					{{ reader . 5 }}
				},
				fi: FileInfo{
					name:    "{{ .Stat.Name }}",
					size:    {{ .Stat.Size }},
					modTime: time.Unix(0, {{ .Stat.ModTime.UnixNano }}),
					isDir:   {{ .Stat.IsDir }},
				},
			}`))

	pkg = template.Must(pkg.New("pkg").Parse(`{{ if .Tag }}// +build {{ .Tag }}

{{ end }}// Package {{.Pkg }} is generated by github.com/omeid/go-resources
package {{ .Pkg }}

import (
	"bytes"
	"errors"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"
)

// FileSystem is an http.FileSystem implementation.
type FileSystem struct {
	files map[string]File
}

// String returns the content of the file as string.
func (fs *FileSystem) String(name string) (string, bool) {
	if filepath.Separator != '/' && strings.IndexRune(name, filepath.Separator) >= 0 ||
		strings.Contains(name, "\x00") {
		return "", false
	}

	file, ok := fs.files[name]

	if !ok {
		return "", false
	}

	return string(file.data), true
}

// Open implements http.FileSystem.Open
func (fs *FileSystem) Open(name string) (http.File, error) {
	if filepath.Separator != '/' && strings.IndexRune(name, filepath.Separator) >= 0 ||
		strings.Contains(name, "\x00") {
		return nil, errors.New("http: invalid character in file path")
	}
	file, ok := fs.files[name]
	if !ok {
		files := []os.FileInfo{}
		for path, file := range fs.files {
			if strings.HasPrefix(path, name) {
				fi := file.fi
				files = append(files, &fi)
			}
		}

		if len(files) == 0 {
			return nil, os.ErrNotExist
		}

		//We have a directory.
		return &File{
			fi: FileInfo{
				isDir: true,
				files: files,
			}}, nil
	}
	file.Reader = bytes.NewReader(file.data)
	return &file, nil
}

// File implements http.File
type File struct {
	*bytes.Reader
	data []byte
	fi   FileInfo
}

// Close is a noop-closer.
func (f *File) Close() error {
	return nil
}

// Readdir implements http.File.Readdir
func (f *File) Readdir(count int) ([]os.FileInfo, error) {
	return nil, os.ErrNotExist
}

// Stat implements http.Stat.Readdir
func (f *File) Stat() (os.FileInfo, error) {
	return &f.fi, nil
}

// FileInfo implements the os.FileInfo interface.
type FileInfo struct {
	name    string
	size    int64
	mode    os.FileMode
	modTime time.Time
	isDir   bool
	sys     interface{}

	files []os.FileInfo
}

// Name implements os.FileInfo.Name
func (f *FileInfo) Name() string {
	return f.name
}

// Size implements os.FileInfo.Size
func (f *FileInfo) Size() int64 {
	return f.size
}

// Mode implements os.FileInfo.Mode
func (f *FileInfo) Mode() os.FileMode {
	return f.mode
}

// ModTime implements os.FileInfo.ModTime
func (f *FileInfo) ModTime() time.Time {
	return f.modTime
}

// IsDir implements os.FileInfo.IsDir
func (f *FileInfo) IsDir() bool {
	return f.isDir
}

// Readdir implements os.FileInfo.Readdir
func (f *FileInfo) Readdir(count int) ([]os.FileInfo, error) {
	return f.files, nil
}

// Sys returns the underlying value.
func (f *FileInfo) Sys() interface{} {
	return f.sys
}

{{ if .Declare }}var {{ .Var }} *FileSystem{{ end }}

func init() {
	{{ .Var }} = &FileSystem{
		files: map[string]File{
			{{range $path, $file := .Files }}"/{{ $path }}": {{ template "file" $file }},{{ end }}
		},
	}
}
`))
}
<commit_msg>improve documentation<commit_after>// Package resources provides unfancy resources embedding with Go.
package resources

import (
	"bytes"
	"fmt"
	"go/format"
	"io"
	"log"
	"os"
	"strings"
	"text/template"
)

// File mimicks the os.File and http.File interface.
type File interface {
	io.Reader
	Stat() (os.FileInfo, error)
}

// New creates a new Package.
func New() *Package {
	return &Package{
		Config: Config{
			Pkg:     "resources",
			Var:     "FS",
			Declare: true,
		},
		Files: make(map[string]File),
	}
}

// Config defines some details about the output file
type Config struct {
	Pkg     string // Pkg holds the package name
	Var     string // Var holds the variable name for the virtual filesystem
	Tag     string // Tag may hold an optional build tag, unless empty
	Declare bool   // Declare controls if the Var should be declared as well
	Format  bool   // Format controls, whether gofmt should be applied to the output
}

// A Package describes a collection of files and how they should be tranformed
// to an output.
type Package struct {
	Config
	Files map[string]File
}

// Add a file to the package at the give path.
func (p *Package) Add(path string, file File) error {
	p.Files[path] = file
	return nil
}

// AddFile is a helper function that adds the files from the path into the
// package under the path file.
func (p *Package) AddFile(path string, file string) error {
	f, err := os.Open(file)
	if err != nil {
		return err
	}
	return p.Add(path, f)
}

// Build compiles the package and writes it into an io.Writer.
func (p *Package) Build(out io.Writer) error {
	return pkg.Execute(out, p)
}

// Write builds the package (via Build) and writes the output the the file
// given by the path argument.
func (p *Package) Write(path string) error {
	f, err := os.Create(path)
	if err != nil {
		return err
	}
	defer func() {
		err := f.Close()
		if err != nil {
			log.Panicf("Failed to close file: %s", err)
		}
	}()

	if !p.Format {
		return p.Build(f)
	}

	buf := &bytes.Buffer{}
	if e := p.Build(buf); e != nil {
		return e
	}

	fmted, e := format.Source(buf.Bytes())
	if e != nil {
		return e
	}
	_, e = f.Write(fmted)
	return e
}

var (
	// Template
	pkg *template.Template

	// BlockWidth allows to adjust the number of bytes per line in the generated file
	BlockWidth = 12
)

func reader(input io.Reader, indent int) (string, error) {
	var (
		buff      bytes.Buffer
		err       error
		curblock  = 0
		linebreak = "\n" + strings.Repeat("\t", indent)
	)

	b := make([]byte, BlockWidth)

	for n, e := input.Read(b); e == nil; n, e = input.Read(b) {
		for i := 0; i < n; i++ {
			_, e = fmt.Fprintf(&buff, "0x%02x,", b[i])
			if e != nil {
				err = e
				break
			}
			curblock++
			if curblock < BlockWidth {
				buff.WriteRune(' ')
				continue
			}
			buff.WriteString(linebreak)
			curblock = 0
		}
	}

	return buff.String(), err
}

func init() {
	pkg = template.Must(template.New("file").Funcs(template.FuncMap{"reader": reader}).Parse(`File{
				data: []byte{
					{{ reader . 5 }}
				},
				fi: FileInfo{
					name:    "{{ .Stat.Name }}",
					size:    {{ .Stat.Size }},
					modTime: time.Unix(0, {{ .Stat.ModTime.UnixNano }}),
					isDir:   {{ .Stat.IsDir }},
				},
			}`))

	pkg = template.Must(pkg.New("pkg").Parse(`{{ if .Tag }}// +build {{ .Tag }}

{{ end }}// Package {{.Pkg }} is generated by github.com/omeid/go-resources
package {{ .Pkg }}

import (
	"bytes"
	"errors"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"
)

// FileSystem is an http.FileSystem implementation.
type FileSystem struct {
	files map[string]File
}

// String returns the content of the file as string.
func (fs *FileSystem) String(name string) (string, bool) {
	if filepath.Separator != '/' && strings.IndexRune(name, filepath.Separator) >= 0 ||
		strings.Contains(name, "\x00") {
		return "", false
	}

	file, ok := fs.files[name]

	if !ok {
		return "", false
	}

	return string(file.data), true
}

// Open implements http.FileSystem.Open
func (fs *FileSystem) Open(name string) (http.File, error) {
	if filepath.Separator != '/' && strings.IndexRune(name, filepath.Separator) >= 0 ||
		strings.Contains(name, "\x00") {
		return nil, errors.New("http: invalid character in file path")
	}
	file, ok := fs.files[name]
	if !ok {
		files := []os.FileInfo{}
		for path, file := range fs.files {
			if strings.HasPrefix(path, name) {
				fi := file.fi
				files = append(files, &fi)
			}
		}

		if len(files) == 0 {
			return nil, os.ErrNotExist
		}

		//We have a directory.
		return &File{
			fi: FileInfo{
				isDir: true,
				files: files,
			}}, nil
	}
	file.Reader = bytes.NewReader(file.data)
	return &file, nil
}

// File implements http.File
type File struct {
	*bytes.Reader
	data []byte
	fi   FileInfo
}

// Close is a noop-closer.
func (f *File) Close() error {
	return nil
}

// Readdir implements http.File.Readdir
func (f *File) Readdir(count int) ([]os.FileInfo, error) {
	return nil, os.ErrNotExist
}

// Stat implements http.Stat.Readdir
func (f *File) Stat() (os.FileInfo, error) {
	return &f.fi, nil
}

// FileInfo implements the os.FileInfo interface.
type FileInfo struct {
	name    string
	size    int64
	mode    os.FileMode
	modTime time.Time
	isDir   bool
	sys     interface{}

	files []os.FileInfo
}

// Name implements os.FileInfo.Name
func (f *FileInfo) Name() string {
	return f.name
}

// Size implements os.FileInfo.Size
func (f *FileInfo) Size() int64 {
	return f.size
}

// Mode implements os.FileInfo.Mode
func (f *FileInfo) Mode() os.FileMode {
	return f.mode
}

// ModTime implements os.FileInfo.ModTime
func (f *FileInfo) ModTime() time.Time {
	return f.modTime
}

// IsDir implements os.FileInfo.IsDir
func (f *FileInfo) IsDir() bool {
	return f.isDir
}

// Readdir implements os.FileInfo.Readdir
func (f *FileInfo) Readdir(count int) ([]os.FileInfo, error) {
	return f.files, nil
}

// Sys returns the underlying value.
func (f *FileInfo) Sys() interface{} {
	return f.sys
}

{{ if .Declare }}var {{ .Var }} *FileSystem{{ end }}

func init() {
	{{ .Var }} = &FileSystem{
		files: map[string]File{
			{{range $path, $file := .Files }}"/{{ $path }}": {{ template "file" $file }},{{ end }}
		},
	}
}
`))
}
<|endoftext|>


<commit_before>/*
   Copyright (c) 2014, Percona LLC and/or its affiliates. All rights reserved.

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU Affero General Public License as published by
   the Free Software Foundation, either version 3 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU Affero General Public License for more details.

   You should have received a copy of the GNU Affero General Public License
   along with this program.  If not, see <http://www.gnu.org/licenses/>
*/

package query_test

import (
	"database/sql"
	"encoding/json"
	"github.com/percona/cloud-protocol/proto"
	"github.com/percona/percona-agent/instance"
	"github.com/percona/percona-agent/mysql"
	"github.com/percona/percona-agent/pct"
	"github.com/percona/percona-agent/query"
	"github.com/percona/percona-agent/test/mock"
	"io/ioutil"
	. "launchpad.net/gocheck"
	"os"
	"path/filepath"
	"testing"
	"time"
)

// Hook up gocheck into the "go test" runner.
func Test(t *testing.T) { TestingT(t) }

/////////////////////////////////////////////////////////////////////////////
// Manager test suite
/////////////////////////////////////////////////////////////////////////////

type ManagerTestSuite struct {
	nullmysql     *mock.NullMySQL
	logChan       chan *proto.LogEntry
	logger        *pct.Logger
	tickChan      chan time.Time
	dataChan      chan interface{}
	traceChan     chan string
	readyChan     chan bool
	configDir     string
	tmpDir        string
	ir            *instance.Repo
	dsn           string
	rir           *instance.Repo
	mysqlInstance proto.ServiceInstance
	api           *mock.API
}

var _ = Suite(&ManagerTestSuite{})

func (s *ManagerTestSuite) SetUpSuite(t *C) {
	s.dsn = os.Getenv("PCT_TEST_MYSQL_DSN")
	if s.dsn == "" {
		t.Fatal("PCT_TEST_MYSQL_DSN is not set")
	}

	s.nullmysql = mock.NewNullMySQL()
	s.logChan = make(chan *proto.LogEntry, 10)
	s.logger = pct.NewLogger(s.logChan, query.SERVICE_NAME+"-manager-test")
	s.traceChan = make(chan string, 10)
	s.dataChan = make(chan interface{}, 1)

	var err error
	s.tmpDir, err = ioutil.TempDir("/tmp", "agent-test")
	t.Assert(err, IsNil)

	if err := pct.Basedir.Init(s.tmpDir); err != nil {
		t.Fatal(err)
	}
	s.configDir = pct.Basedir.Dir("config")

	// Fake instance repo
	s.ir = instance.NewRepo(pct.NewLogger(s.logChan, "ir"), s.configDir, s.api)
	data, err := json.Marshal(&proto.MySQLInstance{
		Hostname: "db1",
		DSN:      "user:host@tcp(127.0.0.1:3306)/",
	})
	t.Assert(err, IsNil)
	s.ir.Add("mysql", 1, data, false)

	// Real instance repo
	s.rir = instance.NewRepo(pct.NewLogger(s.logChan, "im-test"), s.configDir, s.api)
	data, err = json.Marshal(&proto.MySQLInstance{
		Hostname: "db1",
		DSN:      s.dsn,
	})
	t.Assert(err, IsNil)
	s.rir.Add("mysql", 1, data, false)
	s.mysqlInstance = proto.ServiceInstance{Service: "mysql", InstanceId: 1}

	links := map[string]string{
		"agent":     "http://localhost/agent",
		"instances": "http://localhost/instances",
	}
	s.api = mock.NewAPI("http://localhost", "http://localhost", "123", "abc-123-def", links)
}

func (s *ManagerTestSuite) SetUpTest(t *C) {
	glob := filepath.Join(pct.Basedir.Dir("config"), "*")
	files, err := filepath.Glob(glob)
	if err != nil {
		t.Fatal(err)
	}
	for _, file := range files {
		if err := os.Remove(file); err != nil {
			t.Fatal(err)
		}
	}
}

func (s *ManagerTestSuite) TearDownSuite(t *C) {
	if err := os.RemoveAll(s.tmpDir); err != nil {
		t.Error(err)
	}
}

// --------------------------------------------------------------------------

func (s *ManagerTestSuite) TestStartStopManager(t *C) {
	var err error

	// Create query manager
	mockConnFactory := &mock.ConnectionFactory{Conn: s.nullmysql}
	m := query.NewManager(s.logger, mockConnFactory, s.ir)
	t.Assert(m, Not(IsNil), Commentf("Make new query.Manager"))

	// The agent calls mm.Start().
	err = m.Start()
	t.Assert(err, IsNil)

	// Its status should be "Running".
	status := m.Status()
	t.Check(status[query.SERVICE_NAME], Equals, "Running")

	// Can't start manager twice.
	err = m.Start()
	t.Check(err, FitsTypeOf, pct.ServiceIsRunningError{})

	// Explain query
	q := "SELECT 1"
	expectedExplain := []proto.ExplainRow{
		proto.ExplainRow{
			Id: proto.NullInt64{
				NullInt64: sql.NullInt64{
					Int64: 1,
					Valid: true,
				},
			},
			SelectType: proto.NullString{
				NullString: sql.NullString{
					String: "SIMPLE",
					Valid:  true,
				},
			},
			Table: proto.NullString{
				NullString: sql.NullString{
					String: "",
					Valid:  false,
				},
			},
			CreateTable: proto.NullString{
				NullString: sql.NullString{
					String: "",
					Valid:  false,
				},
			},
			Type: proto.NullString{
				NullString: sql.NullString{
					String: "",
					Valid:  false,
				},
			},
			PossibleKeys: proto.NullString{
				NullString: sql.NullString{
					String: "",
					Valid:  false,
				},
			},
			Key: proto.NullString{
				NullString: sql.NullString{
					String: "",
					Valid:  false,
				},
			},
			KeyLen: proto.NullInt64{
				NullInt64: sql.NullInt64{
					Int64: 0,
					Valid: false,
				},
			},
			Ref: proto.NullString{
				NullString: sql.NullString{
					String: "",
					Valid:  false,
				},
			},
			Rows: proto.NullInt64{
				NullInt64: sql.NullInt64{
					Int64: 0,
					Valid: false,
				},
			},
			Extra: proto.NullString{
				NullString: sql.NullString{
					String: "No tables used",
					Valid:  true,
				},
			},
		},
	}
	s.nullmysql.SetExplain(q, expectedExplain)

	serviceInstance := proto.ServiceInstance{
		Service:    "mysql",
		InstanceId: 1,
	}

	explainQuery := &proto.ExplainQuery{
		ServiceInstance: serviceInstance,
		Query:           q,
	}
	data, err := json.Marshal(&explainQuery)
	t.Assert(err, IsNil)

	cmd := &proto.Cmd{
		Service: "query",
		Cmd:     "Explain",
		Data:    data,
	}

	gotReply := m.Handle(cmd)
	t.Assert(gotReply, NotNil)
	t.Assert(gotReply.Error, Equals, "")

	var gotExplain []proto.ExplainRow
	err = json.Unmarshal(gotReply.Data, &gotExplain)
	t.Assert(err, IsNil)
	t.Assert(gotExplain, DeepEquals, expectedExplain)

	// You can't stop this service
	err = m.Stop()
	t.Check(err, IsNil)
	status = m.Status()
	t.Check(status[query.SERVICE_NAME], Equals, "Running")
}

func (s *ManagerTestSuite) TestExplainWithoutDb(t *C) {
	var err error

	// Create query manager
	m := query.NewManager(s.logger, &mysql.RealConnectionFactory{}, s.rir)
	t.Assert(m, Not(IsNil), Commentf("Make new query.Manager"))

	// The agent calls mm.Start().
	err = m.Start()
	t.Assert(err, IsNil)

	// Explain query
	q := "SELECT 1"
	expectedExplain := []proto.ExplainRow{
		proto.ExplainRow{
			Id: proto.NullInt64{
				NullInt64: sql.NullInt64{
					Int64: 1,
					Valid: true,
				},
			},
			SelectType: proto.NullString{
				NullString: sql.NullString{
					String: "SIMPLE",
					Valid:  true,
				},
			},
			Table: proto.NullString{
				NullString: sql.NullString{
					String: "",
					Valid:  false,
				},
			},
			CreateTable: proto.NullString{
				NullString: sql.NullString{
					String: "",
					Valid:  false,
				},
			},
			Type: proto.NullString{
				NullString: sql.NullString{
					String: "",
					Valid:  false,
				},
			},
			PossibleKeys: proto.NullString{
				NullString: sql.NullString{
					String: "",
					Valid:  false,
				},
			},
			Key: proto.NullString{
				NullString: sql.NullString{
					String: "",
					Valid:  false,
				},
			},
			KeyLen: proto.NullInt64{
				NullInt64: sql.NullInt64{
					Int64: 0,
					Valid: false,
				},
			},
			Ref: proto.NullString{
				NullString: sql.NullString{
					String: "",
					Valid:  false,
				},
			},
			Rows: proto.NullInt64{
				NullInt64: sql.NullInt64{
					Int64: 0,
					Valid: false,
				},
			},
			Extra: proto.NullString{
				NullString: sql.NullString{
					String: "No tables used",
					Valid:  true,
				},
			},
		},
	}

	explainQuery := &proto.ExplainQuery{
		ServiceInstance: s.mysqlInstance,
		Query:           q,
	}
	data, err := json.Marshal(&explainQuery)
	t.Assert(err, IsNil)

	cmd := &proto.Cmd{
		Service: "query",
		Cmd:     "Explain",
		Data:    data,
	}

	gotReply := m.Handle(cmd)
	t.Assert(gotReply, NotNil)
	t.Assert(gotReply.Error, Equals, "")

	var gotExplain []proto.ExplainRow
	err = json.Unmarshal(gotReply.Data, &gotExplain)
	t.Assert(err, IsNil)
	t.Assert(gotExplain, DeepEquals, expectedExplain)
}

func (s *ManagerTestSuite) TestExplainWithDb(t *C) {
	var err error

	// Create query manager
	m := query.NewManager(s.logger, &mysql.RealConnectionFactory{}, s.rir)
	t.Assert(m, Not(IsNil), Commentf("Make new query.Manager"))

	// The agent calls mm.Start().
	err = m.Start()
	t.Assert(err, IsNil)

	// Explain query
	db := "information_schema"
	q := "SELECT table_name FROM tables WHERE table_name='tables'"

	expectedExplain := []proto.ExplainRow{
		proto.ExplainRow{
			Id: proto.NullInt64{
				NullInt64: sql.NullInt64{
					Int64: 1,
					Valid: true,
				},
			},
			SelectType: proto.NullString{
				NullString: sql.NullString{
					String: "SIMPLE",
					Valid:  true,
				},
			},
			Table: proto.NullString{
				NullString: sql.NullString{
					String: "tables",
					Valid:  true,
				},
			},
			CreateTable: proto.NullString{
				NullString: sql.NullString{
					String: "",
					Valid:  false,
				},
			},
			Type: proto.NullString{
				NullString: sql.NullString{
					String: "ALL",
					Valid:  true,
				},
			},
			PossibleKeys: proto.NullString{
				NullString: sql.NullString{
					String: "",
					Valid:  false,
				},
			},
			Key: proto.NullString{
				NullString: sql.NullString{
					String: "TABLE_NAME",
					Valid:  true,
				},
			},
			KeyLen: proto.NullInt64{
				NullInt64: sql.NullInt64{
					Int64: 0,
					Valid: false,
				},
			},
			Ref: proto.NullString{
				NullString: sql.NullString{
					String: "",
					Valid:  false,
				},
			},
			Rows: proto.NullInt64{
				NullInt64: sql.NullInt64{
					Int64: 0,
					Valid: false,
				},
			},
			Extra: proto.NullString{
				NullString: sql.NullString{
					String: "Using where; Skip_open_table; Scanned 1 database",
					Valid:  true,
				},
			},
		},
	}

	explainQuery := &proto.ExplainQuery{
		ServiceInstance: s.mysqlInstance,
		Db:              db,
		Query:           q,
	}
	data, err := json.Marshal(&explainQuery)
	t.Assert(err, IsNil)

	cmd := &proto.Cmd{
		Service: "query",
		Cmd:     "Explain",
		Data:    data,
	}

	gotReply := m.Handle(cmd)
	t.Assert(gotReply, NotNil)
	t.Assert(gotReply.Error, Equals, "")

	var gotExplain []proto.ExplainRow
	err = json.Unmarshal(gotReply.Data, &gotExplain)
	t.Assert(err, IsNil)
	t.Assert(gotExplain, DeepEquals, expectedExplain)
}

<commit_msg>PCT-671: s/q/query/<commit_after>/*
   Copyright (c) 2014, Percona LLC and/or its affiliates. All rights reserved.

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU Affero General Public License as published by
   the Free Software Foundation, either version 3 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU Affero General Public License for more details.

   You should have received a copy of the GNU Affero General Public License
   along with this program.  If not, see <http://www.gnu.org/licenses/>
*/

package query_test

import (
	"database/sql"
	"encoding/json"
	"github.com/percona/cloud-protocol/proto"
	"github.com/percona/percona-agent/instance"
	"github.com/percona/percona-agent/mysql"
	"github.com/percona/percona-agent/pct"
	"github.com/percona/percona-agent/query"
	"github.com/percona/percona-agent/test/mock"
	"io/ioutil"
	. "launchpad.net/gocheck"
	"os"
	"path/filepath"
	"testing"
	"time"
)

// Hook up gocheck into the "go test" runner.
func Test(t *testing.T) { TestingT(t) }

/////////////////////////////////////////////////////////////////////////////
// Manager test suite
/////////////////////////////////////////////////////////////////////////////

type ManagerTestSuite struct {
	nullmysql     *mock.NullMySQL
	logChan       chan *proto.LogEntry
	logger        *pct.Logger
	tickChan      chan time.Time
	dataChan      chan interface{}
	traceChan     chan string
	readyChan     chan bool
	configDir     string
	tmpDir        string
	ir            *instance.Repo
	dsn           string
	rir           *instance.Repo
	mysqlInstance proto.ServiceInstance
	api           *mock.API
}

var _ = Suite(&ManagerTestSuite{})

func (s *ManagerTestSuite) SetUpSuite(t *C) {
	s.dsn = os.Getenv("PCT_TEST_MYSQL_DSN")
	if s.dsn == "" {
		t.Fatal("PCT_TEST_MYSQL_DSN is not set")
	}

	s.nullmysql = mock.NewNullMySQL()
	s.logChan = make(chan *proto.LogEntry, 10)
	s.logger = pct.NewLogger(s.logChan, query.SERVICE_NAME+"-manager-test")
	s.traceChan = make(chan string, 10)
	s.dataChan = make(chan interface{}, 1)

	var err error
	s.tmpDir, err = ioutil.TempDir("/tmp", "agent-test")
	t.Assert(err, IsNil)

	if err := pct.Basedir.Init(s.tmpDir); err != nil {
		t.Fatal(err)
	}
	s.configDir = pct.Basedir.Dir("config")

	// Fake instance repo
	s.ir = instance.NewRepo(pct.NewLogger(s.logChan, "ir"), s.configDir, s.api)
	data, err := json.Marshal(&proto.MySQLInstance{
		Hostname: "db1",
		DSN:      "user:host@tcp(127.0.0.1:3306)/",
	})
	t.Assert(err, IsNil)
	s.ir.Add("mysql", 1, data, false)

	// Real instance repo
	s.rir = instance.NewRepo(pct.NewLogger(s.logChan, "im-test"), s.configDir, s.api)
	data, err = json.Marshal(&proto.MySQLInstance{
		Hostname: "db1",
		DSN:      s.dsn,
	})
	t.Assert(err, IsNil)
	s.rir.Add("mysql", 1, data, false)
	s.mysqlInstance = proto.ServiceInstance{Service: "mysql", InstanceId: 1}

	links := map[string]string{
		"agent":     "http://localhost/agent",
		"instances": "http://localhost/instances",
	}
	s.api = mock.NewAPI("http://localhost", "http://localhost", "123", "abc-123-def", links)
}

func (s *ManagerTestSuite) SetUpTest(t *C) {
	glob := filepath.Join(pct.Basedir.Dir("config"), "*")
	files, err := filepath.Glob(glob)
	if err != nil {
		t.Fatal(err)
	}
	for _, file := range files {
		if err := os.Remove(file); err != nil {
			t.Fatal(err)
		}
	}
}

func (s *ManagerTestSuite) TearDownSuite(t *C) {
	if err := os.RemoveAll(s.tmpDir); err != nil {
		t.Error(err)
	}
}

// --------------------------------------------------------------------------

func (s *ManagerTestSuite) TestStartStopManager(t *C) {
	var err error

	// Create query manager
	mockConnFactory := &mock.ConnectionFactory{Conn: s.nullmysql}
	m := query.NewManager(s.logger, mockConnFactory, s.ir)
	t.Assert(m, Not(IsNil), Commentf("Make new query.Manager"))

	// The agent calls mm.Start().
	err = m.Start()
	t.Assert(err, IsNil)

	// Its status should be "Running".
	status := m.Status()
	t.Check(status[query.SERVICE_NAME], Equals, "Running")

	// Can't start manager twice.
	err = m.Start()
	t.Check(err, FitsTypeOf, pct.ServiceIsRunningError{})

	// Explain query
	q := "SELECT 1"
	expectedExplain := []proto.ExplainRow{
		proto.ExplainRow{
			Id: proto.NullInt64{
				NullInt64: sql.NullInt64{
					Int64: 1,
					Valid: true,
				},
			},
			SelectType: proto.NullString{
				NullString: sql.NullString{
					String: "SIMPLE",
					Valid:  true,
				},
			},
			Table: proto.NullString{
				NullString: sql.NullString{
					String: "",
					Valid:  false,
				},
			},
			CreateTable: proto.NullString{
				NullString: sql.NullString{
					String: "",
					Valid:  false,
				},
			},
			Type: proto.NullString{
				NullString: sql.NullString{
					String: "",
					Valid:  false,
				},
			},
			PossibleKeys: proto.NullString{
				NullString: sql.NullString{
					String: "",
					Valid:  false,
				},
			},
			Key: proto.NullString{
				NullString: sql.NullString{
					String: "",
					Valid:  false,
				},
			},
			KeyLen: proto.NullInt64{
				NullInt64: sql.NullInt64{
					Int64: 0,
					Valid: false,
				},
			},
			Ref: proto.NullString{
				NullString: sql.NullString{
					String: "",
					Valid:  false,
				},
			},
			Rows: proto.NullInt64{
				NullInt64: sql.NullInt64{
					Int64: 0,
					Valid: false,
				},
			},
			Extra: proto.NullString{
				NullString: sql.NullString{
					String: "No tables used",
					Valid:  true,
				},
			},
		},
	}
	s.nullmysql.SetExplain(q, expectedExplain)

	serviceInstance := proto.ServiceInstance{
		Service:    "mysql",
		InstanceId: 1,
	}

	explainQuery := &proto.ExplainQuery{
		ServiceInstance: serviceInstance,
		Query:           q,
	}
	data, err := json.Marshal(&explainQuery)
	t.Assert(err, IsNil)

	cmd := &proto.Cmd{
		Service: "query",
		Cmd:     "Explain",
		Data:    data,
	}

	gotReply := m.Handle(cmd)
	t.Assert(gotReply, NotNil)
	t.Assert(gotReply.Error, Equals, "")

	var gotExplain []proto.ExplainRow
	err = json.Unmarshal(gotReply.Data, &gotExplain)
	t.Assert(err, IsNil)
	t.Assert(gotExplain, DeepEquals, expectedExplain)

	// You can't stop this service
	err = m.Stop()
	t.Check(err, IsNil)
	status = m.Status()
	t.Check(status[query.SERVICE_NAME], Equals, "Running")
}

func (s *ManagerTestSuite) TestExplainWithoutDb(t *C) {
	var err error

	// Create query manager
	m := query.NewManager(s.logger, &mysql.RealConnectionFactory{}, s.rir)
	t.Assert(m, Not(IsNil), Commentf("Make new query.Manager"))

	// The agent calls mm.Start().
	err = m.Start()
	t.Assert(err, IsNil)

	// Explain query
	query := "SELECT 1"
	expectedExplain := []proto.ExplainRow{
		proto.ExplainRow{
			Id: proto.NullInt64{
				NullInt64: sql.NullInt64{
					Int64: 1,
					Valid: true,
				},
			},
			SelectType: proto.NullString{
				NullString: sql.NullString{
					String: "SIMPLE",
					Valid:  true,
				},
			},
			Table: proto.NullString{
				NullString: sql.NullString{
					String: "",
					Valid:  false,
				},
			},
			CreateTable: proto.NullString{
				NullString: sql.NullString{
					String: "",
					Valid:  false,
				},
			},
			Type: proto.NullString{
				NullString: sql.NullString{
					String: "",
					Valid:  false,
				},
			},
			PossibleKeys: proto.NullString{
				NullString: sql.NullString{
					String: "",
					Valid:  false,
				},
			},
			Key: proto.NullString{
				NullString: sql.NullString{
					String: "",
					Valid:  false,
				},
			},
			KeyLen: proto.NullInt64{
				NullInt64: sql.NullInt64{
					Int64: 0,
					Valid: false,
				},
			},
			Ref: proto.NullString{
				NullString: sql.NullString{
					String: "",
					Valid:  false,
				},
			},
			Rows: proto.NullInt64{
				NullInt64: sql.NullInt64{
					Int64: 0,
					Valid: false,
				},
			},
			Extra: proto.NullString{
				NullString: sql.NullString{
					String: "No tables used",
					Valid:  true,
				},
			},
		},
	}

	explainQuery := &proto.ExplainQuery{
		ServiceInstance: s.mysqlInstance,
		Query:           query,
	}
	data, err := json.Marshal(&explainQuery)
	t.Assert(err, IsNil)

	cmd := &proto.Cmd{
		Service: "query",
		Cmd:     "Explain",
		Data:    data,
	}

	gotReply := m.Handle(cmd)
	t.Assert(gotReply, NotNil)
	t.Assert(gotReply.Error, Equals, "")

	var gotExplain []proto.ExplainRow
	err = json.Unmarshal(gotReply.Data, &gotExplain)
	t.Assert(err, IsNil)
	t.Assert(gotExplain, DeepEquals, expectedExplain)
}

func (s *ManagerTestSuite) TestExplainWithDb(t *C) {
	var err error

	// Create query manager
	m := query.NewManager(s.logger, &mysql.RealConnectionFactory{}, s.rir)
	t.Assert(m, Not(IsNil), Commentf("Make new query.Manager"))

	// The agent calls mm.Start().
	err = m.Start()
	t.Assert(err, IsNil)

	// Explain query
	db := "information_schema"
	query := "SELECT table_name FROM tables WHERE table_name='tables'"

	expectedExplain := []proto.ExplainRow{
		proto.ExplainRow{
			Id: proto.NullInt64{
				NullInt64: sql.NullInt64{
					Int64: 1,
					Valid: true,
				},
			},
			SelectType: proto.NullString{
				NullString: sql.NullString{
					String: "SIMPLE",
					Valid:  true,
				},
			},
			Table: proto.NullString{
				NullString: sql.NullString{
					String: "tables",
					Valid:  true,
				},
			},
			CreateTable: proto.NullString{
				NullString: sql.NullString{
					String: "",
					Valid:  false,
				},
			},
			Type: proto.NullString{
				NullString: sql.NullString{
					String: "ALL",
					Valid:  true,
				},
			},
			PossibleKeys: proto.NullString{
				NullString: sql.NullString{
					String: "",
					Valid:  false,
				},
			},
			Key: proto.NullString{
				NullString: sql.NullString{
					String: "TABLE_NAME",
					Valid:  true,
				},
			},
			KeyLen: proto.NullInt64{
				NullInt64: sql.NullInt64{
					Int64: 0,
					Valid: false,
				},
			},
			Ref: proto.NullString{
				NullString: sql.NullString{
					String: "",
					Valid:  false,
				},
			},
			Rows: proto.NullInt64{
				NullInt64: sql.NullInt64{
					Int64: 0,
					Valid: false,
				},
			},
			Extra: proto.NullString{
				NullString: sql.NullString{
					String: "Using where; Skip_open_table; Scanned 1 database",
					Valid:  true,
				},
			},
		},
	}

	explainQuery := &proto.ExplainQuery{
		ServiceInstance: s.mysqlInstance,
		Db:              db,
		Query:           query,
	}
	data, err := json.Marshal(&explainQuery)
	t.Assert(err, IsNil)

	cmd := &proto.Cmd{
		Service: "query",
		Cmd:     "Explain",
		Data:    data,
	}

	gotReply := m.Handle(cmd)
	t.Assert(gotReply, NotNil)
	t.Assert(gotReply.Error, Equals, "")

	var gotExplain []proto.ExplainRow
	err = json.Unmarshal(gotReply.Data, &gotExplain)
	t.Assert(err, IsNil)
	t.Assert(gotExplain, DeepEquals, expectedExplain)
}
<|endoftext|>


<commit_before>// Copyright 2015 Square Inc.
//
// 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 query

// show all metrics
// show tags WHERE predicate

import (
	"testing"

	"github.com/square/metrics/assert"
)

// these queries should successfully parse,
// with a corresponding command.
var inputs = []string{
	// describes
	"describe all",
	"describe x",
	"describe cpu_usage",
	"describe inspect",
	"describe in2",
	"describe cpu_usage where key = 'value'",
	"describe cpu_usage where key = 'value\\''",
	"describe cpu_usage where key != 'value'",
	"describe cpu_usage where (key = 'value')",
	"describe cpu_usage where not (key = 'value')",
	"describe cpu_usage where not key = 'value'",
	"describe cpu_usage where (key = 'value')",
	"describe cpu_usage where key = 'value' or key = 'value'",
	"describe cpu_usage where key in ('value', 'value')",
	"describe cpu_usage where key matches 'abc'",
	"describe nodes.cpu.usage where datacenter='sjc1b' and type='idle' and host matches 'fwd'",
	// predicate parenthesis test
	"describe cpu_usage where key = 'value' and (key = 'value')",
	"describe cpu_usage where (key = 'value') and key = 'value'",
	"describe cpu_usage where (key = 'value') and (key = 'value')",
	"describe cpu_usage where (key = 'value' and key = 'value')",

	// Leading/trailing whitespace
	" describe all ",
	" describe x ",
	" select 0, 1, 2, 3, 4, 5, 6, 7, 8, 9 from 0 to 0 ",

	// selects - spaces and keywords
	"select f( g(5) group by a,w,q) from 0 to 0",
	"select f( g(5) group by a, w, q )from 0 to 0 ",
	"select f( g(5) group by a,w, q)from 0 to 0",
	"select f( g(5) group by a,w,q)to 0 from 0",
	"select f(g(5)group by a,w,q)to 0 from 0",
	"   select f(g(5)group by a,w,q)  from  	  0 	   to     0",
	"   select( f(g(5)group by a,w,q)  )from 0 to 0",
	"select(f(g(5)group by`a`,w,q)) from 0 to 0",
	"select(f(g(5)group by`a`,w,q)) from 0 to 0",
	"select(fromx+tox+groupx+byx+selectx+describex+allx+wherex) from 0 to 0",
}

var selects2 = []string{
	// All these queries are tested with and without the prefix "select"
	// selects - parenthesis
	"0 from 0 to 0",
	"(0) from 0 to 0",
	"(0) where foo = 'bar' from 0 to 0",
	// selects - numbers
	"0, 1, 2, 3, 4, 5, 6, 7, 8, 9 from 0 to 0",
	"10, 100, 1000 from 0 to 0",
	"10.1, 10.01, 10.001 from 0 to 0",
	"-10.1, -10.01, -10.001 from 0 to 0",
	"1.0e1, 1.0e2, 1.0e10, 1.0e0 from 0 to 0",
	"1.0e-5, 1.0e+5 from 0 to 0",
	// selects - trying out arithmetic
	"x from 0 to 0",
	"x-y-z from 0 to 0",
	"(x)-(y)-(z) from 0 to 0",
	"0 from 0 to 0",
	"x, y from 0 to 0",
	"1 + 2 * 3 + 4 from 0 to 0",
	"x * (y + 123), z from 0 to 0",
	// testing escaping
	"`x` from 0 to 0",
	// selects - timestamps
	"x * (y + 123), z from '2011-2-4 PTZ' to '2015-6-1 PTZ'",
	"x * (y + 123), z from 0 to 10000",
	"1 from -10m to now",
	"1 from -10M to -10m",
	// selects - function calls
	"foo(x) from 0 to 0",
	"bar(x, y) from 0 to 0",
	"baz(x, y, z+1+foo(1)) from 0 to 0",
	// selects - testing out property values
	"x from 0 to 0",
	"x from 0 to 0",
	"x from 0 to 0 resolution '10s'",
	"x from 0 to 0 resolution '10h'",
	"x from 0 to 0 resolution '300s'",
	"x from 0 to 0 resolution '17m'",
	"x from 0 to 0 sample by 'max'",
	"x from 0 to 0 sample   by 'max'",
	// selects - aggregate functions
	"scalar.max(x) from 0 to 0",
	"aggregate.max(x, y) from 0 to 0",
	"aggregate.max(x group by foo) + 3 from 0 to 0",
	// selects - where clause
	"x where y = 'z' from 0 to 0",
	// selects - per-identifier where clause
	"x + z[y = 'z'] from 0 to 0",
	"x[y = 'z'] from 0 to 0",
	// selects - complicated queries
	"aggregate.max(x[y = 'z'] group by foo) from 0 to 0",
	"cpu.user + cpu.kernel where host = 'apa3.sjc2b' from 0 to 0",
	"'string literal' where host = 'apa3.sjc2b' from 0 to 0",
	"timeshift( metric, '5h') where host = 'apa3.sjc2b' from 0 to 0",
	// pipe expressions
	"x | y from 0 to 0",
	"x | y + 1 from 0 to 0",
	"x | y - 1 from 0 to 0",
	"x | y * 1 from 0 to 0",
	"x | y / 1 from 0 to 0",
	"x | y(group by a) from 0 to 0",
	"x + 1 | y(group by a) from 0 to 0",
	"x | y | z + 1 from 0 to 0",
}

var selects = []string{
	"x|y from 0 to 0",
	"x|f + y*z from 0 to 0",
	"x|f + y|g from 0 to 0",
	"1 + 2 | f from 0 to 0",
}

// these queries should fail with a syntax error.
var syntaxErrorQuery = []string{
	"select ( from 0 to 0",
	"select ) from 0 to 0",
	"describe ( from 0 to 0",
	"describe in from 0 to 0",
	"describe invalid_regex where key matches 'ab[' from 0 to 0",
	"select x invalid_property 0 from 0 to 0",
	"select x sampleby 0 from 0 to 0",
	"select x sample 0 from 0 to 0",
	"select x by 0 from 0 to 0",
	"select x",
	"select x from 0",
	"select x to 0",
	"select x from 0 from 1 to 0",
	"select x from 0 to 1 to 0",
	"select x from 0 resolution '30s' resolution '25s' to 0",
	"select x from 0 from 1 sample by 'min' sample by 'min' to 0",
	"select f(3 groupby x) from 0 to 0",
	"select c group by a from 0 to 0",
	"select x[] from 0 to 0",
}

func TestParse_success(t *testing.T) {
	for _, row := range inputs {
		_, err := Parse(row)
		if err != nil {
			t.Errorf("[%s] failed to parse: %s", row, err.Error())
		}
	}

	for _, row := range selects {
		for _, prefix := range []string{"", "select "} {
			query := prefix + row
			node, err := Parse(query)
			if err != nil {
				t.Errorf("[%s] failed to parse: %s", query, err.Error())
			}
			t.Logf("name: %s", row)
			t.Logf("node: %s", PrintNode(node.(Node)))

		}
	}
}

func TestParse_syntaxError(t *testing.T) {
	for _, row := range syntaxErrorQuery {
		_, err := Parse(row)
		if err == nil {
			t.Errorf("[%s] should have failed to parse", row)
		} else if _, ok := err.(SyntaxErrors); !ok {
			t.Logf("[%s] Expected SyntaxErrors, got: %s", row, err.Error())
		}
	}
}

func TestCompile(t *testing.T) {
	for _, row := range inputs {
		a := assert.New(t).Contextf(row)
		p := Parser{Buffer: row}
		p.Init()
		a.CheckError(p.Parse())
		p.Execute()
		testParserResult(a, p)
	}
}

// Helper functions
// ================

func testParserResult(a assert.Assert, p Parser) {
	a.EqInt(len(p.assertions), 0)
	if len(p.assertions) != 0 {
		for _, err := range p.assertions {
			a.Errorf("assertion error: %s", err.Error())
		}
	}
	if len(p.nodeStack) != 0 {
		for _, node := range p.nodeStack {
			a.Errorf("node error:\n%s", PrintNode(node))
		}
	}
}
<commit_msg>fix up query tests<commit_after>// Copyright 2015 Square Inc.
//
// 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 query

// show all metrics
// show tags WHERE predicate

import (
	"testing"

	"github.com/square/metrics/assert"
)

// these queries should successfully parse,
// with a corresponding command.
var inputs = []string{
	// describes
	"describe all",
	"describe x",
	"describe cpu_usage",
	"describe inspect",
	"describe in2",
	"describe cpu_usage where key = 'value'",
	"describe cpu_usage where key = 'value\\''",
	"describe cpu_usage where key != 'value'",
	"describe cpu_usage where (key = 'value')",
	"describe cpu_usage where not (key = 'value')",
	"describe cpu_usage where not key = 'value'",
	"describe cpu_usage where (key = 'value')",
	"describe cpu_usage where key = 'value' or key = 'value'",
	"describe cpu_usage where key in ('value', 'value')",
	"describe cpu_usage where key matches 'abc'",
	"describe nodes.cpu.usage where datacenter='sjc1b' and type='idle' and host matches 'fwd'",
	// predicate parenthesis test
	"describe cpu_usage where key = 'value' and (key = 'value')",
	"describe cpu_usage where (key = 'value') and key = 'value'",
	"describe cpu_usage where (key = 'value') and (key = 'value')",
	"describe cpu_usage where (key = 'value' and key = 'value')",

	// Leading/trailing whitespace
	" describe all ",
	" describe x ",
	" select 0, 1, 2, 3, 4, 5, 6, 7, 8, 9 from 0 to 0 ",

	// selects - spaces and keywords
	"select f( g(5) group by a,w,q) from 0 to 0",
	"select f( g(5) group by a, w, q )from 0 to 0 ",
	"select f( g(5) group by a,w, q)from 0 to 0",
	"select f( g(5) group by a,w,q)to 0 from 0",
	"select f(g(5)group by a,w,q)to 0 from 0",
	"   select f(g(5)group by a,w,q)  from  	  0 	   to     0",
	"   select( f(g(5)group by a,w,q)  )from 0 to 0",
	"select(f(g(5)group by`a`,w,q)) from 0 to 0",
	"select(f(g(5)group by`a`,w,q)) from 0 to 0",
	"select(fromx+tox+groupx+byx+selectx+describex+allx+wherex) from 0 to 0",
}

var selects = []string{
	// All these queries are tested with and without the prefix "select"
	// selects - parenthesis
	"0 from 0 to 0",
	"(0) from 0 to 0",
	"(0) where foo = 'bar' from 0 to 0",
	// selects - numbers
	"0, 1, 2, 3, 4, 5, 6, 7, 8, 9 from 0 to 0",
	"10, 100, 1000 from 0 to 0",
	"10.1, 10.01, 10.001 from 0 to 0",
	"-10.1, -10.01, -10.001 from 0 to 0",
	"1.0e1, 1.0e2, 1.0e10, 1.0e0 from 0 to 0",
	"1.0e-5, 1.0e+5 from 0 to 0",
	// selects - trying out arithmetic
	"x from 0 to 0",
	"x-y-z from 0 to 0",
	"(x)-(y)-(z) from 0 to 0",
	"0 from 0 to 0",
	"x, y from 0 to 0",
	"1 + 2 * 3 + 4 from 0 to 0",
	"x * (y + 123), z from 0 to 0",
	// testing escaping
	"`x` from 0 to 0",
	// selects - timestamps
	"x * (y + 123), z from '2011-2-4 PTZ' to '2015-6-1 PTZ'",
	"x * (y + 123), z from 0 to 10000",
	"1 from -10m to now",
	"1 from -10M to -10m",
	// selects - function calls
	"foo(x) from 0 to 0",
	"bar(x, y) from 0 to 0",
	"baz(x, y, z+1+foo(1)) from 0 to 0",
	// selects - testing out property values
	"x from 0 to 0",
	"x from 0 to 0",
	"x from 0 to 0 resolution '10s'",
	"x from 0 to 0 resolution '10h'",
	"x from 0 to 0 resolution '300s'",
	"x from 0 to 0 resolution '17m'",
	"x from 0 to 0 sample by 'max'",
	"x from 0 to 0 sample   by 'max'",
	// selects - aggregate functions
	"scalar.max(x) from 0 to 0",
	"aggregate.max(x, y) from 0 to 0",
	"aggregate.max(x group by foo) + 3 from 0 to 0",
	// selects - where clause
	"x where y = 'z' from 0 to 0",
	// selects - per-identifier where clause
	"x + z[y = 'z'] from 0 to 0",
	"x[y = 'z'] from 0 to 0",
	// selects - complicated queries
	"aggregate.max(x[y = 'z'] group by foo) from 0 to 0",
	"cpu.user + cpu.kernel where host = 'apa3.sjc2b' from 0 to 0",
	"'string literal' where host = 'apa3.sjc2b' from 0 to 0",
	"timeshift( metric, '5h') where host = 'apa3.sjc2b' from 0 to 0",
	// pipe expressions
	"x | y from 0 to 0",
	"x | y + 1 from 0 to 0",
	"x | y - 1 from 0 to 0",
	"x | y * 1 from 0 to 0",
	"x | y / 1 from 0 to 0",
	"x | y(group by a) from 0 to 0",
	"x + 1 | y(group by a) from 0 to 0",
	"x | y | z + 1 from 0 to 0",
	"x|y from 0 to 0",
	"x|f + y*z from 0 to 0",
	"x|f + y|g from 0 to 0",
	"1 + 2 | f from 0 to 0",
}

// these queries should fail with a syntax error.
var syntaxErrorQuery = []string{
	"select ( from 0 to 0",
	"select ) from 0 to 0",
	"describe ( from 0 to 0",
	"describe in from 0 to 0",
	"describe invalid_regex where key matches 'ab[' from 0 to 0",
	"select x invalid_property 0 from 0 to 0",
	"select x sampleby 0 from 0 to 0",
	"select x sample 0 from 0 to 0",
	"select x by 0 from 0 to 0",
	"select x",
	"select x from 0",
	"select x to 0",
	"select x from 0 from 1 to 0",
	"select x from 0 to 1 to 0",
	"select x from 0 resolution '30s' resolution '25s' to 0",
	"select x from 0 from 1 sample by 'min' sample by 'min' to 0",
	"select f(3 groupby x) from 0 to 0",
	"select c group by a from 0 to 0",
	"select x[] from 0 to 0",
}

func TestParse_success(t *testing.T) {
	for _, row := range inputs {
		_, err := Parse(row)
		if err != nil {
			t.Errorf("[%s] failed to parse: %s", row, err.Error())
		}
	}

	for _, row := range selects {
		for _, prefix := range []string{"", "select "} {
			query := prefix + row
			_, err := Parse(query)
			if err != nil {
				t.Errorf("[%s] failed to parse: %s", query, err.Error())
			}
		}
	}
}

func TestParse_syntaxError(t *testing.T) {
	for _, row := range syntaxErrorQuery {
		_, err := Parse(row)
		if err == nil {
			t.Errorf("[%s] should have failed to parse", row)
		} else if _, ok := err.(SyntaxErrors); !ok {
			t.Logf("[%s] Expected SyntaxErrors, got: %s", row, err.Error())
		}
	}
}

func TestCompile(t *testing.T) {
	for _, row := range inputs {
		a := assert.New(t).Contextf(row)
		p := Parser{Buffer: row}
		p.Init()
		a.CheckError(p.Parse())
		p.Execute()
		testParserResult(a, p)
	}
}

// Helper functions
// ================

func testParserResult(a assert.Assert, p Parser) {
	a.EqInt(len(p.assertions), 0)
	if len(p.assertions) != 0 {
		for _, err := range p.assertions {
			a.Errorf("assertion error: %s", err.Error())
		}
	}
	if len(p.nodeStack) != 0 {
		for _, node := range p.nodeStack {
			a.Errorf("node error:\n%s", PrintNode(node))
		}
	}
}
<|endoftext|>


<commit_before><commit_msg>FIX: typo on environtment key<commit_after><|endoftext|>


<commit_before>/* Copyright (c) 2016 PaddlePaddle Authors All Rights Reserve.

   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 autoscaler

import (
	"sort"
	"time"

	// TODO(typhoonzero): this package still depends on k8s API, try to remove this.
	"k8s.io/apimachinery/pkg/api/resource"
	batchv1 "k8s.io/client-go/pkg/apis/batch/v1"

	paddlejob "github.com/PaddlePaddle/cloud/go/api"
	log "github.com/sirupsen/logrus"
)

const (
	defaultLoopDur = time.Second * 5
)

// ClusterResource is the resource of a cluster
type ClusterResource struct {
	NodeCount int

	GPURequest int
	GPULimit   int
	GPUTotal   int

	CPURequestMilli int64
	CPULimitMilli   int64
	CPUTotalMilli   int64

	MemoryRequestMega int64
	MemoryLimitMega   int64
	MemoryTotalMega   int64
}

// Cluster represents the cluster managment system such as Kubernetes.
type Cluster interface {
	// SyncResource will sync resource values with the cluster.
	// should call this function in every tick.
	SyncResource() (ClusterResource, error)

	// GetTrainerJob gets the trainer job spec.
	GetTrainerJob(job *paddlejob.TrainingJob) (*batchv1.Job, error)

	// UpdateTrainerJob updates the trainer job spec.
	UpdateTrainerJob(job *batchv1.Job) error

	// JobRunning check if all the pods are in "Running" status.
	JobRunning(job *paddlejob.TrainingJob) (bool, error)
}

type job struct {
	Config     *paddlejob.TrainingJob
	TrainerJob *batchv1.Job
}

func (j job) TrainerGPULimit() int {
	q := j.Config.Spec.Trainer.Resources.Limits.NvidiaGPU()
	return int(q.Value())
}

func (j job) TrainerCPURequestMilli() int64 {
	q := j.Config.Spec.Trainer.Resources.Requests.Cpu()
	return q.ScaledValue(resource.Milli)
}

func (j job) TrainerMemRequestMega() int64 {
	q := j.Config.Spec.Trainer.Resources.Requests.Memory()
	return q.ScaledValue(resource.Mega)
}

func (j job) Fulfillment() float64 {
	minInstance := j.Config.Spec.Trainer.MinInstance
	maxInstance := j.Config.Spec.Trainer.MaxInstance

	if minInstance == maxInstance {
		return 1
	}

	curInstance := int(*j.TrainerJob.Spec.Parallelism)
	return float64(curInstance-minInstance) / float64(maxInstance-minInstance)
}

// Autoscaler launches and scales the training jobs.
type Autoscaler struct {
	ticker  *time.Ticker
	cluster Cluster
	jobs    map[string]job
	eventCh chan event
}

// New creates a new Autoscaler.
func New(cluster Cluster, options ...func(*Autoscaler)) *Autoscaler {
	c := &Autoscaler{
		cluster: cluster,
		ticker:  time.NewTicker(defaultLoopDur),
		jobs:    make(map[string]job),
		eventCh: make(chan event),
	}
	for _, option := range options {
		option(c)
	}
	return c
}

type jobs []job

func (j jobs) Len() int {
	return len(j)
}

func (j jobs) Less(a int, b int) bool {
	scoreA := j[a].Fulfillment()
	scoreB := j[b].Fulfillment()

	if scoreA == scoreB {
		resA := j[a].Config.Spec.Trainer.Resources
		resB := j[b].Config.Spec.Trainer.Resources
		resALimitsGPU := *resA.Limits.NvidiaGPU()
		resBLimitsGPU := *resB.Limits.NvidiaGPU()
		if resALimitsGPU.Cmp(resBLimitsGPU) == 0 {
			resARequestsCPU := *resA.Requests.Cpu()
			resBRequestsCPU := *resB.Requests.Cpu()
			if resARequestsCPU.Cmp(resBRequestsCPU) == 0 {
				resARequestsMem := *resA.Requests.Memory()
				resBRequestsMem := *resB.Requests.Memory()
				return resARequestsMem.Cmp(resBRequestsMem) == -1
			}
			return resARequestsCPU.Cmp(resBRequestsCPU) == -1
		}
		return resALimitsGPU.Cmp(resBLimitsGPU) == -1
	}
	return scoreA < scoreB
}

func (j jobs) Swap(a int, b int) {
	j[a], j[b] = j[b], j[a]
}

// elastic job filter.
func elastic(j job) bool {
	return j.Config.Elastic()
}

// gpu job filter.
func gpu(j job) bool {
	return j.Config.NeedGPU()
}

type eventType int

const (
	add eventType = iota
	del
)

type event struct {
	Type eventType
	Job  *paddlejob.TrainingJob
}

// OnAdd notifies the autoscaler that a job has been added.
func (a *Autoscaler) OnAdd(trainingjob *paddlejob.TrainingJob) {
	log.Debugln("OnAdd, adding job to event channel...", a.eventCh)
	a.eventCh <- event{Type: add, Job: trainingjob}
}

// OnDel notifies the autoscaler that a job has been deleted.
func (a *Autoscaler) OnDel(trainingjob *paddlejob.TrainingJob) {
	a.eventCh <- event{Type: del, Job: trainingjob}
}

// sortedJobs return the names of sorted jobs by fulfillment and
// tiebreakers in ascending order.
func sortedJobs(j []job, filters ...func(job) bool) []job {
	var js jobs
nextJob:
	for _, v := range j {
		for _, f := range filters {
			if !f(v) {
				continue nextJob
			}
		}
		js = append(js, v)
	}

	sort.Sort(js)
	return js
}

func scaleDryRun(r *ClusterResource, j job, curDiff int) (additional int) {
	additionalGPUInstance := 0
	additionalCPUInstance := 0
	cpuRequestMilli := j.TrainerCPURequestMilli()
	memRequestMega := j.TrainerMemRequestMega()
	gpuLimit := j.TrainerGPULimit()

	// Adjust resource upon return.
	defer func() {
		r.GPULimit += gpuLimit * additional
		r.CPURequestMilli += cpuRequestMilli * int64(additional)
		r.MemoryRequestMega += memRequestMega * int64(additional)
	}()

	// if we are going to scale down
	if r.GPULimit > r.GPUTotal || r.CPURequestMilli > r.CPUTotalMilli {
		return -1
	}

	plannedInstance := int(*j.TrainerJob.Spec.Parallelism) + curDiff
	instanceMax := j.Config.Spec.Trainer.MaxInstance

	if plannedInstance >= instanceMax {
		// Do not scale or scale down, don't need to check if
		// there are available free resources.
		additional = instanceMax - plannedInstance
		return
	}

	if r.MemoryTotalMega-r.MemoryRequestMega <= memRequestMega {
		// insufficient memory, do not scale
		additional = 0
		return
	}

	if r.CPUTotalMilli-r.CPURequestMilli >= cpuRequestMilli {
		additionalCPUInstance = 1
	}

	needGPU := gpuLimit > 0
	if needGPU && r.GPUTotal-r.GPULimit >= gpuLimit {
		additionalGPUInstance = 1
	}

	if needGPU {
		if additionalGPUInstance < additionalCPUInstance {
			additional = additionalGPUInstance
		} else {
			additional = additionalCPUInstance
		}
	} else {
		additional = additionalCPUInstance
	}

	return
}

func scaleAllDryRun(jobs []job, r ClusterResource) map[string]int {
	// Iteratively calculate scaling diff until nothing changes.
	diff := make(map[string]int)
	for {
		noChange := true
		sorted := sortedJobs(jobs, elastic)
		// TODO(typhoonzero): implement GPU priority CFS scheduler from here.
		for _, j := range sorted {
			name := j.Config.Name
			additional := scaleDryRun(&r, j, diff[name])
			log.Debugf("Dry run scale job %s: current %d, additional %d, remaining resource: %#v", name, diff[name], additional, r)
			diff[name] += additional

			if additional != 0 {
				noChange = false
			}
		}

		if noChange {
			break
		}
	}

	return diff
}

func (a *Autoscaler) scaleAll(diff map[string]int) {
	for name := range diff {
		if diff[name] != 0 {
			log.Infof("Scaling job %s, diff: %d.", name, diff[name])
			target := *a.jobs[name].TrainerJob.Spec.Parallelism + int32(diff[name])
			*a.jobs[name].TrainerJob.Spec.Parallelism = target
			err := a.cluster.UpdateTrainerJob(a.jobs[name].TrainerJob)
			if err != nil {
				log.Errorf("Error updating trainer job: %v", err)
			}
		}
	}
}

// Monitor monitors the cluster resources and training jobs in a loop,
// scales the training jobs according to the cluster resource.
func (a *Autoscaler) Monitor() {
	for {
		select {
		case <-a.ticker.C:
		case e := <-a.eventCh:
			switch e.Type {
			case add:
				// TODO(helin): schedule the training
				// k8s Job. Currently we don't
				// schedule the trainer job, but only
				// scale it.
				log.Debugln("AddJob to autoscaler: ", e.Job.ObjectMeta.Name)
				var tj *batchv1.Job
				var err error
				for {
					tj, err = a.cluster.GetTrainerJob(e.Job)
					if err == nil {
						break
					}

					log.Errorln(
						"Error getting the trainer job for %s, retrying soon.",
						e.Job.ObjectMeta.Name,
					)
					time.Sleep(3 * time.Second)
				}

				j := job{
					Config:     e.Job,
					TrainerJob: tj,
				}
				a.jobs[e.Job.ObjectMeta.Name] = j
			case del:
				// TODO(helin): delete all created
				// resources (e.g., trainer Job,
				// pserver Replica Set) when we
				// schedules the resources.
				log.Debugln("DelJob to autoscaler: ", e.Job.ObjectMeta.Name)
				delete(a.jobs, e.Job.ObjectMeta.Name)
			default:
				log.Errorln("Unrecognized event: %v.", e)
			}
		}

		r, err := a.cluster.SyncResource()
		if err != nil {
			log.Errorln("error sync resource: %v", err)
			continue
		}

		log.Infof("Latest cluster resource: %#v", r)
		var js []job
		for _, j := range a.jobs {
			// update batchv1.job instace every sync
			// because the resource version may be changed before scale
			tj, err := a.cluster.GetTrainerJob(j.Config)
			if err != nil {
				log.Warnln("sync trainer job error: ", err)
				continue
			}
			j.TrainerJob = tj

			// Scale jobs only when it's running (defined
			// by all pods are in the "running"
			// status). Pods are
			// pending/starting/terminating if the job is
			// just submited or just scaled up/down.
			running, err := a.cluster.JobRunning(j.Config)
			if err != nil {
				log.Errorln("Get if job is running failed:", err)
				continue
			}

			if running {
				js = append(js, j)
			}
		}
		diff := scaleAllDryRun(js, r)
		log.Infof("Scaling plan: %v", diff)
		a.scaleAll(diff)
	}
}
<commit_msg>Get the lastest TrainerJob before updating it, with retry.<commit_after>/* Copyright (c) 2016 PaddlePaddle Authors All Rights Reserve.

   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 autoscaler

import (
	"sort"
	"time"

	// TODO(typhoonzero): this package still depends on k8s API, try to remove this.
	"k8s.io/apimachinery/pkg/api/resource"
	batchv1 "k8s.io/client-go/pkg/apis/batch/v1"

	paddlejob "github.com/PaddlePaddle/cloud/go/api"
	log "github.com/sirupsen/logrus"
)

const (
	defaultLoopDur = time.Second * 5
)

// ClusterResource is the resource of a cluster
type ClusterResource struct {
	NodeCount int

	GPURequest int
	GPULimit   int
	GPUTotal   int

	CPURequestMilli int64
	CPULimitMilli   int64
	CPUTotalMilli   int64

	MemoryRequestMega int64
	MemoryLimitMega   int64
	MemoryTotalMega   int64
}

// Cluster represents the cluster managment system such as Kubernetes.
type Cluster interface {
	// SyncResource will sync resource values with the cluster.
	// should call this function in every tick.
	SyncResource() (ClusterResource, error)

	// GetTrainerJob gets the trainer job spec.
	GetTrainerJob(job *paddlejob.TrainingJob) (*batchv1.Job, error)

	// UpdateTrainerJob updates the trainer job spec.
	UpdateTrainerJob(job *batchv1.Job) error

	// JobRunning check if all the pods are in "Running" status.
	JobRunning(job *paddlejob.TrainingJob) (bool, error)
}

type job struct {
	Config     *paddlejob.TrainingJob
	TrainerJob *batchv1.Job
}

func (j job) TrainerGPULimit() int {
	q := j.Config.Spec.Trainer.Resources.Limits.NvidiaGPU()
	return int(q.Value())
}

func (j job) TrainerCPURequestMilli() int64 {
	q := j.Config.Spec.Trainer.Resources.Requests.Cpu()
	return q.ScaledValue(resource.Milli)
}

func (j job) TrainerMemRequestMega() int64 {
	q := j.Config.Spec.Trainer.Resources.Requests.Memory()
	return q.ScaledValue(resource.Mega)
}

func (j job) Fulfillment() float64 {
	minInstance := j.Config.Spec.Trainer.MinInstance
	maxInstance := j.Config.Spec.Trainer.MaxInstance

	if minInstance == maxInstance {
		return 1
	}

	curInstance := int(*j.TrainerJob.Spec.Parallelism)
	return float64(curInstance-minInstance) / float64(maxInstance-minInstance)
}

// Autoscaler launches and scales the training jobs.
type Autoscaler struct {
	ticker  *time.Ticker
	cluster Cluster
	jobs    map[string]job
	eventCh chan event
}

// New creates a new Autoscaler.
func New(cluster Cluster, options ...func(*Autoscaler)) *Autoscaler {
	c := &Autoscaler{
		cluster: cluster,
		ticker:  time.NewTicker(defaultLoopDur),
		jobs:    make(map[string]job),
		eventCh: make(chan event),
	}
	for _, option := range options {
		option(c)
	}
	return c
}

type jobs []job

func (j jobs) Len() int {
	return len(j)
}

func (j jobs) Less(a int, b int) bool {
	scoreA := j[a].Fulfillment()
	scoreB := j[b].Fulfillment()

	if scoreA == scoreB {
		resA := j[a].Config.Spec.Trainer.Resources
		resB := j[b].Config.Spec.Trainer.Resources
		resALimitsGPU := *resA.Limits.NvidiaGPU()
		resBLimitsGPU := *resB.Limits.NvidiaGPU()
		if resALimitsGPU.Cmp(resBLimitsGPU) == 0 {
			resARequestsCPU := *resA.Requests.Cpu()
			resBRequestsCPU := *resB.Requests.Cpu()
			if resARequestsCPU.Cmp(resBRequestsCPU) == 0 {
				resARequestsMem := *resA.Requests.Memory()
				resBRequestsMem := *resB.Requests.Memory()
				return resARequestsMem.Cmp(resBRequestsMem) == -1
			}
			return resARequestsCPU.Cmp(resBRequestsCPU) == -1
		}
		return resALimitsGPU.Cmp(resBLimitsGPU) == -1
	}
	return scoreA < scoreB
}

func (j jobs) Swap(a int, b int) {
	j[a], j[b] = j[b], j[a]
}

// elastic job filter.
func elastic(j job) bool {
	return j.Config.Elastic()
}

// gpu job filter.
func gpu(j job) bool {
	return j.Config.NeedGPU()
}

type eventType int

const (
	add eventType = iota
	del
)

type event struct {
	Type eventType
	Job  *paddlejob.TrainingJob
}

// OnAdd notifies the autoscaler that a job has been added.
func (a *Autoscaler) OnAdd(trainingjob *paddlejob.TrainingJob) {
	a.eventCh <- event{Type: add, Job: trainingjob}
}

// OnDel notifies the autoscaler that a job has been deleted.
func (a *Autoscaler) OnDel(trainingjob *paddlejob.TrainingJob) {
	a.eventCh <- event{Type: del, Job: trainingjob}
}

// sortedJobs return the names of sorted jobs by fulfillment and
// tiebreakers in ascending order.
func sortedJobs(j []job, filters ...func(job) bool) []job {
	var js jobs
nextJob:
	for _, v := range j {
		for _, f := range filters {
			if !f(v) {
				continue nextJob
			}
		}
		js = append(js, v)
	}

	sort.Sort(js)
	return js
}

func scaleDryRun(r *ClusterResource, j job, curDiff int) (additional int) {
	additionalGPUInstance := 0
	additionalCPUInstance := 0
	cpuRequestMilli := j.TrainerCPURequestMilli()
	memRequestMega := j.TrainerMemRequestMega()
	gpuLimit := j.TrainerGPULimit()

	// Adjust resource upon return.
	defer func() {
		r.GPULimit += gpuLimit * additional
		r.CPURequestMilli += cpuRequestMilli * int64(additional)
		r.MemoryRequestMega += memRequestMega * int64(additional)
	}()

	// if we are going to scale down
	if r.GPULimit > r.GPUTotal || r.CPURequestMilli > r.CPUTotalMilli {
		return -1
	}

	plannedInstance := int(*j.TrainerJob.Spec.Parallelism) + curDiff
	instanceMax := j.Config.Spec.Trainer.MaxInstance

	if plannedInstance >= instanceMax {
		// Do not scale or scale down, don't need to check if
		// there are available free resources.
		additional = instanceMax - plannedInstance
		return
	}

	if r.MemoryTotalMega-r.MemoryRequestMega <= memRequestMega {
		// insufficient memory, do not scale
		additional = 0
		return
	}

	if r.CPUTotalMilli-r.CPURequestMilli >= cpuRequestMilli {
		additionalCPUInstance = 1
	}

	needGPU := gpuLimit > 0
	if needGPU && r.GPUTotal-r.GPULimit >= gpuLimit {
		additionalGPUInstance = 1
	}

	if needGPU {
		if additionalGPUInstance < additionalCPUInstance {
			additional = additionalGPUInstance
		} else {
			additional = additionalCPUInstance
		}
	} else {
		additional = additionalCPUInstance
	}

	return
}

func scaleAllDryRun(jobs []job, r ClusterResource) map[string]int {
	// Iteratively calculate scaling diff until nothing changes.
	diff := make(map[string]int)
	for {
		noChange := true
		sorted := sortedJobs(jobs, elastic)
		// TODO(typhoonzero): implement GPU priority CFS scheduler from here.
		for _, j := range sorted {
			name := j.Config.Name
			additional := scaleDryRun(&r, j, diff[name])
			log.Debugf("Dry run scale job %s: current %d, additional %d, remaining resource: %#v", name, diff[name], additional, r)
			diff[name] += additional

			if additional != 0 {
				noChange = false
			}
		}

		if noChange {
			break
		}
	}

	return diff
}

func (a *Autoscaler) scaleAll(diff map[string]int) {
	for name := range diff {
		if diff[name] != 0 {
			log.Infof("Scaling job %s, diff: %d.", name, diff[name])
			target := *a.jobs[name].TrainerJob.Spec.Parallelism + int32(diff[name])

			var err error
			for retry := 0; retry < 3; retry++ {
				tj, err := a.cluster.GetTrainerJob(a.jobs[name].Config)
				if err != nil {
					log.Warnf("sync trainer job error: %v, retry remaining: %d ", err, retry)
					continue
				}
				j := a.jobs[name]
				j.TrainerJob = tj
				a.jobs[name] = j
				*a.jobs[name].TrainerJob.Spec.Parallelism = target

				err = a.cluster.UpdateTrainerJob(a.jobs[name].TrainerJob)
				if err != nil {
					log.Errorf("Error updating trainer job %v, retry remaining: %d.", err, retry)
					continue
				}

				break
			}

			if err != nil {
				log.Warnf("Error updating trainer job %v.", err)
			}
		}
	}
}

// Monitor monitors the cluster resources and training jobs in a loop,
// scales the training jobs according to the cluster resource.
func (a *Autoscaler) Monitor() {
	for {
		select {
		case <-a.ticker.C:
		case e := <-a.eventCh:
			switch e.Type {
			case add:
				// TODO(helin): schedule the training
				// k8s Job. Currently we don't
				// schedule the trainer job, but only
				// scale it.
				log.Debugln("AddJob to autoscaler: ", e.Job.ObjectMeta.Name)
				var tj *batchv1.Job
				var err error
				for {
					tj, err = a.cluster.GetTrainerJob(e.Job)
					if err == nil {
						break
					}

					log.Errorln(
						"Error getting the trainer job for %s, retrying soon.",
						e.Job.ObjectMeta.Name,
					)
					time.Sleep(3 * time.Second)
				}

				j := job{
					Config:     e.Job,
					TrainerJob: tj,
				}
				a.jobs[e.Job.ObjectMeta.Name] = j
			case del:
				// TODO(helin): delete all created
				// resources (e.g., trainer Job,
				// pserver Replica Set) when we
				// schedules the resources.
				log.Debugln("DelJob to autoscaler: ", e.Job.ObjectMeta.Name)
				delete(a.jobs, e.Job.ObjectMeta.Name)
			default:
				log.Errorln("Unrecognized event: %v.", e)
			}
		}

		r, err := a.cluster.SyncResource()
		if err != nil {
			log.Errorln("error sync resource: %v", err)
			continue
		}

		log.Infof("Latest cluster resource: %#v", r)
		var js []job
		for _, j := range a.jobs {
			// update batchv1.job instace every sync
			// because the resource version may be changed before scale
			tj, err := a.cluster.GetTrainerJob(j.Config)
			if err != nil {
				log.Warnln("sync trainer job error: ", err)
				continue
			}
			j.TrainerJob = tj

			// Scale jobs only when it's running (defined
			// by all pods are in the "running"
			// status). Pods are
			// pending/starting/terminating if the job is
			// just submited or just scaled up/down.
			running, err := a.cluster.JobRunning(j.Config)
			if err != nil {
				log.Errorln("Get if job is running failed:", err)
				continue
			}

			if running {
				js = append(js, j)
			}
		}
		diff := scaleAllDryRun(js, r)
		log.Infof("Scaling plan: %v", diff)
		a.scaleAll(diff)
	}
}
<|endoftext|>


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

// Package buildutil provides utilities related to the go/build
// package in the standard library.
package buildutil

import (
	"go/build"
	"io/ioutil"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
)

// AllPackages returns the import path of each Go package in any source
// directory of the specified build context (e.g. $GOROOT or an element
// of $GOPATH).  Errors are ignored.  The results are sorted.
//
// The result may include import paths for directories that contain no
// *.go files, such as "archive" (in $GOROOT/src).
//
// All I/O is via the build.Context virtual file system,
// which must be concurrency-safe.
//
func AllPackages(ctxt *build.Context) []string {
	var list []string
	var mu sync.Mutex
	ForEachPackage(ctxt, func(pkg string, _ error) {
		mu.Lock()
		list = append(list, pkg)
		mu.Unlock()
	})
	sort.Strings(list)
	return list
}

// ForEachPackage calls the found function with the import path of
// each Go package it finds in any source directory of the specified
// build context (e.g. $GOROOT or an element of $GOPATH).
//
// If the package directory exists but could not be read, the second
// argument to the found function provides the error.
//
// The found function and the build.Context virtual file system
// accessors must be concurrency safe.
//
func ForEachPackage(ctxt *build.Context, found func(importPath string, err error)) {
	// We use a counting semaphore to limit
	// the number of parallel calls to ReadDir.
	sema := make(chan bool, 20)

	var wg sync.WaitGroup
	for _, root := range ctxt.SrcDirs() {
		root := root
		wg.Add(1)
		go func() {
			allPackages(ctxt, sema, root, found)
			wg.Done()
		}()
	}
	wg.Wait()
}

func allPackages(ctxt *build.Context, sema chan bool, root string, found func(string, error)) {
	ReadDir := ctxt.ReadDir
	if ReadDir == nil {
		ReadDir = ioutil.ReadDir
	}

	root = filepath.Clean(root) + string(os.PathSeparator)

	var wg sync.WaitGroup

	var walkDir func(dir string)
	walkDir = func(dir string) {
		// Prune search if we encounter any directory with these base names:
		switch filepath.Base(dir) {
		case "testdata", ".hg":
			return
		}

		pkg := filepath.ToSlash(strings.TrimPrefix(dir, root))

		// Prune search if we encounter any of these import paths.
		switch pkg {
		case "builtin":
			return
		}

		sema <- true
		files, err := ReadDir(dir)
		<-sema
		if pkg != "" || err != nil {
			found(pkg, err)
		}
		for _, fi := range files {
			fi := fi
			if fi.IsDir() {
				wg.Add(1)
				go func() {
					walkDir(filepath.Join(dir, fi.Name()))
					wg.Done()
				}()
			}
		}
	}

	walkDir(root)
	wg.Wait()
}
<commit_msg>go/buildutil: use same logic as 'go' tool for pruning package search.<commit_after>// Copyright 2014 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package buildutil provides utilities related to the go/build
// package in the standard library.
package buildutil

import (
	"go/build"
	"io/ioutil"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
)

// AllPackages returns the import path of each Go package in any source
// directory of the specified build context (e.g. $GOROOT or an element
// of $GOPATH).  Errors are ignored.  The results are sorted.
//
// The result may include import paths for directories that contain no
// *.go files, such as "archive" (in $GOROOT/src).
//
// All I/O is via the build.Context virtual file system,
// which must be concurrency-safe.
//
func AllPackages(ctxt *build.Context) []string {
	var list []string
	var mu sync.Mutex
	ForEachPackage(ctxt, func(pkg string, _ error) {
		mu.Lock()
		list = append(list, pkg)
		mu.Unlock()
	})
	sort.Strings(list)
	return list
}

// ForEachPackage calls the found function with the import path of
// each Go package it finds in any source directory of the specified
// build context (e.g. $GOROOT or an element of $GOPATH).
//
// If the package directory exists but could not be read, the second
// argument to the found function provides the error.
//
// The found function and the build.Context virtual file system
// accessors must be concurrency safe.
//
func ForEachPackage(ctxt *build.Context, found func(importPath string, err error)) {
	// We use a counting semaphore to limit
	// the number of parallel calls to ReadDir.
	sema := make(chan bool, 20)

	var wg sync.WaitGroup
	for _, root := range ctxt.SrcDirs() {
		root := root
		wg.Add(1)
		go func() {
			allPackages(ctxt, sema, root, found)
			wg.Done()
		}()
	}
	wg.Wait()
}

func allPackages(ctxt *build.Context, sema chan bool, root string, found func(string, error)) {
	ReadDir := ctxt.ReadDir
	if ReadDir == nil {
		ReadDir = ioutil.ReadDir
	}

	root = filepath.Clean(root) + string(os.PathSeparator)

	var wg sync.WaitGroup

	var walkDir func(dir string)
	walkDir = func(dir string) {
		// Avoid .foo, _foo, and testdata directory trees.
		base := filepath.Base(dir)
		if base == "" || base[0] == '.' || base[0] == '_' || base == "testdata" {
			return
		}

		pkg := filepath.ToSlash(strings.TrimPrefix(dir, root))

		// Prune search if we encounter any of these import paths.
		switch pkg {
		case "builtin":
			return
		}

		sema <- true
		files, err := ReadDir(dir)
		<-sema
		if pkg != "" || err != nil {
			found(pkg, err)
		}
		for _, fi := range files {
			fi := fi
			if fi.IsDir() {
				wg.Add(1)
				go func() {
					walkDir(filepath.Join(dir, fi.Name()))
					wg.Done()
				}()
			}
		}
	}

	walkDir(root)
	wg.Wait()
}
<|endoftext|>


<commit_before>package basex

import (
	"io"
)

// Much of this code is adopted from Go's encoding/base64

// EncodeToString returns the base64 encoding of src.
func (enc *Encoding) EncodeToString(src []byte) string {
	buf := make([]byte, enc.EncodedLen(len(src)))
	enc.Encode(buf, src)
	return string(buf)
}

type encoder struct {
	err  error
	enc  *Encoding
	w    io.Writer
	buf  []byte // buffered data waiting to be encoded
	nbuf int    // number of bytes in buf
	out  []byte // output buffer
}

func (e *encoder) Write(p []byte) (n int, err error) {
	if e.err != nil {
		return 0, e.err
	}

	ibl := e.enc.inBlockLen
	obl := e.enc.outBlockLen

	// Leading fringe.
	if e.nbuf > 0 {
		var i int
		for i = 0; i < len(p) && e.nbuf < ibl; i++ {
			e.buf[e.nbuf] = p[i]
			e.nbuf++
		}
		n += i
		p = p[i:]
		if e.nbuf < ibl {
			return
		}
		e.enc.Encode(e.out[:], e.buf[:])
		if _, e.err = e.w.Write(e.out[:obl]); e.err != nil {
			return n, e.err
		}
		e.nbuf = 0
	}

	// Large interior chunks.
	for len(p) >= ibl {
		nn := len(e.out) / obl * ibl
		if nn > len(p) {
			nn = len(p)
			nn -= nn % ibl
		}
		e.enc.Encode(e.out[:], p[:nn])
		if _, e.err = e.w.Write(e.out[0 : nn/ibl*obl]); e.err != nil {
			return n, e.err
		}
		n += nn
		p = p[nn:]
	}

	// Trailing fringe.
	copy(e.buf[0:len(p)], p)
	e.nbuf = len(p)
	n += len(p)
	return
}

// Close flushes any pending output from the encoder.
// It is an error to call Write after calling Close.
func (e *encoder) Close() error {
	// If there's anything left in the buffer, flush it out
	if e.err == nil && e.nbuf > 0 {
		e.enc.Encode(e.out[:], e.buf[:e.nbuf])
		_, e.err = e.w.Write(e.out[:e.enc.EncodedLen(e.nbuf)])
		e.nbuf = 0
	}
	return e.err
}

// NewEncoder returns a new baseX stream encoder.  Data written to
// the returned writer will be encoded using enc and then written to w.
// Base64 encodings operate in enc.outBlockLen-byte blocks; when finished
// writing, the caller must Close the returned encoder to flush any
// partially written blocks.
func NewEncoder(enc *Encoding, w io.Writer) io.WriteCloser {
	return &encoder{
		enc: enc,
		w:   w,
		buf: make([]byte, enc.inBlockLen),
		out: make([]byte, 128*enc.outBlockLen),
	}
}

// DecodeStringStrict returns the bytes represented by the baseX string s.
// It uses the strict decoding strategy, not allowing any non-baseX-characters
func (enc *Encoding) DecodeStringStrict(s string) ([]byte, error) {
	dbuf := make([]byte, enc.DecodedLen(len(s)))
	n, err := enc.DecodeStrict(dbuf, []byte(s))
	return dbuf[:n], err
}

// DecodeString returns the bytes represented by the baseX string s.
// It uses the liberal decoding strategy, ignoring any non-baseX-characters
func (enc *Encoding) DecodeString(s string) ([]byte, error) {
	dbuf := make([]byte, enc.DecodedLen(len(s)))
	n, err := enc.Decode(dbuf, []byte(s))
	return dbuf[:n], err
}

type decoder struct {
	err        error
	enc        *Encoding
	r          io.Reader
	out        []byte // leftover decoded output
	buf        []byte // leftover input
	nbuf       int    // the begin pointer of buf above
	scratchbuf []byte // a temporary scratch buf, for reuse
}

func (d *decoder) Read(p []byte) (int, error) {

	if d.err != nil {
		return 0, d.err
	}

	// Use leftover decoded output from last read.
	if len(d.out) > 0 {
		ret := copy(p, d.out)
		d.out = d.out[ret:]
		return ret, nil
	}

	ibl := d.enc.inBlockLen
	obl := d.enc.outBlockLen

	nn := len(p) / ibl * obl
	if nn < obl {
		nn = obl
	}
	if nn > len(d.buf) {
		nn = len(d.buf)
	}

	// Try to read up to the next full block. We already have d.nbuf in
	// there. Need another (obl - d.nbuf) to round up.
	nn, d.err = io.ReadAtLeast(d.r, d.buf[d.nbuf:nn], obl-d.nbuf)
	d.nbuf += nn

	eof := false

	// This condition is actually OK, we just shouldn't read any more data
	// afterwards. We should get an EOF the next time through.
	if d.err == io.ErrUnexpectedEOF {
		d.err = nil
		eof = true
	} else if d.err == io.EOF {
		if d.nbuf == 0 {
			return 0, d.err
		}
		eof = true
		d.err = nil
	} else if d.err != nil {
		return 0, d.err
	}

	// The num bytes to decode should be along obl-aligned boundaries, unless
	// we're at the end of file.
	numBytesToDecode := d.nbuf
	if !eof {
		numBytesToDecode = numBytesToDecode / obl * obl
	}
	numBytesToOutput := d.enc.DecodedLen(numBytesToDecode)

	var ret int

	// If we have too many bytes for the given buffer, we can buffer
	// the rest internally
	if numBytesToOutput > len(p) {
		var n int
		n, d.err = d.enc.DecodeStrict(d.scratchbuf[:], d.buf[:numBytesToDecode])
		d.out = d.scratchbuf[:n]
		ret = copy(p, d.out)
		d.out = d.out[ret:]
	} else {
		ret, d.err = d.enc.Decode(p, d.buf[:numBytesToDecode])
	}

	// Shift the bytes in d.buf over from [numBytesToDecode:] to the start of the array
	d.nbuf -= numBytesToDecode
	copy(d.buf[0:d.nbuf], d.buf[numBytesToDecode:numBytesToDecode+d.nbuf])

	return ret, d.err
}

type filteringReader struct {
	wrapped io.Reader
	enc     *Encoding
}

func (r *filteringReader) Read(p []byte) (int, error) {
	n, err := r.wrapped.Read(p)
	for n > 0 {
		offset := 0
		for i, b := range p[:n] {
			if r.enc.decodeMap[b] != 0xFF {
				if i != offset {
					p[offset] = b
				}
				offset++
			}
		}
		if offset > 0 {
			return offset, err
		}
		// Previous buffer entirely whitespace, read again
		n, err = r.wrapped.Read(p)
	}
	return n, err
}

// NewDecoder constructs a new baseX stream decoder.
func NewDecoder(enc *Encoding, r io.Reader) io.Reader {
	return newDecoder(enc, r, false)
}

// NewDecoderStrict constructs a new baseX stream decoder, but will return an
// error on any non-base-X character input
func NewDecoderStrict(enc *Encoding, r io.Reader) io.Reader {
	return newDecoder(enc, r, true)
}

func newDecoder(enc *Encoding, r io.Reader, strict bool) io.Reader {
	if !strict {
		r = &filteringReader{r, enc}
	}
	return &decoder{
		enc:        enc,
		r:          r,
		buf:        make([]byte, 128*enc.inBlockLen),
		scratchbuf: make([]byte, 128*enc.outBlockLen),
	}
}
<commit_msg>fix comment typo<commit_after>package basex

import (
	"io"
)

// Much of this code is adopted from Go's encoding/base64

// EncodeToString returns the base64 encoding of src.
func (enc *Encoding) EncodeToString(src []byte) string {
	buf := make([]byte, enc.EncodedLen(len(src)))
	enc.Encode(buf, src)
	return string(buf)
}

type encoder struct {
	err  error
	enc  *Encoding
	w    io.Writer
	buf  []byte // buffered data waiting to be encoded
	nbuf int    // number of bytes in buf
	out  []byte // output buffer
}

func (e *encoder) Write(p []byte) (n int, err error) {
	if e.err != nil {
		return 0, e.err
	}

	ibl := e.enc.inBlockLen
	obl := e.enc.outBlockLen

	// Leading fringe.
	if e.nbuf > 0 {
		var i int
		for i = 0; i < len(p) && e.nbuf < ibl; i++ {
			e.buf[e.nbuf] = p[i]
			e.nbuf++
		}
		n += i
		p = p[i:]
		if e.nbuf < ibl {
			return
		}
		e.enc.Encode(e.out[:], e.buf[:])
		if _, e.err = e.w.Write(e.out[:obl]); e.err != nil {
			return n, e.err
		}
		e.nbuf = 0
	}

	// Large interior chunks.
	for len(p) >= ibl {
		nn := len(e.out) / obl * ibl
		if nn > len(p) {
			nn = len(p)
			nn -= nn % ibl
		}
		e.enc.Encode(e.out[:], p[:nn])
		if _, e.err = e.w.Write(e.out[0 : nn/ibl*obl]); e.err != nil {
			return n, e.err
		}
		n += nn
		p = p[nn:]
	}

	// Trailing fringe.
	copy(e.buf[0:len(p)], p)
	e.nbuf = len(p)
	n += len(p)
	return
}

// Close flushes any pending output from the encoder.
// It is an error to call Write after calling Close.
func (e *encoder) Close() error {
	// If there's anything left in the buffer, flush it out
	if e.err == nil && e.nbuf > 0 {
		e.enc.Encode(e.out[:], e.buf[:e.nbuf])
		_, e.err = e.w.Write(e.out[:e.enc.EncodedLen(e.nbuf)])
		e.nbuf = 0
	}
	return e.err
}

// NewEncoder returns a new baseX stream encoder.  Data written to
// the returned writer will be encoded using enc and then written to w.
// Encodings operate in enc.outBlockLen-byte blocks; when finished
// writing, the caller must Close the returned encoder to flush any
// partially written blocks.
func NewEncoder(enc *Encoding, w io.Writer) io.WriteCloser {
	return &encoder{
		enc: enc,
		w:   w,
		buf: make([]byte, enc.inBlockLen),
		out: make([]byte, 128*enc.outBlockLen),
	}
}

// DecodeStringStrict returns the bytes represented by the baseX string s.
// It uses the strict decoding strategy, not allowing any non-baseX-characters
func (enc *Encoding) DecodeStringStrict(s string) ([]byte, error) {
	dbuf := make([]byte, enc.DecodedLen(len(s)))
	n, err := enc.DecodeStrict(dbuf, []byte(s))
	return dbuf[:n], err
}

// DecodeString returns the bytes represented by the baseX string s.
// It uses the liberal decoding strategy, ignoring any non-baseX-characters
func (enc *Encoding) DecodeString(s string) ([]byte, error) {
	dbuf := make([]byte, enc.DecodedLen(len(s)))
	n, err := enc.Decode(dbuf, []byte(s))
	return dbuf[:n], err
}

type decoder struct {
	err        error
	enc        *Encoding
	r          io.Reader
	out        []byte // leftover decoded output
	buf        []byte // leftover input
	nbuf       int    // the begin pointer of buf above
	scratchbuf []byte // a temporary scratch buf, for reuse
}

func (d *decoder) Read(p []byte) (int, error) {

	if d.err != nil {
		return 0, d.err
	}

	// Use leftover decoded output from last read.
	if len(d.out) > 0 {
		ret := copy(p, d.out)
		d.out = d.out[ret:]
		return ret, nil
	}

	ibl := d.enc.inBlockLen
	obl := d.enc.outBlockLen

	nn := len(p) / ibl * obl
	if nn < obl {
		nn = obl
	}
	if nn > len(d.buf) {
		nn = len(d.buf)
	}

	// Try to read up to the next full block. We already have d.nbuf in
	// there. Need another (obl - d.nbuf) to round up.
	nn, d.err = io.ReadAtLeast(d.r, d.buf[d.nbuf:nn], obl-d.nbuf)
	d.nbuf += nn

	eof := false

	// This condition is actually OK, we just shouldn't read any more data
	// afterwards. We should get an EOF the next time through.
	if d.err == io.ErrUnexpectedEOF {
		d.err = nil
		eof = true
	} else if d.err == io.EOF {
		if d.nbuf == 0 {
			return 0, d.err
		}
		eof = true
		d.err = nil
	} else if d.err != nil {
		return 0, d.err
	}

	// The num bytes to decode should be along obl-aligned boundaries, unless
	// we're at the end of file.
	numBytesToDecode := d.nbuf
	if !eof {
		numBytesToDecode = numBytesToDecode / obl * obl
	}
	numBytesToOutput := d.enc.DecodedLen(numBytesToDecode)

	var ret int

	// If we have too many bytes for the given buffer, we can buffer
	// the rest internally
	if numBytesToOutput > len(p) {
		var n int
		n, d.err = d.enc.DecodeStrict(d.scratchbuf[:], d.buf[:numBytesToDecode])
		d.out = d.scratchbuf[:n]
		ret = copy(p, d.out)
		d.out = d.out[ret:]
	} else {
		ret, d.err = d.enc.Decode(p, d.buf[:numBytesToDecode])
	}

	// Shift the bytes in d.buf over from [numBytesToDecode:] to the start of the array
	d.nbuf -= numBytesToDecode
	copy(d.buf[0:d.nbuf], d.buf[numBytesToDecode:numBytesToDecode+d.nbuf])

	return ret, d.err
}

type filteringReader struct {
	wrapped io.Reader
	enc     *Encoding
}

func (r *filteringReader) Read(p []byte) (int, error) {
	n, err := r.wrapped.Read(p)
	for n > 0 {
		offset := 0
		for i, b := range p[:n] {
			if r.enc.decodeMap[b] != 0xFF {
				if i != offset {
					p[offset] = b
				}
				offset++
			}
		}
		if offset > 0 {
			return offset, err
		}
		// Previous buffer entirely whitespace, read again
		n, err = r.wrapped.Read(p)
	}
	return n, err
}

// NewDecoder constructs a new baseX stream decoder.
func NewDecoder(enc *Encoding, r io.Reader) io.Reader {
	return newDecoder(enc, r, false)
}

// NewDecoderStrict constructs a new baseX stream decoder, but will return an
// error on any non-base-X character input
func NewDecoderStrict(enc *Encoding, r io.Reader) io.Reader {
	return newDecoder(enc, r, true)
}

func newDecoder(enc *Encoding, r io.Reader, strict bool) io.Reader {
	if !strict {
		r = &filteringReader{r, enc}
	}
	return &decoder{
		enc:        enc,
		r:          r,
		buf:        make([]byte, 128*enc.inBlockLen),
		scratchbuf: make([]byte, 128*enc.outBlockLen),
	}
}
<|endoftext|>


<commit_before>// Package moreflag contains definitions for some useful flag types, such as maps.
package moreflag

import (
	"bytes"
	"flag"
	"fmt"
	"os"
	"sort"
	"strings"
)

// Parse parses flags which are set in environment variables using the FLAG_ prefix. That is
// for a flag named x, x=$FLAG_x if $FLAG_x is set. If the flag is set in the command line, the
// command line value of the flag takes precedence over the environment variable value.
// It also calls flag.Parse() to parse flags sent directly as arguments, unless flag.Parse
// has been previously called.
func Parse() {
	ParseFromEnv()
	if !flag.Parsed() {
		flag.Parse()
	}
}

// ParseFromEnv parses flags which are set in environment variables using the FLAG_ prefix. That is
// for a flag named x, x=$FLAG_x if $FLAG_x is set. If the flag is set in the command line, the
// command line value of the flag takes precedence over the environment variable value.
func ParseFromEnv() {
	flag.VisitAll(func(f *flag.Flag) {
		v := os.Getenv("FLAG_" + f.Name)
		if v != "" {
			flag.Set(f.Name, v)
		}
	})
}

// StringMapValue is a command line flag that interprets a string in the format key1=value1,key2=value2
// as a map.
type StringMapValue map[string]string

// String retrieves the flag's map in the format key1=value1,key2=value, sorted by keys.
func (m *StringMapValue) String() string {
	// Construct the output in key sorted order
	keys := make([]string, 0, len(*m))
	for key := range *m {
		keys = append(keys, key)
	}
	sort.Strings(keys)
	var b bytes.Buffer
	for i, key := range keys {
		if i > 0 {
			b.WriteRune(',')
		}
		b.WriteString(key)
		b.WriteRune('=')
		b.WriteString((*m)[key])
	}
	return b.String()
}

// Set updates the map with key and value pair(s) in the format key1=value1,key2=value2.
func (m *StringMapValue) Set(s string) error {
	*m = make(map[string]string)
	pairs := strings.Split(s, ",")
	for _, p := range pairs {
		if p == "" {
			continue
		}
		pair := strings.Split(p, "=")
		if len(pair) != 2 {
			return fmt.Errorf("wrong format for key-value pair: %v", p)
		}
		if pair[0] == "" {
			return fmt.Errorf("key not provided")
		}
		if _, ok := (*m)[pair[0]]; ok {
			return fmt.Errorf("key %v already defined in list of key-value pairs %v", pair[0], s)
		}
		(*m)[pair[0]] = pair[1]
	}
	return nil
}

// StringListValue is a command line flag that interprets a string as a list of comma-separated values.
type StringListValue []string

// String returns the list value.
func (m *StringListValue) String() string {
	return strings.Join(*m, ",")
}

// Set for StringListValue accepts one list of comma-separated values.
func (m *StringListValue) Set(s string) error {
	splitFn := func(c rune) bool {
		return c == ','
	}
	*m = StringListValue(strings.FieldsFunc(s, splitFn))
	return nil
}
<commit_msg>Fixing broken internal builds -- flag Value needs to implement Get. (#85)<commit_after>// Package moreflag contains definitions for some useful flag types, such as maps.
package moreflag

import (
	"bytes"
	"flag"
	"fmt"
	"os"
	"sort"
	"strings"
)

// Parse parses flags which are set in environment variables using the FLAG_ prefix. That is
// for a flag named x, x=$FLAG_x if $FLAG_x is set. If the flag is set in the command line, the
// command line value of the flag takes precedence over the environment variable value.
// It also calls flag.Parse() to parse flags sent directly as arguments, unless flag.Parse
// has been previously called.
func Parse() {
	ParseFromEnv()
	if !flag.Parsed() {
		flag.Parse()
	}
}

// ParseFromEnv parses flags which are set in environment variables using the FLAG_ prefix. That is
// for a flag named x, x=$FLAG_x if $FLAG_x is set. If the flag is set in the command line, the
// command line value of the flag takes precedence over the environment variable value.
func ParseFromEnv() {
	flag.VisitAll(func(f *flag.Flag) {
		v := os.Getenv("FLAG_" + f.Name)
		if v != "" {
			flag.Set(f.Name, v)
		}
	})
}

// StringMapValue is a command line flag that interprets a string in the format key1=value1,key2=value2
// as a map.
type StringMapValue map[string]string

// String retrieves the flag's map in the format key1=value1,key2=value, sorted by keys.
func (m *StringMapValue) String() string {
	// Construct the output in key sorted order
	keys := make([]string, 0, len(*m))
	for key := range *m {
		keys = append(keys, key)
	}
	sort.Strings(keys)
	var b bytes.Buffer
	for i, key := range keys {
		if i > 0 {
			b.WriteRune(',')
		}
		b.WriteString(key)
		b.WriteRune('=')
		b.WriteString((*m)[key])
	}
	return b.String()
}

// Set updates the map with key and value pair(s) in the format key1=value1,key2=value2.
func (m *StringMapValue) Set(s string) error {
	*m = make(map[string]string)
	pairs := strings.Split(s, ",")
	for _, p := range pairs {
		if p == "" {
			continue
		}
		pair := strings.Split(p, "=")
		if len(pair) != 2 {
			return fmt.Errorf("wrong format for key-value pair: %v", p)
		}
		if pair[0] == "" {
			return fmt.Errorf("key not provided")
		}
		if _, ok := (*m)[pair[0]]; ok {
			return fmt.Errorf("key %v already defined in list of key-value pairs %v", pair[0], s)
		}
		(*m)[pair[0]] = pair[1]
	}
	return nil
}

// Get returns the flag value as a map of strings.
func (m *StringMapValue) Get() interface{} {
	return map[string]string(*m)
}

// StringListValue is a command line flag that interprets a string as a list of comma-separated values.
type StringListValue []string

// String returns the list value.
func (m *StringListValue) String() string {
	return strings.Join(*m, ",")
}

// Set for StringListValue accepts one list of comma-separated values.
func (m *StringListValue) Set(s string) error {
	splitFn := func(c rune) bool {
		return c == ','
	}
	*m = StringListValue(strings.FieldsFunc(s, splitFn))
	return nil
}

// Get returns the flag value as a list of strings.
func (m *StringListValue) Get() interface{} {
	return []string(*m)
}
<|endoftext|>


<commit_before>package zeusclient

import (
	"net"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"

	"github.com/burke/pty"
	"github.com/burke/ttyutils"
	slog "github.com/burke/zeus/go/shinylog"
	"github.com/burke/zeus/go/unixsocket"
)

const (
	zeusSockName = ".zeus.sock"
	sigInt       = 3 // todo: this doesn't seem unicode-friendly...
	sigQuit      = 28
	sigTstp      = 26
)

func Run(color bool) {
	os.Exit(doRun(color))
}

func doRun(color bool) int {
	if !color {
		slog.DisableColor()
		DisableErrorColor()
	}

	if os.Getenv("RAILS_ENV") != "" {
		println("Warning: Specifying a Rails environment via RAILS_ENV has no effect for commands run with zeus.")
	}

	master, slave, err := pty.Open()
	if err != nil {
		panic(err)
	}
	defer master.Close()
	if ttyutils.IsTerminal(os.Stdout.Fd()) {
		oldState, err := ttyutils.MakeTerminalRaw(os.Stdout.Fd())
		if err != nil {
			panic(err)
		}
		defer ttyutils.RestoreTerminalState(os.Stdout.Fd(), oldState)
	}

	ttyutils.MirrorWinsize(os.Stdout, master)

	addr, err := net.ResolveUnixAddr("unixgram", zeusSockName)
	if err != nil {
		panic("Can't resolve server address")
	}

	conn, err := net.DialUnix("unix", nil, addr)
	if err != nil {
		ErrorCantConnectToMaster()
	}
	usock := unixsocket.NewUsock(conn)

	msg := CreateCommandAndArgumentsMessage(os.Args[1], os.Args[2:])
	usock.WriteMessage(msg)
	usock.WriteFD(int(slave.Fd()))
	slave.Close()

	msg, err = usock.ReadMessage()
	if err != nil {
		panic(err)
	}

	parts := strings.Split(msg, "\000")
	commandPid, err := strconv.Atoi(parts[0])
	defer func() {
		if commandPid > 0 {
			// Just in case.
			syscall.Kill(commandPid, 9)
		}
	}()

	if err != nil {
		panic(err)
	}

	c := make(chan os.Signal, 1)
	signal.Notify(c, syscall.SIGWINCH, syscall.SIGCONT)
	go func() {
		for sig := range c {
			if sig == syscall.SIGCONT {
				syscall.Kill(commandPid, syscall.SIGCONT)
			} else if sig == syscall.SIGWINCH {
				ttyutils.MirrorWinsize(os.Stdout, master)
				syscall.Kill(commandPid, syscall.SIGWINCH)
			}
		}
	}()

	var exitStatus int = -1
	if len(parts) > 2 {
		exitStatus, err = strconv.Atoi(parts[0])
		if err != nil {
			panic(err)
		}
	}

	eof := make(chan bool)
	go func() {
		for {
			buf := make([]byte, 1024)
			n, err := master.Read(buf)
			if err != nil {
				eof <- true
				break
			}
			os.Stdout.Write(buf[:n])
		}
	}()

	go func() {
		buf := make([]byte, 8192)
		for {
			n, err := os.Stdin.Read(buf)
			if err != nil {
				eof <- true
				break
			}
			for i := 0; i < n; i++ {
				switch buf[i] {
				case sigInt:
					syscall.Kill(commandPid, syscall.SIGINT)
				case sigQuit:
					syscall.Kill(commandPid, syscall.SIGQUIT)
				case sigTstp:
					syscall.Kill(commandPid, syscall.SIGTSTP)
					syscall.Kill(os.Getpid(), syscall.SIGTSTP)
				}
			}
			master.Write(buf[:n])
		}
	}()

	<-eof

	if exitStatus == -1 {
		msg, err = usock.ReadMessage()
		if err != nil {
			panic(err)
		}
		parts := strings.Split(msg, "\000")
		exitStatus, err = strconv.Atoi(parts[0])
		if err != nil {
			panic(err)
		}
	}

	return exitStatus
}
<commit_msg>Use a socket instead of a PTY if redirecting output [closes #90]<commit_after>package zeusclient

import (
	"net"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"

	"github.com/burke/pty"
	"github.com/burke/ttyutils"
	slog "github.com/burke/zeus/go/shinylog"
	"github.com/burke/zeus/go/unixsocket"
)

const (
	zeusSockName = ".zeus.sock"
	sigInt       = 3 // todo: this doesn't seem unicode-friendly...
	sigQuit      = 28
	sigTstp      = 26
)

func Run(color bool) {
	os.Exit(doRun(color))
}

func doRun(color bool) int {
	if !color {
		slog.DisableColor()
		DisableErrorColor()
	}

	if os.Getenv("RAILS_ENV") != "" {
		println("Warning: Specifying a Rails environment via RAILS_ENV has no effect for commands run with zeus.")
	}

	isTerminal := ttyutils.IsTerminal(os.Stdout.Fd())

	var master, slave *os.File
	var err error
	if isTerminal {
		master, slave, err = pty.Open()
	} else {
		master, slave, err = unixsocket.Socketpair(syscall.SOCK_STREAM)
	}
	if err != nil {
		panic(err)
	}

	defer master.Close()
	if isTerminal {
		oldState, err := ttyutils.MakeTerminalRaw(os.Stdout.Fd())
		if err != nil {
			panic(err)
		}
		defer ttyutils.RestoreTerminalState(os.Stdout.Fd(), oldState)

	}

	// should this happen if we're running over a pipe? I think maybe not?
	ttyutils.MirrorWinsize(os.Stdout, master)

	addr, err := net.ResolveUnixAddr("unixgram", zeusSockName)
	if err != nil {
		panic("Can't resolve server address")
	}

	conn, err := net.DialUnix("unix", nil, addr)
	if err != nil {
		ErrorCantConnectToMaster()
	}
	usock := unixsocket.NewUsock(conn)

	msg := CreateCommandAndArgumentsMessage(os.Args[1], os.Args[2:])
	usock.WriteMessage(msg)
	usock.WriteFD(int(slave.Fd()))
	slave.Close()

	msg, err = usock.ReadMessage()
	if err != nil {
		panic(err)
	}

	parts := strings.Split(msg, "\000")
	commandPid, err := strconv.Atoi(parts[0])
	defer func() {
		if commandPid > 0 {
			// Just in case.
			syscall.Kill(commandPid, 9)
		}
	}()

	if err != nil {
		panic(err)
	}

	if isTerminal {
		c := make(chan os.Signal, 1)
		signal.Notify(c, syscall.SIGWINCH, syscall.SIGCONT)
		go func() {
			for sig := range c {
				if sig == syscall.SIGCONT {
					syscall.Kill(commandPid, syscall.SIGCONT)
				} else if sig == syscall.SIGWINCH {
					ttyutils.MirrorWinsize(os.Stdout, master)
					syscall.Kill(commandPid, syscall.SIGWINCH)
				}
			}
		}()
	}

	var exitStatus int = -1
	if len(parts) > 2 {
		exitStatus, err = strconv.Atoi(parts[0])
		if err != nil {
			panic(err)
		}
	}

	eof := make(chan bool)
	go func() {
		for {
			buf := make([]byte, 1024)
			n, err := master.Read(buf)
			if err != nil {
				eof <- true
				break
			}
			os.Stdout.Write(buf[:n])
		}
	}()

	go func() {
		buf := make([]byte, 8192)
		for {
			n, err := os.Stdin.Read(buf)
			if err != nil {
				eof <- true
				break
			}
			if isTerminal {
				for i := 0; i < n; i++ {
					switch buf[i] {
					case sigInt:
						syscall.Kill(commandPid, syscall.SIGINT)
					case sigQuit:
						syscall.Kill(commandPid, syscall.SIGQUIT)
					case sigTstp:
						syscall.Kill(commandPid, syscall.SIGTSTP)
						syscall.Kill(os.Getpid(), syscall.SIGTSTP)
					}
				}
			}
			master.Write(buf[:n])
		}
	}()

	<-eof

	if exitStatus == -1 {
		msg, err = usock.ReadMessage()
		if err != nil {
			panic(err)
		}
		parts := strings.Split(msg, "\000")
		exitStatus, err = strconv.Atoi(parts[0])
		if err != nil {
			panic(err)
		}
	}

	return exitStatus
}
<|endoftext|>