<commit_before>package server

import (
	"testing"
	"time"
)

func TestBasicConsumer(t *testing.T) {
	vh := NewVHost("/")

	ch := NewChannel(vh)

	err := ch.ExchangeDeclare("neoway", "topic", nil)

	if err != nil {
		t.Error(err)
		return
	}

	q, err := ch.QueueDeclare("data-queue", nil)

	if err != nil {
		t.Error(err)
		return
	}

	err = ch.QueueBind("data-queue", "process.data", "neoway", nil)

	if err != nil {
		t.Error(err)
		return
	}

	deliveries, err := ch.Consume(
		q.Name(),
		"tag-teste",
		nil,
	)

	if err != nil {
		t.Error(err)
		return
	}

	err = ch.Publish("neoway", "process.data", []byte("teste"), nil)

	if err != nil {
		t.Error(err)
		return
	}

	data := <-deliveries

	if string(data.Body()) != "teste" {
		t.Errorf("Failed to publish message to specified route")
		return
	}
}

func TestUnackedMessagesArentLost(t *testing.T) {
	vh := NewVHost("/")

	ch := NewChannel(vh)

	err := ch.ExchangeDeclare("neoway", "topic", nil)

	if err != nil {
		t.Error(err)
		return
	}

	q, err := ch.QueueDeclare("data-queue", nil)

	if err != nil {
		t.Error(err)
		return
	}

	err = ch.QueueBind("data-queue", "process.data", "neoway", nil)

	if err != nil {
		t.Error(err)
		return
	}

	deliveries, err := ch.Consume(
		q.Name(),
		"tag-teste",
		nil,
	)

	if err != nil {
		t.Error(err)
		return
	}

	err = ch.Publish("neoway", "process.data", []byte("teste"), nil)

	if err != nil {
		t.Error(err)
		return
	}

	done := make(chan bool)

	go func() {
		data := <-deliveries

		if string(data.Body()) != "teste" {
			t.Errorf("Failed to publish message to specified route")
			return
		}

		// Closing the channel without ack'ing
		// The message MUST be enqueued by vhost
		ch.Close()

		done <- true
	}()

	select {
	case <-done:
	}

	// create another channel and get the same data back
	ch2 := NewChannel(vh)

	err = ch2.ExchangeDeclare("neoway", "topic", nil)

	if err != nil {
		t.Error(err)
		return
	}

	q2, err := ch2.QueueDeclare("data-queue", nil)

	if err != nil {
		t.Error(err)
		return
	}

	err = ch2.QueueBind("data-queue", "process.data", "neoway", nil)

	if err != nil {
		t.Error(err)
		return
	}

	deliveries2, err := ch2.Consume(
		q2.Name(),
		"tag-teste",
		nil,
	)

	if err != nil {
		t.Error(err)
		return
	}

	done = make(chan bool)

	go func() {
		data := <-deliveries2

		if string(data.Body()) != "teste" {
			t.Errorf("Failed to publish message to specified route")
			return
		}

		done <- true
	}()

	select {
	case <-done:
	}
}

func TestAckedMessagesAreCommited(t *testing.T) {
	vh := NewVHost("/")

	ch := NewChannel(vh)

	err := ch.ExchangeDeclare("neoway", "topic", nil)

	if err != nil {
		t.Error(err)
		return
	}

	q, err := ch.QueueDeclare("data-queue", nil)

	if err != nil {
		t.Error(err)
		return
	}

	err = ch.QueueBind("data-queue", "process.data", "neoway", nil)

	if err != nil {
		t.Error(err)
		return
	}

	deliveries, err := ch.Consume(
		q.Name(),
		"tag-teste",
		nil,
	)

	if err != nil {
		t.Error(err)
		return
	}

	err = ch.Publish("neoway", "process.data", []byte("teste"), nil)

	if err != nil {
		t.Error(err)
		return
	}

	done := make(chan bool)

	go func() {
		data := <-deliveries

		if string(data.Body()) != "teste" {
			t.Errorf("Failed to publish message to specified route")
			return
		}

		data.Ack(false)

		done <- true
	}()

	select {
	case <-done:
	}

	// Closing the old channel, this should reenqueue everything not ack'ed
	ch.Close()

	// create another channel and get the same data back
	ch2 := NewChannel(vh)

	err = ch2.ExchangeDeclare("neoway", "topic", nil)

	if err != nil {
		t.Error(err)
		return
	}

	q2, err := ch2.QueueDeclare("data-queue", nil)

	if err != nil {
		t.Error(err)
		return
	}

	err = ch2.QueueBind("data-queue", "process.data", "neoway", nil)

	if err != nil {
		t.Error(err)
		return
	}

	deliveries2, err := ch2.Consume(
		q2.Name(),
		"tag-teste",
		nil,
	)

	if err != nil {
		t.Error(err)
		return
	}

	timer := time.After(2 * time.Second)

	go func() {
		data := <-deliveries2

		t.Errorf("Data ack'ed delivered again: %s", string(data.Body()))

		panic("never reach here. Message was ack'ed")
	}()

	select {
	case <-timer:
	}
}

func TestNAckedMessagesAreRequeued(t *testing.T) {
	vh := NewVHost("/")

	ch := NewChannel(vh)

	err := ch.ExchangeDeclare("neoway", "topic", nil)

	if err != nil {
		t.Error(err)
		return
	}

	q, err := ch.QueueDeclare("data-queue", nil)

	if err != nil {
		t.Error(err)
		return
	}

	err = ch.QueueBind("data-queue", "process.data", "neoway", nil)

	if err != nil {
		t.Error(err)
		return
	}

	deliveries, err := ch.Consume(
		q.Name(),
		"tag-teste",
		nil,
	)

	if err != nil {
		t.Error(err)
		return
	}

	err = ch.Publish("neoway", "process.data", []byte("teste"), nil)

	if err != nil {
		t.Error(err)
		return
	}

	done := make(chan bool)

	go func() {
		data := <-deliveries

		if string(data.Body()) != "teste" {
			t.Errorf("Failed to publish message to specified route")
			return
		}

		data.Nack(false, true)
		done <- true
	}()

	select {
	case <-done:
	}

	// Closing the old channel, this should reenqueue everything not ack'ed
	ch.Close()

	// create another channel and get the same data back
	ch2 := NewChannel(vh)

	err = ch2.ExchangeDeclare("neoway", "topic", nil)

	if err != nil {
		t.Error(err)
		return
	}

	q2, err := ch2.QueueDeclare("data-queue", nil)

	if err != nil {
		t.Error(err)
		return
	}

	err = ch2.QueueBind("data-queue", "process.data", "neoway", nil)

	if err != nil {
		t.Error(err)
		return
	}

	deliveries2, err := ch2.Consume(
		q2.Name(),
		"tag-teste",
		nil,
	)

	if err != nil {
		t.Error(err)
		return
	}

	timer := time.After(2 * time.Second)

	go func() {
		data := <-deliveries2

		if string(data.Body()) != "teste" {
			t.Error("wrong nacked message")
		}

		done <- true
	}()

	select {
	case <-done:
	case <-timer:
		t.Errorf("Data nack'ed but not redelivered")
	}
}
<commit_msg>Add tests for Delivery.Reject<commit_after>package server

import (
	"testing"
	"time"
)

func TestBasicConsumer(t *testing.T) {
	vh := NewVHost("/")

	ch := NewChannel(vh)

	err := ch.ExchangeDeclare("neoway", "topic", nil)

	if err != nil {
		t.Error(err)
		return
	}

	q, err := ch.QueueDeclare("data-queue", nil)

	if err != nil {
		t.Error(err)
		return
	}

	err = ch.QueueBind("data-queue", "process.data", "neoway", nil)

	if err != nil {
		t.Error(err)
		return
	}

	deliveries, err := ch.Consume(
		q.Name(),
		"tag-teste",
		nil,
	)

	if err != nil {
		t.Error(err)
		return
	}

	err = ch.Publish("neoway", "process.data", []byte("teste"), nil)

	if err != nil {
		t.Error(err)
		return
	}

	data := <-deliveries

	if string(data.Body()) != "teste" {
		t.Errorf("Failed to publish message to specified route")
		return
	}
}

func TestUnackedMessagesArentLost(t *testing.T) {
	vh := NewVHost("/")

	ch := NewChannel(vh)

	err := ch.ExchangeDeclare("neoway", "topic", nil)

	if err != nil {
		t.Error(err)
		return
	}

	q, err := ch.QueueDeclare("data-queue", nil)

	if err != nil {
		t.Error(err)
		return
	}

	err = ch.QueueBind("data-queue", "process.data", "neoway", nil)

	if err != nil {
		t.Error(err)
		return
	}

	deliveries, err := ch.Consume(
		q.Name(),
		"tag-teste",
		nil,
	)

	if err != nil {
		t.Error(err)
		return
	}

	err = ch.Publish("neoway", "process.data", []byte("teste"), nil)

	if err != nil {
		t.Error(err)
		return
	}

	done := make(chan bool)

	go func() {
		data := <-deliveries

		if string(data.Body()) != "teste" {
			t.Errorf("Failed to publish message to specified route")
			return
		}

		// Closing the channel without ack'ing
		// The message MUST be enqueued by vhost
		ch.Close()

		done <- true
	}()

	select {
	case <-done:
	}

	// create another channel and get the same data back
	ch2 := NewChannel(vh)

	err = ch2.ExchangeDeclare("neoway", "topic", nil)

	if err != nil {
		t.Error(err)
		return
	}

	q2, err := ch2.QueueDeclare("data-queue", nil)

	if err != nil {
		t.Error(err)
		return
	}

	err = ch2.QueueBind("data-queue", "process.data", "neoway", nil)

	if err != nil {
		t.Error(err)
		return
	}

	deliveries2, err := ch2.Consume(
		q2.Name(),
		"tag-teste",
		nil,
	)

	if err != nil {
		t.Error(err)
		return
	}

	done = make(chan bool)

	go func() {
		data := <-deliveries2

		if string(data.Body()) != "teste" {
			t.Errorf("Failed to publish message to specified route")
			return
		}

		done <- true
	}()

	select {
	case <-done:
	}
}

func TestAckedMessagesAreCommited(t *testing.T) {
	vh := NewVHost("/")

	ch := NewChannel(vh)

	err := ch.ExchangeDeclare("neoway", "topic", nil)

	if err != nil {
		t.Error(err)
		return
	}

	q, err := ch.QueueDeclare("data-queue", nil)

	if err != nil {
		t.Error(err)
		return
	}

	err = ch.QueueBind("data-queue", "process.data", "neoway", nil)

	if err != nil {
		t.Error(err)
		return
	}

	deliveries, err := ch.Consume(
		q.Name(),
		"tag-teste",
		nil,
	)

	if err != nil {
		t.Error(err)
		return
	}

	err = ch.Publish("neoway", "process.data", []byte("teste"), nil)

	if err != nil {
		t.Error(err)
		return
	}

	done := make(chan bool)

	go func() {
		data := <-deliveries

		if string(data.Body()) != "teste" {
			t.Errorf("Failed to publish message to specified route")
			return
		}

		data.Ack(false)

		done <- true
	}()

	select {
	case <-done:
	}

	// Closing the old channel, this should reenqueue everything not ack'ed
	ch.Close()

	// create another channel and get the same data back
	ch2 := NewChannel(vh)

	err = ch2.ExchangeDeclare("neoway", "topic", nil)

	if err != nil {
		t.Error(err)
		return
	}

	q2, err := ch2.QueueDeclare("data-queue", nil)

	if err != nil {
		t.Error(err)
		return
	}

	err = ch2.QueueBind("data-queue", "process.data", "neoway", nil)

	if err != nil {
		t.Error(err)
		return
	}

	deliveries2, err := ch2.Consume(
		q2.Name(),
		"tag-teste",
		nil,
	)

	if err != nil {
		t.Error(err)
		return
	}

	timer := time.After(2 * time.Second)

	go func() {
		data := <-deliveries2

		t.Errorf("Data ack'ed delivered again: %s", string(data.Body()))

		panic("never reach here. Message was ack'ed")
	}()

	select {
	case <-timer:
	}
}

