<commit_before>package postgres

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"unicode"

	"github.com/caarlos0/env"
	"github.com/jmoiron/sqlx"
	// Used pg drive on sqlx
	_ "github.com/lib/pq"

	"database/sql"

	"github.com/nuveo/prest/api"
	"github.com/nuveo/prest/config"
)

const (
	pageNumberKey   = "_page"
	pageSizeKey     = "_page_size"
	defaultPageSize = 10
)

// Conn connect on PostgreSQL
// Used sqlx
func Conn() (db *sqlx.DB) {
	cfg := config.Prest{}
	env.Parse(&cfg)

	db, err := sqlx.Connect("postgres", fmt.Sprintf("user=%s dbname=%s host=%s port=%v sslmode=disable", cfg.PGUser, cfg.PGDatabase, cfg.PGHost, cfg.PGPort))
	if err != nil {
		panic(fmt.Sprintf("Unable to connection to database: %v\n", err))
	}
	return
}

// chkInvaidIdentifier return true if identifier is invalid
func chkInvaidIdentifier(identifer string) bool {
	if len(identifer) > 63 ||
		unicode.IsDigit([]rune(identifer)[0]) {
		return true
	}

	for _, v := range identifer {
		if !unicode.IsLetter(v) &&
			!unicode.IsDigit(v) &&
			v != '_' &&
			v != '.' {
			return true
		}
	}
	return false
}

// WhereByRequest create interface for queries + where
func WhereByRequest(r *http.Request, initialPlaceholderID int) (whereSyntax string, values []interface{}, err error) {
	//whereMap := make(map[string]string)
	whereKey := []string{}
	whereValues := []string{}

	u, _ := url.Parse(r.URL.String())
	pid := initialPlaceholderID
	for key, val := range u.Query() {
		if !strings.HasPrefix(key, "_") {
			keyInfo := strings.Split(key, ":")
			if len(keyInfo) > 1 {
				switch keyInfo[1] {
				case "jsonb":
					jsonField := strings.Split(keyInfo[0], "->>")
					if chkInvaidIdentifier(jsonField[0]) ||
						chkInvaidIdentifier(jsonField[1]) {
						err = errors.New("Invalid identifier")
						return
					}
					whereKey = append(whereKey, fmt.Sprintf("%s->>'%s'=$%d", jsonField[0], jsonField[1], pid))
					whereValues = append(whereValues, val[0])
				default:
					if chkInvaidIdentifier(keyInfo[0]) {
						err = errors.New("Invalid identifier")
						return
					}
					whereKey = append(whereKey, fmt.Sprintf("%s=$%d", keyInfo[0], pid))
					whereValues = append(whereValues, val[0])
				}
				continue
			}
			if chkInvaidIdentifier(key) {
				err = errors.New("Invalid identifier")
				return
			}

			whereKey = append(whereKey, fmt.Sprintf("%s=$%d", key, pid))
			whereValues = append(whereValues, val[0])

			pid++
		}
	}

	for i := 0; i < len(whereKey); i++ {
		if whereSyntax == "" {
			whereSyntax += whereKey[i]
		} else {
			whereSyntax += " AND " + whereKey[i]
		}

		values = append(values, whereValues[i])
	}

	return
}

// Query process queries
func Query(SQL string, params ...interface{}) (jsonData []byte, err error) {
	db := Conn()

	prepare, err := db.Prepare(SQL)

	if err != nil {
		return
	}

	rows, err := prepare.Query(params...)
	if err != nil {
		return
	}
	defer rows.Close()

	columns, err := rows.Columns()
	if err != nil {
		return
	}

	count := len(columns)
	tableData := make([]map[string]interface{}, 0)
	values := make([]interface{}, count)
	valuePtrs := make([]interface{}, count)
	for rows.Next() {
		for i := 0; i < count; i++ {
			valuePtrs[i] = &values[i]
		}
		rows.Scan(valuePtrs...)
		entry := make(map[string]interface{})
		for i, col := range columns {
			var v interface{}
			val := values[i]
			b, ok := val.([]byte)
			if ok {
				v = string(b)
			} else {
				v = val
			}
			entry[col] = v
		}
		tableData = append(tableData, entry)
	}
	jsonData, err = json.Marshal(tableData)

	return
}

// PaginateIfPossible func
func PaginateIfPossible(r *http.Request) (paginatedQuery string, err error) {
	u, _ := url.Parse(r.URL.String())
	values := u.Query()
	if _, ok := values[pageNumberKey]; !ok {
		paginatedQuery = ""
		return
	}
	pageNumber, err := strconv.Atoi(values[pageNumberKey][0])
	if err != nil {
		return
	}
	pageSize := defaultPageSize
	if size, ok := values[pageSizeKey]; ok {
		pageSize, err = strconv.Atoi(size[0])
		if err != nil {
			return
		}
	}
	paginatedQuery = fmt.Sprintf("LIMIT %d OFFSET(%d - 1) * %d", pageSize, pageNumber, pageSize)
	return
}

// Insert execute insert sql into a table
func Insert(database, schema, table string, body api.Request) (jsonData []byte, err error) {
	var result sql.Result
	var rowsAffected int64

	fields := make([]string, 0)
	values := make([]string, 0)
	for key, value := range body.Data {
		fields = append(fields, key)
		values = append(values, value)
	}
	colsName := strings.Join(fields, ", ")
	colsValue := strings.Join(values, "', '")
	sql := fmt.Sprintf("INSERT INTO %s.%s.%s (%s) VALUES ('%s')", database, schema, table, colsName, colsValue)

	db := Conn()
	result, err = db.Exec(sql)
	if err != nil {
		return
	}
	rowsAffected, err = result.RowsAffected()
	if err != nil {
		return
	}

	data := make(map[string]interface{})
	data["rows_affected"] = rowsAffected
	jsonData, err = json.Marshal(data)
	return
}

// Delete execute delete sql into a table
func Delete(database, schema, table, where string, whereValues []interface{}) (jsonData []byte, err error) {
	var result sql.Result
	var rowsAffected int64

	sql := fmt.Sprintf("DELETE FROM %s.%s.%s", database, schema, table)
	if where != "" {
		sql = fmt.Sprint(
			sql,
			" WHERE ",
			where)
	}

	db := Conn()
	result, err = db.Exec(sql, whereValues...)
	if err != nil {
		return
	}
	rowsAffected, err = result.RowsAffected()
	if err != nil {
		return
	}

	data := make(map[string]interface{})
	data["rows_affected"] = rowsAffected
	jsonData, err = json.Marshal(data)
	return
}

// Update execute update sql into a table
func Update(database, schema, table, where string, whereValues []interface{}, body api.Request) (jsonData []byte, err error) {
	var result sql.Result
	var rowsAffected int64

	fields := []string{}
	values := make([]interface{}, 0)
	pid := len(whereValues) + 1 // placeholder id
	for key, value := range body.Data {
		fields = append(fields, fmt.Sprintf("%s=$%d", key, pid))
		values = append(values, value)
		pid++
	}
	setSyntax := strings.Join(fields, ", ")

	sql := fmt.Sprintf("UPDATE %s.%s.%s SET %s", database, schema, table, setSyntax)

	if where != "" {
		sql = fmt.Sprint(
			sql,
			" WHERE ",
			where)
		values = append(values, whereValues...)
	}

	db := Conn()
	//result, err = db.Exec(sql, values)
	stmt, err := db.Prepare(sql)
	if err != nil {
		return
	}

	valuesAux := make([]interface{}, 0, len(values))

	for i := 0; i < len(values); i++ {
		valuesAux = append(valuesAux, values[i])
	}

	result, err = stmt.Exec(valuesAux...)
	if err != nil {
		return
	}

	rowsAffected, err = result.RowsAffected()
	if err != nil {
		return
	}

	data := make(map[string]interface{})
	data["rows_affected"] = rowsAffected
	jsonData, err = json.Marshal(data)
	return
}
<commit_msg>remove commented code<commit_after>package postgres

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"unicode"

	"github.com/caarlos0/env"
	"github.com/jmoiron/sqlx"
	// Used pg drive on sqlx
	_ "github.com/lib/pq"

	"database/sql"

	"github.com/nuveo/prest/api"
	"github.com/nuveo/prest/config"
)

const (
	pageNumberKey   = "_page"
	pageSizeKey     = "_page_size"
	defaultPageSize = 10
)

// Conn connect on PostgreSQL
// Used sqlx
func Conn() (db *sqlx.DB) {
	cfg := config.Prest{}
	env.Parse(&cfg)

	db, err := sqlx.Connect("postgres", fmt.Sprintf("user=%s dbname=%s host=%s port=%v sslmode=disable", cfg.PGUser, cfg.PGDatabase, cfg.PGHost, cfg.PGPort))
	if err != nil {
		panic(fmt.Sprintf("Unable to connection to database: %v\n", err))
	}
	return
}

// chkInvaidIdentifier return true if identifier is invalid
func chkInvaidIdentifier(identifer string) bool {
	if len(identifer) > 63 ||
		unicode.IsDigit([]rune(identifer)[0]) {
		return true
	}

	for _, v := range identifer {
		if !unicode.IsLetter(v) &&
			!unicode.IsDigit(v) &&
			v != '_' &&
			v != '.' {
			return true
		}
	}
	return false
}

// WhereByRequest create interface for queries + where
func WhereByRequest(r *http.Request, initialPlaceholderID int) (whereSyntax string, values []interface{}, err error) {
	whereKey := []string{}
	whereValues := []string{}

	u, _ := url.Parse(r.URL.String())
	pid := initialPlaceholderID
	for key, val := range u.Query() {
		if !strings.HasPrefix(key, "_") {
			keyInfo := strings.Split(key, ":")
			if len(keyInfo) > 1 {
				switch keyInfo[1] {
				case "jsonb":
					jsonField := strings.Split(keyInfo[0], "->>")
					if chkInvaidIdentifier(jsonField[0]) ||
						chkInvaidIdentifier(jsonField[1]) {
						err = errors.New("Invalid identifier")
						return
					}
					whereKey = append(whereKey, fmt.Sprintf("%s->>'%s'=$%d", jsonField[0], jsonField[1], pid))
					whereValues = append(whereValues, val[0])
				default:
					if chkInvaidIdentifier(keyInfo[0]) {
						err = errors.New("Invalid identifier")
						return
					}
					whereKey = append(whereKey, fmt.Sprintf("%s=$%d", keyInfo[0], pid))
					whereValues = append(whereValues, val[0])
				}
				continue
			}
			if chkInvaidIdentifier(key) {
				err = errors.New("Invalid identifier")
				return
			}

			whereKey = append(whereKey, fmt.Sprintf("%s=$%d", key, pid))
			whereValues = append(whereValues, val[0])

			pid++
		}
	}

	for i := 0; i < len(whereKey); i++ {
		if whereSyntax == "" {
			whereSyntax += whereKey[i]
		} else {
			whereSyntax += " AND " + whereKey[i]
		}

		values = append(values, whereValues[i])
	}

	return
}

// Query process queries
func Query(SQL string, params ...interface{}) (jsonData []byte, err error) {
	db := Conn()

	prepare, err := db.Prepare(SQL)

	if err != nil {
		return
	}

	rows, err := prepare.Query(params...)
	if err != nil {
		return
	}
	defer rows.Close()

	columns, err := rows.Columns()
	if err != nil {
		return
	}

	count := len(columns)
	tableData := make([]map[string]interface{}, 0)
	values := make([]interface{}, count)
	valuePtrs := make([]interface{}, count)
	for rows.Next() {
		for i := 0; i < count; i++ {
			valuePtrs[i] = &values[i]
		}
		rows.Scan(valuePtrs...)
		entry := make(map[string]interface{})
		for i, col := range columns {
			var v interface{}
			val := values[i]
			b, ok := val.([]byte)
			if ok {
				v = string(b)
			} else {
				v = val
			}
			entry[col] = v
		}
		tableData = append(tableData, entry)
	}
	jsonData, err = json.Marshal(tableData)

	return
}

// PaginateIfPossible func
func PaginateIfPossible(r *http.Request) (paginatedQuery string, err error) {
	u, _ := url.Parse(r.URL.String())
	values := u.Query()
	if _, ok := values[pageNumberKey]; !ok {
		paginatedQuery = ""
		return
	}
	pageNumber, err := strconv.Atoi(values[pageNumberKey][0])
	if err != nil {
		return
	}
	pageSize := defaultPageSize
	if size, ok := values[pageSizeKey]; ok {
		pageSize, err = strconv.Atoi(size[0])
		if err != nil {
			return
		}
	}
	paginatedQuery = fmt.Sprintf("LIMIT %d OFFSET(%d - 1) * %d", pageSize, pageNumber, pageSize)
	return
}

// Insert execute insert sql into a table
func Insert(database, schema, table string, body api.Request) (jsonData []byte, err error) {
	var result sql.Result
	var rowsAffected int64

	fields := make([]string, 0)
	values := make([]string, 0)
	for key, value := range body.Data {
		fields = append(fields, key)
		values = append(values, value)
	}
	colsName := strings.Join(fields, ", ")
	colsValue := strings.Join(values, "', '")
	sql := fmt.Sprintf("INSERT INTO %s.%s.%s (%s) VALUES ('%s')", database, schema, table, colsName, colsValue)

	db := Conn()
	result, err = db.Exec(sql)
	if err != nil {
		return
	}
	rowsAffected, err = result.RowsAffected()
	if err != nil {
		return
	}

	data := make(map[string]interface{})
	data["rows_affected"] = rowsAffected
	jsonData, err = json.Marshal(data)
	return
}

// Delete execute delete sql into a table
func Delete(database, schema, table, where string, whereValues []interface{}) (jsonData []byte, err error) {
	var result sql.Result
	var rowsAffected int64

	sql := fmt.Sprintf("DELETE FROM %s.%s.%s", database, schema, table)
	if where != "" {
		sql = fmt.Sprint(
			sql,
			" WHERE ",
			where)
	}

	db := Conn()
	result, err = db.Exec(sql, whereValues...)
	if err != nil {
		return
	}
	rowsAffected, err = result.RowsAffected()
	if err != nil {
		return
	}

	data := make(map[string]interface{})
	data["rows_affected"] = rowsAffected
	jsonData, err = json.Marshal(data)
	return
}

