<commit_before>// Copyright (c) 2013 Erik St. Martin, Brian Ketelsen. All rights reserved.
// Use of this source code is governed by The MIT License (MIT) that can be
// found in the LICENSE file.

package msg

import (
	"log"
	"net/http"
	"strconv"
	"time"
)

type Service struct {
	UUID        string
	Name        string
	Version     string
	Environment string
	Region      string
	Host        string
	Port        uint16
	TTL         uint32 // Seconds
	Expires     time.Time
	Callback    map[string]Callback `json:"-"` // Callbacks are found by UUID
}

// Returns the amount of time remaining before expiration
func (s *Service) RemainingTTL() uint32 {
	d := s.Expires.Sub(time.Now())
	ttl := uint32(d.Seconds())

	if ttl < 1 {
		return 0
	}
	return ttl
}

// Updates TTL property to the RemainingTTL
func (s *Service) UpdateTTL() {
	s.TTL = s.RemainingTTL()
}

type Callback struct {
	UUID string

	// Name of the service
	Name        string
	Version     string
	Environment string
	Region      string
	Host        string

	Reply string
	Port  uint16
}

// Call calls the callback and performs the HTTP request.
func (c Callback) Call(s Service) {
	go func() {
		log.Println("Performing callback to:", c.Reply, c.Port)
		// TODO(miek): Use DELETE with some data
		resp, err := http.Get("http://" + c.Reply + ":" + strconv.Itoa(int(c.Port)) + "/skydns/callbacks/" + c.UUID)
		if err != nil {
			return
		}
		resp.Body.Close()
		return
	}()
	return
}
<commit_msg>dont run it in a go routine<commit_after>// Copyright (c) 2013 Erik St. Martin, Brian Ketelsen. All rights reserved.
// Use of this source code is governed by The MIT License (MIT) that can be
// found in the LICENSE file.

package msg

import (
	"log"
	"net/http"
	"strconv"
	"time"
)

type Service struct {
	UUID        string
	Name        string
	Version     string
	Environment string
	Region      string
	Host        string
	Port        uint16
	TTL         uint32 // Seconds
	Expires     time.Time
	Callback    map[string]Callback `json:"-"` // Callbacks are found by UUID
}

// Returns the amount of time remaining before expiration
func (s *Service) RemainingTTL() uint32 {
	d := s.Expires.Sub(time.Now())
	ttl := uint32(d.Seconds())

	if ttl < 1 {
		return 0
	}
	return ttl
}

// Updates TTL property to the RemainingTTL
func (s *Service) UpdateTTL() {
	s.TTL = s.RemainingTTL()
}

type Callback struct {
	UUID string

	// Name of the service
	Name        string
	Version     string
	Environment string
	Region      string
	Host        string

	Reply string
	Port  uint16
}

// Call calls the callback and performs the HTTP request.
func (c Callback) Call(s Service) {
		log.Println("Performing callback to:", c.Reply, c.Port)
		// TODO(miek): Use DELETE with some data
		resp, err := http.Get("http://" + c.Reply + ":" + strconv.Itoa(int(c.Port)) + "/skydns/callbacks/" + c.UUID)
		if err != nil {
			return
		}
		resp.Body.Close()
		return
}
<|endoftext|>


<commit_before>package multiconfig

import (
	"errors"
	"fmt"
	"reflect"
	"strconv"
	"strings"

	"github.com/fatih/structs"
)

// Loader loads the configuration from a source. The implementer of Loader is
// responsible of setting the default values of the struct.
type Loader interface {
	// Load loads the source into the config defined by struct s
	Load(s interface{}) error
}

// DefaultLoader implements the Loader interface. It initializes the given
// pointer of struct s with configuration from the default sources. The order
// of load is LoadFile, LoadEnv and lastly LoadFlag.  An error in any step
// stops the loading process. Each step overrides the previous step's config
// (i.e: defining a flag will override previous environment or file config). To
// customize the order use the individual load functions.
type DefaultLoader struct {
	Loader
}

// NewWithPath returns a new instance of Loader to read from the given
// configuration file.
func NewWithPath(path string) *DefaultLoader {
	loaders := []Loader{}

	// Choose what while is passed
	if strings.HasSuffix(path, "toml") {
		loaders = append(loaders, &TOMLLoader{Path: path})
	}

	if strings.HasSuffix(path, "json") {
		loaders = append(loaders, &JSONLoader{Path: path})
	}

	loaders = append(loaders, &EnvironmentLoader{}, &FlagLoader{})
	loader := MultiLoader(loaders...)

	d := &DefaultLoader{}
	d.Loader = loader
	return d
}

// New returns a new instance of DefaultLoader without any file loaders.
func New() *DefaultLoader {
	loader := MultiLoader(
		&EnvironmentLoader{},
		&FlagLoader{},
	)

	d := &DefaultLoader{}
	d.Loader = loader
	return d
}

// MustLoad is like Load but panics if the config cannot be parsed.
func (d *DefaultLoader) MustLoad(conf interface{}) {
	if err := d.Load(conf); err != nil {
		panic(err)
	}
}

// setDefaults parses the struct and reads each field's tag
func setDefaults(s interface{}) error {
	for _, field := range structs.Fields(s) {
		defaultVal := field.Tag("default")
		if defaultVal == "" {
			continue
		}

		err := fieldSet(field, defaultVal)
		if err != nil {
			return err
		}
	}

	return nil
}

// fieldSet sets field value from the given string value. It converts the
// string value in a sane way and is usefulf or environment variables or flags
// which are by nature in string types.
func fieldSet(field *structs.Field, v string) error {
	// TODO: add support for other types
	switch field.Kind() {
	case reflect.Bool:
		val, err := strconv.ParseBool(v)
		if err != nil {
			return err
		}

		if err := field.Set(val); err != nil {
			return err
		}
	case reflect.Int:
		i, err := strconv.Atoi(v)
		if err != nil {
			return err
		}

		if err := field.Set(i); err != nil {
			return err
		}
	case reflect.String:
		field.Set(v)
	case reflect.Slice:
		// TODO add other typed slice support
		if _, ok := field.Value().([]string); !ok {
			return errors.New("can't set on non string slices")
		}

		if err := field.Set(strings.Split(v, ",")); err != nil {
			return err
		}
	case reflect.Float64:
		f, err := strconv.ParseFloat(v, 64)
		if err != nil {
			return err
		}

		if err := field.Set(f); err != nil {
			return err
		}
	default:
		return fmt.Errorf("multiconfig: not supported type: %s", field.Kind())
	}

	return nil
}
<commit_msg>multiconfig: make tag value customizable<commit_after>package multiconfig

import (
	"errors"
	"fmt"
	"reflect"
	"strconv"
	"strings"

	"github.com/fatih/structs"
)

var (
	DefaultDefaultTag = "default"
)

// Loader loads the configuration from a source. The implementer of Loader is
// responsible of setting the default values of the struct.
type Loader interface {
	// Load loads the source into the config defined by struct s
	Load(s interface{}) error
}

// DefaultLoader implements the Loader interface. It initializes the given
// pointer of struct s with configuration from the default sources. The order
// of load is LoadFile, LoadEnv and lastly LoadFlag.  An error in any step
// stops the loading process. Each step overrides the previous step's config
// (i.e: defining a flag will override previous environment or file config). To
// customize the order use the individual load functions.
type DefaultLoader struct {
	Loader
}

// NewWithPath returns a new instance of Loader to read from the given
// configuration file.
func NewWithPath(path string) *DefaultLoader {
	loaders := []Loader{}

	// Choose what while is passed
	if strings.HasSuffix(path, "toml") {
		loaders = append(loaders, &TOMLLoader{Path: path})
	}

	if strings.HasSuffix(path, "json") {
		loaders = append(loaders, &JSONLoader{Path: path})
	}

	loaders = append(loaders, &EnvironmentLoader{}, &FlagLoader{})
	loader := MultiLoader(loaders...)

	d := &DefaultLoader{}
	d.Loader = loader
	return d
}

// New returns a new instance of DefaultLoader without any file loaders.
func New() *DefaultLoader {
	loader := MultiLoader(
		&EnvironmentLoader{},
		&FlagLoader{},
	)

	d := &DefaultLoader{}
	d.Loader = loader
	return d
}

// MustLoad is like Load but panics if the config cannot be parsed.
func (d *DefaultLoader) MustLoad(conf interface{}) {
	if err := d.Load(conf); err != nil {
		panic(err)
	}
}

// setDefaults parses the struct and reads each field's tag
func setDefaults(s interface{}) error {
	for _, field := range structs.Fields(s) {
		defaultVal := field.Tag(DefaultDefaultTag)
		if defaultVal == "" {
			continue
		}

		err := fieldSet(field, defaultVal)
		if err != nil {
			return err
		}
	}

	return nil
}

// fieldSet sets field value from the given string value. It converts the
// string value in a sane way and is usefulf or environment variables or flags
// which are by nature in string types.
func fieldSet(field *structs.Field, v string) error {
	// TODO: add support for other types
	switch field.Kind() {
	case reflect.Bool:
		val, err := strconv.ParseBool(v)
		if err != nil {
			return err
		}

		if err := field.Set(val); err != nil {
			return err
		}
	case reflect.Int:
		i, err := strconv.Atoi(v)
		if err != nil {
			return err
		}

		if err := field.Set(i); err != nil {
			return err
		}
	case reflect.String:
		field.Set(v)
	case reflect.Slice:
		// TODO add other typed slice support
		if _, ok := field.Value().([]string); !ok {
			return errors.New("can't set on non string slices")
		}

		if err := field.Set(strings.Split(v, ",")); err != nil {
			return err
		}
	case reflect.Float64:
		f, err := strconv.ParseFloat(v, 64)
		if err != nil {
			return err
		}

		if err := field.Set(f); err != nil {
			return err
		}
	default:
		return fmt.Errorf("multiconfig: not supported type: %s", field.Kind())
	}

	return nil
}
<|endoftext|>


<commit_before>package pretty

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"strings"
	"text/tabwriter"
	"text/template"

	"github.com/fatih/color"
	ppsclient "github.com/pachyderm/pachyderm/src/client/pps"
	"github.com/pachyderm/pachyderm/src/server/pkg/pretty"
)

// PrintJobHeader prints a job header.
func PrintJobHeader(w io.Writer) {
	fmt.Fprint(w, "ID\tOUTPUT\tSTARTED\tDURATION\tSTATE\t\n")
}

// PrintJobInfo pretty-prints job info.
func PrintJobInfo(w io.Writer, jobInfo *ppsclient.JobInfo) {
	fmt.Fprintf(w, "%s\t", jobInfo.Job.ID)
	if jobInfo.OutputCommit != nil {
		fmt.Fprintf(w, "%s/%s\t", jobInfo.OutputCommit.Repo.Name, jobInfo.OutputCommit.ID)
	} else {
		fmt.Fprintf(w, "-\t")
	}
	fmt.Fprintf(w, "%s\t", pretty.Ago(jobInfo.Started))
	if jobInfo.Finished != nil {
		fmt.Fprintf(w, "%s\t", pretty.Duration(jobInfo.Started, jobInfo.Finished))
	} else {
		fmt.Fprintf(w, "-\t")
	}
	fmt.Fprintf(w, "%s\t\n", jobState(jobInfo.State))
}

// PrintPipelineHeader prints a pipeline header.
func PrintPipelineHeader(w io.Writer) {
	fmt.Fprint(w, "NAME\tINPUT\tOUTPUT\tSTATE\t\n")
}

// PrintPipelineInfo pretty-prints pipeline info.
func PrintPipelineInfo(w io.Writer, pipelineInfo *ppsclient.PipelineInfo) {
	fmt.Fprintf(w, "%s\t", pipelineInfo.Pipeline.Name)
	if len(pipelineInfo.Inputs) == 0 {
		fmt.Fprintf(w, "\t")
	} else {
		for i, input := range pipelineInfo.Inputs {
			fmt.Fprintf(w, "%s", input.Repo.Name)
			if i == len(pipelineInfo.Inputs)-1 {
				fmt.Fprintf(w, "\t")
			} else {
				fmt.Fprintf(w, ", ")
			}
		}
	}
	if pipelineInfo.OutputRepo != nil {
		fmt.Fprintf(w, "%s\t", pipelineInfo.OutputRepo.Name)
	} else {
		fmt.Fprintf(w, "\t")
	}
	fmt.Fprintf(w, "%s\t\n", pipelineState(pipelineInfo.State))
}

// PrintJobInputHeader pretty prints a job input header.
func PrintJobInputHeader(w io.Writer) {
	fmt.Fprint(w, "NAME\tCOMMIT\tPARTITION\tINCREMENTAL\tLAZY\t\n")
}

// PrintJobInput pretty-prints a job input.
func PrintJobInput(w io.Writer, jobInput *ppsclient.JobInput) {
	fmt.Fprintf(w, "%s\t", jobInput.Commit.Repo.Name)
	fmt.Fprintf(w, "%s\t", jobInput.Commit.ID)
	fmt.Fprintf(w, "%s\t", jobInput.Method.Partition)
	fmt.Fprintf(w, "%s\t", jobInput.Method.Incremental)
	fmt.Fprintf(w, "%t\t\n", jobInput.Lazy)
}

