<commit_before>package consul

import (
	"fmt"
	"log"
	"math/rand"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/hashicorp/consul/consul/structs"
	"github.com/hashicorp/consul/lib"
	"github.com/hashicorp/serf/coordinate"
	"github.com/hashicorp/serf/serf"
)

const (
	// clientRPCMinReuseDuration controls the minimum amount of time RPC
	// queries are sent over an established connection to a single server
	clientRPCMinReuseDuration = 120 * time.Second

	// clientRPCJitterFraction determines the amount of jitter added to
	// clientRPCMinReuseDuration before a connection is expired and a new
	// connection is established in order to rebalance load across consul
	// servers.  The cluster-wide number of connections per second from
	// rebalancing is applied after this jitter to ensure the CPU impact
	// is always finite.  See newRebalanceConnsPerSecPerServer's comment
	// for additional commentary.
	//
	// For example, in a 10K consul cluster with 5x servers, this default
	// averages out to ~13 new connections from rebalancing per server
	// per second (each connection is reused for 120s to 180s).
	clientRPCJitterFraction = 2

	// Limit the number of new connections a server receives per second
	// for connection rebalancing.  This limit caps the load caused by
	// continual rebalancing efforts when a cluster is in equilibrium.  A
	// lower value comes at the cost of increased recovery time after a
	// partition.  This parameter begins to take effect when there are
	// more than ~48K clients querying 5x servers or at lower server
	// values when there is a partition.
	//
	// For example, in a 100K consul cluster with 5x servers, it will
	// take ~5min for all servers to rebalance their connections.  If
	// 99,995 agents are in the minority talking to only one server, it
	// will take ~26min for all servers to rebalance.  A 10K cluster in
	// the same scenario will take ~2.6min to rebalance.
	newRebalanceConnsPerSecPerServer = 64

	// clientRPCConnMaxIdle controls how long we keep an idle connection
	// open to a server.  127s was chosen as the first prime above 120s
	// (arbitrarily chose to use a prime) with the intent of reusing
	// connections who are used by once-a-minute cron(8) jobs *and* who
	// use a 60s jitter window (e.g. in vixie cron job execution can
	// drift by up to 59s per job, or 119s for a once-a-minute cron job).
	clientRPCConnMaxIdle = 127 * time.Second

	// clientMaxStreams controls how many idle streams we keep
	// open to a server
	clientMaxStreams = 32
)

// Interface is used to provide either a Client or Server,
// both of which can be used to perform certain common
// Consul methods
type Interface interface {
	RPC(method string, args interface{}, reply interface{}) error
	LANMembers() []serf.Member
	LocalMember() serf.Member
}

// Client is Consul client which uses RPC to communicate with the
// services for service discovery, health checking, and DC forwarding.
type Client struct {
	config *Config

	// Connection pool to consul servers
	connPool *ConnPool

	// consuls tracks the locally known servers
	consuls    []*serverParts
	consulLock sync.RWMutex

	// eventCh is used to receive events from the
	// serf cluster in the datacenter
	eventCh chan serf.Event

	// lastServer is the last server we made an RPC call to,
	// this is used to re-use the last connection
	lastServer  *serverParts
	lastRPCTime time.Time

	// connRebalanceTime is the time at which we should change the server
	// we query for RPC requests.
	connRebalanceTime time.Time

	// Logger uses the provided LogOutput
	logger *log.Logger

	// serf is the Serf cluster maintained inside the DC
	// which contains all the DC nodes
	serf *serf.Serf

	shutdown     bool
	shutdownCh   chan struct{}
	shutdownLock sync.Mutex
}

// NewClient is used to construct a new Consul client from the
// configuration, potentially returning an error
func NewClient(config *Config) (*Client, error) {
	// Check the protocol version
	if err := config.CheckVersion(); err != nil {
		return nil, err
	}

	// Check for a data directory!
	if config.DataDir == "" {
		return nil, fmt.Errorf("Config must provide a DataDir")
	}

	// Sanity check the ACLs
	if err := config.CheckACL(); err != nil {
		return nil, err
	}

	// Ensure we have a log output
	if config.LogOutput == nil {
		config.LogOutput = os.Stderr
	}

	// Create the tls Wrapper
	tlsWrap, err := config.tlsConfig().OutgoingTLSWrapper()
	if err != nil {
		return nil, err
	}

	// Create a logger
	logger := log.New(config.LogOutput, "", log.LstdFlags)

	// Create server
	c := &Client{
		config:     config,
		connPool:   NewPool(config.LogOutput, clientRPCConnMaxIdle, clientMaxStreams, tlsWrap),
		eventCh:    make(chan serf.Event, 256),
		logger:     logger,
		shutdownCh: make(chan struct{}),
	}

	// Start the Serf listeners to prevent a deadlock
	go c.lanEventHandler()

	// Initialize the lan Serf
	c.serf, err = c.setupSerf(config.SerfLANConfig,
		c.eventCh, serfLANSnapshot)
	if err != nil {
		c.Shutdown()
		return nil, fmt.Errorf("Failed to start lan serf: %v", err)
	}
	return c, nil
}

// setupSerf is used to setup and initialize a Serf
func (c *Client) setupSerf(conf *serf.Config, ch chan serf.Event, path string) (*serf.Serf, error) {
	conf.Init()
	conf.NodeName = c.config.NodeName
	conf.Tags["role"] = "node"
	conf.Tags["dc"] = c.config.Datacenter
	conf.Tags["vsn"] = fmt.Sprintf("%d", c.config.ProtocolVersion)
	conf.Tags["vsn_min"] = fmt.Sprintf("%d", ProtocolVersionMin)
	conf.Tags["vsn_max"] = fmt.Sprintf("%d", ProtocolVersionMax)
	conf.Tags["build"] = c.config.Build
	conf.MemberlistConfig.LogOutput = c.config.LogOutput
	conf.LogOutput = c.config.LogOutput
	conf.EventCh = ch
	conf.SnapshotPath = filepath.Join(c.config.DataDir, path)
	conf.ProtocolVersion = protocolVersionMap[c.config.ProtocolVersion]
	conf.RejoinAfterLeave = c.config.RejoinAfterLeave
	conf.Merge = &lanMergeDelegate{dc: c.config.Datacenter}
	conf.DisableCoordinates = c.config.DisableCoordinates
	if err := ensurePath(conf.SnapshotPath, false); err != nil {
		return nil, err
	}
	return serf.Create(conf)
}

// Shutdown is used to shutdown the client
func (c *Client) Shutdown() error {
	c.logger.Printf("[INFO] consul: shutting down client")
	c.shutdownLock.Lock()
	defer c.shutdownLock.Unlock()

	if c.shutdown {
		return nil
	}

	c.shutdown = true
	close(c.shutdownCh)

	if c.serf != nil {
		c.serf.Shutdown()
	}

	// Close the connection pool
	c.connPool.Shutdown()
	return nil
}

// Leave is used to prepare for a graceful shutdown
func (c *Client) Leave() error {
	c.logger.Printf("[INFO] consul: client starting leave")

	// Leave the LAN pool
	if c.serf != nil {
		if err := c.serf.Leave(); err != nil {
			c.logger.Printf("[ERR] consul: Failed to leave LAN Serf cluster: %v", err)
		}
	}
	return nil
}

// JoinLAN is used to have Consul client join the inner-DC pool
// The target address should be another node inside the DC
// listening on the Serf LAN address
func (c *Client) JoinLAN(addrs []string) (int, error) {
	return c.serf.Join(addrs, true)
}

// LocalMember is used to return the local node
func (c *Client) LocalMember() serf.Member {
	return c.serf.LocalMember()
}

// LANMembers is used to return the members of the LAN cluster
func (c *Client) LANMembers() []serf.Member {
	return c.serf.Members()
}

// RemoveFailedNode is used to remove a failed node from the cluster
func (c *Client) RemoveFailedNode(node string) error {
	return c.serf.RemoveFailedNode(node)
}

// KeyManagerLAN returns the LAN Serf keyring manager
func (c *Client) KeyManagerLAN() *serf.KeyManager {
	return c.serf.KeyManager()
}

// Encrypted determines if gossip is encrypted
func (c *Client) Encrypted() bool {
	return c.serf.EncryptionEnabled()
}

// lanEventHandler is used to handle events from the lan Serf cluster
func (c *Client) lanEventHandler() {
	for {
		select {
		case e := <-c.eventCh:
			switch e.EventType() {
			case serf.EventMemberJoin:
				c.nodeJoin(e.(serf.MemberEvent))
			case serf.EventMemberLeave, serf.EventMemberFailed:
				c.nodeFail(e.(serf.MemberEvent))
			case serf.EventUser:
				c.localEvent(e.(serf.UserEvent))
			case serf.EventMemberUpdate: // Ignore
			case serf.EventMemberReap: // Ignore
			case serf.EventQuery: // Ignore
			default:
				c.logger.Printf("[WARN] consul: unhandled LAN Serf Event: %#v", e)
			}
		case <-c.shutdownCh:
			return
		}
	}
}

// nodeJoin is used to handle join events on the serf cluster
func (c *Client) nodeJoin(me serf.MemberEvent) {
	for _, m := range me.Members {
		ok, parts := isConsulServer(m)
		if !ok {
			continue
		}
		if parts.Datacenter != c.config.Datacenter {
			c.logger.Printf("[WARN] consul: server %s for datacenter %s has joined wrong cluster",
				m.Name, parts.Datacenter)
			continue
		}
		c.logger.Printf("[INFO] consul: adding server %s", parts)

		// Check if this server is known
		found := false
		c.consulLock.Lock()
		for idx, existing := range c.consuls {
			if existing.Name == parts.Name {
				c.consuls[idx] = parts
				found = true
				break
			}
		}

		// Add to the list if not known
		if !found {
			c.consuls = append(c.consuls, parts)
		}
		c.consulLock.Unlock()

		// Trigger the callback
		if c.config.ServerUp != nil {
			c.config.ServerUp()
		}
	}
}

// nodeFail is used to handle fail events on the serf cluster
func (c *Client) nodeFail(me serf.MemberEvent) {
	for _, m := range me.Members {
		ok, parts := isConsulServer(m)
		if !ok {
			continue
		}
		c.logger.Printf("[INFO] consul: removing server %s", parts)

		// Remove the server if known
		c.consulLock.Lock()
		n := len(c.consuls)
		for i := 0; i < n; i++ {
			if c.consuls[i].Name == parts.Name {
				c.consuls[i], c.consuls[n-1] = c.consuls[n-1], nil
				c.consuls = c.consuls[:n-1]
				break
			}
		}
		c.consulLock.Unlock()
	}
}

// localEvent is called when we receive an event on the local Serf
func (c *Client) localEvent(event serf.UserEvent) {
	// Handle only consul events
	if !strings.HasPrefix(event.Name, "consul:") {
		return
	}

	switch name := event.Name; {
	case name == newLeaderEvent:
		c.logger.Printf("[INFO] consul: New leader elected: %s", event.Payload)

		// Trigger the callback
		if c.config.ServerUp != nil {
			c.config.ServerUp()
		}
	case isUserEvent(name):
		event.Name = rawUserEventName(name)
		c.logger.Printf("[DEBUG] consul: user event: %s", event.Name)

		// Trigger the callback
		if c.config.UserEventHandler != nil {
			c.config.UserEventHandler(event)
		}
	default:
		c.logger.Printf("[WARN] consul: Unhandled local event: %v", event)
	}
}

// RPC is used to forward an RPC call to a consul server, or fail if no servers
func (c *Client) RPC(method string, args interface{}, reply interface{}) error {
	// Check to make sure we haven't spent too much time querying a
	// single server
	now := time.Now()
	if !c.connRebalanceTime.IsZero() && now.After(c.connRebalanceTime) {
		c.logger.Printf("[DEBUG] consul: connection time to server %s exceeded, rotating server connection", c.lastServer.Addr)
		c.lastServer = nil
	}

	// Allocate these vars on the stack before the goto
	var numConsulServers int
	var clusterWideRebalanceConnsPerSec float64
	var connReuseLowWaterMark time.Duration
	var numLANMembers int

	// Check the last RPC time, continue to reuse cached connection for
	// up to clientRPCMinReuseDuration unless exceeded
	// clientRPCConnMaxIdle
	lastRPCTime := now.Sub(c.lastRPCTime)
	var server *serverParts
	if c.lastServer != nil && lastRPCTime < clientRPCConnMaxIdle {
		server = c.lastServer
		goto TRY_RPC
	}

	// Bail if we can't find any servers
	c.consulLock.RLock()
	numConsulServers = len(c.consuls)
	if numConsulServers == 0 {
		c.consulLock.RUnlock()
		return structs.ErrNoServers
	}

	// Select a random addr
	server = c.consuls[rand.Int31n(int32(numConsulServers))]
	c.consulLock.RUnlock()

	// Limit this connection's life based on the size (and health) of the
	// cluster.  Never rebalance a connection more frequently than
	// connReuseLowWaterMark, and make sure we never exceed
	// clusterWideRebalanceConnsPerSec operations/s across numLANMembers.
	clusterWideRebalanceConnsPerSec = float64(numConsulServers * newRebalanceConnsPerSecPerServer)
	connReuseLowWaterMark = clientRPCMinReuseDuration + lib.RandomStagger(clientRPCMinReuseDuration/clientRPCJitterFraction)
	numLANMembers = len(c.LANMembers())
	c.connRebalanceTime = now.Add(lib.RateScaledInterval(clusterWideRebalanceConnsPerSec, connReuseLowWaterMark, numLANMembers))
	c.logger.Printf("[DEBUG] consul: connection to server %s will expire at %v", server.Addr, c.connRebalanceTime)

	// Forward to remote Consul
TRY_RPC:
	if err := c.connPool.RPC(c.config.Datacenter, server.Addr, server.Version, method, args, reply); err != nil {
		c.connRebalanceTime = time.Time{}
		c.lastRPCTime = time.Time{}
		c.lastServer = nil
		return err
	}

	// Cache the last server
	c.lastServer = server
	c.lastRPCTime = now
	return nil
}

