<commit_before>package main

import (
	"bufio"
	"bytes"
	"context"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"time"

	"github.com/tsavola/gate/run"
	"github.com/tsavola/gate/service"
	_ "github.com/tsavola/gate/service/defaults"
	"github.com/tsavola/gate/service/echo"
	"github.com/tsavola/gate/service/origin"
	"github.com/tsavola/wag"
	"github.com/tsavola/wag/dewag"
	"github.com/tsavola/wag/sections"
)

type readWriteCloser struct {
	io.Reader
	io.WriteCloser
}

type timing struct {
	loading time.Duration
	running time.Duration
	overall time.Duration
}

func init() {
	log.SetFlags(0)
	echo.Default.Log = log.New(os.Stderr, "echo service: ", 0)
}

var (
	stackSize = 16 * 1024 * 1024
	dumpTime  = false
	dumpText  = false
	dumpStack = false
	repeat    = 1
)

func main() {
	var (
		config = run.Config{
			MaxProcs:    run.DefaultMaxProcs,
			LibDir:      "lib",
			CgroupTitle: run.DefaultCgroupTitle,
		}
		addr = ""
	)

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [options] wasm...\nOptions:\n", os.Args[0])
		flag.PrintDefaults()
	}

	flag.IntVar(&config.MaxProcs, "max-procs", config.MaxProcs, "limit number of simultaneous programs")
	flag.StringVar(&config.DaemonSocket, "daemon-socket", config.DaemonSocket, "use containerd via unix socket")
	flag.UintVar(&config.CommonGid, "common-gid", config.CommonGid, "group id for file descriptor sharing")
	flag.UintVar(&config.ContainerCred.Uid, "container-uid", config.ContainerCred.Uid, "user id for bootstrapping executor")
	flag.UintVar(&config.ContainerCred.Gid, "container-gid", config.ContainerCred.Gid, "group id for bootstrapping executor")
	flag.UintVar(&config.ExecutorCred.Uid, "executor-uid", config.ExecutorCred.Uid, "user id for executing code")
	flag.UintVar(&config.ExecutorCred.Gid, "executor-gid", config.ExecutorCred.Gid, "group id for executing code")
	flag.StringVar(&config.LibDir, "libdir", config.LibDir, "path")
	flag.StringVar(&config.CgroupParent, "cgroup-parent", config.CgroupParent, "slice")
	flag.StringVar(&config.CgroupTitle, "cgroup-title", config.CgroupTitle, "prefix of dynamic name")
	flag.IntVar(&stackSize, "stack-size", stackSize, "stack size")
	flag.BoolVar(&dumpTime, "dump-time", dumpTime, "print average timings per program")
	flag.BoolVar(&dumpText, "dump-text", dumpText, "disassemble before running")
	flag.BoolVar(&dumpStack, "dump-stack", dumpStack, "print stacktrace after running")
	flag.IntVar(&repeat, "repeat", repeat, "repeat the program execution(s) multiple times")
	flag.StringVar(&addr, "addr", addr, "I/O socket path (replaces stdio)")

	flag.Parse()

	args := flag.Args()
	if len(args) == 0 {
		flag.Usage()
		os.Exit(2)
	}

	ctx := context.Background()

	if addr != "" {
		os.Remove(addr)
		l, err := net.Listen("unix", addr)
		if err != nil {
			log.Fatal(err)
		}
		conn, err := l.Accept()
		l.Close()
		if err != nil {
			log.Fatal(err)
		}
		defer conn.Close()

		origin.Default.R = conn
		origin.Default.W = conn
	} else {
		origin.Default.R = os.Stdin
		origin.Default.W = os.Stdout
	}

	env, err := run.NewEnvironment(&config)
	if err != nil {
		log.Fatalf("environment: %v", err)
	}
	defer env.Close()

	timings := make([]timing, len(args))

	for round := 0; round < repeat; round++ {
		done := make(chan struct{}, len(args))

		for i, arg := range args {
			var r run.ServiceRegistry

			if i == 0 {
				r = service.Defaults
			} else {
				r = origin.CloneRegistryWith(service.Defaults, nil, os.Stdout)
			}

			go execute(ctx, env, arg, r, &timings[i], done)
		}

		for range args {
			<-done
		}
	}

	if dumpTime {
		for i, arg := range args {
			output := func(title string, sum time.Duration) {
				avg := sum / time.Duration(repeat)
				log.Printf("%s "+title+": %6d.%03dµs", arg, avg/time.Microsecond, avg%time.Microsecond)
			}

			output("loading time", timings[i].loading)
			output("running time", timings[i].running)
			output("overall time", timings[i].overall)
		}
	}
}

func execute(ctx context.Context, env *run.Environment, filename string, services run.ServiceRegistry, timing *timing, done chan<- struct{}) {
	defer func() {
		done <- struct{}{}
	}()

	tBegin := time.Now()

	var payload run.Payload

	err := payload.Init()
	if err != nil {
		log.Fatalf("payload: %v", err)
	}
	defer payload.Close()

	tLoadBegin := tBegin

	var ns sections.NameSection

	m := wag.Module{
		MainSymbol:           "main",
		UnknownSectionLoader: sections.UnknownLoaders{"name": ns.Load}.Load,
	}

	err = load(&m, filename, env)
	if err != nil {
		log.Fatalf("module: %v", err)
	}

	tLoadEnd := time.Now()

	_, memorySize := m.MemoryLimits()

	err = payload.Populate(&m, memorySize, int32(stackSize))
	if err != nil {
		log.Fatalf("payload: %v", err)
	}

	if dumpText {
		dewag.PrintTo(os.Stderr, m.Text(), m.FunctionMap(), &ns)
	}

	tRunBegin := time.Now()

	var proc run.Process

	err = proc.Init(ctx, env, &payload, os.Stderr)
	if err != nil {
		log.Fatal(err)
	}
	defer proc.Close()

	exit, trap, err := run.Run(ctx, env, &proc, &payload, services)
	if err != nil {
		log.Fatal(err)
	}

	tRunEnd := time.Now()
	tEnd := tRunEnd

	if trap != 0 {
		log.Printf("trap: %s", trap)
	} else if exit != 0 {
		log.Printf("exit: %d", exit)
	}

	if dumpStack {
		err := payload.DumpStacktrace(os.Stderr, m.FunctionMap(), m.CallMap(), m.FunctionSignatures(), &ns)
		if err != nil {
			log.Printf("stacktrace: %v", err)
		}
	}

	timing.loading += tLoadEnd.Sub(tLoadBegin)
	timing.running += tRunEnd.Sub(tRunBegin)
	timing.overall += tEnd.Sub(tBegin)
}

func load(m *wag.Module, filename string, env *run.Environment) (err error) {
	f, err := os.Open(filename)
	if err != nil {
		return
	}
	defer f.Close()

	err = m.Load(bufio.NewReader(f), env, new(bytes.Buffer), nil, run.RODataAddr, nil)
	return
}
<commit_msg>runner: initialize process outside run timing<commit_after>package main

import (
	"bufio"
	"bytes"
	"context"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"time"

	"github.com/tsavola/gate/run"
	"github.com/tsavola/gate/service"
	_ "github.com/tsavola/gate/service/defaults"
	"github.com/tsavola/gate/service/echo"
	"github.com/tsavola/gate/service/origin"
	"github.com/tsavola/wag"
	"github.com/tsavola/wag/dewag"
	"github.com/tsavola/wag/sections"
)

type readWriteCloser struct {
	io.Reader
	io.WriteCloser
}

type timing struct {
	loading time.Duration
	running time.Duration
	overall time.Duration
}

func init() {
	log.SetFlags(0)
	echo.Default.Log = log.New(os.Stderr, "echo service: ", 0)
}

var (
	stackSize = 16 * 1024 * 1024
	dumpTime  = false
	dumpText  = false
	dumpStack = false
	repeat    = 1
)

func main() {
	var (
		config = run.Config{
			MaxProcs:    run.DefaultMaxProcs,
			LibDir:      "lib",
			CgroupTitle: run.DefaultCgroupTitle,
		}
		addr = ""
	)

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [options] wasm...\nOptions:\n", os.Args[0])
		flag.PrintDefaults()
	}

	flag.IntVar(&config.MaxProcs, "max-procs", config.MaxProcs, "limit number of simultaneous programs")
	flag.StringVar(&config.DaemonSocket, "daemon-socket", config.DaemonSocket, "use containerd via unix socket")
	flag.UintVar(&config.CommonGid, "common-gid", config.CommonGid, "group id for file descriptor sharing")
	flag.UintVar(&config.ContainerCred.Uid, "container-uid", config.ContainerCred.Uid, "user id for bootstrapping executor")
	flag.UintVar(&config.ContainerCred.Gid, "container-gid", config.ContainerCred.Gid, "group id for bootstrapping executor")
	flag.UintVar(&config.ExecutorCred.Uid, "executor-uid", config.ExecutorCred.Uid, "user id for executing code")
	flag.UintVar(&config.ExecutorCred.Gid, "executor-gid", config.ExecutorCred.Gid, "group id for executing code")
	flag.StringVar(&config.LibDir, "libdir", config.LibDir, "path")
	flag.StringVar(&config.CgroupParent, "cgroup-parent", config.CgroupParent, "slice")
	flag.StringVar(&config.CgroupTitle, "cgroup-title", config.CgroupTitle, "prefix of dynamic name")
	flag.IntVar(&stackSize, "stack-size", stackSize, "stack size")
	flag.BoolVar(&dumpTime, "dump-time", dumpTime, "print average timings per program")
	flag.BoolVar(&dumpText, "dump-text", dumpText, "disassemble before running")
	flag.BoolVar(&dumpStack, "dump-stack", dumpStack, "print stacktrace after running")
	flag.IntVar(&repeat, "repeat", repeat, "repeat the program execution(s) multiple times")
	flag.StringVar(&addr, "addr", addr, "I/O socket path (replaces stdio)")

	flag.Parse()

	args := flag.Args()
	if len(args) == 0 {
		flag.Usage()
		os.Exit(2)
	}

	ctx := context.Background()

	if addr != "" {
		os.Remove(addr)
		l, err := net.Listen("unix", addr)
		if err != nil {
			log.Fatal(err)
		}
		conn, err := l.Accept()
		l.Close()
		if err != nil {
			log.Fatal(err)
		}
		defer conn.Close()

		origin.Default.R = conn
		origin.Default.W = conn
	} else {
		origin.Default.R = os.Stdin
		origin.Default.W = os.Stdout
	}

	env, err := run.NewEnvironment(&config)
	if err != nil {
		log.Fatalf("environment: %v", err)
	}
	defer env.Close()

	timings := make([]timing, len(args))

	for round := 0; round < repeat; round++ {
		done := make(chan struct{}, len(args))

		for i, arg := range args {
			var r run.ServiceRegistry

			if i == 0 {
				r = service.Defaults
			} else {
				r = origin.CloneRegistryWith(service.Defaults, nil, os.Stdout)
			}

			go execute(ctx, env, arg, r, &timings[i], done)
		}

		for range args {
			<-done
		}
	}

	if dumpTime {
		for i, arg := range args {
			output := func(title string, sum time.Duration) {
				avg := sum / time.Duration(repeat)
				log.Printf("%s "+title+": %6d.%03dµs", arg, avg/time.Microsecond, avg%time.Microsecond)
			}

			output("loading time", timings[i].loading)
			output("running time", timings[i].running)
			output("overall time", timings[i].overall)
		}
	}
}

func execute(ctx context.Context, env *run.Environment, filename string, services run.ServiceRegistry, timing *timing, done chan<- struct{}) {
	defer func() {
		done <- struct{}{}
	}()

	tBegin := time.Now()

	var payload run.Payload

	err := payload.Init()
	if err != nil {
		log.Fatalf("payload: %v", err)
	}
	defer payload.Close()

	var proc run.Process

	err = proc.Init(ctx, env, &payload, os.Stderr)
	if err != nil {
		log.Fatalf("process: %v", err)
	}
	defer proc.Close()

	tLoadBegin := tBegin

	var ns sections.NameSection

	m := wag.Module{
		MainSymbol:           "main",
		UnknownSectionLoader: sections.UnknownLoaders{"name": ns.Load}.Load,
	}

	err = load(&m, filename, env)
	if err != nil {
		log.Fatalf("module: %v", err)
	}

	tLoadEnd := time.Now()

	_, memorySize := m.MemoryLimits()

	err = payload.Populate(&m, memorySize, int32(stackSize))
	if err != nil {
		log.Fatalf("payload: %v", err)
	}

	if dumpText {
		dewag.PrintTo(os.Stderr, m.Text(), m.FunctionMap(), &ns)
	}

	tRunBegin := time.Now()

	exit, trap, err := run.Run(ctx, env, &proc, &payload, services)
	if err != nil {
		log.Fatalf("run: %v", err)
	}

	tRunEnd := time.Now()
	tEnd := tRunEnd

	if trap != 0 {
		log.Printf("trap: %s", trap)
	} else if exit != 0 {
		log.Printf("exit: %d", exit)
	}

	if dumpStack {
		err := payload.DumpStacktrace(os.Stderr, m.FunctionMap(), m.CallMap(), m.FunctionSignatures(), &ns)
		if err != nil {
			log.Printf("stacktrace: %v", err)
		}
	}

	timing.loading += tLoadEnd.Sub(tLoadBegin)
	timing.running += tRunEnd.Sub(tRunBegin)
	timing.overall += tEnd.Sub(tBegin)
}