// PrintPipelineInputHeader prints a pipeline input header.
func PrintPipelineInputHeader(w io.Writer) {
	fmt.Fprint(w, "NAME\tPARTITION\tINCREMENTAL\tLAZY\t\n")
}

// PrintPipelineInput pretty-prints a pipeline input.
func PrintPipelineInput(w io.Writer, pipelineInput *ppsclient.PipelineInput) {
	fmt.Fprintf(w, "%s\t", pipelineInput.Repo.Name)
	fmt.Fprintf(w, "%s\t", pipelineInput.Method.Partition)
	fmt.Fprintf(w, "%s\t", pipelineInput.Method.Incremental)
	fmt.Fprintf(w, "%t\t\n", pipelineInput.Lazy)
}

// PrintJobCountsHeader prints a job counts header.
func PrintJobCountsHeader(w io.Writer) {
	fmt.Fprintf(w, strings.ToUpper(jobState(ppsclient.JobState_JOB_CREATING))+"\t")
	fmt.Fprintf(w, strings.ToUpper(jobState(ppsclient.JobState_JOB_RUNNING))+"\t")
	fmt.Fprintf(w, strings.ToUpper(jobState(ppsclient.JobState_JOB_FAILURE))+"\t")
	fmt.Fprintf(w, strings.ToUpper(jobState(ppsclient.JobState_JOB_SUCCESS))+"\t\n")
}

// PrintDetailedJobInfo pretty-prints detailed job info.
func PrintDetailedJobInfo(jobInfo *ppsclient.JobInfo) error {
	template, err := template.New("JobInfo").Funcs(funcMap).Parse(
		`ID: {{.Job.ID}} {{if .ParentJob}}
Parent: {{.ParentJob.ID}} {{end}}
Started: {{prettyAgo .Started}} {{if .Finished}}
Duration: {{prettyDuration .Started .Finished}} {{end}}
State: {{jobState .State}}
ParallelismSpec: {{.ParallelismSpec}}
{{ if .Service }}Service:
	{{ if .Service.InternalPort }}InternalPort: {{ .Service.InternalPort }} {{end}}
	{{ if .Service.ExternalPort }}ExternalPort: {{ .Service.ExternalPort }} {{end}} {{end}}
Inputs:
{{jobInputs .}}Transform:
{{prettyTransform .Transform}}
Output: {{.OutputCommit.ID}}
Chunks:
{{prettyChunks .Chunks}}
`)
	if err != nil {
		return err
	}
	err = template.Execute(os.Stdout, jobInfo)
	if err != nil {
		return err
	}
	return nil
}

// PrintDetailedPipelineInfo pretty-prints detailed pipeline info.
func PrintDetailedPipelineInfo(pipelineInfo *ppsclient.PipelineInfo) error {
	template, err := template.New("PipelineInfo").Funcs(funcMap).Parse(
		`Name: {{.Pipeline.Name}}
Created: {{prettyAgo .CreatedAt}}
State: {{pipelineState .State}}
ParallelismSpec: {{.ParallelismSpec}}
Inputs:
{{pipelineInputs .}}Transform:
{{prettyTransform .Transform}}
{{if .RecentError}} Recent Error: {{.RecentError}} {{end}}
Job Counts:
{{jobCounts .JobCounts}}
`)
	if err != nil {
		return err
	}
	err = template.Execute(os.Stdout, pipelineInfo)
	if err != nil {
		return err
	}
	return nil
}

func podState(podState ppsclient.PodState) string {
	switch podState {
	case ppsclient.PodState_POD_RUNNING:
		return color.New(color.FgYellow).SprintFunc()("running")
	case ppsclient.PodState_POD_SUCCESS:
		return color.New(color.FgGreen).SprintFunc()("success")
	case ppsclient.PodState_POD_FAILED:
		return color.New(color.FgRed).SprintFunc()("failure")
	}
	return "-"
}

func chunkState(chunkState ppsclient.ChunkState) string {
	switch chunkState {
	case ppsclient.ChunkState_CHUNK_UNASSIGNED:
		return color.New(color.FgYellow).SprintFunc()("unassigned")
	case ppsclient.ChunkState_CHUNK_ASSIGNED:
		return color.New(color.FgYellow).SprintFunc()("assigned")
	case ppsclient.ChunkState_CHUNK_SUCCESS:
		return color.New(color.FgGreen).SprintFunc()("success")
	case ppsclient.ChunkState_CHUNK_FAILURE:
		return color.New(color.FgRed).SprintFunc()("failure")
	}
	return "-"
}

func jobState(jobState ppsclient.JobState) string {
	switch jobState {
	case ppsclient.JobState_JOB_CREATING:
		return color.New(color.FgYellow).SprintFunc()("pulling")
	case ppsclient.JobState_JOB_RUNNING:
		return color.New(color.FgYellow).SprintFunc()("running")
	case ppsclient.JobState_JOB_FAILURE:
		return color.New(color.FgRed).SprintFunc()("failure")
	case ppsclient.JobState_JOB_SUCCESS:
		return color.New(color.FgGreen).SprintFunc()("success")
	case ppsclient.JobState_JOB_EMPTY:
		return color.New(color.FgGreen).SprintFunc()("empty")
	}
	return "-"
}

func pipelineState(pipelineState ppsclient.PipelineState) string {
	switch pipelineState {
	case ppsclient.PipelineState_PIPELINE_IDLE:
		return color.New(color.FgYellow).SprintFunc()("idle")
	case ppsclient.PipelineState_PIPELINE_RUNNING:
		return color.New(color.FgGreen).SprintFunc()("running")
	case ppsclient.PipelineState_PIPELINE_RESTARTING:
		return color.New(color.FgYellow).SprintFunc()("restarting")
	case ppsclient.PipelineState_PIPELINE_FAILURE:
		return color.New(color.FgRed).SprintFunc()("failure")
	case ppsclient.PipelineState_PIPELINE_STOPPED:
		return color.New(color.FgYellow).SprintFunc()("stopped")
	}
	return "-"
}

func jobInputs(jobInfo *ppsclient.JobInfo) string {
	var buffer bytes.Buffer
	writer := tabwriter.NewWriter(&buffer, 20, 1, 3, ' ', 0)
	PrintJobInputHeader(writer)
	for _, input := range jobInfo.Inputs {
		PrintJobInput(writer, input)
	}
	// can't error because buffer can't error on Write
	writer.Flush()
	return buffer.String()
}

func pipelineInputs(pipelineInfo *ppsclient.PipelineInfo) string {
	var buffer bytes.Buffer
	writer := tabwriter.NewWriter(&buffer, 20, 1, 3, ' ', 0)
	PrintPipelineInputHeader(writer)
	for _, input := range pipelineInfo.Inputs {
		PrintPipelineInput(writer, input)
	}
	// can't error because buffer can't error on Write
	writer.Flush()
	return buffer.String()
}

func jobCounts(counts map[int32]int32) string {
	var buffer bytes.Buffer
	for i := int32(ppsclient.JobState_JOB_CREATING); i <= int32(ppsclient.JobState_JOB_SUCCESS); i++ {
		fmt.Fprintf(&buffer, "%s: %d\t", jobState(ppsclient.JobState(i)), counts[i])
	}
	return buffer.String()
}

func prettyTransform(transform *ppsclient.Transform) (string, error) {
	result, err := json.MarshalIndent(transform, "", "  ")
	if err != nil {
		return "", err
	}
	return pretty.UnescapeHTML(string(result)), nil
}

func prettyChunks(chunks []*ppsclient.Chunk) (string, error) {
	var buffer bytes.Buffer
	for i, chunk := range chunks {
		fmt.Fprintf(&buffer, "\nChunk %d: %s\n", i+1, chunkState(chunk.State))
		writer := tabwriter.NewWriter(&buffer, 20, 1, 3, ' ', 0)
		fmt.Fprintf(writer, "Pod Name\tOutput Commit\tState\t\n")
		for _, pod := range chunk.Pods {
			fmt.Fprintf(writer, "%s\t%s\t%s\t\n", pod.Name, pod.OutputCommit.ID, podState(pod.State))
		}
		if err := writer.Flush(); err != nil {
			return "", err
		}
	}
	fmt.Fprintln(&buffer)
	return buffer.String(), nil
}

var funcMap = template.FuncMap{
	"pipelineState":   pipelineState,
	"jobState":        jobState,
	"pipelineInputs":  pipelineInputs,
	"jobInputs":       jobInputs,
	"prettyAgo":       pretty.Ago,
	"prettyDuration":  pretty.Duration,
	"jobCounts":       jobCounts,
	"prettyTransform": prettyTransform,
	"prettyChunks":    prettyChunks,
}
<commit_msg>Adds outputs to job/pipeline pretty printing.<commit_after>package pretty

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"strings"
	"text/tabwriter"
	"text/template"

	"github.com/fatih/color"
	ppsclient "github.com/pachyderm/pachyderm/src/client/pps"
	"github.com/pachyderm/pachyderm/src/server/pkg/pretty"
)

// PrintJobHeader prints a job header.
func PrintJobHeader(w io.Writer) {
	fmt.Fprint(w, "ID\tOUTPUT COMMIT\tSTARTED\tDURATION\tSTATE\t\n")
}

// PrintJobInfo pretty-prints job info.
func PrintJobInfo(w io.Writer, jobInfo *ppsclient.JobInfo) {
	fmt.Fprintf(w, "%s\t", jobInfo.Job.ID)
	if jobInfo.OutputCommit != nil {
		fmt.Fprintf(w, "%s/%s\t", jobInfo.OutputCommit.Repo.Name, jobInfo.OutputCommit.ID)
	} else {
		fmt.Fprintf(w, "-\t")
	}
	fmt.Fprintf(w, "%s\t", pretty.Ago(jobInfo.Started))
	if jobInfo.Finished != nil {
		fmt.Fprintf(w, "%s\t", pretty.Duration(jobInfo.Started, jobInfo.Finished))
	} else {
		fmt.Fprintf(w, "-\t")
	}
	fmt.Fprintf(w, "%s\t\n", jobState(jobInfo.State))
}

// PrintPipelineHeader prints a pipeline header.
func PrintPipelineHeader(w io.Writer) {
	fmt.Fprint(w, "NAME\tINPUT\tOUTPUT REPO\tSTATE\t\n")
}

// PrintPipelineInfo pretty-prints pipeline info.
func PrintPipelineInfo(w io.Writer, pipelineInfo *ppsclient.PipelineInfo) {
	fmt.Fprintf(w, "%s\t", pipelineInfo.Pipeline.Name)
	if len(pipelineInfo.Inputs) == 0 {
		fmt.Fprintf(w, "\t")
	} else {
		for i, input := range pipelineInfo.Inputs {
			fmt.Fprintf(w, "%s", input.Repo.Name)
			if i == len(pipelineInfo.Inputs)-1 {
				fmt.Fprintf(w, "\t")
			} else {
				fmt.Fprintf(w, ", ")
			}
		}
	}
	if pipelineInfo.OutputRepo != nil {
		fmt.Fprintf(w, "%s\t", pipelineInfo.OutputRepo.Name)
	} else {
		fmt.Fprintf(w, "\t")
	}
	fmt.Fprintf(w, "%s\t\n", pipelineState(pipelineInfo.State))
}

// PrintJobInputHeader pretty prints a job input header.
func PrintJobInputHeader(w io.Writer) {
	fmt.Fprint(w, "NAME\tCOMMIT\tPARTITION\tINCREMENTAL\tLAZY\t\n")
}

// PrintJobInput pretty-prints a job input.
func PrintJobInput(w io.Writer, jobInput *ppsclient.JobInput) {
	fmt.Fprintf(w, "%s\t", jobInput.Commit.Repo.Name)
	fmt.Fprintf(w, "%s\t", jobInput.Commit.ID)
	fmt.Fprintf(w, "%s\t", jobInput.Method.Partition)
	fmt.Fprintf(w, "%s\t", jobInput.Method.Incremental)
	fmt.Fprintf(w, "%t\t\n", jobInput.Lazy)
}

// PrintPipelineInputHeader prints a pipeline input header.
func PrintPipelineInputHeader(w io.Writer) {
	fmt.Fprint(w, "NAME\tPARTITION\tINCREMENTAL\tLAZY\t\n")
}

// PrintPipelineInput pretty-prints a pipeline input.
func PrintPipelineInput(w io.Writer, pipelineInput *ppsclient.PipelineInput) {
	fmt.Fprintf(w, "%s\t", pipelineInput.Repo.Name)
	fmt.Fprintf(w, "%s\t", pipelineInput.Method.Partition)
	fmt.Fprintf(w, "%s\t", pipelineInput.Method.Incremental)
	fmt.Fprintf(w, "%t\t\n", pipelineInput.Lazy)
}

