<commit_before>package anime

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"

	"github.com/PuerkitoBio/goquery"
	"github.com/sirupsen/logrus"
)

//TimeLine defines bilibili's timeline
type TimeLine struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
	Result  []struct {
		Date      string `json:"date"`
		DateTs    int    `json:"date_ts"`
		DayOfWeek int    `json:"day_of_week"`
		IsToday   int    `json:"is_today"`
		Seasons   []struct {
			Cover        string `json:"cover"`
			Delay        int    `json:"delay"`
			EpID         int    `json:"ep_id"`
			Favorites    int    `json:"favorites"`
			Follow       int    `json:"follow"`
			IsPublished  int    `json:"is_published"`
			PubIndex     string `json:"pub_index"`
			PubTime      string `json:"pub_time"`
			PubTs        int    `json:"pub_ts"`
			SeasonID     int    `json:"season_id"`
			SeasonStatus int    `json:"season_status"`
			SquareCover  string `json:"square_cover"`
			Title        string `json:"title"`
			Badge        string `json:"badge,omitempty"`
		} `json:"seasons"`
	} `json:"result"`
}

//TimeLineCN defines 国创 timeline
type TimeLineCN struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
	Result  []struct {
		Date      string `json:"date"`
		DateTs    int    `json:"date_ts"`
		DayOfWeek int    `json:"day_of_week"`
		IsToday   int    `json:"is_today"`
		Seasons   []struct {
			Cover        string `json:"cover"`
			Delay        int    `json:"delay"`
			EpID         int    `json:"ep_id"`
			Favorites    int    `json:"favorites"`
			Follow       int    `json:"follow"`
			IsPublished  int    `json:"is_published"`
			PubIndex     string `json:"pub_index"`
			PubTime      string `json:"pub_time"`
			PubTs        int    `json:"pub_ts"`
			SeasonID     int    `json:"season_id"`
			SeasonStatus int    `json:"season_status"`
			SquareCover  string `json:"square_cover"`
			Title        string `json:"title"`
		} `json:"seasons"`
	} `json:"result"`
}

//SrcObj defines bangumi obj
type SrcObj struct {
	Src         string
	BangumiName string
	Link        *url.URL
	Pubed       bool
}

const (
	//BilibiliGC B站国创
	BilibiliGC = "https://bangumi.bilibili.com/web_api/timeline_cn"
	//BilibiliJP B站日漫
	BilibiliJP = "https://bangumi.bilibili.com/web_api/timeline_global"
	//Dilidili D站动漫
	Dilidili = "http://www.dilidili.wang"
)

//FormatLinkInMarkdownPreview formats srcobj to Markdown view
func (s *SrcObj) FormatLinkInMarkdownPreview() string {
	name := fmt.Sprintf("[%s From %s]", s.BangumiName, s.Src)
	if s.Pubed {
		linkstr := fmt.Sprintf("(%s)", s.Link.String())
		return fmt.Sprintf("%s%s", name, linkstr)
	}
	return fmt.Sprintf("%s From %s(未更新)", s.BangumiName, s.Src)
}

//GetAllAnimes gets all animes from all src defined.
func GetAllAnimes() (objs []*SrcObj, err error) {
	objs, err = GetAnimeFromBGC()
	if err != nil {
		return
	}
	bbjp, err := GetAnimeFromBJP()
	if err != nil {
		return nil, err
	}
	objs = append(objs, bbjp...)
	d, err := GetAnimeFromD()
	if err != nil {
		return nil, err
	}
	objs = append(objs, d...)
	return
}

// GetAnimeFromB get anime from bilibili
func GetAnimeFromB() (objs []*SrcObj, err error) {
	objs, err = GetAnimeFromBGC()
	if err != nil {
		return
	}
	bbjp, err := GetAnimeFromBJP()
	if err != nil {
		return nil, err
	}
	objs = append(objs, bbjp...)
	return
}

//GetAnimeFromBGC ....
func GetAnimeFromBGC() ([]*SrcObj, error) {
	bgcSrc, err := url.Parse(BilibiliGC)
	if err != nil {
		return nil, err
	}
	objs, err := scrapeBilibiliTimeline(bgcSrc)
	if err != nil {
		return nil, err
	}

	return objs, nil
}

//GetAnimeFromBJP ...
func GetAnimeFromBJP() ([]*SrcObj, error) {
	bjpSrc, err := url.Parse(BilibiliJP)
	if err != nil {
		return nil, err
	}
	objs, err := scrapeBilibiliTimeline(bjpSrc)
	if err != nil {
		return nil, err
	}
	return objs, nil
}

//GetAnimeFromD ...
func GetAnimeFromD() ([]*SrcObj, error) {
	diliURL, err := url.Parse(Dilidili)
	if err != nil {
		return nil, err
	}
	objs, err := scrapeDilidiliTimeLine(diliURL)
	if err != nil {
		return nil, err
	}
	return objs, nil
}

func formatLink(rawLink string) (resLink *url.URL, err error) {
	var resURL string
	if strings.HasPrefix(rawLink, "//") {
		resURL = fmt.Sprintf("https:%s", rawLink)
	} else {
		resURL = rawLink
	}
	resLink, err = url.Parse(resURL)
	if err != nil {
		return
	}
	return
}

func formatNotAbsoluteLink(rawlink string, src string) (resLink *url.URL, err error) {
	resURL := fmt.Sprintf("%s%s", src, rawlink)
	resLink, err = url.Parse(resURL)
	return
}

func scrapeBilibiliTimeline(src *url.URL) ([]*SrcObj, error) {
	req, err := http.Get(src.String())
	if err != nil {
		return nil, err
	}
	var objs []*SrcObj
	body, err := ioutil.ReadAll(req.Body)
	if err != nil {
		return nil, err
	}

	tl := new(TimeLine)
	err = json.Unmarshal(body, tl)
	if err != nil {
		return nil, err
	}
	for _, r := range tl.Result {
		if r.IsToday == 0 {
			continue
		}
		for _, s := range r.Seasons {
			if s.Delay == 1 {
				continue
			}
			obj := new(SrcObj)
			obj.BangumiName = s.Title
			obj.Link, err = formatNotAbsoluteLink(strconv.Itoa(s.EpID),
				"https://www.bilibili.com/bangumi/play/ep")
			if err != nil {
				return nil, err
			}
			obj.Pubed = s.IsPublished > 0
			obj.Src = "bilibili"
			objs = append(objs, obj)
		}
	}

	return objs, nil
}

func scrapeDilidiliTimeLine(src *url.URL) ([]*SrcObj, error) {
	doc, err := goquery.NewDocument(src.String())
	if err != nil {
		return nil, err
	}
	var objs []*SrcObj
	location, err := time.LoadLocation("Asia/Shanghai")
	if err != nil {
		return nil, err
	}

	today := convert2CNWeekDay(int(time.Now().In(location).Weekday()), 6)
	// log.Println(doc.Find(".container-row-1").Find(".two-auto").Find("ul").Find(''))

	doc.Find(".change").Eq(1).Find(".sldr").Find(".wrp > li").Each(func(index int, s *goquery.Selection) {
		if index == today {
			s.Find(".list > li").Each(func(cindex int, cs *goquery.Selection) {
				ele := cs.Find("a")
				obj := new(SrcObj)
				obj.Src = "dilidili"
				obj.BangumiName = ele.Text()
				if ele.Length() > 1 {
					obj.Pubed = true
					link, _ := ele.Eq(1).Attr("href")
					var err error
					obj.Link, err = formatNotAbsoluteLink(link, Dilidili)
					if err != nil {
						logrus.Error("format dilidili url error:%s", err.Error())
					}
				}
				if obj.Link == nil {
					obj.Link, err = url.Parse(Dilidili)
					if err != nil {
						logrus.Error(err)
					}
				}
				objs = append(objs, obj)
			})
		}
	})
	return objs, nil
}

// weekday internationalWeekday    cnWeekday
// Sun     0      6
// Mon     1      0
// Tue     2      1
// Wed     3      2
// Thu     4      3
// Fri     5      4
// Sat     6      5
func convert2CNWeekDay(internationWeekday int, offsetDay int) (cnWeekday int) {
	return (internationWeekday + offsetDay) % 7
}
<commit_msg>plugin: fix dilidili anime link<commit_after>package anime

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"

	"github.com/PuerkitoBio/goquery"
	"github.com/sirupsen/logrus"
)

//TimeLine defines bilibili's timeline
type TimeLine struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
	Result  []struct {
		Date      string `json:"date"`
		DateTs    int    `json:"date_ts"`
		DayOfWeek int    `json:"day_of_week"`
		IsToday   int    `json:"is_today"`
		Seasons   []struct {
			Cover        string `json:"cover"`
			Delay        int    `json:"delay"`
			EpID         int    `json:"ep_id"`
			Favorites    int    `json:"favorites"`
			Follow       int    `json:"follow"`
			IsPublished  int    `json:"is_published"`
			PubIndex     string `json:"pub_index"`
			PubTime      string `json:"pub_time"`
			PubTs        int    `json:"pub_ts"`
			SeasonID     int    `json:"season_id"`
			SeasonStatus int    `json:"season_status"`
			SquareCover  string `json:"square_cover"`
			Title        string `json:"title"`
			Badge        string `json:"badge,omitempty"`
		} `json:"seasons"`
	} `json:"result"`
}

//TimeLineCN defines 国创 timeline
type TimeLineCN struct {
	Code    int    `json:"code"`
	Message string `json:"message"`
	Result  []struct {
		Date      string `json:"date"`
		DateTs    int    `json:"date_ts"`
		DayOfWeek int    `json:"day_of_week"`
		IsToday   int    `json:"is_today"`
		Seasons   []struct {
			Cover        string `json:"cover"`
			Delay        int    `json:"delay"`
			EpID         int    `json:"ep_id"`
			Favorites    int    `json:"favorites"`
			Follow       int    `json:"follow"`
			IsPublished  int    `json:"is_published"`
			PubIndex     string `json:"pub_index"`
			PubTime      string `json:"pub_time"`
			PubTs        int    `json:"pub_ts"`
			SeasonID     int    `json:"season_id"`
			SeasonStatus int    `json:"season_status"`
			SquareCover  string `json:"square_cover"`
			Title        string `json:"title"`
		} `json:"seasons"`
	} `json:"result"`
}

//SrcObj defines bangumi obj
type SrcObj struct {
	Src         string
	BangumiName string
	Link        *url.URL
	Pubed       bool
}

const (
	//BilibiliGC B站国创
	BilibiliGC = "https://bangumi.bilibili.com/web_api/timeline_cn"
	//BilibiliJP B站日漫
	BilibiliJP = "https://bangumi.bilibili.com/web_api/timeline_global"
	//Dilidili D站动漫
	Dilidili = "http://www.dilidili.wang"
)

//FormatLinkInMarkdownPreview formats srcobj to Markdown view
func (s *SrcObj) FormatLinkInMarkdownPreview() string {
	name := fmt.Sprintf("[%s From %s]", s.BangumiName, s.Src)
	if s.Pubed {
		linkstr := fmt.Sprintf("(%s)", s.Link.String())
		return fmt.Sprintf("%s%s", name, linkstr)
	}
	return fmt.Sprintf("%s From %s(未更新)", s.BangumiName, s.Src)
}

//GetAllAnimes gets all animes from all src defined.
func GetAllAnimes() (objs []*SrcObj, err error) {
	objs, err = GetAnimeFromBGC()
	if err != nil {
		return
	}
	bbjp, err := GetAnimeFromBJP()
	if err != nil {
		return nil, err
	}
	objs = append(objs, bbjp...)
	d, err := GetAnimeFromD()
	if err != nil {
		return nil, err
	}
	objs = append(objs, d...)
	return
}

// GetAnimeFromB get anime from bilibili
func GetAnimeFromB() (objs []*SrcObj, err error) {
	objs, err = GetAnimeFromBGC()
	if err != nil {
		return
	}
	bbjp, err := GetAnimeFromBJP()
	if err != nil {
		return nil, err
	}
	objs = append(objs, bbjp...)
	return
}

//GetAnimeFromBGC ....
func GetAnimeFromBGC() ([]*SrcObj, error) {
	bgcSrc, err := url.Parse(BilibiliGC)
	if err != nil {
		return nil, err
	}
	objs, err := scrapeBilibiliTimeline(bgcSrc)
	if err != nil {
		return nil, err
	}

	return objs, nil
}

//GetAnimeFromBJP ...
func GetAnimeFromBJP() ([]*SrcObj, error) {
	bjpSrc, err := url.Parse(BilibiliJP)
	if err != nil {
		return nil, err
	}
	objs, err := scrapeBilibiliTimeline(bjpSrc)
	if err != nil {
		return nil, err
	}
	return objs, nil
}

//GetAnimeFromD ...
func GetAnimeFromD() ([]*SrcObj, error) {
	diliURL, err := url.Parse(Dilidili)
	if err != nil {
		return nil, err
	}
	objs, err := scrapeDilidiliTimeLine(diliURL)
	if err != nil {
		return nil, err
	}
	return objs, nil
}

func formatLink(rawLink string) (resLink *url.URL, err error) {
	var resURL string
	if strings.HasPrefix(rawLink, "//") {
		resURL = fmt.Sprintf("https:%s", rawLink)
	} else {
		resURL = rawLink
	}
	resLink, err = url.Parse(resURL)
	if err != nil {
		return
	}
	return
}

func formatNotAbsoluteLink(uri string, src string) (resLink *url.URL, err error) {
	if strings.Contains(uri, "://") {
		resLink, err = url.Parse(uri)
		return
	}
	resURL := fmt.Sprintf("%s%s", src, uri)
	resLink, err = url.Parse(resURL)
	return
}

