<commit_before>package tchannel

// Copyright (c) 2015 Uber Technologies, Inc.

// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

import (
	"errors"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"sync"
	"time"

	"golang.org/x/net/context"
)

var (
	errAlreadyListening = errors.New("channel already listening")
)

const (
	ephemeralHostPort = "0.0.0.0:0"
)

// ChannelOptions are used to control parameters on a create a TChannel
type ChannelOptions struct {
	// Default Connection options
	DefaultConnectionOptions ConnectionOptions

	// The name of the process, for logging and reporting to peers
	ProcessName string

	// The logger to use for this channel
	Logger Logger

	// The reporter to use for reporting stats for this channel.
	StatsReporter StatsReporter

	// Trace reporter to use for this channel.
	TraceReporter TraceReporter
}

// ChannelState is the state of a channel.
type ChannelState int

const (
	// ChannelClient is a channel that can be used as a client.
	ChannelClient ChannelState = iota + 1

	// ChannelListening is a channel that is listening for new connnections.
	ChannelListening

	// ChannelStartClose is a channel that has received a Close request.
	// The channel is no longer listening, and all new incoming connections are rejected.
	ChannelStartClose

	// ChannelInboundClosed is a channel that has drained all incoming connections, but may
	// have outgoing connections. All incoming calls and new outgoing calls are rejected.
	ChannelInboundClosed

	// ChannelClosed is a channel that has closed completely.
	ChannelClosed
)

//go:generate stringer -type=ChannelState

// A Channel is a bi-directional connection to the peering and routing network.
// Applications can use a Channel to make service calls to remote peers via
// BeginCall, or to listen for incoming calls from peers.  Applications that
// want to receive requests should call one of Serve or ListenAndServe
// TODO(prashant): Shutdown all subchannels + peers when channel is closed.
type Channel struct {
	log               Logger
	commonStatsTags   map[string]string
	statsReporter     StatsReporter
	traceReporter     TraceReporter
	connectionOptions ConnectionOptions
	handlers          *handlerMap
	peers             *PeerList

	// mutable contains all the members of Channel which are mutable.
	mutable struct {
		mut         sync.RWMutex // protects members of the mutable struct.
		closed      bool
		peerInfo    LocalPeerInfo // May be ephemeral if this is a client only channel
		l           net.Listener  // May be nil if this is a client only channel
		subChannels map[string]*SubChannel
		conns       []*Connection
	}
}

// NewChannel creates a new Channel.  The new channel can be used to send outbound requests
// to peers, but will not listen or handling incoming requests until one of ListenAndServe
// or Serve is called. The local service name should be passed to serviceName.
func NewChannel(serviceName string, opts *ChannelOptions) (*Channel, error) {
	if opts == nil {
		opts = &ChannelOptions{}
	}

	logger := opts.Logger
	if logger == nil {
		logger = NullLogger
	}

	traceReporter := opts.TraceReporter
	if traceReporter == nil {
		traceReporter = NullReporter
	}

	processName := opts.ProcessName
	if processName == "" {
		processName = fmt.Sprintf("%s[%d]", filepath.Base(os.Args[0]), os.Getpid())
	}

	statsReporter := opts.StatsReporter
	if statsReporter == nil {
		statsReporter = NullStatsReporter
	}

	ch := &Channel{
		connectionOptions: opts.DefaultConnectionOptions,
		log:               logger,
		statsReporter:     statsReporter,
		traceReporter:     traceReporter,
		handlers:          &handlerMap{},
	}
	ch.mutable.peerInfo = LocalPeerInfo{
		PeerInfo: PeerInfo{
			ProcessName: processName,
			HostPort:    ephemeralHostPort,
		},
		ServiceName: serviceName,
	}
	ch.mutable.subChannels = make(map[string]*SubChannel)
	ch.peers = newPeerList(ch)
	ch.createCommonStats()
	return ch, nil
}

// Serve serves incoming requests using the provided listener.
// The local peer info is set synchronously, but the actual socket listening is done in
// a separate goroutine.
func (ch *Channel) Serve(l net.Listener) error {
	mutable := &ch.mutable
	mutable.mut.Lock()
	defer mutable.mut.Unlock()

	if mutable.l != nil {
		return errAlreadyListening
	}

	mutable.l = l
	mutable.peerInfo.HostPort = l.Addr().String()

	peerInfo := mutable.peerInfo
	ch.log.Debugf("%v (%v) listening on %v", peerInfo.ProcessName, peerInfo.ServiceName, peerInfo.HostPort)
	go ch.serve()
	return nil
}

// ListenAndServe listens on the given address and serves incoming requests.
// The port may be 0, in which case the channel will use an OS assigned port
// This method does not block as the handling of connections is done in a goroutine.
func (ch *Channel) ListenAndServe(hostPort string) error {
	mutable := &ch.mutable
	mutable.mut.RLock()

	if mutable.l != nil {
		mutable.mut.RUnlock()
		return errAlreadyListening
	}

	l, err := net.Listen("tcp", hostPort)
	if err != nil {
		mutable.mut.RUnlock()
		return err
	}

	mutable.mut.RUnlock()
	return ch.Serve(l)
}

// Register registers a handler for a service+operation pair
func (ch *Channel) Register(h Handler, operationName string) {
	ch.handlers.register(h, ch.PeerInfo().ServiceName, operationName)
}

// PeerInfo returns the current peer info for the channel
func (ch *Channel) PeerInfo() LocalPeerInfo {
	ch.mutable.mut.RLock()
	defer ch.mutable.mut.RUnlock()

	return ch.mutable.peerInfo
}

func (ch *Channel) createCommonStats() {
	ch.commonStatsTags = map[string]string{
		"app":     ch.mutable.peerInfo.ProcessName,
		"service": ch.mutable.peerInfo.ServiceName,
	}
	host, err := os.Hostname()
	if err != nil {
		ch.log.Infof("channel failed to get host: %v", err)
		return
	}
	ch.commonStatsTags["host"] = host
	// TODO(prashant): Allow user to pass extra tags (such as cluster, version).
}

func (ch *Channel) registerNewSubChannel(serviceName string) *SubChannel {
	mutable := &ch.mutable
	mutable.mut.Lock()
	defer mutable.mut.Unlock()

	// Recheck for the subchannel under the write lock.
	if sc, ok := mutable.subChannels[serviceName]; ok {
		return sc
	}

	sc := newSubChannel(serviceName, ch.peers)
	mutable.subChannels[serviceName] = sc
	return sc
}

// GetSubChannel returns a SubChannel for the given service name. If the subchannel does not
// exist, it is created.
func (ch *Channel) GetSubChannel(serviceName string) *SubChannel {
	mutable := &ch.mutable
	mutable.mut.RLock()

	if sc, ok := mutable.subChannels[serviceName]; ok {
		mutable.mut.RUnlock()
		return sc
	}

	mutable.mut.RUnlock()
	return ch.registerNewSubChannel(serviceName)
}

// Peers returns the PeerList for the channel.
func (ch *Channel) Peers() *PeerList {
	return ch.peers
}

// BeginCall starts a new call to a remote peer, returning an OutboundCall that can
// be used to write the arguments of the call.
func (ch *Channel) BeginCall(ctx context.Context, hostPort, serviceName, operationName string, callOptions *CallOptions) (*OutboundCall, error) {
	p := ch.peers.GetOrAdd(hostPort)
	return p.BeginCall(ctx, serviceName, operationName, callOptions)
}

// serve runs the listener to accept and manage new incoming connections, blocking
// until the channel is closed.
func (ch *Channel) serve() {
	acceptBackoff := 0 * time.Millisecond

	for {
		netConn, err := ch.mutable.l.Accept()
		if err != nil {
			// Backoff from new accepts if this is a temporary error
			if ne, ok := err.(net.Error); ok && ne.Temporary() {
				if acceptBackoff == 0 {
					acceptBackoff = 5 * time.Millisecond
				} else {
					acceptBackoff *= 2
				}
				if max := 1 * time.Second; acceptBackoff > max {
					acceptBackoff = max
				}
				ch.log.Warnf("accept error: %v; retrying in %v", err, acceptBackoff)
				time.Sleep(acceptBackoff)
				continue
			} else {
				// Only log an error if we are not shutdown.
				if ch.Closed() {
					return
				}
				ch.log.Fatalf("unrecoverable accept error: %v; closing server", err)
				return
			}
		}

		acceptBackoff = 0

		// Register the connection in the peer once the channel is set up.
		onActive := func(c *Connection) {
			ch.log.Debugf("Add connection as an active peer for %v", c.remotePeerInfo.HostPort)
			p := ch.peers.GetOrAdd(c.remotePeerInfo.HostPort)
			p.AddConnection(c)

			ch.mutable.mut.Lock()
			ch.mutable.conns = append(ch.mutable.conns, c)
			ch.mutable.mut.Unlock()
		}
		_, err = ch.newInboundConnection(netConn, onActive, &ch.connectionOptions)
		if err != nil {
			// Server is getting overloaded - begin rejecting new connections
			ch.log.Errorf("could not create new TChannelConnection for incoming conn: %v", err)
			netConn.Close()
			continue
		}
	}
}

// Ping sends a ping message to the given hostPort and waits for a response.
func (ch *Channel) Ping(ctx context.Context, hostPort string) error {
	peer := ch.Peers().GetOrAdd(hostPort)
	conn, err := peer.GetConnection(ctx)
	if err != nil {
		return err
	}

	return conn.ping(ctx)
}

// Logger returns the logger for this channel.
func (ch *Channel) Logger() Logger {
	return ch.log
}

// Connect connects the channel.
func (ch *Channel) Connect(ctx context.Context, hostPort string, connectionOptions *ConnectionOptions) (*Connection, error) {
	c, err := ch.newOutboundConnection(hostPort, connectionOptions)
	if err != nil {
		return nil, err
	}

	if err := c.sendInit(ctx); err != nil {
		return nil, err
	}

	ch.mutable.mut.Lock()
	ch.mutable.conns = append(ch.mutable.conns, c)
	ch.mutable.mut.Unlock()

	return c, err
}

// Closed returns whether this channel has been closed with .Close()
func (ch *Channel) Closed() bool {
	ch.mutable.mut.RLock()
	defer ch.mutable.mut.RUnlock()
	return ch.mutable.closed
}

// Close closes the channel including all connections to any active peers.
func (ch *Channel) Close() {
	ch.mutable.mut.Lock()
	defer ch.mutable.mut.Unlock()

	ch.mutable.closed = true
	if ch.mutable.l != nil {
		ch.mutable.l.Close()
	}
	ch.peers.Close()
}
<commit_msg>Use ChannelState instead of closed bool<commit_after>package tchannel

// Copyright (c) 2015 Uber Technologies, Inc.

// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

import (
	"errors"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"sync"
	"time"

	"golang.org/x/net/context"
)

var (
	errAlreadyListening  = errors.New("channel already listening")
	errInvalidStateForOp = errors.New("channel is in an invalid state for that operation")
)

const (
	ephemeralHostPort = "0.0.0.0:0"
)

// ChannelOptions are used to control parameters on a create a TChannel
type ChannelOptions struct {
	// Default Connection options
	DefaultConnectionOptions ConnectionOptions

	// The name of the process, for logging and reporting to peers
	ProcessName string

	// The logger to use for this channel
	Logger Logger

	// The reporter to use for reporting stats for this channel.
	StatsReporter StatsReporter

	// Trace reporter to use for this channel.
	TraceReporter TraceReporter
}

// ChannelState is the state of a channel.
type ChannelState int

const (
	// ChannelClient is a channel that can be used as a client.
	ChannelClient ChannelState = iota + 1

	// ChannelListening is a channel that is listening for new connnections.
	ChannelListening

	// ChannelStartClose is a channel that has received a Close request.
	// The channel is no longer listening, and all new incoming connections are rejected.
	ChannelStartClose

	// ChannelInboundClosed is a channel that has drained all incoming connections, but may
	// have outgoing connections. All incoming calls and new outgoing calls are rejected.
	ChannelInboundClosed

	// ChannelClosed is a channel that has closed completely.
	ChannelClosed
)

//go:generate stringer -type=ChannelState

// A Channel is a bi-directional connection to the peering and routing network.
// Applications can use a Channel to make service calls to remote peers via
// BeginCall, or to listen for incoming calls from peers.  Applications that
// want to receive requests should call one of Serve or ListenAndServe
// TODO(prashant): Shutdown all subchannels + peers when channel is closed.
type Channel struct {
	log               Logger
	commonStatsTags   map[string]string
	statsReporter     StatsReporter
	traceReporter     TraceReporter
	connectionOptions ConnectionOptions
	handlers          *handlerMap
	peers             *PeerList

	// mutable contains all the members of Channel which are mutable.
	mutable struct {
		mut         sync.RWMutex // protects members of the mutable struct.
		state       ChannelState
		peerInfo    LocalPeerInfo // May be ephemeral if this is a client only channel
		l           net.Listener  // May be nil if this is a client only channel
		subChannels map[string]*SubChannel
		conns       []*Connection
	}
}

