<commit_before>package unilog

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"os"
	"os/exec"
	"os/signal"
	"strings"
	"syscall"
	"text/template"
	"time"

	"golang.org/x/crypto/ssh/terminal"

	"github.com/DataDog/datadog-go/statsd"
	"github.com/getsentry/raven-go"

	"github.com/stripe/unilog/clevels"
	flag "launchpad.net/gnuflag"
)

// hold the argument passed in with "-statstags"
var statstags string

// A filter to be applied to log lines prior to prefixing them
// with a timestamp and logging them.
type Filter func(string) string

// Unilog represents a unilog process. unilog is intended to be used
// as a standalone application, but is exported as a package to allow
// users to perform compile-time configuration to simplify deployment.
type Unilog struct {
	// Sentry DSN for reporting Unilog errors
	// If this is unset, unilog will not report errors to Sentry
	SentryDSN string
	// StatsdAddress for sending metrics
	// If this is unset, it wlil default to "127.0.0.1:8200" -> TODO: is this what we want?
	StatsdAddress string
	// The email address from which unilog will send mail on
	// errors
	MailTo string
	// The email address to which unilog will email breakages. If
	// either MailTo or MailFrom is unset, unilog will not
	// generate email.
	MailFrom string

	// A series of filters which will be applied to each log line
	// in order
	Filters []Filter

	// The version that unilog will report on the command-line and
	// in error emails. Defaults to the toplevel Version constant.
	Version string
	// The number of log lines to buffer in-memory, in case
	// unilog's disk writer falls behind. Note that when talking
	// to unilog over a pipe, the kernel also maintains an
	// in-kernel pipe buffer, sized 64kb on Linux.
	BufferLines int

	Name    string
	Verbose bool

	lines     <-chan string
	errs      <-chan error
	sigReopen <-chan os.Signal
	sigTerm   <-chan os.Signal
	shutdown  chan struct{}
	file      io.WriteCloser
	target    string

	b struct {
		broken bool
		at     time.Time
		count  int
	}
}

func stringFlag(val *string, longname, shortname, init, help string) {
	flag.StringVar(val, longname, init, help)
	flag.StringVar(val, shortname, init, help)
}

func boolFlag(val *bool, longname, shortname string, init bool, help string) {
	flag.BoolVar(val, longname, init, help)
	flag.BoolVar(val, shortname, init, help)
}

func (u *Unilog) fillDefaults() {
	if u.Version == "" {
		u.Version = Version
	}
	if u.BufferLines == 0 {
		u.BufferLines = DefaultBuffer
	}
}

func (u *Unilog) addFlags() {
	stringFlag(&u.Name, "name", "a", "", "Name of logged program")
	boolFlag(&u.Verbose, "verbose", "v", false, "Echo lines to stdout")
	flag.StringVar(&u.MailFrom, "mailfrom", u.MailFrom, "Address to send error emails from")
	flag.StringVar(&u.MailTo, "mailto", u.MailTo, "Address to send error emails to")
	flag.StringVar(&u.SentryDSN, "sentrydsn", u.SentryDSN, "Sentry DSN to send errors to")
	flag.StringVar(&u.StatsdAddress, "statsdaddress", "127.0.0.1:8200", "Address to send statsd metrics to")
	flag.StringVar(&clevels.AusterityFile, "austerityfile", clevels.AusterityFile, "(optional) Location of file to read austerity level from")
	stringFlag(&statstags, "statstags", "s", "", `(optional) tags to include with all statsd metrics (e.g. "foo:bar,baz:quz")`)
}

var emailTemplate = template.Must(template.New("email").Parse(`From: {{.From}}
To: {{.To}}
Subject: [unilog] {{.Name}} could not {{.Action}}

Hi there,

This is unilog reporting from {{.Hostname}}. I'm sad to report that
{{.Name}} is having some troubles writing to its log. I got caught up
trying to log a line to {{.Target}}.

To avoid spamming you, I'm going to shut up for an hour. Please fix me.

{{.Error}}
--
Sent from unilog {{.Version}}
`))

const (
	// Version is the Unilog version. Reported in emails and in
	// response to --version on the command line. Can be overriden
	// by the Version field in a Unilog object.
	Version = "0.3"
	// DefaultBuffer is the default size (in lines) of the
	// in-process line buffer
	DefaultBuffer = 1 << 12
)

var Stats *statsd.Client

func readlines(in io.Reader, bufsize int, shutdown chan struct{}) (<-chan string, <-chan error) {
	linec := make(chan string, bufsize)
	errc := make(chan error, 1)

	u := NewUnilogReader(in, shutdown)
	r := bufio.NewReader(u)

	go func() {
		var err error
		var s string

		for err == nil {
			s, err = r.ReadString('\n')
			if s != "" {
				s = strings.TrimRight(s, "\n")
				linec <- s
				if Stats != nil {
					Stats.Count("unilog.bytes", int64(len(s)), nil, .1)
				}
			}
		}

		if err != io.EOF {
			errc <- err
		}
		close(linec)
	}()

	return linec, errc
}

func (u *Unilog) reopen() error {
	if u.target == "-" {
		u.file = os.Stdout
		return nil
	}

	if u.file != nil {
		u.file.Close()
		u.file = nil
	}

	var e error
	if u.file, e = os.OpenFile(u.target, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0644); e != nil {
		u.file = nil
		return e
	}
	return nil
}

func (u *Unilog) format(line string) string {
	for _, filter := range u.Filters {
		if filter != nil {
			line = filter(line)
		}
	}
	return fmt.Sprintf("[%s] %s\n", time.Now().Format("2006-01-02 15:04:05.000000"), line)
}

func (u *Unilog) logLine(line string) {
	formatted := u.format(line)
	if u.Verbose {
		defer io.WriteString(os.Stdout, formatted)
	}

	var e error
	if u.file == nil {
		e = u.reopen()
	}
	if e != nil {
		u.handleError("reopen_file", e)
		return
	}
	_, e = io.WriteString(u.file, formatted)
	if e != nil {
		u.handleError("write_to_log", e)
	} else {
		u.b.broken = false
	}
}

func (u *Unilog) run() {
	for {
		select {
		case e := <-u.errs:
			if e != nil && e != io.EOF {
				panic(e)
			} else {
				return
			}
		case <-u.sigReopen:
			u.reopen()
		case <-u.sigTerm:
			if u.shutdown != nil {
				close(u.shutdown)
				u.shutdown = nil
			}
		case line, ok := <-u.lines:
			if !ok {
				return
			}
			u.logLine(line)
		}
	}
}

func (u *Unilog) handleError(action string, e error) {
	if !u.b.broken {
		u.b.broken = true
		u.b.at = time.Now()
		u.b.count = 0
	} else if time.Since(u.b.at) > time.Hour {
		u.b.at = time.Now()
		u.b.count = 0
	}

	message := fmt.Sprintf("Could not %s: %s", action, e.Error())

	if terminal.IsTerminal(1) {
		fmt.Printf("%s\n", message)
		return
	}

	if Stats != nil {
		emsg := fmt.Sprintf("err_action:%s", action)
		Stats.Count("unilog.errors_total", 1, []string{emsg}, 1)
	}

	if u.b.count == 0 && u.SentryDSN != "" {
		hostname, _ := os.Hostname()
		keys := map[string]string{
			"Hostname": hostname,
			"Action":   action,
			"Name":     u.Name,
			"Target":   u.target,
			"Error":    e.Error(),
			"Version":  Version,
		}
		raven.CaptureError(e, keys)
	}

	if u.b.count == 0 && u.MailFrom != "" && u.MailTo != "" {
		message := new(bytes.Buffer)
		hostname, _ := os.Hostname()
		emailTemplate.Execute(message, map[string]string{
			"Hostname": hostname,
			"From":     u.MailFrom,
			"To":       u.MailTo,
			"Action":   action,
			"Name":     u.Name,
			"Target":   u.target,
			"Error":    e.Error(),
			"Version":  Version,
		})
		cmd := exec.Command("sendmail", "-t")
		cmd.Stdin = message
		cmd.Run()
	}

	u.b.count++
}

func (u *Unilog) Main() {
	u.fillDefaults()

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [options] dstfile\n", os.Args[0])
		flag.PrintDefaults()
	}

	u.addFlags()
	var flagVersion bool
	boolFlag(&flagVersion, "version", "V", false, "Print the version number and exit")

	flag.Parse(true)

	if flagVersion {
		fmt.Printf("This is unilog v%s\n", Version)
		return
	}
	args := flag.Args()
	if len(args) != 1 {
		flag.Usage()
		os.Exit(1)
	}

	reopen := make(chan os.Signal, 2)
	signal.Notify(reopen, syscall.SIGALRM, syscall.SIGHUP)
	u.sigReopen = reopen

	term := make(chan os.Signal, 2)
	signal.Notify(term, syscall.SIGTERM, syscall.SIGINT)
	u.sigTerm = term

	u.shutdown = make(chan struct{})
	u.target = flag.Arg(0)
	u.reopen()

	u.lines, u.errs = readlines(os.Stdin, u.BufferLines, u.shutdown)

	fileName := u.target

	Stats, _ = statsd.New(u.StatsdAddress)

	Stats.Tags = append(Stats.Tags, fmt.Sprintf("file_name:%s", fileName))
	if statstags != "" {
		Stats.Tags = append(Stats.Tags, strings.Split(statstags, ",")...)
	}

	clevels.Stats = Stats

	_ = raven.SetDSN(u.SentryDSN)

	u.run()
}
<commit_msg>Changed logging to debug flag<commit_after>package unilog

import (
	"bufio"
	"bytes"
	"fmt"
	"io"
	"os"
	"os/exec"
	"os/signal"
	"strings"
	"syscall"
	"text/template"
	"time"

	"github.com/DataDog/datadog-go/statsd"
	"github.com/getsentry/raven-go"

	"github.com/stripe/unilog/clevels"
	flag "launchpad.net/gnuflag"
)

// hold the argument passed in with "-statstags"
var statstags string

// A filter to be applied to log lines prior to prefixing them
// with a timestamp and logging them.
type Filter func(string) string

// Unilog represents a unilog process. unilog is intended to be used
// as a standalone application, but is exported as a package to allow
// users to perform compile-time configuration to simplify deployment.
type Unilog struct {
	// Sentry DSN for reporting Unilog errors
	// If this is unset, unilog will not report errors to Sentry
	SentryDSN string
	// StatsdAddress for sending metrics
	// If this is unset, it wlil default to "127.0.0.1:8200" -> TODO: is this what we want?
	StatsdAddress string
	// The email address from which unilog will send mail on
	// errors
	MailTo string
	// The email address to which unilog will email breakages. If
	// either MailTo or MailFrom is unset, unilog will not
	// generate email.
	MailFrom string

	// A series of filters which will be applied to each log line
	// in order
	Filters []Filter

	// The version that unilog will report on the command-line and
	// in error emails. Defaults to the toplevel Version constant.
	Version string
	// The number of log lines to buffer in-memory, in case
	// unilog's disk writer falls behind. Note that when talking
	// to unilog over a pipe, the kernel also maintains an
	// in-kernel pipe buffer, sized 64kb on Linux.
	BufferLines int

	Name    string
	Verbose bool
	Debug   bool

	lines     <-chan string
	errs      <-chan error
	sigReopen <-chan os.Signal
	sigTerm   <-chan os.Signal
	shutdown  chan struct{}
	file      io.WriteCloser
	target    string

	b struct {
		broken bool
		at     time.Time
		count  int
	}
}