// Update execute update sql into a table
func Update(database, schema, table, where string, whereValues []interface{}, body api.Request) (jsonData []byte, err error) {
	var result sql.Result
	var rowsAffected int64

	fields := []string{}
	values := make([]interface{}, 0)
	pid := len(whereValues) + 1 // placeholder id
	for key, value := range body.Data {
		fields = append(fields, fmt.Sprintf("%s=$%d", key, pid))
		values = append(values, value)
		pid++
	}
	setSyntax := strings.Join(fields, ", ")

	sql := fmt.Sprintf("UPDATE %s.%s.%s SET %s", database, schema, table, setSyntax)

	if where != "" {
		sql = fmt.Sprint(
			sql,
			" WHERE ",
			where)
		values = append(values, whereValues...)
	}

	db := Conn()
	//result, err = db.Exec(sql, values)
	stmt, err := db.Prepare(sql)
	if err != nil {
		return
	}

	valuesAux := make([]interface{}, 0, len(values))

	for i := 0; i < len(values); i++ {
		valuesAux = append(valuesAux, values[i])
	}

	result, err = stmt.Exec(valuesAux...)
	if err != nil {
		return
	}

	rowsAffected, err = result.RowsAffected()
	if err != nil {
		return
	}

	data := make(map[string]interface{})
	data["rows_affected"] = rowsAffected
	jsonData, err = json.Marshal(data)
	return
}
<|endoftext|>


<commit_before>// Copyright 2019 The GoPacket Authors. All rights reserved.
//
// Use of this source code is governed by a BSD-style license that can be found
// in the LICENSE file in the root of the source tree.

package layers

import (
	"bytes"
	"encoding/hex"
	"testing"

	"github.com/google/gopacket"
)

func TestDecodeFromBytes(t *testing.T) {
	b, err := hex.DecodeString("0600ff06")
	if err != nil {
		t.Fatalf("Failed to decode RMCP message")
	}

	rmcp := &RMCP{}
	if err := rmcp.DecodeFromBytes(b, gopacket.NilDecodeFeedback); err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	if !bytes.Equal(rmcp.BaseLayer.Payload, []byte{}) {
		t.Errorf("payload is %v, want %v", rmcp.BaseLayer.Payload, b)
	}
	if !bytes.Equal(rmcp.BaseLayer.Contents, b) {
		t.Errorf("contents is %v, want %v", rmcp.BaseLayer.Contents, b)
	}
	if rmcp.Version != RMCPVersion1 {
		t.Errorf("version is %v, want %v", rmcp.Version, RMCPVersion1)
	}
	if rmcp.Sequence != 0xFF {
		t.Errorf("sequence is %v, want %v", rmcp.Sequence, 0xFF)
	}
	if rmcp.Ack {
		t.Errorf("ack is true, want false")
	}
	if rmcp.Class != RMCPClassASF {
		t.Errorf("class is %v, want %v", rmcp.Class, RMCPClassASF)
	}
}

func serializeRMCP(rmcp *RMCP) ([]byte, error) {
	sb := gopacket.NewSerializeBuffer()
	err := rmcp.SerializeTo(sb, gopacket.SerializeOptions{})
	return sb.Bytes(), err
}

func TestSerializeTo(t *testing.T) {
	table := []struct {
		layer *RMCP
		want  []byte
	}{
		{
			&RMCP{
				Version:  RMCPVersion1,
				Sequence: 1,
				Ack:      false,
				Class:    RMCPClassASF,
			},
			[]byte{0x6, 0x0, 0x1, 0x6},
		},
		{
			&RMCP{
				Version:  RMCPVersion1,
				Sequence: 0xFF,
				Ack:      true,
				Class:    RMCPClassIPMI,
			},
			[]byte{0x6, 0x0, 0xFF, 0x87},
		},
	}
	for _, test := range table {
		b, err := serializeRMCP(test.layer)
		switch {
		case err != nil && test.want != nil:
			t.Errorf("serialize %v failed with %v, wanted %v", test.layer,
				err, test.want)
		case err == nil && !bytes.Equal(b, test.want):
			t.Errorf("serialize %v = %v, want %v", test.layer, b, test.want)
		}
	}
}
<commit_msg>Avoid naming collision with other layers<commit_after>// Copyright 2019 The GoPacket Authors. All rights reserved.
//
// Use of this source code is governed by a BSD-style license that can be found
// in the LICENSE file in the root of the source tree.

package layers

import (
	"bytes"
	"encoding/hex"
	"testing"

	"github.com/google/gopacket"
)

func RMCPTestDecodeFromBytes(t *testing.T) {
	b, err := hex.DecodeString("0600ff06")
	if err != nil {
		t.Fatalf("Failed to decode RMCP message")
	}

	rmcp := &RMCP{}
	if err := rmcp.DecodeFromBytes(b, gopacket.NilDecodeFeedback); err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	if !bytes.Equal(rmcp.BaseLayer.Payload, []byte{}) {
		t.Errorf("payload is %v, want %v", rmcp.BaseLayer.Payload, b)
	}
	if !bytes.Equal(rmcp.BaseLayer.Contents, b) {
		t.Errorf("contents is %v, want %v", rmcp.BaseLayer.Contents, b)
	}
	if rmcp.Version != RMCPVersion1 {
		t.Errorf("version is %v, want %v", rmcp.Version, RMCPVersion1)
	}
	if rmcp.Sequence != 0xFF {
		t.Errorf("sequence is %v, want %v", rmcp.Sequence, 0xFF)
	}
	if rmcp.Ack {
		t.Errorf("ack is true, want false")
	}
	if rmcp.Class != RMCPClassASF {
		t.Errorf("class is %v, want %v", rmcp.Class, RMCPClassASF)
	}
}

func serializeRMCP(rmcp *RMCP) ([]byte, error) {
	sb := gopacket.NewSerializeBuffer()
	err := rmcp.SerializeTo(sb, gopacket.SerializeOptions{})
	return sb.Bytes(), err
}

func RMCPTestSerializeTo(t *testing.T) {
	table := []struct {
		layer *RMCP
		want  []byte
	}{
		{
			&RMCP{
				Version:  RMCPVersion1,
				Sequence: 1,
				Ack:      false,
				Class:    RMCPClassASF,
			},
			[]byte{0x6, 0x0, 0x1, 0x6},
		},
		{
			&RMCP{
				Version:  RMCPVersion1,
				Sequence: 0xFF,
				Ack:      true,
				Class:    RMCPClassIPMI,
			},
			[]byte{0x6, 0x0, 0xFF, 0x87},
		},
	}
	for _, test := range table {
		b, err := serializeRMCP(test.layer)
		switch {
		case err != nil && test.want != nil:
			t.Errorf("serialize %v failed with %v, wanted %v", test.layer,
				err, test.want)
		case err == nil && !bytes.Equal(b, test.want):
			t.Errorf("serialize %v = %v, want %v", test.layer, b, test.want)
		}
	}
}
<|endoftext|>


<commit_before>package cluster

import (
	"crypto/tls"
	"fmt"
	"net"
	"strings"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/samalba/dockerclient"
)

const (
	// Force-refresh the state of the node this often.
	stateRefreshPeriod = 30 * time.Second

	// Timeout for requests sent out to the node.
	requestTimeout = 10 * time.Second
)

func NewNode(addr string) *Node {
	e := &Node{
		Addr:       addr,
		Labels:     make(map[string]string),
		ch:         make(chan bool),
		containers: make(map[string]*Container),
	}
	return e
}

type Node struct {
	sync.RWMutex

	ID     string
	IP     string
	Addr   string
	Name   string
	Cpus   int64
	Memory int64
	Labels map[string]string

	ch           chan bool
	containers   map[string]*Container
	client       dockerclient.Client
	eventHandler EventHandler
}

// Connect will initialize a connection to the Docker daemon running on the
// host, gather machine specs (memory, cpu, ...) and monitor state changes.
func (n *Node) Connect(config *tls.Config) error {
	c, err := dockerclient.NewDockerClientTimeout(n.Addr, config, time.Duration(requestTimeout))
	if err != nil {
		return err
	}

	addr, err := net.ResolveIPAddr("ip4", strings.Split(c.URL.Host, ":")[0])
	if err != nil {
		return err
	}
	n.IP = addr.IP.String()

	return n.connectClient(c)
}

func (n *Node) connectClient(client dockerclient.Client) error {
	n.client = client

	// Fetch the engine labels.
	if err := n.updateSpecs(); err != nil {
		n.client = nil
		return err
	}

	// Force a state update before returning.
	if err := n.refreshContainers(); err != nil {
		n.client = nil
		return err
	}

	// Start the update loop.
	go n.refreshLoop()

	// Start monitoring events from the Node.
	n.client.StartMonitorEvents(n.handler)

	return nil
}

// IsConnected returns true if the engine is connected to a remote docker API
func (e *Node) IsConnected() bool {
	return e.client != nil
}

// Gather node specs (CPU, memory, constraints, ...).
func (n *Node) updateSpecs() error {
	info, err := n.client.Info()
	if err != nil {
		return err
	}
	n.ID = info.ID
	n.Name = info.Name
	n.Cpus = info.NCPU
	n.Memory = info.MemTotal
	n.Labels = map[string]string{
		"storagedriver":   info.Driver,
		"executiondriver": info.ExecutionDriver,
		"kernelversion":   info.KernelVersion,
		"operatingsystem": info.OperatingSystem,
	}
	for _, label := range info.Labels {
		kv := strings.SplitN(label, "=", 2)
		n.Labels[kv[0]] = kv[1]
	}
	return nil
}

// Refresh the list and status of containers running on the node.
func (n *Node) refreshContainers() error {
	containers, err := n.client.ListContainers(true, false, "")
	if err != nil {
		return err
	}

	n.Lock()
	defer n.Unlock()

	merged := make(map[string]*Container)
	for _, c := range containers {
		merged, err = n.updateContainer(c, merged)
		if err != nil {
			log.Errorf("[%s/%s] Unable to update state of %s", n.ID, n.Name, c.Id)
		}
	}

	n.containers = merged

	log.Debugf("[%s/%s] Updated state", n.ID, n.Name)
	return nil
}

// Refresh the status of a container running on the node.
func (n *Node) refreshContainer(ID string) error {
	containers, err := n.client.ListContainers(true, false, fmt.Sprintf("{%q:[%q]}", "id", ID))
	if err != nil {
		return err
	}

	if len(containers) > 1 {
		// We expect one container, if we get more than one, trigger a full refresh.
		return n.refreshContainers()
	}

	n.Lock()
	defer n.Unlock()

	if len(containers) == 0 {
		// The container doesn't exist on the node, remove it.
		delete(n.containers, ID)
		return nil
	}

	_, err = n.updateContainer(containers[0], n.containers)
	return err
}

func (n *Node) updateContainer(c dockerclient.Container, containers map[string]*Container) (map[string]*Container, error) {
	if current, exists := n.containers[c.Id]; exists {
		// The container exists. Update its state.
		current.Container = c
		containers[current.Id] = current
	} else {
		// This is a brand new container.
		container := &Container{}
		container.Container = c
		container.node = n

		info, err := n.client.InspectContainer(c.Id)
		if err != nil {
			return containers, err
		}
		container.Info = *info

		// real CpuShares -> nb of CPUs
		container.Info.Config.CpuShares = container.Info.Config.CpuShares / 100.0 * n.Cpus

		containers[container.Id] = container
	}
	return containers, nil
}

func (n *Node) refreshContainersAsync() {
	n.ch <- true
}

func (n *Node) refreshLoop() {
	for {
		var err error
		select {
		case <-n.ch:
			err = n.refreshContainers()
		case <-time.After(stateRefreshPeriod):
			err = n.refreshContainers()
		}
		if err != nil {
			log.Errorf("[%s] Updated state failed: %v", n.ID, err)
		}
	}
}

// Return the sum of memory reserved by containers.
func (n *Node) ReservedMemory() int64 {
	var r int64 = 0
	n.RLock()
	for _, c := range n.containers {
		r += c.Info.Config.Memory
	}
	n.RUnlock()
	return r
}

// Return the memory availalble on this node.
func (n *Node) AvailableMemory() int64 {
	return n.Memory - n.ReservedMemory()
}

// Return the sum of CPUs reserved by containers.
func (n *Node) ReservedCpus() int64 {
	var r int64 = 0
	n.RLock()
	for _, c := range n.containers {
		r += c.Info.Config.CpuShares
	}
	n.RUnlock()
	return r
}

func (n *Node) AvailalbleCpus() int64 {
	return n.Cpus - n.ReservedCpus()
}

func (n *Node) Create(config *dockerclient.ContainerConfig, name string, pullImage bool) (*Container, error) {
	var (
		err    error
		id     string
		client = n.client
	)

	newConfig := *config

	// nb of CPUs -> real CpuShares
	newConfig.CpuShares = config.CpuShares * 100 / n.Cpus

	if id, err = client.CreateContainer(&newConfig, name); err != nil {
		// If the error is other than not found, abort immediately.
		if err != dockerclient.ErrNotFound {
			return nil, err
		}
		// Otherwise, try to pull the image...
		if err = n.Pull(config.Image); err != nil {
			return nil, err
		}
		// ...And try again.
		if id, err = client.CreateContainer(&newConfig, name); err != nil {
			return nil, err
		}
	}

	// Register the container immediately while waiting for a state refresh.
	// Force a state refresh to pick up the newly created container.
	n.refreshContainer(id)

	n.RLock()
	defer n.RUnlock()

	return n.containers[id], nil
}

