<commit_before>package dns

// A concurrent client implementation.

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

// Order of events:
// *client -> *reply -> Exchange() -> dial()/send()->write()/receive()->read()
// Do I want make this an interface thingy?
type reply struct {
	client         *Client
	addr           string
	req            *Msg
	conn           net.Conn
	tsigRequestMAC string
	tsigTimersOnly bool
	tsigStatus     error
	rtt            time.Duration
	t              time.Time
}

// 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), defaults to 2 * 1e9
	WriteTimeout time.Duration     // the net.Conn.SetWriteTimeout value for new connections (ns), defaults to 2 * 1e9
	TsigSecret   map[string]string // secret(s) for Tsig map[<zonename>]<base64 secret>, zonename must be fully qualified
}

// 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 *dns.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.DialTimeout("udp", w.addr, 5*1e9)
		} else {
			conn, err = net.DialTimeout(w.client.Net, w.addr, 5*1e9)
		}
		if err == nil {
			break
		}
	}
	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)
	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":
		for a := 0; a < attempts; a++ {
			setTimeouts(w)
			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
			if err == nil {
				return n, err
			}
		}
	case "", "udp", "udp4", "udp6":
		for a := 0; a < attempts; a++ {
			setTimeouts(w)
			n, _, err = w.conn.(*net.UDPConn).ReadFromUDP(p)
			if err != nil {
				if e, ok := err.(net.Error); ok && e.Timeout() {
					continue
				}
				return n, err
			}
			if err == nil {
				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
			if err == nil {
				return n, err
			}
		}
	case "", "udp", "udp4", "udp6":
		for a := 0; a < attempts; a++ {
			setTimeouts(w)
			n, err = w.conn.(*net.UDPConn).Write(p)
			if err != nil {
				if e, ok := err.(net.Error); ok && e.Timeout() {
					continue
				}
				return n, err
			}
			if err == nil {
				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>Remove client.Attempts<commit_after>package dns

// A concurrent client implementation.

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

// Order of events:
// *client -> *reply -> Exchange() -> dial()/send()->write()/receive()->read()
// Do I want make this an interface thingy?
type reply struct {
	client         *Client
	addr           string
	req            *Msg
	conn           net.Conn
	tsigRequestMAC string
	tsigTimersOnly bool
	tsigStatus     error
	rtt            time.Duration
	t              time.Time
}

// 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)
	Retry        bool              // retry with TCP
	ReadTimeout  time.Duration     // the net.Conn.SetReadTimeout value for new connections (ns), defaults to 2 * 1e9
	WriteTimeout time.Duration     // the net.Conn.SetWriteTimeout value for new connections (ns), defaults to 2 * 1e9
	TsigSecret   map[string]string // secret(s) for Tsig map[<zonename>]<base64 secret>, zonename must be fully qualified
}

// 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 *dns.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
	if w.client.Net == "" {
		conn, err = net.DialTimeout("udp", w.addr, 5*1e9)
	} else {
		conn, err = net.DialTimeout(w.client.Net, w.addr, 5*1e9)
	}
	if err != nil {
		return err
	}
	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)
	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
	}
	switch w.client.Net {
	case "tcp", "tcp4", "tcp6":
		setTimeouts(w)
		n, err = w.conn.(*net.TCPConn).Read(p[0:2])
		if err != nil || n != 2 {
			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 {
			return n, err
		}
		i := n
		for i < int(l) {
			j, err := w.conn.(*net.TCPConn).Read(p[i:int(l)])
			if err != nil {
				return i, err
			}
			i += j
		}
		n = i
	case "", "udp", "udp4", "udp6":
		setTimeouts(w)
		n, _, err = w.conn.(*net.UDPConn).ReadFromUDP(p)
		if err != nil {
			return n, err
		}
	}
	return n, err
}