func stringFlag(val *string, longname, shortname, init, help string) {
	flag.StringVar(val, longname, init, help)
	flag.StringVar(val, shortname, init, help)
}

func boolFlag(val *bool, longname, shortname string, init bool, help string) {
	flag.BoolVar(val, longname, init, help)
	flag.BoolVar(val, shortname, init, help)
}

func (u *Unilog) fillDefaults() {
	if u.Version == "" {
		u.Version = Version
	}
	if u.BufferLines == 0 {
		u.BufferLines = DefaultBuffer
	}
}

func (u *Unilog) addFlags() {
	stringFlag(&u.Name, "name", "a", "", "Name of logged program")
	boolFlag(&u.Verbose, "verbose", "v", false, "Echo lines to stdout")
	boolFlag(&u.Debug, "debug", "d", false, "Print debug messages")
	flag.StringVar(&u.MailFrom, "mailfrom", u.MailFrom, "Address to send error emails from")
	flag.StringVar(&u.MailTo, "mailto", u.MailTo, "Address to send error emails to")
	flag.StringVar(&u.SentryDSN, "sentrydsn", u.SentryDSN, "Sentry DSN to send errors to")
	flag.StringVar(&u.StatsdAddress, "statsdaddress", "127.0.0.1:8200", "Address to send statsd metrics to")
	flag.StringVar(&clevels.AusterityFile, "austerityfile", clevels.AusterityFile, "(optional) Location of file to read austerity level from")
	stringFlag(&statstags, "statstags", "s", "", `(optional) tags to include with all statsd metrics (e.g. "foo:bar,baz:quz")`)
}

var emailTemplate = template.Must(template.New("email").Parse(`From: {{.From}}
To: {{.To}}
Subject: [unilog] {{.Name}} could not {{.Action}}

Hi there,

This is unilog reporting from {{.Hostname}}. I'm sad to report that
{{.Name}} is having some troubles writing to its log. I got caught up
trying to log a line to {{.Target}}.

To avoid spamming you, I'm going to shut up for an hour. Please fix me.

{{.Error}}
--
Sent from unilog {{.Version}}
`))

const (
	// Version is the Unilog version. Reported in emails and in
	// response to --version on the command line. Can be overriden
	// by the Version field in a Unilog object.
	Version = "0.3"
	// DefaultBuffer is the default size (in lines) of the
	// in-process line buffer
	DefaultBuffer = 1 << 12
)

var Stats *statsd.Client

func readlines(in io.Reader, bufsize int, shutdown chan struct{}) (<-chan string, <-chan error) {
	linec := make(chan string, bufsize)
	errc := make(chan error, 1)

	u := NewUnilogReader(in, shutdown)
	r := bufio.NewReader(u)

	go func() {
		var err error
		var s string

		for err == nil {
			s, err = r.ReadString('\n')
			if s != "" {
				s = strings.TrimRight(s, "\n")
				linec <- s
				if Stats != nil {
					Stats.Count("unilog.bytes", int64(len(s)), nil, .1)
				}
			}
		}

		if err != io.EOF {
			errc <- err
		}
		close(linec)
	}()

	return linec, errc
}

func (u *Unilog) reopen() error {
	if u.target == "-" {
		u.file = os.Stdout
		return nil
	}

	if u.file != nil {
		u.file.Close()
		u.file = nil
	}

	var e error
	if u.file, e = os.OpenFile(u.target, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0644); e != nil {
		u.file = nil
		return e
	}
	return nil
}

func (u *Unilog) format(line string) string {
	for _, filter := range u.Filters {
		if filter != nil {
			line = filter(line)
		}
	}
	return fmt.Sprintf("[%s] %s\n", time.Now().Format("2006-01-02 15:04:05.000000"), line)
}

func (u *Unilog) logLine(line string) {
	formatted := u.format(line)
	if u.Verbose {
		defer io.WriteString(os.Stdout, formatted)
	}

	var e error
	if u.file == nil {
		e = u.reopen()
	}
	if e != nil {
		u.handleError("reopen_file", e)
		return
	}
	_, e = io.WriteString(u.file, formatted)
	if e != nil {
		u.handleError("write_to_log", e)
	} else {
		u.b.broken = false
	}
}

func (u *Unilog) run() {
	for {
		select {
		case e := <-u.errs:
			if e != nil && e != io.EOF {
				panic(e)
			} else {
				return
			}
		case <-u.sigReopen:
			u.reopen()
		case <-u.sigTerm:
			if u.shutdown != nil {
				close(u.shutdown)
				u.shutdown = nil
			}
		case line, ok := <-u.lines:
			if !ok {
				return
			}
			u.logLine(line)
		}
	}
}

func (u *Unilog) handleError(action string, e error) {
	if !u.b.broken {
		u.b.broken = true
		u.b.at = time.Now()
		u.b.count = 0
	} else if time.Since(u.b.at) > time.Hour {
		u.b.at = time.Now()
		u.b.count = 0
	}

	if u.Debug {
		fmt.Fprintf(os.Stderr, "Could not %s: %s\n", action, e.Error())
	}

	if Stats != nil {
		emsg := fmt.Sprintf("err_action:%s", action)
		Stats.Count("unilog.errors_total", 1, []string{emsg}, 1)
	}

	if u.b.count == 0 && u.SentryDSN != "" {
		hostname, _ := os.Hostname()
		keys := map[string]string{
			"Hostname": hostname,
			"Action":   action,
			"Name":     u.Name,
			"Target":   u.target,
			"Error":    e.Error(),
			"Version":  Version,
		}
		raven.CaptureError(e, keys)
	}

	if u.b.count == 0 && u.MailFrom != "" && u.MailTo != "" {
		message := new(bytes.Buffer)
		hostname, _ := os.Hostname()
		emailTemplate.Execute(message, map[string]string{
			"Hostname": hostname,
			"From":     u.MailFrom,
			"To":       u.MailTo,
			"Action":   action,
			"Name":     u.Name,
			"Target":   u.target,
			"Error":    e.Error(),
			"Version":  Version,
		})
		cmd := exec.Command("sendmail", "-t")
		cmd.Stdin = message
		cmd.Run()
	}

	u.b.count++
}

func (u *Unilog) Main() {
	u.fillDefaults()

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [options] dstfile\n", os.Args[0])
		flag.PrintDefaults()
	}

	u.addFlags()
	var flagVersion bool
	boolFlag(&flagVersion, "version", "V", false, "Print the version number and exit")

	flag.Parse(true)

	if flagVersion {
		fmt.Printf("This is unilog v%s\n", Version)
		return
	}
	args := flag.Args()
	if len(args) != 1 {
		flag.Usage()
		os.Exit(1)
	}

	reopen := make(chan os.Signal, 2)
	signal.Notify(reopen, syscall.SIGALRM, syscall.SIGHUP)
	u.sigReopen = reopen

	term := make(chan os.Signal, 2)
	signal.Notify(term, syscall.SIGTERM, syscall.SIGINT)
	u.sigTerm = term

	u.shutdown = make(chan struct{})
	u.target = flag.Arg(0)
	u.reopen()

	u.lines, u.errs = readlines(os.Stdin, u.BufferLines, u.shutdown)

	fileName := u.target

	Stats, _ = statsd.New(u.StatsdAddress)

	Stats.Tags = append(Stats.Tags, fmt.Sprintf("file_name:%s", fileName))
	if statstags != "" {
		Stats.Tags = append(Stats.Tags, strings.Split(statstags, ",")...)
	}

	clevels.Stats = Stats

	_ = raven.SetDSN(u.SentryDSN)

	u.run()
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"io"
	"runtime"
	"sync"
	"time"
)

type M map[string]interface{}

type Request struct {
	Method string
	Token  string
}

type Response struct {
	Token string      `json:"token"`
	Error string      `json:"error"`
	Data  interface{} `json:"data"`
}

type Broker struct {
	served uint
	start  time.Time
	rLck   sync.Mutex
	wLck   sync.Mutex
	r      io.Reader
	w      io.Writer
	in     *bufio.Reader
	out    *json.Encoder
}

func NewBroker(r io.Reader, w io.Writer) *Broker {
	return &Broker{
		r:   r,
		w:   w,
		in:  bufio.NewReader(r),
		out: json.NewEncoder(w),
	}
}

func (b *Broker) Send(resp Response) error {
	err := b.SendNoLog(resp)
	if err != nil {
		logger.Println("Cannot send result", err)
	}
	return err
}

func (b *Broker) SendNoLog(resp Response) error {
	b.wLck.Lock()
	defer b.wLck.Unlock()

	if resp.Data == nil {
		resp.Data = M{}
	}

	s, err := json.Marshal(resp)
	if err != nil {
		return err
	}

	// the only expected write failure are due to broken pipes
	// which usually means the client has gone away so just ignore the error
	b.w.Write(s)
	b.w.Write([]byte{'\n'})
	return nil
}

func (b *Broker) call(req *Request, cl Caller) {
	b.served++

	defer func() {
		err := recover()
		if err != nil {
			logger.Printf("%v#%v PANIC: %v\n", req.Method, req.Token, err)
			b.Send(Response{
				Token: req.Token,
				Error: "broker: " + req.Method + "#" + req.Token + " PANIC",
			})
		}
	}()

	res, err := cl.Call()
	b.Send(Response{
		Token: req.Token,
		Error: err,
		Data:  res,
	})
}

func (b *Broker) accept() (stopLooping bool) {
	line, err := b.in.ReadBytes('\n')

	if err == io.EOF {
		stopLooping = true
	} else if err != nil {
		logger.Println("Cannot read input", err)
		b.Send(Response{
			Error: err.Error(),
		})
		return
	}

	req := &Request{}
	dec := json.NewDecoder(bytes.NewReader(line))
	// if this fails, we are unable to return a useful error(no token to send it to)
	// so we'll simply/implicitly drop the request since it has no method
	// we can safely assume that all such cases will be empty lines and not an actual request
	dec.Decode(&req)

	if req.Method == "" {
		return
	}

	if req.Method == "bye-ni" {
		return true
	}

	m := registry.Lookup(req.Method)
	if m == nil {
		e := "Invalid method " + req.Method
		logger.Println(e)
		b.Send(Response{
			Token: req.Token,
			Error: e,
		})
		return
	}

	cl := m(b)
	err = dec.Decode(cl)
	if err != nil {
		logger.Println("Cannot decode arg", err)
		b.Send(Response{
			Token: req.Token,
			Error: err.Error(),
		})
		return
	}

	go b.call(req, cl)

	return
}

func (b *Broker) Loop(decorate bool) {
	b.start = time.Now()

	if decorate {
		go b.SendNoLog(Response{
			Token: "margo.hello",
			Data: M{
				"time": b.start.String(),
			},
		})
	}

	for {
		stopLooping := b.accept()
		if stopLooping {
			break
		}
		runtime.Gosched()
	}

	if decorate {
		b.SendNoLog(Response{
			Token: "margo.bye-ni",
			Data: M{
				"served": b.served,
				"uptime": time.Now().Sub(b.start).String(),
			},
		})
	}
}
<commit_msg>* allow waiting for requests to complete<commit_after>package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"io"
	"runtime"
	"sync"
	"time"
)

type M map[string]interface{}

type Request struct {
	Method string
	Token  string
}

type Response struct {
	Token string      `json:"token"`
	Error string      `json:"error"`
	Data  interface{} `json:"data"`
}

type Broker struct {
	served uint
	start  time.Time
	rLck   sync.Mutex
	wLck   sync.Mutex
	r      io.Reader
	w      io.Writer
	in     *bufio.Reader
	out    *json.Encoder
	Wg     *sync.WaitGroup
}