// Stats is used to return statistics for debugging and insight
// for various sub-systems
func (c *Client) Stats() map[string]map[string]string {
	toString := func(v uint64) string {
		return strconv.FormatUint(v, 10)
	}
	stats := map[string]map[string]string{
		"consul": map[string]string{
			"server":        "false",
			"known_servers": toString(uint64(len(c.consuls))),
		},
		"serf_lan": c.serf.Stats(),
		"runtime":  runtimeStats(),
	}
	return stats
}

// GetCoordinate returns the network coordinate of the current node, as
// maintained by Serf.
func (c *Client) GetCoordinate() (*coordinate.Coordinate, error) {
	return c.serf.GetCoordinate()
}
<commit_msg>Revert "Merge pull request #1667 from hashicorp/b-redistribute-clients"<commit_after>package consul

import (
	"fmt"
	"log"
	"math/rand"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/hashicorp/consul/consul/structs"
	"github.com/hashicorp/serf/coordinate"
	"github.com/hashicorp/serf/serf"
)

const (
	// clientRPCCache controls how long we keep an idle connection
	// open to a server
	clientRPCCache = 30 * time.Second

	// clientMaxStreams controls how many idle streams we keep
	// open to a server
	clientMaxStreams = 32
)

// Interface is used to provide either a Client or Server,
// both of which can be used to perform certain common
// Consul methods
type Interface interface {
	RPC(method string, args interface{}, reply interface{}) error
	LANMembers() []serf.Member
	LocalMember() serf.Member
}

// Client is Consul client which uses RPC to communicate with the
// services for service discovery, health checking, and DC forwarding.
type Client struct {
	config *Config

	// Connection pool to consul servers
	connPool *ConnPool

	// consuls tracks the locally known servers
	consuls    []*serverParts
	consulLock sync.RWMutex

	// eventCh is used to receive events from the
	// serf cluster in the datacenter
	eventCh chan serf.Event

	// lastServer is the last server we made an RPC call to,
	// this is used to re-use the last connection
	lastServer  *serverParts
	lastRPCTime time.Time

	// Logger uses the provided LogOutput
	logger *log.Logger

	// serf is the Serf cluster maintained inside the DC
	// which contains all the DC nodes
	serf *serf.Serf

	shutdown     bool
	shutdownCh   chan struct{}
	shutdownLock sync.Mutex
}

// NewClient is used to construct a new Consul client from the
// configuration, potentially returning an error
func NewClient(config *Config) (*Client, error) {
	// Check the protocol version
	if err := config.CheckVersion(); err != nil {
		return nil, err
	}

	// Check for a data directory!
	if config.DataDir == "" {
		return nil, fmt.Errorf("Config must provide a DataDir")
	}

	// Sanity check the ACLs
	if err := config.CheckACL(); err != nil {
		return nil, err
	}

	// Ensure we have a log output
	if config.LogOutput == nil {
		config.LogOutput = os.Stderr
	}

	// Create the tls Wrapper
	tlsWrap, err := config.tlsConfig().OutgoingTLSWrapper()
	if err != nil {
		return nil, err
	}

	// Create a logger
	logger := log.New(config.LogOutput, "", log.LstdFlags)

	// Create server
	c := &Client{
		config:     config,
		connPool:   NewPool(config.LogOutput, clientRPCCache, clientMaxStreams, tlsWrap),
		eventCh:    make(chan serf.Event, 256),
		logger:     logger,
		shutdownCh: make(chan struct{}),
	}

	// Start the Serf listeners to prevent a deadlock
	go c.lanEventHandler()

	// Initialize the lan Serf
	c.serf, err = c.setupSerf(config.SerfLANConfig,
		c.eventCh, serfLANSnapshot)
	if err != nil {
		c.Shutdown()
		return nil, fmt.Errorf("Failed to start lan serf: %v", err)
	}
	return c, nil
}

// setupSerf is used to setup and initialize a Serf
func (c *Client) setupSerf(conf *serf.Config, ch chan serf.Event, path string) (*serf.Serf, error) {
	conf.Init()
	conf.NodeName = c.config.NodeName
	conf.Tags["role"] = "node"
	conf.Tags["dc"] = c.config.Datacenter
	conf.Tags["vsn"] = fmt.Sprintf("%d", c.config.ProtocolVersion)
	conf.Tags["vsn_min"] = fmt.Sprintf("%d", ProtocolVersionMin)
	conf.Tags["vsn_max"] = fmt.Sprintf("%d", ProtocolVersionMax)
	conf.Tags["build"] = c.config.Build
	conf.MemberlistConfig.LogOutput = c.config.LogOutput
	conf.LogOutput = c.config.LogOutput
	conf.EventCh = ch
	conf.SnapshotPath = filepath.Join(c.config.DataDir, path)
	conf.ProtocolVersion = protocolVersionMap[c.config.ProtocolVersion]
	conf.RejoinAfterLeave = c.config.RejoinAfterLeave
	conf.Merge = &lanMergeDelegate{dc: c.config.Datacenter}
	conf.DisableCoordinates = c.config.DisableCoordinates
	if err := ensurePath(conf.SnapshotPath, false); err != nil {
		return nil, err
	}
	return serf.Create(conf)
}

// Shutdown is used to shutdown the client
func (c *Client) Shutdown() error {
	c.logger.Printf("[INFO] consul: shutting down client")
	c.shutdownLock.Lock()
	defer c.shutdownLock.Unlock()

	if c.shutdown {
		return nil
	}

	c.shutdown = true
	close(c.shutdownCh)

	if c.serf != nil {
		c.serf.Shutdown()
	}

	// Close the connection pool
	c.connPool.Shutdown()
	return nil
}

// Leave is used to prepare for a graceful shutdown
func (c *Client) Leave() error {
	c.logger.Printf("[INFO] consul: client starting leave")

	// Leave the LAN pool
	if c.serf != nil {
		if err := c.serf.Leave(); err != nil {
			c.logger.Printf("[ERR] consul: Failed to leave LAN Serf cluster: %v", err)
		}
	}
	return nil
}

// JoinLAN is used to have Consul client join the inner-DC pool
// The target address should be another node inside the DC
// listening on the Serf LAN address
func (c *Client) JoinLAN(addrs []string) (int, error) {
	return c.serf.Join(addrs, true)
}

// LocalMember is used to return the local node
func (c *Client) LocalMember() serf.Member {
	return c.serf.LocalMember()
}

// LANMembers is used to return the members of the LAN cluster
func (c *Client) LANMembers() []serf.Member {
	return c.serf.Members()
}

// RemoveFailedNode is used to remove a failed node from the cluster
func (c *Client) RemoveFailedNode(node string) error {
	return c.serf.RemoveFailedNode(node)
}

// KeyManagerLAN returns the LAN Serf keyring manager
func (c *Client) KeyManagerLAN() *serf.KeyManager {
	return c.serf.KeyManager()
}

// Encrypted determines if gossip is encrypted
func (c *Client) Encrypted() bool {
	return c.serf.EncryptionEnabled()
}

// lanEventHandler is used to handle events from the lan Serf cluster
func (c *Client) lanEventHandler() {
	for {
		select {
		case e := <-c.eventCh:
			switch e.EventType() {
			case serf.EventMemberJoin:
				c.nodeJoin(e.(serf.MemberEvent))
			case serf.EventMemberLeave, serf.EventMemberFailed:
				c.nodeFail(e.(serf.MemberEvent))
			case serf.EventUser:
				c.localEvent(e.(serf.UserEvent))
			case serf.EventMemberUpdate: // Ignore
			case serf.EventMemberReap: // Ignore
			case serf.EventQuery: // Ignore
			default:
				c.logger.Printf("[WARN] consul: unhandled LAN Serf Event: %#v", e)
			}
		case <-c.shutdownCh:
			return
		}
	}
}

// nodeJoin is used to handle join events on the serf cluster
func (c *Client) nodeJoin(me serf.MemberEvent) {
	for _, m := range me.Members {
		ok, parts := isConsulServer(m)
		if !ok {
			continue
		}
		if parts.Datacenter != c.config.Datacenter {
			c.logger.Printf("[WARN] consul: server %s for datacenter %s has joined wrong cluster",
				m.Name, parts.Datacenter)
			continue
		}
		c.logger.Printf("[INFO] consul: adding server %s", parts)

		// Check if this server is known
		found := false
		c.consulLock.Lock()
		for idx, existing := range c.consuls {
			if existing.Name == parts.Name {
				c.consuls[idx] = parts
				found = true
				break
			}
		}

		// Add to the list if not known
		if !found {
			c.consuls = append(c.consuls, parts)
		}
		c.consulLock.Unlock()

		// Trigger the callback
		if c.config.ServerUp != nil {
			c.config.ServerUp()
		}
	}
}

// nodeFail is used to handle fail events on the serf cluster
func (c *Client) nodeFail(me serf.MemberEvent) {
	for _, m := range me.Members {
		ok, parts := isConsulServer(m)
		if !ok {
			continue
		}
		c.logger.Printf("[INFO] consul: removing server %s", parts)

		// Remove the server if known
		c.consulLock.Lock()
		n := len(c.consuls)
		for i := 0; i < n; i++ {
			if c.consuls[i].Name == parts.Name {
				c.consuls[i], c.consuls[n-1] = c.consuls[n-1], nil
				c.consuls = c.consuls[:n-1]
				break
			}
		}
		c.consulLock.Unlock()
	}
}

// localEvent is called when we receive an event on the local Serf
func (c *Client) localEvent(event serf.UserEvent) {
	// Handle only consul events
	if !strings.HasPrefix(event.Name, "consul:") {
		return
	}

	switch name := event.Name; {
	case name == newLeaderEvent:
		c.logger.Printf("[INFO] consul: New leader elected: %s", event.Payload)

		// Trigger the callback
		if c.config.ServerUp != nil {
			c.config.ServerUp()
		}
	case isUserEvent(name):
		event.Name = rawUserEventName(name)
		c.logger.Printf("[DEBUG] consul: user event: %s", event.Name)

		// Trigger the callback
		if c.config.UserEventHandler != nil {
			c.config.UserEventHandler(event)
		}
	default:
		c.logger.Printf("[WARN] consul: Unhandled local event: %v", event)
	}
}

// RPC is used to forward an RPC call to a consul server, or fail if no servers
func (c *Client) RPC(method string, args interface{}, reply interface{}) error {
	// Check the last rpc time
	var server *serverParts
	if time.Now().Sub(c.lastRPCTime) < clientRPCCache {
		server = c.lastServer
		if server != nil {
			goto TRY_RPC
		}
	}

	// Bail if we can't find any servers
	c.consulLock.RLock()
	if len(c.consuls) == 0 {
		c.consulLock.RUnlock()
		return structs.ErrNoServers
	}

	// Select a random addr
	server = c.consuls[rand.Int31()%int32(len(c.consuls))]
	c.consulLock.RUnlock()

	// Forward to remote Consul
TRY_RPC:
	if err := c.connPool.RPC(c.config.Datacenter, server.Addr, server.Version, method, args, reply); err != nil {
		c.lastServer = nil
		c.lastRPCTime = time.Time{}
		return err
	}

	// Cache the last server
	c.lastServer = server
	c.lastRPCTime = time.Now()
	return nil
}

// Stats is used to return statistics for debugging and insight
// for various sub-systems
func (c *Client) Stats() map[string]map[string]string {
	toString := func(v uint64) string {
		return strconv.FormatUint(v, 10)
	}
	stats := map[string]map[string]string{
		"consul": map[string]string{
			"server":        "false",
			"known_servers": toString(uint64(len(c.consuls))),
		},
		"serf_lan": c.serf.Stats(),
		"runtime":  runtimeStats(),
	}
	return stats
}

// GetCoordinate returns the network coordinate of the current node, as
// maintained by Serf.
func (c *Client) GetCoordinate() (*coordinate.Coordinate, error) {
	return c.serf.GetCoordinate()
}
<|endoftext|>


<commit_before>package consul

import (
	"fmt"
	"github.com/hashicorp/consul/consul/structs"
	"github.com/hashicorp/serf/serf"
	"log"
	"math/rand"
	"net"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"
)

var (
	clientRPCCache = 30 * time.Second
)

// Interface is used to provide either a Client or Server,
// both of which can be used to perform certain common
// Consul methods
type Interface interface {
	RPC(method string, args interface{}, reply interface{}) error
	LANMembers() []serf.Member
}