func TestNAckedMessagesAreRequeued(t *testing.T) {
	vh := NewVHost("/")

	ch := NewChannel(vh)

	err := ch.ExchangeDeclare("neoway", "topic", nil)

	if err != nil {
		t.Error(err)
		return
	}

	q, err := ch.QueueDeclare("data-queue", nil)

	if err != nil {
		t.Error(err)
		return
	}

	err = ch.QueueBind("data-queue", "process.data", "neoway", nil)

	if err != nil {
		t.Error(err)
		return
	}

	deliveries, err := ch.Consume(
		q.Name(),
		"tag-teste",
		nil,
	)

	if err != nil {
		t.Error(err)
		return
	}

	err = ch.Publish("neoway", "process.data", []byte("teste"), nil)

	if err != nil {
		t.Error(err)
		return
	}

	done := make(chan bool)

	go func() {
		data := <-deliveries

		if string(data.Body()) != "teste" {
			t.Errorf("Failed to publish message to specified route")
			return
		}

		data.Nack(false, true)
		done <- true
	}()

	select {
	case <-done:
	}

	// Closing the old channel, this should reenqueue everything not ack'ed
	ch.Close()

	// create another channel and get the same data back
	ch2 := NewChannel(vh)

	err = ch2.ExchangeDeclare("neoway", "topic", nil)

	if err != nil {
		t.Error(err)
		return
	}

	q2, err := ch2.QueueDeclare("data-queue", nil)

	if err != nil {
		t.Error(err)
		return
	}

	err = ch2.QueueBind("data-queue", "process.data", "neoway", nil)

	if err != nil {
		t.Error(err)
		return
	}

	deliveries2, err := ch2.Consume(
		q2.Name(),
		"tag-teste",
		nil,
	)

	if err != nil {
		t.Error(err)
		return
	}

	timer := time.After(2 * time.Second)

	go func() {
		data := <-deliveries2

		if string(data.Body()) != "teste" {
			t.Error("wrong nacked message")
		}

		done <- true
	}()

	select {
	case <-done:
	case <-timer:
		t.Errorf("Data nack'ed but not redelivered")
	}
}

func TestNAckedMessagesAreRejectedWhenRequested(t *testing.T) {
	vh := NewVHost("/")

	ch := NewChannel(vh)

	err := ch.ExchangeDeclare("neoway", "topic", nil)

	if err != nil {
		t.Error(err)
		return
	}

	q, err := ch.QueueDeclare("data-queue", nil)

	if err != nil {
		t.Error(err)
		return
	}

	err = ch.QueueBind("data-queue", "process.data", "neoway", nil)

	if err != nil {
		t.Error(err)
		return
	}

	deliveries, err := ch.Consume(
		q.Name(),
		"tag-teste",
		nil,
	)

	if err != nil {
		t.Error(err)
		return
	}

	err = ch.Publish("neoway", "process.data", []byte("teste"), nil)

	if err != nil {
		t.Error(err)
		return
	}

	done := make(chan bool)

	go func() {
		data := <-deliveries

		if string(data.Body()) != "teste" {
			t.Errorf("Failed to publish message to specified route")
			return
		}

		data.Nack(false, false)
		done <- true
	}()

	select {
	case <-done:
	}

	// Closing the old channel, this should reenqueue everything not ack'ed
	ch.Close()

	// create another channel and get the same data back
	ch2 := NewChannel(vh)

	err = ch2.ExchangeDeclare("neoway", "topic", nil)

	if err != nil {
		t.Error(err)
		return
	}

	q2, err := ch2.QueueDeclare("data-queue", nil)

	if err != nil {
		t.Error(err)
		return
	}

	err = ch2.QueueBind("data-queue", "process.data", "neoway", nil)

	if err != nil {
		t.Error(err)
		return
	}

	deliveries2, err := ch2.Consume(
		q2.Name(),
		"tag-teste",
		nil,
	)

	if err != nil {
		t.Error(err)
		return
	}

	timer := time.After(2 * time.Second)

	go func() {
		<-deliveries2

		t.Error("Message shall not be redelivered")
		done <- true
	}()

	select {
	case <-done:
	case <-timer:
	}
}

// Reject

func TestRejectedMessagesAreRequeuedWhenRequested(t *testing.T) {
	vh := NewVHost("/")

	ch := NewChannel(vh)

	err := ch.ExchangeDeclare("neoway", "topic", nil)

	if err != nil {
		t.Error(err)
		return
	}

	q, err := ch.QueueDeclare("data-queue", nil)

	if err != nil {
		t.Error(err)
		return
	}

	err = ch.QueueBind("data-queue", "process.data", "neoway", nil)

	if err != nil {
		t.Error(err)
		return
	}

	deliveries, err := ch.Consume(
		q.Name(),
		"tag-teste",
		nil,
	)

	if err != nil {
		t.Error(err)
		return
	}

	err = ch.Publish("neoway", "process.data", []byte("teste"), nil)

	if err != nil {
		t.Error(err)
		return
	}

	done := make(chan bool)

	go func() {
		data := <-deliveries

		if string(data.Body()) != "teste" {
			t.Errorf("Failed to publish message to specified route")
			return
		}

		data.Reject(true)
		done <- true
	}()

	select {
	case <-done:
	}

	// Closing the old channel, this should reenqueue everything not ack'ed
	ch.Close()

	// create another channel and get the same data back
	ch2 := NewChannel(vh)

	err = ch2.ExchangeDeclare("neoway", "topic", nil)

	if err != nil {
		t.Error(err)
		return
	}

	q2, err := ch2.QueueDeclare("data-queue", nil)

	if err != nil {
		t.Error(err)
		return
	}

	err = ch2.QueueBind("data-queue", "process.data", "neoway", nil)

	if err != nil {
		t.Error(err)
		return
	}

	deliveries2, err := ch2.Consume(
		q2.Name(),
		"tag-teste",
		nil,
	)

	if err != nil {
		t.Error(err)
		return
	}

	timer := time.After(2 * time.Second)

	go func() {
		data := <-deliveries2

		if string(data.Body()) != "teste" {
			t.Error("wrong nacked message")
		}

		done <- true
	}()

	select {
	case <-done:
	case <-timer:
		t.Errorf("Data nack'ed but not redelivered")
	}
}

func TestRejectedMessagesAreRejectedWhenRequested(t *testing.T) {
	vh := NewVHost("/")

	ch := NewChannel(vh)

	err := ch.ExchangeDeclare("neoway", "topic", nil)

	if err != nil {
		t.Error(err)
		return
	}

	q, err := ch.QueueDeclare("data-queue", nil)

	if err != nil {
		t.Error(err)
		return
	}

	err = ch.QueueBind("data-queue", "process.data", "neoway", nil)

	if err != nil {
		t.Error(err)
		return
	}

	deliveries, err := ch.Consume(
		q.Name(),
		"tag-teste",
		nil,
	)

	if err != nil {
		t.Error(err)
		return
	}

	err = ch.Publish("neoway", "process.data", []byte("teste"), nil)

	if err != nil {
		t.Error(err)
		return
	}

	done := make(chan bool)

	go func() {
		data := <-deliveries

		if string(data.Body()) != "teste" {
			t.Errorf("Failed to publish message to specified route")
			return
		}

		data.Reject(false)
		done <- true
	}()

	select {
	case <-done:
	}

	// Closing the old channel, this should reenqueue everything not ack'ed
	ch.Close()

	// create another channel and get the same data back
	ch2 := NewChannel(vh)

	err = ch2.ExchangeDeclare("neoway", "topic", nil)

	if err != nil {
		t.Error(err)
		return
	}

	q2, err := ch2.QueueDeclare("data-queue", nil)

	if err != nil {
		t.Error(err)
		return
	}

	err = ch2.QueueBind("data-queue", "process.data", "neoway", nil)

	if err != nil {
		t.Error(err)
		return
	}

	deliveries2, err := ch2.Consume(
		q2.Name(),
		"tag-teste",
		nil,
	)

	if err != nil {
		t.Error(err)
		return
	}

	timer := time.After(2 * time.Second)

	go func() {
		<-deliveries2

		t.Error("Message shall not be redelivered")
		done <- true
	}()

	select {
	case <-done:
	case <-timer:
	}
}
<|endoftext|>


<commit_before>package github

import (
	"net/url"

	"github.com/google/go-github/github"
	"golang.org/x/oauth2"
)

type Config struct {
	Token        string
	Organization string
	BaseURL      string
}

type Organization struct {
	name   string
	client *github.Client
}