func scrapeBilibiliTimeline(src *url.URL) ([]*SrcObj, error) {
	req, err := http.Get(src.String())
	if err != nil {
		return nil, err
	}
	var objs []*SrcObj
	body, err := ioutil.ReadAll(req.Body)
	if err != nil {
		return nil, err
	}

	tl := new(TimeLine)
	err = json.Unmarshal(body, tl)
	if err != nil {
		return nil, err
	}
	for _, r := range tl.Result {
		if r.IsToday == 0 {
			continue
		}
		for _, s := range r.Seasons {
			if s.Delay == 1 {
				continue
			}
			obj := new(SrcObj)
			obj.BangumiName = s.Title
			obj.Link, err = formatNotAbsoluteLink(strconv.Itoa(s.EpID),
				"https://www.bilibili.com/bangumi/play/ep")
			if err != nil {
				return nil, err
			}
			obj.Pubed = s.IsPublished > 0
			obj.Src = "bilibili"
			objs = append(objs, obj)
		}
	}

	return objs, nil
}

func scrapeDilidiliTimeLine(src *url.URL) ([]*SrcObj, error) {
	doc, err := goquery.NewDocument(src.String())
	if err != nil {
		return nil, err
	}
	var objs []*SrcObj
	location, err := time.LoadLocation("Asia/Shanghai")
	if err != nil {
		return nil, err
	}

	today := convert2CNWeekDay(int(time.Now().In(location).Weekday()), 6)
	// log.Println(doc.Find(".container-row-1").Find(".two-auto").Find("ul").Find(''))

	doc.Find(".change").Eq(1).Find(".sldr").Find(".wrp > li").Each(func(index int, s *goquery.Selection) {
		if index == today {
			s.Find(".list > li").Each(func(cindex int, cs *goquery.Selection) {
				ele := cs.Find("a")
				obj := new(SrcObj)
				obj.Src = "dilidili"
				obj.BangumiName = ele.Text()
				if ele.Length() > 1 {
					obj.Pubed = true
					link, _ := ele.Eq(1).Attr("href")
					var err error
					obj.Link, err = formatNotAbsoluteLink(link, Dilidili)
					if err != nil {
						logrus.Error("format dilidili url error:%s", err.Error())
					}
				}
				if obj.Link == nil {
					obj.Link, err = url.Parse(Dilidili)
					if err != nil {
						logrus.Error(err)
					}
				}
				objs = append(objs, obj)
			})
		}
	})
	return objs, nil
}

// weekday internationalWeekday    cnWeekday
// Sun     0      6
// Mon     1      0
// Tue     2      1
// Wed     3      2
// Thu     4      3
// Fri     5      4
// Sat     6      5
func convert2CNWeekDay(internationWeekday int, offsetDay int) (cnWeekday int) {
	return (internationWeekday + offsetDay) % 7
}
<|endoftext|>


<commit_before>package dns

// A concurrent client implementation.

import (
	"io"
	"net"
	"time"
)

// Order of events:
// *client -> *reply -> Exchange*() -> dial()/send()->write()/receive()->read()

type reply struct {
	client         *Client
	addr           string
	req            *Msg
	conn           net.Conn
	tsigRequestMAC string
	tsigTimersOnly bool
	tsigStatus     error
	rtt            time.Duration
	t              time.Time
}

// An Exchange is returned on the channel when calling client.Do or client.DoRtt
type Exchange struct {

}

// A Client defines parameter for a DNS client. A nil
// Client is usable for sending queries.
type Client struct {
	Net          string            // if "tcp" a TCP query will be initiated, otherwise an UDP one (default is "" for UDP)
	Attempts     int               // number of attempts, if not set defaults to 1
	Retry        bool              // retry with TCP
	ReadTimeout  time.Duration     // the net.Conn.SetReadTimeout value for new connections (ns), defauls to 2 * 1e9
	WriteTimeout time.Duration     // the net.Conn.SetWriteTimeout value for new connections (ns), defauls to 2 * 1e9
	TsigSecret   map[string]string // secret(s) for Tsig map[<zonename>]<base64 secret>, zonename must be fully qualified
}

// Do performs an asynchronous query. The msg *Msg is the question to ask, the
// string addr is the address of the nameserver, the parameter data is used
// in the callback function. The call backback function is called with the
// original query, the answer returned from the nameserver an optional error and
// data.
// It calls Exchange.
func (c *Client) Do(msg *Msg, addr string, data interface{}, callback func(*Msg, *Msg, time.Duration, error, interface{})) {
	go func() {
		r, rtt, err := c.Exchange(msg, addr)
		callback(msg, r, rtt, err, data)
	}()
}

// Exchange performs an synchronous query. It sends the message m to the address
// contained in a and waits for an reply. Basic use pattern with a *Client:
//
//	c := new(dns.Client)
//	in, rtt, err := c.Exchange(message, "127.0.0.1:53")
//
func (c *Client) Exchange(m *Msg, a string) (r *Msg, rtt time.Duration, err error) {
	w := new(reply)
	w.client = c
	w.addr = a
	if err = w.dial(); err != nil {
		return nil, 0, err
	}
	if err = w.send(m); err != nil {
		return nil, 0, err
	}
	r, err = w.receive()
	return r, w.rtt, err
}

func (w *reply) RemoteAddr() net.Addr {
	if w.conn != nil {
		return w.conn.RemoteAddr()
	}
	return nil
}

// dial connects to the address addr for the network set in c.Net
func (w *reply) dial() (err error) {
	var conn net.Conn
	attempts := w.client.Attempts
	if attempts == 0 {
		attempts = 1
	}
	for a := 0; a < attempts; a++ {
		if w.client.Net == "" {
			conn, err = net.Dial("udp", w.addr)
		} else {
			conn, err = net.Dial(w.client.Net, w.addr)
		}
		if err != nil {
			// There are no timeouts defined?
			continue
		}
	}
	w.conn = conn
	return
}

func (w *reply) receive() (*Msg, error) {
	var p []byte
	m := new(Msg)
	switch w.client.Net {
	case "tcp", "tcp4", "tcp6":
		p = make([]byte, MaxMsgSize)
	case "", "udp", "udp4", "udp6":
		// OPT! TODO(mg)
		p = make([]byte, DefaultMsgSize)
	}
	n, err := w.read(p)
	if err != nil && n == 0 {
		return nil, err
	}
	p = p[:n]
	if err := m.Unpack(p); err != nil {
		return nil, err
	}
	w.rtt = time.Since(w.t)
	m.Size = n
	if t := m.IsTsig(); t != nil {
		secret := t.Hdr.Name
		if _, ok := w.client.TsigSecret[secret]; !ok {
			w.tsigStatus = ErrSecret
			return m, ErrSecret
		}
		// Need to work on the original message p, as that was used to calculate the tsig.
		w.tsigStatus = TsigVerify(p, w.client.TsigSecret[secret], w.tsigRequestMAC, w.tsigTimersOnly)
	}
	return m, w.tsigStatus
}

func (w *reply) read(p []byte) (n int, err error) {
	if w.conn == nil {
		return 0, ErrConnEmpty
	}
	if len(p) < 2 {
		return 0, io.ErrShortBuffer
	}
	attempts := w.client.Attempts
	if attempts == 0 {
		attempts = 1
	}
	switch w.client.Net {
	case "tcp", "tcp4", "tcp6":
		setTimeouts(w)
		for a := 0; a < attempts; a++ {
			n, err = w.conn.(*net.TCPConn).Read(p[0:2])
			if err != nil || n != 2 {
				if e, ok := err.(net.Error); ok && e.Timeout() {
					continue
				}
				return n, err
			}
			l, _ := unpackUint16(p[0:2], 0)
			if l == 0 {
				return 0, ErrShortRead
			}
			if int(l) > len(p) {
				return int(l), io.ErrShortBuffer
			}
			n, err = w.conn.(*net.TCPConn).Read(p[:l])
			if err != nil {
				if e, ok := err.(net.Error); ok && e.Timeout() {
					continue
				}
				return n, err
			}
			i := n
			for i < int(l) {
				j, err := w.conn.(*net.TCPConn).Read(p[i:int(l)])
				if err != nil {
					if e, ok := err.(net.Error); ok && e.Timeout() {
						// We are half way in our read...
						continue
					}
					return i, err
				}
				i += j
			}
			n = i
		}
	case "", "udp", "udp4", "udp6":
		for a := 0; a < attempts; a++ {
			setTimeouts(w)
			n, _, err = w.conn.(*net.UDPConn).ReadFromUDP(p)
			if err == nil {
				return n, err
			}
			if err != nil {
				if e, ok := err.(net.Error); ok && e.Timeout() {
					continue
				}
				return n, err
			}
		}
	}
	return
}

// send sends a dns msg to the address specified in w.
// If the message m contains a TSIG record the transaction
// signature is calculated.
func (w *reply) send(m *Msg) (err error) {
	var out []byte
	if t := m.IsTsig(); t != nil {
		mac := ""
		name := t.Hdr.Name
		if _, ok := w.client.TsigSecret[name]; !ok {
			return ErrSecret
		}
		out, mac, err = TsigGenerate(m, w.client.TsigSecret[name], w.tsigRequestMAC, w.tsigTimersOnly)
		w.tsigRequestMAC = mac
	} else {
		out, err = m.Pack()
	}
	if err != nil {
		return err
	}
	w.t = time.Now()
	if _, err = w.write(out); err != nil {
		return err
	}
	return nil
}

func (w *reply) write(p []byte) (n int, err error) {
	attempts := w.client.Attempts
	if attempts == 0 {
		attempts = 1
	}
	switch w.client.Net {
	case "tcp", "tcp4", "tcp6":
		if len(p) < 2 {
			return 0, io.ErrShortBuffer
		}
		for a := 0; a < attempts; a++ {
			setTimeouts(w)
			l := make([]byte, 2)
			l[0], l[1] = packUint16(uint16(len(p)))
			p = append(l, p...)
			n, err := w.conn.Write(p)
			if err != nil {
				if e, ok := err.(net.Error); ok && e.Timeout() {
					continue
				}
				return n, err
			}
			i := n
			if i < len(p) {
				j, err := w.conn.Write(p[i:len(p)])
				if err != nil {
					if e, ok := err.(net.Error); ok && e.Timeout() {
						// We are half way in our write...
						continue
					}
					return i, err
				}
				i += j
			}
			n = i
		}
	case "", "udp", "udp4", "udp6":
		for a := 0; a < attempts; a++ {
			setTimeouts(w)
			n, err = w.conn.(*net.UDPConn).Write(p)
			if err == nil {
				return
			}
			if err != nil {
				if e, ok := err.(net.Error); ok && e.Timeout() {
					continue
				}
				return n, err
			}
		}
	}
	return
}

func setTimeouts(w *reply) {
	if w.client.ReadTimeout == 0 {
		w.conn.SetReadDeadline(time.Now().Add(2 * 1e9))
	} else {
		w.conn.SetReadDeadline(time.Now().Add(w.client.ReadTimeout))
	}

	if w.client.WriteTimeout == 0 {
		w.conn.SetWriteDeadline(time.Now().Add(2 * 1e9))
	} else {
		w.conn.SetWriteDeadline(time.Now().Add(w.client.WriteTimeout))
	}
}
<commit_msg>Add exchange structure for the client<commit_after>package dns

// A concurrent client implementation.

import (
	"io"
	"net"
	"time"
)

// Order of events:
// *client -> *reply -> Exchange*() -> dial()/send()->write()/receive()->read()

type reply struct {
	client         *Client
	addr           string
	req            *Msg
	conn           net.Conn
	tsigRequestMAC string
	tsigTimersOnly bool
	tsigStatus     error
	rtt            time.Duration
	t              time.Time
}

// An Exchange is returned on the channel when calling client.Do or client.DoRtt
type Exchange struct {
	Request *Msg          // the outgoing message
	Reply   *Msg          // the reply coming back
	Rtt     time.Duration // Round trip time
	Error   error         // any errors
}

// A Client defines parameter for a DNS client. A nil
// Client is usable for sending queries.
type Client struct {
	Net          string            // if "tcp" a TCP query will be initiated, otherwise an UDP one (default is "" for UDP)
	Attempts     int               // number of attempts, if not set defaults to 1
	Retry        bool              // retry with TCP
	ReadTimeout  time.Duration     // the net.Conn.SetReadTimeout value for new connections (ns), defauls to 2 * 1e9
	WriteTimeout time.Duration     // the net.Conn.SetWriteTimeout value for new connections (ns), defauls to 2 * 1e9
	TsigSecret   map[string]string // secret(s) for Tsig map[<zonename>]<base64 secret>, zonename must be fully qualified
}

// Do performs an asynchronous query. The msg *Msg is the question to ask, the
// string addr is the address of the nameserver, the parameter data is used
// in the callback function. The call backback function is called with the
// original query, the answer returned from the nameserver an optional error and
// data.
// It calls Exchange.
func (c *Client) Do(msg *Msg, addr string, data interface{}, callback func(*Msg, *Msg, time.Duration, error, interface{})) {
	go func() {
		r, rtt, err := c.Exchange(msg, addr)
		callback(msg, r, rtt, err, data)
	}()
}

// Exchange performs an synchronous query. It sends the message m to the address
// contained in a and waits for an reply. Basic use pattern with a *Client:
//
//	c := new(dns.Client)
//	in, rtt, err := c.Exchange(message, "127.0.0.1:53")
//
func (c *Client) Exchange(m *Msg, a string) (r *Msg, rtt time.Duration, err error) {
	w := new(reply)
	w.client = c
	w.addr = a
	if err = w.dial(); err != nil {
		return nil, 0, err
	}
	if err = w.send(m); err != nil {
		return nil, 0, err
	}
	r, err = w.receive()
	return r, w.rtt, err
}

