<commit_before>package main

import "gopkg.in/mgo.v2/bson"

// Create a new ID string at random
func NewId() string {
	return bson.NewObjectId().Hex()
}
<commit_msg>Add list contains util func<commit_after>package main

import "gopkg.in/mgo.v2/bson"

// Create a new ID string at random
func NewId() string {
	return bson.NewObjectId().Hex()
}

// Return if the given key is in the given list
func ListContains(key string, values []string) bool {
	for _, v := range values {
		if v == key {
			return true
		}
	}
	return false
}<|endoftext|>


<commit_before>package server

import (
	"encoding/json"
	"github.com/Comcast/webpa-common/concurrent"
	"github.com/Comcast/webpa-common/logging"
	"sync"
)

// serverExecutor is a local interface describing the set of methods the underlying
// server object must implement. *http.Server, for example, implements this interface.
type serverExecutor interface {
	ListenAndServe() error
	ListenAndServeTLS(certificateFile, keyFile string) error
}

// WebPA represents a server component within the WebPA cluster.  It is used for both
// primary servers (e.g. petasos) and supporting, embedded servers such as pprof.
type WebPA struct {
	name            string
	address         string
	serverExecutor  serverExecutor
	certificateFile string
	keyFile         string
	logger          logging.Logger
	once            sync.Once
}

var _ concurrent.Runnable = (*WebPA)(nil)

func (w *WebPA) Name() string {
	return w.name
}

func (w *WebPA) Address() string {
	return w.address
}

func (w *WebPA) CertificateFile() string {
	return w.certificateFile
}

func (w *WebPA) KeyFile() string {
	return w.keyFile
}

func (w *WebPA) Secure() bool {
	return len(w.certificateFile) > 0 && len(w.keyFile) > 0
}

func (w *WebPA) Logger() logging.Logger {
	return w.logger
}

func (w *WebPA) String() string {
	data, err := w.MarshalJSON()
	if err != nil {
		return err.Error()
	}

	return string(data)
}

func (w *WebPA) MarshalJSON() ([]byte, error) {
	data := struct {
		Name            string `json:"name"`
		Address         string `json:"address"`
		CertificateFile string `json:"cert"`
		KeyFile         string `json:"key"`
	}{
		Name:            w.name,
		Address:         w.address,
		CertificateFile: w.certificateFile,
		KeyFile:         w.keyFile,
	}

	return json.Marshal(&data)
}

// Run executes this WebPA server.  If both certificateFile and keyFile are non-empty, this method will start
// an HTTPS server using the configured certificate and key.  Otherwise, it will
// start an HTTP server.
//
// This method spawns a goroutine that actually executes the appropriate serverExecutor.ListenXXX method.
// The supplied sync.WaitGroup is incremented, and sync.WaitGroup.Done() is called when the
// spawned goroutine exits.
//
// Run is idemptotent.  It can only be execute once, and subsequent invocations have
// no effect.
func (w *WebPA) Run(waitGroup *sync.WaitGroup, shutdown <-chan struct{}) error {
	w.once.Do(func() {
		waitGroup.Add(1)
		go func() {
			defer waitGroup.Done()
			var err error
			w.logger.Info("Starting [%s]", w.name)
			if w.Secure() {
				err = w.serverExecutor.ListenAndServeTLS(w.certificateFile, w.keyFile)
			} else {
				err = w.serverExecutor.ListenAndServe()
			}

			w.logger.Error("%s exiting: %v", w.name, err)
		}()
	})

	return nil
}
<commit_msg>Removed the Add() and Done() calls, as they cause deadlocking<commit_after>package server

import (
	"encoding/json"
	"github.com/Comcast/webpa-common/concurrent"
	"github.com/Comcast/webpa-common/logging"
	"sync"
)

// serverExecutor is a local interface describing the set of methods the underlying
// server object must implement. *http.Server, for example, implements this interface.
type serverExecutor interface {
	ListenAndServe() error
	ListenAndServeTLS(certificateFile, keyFile string) error
}

// WebPA represents a server component within the WebPA cluster.  It is used for both
// primary servers (e.g. petasos) and supporting, embedded servers such as pprof.
type WebPA struct {
	name            string
	address         string
	serverExecutor  serverExecutor
	certificateFile string
	keyFile         string
	logger          logging.Logger
	once            sync.Once
}

var _ concurrent.Runnable = (*WebPA)(nil)

func (w *WebPA) Name() string {
	return w.name
}

func (w *WebPA) Address() string {
	return w.address
}

func (w *WebPA) CertificateFile() string {
	return w.certificateFile
}

func (w *WebPA) KeyFile() string {
	return w.keyFile
}

func (w *WebPA) Secure() bool {
	return len(w.certificateFile) > 0 && len(w.keyFile) > 0
}

func (w *WebPA) Logger() logging.Logger {
	return w.logger
}

func (w *WebPA) String() string {
	data, err := w.MarshalJSON()
	if err != nil {
		return err.Error()
	}

	return string(data)
}

func (w *WebPA) MarshalJSON() ([]byte, error) {
	data := struct {
		Name            string `json:"name"`
		Address         string `json:"address"`
		CertificateFile string `json:"cert"`
		KeyFile         string `json:"key"`
	}{
		Name:            w.name,
		Address:         w.address,
		CertificateFile: w.certificateFile,
		KeyFile:         w.keyFile,
	}

	return json.Marshal(&data)
}

// Run executes this WebPA server.  If both certificateFile and keyFile are non-empty, this method will start
// an HTTPS server using the configured certificate and key.  Otherwise, it will
// start an HTTP server on a separate goroutine and return.
//
// This method allows WebPA to implement concurrent.Runnable.  However, neither parameter to this method
// is used.  Because of the way net/http starts servers, there's no controlled way to shut them down.
//
// Run is idemptotent.  It can only be execute once, and subsequent invocations have
// no effect.
func (w *WebPA) Run(waitGroup *sync.WaitGroup, shutdown <-chan struct{}) error {
	w.once.Do(func() {
		go func() {
			var err error
			w.logger.Info("Starting [%s]", w.name)
			if w.Secure() {
				err = w.serverExecutor.ListenAndServeTLS(w.certificateFile, w.keyFile)
			} else {
				err = w.serverExecutor.ListenAndServe()
			}

			w.logger.Error("%s exiting: %v", w.name, err)
		}()
	})

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2013 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 runtime_test

import (
	"flag"
	"fmt"
	"reflect"
	. "runtime"
	"testing"
	"time"
	"unsafe"
)

func TestMemStats(t *testing.T) {
	// Make sure there's at least one forced GC.
	GC()

	// Test that MemStats has sane values.
	st := new(MemStats)
	ReadMemStats(st)

	nz := func(x interface{}) error {
		if x != reflect.Zero(reflect.TypeOf(x)).Interface() {
			return nil
		}
		return fmt.Errorf("zero value")
	}
	le := func(thresh uint64) func(interface{}) error {
		return func(x interface{}) error {
			if reflect.ValueOf(x).Uint() < thresh {
				return nil
			}
			return fmt.Errorf("insanely high value (overflow?); want <= %d", thresh)
		}
	}
	// Of the uint fields, HeapReleased, HeapIdle, PauseTotalNs, and NumGC can be 0.
	fields := map[string][]func(interface{}) error{
		"Alloc": {nz, le(1e10)}, "TotalAlloc": {nz, le(1e11)}, "Sys": {nz, le(1e10)},
		"Lookups": {nz, le(1e10)}, "Mallocs": {nz, le(1e10)}, "Frees": {nz, le(1e10)},
		"HeapAlloc": {nz, le(1e10)}, "HeapSys": {nz, le(1e10)}, "HeapIdle": {le(1e10)},
		"HeapInuse": {nz, le(1e10)}, "HeapReleased": nil, "HeapObjects": {nz, le(1e10)},
		"StackInuse": {nz, le(1e10)}, "StackSys": {nz, le(1e10)},
		"MSpanInuse": {nz, le(1e10)}, "MSpanSys": {nz, le(1e10)},
		"MCacheInuse": {nz, le(1e10)}, "MCacheSys": {nz, le(1e10)},
		"BuckHashSys": {nz, le(1e10)}, "GCSys": {nz, le(1e10)}, "OtherSys": {nz, le(1e10)},
		"NextGC": {nz, le(1e10)}, "LastGC": nil,
		"PauseTotalNs": {le(1e11)}, "PauseNs": nil, "PauseEnd": nil,
		"NumGC": {le(1e9)}, "NumForcedGC": {nz, le(1e9)},
		"GCCPUFraction": nil, "EnableGC": nil, "DebugGC": nil,
		"BySize": nil,
	}

	rst := reflect.ValueOf(st).Elem()
	for i := 0; i < rst.Type().NumField(); i++ {
		name, val := rst.Type().Field(i).Name, rst.Field(i).Interface()
		checks, ok := fields[name]
		if !ok {
			t.Errorf("unknown MemStats field %s", name)
			continue
		}
		for _, check := range checks {
			if err := check(val); err != nil {
				t.Errorf("%s = %v: %s", name, val, err)
			}
		}
	}

	if st.Sys != st.HeapSys+st.StackSys+st.MSpanSys+st.MCacheSys+
		st.BuckHashSys+st.GCSys+st.OtherSys {
		t.Fatalf("Bad sys value: %+v", *st)
	}

	if st.HeapIdle+st.HeapInuse != st.HeapSys {
		t.Fatalf("HeapIdle(%d) + HeapInuse(%d) should be equal to HeapSys(%d), but isn't.", st.HeapIdle, st.HeapInuse, st.HeapSys)
	}

	if lpe := st.PauseEnd[int(st.NumGC+255)%len(st.PauseEnd)]; st.LastGC != lpe {
		t.Fatalf("LastGC(%d) != last PauseEnd(%d)", st.LastGC, lpe)
	}

	var pauseTotal uint64
	for _, pause := range st.PauseNs {
		pauseTotal += pause
	}
	if int(st.NumGC) < len(st.PauseNs) {
		// We have all pauses, so this should be exact.
		if st.PauseTotalNs != pauseTotal {
			t.Fatalf("PauseTotalNs(%d) != sum PauseNs(%d)", st.PauseTotalNs, pauseTotal)
		}
		for i := int(st.NumGC); i < len(st.PauseNs); i++ {
			if st.PauseNs[i] != 0 {
				t.Fatalf("Non-zero PauseNs[%d]: %+v", i, st)
			}
			if st.PauseEnd[i] != 0 {
				t.Fatalf("Non-zero PauseEnd[%d]: %+v", i, st)
			}
		}
	} else {
		if st.PauseTotalNs < pauseTotal {
			t.Fatalf("PauseTotalNs(%d) < sum PauseNs(%d)", st.PauseTotalNs, pauseTotal)
		}
	}

	if st.NumForcedGC > st.NumGC {
		t.Fatalf("NumForcedGC(%d) > NumGC(%d)", st.NumForcedGC, st.NumGC)
	}
}

func TestStringConcatenationAllocs(t *testing.T) {
	n := testing.AllocsPerRun(1e3, func() {
		b := make([]byte, 10)
		for i := 0; i < 10; i++ {
			b[i] = byte(i) + '0'
		}
		s := "foo" + string(b)
		if want := "foo0123456789"; s != want {
			t.Fatalf("want %v, got %v", want, s)
		}
	})
	// Only string concatenation allocates.
	if n != 1 {
		t.Fatalf("want 1 allocation, got %v", n)
	}
}

func TestTinyAlloc(t *testing.T) {
	const N = 16
	var v [N]unsafe.Pointer
	for i := range v {
		v[i] = unsafe.Pointer(new(byte))
	}

	chunks := make(map[uintptr]bool, N)
	for _, p := range v {
		chunks[uintptr(p)&^7] = true
	}

	if len(chunks) == N {
		t.Fatal("no bytes allocated within the same 8-byte chunk")
	}
}

var mallocSink uintptr

func BenchmarkMalloc8(b *testing.B) {
	var x uintptr
	for i := 0; i < b.N; i++ {
		p := new(int64)
		x ^= uintptr(unsafe.Pointer(p))
	}
	mallocSink = x
}

func BenchmarkMalloc16(b *testing.B) {
	var x uintptr
	for i := 0; i < b.N; i++ {
		p := new([2]int64)
		x ^= uintptr(unsafe.Pointer(p))
	}
	mallocSink = x
}

func BenchmarkMallocTypeInfo8(b *testing.B) {
	var x uintptr
	for i := 0; i < b.N; i++ {
		p := new(struct {
			p [8 / unsafe.Sizeof(uintptr(0))]*int
		})
		x ^= uintptr(unsafe.Pointer(p))
	}
	mallocSink = x
}

func BenchmarkMallocTypeInfo16(b *testing.B) {
	var x uintptr
	for i := 0; i < b.N; i++ {
		p := new(struct {
			p [16 / unsafe.Sizeof(uintptr(0))]*int
		})
		x ^= uintptr(unsafe.Pointer(p))
	}
	mallocSink = x
}

type LargeStruct struct {
	x [16][]byte
}

func BenchmarkMallocLargeStruct(b *testing.B) {
	var x uintptr
	for i := 0; i < b.N; i++ {
		p := make([]LargeStruct, 2)
		x ^= uintptr(unsafe.Pointer(&p[0]))
	}
	mallocSink = x
}

var n = flag.Int("n", 1000, "number of goroutines")

func BenchmarkGoroutineSelect(b *testing.B) {
	quit := make(chan struct{})
	read := func(ch chan struct{}) {
		for {
			select {
			case _, ok := <-ch:
				if !ok {
					return
				}
			case <-quit:
				return
			}
		}
	}
	benchHelper(b, *n, read)
}

func BenchmarkGoroutineBlocking(b *testing.B) {
	read := func(ch chan struct{}) {
		for {
			if _, ok := <-ch; !ok {
				return
			}
		}
	}
	benchHelper(b, *n, read)
}

func BenchmarkGoroutineForRange(b *testing.B) {
	read := func(ch chan struct{}) {
		for range ch {
		}
	}
	benchHelper(b, *n, read)
}

func benchHelper(b *testing.B, n int, read func(chan struct{})) {
	m := make([]chan struct{}, n)
	for i := range m {
		m[i] = make(chan struct{}, 1)
		go read(m[i])
	}
	b.StopTimer()
	b.ResetTimer()
	GC()

	for i := 0; i < b.N; i++ {
		for _, ch := range m {
			if ch != nil {
				ch <- struct{}{}
			}
		}
		time.Sleep(10 * time.Millisecond)
		b.StartTimer()
		GC()
		b.StopTimer()
	}

	for _, ch := range m {
		close(ch)
	}
	time.Sleep(10 * time.Millisecond)
}

func BenchmarkGoroutineIdle(b *testing.B) {
	quit := make(chan struct{})
	fn := func() {
		<-quit
	}
	for i := 0; i < *n; i++ {
		go fn()
	}

	GC()
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		GC()
	}

	b.StopTimer()
	close(quit)
	time.Sleep(10 * time.Millisecond)
}
<commit_msg>runtime: improve TestMemStats checks<commit_after>// Copyright 2013 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 runtime_test

