<commit_before>package splitquery

import (
	"fmt"
	"reflect"
	"testing"

	"github.com/golang/mock/gomock"
	"github.com/youtube/vitess/go/sqltypes"
	"github.com/youtube/vitess/go/vt/sqlparser"
	"github.com/youtube/vitess/go/vt/tabletserver/querytypes"
	"github.com/youtube/vitess/go/vt/tabletserver/splitquery/splitquery_testing"
)

func TestMultipleBoundaries(t *testing.T) {
	mockCtrl := gomock.NewController(t)
	defer mockCtrl.Finish()
	splitParams, err := NewSplitParamsGivenNumRowsPerQueryPart(
		querytypes.BoundQuery{Sql: "select * from test_table where int_col > 5"},
		[]sqlparser.ColIdent{
			sqlparser.NewColIdent("id"),
			sqlparser.NewColIdent("user_id"),
		}, /* splitColumns */
		1000,
		getTestSchema(),
	)
	if err != nil {
		t.Fatalf("NewSplitParamsGivenNumRowsPerQueryPart failed with: %v", err)
	}
	mockSQLExecuter := splitquery_testing.NewMockSQLExecuter(mockCtrl)
	expectedCall1 := mockSQLExecuter.EXPECT().SQLExecute(
		"select id, user_id from test_table force index (`primary`)"+
			" order by id asc, user_id asc"+
			" limit 1000, 1",
		map[string]interface{}{})
	expectedCall1.Return(
		&sqltypes.Result{
			Rows: [][]sqltypes.Value{
				{int64Value(1), int64Value(1)}},
		},
		nil)
	expectedCall2 := mockSQLExecuter.EXPECT().SQLExecute(
		"select id, user_id from test_table force index (`primary`)"+
			" where"+
			" :_splitquery_prev_id < id or"+
			" (:_splitquery_prev_id = id and :_splitquery_prev_user_id <= user_id)"+
			" order by id asc, user_id asc"+
			" limit 1000, 1",
		map[string]interface{}{
			"_splitquery_prev_id":      int64(1),
			"_splitquery_prev_user_id": int64(1),
		})
	expectedCall2.Return(
		&sqltypes.Result{
			Rows: [][]sqltypes.Value{
				{int64Value(2), int64Value(10)}},
		},
		nil)
	expectedCall2.After(expectedCall1)
	expectedCall3 := mockSQLExecuter.EXPECT().SQLExecute(
		"select id, user_id from test_table force index (`primary`)"+
			" where"+
			" :_splitquery_prev_id < id or"+
			" (:_splitquery_prev_id = id and :_splitquery_prev_user_id <= user_id)"+
			" order by id asc, user_id asc"+
			" limit 1000, 1",
		map[string]interface{}{
			"_splitquery_prev_id":      int64(2),
			"_splitquery_prev_user_id": int64(10),
		})
	expectedCall3.Return(
		&sqltypes.Result{Rows: [][]sqltypes.Value{}}, nil)
	expectedCall3.After(expectedCall2)

	algorithm, err := NewFullScanAlgorithm(splitParams, mockSQLExecuter)
	if err != nil {
		t.Fatalf("NewFullScanAlgorithm failed with: %v", err)
	}
	boundaries, err := algorithm.generateBoundaries()
	if err != nil {
		t.Fatalf("FullScanAlgorithm.generateBoundaries() failed with: %v", err)
	}
	expectedBoundaries := []tuple{
		{int64Value(1), int64Value(1)},
		{int64Value(2), int64Value(10)},
	}
	if !reflect.DeepEqual(expectedBoundaries, boundaries) {
		t.Fatalf("expected: %v, got: %v", expectedBoundaries, boundaries)
	}
}

func TestSmallNumberOfRows(t *testing.T) {
	mockCtrl := gomock.NewController(t)
	defer mockCtrl.Finish()
	splitParams, err := NewSplitParamsGivenNumRowsPerQueryPart(
		querytypes.BoundQuery{Sql: "select * from test_table where int_col > 5"},
		[]sqlparser.ColIdent{
			sqlparser.NewColIdent("id"),
			sqlparser.NewColIdent("user_id"),
		}, /* splitColumns */
		1000,
		getTestSchema(),
	)
	if err != nil {
		t.Fatalf("NewSplitParamsGivenNumRowsPerQueryPart failed with: %v", err)
	}
	mockSQLExecuter := splitquery_testing.NewMockSQLExecuter(mockCtrl)
	expectedCall1 := mockSQLExecuter.EXPECT().SQLExecute(
		"select id, user_id from test_table force index (`primary`)"+
			" order by id asc, user_id asc"+
			" limit 1000, 1",
		map[string]interface{}{})
	expectedCall1.Return(
		&sqltypes.Result{Rows: [][]sqltypes.Value{}}, nil)

	algorithm, err := NewFullScanAlgorithm(splitParams, mockSQLExecuter)
	if err != nil {
		t.Fatalf("NewFullScanAlgorithm failed with: %v", err)
	}
	boundaries, err := algorithm.generateBoundaries()
	if err != nil {
		t.Fatalf("FullScanAlgorithm.generateBoundaries() failed with: %v", err)
	}
	expectedBoundaries := []tuple{}
	if !reflect.DeepEqual(expectedBoundaries, boundaries) {
		t.Fatalf("expected: %v, got: %v", expectedBoundaries, boundaries)
	}
}

func TestSQLExecuterReturnsError(t *testing.T) {
	mockCtrl := gomock.NewController(t)
	defer mockCtrl.Finish()
	splitParams, err := NewSplitParamsGivenNumRowsPerQueryPart(
		querytypes.BoundQuery{Sql: "select * from test_table where int_col > 5"},
		[]sqlparser.ColIdent{
			sqlparser.NewColIdent("id"),
			sqlparser.NewColIdent("user_id"),
		}, /* splitColumns */
		1000,
		getTestSchema(),
	)
	if err != nil {
		t.Fatalf("NewSplitParamsGivenNumRowsPerQueryPart failed with: %v", err)
	}
	mockSQLExecuter := splitquery_testing.NewMockSQLExecuter(mockCtrl)
	expectedCall1 := mockSQLExecuter.EXPECT().SQLExecute(
		"select id, user_id from test_table force index (`primary`)"+
			" order by id asc, user_id asc"+
			" limit 1000, 1",
		map[string]interface{}{})
	expectedCall1.Return(
		&sqltypes.Result{
			Rows: [][]sqltypes.Value{
				{int64Value(1), int64Value(1)}},
		},
		nil)
	expectedCall2 := mockSQLExecuter.EXPECT().SQLExecute(
		"select id, user_id from test_table force index (`primary`)"+
			" where"+
			" :_splitquery_prev_id < id or"+
			" (:_splitquery_prev_id = id and :_splitquery_prev_user_id <= user_id)"+
			" order by id asc, user_id asc"+
			" limit 1000, 1",
		map[string]interface{}{
			"_splitquery_prev_id":      int64(1),
			"_splitquery_prev_user_id": int64(1),
		})
	expectedErr := fmt.Errorf("Error accessing database!")
	expectedCall2.Return(nil, expectedErr)
	algorithm, err := NewFullScanAlgorithm(splitParams, mockSQLExecuter)
	if err != nil {
		t.Fatalf("NewFullScanAlgorithm failed with: %v", err)
	}
	boundaries, err := algorithm.generateBoundaries()
	if err != expectedErr {
		t.Fatalf("FullScanAlgorithm.generateBoundaries() did not fail as expected. err: %v", err)
	}
	if boundaries != nil {
		t.Fatalf("boundaries: %v, expected: nil", boundaries)
	}
}
<commit_msg>Fixes full_scan_algorithm_test to use capitalized PRIMARY<commit_after>package splitquery

import (
	"fmt"
	"reflect"
	"testing"

	"github.com/golang/mock/gomock"
	"github.com/youtube/vitess/go/sqltypes"
	"github.com/youtube/vitess/go/vt/sqlparser"
	"github.com/youtube/vitess/go/vt/tabletserver/querytypes"
	"github.com/youtube/vitess/go/vt/tabletserver/splitquery/splitquery_testing"
)

func TestMultipleBoundaries(t *testing.T) {
	mockCtrl := gomock.NewController(t)
	defer mockCtrl.Finish()
	splitParams, err := NewSplitParamsGivenNumRowsPerQueryPart(
		querytypes.BoundQuery{Sql: "select * from test_table where int_col > 5"},
		[]sqlparser.ColIdent{
			sqlparser.NewColIdent("id"),
			sqlparser.NewColIdent("user_id"),
		}, /* splitColumns */
		1000,
		getTestSchema(),
	)
	if err != nil {
		t.Fatalf("NewSplitParamsGivenNumRowsPerQueryPart failed with: %v", err)
	}
	mockSQLExecuter := splitquery_testing.NewMockSQLExecuter(mockCtrl)
	expectedCall1 := mockSQLExecuter.EXPECT().SQLExecute(
		"select id, user_id from test_table force index (`PRIMARY`)"+
			" order by id asc, user_id asc"+
			" limit 1000, 1",
		map[string]interface{}{})
	expectedCall1.Return(
		&sqltypes.Result{
			Rows: [][]sqltypes.Value{
				{int64Value(1), int64Value(1)}},
		},
		nil)
	expectedCall2 := mockSQLExecuter.EXPECT().SQLExecute(
		"select id, user_id from test_table force index (`PRIMARY`)"+
			" where"+
			" :_splitquery_prev_id < id or"+
			" (:_splitquery_prev_id = id and :_splitquery_prev_user_id <= user_id)"+
			" order by id asc, user_id asc"+
			" limit 1000, 1",
		map[string]interface{}{
			"_splitquery_prev_id":      int64(1),
			"_splitquery_prev_user_id": int64(1),
		})
	expectedCall2.Return(
		&sqltypes.Result{
			Rows: [][]sqltypes.Value{
				{int64Value(2), int64Value(10)}},
		},
		nil)
	expectedCall2.After(expectedCall1)
	expectedCall3 := mockSQLExecuter.EXPECT().SQLExecute(
		"select id, user_id from test_table force index (`PRIMARY`)"+
			" where"+
			" :_splitquery_prev_id < id or"+
			" (:_splitquery_prev_id = id and :_splitquery_prev_user_id <= user_id)"+
			" order by id asc, user_id asc"+
			" limit 1000, 1",
		map[string]interface{}{
			"_splitquery_prev_id":      int64(2),
			"_splitquery_prev_user_id": int64(10),
		})
	expectedCall3.Return(
		&sqltypes.Result{Rows: [][]sqltypes.Value{}}, nil)
	expectedCall3.After(expectedCall2)

	algorithm, err := NewFullScanAlgorithm(splitParams, mockSQLExecuter)
	if err != nil {
		t.Fatalf("NewFullScanAlgorithm failed with: %v", err)
	}
	boundaries, err := algorithm.generateBoundaries()
	if err != nil {
		t.Fatalf("FullScanAlgorithm.generateBoundaries() failed with: %v", err)
	}
	expectedBoundaries := []tuple{
		{int64Value(1), int64Value(1)},
		{int64Value(2), int64Value(10)},
	}
	if !reflect.DeepEqual(expectedBoundaries, boundaries) {
		t.Fatalf("expected: %v, got: %v", expectedBoundaries, boundaries)
	}
}

func TestSmallNumberOfRows(t *testing.T) {
	mockCtrl := gomock.NewController(t)
	defer mockCtrl.Finish()
	splitParams, err := NewSplitParamsGivenNumRowsPerQueryPart(
		querytypes.BoundQuery{Sql: "select * from test_table where int_col > 5"},
		[]sqlparser.ColIdent{
			sqlparser.NewColIdent("id"),
			sqlparser.NewColIdent("user_id"),
		}, /* splitColumns */
		1000,
		getTestSchema(),
	)
	if err != nil {
		t.Fatalf("NewSplitParamsGivenNumRowsPerQueryPart failed with: %v", err)
	}
	mockSQLExecuter := splitquery_testing.NewMockSQLExecuter(mockCtrl)
	expectedCall1 := mockSQLExecuter.EXPECT().SQLExecute(
		"select id, user_id from test_table force index (`PRIMARY`)"+
			" order by id asc, user_id asc"+
			" limit 1000, 1",
		map[string]interface{}{})
	expectedCall1.Return(
		&sqltypes.Result{Rows: [][]sqltypes.Value{}}, nil)

	algorithm, err := NewFullScanAlgorithm(splitParams, mockSQLExecuter)
	if err != nil {
		t.Fatalf("NewFullScanAlgorithm failed with: %v", err)
	}
	boundaries, err := algorithm.generateBoundaries()
	if err != nil {
		t.Fatalf("FullScanAlgorithm.generateBoundaries() failed with: %v", err)
	}
	expectedBoundaries := []tuple{}
	if !reflect.DeepEqual(expectedBoundaries, boundaries) {
		t.Fatalf("expected: %v, got: %v", expectedBoundaries, boundaries)
	}
}

func TestSQLExecuterReturnsError(t *testing.T) {
	mockCtrl := gomock.NewController(t)
	defer mockCtrl.Finish()
	splitParams, err := NewSplitParamsGivenNumRowsPerQueryPart(
		querytypes.BoundQuery{Sql: "select * from test_table where int_col > 5"},
		[]sqlparser.ColIdent{
			sqlparser.NewColIdent("id"),
			sqlparser.NewColIdent("user_id"),
		}, /* splitColumns */
		1000,
		getTestSchema(),
	)
	if err != nil {
		t.Fatalf("NewSplitParamsGivenNumRowsPerQueryPart failed with: %v", err)
	}
	mockSQLExecuter := splitquery_testing.NewMockSQLExecuter(mockCtrl)
	expectedCall1 := mockSQLExecuter.EXPECT().SQLExecute(
		"select id, user_id from test_table force index (`PRIMARY`)"+
			" order by id asc, user_id asc"+
			" limit 1000, 1",
		map[string]interface{}{})
	expectedCall1.Return(
		&sqltypes.Result{
			Rows: [][]sqltypes.Value{
				{int64Value(1), int64Value(1)}},
		},
		nil)
	expectedCall2 := mockSQLExecuter.EXPECT().SQLExecute(
		"select id, user_id from test_table force index (`PRIMARY`)"+
			" where"+
			" :_splitquery_prev_id < id or"+
			" (:_splitquery_prev_id = id and :_splitquery_prev_user_id <= user_id)"+
			" order by id asc, user_id asc"+
			" limit 1000, 1",
		map[string]interface{}{
			"_splitquery_prev_id":      int64(1),
			"_splitquery_prev_user_id": int64(1),
		})
	expectedErr := fmt.Errorf("Error accessing database!")
	expectedCall2.Return(nil, expectedErr)
	algorithm, err := NewFullScanAlgorithm(splitParams, mockSQLExecuter)
	if err != nil {
		t.Fatalf("NewFullScanAlgorithm failed with: %v", err)
	}
	boundaries, err := algorithm.generateBoundaries()
	if err != expectedErr {
		t.Fatalf("FullScanAlgorithm.generateBoundaries() did not fail as expected. err: %v", err)
	}
	if boundaries != nil {
		t.Fatalf("boundaries: %v, expected: nil", boundaries)
	}
}
<|endoftext|>