// PrintJobCountsHeader prints a job counts header.
func PrintJobCountsHeader(w io.Writer) {
	fmt.Fprintf(w, strings.ToUpper(jobState(ppsclient.JobState_JOB_CREATING))+"\t")
	fmt.Fprintf(w, strings.ToUpper(jobState(ppsclient.JobState_JOB_RUNNING))+"\t")
	fmt.Fprintf(w, strings.ToUpper(jobState(ppsclient.JobState_JOB_FAILURE))+"\t")
	fmt.Fprintf(w, strings.ToUpper(jobState(ppsclient.JobState_JOB_SUCCESS))+"\t\n")
}

// PrintDetailedJobInfo pretty-prints detailed job info.
func PrintDetailedJobInfo(jobInfo *ppsclient.JobInfo) error {
	template, err := template.New("JobInfo").Funcs(funcMap).Parse(
		`ID: {{.Job.ID}} {{if .ParentJob}}
Parent: {{.ParentJob.ID}} {{end}}
Started: {{prettyAgo .Started}} {{if .Finished}}
Duration: {{prettyDuration .Started .Finished}} {{end}}
State: {{jobState .State}}
ParallelismSpec: {{.ParallelismSpec}}
{{ if .Service }}Service:
	{{ if .Service.InternalPort }}InternalPort: {{ .Service.InternalPort }} {{end}}
	{{ if .Service.ExternalPort }}ExternalPort: {{ .Service.ExternalPort }} {{end}} {{end}}
Inputs:
{{jobInputs .}}Transform:
{{prettyTransform .Transform}}
Output Commit: {{.OutputCommit.ID}}
Output: {{.Output.URL}}
Chunks:
{{prettyChunks .Chunks}}
`)
	if err != nil {
		return err
	}
	err = template.Execute(os.Stdout, jobInfo)
	if err != nil {
		return err
	}
	return nil
}

// PrintDetailedPipelineInfo pretty-prints detailed pipeline info.
func PrintDetailedPipelineInfo(pipelineInfo *ppsclient.PipelineInfo) error {
	template, err := template.New("PipelineInfo").Funcs(funcMap).Parse(
		`Name: {{.Pipeline.Name}}
Created: {{prettyAgo .CreatedAt}}
State: {{pipelineState .State}}
ParallelismSpec: {{.ParallelismSpec}}
Inputs:
{{pipelineInputs .}}Transform:
{{prettyTransform .Transform}}
Output: {{.Output.URL}}
{{if .RecentError}} Recent Error: {{.RecentError}} {{end}}
Job Counts:
{{jobCounts .JobCounts}}
`)
	if err != nil {
		return err
	}
	err = template.Execute(os.Stdout, pipelineInfo)
	if err != nil {
		return err
	}
	return nil
}

func podState(podState ppsclient.PodState) string {
	switch podState {
	case ppsclient.PodState_POD_RUNNING:
		return color.New(color.FgYellow).SprintFunc()("running")
	case ppsclient.PodState_POD_SUCCESS:
		return color.New(color.FgGreen).SprintFunc()("success")
	case ppsclient.PodState_POD_FAILED:
		return color.New(color.FgRed).SprintFunc()("failure")
	}
	return "-"
}

func chunkState(chunkState ppsclient.ChunkState) string {
	switch chunkState {
	case ppsclient.ChunkState_CHUNK_UNASSIGNED:
		return color.New(color.FgYellow).SprintFunc()("unassigned")
	case ppsclient.ChunkState_CHUNK_ASSIGNED:
		return color.New(color.FgYellow).SprintFunc()("assigned")
	case ppsclient.ChunkState_CHUNK_SUCCESS:
		return color.New(color.FgGreen).SprintFunc()("success")
	case ppsclient.ChunkState_CHUNK_FAILURE:
		return color.New(color.FgRed).SprintFunc()("failure")
	}
	return "-"
}

func jobState(jobState ppsclient.JobState) string {
	switch jobState {
	case ppsclient.JobState_JOB_CREATING:
		return color.New(color.FgYellow).SprintFunc()("pulling")
	case ppsclient.JobState_JOB_RUNNING:
		return color.New(color.FgYellow).SprintFunc()("running")
	case ppsclient.JobState_JOB_FAILURE:
		return color.New(color.FgRed).SprintFunc()("failure")
	case ppsclient.JobState_JOB_SUCCESS:
		return color.New(color.FgGreen).SprintFunc()("success")
	case ppsclient.JobState_JOB_EMPTY:
		return color.New(color.FgGreen).SprintFunc()("empty")
	}
	return "-"
}

func pipelineState(pipelineState ppsclient.PipelineState) string {
	switch pipelineState {
	case ppsclient.PipelineState_PIPELINE_IDLE:
		return color.New(color.FgYellow).SprintFunc()("idle")
	case ppsclient.PipelineState_PIPELINE_RUNNING:
		return color.New(color.FgGreen).SprintFunc()("running")
	case ppsclient.PipelineState_PIPELINE_RESTARTING:
		return color.New(color.FgYellow).SprintFunc()("restarting")
	case ppsclient.PipelineState_PIPELINE_FAILURE:
		return color.New(color.FgRed).SprintFunc()("failure")
	case ppsclient.PipelineState_PIPELINE_STOPPED:
		return color.New(color.FgYellow).SprintFunc()("stopped")
	}
	return "-"
}

func jobInputs(jobInfo *ppsclient.JobInfo) string {
	var buffer bytes.Buffer
	writer := tabwriter.NewWriter(&buffer, 20, 1, 3, ' ', 0)
	PrintJobInputHeader(writer)
	for _, input := range jobInfo.Inputs {
		PrintJobInput(writer, input)
	}
	// can't error because buffer can't error on Write
	writer.Flush()
	return buffer.String()
}

func pipelineInputs(pipelineInfo *ppsclient.PipelineInfo) string {
	var buffer bytes.Buffer
	writer := tabwriter.NewWriter(&buffer, 20, 1, 3, ' ', 0)
	PrintPipelineInputHeader(writer)
	for _, input := range pipelineInfo.Inputs {
		PrintPipelineInput(writer, input)
	}
	// can't error because buffer can't error on Write
	writer.Flush()
	return buffer.String()
}

func jobCounts(counts map[int32]int32) string {
	var buffer bytes.Buffer
	for i := int32(ppsclient.JobState_JOB_CREATING); i <= int32(ppsclient.JobState_JOB_SUCCESS); i++ {
		fmt.Fprintf(&buffer, "%s: %d\t", jobState(ppsclient.JobState(i)), counts[i])
	}
	return buffer.String()
}

func prettyTransform(transform *ppsclient.Transform) (string, error) {
	result, err := json.MarshalIndent(transform, "", "  ")
	if err != nil {
		return "", err
	}
	return pretty.UnescapeHTML(string(result)), nil
}

func prettyChunks(chunks []*ppsclient.Chunk) (string, error) {
	var buffer bytes.Buffer
	for i, chunk := range chunks {
		fmt.Fprintf(&buffer, "\nChunk %d: %s\n", i+1, chunkState(chunk.State))
		writer := tabwriter.NewWriter(&buffer, 20, 1, 3, ' ', 0)
		fmt.Fprintf(writer, "Pod Name\tOutput Commit\tState\t\n")
		for _, pod := range chunk.Pods {
			fmt.Fprintf(writer, "%s\t%s\t%s\t\n", pod.Name, pod.OutputCommit.ID, podState(pod.State))
		}
		if err := writer.Flush(); err != nil {
			return "", err
		}
	}
	fmt.Fprintln(&buffer)
	return buffer.String(), nil
}

var funcMap = template.FuncMap{
	"pipelineState":   pipelineState,
	"jobState":        jobState,
	"pipelineInputs":  pipelineInputs,
	"jobInputs":       jobInputs,
	"prettyAgo":       pretty.Ago,
	"prettyDuration":  pretty.Duration,
	"jobCounts":       jobCounts,
	"prettyTransform": prettyTransform,
	"prettyChunks":    prettyChunks,
}
<|endoftext|>


<commit_before>package server

import (
	"context"
	"fmt"
	"path"
	"time"

	log "github.com/sirupsen/logrus"
	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	kube_watch "k8s.io/apimachinery/pkg/watch"
	kube "k8s.io/client-go/kubernetes"

	"github.com/pachyderm/pachyderm/src/client"
	"github.com/pachyderm/pachyderm/src/client/pkg/uuid"
	"github.com/pachyderm/pachyderm/src/client/pps"
	"github.com/pachyderm/pachyderm/src/server/pkg/backoff"
	col "github.com/pachyderm/pachyderm/src/server/pkg/collection"
	"github.com/pachyderm/pachyderm/src/server/pkg/deploy/assets"
	"github.com/pachyderm/pachyderm/src/server/pkg/dlock"
	"github.com/pachyderm/pachyderm/src/server/pkg/util"
	"github.com/pachyderm/pachyderm/src/server/pkg/watch"
	ppsserver "github.com/pachyderm/pachyderm/src/server/pps"
)

const (
	masterLockPath = "_master_lock"
)

var (
	failures = map[string]bool{
		"InvalidImageName": true,
		"ErrImagePull":     true,
	}
)

// The master process is responsible for creating/deleting workers as
// pipelines are created/removed.
func (a *apiServer) master() {
	masterLock := dlock.NewDLock(a.etcdClient, path.Join(a.etcdPrefix, masterLockPath))
	backoff.RetryNotify(func() error {
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		ctx, err := masterLock.Lock(ctx)
		if err != nil {
			return err
		}
		defer masterLock.Unlock(ctx)

		log.Infof("Launching PPS master process")

		pipelineWatcher, err := a.pipelines.ReadOnly(ctx).WatchWithPrev()
		if err != nil {
			return fmt.Errorf("error creating watch: %+v", err)
		}
		defer pipelineWatcher.Close()

		// watchChan will be nil if the Watch call below errors, this means
		// that we won't receive events from k8s and won't be able to detect
		// errors in pods. We could just return that error and retry but that
		// prevents pachyderm from creating pipelines when there's an issue
		// talking to k8s.
		var watchChan <-chan kube_watch.Event
		kubePipelineWatch, err := a.kubeClient.CoreV1().Pods(a.namespace).Watch(
			metav1.ListOptions{
				LabelSelector: metav1.FormatLabelSelector(metav1.SetAsLabelSelector(
					map[string]string{
						"component": "worker",
					})),
				Watch: true,
			})
		if err != nil {
			log.Errorf("failed to watch kuburnetes pods: %v", err)
		} else {
			watchChan = kubePipelineWatch.ResultChan()
			defer kubePipelineWatch.Stop()
		}

		for {
			select {
			case event := <-pipelineWatcher.Watch():
				if event.Err != nil {
					return fmt.Errorf("event err: %+v", event.Err)
				}
				switch event.Type {
				case watch.EventPut:
					var pipelineName string
					var pipelineInfo pps.PipelineInfo
					if err := event.Unmarshal(&pipelineName, &pipelineInfo); err != nil {
						return err
					}

					// This is a bit of a hack that covers for migrations bugs
					// where a field is set by setPipelineDefaults in a newer
					// version of the server but a PipelineInfo which was
					// created with an older version of the server doesn't have
					// that field set because setPipelineDefaults was different
					// when it was created.
					if pipelineInfo.Salt == "" || pipelineInfo.CacheSize == "" {
						if _, err := col.NewSTM(ctx, a.etcdClient, func(stm col.STM) error {
							pipelines := a.pipelines.ReadWrite(stm)
							newPipelineInfo := new(pps.PipelineInfo)
							if err := pipelines.Get(pipelineInfo.Pipeline.Name, newPipelineInfo); err != nil {
								return fmt.Errorf("error getting pipeline %s: %+v", pipelineName, err)
							}
							if newPipelineInfo.Salt == "" {
								newPipelineInfo.Salt = uuid.NewWithoutDashes()
							}
							setPipelineDefaults(newPipelineInfo)
							pipelines.Put(pipelineInfo.Pipeline.Name, newPipelineInfo)
							pipelineInfo = *newPipelineInfo
							return nil
						}); err != nil {
							return err
						}
					}

					var prevPipelineInfo pps.PipelineInfo
					if event.PrevKey != nil {
						if err := event.UnmarshalPrev(&pipelineName, &prevPipelineInfo); err != nil {
							return err
						}
					}

					// If the pipeline has been stopped, delete workers
					if pipelineStateToStopped(pipelineInfo.State) {
						log.Infof("master: deleting workers for pipeline %s", pipelineInfo.Pipeline.Name)
						if err := a.deleteWorkersForPipeline(&pipelineInfo); err != nil {
							return err
						}
					}

					var hasGitInput bool
					pps.VisitInput(pipelineInfo.Input, func(input *pps.Input) {
						if input.Git != nil {
							hasGitInput = true
						}
					})

					// If the pipeline has been restarted, create workers
					if !pipelineStateToStopped(pipelineInfo.State) && event.PrevKey != nil && pipelineStateToStopped(prevPipelineInfo.State) {
						if hasGitInput {
							if err := a.checkOrDeployGithookService(); err != nil {
								return err
							}
						}
						if err := a.upsertWorkersForPipeline(&pipelineInfo); err != nil {
							if err := a.setPipelineFailure(ctx, pipelineInfo.Pipeline.Name, fmt.Sprintf("failed to create workers: %s", err.Error())); err != nil {
								return err
							}
							continue
						}
					}

					// If the pipeline has been updated, create new workers
					if pipelineInfo.Version > prevPipelineInfo.Version && !pipelineStateToStopped(pipelineInfo.State) {
						log.Infof("master: creating/updating workers for pipeline %s", pipelineInfo.Pipeline.Name)
						if event.PrevKey != nil {
							if err := a.deleteWorkersForPipeline(&prevPipelineInfo); err != nil {
								return err
							}
						}
						if hasGitInput {
							if err := a.checkOrDeployGithookService(); err != nil {
								return err
							}
						}
						if err := a.upsertWorkersForPipeline(&pipelineInfo); err != nil {
							if err := a.setPipelineFailure(ctx, pipelineInfo.Pipeline.Name, fmt.Sprintf("failed to create workers: %s", err.Error())); err != nil {
								return err
							}
							continue
						}
					}
				case watch.EventDelete:
					var pipelineName string
					var pipelineInfo pps.PipelineInfo
					if err := event.UnmarshalPrev(&pipelineName, &pipelineInfo); err != nil {
						return err
					}
					if err := a.deleteWorkersForPipeline(&pipelineInfo); err != nil {
						return err
					}
				}
			case event := <-watchChan:
				// if we get an error we restart the watch, k8s watches seem to
				// sometimes get stuck in a loop returning events with Type =
				// "" we treat these as errors since otherwise we get an
				// endless stream of them and can't do anything.
				if event.Type == kube_watch.Error || event.Type == "" {
					if kubePipelineWatch != nil {
						kubePipelineWatch.Stop()
					}
					kubePipelineWatch, err = a.kubeClient.CoreV1().Pods(a.namespace).Watch(
						metav1.ListOptions{
							LabelSelector: metav1.FormatLabelSelector(metav1.SetAsLabelSelector(
								map[string]string{
									"component": "worker",
								})),
							Watch: true,
						})
					if err != nil {
						log.Errorf("failed to watch kuburnetes pods: %v", err)
						watchChan = nil
					} else {
						watchChan = kubePipelineWatch.ResultChan()
						defer kubePipelineWatch.Stop()
					}
				}
				pod, ok := event.Object.(*v1.Pod)
				if !ok {
					continue
				}
				if pod.Status.Phase == v1.PodFailed {
					log.Errorf("pod failed because: %s", pod.Status.Message)
				}
				for _, status := range pod.Status.ContainerStatuses {
					if status.Name == "user" && status.State.Waiting != nil && failures[status.State.Waiting.Reason] {
						if err := a.setPipelineFailure(ctx, pod.ObjectMeta.Annotations["pipelineName"], status.State.Waiting.Message); err != nil {
							return err
						}
					}
				}
			}
		}
	}, backoff.NewInfiniteBackOff(), func(err error, d time.Duration) error {
		log.Errorf("master: error running the master process: %v; retrying in %v", err, d)
		return nil
	})
}