func (n *Node) ListImages() ([]string, error) {
	images, err := n.client.ListImages()
	if err != nil {
		return nil, err
	}

	out := []string{}

	for _, i := range images {
		for _, t := range i.RepoTags {
			out = append(out, t)
		}
	}

	return out, nil
}

func (n *Node) Remove(container *Container, force bool) error {
	if err := n.client.RemoveContainer(container.Id, force); err != nil {
		return err
	}

	// Remove the container from the state. Eventually, the state refresh loop
	// will rewrite this.
	n.Lock()
	defer n.Unlock()
	delete(n.containers, container.Id)

	return nil
}

func (n *Node) Pull(image string) error {
	if err := n.client.PullImage(image); err != nil {
		return err
	}
	return nil
}

// Register an event handler.
func (n *Node) Events(h EventHandler) error {
	if n.eventHandler != nil {
		return fmt.Errorf("event handler already set")
	}
	n.eventHandler = h
	return nil
}

func (n *Node) Containers() []*Container {
	containers := []*Container{}
	n.RLock()
	for _, container := range n.containers {
		containers = append(containers, container)
	}
	n.RUnlock()
	return containers
}

func (n *Node) String() string {
	return fmt.Sprintf("node %s addr %s", n.ID, n.Addr)
}

func (n *Node) handler(ev *dockerclient.Event, args ...interface{}) {
	// Something changed - refresh our internal state.
	n.refreshContainer(ev.Id)

	// If there is no event handler registered, abort right now.
	if n.eventHandler == nil {
		return
	}

	event := &Event{
		NodeName: n.Name,
		NodeID:   n.ID,
		NodeAddr: n.Addr,
		NodeIP:   n.IP,
	}
	event.Event = *ev

	n.eventHandler.Handle(event)
}

// Used only on tests
func (n *Node) AddContainer(container *Container) {
	n.containers[container.Id] = container
}

func (n *Node) CleanupContainers() {
	n.containers = make(map[string]*Container)
}
<commit_msg>Node: Basic health checking.<commit_after>package cluster

import (
	"crypto/tls"
	"fmt"
	"net"
	"strings"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/samalba/dockerclient"
)

const (
	// Force-refresh the state of the node this often.
	stateRefreshPeriod = 30 * time.Second

	// Timeout for requests sent out to the node.
	requestTimeout = 10 * time.Second
)

func NewNode(addr string) *Node {
	e := &Node{
		Addr:       addr,
		Labels:     make(map[string]string),
		ch:         make(chan bool),
		containers: make(map[string]*Container),
		healthy:    true,
	}
	return e
}

type Node struct {
	sync.RWMutex

	ID     string
	IP     string
	Addr   string
	Name   string
	Cpus   int64
	Memory int64
	Labels map[string]string

	ch           chan bool
	containers   map[string]*Container
	client       dockerclient.Client
	eventHandler EventHandler
	healthy      bool
}

// Connect will initialize a connection to the Docker daemon running on the
// host, gather machine specs (memory, cpu, ...) and monitor state changes.
func (n *Node) Connect(config *tls.Config) error {
	c, err := dockerclient.NewDockerClientTimeout(n.Addr, config, time.Duration(requestTimeout))
	if err != nil {
		return err
	}

	addr, err := net.ResolveIPAddr("ip4", strings.Split(c.URL.Host, ":")[0])
	if err != nil {
		return err
	}
	n.IP = addr.IP.String()

	return n.connectClient(c)
}

func (n *Node) connectClient(client dockerclient.Client) error {
	n.client = client

	// Fetch the engine labels.
	if err := n.updateSpecs(); err != nil {
		n.client = nil
		return err
	}

	// Force a state update before returning.
	if err := n.refreshContainers(); err != nil {
		n.client = nil
		return err
	}

	// Start the update loop.
	go n.refreshLoop()

	// Start monitoring events from the Node.
	n.client.StartMonitorEvents(n.handler)

	return nil
}

// IsConnected returns true if the engine is connected to a remote docker API
func (n *Node) IsConnected() bool {
	return n.client != nil
}

func (n *Node) IsHealthy() bool {
	return n.healthy
}

// Gather node specs (CPU, memory, constraints, ...).
func (n *Node) updateSpecs() error {
	info, err := n.client.Info()
	if err != nil {
		return err
	}
	n.ID = info.ID
	n.Name = info.Name
	n.Cpus = info.NCPU
	n.Memory = info.MemTotal
	n.Labels = map[string]string{
		"storagedriver":   info.Driver,
		"executiondriver": info.ExecutionDriver,
		"kernelversion":   info.KernelVersion,
		"operatingsystem": info.OperatingSystem,
	}
	for _, label := range info.Labels {
		kv := strings.SplitN(label, "=", 2)
		n.Labels[kv[0]] = kv[1]
	}
	return nil
}

// Refresh the list and status of containers running on the node.
func (n *Node) refreshContainers() error {
	containers, err := n.client.ListContainers(true, false, "")
	if err != nil {
		return err
	}

	n.Lock()
	defer n.Unlock()

	merged := make(map[string]*Container)
	for _, c := range containers {
		merged, err = n.updateContainer(c, merged)
		if err != nil {
			log.Errorf("[%s/%s] Unable to update state of %s", n.ID, n.Name, c.Id)
		}
	}

	n.containers = merged

	log.Debugf("[%s/%s] Updated state", n.ID, n.Name)
	return nil
}

// Refresh the status of a container running on the node.
func (n *Node) refreshContainer(ID string) error {
	containers, err := n.client.ListContainers(true, false, fmt.Sprintf("{%q:[%q]}", "id", ID))
	if err != nil {
		return err
	}

	if len(containers) > 1 {
		// We expect one container, if we get more than one, trigger a full refresh.
		return n.refreshContainers()
	}

	n.Lock()
	defer n.Unlock()

	if len(containers) == 0 {
		// The container doesn't exist on the node, remove it.
		delete(n.containers, ID)
		return nil
	}

	_, err = n.updateContainer(containers[0], n.containers)
	return err
}

func (n *Node) updateContainer(c dockerclient.Container, containers map[string]*Container) (map[string]*Container, error) {
	if current, exists := n.containers[c.Id]; exists {
		// The container exists. Update its state.
		current.Container = c
		containers[current.Id] = current
	} else {
		// This is a brand new container.
		container := &Container{}
		container.Container = c
		container.node = n

		info, err := n.client.InspectContainer(c.Id)
		if err != nil {
			return containers, err
		}
		container.Info = *info

		// real CpuShares -> nb of CPUs
		container.Info.Config.CpuShares = container.Info.Config.CpuShares / 100.0 * n.Cpus

		containers[container.Id] = container
	}
	return containers, nil
}

func (n *Node) refreshContainersAsync() {
	n.ch <- true
}

func (n *Node) refreshLoop() {
	for {
		var err error
		select {
		case <-n.ch:
			err = n.refreshContainers()
		case <-time.After(stateRefreshPeriod):
			err = n.refreshContainers()
		}

		if err != nil {
			n.healthy = false
			log.Errorf("[%s/%s] Flagging node as dead. Updated state failed: %v", n.ID, n.Name, err)
		} else {
			if !n.healthy {
				log.Infof("[%s/%s] Node came back to life. Hooray!", n.ID, n.Name)
			}
			n.healthy = true
		}
	}
}

// Return the sum of memory reserved by containers.
func (n *Node) ReservedMemory() int64 {
	var r int64 = 0
	n.RLock()
	for _, c := range n.containers {
		r += c.Info.Config.Memory
	}
	n.RUnlock()
	return r
}

// Return the memory availalble on this node.
func (n *Node) AvailableMemory() int64 {
	return n.Memory - n.ReservedMemory()
}

// Return the sum of CPUs reserved by containers.
func (n *Node) ReservedCpus() int64 {
	var r int64 = 0
	n.RLock()
	for _, c := range n.containers {
		r += c.Info.Config.CpuShares
	}
	n.RUnlock()
	return r
}

func (n *Node) AvailalbleCpus() int64 {
	return n.Cpus - n.ReservedCpus()
}

func (n *Node) Create(config *dockerclient.ContainerConfig, name string, pullImage bool) (*Container, error) {
	var (
		err    error
		id     string
		client = n.client
	)

	newConfig := *config

	// nb of CPUs -> real CpuShares
	newConfig.CpuShares = config.CpuShares * 100 / n.Cpus

	if id, err = client.CreateContainer(&newConfig, name); err != nil {
		// If the error is other than not found, abort immediately.
		if err != dockerclient.ErrNotFound {
			return nil, err
		}
		// Otherwise, try to pull the image...
		if err = n.Pull(config.Image); err != nil {
			return nil, err
		}
		// ...And try again.
		if id, err = client.CreateContainer(&newConfig, name); err != nil {
			return nil, err
		}
	}

	// Register the container immediately while waiting for a state refresh.
	// Force a state refresh to pick up the newly created container.
	n.refreshContainer(id)

	n.RLock()
	defer n.RUnlock()

	return n.containers[id], nil
}

func (n *Node) ListImages() ([]string, error) {
	images, err := n.client.ListImages()
	if err != nil {
		return nil, err
	}

	out := []string{}

	for _, i := range images {
		for _, t := range i.RepoTags {
			out = append(out, t)
		}
	}

	return out, nil
}

func (n *Node) Remove(container *Container, force bool) error {
	if err := n.client.RemoveContainer(container.Id, force); err != nil {
		return err
	}

	// Remove the container from the state. Eventually, the state refresh loop
	// will rewrite this.
	n.Lock()
	defer n.Unlock()
	delete(n.containers, container.Id)

	return nil
}

func (n *Node) Pull(image string) error {
	if err := n.client.PullImage(image); err != nil {
		return err
	}
	return nil
}

// Register an event handler.
func (n *Node) Events(h EventHandler) error {
	if n.eventHandler != nil {
		return fmt.Errorf("event handler already set")
	}
	n.eventHandler = h
	return nil
}

func (n *Node) Containers() []*Container {
	containers := []*Container{}
	n.RLock()
	for _, container := range n.containers {
		containers = append(containers, container)
	}
	n.RUnlock()
	return containers
}

func (n *Node) String() string {
	return fmt.Sprintf("node %s addr %s", n.ID, n.Addr)
}

func (n *Node) handler(ev *dockerclient.Event, args ...interface{}) {
	// Something changed - refresh our internal state.
	n.refreshContainer(ev.Id)

	// If there is no event handler registered, abort right now.
	if n.eventHandler == nil {
		return
	}

	event := &Event{
		NodeName: n.Name,
		NodeID:   n.ID,
		NodeAddr: n.Addr,
		NodeIP:   n.IP,
	}
	event.Event = *ev

	n.eventHandler.Handle(event)
}

// Used only on tests
func (n *Node) AddContainer(container *Container) {
	n.containers[container.Id] = container
}

func (n *Node) CleanupContainers() {
	n.containers = make(map[string]*Container)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"

	"github.com/appcelerator/amp/api/client"
	"github.com/appcelerator/amp/cmd/amp/cli"
	"github.com/spf13/cobra"
)

var (
	// Version is set with a linker flag (see Makefile)
	Version string

	// Build is set with a linker flag (see Makefile)
	Build string

	// AMP manages the connection and state for the client
	AMP *client.AMP

	// Config is used by command implementations to access the computed client configuration.
	Config     = &client.Configuration{}
	configFile string
	verbose    bool
	serverAddr string

	// RootCmd is the base command for the CLI.
	RootCmd = &cobra.Command{
		Use:   `amp [OPTIONS] COMMAND [arg...]`,
		Short: "Appcelerator Microservice Platform.",
	}
)

// All main does is process commands and flags and invoke the app
func main() {
	cobra.OnInitialize(func() {
		InitConfig(configFile, Config, verbose, serverAddr)
		if addr := RootCmd.Flag("server").Value.String(); addr != "" {
			Config.ServerAddress = addr
		}
		if Config.ServerAddress == "" {
			Config.ServerAddress = client.DefaultServerAddress
		}
		AMP = client.NewAMP(Config)
		if AMP.Verbose() == false {
			RootCmd.SilenceErrors = true
			RootCmd.SilenceUsage = true
		}
		cli.AtExit(func() {
			if AMP != nil {
				AMP.Disconnect()
			}
		})
	})

	// infoCmd represents the amp information
	infoCmd := &cobra.Command{
		Use:   "info",
		Short: "Display amp version and server information",
		Long:  `Display amp version and server information.`,
		Run: func(cmd *cobra.Command, args []string) {
			fmt.Printf("amp (cli version: %s, build: %s)\n", Version, Build)
			fmt.Printf("Server: %s\n", Config.ServerAddress)
		},
	}
	RootCmd.SetUsageTemplate(usageTemplate)
	RootCmd.SetHelpTemplate(helpTemplate)

	RootCmd.PersistentFlags().StringVar(&configFile, "config", "", "Config file (default is $HOME/.amp.yaml)")
	RootCmd.PersistentFlags().BoolVarP(&verbose, "verbose", "v", false, `Verbose output`)
	RootCmd.PersistentFlags().StringVar(&serverAddr, "server", "", "Server address")
	RootCmd.AddCommand(infoCmd)
	cmd, _, err := RootCmd.Find(os.Args[1:])
	if err != nil {
		fmt.Println(err)
		cli.Exit(1)
	}
	if err := cmd.Execute(); err != nil {
		if AMP.Verbose() == false {
			fmt.Printf("Error during: amp %s, reason: %v\n", cmd.Name(), err)
		}
		cli.Exit(1)
	}
	cli.Exit(0)
}

var usageTemplate = `Usage:	{{if not .HasSubCommands}}{{.UseLine}}{{end}}{{if .HasSubCommands}}{{ .CommandPath}} COMMAND{{end}}

{{ .Short | trim }}{{if gt .Aliases 0}}

Aliases:
  {{.NameAndAliases}}{{end}}{{if .HasExample}}

Examples:
{{ .Example }}{{end}}{{if .HasFlags}}

Options:
{{.Flags.FlagUsages | trimRightSpace}}{{end}}{{ if .HasAvailableSubCommands}}

Commands:{{range .Commands}}{{if .IsAvailableCommand}}
  {{rpad .Name .NamePadding }} {{.Short}}{{end}}{{end}}{{end}}{{ if .HasSubCommands }}

Run '{{.CommandPath}} COMMAND --help' for more information on a command.{{end}}
`