<commit_before>package zenodb

import (
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"sort"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/cloudfoundry/gosigar"
	"github.com/dustin/go-humanize"
	"github.com/getlantern/goexpr/geo"
	"github.com/getlantern/goexpr/isp"
	geredis "github.com/getlantern/goexpr/redis"
	"github.com/getlantern/golog"
	"github.com/getlantern/vtime"
	"github.com/getlantern/wal"
	"github.com/getlantern/zenodb/core"
	"github.com/getlantern/zenodb/planner"
	"github.com/getlantern/zenodb/sql"
	"github.com/rickar/props"
	"gopkg.in/redis.v3"
)

var (
	log = golog.LoggerFor("zenodb")

	systemRAM float64
)

func init() {
	mem := sigar.Mem{}
	err := mem.Get()
	if err != nil {
		panic(fmt.Sprintf("Unable to get system memory: %v", err))
	}
	systemRAM = float64(mem.Total)
}

// DBOpts provides options for configuring the database.
type DBOpts struct {
	// Dir points at the directory that contains the data files.
	Dir string
	// SchemaFile points at a YAML schema file that configures the tables and
	// views in the database.
	SchemaFile string
	// AliasesFile points at a file that contains expression aliases in the form
	// name=template(%v, %v), with one alias per line.
	AliasesFile string
	// EnableGeo enables geolocation functions
	EnableGeo bool
	// ISPProvider configures a provider of ISP lookups. Specify this to allow the
	// use of ISP functions.
	ISPProvider isp.Provider
	// RedisClient provides a connection to redis which enables the use of Redis
	// expressions like HGET.
	RedisClient *redis.Client
	// RedisCacheSize controls the size of redis hash caches
	RedisCacheSize int
	// VirtualTime, if true, tells zenodb to use a virtual clock that advances
	// based on the timestamps of Points received via inserts.
	VirtualTime bool
	// WALSyncInterval governs how frequently to sync the WAL to disk. 0 means
	// it syncs after every write (which is not great for performance).
	WALSyncInterval time.Duration
	// MaxWALAge limits how far back we keep WAL files.
	MaxWALAge time.Duration
	// WALCompressionAge sets a cutoff for the age of WAL files that will be
	// gzipped
	WALCompressionAge time.Duration
	// MaxMemoryRatio caps the maximum memory of this process. When the system
	// comes under memory pressure, it will start flushing table memstores.
	MaxMemoryRatio float64
	// Passthrough flags this node as a passthrough (won't store data in tables,
	// just WAL). Passthrough nodes will also outsource queries to specific
	// partition handlers. Requires that NumPartitions be specified.
	Passthrough bool
	// PartitionBy identifies the dimensions by which to partition, in order of
	// priority. If a datum includes none of these dimensions, we partition on the
	// entire key.
	PartitionBy []string
	// NumPartitions identifies how many partitions to split data from
	// passthrough nodes.
	NumPartitions int
	// Partition identies the partition owned by this follower
	Partition int
	// Follow is a function that allows a follower to request following a stream
	// from a passthrough node.
	Follow                     func(f *Follow, cb func(data []byte, newOffset wal.Offset) error)
	RegisterRemoteQueryHandler func(partition int, query planner.QueryClusterFN)
}

// DB is a zenodb database.
type DB struct {
	opts                *DBOpts
	clock               vtime.Clock
	streams             map[string]*wal.WAL
	tables              map[string]*table
	orderedTables       []*table
	tablesMutex         sync.RWMutex
	isSorting           bool
	nextTableToSort     int
	memory              uint64
	flushMutex          sync.Mutex
	remoteQueryHandlers map[int]chan planner.QueryClusterFN
}

// NewDB creates a database using the given options.
func NewDB(opts *DBOpts) (*DB, error) {
	var err error
	db := &DB{
		opts:                opts,
		clock:               vtime.RealClock,
		tables:              make(map[string]*table),
		streams:             make(map[string]*wal.WAL),
		remoteQueryHandlers: make(map[int]chan planner.QueryClusterFN),
	}
	if opts.VirtualTime {
		db.clock = vtime.NewVirtualClock(time.Time{})
	}
	if opts.MaxWALAge <= 0 {
		opts.MaxWALAge = 24 * time.Hour
	}
	if opts.WALCompressionAge <= 0 {
		opts.WALCompressionAge = opts.MaxWALAge / 10
	}

	// Create db dir
	err = os.MkdirAll(opts.Dir, 0755)
	if err != nil && !os.IsExist(err) {
		return nil, fmt.Errorf("Unable to create db dir at %v: %v", opts.Dir, err)
	}

	if opts.EnableGeo {
		log.Debug("Enabling geolocation functions")
		err = geo.Init(filepath.Join(opts.Dir, "geoip.dat"))
		if err != nil {
			return nil, fmt.Errorf("Unable to initialize geo: %v", err)
		}
	}

	if opts.ISPProvider != nil {
		log.Debugf("Setting ISP provider to %v", opts.ISPProvider)
		isp.SetProvider(opts.ISPProvider)
	}

	if opts.AliasesFile != "" {
		registerAliases(opts.AliasesFile)
	}

	if opts.RedisClient != nil {
		log.Debug("Enabling redis expressions")
		geredis.Configure(opts.RedisClient, opts.RedisCacheSize)
	}

	if opts.SchemaFile != "" {
		err = db.pollForSchema(opts.SchemaFile)
		if err != nil {
			return nil, fmt.Errorf("Unable to apply schema: %v", err)
		}
	}
	log.Debugf("Dir: %v    SchemaFile: %v", opts.Dir, opts.SchemaFile)

	if db.opts.RegisterRemoteQueryHandler != nil {
		go db.opts.RegisterRemoteQueryHandler(db.opts.Partition, db.queryForRemote)
	}

	if db.opts.Passthrough {
		// go db.freshenRemoteQueryHandlers()
		log.Debugf("Partitioning by: %v", strings.Join(db.opts.PartitionBy, ","))
	}

	go db.trackMemStats()

	return db, err
}

func registerAliases(aliasesFile string) {
	log.Debugf("Registering aliases from file at %v", aliasesFile)

	file, err := os.Open(aliasesFile)
	if err != nil {
		log.Errorf("Unable to open aliases file at %v: %v", aliasesFile, err)
		return
	}
	defer file.Close()

	p, err := props.Read(file)
	if err != nil {
		log.Errorf("Unable to read aliases file at %v: %v", aliasesFile, err)
		return
	}

	for _, alias := range p.Names() {
		template := strings.TrimSpace(p.Get(alias))
		alias = strings.TrimSpace(alias)
		sql.RegisterAlias(alias, template)
		log.Debugf("Registered alias %v = %v", alias, template)
	}
}

// TableStats returns the TableStats for the named table.
func (db *DB) TableStats(table string) TableStats {
	t := db.getTable(table)
	if t == nil {
		return TableStats{}
	}
	t.statsMutex.RLock()
	defer t.statsMutex.RUnlock()
	return t.stats
}

// AllTableStats returns all TableStats for all tables, keyed to the table
// names.
func (db *DB) AllTableStats() map[string]TableStats {
	m := make(map[string]TableStats)
	tables := make(map[string]*table, 0)
	db.tablesMutex.RLock()
	for name, t := range db.tables {
		tables[name] = t
	}
	db.tablesMutex.RUnlock()
	for name, t := range tables {
		t.statsMutex.RLock()
		m[name] = t.stats
		t.statsMutex.RUnlock()
	}
	return m
}

// PrintTableStats prints the stats for the named table to a string.
func (db *DB) PrintTableStats(table string) string {
	stats := db.TableStats(table)
	now := db.clock.Now()
	return fmt.Sprintf("%v (%v)\tFiltered: %v    Queued: %v    Inserted: %v    Dropped: %v    Expired: %v",
		table,
		now.In(time.UTC),
		humanize.Comma(stats.FilteredPoints),
		humanize.Comma(stats.QueuedPoints),
		humanize.Comma(stats.InsertedPoints),
		humanize.Comma(stats.DroppedPoints),
		humanize.Comma(stats.ExpiredValues))
}

func (db *DB) getTable(table string) *table {
	db.tablesMutex.RLock()
	t := db.tables[strings.ToLower(table)]
	db.tablesMutex.RUnlock()
	return t
}

func (db *DB) now(table string) time.Time {
	return db.clock.Now()
}

func (db *DB) getFields(table string) (core.Fields, error) {
	t := db.getTable(table)
	if t == nil {
		return nil, fmt.Errorf("Table '%v' not found", table)
	}
	return t.Fields, nil
}

func (db *DB) getFieldsOptional(table string) (core.Fields, error) {
	t := db.getTable(table)
	if t == nil {
		return nil, nil
	}
	return t.Fields, nil
}

func (db *DB) capWALAge(wal *wal.WAL) {
	for {
		time.Sleep(1 * time.Minute)
		err := wal.TruncateBeforeTime(time.Now().Add(-1 * db.opts.MaxWALAge))
		if err != nil {
			log.Errorf("Error truncating WAL: %v", err)
		}
		err = wal.CompressBeforeTime(time.Now().Add(-1 * db.opts.WALCompressionAge))
		if err != nil {
			log.Errorf("Error compressing WAL: %v", err)
		}
	}
}

func (db *DB) trackMemStats() {
	for {
		db.updateMemStats()
		time.Sleep(2 * time.Second)
	}
}

func (db *DB) updateMemStats() {
	memstats := &runtime.MemStats{}
	runtime.ReadMemStats(memstats)
	atomic.StoreUint64(&db.memory, memstats.HeapInuse)
}

func (db *DB) capMemStoreSize() {
	if db.opts.MaxMemoryRatio <= 0 {
		return
	}

	db.tablesMutex.RLock()
	sizes := make(byCurrentSize, 0, len(db.tables))
	for _, table := range db.tables {
		sizes = append(sizes, &memStoreSize{table, table.memStoreSize()})
	}
	db.tablesMutex.RUnlock()

	db.flushMutex.Lock()
	actual := atomic.LoadUint64(&db.memory)
	allowed := db.maxMemoryBytes()
	if actual > allowed {
		// Force flushing on the table with the largest memstore
		sort.Sort(sizes)
		log.Debugf("Memory usage of %v exceeds allowed %v, forcing flush on %v", humanize.Bytes(actual), humanize.Bytes(allowed), sizes[0].t.Name)
		sizes[0].t.forceFlush()
		db.updateMemStats()
		log.Debugf("Done forcing flush on %v", sizes[0].t.Name)
	}
	db.flushMutex.Unlock()
}

func (db *DB) maxMemoryBytes() uint64 {
	return uint64(systemRAM * db.opts.MaxMemoryRatio)
}

type memStoreSize struct {
	t    *table
	size int
}

type byCurrentSize []*memStoreSize

func (a byCurrentSize) Len() int           { return len(a) }
func (a byCurrentSize) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a byCurrentSize) Less(i, j int) bool { return a[i].size > a[j].size }
<commit_msg>Capping memory usage by percentage rather than absolute amount<commit_after>package zenodb

import (
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"sort"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/cloudfoundry/gosigar"
	"github.com/dustin/go-humanize"
	"github.com/getlantern/goexpr/geo"
	"github.com/getlantern/goexpr/isp"
	geredis "github.com/getlantern/goexpr/redis"
	"github.com/getlantern/golog"
	"github.com/getlantern/vtime"
	"github.com/getlantern/wal"
	"github.com/getlantern/zenodb/core"
	"github.com/getlantern/zenodb/planner"
	"github.com/getlantern/zenodb/sql"
	"github.com/rickar/props"
	"gopkg.in/redis.v3"
)

var (
	log = golog.LoggerFor("zenodb")

	systemRAM float64
)

func init() {
	mem := sigar.Mem{}
	err := mem.Get()
	if err != nil {
		panic(fmt.Sprintf("Unable to get system memory: %v", err))
	}
	systemRAM = float64(mem.Total)
}

// DBOpts provides options for configuring the database.
type DBOpts struct {
	// Dir points at the directory that contains the data files.
	Dir string
	// SchemaFile points at a YAML schema file that configures the tables and
	// views in the database.
	SchemaFile string
	// AliasesFile points at a file that contains expression aliases in the form
	// name=template(%v, %v), with one alias per line.
	AliasesFile string
	// EnableGeo enables geolocation functions
	EnableGeo bool
	// ISPProvider configures a provider of ISP lookups. Specify this to allow the
	// use of ISP functions.
	ISPProvider isp.Provider
	// RedisClient provides a connection to redis which enables the use of Redis
	// expressions like HGET.
	RedisClient *redis.Client
	// RedisCacheSize controls the size of redis hash caches
	RedisCacheSize int
	// VirtualTime, if true, tells zenodb to use a virtual clock that advances
	// based on the timestamps of Points received via inserts.
	VirtualTime bool
	// WALSyncInterval governs how frequently to sync the WAL to disk. 0 means
	// it syncs after every write (which is not great for performance).
	WALSyncInterval time.Duration
	// MaxWALAge limits how far back we keep WAL files.
	MaxWALAge time.Duration
	// WALCompressionAge sets a cutoff for the age of WAL files that will be
	// gzipped
	WALCompressionAge time.Duration
	// MaxMemoryRatio caps the maximum memory of this process. When the system
	// comes under memory pressure, it will start flushing table memstores.
	MaxMemoryRatio float64
	// Passthrough flags this node as a passthrough (won't store data in tables,
	// just WAL). Passthrough nodes will also outsource queries to specific
	// partition handlers. Requires that NumPartitions be specified.
	Passthrough bool
	// PartitionBy identifies the dimensions by which to partition, in order of
	// priority. If a datum includes none of these dimensions, we partition on the
	// entire key.
	PartitionBy []string
	// NumPartitions identifies how many partitions to split data from
	// passthrough nodes.
	NumPartitions int
	// Partition identies the partition owned by this follower
	Partition int
	// Follow is a function that allows a follower to request following a stream
	// from a passthrough node.
	Follow                     func(f *Follow, cb func(data []byte, newOffset wal.Offset) error)
	RegisterRemoteQueryHandler func(partition int, query planner.QueryClusterFN)
}

