<commit_before>package ssh2docker

import (
	"fmt"
	"io"
	"os"
	"os/exec"
	"strings"
	"sync"
	"syscall"

	"github.com/Sirupsen/logrus"
	"github.com/kr/pty"
	"golang.org/x/crypto/ssh"
)

// Client is one client connection
type Client struct {
	Conn     *ssh.ServerConn
	Chans    <-chan ssh.NewChannel
	Reqs     <-chan *ssh.Request
	Server   *Server
	Pty, Tty *os.File
	Env      Environment
}

// NewClient initializes a new client
func NewClient(conn *ssh.ServerConn, chans <-chan ssh.NewChannel, reqs <-chan *ssh.Request, server *Server) *Client {
	client := Client{
		Conn:   conn,
		Chans:  chans,
		Reqs:   reqs,
		Server: server,
		Env: Environment{
			"TERM":              os.Getenv("TERM"),
			"DOCKER_HOST":       os.Getenv("DOCKER_HOST"),
			"DOCKER_CERT_PATH":  os.Getenv("DOCKER_CERT_PATH"),
			"DOCKER_TLS_VERIFY": os.Getenv("DOCKER_TLS_VERIFY"),
		},
	}

	logrus.Infof("NewClient: User=%q, ClientVersion=%q", conn.User(), fmt.Sprintf("%x", conn.ClientVersion()))
	return &client
}

// HandleRequests handles SSH requests
func (c *Client) HandleRequests() error {
	go func(in <-chan *ssh.Request) {
		for req := range in {
			logrus.Debugf("HandleRequest: %v", req)
			if req.WantReply {
				req.Reply(false, nil)
			}
		}
	}(c.Reqs)
	return nil
}

// HandleChannels handles SSH channels
func (c *Client) HandleChannels() error {
	for newChannel := range c.Chans {
		if err := c.HandleChannel(newChannel); err != nil {
			return err
		}
	}
	return nil
}

// HandleChannel handles one SSH channel
func (c *Client) HandleChannel(newChannel ssh.NewChannel) error {
	if newChannel.ChannelType() != "session" {
		logrus.Debugf("Unknown channel type: %s", newChannel.ChannelType())
		newChannel.Reject(ssh.UnknownChannelType, "unknown channel type")
		return nil
	}

	channel, requests, err := newChannel.Accept()
	if err != nil {
		logrus.Errorf("newChannel.Accept failed: %v", err)
		return err
	}
	logrus.Debugf("HandleChannel.channel: %v", channel)

	logrus.Debug("Creating pty...")
	f, tty, err := pty.Open()
	if err != nil {
		logrus.Errorf("pty.Open failed: %v", err)
		return nil
	}
	c.Tty = tty
	c.Pty = f

	c.HandleChannelRequests(channel, requests)

	return nil
}

// HandleChannelRequests handles channel requests
func (c *Client) HandleChannelRequests(channel ssh.Channel, requests <-chan *ssh.Request) {
	go func(in <-chan *ssh.Request) {
		for req := range in {
			ok := false
			switch req.Type {
			case "shell":
				logrus.Debugf("HandleChannelRequests.req shell")
				if len(req.Payload) != 0 {
					break
				}
				ok = true

				args := []string{"run", "-it", "--rm", c.Conn.User(), "/bin/sh"}
				logrus.Debugf("Executing 'docker %s'", strings.Join(args, " "))
				cmd := exec.Command("docker", args...)
				cmd.Env = c.Env.List()

				defer c.Tty.Close()
				cmd.Stdout = c.Tty
				cmd.Stdin = c.Tty
				cmd.Stderr = c.Tty
				cmd.SysProcAttr = &syscall.SysProcAttr{
					Setctty: true,
					Setsid:  true,
				}

				err := cmd.Start()
				if err != nil {
					logrus.Warnf("cmd.Start failed: %v", err)
					continue
				}

				var once sync.Once
				close := func() {
					channel.Close()
					logrus.Debug("session closed")
				}

				go func() {
					io.Copy(channel, c.Pty)
					once.Do(close)
				}()

				go func() {
					io.Copy(c.Pty, channel)
					once.Do(close)
				}()

			case "pty-req":
				ok = true
				termLen := req.Payload[3]
				c.Env["TERM"] = string(req.Payload[4 : termLen+4])
				w, h := parseDims(req.Payload[termLen+4:])
				SetWinsize(c.Pty.Fd(), w, h)
				logrus.Debugf("HandleChannelRequests.req pty-req: TERM=%q w=%q h=%q", c.Env["TERM"], int(w), int(h))

			case "window-change":
				w, h := parseDims(req.Payload)
				SetWinsize(c.Pty.Fd(), w, h)
				continue

			case "env":
				keyLen := req.Payload[3]
				key := string(req.Payload[4 : keyLen+4])
				valueLen := req.Payload[keyLen+7]
				value := string(req.Payload[keyLen+8 : keyLen+8+valueLen])
				logrus.Debugf("HandleChannelRequets.req 'env': %s=%q", key, value)
				c.Env[key] = value

			default:
				logrus.Debugf("Unhandled request type: %q: %v", req.Type, req)
			}

			if req.WantReply {
				if !ok {
					logrus.Debugf("Declining %s request...", req.Type)
				}
				req.Reply(ok, nil)
			}
		}
	}(requests)
}
<commit_msg>Added client and channel counters<commit_after>package ssh2docker

import (
	"fmt"
	"io"
	"os"
	"os/exec"
	"strings"
	"sync"
	"syscall"

	"github.com/Sirupsen/logrus"
	"github.com/kr/pty"
	"golang.org/x/crypto/ssh"
)

var clientCounter = 0

// Client is one client connection
type Client struct {
	Idx        int
	ChannelIdx int
	Conn       *ssh.ServerConn
	Chans      <-chan ssh.NewChannel
	Reqs       <-chan *ssh.Request
	Server     *Server
	Pty, Tty   *os.File
	Env        Environment
}

// NewClient initializes a new client
func NewClient(conn *ssh.ServerConn, chans <-chan ssh.NewChannel, reqs <-chan *ssh.Request, server *Server) *Client {
	client := Client{
		Idx:        clientCounter,
		ChannelIdx: 0,
		Conn:       conn,
		Chans:      chans,
		Reqs:       reqs,
		Server:     server,
		Env: Environment{
			"TERM":              os.Getenv("TERM"),
			"DOCKER_HOST":       os.Getenv("DOCKER_HOST"),
			"DOCKER_CERT_PATH":  os.Getenv("DOCKER_CERT_PATH"),
			"DOCKER_TLS_VERIFY": os.Getenv("DOCKER_TLS_VERIFY"),
		},
	}
	clientCounter++

	logrus.Infof("NewClient (%d): User=%q, ClientVersion=%q", client.Idx, conn.User(), fmt.Sprintf("%x", conn.ClientVersion()))
	return &client
}

// HandleRequests handles SSH requests
func (c *Client) HandleRequests() error {
	go func(in <-chan *ssh.Request) {
		for req := range in {
			logrus.Debugf("HandleRequest: %v", req)
			if req.WantReply {
				req.Reply(false, nil)
			}
		}
	}(c.Reqs)
	return nil
}

// HandleChannels handles SSH channels
func (c *Client) HandleChannels() error {
	for newChannel := range c.Chans {
		if err := c.HandleChannel(newChannel); err != nil {
			return err
		}
	}
	return nil
}

// HandleChannel handles one SSH channel
func (c *Client) HandleChannel(newChannel ssh.NewChannel) error {
	if newChannel.ChannelType() != "session" {
		logrus.Debugf("Unknown channel type: %s", newChannel.ChannelType())
		newChannel.Reject(ssh.UnknownChannelType, "unknown channel type")
		return nil
	}

	channel, requests, err := newChannel.Accept()
	if err != nil {
		logrus.Errorf("newChannel.Accept failed: %v", err)
		return err
	}
	c.ChannelIdx++
	logrus.Debugf("HandleChannel.channel (client=%d channel=%d): %v", c.Idx, c.ChannelIdx, channel)

	logrus.Debug("Creating pty...")
	f, tty, err := pty.Open()
	if err != nil {
		logrus.Errorf("pty.Open failed: %v", err)
		return nil
	}
	c.Tty = tty
	c.Pty = f

	c.HandleChannelRequests(channel, requests)

	return nil
}

// HandleChannelRequests handles channel requests
func (c *Client) HandleChannelRequests(channel ssh.Channel, requests <-chan *ssh.Request) {
	go func(in <-chan *ssh.Request) {
		for req := range in {
			ok := false
			switch req.Type {
			case "shell":
				logrus.Debugf("HandleChannelRequests.req shell")
				if len(req.Payload) != 0 {
					break
				}
				ok = true

				args := []string{"run", "-it", "--rm", c.Conn.User(), "/bin/sh"}
				logrus.Debugf("Executing 'docker %s'", strings.Join(args, " "))
				cmd := exec.Command("docker", args...)
				cmd.Env = c.Env.List()

				defer c.Tty.Close()
				cmd.Stdout = c.Tty
				cmd.Stdin = c.Tty
				cmd.Stderr = c.Tty
				cmd.SysProcAttr = &syscall.SysProcAttr{
					Setctty: true,
					Setsid:  true,
				}

				err := cmd.Start()
				if err != nil {
					logrus.Warnf("cmd.Start failed: %v", err)
					continue
				}

				var once sync.Once
				close := func() {
					channel.Close()
					logrus.Debug("session closed")
				}

				go func() {
					io.Copy(channel, c.Pty)
					once.Do(close)
				}()

				go func() {
					io.Copy(c.Pty, channel)
					once.Do(close)
				}()

			case "pty-req":
				ok = true
				termLen := req.Payload[3]
				c.Env["TERM"] = string(req.Payload[4 : termLen+4])
				w, h := parseDims(req.Payload[termLen+4:])
				SetWinsize(c.Pty.Fd(), w, h)
				logrus.Debugf("HandleChannelRequests.req pty-req: TERM=%q w=%q h=%q", c.Env["TERM"], int(w), int(h))

			case "window-change":
				w, h := parseDims(req.Payload)
				SetWinsize(c.Pty.Fd(), w, h)
				continue

			case "env":
				keyLen := req.Payload[3]
				key := string(req.Payload[4 : keyLen+4])
				valueLen := req.Payload[keyLen+7]
				value := string(req.Payload[keyLen+8 : keyLen+8+valueLen])
				logrus.Debugf("HandleChannelRequets.req 'env': %s=%q", key, value)
				c.Env[key] = value

			default:
				logrus.Debugf("Unhandled request type: %q: %v", req.Type, req)
			}

			if req.WantReply {
				if !ok {
					logrus.Debugf("Declining %s request...", req.Type)
				}
				req.Reply(ok, nil)
			}
		}
	}(requests)
}
<|endoftext|>


<commit_before>package gedis

import (
	"net"
)

// A wrapper to net.Conn that handles writing/reading to a Redis
// server
type Client struct {
	conn net.Conn
}

// Connect to a Redis server on address, using the named network
//
// See documentation for net.Dial for more information on named
// networks.
func Dial(network, address string) (c Client, err error) {
	c.conn, err = net.Dial(network, address)
	return
}

// Close the connection to the Redis server
func (c *Client) Close() error {
	return c.conn.Close()
}

// Send a command to the Redis server and receive its reply
func (c *Client) Send(args ...interface{}) (interface{}, error) {
	_, err := Write(c.conn, args...)
	if err != nil {
		return nil, err
	}
	return Read(c.conn)
}
<commit_msg>Add Client.Read func<commit_after>package gedis

import (
	"net"
)

// A wrapper to net.Conn that handles writing/reading to a Redis
// server
type Client struct {
	conn net.Conn
}

// Connect to a Redis server on address, using the named network
//
// See documentation for net.Dial for more information on named
// networks.
func Dial(network, address string) (c Client, err error) {
	c.conn, err = net.Dial(network, address)
	return
}

// Close the connection to the Redis server
func (c *Client) Close() error {
	return c.conn.Close()
}

// Send a command to the Redis server and receive its reply
func (c *Client) Send(args ...interface{}) (interface{}, error) {
	_, err := Write(c.conn, args...)
	if err != nil {
		return nil, err
	}
	return c.Read()
}

// Reads from the client
//
// This is useful for cases like a monitor
func (c *Client) Read() (interface{}, error) {
	return Read(c.conn)
}
<|endoftext|>


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

import (
	"crypto/tls"
	"crypto/x509"
	"errors"
	"io/ioutil"
	"net/http"

	"github.com/mendersoftware/log"
)

var (
	errorLoadingClientCertificate = errors.New("Failed to load certificate and key")
	errorNoServerCertificateFound = errors.New("No server certificate is provided," +
		" use -trusted-certs with a proper certificate.")
	errorAddingServerCertificateToPool = errors.New("Adding server certificate " +
		"to trusted pool failed.")
)

//TODO: this will be hardcoded for now but should be configurable in future
const (
	defaultCertFile   = "/data/certfile.crt"
	defaultCertKey    = "/data/certkey.key"
	defaultServerCert = "/data/server.crt"
)

type authCmdLineArgsType struct {
	// hostname or address to bootstrap to
	bootstrapServer string
	certFile        string
	certKey         string
	serverCert      string
}

func (cred *authCmdLineArgsType) setDefaultKeysAndCerts(clientCert, clientKey,
	serverCert string) {
	if cred.certFile == "" {
		cred.certFile = clientCert
	}
	if cred.certKey == "" {
		cred.certKey = clientKey
	}
	if cred.serverCert == "" {
		cred.serverCert = serverCert
	}
}