var helpTemplate = `
{{if or .Runnable .HasSubCommands}}{{.UsageString}}{{end}}`
<commit_msg>Fix for issue#458 - amp version (#462)<commit_after>package main

import (
	"fmt"
	"os"

	"github.com/appcelerator/amp/api/client"
	"github.com/appcelerator/amp/cmd/amp/cli"
	"github.com/spf13/cobra"
)

var (
	// Version is set with a linker flag (see Makefile)
	Version string

	// Build is set with a linker flag (see Makefile)
	Build string

	// AMP manages the connection and state for the client
	AMP *client.AMP

	// Config is used by command implementations to access the computed client configuration.
	Config     = &client.Configuration{}
	configFile string
	verbose    bool
	serverAddr string
	listVersion = true

	// RootCmd is the base command for the CLI.
	RootCmd = &cobra.Command{
		Use:   `amp [OPTIONS] COMMAND [arg...]`,
		Short: "Appcelerator Microservice Platform.",
		Run: func(cmd *cobra.Command, args []string) {
			if listVersion {
				fmt.Printf("amp (cli version: %s, build: %s)\n", Version, Build)
			}
			fmt.Println(cmd.UsageString())
		},
	}
)

// All main does is process commands and flags and invoke the app
func main() {
	cobra.OnInitialize(func() {
		InitConfig(configFile, Config, verbose, serverAddr)
		if addr := RootCmd.Flag("server").Value.String(); addr != "" {
			Config.ServerAddress = addr
		}
		if Config.ServerAddress == "" {
			Config.ServerAddress = client.DefaultServerAddress
		}
		AMP = client.NewAMP(Config)
		if AMP.Verbose() == false {
			RootCmd.SilenceErrors = true
			RootCmd.SilenceUsage = true
		}
		cli.AtExit(func() {
			if AMP != nil {
				AMP.Disconnect()
			}
		})
	})

	// versionCmd represents the amp version
	versionCmd := &cobra.Command{
		Use:   "version",
		Short: "Display the version number of amp",
		Long:  `Display the version number of amp`,
		Run: func(cmd *cobra.Command, args []string) {
			fmt.Printf("amp (cli version: %s, build: %s)\n", Version, Build)
		},
	}
	RootCmd.AddCommand(versionCmd)

	// infoCmd represents the amp information
	infoCmd := &cobra.Command{
		Use:   "info",
		Short: "Display amp version and server information",
		Long:  `Display amp version and server information.`,
		Run: func(cmd *cobra.Command, args []string) {
			fmt.Printf("amp (cli version: %s, build: %s)\n", Version, Build)
			fmt.Printf("Server: %s\n", Config.ServerAddress)
		},
	}
	RootCmd.SetUsageTemplate(usageTemplate)
	RootCmd.SetHelpTemplate(helpTemplate)

	RootCmd.PersistentFlags().StringVar(&configFile, "config", "", "Config file (default is $HOME/.amp.yaml)")
	RootCmd.PersistentFlags().BoolVarP(&verbose, "verbose", "v", false, `Verbose output`)
	RootCmd.PersistentFlags().StringVar(&serverAddr, "server", "", "Server address")
	RootCmd.Flags().BoolVarP(&listVersion, "version", "V", false, "Version number")
	RootCmd.AddCommand(infoCmd)
	cmd, _, err := RootCmd.Find(os.Args[1:])
	if err != nil {
		fmt.Println(err)
		cli.Exit(1)
	}
	if err := cmd.Execute(); err != nil {
		if AMP.Verbose() == false {
			fmt.Printf("Error during: amp %s, reason: %v\n", cmd.Name(), err)
		}
		cli.Exit(1)
	}
	cli.Exit(0)
}

var usageTemplate = `Usage:	{{if not .HasSubCommands}}{{.UseLine}}{{end}}{{if .HasSubCommands}}{{ .CommandPath}} COMMAND{{end}}

{{ .Short | trim }}{{if gt .Aliases 0}}

Aliases:
  {{.NameAndAliases}}{{end}}{{if .HasExample}}

Examples:
{{ .Example }}{{end}}{{if .HasFlags}}

Options:
{{.Flags.FlagUsages | trimRightSpace}}{{end}}{{ if .HasAvailableSubCommands}}

Commands:{{range .Commands}}{{if .IsAvailableCommand}}
  {{rpad .Name .NamePadding }} {{.Short}}{{end}}{{end}}{{end}}{{ if .HasSubCommands }}

Run '{{.CommandPath}} COMMAND --help' for more information on a command.{{end}}
`

var helpTemplate = `
{{if or .Runnable .HasSubCommands}}{{.UsageString}}{{end}}`
<|endoftext|>


<commit_before>package cmd

import (
	"github.com/codegangsta/cli"
)

const version = "0.0.1"
const author = ""
const support = "liujianping@h2object.io"

func App() *cli.App {
	app := cli.NewApp()

	//! app settings
	app.Name = "scaffold"
	app.Usage = "scaffold database adminstrator portal tool"
	app.Version = version
	app.Author = author
	app.Email = support

	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "template, t",
			Value: "",
			Usage: "template directory",
		},
		cli.BoolFlag{
			Name:  "force, f",
			Usage: "force execute",
		},
	}
	//! app commands
	app.Commands = []cli.Command{
		{
			Name:  "revel",
			Usage: "revel project scaffolding",
			Subcommands: []cli.Command{
				{
					Name:  "init",
					Usage: "revel project scaffolding initialization",
					Action: func(ctx *cli.Context) {
						revelInit(ctx)
					},
				},
				{
					Name:  "index",
					Usage: "revel project scaffolding indexing",
					Flags: []cli.Flag{
						cli.StringFlag{
							Name:  "theme, t",
							Value: "bootstrap",
							Usage: "revel project view theme",
						},
					},
					Action: func(ctx *cli.Context) {
						revelIndex(ctx)
					},
				},
				{
					Name:  "module",
					Usage: "revel project scaffolding module",
					Flags: []cli.Flag{
						cli.StringFlag{
							Name:  "theme, t",
							Value: "bootstrap",
							Usage: "revel project module view theme",
						},
					},
					Action: func(ctx *cli.Context) {
						revelModule(ctx)
					},
				},
				{
					Name:  "controller",
					Usage: "revel project scaffolding controller",
					Flags: []cli.Flag{
						cli.StringFlag{
							Name:  "module, m",
							Value: "*",
							Usage: "revel project module name",
						},
					},
					Action: func(ctx *cli.Context) {
						revelController(ctx)
					},
				},
				{
					Name:  "model",
					Usage: "revel project scaffolding model",
					Action: func(ctx *cli.Context) {
						revelModel(ctx)
					},
				},
				{
					Name:  "view",
					Usage: "revel project scaffolding view",
					Flags: []cli.Flag{
						cli.StringFlag{
							Name:  "theme, t",
							Value: "bootstrap",
							Usage: "revel project module view theme",
						},
					},
					Action: func(ctx *cli.Context) {
						revelView(ctx)
					},
				},
			},
		},
	}

	return app
}
<commit_msg>change contact<commit_after>package cmd

import (
	"github.com/codegangsta/cli"
)

const version = "0.0.1"
const author = ""
const support = "liujianping.itech@qq.com"

func App() *cli.App {
	app := cli.NewApp()

	//! app settings
	app.Name = "scaffold"
	app.Usage = "scaffold database adminstrator portal tool"
	app.Version = version
	app.Author = author
	app.Email = support

	app.Flags = []cli.Flag{
		cli.StringFlag{
			Name:  "template, t",
			Value: "",
			Usage: "template directory",
		},
		cli.BoolFlag{
			Name:  "force, f",
			Usage: "force execute",
		},
	}
	//! app commands
	app.Commands = []cli.Command{
		{
			Name:  "revel",
			Usage: "revel project scaffolding",
			Subcommands: []cli.Command{
				{
					Name:  "init",
					Usage: "revel project scaffolding initialization",
					Action: func(ctx *cli.Context) {
						revelInit(ctx)
					},
				},
				{
					Name:  "index",
					Usage: "revel project scaffolding indexing",
					Flags: []cli.Flag{
						cli.StringFlag{
							Name:  "theme, t",
							Value: "bootstrap",
							Usage: "revel project view theme",
						},
					},
					Action: func(ctx *cli.Context) {
						revelIndex(ctx)
					},
				},
				{
					Name:  "module",
					Usage: "revel project scaffolding module",
					Flags: []cli.Flag{
						cli.StringFlag{
							Name:  "theme, t",
							Value: "bootstrap",
							Usage: "revel project module view theme",
						},
					},
					Action: func(ctx *cli.Context) {
						revelModule(ctx)
					},
				},
				{
					Name:  "controller",
					Usage: "revel project scaffolding controller",
					Flags: []cli.Flag{
						cli.StringFlag{
							Name:  "module, m",
							Value: "*",
							Usage: "revel project module name",
						},
					},
					Action: func(ctx *cli.Context) {
						revelController(ctx)
					},
				},
				{
					Name:  "model",
					Usage: "revel project scaffolding model",
					Action: func(ctx *cli.Context) {
						revelModel(ctx)
					},
				},
				{
					Name:  "view",
					Usage: "revel project scaffolding view",
					Flags: []cli.Flag{
						cli.StringFlag{
							Name:  "theme, t",
							Value: "bootstrap",
							Usage: "revel project module view theme",
						},
					},
					Action: func(ctx *cli.Context) {
						revelView(ctx)
					},
				},
			},
		},
	}

	return app
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"log"
	"os"
	"runtime"
	"strings"

	flag "github.com/bgentry/pflag"
	"github.com/docker/docker/pkg/term"
	"github.com/mgutz/ansi"
	"github.com/remind101/empire/cmd/emp/hkclient"
	"github.com/remind101/empire/pkg/heroku"
)

var (
	apiURL = "http://localhost:8080"
	stdin  = bufio.NewReader(os.Stdin)

	inFd, outFd                 uintptr
	isTerminalIn, isTerminalOut bool
)

type Command struct {
	// args does not include the command name
	Run             func(cmd *Command, args []string)
	Flag            flag.FlagSet
	NeedsApp        bool
	OptionalApp     bool
	NoClient        bool
	OptionalMessage bool

	Usage    string // first word is the command name
	Category string // i.e. "App", "Account", etc.
	Short    string // `emp help` output
	Long     string // `emp help cmd` output
	Alias    string // Optional alias for the command.
	Hidden   bool   // Set to true to hide a command from usage information.
}

func (c *Command) PrintUsage() {
	if c.Runnable() {
		fmt.Fprintf(os.Stderr, "Usage: emp %s\n", c.FullUsage())
	}
	fmt.Fprintf(os.Stderr, "Use 'emp help %s' for more information.\n", c.Name())
}

func (c *Command) PrintLongUsage() {
	if c.Runnable() {
		fmt.Printf("Usage: emp %s\n\n", c.FullUsage())
	}
	fmt.Println(strings.Trim(c.Long, "\n"))
}

func (c *Command) FullUsage() string {
	parts := make([]string, 0, 2)
	if c.NeedsApp || c.OptionalApp {
		parts = append(parts, "[-a <app or remote>]")
	}

	if c.OptionalMessage {
		parts = append(parts, "[-m <message>]")
	}

	if len(parts) > 0 {
		// NB: no space between last two placeholders because of
		// strings.TrimPrefix behavior
		return fmt.Sprintf("%s %s%s", c.Name(), strings.Join(parts, " "), strings.TrimPrefix(c.Usage, c.Name()))
	}
	return c.Usage
}

func (c *Command) Name() string {
	name := c.Usage
	i := strings.Index(name, " ")
	if i >= 0 {
		name = name[:i]
	}
	return name
}

func (c *Command) Runnable() bool {
	return c.Run != nil
}

func (c *Command) Visible() bool {
	return c.Runnable() && !c.Hidden
}

const extra = " (extra)"

func (c *Command) List() bool {
	return c.Short != "" && !strings.HasSuffix(c.Short, extra)
}

func (c *Command) ListAsExtra() bool {
	return c.Short != "" && strings.HasSuffix(c.Short, extra)
}

func (c *Command) ShortExtra() string {
	return c.Short[:len(c.Short)-len(extra)]
}

// Running `emp help` will list commands in this order.
var commands = []*Command{
	cmdCreate,
	cmdApps,
	cmdDynos,
	cmdReleases,
	cmdReleaseInfo,
	cmdRollback,
	cmdScale,
	cmdRestart,
	cmdEnvLoad,
	cmdSet,
	cmdUnset,
	cmdEnv,
	cmdRun,
	cmdLog,
	cmdInfo,
	cmdRename,
	cmdDestroy,
	cmdDomains,
	cmdDomainAdd,
	cmdDomainRemove,
	cmdCertAttach,
	cmdDeploy,
	cmdVersion,
	cmdHelp,

	helpCommands,
	helpEnviron,
	helpMore,
	helpAbout,

	// listed by emp help more
	cmdAPI,
	cmdAuthorize,
	cmdCreds,
	cmdGet,
	cmdLogin,
	cmdLogout,
	cmdSSL,
	cmdSSLCertAdd,
	cmdSSLCertRollback,
	cmdSSLDestroy,
	cmdURL,
	cmdWhichApp,
}

var (
	flagApp     string
	flagMessage string
	client      *heroku.Client
	hkAgent     = "hk/" + Version + " (" + runtime.GOOS + "; " + runtime.GOARCH + ")"
	userAgent   = hkAgent + " " + heroku.DefaultUserAgent
)