// Client configures and returns a fully initialized GithubClient
func (c *Config) Client() (interface{}, error) {
	var org Organization
	org.name = c.Organization
	ts := oauth2.StaticTokenSource(
		&oauth2.Token{AccessToken: c.Token},
	)
	tc := oauth2.NewClient(oauth2.NoContext, ts)

	org.client = github.NewClient(tc)
	if c.BaseURL != "" {
		u, err := url.Parse(c.BaseURL)
		if err != nil {
			return nil, err
		}
		org.client.BaseURL = u
	}
	return &org, nil
}
<commit_msg>provider/github: Log HTTP requests and responses in DEBUG mode (#14363)<commit_after>package github

import (
	"net/url"

	"github.com/google/go-github/github"
	"github.com/hashicorp/terraform/helper/logging"
	"golang.org/x/oauth2"
)

type Config struct {
	Token        string
	Organization string
	BaseURL      string
}

type Organization struct {
	name   string
	client *github.Client
}

// Client configures and returns a fully initialized GithubClient
func (c *Config) Client() (interface{}, error) {
	var org Organization
	org.name = c.Organization
	ts := oauth2.StaticTokenSource(
		&oauth2.Token{AccessToken: c.Token},
	)
	tc := oauth2.NewClient(oauth2.NoContext, ts)

	tc.Transport = logging.NewTransport("Github", tc.Transport)

	org.client = github.NewClient(tc)
	if c.BaseURL != "" {
		u, err := url.Parse(c.BaseURL)
		if err != nil {
			return nil, err
		}
		org.client.BaseURL = u
	}

	return &org, nil
}
<|endoftext|>


<commit_before>package ftmplting

import (
	"strings"
	texttmpl "text/template"
)

var generatorTemplate = texttmpl.Must(texttmpl.New("").Parse(`{{ .GlobalCode }}
// {{ .ErrFuncPrefix }}{{ .FuncName }} evaluates a template {{ .TemplateFile }}
func {{ .ErrFuncPrefix }}{{ .FuncName }}({{ .ArgsJoined }}) (string, error) {
	_template := "{{ .TemplateFile }}"
	_escape := {{ .EscapeFunc }}
	var _ftmpl bytes.Buffer
	_w := func(str string) { _, _ = _ftmpl.WriteString(str) }
	_, _, _ = _template, _escape, _w

{{ range .Lines }}{{ . }}
{{ end }}
	return _ftmpl.String(), nil
}

// {{ .NoerrFuncPrefix }}{{ .FuncName }} evaluates a template {{ .TemplateFile }}
func {{ .NoerrFuncPrefix }}{{ .FuncName }}({{ .ArgsJoined }}) string {
	html, err := {{ .ErrFuncPrefix }}{{ .FuncName }}({{ .ArgNamesJoined }})
	if err != nil {
		_ ,_ = os.Stderr.WriteString("Error running template {{ .TemplateFile }}:" + err.Error())
	}
	return html
}`))

type templateParams struct {
	GlobalCode      string
	ErrFuncPrefix   string
	NoerrFuncPrefix string
	TemplateFile    string
	EscapeFunc      string
	FuncName        string
	Args            []string
	Lines           []string
}

func (tp templateParams) ArgsJoined() string {
	return strings.Join(tp.Args, ", ")
}

func (tp templateParams) ArgNamesJoined() string {
	var argNames []string
	for _, arg := range tp.Args {
		argNames = append(argNames, strings.Split(arg, " ")[0])
	}
	return strings.Join(argNames, ", ")
}

var initFunctionTemplate = `func init() {
	_ = fmt.Sprintf
	_ = errors.New
	_ = os.Stderr
	_ = html.EscapeString
}`

var defaultQuoteTemplate = `func(str string) string {
	return str
}`

var errorTemplate = texttmpl.Must(texttmpl.New("").Parse(`if {{.Expression}} {
	return "", errors.New({{.Message}})
}`))

type errParams struct {
	Expression, Message string
}

var stringTemplate = texttmpl.Must(texttmpl.New("").Parse(`_w(` + "`" + `{{ . }}` + "`" + `)`))

var newlineTemplate = `_w("\\n")`

var patternTemplate = texttmpl.Must(texttmpl.New("").Parse(`_w(fmt.Sprintf(` + "`" + `{{ .Template }}` + "`" + `, {{ .ArgsJoined }}))`))

type patternTemplateParam struct {
	Template string
	Args     []string
}

func (ptp patternTemplateParam) ArgsJoined() string {
	return strings.Join(ptp.Args, ", ")
}
<commit_msg>Not used<commit_after>package ftmplting

import (
	"strings"
	texttmpl "text/template"
)

var generatorTemplate = texttmpl.Must(texttmpl.New("").Parse(`{{ .GlobalCode }}
// {{ .ErrFuncPrefix }}{{ .FuncName }} evaluates a template {{ .TemplateFile }}
func {{ .ErrFuncPrefix }}{{ .FuncName }}({{ .ArgsJoined }}) (string, error) {
	_template := "{{ .TemplateFile }}"
	_escape := {{ .EscapeFunc }}
	var _ftmpl bytes.Buffer
	_w := func(str string) { _, _ = _ftmpl.WriteString(str) }
	_, _, _ = _template, _escape, _w

{{ range .Lines }}{{ . }}
{{ end }}
	return _ftmpl.String(), nil
}

// {{ .NoerrFuncPrefix }}{{ .FuncName }} evaluates a template {{ .TemplateFile }}
func {{ .NoerrFuncPrefix }}{{ .FuncName }}({{ .ArgsJoined }}) string {
	html, err := {{ .ErrFuncPrefix }}{{ .FuncName }}({{ .ArgNamesJoined }})
	if err != nil {
		_ ,_ = os.Stderr.WriteString("Error running template {{ .TemplateFile }}:" + err.Error())
	}
	return html
}`))

type templateParams struct {
	GlobalCode      string
	ErrFuncPrefix   string
	NoerrFuncPrefix string
	TemplateFile    string
	EscapeFunc      string
	FuncName        string
	Args            []string
	Lines           []string
}

func (tp templateParams) ArgsJoined() string {
	return strings.Join(tp.Args, ", ")
}

func (tp templateParams) ArgNamesJoined() string {
	var argNames []string
	for _, arg := range tp.Args {
		argNames = append(argNames, strings.Split(arg, " ")[0])
	}
	return strings.Join(argNames, ", ")
}

var initFunctionTemplate = `func init() {
	_ = fmt.Sprintf
	_ = errors.New
	_ = os.Stderr
	_ = html.EscapeString
}`

var defaultQuoteTemplate = `func(str string) string {
	return str
}`

var errorTemplate = texttmpl.Must(texttmpl.New("").Parse(`if {{.Expression}} {
	return "", errors.New({{.Message}})
}`))

type errParams struct {
	Expression, Message string
}

var stringTemplate = texttmpl.Must(texttmpl.New("").Parse(`_w(` + "`" + `{{ . }}` + "`" + `)`))
<|endoftext|>


<commit_before>package consumer

import (
	"context"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go/service/kinesis"
	"github.com/aws/aws-sdk-go/service/kinesis/kinesisiface"
)

// NewAllGroup returns an intitialized AllGroup for consuming
// all shards on a stream
func NewAllGroup(ksis kinesisiface.KinesisAPI, store Store, streamName string, logger Logger) *AllGroup {
	return &AllGroup{
		ksis:       ksis,
		shards:     make(map[string]*kinesis.Shard),
		streamName: streamName,
		logger:     logger,
		Store:      store,
	}
}

// AllGroup caches a local list of the shards we are already processing
// and routinely polls the stream looking for new shards to process
type AllGroup struct {
	ksis       kinesisiface.KinesisAPI
	streamName string
	logger     Logger
	Store

	shardMu sync.Mutex
	shards  map[string]*kinesis.Shard
}

// Start is a blocking operation which will loop and attempt to find new
// shards on a regular cadence.
func (g *AllGroup) Start(ctx context.Context, shardc chan *kinesis.Shard) {
	var ticker = time.NewTicker(30 * time.Second)
	g.findNewShards(shardc)

	// Note: while ticker is a rather naive approach to this problem,
	// it actually simplies a few things. i.e. If we miss a new shard while
	// AWS is resharding we'll pick it up max 30 seconds later.

	// It might be worth refactoring this flow to allow the consumer to
	// to notify the broker when a shard is closed. However, shards don't
	// necessarily close at the same time, so we could potentially get a
	// thundering heard of notifications from the consumer.

	for {
		select {
		case <-ctx.Done():
			ticker.Stop()
			return
		case <-ticker.C:
			g.findNewShards(shardc)
		}
	}
}

// findNewShards pulls the list of shards from the Kinesis API
// and uses a local cache to determine if we are already processing
// a particular shard.
func (g *AllGroup) findNewShards(shardc chan *kinesis.Shard) {
	g.shardMu.Lock()
	defer g.shardMu.Unlock()

	g.logger.Log("[GROUP]", "fetching shards")

	shards, err := listShards(g.ksis, g.streamName)
	if err != nil {
		g.logger.Log("[GROUP] error:", err)
		return
	}

	for _, shard := range shards {
		if _, ok := g.shards[*shard.ShardId]; ok {
			continue
		}
		g.shards[*shard.ShardId] = shard
		shardc <- shard
	}
}
<commit_msg>Update godoc for allgroup description<commit_after>package consumer

import (
	"context"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go/service/kinesis"
	"github.com/aws/aws-sdk-go/service/kinesis/kinesisiface"
)

// NewAllGroup returns an intitialized AllGroup for consuming
// all shards on a stream
func NewAllGroup(ksis kinesisiface.KinesisAPI, store Store, streamName string, logger Logger) *AllGroup {
	return &AllGroup{
		ksis:       ksis,
		shards:     make(map[string]*kinesis.Shard),
		streamName: streamName,
		logger:     logger,
		Store:      store,
	}
}

// AllGroup is used to consume all shards from a single consumer. It
// caches a local list of the shards we are already processing
// and routinely polls the stream looking for new shards to process.
type AllGroup struct {
	ksis       kinesisiface.KinesisAPI
	streamName string
	logger     Logger
	Store

	shardMu sync.Mutex
	shards  map[string]*kinesis.Shard
}

// Start is a blocking operation which will loop and attempt to find new
// shards on a regular cadence.
func (g *AllGroup) Start(ctx context.Context, shardc chan *kinesis.Shard) {
	var ticker = time.NewTicker(30 * time.Second)
	g.findNewShards(shardc)

	// Note: while ticker is a rather naive approach to this problem,
	// it actually simplies a few things. i.e. If we miss a new shard while
	// AWS is resharding we'll pick it up max 30 seconds later.

	// It might be worth refactoring this flow to allow the consumer to
	// to notify the broker when a shard is closed. However, shards don't
	// necessarily close at the same time, so we could potentially get a
	// thundering heard of notifications from the consumer.

	for {
		select {
		case <-ctx.Done():
			ticker.Stop()
			return
		case <-ticker.C:
			g.findNewShards(shardc)
		}
	}
}

// findNewShards pulls the list of shards from the Kinesis API
// and uses a local cache to determine if we are already processing
// a particular shard.
func (g *AllGroup) findNewShards(shardc chan *kinesis.Shard) {
	g.shardMu.Lock()
	defer g.shardMu.Unlock()

	g.logger.Log("[GROUP]", "fetching shards")

	shards, err := listShards(g.ksis, g.streamName)
	if err != nil {
		g.logger.Log("[GROUP] error:", err)
		return
	}

	for _, shard := range shards {
		if _, ok := g.shards[*shard.ShardId]; ok {
			continue
		}
		g.shards[*shard.ShardId] = shard
		shardc <- shard
	}
}
<|endoftext|>


<commit_before>// Copyright (C) 2017 Google 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 client

import (
	"context"
	"fmt"
	"math"
	"sync"
	"time"

	"github.com/google/gapid/core/app"
	"github.com/google/gapid/core/app/crash"
	"github.com/google/gapid/core/app/status"
	"github.com/google/gapid/core/event/task"
	"github.com/google/gapid/core/log"
	"github.com/google/gapid/core/os/device"
	"github.com/google/gapid/core/os/device/bind"
	"github.com/google/gapid/gapir"
	replaysrv "github.com/google/gapid/gapir/replay_service"
	"google.golang.org/grpc"
)

const (
	// LaunchArgsKey is the bind device property key used to control the command
	// line arguments when launching GAPIR. The property must be of type []string.
	LaunchArgsKey = "gapir-launch-args"
	// gRPCConnectTimeout is the time allowed to establish a gRPC connection.
	gRPCConnectTimeout = time.Second * 10
	// heartbeatInterval is the delay between heartbeat pings.
	heartbeatInterval = time.Second * 2
)

// ReplayExecutor must be implemented by replay executors to handle some live
// interactions with a running replay.
type ReplayExecutor interface {
	// HandlePostData handles the given post data message.
	HandlePostData(context.Context, *gapir.PostData) error
	// HandleNotification handles the given notification message.
	HandleNotification(context.Context, *gapir.Notification) error
	// HandleFinished is notified when the given replay is finished.
	HandleFinished(context.Context, error) error
	// HandleFenceReadyRequest handles when the replayer is waiting for the server
	// to execute the registered FenceReadyRequestCallback for fence ID provided
	// in the FenceReadyRequest.
	HandleFenceReadyRequest(context.Context, *gapir.FenceReadyRequest) error
}

// ReplayerKey is used to uniquely identify a GAPIR instance.
type ReplayerKey struct {
	device bind.Device
	arch   device.Architecture
}

// Client handles multiple GAPIR instances identified by ReplayerKey.
type Client struct {
	// mutex prevents data races when restarting replayers. All exported
	// functions must acquire this mutex upon start, to guard the lookup
	// in the replayers map which might be concurrently updated by a
	// replayer reconnection.
	mutex sync.Mutex
	// replayers stores the informations relater to GAPIR instances.
	replayers map[ReplayerKey]*replayer
}

// New returns a new Client with no replayers.
func New(ctx context.Context) *Client {
	client := &Client{replayers: map[ReplayerKey]*replayer{}}
	app.AddCleanup(ctx, func() {
		client.shutdown(ctx)
	})
	return client
}

// shutdown closes all replayer instances and makes the client invalid.
func (client *Client) shutdown(ctx context.Context) {
	client.mutex.Lock()
	defer client.mutex.Unlock()

	for _, replayer := range client.replayers {
		replayer.closeConnection(ctx)
	}
	client.replayers = nil
}

// Connect starts a GAPIR instance and return its ReplayerKey.
func (client *Client) Connect(ctx context.Context, device bind.Device, abi *device.ABI) (*ReplayerKey, error) {
	client.mutex.Lock()
	defer client.mutex.Unlock()

	ctx = status.Start(ctx, "Connect")
	defer status.Finish(ctx)

	if client.replayers == nil {
		return nil, log.Err(ctx, nil, "Client has been shutdown")
	}

	key := ReplayerKey{device: device, arch: abi.GetArchitecture()}

	if _, ok := client.replayers[key]; ok {
		return &key, nil
	}

	launchArgs, _ := bind.GetRegistry(ctx).DeviceProperty(ctx, device, LaunchArgsKey).([]string)
	newDeviceConnectionInfo, err := initDeviceConnection(ctx, device, abi, launchArgs)
	if err != nil {
		return nil, err
	}

	log.I(ctx, "Waiting for connection to GAPIR...")

	// Create gRPC connection
	conn, err := grpc.Dial(fmt.Sprintf("localhost:%d", newDeviceConnectionInfo.port),
		grpc.WithInsecure(),
		grpc.WithBlock(),
		grpc.WithTimeout(gRPCConnectTimeout),
		grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(math.MaxInt32)))
	if err != nil {
		return nil, log.Err(ctx, err, "Timeout waiting for connection")
	}
	rpcClient := replaysrv.NewGapirClient(conn)

	replayer := &replayer{
		deviceConnectionInfo: *newDeviceConnectionInfo,
		device:               device,
		abi:                  abi,
		conn:                 conn,
		rpcClient:            rpcClient,
	}

	crash.Go(func() { client.heartbeat(ctx, replayer) })
	log.I(ctx, "Heartbeat connection setup done")

	err = replayer.startReplayCommunicationHandler(ctx)
	if err != nil {
		return nil, log.Err(ctx, err, "Error in startReplayCommunicationHandler")
	}

	client.replayers[key] = replayer
	return &key, nil
}

// removeConnection closes and removes a GAPIR instance.
func (client *Client) removeConnection(ctx context.Context, key ReplayerKey) {
	client.mutex.Lock()
	defer client.mutex.Unlock()

	if replayer, ok := client.replayers[key]; ok {
		replayer.closeConnection(ctx)
		delete(client.replayers, key)
	}
}

// reconnect removes the replayer before re-creating it.
func (client *Client) reconnect(ctx context.Context, replayer *replayer) {
	key := ReplayerKey{device: replayer.device, arch: replayer.abi.GetArchitecture()}
	client.removeConnection(ctx, key)
	client.Connect(ctx, replayer.device, replayer.abi)
}

// heartbeat regularly sends a ping to a replayer, and restarts it when it fails to reply.
func (client *Client) heartbeat(ctx context.Context, replayer *replayer) {
	for {
		select {
		case <-task.ShouldStop(ctx):
			return
		case <-time.After(heartbeatInterval):
			err := replayer.ping(ctx)
			if err != nil {
				log.E(ctx, "Error sending keep-alive ping. Error: %v", err)
				client.reconnect(ctx, replayer)
				return
			}
		}
	}
}

// getActiveReplayer returns the replayer identified by key only if this
// replayer has an active connection to a GAPIR instance.
func (client *Client) getActiveReplayer(ctx context.Context, key *ReplayerKey) (*replayer, error) {
	replayer, found := client.replayers[*key]
	if !found {
		return nil, log.Errf(ctx, nil, "Cannot find replayer for this key: %v", key)
	}

	if replayer.rpcClient == nil || replayer.conn == nil || replayer.rpcStream == nil {
		return nil, log.Err(ctx, nil, "Replayer has no active connection")
	}

	return replayer, nil
}

