<commit_before>package honeybadger

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net/http"
	"os"
	"runtime"
	"strings"

	_ "crypto/sha512"
)

type Notifier struct {
	Name     string `json:"name"`
	Url      string `json:"url"`
	Version  string `json:"version"`
	Language string `json:"language"`
}

type BacktraceLine struct {
	Method string `json:"method"`
	File   string `json:"file"`
	Number string `json:"number"`
}

type Error struct {
	Class     string                 `json:"class"`
	Message   string                 `json:"message"`
	Backtrace []*BacktraceLine       `json:"backtrace"`
	Source    map[string]interface{} `json:"source"`
}

type Request struct {
	Url       string                 `json:"url"`
	Component string                 `json:"component"`
	Action    string                 `json:"action"`
	Params    map[string]interface{} `json:"params"`
	Session   map[string]interface{} `json:"session"`
	CgiData   map[string]interface{} `json:"cgi_data"`
	Context   map[string]interface{} `json:"context"`
}

type Server struct {
	ProjectRoot     map[string]interface{} `json:"project_root"`
	EnvironmentName string                 `json:"environment_name"`
	Hostname        string                 `json:"hostname"`
}

type Report struct {
	Notifier *Notifier `json:"notifier"`
	Error    *Error    `json:"error"`
	Request  *Request  `json:"request"`
	Server   *Server   `json:"server"`
}

var (
	ApiKey      string
	Environment string
	Client      *http.Client = http.DefaultClient
	reports     chan *Report
)

func Start(apiKey, environment string) {
	ApiKey = apiKey
	Environment = environment

	if ApiKey == "" {
		log.Println("Honeybadger disabled")
		return
	}

	const numWorkers = 10
	const bufferSize = 1024

	reports = make(chan *Report, bufferSize)
	for i := 0; i < numWorkers; i++ {
		go consume()
	}

	log.Println("Honeybadger client started")
}

func consume() {
	for report := range reports {
		if e := report.Send(); e != nil {
			log.Printf("honeybadger: could not send report: %s", e)
		}
	}
}

// Send sends a report to HB synchronously and returns the error
func Send(err error) error {
	if ApiKey == "" {
		return err
	}
	report, e := NewReport(err)
	if e != nil {
		log.Printf("honeybadger: could not create report: %s", e)
		return err
	}
	if e := report.Send(); e != nil {
		log.Printf("honeybadger: could not send report: %s", e)
	}
	return err
}

// Dispatch sends a report to HB asynchronously and returns the error
func Dispatch(err error) error {
	return DispatchWithContext(err, nil)
}

// DispatchWithContext sends a report to HB asynchronously with context and
// returns the error
func DispatchWithContext(err error, context map[string]interface{}) error {
	if ApiKey == "" {
		return err
	}
	report, rErr := NewReport(err)
	if rErr != nil {
		log.Printf("honeybadger: could not create report: %s", rErr)
		return err
	}
	if context != nil {
		for k, v := range context {
			report.AddContext(k, v)
		}
	}

	select {
	case reports <- report:
	default:
		log.Printf("honeybadger: queue is full, dropping on the floor: %s", err)
	}
	return err
}

// Create a new report using the given error message and current call stack.
func NewReport(msg interface{}) (r *Report, err error) {
	return NewReportWithSkipCallers(msg, 0)
}

func fullMessage(msg interface{}) string {
	return fmt.Sprintf("%s", msg)
}

func exceptionClass(message string) string {
	pieces := strings.Split(message, ":")
	for i := len(pieces) - 1; i >= 0; i-- {
		err := strings.TrimSpace(pieces[i])
		if len(err) > 0 {
			return err
		}
	}
	return ""
}

// Create a new report using the given error message and current call stack.
// Supply an integer indicating how many callers to skip (0 is none).
func NewReportWithSkipCallers(msg interface{}, skipCallers int) (r *Report, err error) {
	var cwd, hostname string

	if ApiKey == "" {
		err = errors.New("You must set an API key first")
		return
	}

	if Environment == "" {
		err = errors.New("You must set an environment first")
		return
	}

	cwd, _ = os.Getwd()
	hostname, _ = os.Hostname()

	fullMessage := fullMessage(msg)
	exceptionClass := exceptionClass(fullMessage)

	r = &Report{
		Notifier: &Notifier{
			Name:     "Honeybadger (Go)",
			Url:      "https://github.com/librato/honeybadger-go",
			Version:  "1.0",
			Language: "Go",
		},
		Error: &Error{
			Class:     exceptionClass,
			Message:   fullMessage,
			Backtrace: make([]*BacktraceLine, 0),
			Source:    make(map[string]interface{}),
		},
		Request: &Request{
			Params:  make(map[string]interface{}),
			Session: make(map[string]interface{}),
			CgiData: make(map[string]interface{}),
			Context: make(map[string]interface{}),
		},
		Server: &Server{
			ProjectRoot:     make(map[string]interface{}),
			EnvironmentName: Environment,
			Hostname:        hostname,
		},
	}

	callers := make([]uintptr, 10)
	runtime.Callers(0, callers)

	for i, pc := range callers {
		if i < 3 { // skip self-inflicted depth
			continue
		}

		fc := runtime.FuncForPC(pc)
		if fc != nil {
			ms := strings.Split(fc.Name(), "/")
			m := ms[len(ms)-1]
			f, n := fc.FileLine(pc)

			// If we reach our origin depth, use the given function name as error class
			if i == skipCallers+2 {
				r.Error.Class = m
			}

			r.Error.Backtrace = append(r.Error.Backtrace, &BacktraceLine{
				Method: m,
				File:   f,
				Number: fmt.Sprintf("%d", n),
			})
		}
	}

	r.Server.ProjectRoot["path"] = cwd

	return
}

// Add a key and given value to the report as context
func (r *Report) AddContext(k string, v interface{}) {
	r.Request.Context[k] = v
}

// Add a key and given value to the report as parameters
func (r *Report) AddParam(k string, v interface{}) {
	r.Request.Params[k] = v
}

// Add a key and given value to the report as session
func (r *Report) AddSession(k string, v interface{}) {
	r.Request.Session[k] = v
}

// Send the report asynchronously
func (r *Report) Dispatch() {
	go func() {
		if err := r.Send(); err != nil {
			log.Printf("honeybadger: could not send report: %s", err)
		}
	}()
}

// Send the report and return an error if present
func (r *Report) Send() (err error) {
	var req *http.Request
	var resp *http.Response
	var payload []byte

	if payload, err = json.MarshalIndent(r, "", "  "); err != nil {
		return
	}

	if req, err = http.NewRequest("POST", "https://api.honeybadger.io/v1/notices", bytes.NewBuffer(payload)); err != nil {
		return
	}

	req.Header.Add("Content-Type", "application/json")
	req.Header.Add("Accept", "application/json")
	req.Header.Add("X-API-Key", ApiKey)

	if resp, err = Client.Do(req); err != nil {
		return
	}
	defer resp.Body.Close()

	if resp.StatusCode < 201 {
		err = errors.New(fmt.Sprintf("unable to send: error %d", resp.StatusCode))
	}
	return
}
<commit_msg>Fix user agent<commit_after>package honeybadger

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net/http"
	"os"
	"runtime"
	"strings"

	_ "crypto/sha512"
)

type Notifier struct {
	Name     string `json:"name"`
	Url      string `json:"url"`
	Version  string `json:"version"`
	Language string `json:"language"`
}

type BacktraceLine struct {
	Method string `json:"method"`
	File   string `json:"file"`
	Number string `json:"number"`
}

type Error struct {
	Class     string                 `json:"class"`
	Message   string                 `json:"message"`
	Backtrace []*BacktraceLine       `json:"backtrace"`
	Source    map[string]interface{} `json:"source"`
}

type Request struct {
	Url       string                 `json:"url"`
	Component string                 `json:"component"`
	Action    string                 `json:"action"`
	Params    map[string]interface{} `json:"params"`
	Session   map[string]interface{} `json:"session"`
	CgiData   map[string]interface{} `json:"cgi_data"`
	Context   map[string]interface{} `json:"context"`
}

type Server struct {
	ProjectRoot     map[string]interface{} `json:"project_root"`
	EnvironmentName string                 `json:"environment_name"`
	Hostname        string                 `json:"hostname"`
}

type Report struct {
	Notifier *Notifier `json:"notifier"`
	Error    *Error    `json:"error"`
	Request  *Request  `json:"request"`
	Server   *Server   `json:"server"`
}

var (
	ApiKey      string
	Environment string
	Client      *http.Client = http.DefaultClient
	reports     chan *Report
)

func Start(apiKey, environment string) {
	ApiKey = apiKey
	Environment = environment

	if ApiKey == "" {
		log.Println("Honeybadger disabled")
		return
	}

	const numWorkers = 10
	const bufferSize = 1024

	reports = make(chan *Report, bufferSize)
	for i := 0; i < numWorkers; i++ {
		go consume()
	}

	log.Println("Honeybadger client started")
}

func consume() {
	for report := range reports {
		if e := report.Send(); e != nil {
			log.Printf("honeybadger: could not send report: %s", e)
		}
	}
}

// Send sends a report to HB synchronously and returns the error
func Send(err error) error {
	if ApiKey == "" {
		return err
	}
	report, e := NewReport(err)
	if e != nil {
		log.Printf("honeybadger: could not create report: %s", e)
		return err
	}
	if e := report.Send(); e != nil {
		log.Printf("honeybadger: could not send report: %s", e)
	}
	return err
}

// Dispatch sends a report to HB asynchronously and returns the error
func Dispatch(err error) error {
	return DispatchWithContext(err, nil)
}

// DispatchWithContext sends a report to HB asynchronously with context and
// returns the error
func DispatchWithContext(err error, context map[string]interface{}) error {
	if ApiKey == "" {
		return err
	}
	report, rErr := NewReport(err)
	if rErr != nil {
		log.Printf("honeybadger: could not create report: %s", rErr)
		return err
	}
	if context != nil {
		for k, v := range context {
			report.AddContext(k, v)
		}
	}

	select {
	case reports <- report:
	default:
		log.Printf("honeybadger: queue is full, dropping on the floor: %s", err)
	}
	return err
}

// Create a new report using the given error message and current call stack.
func NewReport(msg interface{}) (r *Report, err error) {
	return NewReportWithSkipCallers(msg, 0)
}

func fullMessage(msg interface{}) string {
	return fmt.Sprintf("%s", msg)
}

func exceptionClass(message string) string {
	pieces := strings.Split(message, ":")
	for i := len(pieces) - 1; i >= 0; i-- {
		err := strings.TrimSpace(pieces[i])
		if len(err) > 0 {
			return err
		}
	}
	return ""
}

// Create a new report using the given error message and current call stack.
// Supply an integer indicating how many callers to skip (0 is none).
func NewReportWithSkipCallers(msg interface{}, skipCallers int) (r *Report, err error) {
	var cwd, hostname string

	if ApiKey == "" {
		err = errors.New("You must set an API key first")
		return
	}

	if Environment == "" {
		err = errors.New("You must set an environment first")
		return
	}

	cwd, _ = os.Getwd()
	hostname, _ = os.Hostname()

	fullMessage := fullMessage(msg)
	exceptionClass := exceptionClass(fullMessage)

	r = &Report{
		Notifier: &Notifier{
			Name:     "Honeybadger (Go)",
			Url:      "https://github.com/jcoene/honeybadger-go",
			Version:  "1.0",
			Language: "Go",
		},
		Error: &Error{
			Class:     exceptionClass,
			Message:   fullMessage,
			Backtrace: make([]*BacktraceLine, 0),
			Source:    make(map[string]interface{}),
		},
		Request: &Request{
			Params:  make(map[string]interface{}),
			Session: make(map[string]interface{}),
			CgiData: make(map[string]interface{}),
			Context: make(map[string]interface{}),
		},
		Server: &Server{
			ProjectRoot:     make(map[string]interface{}),
			EnvironmentName: Environment,
			Hostname:        hostname,
		},
	}

	callers := make([]uintptr, 10)
	runtime.Callers(0, callers)

	for i, pc := range callers {
		if i < 3 { // skip self-inflicted depth
			continue
		}

		fc := runtime.FuncForPC(pc)
		if fc != nil {
			ms := strings.Split(fc.Name(), "/")
			m := ms[len(ms)-1]
			f, n := fc.FileLine(pc)

			// If we reach our origin depth, use the given function name as error class
			if i == skipCallers+2 {
				r.Error.Class = m
			}

			r.Error.Backtrace = append(r.Error.Backtrace, &BacktraceLine{
				Method: m,
				File:   f,
				Number: fmt.Sprintf("%d", n),
			})
		}
	}

	r.Server.ProjectRoot["path"] = cwd

	return
}

// Add a key and given value to the report as context
func (r *Report) AddContext(k string, v interface{}) {
	r.Request.Context[k] = v
}

// Add a key and given value to the report as parameters
func (r *Report) AddParam(k string, v interface{}) {
	r.Request.Params[k] = v
}

