<commit_before>/*
Copyright 2017 Google Inc.

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 connpool

import (
	"fmt"
	"strings"
	"sync"
	"time"

	log "github.com/golang/glog"
	"golang.org/x/net/context"

	"github.com/youtube/vitess/go/mysql"
	"github.com/youtube/vitess/go/sqltypes"
	"github.com/youtube/vitess/go/sync2"
	"github.com/youtube/vitess/go/trace"
	"github.com/youtube/vitess/go/vt/dbconnpool"
	"github.com/youtube/vitess/go/vt/vttablet/tabletserver/tabletenv"

	querypb "github.com/youtube/vitess/go/vt/proto/query"
)

// BinlogFormat is used for for specifying the binlog format.
type BinlogFormat int

// The following constants specify the possible binlog format values.
const (
	BinlogFormatStatement BinlogFormat = iota
	BinlogFormatRow
	BinlogFormatMixed
)

// DBConn is a db connection for tabletserver.
// It performs automatic reconnects as needed.
// Its Execute function has a timeout that can kill
// its own queries and the underlying connection.
// It will also trigger a CheckMySQL whenever applicable.
type DBConn struct {
	conn    *dbconnpool.DBConnection
	info    *mysql.ConnParams
	dbaPool *dbconnpool.ConnectionPool
	pool    *Pool
	current sync2.AtomicString
}

// NewDBConn creates a new DBConn. It triggers a CheckMySQL if creation fails.
func NewDBConn(
	cp *Pool,
	appParams *mysql.ConnParams) (*DBConn, error) {
	c, err := dbconnpool.NewDBConnection(appParams, tabletenv.MySQLStats)
	if err != nil {
		cp.checker.CheckMySQL()
		return nil, err
	}
	return &DBConn{
		conn: c,
		info: appParams,
		pool: cp,
	}, nil
}

// NewDBConnNoPool creates a new DBConn without a pool.
func NewDBConnNoPool(params *mysql.ConnParams, dbaPool *dbconnpool.ConnectionPool) (*DBConn, error) {
	c, err := dbconnpool.NewDBConnection(params, tabletenv.MySQLStats)
	if err != nil {
		return nil, err
	}
	return &DBConn{
		conn:    c,
		info:    params,
		dbaPool: dbaPool,
		pool:    nil,
	}, nil
}

// Exec executes the specified query. If there is a connection error, it will reconnect
// and retry. A failed reconnect will trigger a CheckMySQL.
func (dbc *DBConn) Exec(ctx context.Context, query string, maxrows int, wantfields bool) (*sqltypes.Result, error) {
	span := trace.NewSpanFromContext(ctx)
	span.StartClient("DBConn.Exec")
	defer span.Finish()

	for attempt := 1; attempt <= 2; attempt++ {
		r, err := dbc.execOnce(ctx, query, maxrows, wantfields)
		switch {
		case err == nil:
			// Success.
			return r, nil
		case !mysql.IsConnErr(err):
			// Not a connection error. Don't retry.
			return nil, err
		case attempt == 2:
			// Reached the retry limit.
			return nil, err
		}

		// Connection error. Retry if context has not expired.
		select {
		case <-ctx.Done():
			return nil, err
		default:
		}

		if reconnectErr := dbc.reconnect(); reconnectErr != nil {
			dbc.pool.checker.CheckMySQL()
			// Return the error of the reconnect and not the original connection error.
			return nil, reconnectErr
		}

		// Reconnect succeeded. Retry query at second attempt.
	}
	panic("unreachable")
}

func (dbc *DBConn) execOnce(ctx context.Context, query string, maxrows int, wantfields bool) (*sqltypes.Result, error) {
	dbc.current.Set(query)
	defer dbc.current.Set("")

	done, wg := dbc.setDeadline(ctx)
	if done != nil {
		defer func() {
			close(done)
			wg.Wait()
		}()
	}
	// Uncomment this line for manual testing.
	// defer time.Sleep(20 * time.Second)
	return dbc.conn.ExecuteFetch(query, maxrows, wantfields)
}

// ExecOnce executes the specified query, but does not retry on connection errors.
func (dbc *DBConn) ExecOnce(ctx context.Context, query string, maxrows int, wantfields bool) (*sqltypes.Result, error) {
	return dbc.execOnce(ctx, query, maxrows, wantfields)
}

// Stream executes the query and streams the results.
func (dbc *DBConn) Stream(ctx context.Context, query string, callback func(*sqltypes.Result) error, streamBufferSize int, includedFields querypb.ExecuteOptions_IncludedFields) error {
	span := trace.NewSpanFromContext(ctx)
	span.StartClient("DBConn.Stream")
	defer span.Finish()

	resultSent := false
	for attempt := 1; attempt <= 2; attempt++ {
		err := dbc.streamOnce(
			ctx,
			query,
			func(r *sqltypes.Result) error {
				if !resultSent {
					resultSent = true
					r = r.StripMetadata(includedFields)
				}
				return callback(r)
			},
			streamBufferSize,
		)
		switch {
		case err == nil:
			// Success.
			return nil
		case !mysql.IsConnErr(err):
			// Not a connection error. Don't retry.
			return err
		case attempt == 2:
			// Reached the retry limit.
			return err
		case resultSent:
			// Don't retry if streaming has started.
			return err
		}

		// Connection error. Retry if context has not expired.
		select {
		case <-ctx.Done():
			return err
		default:
		}
		if reconnectErr := dbc.reconnect(); reconnectErr != nil {
			dbc.pool.checker.CheckMySQL()
			// Return the error of the reconnect and not the original connection error.
			return reconnectErr
		}
	}
	panic("unreachable")
}

func (dbc *DBConn) streamOnce(ctx context.Context, query string, callback func(*sqltypes.Result) error, streamBufferSize int) error {
	dbc.current.Set(query)
	defer dbc.current.Set("")

	done, wg := dbc.setDeadline(ctx)
	if done != nil {
		defer func() {
			close(done)
			wg.Wait()
		}()
	}
	return dbc.conn.ExecuteStreamFetch(query, callback, streamBufferSize)
}

var (
	getModeSQL    = "select @@global.sql_mode"
	getAutocommit = "select @@autocommit"
	showBinlog    = "show variables like 'binlog_format'"
)

// VerifyMode is a helper method to verify mysql is running with
// sql_mode = STRICT_TRANS_TABLES and autocommit=ON. It also returns
// the current binlog format.
func (dbc *DBConn) VerifyMode(strictTransTables bool) (BinlogFormat, error) {
	if strictTransTables {
		qr, err := dbc.conn.ExecuteFetch(getModeSQL, 2, false)
		if err != nil {
			return 0, fmt.Errorf("could not verify mode: %v", err)
		}
		if len(qr.Rows) != 1 {
			return 0, fmt.Errorf("incorrect rowcount received for %s: %d", getModeSQL, len(qr.Rows))
		}
		if !strings.Contains(qr.Rows[0][0].ToString(), "STRICT_TRANS_TABLES") {
			return 0, fmt.Errorf("require sql_mode to be STRICT_TRANS_TABLES: got '%s'", qr.Rows[0][0].ToString())
		}
	}
	qr, err := dbc.conn.ExecuteFetch(getAutocommit, 2, false)
	if err != nil {
		return 0, fmt.Errorf("could not verify mode: %v", err)
	}
	if len(qr.Rows) != 1 {
		return 0, fmt.Errorf("incorrect rowcount received for %s: %d", getAutocommit, len(qr.Rows))
	}
	if !strings.Contains(qr.Rows[0][0].ToString(), "1") {
		return 0, fmt.Errorf("require autocommit to be 1: got %s", qr.Rows[0][0].ToString())
	}
	qr, err = dbc.conn.ExecuteFetch(showBinlog, 10, false)
	if err != nil {
		return 0, fmt.Errorf("could not fetch binlog format: %v", err)
	}
	if len(qr.Rows) != 1 {
		return 0, fmt.Errorf("incorrect rowcount received for %s: %d", showBinlog, len(qr.Rows))
	}
	if len(qr.Rows[0]) != 2 {
		return 0, fmt.Errorf("incorrect column count received for %s: %d", showBinlog, len(qr.Rows[0]))
	}
	switch qr.Rows[0][1].ToString() {
	case "STATEMENT":
		return BinlogFormatStatement, nil
	case "ROW":
		return BinlogFormatRow, nil
	case "MIXED":
		return BinlogFormatMixed, nil
	}
	return 0, fmt.Errorf("unexpected binlog format for %s: %s", showBinlog, qr.Rows[0][1].ToString())
}

// Close closes the DBConn.
func (dbc *DBConn) Close() {
	dbc.conn.Close()
}

// IsClosed returns true if DBConn is closed.
func (dbc *DBConn) IsClosed() bool {
	return dbc.conn.IsClosed()
}

// Recycle returns the DBConn to the pool.
func (dbc *DBConn) Recycle() {
	switch {
	case dbc.pool == nil:
		dbc.Close()
	case dbc.conn.IsClosed():
		dbc.pool.Put(nil)
	default:
		dbc.pool.Put(dbc)
	}
}

// Kill kills the currently executing query both on MySQL side
// and on the connection side. If no query is executing, it's a no-op.
// Kill will also not kill a query more than once.
func (dbc *DBConn) Kill(reason string, elapsed time.Duration) error {
	tabletenv.KillStats.Add("Queries", 1)
	log.Infof("Due to %s, elapsed time: %v, killing query %s", reason, elapsed, dbc.Current())
	// Hack: Most of the times DBConn is created with a pool. There is a snowflake case
	// (NewDBConnNoPool) used for AppDebug user where there is no pool set.
	// In those cases dbaPool will be available in the context directly.
	var dbaPool *dbconnpool.ConnectionPool
	if dbc.pool == nil {
		dbaPool = dbc.dbaPool
	} else {
		dbaPool = dbc.pool.dbaPool
	}
	killConn, err := dbaPool.Get(context.TODO())
	if err != nil {
		log.Warningf("Failed to get conn from dba pool: %v", err)
		return err
	}
	defer killConn.Recycle()
	sql := fmt.Sprintf("kill %d", dbc.conn.ID())
	_, err = killConn.ExecuteFetch(sql, 10000, false)
	if err != nil {
		log.Errorf("Could not kill query %s: %v", dbc.Current(), err)
		return err
	}
	return nil
}

// Current returns the currently executing query.
func (dbc *DBConn) Current() string {
	return dbc.current.Get()
}

// ID returns the connection id.
func (dbc *DBConn) ID() int64 {
	return dbc.conn.ID()
}

func (dbc *DBConn) reconnect() error {
	dbc.conn.Close()
	newConn, err := dbconnpool.NewDBConnection(dbc.info, tabletenv.MySQLStats)
	if err != nil {
		return err
	}
	dbc.conn = newConn
	return nil
}

// setDeadline starts a goroutine that will kill the currently executing query
// if the deadline is exceeded. It returns a channel and a waitgroup. After the
// query is done executing, the caller is required to close the done channel
// and wait for the waitgroup to make sure that the necessary cleanup is done.
func (dbc *DBConn) setDeadline(ctx context.Context) (chan bool, *sync.WaitGroup) {
	if ctx.Done() == nil {
		return nil, nil
	}
	done := make(chan bool)
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		startTime := time.Now()
		select {
		case <-ctx.Done():
			dbc.Kill(ctx.Err().Error(), time.Since(startTime))
		case <-done:
			return
		}
		elapsed := time.Now().Sub(startTime)

		// Give 2x the elapsed time and some buffer as grace period
		// for the query to get killed.
		tmr2 := time.NewTimer(2*elapsed + 5*time.Second)
		defer tmr2.Stop()
		select {
		case <-tmr2.C:
			tabletenv.InternalErrors.Add("HungQuery", 1)
			log.Warningf("Query may be hung: %s", dbc.Current())
		case <-done:
			return
		}
		<-done
		log.Warningf("Hung query returned")
	}()
	return done, &wg
}
<commit_msg>Changes per code review<commit_after>/*
Copyright 2017 Google Inc.

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 connpool

import (
	"fmt"
	"strings"
	"sync"
	"time"

	log "github.com/golang/glog"
	"golang.org/x/net/context"

	"github.com/youtube/vitess/go/mysql"
	"github.com/youtube/vitess/go/sqltypes"
	"github.com/youtube/vitess/go/sync2"
	"github.com/youtube/vitess/go/trace"
	"github.com/youtube/vitess/go/vt/dbconnpool"
	"github.com/youtube/vitess/go/vt/vttablet/tabletserver/tabletenv"

	querypb "github.com/youtube/vitess/go/vt/proto/query"
)

// BinlogFormat is used for for specifying the binlog format.
type BinlogFormat int

// The following constants specify the possible binlog format values.
const (
	BinlogFormatStatement BinlogFormat = iota
	BinlogFormatRow
	BinlogFormatMixed
)

// DBConn is a db connection for tabletserver.
// It performs automatic reconnects as needed.
// Its Execute function has a timeout that can kill
// its own queries and the underlying connection.
// It will also trigger a CheckMySQL whenever applicable.
type DBConn struct {
	conn    *dbconnpool.DBConnection
	info    *mysql.ConnParams
	dbaPool *dbconnpool.ConnectionPool
	pool    *Pool
	current sync2.AtomicString
}

// NewDBConn creates a new DBConn. It triggers a CheckMySQL if creation fails.
func NewDBConn(
	cp *Pool,
	appParams *mysql.ConnParams) (*DBConn, error) {
	c, err := dbconnpool.NewDBConnection(appParams, tabletenv.MySQLStats)
	if err != nil {
		cp.checker.CheckMySQL()
		return nil, err
	}
	return &DBConn{
		conn:    c,
		info:    appParams,
		pool:    cp,
		dbaPool: cp.dbaPool,
	}, nil
}

// NewDBConnNoPool creates a new DBConn without a pool.
func NewDBConnNoPool(params *mysql.ConnParams, dbaPool *dbconnpool.ConnectionPool) (*DBConn, error) {
	c, err := dbconnpool.NewDBConnection(params, tabletenv.MySQLStats)
	if err != nil {
		return nil, err
	}
	return &DBConn{
		conn:    c,
		info:    params,
		dbaPool: dbaPool,
		pool:    nil,
	}, nil
}

// Exec executes the specified query. If there is a connection error, it will reconnect
// and retry. A failed reconnect will trigger a CheckMySQL.
func (dbc *DBConn) Exec(ctx context.Context, query string, maxrows int, wantfields bool) (*sqltypes.Result, error) {
	span := trace.NewSpanFromContext(ctx)
	span.StartClient("DBConn.Exec")
	defer span.Finish()

	for attempt := 1; attempt <= 2; attempt++ {
		r, err := dbc.execOnce(ctx, query, maxrows, wantfields)
		switch {
		case err == nil:
			// Success.
			return r, nil
		case !mysql.IsConnErr(err):
			// Not a connection error. Don't retry.
			return nil, err
		case attempt == 2:
			// Reached the retry limit.
			return nil, err
		}

		// Connection error. Retry if context has not expired.
		select {
		case <-ctx.Done():
			return nil, err
		default:
		}

		if reconnectErr := dbc.reconnect(); reconnectErr != nil {
			dbc.pool.checker.CheckMySQL()
			// Return the error of the reconnect and not the original connection error.
			return nil, reconnectErr
		}

		// Reconnect succeeded. Retry query at second attempt.
	}
	panic("unreachable")
}

func (dbc *DBConn) execOnce(ctx context.Context, query string, maxrows int, wantfields bool) (*sqltypes.Result, error) {
	dbc.current.Set(query)
	defer dbc.current.Set("")

	done, wg := dbc.setDeadline(ctx)
	if done != nil {
		defer func() {
			close(done)
			wg.Wait()
		}()
	}
	// Uncomment this line for manual testing.
	// defer time.Sleep(20 * time.Second)
	return dbc.conn.ExecuteFetch(query, maxrows, wantfields)
}

// ExecOnce executes the specified query, but does not retry on connection errors.
func (dbc *DBConn) ExecOnce(ctx context.Context, query string, maxrows int, wantfields bool) (*sqltypes.Result, error) {
	return dbc.execOnce(ctx, query, maxrows, wantfields)
}

// Stream executes the query and streams the results.
func (dbc *DBConn) Stream(ctx context.Context, query string, callback func(*sqltypes.Result) error, streamBufferSize int, includedFields querypb.ExecuteOptions_IncludedFields) error {
	span := trace.NewSpanFromContext(ctx)
	span.StartClient("DBConn.Stream")
	defer span.Finish()

	resultSent := false
	for attempt := 1; attempt <= 2; attempt++ {
		err := dbc.streamOnce(
			ctx,
			query,
			func(r *sqltypes.Result) error {
				if !resultSent {
					resultSent = true
					r = r.StripMetadata(includedFields)
				}
				return callback(r)
			},
			streamBufferSize,
		)
		switch {
		case err == nil:
			// Success.
			return nil
		case !mysql.IsConnErr(err):
			// Not a connection error. Don't retry.
			return err
		case attempt == 2:
			// Reached the retry limit.
			return err
		case resultSent:
			// Don't retry if streaming has started.
			return err
		}

		// Connection error. Retry if context has not expired.
		select {
		case <-ctx.Done():
			return err
		default:
		}
		if reconnectErr := dbc.reconnect(); reconnectErr != nil {
			dbc.pool.checker.CheckMySQL()
			// Return the error of the reconnect and not the original connection error.
			return reconnectErr
		}
	}
	panic("unreachable")
}

func (dbc *DBConn) streamOnce(ctx context.Context, query string, callback func(*sqltypes.Result) error, streamBufferSize int) error {
	dbc.current.Set(query)
	defer dbc.current.Set("")

	done, wg := dbc.setDeadline(ctx)
	if done != nil {
		defer func() {
			close(done)
			wg.Wait()
		}()
	}
	return dbc.conn.ExecuteStreamFetch(query, callback, streamBufferSize)
}

var (
	getModeSQL    = "select @@global.sql_mode"
	getAutocommit = "select @@autocommit"
	showBinlog    = "show variables like 'binlog_format'"
)

// VerifyMode is a helper method to verify mysql is running with
// sql_mode = STRICT_TRANS_TABLES and autocommit=ON. It also returns
// the current binlog format.
func (dbc *DBConn) VerifyMode(strictTransTables bool) (BinlogFormat, error) {
	if strictTransTables {
		qr, err := dbc.conn.ExecuteFetch(getModeSQL, 2, false)
		if err != nil {
			return 0, fmt.Errorf("could not verify mode: %v", err)
		}
		if len(qr.Rows) != 1 {
			return 0, fmt.Errorf("incorrect rowcount received for %s: %d", getModeSQL, len(qr.Rows))
		}
		if !strings.Contains(qr.Rows[0][0].ToString(), "STRICT_TRANS_TABLES") {
			return 0, fmt.Errorf("require sql_mode to be STRICT_TRANS_TABLES: got '%s'", qr.Rows[0][0].ToString())
		}
	}
	qr, err := dbc.conn.ExecuteFetch(getAutocommit, 2, false)
	if err != nil {
		return 0, fmt.Errorf("could not verify mode: %v", err)
	}
	if len(qr.Rows) != 1 {
		return 0, fmt.Errorf("incorrect rowcount received for %s: %d", getAutocommit, len(qr.Rows))
	}
	if !strings.Contains(qr.Rows[0][0].ToString(), "1") {
		return 0, fmt.Errorf("require autocommit to be 1: got %s", qr.Rows[0][0].ToString())
	}
	qr, err = dbc.conn.ExecuteFetch(showBinlog, 10, false)
	if err != nil {
		return 0, fmt.Errorf("could not fetch binlog format: %v", err)
	}
	if len(qr.Rows) != 1 {
		return 0, fmt.Errorf("incorrect rowcount received for %s: %d", showBinlog, len(qr.Rows))
	}
	if len(qr.Rows[0]) != 2 {
		return 0, fmt.Errorf("incorrect column count received for %s: %d", showBinlog, len(qr.Rows[0]))
	}
	switch qr.Rows[0][1].ToString() {
	case "STATEMENT":
		return BinlogFormatStatement, nil
	case "ROW":
		return BinlogFormatRow, nil
	case "MIXED":
		return BinlogFormatMixed, nil
	}
	return 0, fmt.Errorf("unexpected binlog format for %s: %s", showBinlog, qr.Rows[0][1].ToString())
}

// Close closes the DBConn.
func (dbc *DBConn) Close() {
	dbc.conn.Close()
}

// IsClosed returns true if DBConn is closed.
func (dbc *DBConn) IsClosed() bool {
	return dbc.conn.IsClosed()
}

// Recycle returns the DBConn to the pool.
func (dbc *DBConn) Recycle() {
	switch {
	case dbc.pool == nil:
		dbc.Close()
	case dbc.conn.IsClosed():
		dbc.pool.Put(nil)
	default:
		dbc.pool.Put(dbc)
	}
}

// Kill kills the currently executing query both on MySQL side
// and on the connection side. If no query is executing, it's a no-op.
// Kill will also not kill a query more than once.
func (dbc *DBConn) Kill(reason string, elapsed time.Duration) error {
	tabletenv.KillStats.Add("Queries", 1)
	log.Infof("Due to %s, elapsed time: %v, killing query %s", reason, elapsed, dbc.Current())
	killConn, err := dbc.dbaPool.Get(context.TODO())
	if err != nil {
		log.Warningf("Failed to get conn from dba pool: %v", err)
		return err
	}
	defer killConn.Recycle()
	sql := fmt.Sprintf("kill %d", dbc.conn.ID())
	_, err = killConn.ExecuteFetch(sql, 10000, false)
	if err != nil {
		log.Errorf("Could not kill query %s: %v", dbc.Current(), err)
		return err
	}
	return nil
}

// Current returns the currently executing query.
func (dbc *DBConn) Current() string {
	return dbc.current.Get()
}

// ID returns the connection id.
func (dbc *DBConn) ID() int64 {
	return dbc.conn.ID()
}

func (dbc *DBConn) reconnect() error {
	dbc.conn.Close()
	newConn, err := dbconnpool.NewDBConnection(dbc.info, tabletenv.MySQLStats)
	if err != nil {
		return err
	}
	dbc.conn = newConn
	return nil
}

// setDeadline starts a goroutine that will kill the currently executing query
// if the deadline is exceeded. It returns a channel and a waitgroup. After the
// query is done executing, the caller is required to close the done channel
// and wait for the waitgroup to make sure that the necessary cleanup is done.
func (dbc *DBConn) setDeadline(ctx context.Context) (chan bool, *sync.WaitGroup) {
	if ctx.Done() == nil {
		return nil, nil
	}
	done := make(chan bool)
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		startTime := time.Now()
		select {
		case <-ctx.Done():
			dbc.Kill(ctx.Err().Error(), time.Since(startTime))
		case <-done:
			return
		}
		elapsed := time.Now().Sub(startTime)

		// Give 2x the elapsed time and some buffer as grace period
		// for the query to get killed.
		tmr2 := time.NewTimer(2*elapsed + 5*time.Second)
		defer tmr2.Stop()
		select {
		case <-tmr2.C:
			tabletenv.InternalErrors.Add("HungQuery", 1)
			log.Warningf("Query may be hung: %s", dbc.Current())
		case <-done:
			return
		}
		<-done
		log.Warningf("Hung query returned")
	}()
	return done, &wg
}
<|endoftext|>


<commit_before>package goyaml

// #cgo LDFLAGS: -lm -lpthread
// #cgo windows LDFLAGS: -L. -lyaml
// #cgo windows CFLAGS: -DYAML_DECLARE_STATIC=1
// #cgo CFLAGS: -I. -DHAVE_CONFIG_H=1
//
// #include "helpers.h"
import "C"

import (
	"reflect"
	"strconv"
	"unsafe"
)

const (
	documentNode = 1 << iota
	mappingNode
	sequenceNode
	scalarNode
	aliasNode
)

type node struct {
	kind         int
	line, column int
	tag          string
	value        string
	implicit     bool
	children     []*node
	anchors      map[string]*node
}

func stry(s *C.yaml_char_t) string {
	return C.GoString((*C.char)(unsafe.Pointer(s)))
}

// ----------------------------------------------------------------------------
// Parser, produces a node tree out of a libyaml event stream.

type parser struct {
	parser C.yaml_parser_t
	event  C.yaml_event_t
	doc    *node
}

func newParser(b []byte) *parser {
	p := parser{}
	if C.yaml_parser_initialize(&p.parser) == 0 {
		panic("Failed to initialize YAML emitter")
	}

	if len(b) == 0 {
		b = []byte{'\n'}
	}

	// How unsafe is this really?  Will this break if the GC becomes compacting?
	// Probably not, otherwise that would likely break &parse below as well.
	input := (*C.uchar)(unsafe.Pointer(&b[0]))
	C.yaml_parser_set_input_string(&p.parser, input, (C.size_t)(len(b)))

	p.skip()
	if p.event._type != C.YAML_STREAM_START_EVENT {
		panic("Expected stream start event, got " +
			strconv.Itoa(int(p.event._type)))
	}
	p.skip()
	return &p
}

func (p *parser) destroy() {
	if p.event._type != C.YAML_NO_EVENT {
		C.yaml_event_delete(&p.event)
	}
	C.yaml_parser_delete(&p.parser)
}

func (p *parser) skip() {
	if p.event._type != C.YAML_NO_EVENT {
		if p.event._type == C.YAML_STREAM_END_EVENT {
			panic("Attempted to go past the end of stream. Corrupted value?")
		}
		C.yaml_event_delete(&p.event)
	}
	if C.yaml_parser_parse(&p.parser, &p.event) == 0 {
		p.fail()
	}
}

func (p *parser) fail() {
	var where string
	var line int
	if p.parser.problem_mark.line != 0 {
		line = int(C.int(p.parser.problem_mark.line))
	} else if p.parser.context_mark.line != 0 {
		line = int(C.int(p.parser.context_mark.line))
	}
	if line != 0 {
		where = "line " + strconv.Itoa(line) + ": "
	}
	var msg string
	if p.parser.problem != nil {
		msg = C.GoString(p.parser.problem)
	} else {
		msg = "Unknown problem parsing YAML content"
	}
	panic(where + msg)
}

func (p *parser) anchor(n *node, anchor *C.yaml_char_t) {
	if anchor != nil {
		p.doc.anchors[stry(anchor)] = n
	}
}

func (p *parser) parse() *node {
	switch p.event._type {
	case C.YAML_SCALAR_EVENT:
		return p.scalar()
	case C.YAML_ALIAS_EVENT:
		return p.alias()
	case C.YAML_MAPPING_START_EVENT:
		return p.mapping()
	case C.YAML_SEQUENCE_START_EVENT:
		return p.sequence()
	case C.YAML_DOCUMENT_START_EVENT:
		return p.document()
	case C.YAML_STREAM_END_EVENT:
		// Happens when attempting to decode an empty buffer.
		return nil
	default:
		panic("Attempted to parse unknown event: " +
			strconv.Itoa(int(p.event._type)))
	}
	panic("Unreachable")
}

func (p *parser) node(kind int) *node {
	return &node{kind: kind,
		line:   int(C.int(p.event.start_mark.line)),
		column: int(C.int(p.event.start_mark.column))}
}

func (p *parser) document() *node {
	n := p.node(documentNode)
	n.anchors = make(map[string]*node)
	p.doc = n
	p.skip()
	n.children = append(n.children, p.parse())
	if p.event._type != C.YAML_DOCUMENT_END_EVENT {
		panic("Expected end of document event but got " +
			strconv.Itoa(int(p.event._type)))
	}
	p.skip()
	return n
}

func (p *parser) alias() *node {
	alias := C.event_alias(&p.event)
	n := p.node(aliasNode)
	n.value = stry(alias.anchor)
	p.skip()
	return n
}

func (p *parser) scalar() *node {
	scalar := C.event_scalar(&p.event)
	n := p.node(scalarNode)
	n.value = stry(scalar.value)
	n.tag = stry(scalar.tag)
	n.implicit = (scalar.plain_implicit != 0)
	p.anchor(n, scalar.anchor)
	p.skip()
	return n
}

func (p *parser) sequence() *node {
	n := p.node(sequenceNode)
	p.anchor(n, C.event_sequence_start(&p.event).anchor)
	p.skip()
	for p.event._type != C.YAML_SEQUENCE_END_EVENT {
		n.children = append(n.children, p.parse())
	}
	p.skip()
	return n
}

func (p *parser) mapping() *node {
	n := p.node(mappingNode)
	p.anchor(n, C.event_mapping_start(&p.event).anchor)
	p.skip()
	for p.event._type != C.YAML_MAPPING_END_EVENT {
		n.children = append(n.children, p.parse(), p.parse())
	}
	p.skip()
	return n
}

// ----------------------------------------------------------------------------
// Decoder, unmarshals a node into a provided value.

type decoder struct {
	doc     *node
	aliases map[string]bool
}

func newDecoder() *decoder {
	d := &decoder{}
	d.aliases = make(map[string]bool)
	return d
}

// d.setter deals with setters and pointer dereferencing and initialization.
//
// It's a slightly convoluted case to handle properly:
//
// - nil pointers should be initialized, unless being set to nil
// - we don't know at this point yet what's the value to SetYAML() with.
// - we can't separate pointer deref/init and setter checking, because
//   a setter may be found while going down a pointer chain.
//
// Thus, here is how it takes care of it:
//
// - out is provided as a pointer, so that it can be replaced.
// - when looking at a non-setter ptr, *out=ptr.Elem(), unless tag=!!null
// - when a setter is found, *out=interface{}, and a set() function is
//   returned to call SetYAML() with the value of *out once it's defined.
//
func (d *decoder) setter(tag string, out *reflect.Value, good *bool) (set func()) {
	again := true
	for again {
		again = false
		setter, _ := (*out).Interface().(Setter)
		if tag != "!!null" || setter != nil {
			if pv := (*out); pv.Kind() == reflect.Ptr {
				if pv.IsNil() {
					*out = reflect.New(pv.Type().Elem()).Elem()
					pv.Set((*out).Addr())
				} else {
					*out = pv.Elem()
				}
				setter, _ = pv.Interface().(Setter)
				again = true
			}
		}
		if setter != nil {
			var arg interface{}
			*out = reflect.ValueOf(&arg).Elem()
			return func() {
				*good = setter.SetYAML(tag, arg)
			}
		}
	}
	return nil
}

func (d *decoder) unmarshal(n *node, out reflect.Value) (good bool) {
	switch n.kind {
	case documentNode:
		good = d.document(n, out)
	case scalarNode:
		good = d.scalar(n, out)
	case aliasNode:
		good = d.alias(n, out)
	case mappingNode:
		good = d.mapping(n, out)
	case sequenceNode:
		good = d.sequence(n, out)
	default:
		panic("Internal error: unknown node kind: " + strconv.Itoa(n.kind))
	}
	return
}

func (d *decoder) document(n *node, out reflect.Value) (good bool) {
	if len(n.children) == 1 {
		d.doc = n
		d.unmarshal(n.children[0], out)
		return true
	}
	return false
}

func (d *decoder) alias(n *node, out reflect.Value) (good bool) {
	an, ok := d.doc.anchors[n.value]
	if !ok {
		panic("Unknown anchor '" + n.value + "' referenced")
	}
	if d.aliases[n.value] {
		panic("Anchor '" + n.value + "' value contains itself")
	}
	d.aliases[n.value] = true
	good = d.unmarshal(an, out)
	delete(d.aliases, n.value)
	return good
}

func (d *decoder) scalar(n *node, out reflect.Value) (good bool) {
	var tag string
	var resolved interface{}
	if n.tag == "" && !n.implicit {
		resolved = n.value
	} else {
		tag, resolved = resolve(n.tag, n.value)
		if set := d.setter(tag, &out, &good); set != nil {
			defer set()
		}
	}
	switch out.Kind() {
	case reflect.String:
		out.SetString(n.value)
		good = true
	case reflect.Interface:
		if resolved == nil {
			out.Set(reflect.Zero(out.Type()))
		} else {
			out.Set(reflect.ValueOf(resolved))
		}
		good = true
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		switch resolved := resolved.(type) {
		case int:
			if !out.OverflowInt(int64(resolved)) {
				out.SetInt(int64(resolved))
				good = true
			}
		case int64:
			if !out.OverflowInt(resolved) {
				out.SetInt(resolved)
				good = true
			}
		}
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr:
		switch resolved := resolved.(type) {
		case int:
			if resolved >= 0 {
				out.SetUint(uint64(resolved))
				good = true
			}
		case int64:
			if resolved >= 0 {
				out.SetUint(uint64(resolved))
				good = true
			}
		}
	case reflect.Bool:
		switch resolved := resolved.(type) {
		case bool:
			out.SetBool(resolved)
			good = true
		}
	case reflect.Float32, reflect.Float64:
		switch resolved := resolved.(type) {
		case float64:
			out.SetFloat(resolved)
			good = true
		}
	case reflect.Ptr:
		switch resolved.(type) {
		case nil:
			out.Set(reflect.Zero(out.Type()))
			good = true
		}
	}
	return good
}

func settableValueOf(i interface{}) reflect.Value {
	v := reflect.ValueOf(i)
	sv := reflect.New(v.Type()).Elem()
	sv.Set(v)
	return sv
}

func (d *decoder) sequence(n *node, out reflect.Value) (good bool) {
	if set := d.setter("!!seq", &out, &good); set != nil {
		defer set()
	}
	var iface reflect.Value
	if out.Kind() == reflect.Interface {
		// No type hints. Will have to use a generic sequence.
		iface = out
		out = settableValueOf(make([]interface{}, 0))
	}

	if out.Kind() != reflect.Slice {
		return false
	}
	et := out.Type().Elem()

	l := len(n.children)
	for i := 0; i < l; i++ {
		e := reflect.New(et).Elem()
		if ok := d.unmarshal(n.children[i], e); ok {
			out.Set(reflect.Append(out, e))
		}
	}
	if iface.IsValid() {
		iface.Set(out)
	}
	return true
}

func (d *decoder) mapping(n *node, out reflect.Value) (good bool) {
	if set := d.setter("!!map", &out, &good); set != nil {
		defer set()
	}
	if out.Kind() == reflect.Struct {
		return d.mappingStruct(n, out)
	}

	if out.Kind() == reflect.Interface {
		// No type hints. Will have to use a generic map.
		iface := out
		out = settableValueOf(make(map[interface{}]interface{}))
		iface.Set(out)
	}

	if out.Kind() != reflect.Map {
		return false
	}
	outt := out.Type()
	kt := outt.Key()
	et := outt.Elem()

	if out.IsNil() {
		out.Set(reflect.MakeMap(outt))
	}
	l := len(n.children)
	for i := 0; i < l; i += 2 {
		k := reflect.New(kt).Elem()
		if d.unmarshal(n.children[i], k) {
			e := reflect.New(et).Elem()
			if d.unmarshal(n.children[i+1], e) {
				out.SetMapIndex(k, e)
			}
		}
	}
	return true
}

func (d *decoder) mappingStruct(n *node, out reflect.Value) (good bool) {
	fields, err := getStructFields(out.Type())
	if err != nil {
		panic(err)
	}
	name := settableValueOf("")
	fieldsMap := fields.Map
	l := len(n.children)
	for i := 0; i < l; i += 2 {
		if !d.unmarshal(n.children[i], name) {
			continue
		}
		if info, ok := fieldsMap[name.String()]; ok {
			d.unmarshal(n.children[i+1], out.Field(info.Num))
		}
	}
	return true
}
<commit_msg>You don't have to link -lyaml, just fix the declarations.<commit_after>package goyaml

// #cgo LDFLAGS: -lm -lpthread
// #cgo windows CFLAGS: -DYAML_DECLARE_STATIC=1
// #cgo CFLAGS: -I. -DHAVE_CONFIG_H=1
//
// #include "helpers.h"
import "C"

import (
	"reflect"
	"strconv"
	"unsafe"
)

const (
	documentNode = 1 << iota
	mappingNode
	sequenceNode
	scalarNode
	aliasNode
)

type node struct {
	kind         int
	line, column int
	tag          string
	value        string
	implicit     bool
	children     []*node
	anchors      map[string]*node
}

func stry(s *C.yaml_char_t) string {
	return C.GoString((*C.char)(unsafe.Pointer(s)))
}

// ----------------------------------------------------------------------------
// Parser, produces a node tree out of a libyaml event stream.

type parser struct {
	parser C.yaml_parser_t
	event  C.yaml_event_t
	doc    *node
}

func newParser(b []byte) *parser {
	p := parser{}
	if C.yaml_parser_initialize(&p.parser) == 0 {
		panic("Failed to initialize YAML emitter")
	}

	if len(b) == 0 {
		b = []byte{'\n'}
	}

	// How unsafe is this really?  Will this break if the GC becomes compacting?
	// Probably not, otherwise that would likely break &parse below as well.
	input := (*C.uchar)(unsafe.Pointer(&b[0]))
	C.yaml_parser_set_input_string(&p.parser, input, (C.size_t)(len(b)))

	p.skip()
	if p.event._type != C.YAML_STREAM_START_EVENT {
		panic("Expected stream start event, got " +
			strconv.Itoa(int(p.event._type)))
	}
	p.skip()
	return &p
}

func (p *parser) destroy() {
	if p.event._type != C.YAML_NO_EVENT {
		C.yaml_event_delete(&p.event)
	}
	C.yaml_parser_delete(&p.parser)
}

func (p *parser) skip() {
	if p.event._type != C.YAML_NO_EVENT {
		if p.event._type == C.YAML_STREAM_END_EVENT {
			panic("Attempted to go past the end of stream. Corrupted value?")
		}
		C.yaml_event_delete(&p.event)
	}
	if C.yaml_parser_parse(&p.parser, &p.event) == 0 {
		p.fail()
	}
}

func (p *parser) fail() {
	var where string
	var line int
	if p.parser.problem_mark.line != 0 {
		line = int(C.int(p.parser.problem_mark.line))
	} else if p.parser.context_mark.line != 0 {
		line = int(C.int(p.parser.context_mark.line))
	}
	if line != 0 {
		where = "line " + strconv.Itoa(line) + ": "
	}
	var msg string
	if p.parser.problem != nil {
		msg = C.GoString(p.parser.problem)
	} else {
		msg = "Unknown problem parsing YAML content"
	}
	panic(where + msg)
}

func (p *parser) anchor(n *node, anchor *C.yaml_char_t) {
	if anchor != nil {
		p.doc.anchors[stry(anchor)] = n
	}
}

func (p *parser) parse() *node {
	switch p.event._type {
	case C.YAML_SCALAR_EVENT:
		return p.scalar()
	case C.YAML_ALIAS_EVENT:
		return p.alias()
	case C.YAML_MAPPING_START_EVENT:
		return p.mapping()
	case C.YAML_SEQUENCE_START_EVENT:
		return p.sequence()
	case C.YAML_DOCUMENT_START_EVENT:
		return p.document()
	case C.YAML_STREAM_END_EVENT:
		// Happens when attempting to decode an empty buffer.
		return nil
	default:
		panic("Attempted to parse unknown event: " +
			strconv.Itoa(int(p.event._type)))
	}
	panic("Unreachable")
}

func (p *parser) node(kind int) *node {
	return &node{kind: kind,
		line:   int(C.int(p.event.start_mark.line)),
		column: int(C.int(p.event.start_mark.column))}
}

func (p *parser) document() *node {
	n := p.node(documentNode)
	n.anchors = make(map[string]*node)
	p.doc = n
	p.skip()
	n.children = append(n.children, p.parse())
	if p.event._type != C.YAML_DOCUMENT_END_EVENT {
		panic("Expected end of document event but got " +
			strconv.Itoa(int(p.event._type)))
	}
	p.skip()
	return n
}

func (p *parser) alias() *node {
	alias := C.event_alias(&p.event)
	n := p.node(aliasNode)
	n.value = stry(alias.anchor)
	p.skip()
	return n
}

func (p *parser) scalar() *node {
	scalar := C.event_scalar(&p.event)
	n := p.node(scalarNode)
	n.value = stry(scalar.value)
	n.tag = stry(scalar.tag)
	n.implicit = (scalar.plain_implicit != 0)
	p.anchor(n, scalar.anchor)
	p.skip()
	return n
}

func (p *parser) sequence() *node {
	n := p.node(sequenceNode)
	p.anchor(n, C.event_sequence_start(&p.event).anchor)
	p.skip()
	for p.event._type != C.YAML_SEQUENCE_END_EVENT {
		n.children = append(n.children, p.parse())
	}
	p.skip()
	return n
}

func (p *parser) mapping() *node {
	n := p.node(mappingNode)
	p.anchor(n, C.event_mapping_start(&p.event).anchor)
	p.skip()
	for p.event._type != C.YAML_MAPPING_END_EVENT {
		n.children = append(n.children, p.parse(), p.parse())
	}
	p.skip()
	return n
}

// ----------------------------------------------------------------------------
// Decoder, unmarshals a node into a provided value.

type decoder struct {
	doc     *node
	aliases map[string]bool
}

func newDecoder() *decoder {
	d := &decoder{}
	d.aliases = make(map[string]bool)
	return d
}

// d.setter deals with setters and pointer dereferencing and initialization.
//
// It's a slightly convoluted case to handle properly:
//
// - nil pointers should be initialized, unless being set to nil
// - we don't know at this point yet what's the value to SetYAML() with.
// - we can't separate pointer deref/init and setter checking, because
//   a setter may be found while going down a pointer chain.
//
// Thus, here is how it takes care of it:
//
// - out is provided as a pointer, so that it can be replaced.
// - when looking at a non-setter ptr, *out=ptr.Elem(), unless tag=!!null
// - when a setter is found, *out=interface{}, and a set() function is
//   returned to call SetYAML() with the value of *out once it's defined.
//
func (d *decoder) setter(tag string, out *reflect.Value, good *bool) (set func()) {
	again := true
	for again {
		again = false
		setter, _ := (*out).Interface().(Setter)
		if tag != "!!null" || setter != nil {
			if pv := (*out); pv.Kind() == reflect.Ptr {
				if pv.IsNil() {
					*out = reflect.New(pv.Type().Elem()).Elem()
					pv.Set((*out).Addr())
				} else {
					*out = pv.Elem()
				}
				setter, _ = pv.Interface().(Setter)
				again = true
			}
		}
		if setter != nil {
			var arg interface{}
			*out = reflect.ValueOf(&arg).Elem()
			return func() {
				*good = setter.SetYAML(tag, arg)
			}
		}
	}
	return nil
}

func (d *decoder) unmarshal(n *node, out reflect.Value) (good bool) {
	switch n.kind {
	case documentNode:
		good = d.document(n, out)
	case scalarNode:
		good = d.scalar(n, out)
	case aliasNode:
		good = d.alias(n, out)
	case mappingNode:
		good = d.mapping(n, out)
	case sequenceNode:
		good = d.sequence(n, out)
	default:
		panic("Internal error: unknown node kind: " + strconv.Itoa(n.kind))
	}
	return
}

func (d *decoder) document(n *node, out reflect.Value) (good bool) {
	if len(n.children) == 1 {
		d.doc = n
		d.unmarshal(n.children[0], out)
		return true
	}
	return false
}

func (d *decoder) alias(n *node, out reflect.Value) (good bool) {
	an, ok := d.doc.anchors[n.value]
	if !ok {
		panic("Unknown anchor '" + n.value + "' referenced")
	}
	if d.aliases[n.value] {
		panic("Anchor '" + n.value + "' value contains itself")
	}
	d.aliases[n.value] = true
	good = d.unmarshal(an, out)
	delete(d.aliases, n.value)
	return good
}

func (d *decoder) scalar(n *node, out reflect.Value) (good bool) {
	var tag string
	var resolved interface{}
	if n.tag == "" && !n.implicit {
		resolved = n.value
	} else {
		tag, resolved = resolve(n.tag, n.value)
		if set := d.setter(tag, &out, &good); set != nil {
			defer set()
		}
	}
	switch out.Kind() {
	case reflect.String:
		out.SetString(n.value)
		good = true
	case reflect.Interface:
		if resolved == nil {
			out.Set(reflect.Zero(out.Type()))
		} else {
			out.Set(reflect.ValueOf(resolved))
		}
		good = true
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		switch resolved := resolved.(type) {
		case int:
			if !out.OverflowInt(int64(resolved)) {
				out.SetInt(int64(resolved))
				good = true
			}
		case int64:
			if !out.OverflowInt(resolved) {
				out.SetInt(resolved)
				good = true
			}
		}
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, reflect.Uintptr:
		switch resolved := resolved.(type) {
		case int:
			if resolved >= 0 {
				out.SetUint(uint64(resolved))
				good = true
			}
		case int64:
			if resolved >= 0 {
				out.SetUint(uint64(resolved))
				good = true
			}
		}
	case reflect.Bool:
		switch resolved := resolved.(type) {
		case bool:
			out.SetBool(resolved)
			good = true
		}
	case reflect.Float32, reflect.Float64:
		switch resolved := resolved.(type) {
		case float64:
			out.SetFloat(resolved)
			good = true
		}
	case reflect.Ptr:
		switch resolved.(type) {
		case nil:
			out.Set(reflect.Zero(out.Type()))
			good = true
		}
	}
	return good
}

func settableValueOf(i interface{}) reflect.Value {
	v := reflect.ValueOf(i)
	sv := reflect.New(v.Type()).Elem()
	sv.Set(v)
	return sv
}

func (d *decoder) sequence(n *node, out reflect.Value) (good bool) {
	if set := d.setter("!!seq", &out, &good); set != nil {
		defer set()
	}
	var iface reflect.Value
	if out.Kind() == reflect.Interface {
		// No type hints. Will have to use a generic sequence.
		iface = out
		out = settableValueOf(make([]interface{}, 0))
	}

	if out.Kind() != reflect.Slice {
		return false
	}
	et := out.Type().Elem()

	l := len(n.children)
	for i := 0; i < l; i++ {
		e := reflect.New(et).Elem()
		if ok := d.unmarshal(n.children[i], e); ok {
			out.Set(reflect.Append(out, e))
		}
	}
	if iface.IsValid() {
		iface.Set(out)
	}
	return true
}

func (d *decoder) mapping(n *node, out reflect.Value) (good bool) {
	if set := d.setter("!!map", &out, &good); set != nil {
		defer set()
	}
	if out.Kind() == reflect.Struct {
		return d.mappingStruct(n, out)
	}

	if out.Kind() == reflect.Interface {
		// No type hints. Will have to use a generic map.
		iface := out
		out = settableValueOf(make(map[interface{}]interface{}))
		iface.Set(out)
	}

	if out.Kind() != reflect.Map {
		return false
	}
	outt := out.Type()
	kt := outt.Key()
	et := outt.Elem()

	if out.IsNil() {
		out.Set(reflect.MakeMap(outt))
	}
	l := len(n.children)
	for i := 0; i < l; i += 2 {
		k := reflect.New(kt).Elem()
		if d.unmarshal(n.children[i], k) {
			e := reflect.New(et).Elem()
			if d.unmarshal(n.children[i+1], e) {
				out.SetMapIndex(k, e)
			}
		}
	}
	return true
}

func (d *decoder) mappingStruct(n *node, out reflect.Value) (good bool) {
	fields, err := getStructFields(out.Type())
	if err != nil {
		panic(err)
	}
	name := settableValueOf("")
	fieldsMap := fields.Map
	l := len(n.children)
	for i := 0; i < l; i += 2 {
		if !d.unmarshal(n.children[i], name) {
			continue
		}
		if info, ok := fieldsMap[name.String()]; ok {
			d.unmarshal(n.children[i+1], out.Field(info.Num))
		}
	}
	return true
}
<|endoftext|>


<commit_before>package smpperf

import (
	"log"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"golang.org/x/time/rate"

	"github.com/veoo/go-smpp/smpp"
	"github.com/veoo/go-smpp/smpp/pdu"
	"github.com/veoo/go-smpp/smpp/pdu/pdufield"
	"github.com/veoo/go-smpp/smpp/pdu/pdutext"
)

type SMPPerf struct {
	NumMessages          int
	NumSessions          int
	MessageRate          int
	MessageText          string
	Wait                 int
	User                 string
	Password             string
	Host                 string
	Mode                 string
	Dst                  int
	Src                  string
	Verbose              bool
	transceivers         []*transceiverConn
	counters             *counters
	msgIDToTransceiverID *ConcurrentStringMap
}

type transceiverConn struct {
	transceiver *smpp.Transceiver
	err         error
	m           *sync.RWMutex
}

func (s *SMPPerf) setTransceiverErr(transceiverIndex int, err error) {
	s.transceivers[transceiverIndex].m.Lock()
	defer s.transceivers[transceiverIndex].m.Unlock()
	s.transceivers[transceiverIndex].err = err
	log.Println(err)
}

func (s *SMPPerf) checkTransceiverErr(transceiverIndex int) bool {
	s.transceivers[transceiverIndex].m.Lock()
	defer s.transceivers[transceiverIndex].m.Unlock()
	return s.transceivers[transceiverIndex].err != nil

}

func (s *SMPPerf) submitMessage(transceiverIndex int, req *smpp.ShortMessage) {
	sm, err := s.transceivers[transceiverIndex].transceiver.Submit(req)
	if err != nil {
		if err == smpp.ErrNotConnected {
			go s.counters.connErrorCount.Increment()
		} else {
			go s.counters.sendErrorCount.Increment()
		}
	} else {
		transceiverID := strconv.Itoa(transceiverIndex)
		s.msgIDToTransceiverID.Set(sm.RespID(), transceiverID)
	}
}

type counters struct {
	successCount     *SafeInt
	sendErrorCount   *SafeInt
	unknownRespCount *SafeInt
	connErrorCount   *SafeInt
	submittedCount   *SafeInt
	stateCounters    *ConcurrentIntMap
}

func newCounters() *counters {
	return &counters{
		successCount:     NewSafeInt(0),
		sendErrorCount:   NewSafeInt(0),
		unknownRespCount: NewSafeInt(0),
		connErrorCount:   NewSafeInt(0),
		submittedCount:   NewSafeInt(0),
		stateCounters:    NewConcurrentIntMap(),
	}
}

func closeTransceiverOnSignal(trans *smpp.Transceiver) {
	go func() {
		signalChannel := make(chan os.Signal, 1)
		signal.Notify(signalChannel, os.Interrupt, syscall.SIGTERM)
		sig := <-signalChannel
		log.Println("WARNING:", sig, "signal caught, exiting.")
		trans.Close()
		os.Exit(0)
	}()
}

func getMessageID(p pdu.Body) string {
	tlv := p.TLVFields()
	if tlv == nil {
		return ""
	}
	field := tlv[pdufield.ReceiptedMessageID]
	if field == nil {
		return ""
	}
	return strings.TrimRight(string(field.Bytes()), "\x00")
}

func (s *SMPPerf) getTransceiver() *smpp.Transceiver {
	return &smpp.Transceiver{
		Addr:        s.Host,
		User:        s.User,
		Passwd:      s.Password,
		RespTimeout: 10 * time.Second,
		EnquireLink: 1 * time.Second,
	}
}

func (s *SMPPerf) countState(counterMap *ConcurrentIntMap, state pdufield.MessageStateType) {
	if _, ok := counterMap.Get(state); !ok {
		counterMap.Create(state)
	}
	go counterMap.Increment(state)
	return
}

func (s *SMPPerf) isFinalState(state pdufield.MessageStateType) bool {
	// Expired, Delivered, Undeliverable, Rejected, unsure about Deleted
	return state == pdufield.Expired || state == pdufield.Delivered || state == pdufield.Undeliverable || state == pdufield.Rejected || state == pdufield.Deleted
}

func (s *SMPPerf) SendMessages() {
	s.counters = newCounters()
	s.transceivers = make([]*transceiverConn, s.NumSessions)
	s.msgIDToTransceiverID = NewConcurrentStringMap()

	for i := 0; i < s.NumSessions; i++ {
		transceiverID := strconv.Itoa(i)
		transceiverHandler := func(p pdu.Body) {
			switch p.Header().ID {
			case pdu.DeliverSMID:
				// TODO: check here the resp data is correct
				msgID := getMessageID(p)

				t, ok := s.msgIDToTransceiverID.Get(msgID)
				if !ok {
					log.Printf("ERROR: message %s not found in transceiver %v", msgID, transceiverID)
				} else if t != transceiverID {
					log.Printf("ERROR: message %s was received in wrong transceiver %s", msgID, transceiverID)
				}

				state := pdufield.MessageStateType(p.TLVFields()[pdufield.MessageStateOption].Bytes()[0])
				s.countState(s.counters.stateCounters, state)

				if s.isFinalState(state) {
					go s.counters.successCount.Increment()
				}

			case pdu.UnbindID:
				log.Println("ERROR: They are unbinding me :(")
			case pdu.SubmitSMRespID:
				// Fix something florix?
			default:
				go log.Println(p.Header().ID.String(), p.Header().Status.Error())
				go s.counters.unknownRespCount.Increment()
			}
		}

		transceiver := s.getTransceiver()
		transceiver.Handler = transceiverHandler

		conn := transceiver.Bind() // make persistent connection.
		defer transceiver.Close()
		// make sure connection is alright
		for c := range conn {
			if c.Error() == nil {
				break
			}
			log.Println("ERROR: connection failed:", c.Error())
		}
		// close connection if Interrupted
		closeTransceiverOnSignal(transceiver)

		t := &transceiverConn{
			transceiver: transceiver,
			err:         nil,
			m:           &sync.RWMutex{},
		}

		// report error on failed conn and Increment error count
		go func(transceiverIndex int) {
			for c := range conn {
				s.setTransceiverErr(transceiverIndex, c.Error())
				if c.Error() != nil {
					go s.counters.connErrorCount.Increment()
					log.Printf("ERROR: transciever %v SMPP connection status: %v %v", transceiverIndex, c.Status(), c.Error())
				}
			}
		}(i)

		s.transceivers[i] = t
	}

	go func() {
		now := time.Now()
		burstLimit := 100
		rl := rate.NewLimiter(rate.Limit(s.MessageRate), burstLimit)
		var dest int
		currTransceiver := 0
		i := 0
		for i < s.NumMessages {

			if s.Mode == "dynamic" {
				dest = s.Dst + i
			} else {
				dest = s.Dst
			}

			req := &smpp.ShortMessage{
				Src:      s.Src,
				Dst:      strconv.Itoa(dest),
				Text:     pdutext.Raw(s.MessageText),
				Register: smpp.FinalDeliveryReceipt,
			}

			if s.Verbose == true {
				log.Println("Sending to ", dest)
			}

			r := rl.Reserve()
			if r == nil {
				panic("Something is wrong with rate limiter")
			}
			time.Sleep(r.Delay())
			currTransceiver = (currTransceiver + 1) % s.NumSessions
			if !s.checkTransceiverErr(currTransceiver) {
				go s.submitMessage(currTransceiver, req)
				s.counters.submittedCount.Increment()
				i++
			} else {
				i--
				go s.counters.connErrorCount.Increment()
			}
		}
		log.Println("Time elapsed sending:", time.Since(now))
	}()

	now := time.Now()
	loopTime := 100 * time.Millisecond
	loops := s.Wait * int(time.Second/loopTime)

	for i := 0; i < loops; i += 1 {
		time.Sleep(loopTime)
		if s.counters.successCount.Val()+s.counters.unknownRespCount.Val()+s.counters.sendErrorCount.Val() >= s.NumMessages {
			break
		}
		// Every 10 secs print a progress
		if i%100 == 0 {
			s.printStats(now)
		}
	}
	if s.counters.successCount.Val()+s.counters.unknownRespCount.Val()+s.counters.sendErrorCount.Val() < s.NumMessages {
		log.Println("WARNING: Waiting time is over and didn't receive enough responses.")
	}
	s.printStats(now)

}

func (s *SMPPerf) printStats(now time.Time) {
	log.Println("Time since start:", time.Since(now))
	log.Println("successCount:", s.counters.successCount.Val())
	log.Println("unknownRespCount:", s.counters.unknownRespCount.Val())
	log.Println("sendErrorCount:", s.counters.sendErrorCount.Val())
	log.Println("connErrorCount:", s.counters.connErrorCount.Val())
	log.Println("Submitted Messages:", s.counters.submittedCount.Val())
	for k, v := range s.counters.stateCounters.GetAll() {
		log.Println(k, v)
	}
}

func (s *SMPPerf) Purge() {
	receiptCount := NewSafeInt(0)

	transceiverHandler := func(p pdu.Body) {
		switch p.Header().ID {
		case pdu.DeliverSMID:
			go receiptCount.Increment()
		}
	}

	transceiver := s.getTransceiver()
	transceiver.Handler = transceiverHandler

	conn := transceiver.Bind() // make persistent connection.
	defer transceiver.Close()
	for c := range conn {
		if c.Error() == nil {
			break
		}
		log.Println("ERROR: Error connecting:", c.Error())
	}
	closeTransceiverOnSignal(transceiver)

	time.Sleep(time.Duration(s.Wait) * time.Second)
	log.Println("receiptCount:", receiptCount.Val())
}
<commit_msg>Update MessageStateType to be imported from pdutlv<commit_after>package smpperf

import (
	"log"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"golang.org/x/time/rate"

	"github.com/veoo/go-smpp/smpp"
	"github.com/veoo/go-smpp/smpp/pdu"
	"github.com/veoo/go-smpp/smpp/pdu/pdufield"
	"github.com/veoo/go-smpp/smpp/pdu/pdutext"
)

type SMPPerf struct {
	NumMessages          int
	NumSessions          int
	MessageRate          int
	MessageText          string
	Wait                 int
	User                 string
	Password             string
	Host                 string
	Mode                 string
	Dst                  int
	Src                  string
	Verbose              bool
	transceivers         []*transceiverConn
	counters             *counters
	msgIDToTransceiverID *ConcurrentStringMap
}

type transceiverConn struct {
	transceiver *smpp.Transceiver
	err         error
	m           *sync.RWMutex
}

func (s *SMPPerf) setTransceiverErr(transceiverIndex int, err error) {
	s.transceivers[transceiverIndex].m.Lock()
	defer s.transceivers[transceiverIndex].m.Unlock()
	s.transceivers[transceiverIndex].err = err
	log.Println(err)
}

func (s *SMPPerf) checkTransceiverErr(transceiverIndex int) bool {
	s.transceivers[transceiverIndex].m.Lock()
	defer s.transceivers[transceiverIndex].m.Unlock()
	return s.transceivers[transceiverIndex].err != nil

}

func (s *SMPPerf) submitMessage(transceiverIndex int, req *smpp.ShortMessage) {
	sm, err := s.transceivers[transceiverIndex].transceiver.Submit(req)
	if err != nil {
		if err == smpp.ErrNotConnected {
			go s.counters.connErrorCount.Increment()
		} else {
			go s.counters.sendErrorCount.Increment()
		}
	} else {
		transceiverID := strconv.Itoa(transceiverIndex)
		s.msgIDToTransceiverID.Set(sm.RespID(), transceiverID)
	}
}

type counters struct {
	successCount     *SafeInt
	sendErrorCount   *SafeInt
	unknownRespCount *SafeInt
	connErrorCount   *SafeInt
	submittedCount   *SafeInt
	stateCounters    *ConcurrentIntMap
}

func newCounters() *counters {
	return &counters{
		successCount:     NewSafeInt(0),
		sendErrorCount:   NewSafeInt(0),
		unknownRespCount: NewSafeInt(0),
		connErrorCount:   NewSafeInt(0),
		submittedCount:   NewSafeInt(0),
		stateCounters:    NewConcurrentIntMap(),
	}
}

func closeTransceiverOnSignal(trans *smpp.Transceiver) {
	go func() {
		signalChannel := make(chan os.Signal, 1)
		signal.Notify(signalChannel, os.Interrupt, syscall.SIGTERM)
		sig := <-signalChannel
		log.Println("WARNING:", sig, "signal caught, exiting.")
		trans.Close()
		os.Exit(0)
	}()
}

func getMessageID(p pdu.Body) string {
	tlv := p.TLVFields()
	if tlv == nil {
		return ""
	}
	field := tlv[pdufield.ReceiptedMessageID]
	if field == nil {
		return ""
	}
	return strings.TrimRight(string(field.Bytes()), "\x00")
}

func (s *SMPPerf) getTransceiver() *smpp.Transceiver {
	return &smpp.Transceiver{
		Addr:        s.Host,
		User:        s.User,
		Passwd:      s.Password,
		RespTimeout: 10 * time.Second,
		EnquireLink: 1 * time.Second,
	}
}

func (s *SMPPerf) countState(counterMap *ConcurrentIntMap, state pdufield.MessageStateType) {
	if _, ok := counterMap.Get(state); !ok {
		counterMap.Create(state)
	}
	go counterMap.Increment(state)
	return
}

func (s *SMPPerf) isFinalState(state pdufield.MessageStateType) bool {
	// Expired, Delivered, Undeliverable, Rejected, unsure about Deleted
	return state == pdufield.Expired || state == pdufield.Delivered || state == pdufield.Undeliverable || state == pdufield.Rejected || state == pdufield.Deleted
}

func (s *SMPPerf) SendMessages() {
	s.counters = newCounters()
	s.transceivers = make([]*transceiverConn, s.NumSessions)
	s.msgIDToTransceiverID = NewConcurrentStringMap()

	for i := 0; i < s.NumSessions; i++ {
		transceiverID := strconv.Itoa(i)
		transceiverHandler := func(p pdu.Body) {
			switch p.Header().ID {
			case pdu.DeliverSMID:
				// TODO: check here the resp data is correct
				msgID := getMessageID(p)

				t, ok := s.msgIDToTransceiverID.Get(msgID)
				if !ok {
					log.Printf("ERROR: message %s not found in transceiver %v", msgID, transceiverID)
				} else if t != transceiverID {
					log.Printf("ERROR: message %s was received in wrong transceiver %s", msgID, transceiverID)
				}

				state := pdutlv.MessageStateType(p.TLVFields()[pdutlv.MessageStateOption].Bytes()[0])
				s.countState(s.counters.stateCounters, state)

				if s.isFinalState(state) {
					go s.counters.successCount.Increment()
				}

			case pdu.UnbindID:
				log.Println("ERROR: They are unbinding me :(")
			case pdu.SubmitSMRespID:
				// Fix something florix?
			default:
				go log.Println(p.Header().ID.String(), p.Header().Status.Error())
				go s.counters.unknownRespCount.Increment()
			}
		}

		transceiver := s.getTransceiver()
		transceiver.Handler = transceiverHandler

		conn := transceiver.Bind() // make persistent connection.
		defer transceiver.Close()
		// make sure connection is alright
		for c := range conn {
			if c.Error() == nil {
				break
			}
			log.Println("ERROR: connection failed:", c.Error())
		}
		// close connection if Interrupted
		closeTransceiverOnSignal(transceiver)

		t := &transceiverConn{
			transceiver: transceiver,
			err:         nil,
			m:           &sync.RWMutex{},
		}

		// report error on failed conn and Increment error count
		go func(transceiverIndex int) {
			for c := range conn {
				s.setTransceiverErr(transceiverIndex, c.Error())
				if c.Error() != nil {
					go s.counters.connErrorCount.Increment()
					log.Printf("ERROR: transciever %v SMPP connection status: %v %v", transceiverIndex, c.Status(), c.Error())
				}
			}
		}(i)

		s.transceivers[i] = t
	}

	go func() {
		now := time.Now()
		burstLimit := 100
		rl := rate.NewLimiter(rate.Limit(s.MessageRate), burstLimit)
		var dest int
		currTransceiver := 0
		i := 0
		for i < s.NumMessages {

			if s.Mode == "dynamic" {
				dest = s.Dst + i
			} else {
				dest = s.Dst
			}

			req := &smpp.ShortMessage{
				Src:      s.Src,
				Dst:      strconv.Itoa(dest),
				Text:     pdutext.Raw(s.MessageText),
				Register: smpp.FinalDeliveryReceipt,
			}

			if s.Verbose == true {
				log.Println("Sending to ", dest)
			}

			r := rl.Reserve()
			if r == nil {
				panic("Something is wrong with rate limiter")
			}
			time.Sleep(r.Delay())
			currTransceiver = (currTransceiver + 1) % s.NumSessions
			if !s.checkTransceiverErr(currTransceiver) {
				go s.submitMessage(currTransceiver, req)
				s.counters.submittedCount.Increment()
				i++
			} else {
				i--
				go s.counters.connErrorCount.Increment()
			}
		}
		log.Println("Time elapsed sending:", time.Since(now))
	}()

	now := time.Now()
	loopTime := 100 * time.Millisecond
	loops := s.Wait * int(time.Second/loopTime)

	for i := 0; i < loops; i += 1 {
		time.Sleep(loopTime)
		if s.counters.successCount.Val()+s.counters.unknownRespCount.Val()+s.counters.sendErrorCount.Val() >= s.NumMessages {
			break
		}
		// Every 10 secs print a progress
		if i%100 == 0 {
			s.printStats(now)
		}
	}
	if s.counters.successCount.Val()+s.counters.unknownRespCount.Val()+s.counters.sendErrorCount.Val() < s.NumMessages {
		log.Println("WARNING: Waiting time is over and didn't receive enough responses.")
	}
	s.printStats(now)

}

func (s *SMPPerf) printStats(now time.Time) {
	log.Println("Time since start:", time.Since(now))
	log.Println("successCount:", s.counters.successCount.Val())
	log.Println("unknownRespCount:", s.counters.unknownRespCount.Val())
	log.Println("sendErrorCount:", s.counters.sendErrorCount.Val())
	log.Println("connErrorCount:", s.counters.connErrorCount.Val())
	log.Println("Submitted Messages:", s.counters.submittedCount.Val())
	for k, v := range s.counters.stateCounters.GetAll() {
		log.Println(k, v)
	}
}

func (s *SMPPerf) Purge() {
	receiptCount := NewSafeInt(0)

	transceiverHandler := func(p pdu.Body) {
		switch p.Header().ID {
		case pdu.DeliverSMID:
			go receiptCount.Increment()
		}
	}

	transceiver := s.getTransceiver()
	transceiver.Handler = transceiverHandler

	conn := transceiver.Bind() // make persistent connection.
	defer transceiver.Close()
	for c := range conn {
		if c.Error() == nil {
			break
		}
		log.Println("ERROR: Error connecting:", c.Error())
	}
	closeTransceiverOnSignal(transceiver)

	time.Sleep(time.Duration(s.Wait) * time.Second)
	log.Println("receiptCount:", receiptCount.Val())
}
<|endoftext|>


<commit_before>package dochaincore

import (
	"context"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"strings"
	"time"

	"github.com/digitalocean/godo"
	"golang.org/x/oauth2"
)

type Core struct {
	DropletID   int
	IPv4Address string
	IPv6Address string

	ssh *sshKeyPair
}

type Option func(*options)

func DropletName(name string) Option {
	return func(opt *options) { opt.dropletName = name }
}

func DropletRegion(region string) Option {
	return func(opt *options) { opt.dropletRegion = region }
}

func DropletSize(size string) Option {
	return func(opt *options) { opt.dropletSize = size }
}

func VolumeSizeGB(gb int64) Option {
	return func(opt *options) { opt.volumeSize = gb }
}

type options struct {
	dropletName   string
	dropletRegion string
	dropletSize   string
	volumeSize    int64
}

// Deploy builds and deploys an instance of Chain Core on a DigitalOcean
// droplet. It requires a DigitalOcean access token and optionally takes
// a variadic number of configuration options.
func Deploy(ctx context.Context, accessToken string, opts ...Option) (*Core, error) {
	opt := options{
		dropletName:   "chain-core",
		dropletRegion: "sfo2",
		dropletSize:   "1gb",
		volumeSize:    100,
	}
	for _, o := range opts {
		o(&opt)
	}

	keypair, err := createSSHKeyPair()
	if err != nil {
		return nil, err
	}

	oauthClient := oauth2.NewClient(ctx, oauth2.StaticTokenSource(
		&oauth2.Token{AccessToken: accessToken},
	))
	client := godo.NewClient(oauthClient)

	// Blockchains require storage. Make a volume that we can attach
	// to the droplet. Chain Core will store blockchain data on the volume.
	volume, _, err := client.Storage.CreateVolume(ctx, &godo.VolumeCreateRequest{
		Region:        opt.dropletRegion,
		Name:          fmt.Sprintf("%s-storage", opt.dropletName),
		Description:   "Chain Core storage volume",
		SizeGigaBytes: opt.volumeSize,
	})
	if err != nil {
		return nil, err
	}

	// Query all the SSH keys on the account so we can include them
	// in the droplet.
	sshKeys, _, err := client.Keys.List(ctx, nil)
	if err != nil {
		return nil, err
	}

	// Build user data to initialize the droplet as a Chain Core
	// instance.
	userData, err := buildUserData(keypair)
	if err != nil {
		return nil, err
	}

	// Launch the DigitalOcean droplet.
	createRequest := &godo.DropletCreateRequest{
		Name:     opt.dropletName,
		Region:   opt.dropletRegion,
		Size:     opt.dropletSize,
		IPv6:     true,
		UserData: userData,
		Image: godo.DropletCreateImage{
			Slug: "ubuntu-17-04-x64",
		},
		Volumes: []godo.DropletCreateVolume{
			{ID: volume.ID},
		},
	}
	for _, key := range sshKeys {
		keyToAdd := godo.DropletCreateSSHKey{ID: key.ID}
		createRequest.SSHKeys = append(createRequest.SSHKeys, keyToAdd)
	}

	droplet, _, err := client.Droplets.Create(ctx, createRequest)
	if err != nil {
		return nil, err
	}

	core := &Core{
		DropletID: droplet.ID,
		ssh:       keypair,
	}

	// A just-created droplet won't have any of the network IP addresses
	// quite yet. We have to poll until the droplet is provisioned and
	// they're populated.
	for attempt := 1; core.IPv4Address == "" || core.IPv6Address == ""; attempt++ {
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case <-time.After(time.Duration(attempt) * time.Second): /// linear backoff
		}

		droplet, _, err := client.Droplets.Get(ctx, core.DropletID)
		if err != nil {
			return nil, err
		}

		for _, nv4 := range droplet.Networks.V4 {
			if nv4.IPAddress != "" {
				core.IPv4Address = nv4.IPAddress
			}
		}
		for _, nv6 := range droplet.Networks.V6 {
			if nv6.IPAddress != "" {
				core.IPv6Address = nv6.IPAddress
			}
		}
		if attempt >= 10 {
			return nil, fmt.Errorf("timeout waiting for provisioning of droplet %d", core.DropletID)
		}
	}
	return core, nil
}

// WaitForSSH waits until port 22 on the provided Chain Core's host is opened.
func WaitForSSH(ctx context.Context, c *Core) error {
	return waitForPort(ctx, c.IPv4Address, 22)
}

// WaitForHTTP waits until Chain Core begins listening on port 1999.
func WaitForHTTP(ctx context.Context, c *Core) error {
	return waitForPort(ctx, c.IPv4Address, 1999)
}

func waitForPort(ctx context.Context, host string, port int) (err error) {
	var conn net.Conn
	for conn == nil {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(5 * time.Second):
			conn, err = net.Dial("tcp", fmt.Sprintf("%s:%d", host, port))
		}
	}
	conn.Close()
	return err
}

// CreateClientToken sets up a Chain Core client token for the
// provided Core.
func CreateClientToken(ctx context.Context, c *Core) (string, error) {
	const createClientToken = `
	docker exec dochaincore /usr/bin/chain/corectl create-token do client-readwrite
	`
	// TODO(jackson): remove the ssh key from authorized_keys before
	// closing the SSH session.

	session, err := connect(ctx, c.IPv4Address, c.ssh)
	if err != nil {
		return "", err
	}
	defer session.Close()

	rOut, err := session.StdoutPipe()
	if err != nil {
		return "", err
	}
	rErr, err := session.StderrPipe()
	if err != nil {
		return "", err
	}

	err = session.Start(createClientToken)
	if err != nil {
		return "", err
	}
	outputBytes, err := ioutil.ReadAll(io.MultiReader(rOut, rErr))
	if err != nil {
		return "", err
	}
	output := strings.TrimSpace(string(outputBytes))
	if !strings.HasPrefix(output, "do:") {
		return "", errors.New(output)
	}
	return output, nil
}
<commit_msg>deploy: enable monitoring by default<commit_after>package dochaincore

import (
	"context"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"strings"
	"time"

	"github.com/digitalocean/godo"
	"golang.org/x/oauth2"
)

type Core struct {
	DropletID   int
	IPv4Address string
	IPv6Address string

	ssh *sshKeyPair
}

type Option func(*options)

func DropletName(name string) Option {
	return func(opt *options) { opt.dropletName = name }
}

func DropletRegion(region string) Option {
	return func(opt *options) { opt.dropletRegion = region }
}

func DropletSize(size string) Option {
	return func(opt *options) { opt.dropletSize = size }
}

func VolumeSizeGB(gb int64) Option {
	return func(opt *options) { opt.volumeSize = gb }
}

type options struct {
	dropletName   string
	dropletRegion string
	dropletSize   string
	volumeSize    int64
}

// Deploy builds and deploys an instance of Chain Core on a DigitalOcean
// droplet. It requires a DigitalOcean access token and optionally takes
// a variadic number of configuration options.
func Deploy(ctx context.Context, accessToken string, opts ...Option) (*Core, error) {
	opt := options{
		dropletName:   "chain-core",
		dropletRegion: "sfo2",
		dropletSize:   "1gb",
		volumeSize:    100,
	}
	for _, o := range opts {
		o(&opt)
	}

	keypair, err := createSSHKeyPair()
	if err != nil {
		return nil, err
	}

	oauthClient := oauth2.NewClient(ctx, oauth2.StaticTokenSource(
		&oauth2.Token{AccessToken: accessToken},
	))
	client := godo.NewClient(oauthClient)

	// Blockchains require storage. Make a volume that we can attach
	// to the droplet. Chain Core will store blockchain data on the volume.
	volume, _, err := client.Storage.CreateVolume(ctx, &godo.VolumeCreateRequest{
		Region:        opt.dropletRegion,
		Name:          fmt.Sprintf("%s-storage", opt.dropletName),
		Description:   "Chain Core storage volume",
		SizeGigaBytes: opt.volumeSize,
	})
	if err != nil {
		return nil, err
	}

	// Query all the SSH keys on the account so we can include them
	// in the droplet.
	sshKeys, _, err := client.Keys.List(ctx, nil)
	if err != nil {
		return nil, err
	}

	// Build user data to initialize the droplet as a Chain Core
	// instance.
	userData, err := buildUserData(keypair)
	if err != nil {
		return nil, err
	}

	// Launch the DigitalOcean droplet.
	createRequest := &godo.DropletCreateRequest{
		Name:       opt.dropletName,
		Region:     opt.dropletRegion,
		Size:       opt.dropletSize,
		IPv6:       true,
		Monitoring: true,
		UserData:   userData,
		Image: godo.DropletCreateImage{
			Slug: "ubuntu-17-04-x64",
		},
		Volumes: []godo.DropletCreateVolume{
			{ID: volume.ID},
		},
	}
	for _, key := range sshKeys {
		keyToAdd := godo.DropletCreateSSHKey{ID: key.ID}
		createRequest.SSHKeys = append(createRequest.SSHKeys, keyToAdd)
	}

	droplet, _, err := client.Droplets.Create(ctx, createRequest)
	if err != nil {
		return nil, err
	}

	core := &Core{
		DropletID: droplet.ID,
		ssh:       keypair,
	}

	// A just-created droplet won't have any of the network IP addresses
	// quite yet. We have to poll until the droplet is provisioned and
	// they're populated.
	for attempt := 1; core.IPv4Address == "" || core.IPv6Address == ""; attempt++ {
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		case <-time.After(time.Duration(attempt) * time.Second): /// linear backoff
		}

		droplet, _, err := client.Droplets.Get(ctx, core.DropletID)
		if err != nil {
			return nil, err
		}

		for _, nv4 := range droplet.Networks.V4 {
			if nv4.IPAddress != "" {
				core.IPv4Address = nv4.IPAddress
			}
		}
		for _, nv6 := range droplet.Networks.V6 {
			if nv6.IPAddress != "" {
				core.IPv6Address = nv6.IPAddress
			}
		}
		if attempt >= 10 {
			return nil, fmt.Errorf("timeout waiting for provisioning of droplet %d", core.DropletID)
		}
	}
	return core, nil
}

// WaitForSSH waits until port 22 on the provided Chain Core's host is opened.
func WaitForSSH(ctx context.Context, c *Core) error {
	return waitForPort(ctx, c.IPv4Address, 22)
}

// WaitForHTTP waits until Chain Core begins listening on port 1999.
func WaitForHTTP(ctx context.Context, c *Core) error {
	return waitForPort(ctx, c.IPv4Address, 1999)
}

func waitForPort(ctx context.Context, host string, port int) (err error) {
	var conn net.Conn
	for conn == nil {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(5 * time.Second):
			conn, err = net.Dial("tcp", fmt.Sprintf("%s:%d", host, port))
		}
	}
	conn.Close()
	return err
}

// CreateClientToken sets up a Chain Core client token for the
// provided Core.
func CreateClientToken(ctx context.Context, c *Core) (string, error) {
	const createClientToken = `
	docker exec dochaincore /usr/bin/chain/corectl create-token do client-readwrite
	`
	// TODO(jackson): remove the ssh key from authorized_keys before
	// closing the SSH session.

	session, err := connect(ctx, c.IPv4Address, c.ssh)
	if err != nil {
		return "", err
	}
	defer session.Close()

	rOut, err := session.StdoutPipe()
	if err != nil {
		return "", err
	}
	rErr, err := session.StderrPipe()
	if err != nil {
		return "", err
	}

	err = session.Start(createClientToken)
	if err != nil {
		return "", err
	}
	outputBytes, err := ioutil.ReadAll(io.MultiReader(rOut, rErr))
	if err != nil {
		return "", err
	}
	output := strings.TrimSpace(string(outputBytes))
	if !strings.HasPrefix(output, "do:") {
		return "", errors.New(output)
	}
	return output, nil
}
<|endoftext|>


<commit_before>package deploy

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"regexp"
	"strings"
	"time"

	"github.com/remind101/deploy/Godeps/_workspace/src/github.com/codegangsta/cli"
	"github.com/remind101/deploy/Godeps/_workspace/src/github.com/github/hub/git"
	hub "github.com/remind101/deploy/Godeps/_workspace/src/github.com/github/hub/github"
	"github.com/remind101/deploy/Godeps/_workspace/src/github.com/google/go-github/github"
)

const (
	Name  = "deploy"
	Usage = "A command for creating GitHub deployments"
)

const (
	DefaultRef     = "master"
	DefaultTimeout = 20 * time.Second
)

var errTimeout = errors.New("Timed out waiting for build to start. Did you add a webhook to handle deployment events?")

func init() {
	cli.AppHelpTemplate = `USAGE:
   # Deploy the master branch of remind101/acme-inc to staging
   {{.Name}} --env=staging --ref=master remind101/acme-inc

   # Deploy HEAD of the current branch to staging
   {{.Name}} --env=staging remind101/acme-inc

   # Deploy the current GitHub repo to staging
   {{.Name}} --env=staging
{{if .Flags}}
OPTIONS:
   {{range .Flags}}{{.}}
   {{end}}{{end}}
`
}

var flags = []cli.Flag{
	cli.StringFlag{
		Name:  "ref, branch, commit, tag",
		Value: "",
		Usage: "The git ref to deploy. Can be a git commit, branch or tag.",
	},
	cli.StringFlag{
		Name:  "env, e",
		Value: "",
		Usage: "The environment to deploy to.",
	},
	cli.BoolFlag{
		Name:  "force, f",
		Usage: "Ignore commit status checks.",
	},
	cli.BoolFlag{
		Name:  "detached, d",
		Usage: "Don't wait for the deployment to complete.",
	},
	cli.BoolFlag{
		Name:  "quiet, q",
		Usage: "Silence any output to STDOUT.",
	},
	cli.BoolFlag{
		Name:  "update, u",
		Usage: "Update the binary",
	},
}

var ProtectedEnvironments = map[string]bool{
	"production": true,
	"prod":       true,
}

// NewApp returns a new cli.App for the deploy command.
func NewApp() *cli.App {
	app := cli.NewApp()
	app.Version = Version
	app.Name = Name
	app.Usage = Usage
	app.Flags = flags
	app.Action = func(c *cli.Context) {
		if c.Bool("update") {
			updater := NewUpdater()
			if err := updater.Update(); err != nil {
				fmt.Println(err)
				os.Exit(-1)
			} else {
				os.Exit(0)
			}
		}

		if err := RunDeploy(c); err != nil {
			msg := err.Error()
			if err, ok := err.(*github.ErrorResponse); ok {
				if strings.HasPrefix(err.Message, "Conflict: Commit status checks failed for") {
					msg = "Commit status checks failed. You can bypass commit status checks with the --force flag."
				} else if strings.HasPrefix(err.Message, "No ref found for") {
					msg = fmt.Sprintf("%s. Did you push it to GitHub?", err.Message)
				} else {
					msg = err.Message
				}
			}

			fmt.Println(msg)
			os.Exit(-1)
		}
	}

	return app
}

// RunDeploy performs a deploy.
func RunDeploy(c *cli.Context) error {
	var w io.Writer
	if c.Bool("quiet") {
		w = ioutil.Discard
	} else {
		w = c.App.Writer
	}

	h, err := hub.CurrentConfig().PromptForHost("github.com")
	if err != nil {
		return err
	}

	client, err := newGitHubClient(h)
	if err != nil {
		return err
	}

	nwo, err := Repo(c.Args())
	if err != nil {
		return err
	}

	owner, repo, err := SplitRepo(nwo, os.Getenv("GITHUB_ORGANIZATION"))
	if err != nil {
		return fmt.Errorf("Invalid GitHub repo: %s", nwo)
	}

	displayNewCommits(owner, repo, c, client)

	r, err := newDeploymentRequest(c)
	if err != nil {
		return err
	}

	fmt.Fprintf(w, "Deploying %s/%s@%s to %s...\n", owner, repo, *r.Ref, *r.Environment)

	if c.Bool("detached") {
		return nil
	}

	d, _, err := client.Repositories.CreateDeployment(owner, repo, r)
	if err != nil {
		return err
	}

	started := make(chan *github.DeploymentStatus)
	completed := make(chan *github.DeploymentStatus)

	go func() {
		started <- waitState(pendingStates, owner, repo, *d.ID, client)
	}()

	go func() {
		completed <- waitState(completedStates, owner, repo, *d.ID, client)
	}()

	select {
	case <-time.After(DefaultTimeout):
		return errTimeout
	case status := <-started:
		var url string
		if status.TargetURL != nil {
			url = *status.TargetURL
		}
		fmt.Fprintf(w, "%s\n", url)
	}

	status := <-completed

	if isFailed(*status.State) {
		return errors.New("Failed to deploy")
	}

	return nil
}

func displayNewCommits(owner string, repo string, c *cli.Context, client *github.Client) {
	opt := &github.DeploymentsListOptions{
		Environment: c.String("env"),
	}

	deployments, _, error := client.Repositories.ListDeployments(owner, repo, opt)
	if error == nil {
		sha := *deployments[0].SHA
		compare, _, cmp_error := client.Repositories.CompareCommits(owner, repo, sha, "master")
		if cmp_error == nil && len(compare.Commits) > 0 {
			fmt.Println("Deploying the following commits:\n")
			for _, commit := range compare.Commits {
				message := *commit.Commit.Message
				fmt.Printf("%-20s\t%s\n", *commit.Commit.Author.Name, strings.Split(message, "\n")[0])
			}
			fmt.Printf("\nSee entire diff here: https://github.com/%s/%s/compare/%s...master\n\n", owner, repo, sha)
		}
        }
}

func newDeploymentRequest(c *cli.Context) (*github.DeploymentRequest, error) {
	ref := Ref(c.String("ref"), git.Head)

	env := c.String("env")
	if env == "" {
		return nil, fmt.Errorf("--env flag is required")
	}

	if ProtectedEnvironments[env] {
		yes := askYN(fmt.Sprintf("Are you sure you want to deploy %s to %s?", ref, env))
		if !yes {
			return nil, fmt.Errorf("Deployment aborted.")
		}
	}

	var contexts *[]string
	if c.Bool("force") {
		s := []string{}
		contexts = &s
	}

	return &github.DeploymentRequest{
		Ref:              github.String(ref),
		Task:             github.String("deploy"),
		AutoMerge:        github.Bool(false),
		Environment:      github.String(env),
		RequiredContexts: contexts,
		Payload: map[string]interface{}{
			"force": c.Bool("force"),
		},
		// TODO Description:
	}, nil
}

var (
	pendingStates   = []string{"pending"}
	completedStates = []string{"success", "error", "failure"}
)

func isFailed(state string) bool {
	return state == "error" || state == "failure"
}

// waitState waits for a deployment status that matches the given states, then
// sends on the returned channel.
func waitState(states []string, owner, repo string, deploymentID int, c *github.Client) *github.DeploymentStatus {
	for {
		<-time.After(1 * time.Second)

		statuses, _, err := c.Repositories.ListDeploymentStatuses(owner, repo, deploymentID, nil)
		if err != nil {
			continue
		}

		status := firstStatus(states, statuses)
		if status != nil {
			return status
		}
	}
}

// firstStatus takes a slice of github.DeploymentStatus and returns the
// first status that matches the provided slice of states.
func firstStatus(states []string, statuses []github.DeploymentStatus) *github.DeploymentStatus {
	for _, ds := range statuses {
		for _, s := range states {
			if ds.State != nil && *ds.State == s {
				return &ds
			}
		}
	}

	return nil
}

// refRegex is a regular expression that matches a full git HEAD ref.
var refRegex = regexp.MustCompile(`^refs/heads/(.*)$`)

// Ref attempts to return the proper git ref to deploy. If a ref is provided,
// that will be returned. If not, it will fallback to calling headFunc. If an
// error is returned (not in a git repo), then it will fallback to DefaultRef.
func Ref(ref string, headFunc func() (string, error)) string {
	if ref != "" {
		return ref
	}

	ref, err := headFunc()
	if err != nil {
		// An error means that we're either not in a GitRepo or we're
		// not on a branch. In this case, we just fallback to the
		// DefaultRef.
		return DefaultRef
	}

	// Convert `refs/heads/test-deploy` => `test-deploy`
	return refRegex.ReplaceAllString(ref, "$1")
}

// Repo will determine the correct GitHub repo to deploy to, based on a set of
// arguments.
func Repo(arguments []string) (string, error) {
	if len(arguments) != 0 {
		return arguments[0], nil
	}

	remotes, err := hub.Remotes()
	if err != nil {
		return "", err
	}

	repo := GitHubRepo(remotes)
	if repo == "" {
		return repo, errors.New("no GitHub repo found in .git/config")
	}

	return repo, nil
}

// A regular expression that can convert a URL.Path into a GitHub repo name.
var remoteRegex = regexp.MustCompile(`^/(.*)\.git$`)

// GitHubRepo, given a list of git remotes, will determine what the GitHub repo
// is.
func GitHubRepo(remotes []hub.Remote) string {
	// We only want to look at the `origin` remote.
	remote := findRemote("origin", remotes)
	if remote == nil {
		return ""
	}

	// Remotes that are not pointed at a GitHub repo are not valid.
	if remote.URL.Host != "github.com" {
		return ""
	}

	// Convert `/remind101/acme-inc.git` => `remind101/acme-inc`.
	return remoteRegex.ReplaceAllString(remote.URL.Path, "$1")
}

func findRemote(name string, remotes []hub.Remote) *hub.Remote {
	for _, r := range remotes {
		if r.Name == name {
			return &r
		}
	}

	return nil
}

var errInvalidRepo = errors.New("invalid repo")

// SplitRepo splits a repo string in the form remind101/acme-inc into it's owner
// and repo components.
func SplitRepo(nwo, defaultOrg string) (owner string, repo string, err error) {
	parts := strings.Split(nwo, "/")

	// If we were only given a repo name, and a default organization is set,
	// we'll use the defaultOrg as the owner.
	if len(parts) == 1 && defaultOrg != "" && parts[0] != "" {
		owner = defaultOrg
		repo = parts[0]
		return
	}

	if len(parts) != 2 {
		err = errInvalidRepo
		return
	}

	owner = parts[0]
	repo = parts[1]

	return
}

func askYN(prompt string) bool {
	r := bufio.NewReader(os.Stdin)
	fmt.Printf("%s (y/N)\n", prompt)
	a, _ := r.ReadString('\n')
	return strings.ToUpper(a) == "Y\n"
}
<commit_msg>Use current branch instead of master<commit_after>package deploy

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"regexp"
	"strings"
	"time"

	"github.com/remind101/deploy/Godeps/_workspace/src/github.com/codegangsta/cli"
	"github.com/remind101/deploy/Godeps/_workspace/src/github.com/github/hub/git"
	hub "github.com/remind101/deploy/Godeps/_workspace/src/github.com/github/hub/github"
	"github.com/remind101/deploy/Godeps/_workspace/src/github.com/google/go-github/github"
)

const (
	Name  = "deploy"
	Usage = "A command for creating GitHub deployments"
)

const (
	DefaultRef     = "master"
	DefaultTimeout = 20 * time.Second
)

var errTimeout = errors.New("Timed out waiting for build to start. Did you add a webhook to handle deployment events?")

func init() {
	cli.AppHelpTemplate = `USAGE:
   # Deploy the master branch of remind101/acme-inc to staging
   {{.Name}} --env=staging --ref=master remind101/acme-inc

   # Deploy HEAD of the current branch to staging
   {{.Name}} --env=staging remind101/acme-inc

   # Deploy the current GitHub repo to staging
   {{.Name}} --env=staging
{{if .Flags}}
OPTIONS:
   {{range .Flags}}{{.}}
   {{end}}{{end}}
`
}

var flags = []cli.Flag{
	cli.StringFlag{
		Name:  "ref, branch, commit, tag",
		Value: "",
		Usage: "The git ref to deploy. Can be a git commit, branch or tag.",
	},
	cli.StringFlag{
		Name:  "env, e",
		Value: "",
		Usage: "The environment to deploy to.",
	},
	cli.BoolFlag{
		Name:  "force, f",
		Usage: "Ignore commit status checks.",
	},
	cli.BoolFlag{
		Name:  "detached, d",
		Usage: "Don't wait for the deployment to complete.",
	},
	cli.BoolFlag{
		Name:  "quiet, q",
		Usage: "Silence any output to STDOUT.",
	},
	cli.BoolFlag{
		Name:  "update, u",
		Usage: "Update the binary",
	},
}

var ProtectedEnvironments = map[string]bool{
	"production": true,
	"prod":       true,
}

// NewApp returns a new cli.App for the deploy command.
func NewApp() *cli.App {
	app := cli.NewApp()
	app.Version = Version
	app.Name = Name
	app.Usage = Usage
	app.Flags = flags
	app.Action = func(c *cli.Context) {
		if c.Bool("update") {
			updater := NewUpdater()
			if err := updater.Update(); err != nil {
				fmt.Println(err)
				os.Exit(-1)
			} else {
				os.Exit(0)
			}
		}

		if err := RunDeploy(c); err != nil {
			msg := err.Error()
			if err, ok := err.(*github.ErrorResponse); ok {
				if strings.HasPrefix(err.Message, "Conflict: Commit status checks failed for") {
					msg = "Commit status checks failed. You can bypass commit status checks with the --force flag."
				} else if strings.HasPrefix(err.Message, "No ref found for") {
					msg = fmt.Sprintf("%s. Did you push it to GitHub?", err.Message)
				} else {
					msg = err.Message
				}
			}

			fmt.Println(msg)
			os.Exit(-1)
		}
	}

	return app
}

// RunDeploy performs a deploy.
func RunDeploy(c *cli.Context) error {
	var w io.Writer
	if c.Bool("quiet") {
		w = ioutil.Discard
	} else {
		w = c.App.Writer
	}

	h, err := hub.CurrentConfig().PromptForHost("github.com")
	if err != nil {
		return err
	}

	client, err := newGitHubClient(h)
	if err != nil {
		return err
	}

	nwo, err := Repo(c.Args())
	if err != nil {
		return err
	}

	owner, repo, err := SplitRepo(nwo, os.Getenv("GITHUB_ORGANIZATION"))
	if err != nil {
		return fmt.Errorf("Invalid GitHub repo: %s", nwo)
	}

	displayNewCommits(owner, repo, c, client)

	r, err := newDeploymentRequest(c)
	if err != nil {
		return err
	}

	fmt.Fprintf(w, "Deploying %s/%s@%s to %s...\n", owner, repo, *r.Ref, *r.Environment)

	if c.Bool("detached") {
		return nil
	}

	d, _, err := client.Repositories.CreateDeployment(owner, repo, r)
	if err != nil {
		return err
	}

	started := make(chan *github.DeploymentStatus)
	completed := make(chan *github.DeploymentStatus)

	go func() {
		started <- waitState(pendingStates, owner, repo, *d.ID, client)
	}()

	go func() {
		completed <- waitState(completedStates, owner, repo, *d.ID, client)
	}()

	select {
	case <-time.After(DefaultTimeout):
		return errTimeout
	case status := <-started:
		var url string
		if status.TargetURL != nil {
			url = *status.TargetURL
		}
		fmt.Fprintf(w, "%s\n", url)
	}

	status := <-completed

	if isFailed(*status.State) {
		return errors.New("Failed to deploy")
	}

	return nil
}

func displayNewCommits(owner string, repo string, c *cli.Context, client *github.Client) {
	ref := Ref(c.String("ref"), git.Head)

	opt := &github.DeploymentsListOptions{
		Environment: c.String("env"),
	}

	deployments, _, error := client.Repositories.ListDeployments(owner, repo, opt)
	if error == nil {
		sha := *deployments[0].SHA
		compare, _, cmp_error := client.Repositories.CompareCommits(owner, repo, sha, ref)
		if cmp_error == nil && len(compare.Commits) > 0 {
			fmt.Println("Deploying the following commits:\n")
			for _, commit := range compare.Commits {
				message := *commit.Commit.Message
				fmt.Printf("%-20s\t%s\n", *commit.Commit.Author.Name, strings.Split(message, "\n")[0])
			}
			fmt.Printf("\nSee entire diff here: https://github.com/%s/%s/compare/%s...%s\n\n", owner, repo, sha, ref)
		}
        }
}

func newDeploymentRequest(c *cli.Context) (*github.DeploymentRequest, error) {
	ref := Ref(c.String("ref"), git.Head)

	env := c.String("env")
	if env == "" {
		return nil, fmt.Errorf("--env flag is required")
	}

	if ProtectedEnvironments[env] {
		yes := askYN(fmt.Sprintf("Are you sure you want to deploy %s to %s?", ref, env))
		if !yes {
			return nil, fmt.Errorf("Deployment aborted.")
		}
	}

	var contexts *[]string
	if c.Bool("force") {
		s := []string{}
		contexts = &s
	}

	return &github.DeploymentRequest{
		Ref:              github.String(ref),
		Task:             github.String("deploy"),
		AutoMerge:        github.Bool(false),
		Environment:      github.String(env),
		RequiredContexts: contexts,
		Payload: map[string]interface{}{
			"force": c.Bool("force"),
		},
		// TODO Description:
	}, nil
}

var (
	pendingStates   = []string{"pending"}
	completedStates = []string{"success", "error", "failure"}
)

func isFailed(state string) bool {
	return state == "error" || state == "failure"
}

// waitState waits for a deployment status that matches the given states, then
// sends on the returned channel.
func waitState(states []string, owner, repo string, deploymentID int, c *github.Client) *github.DeploymentStatus {
	for {
		<-time.After(1 * time.Second)

		statuses, _, err := c.Repositories.ListDeploymentStatuses(owner, repo, deploymentID, nil)
		if err != nil {
			continue
		}

		status := firstStatus(states, statuses)
		if status != nil {
			return status
		}
	}
}

// firstStatus takes a slice of github.DeploymentStatus and returns the
// first status that matches the provided slice of states.
func firstStatus(states []string, statuses []github.DeploymentStatus) *github.DeploymentStatus {
	for _, ds := range statuses {
		for _, s := range states {
			if ds.State != nil && *ds.State == s {
				return &ds
			}
		}
	}

	return nil
}

// refRegex is a regular expression that matches a full git HEAD ref.
var refRegex = regexp.MustCompile(`^refs/heads/(.*)$`)

// Ref attempts to return the proper git ref to deploy. If a ref is provided,
// that will be returned. If not, it will fallback to calling headFunc. If an
// error is returned (not in a git repo), then it will fallback to DefaultRef.
func Ref(ref string, headFunc func() (string, error)) string {
	if ref != "" {
		return ref
	}

	ref, err := headFunc()
	if err != nil {
		// An error means that we're either not in a GitRepo or we're
		// not on a branch. In this case, we just fallback to the
		// DefaultRef.
		return DefaultRef
	}

	// Convert `refs/heads/test-deploy` => `test-deploy`
	return refRegex.ReplaceAllString(ref, "$1")
}

// Repo will determine the correct GitHub repo to deploy to, based on a set of
// arguments.
func Repo(arguments []string) (string, error) {
	if len(arguments) != 0 {
		return arguments[0], nil
	}

	remotes, err := hub.Remotes()
	if err != nil {
		return "", err
	}

	repo := GitHubRepo(remotes)
	if repo == "" {
		return repo, errors.New("no GitHub repo found in .git/config")
	}

	return repo, nil
}

// A regular expression that can convert a URL.Path into a GitHub repo name.
var remoteRegex = regexp.MustCompile(`^/(.*)\.git$`)

// GitHubRepo, given a list of git remotes, will determine what the GitHub repo
// is.
func GitHubRepo(remotes []hub.Remote) string {
	// We only want to look at the `origin` remote.
	remote := findRemote("origin", remotes)
	if remote == nil {
		return ""
	}

	// Remotes that are not pointed at a GitHub repo are not valid.
	if remote.URL.Host != "github.com" {
		return ""
	}

	// Convert `/remind101/acme-inc.git` => `remind101/acme-inc`.
	return remoteRegex.ReplaceAllString(remote.URL.Path, "$1")
}

func findRemote(name string, remotes []hub.Remote) *hub.Remote {
	for _, r := range remotes {
		if r.Name == name {
			return &r
		}
	}

	return nil
}

var errInvalidRepo = errors.New("invalid repo")

// SplitRepo splits a repo string in the form remind101/acme-inc into it's owner
// and repo components.
func SplitRepo(nwo, defaultOrg string) (owner string, repo string, err error) {
	parts := strings.Split(nwo, "/")

	// If we were only given a repo name, and a default organization is set,
	// we'll use the defaultOrg as the owner.
	if len(parts) == 1 && defaultOrg != "" && parts[0] != "" {
		owner = defaultOrg
		repo = parts[0]
		return
	}

	if len(parts) != 2 {
		err = errInvalidRepo
		return
	}

	owner = parts[0]
	repo = parts[1]

	return
}

func askYN(prompt string) bool {
	r := bufio.NewReader(os.Stdin)
	fmt.Printf("%s (y/N)\n", prompt)
	a, _ := r.ReadString('\n')
	return strings.ToUpper(a) == "Y\n"
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"io/ioutil"
	"log"
	"sync/atomic"

	sftpclient "github.com/bowlhat/sftp-client"

	"strings"

	"github.com/cheggaaa/pb"
	"gopkg.in/yaml.v2"
)

type configuration struct {
	Connection struct {
		Username string
		Password string
		Hostname string
		Port     int
	}
	Backup struct {
		To   string
		From []string
	}
	Download []sftpclient.FolderMapping
	Upload   []sftpclient.FolderMapping
}

type errorResponse struct {
	Err error
}

var (
	configfilename = flag.String("config", "", "YAML configuration file")
	doBackup       = flag.Bool("backup", true, "Backup files on remote system")
	doDownload     = flag.Bool("download", false, "Download files from remote system")
	doUpload       = flag.Bool("upload", false, "Upload local files to remote system")
	debugging      = flag.Int("debug", 1, "Spew debugging info, e.g. output every file as it's touched. 0=silent, 1=progress, 2=verbose, 3=firehose")

	remoteFiles []string
)

func main() {
	flag.Parse()

	if *configfilename == "" {
		log.Fatalf("config flag indicating configuration file is required")
	}

	configfile, err := ioutil.ReadFile(*configfilename)
	if err != nil {
		log.Fatalf("could not read config file: %v", err)
	}

	config := configuration{}
	if err := yaml.Unmarshal(configfile, &config); err != nil {
		log.Fatalf("YAML error: %v", err)
	}

	upload := *doUpload
	backup := *doBackup
	download := *doDownload

	// start ssh client on tcp connection
	sftp, err := sftpclient.New(
		config.Connection.Hostname,
		config.Connection.Port,
		config.Connection.Username,
		config.Connection.Password)

	if err != nil {
		log.Fatal(err)
	}

	var bar *pb.ProgressBar
	if *debugging <= 1 {
		bar = pb.New(0)
	}

	if backup == true {
		files, err := sftp.FindAllRemoteFiles(config.Backup.From)
		if err != nil {
			log.Fatalln(err)
		}

		if *debugging <= 1 {
			bar.Total = 0
			bar.Prefix("Backing-up... ")
			bar.Start()
		}

		saved, errors := sftp.BackupFiles(config.Backup.To, files)

		go func() {
			for range saved {
				bar.Increment()
			}
		}()

		errorsEncountered := false
		for err := range errors {
			errorsEncountered = true
			log.Println(err.Err)
		}
		if errorsEncountered {
			log.Fatalln("Backup failed. Quitting.")
		}

		if *debugging <= 1 {
			bar.FinishPrint("Backup complete")
		}
	}

	if download == true {
		downloadedFile := make(chan bool)
		downloadDone := make(chan bool)
		errorsChannel := make(chan errorResponse)

		if *debugging <= 1 {
			bar.Total = 0
			bar.Prefix("Downloading... ")
			bar.Start()
		}

		for _, folder := range config.Download {
			go func(f sftpclient.FolderMapping) {
				defer func() {
					downloadDone <- true
				}()

				media, err := sftp.FindAllRemoteFiles([]string{f.Remote})
				if err != nil {
					errorsChannel <- errorResponse{err}
					return
				}
				atomic.AddInt64(&bar.Total, int64(len(media)))
				for _, remotefile := range media {
					trimmed := strings.TrimPrefix(remotefile, f.Remote)
					localfile := strings.Join([]string{f.Local, trimmed}, "/")
					if err := sftp.GetFile(localfile, remotefile); err != nil {
						errorsChannel <- errorResponse{err}
					}
					downloadedFile <- true
				}
			}(folder)
		}

		go func() {
			for range downloadedFile {
				<-downloadedFile
				bar.Increment()
			}
		}()

		go func() {
			for err := range errorsChannel {
				log.Println(err)
			}
		}()

		for range config.Download {
			<-downloadDone
		}
		close(downloadDone)
		close(downloadedFile)
		close(errorsChannel)

		if *debugging <= 1 {
			bar.FinishPrint("Download complete")
		}
	}

	if upload == true {
		if *debugging <= 1 {
			bar.Prefix("Uploading... ")
			bar.Start()
		}

		errorChannel, countChannel, copiedCountChannel := sftp.Upload(config.Upload)

		go func() {
			for range countChannel {
				atomic.AddInt64(&bar.Total, 1)
			}
		}()
		go func() {
			for range copiedCountChannel {
				bar.Increment()
			}
		}()

		for response := range errorChannel {
			log.Println(response.Err)
		}

		if *debugging <= 1 {
			bar.FinishPrint("Upload finished")
		}
	}
}
<commit_msg>ensure progressbar is reset<commit_after>package main

import (
	"flag"
	"io/ioutil"
	"log"
	"sync/atomic"

	sftpclient "github.com/bowlhat/sftp-client"

	"strings"

	"github.com/cheggaaa/pb"
	"gopkg.in/yaml.v2"
)

type configuration struct {
	Connection struct {
		Username string
		Password string
		Hostname string
		Port     int
	}
	Backup struct {
		To   string
		From []string
	}
	Download []sftpclient.FolderMapping
	Upload   []sftpclient.FolderMapping
}

type errorResponse struct {
	Err error
}

var (
	configfilename = flag.String("config", "", "YAML configuration file")
	doBackup       = flag.Bool("backup", true, "Backup files on remote system")
	doDownload     = flag.Bool("download", false, "Download files from remote system")
	doUpload       = flag.Bool("upload", false, "Upload local files to remote system")
	debugging      = flag.Int("debug", 1, "Spew debugging info, e.g. output every file as it's touched. 0=silent, 1=progress, 2=verbose, 3=firehose")

	remoteFiles []string
)

func main() {
	flag.Parse()

	if *configfilename == "" {
		log.Fatalf("config flag indicating configuration file is required")
	}

	configfile, err := ioutil.ReadFile(*configfilename)
	if err != nil {
		log.Fatalf("could not read config file: %v", err)
	}

	config := configuration{}
	if err := yaml.Unmarshal(configfile, &config); err != nil {
		log.Fatalf("YAML error: %v", err)
	}

	upload := *doUpload
	backup := *doBackup
	download := *doDownload

	// start ssh client on tcp connection
	sftp, err := sftpclient.New(
		config.Connection.Hostname,
		config.Connection.Port,
		config.Connection.Username,
		config.Connection.Password)

	if err != nil {
		log.Fatal(err)
	}

	var bar *pb.ProgressBar
	if *debugging <= 1 {
		bar = pb.New(0)
	}

	if backup == true {
		files, err := sftp.FindAllRemoteFiles(config.Backup.From)
		if err != nil {
			log.Fatalln(err)
		}

		if *debugging <= 1 {
			bar.Total = 0
			bar.Prefix("Backing-up... ")
			bar.Start()
		}

		saved, errors := sftp.BackupFiles(config.Backup.To, files)

		go func() {
			for range saved {
				bar.Increment()
			}
		}()

		errorsEncountered := false
		for err := range errors {
			errorsEncountered = true
			log.Println(err.Err)
		}
		if errorsEncountered {
			log.Fatalln("Backup failed. Quitting.")
		}

		if *debugging <= 1 {
			bar.FinishPrint("Backup complete")
		}
	}

	if download == true {
		downloadedFile := make(chan bool)
		downloadDone := make(chan bool)
		errorsChannel := make(chan errorResponse)

		if *debugging <= 1 {
			bar.Total = 0
			bar.Prefix("Downloading... ")
			bar.Start()
		}

		for _, folder := range config.Download {
			go func(f sftpclient.FolderMapping) {
				defer func() {
					downloadDone <- true
				}()

				media, err := sftp.FindAllRemoteFiles([]string{f.Remote})
				if err != nil {
					errorsChannel <- errorResponse{err}
					return
				}
				atomic.AddInt64(&bar.Total, int64(len(media)))
				for _, remotefile := range media {
					trimmed := strings.TrimPrefix(remotefile, f.Remote)
					localfile := strings.Join([]string{f.Local, trimmed}, "/")
					if err := sftp.GetFile(localfile, remotefile); err != nil {
						errorsChannel <- errorResponse{err}
					}
					downloadedFile <- true
				}
			}(folder)
		}

		go func() {
			for range downloadedFile {
				<-downloadedFile
				bar.Increment()
			}
		}()

		go func() {
			for err := range errorsChannel {
				log.Println(err)
			}
		}()

		for range config.Download {
			<-downloadDone
		}
		close(downloadDone)
		close(downloadedFile)
		close(errorsChannel)

		if *debugging <= 1 {
			bar.FinishPrint("Download complete")
		}
	}

	if upload == true {
		if *debugging <= 1 {
			bar.Total = 0
			bar.Prefix("Uploading... ")
			bar.Start()
		}

		errorChannel, countChannel, copiedCountChannel := sftp.Upload(config.Upload)

		go func() {
			for range countChannel {
				atomic.AddInt64(&bar.Total, 1)
			}
		}()
		go func() {
			for range copiedCountChannel {
				bar.Increment()
			}
		}()

		for response := range errorChannel {
			log.Println(response.Err)
		}

		if *debugging <= 1 {
			bar.FinishPrint("Upload finished")
		}
	}
}
<|endoftext|>