// 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) {
	switch w.client.Net {
	case "tcp", "tcp4", "tcp6":
		if len(p) < 2 {
			return 0, io.ErrShortBuffer
		}
		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 {
			return n, err
		}
		i := n
		if i < len(p) {
			j, err := w.conn.Write(p[i:len(p)])
			if err != nil {
				return i, err
			}
			i += j
		}
		n = i
	case "", "udp", "udp4", "udp6":
		setTimeouts(w)
		n, err = w.conn.(*net.UDPConn).Write(p)
		if err != nil {
			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 getter

import (
	"bytes"
	"crypto/md5"
	"crypto/sha1"
	"crypto/sha256"
	"crypto/sha512"
	"encoding/hex"
	"fmt"
	"hash"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	urlhelper "github.com/hashicorp/go-getter/helper/url"
)

// Client is a client for downloading things.
//
// Top-level functions such as Get are shortcuts for interacting with a client.
// Using a client directly allows more fine-grained control over how downloading
// is done, as well as customizing the protocols supported.
type Client struct {
	// Src is the source URL to get.
	//
	// Dst is the path to save the downloaded thing as. If Dir is set to
	// true, then this should be a directory. If the directory doesn't exist,
	// it will be created for you.
	//
	// Pwd is the working directory for detection. If this isn't set, some
	// detection may fail. Client will not default pwd to the current
	// working directory for security reasons.
	Src string
	Dst string
	Pwd string

	// Dir, if true, tells the Client it is downloading a directory (versus
	// a single file). This distinction is necessary since filenames and
	// directory names follow the same format so disambiguating is impossible
	// without knowing ahead of time.
	Dir bool

	// Detectors is the list of detectors that are tried on the source.
	// If this is nil, then the default Detectors will be used.
	Detectors []Detector

	// Decompressors is the map of decompressors supported by this client.
	// If this is nil, then the default value is the Decompressors global.
	Decompressors map[string]Decompressor

	// Getters is the map of protocols supported by this client. If this
	// is nil, then the default Getters variable will be used.
	Getters map[string]Getter
}

// Get downloads the configured source to the destination.
func (c *Client) Get() error {
	// Store this locally since there are cases we swap this
	dir := c.Dir

	// Default decompressor value
	decompressors := c.Decompressors
	if decompressors == nil {
		decompressors = Decompressors
	}

	// Detect the URL. This is safe if it is already detected.
	detectors := c.Detectors
	if detectors == nil {
		detectors = Detectors
	}
	src, err := Detect(c.Src, c.Pwd, detectors)
	if err != nil {
		return err
	}

	// Determine if we have a forced protocol, i.e. "git::http://..."
	force, src := getForcedGetter(src)

	// If there is a subdir component, then we download the root separately
	// and then copy over the proper subdir.
	var realDst string
	dst := c.Dst
	src, subDir := SourceDirSubdir(src)
	if subDir != "" {
		tmpDir, err := ioutil.TempDir("", "tf")
		if err != nil {
			return err
		}
		if err := os.RemoveAll(tmpDir); err != nil {
			return err
		}
		defer os.RemoveAll(tmpDir)

		realDst = dst
		dst = tmpDir
	}

	u, err := urlhelper.Parse(src)
	if err != nil {
		return err
	}
	if force == "" {
		force = u.Scheme
	}

	getters := c.Getters
	if getters == nil {
		getters = Getters
	}

	g, ok := getters[force]
	if !ok {
		return fmt.Errorf(
			"download not supported for scheme '%s'", force)
	}

	// We have magic query parameters that we use to signal different features
	q := u.Query()

	// Determine if we have an archive type
	archiveV := q.Get("archive")
	if archiveV != "" {
		q.Del("archive")
		u.RawQuery = q.Encode()
	}
	if archiveV == "" {
		// We don't appear to... but is it part of the filename?
		for k, _ := range decompressors {
			if strings.HasSuffix(u.Path, k) {
				archiveV = k
				break
			}
		}
	}

	// If we have a decompressor, then we need to change the destination
	// to download to a temporary path. We unarchive this into the final,
	// real path.
	var decompressDst string
	var decompressDir bool
	decompressor := decompressors[archiveV]
	if decompressor != nil {
		// Create a temporary directory to store our archive. We delete
		// this at the end of everything.
		td, err := ioutil.TempDir("", "getter")
		if err != nil {
			return fmt.Errorf(
				"Error creating temporary directory for archive: %s", err)
		}
		defer os.RemoveAll(td)

		// Swap the download directory to be our temporary path and
		// store the old values.
		decompressDst = dst
		decompressDir = dir
		dst = filepath.Join(td, "archive")
		dir = false
	}

	// Determine if we have a checksum
	var checksumHash hash.Hash
	var checksumValue []byte
	if v := q.Get("checksum"); v != "" {
		// Delete the query parameter if we have it.
		q.Del("checksum")
		u.RawQuery = q.Encode()

		// If we're getting a directory, then this is an error. You cannot
		// checksum a directory. TODO: test
		if dir {
			return fmt.Errorf(
				"checksum cannot be specified for directory download")
		}

		// Determine the checksum hash type
		checksumType := ""
		idx := strings.Index(v, ":")
		if idx > -1 {
			checksumType = v[:idx]
		}
		switch checksumType {
		case "md5":
			checksumHash = md5.New()
		case "sha1":
			checksumHash = sha1.New()
		case "sha256":
			checksumHash = sha256.New()
		case "sha512":
			checksumHash = sha512.New()
		default:
			return fmt.Errorf(
				"unsupported checksum type: %s", checksumType)
		}

		// Get the remainder of the value and parse it into bytes
		b, err := hex.DecodeString(v[idx+1:])
		if err != nil {
			return fmt.Errorf("invalid checksum: %s", err)
		}

		// Set our value
		checksumValue = b
	}

	// If we're not downloading a directory, then just download the file
	// and return.
	if !dir {
		err := g.GetFile(dst, u)
		if err != nil {
			return err
		}

		if checksumHash != nil {
			if err := checksum(dst, checksumHash, checksumValue); err != nil {
				return err
			}
		}

		if decompressor != nil {
			// We have a decompressor, so decompress the current destination
			// into the final destination with the proper mode.
			err := decompressor.Decompress(decompressDst, dst, decompressDir)
			if err != nil {
				return err
			}

			// Swap the information back
			dst = decompressDst
			dir = decompressDir
		}

		// We check the dir value again because it can be switched back
		// if we were unarchiving. If we're still only Get-ing a file, then
		// we're done.
		if !dir {
			return nil
		}
	}

	// If we're at this point we're either downloading a directory or we've
	// downloaded and unarchived a directory and we're just checking subdir.
	// In the case we have a decompressor we don't Get because it was Get
	// above.
	if decompressor == nil {
		// We're downloading a directory, which might require a bit more work
		// if we're specifying a subdir.
		err := g.Get(dst, u)
		if err != nil {
			err = fmt.Errorf("error downloading '%s': %s", src, err)
			return err
		}
	}

	// If we have a subdir, copy that over
	if subDir != "" {
		if err := os.RemoveAll(realDst); err != nil {
			return err
		}
		if err := os.MkdirAll(realDst, 0755); err != nil {
			return err
		}

		return copyDir(realDst, filepath.Join(dst, subDir), false)
	}

	return nil
}

// checksum is a simple method to compute the checksum of a source file
// and compare it to the given expected value.
func checksum(source string, h hash.Hash, v []byte) error {
	f, err := os.Open(source)
	if err != nil {
		return fmt.Errorf("Failed to open file for checksum: %s", err)
	}
	defer f.Close()

	if _, err := io.Copy(h, f); err != nil {
		return fmt.Errorf("Failed to hash: %s", err)
	}

	if actual := h.Sum(nil); !bytes.Equal(actual, v) {
		return fmt.Errorf(
			"Checksums did not match.\nExpected: %s\nGot: %s",
			hex.EncodeToString(v),
			hex.EncodeToString(actual))
	}

	return nil
}

var archiveMagic = map[string]struct{}{
	"tar.bz2": struct{}{},
	"tar":     struct{}{},
	"tar.gz":  struct{}{},
	"tbz2":    struct{}{},
	"tgz":     struct{}{},
	"zip":     struct{}{},
}
<commit_msg>find longest matching suffix<commit_after>package getter

import (
	"bytes"
	"crypto/md5"
	"crypto/sha1"
	"crypto/sha256"
	"crypto/sha512"
	"encoding/hex"
	"fmt"
	"hash"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	urlhelper "github.com/hashicorp/go-getter/helper/url"
)

// Client is a client for downloading things.
//
// Top-level functions such as Get are shortcuts for interacting with a client.
// Using a client directly allows more fine-grained control over how downloading
// is done, as well as customizing the protocols supported.
type Client struct {
	// Src is the source URL to get.
	//
	// Dst is the path to save the downloaded thing as. If Dir is set to
	// true, then this should be a directory. If the directory doesn't exist,
	// it will be created for you.
	//
	// Pwd is the working directory for detection. If this isn't set, some
	// detection may fail. Client will not default pwd to the current
	// working directory for security reasons.
	Src string
	Dst string
	Pwd string

	// Dir, if true, tells the Client it is downloading a directory (versus
	// a single file). This distinction is necessary since filenames and
	// directory names follow the same format so disambiguating is impossible
	// without knowing ahead of time.
	Dir bool

	// Detectors is the list of detectors that are tried on the source.
	// If this is nil, then the default Detectors will be used.
	Detectors []Detector

	// Decompressors is the map of decompressors supported by this client.
	// If this is nil, then the default value is the Decompressors global.
	Decompressors map[string]Decompressor

	// Getters is the map of protocols supported by this client. If this
	// is nil, then the default Getters variable will be used.
	Getters map[string]Getter
}

// Get downloads the configured source to the destination.
func (c *Client) Get() error {
	// Store this locally since there are cases we swap this
	dir := c.Dir

	// Default decompressor value
	decompressors := c.Decompressors
	if decompressors == nil {
		decompressors = Decompressors
	}

	// Detect the URL. This is safe if it is already detected.
	detectors := c.Detectors
	if detectors == nil {
		detectors = Detectors
	}
	src, err := Detect(c.Src, c.Pwd, detectors)
	if err != nil {
		return err
	}

	// Determine if we have a forced protocol, i.e. "git::http://..."
	force, src := getForcedGetter(src)

	// If there is a subdir component, then we download the root separately
	// and then copy over the proper subdir.
	var realDst string
	dst := c.Dst
	src, subDir := SourceDirSubdir(src)
	if subDir != "" {
		tmpDir, err := ioutil.TempDir("", "tf")
		if err != nil {
			return err
		}
		if err := os.RemoveAll(tmpDir); err != nil {
			return err
		}
		defer os.RemoveAll(tmpDir)

		realDst = dst
		dst = tmpDir
	}

	u, err := urlhelper.Parse(src)
	if err != nil {
		return err
	}
	if force == "" {
		force = u.Scheme
	}

	getters := c.Getters
	if getters == nil {
		getters = Getters
	}

	g, ok := getters[force]
	if !ok {
		return fmt.Errorf(
			"download not supported for scheme '%s'", force)
	}

	// We have magic query parameters that we use to signal different features
	q := u.Query()

	// Determine if we have an archive type
	archiveV := q.Get("archive")
	if archiveV != "" {
		q.Del("archive")
		u.RawQuery = q.Encode()
	}
	if archiveV == "" {
		// We don't appear to... but is it part of the filename?
		matchingLen := 0
		for k, _ := range decompressors {
			if strings.HasSuffix(u.Path, k) && len(k) > matchingLen {
				archiveV = k
				matchingLen = len(k)
			}
		}
	}

	// If we have a decompressor, then we need to change the destination
	// to download to a temporary path. We unarchive this into the final,
	// real path.
	var decompressDst string
	var decompressDir bool
	decompressor := decompressors[archiveV]
	if decompressor != nil {
		// Create a temporary directory to store our archive. We delete
		// this at the end of everything.
		td, err := ioutil.TempDir("", "getter")
		if err != nil {
			return fmt.Errorf(
				"Error creating temporary directory for archive: %s", err)
		}
		defer os.RemoveAll(td)

		// Swap the download directory to be our temporary path and
		// store the old values.
		decompressDst = dst
		decompressDir = dir
		dst = filepath.Join(td, "archive")
		dir = false
	}

	// Determine if we have a checksum
	var checksumHash hash.Hash
	var checksumValue []byte
	if v := q.Get("checksum"); v != "" {
		// Delete the query parameter if we have it.
		q.Del("checksum")
		u.RawQuery = q.Encode()

		// If we're getting a directory, then this is an error. You cannot
		// checksum a directory. TODO: test
		if dir {
			return fmt.Errorf(
				"checksum cannot be specified for directory download")
		}

		// Determine the checksum hash type
		checksumType := ""
		idx := strings.Index(v, ":")
		if idx > -1 {
			checksumType = v[:idx]
		}
		switch checksumType {
		case "md5":
			checksumHash = md5.New()
		case "sha1":
			checksumHash = sha1.New()
		case "sha256":
			checksumHash = sha256.New()
		case "sha512":
			checksumHash = sha512.New()
		default:
			return fmt.Errorf(
				"unsupported checksum type: %s", checksumType)
		}

		// Get the remainder of the value and parse it into bytes
		b, err := hex.DecodeString(v[idx+1:])
		if err != nil {
			return fmt.Errorf("invalid checksum: %s", err)
		}

		// Set our value
		checksumValue = b
	}

	// If we're not downloading a directory, then just download the file
	// and return.
	if !dir {
		err := g.GetFile(dst, u)
		if err != nil {
			return err
		}

		if checksumHash != nil {
			if err := checksum(dst, checksumHash, checksumValue); err != nil {
				return err
			}
		}

		if decompressor != nil {
			// We have a decompressor, so decompress the current destination
			// into the final destination with the proper mode.
			err := decompressor.Decompress(decompressDst, dst, decompressDir)
			if err != nil {
				return err
			}

			// Swap the information back
			dst = decompressDst
			dir = decompressDir
		}

		// We check the dir value again because it can be switched back
		// if we were unarchiving. If we're still only Get-ing a file, then
		// we're done.
		if !dir {
			return nil
		}
	}

	// If we're at this point we're either downloading a directory or we've
	// downloaded and unarchived a directory and we're just checking subdir.
	// In the case we have a decompressor we don't Get because it was Get
	// above.
	if decompressor == nil {
		// We're downloading a directory, which might require a bit more work
		// if we're specifying a subdir.
		err := g.Get(dst, u)
		if err != nil {
			err = fmt.Errorf("error downloading '%s': %s", src, err)
			return err
		}
	}

	// If we have a subdir, copy that over
	if subDir != "" {
		if err := os.RemoveAll(realDst); err != nil {
			return err
		}
		if err := os.MkdirAll(realDst, 0755); err != nil {
			return err
		}

		return copyDir(realDst, filepath.Join(dst, subDir), false)
	}

	return nil
}

// checksum is a simple method to compute the checksum of a source file
// and compare it to the given expected value.
func checksum(source string, h hash.Hash, v []byte) error {
	f, err := os.Open(source)
	if err != nil {
		return fmt.Errorf("Failed to open file for checksum: %s", err)
	}
	defer f.Close()

	if _, err := io.Copy(h, f); err != nil {
		return fmt.Errorf("Failed to hash: %s", err)
	}

	if actual := h.Sum(nil); !bytes.Equal(actual, v) {
		return fmt.Errorf(
			"Checksums did not match.\nExpected: %s\nGot: %s",
			hex.EncodeToString(v),
			hex.EncodeToString(actual))
	}

	return nil
}

var archiveMagic = map[string]struct{}{
	"tar.bz2": struct{}{},
	"tar":     struct{}{},
	"tar.gz":  struct{}{},
	"tbz2":    struct{}{},
	"tgz":     struct{}{},
	"zip":     struct{}{},
}
<|endoftext|>


<commit_before>package elasticsearch

import (
	"bytes"
	"encoding/json"
	"errors"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
)

// Searcher set the contract to manage indices, synchronize data and request
type Client interface {
	CreateIndex(indexName, mapping string) (*Response, error)
	DeleteIndex(indexName string) (*Response, error)
	UpdateIndexSetting(indexName, mapping string) (*Response, error)
	IndexSettings(indexName string) (Settings, error)
	IndexExists(indexName string) (bool, error)
	Status(indices string) (*Settings, error)
	InsertDocument(indexName, documentType, identifier string, data []byte) (*InsertDocument, error)
	Document(indexName, documentType, identifier string) (*Document, error)
	DeleteDocument(indexName, documentType, identifier string) (*Document, error)
	Bulk(data []byte) (*Bulk, error)
	Search(indexName, documentType, data string, explain bool) (*SearchResult, error)
	MSearch(queries []MSearchQuery) (*MSearchResult, error)
	Suggest(indexName, data string) ([]byte, error)
}

// A SearchClient describes the client configuration to manage an ElasticSearch index.
type client struct {
	Host url.URL
}

// NewSearchClient creates and initializes a new ElasticSearch client, implements core api for Indexing and searching.
func NewClient(scheme, host, port string) Client {
	u := url.URL{
		Scheme: scheme,
		Host:   host + ":" + port,
	}
	return &client{Host: u}
}

// CreateIndex instantiates an index
// http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/indices-create-index.html
func (c *client) CreateIndex(indexName, mapping string) (*Response, error) {
	url := c.Host.String() + "/" + indexName
	reader := bytes.NewBufferString(mapping)
	response, err := sendHTTPRequest("POST", url, reader)
	if err != nil {
		return &Response{}, err
	}

	esResp := &Response{}
	err = json.Unmarshal(response, esResp)
	if err != nil {
		return &Response{}, err
	}

	return esResp, nil
}

// DeleteIndex deletes an existing index.
// http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/indices-delete-index.html
func (c *client) DeleteIndex(indexName string) (*Response, error) {
	url := c.Host.String() + "/" + indexName
	response, err := sendHTTPRequest("DELETE", url, nil)
	if err != nil {
		return &Response{}, err
	}

	esResp := &Response{}
	err = json.Unmarshal(response, esResp)
	if err != nil {
		return &Response{}, err
	}

	return esResp, nil
}

// UpdateIndexSetting changes specific index level settings in real time
// http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/indices-update-settings.html
func (c *client) UpdateIndexSetting(indexName, mapping string) (*Response, error) {
	url := c.Host.String() + "/" + indexName + "/_settings"
	reader := bytes.NewBufferString(mapping)
	response, err := sendHTTPRequest("PUT", url, reader)
	if err != nil {
		return &Response{}, err
	}

	esResp := &Response{}
	err = json.Unmarshal(response, esResp)
	if err != nil {
		return &Response{}, err
	}

	return esResp, nil
}

// IndexSettings allows to retrieve settings of index
// http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/indices-get-settings.html
func (c *client) IndexSettings(indexName string) (Settings, error) {
	url := c.Host.String() + "/" + indexName + "/_settings"
	response, err := sendHTTPRequest("GET", url, nil)
	if err != nil {
		return Settings{}, err
	}

	type settingsArray map[string]Settings
	dec := json.NewDecoder(bytes.NewBuffer(response))
	var info settingsArray
	err = dec.Decode(&info)
	if err != nil {
		return Settings{}, err
	}

	return info[indexName], nil
}

// IndexExists allows to check if the index exists or not.
// http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/indices-exists.html
func (c *client) IndexExists(indexName string) (bool, error) {
	url := c.Host.String() + "/" + indexName
	httpClient := &http.Client{}
	newReq, err := httpClient.Head(url)
	if err != nil {
		return false, err
	}

	return newReq.StatusCode == http.StatusOK, nil
}

// Status allows to get a comprehensive status information
func (c *client) Status(indices string) (*Settings, error) {
	url := c.Host.String() + "/" + indices + "/_status"
	response, err := sendHTTPRequest("GET", url, nil)
	if err != nil {
		return &Settings{}, err
	}

	esResp := &Settings{}
	err = json.Unmarshal(response, esResp)
	if err != nil {
		return &Settings{}, err
	}

	return esResp, nil
}

// InsertDocument adds or updates a typed JSON document in a specific index, making it searchable
// http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/docs-index_.html
func (c *client) InsertDocument(indexName, documentType, identifier string, data []byte) (*InsertDocument, error) {
	url := c.Host.String() + "/" + indexName + "/" + documentType + "/" + identifier
	reader := bytes.NewBuffer(data)
	response, err := sendHTTPRequest("POST", url, reader)
	if err != nil {
		return &InsertDocument{}, err
	}

	esResp := &InsertDocument{}
	err = json.Unmarshal(response, esResp)
	if err != nil {
		return &InsertDocument{}, err
	}

	return esResp, nil
}

// Document gets a typed JSON document from the index based on its id
// http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/docs-get.html
func (c *client) Document(indexName, documentType, identifier string) (*Document, error) {
	url := c.Host.String() + "/" + indexName + "/" + documentType + "/" + identifier
	response, err := sendHTTPRequest("GET", url, nil)
	if err != nil {
		return &Document{}, err
	}

	esResp := &Document{}
	err = json.Unmarshal(response, esResp)
	if err != nil {
		return &Document{}, err
	}

	return esResp, nil
}

// DeleteDocument deletes a typed JSON document from a specific index based on its id
// http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/docs-delete.html
func (c *client) DeleteDocument(indexName, documentType, identifier string) (*Document, error) {
	url := c.Host.String() + "/" + indexName + "/" + documentType + "/" + identifier
	response, err := sendHTTPRequest("DELETE", url, nil)
	if err != nil {
		return &Document{}, err
	}

	esResp := &Document{}
	err = json.Unmarshal(response, esResp)
	if err != nil {
		return &Document{}, err
	}

	return esResp, nil
}

// Bulk makes it possible to perform many index/delete operations in a single API call.
// This can greatly increase the indexing speed.
// http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/docs-bulk.html
func (c *client) Bulk(data []byte) (*Bulk, error) {
	url := c.Host.String() + "/_bulk"
	reader := bytes.NewBuffer(data)
	response, err := sendHTTPRequest("POST", url, reader)
	if err != nil {
		return &Bulk{}, err
	}

	esResp := &Bulk{}
	err = json.Unmarshal(response, esResp)
	if err != nil {
		return &Bulk{}, err
	}

	return esResp, nil
}

// Search allows to execute a search query and get back search hits that match the query
// http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/docs-delete.html
func (c *client) Search(indexName, documentType, data string, explain bool) (*SearchResult, error) {
	if len(documentType) > 0 {
		documentType = documentType + "/"
	}

	url := c.Host.String() + "/" + indexName + "/" + documentType + "/_search"
	if explain {
		url += "?explain"
	}
	reader := bytes.NewBufferString(data)
	response, err := sendHTTPRequest("POST", url, reader)
	if err != nil {
		return &SearchResult{}, err
	}

	esResp := &SearchResult{}
	err = json.Unmarshal(response, esResp)
	if err != nil {
		return &SearchResult{}, err
	}

	return esResp, nil
}

// MSearch allows to execute a multi-search and get back result
// http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/search-multi-search.html
func (c *client) MSearch(queries []MSearchQuery) (*MSearchResult, error) {
	replacer := strings.NewReplacer("\n", " ")
	queriesList := make([]string, len(queries))
	for i, query := range queries {
		queriesList[i] = query.Header + "\n" + replacer.Replace(query.Body)
	}

	mSearchQuery := strings.Join(queriesList, "\n") + "\n" // Don't forget trailing \n
	url := c.Host.String() + "/_msearch"
	reader := bytes.NewBufferString(mSearchQuery)
	response, err := sendHTTPRequest("POST", url, reader)

	if err != nil {
		return &MSearchResult{}, err
	}

	esResp := &MSearchResult{}
	err = json.Unmarshal(response, esResp)
	if err != nil {
		return &MSearchResult{}, err
	}

	return esResp, nil
}

// Suggest allows basic auto-complete functionality.
// http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/search-suggesters-completion.html
func (c *client) Suggest(indexName, data string) ([]byte, error) {
	url := c.Host.String() + "/" + indexName + "/_suggest"
	reader := bytes.NewBufferString(data)
	response, err := sendHTTPRequest("POST", url, reader)
	return response, err
}

func sendHTTPRequest(method, url string, body io.Reader) ([]byte, error) {
	client := &http.Client{}
	req, err := http.NewRequest(method, url, body)
	if err != nil {
		return nil, err
	}

	if method == "POST" || method == "PUT" {
		req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	}

	newReq, err := client.Do(req)
	if err != nil {
		return nil, err
	}

	defer newReq.Body.Close()
	response, err := ioutil.ReadAll(newReq.Body)
	if err != nil {
		return nil, err
	}

	if newReq.StatusCode > http.StatusCreated && newReq.StatusCode < http.StatusNotFound {
		return nil, errors.New(string(response))
	}

	return response, nil
}
<commit_msg>Support constructor with URL<commit_after>package elasticsearch

import (
	"bytes"
	"encoding/json"
	"errors"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"strings"
)

// Searcher set the contract to manage indices, synchronize data and request
type Client interface {
	CreateIndex(indexName, mapping string) (*Response, error)
	DeleteIndex(indexName string) (*Response, error)
	UpdateIndexSetting(indexName, mapping string) (*Response, error)
	IndexSettings(indexName string) (Settings, error)
	IndexExists(indexName string) (bool, error)
	Status(indices string) (*Settings, error)
	InsertDocument(indexName, documentType, identifier string, data []byte) (*InsertDocument, error)
	Document(indexName, documentType, identifier string) (*Document, error)
	DeleteDocument(indexName, documentType, identifier string) (*Document, error)
	Bulk(data []byte) (*Bulk, error)
	Search(indexName, documentType, data string, explain bool) (*SearchResult, error)
	MSearch(queries []MSearchQuery) (*MSearchResult, error)
	Suggest(indexName, data string) ([]byte, error)
}

// A SearchClient describes the client configuration to manage an ElasticSearch index.
type client struct {
	Host url.URL
}

// NewSearchClient creates and initializes a new ElasticSearch client, implements core api for Indexing and searching.
func NewClient(scheme, host, port string) Client {
	u := url.URL{
		Scheme: scheme,
		Host:   host + ":" + port,
	}
	return &client{Host: u}
}

// NewSearchClient creates and initializes a new ElasticSearch client, implements core api for Indexing and searching.
func NewClientFromUrl(rawurl string) Client {
	u, err := url.Parse(rawurl)
	if err != nil {
		log.Fatal(err)
		return nil
	}
	return &client{Host: *u}
}

// CreateIndex instantiates an index
// http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/indices-create-index.html
func (c *client) CreateIndex(indexName, mapping string) (*Response, error) {
	url := c.Host.String() + "/" + indexName
	reader := bytes.NewBufferString(mapping)
	response, err := sendHTTPRequest("POST", url, reader)
	if err != nil {
		return &Response{}, err
	}

	esResp := &Response{}
	err = json.Unmarshal(response, esResp)
	if err != nil {
		return &Response{}, err
	}

	return esResp, nil
}

// DeleteIndex deletes an existing index.
// http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/indices-delete-index.html
func (c *client) DeleteIndex(indexName string) (*Response, error) {
	url := c.Host.String() + "/" + indexName
	response, err := sendHTTPRequest("DELETE", url, nil)
	if err != nil {
		return &Response{}, err
	}

	esResp := &Response{}
	err = json.Unmarshal(response, esResp)
	if err != nil {
		return &Response{}, err
	}

	return esResp, nil
}

// UpdateIndexSetting changes specific index level settings in real time
// http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/indices-update-settings.html
func (c *client) UpdateIndexSetting(indexName, mapping string) (*Response, error) {
	url := c.Host.String() + "/" + indexName + "/_settings"
	reader := bytes.NewBufferString(mapping)
	response, err := sendHTTPRequest("PUT", url, reader)
	if err != nil {
		return &Response{}, err
	}

	esResp := &Response{}
	err = json.Unmarshal(response, esResp)
	if err != nil {
		return &Response{}, err
	}

	return esResp, nil
}

// IndexSettings allows to retrieve settings of index
// http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/indices-get-settings.html
func (c *client) IndexSettings(indexName string) (Settings, error) {
	url := c.Host.String() + "/" + indexName + "/_settings"
	response, err := sendHTTPRequest("GET", url, nil)
	if err != nil {
		return Settings{}, err
	}

	type settingsArray map[string]Settings
	dec := json.NewDecoder(bytes.NewBuffer(response))
	var info settingsArray
	err = dec.Decode(&info)
	if err != nil {
		return Settings{}, err
	}

	return info[indexName], nil
}

// IndexExists allows to check if the index exists or not.
// http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/indices-exists.html
func (c *client) IndexExists(indexName string) (bool, error) {
	url := c.Host.String() + "/" + indexName
	httpClient := &http.Client{}
	newReq, err := httpClient.Head(url)
	if err != nil {
		return false, err
	}

	return newReq.StatusCode == http.StatusOK, nil
}

// Status allows to get a comprehensive status information
func (c *client) Status(indices string) (*Settings, error) {
	url := c.Host.String() + "/" + indices + "/_status"
	response, err := sendHTTPRequest("GET", url, nil)
	if err != nil {
		return &Settings{}, err
	}

	esResp := &Settings{}
	err = json.Unmarshal(response, esResp)
	if err != nil {
		return &Settings{}, err
	}

	return esResp, nil
}

// InsertDocument adds or updates a typed JSON document in a specific index, making it searchable
// http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/docs-index_.html
func (c *client) InsertDocument(indexName, documentType, identifier string, data []byte) (*InsertDocument, error) {
	url := c.Host.String() + "/" + indexName + "/" + documentType + "/" + identifier
	reader := bytes.NewBuffer(data)
	response, err := sendHTTPRequest("POST", url, reader)
	if err != nil {
		return &InsertDocument{}, err
	}

	esResp := &InsertDocument{}
	err = json.Unmarshal(response, esResp)
	if err != nil {
		return &InsertDocument{}, err
	}

	return esResp, nil
}

// Document gets a typed JSON document from the index based on its id
// http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/docs-get.html
func (c *client) Document(indexName, documentType, identifier string) (*Document, error) {
	url := c.Host.String() + "/" + indexName + "/" + documentType + "/" + identifier
	response, err := sendHTTPRequest("GET", url, nil)
	if err != nil {
		return &Document{}, err
	}

	esResp := &Document{}
	err = json.Unmarshal(response, esResp)
	if err != nil {
		return &Document{}, err
	}

	return esResp, nil
}

// DeleteDocument deletes a typed JSON document from a specific index based on its id
// http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/docs-delete.html
func (c *client) DeleteDocument(indexName, documentType, identifier string) (*Document, error) {
	url := c.Host.String() + "/" + indexName + "/" + documentType + "/" + identifier
	response, err := sendHTTPRequest("DELETE", url, nil)
	if err != nil {
		return &Document{}, err
	}

	esResp := &Document{}
	err = json.Unmarshal(response, esResp)
	if err != nil {
		return &Document{}, err
	}

	return esResp, nil
}

// Bulk makes it possible to perform many index/delete operations in a single API call.
// This can greatly increase the indexing speed.
// http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/docs-bulk.html
func (c *client) Bulk(data []byte) (*Bulk, error) {
	url := c.Host.String() + "/_bulk"
	reader := bytes.NewBuffer(data)
	response, err := sendHTTPRequest("POST", url, reader)
	if err != nil {
		return &Bulk{}, err
	}

	esResp := &Bulk{}
	err = json.Unmarshal(response, esResp)
	if err != nil {
		return &Bulk{}, err
	}

	return esResp, nil
}

// Search allows to execute a search query and get back search hits that match the query
// http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/docs-delete.html
func (c *client) Search(indexName, documentType, data string, explain bool) (*SearchResult, error) {
	if len(documentType) > 0 {
		documentType = documentType + "/"
	}

	url := c.Host.String() + "/" + indexName + "/" + documentType + "/_search"
	if explain {
		url += "?explain"
	}
	reader := bytes.NewBufferString(data)
	response, err := sendHTTPRequest("POST", url, reader)
	if err != nil {
		return &SearchResult{}, err
	}

	esResp := &SearchResult{}
	err = json.Unmarshal(response, esResp)
	if err != nil {
		return &SearchResult{}, err
	}

	return esResp, nil
}

// MSearch allows to execute a multi-search and get back result
// http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/search-multi-search.html
func (c *client) MSearch(queries []MSearchQuery) (*MSearchResult, error) {
	replacer := strings.NewReplacer("\n", " ")
	queriesList := make([]string, len(queries))
	for i, query := range queries {
		queriesList[i] = query.Header + "\n" + replacer.Replace(query.Body)
	}

	mSearchQuery := strings.Join(queriesList, "\n") + "\n" // Don't forget trailing \n
	url := c.Host.String() + "/_msearch"
	reader := bytes.NewBufferString(mSearchQuery)
	response, err := sendHTTPRequest("POST", url, reader)

	if err != nil {
		return &MSearchResult{}, err
	}

	esResp := &MSearchResult{}
	err = json.Unmarshal(response, esResp)
	if err != nil {
		return &MSearchResult{}, err
	}

	return esResp, nil
}

// Suggest allows basic auto-complete functionality.
// http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/search-suggesters-completion.html
func (c *client) Suggest(indexName, data string) ([]byte, error) {
	url := c.Host.String() + "/" + indexName + "/_suggest"
	reader := bytes.NewBufferString(data)
	response, err := sendHTTPRequest("POST", url, reader)
	return response, err
}

func sendHTTPRequest(method, url string, body io.Reader) ([]byte, error) {
	client := &http.Client{}
	req, err := http.NewRequest(method, url, body)
	if err != nil {
		return nil, err
	}

	if method == "POST" || method == "PUT" {
		req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	}

	newReq, err := client.Do(req)
	if err != nil {
		return nil, err
	}

	defer newReq.Body.Close()
	response, err := ioutil.ReadAll(newReq.Body)
	if err != nil {
		return nil, err
	}

	if newReq.StatusCode > http.StatusCreated && newReq.StatusCode < http.StatusNotFound {
		return nil, errors.New(string(response))
	}

	return response, nil
}
<|endoftext|>


<commit_before>// Package write implements writing basic Go types as external Erlang terms.
package write

import (
	"bytes"
	"fmt"
	t "github.com/goerlang/etf/types"
	"io"
	"math"
	"math/big"
	"reflect"
)

type ErrUnknownType struct {
	t reflect.Type
}

func (e *ErrUnknownType) Error() string {
	return fmt.Sprintf("write: can't encode type \"%s\"", e.t.Name())
}

func Atom(w io.Writer, atom t.Atom) (err error) {
	switch size := len(atom); {
	case size <= 0xff:
		// $sL…
		if _, err = w.Write([]byte{t.EttSmallAtom, byte(size)}); err == nil {
			_, err = w.Write([]byte(atom))
		}

	case size <= 0xffff:
		// $dLL…
		_, err = w.Write([]byte{byte(t.EttAtom), byte(size >> 8), byte(size)})
		if err == nil {
			_, err = w.Write([]byte(atom))
		}

	default:
		err = fmt.Errorf("atom is too big (%d bytes)", size)
	}

	return
}

func BigInt(w io.Writer, x *big.Int) (err error) {
	sign := 0
	if x.Sign() < 0 {
		sign = 1
	}

	bytes := reverse(new(big.Int).Abs(x).Bytes())

	switch size := len(bytes); {
	case size <= 0xff:
		// $nAS…
		_, err = w.Write([]byte{t.EttSmallBig, byte(size), byte(sign)})

	case int(uint32(size)) == size:
		// $oAAAAS…
		_, err = w.Write([]byte{
			t.EttLargeBig,
			byte(size >> 24), byte(size >> 16), byte(size >> 8), byte(size),
			byte(sign),
		})

	default:
		err = fmt.Errorf("bad big int size (%d)", size)
	}

	if err == nil {
		_, err = w.Write(bytes)
	}

	return
}

func Binary(w io.Writer, bytes []byte) (err error) {
	switch size := len(bytes); {
	case int(uint32(size)) == size:
		// $mLLLL…
		data := []byte{
			t.EttBinary,
			byte(size >> 24), byte(size >> 16), byte(size >> 8), byte(size),
		}
		if _, err = w.Write(data); err == nil {
			_, err = w.Write(bytes)
		}

	default:
		err = fmt.Errorf("bad binary size (%d)", size)
	}

	return
}

func Bool(w io.Writer, b bool) (err error) {
	if b {
		err = Atom(w, t.Atom("true"))
	} else {
		err = Atom(w, t.Atom("false"))
	}

	return
}

func Float(w io.Writer, f float64) (err error) {
	if _, err = w.Write([]byte{t.EttNewFloat}); err == nil {
		fb := math.Float64bits(f)
		_, err = w.Write([]byte{
			byte(fb >> 56), byte(fb >> 48), byte(fb >> 40), byte(fb >> 32),
			byte(fb >> 24), byte(fb >> 16), byte(fb >> 8), byte(fb),
		})
	}
	return
}

func Int(w io.Writer, x int64) (err error) {
	switch {
	case x >= 0 && x <= math.MaxUint8:
		// $aI
		_, err = w.Write([]byte{t.EttSmallInteger, byte(x)})

	case x >= math.MinInt32 && x <= math.MaxInt32:
		// $bIIII
		_, err = w.Write([]byte{
			t.EttInteger,
			byte(x >> 24), byte(x >> 16), byte(x >> 8), byte(x),
		})

	default:
		err = BigInt(w, big.NewInt(x))
	}

	return
}

func Uint(w io.Writer, x uint64) (err error) {
	switch {
	case x <= math.MaxUint8:
		// $aI
		_, err = w.Write([]byte{t.EttSmallInteger, byte(x)})

	case x <= math.MaxInt32:
		// $bIIII
		_, err = w.Write([]byte{
			t.EttInteger,
			byte(x >> 24), byte(x >> 16), byte(x >> 8), byte(x),
		})

	default:
		err = BigInt(w, new(big.Int).SetUint64(x))
	}

	return
}

func Pid(w io.Writer, p t.Pid) (err error) {
	if _, err = w.Write([]byte{t.EttPid}); err != nil {
		return
	} else if err = Atom(w, t.Atom(p.Node)); err != nil {
		return
	}

	_, err = w.Write([]byte{
		0, 0, byte(p.Id >> 8), byte(p.Id),
		byte(p.Serial >> 24),
		byte(p.Serial >> 16),
		byte(p.Serial >> 8),
		byte(p.Serial),
		p.Creation,
	})

	return
}

func String(w io.Writer, s string) (err error) {
	switch size := len(s); {
	case size <= 0xffff:
		// $kLL…
		_, err = w.Write([]byte{t.EttString, byte(size >> 8), byte(size)})
		if err == nil {
			_, err = w.Write([]byte(s))
		}

	default:
		err = fmt.Errorf("string is too big (%d bytes)", size)
	}

	return
}

func List(w io.Writer, l interface{}) (err error) {
	rv := reflect.ValueOf(l)
	n := rv.Len()
	_, err = w.Write([]byte{
		t.EttList,
		byte(n >> 24),
		byte(n >> 16),
		byte(n >> 8),
		byte(n),
	})

	if err != nil {
		return
	}

	for i := 0; i < n; i++ {
		v := rv.Index(i).Interface()
		if err = Term(w, v); err != nil {
			return
		}
	}

	_, err = w.Write([]byte{t.EttNil})

	return
}

func Record(w io.Writer, r interface{}) (err error) {
	rv := reflect.ValueOf(r)
	n := rv.NumField()
	buf := new(bytes.Buffer)
	arity := 0

	for i := 0; i < n; i++ {
		if f := rv.Field(i); f.CanInterface() {
			if err = Term(buf, f.Interface()); err != nil {
				return
			}
			arity++
		}
	}

	if arity <= math.MaxUint8 {
		_, err = w.Write([]byte{t.EttSmallTuple, byte(arity)})
	} else {
		_, err = w.Write([]byte{
			t.EttLargeTuple,
			byte(arity >> 24),
			byte(arity >> 16),
			byte(arity >> 8),
			byte(arity),
		})
	}

	if err == nil {
		_, err = buf.WriteTo(w)
	}

	return
}

func Tuple(w io.Writer, tuple t.Tuple) (err error) {
	n := len(tuple)
	if n <= math.MaxUint8 {
		_, err = w.Write([]byte{t.EttSmallTuple, byte(n)})
	} else {
		_, err = w.Write([]byte{
			t.EttLargeTuple,
			byte(n >> 24),
			byte(n >> 16),
			byte(n >> 8),
			byte(n),
		})
	}

	if err != nil {
		return
	}

	for _, v := range tuple {
		if err = Term(w, v); err != nil {
			return
		}
	}

	return
}

func Term(w io.Writer, term t.Term) (err error) {
	switch v := term.(type) {
	case bool:
		return Bool(w, v)
	case int8, int16, int32, int64, int:
		return Int(w, reflect.ValueOf(term).Int())
	case uint8, uint16, uint32, uint64, uintptr, uint:
		return Uint(w, reflect.ValueOf(term).Uint())
	case string:
		return String(w, v)
	case []byte:
		return Binary(w, v)
	case t.Atom:
		return Atom(w, v)
	case float64:
		return Float(w, v)
	case float32:
		return Float(w, float64(v))
	case t.Pid:
		return Pid(w, v)
	case t.Tuple:
		return Tuple(w, v)
	default:
		rv := reflect.ValueOf(v)
		switch rv.Kind() {
		case reflect.Struct:
			return Record(w, term)
		case reflect.Array, reflect.Slice:
			return List(w, term)
		case reflect.Ptr:
			return Term(w, rv.Elem())
		//case reflect.Map // FIXME
		default:
			err = &ErrUnknownType{rv.Type()}
		}
	}

	return
}

func reverse(b []byte) []byte {
	size := len(b)
	hsize := size >> 1

	for i := 0; i < hsize; i++ {
		b[i], b[size-i-1] = b[size-i-1], b[i]
	}

	return b
}
<commit_msg>write: remove useless cast<commit_after>// Package write implements writing basic Go types as external Erlang terms.
package write

import (
	"bytes"
	"fmt"
	t "github.com/goerlang/etf/types"
	"io"
	"math"
	"math/big"
	"reflect"
)

type ErrUnknownType struct {
	t reflect.Type
}

func (e *ErrUnknownType) Error() string {
	return fmt.Sprintf("write: can't encode type \"%s\"", e.t.Name())
}

func Atom(w io.Writer, atom t.Atom) (err error) {
	switch size := len(atom); {
	case size <= 0xff:
		// $sL…
		if _, err = w.Write([]byte{t.EttSmallAtom, byte(size)}); err == nil {
			_, err = w.Write([]byte(atom))
		}

	case size <= 0xffff:
		// $dLL…
		_, err = w.Write([]byte{byte(t.EttAtom), byte(size >> 8), byte(size)})
		if err == nil {
			_, err = w.Write([]byte(atom))
		}

	default:
		err = fmt.Errorf("atom is too big (%d bytes)", size)
	}

	return
}

func BigInt(w io.Writer, x *big.Int) (err error) {
	sign := 0
	if x.Sign() < 0 {
		sign = 1
	}

	bytes := reverse(new(big.Int).Abs(x).Bytes())

	switch size := len(bytes); {
	case size <= 0xff:
		// $nAS…
		_, err = w.Write([]byte{t.EttSmallBig, byte(size), byte(sign)})

	case int(uint32(size)) == size:
		// $oAAAAS…
		_, err = w.Write([]byte{
			t.EttLargeBig,
			byte(size >> 24), byte(size >> 16), byte(size >> 8), byte(size),
			byte(sign),
		})

	default:
		err = fmt.Errorf("bad big int size (%d)", size)
	}

	if err == nil {
		_, err = w.Write(bytes)
	}

	return
}

func Binary(w io.Writer, bytes []byte) (err error) {
	switch size := len(bytes); {
	case int(uint32(size)) == size:
		// $mLLLL…
		data := []byte{
			t.EttBinary,
			byte(size >> 24), byte(size >> 16), byte(size >> 8), byte(size),
		}
		if _, err = w.Write(data); err == nil {
			_, err = w.Write(bytes)
		}

	default:
		err = fmt.Errorf("bad binary size (%d)", size)
	}

	return
}

func Bool(w io.Writer, b bool) (err error) {
	if b {
		err = Atom(w, t.Atom("true"))
	} else {
		err = Atom(w, t.Atom("false"))
	}

	return
}

func Float(w io.Writer, f float64) (err error) {
	if _, err = w.Write([]byte{t.EttNewFloat}); err == nil {
		fb := math.Float64bits(f)
		_, err = w.Write([]byte{
			byte(fb >> 56), byte(fb >> 48), byte(fb >> 40), byte(fb >> 32),
			byte(fb >> 24), byte(fb >> 16), byte(fb >> 8), byte(fb),
		})
	}
	return
}

func Int(w io.Writer, x int64) (err error) {
	switch {
	case x >= 0 && x <= math.MaxUint8:
		// $aI
		_, err = w.Write([]byte{t.EttSmallInteger, byte(x)})

	case x >= math.MinInt32 && x <= math.MaxInt32:
		// $bIIII
		_, err = w.Write([]byte{
			t.EttInteger,
			byte(x >> 24), byte(x >> 16), byte(x >> 8), byte(x),
		})

	default:
		err = BigInt(w, big.NewInt(x))
	}

	return
}

func Uint(w io.Writer, x uint64) (err error) {
	switch {
	case x <= math.MaxUint8:
		// $aI
		_, err = w.Write([]byte{t.EttSmallInteger, byte(x)})

	case x <= math.MaxInt32:
		// $bIIII
		_, err = w.Write([]byte{
			t.EttInteger,
			byte(x >> 24), byte(x >> 16), byte(x >> 8), byte(x),
		})

	default:
		err = BigInt(w, new(big.Int).SetUint64(x))
	}

	return
}

func Pid(w io.Writer, p t.Pid) (err error) {
	if _, err = w.Write([]byte{t.EttPid}); err != nil {
		return
	} else if err = Atom(w, p.Node); err != nil {
		return
	}

	_, err = w.Write([]byte{
		0, 0, byte(p.Id >> 8), byte(p.Id),
		byte(p.Serial >> 24),
		byte(p.Serial >> 16),
		byte(p.Serial >> 8),
		byte(p.Serial),
		p.Creation,
	})

	return
}

func String(w io.Writer, s string) (err error) {
	switch size := len(s); {
	case size <= 0xffff:
		// $kLL…
		_, err = w.Write([]byte{t.EttString, byte(size >> 8), byte(size)})
		if err == nil {
			_, err = w.Write([]byte(s))
		}

	default:
		err = fmt.Errorf("string is too big (%d bytes)", size)
	}

	return
}

func List(w io.Writer, l interface{}) (err error) {
	rv := reflect.ValueOf(l)
	n := rv.Len()
	_, err = w.Write([]byte{
		t.EttList,
		byte(n >> 24),
		byte(n >> 16),
		byte(n >> 8),
		byte(n),
	})

	if err != nil {
		return
	}

	for i := 0; i < n; i++ {
		v := rv.Index(i).Interface()
		if err = Term(w, v); err != nil {
			return
		}
	}

	_, err = w.Write([]byte{t.EttNil})

	return
}

func Record(w io.Writer, r interface{}) (err error) {
	rv := reflect.ValueOf(r)
	n := rv.NumField()
	buf := new(bytes.Buffer)
	arity := 0

	for i := 0; i < n; i++ {
		if f := rv.Field(i); f.CanInterface() {
			if err = Term(buf, f.Interface()); err != nil {
				return
			}
			arity++
		}
	}

	if arity <= math.MaxUint8 {
		_, err = w.Write([]byte{t.EttSmallTuple, byte(arity)})
	} else {
		_, err = w.Write([]byte{
			t.EttLargeTuple,
			byte(arity >> 24),
			byte(arity >> 16),
			byte(arity >> 8),
			byte(arity),
		})
	}

	if err == nil {
		_, err = buf.WriteTo(w)
	}

	return
}

func Tuple(w io.Writer, tuple t.Tuple) (err error) {
	n := len(tuple)
	if n <= math.MaxUint8 {
		_, err = w.Write([]byte{t.EttSmallTuple, byte(n)})
	} else {
		_, err = w.Write([]byte{
			t.EttLargeTuple,
			byte(n >> 24),
			byte(n >> 16),
			byte(n >> 8),
			byte(n),
		})
	}

	if err != nil {
		return
	}

	for _, v := range tuple {
		if err = Term(w, v); err != nil {
			return
		}
	}

	return
}

func Term(w io.Writer, term t.Term) (err error) {
	switch v := term.(type) {
	case bool:
		return Bool(w, v)
	case int8, int16, int32, int64, int:
		return Int(w, reflect.ValueOf(term).Int())
	case uint8, uint16, uint32, uint64, uintptr, uint:
		return Uint(w, reflect.ValueOf(term).Uint())
	case string:
		return String(w, v)
	case []byte:
		return Binary(w, v)
	case t.Atom:
		return Atom(w, v)
	case float64:
		return Float(w, v)
	case float32:
		return Float(w, float64(v))
	case t.Pid:
		return Pid(w, v)
	case t.Tuple:
		return Tuple(w, v)
	default:
		rv := reflect.ValueOf(v)
		switch rv.Kind() {
		case reflect.Struct:
			return Record(w, term)
		case reflect.Array, reflect.Slice:
			return List(w, term)
		case reflect.Ptr:
			return Term(w, rv.Elem())
		//case reflect.Map // FIXME
		default:
			err = &ErrUnknownType{rv.Type()}
		}
	}

	return
}

func reverse(b []byte) []byte {
	size := len(b)
	hsize := size >> 1

	for i := 0; i < hsize; i++ {
		b[i], b[size-i-1] = b[size-i-1], b[i]
	}

	return b
}
<|endoftext|>


<commit_before>package gorocksdb

// #include "rocksdb/c.h"
import "C"

import (
	"io"
)

// WriteBatch is a batching of Puts, Merges and Deletes.
type WriteBatch struct {
	c *C.rocksdb_writebatch_t
}

// NewWriteBatch create a WriteBatch object.
func NewWriteBatch() *WriteBatch {
	return NewNativeWriteBatch(C.rocksdb_writebatch_create())
}

// NewNativeWriteBatch create a WriteBatch object.
func NewNativeWriteBatch(c *C.rocksdb_writebatch_t) *WriteBatch {
	return &WriteBatch{c}
}

// Put queues a key-value pair.
func (self *WriteBatch) Put(key, value []byte) {
	cKey := byteToChar(key)
	cValue := byteToChar(value)

	C.rocksdb_writebatch_put(self.c, cKey, C.size_t(len(key)), cValue, C.size_t(len(value)))
}

// Merge queues a merge of "value" with the existing value of "key".
func (self *WriteBatch) Merge(key, value []byte) {
	cKey := byteToChar(key)
	cValue := byteToChar(value)

	C.rocksdb_writebatch_merge(self.c, cKey, C.size_t(len(key)), cValue, C.size_t(len(value)))
}

// Delete queues a deletion of the data at key.
func (self *WriteBatch) Delete(key []byte) {
	cKey := byteToChar(key)

	C.rocksdb_writebatch_delete(self.c, cKey, C.size_t(len(key)))
}

// Data returns the serialized version of this batch.
func (self *WriteBatch) Data() []byte {
	var cSize C.size_t
	cValue := C.rocksdb_writebatch_data(self.c, &cSize)

	return charToByte(cValue, cSize)
}

// Count returns the number of updates in the batch.
func (self *WriteBatch) Count() int {
	return int(C.rocksdb_writebatch_count(self.c))
}

func (self *WriteBatch) NewIterator() *WriteBatchIterator {
	data := self.Data()
	if len(data) < 8+4 {
		return &WriteBatchIterator{}
	}

	return &WriteBatchIterator{data: data[12:]}
}

// Clear removes all the enqueued Put and Deletes.
func (self *WriteBatch) Clear() {
	C.rocksdb_writebatch_clear(self.c)
}

// Destroy deallocates the WriteBatch object.
func (self *WriteBatch) Destroy() {
	C.rocksdb_writebatch_destroy(self.c)
	self.c = nil
}

type WriteBatchRecordType byte

const (
	WriteBatchRecordTypeDeletion WriteBatchRecordType = 0x0
	WriteBatchRecordTypeValue    WriteBatchRecordType = 0x1
	WriteBatchRecordTypeMerge    WriteBatchRecordType = 0x2
	WriteBatchRecordTypeLogData  WriteBatchRecordType = 0x3
)

type WriteBatchRecord struct {
	Key   []byte
	Value []byte
	Type  WriteBatchRecordType
}

type WriteBatchIterator struct {
	data   []byte
	record WriteBatchRecord
	err    error
}

func (self *WriteBatchIterator) Next() bool {
	if self.err != nil || len(self.data) == 0 {
		return false
	}

	self.record.Key = nil
	self.record.Value = nil

	recordType := WriteBatchRecordType(self.data[0])
	self.record.Type = recordType
	self.data = self.data[1:]

	x, n := self.decodeVarint(self.data)
	if n == 0 {
		self.err = io.ErrShortBuffer
		return false
	}
	k := n + int(x)
	self.record.Key = self.data[n:k]
	self.data = self.data[k:]

	if recordType == WriteBatchRecordTypeValue || recordType == WriteBatchRecordTypeMerge {
		x, n := self.decodeVarint(self.data)
		if n == 0 {
			self.err = io.ErrShortBuffer
			return false
		}
		k := n + int(x)
		self.record.Value = self.data[n:k]
		self.data = self.data[k:]
	}

	return true
}

func (self *WriteBatchIterator) Record() *WriteBatchRecord {
	return &self.record
}

func (self *WriteBatchIterator) Error() error {
	return self.err
}

func (self *WriteBatchIterator) decodeVarint(buf []byte) (x uint64, n int) {
	// x, n already 0
	for shift := uint(0); shift < 64; shift += 7 {
		if n >= len(buf) {
			return 0, 0
		}
		b := uint64(buf[n])
		n++
		x |= (b & 0x7F) << shift
		if (b & 0x80) == 0 {
			return x, n
		}
	}

	// The number is too large to represent in a 64-bit value.
	return 0, 0
}
<commit_msg>added ability to create from write batch, supporting commit upstream<commit_after>package gorocksdb

// #include "rocksdb/c.h"
import "C"

import (
	"io"
)

// WriteBatch is a batching of Puts, Merges and Deletes.
type WriteBatch struct {
	c *C.rocksdb_writebatch_t
}

// NewWriteBatch create a WriteBatch object.
func NewWriteBatch() *WriteBatch {
	return NewNativeWriteBatch(C.rocksdb_writebatch_create())
}

// NewNativeWriteBatch create a WriteBatch object.
func NewNativeWriteBatch(c *C.rocksdb_writebatch_t) *WriteBatch {
	return &WriteBatch{c}
}

// WriteBatchFrom creates a write batch from a serialized WriteBatch.
func WriteBatchFrom(data []byte) *WriteBatch {
	return NewNativeWriteBatch(C.rocksdb_writebatch_create_from(byteToChar(data), C.size_t(len(data))))
}

// Put queues a key-value pair.
func (self *WriteBatch) Put(key, value []byte) {
	cKey := byteToChar(key)
	cValue := byteToChar(value)

	C.rocksdb_writebatch_put(self.c, cKey, C.size_t(len(key)), cValue, C.size_t(len(value)))
}

// Merge queues a merge of "value" with the existing value of "key".
func (self *WriteBatch) Merge(key, value []byte) {
	cKey := byteToChar(key)
	cValue := byteToChar(value)

	C.rocksdb_writebatch_merge(self.c, cKey, C.size_t(len(key)), cValue, C.size_t(len(value)))
}

// Delete queues a deletion of the data at key.
func (self *WriteBatch) Delete(key []byte) {
	cKey := byteToChar(key)

	C.rocksdb_writebatch_delete(self.c, cKey, C.size_t(len(key)))
}

// Data returns the serialized version of this batch.
func (self *WriteBatch) Data() []byte {
	var cSize C.size_t
	cValue := C.rocksdb_writebatch_data(self.c, &cSize)

	return charToByte(cValue, cSize)
}

// Count returns the number of updates in the batch.
func (self *WriteBatch) Count() int {
	return int(C.rocksdb_writebatch_count(self.c))
}

func (self *WriteBatch) NewIterator() *WriteBatchIterator {
	data := self.Data()
	if len(data) < 8+4 {
		return &WriteBatchIterator{}
	}

	return &WriteBatchIterator{data: data[12:]}
}

// Clear removes all the enqueued Put and Deletes.
func (self *WriteBatch) Clear() {
	C.rocksdb_writebatch_clear(self.c)
}

// Destroy deallocates the WriteBatch object.
func (self *WriteBatch) Destroy() {
	C.rocksdb_writebatch_destroy(self.c)
	self.c = nil
}

type WriteBatchRecordType byte

const (
	WriteBatchRecordTypeDeletion WriteBatchRecordType = 0x0
	WriteBatchRecordTypeValue    WriteBatchRecordType = 0x1
	WriteBatchRecordTypeMerge    WriteBatchRecordType = 0x2
	WriteBatchRecordTypeLogData  WriteBatchRecordType = 0x3
)

type WriteBatchRecord struct {
	Key   []byte
	Value []byte
	Type  WriteBatchRecordType
}

type WriteBatchIterator struct {
	data   []byte
	record WriteBatchRecord
	err    error
}

func (self *WriteBatchIterator) Next() bool {
	if self.err != nil || len(self.data) == 0 {
		return false
	}

	self.record.Key = nil
	self.record.Value = nil

	recordType := WriteBatchRecordType(self.data[0])
	self.record.Type = recordType
	self.data = self.data[1:]

	x, n := self.decodeVarint(self.data)
	if n == 0 {
		self.err = io.ErrShortBuffer
		return false
	}
	k := n + int(x)
	self.record.Key = self.data[n:k]
	self.data = self.data[k:]

	if recordType == WriteBatchRecordTypeValue || recordType == WriteBatchRecordTypeMerge {
		x, n := self.decodeVarint(self.data)
		if n == 0 {
			self.err = io.ErrShortBuffer
			return false
		}
		k := n + int(x)
		self.record.Value = self.data[n:k]
		self.data = self.data[k:]
	}

	return true
}

func (self *WriteBatchIterator) Record() *WriteBatchRecord {
	return &self.record
}

func (self *WriteBatchIterator) Error() error {
	return self.err
}

func (self *WriteBatchIterator) decodeVarint(buf []byte) (x uint64, n int) {
	// x, n already 0
	for shift := uint(0); shift < 64; shift += 7 {
		if n >= len(buf) {
			return 0, 0
		}
		b := uint64(buf[n])
		n++
		x |= (b & 0x7F) << shift
		if (b & 0x80) == 0 {
			return x, n
		}
	}

	// The number is too large to represent in a 64-bit value.
	return 0, 0
}
<|endoftext|>


<commit_before>package cmd

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"os"
	"sort"
	"strconv"
	"time"

	"github.com/erikgeiser/promptkit/confirmation"
	"github.com/erikgeiser/promptkit/selection"
	"github.com/erikgeiser/promptkit/textinput"
	"github.com/howeyc/ledger"
	"github.com/spf13/cobra"
)

type posting struct {
	Name   string
	Amount float64
}

var addDryRun bool

// addCmd represents the stats command
var addCmd = &cobra.Command{
	Use:   "add",
	Short: "Interactive way to add a transaction to ledger",
	Run: func(cmd *cobra.Command, args []string) {
		transactions, terr := cliTransactions()
		if terr != nil {
			log.Fatalln(terr)
		}

		input := textinput.New("Date:")
		input.InitialValue = time.Now().Format("2006/01/02")
		input.Validate = func(v string) error {
			_, err := time.ParseInLocation("2006/01/02", v, time.Local)
			return err
		}

		date, err := input.RunPrompt()
		if err != nil {
			log.Fatalf("Error: %v\n", err)
		}

		payeeSet := make(map[string]struct{})
		accSet := make(map[string]struct{})
		for _, trans := range transactions {
			payeeSet[trans.Payee] = struct{}{}
			for _, p := range trans.AccountChanges {
				accSet[p.Name] = struct{}{}
			}
		}
		var payeeList []string
		for p := range payeeSet {
			payeeList = append(payeeList, p)
		}
		sort.Strings(payeeList)
		var accountList []string
		for a := range accSet {
			accountList = append(accountList, a)
		}
		sort.Strings(accountList)

		sp := selection.New("Payee:", selection.Choices(payeeList))
		payee, err := sp.RunPrompt()
		if err != nil {
			log.Fatalf("Error: %v\n", err)
		}

		var postings []posting

		for {
			p := posting{}

			sa := selection.New("Account:", selection.Choices(accountList))
			accSel, serr := sa.RunPrompt()
			if serr != nil {
				log.Fatalf("Error: %v\n", serr)
			}
			p.Name = accSel.String

			amtin := textinput.New("Amount:")
			amtin.Validate = func(v string) error {
				_, err := strconv.ParseFloat(v, 64)
				return err
			}
			amtStr, serr := amtin.RunPrompt()
			if serr != nil {
				log.Fatalf("Error: %v\n", serr)
			}
			p.Amount, _ = strconv.ParseFloat(amtStr, 64)

			postings = append(postings, p)
			if len(postings) < 2 {
				continue
			}

			another := confirmation.New("Another Posting?", confirmation.Undecided)
			ready, err := another.RunPrompt()
			if err != nil {
				log.Fatalf("Error: %v\n", err)
			}
			if !ready {
				break
			}
		}

		var tbuf bytes.Buffer
		fmt.Fprintln(&tbuf, date, payee.String)
		for _, p := range postings {
			fmt.Fprintf(&tbuf, "    %s  %f\n", p.Name, p.Amount)
		}
		fmt.Fprintln(&tbuf, "")

		/* Check valid transaction is created */
		trans, perr := ledger.ParseLedger(&tbuf)
		if perr != nil {
			log.Fatalf("Error: %v\n", perr)
		}

		var mw io.Writer
		if addDryRun {
			mw = os.Stdout
		} else {
			f, _ := os.OpenFile(ledgerFilePath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0600)
			defer f.Close()
			mw = io.MultiWriter(f, os.Stdout)
		}

		if columnWide {
			columnWidth = 132
		}

		for _, t := range trans {
			WriteTransaction(mw, t, columnWidth)
		}
	},
}

func init() {
	rootCmd.AddCommand(addCmd)

	addCmd.Flags().BoolVarP(&addDryRun, "dry-run", "n", false, "Do not add to ledger file. Display only.")
	addCmd.Flags().IntVar(&columnWidth, "columns", 80, "Set a column width for output.")
	addCmd.Flags().BoolVar(&columnWide, "wide", false, "Wide output (same as --columns=132).")
}
<commit_msg>fix: set page size for add command<commit_after>package cmd

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"os"
	"sort"
	"strconv"
	"time"

	"github.com/erikgeiser/promptkit/confirmation"
	"github.com/erikgeiser/promptkit/selection"
	"github.com/erikgeiser/promptkit/textinput"
	"github.com/howeyc/ledger"
	"github.com/spf13/cobra"
)