func (a *apiServer) setPipelineFailure(ctx context.Context, pipelineName string, reason string) error {
	return util.FailPipeline(ctx, a.etcdClient, a.pipelines, pipelineName, reason)
}

func (a *apiServer) checkOrDeployGithookService() error {
	_, err := getGithookService(a.kubeClient, a.namespace)
	if err != nil {
		if _, ok := err.(*errGithookServiceNotFound); ok {
			svc := assets.GithookService()
			_, err = a.kubeClient.CoreV1().Services(a.namespace).Create(svc)
			return err
		}
		return err
	}
	// service already exists
	return nil
}

func getGithookService(kubeClient *kube.Clientset, namespace string) (*v1.Service, error) {
	labels := map[string]string{
		"app":   "githook",
		"suite": suite,
	}
	serviceList, err := kubeClient.CoreV1().Services(namespace).List(metav1.ListOptions{
		TypeMeta: metav1.TypeMeta{
			Kind:       "ListOptions",
			APIVersion: "v1",
		},
		LabelSelector: metav1.FormatLabelSelector(metav1.SetAsLabelSelector(labels)),
	})
	if err != nil {
		return nil, err
	}
	if len(serviceList.Items) != 1 {
		return nil, &errGithookServiceNotFound{
			fmt.Errorf("githook service not found"),
		}
	}
	return &serviceList.Items[0], nil
}

func (a *apiServer) upsertWorkersForPipeline(pipelineInfo *pps.PipelineInfo) error {
	var errCount int
	return backoff.RetryNotify(func() error {
		parallelism, err := ppsserver.GetExpectedNumWorkers(a.kubeClient, pipelineInfo.ParallelismSpec)
		if err != nil {
			log.Errorf("error getting number of workers, default to 1 worker: %v", err)
			parallelism = 1
		}
		var resourceRequests *v1.ResourceList
		var resourceLimits *v1.ResourceList
		if pipelineInfo.ResourceRequests != nil {
			resourceRequests, err = util.GetRequestsResourceListFromPipeline(pipelineInfo)
			if err != nil {
				return err
			}
		}
		if pipelineInfo.ResourceLimits != nil {
			resourceLimits, err = util.GetLimitsResourceListFromPipeline(pipelineInfo)
			if err != nil {
				return err
			}
		}

		// Retrieve the current state of the RC.  If the RC is scaled down,
		// we want to ensure that it remains scaled down.
		rc := a.kubeClient.CoreV1().ReplicationControllers(a.namespace)
		workerRc, err := rc.Get(
			ppsserver.PipelineRcName(pipelineInfo.Pipeline.Name, pipelineInfo.Version),
			metav1.GetOptions{})
		if err == nil {
			if (workerRc.Spec.Template.Spec.Containers[0].Resources.Requests == nil) && *workerRc.Spec.Replicas == 1 {
				parallelism = 1
				resourceRequests = nil
				resourceLimits = nil
			}
		}

		options := a.getWorkerOptions(
			pipelineInfo.Pipeline.Name,
			ppsserver.PipelineRcName(pipelineInfo.Pipeline.Name, pipelineInfo.Version),
			int32(parallelism),
			resourceRequests,
			resourceLimits,
			pipelineInfo.Transform,
			pipelineInfo.CacheSize,
			pipelineInfo.Service)
		// Set the pipeline name env
		options.workerEnv = append(options.workerEnv, v1.EnvVar{
			Name:  client.PPSPipelineNameEnv,
			Value: pipelineInfo.Pipeline.Name,
		})
		return a.createWorkerRc(options)
	}, backoff.NewInfiniteBackOff(), func(err error, d time.Duration) error {
		errCount++
		if errCount >= 3 {
			return err
		}
		log.Errorf("error creating workers for pipeline %v: %v; retrying in %v", pipelineInfo.Pipeline.Name, err, d)
		return nil
	})
}

func (a *apiServer) deleteWorkersForPipeline(pipelineInfo *pps.PipelineInfo) error {
	rcName := ppsserver.PipelineRcName(pipelineInfo.Pipeline.Name, pipelineInfo.Version)
	if err := a.kubeClient.CoreV1().Services(a.namespace).Delete(
		rcName, &metav1.DeleteOptions{},
	); err != nil {
		if !isNotFoundErr(err) {
			return err
		}
	}
	if pipelineInfo.Service != nil {
		if err := a.kubeClient.CoreV1().Services(a.namespace).Delete(
			rcName+"-user", &metav1.DeleteOptions{},
		); err != nil {
			if !isNotFoundErr(err) {
				return err
			}
		}
	}
	falseVal := false
	deleteOptions := &metav1.DeleteOptions{
		OrphanDependents: &falseVal,
	}
	if err := a.kubeClient.CoreV1().ReplicationControllers(a.namespace).Delete(rcName, deleteOptions); err != nil {
		if !isNotFoundErr(err) {
			return err
		}
	}
	return nil
}
<commit_msg>Update error message based on feedback<commit_after>package server

import (
	"context"
	"fmt"
	"path"
	"time"

	log "github.com/sirupsen/logrus"
	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	kube_watch "k8s.io/apimachinery/pkg/watch"
	kube "k8s.io/client-go/kubernetes"

	"github.com/pachyderm/pachyderm/src/client"
	"github.com/pachyderm/pachyderm/src/client/pkg/uuid"
	"github.com/pachyderm/pachyderm/src/client/pps"
	"github.com/pachyderm/pachyderm/src/server/pkg/backoff"
	col "github.com/pachyderm/pachyderm/src/server/pkg/collection"
	"github.com/pachyderm/pachyderm/src/server/pkg/deploy/assets"
	"github.com/pachyderm/pachyderm/src/server/pkg/dlock"
	"github.com/pachyderm/pachyderm/src/server/pkg/util"
	"github.com/pachyderm/pachyderm/src/server/pkg/watch"
	ppsserver "github.com/pachyderm/pachyderm/src/server/pps"
)

const (
	masterLockPath = "_master_lock"
)

var (
	failures = map[string]bool{
		"InvalidImageName": true,
		"ErrImagePull":     true,
	}
)

// The master process is responsible for creating/deleting workers as
// pipelines are created/removed.
func (a *apiServer) master() {
	masterLock := dlock.NewDLock(a.etcdClient, path.Join(a.etcdPrefix, masterLockPath))
	backoff.RetryNotify(func() error {
		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		ctx, err := masterLock.Lock(ctx)
		if err != nil {
			return err
		}
		defer masterLock.Unlock(ctx)

		log.Infof("Launching PPS master process")

		pipelineWatcher, err := a.pipelines.ReadOnly(ctx).WatchWithPrev()
		if err != nil {
			return fmt.Errorf("error creating watch: %+v", err)
		}
		defer pipelineWatcher.Close()

		// watchChan will be nil if the Watch call below errors, this means
		// that we won't receive events from k8s and won't be able to detect
		// errors in pods. We could just return that error and retry but that
		// prevents pachyderm from creating pipelines when there's an issue
		// talking to k8s.
		var watchChan <-chan kube_watch.Event
		kubePipelineWatch, err := a.kubeClient.CoreV1().Pods(a.namespace).Watch(
			metav1.ListOptions{
				LabelSelector: metav1.FormatLabelSelector(metav1.SetAsLabelSelector(
					map[string]string{
						"component": "worker",
					})),
				Watch: true,
			})
		if err != nil {
			log.Errorf("failed to watch kuburnetes pods: %v", err)
		} else {
			watchChan = kubePipelineWatch.ResultChan()
			defer kubePipelineWatch.Stop()
		}

		for {
			select {
			case event := <-pipelineWatcher.Watch():
				if event.Err != nil {
					return fmt.Errorf("event err: %+v", event.Err)
				}
				switch event.Type {
				case watch.EventPut:
					var pipelineName string
					var pipelineInfo pps.PipelineInfo
					if err := event.Unmarshal(&pipelineName, &pipelineInfo); err != nil {
						return err
					}

					// This is a bit of a hack that covers for migrations bugs
					// where a field is set by setPipelineDefaults in a newer
					// version of the server but a PipelineInfo which was
					// created with an older version of the server doesn't have
					// that field set because setPipelineDefaults was different
					// when it was created.
					if pipelineInfo.Salt == "" || pipelineInfo.CacheSize == "" {
						if _, err := col.NewSTM(ctx, a.etcdClient, func(stm col.STM) error {
							pipelines := a.pipelines.ReadWrite(stm)
							newPipelineInfo := new(pps.PipelineInfo)
							if err := pipelines.Get(pipelineInfo.Pipeline.Name, newPipelineInfo); err != nil {
								return fmt.Errorf("error getting pipeline %s: %+v", pipelineName, err)
							}
							if newPipelineInfo.Salt == "" {
								newPipelineInfo.Salt = uuid.NewWithoutDashes()
							}
							setPipelineDefaults(newPipelineInfo)
							pipelines.Put(pipelineInfo.Pipeline.Name, newPipelineInfo)
							pipelineInfo = *newPipelineInfo
							return nil
						}); err != nil {
							return err
						}
					}

					var prevPipelineInfo pps.PipelineInfo
					if event.PrevKey != nil {
						if err := event.UnmarshalPrev(&pipelineName, &prevPipelineInfo); err != nil {
							return err
						}
					}

					// If the pipeline has been stopped, delete workers
					if pipelineStateToStopped(pipelineInfo.State) {
						log.Infof("master: deleting workers for pipeline %s", pipelineInfo.Pipeline.Name)
						if err := a.deleteWorkersForPipeline(&pipelineInfo); err != nil {
							return err
						}
					}

					var hasGitInput bool
					pps.VisitInput(pipelineInfo.Input, func(input *pps.Input) {
						if input.Git != nil {
							hasGitInput = true
						}
					})

					// If the pipeline has been restarted, create workers
					if !pipelineStateToStopped(pipelineInfo.State) && event.PrevKey != nil && pipelineStateToStopped(prevPipelineInfo.State) {
						if hasGitInput {
							if err := a.checkOrDeployGithookService(); err != nil {
								return err
							}
						}
						if err := a.upsertWorkersForPipeline(&pipelineInfo); err != nil {
							if err := a.setPipelineFailure(ctx, pipelineInfo.Pipeline.Name, fmt.Sprintf("failed to create workers: %s", err.Error())); err != nil {
								return err
							}
							continue
						}
					}

					// If the pipeline has been updated, create new workers
					if pipelineInfo.Version > prevPipelineInfo.Version && !pipelineStateToStopped(pipelineInfo.State) {
						log.Infof("master: creating/updating workers for pipeline %s", pipelineInfo.Pipeline.Name)
						if event.PrevKey != nil {
							if err := a.deleteWorkersForPipeline(&prevPipelineInfo); err != nil {
								return err
							}
						}
						if hasGitInput {
							if err := a.checkOrDeployGithookService(); err != nil {
								return err
							}
						}
						if err := a.upsertWorkersForPipeline(&pipelineInfo); err != nil {
							if err := a.setPipelineFailure(ctx, pipelineInfo.Pipeline.Name, fmt.Sprintf("failed to create workers: %s", err.Error())); err != nil {
								return err
							}
							continue
						}
					}
				case watch.EventDelete:
					var pipelineName string
					var pipelineInfo pps.PipelineInfo
					if err := event.UnmarshalPrev(&pipelineName, &pipelineInfo); err != nil {
						return err
					}
					if err := a.deleteWorkersForPipeline(&pipelineInfo); err != nil {
						return err
					}
				}
			case event := <-watchChan:
				// if we get an error we restart the watch, k8s watches seem to
				// sometimes get stuck in a loop returning events with Type =
				// "" we treat these as errors since otherwise we get an
				// endless stream of them and can't do anything.
				if event.Type == kube_watch.Error || event.Type == "" {
					if kubePipelineWatch != nil {
						kubePipelineWatch.Stop()
					}
					kubePipelineWatch, err = a.kubeClient.CoreV1().Pods(a.namespace).Watch(
						metav1.ListOptions{
							LabelSelector: metav1.FormatLabelSelector(metav1.SetAsLabelSelector(
								map[string]string{
									"component": "worker",
								})),
							Watch: true,
						})
					if err != nil {
						log.Errorf("failed to watch kuburnetes pods: %v", err)
						watchChan = nil
					} else {
						watchChan = kubePipelineWatch.ResultChan()
						defer kubePipelineWatch.Stop()
					}
				}
				pod, ok := event.Object.(*v1.Pod)
				if !ok {
					continue
				}
				if pod.Status.Phase == v1.PodFailed {
					log.Errorf("pod failed because: %s", pod.Status.Message)
				}
				for _, status := range pod.Status.ContainerStatuses {
					if status.Name == "user" && status.State.Waiting != nil && failures[status.State.Waiting.Reason] {
						if err := a.setPipelineFailure(ctx, pod.ObjectMeta.Annotations["pipelineName"], status.State.Waiting.Message); err != nil {
							return err
						}
					}
				}
			}
		}
	}, backoff.NewInfiniteBackOff(), func(err error, d time.Duration) error {
		log.Errorf("master: error running the master process: %v; retrying in %v", err, d)
		return nil
	})
}