import (
	"flag"
	"fmt"
	"reflect"
	. "runtime"
	"testing"
	"time"
	"unsafe"
)

func TestMemStats(t *testing.T) {
	// Make sure there's at least one forced GC.
	GC()

	// Test that MemStats has sane values.
	st := new(MemStats)
	ReadMemStats(st)

	nz := func(x interface{}) error {
		if x != reflect.Zero(reflect.TypeOf(x)).Interface() {
			return nil
		}
		return fmt.Errorf("zero value")
	}
	le := func(thresh float64) func(interface{}) error {
		return func(x interface{}) error {
			if reflect.ValueOf(x).Convert(reflect.TypeOf(thresh)).Float() < thresh {
				return nil
			}
			return fmt.Errorf("insanely high value (overflow?); want <= %v", thresh)
		}
	}
	eq := func(x interface{}) func(interface{}) error {
		return func(y interface{}) error {
			if x == y {
				return nil
			}
			return fmt.Errorf("want %v", x)
		}
	}
	// Of the uint fields, HeapReleased, HeapIdle can be 0.
	// PauseTotalNs can be 0 if timer resolution is poor.
	//
	// TODO: Test that GCCPUFraction is <= 0.99. This currently
	// fails on windows/386. (Issue #19319)
	fields := map[string][]func(interface{}) error{
		"Alloc": {nz, le(1e10)}, "TotalAlloc": {nz, le(1e11)}, "Sys": {nz, le(1e10)},
		"Lookups": {nz, le(1e10)}, "Mallocs": {nz, le(1e10)}, "Frees": {nz, le(1e10)},
		"HeapAlloc": {nz, le(1e10)}, "HeapSys": {nz, le(1e10)}, "HeapIdle": {le(1e10)},
		"HeapInuse": {nz, le(1e10)}, "HeapReleased": {le(1e10)}, "HeapObjects": {nz, le(1e10)},
		"StackInuse": {nz, le(1e10)}, "StackSys": {nz, le(1e10)},
		"MSpanInuse": {nz, le(1e10)}, "MSpanSys": {nz, le(1e10)},
		"MCacheInuse": {nz, le(1e10)}, "MCacheSys": {nz, le(1e10)},
		"BuckHashSys": {nz, le(1e10)}, "GCSys": {nz, le(1e10)}, "OtherSys": {nz, le(1e10)},
		"NextGC": {nz, le(1e10)}, "LastGC": {nz},
		"PauseTotalNs": {le(1e11)}, "PauseNs": nil, "PauseEnd": nil,
		"NumGC": {nz, le(1e9)}, "NumForcedGC": {nz, le(1e9)},
		"GCCPUFraction": nil, "EnableGC": {eq(true)}, "DebugGC": {eq(false)},
		"BySize": nil,
	}

	rst := reflect.ValueOf(st).Elem()
	for i := 0; i < rst.Type().NumField(); i++ {
		name, val := rst.Type().Field(i).Name, rst.Field(i).Interface()
		checks, ok := fields[name]
		if !ok {
			t.Errorf("unknown MemStats field %s", name)
			continue
		}
		for _, check := range checks {
			if err := check(val); err != nil {
				t.Errorf("%s = %v: %s", name, val, err)
			}
		}
	}

	if st.Sys != st.HeapSys+st.StackSys+st.MSpanSys+st.MCacheSys+
		st.BuckHashSys+st.GCSys+st.OtherSys {
		t.Fatalf("Bad sys value: %+v", *st)
	}

	if st.HeapIdle+st.HeapInuse != st.HeapSys {
		t.Fatalf("HeapIdle(%d) + HeapInuse(%d) should be equal to HeapSys(%d), but isn't.", st.HeapIdle, st.HeapInuse, st.HeapSys)
	}

	if lpe := st.PauseEnd[int(st.NumGC+255)%len(st.PauseEnd)]; st.LastGC != lpe {
		t.Fatalf("LastGC(%d) != last PauseEnd(%d)", st.LastGC, lpe)
	}

	var pauseTotal uint64
	for _, pause := range st.PauseNs {
		pauseTotal += pause
	}
	if int(st.NumGC) < len(st.PauseNs) {
		// We have all pauses, so this should be exact.
		if st.PauseTotalNs != pauseTotal {
			t.Fatalf("PauseTotalNs(%d) != sum PauseNs(%d)", st.PauseTotalNs, pauseTotal)
		}
		for i := int(st.NumGC); i < len(st.PauseNs); i++ {
			if st.PauseNs[i] != 0 {
				t.Fatalf("Non-zero PauseNs[%d]: %+v", i, st)
			}
			if st.PauseEnd[i] != 0 {
				t.Fatalf("Non-zero PauseEnd[%d]: %+v", i, st)
			}
		}
	} else {
		if st.PauseTotalNs < pauseTotal {
			t.Fatalf("PauseTotalNs(%d) < sum PauseNs(%d)", st.PauseTotalNs, pauseTotal)
		}
	}

	if st.NumForcedGC > st.NumGC {
		t.Fatalf("NumForcedGC(%d) > NumGC(%d)", st.NumForcedGC, st.NumGC)
	}
}

func TestStringConcatenationAllocs(t *testing.T) {
	n := testing.AllocsPerRun(1e3, func() {
		b := make([]byte, 10)
		for i := 0; i < 10; i++ {
			b[i] = byte(i) + '0'
		}
		s := "foo" + string(b)
		if want := "foo0123456789"; s != want {
			t.Fatalf("want %v, got %v", want, s)
		}
	})
	// Only string concatenation allocates.
	if n != 1 {
		t.Fatalf("want 1 allocation, got %v", n)
	}
}

func TestTinyAlloc(t *testing.T) {
	const N = 16
	var v [N]unsafe.Pointer
	for i := range v {
		v[i] = unsafe.Pointer(new(byte))
	}

	chunks := make(map[uintptr]bool, N)
	for _, p := range v {
		chunks[uintptr(p)&^7] = true
	}

	if len(chunks) == N {
		t.Fatal("no bytes allocated within the same 8-byte chunk")
	}
}

var mallocSink uintptr

func BenchmarkMalloc8(b *testing.B) {
	var x uintptr
	for i := 0; i < b.N; i++ {
		p := new(int64)
		x ^= uintptr(unsafe.Pointer(p))
	}
	mallocSink = x
}

func BenchmarkMalloc16(b *testing.B) {
	var x uintptr
	for i := 0; i < b.N; i++ {
		p := new([2]int64)
		x ^= uintptr(unsafe.Pointer(p))
	}
	mallocSink = x
}

func BenchmarkMallocTypeInfo8(b *testing.B) {
	var x uintptr
	for i := 0; i < b.N; i++ {
		p := new(struct {
			p [8 / unsafe.Sizeof(uintptr(0))]*int
		})
		x ^= uintptr(unsafe.Pointer(p))
	}
	mallocSink = x
}

func BenchmarkMallocTypeInfo16(b *testing.B) {
	var x uintptr
	for i := 0; i < b.N; i++ {
		p := new(struct {
			p [16 / unsafe.Sizeof(uintptr(0))]*int
		})
		x ^= uintptr(unsafe.Pointer(p))
	}
	mallocSink = x
}

type LargeStruct struct {
	x [16][]byte
}

func BenchmarkMallocLargeStruct(b *testing.B) {
	var x uintptr
	for i := 0; i < b.N; i++ {
		p := make([]LargeStruct, 2)
		x ^= uintptr(unsafe.Pointer(&p[0]))
	}
	mallocSink = x
}

var n = flag.Int("n", 1000, "number of goroutines")

func BenchmarkGoroutineSelect(b *testing.B) {
	quit := make(chan struct{})
	read := func(ch chan struct{}) {
		for {
			select {
			case _, ok := <-ch:
				if !ok {
					return
				}
			case <-quit:
				return
			}
		}
	}
	benchHelper(b, *n, read)
}

func BenchmarkGoroutineBlocking(b *testing.B) {
	read := func(ch chan struct{}) {
		for {
			if _, ok := <-ch; !ok {
				return
			}
		}
	}
	benchHelper(b, *n, read)
}

func BenchmarkGoroutineForRange(b *testing.B) {
	read := func(ch chan struct{}) {
		for range ch {
		}
	}
	benchHelper(b, *n, read)
}

func benchHelper(b *testing.B, n int, read func(chan struct{})) {
	m := make([]chan struct{}, n)
	for i := range m {
		m[i] = make(chan struct{}, 1)
		go read(m[i])
	}
	b.StopTimer()
	b.ResetTimer()
	GC()

	for i := 0; i < b.N; i++ {
		for _, ch := range m {
			if ch != nil {
				ch <- struct{}{}
			}
		}
		time.Sleep(10 * time.Millisecond)
		b.StartTimer()
		GC()
		b.StopTimer()
	}

	for _, ch := range m {
		close(ch)
	}
	time.Sleep(10 * time.Millisecond)
}

func BenchmarkGoroutineIdle(b *testing.B) {
	quit := make(chan struct{})
	fn := func() {
		<-quit
	}
	for i := 0; i < *n; i++ {
		go fn()
	}

	GC()
	b.ResetTimer()

	for i := 0; i < b.N; i++ {
		GC()
	}

	b.StopTimer()
	close(quit)
	time.Sleep(10 * time.Millisecond)
}
<|endoftext|>


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

package service

import (
	"fmt"
	"os"
	"reflect"
	"sync"
	"time"

	"go.uber.org/fx/core/config"
	"go.uber.org/fx/core/ulog"
	"go.uber.org/fx/internal/util"
)

const (
	defaultStartupWait = time.Second
)

type host struct {
	serviceCore
	locked   bool
	observer Observer
	modules  []Module
	roles    map[string]bool

	// Shutdown fields.
	shutdownMu     sync.Mutex
	inShutdown     bool  // Protected by shutdownMu
	shutdownReason *Exit // Protected by shutdownMu
	closeChan      chan Exit
	started        bool
}

// TODO(glib): host is both an Owner and a Host?
var _ Host = &host{}
var _ Owner = &host{}

func (s *host) addModule(module Module) error {
	if s.locked {
		return fmt.Errorf("can't add module: service already started")
	}
	s.modules = append(s.modules, module)
	return module.Initialize(s)
}

func (s *host) supportsRole(roles ...string) bool {
	if len(s.roles) == 0 || len(roles) == 0 {
		return true
	}

	for _, role := range roles {
		if found, ok := s.roles[role]; found && ok {
			return true
		}
	}
	return false
}

func (s *host) Modules() []Module {
	mods := make([]Module, len(s.modules))
	copy(mods, s.modules)
	return mods
}

func (s *host) IsRunning() bool {
	return s.closeChan != nil
}

func (s *host) OnCriticalError(err error) {
	shutdown := true
	if s.observer == nil {
		s.Logger().Warn(
			"No observer set to handle lifecycle events. Shutting down.",
			"event", "OnCriticalError",
		)
	} else {
		shutdown = !s.observer.OnCriticalError(err)
	}

	if shutdown {
		if ok, err := s.shutdown(err, "", nil); !ok || err != nil {
			// TODO(ai) verify we flush logs
			s.Logger().Info("Problem shutting down module", "success", ok, "error", err)
		}
	}
}