// BeginReplay sends a replay request to the replayer identified by key.
func (client *Client) BeginReplay(ctx context.Context, key *ReplayerKey, payload string, dependent string) error {
	client.mutex.Lock()
	defer client.mutex.Unlock()

	ctx = log.Enter(ctx, "Starting replay on gapir device")
	replayer, err := client.getActiveReplayer(ctx, key)
	if err != nil {
		return err
	}

	idReq := replaysrv.ReplayRequest{
		Req: &replaysrv.ReplayRequest_Replay{
			Replay: &replaysrv.Replay{
				ReplayId:    payload,
				DependentId: dependent,
			},
		},
	}
	err = replayer.rpcStream.Send(&idReq)
	if err != nil {
		return log.Err(ctx, err, "Sending replay id")
	}

	return nil
}

// SetReplayExecutor assigns a replay executor to the replayer identified by
// key. It returns a cleanup function to remove the executor once the replay
// is finished.
func (client *Client) SetReplayExecutor(ctx context.Context, key *ReplayerKey, executor ReplayExecutor) (func(), error) {
	client.mutex.Lock()
	defer client.mutex.Unlock()

	replayer, err := client.getActiveReplayer(ctx, key)
	if err != nil {
		return nil, err
	}

	if replayer.executor != nil {
		return nil, log.Err(ctx, nil, "Cannot set an executor while one is already present")
	}
	replayer.executor = executor
	return func() { replayer.executor = nil }, nil
}

// PrewarmReplay requests the GAPIR device to get itself into the given state
func (client *Client) PrewarmReplay(ctx context.Context, key *ReplayerKey, payload string, cleanup string) error {
	client.mutex.Lock()
	defer client.mutex.Unlock()

	replayer, err := client.getActiveReplayer(ctx, key)
	if err != nil {
		return log.Err(ctx, err, "Getting replayer replayer")
	}

	PrerunReq := replaysrv.ReplayRequest{
		Req: &replaysrv.ReplayRequest_Prewarm{
			Prewarm: &replaysrv.PrewarmRequest{
				PrerunId:  payload,
				CleanupId: cleanup,
			},
		},
	}
	err = replayer.rpcStream.Send(&PrerunReq)
	if err != nil {
		return log.Err(ctx, err, "Sending replay payload")
	}
	return nil
}
<commit_msg>Restart GAPIR only upon a new replay request (#367)<commit_after>// Copyright (C) 2017 Google 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 client

import (
	"context"
	"fmt"
	"math"
	"sync"
	"time"

	"github.com/google/gapid/core/app"
	"github.com/google/gapid/core/app/crash"
	"github.com/google/gapid/core/app/status"
	"github.com/google/gapid/core/event/task"
	"github.com/google/gapid/core/log"
	"github.com/google/gapid/core/os/device"
	"github.com/google/gapid/core/os/device/bind"
	"github.com/google/gapid/gapir"
	replaysrv "github.com/google/gapid/gapir/replay_service"
	"google.golang.org/grpc"
)

const (
	// LaunchArgsKey is the bind device property key used to control the command
	// line arguments when launching GAPIR. The property must be of type []string.
	LaunchArgsKey = "gapir-launch-args"
	// gRPCConnectTimeout is the time allowed to establish a gRPC connection.
	gRPCConnectTimeout = time.Second * 10
	// heartbeatInterval is the delay between heartbeat pings.
	heartbeatInterval = time.Second * 2
)

// ReplayExecutor must be implemented by replay executors to handle some live
// interactions with a running replay.
type ReplayExecutor interface {
	// HandlePostData handles the given post data message.
	HandlePostData(context.Context, *gapir.PostData) error
	// HandleNotification handles the given notification message.
	HandleNotification(context.Context, *gapir.Notification) error
	// HandleFinished is notified when the given replay is finished.
	HandleFinished(context.Context, error) error
	// HandleFenceReadyRequest handles when the replayer is waiting for the server
	// to execute the registered FenceReadyRequestCallback for fence ID provided
	// in the FenceReadyRequest.
	HandleFenceReadyRequest(context.Context, *gapir.FenceReadyRequest) error
}

// ReplayerKey is used to uniquely identify a GAPIR instance.
type ReplayerKey struct {
	device bind.Device
	arch   device.Architecture
}

// Client handles multiple GAPIR instances identified by ReplayerKey.
type Client struct {
	// mutex prevents data races when restarting replayers. All exported
	// functions must acquire this mutex upon start, to guard the lookup
	// in the replayers map which might be concurrently updated by a
	// replayer reconnection.
	mutex sync.Mutex
	// replayers stores the informations relater to GAPIR instances.
	replayers map[ReplayerKey]*replayer
}

// New returns a new Client with no replayers.
func New(ctx context.Context) *Client {
	client := &Client{replayers: map[ReplayerKey]*replayer{}}
	app.AddCleanup(ctx, func() {
		client.shutdown(ctx)
	})
	return client
}

// shutdown closes all replayer instances and makes the client invalid.
func (client *Client) shutdown(ctx context.Context) {
	client.mutex.Lock()
	defer client.mutex.Unlock()

	for _, replayer := range client.replayers {
		replayer.closeConnection(ctx)
	}
	client.replayers = nil
}

// Connect starts a GAPIR instance and return its ReplayerKey.
func (client *Client) Connect(ctx context.Context, device bind.Device, abi *device.ABI) (*ReplayerKey, error) {
	client.mutex.Lock()
	defer client.mutex.Unlock()

	ctx = status.Start(ctx, "Connect")
	defer status.Finish(ctx)

	if client.replayers == nil {
		return nil, log.Err(ctx, nil, "Client has been shutdown")
	}

	key := ReplayerKey{device: device, arch: abi.GetArchitecture()}

	if _, ok := client.replayers[key]; ok {
		return &key, nil
	}

	launchArgs, _ := bind.GetRegistry(ctx).DeviceProperty(ctx, device, LaunchArgsKey).([]string)
	newDeviceConnectionInfo, err := initDeviceConnection(ctx, device, abi, launchArgs)
	if err != nil {
		return nil, err
	}

	log.I(ctx, "Waiting for connection to GAPIR...")

	// Create gRPC connection
	conn, err := grpc.Dial(fmt.Sprintf("localhost:%d", newDeviceConnectionInfo.port),
		grpc.WithInsecure(),
		grpc.WithBlock(),
		grpc.WithTimeout(gRPCConnectTimeout),
		grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(math.MaxInt32)))
	if err != nil {
		return nil, log.Err(ctx, err, "Timeout waiting for connection")
	}
	rpcClient := replaysrv.NewGapirClient(conn)

	replayer := &replayer{
		deviceConnectionInfo: *newDeviceConnectionInfo,
		device:               device,
		abi:                  abi,
		conn:                 conn,
		rpcClient:            rpcClient,
	}

	crash.Go(func() { client.heartbeat(ctx, replayer) })
	log.I(ctx, "Heartbeat connection setup done")

	err = replayer.startReplayCommunicationHandler(ctx)
	if err != nil {
		return nil, log.Err(ctx, err, "Error in startReplayCommunicationHandler")
	}

	client.replayers[key] = replayer
	return &key, nil
}

// removeReplayer closes and removes a GAPIR instance.
func (client *Client) removeReplayer(ctx context.Context, replayer *replayer) {
	client.mutex.Lock()
	defer client.mutex.Unlock()

	replayer.closeConnection(ctx)
	key := ReplayerKey{device: replayer.device, arch: replayer.abi.GetArchitecture()}
	delete(client.replayers, key)
}

// heartbeat regularly sends a ping to a replayer, and restarts it when it fails to reply.
func (client *Client) heartbeat(ctx context.Context, replayer *replayer) {
	for {
		select {
		case <-task.ShouldStop(ctx):
			return
		case <-time.After(heartbeatInterval):
			err := replayer.ping(ctx)
			if err != nil {
				log.E(ctx, "Error sending keep-alive ping. Error: %v", err)
				client.removeReplayer(ctx, replayer)
				return
			}
		}
	}
}

// getActiveReplayer returns the replayer identified by key only if this
// replayer has an active connection to a GAPIR instance.
func (client *Client) getActiveReplayer(ctx context.Context, key *ReplayerKey) (*replayer, error) {
	replayer, found := client.replayers[*key]
	if !found {
		return nil, log.Errf(ctx, nil, "Cannot find replayer for this key: %v", key)
	}

	if replayer.rpcClient == nil || replayer.conn == nil || replayer.rpcStream == nil {
		return nil, log.Err(ctx, nil, "Replayer has no active connection")
	}

	return replayer, nil
}

// BeginReplay sends a replay request to the replayer identified by key.
func (client *Client) BeginReplay(ctx context.Context, key *ReplayerKey, payload string, dependent string) error {
	client.mutex.Lock()
	defer client.mutex.Unlock()

	ctx = log.Enter(ctx, "Starting replay on gapir device")
	replayer, err := client.getActiveReplayer(ctx, key)
	if err != nil {
		return err
	}

	idReq := replaysrv.ReplayRequest{
		Req: &replaysrv.ReplayRequest_Replay{
			Replay: &replaysrv.Replay{
				ReplayId:    payload,
				DependentId: dependent,
			},
		},
	}
	err = replayer.rpcStream.Send(&idReq)
	if err != nil {
		return log.Err(ctx, err, "Sending replay id")
	}

	return nil
}

// SetReplayExecutor assigns a replay executor to the replayer identified by
// key. It returns a cleanup function to remove the executor once the replay
// is finished.
func (client *Client) SetReplayExecutor(ctx context.Context, key *ReplayerKey, executor ReplayExecutor) (func(), error) {
	client.mutex.Lock()
	defer client.mutex.Unlock()

	replayer, err := client.getActiveReplayer(ctx, key)
	if err != nil {
		return nil, err
	}

	if replayer.executor != nil {
		return nil, log.Err(ctx, nil, "Cannot set an executor while one is already present")
	}
	replayer.executor = executor
	return func() { replayer.executor = nil }, nil
}

// PrewarmReplay requests the GAPIR device to get itself into the given state
func (client *Client) PrewarmReplay(ctx context.Context, key *ReplayerKey, payload string, cleanup string) error {
	client.mutex.Lock()
	defer client.mutex.Unlock()

	replayer, err := client.getActiveReplayer(ctx, key)
	if err != nil {
		return log.Err(ctx, err, "Getting replayer replayer")
	}

	PrerunReq := replaysrv.ReplayRequest{
		Req: &replaysrv.ReplayRequest_Prewarm{
			Prewarm: &replaysrv.PrewarmRequest{
				PrerunId:  payload,
				CleanupId: cleanup,
			},
		},
	}
	err = replayer.rpcStream.Send(&PrerunReq)
	if err != nil {
		return log.Err(ctx, err, "Sending replay payload")
	}
	return nil
}
<|endoftext|>


<commit_before>package pod

import (
	"fmt"
	"io"
	"time"

	"github.com/docker/docker/pkg/stdcopy"

	"github.com/hyperhq/hypercontainer-utils/hlog"
	"github.com/hyperhq/hyperd/utils"
	"github.com/hyperhq/runv/hypervisor"
)

type Exec struct {
	Id        string
	Container string
	Cmds      string
	Terminal  bool
	ExitCode  uint8

	logPrefix string
	finChan   chan bool
}

func (e *Exec) LogPrefix() string {
	return e.logPrefix
}

func (e *Exec) Log(level hlog.LogLevel, args ...interface{}) {
	hlog.HLog(level, e, 1, args...)
}

func (p *XPod) CreateExec(containerId, cmds string, terminal bool) (string, error) {
	c, ok := p.containers[containerId]
	if !ok {
		err := fmt.Errorf("no container available for exec %s", cmds)
		p.Log(ERROR, err)
		return "", err
	}

	if !c.IsAlive() {
		err := fmt.Errorf("container is not available (%v) for exec %s", c.CurrentState(), cmds)
		p.Log(ERROR, err)
		return "", err
	}

	execId := fmt.Sprintf("exec-%s", utils.RandStr(10, "alpha"))

	p.statusLock.Lock()
	p.execs[execId] = &Exec{
		Container: containerId,
		Id:        execId,
		Cmds:      cmds,
		Terminal:  terminal,
		ExitCode:  255,
		logPrefix: fmt.Sprintf("Pod[%s] Con[%s] Exec[%s] ", p.Id(), containerId[:12], execId),
		finChan:   make(chan bool, 1),
	}
	p.statusLock.Unlock()

	return execId, nil
}

type waitClose struct {
	io.ReadCloser
	wait chan bool
}

func (wc *waitClose) Close() error {
	close(wc.wait)
	return wc.ReadCloser.Close()
}