func load(m *wag.Module, filename string, env *run.Environment) (err error) {
	f, err := os.Open(filename)
	if err != nil {
		return
	}
	defer f.Close()

	err = m.Load(bufio.NewReader(f), env, new(bytes.Buffer), nil, run.RODataAddr, nil)
	return
}
<|endoftext|>


<commit_before>package command

import (
	"encoding/json"
	"flag"
	"fmt"
	"strings"

	"github.com/funkygao/gafka/ctx"
	zkr "github.com/funkygao/gafka/registry/zk"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/gofmt"
)

type Kateway struct {
	Ui  cli.Ui
	Cmd string

	zone string
}

func (this *Kateway) Run(args []string) (exitCode int) {
	cmdFlags := flag.NewFlagSet("kateway", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&this.zone, "z", ctx.ZkDefaultZone(), "")
	if err := cmdFlags.Parse(args); err != nil {
		return 2
	}

	zkzone := zk.NewZkZone(zk.DefaultConfig(this.zone, ctx.ZoneZkAddrs(this.zone)))
	instances, _, err := zkzone.Conn().Children(zkr.Root(this.zone))
	if err != nil {
		if err.Error() == "zk: node does not exist" {
			this.Ui.Output("no kateway running")
			return
		} else {
			swallow(err)
		}
	}

	for _, instance := range instances {
		data, stat, err := zkzone.Conn().Get(zkr.Root(this.zone) + "/" + instance)
		swallow(err)

		info := make(map[string]string)
		json.Unmarshal(data, &info)

		this.Ui.Info(fmt.Sprintf("%s id:%-2s up:%s", info["host"], instance,
			gofmt.PrettySince(zk.ZkTimestamp(stat.Ctime).Time())))
		this.Ui.Output(fmt.Sprintf("    ver: %s\n    build: %s\n    pub: %s\n    sub: %s\n    man: %s\n    dbg: %s",
			info["ver"],
			info["build"],
			info["pub"],
			info["sub"],
			info["man"],
			info["debug"],
		))

	}

	return
}

func (*Kateway) Synopsis() string {
	return "List online kateway instances"
}

func (this *Kateway) Help() string {
	help := fmt.Sprintf(`
Usage: %s kateway [options]

    List online kateway instances

Options:

    -z zone

`, this.Cmd)
	return strings.TrimSpace(help)
}
<commit_msg>gk display online kateway consumers<commit_after>package command

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"net/http"
	"strings"

	"github.com/funkygao/gafka/ctx"
	zkr "github.com/funkygao/gafka/registry/zk"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/gofmt"
)

type Kateway struct {
	Ui  cli.Ui
	Cmd string

	zone          string
	showConsumers bool
}

func (this *Kateway) Run(args []string) (exitCode int) {
	cmdFlags := flag.NewFlagSet("kateway", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&this.zone, "z", ctx.ZkDefaultZone(), "")
	cmdFlags.BoolVar(&this.showConsumers, "consumers", false, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 2
	}

	zkzone := zk.NewZkZone(zk.DefaultConfig(this.zone, ctx.ZoneZkAddrs(this.zone)))
	instances, _, err := zkzone.Conn().Children(zkr.Root(this.zone))
	if err != nil {
		if err.Error() == "zk: node does not exist" {
			this.Ui.Output("no kateway running")
			return
		} else {
			swallow(err)
		}
	}

	for _, instance := range instances {
		data, stat, err := zkzone.Conn().Get(zkr.Root(this.zone) + "/" + instance)
		swallow(err)

		info := make(map[string]string)
		json.Unmarshal(data, &info)

		this.Ui.Info(fmt.Sprintf("%s id:%-2s up:%s", info["host"], instance,
			gofmt.PrettySince(zk.ZkTimestamp(stat.Ctime).Time())))
		this.Ui.Output(fmt.Sprintf("    ver: %s\n    build: %s\n    pub: %s\n    sub: %s\n    man: %s\n    dbg: %s",
			info["ver"],
			info["build"],
			info["pub"],
			info["sub"],
			info["man"],
			info["debug"],
		))

		if this.showConsumers {
			// FIXME
			this.showOnlineConsumers("http://" + info["man"] + "/consumers")
		}

	}

	return
}

func (this *Kateway) showOnlineConsumers(url string) {
	response, err := http.Get(url)
	swallow(err)

	b, err := ioutil.ReadAll(response.Body)
	swallow(err)

	response.Body.Close()

	consumers := make(map[string]string)
	json.Unmarshal(b, &consumers)

	for addr, info := range consumers {
		this.Ui.Output(fmt.Sprintf("    consumers: %s  %+v", addr, info))
	}

}

func (*Kateway) Synopsis() string {
	return "List online kateway instances"
}

func (this *Kateway) Help() string {
	help := fmt.Sprintf(`
Usage: %s kateway [options]

    List online kateway instances

Options:

    -z zone

    -consumers
      Display online consumers of kateway

`, this.Cmd)
	return strings.TrimSpace(help)
}
<|endoftext|>


<commit_before>package command

import (
	"bufio"
	"flag"
	"fmt"
	"net"
	"strings"

	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/color"
	"github.com/funkygao/golib/pipestream"
)

// consul members will include:
// - zk cluster as server
// - agents
//   - brokers
//   - kateway
type Members struct {
	Ui  cli.Ui
	Cmd string

	brokerHosts, zkHosts, katewayHosts map[string]struct{}
	nodeHostMap                        map[string]string // consul members node->ip
}