// Client is Consul client which uses RPC to communicate with the
// services for service discovery, health checking, and DC forwarding.
type Client struct {
	config *Config

	// Connection pool to consul servers
	connPool *ConnPool

	// consuls tracks the locally known servers
	consuls    []net.Addr
	consulLock sync.RWMutex

	// eventCh is used to receive events from the
	// serf cluster in the datacenter
	eventCh chan serf.Event

	// lastServer is the last server we made an RPC call to,
	// this is used to re-use the last connection
	lastServer  net.Addr
	lastRPCTime time.Time

	// Logger uses the provided LogOutput
	logger *log.Logger

	// serf is the Serf cluster maintained inside the DC
	// which contains all the DC nodes
	serf *serf.Serf

	shutdown     bool
	shutdownCh   chan struct{}
	shutdownLock sync.Mutex
}

// NewClient is used to construct a new Consul client from the
// configuration, potentially returning an error
func NewClient(config *Config) (*Client, error) {
	// Check the protocol version
	if err := config.CheckVersion(); err != nil {
		return nil, err
	}

	// Check for a data directory!
	if config.DataDir == "" {
		return nil, fmt.Errorf("Config must provide a DataDir")
	}

	// Ensure we have a log output
	if config.LogOutput == nil {
		config.LogOutput = os.Stderr
	}

	// Create a logger
	logger := log.New(config.LogOutput, "", log.LstdFlags)

	// Create server
	c := &Client{
		config:     config,
		connPool:   NewPool(clientRPCCache),
		eventCh:    make(chan serf.Event, 256),
		logger:     logger,
		shutdownCh: make(chan struct{}),
	}

	// Start the Serf listeners to prevent a deadlock
	go c.lanEventHandler()

	// Initialize the lan Serf
	var err error
	c.serf, err = c.setupSerf(config.SerfLANConfig,
		c.eventCh, serfLANSnapshot)
	if err != nil {
		c.Shutdown()
		return nil, fmt.Errorf("Failed to start lan serf: %v", err)
	}
	return c, nil
}

// setupSerf is used to setup and initialize a Serf
func (c *Client) setupSerf(conf *serf.Config, ch chan serf.Event, path string) (*serf.Serf, error) {
	conf.Init()
	conf.NodeName = c.config.NodeName
	conf.Tags["role"] = "node"
	conf.Tags["dc"] = c.config.Datacenter
	conf.Tags["vsn"] = fmt.Sprintf("%d", c.config.ProtocolVersion)
	conf.Tags["vsn_min"] = fmt.Sprintf("%d", ProtocolVersionMin)
	conf.Tags["vsn_max"] = fmt.Sprintf("%d", ProtocolVersionMax)
	conf.MemberlistConfig.LogOutput = c.config.LogOutput
	conf.LogOutput = c.config.LogOutput
	conf.EventCh = ch
	conf.SnapshotPath = filepath.Join(c.config.DataDir, path)
	conf.ProtocolVersion = protocolVersionMap[c.config.ProtocolVersion]
	if err := ensurePath(conf.SnapshotPath, false); err != nil {
		return nil, err
	}
	return serf.Create(conf)
}

// Shutdown is used to shutdown the client
func (c *Client) Shutdown() error {
	c.logger.Printf("[INFO] consul: shutting down client")
	c.shutdownLock.Lock()
	defer c.shutdownLock.Unlock()

	if c.shutdown {
		return nil
	}

	c.shutdown = true
	close(c.shutdownCh)

	if c.serf != nil {
		c.serf.Shutdown()
	}

	// Close the connection pool
	c.connPool.Shutdown()
	return nil
}

// Leave is used to prepare for a graceful shutdown
func (c *Client) Leave() error {
	c.logger.Printf("[INFO] consul: client starting leave")

	// Leave the LAN pool
	if c.serf != nil {
		if err := c.serf.Leave(); err != nil {
			c.logger.Printf("[ERR] consul: Failed to leave LAN Serf cluster: %v", err)
		}
	}
	return nil
}

// JoinLAN is used to have Consul client join the inner-DC pool
// The target address should be another node inside the DC
// listening on the Serf LAN address
func (c *Client) JoinLAN(addrs []string) (int, error) {
	return c.serf.Join(addrs, true)
}

// LANMembers is used to return the members of the LAN cluster
func (c *Client) LANMembers() []serf.Member {
	return c.serf.Members()
}

// RemoveFailedNode is used to remove a failed node from the cluster
func (c *Client) RemoveFailedNode(node string) error {
	return c.serf.RemoveFailedNode(node)
}

// lanEventHandler is used to handle events from the lan Serf cluster
func (c *Client) lanEventHandler() {
	for {
		select {
		case e := <-c.eventCh:
			switch e.EventType() {
			case serf.EventMemberJoin:
				c.nodeJoin(e.(serf.MemberEvent))
			case serf.EventMemberLeave:
				fallthrough
			case serf.EventMemberFailed:
				c.nodeFail(e.(serf.MemberEvent))
			case serf.EventUser:
				c.localEvent(e.(serf.UserEvent))
			case serf.EventMemberUpdate: // Ignore
			case serf.EventMemberReap: // Ignore
			case serf.EventQuery: // Ignore
			default:
				c.logger.Printf("[WARN] consul: unhandled LAN Serf Event: %#v", e)
			}
		case <-c.shutdownCh:
			return
		}
	}
}

// nodeJoin is used to handle join events on the serf cluster
func (c *Client) nodeJoin(me serf.MemberEvent) {
	for _, m := range me.Members {
		ok, parts := isConsulServer(m)
		if !ok {
			continue
		}
		if parts.Datacenter != c.config.Datacenter {
			c.logger.Printf("[WARN] consul: server %s for datacenter %s has joined wrong cluster",
				m.Name, parts.Datacenter)
			continue
		}

		var addr net.Addr = &net.TCPAddr{IP: m.Addr, Port: parts.Port}
		c.logger.Printf("[INFO] consul: adding server for datacenter: %s, addr: %s", parts.Datacenter, addr)

		// Check if this server is known
		found := false
		c.consulLock.Lock()
		for _, c := range c.consuls {
			if c.String() == addr.String() {
				found = true
				break
			}
		}

		// Add to the list if not known
		if !found {
			c.consuls = append(c.consuls, addr)
		}
		c.consulLock.Unlock()

		// Trigger the callback
		if c.config.ServerUp != nil {
			c.config.ServerUp()
		}
	}
}

// nodeFail is used to handle fail events on the serf cluster
func (c *Client) nodeFail(me serf.MemberEvent) {
	for _, m := range me.Members {
		ok, parts := isConsulServer(m)
		if !ok {
			continue
		}
		var addr net.Addr = &net.TCPAddr{IP: m.Addr, Port: parts.Port}
		c.logger.Printf("[INFO] consul: removing server for datacenter: %s, addr: %s", parts.Datacenter, addr)

		// Remove the server if known
		c.consulLock.Lock()
		n := len(c.consuls)
		for i := 0; i < n; i++ {
			if c.consuls[i].String() == addr.String() {
				c.consuls[i], c.consuls[n-1] = c.consuls[n-1], nil
				c.consuls = c.consuls[:n-1]
				break
			}
		}
		c.consulLock.Unlock()
	}
}

// localEvent is called when we receive an event on the local Serf
func (c *Client) localEvent(event serf.UserEvent) {
	// Handle only consul events
	if !strings.HasPrefix(event.Name, "consul:") {
		return
	}

	switch event.Name {
	case newLeaderEvent:
		c.logger.Printf("[INFO] consul: New leader elected: %s", event.Payload)

		// Trigger the callback
		if c.config.ServerUp != nil {
			c.config.ServerUp()
		}
	default:
		c.logger.Printf("[WARN] consul: Unhandled local event: %v", event)
	}
}

// RPC is used to forward an RPC call to a consul server, or fail if no servers
func (c *Client) RPC(method string, args interface{}, reply interface{}) error {
	// Check the last rpc time
	var server net.Addr
	if time.Now().Sub(c.lastRPCTime) < clientRPCCache {
		server = c.lastServer
		if server != nil {
			goto TRY_RPC
		}
	}

	// Bail if we can't find any servers
	c.consulLock.RLock()
	if len(c.consuls) == 0 {
		c.consulLock.RUnlock()
		return structs.ErrNoServers
	}

	// Select a random addr
	server = c.consuls[rand.Int31()%int32(len(c.consuls))]
	c.consulLock.RUnlock()

	// Forward to remote Consul
TRY_RPC:
	if err := c.connPool.RPC(server, method, args, reply); err != nil {
		c.lastServer = nil
		c.lastRPCTime = time.Time{}
		return err
	}

	// Cache the last server
	c.lastServer = server
	c.lastRPCTime = time.Now()
	return nil
}

// Stats is used to return statistics for debugging and insight
// for various sub-systems
func (c *Client) Stats() map[string]map[string]string {
	toString := func(v uint64) string {
		return strconv.FormatUint(v, 10)
	}
	stats := map[string]map[string]string{
		"consul": map[string]string{
			"server":        "false",
			"known_servers": toString(uint64(len(c.consuls))),
		},
		"serf_lan": c.serf.Stats(),
	}
	return stats
}
<commit_msg>consul: Client initializes TLS settings<commit_after>package consul

import (
	"crypto/tls"
	"fmt"
	"github.com/hashicorp/consul/consul/structs"
	"github.com/hashicorp/serf/serf"
	"log"
	"math/rand"
	"net"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"
)

var (
	clientRPCCache = 30 * time.Second
)

// Interface is used to provide either a Client or Server,
// both of which can be used to perform certain common
// Consul methods
type Interface interface {
	RPC(method string, args interface{}, reply interface{}) error
	LANMembers() []serf.Member
}

// Client is Consul client which uses RPC to communicate with the
// services for service discovery, health checking, and DC forwarding.
type Client struct {
	config *Config

	// Connection pool to consul servers
	connPool *ConnPool

	// consuls tracks the locally known servers
	consuls    []net.Addr
	consulLock sync.RWMutex

	// eventCh is used to receive events from the
	// serf cluster in the datacenter
	eventCh chan serf.Event

	// lastServer is the last server we made an RPC call to,
	// this is used to re-use the last connection
	lastServer  net.Addr
	lastRPCTime time.Time

	// Logger uses the provided LogOutput
	logger *log.Logger

	// serf is the Serf cluster maintained inside the DC
	// which contains all the DC nodes
	serf *serf.Serf

	shutdown     bool
	shutdownCh   chan struct{}
	shutdownLock sync.Mutex
}

// NewClient is used to construct a new Consul client from the
// configuration, potentially returning an error
func NewClient(config *Config) (*Client, error) {
	// Check the protocol version
	if err := config.CheckVersion(); err != nil {
		return nil, err
	}

	// Check for a data directory!
	if config.DataDir == "" {
		return nil, fmt.Errorf("Config must provide a DataDir")
	}

	// Ensure we have a log output
	if config.LogOutput == nil {
		config.LogOutput = os.Stderr
	}

	// Create the tlsConfig
	var tlsConfig *tls.Config
	var err error
	if config.VerifyOutgoing {
		if tlsConfig, err = config.OutgoingTLSConfig(); err != nil {
			return nil, err
		}
	}

	// Create a logger
	logger := log.New(config.LogOutput, "", log.LstdFlags)

	// Create server
	c := &Client{
		config:     config,
		connPool:   NewPool(clientRPCCache, tlsConfig),
		eventCh:    make(chan serf.Event, 256),
		logger:     logger,
		shutdownCh: make(chan struct{}),
	}

	// Start the Serf listeners to prevent a deadlock
	go c.lanEventHandler()

	// Initialize the lan Serf
	c.serf, err = c.setupSerf(config.SerfLANConfig,
		c.eventCh, serfLANSnapshot)
	if err != nil {
		c.Shutdown()
		return nil, fmt.Errorf("Failed to start lan serf: %v", err)
	}
	return c, nil
}

// setupSerf is used to setup and initialize a Serf
func (c *Client) setupSerf(conf *serf.Config, ch chan serf.Event, path string) (*serf.Serf, error) {
	conf.Init()
	conf.NodeName = c.config.NodeName
	conf.Tags["role"] = "node"
	conf.Tags["dc"] = c.config.Datacenter
	conf.Tags["vsn"] = fmt.Sprintf("%d", c.config.ProtocolVersion)
	conf.Tags["vsn_min"] = fmt.Sprintf("%d", ProtocolVersionMin)
	conf.Tags["vsn_max"] = fmt.Sprintf("%d", ProtocolVersionMax)
	conf.MemberlistConfig.LogOutput = c.config.LogOutput
	conf.LogOutput = c.config.LogOutput
	conf.EventCh = ch
	conf.SnapshotPath = filepath.Join(c.config.DataDir, path)
	conf.ProtocolVersion = protocolVersionMap[c.config.ProtocolVersion]
	if err := ensurePath(conf.SnapshotPath, false); err != nil {
		return nil, err
	}
	return serf.Create(conf)
}

// Shutdown is used to shutdown the client
func (c *Client) Shutdown() error {
	c.logger.Printf("[INFO] consul: shutting down client")
	c.shutdownLock.Lock()
	defer c.shutdownLock.Unlock()

	if c.shutdown {
		return nil
	}

	c.shutdown = true
	close(c.shutdownCh)

	if c.serf != nil {
		c.serf.Shutdown()
	}

	// Close the connection pool
	c.connPool.Shutdown()
	return nil
}