// DB is a zenodb database.
type DB struct {
	opts                *DBOpts
	clock               vtime.Clock
	streams             map[string]*wal.WAL
	tables              map[string]*table
	orderedTables       []*table
	tablesMutex         sync.RWMutex
	isSorting           bool
	nextTableToSort     int
	memory              uint64
	flushMutex          sync.Mutex
	remoteQueryHandlers map[int]chan planner.QueryClusterFN
}

// NewDB creates a database using the given options.
func NewDB(opts *DBOpts) (*DB, error) {
	var err error
	db := &DB{
		opts:                opts,
		clock:               vtime.RealClock,
		tables:              make(map[string]*table),
		streams:             make(map[string]*wal.WAL),
		remoteQueryHandlers: make(map[int]chan planner.QueryClusterFN),
	}
	if opts.VirtualTime {
		db.clock = vtime.NewVirtualClock(time.Time{})
	}
	if opts.MaxWALAge <= 0 {
		opts.MaxWALAge = 24 * time.Hour
	}
	if opts.WALCompressionAge <= 0 {
		opts.WALCompressionAge = opts.MaxWALAge / 10
	}

	// Create db dir
	err = os.MkdirAll(opts.Dir, 0755)
	if err != nil && !os.IsExist(err) {
		return nil, fmt.Errorf("Unable to create db dir at %v: %v", opts.Dir, err)
	}

	if opts.EnableGeo {
		log.Debug("Enabling geolocation functions")
		err = geo.Init(filepath.Join(opts.Dir, "geoip.dat"))
		if err != nil {
			return nil, fmt.Errorf("Unable to initialize geo: %v", err)
		}
	}

	if opts.ISPProvider != nil {
		log.Debugf("Setting ISP provider to %v", opts.ISPProvider)
		isp.SetProvider(opts.ISPProvider)
	}

	if opts.AliasesFile != "" {
		registerAliases(opts.AliasesFile)
	}

	if opts.RedisClient != nil && opts.RedisCacheSize > 0 {
		log.Debug("Enabling redis expressions")
		geredis.Configure(opts.RedisClient, opts.RedisCacheSize)
	}

	if opts.SchemaFile != "" {
		err = db.pollForSchema(opts.SchemaFile)
		if err != nil {
			return nil, fmt.Errorf("Unable to apply schema: %v", err)
		}
	}
	log.Debugf("Dir: %v    SchemaFile: %v", opts.Dir, opts.SchemaFile)

	if db.opts.RegisterRemoteQueryHandler != nil {
		go db.opts.RegisterRemoteQueryHandler(db.opts.Partition, db.queryForRemote)
	}

	if db.opts.Passthrough {
		// go db.freshenRemoteQueryHandlers()
		log.Debugf("Partitioning by: %v", strings.Join(db.opts.PartitionBy, ","))
	}

	if db.opts.MaxMemoryRatio > 0 {
		log.Debugf("Limiting maximum memstore memory to %v", humanize.Bytes(db.maxMemoryBytes()))
	}
	go db.trackMemStats()

	return db, err
}

func registerAliases(aliasesFile string) {
	log.Debugf("Registering aliases from file at %v", aliasesFile)

	file, err := os.Open(aliasesFile)
	if err != nil {
		log.Errorf("Unable to open aliases file at %v: %v", aliasesFile, err)
		return
	}
	defer file.Close()

	p, err := props.Read(file)
	if err != nil {
		log.Errorf("Unable to read aliases file at %v: %v", aliasesFile, err)
		return
	}

	for _, alias := range p.Names() {
		template := strings.TrimSpace(p.Get(alias))
		alias = strings.TrimSpace(alias)
		sql.RegisterAlias(alias, template)
		log.Debugf("Registered alias %v = %v", alias, template)
	}
}

// TableStats returns the TableStats for the named table.
func (db *DB) TableStats(table string) TableStats {
	t := db.getTable(table)
	if t == nil {
		return TableStats{}
	}
	t.statsMutex.RLock()
	defer t.statsMutex.RUnlock()
	return t.stats
}

// AllTableStats returns all TableStats for all tables, keyed to the table
// names.
func (db *DB) AllTableStats() map[string]TableStats {
	m := make(map[string]TableStats)
	tables := make(map[string]*table, 0)
	db.tablesMutex.RLock()
	for name, t := range db.tables {
		tables[name] = t
	}
	db.tablesMutex.RUnlock()
	for name, t := range tables {
		t.statsMutex.RLock()
		m[name] = t.stats
		t.statsMutex.RUnlock()
	}
	return m
}

// PrintTableStats prints the stats for the named table to a string.
func (db *DB) PrintTableStats(table string) string {
	stats := db.TableStats(table)
	now := db.clock.Now()
	return fmt.Sprintf("%v (%v)\tFiltered: %v    Queued: %v    Inserted: %v    Dropped: %v    Expired: %v",
		table,
		now.In(time.UTC),
		humanize.Comma(stats.FilteredPoints),
		humanize.Comma(stats.QueuedPoints),
		humanize.Comma(stats.InsertedPoints),
		humanize.Comma(stats.DroppedPoints),
		humanize.Comma(stats.ExpiredValues))
}

func (db *DB) getTable(table string) *table {
	db.tablesMutex.RLock()
	t := db.tables[strings.ToLower(table)]
	db.tablesMutex.RUnlock()
	return t
}

func (db *DB) now(table string) time.Time {
	return db.clock.Now()
}

func (db *DB) getFields(table string) (core.Fields, error) {
	t := db.getTable(table)
	if t == nil {
		return nil, fmt.Errorf("Table '%v' not found", table)
	}
	return t.Fields, nil
}

func (db *DB) getFieldsOptional(table string) (core.Fields, error) {
	t := db.getTable(table)
	if t == nil {
		return nil, nil
	}
	return t.Fields, nil
}

func (db *DB) capWALAge(wal *wal.WAL) {
	for {
		time.Sleep(1 * time.Minute)
		err := wal.TruncateBeforeTime(time.Now().Add(-1 * db.opts.MaxWALAge))
		if err != nil {
			log.Errorf("Error truncating WAL: %v", err)
		}
		err = wal.CompressBeforeTime(time.Now().Add(-1 * db.opts.WALCompressionAge))
		if err != nil {
			log.Errorf("Error compressing WAL: %v", err)
		}
	}
}

func (db *DB) trackMemStats() {
	for {
		db.updateMemStats()
		time.Sleep(2 * time.Second)
	}
}

func (db *DB) updateMemStats() {
	memstats := &runtime.MemStats{}
	runtime.ReadMemStats(memstats)
	atomic.StoreUint64(&db.memory, memstats.HeapInuse)
}

func (db *DB) capMemStoreSize() {
	if db.opts.MaxMemoryRatio <= 0 {
		return
	}

	db.tablesMutex.RLock()
	sizes := make(byCurrentSize, 0, len(db.tables))
	for _, table := range db.tables {
		sizes = append(sizes, &memStoreSize{table, table.memStoreSize()})
	}
	db.tablesMutex.RUnlock()

	db.flushMutex.Lock()
	actual := atomic.LoadUint64(&db.memory)
	allowed := db.maxMemoryBytes()
	if actual > allowed {
		// Force flushing on the table with the largest memstore
		sort.Sort(sizes)
		log.Debugf("Memory usage of %v exceeds allowed %v, forcing flush on %v", humanize.Bytes(actual), humanize.Bytes(allowed), sizes[0].t.Name)
		sizes[0].t.forceFlush()
		db.updateMemStats()
		log.Debugf("Done forcing flush on %v", sizes[0].t.Name)
	}
	db.flushMutex.Unlock()
}

func (db *DB) maxMemoryBytes() uint64 {
	return uint64(systemRAM * db.opts.MaxMemoryRatio)
}

type memStoreSize struct {
	t    *table
	size int
}

type byCurrentSize []*memStoreSize

func (a byCurrentSize) Len() int           { return len(a) }
func (a byCurrentSize) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a byCurrentSize) Less(i, j int) bool { return a[i].size > a[j].size }
<|endoftext|>


<commit_before>package flv

import (
	"os"
	"fmt"
	"bytes"
)

type Header struct {
	Version      uint16
}

type Frame struct {
	Stream      uint32
	Dts         uint32
	Type        TagType
	Flavor      Flavor
	Position    int64
	Body        []byte
}

func ReadHeader(inFile *os.File) (*Header, error) {
	header := make([]byte, HEADER_LENGTH)
	_, err := inFile.Read(header)
	if err != nil {
		return nil, err
	}

	sig := header[0:3]
	if bytes.Compare(sig, []byte(SIG)) != 0 {
		return nil, fmt.Errorf("bad file format")
	}
	version := (uint16(header[3]) << 8) | (uint16(header[4]) << 0)
	//skip := header[4:5]
	//offset := header[5:9]

	next_id := make([]byte, 4)
	_, err = inFile.Read(next_id)
	return &Header{Version: version}, nil
}

func ReadTag(inFile *os.File) (fr *Frame, e error) {

	var n int
	var err error

	curPos, _ := inFile.Seek(0, os.SEEK_CUR)

	tagHeaderB := make([]byte, TAG_HEADER_LENGTH)
	n, err = inFile.Read(tagHeaderB)
	if TagSize(n) != TAG_HEADER_LENGTH {
		return nil, fmt.Errorf("bad record")
	}
	if err != nil {
		return nil, err
	}

	tagType := TagType(tagHeaderB[0])
	bodyLen := (uint32(tagHeaderB[1]) << 16) | (uint32(tagHeaderB[2]) << 8) | (uint32(tagHeaderB[3]) << 0)
	ts := (uint32(tagHeaderB[4]) << 16) | (uint32(tagHeaderB[5]) << 8) | (uint32(tagHeaderB[6]) << 0)
	tsExt := uint32(tagHeaderB[7])
	stream := (uint32(tagHeaderB[8]) << 16) | (uint32(tagHeaderB[9]) << 8) | (uint32(tagHeaderB[10]) << 0)

	var dts uint32
	dts = (tsExt << 24) | ts

	bodyBuf := make([]byte, bodyLen)
	n, err = inFile.Read(bodyBuf)
	if err != nil {
		return nil, err
	}

	var flavor Flavor

	switch tagType {
	case TAG_TYPE_META:
		flavor = METADATA
	case TAG_TYPE_VIDEO:
		vft := VideoFrameType(uint(bodyBuf[0]) >> 4)
		switch vft {
		case VIDEO_FRAME_TYPE_KEYFRAME:
			flavor = KEYFRAME
		default:
			flavor = FRAME
		}
	case TAG_TYPE_AUDIO:
		flavor = FRAME
	}

	prevTagSizeB := make([]byte, 4)
	n, err = inFile.Read(prevTagSizeB)
	if err != nil {
		return nil, err
	}
	//prevTagSize := (uint32(prevTagSizeB[0]) << 24) | (uint32(prevTagSizeB[1]) << 16) | (uint32(prevTagSizeB[2]) << 8) | (uint32(prevTagSizeB[3]) << 0)

	return &Frame{
		Stream: stream,
		Dts: dts,
		Type: tagType,
		Flavor: flavor,
		Position: curPos,
		Body: bodyBuf,
	}, nil
}
<commit_msg> * Use reader  * Parse video width and height<commit_after>package flv

import (
	"os"
	"fmt"
	"bytes"
)

type Header struct {
	Version      uint16
}

type Frame interface{
	ToBinary() ([]byte, error)
}

type CFrame struct {
	Stream      uint32
	Dts         uint32
	Type        TagType
	Flavor      Flavor
	Position    int64
	Body        []byte
	PrevTagSize uint32
}

func (f CFrame) ToBinary() ([]byte, error) {
	d := make([]byte, 0)
	buf := bytes.NewBuffer(d)
	buf.WriteByte(byte(f.Type))
	bl := uint32(len(f.Body))
	buf.Write([]byte{byte(bl>>16), byte((bl>>8)&0xFF), byte(bl&0xFF)})
	buf.Write([]byte{byte((f.Dts>>16) & 0xFF), byte((f.Dts>>8) & 0xFF), byte(f.Dts & 0xFF)})
	buf.WriteByte(byte((f.Dts >> 24) & 0xFF))
	buf.Write(f.Body)
	buf.Write([]byte{byte((f.PrevTagSize>>24) & 0xFF), byte((f.PrevTagSize>>16) & 0xFF), byte((f.PrevTagSize>>8) & 0xFF), byte(f.PrevTagSize & 0xFF)})
	return buf.Bytes(), nil
}

type VideoFrame struct {
	CFrame
	CodecId     VideoCodec
	Width       uint16
	Height      uint16
}

func (f VideoFrame) ToBinary() ([]byte, error) {
	return f.ToBinary()
}

type AudioFrame struct {
	CFrame
	CodecId     AudioCodec
    Rate        AudioRate
    BitSize     AudioSize
	Channels    AudioType
}

func (f AudioFrame) ToBinary() ([]byte, error) {
	return f.ToBinary()
}

type MetaFrame struct {
	CFrame
}

func (f MetaFrame) ToBinary() ([]byte, error) {
	return f.ToBinary()
}

type FlvReader struct {
	inFile     *os.File
	width      uint16
	height     uint16
}

func NewReader(inFile *os.File) (*FlvReader) {
	return &FlvReader{
		inFile: inFile,
		width: 0,
		height: 0,
	}
}

func (frReader *FlvReader) ReadHeader() (*Header, error) {
	header := make([]byte, HEADER_LENGTH)
	_, err := frReader.inFile.Read(header)
	if err != nil {
		return nil, err
	}

	sig := header[0:3]
	if bytes.Compare(sig, []byte(SIG)) != 0 {
		return nil, fmt.Errorf("bad file format")
	}
	version := (uint16(header[3]) << 8) | (uint16(header[4]) << 0)
	//skip := header[4:5]
	//offset := header[5:9]

	next_id := make([]byte, 4)
	_, err = frReader.inFile.Read(next_id)
	return &Header{Version: version}, nil
}