// Add a key and given value to the report as session
func (r *Report) AddSession(k string, v interface{}) {
	r.Request.Session[k] = v
}

// Send the report asynchronously
func (r *Report) Dispatch() {
	go func() {
		if err := r.Send(); err != nil {
			log.Printf("honeybadger: could not send report: %s", err)
		}
	}()
}

// Send the report and return an error if present
func (r *Report) Send() (err error) {
	var req *http.Request
	var resp *http.Response
	var payload []byte

	if payload, err = json.MarshalIndent(r, "", "  "); err != nil {
		return
	}

	if req, err = http.NewRequest("POST", "https://api.honeybadger.io/v1/notices", bytes.NewBuffer(payload)); err != nil {
		return
	}

	req.Header.Add("Content-Type", "application/json")
	req.Header.Add("Accept", "application/json")
	req.Header.Add("X-API-Key", ApiKey)

	if resp, err = Client.Do(req); err != nil {
		return
	}
	defer resp.Body.Close()

	if resp.StatusCode < 201 {
		err = errors.New(fmt.Sprintf("unable to send: error %d", resp.StatusCode))
	}
	return
}
<|endoftext|>


<commit_before>package honeybadger

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net/http"
	"os"
	"runtime"
	"strings"

	_ "crypto/sha512"
)

type Notifier struct {
	Name     string `json:"name"`
	Url      string `json:"url"`
	Version  string `json:"version"`
	Language string `json:"language"`
}

type BacktraceLine struct {
	Method string `json:"method"`
	File   string `json:"file"`
	Number string `json:"number"`
}

type Error struct {
	Class     string                 `json:"class"`
	Message   string                 `json:"message"`
	Backtrace []*BacktraceLine       `json:"backtrace"`
	Source    map[string]interface{} `json:"source"`
}

type Request struct {
	Url       string                 `json:"url"`
	Component string                 `json:"component"`
	Action    string                 `json:"action"`
	Params    map[string]interface{} `json:"params"`
	Session   map[string]interface{} `json:"session"`
	CgiData   map[string]interface{} `json:"cgi_data"`
	Context   map[string]interface{} `json:"context"`
}

type Server struct {
	ProjectRoot     map[string]interface{} `json:"project_root"`
	EnvironmentName string                 `json:"environment_name"`
	Hostname        string                 `json:"hostname"`
}

type Report struct {
	Notifier *Notifier `json:"notifier"`
	Error    *Error    `json:"error"`
	Request  *Request  `json:"request"`
	Server   *Server   `json:"server"`
}

var (
	ApiKey      string
	Environment string
	reports     chan *Report
)

func Start(apiKey, environment string) {
	ApiKey = apiKey
	Environment = environment

	if ApiKey == "" {
		log.Println("Honeybadger disabled")
		return
	}

	const numWorkers = 10
	const bufferSize = 1024

	reports = make(chan *Report, bufferSize)
	for i := 0; i < numWorkers; i++ {
		go consume()
	}

	log.Println("Honeybadger client started")
}

func consume() {
	for report := range reports {
		if e := report.Send(); e != nil {
			log.Printf("honeybadger: could not send report: %s", e)
		}
	}
}

// Send sends a report to HB synchronously and returns the error
func Send(err error) error {
	if ApiKey == "" {
		return err
	}
	report, e := NewReport(err)
	if e != nil {
		log.Printf("honeybadger: could not create report: %s", e)
		return err
	}
	if e := report.Send(); e != nil {
		log.Printf("honeybadger: could not send report: %s", e)
	}
	return err
}

// Dispatch sends a report to HB asynchronously and returns the error
func Dispatch(err error) error {
	return DispatchWithContext(err, nil)
}

// DispatchWithContext sends a report to HB asynchronously with context and
// returns the error
func DispatchWithContext(err error, context map[string]interface{}) error {
	if ApiKey == "" {
		return err
	}
	report, rErr := NewReport(err)
	if rErr != nil {
		log.Printf("honeybadger: could not create report: %s", rErr)
		return err
	}
	if context != nil {
		for k, v := range context {
			report.AddContext(k, v)
		}
	}

	select {
	case reports <- report:
	default:
		log.Printf("honeybadger: queue is full, dropping on the floor: %s", err)
	}
	return err
}

// Create a new report using the given error message and current call stack.
func NewReport(msg interface{}) (r *Report, err error) {
	return NewReportWithSkipCallers(msg, 0)
}

func fullMessage(msg interface{}) string {
	return fmt.Sprintf("%s", msg)
}

func exceptionClass(message string) string {
	pieces := strings.Split(message, ":")
	for i := len(pieces)-1; i >= 0; i-- {
		err := strings.TrimSpace(pieces[i])
		if len(err) > 0 {
			return err
		}
	}
	return ""
}

// Create a new report using the given error message and current call stack.
// Supply an integer indicating how many callers to skip (0 is none).
func NewReportWithSkipCallers(msg interface{}, skipCallers int) (r *Report, err error) {
	var cwd, hostname string

	if ApiKey == "" {
		err = errors.New("You must set an API key first")
		return
	}

	if Environment == "" {
		err = errors.New("You must set an environment first")
		return
	}

	cwd, _ = os.Getwd()
	hostname, _ = os.Hostname()

	fullMessage := fullMessage(msg)
	exceptionClass := exceptionClass(fullMessage)

	r = &Report{
		Notifier: &Notifier{
			Name:     "Honeybadger (Go)",
			Url:      "https://github.com/librato/honeybadger-go",
			Version:  "1.0",
			Language: "Go",
		},
		Error: &Error{
			Class:     exceptionClass,
			Message:   fullMessage,
			Backtrace: make([]*BacktraceLine, 0),
			Source:    make(map[string]interface{}),
		},
		Request: &Request{
			Params:  make(map[string]interface{}),
			Session: make(map[string]interface{}),
			CgiData: make(map[string]interface{}),
			Context: make(map[string]interface{}),
		},
		Server: &Server{
			ProjectRoot:     make(map[string]interface{}),
			EnvironmentName: Environment,
			Hostname:        hostname,
		},
	}

	callers := make([]uintptr, 10)
	runtime.Callers(0, callers)

	for i, pc := range callers {
		if i < 3 { // skip self-inflicted depth
			continue
		}

		fc := runtime.FuncForPC(pc)
		if fc != nil {
			ms := strings.Split(fc.Name(), "/")
			m := ms[len(ms)-1]
			f, n := fc.FileLine(pc)

			// If we reach our origin depth, use the given function name as error class
			if i == skipCallers+2 {
				r.Error.Class = m
			}

			r.Error.Backtrace = append(r.Error.Backtrace, &BacktraceLine{
				Method: m,
				File:   f,
				Number: fmt.Sprintf("%d", n),
			})
		}
	}

	r.Server.ProjectRoot["path"] = cwd

	return
}

// Add a key and given value to the report as context
func (r *Report) AddContext(k string, v interface{}) {
	r.Request.Context[k] = v
}

// Add a key and given value to the report as parameters
func (r *Report) AddParam(k string, v interface{}) {
	r.Request.Params[k] = v
}

// Add a key and given value to the report as session
func (r *Report) AddSession(k string, v interface{}) {
	r.Request.Session[k] = v
}

// Send the report asynchronously
func (r *Report) Dispatch() {
	go func() {
		if err := r.Send(); err != nil {
			log.Printf("honeybadger: could not send report: %s", err)
		}
	}()
}

// Send the report and return an error if present
func (r *Report) Send() (err error) {
	var req *http.Request
	var resp *http.Response
	var payload []byte

	if payload, err = json.MarshalIndent(r, "", "  "); err != nil {
		return
	}

	if req, err = http.NewRequest("POST", "https://api.honeybadger.io/v1/notices", bytes.NewBuffer(payload)); err != nil {
		return
	}

	req.Header.Add("Content-Type", "application/json")
	req.Header.Add("Accept", "application/json")
	req.Header.Add("X-API-Key", ApiKey)

	if resp, err = http.DefaultClient.Do(req); err != nil {
		return
	}
	defer resp.Body.Close()

	if resp.StatusCode < 201 {
		err = errors.New(fmt.Sprintf("unable to send: error %d", resp.StatusCode))
	}
	return
}
<commit_msg>Allow user to specify their own http client.<commit_after>package honeybadger

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net/http"
	"os"
	"runtime"
	"strings"

	_ "crypto/sha512"
)

type Notifier struct {
	Name     string `json:"name"`
	Url      string `json:"url"`
	Version  string `json:"version"`
	Language string `json:"language"`
}

type BacktraceLine struct {
	Method string `json:"method"`
	File   string `json:"file"`
	Number string `json:"number"`
}

type Error struct {
	Class     string                 `json:"class"`
	Message   string                 `json:"message"`
	Backtrace []*BacktraceLine       `json:"backtrace"`
	Source    map[string]interface{} `json:"source"`
}

type Request struct {
	Url       string                 `json:"url"`
	Component string                 `json:"component"`
	Action    string                 `json:"action"`
	Params    map[string]interface{} `json:"params"`
	Session   map[string]interface{} `json:"session"`
	CgiData   map[string]interface{} `json:"cgi_data"`
	Context   map[string]interface{} `json:"context"`
}

type Server struct {
	ProjectRoot     map[string]interface{} `json:"project_root"`
	EnvironmentName string                 `json:"environment_name"`
	Hostname        string                 `json:"hostname"`
}

type Report struct {
	Notifier *Notifier `json:"notifier"`
	Error    *Error    `json:"error"`
	Request  *Request  `json:"request"`
	Server   *Server   `json:"server"`
}

var (
	ApiKey      string
	Environment string
	Client      *http.Client = http.DefaultClient
	reports     chan *Report
)

func Start(apiKey, environment string) {
	ApiKey = apiKey
	Environment = environment

	if ApiKey == "" {
		log.Println("Honeybadger disabled")
		return
	}

	const numWorkers = 10
	const bufferSize = 1024

	reports = make(chan *Report, bufferSize)
	for i := 0; i < numWorkers; i++ {
		go consume()
	}

	log.Println("Honeybadger client started")
}

func consume() {
	for report := range reports {
		if e := report.Send(); e != nil {
			log.Printf("honeybadger: could not send report: %s", e)
		}
	}
}

// Send sends a report to HB synchronously and returns the error
func Send(err error) error {
	if ApiKey == "" {
		return err
	}
	report, e := NewReport(err)
	if e != nil {
		log.Printf("honeybadger: could not create report: %s", e)
		return err
	}
	if e := report.Send(); e != nil {
		log.Printf("honeybadger: could not send report: %s", e)
	}
	return err
}

// Dispatch sends a report to HB asynchronously and returns the error
func Dispatch(err error) error {
	return DispatchWithContext(err, nil)
}

// DispatchWithContext sends a report to HB asynchronously with context and
// returns the error
func DispatchWithContext(err error, context map[string]interface{}) error {
	if ApiKey == "" {
		return err
	}
	report, rErr := NewReport(err)
	if rErr != nil {
		log.Printf("honeybadger: could not create report: %s", rErr)
		return err
	}
	if context != nil {
		for k, v := range context {
			report.AddContext(k, v)
		}
	}

	select {
	case reports <- report:
	default:
		log.Printf("honeybadger: queue is full, dropping on the floor: %s", err)
	}
	return err
}

// Create a new report using the given error message and current call stack.
func NewReport(msg interface{}) (r *Report, err error) {
	return NewReportWithSkipCallers(msg, 0)
}

func fullMessage(msg interface{}) string {
	return fmt.Sprintf("%s", msg)
}

func exceptionClass(message string) string {
	pieces := strings.Split(message, ":")
	for i := len(pieces)-1; i >= 0; i-- {
		err := strings.TrimSpace(pieces[i])
		if len(err) > 0 {
			return err
		}
	}
	return ""
}

// Create a new report using the given error message and current call stack.
// Supply an integer indicating how many callers to skip (0 is none).
func NewReportWithSkipCallers(msg interface{}, skipCallers int) (r *Report, err error) {
	var cwd, hostname string

	if ApiKey == "" {
		err = errors.New("You must set an API key first")
		return
	}

	if Environment == "" {
		err = errors.New("You must set an environment first")
		return
	}

	cwd, _ = os.Getwd()
	hostname, _ = os.Hostname()

	fullMessage := fullMessage(msg)
	exceptionClass := exceptionClass(fullMessage)

	r = &Report{
		Notifier: &Notifier{
			Name:     "Honeybadger (Go)",
			Url:      "https://github.com/librato/honeybadger-go",
			Version:  "1.0",
			Language: "Go",
		},
		Error: &Error{
			Class:     exceptionClass,
			Message:   fullMessage,
			Backtrace: make([]*BacktraceLine, 0),
			Source:    make(map[string]interface{}),
		},
		Request: &Request{
			Params:  make(map[string]interface{}),
			Session: make(map[string]interface{}),
			CgiData: make(map[string]interface{}),
			Context: make(map[string]interface{}),
		},
		Server: &Server{
			ProjectRoot:     make(map[string]interface{}),
			EnvironmentName: Environment,
			Hostname:        hostname,
		},
	}

	callers := make([]uintptr, 10)
	runtime.Callers(0, callers)

	for i, pc := range callers {
		if i < 3 { // skip self-inflicted depth
			continue
		}

		fc := runtime.FuncForPC(pc)
		if fc != nil {
			ms := strings.Split(fc.Name(), "/")
			m := ms[len(ms)-1]
			f, n := fc.FileLine(pc)

			// If we reach our origin depth, use the given function name as error class
			if i == skipCallers+2 {
				r.Error.Class = m
			}

			r.Error.Backtrace = append(r.Error.Backtrace, &BacktraceLine{
				Method: m,
				File:   f,
				Number: fmt.Sprintf("%d", n),
			})
		}
	}

	r.Server.ProjectRoot["path"] = cwd

	return
}