type posting struct {
	Name   string
	Amount float64
}

var addDryRun bool

// addCmd represents the stats command
var addCmd = &cobra.Command{
	Use:   "add",
	Short: "Interactive way to add a transaction to ledger",
	Run: func(cmd *cobra.Command, args []string) {
		transactions, terr := cliTransactions()
		if terr != nil {
			log.Fatalln(terr)
		}

		input := textinput.New("Date:")
		input.InitialValue = time.Now().Format("2006/01/02")
		input.Validate = func(v string) error {
			_, err := time.ParseInLocation("2006/01/02", v, time.Local)
			return err
		}

		date, err := input.RunPrompt()
		if err != nil {
			log.Fatalf("Error: %v\n", err)
		}

		payeeSet := make(map[string]struct{})
		accSet := make(map[string]struct{})
		for _, trans := range transactions {
			payeeSet[trans.Payee] = struct{}{}
			for _, p := range trans.AccountChanges {
				accSet[p.Name] = struct{}{}
			}
		}
		var payeeList []string
		for p := range payeeSet {
			payeeList = append(payeeList, p)
		}
		sort.Strings(payeeList)
		var accountList []string
		for a := range accSet {
			accountList = append(accountList, a)
		}
		sort.Strings(accountList)

		sp := selection.New("Payee:", selection.Choices(payeeList))
		sp.PageSize = 10
		sp.Filter = selection.FilterContainsCaseSensitive
		payee, err := sp.RunPrompt()
		if err != nil {
			log.Fatalf("Error: %v\n", err)
		}

		var postings []posting

		for {
			p := posting{}

			sa := selection.New("Account:", selection.Choices(accountList))
			sa.PageSize = 10
			sa.Filter = selection.FilterContainsCaseSensitive
			accSel, serr := sa.RunPrompt()
			if serr != nil {
				log.Fatalf("Error: %v\n", serr)
			}
			p.Name = accSel.String

			amtin := textinput.New("Amount:")
			amtin.Validate = func(v string) error {
				_, err := strconv.ParseFloat(v, 64)
				return err
			}
			amtStr, serr := amtin.RunPrompt()
			if serr != nil {
				log.Fatalf("Error: %v\n", serr)
			}
			p.Amount, _ = strconv.ParseFloat(amtStr, 64)

			postings = append(postings, p)
			if len(postings) < 2 {
				continue
			}

			another := confirmation.New("Another Posting?", confirmation.Undecided)
			ready, err := another.RunPrompt()
			if err != nil {
				log.Fatalf("Error: %v\n", err)
			}
			if !ready {
				break
			}
		}

		var tbuf bytes.Buffer
		fmt.Fprintln(&tbuf, date, payee.String)
		for _, p := range postings {
			fmt.Fprintf(&tbuf, "    %s  %f\n", p.Name, p.Amount)
		}
		fmt.Fprintln(&tbuf, "")

		/* Check valid transaction is created */
		trans, perr := ledger.ParseLedger(&tbuf)
		if perr != nil {
			log.Fatalf("Error: %v\n", perr)
		}

		var mw io.Writer
		if addDryRun {
			mw = os.Stdout
		} else {
			f, _ := os.OpenFile(ledgerFilePath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0600)
			defer f.Close()
			mw = io.MultiWriter(f, os.Stdout)
		}

		if columnWide {
			columnWidth = 132
		}

		for _, t := range trans {
			WriteTransaction(mw, t, columnWidth)
		}
	},
}