// NewChannel creates a new Channel.  The new channel can be used to send outbound requests
// to peers, but will not listen or handling incoming requests until one of ListenAndServe
// or Serve is called. The local service name should be passed to serviceName.
func NewChannel(serviceName string, opts *ChannelOptions) (*Channel, error) {
	if opts == nil {
		opts = &ChannelOptions{}
	}

	logger := opts.Logger
	if logger == nil {
		logger = NullLogger
	}

	traceReporter := opts.TraceReporter
	if traceReporter == nil {
		traceReporter = NullReporter
	}

	processName := opts.ProcessName
	if processName == "" {
		processName = fmt.Sprintf("%s[%d]", filepath.Base(os.Args[0]), os.Getpid())
	}

	statsReporter := opts.StatsReporter
	if statsReporter == nil {
		statsReporter = NullStatsReporter
	}

	ch := &Channel{
		connectionOptions: opts.DefaultConnectionOptions,
		log:               logger,
		statsReporter:     statsReporter,
		traceReporter:     traceReporter,
		handlers:          &handlerMap{},
	}
	ch.mutable.peerInfo = LocalPeerInfo{
		PeerInfo: PeerInfo{
			ProcessName: processName,
			HostPort:    ephemeralHostPort,
		},
		ServiceName: serviceName,
	}
	ch.mutable.subChannels = make(map[string]*SubChannel)
	ch.mutable.state = ChannelClient
	ch.peers = newPeerList(ch)
	ch.createCommonStats()
	return ch, nil
}

// Serve serves incoming requests using the provided listener.
// The local peer info is set synchronously, but the actual socket listening is done in
// a separate goroutine.
func (ch *Channel) Serve(l net.Listener) error {
	mutable := &ch.mutable
	mutable.mut.Lock()
	defer mutable.mut.Unlock()

	if mutable.l != nil {
		return errAlreadyListening
	}
	mutable.l = l

	if mutable.state != ChannelClient {
		return errInvalidStateForOp
	}
	mutable.state = ChannelListening

	mutable.peerInfo.HostPort = l.Addr().String()
	peerInfo := mutable.peerInfo
	ch.log.Debugf("%v (%v) listening on %v", peerInfo.ProcessName, peerInfo.ServiceName, peerInfo.HostPort)
	go ch.serve()
	return nil
}

// ListenAndServe listens on the given address and serves incoming requests.
// The port may be 0, in which case the channel will use an OS assigned port
// This method does not block as the handling of connections is done in a goroutine.
func (ch *Channel) ListenAndServe(hostPort string) error {
	mutable := &ch.mutable
	mutable.mut.RLock()

	if mutable.l != nil {
		mutable.mut.RUnlock()
		return errAlreadyListening
	}

	l, err := net.Listen("tcp", hostPort)
	if err != nil {
		mutable.mut.RUnlock()
		return err
	}

	mutable.mut.RUnlock()
	return ch.Serve(l)
}

// Register registers a handler for a service+operation pair
func (ch *Channel) Register(h Handler, operationName string) {
	ch.handlers.register(h, ch.PeerInfo().ServiceName, operationName)
}

// PeerInfo returns the current peer info for the channel
func (ch *Channel) PeerInfo() LocalPeerInfo {
	ch.mutable.mut.RLock()
	defer ch.mutable.mut.RUnlock()

	return ch.mutable.peerInfo
}

func (ch *Channel) createCommonStats() {
	ch.commonStatsTags = map[string]string{
		"app":     ch.mutable.peerInfo.ProcessName,
		"service": ch.mutable.peerInfo.ServiceName,
	}
	host, err := os.Hostname()
	if err != nil {
		ch.log.Infof("channel failed to get host: %v", err)
		return
	}
	ch.commonStatsTags["host"] = host
	// TODO(prashant): Allow user to pass extra tags (such as cluster, version).
}

func (ch *Channel) registerNewSubChannel(serviceName string) *SubChannel {
	mutable := &ch.mutable
	mutable.mut.Lock()
	defer mutable.mut.Unlock()

	// Recheck for the subchannel under the write lock.
	if sc, ok := mutable.subChannels[serviceName]; ok {
		return sc
	}

	sc := newSubChannel(serviceName, ch.peers)
	mutable.subChannels[serviceName] = sc
	return sc
}

// GetSubChannel returns a SubChannel for the given service name. If the subchannel does not
// exist, it is created.
func (ch *Channel) GetSubChannel(serviceName string) *SubChannel {
	mutable := &ch.mutable
	mutable.mut.RLock()

	if sc, ok := mutable.subChannels[serviceName]; ok {
		mutable.mut.RUnlock()
		return sc
	}

	mutable.mut.RUnlock()
	return ch.registerNewSubChannel(serviceName)
}

// Peers returns the PeerList for the channel.
func (ch *Channel) Peers() *PeerList {
	return ch.peers
}

// BeginCall starts a new call to a remote peer, returning an OutboundCall that can
// be used to write the arguments of the call.
func (ch *Channel) BeginCall(ctx context.Context, hostPort, serviceName, operationName string, callOptions *CallOptions) (*OutboundCall, error) {
	p := ch.peers.GetOrAdd(hostPort)
	return p.BeginCall(ctx, serviceName, operationName, callOptions)
}

// serve runs the listener to accept and manage new incoming connections, blocking
// until the channel is closed.
func (ch *Channel) serve() {
	acceptBackoff := 0 * time.Millisecond

	for {
		netConn, err := ch.mutable.l.Accept()
		if err != nil {
			// Backoff from new accepts if this is a temporary error
			if ne, ok := err.(net.Error); ok && ne.Temporary() {
				if acceptBackoff == 0 {
					acceptBackoff = 5 * time.Millisecond
				} else {
					acceptBackoff *= 2
				}
				if max := 1 * time.Second; acceptBackoff > max {
					acceptBackoff = max
				}
				ch.log.Warnf("accept error: %v; retrying in %v", err, acceptBackoff)
				time.Sleep(acceptBackoff)
				continue
			} else {
				// Only log an error if this didn't happen due to a Close.
				if ch.State() >= ChannelStartClose {
					return
				}
				ch.log.Fatalf("unrecoverable accept error: %v; closing server", err)
				return
			}
		}

		acceptBackoff = 0

		// Register the connection in the peer once the channel is set up.
		onActive := func(c *Connection) {
			ch.log.Debugf("Add connection as an active peer for %v", c.remotePeerInfo.HostPort)
			p := ch.peers.GetOrAdd(c.remotePeerInfo.HostPort)
			p.AddConnection(c)

			ch.mutable.mut.Lock()
			ch.mutable.conns = append(ch.mutable.conns, c)
			ch.mutable.mut.Unlock()
		}
		_, err = ch.newInboundConnection(netConn, onActive, &ch.connectionOptions)
		if err != nil {
			// Server is getting overloaded - begin rejecting new connections
			ch.log.Errorf("could not create new TChannelConnection for incoming conn: %v", err)
			netConn.Close()
			continue
		}
	}
}

// Ping sends a ping message to the given hostPort and waits for a response.
func (ch *Channel) Ping(ctx context.Context, hostPort string) error {
	peer := ch.Peers().GetOrAdd(hostPort)
	conn, err := peer.GetConnection(ctx)
	if err != nil {
		return err
	}

	return conn.ping(ctx)
}

// Logger returns the logger for this channel.
func (ch *Channel) Logger() Logger {
	return ch.log
}

// Connect connects the channel.
func (ch *Channel) Connect(ctx context.Context, hostPort string, connectionOptions *ConnectionOptions) (*Connection, error) {
	c, err := ch.newOutboundConnection(hostPort, connectionOptions)
	if err != nil {
		return nil, err
	}

	if err := c.sendInit(ctx); err != nil {
		return nil, err
	}

	ch.mutable.mut.Lock()
	ch.mutable.conns = append(ch.mutable.conns, c)
	ch.mutable.mut.Unlock()

	return c, err
}

// Closed returns whether this channel has been closed with .Close()
func (ch *Channel) Closed() bool {
	return ch.State() == ChannelClosed
}

// State returns the current channel state.
func (ch *Channel) State() ChannelState {
	ch.mutable.mut.RLock()
	state := ch.mutable.state
	ch.mutable.mut.RUnlock()

	return state
}

// Close closes the channel including all connections to any active peers.
func (ch *Channel) Close() {
	ch.mutable.mut.Lock()

	if ch.mutable.l != nil {
		ch.mutable.l.Close()
	}
	ch.mutable.state = ChannelStartClose
	ch.mutable.mut.Unlock()

	ch.peers.Close()
}
<|endoftext|>


<commit_before>package ssh

import (
	"bytes"
	"code.google.com/p/go.crypto/ssh"
	"errors"
	"fmt"
	"github.com/mitchellh/packer/packer"
	"io"
	"log"
	"net"
	"path/filepath"
)

type comm struct {
	client *ssh.ClientConn
	config *Config
	conn   net.Conn
}

// Config is the structure used to configure the SSH communicator.
type Config struct {
	// The configuration of the Go SSH connection
	SSHConfig *ssh.ClientConfig

	// Connection returns a new connection. The current connection
	// in use will be closed as part of the Close method, or in the
	// case an error occurs.
	Connection func() (net.Conn, error)
}

// Creates a new packer.Communicator implementation over SSH. This takes
// an already existing TCP connection and SSH configuration.
func New(config *Config) (result *comm, err error) {
	// Establish an initial connection and connect
	result = &comm{
		config: config,
	}

	if err = result.reconnect(); err != nil {
		result = nil
		return
	}

	return
}

func (c *comm) Start(cmd *packer.RemoteCmd) (err error) {
	session, err := c.newSession()
	if err != nil {
		return
	}

	// Setup our session
	session.Stdin = cmd.Stdin
	session.Stdout = cmd.Stdout
	session.Stderr = cmd.Stderr

	// Request a PTY
	termModes := ssh.TerminalModes{
		ssh.ECHO:          0,     // do not echo
		ssh.TTY_OP_ISPEED: 14400, // input speed = 14.4kbaud
		ssh.TTY_OP_OSPEED: 14400, // output speed = 14.4kbaud
	}

	if err = session.RequestPty("xterm", 80, 40, termModes); err != nil {
		return
	}

	log.Printf("starting remote command: %s", cmd.Command)
	err = session.Start(cmd.Command + "\n")
	if err != nil {
		return
	}

	// Start a goroutine to wait for the session to end and set the
	// exit boolean and status.
	go func() {
		defer session.Close()
		err := session.Wait()
		cmd.ExitStatus = 0
		if err != nil {
			exitErr, ok := err.(*ssh.ExitError)
			if ok {
				cmd.ExitStatus = exitErr.ExitStatus()
			}
		}

		cmd.Exited = true
	}()

	return
}

func (c *comm) Upload(path string, input io.Reader) error {
	session, err := c.newSession()
	if err != nil {
		return err
	}

	defer session.Close()

	// Get a pipe to stdin so that we can send data down
	w, err := session.StdinPipe()
	if err != nil {
		return err
	}

	// Set stderr/stdout to a bytes buffer
	stderr := new(bytes.Buffer)
	stdout := new(bytes.Buffer)
	session.Stderr = stderr
	session.Stdout = stdout

	// We only want to close once, so we nil w after we close it,
	// and only close in the defer if it hasn't been closed already.
	defer func() {
		if w != nil {
			w.Close()
		}
	}()

	// The target directory and file for talking the SCP protocol
	target_dir := filepath.Dir(path)
	target_file := filepath.Base(path)

	// Start the sink mode on the other side
	// TODO(mitchellh): There are probably issues with shell escaping the path
	log.Println("Starting remote scp process in sink mode")
	if err = session.Start("scp -vt " + target_dir); err != nil {
		return err
	}

	// Determine the length of the upload content by copying it
	// into an in-memory buffer. Note that this means what we upload
	// must fit into memory.
	log.Println("Copying input data into in-memory buffer so we can get the length")
	input_memory := new(bytes.Buffer)
	if _, err = io.Copy(input_memory, input); err != nil {
		return err
	}

	// Start the protocol
	log.Println("Beginning file upload...")
	fmt.Fprintln(w, "C0644", input_memory.Len(), target_file)
	io.Copy(w, input_memory)
	fmt.Fprint(w, "\x00")

	// TODO(mitchellh): Each step above results in a 0/1/2 being sent by
	// the remote side to confirm. We should check for those confirmations.

	// Close the stdin, which sends an EOF, and then set w to nil so that
	// our defer func doesn't close it again since that is unsafe with
	// the Go SSH package.
	log.Println("Upload complete, closing stdin pipe")
	w.Close()
	w = nil

	// Wait for the SCP connection to close, meaning it has consumed all
	// our data and has completed. Or has errored.
	log.Println("Waiting for SSH session to complete")
	err = session.Wait()
	if err != nil {
		if exitErr, ok := err.(*ssh.ExitError); ok {
			// Otherwise, we have an ExitErorr, meaning we can just read
			// the exit status
			log.Printf("non-zero exit status: %d", exitErr.ExitStatus())

			// If we exited with status 127, it means SCP isn't available.
			// Return a more descriptive error for that.
			if exitErr.ExitStatus() == 127 {
				return errors.New(
					"SCP failed to start. This usually means that SCP is not\n" +
						"properly installed on the remote system.")
			}
		}

		return err
	}

	log.Printf("scp stdout (length %d): %#v", stdout.Len(), stdout.Bytes())
	log.Printf("scp stderr (length %d): %s", stderr.Len(), stderr.String())

	return nil
}

func (c *comm) Download(string, io.Writer) error {
	panic("not implemented yet")
}

func (c *comm) newSession() (*ssh.Session, error) {
	log.Println("opening new ssh session")
	session, err := c.client.NewSession()
	if err != nil {
		log.Printf("ssh session open error: '%s', attempting reconnect", err)
		if err := c.reconnect(); err != nil {
			return nil, err
		}

		return c.client.NewSession()
	}

	return session, nil
}