func (frReader *FlvReader) ReadFrame() (fr Frame, e error) {

	var n int
	var err error

	curPos, _ := frReader.inFile.Seek(0, os.SEEK_CUR)

	tagHeaderB := make([]byte, TAG_HEADER_LENGTH)
	n, err = frReader.inFile.Read(tagHeaderB)
	if TagSize(n) != TAG_HEADER_LENGTH {
		return nil, fmt.Errorf("bad record")
	}
	if err != nil {
		return nil, err
	}

	tagType := TagType(tagHeaderB[0])
	bodyLen := (uint32(tagHeaderB[1]) << 16) | (uint32(tagHeaderB[2]) << 8) | (uint32(tagHeaderB[3]) << 0)
	ts := (uint32(tagHeaderB[4]) << 16) | (uint32(tagHeaderB[5]) << 8) | (uint32(tagHeaderB[6]) << 0)
	tsExt := uint32(tagHeaderB[7])
	stream := (uint32(tagHeaderB[8]) << 16) | (uint32(tagHeaderB[9]) << 8) | (uint32(tagHeaderB[10]) << 0)

	var dts uint32
	dts = (tsExt << 24) | ts

	bodyBuf := make([]byte, bodyLen)
	n, err = frReader.inFile.Read(bodyBuf)
	if err != nil {
		return nil, err
	}

	prevTagSizeB := make([]byte, 4)
	n, err = frReader.inFile.Read(prevTagSizeB)
	if err != nil {
		return nil, err
	}
	prevTagSize := (uint32(prevTagSizeB[0]) << 24) | (uint32(prevTagSizeB[1]) << 16) | (uint32(prevTagSizeB[2]) << 8) | (uint32(prevTagSizeB[3]) << 0)

	pFrame := CFrame{
		Stream: stream,
		Dts: dts,
		Type: tagType,
		Position: curPos,
		Body: bodyBuf,
		PrevTagSize: prevTagSize,
	}

	var resFrame Frame

	switch tagType {
	case TAG_TYPE_META:
		pFrame.Flavor = METADATA
		resFrame = MetaFrame{pFrame}
	case TAG_TYPE_VIDEO:
		vft := VideoFrameType(uint8(bodyBuf[0]) >> 4)
		codecId := VideoCodec(uint8(bodyBuf[0]) & 0x0F)
		switch vft {
		case VIDEO_FRAME_TYPE_KEYFRAME:
			pFrame.Flavor = KEYFRAME
			switch codecId {
			case VIDEO_CODEC_ON2VP6:
				hHelper := (uint16(bodyBuf[1]) >> 4) & 0x0F
				wHelper := uint16(bodyBuf[1]) & 0x0F
				w := uint16(bodyBuf[5])
				h := uint16(bodyBuf[6])

				frReader.width = w*16 - wHelper
				frReader.height = h*16 - hHelper
			}
		default:
			pFrame.Flavor = FRAME
		}
		resFrame = VideoFrame{CFrame: pFrame, CodecId: codecId, Width: frReader.width, Height: frReader.height}
	case TAG_TYPE_AUDIO:
		pFrame.Flavor = FRAME
		codecId := AudioCodec(uint8(bodyBuf[0]) >> 4)
		rate := AudioRate((uint8(bodyBuf[0]) >> 2) & 0x03)
		bitSize := AudioSize((uint8(bodyBuf[0]) >> 1) & 0x01)
		channels := AudioType(uint8(bodyBuf[0]) & 0x01)
		resFrame = AudioFrame{CFrame: pFrame, CodecId: codecId, Rate: rate, BitSize: bitSize, Channels: channels}
	}

	return resFrame, nil
}
<|endoftext|>


<commit_before>package api

import (
	"encoding/json"
	"net/http"
	"pkg/ctxutil"
	"pkg/echologger"

	"golang.org/x/net/context"

	"autoscale"

	"github.com/labstack/echo"
	"github.com/labstack/echo/engine"
	"github.com/labstack/echo/engine/standard"
	"github.com/labstack/echo/middleware"
	"github.com/satori/go.uuid"
)

type errorMsg struct {
	Title string `json:"title"`
}

func writeError(w http.ResponseWriter, msg string, code int) {
	em := errorMsg{
		Title: msg,
	}

	w.WriteHeader(code)
	_ = json.NewEncoder(w).Encode(&em)
}

// API is the autoscale API.
type API struct {
	Mux  http.Handler
	repo autoscale.Repository
	ctx  context.Context
}

// New creates an instance of API.
func New(ctx context.Context, repo autoscale.Repository) *API {
	e := echo.New()

	std := standard.WithConfig(engine.Config{})
	std.SetHandler(e)

	a := &API{
		Mux:  std,
		repo: repo,
		ctx:  ctx,
	}

	e.Use(func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			reqID := c.Request().Header().Get("X-Request-Id")
			if reqID == "" {
				reqID = uuid.NewV4().String()
				c.Request().Header().Set("X-Request-Id", reqID)

			}

			newCtx := context.WithValue(c, "RequestID", reqID)
			c.SetNetContext(newCtx)

			return next(c)
		}
	})

	log := ctxutil.LogFromContext(ctx)
	logmw := echologger.NewWithNameAndLogger("autoscale", log)
	e.Use(logmw)

	e.Use(middleware.RecoverWithConfig(middleware.RecoverConfig{
		StackSize: 1 << 10, // 1 KB
	}))

	e.GET("/templates/:id", a.getTemplate)
	e.GET("/templates", a.listTemplates)
	e.POST("/templates", a.createTemplate)
	e.DELETE("/templates/:id", a.deleteTemplate)
	e.GET("/groups", a.listGroups)
	e.GET("/groups/:id", a.getGroup)
	e.POST("/groups", a.createGroup)
	e.DELETE("/groups/:id", a.deleteGroup)
	e.PUT("/groups/:id", a.updateGroup)

	return a
}

func (a *API) listTemplates(c echo.Context) error {
	log := ctxutil.LogFromContext(a.ctx)
	tmpls, err := a.repo.ListTemplates(a.ctx)
	if err != nil {
		log.WithError(err).Error("list templates")

		return echo.NewHTTPError(http.StatusInternalServerError)
	}

	return c.JSON(http.StatusOK, tmpls)
}

func (a *API) getTemplate(c echo.Context) error {
	log := ctxutil.LogFromContext(a.ctx)

	id := c.Param("id")

	tmpl, err := a.repo.GetTemplate(a.ctx, id)
	if err != nil {
		log.WithError(err).Error("retrieve template")
		return echo.ErrNotFound
	}

	return c.JSON(http.StatusOK, tmpl)
}

func (a *API) createTemplate(c echo.Context) error {
	log := ctxutil.LogFromContext(a.ctx)

	var ctr autoscale.CreateTemplateRequest
	if err := c.Bind(&ctr); err != nil {
		return err
	}

	tmpl, err := a.repo.CreateTemplate(a.ctx, ctr)
	if err != nil {
		log.WithError(err).Error("delete template")

		return echo.NewHTTPError(http.StatusBadRequest)
	}

	return c.JSON(http.StatusCreated, tmpl)
}

func (a *API) deleteTemplate(c echo.Context) error {
	log := ctxutil.LogFromContext(a.ctx)

	id := c.Param("id")

	err := a.repo.DeleteTemplate(a.ctx, id)
	if err != nil {
		log.WithError(err).Error("delete template")
		return echo.NewHTTPError(http.StatusBadRequest)
	}

	return c.NoContent(204)
}

func (a *API) listGroups(c echo.Context) error {
	log := ctxutil.LogFromContext(a.ctx)

	groups, err := a.repo.ListGroups(a.ctx)
	if err != nil {
		log.WithError(err).Error("list groups")
		return echo.NewHTTPError(http.StatusInternalServerError)
	}

	return c.JSON(http.StatusOK, groups)
}

func (a *API) getGroup(c echo.Context) error {
	log := ctxutil.LogFromContext(a.ctx)

	id := c.Param("id")

	group, err := a.repo.GetGroup(a.ctx, id)
	if err != nil {
		log.WithError(err).Error("get group")
		return echo.ErrNotFound
	}

	return c.JSON(http.StatusOK, group)
}

func (a *API) createGroup(c echo.Context) error {
	log := ctxutil.LogFromContext(a.ctx)

	var cgr autoscale.CreateGroupRequest
	if err := c.Bind(&cgr); err != nil {
		return err
	}

	g, err := a.repo.CreateGroup(a.ctx, cgr)
	if err != nil {
		log.WithError(err).Error("create group")
		return echo.NewHTTPError(http.StatusBadRequest)
	}

	return c.JSON(http.StatusCreated, g)
}

func (a *API) deleteGroup(c echo.Context) error {
	log := ctxutil.LogFromContext(a.ctx)

	id := c.Param("id")

	err := a.repo.DeleteGroup(a.ctx, id)
	if err != nil {
		log.WithError(err).Error("delete group")
		return echo.NewHTTPError(http.StatusBadRequest)
	}

	return c.NoContent(204)
}

func (a *API) updateGroup(c echo.Context) error {
	log := ctxutil.LogFromContext(a.ctx)

	id := c.Param("id")

	var ugr autoscale.UpdateGroupRequest
	if err := c.Bind(&ugr); err != nil {
		return err
	}

	g, err := a.repo.GetGroup(a.ctx, id)
	if err != nil {
		log.WithError(err).Error("can't retrieve group")
		return echo.NewHTTPError(http.StatusBadRequest)
	}

	err = a.repo.SaveGroup(a.ctx, g)

	if err != nil {
		log.WithError(err).Error("update group")
		return echo.NewHTTPError(http.StatusBadRequest)
	}

	return c.JSON(http.StatusOK, g)
}
<commit_msg>get context from echo<commit_after>package api

import (
	"encoding/json"
	"net/http"
	"pkg/ctxutil"
	"pkg/echologger"

	"golang.org/x/net/context"

	"autoscale"

	"github.com/labstack/echo"
	"github.com/labstack/echo/engine"
	"github.com/labstack/echo/engine/standard"
	"github.com/labstack/echo/middleware"
	"github.com/satori/go.uuid"
)

type errorMsg struct {
	Title string `json:"title"`
}

func writeError(w http.ResponseWriter, msg string, code int) {
	em := errorMsg{
		Title: msg,
	}

	w.WriteHeader(code)
	_ = json.NewEncoder(w).Encode(&em)
}

// API is the autoscale API.
type API struct {
	Mux  http.Handler
	repo autoscale.Repository
	ctx  context.Context
}

// New creates an instance of API.
func New(ctx context.Context, repo autoscale.Repository) *API {
	e := echo.New()

	std := standard.WithConfig(engine.Config{})
	std.SetHandler(e)

	a := &API{
		Mux:  std,
		repo: repo,
		ctx:  ctx,
	}

	log := ctxutil.LogFromContext(ctx)

	e.Use(func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			reqID := c.Request().Header().Get("X-Request-Id")
			if reqID == "" {
				reqID = uuid.NewV4().String()
				c.Request().Header().Set("X-Request-Id", reqID)

			}

			newCtx := context.WithValue(c, "RequestID", reqID)
			newCtx = context.WithValue(newCtx, "log", log)
			c.SetNetContext(newCtx)

			return next(c)
		}
	})

	logmw := echologger.NewWithNameAndLogger("autoscale", log)
	e.Use(logmw)

	e.Use(middleware.RecoverWithConfig(middleware.RecoverConfig{
		StackSize: 1 << 10, // 1 KB
	}))

	e.GET("/templates/:id", a.getTemplate)
	e.GET("/templates", a.listTemplates)
	e.POST("/templates", a.createTemplate)
	e.DELETE("/templates/:id", a.deleteTemplate)
	e.GET("/groups", a.listGroups)
	e.GET("/groups/:id", a.getGroup)
	e.POST("/groups", a.createGroup)
	e.DELETE("/groups/:id", a.deleteGroup)
	e.PUT("/groups/:id", a.updateGroup)

	return a
}

func (a *API) listTemplates(c echo.Context) error {
	log := ctxutil.LogFromContext(c)
	tmpls, err := a.repo.ListTemplates(a.ctx)
	if err != nil {
		log.WithError(err).Error("list templates")

		return echo.NewHTTPError(http.StatusInternalServerError)
	}

	return c.JSON(http.StatusOK, tmpls)
}

func (a *API) getTemplate(c echo.Context) error {
	log := ctxutil.LogFromContext(c)

	id := c.Param("id")

	tmpl, err := a.repo.GetTemplate(a.ctx, id)
	if err != nil {
		log.WithError(err).Error("retrieve template")
		return echo.ErrNotFound
	}

	return c.JSON(http.StatusOK, tmpl)
}

func (a *API) createTemplate(c echo.Context) error {
	log := ctxutil.LogFromContext(c)

	var ctr autoscale.CreateTemplateRequest
	if err := c.Bind(&ctr); err != nil {
		return err
	}

	tmpl, err := a.repo.CreateTemplate(a.ctx, ctr)
	if err != nil {
		log.WithError(err).Error("delete template")

		return echo.NewHTTPError(http.StatusBadRequest)
	}

	return c.JSON(http.StatusCreated, tmpl)
}

func (a *API) deleteTemplate(c echo.Context) error {
	log := ctxutil.LogFromContext(c)

	id := c.Param("id")

	err := a.repo.DeleteTemplate(a.ctx, id)
	if err != nil {
		log.WithError(err).Error("delete template")
		return echo.NewHTTPError(http.StatusBadRequest)
	}

	return c.NoContent(204)
}

func (a *API) listGroups(c echo.Context) error {
	log := ctxutil.LogFromContext(c)

	groups, err := a.repo.ListGroups(a.ctx)
	if err != nil {
		log.WithError(err).Error("list groups")
		return echo.NewHTTPError(http.StatusInternalServerError)
	}

	return c.JSON(http.StatusOK, groups)
}

func (a *API) getGroup(c echo.Context) error {
	log := ctxutil.LogFromContext(c)

	id := c.Param("id")

	group, err := a.repo.GetGroup(a.ctx, id)
	if err != nil {
		log.WithError(err).Error("get group")
		return echo.ErrNotFound
	}

	return c.JSON(http.StatusOK, group)
}

func (a *API) createGroup(c echo.Context) error {
	log := ctxutil.LogFromContext(c)

	var cgr autoscale.CreateGroupRequest
	if err := c.Bind(&cgr); err != nil {
		return err
	}

	g, err := a.repo.CreateGroup(a.ctx, cgr)
	if err != nil {
		log.WithError(err).Error("create group")
		return echo.NewHTTPError(http.StatusBadRequest)
	}

	return c.JSON(http.StatusCreated, g)
}