func (p *XPod) StartExec(stdin io.ReadCloser, stdout io.WriteCloser, containerId, execId string) error {
	p.statusLock.RLock()
	es, ok := p.execs[execId]
	p.statusLock.RUnlock()

	if !ok {
		err := fmt.Errorf("no exec %s exists for container %s", execId, containerId)
		p.Log(ERROR, err)
		return err
	}

	wReader := &waitClose{ReadCloser: stdin, wait: make(chan bool)}
	tty := &hypervisor.TtyIO{
		Stdin:  wReader,
		Stdout: stdout,
	}

	if !es.Terminal && stdout != nil {
		tty.Stderr = stdcopy.NewStdWriter(stdout, stdcopy.Stderr)
		tty.Stdout = stdcopy.NewStdWriter(stdout, stdcopy.Stdout)
	}

	var fin = true
	for fin {
		select {
		case fin = <-es.finChan:
			es.Log(DEBUG, "try to drain the sync chan")
		default:
			fin = false
			es.Log(DEBUG, "the sync chan is empty")
		}
	}

	go func(es *Exec) {
		result := p.sandbox.WaitProcess(false, []string{execId}, -1)
		if result == nil {
			es.Log(ERROR, "can not wait exec")
			return
		}

		r, ok := <-result
		if !ok {
			es.Log(ERROR, "waiting exec interrupted")
			return
		}

		es.Log(DEBUG, "exec terminated at %v with code %d", r.FinishedAt, r.Code)
		es.ExitCode = uint8(r.Code)
		select {
		case es.finChan <- true:
			es.Log(DEBUG, "wake exec stopped chan")
		default:
			es.Log(WARNING, "exec already set as stopped")
		}
	}(es)

	err := p.sandbox.Exec(es.Container, es.Id, es.Cmds, es.Terminal, tty)
	<-wReader.wait
	return err
}

func (p *XPod) GetExecExitCode(containerId, execId string) (uint8, error) {
	p.statusLock.RLock()
	es, ok := p.execs[execId]
	p.statusLock.RUnlock()

	if !ok {
		err := fmt.Errorf("no exec %s exists for container %s", execId, containerId)
		p.Log(ERROR, err)
		return 255, err
	}

	select {
	case <-es.finChan:
		es.finChan <- true
	case <-time.After(time.Second * 10):
		err := fmt.Errorf("wait exec exit code timeout")
		es.Log(ERROR, err)
		return 255, err
	}
	es.Log(INFO, "got exec exit code: %d", es.ExitCode)
	return es.ExitCode, nil
}

func (p *XPod) DeleteExec(containerId, execId string) {
	p.statusLock.Lock()
	delete(p.execs, execId)
	p.statusLock.Unlock()
}

func (p *XPod) CleanupExecs() {
	p.statusLock.Lock()
	p.execs = make(map[string]*Exec)
	p.statusLock.Unlock()
}
<commit_msg>setup process's envs & workdir same as container<commit_after>package pod

import (
	"encoding/json"
	"fmt"
	"io"
	"time"

	"github.com/docker/docker/pkg/stdcopy"

	"github.com/hyperhq/hypercontainer-utils/hlog"
	"github.com/hyperhq/hyperd/utils"
	"github.com/hyperhq/runv/hypervisor"
)

type Exec struct {
	Id        string
	Container string
	Cmds      []string
	Terminal  bool
	ExitCode  uint8

	logPrefix string
	finChan   chan bool
}

func (e *Exec) LogPrefix() string {
	return e.logPrefix
}

func (e *Exec) Log(level hlog.LogLevel, args ...interface{}) {
	hlog.HLog(level, e, 1, args...)
}

func (p *XPod) CreateExec(containerId, cmds string, terminal bool) (string, error) {
	c, ok := p.containers[containerId]
	if !ok {
		err := fmt.Errorf("no container available for exec %s", cmds)
		p.Log(ERROR, err)
		return "", err
	}

	if !c.IsAlive() {
		err := fmt.Errorf("container is not available (%v) for exec %s", c.CurrentState(), cmds)
		p.Log(ERROR, err)
		return "", err
	}
	var command []string
	if err := json.Unmarshal([]byte(cmds), &command); err != nil {
		return "", err
	}

	execId := fmt.Sprintf("exec-%s", utils.RandStr(10, "alpha"))

	p.statusLock.Lock()
	p.execs[execId] = &Exec{
		Container: containerId,
		Id:        execId,
		Cmds:      command,
		Terminal:  terminal,
		ExitCode:  255,
		logPrefix: fmt.Sprintf("Pod[%s] Con[%s] Exec[%s] ", p.Id(), containerId[:12], execId),
		finChan:   make(chan bool, 1),
	}
	p.statusLock.Unlock()

	return execId, nil
}

type waitClose struct {
	io.ReadCloser
	wait chan bool
}

func (wc *waitClose) Close() error {
	close(wc.wait)
	return wc.ReadCloser.Close()
}

func (p *XPod) StartExec(stdin io.ReadCloser, stdout io.WriteCloser, containerId, execId string) error {
	c, ok := p.containers[containerId]
	if !ok {
		err := fmt.Errorf("no container %s available for exec %s", containerId, execId)
		p.Log(ERROR, err)
		return err
	}

	p.statusLock.RLock()
	es, ok := p.execs[execId]
	p.statusLock.RUnlock()

	if !ok {
		err := fmt.Errorf("no exec %s exists for container %s", execId, containerId)
		p.Log(ERROR, err)
		return err
	}

	wReader := &waitClose{ReadCloser: stdin, wait: make(chan bool)}
	tty := &hypervisor.TtyIO{
		Stdin:  wReader,
		Stdout: stdout,
	}

	if !es.Terminal && stdout != nil {
		tty.Stderr = stdcopy.NewStdWriter(stdout, stdcopy.Stderr)
		tty.Stdout = stdcopy.NewStdWriter(stdout, stdcopy.Stdout)
	}

	var fin = true
	for fin {
		select {
		case fin = <-es.finChan:
			es.Log(DEBUG, "try to drain the sync chan")
		default:
			fin = false
			es.Log(DEBUG, "the sync chan is empty")
		}
	}

	go func(es *Exec) {
		result := p.sandbox.WaitProcess(false, []string{execId}, -1)
		if result == nil {
			es.Log(ERROR, "can not wait exec")
			return
		}

		r, ok := <-result
		if !ok {
			es.Log(ERROR, "waiting exec interrupted")
			return
		}

		es.Log(DEBUG, "exec terminated at %v with code %d", r.FinishedAt, r.Code)
		es.ExitCode = uint8(r.Code)
		select {
		case es.finChan <- true:
			es.Log(DEBUG, "wake exec stopped chan")
		default:
			es.Log(WARNING, "exec already set as stopped")
		}
	}(es)

	var envs []string
	for e, v := range c.descript.Envs {
		envs = append(envs, fmt.Sprintf("%s=%s", e, v))
	}

	err := p.sandbox.AddProcess(es.Container, es.Id, es.Terminal, es.Cmds, envs, c.descript.Workdir, tty)
	<-wReader.wait
	return err
}

func (p *XPod) GetExecExitCode(containerId, execId string) (uint8, error) {
	p.statusLock.RLock()
	es, ok := p.execs[execId]
	p.statusLock.RUnlock()

	if !ok {
		err := fmt.Errorf("no exec %s exists for container %s", execId, containerId)
		p.Log(ERROR, err)
		return 255, err
	}

	select {
	case <-es.finChan:
		es.finChan <- true
	case <-time.After(time.Second * 10):
		err := fmt.Errorf("wait exec exit code timeout")
		es.Log(ERROR, err)
		return 255, err
	}
	es.Log(INFO, "got exec exit code: %d", es.ExitCode)
	return es.ExitCode, nil
}

func (p *XPod) DeleteExec(containerId, execId string) {
	p.statusLock.Lock()
	delete(p.execs, execId)
	p.statusLock.Unlock()
}

func (p *XPod) CleanupExecs() {
	p.statusLock.Lock()
	p.execs = make(map[string]*Exec)
	p.statusLock.Unlock()
}
<|endoftext|>


<commit_before>// Copyright 2017, TCN Inc.
// All rights reserved.

// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:

//     * Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//     * Redistributions in binary form must reproduce the above
// copyright notice, this list of conditions and the following disclaimer
// in the documentation and/or other materials provided with the
// distribution.
//     * Neither the name of TCN Inc. nor the names of its
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.

// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

package generator

import (
	"github.com/Sirupsen/logrus"
	"github.com/golang/protobuf/proto"
	"github.com/golang/protobuf/protoc-gen-go/descriptor"
	"github.com/golang/protobuf/protoc-gen-go/plugin"
	"github.com/tcncloud/protoc-gen-persist/persist"
)

type Generator struct {
	OriginalRequest *plugin_go.CodeGeneratorRequest
	// Response        *plugin_go.CodeGeneratorResponse

	// structures that need Value() and Scan() implementation
	ImplementedStructures *StructList
	AllStructures         *StructList

	currentFile    *descriptor.FileDescriptorProto    // current processing file
	currentService *descriptor.ServiceDescriptorProto // current processing service

	files *FileList
}

func NewGenerator(request *plugin_go.CodeGeneratorRequest) *Generator {
	ret := new(Generator)
	ret.OriginalRequest = request
	// ret.Response = new(plugin_go.CodeGeneratorResponse)
	ret.ImplementedStructures = NewStructList()
	ret.AllStructures = NewStructList()
	ret.files = NewFileList()
	return ret
}

func (g *Generator) GetResponse() *plugin_go.CodeGeneratorResponse {
	ret := new(plugin_go.CodeGeneratorResponse)
	for _, fileStruct := range g.files.List {
		// format file Content

		ret.File = append(ret.File, &plugin_go.CodeGeneratorResponse_File{
			Content: proto.String(fileStruct.GetContent()),
			Name:    proto.String(fileStruct.Name),
		})
		ret.Error = proto.String(func() string {
			if ret.Error == nil {
				return fileStruct.ErrorList
			} else {
				return *ret.Error + "\n" + fileStruct.ErrorList
			}
		}())
	}
	logrus.WithField("response", ret).Debug("result")
	return ret
}

func (g *Generator) IsDependency(file *descriptor.FileDescriptorProto) bool {
	for _, f := range g.OriginalRequest.ProtoFile {
		for _, name := range f.GetDependency() {
			if name == file.GetName() {
				return true
			}
		}
	}
	return false
}

func (g *Generator) ProcessType(typ string) {
	struc := g.AllStructures.GetEntry(typ)
	if struc != nil {
		// if the structure is not defined into a dependency file
		if !g.IsDependency(struc.File) {
			outFile := g.files.NewOrGetFile(struc.File)
			if struc.IsMessage() {
				// process inner enums
				for _, e := range struc.MessageDescriptor.GetEnumType() {
					if eStruct := g.AllStructures.GetEnum(e, struc.MessageDescriptor, struc.File); eStruct != nil {
						if !g.ImplementedStructures.ContainEnum(eStruct.EnumDescriptor, eStruct.ParentDescriptor, eStruct.File) {
							outFile.P(eStruct.GetValueFunction())
							outFile.P(eStruct.GetScanFunction())
							g.ImplementedStructures.AddStruct(eStruct)
						}
					}
				}
				logrus.WithField("Is:", g.ImplementedStructures).Info("Status")
				// process inner messages
				for _, m := range struc.MessageDescriptor.GetNestedType() {
					if mStruct := g.AllStructures.GetMessage(m, struc.MessageDescriptor, struc.File); mStruct != nil {
						if !g.ImplementedStructures.ContainMessage(mStruct.MessageDescriptor, mStruct.ParentDescriptor, mStruct.File) {
							outFile.P(mStruct.GetValueFunction())
							outFile.P(mStruct.GetScanFunction())
							g.ImplementedStructures.AddStruct(mStruct)
						}
					}
				}
				logrus.WithField("Is:", g.ImplementedStructures).Info("Status")
				// process fields
				for _, f := range struc.MessageDescriptor.GetField() {
					logrus.WithField("Field", f).Info("processing field")
					if f.GetType() == descriptor.FieldDescriptorProto_TYPE_ENUM ||
						f.GetType() == descriptor.FieldDescriptorProto_TYPE_MESSAGE {
						if t := g.AllStructures.GetEntry(f.GetTypeName()); t != nil {
							if g.ImplementedStructures.GetEntry(f.GetTypeName()) == nil {
								if !g.IsDependency(t.File) {
									out := g.files.NewOrGetFile(t.File)
									out.P(t.GetValueFunction())
									out.P(t.GetScanFunction())
									g.ImplementedStructures.AddStruct(t)
								}
							}
						}
					}
				}
				logrus.WithField("Is:", g.ImplementedStructures).Info("Status")
			} else {
				// don't process structures that are not messages
			}
		}
	}
}

// process input and output types from service method signatures and implement Value() and Scan() mathods
// for messages and enums used inside of those types
func (g *Generator) ProcessStructs() {
	for _, file := range g.OriginalRequest.ProtoFile {
		if !g.IsDependency(file) {
			for _, service := range file.Service {
				for _, method := range service.GetMethod() {
					g.ProcessType(method.GetInputType())
					g.ProcessType(method.GetOutputType())
				}
			}
		}
	}
}

