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

package gocql

import (
	"crypto/tls"
	"crypto/x509"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"math/rand"
	"net"
	"sync"
	"time"
)

// interface to implement to receive the host information
type SetHosts interface {
	SetHosts(hosts []*HostInfo)
}

// interface to implement to receive the partitioner value
type SetPartitioner interface {
	SetPartitioner(partitioner string)
}

func setupTLSConfig(sslOpts *SslOptions) (*tls.Config, error) {
	// ca cert is optional
	if sslOpts.CaPath != "" {
		if sslOpts.RootCAs == nil {
			sslOpts.RootCAs = x509.NewCertPool()
		}

		pem, err := ioutil.ReadFile(sslOpts.CaPath)
		if err != nil {
			return nil, fmt.Errorf("connectionpool: unable to open CA certs: %v", err)
		}

		if !sslOpts.RootCAs.AppendCertsFromPEM(pem) {
			return nil, errors.New("connectionpool: failed parsing or CA certs")
		}
	}

	if sslOpts.CertPath != "" || sslOpts.KeyPath != "" {
		mycert, err := tls.LoadX509KeyPair(sslOpts.CertPath, sslOpts.KeyPath)
		if err != nil {
			return nil, fmt.Errorf("connectionpool: unable to load X509 key pair: %v", err)
		}
		sslOpts.Certificates = append(sslOpts.Certificates, mycert)
	}

	sslOpts.InsecureSkipVerify = !sslOpts.EnableHostVerification

	return &sslOpts.Config, nil
}

type policyConnPool struct {
	session *Session

	port     int
	numConns int
	keyspace string

	mu            sync.RWMutex
	hostPolicy    HostSelectionPolicy
	connPolicy    func() ConnSelectionPolicy
	hostConnPools map[string]*hostConnPool

	endpoints []string
}

func connConfig(session *Session) (*ConnConfig, error) {
	cfg := session.cfg

	var (
		err       error
		tlsConfig *tls.Config
	)

	// TODO(zariel): move tls config setup into session init.
	if cfg.SslOpts != nil {
		tlsConfig, err = setupTLSConfig(cfg.SslOpts)
		if err != nil {
			return nil, err
		}
	}

	return &ConnConfig{
		ProtoVersion:  cfg.ProtoVersion,
		CQLVersion:    cfg.CQLVersion,
		Timeout:       cfg.Timeout,
		Compressor:    cfg.Compressor,
		Authenticator: cfg.Authenticator,
		Keepalive:     cfg.SocketKeepalive,
		tlsConfig:     tlsConfig,
	}, nil
}

func newPolicyConnPool(session *Session, hostPolicy HostSelectionPolicy,
	connPolicy func() ConnSelectionPolicy) *policyConnPool {

	// create the pool
	pool := &policyConnPool{
		session:       session,
		port:          session.cfg.Port,
		numConns:      session.cfg.NumConns,
		keyspace:      session.cfg.Keyspace,
		hostPolicy:    hostPolicy,
		connPolicy:    connPolicy,
		hostConnPools: map[string]*hostConnPool{},
	}

	pool.endpoints = make([]string, len(session.cfg.Hosts))
	copy(pool.endpoints, session.cfg.Hosts)

	return pool
}

func (p *policyConnPool) SetHosts(hosts []*HostInfo) {
	p.mu.Lock()
	defer p.mu.Unlock()

	toRemove := make(map[string]struct{})
	for addr := range p.hostConnPools {
		toRemove[addr] = struct{}{}
	}

	// TODO connect to hosts in parallel, but wait for pools to be
	// created before returning
	for _, host := range hosts {
		pool, exists := p.hostConnPools[host.Peer()]
		if !exists && host.IsUp() {
			// create a connection pool for the host
			pool = newHostConnPool(
				p.session,
				host,
				p.port,
				p.numConns,
				p.keyspace,
				p.connPolicy(),
			)
			p.hostConnPools[host.Peer()] = pool
		} else {
			// still have this host, so don't remove it
			delete(toRemove, host.Peer())
		}
	}

	for addr := range toRemove {
		pool := p.hostConnPools[addr]
		delete(p.hostConnPools, addr)
		pool.Close()
	}

	// update the policy
	p.hostPolicy.SetHosts(hosts)
}

func (p *policyConnPool) SetPartitioner(partitioner string) {
	p.hostPolicy.SetPartitioner(partitioner)
}

func (p *policyConnPool) Size() int {
	p.mu.RLock()
	count := 0
	for _, pool := range p.hostConnPools {
		count += pool.Size()
	}
	p.mu.RUnlock()

	return count
}

func (p *policyConnPool) Pick(qry *Query) (SelectedHost, *Conn) {
	nextHost := p.hostPolicy.Pick(qry)

	var (
		host SelectedHost
		conn *Conn
	)

	p.mu.RLock()
	defer p.mu.RUnlock()
	for conn == nil {
		host = nextHost()
		if host == nil {
			break
		} else if host.Info() == nil {
			panic(fmt.Sprintf("policy %T returned no host info: %+v", p.hostPolicy, host))
		}

		pool, ok := p.hostConnPools[host.Info().Peer()]
		if !ok {
			continue
		}

		conn = pool.Pick(qry)
	}
	return host, conn
}

func (p *policyConnPool) Close() {
	p.mu.Lock()
	defer p.mu.Unlock()

	// remove the hosts from the policy
	p.hostPolicy.SetHosts(nil)

	// close the pools
	for addr, pool := range p.hostConnPools {
		delete(p.hostConnPools, addr)
		pool.Close()
	}
}

func (p *policyConnPool) addHost(host *HostInfo) {
	p.mu.Lock()
	defer p.mu.Unlock()

	pool, ok := p.hostConnPools[host.Peer()]
	if ok {
		go pool.fill()
		return
	}

	pool = newHostConnPool(
		p.session,
		host,
		host.Port(),
		p.numConns,
		p.keyspace,
		p.connPolicy(),
	)

	p.hostConnPools[host.Peer()] = pool
}

func (p *policyConnPool) removeHost(addr string) {
	p.hostPolicy.RemoveHost(addr)
	p.mu.Lock()

	pool, ok := p.hostConnPools[addr]
	if !ok {
		p.mu.Unlock()
		return
	}

	delete(p.hostConnPools, addr)
	p.mu.Unlock()

	pool.Close()
}

func (p *policyConnPool) hostUp(host *HostInfo) {
	// TODO(zariel): have a set of up hosts and down hosts, we can internally
	// detect down hosts, then try to reconnect to them.
	p.addHost(host)
}

func (p *policyConnPool) hostDown(addr string) {
	// TODO(zariel): mark host as down so we can try to connect to it later, for
	// now just treat it has removed.
	p.removeHost(addr)
}

// hostConnPool is a connection pool for a single host.
// Connection selection is based on a provided ConnSelectionPolicy
type hostConnPool struct {
	session  *Session
	host     *HostInfo
	port     int
	addr     string
	size     int
	keyspace string
	policy   ConnSelectionPolicy
	// protection for conns, closed, filling
	mu      sync.RWMutex
	conns   []*Conn
	closed  bool
	filling bool
}

func newHostConnPool(session *Session, host *HostInfo, port, size int,
	keyspace string, policy ConnSelectionPolicy) *hostConnPool {

	pool := &hostConnPool{
		session:  session,
		host:     host,
		port:     port,
		addr:     JoinHostPort(host.Peer(), port),
		size:     size,
		keyspace: keyspace,
		policy:   policy,
		conns:    make([]*Conn, 0, size),
		filling:  false,
		closed:   false,
	}

	// fill the pool with the initial connections before returning
	pool.fill()

	return pool
}

// Pick a connection from this connection pool for the given query.
func (pool *hostConnPool) Pick(qry *Query) *Conn {
	pool.mu.RLock()
	if pool.closed {
		pool.mu.RUnlock()
		return nil
	}

	size := len(pool.conns)
	pool.mu.RUnlock()

	if size < pool.size {
		// try to fill the pool
		go pool.fill()

		if size == 0 {
			return nil
		}
	}

	return pool.policy.Pick(qry)
}

//Size returns the number of connections currently active in the pool
func (pool *hostConnPool) Size() int {
	pool.mu.RLock()
	defer pool.mu.RUnlock()

	return len(pool.conns)
}

//Close the connection pool
func (pool *hostConnPool) Close() {
	pool.mu.Lock()
	defer pool.mu.Unlock()

	if pool.closed {
		return
	}
	pool.closed = true

	// drain, but don't wait
	go pool.drain()
}

// Fill the connection pool
func (pool *hostConnPool) fill() {
	pool.mu.RLock()
	// avoid filling a closed pool, or concurrent filling
	if pool.closed || pool.filling {
		pool.mu.RUnlock()
		return
	}

	// determine the filling work to be done
	startCount := len(pool.conns)
	fillCount := pool.size - startCount

	// avoid filling a full (or overfull) pool
	if fillCount <= 0 {
		pool.mu.RUnlock()
		return
	}

	// switch from read to write lock
	pool.mu.RUnlock()
	pool.mu.Lock()

	// double check everything since the lock was released
	startCount = len(pool.conns)
	fillCount = pool.size - startCount
	if pool.closed || pool.filling || fillCount <= 0 {
		// looks like another goroutine already beat this
		// goroutine to the filling
		pool.mu.Unlock()
		return
	}

	// ok fill the pool
	pool.filling = true

	// allow others to access the pool while filling
	pool.mu.Unlock()
	// only this goroutine should make calls to fill/empty the pool at this
	// point until after this routine or its subordinates calls
	// fillingStopped

	// fill only the first connection synchronously
	if startCount == 0 {
		err := pool.connect()
		pool.logConnectErr(err)

		if err != nil {
			// probably unreachable host
			go pool.fillingStopped()

			// this is calle with the connetion pool mutex held, this call will
			// then recursivly try to lock it again. FIXME
			go pool.session.handleNodeDown(net.ParseIP(pool.host.Peer()), pool.port)
			return
		}

		// filled one
		fillCount--

		// connect all connections to this host in sync
		for fillCount > 0 {
			err := pool.connect()
			pool.logConnectErr(err)

			// decrement, even on error
			fillCount--
		}

		go pool.fillingStopped()
		return
	}

	// fill the rest of the pool asynchronously
	go func() {
		for fillCount > 0 {
			err := pool.connect()
			pool.logConnectErr(err)

			// decrement, even on error
			fillCount--
		}

		// mark the end of filling
		pool.fillingStopped()
	}()
}

func (pool *hostConnPool) logConnectErr(err error) {
	if opErr, ok := err.(*net.OpError); ok && (opErr.Op == "dial" || opErr.Op == "read") {
		// connection refused
		// these are typical during a node outage so avoid log spam.
	} else if err != nil {
		// unexpected error
		log.Printf("error: failed to connect to %s due to error: %v", pool.addr, err)
	}
}

// transition back to a not-filling state.
func (pool *hostConnPool) fillingStopped() {
	// wait for some time to avoid back-to-back filling
	// this provides some time between failed attempts
	// to fill the pool for the host to recover
	time.Sleep(time.Duration(rand.Int31n(100)+31) * time.Millisecond)

	pool.mu.Lock()
	pool.filling = false
	pool.mu.Unlock()
}

// create a new connection to the host and add it to the pool
func (pool *hostConnPool) connect() error {
	// try to connect
	conn, err := pool.session.connect(pool.addr, pool)
	if err != nil {
		return err
	}

	if pool.keyspace != "" {
		// set the keyspace
		if err := conn.UseKeyspace(pool.keyspace); err != nil {
			conn.Close()
			return err
		}
	}

	// add the Conn to the pool
	pool.mu.Lock()
	defer pool.mu.Unlock()

	if pool.closed {
		conn.Close()
		return nil
	}

	pool.conns = append(pool.conns, conn)

	conns := make([]*Conn, len(pool.conns))
	copy(conns, pool.conns)
	pool.policy.SetConns(conns)

	return nil
}

// handle any error from a Conn
func (pool *hostConnPool) HandleError(conn *Conn, err error, closed bool) {
	if !closed {
		// still an open connection, so continue using it
		return
	}

	// TODO: track the number of errors per host and detect when a host is dead,
	// then also have something which can detect when a host comes back.
	pool.mu.Lock()
	defer pool.mu.Unlock()

	if pool.closed {
		// pool closed
		return
	}

	// find the connection index
	for i, candidate := range pool.conns {
		if candidate == conn {
			// remove the connection, not preserving order
			pool.conns[i], pool.conns = pool.conns[len(pool.conns)-1], pool.conns[:len(pool.conns)-1]

			// update the policy
			conns := make([]*Conn, len(pool.conns))
			copy(conns, pool.conns)
			pool.policy.SetConns(conns)

			// lost a connection, so fill the pool
			go pool.fill()
			break
		}
	}
}

// removes and closes all connections from the pool
func (pool *hostConnPool) drain() {
	pool.mu.Lock()
	defer pool.mu.Unlock()

	// empty the pool
	conns := pool.conns
	pool.conns = pool.conns[:0:0]

	// update the policy
	pool.policy.SetConns(nil)

	// close the connections
	for _, conn := range conns {
		conn.Close()
	}
}
<commit_msg>update the host policy from the connection pool<commit_after>// Copyright (c) 2012 The gocql Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package gocql

import (
	"crypto/tls"
	"crypto/x509"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"math/rand"
	"net"
	"sync"
	"time"
)

// interface to implement to receive the host information
type SetHosts interface {
	SetHosts(hosts []*HostInfo)
}

// interface to implement to receive the partitioner value
type SetPartitioner interface {
	SetPartitioner(partitioner string)
}