func init() {
	rootCmd.AddCommand(addCmd)

	addCmd.Flags().BoolVarP(&addDryRun, "dry-run", "n", false, "Do not add to ledger file. Display only.")
	addCmd.Flags().IntVar(&columnWidth, "columns", 80, "Set a column width for output.")
	addCmd.Flags().BoolVar(&columnWide, "wide", false, "Wide output (same as --columns=132).")
}
<|endoftext|>


<commit_before>package horizon

import (
	"github.com/jagregory/halgo"
	"github.com/stellar/go-horizon/db"
	"github.com/stellar/go-horizon/render"
	"github.com/zenazn/goji/web"
	"math"
	"net/http"
	"time"
)

type ledgerResource struct {
	halgo.Links
	Id               string    `json:"id"`
	Hash             string    `json:"hash"`
	PrevHash         string    `json:"prev_hash"`
	Sequence         int32     `json:"sequence"`
	TransactionCount int32     `json:"transaction_count"`
	OperationCount   int32     `json:"operation_count"`
	ClosedAt         time.Time `json:"closed_at"`
}

func (l ledgerResource) FromRecord(record db.LedgerRecord) ledgerResource {
	l.Id = record.LedgerHash
	l.Hash = record.LedgerHash
	l.PrevHash = record.PreviousLedgerHash
	l.Sequence = record.Sequence
	return l
}