// Add a key and given value to the report as context
func (r *Report) AddContext(k string, v interface{}) {
	r.Request.Context[k] = v
}

// Add a key and given value to the report as parameters
func (r *Report) AddParam(k string, v interface{}) {
	r.Request.Params[k] = v
}

// Add a key and given value to the report as session
func (r *Report) AddSession(k string, v interface{}) {
	r.Request.Session[k] = v
}

// Send the report asynchronously
func (r *Report) Dispatch() {
	go func() {
		if err := r.Send(); err != nil {
			log.Printf("honeybadger: could not send report: %s", err)
		}
	}()
}

// Send the report and return an error if present
func (r *Report) Send() (err error) {
	var req *http.Request
	var resp *http.Response
	var payload []byte

	if payload, err = json.MarshalIndent(r, "", "  "); err != nil {
		return
	}

	if req, err = http.NewRequest("POST", "https://api.honeybadger.io/v1/notices", bytes.NewBuffer(payload)); err != nil {
		return
	}

	req.Header.Add("Content-Type", "application/json")
	req.Header.Add("Accept", "application/json")
	req.Header.Add("X-API-Key", ApiKey)

	if resp, err = Client.Do(req); err != nil {
		return
	}
	defer resp.Body.Close()

	if resp.StatusCode < 201 {
		err = errors.New(fmt.Sprintf("unable to send: error %d", resp.StatusCode))
	}
	return
}
<|endoftext|>


<commit_before>// Copyright 2016 The Cockroach 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.
//
// Author: Spencer Kimball (spencer@cockroachlabs.com)

package crdb

import (
	"database/sql"
	"fmt"
	"sync"
	"testing"

	"github.com/cockroachdb/cockroach-go/testserver"
)

func TestExecuteTx(t *testing.T) {
	executeTx := func(db *sql.DB, fn func(tx *sql.Tx) error) (err error) { return ExecuteTx(db, fn) }
	testExecuteTxInner(t, executeTx)
}

func TestExecuteInTx(t *testing.T) {
	executeTx := func(db *sql.DB, fn func(tx *sql.Tx) error) (err error) {
		tx, err := db.Begin()
		if err != nil {
			return err
		}
		return ExecuteInTx(tx, func() error { return fn(tx) })
	}
	testExecuteTxInner(t, executeTx)
}

// TestExecuteTx verifies transaction retry using the classic
// example of write skew in bank account balance transfers.
func testExecuteTxInner(t *testing.T, executeTxFn func(db *sql.DB, fn func(tx *sql.Tx) error) error) {
	db, stop := testserver.NewDBForTest(t)
	defer stop()

	initStmt := `
CREATE DATABASE d;
CREATE TABLE d.t (acct INT PRIMARY KEY, balance INT);
INSERT INTO d.t (acct, balance) VALUES (1, 100), (2, 100);
`
	if _, err := db.Exec(initStmt); err != nil {
		t.Fatal(err)
	}

	type queryI interface {
		Query(string, ...interface{}) (*sql.Rows, error)
	}

	getBalances := func(q queryI) (bal1, bal2 int, err error) {
		var rows *sql.Rows
		rows, err = q.Query(`SELECT balance FROM d.t WHERE acct IN (1, 2);`)
		if err != nil {
			return
		}
		defer rows.Close()
		balances := []*int{&bal1, &bal2}
		i := 0
		for ; rows.Next(); i += 1 {
			if err = rows.Scan(balances[i]); err != nil {
				return
			}
		}
		if i != 2 {
			err = fmt.Errorf("expected two balances; got %d", i)
			return
		}
		return
	}

	runTxn := func(wg *sync.WaitGroup, iter *int) <-chan error {
		errCh := make(chan error, 1)
		go func() {
			*iter = 0
			errCh <- executeTxFn(db, func(tx *sql.Tx) error {
				*iter++
				bal1, bal2, err := getBalances(tx)
				if err != nil {
					return err
				}
				// If this is the first iteration, wait for the other tx to also read.
				if *iter == 1 {
					wg.Done()
					wg.Wait()
				}
				// Now, subtract from one account and give to the other.
				if bal1 > bal2 {
					if _, err := tx.Exec(`
UPDATE d.t SET balance=balance-100 WHERE acct=1;
UPDATE d.t SET balance=balance+100 WHERE acct=2;
`); err != nil {
						return err
					}
				} else {
					if _, err := tx.Exec(`
UPDATE d.t SET balance=balance+100 WHERE acct=1;
UPDATE d.t SET balance=balance-100 WHERE acct=2;
`); err != nil {
						return err
					}
				}
				return nil
			})
		}()
		return errCh
	}

	var wg sync.WaitGroup
	wg.Add(2)
	var iters1, iters2 int
	txn1Err := runTxn(&wg, &iters1)
	txn2Err := runTxn(&wg, &iters2)
	if err := <-txn1Err; err != nil {
		t.Errorf("expected success in txn1; got %s", err)
	}
	if err := <-txn2Err; err != nil {
		t.Errorf("expected success in txn2; got %s", err)
	}
	if iters1+iters2 <= 2 {
		t.Errorf("expected at least one retry between the competing transactions; "+
			"got txn1=%d, txn2=%d", iters1, iters2)
	}
	bal1, bal2, err := getBalances(db)
	if err != nil || bal1 != 100 || bal2 != 100 {
		t.Errorf("expected balances to be restored without error; "+
			"got acct1=%d, acct2=%d: %s", bal1, bal2, err)
	}
}
<commit_msg>Simplify<commit_after>// Copyright 2016 The Cockroach 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.
//
// Author: Spencer Kimball (spencer@cockroachlabs.com)

package crdb

import (
	"database/sql"
	"fmt"
	"sync"
	"testing"

	"github.com/cockroachdb/cockroach-go/testserver"
)

func TestExecuteTx(t *testing.T) {
	testExecuteTxInner(t, ExecuteTx)
}

func TestExecuteInTx(t *testing.T) {
	executeTx := func(db *sql.DB, fn func(tx *sql.Tx) error) (err error) {
		tx, err := db.Begin()
		if err != nil {
			return err
		}
		return ExecuteInTx(tx, func() error { return fn(tx) })
	}
	testExecuteTxInner(t, executeTx)
}

// TestExecuteTx verifies transaction retry using the classic
// example of write skew in bank account balance transfers.
func testExecuteTxInner(t *testing.T, executeTxFn func(db *sql.DB, fn func(tx *sql.Tx) error) error) {
	db, stop := testserver.NewDBForTest(t)
	defer stop()

	initStmt := `
CREATE DATABASE d;
CREATE TABLE d.t (acct INT PRIMARY KEY, balance INT);
INSERT INTO d.t (acct, balance) VALUES (1, 100), (2, 100);
`
	if _, err := db.Exec(initStmt); err != nil {
		t.Fatal(err)
	}

	type queryI interface {
		Query(string, ...interface{}) (*sql.Rows, error)
	}

	getBalances := func(q queryI) (bal1, bal2 int, err error) {
		var rows *sql.Rows
		rows, err = q.Query(`SELECT balance FROM d.t WHERE acct IN (1, 2);`)
		if err != nil {
			return
		}
		defer rows.Close()
		balances := []*int{&bal1, &bal2}
		i := 0
		for ; rows.Next(); i++ {
			if err = rows.Scan(balances[i]); err != nil {
				return
			}
		}
		if i != 2 {
			err = fmt.Errorf("expected two balances; got %d", i)
			return
		}
		return
	}

	runTxn := func(wg *sync.WaitGroup, iter *int) <-chan error {
		errCh := make(chan error, 1)
		go func() {
			*iter = 0
			errCh <- executeTxFn(db, func(tx *sql.Tx) error {
				*iter++
				bal1, bal2, err := getBalances(tx)
				if err != nil {
					return err
				}
				// If this is the first iteration, wait for the other tx to also read.
				if *iter == 1 {
					wg.Done()
					wg.Wait()
				}
				// Now, subtract from one account and give to the other.
				if bal1 > bal2 {
					if _, err := tx.Exec(`
UPDATE d.t SET balance=balance-100 WHERE acct=1;
UPDATE d.t SET balance=balance+100 WHERE acct=2;
`); err != nil {
						return err
					}
				} else {
					if _, err := tx.Exec(`
UPDATE d.t SET balance=balance+100 WHERE acct=1;
UPDATE d.t SET balance=balance-100 WHERE acct=2;
`); err != nil {
						return err
					}
				}
				return nil
			})
		}()
		return errCh
	}

	var wg sync.WaitGroup
	wg.Add(2)
	var iters1, iters2 int
	txn1Err := runTxn(&wg, &iters1)
	txn2Err := runTxn(&wg, &iters2)
	if err := <-txn1Err; err != nil {
		t.Errorf("expected success in txn1; got %s", err)
	}
	if err := <-txn2Err; err != nil {
		t.Errorf("expected success in txn2; got %s", err)
	}
	if iters1+iters2 <= 2 {
		t.Errorf("expected at least one retry between the competing transactions; "+
			"got txn1=%d, txn2=%d", iters1, iters2)
	}
	bal1, bal2, err := getBalances(db)
	if err != nil || bal1 != 100 || bal2 != 100 {
		t.Errorf("expected balances to be restored without error; "+
			"got acct1=%d, acct2=%d: %s", bal1, bal2, err)
	}
}
<|endoftext|>


<commit_before>package client

import (
	"encoding/json"
	"fmt"
	"net/http"
	"sync"
	"time"

	"golang.org/x/oauth2"
)

type client struct {
	// agent is the client's User-Agent in http requests.
	agent string
	// id is the bot's OAuth2 client id.
	id string
	// secret is the bot's OAuth2 client secret.
	secret string
	// user is the bot's username on reddit.
	user string
	// pass is the bot's password on reddit.
	pass string

	// authMu protects authentication fields.
	authMu sync.Mutex
	// cli is the authenticated client to execute requests with.
	cli *http.Client
	// token is the OAuth2 token cli uses to authenticate.
	token *oauth2.Token

	// rateMu protects rate limiting fields.
	rateMu sync.Mutex
	// nextReq is the time at which it is ok to make the next request.
	nextReq time.Time
}

// Do wraps the execution of http requests. It updates authentications and rate
// limits requests to Reddit to comply with the API rules.
func (c *client) Do(r *http.Request, out interface{}) error {
	if time.Now().Before(c.nextReq) {
		<-time.After(c.nextReq.Sub(time.Now()))
	}
	c.nextReq = time.Now().Add(time.Second)
	if !c.token.Valid() {
		var err error
		c.cli, c.token, err = build(c.id, c.secret, c.user, c.pass)
		if err != nil {
			return err
		}
	}
	return c.exec(r, out)
}

// exec executes an http request and parsest the http response into the out
// interface.
func (c *client) exec(r *http.Request, out interface{}) error {
	rawResp, err := c.doRaw(r)
	if err != nil {
		return err
	}

	if rawResp.StatusCode != http.StatusOK {
		return fmt.Errorf("bad response code: %d\n"+
			"request was: %v\n",
			rawResp.StatusCode,
			r)
	}

	defer func() {
		if rawResp.Body != nil {
			rawResp.Body.Close()
		}
	}()

	decoder := json.NewDecoder(rawResp.Body)
	return decoder.Decode(out)
}

// doRaw executes an http Request using an authenticated client, and the configured
// user agent.
func (c *client) doRaw(r *http.Request) (*http.Response, error) {
	if r.Header == nil {
		r.Header = make(http.Header)
	}
	r.Header.Add("User-Agent", c.agent)
	return c.cli.Do(r)
}
<commit_msg>Increase rate limit on requests.<commit_after>package client

import (
	"encoding/json"
	"fmt"
	"net/http"
	"sync"
	"time"

	"golang.org/x/oauth2"
)

const (
	// rateLimit is the wait time between queries.
	rateLimit = 3 * time.Second
)

type client struct {
	// agent is the client's User-Agent in http requests.
	agent string
	// id is the bot's OAuth2 client id.
	id string
	// secret is the bot's OAuth2 client secret.
	secret string
	// user is the bot's username on reddit.
	user string
	// pass is the bot's password on reddit.
	pass string

	// authMu protects authentication fields.
	authMu sync.Mutex
	// cli is the authenticated client to execute requests with.
	cli *http.Client
	// token is the OAuth2 token cli uses to authenticate.
	token *oauth2.Token

	// rateMu protects rate limiting fields.
	rateMu sync.Mutex
	// nextReq is the time at which it is ok to make the next request.
	nextReq time.Time
}