func (c *comm) reconnect() (err error) {
	if c.conn != nil {
		c.conn.Close()
	}

	log.Printf("reconnecting to TCP connection for SSH")
	c.conn, err = c.config.Connection()
	if err != nil {
		log.Printf("reconnection error: %s", err)
		return
	}

	log.Printf("handshaking with SSH")
	c.client, err = ssh.Client(c.conn, c.config.SSHConfig)
	if err != nil {
		log.Printf("handshake error: %s", err)
	}

	return
}
<commit_msg>Return errors from remote scp command<commit_after>package ssh

import (
	"bufio"
	"bytes"
	"code.google.com/p/go.crypto/ssh"
	"errors"
	"fmt"
	"github.com/mitchellh/packer/packer"
	"io"
	"log"
	"net"
	"path/filepath"
)

type comm struct {
	client *ssh.ClientConn
	config *Config
	conn   net.Conn
}

// Config is the structure used to configure the SSH communicator.
type Config struct {
	// The configuration of the Go SSH connection
	SSHConfig *ssh.ClientConfig

	// Connection returns a new connection. The current connection
	// in use will be closed as part of the Close method, or in the
	// case an error occurs.
	Connection func() (net.Conn, error)
}

// Creates a new packer.Communicator implementation over SSH. This takes
// an already existing TCP connection and SSH configuration.
func New(config *Config) (result *comm, err error) {
	// Establish an initial connection and connect
	result = &comm{
		config: config,
	}

	if err = result.reconnect(); err != nil {
		result = nil
		return
	}

	return
}

func (c *comm) Start(cmd *packer.RemoteCmd) (err error) {
	session, err := c.newSession()
	if err != nil {
		return
	}

	// Setup our session
	session.Stdin = cmd.Stdin
	session.Stdout = cmd.Stdout
	session.Stderr = cmd.Stderr

	// Request a PTY
	termModes := ssh.TerminalModes{
		ssh.ECHO:          0,     // do not echo
		ssh.TTY_OP_ISPEED: 14400, // input speed = 14.4kbaud
		ssh.TTY_OP_OSPEED: 14400, // output speed = 14.4kbaud
	}

	if err = session.RequestPty("xterm", 80, 40, termModes); err != nil {
		return
	}

	log.Printf("starting remote command: %s", cmd.Command)
	err = session.Start(cmd.Command + "\n")
	if err != nil {
		return
	}

	// Start a goroutine to wait for the session to end and set the
	// exit boolean and status.
	go func() {
		defer session.Close()
		err := session.Wait()
		cmd.ExitStatus = 0
		if err != nil {
			exitErr, ok := err.(*ssh.ExitError)
			if ok {
				cmd.ExitStatus = exitErr.ExitStatus()
			}
		}

		cmd.Exited = true
	}()

	return
}

func (c *comm) Upload(path string, input io.Reader) error {
	session, err := c.newSession()
	if err != nil {
		return err
	}

	defer session.Close()

	// Get a pipe to stdin so that we can send data down
	w, err := session.StdinPipe()
	if err != nil {
		return err
	}

	// We only want to close once, so we nil w after we close it,
	// and only close in the defer if it hasn't been closed already.
	defer func() {
		if w != nil {
			w.Close()
		}
	}()

	// Get a pipe to stdout so that we can get responses back
	scp_reader, err := session.StdoutPipe()
	if err != nil {
		return err
	}
	r := bufio.NewReader(scp_reader)

	// Set stderr to a bytes buffer
	stderr := new(bytes.Buffer)
	session.Stderr = stderr

	// The target directory and file for talking the SCP protocol
	target_dir := filepath.Dir(path)
	target_file := filepath.Base(path)

	// Start the sink mode on the other side
	// TODO(mitchellh): There are probably issues with shell escaping the path
	log.Println("Starting remote scp process in sink mode")
	if err = session.Start("scp -vt " + target_dir); err != nil {
		return err
	}

	// Determine the length of the upload content by copying it
	// into an in-memory buffer. Note that this means what we upload
	// must fit into memory.
	log.Println("Copying input data into in-memory buffer so we can get the length")
	input_memory := new(bytes.Buffer)
	if _, err = io.Copy(input_memory, input); err != nil {
		return err
	}

	// Start the protocol
	log.Println("Beginning file upload...")
	fmt.Fprintln(w, "C0644", input_memory.Len(), target_file)
	err = check_response(r)
	if err != nil {
		return err
	}

	io.Copy(w, input_memory)
	fmt.Fprint(w, "\x00")
	err = check_response(r)
	if err != nil {
		return err
	}

	// TODO(mitchellh): Each step above results in a 0/1/2 being sent by
	// the remote side to confirm. We should check for those confirmations.

	// Close the stdin, which sends an EOF, and then set w to nil so that
	// our defer func doesn't close it again since that is unsafe with
	// the Go SSH package.
	log.Println("Upload complete, closing stdin pipe")
	w.Close()
	w = nil

	// Wait for the SCP connection to close, meaning it has consumed all
	// our data and has completed. Or has errored.
	log.Println("Waiting for SSH session to complete")
	err = session.Wait()
	if err != nil {
		if exitErr, ok := err.(*ssh.ExitError); ok {
			// Otherwise, we have an ExitErorr, meaning we can just read
			// the exit status
			log.Printf("non-zero exit status: %d", exitErr.ExitStatus())

			// If we exited with status 127, it means SCP isn't available.
			// Return a more descriptive error for that.
			if exitErr.ExitStatus() == 127 {
				return errors.New(
					"SCP failed to start. This usually means that SCP is not\n" +
						"properly installed on the remote system.")
			}
		}

		return err
	}

	log.Printf("scp stderr (length %d): %s", stderr.Len(), stderr.String())

	return nil
}

func (c *comm) Download(string, io.Writer) error {
	panic("not implemented yet")
}

func (c *comm) newSession() (*ssh.Session, error) {
	log.Println("opening new ssh session")
	session, err := c.client.NewSession()
	if err != nil {
		log.Printf("ssh session open error: '%s', attempting reconnect", err)
		if err := c.reconnect(); err != nil {
			return nil, err
		}

		return c.client.NewSession()
	}

	return session, nil
}

func (c *comm) reconnect() (err error) {
	if c.conn != nil {
		c.conn.Close()
	}

	log.Printf("reconnecting to TCP connection for SSH")
	c.conn, err = c.config.Connection()
	if err != nil {
		log.Printf("reconnection error: %s", err)
		return
	}

	log.Printf("handshaking with SSH")
	c.client, err = ssh.Client(c.conn, c.config.SSHConfig)
	if err != nil {
		log.Printf("handshake error: %s", err)
	}

	return
}