func (g *Generator) ProcessAllStructures() {
	for _, file := range g.OriginalRequest.ProtoFile {
		g.currentFile = file
		// scan all messages
		for _, m := range file.GetMessageType() {
			for _, im := range m.GetNestedType() {
				g.AllStructures.AddMessage(im, m, file)
			}
			for _, ie := range m.GetEnumType() {
				g.AllStructures.AddEnum(ie, m, file)
			}
			g.AllStructures.AddMessage(m, nil, file)

		}
		// scan all enums
		for _, e := range file.GetEnumType() {
			g.AllStructures.AddEnum(e, nil, file)
		}
	}
}

func (g *Generator) ProcessServices() {

	for _, file := range g.OriginalRequest.ProtoFile {
		if !g.IsDependency(file) {
			// outFile := g.files.NewOrGetFile(file)
			g.currentFile = file
			// implement file

			for _, service := range file.Service {
				g.currentService = service
				if IsServicePersistEnabled(service) {
					for _, _ = range service.Method {
					}
				}
			}

		}
	}
}

// check if a service has at least one method that has the persist.ql extension defined
func IsServicePersistEnabled(service *descriptor.ServiceDescriptorProto) bool {
	if service.Method != nil {
		for _, method := range service.Method {
			if IsMethodEnabled(method) {
				return true
			}
		}
	}
	return false
}

func IsMethodEnabled(method *descriptor.MethodDescriptorProto) bool {
	if method != nil && method.GetOptions() != nil && proto.HasExtension(method.Options, persist.E_Ql) {
		return true
	}
	return false
}

// Process the request
func (g *Generator) ProcessRequest() {
	g.ProcessAllStructures()
	g.ProcessStructs()
	g.ProcessServices()
}
<commit_msg>added GetMethodOption<commit_after>// Copyright 2017, TCN Inc.
// All rights reserved.

// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:

//     * Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//     * Redistributions in binary form must reproduce the above
// copyright notice, this list of conditions and the following disclaimer
// in the documentation and/or other materials provided with the
// distribution.
//     * Neither the name of TCN Inc. nor the names of its
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.

// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

package generator

import (
	"github.com/Sirupsen/logrus"
	"github.com/golang/protobuf/proto"
	"github.com/golang/protobuf/protoc-gen-go/descriptor"
	"github.com/golang/protobuf/protoc-gen-go/plugin"
	"github.com/tcncloud/protoc-gen-persist/persist"
)

type Generator struct {
	OriginalRequest *plugin_go.CodeGeneratorRequest
	// Response        *plugin_go.CodeGeneratorResponse

	// structures that need Value() and Scan() implementation
	ImplementedStructures *StructList
	AllStructures         *StructList

	currentFile    *descriptor.FileDescriptorProto    // current processing file
	currentService *descriptor.ServiceDescriptorProto // current processing service

	files *FileList
}

func NewGenerator(request *plugin_go.CodeGeneratorRequest) *Generator {
	ret := new(Generator)
	ret.OriginalRequest = request
	// ret.Response = new(plugin_go.CodeGeneratorResponse)
	ret.ImplementedStructures = NewStructList()
	ret.AllStructures = NewStructList()
	ret.files = NewFileList()
	return ret
}

func (g *Generator) GetResponse() *plugin_go.CodeGeneratorResponse {
	ret := new(plugin_go.CodeGeneratorResponse)
	for _, fileStruct := range g.files.List {
		// format file Content

		ret.File = append(ret.File, &plugin_go.CodeGeneratorResponse_File{
			Content: proto.String(fileStruct.GetContent()),
			Name:    proto.String(fileStruct.Name),
		})
		ret.Error = proto.String(func() string {
			if ret.Error == nil {
				return fileStruct.ErrorList
			} else {
				return *ret.Error + "\n" + fileStruct.ErrorList
			}
		}())
	}
	logrus.WithField("response", ret).Debug("result")
	return ret
}

func (g *Generator) IsDependency(file *descriptor.FileDescriptorProto) bool {
	for _, f := range g.OriginalRequest.ProtoFile {
		for _, name := range f.GetDependency() {
			if name == file.GetName() {
				return true
			}
		}
	}
	return false
}

func (g *Generator) ProcessType(typ string) {
	struc := g.AllStructures.GetEntry(typ)
	if struc != nil {
		// if the structure is not defined into a dependency file
		if !g.IsDependency(struc.File) {
			outFile := g.files.NewOrGetFile(struc.File)
			if struc.IsMessage() {
				// process inner enums
				for _, e := range struc.MessageDescriptor.GetEnumType() {
					if eStruct := g.AllStructures.GetEnum(e, struc.MessageDescriptor, struc.File); eStruct != nil {
						if !g.ImplementedStructures.ContainEnum(eStruct.EnumDescriptor, eStruct.ParentDescriptor, eStruct.File) {
							outFile.P(eStruct.GetValueFunction())
							outFile.P(eStruct.GetScanFunction())
							g.ImplementedStructures.AddStruct(eStruct)
						}
					}
				}
				// process inner messages
				for _, m := range struc.MessageDescriptor.GetNestedType() {
					if mStruct := g.AllStructures.GetMessage(m, struc.MessageDescriptor, struc.File); mStruct != nil {
						if !g.ImplementedStructures.ContainMessage(mStruct.MessageDescriptor, mStruct.ParentDescriptor, mStruct.File) {
							outFile.P(mStruct.GetValueFunction())
							outFile.P(mStruct.GetScanFunction())
							g.ImplementedStructures.AddStruct(mStruct)
						}
					}
				}
				// process fields
				for _, f := range struc.MessageDescriptor.GetField() {
					if f.GetType() == descriptor.FieldDescriptorProto_TYPE_ENUM ||
						f.GetType() == descriptor.FieldDescriptorProto_TYPE_MESSAGE {
						if t := g.AllStructures.GetEntry(f.GetTypeName()); t != nil {
							if g.ImplementedStructures.GetEntry(f.GetTypeName()) == nil {
								if !g.IsDependency(t.File) {
									out := g.files.NewOrGetFile(t.File)
									out.P(t.GetValueFunction())
									out.P(t.GetScanFunction())
									g.ImplementedStructures.AddStruct(t)
								}
							}
						}
					}
				}
			} else {
				// don't process structures that are not messages
			}
		}
	}
}

// process input and output types from service method signatures and implement Value() and Scan() mathods
// for messages and enums used inside of those types
func (g *Generator) ProcessStructs() {
	for _, file := range g.OriginalRequest.ProtoFile {
		if !g.IsDependency(file) {
			for _, service := range file.Service {
				for _, method := range service.GetMethod() {
					g.ProcessType(method.GetInputType())
					g.ProcessType(method.GetOutputType())
				}
			}
		}
	}
}

func (g *Generator) ProcessAllStructures() {
	for _, file := range g.OriginalRequest.ProtoFile {
		g.currentFile = file
		// scan all messages
		for _, m := range file.GetMessageType() {
			for _, im := range m.GetNestedType() {
				g.AllStructures.AddMessage(im, m, file)
			}
			for _, ie := range m.GetEnumType() {
				g.AllStructures.AddEnum(ie, m, file)
			}
			g.AllStructures.AddMessage(m, nil, file)

		}
		// scan all enums
		for _, e := range file.GetEnumType() {
			g.AllStructures.AddEnum(e, nil, file)
		}
	}
}

func (g *Generator) ProcessServices() {

	for _, file := range g.OriginalRequest.ProtoFile {
		if !g.IsDependency(file) {
			outFile := g.files.NewOrGetFile(file)

			for _, service := range file.Service {
				g.currentService = service
				if IsServicePersistEnabled(service) {
					srv := NewService(service, file, g.AllStructures, g.files)
					outFile.P(srv.Generate())
				}
			}
		}
	}
}

// check if a service has at least one method that has the persist.ql extension defined
func IsServicePersistEnabled(service *descriptor.ServiceDescriptorProto) bool {
	if service.Method != nil {
		for _, method := range service.Method {
			if IsMethodEnabled(method) {
				return true
			}
		}
	}
	return false
}

func IsMethodEnabled(method *descriptor.MethodDescriptorProto) bool {
	if method != nil && method.GetOptions() != nil && proto.HasExtension(method.Options, persist.E_Ql) {
		return true
	}
	return false
}

func GetMethodOption(method *descriptor.MethodDescriptorProto) *persist.QLImpl {
	if IsMethodEnabled(method) {
		if ret, err := proto.GetExtension(method.Options, persist.E_Ql); err == nil {
			return ret.(*persist.QLImpl)
		}
	}
	return nil
}