// Leave is used to prepare for a graceful shutdown
func (c *Client) Leave() error {
	c.logger.Printf("[INFO] consul: client starting leave")

	// Leave the LAN pool
	if c.serf != nil {
		if err := c.serf.Leave(); err != nil {
			c.logger.Printf("[ERR] consul: Failed to leave LAN Serf cluster: %v", err)
		}
	}
	return nil
}

// JoinLAN is used to have Consul client join the inner-DC pool
// The target address should be another node inside the DC
// listening on the Serf LAN address
func (c *Client) JoinLAN(addrs []string) (int, error) {
	return c.serf.Join(addrs, true)
}

// LANMembers is used to return the members of the LAN cluster
func (c *Client) LANMembers() []serf.Member {
	return c.serf.Members()
}

// RemoveFailedNode is used to remove a failed node from the cluster
func (c *Client) RemoveFailedNode(node string) error {
	return c.serf.RemoveFailedNode(node)
}

// lanEventHandler is used to handle events from the lan Serf cluster
func (c *Client) lanEventHandler() {
	for {
		select {
		case e := <-c.eventCh:
			switch e.EventType() {
			case serf.EventMemberJoin:
				c.nodeJoin(e.(serf.MemberEvent))
			case serf.EventMemberLeave:
				fallthrough
			case serf.EventMemberFailed:
				c.nodeFail(e.(serf.MemberEvent))
			case serf.EventUser:
				c.localEvent(e.(serf.UserEvent))
			case serf.EventMemberUpdate: // Ignore
			case serf.EventMemberReap: // Ignore
			case serf.EventQuery: // Ignore
			default:
				c.logger.Printf("[WARN] consul: unhandled LAN Serf Event: %#v", e)
			}
		case <-c.shutdownCh:
			return
		}
	}
}

// nodeJoin is used to handle join events on the serf cluster
func (c *Client) nodeJoin(me serf.MemberEvent) {
	for _, m := range me.Members {
		ok, parts := isConsulServer(m)
		if !ok {
			continue
		}
		if parts.Datacenter != c.config.Datacenter {
			c.logger.Printf("[WARN] consul: server %s for datacenter %s has joined wrong cluster",
				m.Name, parts.Datacenter)
			continue
		}

		var addr net.Addr = &net.TCPAddr{IP: m.Addr, Port: parts.Port}
		c.logger.Printf("[INFO] consul: adding server for datacenter: %s, addr: %s", parts.Datacenter, addr)

		// Check if this server is known
		found := false
		c.consulLock.Lock()
		for _, c := range c.consuls {
			if c.String() == addr.String() {
				found = true
				break
			}
		}

		// Add to the list if not known
		if !found {
			c.consuls = append(c.consuls, addr)
		}
		c.consulLock.Unlock()

		// Trigger the callback
		if c.config.ServerUp != nil {
			c.config.ServerUp()
		}
	}
}

// nodeFail is used to handle fail events on the serf cluster
func (c *Client) nodeFail(me serf.MemberEvent) {
	for _, m := range me.Members {
		ok, parts := isConsulServer(m)
		if !ok {
			continue
		}
		var addr net.Addr = &net.TCPAddr{IP: m.Addr, Port: parts.Port}
		c.logger.Printf("[INFO] consul: removing server for datacenter: %s, addr: %s", parts.Datacenter, addr)

		// Remove the server if known
		c.consulLock.Lock()
		n := len(c.consuls)
		for i := 0; i < n; i++ {
			if c.consuls[i].String() == addr.String() {
				c.consuls[i], c.consuls[n-1] = c.consuls[n-1], nil
				c.consuls = c.consuls[:n-1]
				break
			}
		}
		c.consulLock.Unlock()
	}
}

// localEvent is called when we receive an event on the local Serf
func (c *Client) localEvent(event serf.UserEvent) {
	// Handle only consul events
	if !strings.HasPrefix(event.Name, "consul:") {
		return
	}

	switch event.Name {
	case newLeaderEvent:
		c.logger.Printf("[INFO] consul: New leader elected: %s", event.Payload)

		// Trigger the callback
		if c.config.ServerUp != nil {
			c.config.ServerUp()
		}
	default:
		c.logger.Printf("[WARN] consul: Unhandled local event: %v", event)
	}
}

// RPC is used to forward an RPC call to a consul server, or fail if no servers
func (c *Client) RPC(method string, args interface{}, reply interface{}) error {
	// Check the last rpc time
	var server net.Addr
	if time.Now().Sub(c.lastRPCTime) < clientRPCCache {
		server = c.lastServer
		if server != nil {
			goto TRY_RPC
		}
	}

	// Bail if we can't find any servers
	c.consulLock.RLock()
	if len(c.consuls) == 0 {
		c.consulLock.RUnlock()
		return structs.ErrNoServers
	}

	// Select a random addr
	server = c.consuls[rand.Int31()%int32(len(c.consuls))]
	c.consulLock.RUnlock()

	// Forward to remote Consul
TRY_RPC:
	if err := c.connPool.RPC(server, method, args, reply); err != nil {
		c.lastServer = nil
		c.lastRPCTime = time.Time{}
		return err
	}

	// Cache the last server
	c.lastServer = server
	c.lastRPCTime = time.Now()
	return nil
}

// Stats is used to return statistics for debugging and insight
// for various sub-systems
func (c *Client) Stats() map[string]map[string]string {
	toString := func(v uint64) string {
		return strconv.FormatUint(v, 10)
	}
	stats := map[string]map[string]string{
		"consul": map[string]string{
			"server":        "false",
			"known_servers": toString(uint64(len(c.consuls))),
		},
		"serf_lan": c.serf.Stats(),
	}
	return stats
}
<|endoftext|>


<commit_before>package gobstore

import (
	"bytes"
	"code.google.com/p/go.crypto/bcrypt"
	"crypto/sha1"
	"encoding/base64"
	"encoding/gob"
	"io/ioutil"
	"os"
)

// The UserGobStore implements the UserStore interface to store Users on disk
// encoded in the gob format.
type UserGobStore struct {
	Path string
}

func NewUserGobStore(path string) *UserGobStore {
	store := UserGobStore{path}
	err := os.MkdirAll(path, 0700)
	if err != nil {
		panic(err)
	}
	return &store
}

func (store UserGobStore) Find(email string) (*User, error) {
	emailSha := generateHash(email)
	userGob, err := ioutil.ReadFile(store.Path + emailSha + ".gob")
	if err != nil {
		return nil, err
	}

	userGobBuf := bytes.NewBuffer(userGob)
	decoder := gob.NewDecoder(userGobBuf)
	user := User{}
	err = decoder.Decode(&user)

	if err != nil {
		return nil, err
	}

	return &user, nil
}

func (store UserGobStore) Save(user User) error {
	emailSha := generateHash(user.Email)
	userGobBuf := new(bytes.Buffer)
	encoder := gob.NewEncoder(userGobBuf)
	encoder.Encode(user)
	return ioutil.WriteFile(store.Path+emailSha+".gob", userGobBuf.Bytes(), 0600)
}

func generateHash(str string) string {
	hasher := sha1.New()
	hasher.Write([]byte(str))
	return base64.URLEncoding.EncodeToString(hasher.Sum(nil))
}
<commit_msg>Extract gobstore.<commit_after>package gobstore

import (
	"bytes"
	"crypto/sha1"
	"encoding/base64"
	"encoding/gob"
	"github.com/jroes/goth/user"
	"io/ioutil"
	"os"
)

// The UserGobStore implements the UserStore interface to store Users on disk
// encoded in the gob format.
type UserGobStore struct {
	Path string
}

func NewUserGobStore(path string) *UserGobStore {
	store := UserGobStore{path}
	err := os.MkdirAll(path, 0700)
	if err != nil {
		panic(err)
	}
	return &store
}

func (store UserGobStore) Find(email string) (*user.User, error) {
	emailSha := generateHash(email)
	userGob, err := ioutil.ReadFile(store.Path + emailSha + ".gob")
	if err != nil {
		return nil, err
	}

	userGobBuf := bytes.NewBuffer(userGob)
	decoder := gob.NewDecoder(userGobBuf)
	user := user.User{}
	err = decoder.Decode(&user)

	if err != nil {
		return nil, err
	}

	return &user, nil
}

func (store UserGobStore) Save(user user.User) error {
	emailSha := generateHash(user.Email)
	userGobBuf := new(bytes.Buffer)
	encoder := gob.NewEncoder(userGobBuf)
	encoder.Encode(user)
	return ioutil.WriteFile(store.Path+emailSha+".gob", userGobBuf.Bytes(), 0600)
}

func generateHash(str string) string {
	hasher := sha1.New()
	hasher.Write([]byte(str))
	return base64.URLEncoding.EncodeToString(hasher.Sum(nil))
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"github.com/bitly/nsq/nsq"
	"log"
	"net"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"
)

var (
	topic            = flag.String("topic", "", "nsq topic")
	channel          = flag.String("channel", "nsq_to_tcp", "nsq channel")
	maxInFlight      = flag.Int("max-in-flight", 200, "max number of messages to allow in flight")
	lookupdHTTPAddrs = flag.String("lookupd-http-address", "127.0.0.1:4161", "lookupd http")
	port             = flag.String("port", ":1514", "log send port")
)

type MsgHandler struct {
	msg_chan chan []byte
}

func (this *MsgHandler) HandleMessage(m *nsq.Message) error {
	this.msg_chan <- m.Body
	return nil
}

func main() {
	flag.Parse()

	if *topic == "" || *channel == "" {
		log.Fatalf("--topic and --channel are required")
	}

	if *maxInFlight < 0 {
		log.Fatalf("--max-in-flight must be > 0")
	}

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP)

	r, err := nsq.NewReader(*topic, *channel)
	if err != nil {
		log.Fatalf(err.Error())
	}
	r.SetMaxInFlight(*maxInFlight)
	msg_handler := MsgHandler{make(chan []byte)}
	r.AddHandler(&msg_handler)
	lookupdlist := strings.Split(*lookupdHTTPAddrs, ",")
	exitchan := make(chan int)
	go tcp_server(*port, msg_handler.msg_chan, exitchan)
	for _, addrString := range lookupdlist {
		log.Printf("lookupd addr %s", addrString)
		err := r.ConnectToLookupd(addrString)
		if err != nil {
			log.Fatalf(err.Error())
		}
	}

	select {
	case <-r.ExitChan:
		log.Println("reader exited")
	case <-sigChan:
		r.Stop()
		log.Println("stop nsq reader")
	}
	time.Sleep(time.Second)
}

func tcp_server(port string, logchan chan []byte, exitchan chan int) {
	server, err := net.Listen("tcp", port)
	if err != nil {
		log.Fatal("server bind failed:", err)
		return
	}
	go func() {
		for {
			fd, err := server.Accept()
			if err != nil &&
				strings.Contains(err.Error(),
					"use of closed network connection") {
				break
			}
			if err != nil {
				log.Fatal("accept error", err)
				time.Sleep(time.Second)
			} else {
				go send_log(fd, logchan)
			}
		}
	}()
	<-exitchan
	server.Close()
	log.Println("tcp server closed")
}

func send_log(fd net.Conn, logchan chan []byte) {
	defer fd.Close()
	var err error
	for {
		msg, ok := <-logchan
		if !ok {
			break
		}
		_, err = fd.Write(msg)
		if err != nil {
			log.Fatal("write failed ", err)
			break
		}
	}
}
<commit_msg>handle sigpipe, which trapped by tcpconn.write<commit_after>package main

import (
	"flag"
	"github.com/bitly/nsq/nsq"
	"log"
	"net"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"
)

var (
	topic            = flag.String("topic", "", "nsq topic")
	channel          = flag.String("channel", "nsq_to_tcp", "nsq channel")
	maxInFlight      = flag.Int("max-in-flight", 200, "max number of messages to allow in flight")
	lookupdHTTPAddrs = flag.String("lookupd-http-address", "127.0.0.1:4161", "lookupd http")
	port             = flag.String("port", ":1514", "log send port")
)

type MsgHandler struct {
	msg_chan chan []byte
}

func (this *MsgHandler) HandleMessage(m *nsq.Message) error {
	this.msg_chan <- m.Body
	return nil
}

func main() {
	flag.Parse()

	if *topic == "" || *channel == "" {
		log.Fatalf("--topic and --channel are required")
	}

	if *maxInFlight < 0 {
		log.Fatalf("--max-in-flight must be > 0")
	}

	sigChan := make(chan os.Signal, 1)
	sig2Chan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP)

	signal.Notify(sig2Chan, syscall.SIGPIPE)

	go func() {
		for {
			<- sig2Chan
			log.Println("tcp epipe")
		}
	}()
	r, err := nsq.NewReader(*topic, *channel)
	if err != nil {
		log.Fatalf(err.Error())
	}
	r.SetMaxInFlight(*maxInFlight)
	msg_handler := MsgHandler{make(chan []byte)}
	r.AddHandler(&msg_handler)
	lookupdlist := strings.Split(*lookupdHTTPAddrs, ",")
	exitchan := make(chan int)
	go tcp_server(*port, msg_handler.msg_chan, exitchan)
	for _, addrString := range lookupdlist {
		log.Printf("lookupd addr %s", addrString)
		err := r.ConnectToLookupd(addrString)
		if err != nil {
			log.Fatalf(err.Error())
		}
	}

	select {
	case <-r.ExitChan:
		log.Println("reader exited")
	case <-sigChan:
		r.Stop()
		exitchan <- 1
		log.Println("stop all")
	}
	time.Sleep(time.Second)
}