func (a *API) deleteGroup(c echo.Context) error {
	log := ctxutil.LogFromContext(c)

	id := c.Param("id")

	err := a.repo.DeleteGroup(a.ctx, id)
	if err != nil {
		log.WithError(err).Error("delete group")
		return echo.NewHTTPError(http.StatusBadRequest)
	}

	return c.NoContent(204)
}

func (a *API) updateGroup(c echo.Context) error {
	log := ctxutil.LogFromContext(c)

	id := c.Param("id")

	var ugr autoscale.UpdateGroupRequest
	if err := c.Bind(&ugr); err != nil {
		return err
	}

	g, err := a.repo.GetGroup(a.ctx, id)
	if err != nil {
		log.WithError(err).Error("can't retrieve group")
		return echo.NewHTTPError(http.StatusBadRequest)
	}

	err = a.repo.SaveGroup(a.ctx, g)

	if err != nil {
		log.WithError(err).Error("update group")
		return echo.NewHTTPError(http.StatusBadRequest)
	}

	return c.JSON(http.StatusOK, g)
}
<|endoftext|>


<commit_before>package cfgo

import (
	"github.com/cloudflare/cloudflare-go"
)

// CloudflareClient the Client that will handle the Cloudflare api
type CloudflareClient struct {
	API *cloudflare.API
}

// NewCloudflareClient creates a new client for the cloudflare api
func NewCloudflareClient(cfKey, cfEmail string) (*CloudflareClient, error) {
	api, err := cloudflare.New(cfKey, cfEmail)
	if err != nil {
		return nil, err
	}

	return &CloudflareClient{API: api}, nil
}

// GetDNSRecord gets the Cloudflare DNS record for the given value
//
// zone is the cloudflare domain you are working on
// eg: domain.com
//
// dns is the address you want get the DNSRecord
// eg: my.domain.com
//
// The first matching DNSRecord will be returned
func (c *CloudflareClient) GetDNSRecord(zone, dns string) (*Domain, error) {
	zoneID, err := c.API.ZoneIDByName(zone)
	if err != nil {
		return &Domain{}, err
	}

	rr := cloudflare.DNSRecord{Name: dns}
	rrs, err := c.API.DNSRecords(zoneID, rr)
	if err != nil {
		return &Domain{}, err
	}
	record := rrs[0]

	return &Domain{
		RecordID: record.ID,
		ZoneID:   record.ZoneID,
		Type:     record.Type,
		Zone:     record.ZoneName,
		DNS:      record.Name,
		IP:       record.Content,
	}, nil
}

// UpdateDNSRecord updates the Cloudflare DNSRecord
func (c *CloudflareClient) UpdateDNSRecord(domain *Domain) error {
	record := cloudflare.DNSRecord{
		ID:      domain.RecordID,
		ZoneID:  domain.ZoneID,
		Content: domain.IP,
	}

	return c.API.UpdateDNSRecord(domain.ZoneID, domain.RecordID, record)
}
<commit_msg>Remove unnecessary variable<commit_after>package cfgo

import (
	"github.com/cloudflare/cloudflare-go"
)

// CloudflareClient the Client that will handle the Cloudflare api
type CloudflareClient struct {
	API *cloudflare.API
}

// NewCloudflareClient creates a new client for the cloudflare api
func NewCloudflareClient(cfKey, cfEmail string) (*CloudflareClient, error) {
	api, err := cloudflare.New(cfKey, cfEmail)
	if err != nil {
		return nil, err
	}

	return &CloudflareClient{API: api}, nil
}

// GetDNSRecord gets the Cloudflare DNS record for the given value
//
// zone is the cloudflare domain you are working on
// eg: domain.com
//
// dns is the address you want get the DNSRecord
// eg: my.domain.com
//
// The first matching DNSRecord will be returned
func (c *CloudflareClient) GetDNSRecord(zone, dns string) (*Domain, error) {
	zoneID, err := c.API.ZoneIDByName(zone)
	if err != nil {
		return &Domain{}, err
	}

	rrs, err := c.API.DNSRecords(zoneID, cloudflare.DNSRecord{Name: dns})
	if err != nil {
		return &Domain{}, err
	}
	record := rrs[0]

	return &Domain{
		RecordID: record.ID,
		ZoneID:   record.ZoneID,
		Type:     record.Type,
		Zone:     record.ZoneName,
		DNS:      record.Name,
		IP:       record.Content,
	}, nil
}

// UpdateDNSRecord updates the Cloudflare DNSRecord
func (c *CloudflareClient) UpdateDNSRecord(domain *Domain) error {
	record := cloudflare.DNSRecord{
		ID:      domain.RecordID,
		ZoneID:  domain.ZoneID,
		Content: domain.IP,
	}

	return c.API.UpdateDNSRecord(domain.ZoneID, domain.RecordID, record)
}
<|endoftext|>


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

package main

import (
	"bufio"
	"exec"
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"unicode"
	"utf8"
)

// Environment for commands.
var (
	XGC       []string // 6g -I _test -o _xtest_.6
	GC        []string // 6g -I _test _testmain.go
	GL        []string // 6l -L _test _testmain.6
	GOARCH    string
	GOROOT    string
	GORUN     string
	O         string
	args      []string // arguments passed to gotest; also passed to the binary
	fileNames []string
	env       = os.Environ()
)

// These strings are created by getTestNames.
var (
	insideFileNames  []string // list of *.go files inside the package.
	outsideFileNames []string // list of *.go files outside the package (in package foo_test).
)

var (
	files      []*File
	importPath string
)

// Flags for our own purposes. We do our own flag processing.
var (
	cFlag bool
	xFlag bool
)

// File represents a file that contains tests.
type File struct {
	name       string
	pkg        string
	file       *os.File
	astFile    *ast.File
	tests      []string // The names of the TestXXXs.
	benchmarks []string // The names of the BenchmarkXXXs.
}

func main() {
	flags()
	needMakefile()
	setEnvironment()
	getTestFileNames()
	parseFiles()
	getTestNames()
	run("gomake", "testpackage-clean")
	run("gomake", "testpackage", fmt.Sprintf("GOTESTFILES=%s", strings.Join(insideFileNames, " ")))
	if len(outsideFileNames) > 0 {
		run(append(XGC, outsideFileNames...)...)
	}
	importPath = runWithStdout("gomake", "-s", "importpath")
	writeTestmainGo()
	run(GC...)
	run(GL...)
	if !cFlag {
		runTestWithArgs("./" + O + ".out")
	}
}

// needMakefile tests that we have a Makefile in this directory.
func needMakefile() {
	if _, err := os.Stat("Makefile"); err != nil {
		Fatalf("please create a Makefile for gotest; see http://golang.org/doc/code.html for details")
	}
}

// Fatalf formats its arguments, prints the message with a final newline, and exits.
func Fatalf(s string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, "gotest: "+s+"\n", args...)
	os.Exit(2)
}

// theChar is the map from architecture to object character.
var theChar = map[string]string{
	"arm":   "5",
	"amd64": "6",
	"386":   "8",
}

// addEnv adds a name=value pair to the environment passed to subcommands.
// If the item is already in the environment, addEnv replaces the value.
func addEnv(name, value string) {
	for i := 0; i < len(env); i++ {
		if strings.HasPrefix(env[i], name+"=") {
			env[i] = name + "=" + value
			return
		}
	}
	env = append(env, name+"="+value)
}

// setEnvironment assembles the configuration for gotest and its subcommands.
func setEnvironment() {
	// Basic environment.
	GOROOT = runtime.GOROOT()
	addEnv("GOROOT", GOROOT)
	GOARCH = runtime.GOARCH
	addEnv("GOARCH", GOARCH)
	O = theChar[GOARCH]
	if O == "" {
		Fatalf("unknown architecture %s", GOARCH)
	}

	// Commands and their flags.
	gc := os.Getenv("GC")
	if gc == "" {
		gc = O + "g"
	}
	XGC = []string{gc, "-I", "_test", "-o", "_xtest_." + O}
	GC = []string{gc, "-I", "_test", "_testmain.go"}
	gl := os.Getenv("GL")
	if gl == "" {
		gl = O + "l"
	}
	GL = []string{gl, "-L", "_test", "_testmain." + O}

	// Silence make on Linux
	addEnv("MAKEFLAGS", "")
	addEnv("MAKELEVEL", "")
}

// getTestFileNames gets the set of files we're looking at.
// If gotest has no arguments, it scans for file names matching "[^.]*_test.go".
func getTestFileNames() {
	names := fileNames
	if len(names) == 0 {
		var err os.Error
		names, err = filepath.Glob("[^.]*_test.go")
		if err != nil {
			Fatalf("Glob pattern error: %s", err)
		}
		if len(names) == 0 {
			Fatalf(`no test files found: no match for "[^.]*_test.go"`)
		}
	}
	for _, n := range names {
		fd, err := os.Open(n)
		if err != nil {
			Fatalf("%s: %s", n, err)
		}
		f := &File{name: n, file: fd}
		files = append(files, f)
	}
}

// parseFiles parses the files and remembers the packages we find.
func parseFiles() {
	fileSet := token.NewFileSet()
	for _, f := range files {
		// Report declaration errors so we can abort if the files are incorrect Go.
		file, err := parser.ParseFile(fileSet, f.name, nil, parser.DeclarationErrors)
		if err != nil {
			Fatalf("parse error: %s", err)
		}
		f.astFile = file
		f.pkg = file.Name.String()
		if f.pkg == "" {
			Fatalf("cannot happen: no package name in %s", f.name)
		}
	}
}

// getTestNames extracts the names of tests and benchmarks.  They are all
// top-level functions that are not methods.
func getTestNames() {
	for _, f := range files {
		for _, d := range f.astFile.Decls {
			n, ok := d.(*ast.FuncDecl)
			if !ok {
				continue
			}
			if n.Recv != nil { // a method, not a function.
				continue
			}
			name := n.Name.String()
			if isTest(name, "Test") {
				f.tests = append(f.tests, name)
			} else if isTest(name, "Benchmark") {
				f.benchmarks = append(f.benchmarks, name)
			}
			// TODO: worth checking the signature? Probably not.
		}
		if strings.HasSuffix(f.pkg, "_test") {
			outsideFileNames = append(outsideFileNames, f.name)
		} else {
			insideFileNames = append(insideFileNames, f.name)
		}
	}
}

// isTest tells whether name looks like a test (or benchmark, according to prefix).
// It is a Test (say) if there is a character after Test that is not a lower-case letter.
// We don't want TesticularCancer.
func isTest(name, prefix string) bool {
	if !strings.HasPrefix(name, prefix) {
		return false
	}
	if len(name) == len(prefix) { // "Test" is ok
		return true
	}
	rune, _ := utf8.DecodeRuneInString(name[len(prefix):])
	return !unicode.IsLower(rune)
}

func run(args ...string) {
	doRun(args, false)
}

// runWithStdout is like run, but returns the text of standard output with the last newline dropped.
func runWithStdout(argv ...string) string {
	s := doRun(argv, true)
	if len(s) == 0 {
		Fatalf("no output from command %s", strings.Join(argv, " "))
	}
	if s[len(s)-1] == '\n' {
		s = s[:len(s)-1]
	}
	return s
}

// runTestWithArgs appends gotest's runs the provided binary with the args passed on the command line.
func runTestWithArgs(binary string) {
	doRun(append([]string{binary}, args...), false)
}

// doRun is the general command runner.  The flag says whether we want to
// retrieve standard output.
func doRun(argv []string, returnStdout bool) string {
	if xFlag {
		fmt.Printf("gotest: %s\n", strings.Join(argv, " "))
	}
	command := argv[0]
	if runtime.GOOS == "windows" && command == "gomake" {
		// gomake is a shell script and it cannot be executed directly on Windows.
		cmd := ""
		for i, v := range argv {
			if i > 0 {
				cmd += " "
			}
			cmd += `"` + v + `"`
		}
		argv = []string{"cmd", "/c", "sh", "-c", cmd}
	}
	var err os.Error
	argv[0], err = exec.LookPath(command)
	if err != nil {
		Fatalf("can't find %s: %s", command, err)
	}
	procAttr := &os.ProcAttr{
		Env: env,
		Files: []*os.File{
			os.Stdin,
			os.Stdout,
			os.Stderr,
		},
	}
	var r, w *os.File
	if returnStdout {
		r, w, err = os.Pipe()
		if err != nil {
			Fatalf("can't create pipe: %s", err)
		}
		procAttr.Files[1] = w
	}
	proc, err := os.StartProcess(argv[0], argv, procAttr)
	if err != nil {
		Fatalf("%s failed to start: %s", command, err)
	}
	if returnStdout {
		defer r.Close()
		w.Close()
	}
	waitMsg, err := proc.Wait(0)
	if err != nil || waitMsg == nil {
		Fatalf("%s failed: %s", command, err)
	}
	if !waitMsg.Exited() || waitMsg.ExitStatus() != 0 {
		Fatalf("%q failed: %s", strings.Join(argv, " "), waitMsg)
	}
	if returnStdout {
		b, err := ioutil.ReadAll(r)
		if err != nil {
			Fatalf("can't read output from command: %s", err)
		}
		return string(b)
	}
	return ""
}

// writeTestmainGo generates the test program to be compiled, "./_testmain.go".
func writeTestmainGo() {
	f, err := os.Create("_testmain.go")
	if err != nil {
		Fatalf("can't create _testmain.go: %s", err)
	}
	defer f.Close()
	b := bufio.NewWriter(f)
	defer b.Flush()

	// Package and imports.
	fmt.Fprint(b, "package main\n\n")
	// Are there tests from a package other than the one we're testing?
	// We can't just use file names because some of the things we compiled
	// contain no tests.
	outsideTests := false
	insideTests := false
	for _, f := range files {
		//println(f.name, f.pkg)
		if len(f.tests) == 0 && len(f.benchmarks) == 0 {
			continue
		}
		if strings.HasSuffix(f.pkg, "_test") {
			outsideTests = true
		} else {
			insideTests = true
		}
	}
	if insideTests {
		switch importPath {
		case "testing":
		case "main":
			// Import path main is reserved, so import with
			// explicit reference to ./_test/main instead.
			// Also, the file we are writing defines a function named main,
			// so rename this import to __main__ to avoid name conflict.
			fmt.Fprintf(b, "import __main__ %q\n", "./_test/main")
		default:
			fmt.Fprintf(b, "import %q\n", importPath)
		}
	}
	if outsideTests {
		fmt.Fprintf(b, "import %q\n", "./_xtest_")
	}
	fmt.Fprintf(b, "import %q\n", "testing")
	fmt.Fprintf(b, "import __os__     %q\n", "os")     // rename in case tested package is called os
	fmt.Fprintf(b, "import __regexp__ %q\n", "regexp") // rename in case tested package is called regexp
	fmt.Fprintln(b)                                    // for gofmt

	// Tests.
	fmt.Fprintln(b, "var tests = []testing.InternalTest{")
	for _, f := range files {
		for _, t := range f.tests {
			fmt.Fprintf(b, "\t{\"%s.%s\", %s.%s},\n", f.pkg, t, notMain(f.pkg), t)
		}
	}
	fmt.Fprintln(b, "}")
	fmt.Fprintln(b)

	// Benchmarks.
	fmt.Fprintln(b, "var benchmarks = []testing.InternalBenchmark{")
	for _, f := range files {
		for _, bm := range f.benchmarks {
			fmt.Fprintf(b, "\t{\"%s.%s\", %s.%s},\n", f.pkg, bm, notMain(f.pkg), bm)
		}
	}
	fmt.Fprintln(b, "}")

	// Body.
	fmt.Fprintln(b, testBody)
}