// Do wraps the execution of http requests. It updates authentications and rate
// limits requests to Reddit to comply with the API rules.
func (c *client) Do(r *http.Request, out interface{}) error {
	if time.Now().Before(c.nextReq) {
		<-time.After(c.nextReq.Sub(time.Now()))
	}
	c.nextReq = time.Now().Add(rateLimit)
	if !c.token.Valid() {
		var err error
		c.cli, c.token, err = build(c.id, c.secret, c.user, c.pass)
		if err != nil {
			return err
		}
	}
	return c.exec(r, out)
}

// exec executes an http request and parsest the http response into the out
// interface.
func (c *client) exec(r *http.Request, out interface{}) error {
	rawResp, err := c.doRaw(r)
	if err != nil {
		return err
	}

	if rawResp.StatusCode != http.StatusOK {
		return fmt.Errorf("bad response code: %d\n"+
			"request was: %v\n",
			rawResp.StatusCode,
			r)
	}

	defer func() {
		if rawResp.Body != nil {
			rawResp.Body.Close()
		}
	}()

	decoder := json.NewDecoder(rawResp.Body)
	return decoder.Decode(out)
}

// doRaw executes an http Request using an authenticated client, and the configured
// user agent.
func (c *client) doRaw(r *http.Request) (*http.Response, error) {
	if r.Header == nil {
		r.Header = make(http.Header)
	}
	r.Header.Add("User-Agent", c.agent)
	return c.cli.Do(r)
}
<|endoftext|>


<commit_before>// Copyright 2012 Dorival de Moraes Pedroso. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build ignore

package main

import (
	"bytes"
	"math"

	"github.com/cpmech/goga"
	"github.com/cpmech/gosl/io"
	"github.com/cpmech/gosl/la"
	"github.com/cpmech/gosl/rnd"
	"github.com/cpmech/gosl/utl"
)

func calc_beta(best *goga.Individual, βref float64, verbose bool) (β float64) {
	xs := make([]float64, best.Nfltgenes)
	for i := 0; i < best.Nfltgenes; i++ {
		xs[i] = best.GetFloat(i)
	}
	β = math.Sqrt(la.VecDot(xs, xs))
	if verbose {
		io.Pf("\nova = %g  oor = %g\n", best.Ova, best.Oor)
		io.Pf("x   = %v\n", xs)
		io.PfYel("β   = %g", β)
		io.Pf(" (%g)\n", βref)
	}
	return
}

func main() {

	rnd.Init(0)

	// standard variables => μ=0 and σ=1 => y = x

	g := func(x []float64) float64 {
		return 0.1*math.Pow(x[0]-x[1], 2.0) - (x[0]+x[1])/math.Sqrt2 + 2.5
	}

	x := make([]float64, 2)
	ovfunc := func(ind *goga.Individual, idIsland, time int, report *bytes.Buffer) (ova, oor float64) {
		x[0], x[1] = ind.GetFloat(0), ind.GetFloat(1)
		fp := utl.GtePenalty(1e-2, math.Abs(g(x)), 1)
		ova = la.VecDot(x, x) + fp
		oor = fp
		return
	}

	// parameters
	C := goga.NewConfParams()
	C.Pll = false
	C.Nisl = 4
	C.Ninds = 20
	C.DoPlot = true
	C.CalcDerived()

	// bingo
	ndim := 2
	vmin, vmax := -1.0, 1.0
	xmin, xmax := utl.DblVals(ndim, vmin), utl.DblVals(ndim, vmax)
	bingo := goga.NewBingoFloats(xmin, xmax)

	// evolver
	βref := 2.5
	evo := goga.NewEvolverFloatChromo(C, xmin, xmax, ovfunc, bingo)

	// for a number of trials
	ntrials := 100
	betas := make([]float64, ntrials)
	for i := 0; i < ntrials; i++ {

		// reset population
		if i > 0 {
			for _, isl := range evo.Islands {
				isl.Pop.GenFloatRandom(C, xmin, xmax)
			}
		}
		pop0 := evo.Islands[0].Pop

		// run
		check := i == ntrials-1
		verbose := check
		doreport := check
		evo.Run(verbose, doreport)
		β := calc_beta(evo.Best, βref, verbose)
		betas[i] = β

		// plot contour
		if check {
			if C.DoPlot {
				xmin := []float64{-1, -1}
				xmax := []float64{5, 5}
				goga.PlotTwoVarsContour("/tmp/goga", "reliability-problem1", pop0, evo.Islands[0].Pop, evo.Best,
					xmin, xmax, 41, true, nil, g, g)
			}
		}
	}

	// analysis
	βmin, βave, βmax, βdev := rnd.StatBasic(betas, true)
	io.Pf("\nβmin = %v\n", βmin)
	io.PfYel("βave = %v\n", βave)
	io.Pf("βmax = %v\n", βmax)
	io.Pf("βdev = %v\n\n", βdev)
	io.Pf(rnd.BuildTextHist(nice_num(βmin-0.005), nice_num(βmax+0.005), 11, betas, "%.3f", 60))
}

func nice_num(x float64) float64 {
	s := io.Sf("%.2f", x)
	return io.Atof(s)
}
<commit_msg>range in example increased making the repeatability better<commit_after>// Copyright 2012 Dorival de Moraes Pedroso. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// +build ignore

package main

import (
	"bytes"
	"math"

	"github.com/cpmech/goga"
	"github.com/cpmech/gosl/io"
	"github.com/cpmech/gosl/la"
	"github.com/cpmech/gosl/rnd"
	"github.com/cpmech/gosl/utl"
)

func calc_beta(best *goga.Individual, βref float64, verbose bool) (β float64) {
	xs := make([]float64, best.Nfltgenes)
	for i := 0; i < best.Nfltgenes; i++ {
		xs[i] = best.GetFloat(i)
	}
	β = math.Sqrt(la.VecDot(xs, xs))
	if verbose {
		io.Pf("\nova = %g  oor = %g\n", best.Ova, best.Oor)
		io.Pf("x   = %v\n", xs)
		io.PfYel("β   = %g", β)
		io.Pf(" (%g)\n", βref)
	}
	return
}

func main() {

	rnd.Init(0)

	// standard variables => μ=0 and σ=1 => y = x

	g := func(x []float64) float64 {
		return 0.1*math.Pow(x[0]-x[1], 2.0) - (x[0]+x[1])/math.Sqrt2 + 2.5
	}

	x := make([]float64, 2)
	ovfunc := func(ind *goga.Individual, idIsland, time int, report *bytes.Buffer) (ova, oor float64) {
		x[0], x[1] = ind.GetFloat(0), ind.GetFloat(1)
		fp := utl.GtePenalty(1e-2, math.Abs(g(x)), 1)
		ova = la.VecDot(x, x) + fp
		oor = fp
		return
	}

	// parameters
	C := goga.NewConfParams()
	C.Pll = false
	C.Nisl = 4
	C.Ninds = 20
	C.DoPlot = true
	C.CalcDerived()

	// bingo
	ndim := 2
	vmin, vmax := -2.0, 2.0
	xmin, xmax := utl.DblVals(ndim, vmin), utl.DblVals(ndim, vmax)
	bingo := goga.NewBingoFloats(xmin, xmax)

	// evolver
	βref := 2.5
	evo := goga.NewEvolverFloatChromo(C, xmin, xmax, ovfunc, bingo)

	// for a number of trials
	ntrials := 100
	betas := make([]float64, ntrials)
	for i := 0; i < ntrials; i++ {

		// reset population
		if i > 0 {
			for _, isl := range evo.Islands {
				isl.Pop.GenFloatRandom(C, xmin, xmax)
			}
		}
		pop0 := evo.Islands[0].Pop

		// run
		check := i == ntrials-1
		verbose := check
		doreport := check
		evo.Run(verbose, doreport)
		β := calc_beta(evo.Best, βref, verbose)
		betas[i] = β

		// plot contour
		if check {
			if C.DoPlot {
				xmin := []float64{-1, -1}
				xmax := []float64{5, 5}
				goga.PlotTwoVarsContour("/tmp/goga", "reliability-problem1", pop0, evo.Islands[0].Pop, evo.Best,
					xmin, xmax, 41, true, nil, g, g)
			}
		}
	}

	// analysis
	βmin, βave, βmax, βdev := rnd.StatBasic(betas, true)
	io.Pf("\nβmin = %v\n", βmin)
	io.PfYel("βave = %v\n", βave)
	io.Pf("βmax = %v\n", βmax)
	io.Pf("βdev = %v\n\n", βdev)
	io.Pf(rnd.BuildTextHist(nice_num(βmin-0.005), nice_num(βmax+0.005), 11, betas, "%.3f", 60))
}

func nice_num(x float64) float64 {
	s := io.Sf("%.2f", x)
	return io.Atof(s)
}
<|endoftext|>


<commit_before>// Copyright 2019 PingCAP, 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 expression_test

import (
	. "github.com/pingcap/check"
	"github.com/pingcap/tidb/util/testkit"
)