func tcp_server(port string, logchan chan []byte, exitchan chan int) {
	server, err := net.Listen("tcp", port)
	if err != nil {
		log.Fatal("server bind failed:", err)
		return
	}
	go func() {
		for {
			fd, err := server.Accept()
			if err != nil &&
				strings.Contains(err.Error(),
					"use of closed network connection") {
				break
			}
			if err != nil {
				log.Fatal("accept error", err)
				time.Sleep(time.Second)
			} else {
				go send_log(fd, logchan)
			}
		}
	}()
	<-exitchan
	server.Close()
	log.Println("tcp server closed")
}

func send_log(fd net.Conn, logchan chan []byte) {
	defer fd.Close()
	var err error
	for {
		msg, ok := <-logchan
		if !ok {
			break
		}
		_, err = fd.Write(msg)
		if err != nil {
			log.Fatal("write failed ", err)
			break
		}
	}
}
<|endoftext|>


<commit_before><commit_msg>netchan: fix comment typo.<commit_after><|endoftext|>


<commit_before>// Copyright 2016 Google Inc. All Rights Reserved.
// This file is available under the Apache license.

package vm

import (
	"fmt"
	"regexp"
	"time"

	"github.com/golang/glog"
	"github.com/google/mtail/metrics"
	"github.com/google/mtail/metrics/datum"
	"github.com/pkg/errors"
)

// compiler is data for the code generator.
type codegen struct {
	name string // Name of the program.

	errors ErrorList // Compile errors.
	obj    object    // The object to return
	l      []int     // jump table

	decos []*decoNode // Decorator stack to unwind
}

// CodeGen is the function that compiles the program to bytecode and data.
func CodeGen(name string, ast astNode) (*object, error) {
	c := &codegen{name: name}
	Walk(c, ast)
	c.writeJumps()
	if len(c.errors) > 0 {
		return nil, c.errors
	}
	return &c.obj, nil
}

func (c *codegen) errorf(pos *position, format string, args ...interface{}) {
	e := "Internal compiler error, aborting compilation: " + fmt.Sprintf(format, args...)
	c.errors.Add(pos, e)
}

func (c *codegen) emit(i instr) {
	c.obj.prog = append(c.obj.prog, i)
}

// newLabel creates a new label to jump to
func (c *codegen) newLabel() (l int) {
	l = len(c.l)
	c.l = append(c.l, -1)
	return
}

// setLabel points a label to the next instruction
func (c *codegen) setLabel(l int) {
	c.l[l] = c.pc() + 1
}

// pc returns the program offset of the last instruction
func (c *codegen) pc() int {
	return len(c.obj.prog) - 1
}

func (c *codegen) VisitBefore(node astNode) Visitor {
	switch n := node.(type) {

	case *declNode:
		var name string
		if n.exportedName != "" {
			name = n.exportedName
		} else {
			name = n.name
		}
		// If the Type is not in the map, then default to metrics.Int.  This is
		// a hack for metrics that no type can be inferred, retaining
		// historical behaviour.
		t := n.Type()
		if IsDimension(t) {
			t = t.(*TypeOperator).Args[len(t.(*TypeOperator).Args)-1]
		}
		var dtyp datum.Type
		switch {
		case Equals(Float, t):
			dtyp = metrics.Float
		default:
			if !IsComplete(t) {
				glog.Infof("Incomplete type %v for %#v", t, n)
			}
			dtyp = metrics.Int
		}
		m := metrics.NewMetric(name, c.name, n.kind, dtyp, n.keys...)
		m.SetSource(n.Pos().String())
		// Scalar counters can be initialized to zero.  Dimensioned counters we
		// don't know the values of the labels yet.  Gauges and Timers we can't
		// assume start at zero.
		if len(n.keys) == 0 && n.kind == metrics.Counter {
			d, err := m.GetDatum()
			if err != nil {
				c.errorf(n.Pos(), "%s", err)
				return nil
			}
			// Initialize to zero at the zero time.
			if dtyp == metrics.Int {
				datum.SetInt(d, 0, time.Unix(0, 0))
			} else {
				datum.SetFloat(d, 0, time.Unix(0, 0))
			}
		}
		m.Hidden = n.hidden
		(*n.sym).Binding = m
		n.sym.Addr = len(c.obj.m)
		c.obj.m = append(c.obj.m, m)
		return nil

	case *condNode:
		lElse := c.newLabel()
		lEnd := c.newLabel()
		if n.cond != nil {
			Walk(c, n.cond)
			c.emit(instr{jnm, lElse})
		}
		// Set matched flag false for children.
		c.emit(instr{setmatched, false})
		Walk(c, n.truthNode)
		// Re-set matched flag to true for rest of current block.
		c.emit(instr{setmatched, true})
		if n.elseNode != nil {
			c.emit(instr{jmp, lEnd})
		}
		c.setLabel(lElse)
		if n.elseNode != nil {
			Walk(c, n.elseNode)
		}
		c.setLabel(lEnd)
		return nil

	case *patternExprNode:
		re, err := regexp.Compile(n.pattern)
		if err != nil {
			c.errorf(n.Pos(), "%s", err)
			return nil
		}
		c.obj.re = append(c.obj.re, re)
		// Store the location of this regular expression in the patterNode
		n.index = len(c.obj.re) - 1
		c.emit(instr{match, n.index})

	case *stringConstNode:
		c.obj.str = append(c.obj.str, n.text)
		c.emit(instr{str, len(c.obj.str) - 1})

	case *intConstNode:
		c.emit(instr{push, n.i})

	case *floatConstNode:
		c.emit(instr{push, n.f})

	case *idNode:
		if n.sym.Kind != VarSymbol {
			break
		}
		if n.sym == nil || n.sym.Binding == nil {
			c.errorf(n.Pos(), "No metric bound to identifier %q", n.name)
			return nil
		}
		c.emit(instr{mload, n.sym.Addr})
		m := n.sym.Binding.(*metrics.Metric)
		c.emit(instr{dload, len(m.Keys)})

	case *caprefNode:
		if n.sym == nil || n.sym.Binding == nil {
			c.errorf(n.Pos(), "No regular expression bound to capref %q", n.name)
			return nil
		}
		rn := n.sym.Binding.(*patternExprNode)
		// rn.index contains the index of the compiled regular expression object
		// in the re slice of the object code
		c.emit(instr{push, rn.index})
		// n.sym.addr is the capture group offset
		c.emit(instr{capref, n.sym.Addr})

	case *decoDefNode:
		// Do nothing, defs are inlined.
		return nil

	case *decoNode:
		// Put the current block on the stack
		c.decos = append(c.decos, n)
		if n.def == nil {
			c.errorf(n.Pos(), "No definition found for decorator %q", n.name)
			return nil
		}
		// then iterate over the decorator's nodes
		Walk(c, n.def.block)
		c.decos = c.decos[:len(c.decos)-1]
		return nil

	case *nextNode:
		// Visit the 'next' block on the decorated block stack
		deco := c.decos[len(c.decos)-1]
		Walk(c, deco.block)
		return nil

	case *otherwiseNode:
		c.emit(instr{op: otherwise})

	case *delNode:
		Walk(c, n.n)
		// overwrite the dload instruction
		pc := c.pc()
		c.obj.prog[pc].op = del

	case *binaryExprNode:
		switch n.op {
		case AND:
			lFalse := c.newLabel()
			lEnd := c.newLabel()
			Walk(c, n.lhs)
			c.emit(instr{jnm, lFalse})
			Walk(c, n.rhs)
			c.emit(instr{jnm, lFalse})
			c.emit(instr{push, true})
			c.emit(instr{jmp, lEnd})
			c.setLabel(lFalse)
			c.emit(instr{push, false})
			c.setLabel(lEnd)
			return nil

		case OR:
			lTrue := c.newLabel()
			lEnd := c.newLabel()
			Walk(c, n.lhs)
			c.emit(instr{jm, lTrue})
			Walk(c, n.rhs)
			c.emit(instr{jm, lTrue})
			c.emit(instr{push, false})
			c.emit(instr{jmp, lEnd})
			c.setLabel(lTrue)
			c.emit(instr{push, true})
			c.setLabel(lEnd)
			return nil

		case ADD_ASSIGN:
			if Equals(n.Type(), Float) {
				// Double-emit the lhs so that it can be assigned to
				Walk(c, n.lhs)
			}

		default:
			// Didn't handle it, let normal walk proceed
			return c
		}

	}

	return c
}

var typedOperators = map[int]map[Type]opcode{
	PLUS: {Int: iadd,
		Float:  fadd,
		String: cat},
	MINUS: {Int: isub,
		Float: fsub},
	MUL: {Int: imul,
		Float: fmul},
	DIV: {Int: idiv,
		Float: fdiv},
	MOD: {Int: imod,
		Float: fmod},
	POW: {Int: ipow,
		Float: fpow},
	ASSIGN: {Int: iset,
		Float: fset},
}

func (c *codegen) VisitAfter(node astNode) {
	switch n := node.(type) {
	case *builtinNode:
		arglen := 0
		if n.args != nil {
			arglen = len(n.args.(*exprlistNode).children)
		}
		switch n.name {
		case "bool":
		// TODO(jaq): Nothing, no support in VM yet.

		case "int", "float", "string":
			// len args should be 1
			if arglen > 1 {
				c.errorf(n.Pos(), "internal error, too many arguments to builtin %q: %#v", n.name, n)
				return
			}
			if err := c.emitConversion(n.args.(*exprlistNode).children[0].Type(), n.Type()); err != nil {
				c.errorf(n.Pos(), "internal error: %s on node %v", err.Error(), n)
				return
			}

		default:
			c.emit(instr{builtin[n.name], arglen})
		}
	case *unaryExprNode:
		switch n.op {
		case INC:
			c.emit(instr{op: inc})
		case NOT:
			c.emit(instr{op: not})
		}
	case *binaryExprNode:
		switch n.op {
		case LT, GT, LE, GE, EQ, NE:
			lFail := c.newLabel()
			lEnd := c.newLabel()
			var cmpArg int
			var jumpOp opcode
			switch n.op {
			case LT:
				cmpArg = -1
				jumpOp = jnm
			case GT:
				cmpArg = 1
				jumpOp = jnm
			case LE:
				cmpArg = 1
				jumpOp = jm
			case GE:
				cmpArg = -1
				jumpOp = jm
			case EQ:
				cmpArg = 0
				jumpOp = jnm
			case NE:
				cmpArg = 0
				jumpOp = jm
			}
			c.emit(instr{cmp, cmpArg})
			c.emit(instr{jumpOp, lFail})
			c.emit(instr{push, true})
			c.emit(instr{jmp, lEnd})
			c.setLabel(lFail)
			c.emit(instr{push, false})
			c.setLabel(lEnd)
		case ADD_ASSIGN:
			// When operand is not nil, inc pops the delta from the stack.
			// TODO(jaq): string concatenation, once datums can hold strings.
			switch {
			case Equals(n.Type(), Int):
				c.emit(instr{inc, 0})
			case Equals(n.Type(), Float):
				// Already walked the lhs and rhs of this expression
				c.emit(instr{fadd, nil})
				// And a second lhs
				c.emit(instr{fset, nil})
			default:
				c.errorf(n.Pos(), "Internal error: invalid type for add-assignment: %v", n.op)
				return
			}
		case PLUS, MINUS, MUL, DIV, MOD, POW, ASSIGN:
			opmap, ok := typedOperators[n.op]
			if !ok {
				c.errorf(n.Pos(), "Internal error: no typed operator for binary expression %v", n.op)
				return
			}
			emitflag := false
			for t, opcode := range opmap {
				if Equals(n.Type(), t) {
					c.emit(instr{op: opcode})
					emitflag = true
					break
				}
			}
			if !emitflag {
				c.errorf(n.Pos(), "Invalid type for binary expression: %v", n.Type())
				return
			}
		case BITAND:
			c.emit(instr{op: and})
		case BITOR:
			c.emit(instr{op: or})
		case XOR:
			c.emit(instr{op: xor})
		case SHL:
			c.emit(instr{op: shl})
		case SHR:
			c.emit(instr{op: shr})
		}

	case *convNode:
		if err := c.emitConversion(n.n.Type(), n.Type()); err != nil {
			c.errorf(n.Pos(), "internal error: %s on node %v", err.Error(), n)
			return
		}
	}
}

func (c *codegen) emitConversion(inType, outType Type) error {
	glog.Infof("Conversion: %q to %q", inType, outType)
	if Equals(Int, inType) && Equals(Float, outType) {
		c.emit(instr{op: i2f})
	} else if Equals(String, inType) && Equals(Float, outType) {
		c.emit(instr{op: s2f})
	} else if Equals(String, inType) && Equals(Int, outType) {
		c.emit(instr{op: s2i})
	} else if Equals(Float, inType) && Equals(String, outType) {
		c.emit(instr{op: f2s})
	} else if Equals(Int, inType) && Equals(String, outType) {
		c.emit(instr{op: i2s})
	} else {
		return errors.Errorf("can't convert %q to %q", inType, outType)
	}
	return nil
}