func ledgerIndexAction(c web.C, w http.ResponseWriter, r *http.Request) {
	ah := &ActionHelper{c: c}
	app := ah.App()
	_, order, limit := ah.GetPagingParams()
	after := ah.GetInt32("after")

	if ah.Err() != nil {
		http.Error(w, ah.Err().Error(), http.StatusBadRequest)
		return
	}

	if after == 0 && order == "desc" {
		after = math.MaxInt32
	}

	query := db.LedgerPageQuery{app.HistoryQuery(), after, order, limit}

	render.Collection(w, r, query, func(record interface{}) interface{} {
		return ledgerResource{}.FromRecord(record.(db.LedgerRecord))
	})
}

func ledgerShowAction(c web.C, w http.ResponseWriter, r *http.Request) {
	ah := &ActionHelper{c: c}
	app := ah.App()
	sequence := ah.GetInt32("id")

	if ah.Err() != nil {
		w.WriteHeader(http.StatusNotFound)
		return
	}

	query := db.LedgerBySequenceQuery{app.HistoryQuery(), sequence}

	render.Single(w, r, query, func(record interface{}) interface{} {
		return ledgerResource{}.FromRecord(record.(db.LedgerRecord))
	})
}
<commit_msg>Extract ledgerRecordToResource<commit_after>package horizon