func setupTLSConfig(sslOpts *SslOptions) (*tls.Config, error) {
	// ca cert is optional
	if sslOpts.CaPath != "" {
		if sslOpts.RootCAs == nil {
			sslOpts.RootCAs = x509.NewCertPool()
		}

		pem, err := ioutil.ReadFile(sslOpts.CaPath)
		if err != nil {
			return nil, fmt.Errorf("connectionpool: unable to open CA certs: %v", err)
		}

		if !sslOpts.RootCAs.AppendCertsFromPEM(pem) {
			return nil, errors.New("connectionpool: failed parsing or CA certs")
		}
	}

	if sslOpts.CertPath != "" || sslOpts.KeyPath != "" {
		mycert, err := tls.LoadX509KeyPair(sslOpts.CertPath, sslOpts.KeyPath)
		if err != nil {
			return nil, fmt.Errorf("connectionpool: unable to load X509 key pair: %v", err)
		}
		sslOpts.Certificates = append(sslOpts.Certificates, mycert)
	}

	sslOpts.InsecureSkipVerify = !sslOpts.EnableHostVerification

	return &sslOpts.Config, nil
}

type policyConnPool struct {
	session *Session

	port     int
	numConns int
	keyspace string

	mu            sync.RWMutex
	hostPolicy    HostSelectionPolicy
	connPolicy    func() ConnSelectionPolicy
	hostConnPools map[string]*hostConnPool

	endpoints []string
}

func connConfig(session *Session) (*ConnConfig, error) {
	cfg := session.cfg

	var (
		err       error
		tlsConfig *tls.Config
	)

	// TODO(zariel): move tls config setup into session init.
	if cfg.SslOpts != nil {
		tlsConfig, err = setupTLSConfig(cfg.SslOpts)
		if err != nil {
			return nil, err
		}
	}

	return &ConnConfig{
		ProtoVersion:  cfg.ProtoVersion,
		CQLVersion:    cfg.CQLVersion,
		Timeout:       cfg.Timeout,
		Compressor:    cfg.Compressor,
		Authenticator: cfg.Authenticator,
		Keepalive:     cfg.SocketKeepalive,
		tlsConfig:     tlsConfig,
	}, nil
}

func newPolicyConnPool(session *Session, hostPolicy HostSelectionPolicy,
	connPolicy func() ConnSelectionPolicy) *policyConnPool {

	// create the pool
	pool := &policyConnPool{
		session:       session,
		port:          session.cfg.Port,
		numConns:      session.cfg.NumConns,
		keyspace:      session.cfg.Keyspace,
		hostPolicy:    hostPolicy,
		connPolicy:    connPolicy,
		hostConnPools: map[string]*hostConnPool{},
	}

	pool.endpoints = make([]string, len(session.cfg.Hosts))
	copy(pool.endpoints, session.cfg.Hosts)

	return pool
}

func (p *policyConnPool) SetHosts(hosts []*HostInfo) {
	p.mu.Lock()
	defer p.mu.Unlock()

	toRemove := make(map[string]struct{})
	for addr := range p.hostConnPools {
		toRemove[addr] = struct{}{}
	}

	// TODO connect to hosts in parallel, but wait for pools to be
	// created before returning
	for _, host := range hosts {
		pool, exists := p.hostConnPools[host.Peer()]
		if !exists && host.IsUp() {
			// create a connection pool for the host
			pool = newHostConnPool(
				p.session,
				host,
				p.port,
				p.numConns,
				p.keyspace,
				p.connPolicy(),
			)
			p.hostConnPools[host.Peer()] = pool
		} else {
			// still have this host, so don't remove it
			delete(toRemove, host.Peer())
		}
	}

	for addr := range toRemove {
		pool := p.hostConnPools[addr]
		delete(p.hostConnPools, addr)
		pool.Close()
	}

	// update the policy
	p.hostPolicy.SetHosts(hosts)
}

func (p *policyConnPool) SetPartitioner(partitioner string) {
	p.hostPolicy.SetPartitioner(partitioner)
}

func (p *policyConnPool) Size() int {
	p.mu.RLock()
	count := 0
	for _, pool := range p.hostConnPools {
		count += pool.Size()
	}
	p.mu.RUnlock()

	return count
}

func (p *policyConnPool) Pick(qry *Query) (SelectedHost, *Conn) {
	nextHost := p.hostPolicy.Pick(qry)

	var (
		host SelectedHost
		conn *Conn
	)

	p.mu.RLock()
	defer p.mu.RUnlock()
	for conn == nil {
		host = nextHost()
		if host == nil {
			break
		} else if host.Info() == nil {
			panic(fmt.Sprintf("policy %T returned no host info: %+v", p.hostPolicy, host))
		}

		pool, ok := p.hostConnPools[host.Info().Peer()]
		if !ok {
			continue
		}

		conn = pool.Pick(qry)
	}
	return host, conn
}

func (p *policyConnPool) Close() {
	p.mu.Lock()
	defer p.mu.Unlock()

	// remove the hosts from the policy
	p.hostPolicy.SetHosts(nil)

	// close the pools
	for addr, pool := range p.hostConnPools {
		delete(p.hostConnPools, addr)
		pool.Close()
	}
}

func (p *policyConnPool) addHost(host *HostInfo) {
	p.mu.Lock()
	defer p.mu.Unlock()

	pool, ok := p.hostConnPools[host.Peer()]
	if ok {
		go pool.fill()
		return
	}

	pool = newHostConnPool(
		p.session,
		host,
		host.Port(),
		p.numConns,
		p.keyspace,
		p.connPolicy(),
	)

	p.hostConnPools[host.Peer()] = pool

	// update policy
	// TODO: policy should not have conns, it should have hosts and return a host
	// iter which the pool will use to serve conns
	p.hostPolicy.AddHost(host)
	p.mu.Unlock()
}

func (p *policyConnPool) removeHost(addr string) {
	p.hostPolicy.RemoveHost(addr)
	p.mu.Lock()

	pool, ok := p.hostConnPools[addr]
	if !ok {
		p.mu.Unlock()
		return
	}

	delete(p.hostConnPools, addr)
	p.mu.Unlock()

	pool.Close()
}

func (p *policyConnPool) hostUp(host *HostInfo) {
	// TODO(zariel): have a set of up hosts and down hosts, we can internally
	// detect down hosts, then try to reconnect to them.
	p.addHost(host)
}

func (p *policyConnPool) hostDown(addr string) {
	// TODO(zariel): mark host as down so we can try to connect to it later, for
	// now just treat it has removed.
	p.removeHost(addr)
}

// hostConnPool is a connection pool for a single host.
// Connection selection is based on a provided ConnSelectionPolicy
type hostConnPool struct {
	session  *Session
	host     *HostInfo
	port     int
	addr     string
	size     int
	keyspace string
	policy   ConnSelectionPolicy
	// protection for conns, closed, filling
	mu      sync.RWMutex
	conns   []*Conn
	closed  bool
	filling bool
}

func newHostConnPool(session *Session, host *HostInfo, port, size int,
	keyspace string, policy ConnSelectionPolicy) *hostConnPool {

	pool := &hostConnPool{
		session:  session,
		host:     host,
		port:     port,
		addr:     JoinHostPort(host.Peer(), port),
		size:     size,
		keyspace: keyspace,
		policy:   policy,
		conns:    make([]*Conn, 0, size),
		filling:  false,
		closed:   false,
	}

	// fill the pool with the initial connections before returning
	pool.fill()

	return pool
}

// Pick a connection from this connection pool for the given query.
func (pool *hostConnPool) Pick(qry *Query) *Conn {
	pool.mu.RLock()
	if pool.closed {
		pool.mu.RUnlock()
		return nil
	}

	size := len(pool.conns)
	pool.mu.RUnlock()

	if size < pool.size {
		// try to fill the pool
		go pool.fill()

		if size == 0 {
			return nil
		}
	}

	return pool.policy.Pick(qry)
}

//Size returns the number of connections currently active in the pool
func (pool *hostConnPool) Size() int {
	pool.mu.RLock()
	defer pool.mu.RUnlock()

	return len(pool.conns)
}

//Close the connection pool
func (pool *hostConnPool) Close() {
	pool.mu.Lock()
	defer pool.mu.Unlock()

	if pool.closed {
		return
	}
	pool.closed = true

	// drain, but don't wait
	go pool.drain()
}

// Fill the connection pool
func (pool *hostConnPool) fill() {
	pool.mu.RLock()
	// avoid filling a closed pool, or concurrent filling
	if pool.closed || pool.filling {
		pool.mu.RUnlock()
		return
	}

	// determine the filling work to be done
	startCount := len(pool.conns)
	fillCount := pool.size - startCount

	// avoid filling a full (or overfull) pool
	if fillCount <= 0 {
		pool.mu.RUnlock()
		return
	}

	// switch from read to write lock
	pool.mu.RUnlock()
	pool.mu.Lock()

	// double check everything since the lock was released
	startCount = len(pool.conns)
	fillCount = pool.size - startCount
	if pool.closed || pool.filling || fillCount <= 0 {
		// looks like another goroutine already beat this
		// goroutine to the filling
		pool.mu.Unlock()
		return
	}

	// ok fill the pool
	pool.filling = true

	// allow others to access the pool while filling
	pool.mu.Unlock()
	// only this goroutine should make calls to fill/empty the pool at this
	// point until after this routine or its subordinates calls
	// fillingStopped

	// fill only the first connection synchronously
	if startCount == 0 {
		err := pool.connect()
		pool.logConnectErr(err)

		if err != nil {
			// probably unreachable host
			go pool.fillingStopped()

			// this is calle with the connetion pool mutex held, this call will
			// then recursivly try to lock it again. FIXME
			go pool.session.handleNodeDown(net.ParseIP(pool.host.Peer()), pool.port)
			return
		}

		// filled one
		fillCount--

		// connect all connections to this host in sync
		for fillCount > 0 {
			err := pool.connect()
			pool.logConnectErr(err)

			// decrement, even on error
			fillCount--
		}

		go pool.fillingStopped()
		return
	}

	// fill the rest of the pool asynchronously
	go func() {
		for fillCount > 0 {
			err := pool.connect()
			pool.logConnectErr(err)

			// decrement, even on error
			fillCount--
		}

		// mark the end of filling
		pool.fillingStopped()
	}()
}

func (pool *hostConnPool) logConnectErr(err error) {
	if opErr, ok := err.(*net.OpError); ok && (opErr.Op == "dial" || opErr.Op == "read") {
		// connection refused
		// these are typical during a node outage so avoid log spam.
	} else if err != nil {
		// unexpected error
		log.Printf("error: failed to connect to %s due to error: %v", pool.addr, err)
	}
}

// transition back to a not-filling state.
func (pool *hostConnPool) fillingStopped() {
	// wait for some time to avoid back-to-back filling
	// this provides some time between failed attempts
	// to fill the pool for the host to recover
	time.Sleep(time.Duration(rand.Int31n(100)+31) * time.Millisecond)

	pool.mu.Lock()
	pool.filling = false
	pool.mu.Unlock()
}

// create a new connection to the host and add it to the pool
func (pool *hostConnPool) connect() error {
	// try to connect
	conn, err := pool.session.connect(pool.addr, pool)
	if err != nil {
		return err
	}

	if pool.keyspace != "" {
		// set the keyspace
		if err := conn.UseKeyspace(pool.keyspace); err != nil {
			conn.Close()
			return err
		}
	}

	// add the Conn to the pool
	pool.mu.Lock()
	defer pool.mu.Unlock()

	if pool.closed {
		conn.Close()
		return nil
	}

	pool.conns = append(pool.conns, conn)

	conns := make([]*Conn, len(pool.conns))
	copy(conns, pool.conns)
	pool.policy.SetConns(conns)

	return nil
}

// handle any error from a Conn
func (pool *hostConnPool) HandleError(conn *Conn, err error, closed bool) {
	if !closed {
		// still an open connection, so continue using it
		return
	}

	// TODO: track the number of errors per host and detect when a host is dead,
	// then also have something which can detect when a host comes back.
	pool.mu.Lock()
	defer pool.mu.Unlock()

	if pool.closed {
		// pool closed
		return
	}

	// find the connection index
	for i, candidate := range pool.conns {
		if candidate == conn {
			// remove the connection, not preserving order
			pool.conns[i], pool.conns = pool.conns[len(pool.conns)-1], pool.conns[:len(pool.conns)-1]

			// update the policy
			conns := make([]*Conn, len(pool.conns))
			copy(conns, pool.conns)
			pool.policy.SetConns(conns)

			// lost a connection, so fill the pool
			go pool.fill()
			break
		}
	}
}

// removes and closes all connections from the pool
func (pool *hostConnPool) drain() {
	pool.mu.Lock()
	defer pool.mu.Unlock()

	// empty the pool
	conns := pool.conns
	pool.conns = pool.conns[:0:0]

	// update the policy
	pool.policy.SetConns(nil)

	// close the connections
	for _, conn := range conns {
		conn.Close()
	}
}
<|endoftext|>


<commit_before>package executor_test

import (
	. "github.com/cloudfoundry-incubator/executor"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Container", func() {
	Describe("HasTags", func() {
		var container Container

		Context("when tags are nil", func() {
			It("returns true if requested tags are nil", func() {
				Ω(container.HasTags(nil)).Should(BeTrue())
			})

			It("returns false if requested tags are not nil", func() {
				Ω(container.HasTags(Tags{"a": "b"})).Should(BeFalse())
			})
		})

		Context("when tags are not nil", func() {
			BeforeEach(func() {
				container = Container{
					Tags: Tags{"a": "b"},
				}
			})

			It("returns true when found", func() {
				Ω(container.HasTags(Tags{"a": "b"})).Should(BeTrue())
			})

			It("returns false when nil", func() {
				Ω(container.HasTags(nil)).Should(BeFalse())
			})

			It("returns false when not found", func() {
				Ω(container.HasTags(Tags{"a": "c"})).Should(BeFalse())
			})
		})
	})
})
<commit_msg>fix order-dependant test<commit_after>package executor_test