func (w *reply) RemoteAddr() net.Addr {
	if w.conn != nil {
		return w.conn.RemoteAddr()
	}
	return nil
}

// dial connects to the address addr for the network set in c.Net
func (w *reply) dial() (err error) {
	var conn net.Conn
	attempts := w.client.Attempts
	if attempts == 0 {
		attempts = 1
	}
	for a := 0; a < attempts; a++ {
		if w.client.Net == "" {
			conn, err = net.Dial("udp", w.addr)
		} else {
			conn, err = net.Dial(w.client.Net, w.addr)
		}
		if err != nil {
			// There are no timeouts defined?
			continue
		}
	}
	w.conn = conn
	return
}

func (w *reply) receive() (*Msg, error) {
	var p []byte
	m := new(Msg)
	switch w.client.Net {
	case "tcp", "tcp4", "tcp6":
		p = make([]byte, MaxMsgSize)
	case "", "udp", "udp4", "udp6":
		// OPT! TODO(mg)
		p = make([]byte, DefaultMsgSize)
	}
	n, err := w.read(p)
	if err != nil && n == 0 {
		return nil, err
	}
	p = p[:n]
	if err := m.Unpack(p); err != nil {
		return nil, err
	}
	w.rtt = time.Since(w.t)
	m.Size = n
	if t := m.IsTsig(); t != nil {
		secret := t.Hdr.Name
		if _, ok := w.client.TsigSecret[secret]; !ok {
			w.tsigStatus = ErrSecret
			return m, ErrSecret
		}
		// Need to work on the original message p, as that was used to calculate the tsig.
		w.tsigStatus = TsigVerify(p, w.client.TsigSecret[secret], w.tsigRequestMAC, w.tsigTimersOnly)
	}
	return m, w.tsigStatus
}

func (w *reply) read(p []byte) (n int, err error) {
	if w.conn == nil {
		return 0, ErrConnEmpty
	}
	if len(p) < 2 {
		return 0, io.ErrShortBuffer
	}
	attempts := w.client.Attempts
	if attempts == 0 {
		attempts = 1
	}
	switch w.client.Net {
	case "tcp", "tcp4", "tcp6":
		setTimeouts(w)
		for a := 0; a < attempts; a++ {
			n, err = w.conn.(*net.TCPConn).Read(p[0:2])
			if err != nil || n != 2 {
				if e, ok := err.(net.Error); ok && e.Timeout() {
					continue
				}
				return n, err
			}
			l, _ := unpackUint16(p[0:2], 0)
			if l == 0 {
				return 0, ErrShortRead
			}
			if int(l) > len(p) {
				return int(l), io.ErrShortBuffer
			}
			n, err = w.conn.(*net.TCPConn).Read(p[:l])
			if err != nil {
				if e, ok := err.(net.Error); ok && e.Timeout() {
					continue
				}
				return n, err
			}
			i := n
			for i < int(l) {
				j, err := w.conn.(*net.TCPConn).Read(p[i:int(l)])
				if err != nil {
					if e, ok := err.(net.Error); ok && e.Timeout() {
						// We are half way in our read...
						continue
					}
					return i, err
				}
				i += j
			}
			n = i
		}
	case "", "udp", "udp4", "udp6":
		for a := 0; a < attempts; a++ {
			setTimeouts(w)
			n, _, err = w.conn.(*net.UDPConn).ReadFromUDP(p)
			if err == nil {
				return n, err
			}
			if err != nil {
				if e, ok := err.(net.Error); ok && e.Timeout() {
					continue
				}
				return n, err
			}
		}
	}
	return
}

// send sends a dns msg to the address specified in w.
// If the message m contains a TSIG record the transaction
// signature is calculated.
func (w *reply) send(m *Msg) (err error) {
	var out []byte
	if t := m.IsTsig(); t != nil {
		mac := ""
		name := t.Hdr.Name
		if _, ok := w.client.TsigSecret[name]; !ok {
			return ErrSecret
		}
		out, mac, err = TsigGenerate(m, w.client.TsigSecret[name], w.tsigRequestMAC, w.tsigTimersOnly)
		w.tsigRequestMAC = mac
	} else {
		out, err = m.Pack()
	}
	if err != nil {
		return err
	}
	w.t = time.Now()
	if _, err = w.write(out); err != nil {
		return err
	}
	return nil
}

func (w *reply) write(p []byte) (n int, err error) {
	attempts := w.client.Attempts
	if attempts == 0 {
		attempts = 1
	}
	switch w.client.Net {
	case "tcp", "tcp4", "tcp6":
		if len(p) < 2 {
			return 0, io.ErrShortBuffer
		}
		for a := 0; a < attempts; a++ {
			setTimeouts(w)
			l := make([]byte, 2)
			l[0], l[1] = packUint16(uint16(len(p)))
			p = append(l, p...)
			n, err := w.conn.Write(p)
			if err != nil {
				if e, ok := err.(net.Error); ok && e.Timeout() {
					continue
				}
				return n, err
			}
			i := n
			if i < len(p) {
				j, err := w.conn.Write(p[i:len(p)])
				if err != nil {
					if e, ok := err.(net.Error); ok && e.Timeout() {
						// We are half way in our write...
						continue
					}
					return i, err
				}
				i += j
			}
			n = i
		}
	case "", "udp", "udp4", "udp6":
		for a := 0; a < attempts; a++ {
			setTimeouts(w)
			n, err = w.conn.(*net.UDPConn).Write(p)
			if err == nil {
				return
			}
			if err != nil {
				if e, ok := err.(net.Error); ok && e.Timeout() {
					continue
				}
				return n, err
			}
		}
	}
	return
}

func setTimeouts(w *reply) {
	if w.client.ReadTimeout == 0 {
		w.conn.SetReadDeadline(time.Now().Add(2 * 1e9))
	} else {
		w.conn.SetReadDeadline(time.Now().Add(w.client.ReadTimeout))
	}

	if w.client.WriteTimeout == 0 {
		w.conn.SetWriteDeadline(time.Now().Add(2 * 1e9))
	} else {
		w.conn.SetWriteDeadline(time.Now().Add(w.client.WriteTimeout))
	}
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"flag"
	"fmt"
	"os"
	"os/exec"
	"strconv"
	"strings"

	mp "github.com/mackerelio/go-mackerel-plugin"
	"github.com/mackerelio/mackerel-agent/logging"
)

var logger = logging.GetLogger("metrics.plugin.jvm")

var graphdef map[string](mp.Graphs) = map[string](mp.Graphs){
	"jvm.gc_events": mp.Graphs{
		Label: "JVM GC events",
		Unit:  "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "YGC", Label: "Young GC event", Diff: true},
			mp.Metrics{Name: "FGC", Label: "Full GC event", Diff: true},
		},
	},
	"jvm.gc_time": mp.Graphs{
		Label: "JVM GC time (msec)",
		Unit:  "float",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "YGCT", Label: "Young GC time", Diff: true},
			mp.Metrics{Name: "FGCT", Label: "Full GC time", Diff: true},
		},
	},
	"jvm.new_space": mp.Graphs{
		Label: "JVM New Space memory (KB)",
		Unit:  "float",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "NGCMX", Label: "New max", Diff: false},
			mp.Metrics{Name: "NGC", Label: "New current", Diff: false},
			mp.Metrics{Name: "EU", Label: "Eden used", Diff: false},
			mp.Metrics{Name: "S0U", Label: "Survivor0 used", Diff: false},
			mp.Metrics{Name: "S1U", Label: "Survivor1 used", Diff: false},
		},
	},
	"jvm.old_space": mp.Graphs{
		Label: "JVM Old Space memory (KB)",
		Unit:  "float",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "OGCMX", Label: "Old max", Diff: false},
			mp.Metrics{Name: "OGC", Label: "Old current", Diff: false},
			mp.Metrics{Name: "OU", Label: "Old used", Diff: false},
		},
	},
	"jvm.perm_space": mp.Graphs{
		Label: "JVM Permanent Space (KB)",
		Unit:  "float",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "PGCMX", Label: "Perm max", Diff: false},
			mp.Metrics{Name: "PGC", Label: "Perm current", Diff: false},
			mp.Metrics{Name: "PU", Label: "Perm used", Diff: false},
		},
	},
}

type JVMPlugin struct {
	Target    string
	Lvmid     string
	JstatPath string
	Tempfile  string
}

// # jps
// 26547 NettyServer
// 6438 Jps
func FetchLvmidByAppname(appname, target, jpsPath string) (string, error) {
	out, err := exec.Command(jpsPath, target).Output()
	if err != nil {
		logger.Errorf("Failed to run exec jps. %s", err)
		return "", err
	}
	for _, line := range strings.Split(string(out), "\n") {
		words := strings.Split(line, " ")
		if len(words) != 2 {
			continue
		}
		lvmid, name := words[0], words[1]
		if name == appname {
			return lvmid, nil
		}
	}
	return "", errors.New(fmt.Sprintf("Cannot get lvmid from %s", appname))
}

func fetchJstatMetrics(lvmid, option, jstatPath string) (map[string]float64, error) {
	out, err := exec.Command(jstatPath, option, lvmid).Output()
	if err != nil {
		logger.Errorf("Failed to run exec jstat. %s", err)
		return nil, err
	}

	lines := strings.Split(string(out), "\n")
	keys := strings.Fields(lines[0])
	values := strings.Fields(lines[1])

	stat := make(map[string]float64)
	for i, key := range keys {
		value, err := strconv.ParseFloat(values[i], 64)
		if err != nil {
			logger.Warningf("Failed to parse value. %s", err)
		}
		stat[key] = value
	}

	return stat, nil
}

func mergeStat(dst, src map[string]float64) {
	for k, v := range src {
		dst[k] = v
	}
}

// # jstat -gc <vmid>
//  S0C    S1C    S0U    S1U      EC       EU        OC         OU       PC     PU    YGC     YGCT    FGC    FGCT     GCT
// 3584.0 3584.0 2528.0  0.0   692224.0 19062.4  1398272.0   485450.1  72704.0 72611.3   3152   30.229   0      0.000   30.229

// # jstat -gccapacity  <vmid>
//  NGCMN    NGCMX     NGC     S0C   S1C       EC      OGCMN      OGCMX       OGC         OC      PGCMN    PGCMX     PGC       PC     YGC    FGC
// 699392.0 699392.0 699392.0 4096.0 4096.0 691200.0  1398272.0  1398272.0  1398272.0  1398272.0  21504.0 524288.0  72704.0  72704.0   4212     0

// # jstat -gcnew  <vmid>
//  S0C    S1C    S0U    S1U   TT MTT  DSS      EC       EU     YGC     YGCT
// 3072.0 3072.0    0.0 2848.0  1  15 3072.0 693248.0 626782.2   3463   33.658

func (m JVMPlugin) FetchMetrics() (map[string]float64, error) {
	gcStat, err := fetchJstatMetrics(m.Lvmid, "-gc", m.JstatPath)
	if err != nil {
		return nil, err
	}
	gcCapacityStat, err := fetchJstatMetrics(m.Lvmid, "-gccapacity", m.JstatPath)
	if err != nil {
		return nil, err
	}
	gcNewStat, err := fetchJstatMetrics(m.Lvmid, "-gcnew", m.JstatPath)
	if err != nil {
		return nil, err
	}
	gcOldStat, err := fetchJstatMetrics(m.Lvmid, "-gcold", m.JstatPath)
	if err != nil {
		return nil, err
	}

	stat := make(map[string]float64)
	mergeStat(stat, gcStat)
	mergeStat(stat, gcCapacityStat)
	mergeStat(stat, gcNewStat)
	mergeStat(stat, gcOldStat)

	return stat, nil
}

func (m JVMPlugin) GraphDefinition() map[string](mp.Graphs) {
	return graphdef
}

func main() {
	optHost := flag.String("host", "localhost", "Hostname")
	optPort := flag.String("port", "1099", "Port")
	optJstatPath := flag.String("jstatpath", "/usr/bin/jstat", "jstat path")
	optJpsPath := flag.String("jpspath", "/usr/bin/jps", "jps path")
	optJavaName := flag.String("javaname", "", "Java app name")
	optTempfile := flag.String("tempfile", "", "Temp file name")
	flag.Parse()

	if *optJavaName == "" {
		logger.Errorf("javaname is required")
		flag.PrintDefaults()
		os.Exit(1)
	}

	var jvm JVMPlugin
	jvm.Target = fmt.Sprintf("%s:%s", *optHost, *optPort)
	lvmid, err := FetchLvmidByAppname(*optJavaName, jvm.Target, *optJpsPath)
	if err != nil {
		logger.Errorf("Failed to fetch lvmid. %s", err)
		os.Exit(1)
	}
	jvm.Lvmid = lvmid
	jvm.JstatPath = *optJstatPath

	helper := mp.NewMackerelPlugin(jvm)
	if *optTempfile != "" {
		helper.Tempfile = *optTempfile
	} else {
		helper.Tempfile = fmt.Sprintf("/tmp/mackerel-plugin-jvm-%s", *optHost)
	}

	if os.Getenv("MACKEREL_AGENT_PLUGIN_META") != "" {
		helper.OutputDefinitions()
	} else {
		helper.OutputValues()
	}
}
<commit_msg>Add 'javaname' specified by user namespace to metric name<commit_after>package main

