<commit_before>package mackerel

import (
	"bytes"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"time"
)

const (
	defaultBaseURL    = "https://api.mackerelio.com/"
	defaultUserAgent  = "mackerel-client-go"
	apiRequestTimeout = 30 * time.Second
)

// Client api client for mackerel
type Client struct {
	BaseURL           *url.URL
	APIKey            string
	Verbose           bool
	UserAgent         string
	AdditionalHeaders http.Header
	HTTPClient        *http.Client
}

// NewClient returns new mackerel.Client
func NewClient(apikey string) *Client {
	u, _ := url.Parse(defaultBaseURL)
	return &Client{u, apikey, false, defaultUserAgent, http.Header{}, &http.Client{}}
}

// NewClientWithOptions returns new mackerel.Client
func NewClientWithOptions(apikey string, rawurl string, verbose bool) (*Client, error) {
	u, err := url.Parse(rawurl)
	if err != nil {
		return nil, err
	}
	return &Client{u, apikey, verbose, defaultUserAgent, http.Header{}, &http.Client{}}, nil
}

func (c *Client) urlFor(path string) *url.URL {
	newURL, err := url.Parse(c.BaseURL.String())
	if err != nil {
		panic("invalid url passed")
	}

	newURL.Path = path

	return newURL
}

func (c *Client) buildReq(req *http.Request) *http.Request {
	for header, values := range c.AdditionalHeaders {
		for _, v := range values {
			req.Header.Add(header, v)
		}
	}
	req.Header.Set("X-Api-Key", c.APIKey)
	req.Header.Set("User-Agent", c.UserAgent)
	return req
}

// Request request to mackerel and receive response
func (c *Client) Request(req *http.Request) (resp *http.Response, err error) {
	req = c.buildReq(req)

	if c.Verbose {
		dump, err := httputil.DumpRequest(req, true)
		if err == nil {
			log.Printf("%s", dump)
		}
	}

	client := c.HTTPClient
	client.Timeout = apiRequestTimeout
	resp, err = client.Do(req)
	if err != nil {
		return nil, err
	}
	if c.Verbose {
		dump, err := httputil.DumpResponse(resp, true)
		if err == nil {
			log.Printf("%s", dump)
		}
	}
	if resp.StatusCode < 200 || resp.StatusCode > 299 {
		message := extractErrorMessage(resp.Body)
		defer resp.Body.Close()
		if message != "" {
			return nil, &APIError{StatusCode: resp.StatusCode, Message: message}
		}
		return nil, fmt.Errorf("API request failed: %s", resp.Status)
	}
	return resp, nil
}

// PostJSON shortcut method for posting json
func (c *Client) PostJSON(path string, payload interface{}) (*http.Response, error) {
	return c.requestJSON("POST", path, payload)
}

// PutJSON shortcut method for putting json
func (c *Client) PutJSON(path string, payload interface{}) (*http.Response, error) {
	return c.requestJSON("PUT", path, payload)
}

func (c *Client) requestJSON(method string, path string, payload interface{}) (*http.Response, error) {
	var body bytes.Buffer
	err := json.NewEncoder(&body).Encode(payload)
	if err != nil {
		return nil, err
	}

	req, err := http.NewRequest(method, c.urlFor(path).String(), &body)
	if err != nil {
		return nil, err
	}
	req.Header.Add("Content-Type", "application/json")
	return c.Request(req)
}

func closeResponse(resp *http.Response) {
	if resp != nil {
		resp.Body.Close()
	}
}
<commit_msg>return APIError after JSON unmarshal fails<commit_after>package mackerel

import (
	"bytes"
	"encoding/json"
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"time"
)

const (
	defaultBaseURL    = "https://api.mackerelio.com/"
	defaultUserAgent  = "mackerel-client-go"
	apiRequestTimeout = 30 * time.Second
)

// Client api client for mackerel
type Client struct {
	BaseURL           *url.URL
	APIKey            string
	Verbose           bool
	UserAgent         string
	AdditionalHeaders http.Header
	HTTPClient        *http.Client
}

// NewClient returns new mackerel.Client
func NewClient(apikey string) *Client {
	u, _ := url.Parse(defaultBaseURL)
	return &Client{u, apikey, false, defaultUserAgent, http.Header{}, &http.Client{}}
}

// NewClientWithOptions returns new mackerel.Client
func NewClientWithOptions(apikey string, rawurl string, verbose bool) (*Client, error) {
	u, err := url.Parse(rawurl)
	if err != nil {
		return nil, err
	}
	return &Client{u, apikey, verbose, defaultUserAgent, http.Header{}, &http.Client{}}, nil
}

func (c *Client) urlFor(path string) *url.URL {
	newURL, err := url.Parse(c.BaseURL.String())
	if err != nil {
		panic("invalid url passed")
	}

	newURL.Path = path

	return newURL
}

func (c *Client) buildReq(req *http.Request) *http.Request {
	for header, values := range c.AdditionalHeaders {
		for _, v := range values {
			req.Header.Add(header, v)
		}
	}
	req.Header.Set("X-Api-Key", c.APIKey)
	req.Header.Set("User-Agent", c.UserAgent)
	return req
}

// Request request to mackerel and receive response
func (c *Client) Request(req *http.Request) (resp *http.Response, err error) {
	req = c.buildReq(req)

	if c.Verbose {
		dump, err := httputil.DumpRequest(req, true)
		if err == nil {
			log.Printf("%s", dump)
		}
	}

	client := c.HTTPClient
	client.Timeout = apiRequestTimeout
	resp, err = client.Do(req)
	if err != nil {
		return nil, err
	}
	if c.Verbose {
		dump, err := httputil.DumpResponse(resp, true)
		if err == nil {
			log.Printf("%s", dump)
		}
	}
	if resp.StatusCode < 200 || resp.StatusCode > 299 {
		message := extractErrorMessage(resp.Body)
		defer resp.Body.Close()
		if message != "" {
			return nil, &APIError{StatusCode: resp.StatusCode, Message: message}
		}
		return nil, &APIError{StatusCode: resp.StatusCode, Message: resp.Status}
	}
	return resp, nil
}

// PostJSON shortcut method for posting json
func (c *Client) PostJSON(path string, payload interface{}) (*http.Response, error) {
	return c.requestJSON("POST", path, payload)
}

// PutJSON shortcut method for putting json
func (c *Client) PutJSON(path string, payload interface{}) (*http.Response, error) {
	return c.requestJSON("PUT", path, payload)
}

func (c *Client) requestJSON(method string, path string, payload interface{}) (*http.Response, error) {
	var body bytes.Buffer
	err := json.NewEncoder(&body).Encode(payload)
	if err != nil {
		return nil, err
	}

	req, err := http.NewRequest(method, c.urlFor(path).String(), &body)
	if err != nil {
		return nil, err
	}
	req.Header.Add("Content-Type", "application/json")
	return c.Request(req)
}

func closeResponse(resp *http.Response) {
	if resp != nil {
		resp.Body.Close()
	}
}
<|endoftext|>


<commit_before>package log

import (
	"fmt"
	"io"
	"log"
	"os"
	"sync"
)

var DefaultLogger *log.Logger
var logChan chan *logContent
var logLevel int
var mutex sync.Mutex

const (
	LevelDebug = iota
	LevelInfo
	LevelWarn
	LevelError
	LevelFatal
)

var levelMap = map[int]string{
	LevelDebug: "[DEBUG] ",
	LevelInfo:  "[INFO] ",
	LevelWarn:  "[WARN] ",
	LevelError: "[ERROR] ",
	LevelFatal: "[FATAL] ",
}

type logContent struct {
	level   int
	content string
}

func output(log *logContent) {
	mutex.Lock()
	defer mutex.Unlock()
	DefaultLogger.SetPrefix(levelMap[log.level])
	switch log.level {
	case LevelDebug, LevelInfo, LevelError:
		if log.level >= logLevel {
			DefaultLogger.Output(3, log.content)
		}
	case LevelFatal:
		DefaultLogger.Output(3, log.content)
		os.Exit(1)
	}
}

func Debug(v ...interface{}) {
	output(&logContent{
		level:   LevelDebug,
		content: fmt.Sprintln(v...),
	})
}

func Info(v ...interface{}) {
	output(&logContent{
		level:   LevelInfo,
		content: fmt.Sprintln(v...),
	})
}

func Warn(v ...interface{}) {
	output(&logContent{
		level:   LevelWarn,
		content: fmt.Sprintln(v...),
	})
}

func Error(v ...interface{}) {
	output(&logContent{
		level:   LevelError,
		content: fmt.Sprintln(v...),
	})
}

func Fatal(v ...interface{}) {
	output(&logContent{
		level:   LevelFatal,
		content: fmt.Sprintln(v...),
	})
}

func Debugf(format string, v ...interface{}) {
	output(&logContent{
		level:   LevelDebug,
		content: fmt.Sprintf(format, v...),
	})
}

func Infof(format string, v ...interface{}) {
	output(&logContent{
		level:   LevelInfo,
		content: fmt.Sprintf(format, v...),
	})
}

func Warnf(format string, v ...interface{}) {
	output(&logContent{
		level:   LevelWarn,
		content: fmt.Sprintf(format, v...),
	})
}

func Errorf(format string, v ...interface{}) {
	output(&logContent{
		level:   LevelError,
		content: fmt.Sprintf(format, v...),
	})
}

func Fatalf(format string, v ...interface{}) {
	output(&logContent{
		level:   LevelFatal,
		content: fmt.Sprintf(format, v...),
	})
}

func SetOptput(w io.Writer) {
	DefaultLogger.SetOutput(w)
}

func SetFlags(flag int) {
	DefaultLogger.SetFlags(flag)
}

func SetLevel(level int) {
	logLevel = level
}

func init() {
	logChan = make(chan *logContent)
	DefaultLogger = log.New(os.Stdout, "", log.LstdFlags|log.Lshortfile)
}
<commit_msg>update log<commit_after>package log

import (
	"fmt"
	"io"
	"log"
	"os"
	"sync"
)

var DefaultLogger *log.Logger
var logChan chan *logContent
var logLevel int
var mutex sync.Mutex

const (
	LevelDebug = iota
	LevelInfo
	LevelWarn
	LevelError
	LevelFatal
)

var levelMap = map[int]string{
	LevelDebug: "[DEBUG] ",
	LevelInfo:  "[INFO] ",
	LevelWarn:  "[WARN] ",
	LevelError: "[ERROR] ",
	LevelFatal: "[FATAL] ",
}

type logContent struct {
	level   int
	content string
}

func output(log *logContent) {
	mutex.Lock()
	defer mutex.Unlock()
	DefaultLogger.SetPrefix(levelMap[log.level])
	switch log.level {
	case LevelDebug, LevelInfo, LevelError:
		if log.level >= logLevel {
			DefaultLogger.Output(3, log.content)
		}
	case LevelFatal:
		DefaultLogger.Output(3, log.content)
		os.Exit(1)
	default:
	}
}

func Debug(v ...interface{}) {
	output(&logContent{
		level:   LevelDebug,
		content: fmt.Sprintln(v...),
	})
}

func Info(v ...interface{}) {
	output(&logContent{
		level:   LevelInfo,
		content: fmt.Sprintln(v...),
	})
}

func Warn(v ...interface{}) {
	output(&logContent{
		level:   LevelWarn,
		content: fmt.Sprintln(v...),
	})
}

func Error(v ...interface{}) {
	output(&logContent{
		level:   LevelError,
		content: fmt.Sprintln(v...),
	})
}

func Fatal(v ...interface{}) {
	output(&logContent{
		level:   LevelFatal,
		content: fmt.Sprintln(v...),
	})
}

func Debugf(format string, v ...interface{}) {
	output(&logContent{
		level:   LevelDebug,
		content: fmt.Sprintf(format, v...),
	})
}

func Infof(format string, v ...interface{}) {
	output(&logContent{
		level:   LevelInfo,
		content: fmt.Sprintf(format, v...),
	})
}

func Warnf(format string, v ...interface{}) {
	output(&logContent{
		level:   LevelWarn,
		content: fmt.Sprintf(format, v...),
	})
}

func Errorf(format string, v ...interface{}) {
	output(&logContent{
		level:   LevelError,
		content: fmt.Sprintf(format, v...),
	})
}

func Fatalf(format string, v ...interface{}) {
	output(&logContent{
		level:   LevelFatal,
		content: fmt.Sprintf(format, v...),
	})
}

func SetOptput(w io.Writer) {
	DefaultLogger.SetOutput(w)
}

func SetFlags(flag int) {
	DefaultLogger.SetFlags(flag)
}

func SetLevel(level int) {
	logLevel = level
}

func init() {
	logChan = make(chan *logContent)
	DefaultLogger = log.New(os.Stdout, "", log.LstdFlags|log.Lshortfile)
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"reflect"
	"strings"
	"text/template"
	"time"

	"github.com/Sirupsen/logrus"
)

// Struct for our apps nested with tasks.
type MarathonApps struct {
	Apps []struct {
		Id              string            `json:"id"`
		Labels          map[string]string `json:"labels"`
		Env             map[string]string `json:"env"`
		HealthChecks    []interface{}     `json:"healthChecks"`
		PortDefinitions []struct {
			Port     int64             `json:"port"`
			Protocol string            `json:"protocol"`
			Labels   map[string]string `json:"labels"`
		} `json:"portDefinitions"`
		Tasks []struct {
			AppId              string `json:"appId"`
			HealthCheckResults []struct {
				Alive bool `json:"alive"`
			} `json:"healthCheckResults"`
			Host         string  `json:"host"`
			Id           string  `json:"id"`
			Ports        []int64 `json:"ports"`
			ServicePorts []int64 `json:"servicePorts"`
			StagedAt     string  `json:"stagedAt"`
			StartedAt    string  `json:"startedAt"`
			Version      string  `json:"version"`
		} `json:"tasks"`
	} `json:"apps"`
}