func initClients() {
	loadNetrc()
	suite, err := hkclient.New(nrc, hkAgent)
	if err != nil {
		printFatal(err.Error())
	}

	client = suite.Client
	apiURL = suite.ApiURL
}

func main() {
	log.SetFlags(0)

	// make sure command is specified, disallow global args
	args := os.Args[1:]
	if len(args) < 1 || strings.IndexRune(args[0], '-') == 0 {
		printUsageTo(os.Stderr)
		os.Exit(2)
	}

	inFd, isTerminalIn = term.GetFdInfo(os.Stdin)
	outFd, isTerminalOut = term.GetFdInfo(os.Stdout)

	if !isTerminalOut {
		ansi.DisableColors(true)
	}

	for _, cmd := range commands {
		if matchesCommand(cmd, args[0]) && cmd.Run != nil {
			defer recoverPanic()

			if !cmd.NoClient {
				initClients()
			}

			cmd.Flag.SetDisableDuplicates(true) // disallow duplicate flag options
			if !gitConfigBool("hk.strict-flag-ordering") {
				cmd.Flag.SetInterspersed(true) // allow flags & non-flag args to mix
			}
			cmd.Flag.Usage = func() {
				cmd.PrintUsage()
			}
			if cmd.NeedsApp || cmd.OptionalApp {
				cmd.Flag.StringVarP(&flagApp, "app", "a", "", "app name")
			}
			if cmd.OptionalMessage {
				cmd.Flag.StringVarP(&flagMessage, "message", "m", "", "message")
			}
			if err := cmd.Flag.Parse(args[1:]); err == flag.ErrHelp {
				cmdHelp.Run(cmdHelp, args[:1])
				return
			} else if err != nil {
				printError(err.Error())
				os.Exit(2)
			}
			if flagApp != "" {
				if gitRemoteApp, err := appFromGitRemote(flagApp); err == nil {
					flagApp = gitRemoteApp
				}
			}
			if cmd.NeedsApp || cmd.OptionalApp {
				a, err := app()
				switch {
				case err == errMultipleHerokuRemotes && cmd.NeedsApp, err == nil && a == "" && cmd.NeedsApp:
					msg := "no app specified"
					if err != nil {
						msg = err.Error()
					}
					printError(msg)
					cmd.PrintUsage()
					os.Exit(2)
				case err != nil && cmd.NeedsApp:
					printFatal(err.Error())
				}
			}
			cmd.Run(cmd, cmd.Flag.Args())
			return
		}
	}

	fmt.Fprintf(os.Stderr, "Unknown command: %s\n", args[0])
	if g := suggest(args[0]); len(g) > 0 {
		fmt.Fprintf(os.Stderr, "Possible alternatives: %v\n", strings.Join(g, " "))
	}
	fmt.Fprintf(os.Stderr, "Run 'emp help' for usage.\n")
	os.Exit(2)
}

func recoverPanic() {
	if Version != "dev" {
		if rec := recover(); rec != nil {
			printFatal("emp encountered and reported an internal client error")
		}
	}
}

func app() (string, error) {
	if flagApp != "" {
		return flagApp, nil
	}

	if app := os.Getenv("EMPAPP"); app != "" {
		return app, nil
	}

	return appFromGitRemote(remoteFromGitConfig())
}

func mustApp() string {
	name, err := app()
	if err != nil {
		printFatal(err.Error())
	}
	return name
}

func getMessage() string {
	return flagMessage
}

// matchesCommand checks if the Command matches the command that we want to run.
func matchesCommand(cmd *Command, want string) bool {
	if cmd.Alias != "" && cmd.Alias == want {
		return true
	}

	return cmd.Name() == want
}
<commit_msg>Don't need to use make<commit_after>package main

import (
	"bufio"
	"fmt"
	"log"
	"os"
	"runtime"
	"strings"

	flag "github.com/bgentry/pflag"
	"github.com/docker/docker/pkg/term"
	"github.com/mgutz/ansi"
	"github.com/remind101/empire/cmd/emp/hkclient"
	"github.com/remind101/empire/pkg/heroku"
)

var (
	apiURL = "http://localhost:8080"
	stdin  = bufio.NewReader(os.Stdin)

	inFd, outFd                 uintptr
	isTerminalIn, isTerminalOut bool
)

type Command struct {
	// args does not include the command name
	Run             func(cmd *Command, args []string)
	Flag            flag.FlagSet
	NeedsApp        bool
	OptionalApp     bool
	NoClient        bool
	OptionalMessage bool

	Usage    string // first word is the command name
	Category string // i.e. "App", "Account", etc.
	Short    string // `emp help` output
	Long     string // `emp help cmd` output
	Alias    string // Optional alias for the command.
	Hidden   bool   // Set to true to hide a command from usage information.
}

func (c *Command) PrintUsage() {
	if c.Runnable() {
		fmt.Fprintf(os.Stderr, "Usage: emp %s\n", c.FullUsage())
	}
	fmt.Fprintf(os.Stderr, "Use 'emp help %s' for more information.\n", c.Name())
}

func (c *Command) PrintLongUsage() {
	if c.Runnable() {
		fmt.Printf("Usage: emp %s\n\n", c.FullUsage())
	}
	fmt.Println(strings.Trim(c.Long, "\n"))
}

func (c *Command) FullUsage() string {
	var parts []string
	if c.NeedsApp || c.OptionalApp {
		parts = append(parts, "[-a <app or remote>]")
	}

	if c.OptionalMessage {
		parts = append(parts, "[-m <message>]")
	}

	if len(parts) > 0 {
		// NB: no space between last two placeholders because of
		// strings.TrimPrefix behavior
		return fmt.Sprintf("%s %s%s", c.Name(), strings.Join(parts, " "), strings.TrimPrefix(c.Usage, c.Name()))
	}
	return c.Usage
}

func (c *Command) Name() string {
	name := c.Usage
	i := strings.Index(name, " ")
	if i >= 0 {
		name = name[:i]
	}
	return name
}

func (c *Command) Runnable() bool {
	return c.Run != nil
}

func (c *Command) Visible() bool {
	return c.Runnable() && !c.Hidden
}

const extra = " (extra)"

func (c *Command) List() bool {
	return c.Short != "" && !strings.HasSuffix(c.Short, extra)
}

func (c *Command) ListAsExtra() bool {
	return c.Short != "" && strings.HasSuffix(c.Short, extra)
}

func (c *Command) ShortExtra() string {
	return c.Short[:len(c.Short)-len(extra)]
}

// Running `emp help` will list commands in this order.
var commands = []*Command{
	cmdCreate,
	cmdApps,
	cmdDynos,
	cmdReleases,
	cmdReleaseInfo,
	cmdRollback,
	cmdScale,
	cmdRestart,
	cmdEnvLoad,
	cmdSet,
	cmdUnset,
	cmdEnv,
	cmdRun,
	cmdLog,
	cmdInfo,
	cmdRename,
	cmdDestroy,
	cmdDomains,
	cmdDomainAdd,
	cmdDomainRemove,
	cmdCertAttach,
	cmdDeploy,
	cmdVersion,
	cmdHelp,

	helpCommands,
	helpEnviron,
	helpMore,
	helpAbout,

	// listed by emp help more
	cmdAPI,
	cmdAuthorize,
	cmdCreds,
	cmdGet,
	cmdLogin,
	cmdLogout,
	cmdSSL,
	cmdSSLCertAdd,
	cmdSSLCertRollback,
	cmdSSLDestroy,
	cmdURL,
	cmdWhichApp,
}

var (
	flagApp     string
	flagMessage string
	client      *heroku.Client
	hkAgent     = "hk/" + Version + " (" + runtime.GOOS + "; " + runtime.GOARCH + ")"
	userAgent   = hkAgent + " " + heroku.DefaultUserAgent
)

func initClients() {
	loadNetrc()
	suite, err := hkclient.New(nrc, hkAgent)
	if err != nil {
		printFatal(err.Error())
	}

	client = suite.Client
	apiURL = suite.ApiURL
}

func main() {
	log.SetFlags(0)

	// make sure command is specified, disallow global args
	args := os.Args[1:]
	if len(args) < 1 || strings.IndexRune(args[0], '-') == 0 {
		printUsageTo(os.Stderr)
		os.Exit(2)
	}

	inFd, isTerminalIn = term.GetFdInfo(os.Stdin)
	outFd, isTerminalOut = term.GetFdInfo(os.Stdout)

	if !isTerminalOut {
		ansi.DisableColors(true)
	}

	for _, cmd := range commands {
		if matchesCommand(cmd, args[0]) && cmd.Run != nil {
			defer recoverPanic()

			if !cmd.NoClient {
				initClients()
			}

			cmd.Flag.SetDisableDuplicates(true) // disallow duplicate flag options
			if !gitConfigBool("hk.strict-flag-ordering") {
				cmd.Flag.SetInterspersed(true) // allow flags & non-flag args to mix
			}
			cmd.Flag.Usage = func() {
				cmd.PrintUsage()
			}
			if cmd.NeedsApp || cmd.OptionalApp {
				cmd.Flag.StringVarP(&flagApp, "app", "a", "", "app name")
			}
			if cmd.OptionalMessage {
				cmd.Flag.StringVarP(&flagMessage, "message", "m", "", "message")
			}
			if err := cmd.Flag.Parse(args[1:]); err == flag.ErrHelp {
				cmdHelp.Run(cmdHelp, args[:1])
				return
			} else if err != nil {
				printError(err.Error())
				os.Exit(2)
			}
			if flagApp != "" {
				if gitRemoteApp, err := appFromGitRemote(flagApp); err == nil {
					flagApp = gitRemoteApp
				}
			}
			if cmd.NeedsApp || cmd.OptionalApp {
				a, err := app()
				switch {
				case err == errMultipleHerokuRemotes && cmd.NeedsApp, err == nil && a == "" && cmd.NeedsApp:
					msg := "no app specified"
					if err != nil {
						msg = err.Error()
					}
					printError(msg)
					cmd.PrintUsage()
					os.Exit(2)
				case err != nil && cmd.NeedsApp:
					printFatal(err.Error())
				}
			}
			cmd.Run(cmd, cmd.Flag.Args())
			return
		}
	}

	fmt.Fprintf(os.Stderr, "Unknown command: %s\n", args[0])
	if g := suggest(args[0]); len(g) > 0 {
		fmt.Fprintf(os.Stderr, "Possible alternatives: %v\n", strings.Join(g, " "))
	}
	fmt.Fprintf(os.Stderr, "Run 'emp help' for usage.\n")
	os.Exit(2)
}

func recoverPanic() {
	if Version != "dev" {
		if rec := recover(); rec != nil {
			printFatal("emp encountered and reported an internal client error")
		}
	}
}

func app() (string, error) {
	if flagApp != "" {
		return flagApp, nil
	}

	if app := os.Getenv("EMPAPP"); app != "" {
		return app, nil
	}

	return appFromGitRemote(remoteFromGitConfig())
}

func mustApp() string {
	name, err := app()
	if err != nil {
		printFatal(err.Error())
	}
	return name
}

func getMessage() string {
	return flagMessage
}

// matchesCommand checks if the Command matches the command that we want to run.
func matchesCommand(cmd *Command, want string) bool {
	if cmd.Alias != "" && cmd.Alias == want {
		return true
	}

	return cmd.Name() == want
}
<|endoftext|>


<commit_before>package gateway

import (
	"bufio"
	"time"

	"github.com/bjyoungblood/gozw/common"
	"github.com/bjyoungblood/gozw/zwave"
	"github.com/tarm/serial"
)

// AckCallback is a function callback to be executed when a frame is transmitted.
// status will be one of zwave.FrameHeader*
type AckCallback func(status int)

// Request represents a ZFrame queued for transmission to the controller
type Request struct {
	frame    *zwave.ZFrame
	callback AckCallback
	attempts int
}

// SerialPort is a container/wrapper for the actual serial port, with some
// extra protection to ensure proper connection state with the controller
type SerialPort struct {
	port            *serial.Port
	incomingPrivate chan *zwave.ZFrame
	requestQueue    chan Request
	requestInFlight Request
	Incoming        chan *zwave.ZFrame
}

// NewSerialPort Open a(n actual) serial port and create some supporting channels
func NewSerialPort(config *common.GozwConfig) (*SerialPort, error) {
	// Open the serial port with the given device and baud rate
	// Note: could probably consider inlining the baud rate, since it should
	// always be 115200
	port, err := serial.OpenPort(&serial.Config{
		Name: config.Device,
		Baud: config.Baud,
	})

	if err != nil {
		return nil, err
	}

	// Channel for Z-Wave commands we need to queue up
	requestQueue := make(chan Request, 1)

	// Channel for frames we receive from the Z-Wave controller, but don't necessarily
	// want to make public (yet), since this can include ACKs, NAKs, and CANs
	incomingPrivate := make(chan *zwave.ZFrame, 1)

	// Channel for frames we want to release into the wild
	incomingPublic := make(chan *zwave.ZFrame, 1)

	serialPort := SerialPort{
		port:            port,
		incomingPrivate: incomingPrivate,
		requestQueue:    requestQueue,
		Incoming:        incomingPublic,
	}

	return &serialPort, nil
}

// Initialize We need to do some initial setup on the device before we are able
// to enter our normal handler loop
func (s *SerialPort) Initialize() {

	// According to 6.1 in the Serial API guide, we're supposed to start up by
	// sending a NAK, then doing a hard or soft reset. Soft reset isn't implemented
	// yet, and I don't know if hard reset is possible with a USB controller
	err := s.sendNak()
	if err != nil {
		panic(err)
	}

	// Read frames from the serial port in a goroutine, and make them available on the
	// incomingPrivate channel
	go readFrames(bufio.NewReader(s.port), s.incomingPrivate)

	// This block will block to receive incoming frames and continue to do so until
	// 2 seconds after it has received the last frame. We do this because if we
	// previously crashed, a quick startup could bring us up while the controller is
	// still retransmitting frames we haven't ACKed, and we might not know what to do
	// with them
	for {
		select {
		case frame := <-s.incomingPrivate:
			// this runs in a goroutine in case nothing is listening to s.Incoming yet
			// the goroutine basically just blocks until something listens.
			go func(frame *zwave.ZFrame) {
				s.Incoming <- frame
			}(frame)
		case <-time.After(time.Second * 2):
			// after 2 seconds of not receiving any frames, return
			return
		}
	}
}