import (
	"errors"
	"flag"
	"fmt"
	"os"
	"os/exec"
	"strconv"
	"strings"

	mp "github.com/mackerelio/go-mackerel-plugin"
	"github.com/mackerelio/mackerel-agent/logging"
)

var logger = logging.GetLogger("metrics.plugin.jvm")

var graphdef map[string](mp.Graphs) = map[string](mp.Graphs){
	"jvm.gc_events": mp.Graphs{
		Label: "JVM GC events",
		Unit:  "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "YGC", Label: "Young GC event", Diff: true},
			mp.Metrics{Name: "FGC", Label: "Full GC event", Diff: true},
		},
	},
	"jvm.gc_time": mp.Graphs{
		Label: "JVM GC time (msec)",
		Unit:  "float",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "YGCT", Label: "Young GC time", Diff: true},
			mp.Metrics{Name: "FGCT", Label: "Full GC time", Diff: true},
		},
	},
	"jvm.new_space": mp.Graphs{
		Label: "JVM New Space memory (KB)",
		Unit:  "float",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "NGCMX", Label: "New max", Diff: false},
			mp.Metrics{Name: "NGC", Label: "New current", Diff: false},
			mp.Metrics{Name: "EU", Label: "Eden used", Diff: false},
			mp.Metrics{Name: "S0U", Label: "Survivor0 used", Diff: false},
			mp.Metrics{Name: "S1U", Label: "Survivor1 used", Diff: false},
		},
	},
	"jvm.old_space": mp.Graphs{
		Label: "JVM Old Space memory (KB)",
		Unit:  "float",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "OGCMX", Label: "Old max", Diff: false},
			mp.Metrics{Name: "OGC", Label: "Old current", Diff: false},
			mp.Metrics{Name: "OU", Label: "Old used", Diff: false},
		},
	},
	"jvm.perm_space": mp.Graphs{
		Label: "JVM Permanent Space (KB)",
		Unit:  "float",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "PGCMX", Label: "Perm max", Diff: false},
			mp.Metrics{Name: "PGC", Label: "Perm current", Diff: false},
			mp.Metrics{Name: "PU", Label: "Perm used", Diff: false},
		},
	},
}

type JVMPlugin struct {
	Target    string
	Lvmid     string
	JstatPath string
	JavaName  string
	Tempfile  string
}

// # jps
// 26547 NettyServer
// 6438 Jps
func FetchLvmidByAppname(appname, target, jpsPath string) (string, error) {
	out, err := exec.Command(jpsPath, target).Output()
	if err != nil {
		logger.Errorf("Failed to run exec jps. %s", err)
		return "", err
	}
	for _, line := range strings.Split(string(out), "\n") {
		words := strings.Split(line, " ")
		if len(words) != 2 {
			continue
		}
		lvmid, name := words[0], words[1]
		if name == appname {
			return lvmid, nil
		}
	}
	return "", errors.New(fmt.Sprintf("Cannot get lvmid from %s", appname))
}

func fetchJstatMetrics(lvmid, option, jstatPath string) (map[string]float64, error) {
	out, err := exec.Command(jstatPath, option, lvmid).Output()
	if err != nil {
		logger.Errorf("Failed to run exec jstat. %s", err)
		return nil, err
	}

	lines := strings.Split(string(out), "\n")
	keys := strings.Fields(lines[0])
	values := strings.Fields(lines[1])

	stat := make(map[string]float64)
	for i, key := range keys {
		value, err := strconv.ParseFloat(values[i], 64)
		if err != nil {
			logger.Warningf("Failed to parse value. %s", err)
		}
		stat[key] = value
	}

	return stat, nil
}

func mergeStat(dst, src map[string]float64) {
	for k, v := range src {
		dst[k] = v
	}
}

// # jstat -gc <vmid>
//  S0C    S1C    S0U    S1U      EC       EU        OC         OU       PC     PU    YGC     YGCT    FGC    FGCT     GCT
// 3584.0 3584.0 2528.0  0.0   692224.0 19062.4  1398272.0   485450.1  72704.0 72611.3   3152   30.229   0      0.000   30.229

// # jstat -gccapacity  <vmid>
//  NGCMN    NGCMX     NGC     S0C   S1C       EC      OGCMN      OGCMX       OGC         OC      PGCMN    PGCMX     PGC       PC     YGC    FGC
// 699392.0 699392.0 699392.0 4096.0 4096.0 691200.0  1398272.0  1398272.0  1398272.0  1398272.0  21504.0 524288.0  72704.0  72704.0   4212     0

// # jstat -gcnew  <vmid>
//  S0C    S1C    S0U    S1U   TT MTT  DSS      EC       EU     YGC     YGCT
// 3072.0 3072.0    0.0 2848.0  1  15 3072.0 693248.0 626782.2   3463   33.658

func (m JVMPlugin) FetchMetrics() (map[string]float64, error) {
	gcStat, err := fetchJstatMetrics(m.Lvmid, "-gc", m.JstatPath)
	if err != nil {
		return nil, err
	}
	gcCapacityStat, err := fetchJstatMetrics(m.Lvmid, "-gccapacity", m.JstatPath)
	if err != nil {
		return nil, err
	}
	gcNewStat, err := fetchJstatMetrics(m.Lvmid, "-gcnew", m.JstatPath)
	if err != nil {
		return nil, err
	}
	gcOldStat, err := fetchJstatMetrics(m.Lvmid, "-gcold", m.JstatPath)
	if err != nil {
		return nil, err
	}

	stat := make(map[string]float64)
	mergeStat(stat, gcStat)
	mergeStat(stat, gcCapacityStat)
	mergeStat(stat, gcNewStat)
	mergeStat(stat, gcOldStat)

	return stat, nil
}

func (m JVMPlugin) GraphDefinition() map[string](mp.Graphs) {
	return map[string](mp.Graphs){
		fmt.Sprintf("jvm.%s.gc_events", m.JavaName): mp.Graphs{
			Label: "JVM GC events",
			Unit:  "integer",
			Metrics: [](mp.Metrics){
				mp.Metrics{Name: "YGC", Label: "Young GC event", Diff: true},
				mp.Metrics{Name: "FGC", Label: "Full GC event", Diff: true},
			},
		},
		fmt.Sprintf("jvm.%s.gc_time", m.JavaName): mp.Graphs{
			Label: "JVM GC time (msec)",
			Unit:  "float",
			Metrics: [](mp.Metrics){
				mp.Metrics{Name: "YGCT", Label: "Young GC time", Diff: true},
				mp.Metrics{Name: "FGCT", Label: "Full GC time", Diff: true},
			},
		},
		fmt.Sprintf("jvm.%s.new_space", m.JavaName): mp.Graphs{
			Label: "JVM New Space memory (KB)",
			Unit:  "float",
			Metrics: [](mp.Metrics){
				mp.Metrics{Name: "NGCMX", Label: "New max", Diff: false},
				mp.Metrics{Name: "NGC", Label: "New current", Diff: false},
				mp.Metrics{Name: "EU", Label: "Eden used", Diff: false},
				mp.Metrics{Name: "S0U", Label: "Survivor0 used", Diff: false},
				mp.Metrics{Name: "S1U", Label: "Survivor1 used", Diff: false},
			},
		},
		fmt.Sprintf("jvm.%s.old_space", m.JavaName): mp.Graphs{
			Label: "JVM Old Space memory (KB)",
			Unit:  "float",
			Metrics: [](mp.Metrics){
				mp.Metrics{Name: "OGCMX", Label: "Old max", Diff: false},
				mp.Metrics{Name: "OGC", Label: "Old current", Diff: false},
				mp.Metrics{Name: "OU", Label: "Old used", Diff: false},
			},
		},
		fmt.Sprintf("jvm.%s.perm_space", m.JavaName): mp.Graphs{
			Label: "JVM Permanent Space (KB)",
			Unit:  "float",
			Metrics: [](mp.Metrics){
				mp.Metrics{Name: "PGCMX", Label: "Perm max", Diff: false},
				mp.Metrics{Name: "PGC", Label: "Perm current", Diff: false},
				mp.Metrics{Name: "PU", Label: "Perm used", Diff: false},
			},
		},
	}
}

func main() {
	optHost := flag.String("host", "localhost", "Hostname")
	optPort := flag.String("port", "1099", "Port")
	optJstatPath := flag.String("jstatpath", "/usr/bin/jstat", "jstat path")
	optJpsPath := flag.String("jpspath", "/usr/bin/jps", "jps path")
	optJavaName := flag.String("javaname", "", "Java app name")
	optTempfile := flag.String("tempfile", "", "Temp file name")
	flag.Parse()

	if *optJavaName == "" {
		logger.Errorf("javaname is required")
		flag.PrintDefaults()
		os.Exit(1)
	}

	var jvm JVMPlugin
	jvm.Target = fmt.Sprintf("%s:%s", *optHost, *optPort)
	lvmid, err := FetchLvmidByAppname(*optJavaName, jvm.Target, *optJpsPath)
	if err != nil {
		logger.Errorf("Failed to fetch lvmid. %s", err)
		os.Exit(1)
	}
	jvm.Lvmid = lvmid
	jvm.JstatPath = *optJstatPath
	jvm.JavaName = strings.ToLower(*optJavaName)

	helper := mp.NewMackerelPlugin(jvm)
	if *optTempfile != "" {
		helper.Tempfile = *optTempfile
	} else {
		helper.Tempfile = fmt.Sprintf("/tmp/mackerel-plugin-jvm-%s", *optHost)
	}

	if os.Getenv("MACKEREL_AGENT_PLUGIN_META") != "" {
		helper.OutputDefinitions()
	} else {
		helper.OutputValues()
	}
}
<|endoftext|>


<commit_before>package slack_rtm

import (
	"encoding/json"
	"errors"
	ws "github.com/gorilla/websocket"
	"io/ioutil"
	"net/http"
)

var slackAddr = "https://slack.com/api/rtm.start"

var (
	errInvalidEvent = errors.New("slackClient: message received but no type specified")
	errTypeNotFound = errors.New("slackClient: message received but type unrecognized")
)

type SlackClient struct {
	slackData  SlackData
	dispatcher *slackDispatcher
	conn       *ws.Conn
}

func New(token string) (*SlackClient, error) {

	resp, err := http.Get(slackAddr + "?token=" + token)
	defer resp.Body.Close()
	if err != nil {
		return nil, err
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	s := &SlackClient{
		dispatcher: &slackDispatcher{},
	}
	err = json.Unmarshal(body, &s.slackData)
	if err != nil {
		return nil, err
	}

	return s, nil
}

func (s *SlackClient) Run(h ...HelloHandler) error {

	conn, _, err := ws.DefaultDialer.Dial(s.slackData.Url, nil)
	if err != nil {
		return err
	}

	s.conn = conn

	if len(h) > 0 {
		s.dispatcher.addHelloListener(h[0])
	}

	s.startReader()
	return nil
}

func (s *SlackClient) AddListener(eType EventType, v interface{}) {

	switch eType {
	case HelloEvent:
		s.dispatcher.addHelloListener(v.(HelloHandler))
	case MessageEvent:
		s.dispatcher.addMessageListener(v.(MessageHandler))
	}

}

type Event struct {
	t    EventType
	data interface{}
}

func (s *SlackClient) startReader() {

	for {
		_, data, err := s.conn.ReadMessage()
		if err != nil {
			continue
		}

		var event AbstractEvent
		err = json.Unmarshal(data, &event)
		if err != nil {
			continue
		}

		ctx := &SlackContext{s}

		d := s.dispatcher
		switch event.Type {
		case "hello":
			d.dispatchHello(ctx)
		case "message":
			m := &MessageType{}
			json.Unmarshal(data, &m)
			d.dispatchMessage(ctx, m)

		default:
			continue
		}
	}
}

func (s *SlackClient) WriteMessage(v interface{}) error {
	return s.conn.WriteJSON(v)
}
<commit_msg>defer only on connection success<commit_after>package slack_rtm

import (
	"encoding/json"
	"errors"
	ws "github.com/gorilla/websocket"
	"io/ioutil"
	"net/http"
)

var slackAddr = "https://slack.com/api/rtm.start"

var (
	errInvalidEvent = errors.New("slackClient: message received but no type specified")
	errTypeNotFound = errors.New("slackClient: message received but type unrecognized")
)

type SlackClient struct {
	slackData  SlackData
	dispatcher *slackDispatcher
	conn       *ws.Conn
}

func New(token string) (*SlackClient, error) {

	resp, err := http.Get(slackAddr + "?token=" + token)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	s := &SlackClient{
		dispatcher: &slackDispatcher{},
	}
	err = json.Unmarshal(body, &s.slackData)
	if err != nil {
		return nil, err
	}

	return s, nil
}

func (s *SlackClient) Run(h ...HelloHandler) error {

	conn, _, err := ws.DefaultDialer.Dial(s.slackData.Url, nil)
	if err != nil {
		return err
	}

	s.conn = conn

	if len(h) > 0 {
		s.dispatcher.addHelloListener(h[0])
	}

	s.startReader()
	return nil
}

func (s *SlackClient) AddListener(eType EventType, v interface{}) {

	switch eType {
	case HelloEvent:
		s.dispatcher.addHelloListener(v.(HelloHandler))
	case MessageEvent:
		s.dispatcher.addMessageListener(v.(MessageHandler))
	}

}

type Event struct {
	t    EventType
	data interface{}
}

func (s *SlackClient) startReader() {

	for {
		_, data, err := s.conn.ReadMessage()
		if err != nil {
			continue
		}

		var event AbstractEvent
		err = json.Unmarshal(data, &event)
		if err != nil {
			continue
		}

		ctx := &SlackContext{s}

		d := s.dispatcher
		switch event.Type {
		case "hello":
			d.dispatchHello(ctx)
		case "message":
			m := &MessageType{}
			json.Unmarshal(data, &m)
			d.dispatchMessage(ctx, m)

		default:
			continue
		}
	}
}

func (s *SlackClient) WriteMessage(v interface{}) error {
	return s.conn.WriteJSON(v)
}
<|endoftext|>


<commit_before>// Copyright 2015 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package maps

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"time"

	"golang.org/x/net/context"
	"golang.org/x/net/context/ctxhttp"
	"googlemaps.github.io/maps/internal"
)