// Process the request
func (g *Generator) ProcessRequest() {
	g.ProcessAllStructures()
	g.ProcessStructs()
	g.ProcessServices()
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package daemon

import (
	"fmt"

	"github.com/golang/glog"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	utilerrors "k8s.io/apimachinery/pkg/util/errors"
	intstrutil "k8s.io/apimachinery/pkg/util/intstr"
	"k8s.io/kubernetes/pkg/api/v1"
	podutil "k8s.io/kubernetes/pkg/api/v1/pod"
	extensions "k8s.io/kubernetes/pkg/apis/extensions/v1beta1"
	"k8s.io/kubernetes/pkg/controller/daemon/util"
)

// rollingUpdate deletes old daemon set pods making sure that no more than
// ds.Spec.UpdateStrategy.RollingUpdate.MaxUnavailable pods are unavailable
func (dsc *DaemonSetsController) rollingUpdate(ds *extensions.DaemonSet) error {
	nodeToDaemonPods, err := dsc.getNodesToDaemonPods(ds)
	if err != nil {
		return fmt.Errorf("couldn't get node to daemon pod mapping for daemon set %q: %v", ds.Name, err)
	}

	_, oldPods, err := dsc.getAllDaemonSetPods(ds, nodeToDaemonPods)
	maxUnavailable, numUnavailable, err := dsc.getUnavailableNumbers(ds, nodeToDaemonPods)
	if err != nil {
		return fmt.Errorf("Couldn't get unavailable numbers: %v", err)
	}
	oldAvailablePods, oldUnavailablePods := util.SplitByAvailablePods(ds.Spec.MinReadySeconds, oldPods)

	// for oldPods delete all not running pods
	var podsToDelete []string
	glog.V(4).Infof("Marking all unavailable old pods for deletion")
	for _, pod := range oldUnavailablePods {
		// Skip terminating pods. We won't delete them again
		if pod.DeletionTimestamp != nil {
			continue
		}
		glog.V(4).Infof("Marking pod %s/%s for deletion", ds.Name, pod.Name)
		podsToDelete = append(podsToDelete, pod.Name)
	}

	glog.V(4).Infof("Marking old pods for deletion")
	for _, pod := range oldAvailablePods {
		if numUnavailable >= maxUnavailable {
			glog.V(4).Infof("Number of unavailable DaemonSet pods: %d, is equal to or exceeds allowed maximum: %d", numUnavailable, maxUnavailable)
			break
		}
		glog.V(4).Infof("Marking pod %s/%s for deletion", ds.Name, pod.Name)
		podsToDelete = append(podsToDelete, pod.Name)
		numUnavailable++
	}
	errors := dsc.syncNodes(ds, podsToDelete, []string{})
	return utilerrors.NewAggregate(errors)
}

func (dsc *DaemonSetsController) getAllDaemonSetPods(ds *extensions.DaemonSet, nodeToDaemonPods map[string][]*v1.Pod) ([]*v1.Pod, []*v1.Pod, error) {
	var newPods []*v1.Pod
	var oldPods []*v1.Pod

	for _, pods := range nodeToDaemonPods {
		for _, pod := range pods {
			if util.IsPodUpdated(ds.Spec.TemplateGeneration, pod) {
				newPods = append(newPods, pod)
			} else {
				oldPods = append(oldPods, pod)
			}
		}
	}
	return newPods, oldPods, nil
}

func (dsc *DaemonSetsController) getUnavailableNumbers(ds *extensions.DaemonSet, nodeToDaemonPods map[string][]*v1.Pod) (int, int, error) {
	glog.V(4).Infof("Getting unavailable numbers")
	// TODO: get nodeList once in syncDaemonSet and pass it to other functions
	nodeList, err := dsc.nodeLister.List(labels.Everything())
	if err != nil {
		return -1, -1, fmt.Errorf("couldn't get list of nodes during rolling update of daemon set %#v: %v", ds, err)
	}

	var numUnavailable, desiredNumberScheduled int
	for i := range nodeList {
		node := nodeList[i]
		wantToRun, _, _, err := dsc.nodeShouldRunDaemonPod(node, ds)
		if err != nil {
			return -1, -1, err
		}
		if !wantToRun {
			continue
		}
		desiredNumberScheduled++
		daemonPods, exists := nodeToDaemonPods[node.Name]
		if !exists {
			numUnavailable++
			continue
		}
		available := false
		for _, pod := range daemonPods {
			if podutil.IsPodAvailable(pod, ds.Spec.MinReadySeconds, metav1.Now()) {
				available = true
				break
			}
		}
		if !available {
			numUnavailable++
		}
	}
	maxUnavailable, err := intstrutil.GetValueFromIntOrPercent(ds.Spec.UpdateStrategy.RollingUpdate.MaxUnavailable, desiredNumberScheduled, true)
	if err != nil {
		return -1, -1, fmt.Errorf("Invalid value for MaxUnavailable: %v", err)
	}
	glog.V(4).Infof(" DaemonSet %s/%s, maxUnavailable: %d, numUnavailable: %d", ds.Namespace, ds.Name, maxUnavailable, numUnavailable)
	return maxUnavailable, numUnavailable, nil
}
<commit_msg>delete err when return _<commit_after>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package daemon

import (
	"fmt"

	"github.com/golang/glog"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	utilerrors "k8s.io/apimachinery/pkg/util/errors"
	intstrutil "k8s.io/apimachinery/pkg/util/intstr"
	"k8s.io/kubernetes/pkg/api/v1"
	podutil "k8s.io/kubernetes/pkg/api/v1/pod"
	extensions "k8s.io/kubernetes/pkg/apis/extensions/v1beta1"
	"k8s.io/kubernetes/pkg/controller/daemon/util"
)

// rollingUpdate deletes old daemon set pods making sure that no more than
// ds.Spec.UpdateStrategy.RollingUpdate.MaxUnavailable pods are unavailable
func (dsc *DaemonSetsController) rollingUpdate(ds *extensions.DaemonSet) error {
	nodeToDaemonPods, err := dsc.getNodesToDaemonPods(ds)
	if err != nil {
		return fmt.Errorf("couldn't get node to daemon pod mapping for daemon set %q: %v", ds.Name, err)
	}

	_, oldPods := dsc.getAllDaemonSetPods(ds, nodeToDaemonPods)
	maxUnavailable, numUnavailable, err := dsc.getUnavailableNumbers(ds, nodeToDaemonPods)
	if err != nil {
		return fmt.Errorf("Couldn't get unavailable numbers: %v", err)
	}
	oldAvailablePods, oldUnavailablePods := util.SplitByAvailablePods(ds.Spec.MinReadySeconds, oldPods)

	// for oldPods delete all not running pods
	var podsToDelete []string
	glog.V(4).Infof("Marking all unavailable old pods for deletion")
	for _, pod := range oldUnavailablePods {
		// Skip terminating pods. We won't delete them again
		if pod.DeletionTimestamp != nil {
			continue
		}
		glog.V(4).Infof("Marking pod %s/%s for deletion", ds.Name, pod.Name)
		podsToDelete = append(podsToDelete, pod.Name)
	}

	glog.V(4).Infof("Marking old pods for deletion")
	for _, pod := range oldAvailablePods {
		if numUnavailable >= maxUnavailable {
			glog.V(4).Infof("Number of unavailable DaemonSet pods: %d, is equal to or exceeds allowed maximum: %d", numUnavailable, maxUnavailable)
			break
		}
		glog.V(4).Infof("Marking pod %s/%s for deletion", ds.Name, pod.Name)
		podsToDelete = append(podsToDelete, pod.Name)
		numUnavailable++
	}
	errors := dsc.syncNodes(ds, podsToDelete, []string{})
	return utilerrors.NewAggregate(errors)
}

func (dsc *DaemonSetsController) getAllDaemonSetPods(ds *extensions.DaemonSet, nodeToDaemonPods map[string][]*v1.Pod) ([]*v1.Pod, []*v1.Pod) {
	var newPods []*v1.Pod
	var oldPods []*v1.Pod

	for _, pods := range nodeToDaemonPods {
		for _, pod := range pods {
			if util.IsPodUpdated(ds.Spec.TemplateGeneration, pod) {
				newPods = append(newPods, pod)
			} else {
				oldPods = append(oldPods, pod)
			}
		}
	}
	return newPods, oldPods
}

func (dsc *DaemonSetsController) getUnavailableNumbers(ds *extensions.DaemonSet, nodeToDaemonPods map[string][]*v1.Pod) (int, int, error) {
	glog.V(4).Infof("Getting unavailable numbers")
	// TODO: get nodeList once in syncDaemonSet and pass it to other functions
	nodeList, err := dsc.nodeLister.List(labels.Everything())
	if err != nil {
		return -1, -1, fmt.Errorf("couldn't get list of nodes during rolling update of daemon set %#v: %v", ds, err)
	}

	var numUnavailable, desiredNumberScheduled int
	for i := range nodeList {
		node := nodeList[i]
		wantToRun, _, _, err := dsc.nodeShouldRunDaemonPod(node, ds)
		if err != nil {
			return -1, -1, err
		}
		if !wantToRun {
			continue
		}
		desiredNumberScheduled++
		daemonPods, exists := nodeToDaemonPods[node.Name]
		if !exists {
			numUnavailable++
			continue
		}
		available := false
		for _, pod := range daemonPods {
			if podutil.IsPodAvailable(pod, ds.Spec.MinReadySeconds, metav1.Now()) {
				available = true
				break
			}
		}
		if !available {
			numUnavailable++
		}
	}
	maxUnavailable, err := intstrutil.GetValueFromIntOrPercent(ds.Spec.UpdateStrategy.RollingUpdate.MaxUnavailable, desiredNumberScheduled, true)
	if err != nil {
		return -1, -1, fmt.Errorf("Invalid value for MaxUnavailable: %v", err)
	}
	glog.V(4).Infof(" DaemonSet %s/%s, maxUnavailable: %d, numUnavailable: %d", ds.Namespace, ds.Name, maxUnavailable, numUnavailable)
	return maxUnavailable, numUnavailable, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2017 Google 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 sqlparser

// analyzer.go contains utility analysis functions.

import (
	"errors"
	"fmt"
	"strconv"
	"strings"
	"unicode"

	"github.com/xwb1989/sqlparser/dependency/sqltypes"
)

// These constants are used to identify the SQL statement type.
const (
	StmtSelect = iota
	StmtInsert
	StmtReplace
	StmtUpdate
	StmtDelete
	StmtDDL
	StmtBegin
	StmtCommit
	StmtRollback
	StmtSet
	StmtShow
	StmtUse
	StmtOther
	StmtUnknown
	StmtComment
)

// Preview analyzes the beginning of the query using a simpler and faster
// textual comparison to identify the statement type.
func Preview(sql string) int {
	trimmed := StripLeadingComments(sql)
	trimmedNoComments, _ := SplitTrailingComments(trimmed)

	firstWord := trimmed
	if end := strings.IndexFunc(trimmed, unicode.IsSpace); end != -1 {
		firstWord = trimmed[:end]
	}

	// Comparison is done in order of priority.
	loweredFirstWord := strings.ToLower(firstWord)
	switch loweredFirstWord {
	case "select":
		return StmtSelect
	case "insert":
		return StmtInsert
	case "replace":
		return StmtReplace
	case "update":
		return StmtUpdate
	case "delete":
		return StmtDelete
	}
	switch strings.ToLower(trimmedNoComments) {
	case "begin", "start transaction":
		return StmtBegin
	case "commit":
		return StmtCommit
	case "rollback":
		return StmtRollback
	}
	switch loweredFirstWord {
	case "create", "alter", "rename", "drop", "truncate":
		return StmtDDL
	case "set":
		return StmtSet
	case "show":
		return StmtShow
	case "use":
		return StmtUse
	case "analyze", "describe", "desc", "explain", "repair", "optimize":
		return StmtOther
	}
	if strings.Index(trimmed, "/*!") == 0 {
		return StmtComment
	}
	return StmtUnknown
}

// StmtType returns the statement type as a string
func StmtType(stmtType int) string {
	switch stmtType {
	case StmtSelect:
		return "SELECT"
	case StmtInsert:
		return "INSERT"
	case StmtReplace:
		return "REPLACE"
	case StmtUpdate:
		return "UPDATE"
	case StmtDelete:
		return "DELETE"
	case StmtDDL:
		return "DDL"
	case StmtBegin:
		return "BEGIN"
	case StmtCommit:
		return "COMMIT"
	case StmtRollback:
		return "ROLLBACK"
	case StmtSet:
		return "SET"
	case StmtShow:
		return "SHOW"
	case StmtUse:
		return "USE"
	case StmtOther:
		return "OTHER"
	default:
		return "UNKNOWN"
	}
}

// IsDML returns true if the query is an INSERT, UPDATE or DELETE statement.
func IsDML(sql string) bool {
	switch Preview(sql) {
	case StmtInsert, StmtReplace, StmtUpdate, StmtDelete:
		return true
	}
	return false
}

// GetTableName returns the table name from the SimpleTableExpr
// only if it's a simple expression. Otherwise, it returns "".
func GetTableName(node SimpleTableExpr) TableIdent {
	if n, ok := node.(TableName); ok && n.Qualifier.IsEmpty() {
		return n.Name
	}
	// sub-select or '.' expression
	return NewTableIdent("")
}

// IsColName returns true if the Expr is a *ColName.
func IsColName(node Expr) bool {
	_, ok := node.(*ColName)
	return ok
}

// IsValue returns true if the Expr is a string, integral or value arg.
// NULL is not considered to be a value.
func IsValue(node Expr) bool {
	switch v := node.(type) {
	case *SQLVal:
		switch v.Type {
		case StrVal, HexVal, IntVal, ValArg:
			return true
		}
	}
	return false
}

// IsNull returns true if the Expr is SQL NULL
func IsNull(node Expr) bool {
	switch node.(type) {
	case *NullVal:
		return true
	}
	return false
}

// IsSimpleTuple returns true if the Expr is a ValTuple that
// contains simple values or if it's a list arg.
func IsSimpleTuple(node Expr) bool {
	switch vals := node.(type) {
	case ValTuple:
		for _, n := range vals {
			if !IsValue(n) {
				return false
			}
		}
		return true
	case ListArg:
		return true
	}
	// It's a subquery
	return false
}

// NewPlanValue builds a sqltypes.PlanValue from an Expr.
func NewPlanValue(node Expr) (sqltypes.PlanValue, error) {
	switch node := node.(type) {
	case *SQLVal:
		switch node.Type {
		case ValArg:
			return sqltypes.PlanValue{Key: string(node.Val[1:])}, nil
		case IntVal:
			n, err := sqltypes.NewIntegral(string(node.Val))
			if err != nil {
				return sqltypes.PlanValue{}, fmt.Errorf("%v", err)
			}
			return sqltypes.PlanValue{Value: n}, nil
		case StrVal:
			return sqltypes.PlanValue{Value: sqltypes.MakeTrusted(sqltypes.VarBinary, node.Val)}, nil
		case HexVal:
			v, err := node.HexDecode()
			if err != nil {
				return sqltypes.PlanValue{}, fmt.Errorf("%v", err)
			}
			return sqltypes.PlanValue{Value: sqltypes.MakeTrusted(sqltypes.VarBinary, v)}, nil
		}
	case ListArg:
		return sqltypes.PlanValue{ListKey: string(node[2:])}, nil
	case ValTuple:
		pv := sqltypes.PlanValue{
			Values: make([]sqltypes.PlanValue, 0, len(node)),
		}
		for _, val := range node {
			innerpv, err := NewPlanValue(val)
			if err != nil {
				return sqltypes.PlanValue{}, err
			}
			if innerpv.ListKey != "" || innerpv.Values != nil {
				return sqltypes.PlanValue{}, errors.New("unsupported: nested lists")
			}
			pv.Values = append(pv.Values, innerpv)
		}
		return pv, nil
	case *NullVal:
		return sqltypes.PlanValue{}, nil
	}
	return sqltypes.PlanValue{}, fmt.Errorf("expression is too complex '%v'", String(node))
}

// StringIn is a convenience function that returns
// true if str matches any of the values.
func StringIn(str string, values ...string) bool {
	for _, val := range values {
		if str == val {
			return true
		}
	}
	return false
}

// ExtractSetValues returns a map of key-value pairs
// if the query is a SET statement. Values can be int64 or string.
// Since set variable names are case insensitive, all keys are returned
// as lower case.
func ExtractSetValues(sql string) (keyValues map[string]interface{}, charset string, scope string, err error) {
	stmt, err := Parse(sql)
	if err != nil {
		return nil, "", "", err
	}
	setStmt, ok := stmt.(*Set)
	if !ok {
		return nil, "", "", fmt.Errorf("ast did not yield *sqlparser.Set: %T", stmt)
	}
	result := make(map[string]interface{})
	for _, expr := range setStmt.Exprs {
		if !expr.Name.Qualifier.IsEmpty() {
			return nil, "", "", fmt.Errorf("invalid syntax: %v", String(expr.Name))
		}
		key := expr.Name.Name.Lowered()

		switch expr := expr.Expr.(type) {
		case *SQLVal:
			switch expr.Type {
			case StrVal:
				result[key] = string(expr.Val)
			case IntVal:
				num, err := strconv.ParseInt(string(expr.Val), 0, 64)
				if err != nil {
					return nil, "", "", err
				}
				result[key] = num
			default:
				return nil, "", "", fmt.Errorf("invalid value type: %v", String(expr))
			}
		case *NullVal:
			result[key] = nil
		case *Default:
			result[key] = "default"
		default:
			return nil, "", "", fmt.Errorf("invalid syntax: %s", String(expr))
		}
	}
	return result, setStmt.Charset.Lowered(), strings.ToLower(setStmt.Scope), nil
}
<commit_msg>nit - only get trimmed comments when needed<commit_after>/*
Copyright 2017 Google 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 sqlparser

// analyzer.go contains utility analysis functions.

import (
	"errors"
	"fmt"
	"strconv"
	"strings"
	"unicode"

	"github.com/xwb1989/sqlparser/dependency/sqltypes"
)

// These constants are used to identify the SQL statement type.
const (
	StmtSelect = iota
	StmtInsert
	StmtReplace
	StmtUpdate
	StmtDelete
	StmtDDL
	StmtBegin
	StmtCommit
	StmtRollback
	StmtSet
	StmtShow
	StmtUse
	StmtOther
	StmtUnknown
	StmtComment
)

// Preview analyzes the beginning of the query using a simpler and faster
// textual comparison to identify the statement type.
func Preview(sql string) int {
	trimmed := StripLeadingComments(sql)

	firstWord := trimmed
	if end := strings.IndexFunc(trimmed, unicode.IsSpace); end != -1 {
		firstWord = trimmed[:end]
	}

	// Comparison is done in order of priority.
	loweredFirstWord := strings.ToLower(firstWord)
	switch loweredFirstWord {
	case "select":
		return StmtSelect
	case "insert":
		return StmtInsert
	case "replace":
		return StmtReplace
	case "update":
		return StmtUpdate
	case "delete":
		return StmtDelete
	}
	trimmedNoComments, _ := SplitTrailingComments(trimmed)
	switch strings.ToLower(trimmedNoComments) {
	case "begin", "start transaction":
		return StmtBegin
	case "commit":
		return StmtCommit
	case "rollback":
		return StmtRollback
	}
	switch loweredFirstWord {
	case "create", "alter", "rename", "drop", "truncate":
		return StmtDDL
	case "set":
		return StmtSet
	case "show":
		return StmtShow
	case "use":
		return StmtUse
	case "analyze", "describe", "desc", "explain", "repair", "optimize":
		return StmtOther
	}
	if strings.Index(trimmed, "/*!") == 0 {
		return StmtComment
	}
	return StmtUnknown
}

// StmtType returns the statement type as a string
func StmtType(stmtType int) string {
	switch stmtType {
	case StmtSelect:
		return "SELECT"
	case StmtInsert:
		return "INSERT"
	case StmtReplace:
		return "REPLACE"
	case StmtUpdate:
		return "UPDATE"
	case StmtDelete:
		return "DELETE"
	case StmtDDL:
		return "DDL"
	case StmtBegin:
		return "BEGIN"
	case StmtCommit:
		return "COMMIT"
	case StmtRollback:
		return "ROLLBACK"
	case StmtSet:
		return "SET"
	case StmtShow:
		return "SHOW"
	case StmtUse:
		return "USE"
	case StmtOther:
		return "OTHER"
	default:
		return "UNKNOWN"
	}
}

// IsDML returns true if the query is an INSERT, UPDATE or DELETE statement.
func IsDML(sql string) bool {
	switch Preview(sql) {
	case StmtInsert, StmtReplace, StmtUpdate, StmtDelete:
		return true
	}
	return false
}

// GetTableName returns the table name from the SimpleTableExpr
// only if it's a simple expression. Otherwise, it returns "".
func GetTableName(node SimpleTableExpr) TableIdent {
	if n, ok := node.(TableName); ok && n.Qualifier.IsEmpty() {
		return n.Name
	}
	// sub-select or '.' expression
	return NewTableIdent("")
}

// IsColName returns true if the Expr is a *ColName.
func IsColName(node Expr) bool {
	_, ok := node.(*ColName)
	return ok
}

// IsValue returns true if the Expr is a string, integral or value arg.
// NULL is not considered to be a value.
func IsValue(node Expr) bool {
	switch v := node.(type) {
	case *SQLVal:
		switch v.Type {
		case StrVal, HexVal, IntVal, ValArg:
			return true
		}
	}
	return false
}

// IsNull returns true if the Expr is SQL NULL
func IsNull(node Expr) bool {
	switch node.(type) {
	case *NullVal:
		return true
	}
	return false
}

// IsSimpleTuple returns true if the Expr is a ValTuple that
// contains simple values or if it's a list arg.
func IsSimpleTuple(node Expr) bool {
	switch vals := node.(type) {
	case ValTuple:
		for _, n := range vals {
			if !IsValue(n) {
				return false
			}
		}
		return true
	case ListArg:
		return true
	}
	// It's a subquery
	return false
}

// NewPlanValue builds a sqltypes.PlanValue from an Expr.
func NewPlanValue(node Expr) (sqltypes.PlanValue, error) {
	switch node := node.(type) {
	case *SQLVal:
		switch node.Type {
		case ValArg:
			return sqltypes.PlanValue{Key: string(node.Val[1:])}, nil
		case IntVal:
			n, err := sqltypes.NewIntegral(string(node.Val))
			if err != nil {
				return sqltypes.PlanValue{}, fmt.Errorf("%v", err)
			}
			return sqltypes.PlanValue{Value: n}, nil
		case StrVal:
			return sqltypes.PlanValue{Value: sqltypes.MakeTrusted(sqltypes.VarBinary, node.Val)}, nil
		case HexVal:
			v, err := node.HexDecode()
			if err != nil {
				return sqltypes.PlanValue{}, fmt.Errorf("%v", err)
			}
			return sqltypes.PlanValue{Value: sqltypes.MakeTrusted(sqltypes.VarBinary, v)}, nil
		}
	case ListArg:
		return sqltypes.PlanValue{ListKey: string(node[2:])}, nil
	case ValTuple:
		pv := sqltypes.PlanValue{
			Values: make([]sqltypes.PlanValue, 0, len(node)),
		}
		for _, val := range node {
			innerpv, err := NewPlanValue(val)
			if err != nil {
				return sqltypes.PlanValue{}, err
			}
			if innerpv.ListKey != "" || innerpv.Values != nil {
				return sqltypes.PlanValue{}, errors.New("unsupported: nested lists")
			}
			pv.Values = append(pv.Values, innerpv)
		}
		return pv, nil
	case *NullVal:
		return sqltypes.PlanValue{}, nil
	}
	return sqltypes.PlanValue{}, fmt.Errorf("expression is too complex '%v'", String(node))
}

// StringIn is a convenience function that returns
// true if str matches any of the values.
func StringIn(str string, values ...string) bool {
	for _, val := range values {
		if str == val {
			return true
		}
	}
	return false
}

// ExtractSetValues returns a map of key-value pairs
// if the query is a SET statement. Values can be int64 or string.
// Since set variable names are case insensitive, all keys are returned
// as lower case.
func ExtractSetValues(sql string) (keyValues map[string]interface{}, charset string, scope string, err error) {
	stmt, err := Parse(sql)
	if err != nil {
		return nil, "", "", err
	}
	setStmt, ok := stmt.(*Set)
	if !ok {
		return nil, "", "", fmt.Errorf("ast did not yield *sqlparser.Set: %T", stmt)
	}
	result := make(map[string]interface{})
	for _, expr := range setStmt.Exprs {
		if !expr.Name.Qualifier.IsEmpty() {
			return nil, "", "", fmt.Errorf("invalid syntax: %v", String(expr.Name))
		}
		key := expr.Name.Name.Lowered()

		switch expr := expr.Expr.(type) {
		case *SQLVal:
			switch expr.Type {
			case StrVal:
				result[key] = string(expr.Val)
			case IntVal:
				num, err := strconv.ParseInt(string(expr.Val), 0, 64)
				if err != nil {
					return nil, "", "", err
				}
				result[key] = num
			default:
				return nil, "", "", fmt.Errorf("invalid value type: %v", String(expr))
			}
		case *NullVal:
			result[key] = nil
		case *Default:
			result[key] = "default"
		default:
			return nil, "", "", fmt.Errorf("invalid syntax: %s", String(expr))
		}
	}
	return result, setStmt.Charset.Lowered(), strings.ToLower(setStmt.Scope), nil
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package slice

import (
	"sort"
)

// Int64Slice attaches the methods of Interface to []int64,
// sorting in increasing order.
type Int64Slice []int64

func (p Int64Slice) Len() int           { return len(p) }
func (p Int64Slice) Less(i, j int) bool { return p[i] < p[j] }
func (p Int64Slice) Swap(i, j int)      { p[i], p[j] = p[j], p[i] }

// Sorts []int64 in increasing order
func SortInts64(a []int64) { sort.Sort(Int64Slice(a)) }
<commit_msg>Simplify the sorting codes<commit_after>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package slice

import (
	"sort"
)

// Sorts []int64 in increasing order
func SortInts64(a []int64) { sort.Slice(a, func(i, j int) bool { return a[i] < a[j] }) }
<|endoftext|>


<commit_before>/*
Copyright 2020 The Knative Authors

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

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

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

*/

package environment

import (
	"flag"
	"os"
	"os/user"
	"path"

	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"
)

type Cluster struct {
	Name       string // K8s cluster (defaults to cluster in kubeconfig)
	KubeConfig string // Path to kubeconfig (defaults to ./kube/config)
}

func (s *Cluster) AddFlags(fs *flag.FlagSet) {
	fs.StringVar(&s.Name, "env.cluster.name", "",
		"Provide the cluster to test against. Defaults to the current cluster in kubeconfig.")

	// Use KUBECONFIG if available
	defaultKubeconfig := os.Getenv("KUBECONFIG")

	// If KUBECONFIG env var isn't set then look for $HOME/.kube/config
	if defaultKubeconfig == "" {
		if usr, err := user.Current(); err == nil {
			defaultKubeconfig = path.Join(usr.HomeDir, ".kube/config")
		}
	}

	// Allow for --kubeconfig on the cmd line to override the above logic
	fs.StringVar(&s.KubeConfig, "env.cluster.kubeconfig", defaultKubeconfig,
		"Provide the path to the `kubeconfig` file you'd like to use for these tests. The `current-context` will be used.")
}

func (c *Cluster) ClientConfig() *rest.Config {
	overrides := &clientcmd.ConfigOverrides{}
	overrides.Context.Cluster = c.Name

	loader := &clientcmd.ClientConfigLoadingRules{ExplicitPath: c.KubeConfig}

	conf, err := clientcmd.NewNonInteractiveDeferredLoadingClientConfig(loader, overrides).ClientConfig()
	if err != nil {
		panic(err)
	}
	return conf
}
<commit_msg>support configuring cluster domains (#20)<commit_after>/*
Copyright 2020 The Knative Authors

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

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

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

*/

package environment

import (
	"flag"
	"os"
	"os/user"
	"path"

	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"
)

type Cluster struct {
	Name       string // K8s cluster (defaults to cluster in kubeconfig)
	KubeConfig string // Path to kubeconfig (defaults to ./kube/config)
	DomainName string // ie. cluster.local
}

func (s *Cluster) AddFlags(fs *flag.FlagSet) {
	fs.StringVar(&s.Name, "env.cluster.name", "",
		"Provide the cluster to test against. Defaults to the current cluster in kubeconfig.")

	// Allow for --env.cluster.kubeconfig to override the KUBECONFIG env var
	fs.StringVar(&s.KubeConfig, "env.cluster.kubeconfig", defaultKubeConfig(),
		"Provide the path to the `kubeconfig` file you'd like to use for these tests. The `current-context` will be used.")

	// Allows the test to use a different cluster domain - ie. if it's not cluster.local
	fs.StringVar(&s.DomainName, "env.cluster.domain", "cluster.local",
		"Provide the path to the `kubeconfig` file you'd like to use for these tests. The `current-context` will be used.")
}

func defaultKubeConfig() string {
	// Use KUBECONFIG if available
	kubeconfig := os.Getenv("KUBECONFIG")

	// If KUBECONFIG env var isn't set then look for $HOME/.kube/config
	if kubeconfig == "" {
		if usr, err := user.Current(); err == nil {
			kubeconfig = path.Join(usr.HomeDir, ".kube/config")
		}
	}
	return kubeconfig
}

func (c *Cluster) ClientConfig() *rest.Config {
	overrides := &clientcmd.ConfigOverrides{}
	overrides.Context.Cluster = c.Name

	loader := &clientcmd.ClientConfigLoadingRules{ExplicitPath: c.KubeConfig}

	conf, err := clientcmd.NewNonInteractiveDeferredLoadingClientConfig(loader, overrides).ClientConfig()
	if err != nil {
		panic(err)
	}
	return conf
}
<|endoftext|>