func (s *testIntegrationSuite) TestFoldIfNull(c *C) {
	tk := testkit.NewTestKit(c, s.store)
	tk.MustExec(`use test;`)
	tk.MustExec(`drop table if exists t;`)
	tk.MustExec(`create table t(a bigint, b bigint);`)
	tk.MustExec(`insert into t values(1, 1);`)
	tk.MustQuery(`desc select ifnull("aaaa", a) from t;`).Check(testkit.Rows(
		`Projection_3 10000.00 root  aaaa->Column#4`,
		`└─TableReader_5 10000.00 root  data:TableFullScan_4`,
		`  └─TableFullScan_4 10000.00 cop[tikv] table:t keep order:false, stats:pseudo`,
	))
	tk.MustQuery(`show warnings;`).Check(testkit.Rows())
	tk.MustQuery(`select ifnull("aaaa", a) from t;`).Check(testkit.Rows("aaaa"))
	tk.MustQuery(`show warnings;`).Check(testkit.Rows())
}
<commit_msg>expression: migrate test-infra to testify for constant_fold_test.go (#30424)<commit_after>// Copyright 2019 PingCAP, 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 expression_test

import (
	"testing"

	"github.com/pingcap/tidb/testkit"
)

func TestFoldIfNull(t *testing.T) {
	t.Parallel()

	store, clean := testkit.CreateMockStore(t)
	defer clean()

	tk := testkit.NewTestKit(t, store)
	tk.MustExec(`use test;`)
	tk.MustExec(`drop table if exists t;`)
	tk.MustExec(`create table t(a bigint, b bigint);`)
	tk.MustExec(`insert into t values(1, 1);`)
	tk.MustQuery(`desc select ifnull("aaaa", a) from t;`).Check(testkit.Rows(
		`Projection_3 10000.00 root  aaaa->Column#4`,
		`└─TableReader_5 10000.00 root  data:TableFullScan_4`,
		`  └─TableFullScan_4 10000.00 cop[tikv] table:t keep order:false, stats:pseudo`,
	))
	tk.MustQuery(`show warnings;`).Check(testkit.Rows())
	tk.MustQuery(`select ifnull("aaaa", a) from t;`).Check(testkit.Rows("aaaa"))
	tk.MustQuery(`show warnings;`).Check(testkit.Rows())
}
<|endoftext|>


<commit_before>package filemanager

import (
	"errors"
	"mime"
	"net/http"
	"path/filepath"
	"strings"
)

// assetsURL is the url where static assets are served.
const assetsURL = "/_internal"

// Serve provides the needed assets for the front-end
func serveAssets(w http.ResponseWriter, r *http.Request, m *FileManager) (int, error) {
	// gets the filename to be used with Assets function
	filename := strings.Replace(r.URL.Path, m.BaseURL+assetsURL, "", 1)

	var file []byte
	var err error

	switch {
	case strings.HasPrefix(filename, "/css"):
		filename = strings.Replace(filename, "/css/", "", 1)

		if m.Assets.CSS != nil {
			file, err = m.Assets.JS.Bytes(filename)
			if err == nil {
				break
			}
		}

		file, err = m.Assets.baseCSS.Bytes(filename)
	case strings.HasPrefix(filename, "/js"):
		filename = strings.Replace(filename, "/js/", "", 1)
		file, err = m.Assets.requiredJS.Bytes(filename)
	case strings.HasPrefix(filename, "/vendor"):
		if m.Assets.JS != nil {
			filename = strings.Replace(filename, "/vendor/", "", 1)
			file, err = m.Assets.JS.Bytes(filename)
			break
		}

		fallthrough
	default:
		err = errors.New("not found")
	}

	if err != nil {
		return http.StatusNotFound, nil
	}

	// Get the file extension and its mimetype
	extension := filepath.Ext(filename)
	mediatype := mime.TypeByExtension(extension)

	// Write the header with the Content-Type and write the file
	// content to the buffer
	w.Header().Set("Content-Type", mediatype)
	w.Write(file)
	return 200, nil
}
<commit_msg>fix css bug<commit_after>package filemanager

import (
	"errors"
	"mime"
	"net/http"
	"path/filepath"
	"strings"
)

// assetsURL is the url where static assets are served.
const assetsURL = "/_internal"

// Serve provides the needed assets for the front-end
func serveAssets(w http.ResponseWriter, r *http.Request, m *FileManager) (int, error) {
	// gets the filename to be used with Assets function
	filename := strings.Replace(r.URL.Path, m.BaseURL+assetsURL, "", 1)

	var file []byte
	var err error

	switch {
	case strings.HasPrefix(filename, "/css"):
		filename = strings.Replace(filename, "/css/", "", 1)

		if m.Assets.CSS != nil {
			file, err = m.Assets.CSS.Bytes(filename)
			if err == nil {
				break
			}
		}

		file, err = m.Assets.baseCSS.Bytes(filename)
	case strings.HasPrefix(filename, "/js"):
		filename = strings.Replace(filename, "/js/", "", 1)
		file, err = m.Assets.requiredJS.Bytes(filename)
	case strings.HasPrefix(filename, "/vendor"):
		if m.Assets.JS != nil {
			filename = strings.Replace(filename, "/vendor/", "", 1)
			file, err = m.Assets.JS.Bytes(filename)
			break
		}

		fallthrough
	default:
		err = errors.New("not found")
	}

	if err != nil {
		return http.StatusNotFound, nil
	}

	// Get the file extension and its mimetype
	extension := filepath.Ext(filename)
	mediatype := mime.TypeByExtension(extension)

	// Write the header with the Content-Type and write the file
	// content to the buffer
	w.Header().Set("Content-Type", mediatype)
	w.Write(file)
	return 200, nil
}
<|endoftext|>


<commit_before>package httpu

import (
	"bufio"
	"bytes"
	"fmt"
	"log"
	"net"
	"net/http"
	"sync"
	"time"
)

// HTTPUClient is a client for dealing with HTTPU (HTTP over UDP). Its typical
// function is for HTTPMU, and particularly SSDP.
type HTTPUClient struct {
	connLock sync.Mutex // Protects use of conn.
	conn     net.PacketConn
}

// NewHTTPUClient creates a new HTTPUClient, opening up a new UDP socket for the
// purpose.
func NewHTTPUClient() (*HTTPUClient, error) {
	conn, err := net.ListenPacket("udp", ":0")
	if err != nil {
		return nil, err
	}
	return &HTTPUClient{conn: conn}, nil
}

// Close shuts down the client. The client will no longer be useful following
// this.
func (httpu *HTTPUClient) Close() error {
	httpu.connLock.Lock()
	defer httpu.connLock.Unlock()
	return httpu.conn.Close()
}

// Do performs a request. The timeout is how long to wait for before returning
// the responses that were received. An error is only returned for failing to
// send the request. Failures in receipt simply do not add to the resulting
// responses.
//
// Note that at present only one concurrent connection will happen per
// HTTPUClient.
func (httpu *HTTPUClient) Do(req *http.Request, timeout time.Duration, numSends int) ([]*http.Response, error) {
	httpu.connLock.Lock()
	defer httpu.connLock.Unlock()

	var requestBuf bytes.Buffer
	if err := req.Write(&requestBuf); err != nil {
		return nil, err
	}
	destAddr, err := net.ResolveUDPAddr("udp", req.Host)
	if err != nil {
		return nil, err
	}
	if err = httpu.conn.SetDeadline(time.Now().Add(timeout)); err != nil {
		return nil, err
	}

	// Send request.
	for i := 0; i < numSends; i++ {
		if n, err := httpu.conn.WriteTo(requestBuf.Bytes(), destAddr); err != nil {
			return nil, err
		} else if n < len(requestBuf.Bytes()) {
			return nil, fmt.Errorf("httpu: wrote %d bytes rather than full %d in request",
				n, len(requestBuf.Bytes()))
		}
		time.Sleep(5 * time.Millisecond)
	}

	// Await responses until timeout.
	var responses []*http.Response
	responseBytes := make([]byte, 2048)
	for {
		// 2048 bytes should be sufficient for most networks.
		n, _, err := httpu.conn.ReadFrom(responseBytes)
		if err != nil {
			if err, ok := err.(net.Error); ok && err.Timeout() {
				break
			}
			log.Print("httpu: error while receiving response: %v", err)
			// Sleep in case this is a persistent error to avoid pegging CPU until deadline.
			time.Sleep(10 * time.Millisecond)
			continue
		}

		// Parse response.
		response, err := http.ReadResponse(bufio.NewReader(bytes.NewBuffer(responseBytes[:n])), req)
		if err != nil {
			log.Print("httpu: error while parsing response: %v", err)
			continue
		}

		responses = append(responses, response)
	}
	return responses, err
}
<commit_msg>Reduce headers written by httpu, in case this confuses some devices.<commit_after>package httpu

import (
	"bufio"
	"bytes"
	"fmt"
	"log"
	"net"
	"net/http"
	"sync"
	"time"
)

// HTTPUClient is a client for dealing with HTTPU (HTTP over UDP). Its typical
// function is for HTTPMU, and particularly SSDP.
type HTTPUClient struct {
	connLock sync.Mutex // Protects use of conn.
	conn     net.PacketConn
}

// NewHTTPUClient creates a new HTTPUClient, opening up a new UDP socket for the
// purpose.
func NewHTTPUClient() (*HTTPUClient, error) {
	conn, err := net.ListenPacket("udp", ":0")
	if err != nil {
		return nil, err
	}
	return &HTTPUClient{conn: conn}, nil
}

// Close shuts down the client. The client will no longer be useful following
// this.
func (httpu *HTTPUClient) Close() error {
	httpu.connLock.Lock()
	defer httpu.connLock.Unlock()
	return httpu.conn.Close()
}

// Do performs a request. The timeout is how long to wait for before returning
// the responses that were received. An error is only returned for failing to
// send the request. Failures in receipt simply do not add to the resulting
// responses.
//
// Note that at present only one concurrent connection will happen per
// HTTPUClient.
func (httpu *HTTPUClient) Do(req *http.Request, timeout time.Duration, numSends int) ([]*http.Response, error) {
	httpu.connLock.Lock()
	defer httpu.connLock.Unlock()

	// Create the request. This is a subset of what http.Request.Write does
	// deliberately to avoid creating extra fields which may confuse some
	// devices.
	var requestBuf bytes.Buffer
	method := req.Method
	if method == "" {
		method = "GET"
	}
	if _, err := fmt.Fprintf(&requestBuf, "%s %s HTTP/1.1\r\n", method, req.URL.RequestURI()); err != nil {
		return nil, err
	}
	if err := req.Header.Write(&requestBuf); err != nil {
		return nil, err
	}
	if _, err := requestBuf.Write([]byte{'\r', '\n'}); err != nil {
		return nil, err
	}

	destAddr, err := net.ResolveUDPAddr("udp", req.Host)
	if err != nil {
		return nil, err
	}
	if err = httpu.conn.SetDeadline(time.Now().Add(timeout)); err != nil {
		return nil, err
	}

	// Send request.
	for i := 0; i < numSends; i++ {
		if n, err := httpu.conn.WriteTo(requestBuf.Bytes(), destAddr); err != nil {
			return nil, err
		} else if n < len(requestBuf.Bytes()) {
			return nil, fmt.Errorf("httpu: wrote %d bytes rather than full %d in request",
				n, len(requestBuf.Bytes()))
		}
		time.Sleep(5 * time.Millisecond)
	}

	// Await responses until timeout.
	var responses []*http.Response
	responseBytes := make([]byte, 2048)
	for {
		// 2048 bytes should be sufficient for most networks.
		n, _, err := httpu.conn.ReadFrom(responseBytes)
		if err != nil {
			if err, ok := err.(net.Error); ok && err.Timeout() {
				break
			}
			log.Print("httpu: error while receiving response: %v", err)
			// Sleep in case this is a persistent error to avoid pegging CPU until deadline.
			time.Sleep(10 * time.Millisecond)
			continue
		}

		// Parse response.
		response, err := http.ReadResponse(bufio.NewReader(bytes.NewBuffer(responseBytes[:n])), req)
		if err != nil {
			log.Print("httpu: error while parsing response: %v", err)
			continue
		}

		responses = append(responses, response)
	}
	return responses, err
}
<|endoftext|>


<commit_before>package warserver

import (
    "flag"
    "warserver/logger"
    "warserver/PortMgmt"
    "net/http"
    "os"
    "strings"
)

func setupLogger(path string) {
    if strings.Contains(path, "/dev/stderr") {
        logfile := os.Stderr
        logger.SetupLogger(logger.DEBUG, logger.USUAL, logfile)
    } else {
        reallog, err := os.Create(path)
        if err != nil {
            panic(err)
        }
        logger.SetupLogger(logger.DEBUG, logger.USUAL, reallog)
    }
}

var port PortMgmt.PortInfo
var tcpPort PortMgmt.PortInfo
var porterPort PortMgmt.PortInfo
var porterIP PortMgmt.IPString

func Main() {
    porterIPString := flag.String("porterIP", "127.0.0.1", "Porter Server IP")
    porterPortString := flag.String("porterPort", ":5269", "Porter Server Port")
    portstring := flag.String("port", ":8888", "Server port")
    tcpportstring := flag.String("tcpport", ":11199", "TCP socket port")
    logpath := flag.String("logpath", "/dev/stderr", "Logging location")
    flag.Parse()

    port = PortMgmt.NewPortInfo(*portstring)
    tcpPort = PortMgmt.NewPortInfo(*tcpportstring)
    porterPort = PortMgmt.NewPortInfo(*porterPortString)
    porterIP = PortMgmt.IPString(*porterIPString)

    setupLogger(*logpath)
    setupGamehub()

    go gamehub.handleConnections()

    go socketListen(tcpPort.Port)
    logger.Debugf("Socket listening on port %s", tcpPort)

    http.HandleFunc("/ws", serveWs)

    logger.Debugf("Http server listening on port %s", port);
    err := http.ListenAndServe(port.Port, nil)
    if err != nil {
        logger.Fatalf("ListenAndServe: %s", err.Error())
    }
}
<commit_msg>Allow logging to stdout<commit_after>package warserver

import (
    "flag"
    "warserver/logger"
    "warserver/PortMgmt"
    "net/http"
    "os"
    "strings"
)

func setupLogger(path string) {
    if strings.Contains(path, "/dev/stderr") {
        logfile := os.Stderr
        logger.SetupLogger(logger.DEBUG, logger.USUAL, logfile)
    } else if strings.Contains(path, "dev/stdout") {
        logfile := os.Stdout
        logger.SetupLogger(logger.DEBUG, logger.USUAL, logfile)
    } else {
        reallog, err := os.Create(path)
        if err != nil {
            panic(err)
        }
        logger.SetupLogger(logger.DEBUG, logger.USUAL, reallog)
    }
}

var port PortMgmt.PortInfo
var tcpPort PortMgmt.PortInfo
var porterPort PortMgmt.PortInfo
var porterIP PortMgmt.IPString

func Main() {
    porterIPString := flag.String("porterIP", "127.0.0.1", "Porter Server IP")
    porterPortString := flag.String("porterPort", ":5269", "Porter Server Port")
    portstring := flag.String("port", ":8888", "Server port")
    tcpportstring := flag.String("tcpport", ":11199", "TCP socket port")
    logpath := flag.String("logpath", "/dev/stderr", "Logging location")
    flag.Parse()

    port = PortMgmt.NewPortInfo(*portstring)
    tcpPort = PortMgmt.NewPortInfo(*tcpportstring)
    porterPort = PortMgmt.NewPortInfo(*porterPortString)
    porterIP = PortMgmt.IPString(*porterIPString)

    setupLogger(*logpath)
    setupGamehub()

    go gamehub.handleConnections()

    go socketListen(tcpPort.Port)
    logger.Debugf("Socket listening on port %s", tcpPort)

    http.HandleFunc("/ws", serveWs)

    logger.Debugf("Http server listening on port %s", port);
    err := http.ListenAndServe(port.Port, nil)
    if err != nil {
        logger.Fatalf("ListenAndServe: %s", err.Error())
    }
}
<|endoftext|>


<commit_before>/*
Copyright 2014 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 framework

import (
	"fmt"
	"os"
	"path"
	"path/filepath"
	"strings"
	"time"

	"k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/apimachinery/pkg/util/wait"
	clientset "k8s.io/client-go/kubernetes"
	imageutils "k8s.io/kubernetes/test/utils/image"
)

func EtcdUpgrade(target_storage, target_version string) error {
	switch TestContext.Provider {
	case "gce":
		return etcdUpgradeGCE(target_storage, target_version)
	default:
		return fmt.Errorf("EtcdUpgrade() is not implemented for provider %s", TestContext.Provider)
	}
}

func MasterUpgrade(v string) error {
	switch TestContext.Provider {
	case "gce":
		return masterUpgradeGCE(v, false)
	case "gke":
		return masterUpgradeGKE(v)
	case "kubernetes-anywhere":
		return masterUpgradeKubernetesAnywhere(v)
	default:
		return fmt.Errorf("MasterUpgrade() is not implemented for provider %s", TestContext.Provider)
	}
}

func etcdUpgradeGCE(target_storage, target_version string) error {
	env := append(
		os.Environ(),
		"TEST_ETCD_VERSION="+target_version,
		"STORAGE_BACKEND="+target_storage,
		"TEST_ETCD_IMAGE=3.0.17")

	_, _, err := RunCmdEnv(env, gceUpgradeScript(), "-l", "-M")
	return err
}

// TODO(mrhohn): Remove this function when kube-proxy is run as a DaemonSet by default.
func MasterUpgradeGCEWithKubeProxyDaemonSet(v string, enableKubeProxyDaemonSet bool) error {
	return masterUpgradeGCE(v, enableKubeProxyDaemonSet)
}

// TODO(mrhohn): Remove 'enableKubeProxyDaemonSet' when kube-proxy is run as a DaemonSet by default.
func masterUpgradeGCE(rawV string, enableKubeProxyDaemonSet bool) error {
	env := append(os.Environ(), fmt.Sprintf("KUBE_PROXY_DAEMONSET=%v", enableKubeProxyDaemonSet))
	// TODO: Remove these variables when they're no longer needed for downgrades.
	if TestContext.EtcdUpgradeVersion != "" && TestContext.EtcdUpgradeStorage != "" {
		env = append(env,
			"TEST_ETCD_VERSION="+TestContext.EtcdUpgradeVersion,
			"STORAGE_BACKEND="+TestContext.EtcdUpgradeStorage,
			"TEST_ETCD_IMAGE=3.0.17")
	}

	v := "v" + rawV
	_, _, err := RunCmdEnv(env, gceUpgradeScript(), "-M", v)
	return err
}

func masterUpgradeGKE(v string) error {
	Logf("Upgrading master to %q", v)
	_, _, err := RunCmd("gcloud", "container",
		"clusters",
		fmt.Sprintf("--project=%s", TestContext.CloudConfig.ProjectID),
		fmt.Sprintf("--zone=%s", TestContext.CloudConfig.Zone),
		"upgrade",
		TestContext.CloudConfig.Cluster,
		"--master",
		fmt.Sprintf("--cluster-version=%s", v),
		"--quiet")
	if err != nil {
		return err
	}

	waitForSSHTunnels()

	return nil
}

func masterUpgradeKubernetesAnywhere(v string) error {
	Logf("Upgrading master to %q", v)

	kaPath := TestContext.KubernetesAnywherePath
	originalConfigPath := filepath.Join(kaPath, ".config")
	backupConfigPath := filepath.Join(kaPath, ".config.bak")
	updatedConfigPath := filepath.Join(kaPath, fmt.Sprintf(".config-%s", v))

	// modify config with specified k8s version
	if _, _, err := RunCmd("sed",
		"-i.bak", // writes original to .config.bak
		fmt.Sprintf(`s/kubernetes_version=.*$/kubernetes_version=%s/`, v),
		originalConfigPath); err != nil {
		return err
	}

	defer func() {
		// revert .config.bak to .config
		if err := os.Rename(backupConfigPath, originalConfigPath); err != nil {
			Logf("Could not rename %s back to %s", backupConfigPath, originalConfigPath)
		}
	}()

	// invoke ka upgrade
	if _, _, err := RunCmd("make", "-C", TestContext.KubernetesAnywherePath,
		"WAIT_FOR_KUBECONFIG=y", "upgrade-master"); err != nil {
		return err
	}

	// move .config to .config.<version>
	if err := os.Rename(originalConfigPath, updatedConfigPath); err != nil {
		return err
	}

	return nil
}

func NodeUpgrade(f *Framework, v string, img string) error {
	// Perform the upgrade.
	var err error
	switch TestContext.Provider {
	case "gce":
		err = nodeUpgradeGCE(v, img, false)
	case "gke":
		err = nodeUpgradeGKE(v, img)
	default:
		err = fmt.Errorf("NodeUpgrade() is not implemented for provider %s", TestContext.Provider)
	}
	if err != nil {
		return err
	}

	// Wait for it to complete and validate nodes are healthy.
	//
	// TODO(ihmccreery) We shouldn't have to wait for nodes to be ready in
	// GKE; the operation shouldn't return until they all are.
	Logf("Waiting up to %v for all nodes to be ready after the upgrade", RestartNodeReadyAgainTimeout)
	if _, err := CheckNodesReady(f.ClientSet, RestartNodeReadyAgainTimeout, TestContext.CloudConfig.NumNodes); err != nil {
		return err
	}
	return nil
}

// TODO(mrhohn): Remove this function when kube-proxy is run as a DaemonSet by default.
func NodeUpgradeGCEWithKubeProxyDaemonSet(f *Framework, v string, img string, enableKubeProxyDaemonSet bool) error {
	// Perform the upgrade.
	if err := nodeUpgradeGCE(v, img, enableKubeProxyDaemonSet); err != nil {
		return err
	}
	// Wait for it to complete and validate nodes are healthy.
	Logf("Waiting up to %v for all nodes to be ready after the upgrade", RestartNodeReadyAgainTimeout)
	if _, err := CheckNodesReady(f.ClientSet, RestartNodeReadyAgainTimeout, TestContext.CloudConfig.NumNodes); err != nil {
		return err
	}
	return nil
}

// TODO(mrhohn): Remove 'enableKubeProxyDaemonSet' when kube-proxy is run as a DaemonSet by default.
func nodeUpgradeGCE(rawV, img string, enableKubeProxyDaemonSet bool) error {
	v := "v" + rawV
	env := append(os.Environ(), fmt.Sprintf("KUBE_PROXY_DAEMONSET=%v", enableKubeProxyDaemonSet))
	if img != "" {
		env = append(env, "KUBE_NODE_OS_DISTRIBUTION="+img)
		_, _, err := RunCmdEnv(env, gceUpgradeScript(), "-N", "-o", v)
		return err
	}
	_, _, err := RunCmdEnv(env, gceUpgradeScript(), "-N", v)
	return err
}

func nodeUpgradeGKE(v string, img string) error {
	Logf("Upgrading nodes to version %q and image %q", v, img)
	args := []string{
		"container",
		"clusters",
		fmt.Sprintf("--project=%s", TestContext.CloudConfig.ProjectID),
		fmt.Sprintf("--zone=%s", TestContext.CloudConfig.Zone),
		"upgrade",
		TestContext.CloudConfig.Cluster,
		fmt.Sprintf("--cluster-version=%s", v),
		"--quiet",
	}
	if len(img) > 0 {
		args = append(args, fmt.Sprintf("--image-type=%s", img))
	}
	_, _, err := RunCmd("gcloud", args...)

	if err != nil {
		return err
	}

	waitForSSHTunnels()

	return nil
}

// CheckNodesReady waits up to nt for expect nodes accessed by c to be ready,
// returning an error if this doesn't happen in time. It returns the names of
// nodes it finds.
func CheckNodesReady(c clientset.Interface, nt time.Duration, expect int) ([]string, error) {
	// First, keep getting all of the nodes until we get the number we expect.
	var nodeList *v1.NodeList
	var errLast error
	start := time.Now()
	found := wait.Poll(Poll, nt, func() (bool, error) {
		// A rolling-update (GCE/GKE implementation of restart) can complete before the apiserver
		// knows about all of the nodes. Thus, we retry the list nodes call
		// until we get the expected number of nodes.
		nodeList, errLast = c.Core().Nodes().List(metav1.ListOptions{
			FieldSelector: fields.Set{"spec.unschedulable": "false"}.AsSelector().String()})
		if errLast != nil {
			return false, nil
		}
		if len(nodeList.Items) != expect {
			errLast = fmt.Errorf("expected to find %d nodes but found only %d (%v elapsed)",
				expect, len(nodeList.Items), time.Since(start))
			Logf("%v", errLast)
			return false, nil
		}
		return true, nil
	}) == nil
	nodeNames := make([]string, len(nodeList.Items))
	for i, n := range nodeList.Items {
		nodeNames[i] = n.ObjectMeta.Name
	}
	if !found {
		return nodeNames, fmt.Errorf("couldn't find %d nodes within %v; last error: %v",
			expect, nt, errLast)
	}
	Logf("Successfully found %d nodes", expect)

	// Next, ensure in parallel that all the nodes are ready. We subtract the
	// time we spent waiting above.
	timeout := nt - time.Since(start)
	result := make(chan bool, len(nodeList.Items))
	for _, n := range nodeNames {
		n := n
		go func() { result <- WaitForNodeToBeReady(c, n, timeout) }()
	}
	failed := false
	// TODO(mbforbes): Change to `for range` syntax once we support only Go
	// >= 1.4.
	for i := range nodeList.Items {
		_ = i
		if !<-result {
			failed = true
		}
	}
	if failed {
		return nodeNames, fmt.Errorf("at least one node failed to be ready")
	}
	return nodeNames, nil
}

// MigTemplate (GCE-only) returns the name of the MIG template that the
// nodes of the cluster use.
func MigTemplate() (string, error) {
	var errLast error
	var templ string
	key := "instanceTemplate"
	if wait.Poll(Poll, SingleCallTimeout, func() (bool, error) {
		// TODO(mikedanese): make this hit the compute API directly instead of
		// shelling out to gcloud.
		// An `instance-groups managed describe` call outputs what we want to stdout.
		output, _, err := retryCmd("gcloud", "compute", "instance-groups", "managed",
			fmt.Sprintf("--project=%s", TestContext.CloudConfig.ProjectID),
			"describe",
			fmt.Sprintf("--zone=%s", TestContext.CloudConfig.Zone),
			TestContext.CloudConfig.NodeInstanceGroup)
		if err != nil {
			errLast = fmt.Errorf("gcloud compute instance-groups managed describe call failed with err: %v", err)
			return false, nil
		}

		// The 'describe' call probably succeeded; parse the output and try to
		// find the line that looks like "instanceTemplate: url/to/<templ>" and
		// return <templ>.
		if val := ParseKVLines(output, key); len(val) > 0 {
			url := strings.Split(val, "/")
			templ = url[len(url)-1]
			Logf("MIG group %s using template: %s", TestContext.CloudConfig.NodeInstanceGroup, templ)
			return true, nil
		}
		errLast = fmt.Errorf("couldn't find %s in output to get MIG template. Output: %s", key, output)
		return false, nil
	}) != nil {
		return "", fmt.Errorf("MigTemplate() failed with last error: %v", errLast)
	}
	return templ, nil
}

func gceUpgradeScript() string {
	if len(TestContext.GCEUpgradeScript) == 0 {
		return path.Join(TestContext.RepoRoot, "cluster/gce/upgrade.sh")
	}
	return TestContext.GCEUpgradeScript
}

func waitForSSHTunnels() {
	Logf("Waiting for SSH tunnels to establish")
	RunKubectl("run", "ssh-tunnel-test",
		"--image="+imageutils.GetBusyBoxImage(),
		"--restart=Never",
		"--command", "--",
		"echo", "Hello")
	defer RunKubectl("delete", "pod", "ssh-tunnel-test")

	// allow up to a minute for new ssh tunnels to establish
	wait.PollImmediate(5*time.Second, time.Minute, func() (bool, error) {
		_, err := RunKubectl("logs", "ssh-tunnel-test")
		return err == nil, nil
	})
}
<commit_msg>Version should be quoted so jq doesn't interpret it as numeric<commit_after>/*
Copyright 2014 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 framework

import (
	"fmt"
	"os"
	"path"
	"path/filepath"
	"strings"
	"time"

	"k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/apimachinery/pkg/util/wait"
	clientset "k8s.io/client-go/kubernetes"
	imageutils "k8s.io/kubernetes/test/utils/image"
)

func EtcdUpgrade(target_storage, target_version string) error {
	switch TestContext.Provider {
	case "gce":
		return etcdUpgradeGCE(target_storage, target_version)
	default:
		return fmt.Errorf("EtcdUpgrade() is not implemented for provider %s", TestContext.Provider)
	}
}

func MasterUpgrade(v string) error {
	switch TestContext.Provider {
	case "gce":
		return masterUpgradeGCE(v, false)
	case "gke":
		return masterUpgradeGKE(v)
	case "kubernetes-anywhere":
		return masterUpgradeKubernetesAnywhere(v)
	default:
		return fmt.Errorf("MasterUpgrade() is not implemented for provider %s", TestContext.Provider)
	}
}

func etcdUpgradeGCE(target_storage, target_version string) error {
	env := append(
		os.Environ(),
		"TEST_ETCD_VERSION="+target_version,
		"STORAGE_BACKEND="+target_storage,
		"TEST_ETCD_IMAGE=3.0.17")

	_, _, err := RunCmdEnv(env, gceUpgradeScript(), "-l", "-M")
	return err
}

// TODO(mrhohn): Remove this function when kube-proxy is run as a DaemonSet by default.
func MasterUpgradeGCEWithKubeProxyDaemonSet(v string, enableKubeProxyDaemonSet bool) error {
	return masterUpgradeGCE(v, enableKubeProxyDaemonSet)
}

// TODO(mrhohn): Remove 'enableKubeProxyDaemonSet' when kube-proxy is run as a DaemonSet by default.
func masterUpgradeGCE(rawV string, enableKubeProxyDaemonSet bool) error {
	env := append(os.Environ(), fmt.Sprintf("KUBE_PROXY_DAEMONSET=%v", enableKubeProxyDaemonSet))
	// TODO: Remove these variables when they're no longer needed for downgrades.
	if TestContext.EtcdUpgradeVersion != "" && TestContext.EtcdUpgradeStorage != "" {
		env = append(env,
			"TEST_ETCD_VERSION="+TestContext.EtcdUpgradeVersion,
			"STORAGE_BACKEND="+TestContext.EtcdUpgradeStorage,
			"TEST_ETCD_IMAGE=3.0.17")
	}

	v := "v" + rawV
	_, _, err := RunCmdEnv(env, gceUpgradeScript(), "-M", v)
	return err
}

func masterUpgradeGKE(v string) error {
	Logf("Upgrading master to %q", v)
	_, _, err := RunCmd("gcloud", "container",
		"clusters",
		fmt.Sprintf("--project=%s", TestContext.CloudConfig.ProjectID),
		fmt.Sprintf("--zone=%s", TestContext.CloudConfig.Zone),
		"upgrade",
		TestContext.CloudConfig.Cluster,
		"--master",
		fmt.Sprintf("--cluster-version=%s", v),
		"--quiet")
	if err != nil {
		return err
	}

	waitForSSHTunnels()

	return nil
}

func masterUpgradeKubernetesAnywhere(v string) error {
	Logf("Upgrading master to %q", v)

	kaPath := TestContext.KubernetesAnywherePath
	originalConfigPath := filepath.Join(kaPath, ".config")
	backupConfigPath := filepath.Join(kaPath, ".config.bak")
	updatedConfigPath := filepath.Join(kaPath, fmt.Sprintf(".config-%s", v))

	// modify config with specified k8s version
	if _, _, err := RunCmd("sed",
		"-i.bak", // writes original to .config.bak
		fmt.Sprintf(`s/kubernetes_version=.*$/kubernetes_version=%q/`, v),
		originalConfigPath); err != nil {
		return err
	}

	defer func() {
		// revert .config.bak to .config
		if err := os.Rename(backupConfigPath, originalConfigPath); err != nil {
			Logf("Could not rename %s back to %s", backupConfigPath, originalConfigPath)
		}
	}()

	// invoke ka upgrade
	if _, _, err := RunCmd("make", "-C", TestContext.KubernetesAnywherePath,
		"WAIT_FOR_KUBECONFIG=y", "upgrade-master"); err != nil {
		return err
	}

	// move .config to .config.<version>
	if err := os.Rename(originalConfigPath, updatedConfigPath); err != nil {
		return err
	}

	return nil
}

func NodeUpgrade(f *Framework, v string, img string) error {
	// Perform the upgrade.
	var err error
	switch TestContext.Provider {
	case "gce":
		err = nodeUpgradeGCE(v, img, false)
	case "gke":
		err = nodeUpgradeGKE(v, img)
	default:
		err = fmt.Errorf("NodeUpgrade() is not implemented for provider %s", TestContext.Provider)
	}
	if err != nil {
		return err
	}

	// Wait for it to complete and validate nodes are healthy.
	//
	// TODO(ihmccreery) We shouldn't have to wait for nodes to be ready in
	// GKE; the operation shouldn't return until they all are.
	Logf("Waiting up to %v for all nodes to be ready after the upgrade", RestartNodeReadyAgainTimeout)
	if _, err := CheckNodesReady(f.ClientSet, RestartNodeReadyAgainTimeout, TestContext.CloudConfig.NumNodes); err != nil {
		return err
	}
	return nil
}

// TODO(mrhohn): Remove this function when kube-proxy is run as a DaemonSet by default.
func NodeUpgradeGCEWithKubeProxyDaemonSet(f *Framework, v string, img string, enableKubeProxyDaemonSet bool) error {
	// Perform the upgrade.
	if err := nodeUpgradeGCE(v, img, enableKubeProxyDaemonSet); err != nil {
		return err
	}
	// Wait for it to complete and validate nodes are healthy.
	Logf("Waiting up to %v for all nodes to be ready after the upgrade", RestartNodeReadyAgainTimeout)
	if _, err := CheckNodesReady(f.ClientSet, RestartNodeReadyAgainTimeout, TestContext.CloudConfig.NumNodes); err != nil {
		return err
	}
	return nil
}

// TODO(mrhohn): Remove 'enableKubeProxyDaemonSet' when kube-proxy is run as a DaemonSet by default.
func nodeUpgradeGCE(rawV, img string, enableKubeProxyDaemonSet bool) error {
	v := "v" + rawV
	env := append(os.Environ(), fmt.Sprintf("KUBE_PROXY_DAEMONSET=%v", enableKubeProxyDaemonSet))
	if img != "" {
		env = append(env, "KUBE_NODE_OS_DISTRIBUTION="+img)
		_, _, err := RunCmdEnv(env, gceUpgradeScript(), "-N", "-o", v)
		return err
	}
	_, _, err := RunCmdEnv(env, gceUpgradeScript(), "-N", v)
	return err
}

func nodeUpgradeGKE(v string, img string) error {
	Logf("Upgrading nodes to version %q and image %q", v, img)
	args := []string{
		"container",
		"clusters",
		fmt.Sprintf("--project=%s", TestContext.CloudConfig.ProjectID),
		fmt.Sprintf("--zone=%s", TestContext.CloudConfig.Zone),
		"upgrade",
		TestContext.CloudConfig.Cluster,
		fmt.Sprintf("--cluster-version=%s", v),
		"--quiet",
	}
	if len(img) > 0 {
		args = append(args, fmt.Sprintf("--image-type=%s", img))
	}
	_, _, err := RunCmd("gcloud", args...)

	if err != nil {
		return err
	}

	waitForSSHTunnels()

	return nil
}

// CheckNodesReady waits up to nt for expect nodes accessed by c to be ready,
// returning an error if this doesn't happen in time. It returns the names of
// nodes it finds.
func CheckNodesReady(c clientset.Interface, nt time.Duration, expect int) ([]string, error) {
	// First, keep getting all of the nodes until we get the number we expect.
	var nodeList *v1.NodeList
	var errLast error
	start := time.Now()
	found := wait.Poll(Poll, nt, func() (bool, error) {
		// A rolling-update (GCE/GKE implementation of restart) can complete before the apiserver
		// knows about all of the nodes. Thus, we retry the list nodes call
		// until we get the expected number of nodes.
		nodeList, errLast = c.Core().Nodes().List(metav1.ListOptions{
			FieldSelector: fields.Set{"spec.unschedulable": "false"}.AsSelector().String()})
		if errLast != nil {
			return false, nil
		}
		if len(nodeList.Items) != expect {
			errLast = fmt.Errorf("expected to find %d nodes but found only %d (%v elapsed)",
				expect, len(nodeList.Items), time.Since(start))
			Logf("%v", errLast)
			return false, nil
		}
		return true, nil
	}) == nil
	nodeNames := make([]string, len(nodeList.Items))
	for i, n := range nodeList.Items {
		nodeNames[i] = n.ObjectMeta.Name
	}
	if !found {
		return nodeNames, fmt.Errorf("couldn't find %d nodes within %v; last error: %v",
			expect, nt, errLast)
	}
	Logf("Successfully found %d nodes", expect)

	// Next, ensure in parallel that all the nodes are ready. We subtract the
	// time we spent waiting above.
	timeout := nt - time.Since(start)
	result := make(chan bool, len(nodeList.Items))
	for _, n := range nodeNames {
		n := n
		go func() { result <- WaitForNodeToBeReady(c, n, timeout) }()
	}
	failed := false
	// TODO(mbforbes): Change to `for range` syntax once we support only Go
	// >= 1.4.
	for i := range nodeList.Items {
		_ = i
		if !<-result {
			failed = true
		}
	}
	if failed {
		return nodeNames, fmt.Errorf("at least one node failed to be ready")
	}
	return nodeNames, nil
}

// MigTemplate (GCE-only) returns the name of the MIG template that the
// nodes of the cluster use.
func MigTemplate() (string, error) {
	var errLast error
	var templ string
	key := "instanceTemplate"
	if wait.Poll(Poll, SingleCallTimeout, func() (bool, error) {
		// TODO(mikedanese): make this hit the compute API directly instead of
		// shelling out to gcloud.
		// An `instance-groups managed describe` call outputs what we want to stdout.
		output, _, err := retryCmd("gcloud", "compute", "instance-groups", "managed",
			fmt.Sprintf("--project=%s", TestContext.CloudConfig.ProjectID),
			"describe",
			fmt.Sprintf("--zone=%s", TestContext.CloudConfig.Zone),
			TestContext.CloudConfig.NodeInstanceGroup)
		if err != nil {
			errLast = fmt.Errorf("gcloud compute instance-groups managed describe call failed with err: %v", err)
			return false, nil
		}

		// The 'describe' call probably succeeded; parse the output and try to
		// find the line that looks like "instanceTemplate: url/to/<templ>" and
		// return <templ>.
		if val := ParseKVLines(output, key); len(val) > 0 {
			url := strings.Split(val, "/")
			templ = url[len(url)-1]
			Logf("MIG group %s using template: %s", TestContext.CloudConfig.NodeInstanceGroup, templ)
			return true, nil
		}
		errLast = fmt.Errorf("couldn't find %s in output to get MIG template. Output: %s", key, output)
		return false, nil
	}) != nil {
		return "", fmt.Errorf("MigTemplate() failed with last error: %v", errLast)
	}
	return templ, nil
}

func gceUpgradeScript() string {
	if len(TestContext.GCEUpgradeScript) == 0 {
		return path.Join(TestContext.RepoRoot, "cluster/gce/upgrade.sh")
	}
	return TestContext.GCEUpgradeScript
}

func waitForSSHTunnels() {
	Logf("Waiting for SSH tunnels to establish")
	RunKubectl("run", "ssh-tunnel-test",
		"--image="+imageutils.GetBusyBoxImage(),
		"--restart=Never",
		"--command", "--",
		"echo", "Hello")
	defer RunKubectl("delete", "pod", "ssh-tunnel-test")

	// allow up to a minute for new ssh tunnels to establish
	wait.PollImmediate(5*time.Second, time.Minute, func() (bool, error) {
		_, err := RunKubectl("logs", "ssh-tunnel-test")
		return err == nil, nil
	})
}
<|endoftext|>


<commit_before>package iso

import (
	"bytes"
	"fmt"
	"github.com/mitchellh/multistep"
	"github.com/mitchellh/packer/packer"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
)

type StepExport struct {
	Format     string
	SkipExport bool
}

func (s *StepExport) generateArgs(c *Config, outputPath string, hidePassword bool) []string {
	password := url.QueryEscape(c.RemotePassword)
	if hidePassword {
		password = "****"
	}
	return []string{
		"--noSSLVerify=true",
		"--skipManifestCheck",
		"-tt=" + s.Format,
		"vi://" + c.RemoteUser + ":" + password + "@" + c.RemoteHost + "/" + c.VMName,
		outputPath,
	}
}

func (s *StepExport) Run(state multistep.StateBag) multistep.StepAction {
	c := state.Get("config").(*Config)
	ui := state.Get("ui").(packer.Ui)

	// Skip export if requested
	if c.SkipExport {
		ui.Say("Skipping export of virtual machine...")
		return multistep.ActionContinue
	}

	if c.RemoteType != "esx5" || s.Format == "" {
		return multistep.ActionContinue
	}

	ovftool := "ovftool"
	if runtime.GOOS == "windows" {
		ovftool = "ovftool.exe"
	}

	if _, err := exec.LookPath(ovftool); err != nil {
		err := fmt.Errorf("Error %s not found: %s", ovftool, err)
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	// Export the VM
	outputPath := filepath.Join(c.VMName, c.VMName+"."+s.Format)

	if s.Format == "ova" {
		os.MkdirAll(outputPath, 0755)
	}

	ui.Say("Exporting virtual machine...")
	ui.Message(fmt.Sprintf("Executing: %s %s", ovftool, strings.Join(s.generateArgs(c, outputPath, true), " ")))
	var out bytes.Buffer
	cmd := exec.Command(ovftool, s.generateArgs(c, outputPath, false)...)
	cmd.Stdout = &out
	if err := cmd.Run(); err != nil {
		err := fmt.Errorf("Error exporting virtual machine: %s\n%s\n", err, out.String())
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	ui.Message(fmt.Sprintf("%s", out.String()))

	state.Put("exportPath", outputPath)

	return multistep.ActionContinue
}

func (s *StepExport) Cleanup(state multistep.StateBag) {}
<commit_msg>builder/vmware-iso: set ovftool output path<commit_after>package iso

import (
	"bytes"
	"fmt"
	"github.com/mitchellh/multistep"
	"github.com/mitchellh/packer/packer"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
)

type StepExport struct {
	Format     string
	SkipExport bool
}

func (s *StepExport) generateArgs(c *Config, outputPath string, hidePassword bool) []string {
	password := url.QueryEscape(c.RemotePassword)
	if hidePassword {
		password = "****"
	}
	return []string{
		"--noSSLVerify=true",
		"--skipManifestCheck",
		"-tt=" + s.Format,
		"vi://" + c.RemoteUser + ":" + password + "@" + c.RemoteHost + "/" + c.VMName,
		outputPath,
	}
}

func (s *StepExport) Run(state multistep.StateBag) multistep.StepAction {
	c := state.Get("config").(*Config)
	ui := state.Get("ui").(packer.Ui)

	// Skip export if requested
	if c.SkipExport {
		ui.Say("Skipping export of virtual machine...")
		return multistep.ActionContinue
	}

	if c.RemoteType != "esx5" || s.Format == "" {
		return multistep.ActionContinue
	}

	ovftool := "ovftool"
	if runtime.GOOS == "windows" {
		ovftool = "ovftool.exe"
	}

	if _, err := exec.LookPath(ovftool); err != nil {
		err := fmt.Errorf("Error %s not found: %s", ovftool, err)
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	// Export the VM
	outputPath := filepath.Join(c.OutputDir, c.VMName+"."+s.Format)

	if s.Format == "ova" {
		os.MkdirAll(outputPath, 0755)
	}

	ui.Say("Exporting virtual machine...")
	ui.Message(fmt.Sprintf("Executing: %s %s", ovftool, strings.Join(s.generateArgs(c, outputPath, true), " ")))
	var out bytes.Buffer
	cmd := exec.Command(ovftool, s.generateArgs(c, outputPath, false)...)
	cmd.Stdout = &out
	if err := cmd.Run(); err != nil {
		err := fmt.Errorf("Error exporting virtual machine: %s\n%s\n", err, out.String())
		state.Put("error", err)
		ui.Error(err.Error())
		return multistep.ActionHalt
	}

	ui.Message(fmt.Sprintf("%s", out.String()))

	state.Put("exportPath", outputPath)

	return multistep.ActionContinue
}

func (s *StepExport) Cleanup(state multistep.StateBag) {}
<|endoftext|>


<commit_before>package common

import (
	"errors"
	"fmt"
	"github.com/mitchellh/goamz/ec2"
	"github.com/mitchellh/multistep"
	"log"
	"time"
)

// StateRefreshFunc is a function type used for StateChangeConf that is
// responsible for refreshing the item being watched for a state change.
//
// It returns three results. `result` is any object that will be returned
// as the final object after waiting for state change. This allows you to
// return the final updated object, for example an EC2 instance after refreshing
// it.
//
// `state` is the latest state of that object. And `err` is any error that
// may have happened while refreshing the state.
type StateRefreshFunc func() (result interface{}, state string, err error)

// StateChangeConf is the configuration struct used for `WaitForState`.
type StateChangeConf struct {
	Pending   []string
	Refresh   StateRefreshFunc
	StepState multistep.StateBag
	Target    string
}

// AMIStateRefreshFunc returns a StateRefreshFunc that is used to watch
// an AMI for state changes.
func AMIStateRefreshFunc(conn *ec2.EC2, imageId string) StateRefreshFunc {
	return func() (interface{}, string, error) {
		resp, err := conn.Images([]string{imageId}, ec2.NewFilter())
		if err != nil {
			if ec2err, ok := err.(*ec2.Error); ok && ec2err.Code == "InvalidAMIID.NotFound" {
				// Set this to nil as if we didn't find anything.
				resp = nil
			} else {
				log.Printf("Error on AMIStateRefresh: %s", err)
				return nil, "", err
			}
		}

		if resp == nil || len(resp.Images) == 0 {
			// Sometimes AWS has consistency issues and doesn't see the
			// AMI. Return an empty state.
			return nil, "", nil
		}

		i := resp.Images[0]
		return i, i.State, nil
	}
}

// InstanceStateRefreshFunc returns a StateRefreshFunc that is used to watch
// an EC2 instance.
func InstanceStateRefreshFunc(conn *ec2.EC2, i *ec2.Instance) StateRefreshFunc {
	return func() (interface{}, string, error) {
		resp, err := conn.Instances([]string{i.InstanceId}, ec2.NewFilter())
		if err != nil {
			if ec2err, ok := err.(*ec2.Error); ok && ec2err.Code == "InvalidInstanceID.NotFound" {
				// Set this to nil as if we didn't find anything.
				resp = nil
			} else {
				log.Printf("Error on InstanceStateRefresh: %s", err)
				return nil, "", err
			}
		}

		if resp == nil || len(resp.Reservations) == 0 || len(resp.Reservations[0].Instances) == 0 {
			// Sometimes AWS just has consistency issues and doesn't see
			// our instance yet. Return an empty state.
			return nil, "", nil
		}

		i = &resp.Reservations[0].Instances[0]
		return i, i.State.Name, nil
	}
}

// SpotRequestStateRefreshFunc returns a StateRefreshFunc that is used to watch
// a spot request for state changes.
func SpotRequestStateRefreshFunc(conn *ec2.EC2, spotRequestId string) StateRefreshFunc {
	return func() (interface{}, string, error) {
		resp, err := conn.DescribeSpotRequests([]string{spotRequestId}, ec2.NewFilter())
		if err != nil {
			if ec2err, ok := err.(*ec2.Error); ok && ec2err.Code == "InvalidSpotInstanceRequestID.NotFound" {
				// Set this to nil as if we didn't find anything.
				resp = nil
			} else {
				log.Printf("Error on SpotRequestStateRefresh: %s", err)
				return nil, "", err
			}
		}

		if resp == nil || len(resp.SpotRequestResults) == 0 {
			// Sometimes AWS has consistency issues and doesn't see the
			// SpotRequest. Return an empty state.
			return nil, "", nil
		}

		i := resp.SpotRequestResults[0]
		return i, i.State, nil
	}
}

// WaitForState watches an object and waits for it to achieve a certain
// state.
func WaitForState(conf *StateChangeConf) (i interface{}, err error) {
	log.Printf("Waiting for state to become: %s", conf.Target)

	notfoundTick := 0

	for {
		var currentState string
		i, currentState, err = conf.Refresh()
		if err != nil {
			return
		}

		if i == nil {
			// If we didn't find the resource, check if we have been
			// not finding it for awhile, and if so, report an error.
			notfoundTick += 1
			if notfoundTick > 20 {
				return nil, errors.New("couldn't find resource")
			}
		} else {
			// Reset the counter for when a resource isn't found
			notfoundTick = 0

			if currentState == conf.Target {
				return
			}

			if conf.StepState != nil {
				if _, ok := conf.StepState.GetOk(multistep.StateCancelled); ok {
					return nil, errors.New("interrupted")
				}
			}

			found := false
			for _, allowed := range conf.Pending {
				if currentState == allowed {
					found = true
					break
				}
			}

			if !found {
				err := fmt.Errorf("unexpected state '%s', wanted target '%s'", currentState, conf.Target)
				return nil, err
			}
		}

		time.Sleep(2 * time.Second)
	}

	return
}
<commit_msg>fix: Ignore transient network errors when fetching AWS state.<commit_after>package common

import (
	"errors"
	"fmt"
	"github.com/mitchellh/goamz/ec2"
	"github.com/mitchellh/multistep"
	"log"
	"net"
	"time"
)

// StateRefreshFunc is a function type used for StateChangeConf that is
// responsible for refreshing the item being watched for a state change.
//
// It returns three results. `result` is any object that will be returned
// as the final object after waiting for state change. This allows you to
// return the final updated object, for example an EC2 instance after refreshing
// it.
//
// `state` is the latest state of that object. And `err` is any error that
// may have happened while refreshing the state.
type StateRefreshFunc func() (result interface{}, state string, err error)

// StateChangeConf is the configuration struct used for `WaitForState`.
type StateChangeConf struct {
	Pending   []string
	Refresh   StateRefreshFunc
	StepState multistep.StateBag
	Target    string
}

func isTransientNetworkError(err error) bool {
	ret := false
	if nerr, ok := err.(net.Error); ok && nerr.Temporary() {
		ret = true
	}
	return ret
}

// AMIStateRefreshFunc returns a StateRefreshFunc that is used to watch
// an AMI for state changes.
func AMIStateRefreshFunc(conn *ec2.EC2, imageId string) StateRefreshFunc {
	return func() (interface{}, string, error) {
		resp, err := conn.Images([]string{imageId}, ec2.NewFilter())
		if err != nil {
			if ec2err, ok := err.(*ec2.Error); ok && ec2err.Code == "InvalidAMIID.NotFound" {
				// Set this to nil as if we didn't find anything.
				resp = nil
			} else if isTransientNetworkError(err) {
				// Transient network error, treat it as if we didn't find anything
				resp = nil
			} else {
				log.Printf("Error on AMIStateRefresh: %s", err)
				return nil, "", err
			}
		}

		if resp == nil || len(resp.Images) == 0 {
			// Sometimes AWS has consistency issues and doesn't see the
			// AMI. Return an empty state.
			return nil, "", nil
		}

		i := resp.Images[0]
		return i, i.State, nil
	}
}

// InstanceStateRefreshFunc returns a StateRefreshFunc that is used to watch
// an EC2 instance.
func InstanceStateRefreshFunc(conn *ec2.EC2, i *ec2.Instance) StateRefreshFunc {
	return func() (interface{}, string, error) {
		resp, err := conn.Instances([]string{i.InstanceId}, ec2.NewFilter())
		if err != nil {
			if ec2err, ok := err.(*ec2.Error); ok && ec2err.Code == "InvalidInstanceID.NotFound" {
				// Set this to nil as if we didn't find anything.
				resp = nil
			} else if isTransientNetworkError(err) {
				// Transient network error, treat it as if we didn't find anything
				resp = nil
			} else {
				log.Printf("Error on InstanceStateRefresh: %s", err)
				return nil, "", err
			}
		}

		if resp == nil || len(resp.Reservations) == 0 || len(resp.Reservations[0].Instances) == 0 {
			// Sometimes AWS just has consistency issues and doesn't see
			// our instance yet. Return an empty state.
			return nil, "", nil
		}

		i = &resp.Reservations[0].Instances[0]
		return i, i.State.Name, nil
	}
}

// SpotRequestStateRefreshFunc returns a StateRefreshFunc that is used to watch
// a spot request for state changes.
func SpotRequestStateRefreshFunc(conn *ec2.EC2, spotRequestId string) StateRefreshFunc {
	return func() (interface{}, string, error) {
		resp, err := conn.DescribeSpotRequests([]string{spotRequestId}, ec2.NewFilter())
		if err != nil {
			if ec2err, ok := err.(*ec2.Error); ok && ec2err.Code == "InvalidSpotInstanceRequestID.NotFound" {
				// Set this to nil as if we didn't find anything.
				resp = nil
			} else if isTransientNetworkError(err) {
				// Transient network error, treat it as if we didn't find anything
				resp = nil
			} else {
				log.Printf("Error on SpotRequestStateRefresh: %s", err)
				return nil, "", err
			}
		}

		if resp == nil || len(resp.SpotRequestResults) == 0 {
			// Sometimes AWS has consistency issues and doesn't see the
			// SpotRequest. Return an empty state.
			return nil, "", nil
		}

		i := resp.SpotRequestResults[0]
		return i, i.State, nil
	}
}

// WaitForState watches an object and waits for it to achieve a certain
// state.
func WaitForState(conf *StateChangeConf) (i interface{}, err error) {
	log.Printf("Waiting for state to become: %s", conf.Target)

	notfoundTick := 0

	for {
		var currentState string
		i, currentState, err = conf.Refresh()
		if err != nil {
			return
		}

		if i == nil {
			// If we didn't find the resource, check if we have been
			// not finding it for awhile, and if so, report an error.
			notfoundTick += 1
			if notfoundTick > 20 {
				return nil, errors.New("couldn't find resource")
			}
		} else {
			// Reset the counter for when a resource isn't found
			notfoundTick = 0

			if currentState == conf.Target {
				return
			}

			if conf.StepState != nil {
				if _, ok := conf.StepState.GetOk(multistep.StateCancelled); ok {
					return nil, errors.New("interrupted")
				}
			}

			found := false
			for _, allowed := range conf.Pending {
				if currentState == allowed {
					found = true
					break
				}
			}

			if !found {
				err := fmt.Errorf("unexpected state '%s', wanted target '%s'", currentState, conf.Target)
				return nil, err
			}
		}

		time.Sleep(2 * time.Second)
	}

	return
}
<|endoftext|>