func (s *host) shutdown(err error, reason string, exitCode *int) (bool, error) {
	s.shutdownMu.Lock()
	s.inShutdown = true
	defer func() {
		s.inShutdown = false
		s.shutdownMu.Unlock()
	}()

	if s.shutdownReason != nil || !s.IsRunning() {
		return false, nil
	}

	s.shutdownReason = &Exit{
		Reason:   reason,
		Error:    err,
		ExitCode: 0,
	}

	if err != nil {
		if reason != "" {
			s.shutdownReason.Reason = err.Error()
		}
		s.shutdownReason.ExitCode = 1
	}

	if exitCode != nil {
		s.shutdownReason.ExitCode = *exitCode
	}

	// Log the module shutdown errors
	errs := s.stopModules()
	if len(errs) > 0 {
		for k, v := range errs {
			ulog.Logger().Error("Failure to shut down module", "name", k.Name(), "error", v.Error())
		}
	}

	// Stop the metrics reporting
	if s.scope != nil {
		s.scope.Close()
	}

	// report that we shutdown.
	s.closeChan <- *s.shutdownReason
	close(s.closeChan)

	if s.observer != nil {
		s.observer.OnShutdown(*s.shutdownReason)
	}
	return true, err
}

// AddModules adds the given modules to a service host
func (s *host) AddModules(modules ...ModuleCreateFunc) error {
	for _, mcf := range modules {
		mi := ModuleCreateInfo{
			Host:  s,
			Items: make(map[string]interface{}),
		}

		mods, err := mcf(mi)
		if err != nil {
			return err
		}

		if !s.supportsRole(mi.Roles...) {
			s.Logger().Info(
				"module will not be added due to selected roles",
				"roles", mi.Roles,
			)
		}

		for _, mod := range mods {
			err = s.addModule(mod)
		}
	}

	return nil
}

// Start runs the serve loop. If Shutdown() was called then the shutdown reason
// will be returned.
func (s *host) Start(waitForShutdown bool) (<-chan Exit, <-chan struct{}, error) {
	var err error
	s.locked = true
	s.shutdownMu.Lock()
	defer s.shutdownMu.Unlock()

	readyCh := make(chan struct{}, 1)
	defer func() {
		readyCh <- struct{}{}
	}()
	if s.inShutdown {
		return nil, readyCh, fmt.Errorf("errShuttingDown")
	} else if s.IsRunning() {
		return s.closeChan, readyCh, nil
	} else {
		if s.observer != nil {
			if err := s.observer.OnInit(s); err != nil {
				return nil, readyCh, err
			}
		}
		s.shutdownReason = nil
		s.closeChan = make(chan Exit, 1)
		errs := s.startModules()
		if len(errs) > 0 {
			// grab the first error, shut down the service
			// and return the error
			for _, e := range errs {
				errChan := make(chan Exit, 1)
				errChan <- Exit{
					Error:    e,
					Reason:   "Module start failed",
					ExitCode: 4,
				}

				s.shutdownMu.Unlock()
				if _, err := s.shutdown(e, "", nil); err != nil {
					ulog.Logger().Error("Unable to shut down modules", "initialError", e, "shutdownError", err)
				}
				s.shutdownMu.Lock()

				return errChan, readyCh, e
			}
		}
	}

	if waitForShutdown {
		s.WaitForShutdown(nil)
	}

	return s.closeChan, readyCh, err
}

// Stop shuts down the service.
func (s *host) Stop(reason string, exitCode int) error {
	ec := &exitCode
	_, err := s.shutdown(nil, reason, ec)
	return err
}

func (s *host) startModules() map[Module]error {
	results := map[Module]error{}
	wg := sync.WaitGroup{}

	// make sure we wait for all the start
	// calls to return
	wg.Add(len(s.modules))
	for _, mod := range s.modules {
		go func(m Module) {
			if !m.IsRunning() {
				readyCh := make(chan struct{}, 1)
				startResult := m.Start(readyCh)

				select {
				case <-readyCh:
					s.Logger().Info("Module started up cleanly", "module", m.Name())
				case <-time.After(defaultStartupWait):
					results[m] = fmt.Errorf("module didn't start after %v", defaultStartupWait)
				}

				if startError := <-startResult; startError != nil {
					s.Logger().Error("Error received while starting module", "module", m.Name(), "error", startError)
					results[m] = startError
				}
			}
			wg.Done()
		}(mod)
	}

	// wait for the modules to all start
	wg.Wait()
	return results
}

func (s *host) stopModules() map[Module]error {
	results := map[Module]error{}
	wg := sync.WaitGroup{}
	wg.Add(len(s.modules))
	for _, mod := range s.modules {
		go func(m Module) {
			if !m.IsRunning() {
				// TODO: have a timeout here so a bad shutdown
				// doesn't block everyone
				if err := m.Stop(); err != nil {
					results[m] = err
				}
			}
			wg.Done()
		}(mod)
	}
	wg.Wait()
	return results
}

// A ExitCallback is a function to handle a service shutdown and provide
// an exit code
type ExitCallback func(shutdown Exit) int

func (s *host) WaitForShutdown(exitCallback ExitCallback) {
	shutdown := <-s.closeChan
	s.Logger().Info("Shutting down", "reason", shutdown.Reason)

	exit := 0
	if exitCallback != nil {
		exit = exitCallback(shutdown)
	} else if shutdown.Error != nil {
		exit = 1
	}
	os.Exit(exit)
}

func (s *host) transitionState(to State) {
	// TODO(ai) this isn't used yet
	if to < s.state {
		s.Logger().Fatal("Can't down from state", "from", s.state, "to", to, "service", s.Name())
	}

	for s.state < to {
		old := s.state
		newState := s.state
		switch s.state {
		case Uninitialized:
			newState = Initialized
		case Initialized:
			newState = Starting
		case Starting:
			newState = Running
		case Running:
			newState = Stopping
		case Stopping:
			newState = Stopped
		case Stopped:
		}
		s.state = newState
		if s.observer != nil {
			s.observer.OnStateChange(old, newState)
		}
	}
}

const instanceConfigName = "ServiceConfig"

func loadInstanceConfig(cfg config.ConfigurationProvider, key string, instance interface{}) bool {
	fieldName := instanceConfigName
	if field, found := util.FindField(instance, &fieldName, nil); found {

		configValue := reflect.New(field.Type())

		// Try to load the service config
		if cfg.GetValue(key).PopulateStruct(configValue.Interface()) {
			instanceValue := reflect.ValueOf(instance).Elem()
			instanceValue.FieldByName(fieldName).Set(configValue.Elem())
			return true
		}
	}
	return false
}
<commit_msg>Add signal handlers for SIGINT and SIGTERM (#20)<commit_after>// Copyright (c) 2016 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package service

import (
	"fmt"
	"os"
	"os/signal"
	"reflect"
	"sync"
	"syscall"
	"time"

	"go.uber.org/fx/core/config"
	"go.uber.org/fx/core/ulog"
	"go.uber.org/fx/internal/util"
)

const (
	defaultStartupWait = time.Second
)

type host struct {
	serviceCore
	locked   bool
	observer Observer
	modules  []Module
	roles    map[string]bool

	// Shutdown fields.
	shutdownMu     sync.Mutex
	inShutdown     bool  // Protected by shutdownMu
	shutdownReason *Exit // Protected by shutdownMu
	closeChan      chan Exit
	started        bool
}

// TODO(glib): host is both an Owner and a Host?
var _ Host = &host{}
var _ Owner = &host{}

func (s *host) addModule(module Module) error {
	if s.locked {
		return fmt.Errorf("can't add module: service already started")
	}
	s.modules = append(s.modules, module)
	return module.Initialize(s)
}

func (s *host) supportsRole(roles ...string) bool {
	if len(s.roles) == 0 || len(roles) == 0 {
		return true
	}

	for _, role := range roles {
		if found, ok := s.roles[role]; found && ok {
			return true
		}
	}
	return false
}

func (s *host) Modules() []Module {
	mods := make([]Module, len(s.modules))
	copy(mods, s.modules)
	return mods
}

func (s *host) IsRunning() bool {
	return s.closeChan != nil
}

func (s *host) OnCriticalError(err error) {
	shutdown := true
	if s.observer == nil {
		s.Logger().Warn(
			"No observer set to handle lifecycle events. Shutting down.",
			"event", "OnCriticalError",
		)
	} else {
		shutdown = !s.observer.OnCriticalError(err)
	}

	if shutdown {
		if ok, err := s.shutdown(err, "", nil); !ok || err != nil {
			// TODO(ai) verify we flush logs
			s.Logger().Info("Problem shutting down module", "success", ok, "error", err)
		}
	}
}

func (s *host) shutdown(err error, reason string, exitCode *int) (bool, error) {
	s.shutdownMu.Lock()
	s.inShutdown = true
	defer func() {
		s.inShutdown = false
		s.shutdownMu.Unlock()
	}()

	if s.shutdownReason != nil || !s.IsRunning() {
		return false, nil
	}

	s.shutdownReason = &Exit{
		Reason:   reason,
		Error:    err,
		ExitCode: 0,
	}

	if err != nil {
		if reason != "" {
			s.shutdownReason.Reason = err.Error()
		}
		s.shutdownReason.ExitCode = 1
	}

	if exitCode != nil {
		s.shutdownReason.ExitCode = *exitCode
	}

	// Log the module shutdown errors
	errs := s.stopModules()
	if len(errs) > 0 {
		for k, v := range errs {
			ulog.Logger().Error("Failure to shut down module", "name", k.Name(), "error", v.Error())
		}
	}

	// Stop the metrics reporting
	if s.scope != nil {
		s.scope.Close()
	}

	// report that we shutdown.
	s.closeChan <- *s.shutdownReason
	close(s.closeChan)

	if s.observer != nil {
		s.observer.OnShutdown(*s.shutdownReason)
	}
	return true, err
}

// AddModules adds the given modules to a service host
func (s *host) AddModules(modules ...ModuleCreateFunc) error {
	for _, mcf := range modules {
		mi := ModuleCreateInfo{
			Host:  s,
			Items: make(map[string]interface{}),
		}

		mods, err := mcf(mi)
		if err != nil {
			return err
		}

		if !s.supportsRole(mi.Roles...) {
			s.Logger().Info(
				"module will not be added due to selected roles",
				"roles", mi.Roles,
			)
		}

		for _, mod := range mods {
			err = s.addModule(mod)
		}
	}

	return nil
}

// Start runs the serve loop. If Shutdown() was called then the shutdown reason
// will be returned.
func (s *host) Start(waitForShutdown bool) (<-chan Exit, <-chan struct{}, error) {
	var err error
	s.locked = true
	s.shutdownMu.Lock()

	readyCh := make(chan struct{}, 1)
	defer func() {
		readyCh <- struct{}{}
	}()
	if s.inShutdown {
		return nil, readyCh, fmt.Errorf("errShuttingDown")
	} else if s.IsRunning() {
		return s.closeChan, readyCh, nil
	} else {
		if s.observer != nil {
			if err := s.observer.OnInit(s); err != nil {
				s.shutdownMu.Unlock()
				return nil, readyCh, err
			}
		}
		s.shutdownReason = nil
		s.closeChan = make(chan Exit, 1)
		errs := s.startModules()
		if len(errs) > 0 {
			// grab the first error, shut down the service
			// and return the error
			for _, e := range errs {
				errChan := make(chan Exit, 1)
				errChan <- Exit{
					Error:    e,
					Reason:   "Module start failed",
					ExitCode: 4,
				}

				s.shutdownMu.Unlock()
				if _, err := s.shutdown(e, "", nil); err != nil {
					ulog.Logger().Error("Unable to shut down modules", "initialError", e, "shutdownError", err)
				}
				return errChan, readyCh, e
			}
		}
	}

	s.shutdownMu.Unlock()
	s.registerSignalHandlers()

	if waitForShutdown {
		s.WaitForShutdown(nil)
	}

	return s.closeChan, readyCh, err
}

func (s *host) registerSignalHandlers() {
	ch := make(chan os.Signal)
	signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		sig := <-ch
		s.Logger().Warn("Received shutdown signal", "signal", sig.String())
		if err := s.Stop("Received syscall", 0); err != nil {
			s.Logger().Error("Error shutting down", "error", err.Error())
		}
	}()
}

// Stop shuts down the service.
func (s *host) Stop(reason string, exitCode int) error {
	ec := &exitCode
	_, err := s.shutdown(nil, reason, ec)
	return err
}

func (s *host) startModules() map[Module]error {
	results := map[Module]error{}
	wg := sync.WaitGroup{}

	// make sure we wait for all the start
	// calls to return
	wg.Add(len(s.modules))
	for _, mod := range s.modules {
		go func(m Module) {
			if !m.IsRunning() {
				readyCh := make(chan struct{}, 1)
				startResult := m.Start(readyCh)

				select {
				case <-readyCh:
					s.Logger().Info("Module started up cleanly", "module", m.Name())
				case <-time.After(defaultStartupWait):
					results[m] = fmt.Errorf("module didn't start after %v", defaultStartupWait)
				}

				if startError := <-startResult; startError != nil {
					s.Logger().Error("Error received while starting module", "module", m.Name(), "error", startError)
					results[m] = startError
				}
			}
			wg.Done()
		}(mod)
	}

	// wait for the modules to all start
	wg.Wait()
	return results
}