func eventStream() {
	go func() {
		client := &http.Client{
			Timeout:   0 * time.Second,
			Transport: tr,
		}
		ticker := time.NewTicker(1 * time.Second)
		for _ = range ticker.C {
			var endpoint string
			for _, es := range health.Endpoints {
				if es.Healthy == true {
					endpoint = es.Endpoint
					break
				}
			}
			if endpoint == "" {
				logger.Error("all endpoints are down")
				continue
			}
			req, err := http.NewRequest("GET", endpoint+"/v2/events", nil)
			if err != nil {
				logger.WithFields(logrus.Fields{
					"error":    err.Error(),
					"endpoint": endpoint,
				}).Error("unable to create event stream request")
				continue
			}
			req.Header.Set("Accept", "text/event-stream")
			if config.User != "" {
				req.SetBasicAuth(config.User, config.Pass)
			}
			// Using new context package from Go 1.7
			ctx, cancel := context.WithCancel(context.TODO())
			// initial request cancellation timer of 15s
			timer := time.AfterFunc(15*time.Second, func() {
				cancel()
				logger.Warn("No data for 15s, event stream request was cancelled")
			})
			req = req.WithContext(ctx)
			resp, err := client.Do(req)
			if err != nil {
				logger.WithFields(logrus.Fields{
					"error":    err.Error(),
					"endpoint": endpoint,
				}).Error("unable to access Marathon event stream")
				// expire request cancellation timer immediately
				timer.Reset(100 * time.Millisecond)
				continue
			}
			reader := bufio.NewReader(resp.Body)
			for {
				// reset request cancellation timer to 15s (should be >10s to avoid unnecessary reconnects
				// since ~10s seems to be the rate for dummy/keepalive events on the marathon event stream
				timer.Reset(15 * time.Second)
				line, err := reader.ReadString('\n')
				if err != nil {
					logger.WithFields(logrus.Fields{
						"error":    err.Error(),
						"endpoint": endpoint,
					}).Error("error reading Marathon event stream")
					resp.Body.Close()
					break
				}
				if !strings.HasPrefix(line, "event: ") {
					continue
				}
				logger.WithFields(logrus.Fields{
					"event":    strings.TrimSpace(line[6:]),
					"endpoint": endpoint,
				}).Info("marathon event received")
				select {
				case eventqueue <- true: // add reload to our queue channel, unless it is full of course.
				default:
					logger.Warn("queue is full")
				}
			}
			resp.Body.Close()
			logger.Warn("event stream connection was closed, re-opening")
		}
	}()
}

func endpointHealth() {
	go func() {
		ticker := time.NewTicker(2 * time.Second)
		for {
			select {
			case <-ticker.C:
				for i, es := range health.Endpoints {
					client := &http.Client{
						Timeout:   5 * time.Second,
						Transport: tr,
					}
					req, err := http.NewRequest("GET", es.Endpoint+"/ping", nil)
					if err != nil {
						logger.WithFields(logrus.Fields{
							"error":    err.Error(),
							"endpoint": es.Endpoint,
						}).Error("an error occurred creating endpoint health request")
						health.Endpoints[i].Healthy = false
						health.Endpoints[i].Message = err.Error()
						continue
					}
					if config.User != "" {
						req.SetBasicAuth(config.User, config.Pass)
					}
					resp, err := client.Do(req)
					if err != nil {
						logger.WithFields(logrus.Fields{
							"error":    err.Error(),
							"endpoint": es.Endpoint,
						}).Error("endpoint is down")
						health.Endpoints[i].Healthy = false
						health.Endpoints[i].Message = err.Error()
						continue
					}
					resp.Body.Close()
					if resp.StatusCode != 200 {
						logger.WithFields(logrus.Fields{
							"status":   resp.StatusCode,
							"endpoint": es.Endpoint,
						}).Error("endpoint check failed")
						health.Endpoints[i].Healthy = false
						health.Endpoints[i].Message = resp.Status
						continue
					}
					health.Endpoints[i].Healthy = true
					health.Endpoints[i].Message = "OK"
				}
			}
		}
	}()
}

func eventWorker() {
	go func() {
		// a ticker channel to limit reloads to marathon, 1s is enough for now.
		ticker := time.NewTicker(1 * time.Second)
		for {
			select {
			case <-ticker.C:
				<-eventqueue
				reload()
			}
		}
	}()
}

func fetchApps(jsonapps *MarathonApps) error {
	var endpoint string
	for _, es := range health.Endpoints {
		if es.Healthy == true {
			endpoint = es.Endpoint
			break
		}
	}
	if endpoint == "" {
		err := errors.New("all endpoints are down")
		return err
	}
	client := &http.Client{
		Timeout:   5 * time.Second,
		Transport: tr,
	}
	// fetch all apps and tasks with a single request.
	req, err := http.NewRequest("GET", endpoint+"/v2/apps?embed=apps.tasks", nil)
	if err != nil {
		return err
	}
	req.Header.Set("Accept", "application/json")
	if config.User != "" {
		req.SetBasicAuth(config.User, config.Pass)
	}
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	decoder := json.NewDecoder(resp.Body)
	err = decoder.Decode(&jsonapps)
	if err != nil {
		return err
	}
	return nil
}

func syncApps(jsonapps *MarathonApps) bool {
	config.Lock()
	defer config.Unlock()
	apps := make(map[string]App)
	for _, app := range jsonapps.Apps {
		var newapp = App{}
		for _, task := range app.Tasks {
			// lets skip tasks that does not expose any ports.
			if len(task.Ports) == 0 {
				continue
			}
			if len(app.HealthChecks) > 0 {
				if len(task.HealthCheckResults) == 0 {
					// this means tasks is being deployed but not yet monitored as alive. Assume down.
					continue
				}
				alive := true
				for _, health := range task.HealthCheckResults {
					// check if health check is alive
					if health.Alive == false {
						alive = false
					}
				}
				if alive != true {
					// at least one health check has failed. Assume down.
					continue
				}
			}
			var newtask = Task{}
			newtask.Host = task.Host
			newtask.Ports = task.Ports
			newtask.ServicePorts = task.ServicePorts
			newtask.StagedAt = task.StagedAt
			newtask.StartedAt = task.StartedAt
			newtask.Version = task.Version
			newapp.Tasks = append(newapp.Tasks, newtask)
		}
		// Lets ignore apps if no tasks are available
		if len(newapp.Tasks) > 0 {
			if s, ok := app.Labels["subdomain"]; ok {
				hosts := strings.Split(s, " ")
				for _, host := range hosts {
					newapp.Hosts = append(newapp.Hosts, host)
				}
			} else if s, ok := app.Labels["moxy_subdomain"]; ok {
				// to be compatible with moxy
				hosts := strings.Split(s, " ")
				for _, host := range hosts {
					newapp.Hosts = append(newapp.Hosts, host)
				}
			} else {
				// If directories are used lets use them as subdomain dividers.
				// Ex: /project/app becomes app.project
				// Ex: /app becomes just app
				domains := strings.Split(app.Id[1:], "/")
				for i, j := 0, len(domains)-1; i < j; i, j = i+1, j-1 {
					domains[i], domains[j] = domains[j], domains[i]
				}
				host := strings.Join(domains, ".")
				newapp.Hosts = append(newapp.Hosts, host)
			}
			// Check for duplicated subdomain labels
			for _, confapp := range apps {
				for _, host := range confapp.Hosts {
					for _, newhost := range newapp.Hosts {
						if newhost == host {
							logger.WithFields(logrus.Fields{
								"app":       app.Id,
								"subdomain": host,
							}).Warn("duplicate subdomain label")
							// reset hosts if duplicate.
							newapp.Hosts = nil
						}
					}
				}
			}
			// Got duplicated subdomains, lets ignore this one.
			if len(newapp.Hosts) == 0 {
				continue
			}
			newapp.Labels = app.Labels
			newapp.Env = app.Env
			for _, pds := range app.PortDefinitions {
				pd := PortDefinitions{
					Port:     pds.Port,
					Protocol: pds.Protocol,
					Labels:   pds.Labels,
				}
				newapp.PortDefinitions = append(newapp.PortDefinitions, pd)
			}
			apps[app.Id] = newapp
		}
	}
	// Not all events bring changes, so lets see if anything is new.
	eq := reflect.DeepEqual(apps, config.Apps)
	if eq {
		return true
	} else {
		config.Apps = apps
		return false
	}
}

func writeConf() error {
	config.RLock()
	defer config.RUnlock()
	template, err := template.New(filepath.Base(config.Nginx_template)).ParseFiles(config.Nginx_template)
	if err != nil {
		return err
	}

	parent := filepath.Dir(config.Nginx_config)
	tmpFile, err := ioutil.TempFile(parent, ".nginx.conf.tmp-")
	defer tmpFile.Close()

	err = template.Execute(tmpFile, config)
	if err != nil {
		return err
	}
	config.LastUpdates.LastConfigRendered = time.Now()
	err = checkConf(tmpFile.Name())
	if err != nil {
		return err
	}
	err = os.Rename(tmpFile.Name(), config.Nginx_config)
	if err != nil {
		return err
	}
	return nil
}

func checkTmpl() error {
	config.RLock()
	defer config.RUnlock()
	t, err := template.New(filepath.Base(config.Nginx_template)).ParseFiles(config.Nginx_template)
	if err != nil {
		return err
	}
	err = t.Execute(ioutil.Discard, config)
	if err != nil {
		return err
	}
	return nil
}

func checkConf(path string) error {
	// This is to allow arguments as well. Example "docker exec nginx..."
	args := strings.Fields(config.Nginx_cmd)
	head := args[0]
	args = args[1:len(args)]
	args = append(args, "-c")
	args = append(args, path)
	args = append(args, "-t")
	cmd := exec.Command(head, args...)
	//cmd := exec.Command(parts..., "-c", path, "-t")
	var stderr bytes.Buffer
	cmd.Stderr = &stderr
	err := cmd.Run() // will wait for command to return
	if err != nil {
		msg := fmt.Sprint(err) + ": " + stderr.String()
		errstd := errors.New(msg)
		return errstd
	}
	return nil
}

func reloadNginx() error {
	// This is to allow arguments as well. Example "docker exec nginx..."
	args := strings.Fields(config.Nginx_cmd)
	head := args[0]
	args = args[1:len(args)]
	args = append(args, "-s")
	args = append(args, "reload")
	cmd := exec.Command(head, args...)
	var stderr bytes.Buffer
	cmd.Stderr = &stderr
	err := cmd.Run() // will wait for command to return
	if err != nil {
		msg := fmt.Sprint(err) + ": " + stderr.String()
		errstd := errors.New(msg)
		return errstd
	}
	return nil
}

func reload() {
	start := time.Now()
	jsonapps := MarathonApps{}
	err := fetchApps(&jsonapps)
	if err != nil {
		logger.WithFields(logrus.Fields{
			"error": err.Error(),
		}).Error("unable to sync from marathon")
		go statsCount("reload.failed", 1)
		return
	}
	equal := syncApps(&jsonapps)
	if equal {
		logger.Info("no config changes")
		return
	}
	config.LastUpdates.LastSync = time.Now()
	err = writeConf()
	if err != nil {
		logger.WithFields(logrus.Fields{
			"error": err.Error(),
		}).Error("unable to generate nginx config")
		go statsCount("reload.failed", 1)
		return
	}
	config.LastUpdates.LastConfigValid = time.Now()
	err = reloadNginx()
	if err != nil {
		logger.WithFields(logrus.Fields{
			"error": err.Error(),
		}).Error("unable to reload nginx")
		go statsCount("reload.failed", 1)
		return
	}
	elapsed := time.Since(start)
	logger.WithFields(logrus.Fields{
		"took": elapsed,
	}).Info("config updated")
	go statsCount("reload.success", 1)
	go statsTiming("reload.time", elapsed)
	config.LastUpdates.LastNginxReload = time.Now()
	return
}
<commit_msg>Add commonly used string methods and getenv method to make templates more configurable.<commit_after>package main

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"reflect"
	"strings"
	"text/template"
	"time"

	"github.com/Sirupsen/logrus"
)

// Struct for our apps nested with tasks.
type MarathonApps struct {
	Apps []struct {
		Id              string            `json:"id"`
		Labels          map[string]string `json:"labels"`
		Env             map[string]string `json:"env"`
		HealthChecks    []interface{}     `json:"healthChecks"`
		PortDefinitions []struct {
			Port     int64             `json:"port"`
			Protocol string            `json:"protocol"`
			Labels   map[string]string `json:"labels"`
		} `json:"portDefinitions"`
		Tasks []struct {
			AppId              string `json:"appId"`
			HealthCheckResults []struct {
				Alive bool `json:"alive"`
			} `json:"healthCheckResults"`
			Host         string  `json:"host"`
			Id           string  `json:"id"`
			Ports        []int64 `json:"ports"`
			ServicePorts []int64 `json:"servicePorts"`
			StagedAt     string  `json:"stagedAt"`
			StartedAt    string  `json:"startedAt"`
			Version      string  `json:"version"`
		} `json:"tasks"`
	} `json:"apps"`
}