func NewBroker(r io.Reader, w io.Writer) *Broker {
	return &Broker{
		r:   r,
		w:   w,
		in:  bufio.NewReader(r),
		out: json.NewEncoder(w),
		Wg:  &sync.WaitGroup{},
	}
}

func (b *Broker) Send(resp Response) error {
	err := b.SendNoLog(resp)
	if err != nil {
		logger.Println("Cannot send result", err)
	}
	return err
}

func (b *Broker) SendNoLog(resp Response) error {
	b.wLck.Lock()
	defer b.wLck.Unlock()

	if resp.Data == nil {
		resp.Data = M{}
	}

	s, err := json.Marshal(resp)
	if err != nil {
		return err
	}

	// the only expected write failure are due to broken pipes
	// which usually means the client has gone away so just ignore the error
	b.w.Write(s)
	b.w.Write([]byte{'\n'})
	return nil
}

func (b *Broker) call(req *Request, cl Caller) {
	defer b.Wg.Done()
	b.served++

	defer func() {
		err := recover()
		if err != nil {
			logger.Printf("%v#%v PANIC: %v\n", req.Method, req.Token, err)
			b.Send(Response{
				Token: req.Token,
				Error: "broker: " + req.Method + "#" + req.Token + " PANIC",
			})
		}
	}()

	res, err := cl.Call()
	b.Send(Response{
		Token: req.Token,
		Error: err,
		Data:  res,
	})
}

func (b *Broker) accept() (stopLooping bool) {
	line, err := b.in.ReadBytes('\n')

	if err == io.EOF {
		stopLooping = true
	} else if err != nil {
		logger.Println("Cannot read input", err)
		b.Send(Response{
			Error: err.Error(),
		})
		return
	}

	req := &Request{}
	dec := json.NewDecoder(bytes.NewReader(line))
	// if this fails, we are unable to return a useful error(no token to send it to)
	// so we'll simply/implicitly drop the request since it has no method
	// we can safely assume that all such cases will be empty lines and not an actual request
	dec.Decode(&req)

	if req.Method == "" {
		return
	}

	if req.Method == "bye-ni" {
		return true
	}

	m := registry.Lookup(req.Method)
	if m == nil {
		e := "Invalid method " + req.Method
		logger.Println(e)
		b.Send(Response{
			Token: req.Token,
			Error: e,
		})
		return
	}

	cl := m(b)
	err = dec.Decode(cl)
	if err != nil {
		logger.Println("Cannot decode arg", err)
		b.Send(Response{
			Token: req.Token,
			Error: err.Error(),
		})
		return
	}

	b.Wg.Add(1)
	go b.call(req, cl)

	return
}

func (b *Broker) Loop(decorate bool) {
	b.start = time.Now()

	if decorate {
		go b.SendNoLog(Response{
			Token: "margo.hello",
			Data: M{
				"time": b.start.String(),
			},
		})
	}

	for {
		stopLooping := b.accept()
		if stopLooping {
			break
		}
		runtime.Gosched()
	}

	if decorate {
		b.SendNoLog(Response{
			Token: "margo.bye-ni",
			Data: M{
				"served": b.served,
				"uptime": time.Now().Sub(b.start).String(),
			},
		})
	}
}
<|endoftext|>


<commit_before>// Copyright ©2013 The gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package mat64

import (
	//"fmt"
	check "launchpad.net/gocheck"
	"math"

	"github.com/gonum/floats"
)

func isUpperTriangular(a *Dense) bool {
	rows, cols := a.Dims()
	for c := 0; c < cols-1; c++ {
		for r := c + 1; r < rows; r++ {
			if math.Abs(a.At(r, c)) > 1e-14 {
				return false
			}
		}
	}
	return true
}

func isOrthogonal(a *Dense) bool {
	rows, cols := a.Dims()
	col1 := make([]float64, rows)
	col2 := make([]float64, rows)
	for i := 0; i < cols-1; i++ {
		for j := i + 1; j < cols; j++ {
			a.Col(col1, i)
			a.Col(col2, j)
			dot := floats.Dot(col1, col2)
			if math.Abs(dot) > 1e-14 {
				return false
			}
		}
	}
	return true
}

func (s *S) TestQRD(c *check.C) {
	for _, test := range []struct {
		a    [][]float64
		name string
	}{
		{
			name: "Square",
			a: [][]float64{
				{1.3, 2.4, 8.9},
				{-2.6, 8.7, 9.1},
				{5.6, 5.8, 2.1},
			},
		},
		{
			name: "Skinny",
			a: [][]float64{
				{1.3, 2.4, 8.9},
				{-2.6, 8.7, 9.1},
				{5.6, 5.8, 2.1},
				{19.4, 5.2, -26.1},
			},
		},
		{
			name: "Fat",
			a: [][]float64{
				{1.3, 2.4, 8.9, 19.4},
				{-2.6, 8.7, 9.1, 5.2},
				{5.6, 5.8, 2.1, -26.1},
			},
		},
	} {

		a := NewDense(flatten(test.a))
		qf := QR(DenseCopyOf(a))
		r := qf.R()
		q := qf.Q()

		rows, cols := a.Dims()
		newA := NewDense(rows, cols, nil)
		newA.Mul(q, r)

		c.Check(isOrthogonal(q), check.Equals, true, check.Commentf("Test %v: Q not orthogonal", test.name))
		c.Check(isUpperTriangular(r), check.Equals, true, check.Commentf("Test %v: R not upper triangular", test.name))
		c.Check(a.EqualsApprox(newA, 1e-14), check.Equals, true, check.Commentf("Test %v: Q*R != A", test.name))
	}
}
<commit_msg>Removed fat test, algorithm doesn't work in that case<commit_after>// Copyright ©2013 The gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package mat64

import (
	//"fmt"
	check "launchpad.net/gocheck"
	"math"

	"github.com/gonum/floats"
)

func isUpperTriangular(a *Dense) bool {
	rows, cols := a.Dims()
	for c := 0; c < cols-1; c++ {
		for r := c + 1; r < rows; r++ {
			if math.Abs(a.At(r, c)) > 1e-14 {
				return false
			}
		}
	}
	return true
}

func isOrthogonal(a *Dense) bool {
	rows, cols := a.Dims()
	col1 := make([]float64, rows)
	col2 := make([]float64, rows)
	for i := 0; i < cols-1; i++ {
		for j := i + 1; j < cols; j++ {
			a.Col(col1, i)
			a.Col(col2, j)
			dot := floats.Dot(col1, col2)
			if math.Abs(dot) > 1e-14 {
				return false
			}
		}
	}
	return true
}

func (s *S) TestQRD(c *check.C) {
	for _, test := range []struct {
		a    [][]float64
		name string
	}{
		{
			name: "Square",
			a: [][]float64{
				{1.3, 2.4, 8.9},
				{-2.6, 8.7, 9.1},
				{5.6, 5.8, 2.1},
			},
		},
		{
			name: "Skinny",
			a: [][]float64{
				{1.3, 2.4, 8.9},
				{-2.6, 8.7, 9.1},
				{5.6, 5.8, 2.1},
				{19.4, 5.2, -26.1},
			},
		},
		/*
			{
				name: "Fat",
				a: [][]float64{
					{1.3, 2.4, 8.9, 19.4},
					{-2.6, 8.7, 9.1, 5.2},
					{5.6, 5.8, 2.1, -26.1},
				},
			},
		*/
	} {

		a := NewDense(flatten(test.a))
		qf := QR(DenseCopyOf(a))
		r := qf.R()
		q := qf.Q()

		rows, cols := a.Dims()
		newA := NewDense(rows, cols, nil)
		newA.Mul(q, r)

		c.Check(isOrthogonal(q), check.Equals, true, check.Commentf("Test %v: Q not orthogonal", test.name))
		c.Check(isUpperTriangular(r), check.Equals, true, check.Commentf("Test %v: R not upper triangular", test.name))
		c.Check(a.EqualsApprox(newA, 1e-14), check.Equals, true, check.Commentf("Test %v: Q*R != A", test.name))
	}
}
<|endoftext|>


<commit_before>package main

import (
	"compress/gzip"
	"crypto/sha1"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"syscall"
	"time"

	"github.com/dickeyxxx/golock"
)

var updateTopic = &Topic{
	Name:        "update",
	Description: "update heroku-cli",
}

var updateCmd = &Command{
	Topic:       "update",
	Description: "updates heroku-cli",
	Args:        []Arg{{Name: "channel", Optional: true}},
	Run: func(ctx *Context) {
		channel := ctx.Args.(map[string]string)["channel"]
		if channel == "" {
			channel = Channel
		}
		Update(channel)
	},
}

var binPath = filepath.Join(AppDir, "heroku-cli")
var updateLockPath = filepath.Join(AppDir, "updating.lock")
var autoupdateFile = filepath.Join(AppDir, "autoupdate")

func init() {
	if runtime.GOOS == "windows" {
		binPath = binPath + ".exe"
	}
}

// Update updates the CLI and plugins
func Update(channel string) {
	touchAutoupdateFile()
	if err := golock.Lock(updateLockPath); err != nil {
		panic(err)
	}
	defer golock.Unlock(updateLockPath)
	done := make(chan bool)
	go func() {
		updateCLI(channel)
		updatePlugins()
		done <- true
	}()
	select {
	case <-time.After(time.Second * 120):
		Errln("Timed out while updating")
	case <-done:
	}
}

func updatePlugins() {
	Err("updating plugins... ")
	b, _ := node.UpdatePackages()
	Errln("done")
	if len(b) > 0 {
		Err("rebuilding plugins cache... ")
		ClearPluginCache()
		WritePluginCache(GetPlugins())
		Errln("done")
	}
}

func updateCLI(channel string) {
	Err("checking for CLI update... ")
	manifest := getUpdateManifest(channel)
	if manifest.Version == Version && manifest.Channel == Channel {
		Errf("already on latest version %s\n", Version)
		return
	}
	if !updatable() {
		Errf("Out of date: You are running %s but %s is out\n", Version, manifest.Version)
		return
	}
	Errf("updating from %s to %s (%s)... ", Version, manifest.Version, manifest.Channel)
	build := manifest.Builds[runtime.GOOS][runtime.GOARCH]
	// on windows we can't remove an existing file or remove the running binary
	// so we download the file to binName.new
	// move the running binary to binName.old (deleting any existing file first)
	// rename the downloaded file to binName
	if err := downloadBin(binPath+".new", build.URL); err != nil {
		panic(err)
	}
	if fileSha1(binPath+".new") != build.Sha1 {
		panic("SHA mismatch")
	}
	os.Remove(binPath + ".old")
	if err := os.Rename(binPath, binPath+".old"); err != nil {
		panic(err)
	}
	if err := os.Rename(binPath+".new", binPath); err != nil {
		panic(err)
	}
	Errln("done")
	reexecBin()
}

// IsUpdateNeeded checks if an update is available
func IsUpdateNeeded(t string) bool {
	f, err := os.Stat(autoupdateFile)
	if err != nil {
		return true
	}
	if t == "soft" {
		return time.Since(f.ModTime()) > 1*time.Hour
	}
	return time.Since(f.ModTime()) > 168*time.Hour
}

func touchAutoupdateFile() {
	out, err := os.OpenFile(autoupdateFile, os.O_WRONLY|os.O_CREATE, 0644)
	if err != nil {
		panic(err)
	}
	out.WriteString(time.Now().String())
}

type manifest struct {
	Channel, Version string
	Builds           map[string]map[string]struct {
		URL, Sha1 string
	}
}