// Client may be used to make requests to the Google Maps WebService APIs
type Client struct {
	httpClient        *http.Client
	apiKey            string
	baseURL           string
	clientID          string
	signature         []byte
	requestsPerSecond int
	rateLimiter       chan int
}

// ClientOption is the type of constructor options for NewClient(...).
type ClientOption func(*Client) error

var defaultRequestsPerSecond = 10

// NewClient constructs a new Client which can make requests to the Google Maps WebService APIs.
func NewClient(options ...ClientOption) (*Client, error) {
	c := &Client{requestsPerSecond: defaultRequestsPerSecond}
	WithHTTPClient(&http.Client{})(c)
	for _, option := range options {
		err := option(c)
		if err != nil {
			return nil, err
		}
	}
	if c.apiKey == "" && (c.clientID == "" || len(c.signature) == 0) {
		return nil, errors.New("maps: API Key or Maps for Work credentials missing")
	}

	// Implement a bursty rate limiter.
	// Allow up to 1 second worth of requests to be made at once.
	c.rateLimiter = make(chan int, c.requestsPerSecond)
	// Prefill rateLimiter with 1 seconds worth of requests.
	for i := 0; i < c.requestsPerSecond; i++ {
		c.rateLimiter <- 1
	}
	go func() {
		// Wait a second for pre-filled quota to drain
		time.Sleep(time.Second)
		// Then, refill rateLimiter continuously
		for _ = range time.Tick(time.Second / time.Duration(c.requestsPerSecond)) {
			c.rateLimiter <- 1
		}
	}()

	return c, nil
}

// WithHTTPClient configures a Maps API client with a http.Client to make requests over.
func WithHTTPClient(c *http.Client) ClientOption {
	return func(client *Client) error {
		if _, ok := c.Transport.(*transport); !ok {
			t := c.Transport
			if t != nil {
				c.Transport = &transport{Base: t}
			} else {
				c.Transport = &transport{Base: http.DefaultTransport}
			}
		}
		client.httpClient = c
		return nil
	}
}

// WithAPIKey configures a Maps API client with an API Key
func WithAPIKey(apiKey string) ClientOption {
	return func(c *Client) error {
		c.apiKey = apiKey
		return nil
	}
}

// WithClientIDAndSignature configures a Maps API client for a Maps for Work application
// The signature is assumed to be URL modified Base64 encoded
func WithClientIDAndSignature(clientID, signature string) ClientOption {
	return func(c *Client) error {
		c.clientID = clientID
		decoded, err := base64.URLEncoding.DecodeString(signature)
		if err != nil {
			return err
		}
		c.signature = decoded
		return nil
	}
}

// WithRateLimit configures the rate limit for back end requests.
// Default is to limit to 10 requests per second.
func WithRateLimit(requestsPerSecond int) ClientOption {
	return func(c *Client) error {
		c.requestsPerSecond = requestsPerSecond
		return nil
	}
}

type apiConfig struct {
	host            string
	path            string
	acceptsClientID bool
}

type apiRequest interface {
	params() url.Values
}

func (c *Client) get(ctx context.Context, config *apiConfig, apiReq apiRequest) (*http.Response, error) {
	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case <-c.rateLimiter:
		// Execute request.
	}

	host := config.host
	if c.baseURL != "" {
		host = c.baseURL
	}
	req, err := http.NewRequest("GET", host+config.path, nil)
	if err != nil {
		return nil, err
	}
	q, err := c.generateAuthQuery(config.path, apiReq.params(), config.acceptsClientID)
	if err != nil {
		return nil, err
	}
	req.URL.RawQuery = q
	return ctxhttp.Do(ctx, c.httpClient, req)
}

func (c *Client) post(ctx context.Context, config *apiConfig, apiReq interface{}) (*http.Response, error) {
	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case <-c.rateLimiter:
		// Execute request.
	}

	host := config.host
	if c.baseURL != "" {
		host = c.baseURL
	}

	body, err := json.Marshal(apiReq)
	if err != nil {
		return nil, err
	}
	req, err := http.NewRequest("POST", host+config.path, bytes.NewBuffer(body))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")
	q, err := c.generateAuthQuery(config.path, url.Values{}, config.acceptsClientID)
	if err != nil {
		return nil, err
	}

	req.URL.RawQuery = q
	return ctxhttp.Do(ctx, c.httpClient, req)
}

func (c *Client) getJSON(ctx context.Context, config *apiConfig, apiReq apiRequest, resp interface{}) error {
	httpResp, err := c.get(ctx, config, apiReq)
	if err != nil {
		return err
	}
	defer httpResp.Body.Close()

	return json.NewDecoder(httpResp.Body).Decode(resp)
}

func (c *Client) postJSON(ctx context.Context, config *apiConfig, apiReq interface{}, resp interface{}) error {
	httpResp, err := c.post(ctx, config, apiReq)
	if err != nil {
		return err
	}
	defer httpResp.Body.Close()

	return json.NewDecoder(httpResp.Body).Decode(resp)
}

type binaryResponse struct {
	statusCode  int
	contentType string
	data        io.ReadCloser
}

func (c *Client) getBinary(ctx context.Context, config *apiConfig, apiReq apiRequest) (binaryResponse, error) {
	httpResp, err := c.get(ctx, config, apiReq)
	if err != nil {
		return binaryResponse{}, err
	}

	return binaryResponse{httpResp.StatusCode, httpResp.Header.Get("Content-Type"), httpResp.Body}, nil
}

func (c *Client) generateAuthQuery(path string, q url.Values, acceptClientID bool) (string, error) {
	if c.apiKey != "" {
		q.Set("key", c.apiKey)
		return q.Encode(), nil
	}
	if acceptClientID {
		return internal.SignURL(path, c.clientID, c.signature, q)
	}
	return "", errors.New("maps: API Key missing")
}

// commonResponse contains the common response fields to most API calls inside
// the Google Maps APIs. This is used internally.
type commonResponse struct {
	// Status contains the status of the request, and may contain debugging
	// information to help you track down why the call failed.
	Status string `json:"status"`

	// ErrorMessage is the explanatory field added when Status is an error.
	ErrorMessage string `json:"error_message"`
}

// StatusError returns an error iff this object has a non-OK Status.
func (c *commonResponse) StatusError() error {
	if c.Status != "OK" {
		return fmt.Errorf("maps: %s - %s", c.Status, c.ErrorMessage)
	}
	return nil
}
<commit_msg>Makes client-level rate-limiting optional<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package maps

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"time"

	"golang.org/x/net/context"
	"golang.org/x/net/context/ctxhttp"
	"googlemaps.github.io/maps/internal"
)

// Client may be used to make requests to the Google Maps WebService APIs
type Client struct {
	httpClient        *http.Client
	apiKey            string
	baseURL           string
	clientID          string
	signature         []byte
	requestsPerSecond int
	rateLimiter       chan int
}

// ClientOption is the type of constructor options for NewClient(...).
type ClientOption func(*Client) error

var defaultRequestsPerSecond = 10

// NewClient constructs a new Client which can make requests to the Google Maps WebService APIs.
func NewClient(options ...ClientOption) (*Client, error) {
	c := &Client{requestsPerSecond: defaultRequestsPerSecond}
	WithHTTPClient(&http.Client{})(c)
	for _, option := range options {
		err := option(c)
		if err != nil {
			return nil, err
		}
	}
	if c.apiKey == "" && (c.clientID == "" || len(c.signature) == 0) {
		return nil, errors.New("maps: API Key or Maps for Work credentials missing")
	}

	if c.requestsPerSecond > 0 {
		// Implement a bursty rate limiter.
		// Allow up to 1 second worth of requests to be made at once.
		c.rateLimiter = make(chan int, c.requestsPerSecond)
		// Prefill rateLimiter with 1 seconds worth of requests.
		for i := 0; i < c.requestsPerSecond; i++ {
			c.rateLimiter <- 1
		}
		go func() {
			// Wait a second for pre-filled quota to drain
			time.Sleep(time.Second)
			// Then, refill rateLimiter continuously
			for _ = range time.Tick(time.Second / time.Duration(c.requestsPerSecond)) {
				c.rateLimiter <- 1
			}
		}()
	}

	return c, nil
}

// WithHTTPClient configures a Maps API client with a http.Client to make requests over.
func WithHTTPClient(c *http.Client) ClientOption {
	return func(client *Client) error {
		if _, ok := c.Transport.(*transport); !ok {
			t := c.Transport
			if t != nil {
				c.Transport = &transport{Base: t}
			} else {
				c.Transport = &transport{Base: http.DefaultTransport}
			}
		}
		client.httpClient = c
		return nil
	}
}

// WithAPIKey configures a Maps API client with an API Key
func WithAPIKey(apiKey string) ClientOption {
	return func(c *Client) error {
		c.apiKey = apiKey
		return nil
	}
}

// WithClientIDAndSignature configures a Maps API client for a Maps for Work application
// The signature is assumed to be URL modified Base64 encoded
func WithClientIDAndSignature(clientID, signature string) ClientOption {
	return func(c *Client) error {
		c.clientID = clientID
		decoded, err := base64.URLEncoding.DecodeString(signature)
		if err != nil {
			return err
		}
		c.signature = decoded
		return nil
	}
}

// WithRateLimit configures the rate limit for back end requests. Default is to
// limit to 10 requests per second. A value of zero disables rate limiting.
func WithRateLimit(requestsPerSecond int) ClientOption {
	return func(c *Client) error {
		c.requestsPerSecond = requestsPerSecond
		return nil
	}
}

type apiConfig struct {
	host            string
	path            string
	acceptsClientID bool
}

type apiRequest interface {
	params() url.Values
}

func (c *Client) awaitRateLimiter(ctx context.Context) error {
	if c.rateLimiter == nil {
		return nil
	}
	select {
	case <-ctx.Done():
		return ctx.Err()
	case <-c.rateLimiter:
		// Execute request.
		return nil
	}
}

func (c *Client) get(ctx context.Context, config *apiConfig, apiReq apiRequest) (*http.Response, error) {
	if err := c.awaitRateLimiter(ctx); err != nil {
		return nil, err
	}

	host := config.host
	if c.baseURL != "" {
		host = c.baseURL
	}
	req, err := http.NewRequest("GET", host+config.path, nil)
	if err != nil {
		return nil, err
	}
	q, err := c.generateAuthQuery(config.path, apiReq.params(), config.acceptsClientID)
	if err != nil {
		return nil, err
	}
	req.URL.RawQuery = q
	return ctxhttp.Do(ctx, c.httpClient, req)
}

func (c *Client) post(ctx context.Context, config *apiConfig, apiReq interface{}) (*http.Response, error) {
	if err := c.awaitRateLimiter(ctx); err != nil {
		return nil, err
	}

	host := config.host
	if c.baseURL != "" {
		host = c.baseURL
	}

	body, err := json.Marshal(apiReq)
	if err != nil {
		return nil, err
	}
	req, err := http.NewRequest("POST", host+config.path, bytes.NewBuffer(body))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")
	q, err := c.generateAuthQuery(config.path, url.Values{}, config.acceptsClientID)
	if err != nil {
		return nil, err
	}

	req.URL.RawQuery = q
	return ctxhttp.Do(ctx, c.httpClient, req)
}

func (c *Client) getJSON(ctx context.Context, config *apiConfig, apiReq apiRequest, resp interface{}) error {
	httpResp, err := c.get(ctx, config, apiReq)
	if err != nil {
		return err
	}
	defer httpResp.Body.Close()

	return json.NewDecoder(httpResp.Body).Decode(resp)
}

func (c *Client) postJSON(ctx context.Context, config *apiConfig, apiReq interface{}, resp interface{}) error {
	httpResp, err := c.post(ctx, config, apiReq)
	if err != nil {
		return err
	}
	defer httpResp.Body.Close()

	return json.NewDecoder(httpResp.Body).Decode(resp)
}

type binaryResponse struct {
	statusCode  int
	contentType string
	data        io.ReadCloser
}

func (c *Client) getBinary(ctx context.Context, config *apiConfig, apiReq apiRequest) (binaryResponse, error) {
	httpResp, err := c.get(ctx, config, apiReq)
	if err != nil {
		return binaryResponse{}, err
	}

	return binaryResponse{httpResp.StatusCode, httpResp.Header.Get("Content-Type"), httpResp.Body}, nil
}

func (c *Client) generateAuthQuery(path string, q url.Values, acceptClientID bool) (string, error) {
	if c.apiKey != "" {
		q.Set("key", c.apiKey)
		return q.Encode(), nil
	}
	if acceptClientID {
		return internal.SignURL(path, c.clientID, c.signature, q)
	}
	return "", errors.New("maps: API Key missing")
}

// commonResponse contains the common response fields to most API calls inside
// the Google Maps APIs. This is used internally.
type commonResponse struct {
	// Status contains the status of the request, and may contain debugging
	// information to help you track down why the call failed.
	Status string `json:"status"`

	// ErrorMessage is the explanatory field added when Status is an error.
	ErrorMessage string `json:"error_message"`
}