import (
	"github.com/jagregory/halgo"
	"github.com/stellar/go-horizon/db"
	"github.com/stellar/go-horizon/render"
	"github.com/zenazn/goji/web"
	"math"
	"net/http"
	"time"
)

type ledgerResource struct {
	halgo.Links
	Id               string    `json:"id"`
	Hash             string    `json:"hash"`
	PrevHash         string    `json:"prev_hash"`
	Sequence         int32     `json:"sequence"`
	TransactionCount int32     `json:"transaction_count"`
	OperationCount   int32     `json:"operation_count"`
	ClosedAt         time.Time `json:"closed_at"`
}

func (l ledgerResource) FromRecord(record db.LedgerRecord) ledgerResource {
	l.Id = record.LedgerHash
	l.Hash = record.LedgerHash
	l.PrevHash = record.PreviousLedgerHash
	l.Sequence = record.Sequence
	return l
}

func ledgerIndexAction(c web.C, w http.ResponseWriter, r *http.Request) {
	ah := &ActionHelper{c: c}
	app := ah.App()
	_, order, limit := ah.GetPagingParams()
	after := ah.GetInt32("after")

	if ah.Err() != nil {
		http.Error(w, ah.Err().Error(), http.StatusBadRequest)
		return
	}

	if after == 0 && order == "desc" {
		after = math.MaxInt32
	}

	query := db.LedgerPageQuery{app.HistoryQuery(), after, order, limit}

	render.Collection(w, r, query, ledgerRecordToResource)
}