func eventStream() {
	go func() {
		client := &http.Client{
			Timeout:   0 * time.Second,
			Transport: tr,
		}
		ticker := time.NewTicker(1 * time.Second)
		for _ = range ticker.C {
			var endpoint string
			for _, es := range health.Endpoints {
				if es.Healthy == true {
					endpoint = es.Endpoint
					break
				}
			}
			if endpoint == "" {
				logger.Error("all endpoints are down")
				continue
			}
			req, err := http.NewRequest("GET", endpoint+"/v2/events", nil)
			if err != nil {
				logger.WithFields(logrus.Fields{
					"error":    err.Error(),
					"endpoint": endpoint,
				}).Error("unable to create event stream request")
				continue
			}
			req.Header.Set("Accept", "text/event-stream")
			if config.User != "" {
				req.SetBasicAuth(config.User, config.Pass)
			}
			// Using new context package from Go 1.7
			ctx, cancel := context.WithCancel(context.TODO())
			// initial request cancellation timer of 15s
			timer := time.AfterFunc(15*time.Second, func() {
				cancel()
				logger.Warn("No data for 15s, event stream request was cancelled")
			})
			req = req.WithContext(ctx)
			resp, err := client.Do(req)
			if err != nil {
				logger.WithFields(logrus.Fields{
					"error":    err.Error(),
					"endpoint": endpoint,
				}).Error("unable to access Marathon event stream")
				// expire request cancellation timer immediately
				timer.Reset(100 * time.Millisecond)
				continue
			}
			reader := bufio.NewReader(resp.Body)
			for {
				// reset request cancellation timer to 15s (should be >10s to avoid unnecessary reconnects
				// since ~10s seems to be the rate for dummy/keepalive events on the marathon event stream
				timer.Reset(15 * time.Second)
				line, err := reader.ReadString('\n')
				if err != nil {
					logger.WithFields(logrus.Fields{
						"error":    err.Error(),
						"endpoint": endpoint,
					}).Error("error reading Marathon event stream")
					resp.Body.Close()
					break
				}
				if !strings.HasPrefix(line, "event: ") {
					continue
				}
				logger.WithFields(logrus.Fields{
					"event":    strings.TrimSpace(line[6:]),
					"endpoint": endpoint,
				}).Info("marathon event received")
				select {
				case eventqueue <- true: // add reload to our queue channel, unless it is full of course.
				default:
					logger.Warn("queue is full")
				}
			}
			resp.Body.Close()
			logger.Warn("event stream connection was closed, re-opening")
		}
	}()
}

func endpointHealth() {
	go func() {
		ticker := time.NewTicker(2 * time.Second)
		for {
			select {
			case <-ticker.C:
				for i, es := range health.Endpoints {
					client := &http.Client{
						Timeout:   5 * time.Second,
						Transport: tr,
					}
					req, err := http.NewRequest("GET", es.Endpoint+"/ping", nil)
					if err != nil {
						logger.WithFields(logrus.Fields{
							"error":    err.Error(),
							"endpoint": es.Endpoint,
						}).Error("an error occurred creating endpoint health request")
						health.Endpoints[i].Healthy = false
						health.Endpoints[i].Message = err.Error()
						continue
					}
					if config.User != "" {
						req.SetBasicAuth(config.User, config.Pass)
					}
					resp, err := client.Do(req)
					if err != nil {
						logger.WithFields(logrus.Fields{
							"error":    err.Error(),
							"endpoint": es.Endpoint,
						}).Error("endpoint is down")
						health.Endpoints[i].Healthy = false
						health.Endpoints[i].Message = err.Error()
						continue
					}
					resp.Body.Close()
					if resp.StatusCode != 200 {
						logger.WithFields(logrus.Fields{
							"status":   resp.StatusCode,
							"endpoint": es.Endpoint,
						}).Error("endpoint check failed")
						health.Endpoints[i].Healthy = false
						health.Endpoints[i].Message = resp.Status
						continue
					}
					health.Endpoints[i].Healthy = true
					health.Endpoints[i].Message = "OK"
				}
			}
		}
	}()
}

func eventWorker() {
	go func() {
		// a ticker channel to limit reloads to marathon, 1s is enough for now.
		ticker := time.NewTicker(1 * time.Second)
		for {
			select {
			case <-ticker.C:
				<-eventqueue
				reload()
			}
		}
	}()
}

func fetchApps(jsonapps *MarathonApps) error {
	var endpoint string
	for _, es := range health.Endpoints {
		if es.Healthy == true {
			endpoint = es.Endpoint
			break
		}
	}
	if endpoint == "" {
		err := errors.New("all endpoints are down")
		return err
	}
	client := &http.Client{
		Timeout:   5 * time.Second,
		Transport: tr,
	}
	// fetch all apps and tasks with a single request.
	req, err := http.NewRequest("GET", endpoint+"/v2/apps?embed=apps.tasks", nil)
	if err != nil {
		return err
	}
	req.Header.Set("Accept", "application/json")
	if config.User != "" {
		req.SetBasicAuth(config.User, config.Pass)
	}
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	decoder := json.NewDecoder(resp.Body)
	err = decoder.Decode(&jsonapps)
	if err != nil {
		return err
	}
	return nil
}

func syncApps(jsonapps *MarathonApps) bool {
	config.Lock()
	defer config.Unlock()
	apps := make(map[string]App)
	for _, app := range jsonapps.Apps {
		var newapp = App{}
		for _, task := range app.Tasks {
			// lets skip tasks that does not expose any ports.
			if len(task.Ports) == 0 {
				continue
			}
			if len(app.HealthChecks) > 0 {
				if len(task.HealthCheckResults) == 0 {
					// this means tasks is being deployed but not yet monitored as alive. Assume down.
					continue
				}
				alive := true
				for _, health := range task.HealthCheckResults {
					// check if health check is alive
					if health.Alive == false {
						alive = false
					}
				}
				if alive != true {
					// at least one health check has failed. Assume down.
					continue
				}
			}
			var newtask = Task{}
			newtask.Host = task.Host
			newtask.Ports = task.Ports
			newtask.ServicePorts = task.ServicePorts
			newtask.StagedAt = task.StagedAt
			newtask.StartedAt = task.StartedAt
			newtask.Version = task.Version
			newapp.Tasks = append(newapp.Tasks, newtask)
		}
		// Lets ignore apps if no tasks are available
		if len(newapp.Tasks) > 0 {
			if s, ok := app.Labels["subdomain"]; ok {
				hosts := strings.Split(s, " ")
				for _, host := range hosts {
					newapp.Hosts = append(newapp.Hosts, host)
				}
			} else if s, ok := app.Labels["moxy_subdomain"]; ok {
				// to be compatible with moxy
				hosts := strings.Split(s, " ")
				for _, host := range hosts {
					newapp.Hosts = append(newapp.Hosts, host)
				}
			} else {
				// If directories are used lets use them as subdomain dividers.
				// Ex: /project/app becomes app.project
				// Ex: /app becomes just app
				domains := strings.Split(app.Id[1:], "/")
				for i, j := 0, len(domains)-1; i < j; i, j = i+1, j-1 {
					domains[i], domains[j] = domains[j], domains[i]
				}
				host := strings.Join(domains, ".")
				newapp.Hosts = append(newapp.Hosts, host)
			}
			// Check for duplicated subdomain labels
			for _, confapp := range apps {
				for _, host := range confapp.Hosts {
					for _, newhost := range newapp.Hosts {
						if newhost == host {
							logger.WithFields(logrus.Fields{
								"app":       app.Id,
								"subdomain": host,
							}).Warn("duplicate subdomain label")
							// reset hosts if duplicate.
							newapp.Hosts = nil
						}
					}
				}
			}
			// Got duplicated subdomains, lets ignore this one.
			if len(newapp.Hosts) == 0 {
				continue
			}
			newapp.Labels = app.Labels
			newapp.Env = app.Env
			for _, pds := range app.PortDefinitions {
				pd := PortDefinitions{
					Port:     pds.Port,
					Protocol: pds.Protocol,
					Labels:   pds.Labels,
				}
				newapp.PortDefinitions = append(newapp.PortDefinitions, pd)
			}
			apps[app.Id] = newapp
		}
	}
	// Not all events bring changes, so lets see if anything is new.
	eq := reflect.DeepEqual(apps, config.Apps)
	if eq {
		return true
	} else {
		config.Apps = apps
		return false
	}
}

func writeConf() error {
	config.RLock()
	defer config.RUnlock()
	template, err := getTmpl()
	if err != nil {
		return err
	}

	parent := filepath.Dir(config.Nginx_config)
	tmpFile, err := ioutil.TempFile(parent, ".nginx.conf.tmp-")
	defer tmpFile.Close()

	err = template.Execute(tmpFile, config)
	if err != nil {
		return err
	}
	config.LastUpdates.LastConfigRendered = time.Now()
	err = checkConf(tmpFile.Name())
	if err != nil {
		return err
	}
	err = os.Rename(tmpFile.Name(), config.Nginx_config)
	if err != nil {
		return err
	}
	return nil
}

func checkTmpl() error {
	config.RLock()
	defer config.RUnlock()
	t, err := getTmpl()
	if err != nil {
		return err
	}
	err = t.Execute(ioutil.Discard, config)
	if err != nil {
		return err
	}
	return nil
}

func getTmpl() (*template.Template, error) {
	return template.New(filepath.Base(config.Nginx_template)).
		Funcs(template.FuncMap{
			"split": strings.Split,
			"join": strings.Join,
			"trim": strings.Trim,
			"replace": strings.Replace,
			"getenv": os.Getenv}).
		ParseFiles(config.Nginx_template)
}

func checkConf(path string) error {
	// This is to allow arguments as well. Example "docker exec nginx..."
	args := strings.Fields(config.Nginx_cmd)
	head := args[0]
	args = args[1:len(args)]
	args = append(args, "-c")
	args = append(args, path)
	args = append(args, "-t")
	cmd := exec.Command(head, args...)
	//cmd := exec.Command(parts..., "-c", path, "-t")
	var stderr bytes.Buffer
	cmd.Stderr = &stderr
	err := cmd.Run() // will wait for command to return
	if err != nil {
		msg := fmt.Sprint(err) + ": " + stderr.String()
		errstd := errors.New(msg)
		return errstd
	}
	return nil
}

func reloadNginx() error {
	// This is to allow arguments as well. Example "docker exec nginx..."
	args := strings.Fields(config.Nginx_cmd)
	head := args[0]
	args = args[1:len(args)]
	args = append(args, "-s")
	args = append(args, "reload")
	cmd := exec.Command(head, args...)
	var stderr bytes.Buffer
	cmd.Stderr = &stderr
	err := cmd.Run() // will wait for command to return
	if err != nil {
		msg := fmt.Sprint(err) + ": " + stderr.String()
		errstd := errors.New(msg)
		return errstd
	}
	return nil
}

func reload() {
	start := time.Now()
	jsonapps := MarathonApps{}
	err := fetchApps(&jsonapps)
	if err != nil {
		logger.WithFields(logrus.Fields{
			"error": err.Error(),
		}).Error("unable to sync from marathon")
		go statsCount("reload.failed", 1)
		return
	}
	equal := syncApps(&jsonapps)
	if equal {
		logger.Info("no config changes")
		return
	}
	config.LastUpdates.LastSync = time.Now()
	err = writeConf()
	if err != nil {
		logger.WithFields(logrus.Fields{
			"error": err.Error(),
		}).Error("unable to generate nginx config")
		go statsCount("reload.failed", 1)
		return
	}
	config.LastUpdates.LastConfigValid = time.Now()
	err = reloadNginx()
	if err != nil {
		logger.WithFields(logrus.Fields{
			"error": err.Error(),
		}).Error("unable to reload nginx")
		go statsCount("reload.failed", 1)
		return
	}
	elapsed := time.Since(start)
	logger.WithFields(logrus.Fields{
		"took": elapsed,
	}).Info("config updated")
	go statsCount("reload.success", 1)
	go statsTiming("reload.time", elapsed)
	config.LastUpdates.LastNginxReload = time.Now()
	return
}
<|endoftext|>


<commit_before>package peco

import (
	"fmt"
	"os/exec"
	"regexp"
	"sort"
	"strings"
)

// Global var used to strips ansi sequences
var ansiStrips = regexp.MustCompile("\x1B\\[(?:[0-9]{1,2}(?:;[0-9]{1,2})?)*[a-zA-Z]")

// Function who strips ansi sequences
func stripANSISequence(s string) string {
	return ansiStrips.ReplaceAllString(s, "")
}

// Match defines the interface for matches. Note that to make drawing easier,
// we have a DidMatch and NoMatch types instead of using []Match and []string.
type Match interface {
	Buffer() string // Raw buffer, may contain null
	Line() string   // Line to be displayed
	Output() string // Output string to be displayed after peco is done
	Indices() [][]int
}

type matchString struct {
	buf         string
	sepLoc      int
	displayLine string
}

func newMatchString(v string, enableSep bool) *matchString {
	m := &matchString{
		v,
		-1,
		"",
	}
	if !enableSep {
		return m
	}

	// XXX This may be silly, but we're avoiding using strings.IndexByte()
	// here because it doesn't exist on go1.1. Let's remove support for
	// 1.1 when 1.4 comes out (or something)
	for i := 0; i < len(m.buf); i++ {
		if m.buf[i] == '\000' {
			m.sepLoc = i
		}
	}
	return m
}

func (m matchString) Buffer() string {
	return m.buf
}

func (m matchString) Line() string {
	if m.displayLine != "" {
		return m.displayLine
	}

	if i := m.sepLoc; i > -1 {
		m.displayLine = stripANSISequence(m.buf[:i])
	} else {
		m.displayLine = stripANSISequence(m.buf)
	}
	return m.displayLine
}

func (m matchString) Output() string {
	if i := m.sepLoc; i > -1 {
		return m.buf[i+1:]
	}
	return m.buf
}

// NoMatch is actually an alias to a regular string. It implements the
// Match interface, but just returns the underlying string with no matches
type NoMatch struct {
	*matchString
}

// NewNoMatch creates a NoMatch struct
func NewNoMatch(v string, enableSep bool) *NoMatch {
	return &NoMatch{newMatchString(v, enableSep)}
}

// Indices always returns nil
func (m NoMatch) Indices() [][]int {
	return nil
}

// DidMatch contains the actual match, and the indices to the matches
// in the line
type DidMatch struct {
	*matchString
	matches [][]int
}

// NewDidMatch creates a new DidMatch struct
func NewDidMatch(v string, enableSep bool, m [][]int) *DidMatch {
	return &DidMatch{newMatchString(v, enableSep), m}
}

// Indices returns the indices in the buffer that matched
func (d DidMatch) Indices() [][]int {
	return d.matches
}