type authCredsType struct {
	// Cert+privkey that authenticates this client
	clientCert tls.Certificate
	// Trusted server certificates
	trustedCerts x509.CertPool
}

// Client represents the http(s) client used for network communication.
//
type Client struct {
	BaseURL string
	authCredsType
	HTTPClient *http.Client
}

// Client initialization

// NewClient returnes a poiter to newly initialized client
func NewClient(args authCmdLineArgsType) (*Client, error) {
	var httpsClient Client
	args.setDefaultKeysAndCerts(defaultCertFile, defaultCertKey, defaultServerCert)

	if err := httpsClient.initServerTrust(args); err != nil {
		return nil, err
	}
	if err := httpsClient.initClientCert(args); err != nil {
		return nil, err
	}

	tlsConf := tls.Config{
		RootCAs:      &httpsClient.trustedCerts,
		Certificates: []tls.Certificate{httpsClient.clientCert},
	}

	transport := http.Transport{
		TLSClientConfig: &tlsConf,
	}

	httpsClient.HTTPClient = &http.Client{
		Transport: &transport,
	}

	return &httpsClient, nil
}

func (c *Client) initServerTrust(args authCmdLineArgsType) error {

	if args.serverCert == "" {
		panic("certificate should be replaced with default one")
	}
	trustedCerts := *x509.NewCertPool()
	certPoolAppendCertsFromFile(&trustedCerts, args.serverCert)

	if len(trustedCerts.Subjects()) == 0 {
		return errorAddingServerCertificateToPool
	}
	c.trustedCerts = trustedCerts
	return nil
}

func (c *Client) initClientCert(args authCmdLineArgsType) error {
	clientCert, err := tls.LoadX509KeyPair(args.certFile, args.certKey)
	if err != nil {
		return errorLoadingClientCertificate
	}
	c.clientCert = clientCert
	return nil
}

func certPoolAppendCertsFromFile(s *x509.CertPool, f string) bool {
	cacert, err := ioutil.ReadFile(f)
	if err != nil {
		log.Warnln("Error reading certificate file ", err)
		return false
	}

	return s.AppendCertsFromPEM(cacert)
}

// Client request sending and parsing

type clientRequestType struct {
	reqType string
	request string
}

type clientRequester interface {
	formatRequest() clientRequestType
	actOnResponse(http.Response, []byte) error
	getClient() Client
}

func makeJobDone(req clientRequester) error {
	request := req.formatRequest()
	client := req.getClient()

	response, data, err := client.sendRequest(request.reqType, request.request)
	if err != nil {
		return err
	}

	return req.actOnResponse(*response, data)
}

func (c *Client) sendRequest(reqType string, request string) (*http.Response, []byte, error) {

	switch reqType {
	//TODO: in future we can use different request types
	case http.MethodGet:
		log.Debug("Sending HTTP GET: ", request)

		response, err := c.HTTPClient.Get(request)
		if err != nil {
			return nil, nil, err
		}
		defer response.Body.Close()

		log.Debug("Received headers:", response.Header)
		log.Debug("Received response: ", response.Status)

		respData, err := ioutil.ReadAll(response.Body)
		if err != nil {
			return nil, nil, err
		}
		log.Debug("Received response body: ", string(respData))

		return response, respData, nil
	}
	panic("unknown http request")
}
<commit_msg>More review fixes.<commit_after>// Copyright 2016 Mender Software AS
//
//    Licensed under the Apache License, Version 2.0 (the "License");
//    you may not use this file except in compliance with the License.
//    You may obtain a copy of the License at
//
//        http://www.apache.org/licenses/LICENSE-2.0
//
//    Unless required by applicable law or agreed to in writing, software
//    distributed under the License is distributed on an "AS IS" BASIS,
//    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//    See the License for the specific language governing permissions and
//    limitations under the License.
package main

import (
	"crypto/tls"
	"crypto/x509"
	"errors"
	"io/ioutil"
	"net/http"

	"github.com/mendersoftware/log"
)

var (
	errorLoadingClientCertificate = errors.New("Failed to load certificate and key")
	errorNoServerCertificateFound = errors.New("No server certificate is provided," +
		" use -trusted-certs with a proper certificate.")
	errorAddingServerCertificateToPool = errors.New("Adding server certificate " +
		"to trusted pool failed.")
)

//TODO: this will be hardcoded for now but should be configurable in future
const (
	defaultCertFile   = "/data/certfile.crt"
	defaultCertKey    = "/data/certkey.key"
	defaultServerCert = "/data/server.crt"
)

type authCmdLineArgsType struct {
	// hostname or address to bootstrap to
	bootstrapServer string
	certFile        string
	certKey         string
	serverCert      string
}

func (cred *authCmdLineArgsType) setDefaultKeysAndCerts(clientCert, clientKey,
	serverCert string) {
	if cred.certFile == "" {
		cred.certFile = clientCert
	}
	if cred.certKey == "" {
		cred.certKey = clientKey
	}
	if cred.serverCert == "" {
		cred.serverCert = serverCert
	}
}

type authCredsType struct {
	// Cert+privkey that authenticates this client
	clientCert tls.Certificate
	// Trusted server certificates
	trustedCerts x509.CertPool
}

// Client represents the http(s) client used for network communication.
//
type Client struct {
	BaseURL string
	authCredsType
	HTTPClient *http.Client
}

// Client initialization

// NewClient returnes a poiter to newly initialized client
func NewClient(args authCmdLineArgsType) (*Client, error) {
	var httpsClient Client
	args.setDefaultKeysAndCerts(defaultCertFile, defaultCertKey, defaultServerCert)

	if err := httpsClient.initServerTrust(args); err != nil {
		return nil, err
	}
	if err := httpsClient.initClientCert(args); err != nil {
		return nil, err
	}

	tlsConf := tls.Config{
		RootCAs:      &httpsClient.trustedCerts,
		Certificates: []tls.Certificate{httpsClient.clientCert},
	}

	transport := http.Transport{
		TLSClientConfig: &tlsConf,
	}

	httpsClient.HTTPClient = &http.Client{
		Transport: &transport,
	}

	return &httpsClient, nil
}

func (c *Client) initServerTrust(args authCmdLineArgsType) error {

	if args.serverCert == "" {
		return errors.New("Can not read server certificate")
	}
	trustedCerts := *x509.NewCertPool()
	certPoolAppendCertsFromFile(&trustedCerts, args.serverCert)

	if len(trustedCerts.Subjects()) == 0 {
		return errorAddingServerCertificateToPool
	}
	c.trustedCerts = trustedCerts
	return nil
}

func (c *Client) initClientCert(args authCmdLineArgsType) error {
	clientCert, err := tls.LoadX509KeyPair(args.certFile, args.certKey)
	if err != nil {
		return errorLoadingClientCertificate
	}
	c.clientCert = clientCert
	return nil
}

func certPoolAppendCertsFromFile(s *x509.CertPool, f string) bool {
	cacert, err := ioutil.ReadFile(f)
	if err != nil {
		log.Warnln("Error reading certificate file ", err)
		return false
	}

	return s.AppendCertsFromPEM(cacert)
}

// Client request sending and parsing

type clientRequestType struct {
	reqType string
	request string
}

type clientRequester interface {
	formatRequest() clientRequestType
	actOnResponse(http.Response, []byte) error
	getClient() Client
}

func makeJobDone(req clientRequester) error {
	request := req.formatRequest()
	client := req.getClient()

	response, err := client.sendRequest(request.reqType, request.request)
	if err != nil {
		return err
	}

	defer response.Body.Close()

	respBody, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return err
	}
	log.Debug("Received response body: ", string(respBody))

	return req.actOnResponse(*response, respBody)
}

func (c *Client) sendRequest(reqType string, request string) (*http.Response, error) {

	switch reqType {
	//TODO: in future we can use different request types
	case http.MethodGet:
		log.Debug("Sending HTTP GET: ", request)

		response, err := c.HTTPClient.Get(request)
		if err != nil {
			return nil, err
		}

		log.Debug("Received headers:", response.Header)
		log.Debug("Received response: ", response.Status)

		return response, nil
	}
	return nil, errors.New("trying to send unsupported request of type " + reqType)
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"regexp"
	"sort"
	"strconv"
	"strings"

	statsd "github.com/DataDog/datadog-go/statsd"
	log "github.com/Sirupsen/logrus"
)

const sampleRate = 1.0

const (
	routerMsg int = iota
	scalingMsg
	sampleMsg
	metricsTag
	releaseMsg
)

var routerMetricsKeys = []string{"dyno", "method", "status", "path", "host", "code", "desc", "at"}
var sampleMetricsKeys = []string{"source"}
var scalingMetricsKeys = []string{"mailer", "web"}
var customMetricsKeys = []string{"media_type", "output_type", "path", "version", "name"}

type Client struct {
	*statsd.Client
	ExcludedTags map[string]bool
}

var statusCode *regexp.Regexp = regexp.MustCompile(`^(?P<Family>\d)\d\d`)

func statsdClient(addr string) (*Client, error) {

	c, err := statsd.New(addr)
	return &Client{c, make(map[string]bool)}, err
}

func (c *Client) sendToStatsd(in chan *logMetrics) {

	var data *logMetrics
	var ok bool
	for {
		data, ok = <-in

		if !ok { //Exit, channel was closed
			return
		}

		log.WithFields(log.Fields{
			"type":   data.typ,
			"app":    data.app,
			"tags":   data.tags,
			"prefix": data.prefix,
		}).Debug("logMetrics received")

		if data.typ == routerMsg {
			c.sendRouterMsg(data)
		} else if data.typ == sampleMsg {
			c.sendSampleMsg(data)
		} else if data.typ == scalingMsg {
			c.sendEvents(*data.app, "heroku", data.events, *data.tags)
			c.sendScalingMsg(data)
		} else if data.typ == metricsTag {
			c.sendMetricsWithTags(data)
		} else if data.typ == releaseMsg {
			c.sendEvents(*data.app, "app", data.events, *data.tags)
		} else {
			log.WithField("type", data.typ).Warn("Unknown log message")
		}
	}
}

func (c *Client) sendEvents(app string, namespace string, events []string, tags []string) {
	for _, v := range events {
		event := statsd.NewEvent(namespace+"/api: "+app, v)
		event.Tags = tags
		c.Event(event)
		log.WithFields(log.Fields{
			"type":  "event",
			"app":   app,
			"value": v,
		}).Info("Event sent")
	}
}

func (c *Client) extractTags(tags []string, permittedTags []string, metrics map[string]logValue) []string {
	for _, mk := range permittedTags {
		if c.ExcludedTags[mk] {
			continue
		}
		if v, ok := metrics[mk]; ok {
			tags = append(tags, mk+":"+v.Val)
		}
	}
	sort.Strings(tags)
	return tags
}

func addStatusFamilyToTags(data *logMetrics, tags []string) []string {
	if val, ok := data.metrics["status"]; ok {
		match := statusCode.FindStringSubmatch(val.Val)
		if len(match) > 1 {
			tags = append(tags, "statusFamily:"+match[1]+"xx")
		}
	}
	return tags
}

func (c *Client) sendRouterMsg(data *logMetrics) {
	tags := c.extractTags(*data.tags, routerMetricsKeys, data.metrics)
	tags = addStatusFamilyToTags(data, tags)

	log.WithFields(log.Fields{
		"app":    *data.app,
		"tags":   *data.tags,
		"prefix": *data.prefix,
	}).Debug("sendRouterMsg")

	conn, err := strconv.ParseFloat(data.metrics["connect"].Val, 10)
	if err != nil {
		log.WithFields(log.Fields{
			"type":   "router",
			"err":    err,
			"metric": "connect",
		}).Info("Could not parse metric value")
		return
	}
	serv, err := strconv.ParseFloat(data.metrics["service"].Val, 10)
	if err != nil {
		log.WithFields(log.Fields{
			"type":   "router",
			"metric": "service",
			"err":    err,
		}).Info("Could not parse metric value")
		return
	}

	err = c.Histogram(*data.prefix+"heroku.router.request.connect", conn, tags, sampleRate)
	if err != nil {
		log.WithField("error", err).Info("Failed to send Histogram")
	}
	err = c.Histogram(*data.prefix+"heroku.router.request.service", serv, tags, sampleRate)
	if err != nil {
		log.WithField("error", err).Info("Failed to send Histogram")
	}
	if data.metrics["at"].Val == "error" {
		err = c.Count(*data.prefix+"heroku.router.error", 1, tags, 0.1)
		if err != nil {
			log.WithField("error", err).Info("Failed to send Count")
		}
	}
}

func (c *Client) sendSampleMsg(data *logMetrics) {
	tags := c.extractTags(*data.tags, sampleMetricsKeys, data.metrics)

	log.WithFields(log.Fields{
		"app":    *data.app,
		"tags":   tags,
		"prefix": *data.prefix,
	}).Debug("sendSampleMsg")

	for k, v := range data.metrics {
		if strings.Index(k, "#") != -1 {
			m := strings.Replace(strings.Split(k, "#")[1], "_", ".", -1)
			vnum, err := strconv.ParseFloat(v.Val, 10)
			if err == nil {
				err = c.Gauge(*data.prefix+"heroku.dyno."+m, vnum, tags, sampleRate)
				if err != nil {
					log.WithField("error", err).Info("Failed to send Gauge")
				}
			} else {
				log.WithFields(log.Fields{
					"type":   "sample",
					"metric": k,
					"err":    err,
				}).Info("Could not parse metric value")
			}
		}
	}
}