// Run Handles unsolicited incoming frames and transmits outgoing frames queued
// using SendFrame
func (s *SerialPort) Run() {
	for {
		select {
		case incoming := <-s.incomingPrivate:
			err := incoming.VerifyChecksum()
			if err != nil {
				s.sendNak()
				continue
			} else if incoming.IsData() {
				s.sendAck()
				continue
			}

			// If the next frame is an ACK or a NAK, respond with it
			if incoming.IsAck() || incoming.IsNak() {
				if &s.requestInFlight != nil && s.requestInFlight.callback != nil {
					s.requestInFlight.callback(int(incoming.Header))
					s.requestInFlight = Request{}
					continue
				}
			}

		case request := <-s.requestQueue:
			s.requestInFlight = request
			_, err := s.port.Write(request.frame.Marshal())
			if err != nil {
				panic(err)
			}
		}
	}
}

// SendFrameSync wraps SendFrame with some magic that blocks until the result
// arrives
func (s *SerialPort) SendFrameSync(frame *zwave.ZFrame) int {
	// Make a channel we can block on
	await := make(chan int, 1)

	// All our callback needs to do is publish the response frame back to the channel
	callback := func(response int) {
		await <- response
	}

	// Send the frame in a goroutine, since we don't want to block on this
	go s.SendFrame(frame, callback)

	// Block until the channel emits a value for us, and then return that value
	return <-await
}

// SendFrame queues a frame to be sent to the controller
func (s *SerialPort) SendFrame(frame *zwave.ZFrame, callback AckCallback) {
	go func(frame *zwave.ZFrame, callback AckCallback) {
		s.requestQueue <- Request{
			frame:    frame,
			callback: callback,
		}
	}(frame, callback)
}

// Close the serial port
func (s *SerialPort) Close() error {
	return s.port.Close()
}

func (s *SerialPort) sendAck() error {
	_, err := s.port.Write(zwave.NewAckFrame().Marshal())
	return err
}

func (s *SerialPort) sendNak() error {
	_, err := s.port.Write(zwave.NewNakFrame().Marshal())
	return err
}
<commit_msg>fix frame emit flow<commit_after>package gateway

import (
	"bufio"
	"time"

	"github.com/bjyoungblood/gozw/common"
	"github.com/bjyoungblood/gozw/zwave"
	"github.com/tarm/serial"
)

// AckCallback is a function callback to be executed when a frame is transmitted.
// status will be one of zwave.FrameHeader*
type AckCallback func(status int)

// Request represents a ZFrame queued for transmission to the controller
type Request struct {
	frame    *zwave.ZFrame
	callback AckCallback
	attempts int
}

// SerialPort is a container/wrapper for the actual serial port, with some
// extra protection to ensure proper connection state with the controller
type SerialPort struct {
	port            *serial.Port
	incomingPrivate chan *zwave.ZFrame
	requestQueue    chan Request
	requestInFlight Request
	Incoming        chan *zwave.ZFrame
}

// NewSerialPort Open a(n actual) serial port and create some supporting channels
func NewSerialPort(config *common.GozwConfig) (*SerialPort, error) {
	// Open the serial port with the given device and baud rate
	// Note: could probably consider inlining the baud rate, since it should
	// always be 115200
	port, err := serial.OpenPort(&serial.Config{
		Name: config.Device,
		Baud: config.Baud,
	})

	if err != nil {
		return nil, err
	}

	// Channel for Z-Wave commands we need to queue up
	requestQueue := make(chan Request, 1)

	// Channel for frames we receive from the Z-Wave controller, but don't necessarily
	// want to make public (yet), since this can include ACKs, NAKs, and CANs
	incomingPrivate := make(chan *zwave.ZFrame, 1)

	// Channel for frames we want to release into the wild
	incomingPublic := make(chan *zwave.ZFrame, 1)

	serialPort := SerialPort{
		port:            port,
		incomingPrivate: incomingPrivate,
		requestQueue:    requestQueue,
		Incoming:        incomingPublic,
	}

	return &serialPort, nil
}

// Initialize We need to do some initial setup on the device before we are able
// to enter our normal handler loop
func (s *SerialPort) Initialize() {

	// According to 6.1 in the Serial API guide, we're supposed to start up by
	// sending a NAK, then doing a hard or soft reset. Soft reset isn't implemented
	// yet, and I don't know if hard reset is possible with a USB controller
	err := s.sendNak()
	if err != nil {
		panic(err)
	}

	// Read frames from the serial port in a goroutine, and make them available on the
	// incomingPrivate channel
	go readFrames(bufio.NewReader(s.port), s.incomingPrivate)

	// This block will block to receive incoming frames and continue to do so until
	// 2 seconds after it has received the last frame. We do this because if we
	// previously crashed, a quick startup could bring us up while the controller is
	// still retransmitting frames we haven't ACKed, and we might not know what to do
	// with them
	for {
		select {
		case frame := <-s.incomingPrivate:
			// this runs in a goroutine in case nothing is listening to s.Incoming yet
			// the goroutine basically just blocks until something listens.
			go func(frame *zwave.ZFrame) {
				s.Incoming <- frame
			}(frame)
		case <-time.After(time.Second * 2):
			// after 2 seconds of not receiving any frames, return
			return
		}
	}
}

// Run Handles unsolicited incoming frames and transmits outgoing frames queued
// using SendFrame
func (s *SerialPort) Run() {
	for {
		select {
		case incoming := <-s.incomingPrivate:
			err := incoming.VerifyChecksum()
			if err != nil {
				s.sendNak()
				continue
			} else if incoming.IsAck() || incoming.IsNak() {
				if &s.requestInFlight != nil && s.requestInFlight.callback != nil {
					s.requestInFlight.callback(int(incoming.Header))
					s.requestInFlight = Request{}
					continue
				}
			} else if incoming.IsData() {
				// If everything else has been processed, then release it into the wild
				s.sendAck()
				s.Incoming <- incoming
			}

		case request := <-s.requestQueue:
			s.requestInFlight = request
			_, err := s.port.Write(request.frame.Marshal())
			if err != nil {
				panic(err)
			}
		}
	}
}

// SendFrameSync wraps SendFrame with some magic that blocks until the result
// arrives
func (s *SerialPort) SendFrameSync(frame *zwave.ZFrame) int {
	// Make a channel we can block on
	await := make(chan int, 1)

	// All our callback needs to do is publish the response frame back to the channel
	callback := func(response int) {
		await <- response
	}

	// Send the frame in a goroutine, since we don't want to block on this
	go s.SendFrame(frame, callback)

	// Block until the channel emits a value for us, and then return that value
	return <-await
}

// SendFrame queues a frame to be sent to the controller
func (s *SerialPort) SendFrame(frame *zwave.ZFrame, callback AckCallback) {
	go func(frame *zwave.ZFrame, callback AckCallback) {
		s.requestQueue <- Request{
			frame:    frame,
			callback: callback,
		}
	}(frame, callback)
}

// Close the serial port
func (s *SerialPort) Close() error {
	return s.port.Close()
}

func (s *SerialPort) sendAck() error {
	_, err := s.port.Write(zwave.NewAckFrame().Marshal())
	return err
}

func (s *SerialPort) sendNak() error {
	_, err := s.port.Write(zwave.NewNakFrame().Marshal())
	return err
}
<|endoftext|>


<commit_before>package broker

import (
	"context"
	"sync"
	"time"

	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"

	"github.com/celrenheit/sandflake"
	"github.com/gogo/protobuf/proto"

	"github.com/celrenheit/sandglass-grpc/go/sgproto"
	"golang.org/x/sync/errgroup"
)

var (
	// TODO: make these variables configurable
	RedeliveryTimeout  = 10 * time.Second
	MaxRedeliveryCount = 5
)

type ConsumerGroup struct {
	broker    *Broker
	topic     string
	partition string
	name      string
	mu        sync.RWMutex
	receivers []*receiver
}

func NewConsumerGroup(b *Broker, topic, partition, name string) *ConsumerGroup {
	return &ConsumerGroup{
		broker:    b,
		name:      name,
		topic:     topic,
		partition: partition,
	}
}

type receiver struct {
	name   string
	msgCh  chan *sgproto.Message
	doneCh chan struct{}
}

func (c *ConsumerGroup) register(consumerName string) *receiver {
	r := c.getReceiver(consumerName)
	if r != nil {
		return r
	}

	c.mu.Lock()
	defer c.mu.Unlock()

	r = &receiver{
		name:   consumerName,
		msgCh:  make(chan *sgproto.Message),
		doneCh: make(chan struct{}),
	}
	c.receivers = append(c.receivers, r)

	if len(c.receivers) == 1 {
		go c.consumeLoop()
	}

	return r
}

func (c *ConsumerGroup) consumeLoop() {
	defer func() { // close receivers for whatever reason
		c.mu.Lock()
		for _, r := range c.receivers {
			close(r.msgCh)
			close(r.doneCh)
		}
		c.receivers = c.receivers[:0]
		c.mu.Unlock()
	}()

	lastCommited, err := c.broker.LastOffset(context.TODO(), c.topic, c.partition, c.name, "", sgproto.MarkKind_Commited)
	if err != nil {
		c.broker.Debug("got error when fetching last committed offset: %v ", err)
		return
	}

	from, err := c.broker.LastOffset(context.TODO(), c.topic, c.partition, c.name, "", sgproto.MarkKind_Consumed)
	if err != nil {
		c.broker.Debug("got error when fetching last committed offset: %v ", err)
		return
	}

	msgCh := make(chan *sgproto.Message)
	var group errgroup.Group

	if !lastCommited.Equal(from) {
		group.Go(func() error {
			var (
				lastMessage *sgproto.Message
				committed   = false
			)
			req := &sgproto.FetchRangeRequest{
				Topic:     c.topic,
				Partition: c.partition,
				From:      lastCommited,
				To:        from,
			}

			commit := func(offset sandflake.ID) {
				_, err := c.broker.Commit(context.TODO(), c.topic, c.partition, c.name, "", lastMessage.Offset)
				if err != nil {
					c.broker.Debug("unable to commit")
				}
			}

			i := 0
			err := c.broker.FetchRange(context.TODO(), req, func(m *sgproto.Message) error {
				if m.Offset.Equal(lastCommited) { // skip first item, since it is already committed
					lastMessage = m
					return nil
				}
				i++

				msg, err := c.broker.GetMarkStateMessage(context.TODO(), c.topic, c.partition, c.name, "", m.Offset)
				if err != nil {
					s, ok := status.FromError(err)
					if !ok || s.Code() != codes.NotFound {
						return err
					}
				}

				var state sgproto.MarkState
				if msg != nil {
					err := proto.Unmarshal(msg.Value, &state)
					if err != nil {
						return err
					}
				}

				// advance commit offset
				// if we only got acked messages before
				if !committed && lastMessage != nil {
					if state.Kind != sgproto.MarkKind_Acknowledged {
						// we might commit in a goroutine, we can redo this the next time we consume
						if !lastMessage.Offset.Equal(lastCommited) {
							commit(lastMessage.Offset)
						}
						committed = true
					} else if i%10000 == 0 {
						go commit(lastMessage.Offset)
					}
				}
				lastMessage = m

				if shouldRedeliver(m.Index, state) {
					msgCh <- m // deliver

					// those calls should be batched
					if state.Kind == sgproto.MarkKind_Unknown {
						// TODO: Should this be nacked?
						_, err := c.broker.NotAcknowledge(context.Background(), c.topic, c.partition, c.name, "NOT SET", m.Offset)
						if err != nil {
							c.broker.Debug("error while acking message for the first redilvery", err)
							return err
						}
					} else {
						state.DeliveryCount++

						if int(state.DeliveryCount) >= MaxRedeliveryCount {
							// Mark the message as ACKed
							// TODO: produce this a dead letter queue
							state.Kind = sgproto.MarkKind_Acknowledged
						}

						msg.Value, err = proto.Marshal(&state)
						if err != nil {
							return err
						}

						// TODO: Should handle this in higher level method
						t := c.broker.GetTopic(ConsumerOffsetTopicName)
						p := t.ChoosePartitionForKey(msg.Key)
						msg.ClusteringKey = generateClusterKey(msg.Offset, state.Kind)

						if _, err := c.broker.Produce(context.TODO(), &sgproto.ProduceMessageRequest{
							Topic:     ConsumerOffsetTopicName,
							Partition: p.Id,
							Messages:  []*sgproto.Message{msg},
						}); err != nil {
							return err
						}
					}
				}

				return nil
			})
			if err != nil {
				return err
			}

			if !committed && lastMessage != nil {
				commit(lastMessage.Offset)
			}

			return nil
		})
	}
	group.Go(func() error {
		now := sandflake.NewID(time.Now().UTC(), sandflake.MaxID.WorkerID(), sandflake.MaxID.Sequence(), sandflake.MaxID.RandomBytes())
		req := &sgproto.FetchRangeRequest{
			Topic:     c.topic,
			Partition: c.partition,
			From:      from,
			To:        now,
		}

		return c.broker.FetchRange(context.TODO(), req, func(m *sgproto.Message) error {
			// skip the first if it is the same as the starting point
			if from == m.Offset {
				return nil
			}

			msgCh <- m

			return nil
		})
	})

	go func() {
		err := group.Wait()
		if err != nil {
			c.broker.Info("error in consumeLoop: %v", err)
		}
		close(msgCh)
	}()

	var i int
	var m *sgproto.Message
loop:
	for m = range msgCh {
		// select receiver
	selectreceiver:
		i++
		c.mu.RLock()
		r := c.receivers[i%len(c.receivers)]
		c.mu.RUnlock()

		select {
		case <-r.doneCh:
			if c.removeConsumer(r.name) {
				c.mu.RLock()
				l := len(c.receivers)
				c.mu.RUnlock()

				if l == 0 {
					break loop
				}

				goto selectreceiver // select another receiver
			}
		case r.msgCh <- m:
		}
	}

	if m != nil && !m.Offset.Equal(from) {
		_, err := c.broker.MarkConsumed(context.TODO(), c.topic, c.partition, c.name, "REMOVE THIS", m.Offset)
		if err != nil {
			c.broker.Debug("unable to mark as consumed: %v", err)
		}
	}
}