// Matcher interface defines the API for things that want to
// match against the buffer
type Matcher interface {
	// Match takes in three parameters.
	//
	// The first chan is the channel where cancel requests are sent.
	// If you receive a request here, you should stop running your query.
	//
	// The second is the query. Do what you want with it
	//
	// The third is the buffer in which to match the query against.
	Match(chan struct{}, string, []Match) []Match
	String() string

	// This is fugly. We just added a method only for CustomMatcner.
	// Must think about this again
	Verify() error
}

// These are used as keys in the config file
const (
	IgnoreCaseMatch    = "IgnoreCase"
	CaseSensitiveMatch = "CaseSensitive"
	RegexpMatch        = "Regexp"
)

// RegexpMatcher is the most basic matcher
type RegexpMatcher struct {
	enableSep bool
	flags     []string
	quotemeta bool
}

// CaseSensitiveMatcher extends the RegxpMatcher, but always
// turns off the ignore-case flag in the regexp
type CaseSensitiveMatcher struct {
	*RegexpMatcher
}

// IgnoreCaseMatcher extends the RegexpMatcher, and always
// turns ON the ignore-case flag in the regexp
type IgnoreCaseMatcher struct {
	*RegexpMatcher
}

// CustomMatcher spawns a new process to filter the buffer
// in peco, and uses the output in its Stdout to figure
// out what to display
type CustomMatcher struct {
	enableSep bool
	name      string
	args      []string
}

// NewCaseSensitiveMatcher creates a new CaseSensitiveMatcher
func NewCaseSensitiveMatcher(enableSep bool) *CaseSensitiveMatcher {
	m := &CaseSensitiveMatcher{NewRegexpMatcher(enableSep)}
	m.quotemeta = true
	return m
}

// NewIgnoreCaseMatcher creates a new IgnoreCaseMatcher
func NewIgnoreCaseMatcher(enableSep bool) *IgnoreCaseMatcher {
	m := &IgnoreCaseMatcher{NewRegexpMatcher(enableSep)}
	m.flags = []string{"i"}
	m.quotemeta = true
	return m
}

// NewRegexpMatcher creates a new RegexpMatcher
func NewRegexpMatcher(enableSep bool) *RegexpMatcher {
	return &RegexpMatcher{
		enableSep,
		[]string{},
		false,
	}
}

// Verify always returns nil
func (m *RegexpMatcher) Verify() error {
	return nil
}

// NewCustomMatcher creates a new CustomMatcher
func NewCustomMatcher(enableSep bool, name string, args []string) *CustomMatcher {
	return &CustomMatcher{enableSep, name, args}
}

// Verify checks to see that the executable given to CustomMatcher
// is actual found and is executable via exec.LookPath
func (m *CustomMatcher) Verify() error {
	if _, err := exec.LookPath(m.args[0]); err != nil {
		return err
	}
	return nil
}

func regexpFor(q string, flags []string, quotemeta bool) (*regexp.Regexp, error) {
	reTxt := q
	if quotemeta {
		reTxt = regexp.QuoteMeta(q)
	}

	if flags != nil && len(flags) > 0 {
		reTxt = fmt.Sprintf("(?%s)%s", strings.Join(flags, ""), reTxt)
	}

	re, err := regexp.Compile(reTxt)
	if err != nil {
		return nil, err
	}
	return re, nil
}

func (m *RegexpMatcher) queryToRegexps(query string) ([]*regexp.Regexp, error) {
	queries := strings.Split(strings.TrimSpace(query), " ")
	regexps := make([]*regexp.Regexp, 0)

	for _, q := range queries {
		re, err := regexpFor(q, m.flags, m.quotemeta)
		if err != nil {
			return nil, err
		}
		regexps = append(regexps, re)
	}

	return regexps, nil
}

func (m *RegexpMatcher) String() string {
	return "Regexp"
}

func (m *CaseSensitiveMatcher) String() string {
	return "CaseSensitive"
}

func (m *IgnoreCaseMatcher) String() string {
	return "IgnoreCase"
}

func (m *CustomMatcher) String() string {
	return m.name
}

// sort related stuff
type byStart [][]int

func (m byStart) Len() int {
	return len(m)
}

func (m byStart) Swap(i, j int) {
	m[i], m[j] = m[j], m[i]
}

func (m byStart) Less(i, j int) bool {
	return m[i][0] < m[j][0]
}

// Match does the heavy lifting, and matches `q` against `buffer`.
// While it is doing the match, it also listens for messages
// via `quit`. If anything is received via `quit`, the match
// is halted.
func (m *RegexpMatcher) Match(quit chan struct{}, q string, buffer []Match) []Match {
	results := []Match{}
	regexps, err := m.queryToRegexps(q)
	if err != nil {
		return results
	}

	// The actual matching is done in a separate goroutine
	iter := make(chan Match, len(buffer))
	go func() {
		// This protects us from panics, caused when we cancel the
		// query and forcefully close the channel (and thereby
		// causing a "close of a closed channel"
		defer func() { recover() }()

		// This must be here to make sure the channel is properly
		// closed in normal cases
		defer close(iter)

		// Iterate through the lines, and do the match.
		// Upon success, send it through the channel
		for _, match := range buffer {
			ms := m.MatchAllRegexps(regexps, match.Line())
			if ms == nil {
				continue
			}
			iter <- NewDidMatch(match.Buffer(), m.enableSep, ms)
		}
		iter <- nil
	}()

MATCH:
	for {
		select {
		case <-quit:
			// If we recieved a cancel request, we immediately bail out.
			// It's a little dirty, but we focefully terminate the other
			// goroutine by closing the channel, and invoking a panic in the
			// goroutine above

			// There's a possibility that the match fails early and the
			// cancel happens after iter has been closed. It's totally okay
			// for us to try to close iter, but trying to detect if the
			// channel can be closed safely synchronously is really hard
			// so we punt it by letting the close() happen at a separate
			// goroutine, protected by a defer recover()
			go func() {
				defer func() { recover() }()
				close(iter)
			}()
			break MATCH
		case match := <-iter:
			// Receive elements from the goroutine performing the match
			if match == nil {
				break MATCH
			}

			results = append(results, match)
		}
	}
	return results
}

// MatchAllRegexps matches all the regexps in `regexps` against line
func (m *RegexpMatcher) MatchAllRegexps(regexps []*regexp.Regexp, line string) [][]int {
	matches := make([][]int, 0)

	allMatched := true
Match:
	for _, re := range regexps {
		match := re.FindAllStringSubmatchIndex(line, -1)
		if match == nil {
			allMatched = false
			break Match
		}

		for _, ma := range match {
			start, end := ma[0], ma[1]
			for _, m := range matches {
				if start >= m[0] && start < m[1] {
					continue Match
				}

				if start < m[0] && end >= m[0] {
					continue Match
				}
			}
			matches = append(matches, ma)
		}
	}

	if !allMatched {
		return nil
	}

	sort.Sort(byStart(matches))

	return matches
}

// Match matches `q` aginst `buffer`
func (m *CustomMatcher) Match(quit chan struct{}, q string, buffer []Match) []Match {
	if len(m.args) < 1 {
		return []Match{}
	}

	results := []Match{}
	if q == "" {
		for _, match := range buffer {
			results = append(results, NewDidMatch(match.Buffer(), m.enableSep, nil))
		}
		return results
	}

	// Receive elements from the goroutine performing the match
	lines := []Match{}
	matcherInput := ""
	for _, match := range buffer {
		matcherInput += match.Line() + "\n"
		lines = append(lines, match)
	}
	args := []string{}
	for _, arg := range m.args {
		if arg == "$QUERY" {
			arg = q
		}
		args = append(args, arg)
	}
	cmd := exec.Command(args[0], args[1:]...)
	cmd.Stdin = strings.NewReader(matcherInput)

	// See RegexpMatcher.Match() for explanation of constructs
	iter := make(chan Match, len(buffer))
	go func() {
		defer func() { recover() }()
		defer func() {
			if p := cmd.Process; p != nil {
				p.Kill()
			}
			close(iter)
		}()
		b, err := cmd.Output()
		if err != nil {
			iter <- nil
		}
		for _, line := range strings.Split(string(b), "\n") {
			if len(line) > 0 {
				iter <- NewDidMatch(line, m.enableSep, nil)
			}
		}
		iter <- nil
	}()
MATCH:
	for {
		select {
		case <-quit:
			go func() {
				defer func() { recover() }()
				close(iter)
			}()
			break MATCH
		case match := <-iter:
			if match == nil {
				break MATCH
			}
			results = append(results, match)
		}
	}

	return results
}
<commit_msg>rename<commit_after>package peco

import (
	"fmt"
	"os/exec"
	"regexp"
	"sort"
	"strings"
)

// Global var used to strips ansi sequences
var reANSIEscapeChars = regexp.MustCompile("\x1B\\[(?:[0-9]{1,2}(?:;[0-9]{1,2})?)*[a-zA-Z]")

// Function who strips ansi sequences
func stripANSISequence(s string) string {
	return reANSIEscapeChars.ReplaceAllString(s, "")
}

// Match defines the interface for matches. Note that to make drawing easier,
// we have a DidMatch and NoMatch types instead of using []Match and []string.
type Match interface {
	Buffer() string // Raw buffer, may contain null
	Line() string   // Line to be displayed
	Output() string // Output string to be displayed after peco is done
	Indices() [][]int
}

type matchString struct {
	buf         string
	sepLoc      int
	displayLine string
}

func newMatchString(v string, enableSep bool) *matchString {
	m := &matchString{
		v,
		-1,
		"",
	}
	if !enableSep {
		return m
	}

	// XXX This may be silly, but we're avoiding using strings.IndexByte()
	// here because it doesn't exist on go1.1. Let's remove support for
	// 1.1 when 1.4 comes out (or something)
	for i := 0; i < len(m.buf); i++ {
		if m.buf[i] == '\000' {
			m.sepLoc = i
		}
	}
	return m
}

func (m matchString) Buffer() string {
	return m.buf
}

func (m matchString) Line() string {
	if m.displayLine != "" {
		return m.displayLine
	}

	if i := m.sepLoc; i > -1 {
		m.displayLine = stripANSISequence(m.buf[:i])
	} else {
		m.displayLine = stripANSISequence(m.buf)
	}
	return m.displayLine
}

func (m matchString) Output() string {
	if i := m.sepLoc; i > -1 {
		return m.buf[i+1:]
	}
	return m.buf
}

// NoMatch is actually an alias to a regular string. It implements the
// Match interface, but just returns the underlying string with no matches
type NoMatch struct {
	*matchString
}

// NewNoMatch creates a NoMatch struct
func NewNoMatch(v string, enableSep bool) *NoMatch {
	return &NoMatch{newMatchString(v, enableSep)}
}

// Indices always returns nil
func (m NoMatch) Indices() [][]int {
	return nil
}

// DidMatch contains the actual match, and the indices to the matches
// in the line
type DidMatch struct {
	*matchString
	matches [][]int
}

// NewDidMatch creates a new DidMatch struct
func NewDidMatch(v string, enableSep bool, m [][]int) *DidMatch {
	return &DidMatch{newMatchString(v, enableSep), m}
}

// Indices returns the indices in the buffer that matched
func (d DidMatch) Indices() [][]int {
	return d.matches
}

// Matcher interface defines the API for things that want to
// match against the buffer
type Matcher interface {
	// Match takes in three parameters.
	//
	// The first chan is the channel where cancel requests are sent.
	// If you receive a request here, you should stop running your query.
	//
	// The second is the query. Do what you want with it
	//
	// The third is the buffer in which to match the query against.
	Match(chan struct{}, string, []Match) []Match
	String() string

	// This is fugly. We just added a method only for CustomMatcner.
	// Must think about this again
	Verify() error
}

// These are used as keys in the config file
const (
	IgnoreCaseMatch    = "IgnoreCase"
	CaseSensitiveMatch = "CaseSensitive"
	RegexpMatch        = "Regexp"
)

// RegexpMatcher is the most basic matcher
type RegexpMatcher struct {
	enableSep bool
	flags     []string
	quotemeta bool
}

// CaseSensitiveMatcher extends the RegxpMatcher, but always
// turns off the ignore-case flag in the regexp
type CaseSensitiveMatcher struct {
	*RegexpMatcher
}

// IgnoreCaseMatcher extends the RegexpMatcher, and always
// turns ON the ignore-case flag in the regexp
type IgnoreCaseMatcher struct {
	*RegexpMatcher
}

// CustomMatcher spawns a new process to filter the buffer
// in peco, and uses the output in its Stdout to figure
// out what to display
type CustomMatcher struct {
	enableSep bool
	name      string
	args      []string
}

// NewCaseSensitiveMatcher creates a new CaseSensitiveMatcher
func NewCaseSensitiveMatcher(enableSep bool) *CaseSensitiveMatcher {
	m := &CaseSensitiveMatcher{NewRegexpMatcher(enableSep)}
	m.quotemeta = true
	return m
}

// NewIgnoreCaseMatcher creates a new IgnoreCaseMatcher
func NewIgnoreCaseMatcher(enableSep bool) *IgnoreCaseMatcher {
	m := &IgnoreCaseMatcher{NewRegexpMatcher(enableSep)}
	m.flags = []string{"i"}
	m.quotemeta = true
	return m
}

// NewRegexpMatcher creates a new RegexpMatcher
func NewRegexpMatcher(enableSep bool) *RegexpMatcher {
	return &RegexpMatcher{
		enableSep,
		[]string{},
		false,
	}
}

// Verify always returns nil
func (m *RegexpMatcher) Verify() error {
	return nil
}

// NewCustomMatcher creates a new CustomMatcher
func NewCustomMatcher(enableSep bool, name string, args []string) *CustomMatcher {
	return &CustomMatcher{enableSep, name, args}
}