// StatusError returns an error iff this object has a non-OK Status.
func (c *commonResponse) StatusError() error {
	if c.Status != "OK" {
		return fmt.Errorf("maps: %s - %s", c.Status, c.ErrorMessage)
	}
	return nil
}
<|endoftext|>


<commit_before>package apns

import (
	"log"
	"net/rpc"
	"strconv"
)

type Client struct {
	Client *rpc.Client
}

func (c *Client) Configure(port int) {
	portString := strconv.Itoa(port)
	client, err := rpc.DialHTTP("tcp", ":"+portString)

	if err != nil {
		log.Fatal("unable to open client connection on localhost:" + portString)
	}

	c.Client = client
}

func (c *Client) Provision(appId string, certificatePath string, environment string) {
	if c.Client == nil {
		log.Fatal("configuration needs to be called first")
	}

	var reply int
	err := c.Client.Call("Server.Provision", certificatePath, &reply)

	if err != nil {
		log.Fatal("provisioning was unsuccessful")
	}
}

func (c *Client) Notify(appId string, notification *Notification) {
	if c.Client == nil {
		log.Fatal("configuration needs to be called first")
	}

	var reply int
	err := c.Client.Call("Server.Notify", notification, &reply)

	if err != nil {
		log.Fatal("notification was unsuccessful")
	}
}
<commit_msg>Differentiate RPC client, Client is a basic wrapper around init<commit_after>package apns

import (
	"log"
	"net/rpc"
	"strconv"
)

type Client struct {
	RpcClient *rpc.Client
}

func (c *Client) Configure(port int) {
	portString := strconv.Itoa(port)
	client, err := rpc.DialHTTP("tcp", ":"+portString)

	if err != nil {
		log.Fatal("unable to open client connection on localhost:" + portString)
	}

	c.RpcClient = client
}

func (c *Client) Provision(appId string, certificatePath string, environment string) {
	if c.RpcClient == nil {
		log.Fatal("configuration needs to be called first")
	}

	var reply int
	err := c.RpcClient.Call("Server.Provision", certificatePath, &reply)

	if err != nil {
		log.Fatal("provisioning was unsuccessful")
	}
}

func (c *Client) Notify(appId string, notification *Notification) {
	if c.RpcClient == nil {
		log.Fatal("configuration needs to be called first")
	}

	var reply int
	err := c.RpcClient.Call("Server.Notify", notification, &reply)

	if err != nil {
		log.Fatal("notification was unsuccessful")
	}
}
<|endoftext|>


<commit_before>package honeybadger

import (
	"net/http"
	"strings"
	"time"
)

// The Payload interface is implemented by any type which can be handled by the
// Backend interface.
type Payload interface {
	toJSON() []byte
}

// The Backend interface is implemented by the server type by default, but a
// custom implementation may be configured by the user.
type Backend interface {
	Notify(feature Feature, payload Payload) error
}

type noticeHandler func(*Notice) error

// Client is the manager for interacting with the Honeybadger service. It holds
// the configuration and implements the public API.
type Client struct {
	Config               *Configuration
	context              *Context
	worker               worker
	beforeNotifyHandlers []noticeHandler
}

// Configure updates the client configuration with the supplied config.
func (client *Client) Configure(config Configuration) {
	client.Config.update(&config)
}

// SetContext updates the client context with supplied context.
func (client *Client) SetContext(context Context) {
	client.context.Update(context)
}

// Flush blocks until the worker has processed its queue.
func (client *Client) Flush() {
	client.worker.Flush()
}

// BeforeNotify adds a callback function which is run before a notice is
// reported to Honeybadger. If any function returns an error the notification
// will be skipped, otherwise it will be sent.
func (client *Client) BeforeNotify(handler func(notice *Notice) error) {
	client.beforeNotifyHandlers = append(client.beforeNotifyHandlers, handler)
}

// Notify reports the error err to the Honeybadger service.
func (client *Client) Notify(err interface{}, extra ...interface{}) (string, error) {
	extra = append([]interface{}{*client.context}, extra...)
	notice := newNotice(client.Config, newError(err, 2), extra...)
	for _, handler := range client.beforeNotifyHandlers {
		if err := handler(notice); err != nil {
			return "", err
		}
	}
	workerErr := client.worker.Push(func() error {
		if err := client.Config.Backend.Notify(Notices, notice); err != nil {
			return err
		}
		return nil
	})
	if workerErr != nil {
		client.Config.Logger.Printf("worker error: %v\n", workerErr)
		return "", workerErr
	}
	return notice.Token, nil
}

// Monitor automatically reports panics which occur in the function it's called
// from. Must be deferred.
func (client *Client) Monitor() {
	if err := recover(); err != nil {
		client.Notify(newError(err, 2))
		panic(err)
	}
}

// Handler returns an http.Handler function which automatically reports panics
// to Honeybadger and then re-panics.
func (client *Client) Handler(h http.Handler) http.Handler {
	if h == nil {
		h = http.DefaultServeMux
	}
	fn := func(w http.ResponseWriter, r *http.Request) {
		defer func() {
			if err := recover(); err != nil {
				client.Notify(newError(err, 2), Params(r.Form), getCGIData(r), *r.URL)
				panic(err)
			}
		}()
		h.ServeHTTP(w, r)
	}
	return http.HandlerFunc(fn)
}

// MetricsHandler is deprecated.
func (client *Client) MetricsHandler(h http.Handler) http.Handler {
	client.Config.Logger.Printf("DEPRECATION WARNING: honeybadger.MetricsHandler() has no effect and will be removed.")
	if h == nil {
		h = http.DefaultServeMux
	}
	fn := func(w http.ResponseWriter, r *http.Request) {
		h.ServeHTTP(w, r)
	}
	return http.HandlerFunc(fn)
}

// Increment is deprecated.
func (client *Client) Increment(metric string, value int) {
	client.Config.Logger.Printf("DEPRECATION WARNING: honeybadger.Increment() has no effect and will be removed.")
}

// Timing is deprecated.
func (client *Client) Timing(metric string, value time.Duration) {
	client.Config.Logger.Printf("DEPRECATION WARNING: honeybadger.Timing() has no effect and will be removed.")
}

// New returns a new instance of Client.
func New(c Configuration) *Client {
	config := newConfig(c)
	worker := newBufferedWorker(config)

	client := Client{
		Config:  config,
		worker:  worker,
		context: &Context{},
	}

	return &client
}

func getCGIData(request *http.Request) CGIData {
	cgiData := CGIData{}
	replacer := strings.NewReplacer("-", "_")
	for k, v := range request.Header {
		key := "HTTP_" + replacer.Replace(strings.ToUpper(k))
		cgiData[key] = v[0]
	}
	return cgiData
}
<commit_msg>Flush notices in client monitor when catching a panic<commit_after>package honeybadger

import (
	"net/http"
	"strings"
	"time"
)

// The Payload interface is implemented by any type which can be handled by the
// Backend interface.
type Payload interface {
	toJSON() []byte
}

// The Backend interface is implemented by the server type by default, but a
// custom implementation may be configured by the user.
type Backend interface {
	Notify(feature Feature, payload Payload) error
}

type noticeHandler func(*Notice) error

// Client is the manager for interacting with the Honeybadger service. It holds
// the configuration and implements the public API.
type Client struct {
	Config               *Configuration
	context              *Context
	worker               worker
	beforeNotifyHandlers []noticeHandler
}

// Configure updates the client configuration with the supplied config.
func (client *Client) Configure(config Configuration) {
	client.Config.update(&config)
}

// SetContext updates the client context with supplied context.
func (client *Client) SetContext(context Context) {
	client.context.Update(context)
}

// Flush blocks until the worker has processed its queue.
func (client *Client) Flush() {
	client.worker.Flush()
}

// BeforeNotify adds a callback function which is run before a notice is
// reported to Honeybadger. If any function returns an error the notification
// will be skipped, otherwise it will be sent.
func (client *Client) BeforeNotify(handler func(notice *Notice) error) {
	client.beforeNotifyHandlers = append(client.beforeNotifyHandlers, handler)
}

// Notify reports the error err to the Honeybadger service.
func (client *Client) Notify(err interface{}, extra ...interface{}) (string, error) {
	extra = append([]interface{}{*client.context}, extra...)
	notice := newNotice(client.Config, newError(err, 2), extra...)
	for _, handler := range client.beforeNotifyHandlers {
		if err := handler(notice); err != nil {
			return "", err
		}
	}
	workerErr := client.worker.Push(func() error {
		if err := client.Config.Backend.Notify(Notices, notice); err != nil {
			return err
		}
		return nil
	})
	if workerErr != nil {
		client.Config.Logger.Printf("worker error: %v\n", workerErr)
		return "", workerErr
	}
	return notice.Token, nil
}

// Monitor automatically reports panics which occur in the function it's called
// from. Must be deferred.
func (client *Client) Monitor() {
	if err := recover(); err != nil {
		client.Notify(newError(err, 2))
		client.Flush()
		panic(err)
	}
}

// Handler returns an http.Handler function which automatically reports panics
// to Honeybadger and then re-panics.
func (client *Client) Handler(h http.Handler) http.Handler {
	if h == nil {
		h = http.DefaultServeMux
	}
	fn := func(w http.ResponseWriter, r *http.Request) {
		defer func() {
			if err := recover(); err != nil {
				client.Notify(newError(err, 2), Params(r.Form), getCGIData(r), *r.URL)
				panic(err)
			}
		}()
		h.ServeHTTP(w, r)
	}
	return http.HandlerFunc(fn)
}

// MetricsHandler is deprecated.
func (client *Client) MetricsHandler(h http.Handler) http.Handler {
	client.Config.Logger.Printf("DEPRECATION WARNING: honeybadger.MetricsHandler() has no effect and will be removed.")
	if h == nil {
		h = http.DefaultServeMux
	}
	fn := func(w http.ResponseWriter, r *http.Request) {
		h.ServeHTTP(w, r)
	}
	return http.HandlerFunc(fn)
}

// Increment is deprecated.
func (client *Client) Increment(metric string, value int) {
	client.Config.Logger.Printf("DEPRECATION WARNING: honeybadger.Increment() has no effect and will be removed.")
}

// Timing is deprecated.
func (client *Client) Timing(metric string, value time.Duration) {
	client.Config.Logger.Printf("DEPRECATION WARNING: honeybadger.Timing() has no effect and will be removed.")
}

// New returns a new instance of Client.
func New(c Configuration) *Client {
	config := newConfig(c)
	worker := newBufferedWorker(config)

	client := Client{
		Config:  config,
		worker:  worker,
		context: &Context{},
	}

	return &client
}

func getCGIData(request *http.Request) CGIData {
	cgiData := CGIData{}
	replacer := strings.NewReplacer("-", "_")
	for k, v := range request.Header {
		key := "HTTP_" + replacer.Replace(strings.ToUpper(k))
		cgiData[key] = v[0]
	}
	return cgiData
}
<|endoftext|>


<commit_before>package router

import (
	"testing"
)

type cargs map[string]string

func TestStringConverter(t *testing.T) {
	var stringConverterTests = []struct {
		args        cargs
		regexp      string
		toGoParam   string
		toGoResult  string
		toUrlParam  string
		toUrlResult string
	}{
		{cargs{}, `[^/]{1,}`, "test", "test", "test", "test"},
		{cargs{"minLength": "1"}, `[^/]{1,}`, "test", "test", "test", "test"},
		{cargs{"minLength": "1", "maxLength": "4"}, `[^/]{1,4}`, "test", "test", "test", "test"},
		{cargs{"minLength": "1", "maxLength": "2", "length": "4"}, `[^/]{4}`, "test", "test", "test", "test"},
	}

	for _, tt := range stringConverterTests {
		c := NewStringConverter(tt.args)

		if regexp := c.Regexp(); regexp != tt.regexp {
			t.Errorf("StringConverter regexp expected `%v` but got `%v`",
				tt.regexp, regexp)
		}

		toGoResult, err := c.ToGo(tt.toGoParam)
		if err != nil {
			t.Errorf("StringConverter ToGo(%q) unexpected error: %v",
				tt.toGoParam, err)
		}
		if toGoResult != tt.toGoResult {
			t.Errorf("StringConverter ToGo(%q) expected %q but got %q",
				tt.toGoParam, tt.toGoResult, toGoResult)
		}

		toUrlResult, err := c.ToUrl(tt.toUrlParam)
		if err != nil {
			t.Errorf("StringConverter ToUrl(%q) unexpected error: %v",
				tt.toUrlParam, err)
		}
		if toUrlResult != tt.toUrlResult {
			t.Errorf("StringConverter ToUrl(%q) expected %q but got %q",
				tt.toUrlParam, tt.toUrlResult, toUrlResult)
		}
	}
}

func TestPathConverter(t *testing.T) {
	var pathConverterTests = []struct {
		toGoParam   string
		toGoResult  string
		toUrlParam  string
		toUrlResult string
	}{
		{"foo", "foo", "foo", "foo"},
		{"foo/bar", "foo/bar", "foo/bar", "foo/bar"},
	}

	args := cargs{}
	for _, tt := range pathConverterTests {
		c := NewPathConverter(args)

		toGoResult, err := c.ToGo(tt.toGoParam)
		if err != nil {
			t.Errorf("PathConverter ToGo(%q) unexpected error: %v",
				tt.toGoParam, err)
		}
		if toGoResult != tt.toGoResult {
			t.Errorf("PathConverter ToGo(%q) expected %q but got %q",
				tt.toGoParam, tt.toGoResult, toGoResult)
		}

		toUrlResult, err := c.ToUrl(tt.toUrlParam)
		if err != nil {
			t.Errorf("PathConverter ToUrl(%q) unexpected error: %v",
				tt.toUrlParam, err)
		}
		if toUrlResult != tt.toUrlResult {
			t.Errorf("PathConverter ToUrl(%q) expected %q but got %q",
				tt.toUrlParam, tt.toUrlResult, toUrlResult)
		}
	}
}