func getUpdateManifest(channel string) manifest {
	res, err := http.Get("https://d1gvo455cekpjp.cloudfront.net/" + channel + "/manifest.json")
	if err != nil {
		panic(err)
	}
	var m manifest
	json.NewDecoder(res.Body).Decode(&m)
	return m
}

func updatable() bool {
	path, err := filepath.Abs(os.Args[0])
	if err != nil {
		Errln(err)
	}
	return path == binPath
}

func downloadBin(path, url string) error {
	out, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0755)
	if err != nil {
		return err
	}
	client := &http.Client{}
	req, err := http.NewRequest("GET", url+".gz", nil)
	if err != nil {
		return err
	}
	req.Header.Add("Accept-Encoding", "gzip")
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	uncompressed, err := gzip.NewReader(resp.Body)
	if err != nil {
		return err
	}
	_, err = io.Copy(out, uncompressed)
	if err != nil {
		return err
	}
	return out.Close()
}

func fileSha1(path string) string {
	data, err := ioutil.ReadFile(path)
	if err != nil {
		panic(err)
	}
	return fmt.Sprintf("%x", sha1.Sum(data))
}

func reexecBin() {
	if runtime.GOOS == "windows" {
		cmd := exec.Command(binPath, os.Args[1:]...)
		cmd.Stdin = os.Stdin
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		cmd.Run()
		os.Exit(0)
	} else {
		if err := syscall.Exec(binPath, os.Args, os.Environ()); err != nil {
			panic(err)
		}
		os.Exit(99) // should never happen
	}
}

// TriggerBackgroundUpdate will trigger an update to the client in the background
func TriggerBackgroundUpdate() {
	if err := exec.Command(binPath, "update").Start(); err != nil {
		panic(err)
	}
}
<commit_msg>unlock lock file before reexecing<commit_after>package main

import (
	"compress/gzip"
	"crypto/sha1"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"syscall"
	"time"

	"github.com/dickeyxxx/golock"
)

var updateTopic = &Topic{
	Name:        "update",
	Description: "update heroku-cli",
}

var updateCmd = &Command{
	Topic:       "update",
	Description: "updates heroku-cli",
	Args:        []Arg{{Name: "channel", Optional: true}},
	Run: func(ctx *Context) {
		channel := ctx.Args.(map[string]string)["channel"]
		if channel == "" {
			channel = Channel
		}
		Update(channel)
	},
}

var binPath = filepath.Join(AppDir, "heroku-cli")
var updateLockPath = filepath.Join(AppDir, "updating.lock")
var autoupdateFile = filepath.Join(AppDir, "autoupdate")

func init() {
	if runtime.GOOS == "windows" {
		binPath = binPath + ".exe"
	}
}

// Update updates the CLI and plugins
func Update(channel string) {
	touchAutoupdateFile()
	if err := golock.Lock(updateLockPath); err != nil {
		panic(err)
	}
	defer golock.Unlock(updateLockPath)
	done := make(chan bool)
	go func() {
		updateCLI(channel)
		updatePlugins()
		done <- true
	}()
	select {
	case <-time.After(time.Second * 120):
		Errln("Timed out while updating")
	case <-done:
	}
}

func updatePlugins() {
	Err("updating plugins... ")
	b, _ := node.UpdatePackages()
	Errln("done")
	if len(b) > 0 {
		Err("rebuilding plugins cache... ")
		ClearPluginCache()
		WritePluginCache(GetPlugins())
		Errln("done")
	}
}

func updateCLI(channel string) {
	Err("checking for CLI update... ")
	manifest := getUpdateManifest(channel)
	if manifest.Version == Version && manifest.Channel == Channel {
		Errf("already on latest version %s\n", Version)
		return
	}
	if !updatable() {
		Errf("Out of date: You are running %s but %s is out\n", Version, manifest.Version)
		return
	}
	Errf("updating from %s to %s (%s)... ", Version, manifest.Version, manifest.Channel)
	build := manifest.Builds[runtime.GOOS][runtime.GOARCH]
	// on windows we can't remove an existing file or remove the running binary
	// so we download the file to binName.new
	// move the running binary to binName.old (deleting any existing file first)
	// rename the downloaded file to binName
	if err := downloadBin(binPath+".new", build.URL); err != nil {
		panic(err)
	}
	if fileSha1(binPath+".new") != build.Sha1 {
		panic("SHA mismatch")
	}
	os.Remove(binPath + ".old")
	if err := os.Rename(binPath, binPath+".old"); err != nil {
		panic(err)
	}
	if err := os.Rename(binPath+".new", binPath); err != nil {
		panic(err)
	}
	Errln("done")
	golock.Unlock(updateLockPath)
	reexecBin()
}

// IsUpdateNeeded checks if an update is available
func IsUpdateNeeded(t string) bool {
	f, err := os.Stat(autoupdateFile)
	if err != nil {
		return true
	}
	if t == "soft" {
		return time.Since(f.ModTime()) > 1*time.Hour
	}
	return time.Since(f.ModTime()) > 168*time.Hour
}

func touchAutoupdateFile() {
	out, err := os.OpenFile(autoupdateFile, os.O_WRONLY|os.O_CREATE, 0644)
	if err != nil {
		panic(err)
	}
	out.WriteString(time.Now().String())
}

type manifest struct {
	Channel, Version string
	Builds           map[string]map[string]struct {
		URL, Sha1 string
	}
}

func getUpdateManifest(channel string) manifest {
	res, err := http.Get("https://d1gvo455cekpjp.cloudfront.net/" + channel + "/manifest.json")
	if err != nil {
		panic(err)
	}
	var m manifest
	json.NewDecoder(res.Body).Decode(&m)
	return m
}

func updatable() bool {
	path, err := filepath.Abs(os.Args[0])
	if err != nil {
		Errln(err)
	}
	return path == binPath
}

func downloadBin(path, url string) error {
	out, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0755)
	if err != nil {
		return err
	}
	client := &http.Client{}
	req, err := http.NewRequest("GET", url+".gz", nil)
	if err != nil {
		return err
	}
	req.Header.Add("Accept-Encoding", "gzip")
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	uncompressed, err := gzip.NewReader(resp.Body)
	if err != nil {
		return err
	}
	_, err = io.Copy(out, uncompressed)
	if err != nil {
		return err
	}
	return out.Close()
}

func fileSha1(path string) string {
	data, err := ioutil.ReadFile(path)
	if err != nil {
		panic(err)
	}
	return fmt.Sprintf("%x", sha1.Sum(data))
}

func reexecBin() {
	if runtime.GOOS == "windows" {
		cmd := exec.Command(binPath, os.Args[1:]...)
		cmd.Stdin = os.Stdin
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		cmd.Run()
		os.Exit(0)
	} else {
		if err := syscall.Exec(binPath, os.Args, os.Environ()); err != nil {
			panic(err)
		}
		os.Exit(99) // should never happen
	}
}

// TriggerBackgroundUpdate will trigger an update to the client in the background
func TriggerBackgroundUpdate() {
	if err := exec.Command(binPath, "update").Start(); err != nil {
		panic(err)
	}
}
<|endoftext|>


<commit_before>/*
Package update allows a program to "self-update", replacing its executable file
with new bytes.

Package update provides the facility to create user experiences like auto-updating
or user-approved updates which manifest as user prompts in commercial applications
with copy similar to "Restart to being using the new version of X".

Updating your program to a new version is as easy as:

	err := update.FromUrl("http://release.example.com/2.0/myprogram")
	if err != nil {
		fmt.Printf("Update failed: %v", err)
	}

The most low-level API is FromStream() which updates the current executable
with the bytes read from an io.Reader.

Additional APIs are provided for common update strategies which include
updating from a file with FromFile() and updating from the internet with
FromUrl().

Using the more advaced Download.UpdateFromUrl() API gives you the ability
to resume an interrupted download to enable large updates to complete even
over intermittent or slow connections. This API also enables more fine-grained
control over how the update is downloaded from the internet as well as access to
download progress,
*/
package update

import (
	"compress/gzip"
	"fmt"
	execpath "github.com/inconshreveable/go-execpath"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"runtime"
)

var (
	// Returned when the remote server indicates that no download is available
	UpdateUnavailable error
)

func init() {
	UpdateUnavailable = fmt.Errorf("204 server response indicates no available update")
}

type MeteredReader struct {
	rd        io.ReadCloser
	totalSize int64
	progress  chan int
	totalRead int64
	ticks     int64
}

func (m *MeteredReader) Close() error {
	return m.rd.Close()
}

func (m *MeteredReader) Read(b []byte) (n int, err error) {
	chunkSize := (m.totalSize / 100) + 1
	lenB := int64(len(b))

	var nChunk int
	for start := int64(0); start < lenB; start += int64(nChunk) {
		end := start + chunkSize
		if end > lenB {
			end = lenB
		}

		nChunk, err = m.rd.Read(b[start:end])

		n += nChunk
		m.totalRead += int64(nChunk)

		if m.totalRead > (m.ticks * chunkSize) {
			m.ticks += 1
			// try to send on channel, but don't block if it's full
			select {
			case m.progress <- int(m.ticks + 1):
			default:
			}

			// give the progress channel consumer a chance to run
			runtime.Gosched()
		}

		if err != nil {
			return
		}
	}

	return
}

// We wrap the round tripper when making requests
// because we need to add headers to the requests we make
// even when they are requests made after a redirect
type RoundTripper struct {
	RoundTripFn func(*http.Request) (*http.Response, error)
}

func (rt *RoundTripper) RoundTrip(r *http.Request) (*http.Response, error) {
	return rt.RoundTripFn(r)
}

// Type Download encapsulates the necessary parameters and state
// needed to download an update from the internet. Create an instance
// with the NewDownload() factory function.
type Download struct {
	// net/http.Client to use when downloading the update.
	// If nil, a default http.Client is used
	HttpClient *http.Client

	// Path on the file system to dowload the update to
	// If empty, a temporary file is used.
	// After the download begins, this path will be set
	// so that the client can use it to resume aborted
	// downloads
	Path string

	// Progress returns the percentage of the download
	// completed as an integer between 0 and 100
	Progress chan (int)

	// HTTP Method to use in the download request. Default is "GET"
	Method string
}

// NewDownload initializes a new Download object
func NewDownload() *Download {
	return &Download{
		HttpClient: new(http.Client),
		Progress:   make(chan int),
		Method:     "GET",
	}
}