func (a *apiServer) setPipelineFailure(ctx context.Context, pipelineName string, reason string) error {
	return util.FailPipeline(ctx, a.etcdClient, a.pipelines, pipelineName, reason)
}

func (a *apiServer) checkOrDeployGithookService() error {
	_, err := getGithookService(a.kubeClient, a.namespace)
	if err != nil {
		if _, ok := err.(*errGithookServiceNotFound); ok {
			svc := assets.GithookService()
			_, err = a.kubeClient.CoreV1().Services(a.namespace).Create(svc)
			return err
		}
		return err
	}
	// service already exists
	return nil
}

func getGithookService(kubeClient *kube.Clientset, namespace string) (*v1.Service, error) {
	labels := map[string]string{
		"app":   "githook",
		"suite": suite,
	}
	serviceList, err := kubeClient.CoreV1().Services(namespace).List(metav1.ListOptions{
		TypeMeta: metav1.TypeMeta{
			Kind:       "ListOptions",
			APIVersion: "v1",
		},
		LabelSelector: metav1.FormatLabelSelector(metav1.SetAsLabelSelector(labels)),
	})
	if err != nil {
		return nil, err
	}
	if len(serviceList.Items) != 1 {
		return nil, &errGithookServiceNotFound{
			fmt.Errorf("expected 1 githook service but found %v", len(serviceList.Items)),
		}
	}
	return &serviceList.Items[0], nil
}

func (a *apiServer) upsertWorkersForPipeline(pipelineInfo *pps.PipelineInfo) error {
	var errCount int
	return backoff.RetryNotify(func() error {
		parallelism, err := ppsserver.GetExpectedNumWorkers(a.kubeClient, pipelineInfo.ParallelismSpec)
		if err != nil {
			log.Errorf("error getting number of workers, default to 1 worker: %v", err)
			parallelism = 1
		}
		var resourceRequests *v1.ResourceList
		var resourceLimits *v1.ResourceList
		if pipelineInfo.ResourceRequests != nil {
			resourceRequests, err = util.GetRequestsResourceListFromPipeline(pipelineInfo)
			if err != nil {
				return err
			}
		}
		if pipelineInfo.ResourceLimits != nil {
			resourceLimits, err = util.GetLimitsResourceListFromPipeline(pipelineInfo)
			if err != nil {
				return err
			}
		}

		// Retrieve the current state of the RC.  If the RC is scaled down,
		// we want to ensure that it remains scaled down.
		rc := a.kubeClient.CoreV1().ReplicationControllers(a.namespace)
		workerRc, err := rc.Get(
			ppsserver.PipelineRcName(pipelineInfo.Pipeline.Name, pipelineInfo.Version),
			metav1.GetOptions{})
		if err == nil {
			if (workerRc.Spec.Template.Spec.Containers[0].Resources.Requests == nil) && *workerRc.Spec.Replicas == 1 {
				parallelism = 1
				resourceRequests = nil
				resourceLimits = nil
			}
		}

		options := a.getWorkerOptions(
			pipelineInfo.Pipeline.Name,
			ppsserver.PipelineRcName(pipelineInfo.Pipeline.Name, pipelineInfo.Version),
			int32(parallelism),
			resourceRequests,
			resourceLimits,
			pipelineInfo.Transform,
			pipelineInfo.CacheSize,
			pipelineInfo.Service)
		// Set the pipeline name env
		options.workerEnv = append(options.workerEnv, v1.EnvVar{
			Name:  client.PPSPipelineNameEnv,
			Value: pipelineInfo.Pipeline.Name,
		})
		return a.createWorkerRc(options)
	}, backoff.NewInfiniteBackOff(), func(err error, d time.Duration) error {
		errCount++
		if errCount >= 3 {
			return err
		}
		log.Errorf("error creating workers for pipeline %v: %v; retrying in %v", pipelineInfo.Pipeline.Name, err, d)
		return nil
	})
}