func (c *Client) sendScalingMsg(data *logMetrics) {
	tags := *data.tags

	log.WithFields(log.Fields{
		"app":    *data.app,
		"tags":   tags,
		"prefix": *data.prefix,
	}).Debug("sendScalingMsg")

	for _, mk := range scalingMetricsKeys {
		if v, ok := data.metrics[mk]; ok {
			vnum, err := strconv.ParseFloat(v.Val, 10)
			if err == nil {
				err = c.Gauge(*data.prefix+"heroku.dyno."+mk, vnum, tags, sampleRate)
				if err != nil {
					log.WithField("error", err).Info("Failed to send Gauge")
				}
			} else {
				log.WithFields(log.Fields{
					"type":   "scaling",
					"metric": mk,
					"err":    err,
				}).Info("Could not parse metric value")
			}
		}
	}
}

func (c *Client) sendMetric(metricType string, metricName string, value float64, tags []string) error {
	switch metricType {
	case "metric", "sample":
		return c.Gauge(metricName, value, tags, sampleRate)
	case "measure":
		return c.Histogram(metricName, value, tags, sampleRate)
	case "count":
		return c.Count(metricName, int64(value), tags, sampleRate)
	default:
		return errors.New("Unknown metric type" + metricType)
	}
}

func (c *Client) sendMetricsWithTags(data *logMetrics) {
	tags := *data.tags

Tags:
	for k, v := range data.metrics {
		if strings.Index(k, "tag#") != -1 {
			if _, err := strconv.Atoi(v.Val); err != nil {
				m := strings.Replace(strings.Split(k, "tag#")[1], "_", ".", -1)
				for _, mk := range customMetricsKeys {
					if m == mk {
						tags = append(tags, mk+":"+v.Val)
						continue Tags
					}
				}
			}
		}
	}

	log.WithFields(log.Fields{
		"app":    *data.app,
		"tags":   tags,
		"prefix": *data.prefix,
	}).Debug("sendMetricTag")

	for k, v := range data.metrics {
		if strings.Index(k, "#") != -1 {
			if vnum, err := strconv.ParseFloat(v.Val, 10); err == nil {
				keySplit := strings.Split(k, "#")
				metricType := keySplit[0]
				m := strings.Replace(keySplit[1], "_", ".", -1)
				err = c.sendMetric(metricType, *data.prefix+m, vnum, tags)
				if err != nil {
					log.WithField("error", err).Warning("Failed to send Gauge")
				}
			} else {
				log.WithFields(log.Fields{
					"type":   "metrics",
					"metric": k,
					"err":    err,
				}).Debug("Could not parse metric value")
			}
		}
	}
}
<commit_msg>using count for custom metrics<commit_after>package main

import (
	"errors"
	"regexp"
	"sort"
	"strconv"
	"strings"

	statsd "github.com/DataDog/datadog-go/statsd"
	log "github.com/Sirupsen/logrus"
)

const sampleRate = 1.0

const (
	routerMsg int = iota
	scalingMsg
	sampleMsg
	metricsTag
	releaseMsg
)

var routerMetricsKeys = []string{"dyno", "method", "status", "path", "host", "code", "desc", "at"}
var sampleMetricsKeys = []string{"source"}
var scalingMetricsKeys = []string{"mailer", "web"}
var customMetricsKeys = []string{"media_type", "output_type", "path", "version", "name"}

type Client struct {
	*statsd.Client
	ExcludedTags map[string]bool
}

var statusCode *regexp.Regexp = regexp.MustCompile(`^(?P<Family>\d)\d\d`)

func statsdClient(addr string) (*Client, error) {

	c, err := statsd.New(addr)
	return &Client{c, make(map[string]bool)}, err
}

func (c *Client) sendToStatsd(in chan *logMetrics) {

	var data *logMetrics
	var ok bool
	for {
		data, ok = <-in

		if !ok { //Exit, channel was closed
			return
		}

		log.WithFields(log.Fields{
			"type":   data.typ,
			"app":    data.app,
			"tags":   data.tags,
			"prefix": data.prefix,
		}).Debug("logMetrics received")

		if data.typ == routerMsg {
			c.sendRouterMsg(data)
		} else if data.typ == sampleMsg {
			c.sendSampleMsg(data)
		} else if data.typ == scalingMsg {
			c.sendEvents(*data.app, "heroku", data.events, *data.tags)
			c.sendScalingMsg(data)
		} else if data.typ == metricsTag {
			c.sendMetricsWithTags(data)
		} else if data.typ == releaseMsg {
			c.sendEvents(*data.app, "app", data.events, *data.tags)
		} else {
			log.WithField("type", data.typ).Warn("Unknown log message")
		}
	}
}

func (c *Client) sendEvents(app string, namespace string, events []string, tags []string) {
	for _, v := range events {
		event := statsd.NewEvent(namespace+"/api: "+app, v)
		event.Tags = tags
		c.Event(event)
		log.WithFields(log.Fields{
			"type":  "event",
			"app":   app,
			"value": v,
		}).Info("Event sent")
	}
}

func (c *Client) extractTags(tags []string, permittedTags []string, metrics map[string]logValue) []string {
	for _, mk := range permittedTags {
		if c.ExcludedTags[mk] {
			continue
		}
		if v, ok := metrics[mk]; ok {
			tags = append(tags, mk+":"+v.Val)
		}
	}
	sort.Strings(tags)
	return tags
}

func addStatusFamilyToTags(data *logMetrics, tags []string) []string {
	if val, ok := data.metrics["status"]; ok {
		match := statusCode.FindStringSubmatch(val.Val)
		if len(match) > 1 {
			tags = append(tags, "statusFamily:"+match[1]+"xx")
		}
	}
	return tags
}

func (c *Client) sendRouterMsg(data *logMetrics) {
	tags := c.extractTags(*data.tags, routerMetricsKeys, data.metrics)
	tags = addStatusFamilyToTags(data, tags)

	log.WithFields(log.Fields{
		"app":    *data.app,
		"tags":   *data.tags,
		"prefix": *data.prefix,
	}).Debug("sendRouterMsg")

	conn, err := strconv.ParseFloat(data.metrics["connect"].Val, 10)
	if err != nil {
		log.WithFields(log.Fields{
			"type":   "router",
			"err":    err,
			"metric": "connect",
		}).Info("Could not parse metric value")
		return
	}
	serv, err := strconv.ParseFloat(data.metrics["service"].Val, 10)
	if err != nil {
		log.WithFields(log.Fields{
			"type":   "router",
			"metric": "service",
			"err":    err,
		}).Info("Could not parse metric value")
		return
	}

	err = c.Histogram(*data.prefix+"heroku.router.request.connect", conn, tags, sampleRate)
	if err != nil {
		log.WithField("error", err).Info("Failed to send Histogram")
	}
	err = c.Histogram(*data.prefix+"heroku.router.request.service", serv, tags, sampleRate)
	if err != nil {
		log.WithField("error", err).Info("Failed to send Histogram")
	}
	if data.metrics["at"].Val == "error" {
		err = c.Count(*data.prefix+"heroku.router.error", 1, tags, 0.1)
		if err != nil {
			log.WithField("error", err).Info("Failed to send Count")
		}
	}
}

func (c *Client) sendSampleMsg(data *logMetrics) {
	tags := c.extractTags(*data.tags, sampleMetricsKeys, data.metrics)

	log.WithFields(log.Fields{
		"app":    *data.app,
		"tags":   tags,
		"prefix": *data.prefix,
	}).Debug("sendSampleMsg")

	for k, v := range data.metrics {
		if strings.Index(k, "#") != -1 {
			m := strings.Replace(strings.Split(k, "#")[1], "_", ".", -1)
			vnum, err := strconv.ParseFloat(v.Val, 10)
			if err == nil {
				err = c.Gauge(*data.prefix+"heroku.dyno."+m, vnum, tags, sampleRate)
				if err != nil {
					log.WithField("error", err).Info("Failed to send Gauge")
				}
			} else {
				log.WithFields(log.Fields{
					"type":   "sample",
					"metric": k,
					"err":    err,
				}).Info("Could not parse metric value")
			}
		}
	}
}

func (c *Client) sendScalingMsg(data *logMetrics) {
	tags := *data.tags

	log.WithFields(log.Fields{
		"app":    *data.app,
		"tags":   tags,
		"prefix": *data.prefix,
	}).Debug("sendScalingMsg")

	for _, mk := range scalingMetricsKeys {
		if v, ok := data.metrics[mk]; ok {
			vnum, err := strconv.ParseFloat(v.Val, 10)
			if err == nil {
				err = c.Gauge(*data.prefix+"heroku.dyno."+mk, vnum, tags, sampleRate)
				if err != nil {
					log.WithField("error", err).Info("Failed to send Gauge")
				}
			} else {
				log.WithFields(log.Fields{
					"type":   "scaling",
					"metric": mk,
					"err":    err,
				}).Info("Could not parse metric value")
			}
		}
	}
}

func (c *Client) sendMetric(metricType string, metricName string, value float64, tags []string) error {
	switch metricType {
	case "sample":
		return c.Gauge(metricName, value, tags, sampleRate)
	case "measure":
		return c.Histogram(metricName, value, tags, sampleRate)
	case "metric", "count":
		return c.Count(metricName, int64(value), tags, sampleRate)
	default:
		return errors.New("Unknown metric type" + metricType)
	}
}

func (c *Client) sendMetricsWithTags(data *logMetrics) {
	tags := *data.tags

Tags:
	for k, v := range data.metrics {
		if strings.Index(k, "tag#") != -1 {
			if _, err := strconv.Atoi(v.Val); err != nil {
				m := strings.Replace(strings.Split(k, "tag#")[1], "_", ".", -1)
				for _, mk := range customMetricsKeys {
					if m == mk {
						tags = append(tags, mk+":"+v.Val)
						continue Tags
					}
				}
			}
		}
	}

	log.WithFields(log.Fields{
		"app":    *data.app,
		"tags":   tags,
		"prefix": *data.prefix,
	}).Debug("sendMetricTag")

	for k, v := range data.metrics {
		if strings.Index(k, "#") != -1 {
			if vnum, err := strconv.ParseFloat(v.Val, 10); err == nil {
				keySplit := strings.Split(k, "#")
				metricType := keySplit[0]
				m := strings.Replace(keySplit[1], "_", ".", -1)
				err = c.sendMetric(metricType, *data.prefix+m, vnum, tags)
				if err != nil {
					log.WithField("error", err).Warning("Failed to send Gauge")
				}
			} else {
				log.WithFields(log.Fields{
					"type":   "metrics",
					"metric": k,
					"err":    err,
				}).Debug("Could not parse metric value")
			}
		}
	}
}
<|endoftext|>


<commit_before>package apns

import (
	"crypto/tls"
	"errors"
	"fmt"
	"log"
	"net"
	"time"
)

// Client contains the fields necessary to communicate
// with Apple, such as the gateway to use and your
// certificate contents.
//
// You'll need to provide your own CertificateFile
// and KeyFile to send notifications. Ideally, you'll
// just set the CertificateFile and KeyFile fields to
// a location on drive where the certs can be loaded,
// but if you prefer you can use the CertificateBase64
// and KeyBase64 fields to store the actual contents.
type Client struct {
	Gateway           string
	CertificateFile   string
	CertificateBase64 string
	KeyFile           string
	KeyBase64         string
	certificate       tls.Certificate
	apnsConnection    *tls.Conn
}

// BareClient can be used to set the contents of your
// certificate and key blocks manually.
func BareClient(gateway, certificateBase64, keyBase64 string) (c *Client) {
	c = new(Client)
	c.Gateway = gateway
	c.CertificateBase64 = certificateBase64
	c.KeyBase64 = keyBase64
	return
}

// NewClient assumes you'll be passing in paths that
// point to your certificate and key.
func NewClient(gateway, certificateFile, keyFile string) (c *Client) {
	c = new(Client)
	c.Gateway = gateway
	c.CertificateFile = certificateFile
	c.KeyFile = keyFile
	return
}

// Send connects to the APN service and sends your push notification.
// Remember that if the submission is successful, Apple won't reply.
func (client *Client) Send(pn *PushNotification) (resp *PushNotificationResponse) {
	resp = new(PushNotificationResponse)

	payload, err := pn.ToBytes()
	if err != nil {
		resp.Success = false
		resp.Error = err
		return
	}

	err = client.ConnectAndWrite(resp, payload)
	if err != nil {
		resp.Success = false
		resp.Error = err
		return
	}

	resp.Success = true
	resp.Error = nil

	return
}