// notMain returns the package, renaming as appropriate if it's "main".
func notMain(pkg string) string {
	if pkg == "main" {
		return "__main__"
	}
	return pkg
}

// testBody is just copied to the output. It's the code that runs the tests.
var testBody = `
var matchPat string
var matchRe *__regexp__.Regexp

func matchString(pat, str string) (result bool, err __os__.Error) {
	if matchRe == nil || matchPat != pat {
		matchPat = pat
		matchRe, err = __regexp__.Compile(matchPat)
		if err != nil {
			return
		}
	}
	return matchRe.MatchString(str), nil
}

func main() {
	testing.Main(matchString, tests, benchmarks)
}`
<commit_msg>gotest: fix windows build.<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bufio"
	"exec"
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"io/ioutil"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"unicode"
	"utf8"
)

// Environment for commands.
var (
	XGC       []string // 6g -I _test -o _xtest_.6
	GC        []string // 6g -I _test _testmain.go
	GL        []string // 6l -L _test _testmain.6
	GOARCH    string
	GOROOT    string
	GORUN     string
	O         string
	args      []string // arguments passed to gotest; also passed to the binary
	fileNames []string
	env       = os.Environ()
)

// These strings are created by getTestNames.
var (
	insideFileNames  []string // list of *.go files inside the package.
	outsideFileNames []string // list of *.go files outside the package (in package foo_test).
)

var (
	files      []*File
	importPath string
)

// Flags for our own purposes. We do our own flag processing.
var (
	cFlag bool
	xFlag bool
)

// File represents a file that contains tests.
type File struct {
	name       string
	pkg        string
	file       *os.File
	astFile    *ast.File
	tests      []string // The names of the TestXXXs.
	benchmarks []string // The names of the BenchmarkXXXs.
}

func main() {
	flags()
	needMakefile()
	setEnvironment()
	getTestFileNames()
	parseFiles()
	getTestNames()
	run("gomake", "testpackage-clean")
	run("gomake", "testpackage", fmt.Sprintf("GOTESTFILES=%s", strings.Join(insideFileNames, " ")))
	if len(outsideFileNames) > 0 {
		run(append(XGC, outsideFileNames...)...)
	}
	importPath = runWithStdout("gomake", "-s", "importpath")
	writeTestmainGo()
	run(GC...)
	run(GL...)
	if !cFlag {
		runTestWithArgs("./" + O + ".out")
	}
}

// needMakefile tests that we have a Makefile in this directory.
func needMakefile() {
	if _, err := os.Stat("Makefile"); err != nil {
		Fatalf("please create a Makefile for gotest; see http://golang.org/doc/code.html for details")
	}
}

// Fatalf formats its arguments, prints the message with a final newline, and exits.
func Fatalf(s string, args ...interface{}) {
	fmt.Fprintf(os.Stderr, "gotest: "+s+"\n", args...)
	os.Exit(2)
}

// theChar is the map from architecture to object character.
var theChar = map[string]string{
	"arm":   "5",
	"amd64": "6",
	"386":   "8",
}

// addEnv adds a name=value pair to the environment passed to subcommands.
// If the item is already in the environment, addEnv replaces the value.
func addEnv(name, value string) {
	for i := 0; i < len(env); i++ {
		if strings.HasPrefix(env[i], name+"=") {
			env[i] = name + "=" + value
			return
		}
	}
	env = append(env, name+"="+value)
}

// setEnvironment assembles the configuration for gotest and its subcommands.
func setEnvironment() {
	// Basic environment.
	GOROOT = runtime.GOROOT()
	addEnv("GOROOT", GOROOT)
	GOARCH = runtime.GOARCH
	addEnv("GOARCH", GOARCH)
	O = theChar[GOARCH]
	if O == "" {
		Fatalf("unknown architecture %s", GOARCH)
	}

	// Commands and their flags.
	gc := os.Getenv("GC")
	if gc == "" {
		gc = O + "g"
	}
	XGC = []string{gc, "-I", "_test", "-o", "_xtest_." + O}
	GC = []string{gc, "-I", "_test", "_testmain.go"}
	gl := os.Getenv("GL")
	if gl == "" {
		gl = O + "l"
	}
	GL = []string{gl, "-L", "_test", "_testmain." + O}

	// Silence make on Linux
	addEnv("MAKEFLAGS", "")
	addEnv("MAKELEVEL", "")
}

// getTestFileNames gets the set of files we're looking at.
// If gotest has no arguments, it scans for file names matching "[^.]*_test.go".
func getTestFileNames() {
	names := fileNames
	if len(names) == 0 {
		var err os.Error
		names, err = filepath.Glob("[^.]*_test.go")
		if err != nil {
			Fatalf("Glob pattern error: %s", err)
		}
		if len(names) == 0 {
			Fatalf(`no test files found: no match for "[^.]*_test.go"`)
		}
	}
	for _, n := range names {
		fd, err := os.Open(n)
		if err != nil {
			Fatalf("%s: %s", n, err)
		}
		f := &File{name: n, file: fd}
		files = append(files, f)
	}
}

// parseFiles parses the files and remembers the packages we find.
func parseFiles() {
	fileSet := token.NewFileSet()
	for _, f := range files {
		// Report declaration errors so we can abort if the files are incorrect Go.
		file, err := parser.ParseFile(fileSet, f.name, nil, parser.DeclarationErrors)
		if err != nil {
			Fatalf("parse error: %s", err)
		}
		f.astFile = file
		f.pkg = file.Name.String()
		if f.pkg == "" {
			Fatalf("cannot happen: no package name in %s", f.name)
		}
	}
}

// getTestNames extracts the names of tests and benchmarks.  They are all
// top-level functions that are not methods.
func getTestNames() {
	for _, f := range files {
		for _, d := range f.astFile.Decls {
			n, ok := d.(*ast.FuncDecl)
			if !ok {
				continue
			}
			if n.Recv != nil { // a method, not a function.
				continue
			}
			name := n.Name.String()
			if isTest(name, "Test") {
				f.tests = append(f.tests, name)
			} else if isTest(name, "Benchmark") {
				f.benchmarks = append(f.benchmarks, name)
			}
			// TODO: worth checking the signature? Probably not.
		}
		if strings.HasSuffix(f.pkg, "_test") {
			outsideFileNames = append(outsideFileNames, f.name)
		} else {
			insideFileNames = append(insideFileNames, f.name)
		}
	}
}

// isTest tells whether name looks like a test (or benchmark, according to prefix).
// It is a Test (say) if there is a character after Test that is not a lower-case letter.
// We don't want TesticularCancer.
func isTest(name, prefix string) bool {
	if !strings.HasPrefix(name, prefix) {
		return false
	}
	if len(name) == len(prefix) { // "Test" is ok
		return true
	}
	rune, _ := utf8.DecodeRuneInString(name[len(prefix):])
	return !unicode.IsLower(rune)
}

func run(args ...string) {
	doRun(args, false)
}

// runWithStdout is like run, but returns the text of standard output with the last newline dropped.
func runWithStdout(argv ...string) string {
	s := doRun(argv, true)
	if len(s) == 0 {
		Fatalf("no output from command %s", strings.Join(argv, " "))
	}
	if s[len(s)-1] == '\n' {
		s = s[:len(s)-1]
	}
	return s
}

// runTestWithArgs appends gotest's runs the provided binary with the args passed on the command line.
func runTestWithArgs(binary string) {
	doRun(append([]string{binary}, args...), false)
}

// doRun is the general command runner.  The flag says whether we want to
// retrieve standard output.
func doRun(argv []string, returnStdout bool) string {
	if xFlag {
		fmt.Printf("gotest: %s\n", strings.Join(argv, " "))
	}
	command := argv[0]
	if runtime.GOOS == "windows" && command == "gomake" {
		// gomake is a shell script and it cannot be executed directly on Windows.
		cmd := ""
		for i, v := range argv {
			if i > 0 {
				cmd += " "
			}
			cmd += `"` + v + `"`
		}
		argv = []string{"cmd", "/c", "sh", "-c", cmd}
	}
	var err os.Error
	argv[0], err = exec.LookPath(argv[0])
	if err != nil {
		Fatalf("can't find %s: %s", command, err)
	}
	procAttr := &os.ProcAttr{
		Env: env,
		Files: []*os.File{
			os.Stdin,
			os.Stdout,
			os.Stderr,
		},
	}
	var r, w *os.File
	if returnStdout {
		r, w, err = os.Pipe()
		if err != nil {
			Fatalf("can't create pipe: %s", err)
		}
		procAttr.Files[1] = w
	}
	proc, err := os.StartProcess(argv[0], argv, procAttr)
	if err != nil {
		Fatalf("%s failed to start: %s", command, err)
	}
	if returnStdout {
		defer r.Close()
		w.Close()
	}
	waitMsg, err := proc.Wait(0)
	if err != nil || waitMsg == nil {
		Fatalf("%s failed: %s", command, err)
	}
	if !waitMsg.Exited() || waitMsg.ExitStatus() != 0 {
		Fatalf("%q failed: %s", strings.Join(argv, " "), waitMsg)
	}
	if returnStdout {
		b, err := ioutil.ReadAll(r)
		if err != nil {
			Fatalf("can't read output from command: %s", err)
		}
		return string(b)
	}
	return ""
}

// writeTestmainGo generates the test program to be compiled, "./_testmain.go".
func writeTestmainGo() {
	f, err := os.Create("_testmain.go")
	if err != nil {
		Fatalf("can't create _testmain.go: %s", err)
	}
	defer f.Close()
	b := bufio.NewWriter(f)
	defer b.Flush()

	// Package and imports.
	fmt.Fprint(b, "package main\n\n")
	// Are there tests from a package other than the one we're testing?
	// We can't just use file names because some of the things we compiled
	// contain no tests.
	outsideTests := false
	insideTests := false
	for _, f := range files {
		//println(f.name, f.pkg)
		if len(f.tests) == 0 && len(f.benchmarks) == 0 {
			continue
		}
		if strings.HasSuffix(f.pkg, "_test") {
			outsideTests = true
		} else {
			insideTests = true
		}
	}
	if insideTests {
		switch importPath {
		case "testing":
		case "main":
			// Import path main is reserved, so import with
			// explicit reference to ./_test/main instead.
			// Also, the file we are writing defines a function named main,
			// so rename this import to __main__ to avoid name conflict.
			fmt.Fprintf(b, "import __main__ %q\n", "./_test/main")
		default:
			fmt.Fprintf(b, "import %q\n", importPath)
		}
	}
	if outsideTests {
		fmt.Fprintf(b, "import %q\n", "./_xtest_")
	}
	fmt.Fprintf(b, "import %q\n", "testing")
	fmt.Fprintf(b, "import __os__     %q\n", "os")     // rename in case tested package is called os
	fmt.Fprintf(b, "import __regexp__ %q\n", "regexp") // rename in case tested package is called regexp
	fmt.Fprintln(b)                                    // for gofmt

	// Tests.
	fmt.Fprintln(b, "var tests = []testing.InternalTest{")
	for _, f := range files {
		for _, t := range f.tests {
			fmt.Fprintf(b, "\t{\"%s.%s\", %s.%s},\n", f.pkg, t, notMain(f.pkg), t)
		}
	}
	fmt.Fprintln(b, "}")
	fmt.Fprintln(b)

	// Benchmarks.
	fmt.Fprintln(b, "var benchmarks = []testing.InternalBenchmark{")
	for _, f := range files {
		for _, bm := range f.benchmarks {
			fmt.Fprintf(b, "\t{\"%s.%s\", %s.%s},\n", f.pkg, bm, notMain(f.pkg), bm)
		}
	}
	fmt.Fprintln(b, "}")

	// Body.
	fmt.Fprintln(b, testBody)
}

// notMain returns the package, renaming as appropriate if it's "main".
func notMain(pkg string) string {
	if pkg == "main" {
		return "__main__"
	}
	return pkg
}

// testBody is just copied to the output. It's the code that runs the tests.
var testBody = `
var matchPat string
var matchRe *__regexp__.Regexp

func matchString(pat, str string) (result bool, err __os__.Error) {
	if matchRe == nil || matchPat != pat {
		matchPat = pat
		matchRe, err = __regexp__.Compile(matchPat)
		if err != nil {
			return
		}
	}
	return matchRe.MatchString(str), nil
}

func main() {
	testing.Main(matchString, tests, benchmarks)
}`
<|endoftext|>


<commit_before>package marvel_test

import (
	"strings"
	"testing"
	"time"

	"github.com/dustinrc/marvel"
	"github.com/stretchr/testify/assert"
)

func TestCharactersAll(t *testing.T) {
	c := newTestClient(t, "characters_all")
	defer c.stopRecorder()

	chars, err := c.Characters.All(nil)
	assert.NoError(t, err, "Characters.All({}) returned an error")
	assert.NotEmpty(t, chars, "Characters.All({}) returned empty character list")
}

func TestCharactersAllName(t *testing.T) {
	c := newTestClient(t, "characters_all_name")
	defer c.stopRecorder()

	params := &marvel.CharacterParams{Name: "Spider-Man"}
	chars, err := c.Characters.All(params)
	assert.NoError(t, err, "Characters.All({}) returned an error")
	assert.Equal(t, 1009610, chars[0].ID, "Incorrect ID")
	assert.Contains(t, strings.ToLower(chars[0].Name), "spider-man", "Incorrect Name")
}