func (s *host) stopModules() map[Module]error {
	results := map[Module]error{}
	wg := sync.WaitGroup{}
	wg.Add(len(s.modules))
	for _, mod := range s.modules {
		go func(m Module) {
			if !m.IsRunning() {
				// TODO: have a timeout here so a bad shutdown
				// doesn't block everyone
				if err := m.Stop(); err != nil {
					results[m] = err
				}
			}
			wg.Done()
		}(mod)
	}
	wg.Wait()
	return results
}

// A ExitCallback is a function to handle a service shutdown and provide
// an exit code
type ExitCallback func(shutdown Exit) int

func (s *host) WaitForShutdown(exitCallback ExitCallback) {
	shutdown := <-s.closeChan
	s.Logger().Info("Shutting down", "reason", shutdown.Reason)

	exit := 0
	if exitCallback != nil {
		exit = exitCallback(shutdown)
	} else if shutdown.Error != nil {
		exit = 1
	}
	os.Exit(exit)
}

func (s *host) transitionState(to State) {
	// TODO(ai) this isn't used yet
	if to < s.state {
		s.Logger().Fatal("Can't down from state", "from", s.state, "to", to, "service", s.Name())
	}

	for s.state < to {
		old := s.state
		newState := s.state
		switch s.state {
		case Uninitialized:
			newState = Initialized
		case Initialized:
			newState = Starting
		case Starting:
			newState = Running
		case Running:
			newState = Stopping
		case Stopping:
			newState = Stopped
		case Stopped:
		}
		s.state = newState
		if s.observer != nil {
			s.observer.OnStateChange(old, newState)
		}
	}
}

const instanceConfigName = "ServiceConfig"