func TestPathConverterNil(t *testing.T) {
	args := cargs{"key": "value"}
	c := NewPathConverter(args)
	if c != nil {
		t.Errorf("NewPathConverter(%v) = %v, want <nil>", args, c)
	}
}

func TestPathConverterRegexp(t *testing.T) {
	args := cargs{}
	expectedRegexp := `[^/].*?`
	c := NewPathConverter(args)
	if regexp := c.Regexp(); regexp != expectedRegexp {
		t.Errorf("PathConverter regexp expected `%v` but got `%v`",
			expectedRegexp, regexp)
	}
}

func TestIntConverter(t *testing.T) {
	var intConverterTests = []struct {
		args        cargs
		regexp      string
		toGoParam   string
		toGoResult  int
		toUrlParam  int
		toUrlResult string
	}{
		{cargs{}, `\d+`, "4", 4, 4, "4"},
		{cargs{"digits": "2"}, `\d+`, "44", 44, 44, "44"},
		{cargs{"digits": "2"}, `\d+`, "04", 4, 4, "04"},
		{cargs{"digits": "2"}, `\d+`, "4", -1, 4, "04"},
		{cargs{"min": "3"}, `\d+`, "4", 4, 4, "4"},
		{cargs{"min": "4"}, `\d+`, "4", 4, 4, "4"},
		{cargs{"min": "5"}, `\d+`, "4", -1, 4, "4"},
		{cargs{"max": "5"}, `\d+`, "4", 4, 4, "4"},
		{cargs{"max": "4"}, `\d+`, "4", 4, 4, "4"},
		{cargs{"max": "3"}, `\d+`, "4", -1, 4, "4"},
		{cargs{"min": "3", "max": "5"}, `\d+`, "4", 4, 4, "4"},
		{cargs{"min": "4", "max": "5"}, `\d+`, "4", 4, 4, "4"},
		{cargs{"min": "5", "max": "5"}, `\d+`, "4", -1, 4, "4"},
		{cargs{"min": "3", "max": "5"}, `\d+`, "4", 4, 4, "4"},
		{cargs{"min": "3", "max": "4"}, `\d+`, "4", 4, 4, "4"},
		{cargs{"min": "3", "max": "3"}, `\d+`, "4", -1, 4, "4"},
		{cargs{"digits": "2", "min": "3", "max": "4"}, `\d+`, "04", 4, 4, "04"},
		{cargs{"digits": "2", "min": "3", "max": "4"}, `\d+`, "05", -1, 5, "05"},
	}

	for _, tt := range intConverterTests {
		c := NewIntConverter(tt.args)
		if _, ok := c.(*IntConverter); !ok {
			t.Errorf("NewIntConverter(%v) got <nil>", tt.args)
			continue
		}

		if regexp := c.Regexp(); regexp != tt.regexp {
			t.Errorf("IntConverter regexp expected `%v` but got `%v`",
				tt.regexp, regexp)
		}

		toGoResult, err := c.ToGo(tt.toGoParam)
		if err != nil && tt.toGoResult != -1 {
			t.Errorf("IntConverter ToGo(%q) unexpected error: %v",
				tt.toGoParam, err)
		}
		if toGoResult != tt.toGoResult {
			t.Errorf("IntConverter ToGo(%q) expected %v but got %v",
				tt.toGoParam, tt.toGoResult, toGoResult)
		}

		toUrlResult, err := c.ToUrl(tt.toUrlParam)
		if err != nil {
			t.Errorf("IntConverter ToUrl(%v) unexpected error: %v",
				tt.toUrlParam, err)
		}
		if toUrlResult != tt.toUrlResult {
			t.Errorf("IntConverter ToUrl(%v) expected %v but got %q",
				tt.toUrlParam, tt.toUrlResult, toUrlResult)
		}
	}
}
<commit_msg>Clean up converter tests.<commit_after>package router

import (
	"testing"
)

type cargs map[string]string

func TestStringConverter(t *testing.T) {
	var stringConverterTests = []struct {
		args   cargs
		regexp string
	}{
		{cargs{}, `[^/]{1,}`},
		{cargs{"minLength": "1"}, `[^/]{1,}`},
		{cargs{"minLength": "1", "maxLength": "4"}, `[^/]{1,4}`},
		{cargs{"minLength": "1", "maxLength": "2", "length": "4"}, `[^/]{4}`},
	}

	for _, tt := range stringConverterTests {
		c := NewStringConverter(tt.args)
		str := "test"

		if regexp := c.Regexp(); regexp != tt.regexp {
			t.Errorf("StringConverter regexp expected `%v` but got `%v`",
				tt.regexp, regexp)
		}

		toGoResult, err := c.ToGo(str)
		if err != nil {
			t.Errorf("StringConverter ToGo(%q) unexpected error: %v", str, err)
		}
		if toGoResult != str {
			t.Errorf("StringConverter ToGo(%q)\nhave %q\nwant %q", str, toGoResult, str)
		}

		toUrlResult, err := c.ToUrl(str)
		if err != nil {
			t.Errorf("StringConverter ToUrl(%q) unexpected error: %v", str, err)
		}
		if toUrlResult != str {
			t.Errorf("StringConverter ToUrl(%q)\nhave %q\nwant %q", str, toUrlResult, str)
		}
	}
}

func TestPathConverter(t *testing.T) {
	var pathConverterTests = []string{
		"foo",
		"foo/bar",
	}

	args := cargs{}
	for _, path := range pathConverterTests {
		c := NewPathConverter(args)

		toGoResult, err := c.ToGo(path)
		if err != nil {
			t.Errorf("PathConverter ToGo(%q) unexpected error: %v", path, err)
		}
		if toGoResult != path {
			t.Errorf("PathConverter ToGo(%q)\nhave %q\nwant %q", path, toGoResult, path)
		}

		toUrlResult, err := c.ToUrl(path)
		if err != nil {
			t.Errorf("PathConverter ToUrl(%q) unexpected error: %v", path, err)
		}
		if toUrlResult != path {
			t.Errorf("PathConverter ToUrl(%q)\nhave %q\nwant %q", path, toUrlResult, path)
		}
	}
}

func TestPathConverterNil(t *testing.T) {
	args := cargs{"key": "value"}
	c := NewPathConverter(args)
	if c != nil {
		t.Errorf("NewPathConverter(%v) = %v, want <nil>", args, c)
	}
}

func TestPathConverterRegexp(t *testing.T) {
	args := cargs{}
	expectedRegexp := `[^/].*?`
	c := NewPathConverter(args)
	if regexp := c.Regexp(); regexp != expectedRegexp {
		t.Errorf("PathConverter regexp\nhave `%v`\nwant `%v`", regexp, expectedRegexp)
	}
}

func TestIntConverter(t *testing.T) {
	var intConverterTests = []struct {
		args        cargs
		regexp      string
		toGoParam   string
		toGoResult  int
		toUrlParam  int
		toUrlResult string
	}{
		{cargs{}, `\d+`, "4", 4, 4, "4"},
		{cargs{"digits": "2"}, `\d+`, "44", 44, 44, "44"},
		{cargs{"digits": "2"}, `\d+`, "04", 4, 4, "04"},
		{cargs{"digits": "2"}, `\d+`, "4", -1, 4, "04"},
		{cargs{"min": "3"}, `\d+`, "4", 4, 4, "4"},
		{cargs{"min": "4"}, `\d+`, "4", 4, 4, "4"},
		{cargs{"min": "5"}, `\d+`, "4", -1, 4, "4"},
		{cargs{"max": "5"}, `\d+`, "4", 4, 4, "4"},
		{cargs{"max": "4"}, `\d+`, "4", 4, 4, "4"},
		{cargs{"max": "3"}, `\d+`, "4", -1, 4, "4"},
		{cargs{"min": "3", "max": "5"}, `\d+`, "4", 4, 4, "4"},
		{cargs{"min": "4", "max": "5"}, `\d+`, "4", 4, 4, "4"},
		{cargs{"min": "5", "max": "5"}, `\d+`, "4", -1, 4, "4"},
		{cargs{"min": "3", "max": "5"}, `\d+`, "4", 4, 4, "4"},
		{cargs{"min": "3", "max": "4"}, `\d+`, "4", 4, 4, "4"},
		{cargs{"min": "3", "max": "3"}, `\d+`, "4", -1, 4, "4"},
		{cargs{"digits": "2", "min": "3", "max": "4"}, `\d+`, "04", 4, 4, "04"},
		{cargs{"digits": "2", "min": "3", "max": "4"}, `\d+`, "05", -1, 5, "05"},
	}

	for _, tt := range intConverterTests {
		c := NewIntConverter(tt.args)
		if _, ok := c.(*IntConverter); !ok {
			t.Errorf("NewIntConverter(%v) got <nil>", tt.args)
			continue
		}

		if regexp := c.Regexp(); regexp != tt.regexp {
			t.Errorf("IntConverter regexp\nhave `%v`\nwant `%v`", regexp, tt.regexp)
		}

		toGoResult, err := c.ToGo(tt.toGoParam)
		if err != nil && tt.toGoResult != -1 {
			t.Errorf("IntConverter ToGo(%q) unexpected error: %v",
				tt.toGoParam, err)
		}
		if toGoResult != tt.toGoResult {
			t.Errorf("IntConverter ToGo(%q)\nhave %v\nwant %v",
				tt.toGoParam, toGoResult, tt.toGoResult)
		}

		toUrlResult, err := c.ToUrl(tt.toUrlParam)
		if err != nil {
			t.Errorf("IntConverter ToUrl(%v) unexpected error: %v",
				tt.toUrlParam, err)
		}
		if toUrlResult != tt.toUrlResult {
			t.Errorf("IntConverter ToUrl(%v)\nhave %v\nwant %v",
				tt.toUrlParam, toUrlResult, tt.toUrlResult)
		}
	}
}
<|endoftext|>


<commit_before>package mpb

import (
	"bytes"
	"io"

	"github.com/vbauerster/mpb/v5/decor"
)

// BarOption is a function option which changes the default behavior of a bar.
type BarOption func(*bState)

func (s *bState) addDecorators(dest *[]decor.Decorator, decorators ...decor.Decorator) {
	type mergeWrapper interface {
		MergeUnwrap() []decor.Decorator
	}
	for _, decorator := range decorators {
		if mw, ok := decorator.(mergeWrapper); ok {
			*dest = append(*dest, mw.MergeUnwrap()...)
		}
		*dest = append(*dest, decorator)
	}
}

// AppendDecorators let you inject decorators to the bar's right side.
func AppendDecorators(decorators ...decor.Decorator) BarOption {
	return func(s *bState) {
		s.addDecorators(&s.aDecorators, decorators...)
	}
}

// PrependDecorators let you inject decorators to the bar's left side.
func PrependDecorators(decorators ...decor.Decorator) BarOption {
	return func(s *bState) {
		s.addDecorators(&s.pDecorators, decorators...)
	}
}

// BarID sets bar id.
func BarID(id int) BarOption {
	return func(s *bState) {
		s.id = id
	}
}

// BarWidth sets bar width independent of the container.
func BarWidth(width int) BarOption {
	return func(s *bState) {
		s.reqWidth = width
	}
}

// BarQueueAfter queues this (being constructed) bar to relplace
// runningBar after it has been completed.
func BarQueueAfter(runningBar *Bar) BarOption {
	if runningBar == nil {
		return nil
	}
	return func(s *bState) {
		s.runningBar = runningBar
	}
}

// BarRemoveOnComplete removes both bar's filler and its decorators
// on complete event.
func BarRemoveOnComplete() BarOption {
	return func(s *bState) {
		s.dropOnComplete = true
	}
}

// BarFillerClearOnComplete clears bar's filler on complete event.
// It's shortcut for BarFillerOnComplete("").
func BarFillerClearOnComplete() BarOption {
	return BarFillerOnComplete("")
}

// BarFillerOnComplete replaces bar's filler with message, on complete event.
func BarFillerOnComplete(message string) BarOption {
	return BarFillerMiddleware(func(base BarFiller) BarFiller {
		return BarFillerFunc(func(w io.Writer, reqWidth int, st decor.Statistics) {
			if st.Completed {
				io.WriteString(w, message)
			} else {
				base.Fill(w, reqWidth, st)
			}
		})
	})
}

// BarFillerMiddleware provides a way to augment default BarFiller.
func BarFillerMiddleware(middle func(BarFiller) BarFiller) BarOption {
	return func(s *bState) {
		s.middleware = middle
	}
}

// BarPriority sets bar's priority. Zero is highest priority, i.e. bar
// will be on top. If `BarReplaceOnComplete` option is supplied, this
// option is ignored.
func BarPriority(priority int) BarOption {
	return func(s *bState) {
		s.priority = priority
	}
}

// BarExtender is an option to extend bar to the next new line, with
// arbitrary output.
func BarExtender(filler BarFiller) BarOption {
	if filler == nil {
		return nil
	}
	return func(s *bState) {
		s.extender = makeExtFunc(filler)
	}
}

func makeExtFunc(filler BarFiller) extFunc {
	buf := new(bytes.Buffer)
	return func(r io.Reader, reqWidth int, st decor.Statistics) (io.Reader, int) {
		filler.Fill(buf, reqWidth, st)
		return io.MultiReader(r, buf), bytes.Count(buf.Bytes(), []byte("\n"))
	}
}