func TestCharactersAllModifiedSince(t *testing.T) {
	c := newTestClient(t, "characters_all_modified_since")
	defer c.stopRecorder()

	modDate := time.Date(2016, time.August, 17, 17, 46, 57, 123, time.UTC)
	params := &marvel.CharacterParams{ModifiedSince: modDate}
	chars, err := c.Characters.All(params)
	assert.NoError(t, err, "Characters.All({}) returned an error")
	assert.Equal(t, 1009268, chars[2].ID, "Incorrect ID")
	assert.Contains(t, strings.ToLower(chars[2].Name), "deadpool", "Incorrect Name")
}

func TestCharactersAllBadParam(t *testing.T) {
	c := newTestClient(t, "characters_all_bad_param")
	defer c.stopRecorder()

	params := &marvel.CharacterParams{OrderBy: "superpower"}
	chars, err := c.Characters.All(params)
	assert.Error(t, err)
	assert.Empty(t, chars, "Character list should have been empty")
}

func TestCharactersGetWrapped(t *testing.T) {
	c := newTestClient(t, "characters_get")
	defer c.stopRecorder()

	wrap, resp, err := c.Characters.GetWrapped(1017575)
	assert.NoError(t, err, "Characters.GetWrapped() returned an error")
	assert.NotNil(t, resp)

	t.Run("CharacterDataWrapper is correct", func(t *testing.T) {
		assert.Equal(t, 200, wrap.Code, "Incorrect Code")
		assert.Equal(t, "ok", strings.ToLower(wrap.Status), "Incorrect Status")
		assert.Contains(t, strings.ToLower(wrap.Copyright), "marvel", "Incorrect Copyright")
		assert.Contains(t, strings.ToLower(wrap.AttributionText), "marvel", "Incorrect AttributionText")
		assert.Contains(t, strings.ToLower(wrap.AttributionHTML), "marvel", "Incorrect AttributionHTML")
		assert.NotEmpty(t, wrap.ETag, "Empty ETag")
	})
	t.Run("CharacterDataContainer is correct", func(t *testing.T) {
		assert.Equal(t, 0, wrap.Data.Offset)
		assert.Equal(t, 20, wrap.Data.Limit)
		assert.Equal(t, 1, wrap.Data.Total)
		assert.Equal(t, 1, wrap.Data.Count)
	})
}

func TestCharactersGet(t *testing.T) {
	c := newTestClient(t, "characters_get")
	defer c.stopRecorder()

	char, err := c.Characters.Get(1017575)
	assert.NoError(t, err, "Characters.Get() returned an error for 1017575")

	assert.Equal(t, 1017575, char.ID, "Incorrect ID")
	assert.Contains(t, strings.ToLower(char.Name), "wilson", "Incorrect Name")
	assert.Contains(t, strings.ToLower(char.Description), "falcon", "Incorrect Description")
	assert.True(t, time.Now().After(char.Modified.Time), "Incorrect Modified time")
	assert.NotEmpty(t, char.Thumbnail, "Incorrect Thumbnail")
	assert.True(t, strings.HasSuffix(strings.ToLower(char.ResourceURI), "characters/1017575"), "Incorrect ResourceURI")
	assert.Equal(t, 25, char.Comics.Available)
	assert.True(t, strings.HasSuffix(strings.ToLower(char.Comics.CollectionURI), "characters/1017575/comics"), "Incorrect comics CollectionURI")
	assert.Equal(t, 9, char.Series.Available)
	assert.True(t, strings.HasSuffix(strings.ToLower(char.Series.CollectionURI), "characters/1017575/series"), "Incorrect series CollectionURI")
	assert.Equal(t, 25, char.Stories.Available)
	assert.True(t, strings.HasSuffix(strings.ToLower(char.Stories.CollectionURI), "characters/1017575/stories"), "Incorrect stories CollectionURI")
	assert.Equal(t, 1, char.Events.Available)
	assert.True(t, strings.HasSuffix(strings.ToLower(char.Events.CollectionURI), "characters/1017575/events"), "Incorrect events CollectionURI")
	assert.NotEmpty(t, char.URLs, "Incorrect URLs")
}

func TestCharacterGetBadID(t *testing.T) {
	c := newTestClient(t, "characters_get_bad_id")
	defer c.stopRecorder()

	char, err := c.Characters.Get(-1)
	assert.Error(t, err)
	assert.Nil(t, char, "Character should have been nil")
}
<commit_msg>Verify slice of IDs works in option params<commit_after>package marvel_test

import (
	"strings"
	"testing"
	"time"

	"github.com/dustinrc/marvel"
	"github.com/stretchr/testify/assert"
)

func TestCharactersAll(t *testing.T) {
	c := newTestClient(t, "characters_all")
	defer c.stopRecorder()

	chars, err := c.Characters.All(nil)
	assert.NoError(t, err, "Characters.All({}) returned an error")
	assert.NotEmpty(t, chars, "Characters.All({}) returned empty character list")
}

func TestCharactersAllName(t *testing.T) {
	c := newTestClient(t, "characters_all_name")
	defer c.stopRecorder()

	params := &marvel.CharacterParams{Name: "Spider-Man"}
	chars, err := c.Characters.All(params)
	assert.NoError(t, err, "Characters.All({}) returned an error")
	assert.Equal(t, 1009610, chars[0].ID, "Incorrect ID")
	assert.Contains(t, strings.ToLower(chars[0].Name), "spider-man", "Incorrect Name")
}

func TestCharactersAllModifiedSince(t *testing.T) {
	c := newTestClient(t, "characters_all_modified_since")
	defer c.stopRecorder()

	modDate := time.Date(2016, time.August, 17, 17, 46, 57, 123, time.UTC)
	params := &marvel.CharacterParams{ModifiedSince: modDate}
	chars, err := c.Characters.All(params)
	assert.NoError(t, err, "Characters.All({}) returned an error")
	assert.Equal(t, 1009268, chars[2].ID, "Incorrect ID")
	assert.Contains(t, strings.ToLower(chars[2].Name), "deadpool", "Incorrect Name")
}

func TestCharactersAllOnlyInComics(t *testing.T) {
	c := newTestClient(t, "characters_all_only_in_comics")
	defer c.stopRecorder()

	params := &marvel.CharacterParams{Comics: []int{11200, 22222}}
	chars, err := c.Characters.All(params)
	assert.NoError(t, err, "Characters.All({}) returned an error")
	assert.Equal(t, 1009515, chars[0].ID, "Incorrect ID")
	assert.Equal(t, 1010791, chars[1].ID, "Incorrect ID")
}

func TestCharactersAllBadParam(t *testing.T) {
	c := newTestClient(t, "characters_all_bad_param")
	defer c.stopRecorder()

	params := &marvel.CharacterParams{OrderBy: "superpower"}
	chars, err := c.Characters.All(params)
	assert.Error(t, err)
	assert.Empty(t, chars, "Character list should have been empty")
}

func TestCharactersGetWrapped(t *testing.T) {
	c := newTestClient(t, "characters_get")
	defer c.stopRecorder()

	wrap, resp, err := c.Characters.GetWrapped(1017575)
	assert.NoError(t, err, "Characters.GetWrapped() returned an error")
	assert.NotNil(t, resp)

	t.Run("CharacterDataWrapper is correct", func(t *testing.T) {
		assert.Equal(t, 200, wrap.Code, "Incorrect Code")
		assert.Equal(t, "ok", strings.ToLower(wrap.Status), "Incorrect Status")
		assert.Contains(t, strings.ToLower(wrap.Copyright), "marvel", "Incorrect Copyright")
		assert.Contains(t, strings.ToLower(wrap.AttributionText), "marvel", "Incorrect AttributionText")
		assert.Contains(t, strings.ToLower(wrap.AttributionHTML), "marvel", "Incorrect AttributionHTML")
		assert.NotEmpty(t, wrap.ETag, "Empty ETag")
	})
	t.Run("CharacterDataContainer is correct", func(t *testing.T) {
		assert.Equal(t, 0, wrap.Data.Offset)
		assert.Equal(t, 20, wrap.Data.Limit)
		assert.Equal(t, 1, wrap.Data.Total)
		assert.Equal(t, 1, wrap.Data.Count)
	})
}

func TestCharactersGet(t *testing.T) {
	c := newTestClient(t, "characters_get")
	defer c.stopRecorder()

	char, err := c.Characters.Get(1017575)
	assert.NoError(t, err, "Characters.Get() returned an error for 1017575")

	assert.Equal(t, 1017575, char.ID, "Incorrect ID")
	assert.Contains(t, strings.ToLower(char.Name), "wilson", "Incorrect Name")
	assert.Contains(t, strings.ToLower(char.Description), "falcon", "Incorrect Description")
	assert.True(t, time.Now().After(char.Modified.Time), "Incorrect Modified time")
	assert.NotEmpty(t, char.Thumbnail, "Incorrect Thumbnail")
	assert.True(t, strings.HasSuffix(strings.ToLower(char.ResourceURI), "characters/1017575"), "Incorrect ResourceURI")
	assert.Equal(t, 25, char.Comics.Available)
	assert.True(t, strings.HasSuffix(strings.ToLower(char.Comics.CollectionURI), "characters/1017575/comics"), "Incorrect comics CollectionURI")
	assert.Equal(t, 9, char.Series.Available)
	assert.True(t, strings.HasSuffix(strings.ToLower(char.Series.CollectionURI), "characters/1017575/series"), "Incorrect series CollectionURI")
	assert.Equal(t, 25, char.Stories.Available)
	assert.True(t, strings.HasSuffix(strings.ToLower(char.Stories.CollectionURI), "characters/1017575/stories"), "Incorrect stories CollectionURI")
	assert.Equal(t, 1, char.Events.Available)
	assert.True(t, strings.HasSuffix(strings.ToLower(char.Events.CollectionURI), "characters/1017575/events"), "Incorrect events CollectionURI")
	assert.NotEmpty(t, char.URLs, "Incorrect URLs")
}

func TestCharacterGetBadID(t *testing.T) {
	c := newTestClient(t, "characters_get_bad_id")
	defer c.stopRecorder()

	char, err := c.Characters.Get(-1)
	assert.Error(t, err)
	assert.Nil(t, char, "Character should have been nil")
}
<|endoftext|>


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

// +build !static

package inspectors

import (
	"flag"

	log "github.com/cihub/seelog"
	"github.com/mitchellh/cli"

	inspector "github.com/osrg/namazu/nmz/inspector/ethernet"
)

type etherFlags struct {
	commonFlags
	HookSwitchZMQAddr string
	NFQNumber         int
}

var (
	etherFlagset = flag.NewFlagSet("ethernet", flag.ExitOnError)
	_etherFlags  = etherFlags{}
)

func init() {
	initCommon(etherFlagset, &_etherFlags.commonFlags, "_namazu_ethernet_inspector")
	etherFlagset.StringVar(&_etherFlags.HookSwitchZMQAddr, "hookswitch",
		"ipc:///tmp/namazu-hookswitch-zmq", "HookSwitch ZeroMQ addr")
	etherFlagset.IntVar(&_etherFlags.NFQNumber, "nfq-number",
		-1, "netfilter_queue number")
}

type etherCmd struct {
}

func EtherCommandFactory() (cli.Command, error) {
	return etherCmd{}, nil
}

func (cmd etherCmd) Help() string {
	return "Please run `nmz --help inspectors` instead"
}

func (cmd etherCmd) Synopsis() string {
	return "Start Ethernet inspector"
}