// Verify checks to see that the executable given to CustomMatcher
// is actual found and is executable via exec.LookPath
func (m *CustomMatcher) Verify() error {
	if _, err := exec.LookPath(m.args[0]); err != nil {
		return err
	}
	return nil
}

func regexpFor(q string, flags []string, quotemeta bool) (*regexp.Regexp, error) {
	reTxt := q
	if quotemeta {
		reTxt = regexp.QuoteMeta(q)
	}

	if flags != nil && len(flags) > 0 {
		reTxt = fmt.Sprintf("(?%s)%s", strings.Join(flags, ""), reTxt)
	}

	re, err := regexp.Compile(reTxt)
	if err != nil {
		return nil, err
	}
	return re, nil
}

func (m *RegexpMatcher) queryToRegexps(query string) ([]*regexp.Regexp, error) {
	queries := strings.Split(strings.TrimSpace(query), " ")
	regexps := make([]*regexp.Regexp, 0)

	for _, q := range queries {
		re, err := regexpFor(q, m.flags, m.quotemeta)
		if err != nil {
			return nil, err
		}
		regexps = append(regexps, re)
	}

	return regexps, nil
}

func (m *RegexpMatcher) String() string {
	return "Regexp"
}

func (m *CaseSensitiveMatcher) String() string {
	return "CaseSensitive"
}

func (m *IgnoreCaseMatcher) String() string {
	return "IgnoreCase"
}

func (m *CustomMatcher) String() string {
	return m.name
}

// sort related stuff
type byStart [][]int

func (m byStart) Len() int {
	return len(m)
}

func (m byStart) Swap(i, j int) {
	m[i], m[j] = m[j], m[i]
}

func (m byStart) Less(i, j int) bool {
	return m[i][0] < m[j][0]
}

// Match does the heavy lifting, and matches `q` against `buffer`.
// While it is doing the match, it also listens for messages
// via `quit`. If anything is received via `quit`, the match
// is halted.
func (m *RegexpMatcher) Match(quit chan struct{}, q string, buffer []Match) []Match {
	results := []Match{}
	regexps, err := m.queryToRegexps(q)
	if err != nil {
		return results
	}

	// The actual matching is done in a separate goroutine
	iter := make(chan Match, len(buffer))
	go func() {
		// This protects us from panics, caused when we cancel the
		// query and forcefully close the channel (and thereby
		// causing a "close of a closed channel"
		defer func() { recover() }()

		// This must be here to make sure the channel is properly
		// closed in normal cases
		defer close(iter)

		// Iterate through the lines, and do the match.
		// Upon success, send it through the channel
		for _, match := range buffer {
			ms := m.MatchAllRegexps(regexps, match.Line())
			if ms == nil {
				continue
			}
			iter <- NewDidMatch(match.Buffer(), m.enableSep, ms)
		}
		iter <- nil
	}()

MATCH:
	for {
		select {
		case <-quit:
			// If we recieved a cancel request, we immediately bail out.
			// It's a little dirty, but we focefully terminate the other
			// goroutine by closing the channel, and invoking a panic in the
			// goroutine above

			// There's a possibility that the match fails early and the
			// cancel happens after iter has been closed. It's totally okay
			// for us to try to close iter, but trying to detect if the
			// channel can be closed safely synchronously is really hard
			// so we punt it by letting the close() happen at a separate
			// goroutine, protected by a defer recover()
			go func() {
				defer func() { recover() }()
				close(iter)
			}()
			break MATCH
		case match := <-iter:
			// Receive elements from the goroutine performing the match
			if match == nil {
				break MATCH
			}

			results = append(results, match)
		}
	}
	return results
}

// MatchAllRegexps matches all the regexps in `regexps` against line
func (m *RegexpMatcher) MatchAllRegexps(regexps []*regexp.Regexp, line string) [][]int {
	matches := make([][]int, 0)

	allMatched := true
Match:
	for _, re := range regexps {
		match := re.FindAllStringSubmatchIndex(line, -1)
		if match == nil {
			allMatched = false
			break Match
		}

		for _, ma := range match {
			start, end := ma[0], ma[1]
			for _, m := range matches {
				if start >= m[0] && start < m[1] {
					continue Match
				}

				if start < m[0] && end >= m[0] {
					continue Match
				}
			}
			matches = append(matches, ma)
		}
	}

	if !allMatched {
		return nil
	}

	sort.Sort(byStart(matches))

	return matches
}

// Match matches `q` aginst `buffer`
func (m *CustomMatcher) Match(quit chan struct{}, q string, buffer []Match) []Match {
	if len(m.args) < 1 {
		return []Match{}
	}

	results := []Match{}
	if q == "" {
		for _, match := range buffer {
			results = append(results, NewDidMatch(match.Buffer(), m.enableSep, nil))
		}
		return results
	}

	// Receive elements from the goroutine performing the match
	lines := []Match{}
	matcherInput := ""
	for _, match := range buffer {
		matcherInput += match.Line() + "\n"
		lines = append(lines, match)
	}
	args := []string{}
	for _, arg := range m.args {
		if arg == "$QUERY" {
			arg = q
		}
		args = append(args, arg)
	}
	cmd := exec.Command(args[0], args[1:]...)
	cmd.Stdin = strings.NewReader(matcherInput)

	// See RegexpMatcher.Match() for explanation of constructs
	iter := make(chan Match, len(buffer))
	go func() {
		defer func() { recover() }()
		defer func() {
			if p := cmd.Process; p != nil {
				p.Kill()
			}
			close(iter)
		}()
		b, err := cmd.Output()
		if err != nil {
			iter <- nil
		}
		for _, line := range strings.Split(string(b), "\n") {
			if len(line) > 0 {
				iter <- NewDidMatch(line, m.enableSep, nil)
			}
		}
		iter <- nil
	}()
MATCH:
	for {
		select {
		case <-quit:
			go func() {
				defer func() { recover() }()
				close(iter)
			}()
			break MATCH
		case match := <-iter:
			if match == nil {
				break MATCH
			}
			results = append(results, match)
		}
	}

	return results
}
<|endoftext|>


<commit_before>// Code generated by protoc-gen-go.
// source: apartment.proto
// DO NOT EDIT!

/*
Package apartment is a generated protocol buffer package.

It is generated from these files:
	apartment.proto

It has these top-level messages:
	Device
	ListDevicesRequest
	ListDevicesResponse
	GetDeviceRequest
	UpdateDeviceRequest
*/
package apartment

import proto "github.com/golang/protobuf/proto"
import fmt "fmt"
import math "math"

import (
	context "golang.org/x/net/context"
	grpc "google.golang.org/grpc"
)

// Reference imports to suppress errors if they are not otherwise used.
var _ = proto.Marshal
var _ = fmt.Errorf
var _ = math.Inf

// This is a compile-time assertion to ensure that this generated file
// is compatible with the proto package it is being compiled against.
// A compilation error at this line likely means your copy of the
// proto package needs to be updated.
const _ = proto.ProtoPackageIsVersion2 // please upgrade the proto package

type Device struct {
	Name         string `protobuf:"bytes,1,opt,name=name" json:"name,omitempty"`
	FriendlyName string `protobuf:"bytes,2,opt,name=friendly_name,json=friendlyName" json:"friendly_name,omitempty"`
	State        bool   `protobuf:"varint,3,opt,name=state" json:"state,omitempty"`
}

func (m *Device) Reset()                    { *m = Device{} }
func (m *Device) String() string            { return proto.CompactTextString(m) }
func (*Device) ProtoMessage()               {}
func (*Device) Descriptor() ([]byte, []int) { return fileDescriptor0, []int{0} }

type ListDevicesRequest struct {
}

func (m *ListDevicesRequest) Reset()                    { *m = ListDevicesRequest{} }
func (m *ListDevicesRequest) String() string            { return proto.CompactTextString(m) }
func (*ListDevicesRequest) ProtoMessage()               {}
func (*ListDevicesRequest) Descriptor() ([]byte, []int) { return fileDescriptor0, []int{1} }

type ListDevicesResponse struct {
	Device []*Device `protobuf:"bytes,1,rep,name=device" json:"device,omitempty"`
}

func (m *ListDevicesResponse) Reset()                    { *m = ListDevicesResponse{} }
func (m *ListDevicesResponse) String() string            { return proto.CompactTextString(m) }
func (*ListDevicesResponse) ProtoMessage()               {}
func (*ListDevicesResponse) Descriptor() ([]byte, []int) { return fileDescriptor0, []int{2} }

func (m *ListDevicesResponse) GetDevice() []*Device {
	if m != nil {
		return m.Device
	}
	return nil
}

type GetDeviceRequest struct {
	Name string `protobuf:"bytes,1,opt,name=name" json:"name,omitempty"`
}

func (m *GetDeviceRequest) Reset()                    { *m = GetDeviceRequest{} }
func (m *GetDeviceRequest) String() string            { return proto.CompactTextString(m) }
func (*GetDeviceRequest) ProtoMessage()               {}
func (*GetDeviceRequest) Descriptor() ([]byte, []int) { return fileDescriptor0, []int{3} }

type UpdateDeviceRequest struct {
	Device *Device `protobuf:"bytes,1,opt,name=device" json:"device,omitempty"`
}

func (m *UpdateDeviceRequest) Reset()                    { *m = UpdateDeviceRequest{} }
func (m *UpdateDeviceRequest) String() string            { return proto.CompactTextString(m) }
func (*UpdateDeviceRequest) ProtoMessage()               {}
func (*UpdateDeviceRequest) Descriptor() ([]byte, []int) { return fileDescriptor0, []int{4} }

func (m *UpdateDeviceRequest) GetDevice() *Device {
	if m != nil {
		return m.Device
	}
	return nil
}

func init() {
	proto.RegisterType((*Device)(nil), "apartment.Device")
	proto.RegisterType((*ListDevicesRequest)(nil), "apartment.ListDevicesRequest")
	proto.RegisterType((*ListDevicesResponse)(nil), "apartment.ListDevicesResponse")
	proto.RegisterType((*GetDeviceRequest)(nil), "apartment.GetDeviceRequest")
	proto.RegisterType((*UpdateDeviceRequest)(nil), "apartment.UpdateDeviceRequest")
}

// Reference imports to suppress errors if they are not otherwise used.
var _ context.Context
var _ grpc.ClientConn

// This is a compile-time assertion to ensure that this generated file
// is compatible with the grpc package it is being compiled against.
const _ = grpc.SupportPackageIsVersion3

// Client API for Apartment service

type ApartmentClient interface {
	ListDevices(ctx context.Context, in *ListDevicesRequest, opts ...grpc.CallOption) (*ListDevicesResponse, error)
	GetDevice(ctx context.Context, in *GetDeviceRequest, opts ...grpc.CallOption) (*Device, error)
	UpdateDevice(ctx context.Context, in *UpdateDeviceRequest, opts ...grpc.CallOption) (*Device, error)
}

type apartmentClient struct {
	cc *grpc.ClientConn
}

func NewApartmentClient(cc *grpc.ClientConn) ApartmentClient {
	return &apartmentClient{cc}
}

func (c *apartmentClient) ListDevices(ctx context.Context, in *ListDevicesRequest, opts ...grpc.CallOption) (*ListDevicesResponse, error) {
	out := new(ListDevicesResponse)
	err := grpc.Invoke(ctx, "/apartment.Apartment/ListDevices", in, out, c.cc, opts...)
	if err != nil {
		return nil, err
	}
	return out, nil
}

func (c *apartmentClient) GetDevice(ctx context.Context, in *GetDeviceRequest, opts ...grpc.CallOption) (*Device, error) {
	out := new(Device)
	err := grpc.Invoke(ctx, "/apartment.Apartment/GetDevice", in, out, c.cc, opts...)
	if err != nil {
		return nil, err
	}
	return out, nil
}

func (c *apartmentClient) UpdateDevice(ctx context.Context, in *UpdateDeviceRequest, opts ...grpc.CallOption) (*Device, error) {
	out := new(Device)
	err := grpc.Invoke(ctx, "/apartment.Apartment/UpdateDevice", in, out, c.cc, opts...)
	if err != nil {
		return nil, err
	}
	return out, nil
}

// Server API for Apartment service

type ApartmentServer interface {
	ListDevices(context.Context, *ListDevicesRequest) (*ListDevicesResponse, error)
	GetDevice(context.Context, *GetDeviceRequest) (*Device, error)
	UpdateDevice(context.Context, *UpdateDeviceRequest) (*Device, error)
}

func RegisterApartmentServer(s *grpc.Server, srv ApartmentServer) {
	s.RegisterService(&_Apartment_serviceDesc, srv)
}

func _Apartment_ListDevices_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
	in := new(ListDevicesRequest)
	if err := dec(in); err != nil {
		return nil, err
	}
	if interceptor == nil {
		return srv.(ApartmentServer).ListDevices(ctx, in)
	}
	info := &grpc.UnaryServerInfo{
		Server:     srv,
		FullMethod: "/apartment.Apartment/ListDevices",
	}
	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
		return srv.(ApartmentServer).ListDevices(ctx, req.(*ListDevicesRequest))
	}
	return interceptor(ctx, in, info, handler)
}

func _Apartment_GetDevice_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
	in := new(GetDeviceRequest)
	if err := dec(in); err != nil {
		return nil, err
	}
	if interceptor == nil {
		return srv.(ApartmentServer).GetDevice(ctx, in)
	}
	info := &grpc.UnaryServerInfo{
		Server:     srv,
		FullMethod: "/apartment.Apartment/GetDevice",
	}
	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
		return srv.(ApartmentServer).GetDevice(ctx, req.(*GetDeviceRequest))
	}
	return interceptor(ctx, in, info, handler)
}

func _Apartment_UpdateDevice_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
	in := new(UpdateDeviceRequest)
	if err := dec(in); err != nil {
		return nil, err
	}
	if interceptor == nil {
		return srv.(ApartmentServer).UpdateDevice(ctx, in)
	}
	info := &grpc.UnaryServerInfo{
		Server:     srv,
		FullMethod: "/apartment.Apartment/UpdateDevice",
	}
	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
		return srv.(ApartmentServer).UpdateDevice(ctx, req.(*UpdateDeviceRequest))
	}
	return interceptor(ctx, in, info, handler)
}