// UpdateFromUrl downloads the given url from the internet to a file on disk
// and then calls FromStream() to update the current program's executable file
// with the contents of that file.
//
// If the update is successful, the downloaded file will be erased from disk.
// Otherwise, it will remain in d.Path to allow the download to resume later
// or be skipped entirely.
//
// Only HTTP/1.1 servers that implement the Range header are supported.
//
// UpdateFromUrl() uses HTTP status codes to determine what action to take.
//
// - The HTTP server should return 200 or 206 for the update to be downloaded.
//
// - The HTTP server should return 204 if no update is available at this time.
// This will cause UpdateFromUrl to return the error UpdateUnavailable.
//
// - If the HTTP server returns a 3XX redirect, it will be followed
// according to d.HttpClient's redirect policy.
//
// - Any other HTTP status code will cause UpdateFromUrl to return an error.
func (d *Download) UpdateFromUrl(url string) (err error) {
	var offset int64 = 0
	var fp *os.File

	// Close the progress channel whenever this function completes
	defer close(d.Progress)

	// open a file where we will stream the downloaded update to
	// we do this first because if the caller specified a non-empty dlpath
	// we need to determine how large it is in order to resume the download
	if d.Path == "" {
		// no dlpath specified, use a random tempfile
		fp, err = ioutil.TempFile("", "update")
		if err != nil {
			return
		}
		defer fp.Close()

		// remember the path
		d.Path = fp.Name()
	} else {
		fp, err = os.OpenFile(d.Path, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0600)
		if err != nil {
			return
		}
		defer fp.Close()

		// determine the file size so we can resume the download, if possible
		var fi os.FileInfo
		fi, err = fp.Stat()
		if err != nil {
			return
		}

		offset = fi.Size()
	}

	// create the download request
	req, err := http.NewRequest(d.Method, url, nil)
	if err != nil {
		return
	}

	// we have to add headers like this so they get used across redirects
	trans := d.HttpClient.Transport
	if trans == nil {
		trans = http.DefaultTransport
	}

	d.HttpClient.Transport = &RoundTripper{
		RoundTripFn: func(r *http.Request) (*http.Response, error) {
			// add header for download continuation
			if offset > 0 {
				r.Header.Add("Range", fmt.Sprintf("%d-", offset))
			}

			// ask for gzipped content so that net/http won't unzip it for us
			// and destroy the content length header we need for progress calculations
			r.Header.Add("Accept-Encoding", "gzip")

			return trans.RoundTrip(r)
		},
	}

	// start downloading the file
	resp, err := d.HttpClient.Do(req)
	if err != nil {
		return
	}
	defer resp.Body.Close()

	switch resp.StatusCode {
	// ok
	case 200:
	case 206:

	// no update available
	case 204:
		err = UpdateUnavailable
		return

	// server error
	default:
		err = fmt.Errorf("Non 2XX response when downloading update: %s", resp.Status)
		return
	}

	// Determine how much we have to download
	// net/http sets this to -1 when it is unknown
	clength := resp.ContentLength

	// Read the content from the response body
	rd := resp.Body

	// meter the rate at which we download content for
	// progress reporting if we know how much to expect
	if clength > 0 {
		rd = &MeteredReader{rd: rd, totalSize: clength, progress: d.Progress}
	}

	// Decompress the content if necessary
	if resp.Header.Get("Content-Encoding") == "gzip" {
		rd, err = gzip.NewReader(rd)
		if err != nil {
			return
		}
	}

	// Download the update
	_, err = io.Copy(fp, rd)
	if err != nil {
		return
	}

	// Seek to the beginning of the file before we pass fp to FromStream()
	_, err = fp.Seek(0, os.SEEK_SET)
	if err != nil {
		return
	}

	// Perform the update
	err = FromStream(fp)
	if err != nil {
		return
	}

	// remove the downloaded binary after it's been installed
	os.Remove(d.Path)

	return
}

// FromUrl downloads the contents of the given url and uses them to update
// the current program's executable file. It is a convenience function which is equivalent to
//
// 	NewDownload().UpdateFromUrl(url)
//
// See Download.UpdateFromUrl for more details.
func FromUrl(url string) error {
	return NewDownload().UpdateFromUrl(url)
}

// FromFile reads the contents of the given file and uses them
// to update the current program's executable file by calling FromStream().
func FromFile(filepath string) (err error) {
	// open the new binary
	fp, err := os.Open(filepath)
	if err != nil {
		return
	}
	defer fp.Close()

	// do the update
	err = FromStream(fp)
	return
}

// FromStream reads the contents of the supplied io.Reader newBinary
// and uses them to update the current program's executable file.
//
// FromStream performs the following actions to ensure a cross-platform safe
// update:
//
// - Creates a new file, /path/to/.program-name.new with mode 0755 and copies
// the contents of newBinary into the file
//
// - Renames the current program's executable file from /path/to/program-name
// to /path/to/.program-name.old
//
// - Renames /path/to/.program-name.new to /path/to/program-name
//
// - If the rename is successful, it erases /path/to/.program.old. If this operation
// fails, no error is reported.
//
// - If the rename is unsuccessful, it attempts to rename /path/to/.program-name.old
// back to /path/to/program-name. If this operation fails, the error is not reported
// in order to not mask the error that caused the rename recovery attempt.
func FromStream(newBinary io.Reader) (err error) {
	// get the path to the executable
	thisExecPath, err := execpath.Get()
	if err != nil {
		return
	}

	// get the directory the executable exists in
	execDir := filepath.Dir(thisExecPath)
	execName := filepath.Base(thisExecPath)

	// Copy the contents of of newbinary to a the new executable file
	newExecPath := filepath.Join(execDir, fmt.Sprintf(".%s.new", execName))
	fp, err := os.OpenFile(newExecPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0755)
	if err != nil {
		return
	}
	defer fp.Close()
	_, err = io.Copy(fp, newBinary)

	// if we don't call fp.Close(), windows won't let us move the new executable
	// because the file will still be "in use"
	fp.Close()

	// move the existing executable to a new file in the same directory
	oldExecPath := filepath.Join(execDir, fmt.Sprintf(".%s.old", execName))
	err = os.Rename(thisExecPath, oldExecPath)
	if err != nil {
		return
	}

	// move the new exectuable in to become the new program
	err = os.Rename(newExecPath, thisExecPath)

	if err != nil {
		// copy unsuccessful
		_ = os.Rename(oldExecPath, thisExecPath)
	} else {
		// copy successful, remove the old binary
		_ = os.Remove(oldExecPath)
	}

	return
}
<commit_msg>add new Available attribute on Download that indicates whether the server has a new update<commit_after>/*
Package update allows a program to "self-update", replacing its executable file
with new bytes.

Package update provides the facility to create user experiences like auto-updating
or user-approved updates which manifest as user prompts in commercial applications
with copy similar to "Restart to being using the new version of X".

Updating your program to a new version is as easy as:

	err := update.FromUrl("http://release.example.com/2.0/myprogram")
	if err != nil {
		fmt.Printf("Update failed: %v", err)
	}

The most low-level API is FromStream() which updates the current executable
with the bytes read from an io.Reader.

Additional APIs are provided for common update strategies which include
updating from a file with FromFile() and updating from the internet with
FromUrl().

Using the more advaced Download.UpdateFromUrl() API gives you the ability
to resume an interrupted download to enable large updates to complete even
over intermittent or slow connections. This API also enables more fine-grained
control over how the update is downloaded from the internet as well as access to
download progress,
*/
package update

import (
	"compress/gzip"
	"fmt"
	execpath "github.com/inconshreveable/go-execpath"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"runtime"
)

type MeteredReader struct {
	rd        io.ReadCloser
	totalSize int64
	progress  chan int
	totalRead int64
	ticks     int64
}

func (m *MeteredReader) Close() error {
	return m.rd.Close()
}

func (m *MeteredReader) Read(b []byte) (n int, err error) {
	chunkSize := (m.totalSize / 100) + 1
	lenB := int64(len(b))

	var nChunk int
	for start := int64(0); start < lenB; start += int64(nChunk) {
		end := start + chunkSize
		if end > lenB {
			end = lenB
		}

		nChunk, err = m.rd.Read(b[start:end])

		n += nChunk
		m.totalRead += int64(nChunk)

		if m.totalRead > (m.ticks * chunkSize) {
			m.ticks += 1
			// try to send on channel, but don't block if it's full
			select {
			case m.progress <- int(m.ticks + 1):
			default:
			}

			// give the progress channel consumer a chance to run
			runtime.Gosched()
		}

		if err != nil {
			return
		}
	}

	return
}

// We wrap the round tripper when making requests
// because we need to add headers to the requests we make
// even when they are requests made after a redirect
type RoundTripper struct {
	RoundTripFn func(*http.Request) (*http.Response, error)
}

func (rt *RoundTripper) RoundTrip(r *http.Request) (*http.Response, error) {
	return rt.RoundTripFn(r)
}

// Type Download encapsulates the necessary parameters and state
// needed to download an update from the internet. Create an instance
// with the NewDownload() factory function.
//
// You may only use a Download once,
type Download struct {
	// net/http.Client to use when downloading the update.
	// If nil, a default http.Client is used
	HttpClient *http.Client

	// Path on the file system to dowload the update to
	// If empty, a temporary file is used.
	// After the download begins, this path will be set
	// so that the client can use it to resume aborted
	// downloads
	Path string

	// Progress returns the percentage of the download
	// completed as an integer between 0 and 100
	Progress chan (int)

	// HTTP Method to use in the download request. Default is "GET"
	Method string

	// Set to true when the server confirms a new version is available
	// even if the updating process encounters an error later on
	Available bool
}

// NewDownload initializes a new Download object
func NewDownload() *Download {
	return &Download{
		HttpClient: new(http.Client),
		Progress:   make(chan int),
		Method:     "GET",
	}
}

// UpdateFromUrl downloads the given url from the internet to a file on disk
// and then calls FromStream() to update the current program's executable file
// with the contents of that file.
//
// If the update is successful, the downloaded file will be erased from disk.
// Otherwise, it will remain in d.Path to allow the download to resume later
// or be skipped entirely.
//
// Only HTTP/1.1 servers that implement the Range header are supported.
//
// UpdateFromUrl() uses HTTP status codes to determine what action to take.
//
// - The HTTP server should return 200 or 206 for the update to be downloaded.
//
// - The HTTP server should return 204 if no update is available at this time.
//
// - If the HTTP server returns a 3XX redirect, it will be followed
// according to d.HttpClient's redirect policy.
//
// - Any other HTTP status code will cause UpdateFromUrl to return an error.
func (d *Download) UpdateFromUrl(url string) (err error) {
	var offset int64 = 0
	var fp *os.File

	// Close the progress channel whenever this function completes
	defer close(d.Progress)

	// open a file where we will stream the downloaded update to
	// we do this first because if the caller specified a non-empty dlpath
	// we need to determine how large it is in order to resume the download
	if d.Path == "" {
		// no dlpath specified, use a random tempfile
		fp, err = ioutil.TempFile("", "update")
		if err != nil {
			return
		}
		defer fp.Close()

		// remember the path
		d.Path = fp.Name()
	} else {
		fp, err = os.OpenFile(d.Path, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0600)
		if err != nil {
			return
		}
		defer fp.Close()

		// determine the file size so we can resume the download, if possible
		var fi os.FileInfo
		fi, err = fp.Stat()
		if err != nil {
			return
		}

		offset = fi.Size()
	}

	// create the download request
	req, err := http.NewRequest(d.Method, url, nil)
	if err != nil {
		return
	}

	// we have to add headers like this so they get used across redirects
	trans := d.HttpClient.Transport
	if trans == nil {
		trans = http.DefaultTransport
	}

	d.HttpClient.Transport = &RoundTripper{
		RoundTripFn: func(r *http.Request) (*http.Response, error) {
			// add header for download continuation
			if offset > 0 {
				r.Header.Add("Range", fmt.Sprintf("%d-", offset))
			}

			// ask for gzipped content so that net/http won't unzip it for us
			// and destroy the content length header we need for progress calculations
			r.Header.Add("Accept-Encoding", "gzip")

			return trans.RoundTrip(r)
		},
	}

	// start downloading the file
	resp, err := d.HttpClient.Do(req)
	if err != nil {
		return
	}
	defer resp.Body.Close()

	switch resp.StatusCode {
	// ok
	case 200, 206:
		d.Available = true

	// no update available
	case 204:
		err = fmt.Errorf("204 server response indicates no available update")
		return

	// server error
	default:
		err = fmt.Errorf("Non 2XX response when downloading update: %s", resp.Status)
		return
	}

	// Determine how much we have to download
	// net/http sets this to -1 when it is unknown
	clength := resp.ContentLength

	// Read the content from the response body
	rd := resp.Body

	// meter the rate at which we download content for
	// progress reporting if we know how much to expect
	if clength > 0 {
		rd = &MeteredReader{rd: rd, totalSize: clength, progress: d.Progress}
	}

	// Decompress the content if necessary
	if resp.Header.Get("Content-Encoding") == "gzip" {
		rd, err = gzip.NewReader(rd)
		if err != nil {
			return
		}
	}

	// Download the update
	_, err = io.Copy(fp, rd)
	if err != nil {
		return
	}

	// Seek to the beginning of the file before we pass fp to FromStream()
	_, err = fp.Seek(0, os.SEEK_SET)
	if err != nil {
		return
	}

	// Perform the update
	err = FromStream(fp)
	if err != nil {
		return
	}

	// remove the downloaded binary after it's been installed
	os.Remove(d.Path)

	return
}