// BarFillerTrim bar filler is rendered with leading and trailing space
// like ' [===] ' by default. With this option leading and trailing
// space will be removed.
func BarFillerTrim() BarOption {
	return func(s *bState) {
		s.trimSpace = true
	}
}

// BarNoPop disables bar pop out of container. Effective when
// PopCompletedMode of container is enabled.
func BarNoPop() BarOption {
	return func(s *bState) {
		s.noPop = true
	}
}

// BarOptOn returns option when condition evaluates to true.
func BarOptOn(option BarOption, condition func() bool) BarOption {
	if condition() {
		return option
	}
	return nil
}
<commit_msg>BarFillerTrim godoc update<commit_after>package mpb

import (
	"bytes"
	"io"

	"github.com/vbauerster/mpb/v5/decor"
)

// BarOption is a function option which changes the default behavior of a bar.
type BarOption func(*bState)

func (s *bState) addDecorators(dest *[]decor.Decorator, decorators ...decor.Decorator) {
	type mergeWrapper interface {
		MergeUnwrap() []decor.Decorator
	}
	for _, decorator := range decorators {
		if mw, ok := decorator.(mergeWrapper); ok {
			*dest = append(*dest, mw.MergeUnwrap()...)
		}
		*dest = append(*dest, decorator)
	}
}

// AppendDecorators let you inject decorators to the bar's right side.
func AppendDecorators(decorators ...decor.Decorator) BarOption {
	return func(s *bState) {
		s.addDecorators(&s.aDecorators, decorators...)
	}
}

// PrependDecorators let you inject decorators to the bar's left side.
func PrependDecorators(decorators ...decor.Decorator) BarOption {
	return func(s *bState) {
		s.addDecorators(&s.pDecorators, decorators...)
	}
}

// BarID sets bar id.
func BarID(id int) BarOption {
	return func(s *bState) {
		s.id = id
	}
}

// BarWidth sets bar width independent of the container.
func BarWidth(width int) BarOption {
	return func(s *bState) {
		s.reqWidth = width
	}
}

// BarQueueAfter queues this (being constructed) bar to relplace
// runningBar after it has been completed.
func BarQueueAfter(runningBar *Bar) BarOption {
	if runningBar == nil {
		return nil
	}
	return func(s *bState) {
		s.runningBar = runningBar
	}
}

// BarRemoveOnComplete removes both bar's filler and its decorators
// on complete event.
func BarRemoveOnComplete() BarOption {
	return func(s *bState) {
		s.dropOnComplete = true
	}
}

// BarFillerClearOnComplete clears bar's filler on complete event.
// It's shortcut for BarFillerOnComplete("").
func BarFillerClearOnComplete() BarOption {
	return BarFillerOnComplete("")
}

// BarFillerOnComplete replaces bar's filler with message, on complete event.
func BarFillerOnComplete(message string) BarOption {
	return BarFillerMiddleware(func(base BarFiller) BarFiller {
		return BarFillerFunc(func(w io.Writer, reqWidth int, st decor.Statistics) {
			if st.Completed {
				io.WriteString(w, message)
			} else {
				base.Fill(w, reqWidth, st)
			}
		})
	})
}

// BarFillerMiddleware provides a way to augment default BarFiller.
func BarFillerMiddleware(middle func(BarFiller) BarFiller) BarOption {
	return func(s *bState) {
		s.middleware = middle
	}
}

// BarPriority sets bar's priority. Zero is highest priority, i.e. bar
// will be on top. If `BarReplaceOnComplete` option is supplied, this
// option is ignored.
func BarPriority(priority int) BarOption {
	return func(s *bState) {
		s.priority = priority
	}
}

// BarExtender is an option to extend bar to the next new line, with
// arbitrary output.
func BarExtender(filler BarFiller) BarOption {
	if filler == nil {
		return nil
	}
	return func(s *bState) {
		s.extender = makeExtFunc(filler)
	}
}

func makeExtFunc(filler BarFiller) extFunc {
	buf := new(bytes.Buffer)
	return func(r io.Reader, reqWidth int, st decor.Statistics) (io.Reader, int) {
		filler.Fill(buf, reqWidth, st)
		return io.MultiReader(r, buf), bytes.Count(buf.Bytes(), []byte("\n"))
	}
}

// BarFillerTrim removes leading and trailing space around the underlying BarFiller.
func BarFillerTrim() BarOption {
	return func(s *bState) {
		s.trimSpace = true
	}
}

// BarNoPop disables bar pop out of container. Effective when
// PopCompletedMode of container is enabled.
func BarNoPop() BarOption {
	return func(s *bState) {
		s.noPop = true
	}
}

// BarOptOn returns option when condition evaluates to true.
func BarOptOn(option BarOption, condition func() bool) BarOption {
	if condition() {
		return option
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package main

import (
	"fmt"
	"os"
	"os/exec"
	"strings"

	"launchpad.net/loggo"

	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/instance"
	"launchpad.net/juju-core/juju"
	_ "launchpad.net/juju-core/provider/all"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/utils"
)

var logger = loggo.GetLogger("juju.plugins.updatebootstrap")

const updateBootstrapDoc = `
Patches all machines after state server has been restored from backup, to
update state server address to new location.
`

type updateBootstrapCommand struct {
	cmd.EnvCommandBase
}

func (c *updateBootstrapCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "juju-update-bootstrap",
		Purpose: "update all machines after recovering state server",
		Doc:     updateBootstrapDoc,
	}
}

func (c *updateBootstrapCommand) Run(ctx *cmd.Context) error {
	conn, err := juju.NewConnFromName(c.EnvName)
	if err != nil {
		return err
	}
	defer conn.Close()
	stateAddr, err := GetStateAddress(conn.Environ)
	if err != nil {
		return err
	}
	fmt.Printf("using state address %v\n", stateAddr)
	return updateAllMachines(conn, stateAddr)
}

// GetStateAddress returns the address of one state server
func GetStateAddress(environ environs.Environ) (string, error) {
	// XXX: Can easily look up state server address using api instead
	stateInfo, _, err := environ.StateInfo()
	if err != nil {
		return "", err
	}
	return strings.Split(stateInfo.Addrs[0], ":")[0], nil
}

var agentAddressTemplate = `
set -exu
cd /var/lib/juju/agents
for agent in *
do
	initctl stop jujud-$agent
	sed -i.old -r "/^(stateaddresses|apiaddresses):/{
		n
		s/- .*(:[0-9]+)/- $ADDR\1/
	}" $agent/agent.conf
	if [[ $agent = unit-* ]]
	then
		sed -i -r 's/change-version: [0-9]+$/change-version: 0/' $agent/state/relations/*/*
	fi
	initctl start jujud-$agent
done
sed -i -r 's/^(:syslogtag, startswith, "juju-" @)(.*)(:[0-9]+.*)$/\1'$ADDR'\3/' /etc/rsyslog.d/*-juju*.conf
`

// renderScriptArg generates an ssh script argument to update state addresses
func renderScriptArg(stateAddr string) string {
	script := strings.Replace(agentAddressTemplate, "$ADDR", stateAddr, -1)
	return "sudo bash -c " + utils.ShQuote(script)
}

// runMachineUpdate connects via ssh to the machine and runs the update script
func runMachineUpdate(m *state.Machine, sshArg string) error {
	logger.Infof("updating machine: %v\n", m)
	addr := instance.SelectPublicAddress(m.Addresses())
	if addr == "" {
		return fmt.Errorf("no appropriate public address found")
	}
	args := []string{
		"-l", "ubuntu",
		"-T",
		"-o", "StrictHostKeyChecking no",
		"-o", "PasswordAuthentication no",
		addr,
		sshArg,
	}
	c := exec.Command("ssh", args...)
	if data, err := c.CombinedOutput(); err != nil {
		return fmt.Errorf("ssh command failed: %v (%q)", err, data)
	}
	return nil
}

// updateAllMachines finds all machines resets the stored state address
func updateAllMachines(conn *juju.Conn, stateAddr string) error {
	machines, err := conn.State.AllMachines()
	if err != nil {
		return err
	}
	pendingMachineCount := 0
	done := make(chan error)
	for _, machine := range machines {
		// A newly resumed state server requires no updating, and more
		// than one state server is not yet support by this plugin.
		if machine.IsManager() || machine.Life() != state.Alive {
			continue
		}
		pendingMachineCount += 1
		machine := machine
		go func() {
			err := runMachineUpdate(machine, renderScriptArg(stateAddr))
			if err != nil {
				logger.Errorf("failed to update machine %s: %v", machine, err)
			} else {
				logger.Infof("updated machine %s", machine)
			}
			done <- err
		}()
	}
	err = nil
	for ; pendingMachineCount > 0; pendingMachineCount-- {
		if updateErr := <-done; updateErr != nil && err == nil {
			err = fmt.Errorf("machine update failed")
		}
	}
	return err
}

func Main(args []string) {
	if err := juju.InitJujuHome(); err != nil {
		fmt.Fprintf(os.Stderr, "error: %s\n", err)
		os.Exit(2)
	}
	command := updateBootstrapCommand{}
	os.Exit(cmd.Main(&command, cmd.DefaultContext(), args[1:]))
}

func main() {
	Main(os.Args)
}
<commit_msg>[r=gz] Merge 1.16 branch into trunk<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package main

import (
	"fmt"
	"os"
	"os/exec"
	"strings"

	"launchpad.net/loggo"

	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/instance"
	"launchpad.net/juju-core/juju"
	_ "launchpad.net/juju-core/provider/all"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/utils"
)

var logger = loggo.GetLogger("juju.plugins.updatebootstrap")

const updateBootstrapDoc = `
Patches all machines after state server has been restored from backup, to
update state server address to new location.
`

type updateBootstrapCommand struct {
	cmd.EnvCommandBase
}

func (c *updateBootstrapCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "juju-update-bootstrap",
		Purpose: "update all machines after recovering state server",
		Doc:     updateBootstrapDoc,
	}
}

func (c *updateBootstrapCommand) Run(ctx *cmd.Context) error {
	conn, err := juju.NewConnFromName(c.EnvName)
	if err != nil {
		return err
	}
	defer conn.Close()
	stateAddr, err := GetStateAddress(conn.Environ)
	if err != nil {
		return err
	}
	fmt.Printf("using state address %v\n", stateAddr)
	return updateAllMachines(conn, stateAddr)
}

// GetStateAddress returns the address of one state server
func GetStateAddress(environ environs.Environ) (string, error) {
	// XXX: Can easily look up state server address using api instead
	stateInfo, _, err := environ.StateInfo()
	if err != nil {
		return "", err
	}
	return strings.Split(stateInfo.Addrs[0], ":")[0], nil
}

var agentAddressTemplate = `
set -exu
cd /var/lib/juju/agents
for agent in *
do
	initctl stop jujud-$agent
	sed -i.old -r "/^(stateaddresses|apiaddresses):/{
		n
		s/- .*(:[0-9]+)/- $ADDR\1/
	}" $agent/agent.conf
	if [[ $agent = unit-* ]]
	then
		sed -i -r 's/change-version: [0-9]+$/change-version: 0/' $agent/state/relations/*/* || true
	fi
	initctl start jujud-$agent
done
sed -i -r 's/^(:syslogtag, startswith, "juju-" @)(.*)(:[0-9]+.*)$/\1'$ADDR'\3/' /etc/rsyslog.d/*-juju*.conf
`

// renderScriptArg generates an ssh script argument to update state addresses
func renderScriptArg(stateAddr string) string {
	script := strings.Replace(agentAddressTemplate, "$ADDR", stateAddr, -1)
	return "sudo bash -c " + utils.ShQuote(script)
}

// runMachineUpdate connects via ssh to the machine and runs the update script
func runMachineUpdate(m *state.Machine, sshArg string) error {
	logger.Infof("updating machine: %v\n", m)
	addr := instance.SelectPublicAddress(m.Addresses())
	if addr == "" {
		return fmt.Errorf("no appropriate public address found")
	}
	args := []string{
		"-l", "ubuntu",
		"-T",
		"-o", "StrictHostKeyChecking no",
		"-o", "PasswordAuthentication no",
		addr,
		sshArg,
	}
	c := exec.Command("ssh", args...)
	if data, err := c.CombinedOutput(); err != nil {
		return fmt.Errorf("ssh command failed: %v (%q)", err, data)
	}
	return nil
}

// updateAllMachines finds all machines resets the stored state address
func updateAllMachines(conn *juju.Conn, stateAddr string) error {
	machines, err := conn.State.AllMachines()
	if err != nil {
		return err
	}
	pendingMachineCount := 0
	done := make(chan error)
	for _, machine := range machines {
		// A newly resumed state server requires no updating, and more
		// than one state server is not yet support by this plugin.
		if machine.IsManager() || machine.Life() == state.Dead {
			continue
		}
		pendingMachineCount++
		machine := machine
		go func() {
			err := runMachineUpdate(machine, renderScriptArg(stateAddr))
			if err != nil {
				logger.Errorf("failed to update machine %s: %v", machine, err)
			} else {
				logger.Infof("updated machine %s", machine)
			}
			done <- err
		}()
	}
	err = nil
	for ; pendingMachineCount > 0; pendingMachineCount-- {
		if updateErr := <-done; updateErr != nil && err == nil {
			err = fmt.Errorf("machine update failed")
		}
	}
	return err
}

func Main(args []string) {
	if err := juju.InitJujuHome(); err != nil {
		fmt.Fprintf(os.Stderr, "error: %s\n", err)
		os.Exit(2)
	}
	command := updateBootstrapCommand{}
	os.Exit(cmd.Main(&command, cmd.DefaultContext(), args[1:]))
}

func main() {
	Main(os.Args)
}
<|endoftext|>