func (c *codegen) writeJumps() {
	for j, i := range c.obj.prog {
		switch i.op {
		case jmp, jm, jnm:
			index := i.opnd.(int)
			if index > len(c.l) {
				c.errorf(nil, "internal error: no jump at label %v, table is %v", i.opnd, c.l)
				continue
			}
			offset := c.l[index]
			if offset < 0 {
				c.errorf(nil, "internal error: offset for label %v is negative, table is %v", i.opnd, c.l)
				continue
			}
			c.obj.prog[j].opnd = c.l[index]
		}
	}
}
<commit_msg>Clean up error messages, report on unexpected op.<commit_after>// Copyright 2016 Google Inc. All Rights Reserved.
// This file is available under the Apache license.

package vm

import (
	"fmt"
	"regexp"
	"time"

	"github.com/golang/glog"
	"github.com/google/mtail/metrics"
	"github.com/google/mtail/metrics/datum"
	"github.com/pkg/errors"
)

// compiler is data for the code generator.
type codegen struct {
	name string // Name of the program.

	errors ErrorList // Compile errors.
	obj    object    // The object to return
	l      []int     // jump table

	decos []*decoNode // Decorator stack to unwind
}

// CodeGen is the function that compiles the program to bytecode and data.
func CodeGen(name string, ast astNode) (*object, error) {
	c := &codegen{name: name}
	Walk(c, ast)
	c.writeJumps()
	if len(c.errors) > 0 {
		return nil, c.errors
	}
	return &c.obj, nil
}

func (c *codegen) errorf(pos *position, format string, args ...interface{}) {
	e := "Internal compiler error, aborting compilation: " + fmt.Sprintf(format, args...)
	c.errors.Add(pos, e)
}

func (c *codegen) emit(i instr) {
	c.obj.prog = append(c.obj.prog, i)
}

// newLabel creates a new label to jump to
func (c *codegen) newLabel() (l int) {
	l = len(c.l)
	c.l = append(c.l, -1)
	return
}

// setLabel points a label to the next instruction
func (c *codegen) setLabel(l int) {
	c.l[l] = c.pc() + 1
}

// pc returns the program offset of the last instruction
func (c *codegen) pc() int {
	return len(c.obj.prog) - 1
}

func (c *codegen) VisitBefore(node astNode) Visitor {
	switch n := node.(type) {

	case *declNode:
		var name string
		if n.exportedName != "" {
			name = n.exportedName
		} else {
			name = n.name
		}
		// If the Type is not in the map, then default to metrics.Int.  This is
		// a hack for metrics that no type can be inferred, retaining
		// historical behaviour.
		t := n.Type()
		if IsDimension(t) {
			t = t.(*TypeOperator).Args[len(t.(*TypeOperator).Args)-1]
		}
		var dtyp datum.Type
		switch {
		case Equals(Float, t):
			dtyp = metrics.Float
		default:
			if !IsComplete(t) {
				glog.Infof("Incomplete type %v for %#v", t, n)
			}
			dtyp = metrics.Int
		}
		m := metrics.NewMetric(name, c.name, n.kind, dtyp, n.keys...)
		m.SetSource(n.Pos().String())
		// Scalar counters can be initialized to zero.  Dimensioned counters we
		// don't know the values of the labels yet.  Gauges and Timers we can't
		// assume start at zero.
		if len(n.keys) == 0 && n.kind == metrics.Counter {
			d, err := m.GetDatum()
			if err != nil {
				c.errorf(n.Pos(), "%s", err)
				return nil
			}
			// Initialize to zero at the zero time.
			if dtyp == metrics.Int {
				datum.SetInt(d, 0, time.Unix(0, 0))
			} else {
				datum.SetFloat(d, 0, time.Unix(0, 0))
			}
		}
		m.Hidden = n.hidden
		(*n.sym).Binding = m
		n.sym.Addr = len(c.obj.m)
		c.obj.m = append(c.obj.m, m)
		return nil

	case *condNode:
		lElse := c.newLabel()
		lEnd := c.newLabel()
		if n.cond != nil {
			Walk(c, n.cond)
			c.emit(instr{jnm, lElse})
		}
		// Set matched flag false for children.
		c.emit(instr{setmatched, false})
		Walk(c, n.truthNode)
		// Re-set matched flag to true for rest of current block.
		c.emit(instr{setmatched, true})
		if n.elseNode != nil {
			c.emit(instr{jmp, lEnd})
		}
		c.setLabel(lElse)
		if n.elseNode != nil {
			Walk(c, n.elseNode)
		}
		c.setLabel(lEnd)
		return nil

	case *patternExprNode:
		re, err := regexp.Compile(n.pattern)
		if err != nil {
			c.errorf(n.Pos(), "%s", err)
			return nil
		}
		c.obj.re = append(c.obj.re, re)
		// Store the location of this regular expression in the patterNode
		n.index = len(c.obj.re) - 1
		c.emit(instr{match, n.index})

	case *stringConstNode:
		c.obj.str = append(c.obj.str, n.text)
		c.emit(instr{str, len(c.obj.str) - 1})

	case *intConstNode:
		c.emit(instr{push, n.i})

	case *floatConstNode:
		c.emit(instr{push, n.f})

	case *idNode:
		if n.sym.Kind != VarSymbol {
			break
		}
		if n.sym == nil || n.sym.Binding == nil {
			c.errorf(n.Pos(), "No metric bound to identifier %q", n.name)
			return nil
		}
		c.emit(instr{mload, n.sym.Addr})
		m := n.sym.Binding.(*metrics.Metric)
		c.emit(instr{dload, len(m.Keys)})

	case *caprefNode:
		if n.sym == nil || n.sym.Binding == nil {
			c.errorf(n.Pos(), "No regular expression bound to capref %q", n.name)
			return nil
		}
		rn := n.sym.Binding.(*patternExprNode)
		// rn.index contains the index of the compiled regular expression object
		// in the re slice of the object code
		c.emit(instr{push, rn.index})
		// n.sym.addr is the capture group offset
		c.emit(instr{capref, n.sym.Addr})

	case *decoDefNode:
		// Do nothing, defs are inlined.
		return nil

	case *decoNode:
		// Put the current block on the stack
		c.decos = append(c.decos, n)
		if n.def == nil {
			c.errorf(n.Pos(), "No definition found for decorator %q", n.name)
			return nil
		}
		// then iterate over the decorator's nodes
		Walk(c, n.def.block)
		c.decos = c.decos[:len(c.decos)-1]
		return nil

	case *nextNode:
		// Visit the 'next' block on the decorated block stack
		deco := c.decos[len(c.decos)-1]
		Walk(c, deco.block)
		return nil

	case *otherwiseNode:
		c.emit(instr{op: otherwise})

	case *delNode:
		Walk(c, n.n)
		// overwrite the dload instruction
		pc := c.pc()
		c.obj.prog[pc].op = del

	case *binaryExprNode:
		switch n.op {
		case AND:
			lFalse := c.newLabel()
			lEnd := c.newLabel()
			Walk(c, n.lhs)
			c.emit(instr{jnm, lFalse})
			Walk(c, n.rhs)
			c.emit(instr{jnm, lFalse})
			c.emit(instr{push, true})
			c.emit(instr{jmp, lEnd})
			c.setLabel(lFalse)
			c.emit(instr{push, false})
			c.setLabel(lEnd)
			return nil

		case OR:
			lTrue := c.newLabel()
			lEnd := c.newLabel()
			Walk(c, n.lhs)
			c.emit(instr{jm, lTrue})
			Walk(c, n.rhs)
			c.emit(instr{jm, lTrue})
			c.emit(instr{push, false})
			c.emit(instr{jmp, lEnd})
			c.setLabel(lTrue)
			c.emit(instr{push, true})
			c.setLabel(lEnd)
			return nil

		case ADD_ASSIGN:
			if Equals(n.Type(), Float) {
				// Double-emit the lhs so that it can be assigned to
				Walk(c, n.lhs)
			}

		default:
			// Didn't handle it, let normal walk proceed
			return c
		}

	}

	return c
}

var typedOperators = map[int]map[Type]opcode{
	PLUS: {Int: iadd,
		Float:  fadd,
		String: cat},
	MINUS: {Int: isub,
		Float: fsub},
	MUL: {Int: imul,
		Float: fmul},
	DIV: {Int: idiv,
		Float: fdiv},
	MOD: {Int: imod,
		Float: fmod},
	POW: {Int: ipow,
		Float: fpow},
	ASSIGN: {Int: iset,
		Float: fset},
}

func (c *codegen) VisitAfter(node astNode) {
	switch n := node.(type) {
	case *builtinNode:
		arglen := 0
		if n.args != nil {
			arglen = len(n.args.(*exprlistNode).children)
		}
		switch n.name {
		case "bool":
		// TODO(jaq): Nothing, no support in VM yet.

		case "int", "float", "string":
			// len args should be 1
			if arglen > 1 {
				c.errorf(n.Pos(), "too many arguments to builtin %q: %#v", n.name, n)
				return
			}
			if err := c.emitConversion(n.args.(*exprlistNode).children[0].Type(), n.Type()); err != nil {
				c.errorf(n.Pos(), "%s on node %v", err.Error(), n)
				return
			}

		default:
			c.emit(instr{builtin[n.name], arglen})
		}
	case *unaryExprNode:
		switch n.op {
		case INC:
			c.emit(instr{op: inc})
		case NOT:
			c.emit(instr{op: not})
		}
	case *binaryExprNode:
		switch n.op {
		case LT, GT, LE, GE, EQ, NE:
			lFail := c.newLabel()
			lEnd := c.newLabel()
			var cmpArg int
			var jumpOp opcode
			switch n.op {
			case LT:
				cmpArg = -1
				jumpOp = jnm
			case GT:
				cmpArg = 1
				jumpOp = jnm
			case LE:
				cmpArg = 1
				jumpOp = jm
			case GE:
				cmpArg = -1
				jumpOp = jm
			case EQ:
				cmpArg = 0
				jumpOp = jnm
			case NE:
				cmpArg = 0
				jumpOp = jm
			}
			c.emit(instr{cmp, cmpArg})
			c.emit(instr{jumpOp, lFail})
			c.emit(instr{push, true})
			c.emit(instr{jmp, lEnd})
			c.setLabel(lFail)
			c.emit(instr{push, false})
			c.setLabel(lEnd)
		case ADD_ASSIGN:
			// When operand is not nil, inc pops the delta from the stack.
			// TODO(jaq): string concatenation, once datums can hold strings.
			switch {
			case Equals(n.Type(), Int):
				c.emit(instr{inc, 0})
			case Equals(n.Type(), Float):
				// Already walked the lhs and rhs of this expression
				c.emit(instr{fadd, nil})
				// And a second lhs
				c.emit(instr{fset, nil})
			default:
				c.errorf(n.Pos(), "invalid type for add-assignment: %v", n.op)
				return
			}
		case PLUS, MINUS, MUL, DIV, MOD, POW, ASSIGN:
			opmap, ok := typedOperators[n.op]
			if !ok {
				c.errorf(n.Pos(), "no typed operator for binary expression %v", n.op)
				return
			}
			emitflag := false
			for t, opcode := range opmap {
				if Equals(n.Type(), t) {
					c.emit(instr{op: opcode})
					emitflag = true
					break
				}
			}
			if !emitflag {
				c.errorf(n.Pos(), "invalid type for binary expression: %v", n.Type())
				return
			}
		case BITAND:
			c.emit(instr{op: and})
		case BITOR:
			c.emit(instr{op: or})
		case XOR:
			c.emit(instr{op: xor})
		case SHL:
			c.emit(instr{op: shl})
		case SHR:
			c.emit(instr{op: shr})
		default:
			c.errorf(n.Pos(), "unexpected op %v", n.op)
		}

	case *convNode:
		if err := c.emitConversion(n.n.Type(), n.Type()); err != nil {
			c.errorf(n.Pos(), "internal error: %s on node %v", err.Error(), n)
			return
		}
	}
}

func (c *codegen) emitConversion(inType, outType Type) error {
	glog.Infof("Conversion: %q to %q", inType, outType)
	if Equals(Int, inType) && Equals(Float, outType) {
		c.emit(instr{op: i2f})
	} else if Equals(String, inType) && Equals(Float, outType) {
		c.emit(instr{op: s2f})
	} else if Equals(String, inType) && Equals(Int, outType) {
		c.emit(instr{op: s2i})
	} else if Equals(Float, inType) && Equals(String, outType) {
		c.emit(instr{op: f2s})
	} else if Equals(Int, inType) && Equals(String, outType) {
		c.emit(instr{op: i2s})
	} else {
		return errors.Errorf("can't convert %q to %q", inType, outType)
	}
	return nil
}

func (c *codegen) writeJumps() {
	for j, i := range c.obj.prog {
		switch i.op {
		case jmp, jm, jnm:
			index := i.opnd.(int)
			if index > len(c.l) {
				c.errorf(nil, "no jump at label %v, table is %v", i.opnd, c.l)
				continue
			}
			offset := c.l[index]
			if offset < 0 {
				c.errorf(nil, "offset for label %v is negative, table is %v", i.opnd, c.l)
				continue
			}
			c.obj.prog[j].opnd = c.l[index]
		}
	}
}
<|endoftext|>


<commit_before>package vm