// FromUrl downloads the contents of the given url and uses them to update
// the current program's executable file. It is a convenience function which is equivalent to
//
// 	NewDownload().UpdateFromUrl(url)
//
// See Download.UpdateFromUrl for more details.
func FromUrl(url string) error {
	return NewDownload().UpdateFromUrl(url)
}

// FromFile reads the contents of the given file and uses them
// to update the current program's executable file by calling FromStream().
func FromFile(filepath string) (err error) {
	// open the new binary
	fp, err := os.Open(filepath)
	if err != nil {
		return
	}
	defer fp.Close()

	// do the update
	err = FromStream(fp)
	return
}

// FromStream reads the contents of the supplied io.Reader newBinary
// and uses them to update the current program's executable file.
//
// FromStream performs the following actions to ensure a cross-platform safe
// update:
//
// - Creates a new file, /path/to/.program-name.new with mode 0755 and copies
// the contents of newBinary into the file
//
// - Renames the current program's executable file from /path/to/program-name
// to /path/to/.program-name.old
//
// - Renames /path/to/.program-name.new to /path/to/program-name
//
// - If the rename is successful, it erases /path/to/.program.old. If this operation
// fails, no error is reported.
//
// - If the rename is unsuccessful, it attempts to rename /path/to/.program-name.old
// back to /path/to/program-name. If this operation fails, the error is not reported
// in order to not mask the error that caused the rename recovery attempt.
func FromStream(newBinary io.Reader) (err error) {
	// get the path to the executable
	thisExecPath, err := execpath.Get()
	if err != nil {
		return
	}

	// get the directory the executable exists in
	execDir := filepath.Dir(thisExecPath)
	execName := filepath.Base(thisExecPath)

	// Copy the contents of of newbinary to a the new executable file
	newExecPath := filepath.Join(execDir, fmt.Sprintf(".%s.new", execName))
	fp, err := os.OpenFile(newExecPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0755)
	if err != nil {
		return
	}
	defer fp.Close()
	_, err = io.Copy(fp, newBinary)

	// if we don't call fp.Close(), windows won't let us move the new executable
	// because the file will still be "in use"
	fp.Close()

	// move the existing executable to a new file in the same directory
	oldExecPath := filepath.Join(execDir, fmt.Sprintf(".%s.old", execName))
	err = os.Rename(thisExecPath, oldExecPath)
	if err != nil {
		return
	}

	// move the new exectuable in to become the new program
	err = os.Rename(newExecPath, thisExecPath)

	if err != nil {
		// copy unsuccessful
		_ = os.Rename(oldExecPath, thisExecPath)
	} else {
		// copy successful, remove the old binary
		_ = os.Remove(oldExecPath)
	}

	return
}
<|endoftext|>


<commit_before>package wmi

import "testing"

func TestQuery(t *testing.T) {
	var dst []Win32_Process
	err := Query("SELECT * FROM Win32_Process", &dst)
	if err != nil {
		t.Fatal(err)
	}
}

func TestFieldMismatch(t *testing.T) {
	type s struct {
		Name        string
		HandleCount uint32
		Blah        uint32
	}
	var dst []s
	err := Query("SELECT Name, HandleCount FROM Win32_Process", &dst)
	if err == nil || err.Error() != `wmi: cannot load field "Blah" into a "uint32": no such struct field` {
		t.Error("Expected err field mismatch")
	}
}

type Win32_Process struct {
	CSCreationClassName        string
	CSName                     string
	Caption                    string
	CommandLine                string
	CreationClassName          string
	CreationDate               time.Time
	Description                string
	ExecutablePath             string
	ExecutionState             uint16
	Handle                     string
	HandleCount                uint32
	InstallDate                time.Time
	KernelModeTime             uint64
	MaximumWorkingSetSize      uint32
	MinimumWorkingSetSize      uint32
	Name                       string
	OSCreationClassName        string
	OSName                     string
	OtherOperationCount        uint64
	OtherTransferCount         uint64
	PageFaults                 uint32
	PageFileUsage              uint32
	ParentProcessId            uint32
	PeakPageFileUsage          uint32
	PeakVirtualSize            uint64
	PeakWorkingSetSize         uint32
	Priority                   uint32
	PrivatePageCount           uint64
	ProcessId                  uint32
	QuotaNonPagedPoolUsage     uint32
	QuotaPagedPoolUsage        uint32
	QuotaPeakNonPagedPoolUsage uint32
	QuotaPeakPagedPoolUsage    uint32
	ReadOperationCount         uint64
	ReadTransferCount          uint64
	SessionId                  uint32
	Status                     string
	TerminationDate            time.Time
	ThreadCount                uint32
	UserModeTime               uint64
	VirtualSize                uint64
	WindowsVersion             string
	WorkingSetSize             uint64
	WriteOperationCount        uint64
	WriteTransferCount         uint64
}
<commit_msg>Crashing test<commit_after>package wmi

import (
	"fmt"
	"testing"
	"time"
)

func TestQuery(t *testing.T) {
	var dst []Win32_Process
	err := Query("SELECT * FROM Win32_Process", &dst)
	if err != nil {
		t.Fatal(err)
	}
}

func TestFieldMismatch(t *testing.T) {
	type s struct {
		Name        string
		HandleCount uint32
		Blah        uint32
	}
	var dst []s
	err := Query("SELECT Name, HandleCount FROM Win32_Process", &dst)
	if err == nil || err.Error() != `wmi: cannot load field "Blah" into a "uint32": no such struct field` {
		t.Error("Expected err field mismatch")
	}
}

func _TestMany(t *testing.T) {
	go func() {
		for i := 0; true; i++ {
			fmt.Println(i)
			var dst []Win32_PerfRawData_PerfDisk_LogicalDisk
			err := Query("SELECT * FROM Win32_PerfRawData_PerfDisk_LogicalDisk", &dst)
			if err != nil {
				l.Fatal("disk", err)
			}
		}
	}()
	go func() {
		for i := 0; true; i-- {
			fmt.Println(i)
			var dst []Win32_OperatingSystem
			err := Query("SELECT * FROM Win32_OperatingSystem", &dst)
			if err != nil {
				l.Fatal("OS", err)
			}
		}
	}()
	select {}
}

type Win32_Process struct {
	CSCreationClassName        string
	CSName                     string
	Caption                    string
	CommandLine                string
	CreationClassName          string
	CreationDate               time.Time
	Description                string
	ExecutablePath             string
	ExecutionState             uint16
	Handle                     string
	HandleCount                uint32
	InstallDate                time.Time
	KernelModeTime             uint64
	MaximumWorkingSetSize      uint32
	MinimumWorkingSetSize      uint32
	Name                       string
	OSCreationClassName        string
	OSName                     string
	OtherOperationCount        uint64
	OtherTransferCount         uint64
	PageFaults                 uint32
	PageFileUsage              uint32
	ParentProcessId            uint32
	PeakPageFileUsage          uint32
	PeakVirtualSize            uint64
	PeakWorkingSetSize         uint32
	Priority                   uint32
	PrivatePageCount           uint64
	ProcessId                  uint32
	QuotaNonPagedPoolUsage     uint32
	QuotaPagedPoolUsage        uint32
	QuotaPeakNonPagedPoolUsage uint32
	QuotaPeakPagedPoolUsage    uint32
	ReadOperationCount         uint64
	ReadTransferCount          uint64
	SessionId                  uint32
	Status                     string
	TerminationDate            time.Time
	ThreadCount                uint32
	UserModeTime               uint64
	VirtualSize                uint64
	WindowsVersion             string
	WorkingSetSize             uint64
	WriteOperationCount        uint64
	WriteTransferCount         uint64
}

type Win32_PerfRawData_PerfDisk_LogicalDisk struct {
	AvgDiskBytesPerRead          uint64
	AvgDiskBytesPerRead_Base     uint32
	AvgDiskBytesPerTransfer      uint64
	AvgDiskBytesPerTransfer_Base uint32
	AvgDiskBytesPerWrite         uint64
	AvgDiskBytesPerWrite_Base    uint32
	AvgDiskQueueLength           uint64
	AvgDiskReadQueueLength       uint64
	AvgDiskSecPerRead            uint32
	AvgDiskSecPerRead_Base       uint32
	AvgDiskSecPerTransfer        uint32
	AvgDiskSecPerTransfer_Base   uint32
	AvgDiskSecPerWrite           uint32
	AvgDiskSecPerWrite_Base      uint32
	AvgDiskWriteQueueLength      uint64
	Caption                      string
	CurrentDiskQueueLength       uint32
	Description                  string
	DiskBytesPerSec              uint64
	DiskReadBytesPerSec          uint64
	DiskReadsPerSec              uint32
	DiskTransfersPerSec          uint32
	DiskWriteBytesPerSec         uint64
	DiskWritesPerSec             uint32
	FreeMegabytes                uint32
	Frequency_Object             uint64
	Frequency_PerfTime           uint64
	Frequency_Sys100NS           uint64
	Name                         string
	PercentDiskReadTime          uint64
	PercentDiskReadTime_Base     uint64
	PercentDiskTime              uint64
	PercentDiskTime_Base         uint64
	PercentDiskWriteTime         uint64
	PercentDiskWriteTime_Base    uint64
	PercentFreeSpace             uint32
	PercentFreeSpace_Base        uint32
	PercentIdleTime              uint64
	PercentIdleTime_Base         uint64
	SplitIOPerSec                uint32
	Timestamp_Object             uint64
	Timestamp_PerfTime           uint64
	Timestamp_Sys100NS           uint64
}

type Win32_OperatingSystem struct {
	BootDevice                                string
	BuildNumber                               string
	BuildType                                 string
	Caption                                   string
	CodeSet                                   string
	CountryCode                               string
	CreationClassName                         string
	CSCreationClassName                       string
	CSDVersion                                string
	CSName                                    string
	CurrentTimeZone                           int16
	DataExecutionPrevention_Available         bool
	DataExecutionPrevention_32BitApplications bool
	DataExecutionPrevention_Drivers           bool
	DataExecutionPrevention_SupportPolicy     uint8
	Debug                                     bool
	Description                               string
	Distributed                               bool
	EncryptionLevel                           uint32
	ForegroundApplicationBoost                uint8
	FreePhysicalMemory                        uint64
	FreeSpaceInPagingFiles                    uint64
	FreeVirtualMemory                         uint64
	InstallDate                               time.Time
	LargeSystemCache                          uint32
	LastBootUpTime                            time.Time
	LocalDateTime                             time.Time
	Locale                                    string
	Manufacturer                              string
	MaxNumberOfProcesses                      uint32
	MaxProcessMemorySize                      uint64
	MUILanguages                              []string
	Name                                      string
	NumberOfLicensedUsers                     uint32
	NumberOfProcesses                         uint32
	NumberOfUsers                             uint32
	OperatingSystemSKU                        uint32
	Organization                              string
	OSArchitecture                            string
	OSLanguage                                uint32
	OSProductSuite                            uint32
	OSType                                    uint16
	OtherTypeDescription                      string
	PAEEnabled                                bool
	PlusProductID                             string
	PlusVersionNumber                         string
	PortableOperatingSystem                   bool
	Primary                                   bool
	ProductType                               uint32
	RegisteredUser                            string
	SerialNumber                              string
	ServicePackMajorVersion                   uint16
	ServicePackMinorVersion                   uint16
	SizeStoredInPagingFiles                   uint64
	Status                                    string
	SuiteMask                                 uint32
	SystemDevice                              string
	SystemDirectory                           string
	SystemDrive                               string
	TotalSwapSpaceSize                        uint64
	TotalVirtualMemorySize                    uint64
	TotalVisibleMemorySize                    uint64
	Version                                   string
	WindowsDirectory                          string
}
<|endoftext|>