var _Apartment_serviceDesc = grpc.ServiceDesc{
	ServiceName: "apartment.Apartment",
	HandlerType: (*ApartmentServer)(nil),
	Methods: []grpc.MethodDesc{
		{
			MethodName: "ListDevices",
			Handler:    _Apartment_ListDevices_Handler,
		},
		{
			MethodName: "GetDevice",
			Handler:    _Apartment_GetDevice_Handler,
		},
		{
			MethodName: "UpdateDevice",
			Handler:    _Apartment_UpdateDevice_Handler,
		},
	},
	Streams:  []grpc.StreamDesc{},
	Metadata: fileDescriptor0,
}

func init() { proto.RegisterFile("apartment.proto", fileDescriptor0) }

var fileDescriptor0 = []byte{
	// 258 bytes of a gzipped FileDescriptorProto
	0x1f, 0x8b, 0x08, 0x00, 0x00, 0x09, 0x6e, 0x88, 0x02, 0xff, 0xe2, 0xe2, 0x4f, 0x2c, 0x48, 0x2c,
	0x2a, 0xc9, 0x4d, 0xcd, 0x2b, 0xd1, 0x2b, 0x28, 0xca, 0x2f, 0xc9, 0x17, 0xe2, 0x84, 0x0b, 0x28,
	0x85, 0x73, 0xb1, 0xb9, 0xa4, 0x96, 0x65, 0x26, 0xa7, 0x0a, 0x09, 0x71, 0xb1, 0xe4, 0x25, 0xe6,
	0xa6, 0x4a, 0x30, 0x2a, 0x30, 0x6a, 0x70, 0x06, 0x81, 0xd9, 0x42, 0xca, 0x5c, 0xbc, 0x69, 0x45,
	0x99, 0xa9, 0x79, 0x29, 0x39, 0x95, 0xf1, 0x60, 0x49, 0x26, 0xb0, 0x24, 0x0f, 0x4c, 0xd0, 0x0f,
	0xa4, 0x48, 0x84, 0x8b, 0xb5, 0xb8, 0x24, 0xb1, 0x24, 0x55, 0x82, 0x19, 0x28, 0xc9, 0x11, 0x04,
	0xe1, 0x28, 0x89, 0x70, 0x09, 0xf9, 0x64, 0x16, 0x97, 0x40, 0x0c, 0x2f, 0x0e, 0x4a, 0x2d, 0x2c,
	0x4d, 0x2d, 0x2e, 0x51, 0x72, 0xe0, 0x12, 0x46, 0x11, 0x2d, 0x2e, 0xc8, 0xcf, 0x2b, 0x4e, 0x15,
	0xd2, 0xe4, 0x62, 0x4b, 0x01, 0x0b, 0x01, 0x6d, 0x67, 0xd6, 0xe0, 0x36, 0x12, 0xd4, 0x43, 0x38,
	0x19, 0xa2, 0x36, 0x08, 0xaa, 0x40, 0x49, 0x8d, 0x4b, 0xc0, 0x3d, 0x15, 0x6a, 0x00, 0xd4, 0x54,
	0x6c, 0x4e, 0x07, 0xd9, 0x14, 0x5a, 0x90, 0x02, 0x74, 0x09, 0xaa, 0x52, 0x64, 0x9b, 0x18, 0xf1,
	0xda, 0x64, 0x74, 0x9f, 0x91, 0x8b, 0xd3, 0x11, 0x26, 0x29, 0xe4, 0xc7, 0xc5, 0x8d, 0xe4, 0x72,
	0x21, 0x59, 0x24, 0x7d, 0x98, 0xfe, 0x94, 0x92, 0xc3, 0x25, 0x0d, 0xf1, 0xb0, 0x12, 0x83, 0x90,
	0x2d, 0x17, 0x27, 0xdc, 0x1f, 0x42, 0xd2, 0x48, 0xca, 0xd1, 0x7d, 0x27, 0x85, 0xe9, 0x44, 0xa0,
	0x76, 0x67, 0x2e, 0x1e, 0x64, 0xef, 0x09, 0x21, 0x5b, 0x88, 0xc5, 0xdf, 0x58, 0x0d, 0x49, 0x62,
	0x03, 0x27, 0x07, 0x63, 0x40, 0x00, 0x00, 0x00, 0xff, 0xff, 0x1a, 0x36, 0x21, 0x5d, 0x21, 0x02,
	0x00, 0x00,
}
<commit_msg>Update protoc generated code.<commit_after>// Code generated by protoc-gen-go.
// source: apartment.proto
// DO NOT EDIT!

/*
Package apartment is a generated protocol buffer package.

It is generated from these files:
	apartment.proto

It has these top-level messages:
	Device
	ListDevicesRequest
	ListDevicesResponse
	GetDeviceRequest
	UpdateDeviceRequest
*/
package apartment

import proto "github.com/golang/protobuf/proto"
import fmt "fmt"
import math "math"

import (
	context "golang.org/x/net/context"
	grpc "google.golang.org/grpc"
)

// Reference imports to suppress errors if they are not otherwise used.
var _ = proto.Marshal
var _ = fmt.Errorf
var _ = math.Inf

// This is a compile-time assertion to ensure that this generated file
// is compatible with the proto package it is being compiled against.
// A compilation error at this line likely means your copy of the
// proto package needs to be updated.
const _ = proto.ProtoPackageIsVersion2 // please upgrade the proto package

type Device struct {
	Name         string `protobuf:"bytes,1,opt,name=name" json:"name,omitempty"`
	FriendlyName string `protobuf:"bytes,2,opt,name=friendly_name,json=friendlyName" json:"friendly_name,omitempty"`
	State        bool   `protobuf:"varint,3,opt,name=state" json:"state,omitempty"`
}

func (m *Device) Reset()                    { *m = Device{} }
func (m *Device) String() string            { return proto.CompactTextString(m) }
func (*Device) ProtoMessage()               {}
func (*Device) Descriptor() ([]byte, []int) { return fileDescriptor0, []int{0} }

func (m *Device) GetName() string {
	if m != nil {
		return m.Name
	}
	return ""
}

func (m *Device) GetFriendlyName() string {
	if m != nil {
		return m.FriendlyName
	}
	return ""
}

func (m *Device) GetState() bool {
	if m != nil {
		return m.State
	}
	return false
}

type ListDevicesRequest struct {
}

func (m *ListDevicesRequest) Reset()                    { *m = ListDevicesRequest{} }
func (m *ListDevicesRequest) String() string            { return proto.CompactTextString(m) }
func (*ListDevicesRequest) ProtoMessage()               {}
func (*ListDevicesRequest) Descriptor() ([]byte, []int) { return fileDescriptor0, []int{1} }

type ListDevicesResponse struct {
	Device []*Device `protobuf:"bytes,1,rep,name=device" json:"device,omitempty"`
}

func (m *ListDevicesResponse) Reset()                    { *m = ListDevicesResponse{} }
func (m *ListDevicesResponse) String() string            { return proto.CompactTextString(m) }
func (*ListDevicesResponse) ProtoMessage()               {}
func (*ListDevicesResponse) Descriptor() ([]byte, []int) { return fileDescriptor0, []int{2} }

func (m *ListDevicesResponse) GetDevice() []*Device {
	if m != nil {
		return m.Device
	}
	return nil
}

type GetDeviceRequest struct {
	Name string `protobuf:"bytes,1,opt,name=name" json:"name,omitempty"`
}

func (m *GetDeviceRequest) Reset()                    { *m = GetDeviceRequest{} }
func (m *GetDeviceRequest) String() string            { return proto.CompactTextString(m) }
func (*GetDeviceRequest) ProtoMessage()               {}
func (*GetDeviceRequest) Descriptor() ([]byte, []int) { return fileDescriptor0, []int{3} }

func (m *GetDeviceRequest) GetName() string {
	if m != nil {
		return m.Name
	}
	return ""
}

type UpdateDeviceRequest struct {
	Device *Device `protobuf:"bytes,1,opt,name=device" json:"device,omitempty"`
}

func (m *UpdateDeviceRequest) Reset()                    { *m = UpdateDeviceRequest{} }
func (m *UpdateDeviceRequest) String() string            { return proto.CompactTextString(m) }
func (*UpdateDeviceRequest) ProtoMessage()               {}
func (*UpdateDeviceRequest) Descriptor() ([]byte, []int) { return fileDescriptor0, []int{4} }

func (m *UpdateDeviceRequest) GetDevice() *Device {
	if m != nil {
		return m.Device
	}
	return nil
}

func init() {
	proto.RegisterType((*Device)(nil), "apartment.Device")
	proto.RegisterType((*ListDevicesRequest)(nil), "apartment.ListDevicesRequest")
	proto.RegisterType((*ListDevicesResponse)(nil), "apartment.ListDevicesResponse")
	proto.RegisterType((*GetDeviceRequest)(nil), "apartment.GetDeviceRequest")
	proto.RegisterType((*UpdateDeviceRequest)(nil), "apartment.UpdateDeviceRequest")
}

// Reference imports to suppress errors if they are not otherwise used.
var _ context.Context
var _ grpc.ClientConn

// This is a compile-time assertion to ensure that this generated file
// is compatible with the grpc package it is being compiled against.
const _ = grpc.SupportPackageIsVersion4

// Client API for Apartment service

type ApartmentClient interface {
	ListDevices(ctx context.Context, in *ListDevicesRequest, opts ...grpc.CallOption) (*ListDevicesResponse, error)
	GetDevice(ctx context.Context, in *GetDeviceRequest, opts ...grpc.CallOption) (*Device, error)
	UpdateDevice(ctx context.Context, in *UpdateDeviceRequest, opts ...grpc.CallOption) (*Device, error)
}

type apartmentClient struct {
	cc *grpc.ClientConn
}

func NewApartmentClient(cc *grpc.ClientConn) ApartmentClient {
	return &apartmentClient{cc}
}

func (c *apartmentClient) ListDevices(ctx context.Context, in *ListDevicesRequest, opts ...grpc.CallOption) (*ListDevicesResponse, error) {
	out := new(ListDevicesResponse)
	err := grpc.Invoke(ctx, "/apartment.Apartment/ListDevices", in, out, c.cc, opts...)
	if err != nil {
		return nil, err
	}
	return out, nil
}

func (c *apartmentClient) GetDevice(ctx context.Context, in *GetDeviceRequest, opts ...grpc.CallOption) (*Device, error) {
	out := new(Device)
	err := grpc.Invoke(ctx, "/apartment.Apartment/GetDevice", in, out, c.cc, opts...)
	if err != nil {
		return nil, err
	}
	return out, nil
}

func (c *apartmentClient) UpdateDevice(ctx context.Context, in *UpdateDeviceRequest, opts ...grpc.CallOption) (*Device, error) {
	out := new(Device)
	err := grpc.Invoke(ctx, "/apartment.Apartment/UpdateDevice", in, out, c.cc, opts...)
	if err != nil {
		return nil, err
	}
	return out, nil
}

// Server API for Apartment service

type ApartmentServer interface {
	ListDevices(context.Context, *ListDevicesRequest) (*ListDevicesResponse, error)
	GetDevice(context.Context, *GetDeviceRequest) (*Device, error)
	UpdateDevice(context.Context, *UpdateDeviceRequest) (*Device, error)
}

func RegisterApartmentServer(s *grpc.Server, srv ApartmentServer) {
	s.RegisterService(&_Apartment_serviceDesc, srv)
}

func _Apartment_ListDevices_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
	in := new(ListDevicesRequest)
	if err := dec(in); err != nil {
		return nil, err
	}
	if interceptor == nil {
		return srv.(ApartmentServer).ListDevices(ctx, in)
	}
	info := &grpc.UnaryServerInfo{
		Server:     srv,
		FullMethod: "/apartment.Apartment/ListDevices",
	}
	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
		return srv.(ApartmentServer).ListDevices(ctx, req.(*ListDevicesRequest))
	}
	return interceptor(ctx, in, info, handler)
}

func _Apartment_GetDevice_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
	in := new(GetDeviceRequest)
	if err := dec(in); err != nil {
		return nil, err
	}
	if interceptor == nil {
		return srv.(ApartmentServer).GetDevice(ctx, in)
	}
	info := &grpc.UnaryServerInfo{
		Server:     srv,
		FullMethod: "/apartment.Apartment/GetDevice",
	}
	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
		return srv.(ApartmentServer).GetDevice(ctx, req.(*GetDeviceRequest))
	}
	return interceptor(ctx, in, info, handler)
}

func _Apartment_UpdateDevice_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
	in := new(UpdateDeviceRequest)
	if err := dec(in); err != nil {
		return nil, err
	}
	if interceptor == nil {
		return srv.(ApartmentServer).UpdateDevice(ctx, in)
	}
	info := &grpc.UnaryServerInfo{
		Server:     srv,
		FullMethod: "/apartment.Apartment/UpdateDevice",
	}
	handler := func(ctx context.Context, req interface{}) (interface{}, error) {
		return srv.(ApartmentServer).UpdateDevice(ctx, req.(*UpdateDeviceRequest))
	}
	return interceptor(ctx, in, info, handler)
}

var _Apartment_serviceDesc = grpc.ServiceDesc{
	ServiceName: "apartment.Apartment",
	HandlerType: (*ApartmentServer)(nil),
	Methods: []grpc.MethodDesc{
		{
			MethodName: "ListDevices",
			Handler:    _Apartment_ListDevices_Handler,
		},
		{
			MethodName: "GetDevice",
			Handler:    _Apartment_GetDevice_Handler,
		},
		{
			MethodName: "UpdateDevice",
			Handler:    _Apartment_UpdateDevice_Handler,
		},
	},
	Streams:  []grpc.StreamDesc{},
	Metadata: "apartment.proto",
}