func (cmd etherCmd) Run(args []string) int {
	if err := etherFlagset.Parse(args); err != nil {
		log.Critical(err)
		return 1
	}

	useHookSwitch := _etherFlags.NFQNumber < 0

	if useHookSwitch && _etherFlags.HookSwitchZMQAddr == "" {
		log.Critical("hookswitch is invalid")
		return 1
	}
	if !useHookSwitch && _etherFlags.NFQNumber > 0xFFFF {
		log.Critical("nfq-number is invalid")
		return 1
	}

	autopilot, err := conditionalStartAutopilotOrchestrator(_fsFlags.commonFlags)
	if err != nil {
		log.Critical(err)
		return 1
	}
	log.Infof("Autopilot-mode: %t", autopilot)

	var etherInspector inspector.EthernetInspector
	if useHookSwitch {
		log.Infof("Using hookswitch %s", _etherFlags.HookSwitchZMQAddr)
		etherInspector = &inspector.HookSwitchInspector{
			OrchestratorURL:   _etherFlags.OrchestratorURL,
			EntityID:          _etherFlags.EntityID,
			HookSwitchZMQAddr: _etherFlags.HookSwitchZMQAddr,
			EnableTCPWatcher:  true,
		}
	} else {
		log.Infof("Using NFQ %s", _etherFlags.HookSwitchZMQAddr)
		etherInspector = &inspector.NFQInspector{
			OrchestratorURL:  _etherFlags.OrchestratorURL,
			EntityID:         _etherFlags.EntityID,
			NFQNumber:        uint16(_etherFlags.NFQNumber),
			EnableTCPWatcher: true,
		}
	}

	if err := etherInspector.Serve(); err != nil {
		panic(log.Critical(err))
	}

	// NOTREACHED
	return 0
}
<commit_msg>cli:fix ethernet inspector cannot init commonFlags (#173)<commit_after>// Copyright (C) 2015 Nippon Telegraph and Telephone Corporation.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// +build !static

package inspectors

import (
	"flag"

	log "github.com/cihub/seelog"
	"github.com/mitchellh/cli"

	inspector "github.com/osrg/namazu/nmz/inspector/ethernet"
)

type etherFlags struct {
	commonFlags
	HookSwitchZMQAddr string
	NFQNumber         int
}

var (
	etherFlagset = flag.NewFlagSet("ethernet", flag.ExitOnError)
	_etherFlags  = etherFlags{}
)

func init() {
	initCommon(etherFlagset, &_etherFlags.commonFlags, "_namazu_ethernet_inspector")
	etherFlagset.StringVar(&_etherFlags.HookSwitchZMQAddr, "hookswitch",
		"ipc:///tmp/namazu-hookswitch-zmq", "HookSwitch ZeroMQ addr")
	etherFlagset.IntVar(&_etherFlags.NFQNumber, "nfq-number",
		-1, "netfilter_queue number")
}

type etherCmd struct {
}

func EtherCommandFactory() (cli.Command, error) {
	return etherCmd{}, nil
}

func (cmd etherCmd) Help() string {
	return "Please run `nmz --help inspectors` instead"
}

func (cmd etherCmd) Synopsis() string {
	return "Start Ethernet inspector"
}

func (cmd etherCmd) Run(args []string) int {
	if err := etherFlagset.Parse(args); err != nil {
		log.Critical(err)
		return 1
	}

	useHookSwitch := _etherFlags.NFQNumber < 0

	if useHookSwitch && _etherFlags.HookSwitchZMQAddr == "" {
		log.Critical("hookswitch is invalid")
		return 1
	}
	if !useHookSwitch && _etherFlags.NFQNumber > 0xFFFF {
		log.Critical("nfq-number is invalid")
		return 1
	}

	autopilot, err := conditionalStartAutopilotOrchestrator(_etherFlags.commonFlags)
	if err != nil {
		log.Critical(err)
		return 1
	}
	log.Infof("Autopilot-mode: %t", autopilot)

	var etherInspector inspector.EthernetInspector
	if useHookSwitch {
		log.Infof("Using hookswitch %s", _etherFlags.HookSwitchZMQAddr)
		etherInspector = &inspector.HookSwitchInspector{
			OrchestratorURL:   _etherFlags.OrchestratorURL,
			EntityID:          _etherFlags.EntityID,
			HookSwitchZMQAddr: _etherFlags.HookSwitchZMQAddr,
			EnableTCPWatcher:  true,
		}
	} else {
		log.Infof("Using NFQ %s", _etherFlags.HookSwitchZMQAddr)
		etherInspector = &inspector.NFQInspector{
			OrchestratorURL:  _etherFlags.OrchestratorURL,
			EntityID:         _etherFlags.EntityID,
			NFQNumber:        uint16(_etherFlags.NFQNumber),
			EnableTCPWatcher: true,
		}
	}

	if err := etherInspector.Serve(); err != nil {
		panic(log.Critical(err))
	}

	// NOTREACHED
	return 0
}
<|endoftext|>


<commit_before>// Package serve implements the serve command for CFSSL's API.
package serve

import (
	"errors"
	"net"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"

	rice "github.com/GeertJohan/go.rice"
	"github.com/cloudflare/cfssl/api/bundle"
	"github.com/cloudflare/cfssl/api/certinfo"
	"github.com/cloudflare/cfssl/api/generator"
	"github.com/cloudflare/cfssl/api/info"
	"github.com/cloudflare/cfssl/api/initca"
	apiocsp "github.com/cloudflare/cfssl/api/ocsp"
	"github.com/cloudflare/cfssl/api/scan"
	apisign "github.com/cloudflare/cfssl/api/sign"
	"github.com/cloudflare/cfssl/bundler"
	"github.com/cloudflare/cfssl/cli"
	ocspsign "github.com/cloudflare/cfssl/cli/ocspsign"
	"github.com/cloudflare/cfssl/cli/sign"
	"github.com/cloudflare/cfssl/log"
	"github.com/cloudflare/cfssl/ocsp"
	"github.com/cloudflare/cfssl/signer"
	"github.com/cloudflare/cfssl/ubiquity"
)

// Usage text of 'cfssl serve'
var serverUsageText = `cfssl serve -- set up a HTTP server handles CF SSL requests

Usage of serve:
        cfssl serve [-address address] [-ca cert] [-ca-bundle bundle] \
                    [-ca-key key] [-int-bundle bundle] [-int-dir dir] [-port port] \
                    [-metadata file] [-remote remote_host] [-config config] \
                    [-responder cert] [-responder-key key]

Flags:
`

// Flags used by 'cfssl serve'
var serverFlags = []string{"address", "port", "ca", "ca-key", "ca-bundle", "int-bundle", "int-dir", "metadata", "remote", "config", "responder", "responder-key"}

var (
	conf       cli.Config
	s          signer.Signer
	ocspSigner ocsp.Signer
)

// V1APIPrefix is the prefix of all CFSSL V1 API Endpoints.
var V1APIPrefix = "/api/v1/cfssl/"

// v1APIPath prepends the V1 API prefix to endpoints not beginning with "/"
func v1APIPath(path string) string {
	if !strings.HasPrefix(path, "/") {
		path = V1APIPrefix + path
	}
	return (&url.URL{Path: path}).String()
}

// httpBox implements http.FileSystem which allows the use of Box with a http.FileServer.
// Atempting to Open an API endpoint will result in an error.
type httpBox struct {
	*rice.Box
	redirects map[string]string
}

// Open returns a File for non-API enpoints using the http.File interface.
func (hb *httpBox) Open(name string) (http.File, error) {
	if strings.HasPrefix(name, V1APIPrefix) {
		return nil, os.ErrNotExist
	}

	if location, ok := hb.redirects[name]; ok {
		return hb.Box.Open(location)
	}

	return hb.Box.Open(name)
}

// staticBox is the box containing all static assets.
var staticBox = &httpBox{
	Box: rice.MustFindBox("static"),
	redirects: map[string]string{
		"/scan":   "/index.html",
		"/bundle": "/index.html",
	},
}

var errBadSigner = errors.New("signer not initialized")

var endpoints = map[string]func() (http.Handler, error){
	"sign": func() (http.Handler, error) {
		if s == nil {
			return nil, errBadSigner
		}
		return apisign.NewHandlerFromSigner(s)
	},

	"authsign": func() (http.Handler, error) {
		if s == nil {
			return nil, errBadSigner
		}
		return apisign.NewAuthHandlerFromSigner(s)
	},

	"info": func() (http.Handler, error) {
		if s == nil {
			return nil, errBadSigner
		}
		return info.NewHandler(s)
	},

	"newcert": func() (http.Handler, error) {
		if s == nil {
			return nil, errBadSigner
		}
		return generator.NewCertGeneratorHandlerFromSigner(generator.CSRValidate, s), nil
	},

	"bundle": func() (http.Handler, error) {
		return bundle.NewHandler(conf.CABundleFile, conf.IntBundleFile)
	},

	"newkey": func() (http.Handler, error) {
		return generator.NewHandler(generator.CSRValidate)
	},

	"init_ca": func() (http.Handler, error) {
		return initca.NewHandler(), nil
	},

	"scan": func() (http.Handler, error) {
		return scan.NewHandler(conf.CABundleFile)
	},

	"scaninfo": func() (http.Handler, error) {
		return scan.NewInfoHandler(), nil
	},

	"certinfo": func() (http.Handler, error) {
		return certinfo.NewHandler(), nil
	},

	"ocspsign": func() (http.Handler, error) {
		if ocspSigner == nil {
			return nil, errBadSigner
		}
		return apiocsp.NewHandler(ocspSigner), nil
	},

	"/": func() (http.Handler, error) {
		return http.FileServer(staticBox), nil
	},
}

// registerHandlers instantiates various handlers and associate them to corresponding endpoints.
func registerHandlers() {
	for path, getHandler := range endpoints {
		path = v1APIPath(path)
		log.Infof("Setting up '%s' endpoint", path)
		if handler, err := getHandler(); err != nil {
			log.Warningf("endpoint '%s' is disabled: %v", path, err)
		} else {
			http.Handle(path, handler)
		}
	}

	log.Info("Handler set up complete.")
}

// serverMain is the command line entry point to the API server. It sets up a
// new HTTP server to handle sign, bundle, and validate requests.
func serverMain(args []string, c cli.Config) error {
	conf = c
	// serve doesn't support arguments.
	if len(args) > 0 {
		return errors.New("argument is provided but not defined; please refer to the usage by flag -h")
	}

	bundler.IntermediateStash = conf.IntDir
	var err error

	if err = ubiquity.LoadPlatforms(conf.Metadata); err != nil {
		return err
	}

	log.Info("Initializing signer")
	if s, err = sign.SignerFromConfig(c); err != nil {
		log.Warningf("couldn't initialize signer: %v", err)
	}

	if ocspSigner, err = ocspsign.SignerFromConfig(c); err != nil {
		log.Warningf("couldn't initialize ocsp signer: %v", err)
	}

	registerHandlers()

	addr := net.JoinHostPort(conf.Address, strconv.Itoa(conf.Port))
	log.Info("Now listening on ", addr)
	return http.ListenAndServe(addr, nil)
}

// Command assembles the definition of Command 'serve'
var Command = &cli.Command{UsageText: serverUsageText, Flags: serverFlags, Main: serverMain}
<commit_msg>Update rice FindBox to run after init-time<commit_after>// Package serve implements the serve command for CFSSL's API.
package serve

import (
	"errors"
	"net"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"

	rice "github.com/GeertJohan/go.rice"
	"github.com/cloudflare/cfssl/api/bundle"
	"github.com/cloudflare/cfssl/api/certinfo"
	"github.com/cloudflare/cfssl/api/generator"
	"github.com/cloudflare/cfssl/api/info"
	"github.com/cloudflare/cfssl/api/initca"
	apiocsp "github.com/cloudflare/cfssl/api/ocsp"
	"github.com/cloudflare/cfssl/api/scan"
	apisign "github.com/cloudflare/cfssl/api/sign"
	"github.com/cloudflare/cfssl/bundler"
	"github.com/cloudflare/cfssl/cli"
	ocspsign "github.com/cloudflare/cfssl/cli/ocspsign"
	"github.com/cloudflare/cfssl/cli/sign"
	"github.com/cloudflare/cfssl/log"
	"github.com/cloudflare/cfssl/ocsp"
	"github.com/cloudflare/cfssl/signer"
	"github.com/cloudflare/cfssl/ubiquity"
)

// Usage text of 'cfssl serve'
var serverUsageText = `cfssl serve -- set up a HTTP server handles CF SSL requests

Usage of serve:
        cfssl serve [-address address] [-ca cert] [-ca-bundle bundle] \
                    [-ca-key key] [-int-bundle bundle] [-int-dir dir] [-port port] \
                    [-metadata file] [-remote remote_host] [-config config] \
                    [-responder cert] [-responder-key key]

Flags:
`

// Flags used by 'cfssl serve'
var serverFlags = []string{"address", "port", "ca", "ca-key", "ca-bundle", "int-bundle", "int-dir", "metadata", "remote", "config", "responder", "responder-key"}

var (
	conf       cli.Config
	s          signer.Signer
	ocspSigner ocsp.Signer
)

// V1APIPrefix is the prefix of all CFSSL V1 API Endpoints.
var V1APIPrefix = "/api/v1/cfssl/"

// v1APIPath prepends the V1 API prefix to endpoints not beginning with "/"
func v1APIPath(path string) string {
	if !strings.HasPrefix(path, "/") {
		path = V1APIPrefix + path
	}
	return (&url.URL{Path: path}).String()
}

// httpBox implements http.FileSystem which allows the use of Box with a http.FileServer.
// Atempting to Open an API endpoint will result in an error.
type httpBox struct {
	*rice.Box
	redirects map[string]string
}

func (hb *httpBox) findBox(name string) (err error) {
	hb.Box, err = rice.FindBox(name)
	return
}

// Open returns a File for non-API enpoints using the http.File interface.
func (hb *httpBox) Open(name string) (http.File, error) {
	if strings.HasPrefix(name, V1APIPrefix) {
		return nil, os.ErrNotExist
	}

	if location, ok := hb.redirects[name]; ok {
		return hb.Box.Open(location)
	}

	return hb.Box.Open(name)
}

// staticBox is the box containing all static assets.
var staticBox = &httpBox{
	redirects: map[string]string{
		"/scan":   "/index.html",
		"/bundle": "/index.html",
	},
}

var errBadSigner = errors.New("signer not initialized")

var endpoints = map[string]func() (http.Handler, error){
	"sign": func() (http.Handler, error) {
		if s == nil {
			return nil, errBadSigner
		}
		return apisign.NewHandlerFromSigner(s)
	},

	"authsign": func() (http.Handler, error) {
		if s == nil {
			return nil, errBadSigner
		}
		return apisign.NewAuthHandlerFromSigner(s)
	},

	"info": func() (http.Handler, error) {
		if s == nil {
			return nil, errBadSigner
		}
		return info.NewHandler(s)
	},

	"newcert": func() (http.Handler, error) {
		if s == nil {
			return nil, errBadSigner
		}
		return generator.NewCertGeneratorHandlerFromSigner(generator.CSRValidate, s), nil
	},

	"bundle": func() (http.Handler, error) {
		return bundle.NewHandler(conf.CABundleFile, conf.IntBundleFile)
	},

	"newkey": func() (http.Handler, error) {
		return generator.NewHandler(generator.CSRValidate)
	},

	"init_ca": func() (http.Handler, error) {
		return initca.NewHandler(), nil
	},

	"scan": func() (http.Handler, error) {
		return scan.NewHandler(conf.CABundleFile)
	},

	"scaninfo": func() (http.Handler, error) {
		return scan.NewInfoHandler(), nil
	},

	"certinfo": func() (http.Handler, error) {
		return certinfo.NewHandler(), nil
	},

	"ocspsign": func() (http.Handler, error) {
		if ocspSigner == nil {
			return nil, errBadSigner
		}
		return apiocsp.NewHandler(ocspSigner), nil
	},

	"/": func() (http.Handler, error) {
		if err := staticBox.findBox("static"); err != nil {
			return nil, err
		}

		return http.FileServer(staticBox), nil
	},
}

// registerHandlers instantiates various handlers and associate them to corresponding endpoints.
func registerHandlers() {
	for path, getHandler := range endpoints {
		path = v1APIPath(path)
		log.Infof("Setting up '%s' endpoint", path)
		if handler, err := getHandler(); err != nil {
			log.Warningf("endpoint '%s' is disabled: %v", path, err)
		} else {
			http.Handle(path, handler)
		}
	}

	log.Info("Handler set up complete.")
}

// serverMain is the command line entry point to the API server. It sets up a
// new HTTP server to handle sign, bundle, and validate requests.
func serverMain(args []string, c cli.Config) error {
	conf = c
	// serve doesn't support arguments.
	if len(args) > 0 {
		return errors.New("argument is provided but not defined; please refer to the usage by flag -h")
	}

	bundler.IntermediateStash = conf.IntDir
	var err error

	if err = ubiquity.LoadPlatforms(conf.Metadata); err != nil {
		return err
	}

	log.Info("Initializing signer")
	if s, err = sign.SignerFromConfig(c); err != nil {
		log.Warningf("couldn't initialize signer: %v", err)
	}

	if ocspSigner, err = ocspsign.SignerFromConfig(c); err != nil {
		log.Warningf("couldn't initialize ocsp signer: %v", err)
	}

	registerHandlers()

	addr := net.JoinHostPort(conf.Address, strconv.Itoa(conf.Port))
	log.Info("Now listening on ", addr)
	return http.ListenAndServe(addr, nil)
}

// Command assembles the definition of Command 'serve'
var Command = &cli.Command{UsageText: serverUsageText, Flags: serverFlags, Main: serverMain}
<|endoftext|>