// ConnectAndWrite establishes the connection to Apple and handles the
// transmission of your push notification, as well as waiting for a reply.
//
// In lieu of a timeout (which would be available in Go 1.1)
// we use a timeout channel pattern instead. We start two goroutines,
// one of which just sleeps for TimeoutSeconds seconds, while the other
// waits for a response from the Apple servers.
//
// Whichever channel puts data on first is the "winner". As such, it's
// possible to get a false positive if Apple takes a long time to respond.
// It's probably not a deal-breaker, but something to be aware of.
func (client *Client) ConnectAndWrite(resp *PushNotificationResponse, payload []byte) error {
	var bytesWritten int
	var err error

	if client.apnsConnection == nil {
		err = client.openConnection()
		if err != nil {
			return err
		}
	}

	bytesWritten, err = client.apnsConnection.Write(payload)
	if err != nil {
		if err.Error() != "use of closed network connection" {
			return err
		}

		// If the connection is closed, reconnect
		err = client.openConnection()
		if err != nil {
			return err
		}

		bytesWritten, err = client.apnsConnection.Write(payload)
		if err != nil {
			return err
		}
		if bytesWritten == 0 {
			client.apnsConnection.Close()
			return fmt.Errorf("Could not open connection to %s.  Please try again.", client.Gateway)
		}
	}

	// Create one channel that will serve to handle
	// timeouts when the notification succeeds.
	timeoutChannel := make(chan bool, 1)
	go func() {
		time.Sleep(time.Second * TimeoutSeconds)
		timeoutChannel <- true
	}()

	// This channel will contain the binary response
	// from Apple in the event of a failure.
	responseChannel := make(chan []byte, 1)
	go func() {
		buffer := make([]byte, 6, 6)
		client.apnsConnection.Read(buffer)
		responseChannel <- buffer
	}()

	// First one back wins!
	// The data structure for an APN response is as follows:
	//
	// command    -> 1 byte
	// status     -> 1 byte
	// identifier -> 4 bytes
	//
	// The first byte will always be set to 8.
	select {
	case r := <-responseChannel:
		resp.Success = false
		resp.AppleResponse = ApplePushResponses[r[1]]
		err = errors.New(resp.AppleResponse)
	case <-timeoutChannel:
		resp.Success = true
	}

	return err
}

// Opens a connection to the Apple APNS server
// The connection is created and persisted to the client's apnsConnection property
//	to save on the overhead of the crypto libraries.
func (client *Client) openConnection() error {
	log.Println("APNS - openConnection")
	err := client.getCertificate()
	if err != nil {
		return err
	}

	conf := &tls.Config{
		Certificates: []tls.Certificate{client.certificate},
	}

	conn, err := net.Dial("tcp", client.Gateway)
	if err != nil {
		return err
	}

	tlsConn := tls.Client(conn, conf)
	err = tlsConn.Handshake()
	if err != nil {
		return err
	}

	client.apnsConnection = tlsConn
	return nil
}

// Returns a certificate to use to send the notification.
// The certificate is only created once to save on
// the overhead of the crypto libraries.
func (client *Client) getCertificate() error {
	var err error

	if client.certificate.PrivateKey == nil {
		log.Println("APNS - getCertificate")
		if len(client.CertificateBase64) == 0 && len(client.KeyBase64) == 0 {
			// The user did not specify raw block contents, so check the filesystem.
			client.certificate, err = tls.LoadX509KeyPair(client.CertificateFile, client.KeyFile)
		} else {
			// The user provided the raw block contents, so use that.
			client.certificate, err = tls.X509KeyPair([]byte(client.CertificateBase64), []byte(client.KeyBase64))
		}
	}

	return err
}
<commit_msg>Added debugging to see where iOS send failed with EOF error<commit_after>package apns

import (
	"crypto/tls"
	"errors"
	"fmt"
	"log"
	"net"
	"time"
)

// Client contains the fields necessary to communicate
// with Apple, such as the gateway to use and your
// certificate contents.
//
// You'll need to provide your own CertificateFile
// and KeyFile to send notifications. Ideally, you'll
// just set the CertificateFile and KeyFile fields to
// a location on drive where the certs can be loaded,
// but if you prefer you can use the CertificateBase64
// and KeyBase64 fields to store the actual contents.
type Client struct {
	Gateway           string
	CertificateFile   string
	CertificateBase64 string
	KeyFile           string
	KeyBase64         string
	certificate       tls.Certificate
	apnsConnection    *tls.Conn
}

// BareClient can be used to set the contents of your
// certificate and key blocks manually.
func BareClient(gateway, certificateBase64, keyBase64 string) (c *Client) {
	c = new(Client)
	c.Gateway = gateway
	c.CertificateBase64 = certificateBase64
	c.KeyBase64 = keyBase64
	return
}

// NewClient assumes you'll be passing in paths that
// point to your certificate and key.
func NewClient(gateway, certificateFile, keyFile string) (c *Client) {
	c = new(Client)
	c.Gateway = gateway
	c.CertificateFile = certificateFile
	c.KeyFile = keyFile
	return
}

// Send connects to the APN service and sends your push notification.
// Remember that if the submission is successful, Apple won't reply.
func (client *Client) Send(pn *PushNotification) (resp *PushNotificationResponse) {
	resp = new(PushNotificationResponse)

	payload, err := pn.ToBytes()
	if err != nil {
		resp.Success = false
		resp.Error = err
		return
	}

	err = client.ConnectAndWrite(resp, payload)
	if err != nil {
		resp.Success = false
		resp.Error = err
		return
	}

	resp.Success = true
	resp.Error = nil

	return
}

// ConnectAndWrite establishes the connection to Apple and handles the
// transmission of your push notification, as well as waiting for a reply.
//
// In lieu of a timeout (which would be available in Go 1.1)
// we use a timeout channel pattern instead. We start two goroutines,
// one of which just sleeps for TimeoutSeconds seconds, while the other
// waits for a response from the Apple servers.
//
// Whichever channel puts data on first is the "winner". As such, it's
// possible to get a false positive if Apple takes a long time to respond.
// It's probably not a deal-breaker, but something to be aware of.
func (client *Client) ConnectAndWrite(resp *PushNotificationResponse, payload []byte) error {
	log.Println("APNS - ConnectAndWrite")

	var bytesWritten int
	var err error

	if client.apnsConnection == nil {
		err = client.openConnection()
		if err != nil {
			log.Printf("APNS - ConnectAndWrite: open connection (#1) error = %s\n", err.Error())
			return err
		}
	}

	bytesWritten, err = client.apnsConnection.Write(payload)
	if err != nil {
		if err.Error() != "use of closed network connection" {
			log.Printf("APNS - ConnectAndWrite: write (#1) error = %s\n", err.Error())
			return err
		}

		// If the connection is closed, reconnect
		err = client.openConnection()
		if err != nil {
			log.Printf("APNS - ConnectAndWrite: open connection (#2) error = %s\n", err.Error())
			return err
		}

		bytesWritten, err = client.apnsConnection.Write(payload)
		if err != nil {
			log.Printf("APNS - ConnectAndWrite: write (#2) error = %s\n", err.Error())
			return err
		}
		if bytesWritten == 0 {
			log.Printf("APNS - ConnectAndWrite: write (#2), bytes written = 0\n")
			client.apnsConnection.Close()
			return fmt.Errorf("Could not open connection to %s.  Please try again.", client.Gateway)
		}
	}

	// Create one channel that will serve to handle
	// timeouts when the notification succeeds.
	timeoutChannel := make(chan bool, 1)
	go func() {
		time.Sleep(time.Second * TimeoutSeconds)
		timeoutChannel <- true
	}()

	// This channel will contain the binary response
	// from Apple in the event of a failure.
	responseChannel := make(chan []byte, 1)
	go func() {
		buffer := make([]byte, 6, 6)
		client.apnsConnection.Read(buffer)
		responseChannel <- buffer
	}()

	// First one back wins!
	// The data structure for an APN response is as follows:
	//
	// command    -> 1 byte
	// status     -> 1 byte
	// identifier -> 4 bytes
	//
	// The first byte will always be set to 8.
	select {
	case r := <-responseChannel:
		resp.Success = false
		resp.AppleResponse = ApplePushResponses[r[1]]
		err = errors.New(resp.AppleResponse)
	case <-timeoutChannel:
		resp.Success = true
	}

	return err
}

// Opens a connection to the Apple APNS server
// The connection is created and persisted to the client's apnsConnection property
//	to save on the overhead of the crypto libraries.
func (client *Client) openConnection() error {
	log.Println("APNS - openConnection")
	err := client.getCertificate()
	if err != nil {
		return err
	}

	conf := &tls.Config{
		Certificates: []tls.Certificate{client.certificate},
	}

	conn, err := net.Dial("tcp", client.Gateway)
	if err != nil {
		return err
	}

	tlsConn := tls.Client(conn, conf)
	err = tlsConn.Handshake()
	if err != nil {
		return err
	}

	client.apnsConnection = tlsConn
	return nil
}

// Returns a certificate to use to send the notification.
// The certificate is only created once to save on
// the overhead of the crypto libraries.
func (client *Client) getCertificate() error {
	var err error

	if client.certificate.PrivateKey == nil {
		log.Println("APNS - getCertificate")
		if len(client.CertificateBase64) == 0 && len(client.KeyBase64) == 0 {
			// The user did not specify raw block contents, so check the filesystem.
			client.certificate, err = tls.LoadX509KeyPair(client.CertificateFile, client.KeyFile)
		} else {
			// The user provided the raw block contents, so use that.
			client.certificate, err = tls.X509KeyPair([]byte(client.CertificateBase64), []byte(client.KeyBase64))
		}
	}

	return err
}
<|endoftext|>


<commit_before>package amazonpa

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/xml"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"strings"
	"time"
)

// ItemLookupQuery describes the allowed parameters for a ItemLookup request
type ItemLookupQuery struct {
	ItemIDs        []string
	IDType         string
	MerchantID     string
	ResponseGroups []string
}

// Client provides the functions to interact with the API
type Client struct {
	accessKeyID  string
	secretKey    string
	associateTag string
	region       string
	secure       bool
}

// NewClient returns a new Client
func NewClient(accessKeyID string, secretKey string, associateTag string, region string, secure bool) *Client {
	c := Client{accessKeyID, secretKey, associateTag, region, secure}

	return &c
}

// NewRequest returns a request with basic parameters
func (client Client) NewRequest(operation string) *Request {

	request := Request{}

	if client.secure {
		request.scheme = "https"
	} else {
		request.scheme = "http"
	}

	request.endpoint = Endpoints[client.region]
	request.endpointURI = EndpointURI

	request.parameters = map[string]string{
		"Service":        "AWSECommerceService",
		"AWSAccessKeyId": client.accessKeyID,
		"AssociateTag":   client.associateTag,
		"Version":        "2013-08-01",
		"Operation":      operation,
		"Timestamp":      time.Now().Format(time.RFC3339),
	}

	return &request
}

// SignRequest produces the signature for the given query string
func (client Client) SignRequest(request *Request) {
	signable := fmt.Sprintf("GET\n%s\n%s\n%s", request.endpoint, request.endpointURI, request.QueryString())

	hasher := hmac.New(sha256.New, []byte(client.secretKey))
	hasher.Write([]byte(signable))

	request.signature = base64.StdEncoding.EncodeToString(hasher.Sum(nil))
}

// ProcessRequest takes a request and queries the API
func (client Client) ProcessRequest(request *Request) ([]byte, error) {

	// Sign the request
	client.SignRequest(request)

	// Execute the HTTP request
	var httpResponse *http.Response
	var err error
	var contents []byte

	requestURL, err := request.SignedURL()

	if err != nil {
		return nil, errors.New("amazonpa: cannot get the signed request URL")
	}

	httpResponse, err = http.Get(requestURL)

	if err != nil {
		return nil, errors.New("amazonpa: error processing the http request")
	}

	contents, err = ioutil.ReadAll(httpResponse.Body)
	httpResponse.Body.Close()

	if err != nil {
		return nil, errors.New("amazonpa: error while reading the server response")
	}

	return contents, nil
}

// ItemLookup performs an ItemLookup request
func (client Client) ItemLookup(query ItemLookupQuery) (*ItemLookupResponse, error) {

	request := client.NewRequest("ItemLookup")

	request.SetParameter("ItemId", strings.Join(query.ItemIDs, ","))
	request.SetParameter("ResponseGroup", strings.Join(query.ResponseGroups, ","))
	request.SetParameter("IdType", query.IDType)
	request.SetParameter("MerchantId", query.MerchantID)

	xmlData, err := client.ProcessRequest(request)

	if err != nil {
		return nil, err
	}

	var response ItemLookupResponse
	xml.Unmarshal(xmlData, &response)

	if response.Items.Request.IsValid != true {
		return nil, errors.New("amazonpa: request is invalid")
	}

	return &response, nil
}
<commit_msg>Return invalid response struct<commit_after>package amazonpa

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"encoding/xml"
	"errors"
	"fmt"
	"io/ioutil"
	"net/http"
	"strings"
	"time"
)

// ItemLookupQuery describes the allowed parameters for a ItemLookup request
type ItemLookupQuery struct {
	ItemIDs        []string
	IDType         string
	MerchantID     string
	ResponseGroups []string
}

// Client provides the functions to interact with the API
type Client struct {
	accessKeyID  string
	secretKey    string
	associateTag string
	region       string
	secure       bool
}

// NewClient returns a new Client
func NewClient(accessKeyID string, secretKey string, associateTag string, region string, secure bool) *Client {
	c := Client{accessKeyID, secretKey, associateTag, region, secure}

	return &c
}

// NewRequest returns a request with basic parameters
func (client Client) NewRequest(operation string) *Request {

	request := Request{}

	if client.secure {
		request.scheme = "https"
	} else {
		request.scheme = "http"
	}

	request.endpoint = Endpoints[client.region]
	request.endpointURI = EndpointURI

	request.parameters = map[string]string{
		"Service":        "AWSECommerceService",
		"AWSAccessKeyId": client.accessKeyID,
		"AssociateTag":   client.associateTag,
		"Version":        "2013-08-01",
		"Operation":      operation,
		"Timestamp":      time.Now().Format(time.RFC3339),
	}

	return &request
}