func init() { proto.RegisterFile("apartment.proto", fileDescriptor0) }

var fileDescriptor0 = []byte{
	// 261 bytes of a gzipped FileDescriptorProto
	0x1f, 0x8b, 0x08, 0x00, 0x00, 0x09, 0x6e, 0x88, 0x02, 0xff, 0xe2, 0xe2, 0x4f, 0x2c, 0x48, 0x2c,
	0x2a, 0xc9, 0x4d, 0xcd, 0x2b, 0xd1, 0x2b, 0x28, 0xca, 0x2f, 0xc9, 0x17, 0xe2, 0x84, 0x0b, 0x28,
	0x85, 0x73, 0xb1, 0xb9, 0xa4, 0x96, 0x65, 0x26, 0xa7, 0x0a, 0x09, 0x71, 0xb1, 0xe4, 0x25, 0xe6,
	0xa6, 0x4a, 0x30, 0x2a, 0x30, 0x6a, 0x70, 0x06, 0x81, 0xd9, 0x42, 0xca, 0x5c, 0xbc, 0x69, 0x45,
	0x99, 0xa9, 0x79, 0x29, 0x39, 0x95, 0xf1, 0x60, 0x49, 0x26, 0xb0, 0x24, 0x0f, 0x4c, 0xd0, 0x0f,
	0xa4, 0x48, 0x84, 0x8b, 0xb5, 0xb8, 0x24, 0xb1, 0x24, 0x55, 0x82, 0x59, 0x81, 0x51, 0x83, 0x23,
	0x08, 0xc2, 0x51, 0x12, 0xe1, 0x12, 0xf2, 0xc9, 0x2c, 0x2e, 0x81, 0x18, 0x5e, 0x1c, 0x94, 0x5a,
	0x58, 0x9a, 0x5a, 0x5c, 0xa2, 0xe4, 0xc0, 0x25, 0x8c, 0x22, 0x5a, 0x5c, 0x90, 0x9f, 0x57, 0x9c,
	0x2a, 0xa4, 0xc9, 0xc5, 0x96, 0x02, 0x16, 0x92, 0x60, 0x54, 0x60, 0xd6, 0xe0, 0x36, 0x12, 0xd4,
	0x43, 0x38, 0x19, 0xa2, 0x36, 0x08, 0xaa, 0x40, 0x49, 0x8d, 0x4b, 0xc0, 0x3d, 0x15, 0x6a, 0x00,
	0xd4, 0x54, 0x6c, 0x4e, 0x07, 0xd9, 0x14, 0x5a, 0x90, 0x92, 0x58, 0x92, 0x8a, 0xaa, 0x14, 0xd9,
	0x26, 0x46, 0xbc, 0x36, 0x19, 0xdd, 0x67, 0xe4, 0xe2, 0x74, 0x84, 0x49, 0x0a, 0xf9, 0x71, 0x71,
	0x23, 0xb9, 0x5c, 0x48, 0x16, 0x49, 0x1f, 0xa6, 0x3f, 0xa5, 0xe4, 0x70, 0x49, 0x43, 0x3c, 0xac,
	0xc4, 0x20, 0x64, 0xcb, 0xc5, 0x09, 0xf7, 0x87, 0x90, 0x34, 0x92, 0x72, 0x74, 0xdf, 0x49, 0x61,
	0x3a, 0x51, 0x89, 0x41, 0xc8, 0x99, 0x8b, 0x07, 0xd9, 0x7b, 0x42, 0xc8, 0x16, 0x62, 0xf1, 0x37,
	0x56, 0x43, 0x92, 0xd8, 0xc0, 0xc9, 0xc1, 0x18, 0x10, 0x00, 0x00, 0xff, 0xff, 0x1a, 0x36, 0x21,
	0x5d, 0x21, 0x02, 0x00, 0x00,
}
<|endoftext|>


<commit_before>package main

import "fmt"

func main() {
  fmt.Println("Hello")
  fmt.Println("Hallo")
  fmt.Println("Hola")
  fmt.Println("End my suffering")
  fmt.Println("Heisann")
}
<commit_msg>Hei på fransk<commit_after>package main

import "fmt"

func main() {
  fmt.Println("Hello")
  fmt.Println("Hallo")
  fmt.Println("Hola")
  fmt.Println("End my suffering")
  fmt.Println("Heisann")
  fmt.Println("Bonjour")
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/ast"
	"go/format"
	"go/parser"
	"go/token"
	"os"
	"path/filepath"
	"strings"
)

type (
	options struct {
		InputFile     string
		OutputFile    string
		InterfaceName string
		StructName    string
		Package       string
	}

	generator struct {
		options
		content      *bytes.Buffer
		fileset      *token.FileSet
		err          error
		importsPaths map[string]string
		importsUsage map[string]bool
		pkg          string
	}
)

func main() {
	opts := processFlags()
	if opts == nil {
		os.Exit(1)
	}

	g := newGenerator(opts)

	generatedContent, err := g.generate()
	if err != nil {
		die(err)
	}

	f, err := os.Create(opts.OutputFile)
	if err != nil {
		die(err)
	}

	if _, err := f.Write(generatedContent); err != nil {
		die(err)
	}
}

func die(err error) {
	fmt.Fprintf(os.Stderr, "%v\n", err)
	os.Exit(1)
}

func processFlags() *options {
	var (
		input  = flag.String("f", "", "input file containing interface declaration")
		name   = flag.String("i", "", "interface name")
		output = flag.String("o", "", "destination file for interface implementation")
		pkg    = flag.String("p", "", "destination package name")
		sname  = flag.String("t", "", "target struct name, default: <interface name>Mock")
	)

	flag.Parse()

	if *pkg == "" || *input == "" || *output == "" || *name == "" || !strings.HasSuffix(*input, ".go") || !strings.HasSuffix(*output, ".go") {
		flag.Usage()
		return nil
	}

	if *sname == "" {
		*sname = *name + "Mock"
	}

	return &options{
		InputFile:     *input,
		OutputFile:    *output,
		InterfaceName: *name,
		Package:       *pkg,
		StructName:    *sname,
	}
}

func newGenerator(o *options) *generator {
	return &generator{
		options:      *o,
		content:      new(bytes.Buffer),
		importsPaths: map[string]string{},
		importsUsage: map[string]bool{},
	}
}

func (g *generator) generate() ([]byte, error) {
	g.fileset = token.NewFileSet()

	parsedFile, err := parser.ParseFile(g.fileset, g.InputFile, nil, 0)
	if err != nil {
		return nil, err
	}

	g.pkg = parsedFile.Name.Name

	inputFilePath, err := filepath.Abs(g.InputFile)
	if err != nil {
		return nil, err
	}

	outputFilePath, err := filepath.Abs(g.OutputFile)
	if err != nil {
		return nil, err
	}

	if g.pkg == "main" && inputFilePath != outputFilePath {
		return nil, fmt.Errorf("file %s is a part of the main package, output file should be in the same dir", g.InputFile)
	}

	g.importsPaths = processImports(parsedFile)

	//we have selector in input file that is equal to input file's package name
	if _, ok := g.importsPaths[g.pkg]; ok {
		g.pkg = "_root_pkg"
	}

	g.importsPaths[g.pkg] = getRelativePkgPath(inputFilePath)
	ast.Inspect(parsedFile, g.interfaceInspector)

	if g.err != nil {
		return nil, g.err
	}

	return g.getGeneratedContent()
}

func (g *generator) interfaceInspector(node ast.Node) bool {
	typeSpec, ok := node.(*ast.TypeSpec)
	if ok && typeSpec.Name.Name == g.InterfaceName {
		if interfaceType, ok := typeSpec.Type.(*ast.InterfaceType); !ok {
			g.err = fmt.Errorf("%s: type %s is not an interface", g.InputFile, g.InterfaceName)
		} else {
			g.processInterfaceDeclaration(interfaceType)
		}
	}

	return !ok
}

const (
	templateForHeader = `
		package {{package}}

		import "testing"
		import "sync"
		{{imports}}

		type {{struct}} struct {
			t *testing.T
			m *sync.Mutex

			{{methods}}

			{{counters}}
		}

		func New{{struct}}(t *testing.T) *{{struct}} {
			return &{{struct}}{t: t, m: &sync.Mutex{}}
		}
		`

	templateForMethod = `
		func (m *{{struct}}) {{method}}{{inParams}} {{outParams}} {
			m.m.Lock()
			m.{{method}}Counter += 1
			m.m.Unlock()

			if m.{{method}}Func == nil {
				m.t.Fatalf("Unexpected call to {{struct}}.{{method}}")
			}

			{{return}}m.{{method}}Func({{paramsList}})
		}
		`
	templateValidateHeader = `
		func (m *{{struct}}) ValidateCallCounters() {
	`

	templateValidateBody = `
		if m.{{method}}Func != nil && m.{{method}}Counter == 0 {
			m.t.Error("Expected call to {{struct}}.{{method}}")
		}

	`
	templateValidateFooter = `
		}
	`
)

func (g *generator) processTemplate(template string, vars map[string]string) {
	for k, v := range vars {
		template = strings.Replace(template, "{{"+k+"}}", v, -1)
	}

	fmt.Fprintf(g.content, template)
}

func (g *generator) processInterfaceDeclaration(it *ast.InterfaceType) error {
	methods := []string{}
	counters := []string{}
	names := []string{}

	for i, method := range it.Methods.List {
		methodName := it.Methods.List[i].Names[0].Name
		names = append(names, methodName)

		signature := fmt.Sprintf("  %sFunc %s", methodName, g.getMethodSignature(method.Type))

		methods = append(methods, signature)
		counters = append(counters, methodName+"Counter int")

		f := method.Type.(*ast.FuncType)

		g.processUsedImports(f.Params)
		g.processUsedImports(f.Results)
	}

	vars := map[string]string{
		"struct":   g.StructName,
		"package":  g.Package,
		"methods":  strings.Join(methods, "\n"),
		"counters": strings.Join(counters, "\n"),
		"imports":  g.getImports(),
	}

	g.processTemplate(templateForHeader, vars)

	for _, method := range it.Methods.List {
		f := method.Type.(*ast.FuncType)

		vars["method"] = method.Names[0].Name
		vars["inParams"] = g.getInParams(f)
		vars["outParams"] = g.getOutParams(f)

		vars["paramsList"] = strings.Join(g.getParamsList(f.Params.List), ", ")
		if f.Results == nil {
			vars["return"] = ""
		} else {
			vars["return"] = "return "
		}

		g.processTemplate(templateForMethod, vars)
	}

	g.processTemplate(templateValidateHeader, vars)
	for _, name := range names {
		vars["method"] = name
		g.processTemplate(templateValidateBody, vars)
	}
	g.processTemplate(templateValidateFooter, vars)

	return nil
}

func processImports(f *ast.File) map[string]string {
	result := map[string]string{}

	for _, imp := range f.Imports {
		selector := ""
		path := imp.Path.Value[1 : len(imp.Path.Value)-1]

		if imp.Name == nil {
			selector = filepath.Base(path)
		} else {
			selector = strings.SplitN(imp.Name.Name, ".", 1)[0]
		}

		result[selector] = path
	}

	return result
}

func (g *generator) getImports() string {
	imports := []string{}

	for selector, _ := range g.importsUsage {
		path := g.importsPaths[selector]
		importSpec := ""
		if path == selector || strings.HasSuffix(path, "/"+selector) {
			importSpec = fmt.Sprintf("import %q", path)
		} else {
			importSpec = fmt.Sprintf("import %s %q", selector, path)
		}

		imports = append(imports, importSpec)
	}

	return strings.Join(imports, "\n")
}

func (g *generator) getParamsList(fields []*ast.Field) []string {
	result := []string{}

	for i, f := range fields {
		name := ""
		if len(f.Names) == 0 {
			name = fmt.Sprintf("p%d_0%s", i, variadicSuffix(f))
			result = append(result, name)
		} else {
			for _, n := range f.Names {
				name = n.Name + variadicSuffix(f)
				result = append(result, name)
			}
		}
	}

	return result
}

func (g *generator) processUsedImports(fields *ast.FieldList) {
	if fields == nil {
		return
	}

	for _, f := range fields.List {
		if selector := g.getSelector(f.Type); selector != "" {
			g.importsUsage[selector] = true
		}
	}
}

func (g *generator) getInParams(f *ast.FuncType) string {
	params := []string{}

	for i, f := range f.Params.List {
		names := []string{}
		if len(f.Names) == 0 {
			names = append(names, fmt.Sprintf("p%d_0", i))
		} else {
			for _, n := range f.Names {
				names = append(names, n.Name)
			}
		}

		params = append(params, strings.Join(names, ", ")+" "+g.typeToString(f.Type))
	}

	return "(" + strings.Join(params, ", ") + ")"
}

func (g *generator) getOutParams(f *ast.FuncType) string {
	if f.Results == nil {
		return ""
	}

	params := []string{}

	for _, f := range f.Results.List {
		types := []string{}
		if len(f.Names) == 0 {
			types = append(types, variadicSuffix(f)+g.typeToString(f.Type))
		} else {
			for range f.Names {
				types = append(types, variadicSuffix(f)+g.typeToString(f.Type))
			}
		}

		params = append(params, strings.Join(types, ", "))
	}

	if len(params) < 2 {
		return strings.Join(params, ", ")
	}

	return "(" + strings.Join(params, ", ") + ")"
}