func (this *Members) Run(args []string) (exitCode int) {
	var (
		zone        string
		showLoadAvg bool
	)
	cmdFlags := flag.NewFlagSet("members", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&zone, "z", ctx.ZkDefaultZone(), "")
	cmdFlags.BoolVar(&showLoadAvg, "l", false, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	zkzone := zk.NewZkZone(zk.DefaultConfig(zone, ctx.ZoneZkAddrs(zone)))
	this.fillTheHosts(zkzone)

	consulLiveNode, consulDeadNodes := this.consulMembers()
	for _, node := range consulDeadNodes {
		this.Ui.Error(fmt.Sprintf("%s consul dead", node))
	}

	consulLiveMap := make(map[string]struct{})
	brokerN, zkN, katewayN, unknownN := 0, 0, 0, 0
	for _, node := range consulLiveNode {
		_, presentInBroker := this.brokerHosts[node]
		_, presentInZk := this.zkHosts[node]
		_, presentInKateway := this.katewayHosts[node]
		if presentInBroker {
			brokerN++
		}
		if presentInZk {
			zkN++
		}
		if presentInKateway {
			katewayN++
		}

		if !presentInBroker && !presentInZk && !presentInKateway {
			unknownN++

			this.Ui.Info(fmt.Sprintf("? %s", node))
		}

		consulLiveMap[node] = struct{}{}
	}

	// all brokers should run consul
	for broker, _ := range this.brokerHosts {
		if _, present := consulLiveMap[broker]; !present {
			this.Ui.Warn(fmt.Sprintf("- %s", broker))
		}
	}

	if showLoadAvg {
		this.displayLoadAvg()
	}

	this.Ui.Output(fmt.Sprintf("zk:%s broker:%s kateway:%s ?:%s",
		color.Magenta("%d", zkN),
		color.Magenta("%d", brokerN),
		color.Magenta("%d", katewayN),
		color.Green("%d", unknownN)))

	return
}

func (this *Members) fillTheHosts(zkzone *zk.ZkZone) {
	this.brokerHosts = make(map[string]struct{})
	zkzone.ForSortedBrokers(func(cluster string, brokers map[string]*zk.BrokerZnode) {
		for _, brokerInfo := range brokers {
			this.brokerHosts[brokerInfo.Host] = struct{}{}
		}
	})

	this.zkHosts = make(map[string]struct{})
	for _, addr := range zkzone.ZkAddrList() {
		zkNode, _, err := net.SplitHostPort(addr)
		swallow(err)
		this.zkHosts[zkNode] = struct{}{}
	}

	this.katewayHosts = make(map[string]struct{})
	kws, err := zkzone.KatewayInfos()
	swallow(err)
	for _, kw := range kws {
		host, _, err := net.SplitHostPort(kw.PubAddr)
		swallow(err)
		this.katewayHosts[host] = struct{}{}
	}
}

func (this *Members) displayLoadAvg() {
	cmd := pipestream.New("consul", "exec",
		"uptime", "|", "grep", "load")
	err := cmd.Open()
	swallow(err)
	defer cmd.Close()

	scanner := bufio.NewScanner(cmd.Reader())
	scanner.Split(bufio.ScanLines)
	for scanner.Scan() {
		line := scanner.Text()
		fields := strings.Fields(line)
		node := fields[0]
		parts := strings.Split(line, "load average:")
		if len(parts) < 2 {
			continue
		}
		if strings.HasSuffix(node, ":") {
			node = strings.TrimRight(node, ":")
		}

		host := this.nodeHostMap[node]
		this.Ui.Output(fmt.Sprintf("%35s %s %s", node, this.roleOfHost(host), parts[1]))
	}
}

func (this *Members) roleOfHost(host string) string {
	if _, present := this.brokerHosts[host]; present {
		return "B"
	}
	if _, present := this.zkHosts[host]; present {
		return "Z"
	}
	if _, present := this.katewayHosts[host]; present {
		return "K"
	}
	return "?"
}

func (this *Members) consulMembers() ([]string, []string) {
	cmd := pipestream.New("consul", "members")
	err := cmd.Open()
	swallow(err)
	defer cmd.Close()

	liveHosts, deadHosts := []string{}, []string{}
	scanner := bufio.NewScanner(cmd.Reader())
	scanner.Split(bufio.ScanLines)
	this.nodeHostMap = make(map[string]string)
	for scanner.Scan() {
		if strings.Contains(scanner.Text(), "Protocol") {
			// the header
			continue
		}

		fields := strings.Fields(scanner.Text())
		node, addr, alive := fields[0], fields[1], fields[2]
		host, _, err := net.SplitHostPort(addr)
		swallow(err)

		this.nodeHostMap[node] = host

		if alive == "alive" {
			liveHosts = append(liveHosts, host)
		} else {
			deadHosts = append(deadHosts, host)
		}
	}

	return liveHosts, deadHosts
}

func (*Members) Synopsis() string {
	return "Verify consul members match kafka zone"
}

func (this *Members) Help() string {
	help := fmt.Sprintf(`
Usage: %s members [options]

    Verify consul members match kafka zone

    -z zone

    -l
      Display each member load average

`, this.Cmd)
	return strings.TrimSpace(help)
}
<commit_msg>format the output<commit_after>package command

import (
	"bufio"
	"flag"
	"fmt"
	"net"
	"strings"

	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/color"
	"github.com/funkygao/golib/pipestream"
	"github.com/ryanuber/columnize"
)

// consul members will include:
// - zk cluster as server
// - agents
//   - brokers
//   - kateway
type Members struct {
	Ui  cli.Ui
	Cmd string

	brokerHosts, zkHosts, katewayHosts map[string]struct{}
	nodeHostMap                        map[string]string // consul members node->ip
}

func (this *Members) Run(args []string) (exitCode int) {
	var (
		zone        string
		showLoadAvg bool
	)
	cmdFlags := flag.NewFlagSet("members", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&zone, "z", ctx.ZkDefaultZone(), "")
	cmdFlags.BoolVar(&showLoadAvg, "l", false, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	zkzone := zk.NewZkZone(zk.DefaultConfig(zone, ctx.ZoneZkAddrs(zone)))
	this.fillTheHosts(zkzone)

	consulLiveNode, consulDeadNodes := this.consulMembers()
	for _, node := range consulDeadNodes {
		this.Ui.Error(fmt.Sprintf("%s consul dead", node))
	}

	consulLiveMap := make(map[string]struct{})
	brokerN, zkN, katewayN, unknownN := 0, 0, 0, 0
	for _, node := range consulLiveNode {
		_, presentInBroker := this.brokerHosts[node]
		_, presentInZk := this.zkHosts[node]
		_, presentInKateway := this.katewayHosts[node]
		if presentInBroker {
			brokerN++
		}
		if presentInZk {
			zkN++
		}
		if presentInKateway {
			katewayN++
		}

		if !presentInBroker && !presentInZk && !presentInKateway {
			unknownN++

			this.Ui.Info(fmt.Sprintf("? %s", node))
		}

		consulLiveMap[node] = struct{}{}
	}

	// all brokers should run consul
	for broker, _ := range this.brokerHosts {
		if _, present := consulLiveMap[broker]; !present {
			this.Ui.Warn(fmt.Sprintf("- %s", broker))
		}
	}

	if showLoadAvg {
		this.displayLoadAvg()
	}

	this.Ui.Output(fmt.Sprintf("zk:%s broker:%s kateway:%s ?:%s",
		color.Magenta("%d", zkN),
		color.Magenta("%d", brokerN),
		color.Magenta("%d", katewayN),
		color.Green("%d", unknownN)))

	return
}

func (this *Members) fillTheHosts(zkzone *zk.ZkZone) {
	this.brokerHosts = make(map[string]struct{})
	zkzone.ForSortedBrokers(func(cluster string, brokers map[string]*zk.BrokerZnode) {
		for _, brokerInfo := range brokers {
			this.brokerHosts[brokerInfo.Host] = struct{}{}
		}
	})

	this.zkHosts = make(map[string]struct{})
	for _, addr := range zkzone.ZkAddrList() {
		zkNode, _, err := net.SplitHostPort(addr)
		swallow(err)
		this.zkHosts[zkNode] = struct{}{}
	}

	this.katewayHosts = make(map[string]struct{})
	kws, err := zkzone.KatewayInfos()
	swallow(err)
	for _, kw := range kws {
		host, _, err := net.SplitHostPort(kw.PubAddr)
		swallow(err)
		this.katewayHosts[host] = struct{}{}
	}
}

func (this *Members) displayLoadAvg() {
	cmd := pipestream.New("consul", "exec",
		"uptime", "|", "grep", "load")
	err := cmd.Open()
	swallow(err)
	defer cmd.Close()

	lines := make([]string, 0)
	header := "Node|Host|Role|Load Avg"
	lines = append(lines, header)

	scanner := bufio.NewScanner(cmd.Reader())
	scanner.Split(bufio.ScanLines)
	for scanner.Scan() {
		line := scanner.Text()
		fields := strings.Fields(line)
		node := fields[0]
		parts := strings.Split(line, "load average:")
		if len(parts) < 2 {
			continue
		}
		if strings.HasSuffix(node, ":") {
			node = strings.TrimRight(node, ":")
		}

		host := this.nodeHostMap[node]
		lines = append(lines, fmt.Sprintf("%s|%s|%s|%s", node, host, this.roleOfHost(host), parts[1]))
	}

	if len(lines) > 1 {
		this.Ui.Output(columnize.SimpleFormat(lines))
	}
}

func (this *Members) roleOfHost(host string) string {
	if _, present := this.brokerHosts[host]; present {
		return "B"
	}
	if _, present := this.zkHosts[host]; present {
		return "Z"
	}
	if _, present := this.katewayHosts[host]; present {
		return "K"
	}
	return "?"
}

func (this *Members) consulMembers() ([]string, []string) {
	cmd := pipestream.New("consul", "members")
	err := cmd.Open()
	swallow(err)
	defer cmd.Close()

	liveHosts, deadHosts := []string{}, []string{}
	scanner := bufio.NewScanner(cmd.Reader())
	scanner.Split(bufio.ScanLines)
	this.nodeHostMap = make(map[string]string)
	for scanner.Scan() {
		if strings.Contains(scanner.Text(), "Protocol") {
			// the header
			continue
		}

		fields := strings.Fields(scanner.Text())
		node, addr, alive := fields[0], fields[1], fields[2]
		host, _, err := net.SplitHostPort(addr)
		swallow(err)

		this.nodeHostMap[node] = host

		if alive == "alive" {
			liveHosts = append(liveHosts, host)
		} else {
			deadHosts = append(deadHosts, host)
		}
	}

	return liveHosts, deadHosts
}

func (*Members) Synopsis() string {
	return "Verify consul members match kafka zone"
}

func (this *Members) Help() string {
	help := fmt.Sprintf(`
Usage: %s members [options]

    Verify consul members match kafka zone

    -z zone

    -l
      Display each member load average

`, this.Cmd)
	return strings.TrimSpace(help)
}
<|endoftext|>


<commit_before>// Copyright ©2017 The Gonum 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 flow

import (
	"flag"
	"fmt"
	"io/ioutil"
	"math"
	"path/filepath"
	"strings"
	"testing"

	"golang.org/x/exp/rand"

	"gonum.org/v1/gonum/graph"
	"gonum.org/v1/gonum/graph/encoding"
	"gonum.org/v1/gonum/graph/encoding/dot"
	"gonum.org/v1/gonum/graph/graphs/gen"
	"gonum.org/v1/gonum/graph/iterator"
	"gonum.org/v1/gonum/graph/simple"
	"gonum.org/v1/gonum/graph/topo"
)

var slta = flag.Bool("slta", false, "specify DominatorsSLT benchmark")

func BenchmarkDominators(b *testing.B) {
	testdata := filepath.FromSlash("./testdata/flow")

	fis, err := ioutil.ReadDir(testdata)
	if err != nil {
		b.Fatalf("failed to open control flow testdata: %v", err)
	}
	for _, fi := range fis {
		name := fi.Name()
		ext := filepath.Ext(name)
		if ext != ".dot" {
			continue
		}
		test := name[:len(name)-len(ext)]

		data, err := ioutil.ReadFile(filepath.Join(testdata, name))
		if err != nil {
			b.Errorf("failed to open control flow case: %v", err)
			continue
		}
		g := &labeled{DirectedGraph: simple.NewDirectedGraph()}
		err = dot.Unmarshal(data, g)
		if err != nil {
			b.Errorf("failed to unmarshal graph data: %v", err)
			continue
		}
		want := g.root
		if want == nil {
			b.Error("no entry node label for graph")
			continue
		}

		if *slta {
			b.Run(test, func(b *testing.B) {
				for i := 0; i < b.N; i++ {
					d := DominatorsSLT(g.root, g)
					if got := d.Root(); got.ID() != want.ID() {
						b.Fatalf("unexpected root node: got:%d want:%d", got.ID(), want.ID())
					}
				}
			})
		} else {
			b.Run(test, func(b *testing.B) {
				for i := 0; i < b.N; i++ {
					d := Dominators(g.root, g)
					if got := d.Root(); got.ID() != want.ID() {
						b.Fatalf("unexpected root node: got:%d want:%d", got.ID(), want.ID())
					}
				}
			})
		}
	}
}

type labeled struct {
	*simple.DirectedGraph

	root *node
}

func (g *labeled) NewNode() graph.Node {
	return &node{Node: g.DirectedGraph.NewNode(), g: g}
}

func (g *labeled) SetEdge(e graph.Edge) {
	if e.To().ID() == e.From().ID() {
		// Do not attempt to add self edges.
		return
	}
	g.DirectedGraph.SetEdge(e)
}

type node struct {
	graph.Node
	name string
	g    *labeled
}

func (n *node) SetDOTID(id string) {
	n.name = id
}

func (n *node) SetAttribute(attr encoding.Attribute) error {
	if attr.Key != "label" {
		return nil
	}
	switch attr.Value {
	default:
		if attr.Value != `"{%0}"` && !strings.HasPrefix(attr.Value, `"{%0|`) {
			return nil
		}
		fallthrough
	case "entry", "root":
		if n.g.root != nil {
			return fmt.Errorf("set root for graph with existing root: old=%q new=%q", n.g.root.name, n.name)
		}
		n.g.root = n
	}
	return nil
}

func BenchmarkRandomGraphDominators(b *testing.B) {
	tests := []struct {
		name string
		g    func() *simple.DirectedGraph
	}{
		{name: "gnm-n=1e3-m=1e3", g: gnm(1e3, 1e3)},
		{name: "gnm-n=1e3-m=3e3", g: gnm(1e3, 3e3)},
		{name: "gnm-n=1e3-m=1e4", g: gnm(1e3, 1e4)},
		{name: "gnm-n=1e3-m=3e4", g: gnm(1e3, 3e4)},

		{name: "gnm-n=1e4-m=1e4", g: gnm(1e4, 1e4)},
		{name: "gnm-n=1e4-m=3e4", g: gnm(1e4, 3e4)},
		{name: "gnm-n=1e4-m=1e5", g: gnm(1e4, 1e5)},
		{name: "gnm-n=1e4-m=3e5", g: gnm(1e4, 3e5)},

		{name: "gnm-n=1e5-m=1e5", g: gnm(1e5, 1e5)},
		{name: "gnm-n=1e5-m=3e5", g: gnm(1e5, 3e5)},
		{name: "gnm-n=1e5-m=1e6", g: gnm(1e5, 1e6)},
		{name: "gnm-n=1e5-m=3e6", g: gnm(1e5, 3e6)},

		{name: "gnm-n=1e6-m=1e6", g: gnm(1e6, 1e6)},
		{name: "gnm-n=1e6-m=3e6", g: gnm(1e6, 3e6)},
		{name: "gnm-n=1e6-m=1e7", g: gnm(1e6, 1e7)},
		{name: "gnm-n=1e6-m=3e7", g: gnm(1e6, 3e7)},

		{name: "dup-n=1e3-d=0.8-a=0.1", g: duplication(1e3, 0.8, 0.1, math.NaN())},
		{name: "dup-n=1e3-d=0.5-a=0.2", g: duplication(1e3, 0.5, 0.2, math.NaN())},

		{name: "dup-n=1e4-d=0.8-a=0.1", g: duplication(1e4, 0.8, 0.1, math.NaN())},
		{name: "dup-n=1e4-d=0.5-a=0.2", g: duplication(1e4, 0.5, 0.2, math.NaN())},

		{name: "dup-n=1e5-d=0.8-a=0.1", g: duplication(1e5, 0.8, 0.1, math.NaN())},
		{name: "dup-n=1e5-d=0.5-a=0.2", g: duplication(1e5, 0.5, 0.2, math.NaN())},
	}

	for _, test := range tests {
		rnd := rand.New(rand.NewSource(1))
		g := test.g()

		// Guess a maximally expensive entry to the graph.
		sort, err := topo.Sort(g)
		root := sort[0]
		if root == nil {
			// If we did not get a node in the first position
			// then there must be an unorderable set of nodes
			// in the first position of the error. Pick one
			// of the nodes at random.
			unordered := err.(topo.Unorderable)
			root = unordered[0][rnd.Intn(len(unordered[0]))]
		}
		if root == nil {
			b.Error("no entry node label for graph")
			continue
		}

		if len(sort) > 1 {
			// Ensure that the graph has a complete path
			// through the sorted nodes.

			// unordered will only be accessed if there is
			// a sort element that is nil, in which case
			// unordered will contain a set of nodes from
			// an SCC.
			unordered, _ := err.(topo.Unorderable)

			var ui int
			for i, v := range sort[1:] {
				u := sort[i]
				if u == nil {
					u = unordered[ui][rnd.Intn(len(unordered[ui]))]
					ui++
				}
				if v == nil {
					v = unordered[ui][rnd.Intn(len(unordered[ui]))]
				}
				if !g.HasEdgeFromTo(u.ID(), v.ID()) {
					g.SetEdge(g.NewEdge(u, v))
				}
			}
		}

		b.Run(test.name, func(b *testing.B) {
			for i := 0; i < b.N; i++ {
				d := Dominators(root, g)
				if got := d.Root(); got.ID() != root.ID() {
					b.Fatalf("unexpected root node: got:%d want:%d", got.ID(), root.ID())
				}
			}
		})
	}
}

// gnm returns a directed G(n,m) Erdõs-Rényi graph.
func gnm(n, m int) func() *simple.DirectedGraph {
	return func() *simple.DirectedGraph {
		dg := simple.NewDirectedGraph()
		err := gen.Gnm(dg, n, m, rand.New(rand.NewSource(1)))
		if err != nil {
			panic(err)
		}
		return dg
	}
}

// duplication returns an edge-induced directed subgraph of a
// duplication graph.
func duplication(n int, delta, alpha, sigma float64) func() *simple.DirectedGraph {
	return func() *simple.DirectedGraph {
		g := undirected{simple.NewDirectedGraph()}
		rnd := rand.New(rand.NewSource(1))
		err := gen.Duplication(g, n, delta, alpha, sigma, rnd)
		if err != nil {
			panic(err)
		}
		for _, e := range graph.EdgesOf(g.Edges()) {
			if rnd.Intn(2) == 0 {
				g.RemoveEdge(e.From().ID(), e.To().ID())
			}
		}
		return g.DirectedGraph
	}
}

type undirected struct {
	*simple.DirectedGraph
}

func (g undirected) From(id int64) graph.Nodes {
	return iterator.NewOrderedNodes(append(
		graph.NodesOf(g.DirectedGraph.From(id)),
		graph.NodesOf(g.DirectedGraph.To(id))...))
}

func (g undirected) HasEdgeBetween(xid, yid int64) bool {
	return g.DirectedGraph.HasEdgeFromTo(xid, yid)
}

func (g undirected) EdgeBetween(xid, yid int64) graph.Edge {
	return g.DirectedGraph.Edge(xid, yid)
}

func (g undirected) SetEdge(e graph.Edge) {
	g.DirectedGraph.SetEdge(e)
	g.DirectedGraph.SetEdge(g.DirectedGraph.NewEdge(e.To(), e.From()))
}
<commit_msg>graph/flow: skip rather than erroring when there is no benchmark data<commit_after>// Copyright ©2017 The Gonum 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 flow

import (
	"flag"
	"fmt"
	"io/ioutil"
	"math"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"golang.org/x/exp/rand"

	"gonum.org/v1/gonum/graph"
	"gonum.org/v1/gonum/graph/encoding"
	"gonum.org/v1/gonum/graph/encoding/dot"
	"gonum.org/v1/gonum/graph/graphs/gen"
	"gonum.org/v1/gonum/graph/iterator"
	"gonum.org/v1/gonum/graph/simple"
	"gonum.org/v1/gonum/graph/topo"
)

var slta = flag.Bool("slta", false, "specify DominatorsSLT benchmark")

func BenchmarkDominators(b *testing.B) {
	testdata := filepath.FromSlash("./testdata/flow")

	fis, err := ioutil.ReadDir(testdata)
	if err != nil {
		if os.IsNotExist(err) {
			b.Skipf("no control flow testdata: %v", err)
		}
		b.Fatalf("failed to open control flow testdata: %v", err)
	}
	for _, fi := range fis {
		name := fi.Name()
		ext := filepath.Ext(name)
		if ext != ".dot" {
			continue
		}
		test := name[:len(name)-len(ext)]

		data, err := ioutil.ReadFile(filepath.Join(testdata, name))
		if err != nil {
			b.Errorf("failed to open control flow case: %v", err)
			continue
		}
		g := &labeled{DirectedGraph: simple.NewDirectedGraph()}
		err = dot.Unmarshal(data, g)
		if err != nil {
			b.Errorf("failed to unmarshal graph data: %v", err)
			continue
		}
		want := g.root
		if want == nil {
			b.Error("no entry node label for graph")
			continue
		}

		if *slta {
			b.Run(test, func(b *testing.B) {
				for i := 0; i < b.N; i++ {
					d := DominatorsSLT(g.root, g)
					if got := d.Root(); got.ID() != want.ID() {
						b.Fatalf("unexpected root node: got:%d want:%d", got.ID(), want.ID())
					}
				}
			})
		} else {
			b.Run(test, func(b *testing.B) {
				for i := 0; i < b.N; i++ {
					d := Dominators(g.root, g)
					if got := d.Root(); got.ID() != want.ID() {
						b.Fatalf("unexpected root node: got:%d want:%d", got.ID(), want.ID())
					}
				}
			})
		}
	}
}

type labeled struct {
	*simple.DirectedGraph

	root *node
}

func (g *labeled) NewNode() graph.Node {
	return &node{Node: g.DirectedGraph.NewNode(), g: g}
}

func (g *labeled) SetEdge(e graph.Edge) {
	if e.To().ID() == e.From().ID() {
		// Do not attempt to add self edges.
		return
	}
	g.DirectedGraph.SetEdge(e)
}

type node struct {
	graph.Node
	name string
	g    *labeled
}

func (n *node) SetDOTID(id string) {
	n.name = id
}

func (n *node) SetAttribute(attr encoding.Attribute) error {
	if attr.Key != "label" {
		return nil
	}
	switch attr.Value {
	default:
		if attr.Value != `"{%0}"` && !strings.HasPrefix(attr.Value, `"{%0|`) {
			return nil
		}
		fallthrough
	case "entry", "root":
		if n.g.root != nil {
			return fmt.Errorf("set root for graph with existing root: old=%q new=%q", n.g.root.name, n.name)
		}
		n.g.root = n
	}
	return nil
}

func BenchmarkRandomGraphDominators(b *testing.B) {
	tests := []struct {
		name string
		g    func() *simple.DirectedGraph
	}{
		{name: "gnm-n=1e3-m=1e3", g: gnm(1e3, 1e3)},
		{name: "gnm-n=1e3-m=3e3", g: gnm(1e3, 3e3)},
		{name: "gnm-n=1e3-m=1e4", g: gnm(1e3, 1e4)},
		{name: "gnm-n=1e3-m=3e4", g: gnm(1e3, 3e4)},

		{name: "gnm-n=1e4-m=1e4", g: gnm(1e4, 1e4)},
		{name: "gnm-n=1e4-m=3e4", g: gnm(1e4, 3e4)},
		{name: "gnm-n=1e4-m=1e5", g: gnm(1e4, 1e5)},
		{name: "gnm-n=1e4-m=3e5", g: gnm(1e4, 3e5)},

		{name: "gnm-n=1e5-m=1e5", g: gnm(1e5, 1e5)},
		{name: "gnm-n=1e5-m=3e5", g: gnm(1e5, 3e5)},
		{name: "gnm-n=1e5-m=1e6", g: gnm(1e5, 1e6)},
		{name: "gnm-n=1e5-m=3e6", g: gnm(1e5, 3e6)},

		{name: "gnm-n=1e6-m=1e6", g: gnm(1e6, 1e6)},
		{name: "gnm-n=1e6-m=3e6", g: gnm(1e6, 3e6)},
		{name: "gnm-n=1e6-m=1e7", g: gnm(1e6, 1e7)},
		{name: "gnm-n=1e6-m=3e7", g: gnm(1e6, 3e7)},

		{name: "dup-n=1e3-d=0.8-a=0.1", g: duplication(1e3, 0.8, 0.1, math.NaN())},
		{name: "dup-n=1e3-d=0.5-a=0.2", g: duplication(1e3, 0.5, 0.2, math.NaN())},

		{name: "dup-n=1e4-d=0.8-a=0.1", g: duplication(1e4, 0.8, 0.1, math.NaN())},
		{name: "dup-n=1e4-d=0.5-a=0.2", g: duplication(1e4, 0.5, 0.2, math.NaN())},

		{name: "dup-n=1e5-d=0.8-a=0.1", g: duplication(1e5, 0.8, 0.1, math.NaN())},
		{name: "dup-n=1e5-d=0.5-a=0.2", g: duplication(1e5, 0.5, 0.2, math.NaN())},
	}

	for _, test := range tests {
		rnd := rand.New(rand.NewSource(1))
		g := test.g()

		// Guess a maximally expensive entry to the graph.
		sort, err := topo.Sort(g)
		root := sort[0]
		if root == nil {
			// If we did not get a node in the first position
			// then there must be an unorderable set of nodes
			// in the first position of the error. Pick one
			// of the nodes at random.
			unordered := err.(topo.Unorderable)
			root = unordered[0][rnd.Intn(len(unordered[0]))]
		}
		if root == nil {
			b.Error("no entry node label for graph")
			continue
		}

		if len(sort) > 1 {
			// Ensure that the graph has a complete path
			// through the sorted nodes.

			// unordered will only be accessed if there is
			// a sort element that is nil, in which case
			// unordered will contain a set of nodes from
			// an SCC.
			unordered, _ := err.(topo.Unorderable)

			var ui int
			for i, v := range sort[1:] {
				u := sort[i]
				if u == nil {
					u = unordered[ui][rnd.Intn(len(unordered[ui]))]
					ui++
				}
				if v == nil {
					v = unordered[ui][rnd.Intn(len(unordered[ui]))]
				}
				if !g.HasEdgeFromTo(u.ID(), v.ID()) {
					g.SetEdge(g.NewEdge(u, v))
				}
			}
		}

		b.Run(test.name, func(b *testing.B) {
			for i := 0; i < b.N; i++ {
				d := Dominators(root, g)
				if got := d.Root(); got.ID() != root.ID() {
					b.Fatalf("unexpected root node: got:%d want:%d", got.ID(), root.ID())
				}
			}
		})
	}
}

// gnm returns a directed G(n,m) Erdõs-Rényi graph.
func gnm(n, m int) func() *simple.DirectedGraph {
	return func() *simple.DirectedGraph {
		dg := simple.NewDirectedGraph()
		err := gen.Gnm(dg, n, m, rand.New(rand.NewSource(1)))
		if err != nil {
			panic(err)
		}
		return dg
	}
}

// duplication returns an edge-induced directed subgraph of a
// duplication graph.
func duplication(n int, delta, alpha, sigma float64) func() *simple.DirectedGraph {
	return func() *simple.DirectedGraph {
		g := undirected{simple.NewDirectedGraph()}
		rnd := rand.New(rand.NewSource(1))
		err := gen.Duplication(g, n, delta, alpha, sigma, rnd)
		if err != nil {
			panic(err)
		}
		for _, e := range graph.EdgesOf(g.Edges()) {
			if rnd.Intn(2) == 0 {
				g.RemoveEdge(e.From().ID(), e.To().ID())
			}
		}
		return g.DirectedGraph
	}
}

type undirected struct {
	*simple.DirectedGraph
}

func (g undirected) From(id int64) graph.Nodes {
	return iterator.NewOrderedNodes(append(
		graph.NodesOf(g.DirectedGraph.From(id)),
		graph.NodesOf(g.DirectedGraph.To(id))...))
}

func (g undirected) HasEdgeBetween(xid, yid int64) bool {
	return g.DirectedGraph.HasEdgeFromTo(xid, yid)
}

func (g undirected) EdgeBetween(xid, yid int64) graph.Edge {
	return g.DirectedGraph.Edge(xid, yid)
}

func (g undirected) SetEdge(e graph.Edge) {
	g.DirectedGraph.SetEdge(e)
	g.DirectedGraph.SetEdge(g.DirectedGraph.NewEdge(e.To(), e.From()))
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"
	"time"

	"github.com/Symantec/Dominator/lib/constants"
	"github.com/Symantec/Dominator/lib/flags/loadflags"
	"github.com/Symantec/Dominator/lib/log"
	"github.com/Symantec/Dominator/lib/log/cmdlogger"
	"github.com/Symantec/Dominator/lib/srpc/setupclient"
	"github.com/Symantec/Dominator/lib/tags"
)

var (
	fleetManagerHostname = flag.String("fleetManagerHostname", "",
		"Hostname of Fleet Manager")
	fleetManagerPortNum = flag.Uint("fleetManagerPortNum",
		constants.FleetManagerPortNumber,
		"Port number of Fleet Resource Manager")
	hypervisorHostname = flag.String("hypervisorHostname", "",
		"Hostname of hypervisor")
	hypervisorPortNum = flag.Uint("hypervisorPortNum",
		constants.HypervisorPortNumber, "Port number of hypervisor")
	hypervisorTags      tags.Tags
	imageServerHostname = flag.String("imageServerHostname", "localhost",
		"Hostname of image server")
	imageServerPortNum = flag.Uint("imageServerPortNum",
		constants.ImageServerPortNumber,
		"Port number of image server")
	location = flag.String("location", "",
		"Location to search for hypervisors")
	offerTimeout = flag.Duration("offerTimeout", time.Minute+time.Second,
		"How long to offer DHCP OFFERs and ACKs")
	netbootFiles        tags.Tags
	netbootFilesTimeout = flag.Duration("netbootFilesTimeout",
		time.Minute+time.Second, "How long to provide extra files via TFTP")
	netbootTimeout = flag.Duration("netbootTimeout", time.Minute,
		"Time to wait for DHCP ACKs to be sent")
	numAcknowledgementsToWaitFor = flag.Uint("numAcknowledgementsToWaitFor",
		2, "Number of DHCP ACKs to wait for")
	topologyDir = flag.String("topologyDir", "",
		"Name of local topology directory in Git repository")

	logger log.DebugLogger
)

func init() {
	flag.Var(&hypervisorTags, "hypervisorTags", "Tags to apply to Hypervisor")
	flag.Var(&netbootFiles, "netbootFiles", "Extra files served by TFTP server")
}

func printUsage() {
	fmt.Fprintln(os.Stderr,
		"Usage: hyper-control [flags...] command [args...]")
	fmt.Fprintln(os.Stderr, "Common flags:")
	flag.PrintDefaults()
	fmt.Fprintln(os.Stderr, "Commands:")
	fmt.Fprintln(os.Stderr, "  add-address MACaddr [IPaddr]")
	fmt.Fprintln(os.Stderr, "  add-subnet ID IPgateway IPmask DNSserver...")
	fmt.Fprintln(os.Stderr, "  change-tags")
	fmt.Fprintln(os.Stderr, "  get-machine-info hostname")
	fmt.Fprintln(os.Stderr, "  get-updates")
	fmt.Fprintln(os.Stderr, "  netboot-host hostname")
	fmt.Fprintln(os.Stderr, "  netboot-machine MACaddr IPaddr [hostname]")
	fmt.Fprintln(os.Stderr, "  remove-excess-addresses MaxFreeAddr")
	fmt.Fprintln(os.Stderr, "  rollout-image name")
}

type commandFunc func([]string, log.DebugLogger)

type subcommand struct {
	command string
	minArgs int
	maxArgs int
	cmdFunc commandFunc
}

var subcommands = []subcommand{
	{"add-address", 1, 2, addAddressSubcommand},
	{"add-subnet", 4, -1, addSubnetSubcommand},
	{"change-tags", 0, 0, changeTagsSubcommand},
	{"get-machine-info", 1, 1, getMachineInfoSubcommand},
	{"get-updates", 0, 0, getUpdatesSubcommand},
	{"netboot-host", 1, 1, netbootHostSubcommand},
	{"netboot-machine", 2, 3, netbootMachineSubcommand},
	{"remove-excess-addresses", 1, 1, removeExcessAddressesSubcommand},
	{"rollout-image", 1, 1, rolloutImageSubcommand},
}

func main() {
	if err := loadflags.LoadForCli("hyper-control"); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
	flag.Usage = printUsage
	flag.Parse()
	if flag.NArg() < 1 {
		printUsage()
		os.Exit(2)
	}
	logger = cmdlogger.New()
	if err := setupclient.SetupTls(false); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
	numSubcommandArgs := flag.NArg() - 1
	for _, subcommand := range subcommands {
		if flag.Arg(0) == subcommand.command {
			if numSubcommandArgs < subcommand.minArgs ||
				(subcommand.maxArgs >= 0 &&
					numSubcommandArgs > subcommand.maxArgs) {
				printUsage()
				os.Exit(2)
			}
			subcommand.cmdFunc(flag.Args()[1:], logger)
			os.Exit(3)
		}
	}
	printUsage()
	os.Exit(2)
}
<commit_msg>Update documenation for -netbootFilesTimeout option in hyper-control tool.<commit_after>package main

import (
	"flag"
	"fmt"
	"os"
	"time"

	"github.com/Symantec/Dominator/lib/constants"
	"github.com/Symantec/Dominator/lib/flags/loadflags"
	"github.com/Symantec/Dominator/lib/log"
	"github.com/Symantec/Dominator/lib/log/cmdlogger"
	"github.com/Symantec/Dominator/lib/srpc/setupclient"
	"github.com/Symantec/Dominator/lib/tags"
)

var (
	fleetManagerHostname = flag.String("fleetManagerHostname", "",
		"Hostname of Fleet Manager")
	fleetManagerPortNum = flag.Uint("fleetManagerPortNum",
		constants.FleetManagerPortNumber,
		"Port number of Fleet Resource Manager")
	hypervisorHostname = flag.String("hypervisorHostname", "",
		"Hostname of hypervisor")
	hypervisorPortNum = flag.Uint("hypervisorPortNum",
		constants.HypervisorPortNumber, "Port number of hypervisor")
	hypervisorTags      tags.Tags
	imageServerHostname = flag.String("imageServerHostname", "localhost",
		"Hostname of image server")
	imageServerPortNum = flag.Uint("imageServerPortNum",
		constants.ImageServerPortNumber,
		"Port number of image server")
	location = flag.String("location", "",
		"Location to search for hypervisors")
	offerTimeout = flag.Duration("offerTimeout", time.Minute+time.Second,
		"How long to offer DHCP OFFERs and ACKs")
	netbootFiles        tags.Tags
	netbootFilesTimeout = flag.Duration("netbootFilesTimeout",
		time.Minute+time.Second,
		"How long to provide files via TFTP after last DHCP ACK")
	netbootTimeout = flag.Duration("netbootTimeout", time.Minute,
		"Time to wait for DHCP ACKs to be sent")
	numAcknowledgementsToWaitFor = flag.Uint("numAcknowledgementsToWaitFor",
		2, "Number of DHCP ACKs to wait for")
	topologyDir = flag.String("topologyDir", "",
		"Name of local topology directory in Git repository")

	logger log.DebugLogger
)

func init() {
	flag.Var(&hypervisorTags, "hypervisorTags", "Tags to apply to Hypervisor")
	flag.Var(&netbootFiles, "netbootFiles", "Extra files served by TFTP server")
}

func printUsage() {
	fmt.Fprintln(os.Stderr,
		"Usage: hyper-control [flags...] command [args...]")
	fmt.Fprintln(os.Stderr, "Common flags:")
	flag.PrintDefaults()
	fmt.Fprintln(os.Stderr, "Commands:")
	fmt.Fprintln(os.Stderr, "  add-address MACaddr [IPaddr]")
	fmt.Fprintln(os.Stderr, "  add-subnet ID IPgateway IPmask DNSserver...")
	fmt.Fprintln(os.Stderr, "  change-tags")
	fmt.Fprintln(os.Stderr, "  get-machine-info hostname")
	fmt.Fprintln(os.Stderr, "  get-updates")
	fmt.Fprintln(os.Stderr, "  netboot-host hostname")
	fmt.Fprintln(os.Stderr, "  netboot-machine MACaddr IPaddr [hostname]")
	fmt.Fprintln(os.Stderr, "  remove-excess-addresses MaxFreeAddr")
	fmt.Fprintln(os.Stderr, "  rollout-image name")
}

type commandFunc func([]string, log.DebugLogger)

type subcommand struct {
	command string
	minArgs int
	maxArgs int
	cmdFunc commandFunc
}

var subcommands = []subcommand{
	{"add-address", 1, 2, addAddressSubcommand},
	{"add-subnet", 4, -1, addSubnetSubcommand},
	{"change-tags", 0, 0, changeTagsSubcommand},
	{"get-machine-info", 1, 1, getMachineInfoSubcommand},
	{"get-updates", 0, 0, getUpdatesSubcommand},
	{"netboot-host", 1, 1, netbootHostSubcommand},
	{"netboot-machine", 2, 3, netbootMachineSubcommand},
	{"remove-excess-addresses", 1, 1, removeExcessAddressesSubcommand},
	{"rollout-image", 1, 1, rolloutImageSubcommand},
}

func main() {
	if err := loadflags.LoadForCli("hyper-control"); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
	flag.Usage = printUsage
	flag.Parse()
	if flag.NArg() < 1 {
		printUsage()
		os.Exit(2)
	}
	logger = cmdlogger.New()
	if err := setupclient.SetupTls(false); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
	numSubcommandArgs := flag.NArg() - 1
	for _, subcommand := range subcommands {
		if flag.Arg(0) == subcommand.command {
			if numSubcommandArgs < subcommand.minArgs ||
				(subcommand.maxArgs >= 0 &&
					numSubcommandArgs > subcommand.maxArgs) {
				printUsage()
				os.Exit(2)
			}
			subcommand.cmdFunc(flag.Args()[1:], logger)
			os.Exit(3)
		}
	}
	printUsage()
	os.Exit(2)
}
<|endoftext|>


<commit_before>package main

import (
	"launchpad.net/gnuflag"
	"launchpad.net/juju/go/cmd"
	"launchpad.net/juju/go/environs"
	"launchpad.net/juju/go/log"
	"launchpad.net/juju/go/state"
	"launchpad.net/tomb"

	// register providers
	_ "launchpad.net/juju/go/environs/dummy"
	_ "launchpad.net/juju/go/environs/ec2"
)

// ProvisioningAgent is a cmd.Command responsible for running a provisioning agent.
type ProvisioningAgent struct {
	Conf AgentConf
}

// Info returns usage information for the command.
func (a *ProvisioningAgent) Info() *cmd.Info {
	return &cmd.Info{"provisioning", "", "run a juju provisioning agent", ""}
}

// Init initializes the command for running.
func (a *ProvisioningAgent) Init(f *gnuflag.FlagSet, args []string) error {
	a.Conf.addFlags(f)
	if err := f.Parse(true, args); err != nil {
		return err
	}
	return a.Conf.checkArgs(f.Args())
}

// Run runs a provisioning agent.
func (a *ProvisioningAgent) Run(_ *cmd.Context) error {
	st, err := state.Open(&a.Conf.StateInfo)
	if err != nil {
		return err
	}
	p := NewProvisioner(st)
	return p.tomb.Wait()
}

type Provisioner struct {
	st      *state.State
	environ environs.Environ
	tomb    tomb.Tomb

	environWatcher  *state.ConfigWatcher
	machinesWatcher *state.MachinesWatcher
}

// environChanges returns a channel that will receive the new *ConfigNode
// when a change is detected.
func (p *Provisioner) environChanges() <-chan *state.ConfigNode {
	if p.environWatcher == nil {
		p.environWatcher = p.st.WatchEnvironConfig()
	}
	return p.environWatcher.Changes()
}

// stopEnvironWatcher stops and invalidates the current environWatcher.
func (p *Provisioner) stopEnvironWatcher() (err error) {
	if p.environWatcher != nil {
		err = p.environWatcher.Stop()
	}
	p.environWatcher = nil
	return
}

// changes returns a channel that will receive the new *ConfigNode when a
// change is detected.
func (p *Provisioner) machinesChanges() <-chan *state.MachinesChange {
	if p.machinesWatcher == nil {
		p.machinesWatcher = p.st.WatchMachines()
	}
	return p.machinesWatcher.Changes()
}

// stopMachinesWatcher stops and invalidates the current machinesWatcher..
func (p *Provisioner) stopMachinesWatcher() (err error) {
	if p.machinesWatcher != nil {
		err = p.machinesWatcher.Stop()
	}
	p.machinesWatcher = nil
	return
}

// NewProvisioner returns a Provisioner.
func NewProvisioner(st *state.State) *Provisioner {
	p := &Provisioner{
		st: st,
	}
	go p.loop()
	return p
}

func (p *Provisioner) loop() {
	defer p.tomb.Done()
	// TODO(dfc) we need a method like state.IsValid() here to exit cleanly if
	// there is a connection problem.
	for {
		select {
		case <-p.tomb.Dying():
			return
		case config, ok := <-p.environChanges():
			if !ok {
				p.stopEnvironWatcher()
				continue
			}
			var err error
			p.environ, err = environs.NewEnviron(config.Map())
			if err != nil {
				log.Printf("provisioner: unable to create environment from supplied configuration: %v", err)
				continue
			}
			log.Printf("provisioner: valid environment configured")
			p.innerLoop()
		}
	}
}

func (p *Provisioner) innerLoop() {
	// TODO(dfc) we need a method like state.IsValid() here to exit cleanly if
	// there is a connection problem.
	for {
		select {
		case <-p.tomb.Dying():
			return
		case change, ok := <-p.environChanges():
			if !ok {
				p.stopEnvironWatcher()
				continue
			}
			config, err := environs.NewConfig(change.Map())
			if err != nil {
				log.Printf("provisioner: new configuration received, but was not valid: %v", err)
				continue
			}
			p.environ.SetConfig(config)
			log.Printf("provisioner: new environment configuartion applied")
		case machines, ok := <-p.machinesChanges():
			if !ok {
				p.stopMachinesWatcher()
				continue
			}
			p.processMachines(machines)
		}
	}
}

// Stop stops the Provisioner and returns any error encountered while
// provisioning.
func (p *Provisioner) Stop() error {
	p.tomb.Kill(nil)
	err1 := p.stopEnvironWatcher()
	err2 := p.stopMachinesWatcher()
	err3 := p.tomb.Wait()
	for _, err := range []error{err3, err2, err1} {
		if err != nil {
			return err
		}
	}
	return nil
}

func (p *Provisioner) processMachines(changes *state.MachinesChange) {}
<commit_msg>log stop errors<commit_after>package main

import (
	"launchpad.net/gnuflag"
	"launchpad.net/juju/go/cmd"
	"launchpad.net/juju/go/environs"
	"launchpad.net/juju/go/log"
	"launchpad.net/juju/go/state"
	"launchpad.net/tomb"

	// register providers
	_ "launchpad.net/juju/go/environs/dummy"
	_ "launchpad.net/juju/go/environs/ec2"
)

// ProvisioningAgent is a cmd.Command responsible for running a provisioning agent.
type ProvisioningAgent struct {
	Conf AgentConf
}

// Info returns usage information for the command.
func (a *ProvisioningAgent) Info() *cmd.Info {
	return &cmd.Info{"provisioning", "", "run a juju provisioning agent", ""}
}

// Init initializes the command for running.
func (a *ProvisioningAgent) Init(f *gnuflag.FlagSet, args []string) error {
	a.Conf.addFlags(f)
	if err := f.Parse(true, args); err != nil {
		return err
	}
	return a.Conf.checkArgs(f.Args())
}

// Run runs a provisioning agent.
func (a *ProvisioningAgent) Run(_ *cmd.Context) error {
	st, err := state.Open(&a.Conf.StateInfo)
	if err != nil {
		return err
	}
	p := NewProvisioner(st)
	return p.tomb.Wait()
}

type Provisioner struct {
	st      *state.State
	environ environs.Environ
	tomb    tomb.Tomb

	environWatcher  *state.ConfigWatcher
	machinesWatcher *state.MachinesWatcher
}

// environChanges returns a channel that will receive the new *ConfigNode
// when a change is detected.
func (p *Provisioner) environChanges() <-chan *state.ConfigNode {
	if p.environWatcher == nil {
		p.environWatcher = p.st.WatchEnvironConfig()
	}
	return p.environWatcher.Changes()
}

// stopEnvironWatcher stops and invalidates the current environWatcher.
func (p *Provisioner) stopEnvironWatcher() (err error) {
	if p.environWatcher != nil {
		if err = p.environWatcher.Stop(); err != nil {
			log.Printf("provisioning: environWatcher reported error on Stop: %v", err)
		}
	}
	p.environWatcher = nil
	return
}

// changes returns a channel that will receive the new *ConfigNode when a
// change is detected.
func (p *Provisioner) machinesChanges() <-chan *state.MachinesChange {
	if p.machinesWatcher == nil {
		p.machinesWatcher = p.st.WatchMachines()
	}
	return p.machinesWatcher.Changes()
}

// stopMachinesWatcher stops and invalidates the current machinesWatcher..
func (p *Provisioner) stopMachinesWatcher() (err error) {
	if p.machinesWatcher != nil {
		if err = p.machinesWatcher.Stop(); err != nil {
			log.Printf("provisioning: machinesWatcehr reported error on Stop: %v", err)
		}
	}
	p.machinesWatcher = nil
	return
}

// NewProvisioner returns a Provisioner.
func NewProvisioner(st *state.State) *Provisioner {
	p := &Provisioner{
		st: st,
	}
	go p.loop()
	return p
}

func (p *Provisioner) loop() {
	defer p.tomb.Done()
	// TODO(dfc) we need a method like state.IsValid() here to exit cleanly if
	// there is a connection problem.
	for {
		select {
		case <-p.tomb.Dying():
			return
		case config, ok := <-p.environChanges():
			if !ok {
				p.stopEnvironWatcher()
				continue
			}
			var err error
			p.environ, err = environs.NewEnviron(config.Map())
			if err != nil {
				log.Printf("provisioner: unable to create environment from supplied configuration: %v", err)
				continue
			}
			log.Printf("provisioner: valid environment configured")
			p.innerLoop()
		}
	}
}

func (p *Provisioner) innerLoop() {
	// TODO(dfc) we need a method like state.IsValid() here to exit cleanly if
	// there is a connection problem.
	for {
		select {
		case <-p.tomb.Dying():
			return
		case change, ok := <-p.environChanges():
			if !ok {
				p.stopEnvironWatcher()
				continue
			}
			config, err := environs.NewConfig(change.Map())
			if err != nil {
				log.Printf("provisioner: new configuration received, but was not valid: %v", err)
				continue
			}
			p.environ.SetConfig(config)
			log.Printf("provisioner: new environment configuartion applied")
		case machines, ok := <-p.machinesChanges():
			if !ok {
				p.stopMachinesWatcher()
				continue
			}
			p.processMachines(machines)
		}
	}
}

// Stop stops the Provisioner and returns any error encountered while
// provisioning.
func (p *Provisioner) Stop() error {
	p.tomb.Kill(nil)
	err1 := p.stopEnvironWatcher()
	err2 := p.stopMachinesWatcher()
	err3 := p.tomb.Wait()
	for _, err := range []error{err3, err2, err1} {
		if err != nil {
			return err
		}
	}
	return nil
}

func (p *Provisioner) processMachines(changes *state.MachinesChange) {}
<|endoftext|>


<commit_before>/*
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements.  See the NOTICE file
distributed with this work for additional information
regarding copyright ownership.  The ASF licenses this file
to you 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 gremlingo

import (
	"fmt"
	"github.com/cucumber/godog"
	gremlingo "github.com/lyndonb-bq/tinkerpop/gremlin-go/driver"
	"os"
	"reflect"
	"strconv"
)

type TinkerPopWorld struct {
	scenario     *godog.Scenario
	g            *gremlingo.GraphTraversalSource
	graphName    string
	traversal    *gremlingo.GraphTraversal
	result       []interface{}
	graphDataMap map[string]*DataGraph
	parameters   map[string]interface{}
}

type DataGraph struct {
	name       string
	connection *gremlingo.DriverRemoteConnection
	vertices   map[string]*gremlingo.Vertex
	edges      map[string]*gremlingo.Edge
}

func getEnvOrDefaultString(key string, defaultValue string) string {
	// Missing value is returned as "".
	value := os.Getenv(key)
	if len(value) != 0 {
		return value
	}
	return defaultValue
}

func getEnvOrDefaultInt(key string, defaultValue int) int {
	value := getEnvOrDefaultString(key, "")
	if len(value) != 0 {
		intValue, err := strconv.Atoi(value)
		if err == nil {
			return intValue
		}
	}
	return defaultValue
}

func scenarioUrl() string {
	return getEnvOrDefaultString("GREMLIN_SERVER_URL", "ws://localhost:8182/gremlin")
}

func NewTinkerPopWorld() *TinkerPopWorld {
	return &TinkerPopWorld{
		scenario:     nil,
		g:            nil,
		graphName:    "",
		traversal:    nil,
		result:       nil,
		graphDataMap: make(map[string]*DataGraph),
		parameters:   make(map[string]interface{}),
	}
}

var graphNames = []string{"modern", "classic", "crew", "grateful", "sink", "empty"}

func (t *TinkerPopWorld) getDataGraphFromMap(name string) *DataGraph {
	if val, ok := t.graphDataMap[name]; ok {
		return val
	} else {
		return nil
	}
}

func (t *TinkerPopWorld) loadAllDataGraph() {
	for _, name := range graphNames {
		if name == "empty" {
			t.loadEmptyDataGraph()
		} else {
			connection, err := gremlingo.NewDriverRemoteConnection(scenarioUrl(),
				func(settings *gremlingo.DriverRemoteConnectionSettings) {
					settings.TraversalSource = "g" + name
				})
			if err != nil {
				panic(fmt.Sprintf("Failed to create connection '%v'", err))
			}
			g := gremlingo.Traversal_().WithRemote(connection)
			t.graphDataMap[name] = &DataGraph{
				name:       name,
				connection: connection,
				vertices:   getVertices(g),
				edges:      getEdges(g),
			}
		}
	}
}

func (t *TinkerPopWorld) loadEmptyDataGraph() {
	connection, _ := gremlingo.NewDriverRemoteConnection(scenarioUrl(), func(settings *gremlingo.DriverRemoteConnectionSettings) {
		settings.TraversalSource = "ggraph"
	})
	t.graphDataMap["empty"] = &DataGraph{connection: connection}
}

func (t *TinkerPopWorld) reloadEmptyData() {
	graphData := t.getDataGraphFromMap("empty")
	g := gremlingo.Traversal_().WithRemote(graphData.connection)
	graphData.vertices = getVertices(g)
	graphData.edges = getEdges(g)
}

func (t *TinkerPopWorld) cleanEmptyDataGraph(g *gremlingo.GraphTraversalSource) error {
	_, future, err := g.V().Drop().Iterate()
	if err != nil {
		return err
	}
	<-future
	return nil
}

func getVertices(g *gremlingo.GraphTraversalSource) map[string]*gremlingo.Vertex {
	vertexMap := make(map[string]*gremlingo.Vertex)
	res, err := g.V().Group().By("name").By(gremlingo.T__.Tail()).Next()
	if res == nil {
		return nil
	}
	if err != nil {
		return nil
	}
	v := reflect.ValueOf(res.GetInterface())
	if v.Kind() != reflect.Map {
		fmt.Printf("Expecting to get a map as a result, got %v instead.", v.Kind())
		return nil
	}
	keys := v.MapKeys()
	for _, k := range keys {
		convKey := k.Convert(v.Type().Key())
		val := v.MapIndex(convKey)
		vertexMap[k.Interface().(string)] = val.Interface().(*gremlingo.Vertex)
	}
	return vertexMap
}

func getEdges(g *gremlingo.GraphTraversalSource) map[string]*gremlingo.Edge {
	edgeMap := make(map[string]*gremlingo.Edge)
	resE, err := g.E().Group().By(gremlingo.T__.Project("o", "l", "i").
		By(gremlingo.T__.OutV().Values("name")).By(gremlingo.T__.Label()).By(gremlingo.T__.InV().Values("name"))).
		By(gremlingo.T__.Tail()).Next()
	if err != nil {
		return nil
	}
	valMap := reflect.ValueOf(resE.GetInterface())
	if valMap.Kind() != reflect.Map {
		fmt.Printf("Expecting to get a map as a result, got %v instead.", valMap.Kind())
		return nil
	}
	keys := valMap.MapKeys()
	for _, k := range keys {
		convKey := k.Convert(valMap.Type().Key())
		val := valMap.MapIndex(convKey)
		keyMap := reflect.ValueOf(k.Interface()).Elem().Interface().(map[interface{}]interface{})
		edgeMap[getEdgeKey(keyMap)] = val.Interface().(*gremlingo.Edge)
	}
	return edgeMap
}

func getEdgeKey(edgeKeyMap map[interface{}]interface{}) string {
	return fmt.Sprint(edgeKeyMap["o"], "-", edgeKeyMap["l"], "->", edgeKeyMap["i"])
}

// This function is used to isolate connection problems to each scenario, and used in the Before context hook to prevent
// a failing test in one scenario closing the shared connection that leads to failing subsequent scenario tests.
// This function can be removed once all pending tests pass.
func (t *TinkerPopWorld) recreateAllDataGraphConnection() error {
	var err error
	for _, name := range graphNames {
		if name == "empty" {
			t.getDataGraphFromMap(name).connection, err = gremlingo.NewDriverRemoteConnection(scenarioUrl(), func(settings *gremlingo.DriverRemoteConnectionSettings) {
				settings.TraversalSource = "ggraph"
			})
		} else {
			t.getDataGraphFromMap(name).connection, err = gremlingo.NewDriverRemoteConnection(scenarioUrl(), func(settings *gremlingo.DriverRemoteConnectionSettings) {
				settings.TraversalSource = "g" + name
			})
		}
	}
	return err
}

func (t *TinkerPopWorld) closeAllDataGraphConnection() error {
	for _, name := range graphNames {
		err := t.getDataGraphFromMap(name).connection.Close()
		if err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>AN-1063 reuse connection for loadEmptyDataGraph<commit_after>/*
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements.  See the NOTICE file
distributed with this work for additional information
regarding copyright ownership.  The ASF licenses this file
to you 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 gremlingo

import (
	"fmt"
	"github.com/cucumber/godog"
	gremlingo "github.com/lyndonb-bq/tinkerpop/gremlin-go/driver"
	"os"
	"reflect"
	"strconv"
)

type TinkerPopWorld struct {
	scenario     *godog.Scenario
	g            *gremlingo.GraphTraversalSource
	graphName    string
	traversal    *gremlingo.GraphTraversal
	result       []interface{}
	graphDataMap map[string]*DataGraph
	parameters   map[string]interface{}
}

type DataGraph struct {
	name       string
	connection *gremlingo.DriverRemoteConnection
	vertices   map[string]*gremlingo.Vertex
	edges      map[string]*gremlingo.Edge
}

func getEnvOrDefaultString(key string, defaultValue string) string {
	// Missing value is returned as "".
	value := os.Getenv(key)
	if len(value) != 0 {
		return value
	}
	return defaultValue
}

func getEnvOrDefaultInt(key string, defaultValue int) int {
	value := getEnvOrDefaultString(key, "")
	if len(value) != 0 {
		intValue, err := strconv.Atoi(value)
		if err == nil {
			return intValue
		}
	}
	return defaultValue
}

func scenarioUrl() string {
	return getEnvOrDefaultString("GREMLIN_SERVER_URL", "ws://localhost:8182/gremlin")
}

func NewTinkerPopWorld() *TinkerPopWorld {
	return &TinkerPopWorld{
		scenario:     nil,
		g:            nil,
		graphName:    "",
		traversal:    nil,
		result:       nil,
		graphDataMap: make(map[string]*DataGraph),
		parameters:   make(map[string]interface{}),
	}
}

var graphNames = []string{"modern", "classic", "crew", "grateful", "sink", "empty"}

func (t *TinkerPopWorld) getDataGraphFromMap(name string) *DataGraph {
	if val, ok := t.graphDataMap[name]; ok {
		return val
	} else {
		return nil
	}
}

func (t *TinkerPopWorld) loadAllDataGraph() {
	for _, name := range graphNames {
		if name == "empty" {
			t.loadEmptyDataGraph()
		} else {
			connection, err := gremlingo.NewDriverRemoteConnection(scenarioUrl(),
				func(settings *gremlingo.DriverRemoteConnectionSettings) {
					settings.TraversalSource = "g" + name
				})
			if err != nil {
				panic(fmt.Sprintf("Failed to create connection '%v'", err))
			}
			g := gremlingo.Traversal_().WithRemote(connection)
			t.graphDataMap[name] = &DataGraph{
				name:       name,
				connection: connection,
				vertices:   getVertices(g),
				edges:      getEdges(g),
			}
		}
	}
}

func (t *TinkerPopWorld) loadEmptyDataGraph() {
	if _, ok := t.graphDataMap["empty"]; ok {
		return
	}

	connection, _ := gremlingo.NewDriverRemoteConnection(scenarioUrl(), func(settings *gremlingo.DriverRemoteConnectionSettings) {
		settings.TraversalSource = "ggraph"
	})
	t.graphDataMap["empty"] = &DataGraph{connection: connection}
}

func (t *TinkerPopWorld) reloadEmptyData() {
	graphData := t.getDataGraphFromMap("empty")
	g := gremlingo.Traversal_().WithRemote(graphData.connection)
	graphData.vertices = getVertices(g)
	graphData.edges = getEdges(g)
}

func (t *TinkerPopWorld) cleanEmptyDataGraph(g *gremlingo.GraphTraversalSource) error {
	_, future, err := g.V().Drop().Iterate()
	if err != nil {
		return err
	}
	<-future
	return nil
}

func getVertices(g *gremlingo.GraphTraversalSource) map[string]*gremlingo.Vertex {
	vertexMap := make(map[string]*gremlingo.Vertex)
	res, err := g.V().Group().By("name").By(gremlingo.T__.Tail()).Next()
	if res == nil {
		return nil
	}
	if err != nil {
		return nil
	}
	v := reflect.ValueOf(res.GetInterface())
	if v.Kind() != reflect.Map {
		fmt.Printf("Expecting to get a map as a result, got %v instead.", v.Kind())
		return nil
	}
	keys := v.MapKeys()
	for _, k := range keys {
		convKey := k.Convert(v.Type().Key())
		val := v.MapIndex(convKey)
		vertexMap[k.Interface().(string)] = val.Interface().(*gremlingo.Vertex)
	}
	return vertexMap
}

func getEdges(g *gremlingo.GraphTraversalSource) map[string]*gremlingo.Edge {
	edgeMap := make(map[string]*gremlingo.Edge)
	resE, err := g.E().Group().By(gremlingo.T__.Project("o", "l", "i").
		By(gremlingo.T__.OutV().Values("name")).By(gremlingo.T__.Label()).By(gremlingo.T__.InV().Values("name"))).
		By(gremlingo.T__.Tail()).Next()
	if err != nil {
		return nil
	}
	valMap := reflect.ValueOf(resE.GetInterface())
	if valMap.Kind() != reflect.Map {
		fmt.Printf("Expecting to get a map as a result, got %v instead.", valMap.Kind())
		return nil
	}
	keys := valMap.MapKeys()
	for _, k := range keys {
		convKey := k.Convert(valMap.Type().Key())
		val := valMap.MapIndex(convKey)
		keyMap := reflect.ValueOf(k.Interface()).Elem().Interface().(map[interface{}]interface{})
		edgeMap[getEdgeKey(keyMap)] = val.Interface().(*gremlingo.Edge)
	}
	return edgeMap
}

func getEdgeKey(edgeKeyMap map[interface{}]interface{}) string {
	return fmt.Sprint(edgeKeyMap["o"], "-", edgeKeyMap["l"], "->", edgeKeyMap["i"])
}

// This function is used to isolate connection problems to each scenario, and used in the Before context hook to prevent
// a failing test in one scenario closing the shared connection that leads to failing subsequent scenario tests.
// This function can be removed once all pending tests pass.
func (t *TinkerPopWorld) recreateAllDataGraphConnection() error {
	var err error
	for _, name := range graphNames {
		if name == "empty" {
			t.getDataGraphFromMap(name).connection, err = gremlingo.NewDriverRemoteConnection(scenarioUrl(), func(settings *gremlingo.DriverRemoteConnectionSettings) {
				settings.TraversalSource = "ggraph"
			})
		} else {
			t.getDataGraphFromMap(name).connection, err = gremlingo.NewDriverRemoteConnection(scenarioUrl(), func(settings *gremlingo.DriverRemoteConnectionSettings) {
				settings.TraversalSource = "g" + name
			})
		}
	}
	return err
}

func (t *TinkerPopWorld) closeAllDataGraphConnection() error {
	for _, name := range graphNames {
		err := t.getDataGraphFromMap(name).connection.Close()
		if err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2020 Google LLC

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 (
	"encoding/base64"
	"flag"

	"github.com/grafeas/kritis/pkg/kritis/apis/kritis/v1beta1"
	"github.com/grafeas/kritis/pkg/kritis/crd/vulnzsigningpolicy"
	"github.com/grafeas/kritis/pkg/kritis/metadata/containeranalysis"
	"github.com/grafeas/kritis/pkg/kritis/secrets"
	"github.com/grafeas/kritis/pkg/kritis/signer"
	"google.golang.org/api/option"
	"io/ioutil"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	"github.com/golang/glog"
	"gopkg.in/yaml.v2"
)

func main() {
	var image, json_key_path, pri_key_path, policy_path string

	flag.StringVar(&image, "image", "", "image url, e.g., gcr.io/foo/bar@sha256:abcd")
	flag.StringVar(&json_key_path, "credentials", "", "json credentials file path, e.g., ./key.json")
	flag.StringVar(&pri_key_path, "private_key", "", "signer private key path, e.g., /dev/shm/key.pgp")
	flag.StringVar(&policy_path, "policy", "", "vulnerability signing policy file path, e.g., /tmp/vulnz_signing_policy.yaml")
	flag.Parse()

	glog.Infof("image: %s, json_path: %s, s_key: %s, policy: %s", image, json_key_path, pri_key_path, policy_path)

	signerKey, err := ioutil.ReadFile(pri_key_path)
	if err != nil {
		glog.Fatalf("Fail to read signer key: %v", err)
	}

	policyFile, err := ioutil.ReadFile(policy_path)
	if err != nil {
		glog.Fatalf("Fail to read vulnz signing policy: %v", err)
	} else {
		glog.Infof("Policy file: %v\n", string(policyFile))
	}

	// Parse the vulnz signing policy
	policy := v1beta1.VulnzSigningPolicy{}

	err = yaml.Unmarshal(policyFile, &policy)
	if err != nil {
		glog.Fatalf("Fail to parse policy file: %v", err)
	} else {
		glog.Infof("Policy noteReference: %v\n", policy.Spec.NoteReference)
	}

	// Read the vulnz scanning events
	if image == "" {
		glog.Fatalf("Image url is empty: %s", image)
	}

	d, err := containeranalysis.New(option.WithCredentialsFile(json_key_path))
	if err != nil {
		glog.Fatalf("Could not initialize the client %v", err)
	}
	vulnz, err := d.Vulnerabilities(image)
	if err != nil {
		glog.Fatalf("Found err %s", err)
	}
	if vulnz == nil {
		glog.Fatalf("Expected some vulnerabilities. Nil found")
	}

	//fmt.Printf("policy noteReference %s\n", policy.Spec.NoteReference)
	// fmt.Printf("signer_key %v\n", signerKey)

	// Run the signer
	client, err := containeranalysis.NewCache(option.WithCredentialsFile(json_key_path))
	if err != nil {
		glog.Fatalf("Error getting Container Analysis client: %v", err)
	}

	// Create pgp key
	// Assume empty passphrase
	// TODO: support non-empty passphrase
	passphrase := ""
	pgpKey, err := secrets.NewPgpSigningKey(string(signerKey), passphrase)
	if err != nil {
		glog.Fatalf("Creating pgp key from files fail: %v\nprivate key:\n%s\n", err, string(signerKey))
	}
	// Create AA
	// Create an AttestaionAuthority to help create noteOcurrences.
	// This is quite hacky.
	// TODO: refactor out the authority code
	authority := v1beta1.AttestationAuthority{
		ObjectMeta: metav1.ObjectMeta{Name: "signing-aa"},
		Spec: v1beta1.AttestationAuthoritySpec{
			NoteReference: policy.Spec.NoteReference,
			PublicKeyList: []string{base64.StdEncoding.EncodeToString([]byte(""))},
		},
	}

	r := signer.New(client, &signer.Config{
		Validate:  vulnzsigningpolicy.ValidateVulnzSigningPolicy,
		PgpKey:    pgpKey,
		Authority: authority,
		Project:   policy.Spec.Project,
	})
	imageVulnz := signer.ImageVulnerabilities{
		ImageRef:        image,
		Vulnerabilities: vulnz,
	}

	if err := r.ValidateAndSign(imageVulnz, policy); err != nil {
		glog.Fatalf("Error creating signature: %v", err)
	}
}
<commit_msg>Add passphrase support for signing key.<commit_after>/*
Copyright 2020 Google LLC

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 (
	"encoding/base64"
	"flag"

	"github.com/grafeas/kritis/pkg/kritis/apis/kritis/v1beta1"
	"github.com/grafeas/kritis/pkg/kritis/crd/vulnzsigningpolicy"
	"github.com/grafeas/kritis/pkg/kritis/metadata/containeranalysis"
	"github.com/grafeas/kritis/pkg/kritis/secrets"
	"github.com/grafeas/kritis/pkg/kritis/signer"
	"google.golang.org/api/option"
	"io/ioutil"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	"github.com/golang/glog"
	"gopkg.in/yaml.v2"
)

func main() {
	var image, json_key_path, passphrase, pri_key_path, policy_path string

	flag.StringVar(&image, "image", "", "image url, e.g., gcr.io/foo/bar@sha256:abcd")
	flag.StringVar(&json_key_path, "credentials", "", "json credentials file path, e.g., ./key.json")
	flag.StringVar(&pri_key_path, "private_key", "", "signer private key path, e.g., /dev/shm/key.pgp")
	flag.StringVar(&passphrase, "passphrase", "", "passphrase for private key, if any")
	flag.StringVar(&policy_path, "policy", "", "vulnerability signing policy file path, e.g., /tmp/vulnz_signing_policy.yaml")
	flag.Parse()

	glog.Infof("image: %s, json_path: %s, s_key: %s, policy: %s", image, json_key_path, pri_key_path, policy_path)

	signerKey, err := ioutil.ReadFile(pri_key_path)
	if err != nil {
		glog.Fatalf("Fail to read signer key: %v", err)
	}

	policyFile, err := ioutil.ReadFile(policy_path)
	if err != nil {
		glog.Fatalf("Fail to read vulnz signing policy: %v", err)
	} else {
		glog.Infof("Policy file: %v\n", string(policyFile))
	}

	// Parse the vulnz signing policy
	policy := v1beta1.VulnzSigningPolicy{}

	err = yaml.Unmarshal(policyFile, &policy)
	if err != nil {
		glog.Fatalf("Fail to parse policy file: %v", err)
	} else {
		glog.Infof("Policy noteReference: %v\n", policy.Spec.NoteReference)
	}

	// Read the vulnz scanning events
	if image == "" {
		glog.Fatalf("Image url is empty: %s", image)
	}

	d, err := containeranalysis.New(option.WithCredentialsFile(json_key_path))
	if err != nil {
		glog.Fatalf("Could not initialize the client %v", err)
	}
	vulnz, err := d.Vulnerabilities(image)
	if err != nil {
		glog.Fatalf("Found err %s", err)
	}
	if vulnz == nil {
		glog.Fatalf("Expected some vulnerabilities. Nil found")
	}

	//fmt.Printf("policy noteReference %s\n", policy.Spec.NoteReference)
	// fmt.Printf("signer_key %v\n", signerKey)

	// Run the signer
	client, err := containeranalysis.NewCache(option.WithCredentialsFile(json_key_path))
	if err != nil {
		glog.Fatalf("Error getting Container Analysis client: %v", err)
	}

	// Create pgp key
	pgpKey, err := secrets.NewPgpSigningKey(string(signerKey), passphrase)
	if err != nil {
		glog.Fatalf("Creating pgp key from files fail: %v\nprivate key:\n%s\n", err, string(signerKey))
	}
	// Create AA
	// Create an AttestaionAuthority to help create noteOcurrences.
	// This is quite hacky.
	// TODO: refactor out the authority code
	authority := v1beta1.AttestationAuthority{
		ObjectMeta: metav1.ObjectMeta{Name: "signing-aa"},
		Spec: v1beta1.AttestationAuthoritySpec{
			NoteReference: policy.Spec.NoteReference,
			PublicKeyList: []string{base64.StdEncoding.EncodeToString([]byte(""))},
		},
	}

	r := signer.New(client, &signer.Config{
		Validate:  vulnzsigningpolicy.ValidateVulnzSigningPolicy,
		PgpKey:    pgpKey,
		Authority: authority,
		Project:   policy.Spec.Project,
	})
	imageVulnz := signer.ImageVulnerabilities{
		ImageRef:        image,
		Vulnerabilities: vulnz,
	}

	if err := r.ValidateAndSign(imageVulnz, policy); err != nil {
		glog.Fatalf("Error creating signature: %v", err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	_ "expvar"
	"flag"
	"fmt"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"runtime"

	"github.com/docker/distribution/health"
	"github.com/sirupsen/logrus"
	"github.com/theupdateframework/notary/server"
	"github.com/theupdateframework/notary/utils"
	"github.com/theupdateframework/notary/version"
)

// DebugAddress is the debug server address to listen on
const (
	jsonLogFormat = "json"
	DebugAddress  = "localhost:8080"
	envPrefix     = "NOTARY_SERVER"
)

type cmdFlags struct {
	debug       bool
	logFormat   string
	configFile  string
	doBootstrap bool
	version     bool
}

func setupFlags(flagStorage *cmdFlags) {
	// Setup flags
	flag.StringVar(&flagStorage.configFile, "config", "", "Path to configuration file")
	flag.BoolVar(&flagStorage.debug, "debug", false, "Enable the debugging server on localhost:8080")
	flag.StringVar(&flagStorage.logFormat, "logf", "json", "Set the format of the logs. Only 'json' and 'logfmt' are supported at the moment.")
	flag.BoolVar(&flagStorage.doBootstrap, "bootstrap", false, "Do any necessary setup of configured backend storage services")
	flag.BoolVar(&flagStorage.version, "version", false, "Print the version number of notary-server")

	// this needs to be in init so that _ALL_ logs are in the correct format
	if flagStorage.logFormat == jsonLogFormat {
		logrus.SetFormatter(new(logrus.JSONFormatter))
	}

	flag.Usage = usage
}

func main() {
	flagStorage := cmdFlags{}
	setupFlags(&flagStorage)

	flag.Parse()

	if flagStorage.version {
		fmt.Println("notary-server " + getVersion())
		os.Exit(0)
	}

	if flagStorage.debug {
		go debugServer(DebugAddress)
	}

	// when the server starts print the version for debugging and issue logs later
	logrus.Info(getVersion())

	ctx, serverConfig, err := parseServerConfig(flagStorage.configFile, health.RegisterPeriodicFunc, flagStorage.doBootstrap)
	if err != nil {
		logrus.Fatal(err.Error())
	}

	c := utils.SetupSignalTrap(utils.LogLevelSignalHandle)
	if c != nil {
		defer signal.Stop(c)
	}

	if flagStorage.doBootstrap {
		err = bootstrap(ctx)
	} else {
		logrus.Info("Starting Server")
		err = server.Run(ctx, serverConfig)
	}

	if err != nil {
		logrus.Fatal(err.Error())
	}
	return
}

func usage() {
	fmt.Println("usage:", os.Args[0])
	flag.PrintDefaults()
}

func getVersion() string {
	return fmt.Sprintf("Version: %s, Git commit: %s, Go version: %s", version.NotaryVersion, version.GitCommit, runtime.Version())
}

// debugServer starts the debug server with pprof, expvar among other
// endpoints. The addr should not be exposed externally. For most of these to
// work, tls cannot be enabled on the endpoint, so it is generally separate.
func debugServer(addr string) {
	logrus.Infof("Debug server listening on %s", addr)
	if err := http.ListenAndServe(addr, nil); err != nil {
		logrus.Fatalf("error listening on debug interface: %v", err)
	}
}
<commit_msg>gosec: ignore  G108: Profiling endpoint automatically exposed<commit_after>package main

import (
	_ "expvar"
	"flag"
	"fmt"
	"net/http"
	_ "net/http/pprof" // #nosec G108 // false positive as it's only listening through debugServer()
	"os"
	"os/signal"
	"runtime"

	"github.com/docker/distribution/health"
	"github.com/sirupsen/logrus"
	"github.com/theupdateframework/notary/server"
	"github.com/theupdateframework/notary/utils"
	"github.com/theupdateframework/notary/version"
)

// DebugAddress is the debug server address to listen on
const (
	jsonLogFormat = "json"
	DebugAddress  = "localhost:8080"
	envPrefix     = "NOTARY_SERVER"
)

type cmdFlags struct {
	debug       bool
	logFormat   string
	configFile  string
	doBootstrap bool
	version     bool
}

func setupFlags(flagStorage *cmdFlags) {
	// Setup flags
	flag.StringVar(&flagStorage.configFile, "config", "", "Path to configuration file")
	flag.BoolVar(&flagStorage.debug, "debug", false, "Enable the debugging server on localhost:8080")
	flag.StringVar(&flagStorage.logFormat, "logf", "json", "Set the format of the logs. Only 'json' and 'logfmt' are supported at the moment.")
	flag.BoolVar(&flagStorage.doBootstrap, "bootstrap", false, "Do any necessary setup of configured backend storage services")
	flag.BoolVar(&flagStorage.version, "version", false, "Print the version number of notary-server")

	// this needs to be in init so that _ALL_ logs are in the correct format
	if flagStorage.logFormat == jsonLogFormat {
		logrus.SetFormatter(new(logrus.JSONFormatter))
	}

	flag.Usage = usage
}

func main() {
	flagStorage := cmdFlags{}
	setupFlags(&flagStorage)

	flag.Parse()

	if flagStorage.version {
		fmt.Println("notary-server " + getVersion())
		os.Exit(0)
	}

	if flagStorage.debug {
		go debugServer(DebugAddress)
	}

	// when the server starts print the version for debugging and issue logs later
	logrus.Info(getVersion())

	ctx, serverConfig, err := parseServerConfig(flagStorage.configFile, health.RegisterPeriodicFunc, flagStorage.doBootstrap)
	if err != nil {
		logrus.Fatal(err.Error())
	}

	c := utils.SetupSignalTrap(utils.LogLevelSignalHandle)
	if c != nil {
		defer signal.Stop(c)
	}

	if flagStorage.doBootstrap {
		err = bootstrap(ctx)
	} else {
		logrus.Info("Starting Server")
		err = server.Run(ctx, serverConfig)
	}

	if err != nil {
		logrus.Fatal(err.Error())
	}
	return
}

func usage() {
	fmt.Println("usage:", os.Args[0])
	flag.PrintDefaults()
}

func getVersion() string {
	return fmt.Sprintf("Version: %s, Git commit: %s, Go version: %s", version.NotaryVersion, version.GitCommit, runtime.Version())
}

// debugServer starts the debug server with pprof, expvar among other
// endpoints. The addr should not be exposed externally. For most of these to
// work, tls cannot be enabled on the endpoint, so it is generally separate.
func debugServer(addr string) {
	logrus.Infof("Debug server listening on %s", addr)
	if err := http.ListenAndServe(addr, nil); err != nil {
		logrus.Fatalf("error listening on debug interface: %v", err)
	}
}
<|endoftext|>


<commit_before>// Jamie: This contains the go source code that will become shield.

package main

import (
	"fmt"
	"github.com/starkandwayne/shield/db"
	"github.com/voxelbrain/goptions"
	"os"

	// sql drivers
	_ "github.com/mattn/go-sqlite3"
)

func main() {
	fmt.Printf("starting up...\n")

	options := struct {
		Driver string `goptions:"-t,--type, obligatory, description='Type of database backend'"`
		DSN    string `goptions:"-d,--database, obligatory, description='DSN of the database backend'"`
	}{
	// No defaults
	}
	goptions.ParseAndFail(&options)

	database := &db.DB{
		Driver: options.Driver,
		DSN:    options.DSN,
	}

	fmt.Fprintf(os.Stderr, "connecting to %s database at %s\n", database.Driver, database.DSN)
	if err := database.Connect(); err != nil {
		fmt.Fprintf(os.Stderr, "failed to connect to %s database at %s: %s\n",
			database.Driver, database.DSN, err)
	}

	if err := database.Setup(); err != nil {
		fmt.Fprintf(os.Stderr, "failed to set up schema in %s database at %s: %s\n",
			database.Driver, database.DSN, err)
		return
	}

	fmt.Printf("deployed schema version %d\n", db.CurrentSchema)
}
<commit_msg>Enable Postgres support in shield-schema<commit_after>// Jamie: This contains the go source code that will become shield.

package main

import (
	"fmt"
	"github.com/starkandwayne/shield/db"
	"github.com/voxelbrain/goptions"
	"os"

	// sql drivers
	_ "github.com/mattn/go-sqlite3"
	_ "github.com/lib/pq"
)

func main() {
	fmt.Printf("starting up...\n")

	options := struct {
		Driver string `goptions:"-t,--type, obligatory, description='Type of database backend'"`
		DSN    string `goptions:"-d,--database, obligatory, description='DSN of the database backend'"`
	}{
	// No defaults
	}
	goptions.ParseAndFail(&options)

	database := &db.DB{
		Driver: options.Driver,
		DSN:    options.DSN,
	}

	fmt.Fprintf(os.Stderr, "connecting to %s database at %s\n", database.Driver, database.DSN)
	if err := database.Connect(); err != nil {
		fmt.Fprintf(os.Stderr, "failed to connect to %s database at %s: %s\n",
			database.Driver, database.DSN, err)
	}

	if err := database.Setup(); err != nil {
		fmt.Fprintf(os.Stderr, "failed to set up schema in %s database at %s: %s\n",
			database.Driver, database.DSN, err)
		return
	}

	fmt.Printf("deployed schema version %d\n", db.CurrentSchema)
}
<|endoftext|>


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

package main

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"
)

// Round 1.
const (
	secretStr = "pibud-sijat-ponam-zizaz.kudol-visin-vakok-jinok"
	publicKey = `p256
20475414006091125411730282854763965332579614918776190347990649355528840360162
41618798560597642013440926161855187887081385971895806061707694318148863738083
`

	privateKey = "103735382135370212717736500933863354513183407328603457343387144070761075604179\n"
)

// Round 2.
const (
	secretStr2 = "ponam-sijat-pibud-zizaz.vakok-visin-kudol-jinok"

	public2Key = `p256
12123951731103713463684549996372980607529899544492565862335281576402253981962
115773437512973772745017348056058736133905153504498008582636958858837556112862
`

	private2Key = "47513031334211958720530809945101319621000940818220052385943490959145831252616\n"

	archive2Key = `# EE
p256
20475414006091125411730282854763965332579614918776190347990649355528840360162
41618798560597642013440926161855187887081385971895806061707694318148863738083
103735382135370212717736500933863354513183407328603457343387144070761075604179
`
)

func TestSaveKeygen(t *testing.T) {
	public, private, _, err := createKeys("p256", secretStr)
	if err != nil {
		t.Fatalf("creating keys: %v", err)
	}
	if public != publicKey {
		t.Errorf("round 1: got public key %q; want %q", public, publicKey)
	}
	if private != privateKey {
		t.Errorf("round 1: got private key %q; want %q", private, privateKey)
	}

	// Write them to a file.
	dir, err := ioutil.TempDir("", "keygen")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(dir)
	err = writeKeys(dir, public, private, "")
	if err != nil {
		t.Fatalf("writing keys: %v", err)
	}

	// Read them back.
	data, err := ioutil.ReadFile(filepath.Join(dir, "public.upspinkey"))
	if err != nil {
		t.Fatal(err)
	}
	if string(data) != public {
		t.Fatalf("reading public key: got %q; want %q", data, public)
	}
	data, err = ioutil.ReadFile(filepath.Join(dir, "secret.upspinkey"))
	if err != nil {
		t.Fatal(err)
	}
	if string(data) != private {
		t.Fatalf("reading secret key: got %q; want %q", data, private)
	}

	// Generate again.
	public, private, _, err = createKeys("p256", secretStr2)
	if err != nil {
		t.Fatalf("creating keys: %v", err)
	}
	if public != public2Key {
		t.Errorf("round 2: got public key %q; want %q", public, publicKey)
	}
	if private != private2Key {
		t.Errorf("round 2: got private key %q; want %q", private, private2Key)
	}

	// Save, then write them to a file again.
	err = saveKeys(dir)
	if err != nil {
		t.Fatalf("saving keys: %v", err)
	}
	err = writeKeys(dir, public, private, "")
	if err != nil {
		t.Fatalf("writing keys: %v", err)
	}

	// Read them back.
	data, err = ioutil.ReadFile(filepath.Join(dir, "public.upspinkey"))
	if err != nil {
		t.Fatal(err)
	}
	if string(data) != public {
		t.Fatalf("reading public key: got %q; want %q", data, public)
	}
	data, err = ioutil.ReadFile(filepath.Join(dir, "secret.upspinkey"))
	if err != nil {
		t.Fatal(err)
	}
	if string(data) != private {
		t.Fatalf("reading secret key: got %q; want %q", data, private)
	}

	// Now check the archive.
	data, err = ioutil.ReadFile(filepath.Join(dir, "secret2.upspinkey"))
	if err != nil {
		t.Fatal(err)
	}
	if string(data) != archive2Key {
		t.Fatalf("reading archive key: got %q; want %q", data, archive2Key)
	}
}
<commit_msg>cmd/upspin: fix keygen test<commit_after>// Copyright 2016 The Upspin Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"io/ioutil"
	"os"
	"path/filepath"
	"testing"
)

// Round 1.
const (
	secretStr = "pibud-sijat-ponam-zizaz.kudol-visin-vakok-jinok"
	publicKey = `p256
20475414006091125411730282854763965332579614918776190347990649355528840360162
41618798560597642013440926161855187887081385971895806061707694318148863738083
`

	privateKey = "103735382135370212717736500933863354513183407328603457343387144070761075604179\n"
)

// Round 2.
const (
	secretStr2 = "ponam-sijat-pibud-zizaz.vakok-visin-kudol-jinok"

	public2Key = `p256
12123951731103713463684549996372980607529899544492565862335281576402253981962
115773437512973772745017348056058736133905153504498008582636958858837556112862
`

	private2Key = "47513031334211958720530809945101319621000940818220052385943490959145831252616\n"

	archive2Key = `# EE
p256
20475414006091125411730282854763965332579614918776190347990649355528840360162
41618798560597642013440926161855187887081385971895806061707694318148863738083
103735382135370212717736500933863354513183407328603457343387144070761075604179
`
)

func TestSaveKeygen(t *testing.T) {
	public, private, _, err := createKeys("p256", secretStr)
	if err != nil {
		t.Fatalf("creating keys: %v", err)
	}
	if public != publicKey {
		t.Errorf("round 1: got public key %q; want %q", public, publicKey)
	}
	if private != privateKey {
		t.Errorf("round 1: got private key %q; want %q", private, privateKey)
	}

	// Write them to a file.
	dir, err := ioutil.TempDir("", "keygen")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(dir)
	err = writeKeys(dir, public, private)
	if err != nil {
		t.Fatalf("writing keys: %v", err)
	}

	// Read them back.
	data, err := ioutil.ReadFile(filepath.Join(dir, "public.upspinkey"))
	if err != nil {
		t.Fatal(err)
	}
	if string(data) != public {
		t.Fatalf("reading public key: got %q; want %q", data, public)
	}
	data, err = ioutil.ReadFile(filepath.Join(dir, "secret.upspinkey"))
	if err != nil {
		t.Fatal(err)
	}
	if string(data) != private {
		t.Fatalf("reading secret key: got %q; want %q", data, private)
	}

	// Generate again.
	public, private, _, err = createKeys("p256", secretStr2)
	if err != nil {
		t.Fatalf("creating keys: %v", err)
	}
	if public != public2Key {
		t.Errorf("round 2: got public key %q; want %q", public, publicKey)
	}
	if private != private2Key {
		t.Errorf("round 2: got private key %q; want %q", private, private2Key)
	}

	// Save, then write them to a file again.
	err = saveKeys(dir)
	if err != nil {
		t.Fatalf("saving keys: %v", err)
	}
	err = writeKeys(dir, public, private)
	if err != nil {
		t.Fatalf("writing keys: %v", err)
	}

	// Read them back.
	data, err = ioutil.ReadFile(filepath.Join(dir, "public.upspinkey"))
	if err != nil {
		t.Fatal(err)
	}
	if string(data) != public {
		t.Fatalf("reading public key: got %q; want %q", data, public)
	}
	data, err = ioutil.ReadFile(filepath.Join(dir, "secret.upspinkey"))
	if err != nil {
		t.Fatal(err)
	}
	if string(data) != private {
		t.Fatalf("reading secret key: got %q; want %q", data, private)
	}

	// Now check the archive.
	data, err = ioutil.ReadFile(filepath.Join(dir, "secret2.upspinkey"))
	if err != nil {
		t.Fatal(err)
	}
	if string(data) != archive2Key {
		t.Fatalf("reading archive key: got %q; want %q", data, archive2Key)
	}
}
<|endoftext|>