import (
	. "github.com/cloudfoundry-incubator/executor"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Container", func() {
	Describe("HasTags", func() {
		var container Container

		Context("when tags are nil", func() {
			BeforeEach(func() {
				container = Container{
					Tags: nil,
				}
			})

			It("returns true if requested tags are nil", func() {
				Ω(container.HasTags(nil)).Should(BeTrue())
			})

			It("returns false if requested tags are not nil", func() {
				Ω(container.HasTags(Tags{"a": "b"})).Should(BeFalse())
			})
		})

		Context("when tags are not nil", func() {
			BeforeEach(func() {
				container = Container{
					Tags: Tags{"a": "b"},
				}
			})

			It("returns true when found", func() {
				Ω(container.HasTags(Tags{"a": "b"})).Should(BeTrue())
			})

			It("returns false when nil", func() {
				Ω(container.HasTags(nil)).Should(BeFalse())
			})

			It("returns false when not found", func() {
				Ω(container.HasTags(Tags{"a": "c"})).Should(BeFalse())
			})
		})
	})
})
<|endoftext|>


<commit_before>// Copyright 2014 struktur AG. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package phoenix

import (
	"log"
	"os"
	"testing"
)

func newTestContainer(name, version string) Container {
	return &container{
		name,
		version,
		nil,
		nil,
		nil,
	}
}

func Test_Container_Name_UsesTheGivenValue(t *testing.T) {
	name := "spreed-app"
	container := newTestContainer(name, "")
	if actualName := container.Name(); actualName != name {
		t.Errorf("Expected app name to be '%s', but was '%s'", name, actualName)
	}
}

func Test_Container_Name_DefaultstoAppIfUnset(t *testing.T) {
	container := newTestContainer("", "")
	if expected, actual := "app", container.Name(); expected != actual {
		t.Errorf("Expected app name to be '%s', but was '%s'", expected, actual)
	}
}

func Test_Container_Version_UsesTheGivenValue(t *testing.T) {
	version := "0.9.4b1"
	container := newTestContainer("", version)
	if actualVersion := container.Version(); actualVersion != version {
		t.Errorf("Expected app name to be '%s', but was '%s'", version, actualVersion)
	}
}

func Test_Container_Name_DefaultstoUnreleasedIfUnset(t *testing.T) {
	container := newTestContainer("", "")
	if expected, actual := "unreleased", container.Version(); expected != actual {
		t.Errorf("Expected app version to be '%s', but was '%s'", expected, actual)
	}
}

func Test_Container_Syslog(t *testing.T) {
	logFilename := "syslog"
	container, err := newContainer("test", "", &logFilename, nil)
	if err != nil {
		t.Fatalf("Could not create container: '%s'", err)
	}
	if err := container.Close(); err != nil {
		t.Errorf("Closing the syslog container returned '%s'", err)
	}
	log.Println("Testing")
	fp, err := os.Open(logFilename)
	if fp != nil {
		fp.Close()
	}
	if err == nil {
		t.Errorf("Logging created a file '%s' but should have logged to syslog", logFilename)
	} else if !os.IsNotExist(err) {
		t.Errorf("Logfile '%s' could not be opened but should not exist: '%s'", logFilename, err)
	}
}
<commit_msg>Log warning and skip test if connection to syslog can not be established. This can happen when running on Travis infrastructure.<commit_after>// Copyright 2014 struktur AG. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package phoenix

import (
	"fmt"
	"log"
	"os"
	"testing"
)

func newTestContainer(name, version string) Container {
	return &container{
		name,
		version,
		nil,
		nil,
		nil,
	}
}

func Test_Container_Name_UsesTheGivenValue(t *testing.T) {
	name := "spreed-app"
	container := newTestContainer(name, "")
	if actualName := container.Name(); actualName != name {
		t.Errorf("Expected app name to be '%s', but was '%s'", name, actualName)
	}
}

func Test_Container_Name_DefaultstoAppIfUnset(t *testing.T) {
	container := newTestContainer("", "")
	if expected, actual := "app", container.Name(); expected != actual {
		t.Errorf("Expected app name to be '%s', but was '%s'", expected, actual)
	}
}

func Test_Container_Version_UsesTheGivenValue(t *testing.T) {
	version := "0.9.4b1"
	container := newTestContainer("", version)
	if actualVersion := container.Version(); actualVersion != version {
		t.Errorf("Expected app name to be '%s', but was '%s'", version, actualVersion)
	}
}

func Test_Container_Name_DefaultstoUnreleasedIfUnset(t *testing.T) {
	container := newTestContainer("", "")
	if expected, actual := "unreleased", container.Version(); expected != actual {
		t.Errorf("Expected app version to be '%s', but was '%s'", expected, actual)
	}
}