import (
	"fmt"
	"github.com/goby-lang/goby/compiler"
	"github.com/goby-lang/goby/compiler/bytecode"
	"github.com/goby-lang/goby/compiler/lexer"
	"github.com/goby-lang/goby/compiler/parser"
	"github.com/goby-lang/goby/vm/errors"
	"os"
	"runtime"
	"testing"
)

func TestVM_REPLExec(t *testing.T) {
	tests := []struct {
		inputs   []string
		expected interface{}
	}{
		{
			[]string{
				`
				a, b = [3, 6]
				a + b
				`,
			},
			9,
		}, {
			[]string{
				`
				a, _, c = [1, 2, 3]
				c
				`,
			},
			3,
		}, {
			[]string{`
def foo(x)
  yield(x + 10)
end
def bar(y)
  foo(y) do |f|
	yield(f)
  end
end
def baz(z)
  bar(z + 100) do |b|
	yield(b)
  end
end
a = 0
baz(100) do |b|
  a = b
end
a
			`,
				`
class Foo
  def bar
	100
  end
end
module Baz
  class Bar
	def bar
	  Foo.new.bar
	end
  end
end
				`,
				`
Baz::Bar.new.bar + a
`,
			},
			310,
		},
		{
			[]string{`

def foo
  123
end
`, `
foo
`},
			123},
		{
			[]string{
				`
class Foo
  def bar(x)
    x + 10
  end
end
`,
				`
Foo.new.bar(90)
`},
			100},
		{
			[]string{
				`
def foo
  123
end
`,
				`
foo
`,
				`
def foo
  345
end
`,

				`
foo
`,
			}, 345},
	}

	for i, test := range tests {
		v := initTestVM()
		v.InitForREPL()

		// Initialize parser, lexer is not important here
		p := parser.New(lexer.New(""))
		p.Mode = parser.REPLMode

		program, _ := p.ParseProgram()

		// Initialize code generator, and it will behavior a little different in REPL mode.
		g := bytecode.NewGenerator()
		g.REPL = true
		g.InitTopLevelScope(program)

		for _, input := range test.inputs {
			p := parser.New(lexer.New(input))
			p.Mode = parser.REPLMode

			program, _ := p.ParseProgram()
			sets := g.GenerateInstructions(program.Statements)

			v.REPLExec(sets)
		}

		evaluated := v.GetExecResult()
		verifyExpected(t, i, evaluated, test.expected)
		// Because REPL should maintain a base call frame so that the whole program won't exit
		v.checkCFP(t, i, 1)
	}
}

func TestVM_REPLExecFail(t *testing.T) {

	tests := []struct {
		inputs   []string
		expected string
	}{
		{
			[]string{
				`raise ArgumentError`,
			},
			fmt.Sprintf("InternalError: '%s'", errors.ArgumentError),
		},
		{
			[]string{
				"NonExistentBuiltinMethod",
			},
			"NameError: uninitialized constant NonExistentBuiltinMethod",
		},
		{
			[]string{
				"Hash.notExist",
			},
			"UndefinedMethodError: Undefined Method 'notExist' for Hash",
		},
	}

	for i, test := range tests {
		v := initTestVM()
		v.InitForREPL()

		// Initialize parser, lexer is not important here
		p := parser.New(lexer.New(""))
		p.Mode = parser.REPLMode

		program, _ := p.ParseProgram()

		// Initialize code generator, and it will behavior a little different in REPL mode.
		g := bytecode.NewGenerator()
		g.REPL = true
		g.InitTopLevelScope(program)

		for _, input := range test.inputs {
			p := parser.New(lexer.New(input))
			p.Mode = parser.REPLMode

			// prevent parse errors from panicking tests
			program, Err := p.ParseProgram()
			if Err != nil {
				t.Fatalf("At case %d unexpected parse error %q", i, Err.Message)
			}
			sets := g.GenerateInstructions(program.Statements)

			v.REPLExec(sets)
		}

		evaluated := v.GetExecResult()

		if evaluated.toString() != test.expected {
			t.Fatalf("At case %d expected %s got %s", i,
				test.expected, evaluated.toString())
		}

		// Because REPL should maintain a base call frame so that the whole program won't exit
		v.checkCFP(t, i, 1)
	}
}

func initTestVM() *VM {
	fn, err := os.Getwd()

	if err != nil {
		panic(err)
	}

	v, err := New(fn, []string{})

	if err != nil {
		panic(err)
	}

	v.mode = TestMode
	return v
}

func (v *VM) testEval(t *testing.T, input, filepath string) Object {
	iss, err := compiler.CompileToInstructions(input, parser.TestMode)

	if err != nil {
		t.Helper()
		t.Errorf("Error when compiling input: %s", input)
		t.Fatal(err.Error())
	}

	v.ExecInstructions(iss, filepath)

	return v.mainThread.stack.top().Target
}

func (v *VM) checkCFP(t *testing.T, index, expectedCFP int) {
	t.Helper()
	if v.mainThread.cfp != expectedCFP {
		t.Errorf("At case %d expect main thread's cfp to be %d. got: %d", index, expectedCFP, v.mainThread.cfp)
	}
}

func (v *VM) checkSP(t *testing.T, index, expectedSp int) {
	t.Helper()
	if v.mainThread.sp != expectedSp {
		fmt.Println(v.mainThread.stack.inspect())
		t.Errorf("At case %d expect main thread's sp to be %d. got: %d", index, expectedSp, v.mainThread.sp)
	}

}

// Verification helpers

func verifyExpected(t *testing.T, i int, evaluated Object, expected interface{}) {
	t.Helper()
	if isError(evaluated) {
		t.Errorf("At test case %d: %s", i, evaluated.toString())
		return
	}

	switch expected := expected.(type) {
	case int:
		verifyIntegerObject(t, i, evaluated, expected)
	case float64:
		verifyFloatObject(t, i, evaluated, expected)
	case string:
		verifyStringObject(t, i, evaluated, expected)
	case bool:
		verifyBooleanObject(t, i, evaluated, expected)
	case []interface{}:
		verifyArrayObject(t, i, evaluated, expected)
	case nil:
		verifyNullObject(t, i, evaluated)
	default:
		t.Errorf("Unknown type %T at case %d", expected, i)
	}
}

func verifyIntegerObject(t *testing.T, i int, obj Object, expected int) bool {
	t.Helper()
	switch result := obj.(type) {
	case *IntegerObject:
		if result.value != expected {
			t.Errorf("At test case %d: object has wrong value. expect=%d, got=%d", i, expected, result.value)
			return false
		}

		return true
	case *Error:
		t.Errorf("At test case %d: %s", i, result.Message())
		return false
	default:
		t.Errorf("At test case %d: object is not Integer. got=%s (%+v).", i, obj.Class().Name, obj)
		return false
	}
}

func verifyFloatObject(t *testing.T, i int, obj Object, expected float64) bool {
	t.Helper()
	switch result := obj.(type) {
	case *FloatObject:
		if result.value != expected {
			t.Errorf("At test case %d: object has wrong value. expect=%f, got=%f", i, expected, result.value)
			return false
		}

		return true
	case *Error:
		t.Errorf("At test case %d: %s", i, result.Message())
		return false
	default:
		t.Errorf("At test case %d: object is not Float. got=%s (%+v).", i, obj.Class().Name, obj)
		return false
	}
}

func verifyNullObject(t *testing.T, i int, obj Object) bool {
	t.Helper()
	switch result := obj.(type) {
	case *NullObject:
		return true
	case *Error:
		t.Errorf("At test case %d: %s", i, result.Message())
		return false
	default:
		t.Errorf("At test case %d: object is not NULL. got=%s (%+v)", i, obj.Class().Name, obj)
		return false
	}
}

func verifyStringObject(t *testing.T, i int, obj Object, expected string) bool {
	t.Helper()
	switch result := obj.(type) {
	case *StringObject:
		if result.value != expected {
			t.Errorf("At test case %d: object has wrong value. expect=%q, got=%q", i, expected, result.value)
			return false
		}

		return true
	case *Error:
		t.Errorf(result.Message())
		return false
	default:
		t.Errorf("At test case %d: object is not String. got=%s (%+v).", i, obj.Class().Name, obj)
		return false
	}
}

func verifyBooleanObject(t *testing.T, i int, obj Object, expected bool) bool {
	t.Helper()
	switch result := obj.(type) {
	case *BooleanObject:
		if result.value != expected {
			t.Errorf("At test case %d: object has wrong value. expect=%t, got=%t", i, expected, result.value)
			return false
		}

		return true
	case *Error:
		t.Errorf(result.Message())
		return false
	default:
		t.Errorf("At test case %d: object is not Boolean. got=%s (%+v).", i, obj.Class().Name, obj)
		return false
	}
}

func verifyArrayObject(t *testing.T, index int, obj Object, expected []interface{}) bool {
	t.Helper()
	result, ok := obj.(*ArrayObject)
	if !ok {
		t.Errorf("At test case %d: object is not Array. got=%s (%+v)", index, obj.Class().Name, obj)
		return false
	}

	if len(result.Elements) != len(expected) {
		t.Errorf("Don't equals length of array. Expect %d, got=%d", len(expected), len(result.Elements))
	}

	for i := 0; i < len(result.Elements); i++ {
		verifyExpected(t, index, result.Elements[i], expected[i])
	}

	return true
}

// Same as testHashObject(), but expects a ConcurrentArray.
func verifyConcurrentArrayObject(t *testing.T, index int, obj Object, expected []interface{}) bool {
	t.Helper()
	result, ok := obj.(*ConcurrentArrayObject)
	if !ok {
		t.Errorf("At test case %d: object is not ConcurrentArray. got=%s (%+v)", index, obj.Class().Name, obj)
		return false
	}

	if len(result.InternalArray.Elements) != len(expected) {
		t.Errorf("Don't equals length of array. Expect %d, got=%d", len(expected), len(result.InternalArray.Elements))
	}

	for i := 0; i < len(result.InternalArray.Elements); i++ {
		verifyExpected(t, i, result.InternalArray.Elements[i], expected[i])
	}

	return true
}

// Same as testHashObject(), but expects a ConcurrentHash.
//
func verifyConcurrentHashObject(t *testing.T, index int, objectResult Object, expected map[string]interface{}) bool {
	t.Helper()
	result, ok := objectResult.(*ConcurrentHashObject)

	if !ok {
		t.Errorf("At test case %d: result is not ConcurrentHash. got=%s", index, objectResult.Class().Name)
		return false
	}

	pairs := make(map[string]Object)

	iterator := func(key, value interface{}) bool {
		pairs[key.(string)] = value.(Object)
		return true
	}

	result.internalMap.Range(iterator)

	return _checkHashPairs(t, pairs, expected)
}

// Tests a Hash Object, with a few limitations:
//
// - the tested hash must be shallow (no nested objects as values);
// - the test hash must have strings as keys;
// - the error message won't mention the key - only the value.
//
// The second limitation is currently the only Hash format in Goby, anyway.
//
func verifyHashObject(t *testing.T, index int, objectResult Object, expected map[string]interface{}) bool {
	t.Helper()
	result, ok := objectResult.(*HashObject)

	if !ok {
		t.Errorf("At test case %d: result is not Hash. got=%s", index, objectResult.Class().Name)
		return false
	}

	return _checkHashPairs(t, result.Pairs, expected)
}

// Testing API like testArrayObject(), but performed on bidimensional arrays.
//
// Input example:
//
//		evaluated = '[["a", 1], ["b", "2"]]'
//		expected = [][]interface{}{{"a", 1}, {"b", "2"}}
//		testBidimensionalArrayObject(t, i, evaluated, expected)
//
func verifyBidimensionalArrayObject(t *testing.T, index int, obj Object, expected [][]interface{}) bool {
	t.Helper()
	result, ok := obj.(*ArrayObject)
	if !ok {
		t.Errorf("At test case %d: object is not Array. got=%T (%+v)", index, obj, obj)
		return false
	}

	if len(result.Elements) != len(expected) {
		t.Errorf("Unexpected result size. Expect %d, got=%d", len(expected), len(result.Elements))
	}

	for i := 0; i < len(result.Elements); i++ {
		resultRow := result.Elements[i]
		expectedRow := expected[i]

		verifyArrayObject(t, index, resultRow, expectedRow)
	}

	return true
}

func isError(obj Object) bool {
	if obj != nil {
		_, ok := obj.(*Error)
		return ok
	}
	return false
}

// Internal helpers -----------------------------------------------------

func _checkHashPairs(t *testing.T, actual map[string]Object, expected map[string]interface{}) bool {
	if len(actual) != len(expected) {
		t.Errorf("Unexpected result size. Expected %d, got=%d", len(expected), len(actual))
	}

	for expectedKey, expectedValue := range expected {
		resultValue := actual[expectedKey]

		verifyExpected(t, i, resultValue, expectedValue)
	}

	return true
}

func getFilename() string {
	_, filename, _, _ := runtime.Caller(1)
	return filename
}
<commit_msg>Add test case for missing library file<commit_after>package vm

import (
	"fmt"
	"github.com/goby-lang/goby/compiler"
	"github.com/goby-lang/goby/compiler/bytecode"
	"github.com/goby-lang/goby/compiler/lexer"
	"github.com/goby-lang/goby/compiler/parser"
	"github.com/goby-lang/goby/vm/errors"
	"os"
	"path/filepath"
	"runtime"
	"testing"
)