func check_response(r *bufio.Reader) (err error) {
	scp_status_code, err := r.ReadByte()
	if err != nil {
		return err
	}
	if scp_status_code != 0 {
		// Treat any non-zero (really 1 and 2) as fatal errors
		error_message, _, err := r.ReadLine()
		err = fmt.Errorf(string(error_message[:]))
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package benchmarkbbs_test

import (
	"context"
	"crypto/rand"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"runtime"
	"strings"
	"testing"
	"time"

	"code.cloudfoundry.org/bbs"
	"code.cloudfoundry.org/bbs/db"
	"code.cloudfoundry.org/bbs/db/sqldb"
	"code.cloudfoundry.org/bbs/db/sqldb/helpers"
	"code.cloudfoundry.org/bbs/db/sqldb/helpers/monitor"
	"code.cloudfoundry.org/bbs/encryption"
	"code.cloudfoundry.org/bbs/guidprovider"
	"code.cloudfoundry.org/bbs/models"
	benchmarkconfig "code.cloudfoundry.org/benchmarkbbs/config"
	"code.cloudfoundry.org/benchmarkbbs/generator"
	"code.cloudfoundry.org/benchmarkbbs/reporter"
	"code.cloudfoundry.org/clock"
	fakes "code.cloudfoundry.org/diego-logging-client/testhelpers"
	"code.cloudfoundry.org/executor"
	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/locket"
	locketmodels "code.cloudfoundry.org/locket/models"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	awssession "github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3/s3manager"
	_ "github.com/go-sql-driver/mysql"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/zorkian/go-datadog-api"
)

const LongTermTTL = int64(864000) // 10 days

var (
	expectedLRPCount int

	expectedActualLRPCounts map[string]int
	config                  benchmarkconfig.BenchmarkBBSConfig

	logger          lager.Logger
	locketClient    locketmodels.LocketClient
	sqlDB           *sqldb.SQLDB
	activeDB        db.DB
	bbsClient       bbs.InternalClient
	dataDogClient   *datadog.Client
	dataDogReporter reporter.DataDogReporter
	reporters       []Reporter
)

const (
	DEBUG = "debug"
	INFO  = "info"
	ERROR = "error"
	FATAL = "fatal"
)

func init() {
	runtime.GOMAXPROCS(runtime.NumCPU())

	configPath, ok := os.LookupEnv("BENCHMARK_BBS_CONFIG_PATH")
	if !ok {
		log.Fatal("error reading config file: BENCHMARK_BBS_CONFIG_PATH env var not set")
	}

	var err error
	config, err = benchmarkconfig.NewBenchmarkBBSConfig(configPath)
	if err != nil {
		panic(err)
	}

	if config.BBSAddress == "" {
		log.Fatal("bbsAddress is required")
	}

	BenchmarkTests(config.NumReps, config.NumTrials, config.LocalRouteEmitters)
}

func TestBenchmarkBbs(t *testing.T) {
	var lagerLogLevel lager.LogLevel
	switch config.LogLevel {
	case DEBUG:
		lagerLogLevel = lager.DEBUG
	case INFO:
		lagerLogLevel = lager.INFO
	case ERROR:
		lagerLogLevel = lager.ERROR
	case FATAL:
		lagerLogLevel = lager.FATAL
	default:
		panic(fmt.Errorf("unknown log level: %s", config.LogLevel))
	}

	var logWriter io.Writer
	if config.LogFilename == "" {
		logWriter = GinkgoWriter
	} else {
		logFile, err := os.OpenFile(config.LogFilename, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
		if err != nil {
			panic("Error opening file, err: " + err.Error())
		}
		defer logFile.Close()

		logWriter = logFile
	}

	logger = lager.NewLogger("bbs-benchmarks-test")
	logger.RegisterSink(lager.NewWriterSink(logWriter, lagerLogLevel))

	reporters = []Reporter{}

	if config.DataDogAPIKey != "" && config.DataDogAppKey != "" {
		dataDogClient = datadog.NewClient(config.DataDogAPIKey, config.DataDogAppKey)
		dataDogReporter = reporter.NewDataDogReporter(logger, config.MetricPrefix, dataDogClient)
		reporters = append(reporters, &dataDogReporter)
	}

	if config.AwsAccessKeyID != "" && config.AwsSecretAccessKey != "" && config.AwsBucketName != "" {
		uploader := s3manager.NewUploader(awssession.New(&aws.Config{
			Region:      &config.AwsRegion,
			Credentials: credentials.NewStaticCredentials(config.AwsAccessKeyID, config.AwsSecretAccessKey, ""),
		}))
		reporter := reporter.NewS3Reporter(logger, config.AwsBucketName, uploader)
		reporters = append(reporters, &reporter)
	}

	RegisterFailHandler(Fail)
	RunSpecsWithDefaultAndCustomReporters(t, "Benchmark BBS Suite", reporters)
}

func initializeActiveDB() *sql.DB {
	if activeDB != nil {
		return nil
	}

	if config.DatabaseConnectionString == "" {
		logger.Fatal("no-sql-configuration", errors.New("no-sql-configuration"))
	}

	if config.DatabaseDriver == "postgres" && !strings.Contains(config.DatabaseConnectionString, "sslmode") {
		config.DatabaseConnectionString = fmt.Sprintf("%s?sslmode=disable", config.DatabaseConnectionString)
	}

	sqlConn, err := sql.Open(config.DatabaseDriver, config.DatabaseConnectionString)
	if err != nil {
		logger.Fatal("failed-to-open-sql", err)
	}
	sqlConn.SetMaxOpenConns(1)
	sqlConn.SetMaxIdleConns(1)

	err = sqlConn.Ping()
	Expect(err).NotTo(HaveOccurred())

	wrappedDB := helpers.NewMonitoredDB(sqlConn, monitor.New())
	sqlDB = initializeSQLDB(logger, wrappedDB)
	activeDB = sqlDB
	return sqlConn
}

var _ = BeforeSuite(func() {
	bbsClient = initializeBBSClient(logger, time.Duration(config.BBSClientHTTPTimeout))

	conn := initializeActiveDB()
	initializeLocketClient()

	resp, err := locketClient.FetchAll(context.Background(), &locketmodels.FetchAllRequest{TypeCode: locketmodels.PRESENCE})
	Expect(err).NotTo(HaveOccurred())

	cells := make(map[string]struct{})
	for i := 0; i < config.NumReps; i++ {
		cells[fmt.Sprintf("cell-%d", i)] = struct{}{}
	}

	if config.ReseedDatabase || config.NumReps != len(resp.Resources) {
		var err error

		cleanupSQLDB(conn)
		resetLocketDB(locketClient, cells, resp)

		if config.DesiredLRPs > 0 {
			desiredLRPGenerator := generator.NewDesiredLRPGenerator(
				config.ErrorTolerance,
				config.MetricPrefix,
				config.NumPopulateWorkers,
				bbsClient,
				dataDogClient,
			)

			expectedLRPCount, expectedActualLRPCounts, err = desiredLRPGenerator.Generate(logger, config.NumReps, config.DesiredLRPs)
			Expect(err).NotTo(HaveOccurred())
		}
	} else {
		resetUnclaimedActualLRPs(conn, cells)

		query := `
			SELECT
				COUNT(*)
			FROM desired_lrps
		`
		res := conn.QueryRow(query)
		err := res.Scan(&expectedLRPCount)
		Expect(err).NotTo(HaveOccurred())

		expectedActualLRPCounts = make(map[string]int)
		query = `
			SELECT
				cell_id, COUNT(*)
			FROM actual_lrps
			GROUP BY cell_id
		`
		rows, err := conn.Query(query)
		Expect(err).NotTo(HaveOccurred())

		defer rows.Close()
		for rows.Next() {
			var cellId string
			var count int
			err = rows.Scan(&cellId, &count)
			Expect(err).NotTo(HaveOccurred())

			expectedActualLRPCounts[cellId] = count
		}
	}

	if float64(expectedLRPCount) < float64(config.DesiredLRPs)*config.ErrorTolerance {
		Fail(fmt.Sprintf("Error rate of %.3f for actuals exceeds tolerance of %.3f", float64(expectedLRPCount)/float64(config.DesiredLRPs), config.ErrorTolerance))
	}
})

func initializeSQLDB(logger lager.Logger, sqlConn helpers.QueryableDB) *sqldb.SQLDB {
	key, keys, err := config.EncryptionConfig.Parse()
	if err != nil {
		logger.Fatal("cannot-setup-encryption", err)
	}
	keyManager, err := encryption.NewKeyManager(key, keys)
	if err != nil {
		logger.Fatal("cannot-setup-encryption", err)
	}
	cryptor := encryption.NewCryptor(keyManager, rand.Reader)

	return sqldb.NewSQLDB(
		sqlConn,
		1000,
		1000,
		cryptor,
		guidprovider.DefaultGuidProvider,
		clock.NewClock(),
		config.DatabaseDriver,
		&fakes.FakeIngressClient{},
	)
}

func initializeLocketClient() {
	var err error
	locketClient, err = locket.NewClient(logger, config.ClientLocketConfig)
	Expect(err).NotTo(HaveOccurred())
}

func initializeBBSClient(logger lager.Logger, bbsClientHTTPTimeout time.Duration) bbs.InternalClient {
	bbsClient, err := bbs.NewClientWithConfig(bbs.ClientConfig{
		URL:                    config.BBSAddress,
		IsTLS:                  true,
		CAFile:                 config.BBSCACert,
		CertFile:               config.BBSClientCert,
		KeyFile:                config.BBSClientKey,
		ClientSessionCacheSize: 1,
		MaxIdleConnsPerHost:    25000,
		RequestTimeout:         bbsClientHTTPTimeout,
	})
	if err != nil {
		logger.Fatal("Failed to configure secure BBS client", err)
	}
	return bbsClient
}

func cleanupSQLDB(conn *sql.DB) {
	_, err := conn.Exec("TRUNCATE actual_lrps")
	Expect(err).NotTo(HaveOccurred())
	_, err = conn.Exec("TRUNCATE desired_lrps")
	Expect(err).NotTo(HaveOccurred())
}

func resetLocketDB(locketClient locketmodels.LocketClient, cells map[string]struct{}, resp *locketmodels.FetchAllResponse) {
	for _, resource := range resp.Resources {
		if _, ok := cells[resource.Owner]; !ok {
			_, err := locketClient.Release(context.Background(), &locketmodels.ReleaseRequest{Resource: resource})
			Expect(err).NotTo(HaveOccurred())
		}
	}

	for cellId, _ := range cells {
		_, err := locketClient.Lock(context.Background(), &locketmodels.LockRequest{Resource: lockResource(cellId), TtlInSeconds: LongTermTTL})
		Expect(err).NotTo(HaveOccurred())
	}
}

func lockResource(cellID string) *locketmodels.Resource {
	resources := executor.ExecutorResources{
		MemoryMB:   1000,
		DiskMB:     2000,
		Containers: 300,
	}
	cellCapacity := models.NewCellCapacity(int32(resources.MemoryMB), int32(resources.DiskMB), int32(resources.Containers))
	cellPresence := models.NewCellPresence(cellID, cellID+".address", cellID+".url",
		"z1", cellCapacity, []string{"providers"},
		[]string{"cflinuxfs9"}, []string{}, []string{})

	payload, err := json.Marshal(cellPresence)
	Expect(err).NotTo(HaveOccurred())

	return &locketmodels.Resource{
		Key:   cellID,
		Owner: cellID,
		Value: string(payload),
		Type:  locketmodels.PresenceType,
	}
}

func resetUnclaimedActualLRPs(conn *sql.DB, cells map[string]struct{}) {
	existingCellIds := make(map[string]struct{})

	query := `
			SELECT
				cell_id
			FROM actual_lrps
			GROUP BY cell_id
		`
	rows, err := conn.Query(query)
	Expect(err).NotTo(HaveOccurred())

	defer rows.Close()
	for rows.Next() {
		var cellId string
		err = rows.Scan(&cellId)
		Expect(err).NotTo(HaveOccurred())

		existingCellIds[cellId] = struct{}{}
	}

	missingCells := findMissingCells(existingCellIds, cells)

	query = `
				SELECT
					process_guid, domain
				FROM actual_lrps WHERE cell_id = ''
			`

	type guidAndDomain struct {
		ProcessGuid string
		Domain      string
	}
	var lrps []guidAndDomain

	rows, err = conn.Query(query)
	Expect(err).NotTo(HaveOccurred())

	defer rows.Close()
	for rows.Next() {
		var guid, domain string
		err = rows.Scan(&guid, &domain)
		Expect(err).NotTo(HaveOccurred())
		lrps = append(lrps, guidAndDomain{ProcessGuid: guid, Domain: domain})
	}

	for i, lrp := range lrps {
		cellID := missingCells[i%len(missingCells)]
		actualLRPInstanceKey := &models.ActualLRPInstanceKey{InstanceGuid: lrp.ProcessGuid + "-i", CellId: cellID}
		netInfo := models.NewActualLRPNetInfo("1.2.3.4", "2.2.2.2", models.ActualLRPNetInfo_PreferredAddressHost, models.NewPortMapping(61999, 8080))
		err := bbsClient.StartActualLRP(logger, &models.ActualLRPKey{Domain: lrp.Domain, ProcessGuid: lrp.ProcessGuid, Index: 0}, actualLRPInstanceKey, &netInfo)
		Expect(err).NotTo(HaveOccurred())
	}
}

func findMissingCells(existing map[string]struct{}, expected map[string]struct{}) []string {
	var missing []string
	for cell, _ := range expected {
		if _, ok := existing[cell]; !ok {
			missing = append(missing, cell)
		}
	}

	return missing
}
<commit_msg>Use jackc/pgx instead of lib/pq<commit_after>package benchmarkbbs_test

import (
	"context"
	"crypto/rand"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"runtime"
	"strings"
	"testing"
	"time"

	"code.cloudfoundry.org/bbs"
	"code.cloudfoundry.org/bbs/db"
	"code.cloudfoundry.org/bbs/db/sqldb"
	"code.cloudfoundry.org/bbs/db/sqldb/helpers"
	"code.cloudfoundry.org/bbs/db/sqldb/helpers/monitor"
	"code.cloudfoundry.org/bbs/encryption"
	"code.cloudfoundry.org/bbs/guidprovider"
	"code.cloudfoundry.org/bbs/models"
	benchmarkconfig "code.cloudfoundry.org/benchmarkbbs/config"
	"code.cloudfoundry.org/benchmarkbbs/generator"
	"code.cloudfoundry.org/benchmarkbbs/reporter"
	"code.cloudfoundry.org/clock"
	fakes "code.cloudfoundry.org/diego-logging-client/testhelpers"
	"code.cloudfoundry.org/executor"
	"code.cloudfoundry.org/lager"
	"code.cloudfoundry.org/locket"
	locketmodels "code.cloudfoundry.org/locket/models"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	awssession "github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3/s3manager"
	_ "github.com/go-sql-driver/mysql"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/zorkian/go-datadog-api"
)

const LongTermTTL = int64(864000) // 10 days

var (
	expectedLRPCount int

	expectedActualLRPCounts map[string]int
	config                  benchmarkconfig.BenchmarkBBSConfig

	logger          lager.Logger
	locketClient    locketmodels.LocketClient
	sqlDB           *sqldb.SQLDB
	activeDB        db.DB
	bbsClient       bbs.InternalClient
	dataDogClient   *datadog.Client
	dataDogReporter reporter.DataDogReporter
	reporters       []Reporter
)

const (
	DEBUG = "debug"
	INFO  = "info"
	ERROR = "error"
	FATAL = "fatal"
)

func init() {
	runtime.GOMAXPROCS(runtime.NumCPU())

	configPath, ok := os.LookupEnv("BENCHMARK_BBS_CONFIG_PATH")
	if !ok {
		log.Fatal("error reading config file: BENCHMARK_BBS_CONFIG_PATH env var not set")
	}

	var err error
	config, err = benchmarkconfig.NewBenchmarkBBSConfig(configPath)
	if err != nil {
		panic(err)
	}

	if config.BBSAddress == "" {
		log.Fatal("bbsAddress is required")
	}

	BenchmarkTests(config.NumReps, config.NumTrials, config.LocalRouteEmitters)
}

func TestBenchmarkBbs(t *testing.T) {
	var lagerLogLevel lager.LogLevel
	switch config.LogLevel {
	case DEBUG:
		lagerLogLevel = lager.DEBUG
	case INFO:
		lagerLogLevel = lager.INFO
	case ERROR:
		lagerLogLevel = lager.ERROR
	case FATAL:
		lagerLogLevel = lager.FATAL
	default:
		panic(fmt.Errorf("unknown log level: %s", config.LogLevel))
	}

	var logWriter io.Writer
	if config.LogFilename == "" {
		logWriter = GinkgoWriter
	} else {
		logFile, err := os.OpenFile(config.LogFilename, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
		if err != nil {
			panic("Error opening file, err: " + err.Error())
		}
		defer logFile.Close()

		logWriter = logFile
	}

	logger = lager.NewLogger("bbs-benchmarks-test")
	logger.RegisterSink(lager.NewWriterSink(logWriter, lagerLogLevel))

	reporters = []Reporter{}

	if config.DataDogAPIKey != "" && config.DataDogAppKey != "" {
		dataDogClient = datadog.NewClient(config.DataDogAPIKey, config.DataDogAppKey)
		dataDogReporter = reporter.NewDataDogReporter(logger, config.MetricPrefix, dataDogClient)
		reporters = append(reporters, &dataDogReporter)
	}

	if config.AwsAccessKeyID != "" && config.AwsSecretAccessKey != "" && config.AwsBucketName != "" {
		uploader := s3manager.NewUploader(awssession.New(&aws.Config{
			Region:      &config.AwsRegion,
			Credentials: credentials.NewStaticCredentials(config.AwsAccessKeyID, config.AwsSecretAccessKey, ""),
		}))
		reporter := reporter.NewS3Reporter(logger, config.AwsBucketName, uploader)
		reporters = append(reporters, &reporter)
	}

	RegisterFailHandler(Fail)
	RunSpecsWithDefaultAndCustomReporters(t, "Benchmark BBS Suite", reporters)
}

func initializeActiveDB() *sql.DB {
	if activeDB != nil {
		return nil
	}

	if config.DatabaseConnectionString == "" {
		logger.Fatal("no-sql-configuration", errors.New("no-sql-configuration"))
	}

	if config.DatabaseDriver == "postgres" && !strings.Contains(config.DatabaseConnectionString, "sslmode") {
		config.DatabaseConnectionString = fmt.Sprintf("%s?sslmode=disable", config.DatabaseConnectionString)
	}

	sqlConn, err := helpers.Connect(logger, config.DatabaseDriver, config.DatabaseConnectionString, "", false)
	if err != nil {
		logger.Fatal("failed-to-open-sql", err)
	}
	sqlConn.SetMaxOpenConns(1)
	sqlConn.SetMaxIdleConns(1)

	err = sqlConn.Ping()
	Expect(err).NotTo(HaveOccurred())

	wrappedDB := helpers.NewMonitoredDB(sqlConn, monitor.New())
	sqlDB = initializeSQLDB(logger, wrappedDB)
	activeDB = sqlDB
	return sqlConn
}

var _ = BeforeSuite(func() {
	bbsClient = initializeBBSClient(logger, time.Duration(config.BBSClientHTTPTimeout))

	conn := initializeActiveDB()
	initializeLocketClient()

	resp, err := locketClient.FetchAll(context.Background(), &locketmodels.FetchAllRequest{TypeCode: locketmodels.PRESENCE})
	Expect(err).NotTo(HaveOccurred())

	cells := make(map[string]struct{})
	for i := 0; i < config.NumReps; i++ {
		cells[fmt.Sprintf("cell-%d", i)] = struct{}{}
	}

	if config.ReseedDatabase || config.NumReps != len(resp.Resources) {
		var err error

		cleanupSQLDB(conn)
		resetLocketDB(locketClient, cells, resp)

		if config.DesiredLRPs > 0 {
			desiredLRPGenerator := generator.NewDesiredLRPGenerator(
				config.ErrorTolerance,
				config.MetricPrefix,
				config.NumPopulateWorkers,
				bbsClient,
				dataDogClient,
			)

			expectedLRPCount, expectedActualLRPCounts, err = desiredLRPGenerator.Generate(logger, config.NumReps, config.DesiredLRPs)
			Expect(err).NotTo(HaveOccurred())
		}
	} else {
		resetUnclaimedActualLRPs(conn, cells)

		query := `
			SELECT
				COUNT(*)
			FROM desired_lrps
		`
		res := conn.QueryRow(query)
		err := res.Scan(&expectedLRPCount)
		Expect(err).NotTo(HaveOccurred())

		expectedActualLRPCounts = make(map[string]int)
		query = `
			SELECT
				cell_id, COUNT(*)
			FROM actual_lrps
			GROUP BY cell_id
		`
		rows, err := conn.Query(query)
		Expect(err).NotTo(HaveOccurred())

		defer rows.Close()
		for rows.Next() {
			var cellId string
			var count int
			err = rows.Scan(&cellId, &count)
			Expect(err).NotTo(HaveOccurred())

			expectedActualLRPCounts[cellId] = count
		}
	}

	if float64(expectedLRPCount) < float64(config.DesiredLRPs)*config.ErrorTolerance {
		Fail(fmt.Sprintf("Error rate of %.3f for actuals exceeds tolerance of %.3f", float64(expectedLRPCount)/float64(config.DesiredLRPs), config.ErrorTolerance))
	}
})

func initializeSQLDB(logger lager.Logger, sqlConn helpers.QueryableDB) *sqldb.SQLDB {
	key, keys, err := config.EncryptionConfig.Parse()
	if err != nil {
		logger.Fatal("cannot-setup-encryption", err)
	}
	keyManager, err := encryption.NewKeyManager(key, keys)
	if err != nil {
		logger.Fatal("cannot-setup-encryption", err)
	}
	cryptor := encryption.NewCryptor(keyManager, rand.Reader)

	return sqldb.NewSQLDB(
		sqlConn,
		1000,
		1000,
		cryptor,
		guidprovider.DefaultGuidProvider,
		clock.NewClock(),
		config.DatabaseDriver,
		&fakes.FakeIngressClient{},
	)
}

func initializeLocketClient() {
	var err error
	locketClient, err = locket.NewClient(logger, config.ClientLocketConfig)
	Expect(err).NotTo(HaveOccurred())
}

func initializeBBSClient(logger lager.Logger, bbsClientHTTPTimeout time.Duration) bbs.InternalClient {
	bbsClient, err := bbs.NewClientWithConfig(bbs.ClientConfig{
		URL:                    config.BBSAddress,
		IsTLS:                  true,
		CAFile:                 config.BBSCACert,
		CertFile:               config.BBSClientCert,
		KeyFile:                config.BBSClientKey,
		ClientSessionCacheSize: 1,
		MaxIdleConnsPerHost:    25000,
		RequestTimeout:         bbsClientHTTPTimeout,
	})
	if err != nil {
		logger.Fatal("Failed to configure secure BBS client", err)
	}
	return bbsClient
}

func cleanupSQLDB(conn *sql.DB) {
	_, err := conn.Exec("TRUNCATE actual_lrps")
	Expect(err).NotTo(HaveOccurred())
	_, err = conn.Exec("TRUNCATE desired_lrps")
	Expect(err).NotTo(HaveOccurred())
}

func resetLocketDB(locketClient locketmodels.LocketClient, cells map[string]struct{}, resp *locketmodels.FetchAllResponse) {
	for _, resource := range resp.Resources {
		if _, ok := cells[resource.Owner]; !ok {
			_, err := locketClient.Release(context.Background(), &locketmodels.ReleaseRequest{Resource: resource})
			Expect(err).NotTo(HaveOccurred())
		}
	}

	for cellId, _ := range cells {
		_, err := locketClient.Lock(context.Background(), &locketmodels.LockRequest{Resource: lockResource(cellId), TtlInSeconds: LongTermTTL})
		Expect(err).NotTo(HaveOccurred())
	}
}

func lockResource(cellID string) *locketmodels.Resource {
	resources := executor.ExecutorResources{
		MemoryMB:   1000,
		DiskMB:     2000,
		Containers: 300,
	}
	cellCapacity := models.NewCellCapacity(int32(resources.MemoryMB), int32(resources.DiskMB), int32(resources.Containers))
	cellPresence := models.NewCellPresence(cellID, cellID+".address", cellID+".url",
		"z1", cellCapacity, []string{"providers"},
		[]string{"cflinuxfs9"}, []string{}, []string{})

	payload, err := json.Marshal(cellPresence)
	Expect(err).NotTo(HaveOccurred())

	return &locketmodels.Resource{
		Key:   cellID,
		Owner: cellID,
		Value: string(payload),
		Type:  locketmodels.PresenceType,
	}
}

func resetUnclaimedActualLRPs(conn *sql.DB, cells map[string]struct{}) {
	existingCellIds := make(map[string]struct{})

	query := `
			SELECT
				cell_id
			FROM actual_lrps
			GROUP BY cell_id
		`
	rows, err := conn.Query(query)
	Expect(err).NotTo(HaveOccurred())

	defer rows.Close()
	for rows.Next() {
		var cellId string
		err = rows.Scan(&cellId)
		Expect(err).NotTo(HaveOccurred())

		existingCellIds[cellId] = struct{}{}
	}

	missingCells := findMissingCells(existingCellIds, cells)

	query = `
				SELECT
					process_guid, domain
				FROM actual_lrps WHERE cell_id = ''
			`

	type guidAndDomain struct {
		ProcessGuid string
		Domain      string
	}
	var lrps []guidAndDomain

	rows, err = conn.Query(query)
	Expect(err).NotTo(HaveOccurred())

	defer rows.Close()
	for rows.Next() {
		var guid, domain string
		err = rows.Scan(&guid, &domain)
		Expect(err).NotTo(HaveOccurred())
		lrps = append(lrps, guidAndDomain{ProcessGuid: guid, Domain: domain})
	}

	for i, lrp := range lrps {
		cellID := missingCells[i%len(missingCells)]
		actualLRPInstanceKey := &models.ActualLRPInstanceKey{InstanceGuid: lrp.ProcessGuid + "-i", CellId: cellID}
		netInfo := models.NewActualLRPNetInfo("1.2.3.4", "2.2.2.2", models.ActualLRPNetInfo_PreferredAddressHost, models.NewPortMapping(61999, 8080))
		err := bbsClient.StartActualLRP(logger, &models.ActualLRPKey{Domain: lrp.Domain, ProcessGuid: lrp.ProcessGuid, Index: 0}, actualLRPInstanceKey, &netInfo)
		Expect(err).NotTo(HaveOccurred())
	}
}

func findMissingCells(existing map[string]struct{}, expected map[string]struct{}) []string {
	var missing []string
	for cell, _ := range expected {
		if _, ok := existing[cell]; !ok {
			missing = append(missing, cell)
		}
	}

	return missing
}
<|endoftext|>


<commit_before>package frame

import (
	"image"
	"image/draw"
)

// Insert inserts the contents of s at index p0 in
// the frame and returns the number of characters
// written.
func (f *Frame) Insert(s []byte, p0 int64) (wrote int) {
	if p0 > f.Nchars || len(s) == 0 || f.b == nil {
		return
	}

	// find p0, it's box, and its point in the box its in
	b0 := f.Find(0, 0, p0)
	//	eb := f.StartCell(b0)
	//ob0 := b0
	cb0 := p0
	b1 := b0
	pt0 := f.pointOfBox(p0, b0)

	// find p1
	ppt0, pt1 := f.boxscan(s, pt0)
	opt0 := pt0
	ppt1 := pt1
	// Line wrap
	if b0 < f.Nbox {
		b := &f.Box[b0]
		pt0 = f.wrapMax(pt0, b)
		ppt1 = f.wrapMin(ppt1, b)
	}
	f.modified = true

	if f.p0 == f.p1 {
		f.tickat(f.PointOf(int64(f.p0)), false)
	}

	cb0, b0, pt0, pt1 = f.boxalign(cb0, b0, pt0, pt1)
	f.boxpush(p0, b0, b1, pt0, pt1, ppt1)
	f.bitblt(cb0, b0, pt0, pt1, opt0)
	text, back := f.pick(p0, f.p0+1, f.p1+1)
	f.Paint(ppt0, ppt1, back)
	f.redrawRun0(f.ir, ppt0, text, back)
	f.Run.Combine(f.ir, b1)
	if b1 > 0 && f.Box[b1-1].Nrune >= 0 && ppt0.X-f.Box[b1-1].Width >= f.r.Min.X {
		b1--
		ppt0.X -= f.Box[b1].Width
	}
	b0 += f.ir.Nbox
	if b0 < f.Nbox-1 {
		b0++
	}
	f.clean(ppt0, b1, b0)
	f.Nchars += f.ir.Nchars
	f.p0, f.p1 = coInsert(p0, p0+f.Nchars, f.p0, f.p1)
	if f.p0 == f.p1 {
		f.tickat(f.PointOf(f.p0), true)
	}
	f.badElasticAlg()
	return int(f.ir.Nchars)
}

func (f *Frame) Draw(dst draw.Image, r image.Rectangle, src image.Image, sp image.Point, op draw.Op) {
	if f == nil {
		panic("nil frame")
	}
	if f.Drawer == nil {
		panic("nil drawer")
	}
	f.Drawer.Draw(dst, r, src, sp, op)
	f.Flush(r)
}
func (f *Frame) badElasticAlg() {
	if f.elastic() {
		if f.Nbox <= 1 {
			return
		}
		b := 0
		b1 := 0
		for b != f.Nbox {
			b1 = b
			b = f.NextCell(b)
			if b == 0 || b == b1 {
				break
			}
		}
		f.Stretch(f.Nbox)
		b = b1
		for ; b1 > 1; b1 = f.Stretch(b1) {
			if b == b1 {
				break
			}
			b = b1
		}
		f.Stretch(b1)
		f.Refresh()
	}
}

// Mark marks the frame as dirty
func (f *Frame) Mark() {
	f.modified = true
}

// boxalign collects a list of pts of each box
// on the frame before and after an insertion occurs
func (f *Frame) boxalign(cb0 int64, b0 int, pt0, pt1 image.Point) (int64, int, image.Point, image.Point) {
	type Pts [2]image.Point
	f.pts = f.pts[:0]

	// collect the start pts for each box on the plane
	// before and after the insertion
	for {
		if pt0.X == pt1.X || pt1.Y == f.r.Max.Y || b0 == f.Nbox {
			break
		}
		b := &f.Box[b0]
		pt0 = f.wrapMax(pt0, b)
		pt1 = f.wrapMin(pt1, b)
		if b.Nrune > 0 {
			if n := f.fits(pt1, b); n != b.Nrune {
				f.Split(b0, n)
				b = &f.Box[b0]
			}
		}

		// early exit - point went off the frame
		if pt1.Y == f.r.Max.Y {
			break
		}

		f.pts = append(f.pts, Pts{pt0, pt1})

		pt0 = f.advance(pt0, b)
		pt1.X += f.plot(pt1, b)

		cb0 += int64(b.Len())
		b0++
	}
	return cb0, b0, pt0, pt1
}

// boxpush moves boxes down the frame to make room for an insertion
// from pt0 to pt1
func (f *Frame) boxpush(p0 int64, b0, b1 int, pt0, pt1, ppt1 image.Point) {
	h := f.Face.Dy()
	// delete boxes that ran off the frame
	// and update the char count
	if pt1.Y == f.r.Max.Y && b0 < f.Nbox {
		f.Nchars -= f.Count(b0)
		f.Run.Delete(b0, f.Nbox-1)
	}

	// update the line count
	if b0 == f.Nbox {
		f.Nlines = (pt1.Y - f.r.Min.Y) / h
		if pt1.X > f.r.Min.X {
			f.Nlines++
		}
		return
	}

	if pt1.Y == pt0.Y {
		// insertion won't propagate down
		return
	}

	qt0 := pt0.Y + h
	qt1 := pt1.Y + h
	f.Nlines += (qt1 - qt0) / h
	if f.Nlines > f.maxlines {
		f.trim(ppt1, p0, b1)
	}

	// shift down the existing boxes
	// on the bitmap
	if r := f.r; pt1.Y < r.Max.Y {
		r.Min.Y = qt1

		// rectangular group of boxes
		if qt1 < f.r.Max.Y {
			f.Draw(f.b, r, f.b, image.Pt(f.r.Min.X, qt0), f.op)
		}

		// partial line
		r.Min = pt1
		r.Max.X = pt1.X + (f.r.Max.X - pt0.X)
		r.Max.Y = qt1
		f.Draw(f.b, r, f.b, pt0, f.op)
	}
}

func (f *Frame) bitblt(cb0 int64, b0 int, pt0, pt1, opt0 image.Point) (res image.Rectangle) {
	h := f.Face.Dy()
	y := 0
	if pt1.Y == f.r.Max.Y {
		y = pt1.Y
	}
	x := len(f.pts)
	if x != 0 {
		res = image.Rectangle{pt0, pt1}
		res.Canon().Max.Add(f.pts[x-1][0])
	}
	run := f.Box[b0-x:]
	x--
	_, back := f.pick(cb0, f.p0, f.p1)
	for ; x >= 0; x-- {
		b := &run[x]
		br := image.Rect(0, 0, b.Width, h)
		pt := f.pts[x]
		if b.Nrune > 0 {
			f.Draw(f.b, br.Add(pt[1]), f.b, pt[0], f.op)
			// clear bit hanging off right
			if x == 0 && pt[1].Y > pt0.Y {

				// new char was wider than the old
				// one so the line wrapped anyway

				_, back = f.pick(cb0, f.p0, f.p1)
				r := br.Add(opt0)
				r.Max.X = f.r.Max.X
				f.Draw(f.b, r, back, r.Min, f.op)

			} else if pt[1].Y < y {

				// copy from left to right, bottom to top

				_, back = f.pick(cb0, f.p0, f.p1)
				r := image.ZR.Add(pt[1])
				r.Min.X += b.Width
				r.Max.X += f.r.Max.X
				r.Max.Y += h
				f.Draw(f.b, r, back, r.Min, f.op)

			}
			y = pt[1].Y
			cb0 -= int64(b.Nrune)
		} else {
			r := br.Add(pt[1])
			if r.Max.X >= f.r.Max.X {
				r.Max.X = f.r.Max.X
			}
			cb0--
			_, back = f.pick(cb0, f.p0, f.p1)
			f.Draw(f.b, r, back, r.Min, f.op)
			y = 0
			if pt[1].X == f.r.Min.X {
				y = pt[1].Y
			}
		}
	}
	return res
}

func (f *Frame) pick(c, p0, p1 int64) (text, back image.Image) {
	if p0 <= c && c < p1 {
		return f.Color.Hi.Text, f.Color.Hi.Back
	}
	return f.Color.Text, f.Color.Back
}

/*
//
// Below ideas

type offset struct {
	p0   int64
	b0   int
	cb0  int64
	b1   int64
	pt0  image.Point
	pt1  image.Point
	opt0 image.Point
}

func (f *Frame) zInsertElastic(s []byte, p0 int64) (wrote int) {
	type Pts [2]image.Point
	if p0 > f.Nchars || len(s) == 0 || f.b == nil {
		return
	}

	// find p0, it's box, and its point in the box its in
	b0 := f.Find(0, 0, p0)
	cb0 := p0

	eb0 := f.StartCell(b0)
	pt0 := f.pointOfBox(p0, b0)
	ppt0, pt1 := f.boxscan(s, pt0)
	//	ept0 := f.pointOfBox(p0, eb0)
	f.Box = append(f.Box[:b0], append(f.ir.Box[:f.ir.Nbox], f.Box[b0:]...)...)
	f.Nbox += f.ir.Nbox

	b1 := b0

	for bn := f.NextCell(b1 + f.ir.Nbox); bn > eb0; bn = f.Stretch(bn) {
	}
	f.Stretch(eb0)

	//		pt0 = f.pointOfBox(p0, b0)
	//		pt1 = f.pointOfBox(65535, b1)

	opt0 := pt0
	ppt1 := pt1
	// Line wrap
	if b0 < f.Nbox {
		b := &f.Box[b0]
		pt0 = f.wrapMax(pt0, b)
		ppt1 = f.wrapMin(ppt1, b)
	}
	f.modified = true

	if f.p0 == f.p1 {
		f.tickat(f.PointOf(int64(f.p0)), false)
	}

	cb0, b0, pt0, pt1 = f.boxalign(cb0, b0+f.ir.Nbox, pt0, pt1)
	f.boxpush(p0, b0+f.ir.Nbox, b1+f.ir.Nbox, pt0, pt1, ppt1)
	f.bitblt(cb0, b0, pt0, pt1, opt0)
	text, back := f.pick(p0, f.p0+1, f.p1+1)
	f.Paint(ppt0, ppt1, back)
	f.redrawRun0(f.ir, ppt0, text, back)

	b1 = b0
	if b1 > 0 && f.Box[b1-1].Nrune >= 0 && ppt0.X-f.Box[b1-1].Width >= f.r.Min.X {
		b1--
		ppt0.X -= f.Box[b1].Width
	}

	b0 += f.ir.Nbox
	if b0 < f.Nbox-1 {
		b0++
	}
	f.clean(f.pointOfBox(p0, b1), b0, b1)
	f.Nchars += f.ir.Nchars

	f.p0, f.p1 = coInsert(p0, p0+f.Nchars, f.p0, f.p1)
	if f.p0 == f.p1 {
		f.tickat(f.PointOf(f.p0), true)
	}
	return int(f.ir.Nchars)

}
*/
<commit_msg>frame: insert: fixes #12<commit_after>package frame

import (
	"image"
	"image/draw"
)

// Insert inserts the contents of s at index p0 in
// the frame and returns the number of characters
// written.
func (f *Frame) Insert(s []byte, p0 int64) (wrote int) {
	if p0 > f.Nchars || len(s) == 0 || f.b == nil {
		return
	}

	// find p0, it's box, and its point in the box its in
	b0 := f.Find(0, 0, p0)
	//	eb := f.StartCell(b0)
	//ob0 := b0
	cb0 := p0
	b1 := b0
	pt0 := f.pointOfBox(p0, b0)

	// find p1
	ppt0, pt1 := f.boxscan(s, pt0)
	opt0 := pt0
	ppt1 := pt1
	// Line wrap
	if b0 < f.Nbox {
		b := &f.Box[b0]
		pt0 = f.wrapMax(pt0, b)
		ppt1 = f.wrapMin(ppt1, b)
	}
	f.modified = true

	if f.p0 == f.p1 {
		f.tickat(f.PointOf(int64(f.p0)), false)
	}

	cb0, b0, pt0, pt1 = f.boxalign(cb0, b0, pt0, pt1)
	f.boxpush(p0, b0, b1, pt0, pt1, ppt1)
	f.bitblt(cb0, b0, pt0, pt1, opt0)
	text, back := f.pick(p0, f.p0+1, f.p1+1)
	f.Paint(ppt0, ppt1, back)
	f.redrawRun0(f.ir, ppt0, text, back)
	f.Run.Combine(f.ir, b1)
	if b1 > 0 && f.Box[b1-1].Nrune >= 0 && ppt0.X-f.Box[b1-1].Width >= f.r.Min.X {
		b1--
		ppt0.X -= f.Box[b1].Width
	}
	b0 += f.ir.Nbox
	if b0 < f.Nbox-1 {
		b0++
	}
	f.clean(ppt0, b1, b0)
	f.Nchars += f.ir.Nchars
	if p0 <= f.p0{
		f.p0+=f.ir.Nchars
		f.p1+=f.ir.Nchars
	} else if p0 < f.p1{
		f.p1+=f.ir.Nchars
	}
	if f.p0 == f.p1 {
		f.tickat(f.PointOf(f.p0), true)
	}
	f.badElasticAlg()
	return int(f.ir.Nchars)
}

func (f *Frame) Draw(dst draw.Image, r image.Rectangle, src image.Image, sp image.Point, op draw.Op) {
	if f == nil {
		panic("nil frame")
	}
	if f.Drawer == nil {
		panic("nil drawer")
	}
	f.Drawer.Draw(dst, r, src, sp, op)
	f.Flush(r)
}
func (f *Frame) badElasticAlg() {
	if f.elastic() {
		if f.Nbox <= 1 {
			return
		}
		b := 0
		b1 := 0
		for b != f.Nbox {
			b1 = b
			b = f.NextCell(b)
			if b == 0 || b == b1 {
				break
			}
		}
		f.Stretch(f.Nbox)
		b = b1
		for ; b1 > 1; b1 = f.Stretch(b1) {
			if b == b1 {
				break
			}
			b = b1
		}
		f.Stretch(b1)
		f.Refresh()
	}
}

// Mark marks the frame as dirty
func (f *Frame) Mark() {
	f.modified = true
}

// boxalign collects a list of pts of each box
// on the frame before and after an insertion occurs
func (f *Frame) boxalign(cb0 int64, b0 int, pt0, pt1 image.Point) (int64, int, image.Point, image.Point) {
	type Pts [2]image.Point
	f.pts = f.pts[:0]

	// collect the start pts for each box on the plane
	// before and after the insertion
	for {
		if pt0.X == pt1.X || pt1.Y == f.r.Max.Y || b0 == f.Nbox {
			break
		}
		b := &f.Box[b0]
		pt0 = f.wrapMax(pt0, b)
		pt1 = f.wrapMin(pt1, b)
		if b.Nrune > 0 {
			if n := f.fits(pt1, b); n != b.Nrune {
				f.Split(b0, n)
				b = &f.Box[b0]
			}
		}

		// early exit - point went off the frame
		if pt1.Y == f.r.Max.Y {
			break
		}

		f.pts = append(f.pts, Pts{pt0, pt1})

		pt0 = f.advance(pt0, b)
		pt1.X += f.plot(pt1, b)

		cb0 += int64(b.Len())
		b0++
	}
	return cb0, b0, pt0, pt1
}

// boxpush moves boxes down the frame to make room for an insertion
// from pt0 to pt1
func (f *Frame) boxpush(p0 int64, b0, b1 int, pt0, pt1, ppt1 image.Point) {
	h := f.Face.Dy()
	// delete boxes that ran off the frame
	// and update the char count
	if pt1.Y == f.r.Max.Y && b0 < f.Nbox {
		f.Nchars -= f.Count(b0)
		f.Run.Delete(b0, f.Nbox-1)
	}

	// update the line count
	if b0 == f.Nbox {
		f.Nlines = (pt1.Y - f.r.Min.Y) / h
		if pt1.X > f.r.Min.X {
			f.Nlines++
		}
		return
	}

	if pt1.Y == pt0.Y {
		// insertion won't propagate down
		return
	}

	qt0 := pt0.Y + h
	qt1 := pt1.Y + h
	f.Nlines += (qt1 - qt0) / h
	if f.Nlines > f.maxlines {
		f.trim(ppt1, p0, b1)
	}

	// shift down the existing boxes
	// on the bitmap
	if r := f.r; pt1.Y < r.Max.Y {
		r.Min.Y = qt1

		// rectangular group of boxes
		if qt1 < f.r.Max.Y {
			f.Draw(f.b, r, f.b, image.Pt(f.r.Min.X, qt0), f.op)
		}

		// partial line
		r.Min = pt1
		r.Max.X = pt1.X + (f.r.Max.X - pt0.X)
		r.Max.Y = qt1
		f.Draw(f.b, r, f.b, pt0, f.op)
	}
}

func (f *Frame) bitblt(cb0 int64, b0 int, pt0, pt1, opt0 image.Point) (res image.Rectangle) {
	h := f.Face.Dy()
	y := 0
	if pt1.Y == f.r.Max.Y {
		y = pt1.Y
	}
	x := len(f.pts)
	if x != 0 {
		res = image.Rectangle{pt0, pt1}
		res.Canon().Max.Add(f.pts[x-1][0])
	}
	run := f.Box[b0-x:]
	x--
	_, back := f.pick(cb0, f.p0, f.p1)
	for ; x >= 0; x-- {
		b := &run[x]
		br := image.Rect(0, 0, b.Width, h)
		pt := f.pts[x]
		if b.Nrune > 0 {
			f.Draw(f.b, br.Add(pt[1]), f.b, pt[0], f.op)
			// clear bit hanging off right
			if x == 0 && pt[1].Y > pt0.Y {

				// new char was wider than the old
				// one so the line wrapped anyway

				_, back = f.pick(cb0, f.p0, f.p1)
				r := br.Add(opt0)
				r.Max.X = f.r.Max.X
				f.Draw(f.b, r, back, r.Min, f.op)

			} else if pt[1].Y < y {

				// copy from left to right, bottom to top

				_, back = f.pick(cb0, f.p0, f.p1)
				r := image.ZR.Add(pt[1])
				r.Min.X += b.Width
				r.Max.X += f.r.Max.X
				r.Max.Y += h
				f.Draw(f.b, r, back, r.Min, f.op)

			}
			y = pt[1].Y
			cb0 -= int64(b.Nrune)
		} else {
			r := br.Add(pt[1])
			if r.Max.X >= f.r.Max.X {
				r.Max.X = f.r.Max.X
			}
			cb0--
			_, back = f.pick(cb0, f.p0, f.p1)
			f.Draw(f.b, r, back, r.Min, f.op)
			y = 0
			if pt[1].X == f.r.Min.X {
				y = pt[1].Y
			}
		}
	}
	return res
}

func (f *Frame) pick(c, p0, p1 int64) (text, back image.Image) {
	if p0 <= c && c < p1 {
		return f.Color.Hi.Text, f.Color.Hi.Back
	}
	return f.Color.Text, f.Color.Back
}

/*
//
// Below ideas

type offset struct {
	p0   int64
	b0   int
	cb0  int64
	b1   int64
	pt0  image.Point
	pt1  image.Point
	opt0 image.Point
}

func (f *Frame) zInsertElastic(s []byte, p0 int64) (wrote int) {
	type Pts [2]image.Point
	if p0 > f.Nchars || len(s) == 0 || f.b == nil {
		return
	}

	// find p0, it's box, and its point in the box its in
	b0 := f.Find(0, 0, p0)
	cb0 := p0

	eb0 := f.StartCell(b0)
	pt0 := f.pointOfBox(p0, b0)
	ppt0, pt1 := f.boxscan(s, pt0)
	//	ept0 := f.pointOfBox(p0, eb0)
	f.Box = append(f.Box[:b0], append(f.ir.Box[:f.ir.Nbox], f.Box[b0:]...)...)
	f.Nbox += f.ir.Nbox

	b1 := b0

	for bn := f.NextCell(b1 + f.ir.Nbox); bn > eb0; bn = f.Stretch(bn) {
	}
	f.Stretch(eb0)

	//		pt0 = f.pointOfBox(p0, b0)
	//		pt1 = f.pointOfBox(65535, b1)

	opt0 := pt0
	ppt1 := pt1
	// Line wrap
	if b0 < f.Nbox {
		b := &f.Box[b0]
		pt0 = f.wrapMax(pt0, b)
		ppt1 = f.wrapMin(ppt1, b)
	}
	f.modified = true

	if f.p0 == f.p1 {
		f.tickat(f.PointOf(int64(f.p0)), false)
	}

	cb0, b0, pt0, pt1 = f.boxalign(cb0, b0+f.ir.Nbox, pt0, pt1)
	f.boxpush(p0, b0+f.ir.Nbox, b1+f.ir.Nbox, pt0, pt1, ppt1)
	f.bitblt(cb0, b0, pt0, pt1, opt0)
	text, back := f.pick(p0, f.p0+1, f.p1+1)
	f.Paint(ppt0, ppt1, back)
	f.redrawRun0(f.ir, ppt0, text, back)

	b1 = b0
	if b1 > 0 && f.Box[b1-1].Nrune >= 0 && ppt0.X-f.Box[b1-1].Width >= f.r.Min.X {
		b1--
		ppt0.X -= f.Box[b1].Width
	}

	b0 += f.ir.Nbox
	if b0 < f.Nbox-1 {
		b0++
	}
	f.clean(f.pointOfBox(p0, b1), b0, b1)
	f.Nchars += f.ir.Nchars

	f.p0, f.p1 = coInsert(p0, p0+f.Nchars, f.p0, f.p1)
	if f.p0 == f.p1 {
		f.tickat(f.PointOf(f.p0), true)
	}
	return int(f.ir.Nchars)

}
*/
<|endoftext|>


<commit_before>package main

import "bufio"
import "encoding/binary"
import "fmt"
import "io"
import "log"
import "os"
import "runtime"
import "sync"
import "sync/atomic"

type Edge struct {
	From uint32
	To   uint32
}

func processEdgeStore(edgeStore []uint64, f func(uint32, uint32)) {
	for _, e := range edgeStore {
		// Seperate the two 32 bit nodes from the 64 bit edge
		eFrom := uint32(e >> 32)
		// Converting uint64 to uint32 drops the top 32 bits
		// Had (edge & 0xFFFFFFFF) for clarity but Go compiler doesn't optimize it away ...
		eTo := uint32(e)
		// Edges are distributed across workers according to either source or destination node
		f(eFrom, eTo)
	}
}

func sendEdges(filename string, f func(uint32, uint32)) {
	file, _ := os.Open(filename)
	defer file.Close()
	// Adds the ReadByte method requird by io.ByteReader interface
	wrappedByteReader := bufio.NewReader(file)
	edge := uint64(0)
	edgeStore := make([]uint64, 0, 524288)
	for {
		// Read the variable integer and undo the delta encoding by adding the previous edge
		rawEdge, err := binary.ReadUvarint(wrappedByteReader)
		edge += rawEdge
		if err == io.EOF {
			break
		}
		if err != nil {
			panic(err)
		}
		if len(edgeStore) == cap(edgeStore) {
			processEdgeStore(edgeStore, f)
			// Empty the edge store
			edgeStore = edgeStore[0:0]
		}
		edgeStore = append(edgeStore, edge)
	}
	processEdgeStore(edgeStore, f)
}

func applyFunctionToEdges(f func(uint32, uint32)) {
	var senderGroup sync.WaitGroup
	totalParts := 4
	for i := 0; i < totalParts; i++ {
		senderGroup.Add(1)
		go func(i int) {
			sendEdges(fmt.Sprintf("pld-arc.%d.bin", i), f)
			log.Printf("Completed processing part %d of %d\n", i, totalParts)
			senderGroup.Done()
		}(i)
	}
	//
	senderGroup.Wait()
}

func main() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	// We can either have total nodes supplied by the user or perform a full traversal of the data
	total := uint32(42889799 + 1)
	alpha := float64(0.85)
	saveResults := false
	//
	// floats must be 64 bit else we lose probability mass to floating point errors
	// This is annoying as it substantially increases our in-memory requirements ...
	src := make([]float64, total, total)
	dest := make([]float64, total, total)
	degree := make([]uint32, total, total)
	//
	log.Printf("Calculating degree of each source node\n")
	applyFunctionToEdges(func(from uint32, to uint32) {
		// Atomic is necessary here as each file is partitioned on the Edge(from, to)'s "to"
		// Different workers may try to update the same "from" in a non-atomic fashion
		atomic.AddUint32(&degree[from], 1)
	})
	//
	// The first loop should have alpha set, else it's a noop as all src is equal to zero
	// We also distribute the starting probability mass s.t. it totals one
	for i := range dest {
		dest[i] = 1 / float64(total)
	}
	//
	for iter := 0; iter < 20; iter++ {
		log.Printf("PageRank Iteration: %d\n", iter+1)
		log.Printf("Calculating the source and destination vectors\n")
		// Calculate the probability mass that will be lost via dangling nodes
		missingProb := float64(0)
		for i := range degree {
			if degree[i] == 0 {
				missingProb += dest[i]
			}
		}
		// Calculate the starting values
		for i := range dest {
			// If the node is dangling, src will equal +Inf due to degree being zero
			// As the result is not used elsewhere, this isn't so much a problem
			src[i] = alpha * (dest[i] / float64(degree[i]))
			dest[i] = ((1 - alpha) / float64(total))
		}
		// Distribute the probability mass according to the edges
		log.Printf("Calculating the probability mass gifted by incoming edges\n")
		applyFunctionToEdges(func(from uint32, to uint32) {
			dest[to] += src[from]
		})
		// Replace missing probability mass from dangling nodes
		// (assumption is that they were equally distributed to all nodes)
		for i := range dest {
			dest[i] += alpha * (missingProb / float64(total))
		}
	}
	// Write result
	if saveResults {
		log.Printf("Saving results\n")
		outf, _ := os.Create("result.txt")
		defer outf.Close()
		w := bufio.NewWriter(outf)
		defer w.Flush()
		for i, v := range dest {
			w.WriteString(fmt.Sprintf("%d\t%.12f\n", i, v))
		}
		log.Printf("Saved results\n")
	}
	//
	totalProb := float64(0)
	for _, v := range dest {
		totalProb += v
	}
	log.Printf("Total probability mass: %f\n", totalProb)
}
<commit_msg>Add HTTP based pprof support for CPU profiling<commit_after>package main

import "bufio"
import "encoding/binary"
import "fmt"
import "net/http"
import "io"
import "log"
import "os"
import "runtime"
import "sync"
import "sync/atomic"
import _ "net/http/pprof"

type Edge struct {
	From uint32
	To   uint32
}

func processEdgeStore(edgeStore []uint64, f func(uint32, uint32)) {
	for _, e := range edgeStore {
		// Seperate the two 32 bit nodes from the 64 bit edge
		eFrom := uint32(e >> 32)
		// Converting uint64 to uint32 drops the top 32 bits
		// Had (edge & 0xFFFFFFFF) for clarity but Go compiler doesn't optimize it away ...
		eTo := uint32(e)
		// Edges are distributed across workers according to either source or destination node
		f(eFrom, eTo)
	}
}

func sendEdges(filename string, f func(uint32, uint32)) {
	file, _ := os.Open(filename)
	defer file.Close()
	// Adds the ReadByte method requird by io.ByteReader interface
	wrappedByteReader := bufio.NewReader(file)
	edge := uint64(0)
	edgeStore := make([]uint64, 0, 524288)
	for {
		// Read the variable integer and undo the delta encoding by adding the previous edge
		rawEdge, err := binary.ReadUvarint(wrappedByteReader)
		edge += rawEdge
		if err == io.EOF {
			break
		}
		if err != nil {
			panic(err)
		}
		if len(edgeStore) == cap(edgeStore) {
			processEdgeStore(edgeStore, f)
			// Empty the edge store
			edgeStore = edgeStore[0:0]
		}
		edgeStore = append(edgeStore, edge)
	}
	processEdgeStore(edgeStore, f)
}

func applyFunctionToEdges(f func(uint32, uint32)) {
	var senderGroup sync.WaitGroup
	totalParts := 4
	for i := 0; i < totalParts; i++ {
		senderGroup.Add(1)
		go func(i int) {
			sendEdges(fmt.Sprintf("pld-arc.%d.bin", i), f)
			log.Printf("Completed processing part %d of %d\n", i, totalParts)
			senderGroup.Done()
		}(i)
	}
	//
	senderGroup.Wait()
}

func main() {
	// Start a server for pprof so we can get CPU profiles whilst running
	go func() {
		log.Println("Starting pprof server at localhost:6060...")
		log.Println(http.ListenAndServe("localhost:6060", nil))
	}()
	//
	runtime.GOMAXPROCS(runtime.NumCPU())
	// We can either have total nodes supplied by the user or perform a full traversal of the data
	total := uint32(42889799 + 1)
	alpha := float64(0.85)
	saveResults := false
	//
	// floats must be 64 bit else we lose probability mass to floating point errors
	// This is annoying as it substantially increases our in-memory requirements ...
	src := make([]float64, total, total)
	dest := make([]float64, total, total)
	degree := make([]uint32, total, total)
	//
	log.Printf("Calculating degree of each source node\n")
	applyFunctionToEdges(func(from uint32, to uint32) {
		// Atomic is necessary here as each file is partitioned on the Edge(from, to)'s "to"
		// Different workers may try to update the same "from" in a non-atomic fashion
		atomic.AddUint32(&degree[from], 1)
	})
	//
	// The first loop should have alpha set, else it's a noop as all src is equal to zero
	// We also distribute the starting probability mass s.t. it totals one
	for i := range dest {
		dest[i] = 1 / float64(total)
	}
	//
	for iter := 0; iter < 20; iter++ {
		log.Printf("PageRank Iteration: %d\n", iter+1)
		log.Printf("Calculating the source and destination vectors\n")
		// Calculate the probability mass that will be lost via dangling nodes
		missingProb := float64(0)
		for i := range degree {
			if degree[i] == 0 {
				missingProb += dest[i]
			}
		}
		// Calculate the starting values
		for i := range dest {
			// If the node is dangling, src will equal +Inf due to degree being zero
			// As the result is not used elsewhere, this isn't so much a problem
			src[i] = alpha * (dest[i] / float64(degree[i]))
			dest[i] = ((1 - alpha) / float64(total))
		}
		// Distribute the probability mass according to the edges
		log.Printf("Calculating the probability mass gifted by incoming edges\n")
		applyFunctionToEdges(func(from uint32, to uint32) {
			dest[to] += src[from]
		})
		// Replace missing probability mass from dangling nodes
		// (assumption is that they were equally distributed to all nodes)
		for i := range dest {
			dest[i] += alpha * (missingProb / float64(total))
		}
	}
	// Write result
	if saveResults {
		log.Printf("Saving results\n")
		outf, _ := os.Create("result.txt")
		defer outf.Close()
		w := bufio.NewWriter(outf)
		defer w.Flush()
		for i, v := range dest {
			w.WriteString(fmt.Sprintf("%d\t%.12f\n", i, v))
		}
		log.Printf("Saved results\n")
	}
	//
	totalProb := float64(0)
	for _, v := range dest {
		totalProb += v
	}
	log.Printf("Total probability mass: %f\n", totalProb)
}
<|endoftext|>


<commit_before>package main

import (
	"code.google.com/p/go.crypto/ssh"
	"code.google.com/p/goauth2/oauth"
	"crypto"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"github.com/google/go-github/github"
	"github.com/gorilla/mux"
	"github.com/kylelemons/go-gypsy/yaml"
	"html/template"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"os/user"
	"path"
	"reflect"
	"strings"
	"sync"
)

var (
	port       = "8080"
	sshPort    = "22"
	configFile = "config.yml"
	projects   []Project
)

type Host struct {
	URI          string
	LatestCommit string
}

type Environment struct {
	Name               string
	RepoPath           string
	Hosts              []Host
	Branch             string
	LatestGitHubCommit string
}

type Project struct {
	Name         string
	GitHubURL    string
	RepoName     string
	RepoOwner    string
	Environments []Environment
}

/*
Thanks to Russ Cox: https://groups.google.com/d/msg/golang-nuts/OEdSDgEC7js/iyhU9DW_IKcJ
eq reports whether the first argument is equal to
any of the remaining arguments. (for use in templates)
*/
func eq(args ...interface{}) bool {
	if len(args) == 0 {
		return false
	}
	x := args[0]
	switch x := x.(type) {
	case string, int, int64, byte, float32, float64:
		for _, y := range args[1:] {
			if x == y {
				return true
			}
		}
		return false
	}

	for _, y := range args[1:] {
		if reflect.DeepEqual(x, y) {
			return true
		}
	}
	return false
}

func (h *Host) ShortCommitHash() string {
	return h.LatestCommit[:7]
}

func getPrivateKey(filename string) []byte {
	content, err := ioutil.ReadFile(filename)
	if err != nil {
		log.Panic("Failed to open private key file: " + err.Error())
	}
	return content
}

type keychain struct {
	key *rsa.PrivateKey
}

func (k *keychain) Key(i int) (interface{}, error) {
	if i != 0 {
		return nil, nil
	}
	return &k.key.PublicKey, nil
}

func (k *keychain) Sign(i int, rand io.Reader, data []byte) (sig []byte, err error) {
	hashFunc := crypto.SHA1
	h := hashFunc.New()
	h.Write(data)
	digest := h.Sum(nil)
	return rsa.SignPKCS1v15(rand, k.key, hashFunc, digest)
}

func remoteCmdOutput(username, hostname, privateKey, cmd string) []byte {
	block, _ := pem.Decode([]byte(privateKey))
	rsakey, _ := x509.ParsePKCS1PrivateKey(block.Bytes)
	clientKey := &keychain{rsakey}
	clientConfig := &ssh.ClientConfig{
		User: username,
		Auth: []ssh.ClientAuth{
			ssh.ClientAuthKeyring(clientKey),
		},
	}
	client, err := ssh.Dial("tcp", hostname, clientConfig)
	if err != nil {
		log.Panic("Failed to dial: " + err.Error())
	}
	session, err := client.NewSession()
	if err != nil {
		log.Panic("Failed to create session: " + err.Error())
	}
	defer session.Close()
	output, err := session.Output(cmd)
	if err != nil {
		log.Panic("Failed to run cmd: " + err.Error())
	}
	return output
}

func latestDeployedCommit(username, hostname string, e Environment) []byte {
	usr, err := user.Current()
	if err != nil {
		log.Fatal(err)
	}
	privateKey := string(getPrivateKey(path.Join(usr.HomeDir, "/.ssh/id_rsa")))
	output := remoteCmdOutput(username, hostname, privateKey, fmt.Sprintf("git --git-dir=%s rev-parse %s", e.RepoPath, e.Branch))

	return output
}

func getYAMLString(n yaml.Node, key string) string {
	return n.(yaml.Map)[key].(yaml.Scalar).String()
}

func parseYAMLEnvironment(m yaml.Node) Environment {
	e := Environment{}
	for k, v := range m.(yaml.Map) {
		e.Name = k
		e.Branch = getYAMLString(v, "branch")
		e.RepoPath = getYAMLString(v, "repo_path")
		for _, v := range v.(yaml.Map)["hosts"].(yaml.List) {
			h := Host{URI: v.(yaml.Scalar).String()}
			e.Hosts = append(e.Hosts, h)
		}
	}
	return e
}

func parseYAML() (allProjects []Project, deployUser string) {
	config, err := yaml.ReadFile(configFile)
	if err != nil {
		log.Fatal(err)
	}
	deployUser, err = config.Get("deploy_user")
	if err != nil {
		log.Fatal("config.yml is missing deploy_user: " + err.Error())
	}
	configRoot, _ := config.Root.(yaml.Map)
	projects, _ := configRoot["projects"].(yaml.List)
	allProjects = []Project{}
	for _, p := range projects {
		for _, v := range p.(yaml.Map) {
			proj := Project{Name: getYAMLString(v, "project_name"),
				GitHubURL: getYAMLString(v, "github_url"),
				RepoName:  getYAMLString(v, "repo_name"),
				RepoOwner: getYAMLString(v, "repo_owner")}
			for _, v := range v.(yaml.Map)["environments"].(yaml.List) {
				proj.Environments = append(proj.Environments, parseYAMLEnvironment(v))
			}
			allProjects = append(allProjects, proj)
		}
	}
	return allProjects, deployUser
}

func getCommit(wg *sync.WaitGroup, projects []Project, env Environment, host Host, deployUser string, i, j, k int) {
	defer wg.Done()
	host.LatestCommit = string(latestDeployedCommit(deployUser, host.URI+":"+sshPort, env))
	host.LatestCommit = strings.Trim(host.LatestCommit, "\n\r")
	projects[i].Environments[j].Hosts[k] = host
}

//  Get the most recent commit hash on a given branch from GitHub
func getLatestGitHubCommit(wg *sync.WaitGroup, projects []Project, environment Environment, c *github.Client, repoOwner, repoName string, i, j int) {
	defer wg.Done()
	opts := &github.CommitsListOptions{SHA: environment.Branch}
	commits, _, err := c.Repositories.ListCommits(repoOwner, repoName, opts)
	if err != nil {
		log.Panic(err)
	}
	environment.LatestGitHubCommit = *commits[0].SHA
	projects[i].Environments[j] = environment
}

func retrieveCommits(projects []Project, deployUser string) []Project {
	// define a wait group to wait for all goroutines to finish
	var wg sync.WaitGroup
	githubToken := os.Getenv("GITHUB_API_TOKEN")
	t := &oauth.Transport{
		Token: &oauth.Token{AccessToken: githubToken},
	}
	client := github.NewClient(t.Client())
	for i, project := range projects {
		for j, environment := range project.Environments {
			for k, host := range environment.Hosts {
				// start a goroutine for SSH-ing on to the machine
				wg.Add(1)
				go getCommit(&wg, projects, environment, host, deployUser, i, j, k)
			}
			wg.Add(1)
			go getLatestGitHubCommit(&wg, projects, environment, client, project.RepoOwner, project.RepoName, i, j)
		}
	}
	// wait for goroutines to finish
	wg.Wait()
	return projects
}

func HomeHandler(w http.ResponseWriter, r *http.Request) {
	var deployUser string
	projects, deployUser = parseYAML()
	projects = retrieveCommits(projects, deployUser)
	t, err := template.New("index.html").Funcs(template.FuncMap{"eq": eq}).ParseFiles("templates/index.html")
	if err != nil {
		log.Panic(err)
	}
	err = t.Execute(w, map[string]interface{}{"Projects": projects})
	if err != nil {
		log.Panic(err)
	}
}

func DeployHandler(w http.ResponseWriter, r *http.Request) {
	project := r.FormValue("project")
	environment := r.FormValue("environment")
	command := strings.Split(projects[project][environment]["deploy_command"], " ")
	fmt.Println(command)
	cmd := exec.Command(command)
	if err != nil {
		log.Fatal(err)
	}
}

func main() {
	r := mux.NewRouter()
	r.HandleFunc("/", HomeHandler)
	r.HandleFunc("/deploy", DeployHandler)
	fmt.Println("Running on localhost:" + port)
	log.Fatal(http.ListenAndServe(":"+port, r))
}
<commit_msg>add handler for /deploy<commit_after>package main

import (
	"bytes"
	"code.google.com/p/go.crypto/ssh"
	"code.google.com/p/goauth2/oauth"
	"crypto"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"github.com/google/go-github/github"
	"github.com/gorilla/mux"
	"github.com/kylelemons/go-gypsy/yaml"
	"html/template"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"os/exec"
	"os/user"
	"path"
	"reflect"
	"strings"
	"sync"
)

var (
	port       = "8080"
	sshPort    = "22"
	configFile = "config.yml"
	projects   []Project
)

type Host struct {
	URI          string
	LatestCommit string
}

type Environment struct {
	Name               string
	Deploy             string
	RepoPath           string
	Hosts              []Host
	Branch             string
	LatestGitHubCommit string
}

type Project struct {
	Name         string
	GitHubURL    string
	RepoName     string
	RepoOwner    string
	Environments []Environment
}

/*
Thanks to Russ Cox: https://groups.google.com/d/msg/golang-nuts/OEdSDgEC7js/iyhU9DW_IKcJ
eq reports whether the first argument is equal to
any of the remaining arguments. (for use in templates)
*/
func eq(args ...interface{}) bool {
	if len(args) == 0 {
		return false
	}
	x := args[0]
	switch x := x.(type) {
	case string, int, int64, byte, float32, float64:
		for _, y := range args[1:] {
			if x == y {
				return true
			}
		}
		return false
	}

	for _, y := range args[1:] {
		if reflect.DeepEqual(x, y) {
			return true
		}
	}
	return false
}

func (h *Host) ShortCommitHash() string {
	return h.LatestCommit[:7]
}

func getPrivateKey(filename string) []byte {
	content, err := ioutil.ReadFile(filename)
	if err != nil {
		log.Panic("Failed to open private key file: " + err.Error())
	}
	return content
}

type keychain struct {
	key *rsa.PrivateKey
}

func (k *keychain) Key(i int) (interface{}, error) {
	if i != 0 {
		return nil, nil
	}
	return &k.key.PublicKey, nil
}

func (k *keychain) Sign(i int, rand io.Reader, data []byte) (sig []byte, err error) {
	hashFunc := crypto.SHA1
	h := hashFunc.New()
	h.Write(data)
	digest := h.Sum(nil)
	return rsa.SignPKCS1v15(rand, k.key, hashFunc, digest)
}

func remoteCmdOutput(username, hostname, privateKey, cmd string) []byte {
	block, _ := pem.Decode([]byte(privateKey))
	rsakey, _ := x509.ParsePKCS1PrivateKey(block.Bytes)
	clientKey := &keychain{rsakey}
	clientConfig := &ssh.ClientConfig{
		User: username,
		Auth: []ssh.ClientAuth{
			ssh.ClientAuthKeyring(clientKey),
		},
	}
	client, err := ssh.Dial("tcp", hostname, clientConfig)
	if err != nil {
		log.Panic("Failed to dial: " + err.Error())
	}
	session, err := client.NewSession()
	if err != nil {
		log.Panic("Failed to create session: " + err.Error())
	}
	defer session.Close()
	output, err := session.Output(cmd)
	if err != nil {
		log.Panic("Failed to run cmd: " + err.Error())
	}
	return output
}

func latestDeployedCommit(username, hostname string, e Environment) []byte {
	usr, err := user.Current()
	if err != nil {
		log.Fatal(err)
	}
	privateKey := string(getPrivateKey(path.Join(usr.HomeDir, "/.ssh/id_rsa")))
	output := remoteCmdOutput(username, hostname, privateKey, fmt.Sprintf("git --git-dir=%s rev-parse %s", e.RepoPath, e.Branch))

	return output
}

func getYAMLString(n yaml.Node, key string) string {
	return n.(yaml.Map)[key].(yaml.Scalar).String()
}

func parseYAMLEnvironment(m yaml.Node) Environment {
	e := Environment{}
	for k, v := range m.(yaml.Map) {
		e.Name = k
		e.Branch = getYAMLString(v, "branch")
		e.RepoPath = getYAMLString(v, "repo_path")
		e.Deploy = getYAMLString(v, "deploy")
		for _, v := range v.(yaml.Map)["hosts"].(yaml.List) {
			h := Host{URI: v.(yaml.Scalar).String()}
			e.Hosts = append(e.Hosts, h)
		}
	}
	return e
}

func parseYAML() (allProjects []Project, deployUser string) {
	config, err := yaml.ReadFile(configFile)
	if err != nil {
		log.Fatal(err)
	}
	deployUser, err = config.Get("deploy_user")
	if err != nil {
		log.Fatal("config.yml is missing deploy_user: " + err.Error())
	}
	configRoot, _ := config.Root.(yaml.Map)
	projects, _ := configRoot["projects"].(yaml.List)
	allProjects = []Project{}
	for _, p := range projects {
		for _, v := range p.(yaml.Map) {
			proj := Project{Name: getYAMLString(v, "project_name"),
				GitHubURL: getYAMLString(v, "github_url"),
				RepoName:  getYAMLString(v, "repo_name"),
				RepoOwner: getYAMLString(v, "repo_owner")}
			for _, v := range v.(yaml.Map)["environments"].(yaml.List) {
				proj.Environments = append(proj.Environments, parseYAMLEnvironment(v))
			}
			allProjects = append(allProjects, proj)
		}
	}
	return allProjects, deployUser
}

func getCommit(wg *sync.WaitGroup, projects []Project, env Environment, host Host, deployUser string, i, j, k int) {
	defer wg.Done()
	host.LatestCommit = string(latestDeployedCommit(deployUser, host.URI+":"+sshPort, env))
	host.LatestCommit = strings.Trim(host.LatestCommit, "\n\r")
	projects[i].Environments[j].Hosts[k] = host
}

//  Get the most recent commit hash on a given branch from GitHub
func getLatestGitHubCommit(wg *sync.WaitGroup, projects []Project, environment Environment, c *github.Client, repoOwner, repoName string, i, j int) {
	defer wg.Done()
	opts := &github.CommitsListOptions{SHA: environment.Branch}
	commits, _, err := c.Repositories.ListCommits(repoOwner, repoName, opts)
	if err != nil {
		log.Panic(err)
	}
	environment.LatestGitHubCommit = *commits[0].SHA
	projects[i].Environments[j] = environment
}

func retrieveCommits(projects []Project, deployUser string) []Project {
	// define a wait group to wait for all goroutines to finish
	var wg sync.WaitGroup
	githubToken := os.Getenv("GITHUB_API_TOKEN")
	t := &oauth.Transport{
		Token: &oauth.Token{AccessToken: githubToken},
	}
	client := github.NewClient(t.Client())
	for i, project := range projects {
		for j, environment := range project.Environments {
			for k, host := range environment.Hosts {
				// start a goroutine for SSH-ing on to the machine
				wg.Add(1)
				go getCommit(&wg, projects, environment, host, deployUser, i, j, k)
			}
			wg.Add(1)
			go getLatestGitHubCommit(&wg, projects, environment, client, project.RepoOwner, project.RepoName, i, j)
		}
	}
	// wait for goroutines to finish
	wg.Wait()
	return projects
}

func HomeHandler(w http.ResponseWriter, r *http.Request) {
	var deployUser string
	projects, deployUser = parseYAML()
	projects = retrieveCommits(projects, deployUser)
	t, err := template.New("index.html").Funcs(template.FuncMap{"eq": eq}).ParseFiles("templates/index.html")
	if err != nil {
		log.Panic(err)
	}
	err = t.Execute(w, map[string]interface{}{"Projects": projects})
	if err != nil {
		log.Panic(err)
	}
}

func DeployHandler(w http.ResponseWriter, r *http.Request) {
	var command []string
	p := r.FormValue("project")
	env := r.FormValue("environment")
	for i, project := range projects {
		if project.Name == p {
			for j, environment := range project.Environments {
				if environment.Name == env {
					command = strings.Split(projects[i].Environments[j].Deploy, " ")
				}
			}
		}
	}
	var out bytes.Buffer
	cmd := exec.Command(command[0], command[1:]...)
	cmd.Stdout = &out
	err := cmd.Run()
	if err != nil {
		log.Fatal(err)
	}
	fmt.Println(out.String())
}

func main() {
	r := mux.NewRouter()
	r.HandleFunc("/", HomeHandler)
	r.HandleFunc("/deploy", DeployHandler)
	fmt.Println("Running on localhost:" + port)
	log.Fatal(http.ListenAndServe(":"+port, r))
}
<|endoftext|>