// SignRequest produces the signature for the given query string
func (client Client) SignRequest(request *Request) {
	signable := fmt.Sprintf("GET\n%s\n%s\n%s", request.endpoint, request.endpointURI, request.QueryString())

	hasher := hmac.New(sha256.New, []byte(client.secretKey))
	hasher.Write([]byte(signable))

	request.signature = base64.StdEncoding.EncodeToString(hasher.Sum(nil))
}

// ProcessRequest takes a request and queries the API
func (client Client) ProcessRequest(request *Request) ([]byte, error) {

	// Sign the request
	client.SignRequest(request)

	// Execute the HTTP request
	var httpResponse *http.Response
	var err error
	var contents []byte

	requestURL, err := request.SignedURL()

	if err != nil {
		return nil, errors.New("amazonpa: cannot get the signed request URL")
	}

	httpResponse, err = http.Get(requestURL)

	if err != nil {
		return nil, errors.New("amazonpa: error processing the http request")
	}

	contents, err = ioutil.ReadAll(httpResponse.Body)
	httpResponse.Body.Close()

	if err != nil {
		return nil, errors.New("amazonpa: error while reading the server response")
	}

	return contents, nil
}

// ItemLookup performs an ItemLookup request
func (client Client) ItemLookup(query ItemLookupQuery) (*ItemLookupResponse, error) {

	request := client.NewRequest("ItemLookup")

	request.SetParameter("ItemId", strings.Join(query.ItemIDs, ","))
	request.SetParameter("ResponseGroup", strings.Join(query.ResponseGroups, ","))
	request.SetParameter("IdType", query.IDType)
	request.SetParameter("MerchantId", query.MerchantID)

	xmlData, err := client.ProcessRequest(request)

	if err != nil {
		return nil, err
	}

	var response ItemLookupResponse
	xml.Unmarshal(xmlData, &response)

	if response.Items.Request.IsValid != true {
		return &response, errors.New("amazonpa: request is invalid")
	}

	return &response, nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2012, Suryandaru Triandana <syndtr@gmail.com>
// All rights reserved.
//
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

// This LevelDB Go implementation is based on LevelDB C++ implementation.
// Which contains the following header:
//   Copyright (c) 2011 The LevelDB Authors. All rights reserved.
//   Use of this source code is governed by a BSD-style license that can be
//   found in the LEVELDBCPP_LICENSE file. See the LEVELDBCPP_AUTHORS file
//   for names of contributors.

package db

import (
	"leveldb"
	"sync/atomic"
)

var levelMaxSize [kNumLevels]float64

func init() {
	// Precompute max size of each level
	for level := range levelMaxSize {
		res := float64(10 * 1048576)
		for n := level; n > 1; n-- {
			res *= 10
		}
		levelMaxSize[level] = res
	}
}

type version struct {
	s *session

	tables [kNumLevels]tFiles

	// Level that should be compacted next and its compaction score.
	// Score < 1 means compaction is not strictly needed.  These fields
	// are initialized by ComputeCompaction()
	compactionLevel int
	compactionScore float64

	seekCompactionLevel int
	seekCompactionTable *tFile
}

func (v *version) get(key iKey, ro *leveldb.ReadOptions) (value []byte, cState bool, err error) {
	s := v.s
	ukey := key.ukey()

	var fLevel int
	var fTable *tFile
	fSeek := true

	// We can search level-by-level since entries never hop across
	// levels. Therefore we are guaranteed that if we find data
	// in an smaller level, later levels are irrelevant.
	for level, ts := range v.tables {
		if len(ts) == 0 {
			continue
		}

		if level == 0 {
			// Level-0 files may overlap each other. Find all files that
			// overlap user_key and process them in order from newest to
			var tmp tFiles
			for _, t := range ts {
				if s.cmp.Compare(ukey, t.smallest.ukey()) >= 0 &&
					s.cmp.Compare(ukey, t.largest.ukey()) <= 0 {
					tmp = append(tmp, t)
				}
			}

			if len(tmp) == 0 {
				continue
			}

			tmp.sort(tFileSorterNewest(nil))
			ts = tmp
		} else {
			i := ts.search(key, s.icmp)
			if i >= len(ts) || s.cmp.Compare(ukey, ts[i].smallest.ukey()) < 0 {
				continue
			}

			ts = ts[i : i+1]
		}

		for _, t := range ts {
			if fSeek {
				if fTable != nil {
					if atomic.AddInt32(&fTable.seekLeft, -1) <= 0 {
						s.st.Lock()
						if v.seekCompactionTable == nil {
							v.seekCompactionLevel = fLevel
							v.seekCompactionTable = fTable
							cState = true
						}
						s.st.Unlock()
					}

					fSeek = false
				} else {
					fLevel = level
					fTable = t
				}
			}

			var rkey, rval []byte
			rkey, rval, err = s.tops.get(t, key, ro)
			if err == leveldb.ErrNotFound {
				continue
			} else if err != nil {
				return
			}

			pk := iKey(rkey).parse()
			if pk == nil {
				err = leveldb.ErrCorrupt("internal key corrupted")
				return
			}
			if s.cmp.Compare(ukey, pk.ukey) == 0 {
				switch pk.vtype {
				case tVal:
					value = rval
				case tDel:
					err = leveldb.ErrNotFound
				default:
					panic("not reached")
				}
				return
			}
		}
	}

	err = leveldb.ErrNotFound
	return
}

func (v *version) newStaging() *versionStaging {
	return &versionStaging{base: v}
}

// Spawn a new version based on this version.
func (v *version) spawn(r *sessionRecord) *version {
	staging := v.newStaging()
	staging.commit(r)
	return staging.finish()
}

func (v *version) fillRecord(r *sessionRecord) {
	for level, ts := range v.tables {
		for _, t := range ts {
			r.addTableFile(level, t)
		}
	}
}

func (v *version) tLen(level int) int {
	return len(v.tables[0])
}

func (v *version) pickLevel(min, max []byte) (level int) {
	icmp := v.s.icmp
	ucmp := icmp.cmp

	if !v.tables[0].isOverlaps(min, max, false, icmp) {
		var r tFiles
		for ; level < kMaxMemCompactLevel; level++ {
			if v.tables[level+1].isOverlaps(min, max, true, icmp) {
				break
			}
			v.tables[level+2].getOverlaps(min, max, &r, true, ucmp)
			if r.size() > kMaxGrandParentOverlapBytes {
				break
			}
		}
	}

	return
}

func (v *version) computeCompaction() {
	// Precomputed best level for next compaction
	var bestLevel int = -1
	var bestScore float64 = -1

	for level, ff := range v.tables {
		var score float64
		if level == 0 {
			// We treat level-0 specially by bounding the number of files
			// instead of number of bytes for two reasons:
			//
			// (1) With larger write-buffer sizes, it is nice not to do too
			// many level-0 compactions.
			//
			// (2) The files in level-0 are merged on every read and
			// therefore we wish to avoid too many files when the individual
			// file size is small (perhaps because of a small write-buffer
			// setting, or very high compression ratios, or lots of
			// overwrites/deletions).
			score = float64(len(ff)) / kL0_CompactionTrigger
		} else {
			score = float64(ff.size()) / levelMaxSize[level]
		}

		if score > bestScore {
			bestLevel = level
			bestScore = score
		}
	}

	v.compactionLevel = bestLevel
	v.compactionScore = bestScore
}

type versionStaging struct {
	base   *version
	tables [kNumLevels]struct {
		added   map[uint64]ntRecord
		deleted map[uint64]struct{}
	}
}

func (p *versionStaging) commit(r *sessionRecord) {
	// deleted tables
	for _, tr := range r.deletedTables {
		tm := &(p.tables[tr.level])
		if tm.deleted == nil {
			tm.deleted = make(map[uint64]struct{})
		}
		tm.deleted[tr.num] = struct{}{}
		delete(tm.added, tr.num)
	}

	// new tables
	for _, tr := range r.newTables {
		tm := &(p.tables[tr.level])
		if tm.added == nil {
			tm.added = make(map[uint64]ntRecord)
		}
		tm.added[tr.num] = tr
		delete(tm.deleted, tr.num)
	}
}

func (p *versionStaging) finish() *version {
	// build new version
	s := p.base.s
	nv := &version{s: s}
	sorter := &tFileSorterKey{cmp: s.icmp}
	for level, tm := range p.tables {
		ot := p.base.tables[level]
		nt := make(tFiles, 0, len(tm.added)+len(ot)-len(tm.deleted))
		// old tables
		for _, t := range ot {
			if _, ok := tm.deleted[t.file.Number()]; ok {
				continue
			}
			if _, ok := tm.added[t.file.Number()]; ok {
				continue
			}
			nt = append(nt, t)
		}
		// new tables
		for _, tr := range tm.added {
			nt = append(nt, tr.makeFile(s))
		}
		// sort tables
		nt.sort(sorter)
		nv.tables[level] = nt
	}

	// compute compaction score for new version
	nv.computeCompaction()

	return nv
}
<commit_msg>db: version: *versionStaging.commit(): make sure tm.[added|deleted] not nil<commit_after>// Copyright (c) 2012, Suryandaru Triandana <syndtr@gmail.com>
// All rights reserved.
//
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

// This LevelDB Go implementation is based on LevelDB C++ implementation.
// Which contains the following header:
//   Copyright (c) 2011 The LevelDB Authors. All rights reserved.
//   Use of this source code is governed by a BSD-style license that can be
//   found in the LEVELDBCPP_LICENSE file. See the LEVELDBCPP_AUTHORS file
//   for names of contributors.

package db

import (
	"leveldb"
	"sync/atomic"
)

var levelMaxSize [kNumLevels]float64

func init() {
	// Precompute max size of each level
	for level := range levelMaxSize {
		res := float64(10 * 1048576)
		for n := level; n > 1; n-- {
			res *= 10
		}
		levelMaxSize[level] = res
	}
}

type version struct {
	s *session

	tables [kNumLevels]tFiles

	// Level that should be compacted next and its compaction score.
	// Score < 1 means compaction is not strictly needed.  These fields
	// are initialized by ComputeCompaction()
	compactionLevel int
	compactionScore float64

	seekCompactionLevel int
	seekCompactionTable *tFile
}

func (v *version) get(key iKey, ro *leveldb.ReadOptions) (value []byte, cState bool, err error) {
	s := v.s
	ukey := key.ukey()

	var fLevel int
	var fTable *tFile
	fSeek := true

	// We can search level-by-level since entries never hop across
	// levels. Therefore we are guaranteed that if we find data
	// in an smaller level, later levels are irrelevant.
	for level, ts := range v.tables {
		if len(ts) == 0 {
			continue
		}

		if level == 0 {
			// Level-0 files may overlap each other. Find all files that
			// overlap user_key and process them in order from newest to
			var tmp tFiles
			for _, t := range ts {
				if s.cmp.Compare(ukey, t.smallest.ukey()) >= 0 &&
					s.cmp.Compare(ukey, t.largest.ukey()) <= 0 {
					tmp = append(tmp, t)
				}
			}

			if len(tmp) == 0 {
				continue
			}

			tmp.sort(tFileSorterNewest(nil))
			ts = tmp
		} else {
			i := ts.search(key, s.icmp)
			if i >= len(ts) || s.cmp.Compare(ukey, ts[i].smallest.ukey()) < 0 {
				continue
			}

			ts = ts[i : i+1]
		}

		for _, t := range ts {
			if fSeek {
				if fTable != nil {
					if atomic.AddInt32(&fTable.seekLeft, -1) <= 0 {
						s.st.Lock()
						if v.seekCompactionTable == nil {
							v.seekCompactionLevel = fLevel
							v.seekCompactionTable = fTable
							cState = true
						}
						s.st.Unlock()
					}

					fSeek = false
				} else {
					fLevel = level
					fTable = t
				}
			}

			var rkey, rval []byte
			rkey, rval, err = s.tops.get(t, key, ro)
			if err == leveldb.ErrNotFound {
				continue
			} else if err != nil {
				return
			}

			pk := iKey(rkey).parse()
			if pk == nil {
				err = leveldb.ErrCorrupt("internal key corrupted")
				return
			}
			if s.cmp.Compare(ukey, pk.ukey) == 0 {
				switch pk.vtype {
				case tVal:
					value = rval
				case tDel:
					err = leveldb.ErrNotFound
				default:
					panic("not reached")
				}
				return
			}
		}
	}

	err = leveldb.ErrNotFound
	return
}

func (v *version) newStaging() *versionStaging {
	return &versionStaging{base: v}
}

// Spawn a new version based on this version.
func (v *version) spawn(r *sessionRecord) *version {
	staging := v.newStaging()
	staging.commit(r)
	return staging.finish()
}

func (v *version) fillRecord(r *sessionRecord) {
	for level, ts := range v.tables {
		for _, t := range ts {
			r.addTableFile(level, t)
		}
	}
}

func (v *version) tLen(level int) int {
	return len(v.tables[0])
}

func (v *version) pickLevel(min, max []byte) (level int) {
	icmp := v.s.icmp
	ucmp := icmp.cmp

	if !v.tables[0].isOverlaps(min, max, false, icmp) {
		var r tFiles
		for ; level < kMaxMemCompactLevel; level++ {
			if v.tables[level+1].isOverlaps(min, max, true, icmp) {
				break
			}
			v.tables[level+2].getOverlaps(min, max, &r, true, ucmp)
			if r.size() > kMaxGrandParentOverlapBytes {
				break
			}
		}
	}

	return
}

func (v *version) computeCompaction() {
	// Precomputed best level for next compaction
	var bestLevel int = -1
	var bestScore float64 = -1

	for level, ff := range v.tables {
		var score float64
		if level == 0 {
			// We treat level-0 specially by bounding the number of files
			// instead of number of bytes for two reasons:
			//
			// (1) With larger write-buffer sizes, it is nice not to do too
			// many level-0 compactions.
			//
			// (2) The files in level-0 are merged on every read and
			// therefore we wish to avoid too many files when the individual
			// file size is small (perhaps because of a small write-buffer
			// setting, or very high compression ratios, or lots of
			// overwrites/deletions).
			score = float64(len(ff)) / kL0_CompactionTrigger
		} else {
			score = float64(ff.size()) / levelMaxSize[level]
		}

		if score > bestScore {
			bestLevel = level
			bestScore = score
		}
	}

	v.compactionLevel = bestLevel
	v.compactionScore = bestScore
}

type versionStaging struct {
	base   *version
	tables [kNumLevels]struct {
		added   map[uint64]ntRecord
		deleted map[uint64]struct{}
	}
}

func (p *versionStaging) commit(r *sessionRecord) {
	// deleted tables
	for _, tr := range r.deletedTables {
		tm := &(p.tables[tr.level])
		if tm.deleted == nil {
			tm.deleted = make(map[uint64]struct{})
		}
		tm.deleted[tr.num] = struct{}{}
		if tm.added != nil {
			delete(tm.added, tr.num)
		}
	}

	// new tables
	for _, tr := range r.newTables {
		tm := &(p.tables[tr.level])
		if tm.added == nil {
			tm.added = make(map[uint64]ntRecord)
		}
		tm.added[tr.num] = tr
		if tm.deleted != nil {
			delete(tm.deleted, tr.num)
		}
	}
}

func (p *versionStaging) finish() *version {
	// build new version
	s := p.base.s
	nv := &version{s: s}
	sorter := &tFileSorterKey{cmp: s.icmp}
	for level, tm := range p.tables {
		ot := p.base.tables[level]
		nt := make(tFiles, 0, len(tm.added)+len(ot)-len(tm.deleted))
		// old tables
		for _, t := range ot {
			if _, ok := tm.deleted[t.file.Number()]; ok {
				continue
			}
			if _, ok := tm.added[t.file.Number()]; ok {
				continue
			}
			nt = append(nt, t)
		}
		// new tables
		for _, tr := range tm.added {
			nt = append(nt, tr.makeFile(s))
		}
		// sort tables
		nt.sort(sorter)
		nv.tables[level] = nt
	}

	// compute compaction score for new version
	nv.computeCompaction()

	return nv
}
<|endoftext|>


<commit_before>package lfs

import (
	"sync"
	"sync/atomic"

	"github.com/github/git-lfs/git"
	"github.com/github/git-lfs/vendor/_nuts/github.com/rubyist/tracerx"
)

const (
	batchSize = 100
)

type Transferable interface {
	Check() (*objectResource, error)
	Transfer(CopyCallback) error
	Object() *objectResource
	Oid() string
	Size() int64
	Name() string
	SetObject(*objectResource)
}

// TransferQueue provides a queue that will allow concurrent transfers.
type TransferQueue struct {
	retrying      uint32
	meter         *ProgressMeter
	workers       int // Number of transfer workers to spawn
	transferKind  string
	errors        []error
	transferables map[string]Transferable
	retries       []Transferable
	batcher       *Batcher
	apic          chan Transferable // Channel for processing individual API requests
	transferc     chan Transferable // Channel for processing transfers
	retriesc      chan Transferable // Channel for processing retries
	errorc        chan error        // Channel for processing errors
	watchers      []chan string
	errorwait     sync.WaitGroup
	retrywait     sync.WaitGroup
	wait          sync.WaitGroup
}

// newTransferQueue builds a TransferQueue, allowing `workers` concurrent transfers.
func newTransferQueue(files int, size int64, dryRun bool) *TransferQueue {
	q := &TransferQueue{
		meter:         NewProgressMeter(files, size, dryRun),
		apic:          make(chan Transferable, batchSize),
		transferc:     make(chan Transferable, batchSize),
		retriesc:      make(chan Transferable, batchSize),
		errorc:        make(chan error),
		workers:       Config.ConcurrentTransfers(),
		transferables: make(map[string]Transferable),
	}

	q.errorwait.Add(1)
	q.retrywait.Add(1)

	q.run()

	return q
}

// Add adds a Transferable to the transfer queue.
func (q *TransferQueue) Add(t Transferable) {
	q.wait.Add(1)
	q.transferables[t.Oid()] = t

	if q.batcher != nil {
		q.batcher.Add(t)
		return
	}

	q.apic <- t
}

// Wait waits for the queue to finish processing all transfers. Once Wait is
// called, Add will no longer add transferables to the queue. Any failed
// transfers will be automatically retried once.
func (q *TransferQueue) Wait() {
	if q.batcher != nil {
		q.batcher.Exit()
	}

	q.wait.Wait()

	// Handle any retries
	close(q.retriesc)
	q.retrywait.Wait()
	atomic.StoreUint32(&q.retrying, 1)

	if len(q.retries) > 0 && q.batcher != nil {
		tracerx.Printf("tq: retrying %d failed transfers", len(q.retries))
		for _, t := range q.retries {
			q.Add(t)
		}
		q.batcher.Exit()
		q.wait.Wait()
	}

	atomic.StoreUint32(&q.retrying, 0)

	close(q.apic)
	close(q.transferc)
	close(q.errorc)

	for _, watcher := range q.watchers {
		close(watcher)
	}

	q.meter.Finish()
	q.errorwait.Wait()
}

// Watch returns a channel where the queue will write the OID of each transfer
// as it completes. The channel will be closed when the queue finishes processing.
func (q *TransferQueue) Watch() chan string {
	c := make(chan string, batchSize)
	q.watchers = append(q.watchers, c)
	return c
}

// individualApiRoutine processes the queue of transfers one at a time by making
// a POST call for each object, feeding the results to the transfer workers.
// If configured, the object transfers can still happen concurrently, the
// sequential nature here is only for the meta POST calls.
func (q *TransferQueue) individualApiRoutine(apiWaiter chan interface{}) {
	for t := range q.apic {
		obj, err := t.Check()
		if err != nil {
			if q.canRetry(err) {
				q.retry(t)
			} else {
				q.errorc <- err
			}
			q.wait.Done()
			continue
		}

		if apiWaiter != nil { // Signal to launch more individual api workers
			q.meter.Start()
			select {
			case apiWaiter <- 1:
			default:
			}
		}

		if obj != nil {
			t.SetObject(obj)
			q.meter.Add(t.Name())
			q.transferc <- t
		} else {
			q.meter.Skip(t.Size())
			q.wait.Done()
		}
	}
}

// legacyFallback is used when a batch request is made to a server that does
// not support the batch endpoint. When this happens, the Transferables are
// fed from the batcher into apic to be processed individually.
func (q *TransferQueue) legacyFallback(failedBatch []Transferable) {
	tracerx.Printf("tq: batch api not implemented, falling back to individual")

	q.launchIndividualApiRoutines()

	for _, t := range failedBatch {
		q.apic <- t
	}

	for {
		batch := q.batcher.Next()
		if batch == nil {
			break
		}

		for _, t := range batch {
			q.apic <- t
		}
	}
}

// batchApiRoutine processes the queue of transfers using the batch endpoint,
// making only one POST call for all objects. The results are then handed
// off to the transfer workers.
func (q *TransferQueue) batchApiRoutine() {
	var startProgress sync.Once

	for {
		batch := q.batcher.Next()
		if batch == nil {
			break
		}

		tracerx.Printf("tq: sending batch of size %d", len(batch))

		transfers := make([]*objectResource, 0, len(batch))
		for _, t := range batch {
			transfers = append(transfers, &objectResource{Oid: t.Oid(), Size: t.Size()})
		}

		objects, err := Batch(transfers, q.transferKind)
		if err != nil {
			if IsNotImplementedError(err) {
				git.Config.SetLocal("", "lfs.batch", "false")

				go q.legacyFallback(batch)
				return
			}

			if q.canRetry(err) {
				for _, t := range batch {
					q.retry(t)
				}
			} else {
				q.errorc <- err
			}

			q.wait.Add(-len(transfers))
			continue
		}

		startProgress.Do(q.meter.Start)

		for _, o := range objects {
			if _, ok := o.Rel(q.transferKind); ok {
				// This object has an error
				if o.Error != nil {
					q.errorc <- Error(o.Error)
					q.meter.Skip(o.Size)
					q.wait.Done()
					continue
				}

				// This object needs to be transferred
				if transfer, ok := q.transferables[o.Oid]; ok {
					transfer.SetObject(o)
					q.meter.Add(transfer.Name())
					q.transferc <- transfer
				} else {
					q.meter.Skip(transfer.Size())
					q.wait.Done()
				}
			} else {
				q.meter.Skip(o.Size)
				q.wait.Done()
			}
		}
	}
}

// This goroutine collects errors returned from transfers
func (q *TransferQueue) errorCollector() {
	for err := range q.errorc {
		q.errors = append(q.errors, err)
	}
	q.errorwait.Done()
}

func (q *TransferQueue) retryCollector() {
	for t := range q.retriesc {
		q.retries = append(q.retries, t)
	}
	q.retrywait.Done()
}

func (q *TransferQueue) transferWorker() {
	for transfer := range q.transferc {
		cb := func(total, read int64, current int) error {
			q.meter.TransferBytes(q.transferKind, transfer.Name(), read, total, current)
			return nil
		}

		if err := transfer.Transfer(cb); err != nil {
			if q.canRetry(err) {
				tracerx.Printf("tq: retrying object %s", transfer.Oid())
				q.retry(transfer)
			} else {
				q.errorc <- err
			}
		} else {
			oid := transfer.Oid()
			for _, c := range q.watchers {
				c <- oid
			}
		}

		q.meter.FinishTransfer(transfer.Name())

		q.wait.Done()
	}
}

// launchIndividualApiRoutines first launches a single api worker. When it
// receives the first successful api request it launches workers - 1 more
// workers. This prevents being prompted for credentials multiple times at once
// when they're needed.
func (q *TransferQueue) launchIndividualApiRoutines() {
	go func() {
		apiWaiter := make(chan interface{})
		go q.individualApiRoutine(apiWaiter)

		<-apiWaiter

		for i := 0; i < q.workers-1; i++ {
			go q.individualApiRoutine(nil)
		}
	}()
}

// run starts the transfer queue, doing individual or batch transfers depending
// on the Config.BatchTransfer() value. run will transfer files sequentially or
// concurrently depending on the Config.ConcurrentTransfers() value.
func (q *TransferQueue) run() {
	go q.errorCollector()
	go q.retryCollector()

	tracerx.Printf("tq: starting %d transfer workers", q.workers)
	for i := 0; i < q.workers; i++ {
		go q.transferWorker()
	}

	if Config.BatchTransfer() {
		tracerx.Printf("tq: running as batched queue, batch size of %d", batchSize)
		q.batcher = NewBatcher(batchSize)
		go q.batchApiRoutine()
	} else {
		tracerx.Printf("tq: running as individual queue")
		q.launchIndividualApiRoutines()
	}
}

func (q *TransferQueue) retry(t Transferable) {
	q.retriesc <- t
}

func (q *TransferQueue) canRetry(err error) bool {
	if !IsRetriableError(err) || atomic.LoadUint32(&q.retrying) == 1 {
		return false
	}

	return true
}

// Errors returns any errors encountered during transfer.
func (q *TransferQueue) Errors() []error {
	return q.errors
}
<commit_msg>ンンン ンンン ンン<commit_after>package lfs

import (
	"sync"
	"sync/atomic"

	"github.com/github/git-lfs/git"
	"github.com/github/git-lfs/vendor/_nuts/github.com/rubyist/tracerx"
)

const (
	batchSize = 100
)

type Transferable interface {
	Check() (*objectResource, error)
	Transfer(CopyCallback) error
	Object() *objectResource
	Oid() string
	Size() int64
	Name() string
	SetObject(*objectResource)
}

// TransferQueue provides a queue that will allow concurrent transfers.
type TransferQueue struct {
	retrying      uint32
	meter         *ProgressMeter
	workers       int // Number of transfer workers to spawn
	transferKind  string
	errors        []error
	transferables map[string]Transferable
	retries       []Transferable
	batcher       *Batcher
	apic          chan Transferable // Channel for processing individual API requests
	transferc     chan Transferable // Channel for processing transfers
	retriesc      chan Transferable // Channel for processing retries
	errorc        chan error        // Channel for processing errors
	watchers      []chan string
	errorwait     sync.WaitGroup
	retrywait     sync.WaitGroup
	wait          sync.WaitGroup
}

// newTransferQueue builds a TransferQueue, allowing `workers` concurrent transfers.
func newTransferQueue(files int, size int64, dryRun bool) *TransferQueue {
	q := &TransferQueue{
		meter:         NewProgressMeter(files, size, dryRun),
		apic:          make(chan Transferable, batchSize),
		transferc:     make(chan Transferable, batchSize),
		retriesc:      make(chan Transferable, batchSize),
		errorc:        make(chan error),
		workers:       Config.ConcurrentTransfers(),
		transferables: make(map[string]Transferable),
	}

	q.errorwait.Add(1)
	q.retrywait.Add(1)

	q.run()

	return q
}

// Add adds a Transferable to the transfer queue.
func (q *TransferQueue) Add(t Transferable) {
	q.wait.Add(1)
	q.transferables[t.Oid()] = t

	if q.batcher != nil {
		q.batcher.Add(t)
		return
	}

	q.apic <- t
}

// Wait waits for the queue to finish processing all transfers. Once Wait is
// called, Add will no longer add transferables to the queue. Any failed
// transfers will be automatically retried once.
func (q *TransferQueue) Wait() {
	if q.batcher != nil {
		q.batcher.Exit()
	}

	q.wait.Wait()

	// Handle any retries
	close(q.retriesc)
	q.retrywait.Wait()
	atomic.StoreUint32(&q.retrying, 1)

	if len(q.retries) > 0 {
		tracerx.Printf("tq: retrying %d failed transfers", len(q.retries))
		for _, t := range q.retries {
			q.Add(t)
		}
		if q.batcher != nil {
			q.batcher.Exit()
		}
		q.wait.Wait()
	}

	atomic.StoreUint32(&q.retrying, 0)

	close(q.apic)
	close(q.transferc)
	close(q.errorc)

	for _, watcher := range q.watchers {
		close(watcher)
	}

	q.meter.Finish()
	q.errorwait.Wait()
}

// Watch returns a channel where the queue will write the OID of each transfer
// as it completes. The channel will be closed when the queue finishes processing.
func (q *TransferQueue) Watch() chan string {
	c := make(chan string, batchSize)
	q.watchers = append(q.watchers, c)
	return c
}

// individualApiRoutine processes the queue of transfers one at a time by making
// a POST call for each object, feeding the results to the transfer workers.
// If configured, the object transfers can still happen concurrently, the
// sequential nature here is only for the meta POST calls.
func (q *TransferQueue) individualApiRoutine(apiWaiter chan interface{}) {
	for t := range q.apic {
		obj, err := t.Check()
		if err != nil {
			if q.canRetry(err) {
				q.retry(t)
			} else {
				q.errorc <- err
			}
			q.wait.Done()
			continue
		}

		if apiWaiter != nil { // Signal to launch more individual api workers
			q.meter.Start()
			select {
			case apiWaiter <- 1:
			default:
			}
		}

		if obj != nil {
			t.SetObject(obj)
			q.meter.Add(t.Name())
			q.transferc <- t
		} else {
			q.meter.Skip(t.Size())
			q.wait.Done()
		}
	}
}

// legacyFallback is used when a batch request is made to a server that does
// not support the batch endpoint. When this happens, the Transferables are
// fed from the batcher into apic to be processed individually.
func (q *TransferQueue) legacyFallback(failedBatch []Transferable) {
	tracerx.Printf("tq: batch api not implemented, falling back to individual")

	q.launchIndividualApiRoutines()

	for _, t := range failedBatch {
		q.apic <- t
	}

	for {
		batch := q.batcher.Next()
		if batch == nil {
			break
		}

		for _, t := range batch {
			q.apic <- t
		}
	}
}

// batchApiRoutine processes the queue of transfers using the batch endpoint,
// making only one POST call for all objects. The results are then handed
// off to the transfer workers.
func (q *TransferQueue) batchApiRoutine() {
	var startProgress sync.Once

	for {
		batch := q.batcher.Next()
		if batch == nil {
			break
		}

		tracerx.Printf("tq: sending batch of size %d", len(batch))

		transfers := make([]*objectResource, 0, len(batch))
		for _, t := range batch {
			transfers = append(transfers, &objectResource{Oid: t.Oid(), Size: t.Size()})
		}

		objects, err := Batch(transfers, q.transferKind)
		if err != nil {
			if IsNotImplementedError(err) {
				git.Config.SetLocal("", "lfs.batch", "false")

				go q.legacyFallback(batch)
				return
			}

			if q.canRetry(err) {
				for _, t := range batch {
					q.retry(t)
				}
			} else {
				q.errorc <- err
			}

			q.wait.Add(-len(transfers))
			continue
		}

		startProgress.Do(q.meter.Start)

		for _, o := range objects {
			if _, ok := o.Rel(q.transferKind); ok {
				// This object has an error
				if o.Error != nil {
					q.errorc <- Error(o.Error)
					q.meter.Skip(o.Size)
					q.wait.Done()
					continue
				}

				// This object needs to be transferred
				if transfer, ok := q.transferables[o.Oid]; ok {
					transfer.SetObject(o)
					q.meter.Add(transfer.Name())
					q.transferc <- transfer
				} else {
					q.meter.Skip(transfer.Size())
					q.wait.Done()
				}
			} else {
				q.meter.Skip(o.Size)
				q.wait.Done()
			}
		}
	}
}

// This goroutine collects errors returned from transfers
func (q *TransferQueue) errorCollector() {
	for err := range q.errorc {
		q.errors = append(q.errors, err)
	}
	q.errorwait.Done()
}

func (q *TransferQueue) retryCollector() {
	for t := range q.retriesc {
		q.retries = append(q.retries, t)
	}
	q.retrywait.Done()
}

func (q *TransferQueue) transferWorker() {
	for transfer := range q.transferc {
		cb := func(total, read int64, current int) error {
			q.meter.TransferBytes(q.transferKind, transfer.Name(), read, total, current)
			return nil
		}

		if err := transfer.Transfer(cb); err != nil {
			if q.canRetry(err) {
				tracerx.Printf("tq: retrying object %s", transfer.Oid())
				q.retry(transfer)
			} else {
				q.errorc <- err
			}
		} else {
			oid := transfer.Oid()
			for _, c := range q.watchers {
				c <- oid
			}
		}

		q.meter.FinishTransfer(transfer.Name())

		q.wait.Done()
	}
}

// launchIndividualApiRoutines first launches a single api worker. When it
// receives the first successful api request it launches workers - 1 more
// workers. This prevents being prompted for credentials multiple times at once
// when they're needed.
func (q *TransferQueue) launchIndividualApiRoutines() {
	go func() {
		apiWaiter := make(chan interface{})
		go q.individualApiRoutine(apiWaiter)

		<-apiWaiter

		for i := 0; i < q.workers-1; i++ {
			go q.individualApiRoutine(nil)
		}
	}()
}

// run starts the transfer queue, doing individual or batch transfers depending
// on the Config.BatchTransfer() value. run will transfer files sequentially or
// concurrently depending on the Config.ConcurrentTransfers() value.
func (q *TransferQueue) run() {
	go q.errorCollector()
	go q.retryCollector()

	tracerx.Printf("tq: starting %d transfer workers", q.workers)
	for i := 0; i < q.workers; i++ {
		go q.transferWorker()
	}

	if Config.BatchTransfer() {
		tracerx.Printf("tq: running as batched queue, batch size of %d", batchSize)
		q.batcher = NewBatcher(batchSize)
		go q.batchApiRoutine()
	} else {
		tracerx.Printf("tq: running as individual queue")
		q.launchIndividualApiRoutines()
	}
}

func (q *TransferQueue) retry(t Transferable) {
	q.retriesc <- t
}

func (q *TransferQueue) canRetry(err error) bool {
	if !IsRetriableError(err) || atomic.LoadUint32(&q.retrying) == 1 {
		return false
	}

	return true
}

// Errors returns any errors encountered during transfer.
func (q *TransferQueue) Errors() []error {
	return q.errors
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/binary"
	"flag"
	"github.com/goerlang/epmd"
	"log"
	"net"
	"os"
	"os/signal"
	"runtime/pprof"
	"strconv"
	"time"
)

var noEpmd bool
var listenPort string
var regLimit int
var unregTTL int
var cpuProfile string

func init() {
	flag.StringVar(&listenPort, "port", "4369", "listen port")
	flag.BoolVar(&noEpmd, "no-epmd", false, "disable epmd")
	flag.IntVar(&regLimit, "nodes-limit", 1000, "limit size of registration table to prune unregistered nodes")
	flag.IntVar(&unregTTL, "unreg-ttl", 10, "prune unregistered nodes if unregistration older than this value in minutes")
	flag.StringVar(&cpuProfile, "profile-cpu", "", "profile CPU to file")
}

type regAns struct {
	reply   []byte
	isClose bool
}

type regReq struct {
	buf     []byte
	replyTo chan regAns
	conn    net.Conn
}

func main() {
	flag.Parse()
	if cpuProfile != "" {
		f, err := os.Create(cpuProfile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}
	stopCh := make(chan bool)

	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	go func() {
		for sig := range c {
			log.Printf("Signal %#v", sig)
			stopCh <- true
		}
	}()

	if !cliEnabled() {
		var err error
		var l net.Listener
		if !noEpmd {
			l, err = net.Listen("tcp", net.JoinHostPort("", listenPort))
		}
		if err != nil || noEpmd {
			// Cannot bind, eclus instance already running, connect to it
			eclusCli()
		} else {
			if !noEpmd {
				epm := make(chan regReq, 10)
				go epmReg(epm)
				go func() {
					for {
						conn, err := l.Accept()
						log.Printf("Accept new")
						if err != nil {
							log.Printf(err.Error())
						} else {
							go mLoop(conn, epm)
						}
					}
				}()
			}
			if nodeEnabled() {
				go runNode()
			}
			<-stopCh
		}
	} else {
		eclusCli()
	}
}

type nodeRec struct {
	*epmd.NodeInfo
	Time   time.Time
	Active bool
	conn   net.Conn
}

func epmReg(in <-chan regReq) {
	var nReg = make(map[string]*nodeRec)
	for {
		select {
		case req := <-in:
			buf := req.buf
			if len(buf) == 0 {
				rs := len(nReg)
				log.Printf("REG %d records", rs)
				now := time.Now()

				for node, rec := range nReg {
					if rec.conn == req.conn {
						log.Printf("Connection for %s dropped", node)
						nReg[node].Active = false
						nReg[node].Time = now
						nReg[node].conn = nil
					} else if rs > regLimit && !rec.Active && now.Sub(rec.Time).Minutes() > float64(unregTTL) {
						log.Printf("REG prune %s:%+v", node, rec)
						delete(nReg, node)
					}
				}
				continue
			}
			replyTo := req.replyTo
			log.Printf("IN: %v", buf)
			switch epmd.MessageId(buf[0]) {
			case epmd.ALIVE2_REQ:
				nConn := req.conn
				nInfo := epmd.Read_ALIVE2_REQ(buf)
				log.Printf("NodeInfo: %+v", nInfo)
				var reply []byte

				if rec, ok := nReg[nInfo.Name]; ok {
					log.Printf("Node %s found", nInfo.Name)
					if rec.Active {
						log.Printf("Node %s is running", nInfo.Name)
						reply = epmd.Compose_ALIVE2_RESP(false)
					} else {
						log.Printf("Node %s is not running", nInfo.Name)
						rec.conn = nConn

						nInfo.Creation = (rec.Creation % 3) + 1
						rec.NodeInfo = nInfo
						rec.Active = true
						reply = epmd.Compose_ALIVE2_RESP(true, rec.NodeInfo)
					}
				} else {
					log.Printf("New node %s", nInfo.Name)
					nInfo.Creation = 1
					rec := &nodeRec{
						NodeInfo: nInfo,
						conn:     nConn,
						Time:     time.Now(),
						Active:   true,
					}
					nReg[nInfo.Name] = rec
					reply = epmd.Compose_ALIVE2_RESP(true, rec.NodeInfo)
				}
				replyTo <- regAns{reply: reply, isClose: false}
			case epmd.PORT_PLEASE2_REQ:
				nName := epmd.Read_PORT_PLEASE2_REQ(buf)
				var reply []byte
				if rec, ok := nReg[nName]; ok && rec.Active {
					reply = epmd.Compose_PORT2_RESP(rec.NodeInfo)
				} else {
					reply = epmd.Compose_PORT2_RESP(nil)
				}
				replyTo <- regAns{reply: reply, isClose: true}
			case epmd.STOP_REQ:
				nName := epmd.Read_STOP_REQ(buf)
				var reply []byte
				if rec, ok := nReg[nName]; ok && rec.Active {
					// TODO: stop node
					reply = epmd.Compose_STOP_RESP(true)
				} else {
					reply = epmd.Compose_STOP_RESP(false)
				}
				replyTo <- regAns{reply: reply, isClose: true}
			case epmd.NAMES_REQ, epmd.DUMP_REQ:
				lp, err := strconv.Atoi(listenPort)
				if err != nil {
					log.Printf("Cannot convert %s to integer", listenPort)
					replyTo <- regAns{reply: nil, isClose: true}
				} else {
					replyB := new(bytes.Buffer)
					epmd.Compose_START_NAMES_RESP(replyB, lp)
					for _, rec := range nReg {
						if epmd.MessageId(buf[0]) == epmd.NAMES_REQ {
							if rec.Active {
								epmd.Append_NAMES_RESP(replyB, rec.NodeInfo)
							} else {
								if rec.Active {
									epmd.Append_DUMP_RESP_ACTIVE(replyB, rec.NodeInfo)
								} else {
									epmd.Append_DUMP_RESP_UNUSED(replyB, rec.NodeInfo)
								}
							}
						}
					}
					replyTo <- regAns{reply: replyB.Bytes(), isClose: true}
				}
			case epmd.KILL_REQ:
				reply := epmd.Compose_KILL_RESP()
				replyTo <- regAns{reply: reply, isClose: true}
			default:
				switch cliMessageId(buf[0]) {
				case REQ_NAMES:
					reply := ansNames(nReg)
					replyTo <- regAns{reply: reply, isClose: true}
				default:
					replyTo <- regAns{reply: nil, isClose: true}
				}
			}
		}
	}
}

func mLoop(c net.Conn, epm chan regReq) {
	buf := make([]byte, 1024)
	for {
		n, err := c.Read(buf)
		if err != nil {
			c.Close()
			log.Printf("Stop loop: %v", err)
			epm <- regReq{buf: []byte{}, conn: c}
			return
		}
		length := binary.BigEndian.Uint16(buf[0:2])
		if length != uint16(n-2) {
			log.Printf("Incomplete packet from erlang node to epmd: %d from %d", n, length)
			break
		}
		log.Printf("Read %d, %d: %v", n, length, buf[2:n])
		if isClose := handleMsg(c, buf[2:n], epm); isClose {
			break
		}
	}
	c.Close()
}

func handleMsg(c net.Conn, buf []byte, epm chan regReq) bool {
	myChan := make(chan regAns)
	epm <- regReq{buf: buf, replyTo: myChan, conn: c}
	select {
	case ans := <-myChan:
		log.Printf("Got reply: %+v", ans)
		if ans.reply != nil {
			c.Write(ans.reply)
		}
		return ans.isClose
	}
	return true
}
<commit_msg>Run if another epmd instance is up<commit_after>package main

import (
	"bytes"
	"encoding/binary"
	"flag"
	"github.com/goerlang/epmd"
	"log"
	"net"
	"os"
	"os/signal"
	"runtime/pprof"
	"strconv"
	"time"
)

var noEpmd bool
var listenPort string
var regLimit int
var unregTTL int
var cpuProfile string

func init() {
	flag.StringVar(&listenPort, "port", "4369", "listen port")
	flag.BoolVar(&noEpmd, "no-epmd", false, "disable epmd")
	flag.IntVar(&regLimit, "nodes-limit", 1000, "limit size of registration table to prune unregistered nodes")
	flag.IntVar(&unregTTL, "unreg-ttl", 10, "prune unregistered nodes if unregistration older than this value in minutes")
	flag.StringVar(&cpuProfile, "profile-cpu", "", "profile CPU to file")
}

type regAns struct {
	reply   []byte
	isClose bool
}

type regReq struct {
	buf     []byte
	replyTo chan regAns
	conn    net.Conn
}

func main() {
	flag.Parse()
	if cpuProfile != "" {
		f, err := os.Create(cpuProfile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}
	stopCh := make(chan bool)

	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	go func() {
		for sig := range c {
			log.Printf("Signal %#v", sig)
			stopCh <- true
		}
	}()

	if !cliEnabled() {
		var err error
		var l net.Listener
		if !noEpmd {
			l, err = net.Listen("tcp", net.JoinHostPort("", listenPort))
			if err != nil || noEpmd {
				// Cannot bind, eclus instance already running, connect to it
				eclusCli()
			}
		} else {
			if !noEpmd {
				epm := make(chan regReq, 10)
				go epmReg(epm)
				go func() {
					for {
						conn, err := l.Accept()
						log.Printf("Accept new")
						if err != nil {
							log.Printf(err.Error())
						} else {
							go mLoop(conn, epm)
						}
					}
				}()
			}
			if nodeEnabled() {
				go runNode()
			}
			<-stopCh
		}
	} else {
		eclusCli()
	}
}

type nodeRec struct {
	*epmd.NodeInfo
	Time   time.Time
	Active bool
	conn   net.Conn
}

func epmReg(in <-chan regReq) {
	var nReg = make(map[string]*nodeRec)
	for {
		select {
		case req := <-in:
			buf := req.buf
			if len(buf) == 0 {
				rs := len(nReg)
				log.Printf("REG %d records", rs)
				now := time.Now()

				for node, rec := range nReg {
					if rec.conn == req.conn {
						log.Printf("Connection for %s dropped", node)
						nReg[node].Active = false
						nReg[node].Time = now
						nReg[node].conn = nil
					} else if rs > regLimit && !rec.Active && now.Sub(rec.Time).Minutes() > float64(unregTTL) {
						log.Printf("REG prune %s:%+v", node, rec)
						delete(nReg, node)
					}
				}
				continue
			}
			replyTo := req.replyTo
			log.Printf("IN: %v", buf)
			switch epmd.MessageId(buf[0]) {
			case epmd.ALIVE2_REQ:
				nConn := req.conn
				nInfo := epmd.Read_ALIVE2_REQ(buf)
				log.Printf("NodeInfo: %+v", nInfo)
				var reply []byte

				if rec, ok := nReg[nInfo.Name]; ok {
					log.Printf("Node %s found", nInfo.Name)
					if rec.Active {
						log.Printf("Node %s is running", nInfo.Name)
						reply = epmd.Compose_ALIVE2_RESP(false)
					} else {
						log.Printf("Node %s is not running", nInfo.Name)
						rec.conn = nConn

						nInfo.Creation = (rec.Creation % 3) + 1
						rec.NodeInfo = nInfo
						rec.Active = true
						reply = epmd.Compose_ALIVE2_RESP(true, rec.NodeInfo)
					}
				} else {
					log.Printf("New node %s", nInfo.Name)
					nInfo.Creation = 1
					rec := &nodeRec{
						NodeInfo: nInfo,
						conn:     nConn,
						Time:     time.Now(),
						Active:   true,
					}
					nReg[nInfo.Name] = rec
					reply = epmd.Compose_ALIVE2_RESP(true, rec.NodeInfo)
				}
				replyTo <- regAns{reply: reply, isClose: false}
			case epmd.PORT_PLEASE2_REQ:
				nName := epmd.Read_PORT_PLEASE2_REQ(buf)
				var reply []byte
				if rec, ok := nReg[nName]; ok && rec.Active {
					reply = epmd.Compose_PORT2_RESP(rec.NodeInfo)
				} else {
					reply = epmd.Compose_PORT2_RESP(nil)
				}
				replyTo <- regAns{reply: reply, isClose: true}
			case epmd.STOP_REQ:
				nName := epmd.Read_STOP_REQ(buf)
				var reply []byte
				if rec, ok := nReg[nName]; ok && rec.Active {
					// TODO: stop node
					reply = epmd.Compose_STOP_RESP(true)
				} else {
					reply = epmd.Compose_STOP_RESP(false)
				}
				replyTo <- regAns{reply: reply, isClose: true}
			case epmd.NAMES_REQ, epmd.DUMP_REQ:
				lp, err := strconv.Atoi(listenPort)
				if err != nil {
					log.Printf("Cannot convert %s to integer", listenPort)
					replyTo <- regAns{reply: nil, isClose: true}
				} else {
					replyB := new(bytes.Buffer)
					epmd.Compose_START_NAMES_RESP(replyB, lp)
					for _, rec := range nReg {
						if epmd.MessageId(buf[0]) == epmd.NAMES_REQ {
							if rec.Active {
								epmd.Append_NAMES_RESP(replyB, rec.NodeInfo)
							} else {
								if rec.Active {
									epmd.Append_DUMP_RESP_ACTIVE(replyB, rec.NodeInfo)
								} else {
									epmd.Append_DUMP_RESP_UNUSED(replyB, rec.NodeInfo)
								}
							}
						}
					}
					replyTo <- regAns{reply: replyB.Bytes(), isClose: true}
				}
			case epmd.KILL_REQ:
				reply := epmd.Compose_KILL_RESP()
				replyTo <- regAns{reply: reply, isClose: true}
			default:
				switch cliMessageId(buf[0]) {
				case REQ_NAMES:
					reply := ansNames(nReg)
					replyTo <- regAns{reply: reply, isClose: true}
				default:
					replyTo <- regAns{reply: nil, isClose: true}
				}
			}
		}
	}
}

func mLoop(c net.Conn, epm chan regReq) {
	buf := make([]byte, 1024)
	for {
		n, err := c.Read(buf)
		if err != nil {
			c.Close()
			log.Printf("Stop loop: %v", err)
			epm <- regReq{buf: []byte{}, conn: c}
			return
		}
		length := binary.BigEndian.Uint16(buf[0:2])
		if length != uint16(n-2) {
			log.Printf("Incomplete packet from erlang node to epmd: %d from %d", n, length)
			break
		}
		log.Printf("Read %d, %d: %v", n, length, buf[2:n])
		if isClose := handleMsg(c, buf[2:n], epm); isClose {
			break
		}
	}
	c.Close()
}

func handleMsg(c net.Conn, buf []byte, epm chan regReq) bool {
	myChan := make(chan regAns)
	epm <- regReq{buf: buf, replyTo: myChan, conn: c}
	select {
	case ans := <-myChan:
		log.Printf("Got reply: %+v", ans)
		if ans.reply != nil {
			c.Write(ans.reply)
		}
		return ans.isClose
	}
	return true
}
<|endoftext|>


<commit_before>package main

import ("log"
	"fmt"
	"flag"
	"Bowery/prompt"
	"boltdb/bolt"
)

// Parse the given block count file, populate and return a new database
func parse2db(countFile string, newDB string) error {
}

// open an interactive command line to query the provided db
func dbInteract(dbName string) error {
}

func main() {
	// define flags
	idb := flag.String("i", "", "database for interactive queries")
	odb := flag.String("wdb", "", "new database which will be populated with count file data")
	cFile := flag.String("cf", "", "count file to parse and store in a new database")

	// parse command line flags
	flag.Parse()

	// check flags
	if (*idb != "" && ((*odb != "") || (*cFile != "")))  {
		log.Fatal("-wdb and -cf may not be specified in interactive mode")
	}
	if (*cFile != "" && *odb == "") || (*cFile == "" && *odb != ""){
		log.Fatal("-wdb and -cf must be specified together")
	}

	// launch interactive query session
	if *idb != "" {
		err := dbInteract(*idb)
		if err != nil {
			log.Fatalf("Error from interactive query session for database %s: %v\n", *idb, err)
		}
	} else {
		// create new database from count file
		if *cFile != nil {
			err := parse2db(*cFile, *odb)
			if err != nil {
				log.Fatalf("Error creating database file %s from count file %s: %v\n", *odb, *cFile, err)
			}
		} else {
			fmt.Println("Usage: cnt2db -i <database> or cnt2db -cf <countFile> -wdb <new database>")
		}
}
<commit_msg>still work in progress<commit_after>package main

import ("log"
	"fmt"
	"flag"
	"bufio"
	"regexp"
	"Bowery/prompt"
	"boltdb/bolt"
)

// Parse the given block count file, populate and return a new database
func parse2db(countFile string, newDB string) error {

	// Create and open the database
	db, err := bolt.Open(newDB, 0600, nil)
	if err != nil {
		return fmt.Errorf("parse2db error opening new database file %s: %v", newDB, err)
	}

	// open the count file for read
	cfd, ferr := os.Open(countFile)
	if ferr != nil {
		return fmt.Errorf("Error opening count file %s: %v", countFile, err)
	}
	defer cfd.Close()

	// create the line parsing regexps
	blre := regexp.MustCompile(`^\s*block\s*:\s*(\w+)\s*$`) // blocknames - block: <blockname>
	clre := regexp.MustCompile(`^\s*(\w+)\s*:\s*(\d+)\s*$`) // device counts - <deviceName>: <count>

	terr := db.Update(func(tx *bolt.Tx) error {

		// parse the file by line
		scanner := bufio.NewScanner(os.Stdin)
		block := ""
		var bucket = bolt.Bucket{}
		for scanner.Scan() {
			line := scanner.Text()
			bsm := blre.FindStringSubmatch(line) // look for block starts
			if bsm != nil { // start of a new block
				block := bsm[1] // 0 is the whole match
				berr, bucket := tx.CreateBucket([]byte(block))
			}
		}
		if serr := scanner.Err(); serr != nil {
			return fmt.Errorf("Error reading file %s: %v", countFile, serr)
		}

		return nil
	})
	if terr != nil {
		return fmt.Errorf("Error creating read/write db transaction for db file %s and count file %s: %v",
			newDB, countFile, err)
	}

	Defer db.Close()
}

// open an interactive command line to query the provided db
func dbInteract(dbName string) error {
}

func main() {
	// define flags
	idb := flag.String("i", "", "database for interactive queries")
	odb := flag.String("wdb", "", "new database which will be populated with count file data")
	cFile := flag.String("cf", "", "count file to parse and store in a new database")

	// parse command line flags
	flag.Parse()

	// check flags
	if (*idb != "" && ((*odb != "") || (*cFile != "")))  {
		log.Fatal("-wdb and -cf may not be specified in interactive mode")
	}
	if (*cFile != "" && *odb == "") || (*cFile == "" && *odb != ""){
		log.Fatal("-wdb and -cf must be specified together")
	}

	// launch interactive query session
	if *idb != "" {
		err := dbInteract(*idb)
		if err != nil {
			log.Fatalf("Error from interactive query session for database %s: %v\n", *idb, err)
		}
	} else {
		// create new database from count file
		if *cFile != nil {
			err := parse2db(*cFile, *odb)
			if err != nil {
				log.Fatalf("Error creating database file %s from count file %s: %v\n", *odb, *cFile, err)
			}
		} else {
			fmt.Println("Usage: cnt2db -i <database> or cnt2db -cf <countFile> -wdb <new database>")
		}
}
<|endoftext|>