func ledgerShowAction(c web.C, w http.ResponseWriter, r *http.Request) {
	ah := &ActionHelper{c: c}
	app := ah.App()
	sequence := ah.GetInt32("id")

	if ah.Err() != nil {
		w.WriteHeader(http.StatusNotFound)
		return
	}

	query := db.LedgerBySequenceQuery{app.HistoryQuery(), sequence}

	render.Single(w, r, query, ledgerRecordToResource)
}

func ledgerRecordToResource(record interface{}) (interface{}, error) {
	return ledgerResource{}.FromRecord(record.(db.LedgerRecord)), nil
}
<|endoftext|>


<commit_before>package dalga

import (
	"bytes"
	"fmt"
	"time"
)

type Job struct {
	RoutingKey string
	Body       []byte
	Interval   time.Duration
	NextRun    time.Time
}

func NewJob(routingKey string, body []byte, interval uint32) *Job {
	job := Job{
		RoutingKey: routingKey,
		Body:       body,
		Interval:   time.Duration(interval) * time.Second,
	}
	job.SetNewNextRun()
	return &job
}

func (j *Job) String() string {
	return fmt.Sprintf("Job{%q, %q, %d, %s}", j.RoutingKey, string(j.Body), j.Interval/time.Second, j.NextRun.String()[:23])
}

// Remaining returns the duration until the job's next scheduled time.
func (j *Job) Remaining() time.Duration {
	return -time.Since(j.NextRun)
}

// SetNewNextRun calculates the new run time according to current time and sets it on the job.
func (j *Job) SetNewNextRun() {
	j.NextRun = time.Now().UTC().Add(j.Interval)
}

func (j *Job) Equals(k *Job) bool {
	return (j.RoutingKey == k.RoutingKey) && (bytes.Compare(j.Body, k.Body) == 0)
}
<commit_msg>add comment<commit_after>package dalga

import (
	"bytes"
	"fmt"
	"time"
)

type Job struct {
	RoutingKey string
	Body       []byte
	Interval   time.Duration
	NextRun    time.Time
}

func NewJob(routingKey string, body []byte, interval uint32) *Job {
	job := Job{
		RoutingKey: routingKey,
		Body:       body,
		Interval:   time.Duration(interval) * time.Second,
	}
	job.SetNewNextRun()
	return &job
}

// String implements Stringer interface. Returns the job in human-readable form.
func (j *Job) String() string {
	return fmt.Sprintf("Job{%q, %q, %d, %s}", j.RoutingKey, string(j.Body), j.Interval/time.Second, j.NextRun.String()[:23])
}

// Remaining returns the duration until the job's next scheduled time.
func (j *Job) Remaining() time.Duration {
	return -time.Since(j.NextRun)
}

// SetNewNextRun calculates the new run time according to current time and sets it on the job.
func (j *Job) SetNewNextRun() {
	j.NextRun = time.Now().UTC().Add(j.Interval)
}

func (j *Job) Equals(k *Job) bool {
	return (j.RoutingKey == k.RoutingKey) && (bytes.Compare(j.Body, k.Body) == 0)
}
<|endoftext|>


<commit_before>package pool

import (
	"errors"
	"runtime"
	"sync"
	"sync/atomic"
	"time"
)

const (
	addRoutine = 1
	rmvRoutine = 2
)

// Set of error variables for start up.
var (
	ErrNilMinRoutines        = errors.New("Invalid (nil) minimum number of routines")
	ErrNilMaxRoutines        = errors.New("Invalid (nil) maximum number of routines")
	ErrInvalidMinRoutines    = errors.New("Invalid minimum number of routines")
	ErrInvalidMaxRoutines    = errors.New("Invalid maximum number of routines")
	ErrInvalidAdd            = errors.New("Invalid number of routines to add")
	ErrInvalidMetricHandler  = errors.New("Invalid metric handler")
	ErrInvalidMetricInterval = errors.New("Invalid metric interval")
)

//==============================================================================

// Worker must be implemented by types that want to use
// this worker processes.
type Worker interface {
	Work(context interface{}, id int)
}

// doWork is used internally to route work to the pool.
type doWork struct {
	context interface{}
	do      Worker
}

// Stat contains information about the pool.
type Stat struct {
	Routines    int64 // Current number of routines.
	Pending     int64 // Pending number of routines waiting to submit work.
	Active      int64 // Active number of routines in the work pool.
	Executed    int64 // Number of pieces of work executed.
	MaxRoutines int64 // High water mark of routines the pool has been at.
}

//==============================================================================

// OptEvent defines an handler used to provide events.
type OptEvent struct {
	Event func(context interface{}, event string, format string, a ...interface{})
}

// Config provides configuration for the pool.
type Config struct {
	MinRoutines func() int // Initial and minimum number of routines always in the pool.
	MaxRoutines func() int // Maximum number of routines we will ever grow the pool to.

	// *************************************************************************
	// ** Not Required, optional                                              **
	// *************************************************************************

	OptEvent
}

// Event fires events back to the user for important events.
func (cfg *Config) Event(context interface{}, event string, format string, a ...interface{}) {
	if cfg.OptEvent.Event != nil {
		cfg.OptEvent.Event(context, event, format, a...)
	}
}

//==============================================================================

// Pool provides a pool of routines that can execute any Worker
// tasks that are submitted.
type Pool struct {
	Config
	Name string // Name of this pool.

	tasks    chan doWork    // Unbuffered channel that work is sent into.
	control  chan int       // Unbuffered channel that work for the manager is send into.
	kill     chan struct{}  // Unbuffered channel to signal for a goroutine to die.
	shutdown chan struct{}  // Closed when the Work pool is being shutdown.
	wg       sync.WaitGroup // Manages the number of routines for shutdown.

	counter       int64 // Maintains a count of goroutines ever created to use as an id.
	updatePending int64 // Used to indicate a change to the pool is pending.

	muHealth sync.Mutex // Mutex used to check the health of the system safely.

	routines    int64 // Current number of routines.
	pending     int64 // Pending number of routines waiting to submit work.
	active      int64 // Active number of routines in the work pool.
	executed    int64 // Number of pieces of work executed.
	maxRoutines int64 // High water mark of routines the pool has been at.
}

// New creates a new Pool.
func New(context interface{}, name string, cfg Config) (*Pool, error) {
	if cfg.MinRoutines == nil {
		return nil, ErrNilMinRoutines
	}
	if cfg.MinRoutines() <= 0 {
		return nil, ErrInvalidMinRoutines
	}

	if cfg.MaxRoutines == nil {
		return nil, ErrNilMaxRoutines
	}
	if cfg.MaxRoutines() < cfg.MinRoutines() {
		return nil, ErrInvalidMaxRoutines
	}

	p := Pool{
		Config: cfg,
		Name:   name,

		tasks:    make(chan doWork),
		control:  make(chan int),
		kill:     make(chan struct{}),
		shutdown: make(chan struct{}),
	}

	p.manager(context)
	p.add(context, cfg.MinRoutines())

	return &p, nil
}

// Shutdown waits for all the workers to finish.
func (p *Pool) Shutdown(context interface{}) {
	// If a reset or change is being made, we need to wait.
	for atomic.LoadInt64(&p.updatePending) > 0 {
		time.Sleep(time.Second)
	}

	close(p.shutdown)
	p.wg.Wait()
}

// Do waits for the goroutine pool to take the work to be executed.
func (p *Pool) Do(context interface{}, work Worker) {
	dw := doWork{
		context: context,
		do:      work,
	}

	p.measureHealth()

	atomic.AddInt64(&p.pending, 1)
	p.tasks <- dw
	atomic.AddInt64(&p.pending, -1)
}

// DoWait waits for the goroutine pool to take the work to be executed or gives
// up after the alloted duration. Only use when you want to throw away work and
// not push back.
func (p *Pool) DoWait(context interface{}, work Worker, duration time.Duration) error {
	dw := doWork{
		context: context,
		do:      work,
	}

	p.measureHealth()

	atomic.AddInt64(&p.pending, 1)

	select {
	case p.tasks <- dw:
		atomic.AddInt64(&p.pending, -1)
		return nil

	case <-time.After(duration):
		atomic.AddInt64(&p.pending, -1)
		return errors.New("Timedout waiting to post work")
	}
}

// Stats returns the current snapshot of the pool stats.
func (p *Pool) Stats() Stat {
	return Stat{
		Routines:    atomic.LoadInt64(&p.routines),
		Pending:     atomic.LoadInt64(&p.pending),
		Active:      atomic.LoadInt64(&p.active),
		Executed:    atomic.LoadInt64(&p.executed),
		MaxRoutines: atomic.LoadInt64(&p.maxRoutines),
	}
}

// add creates routines to process work or sets a count for
// routines to terminate.
// NOTE: since our pools are auto-adjustable, we will not give the user ability
// to add routines.
func (p *Pool) add(context interface{}, routines int) error {
	if routines == 0 {
		return ErrInvalidAdd
	}

	cmd := addRoutine
	if routines < 0 {
		routines = routines * -1
		cmd = rmvRoutine
	}

	// Mark the number of adds or removes we are going to perform.
	atomic.AddInt64(&p.updatePending, int64(routines))

	for i := 0; i < routines; i++ {
		p.control <- cmd
	}

	return nil
}

// Reset re-adjusts the pool to match the specified number of routines.
// NOTE: since our pools are auto-adjustable, we will not give the user ability
// to reset the number of routines.
func (p *Pool) reset(context interface{}, routines int) {
	if routines < 0 {
		routines = 0
	}

	current := int(atomic.LoadInt64(&p.routines))
	p.add(context, routines-current)
}

// work performs the users work and keeps stats.
func (p *Pool) work(id int) {

	// Increment the number of routines.
	value := atomic.AddInt64(&p.routines, 1)

	// We need to check and set the high water mark.
	if value > atomic.LoadInt64(&p.maxRoutines) {
		atomic.StoreInt64(&p.maxRoutines, value)
	}

	// Decrement that the add command is complete.
	atomic.AddInt64(&p.updatePending, -1)

done:
	for {
		select {
		case dw := <-p.tasks:
			atomic.AddInt64(&p.active, 1)

			p.execute(id, dw)

			atomic.AddInt64(&p.active, -1)
			atomic.AddInt64(&p.executed, 1)

		case <-p.kill:
			break done
		}
	}

	// Decrement the number of routines.
	atomic.AddInt64(&p.routines, -1)

	// Decrement that the rmv command is complete.
	atomic.AddInt64(&p.updatePending, -1)

	p.wg.Done()
}

// execute performs the work in a recoverable way.
func (p *Pool) execute(id int, dw doWork) {
	defer func() {
		if r := recover(); r != nil {

			// Capture the stack trace
			buf := make([]byte, 10000)
			runtime.Stack(buf, false)

			p.Event(dw.context, "execute", "ERROR : %s", string(buf))
		}
	}()

	// Perform the work.
	dw.do.Work(dw.context, id)
}