<commit_before>package upsert

import (
	"bytes"
	"errors"
	"fmt"
	"log"
	"reflect"
	"strings"
	"time"

	"github.com/jmoiron/sqlx"
)

const (
	Inserted = iota
	Updated
	NoChange
)

var (
	ErrNoIDReturned = errors.New("no id returned")

	// LongQuery will log long queries if set to a non-zero time
	LongQuery time.Duration
)

// Upserter is an interface specific to sqlx and PostgreSQL that can save a
// single row of data via Upsert(), Update() or Insert(). It doesn't try to
// know anything about relationships between tables. The behavior of Upserter
// depends on three struct tags.
//
//  * db: As with sqlx, this tag is the database column name for the field.
//     If db is not defined, the default is the lowercase value of the field
//     name.
//
//  * upsert: This may either be "key" or "omit". If it's "key", the
//     field/column is part of the where clause when attempting to update
//     an existing column. If it's "omit", the field is ignored completely.
//     By default, the field is considered a non-key value that should be
//     updated/set in the db.
//
//  * upsert_value: This is the placeholder for the value of the field for
//     use by sqlx.NamedExec(). By default, this is :column_name and typically
//     doesn't need to be changed. However, if the field needs to be
//     transformed by an SQL function before storing in the database,
//     this tag can be set. For example, if you had "lat" and "lon" columns,
//     you wouldn't want to store them in the db. Instead you'd want a
//     "location" column tagged with `upsert_value:"ll_to_earth(:lat, :lon)`
//
type Upserter interface {
	// Table returns table name we should save to
	Table() string
}

type columnSpec struct {
	name  string
	value string
}

func newColumnSpec(fieldName string, tag reflect.StructTag) columnSpec {
	cs := columnSpec{}

	// The name of the column is either the value of the "db" struct tag
	// or a lowercase version of the field name.
	dbTag := tag.Get("db")
	if len(dbTag) > 0 {
		cs.name = dbTag
	} else {
		cs.name = strings.ToLower(fieldName)
	}

	// The value placeholder of the column is typically just ":column_name"
	// but can be overriden with upsert_value.
	val := tag.Get("upsert_value")
	if len(val) > 0 {
		cs.value = val
	} else {
		cs.value = ":" + cs.name
	}

	return cs
}

// updateColumns returns the fields that are read from the struct and set
// on upserting in the db. Typically this should include everything except the
// key fields and any composite (array, nested struct) types or any
// field that doesn't map directly into a db column. Tag a field with
// `upsert:"omit"` to explicitly exclude from this list.
func updateColumns(u Upserter) (columns []columnSpec) {
	ut := reflect.TypeOf(u)

	if ut.Kind() == reflect.Ptr {
		ut = ut.Elem()
	}

	if ut.Kind() != reflect.Struct {
		return
	}

	for i := 0; i < ut.NumField(); i++ {
		field := ut.Field(i)
		tag := field.Tag

		// Include any column that isn't tagged with upsert:omit
		if !strings.Contains(tag.Get("upsert"), "omit") {
			columns = append(columns, newColumnSpec(field.Name, tag))
		}
	}

	return
}

// uniqueKeyColumns returns the fields of the struct that together are
// naturally unique. For example, an md5 hash of the content. Or a
// foreign key plus an internal value. This is used in where clause
// when trying to find existing rows. Tag a field with `"upsert:"key"`
// to include in the unique key.
func uniqueKeyColumns(u Upserter) (columns []columnSpec) {
	ut := reflect.TypeOf(u)

	if ut.Kind() == reflect.Ptr {
		ut = ut.Elem()
	}

	if ut.Kind() != reflect.Struct {
		return
	}

	for i := 0; i < ut.NumField(); i++ {
		field := ut.Field(i)
		tag := field.Tag
		// Check if upsert tag contains "key". This wouldn't work
		// if possible options were substrings of one another. For a
		// better implementation, look at src/encoding/json/tags.go
		if strings.Contains(tag.Get("upsert"), "key") {
			columns = append(columns, newColumnSpec(field.Name, tag))
		}
	}

	return
}

// set returns a string like "SET "col1" = :col1, "col2" = :col2" for
// use with sqlx.NamedExec() and friends.
func set(u Upserter) string {
	cols := updateColumns(u)
	n := len(cols)

	b := bytes.Buffer{}

	b.WriteString("SET ")
	for i := 0; i < n; i++ {
		fmt.Fprintf(&b, `"%s" = %s`, cols[i].name, cols[i].value)

		// If we are not at the last value, add a comma
		if i < n-1 {
			b.WriteRune(',')
		}
	}

	return b.String()
}

// values returns a string like `("col1", "col2") VALUES(:col1, :col2)`
// for use with sqlx.NamedExec() etc.
func values(u Upserter) string {
	cols := updateColumns(u)
	n := len(cols)

	b := bytes.Buffer{}

	b.WriteRune('(')
	for i := 0; i < n; i++ {
		fmt.Fprintf(&b, `"%s"`, cols[i].name)

		// If we are not at the last value, add a comma
		if i < n-1 {
			b.WriteRune(',')
		}
	}
	b.WriteRune(')')

	b.WriteString("VALUES (")
	for i := 0; i < n; i++ {
		fmt.Fprintf(&b, `%s`, cols[i].value)

		// If we are not at the last value, add a comma
		if i < n-1 {
			b.WriteRune(',')
		}
	}
	b.WriteRune(')')

	return b.String()
}

// where returns an SQL where clause with all the key columns of
// this Upserter
func where(u Upserter) string {
	keycols := uniqueKeyColumns(u)
	b := bytes.Buffer{}
	n := len(keycols)

	fmt.Fprintf(&b, "WHERE ")

	for i := 0; i < n; i++ {
		// If we need to support NULLs here, the best option may be
		// something like (x = y OR (x is null and y is null))
		// rather than "IS NOT DISTINCT FROM" which doesn't use indexes
		// it seems
		fmt.Fprintf(&b, `%s = %s`, keycols[i].name, keycols[i].value)

		if i < n-1 {
			fmt.Fprint(&b, " AND ")
		}
	}

	return b.String()
}

// updateSQL returns a full SQL command to update this Upserter u
func updateSQL(u Upserter) string {
	q := fmt.Sprintf(`UPDATE "%s" %s %s RETURNING *`,
		u.Table(), set(u), where(u))

	return q
}

// insertSQL returns a full SQL command to insert this Upserter u
func insertSQL(u Upserter) string {
	q := fmt.Sprintf(`INSERT INTO "%s" %s RETURNING *`, u.Table(), values(u))

	return q
}

// getSQL returns a full SQL command to retrieve this Upserter u
func getSQL(u Upserter) string {
	q := fmt.Sprintf(`SELECT * FROM %s %s`, u.Table(), where(u))

	return q
}

func Update(ext sqlx.Ext, u Upserter) (status int, err error) {
	q := updateSQL(u)

	if LongQuery > time.Duration(0) {
		t1 := time.Now()
		defer func() {
			t2 := time.Now()
			if t2.Sub(t1) > LongQuery {
				log.Println(t2.Sub(t1), q, u)
			}
		}()
	}

	otherPtr := reflect.New(reflect.TypeOf(u).Elem())
	other := reflect.Indirect(otherPtr)
	otherInterface := other.Addr().Interface()

	rows, err := sqlx.NamedQuery(ext, getSQL(u), u)
	if err != nil {
		log.Println("error getting", err)
		return
	}

	if rows.Next() {
		err = rows.StructScan(otherInterface)
		if err != nil {
			log.Println("error scanning", err)
			return
		}

		if reflect.DeepEqual(otherInterface, u) {
			log.Println("no change, ignoring update")
			status = NoChange
			return
		} else {
			log.Println("yes there is a change", otherInterface, u)
		}
	}
	rows.Close()

	status = Updated

	// Try to update an existing row
	rows, err = sqlx.NamedQuery(ext, q, u)
	if err != nil {
		log.Println(updateSQL(u), err)
		return
	}
	defer rows.Close()

	if rows.Next() {
		err = rows.StructScan(u)
		if err != nil {
			log.Println(err)
		}
	} else {
		// We could not find anything to update.
		err = ErrNoIDReturned
		return
	}

	return
}

// Insert takes either an sqlx.DB or sqlx.Tx as ext, along with a value
// that implements the Upserter() interface. We attempt to insert it
// and set its primary key id value.
func Insert(ext sqlx.Ext, u Upserter) (err error) {
	q := insertSQL(u)

	if LongQuery > time.Duration(0) {
		t1 := time.Now()
		defer func() {
			t2 := time.Now()
			if t2.Sub(t1) > LongQuery {
				log.Println(t2.Sub(t1), q, u)
			}
		}()
	}

	// Try to insert a row
	rows, err := sqlx.NamedQuery(ext, q, u)
	if err != nil {
		log.Println(err)
		return
	}
	defer rows.Close()

	if rows.Next() {
		err = rows.StructScan(u)
		if err != nil {
			log.Println(err)
		}
	} else {
		// No rows were returned but no SQL error. Weird, return generic
		// error.
		err = ErrNoIDReturned
		return
	}

	return
}

// Upsert takes either an sqlx.DB or sqlx.Tx as ext, along with a value
// that implements the Upserter() interface. We attempt to insert/update it
// and set the new primary key id if that succeeds. inserted returns true
// if a new row was inserted. The client is responsible for wrapping
// in a transaction when needed. This can be used when running a transaction
// at a higher level (upserting multiple items).
func Upsert(ext sqlx.Ext, u Upserter) (status int, err error) {
	// Try to update, return immediately if succcesful
	status, err = Update(ext, u)
	if err == nil {
		return
	}

	// Can't update? Try insert
	err = Insert(ext, u)
	if err != nil {
		log.Println(err)
		return
	}

	status = Inserted

	return
}

// UpsertTx takes only an sqlx.DB and wraps the upsert attempt into a
// a transaction.
func UpsertTx(db *sqlx.DB, u Upserter) (status int, err error) {
	tx, err := db.Beginx()
	if err != nil {
		log.Println("can't start transaction", err)
		return
	}
	defer func() {
		if err != nil {
			tx.Rollback()
		} else {
			tx.Commit()
		}
	}()

	// Try to update
	status, err = Update(tx, u)

	// If we have a nil error, we successfully updated. If we have
	// an err other than ErrNoIDReturned, we couldn't update for an
	// unexpected reason. In either case return.
	if err != ErrNoIDReturned {
		return
	}

	// No ID returned in the update? Try insert
	err = Insert(tx, u)
	if err != nil {
		log.Println(err)
		return
	}

	status = Inserted

	return
}