func shouldRedeliver(index sandflake.ID, state sgproto.MarkState) bool {
	switch state.Kind {
	case sgproto.MarkKind_NotAcknowledged:
		return true
	case sgproto.MarkKind_Consumed, sgproto.MarkKind_Unknown: // inflight
		return index.Time().Add(RedeliveryTimeout).Before(time.Now().UTC())
	case sgproto.MarkKind_Acknowledged, sgproto.MarkKind_Commited:
		return false
	default:
		panic("unknown markkind: " + state.Kind.String())
	}

	return false
}

func (c *ConsumerGroup) removeConsumer(name string) bool {
	c.mu.Lock()
	defer c.mu.Unlock()

	for i, r := range c.receivers {
		if r.name == name {
			c.receivers = append(c.receivers[:i], c.receivers[i+1:]...)
			return true
		}
	}

	return false
}

func (c *ConsumerGroup) getReceiver(consumerName string) *receiver {
	c.mu.RLock()
	defer c.mu.RUnlock()
	for _, r := range c.receivers {
		if r.name == consumerName {
			return r
		}
	}

	return nil
}

func (c *ConsumerGroup) Consume(consumerName string) (<-chan *sgproto.Message, chan<- struct{}, error) {
	r := c.register(consumerName)

	return r.msgCh, r.doneCh, nil
}
<commit_msg>use the correct offset for clustering key, mark consumed instead of nack and adjust delay for redelivery<commit_after>package broker

import (
	"context"
	"sync"
	"time"

	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"

	"github.com/celrenheit/sandflake"
	"github.com/gogo/protobuf/proto"

	"github.com/celrenheit/sandglass-grpc/go/sgproto"
	"golang.org/x/sync/errgroup"
)

var (
	// TODO: make these variables configurable
	RedeliveryTimeout  = 10 * time.Second
	MaxRedeliveryCount = 5
)

type ConsumerGroup struct {
	broker    *Broker
	topic     string
	partition string
	name      string
	mu        sync.RWMutex
	receivers []*receiver
}

func NewConsumerGroup(b *Broker, topic, partition, name string) *ConsumerGroup {
	return &ConsumerGroup{
		broker:    b,
		name:      name,
		topic:     topic,
		partition: partition,
	}
}

type receiver struct {
	name   string
	msgCh  chan *sgproto.Message
	doneCh chan struct{}
}

func (c *ConsumerGroup) register(consumerName string) *receiver {
	r := c.getReceiver(consumerName)
	if r != nil {
		return r
	}

	c.mu.Lock()
	defer c.mu.Unlock()

	r = &receiver{
		name:   consumerName,
		msgCh:  make(chan *sgproto.Message),
		doneCh: make(chan struct{}),
	}
	c.receivers = append(c.receivers, r)

	if len(c.receivers) == 1 {
		go c.consumeLoop()
	}

	return r
}

func (c *ConsumerGroup) consumeLoop() {
	defer func() { // close receivers for whatever reason
		c.mu.Lock()
		for _, r := range c.receivers {
			close(r.msgCh)
			close(r.doneCh)
		}
		c.receivers = c.receivers[:0]
		c.mu.Unlock()
	}()

	lastCommited, err := c.broker.LastOffset(context.TODO(), c.topic, c.partition, c.name, "", sgproto.MarkKind_Commited)
	if err != nil {
		c.broker.Debug("got error when fetching last committed offset: %v ", err)
		return
	}

	from, err := c.broker.LastOffset(context.TODO(), c.topic, c.partition, c.name, "", sgproto.MarkKind_Consumed)
	if err != nil {
		c.broker.Debug("got error when fetching last committed offset: %v ", err)
		return
	}

	msgCh := make(chan *sgproto.Message)
	var group errgroup.Group

	if !lastCommited.Equal(from) {
		group.Go(func() error {
			var (
				lastMessage *sgproto.Message
				committed   = false
			)
			req := &sgproto.FetchRangeRequest{
				Topic:     c.topic,
				Partition: c.partition,
				From:      lastCommited,
				To:        from,
			}

			commit := func(offset sandflake.ID) {
				_, err := c.broker.Commit(context.TODO(), c.topic, c.partition, c.name, "", lastMessage.Offset)
				if err != nil {
					c.broker.Debug("unable to commit")
				}
			}

			i := 0
			err := c.broker.FetchRange(context.TODO(), req, func(m *sgproto.Message) error {
				if m.Offset.Equal(lastCommited) { // skip first item, since it is already committed
					lastMessage = m
					return nil
				}
				i++

				msg, err := c.broker.GetMarkStateMessage(context.TODO(), c.topic, c.partition, c.name, "", m.Offset)
				if err != nil {
					s, ok := status.FromError(err)
					if !ok || s.Code() != codes.NotFound {
						return err
					}
				}

				var state sgproto.MarkState
				if msg != nil {
					err := proto.Unmarshal(msg.Value, &state)
					if err != nil {
						return err
					}
				}

				// advance commit offset
				// if we only got acked messages before
				if !committed && lastMessage != nil {
					if state.Kind != sgproto.MarkKind_Acknowledged {
						// we might commit in a goroutine, we can redo this the next time we consume
						if !lastMessage.Offset.Equal(lastCommited) {
							commit(lastMessage.Offset)
						}
						committed = true
					} else if i%10000 == 0 {
						go commit(lastMessage.Offset)
					}
				}
				lastMessage = m

				if shouldRedeliver(m.Index, state) {
					msgCh <- m // deliver

					// those calls should be batched
					if state.Kind == sgproto.MarkKind_Unknown {
						// TODO: Should we mark this consumed?
						_, err := c.broker.MarkConsumed(context.Background(), c.topic, c.partition, c.name, "NOT SET", m.Offset)
						if err != nil {
							c.broker.Debug("error while acking message for the first redilvery", err)
							return err
						}
					} else {
						state.DeliveryCount++

						if int(state.DeliveryCount) >= MaxRedeliveryCount {
							// Mark the message as ACKed
							// TODO: produce this a dead letter queue
							state.Kind = sgproto.MarkKind_Acknowledged
						}

						msg.Value, err = proto.Marshal(&state)
						if err != nil {
							return err
						}

						// TODO: Should handle this in higher level method
						t := c.broker.GetTopic(ConsumerOffsetTopicName)
						p := t.ChoosePartitionForKey(msg.Key)
						msg.ClusteringKey = generateClusterKey(m.Offset, state.Kind)

						if _, err := c.broker.Produce(context.TODO(), &sgproto.ProduceMessageRequest{
							Topic:     ConsumerOffsetTopicName,
							Partition: p.Id,
							Messages:  []*sgproto.Message{msg},
						}); err != nil {
							return err
						}
					}
				}

				return nil
			})
			if err != nil {
				return err
			}

			if !committed && lastMessage != nil {
				commit(lastMessage.Offset)
			}

			return nil
		})
	}
	group.Go(func() error {
		now := sandflake.NewID(time.Now().UTC(), sandflake.MaxID.WorkerID(), sandflake.MaxID.Sequence(), sandflake.MaxID.RandomBytes())
		req := &sgproto.FetchRangeRequest{
			Topic:     c.topic,
			Partition: c.partition,
			From:      from,
			To:        now,
		}

		return c.broker.FetchRange(context.TODO(), req, func(m *sgproto.Message) error {
			// skip the first if it is the same as the starting point
			if from == m.Offset {
				return nil
			}

			msgCh <- m

			return nil
		})
	})

	go func() {
		err := group.Wait()
		if err != nil {
			c.broker.Info("error in consumeLoop: %v", err)
		}
		close(msgCh)
	}()

	var i int
	var m *sgproto.Message
loop:
	for m = range msgCh {
		// select receiver
	selectreceiver:
		i++
		c.mu.RLock()
		r := c.receivers[i%len(c.receivers)]
		c.mu.RUnlock()

		select {
		case <-r.doneCh:
			if c.removeConsumer(r.name) {
				c.mu.RLock()
				l := len(c.receivers)
				c.mu.RUnlock()

				if l == 0 {
					break loop
				}

				goto selectreceiver // select another receiver
			}
		case r.msgCh <- m:
		}
	}

	if m != nil && !m.Offset.Equal(from) {
		_, err := c.broker.MarkConsumed(context.TODO(), c.topic, c.partition, c.name, "REMOVE THIS", m.Offset)
		if err != nil {
			c.broker.Debug("unable to mark as consumed: %v", err)
		}
	}
}

func shouldRedeliver(index sandflake.ID, state sgproto.MarkState) bool {
	switch state.Kind {
	case sgproto.MarkKind_NotAcknowledged:
		return true
	case sgproto.MarkKind_Consumed, sgproto.MarkKind_Unknown: // inflight
		return index.Time().Add(time.Duration(state.DeliveryCount) * RedeliveryTimeout).Before(time.Now().UTC())
	case sgproto.MarkKind_Acknowledged, sgproto.MarkKind_Commited:
		return false
	default:
		panic("unknown markkind: " + state.Kind.String())
	}

	return false
}

func (c *ConsumerGroup) removeConsumer(name string) bool {
	c.mu.Lock()
	defer c.mu.Unlock()

	for i, r := range c.receivers {
		if r.name == name {
			c.receivers = append(c.receivers[:i], c.receivers[i+1:]...)
			return true
		}
	}

	return false
}

func (c *ConsumerGroup) getReceiver(consumerName string) *receiver {
	c.mu.RLock()
	defer c.mu.RUnlock()
	for _, r := range c.receivers {
		if r.name == consumerName {
			return r
		}
	}

	return nil
}

func (c *ConsumerGroup) Consume(consumerName string) (<-chan *sgproto.Message, chan<- struct{}, error) {
	r := c.register(consumerName)

	return r.msgCh, r.doneCh, nil
}
<|endoftext|>


<commit_before>package main

import (
	"../nsq"
	"github.com/bmizerany/assert"
	"io/ioutil"
	"log"
	"os"
	"strconv"
	"testing"
	"time"
)

// ensure that we can push a message through a topic and get it out of a channel
func TestPutMessage(t *testing.T) {
	log.SetOutput(ioutil.Discard)
	defer log.SetOutput(os.Stdout)

	nsqd := NewNSQd(1)

	topicName := "test_put_message" + strconv.Itoa(int(time.Now().Unix()))
	topic := nsqd.GetTopic(topicName)
	channel1 := topic.GetChannel("ch")

	msg := nsq.NewMessage([]byte("abcdefghijklmnop"), []byte("test"))
	topic.PutMessage(msg)

	outputMsg := <-channel1.clientMessageChan
	assert.Equal(t, msg.Id, outputMsg.Id)
	assert.Equal(t, msg.Body, outputMsg.Body)
}

// ensure that both channels get the same message
func TestPutMessage2Chan(t *testing.T) {
	log.SetOutput(ioutil.Discard)
	defer log.SetOutput(os.Stdout)

	nsqd := NewNSQd(1)

	topicName := "test_put_message_2chan" + strconv.Itoa(int(time.Now().Unix()))
	topic := nsqd.GetTopic(topicName)
	channel1 := topic.GetChannel("ch1")
	channel2 := topic.GetChannel("ch2")

	msg := nsq.NewMessage([]byte("abcdefghijklmnop"), []byte("test"))
	topic.PutMessage(msg)

	outputMsg1 := <-channel1.clientMessageChan
	assert.Equal(t, msg.Id, outputMsg1.Id)
	assert.Equal(t, msg.Body, outputMsg1.Body)

	outputMsg2 := <-channel2.clientMessageChan
	assert.Equal(t, msg.Id, outputMsg2.Id)
	assert.Equal(t, msg.Body, outputMsg2.Body)
}
<commit_msg>add in flight worker test<commit_after>package main

import (
	"../nsq"
	"github.com/bmizerany/assert"
	"io/ioutil"
	"log"
	"os"
	"strconv"
	"testing"
	"time"
)

// ensure that we can push a message through a topic and get it out of a channel
func TestPutMessage(t *testing.T) {
	log.SetOutput(ioutil.Discard)
	defer log.SetOutput(os.Stdout)

	nsqd := NewNSQd(1)

	topicName := "test_put_message" + strconv.Itoa(int(time.Now().Unix()))
	topic := nsqd.GetTopic(topicName)
	channel1 := topic.GetChannel("ch")

	msg := nsq.NewMessage([]byte("abcdefghijklmnop"), []byte("test"))
	topic.PutMessage(msg)

	outputMsg := <-channel1.clientMessageChan
	assert.Equal(t, msg.Id, outputMsg.Id)
	assert.Equal(t, msg.Body, outputMsg.Body)
}

// ensure that both channels get the same message
func TestPutMessage2Chan(t *testing.T) {
	log.SetOutput(ioutil.Discard)
	defer log.SetOutput(os.Stdout)

	nsqd := NewNSQd(1)

	topicName := "test_put_message_2chan" + strconv.Itoa(int(time.Now().Unix()))
	topic := nsqd.GetTopic(topicName)
	channel1 := topic.GetChannel("ch1")
	channel2 := topic.GetChannel("ch2")

	msg := nsq.NewMessage([]byte("abcdefghijklmnop"), []byte("test"))
	topic.PutMessage(msg)

	outputMsg1 := <-channel1.clientMessageChan
	assert.Equal(t, msg.Id, outputMsg1.Id)
	assert.Equal(t, msg.Body, outputMsg1.Body)

	outputMsg2 := <-channel2.clientMessageChan
	assert.Equal(t, msg.Id, outputMsg2.Id)
	assert.Equal(t, msg.Body, outputMsg2.Body)
}