// measureHealth calculates the health of the work pool.
func (p *Pool) measureHealth() {

	// If there are values pending to be updated, just
	// leave. We need those to finish first.
	if atomic.LoadInt64(&p.updatePending) > 0 {
		return
	}

	p.muHealth.Lock()
	defer p.muHealth.Unlock()

	stats := p.Stats()

	// We are not performing any work at all and we have more routines than min.
	if stats.Pending == 0 && stats.Active == 0 && (stats.Routines > int64(p.MinRoutines())) {

		// Reset the pool back to the min value.
		p.reset(p.Name, p.MinRoutines())
		return
	}

	// If we have no available routines at the moment and we have room to grow.
	if (stats.Routines == stats.Active) && (stats.Routines < int64(p.MaxRoutines())) {

		// Calculate the number of goroutines to add.
		add := int(float64(stats.Routines) * .20)

		// Check if we calculated a 0 grow.
		if add == 0 {
			add = 1
		}

		// Check if we will go over max.
		if (int(stats.Routines) + add) > p.MaxRoutines() {
			add = p.MaxRoutines() - int(stats.Routines)
		}

		// Request this number to be added.
		p.add(p.Name, add)
	}
}

// manager controls changes to the work pool including stats and shutting down.
func (p *Pool) manager(context interface{}) {
	p.wg.Add(1)

	go func() {
		for {
			select {
			case <-p.shutdown:

				// Capture the current number of routines.
				routines := int(atomic.LoadInt64(&p.routines))

				// Send a kill to all the existing routines.
				for i := 0; i < routines; i++ {
					p.kill <- struct{}{}
				}

				// Decrement the waitgroup and kill the manager.
				p.wg.Done()
				return

			case c := <-p.control:
				switch c {
				case addRoutine:

					// Capture the number of routines.
					routines := int(atomic.LoadInt64(&p.routines))

					// Is there room to add goroutines.
					if routines == p.MaxRoutines() {
						break
					}

					// Increment the total number of routines ever created.
					counter := atomic.AddInt64(&p.counter, 1)

					// Create the routine.
					p.wg.Add(1)
					go p.work(int(counter))

				case rmvRoutine:

					// Capture the number of routines.
					routines := int(atomic.LoadInt64(&p.routines))

					// Are there routines to remove.
					if routines <= p.MinRoutines() {
						break
					}

					// Send a kill signal to remove a routine.
					p.kill <- struct{}{}
				}
			}
		}
	}()
}
<commit_msg>Better design for the channel to allow the user to manage the cost of garbage.<commit_after>package pool

import (
	"errors"
	"runtime"
	"sync"
	"sync/atomic"
	"time"
)

const (
	addRoutine = 1
	rmvRoutine = 2
)

// Set of error variables for start up.
var (
	ErrNilMinRoutines        = errors.New("Invalid (nil) minimum number of routines")
	ErrNilMaxRoutines        = errors.New("Invalid (nil) maximum number of routines")
	ErrInvalidMinRoutines    = errors.New("Invalid minimum number of routines")
	ErrInvalidMaxRoutines    = errors.New("Invalid maximum number of routines")
	ErrInvalidAdd            = errors.New("Invalid number of routines to add")
	ErrInvalidMetricHandler  = errors.New("Invalid metric handler")
	ErrInvalidMetricInterval = errors.New("Invalid metric interval")
)

//==============================================================================

// Worker must be implemented by types that want to use
// this worker processes.
type Worker interface {
	Work(context interface{}, id int)
}

// doWork is used internally to route work to the pool.
type doWork struct {
	context interface{}
	do      Worker
}

// Stat contains information about the pool.
type Stat struct {
	Routines    int64 // Current number of routines.
	Pending     int64 // Pending number of routines waiting to submit work.
	Active      int64 // Active number of routines in the work pool.
	Executed    int64 // Number of pieces of work executed.
	MaxRoutines int64 // High water mark of routines the pool has been at.
}

//==============================================================================

// OptEvent defines an handler used to provide events.
type OptEvent struct {
	Event func(context interface{}, event string, format string, a ...interface{})
}

// Config provides configuration for the pool.
type Config struct {
	MinRoutines func() int // Initial and minimum number of routines always in the pool.
	MaxRoutines func() int // Maximum number of routines we will ever grow the pool to.

	// *************************************************************************
	// ** Not Required, optional                                              **
	// *************************************************************************

	OptEvent
}

// Event fires events back to the user for important events.
func (cfg *Config) Event(context interface{}, event string, format string, a ...interface{}) {
	if cfg.OptEvent.Event != nil {
		cfg.OptEvent.Event(context, event, format, a...)
	}
}

//==============================================================================

// Pool provides a pool of routines that can execute any Worker
// tasks that are submitted.
type Pool struct {
	Config
	Name string // Name of this pool.

	tasks    chan doWork    // Unbuffered channel that work is sent into.
	control  chan int       // Unbuffered channel that work for the manager is send into.
	kill     chan struct{}  // Unbuffered channel to signal for a goroutine to die.
	shutdown chan struct{}  // Closed when the Work pool is being shutdown.
	wg       sync.WaitGroup // Manages the number of routines for shutdown.

	counter       int64 // Maintains a count of goroutines ever created to use as an id.
	updatePending int64 // Used to indicate a change to the pool is pending.

	muHealth sync.Mutex // Mutex used to check the health of the system safely.

	routines    int64 // Current number of routines.
	pending     int64 // Pending number of routines waiting to submit work.
	active      int64 // Active number of routines in the work pool.
	executed    int64 // Number of pieces of work executed.
	maxRoutines int64 // High water mark of routines the pool has been at.
}

// New creates a new Pool.
func New(context interface{}, name string, cfg Config) (*Pool, error) {
	if cfg.MinRoutines == nil {
		return nil, ErrNilMinRoutines
	}
	if cfg.MinRoutines() <= 0 {
		return nil, ErrInvalidMinRoutines
	}

	if cfg.MaxRoutines == nil {
		return nil, ErrNilMaxRoutines
	}
	if cfg.MaxRoutines() < cfg.MinRoutines() {
		return nil, ErrInvalidMaxRoutines
	}

	p := Pool{
		Config: cfg,
		Name:   name,

		tasks:    make(chan doWork),
		control:  make(chan int),
		kill:     make(chan struct{}),
		shutdown: make(chan struct{}),
	}

	p.manager(context)
	p.add(context, cfg.MinRoutines())

	return &p, nil
}

// Shutdown waits for all the workers to finish.
func (p *Pool) Shutdown(context interface{}) {
	// If a reset or change is being made, we need to wait.
	for atomic.LoadInt64(&p.updatePending) > 0 {
		time.Sleep(time.Second)
	}

	close(p.shutdown)
	p.wg.Wait()
}

// Do waits for the goroutine pool to take the work to be executed.
func (p *Pool) Do(context interface{}, work Worker) {
	dw := doWork{
		context: context,
		do:      work,
	}

	p.measureHealth()

	atomic.AddInt64(&p.pending, 1)
	p.tasks <- dw
	atomic.AddInt64(&p.pending, -1)
}

// DoWait waits for the goroutine pool to take the work to be executed or gives
// up after the alloted duration. Only use when you want to throw away work and
// not push back.
func (p *Pool) DoWait(context interface{}, work Worker, duration <-chan time.Time) error {
	dw := doWork{
		context: context,
		do:      work,
	}

	p.measureHealth()

	atomic.AddInt64(&p.pending, 1)

	select {
	case p.tasks <- dw:
		atomic.AddInt64(&p.pending, -1)
		return nil

	case <-duration:
		atomic.AddInt64(&p.pending, -1)
		return errors.New("Timedout waiting to post work")
	}
}

// Stats returns the current snapshot of the pool stats.
func (p *Pool) Stats() Stat {
	return Stat{
		Routines:    atomic.LoadInt64(&p.routines),
		Pending:     atomic.LoadInt64(&p.pending),
		Active:      atomic.LoadInt64(&p.active),
		Executed:    atomic.LoadInt64(&p.executed),
		MaxRoutines: atomic.LoadInt64(&p.maxRoutines),
	}
}

// add creates routines to process work or sets a count for
// routines to terminate.
// NOTE: since our pools are auto-adjustable, we will not give the user ability
// to add routines.
func (p *Pool) add(context interface{}, routines int) error {
	if routines == 0 {
		return ErrInvalidAdd
	}

	cmd := addRoutine
	if routines < 0 {
		routines = routines * -1
		cmd = rmvRoutine
	}

	// Mark the number of adds or removes we are going to perform.
	atomic.AddInt64(&p.updatePending, int64(routines))

	for i := 0; i < routines; i++ {
		p.control <- cmd
	}

	return nil
}

// Reset re-adjusts the pool to match the specified number of routines.
// NOTE: since our pools are auto-adjustable, we will not give the user ability
// to reset the number of routines.
func (p *Pool) reset(context interface{}, routines int) {
	if routines < 0 {
		routines = 0
	}

	current := int(atomic.LoadInt64(&p.routines))
	p.add(context, routines-current)
}

// work performs the users work and keeps stats.
func (p *Pool) work(id int) {

	// Increment the number of routines.
	value := atomic.AddInt64(&p.routines, 1)

	// We need to check and set the high water mark.
	if value > atomic.LoadInt64(&p.maxRoutines) {
		atomic.StoreInt64(&p.maxRoutines, value)
	}

	// Decrement that the add command is complete.
	atomic.AddInt64(&p.updatePending, -1)

done:
	for {
		select {
		case dw := <-p.tasks:
			atomic.AddInt64(&p.active, 1)

			p.execute(id, dw)

			atomic.AddInt64(&p.active, -1)
			atomic.AddInt64(&p.executed, 1)

		case <-p.kill:
			break done
		}
	}

	// Decrement the number of routines.
	atomic.AddInt64(&p.routines, -1)

	// Decrement that the rmv command is complete.
	atomic.AddInt64(&p.updatePending, -1)

	p.wg.Done()
}

// execute performs the work in a recoverable way.
func (p *Pool) execute(id int, dw doWork) {
	defer func() {
		if r := recover(); r != nil {

			// Capture the stack trace
			buf := make([]byte, 10000)
			runtime.Stack(buf, false)

			p.Event(dw.context, "execute", "ERROR : %s", string(buf))
		}
	}()

	// Perform the work.
	dw.do.Work(dw.context, id)
}

// measureHealth calculates the health of the work pool.
func (p *Pool) measureHealth() {

	// If there are values pending to be updated, just
	// leave. We need those to finish first.
	if atomic.LoadInt64(&p.updatePending) > 0 {
		return
	}

	p.muHealth.Lock()
	defer p.muHealth.Unlock()

	stats := p.Stats()

	// We are not performing any work at all and we have more routines than min.
	if stats.Pending == 0 && stats.Active == 0 && (stats.Routines > int64(p.MinRoutines())) {

		// Reset the pool back to the min value.
		p.reset(p.Name, p.MinRoutines())
		return
	}

	// If we have no available routines at the moment and we have room to grow.
	if (stats.Routines == stats.Active) && (stats.Routines < int64(p.MaxRoutines())) {

		// Calculate the number of goroutines to add.
		add := int(float64(stats.Routines) * .20)

		// Check if we calculated a 0 grow.
		if add == 0 {
			add = 1
		}

		// Check if we will go over max.
		if (int(stats.Routines) + add) > p.MaxRoutines() {
			add = p.MaxRoutines() - int(stats.Routines)
		}

		// Request this number to be added.
		p.add(p.Name, add)
	}
}

// manager controls changes to the work pool including stats and shutting down.
func (p *Pool) manager(context interface{}) {
	p.wg.Add(1)

	go func() {
		for {
			select {
			case <-p.shutdown:

				// Capture the current number of routines.
				routines := int(atomic.LoadInt64(&p.routines))

				// Send a kill to all the existing routines.
				for i := 0; i < routines; i++ {
					p.kill <- struct{}{}
				}

				// Decrement the waitgroup and kill the manager.
				p.wg.Done()
				return

			case c := <-p.control:
				switch c {
				case addRoutine:

					// Capture the number of routines.
					routines := int(atomic.LoadInt64(&p.routines))

					// Is there room to add goroutines.
					if routines == p.MaxRoutines() {
						break
					}

					// Increment the total number of routines ever created.
					counter := atomic.AddInt64(&p.counter, 1)

					// Create the routine.
					p.wg.Add(1)
					go p.work(int(counter))

				case rmvRoutine:

					// Capture the number of routines.
					routines := int(atomic.LoadInt64(&p.routines))

					// Are there routines to remove.
					if routines <= p.MinRoutines() {
						break
					}

					// Send a kill signal to remove a routine.
					p.kill <- struct{}{}
				}
			}
		}
	}()
}
<|endoftext|>