<commit_before>package main

import (
	"net/http"
	"strconv"

	"github.com/pilu/traffic"
	"google.golang.org/appengine"
	"google.golang.org/appengine/log"
	"google.golang.org/appengine/search"
	"google.golang.org/appengine/user"
)

// https://cloud.google.com/appengine/docs/go/search/query_strings
// https://godoc.org/google.golang.org/appengine/search
// https://cloud.google.com/appengine/docs/go/search

type SearchData struct {
	Count   int
	IsAdmin bool
}

func SearchHandler(w traffic.ResponseWriter, r *traffic.Request) {
	c := appengine.NewContext(r.Request)
	s_val := r.Request.FormValue("s")

	count := 0
	if s_val != "" {
		index, err := search.Open("entries")
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		iter := index.Search(c, s_val, nil)
		count = iter.Count()
		log.Infof(c, "Search Results for %+v: %+v", s_val, iter)
	}

	w.Render("search", &SearchData{
		Count:   count,
		IsAdmin: user.IsAdmin(c),
	})
}

func SearchWorkHandler(w traffic.ResponseWriter, r *traffic.Request) {
	c := appengine.NewContext(r.Request)
	entries, err := AllPosts(c)
	if err != nil {
		log.Errorf(c, err.Error())
		http.Error(w, err.Error(), 500)
		return
	}

	index, err := search.Open("entries")
	if err != nil {
		log.Errorf(c, err.Error())
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	for _, v := range *entries {
		_, err = index.Put(c, strconv.FormatInt(v.Id, 10), v.SearchDoc())
		log.Infof(c, "Trying to put %+v: %+v", v.Id, strconv.FormatInt(v.Id, 10))
		if err != nil {
			log.Errorf(c, err.Error())
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
	}
}
<commit_msg>working search handler<commit_after>package main

import (
	"net/http"
	"strconv"

	"github.com/pilu/traffic"
	"google.golang.org/appengine"
	"google.golang.org/appengine/log"
	"google.golang.org/appengine/search"
	"google.golang.org/appengine/user"
)

// https://cloud.google.com/appengine/docs/go/search/query_strings
// https://godoc.org/google.golang.org/appengine/search
// https://cloud.google.com/appengine/docs/go/search

type SearchData struct {
	Count   int
	Results *[]Entry
	IsAdmin bool
}

func SearchHandler(w traffic.ResponseWriter, r *traffic.Request) {
	c := appengine.NewContext(r.Request)
	s_val := r.Request.FormValue("s")

	results := []Entry{}

	if s_val != "" {
		index, err := search.Open("entries")
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		iter := index.Search(c, s_val, nil)
		log.Infof(c, "Search Results for %+v: %+v", s_val, iter)

		for t := iter; ; {
			var doc EntrySearch
			id, err := t.Next(&doc)
			if err == search.Done {
				break
			}
			if err != nil {
				log.Errorf(c, "Error iterating: %+v", err)
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}

			log.Debugf(c, "%s -> %#v\n", id, doc)

			entry, err := GetEntry(c, int64(doc.Id))
			if err != nil {
				log.Errorf(c, "Error getting entry: %+v", err)
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}
			results = append(results, *entry)
		}
	}

	w.Render("search", &SearchData{
		Count:   len(results),
		IsAdmin: user.IsAdmin(c),
		Results: &results,
	})
}

func SearchWorkHandler(w traffic.ResponseWriter, r *traffic.Request) {
	c := appengine.NewContext(r.Request)
	entries, err := AllPosts(c)
	if err != nil {
		log.Errorf(c, err.Error())
		http.Error(w, err.Error(), 500)
		return
	}

	index, err := search.Open("entries")
	if err != nil {
		log.Errorf(c, err.Error())
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	for _, v := range *entries {
		_, err = index.Put(c, strconv.FormatInt(v.Id, 10), v.SearchDoc())
		log.Infof(c, "Trying to put %+v: %+v", v.Id, strconv.FormatInt(v.Id, 10))
		if err != nil {
			log.Errorf(c, err.Error())
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"log"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/inominate/apicache"
	"github.com/kr/pretty"
)

var apiClient apicache.Client

// worker tracking
var activeWorkerCount, workerCount int32
var workCount []int32

type apiReq struct {
	apiReq  *apicache.Request
	apiResp *apicache.Response

	apiErr apicache.APIError

	expires time.Time

	worker   int
	httpCode int
	err      error
	respChan chan apiReq
}

// Channel for sending jobs to workers
var workChan chan apiReq

func APIReq(url string, params map[string]string) (*apicache.Response, error) {
	var errorStr string

	if atomic.LoadInt32(&workerCount) <= 0 {
		panic("No workers!")
	}

	// Build the request
	apireq := apicache.NewRequest(url)
	for k, v := range params {
		switch k {
		case "force":
			if v != "" {
				apireq.Force = true
			}
		default:
			apireq.Set(k, v)
		}
	}

	workerID := "C"
	// Don't send it to a worker if we can just yank it fromm the cache
	apiResp, err := apireq.GetCached()
	// 221's are the bane of my existence.
	if err != nil || apireq.Force || apiResp.Error.ErrorCode == 221 {
		for i := 0; i < conf.Retries; i++ {
			respChan := make(chan apiReq)
			req := apiReq{apiReq: apireq, respChan: respChan}
			workChan <- req

			resp := <-respChan
			close(respChan)

			apiResp = resp.apiResp
			err = resp.err
			workerID = fmt.Sprintf("%d", resp.worker)

			// Attempt to recover from server issues, invalidate flag means we
			// believe this is not a server failure.
			if apiResp.Invalidate {
				break
			}
			time.Sleep(2 * time.Second)
			apireq.Force = true
		}
	}
	var singleDebug bool
	if apiResp.Error.ErrorCode == 221 {
		log.Printf("Failed to recover from error 221.")
		singleDebug = true
	}

	// This is similar to the request log, but knows more about where it came from.
	if debug || singleDebug {
		if apiResp.Error.ErrorCode != 0 {
			errorStr = fmt.Sprintf(" Error %d: %s", apiResp.Error.ErrorCode, apiResp.Error.ErrorText)
		}
		logParams := ""
		var paramVal string
		for k, _ := range params {
			if conf.Logging.CensorLog && strings.ToLower(k) == "vcode" && len(params[k]) > 8 {
				paramVal = params[k][0:8] + "..."
			} else {
				paramVal = params[k]
			}
			logParams = fmt.Sprintf("%s&%s=%s", logParams, k, paramVal)
		}
		if logParams != "" {
			logParams = "?" + logParams[1:]
		}

		debugLog.Printf("w%s: %s%s HTTP: %d Expires: %s%s", workerID, url, logParams, apiResp.HTTPCode, apiResp.Expires.Format("2006-01-02 15:04:05"), errorStr)
	}
	return apiResp, err
}

func worker(reqChan chan apiReq, workerID int) {
	atomic.AddInt32(&workerCount, 1)
	var err error

	var eErr error
	var rErr error

	for req := range reqChan {
		var errStr string

		atomic.AddInt32(&activeWorkerCount, 1)

		// Run both of the error limiters simultaneously rather than in
		// sequence. Still need both before we continue.
		errorLimiter := make(chan error)
		rpsLimiter := make(chan error)
		go func() {
			err := errorRateLimiter.Start(30 * time.Second)
			errorLimiter <- err
		}()
		go func() {
			err := rateLimiter.Start(30 * time.Second)
			rpsLimiter <- err
		}()
		eErr = <-errorLimiter
		rErr = <-rpsLimiter

		// Check the error limiter for timeouts
		if eErr != nil {
			err = eErr
			errStr = "error throttling"

			// If the rate limiter didn't timeout be sure to signal it that we
			// didn't do anything.
			if rErr == nil {
				rateLimiter.Finish(true)
			}
		}
		if rErr != nil {
			err = rErr
			if errStr == "" {
				errStr = "rate limiting"
			} else {
				errStr += " and rate limiting"
			}

			// If the error limiter didn't also timeout be sure to signal it that we
			// didn't do anything.
			if eErr == nil {
				errorRateLimiter.Finish(true)
			}
		}
		// We're left with a single err and errStr for returning an error to the client.
		if err != nil {
			log.Printf("Rate Limit Error: %s - %s", errStr, err)
			log.Printf("RPS Events: %d Outstanding: %d", rateLimiter.Count(), rateLimiter.Outstanding())
			log.Printf("Errors Events: %d Outstanding: %d", errorRateLimiter.Count(), errorRateLimiter.Outstanding())
			log.Printf("RPSLimit: %# v", pretty.Formatter(rateLimiter))
			log.Printf("ErrorLimit: %# v", pretty.Formatter(errorRateLimiter))

			req.apiResp = &apicache.Response{
				Data: apicache.SynthesizeAPIError(500,
					fmt.Sprintf("APIProxy Error: Proxy timeout due to %s.", errStr),
					5*time.Minute),
				Expires: time.Now().Add(5 * time.Minute),
				Error: apicache.APIError{500,
					fmt.Sprintf("APIProxy Error: Proxy timeout due to %s.", errStr)},
				HTTPCode: 504,
			}
			req.err = err
		} else {
			resp, err := req.apiReq.Do()
			req.apiResp = resp
			req.err = err
			if resp.Error.ErrorCode == 0 || resp.HTTPCode == 504 || resp.HTTPCode == 418 {
				// 418 means we are currently tempbanned from the API.
				// 504 means the API proxy had some kind of internal or network error.
				//
				// We do not treat these as an error for rate limiting because
				// the apicache library handles it for us, these requests are
				// not actually making it to the CCP API.

				// Finish, but skip recording the event in the rate limiter
				// when there is no error.
				errorRateLimiter.Finish(true)
			} else {
				errorRateLimiter.Finish(false)
			}
			rateLimiter.Finish(false)
		}

		req.worker = workerID
		req.respChan <- req
		atomic.AddInt32(&workCount[workerID], 1)
		atomic.AddInt32(&activeWorkerCount, -1)
	}
	atomic.AddInt32(&workerCount, -1)
}

var startWorkersOnce = &sync.Once{}

func startWorkers() {
	startWorkersOnce.Do(realStartWorkers)
}

func realStartWorkers() {
	log.Printf("Starting %d Workers...", conf.Workers)
	workChan = make(chan apiReq)
	workCount = make([]int32, conf.Workers+1)

	for i := 1; i <= conf.Workers; i++ {
		debugLog.Printf("Starting worker #%d.", i)
		go worker(workChan, i)
	}
}

func stopWorkers() {
	close(workChan)
	for atomic.LoadInt32(&workerCount) > 0 {
		time.Sleep(10 * time.Millisecond)
	}
	startWorkersOnce = &sync.Once{}
}

func PrintWorkerStats(w io.Writer) {
	active, loaded := GetWorkerStats()
	fmt.Fprintf(w, "%d workers idle, %d workers active.\n", loaded-active, active)

	rateCount := rateLimiter.Count()
	rateOutstanding := rateLimiter.Outstanding()

	errorCount := errorRateLimiter.Count()
	errorOutstanding := errorRateLimiter.Outstanding()

	fmt.Fprintf(w, "%d requests in the last second. %d requests outstanding.\n", rateCount, rateOutstanding)
	fmt.Fprintf(w, "%d errors over last %d seconds. %d errors outstanding.\n", errorCount, conf.ErrorPeriod, errorOutstanding)

	for i := int32(1); i <= atomic.LoadInt32(&workerCount); i++ {
		count := atomic.LoadInt32(&workCount[i])
		fmt.Fprintf(w, "   %d: %d\n", i, count)
	}
}

func GetWorkerStats() (int32, int32) {
	return atomic.LoadInt32(&activeWorkerCount), atomic.LoadInt32(&workerCount)
}
<commit_msg>minor retry modification<commit_after>package main

import (
	"fmt"
	"io"
	"log"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/inominate/apicache"
	"github.com/kr/pretty"
)

var apiClient apicache.Client

// worker tracking
var activeWorkerCount, workerCount int32
var workCount []int32

type apiReq struct {
	apiReq  *apicache.Request
	apiResp *apicache.Response

	apiErr apicache.APIError

	expires time.Time

	worker   int
	httpCode int
	err      error
	respChan chan apiReq
}

// Channel for sending jobs to workers
var workChan chan apiReq

func APIReq(url string, params map[string]string) (*apicache.Response, error) {
	var errorStr string

	if atomic.LoadInt32(&workerCount) <= 0 {
		panic("No workers!")
	}

	// Build the request
	apireq := apicache.NewRequest(url)
	for k, v := range params {
		switch k {
		case "force":
			if v != "" {
				apireq.Force = true
			}
		default:
			apireq.Set(k, v)
		}
	}

	workerID := "C"
	// Don't send it to a worker if we can just yank it fromm the cache
	apiResp, err := apireq.GetCached()
	// 221's are the bane of my existence.
	if err != nil || apireq.Force || apiResp.Error.ErrorCode == 221 {
		for i := 0; i < conf.Retries; i++ {
			respChan := make(chan apiReq)
			req := apiReq{apiReq: apireq, respChan: respChan}
			workChan <- req

			resp := <-respChan
			close(respChan)

			apiResp = resp.apiResp
			err = resp.err
			workerID = fmt.Sprintf("%d", resp.worker)

			// Attempt to recover from server issues, invalidate flag means we
			// believe this is not a server failure.
			// 418 is the tempban code
			// 500/900 are panic codes
			if apiResp.Invalidate || apiResp.HTTPCode == 418 || apiResp.HTTPCode == 500 || apiResp.HTTPCode == 900 {
				break
			}
			time.Sleep(2 * time.Second)
			apireq.Force = true
		}
	}
	var singleDebug bool
	if apiResp.Error.ErrorCode == 221 {
		log.Printf("Failed to recover from error 221.")
		singleDebug = true
	}

	// This is similar to the request log, but knows more about where it came from.
	if debug || singleDebug {
		if apiResp.Error.ErrorCode != 0 {
			errorStr = fmt.Sprintf(" Error %d: %s", apiResp.Error.ErrorCode, apiResp.Error.ErrorText)
		}
		logParams := ""
		var paramVal string
		for k, _ := range params {
			if conf.Logging.CensorLog && strings.ToLower(k) == "vcode" && len(params[k]) > 8 {
				paramVal = params[k][0:8] + "..."
			} else {
				paramVal = params[k]
			}
			logParams = fmt.Sprintf("%s&%s=%s", logParams, k, paramVal)
		}
		if logParams != "" {
			logParams = "?" + logParams[1:]
		}

		debugLog.Printf("w%s: %s%s HTTP: %d Expires: %s%s", workerID, url, logParams, apiResp.HTTPCode, apiResp.Expires.Format("2006-01-02 15:04:05"), errorStr)
	}
	return apiResp, err
}

func worker(reqChan chan apiReq, workerID int) {
	atomic.AddInt32(&workerCount, 1)
	var err error

	var eErr error
	var rErr error

	for req := range reqChan {
		var errStr string

		atomic.AddInt32(&activeWorkerCount, 1)

		// Run both of the error limiters simultaneously rather than in
		// sequence. Still need both before we continue.
		errorLimiter := make(chan error)
		rpsLimiter := make(chan error)
		go func() {
			err := errorRateLimiter.Start(30 * time.Second)
			errorLimiter <- err
		}()
		go func() {
			err := rateLimiter.Start(30 * time.Second)
			rpsLimiter <- err
		}()
		eErr = <-errorLimiter
		rErr = <-rpsLimiter

		// Check the error limiter for timeouts
		if eErr != nil {
			err = eErr
			errStr = "error throttling"

			// If the rate limiter didn't timeout be sure to signal it that we
			// didn't do anything.
			if rErr == nil {
				rateLimiter.Finish(true)
			}
		}
		if rErr != nil {
			err = rErr
			if errStr == "" {
				errStr = "rate limiting"
			} else {
				errStr += " and rate limiting"
			}

			// If the error limiter didn't also timeout be sure to signal it that we
			// didn't do anything.
			if eErr == nil {
				errorRateLimiter.Finish(true)
			}
		}
		// We're left with a single err and errStr for returning an error to the client.
		if err != nil {
			log.Printf("Rate Limit Error: %s - %s", errStr, err)
			log.Printf("RPS Events: %d Outstanding: %d", rateLimiter.Count(), rateLimiter.Outstanding())
			log.Printf("Errors Events: %d Outstanding: %d", errorRateLimiter.Count(), errorRateLimiter.Outstanding())
			log.Printf("RPSLimit: %# v", pretty.Formatter(rateLimiter))
			log.Printf("ErrorLimit: %# v", pretty.Formatter(errorRateLimiter))

			req.apiResp = &apicache.Response{
				Data: apicache.SynthesizeAPIError(500,
					fmt.Sprintf("APIProxy Error: Proxy timeout due to %s.", errStr),
					5*time.Minute),
				Expires: time.Now().Add(5 * time.Minute),
				Error: apicache.APIError{500,
					fmt.Sprintf("APIProxy Error: Proxy timeout due to %s.", errStr)},
				HTTPCode: 504,
			}
			req.err = err
		} else {
			resp, err := req.apiReq.Do()
			req.apiResp = resp
			req.err = err
			if resp.Error.ErrorCode == 0 || resp.HTTPCode == 504 || resp.HTTPCode == 418 {
				// 418 means we are currently tempbanned from the API.
				// 504 means the API proxy had some kind of internal or network error.
				//
				// We do not treat these as an error for rate limiting because
				// the apicache library handles it for us, these requests are
				// not actually making it to the CCP API.

				// Finish, but skip recording the event in the rate limiter
				// when there is no error.
				errorRateLimiter.Finish(true)
			} else {
				errorRateLimiter.Finish(false)
			}
			rateLimiter.Finish(false)
		}

		req.worker = workerID
		req.respChan <- req
		atomic.AddInt32(&workCount[workerID], 1)
		atomic.AddInt32(&activeWorkerCount, -1)
	}
	atomic.AddInt32(&workerCount, -1)
}

var startWorkersOnce = &sync.Once{}

func startWorkers() {
	startWorkersOnce.Do(realStartWorkers)
}

func realStartWorkers() {
	log.Printf("Starting %d Workers...", conf.Workers)
	workChan = make(chan apiReq)
	workCount = make([]int32, conf.Workers+1)

	for i := 1; i <= conf.Workers; i++ {
		debugLog.Printf("Starting worker #%d.", i)
		go worker(workChan, i)
	}
}

func stopWorkers() {
	close(workChan)
	for atomic.LoadInt32(&workerCount) > 0 {
		time.Sleep(10 * time.Millisecond)
	}
	startWorkersOnce = &sync.Once{}
}

func PrintWorkerStats(w io.Writer) {
	active, loaded := GetWorkerStats()
	fmt.Fprintf(w, "%d workers idle, %d workers active.\n", loaded-active, active)

	rateCount := rateLimiter.Count()
	rateOutstanding := rateLimiter.Outstanding()

	errorCount := errorRateLimiter.Count()
	errorOutstanding := errorRateLimiter.Outstanding()

	fmt.Fprintf(w, "%d requests in the last second. %d requests outstanding.\n", rateCount, rateOutstanding)
	fmt.Fprintf(w, "%d errors over last %d seconds. %d errors outstanding.\n", errorCount, conf.ErrorPeriod, errorOutstanding)

	for i := int32(1); i <= atomic.LoadInt32(&workerCount); i++ {
		count := atomic.LoadInt32(&workCount[i])
		fmt.Fprintf(w, "   %d: %d\n", i, count)
	}
}

func GetWorkerStats() (int32, int32) {
	return atomic.LoadInt32(&activeWorkerCount), atomic.LoadInt32(&workerCount)
}
<|endoftext|>


<commit_before>// Copyright (C) 2010, Kyle Lemons <kyle@kylelemons.net>.  All rights reserved.

package log4go

import (
	"errors"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

var (
	Global Logger
)

func init() {
	// auto load config from default position
	Global = NewDefaultLogger(DEBUG)
	file, _ := exec.LookPath(os.Args[0])
	dir := filepath.Dir(file)
	if _, err := os.Stat("log4go.xml"); !os.IsNotExist(err) {
		Global.LoadConfiguration("log4go.xml")
	} else if _, err := os.Stat(filepath.Join(dir, "/log4go.xml")); !os.IsNotExist(err) {
		Global.LoadConfiguration(filepath.Join(dir, "log4go.xml"))
	} else if _, err := os.Stat(filepath.Join(dir, "/conf/log4go.xml")); !os.IsNotExist(err) {
		Global.LoadConfiguration(filepath.Join(dir, "/conf/log4go.xml"))
	} else {
		fmt.Fprintf(os.Stdout, "log4go config not found, exec dir is: %s, u need to load it by yourself.\n", dir)
	}
}

// setup by config string, not a config file
func Setup(config []byte) {
	Global.Config(config)
}

// Wrapper for (*Logger).LoadConfiguration
func LoadConfiguration(filename string) {
	Global.LoadConfiguration(filename)
}

// Wrapper for (*Logger).AddFilter
func AddFilter(name string, lvl Level, writer LogWriter) {
	Global.AddFilter(name, lvl, writer)
}

// Wrapper for (*Logger).Close (closes and removes all logwriters)
func Close() {
	Global.Close()
}

func Crash(args ...interface{}) {
	if len(args) > 0 {
		Global.intLogf(CRITICAL, strings.Repeat(" %v", len(args))[1:], args...)
	}
	panic(args)
}

// Logs the given message and crashes the program
func Crashf(format string, args ...interface{}) {
	Global.intLogf(CRITICAL, format, args...)
	Global.Close() // so that hopefully the messages get logged
	panic(fmt.Sprintf(format, args...))
}

// Compatibility with `log`
func Exit(args ...interface{}) {
	if len(args) > 0 {
		Global.intLogf(ERROR, strings.Repeat(" %v", len(args))[1:], args...)
	}
	Global.Close() // so that hopefully the messages get logged
	os.Exit(0)
}

// Compatibility with `log`
func Exitf(format string, args ...interface{}) {
	Global.intLogf(ERROR, format, args...)
	Global.Close() // so that hopefully the messages get logged
	os.Exit(0)
}

// Compatibility with `log`
func Stderr(args ...interface{}) {
	if len(args) > 0 {
		Global.intLogf(ERROR, strings.Repeat(" %v", len(args))[1:], args...)
	}
}

// Compatibility with `log`
func Stderrf(format string, args ...interface{}) {
	Global.intLogf(ERROR, format, args...)
}

// Compatibility with `log`
func Stdout(args ...interface{}) {
	if len(args) > 0 {
		Global.intLogf(INFO, strings.Repeat(" %v", len(args))[1:], args...)
	}
}

// Compatibility with `log`
func Stdoutf(format string, args ...interface{}) {
	Global.intLogf(INFO, format, args...)
}

// Send a log message manually
// Wrapper for (*Logger).Log
func Log(lvl Level, source, message string) {
	Global.Log(lvl, source, message)
}

// Send a formatted log message easily
// Wrapper for (*Logger).Logf
func Logf(lvl Level, format string, args ...interface{}) {
	Global.intLogf(lvl, format, args...)
}

// Send a closure log message
// Wrapper for (*Logger).Logc
func Logc(lvl Level, closure func() string) {
	Global.intLogc(lvl, closure)
}

// Utility for finest log messages (see Debug() for parameter explanation)
// Wrapper for (*Logger).Finest
func Finest(arg0 interface{}, args ...interface{}) {
	const (
		lvl = FINEST
	)
	switch first := arg0.(type) {
	case string:
		// Use the string as a format string
		Global.intLogf(lvl, first, args...)
	case func() string:
		// Log the closure (no other arguments used)
		Global.intLogc(lvl, first)
	default:
		// Build a format string so that it will be similar to Sprint
		Global.intLogf(lvl, fmt.Sprint(arg0)+strings.Repeat(" %v", len(args)), args...)
	}
}

// Utility for fine log messages (see Debug() for parameter explanation)
// Wrapper for (*Logger).Fine
func Fine(arg0 interface{}, args ...interface{}) {
	const (
		lvl = FINE
	)
	switch first := arg0.(type) {
	case string:
		// Use the string as a format string
		Global.intLogf(lvl, first, args...)
	case func() string:
		// Log the closure (no other arguments used)
		Global.intLogc(lvl, first)
	default:
		// Build a format string so that it will be similar to Sprint
		Global.intLogf(lvl, fmt.Sprint(arg0)+strings.Repeat(" %v", len(args)), args...)
	}
}

// Utility for debug log messages
// When given a string as the first argument, this behaves like Logf but with the DEBUG log level (e.g. the first argument is interpreted as a format for the latter arguments)
// When given a closure of type func()string, this logs the string returned by the closure iff it will be logged.  The closure runs at most one time.
// When given anything else, the log message will be each of the arguments formatted with %v and separated by spaces (ala Sprint).
// Wrapper for (*Logger).Debug
func Debug(arg0 interface{}, args ...interface{}) {
	const (
		lvl = DEBUG
	)
	switch first := arg0.(type) {
	case string:
		// Use the string as a format string
		Global.intLogf(lvl, first, args...)
	case func() string:
		// Log the closure (no other arguments used)
		Global.intLogc(lvl, first)
	default:
		// Build a format string so that it will be similar to Sprint
		Global.intLogf(lvl, fmt.Sprint(arg0)+strings.Repeat(" %v", len(args)), args...)
	}
}

// Utility for trace log messages (see Debug() for parameter explanation)
// Wrapper for (*Logger).Trace
func Trace(arg0 interface{}, args ...interface{}) {
	const (
		lvl = TRACE
	)
	switch first := arg0.(type) {
	case string:
		// Use the string as a format string
		Global.intLogf(lvl, first, args...)
	case func() string:
		// Log the closure (no other arguments used)
		Global.intLogc(lvl, first)
	default:
		// Build a format string so that it will be similar to Sprint
		Global.intLogf(lvl, fmt.Sprint(arg0)+strings.Repeat(" %v", len(args)), args...)
	}
}

// Utility for info log messages (see Debug() for parameter explanation)
// Wrapper for (*Logger).Info
func Info(arg0 interface{}, args ...interface{}) {
	const (
		lvl = INFO
	)
	switch first := arg0.(type) {
	case string:
		// Use the string as a format string
		Global.intLogf(lvl, first, args...)
	case func() string:
		// Log the closure (no other arguments used)
		Global.intLogc(lvl, first)
	default:
		// Build a format string so that it will be similar to Sprint
		Global.intLogf(lvl, fmt.Sprint(arg0)+strings.Repeat(" %v", len(args)), args...)
	}
}

// Utility for Access log messages (see Debug() for parameter explanation)
// Wrapper for (*Logger).Info
func Access(arg0 interface{}, args ...interface{}) {
	const (
		lvl = ACCESS
	)
	switch first := arg0.(type) {
	case string:
		// Use the string as a format string
		Global.intLogf(lvl, first, args...)
	case func() string:
		// Log the closure (no other arguments used)
		Global.intLogc(lvl, first)
	default:
		// Build a format string so that it will be similar to Sprint
		Global.intLogf(lvl, fmt.Sprint(arg0)+strings.Repeat(" %v", len(args)), args...)
	}
}

// Utility for warn log messages (returns an error for easy function returns) (see Debug() for parameter explanation)
// These functions will execute a closure exactly once, to build the error message for the return
// Wrapper for (*Logger).Warn
func Warn(arg0 interface{}, args ...interface{}) error {
	const (
		lvl = WARNING
	)
	switch first := arg0.(type) {
	case string:
		// Use the string as a format string
		Global.intLogf(lvl, first, args...)
		return errors.New(fmt.Sprintf(first, args...))
	case func() string:
		// Log the closure (no other arguments used)
		str := first()
		Global.intLogf(lvl, "%s", str)
		return errors.New(str)
	default:
		// Build a format string so that it will be similar to Sprint
		Global.intLogf(lvl, fmt.Sprint(first)+strings.Repeat(" %v", len(args)), args...)
		return errors.New(fmt.Sprint(first) + fmt.Sprintf(strings.Repeat(" %v", len(args)), args...))
	}
	return nil
}

// Utility for error log messages (returns an error for easy function returns) (see Debug() for parameter explanation)
// These functions will execute a closure exactly once, to build the error message for the return
// Wrapper for (*Logger).Error
func Error(arg0 interface{}, args ...interface{}) error {
	const (
		lvl = ERROR
	)
	switch first := arg0.(type) {
	case string:
		// Use the string as a format string
		Global.intLogf(lvl, first, args...)
		return errors.New(fmt.Sprintf(first, args...))
	case func() string:
		// Log the closure (no other arguments used)
		str := first()
		Global.intLogf(lvl, "%s", str)
		return errors.New(str)
	default:
		// Build a format string so that it will be similar to Sprint
		Global.intLogf(lvl, fmt.Sprint(first)+strings.Repeat(" %v", len(args)), args...)
		return errors.New(fmt.Sprint(first) + fmt.Sprintf(strings.Repeat(" %v", len(args)), args...))
	}
	return nil
}

// Utility for critical log messages (returns an error for easy function returns) (see Debug() for parameter explanation)
// These functions will execute a closure exactly once, to build the error message for the return
// Wrapper for (*Logger).Critical. This method will log the call stack
func Critical(arg0 interface{}, args ...interface{}) error {
	const (
		lvl = CRITICAL
	)
	switch first := arg0.(type) {
	case string:
		// Use the string as a format string
		msg := fmt.Sprintf("%s\n%s", fmt.Sprintf(first, args...), CallStack(3))
		Global.intLogf(lvl, msg)
		//Global.intLogf(lvl, "%s", CallStack(3))
		return errors.New(fmt.Sprintf(first, args...))
	case func() string:
		// Log the closure (no other arguments used)
		str := first()
		Global.intLogf(lvl, "%s\n%s", str, CallStack(3))
		//Global.intLogf(lvl, "%s", CallStack(3))
		return errors.New(str)
	case func(interface{}) string:
		str := first(args[0])
		Global.intLogf(lvl, "%s\n%s", str, CallStack(3))
		return errors.New(str)
	default:
		// Build a format string so that it will be similar to Sprint
		msg := fmt.Sprintf("%s\n%s", fmt.Sprint(first) + fmt.Sprintf(strings.Repeat(" %v", len(args)), args...), CallStack(3))
		Global.intLogf(lvl, msg)
		return errors.New(fmt.Sprint(first) + fmt.Sprintf(strings.Repeat(" %v", len(args)), args...))
	}
	return nil
}

// Recover used to log the stack when panic occur.
// usage: defer log4go.Recover("this is a msg: %v", "msg")
// or:
//      defer log4go.Recover(func(err interface{}) string {
//          // ... your code here, return the error message
//          return fmt.Sprintf("recover..v1=%v;v2=%v;err=%v", 1, 2, err)
//      })
func Recover(arg0 interface{}, args ...interface{}) {
	if err := recover(); err != nil {
		switch a := arg0.(type) {
		case func(interface{}) string:
			// the recovered err will pass to this func
			Critical(arg0, append([]interface{}{err}, args)...)
		case string:
			Critical(a+"\n%v", append(args, err)...)
		default:
			Critical(arg0, append(args, err)...)
		}
	}
}
<commit_msg>update log print<commit_after>// Copyright (C) 2010, Kyle Lemons <kyle@kylelemons.net>.  All rights reserved.

package log4go

import (
	"errors"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

var (
	Global Logger
)

func init() {
	// auto load config from default position
	Global = NewDefaultLogger(DEBUG)
	file, _ := exec.LookPath(os.Args[0])
	dir := filepath.Dir(file)
	if _, err := os.Stat("log4go.xml"); !os.IsNotExist(err) {
		Global.LoadConfiguration("log4go.xml")
	} else if _, err := os.Stat(filepath.Join(dir, "/log4go.xml")); !os.IsNotExist(err) {
		Global.LoadConfiguration(filepath.Join(dir, "log4go.xml"))
	} else if _, err := os.Stat(filepath.Join(dir, "/conf/log4go.xml")); !os.IsNotExist(err) {
		Global.LoadConfiguration(filepath.Join(dir, "/conf/log4go.xml"))
	} else {
		fmt.Fprintf(os.Stderr, "log4go config not found, exec dir is: %s, u need to load it by yourself.\n", dir)
	}
}

// setup by config string, not a config file
func Setup(config []byte) {
	Global.Config(config)
}

// Wrapper for (*Logger).LoadConfiguration
func LoadConfiguration(filename string) {
	Global.LoadConfiguration(filename)
}

// Wrapper for (*Logger).AddFilter
func AddFilter(name string, lvl Level, writer LogWriter) {
	Global.AddFilter(name, lvl, writer)
}

// Wrapper for (*Logger).Close (closes and removes all logwriters)
func Close() {
	Global.Close()
}

func Crash(args ...interface{}) {
	if len(args) > 0 {
		Global.intLogf(CRITICAL, strings.Repeat(" %v", len(args))[1:], args...)
	}
	panic(args)
}

// Logs the given message and crashes the program
func Crashf(format string, args ...interface{}) {
	Global.intLogf(CRITICAL, format, args...)
	Global.Close() // so that hopefully the messages get logged
	panic(fmt.Sprintf(format, args...))
}

// Compatibility with `log`
func Exit(args ...interface{}) {
	if len(args) > 0 {
		Global.intLogf(ERROR, strings.Repeat(" %v", len(args))[1:], args...)
	}
	Global.Close() // so that hopefully the messages get logged
	os.Exit(0)
}

// Compatibility with `log`
func Exitf(format string, args ...interface{}) {
	Global.intLogf(ERROR, format, args...)
	Global.Close() // so that hopefully the messages get logged
	os.Exit(0)
}

// Compatibility with `log`
func Stderr(args ...interface{}) {
	if len(args) > 0 {
		Global.intLogf(ERROR, strings.Repeat(" %v", len(args))[1:], args...)
	}
}

// Compatibility with `log`
func Stderrf(format string, args ...interface{}) {
	Global.intLogf(ERROR, format, args...)
}

// Compatibility with `log`
func Stdout(args ...interface{}) {
	if len(args) > 0 {
		Global.intLogf(INFO, strings.Repeat(" %v", len(args))[1:], args...)
	}
}

// Compatibility with `log`
func Stdoutf(format string, args ...interface{}) {
	Global.intLogf(INFO, format, args...)
}

// Send a log message manually
// Wrapper for (*Logger).Log
func Log(lvl Level, source, message string) {
	Global.Log(lvl, source, message)
}

// Send a formatted log message easily
// Wrapper for (*Logger).Logf
func Logf(lvl Level, format string, args ...interface{}) {
	Global.intLogf(lvl, format, args...)
}

// Send a closure log message
// Wrapper for (*Logger).Logc
func Logc(lvl Level, closure func() string) {
	Global.intLogc(lvl, closure)
}

// Utility for finest log messages (see Debug() for parameter explanation)
// Wrapper for (*Logger).Finest
func Finest(arg0 interface{}, args ...interface{}) {
	const (
		lvl = FINEST
	)
	switch first := arg0.(type) {
	case string:
		// Use the string as a format string
		Global.intLogf(lvl, first, args...)
	case func() string:
		// Log the closure (no other arguments used)
		Global.intLogc(lvl, first)
	default:
		// Build a format string so that it will be similar to Sprint
		Global.intLogf(lvl, fmt.Sprint(arg0)+strings.Repeat(" %v", len(args)), args...)
	}
}

// Utility for fine log messages (see Debug() for parameter explanation)
// Wrapper for (*Logger).Fine
func Fine(arg0 interface{}, args ...interface{}) {
	const (
		lvl = FINE
	)
	switch first := arg0.(type) {
	case string:
		// Use the string as a format string
		Global.intLogf(lvl, first, args...)
	case func() string:
		// Log the closure (no other arguments used)
		Global.intLogc(lvl, first)
	default:
		// Build a format string so that it will be similar to Sprint
		Global.intLogf(lvl, fmt.Sprint(arg0)+strings.Repeat(" %v", len(args)), args...)
	}
}

// Utility for debug log messages
// When given a string as the first argument, this behaves like Logf but with the DEBUG log level (e.g. the first argument is interpreted as a format for the latter arguments)
// When given a closure of type func()string, this logs the string returned by the closure iff it will be logged.  The closure runs at most one time.
// When given anything else, the log message will be each of the arguments formatted with %v and separated by spaces (ala Sprint).
// Wrapper for (*Logger).Debug
func Debug(arg0 interface{}, args ...interface{}) {
	const (
		lvl = DEBUG
	)
	switch first := arg0.(type) {
	case string:
		// Use the string as a format string
		Global.intLogf(lvl, first, args...)
	case func() string:
		// Log the closure (no other arguments used)
		Global.intLogc(lvl, first)
	default:
		// Build a format string so that it will be similar to Sprint
		Global.intLogf(lvl, fmt.Sprint(arg0)+strings.Repeat(" %v", len(args)), args...)
	}
}

// Utility for trace log messages (see Debug() for parameter explanation)
// Wrapper for (*Logger).Trace
func Trace(arg0 interface{}, args ...interface{}) {
	const (
		lvl = TRACE
	)
	switch first := arg0.(type) {
	case string:
		// Use the string as a format string
		Global.intLogf(lvl, first, args...)
	case func() string:
		// Log the closure (no other arguments used)
		Global.intLogc(lvl, first)
	default:
		// Build a format string so that it will be similar to Sprint
		Global.intLogf(lvl, fmt.Sprint(arg0)+strings.Repeat(" %v", len(args)), args...)
	}
}

// Utility for info log messages (see Debug() for parameter explanation)
// Wrapper for (*Logger).Info
func Info(arg0 interface{}, args ...interface{}) {
	const (
		lvl = INFO
	)
	switch first := arg0.(type) {
	case string:
		// Use the string as a format string
		Global.intLogf(lvl, first, args...)
	case func() string:
		// Log the closure (no other arguments used)
		Global.intLogc(lvl, first)
	default:
		// Build a format string so that it will be similar to Sprint
		Global.intLogf(lvl, fmt.Sprint(arg0)+strings.Repeat(" %v", len(args)), args...)
	}
}

// Utility for Access log messages (see Debug() for parameter explanation)
// Wrapper for (*Logger).Info
func Access(arg0 interface{}, args ...interface{}) {
	const (
		lvl = ACCESS
	)
	switch first := arg0.(type) {
	case string:
		// Use the string as a format string
		Global.intLogf(lvl, first, args...)
	case func() string:
		// Log the closure (no other arguments used)
		Global.intLogc(lvl, first)
	default:
		// Build a format string so that it will be similar to Sprint
		Global.intLogf(lvl, fmt.Sprint(arg0)+strings.Repeat(" %v", len(args)), args...)
	}
}

// Utility for warn log messages (returns an error for easy function returns) (see Debug() for parameter explanation)
// These functions will execute a closure exactly once, to build the error message for the return
// Wrapper for (*Logger).Warn
func Warn(arg0 interface{}, args ...interface{}) error {
	const (
		lvl = WARNING
	)
	switch first := arg0.(type) {
	case string:
		// Use the string as a format string
		Global.intLogf(lvl, first, args...)
		return errors.New(fmt.Sprintf(first, args...))
	case func() string:
		// Log the closure (no other arguments used)
		str := first()
		Global.intLogf(lvl, "%s", str)
		return errors.New(str)
	default:
		// Build a format string so that it will be similar to Sprint
		Global.intLogf(lvl, fmt.Sprint(first)+strings.Repeat(" %v", len(args)), args...)
		return errors.New(fmt.Sprint(first) + fmt.Sprintf(strings.Repeat(" %v", len(args)), args...))
	}
	return nil
}

// Utility for error log messages (returns an error for easy function returns) (see Debug() for parameter explanation)
// These functions will execute a closure exactly once, to build the error message for the return
// Wrapper for (*Logger).Error
func Error(arg0 interface{}, args ...interface{}) error {
	const (
		lvl = ERROR
	)
	switch first := arg0.(type) {
	case string:
		// Use the string as a format string
		Global.intLogf(lvl, first, args...)
		return errors.New(fmt.Sprintf(first, args...))
	case func() string:
		// Log the closure (no other arguments used)
		str := first()
		Global.intLogf(lvl, "%s", str)
		return errors.New(str)
	default:
		// Build a format string so that it will be similar to Sprint
		Global.intLogf(lvl, fmt.Sprint(first)+strings.Repeat(" %v", len(args)), args...)
		return errors.New(fmt.Sprint(first) + fmt.Sprintf(strings.Repeat(" %v", len(args)), args...))
	}
	return nil
}

// Utility for critical log messages (returns an error for easy function returns) (see Debug() for parameter explanation)
// These functions will execute a closure exactly once, to build the error message for the return
// Wrapper for (*Logger).Critical. This method will log the call stack
func Critical(arg0 interface{}, args ...interface{}) error {
	const (
		lvl = CRITICAL
	)
	switch first := arg0.(type) {
	case string:
		// Use the string as a format string
		msg := fmt.Sprintf("%s\n%s", fmt.Sprintf(first, args...), CallStack(3))
		Global.intLogf(lvl, msg)
		//Global.intLogf(lvl, "%s", CallStack(3))
		return errors.New(fmt.Sprintf(first, args...))
	case func() string:
		// Log the closure (no other arguments used)
		str := first()
		Global.intLogf(lvl, "%s\n%s", str, CallStack(3))
		//Global.intLogf(lvl, "%s", CallStack(3))
		return errors.New(str)
	case func(interface{}) string:
		str := first(args[0])
		Global.intLogf(lvl, "%s\n%s", str, CallStack(3))
		return errors.New(str)
	default:
		// Build a format string so that it will be similar to Sprint
		msg := fmt.Sprintf("%s\n%s", fmt.Sprint(first) + fmt.Sprintf(strings.Repeat(" %v", len(args)), args...), CallStack(3))
		Global.intLogf(lvl, msg)
		return errors.New(fmt.Sprint(first) + fmt.Sprintf(strings.Repeat(" %v", len(args)), args...))
	}
	return nil
}

// Recover used to log the stack when panic occur.
// usage: defer log4go.Recover("this is a msg: %v", "msg")
// or:
//      defer log4go.Recover(func(err interface{}) string {
//          // ... your code here, return the error message
//          return fmt.Sprintf("recover..v1=%v;v2=%v;err=%v", 1, 2, err)
//      })
func Recover(arg0 interface{}, args ...interface{}) {
	if err := recover(); err != nil {
		switch a := arg0.(type) {
		case func(interface{}) string:
			// the recovered err will pass to this func
			Critical(arg0, append([]interface{}{err}, args)...)
		case string:
			Critical(a+"\n%v", append(args, err)...)
		default:
			Critical(arg0, append(args, err)...)
		}
	}
}
<|endoftext|>


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

package ws

import (
	"fmt"
	"io"
	"net"
	"net/http"
	"time"

	"github.com/jseun/gopivo"
	"github.com/gorilla/websocket"
)

const (
	defaultBacklogSize = 64
	defaultPingTimeout = 60 * time.Second

	writeWaitTime = 10 * time.Second
)

var upgrader = &websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

type Conn struct {
	backlog int
	port    chan []byte

	pingTimeout time.Duration

	ws *websocket.Conn
}

func NewConn() *Conn {
	return &Conn{
		backlog:     defaultBacklogSize,
		pingTimeout: defaultPingTimeout,
		ws:          &websocket.Conn{},
	}
}

func (c *Conn) ping() error {
	return c.write(websocket.PingMessage, []byte{})
}

func (c *Conn) send(buf []byte) (err error) {
	return c.write(websocket.TextMessage, buf)
}

func (c *Conn) write(t int, buf []byte) error {
	c.ws.SetWriteDeadline(time.Now().Add(writeWaitTime))
	return c.ws.WriteMessage(t, buf)
}

func (c *Conn) Close(err error) error {
	defer c.ws.Close()
	code := websocket.CloseNormalClosure
	msg := websocket.FormatCloseMessage(code, fmt.Sprint(err))
	wait := time.Now().Add(writeWaitTime)
	return c.ws.WriteControl(websocket.CloseMessage, msg, wait)
}

func (c *Conn) Dial(url string, h http.Header) (*Conn, *http.Response, error) {
	var dialer = &websocket.Dialer{}
	ws, r, err := dialer.Dial(url, h)
	if err != nil {
		return nil, r, err
	}
	c.ws = ws
	return c, r, nil
}

func (c *Conn) Receiver(rc gopivo.OnReadCloser) error {
	defer c.ws.Close()
	c.ws.SetReadDeadline(time.Now().Add(c.pingTimeout))
	c.ws.SetPongHandler(func(string) error {
		c.ws.SetReadDeadline(time.Now().Add(c.pingTimeout))
		return nil
	})

	for {
		msgt, data, err := c.ws.ReadMessage()
		switch {
		case err == io.EOF:
			return rc.OnClose(nil)
		case err != nil:
			return rc.OnClose(err)
		case msgt == websocket.BinaryMessage:
			if err := rc.OnReadBinary(data); err != nil {
				return rc.OnClose(err)
			}
		case msgt == websocket.TextMessage:
			if err := rc.OnReadText(data); err != nil {
				return rc.OnClose(err)
			}
		}
	}
}

func (c *Conn) RemoteAddr() net.Addr {
	return c.ws.RemoteAddr()
}

func (c *Conn) Sender() chan []byte {
	c.port = make(chan []byte, c.backlog)
	pingInterval := (c.pingTimeout * 9) / 10
	pinger := time.NewTicker(pingInterval)
	go func() {
		defer pinger.Stop()
		for {
			select {
			case msg, ok := <-c.port:
				if !ok {
					return
				}

				if err := c.send(msg); err != nil {
					c.ws.Close()
					return
				}
			case <-pinger.C:
				if err := c.ping(); err != nil {
					c.ws.Close()
					return
				}
			}
		}
	}()
	return c.port
}

func (c *Conn) Upgrade(w http.ResponseWriter, r *http.Request, h http.Header) error {
	ws, err := upgrader.Upgrade(w, r, h)
	if err != nil {
		return err
	}
	c.ws = ws
	return nil
}
<commit_msg>Close() should not close the websocket right away<commit_after>// Copyright 2015 The GoPivo Authors. All rights reserved.
// Use of this source code is governed by a Simplified BSD
// license that can be found in the LICENSE file.

package ws

import (
	"fmt"
	"io"
	"net"
	"net/http"
	"time"

	"github.com/jseun/gopivo"
	"github.com/gorilla/websocket"
)

const (
	defaultBacklogSize = 64
	defaultPingTimeout = 60 * time.Second

	writeWaitTime = 10 * time.Second
)

var upgrader = &websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

type Conn struct {
	backlog int
	port    chan []byte

	pingTimeout time.Duration

	ws *websocket.Conn
}

func NewConn() *Conn {
	return &Conn{
		backlog:     defaultBacklogSize,
		pingTimeout: defaultPingTimeout,
		ws:          &websocket.Conn{},
	}
}

func (c *Conn) ping() error {
	return c.write(websocket.PingMessage, []byte{})
}

func (c *Conn) send(buf []byte) (err error) {
	return c.write(websocket.TextMessage, buf)
}

func (c *Conn) write(t int, buf []byte) error {
	c.ws.SetWriteDeadline(time.Now().Add(writeWaitTime))
	return c.ws.WriteMessage(t, buf)
}

func (c *Conn) Close(err error) error {
	code := websocket.CloseNormalClosure
	msg := websocket.FormatCloseMessage(code, fmt.Sprint(err))
	wait := time.Now().Add(writeWaitTime)
	return c.ws.WriteControl(websocket.CloseMessage, msg, wait)
}

func (c *Conn) Dial(url string, h http.Header) (*Conn, *http.Response, error) {
	var dialer = &websocket.Dialer{}
	ws, r, err := dialer.Dial(url, h)
	if err != nil {
		return nil, r, err
	}
	c.ws = ws
	return c, r, nil
}

func (c *Conn) Receiver(rc gopivo.OnReadCloser) error {
	defer c.ws.Close()
	c.ws.SetReadDeadline(time.Now().Add(c.pingTimeout))
	c.ws.SetPongHandler(func(string) error {
		c.ws.SetReadDeadline(time.Now().Add(c.pingTimeout))
		return nil
	})

	for {
		msgt, data, err := c.ws.ReadMessage()
		switch {
		case err == io.EOF:
			return rc.OnClose(nil)
		case err != nil:
			return rc.OnClose(err)
		case msgt == websocket.BinaryMessage:
			if err := rc.OnReadBinary(data); err != nil {
				return rc.OnClose(err)
			}
		case msgt == websocket.TextMessage:
			if err := rc.OnReadText(data); err != nil {
				return rc.OnClose(err)
			}
		}
	}
}

func (c *Conn) RemoteAddr() net.Addr {
	return c.ws.RemoteAddr()
}

func (c *Conn) Sender() chan []byte {
	c.port = make(chan []byte, c.backlog)
	pingInterval := (c.pingTimeout * 9) / 10
	pinger := time.NewTicker(pingInterval)
	go func() {
		defer pinger.Stop()
		for {
			select {
			case msg, ok := <-c.port:
				if !ok {
					return
				}

				if err := c.send(msg); err != nil {
					c.ws.Close()
					return
				}
			case <-pinger.C:
				if err := c.ping(); err != nil {
					c.ws.Close()
					return
				}
			}
		}
	}()
	return c.port
}

func (c *Conn) Upgrade(w http.ResponseWriter, r *http.Request, h http.Header) error {
	ws, err := upgrader.Upgrade(w, r, h)
	if err != nil {
		return err
	}
	c.ws = ws
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2014 Google Inc. All rights reserved.
//
// 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.

// The Blueprint bootstrapping mechanism is intended to enable building a source
// tree using a Blueprint-based build system that is embedded (as source) in
// that source tree.  The only prerequisites for performing such a build are:
//
//   1. A Ninja binary
//   2. A script interpreter (e.g. Bash or Python)
//   3. A Go toolchain
//
// The Primary Builder
//
// As part of the bootstrapping process, a binary called the "primary builder"
// is created.  This primary builder is the binary that includes both the core
// Blueprint library and the build logic specific to the source tree.  It is
// used to generate the Ninja file that describes how to build the entire source
// tree.
//
// The primary builder must be a pure Go (i.e. no cgo) module built with the
// module type 'bootstrap_go_binary'.  It should have the 'primaryBuilder'
// module property set to true in its Blueprints file.  If more than one module
// sets primaryBuilder to true the build will fail.
//
// The primary builder main function should look something like:
//
//   package main
//
//   import (
//       "flag"
//       "github.com/google/blueprint"
//       "github.com/google/blueprint/bootstrap"
//       "path/filepath"
//
//       "my/custom/build/logic"
//   )
//
//   func main() {
//       // The primary builder should use the global flag set because the
//       // bootstrap package registers its own flags there.
//       flag.Parse()
//
//       // The top-level Blueprints file is passed as the first argument.
//       srcDir := filepath.Dir(flag.Arg(0))
//
//       // Create the build context.
//       ctx := blueprint.NewContext()
//
//       // Register custom module types
//       ctx.RegisterModuleType("foo", logic.FooModule)
//       ctx.RegisterModuleType("bar", logic.BarModule)
//
//       // Register custom singletons
//       ctx.RegisterSingleton("baz", logic.NewBazSingleton())
//
//       // Create and initialize the custom Config object.
//       config := logic.NewConfig(srcDir)
//
//       // This call never returns
//       bootstrap.Main(ctx, config)
//   }
//
// Required Source Files
//
// There are three files that must be included in the source tree to facilitate
// the build bootstrapping:
//
//   1. The top-level Blueprints file
//   2. The bootstrap Ninja file template
//   3. The bootstrap script
//
// The top-level Blueprints file describes how the entire source tree should be
// built.  It must have a 'subdirs' assignment that includes both the core
// Blueprint library and the custom build logic for the source tree.  It should
// also include (either directly or through a subdirs entry) describe all the
// modules to be built in the source tree.
//
// The bootstrap Ninja file template describes the build actions necessary to
// build the primary builder for the source tree.  This template contains a set
// of placeholder Ninja variable values that get filled in by the bootstrap
// script to create a usable Ninja file.  It can be created by running the
// minibp binary that gets created as part of the standalone Blueprint build.
// Passing minibp the path to the top-level Blueprints file will cause it to
// create a bootstrap Ninja file template named 'build.ninja.in'.
//
// The bootstrap script is a small script (or theoretically a compiled binary)
// that is included in the source tree to begin the bootstrapping process.  It
// is responsible for filling in the bootstrap Ninja file template with some
// basic information about the Go build environemnt and the path to the root
// source directory.  It does this by performing a simple string substitution on
// the template file to produce a usable build.ninja file.
//
// The Bootstrapping Process
//
// There are three stages to the bootstrapping process, each with a
// corresponding Ninja file. The stages are referred to as the "bootstrap",
// "primary", and "main" stages. Each stage builds the next stage's Ninja file.
//
// The bootstrapping process begins with the user running the bootstrap script
// to initialize a new build directory.  The script is run from the build
// directory, and when run with no arguments it copies the source bootstrap
// Ninja file into the build directory as ".minibootstrap/build.ninja".  It
// also performs a set of string substitutions on the file to configure it for
// the user's build environment. Specifically, the following strings are
// substituted in the file:
//
//   @@SrcDir@@            - The path to the root source directory (either
//                           absolute or relative to the build dir)
//   @@BuildDir@@          - The path to the build directory
//   @@GoRoot@@            - The path to the root directory of the Go toolchain
//   @@GoCompile@@         - The path to the Go compiler (6g or compile)
//   @@GoLink@@            - The path to the Go linker (6l or link)
//   @@Bootstrap@@         - The path to the bootstrap script
//   @@BootstrapManifest@@ - The path to the source bootstrap Ninja file
//
// Once the script completes the build directory is initialized and ready to run
// a build. A wrapper script (blueprint.bash by default) has been installed in
// order to run a build. It iterates through the three stages of the build:
//
//      - Checks to see if the source bootstrap Ninja file is newer than the
//        one that is in the build directory, if so, update the build dir copy.
//      - Run the Bootstrap stage
//      - Run the Primary stage
//      - Run the Main stage
//
// Previously, we were keeping track of the "state" of the build directory and
// only going back to previous stages when something had changed. But that
// added complexity, and failed when there was a build error in the Primary
// stage that would have been fixed if the Bootstrap stage was re-run (we
// would only evaluate which stages needed to be run at the end of the stage).
// So now we always run through each stage, and the first two stages will do
// nothing when nothing has changed.
//
// During the Bootstrap stage, <builddir>/.minibootstrap/build.ninja, the
// following actions are taken, if necessary:
//
//      - Build all bootstrap_core_go_binary rules, and dependencies --
//        minibp and some test helpers.
//      - Run minibp to generate .bootstrap/build.ninja (Primary stage)
//      - Run minibp to generate .minibootstrap/build.ninja.in
//      - Restart if .minibootstrap/build.ninja.in has changed
//
// During the Primary stage, <builddir>/.bootstrap/build.ninja, the following
// actions are taken, if necessary:
//
//      - Build any bootstrap_go_binary rules and dependencies -- usually the
//        primary builder and any build or runtime dependencies.
//      - Run the primary builder to generate build.ninja
//      - Run the primary builder to extract documentation
//
// Then the main stage is at <builddir>/build.ninja, and will contain all the
// rules generated by the primary builder. In addition, the bootstrap code
// adds a phony rule "blueprint_tools" that depends on all blueprint_go_binary
// rules (bpfmt, bpmodify, etc).
//
// Updating the Bootstrap Ninja File Template
//
// The main purpose of the bootstrap stage is to generate the Ninja file for the
// primary stage.  The one additional thing it does is generate a new bootstrap
// Ninja file template at .minibootstrap/build.ninja.in.  When generating this
// file, minibp will compare the new bootstrap Ninja file contents with the
// original (in the source tree).
//
// This scheme ensures that updates to the source tree are always incorporated
// into the build process.
//
package bootstrap
<commit_msg>Update documentation<commit_after>// Copyright 2014 Google Inc. All rights reserved.
//
// 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.

// The Blueprint bootstrapping mechanism is intended to enable building a
// source tree with minimal prebuilts.  The only prerequisites for performing
// such a build are:
//
//   1. A Ninja binary
//   2. A script interpreter (e.g. Bash or Python)
//   3. A Go toolchain
//
// The Primary Builder
//
// As part of the bootstrapping process, a binary called the "primary builder"
// is created.  This primary builder is the binary that includes both the core
// Blueprint library and the build logic specific to the source tree.  It is
// used to generate the Ninja file that describes how to build the entire source
// tree.
//
// The primary builder must be a pure Go (i.e. no cgo) module built with the
// module type 'bootstrap_go_binary'.  It should have the 'primaryBuilder'
// module property set to true in its Blueprints file.  If more than one module
// sets primaryBuilder to true the build will fail.
//
// The primary builder main function should look something like:
//
//   package main
//
//   import (
//       "flag"
//       "github.com/google/blueprint"
//       "github.com/google/blueprint/bootstrap"
//       "path/filepath"
//
//       "my/custom/build/logic"
//   )
//
//   func main() {
//       // The primary builder should use the global flag set because the
//       // bootstrap package registers its own flags there.
//       flag.Parse()
//
//       // The top-level Blueprints file is passed as the first argument.
//       srcDir := filepath.Dir(flag.Arg(0))
//
//       // Create the build context.
//       ctx := blueprint.NewContext()
//
//       // Register custom module types
//       ctx.RegisterModuleType("foo", logic.FooModule)
//       ctx.RegisterModuleType("bar", logic.BarModule)
//
//       // Register custom singletons
//       ctx.RegisterSingleton("baz", logic.NewBazSingleton())
//
//       // Create and initialize the custom Config object.
//       config := logic.NewConfig(srcDir)
//
//       // This call never returns
//       bootstrap.Main(ctx, config)
//   }
//
// Required Source Files
//
// There are three files that must be included in the source tree to facilitate
// the build bootstrapping:
//
//   1. The top-level Blueprints file
//   2. The bootstrap script
//   3. The build wrapper script
//
// The top-level Blueprints file describes how the entire source tree should be
// built.  It must have a 'subdirs' assignment that includes both the core
// Blueprint library and the custom build logic for the source tree.  It should
// also include (either directly or through a subdirs entry) describe all the
// modules to be built in the source tree.
//
// The bootstrap script is a small script to setup the build directory, writing
// a couple configuration files (including the path the source directory,
// information about the Go build environment, etc), then copying the build
// wrapper into the build directory.
//
// The Bootstrapping Process
//
// There are three stages to the bootstrapping process, each with a
// corresponding Ninja file. The stages are referred to as the "bootstrap",
// "primary", and "main" stages. Each stage builds the next stage's Ninja file.
//
// The bootstrapping process begins with the user running the bootstrap script
// to initialize a new build directory.  The script is run from the build
// directory, and creates a ".minibootstrap/build.ninja" file that sets a few
// variables then includes blueprint's "bootstrap/build.ninja". It also writes
// out a ".blueprint.bootstrap" file that contains a few variables for later use:
//
//   BLUEPRINT_BOOTSTRAP_VERSION - Used to detect when a user needs to run
//                                 bootstrap.bash again
//
//   SRCDIR         - The path to the source directory
//   BLUEPRINTDIR   - The path to the blueprints directory (includes $SRCDIR)
//   GOROOT         - The path to the root directory of the Go toolchain
//   NINJA_BUILDDIR - The path to store .ninja_log, .ninja_deps
//
// Once the script completes the build directory is initialized and ready to run
// a build. A wrapper script (blueprint.bash by default) has been installed in
// order to run a build. It iterates through the three stages of the build:
//
//      - Runs microfactory.bash to build minibp
//      - Runs the .minibootstrap/build.ninja to build .bootstrap/build.ninja
//      - Runs .bootstrap/build.ninja to build and run the primary builder
//      - Runs build.ninja to build your code
//
// Microfactory takes care of building an up to date version of `minibp` under
// the .minibootstrap/ directory.
//
// During <builddir>/.minibootstrap/build.ninja, the following actions are
// taken, if necessary:
//
//      - Run minibp to generate .bootstrap/build.ninja (Primary stage)
//
// During the <builddir>/.bootstrap/build.ninja, the following actions are
// taken, if necessary:
//
//      - Rebuild .bootstrap/build.ninja, usually due to globs changing --
//        other dependencies will trigger it to be built during minibootstrap
//      - Build the primary builder, anything marked `default: true`, and
//        any dependencies.
//      - Run the primary builder to generate build.ninja
//      - Run the primary builder to extract documentation
//
// Then the main stage is at <builddir>/build.ninja, and will contain all the
// rules generated by the primary builder. In addition, the bootstrap code
// adds a phony rule "blueprint_tools" that depends on all blueprint_go_binary
// rules (bpfmt, bpmodify, etc).
//
package bootstrap
<|endoftext|>


<commit_before>package srvproxy

import (
	"log"
	"time"
)

// BreakerProxy acts as a SRV proxy with a built-in circuit breaker.
type BreakerProxy struct {
	c chan Endpoint
	f chan Endpoint
	q chan chan struct{}
}

// NewBreakerProxy constructs and returns a BreakerProxy, performing regular
// resolve lookups against a defined name string, and caching the list of
// returned endpoints.
func NewBreakerProxy(resolve Resolver, name string, pollInterval time.Duration) (*BreakerProxy, error) {
	endpoints, err := resolve(name)
	if err != nil {
		return nil, err
	}
	p := &BreakerProxy{
		c: make(chan Endpoint),
		f: make(chan Endpoint),
		q: make(chan chan struct{}),
	}
	go p.loop(resolve, name, updateControllers(endpoints, map[Endpoint]*controller{}, p.c), pollInterval)
	return p, nil
}

func (p *BreakerProxy) loop(resolve Resolver, name string, controllers map[Endpoint]*controller, pollInterval time.Duration) {
	tick := time.Tick(pollInterval)
	for {
		select {
		case q := <-p.q:
			for _, controller := range controllers {
				controller.stop()
			}
			close(q)
			return

		case endpoint := <-p.f:
			if controller, ok := controllers[endpoint]; ok {
				controller.failure()
			}

		case <-tick:
			endpoints, err := resolve(name)
			if err != nil {
				log.Printf("srvproxy: poll: %s", err)
				continue // don't replace good endpoints with bad
			}
			controllers = updateControllers(endpoints, controllers, p.c)
		}
	}
}

// Close stops all goroutines associated with this BreakerProxy. Closed
// proxies must not be accessed.
func (p *BreakerProxy) Close() {
	q := make(chan struct{})
	p.q <- q
	<-q
}

// With chooses an endpoint and yields it to the passed closure. If the
// closure returns a non-nil error, With reports it as a failure to the
// associated endpoint, which will trip that endpoint's circuit breaker,
// rendering it unavailable for one second.
func (p BreakerProxy) With(f func(Endpoint) error) error {
	var endpoint Endpoint
	select {
	case endpoint = <-p.c:
	default:
		return errNoEndpoints
	}
	err := f(endpoint)
	if err != nil {
		p.f <- endpoint
	}
	return err
}

type controller struct {
	out  chan Endpoint
	fail chan struct{}
	quit chan chan struct{}
}

type states int

const (
	closed states = iota
	tripped
	open
)

func newController(out chan Endpoint, endpoint Endpoint) *controller {
	c := &controller{
		out:  make(chan Endpoint),
		fail: make(chan struct{}),
		quit: make(chan chan struct{}),
	}
	go c.loop(endpoint)
	return c
}

func (c *controller) loop(endpoint Endpoint) {
	var (
		state   states
		timeout <-chan time.Time
	)
	for {
		switch state {
		case closed:
			select {
			case q := <-c.quit:
				close(q)
				return
			case c.out <- endpoint:
				// continue
			case <-c.fail:
				log.Printf("%v: failed, tripping circuit", endpoint)
				state = tripped
			}

		case tripped:
			timeout = time.After(1 * time.Second)
			state = open

		case open:
			select {
			case q := <-c.quit:
				close(q)
				return
			case <-c.fail:
				// ignore
			case <-timeout:
				log.Printf("%v: cooldown expired, closing circuit", endpoint)
				state = closed
			}
		}
	}
}

func (c *controller) failure() {
	c.fail <- struct{}{}
}

func (c *controller) stop() {
	q := make(chan struct{})
	c.quit <- q
	<-q
}

func updateControllers(endpoints []Endpoint, existing map[Endpoint]*controller, c chan Endpoint) map[Endpoint]*controller {
	updated := map[Endpoint]*controller{}
	for _, endpoint := range endpoints {
		controller, ok := existing[endpoint]
		if ok {
			delete(existing, endpoint) // mv
		} else {
			controller = newController(c, endpoint)
		}
		updated[endpoint] = controller
	}
	for _, controller := range existing {
		controller.stop()
	}
	return updated
}
<commit_msg>Fix bug in breaker proxy<commit_after>package srvproxy

import (
	"log"
	"time"
)

// BreakerProxy acts as a SRV proxy with a built-in circuit breaker.
type BreakerProxy struct {
	c chan Endpoint
	f chan Endpoint
	q chan chan struct{}
}

// NewBreakerProxy constructs and returns a BreakerProxy, performing regular
// resolve lookups against a defined name string, and caching the list of
// returned endpoints.
func NewBreakerProxy(resolve Resolver, name string, pollInterval time.Duration) (*BreakerProxy, error) {
	endpoints, err := resolve(name)
	if err != nil {
		return nil, err
	}
	p := &BreakerProxy{
		c: make(chan Endpoint),
		f: make(chan Endpoint),
		q: make(chan chan struct{}),
	}
	controllers := updateControllers(endpoints, map[Endpoint]*controller{}, p.c)
	go p.loop(resolve, name, controllers, pollInterval)
	return p, nil
}

func (p *BreakerProxy) loop(resolve Resolver, name string, controllers map[Endpoint]*controller, pollInterval time.Duration) {
	tick := time.Tick(pollInterval)
	for {
		select {
		case q := <-p.q:
			for _, controller := range controllers {
				controller.stop()
			}
			close(q)
			return

		case endpoint := <-p.f:
			if controller, ok := controllers[endpoint]; ok {
				controller.failure()
			}

		case <-tick:
			endpoints, err := resolve(name)
			if err != nil {
				log.Printf("srvproxy: resolve: %s", err)
				continue // don't replace good endpoints with bad
			}
			controllers = updateControllers(endpoints, controllers, p.c)
		}
	}
}

// Close stops all goroutines associated with this BreakerProxy. Closed
// proxies must not be accessed.
func (p *BreakerProxy) Close() {
	q := make(chan struct{})
	p.q <- q
	<-q
}

// With chooses an endpoint and yields it to the passed closure. If the
// closure returns a non-nil error, With reports it as a failure to the
// associated endpoint, which will trip that endpoint's circuit breaker,
// rendering it unavailable for one second.
func (p BreakerProxy) With(f func(Endpoint) error) error {
	var endpoint Endpoint
	select {
	case endpoint = <-p.c:
	default:
		return errNoEndpoints
	}
	err := f(endpoint)
	if err != nil {
		p.f <- endpoint
	}
	return err
}

type controller struct {
	out  chan Endpoint
	fail chan struct{}
	quit chan chan struct{}
}

type states int

const (
	closed states = iota
	tripped
	open
)

func newController(out chan Endpoint, endpoint Endpoint) *controller {
	c := &controller{
		out:  out,
		fail: make(chan struct{}),
		quit: make(chan chan struct{}),
	}
	go c.loop(endpoint)
	return c
}

func (c *controller) loop(endpoint Endpoint) {
	var (
		state   states
		timeout <-chan time.Time
	)
	for {
		switch state {
		case closed:
			select {
			case q := <-c.quit:
				close(q)
				return
			case c.out <- endpoint:
				// continue
			case <-c.fail:
				log.Printf("srvproxy: %v: failed, tripping circuit", endpoint)
				state = tripped
			}

		case tripped:
			timeout = time.After(1 * time.Second)
			state = open

		case open:
			select {
			case q := <-c.quit:
				close(q)
				return
			case <-c.fail:
				// ignore
			case <-timeout:
				log.Printf("srvproxy: %v: cooldown expired, closing circuit", endpoint)
				state = closed
			}
		}
	}
}

func (c *controller) failure() {
	c.fail <- struct{}{}
}

func (c *controller) stop() {
	q := make(chan struct{})
	c.quit <- q
	<-q
}

func updateControllers(endpoints []Endpoint, existing map[Endpoint]*controller, c chan Endpoint) map[Endpoint]*controller {
	updated := map[Endpoint]*controller{}
	for _, endpoint := range endpoints {
		controller, ok := existing[endpoint]
		if ok {
			delete(existing, endpoint) // mv
		} else {
			log.Printf("srvproxy: %v: added", endpoint)
			controller = newController(c, endpoint)
		}
		updated[endpoint] = controller
	}
	for endpoint, controller := range existing {
		log.Printf("srvproxy: %v: removed", endpoint)
		controller.stop()
	}
	return updated
}
<|endoftext|>


<commit_before>/*
	Package broker reserves jobs from beanstalkd, spawns worker processes,
	and manages the interaction between the two.
*/
package broker

import (
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"syscall"
	"time"

	"github.com/kr/beanstalk"
)

const (
	// ttrMargin compensates for beanstalkd's integer precision.
	// e.g. reserving a TTR=1 job will show time-left=0.
	// We need to set our SIGTERM timer to time-left + ttrMargin.
	ttrMargin = 1 * time.Second

	// deadlineSoonDelay defines a period to sleep between receiving
	// DEADLINE_SOON in response to reserve, and re-attempting the reserve.
	deadlineSoonDelay = 1 * time.Second
)

type Broker struct {

	// Address of the beanstalkd server.
	Address string

	// The shell command to execute for each job.
	Cmd string

	// Tube name this broker will service.
	Tube string

	log     *log.Logger
	results chan<- *JobResult
}

type JobResult struct {

	// Buried is true if the job was buried.
	Buried bool

	// Executed is true if the job command was executed (or attempted).
	Executed bool

	// ExitStatus of the command; 0 for success.
	ExitStatus int

	// JobId from beanstalkd.
	JobId uint64

	// Stdout of the command.
	Stdout []byte

	// TimedOut indicates the worker exceeded TTR for the job.
	// Note this is tracked by a timer, separately to beanstalkd.
	TimedOut bool

	// Error raised while attempting to handle the job.
	Error error
}

// New broker instance.
func New(address, tube string, cmd string, results chan<- *JobResult) (b Broker) {
	b.Address = address
	b.Tube = tube
	b.Cmd = cmd

	b.log = log.New(os.Stdout, fmt.Sprintf("[%s] ", tube), log.LstdFlags)
	b.results = results
	return
}

// reserve-with-timeout until there's a job or something panic-worthy.
func (b *Broker) mustReserveWithoutTimeout(ts *beanstalk.TubeSet) (id uint64, body []byte) {
	var err error
	for {
		id, body, err = ts.Reserve(1 * time.Hour)
		if err == nil {
			return
		} else if err.(beanstalk.ConnError).Err == beanstalk.ErrTimeout {
			continue
		} else if err.(beanstalk.ConnError).Err == beanstalk.ErrDeadline {
			b.log.Printf("%v (retrying in %v)", err, deadlineSoonDelay)
			time.Sleep(deadlineSoonDelay)
			continue
		} else {
			panic(err)
		}
	}
}

// Run connects to beanstalkd and starts broking.
// If ticks channel is present, one job is processed per tick.
func (b *Broker) Run(ticks chan bool) {
	b.log.Println("command:", b.Cmd)
	b.log.Println("connecting to", b.Address)
	conn, err := beanstalk.Dial("tcp", b.Address)
	if err != nil {
		panic(err)
	}

	b.log.Println("watching", b.Tube)
	ts := beanstalk.NewTubeSet(conn, b.Tube)

	for {
		if ticks != nil {
			if _, ok := <-ticks; !ok {
				break
			}
		}

		b.log.Println("reserve (waiting for job)")
		id, body := b.mustReserveWithoutTimeout(ts)
		job := &job{id: id, body: body, conn: conn}

		t, err := job.timeouts()
		if err != nil {
			log.Panic(err)
		}
		if t > 0 {
			b.log.Printf("job %d has %d timeouts, burying", job.id, t)
			job.bury()
			if b.results != nil {
				b.results <- &JobResult{JobId: job.id, Buried: true}
			}
			continue
		}

		b.log.Printf("executing job %d", job.id)
		result, err := b.executeJob(job, b.Cmd)
		if err != nil {
			log.Panic(err)
		}

		err = b.handleResult(job, result)
		if err != nil {
			log.Panic(err)
		}

		if result.Error != nil {
			b.log.Println("result had error:", result.Error)
		}

		if b.results != nil {
			b.results <- result
		}
	}

	b.log.Println("broker finished")
}

func (b *Broker) executeJob(job *job, shellCmd string) (result *JobResult, err error) {
	result = &JobResult{JobId: job.id, Executed: true}

	ttr, err := job.timeLeft()
	timer := time.NewTimer(ttr + ttrMargin)
	if err != nil {
		return
	}

	cmd, stdout, err := startCommand(shellCmd, job.body)
	if err != nil {
		return
	}

	stdoutC := readerToChannel(stdout)

	// TODO: end loop when stdout closes
stdoutReader:
	for {
		select {
		case <-timer.C:
			b.killWorker(cmd.Process)
			result.TimedOut = true
		case data, ok := <-stdoutC:
			if !ok {
				break stdoutReader
			}
			b.log.Printf("stdout: %s", data)
			result.Stdout = append(result.Stdout, data...)
		}
	}

	waitC := make(chan error)
	go func() {
		waitC <- cmd.Wait()
	}()

waitLoop:
	for {
		select {
		case err = <-waitC:
			timer.Stop()
			if e1, ok := err.(*exec.ExitError); ok {
				result.ExitStatus = e1.Sys().(syscall.WaitStatus).ExitStatus()
				b.log.Printf("set exit status to %d", result.ExitStatus)
				err = nil // not a executeJob error
			}
			break waitLoop
		case <-timer.C:
			b.killWorker(cmd.Process)
			result.TimedOut = true
		}
	}

	return
}

func (b *Broker) killWorker(p *os.Process) {
	b.log.Printf("Sending SIGTERM to worker PID %d", p.Pid)
	p.Signal(syscall.SIGTERM)
	// TODO: follow up with SIGKILL if still running.
}

func (b *Broker) handleResult(job *job, result *JobResult) (err error) {
	if result.TimedOut {
		b.log.Printf("job %d timed out", job.id)
		return
	}
	b.log.Printf("job %d finished with exit(%d)", job.id, result.ExitStatus)
	switch result.ExitStatus {
	case 0:
		b.log.Printf("deleting job %d", job.id)
		err = job.delete()
	case 1:
		b.log.Printf("releasing job %d", job.id)
		err = job.release()
	default:
		err = fmt.Errorf("Unhandled exit status %d", result.ExitStatus)
	}
	return
}

func startCommand(shellCmd string, input []byte) (cmd *exec.Cmd, stdout io.ReadCloser, err error) {
	cmd = exec.Command("/bin/bash", "-c", shellCmd)

	stdin, err := cmd.StdinPipe()
	if err != nil {
		return
	}

	stdout, err = cmd.StdoutPipe()
	if err != nil {
		return
	}

	cmd.Stderr = os.Stderr

	err = cmd.Start()
	if err != nil {
		return
	}

	_, err = stdin.Write(input)
	if err != nil {
		return
	}
	stdin.Close()

	return
}

func readerToChannel(reader io.Reader) <-chan []byte {
	c := make(chan []byte)
	go func() {
		buf := make([]byte, 1024)
		for {
			n, err := reader.Read(buf)
			if n > 0 {
				res := make([]byte, n)
				copy(res, buf[:n])
				c <- res
			}
			if err != nil {
				close(c)
				break
			}
		}
	}()
	return c
}
<commit_msg>waitChan extracted from executeJob.<commit_after>/*
	Package broker reserves jobs from beanstalkd, spawns worker processes,
	and manages the interaction between the two.
*/
package broker

import (
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"syscall"
	"time"

	"github.com/kr/beanstalk"
)

const (
	// ttrMargin compensates for beanstalkd's integer precision.
	// e.g. reserving a TTR=1 job will show time-left=0.
	// We need to set our SIGTERM timer to time-left + ttrMargin.
	ttrMargin = 1 * time.Second

	// deadlineSoonDelay defines a period to sleep between receiving
	// DEADLINE_SOON in response to reserve, and re-attempting the reserve.
	deadlineSoonDelay = 1 * time.Second
)

type Broker struct {

	// Address of the beanstalkd server.
	Address string

	// The shell command to execute for each job.
	Cmd string

	// Tube name this broker will service.
	Tube string

	log     *log.Logger
	results chan<- *JobResult
}

type JobResult struct {

	// Buried is true if the job was buried.
	Buried bool

	// Executed is true if the job command was executed (or attempted).
	Executed bool

	// ExitStatus of the command; 0 for success.
	ExitStatus int

	// JobId from beanstalkd.
	JobId uint64

	// Stdout of the command.
	Stdout []byte

	// TimedOut indicates the worker exceeded TTR for the job.
	// Note this is tracked by a timer, separately to beanstalkd.
	TimedOut bool

	// Error raised while attempting to handle the job.
	Error error
}

// New broker instance.
func New(address, tube string, cmd string, results chan<- *JobResult) (b Broker) {
	b.Address = address
	b.Tube = tube
	b.Cmd = cmd

	b.log = log.New(os.Stdout, fmt.Sprintf("[%s] ", tube), log.LstdFlags)
	b.results = results
	return
}

// reserve-with-timeout until there's a job or something panic-worthy.
func (b *Broker) mustReserveWithoutTimeout(ts *beanstalk.TubeSet) (id uint64, body []byte) {
	var err error
	for {
		id, body, err = ts.Reserve(1 * time.Hour)
		if err == nil {
			return
		} else if err.(beanstalk.ConnError).Err == beanstalk.ErrTimeout {
			continue
		} else if err.(beanstalk.ConnError).Err == beanstalk.ErrDeadline {
			b.log.Printf("%v (retrying in %v)", err, deadlineSoonDelay)
			time.Sleep(deadlineSoonDelay)
			continue
		} else {
			panic(err)
		}
	}
}

// Run connects to beanstalkd and starts broking.
// If ticks channel is present, one job is processed per tick.
func (b *Broker) Run(ticks chan bool) {
	b.log.Println("command:", b.Cmd)
	b.log.Println("connecting to", b.Address)
	conn, err := beanstalk.Dial("tcp", b.Address)
	if err != nil {
		panic(err)
	}

	b.log.Println("watching", b.Tube)
	ts := beanstalk.NewTubeSet(conn, b.Tube)

	for {
		if ticks != nil {
			if _, ok := <-ticks; !ok {
				break
			}
		}

		b.log.Println("reserve (waiting for job)")
		id, body := b.mustReserveWithoutTimeout(ts)
		job := &job{id: id, body: body, conn: conn}

		t, err := job.timeouts()
		if err != nil {
			log.Panic(err)
		}
		if t > 0 {
			b.log.Printf("job %d has %d timeouts, burying", job.id, t)
			job.bury()
			if b.results != nil {
				b.results <- &JobResult{JobId: job.id, Buried: true}
			}
			continue
		}

		b.log.Printf("executing job %d", job.id)
		result, err := b.executeJob(job, b.Cmd)
		if err != nil {
			log.Panic(err)
		}

		err = b.handleResult(job, result)
		if err != nil {
			log.Panic(err)
		}

		if result.Error != nil {
			b.log.Println("result had error:", result.Error)
		}

		if b.results != nil {
			b.results <- result
		}
	}

	b.log.Println("broker finished")
}

func (b *Broker) executeJob(job *job, shellCmd string) (result *JobResult, err error) {
	result = &JobResult{JobId: job.id, Executed: true}

	ttr, err := job.timeLeft()
	timer := time.NewTimer(ttr + ttrMargin)
	if err != nil {
		return
	}

	cmd, stdout, err := startCommand(shellCmd, job.body)
	if err != nil {
		return
	}

	stdoutC := readerToChannel(stdout)

	// TODO: end loop when stdout closes
stdoutReader:
	for {
		select {
		case <-timer.C:
			b.killWorker(cmd.Process)
			result.TimedOut = true
		case data, ok := <-stdoutC:
			if !ok {
				break stdoutReader
			}
			b.log.Printf("stdout: %s", data)
			result.Stdout = append(result.Stdout, data...)
		}
	}

	waitC := waitChan(cmd)

waitLoop:
	for {
		select {
		case wr := <-waitC:
			timer.Stop()
			if wr.err == nil {
				err = wr.err
			}
			result.ExitStatus = wr.status
			break waitLoop
		case <-timer.C:
			b.killWorker(cmd.Process)
			result.TimedOut = true
		}
	}

	return
}

type waitResult struct {
	status int
	err    error
}

// Given a command, waits and sends the exit status over the returned channel.
func waitChan(cmd *exec.Cmd) <-chan waitResult {
	c := make(chan waitResult)
	go func() {
		err := cmd.Wait()
		if err == nil {
			c <- waitResult{0, nil}
		} else if e1, ok := err.(*exec.ExitError); ok {
			status := e1.Sys().(syscall.WaitStatus).ExitStatus()
			c <- waitResult{status, nil}
		} else {
			c <- waitResult{-1, err}
		}
	}()
	return c
}

func (b *Broker) killWorker(p *os.Process) {
	b.log.Printf("Sending SIGTERM to worker PID %d", p.Pid)
	p.Signal(syscall.SIGTERM)
	// TODO: follow up with SIGKILL if still running.
}

func (b *Broker) handleResult(job *job, result *JobResult) (err error) {
	if result.TimedOut {
		b.log.Printf("job %d timed out", job.id)
		return
	}
	b.log.Printf("job %d finished with exit(%d)", job.id, result.ExitStatus)
	switch result.ExitStatus {
	case 0:
		b.log.Printf("deleting job %d", job.id)
		err = job.delete()
	case 1:
		b.log.Printf("releasing job %d", job.id)
		err = job.release()
	default:
		err = fmt.Errorf("Unhandled exit status %d", result.ExitStatus)
	}
	return
}

func startCommand(shellCmd string, input []byte) (cmd *exec.Cmd, stdout io.ReadCloser, err error) {
	cmd = exec.Command("/bin/bash", "-c", shellCmd)

	stdin, err := cmd.StdinPipe()
	if err != nil {
		return
	}

	stdout, err = cmd.StdoutPipe()
	if err != nil {
		return
	}

	cmd.Stderr = os.Stderr

	err = cmd.Start()
	if err != nil {
		return
	}

	_, err = stdin.Write(input)
	if err != nil {
		return
	}
	stdin.Close()

	return
}

func readerToChannel(reader io.Reader) <-chan []byte {
	c := make(chan []byte)
	go func() {
		buf := make([]byte, 1024)
		for {
			n, err := reader.Read(buf)
			if n > 0 {
				res := make([]byte, n)
				copy(res, buf[:n])
				c <- res
			}
			if err != nil {
				close(c)
				break
			}
		}
	}()
	return c
}
<|endoftext|>


<commit_before><commit_msg>Excluded a few IAM resources from Terraform Validator generation (#4925) (#734)<commit_after><|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"log"
	"strconv"
	"strings"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/elb"
	"github.com/hashicorp/terraform-plugin-sdk/helper/schema"
)

func resourceAwsProxyProtocolPolicy() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsProxyProtocolPolicyCreate,
		Read:   resourceAwsProxyProtocolPolicyRead,
		Update: resourceAwsProxyProtocolPolicyUpdate,
		Delete: resourceAwsProxyProtocolPolicyDelete,

		Schema: map[string]*schema.Schema{
			"load_balancer": {
				Type:     schema.TypeString,
				Required: true,
			},

			"instance_ports": {
				Type:     schema.TypeSet,
				Elem:     &schema.Schema{Type: schema.TypeString},
				Required: true,
				Set:      schema.HashString,
			},
		},
	}
}

func resourceAwsProxyProtocolPolicyCreate(d *schema.ResourceData, meta interface{}) error {
	elbconn := meta.(*AWSClient).elbconn
	elbname := aws.String(d.Get("load_balancer").(string))

	input := &elb.CreateLoadBalancerPolicyInput{
		LoadBalancerName: elbname,
		PolicyAttributes: []*elb.PolicyAttribute{
			{
				AttributeName:  aws.String("ProxyProtocol"),
				AttributeValue: aws.String("True"),
			},
		},
		PolicyName:     aws.String("TFEnableProxyProtocol"),
		PolicyTypeName: aws.String("ProxyProtocolPolicyType"),
	}

	// Create a policy
	log.Printf("[DEBUG] ELB create a policy %s from policy type %s",
		*input.PolicyName, *input.PolicyTypeName)

	if _, err := elbconn.CreateLoadBalancerPolicy(input); err != nil {
		return fmt.Errorf("Error creating a policy %s: %s",
			*input.PolicyName, err)
	}

	// Assign the policy name for use later
	d.Partial(true)
	d.SetId(fmt.Sprintf("%s:%s", *elbname, *input.PolicyName))
	d.SetPartial("load_balancer")
	log.Printf("[INFO] ELB PolicyName: %s", *input.PolicyName)

	return resourceAwsProxyProtocolPolicyUpdate(d, meta)
}

func resourceAwsProxyProtocolPolicyRead(d *schema.ResourceData, meta interface{}) error {
	elbconn := meta.(*AWSClient).elbconn
	elbname := aws.String(d.Get("load_balancer").(string))

	// Retrieve the current ELB policies for updating the state
	req := &elb.DescribeLoadBalancersInput{
		LoadBalancerNames: []*string{elbname},
	}
	resp, err := elbconn.DescribeLoadBalancers(req)
	if err != nil {
		if isLoadBalancerNotFound(err) {
			// The ELB is gone now, so just remove it from the state
			d.SetId("")
			return nil
		}
		return fmt.Errorf("Error retrieving ELB attributes: %s", err)
	}

	backends := flattenBackendPolicies(resp.LoadBalancerDescriptions[0].BackendServerDescriptions)

	ports := []*string{}
	for ip := range backends {
		ipstr := strconv.Itoa(int(ip))
		ports = append(ports, &ipstr)
	}
	d.Set("instance_ports", ports)
	d.Set("load_balancer", *elbname)
	return nil
}

func resourceAwsProxyProtocolPolicyUpdate(d *schema.ResourceData, meta interface{}) error {
	elbconn := meta.(*AWSClient).elbconn
	elbname := aws.String(d.Get("load_balancer").(string))

	// Retrieve the current ELB policies for updating the state
	req := &elb.DescribeLoadBalancersInput{
		LoadBalancerNames: []*string{elbname},
	}
	resp, err := elbconn.DescribeLoadBalancers(req)
	if err != nil {
		if isLoadBalancerNotFound(err) {
			// The ELB is gone now, so just remove it from the state
			d.SetId("")
			return nil
		}
		return fmt.Errorf("Error retrieving ELB attributes: %s", err)
	}

	backends := flattenBackendPolicies(resp.LoadBalancerDescriptions[0].BackendServerDescriptions)
	policyName := resourceAwsProxyProtocolPolicyParseId(d.Id())

	d.Partial(true)
	if d.HasChange("instance_ports") {
		o, n := d.GetChange("instance_ports")
		os := o.(*schema.Set)
		ns := n.(*schema.Set)
		remove := os.Difference(ns).List()
		add := ns.Difference(os).List()

		inputs := []*elb.SetLoadBalancerPoliciesForBackendServerInput{}

		i, err := resourceAwsProxyProtocolPolicyRemove(policyName, remove, backends)
		if err != nil {
			return err
		}
		inputs = append(inputs, i...)

		i, err = resourceAwsProxyProtocolPolicyAdd(policyName, add, backends)
		if err != nil {
			return err
		}
		inputs = append(inputs, i...)

		for _, input := range inputs {
			input.LoadBalancerName = elbname
			if _, err := elbconn.SetLoadBalancerPoliciesForBackendServer(input); err != nil {
				return fmt.Errorf("Error setting policy for backend: %s", err)
			}
		}

		d.SetPartial("instance_ports")
	}

	return resourceAwsProxyProtocolPolicyRead(d, meta)
}

func resourceAwsProxyProtocolPolicyDelete(d *schema.ResourceData, meta interface{}) error {
	elbconn := meta.(*AWSClient).elbconn
	elbname := aws.String(d.Get("load_balancer").(string))

	// Retrieve the current ELB policies for updating the state
	req := &elb.DescribeLoadBalancersInput{
		LoadBalancerNames: []*string{elbname},
	}
	var err error
	resp, err := elbconn.DescribeLoadBalancers(req)
	if err != nil {
		if isLoadBalancerNotFound(err) {
			return nil
		}
		return fmt.Errorf("Error retrieving ELB attributes: %s", err)
	}

	backends := flattenBackendPolicies(resp.LoadBalancerDescriptions[0].BackendServerDescriptions)
	ports := d.Get("instance_ports").(*schema.Set).List()
	policyName := resourceAwsProxyProtocolPolicyParseId(d.Id())

	inputs, err := resourceAwsProxyProtocolPolicyRemove(policyName, ports, backends)
	if err != nil {
		return fmt.Errorf("Error detaching a policy from backend: %s", err)
	}
	for _, input := range inputs {
		input.LoadBalancerName = elbname
		if _, err := elbconn.SetLoadBalancerPoliciesForBackendServer(input); err != nil {
			return fmt.Errorf("Error setting policy for backend: %s", err)
		}
	}

	pOpt := &elb.DeleteLoadBalancerPolicyInput{
		LoadBalancerName: elbname,
		PolicyName:       aws.String(policyName),
	}
	if _, err := elbconn.DeleteLoadBalancerPolicy(pOpt); err != nil {
		return fmt.Errorf("Error removing a policy from load balancer: %s", err)
	}

	return nil
}

func resourceAwsProxyProtocolPolicyRemove(policyName string, ports []interface{}, backends map[int64][]string) ([]*elb.SetLoadBalancerPoliciesForBackendServerInput, error) {
	inputs := make([]*elb.SetLoadBalancerPoliciesForBackendServerInput, 0, len(ports))
	for _, p := range ports {
		ip, err := strconv.ParseInt(p.(string), 10, 64)
		if err != nil {
			return nil, fmt.Errorf("Error detaching the policy: %s", err)
		}

		newPolicies := []*string{}
		curPolicies, found := backends[ip]
		if !found {
			// No policy for this instance port found, just skip it.
			continue
		}

		for _, policy := range curPolicies {
			if policy == policyName {
				// remove the policy
				continue
			}
			newPolicies = append(newPolicies, &policy)
		}

		inputs = append(inputs, &elb.SetLoadBalancerPoliciesForBackendServerInput{
			InstancePort: &ip,
			PolicyNames:  newPolicies,
		})
	}
	return inputs, nil
}

func resourceAwsProxyProtocolPolicyAdd(policyName string, ports []interface{}, backends map[int64][]string) ([]*elb.SetLoadBalancerPoliciesForBackendServerInput, error) {
	inputs := make([]*elb.SetLoadBalancerPoliciesForBackendServerInput, 0, len(ports))
	for _, p := range ports {
		ip, err := strconv.ParseInt(p.(string), 10, 64)
		if err != nil {
			return nil, fmt.Errorf("Error attaching the policy: %s", err)
		}

		newPolicies := []*string{}
		curPolicies := backends[ip]
		for _, p := range curPolicies {
			if p == policyName {
				// Just remove it for now. It will be back later.
				continue
			} else {
				newPolicies = append(newPolicies, &p)
			}
		}
		newPolicies = append(newPolicies, aws.String(policyName))

		inputs = append(inputs, &elb.SetLoadBalancerPoliciesForBackendServerInput{
			InstancePort: &ip,
			PolicyNames:  newPolicies,
		})
	}
	return inputs, nil
}

// resourceAwsProxyProtocolPolicyParseId takes an ID and parses it into
// it's constituent parts. You need two axes (LB name, policy name)
// to create or identify a proxy protocol policy in AWS's API.
func resourceAwsProxyProtocolPolicyParseId(id string) string {
	parts := strings.SplitN(id, ":", 2)
	// We currently omit the ELB name as it is not currently used anywhere
	return parts[1]
}
<commit_msg>resource/aws_proxy_protocol_policy: Fixes for tfproviderlint R002 (#12030)<commit_after>package aws

import (
	"fmt"
	"log"
	"strconv"
	"strings"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/elb"
	"github.com/hashicorp/terraform-plugin-sdk/helper/schema"
)

func resourceAwsProxyProtocolPolicy() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsProxyProtocolPolicyCreate,
		Read:   resourceAwsProxyProtocolPolicyRead,
		Update: resourceAwsProxyProtocolPolicyUpdate,
		Delete: resourceAwsProxyProtocolPolicyDelete,

		Schema: map[string]*schema.Schema{
			"load_balancer": {
				Type:     schema.TypeString,
				Required: true,
			},

			"instance_ports": {
				Type:     schema.TypeSet,
				Elem:     &schema.Schema{Type: schema.TypeString},
				Required: true,
				Set:      schema.HashString,
			},
		},
	}
}

func resourceAwsProxyProtocolPolicyCreate(d *schema.ResourceData, meta interface{}) error {
	elbconn := meta.(*AWSClient).elbconn
	elbname := aws.String(d.Get("load_balancer").(string))

	input := &elb.CreateLoadBalancerPolicyInput{
		LoadBalancerName: elbname,
		PolicyAttributes: []*elb.PolicyAttribute{
			{
				AttributeName:  aws.String("ProxyProtocol"),
				AttributeValue: aws.String("True"),
			},
		},
		PolicyName:     aws.String("TFEnableProxyProtocol"),
		PolicyTypeName: aws.String("ProxyProtocolPolicyType"),
	}

	// Create a policy
	log.Printf("[DEBUG] ELB create a policy %s from policy type %s",
		*input.PolicyName, *input.PolicyTypeName)

	if _, err := elbconn.CreateLoadBalancerPolicy(input); err != nil {
		return fmt.Errorf("Error creating a policy %s: %s",
			*input.PolicyName, err)
	}

	// Assign the policy name for use later
	d.Partial(true)
	d.SetId(fmt.Sprintf("%s:%s", *elbname, *input.PolicyName))
	d.SetPartial("load_balancer")
	log.Printf("[INFO] ELB PolicyName: %s", *input.PolicyName)

	return resourceAwsProxyProtocolPolicyUpdate(d, meta)
}

func resourceAwsProxyProtocolPolicyRead(d *schema.ResourceData, meta interface{}) error {
	elbconn := meta.(*AWSClient).elbconn
	elbname := d.Get("load_balancer").(string)

	// Retrieve the current ELB policies for updating the state
	req := &elb.DescribeLoadBalancersInput{
		LoadBalancerNames: []*string{aws.String(elbname)},
	}
	resp, err := elbconn.DescribeLoadBalancers(req)
	if err != nil {
		if isLoadBalancerNotFound(err) {
			// The ELB is gone now, so just remove it from the state
			d.SetId("")
			return nil
		}
		return fmt.Errorf("Error retrieving ELB attributes: %s", err)
	}

	backends := flattenBackendPolicies(resp.LoadBalancerDescriptions[0].BackendServerDescriptions)

	ports := []*string{}
	for ip := range backends {
		ipstr := strconv.Itoa(int(ip))
		ports = append(ports, &ipstr)
	}
	d.Set("instance_ports", ports)
	d.Set("load_balancer", elbname)
	return nil
}

func resourceAwsProxyProtocolPolicyUpdate(d *schema.ResourceData, meta interface{}) error {
	elbconn := meta.(*AWSClient).elbconn
	elbname := aws.String(d.Get("load_balancer").(string))

	// Retrieve the current ELB policies for updating the state
	req := &elb.DescribeLoadBalancersInput{
		LoadBalancerNames: []*string{elbname},
	}
	resp, err := elbconn.DescribeLoadBalancers(req)
	if err != nil {
		if isLoadBalancerNotFound(err) {
			// The ELB is gone now, so just remove it from the state
			d.SetId("")
			return nil
		}
		return fmt.Errorf("Error retrieving ELB attributes: %s", err)
	}

	backends := flattenBackendPolicies(resp.LoadBalancerDescriptions[0].BackendServerDescriptions)
	policyName := resourceAwsProxyProtocolPolicyParseId(d.Id())

	d.Partial(true)
	if d.HasChange("instance_ports") {
		o, n := d.GetChange("instance_ports")
		os := o.(*schema.Set)
		ns := n.(*schema.Set)
		remove := os.Difference(ns).List()
		add := ns.Difference(os).List()

		inputs := []*elb.SetLoadBalancerPoliciesForBackendServerInput{}

		i, err := resourceAwsProxyProtocolPolicyRemove(policyName, remove, backends)
		if err != nil {
			return err
		}
		inputs = append(inputs, i...)

		i, err = resourceAwsProxyProtocolPolicyAdd(policyName, add, backends)
		if err != nil {
			return err
		}
		inputs = append(inputs, i...)

		for _, input := range inputs {
			input.LoadBalancerName = elbname
			if _, err := elbconn.SetLoadBalancerPoliciesForBackendServer(input); err != nil {
				return fmt.Errorf("Error setting policy for backend: %s", err)
			}
		}

		d.SetPartial("instance_ports")
	}

	return resourceAwsProxyProtocolPolicyRead(d, meta)
}

func resourceAwsProxyProtocolPolicyDelete(d *schema.ResourceData, meta interface{}) error {
	elbconn := meta.(*AWSClient).elbconn
	elbname := aws.String(d.Get("load_balancer").(string))

	// Retrieve the current ELB policies for updating the state
	req := &elb.DescribeLoadBalancersInput{
		LoadBalancerNames: []*string{elbname},
	}
	var err error
	resp, err := elbconn.DescribeLoadBalancers(req)
	if err != nil {
		if isLoadBalancerNotFound(err) {
			return nil
		}
		return fmt.Errorf("Error retrieving ELB attributes: %s", err)
	}

	backends := flattenBackendPolicies(resp.LoadBalancerDescriptions[0].BackendServerDescriptions)
	ports := d.Get("instance_ports").(*schema.Set).List()
	policyName := resourceAwsProxyProtocolPolicyParseId(d.Id())

	inputs, err := resourceAwsProxyProtocolPolicyRemove(policyName, ports, backends)
	if err != nil {
		return fmt.Errorf("Error detaching a policy from backend: %s", err)
	}
	for _, input := range inputs {
		input.LoadBalancerName = elbname
		if _, err := elbconn.SetLoadBalancerPoliciesForBackendServer(input); err != nil {
			return fmt.Errorf("Error setting policy for backend: %s", err)
		}
	}

	pOpt := &elb.DeleteLoadBalancerPolicyInput{
		LoadBalancerName: elbname,
		PolicyName:       aws.String(policyName),
	}
	if _, err := elbconn.DeleteLoadBalancerPolicy(pOpt); err != nil {
		return fmt.Errorf("Error removing a policy from load balancer: %s", err)
	}

	return nil
}

func resourceAwsProxyProtocolPolicyRemove(policyName string, ports []interface{}, backends map[int64][]string) ([]*elb.SetLoadBalancerPoliciesForBackendServerInput, error) {
	inputs := make([]*elb.SetLoadBalancerPoliciesForBackendServerInput, 0, len(ports))
	for _, p := range ports {
		ip, err := strconv.ParseInt(p.(string), 10, 64)
		if err != nil {
			return nil, fmt.Errorf("Error detaching the policy: %s", err)
		}

		newPolicies := []*string{}
		curPolicies, found := backends[ip]
		if !found {
			// No policy for this instance port found, just skip it.
			continue
		}

		for _, policy := range curPolicies {
			if policy == policyName {
				// remove the policy
				continue
			}
			newPolicies = append(newPolicies, &policy)
		}

		inputs = append(inputs, &elb.SetLoadBalancerPoliciesForBackendServerInput{
			InstancePort: &ip,
			PolicyNames:  newPolicies,
		})
	}
	return inputs, nil
}

func resourceAwsProxyProtocolPolicyAdd(policyName string, ports []interface{}, backends map[int64][]string) ([]*elb.SetLoadBalancerPoliciesForBackendServerInput, error) {
	inputs := make([]*elb.SetLoadBalancerPoliciesForBackendServerInput, 0, len(ports))
	for _, p := range ports {
		ip, err := strconv.ParseInt(p.(string), 10, 64)
		if err != nil {
			return nil, fmt.Errorf("Error attaching the policy: %s", err)
		}

		newPolicies := []*string{}
		curPolicies := backends[ip]
		for _, p := range curPolicies {
			if p == policyName {
				// Just remove it for now. It will be back later.
				continue
			} else {
				newPolicies = append(newPolicies, &p)
			}
		}
		newPolicies = append(newPolicies, aws.String(policyName))

		inputs = append(inputs, &elb.SetLoadBalancerPoliciesForBackendServerInput{
			InstancePort: &ip,
			PolicyNames:  newPolicies,
		})
	}
	return inputs, nil
}

// resourceAwsProxyProtocolPolicyParseId takes an ID and parses it into
// it's constituent parts. You need two axes (LB name, policy name)
// to create or identify a proxy protocol policy in AWS's API.
func resourceAwsProxyProtocolPolicyParseId(id string) string {
	parts := strings.SplitN(id, ":", 2)
	// We currently omit the ELB name as it is not currently used anywhere
	return parts[1]
}
<|endoftext|>