func Delete(ext sqlx.Ext, u Upserter) (err error) {
	q := fmt.Sprintf(`DELETE FROM "%s" %s`,
		u.Table(), where(u))
	_, err = sqlx.NamedExec(ext, q, u)

	if err != nil {
		log.Println("can't delete", err)
		return
	}

	return
}
<commit_msg>final fixes for getsql upsert<commit_after>package upsert

import (
	"bytes"
	"errors"
	"fmt"
	"log"
	"reflect"
	"strings"
	"time"

	"github.com/jmoiron/sqlx"
)

const (
	Inserted = iota
	Updated
	NoChange
)

func statusToText(status int) string {
	switch status {
	case Inserted:
		return "inserted"
	case Updated:
		return "updated"
	case NoChange:
		return "no change"
	}

	return "invalid status"
}

var (
	ErrNoIDReturned = errors.New("no id returned")

	// LongQuery will log long queries if set to a non-zero time
	LongQuery time.Duration
	Debug     = false
)

// Upserter is an interface specific to sqlx and PostgreSQL that can save a
// single row of data via Upsert(), Update() or Insert(). It doesn't try to
// know anything about relationships between tables. The behavior of Upserter
// depends on three struct tags.
//
//  * db: As with sqlx, this tag is the database column name for the field.
//     If db is not defined, the default is the lowercase value of the field
//     name.
//
//  * upsert: This may either be "key" or "omit". If it's "key", the
//     field/column is part of the where clause when attempting to update
//     an existing column. If it's "omit", the field is ignored completely.
//     By default, the field is considered a non-key value that should be
//     updated/set in the db.
//
//  * upsert_value: This is the placeholder for the value of the field for
//     use by sqlx.NamedExec(). By default, this is :column_name and typically
//     doesn't need to be changed. However, if the field needs to be
//     transformed by an SQL function before storing in the database,
//     this tag can be set. For example, if you had "lat" and "lon" columns,
//     you wouldn't want to store them in the db. Instead you'd want a
//     "location" column tagged with `upsert_value:"ll_to_earth(:lat, :lon)`
//
type Upserter interface {
	// Table returns table name we should save to
	Table() string
}

type columnSpec struct {
	name  string
	value string
}

func newColumnSpec(fieldName string, tag reflect.StructTag) columnSpec {
	cs := columnSpec{}

	// The name of the column is either the value of the "db" struct tag
	// or a lowercase version of the field name.
	dbTag := tag.Get("db")
	if len(dbTag) > 0 {
		cs.name = dbTag
	} else {
		cs.name = strings.ToLower(fieldName)
	}

	// The value placeholder of the column is typically just ":column_name"
	// but can be overriden with upsert_value.
	val := tag.Get("upsert_value")
	if len(val) > 0 {
		cs.value = val
	} else {
		cs.value = ":" + cs.name
	}

	return cs
}

// updateColumns returns the fields that are read from the struct and set
// on upserting in the db. Typically this should include everything except the
// key fields and any composite (array, nested struct) types or any
// field that doesn't map directly into a db column. Tag a field with
// `upsert:"omit"` to explicitly exclude from this list.
func updateColumns(u interface{}) (columns []columnSpec) {
	ut := reflect.TypeOf(u)

	if ut.Kind() == reflect.Ptr {
		ut = ut.Elem()
	}

	if ut.Kind() != reflect.Struct {
		return
	}

	for i := 0; i < ut.NumField(); i++ {
		field := ut.Field(i)
		tag := field.Tag

		// Include any column that isn't tagged with upsert:omit
		if !strings.Contains(tag.Get("upsert"), "omit") {
			columns = append(columns, newColumnSpec(field.Name, tag))
		}
	}

	return
}

// uniqueKeyColumns returns the fields of the struct that together are
// naturally unique. For example, an md5 hash of the content. Or a
// foreign key plus an internal value. This is used in where clause
// when trying to find existing rows. Tag a field with `"upsert:"key"`
// to include in the unique key.
func uniqueKeyColumns(u interface{}) (columns []columnSpec) {
	ut := reflect.TypeOf(u)

	if ut.Kind() == reflect.Ptr {
		ut = ut.Elem()
	}

	if ut.Kind() != reflect.Struct {
		return
	}

	for i := 0; i < ut.NumField(); i++ {
		field := ut.Field(i)
		tag := field.Tag
		// Check if upsert tag contains "key". This wouldn't work
		// if possible options were substrings of one another. For a
		// better implementation, look at src/encoding/json/tags.go
		if strings.Contains(tag.Get("upsert"), "key") {
			columns = append(columns, newColumnSpec(field.Name, tag))
		}
	}

	return
}

// set returns a string like "SET "col1" = :col1, "col2" = :col2" for
// use with sqlx.NamedExec() and friends.
func set(u Upserter) string {
	cols := updateColumns(u)
	n := len(cols)

	b := bytes.Buffer{}

	b.WriteString("SET ")
	for i := 0; i < n; i++ {
		fmt.Fprintf(&b, `"%s" = %s`, cols[i].name, cols[i].value)

		// If we are not at the last value, add a comma
		if i < n-1 {
			b.WriteRune(',')
		}
	}

	return b.String()
}

// values returns a string like `("col1", "col2") VALUES(:col1, :col2)`
// for use with sqlx.NamedExec() etc.
func values(u interface{}) string {
	cols := updateColumns(u)
	n := len(cols)

	b := bytes.Buffer{}

	b.WriteRune('(')
	for i := 0; i < n; i++ {
		fmt.Fprintf(&b, `"%s"`, cols[i].name)

		// If we are not at the last value, add a comma
		if i < n-1 {
			b.WriteRune(',')
		}
	}
	b.WriteRune(')')

	b.WriteString("VALUES (")
	for i := 0; i < n; i++ {
		fmt.Fprintf(&b, `%s`, cols[i].value)

		// If we are not at the last value, add a comma
		if i < n-1 {
			b.WriteRune(',')
		}
	}
	b.WriteRune(')')

	return b.String()
}

// where returns an SQL where clause with all the key columns of
// this Upserter
func where(u Upserter) string {
	keycols := uniqueKeyColumns(u)
	b := bytes.Buffer{}
	n := len(keycols)

	fmt.Fprintf(&b, "WHERE ")

	for i := 0; i < n; i++ {
		// If we need to support NULLs here, the best option may be
		// something like (x = y OR (x is null and y is null))
		// rather than "IS NOT DISTINCT FROM" which doesn't use indexes
		// it seems
		fmt.Fprintf(&b, `%s = %s`, keycols[i].name, keycols[i].value)

		if i < n-1 {
			fmt.Fprint(&b, " AND ")
		}
	}

	return b.String()
}

// updateSQL returns a full SQL command to update this Upserter u
func updateSQL(u Upserter) string {
	q := fmt.Sprintf(`UPDATE "%s" %s %s RETURNING *`,
		u.Table(), set(u), where(u))

	return q
}

// insertSQL returns a full SQL command to insert this Upserter u
func insertSQL(u Upserter) string {
	q := fmt.Sprintf(`INSERT INTO "%s" %s RETURNING *`, u.Table(), values(u))

	return q
}

// getSQL returns a full SQL command to retrieve this Upserter u
func getSQL(u Upserter) string {
	q := fmt.Sprintf(`SELECT * FROM %s %s`, u.Table(), where(u))

	return q
}

func Update(ext sqlx.Ext, u Upserter) (status int, err error) {
	q := updateSQL(u)

	if LongQuery > time.Duration(0) {
		t1 := time.Now()
		defer func() {
			t2 := time.Now()
			if t2.Sub(t1) > LongQuery {
				log.Println(t2.Sub(t1), q, u)
			}
		}()
	}

	otherPtr := reflect.New(reflect.TypeOf(u).Elem())
	other := reflect.Indirect(otherPtr)
	otherInterface := other.Addr().Interface()

	rows, err := sqlx.NamedQuery(ext, getSQL(u), u)
	if err != nil {
		log.Println("error getting", err)
		return
	}

	if rows.Next() {
		err = rows.StructScan(otherInterface)
		if err != nil {
			log.Println("error scanning", err)
			return
		}

		if values(otherInterface) == values(u) {
			status = NoChange
			rows.Close()
			return
		}
	}
	rows.Close()

	status = Updated

	// Try to update an existing row
	rows, err = sqlx.NamedQuery(ext, q, u)
	if err != nil {
		log.Println(updateSQL(u), err)
		return
	}
	defer rows.Close()

	if rows.Next() {
		err = rows.StructScan(u)
		if err != nil {
			log.Println(err)
		}
	} else {
		// We could not find anything to update.
		err = ErrNoIDReturned
		return
	}

	return
}

// Insert takes either an sqlx.DB or sqlx.Tx as ext, along with a value
// that implements the Upserter() interface. We attempt to insert it
// and set its primary key id value.
func Insert(ext sqlx.Ext, u Upserter) (err error) {
	q := insertSQL(u)

	if LongQuery > time.Duration(0) {
		t1 := time.Now()
		defer func() {
			t2 := time.Now()
			if t2.Sub(t1) > LongQuery {
				log.Println(t2.Sub(t1), q, u)
			}
		}()
	}

	// Try to insert a row
	rows, err := sqlx.NamedQuery(ext, q, u)
	if err != nil {
		log.Println(err)
		return
	}
	defer rows.Close()

	if rows.Next() {
		err = rows.StructScan(u)
		if err != nil {
			log.Println(err)
		}
	} else {
		// No rows were returned but no SQL error. Weird, return generic
		// error.
		err = ErrNoIDReturned
		return
	}

	return
}

// Upsert takes either an sqlx.DB or sqlx.Tx as ext, along with a value
// that implements the Upserter() interface. We attempt to insert/update it
// and set the new primary key id if that succeeds. inserted returns true
// if a new row was inserted. The client is responsible for wrapping
// in a transaction when needed. This can be used when running a transaction
// at a higher level (upserting multiple items).
func Upsert(ext sqlx.Ext, u Upserter) (status int, err error) {
	defer func() {
		if Debug {
			log.Println(statusToText(status), u)
		}
	}()

	// Try to update, return immediately if succcesful
	status, err = Update(ext, u)
	if err == nil {
		return
	}

	// Can't update? Try insert
	err = Insert(ext, u)
	if err != nil {
		log.Println(err)
		return
	}

	status = Inserted

	return
}

// UpsertTx takes only an sqlx.DB and wraps the upsert attempt into a
// a transaction.
func UpsertTx(db *sqlx.DB, u Upserter) (status int, err error) {
	defer func() {
		if Debug {
			log.Println(statusToText(status), u)
		}
	}()

	tx, err := db.Beginx()
	if err != nil {
		log.Println("can't start transaction", err)
		return
	}
	defer func() {
		if err != nil {
			tx.Rollback()
		} else {
			tx.Commit()
		}
	}()

	// Try to update
	status, err = Update(tx, u)

	// If we have a nil error, we successfully updated. If we have
	// an err other than ErrNoIDReturned, we couldn't update for an
	// unexpected reason. In either case return.
	if err != ErrNoIDReturned {
		return
	}

	// No ID returned in the update? Try insert
	err = Insert(tx, u)
	if err != nil {
		log.Println(err)
		return
	}

	status = Inserted

	if Debug {
		log.Println(statusToText(status), u)
	}

	return
}

func Delete(ext sqlx.Ext, u Upserter) (err error) {
	q := fmt.Sprintf(`DELETE FROM "%s" %s`,
		u.Table(), where(u))
	_, err = sqlx.NamedExec(ext, q, u)

	if err != nil {
		log.Println("can't delete", err)
		return
	}

	return
}
<|endoftext|>