func loadInstanceConfig(cfg config.ConfigurationProvider, key string, instance interface{}) bool {
	fieldName := instanceConfigName
	if field, found := util.FindField(instance, &fieldName, nil); found {

		configValue := reflect.New(field.Type())

		// Try to load the service config
		if cfg.GetValue(key).PopulateStruct(configValue.Interface()) {
			instanceValue := reflect.ValueOf(instance).Elem()
			instanceValue.FieldByName(fieldName).Set(configValue.Elem())
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>package scheduler

import (
	"container/heap"
	"fmt"
)

var _ = fmt.Println

const (
	defaultQueueSize = 1024
)

type Result struct {
	Id     int
	Demand int
	Alloc  map[int]int
}

type Scheduler struct {
	TotalResource int // total resource count
	FreeResource  int

	Resource map[int]int // single resource, nodeId -> resourceNumber

	NodeList []int // nodelist
	NextNode int   // index for RR

	// request heap
	Requestheap *RequestHeap

	ResourceChan chan *Resource
	RequestChan  chan *Request

	ScheduleResult chan *Result

	TerminateChan chan bool
}

func New() *Scheduler {
	s := &Scheduler{
		Resource:       make(map[int]int),
		Requestheap:    new(RequestHeap),
		ResourceChan:   make(chan *Resource, defaultQueueSize),
		RequestChan:    make(chan *Request, defaultQueueSize),
		TerminateChan:  make(chan bool),
		ScheduleResult: make(chan *Result, defaultQueueSize),
	}

	heap.Init(s.Requestheap)

	return s
}

// Start the event-loop
func (s *Scheduler) Start() {
	for {
		select {
		case res := <-s.ResourceChan:
			s.handleNewResource(res)
		case req := <-s.RequestChan:
			s.handleNewRequest(req)
		case <-s.TerminateChan:
			// terminate the scheduler
			return
		}
	}
}

// Nonblocking start
func (s *Scheduler) GoStart() {
	go s.Start()
}

func (s *Scheduler) handleNewResource(res *Resource) {
	// add new node
	if _, ok := s.Resource[res.nodeId]; !ok {
		s.NodeList = append(s.NodeList, res.nodeId)
	}

	// increase the resource count
	s.Resource[res.nodeId] = s.Resource[res.nodeId] + res.resourceCount
	s.TotalResource = s.TotalResource + res.resourceCount
	s.FreeResource = s.FreeResource + res.resourceCount

	s.schedule()
}

func (s *Scheduler) handleNewRequest(req *Request) {
	heap.Push(s.Requestheap, req)
}

// return a list of node: resource
func (s *Scheduler) schedule() {
	if s.FreeResource == 0 || s.Requestheap.Len() == 0 {
		return
	}

	req := s.Requestheap.Peek()
	if s.FreeResource < req.Demand { // cluster full
		return
	}

	scheduledResourceMap := make(map[int]int)
	s.GetResource(req.Demand, scheduledResourceMap)

	heap.Pop(s.Requestheap)

	// sendback result
	s.ScheduleResult <- &Result{req.id, req.Demand, scheduledResourceMap}
}

func (s *Scheduler) GetResource(demand int, scheduledResourceMap map[int]int) {
	// TODO: wait a while for locality, Now just RR
	for {
		nodeId := s.NodeList[s.NextNode]

		if s.Resource[nodeId] > 0 {
			delta := Min(s.Resource[nodeId], demand)

			s.Resource[nodeId] = s.Resource[nodeId] - delta
			s.FreeResource = s.FreeResource - delta
			demand = demand - delta

			scheduledResourceMap[nodeId] = delta
			if demand == 0 {
				break
			}
		}
		s.NextNode = (s.NextNode + 1) % len(s.NodeList) // TODO: delete empty nodes
	}
}

func (s *Scheduler) Terminate() {
	close(s.TerminateChan)
}

// return the min of the two value
func Min(a, b int) int {
	if a < b {
		return a
	}
	return b
}
<commit_msg>more readable..<commit_after>package scheduler

import (
	"container/heap"
	"container/list"
	"fmt"
)

var _ = fmt.Println

const (
	defaultQueueSize = 1024
)

type Result struct {
	Id     int
	Demand int
	Alloc  map[int]int
}

type NodeInfo struct {
	NodeId        int
	ResourceCount int
}

func (n *NodeInfo) increaseResource(quota int) {
	n.ResourceCount = n.ResourceCount + quota
}

func (n *NodeInfo) decreaseResource(quota int) {
	n.ResourceCount = n.ResourceCount - quota
}

func (n *NodeInfo) inFreeNodes() bool {
	return n.ResourceCount != 0
}

type Scheduler struct {
	TotalResource int // total resource count
	FreeResource  int

	NodeInfos map[int]*NodeInfo // NodeId -> NodeInfo
	FreeNodes *list.List

	// request heap
	Requestheap *RequestHeap

	// signal channels
	ResourceChan   chan *Resource
	RequestChan    chan *Request
	ScheduleResult chan *Result
	TerminateChan  chan bool
}

func New() *Scheduler {
	s := &Scheduler{
		NodeInfos:      make(map[int]*NodeInfo),
		FreeNodes:      list.New(),
		Requestheap:    new(RequestHeap),
		ResourceChan:   make(chan *Resource, defaultQueueSize),
		RequestChan:    make(chan *Request, defaultQueueSize),
		TerminateChan:  make(chan bool),
		ScheduleResult: make(chan *Result, defaultQueueSize),
	}

	heap.Init(s.Requestheap)

	return s
}

// Start the event-loop
func (s *Scheduler) Start() {
	for {
		select {
		case res := <-s.ResourceChan:
			s.handleNewResource(res)
		case req := <-s.RequestChan:
			s.handleNewRequest(req)
		case <-s.TerminateChan:
			// terminate the scheduler
			return
		}
	}
}

// Nonblocking start
func (s *Scheduler) GoStart() {
	go s.Start()
}

func (s *Scheduler) handleNewResource(res *Resource) {
	s.TotalResource = s.TotalResource + res.resourceCount
	s.FreeResource = s.FreeResource + res.resourceCount

	if _, ok := s.NodeInfos[res.nodeId]; !ok {
		// add the new node to NodeInfos and FreeNodes
		newNode := &NodeInfo{
			NodeId: res.nodeId,
		}
		s.NodeInfos[res.nodeId] = newNode

	}

	node := s.NodeInfos[res.nodeId]

	if !node.inFreeNodes() {
		s.FreeNodes.PushBack(node.NodeId)
	}
	node.increaseResource(res.resourceCount)

	s.schedule()
}

func (s *Scheduler) handleNewRequest(req *Request) {
	heap.Push(s.Requestheap, req)
}

// return a list of node: resource
func (s *Scheduler) schedule() {
	if s.FreeResource == 0 || s.Requestheap.Len() == 0 {
		return
	}

	req := s.Requestheap.Peek()
	if s.FreeResource < req.Demand { // cluster full
		return
	}

	// create a map for storing allocation result
	result := make(map[int]int)
	s.GetResource(req.Demand, result)

	heap.Pop(s.Requestheap)

	// sendback result
	s.ScheduleResult <- &Result{req.id, req.Demand, result}
}

func (s *Scheduler) GetResource(demand int, result map[int]int) {
	// TODO: wait a while for locality, Now just RR
	for {
		e := s.FreeNodes.Front()
		nodeId := e.Value.(int)

		quota := Min(s.NodeInfos[nodeId].ResourceCount, demand)

		s.NodeInfos[nodeId].decreaseResource(quota)
		s.FreeResource = s.FreeResource - quota
		demand = demand - quota

		result[nodeId] = quota
		if demand == 0 {
			break
		}
		s.FreeNodes.Remove(e)
	}
}

func (s *Scheduler) Terminate() {
	close(s.TerminateChan)
}

// return the min of the two value
func Min(a, b int) int {
	if a < b {
		return a
	}
	return b
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"os/exec"

	"github.com/flynn/flynn/Godeps/_workspace/src/github.com/flynn/go-docopt"
	"github.com/flynn/flynn/controller/client"
	ct "github.com/flynn/flynn/controller/types"
)

func init() {
	register("create", runCreate, `
usage: flynn create [<name>]

Create an application in Flynn.

If a name is not provided, a random name will be generated.

If run from a git repository, a 'flynn' remote will be created or replaced that
allows deploying the application via git.

Examples:

	$ flynn create
	Created turkeys-stupefy-perry
`)

	register("delete", runDelete, `
usage: flynn delete

Delete an app.

If run from a git repository with a 'flynn' remote for the app, it will be
removed.

Examples:

	$ flynn -a turkeys-stupefy-perry delete
	Are you sure you want to delete the app "turkeys-stupefy-perry"? (yes/no): yes
	Deleted turkeys-stupefy-perry
`)
	register("apps", runApps, `
usage: flynn apps

List all apps.

Examples:

	$ flynn apps
	ID                                NAME
	f1e85f5392454a329929e3f27f7a5644  gitreceive
	4c6325c1f13547059e5496c91a6a97dd  router
	8cfd94d040b14bd8aecc086c8f5f5e0d  blobstore
	f488cfb478f54edea497bf6347c2eb80  postgres
	9d5be7be873c41b9898032c08aa87597  controller
`)
}

func runCreate(args *docopt.Args, client *controller.Client) error {
	app := &ct.App{}
	app.Name = args.String["<name>"]

	if err := client.CreateApp(app); err != nil {
		return err
	}

	exec.Command("git", "remote", "remove", "flynn").Run()
	exec.Command("git", "remote", "add", "flynn", gitURLPre(clusterConf.GitHost)+app.Name+gitURLSuf).Run()
	log.Printf("Created %s", app.Name)
	return nil
}

func runDelete(args *docopt.Args, client *controller.Client) error {
	appName := mustApp()

	fmt.Printf("Are you sure you want to delete the app %q? (yes/no): ", appName)
loop:
	for {
		var answer string
		fmt.Scanln(&answer)
		switch answer {
		case "y", "yes":
			break loop
		case "n", "no":
			return nil
		default:
			fmt.Print("Please type 'yes' or 'no': ")
		}
	}

	if err := client.DeleteApp(appName); err != nil {
		return err
	}

	if remotes, err := gitRemotes(); err == nil {
		if app, ok := remotes["flynn"]; ok && app.Name == appName {
			exec.Command("git", "remote", "remove", "flynn").Run()
		}
	}

	log.Printf("Deleted %s", appName)
	return nil
}

func runApps(args *docopt.Args, client *controller.Client) error {
	apps, err := client.AppList()
	if err != nil {
		return err
	}

	w := tabWriter()
	defer w.Flush()

	listRec(w, "ID", "NAME")
	for _, a := range apps {
		listRec(w, a.ID, a.Name)
	}
	return nil
}
<commit_msg>cli: Add a --yes option to `app delete`.<commit_after>package main

import (
	"fmt"
	"log"
	"os/exec"

	"github.com/flynn/flynn/Godeps/_workspace/src/github.com/flynn/go-docopt"
	"github.com/flynn/flynn/controller/client"
	ct "github.com/flynn/flynn/controller/types"
)

func init() {
	register("create", runCreate, `
usage: flynn create [<name>]

Create an application in Flynn.

If a name is not provided, a random name will be generated.

If run from a git repository, a 'flynn' remote will be created or replaced that
allows deploying the application via git.

Examples:

	$ flynn create
	Created turkeys-stupefy-perry
`)

	register("delete", runDelete, `
usage: flynn delete [-y]

Delete an app.

If run from a git repository with a 'flynn' remote for the app, it will be
removed.

Options:
	-y, --yes  Skip the confirmation prompt.

Examples:

	$ flynn -a turkeys-stupefy-perry delete
	Are you sure you want to delete the app "turkeys-stupefy-perry"? (yes/no): yes
	Deleted turkeys-stupefy-perry
`)
	register("apps", runApps, `
usage: flynn apps

List all apps.

Examples:

	$ flynn apps
	ID                                NAME
	f1e85f5392454a329929e3f27f7a5644  gitreceive
	4c6325c1f13547059e5496c91a6a97dd  router
	8cfd94d040b14bd8aecc086c8f5f5e0d  blobstore
	f488cfb478f54edea497bf6347c2eb80  postgres
	9d5be7be873c41b9898032c08aa87597  controller
`)
}

func runCreate(args *docopt.Args, client *controller.Client) error {
	app := &ct.App{}
	app.Name = args.String["<name>"]

	if err := client.CreateApp(app); err != nil {
		return err
	}

	exec.Command("git", "remote", "remove", "flynn").Run()
	exec.Command("git", "remote", "add", "flynn", gitURLPre(clusterConf.GitHost)+app.Name+gitURLSuf).Run()
	log.Printf("Created %s", app.Name)
	return nil
}

func runDelete(args *docopt.Args, client *controller.Client) error {
	appName := mustApp()

	if !args.Bool["--yes"] {
		fmt.Printf("Are you sure you want to delete the app %q? (yes/no): ", appName)
	loop:
		for {
			var answer string
			fmt.Scanln(&answer)
			switch answer {
			case "y", "yes":
				break loop
			case "n", "no":
				return nil
			default:
				fmt.Print("Please type 'yes' or 'no': ")
			}
		}
	}

	if err := client.DeleteApp(appName); err != nil {
		return err
	}

	if remotes, err := gitRemotes(); err == nil {
		if app, ok := remotes["flynn"]; ok && app.Name == appName {
			exec.Command("git", "remote", "remove", "flynn").Run()
		}
	}

	log.Printf("Deleted %s", appName)
	return nil
}

func runApps(args *docopt.Args, client *controller.Client) error {
	apps, err := client.AppList()
	if err != nil {
		return err
	}

	w := tabWriter()
	defer w.Flush()

	listRec(w, "ID", "NAME")
	for _, a := range apps {
		listRec(w, a.ID, a.Name)
	}
	return nil
}
<|endoftext|>


<commit_before>/* Client functions */

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

package main

import (
	"net/http"
	"encoding/json"
	"github.com/ctdk/goiardi/actor"
	"github.com/ctdk/goiardi/util"
	"fmt"
	"strings"
)

func actor_handler(w http.ResponseWriter, r *http.Request){
	w.Header().Set("Content-Type", "application/json")
	path := SplitPath(r.URL.Path)
	op := path[0]
	client_name := path[1]
	opUser, oerr := actor.GetReqUser(r.Header.Get("X-OPS-USERID"))

	chef_type := strings.TrimSuffix(op, "s")
	/* Make sure we aren't trying anything with a user */
	c_chk, _ := actor.Get(client_name)
	if (c_chk != nil && c_chk.ChefType != chef_type){
		var other_actor_type string
		if chef_type == "client" {
			other_actor_type = "user"
		} else {
			other_actor_type = "client"
		}
		err_user := fmt.Errorf("%s is a %s, not a client.", client_name, chef_type, other_actor_type)
		JsonErrorReport(w, r, err_user.Error(), http.StatusBadRequest)
		return
	}
	switch r.Method {
		case "DELETE":
			if oerr != nil {
				JsonErrorReport(w, r, oerr.Error(), oerr.Status())
				return
			}
			/* no response body here */
			chef_client, err := actor.Get(client_name)
			if err != nil {
				JsonErrorReport(w, r, err.Error(), http.StatusNotFound)
				return
			}
			if !opUser.IsAdmin() && !opUser.IsSelf(chef_client) {
				JsonErrorReport(w, r, "Deleting that client is forbidden", http.StatusForbidden)
				return
			}
			err = chef_client.Delete()
			if err != nil {
				JsonErrorReport(w, r, err.Error(), http.StatusInternalServerError)
				return
			}
			/* Otherwise, we don't actually do anything. */
		case "GET":
			if oerr != nil {
				JsonErrorReport(w, r, oerr.Error(), oerr.Status())
				return
			}
			chef_client, err := actor.Get(client_name)

			if err != nil {
				JsonErrorReport(w, r, err.Error(), http.StatusNotFound)
				return
			}
			if !opUser.IsAdmin() && !opUser.IsSelf(chef_client) {
				JsonErrorReport(w, r, "You are not allowed to perform that action.", http.StatusForbidden)
				return
			}

			/* API docs are wrong here re: public_key vs.
			 * certificate. Also orgname (at least w/ open source)
			 * and clientname, and it wants chef_type and
			 * json_class
			 */
			json_client := map[string]interface{}{
				"name": chef_client.Name, // set same as above
							  // for now
				"chef_type": chef_client.ChefType,
				"admin": chef_client.Admin,
				//"orgname": chef_client.Orgname,
				"public_key": chef_client.PublicKey,
			}
			if op != "users" {
				json_client["validator"] = chef_client.Validator
				json_client["json_class"] = chef_client.JsonClass
			}
			enc := json.NewEncoder(w)
			if err = enc.Encode(&json_client); err != nil{
				JsonErrorReport(w, r, err.Error(), http.StatusInternalServerError)
				return
			}
		case "PUT":
			if oerr != nil {
				JsonErrorReport(w, r, oerr.Error(), oerr.Status())
				return
			}
			client_data, jerr := ParseObjJson(r.Body)
			if jerr != nil {
				JsonErrorReport(w, r, jerr.Error(), http.StatusBadRequest)
				return
			}
			chef_client, err := actor.Get(client_name)
			if err != nil {
				JsonErrorReport(w, r, err.Error(), http.StatusNotFound)
				return
			}

			json_name, sterr := util.ValidateAsString(client_data["name"])
			if sterr != nil {
				JsonErrorReport(w, r, sterr.Error(), http.StatusBadRequest)
				return
			}

			/* If client_name and client_data["name"] aren't the
			 * same, we're renaming. Check the new name doesn't
			 * already exist. */
			json_client := make(map[string]interface{})
			if client_name != json_name {
				err := chef_client.Rename(json_name)
				if err != nil {
					JsonErrorReport(w, r, err.Error(), err.Status())
					return
				} else {
					w.WriteHeader(http.StatusCreated)
				}
			}
			if uerr := chef_client.UpdateFromJson(client_data, chef_type); uerr != nil {
				JsonErrorReport(w, r, uerr.Error(), uerr.Status())
				return
			}

			if pk, pkfound := client_data["public_key"]; pkfound {
				switch pk := pk.(type){
					case string:
						if pkok, pkerr := actor.ValidatePublicKey(pk); !pkok {
							JsonErrorReport(w, r, pkerr.Error(), http.StatusBadRequest)
							return
						}
						chef_client.PublicKey = pk
					case nil:
						//show_public_key = false
						;
					default:
						JsonErrorReport(w, r, "Bad request", http.StatusBadRequest)
						return
				}
			}

			if p, pfound := client_data["private_key"]; pfound {
				switch p := p.(type) {
					case bool:
						if p {
							if json_client["private_key"], err = chef_client.GenerateKeys(); err != nil {
								JsonErrorReport(w, r, err.Error(), http.StatusInternalServerError)
								return
							}
						}
					default:
						JsonErrorReport(w, r, "Bad request", http.StatusBadRequest)
						return
				}
			}

			chef_client.Save()
			json_client["name"] = chef_client.Name
			json_client["admin"] = chef_client.Admin
			if op != "users" {
				json_client["validator"] = chef_client.Validator
				json_client["public_key"] = chef_client.PublicKey
			}
			enc := json.NewEncoder(w)
			if err = enc.Encode(&json_client); err != nil{
				JsonErrorReport(w, r, err.Error(), http.StatusInternalServerError)
				return
			}
		default:
			JsonErrorReport(w, r, "Unrecognized method for client!", http.StatusMethodNotAllowed)
	}
}
<commit_msg>Fix client delete behavior; the docs claimed it had no response body, but knife expects the response to contain the JSON representation of the deleted client<commit_after>/* Client functions */

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

package main

import (
	"net/http"
	"encoding/json"
	"github.com/ctdk/goiardi/actor"
	"github.com/ctdk/goiardi/util"
	"fmt"
	"strings"
)

func actor_handler(w http.ResponseWriter, r *http.Request){
	w.Header().Set("Content-Type", "application/json")
	path := SplitPath(r.URL.Path)
	op := path[0]
	client_name := path[1]
	opUser, oerr := actor.GetReqUser(r.Header.Get("X-OPS-USERID"))

	chef_type := strings.TrimSuffix(op, "s")
	/* Make sure we aren't trying anything with a user */
	c_chk, _ := actor.Get(client_name)
	if (c_chk != nil && c_chk.ChefType != chef_type){
		var other_actor_type string
		if chef_type == "client" {
			other_actor_type = "user"
		} else {
			other_actor_type = "client"
		}
		err_user := fmt.Errorf("%s is a %s, not a client.", client_name, chef_type, other_actor_type)
		JsonErrorReport(w, r, err_user.Error(), http.StatusBadRequest)
		return
	}
	switch r.Method {
		case "DELETE":
			if oerr != nil {
				JsonErrorReport(w, r, oerr.Error(), oerr.Status())
				return
			}
			/* no response body here */
			chef_client, err := actor.Get(client_name)
			if err != nil {
				JsonErrorReport(w, r, err.Error(), http.StatusNotFound)
				return
			}
			if !opUser.IsAdmin() && !opUser.IsSelf(chef_client) {
				JsonErrorReport(w, r, "Deleting that client is forbidden", http.StatusForbidden)
				return
			}
			/* Docs were incorrect. It does want the body of the
			 * deleted object. */
			enc := json.NewEncoder(w)
			if err = enc.Encode(&chef_client); err != nil{
				JsonErrorReport(w, r, err.Error(), http.StatusInternalServerError)
				return
			}
			err = chef_client.Delete()
			if err != nil {
				JsonErrorReport(w, r, err.Error(), http.StatusInternalServerError)
				return
			}

		case "GET":
			if oerr != nil {
				JsonErrorReport(w, r, oerr.Error(), oerr.Status())
				return
			}
			chef_client, err := actor.Get(client_name)

			if err != nil {
				JsonErrorReport(w, r, err.Error(), http.StatusNotFound)
				return
			}
			if !opUser.IsAdmin() && !opUser.IsSelf(chef_client) {
				JsonErrorReport(w, r, "You are not allowed to perform that action.", http.StatusForbidden)
				return
			}

			/* API docs are wrong here re: public_key vs.
			 * certificate. Also orgname (at least w/ open source)
			 * and clientname, and it wants chef_type and
			 * json_class
			 */
			json_client := map[string]interface{}{
				"name": chef_client.Name, // set same as above
							  // for now
				"chef_type": chef_client.ChefType,
				"admin": chef_client.Admin,
				//"orgname": chef_client.Orgname,
				"public_key": chef_client.PublicKey,
			}
			if op != "users" {
				json_client["validator"] = chef_client.Validator
				json_client["json_class"] = chef_client.JsonClass
			}
			enc := json.NewEncoder(w)
			if err = enc.Encode(&json_client); err != nil{
				JsonErrorReport(w, r, err.Error(), http.StatusInternalServerError)
				return
			}
		case "PUT":
			if oerr != nil {
				JsonErrorReport(w, r, oerr.Error(), oerr.Status())
				return
			}
			client_data, jerr := ParseObjJson(r.Body)
			if jerr != nil {
				JsonErrorReport(w, r, jerr.Error(), http.StatusBadRequest)
				return
			}
			chef_client, err := actor.Get(client_name)
			if err != nil {
				JsonErrorReport(w, r, err.Error(), http.StatusNotFound)
				return
			}

			json_name, sterr := util.ValidateAsString(client_data["name"])
			if sterr != nil {
				JsonErrorReport(w, r, sterr.Error(), http.StatusBadRequest)
				return
			}

			/* If client_name and client_data["name"] aren't the
			 * same, we're renaming. Check the new name doesn't
			 * already exist. */
			json_client := make(map[string]interface{})
			if client_name != json_name {
				err := chef_client.Rename(json_name)
				if err != nil {
					JsonErrorReport(w, r, err.Error(), err.Status())
					return
				} else {
					w.WriteHeader(http.StatusCreated)
				}
			}
			if uerr := chef_client.UpdateFromJson(client_data, chef_type); uerr != nil {
				JsonErrorReport(w, r, uerr.Error(), uerr.Status())
				return
			}

			if pk, pkfound := client_data["public_key"]; pkfound {
				switch pk := pk.(type){
					case string:
						if pkok, pkerr := actor.ValidatePublicKey(pk); !pkok {
							JsonErrorReport(w, r, pkerr.Error(), http.StatusBadRequest)
							return
						}
						chef_client.PublicKey = pk
					case nil:
						//show_public_key = false
						;
					default:
						JsonErrorReport(w, r, "Bad request", http.StatusBadRequest)
						return
				}
			}

			if p, pfound := client_data["private_key"]; pfound {
				switch p := p.(type) {
					case bool:
						if p {
							if json_client["private_key"], err = chef_client.GenerateKeys(); err != nil {
								JsonErrorReport(w, r, err.Error(), http.StatusInternalServerError)
								return
							}
						}
					default:
						JsonErrorReport(w, r, "Bad request", http.StatusBadRequest)
						return
				}
			}

			chef_client.Save()
			json_client["name"] = chef_client.Name
			json_client["admin"] = chef_client.Admin
			if op != "users" {
				json_client["validator"] = chef_client.Validator
				json_client["public_key"] = chef_client.PublicKey
			}
			enc := json.NewEncoder(w)
			if err = enc.Encode(&json_client); err != nil{
				JsonErrorReport(w, r, err.Error(), http.StatusInternalServerError)
				return
			}
		default:
			JsonErrorReport(w, r, "Unrecognized method for client!", http.StatusMethodNotAllowed)
	}
}
<|endoftext|>


<commit_before>package main

import(
	"log"
	"net"
)

const(
	RECV_BUF_LEN = 1024
)

func main() {
	log.Print("Starting the server")
	listener, err := net.Listen("tcp", "127.0.0.1:8377")
	if err != nil {
		log.Fatal("Listen: ", err.Error())
	}

	for {
		conn, err := listener.Accept()
		if err != nil {
			log.Print("Accept: ", err.Error())
			return
		}

		go handleConnection(conn)
	}
}

func handleConnection(conn net.Conn) {
	buf := make([]byte, RECV_BUF_LEN)
	n, err := conn.Read(buf)
	if err != nil {
		log.Print("Read: ", err.Error())
		return
	}

	log.Print("Received ", n, " bytes of data")
	_, err = conn.Write(buf)
	if err != nil {
		log.Print("Write: ", err.Error())
		return
	} else {
		log.Print("Reply echoed")
		conn.Close()
	}
}
<commit_msg>Simplify handleConnection() by using io.Copy<commit_after>package main

import(
	"io"
	"log"
	"net"
)

func main() {
	log.Print("Starting the server")
	listener, err := net.Listen("tcp", "127.0.0.1:8377")
	if err != nil {
		log.Fatal(err.Error())
	}

	for {
		conn, err := listener.Accept()
		if err != nil {
			log.Print(err.Error())
			return
		}

		go handleConnection(conn)
	}
}

func handleConnection(conn net.Conn) {
	copied, err := io.Copy(conn, conn)
	if err != nil {
		log.Print(err.Error())
	} else {
		log.Print("Echoed ", copied, " bytes")
	}
	conn.Close()
	log.Print("Connection closed")
}
<|endoftext|>


<commit_before>// Package cmd is essentially the main() of the secret volume service.
// It is a separate package to allow convenient use of Go build tags to control
// debug logging and system calls.
package cmd

import (
	"os"
	"path/filepath"

	"github.com/negz/secret-volume/api"
	"github.com/negz/secret-volume/secrets"
	"github.com/negz/secret-volume/server"
	"github.com/negz/secret-volume/volume"

	"github.com/benschw/srv-lb/dns"
	"github.com/benschw/srv-lb/lb"
	"github.com/benschw/srv-lb/strategy/random"
	"github.com/facebookgo/httpdown"

	"gopkg.in/alecthomas/kingpin.v2"
)

func setupLb(ns, srv string) lb.LoadBalancer {
	var lib dns.Lookup
	if ns == "" {
		// TODO(negz): Handle error if/when https://github.com/benschw/srv-lb/pull/5 is merged
		lib = dns.NewDefaultLookupLib()
	} else {
		lib = dns.NewLookupLib(ns)
	}

	return lb.New(&lb.Config{Dns: lib, Strategy: random.RandomStrategy}, srv)
}

// Run is effectively the main() of the secretvolume binary.
// It lives here in its own package to allow convenient use of Go build tags
// to control debug logging and system calls.
func Run() {
	var (
		app = kingpin.New(filepath.Base(os.Args[0]), "Manages sets of files containing secrets.").DefaultEnvars()

		srv    = app.Flag("talos-srv", "SRV record at which to lookup Talos").String()
		addr   = app.Flag("addr", "Address at which to serve requests (host:port)").Default(":10002").String()
		ns     = app.Flag("ns", "DNS server to use to lookup SRV records (host:port)").String()
		parent = app.Flag("parent", "Directory under which to mount secret volumes").Short('p').Default("/secrets").ExistingDir()
		virt   = app.Flag("virtual", "Use an in-memory filesystem and a no-op parenter for testing").Bool()
		stop   = app.Flag("close-after", "Wait this long at shutdown before closing HTTP connections").Default("1m").Duration()
		kill   = app.Flag("kill-after", "Wait this long at shutdown before exiting").Default("2m").Duration()
	)

	kingpin.MustParse(app.Parse(os.Args[1:]))

	sp, err := secrets.NewTalosProducer(setupLb(*ns, *srv))
	kingpin.FatalIfError(err, "cannot setup Talos secret producer")

	m, fs, err := setupFs(*virt, *parent)
	kingpin.FatalIfError(err, "cannot setup filesystem and parenter")

	sps := map[api.SecretSource]secrets.Producer{api.Talos: sp}
	vm, err := volume.NewManager(m, sps, volume.Filesystem(fs))
	kingpin.FatalIfError(err, "cannot setup secret volume manager")

	handlers, err := server.NewHTTPHandlers(vm)
	kingpin.FatalIfError(err, "cannot setup HTTP handlers")

	hd := &httpdown.HTTP{StopTimeout: *stop, KillTimeout: *kill}
	http := handlers.HTTPServer(*addr)
	kingpin.FatalIfError(httpdown.ListenAndServe(http, hd), "HTTP server error")
}
<commit_msg>Make Talos secret provider optional<commit_after>// Package cmd is essentially the main() of the secret volume service.
// It is a separate package to allow convenient use of Go build tags to control
// debug logging and system calls.
package cmd

import (
	"os"
	"path/filepath"

	"github.com/negz/secret-volume/api"
	"github.com/negz/secret-volume/secrets"
	"github.com/negz/secret-volume/server"
	"github.com/negz/secret-volume/volume"

	"github.com/benschw/srv-lb/dns"
	"github.com/benschw/srv-lb/lb"
	"github.com/benschw/srv-lb/strategy/random"
	"github.com/facebookgo/httpdown"

	"gopkg.in/alecthomas/kingpin.v2"
)

func setupTalosLb(ns, srv string) lb.LoadBalancer {
	var lib dns.Lookup
	if ns == "" {
		// TODO(negz): Handle error if/when https://github.com/benschw/srv-lb/pull/5 is merged
		lib = dns.NewDefaultLookupLib()
	} else {
		lib = dns.NewLookupLib(ns)
	}

	return lb.New(&lb.Config{Dns: lib, Strategy: random.RandomStrategy}, srv)
}

// Run is effectively the main() of the secretvolume binary.
// It lives here in its own package to allow convenient use of Go build tags
// to control debug logging and system calls.
func Run() {
	var (
		app = kingpin.New(filepath.Base(os.Args[0]), "Manages sets of files containing secrets.").DefaultEnvars()

		talos  = app.Flag("talos-srv", "Enables Talos by providing an SRV record at which to find it.").String()
		addr   = app.Flag("addr", "Address at which to serve requests (host:port).").Default(":10002").String()
		ns     = app.Flag("ns", "DNS server to use to lookup SRV records (host:port).").String()
		parent = app.Flag("parent", "Directory under which to mount secret volumes.").Default("/secrets").String()
		virt   = app.Flag("virtual", "Use an in-memory filesystem and a no-op parenter for testing.").Bool()
		stop   = app.Flag("close-after", "Wait this long at shutdown before closing HTTP connections.").Default("1m").Duration()
		kill   = app.Flag("kill-after", "Wait this long at shutdown before exiting.").Default("2m").Duration()
	)

	kingpin.MustParse(app.Parse(os.Args[1:]))

	m, fs, err := setupFs(*virt, *parent)
	kingpin.FatalIfError(err, "cannot setup filesystem and parenter")

	sps := make(map[api.SecretSource]secrets.Producer)
	if *talos != "" {
		sp, terr := secrets.NewTalosProducer(setupTalosLb(*ns, *talos))
		kingpin.FatalIfError(terr, "cannot setup Talos secret producer")
		sps[api.Talos] = sp
	}

	vm, err := volume.NewManager(m, sps, volume.Filesystem(fs))
	kingpin.FatalIfError(err, "cannot setup secret volume manager")

	handlers, err := server.NewHTTPHandlers(vm)
	kingpin.FatalIfError(err, "cannot setup HTTP handlers")

	hd := &httpdown.HTTP{StopTimeout: *stop, KillTimeout: *kill}
	http := handlers.HTTPServer(*addr)
	kingpin.FatalIfError(httpdown.ListenAndServe(http, hd), "HTTP server error")
}
<|endoftext|>


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

package cmd

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net/url"
	"os"
	"regexp"
	"strings"

	"github.com/spf13/cobra"

	"github.com/plusvic/go-ansi"

	"github.com/fatih/color"
	"github.com/spf13/pflag"
	"github.com/spf13/viper"

	"github.com/VirusTotal/vt-cli/utils"
	"github.com/VirusTotal/vt-cli/yaml"
	"github.com/VirusTotal/vt-go/vt"
)

var colorScheme = yaml.Colors{
	KeyColor:     color.New(color.FgYellow),
	ValueColor:   color.New(color.FgHiGreen),
	CommentColor: color.New(color.Faint)}

func addAPIKeyFlag(flags *pflag.FlagSet) {
	flags.StringP(
		"apikey", "k", "",
		"api key")
}

func addHostFlag(flags *pflag.FlagSet) {
	flags.String(
		"host", "www.virustotal.com",
		"API host name")
	flags.MarkHidden("host")
}

func addIncludeExcludeFlags(flags *pflag.FlagSet) {
	flags.StringSliceP(
		"include", "i", []string{"**"},
		"include fields matching the provided pattern")

	flags.StringSliceP(
		"exclude", "x", []string{},
		"exclude fields matching the provided pattern")
}

func addThreadsFlag(flags *pflag.FlagSet) {
	flags.IntP(
		"threads", "t", 5,
		"number of threads working in parallel")
}

func addIDOnlyFlag(flags *pflag.FlagSet) {
	flags.BoolP(
		"identifiers-only", "I", false,
		"print identifiers only")
}

func addLimitFlag(flags *pflag.FlagSet) {
	flags.IntP(
		"limit", "n", 10,
		"maximum number of results")
}

func addCursorFlag(flags *pflag.FlagSet) {
	flags.StringP(
		"cursor", "c", "",
		"cursor")
}

func addOutputFlag(flags *pflag.FlagSet) {
	flags.StringP(
		"output", "o", ".",
		"directory where downloaded files are put")
}

func addFilterFlag(flags *pflag.FlagSet) {
	flags.StringP(
		"filter", "f", "",
		"filter")
}

func addVerboseFlag(flags *pflag.FlagSet) {
	flags.BoolP(
		"verbose", "v", false,
		"verbose output")
}

func addYAMLFlag(flags *pflag.FlagSet) {
	flags.BoolP(
		"yaml", "y", false,
		"output in YAML format")
}

// ReadFile reads the specified file and returns its content. If filename is "-"
// the data is read from stdin.
func ReadFile(filename string) ([]byte, error) {
	if filename == "-" {
		return ioutil.ReadAll(os.Stdin)
	}
	return ioutil.ReadFile(filename)
}

// PrintCommandLineWithCursor prints the same command-line that was used for
// executing the program but adding or replacing the --cursor flag with
// the current cursor for the given iterator.
func PrintCommandLineWithCursor(cmd *cobra.Command, it *vt.Iterator) {
	if cursor := it.Cursor(); cursor != "" {
		args := cmd.Flags().Args()
		cursorFound := false
		for i, arg := range args {
			if arg == "-c" {
				args[i+1] = cursor
				cursorFound = true
				break
			} else if strings.HasPrefix(arg, "--cursor=") {
				args[i] = fmt.Sprintf("--cursor=%s", cursor)
				cursorFound = true
				break
			}
			args[i] = fmt.Sprintf("'%s'", arg)
		}
		if !cursorFound {
			args = append(args, fmt.Sprintf("--cursor=%s", cursor))
		}
		color.New(color.Faint).Fprintf(
			ansi.NewAnsiStderr(), "\nMORE WITH:\n%s %s\n",
			cmd.CommandPath(), strings.Join(args, " "))
	}
}

// NewAPIClient returns a new utils.APIClient with the API key specified via
// command-line argument or config file.
func NewAPIClient() (*utils.APIClient, error) {
	apikey := viper.GetString("apikey")
	if apikey == "" {
		return nil, errors.New("An API key is needed. Either use the --apikey flag or run \"vt init\" to set up your API key")
	}
	return utils.NewAPIClient(apikey, fmt.Sprintf("vt-cli %s", Version))
}

// ObjectPrinter ...
type ObjectPrinter struct {
	client *utils.APIClient
	cmd    *cobra.Command
}

// NewObjectPrinter ...
func NewObjectPrinter(cmd *cobra.Command) (*ObjectPrinter, error) {
	client, err := NewAPIClient()
	if err != nil {
		return nil, err
	}
	return &ObjectPrinter{client: client}, nil
}

// Print ...
func (p *ObjectPrinter) Print(objType string, args []string, argRe *regexp.Regexp) error {

	var r utils.StringReader

	if len(args) == 1 && args[0] == "-" {
		r = utils.NewStringIOReader(os.Stdin)
	} else {
		r = utils.NewStringArrayReader(args)
	}

	if argRe != nil {
		r = utils.NewFilteredStringReader(r, argRe)
	}

	filteredArgs := make([]string, 0)
	for s, err := r.ReadString(); s != "" || err == nil; s, err = r.ReadString() {
		filteredArgs = append(filteredArgs, s)
	}

	objectsCh := make(chan *vt.Object)
	errorsCh := make(chan error, len(filteredArgs))

	go p.client.RetrieveObjects(objType, filteredArgs, objectsCh, errorsCh)

	objs := make([]*vt.Object, 0)

	for obj := range objectsCh {
		if viper.GetBool("identifiers-only") {
			fmt.Printf("%s\n", obj.ID)
		} else {
			objs = append(objs, obj)
		}
	}

	if len(objs) > 0 {
		if err := p.PrintObjects(objs); err != nil {
			return err
		}
	}

	for err := range errorsCh {
		fmt.Fprintln(os.Stderr, err)
	}

	return nil
}

func (p *ObjectPrinter) PrintCollection(collection *url.URL) error {
	it, err := p.client.Iterator(collection,
		vt.IteratorOptions{
			Limit:  viper.GetInt("limit"),
			Cursor: viper.GetString("cursor"),
			Filter: viper.GetString("filter"),
		})
	if err != nil {
		return err
	}
	return p.PrintIter(it)
}

func (p *ObjectPrinter) PrintIter(it *vt.Iterator) error {

	objs := make([]*vt.Object, 0)
	for it.Next() {
		obj := it.Get()
		if viper.GetBool("identifiers-only") {
			fmt.Printf("%s\n", obj.ID)
		} else {
			objs = append(objs, obj)
		}
	}

	if err := it.Error(); err != nil {
		return err
	}

	if len(objs) > 0 {
		if err := p.PrintObjects(objs); err != nil {
			return err
		}
	}

	PrintCommandLineWithCursor(p.cmd, it)
	return nil
}

func (p *ObjectPrinter) PrintObject(obj *vt.Object) error {
	objs := make([]*vt.Object, 1)
	objs[0] = obj
	return p.PrintObjects(objs)
}

func (p *ObjectPrinter) PrintObjects(objs []*vt.Object) error {

	list := make([]map[string]interface{}, 0)

	for _, obj := range objs {
		m := obj.Attributes
		if viper.IsSet("include") && viper.IsSet("exclude") {
			m = utils.FilterMap(
				m, viper.GetStringSlice("include"), viper.GetStringSlice("exclude"))
		}
		for name, r := range obj.Relationships {
			if r.IsOneToOne && len(r.RelatedObjects) > 0 {
				m[name] = r.RelatedObjects[0].ID
			} else {
				l := make([]string, 0)
				for _, obj := range r.RelatedObjects {
					l = append(l, obj.ID)
				}
				m[name] = l
			}
		}
		key := fmt.Sprintf("%s <%s>", obj.Type, obj.ID)
		list = append(list, map[string]interface{}{key: m})
	}

	return yaml.NewColorEncoder(ansi.NewAnsiStdout(), colorScheme).Encode(list)
}
<commit_msg>Fix nil pointer dereference.<commit_after>// Copyright © 2017 The VirusTotal CLI authors. All Rights Reserved.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package cmd

import (
	"errors"
	"fmt"
	"io/ioutil"
	"net/url"
	"os"
	"regexp"
	"strings"

	"github.com/spf13/cobra"

	"github.com/plusvic/go-ansi"

	"github.com/fatih/color"
	"github.com/spf13/pflag"
	"github.com/spf13/viper"

	"github.com/VirusTotal/vt-cli/utils"
	"github.com/VirusTotal/vt-cli/yaml"
	"github.com/VirusTotal/vt-go/vt"
)

var colorScheme = yaml.Colors{
	KeyColor:     color.New(color.FgYellow),
	ValueColor:   color.New(color.FgHiGreen),
	CommentColor: color.New(color.Faint)}

func addAPIKeyFlag(flags *pflag.FlagSet) {
	flags.StringP(
		"apikey", "k", "",
		"api key")
}

func addHostFlag(flags *pflag.FlagSet) {
	flags.String(
		"host", "www.virustotal.com",
		"API host name")
	flags.MarkHidden("host")
}

func addIncludeExcludeFlags(flags *pflag.FlagSet) {
	flags.StringSliceP(
		"include", "i", []string{"**"},
		"include fields matching the provided pattern")

	flags.StringSliceP(
		"exclude", "x", []string{},
		"exclude fields matching the provided pattern")
}

func addThreadsFlag(flags *pflag.FlagSet) {
	flags.IntP(
		"threads", "t", 5,
		"number of threads working in parallel")
}

func addIDOnlyFlag(flags *pflag.FlagSet) {
	flags.BoolP(
		"identifiers-only", "I", false,
		"print identifiers only")
}

func addLimitFlag(flags *pflag.FlagSet) {
	flags.IntP(
		"limit", "n", 10,
		"maximum number of results")
}

func addCursorFlag(flags *pflag.FlagSet) {
	flags.StringP(
		"cursor", "c", "",
		"cursor")
}

func addOutputFlag(flags *pflag.FlagSet) {
	flags.StringP(
		"output", "o", ".",
		"directory where downloaded files are put")
}

func addFilterFlag(flags *pflag.FlagSet) {
	flags.StringP(
		"filter", "f", "",
		"filter")
}

func addVerboseFlag(flags *pflag.FlagSet) {
	flags.BoolP(
		"verbose", "v", false,
		"verbose output")
}

func addYAMLFlag(flags *pflag.FlagSet) {
	flags.BoolP(
		"yaml", "y", false,
		"output in YAML format")
}

// ReadFile reads the specified file and returns its content. If filename is "-"
// the data is read from stdin.
func ReadFile(filename string) ([]byte, error) {
	if filename == "-" {
		return ioutil.ReadAll(os.Stdin)
	}
	return ioutil.ReadFile(filename)
}

// PrintCommandLineWithCursor prints the same command-line that was used for
// executing the program but adding or replacing the --cursor flag with
// the current cursor for the given iterator.
func PrintCommandLineWithCursor(cmd *cobra.Command, it *vt.Iterator) {
	if cursor := it.Cursor(); cursor != "" {
		args := cmd.Flags().Args()
		cursorFound := false
		for i, arg := range args {
			if arg == "-c" {
				args[i+1] = cursor
				cursorFound = true
				break
			} else if strings.HasPrefix(arg, "--cursor=") {
				args[i] = fmt.Sprintf("--cursor=%s", cursor)
				cursorFound = true
				break
			}
			args[i] = fmt.Sprintf("'%s'", arg)
		}
		if !cursorFound {
			args = append(args, fmt.Sprintf("--cursor=%s", cursor))
		}
		color.New(color.Faint).Fprintf(
			ansi.NewAnsiStderr(), "\nMORE WITH:\n%s %s\n",
			cmd.CommandPath(), strings.Join(args, " "))
	}
}

// NewAPIClient returns a new utils.APIClient with the API key specified via
// command-line argument or config file.
func NewAPIClient() (*utils.APIClient, error) {
	apikey := viper.GetString("apikey")
	if apikey == "" {
		return nil, errors.New("An API key is needed. Either use the --apikey flag or run \"vt init\" to set up your API key")
	}
	return utils.NewAPIClient(apikey, fmt.Sprintf("vt-cli %s", Version))
}

// ObjectPrinter ...
type ObjectPrinter struct {
	client *utils.APIClient
	cmd    *cobra.Command
}

// NewObjectPrinter ...
func NewObjectPrinter(cmd *cobra.Command) (*ObjectPrinter, error) {
	client, err := NewAPIClient()
	if err != nil {
		return nil, err
	}
	return &ObjectPrinter{client: client, cmd: cmd}, nil
}

// Print ...
func (p *ObjectPrinter) Print(objType string, args []string, argRe *regexp.Regexp) error {

	var r utils.StringReader

	if len(args) == 1 && args[0] == "-" {
		r = utils.NewStringIOReader(os.Stdin)
	} else {
		r = utils.NewStringArrayReader(args)
	}

	if argRe != nil {
		r = utils.NewFilteredStringReader(r, argRe)
	}

	filteredArgs := make([]string, 0)
	for s, err := r.ReadString(); s != "" || err == nil; s, err = r.ReadString() {
		filteredArgs = append(filteredArgs, s)
	}

	objectsCh := make(chan *vt.Object)
	errorsCh := make(chan error, len(filteredArgs))

	go p.client.RetrieveObjects(objType, filteredArgs, objectsCh, errorsCh)

	objs := make([]*vt.Object, 0)

	for obj := range objectsCh {
		if viper.GetBool("identifiers-only") {
			fmt.Printf("%s\n", obj.ID)
		} else {
			objs = append(objs, obj)
		}
	}

	if len(objs) > 0 {
		if err := p.PrintObjects(objs); err != nil {
			return err
		}
	}

	for err := range errorsCh {
		fmt.Fprintln(os.Stderr, err)
	}

	return nil
}

func (p *ObjectPrinter) PrintCollection(collection *url.URL) error {
	it, err := p.client.Iterator(collection,
		vt.IteratorOptions{
			Limit:  viper.GetInt("limit"),
			Cursor: viper.GetString("cursor"),
			Filter: viper.GetString("filter"),
		})
	if err != nil {
		return err
	}
	return p.PrintIter(it)
}

func (p *ObjectPrinter) PrintIter(it *vt.Iterator) error {

	objs := make([]*vt.Object, 0)
	for it.Next() {
		obj := it.Get()
		if viper.GetBool("identifiers-only") {
			fmt.Printf("%s\n", obj.ID)
		} else {
			objs = append(objs, obj)
		}
	}

	if err := it.Error(); err != nil {
		return err
	}

	if len(objs) > 0 {
		if err := p.PrintObjects(objs); err != nil {
			return err
		}
	}

	PrintCommandLineWithCursor(p.cmd, it)
	return nil
}

func (p *ObjectPrinter) PrintObject(obj *vt.Object) error {
	objs := make([]*vt.Object, 1)
	objs[0] = obj
	return p.PrintObjects(objs)
}

func (p *ObjectPrinter) PrintObjects(objs []*vt.Object) error {

	list := make([]map[string]interface{}, 0)

	for _, obj := range objs {
		m := obj.Attributes
		if viper.IsSet("include") && viper.IsSet("exclude") {
			m = utils.FilterMap(
				m, viper.GetStringSlice("include"), viper.GetStringSlice("exclude"))
		}
		for name, r := range obj.Relationships {
			if r.IsOneToOne && len(r.RelatedObjects) > 0 {
				m[name] = r.RelatedObjects[0].ID
			} else {
				l := make([]string, 0)
				for _, obj := range r.RelatedObjects {
					l = append(l, obj.ID)
				}
				m[name] = l
			}
		}
		key := fmt.Sprintf("%s <%s>", obj.Type, obj.ID)
		list = append(list, map[string]interface{}{key: m})
	}

	return yaml.NewColorEncoder(ansi.NewAnsiStdout(), colorScheme).Encode(list)
}
<|endoftext|>


<commit_before>package cmd

import (
	"flag"
	"fmt"
	"io"
	"os"
	"strings"
	"text/tabwriter"
	"text/template"

	"github.com/codegangsta/cli"
	"github.com/myodc/go-micro/broker"
	"github.com/myodc/go-micro/client"
	"github.com/myodc/go-micro/registry"
	"github.com/myodc/go-micro/server"
	"github.com/myodc/go-micro/transport"

	// brokers
	"github.com/myodc/go-micro/broker/http"
	"github.com/myodc/go-micro/broker/nats"
	"github.com/myodc/go-micro/broker/rabbitmq"

	// registries
	"github.com/myodc/go-micro/registry/consul"
	"github.com/myodc/go-micro/registry/etcd"
	"github.com/myodc/go-micro/registry/memory"

	// transport
	thttp "github.com/myodc/go-micro/transport/http"
	tnats "github.com/myodc/go-micro/transport/nats"
	trmq "github.com/myodc/go-micro/transport/rabbitmq"
)

var (
	Flags = []cli.Flag{
		cli.StringFlag{
			Name:   "server_name",
			EnvVar: "MICRO_SERVER_NAME",
			Usage:  "Name of the server. go.micro.srv.example",
		},
		cli.StringFlag{
			Name:   "server_id",
			EnvVar: "MICRO_SERVER_ID",
			Usage:  "Id of the server. Auto-generated if not specified",
		},
		cli.StringFlag{
			Name:   "server_address",
			EnvVar: "MICRO_SERVER_ADDRESS",
			Value:  ":0",
			Usage:  "Bind address for the server. 127.0.0.1:8080",
		},
		cli.StringSliceFlag{
			Name:   "server_metadata",
			EnvVar: "MICRO_SERVER_METADATA",
			Value:  &cli.StringSlice{},
			Usage:  "A list of key-value pairs defining metadata. version=1.0.0",
		},
		cli.StringFlag{
			Name:   "broker",
			EnvVar: "MICRO_BROKER",
			Value:  "http",
			Usage:  "Broker for pub/sub. http, nats, rabbitmq",
		},
		cli.StringFlag{
			Name:   "broker_address",
			EnvVar: "MICRO_BROKER_ADDRESS",
			Usage:  "Comma-separated list of broker addresses",
		},
		cli.StringFlag{
			Name:   "registry",
			EnvVar: "MICRO_REGISTRY",
			Value:  "consul",
			Usage:  "Registry for discovery. memory, consul, etcd, kubernetes",
		},
		cli.StringFlag{
			Name:   "registry_address",
			EnvVar: "MICRO_REGISTRY_ADDRESS",
			Usage:  "Comma-separated list of registry addresses",
		},
		cli.StringFlag{
			Name:   "transport",
			EnvVar: "MICRO_TRANSPORT",
			Value:  "http",
			Usage:  "Transport mechanism used; http, rabbitmq, nats",
		},
		cli.StringFlag{
			Name:   "transport_address",
			EnvVar: "MICRO_TRANSPORT_ADDRESS",
			Usage:  "Comma-separated list of transport addresses",
		},

		// logging flags
		cli.BoolFlag{
			Name:  "logtostderr",
			Usage: "log to standard error instead of files",
		},
		cli.BoolFlag{
			Name:  "alsologtostderr",
			Usage: "log to standard error as well as files",
		},
		cli.StringFlag{
			Name:  "log_dir",
			Usage: "log files will be written to this directory instead of the default temporary directory",
		},
		cli.StringFlag{
			Name:  "stderrthreshold",
			Usage: "logs at or above this threshold go to stderr",
		},
		cli.StringFlag{
			Name:  "v",
			Usage: "log level for V logs",
		},
		cli.StringFlag{
			Name:  "vmodule",
			Usage: "comma-separated list of pattern=N settings for file-filtered logging",
		},
		cli.StringFlag{
			Name:  "log_backtrace_at",
			Usage: "when logging hits line file:N, emit a stack trace",
		},
	}

	Brokers = map[string]func([]string, ...broker.Option) broker.Broker{
		"http":     http.NewBroker,
		"nats":     nats.NewBroker,
		"rabbitmq": rabbitmq.NewBroker,
	}

	Registries = map[string]func([]string, ...registry.Option) registry.Registry{
		"consul":     consul.NewRegistry,
		"etcd":       etcd.NewRegistry,
		"memory":     memory.NewRegistry,
	}

	Transports = map[string]func([]string, ...transport.Option) transport.Transport{
		"http":     thttp.NewTransport,
		"rabbitmq": trmq.NewTransport,
		"nats":     tnats.NewTransport,
	}
)

func Setup(c *cli.Context) error {
	os.Args = os.Args[:1]

	flag.Set("logtostderr", fmt.Sprintf("%v", c.Bool("logtostderr")))
	flag.Set("alsologtostderr", fmt.Sprintf("%v", c.Bool("alsologtostderr")))
	flag.Set("stderrthreshold", c.String("stderrthreshold"))
	flag.Set("log_backtrace_at", c.String("log_backtrace_at"))
	flag.Set("log_dir", c.String("log_dir"))
	flag.Set("vmodule", c.String("vmodule"))
	flag.Set("v", c.String("v"))

	flag.Parse()

	if b, ok := Brokers[c.String("broker")]; ok {
		broker.DefaultBroker = b(strings.Split(c.String("broker_address"), ","))
	}

	if r, ok := Registries[c.String("registry")]; ok {
		registry.DefaultRegistry = r(strings.Split(c.String("registry_address"), ","))
	}

	if t, ok := Transports[c.String("transport")]; ok {
		transport.DefaultTransport = t(strings.Split(c.String("transport_address"), ","))
	}

	metadata := make(map[string]string)
	for _, d := range c.StringSlice("server_metadata") {
		var key, val string
		parts := strings.Split(d, "=")
		key = parts[0]
		if len(parts) > 1 {
			val = strings.Join(parts[1:], "=")
		}
		metadata[key] = val
	}

	server.DefaultServer = server.NewServer(
		server.Name(c.String("server_name")),
		server.Id(c.String("server_id")),
		server.Address(c.String("server_address")),
		server.Metadata(metadata),
	)

	client.DefaultClient = client.NewClient()

	return nil
}

func Init() {
	cli.AppHelpTemplate = `
GLOBAL OPTIONS:
   {{range .Flags}}{{.}}
   {{end}}
`

	cli.HelpPrinter = func(writer io.Writer, templ string, data interface{}) {
		w := tabwriter.NewWriter(os.Stdout, 0, 8, 1, '\t', 0)
		t := template.Must(template.New("help").Parse(templ))
		err := t.Execute(w, data)
		if err != nil {
			panic(err)
		}
		w.Flush()
		os.Exit(2)
	}

	app := cli.NewApp()
	app.HideVersion = true
	app.Usage = "a go micro app"
	app.Action = func(c *cli.Context) {}
	app.Before = Setup
	app.Flags = Flags
	app.RunAndExitOnError()
}
<commit_msg>go fmt<commit_after>package cmd

import (
	"flag"
	"fmt"
	"io"
	"os"
	"strings"
	"text/tabwriter"
	"text/template"

	"github.com/codegangsta/cli"
	"github.com/myodc/go-micro/broker"
	"github.com/myodc/go-micro/client"
	"github.com/myodc/go-micro/registry"
	"github.com/myodc/go-micro/server"
	"github.com/myodc/go-micro/transport"

	// brokers
	"github.com/myodc/go-micro/broker/http"
	"github.com/myodc/go-micro/broker/nats"
	"github.com/myodc/go-micro/broker/rabbitmq"

	// registries
	"github.com/myodc/go-micro/registry/consul"
	"github.com/myodc/go-micro/registry/etcd"
	"github.com/myodc/go-micro/registry/memory"

	// transport
	thttp "github.com/myodc/go-micro/transport/http"
	tnats "github.com/myodc/go-micro/transport/nats"
	trmq "github.com/myodc/go-micro/transport/rabbitmq"
)

var (
	Flags = []cli.Flag{
		cli.StringFlag{
			Name:   "server_name",
			EnvVar: "MICRO_SERVER_NAME",
			Usage:  "Name of the server. go.micro.srv.example",
		},
		cli.StringFlag{
			Name:   "server_id",
			EnvVar: "MICRO_SERVER_ID",
			Usage:  "Id of the server. Auto-generated if not specified",
		},
		cli.StringFlag{
			Name:   "server_address",
			EnvVar: "MICRO_SERVER_ADDRESS",
			Value:  ":0",
			Usage:  "Bind address for the server. 127.0.0.1:8080",
		},
		cli.StringSliceFlag{
			Name:   "server_metadata",
			EnvVar: "MICRO_SERVER_METADATA",
			Value:  &cli.StringSlice{},
			Usage:  "A list of key-value pairs defining metadata. version=1.0.0",
		},
		cli.StringFlag{
			Name:   "broker",
			EnvVar: "MICRO_BROKER",
			Value:  "http",
			Usage:  "Broker for pub/sub. http, nats, rabbitmq",
		},
		cli.StringFlag{
			Name:   "broker_address",
			EnvVar: "MICRO_BROKER_ADDRESS",
			Usage:  "Comma-separated list of broker addresses",
		},
		cli.StringFlag{
			Name:   "registry",
			EnvVar: "MICRO_REGISTRY",
			Value:  "consul",
			Usage:  "Registry for discovery. memory, consul, etcd, kubernetes",
		},
		cli.StringFlag{
			Name:   "registry_address",
			EnvVar: "MICRO_REGISTRY_ADDRESS",
			Usage:  "Comma-separated list of registry addresses",
		},
		cli.StringFlag{
			Name:   "transport",
			EnvVar: "MICRO_TRANSPORT",
			Value:  "http",
			Usage:  "Transport mechanism used; http, rabbitmq, nats",
		},
		cli.StringFlag{
			Name:   "transport_address",
			EnvVar: "MICRO_TRANSPORT_ADDRESS",
			Usage:  "Comma-separated list of transport addresses",
		},

		// logging flags
		cli.BoolFlag{
			Name:  "logtostderr",
			Usage: "log to standard error instead of files",
		},
		cli.BoolFlag{
			Name:  "alsologtostderr",
			Usage: "log to standard error as well as files",
		},
		cli.StringFlag{
			Name:  "log_dir",
			Usage: "log files will be written to this directory instead of the default temporary directory",
		},
		cli.StringFlag{
			Name:  "stderrthreshold",
			Usage: "logs at or above this threshold go to stderr",
		},
		cli.StringFlag{
			Name:  "v",
			Usage: "log level for V logs",
		},
		cli.StringFlag{
			Name:  "vmodule",
			Usage: "comma-separated list of pattern=N settings for file-filtered logging",
		},
		cli.StringFlag{
			Name:  "log_backtrace_at",
			Usage: "when logging hits line file:N, emit a stack trace",
		},
	}

	Brokers = map[string]func([]string, ...broker.Option) broker.Broker{
		"http":     http.NewBroker,
		"nats":     nats.NewBroker,
		"rabbitmq": rabbitmq.NewBroker,
	}

	Registries = map[string]func([]string, ...registry.Option) registry.Registry{
		"consul": consul.NewRegistry,
		"etcd":   etcd.NewRegistry,
		"memory": memory.NewRegistry,
	}

	Transports = map[string]func([]string, ...transport.Option) transport.Transport{
		"http":     thttp.NewTransport,
		"rabbitmq": trmq.NewTransport,
		"nats":     tnats.NewTransport,
	}
)

func Setup(c *cli.Context) error {
	os.Args = os.Args[:1]

	flag.Set("logtostderr", fmt.Sprintf("%v", c.Bool("logtostderr")))
	flag.Set("alsologtostderr", fmt.Sprintf("%v", c.Bool("alsologtostderr")))
	flag.Set("stderrthreshold", c.String("stderrthreshold"))
	flag.Set("log_backtrace_at", c.String("log_backtrace_at"))
	flag.Set("log_dir", c.String("log_dir"))
	flag.Set("vmodule", c.String("vmodule"))
	flag.Set("v", c.String("v"))

	flag.Parse()

	if b, ok := Brokers[c.String("broker")]; ok {
		broker.DefaultBroker = b(strings.Split(c.String("broker_address"), ","))
	}

	if r, ok := Registries[c.String("registry")]; ok {
		registry.DefaultRegistry = r(strings.Split(c.String("registry_address"), ","))
	}

	if t, ok := Transports[c.String("transport")]; ok {
		transport.DefaultTransport = t(strings.Split(c.String("transport_address"), ","))
	}

	metadata := make(map[string]string)
	for _, d := range c.StringSlice("server_metadata") {
		var key, val string
		parts := strings.Split(d, "=")
		key = parts[0]
		if len(parts) > 1 {
			val = strings.Join(parts[1:], "=")
		}
		metadata[key] = val
	}

	server.DefaultServer = server.NewServer(
		server.Name(c.String("server_name")),
		server.Id(c.String("server_id")),
		server.Address(c.String("server_address")),
		server.Metadata(metadata),
	)

	client.DefaultClient = client.NewClient()

	return nil
}

func Init() {
	cli.AppHelpTemplate = `
GLOBAL OPTIONS:
   {{range .Flags}}{{.}}
   {{end}}
`

	cli.HelpPrinter = func(writer io.Writer, templ string, data interface{}) {
		w := tabwriter.NewWriter(os.Stdout, 0, 8, 1, '\t', 0)
		t := template.Must(template.New("help").Parse(templ))
		err := t.Execute(w, data)
		if err != nil {
			panic(err)
		}
		w.Flush()
		os.Exit(2)
	}

	app := cli.NewApp()
	app.HideVersion = true
	app.Usage = "a go micro app"
	app.Action = func(c *cli.Context) {}
	app.Before = Setup
	app.Flags = Flags
	app.RunAndExitOnError()
}
<|endoftext|>