func (a *apiServer) deleteWorkersForPipeline(pipelineInfo *pps.PipelineInfo) error {
	rcName := ppsserver.PipelineRcName(pipelineInfo.Pipeline.Name, pipelineInfo.Version)
	if err := a.kubeClient.CoreV1().Services(a.namespace).Delete(
		rcName, &metav1.DeleteOptions{},
	); err != nil {
		if !isNotFoundErr(err) {
			return err
		}
	}
	if pipelineInfo.Service != nil {
		if err := a.kubeClient.CoreV1().Services(a.namespace).Delete(
			rcName+"-user", &metav1.DeleteOptions{},
		); err != nil {
			if !isNotFoundErr(err) {
				return err
			}
		}
	}
	falseVal := false
	deleteOptions := &metav1.DeleteOptions{
		OrphanDependents: &falseVal,
	}
	if err := a.kubeClient.CoreV1().ReplicationControllers(a.namespace).Delete(rcName, deleteOptions); err != nil {
		if !isNotFoundErr(err) {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2018 fatedier, fatedier@gmail.com
//
// 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 mux

import (
	"fmt"
	"io"
	"net"
	"sort"
	"sync"
	"time"

	"github.com/fatedier/golib/errors"
	gnet "github.com/fatedier/golib/net"
)

const (
	// DefaultTimeout is the default length of time to wait for bytes we need.
	DefaultTimeout = 10 * time.Second
)

type Mux struct {
	ln net.Listener

	defaultLn *listener

	// sorted by priority
	lns             []*listener
	maxNeedBytesNum uint32

	mu sync.RWMutex
}

func NewMux(ln net.Listener) (mux *Mux) {
	mux = &Mux{
		ln:  ln,
		lns: make([]*listener, 0),
	}
	return
}

// priority
func (mux *Mux) Listen(priority int, needBytesNum uint32, fn MatchFunc) net.Listener {
	ln := &listener{
		c:            make(chan net.Conn),
		mux:          mux,
		priority:     priority,
		needBytesNum: needBytesNum,
		matchFn:      fn,
	}

	mux.mu.Lock()
	defer mux.mu.Unlock()
	if needBytesNum > mux.maxNeedBytesNum {
		mux.maxNeedBytesNum = needBytesNum
	}

	newlns := append(mux.copyLns(), ln)
	sort.Slice(newlns, func(i, j int) bool {
		if newlns[i].priority == newlns[j].priority {
			return newlns[i].needBytesNum < newlns[j].needBytesNum
		}
		return newlns[i].priority < newlns[j].priority
	})
	mux.lns = newlns
	return ln
}

func (mux *Mux) ListenHttp(priority int) net.Listener {
	return mux.Listen(priority, HttpNeedBytesNum, HttpMatchFunc)
}

func (mux *Mux) ListenHttps(priority int) net.Listener {
	return mux.Listen(priority, HttpsNeedBytesNum, HttpsMatchFunc)
}

func (mux *Mux) DefaultListener() net.Listener {
	mux.mu.Lock()
	defer mux.mu.Unlock()
	if mux.defaultLn == nil {
		mux.defaultLn = &listener{
			c:   make(chan net.Conn),
			mux: mux,
		}
	}
	return mux.defaultLn
}

func (mux *Mux) release(ln *listener) bool {
	result := false
	mux.mu.Lock()
	defer mux.mu.Unlock()
	lns := mux.copyLns()

	for i, l := range lns {
		if l == ln {
			lns = append(lns[:i], lns[i+1:]...)
			result = true
			break
		}
	}
	mux.lns = lns
	return result
}

func (mux *Mux) copyLns() []*listener {
	lns := make([]*listener, 0, len(mux.lns))
	for _, l := range mux.lns {
		lns = append(lns, l)
	}
	return lns
}

// Serve handles connections from ln and multiplexes then across registered listeners.
func (mux *Mux) Serve() error {
	var tempDelay time.Duration // how long to sleep on accept failure

	for {
		// Wait for the next connection.
		// If it returns a temporary error then simply retry.
		// If it returns any other error then exit immediately.
		conn, err := mux.ln.Accept()
		if err, ok := err.(interface {
			Temporary() bool
		}); ok && err.Temporary() {
			if tempDelay == 0 {
				tempDelay = 5 * time.Millisecond
			} else {
				tempDelay *= 2
			}
			if max := 1 * time.Second; tempDelay > max {
				tempDelay = max
			}
			time.Sleep(tempDelay)
			continue
		}

		if err != nil {
			return err
		}

		go mux.handleConn(conn)
	}
}

func (mux *Mux) handleConn(conn net.Conn) {
	mux.mu.RLock()
	maxNeedBytesNum := mux.maxNeedBytesNum
	lns := mux.lns
	defaultLn := mux.defaultLn
	mux.mu.RUnlock()

	sharedConn, rd := gnet.NewSharedConnSize(conn, int(maxNeedBytesNum))
	data := make([]byte, maxNeedBytesNum)

	conn.SetReadDeadline(time.Now().Add(DefaultTimeout))
	_, err := io.ReadFull(rd, data)
	if err != nil {
		conn.Close()
		return
	}
	conn.SetReadDeadline(time.Time{})

	for _, ln := range lns {
		if match := ln.matchFn(data); match {
			err = errors.PanicToError(func() {
				ln.c <- sharedConn
			})
			if err != nil {
				conn.Close()
			}
			return
		}
	}

	// No match listeners
	if defaultLn != nil {
		err = errors.PanicToError(func() {
			defaultLn.c <- sharedConn
		})
		if err != nil {
			conn.Close()
		}
		return
	}

	// No listeners for this connection, close it.
	conn.Close()
	return
}

type listener struct {
	mux *Mux

	priority     int
	needBytesNum uint32
	matchFn      MatchFunc

	c  chan net.Conn
	mu sync.RWMutex
}

// Accept waits for and returns the next connection to the listener.
func (ln *listener) Accept() (net.Conn, error) {
	conn, ok := <-ln.c
	if !ok {
		return nil, fmt.Errorf("network connection closed")
	}
	return conn, nil
}

// Close removes this listener from the parent mux and closes the channel.
func (ln *listener) Close() error {
	if ok := ln.mux.release(ln); ok {
		// Close done to signal to any RLock holders to release their lock.
		close(ln.c)
	}
	return nil
}

func (ln *listener) Addr() net.Addr {
	if ln.mux == nil {
		return nil
	}
	ln.mux.mu.RLock()
	defer ln.mux.mu.RUnlock()
	if ln.mux.ln == nil {
		return nil
	}
	return ln.mux.ln.Addr()
}
<commit_msg>update mux<commit_after>// Copyright 2018 fatedier, fatedier@gmail.com
//
// 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 mux

import (
	"fmt"
	"io"
	"net"
	"sort"
	"sync"
	"time"

	"github.com/fatedier/golib/errors"
	gnet "github.com/fatedier/golib/net"
)

const (
	// DefaultTimeout is the default length of time to wait for bytes we need.
	DefaultTimeout = 10 * time.Second
)

type Mux struct {
	ln net.Listener

	defaultLn *listener

	// sorted by priority
	lns             []*listener
	maxNeedBytesNum uint32

	mu sync.RWMutex
}

func NewMux(ln net.Listener) (mux *Mux) {
	mux = &Mux{
		ln:  ln,
		lns: make([]*listener, 0),
	}
	return
}

// priority
func (mux *Mux) Listen(priority int, needBytesNum uint32, fn MatchFunc) net.Listener {
	ln := &listener{
		c:            make(chan net.Conn),
		mux:          mux,
		priority:     priority,
		needBytesNum: needBytesNum,
		matchFn:      fn,
	}

	mux.mu.Lock()
	defer mux.mu.Unlock()
	if needBytesNum > mux.maxNeedBytesNum {
		mux.maxNeedBytesNum = needBytesNum
	}

	newlns := append(mux.copyLns(), ln)
	sort.Slice(newlns, func(i, j int) bool {
		if newlns[i].priority == newlns[j].priority {
			return newlns[i].needBytesNum < newlns[j].needBytesNum
		}
		return newlns[i].priority < newlns[j].priority
	})
	mux.lns = newlns
	return ln
}

func (mux *Mux) ListenHttp(priority int) net.Listener {
	return mux.Listen(priority, HttpNeedBytesNum, HttpMatchFunc)
}

func (mux *Mux) ListenHttps(priority int) net.Listener {
	return mux.Listen(priority, HttpsNeedBytesNum, HttpsMatchFunc)
}

func (mux *Mux) DefaultListener() net.Listener {
	mux.mu.Lock()
	defer mux.mu.Unlock()
	if mux.defaultLn == nil {
		mux.defaultLn = &listener{
			c:   make(chan net.Conn),
			mux: mux,
		}
	}
	return mux.defaultLn
}

func (mux *Mux) release(ln *listener) bool {
	result := false
	mux.mu.Lock()
	defer mux.mu.Unlock()
	lns := mux.copyLns()

	for i, l := range lns {
		if l == ln {
			lns = append(lns[:i], lns[i+1:]...)
			result = true
			break
		}
	}
	mux.lns = lns
	return result
}

func (mux *Mux) copyLns() []*listener {
	lns := make([]*listener, 0, len(mux.lns))
	for _, l := range mux.lns {
		lns = append(lns, l)
	}
	return lns
}

// Serve handles connections from ln and multiplexes then across registered listeners.
func (mux *Mux) Serve() error {
	var tempDelay time.Duration // how long to sleep on accept failure

	for {
		// Wait for the next connection.
		// If it returns a temporary error then simply retry.
		// If it returns any other error then exit immediately.
		conn, err := mux.ln.Accept()
		if err, ok := err.(interface {
			Temporary() bool
		}); ok && err.Temporary() {
			if tempDelay == 0 {
				tempDelay = 5 * time.Millisecond
			} else {
				tempDelay *= 2
			}
			if max := 1 * time.Second; tempDelay > max {
				tempDelay = max
			}
			time.Sleep(tempDelay)
			continue
		}

		if err != nil {
			return err
		}
		tempDelay = 0

		go mux.handleConn(conn)
	}
}

func (mux *Mux) handleConn(conn net.Conn) {
	mux.mu.RLock()
	maxNeedBytesNum := mux.maxNeedBytesNum
	lns := mux.lns
	defaultLn := mux.defaultLn
	mux.mu.RUnlock()

	sharedConn, rd := gnet.NewSharedConnSize(conn, int(maxNeedBytesNum))
	data := make([]byte, maxNeedBytesNum)

	conn.SetReadDeadline(time.Now().Add(DefaultTimeout))
	_, err := io.ReadFull(rd, data)
	if err != nil {
		conn.Close()
		return
	}
	conn.SetReadDeadline(time.Time{})

	for _, ln := range lns {
		if match := ln.matchFn(data); match {
			err = errors.PanicToError(func() {
				ln.c <- sharedConn
			})
			if err != nil {
				conn.Close()
			}
			return
		}
	}

	// No match listeners
	if defaultLn != nil {
		err = errors.PanicToError(func() {
			defaultLn.c <- sharedConn
		})
		if err != nil {
			conn.Close()
		}
		return
	}

	// No listeners for this connection, close it.
	conn.Close()
	return
}

type listener struct {
	mux *Mux

	priority     int
	needBytesNum uint32
	matchFn      MatchFunc

	c  chan net.Conn
	mu sync.RWMutex
}

// Accept waits for and returns the next connection to the listener.
func (ln *listener) Accept() (net.Conn, error) {
	conn, ok := <-ln.c
	if !ok {
		return nil, fmt.Errorf("network connection closed")
	}
	return conn, nil
}

// Close removes this listener from the parent mux and closes the channel.
func (ln *listener) Close() error {
	if ok := ln.mux.release(ln); ok {
		// Close done to signal to any RLock holders to release their lock.
		close(ln.c)
	}
	return nil
}

func (ln *listener) Addr() net.Addr {
	if ln.mux == nil {
		return nil
	}
	ln.mux.mu.RLock()
	defer ln.mux.mu.RUnlock()
	if ln.mux.ln == nil {
		return nil
	}
	return ln.mux.ln.Addr()
}
<|endoftext|>


<commit_before>// GoMySQL - A MySQL client library for Go
//
// Copyright 2010-2011 Phil Bayfield. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package mysql

import (
	"bytes"
	"fmt"
	"os"
)

// Packet type identifier
type packetType uint32

// Packet types
const (
	PACKET_INIT        packetType = 1 << iota
	PACKET_AUTH        packetType = 1 << iota
	PACKET_OK          packetType = 1 << iota
	PACKET_ERROR       packetType = 1 << iota
	PACKET_CMD         packetType = 1 << iota
	PACKET_RESULT      packetType = 1 << iota
	PACKET_FIELD       packetType = 1 << iota
	PACKET_ROW         packetType = 1 << iota
	PACKET_EOF_40      packetType = 1 << iota
	PACKET_EOF_41      packetType = 1 << iota
	PACKET_OK_PREPARED packetType = 1 << iota
	PACKET_PARAM       packetType = 1 << iota
	PACKET_LONG_DATA   packetType = 1 << iota
	PACKET_EXECUTE     packetType = 1 << iota
	PACKET_ROW_BINARY  packetType = 1 << iota
)

// Readable packet interface
type packetReadable interface {
	read(data []byte) (err os.Error)
}

// Writable packet interface
type packetWritable interface {
	write() (data []byte, err os.Error)
}

// Generic packet interface (read/writable)
type packet interface {
	packetReadable
	packetWritable
}

// Packet base struct
type packetBase struct {
	protocol uint8
	sequence uint8
}

// Read a slice from the data
func (p *packetBase) readSlice(data []byte, delim byte) (slice []byte, err os.Error) {
	pos := bytes.IndexByte(data, delim)
	if pos > -1 {
		slice = data[:pos]
	} else {
		slice = data
		err = os.EOF
	}
	return
}

// Read length coded binary
func (p *packetBase) readLengthCodedBinary(data []byte) (num uint64, n int, err os.Error) {
	switch {
	// 0-250 = value of first byte
	case data[0] <= 250:
		num = uint64(data[0])
		n = 1
		return
	// 251 column value = NULL
	case data[0] == 251:
		num = 0
		n = 1
		return
	// 252 following 2 = value of following 16-bit word
	case data[0] == 252:
		n = 3
	// 253 following 3 = value of following 24-bit word
	case data[0] == 253:
		n = 4
	// 254 following 8 = value of following 64-bit word
	case data[0] == 254:
		n = 9
	}
	// Check there are enough bytes
	if len(data) < n {
		err = os.EOF
		return
	}
	// Unpack number
	num = p.unpackNumber(data[1:n])
	return
}

// Convert byte array into a number
func (p *packetBase) unpackNumber(data []byte) (num uint64) {
	for i := uint8(0); i < uint8(len(data)); i++ {
		num |= uint64(data[i]) << (i * 8)
	}
	return
}

// Convert number into a byte array
func (p *packetBase) packNumber(num uint64, n uint8) (bytes []byte) {
	bytes = make([]byte, n)
	for i := uint8(0); i < n; i++ {
		bytes[i] = byte(num >> (i * 8))
	}
	return
}

// Prepend packet data with header info
func (p *packetBase) addHeader(data []byte) (pkt []byte) {
	pkt = p.packNumber(uint64(len(data)), 3)
	pkt = append(pkt, p.sequence)
	pkt = append(pkt, data...)
	return
}

// Init packet
type packetInit struct {
	packetBase
	protocolVersion uint8
	serverVersion   string
	threadId        uint32
	scrambleBuff    []byte
	serverCaps      uint16
	serverLanguage  uint8
	serverStatus    uint16
}

// Init packet reader
func (p *packetInit) read(data []byte) (err os.Error) {
	// Recover errors
	defer func() {
		if e := recover(); e != nil {
			err = os.NewError(fmt.Sprintf("%s", e))
		}
	}()
	// Position
	pos := 0
	// Protocol version [8 bit uint]
	p.protocolVersion = data[pos]
	pos++
	// Server version [null terminated string]
	slice, err := p.readSlice(data[pos:], 0x00)
	if err != nil {
		return
	}
	p.serverVersion = string(slice)
	pos += len(slice) + 1
	// Thread id [32 bit uint]
	p.threadId = uint32(p.unpackNumber(data[pos : pos+4]))
	pos += 4
	// First part of scramble buffer [8 bytes]
	p.scrambleBuff = make([]byte, 8)
	p.scrambleBuff = data[pos : pos+8]
	pos += 9
	// Server capabilities [16 bit uint]
	p.serverCaps = uint16(p.unpackNumber(data[pos : pos+2]))
	pos += 2
	// Server language [8 bit uint]
	p.serverLanguage = data[pos]
	pos++
	// Server status [16 bit uint]
	p.serverStatus = uint16(p.unpackNumber(data[pos : pos+2]))
	pos += 15
	// Second part of scramble buffer, if exists (4.1+) [13 bytes]
	if ClientFlag(p.serverCaps)&CLIENT_PROTOCOL_41 > 0 {
		p.scrambleBuff = append(p.scrambleBuff, data[pos:pos+12]...)
	}
	return
}

// Auth packet
type packetAuth struct {
	packetBase
	clientFlags   uint32
	maxPacketSize uint32
	charsetNumber uint8
	user          string
	scrambleBuff  []byte
	database      string
}

// Auth packet writer
func (p *packetAuth) write() (data []byte, err os.Error) {
	// For MySQL 4.1+
	if p.protocol == PROTOCOL_41 {
		// Client flags
		data = p.packNumber(uint64(p.clientFlags), 4)
		// Max packet size
		data = append(data, p.packNumber(uint64(p.maxPacketSize), 4)...)
		// Charset
		data = append(data, p.charsetNumber)
		// Filler
		data = append(data, make([]byte, 23)...)
		// User
		if len(p.user) > 0 {
			data = append(data, []byte(p.user)...)
		}
		// Terminator
		data = append(data, 0x00)
		// Scramble buffer
		data = append(data, byte(len(p.scrambleBuff)))
		if len(p.scrambleBuff) > 0 {
			data = append(data, p.scrambleBuff...)
		}
		// Database name
		if len(p.database) > 0 {
			data = append(data, []byte(p.database)...)
			// Terminator
			data = append(data, 0x00)
		}
		// For MySQL < 4.1
	} else {
		// Client flags
		data = p.packNumber(uint64(p.clientFlags), 2)
		// Max packet size
		data = append(data, p.packNumber(uint64(p.maxPacketSize), 3)...)
		// User
		if len(p.user) > 0 {
			data = append(data, []byte(p.user)...)
		}
		// Terminator
		data = append(data, 0x00)
		// Scramble buffer
		if len(p.scrambleBuff) > 0 {
			data = append(data, p.scrambleBuff...)
		}
		// Padding
		data = append(data, 0x00)
	}
	// Add the packet header
	data = p.addHeader(data)
	return
}

// Ok packet struct
type packetOK struct {
	packetBase
	affectedRows uint64
	insertId     uint64
	serverStatus uint16
	warningCount uint16
	message      string
}

// OK packet reader
func (p *packetOK) read(data []byte) (err os.Error) {
	// Recover errors
	defer func() {
		if e := recover(); e != nil {
			err = os.NewError(fmt.Sprintf("%s", e))
		}
	}()
	// Position (skip first byte/field count)
	pos := 1
	// Affected rows [length coded binary]
	num, n, err := p.readLengthCodedBinary(data[pos:])
	if err != nil {
		return
	}
	p.affectedRows = num
	pos += n
	// Insert id [length coded binary]
	num, n, err = p.readLengthCodedBinary(data[pos:])
	if err != nil {
		return
	}
	p.insertId = num
	pos += n
	// Server status [16 bit uint]
	p.serverStatus = uint16(p.unpackNumber(data[pos : pos+2]))
	pos += 2
	// Warning (4.1 only) [16 bit uint]
	if p.protocol == PROTOCOL_41 {
		p.warningCount = uint16(p.unpackNumber(data[pos : pos+2]))
		pos += 2
	}
	// Message (optional) [string]
	if pos < len(data) {
		p.message = string(data[pos:])
	}
	return
}

// Error packet struct
type packetError struct {
	packetBase
	errno uint16
	state string
	error string
}

// Error packet reader
func (p *packetError) read(data []byte) (err os.Error) {
	// Recover errors
	defer func() {
		if e := recover(); e != nil {
			err = os.NewError(fmt.Sprintf("%s", e))
		}
	}()
	// Position
	pos := 1
	// Error number [16 bit uint]
	p.errno = uint16(p.unpackNumber(data[pos : pos+2]))
	pos += 2
	// State (4.1 only) [string]
	if p.protocol == PROTOCOL_41 {
		pos++
		p.state = string(data[pos : pos+5])
		pos += 5
	}
	// Message [string]
	p.error = string(data[pos:])
	return
}

// Command packet struct
type packetCommand struct {
	packetBase
	command command
	args    []interface{}
}

// Command packet writer
func (p *packetCommand) write() (data []byte, err os.Error) {
	// Recover errors (if wrong param type supplied)
	defer func() {
		if e := recover(); e != nil {
			err = os.NewError(fmt.Sprintf("%s", e))
		}
	}()
	// Make slice from command byte
	data = []byte{byte(p.command)}
	// Add args to requests
	switch p.command {
	// Commands with 1 arg unterminated string
	case COM_INIT_DB, COM_QUERY, COM_CREATE_DB, COM_DROP_DB, COM_STMT_PREPARE:
		data = append(data, []byte(p.args[0].(string))...)
	// Commands with 1 arg 32 bit uint
	case COM_PROCESS_KILL, COM_STMT_CLOSE, COM_STMT_RESET:
		data = append(data, p.packNumber(uint64(p.args[0].(uint32)), 4)...)
	// Field list command
	case COM_FIELD_LIST:
		// Table name
		data = append(data, []byte(p.args[0].(string))...)
		// Terminator
		data = append(data, 0x00)
		// Column name
		if len(p.args) > 1 {
			data = append(data, []byte(p.args[1].(string))...)
		}
	// Refresh command
	case COM_REFRESH:
		data = append(data, byte(p.args[0].(Refresh)))
	// Shutdown command
	case COM_SHUTDOWN:
		data = append(data, byte(p.args[0].(Shutdown)))
	// Change user command
	case COM_CHANGE_USER:
		// User
		data = append(data, []byte(p.args[0].(string))...)
		// Terminator
		data = append(data, 0x00)
		// Scramble length for 4.1
		if p.protocol == PROTOCOL_41 {
			data = append(data, byte(len(p.args[1].([]byte))))
		}
		// Scramble buffer
		if len(p.args[1].([]byte)) > 0 {
			data = append(data, p.args[1].([]byte)...)
		}
		// Temrminator for 3.23
		if p.protocol == PROTOCOL_40 {
			data = append(data, 0x00)
		}
		// Database name
		if len(p.args[2].(string)) > 0 {
			data = append(data, []byte(p.args[2].(string))...)
		}
		// Terminator
		data = append(data, 0x00)
		// Character set number (5.1.23+ needs testing with earlier versions)
		data = append(data, p.packNumber(uint64(p.args[3].(uint16)), 2)...)
	// Fetch statement command
	case COM_STMT_FETCH:
		// Statement id
		data = append(data, p.packNumber(uint64(p.args[0].(uint32)), 4)...)
		// Number of rows
		data = append(data, p.packNumber(uint64(p.args[1].(uint32)), 4)...)
	}
	return
}
<commit_msg>added missing packet header to command packets<commit_after>// GoMySQL - A MySQL client library for Go
//
// Copyright 2010-2011 Phil Bayfield. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package mysql

import (
	"bytes"
	"fmt"
	"os"
)

// Packet type identifier
type packetType uint32

// Packet types
const (
	PACKET_INIT        packetType = 1 << iota
	PACKET_AUTH        packetType = 1 << iota
	PACKET_OK          packetType = 1 << iota
	PACKET_ERROR       packetType = 1 << iota
	PACKET_CMD         packetType = 1 << iota
	PACKET_RESULT      packetType = 1 << iota
	PACKET_FIELD       packetType = 1 << iota
	PACKET_ROW         packetType = 1 << iota
	PACKET_EOF_40      packetType = 1 << iota
	PACKET_EOF_41      packetType = 1 << iota
	PACKET_OK_PREPARED packetType = 1 << iota
	PACKET_PARAM       packetType = 1 << iota
	PACKET_LONG_DATA   packetType = 1 << iota
	PACKET_EXECUTE     packetType = 1 << iota
	PACKET_ROW_BINARY  packetType = 1 << iota
)

// Readable packet interface
type packetReadable interface {
	read(data []byte) (err os.Error)
}

// Writable packet interface
type packetWritable interface {
	write() (data []byte, err os.Error)
}

// Generic packet interface (read/writable)
type packet interface {
	packetReadable
	packetWritable
}

// Packet base struct
type packetBase struct {
	protocol uint8
	sequence uint8
}

// Read a slice from the data
func (p *packetBase) readSlice(data []byte, delim byte) (slice []byte, err os.Error) {
	pos := bytes.IndexByte(data, delim)
	if pos > -1 {
		slice = data[:pos]
	} else {
		slice = data
		err = os.EOF
	}
	return
}

// Read length coded binary
func (p *packetBase) readLengthCodedBinary(data []byte) (num uint64, n int, err os.Error) {
	switch {
	// 0-250 = value of first byte
	case data[0] <= 250:
		num = uint64(data[0])
		n = 1
		return
	// 251 column value = NULL
	case data[0] == 251:
		num = 0
		n = 1
		return
	// 252 following 2 = value of following 16-bit word
	case data[0] == 252:
		n = 3
	// 253 following 3 = value of following 24-bit word
	case data[0] == 253:
		n = 4
	// 254 following 8 = value of following 64-bit word
	case data[0] == 254:
		n = 9
	}
	// Check there are enough bytes
	if len(data) < n {
		err = os.EOF
		return
	}
	// Unpack number
	num = p.unpackNumber(data[1:n])
	return
}

// Convert byte array into a number
func (p *packetBase) unpackNumber(data []byte) (num uint64) {
	for i := uint8(0); i < uint8(len(data)); i++ {
		num |= uint64(data[i]) << (i * 8)
	}
	return
}

// Convert number into a byte array
func (p *packetBase) packNumber(num uint64, n uint8) (bytes []byte) {
	bytes = make([]byte, n)
	for i := uint8(0); i < n; i++ {
		bytes[i] = byte(num >> (i * 8))
	}
	return
}

// Prepend packet data with header info
func (p *packetBase) addHeader(data []byte) (pkt []byte) {
	pkt = p.packNumber(uint64(len(data)), 3)
	pkt = append(pkt, p.sequence)
	pkt = append(pkt, data...)
	return
}

// Init packet
type packetInit struct {
	packetBase
	protocolVersion uint8
	serverVersion   string
	threadId        uint32
	scrambleBuff    []byte
	serverCaps      uint16
	serverLanguage  uint8
	serverStatus    uint16
}

// Init packet reader
func (p *packetInit) read(data []byte) (err os.Error) {
	// Recover errors
	defer func() {
		if e := recover(); e != nil {
			err = os.NewError(fmt.Sprintf("%s", e))
		}
	}()
	// Position
	pos := 0
	// Protocol version [8 bit uint]
	p.protocolVersion = data[pos]
	pos++
	// Server version [null terminated string]
	slice, err := p.readSlice(data[pos:], 0x00)
	if err != nil {
		return
	}
	p.serverVersion = string(slice)
	pos += len(slice) + 1
	// Thread id [32 bit uint]
	p.threadId = uint32(p.unpackNumber(data[pos : pos+4]))
	pos += 4
	// First part of scramble buffer [8 bytes]
	p.scrambleBuff = make([]byte, 8)
	p.scrambleBuff = data[pos : pos+8]
	pos += 9
	// Server capabilities [16 bit uint]
	p.serverCaps = uint16(p.unpackNumber(data[pos : pos+2]))
	pos += 2
	// Server language [8 bit uint]
	p.serverLanguage = data[pos]
	pos++
	// Server status [16 bit uint]
	p.serverStatus = uint16(p.unpackNumber(data[pos : pos+2]))
	pos += 15
	// Second part of scramble buffer, if exists (4.1+) [13 bytes]
	if ClientFlag(p.serverCaps)&CLIENT_PROTOCOL_41 > 0 {
		p.scrambleBuff = append(p.scrambleBuff, data[pos:pos+12]...)
	}
	return
}

// Auth packet
type packetAuth struct {
	packetBase
	clientFlags   uint32
	maxPacketSize uint32
	charsetNumber uint8
	user          string
	scrambleBuff  []byte
	database      string
}

// Auth packet writer
func (p *packetAuth) write() (data []byte, err os.Error) {
	// For MySQL 4.1+
	if p.protocol == PROTOCOL_41 {
		// Client flags
		data = p.packNumber(uint64(p.clientFlags), 4)
		// Max packet size
		data = append(data, p.packNumber(uint64(p.maxPacketSize), 4)...)
		// Charset
		data = append(data, p.charsetNumber)
		// Filler
		data = append(data, make([]byte, 23)...)
		// User
		if len(p.user) > 0 {
			data = append(data, []byte(p.user)...)
		}
		// Terminator
		data = append(data, 0x00)
		// Scramble buffer
		data = append(data, byte(len(p.scrambleBuff)))
		if len(p.scrambleBuff) > 0 {
			data = append(data, p.scrambleBuff...)
		}
		// Database name
		if len(p.database) > 0 {
			data = append(data, []byte(p.database)...)
			// Terminator
			data = append(data, 0x00)
		}
		// For MySQL < 4.1
	} else {
		// Client flags
		data = p.packNumber(uint64(p.clientFlags), 2)
		// Max packet size
		data = append(data, p.packNumber(uint64(p.maxPacketSize), 3)...)
		// User
		if len(p.user) > 0 {
			data = append(data, []byte(p.user)...)
		}
		// Terminator
		data = append(data, 0x00)
		// Scramble buffer
		if len(p.scrambleBuff) > 0 {
			data = append(data, p.scrambleBuff...)
		}
		// Padding
		data = append(data, 0x00)
	}
	// Add the packet header
	data = p.addHeader(data)
	return
}

// Ok packet struct
type packetOK struct {
	packetBase
	affectedRows uint64
	insertId     uint64
	serverStatus uint16
	warningCount uint16
	message      string
}

// OK packet reader
func (p *packetOK) read(data []byte) (err os.Error) {
	// Recover errors
	defer func() {
		if e := recover(); e != nil {
			err = os.NewError(fmt.Sprintf("%s", e))
		}
	}()
	// Position (skip first byte/field count)
	pos := 1
	// Affected rows [length coded binary]
	num, n, err := p.readLengthCodedBinary(data[pos:])
	if err != nil {
		return
	}
	p.affectedRows = num
	pos += n
	// Insert id [length coded binary]
	num, n, err = p.readLengthCodedBinary(data[pos:])
	if err != nil {
		return
	}
	p.insertId = num
	pos += n
	// Server status [16 bit uint]
	p.serverStatus = uint16(p.unpackNumber(data[pos : pos+2]))
	pos += 2
	// Warning (4.1 only) [16 bit uint]
	if p.protocol == PROTOCOL_41 {
		p.warningCount = uint16(p.unpackNumber(data[pos : pos+2]))
		pos += 2
	}
	// Message (optional) [string]
	if pos < len(data) {
		p.message = string(data[pos:])
	}
	return
}

// Error packet struct
type packetError struct {
	packetBase
	errno uint16
	state string
	error string
}

// Error packet reader
func (p *packetError) read(data []byte) (err os.Error) {
	// Recover errors
	defer func() {
		if e := recover(); e != nil {
			err = os.NewError(fmt.Sprintf("%s", e))
		}
	}()
	// Position
	pos := 1
	// Error number [16 bit uint]
	p.errno = uint16(p.unpackNumber(data[pos : pos+2]))
	pos += 2
	// State (4.1 only) [string]
	if p.protocol == PROTOCOL_41 {
		pos++
		p.state = string(data[pos : pos+5])
		pos += 5
	}
	// Message [string]
	p.error = string(data[pos:])
	return
}

// Command packet struct
type packetCommand struct {
	packetBase
	command command
	args    []interface{}
}

// Command packet writer
func (p *packetCommand) write() (data []byte, err os.Error) {
	// Recover errors (if wrong param type supplied)
	defer func() {
		if e := recover(); e != nil {
			err = os.NewError(fmt.Sprintf("%s", e))
		}
	}()
	// Make slice from command byte
	data = []byte{byte(p.command)}
	// Add args to requests
	switch p.command {
	// Commands with 1 arg unterminated string
	case COM_INIT_DB, COM_QUERY, COM_CREATE_DB, COM_DROP_DB, COM_STMT_PREPARE:
		data = append(data, []byte(p.args[0].(string))...)
	// Commands with 1 arg 32 bit uint
	case COM_PROCESS_KILL, COM_STMT_CLOSE, COM_STMT_RESET:
		data = append(data, p.packNumber(uint64(p.args[0].(uint32)), 4)...)
	// Field list command
	case COM_FIELD_LIST:
		// Table name
		data = append(data, []byte(p.args[0].(string))...)
		// Terminator
		data = append(data, 0x00)
		// Column name
		if len(p.args) > 1 {
			data = append(data, []byte(p.args[1].(string))...)
		}
	// Refresh command
	case COM_REFRESH:
		data = append(data, byte(p.args[0].(Refresh)))
	// Shutdown command
	case COM_SHUTDOWN:
		data = append(data, byte(p.args[0].(Shutdown)))
	// Change user command
	case COM_CHANGE_USER:
		// User
		data = append(data, []byte(p.args[0].(string))...)
		// Terminator
		data = append(data, 0x00)
		// Scramble length for 4.1
		if p.protocol == PROTOCOL_41 {
			data = append(data, byte(len(p.args[1].([]byte))))
		}
		// Scramble buffer
		if len(p.args[1].([]byte)) > 0 {
			data = append(data, p.args[1].([]byte)...)
		}
		// Temrminator for 3.23
		if p.protocol == PROTOCOL_40 {
			data = append(data, 0x00)
		}
		// Database name
		if len(p.args[2].(string)) > 0 {
			data = append(data, []byte(p.args[2].(string))...)
		}
		// Terminator
		data = append(data, 0x00)
		// Character set number (5.1.23+ needs testing with earlier versions)
		data = append(data, p.packNumber(uint64(p.args[3].(uint16)), 2)...)
	// Fetch statement command
	case COM_STMT_FETCH:
		// Statement id
		data = append(data, p.packNumber(uint64(p.args[0].(uint32)), 4)...)
		// Number of rows
		data = append(data, p.packNumber(uint64(p.args[1].(uint32)), 4)...)
	}
	// Add the packet header
	data = p.addHeader(data)
	return
}
<|endoftext|>


<commit_before>package nlopt

import "C"
import "unsafe"

//export nloptFunc
func nloptFunc(n uint, x *C.double, gradient *C.double, fData unsafe.Pointer) C.double {
	f := getFunc(uintptr(fData))
	cX := (*[1 << 30]C.double)(unsafe.Pointer(x))[:n:n]
	goX := toGoArray(cX)

	var goGrad []float64
	var cGrad []C.double
	if gradient != nil {
		cGrad = (*[1 << 30]C.double)(unsafe.Pointer(gradient))[:n:n]
		goGrad = toGoArray(cGrad)
	}
	v := (C.double)(f(goX, goGrad))
	if gradient != nil {
		for i := 0; i < len(cGrad); i++ {
			cGrad[i] = (C.double)((goGrad)[i])
		}
	}
	return v
}

//export nloptMfunc
func nloptMfunc(m uint, result *C.double, n uint, x *C.double, gradient *C.double, fData unsafe.Pointer) {
	f := getMfunc(uintptr(fData))
	cX := (*[1 << 30]C.double)(unsafe.Pointer(x))[:n:n]
	goX := toGoArray(cX)

	var goGrad []float64
	var cGrad []C.double
	if gradient != nil {
		cGrad = (*[1 << 30]C.double)(unsafe.Pointer(gradient))[:n:n]
		goGrad = toGoArray(cGrad)
	}

	var goResult []float64
	var cResult []C.double
	if result != nil {
		cResult = (*[1 << 30]C.double)(unsafe.Pointer(result))[: m*n : m*n]
		goResult = toGoArray(cResult)
	}

	f(goResult, goX, goGrad)

	if gradient != nil {
		for i := 0; i < len(cGrad); i++ {
			cGrad[i] = (C.double)((goGrad)[i])
		}
	}

	if result != nil {
		for i := 0; i < len(cResult); i++ {
			cResult[i] = (C.double)((goResult)[i])
		}
	}

}
<commit_msg>Workaround for Linux 32-bit<commit_after>package nlopt

import "C"
import "unsafe"

//export nloptFunc
func nloptFunc(n uint, x *C.double, gradient *C.double, fData unsafe.Pointer) C.double {
	f := getFunc(uintptr(fData))
	cX := (*[(1 << 29) - 1]C.double)(unsafe.Pointer(x))[:n:n]
	goX := toGoArray(cX)

	var goGrad []float64
	var cGrad []C.double
	if gradient != nil {
		cGrad = (*[(1 << 29) - 1]C.double)(unsafe.Pointer(gradient))[:n:n]
		goGrad = toGoArray(cGrad)
	}
	v := (C.double)(f(goX, goGrad))
	if gradient != nil {
		for i := 0; i < len(cGrad); i++ {
			cGrad[i] = (C.double)((goGrad)[i])
		}
	}
	return v
}

//export nloptMfunc
func nloptMfunc(m uint, result *C.double, n uint, x *C.double, gradient *C.double, fData unsafe.Pointer) {
	f := getMfunc(uintptr(fData))
	cX := (*[(1 << 29) - 1]C.double)(unsafe.Pointer(x))[:n:n]
	goX := toGoArray(cX)

	var goGrad []float64
	var cGrad []C.double
	if gradient != nil {
		cGrad = (*[(1 << 29) - 1]C.double)(unsafe.Pointer(gradient))[:n:n]
		goGrad = toGoArray(cGrad)
	}

	var goResult []float64
	var cResult []C.double
	if result != nil {
		cResult = (*[(1 << 29) - 1]C.double)(unsafe.Pointer(result))[: m*n : m*n]
		goResult = toGoArray(cResult)
	}

	f(goResult, goX, goGrad)

	if gradient != nil {
		for i := 0; i < len(cGrad); i++ {
			cGrad[i] = (C.double)((goGrad)[i])
		}
	}

	if result != nil {
		for i := 0; i < len(cResult); i++ {
			cResult[i] = (C.double)((goResult)[i])
		}
	}

}
<|endoftext|>


<commit_before>package yagnats

import (
	"bufio"
	"errors"
	"fmt"
	"regexp"
	"strconv"
)

type Parser func(*bufio.Reader) (Packet, error)

var PARSERS = map[string]Parser{
	// PING\s*\r\n
	"PING": func(io *bufio.Reader) (Packet, error) {
		io.ReadBytes('\n')
		return &PingPacket{}, nil
	},

	// PONG\s*\r\n
	"PONG": func(io *bufio.Reader) (Packet, error) {
		io.ReadBytes('\n')
		return &PongPacket{}, nil
	},

	// +OK\s*\r\n
	"+OK": func(io *bufio.Reader) (Packet, error) {
		io.ReadBytes('\n')
		return &OKPacket{}, nil
	},

	// -ERR '(message)'\r\n
	"-ERR": func(io *bufio.Reader) (Packet, error) {
		bytes, _ := io.ReadBytes('\n')
		re := regexp.MustCompile(`\s*'(.*)'\r\n`)
		match := re.FindSubmatchIndex(bytes)
		return &ERRPacket{Message: string(bytes[match[2]:match[3]])}, nil
	},

	// INFO (payload)\r\n
	"INFO": func(io *bufio.Reader) (Packet, error) {
		bytes, _ := io.ReadBytes('\n')
		re := regexp.MustCompile(`\s*([^\s]+)\r\n`)
		match := re.FindSubmatchIndex(bytes)
		return &InfoPacket{Payload: string(bytes[match[2]:match[3]])}, nil
	},

	// MSG (subject) (subscriber-id) (reply)? (length)\r\n(byte * length)\r\n
	"MSG": func(io *bufio.Reader) (Packet, error) {
		bytes, _ := io.ReadBytes('\n')
		re := regexp.MustCompile(`\s*([^\s]+)\s+(\d+)\s+(([^\s]+)\s+)?(\d+)\r\n`)
		matches := re.FindStringSubmatch(string(bytes))

		subID, _ := strconv.Atoi(matches[2])
		payloadLen, _ := strconv.Atoi(matches[5])
		payload := make([]byte, payloadLen)
		io.Read(payload)
		io.ReadBytes('\n')

		return &MsgPacket{
			Subject: matches[1],
			SubID:   subID,
			ReplyTo: matches[4],
			Payload: string(payload),
		}, nil
	},
}

func Parse(io *bufio.Reader) (val Packet, err error) {
	header, err := readWord(io)
	if err != nil {
		return nil, err
	}

	parser := PARSERS[string(header)]
	if parser == nil {
		return nil, errors.New(fmt.Sprintf("Unknown header: %s", header))
	}

	return parser(io)
}

func readWord(io *bufio.Reader) ([]byte, error) {
	length, err := wordEnd(io)
	if err != nil {
		return nil, err
	}

	word := make([]byte, length)
	io.Read(word)

	return word, nil
}

func wordEnd(io *bufio.Reader) (int, error) {
	length := 0

	for {
		next, err := io.Peek(length + 1)
		if err != nil {
			return length, err
		}

		switch next[length] {
		case ' ', '\r', '\t', '\n':
			return length, nil

		default:
			length += 1
		}
	}

	return length, nil
}
<commit_msg>simplify readWord using bytes.Buffer<commit_after>package yagnats

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"regexp"
	"strconv"
)

type Parser func(*bufio.Reader) (Packet, error)

var PARSERS = map[string]Parser{
	// PING\s*\r\n
	"PING": func(io *bufio.Reader) (Packet, error) {
		io.ReadBytes('\n')
		return &PingPacket{}, nil
	},

	// PONG\s*\r\n
	"PONG": func(io *bufio.Reader) (Packet, error) {
		io.ReadBytes('\n')
		return &PongPacket{}, nil
	},

	// +OK\s*\r\n
	"+OK": func(io *bufio.Reader) (Packet, error) {
		io.ReadBytes('\n')
		return &OKPacket{}, nil
	},

	// -ERR '(message)'\r\n
	"-ERR": func(io *bufio.Reader) (Packet, error) {
		bytes, _ := io.ReadBytes('\n')
		re := regexp.MustCompile(`\s*'(.*)'\r\n`)
		match := re.FindSubmatchIndex(bytes)
		return &ERRPacket{Message: string(bytes[match[2]:match[3]])}, nil
	},

	// INFO (payload)\r\n
	"INFO": func(io *bufio.Reader) (Packet, error) {
		bytes, _ := io.ReadBytes('\n')
		re := regexp.MustCompile(`\s*([^\s]+)\r\n`)
		match := re.FindSubmatchIndex(bytes)
		return &InfoPacket{Payload: string(bytes[match[2]:match[3]])}, nil
	},

	// MSG (subject) (subscriber-id) (reply)? (length)\r\n(byte * length)\r\n
	"MSG": func(io *bufio.Reader) (Packet, error) {
		bytes, _ := io.ReadBytes('\n')
		re := regexp.MustCompile(`\s*([^\s]+)\s+(\d+)\s+(([^\s]+)\s+)?(\d+)\r\n`)
		matches := re.FindStringSubmatch(string(bytes))

		subID, _ := strconv.Atoi(matches[2])
		payloadLen, _ := strconv.Atoi(matches[5])
		payload := make([]byte, payloadLen)
		io.Read(payload)
		io.ReadBytes('\n')

		return &MsgPacket{
			Subject: matches[1],
			SubID:   subID,
			ReplyTo: matches[4],
			Payload: string(payload),
		}, nil
	},
}

func Parse(io *bufio.Reader) (val Packet, err error) {
	header, err := readWord(io)
	if err != nil {
		return nil, err
	}

	parser := PARSERS[string(header)]
	if parser == nil {
		return nil, errors.New(fmt.Sprintf("Unknown header: %s", header))
	}

	return parser(io)
}

func readWord(io *bufio.Reader) ([]byte, error) {
	word := new(bytes.Buffer)

	for {
		next, err := io.ReadByte()
		if err != nil {
			return nil, err
		}

		switch next {
		case ' ', '\t', '\r', '\n':
			return word.Bytes(), nil

		default:
			word.WriteByte(next)
		}
	}
}
<|endoftext|>