func (g *generator) typeToString(t ast.Expr) string {
	switch e := t.(type) {
	case *ast.StarExpr:
		return "*" + g.typeToString(e.X)
	case *ast.ArrayType:
		return "[]" + g.typeToString(e.Elt)
	case *ast.MapType:
		return "map[" + g.typeToString(e.Key) + "]" + g.typeToString(e.Value)
	case *ast.SelectorExpr:
		return g.getSelector(e) + "." + e.Sel.Name
	case *ast.InterfaceType:
		return "interface{}"
	case *ast.ChanType:
		switch e.Dir {
		case ast.SEND:
			return "chan <- " + g.typeToString(e.Value)
		case ast.RECV:
			return "<- chan " + g.typeToString(e.Value)
		default:
			return "chan " + g.typeToString(e.Value)
		}
	case *ast.Ellipsis:
		return "..." + g.typeToString(e.Elt)
	case *ast.Ident:
		name := e.Name
		if strings.ToUpper(name)[0] == name[0] {
			return g.pkg + "." + name
		} else {
			return name
		}
	}

	panic(fmt.Sprintf("Unexpected expression of type %T: %+v", t, t))
}

func (g *generator) getSelector(e ast.Expr) string {
	switch t := e.(type) {
	case *ast.SelectorExpr:
		return t.X.(*ast.Ident).Name
	case *ast.ArrayType:
		return g.getSelector(t.Elt)
	case *ast.StarExpr:
		return g.getSelector(t.X)
	case *ast.Ident:
		name := t.Name
		if strings.ToUpper(name)[0] == name[0] {
			return g.pkg
		}
	}

	return ""
}

func (g *generator) getGeneratedContent() ([]byte, error) {
	formatted, err := format.Source(g.content.Bytes())
	if err != nil {
		fmt.Printf(string(g.content.Bytes()))
		return nil, err
	}

	return formatted, nil
}

func (g *generator) getMethodSignature(t ast.Expr) string {
	return "func" + g.getInParams(t.(*ast.FuncType)) + " " + g.getOutParams(t.(*ast.FuncType))
}

func getRelativePkgPath(inputFilePath string) string {
	pkgDir := filepath.Dir(inputFilePath)

	paths := strings.Split(os.Getenv("GOPATH"), ":")
	for _, p := range paths {
		prefix := p + "/src/"
		if strings.HasPrefix(pkgDir, prefix) {
			return strings.Replace(pkgDir, prefix, "", 1)
		}
	}

	return pkgDir
}

func variadicSuffix(f *ast.Field) string {
	if _, ok := f.Type.(*ast.Ellipsis); ok {
		return "..."
	}

	return ""
}
<commit_msg>Grouping all imports into one<commit_after>package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/ast"
	"go/format"
	"go/parser"
	"go/token"
	"os"
	"path/filepath"
	"strings"
)

type (
	options struct {
		InputFile     string
		OutputFile    string
		InterfaceName string
		StructName    string
		Package       string
	}

	generator struct {
		options
		content      *bytes.Buffer
		fileset      *token.FileSet
		err          error
		importsPaths map[string]string
		importsUsage map[string]bool
		pkg          string
	}
)

func main() {
	opts := processFlags()
	if opts == nil {
		os.Exit(1)
	}

	g := newGenerator(opts)

	generatedContent, err := g.generate()
	if err != nil {
		die(err)
	}

	f, err := os.Create(opts.OutputFile)
	if err != nil {
		die(err)
	}

	if _, err := f.Write(generatedContent); err != nil {
		die(err)
	}
}

func die(err error) {
	fmt.Fprintf(os.Stderr, "%v\n", err)
	os.Exit(1)
}

func processFlags() *options {
	var (
		input  = flag.String("f", "", "input file containing interface declaration")
		name   = flag.String("i", "", "interface name")
		output = flag.String("o", "", "destination file for interface implementation")
		pkg    = flag.String("p", "", "destination package name")
		sname  = flag.String("t", "", "target struct name, default: <interface name>Mock")
	)

	flag.Parse()

	if *pkg == "" || *input == "" || *output == "" || *name == "" || !strings.HasSuffix(*input, ".go") || !strings.HasSuffix(*output, ".go") {
		flag.Usage()
		return nil
	}

	if *sname == "" {
		*sname = *name + "Mock"
	}

	return &options{
		InputFile:     *input,
		OutputFile:    *output,
		InterfaceName: *name,
		Package:       *pkg,
		StructName:    *sname,
	}
}

func newGenerator(o *options) *generator {
	return &generator{
		options:      *o,
		content:      new(bytes.Buffer),
		importsPaths: map[string]string{},
		importsUsage: map[string]bool{},
	}
}

func (g *generator) generate() ([]byte, error) {
	g.fileset = token.NewFileSet()

	parsedFile, err := parser.ParseFile(g.fileset, g.InputFile, nil, 0)
	if err != nil {
		return nil, err
	}

	g.pkg = parsedFile.Name.Name

	inputFilePath, err := filepath.Abs(g.InputFile)
	if err != nil {
		return nil, err
	}

	outputFilePath, err := filepath.Abs(g.OutputFile)
	if err != nil {
		return nil, err
	}

	if g.pkg == "main" && inputFilePath != outputFilePath {
		return nil, fmt.Errorf("file %s is a part of the main package, output file should be in the same dir", g.InputFile)
	}

	g.importsPaths = processImports(parsedFile)

	//we have selector in input file that is equal to input file's package name
	if _, ok := g.importsPaths[g.pkg]; ok {
		g.pkg = "_root_pkg"
	}

	g.importsPaths[g.pkg] = getRelativePkgPath(inputFilePath)
	ast.Inspect(parsedFile, g.interfaceInspector)

	if g.err != nil {
		return nil, g.err
	}

	return g.getGeneratedContent()
}

func (g *generator) interfaceInspector(node ast.Node) bool {
	typeSpec, ok := node.(*ast.TypeSpec)
	if ok && typeSpec.Name.Name == g.InterfaceName {
		if interfaceType, ok := typeSpec.Type.(*ast.InterfaceType); !ok {
			g.err = fmt.Errorf("%s: type %s is not an interface", g.InputFile, g.InterfaceName)
		} else {
			g.processInterfaceDeclaration(interfaceType)
		}
	}

	return !ok
}

const (
	templateForHeader = `
		package {{package}}

		import (
			"testing"
			"sync"

			{{imports}}
	    )

		type {{struct}} struct {
			t *testing.T
			m *sync.Mutex

			{{methods}}

			{{counters}}
		}

		func New{{struct}}(t *testing.T) *{{struct}} {
			return &{{struct}}{t: t, m: &sync.Mutex{}}
		}
		`

	templateForMethod = `
		func (m *{{struct}}) {{method}}{{inParams}} {{outParams}} {
			m.m.Lock()
			m.{{method}}Counter += 1
			m.m.Unlock()

			if m.{{method}}Func == nil {
				m.t.Fatalf("Unexpected call to {{struct}}.{{method}}")
			}

			{{return}}m.{{method}}Func({{paramsList}})
		}
		`
	templateValidateHeader = `
		func (m *{{struct}}) ValidateCallCounters() {
	`

	templateValidateBody = `
		if m.{{method}}Func != nil && m.{{method}}Counter == 0 {
			m.t.Error("Expected call to {{struct}}.{{method}}")
		}

	`
	templateValidateFooter = `
		}
	`
)

func (g *generator) processTemplate(template string, vars map[string]string) {
	for k, v := range vars {
		template = strings.Replace(template, "{{"+k+"}}", v, -1)
	}

	fmt.Fprintf(g.content, template)
}

func (g *generator) processInterfaceDeclaration(it *ast.InterfaceType) error {
	methods := []string{}
	counters := []string{}
	names := []string{}

	for i, method := range it.Methods.List {
		methodName := it.Methods.List[i].Names[0].Name
		names = append(names, methodName)

		signature := fmt.Sprintf("  %sFunc %s", methodName, g.getMethodSignature(method.Type))

		methods = append(methods, signature)
		counters = append(counters, methodName+"Counter int")

		f := method.Type.(*ast.FuncType)

		g.processUsedImports(f.Params)
		g.processUsedImports(f.Results)
	}

	vars := map[string]string{
		"struct":   g.StructName,
		"package":  g.Package,
		"methods":  strings.Join(methods, "\n"),
		"counters": strings.Join(counters, "\n"),
		"imports":  g.getImports(),
	}

	g.processTemplate(templateForHeader, vars)

	for _, method := range it.Methods.List {
		f := method.Type.(*ast.FuncType)

		vars["method"] = method.Names[0].Name
		vars["inParams"] = g.getInParams(f)
		vars["outParams"] = g.getOutParams(f)

		vars["paramsList"] = strings.Join(g.getParamsList(f.Params.List), ", ")
		if f.Results == nil {
			vars["return"] = ""
		} else {
			vars["return"] = "return "
		}

		g.processTemplate(templateForMethod, vars)
	}

	g.processTemplate(templateValidateHeader, vars)
	for _, name := range names {
		vars["method"] = name
		g.processTemplate(templateValidateBody, vars)
	}
	g.processTemplate(templateValidateFooter, vars)

	return nil
}

func processImports(f *ast.File) map[string]string {
	result := map[string]string{}

	for _, imp := range f.Imports {
		selector := ""
		path := imp.Path.Value[1 : len(imp.Path.Value)-1]

		if imp.Name == nil {
			selector = filepath.Base(path)
		} else {
			selector = strings.SplitN(imp.Name.Name, ".", 1)[0]
		}

		result[selector] = path
	}

	return result
}

func (g *generator) getImports() string {
	imports := []string{}

	for selector, _ := range g.importsUsage {
		path := g.importsPaths[selector]
		importSpec := ""
		if path == selector || strings.HasSuffix(path, "/"+selector) {
			importSpec = fmt.Sprintf("%q", path)
		} else {
			importSpec = fmt.Sprintf("%s %q", selector, path)
		}

		imports = append(imports, importSpec)
	}

	return strings.Join(imports, "\n")
}

func (g *generator) getParamsList(fields []*ast.Field) []string {
	result := []string{}

	for i, f := range fields {
		name := ""
		if len(f.Names) == 0 {
			name = fmt.Sprintf("p%d_0%s", i, variadicSuffix(f))
			result = append(result, name)
		} else {
			for _, n := range f.Names {
				name = n.Name + variadicSuffix(f)
				result = append(result, name)
			}
		}
	}

	return result
}

func (g *generator) processUsedImports(fields *ast.FieldList) {
	if fields == nil {
		return
	}

	for _, f := range fields.List {
		if selector := g.getSelector(f.Type); selector != "" {
			g.importsUsage[selector] = true
		}
	}
}

func (g *generator) getInParams(f *ast.FuncType) string {
	params := []string{}

	for i, f := range f.Params.List {
		names := []string{}
		if len(f.Names) == 0 {
			names = append(names, fmt.Sprintf("p%d_0", i))
		} else {
			for _, n := range f.Names {
				names = append(names, n.Name)
			}
		}

		params = append(params, strings.Join(names, ", ")+" "+g.typeToString(f.Type))
	}

	return "(" + strings.Join(params, ", ") + ")"
}

func (g *generator) getOutParams(f *ast.FuncType) string {
	if f.Results == nil {
		return ""
	}

	params := []string{}

	for _, f := range f.Results.List {
		types := []string{}
		if len(f.Names) == 0 {
			types = append(types, variadicSuffix(f)+g.typeToString(f.Type))
		} else {
			for range f.Names {
				types = append(types, variadicSuffix(f)+g.typeToString(f.Type))
			}
		}

		params = append(params, strings.Join(types, ", "))
	}

	if len(params) < 2 {
		return strings.Join(params, ", ")
	}

	return "(" + strings.Join(params, ", ") + ")"
}

func (g *generator) typeToString(t ast.Expr) string {
	switch e := t.(type) {
	case *ast.StarExpr:
		return "*" + g.typeToString(e.X)
	case *ast.ArrayType:
		return "[]" + g.typeToString(e.Elt)
	case *ast.MapType:
		return "map[" + g.typeToString(e.Key) + "]" + g.typeToString(e.Value)
	case *ast.SelectorExpr:
		return g.getSelector(e) + "." + e.Sel.Name
	case *ast.InterfaceType:
		return "interface{}"
	case *ast.ChanType:
		switch e.Dir {
		case ast.SEND:
			return "chan <- " + g.typeToString(e.Value)
		case ast.RECV:
			return "<- chan " + g.typeToString(e.Value)
		default:
			return "chan " + g.typeToString(e.Value)
		}
	case *ast.Ellipsis:
		return "..." + g.typeToString(e.Elt)
	case *ast.Ident:
		name := e.Name
		if strings.ToUpper(name)[0] == name[0] {
			return g.pkg + "." + name
		} else {
			return name
		}
	}

	panic(fmt.Sprintf("Unexpected expression of type %T: %+v", t, t))
}

func (g *generator) getSelector(e ast.Expr) string {
	switch t := e.(type) {
	case *ast.SelectorExpr:
		return t.X.(*ast.Ident).Name
	case *ast.ArrayType:
		return g.getSelector(t.Elt)
	case *ast.StarExpr:
		return g.getSelector(t.X)
	case *ast.Ident:
		name := t.Name
		if strings.ToUpper(name)[0] == name[0] {
			return g.pkg
		}
	}

	return ""
}

func (g *generator) getGeneratedContent() ([]byte, error) {
	formatted, err := format.Source(g.content.Bytes())
	if err != nil {
		fmt.Printf(string(g.content.Bytes()))
		return nil, err
	}

	return formatted, nil
}

func (g *generator) getMethodSignature(t ast.Expr) string {
	return "func" + g.getInParams(t.(*ast.FuncType)) + " " + g.getOutParams(t.(*ast.FuncType))
}

func getRelativePkgPath(inputFilePath string) string {
	pkgDir := filepath.Dir(inputFilePath)

	paths := strings.Split(os.Getenv("GOPATH"), ":")
	for _, p := range paths {
		prefix := p + "/src/"
		if strings.HasPrefix(pkgDir, prefix) {
			return strings.Replace(pkgDir, prefix, "", 1)
		}
	}

	return pkgDir
}

func variadicSuffix(f *ast.Field) string {
	if _, ok := f.Type.(*ast.Ellipsis); ok {
		return "..."
	}

	return ""
}
<|endoftext|>