func TestVM_REPLExec(t *testing.T) {
	tests := []struct {
		inputs   []string
		expected interface{}
	}{
		{
			[]string{
				`
				a, b = [3, 6]
				a + b
				`,
			},
			9,
		}, {
			[]string{
				`
				a, _, c = [1, 2, 3]
				c
				`,
			},
			3,
		}, {
			[]string{`
def foo(x)
  yield(x + 10)
end
def bar(y)
  foo(y) do |f|
	yield(f)
  end
end
def baz(z)
  bar(z + 100) do |b|
	yield(b)
  end
end
a = 0
baz(100) do |b|
  a = b
end
a
			`,
				`
class Foo
  def bar
	100
  end
end
module Baz
  class Bar
	def bar
	  Foo.new.bar
	end
  end
end
				`,
				`
Baz::Bar.new.bar + a
`,
			},
			310,
		},
		{
			[]string{`

def foo
  123
end
`, `
foo
`},
			123},
		{
			[]string{
				`
class Foo
  def bar(x)
    x + 10
  end
end
`,
				`
Foo.new.bar(90)
`},
			100},
		{
			[]string{
				`
def foo
  123
end
`,
				`
foo
`,
				`
def foo
  345
end
`,

				`
foo
`,
			}, 345},
	}

	for i, test := range tests {
		v := initTestVM()
		v.InitForREPL()

		// Initialize parser, lexer is not important here
		p := parser.New(lexer.New(""))
		p.Mode = parser.REPLMode

		program, _ := p.ParseProgram()

		// Initialize code generator, and it will behavior a little different in REPL mode.
		g := bytecode.NewGenerator()
		g.REPL = true
		g.InitTopLevelScope(program)

		for _, input := range test.inputs {
			p := parser.New(lexer.New(input))
			p.Mode = parser.REPLMode

			program, _ := p.ParseProgram()
			sets := g.GenerateInstructions(program.Statements)

			v.REPLExec(sets)
		}

		evaluated := v.GetExecResult()
		verifyExpected(t, i, evaluated, test.expected)
		// Because REPL should maintain a base call frame so that the whole program won't exit
		v.checkCFP(t, i, 1)
	}
}

func TestVM_REPLExecFail(t *testing.T) {

	tests := []struct {
		inputs   []string
		expected string
	}{
		{
			[]string{
				`raise ArgumentError`,
			},
			fmt.Sprintf("InternalError: '%s'", errors.ArgumentError),
		},
		{
			[]string{
				"NonExistentBuiltinMethod",
			},
			"NameError: uninitialized constant NonExistentBuiltinMethod",
		},
		{
			[]string{
				"Hash.notExist",
			},
			"UndefinedMethodError: Undefined Method 'notExist' for Hash",
		},
	}

	for i, test := range tests {
		v := initTestVM()
		v.InitForREPL()

		// Initialize parser, lexer is not important here
		p := parser.New(lexer.New(""))
		p.Mode = parser.REPLMode

		program, _ := p.ParseProgram()

		// Initialize code generator, and it will behavior a little different in REPL mode.
		g := bytecode.NewGenerator()
		g.REPL = true
		g.InitTopLevelScope(program)

		for _, input := range test.inputs {
			p := parser.New(lexer.New(input))
			p.Mode = parser.REPLMode

			// prevent parse errors from panicking tests
			program, Err := p.ParseProgram()
			if Err != nil {
				t.Fatalf("At case %d unexpected parse error %q", i, Err.Message)
			}
			sets := g.GenerateInstructions(program.Statements)

			v.REPLExec(sets)
		}

		evaluated := v.GetExecResult()

		if evaluated.toString() != test.expected {
			t.Fatalf("At case %d expected %s got %s", i,
				test.expected, evaluated.toString())
		}

		// Because REPL should maintain a base call frame so that the whole program won't exit
		v.checkCFP(t, i, 1)
	}
}

func initTestVM() *VM {
	fn, err := os.Getwd()

	if err != nil {
		panic(err)
	}

	v, err := New(fn, []string{})

	if err != nil {
		panic(err)
	}

	v.mode = TestMode
	return v
}

func TestLoadingGobyLibraryFail(t *testing.T) {
	vm := initTestVM()

	libFileFullPath := filepath.Join(vm.projectRoot, "lib/_library_not_existing.gb")
	expectedErrorMessage := fmt.Sprintf("open %s: no such file or directory", libFileFullPath)

	err := vm.mainThread.execGobyLib("_library_not_existing.gb")

	if err.Error() != expectedErrorMessage {
		t.Errorf("Unexpected error message: %s", err.Error())
	}
}

func (v *VM) testEval(t *testing.T, input, filepath string) Object {
	iss, err := compiler.CompileToInstructions(input, parser.TestMode)

	if err != nil {
		t.Helper()
		t.Errorf("Error when compiling input: %s", input)
		t.Fatal(err.Error())
	}

	v.ExecInstructions(iss, filepath)

	return v.mainThread.stack.top().Target
}

func (v *VM) checkCFP(t *testing.T, index, expectedCFP int) {
	t.Helper()
	if v.mainThread.cfp != expectedCFP {
		t.Errorf("At case %d expect main thread's cfp to be %d. got: %d", index, expectedCFP, v.mainThread.cfp)
	}
}

func (v *VM) checkSP(t *testing.T, index, expectedSp int) {
	t.Helper()
	if v.mainThread.sp != expectedSp {
		fmt.Println(v.mainThread.stack.inspect())
		t.Errorf("At case %d expect main thread's sp to be %d. got: %d", index, expectedSp, v.mainThread.sp)
	}

}

// Verification helpers

func verifyExpected(t *testing.T, i int, evaluated Object, expected interface{}) {
	t.Helper()
	if isError(evaluated) {
		t.Errorf("At test case %d: %s", i, evaluated.toString())
		return
	}

	switch expected := expected.(type) {
	case int:
		verifyIntegerObject(t, i, evaluated, expected)
	case float64:
		verifyFloatObject(t, i, evaluated, expected)
	case string:
		verifyStringObject(t, i, evaluated, expected)
	case bool:
		verifyBooleanObject(t, i, evaluated, expected)
	case []interface{}:
		verifyArrayObject(t, i, evaluated, expected)
	case nil:
		verifyNullObject(t, i, evaluated)
	default:
		t.Errorf("Unknown type %T at case %d", expected, i)
	}
}

func verifyIntegerObject(t *testing.T, i int, obj Object, expected int) bool {
	t.Helper()
	switch result := obj.(type) {
	case *IntegerObject:
		if result.value != expected {
			t.Errorf("At test case %d: object has wrong value. expect=%d, got=%d", i, expected, result.value)
			return false
		}

		return true
	case *Error:
		t.Errorf("At test case %d: %s", i, result.Message())
		return false
	default:
		t.Errorf("At test case %d: object is not Integer. got=%s (%+v).", i, obj.Class().Name, obj)
		return false
	}
}

func verifyFloatObject(t *testing.T, i int, obj Object, expected float64) bool {
	t.Helper()
	switch result := obj.(type) {
	case *FloatObject:
		if result.value != expected {
			t.Errorf("At test case %d: object has wrong value. expect=%f, got=%f", i, expected, result.value)
			return false
		}

		return true
	case *Error:
		t.Errorf("At test case %d: %s", i, result.Message())
		return false
	default:
		t.Errorf("At test case %d: object is not Float. got=%s (%+v).", i, obj.Class().Name, obj)
		return false
	}
}

func verifyNullObject(t *testing.T, i int, obj Object) bool {
	t.Helper()
	switch result := obj.(type) {
	case *NullObject:
		return true
	case *Error:
		t.Errorf("At test case %d: %s", i, result.Message())
		return false
	default:
		t.Errorf("At test case %d: object is not NULL. got=%s (%+v)", i, obj.Class().Name, obj)
		return false
	}
}

func verifyStringObject(t *testing.T, i int, obj Object, expected string) bool {
	t.Helper()
	switch result := obj.(type) {
	case *StringObject:
		if result.value != expected {
			t.Errorf("At test case %d: object has wrong value. expect=%q, got=%q", i, expected, result.value)
			return false
		}

		return true
	case *Error:
		t.Errorf(result.Message())
		return false
	default:
		t.Errorf("At test case %d: object is not String. got=%s (%+v).", i, obj.Class().Name, obj)
		return false
	}
}

func verifyBooleanObject(t *testing.T, i int, obj Object, expected bool) bool {
	t.Helper()
	switch result := obj.(type) {
	case *BooleanObject:
		if result.value != expected {
			t.Errorf("At test case %d: object has wrong value. expect=%t, got=%t", i, expected, result.value)
			return false
		}

		return true
	case *Error:
		t.Errorf(result.Message())
		return false
	default:
		t.Errorf("At test case %d: object is not Boolean. got=%s (%+v).", i, obj.Class().Name, obj)
		return false
	}
}

func verifyArrayObject(t *testing.T, index int, obj Object, expected []interface{}) bool {
	t.Helper()
	result, ok := obj.(*ArrayObject)
	if !ok {
		t.Errorf("At test case %d: object is not Array. got=%s (%+v)", index, obj.Class().Name, obj)
		return false
	}

	if len(result.Elements) != len(expected) {
		t.Errorf("Don't equals length of array. Expect %d, got=%d", len(expected), len(result.Elements))
	}

	for i := 0; i < len(result.Elements); i++ {
		verifyExpected(t, index, result.Elements[i], expected[i])
	}

	return true
}

// Same as testHashObject(), but expects a ConcurrentArray.
func verifyConcurrentArrayObject(t *testing.T, index int, obj Object, expected []interface{}) bool {
	t.Helper()
	result, ok := obj.(*ConcurrentArrayObject)
	if !ok {
		t.Errorf("At test case %d: object is not ConcurrentArray. got=%s (%+v)", index, obj.Class().Name, obj)
		return false
	}

	if len(result.InternalArray.Elements) != len(expected) {
		t.Errorf("Don't equals length of array. Expect %d, got=%d", len(expected), len(result.InternalArray.Elements))
	}

	for i := 0; i < len(result.InternalArray.Elements); i++ {
		verifyExpected(t, i, result.InternalArray.Elements[i], expected[i])
	}

	return true
}

// Same as testHashObject(), but expects a ConcurrentHash.
//
func verifyConcurrentHashObject(t *testing.T, index int, objectResult Object, expected map[string]interface{}) bool {
	t.Helper()
	result, ok := objectResult.(*ConcurrentHashObject)

	if !ok {
		t.Errorf("At test case %d: result is not ConcurrentHash. got=%s", index, objectResult.Class().Name)
		return false
	}

	pairs := make(map[string]Object)

	iterator := func(key, value interface{}) bool {
		pairs[key.(string)] = value.(Object)
		return true
	}

	result.internalMap.Range(iterator)

	return _checkHashPairs(t, pairs, expected)
}

// Tests a Hash Object, with a few limitations:
//
// - the tested hash must be shallow (no nested objects as values);
// - the test hash must have strings as keys;
// - the error message won't mention the key - only the value.
//
// The second limitation is currently the only Hash format in Goby, anyway.
//
func verifyHashObject(t *testing.T, index int, objectResult Object, expected map[string]interface{}) bool {
	t.Helper()
	result, ok := objectResult.(*HashObject)

	if !ok {
		t.Errorf("At test case %d: result is not Hash. got=%s", index, objectResult.Class().Name)
		return false
	}

	return _checkHashPairs(t, result.Pairs, expected)
}

// Testing API like testArrayObject(), but performed on bidimensional arrays.
//
// Input example:
//
//		evaluated = '[["a", 1], ["b", "2"]]'
//		expected = [][]interface{}{{"a", 1}, {"b", "2"}}
//		testBidimensionalArrayObject(t, i, evaluated, expected)
//
func verifyBidimensionalArrayObject(t *testing.T, index int, obj Object, expected [][]interface{}) bool {
	t.Helper()
	result, ok := obj.(*ArrayObject)
	if !ok {
		t.Errorf("At test case %d: object is not Array. got=%T (%+v)", index, obj, obj)
		return false
	}

	if len(result.Elements) != len(expected) {
		t.Errorf("Unexpected result size. Expect %d, got=%d", len(expected), len(result.Elements))
	}

	for i := 0; i < len(result.Elements); i++ {
		resultRow := result.Elements[i]
		expectedRow := expected[i]

		verifyArrayObject(t, index, resultRow, expectedRow)
	}

	return true
}

func isError(obj Object) bool {
	if obj != nil {
		_, ok := obj.(*Error)
		return ok
	}
	return false
}

// Internal helpers -----------------------------------------------------

func _checkHashPairs(t *testing.T, actual map[string]Object, expected map[string]interface{}) bool {
	if len(actual) != len(expected) {
		t.Errorf("Unexpected result size. Expected %d, got=%d", len(expected), len(actual))
	}

	for expectedKey, expectedValue := range expected {
		resultValue := actual[expectedKey]

		verifyExpected(t, i, resultValue, expectedValue)
	}

	return true
}

func getFilename() string {
	_, filename, _, _ := runtime.Caller(1)
	return filename
}
<|endoftext|>


<commit_before><commit_msg>genlib.go: use go/format" instead of go fmt command if goreturns is disable<commit_after><|endoftext|>