func TestInFlightWorker(t *testing.T) {
	log.SetOutput(ioutil.Discard)
	defer log.SetOutput(os.Stdout)

	nsqd := NewNSQd(1)
	nsqd.msgTimeout = 300 * time.Millisecond

	topic := nsqd.GetTopic("topic")
	channel := topic.GetChannel("channel")

	for i := 0; i < 1000; i++ {
		msg := nsq.NewMessage(<-nsqd.idChan, []byte("test"))
		channel.StartInFlightTimeout(msg, NewClientV2(nil))
	}

	assert.Equal(t, len(channel.inFlightMessages), 1000)
	assert.Equal(t, len(channel.inFlightPQ), 1000)

	time.Sleep(350 * time.Millisecond)

	assert.Equal(t, len(channel.inFlightMessages), 0)
	assert.Equal(t, len(channel.inFlightPQ), 0)
}
<|endoftext|>


<commit_before>package main

import (
	"io/ioutil"
	"log"
	"os"
	"strconv"
	"testing"
	"time"

	"github.com/bitly/go-nsq"
	"github.com/bmizerany/assert"
)

// ensure that we can push a message through a topic and get it out of a channel
func TestPutMessage(t *testing.T) {
	log.SetOutput(ioutil.Discard)
	defer log.SetOutput(os.Stdout)

	_, _, nsqd := mustStartNSQD(NewNSQDOptions())
	defer nsqd.Exit()

	topicName := "test_put_message" + strconv.Itoa(int(time.Now().Unix()))
	topic := nsqd.GetTopic(topicName)
	channel1 := topic.GetChannel("ch")

	var id nsq.MessageID
	msg := nsq.NewMessage(id, []byte("test"))
	topic.PutMessage(msg)

	outputMsg := <-channel1.clientMsgChan
	assert.Equal(t, msg.Id, outputMsg.Id)
	assert.Equal(t, msg.Body, outputMsg.Body)
}

// ensure that both channels get the same message
func TestPutMessage2Chan(t *testing.T) {
	log.SetOutput(ioutil.Discard)
	defer log.SetOutput(os.Stdout)

	_, _, nsqd := mustStartNSQD(NewNSQDOptions())
	defer nsqd.Exit()

	topicName := "test_put_message_2chan" + strconv.Itoa(int(time.Now().Unix()))
	topic := nsqd.GetTopic(topicName)
	channel1 := topic.GetChannel("ch1")
	channel2 := topic.GetChannel("ch2")

	var id nsq.MessageID
	msg := nsq.NewMessage(id, []byte("test"))
	topic.PutMessage(msg)

	outputMsg1 := <-channel1.clientMsgChan
	assert.Equal(t, msg.Id, outputMsg1.Id)
	assert.Equal(t, msg.Body, outputMsg1.Body)

	outputMsg2 := <-channel2.clientMsgChan
	assert.Equal(t, msg.Id, outputMsg2.Id)
	assert.Equal(t, msg.Body, outputMsg2.Body)
}

func TestInFlightWorker(t *testing.T) {
	log.SetOutput(ioutil.Discard)
	defer log.SetOutput(os.Stdout)

	options := NewNSQDOptions()
	options.MsgTimeout = 200 * time.Millisecond
	_, _, nsqd := mustStartNSQD(options)
	defer nsqd.Exit()

	topicName := "test_in_flight_worker" + strconv.Itoa(int(time.Now().Unix()))
	topic := nsqd.GetTopic(topicName)
	channel := topic.GetChannel("channel")

	for i := 0; i < 1000; i++ {
		msg := nsq.NewMessage(<-nsqd.idChan, []byte("test"))
		channel.StartInFlightTimeout(msg, 0, options.MsgTimeout)
	}

	channel.Lock()
	inFlightMsgs := len(channel.inFlightMessages)
	channel.Unlock()
	assert.Equal(t, inFlightMsgs, 1000)

	channel.inFlightMutex.Lock()
	inFlightPQMsgs := len(channel.inFlightPQ)
	channel.inFlightMutex.Unlock()
	assert.Equal(t, inFlightPQMsgs, 1000)

	// the in flight worker has a resolution of 100ms so we need to wait
	// at least that much longer than our msgTimeout (in worst case)
	time.Sleep(2 * options.MsgTimeout)

	channel.Lock()
	inFlightMsgs = len(channel.inFlightMessages)
	channel.Unlock()
	assert.Equal(t, inFlightMsgs, 0)

	channel.inFlightMutex.Lock()
	inFlightPQMsgs = len(channel.inFlightPQ)
	channel.inFlightMutex.Unlock()
	assert.Equal(t, inFlightPQMsgs, 0)
}

func TestChannelEmpty(t *testing.T) {
	log.SetOutput(ioutil.Discard)
	defer log.SetOutput(os.Stdout)

	options := NewNSQDOptions()
	_, _, nsqd := mustStartNSQD(options)
	defer nsqd.Exit()

	topicName := "test_channel_empty" + strconv.Itoa(int(time.Now().Unix()))
	topic := nsqd.GetTopic(topicName)
	channel := topic.GetChannel("channel")

	msgs := make([]*nsq.Message, 0, 25)
	for i := 0; i < 25; i++ {
		msg := nsq.NewMessage(<-nsqd.idChan, []byte("test"))
		channel.StartInFlightTimeout(msg, 0, options.MsgTimeout)
		msgs = append(msgs, msg)
	}

	channel.RequeueMessage(0, msgs[len(msgs)-1].Id, 100*time.Millisecond)
	assert.Equal(t, len(channel.inFlightMessages), 24)
	assert.Equal(t, len(channel.inFlightPQ), 24)
	assert.Equal(t, len(channel.deferredMessages), 1)
	assert.Equal(t, len(channel.deferredPQ), 1)

	channel.Empty()

	assert.Equal(t, len(channel.inFlightMessages), 0)
	assert.Equal(t, len(channel.inFlightPQ), 0)
	assert.Equal(t, len(channel.deferredMessages), 0)
	assert.Equal(t, len(channel.deferredPQ), 0)
	assert.Equal(t, channel.Depth(), int64(0))
}

func TestChannelEmptyConsumer(t *testing.T) {
	log.SetOutput(ioutil.Discard)
	defer log.SetOutput(os.Stdout)

	options := NewNSQDOptions()
	tcpAddr, _, nsqd := mustStartNSQD(options)
	defer nsqd.Exit()
	conn, _ := mustConnectNSQD(tcpAddr)

	topicName := "test_channel_empty" + strconv.Itoa(int(time.Now().Unix()))
	topic := nsqd.GetTopic(topicName)
	channel := topic.GetChannel("channel")
	client := NewClientV2(0, conn, &Context{nsqd})
	client.SetReadyCount(25)
	channel.AddClient(client.ID, client)

	for i := 0; i < 25; i++ {
		msg := nsq.NewMessage(<-nsqd.idChan, []byte("test"))
		channel.StartInFlightTimeout(msg, 0, options.MsgTimeout)
		client.SendingMessage()
	}

	for _, cl := range channel.clients {
		stats := cl.Stats()
		assert.Equal(t, stats.InFlightCount, int64(25))
	}

	channel.Empty()

	for _, cl := range channel.clients {
		stats := cl.Stats()
		assert.Equal(t, stats.InFlightCount, int64(0))
	}

}
<commit_msg>nsqd: deflake TestInFlightWorker<commit_after>package main

import (
	"io/ioutil"
	"log"
	"os"
	"strconv"
	"testing"
	"time"

	"github.com/bitly/go-nsq"
	"github.com/bmizerany/assert"
)

// ensure that we can push a message through a topic and get it out of a channel
func TestPutMessage(t *testing.T) {
	log.SetOutput(ioutil.Discard)
	defer log.SetOutput(os.Stdout)

	_, _, nsqd := mustStartNSQD(NewNSQDOptions())
	defer nsqd.Exit()

	topicName := "test_put_message" + strconv.Itoa(int(time.Now().Unix()))
	topic := nsqd.GetTopic(topicName)
	channel1 := topic.GetChannel("ch")

	var id nsq.MessageID
	msg := nsq.NewMessage(id, []byte("test"))
	topic.PutMessage(msg)

	outputMsg := <-channel1.clientMsgChan
	assert.Equal(t, msg.Id, outputMsg.Id)
	assert.Equal(t, msg.Body, outputMsg.Body)
}

// ensure that both channels get the same message
func TestPutMessage2Chan(t *testing.T) {
	log.SetOutput(ioutil.Discard)
	defer log.SetOutput(os.Stdout)

	_, _, nsqd := mustStartNSQD(NewNSQDOptions())
	defer nsqd.Exit()

	topicName := "test_put_message_2chan" + strconv.Itoa(int(time.Now().Unix()))
	topic := nsqd.GetTopic(topicName)
	channel1 := topic.GetChannel("ch1")
	channel2 := topic.GetChannel("ch2")

	var id nsq.MessageID
	msg := nsq.NewMessage(id, []byte("test"))
	topic.PutMessage(msg)

	outputMsg1 := <-channel1.clientMsgChan
	assert.Equal(t, msg.Id, outputMsg1.Id)
	assert.Equal(t, msg.Body, outputMsg1.Body)

	outputMsg2 := <-channel2.clientMsgChan
	assert.Equal(t, msg.Id, outputMsg2.Id)
	assert.Equal(t, msg.Body, outputMsg2.Body)
}

func TestInFlightWorker(t *testing.T) {
	log.SetOutput(ioutil.Discard)
	defer log.SetOutput(os.Stdout)

	count := 250

	options := NewNSQDOptions()
	options.MsgTimeout = 100 * time.Millisecond
	_, _, nsqd := mustStartNSQD(options)
	defer nsqd.Exit()

	topicName := "test_in_flight_worker" + strconv.Itoa(int(time.Now().Unix()))
	topic := nsqd.GetTopic(topicName)
	channel := topic.GetChannel("channel")

	for i := 0; i < count; i++ {
		msg := nsq.NewMessage(<-nsqd.idChan, []byte("test"))
		channel.StartInFlightTimeout(msg, 0, options.MsgTimeout)
	}

	channel.Lock()
	inFlightMsgs := len(channel.inFlightMessages)
	channel.Unlock()
	assert.Equal(t, inFlightMsgs, count)

	channel.inFlightMutex.Lock()
	inFlightPQMsgs := len(channel.inFlightPQ)
	channel.inFlightMutex.Unlock()
	assert.Equal(t, inFlightPQMsgs, count)

	// the in flight worker has a resolution of 100ms so we need to wait
	// at least that much longer than our msgTimeout (in worst case)
	time.Sleep(4 * options.MsgTimeout)

	channel.Lock()
	inFlightMsgs = len(channel.inFlightMessages)
	channel.Unlock()
	assert.Equal(t, inFlightMsgs, 0)

	channel.inFlightMutex.Lock()
	inFlightPQMsgs = len(channel.inFlightPQ)
	channel.inFlightMutex.Unlock()
	assert.Equal(t, inFlightPQMsgs, 0)
}

func TestChannelEmpty(t *testing.T) {
	log.SetOutput(ioutil.Discard)
	defer log.SetOutput(os.Stdout)

	options := NewNSQDOptions()
	_, _, nsqd := mustStartNSQD(options)
	defer nsqd.Exit()

	topicName := "test_channel_empty" + strconv.Itoa(int(time.Now().Unix()))
	topic := nsqd.GetTopic(topicName)
	channel := topic.GetChannel("channel")

	msgs := make([]*nsq.Message, 0, 25)
	for i := 0; i < 25; i++ {
		msg := nsq.NewMessage(<-nsqd.idChan, []byte("test"))
		channel.StartInFlightTimeout(msg, 0, options.MsgTimeout)
		msgs = append(msgs, msg)
	}

	channel.RequeueMessage(0, msgs[len(msgs)-1].Id, 100*time.Millisecond)
	assert.Equal(t, len(channel.inFlightMessages), 24)
	assert.Equal(t, len(channel.inFlightPQ), 24)
	assert.Equal(t, len(channel.deferredMessages), 1)
	assert.Equal(t, len(channel.deferredPQ), 1)

	channel.Empty()

	assert.Equal(t, len(channel.inFlightMessages), 0)
	assert.Equal(t, len(channel.inFlightPQ), 0)
	assert.Equal(t, len(channel.deferredMessages), 0)
	assert.Equal(t, len(channel.deferredPQ), 0)
	assert.Equal(t, channel.Depth(), int64(0))
}

func TestChannelEmptyConsumer(t *testing.T) {
	log.SetOutput(ioutil.Discard)
	defer log.SetOutput(os.Stdout)

	options := NewNSQDOptions()
	tcpAddr, _, nsqd := mustStartNSQD(options)
	defer nsqd.Exit()
	conn, _ := mustConnectNSQD(tcpAddr)

	topicName := "test_channel_empty" + strconv.Itoa(int(time.Now().Unix()))
	topic := nsqd.GetTopic(topicName)
	channel := topic.GetChannel("channel")
	client := NewClientV2(0, conn, &Context{nsqd})
	client.SetReadyCount(25)
	channel.AddClient(client.ID, client)

	for i := 0; i < 25; i++ {
		msg := nsq.NewMessage(<-nsqd.idChan, []byte("test"))
		channel.StartInFlightTimeout(msg, 0, options.MsgTimeout)
		client.SendingMessage()
	}

	for _, cl := range channel.clients {
		stats := cl.Stats()
		assert.Equal(t, stats.InFlightCount, int64(25))
	}

	channel.Empty()

	for _, cl := range channel.clients {
		stats := cl.Stats()
		assert.Equal(t, stats.InFlightCount, int64(0))
	}

}
<|endoftext|>