func Test_Container_Syslog(t *testing.T) {
	logFilename := "syslog"
	container, err := newContainer("test", "", &logFilename, nil)
	if err != nil {
		fmt.Fprintf(os.Stderr, "WARNING: Could not create container, test will be skipped: '%v'\n", err)
		return
	}
	if err := container.Close(); err != nil {
		t.Errorf("Closing the syslog container returned '%s'", err)
	}
	log.Println("Testing")
	fp, err := os.Open(logFilename)
	if fp != nil {
		fp.Close()
	}
	if err == nil {
		t.Errorf("Logging created a file '%s' but should have logged to syslog", logFilename)
	} else if !os.IsNotExist(err) {
		t.Errorf("Logfile '%s' could not be opened but should not exist: '%s'", logFilename, err)
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package install

import (
	"k8s.io/apimachinery/pkg/apimachinery/announced"
	"k8s.io/apimachinery/pkg/apimachinery/registered"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/kubernetes/pkg/apis/admissionregistration"
	"k8s.io/kubernetes/pkg/apis/admissionregistration/v1alpha1"
)

// Install registers the API group and adds types to a scheme
func Install(groupFactoryRegistry announced.APIGroupFactoryRegistry, registry *registered.APIRegistrationManager, scheme *runtime.Scheme) {
	if err := announced.NewGroupMetaFactory(
		&announced.GroupMetaFactoryArgs{
			GroupName:                  admissionregistration.GroupName,
			RootScopedKinds:            sets.NewString("InitializerConfiguration", "ExternalAdmissionHookConfiguration"),
			VersionPreferenceOrder:     []string{v1alpha1.SchemeGroupVersion.Version},
			ImportPrefix:               "k8s.io/kubernetes/pkg/apis/admissionregistration",
			AddInternalObjectsToScheme: admissionregistration.AddToScheme,
		},
		announced.VersionToSchemeFunc{
			v1alpha1.SchemeGroupVersion.Version: v1alpha1.AddToScheme,
		},
	).Announce(groupFactoryRegistry).RegisterAndEnable(registry, scheme); err != nil {
		panic(err)
	}
}
<commit_msg>serve the api in kube-apiserver<commit_after>/*
Copyright 2017 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package install

import (
	"k8s.io/apimachinery/pkg/apimachinery/announced"
	"k8s.io/apimachinery/pkg/apimachinery/registered"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/util/sets"
	"k8s.io/kubernetes/pkg/api"
	"k8s.io/kubernetes/pkg/apis/admissionregistration"
	"k8s.io/kubernetes/pkg/apis/admissionregistration/v1alpha1"
)

func init() {
	Install(api.GroupFactoryRegistry, api.Registry, api.Scheme)
}

// Install registers the API group and adds types to a scheme
func Install(groupFactoryRegistry announced.APIGroupFactoryRegistry, registry *registered.APIRegistrationManager, scheme *runtime.Scheme) {
	if err := announced.NewGroupMetaFactory(
		&announced.GroupMetaFactoryArgs{
			GroupName:                  admissionregistration.GroupName,
			RootScopedKinds:            sets.NewString("InitializerConfiguration", "ExternalAdmissionHookConfiguration"),
			VersionPreferenceOrder:     []string{v1alpha1.SchemeGroupVersion.Version},
			ImportPrefix:               "k8s.io/kubernetes/pkg/apis/admissionregistration",
			AddInternalObjectsToScheme: admissionregistration.AddToScheme,
		},
		announced.VersionToSchemeFunc{
			v1alpha1.SchemeGroupVersion.Version: v1alpha1.AddToScheme,
		},
	).Announce(groupFactoryRegistry).RegisterAndEnable(registry, scheme); err != nil {
		panic(err)
	}
}
<|endoftext|>


<commit_before>package state

import (
	"fmt"
	"sync"

	"github.com/Dataman-Cloud/swan/src/mesosproto/mesos"
	"github.com/Dataman-Cloud/swan/src/types"

	"github.com/Sirupsen/logrus"
)

//  TASK_STAGING = 6;  // Initial state. Framework status updates should not use.
//  TASK_STARTING = 0; // The task is being launched by the executor.
//  TASK_RUNNING = 1;
//  TASK_KILLING = 8;  // The task is being killed by the executor.
//  TASK_FINISHED = 2; // TERMINAL: The task finished successfully.
//  TASK_FAILED = 3;   // TERMINAL: The task failed to finish successfully.
//  TASK_KILLED = 4;   // TERMINAL: The task was killed by the executor.
//  TASK_ERROR = 7;    // TERMINAL: The task description contains an error.
//  TASK_LOST = 5;     // TERMINAL: The task failed but can be rescheduled.
//  TASK_DROPPED = 9;  // TERMINAL.
//  TASK_UNREACHABLE = 10;
//  TASK_GONE = 11;    // TERMINAL.
//  TASK_GONE_BY_OPERATOR = 12;
//  TASK_UNKNOWN = 13;

const (
	SLOT_STATE_PENDING_OFFER = "slot_task_pending_offer"
	SLOT_STATE_PENDING_KILL  = "slot_task_pending_killed"

	SLOT_STATE_TASK_STAGING          = "slot_task_staging"
	SLOT_STATE_TASK_STARTING         = "slot_task_starting"
	SLOT_STATE_TASK_RUNNING          = "slot_task_running"
	SLOT_STATE_TASK_KILLING          = "slot_task_killing"
	SLOT_STATE_TASK_FINISHED         = "slot_task_finished"
	SLOT_STATE_TASK_FAILED           = "slot_task_failed"
	SLOT_STATE_TASK_KILLED           = "slot_task_killed"
	SLOT_STATE_TASK_ERROR            = "slot_task_error"
	SLOT_STATE_TASK_LOST             = "slot_task_lost"
	SLOT_STATE_TASK_DROPPED          = "slot_task_dropped"
	SLOT_STATE_TASK_UNREACHABLE      = "slot_task_unreachable"
	SLOT_STATE_TASK_GONE             = "slot_task_gone"
	SLOT_STATE_TASK_GONE_BY_OPERATOR = "slot_task_gone_by_operator"
	SLOT_STATE_TASK_UNKNOWN          = "slot_task_unknown"
)

type SlotStateCallbackFuncs func(slot *Slot)

type Slot struct {
	Index           int
	Id              string
	App             *App
	Version         *types.Version
	State           string
	StatesCallbacks map[string][]SlotStateCallbackFuncs

	CurrentTask *Task
	TaskHistory []*Task

	OfferId       string
	AgentId       string
	Ip            string
	AgentHostName string

	resourceReservationLock sync.Mutex

	MarkForDeletion      bool
	MarkForRollingUpdate bool
}

func NewSlot(app *App, version *types.Version, index int) *Slot {
	slot := &Slot{
		Index:       index,
		App:         app,
		Version:     version,
		TaskHistory: make([]*Task, 0),
		Id:          fmt.Sprintf("%d-%s-%s-%s", index, version.AppId, version.RunAs, app.Scheduler.ClusterId), // should be app.AppId

		resourceReservationLock: sync.Mutex{},

		MarkForRollingUpdate: false,
		MarkForDeletion:      false,
		StatesCallbacks:      make(map[string][]SlotStateCallbackFuncs),
	}

	slot.DispatchNewTask(slot.Version)

	return slot
}

// kill task doesn't need cleanup slot from app.Slots
func (slot *Slot) KillTask() {
	slot.SetState(SLOT_STATE_PENDING_KILL)
	slot.CurrentTask.Kill()
}

// kill task and make slot sweeped after successfully kill task
func (slot *Slot) Kill() {
	slot.MarkForDeletion = true
	slot.SetState(SLOT_STATE_PENDING_KILL)

	slot.CurrentTask.Kill()
}

func (slot *Slot) Archive() {
	slot.TaskHistory = append(slot.TaskHistory, slot.CurrentTask)
}

func (slot *Slot) DispatchNewTask(version *types.Version) {
	slot.Version = version
	slot.CurrentTask = NewTask(slot.App, slot.Version, slot)
	slot.SetState(SLOT_STATE_PENDING_OFFER)
	slot.App.OfferAllocatorRef.PutSlotBackToPendingQueue(slot)

	//onSlotRunning := func(slot *Slot) {
	//if slot.MarkForRollingUpdate {
	//}
	//}
	//slot.RegisterStateCallbacks(SLOT_STATE_TASK_RUNNING, onSlotRunning)

}

func (slot *Slot) Update(version *types.Version) {
	logrus.Infof("update slot %s with version ID %s", slot.Id, version.ID)

	slot.MarkForRollingUpdate = true // mark as in progress of rolling update

	onSlotFinished := func(slot *Slot) {
		logrus.Infof("onSlotFinished func")
		slot.Archive()
		slot.DispatchNewTask(version)
	}
	slot.RegisterStateCallbacks(SLOT_STATE_TASK_FINISHED, onSlotFinished)

	slot.KillTask() // kill task but doesn't clean slot
}

func (slot *Slot) TestOfferMatch(ow *OfferWrapper) bool {
	return ow.CpuRemain() > slot.Version.Cpus &&
		ow.MemRemain() > slot.Version.Mem &&
		ow.DiskRemain() > slot.Version.Disk
}

func (slot *Slot) ReserveOfferAndPrepareTaskInfo(ow *OfferWrapper) (*OfferWrapper, *mesos.TaskInfo) {
	slot.resourceReservationLock.Lock()
	defer slot.resourceReservationLock.Unlock()

	ow.CpusUsed += slot.Version.Cpus
	ow.MemUsed += slot.Version.Mem
	ow.DiskUsed += slot.Version.Disk

	return ow, slot.CurrentTask.PrepareTaskInfo(ow.Offer)
}

func (slot *Slot) Resources() []*mesos.Resource {
	var resources = []*mesos.Resource{}

	if slot.Version.Cpus > 0 {
		resources = append(resources, createScalarResource("cpus", slot.Version.Cpus))
	}

	if slot.Version.Mem > 0 {
		resources = append(resources, createScalarResource("mem", slot.Version.Cpus))
	}

	if slot.Version.Disk > 0 {
		resources = append(resources, createScalarResource("disk", slot.Version.Disk))
	}

	return resources
}

func (slot *Slot) StateIs(state string) bool {
	return slot.State == state
}

func (slot *Slot) SetState(state string) error {
	slot.State = state

	// handle callback
	if len(slot.StatesCallbacks[slot.State]) > 0 {
		for _, cb := range slot.StatesCallbacks[slot.State] {
			cb(slot)
		}
	}
	slot.RemoveStateCallbacks(slot.State)

	logrus.Infof("setting state for slot %s to %s", slot.Id, slot.State)

	switch slot.State {
	case SLOT_STATE_PENDING_KILL:
	case SLOT_STATE_TASK_KILLED:
	case SLOT_STATE_TASK_FINISHED:
	case SLOT_STATE_TASK_RUNNING:
	case SLOT_STATE_TASK_FAILED:
		// restart if needed
	default:
	}

	// skip app invalidation if slot state is not mesos driven
	if (slot.State != SLOT_STATE_PENDING_OFFER) ||
		(slot.State != SLOT_STATE_PENDING_KILL) {
		slot.App.InvalidateSlots()
	}
	return nil
}

func (slot *Slot) AppendStateCallbacks(state string, callback SlotStateCallbackFuncs) {
	slot.StatesCallbacks[state] = append(slot.StatesCallbacks[state], callback)
}

// empty callback queue first
func (slot *Slot) RegisterStateCallbacks(state string, callback SlotStateCallbackFuncs) {
	slot.RemoveStateCallbacks(state)
	slot.StatesCallbacks[state] = append(slot.StatesCallbacks[state], callback)
}

func (slot *Slot) RemoveStateCallbacks(state string) {
	slot.StatesCallbacks[state] = make([]SlotStateCallbackFuncs, 0)
}
<commit_msg>remove unsed code<commit_after>package state

import (
	"fmt"
	"sync"

	"github.com/Dataman-Cloud/swan/src/mesosproto/mesos"
	"github.com/Dataman-Cloud/swan/src/types"

	"github.com/Sirupsen/logrus"
)

//  TASK_STAGING = 6;  // Initial state. Framework status updates should not use.
//  TASK_STARTING = 0; // The task is being launched by the executor.
//  TASK_RUNNING = 1;
//  TASK_KILLING = 8;  // The task is being killed by the executor.
//  TASK_FINISHED = 2; // TERMINAL: The task finished successfully.
//  TASK_FAILED = 3;   // TERMINAL: The task failed to finish successfully.
//  TASK_KILLED = 4;   // TERMINAL: The task was killed by the executor.
//  TASK_ERROR = 7;    // TERMINAL: The task description contains an error.
//  TASK_LOST = 5;     // TERMINAL: The task failed but can be rescheduled.
//  TASK_DROPPED = 9;  // TERMINAL.
//  TASK_UNREACHABLE = 10;
//  TASK_GONE = 11;    // TERMINAL.
//  TASK_GONE_BY_OPERATOR = 12;
//  TASK_UNKNOWN = 13;

const (
	SLOT_STATE_PENDING_OFFER = "slot_task_pending_offer"
	SLOT_STATE_PENDING_KILL  = "slot_task_pending_killed"

	SLOT_STATE_TASK_STAGING          = "slot_task_staging"
	SLOT_STATE_TASK_STARTING         = "slot_task_starting"
	SLOT_STATE_TASK_RUNNING          = "slot_task_running"
	SLOT_STATE_TASK_KILLING          = "slot_task_killing"
	SLOT_STATE_TASK_FINISHED         = "slot_task_finished"
	SLOT_STATE_TASK_FAILED           = "slot_task_failed"
	SLOT_STATE_TASK_KILLED           = "slot_task_killed"
	SLOT_STATE_TASK_ERROR            = "slot_task_error"
	SLOT_STATE_TASK_LOST             = "slot_task_lost"
	SLOT_STATE_TASK_DROPPED          = "slot_task_dropped"
	SLOT_STATE_TASK_UNREACHABLE      = "slot_task_unreachable"
	SLOT_STATE_TASK_GONE             = "slot_task_gone"
	SLOT_STATE_TASK_GONE_BY_OPERATOR = "slot_task_gone_by_operator"
	SLOT_STATE_TASK_UNKNOWN          = "slot_task_unknown"
)

type SlotStateCallbackFuncs func(slot *Slot)

type Slot struct {
	Index           int
	Id              string
	App             *App
	Version         *types.Version
	State           string
	StatesCallbacks map[string][]SlotStateCallbackFuncs

	CurrentTask *Task
	TaskHistory []*Task

	OfferId       string
	AgentId       string
	Ip            string
	AgentHostName string

	resourceReservationLock sync.Mutex

	MarkForDeletion      bool
	MarkForRollingUpdate bool
}

func NewSlot(app *App, version *types.Version, index int) *Slot {
	slot := &Slot{
		Index:       index,
		App:         app,
		Version:     version,
		TaskHistory: make([]*Task, 0),
		Id:          fmt.Sprintf("%d-%s-%s-%s", index, version.AppId, version.RunAs, app.Scheduler.ClusterId), // should be app.AppId

		resourceReservationLock: sync.Mutex{},

		MarkForRollingUpdate: false,
		MarkForDeletion:      false,
		StatesCallbacks:      make(map[string][]SlotStateCallbackFuncs),
	}

	slot.DispatchNewTask(slot.Version)

	return slot
}

// kill task doesn't need cleanup slot from app.Slots
func (slot *Slot) KillTask() {
	slot.SetState(SLOT_STATE_PENDING_KILL)
	slot.CurrentTask.Kill()
}

// kill task and make slot sweeped after successfully kill task
func (slot *Slot) Kill() {
	slot.MarkForDeletion = true
	slot.SetState(SLOT_STATE_PENDING_KILL)

	slot.CurrentTask.Kill()
}

func (slot *Slot) Archive() {
	slot.TaskHistory = append(slot.TaskHistory, slot.CurrentTask)
}

func (slot *Slot) DispatchNewTask(version *types.Version) {
	slot.Version = version
	slot.CurrentTask = NewTask(slot.App, slot.Version, slot)
	slot.SetState(SLOT_STATE_PENDING_OFFER)

	slot.App.OfferAllocatorRef.PutSlotBackToPendingQueue(slot)
}

func (slot *Slot) Update(version *types.Version) {
	logrus.Infof("update slot %s with version ID %s", slot.Id, version.ID)

	slot.MarkForRollingUpdate = true // mark as in progress of rolling update

	onSlotFinished := func(slot *Slot) {
		logrus.Infof("onSlotFinished func")
		slot.Archive()
		slot.DispatchNewTask(version)
	}
	slot.RegisterStateCallbacks(SLOT_STATE_TASK_FINISHED, onSlotFinished)

	slot.KillTask() // kill task but doesn't clean slot
}

func (slot *Slot) TestOfferMatch(ow *OfferWrapper) bool {
	return ow.CpuRemain() > slot.Version.Cpus &&
		ow.MemRemain() > slot.Version.Mem &&
		ow.DiskRemain() > slot.Version.Disk
}

func (slot *Slot) ReserveOfferAndPrepareTaskInfo(ow *OfferWrapper) (*OfferWrapper, *mesos.TaskInfo) {
	slot.resourceReservationLock.Lock()
	defer slot.resourceReservationLock.Unlock()

	ow.CpusUsed += slot.Version.Cpus
	ow.MemUsed += slot.Version.Mem
	ow.DiskUsed += slot.Version.Disk

	return ow, slot.CurrentTask.PrepareTaskInfo(ow.Offer)
}

func (slot *Slot) Resources() []*mesos.Resource {
	var resources = []*mesos.Resource{}

	if slot.Version.Cpus > 0 {
		resources = append(resources, createScalarResource("cpus", slot.Version.Cpus))
	}

	if slot.Version.Mem > 0 {
		resources = append(resources, createScalarResource("mem", slot.Version.Cpus))
	}

	if slot.Version.Disk > 0 {
		resources = append(resources, createScalarResource("disk", slot.Version.Disk))
	}

	return resources
}

func (slot *Slot) StateIs(state string) bool {
	return slot.State == state
}

func (slot *Slot) SetState(state string) error {
	slot.State = state

	// handle callback
	if len(slot.StatesCallbacks[slot.State]) > 0 {
		for _, cb := range slot.StatesCallbacks[slot.State] {
			cb(slot)
		}
	}
	slot.RemoveStateCallbacks(slot.State)

	logrus.Infof("setting state for slot %s to %s", slot.Id, slot.State)

	switch slot.State {
	case SLOT_STATE_PENDING_KILL:
	case SLOT_STATE_TASK_KILLED:
	case SLOT_STATE_TASK_FINISHED:
	case SLOT_STATE_TASK_RUNNING:
	case SLOT_STATE_TASK_FAILED:
		// restart if needed
	default:
	}

	// skip app invalidation if slot state is not mesos driven
	if (slot.State != SLOT_STATE_PENDING_OFFER) ||
		(slot.State != SLOT_STATE_PENDING_KILL) {
		slot.App.InvalidateSlots()
	}
	return nil
}

func (slot *Slot) AppendStateCallbacks(state string, callback SlotStateCallbackFuncs) {
	slot.StatesCallbacks[state] = append(slot.StatesCallbacks[state], callback)
}

// empty callback queue first
func (slot *Slot) RegisterStateCallbacks(state string, callback SlotStateCallbackFuncs) {
	slot.RemoveStateCallbacks(state)
	slot.StatesCallbacks[state] = append(slot.StatesCallbacks[state], callback)
}

func (slot *Slot) RemoveStateCallbacks(state string) {
	slot.StatesCallbacks[state] = make([]SlotStateCallbackFuncs, 0)
}
<|endoftext|>


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

package exec

import (
	"bytes"
	"fmt"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"
	"io/ioutil"
	"net"
	"os"
	"time"
)

type Options struct {
	machines []string
	port     string
	user     string
	cmd      string
	key      string
	useAgent bool
}

type executeResult struct {
	result string
	err    error
}

func User(u string) func(*Options) {
	return func(e *Options) {
		e.user = u
	}
}

func Port(p string) func(*Options) {
	return func(e *Options) {
		e.port = p
	}
}

func Cmd(c string) func(*Options) {
	return func(e *Options) {
		e.cmd = c
	}
}

func Machines(m []string) func(*Options) {
	return func(e *Options) {
		e.machines = m
	}
}

func Key(k string) func(*Options) {
	return func(e *Options) {
		e.key = k
	}
}

func UseAgent(u bool) func(*Options) {
	return func(e *Options) {
		e.useAgent = u
	}
}

func makeSigner(keyname string) (signer ssh.Signer, err error) {
	fp, err := os.Open(keyname)
	if err != nil {
		return
	}
	defer fp.Close()

	buf, _ := ioutil.ReadAll(fp)
	signer, _ = ssh.ParsePrivateKey(buf)
	return
}

func makeKeyring(key string, useAgent bool) ssh.AuthMethod {
	signers := []ssh.Signer{}

	if useAgent == true {
		aConn, _ := net.Dial("unix", os.Getenv("SSH_AUTH_SOCK"))
		sshAgent := agent.NewClient(aConn)
		aSigners, _ := sshAgent.Signers()
		for _, signer := range aSigners {
			signers = append(signers, signer)
		}
	}

	keys := []string{key}

	for _, keyname := range keys {
		signer, err := makeSigner(keyname)
		if err == nil {
			signers = append(signers, signer)
		}
	}
	return ssh.PublicKeys(signers...)
}

func executeCmd(opt Options, hostname string, config *ssh.ClientConfig) executeResult {

	conn, err := ssh.Dial("tcp", hostname+":"+opt.port, config)

	if err != nil {
		return executeResult{result: "",
			err: err}
	}

	session, _ := conn.NewSession()
	defer session.Close()

	var stdoutBuf bytes.Buffer
	session.Stdout = &stdoutBuf
	err = session.Run(opt.cmd)

	return executeResult{result: hostname + ":\n" + stdoutBuf.String(),
		err: err}
}

func Run(options ...func(*Options)) bool {
	// in 5 seconds the message will come to timeout channel
	opt := Options{}
	for _, option := range options {
		option(&opt)
	}

	timeout := time.After(5 * time.Second)
	results := make(chan executeResult, len(opt.machines))

	config := &ssh.ClientConfig{
		User: opt.user,
		Auth: []ssh.AuthMethod{
			makeKeyring(opt.key, opt.useAgent),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
	}

	for _, m := range opt.machines {
		go func(hostname string) {
			// we’ll write results into the buffered channel of strings
			results <- executeCmd(opt, hostname, config)
		}(m)
	}

	for i := 0; i < len(opt.machines); i++ {
		select {
		case res := <-results:
			if res.err == nil {
				fmt.Print(res.result)
			} else {
				fmt.Println(res.err)
				return false
			}
		case <-timeout:
			fmt.Println("Timed out!")
			return false
		}
	}
	return true
}
<commit_msg>Minor update to comment (#18)<commit_after>// Copyright © 2017 Ricardo Aravena <raravena@branch.io>
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package exec

import (
	"bytes"
	"fmt"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"
	"io/ioutil"
	"net"
	"os"
	"time"
)

type Options struct {
	machines []string
	port     string
	user     string
	cmd      string
	key      string
	useAgent bool
}

type executeResult struct {
	result string
	err    error
}

func User(u string) func(*Options) {
	return func(e *Options) {
		e.user = u
	}
}

func Port(p string) func(*Options) {
	return func(e *Options) {
		e.port = p
	}
}

func Cmd(c string) func(*Options) {
	return func(e *Options) {
		e.cmd = c
	}
}

func Machines(m []string) func(*Options) {
	return func(e *Options) {
		e.machines = m
	}
}

func Key(k string) func(*Options) {
	return func(e *Options) {
		e.key = k
	}
}

func UseAgent(u bool) func(*Options) {
	return func(e *Options) {
		e.useAgent = u
	}
}

func makeSigner(keyname string) (signer ssh.Signer, err error) {
	fp, err := os.Open(keyname)
	if err != nil {
		return
	}
	defer fp.Close()

	buf, _ := ioutil.ReadAll(fp)
	signer, _ = ssh.ParsePrivateKey(buf)
	return
}

func makeKeyring(key string, useAgent bool) ssh.AuthMethod {
	signers := []ssh.Signer{}

	if useAgent == true {
		aConn, _ := net.Dial("unix", os.Getenv("SSH_AUTH_SOCK"))
		sshAgent := agent.NewClient(aConn)
		aSigners, _ := sshAgent.Signers()
		for _, signer := range aSigners {
			signers = append(signers, signer)
		}
	}

	keys := []string{key}

	for _, keyname := range keys {
		signer, err := makeSigner(keyname)
		if err == nil {
			signers = append(signers, signer)
		}
	}
	return ssh.PublicKeys(signers...)
}

func executeCmd(opt Options, hostname string, config *ssh.ClientConfig) executeResult {

	conn, err := ssh.Dial("tcp", hostname+":"+opt.port, config)

	if err != nil {
		return executeResult{result: "",
			err: err}
	}

	session, _ := conn.NewSession()
	defer session.Close()

	var stdoutBuf bytes.Buffer
	session.Stdout = &stdoutBuf
	err = session.Run(opt.cmd)

	return executeResult{result: hostname + ":\n" + stdoutBuf.String(),
		err: err}
}

func Run(options ...func(*Options)) bool {
	opt := Options{}
	for _, option := range options {
		option(&opt)
	}

	// in 5 seconds the message will come to timeout channel
	timeout := time.After(5 * time.Second)
	results := make(chan executeResult, len(opt.machines))

	config := &ssh.ClientConfig{
		User: opt.user,
		Auth: []ssh.AuthMethod{
			makeKeyring(opt.key, opt.useAgent),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
	}

	for _, m := range opt.machines {
		go func(hostname string) {
			// we’ll write results into the buffered channel of strings
			results <- executeCmd(opt, hostname, config)
		}(m)
	}

	for i := 0; i < len(opt.machines); i++ {
		select {
		case res := <-results:
			if res.err == nil {
				fmt.Print(res.result)
			} else {
				fmt.Println(res.err)
				return false
			}
		case <-timeout:
			fmt.Println("Timed out!")
			return false
		}
	}
	return true
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io"
	"os"

	"github.com/apex/apex/function"
	"github.com/apex/apex/project"
	"github.com/apex/log"
	"github.com/apex/log/handlers/text"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/lambda"
	"github.com/segmentio/go-prompt"
	"github.com/tj/docopt"
)

var version = "0.1.0"

const usage = `
  Usage:
    apex deploy [options] [<name>...]
    apex delete [options] [<name>...]
    apex invoke [options] <name> [--async] [-v]
    apex rollback [options] <name> [<version>]
    apex build [options] <name>
    apex list [options]
    apex -h | --help
    apex --version

  Options:
    -l, --log-level level   Log severity level [default: info]
    -a, --async             Async invocation
    -C, --chdir path        Working directory
    -y, --yes               Automatic yes to prompts
    -h, --help              Output help information
    -v, --verbose           Output verbose logs
    -V, --version           Output version

  Examples:
    Deploy all functions
    $ apex deploy

    Deploy specific functions
    $ apex deploy foo bar

    Delete all functions
    $ apex delete

    Delete specified functions
    $ apex delete foo bar

    Invoke a function with input json
    $ apex invoke foo < request.json

    Rollback a function to the previous version
    $ apex rollback foo

    Rollback a function to the specified version
    $ apex rollback bar 3

    Deploy functions in a different project
    $ apex deploy -C ~/dev/myapp

    Build zip output for a function
    $ apex build foo > /tmp/out.zip
`

func main() {
	args, err := docopt.Parse(usage, nil, true, version, false)
	if err != nil {
		log.Fatalf("error: %s", err)
	}

	log.SetHandler(text.New(os.Stderr))

	if l, err := log.ParseLevel(args["--log-level"].(string)); err == nil {
		log.SetLevel(l)
	}

	project := &project.Project{
		Service: lambda.New(session.New(aws.NewConfig())),
		Log:     log.Log,
		Path:    ".",
	}

	if dir, ok := args["--chdir"].(string); ok {
		if err := os.Chdir(dir); err != nil {
			log.Fatalf("error: %s", err)
		}
	}

	if err := project.Open(); err != nil {
		log.Fatalf("error opening project: %s", err)
	}

	switch {
	case args["list"].(bool):
		list(project)
	case args["deploy"].(bool):
		deploy(project, args["<name>"].([]string))
	case args["delete"].(bool):
		delete(project, args["<name>"].([]string), args["--yes"].(bool))
	case args["invoke"].(bool):
		invoke(project, args["<name>"].([]string), args["--verbose"].(bool), args["--async"].(bool))
	case args["rollback"].(bool):
		rollback(project, args["<name>"].([]string), args["<version>"])
	case args["build"].(bool):
		build(project, args["<name>"].([]string))
	}
}

// list functions.
func list(project *project.Project) {
	// TODO(tj): more informative output
	fmt.Println()
	for _, fn := range project.Functions {
		fmt.Printf("  - %s (%s)\n", fn.Name, fn.Runtime)
	}
	fmt.Println()
}

// invoke reads request json from stdin and outputs the responses.
func invoke(project *project.Project, name []string, verbose, async bool) {
	dec := json.NewDecoder(os.Stdin)
	kind := function.RequestResponse

	if async {
		kind = function.Event
	}

	fn, err := project.FunctionByName(name[0])
	if err != nil {
		log.Fatalf("error: %s", err)
	}

	for {
		var v struct {
			Event   interface{}
			Context interface{}
		}

		err := dec.Decode(&v)

		if err == io.EOF {
			break
		}

		if err != nil {
			log.Fatalf("error: %s", err)
		}

		reply, logs, err := fn.Invoke(v.Event, v.Context, kind)
		if err != nil {
			log.Fatalf("error: %s", err)
		}

		// TODO(tj) rename flag to --with-logs or --logs
		if verbose {
			io.Copy(os.Stderr, logs)
		}

		io.Copy(os.Stdout, reply)
		fmt.Fprintf(os.Stdout, "\n")
	}
}

// deploy changes.
func deploy(project *project.Project, names []string) {
	var err error

	if len(names) == 0 {
		names, err = project.FunctionNames()
	}

	if err != nil {
		log.Fatalf("error: %s", err)
	}

	if err := project.DeployAndClean(names); err != nil {
		log.Fatalf("error: %s", err)
	}
}

// delete the resources.
func delete(project *project.Project, names []string, force bool) {
	var err error

	if len(names) == 0 {
		names, err = project.FunctionNames()
	}

	if err != nil {
		log.Fatalf("error: %s", err)
	}

	if !force && len(names) > 1 {
		fmt.Printf("The following will be deleted:\n\n")
		for _, name := range names {
			fmt.Printf("  - %s\n", name)
		}
		fmt.Printf("\n")
	}

	if !force && !prompt.Confirm("Are you sure? (yes/no)") {
		return
	}

	if err := project.Delete(names); err != nil {
		log.Fatalf("error: %s", err)
	}
}

// rollback the function.
func rollback(project *project.Project, name []string, version interface{}) {
	fn, err := project.FunctionByName(name[0])
	if err != nil {
		log.Fatalf("error: %s", err)
	}

	if version == nil {
		err = fn.Rollback()
	} else {
		err = fn.Rollback(version.(string))
	}

	if err != nil {
		log.Fatalf("error: %s", err)
	}
}

// build outputs the generated archive to stdout.
func build(project *project.Project, name []string) {
	fn, err := project.FunctionByName(name[0])
	if err != nil {
		log.Fatalf("error: %s", err)
	}

	zip, err := fn.Zip()
	if err != nil {
		log.Fatalf("error: %s", err)
	}

	_, err = io.Copy(os.Stdout, zip)
	if err != nil {
		log.Fatalf("error: %s", err)
	}
}
<commit_msg>tweak a few comments<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io"
	"os"

	"github.com/apex/apex/function"
	"github.com/apex/apex/project"
	"github.com/apex/log"
	"github.com/apex/log/handlers/text"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/lambda"
	"github.com/segmentio/go-prompt"
	"github.com/tj/docopt"
)

var version = "0.1.0"

const usage = `
  Usage:
    apex deploy [options] [<name>...]
    apex delete [options] [<name>...]
    apex invoke [options] <name> [--async] [-v]
    apex rollback [options] <name> [<version>]
    apex build [options] <name>
    apex list [options]
    apex -h | --help
    apex --version

  Options:
    -l, --log-level level   Log severity level [default: info]
    -a, --async             Async invocation
    -C, --chdir path        Working directory
    -y, --yes               Automatic yes to prompts
    -h, --help              Output help information
    -v, --verbose           Output verbose logs
    -V, --version           Output version

  Examples:
    Deploy all functions
    $ apex deploy

    Deploy specific functions
    $ apex deploy foo bar

    Delete all functions
    $ apex delete

    Delete specified functions
    $ apex delete foo bar

    Invoke a function with input json
    $ apex invoke foo < request.json

    Rollback a function to the previous version
    $ apex rollback foo

    Rollback a function to the specified version
    $ apex rollback bar 3

    Deploy functions in a different project
    $ apex deploy -C ~/dev/myapp

    Build zip output for a function
    $ apex build foo > /tmp/out.zip
`

func main() {
	args, err := docopt.Parse(usage, nil, true, version, false)
	if err != nil {
		log.Fatalf("error: %s", err)
	}

	log.SetHandler(text.New(os.Stderr))

	if l, err := log.ParseLevel(args["--log-level"].(string)); err == nil {
		log.SetLevel(l)
	}

	project := &project.Project{
		Service: lambda.New(session.New(aws.NewConfig())),
		Log:     log.Log,
		Path:    ".",
	}

	if dir, ok := args["--chdir"].(string); ok {
		if err := os.Chdir(dir); err != nil {
			log.Fatalf("error: %s", err)
		}
	}

	if err := project.Open(); err != nil {
		log.Fatalf("error opening project: %s", err)
	}

	switch {
	case args["list"].(bool):
		list(project)
	case args["deploy"].(bool):
		deploy(project, args["<name>"].([]string))
	case args["delete"].(bool):
		delete(project, args["<name>"].([]string), args["--yes"].(bool))
	case args["invoke"].(bool):
		invoke(project, args["<name>"].([]string), args["--verbose"].(bool), args["--async"].(bool))
	case args["rollback"].(bool):
		rollback(project, args["<name>"].([]string), args["<version>"])
	case args["build"].(bool):
		build(project, args["<name>"].([]string))
	}
}

// list functions.
func list(project *project.Project) {
	// TODO(tj): more informative output
	fmt.Println()
	for _, fn := range project.Functions {
		fmt.Printf("  - %s (%s)\n", fn.Name, fn.Runtime)
	}
	fmt.Println()
}

// invoke reads request json from stdin and outputs the responses.
func invoke(project *project.Project, name []string, verbose, async bool) {
	dec := json.NewDecoder(os.Stdin)
	kind := function.RequestResponse

	if async {
		kind = function.Event
	}

	fn, err := project.FunctionByName(name[0])
	if err != nil {
		log.Fatalf("error: %s", err)
	}

	for {
		var v struct {
			Event   interface{}
			Context interface{}
		}

		err := dec.Decode(&v)

		if err == io.EOF {
			break
		}

		if err != nil {
			log.Fatalf("error: %s", err)
		}

		reply, logs, err := fn.Invoke(v.Event, v.Context, kind)
		if err != nil {
			log.Fatalf("error: %s", err)
		}

		// TODO(tj) rename flag to --with-logs or --logs
		if verbose {
			io.Copy(os.Stderr, logs)
		}

		io.Copy(os.Stdout, reply)
		fmt.Fprintf(os.Stdout, "\n")
	}
}

// deploy code and config changes.
func deploy(project *project.Project, names []string) {
	var err error

	if len(names) == 0 {
		names, err = project.FunctionNames()
	}

	if err != nil {
		log.Fatalf("error: %s", err)
	}

	if err := project.DeployAndClean(names); err != nil {
		log.Fatalf("error: %s", err)
	}
}

// delete the functions.
func delete(project *project.Project, names []string, force bool) {
	var err error

	if len(names) == 0 {
		names, err = project.FunctionNames()
	}

	if err != nil {
		log.Fatalf("error: %s", err)
	}

	if !force && len(names) > 1 {
		fmt.Printf("The following will be deleted:\n\n")
		for _, name := range names {
			fmt.Printf("  - %s\n", name)
		}
		fmt.Printf("\n")
	}

	if !force && !prompt.Confirm("Are you sure? (yes/no)") {
		return
	}

	if err := project.Delete(names); err != nil {
		log.Fatalf("error: %s", err)
	}
}

// rollback the function with optional version.
func rollback(project *project.Project, name []string, version interface{}) {
	fn, err := project.FunctionByName(name[0])
	if err != nil {
		log.Fatalf("error: %s", err)
	}

	if version == nil {
		err = fn.Rollback()
	} else {
		err = fn.Rollback(version.(string))
	}

	if err != nil {
		log.Fatalf("error: %s", err)
	}
}

// build outputs the generated archive to stdout.
func build(project *project.Project, name []string) {
	fn, err := project.FunctionByName(name[0])
	if err != nil {
		log.Fatalf("error: %s", err)
	}

	zip, err := fn.Zip()
	if err != nil {
		log.Fatalf("error: %s", err)
	}

	_, err = io.Copy(os.Stdout, zip)
	if err != nil {
		log.Fatalf("error: %s", err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/signal"
	"sort"
	"strings"
	"syscall"

	"github.com/ianremmler/clac"
	"golang.org/x/crypto/ssh/terminal"
)

const usageStr = `usage:

Interactive:  clac [-i <input>]
Command line: clac [-x | -p <precision>] <input>

Command line mode requires input from arguments (without -i) and/or stdin.
`

var (
	// flags
	doInitStack      = false
	doHexOut         = false
	cliPrec     uint = 12

	trm         *terminal.Terminal
	oldTrmState *terminal.State
	lastErr     error
	cl          = clac.New()
	cmdList     = []string{}
)

var cmdMap = map[string]func() error{
	"neg":    cl.Neg,
	"abs":    cl.Abs,
	"inv":    cl.Inv,
	"+":      cl.Add,
	"-":      cl.Sub,
	"*":      cl.Mul,
	"/":      cl.Div,
	"div":    cl.IntDiv,
	"mod":    cl.Mod,
	"exp":    cl.Exp,
	"^":      cl.Pow,
	"2^":     cl.Pow2,
	"10^":    cl.Pow10,
	"logn":   cl.LogN,
	"ln":     cl.Ln,
	"log":    cl.Log,
	"lg":     cl.Lg,
	"sqrt":   cl.Sqrt,
	"!":      cl.Factorial,
	"comb":   cl.Comb,
	"perm":   cl.Perm,
	"sin":    cl.Sin,
	"cos":    cl.Cos,
	"tan":    cl.Tan,
	"asin":   cl.Asin,
	"acos":   cl.Acos,
	"atan":   cl.Atan,
	"atan2":  cl.Atan2,
	"dtor":   cl.DegToRad,
	"rtod":   cl.RadToDeg,
	"rtop":   cl.RectToPolar,
	"ptor":   cl.PolarToRect,
	"floor":  cl.Floor,
	"ceil":   cl.Ceil,
	"trunc":  cl.Trunc,
	"and":    cl.And,
	"or":     cl.Or,
	"xor":    cl.Xor,
	"not":    cl.Not,
	"andn":   cl.AndN,
	"orn":    cl.OrN,
	"xorn":   cl.XorN,
	"sum":    cl.Sum,
	"avg":    cl.Avg,
	"drop":   cl.Drop,
	"dropn":  cl.DropN,
	"dropr":  cl.DropR,
	"dup":    cl.Dup,
	"dupn":   cl.DupN,
	"dupr":   cl.DupR,
	"pick":   cl.Pick,
	"swap":   cl.Swap,
	"depth":  cl.Depth,
	"min":    cl.Min,
	"max":    cl.Max,
	"minn":   cl.MinN,
	"maxn":   cl.MaxN,
	"rot":    cl.Rot,
	"rotr":   cl.RotR,
	"unrot":  cl.Unrot,
	"unrotr": cl.UnrotR,
	"mag":    cl.Mag,
	"hyp":    cl.Hypot,
	"dot":    cl.Dot,
	"dot3":   cl.Dot3,
	"cross":  cl.Cross,
	"pi":     func() error { return cl.Push(clac.Pi) },
	"e":      func() error { return cl.Push(clac.E) },
	"phi":    func() error { return cl.Push(clac.Phi) },
}

var interactiveCmdMap = map[string]func() error{
	"undo":  cl.Undo,
	"redo":  cl.Redo,
	"clear": cl.Clear,
	"reset": func() error { return cl.Reset() },
	"help":  func() error { help(); return clac.ErrNoHistUpdate },
	"quit":  func() error { exit(); return nil },
}

type term struct {
	io.Reader
	io.Writer
}

func init() {
	log.SetFlags(0)
	log.SetPrefix("clac: ")
	flag.BoolVar(&doHexOut, "x", doHexOut,
		"Command line mode: hexidecimal output")
	flag.UintVar(&cliPrec, "p", cliPrec,
		"Command line mode: output precision")
	flag.BoolVar(&doInitStack, "i", doInitStack,
		"Initialize with input from command line arguments")
	flag.Usage = func() {
		fmt.Fprintln(os.Stderr, usageStr)
		flag.PrintDefaults()
	}
}

func main() {
	flag.Parse()
	if !processCmdLine() {
		printCmdLineStack(cl.Stack())
		os.Exit(0)
	}

	interactiveSetup()

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt)
	go func() {
		<-sigChan
		exit()
	}()

	repl()
}

func interactiveSetup() {
	if !terminal.IsTerminal(syscall.Stdin) {
		log.Fatalln("this doesn't look like an interactive terminal")
	}
	var err error
	oldTrmState, err = terminal.MakeRaw(syscall.Stdin)
	if err != nil {
		log.Fatalln(err)
	}
	trm = terminal.NewTerminal(term{os.Stdin, os.Stdout}, "")

	for cmd, fn := range interactiveCmdMap {
		cmdMap[cmd] = fn
	}
	for cmd := range cmdMap {
		cmdList = append(cmdList, cmd)
	}
	sort.Strings(cmdList)
}

func repl() {
	for {
		printStack(cl.Stack())
		input, err := trm.ReadLine()
		lastErr = nil
		if err == io.EOF {
			exit()
		}
		if err != nil {
			continue
		}
		processInput(input, true)
	}
}

func processCmdLine() bool {
	input := ""
	if stat, err := os.Stdin.Stat(); err == nil && stat.Mode()&os.ModeNamedPipe != 0 {
		if pipeInput, err := ioutil.ReadAll(os.Stdin); err == nil {
			input = string(pipeInput)
		}
	}
	if len(flag.Args()) > 0 {
		input += " " + strings.Join(flag.Args(), " ")
	}
	isInteractive := doInitStack || (input == "")
	if isInteractive {
		cl.EnableHistory(isInteractive)
	}
	processInput(string(input), false)
	return isInteractive
}

func printCmdLineStack(stack clac.Stack) {
	if doHexOut {
		clac.SetFormat("%#x")
	} else {
		clac.SetFormat(fmt.Sprintf("%%.%dg", cliPrec))
	}
	for i := range stack {
		val := stack[len(stack)-i-1]
		var err error
		if doHexOut {
			val, err = clac.Trunc(val)
		}
		if err != nil {
			fmt.Print("error")
		} else {
			fmt.Print(clac.Sprint(val))
		}
		if i < len(stack)-1 {
			fmt.Print(" ")
		}
	}
	fmt.Println()
}

func exit() {
	terminal.Restore(syscall.Stdin, oldTrmState)
	fmt.Println()
	os.Exit(0)
}

func help() {
	clearScreen()
	for i := range cmdList {
		fmt.Printf("%-8s", cmdList[i])
		if (i+1)%5 == 0 {
			fmt.Println("\r")
		}
	}
	if len(cmdList)%5 != 0 {
		fmt.Println("\r")
	}
	fmt.Print("\n[Press any key to continue]")
	waitKey()
}

func processInput(input string, isInteractive bool) {
	errorHandler := func(err error) { lastErr = err }
	if !isInteractive {
		errorHandler = func(err error) { log.Println(err) }
	}
	scanner := bufio.NewScanner(strings.NewReader(input))
	scanner.Split(bufio.ScanWords)
	for scanner.Scan() {
		tok := scanner.Text()
		if num, err := clac.ParseNum(tok); err == nil {
			if err = cl.Exec(func() error { return cl.Push(num) }); err != nil {
				errorHandler(fmt.Errorf("push: %s", err))
			}
			continue
		}
		if cmd, ok := cmdMap[tok]; ok {
			if err := cl.Exec(cmd); err != nil {
				errorHandler(fmt.Errorf("%s: %s", tok, err))
			}
			continue
		}
		errorHandler(fmt.Errorf("%s: invalid input", tok))
	}
	if err := scanner.Err(); err != nil {
		errorHandler(err)
	}
}

func printStack(stack clac.Stack) {
	cols, rows, err := terminal.GetSize(syscall.Stdout)
	if err != nil {
		rows = len(stack) + 1
	}
	// ensure sane width
	if cols < 20 {
		cols = 20
	}
	clearScreen()

	dataCols := cols - 4
	hexCols := dataCols / 2
	floatCols := dataCols - hexCols
	floatFmt := fmt.Sprintf("%%%d.%dg", floatCols-1, floatCols-8)
	hexFmt := fmt.Sprintf("%%#%dx", hexCols-3)
	for i := rows - 3; i >= 0; i-- {
		line := fmt.Sprintf("%02d:", i)
		if i < len(stack) {
			clac.SetFormat(floatFmt)
			line += fmt.Sprintf(fmt.Sprintf(" %%%ds", floatCols), clac.Sprint(stack[i]))
			if val, err := clac.Trunc(stack[i]); err == nil {
				clac.SetFormat(hexFmt)
				hexStr := fmt.Sprintf(fmt.Sprintf(" %%%ds", hexCols-1), clac.Sprint(val))
				if len(hexStr) > hexCols {
					hexStr = hexStr[:hexCols-1] + "…"
				}
				line += hexStr
			}
		}
		fmt.Println(line + "\r")
	}
	if lastErr == nil {
		fmt.Println(strings.Repeat("-", cols))
	} else {
		fmt.Println("Error:", lastErr)
	}
	fmt.Print("\r")
}

func clearScreen() {
	fmt.Print("\033[2J\033[H")
}

func waitKey() {
	bufio.NewReader(os.Stdin).ReadByte()
}
<commit_msg>Single char aliases for common commands.<commit_after>package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/signal"
	"sort"
	"strings"
	"syscall"

	"github.com/ianremmler/clac"
	"golang.org/x/crypto/ssh/terminal"
	"robpike.io/ivy/value"
)

const usageStr = `usage:

Interactive:  clac [-i <input>]
Command line: clac [-x | -p <precision>] <input>

Command line mode requires input from arguments (without -i) and/or stdin.
`

var (
	// flags
	doInitStack      = false
	doHexOut         = false
	cliPrec     uint = 12

	trm         *terminal.Terminal
	oldTrmState *terminal.State
	lastErr     error
	cl          = clac.New()
	cmdList     = []string{}
)

var cmdMap = map[string]func() error{
	"neg":    cl.Neg,
	"n":      cl.Neg,
	"abs":    cl.Abs,
	"a":      cl.Abs,
	"inv":    cl.Inv,
	"i":      cl.Inv,
	"+":      cl.Add,
	"-":      cl.Sub,
	"*":      cl.Mul,
	"x":      cl.Mul,
	"/":      cl.Div,
	"div":    cl.IntDiv,
	"%":      cl.Mod,
	"exp":    cl.Exp,
	"^":      cl.Pow,
	"2^":     cl.Pow2,
	"10^":    cl.Pow10,
	"logn":   cl.LogN,
	"ln":     cl.Ln,
	"log":    cl.Log,
	"lg":     cl.Lg,
	"sqrt":   cl.Sqrt,
	"!":      cl.Factorial,
	"comb":   cl.Comb,
	"perm":   cl.Perm,
	"sin":    cl.Sin,
	"cos":    cl.Cos,
	"tan":    cl.Tan,
	"asin":   cl.Asin,
	"acos":   cl.Acos,
	"atan":   cl.Atan,
	"atan2":  cl.Atan2,
	"dtor":   cl.DegToRad,
	"rtod":   cl.RadToDeg,
	"rtop":   cl.RectToPolar,
	"ptor":   cl.PolarToRect,
	"floor":  cl.Floor,
	"ceil":   cl.Ceil,
	"trunc":  cl.Trunc,
	"and":    cl.And,
	"or":     cl.Or,
	"xor":    cl.Xor,
	"not":    cl.Not,
	"andn":   cl.AndN,
	"orn":    cl.OrN,
	"xorn":   cl.XorN,
	"sum":    cl.Sum,
	"avg":    cl.Avg,
	"drop":   cl.Drop,
	"k":      cl.Drop,
	"dropn":  cl.DropN,
	"dropr":  cl.DropR,
	"dup":    cl.Dup,
	"d":      cl.Dup,
	"dupn":   cl.DupN,
	"dupr":   cl.DupR,
	"pick":   cl.Pick,
	"p":      cl.Pick,
	"swap":   cl.Swap,
	"s":      cl.Swap,
	"depth":  cl.Depth,
	"min":    cl.Min,
	"max":    cl.Max,
	"minn":   cl.MinN,
	"maxn":   cl.MaxN,
	"rot":    cl.Rot,
	"r":      cl.Rot,
	"rotr":   cl.RotR,
	"unrot":  cl.Unrot,
	"unrotr": cl.UnrotR,
	"mag":    cl.Mag,
	"hyp":    cl.Hypot,
	"dot":    cl.Dot,
	"dot3":   cl.Dot3,
	"cross":  cl.Cross,
	"pi":     constant(clac.Pi),
	"e":      constant(clac.E),
	"phi":    constant(clac.Phi),
}

var interactiveCmdMap = map[string]func() error{
	"undo":  cl.Undo,
	"u":     cl.Undo,
	"redo":  cl.Redo,
	"r":     cl.Redo,
	"clear": cl.Clear,
	"reset": cl.Reset,
	"help":  help,
	"quit":  exit,
}

type term struct {
	io.Reader
	io.Writer
}

func init() {
	log.SetFlags(0)
	log.SetPrefix("clac: ")
	flag.BoolVar(&doHexOut, "x", doHexOut,
		"Command line mode: hexidecimal output")
	flag.UintVar(&cliPrec, "p", cliPrec,
		"Command line mode: output precision")
	flag.BoolVar(&doInitStack, "i", doInitStack,
		"Initialize with input from command line arguments")
	flag.Usage = func() {
		fmt.Fprintln(os.Stderr, usageStr)
		flag.PrintDefaults()
	}
}

func main() {
	flag.Parse()
	if !processCmdLine() {
		printCmdLineStack(cl.Stack())
		os.Exit(0)
	}

	interactiveSetup()

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, os.Interrupt)
	go func() {
		<-sigChan
		exit()
	}()

	repl()
}

func constant(v value.Value) func() error {
	return func() error { return cl.Push(v) }
}

func interactiveSetup() {
	if !terminal.IsTerminal(syscall.Stdin) {
		log.Fatalln("this doesn't look like an interactive terminal")
	}
	var err error
	oldTrmState, err = terminal.MakeRaw(syscall.Stdin)
	if err != nil {
		log.Fatalln(err)
	}
	trm = terminal.NewTerminal(term{os.Stdin, os.Stdout}, "")

	for cmd, fn := range interactiveCmdMap {
		cmdMap[cmd] = fn
	}
	for cmd := range cmdMap {
		cmdList = append(cmdList, cmd)
	}
	sort.Strings(cmdList)
}

func repl() {
	for {
		printStack(cl.Stack())
		input, err := trm.ReadLine()
		lastErr = nil
		if err == io.EOF {
			exit()
		}
		if err != nil {
			continue
		}
		processInput(input, true)
	}
}

func processCmdLine() bool {
	input := ""
	if stat, err := os.Stdin.Stat(); err == nil && stat.Mode()&os.ModeNamedPipe != 0 {
		if pipeInput, err := ioutil.ReadAll(os.Stdin); err == nil {
			input = string(pipeInput)
		}
	}
	if len(flag.Args()) > 0 {
		input += " " + strings.Join(flag.Args(), " ")
	}
	isInteractive := doInitStack || (input == "")
	if isInteractive {
		cl.EnableHistory(isInteractive)
	}
	processInput(string(input), false)
	return isInteractive
}

func printCmdLineStack(stack clac.Stack) {
	if doHexOut {
		clac.SetFormat("%#x")
	} else {
		clac.SetFormat(fmt.Sprintf("%%.%dg", cliPrec))
	}
	for i := range stack {
		val := stack[len(stack)-i-1]
		var err error
		if doHexOut {
			val, err = clac.Trunc(val)
		}
		if err != nil {
			fmt.Print("error")
		} else {
			fmt.Print(clac.Sprint(val))
		}
		if i < len(stack)-1 {
			fmt.Print(" ")
		}
	}
	fmt.Println()
}

func exit() error {
	terminal.Restore(syscall.Stdin, oldTrmState)
	fmt.Println()
	os.Exit(0)
	return nil
}

func help() error {
	clearScreen()
	for i := range cmdList {
		fmt.Printf("%-8s", cmdList[i])
		if (i+1)%5 == 0 {
			fmt.Println("\r")
		}
	}
	if len(cmdList)%5 != 0 {
		fmt.Println("\r")
	}
	fmt.Print("\n[Press any key to continue]")
	waitKey()
	return clac.ErrNoHistUpdate
}

func processInput(input string, isInteractive bool) {
	errorHandler := func(err error) { lastErr = err }
	if !isInteractive {
		errorHandler = func(err error) { log.Println(err) }
	}
	scanner := bufio.NewScanner(strings.NewReader(input))
	scanner.Split(bufio.ScanWords)
	for scanner.Scan() {
		tok := scanner.Text()
		if num, err := clac.ParseNum(tok); err == nil {
			if err = cl.Exec(func() error { return cl.Push(num) }); err != nil {
				errorHandler(fmt.Errorf("push: %s", err))
			}
			continue
		}
		if cmd, ok := cmdMap[tok]; ok {
			if err := cl.Exec(cmd); err != nil {
				errorHandler(fmt.Errorf("%s: %s", tok, err))
			}
			continue
		}
		errorHandler(fmt.Errorf("%s: invalid input", tok))
	}
	if err := scanner.Err(); err != nil {
		errorHandler(err)
	}
}

func printStack(stack clac.Stack) {
	cols, rows, err := terminal.GetSize(syscall.Stdout)
	if err != nil {
		rows = len(stack) + 1
	}
	// ensure sane width
	if cols < 20 {
		cols = 20
	}
	clearScreen()

	dataCols := cols - 4
	hexCols := dataCols / 2
	floatCols := dataCols - hexCols
	floatFmt := fmt.Sprintf("%%%d.%dg", floatCols-1, floatCols-8)
	hexFmt := fmt.Sprintf("%%#%dx", hexCols-3)
	for i := rows - 3; i >= 0; i-- {
		line := fmt.Sprintf("%02d:", i)
		if i < len(stack) {
			clac.SetFormat(floatFmt)
			line += fmt.Sprintf(fmt.Sprintf(" %%%ds", floatCols), clac.Sprint(stack[i]))
			if val, err := clac.Trunc(stack[i]); err == nil {
				clac.SetFormat(hexFmt)
				hexStr := fmt.Sprintf(fmt.Sprintf(" %%%ds", hexCols-1), clac.Sprint(val))
				if len(hexStr) > hexCols {
					hexStr = hexStr[:hexCols-1] + "…"
				}
				line += hexStr
			}
		}
		fmt.Println(line + "\r")
	}
	if lastErr == nil {
		fmt.Println(strings.Repeat("-", cols))
	} else {
		fmt.Println("Error:", lastErr)
	}
	fmt.Print("\r")
}

func clearScreen() {
	fmt.Print("\033[2J\033[H")
}

func waitKey() {
	bufio.NewReader(os.Stdin).ReadByte()
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"crypto/rand"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"runtime"
	"strconv"
	"strings"
	"syscall"
	"time"

	"code.google.com/p/jra-go/jpu"
)

const intro = `
It's a moonless night, you're lucky. Cutting your way through
the barbed wire fence, you've made your way into the compound.
Your target up ahead, just a 10 meter sprint across open ground:
a service door and the control panel next to it with its slowly
blinking red LED. You concentrate on the red dot and make a
run for it.

Kneeling beneath the control panel, you unscrew the cover and
find it right where they told you to look for it, the DIAG_IO
port. You hook up Tx, Rx, and ground then boot up your portable
terminal. Propping your back against the wall, you settle in
for a little hacking...
`

const help = `
You have successfully connected your terminal to the DIAG_IO port
of the door control computer. You can send the following commands
to its debugger:

	load address byte byte byte ...
		Load the bytes into memory, starting at address.
	dump
		Dump memory
	reg
		See the registers
	reg x y
		Set register x to y
	step
		Run one step
	go
		Run until halt
	quiet
		Turn off instruction logging
	init
		Reinitialize (same as exit and run the program again)
	exit
		Disconnect the leads and scurry back through the fence.

As a shortcut, you can type just the first letter of any command.
`

type logger struct{}

func (l logger) Log(msg string) {
	fmt.Println(msg)
}

const includeEncode = false
const includeAssemble = true

var restart *bool = flag.Bool("restart", false,
	"for internal use only")
var havingFun *bool = flag.Bool("havingFun", true,
	"set this to false if you're not having fun anymore")
var assemble *string
var encode *string

func init() {
	if includeEncode {
		encode = flag.String("encode", "", "string to hide")
	}
	if includeAssemble {
		assemble = flag.String("assemble", "", "program to assemble")
	}
}

func main() {
	// panic handler
	defer func() {
		if r := recover(); r != nil {
			fmt.Println("Error:", r)
		}
	}()

	// override the normal usage to prevent them from easily seeing
	// the havingFun flag (still might see it with strings, good on 'em)
	flag.Usage = func() { fmt.Fprintf(os.Stderr, "Solve this clue. Duh.\n") }
	flag.Parse()

	if includeAssemble && *assemble != "" {
		f, err := os.Open(*assemble)
		if err != nil {
			fmt.Print("Error: ", err)
			os.Exit(1)
		}
		all, err := ioutil.ReadAll(f)
		if err != nil {
			fmt.Print("Error: ", err)
			os.Exit(1)
		}
		res, org := jpu.Assemble(string(all))

		if *assemble == "bigmac.src" {
		fmt.Printf("%#v\n", res)
		} else {
		fmt.Printf("load %d ", org)
		for _, x := range res {
			fmt.Printf("%d ", x)
		}
		fmt.Println()
		}

		fmt.Printf("reg 0 %d\n", org)
		fmt.Println("quiet")
		fmt.Println("go")
		return
	}

	// this feature is only compiled in when includeEncode is set
	// it creates an encrypted version of the answer using a one-time-pad
	// so that strings on the binary can't find the answer.
	// the answer is decrypted before it is loaded into jpu1's RAM
	if includeEncode && *encode != "" {
		in := []byte(*encode)
		mask := make([]byte, len(in))
		text := make([]byte, len(in))

		_, _ = rand.Read(mask)
		for i, _ := range in {
			for {
				text[i] = mask[i] ^ in[i]
				// do not allow zeros in the encrypted results
				if text[i] == 0 {
					mask[i]++
					// go try again
				} else {
					break
				}
			}
		}

		// add terminators
		mask = append(mask, 0)
		text = append(text, 0)

		fmt.Printf("mask := %#v\n", mask)
		fmt.Printf("text := %#v\n", text)
		for i, _ := range in {
			text[i] = mask[i] ^ text[i]
		}
		fmt.Printf("//check: %s\n", string(text))
		return
	}

	logger := logger{}
	jpu1 := jpu.NewProcessor(1000)
	jpu2 := jpu.NewProcessor(300)

	if !*havingFun {
		jpu1.Trace(" BIGMAC", logger)
	}
	jpu2.Trace("DoorCtl", logger)

	var theBuffer byte

	f := func(where jpu.Address, what byte) { theBuffer = what }
	jpu1.RegisterOut(f, 1)
	jpu2.RegisterOut(f, 1)

	g := func(where jpu.Address) byte { return theBuffer }
	jpu1.RegisterIn(g, 1)
	jpu2.RegisterIn(g, 1)

	// an output port for jpu2
	jpu2.RegisterOut(func(where jpu.Address, what byte) { fmt.Fprintf(os.Stderr, "%c", what) }, 0)

	// from assembling bigmac.src
	prog := []byte{0x4, 0x0, 0x1, 0x1, 0x4, 0x0, 0xc8, 0x3, 0x4, 0x0, 0x0, 0x6, 0x2, 0x3, 0x4, 0x4, 0x0, 0x0, 0x8, 0x9, 0x0, 0x92, 0x4, 0x8, 0x4, 0x0, 0x40, 0x8, 0x7, 0x8, 0x4, 0x5, 0x4, 0x5, 0x6, 0x5, 0x4, 0x6, 0x5, 0x4, 0x6, 0x6, 0x4, 0x5, 0x5, 0x6, 0x3, 0x4, 0x1, 0x4, 0x0, 0x1, 0x8, 0x6, 0x8, 0x3, 0x4, 0x0, 0xff, 0x8, 0x2, 0x1, 0x9, 0xa, 0x0, 0xa0, 0x9, 0x8, 0x4, 0x0, 0x0, 0x8, 0x9, 0x0, 0x68, 0x4, 0x8, 0x4, 0x0, 0x70, 0x0}
fmt.Println(prog)
os.Exit(0)
	jpu1.LoadMem(prog, 100)
	jpu1.Reg[0] = 100

	// compile with includeEncode and use ./clue -encode="XXX" to make these
	mask := []byte{0xc2, 0x87, 0x9c, 0x2a, 0xdf, 0xd2, 0x83, 0x4c, 0xa2, 0x2f, 0xd7, 0x1c, 0xa2, 0xb5, 0x9, 0x8e, 0x52, 0xf, 0xf6, 0x4f, 0xd4, 0x77, 0xa, 0xa5, 0x16, 0x8b, 0x20, 0x36, 0x21, 0xc5, 0x74, 0x85, 0x95, 0xa2, 0xcb, 0xa1, 0x48, 0xcf, 0x9b, 0x0}
	text := []byte{0x96, 0xcf, 0xd9, 0x6a, 0x9c, 0x9d, 0xc7, 0x9, 0xe2, 0x7b, 0x98, 0x5c, 0xed, 0xe5, 0x4c, 0xc0, 0x12, 0x5b, 0xbe, 0xa, 0x94, 0x33, 0x45, 0xea, 0x44, 0xcb, 0x69, 0x65, 0x61, 0x88, 0x3d, 0xd6, 0xc6, 0xf2, 0x82, 0xe6, 0xf, 0x96, 0xdb, 0x0}

	// decrypt the text before putting it in the RAM of the sending machine
	for i := 0; i < len(text); i++ {
		text[i] = text[i] ^ mask[i]
	}
	jpu1.LoadMem(text, 200)

	// start backend machine running indefinitely
	go func() {
		defer func() {
			if r := recover(); r != nil {
				fmt.Println("Error in BIGMAC:", r)
				os.Exit(0)
			}
		}()

		for jpu1.Step() {
			time.Sleep(10 * time.Millisecond)
		}
	}()

	if !*restart {
		fmt.Print(intro)
		fmt.Print(help)
	}

	// process commands
	r := bufio.NewReader(os.Stdin)
loop:
	for {
		// prompt and get input
		fmt.Fprintf(os.Stdout, "> ")
		cmd, err := r.ReadString('\n')
		if err == io.EOF {
			break
		}
		if err != nil {
			log.Fatal("Error:", err)
		}

		// process it
		cmd = strings.ToLower(cmd)
		tok := strings.Fields(cmd)

		if len(tok) < 1 {
			fmt.Print(help)
			continue
		}

		switch tok[0] {
		default:
			fmt.Println("Command not recognized. Type help for a reminder.")
		case "h":
			fallthrough
		case "help":
			fmt.Print(help)
		case "i":
			fallthrough
		case "init":
			fmt.Println("Time warp: ...you settle in for a little hacking...")
			syscall.Exec(os.Args[0], []string{os.Args[0], "-restart"}, nil)
			panic("exec failed")
		case "e":
			fallthrough
		case "exit":
			break loop
		case "q":
			fallthrough
		case "quiet":
			jpu1.Trace("", nil)
			jpu2.Trace("", nil)
		case "d":
			fallthrough
		case "dump":
			row := 15
			for i := 0; i < jpu2.Top/row; i++ {
				fmt.Printf("%3d: ", i*row)
				for j := 0; j < row; j++ {
					fmt.Printf("%3d ", jpu2.Peek(jpu.Address(i*row+j)))
				}
				fmt.Println("")
			}
		case "l":
			fallthrough
		case "load":
			if len(tok) < 2 {
				fmt.Println("Load needs at least one arg.")
			} else {
				a, _ := strconv.ParseInt(tok[1], 0, 16)
				addr := jpu.Address(a)
				for i := 2; i < len(tok); i++ {
					val, err := strconv.ParseInt(tok[i], 0, 16)
					if err == nil {
						jpu2.Poke(addr, byte(val))
						addr++
					} else {
						fmt.Println(err)
					}
				}
			}
		case "r":
			fallthrough
		case "reg":
			if len(tok) > 1 {
				if len(tok) != 3 {
					fmt.Println("Expected 2 args.")
				} else {
					reg, _ := strconv.ParseInt(tok[1], 0, 16)
					val, _ := strconv.ParseInt(tok[2], 0, 16)
					if int(reg) < len(jpu2.Reg) {
						jpu2.Reg[int(reg)] = jpu.Address(val)
					} else {
						fmt.Println("No such register.")
					}
				}
			} else {
				// print all
				for j := 0; j < len(jpu2.Reg); j++ {
					fmt.Printf("%3d ", j)
				}
				fmt.Println("")
				for j := 0; j < len(jpu2.Reg); j++ {
					fmt.Printf("%3d ", jpu2.Reg[j])
				}
				fmt.Println("")
			}
		case "s":
			fallthrough
		case "step":
			running := jpu2.Step()
			runtime.Gosched()
			if !running {
				fmt.Println("Halted.")
			}
		case "g":
			fallthrough
		case "go":
			for jpu2.Step() {
				// make the door run 1/5 as fast, so that BIGMAC beats it to
				// the first spin loop
				time.Sleep(50 * time.Millisecond)
				runtime.Gosched()
			}
			fmt.Println("Halted.")
		}
	}
}

// ex:ts=2
<commit_msg>Removed a hack, go fmt.<commit_after>package main

import (
	"bufio"
	"crypto/rand"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"runtime"
	"strconv"
	"strings"
	"syscall"
	"time"

	"code.google.com/p/jra-go/jpu"
)

const intro = `
It's a moonless night, you're lucky. Cutting your way through
the barbed wire fence, you've made your way into the compound.
Your target up ahead, just a 10 meter sprint across open ground:
a service door and the control panel next to it with its slowly
blinking red LED. You concentrate on the red dot and make a
run for it.

Kneeling beneath the control panel, you unscrew the cover and
find it right where they told you to look for it, the DIAG_IO
port. You hook up Tx, Rx, and ground then boot up your portable
terminal. Propping your back against the wall, you settle in
for a little hacking...
`

const help = `
You have successfully connected your terminal to the DIAG_IO port
of the door control computer. You can send the following commands
to its debugger:

	load address byte byte byte ...
		Load the bytes into memory, starting at address.
	dump
		Dump memory
	reg
		See the registers
	reg x y
		Set register x to y
	step
		Run one step
	go
		Run until halt
	quiet
		Turn off instruction logging
	init
		Reinitialize (same as exit and run the program again)
	exit
		Disconnect the leads and scurry back through the fence.

As a shortcut, you can type just the first letter of any command.
`

type logger struct{}

func (l logger) Log(msg string) {
	fmt.Println(msg)
}

const includeEncode = false
const includeAssemble = true

var restart *bool = flag.Bool("restart", false,
	"for internal use only")
var havingFun *bool = flag.Bool("havingFun", true,
	"set this to false if you're not having fun anymore")
var assemble *string
var encode *string

func init() {
	if includeEncode {
		encode = flag.String("encode", "", "string to hide")
	}
	if includeAssemble {
		assemble = flag.String("assemble", "", "program to assemble")
	}
}

func main() {
	// panic handler
	defer func() {
		if r := recover(); r != nil {
			fmt.Println("Error:", r)
		}
	}()

	// override the normal usage to prevent them from easily seeing
	// the havingFun flag (still might see it with strings, good on 'em)
	flag.Usage = func() { fmt.Fprintf(os.Stderr, "Solve this clue. Duh.\n") }
	flag.Parse()

	if includeAssemble && *assemble != "" {
		f, err := os.Open(*assemble)
		if err != nil {
			fmt.Print("Error: ", err)
			os.Exit(1)
		}
		all, err := ioutil.ReadAll(f)
		if err != nil {
			fmt.Print("Error: ", err)
			os.Exit(1)
		}
		res, org := jpu.Assemble(string(all))

		if *assemble == "bigmac.src" {
			fmt.Printf("%#v\n", res)
		} else {
			fmt.Printf("load %d ", org)
			for _, x := range res {
				fmt.Printf("%d ", x)
			}
			fmt.Println()
		}

		fmt.Printf("reg 0 %d\n", org)
		fmt.Println("quiet")
		fmt.Println("go")
		return
	}

	// this feature is only compiled in when includeEncode is set
	// it creates an encrypted version of the answer using a one-time-pad
	// so that strings on the binary can't find the answer.
	// the answer is decrypted before it is loaded into jpu1's RAM
	if includeEncode && *encode != "" {
		in := []byte(*encode)
		mask := make([]byte, len(in))
		text := make([]byte, len(in))

		_, _ = rand.Read(mask)
		for i, _ := range in {
			for {
				text[i] = mask[i] ^ in[i]
				// do not allow zeros in the encrypted results
				if text[i] == 0 {
					mask[i]++
					// go try again
				} else {
					break
				}
			}
		}

		// add terminators
		mask = append(mask, 0)
		text = append(text, 0)

		fmt.Printf("mask := %#v\n", mask)
		fmt.Printf("text := %#v\n", text)
		for i, _ := range in {
			text[i] = mask[i] ^ text[i]
		}
		fmt.Printf("//check: %s\n", string(text))
		return
	}

	logger := logger{}
	jpu1 := jpu.NewProcessor(1000)
	jpu2 := jpu.NewProcessor(300)

	if !*havingFun {
		jpu1.Trace(" BIGMAC", logger)
	}
	jpu2.Trace("DoorCtl", logger)

	var theBuffer byte

	f := func(where jpu.Address, what byte) { theBuffer = what }
	jpu1.RegisterOut(f, 1)
	jpu2.RegisterOut(f, 1)

	g := func(where jpu.Address) byte { return theBuffer }
	jpu1.RegisterIn(g, 1)
	jpu2.RegisterIn(g, 1)

	// an output port for jpu2
	jpu2.RegisterOut(func(where jpu.Address, what byte) { fmt.Fprintf(os.Stderr, "%c", what) }, 0)

	// from assembling bigmac.src
	prog := []byte{0x4, 0x0, 0x1, 0x1, 0x4, 0x0, 0xc8, 0x3, 0x4, 0x0, 0x0, 0x6, 0x2, 0x3, 0x4, 0x4, 0x0, 0x0, 0x8, 0x9, 0x0, 0x92, 0x4, 0x8, 0x4, 0x0, 0x40, 0x8, 0x7, 0x8, 0x4, 0x5, 0x4, 0x5, 0x6, 0x5, 0x4, 0x6, 0x5, 0x4, 0x6, 0x6, 0x4, 0x5, 0x5, 0x6, 0x3, 0x4, 0x1, 0x4, 0x0, 0x1, 0x8, 0x6, 0x8, 0x3, 0x4, 0x0, 0xff, 0x8, 0x2, 0x1, 0x9, 0xa, 0x0, 0xa0, 0x9, 0x8, 0x4, 0x0, 0x0, 0x8, 0x9, 0x0, 0x68, 0x4, 0x8, 0x4, 0x0, 0x70, 0x0}
	jpu1.LoadMem(prog, 100)
	jpu1.Reg[0] = 100

	// compile with includeEncode and use ./clue -encode="XXX" to make these
	mask := []byte{0xc2, 0x87, 0x9c, 0x2a, 0xdf, 0xd2, 0x83, 0x4c, 0xa2, 0x2f, 0xd7, 0x1c, 0xa2, 0xb5, 0x9, 0x8e, 0x52, 0xf, 0xf6, 0x4f, 0xd4, 0x77, 0xa, 0xa5, 0x16, 0x8b, 0x20, 0x36, 0x21, 0xc5, 0x74, 0x85, 0x95, 0xa2, 0xcb, 0xa1, 0x48, 0xcf, 0x9b, 0x0}
	text := []byte{0x96, 0xcf, 0xd9, 0x6a, 0x9c, 0x9d, 0xc7, 0x9, 0xe2, 0x7b, 0x98, 0x5c, 0xed, 0xe5, 0x4c, 0xc0, 0x12, 0x5b, 0xbe, 0xa, 0x94, 0x33, 0x45, 0xea, 0x44, 0xcb, 0x69, 0x65, 0x61, 0x88, 0x3d, 0xd6, 0xc6, 0xf2, 0x82, 0xe6, 0xf, 0x96, 0xdb, 0x0}

	// decrypt the text before putting it in the RAM of the sending machine
	for i := 0; i < len(text); i++ {
		text[i] = text[i] ^ mask[i]
	}
	jpu1.LoadMem(text, 200)

	// start backend machine running indefinitely
	go func() {
		defer func() {
			if r := recover(); r != nil {
				fmt.Println("Error in BIGMAC:", r)
				os.Exit(0)
			}
		}()

		for jpu1.Step() {
			time.Sleep(10 * time.Millisecond)
		}
	}()

	if !*restart {
		fmt.Print(intro)
		fmt.Print(help)
	}

	// process commands
	r := bufio.NewReader(os.Stdin)
loop:
	for {
		// prompt and get input
		fmt.Fprintf(os.Stdout, "> ")
		cmd, err := r.ReadString('\n')
		if err == io.EOF {
			break
		}
		if err != nil {
			log.Fatal("Error:", err)
		}

		// process it
		cmd = strings.ToLower(cmd)
		tok := strings.Fields(cmd)

		if len(tok) < 1 {
			fmt.Print(help)
			continue
		}

		switch tok[0] {
		default:
			fmt.Println("Command not recognized. Type help for a reminder.")
		case "h":
			fallthrough
		case "help":
			fmt.Print(help)
		case "i":
			fallthrough
		case "init":
			fmt.Println("Time warp: ...you settle in for a little hacking...")
			syscall.Exec(os.Args[0], []string{os.Args[0], "-restart"}, nil)
			panic("exec failed")
		case "e":
			fallthrough
		case "exit":
			break loop
		case "q":
			fallthrough
		case "quiet":
			jpu1.Trace("", nil)
			jpu2.Trace("", nil)
		case "d":
			fallthrough
		case "dump":
			row := 15
			for i := 0; i < jpu2.Top/row; i++ {
				fmt.Printf("%3d: ", i*row)
				for j := 0; j < row; j++ {
					fmt.Printf("%3d ", jpu2.Peek(jpu.Address(i*row+j)))
				}
				fmt.Println("")
			}
		case "l":
			fallthrough
		case "load":
			if len(tok) < 2 {
				fmt.Println("Load needs at least one arg.")
			} else {
				a, _ := strconv.ParseInt(tok[1], 0, 16)
				addr := jpu.Address(a)
				for i := 2; i < len(tok); i++ {
					val, err := strconv.ParseInt(tok[i], 0, 16)
					if err == nil {
						jpu2.Poke(addr, byte(val))
						addr++
					} else {
						fmt.Println(err)
					}
				}
			}
		case "r":
			fallthrough
		case "reg":
			if len(tok) > 1 {
				if len(tok) != 3 {
					fmt.Println("Expected 2 args.")
				} else {
					reg, _ := strconv.ParseInt(tok[1], 0, 16)
					val, _ := strconv.ParseInt(tok[2], 0, 16)
					if int(reg) < len(jpu2.Reg) {
						jpu2.Reg[int(reg)] = jpu.Address(val)
					} else {
						fmt.Println("No such register.")
					}
				}
			} else {
				// print all
				for j := 0; j < len(jpu2.Reg); j++ {
					fmt.Printf("%3d ", j)
				}
				fmt.Println("")
				for j := 0; j < len(jpu2.Reg); j++ {
					fmt.Printf("%3d ", jpu2.Reg[j])
				}
				fmt.Println("")
			}
		case "s":
			fallthrough
		case "step":
			running := jpu2.Step()
			runtime.Gosched()
			if !running {
				fmt.Println("Halted.")
			}
		case "g":
			fallthrough
		case "go":
			for jpu2.Step() {
				// make the door run 1/5 as fast, so that BIGMAC beats it to
				// the first spin loop
				time.Sleep(50 * time.Millisecond)
				runtime.Gosched()
			}
			fmt.Println("Halted.")
		}
	}
}

// ex:ts=2
<|endoftext|>