<commit_before>package main

import (
	"flag"
	"fmt"
	"doozer/client"
	"doozer/lock"
	"doozer/member"
	"doozer/mon"
	"doozer/paxos"
	"doozer/server"
	"doozer/session"
	"doozer/store"
	"doozer/util"
	"doozer/web"
	"net"
	"os"
	"time"
)

const (
	alpha    = 50
	interval = 1e9 // ns == 1s
)

// Flags
var (
	listenAddr  = flag.String("l", "127.0.0.1:8046", "The address to bind to.")
	attachAddr  = flag.String("a", "", "The address of another node to attach to.")
	webAddr     = flag.String("w", ":8080", "Serve web requests on this address.")
	clusterName = flag.String("c", "local", "The non-empty cluster name.")
)

func activate(st *store.Store, self, prefix string, c *client.Client, cal chan int) {
	logger := util.NewLogger("activate")
	ch := make(chan store.Event)
	st.GetDirAndWatch("/doozer/slot", ch)
	for ev := range ch {
		// TODO ev.IsEmpty()
		if ev.IsSet() && ev.Body == "" {
			_, err := c.Set(prefix+ev.Path, self, ev.Cas)
			if err != nil {
				logger.Println(err)
				continue
			}
			cal <- 1
			cal <- 1
			close(ch)
		}
	}
}

func advanceUntil(cl *client.Client, done chan int) {
	for _, ok := <-done; !ok; _, ok = <-done {
		cl.Noop()
	}
}

func Usage() {
	fmt.Fprintf(os.Stderr, "Usage: %s [OPTIONS] <cluster-name>\n", os.Args[0])
	fmt.Fprintf(os.Stderr, "\nOptions:\n")
	flag.PrintDefaults()
}

func main() {
	util.LogWriter = os.Stderr
	logger := util.NewLogger("main")

	flag.Parse()
	flag.Usage = Usage

	prefix := "/d/" + *clusterName

	if *listenAddr == "" {
		logger.Println("require a listen address")
		flag.Usage()
		os.Exit(1)
	}

	var webListener net.Listener
	if *webAddr != "" {
		wl, err := net.Listen("tcp", *webAddr)
		if err != nil {
			panic(err)
		}
		webListener = wl
	}

	listener, err := net.Listen("tcp", *listenAddr)
	if err != nil {
		panic(err)
	}

	outs := make(paxos.ChanPutCloserTo)

	cal := make(chan int, 2)

	var cl *client.Client
	self := util.RandId()
	st := store.New()
	seqn := uint64(0)
	if *attachAddr == "" { // we are the only node in a new cluster
		seqn = addPublicAddr(st, seqn+1, self, *listenAddr)
		seqn = addHostname(st, seqn+1, self, os.Getenv("HOSTNAME"))
		seqn = addMember(st, seqn+1, self, *listenAddr)
		seqn = claimSlot(st, seqn+1, "1", self)
		seqn = claimLeader(st, seqn+1, self)
		//seqn = claimSlot(st, seqn+1, "2", "")
		//seqn = claimSlot(st, seqn+1, "3", "")
		//seqn = claimSlot(st, seqn+1, "4", "")
		//seqn = claimSlot(st, seqn+1, "5", "")
		seqn = addPing(st, seqn+1, "pong")

		cal <- 1
		cal <- 1

		cl, err = client.Dial(*listenAddr)
		if err != nil {
			panic(err)
		}
	} else {
		cl, err = client.Dial(*attachAddr)
		if err != nil {
			panic(err)
		}

		path := prefix + "/doozer/info/" + self + "/public-addr"
		_, err = cl.Set(path, *listenAddr, store.Clobber)
		if err != nil {
			panic(err)
		}

		path = prefix + "/doozer/info/" + self + "/hostname"
		_, err = cl.Set(path, os.Getenv("HOSTNAME"), store.Clobber)
		if err != nil {
			panic(err)
		}

		var snap string
		seqn, snap, err = cl.Join(self, *listenAddr)
		if err != nil {
			panic(err)
		}

		done := make(chan int)
		st.Ops <- store.Op{1, snap}

		go advanceUntil(cl, done)

		go func() {
			st.Sync(seqn + alpha)
			close(done)
			activate(st, self, prefix, cl, cal)
		}()

		// TODO sink needs a way to pick up missing values if there are any
		// gaps in its sequence
	}
	mg := paxos.NewManager(self, seqn, alpha, st, st.Ops, outs)

	if *attachAddr == "" {
		// Skip ahead alpha steps so that the registrar can provide a
		// meaningful cluster.
		for i := seqn + 1; i < seqn+alpha; i++ {
			st.Ops <- store.Op{i, store.Nop}
		}
	}

	go func() {
		<-cal
		go lock.Clean(st, mg)
		go session.Clean(st, mg)
		go member.Clean(st, mg)
	}()

	sv := &server.Server{*listenAddr, st, mg, self, prefix}

	go func() {
		cas := store.Missing
		for _ = range time.Tick(interval) {
			_, cas, err = cl.Checkin(self, cas)
			if err != nil {
				logger.Println(err)
			}
		}
	}()

	go func() {
		panic(mon.Monitor(self, prefix, st, cl))
	}()

	go func() {
		panic(sv.Serve(listener, cal))
	}()

	if webListener != nil {
		web.Store = st
		web.ClusterName = *clusterName
		go web.Serve(webListener)
	}

	panic(sv.ListenAndServeUdp(outs))
}

func addPublicAddr(st *store.Store, seqn uint64, self, addr string) uint64 {
	// TODO pull out path as a const
	path := "/doozer/info/" + self + "/public-addr"
	mx, err := store.EncodeSet(path, addr, store.Missing)
	if err != nil {
		panic(err)
	}
	st.Ops <- store.Op{seqn, mx}
	return seqn
}

func addHostname(st *store.Store, seqn uint64, self, addr string) uint64 {
	// TODO pull out path as a const
	path := "/doozer/info/" + self + "/hostname"
	mx, err := store.EncodeSet(path, addr, store.Missing)
	if err != nil {
		panic(err)
	}
	st.Ops <- store.Op{seqn, mx}
	return seqn
}

func addMember(st *store.Store, seqn uint64, self, addr string) uint64 {
	// TODO pull out path as a const
	mx, err := store.EncodeSet("/doozer/members/"+self, addr, store.Missing)
	if err != nil {
		panic(err)
	}
	st.Ops <- store.Op{seqn, mx}
	return seqn
}

func claimSlot(st *store.Store, seqn uint64, slot, self string) uint64 {
	// TODO pull out path as a const
	mx, err := store.EncodeSet("/doozer/slot/"+slot, self, store.Missing)
	if err != nil {
		panic(err)
	}
	st.Ops <- store.Op{seqn, mx}
	return seqn
}

func claimLeader(st *store.Store, seqn uint64, self string) uint64 {
	// TODO pull out path as a const
	mx, err := store.EncodeSet("/doozer/leader", self, store.Missing)
	if err != nil {
		panic(err)
	}
	st.Ops <- store.Op{seqn, mx}
	return seqn
}

func addPing(st *store.Store, seqn uint64, v string) uint64 {
	// TODO pull out path as a const
	mx, err := store.EncodeSet("/ping", v, store.Missing)
	if err != nil {
		panic(err)
	}
	st.Ops <- store.Op{seqn, mx}
	return seqn
}
<commit_msg>pulsing<commit_after>package main

import (
	"flag"
	"fmt"
	"doozer/client"
	"doozer/lock"
	"doozer/gc"
	"doozer/member"
	"doozer/mon"
	"doozer/paxos"
	"doozer/server"
	"doozer/session"
	"doozer/store"
	"doozer/util"
	"doozer/web"
	"net"
	"os"
	"time"
)

const (
	alpha    = 50
	checkinInterval = 1e9 // ns == 1s
	pulseInterval = 1e9
)

// Flags
var (
	listenAddr  = flag.String("l", "127.0.0.1:8046", "The address to bind to.")
	attachAddr  = flag.String("a", "", "The address of another node to attach to.")
	webAddr     = flag.String("w", ":8080", "Serve web requests on this address.")
	clusterName = flag.String("c", "local", "The non-empty cluster name.")
)

func activate(st *store.Store, self, prefix string, c *client.Client, cal chan int) {
	logger := util.NewLogger("activate")
	ch := make(chan store.Event)
	st.GetDirAndWatch("/doozer/slot", ch)
	for ev := range ch {
		// TODO ev.IsEmpty()
		if ev.IsSet() && ev.Body == "" {
			_, err := c.Set(prefix+ev.Path, self, ev.Cas)
			if err != nil {
				logger.Println(err)
				continue
			}
			cal <- 1
			cal <- 1
			close(ch)
		}
	}
}

func advanceUntil(cl *client.Client, done chan int) {
	for _, ok := <-done; !ok; _, ok = <-done {
		cl.Noop()
	}
}

func Usage() {
	fmt.Fprintf(os.Stderr, "Usage: %s [OPTIONS] <cluster-name>\n", os.Args[0])
	fmt.Fprintf(os.Stderr, "\nOptions:\n")
	flag.PrintDefaults()
}

func main() {
	util.LogWriter = os.Stderr
	logger := util.NewLogger("main")

	flag.Parse()
	flag.Usage = Usage

	prefix := "/d/" + *clusterName

	if *listenAddr == "" {
		logger.Println("require a listen address")
		flag.Usage()
		os.Exit(1)
	}

	var webListener net.Listener
	if *webAddr != "" {
		wl, err := net.Listen("tcp", *webAddr)
		if err != nil {
			panic(err)
		}
		webListener = wl
	}

	listener, err := net.Listen("tcp", *listenAddr)
	if err != nil {
		panic(err)
	}

	outs := make(paxos.ChanPutCloserTo)

	cal := make(chan int, 2)

	var cl *client.Client
	self := util.RandId()
	st := store.New()
	seqn := uint64(0)
	if *attachAddr == "" { // we are the only node in a new cluster
		seqn = addPublicAddr(st, seqn+1, self, *listenAddr)
		seqn = addHostname(st, seqn+1, self, os.Getenv("HOSTNAME"))
		seqn = addMember(st, seqn+1, self, *listenAddr)
		seqn = claimSlot(st, seqn+1, "1", self)
		seqn = claimLeader(st, seqn+1, self)
		//seqn = claimSlot(st, seqn+1, "2", "")
		//seqn = claimSlot(st, seqn+1, "3", "")
		//seqn = claimSlot(st, seqn+1, "4", "")
		//seqn = claimSlot(st, seqn+1, "5", "")
		seqn = addPing(st, seqn+1, "pong")

		cal <- 1
		cal <- 1

		cl, err = client.Dial(*listenAddr)
		if err != nil {
			panic(err)
		}
	} else {
		cl, err = client.Dial(*attachAddr)
		if err != nil {
			panic(err)
		}

		path := prefix + "/doozer/info/" + self + "/public-addr"
		_, err = cl.Set(path, *listenAddr, store.Clobber)
		if err != nil {
			panic(err)
		}

		path = prefix + "/doozer/info/" + self + "/hostname"
		_, err = cl.Set(path, os.Getenv("HOSTNAME"), store.Clobber)
		if err != nil {
			panic(err)
		}

		var snap string
		seqn, snap, err = cl.Join(self, *listenAddr)
		if err != nil {
			panic(err)
		}

		done := make(chan int)
		st.Ops <- store.Op{1, snap}

		go advanceUntil(cl, done)

		go func() {
			st.Sync(seqn + alpha)
			close(done)
			activate(st, self, prefix, cl, cal)
		}()

		// TODO sink needs a way to pick up missing values if there are any
		// gaps in its sequence
	}
	mg := paxos.NewManager(self, seqn, alpha, st, st.Ops, outs)

	if *attachAddr == "" {
		// Skip ahead alpha steps so that the registrar can provide a
		// meaningful cluster.
		for i := seqn + 1; i < seqn+alpha; i++ {
			st.Ops <- store.Op{i, store.Nop}
		}
	}

	go func() {
		<-cal
		go lock.Clean(st, mg)
		go session.Clean(st, mg)
		go member.Clean(st, mg)
		go gc.Pulse(self, st.Seqns, cl, pulseInterval)
	}()

	sv := &server.Server{*listenAddr, st, mg, self, prefix}

	go func() {
		cas := store.Missing
		for _ = range time.Tick(checkinInterval) {
			_, cas, err = cl.Checkin(self, cas)
			if err != nil {
				logger.Println(err)
			}
		}
	}()

	go func() {
		panic(mon.Monitor(self, prefix, st, cl))
	}()

	go func() {
		panic(sv.Serve(listener, cal))
	}()

	if webListener != nil {
		web.Store = st
		web.ClusterName = *clusterName
		go web.Serve(webListener)
	}

	panic(sv.ListenAndServeUdp(outs))
}

func addPublicAddr(st *store.Store, seqn uint64, self, addr string) uint64 {
	// TODO pull out path as a const
	path := "/doozer/info/" + self + "/public-addr"
	mx, err := store.EncodeSet(path, addr, store.Missing)
	if err != nil {
		panic(err)
	}
	st.Ops <- store.Op{seqn, mx}
	return seqn
}

func addHostname(st *store.Store, seqn uint64, self, addr string) uint64 {
	// TODO pull out path as a const
	path := "/doozer/info/" + self + "/hostname"
	mx, err := store.EncodeSet(path, addr, store.Missing)
	if err != nil {
		panic(err)
	}
	st.Ops <- store.Op{seqn, mx}
	return seqn
}

func addMember(st *store.Store, seqn uint64, self, addr string) uint64 {
	// TODO pull out path as a const
	mx, err := store.EncodeSet("/doozer/members/"+self, addr, store.Missing)
	if err != nil {
		panic(err)
	}
	st.Ops <- store.Op{seqn, mx}
	return seqn
}

func claimSlot(st *store.Store, seqn uint64, slot, self string) uint64 {
	// TODO pull out path as a const
	mx, err := store.EncodeSet("/doozer/slot/"+slot, self, store.Missing)
	if err != nil {
		panic(err)
	}
	st.Ops <- store.Op{seqn, mx}
	return seqn
}

func claimLeader(st *store.Store, seqn uint64, self string) uint64 {
	// TODO pull out path as a const
	mx, err := store.EncodeSet("/doozer/leader", self, store.Missing)
	if err != nil {
		panic(err)
	}
	st.Ops <- store.Op{seqn, mx}
	return seqn
}

func addPing(st *store.Store, seqn uint64, v string) uint64 {
	// TODO pull out path as a const
	mx, err := store.EncodeSet("/ping", v, store.Missing)
	if err != nil {
		panic(err)
	}
	st.Ops <- store.Op{seqn, mx}
	return seqn
}
<|endoftext|>


<commit_before>package command

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

	"github.com/funkygao/columnize"
	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/gofmt"
	zklib "github.com/samuel/go-zookeeper/zk"
)

type Dbus struct {
	Ui  cli.Ui
	Cmd string
}

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

	zkzone := zk.NewZkZone(zk.DefaultConfig(zone, ctx.ZoneZkAddrs(zone)))
	for {
		this.checkMyslave(zkzone)
		if !topMode {
			break
		}

		time.Sleep(time.Second * 3)
		refreshScreen()
	}

	return
}

type binlogCheckpoint struct {
	File   string `json:"file"`
	Offset int64  `json:"offset"`
	Owner  string `json:"owner"`
}

func (this *Dbus) checkMyslave(zkzone *zk.ZkZone) {
	lines := []string{"Mysql|File|Offset|dbus|ver|pid|uptime|msince"}
	root := "/dbus/myslave"
	dbs, _, err := zkzone.Conn().Children(root)
	swallow(err)

	sort.Strings(dbs)

	for _, db := range dbs {
		dbRoot := fmt.Sprintf("%s/%s", root, db)
		data, _, err := zkzone.Conn().Get(dbRoot)
		var v binlogCheckpoint
		swallow(json.Unmarshal(data, &v))

		ownerPath := fmt.Sprintf("%s/owner", dbRoot)
		owner, ownerStat, err := zkzone.Conn().Get(ownerPath)
		if err != nil {
			if err != zklib.ErrNoNode {
				panic(err)
			}

			// an orphan binlog stream: no dbus consuming it
			lines = append(lines, fmt.Sprintf("%s| | |%s|%s|%s|%s| ", db))
			continue
		}

		idsPath := fmt.Sprintf("%s/ids", dbRoot)
		runners, _, err := zkzone.Conn().Children(idsPath)
		swallow(err)

		for _, r := range runners {
			ver, stat, err := zkzone.Conn().Get(fmt.Sprintf("%s/%s", idsPath, r))
			swallow(err)

			i := strings.LastIndex(r, "-")
			instance := r[:i]
			pid := r[i+1:]
			uptime := zk.ZkTimestamp(stat.Ctime).Time()

			if r == string(owner) {
				ownerTime := zk.ZkTimestamp(ownerStat.Mtime).Time()
				lines = append(lines, fmt.Sprintf("%s|%s|%s|%s|%s|%s|%s|%s",
					db, v.File, gofmt.Comma(v.Offset), instance, string(ver), pid,
					gofmt.PrettySince(uptime), gofmt.PrettySince(ownerTime)))
			} else {
				lines = append(lines, fmt.Sprintf("%s| | |%s|%s|%s|%s| ",
					db, instance, string(ver), pid, gofmt.PrettySince(uptime)))
			}
		}

	}

	this.Ui.Output(columnize.SimpleFormat(lines))
}

func (*Dbus) Synopsis() string {
	return "Monitor database system runtime"
}

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

    %s

Options:

    -z zone

    -top
      Run in top mode.

`, this.Cmd, this.Synopsis())
	return strings.TrimSpace(help)
}
<commit_msg>fmt error<commit_after>package command

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

	"github.com/funkygao/columnize"
	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/gofmt"
	zklib "github.com/samuel/go-zookeeper/zk"
)

type Dbus struct {
	Ui  cli.Ui
	Cmd string
}

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

	zkzone := zk.NewZkZone(zk.DefaultConfig(zone, ctx.ZoneZkAddrs(zone)))
	for {
		this.checkMyslave(zkzone)
		if !topMode {
			break
		}

		time.Sleep(time.Second * 3)
		refreshScreen()
	}

	return
}

type binlogCheckpoint struct {
	File   string `json:"file"`
	Offset int64  `json:"offset"`
	Owner  string `json:"owner"`
}

func (this *Dbus) checkMyslave(zkzone *zk.ZkZone) {
	lines := []string{"Mysql|File|Offset|dbus|ver|pid|uptime|msince"}
	root := "/dbus/myslave"
	dbs, _, err := zkzone.Conn().Children(root)
	swallow(err)

	sort.Strings(dbs)

	for _, db := range dbs {
		dbRoot := fmt.Sprintf("%s/%s", root, db)
		data, _, err := zkzone.Conn().Get(dbRoot)
		var v binlogCheckpoint
		swallow(json.Unmarshal(data, &v))

		ownerPath := fmt.Sprintf("%s/owner", dbRoot)
		owner, ownerStat, err := zkzone.Conn().Get(ownerPath)
		if err != nil {
			if err != zklib.ErrNoNode {
				panic(err)
			}

			// an orphan binlog stream: no dbus consuming it
			lines = append(lines, fmt.Sprintf("%s|-|-|-|-|-|-|-", db))
			continue
		}

		idsPath := fmt.Sprintf("%s/ids", dbRoot)
		runners, _, err := zkzone.Conn().Children(idsPath)
		swallow(err)

		for _, r := range runners {
			ver, stat, err := zkzone.Conn().Get(fmt.Sprintf("%s/%s", idsPath, r))
			swallow(err)

			i := strings.LastIndex(r, "-")
			instance := r[:i]
			pid := r[i+1:]
			uptime := zk.ZkTimestamp(stat.Ctime).Time()

			if r == string(owner) {
				ownerTime := zk.ZkTimestamp(ownerStat.Mtime).Time()
				lines = append(lines, fmt.Sprintf("%s|%s|%s|%s|%s|%s|%s|%s",
					db, v.File, gofmt.Comma(v.Offset), instance, string(ver), pid,
					gofmt.PrettySince(uptime), gofmt.PrettySince(ownerTime)))
			} else {
				lines = append(lines, fmt.Sprintf("%s| | |%s|%s|%s|%s| ",
					db, instance, string(ver), pid, gofmt.PrettySince(uptime)))
			}
		}

	}

	this.Ui.Output(columnize.SimpleFormat(lines))
}

func (*Dbus) Synopsis() string {
	return "Monitor database system runtime"
}

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

    %s

Options:

    -z zone

    -top
      Run in top mode.

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


<commit_before>package command

import (
	"flag"
	"fmt"
	"sort"
	"strings"
	"time"

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

type Lags struct {
	Ui  cli.Ui
	Cmd string

	onlineOnly      bool
	groupPattern    string
	topicPattern    string
	watchMode       bool
	problematicMode bool
	tableFmt        bool
	lagThreshold    int
}

func (this *Lags) Run(args []string) (exitCode int) {
	const secondsInMinute = 60
	var (
		cluster string
		zone    string
	)
	cmdFlags := flag.NewFlagSet("lags", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&zone, "z", ctx.ZkDefaultZone(), "")
	cmdFlags.StringVar(&cluster, "c", "", "")
	cmdFlags.BoolVar(&this.onlineOnly, "l", false, "")
	cmdFlags.BoolVar(&this.problematicMode, "p", false, "")
	cmdFlags.StringVar(&this.groupPattern, "g", "", "")
	cmdFlags.StringVar(&this.topicPattern, "t", "", "")
	cmdFlags.BoolVar(&this.tableFmt, "table", false, "")
	cmdFlags.BoolVar(&this.watchMode, "w", false, "")
	cmdFlags.IntVar(&this.lagThreshold, "lag", 5000, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	if this.watchMode {
		refreshScreen()
	}

	if this.problematicMode {
		this.onlineOnly = true
	}

	zkzone := zk.NewZkZone(zk.DefaultConfig(zone, ctx.ZoneZkAddrs(zone)))
	bar := progress.New(secondsInMinute)
	if cluster == "" {
		for {
			zkzone.ForSortedClusters(func(zkcluster *zk.ZkCluster) {
				this.Ui.Output(zkcluster.Name())
				if this.tableFmt {
					this.printConsumersLagTable(zkcluster)
				} else {
					this.printConsumersLag(zkcluster)
				}
			})

			if this.watchMode {
				for i := 1; i <= secondsInMinute; i++ {
					bar.ShowProgress(i)
					time.Sleep(time.Second)
				}
			} else {
				break
			}

			printSwallowedErrors(this.Ui, zkzone)
		}

		return
	}

	for {
		this.Ui.Output(cluster)
		zkcluster := zkzone.NewCluster(cluster) // panic if invalid cluster
		if this.tableFmt {
			this.printConsumersLagTable(zkcluster)
		} else {
			this.printConsumersLag(zkcluster)
		}

		if this.watchMode {
			for i := 1; i <= secondsInMinute; i++ {
				bar.ShowProgress(i)
				time.Sleep(time.Second)
			}
			//refreshScreen()
		} else {
			break
		}

		printSwallowedErrors(this.Ui, zkzone)
	}

	return
}

func (this *Lags) printConsumersLagTable(zkcluster *zk.ZkCluster) {
	lines := make([]string, 0)
	header := "ConsumerGroup|Topic/Partition|Produced|Consumed|Lag|Committed|Uptime"
	lines = append(lines, header)

	// sort by group name
	consumersByGroup := zkcluster.ConsumersByGroup(this.groupPattern)
	sortedGroups := make([]string, 0, len(consumersByGroup))
	for group, _ := range consumersByGroup {
		sortedGroups = append(sortedGroups, group)
	}
	sort.Strings(sortedGroups)

	for _, group := range sortedGroups {
		if !patternMatched(group, this.groupPattern) {
			continue
		}

		sortedTopicAndPartitionIds := make([]string, 0, len(consumersByGroup[group]))
		consumers := make(map[string]zk.ConsumerMeta)
		for _, t := range consumersByGroup[group] {
			key := fmt.Sprintf("%s:%s", t.Topic, t.PartitionId)
			sortedTopicAndPartitionIds = append(sortedTopicAndPartitionIds, key)

			consumers[key] = t
		}
		sort.Strings(sortedTopicAndPartitionIds)

		for _, topicAndPartitionId := range sortedTopicAndPartitionIds {
			consumer := consumers[topicAndPartitionId]

			if !patternMatched(consumer.Topic, this.topicPattern) {
				continue
			}
			if !consumer.Online {
				continue
			}

			lines = append(lines,
				fmt.Sprintf("%s|%s/%s|%s|%s|%s|%s|%s",
					group,
					consumer.Topic, consumer.PartitionId,
					gofmt.Comma(consumer.ProducerOffset),
					gofmt.Comma(consumer.ConsumerOffset),
					gofmt.Comma(consumer.Lag),
					gofmt.PrettySince(consumer.Mtime.Time()),
					gofmt.PrettySince(consumer.ConsumerZnode.Uptime())))
		}
	}

	this.Ui.Output(columnize.SimpleFormat(lines))
}

func (this *Lags) printConsumersLag(zkcluster *zk.ZkCluster) {
	// sort by group name
	consumersByGroup := zkcluster.ConsumersByGroup(this.groupPattern)
	sortedGroups := make([]string, 0, len(consumersByGroup))
	for group, _ := range consumersByGroup {
		sortedGroups = append(sortedGroups, group)
	}
	sort.Strings(sortedGroups)

	for _, group := range sortedGroups {
		this.Ui.Output(strings.Repeat(" ", 4) + group)

		sortedTopicAndPartitionIds := make([]string, 0)
		consumers := make(map[string]zk.ConsumerMeta)
		for _, t := range consumersByGroup[group] {
			key := fmt.Sprintf("%s:%s", t.Topic, t.PartitionId)
			sortedTopicAndPartitionIds = append(sortedTopicAndPartitionIds, key)

			consumers[key] = t
		}
		sort.Strings(sortedTopicAndPartitionIds)

		for _, topicAndPartitionId := range sortedTopicAndPartitionIds {
			consumer := consumers[topicAndPartitionId]

			if !patternMatched(consumer.Topic, this.topicPattern) {
				continue
			}

			var (
				lagOutput string
				symbol    string
			)
			if consumer.Lag > int64(this.lagThreshold) {
				lagOutput = color.Red("%15s", gofmt.Comma(consumer.Lag))
				if consumer.Online {
					symbol = color.Yellow("⚠")
				} else {
					symbol = color.Yellow("◎")
				}
			} else {
				lagOutput = color.Blue("%15s", gofmt.Comma(consumer.Lag))
				if consumer.Online {
					symbol = color.Green("◉")
				} else {
					symbol = color.Yellow("◎")
				}
			}

			if consumer.Online {
				if this.problematicMode && consumer.Lag <= int64(this.lagThreshold) {
					continue
				}

				var (
					host   string
					uptime string
				)
				if consumer.ConsumerZnode == nil {
					host = "unrecognized"
					uptime = "-"
				} else {
					host = color.Green("%90s", consumer.ConsumerZnode.Host())
					uptime = gofmt.PrettySince(consumer.ConsumerZnode.Uptime())
				}

				this.Ui.Output(fmt.Sprintf("\t%s %35s/%-2s %12s -> %-12s %s %s\n%s %s",
					symbol,
					consumer.Topic, consumer.PartitionId,
					gofmt.Comma(consumer.ProducerOffset),
					gofmt.Comma(consumer.ConsumerOffset),
					lagOutput,
					gofmt.PrettySince(consumer.Mtime.Time()),
					host, uptime))
			} else if !this.onlineOnly {
				this.Ui.Output(fmt.Sprintf("\t%s %35s/%-2s %12s -> %-12s %s %s",
					symbol,
					consumer.Topic, consumer.PartitionId,
					gofmt.Comma(consumer.ProducerOffset),
					gofmt.Comma(consumer.ConsumerOffset),
					lagOutput,
					gofmt.PrettySince(consumer.Mtime.Time())))
			}
		}
	}
}

func (*Lags) Synopsis() string {
	return "Display online high level consumers lag on a topic"
}

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

    Display online high level consumers lag on a topic

Options:

    -z zone
      Default %s

    -c cluster

    -g group name pattern

    -t topic name pattern

    -w
      Watch mode. Keep printing out lags with 60s interval refresh.

    -l
      Only show online consumers lag.

    -lag threshold
      Default 5000.

    -p
      Only show problematic consumers.

    -table
      Display in table format.

`, this.Cmd, ctx.ZkDefaultZone())
	return strings.TrimSpace(help)
}
<commit_msg>tweak of 'gk lags'<commit_after>package command

import (
	"flag"
	"fmt"
	"sort"
	"strings"
	"time"

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

type Lags struct {
	Ui  cli.Ui
	Cmd string

	onlineOnly      bool
	groupPattern    string
	topicPattern    string
	watchMode       bool
	problematicMode bool
	tableFmt        bool
	lagThreshold    int
}

func (this *Lags) Run(args []string) (exitCode int) {
	const secondsInMinute = 60
	var (
		cluster string
		zone    string
	)
	cmdFlags := flag.NewFlagSet("lags", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&zone, "z", ctx.ZkDefaultZone(), "")
	cmdFlags.StringVar(&cluster, "c", "", "")
	cmdFlags.BoolVar(&this.onlineOnly, "l", false, "")
	cmdFlags.BoolVar(&this.problematicMode, "p", false, "")
	cmdFlags.StringVar(&this.groupPattern, "g", "", "")
	cmdFlags.StringVar(&this.topicPattern, "t", "", "")
	cmdFlags.BoolVar(&this.tableFmt, "table", false, "")
	cmdFlags.BoolVar(&this.watchMode, "w", false, "")
	cmdFlags.IntVar(&this.lagThreshold, "lag", 5000, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	if this.watchMode {
		refreshScreen()
	}

	if this.problematicMode {
		this.onlineOnly = true
	}

	zkzone := zk.NewZkZone(zk.DefaultConfig(zone, ctx.ZoneZkAddrs(zone)))
	bar := progress.New(secondsInMinute)
	if cluster == "" {
		for {
			zkzone.ForSortedClusters(func(zkcluster *zk.ZkCluster) {
				if this.tableFmt {
					this.printConsumersLagTable(zkcluster)
				} else {
					this.printConsumersLag(zkcluster)
				}
			})

			if this.watchMode {
				for i := 1; i <= secondsInMinute; i++ {
					bar.ShowProgress(i)
					time.Sleep(time.Second)
				}
			} else {
				break
			}

			printSwallowedErrors(this.Ui, zkzone)
		}

		return
	}

	// display a single cluster
	for {
		zkcluster := zkzone.NewCluster(cluster) // panic if invalid cluster
		if this.tableFmt {
			this.printConsumersLagTable(zkcluster)
		} else {
			this.printConsumersLag(zkcluster)
		}

		if this.watchMode {
			for i := 1; i <= secondsInMinute; i++ {
				bar.ShowProgress(i)
				time.Sleep(time.Second)
			}
			//refreshScreen()
		} else {
			break
		}

		printSwallowedErrors(this.Ui, zkzone)
	}

	return
}

func (this *Lags) printConsumersLagTable(zkcluster *zk.ZkCluster) {
	lines := make([]string, 0)
	header := "ConsumerGroup|Topic/Partition|Produced|Consumed|Lag|Committed|Uptime"
	lines = append(lines, header)

	// sort by group name
	consumersByGroup := zkcluster.ConsumersByGroup(this.groupPattern)
	sortedGroups := make([]string, 0, len(consumersByGroup))
	for group, _ := range consumersByGroup {
		sortedGroups = append(sortedGroups, group)
	}
	sort.Strings(sortedGroups)

	for _, group := range sortedGroups {
		if !patternMatched(group, this.groupPattern) {
			continue
		}

		sortedTopicAndPartitionIds := make([]string, 0, len(consumersByGroup[group]))
		consumers := make(map[string]zk.ConsumerMeta)
		for _, t := range consumersByGroup[group] {
			key := fmt.Sprintf("%s:%s", t.Topic, t.PartitionId)
			sortedTopicAndPartitionIds = append(sortedTopicAndPartitionIds, key)

			consumers[key] = t
		}
		sort.Strings(sortedTopicAndPartitionIds)

		for _, topicAndPartitionId := range sortedTopicAndPartitionIds {
			consumer := consumers[topicAndPartitionId]

			if !patternMatched(consumer.Topic, this.topicPattern) {
				continue
			}
			if !consumer.Online {
				continue
			}
			if this.problematicMode && consumer.Lag <= int64(this.lagThreshold) {
				continue
			}
			if consumer.ConsumerZnode == nil {
				this.Ui.Warn(fmt.Sprintf("%+v has no znode", consumer))
				continue
			}

			lines = append(lines,
				fmt.Sprintf("%s|%s/%s|%s|%s|%s|%s|%s",
					group,
					consumer.Topic, consumer.PartitionId,
					gofmt.Comma(consumer.ProducerOffset),
					gofmt.Comma(consumer.ConsumerOffset),
					gofmt.Comma(consumer.Lag),
					gofmt.PrettySince(consumer.Mtime.Time()),
					gofmt.PrettySince(consumer.ConsumerZnode.Uptime())))
		}
	}

	if len(lines) > 1 {
		this.Ui.Info(fmt.Sprintf("%s ▾", zkcluster.Name()))
		this.Ui.Output(columnize.SimpleFormat(lines))
	}
}

func (this *Lags) printConsumersLag(zkcluster *zk.ZkCluster) {
	// sort by group name
	consumersByGroup := zkcluster.ConsumersByGroup(this.groupPattern)
	sortedGroups := make([]string, 0, len(consumersByGroup))
	for group, _ := range consumersByGroup {
		sortedGroups = append(sortedGroups, group)
	}
	sort.Strings(sortedGroups)

	lines := make([]string, 0, 100)
	for _, group := range sortedGroups {
		sortedTopicAndPartitionIds := make([]string, 0)
		consumers := make(map[string]zk.ConsumerMeta)
		for _, t := range consumersByGroup[group] {
			key := fmt.Sprintf("%s:%s", t.Topic, t.PartitionId)
			sortedTopicAndPartitionIds = append(sortedTopicAndPartitionIds, key)

			consumers[key] = t
		}
		sort.Strings(sortedTopicAndPartitionIds)

		for _, topicAndPartitionId := range sortedTopicAndPartitionIds {
			consumer := consumers[topicAndPartitionId]

			if !patternMatched(consumer.Topic, this.topicPattern) {
				continue
			}

			var (
				lagOutput string
				symbol    string
			)
			if consumer.Lag > int64(this.lagThreshold) {
				lagOutput = color.Red("%15s", gofmt.Comma(consumer.Lag))
				if consumer.Online {
					symbol = color.Yellow("⚠")
				} else {
					symbol = color.Yellow("◎")
				}
			} else {
				lagOutput = color.Blue("%15s", gofmt.Comma(consumer.Lag))
				if consumer.Online {
					symbol = color.Green("◉")
				} else {
					symbol = color.Yellow("◎")
				}
			}

			if consumer.Online {
				if this.problematicMode && consumer.Lag <= int64(this.lagThreshold) {
					continue
				}

				var (
					host   string
					uptime string
				)
				if consumer.ConsumerZnode == nil {
					host = "unrecognized"
					uptime = "-"
				} else {
					host = color.Green("%90s", consumer.ConsumerZnode.Host())
					uptime = gofmt.PrettySince(consumer.ConsumerZnode.Uptime())
				}

				lines = append(lines, fmt.Sprintf("\t%s %35s/%-2s %12s -> %-12s %s %s\n%s %s",
					symbol,
					consumer.Topic, consumer.PartitionId,
					gofmt.Comma(consumer.ProducerOffset),
					gofmt.Comma(consumer.ConsumerOffset),
					lagOutput,
					gofmt.PrettySince(consumer.Mtime.Time()),
					host, uptime))
			} else if !this.onlineOnly {
				lines = append(lines, fmt.Sprintf("\t%s %35s/%-2s %12s -> %-12s %s %s",
					symbol,
					consumer.Topic, consumer.PartitionId,
					gofmt.Comma(consumer.ProducerOffset),
					gofmt.Comma(consumer.ConsumerOffset),
					lagOutput,
					gofmt.PrettySince(consumer.Mtime.Time())))
			}
		}

		if len(lines) > 0 {
			this.Ui.Output(strings.Repeat(" ", 4) + group)
			for _, l := range lines {
				this.Ui.Output(l)
			}
		}
	}
}

func (*Lags) Synopsis() string {
	return "Display online high level consumers lag on a topic"
}

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

    Display online high level consumers lag on a topic

Options:

    -z zone
      Default %s

    -c cluster

    -g group name pattern

    -t topic name pattern

    -w
      Watch mode. Keep printing out lags with 60s interval refresh.

    -l
      Only show online consumers lag.

    -lag threshold
      Default 5000.

    -p
      Only show problematic consumers.

    -table
      Display in table format.

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


<commit_before>package main

import (
	"bufio"
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"sort"
	"strings"
	"text/tabwriter"
	"time"

	"net/http"
	_ "net/http/pprof"

	"github.com/influxdb/influxdb/cmd/influx_tsm/b1"
	"github.com/influxdb/influxdb/cmd/influx_tsm/bz1"
	"github.com/influxdb/influxdb/cmd/influx_tsm/tsdb"
)

type ShardReader interface {
	KeyIterator
	Open() error
	Close() error
}

const (
	tsmExt = "tsm"
)

var description = `
Convert a database from b1 or bz1 format to tsm1 format.

This tool will backup the directories before conversion (if not disabled).
The backed-up files must be removed manually, generally after starting up the
node again to make sure all of data has been converted correctly.

To restore a backup:
  Shut down the node, remove the converted directory, and
  copy the backed-up directory to the original location.`

type options struct {
	DataPath       string
	BackupPath     string
	DBs            []string
	DebugAddr      string
	TSMSize        uint64
	Parallel       bool
	SkipBackup     bool
	UpdateInterval time.Duration
	// Quiet          bool
}

func (o *options) Parse() error {
	fs := flag.NewFlagSet(os.Args[0], flag.ExitOnError)

	var dbs string

	fs.StringVar(&dbs, "dbs", "", "Comma-delimited list of databases to convert. Default is to convert all databases.")
	fs.Uint64Var(&opts.TSMSize, "sz", maxTSMSz, "Maximum size of individual TSM files.")
	fs.BoolVar(&opts.Parallel, "parallel", false, "Perform parallel conversion. (up to GOMAXPROCS shards at once)")
	fs.BoolVar(&opts.SkipBackup, "nobackup", false, "Disable database backups. Not recommended.")
	fs.StringVar(&opts.BackupPath, "backup", "", "The location to backup up the current databases. Must not be within the data directoryi.")
	// fs.BoolVar(&opts.Quiet, "quiet", false, "Suppresses the regular status updates.")
	fs.StringVar(&opts.DebugAddr, "debug", "", "If set, http debugging endpoints will be enabled on the given address")
	fs.DurationVar(&opts.UpdateInterval, "interval", 5*time.Second, "How often status updates are printed.")
	fs.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %v [options] <data-path> \n", os.Args[0])
		fmt.Fprintf(os.Stderr, "%v\n\nOptions:\n", description)
		fs.PrintDefaults()
		fmt.Fprintf(os.Stderr, "\n")
	}

	if err := fs.Parse(os.Args[1:]); err != nil {
		return err
	}

	if len(fs.Args()) < 1 {
		return errors.New("no data directory specified")
	}
	var err error
	if o.DataPath, err = filepath.Abs(fs.Args()[0]); err != nil {
		return err
	}
	if o.DataPath, err = filepath.EvalSymlinks(filepath.Clean(o.DataPath)); err != nil {
		return err
	}

	if o.TSMSize > maxTSMSz {
		return fmt.Errorf("bad TSM file size, maximum TSM file size is %d", maxTSMSz)
	}

	// Check if specific databases were requested.
	o.DBs = strings.Split(dbs, ",")
	if len(o.DBs) == 1 && o.DBs[0] == "" {
		o.DBs = nil
	}

	if !o.SkipBackup {
		if o.BackupPath == "" {
			return errors.New("either -nobackup or -backup DIR must be set")
		}
		if o.BackupPath, err = filepath.Abs(o.BackupPath); err != nil {
			return err
		}
		if o.BackupPath, err = filepath.EvalSymlinks(filepath.Clean(o.BackupPath)); err != nil {
			if os.IsNotExist(err) {
				return errors.New("backup directory must already exist")
			}
			return err
		}

		if strings.HasPrefix(o.BackupPath, o.DataPath) {
			fmt.Println(o.BackupPath, o.DataPath)
			return errors.New("backup directory cannot be contained within data directory.")
		}
	}

	if o.DebugAddr != "" {
		log.Printf("Starting debugging server on http://%v", o.DebugAddr)
		go func() {
			log.Fatal(http.ListenAndServe(o.DebugAddr, nil))
		}()
	}

	return nil
}

var opts options

const maxTSMSz uint64 = 2 * 1024 * 1024 * 1024

func init() {
	log.SetOutput(os.Stderr)
	log.SetFlags(log.Ldate | log.Ltime | log.Lmicroseconds)
}

func main() {
	if err := opts.Parse(); err != nil {
		log.Fatal(err)
	}

	// Determine the list of databases
	dbs, err := ioutil.ReadDir(opts.DataPath)
	if err != nil {
		log.Fatalf("failed to access data directory at %v: %v\n", opts.DataPath, err)
	}
	fmt.Println() // Cleanly separate output from start of program.

	if opts.Parallel {
		if !isEnvSet("GOMAXPROCS") {
			// Only modify GOMAXPROCS if it wasn't set in the environment
			// This means 'GOMAXPROCS=1 influx_tsm -parallel' will not actually
			// run in parallel
			runtime.GOMAXPROCS(runtime.NumCPU())
		}
	}

	var badUser string
	if opts.SkipBackup {
		badUser = "(NOT RECOMMENDED)"
	}

	// Dump summary of what is about to happen.
	fmt.Println("b1 and bz1 shard conversion.")
	fmt.Println("-----------------------------------")
	fmt.Println("Data directory is:       ", opts.DataPath)
	fmt.Println("Backup directory is:     ", opts.BackupPath)
	fmt.Println("Databases specified:     ", allDBs(opts.DBs))
	fmt.Println("Database backups enabled:", yesno(!opts.SkipBackup), badUser)
	fmt.Println("Parallel mode enabled:   ", yesno(opts.Parallel), runtime.GOMAXPROCS(0))
	fmt.Println()

	shards := collectShards(dbs)

	// Anything to convert?
	fmt.Printf("\nFound %d shards that will be converted.\n", len(shards))
	if len(shards) == 0 {
		fmt.Println("Nothing to do.")
		return
	}

	// Display list of convertible shards.
	fmt.Println()
	w := new(tabwriter.Writer)
	w.Init(os.Stdout, 0, 8, 1, '\t', 0)
	fmt.Fprintln(w, "Database\tRetention\tPath\tEngine\tSize")
	for _, si := range shards {
		fmt.Fprintf(w, "%v\t%v\t%v\t%v\t%d\n", si.Database, si.RetentionPolicy, si.FullPath(opts.DataPath), si.FormatAsString(), si.Size)
	}
	w.Flush()

	// Get confirmation from user.
	fmt.Printf("\nThese shards will be converted. Proceed? y/N: ")
	liner := bufio.NewReader(os.Stdin)
	yn, err := liner.ReadString('\n')
	if err != nil {
		log.Fatalf("failed to read response: %v", err)
	}
	yn = strings.TrimRight(strings.ToLower(yn), "\n")
	if yn != "y" {
		log.Fatal("Conversion aborted.")
	}
	fmt.Println("Conversion starting....")

	tr := newTracker(shards, opts)

	if err := tr.Run(); err != nil {
		log.Fatalf("Error occurred preventing completion: %v\n", err)
	}

	tr.PrintStats()
}

func collectShards(dbs []os.FileInfo) tsdb.ShardInfos {
	// Get the list of shards for conversion.
	var shards tsdb.ShardInfos
	for _, db := range dbs {
		d := tsdb.NewDatabase(filepath.Join(opts.DataPath, db.Name()))
		shs, err := d.Shards()
		if err != nil {
			log.Fatalf("Failed to access shards for database %v: %v\n", d.Name(), err)
		}
		shards = append(shards, shs...)
	}

	sort.Sort(shards)
	shards = shards.FilterFormat(tsdb.TSM1)
	if len(dbs) > 0 {
		shards = shards.ExclusiveDatabases(opts.DBs)
	}

	return shards
}

// backupDatabase backs up the database named db
func backupDatabase(db string) error {
	copyFile := func(path string, info os.FileInfo, err error) error {
		// Strip the DataPath from the path and replace with BackupPath.
		toPath := strings.Replace(path, opts.DataPath, opts.BackupPath, 1)

		if info.IsDir() {
			return os.MkdirAll(toPath, info.Mode())
		}

		in, err := os.Open(path)
		if err != nil {
			return err
		}
		defer in.Close()

		srcInfo, err := os.Stat(path)
		if err != nil {
			return err
		}

		out, err := os.OpenFile(toPath, os.O_CREATE|os.O_WRONLY, info.Mode())
		if err != nil {
			return err
		}
		defer out.Close()

		dstInfo, err := os.Stat(toPath)
		if err != nil {
			return err
		}

		if dstInfo.Size() == srcInfo.Size() {
			log.Printf("Backup file already found for %v with correct size, skipping.", path)
			return nil
		}

		if dstInfo.Size() > srcInfo.Size() {
			log.Printf("Invalid backup file found for %v, replacing with good copy.", path)
			if err := out.Truncate(0); err != nil {
				return err
			}
			if _, err := out.Seek(0, os.SEEK_SET); err != nil {
				return err
			}
		}

		if dstInfo.Size() > 0 {
			log.Printf("Resuming backup of file %v, starting at %v bytes", path, dstInfo.Size())
		}

		off, err := out.Seek(0, os.SEEK_END)
		if err != nil {
			return err
		}
		if _, err := in.Seek(off, os.SEEK_SET); err != nil {
			return err
		}

		log.Printf("Backing up file %v", path)

		_, err = io.Copy(out, in)

		return err
	}

	return filepath.Walk(filepath.Join(opts.DataPath, db), copyFile)
}

// convertShard converts the shard in-place.
func convertShard(si *tsdb.ShardInfo, tr *tracker) error {
	src := si.FullPath(opts.DataPath)
	dst := fmt.Sprintf("%v.%v", src, tsmExt)

	var reader ShardReader
	switch si.Format {
	case tsdb.BZ1:
		reader = bz1.NewReader(src)
	case tsdb.B1:
		reader = b1.NewReader(src)
	default:
		return fmt.Errorf("Unsupported shard format: %v", si.FormatAsString())
	}
	defer reader.Close()

	// Open the shard, and create a converter.
	if err := reader.Open(); err != nil {
		return fmt.Errorf("Failed to open %v for conversion: %v", src, err)
	}
	converter := NewConverter(dst, uint32(opts.TSMSize), tr)

	// Perform the conversion.
	if err := converter.Process(reader); err != nil {
		return fmt.Errorf("Conversion of %v failed: %v", src, err)
	}

	// Delete source shard, and rename new tsm1 shard.
	if err := reader.Close(); err != nil {
		return fmt.Errorf("Conversion of %v failed due to close: %v", src, err)
	}

	if err := os.RemoveAll(si.FullPath(opts.DataPath)); err != nil {
		return fmt.Errorf("Deletion of %v failed: %v", src, err)
	}
	if err := os.Rename(dst, src); err != nil {
		return fmt.Errorf("Rename of %v to %v failed: %v", dst, src, err)
	}

	return nil
}

// ParallelGroup allows the maximum parrallelism of a set of operations to be controlled.
type ParallelGroup chan struct{}

// NewParallelGroup returns a group which allows n operations to run in parallel. A value of 0
// means no operations will ever run.
func NewParallelGroup(n int) ParallelGroup {
	return make(chan struct{}, n)
}

func (p ParallelGroup) Do(f func()) {
	p <- struct{}{} // acquire working slot
	defer func() { <-p }()

	f()
}

// yesno returns "yes" for true, "no" for false.
func yesno(b bool) string {
	if b {
		return "yes"
	}
	return "no"
}

// allDBs returns "all" if all databases are requested for conversion.
func allDBs(dbs []string) string {
	if dbs == nil {
		return "all"
	}
	return fmt.Sprintf("%v", dbs)
}

// isEnvSet checks to see if a variable was set in the environment
func isEnvSet(name string) bool {
	for _, s := range os.Environ() {
		if strings.SplitN(s, "=", 2)[0] == name {
			return true
		}
	}
	return false
}
<commit_msg>typo: remove trailing i.<commit_after>package main

import (
	"bufio"
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"runtime"
	"sort"
	"strings"
	"text/tabwriter"
	"time"

	"net/http"
	_ "net/http/pprof"

	"github.com/influxdb/influxdb/cmd/influx_tsm/b1"
	"github.com/influxdb/influxdb/cmd/influx_tsm/bz1"
	"github.com/influxdb/influxdb/cmd/influx_tsm/tsdb"
)

type ShardReader interface {
	KeyIterator
	Open() error
	Close() error
}

const (
	tsmExt = "tsm"
)

var description = `
Convert a database from b1 or bz1 format to tsm1 format.

This tool will backup the directories before conversion (if not disabled).
The backed-up files must be removed manually, generally after starting up the
node again to make sure all of data has been converted correctly.

To restore a backup:
  Shut down the node, remove the converted directory, and
  copy the backed-up directory to the original location.`

type options struct {
	DataPath       string
	BackupPath     string
	DBs            []string
	DebugAddr      string
	TSMSize        uint64
	Parallel       bool
	SkipBackup     bool
	UpdateInterval time.Duration
	// Quiet          bool
}

func (o *options) Parse() error {
	fs := flag.NewFlagSet(os.Args[0], flag.ExitOnError)

	var dbs string

	fs.StringVar(&dbs, "dbs", "", "Comma-delimited list of databases to convert. Default is to convert all databases.")
	fs.Uint64Var(&opts.TSMSize, "sz", maxTSMSz, "Maximum size of individual TSM files.")
	fs.BoolVar(&opts.Parallel, "parallel", false, "Perform parallel conversion. (up to GOMAXPROCS shards at once)")
	fs.BoolVar(&opts.SkipBackup, "nobackup", false, "Disable database backups. Not recommended.")
	fs.StringVar(&opts.BackupPath, "backup", "", "The location to backup up the current databases. Must not be within the data directory.")
	// fs.BoolVar(&opts.Quiet, "quiet", false, "Suppresses the regular status updates.")
	fs.StringVar(&opts.DebugAddr, "debug", "", "If set, http debugging endpoints will be enabled on the given address")
	fs.DurationVar(&opts.UpdateInterval, "interval", 5*time.Second, "How often status updates are printed.")
	fs.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %v [options] <data-path> \n", os.Args[0])
		fmt.Fprintf(os.Stderr, "%v\n\nOptions:\n", description)
		fs.PrintDefaults()
		fmt.Fprintf(os.Stderr, "\n")
	}

	if err := fs.Parse(os.Args[1:]); err != nil {
		return err
	}

	if len(fs.Args()) < 1 {
		return errors.New("no data directory specified")
	}
	var err error
	if o.DataPath, err = filepath.Abs(fs.Args()[0]); err != nil {
		return err
	}
	if o.DataPath, err = filepath.EvalSymlinks(filepath.Clean(o.DataPath)); err != nil {
		return err
	}

	if o.TSMSize > maxTSMSz {
		return fmt.Errorf("bad TSM file size, maximum TSM file size is %d", maxTSMSz)
	}

	// Check if specific databases were requested.
	o.DBs = strings.Split(dbs, ",")
	if len(o.DBs) == 1 && o.DBs[0] == "" {
		o.DBs = nil
	}

	if !o.SkipBackup {
		if o.BackupPath == "" {
			return errors.New("either -nobackup or -backup DIR must be set")
		}
		if o.BackupPath, err = filepath.Abs(o.BackupPath); err != nil {
			return err
		}
		if o.BackupPath, err = filepath.EvalSymlinks(filepath.Clean(o.BackupPath)); err != nil {
			if os.IsNotExist(err) {
				return errors.New("backup directory must already exist")
			}
			return err
		}

		if strings.HasPrefix(o.BackupPath, o.DataPath) {
			fmt.Println(o.BackupPath, o.DataPath)
			return errors.New("backup directory cannot be contained within data directory.")
		}
	}

	if o.DebugAddr != "" {
		log.Printf("Starting debugging server on http://%v", o.DebugAddr)
		go func() {
			log.Fatal(http.ListenAndServe(o.DebugAddr, nil))
		}()
	}

	return nil
}

var opts options

const maxTSMSz uint64 = 2 * 1024 * 1024 * 1024

func init() {
	log.SetOutput(os.Stderr)
	log.SetFlags(log.Ldate | log.Ltime | log.Lmicroseconds)
}

func main() {
	if err := opts.Parse(); err != nil {
		log.Fatal(err)
	}

	// Determine the list of databases
	dbs, err := ioutil.ReadDir(opts.DataPath)
	if err != nil {
		log.Fatalf("failed to access data directory at %v: %v\n", opts.DataPath, err)
	}
	fmt.Println() // Cleanly separate output from start of program.

	if opts.Parallel {
		if !isEnvSet("GOMAXPROCS") {
			// Only modify GOMAXPROCS if it wasn't set in the environment
			// This means 'GOMAXPROCS=1 influx_tsm -parallel' will not actually
			// run in parallel
			runtime.GOMAXPROCS(runtime.NumCPU())
		}
	}

	var badUser string
	if opts.SkipBackup {
		badUser = "(NOT RECOMMENDED)"
	}

	// Dump summary of what is about to happen.
	fmt.Println("b1 and bz1 shard conversion.")
	fmt.Println("-----------------------------------")
	fmt.Println("Data directory is:       ", opts.DataPath)
	fmt.Println("Backup directory is:     ", opts.BackupPath)
	fmt.Println("Databases specified:     ", allDBs(opts.DBs))
	fmt.Println("Database backups enabled:", yesno(!opts.SkipBackup), badUser)
	fmt.Println("Parallel mode enabled:   ", yesno(opts.Parallel), runtime.GOMAXPROCS(0))
	fmt.Println()

	shards := collectShards(dbs)

	// Anything to convert?
	fmt.Printf("\nFound %d shards that will be converted.\n", len(shards))
	if len(shards) == 0 {
		fmt.Println("Nothing to do.")
		return
	}

	// Display list of convertible shards.
	fmt.Println()
	w := new(tabwriter.Writer)
	w.Init(os.Stdout, 0, 8, 1, '\t', 0)
	fmt.Fprintln(w, "Database\tRetention\tPath\tEngine\tSize")
	for _, si := range shards {
		fmt.Fprintf(w, "%v\t%v\t%v\t%v\t%d\n", si.Database, si.RetentionPolicy, si.FullPath(opts.DataPath), si.FormatAsString(), si.Size)
	}
	w.Flush()

	// Get confirmation from user.
	fmt.Printf("\nThese shards will be converted. Proceed? y/N: ")
	liner := bufio.NewReader(os.Stdin)
	yn, err := liner.ReadString('\n')
	if err != nil {
		log.Fatalf("failed to read response: %v", err)
	}
	yn = strings.TrimRight(strings.ToLower(yn), "\n")
	if yn != "y" {
		log.Fatal("Conversion aborted.")
	}
	fmt.Println("Conversion starting....")

	tr := newTracker(shards, opts)

	if err := tr.Run(); err != nil {
		log.Fatalf("Error occurred preventing completion: %v\n", err)
	}

	tr.PrintStats()
}

func collectShards(dbs []os.FileInfo) tsdb.ShardInfos {
	// Get the list of shards for conversion.
	var shards tsdb.ShardInfos
	for _, db := range dbs {
		d := tsdb.NewDatabase(filepath.Join(opts.DataPath, db.Name()))
		shs, err := d.Shards()
		if err != nil {
			log.Fatalf("Failed to access shards for database %v: %v\n", d.Name(), err)
		}
		shards = append(shards, shs...)
	}

	sort.Sort(shards)
	shards = shards.FilterFormat(tsdb.TSM1)
	if len(dbs) > 0 {
		shards = shards.ExclusiveDatabases(opts.DBs)
	}

	return shards
}

// backupDatabase backs up the database named db
func backupDatabase(db string) error {
	copyFile := func(path string, info os.FileInfo, err error) error {
		// Strip the DataPath from the path and replace with BackupPath.
		toPath := strings.Replace(path, opts.DataPath, opts.BackupPath, 1)

		if info.IsDir() {
			return os.MkdirAll(toPath, info.Mode())
		}

		in, err := os.Open(path)
		if err != nil {
			return err
		}
		defer in.Close()

		srcInfo, err := os.Stat(path)
		if err != nil {
			return err
		}

		out, err := os.OpenFile(toPath, os.O_CREATE|os.O_WRONLY, info.Mode())
		if err != nil {
			return err
		}
		defer out.Close()

		dstInfo, err := os.Stat(toPath)
		if err != nil {
			return err
		}

		if dstInfo.Size() == srcInfo.Size() {
			log.Printf("Backup file already found for %v with correct size, skipping.", path)
			return nil
		}

		if dstInfo.Size() > srcInfo.Size() {
			log.Printf("Invalid backup file found for %v, replacing with good copy.", path)
			if err := out.Truncate(0); err != nil {
				return err
			}
			if _, err := out.Seek(0, os.SEEK_SET); err != nil {
				return err
			}
		}

		if dstInfo.Size() > 0 {
			log.Printf("Resuming backup of file %v, starting at %v bytes", path, dstInfo.Size())
		}

		off, err := out.Seek(0, os.SEEK_END)
		if err != nil {
			return err
		}
		if _, err := in.Seek(off, os.SEEK_SET); err != nil {
			return err
		}

		log.Printf("Backing up file %v", path)

		_, err = io.Copy(out, in)

		return err
	}

	return filepath.Walk(filepath.Join(opts.DataPath, db), copyFile)
}

// convertShard converts the shard in-place.
func convertShard(si *tsdb.ShardInfo, tr *tracker) error {
	src := si.FullPath(opts.DataPath)
	dst := fmt.Sprintf("%v.%v", src, tsmExt)

	var reader ShardReader
	switch si.Format {
	case tsdb.BZ1:
		reader = bz1.NewReader(src)
	case tsdb.B1:
		reader = b1.NewReader(src)
	default:
		return fmt.Errorf("Unsupported shard format: %v", si.FormatAsString())
	}
	defer reader.Close()

	// Open the shard, and create a converter.
	if err := reader.Open(); err != nil {
		return fmt.Errorf("Failed to open %v for conversion: %v", src, err)
	}
	converter := NewConverter(dst, uint32(opts.TSMSize), tr)

	// Perform the conversion.
	if err := converter.Process(reader); err != nil {
		return fmt.Errorf("Conversion of %v failed: %v", src, err)
	}

	// Delete source shard, and rename new tsm1 shard.
	if err := reader.Close(); err != nil {
		return fmt.Errorf("Conversion of %v failed due to close: %v", src, err)
	}

	if err := os.RemoveAll(si.FullPath(opts.DataPath)); err != nil {
		return fmt.Errorf("Deletion of %v failed: %v", src, err)
	}
	if err := os.Rename(dst, src); err != nil {
		return fmt.Errorf("Rename of %v to %v failed: %v", dst, src, err)
	}

	return nil
}

// ParallelGroup allows the maximum parrallelism of a set of operations to be controlled.
type ParallelGroup chan struct{}

// NewParallelGroup returns a group which allows n operations to run in parallel. A value of 0
// means no operations will ever run.
func NewParallelGroup(n int) ParallelGroup {
	return make(chan struct{}, n)
}

func (p ParallelGroup) Do(f func()) {
	p <- struct{}{} // acquire working slot
	defer func() { <-p }()

	f()
}

// yesno returns "yes" for true, "no" for false.
func yesno(b bool) string {
	if b {
		return "yes"
	}
	return "no"
}

// allDBs returns "all" if all databases are requested for conversion.
func allDBs(dbs []string) string {
	if dbs == nil {
		return "all"
	}
	return fmt.Sprintf("%v", dbs)
}

// isEnvSet checks to see if a variable was set in the environment
func isEnvSet(name string) bool {
	for _, s := range os.Environ() {
		if strings.SplitN(s, "=", 2)[0] == name {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>// Command interfaces
package main

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

	"github.com/rjeczalik/interfaces"
)

var (
	query  = flag.String("for", "", "Type to generate an interface for.")
	as     = flag.String("as", "main.Interface", `Generated interface name.`)
	output = flag.String("o", "-", "Output file.")
	all    = flag.Bool("all", false, "Include also unexported methods.")
)

var tmpl = template.Must(template.New("").Parse(`// Created by interfacer; DO NOT EDIT

package {{.PackageName}}

import (
{{range .Deps}}	"{{.}}"
{{end}})

// {{.InterfaceName}} is an interface generated for {{.Type}}.
type {{.InterfaceName}} interface {
{{range .Interface}}	{{.}}
{{end}}}
`))

type vars struct {
	PackageName   string
	InterfaceName string
	Type          string
	Deps          []string
	Interface     interfaces.Interface
}

func nonil(err ...error) error {
	for _, e := range err {
		if e != nil {
			return e
		}
	}
	return nil
}

func die(v interface{}) {
	fmt.Fprintln(os.Stderr, v)
	os.Exit(1)
}

func main() {
	flag.Parse()
	if *query == "" {
		die("empty -for flag value; see -help for details")
	}
	if *output == "" {
		die("empty -o flag value; see -help for details")
	}
	q, err := interfaces.ParseQuery(*query)
	if err != nil {
		die(err)
	}
	opts := &interfaces.Options{
		Query:      q,
		Unexported: *all,
	}
	i, err := interfaces.NewWithOptions(opts)
	if err != nil {
		die(err)
	}
	v := &vars{
		Type:      fmt.Sprintf(`"%s"`, *query),
		Deps:      i.Deps(),
		Interface: i,
	}
	if i := strings.IndexRune(*as, '.'); i != -1 {
		v.PackageName = (*as)[:i]
		v.InterfaceName = (*as)[i+1:]
	} else {
		v.InterfaceName = *as
	}
	f := os.Stdout
	if *output != "-" {
		f, err = os.OpenFile(*output, os.O_TRUNC|os.O_CREATE|os.O_WRONLY, 0755)
		if err != nil {
			die(err)
		}
	}
	err = nonil(tmpl.Execute(f, v), f.Close())
	if err != nil {
		die(err)
	}
}
<commit_msg>fix code comment<commit_after>// Command interfaces
package main

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

	"github.com/rjeczalik/interfaces"
)

var (
	query  = flag.String("for", "", "Type to generate an interface for.")
	as     = flag.String("as", "main.Interface", `Generated interface name.`)
	output = flag.String("o", "-", "Output file.")
	all    = flag.Bool("all", false, "Include also unexported methods.")
)

var tmpl = template.Must(template.New("").Parse(`// Code generated by interfacer; DO NOT EDIT

package {{.PackageName}}

import (
{{range .Deps}}	"{{.}}"
{{end}})

// {{.InterfaceName}} is an interface generated for {{.Type}}.
type {{.InterfaceName}} interface {
{{range .Interface}}	{{.}}
{{end}}}
`))

type vars struct {
	PackageName   string
	InterfaceName string
	Type          string
	Deps          []string
	Interface     interfaces.Interface
}

func nonil(err ...error) error {
	for _, e := range err {
		if e != nil {
			return e
		}
	}
	return nil
}

func die(v interface{}) {
	fmt.Fprintln(os.Stderr, v)
	os.Exit(1)
}

func main() {
	flag.Parse()
	if *query == "" {
		die("empty -for flag value; see -help for details")
	}
	if *output == "" {
		die("empty -o flag value; see -help for details")
	}
	q, err := interfaces.ParseQuery(*query)
	if err != nil {
		die(err)
	}
	opts := &interfaces.Options{
		Query:      q,
		Unexported: *all,
	}
	i, err := interfaces.NewWithOptions(opts)
	if err != nil {
		die(err)
	}
	v := &vars{
		Type:      fmt.Sprintf(`"%s"`, *query),
		Deps:      i.Deps(),
		Interface: i,
	}
	if i := strings.IndexRune(*as, '.'); i != -1 {
		v.PackageName = (*as)[:i]
		v.InterfaceName = (*as)[i+1:]
	} else {
		v.InterfaceName = *as
	}
	f := os.Stdout
	if *output != "-" {
		f, err = os.OpenFile(*output, os.O_TRUNC|os.O_CREATE|os.O_WRONLY, 0755)
		if err != nil {
			die(err)
		}
	}
	err = nonil(tmpl.Execute(f, v), f.Close())
	if err != nil {
		die(err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"io"
	"launchpad.net/gnuflag"
	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/log"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/version"
)

// SyncToolsCommand copies all the tools from the us-east-1 bucket to the local
// bucket.
type SyncToolsCommand struct {
	EnvCommandBase
	allVersions bool
	dryRun      bool
}

var _ cmd.Command = (*SyncToolsCommand)(nil)

func (c *SyncToolsCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "sync-tools",
		Purpose: "copy tools from the official bucket into a local environment",
		Doc: `
This copies the Juju tools tarball from the official bucket into
your environment. This is generally done when you want Juju to be able
to run without having to access Amazon. Sometimes this is because the
environment does not have public access, and sometimes you just want
to avoid having to access data outside of the local cloud.
`,
	}
}

func (c *SyncToolsCommand) SetFlags(f *gnuflag.FlagSet) {
	c.EnvCommandBase.SetFlags(f)
	f.BoolVar(&c.allVersions, "all", false, "copy all versions, not just the latest")
	f.BoolVar(&c.dryRun, "dry-run", false, "don't copy, just print what would be copied")
	// BUG(lp:1163164)  jam 2013-04-2 we would like to add a "source"
	// location, rather than only copying from us-east-1
}

func (c *SyncToolsCommand) Init(args []string) error {
	return cmd.CheckEmpty(args)
}

var officialBucketAttrs = map[string]interface{}{
	"name":            "juju-public",
	"type":            "ec2",
	"control-bucket":  "juju-dist",
	"access-key":      "",
	"secret-key":      "",
	"authorized-keys": "not-really", // We shouldn't need ssh access
}

// Find the set of tools at the 'latest' version
func findNewest(fullTools []*state.Tools) []*state.Tools {
	var curBest *state.Tools = nil
	var res []*state.Tools = nil
	for _, tool := range fullTools {
		var add = false
		if curBest == nil || curBest.Number.Less(tool.Number) {
			// This best is clearly better than all existing
			// entries, so reset the list
			res = make([]*state.Tools, 0, 1)
			add = true
			curBest = tool
		}
		if curBest.Number == tool.Number {
			add = true
		}
		if add {
			res = append(res, tool)
		}
	}
	return res
}

// Find tools that aren't present in target
func findMissing(sourceTools, targetTools []*state.Tools) []*state.Tools {
	var target = make(map[version.Binary]bool, len(targetTools))
	for _, tool := range targetTools {
		target[tool.Binary] = true
	}
	res := make([]*state.Tools, 0)
	for _, tool := range sourceTools {
		if present := target[tool.Binary]; !present {
			res = append(res, tool)
		}
	}
	return res
}

func copyOne(
	tool *state.Tools, source environs.StorageReader,
	target environs.Storage, ctx *cmd.Context,
) error {
	toolsPath := environs.ToolsStoragePath(tool.Binary)
	fmt.Fprintf(ctx.Stdout, "copying %v", toolsPath)
	srcFile, err := source.Get(toolsPath)
	if err != nil {
		return err
	}
	defer srcFile.Close()
	// We have to buffer the content, because Put requires the content
	// length, but Get only returns us a ReadCloser
	buf := bytes.NewBuffer(nil)
	nBytes, err := io.Copy(buf, srcFile)
	if err != nil {
		return err
	}
	log.Infof("cmd/juju: downloaded %v (%dkB), uploading", toolsPath, (nBytes+512)/1024)
	fmt.Fprintf(ctx.Stdout, ", download %dkB, uploading\n", (nBytes+512)/1024)

	if err := target.Put(toolsPath, buf, nBytes); err != nil {
		return err
	}
	return nil
}

func copyTools(
	tools []*state.Tools, source environs.StorageReader,
	target environs.Storage, dryRun bool, ctx *cmd.Context,
) error {
	for _, tool := range tools {
		log.Infof("cmd/juju: copying %s from %s", tool.Binary, tool.URL)
		if dryRun {
			continue
		}
		if err := copyOne(tool, source, target, ctx); err != nil {
			return err
		}
	}
	return nil
}

func (c *SyncToolsCommand) Run(ctx *cmd.Context) error {
	officialEnviron, err := environs.NewFromAttrs(officialBucketAttrs)
	if err != nil {
		log.Errorf("cmd/juju: failed to initialize the official bucket environment")
		return err
	}
	fmt.Fprintf(ctx.Stdout, "listing the source bucket\n")
	sourceToolsList, err := environs.ListTools(officialEnviron, version.Current.Major)
	if err != nil {
		return err
	}
	targetEnv, err := environs.NewFromName(c.EnvName)
	if err != nil {
		log.Errorf("cmd/juju: unable to read %q from environment", c.EnvName)
		return err
	}
	toolsToCopy := sourceToolsList.Public
	if !c.allVersions {
		toolsToCopy = findNewest(toolsToCopy)
	}
	fmt.Fprintf(ctx.Stdout, "found %d tools in source (%d recent ones)\n",
		len(sourceToolsList.Public), len(toolsToCopy))
	for _, tool := range toolsToCopy {
		log.Debugf("cmd/juju: found source tool: %s", tool)
	}
	fmt.Fprintf(ctx.Stdout, "listing target bucket\n")
	targetToolsList, err := environs.ListTools(targetEnv, version.Current.Major)
	if err != nil {
		return err
	}
	for _, tool := range targetToolsList.Private {
		log.Debugf("cmd/juju: found target tool: %s", tool)
	}
	missing := findMissing(toolsToCopy, targetToolsList.Private)
	fmt.Fprintf(ctx.Stdout, "found %d tools in target; %d tools to be copied\n",
		len(targetToolsList.Private), len(missing))
	err = copyTools(missing, officialEnviron.PublicStorage(), targetEnv.Storage(), c.dryRun, ctx)
	if err != nil {
		return err
	}
	fmt.Fprintf(ctx.Stdout, "copied %d tools\n", len(missing))
	return nil
}
<commit_msg>Some cleanups suggested by Tim Penhey<commit_after>package main

import (
	"bytes"
	"fmt"
	"io"
	"launchpad.net/gnuflag"
	"launchpad.net/juju-core/cmd"
	"launchpad.net/juju-core/environs"
	"launchpad.net/juju-core/log"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/version"
)

// SyncToolsCommand copies all the tools from the us-east-1 bucket to the local
// bucket.
type SyncToolsCommand struct {
	EnvCommandBase
	allVersions bool
	dryRun      bool
}

var _ cmd.Command = (*SyncToolsCommand)(nil)

func (c *SyncToolsCommand) Info() *cmd.Info {
	return &cmd.Info{
		Name:    "sync-tools",
		Purpose: "copy tools from the official bucket into a local environment",
		Doc: `
This copies the Juju tools tarball from the official bucket into
your environment. This is generally done when you want Juju to be able
to run without having to access Amazon. Sometimes this is because the
environment does not have public access, and sometimes you just want
to avoid having to access data outside of the local cloud.
`,
	}
}

func (c *SyncToolsCommand) SetFlags(f *gnuflag.FlagSet) {
	c.EnvCommandBase.SetFlags(f)
	f.BoolVar(&c.allVersions, "all", false, "copy all versions, not just the latest")
	f.BoolVar(&c.dryRun, "dry-run", false, "don't copy, just print what would be copied")
	// BUG(lp:1163164)  jam 2013-04-2 we would like to add a "source"
	// location, rather than only copying from us-east-1
}

func (c *SyncToolsCommand) Init(args []string) error {
	return cmd.CheckEmpty(args)
}

var officialBucketAttrs = map[string]interface{}{
	"name":            "juju-public",
	"type":            "ec2",
	"control-bucket":  "juju-dist",
	"access-key":      "",
	"secret-key":      "",
	"authorized-keys": "not-really", // We shouldn't need ssh access
}

// Find the set of tools at the 'latest' version
func findNewest(fullTools []*state.Tools) []*state.Tools {
	// This assumes the nil version of Number is always less than a real
	// number, but we don't have negative versions, so this should be fine
	var curBest version.Number
	var res []*state.Tools = nil
	for _, tool := range fullTools {
		if curBest.Less(tool.Number) {
			// This tool is newer than our current best,
			// so reset the list
			res = []*state.Tools{tool}
			curBest = tool.Number
		} else if curBest == tool.Number {
			res = append(res, tool)
		}
	}
	return res
}

// Find tools that aren't present in target
func findMissing(sourceTools, targetTools []*state.Tools) []*state.Tools {
	var target = make(map[version.Binary]bool, len(targetTools))
	for _, tool := range targetTools {
		target[tool.Binary] = true
	}
	res := make([]*state.Tools, 0)
	for _, tool := range sourceTools {
		if present := target[tool.Binary]; !present {
			res = append(res, tool)
		}
	}
	return res
}

func copyOne(
	tool *state.Tools, source environs.StorageReader,
	target environs.Storage, ctx *cmd.Context,
) error {
	toolsPath := environs.ToolsStoragePath(tool.Binary)
	fmt.Fprintf(ctx.Stdout, "copying %v", toolsPath)
	srcFile, err := source.Get(toolsPath)
	if err != nil {
		return err
	}
	defer srcFile.Close()
	// We have to buffer the content, because Put requires the content
	// length, but Get only returns us a ReadCloser
	buf := bytes.NewBuffer(nil)
	nBytes, err := io.Copy(buf, srcFile)
	if err != nil {
		return err
	}
	log.Infof("cmd/juju: downloaded %v (%dkB), uploading", toolsPath, (nBytes+512)/1024)
	fmt.Fprintf(ctx.Stdout, ", download %dkB, uploading\n", (nBytes+512)/1024)

	if err := target.Put(toolsPath, buf, nBytes); err != nil {
		return err
	}
	return nil
}

func copyTools(
	tools []*state.Tools, source environs.StorageReader,
	target environs.Storage, dryRun bool, ctx *cmd.Context,
) error {
	for _, tool := range tools {
		log.Infof("cmd/juju: copying %s from %s", tool.Binary, tool.URL)
		if dryRun {
			continue
		}
		if err := copyOne(tool, source, target, ctx); err != nil {
			return err
		}
	}
	return nil
}

func (c *SyncToolsCommand) Run(ctx *cmd.Context) error {
	officialEnviron, err := environs.NewFromAttrs(officialBucketAttrs)
	if err != nil {
		log.Errorf("cmd/juju: failed to initialize the official bucket environment")
		return err
	}
	fmt.Fprintf(ctx.Stdout, "listing the source bucket\n")
	sourceToolsList, err := environs.ListTools(officialEnviron, version.Current.Major)
	if err != nil {
		return err
	}
	targetEnv, err := environs.NewFromName(c.EnvName)
	if err != nil {
		log.Errorf("cmd/juju: unable to read %q from environment", c.EnvName)
		return err
	}
	toolsToCopy := sourceToolsList.Public
	if !c.allVersions {
		toolsToCopy = findNewest(toolsToCopy)
	}
	fmt.Fprintf(ctx.Stdout, "found %d tools in source (%d recent ones)\n",
		len(sourceToolsList.Public), len(toolsToCopy))
	for _, tool := range toolsToCopy {
		log.Debugf("cmd/juju: found source tool: %s", tool)
	}
	fmt.Fprintf(ctx.Stdout, "listing target bucket\n")
	targetToolsList, err := environs.ListTools(targetEnv, version.Current.Major)
	if err != nil {
		return err
	}
	for _, tool := range targetToolsList.Private {
		log.Debugf("cmd/juju: found target tool: %s", tool)
	}
	missing := findMissing(toolsToCopy, targetToolsList.Private)
	fmt.Fprintf(ctx.Stdout, "found %d tools in target; %d tools to be copied\n",
		len(targetToolsList.Private), len(missing))
	err = copyTools(missing, officialEnviron.PublicStorage(), targetEnv.Storage(), c.dryRun, ctx)
	if err != nil {
		return err
	}
	fmt.Fprintf(ctx.Stdout, "copied %d tools\n", len(missing))
	return nil
}
<|endoftext|>


<commit_before>//go:generate go-bindata -prefix ../../migrations/ -pkg migrations -o ../../migrations/migrations.go ../../migrations/
//go:generate go-bindata -prefix ../../static/ -pkg static -o ../../static/static.go ../../static/...

package main

import (
	"net/http"
	"os"
	"os/signal"
	"syscall"

	log "github.com/Sirupsen/logrus"
	"github.com/brocaar/loraserver"
	application "github.com/brocaar/loraserver/application/mqttpubsub"
	gateway "github.com/brocaar/loraserver/gateway/mqttpubsub"
	"github.com/brocaar/loraserver/migrations"
	"github.com/brocaar/loraserver/static"
	"github.com/brocaar/lorawan"
	"github.com/codegangsta/cli"
	"github.com/elazarl/go-bindata-assetfs"
	_ "github.com/lib/pq"
	"github.com/rubenv/sql-migrate"
)

var version string // set by the compiler

func run(c *cli.Context) {
	// parse the NetID
	var netID lorawan.NetID
	if err := netID.UnmarshalText([]byte(c.String("net-id"))); err != nil {
		log.Fatalf("could not parse NetID: %s", err)
	}

	// connect to the database
	log.Info("connecting to postgresql")
	db, err := loraserver.OpenDatabase(c.String("postgres-dsn"))
	if err != nil {
		log.Fatalf("could not connect to the database: %s", err)
	}

	// setup redis pool
	log.Info("setup redis connection pool")
	rp := loraserver.NewRedisPool(c.String("redis-url"))

	// setup gateway backend
	gw, err := gateway.NewBackend(c.String("gw-mqtt-server"), c.String("gw-mqtt-username"), c.String("gw-mqtt-password"))
	if err != nil {
		log.Fatalf("could not setup gateway backend: %s", err)
	}

	// setup application backend
	app, err := application.NewBackend(c.String("app-mqtt-server"), c.String("app-mqtt-username"), c.String("app-mqtt-password"))
	if err != nil {
		log.Fatalf("could not setup application backend: %s", err)
	}

	// auto-migrate the database
	if c.Bool("db-automigrate") {
		log.Info("applying database migrations")
		m := &migrate.AssetMigrationSource{
			Asset:    migrations.Asset,
			AssetDir: migrations.AssetDir,
			Dir:      "",
		}
		n, err := migrate.Exec(db.DB, "postgres", m, migrate.Up)
		if err != nil {
			log.Fatalf("migrations failed: %s", err)
		}
		log.WithField("count", n).Info("migrations applied")
	}

	ctx := loraserver.Context{
		DB:          db,
		RedisPool:   rp,
		Gateway:     gw,
		Application: app,
		NetID:       netID,
	}

	// start the loraserver
	server := loraserver.NewServer(ctx)
	if err := server.Start(); err != nil {
		log.Fatal(err)
	}

	// setup json-rpc api handler
	apiHandler, err := loraserver.NewJSONRPCHandler(
		loraserver.NewApplicationAPI(ctx),
		loraserver.NewNodeAPI(ctx),
		loraserver.NewNodeSessionAPI(ctx),
	)
	if err != nil {
		log.Fatal(err)
	}
	log.WithField("path", "/rpc").Info("registering json-rpc handler")
	http.Handle("/rpc", apiHandler)

	// setup static file server (for the gui)
	log.WithField("path", "/").Info("registering gui handler")
	http.Handle("/", http.FileServer(&assetfs.AssetFS{
		Asset:     static.Asset,
		AssetDir:  static.AssetDir,
		AssetInfo: static.AssetInfo,
		Prefix:    "",
	}))

	// start the http server
	go func() {
		log.WithField("bind", c.String("http-bind")).Info("starting http server")
		log.Fatal(http.ListenAndServe(c.String("http-bind"), nil))
	}()

	sigChan := make(chan os.Signal)
	signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
	log.WithField("signal", <-sigChan).Info("signal received")
	log.Warning("loraserver is shutting down")
	if err := server.Stop(); err != nil {
		log.Fatal(err)
	}
}

func main() {
	app := cli.NewApp()
	app.Name = "loraserver"
	app.Usage = "network-server for LoRaWAN networks"
	app.Version = version
	app.Copyright = "See http://github.com/brocaar/loraserver for copyright information"
	app.Action = run
	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:   "db-automigrate",
			Usage:  "automatically apply database migrations",
			EnvVar: "DB_AUTOMIGRATE",
		},
		cli.StringFlag{
			Name:   "net-id",
			Usage:  "network identifier (NetID, 3 bytes) encoded as HEX (e.g. 010203)",
			EnvVar: "NET_ID",
		},
		cli.StringFlag{
			Name:   "postgres-dsn",
			Usage:  "postgresql dsn (e.g.: postgres://user:password@hostname/database?sslmode=disable)",
			Value:  "postgres://localhost/loraserver?sslmode=disable",
			EnvVar: "POSTGRES_DSN",
		},
		cli.StringFlag{
			Name:   "redis-url",
			Usage:  "redis url",
			Value:  "redis://localhost:6379",
			EnvVar: "REDIS_URL",
		},
		cli.StringFlag{
			Name:   "http-bind",
			Usage:  "ip:port to bind the http api server to",
			Value:  "0.0.0.0:8000",
			EnvVar: "HTTP_BIND",
		},
		cli.StringFlag{
			Name:   "gw-mqtt-server",
			Usage:  "Gateway-backend MQTT server",
			Value:  "tcp://localhost:1883",
			EnvVar: "GW_MQTT_SERVER",
		},
		cli.StringFlag{
			Name:   "gw-mqtt-username",
			Usage:  "Gateway-backend MQTT username",
			EnvVar: "GW_MQTT_USERNAME",
		},
		cli.StringFlag{
			Name:   "gw-mqtt-password",
			Usage:  "Gateway-backend MQTT password",
			EnvVar: "GW_MQTT_PASSWORD",
		},
		cli.StringFlag{
			Name:   "app-mqtt-server",
			Usage:  "Application-backend MQTT server",
			Value:  "tcp://localhost:1883",
			EnvVar: "APP_MQTT_SERVER",
		},
		cli.StringFlag{
			Name:   "app-mqtt-username",
			Usage:  "Application-backend MQTT username",
			EnvVar: "APP_MQTT_USERNAME",
		},
		cli.StringFlag{
			Name:   "app-mqtt-password",
			Usage:  "Application-backend MQTT password",
			EnvVar: "APP_MQTT_PASSWORD",
		},
	}
	app.Run(os.Args)
}
<commit_msg>Force stop on second received signal.<commit_after>//go:generate go-bindata -prefix ../../migrations/ -pkg migrations -o ../../migrations/migrations.go ../../migrations/
//go:generate go-bindata -prefix ../../static/ -pkg static -o ../../static/static.go ../../static/...

package main

import (
	"net/http"
	"os"
	"os/signal"
	"syscall"

	log "github.com/Sirupsen/logrus"
	"github.com/brocaar/loraserver"
	application "github.com/brocaar/loraserver/application/mqttpubsub"
	gateway "github.com/brocaar/loraserver/gateway/mqttpubsub"
	"github.com/brocaar/loraserver/migrations"
	"github.com/brocaar/loraserver/static"
	"github.com/brocaar/lorawan"
	"github.com/codegangsta/cli"
	"github.com/elazarl/go-bindata-assetfs"
	_ "github.com/lib/pq"
	"github.com/rubenv/sql-migrate"
)

var version string // set by the compiler

func run(c *cli.Context) {
	// parse the NetID
	var netID lorawan.NetID
	if err := netID.UnmarshalText([]byte(c.String("net-id"))); err != nil {
		log.Fatalf("could not parse NetID: %s", err)
	}

	// connect to the database
	log.Info("connecting to postgresql")
	db, err := loraserver.OpenDatabase(c.String("postgres-dsn"))
	if err != nil {
		log.Fatalf("could not connect to the database: %s", err)
	}

	// setup redis pool
	log.Info("setup redis connection pool")
	rp := loraserver.NewRedisPool(c.String("redis-url"))

	// setup gateway backend
	gw, err := gateway.NewBackend(c.String("gw-mqtt-server"), c.String("gw-mqtt-username"), c.String("gw-mqtt-password"))
	if err != nil {
		log.Fatalf("could not setup gateway backend: %s", err)
	}

	// setup application backend
	app, err := application.NewBackend(c.String("app-mqtt-server"), c.String("app-mqtt-username"), c.String("app-mqtt-password"))
	if err != nil {
		log.Fatalf("could not setup application backend: %s", err)
	}

	// auto-migrate the database
	if c.Bool("db-automigrate") {
		log.Info("applying database migrations")
		m := &migrate.AssetMigrationSource{
			Asset:    migrations.Asset,
			AssetDir: migrations.AssetDir,
			Dir:      "",
		}
		n, err := migrate.Exec(db.DB, "postgres", m, migrate.Up)
		if err != nil {
			log.Fatalf("migrations failed: %s", err)
		}
		log.WithField("count", n).Info("migrations applied")
	}

	ctx := loraserver.Context{
		DB:          db,
		RedisPool:   rp,
		Gateway:     gw,
		Application: app,
		NetID:       netID,
	}

	// start the loraserver
	server := loraserver.NewServer(ctx)
	if err := server.Start(); err != nil {
		log.Fatal(err)
	}

	// setup json-rpc api handler
	apiHandler, err := loraserver.NewJSONRPCHandler(
		loraserver.NewApplicationAPI(ctx),
		loraserver.NewNodeAPI(ctx),
		loraserver.NewNodeSessionAPI(ctx),
	)
	if err != nil {
		log.Fatal(err)
	}
	log.WithField("path", "/rpc").Info("registering json-rpc handler")
	http.Handle("/rpc", apiHandler)

	// setup static file server (for the gui)
	log.WithField("path", "/").Info("registering gui handler")
	http.Handle("/", http.FileServer(&assetfs.AssetFS{
		Asset:     static.Asset,
		AssetDir:  static.AssetDir,
		AssetInfo: static.AssetInfo,
		Prefix:    "",
	}))

	// start the http server
	go func() {
		log.WithField("bind", c.String("http-bind")).Info("starting http server")
		log.Fatal(http.ListenAndServe(c.String("http-bind"), nil))
	}()

	sigChan := make(chan os.Signal)
	exitChan := make(chan struct{})
	signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
	log.WithField("signal", <-sigChan).Info("signal received")
	go func() {
		log.Warning("stopping loraserver")
		if err := server.Stop(); err != nil {
			log.Fatal(err)
		}
		exitChan <- struct{}{}
	}()
	select {
	case <-exitChan:
	case s := <-sigChan:
		log.WithField("signal", s).Info("signal received, stopping immediately")
	}
}

func main() {
	app := cli.NewApp()
	app.Name = "loraserver"
	app.Usage = "network-server for LoRaWAN networks"
	app.Version = version
	app.Copyright = "See http://github.com/brocaar/loraserver for copyright information"
	app.Action = run
	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:   "db-automigrate",
			Usage:  "automatically apply database migrations",
			EnvVar: "DB_AUTOMIGRATE",
		},
		cli.StringFlag{
			Name:   "net-id",
			Usage:  "network identifier (NetID, 3 bytes) encoded as HEX (e.g. 010203)",
			EnvVar: "NET_ID",
		},
		cli.StringFlag{
			Name:   "postgres-dsn",
			Usage:  "postgresql dsn (e.g.: postgres://user:password@hostname/database?sslmode=disable)",
			Value:  "postgres://localhost/loraserver?sslmode=disable",
			EnvVar: "POSTGRES_DSN",
		},
		cli.StringFlag{
			Name:   "redis-url",
			Usage:  "redis url",
			Value:  "redis://localhost:6379",
			EnvVar: "REDIS_URL",
		},
		cli.StringFlag{
			Name:   "http-bind",
			Usage:  "ip:port to bind the http api server to",
			Value:  "0.0.0.0:8000",
			EnvVar: "HTTP_BIND",
		},
		cli.StringFlag{
			Name:   "gw-mqtt-server",
			Usage:  "Gateway-backend MQTT server",
			Value:  "tcp://localhost:1883",
			EnvVar: "GW_MQTT_SERVER",
		},
		cli.StringFlag{
			Name:   "gw-mqtt-username",
			Usage:  "Gateway-backend MQTT username",
			EnvVar: "GW_MQTT_USERNAME",
		},
		cli.StringFlag{
			Name:   "gw-mqtt-password",
			Usage:  "Gateway-backend MQTT password",
			EnvVar: "GW_MQTT_PASSWORD",
		},
		cli.StringFlag{
			Name:   "app-mqtt-server",
			Usage:  "Application-backend MQTT server",
			Value:  "tcp://localhost:1883",
			EnvVar: "APP_MQTT_SERVER",
		},
		cli.StringFlag{
			Name:   "app-mqtt-username",
			Usage:  "Application-backend MQTT username",
			EnvVar: "APP_MQTT_USERNAME",
		},
		cli.StringFlag{
			Name:   "app-mqtt-password",
			Usage:  "Application-backend MQTT password",
			EnvVar: "APP_MQTT_PASSWORD",
		},
	}
	app.Run(os.Args)
}
<|endoftext|>


<commit_before>package network

import (
	"errors"
	"fmt"
	"net"
	"os"
	"os/exec"

	log "github.com/Sirupsen/logrus"

	"github.com/rancherio/os/config"
	"github.com/rancherio/os/docker"
	"github.com/ryanuber/go-glob"
	"github.com/vishvananda/netlink"
)

func Main() {
	args := os.Args
	if len(args) > 1 {
		fmt.Println("call " + args[0] + " to load network config from rancher.yml config file")
		return
	}
	cfg, err := config.LoadConfig()
	if err != nil {
		log.Fatal(err)
	}
	applyNetworkConfigs(cfg)
}

func applyNetworkConfigs(cfg *config.Config) error {
	links, err := netlink.LinkList()
	if err != nil {
		return err
	}

	//apply network config
	for _, link := range links {
		linkName := link.Attrs().Name
		var match config.InterfaceConfig

		for key, netConf := range cfg.Network.Interfaces {
			if netConf.Match == "" {
				netConf.Match = key
			}

			if netConf.Match == "" {
				continue
			}

			// "" means match has not been found
			if match.Match == "" && matches(linkName, netConf.Match) {
				match = netConf
			}

			if netConf.Match == linkName {
				// Found exact match, use it over wildcard match
				match = netConf
			}
		}

		if match.Match != "" {
			err = applyNetConf(link, match)
			if err != nil {
				log.Errorf("Failed to apply settings to %s : %v", linkName, err)
			}
		}
	}

	if err != nil {
		return err
	}

	//post run
	if cfg.Network.PostRun != nil {
		return docker.StartAndWait(config.DOCKER_HOST, cfg.Network.PostRun)
	}
	return nil
}

func applyNetConf(link netlink.Link, netConf config.InterfaceConfig) error {
	if netConf.DHCP {
		log.Infof("Running DHCP on %s", link.Attrs().Name)
		cmd := exec.Command("udhcpc", "-i", link.Attrs().Name, "-t", "20", "-n")
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		if err := cmd.Run(); err != nil {
			log.Error(err)
		}
	} else if netConf.Address == "" {
		return nil
	} else {
		addr, err := netlink.ParseAddr(netConf.Address)
		if err != nil {
			return err
		}
		if err := netlink.AddrAdd(link, addr); err != nil {
			log.Error("addr add failed")
			return err
		}
		log.Infof("Set %s on %s", netConf.Address, link.Attrs().Name)
	}

	if netConf.MTU > 0 {
		if err := netlink.LinkSetMTU(link, netConf.MTU); err != nil {
			log.Error("set MTU Failed")
			return err
		}
	}

	if err := netlink.LinkSetUp(link); err != nil {
		log.Error("failed to setup link")
		return err
	}

	if netConf.Gateway != "" {
		gatewayIp := net.ParseIP(netConf.Gateway)
		if gatewayIp == nil {
			return errors.New("Invalid gateway address " + netConf.Gateway)
		}

		route := netlink.Route{
			Scope: netlink.SCOPE_UNIVERSE,
			Gw:    net.ParseIP(netConf.Gateway),
		}
		if err := netlink.RouteAdd(&route); err != nil {
			log.Error("gateway set failed")
			return err
		}

		log.Infof("Set default gateway %s", netConf.Gateway)
	}

	return nil
}

func matches(link, conf string) bool {
	return glob.Glob(conf, link)
}
<commit_msg>Make the post hook use the system docker<commit_after>package network

import (
	"errors"
	"fmt"
	"net"
	"os"
	"os/exec"

	log "github.com/Sirupsen/logrus"

	"github.com/rancherio/os/config"
	"github.com/rancherio/os/docker"
	"github.com/ryanuber/go-glob"
	"github.com/vishvananda/netlink"
)

func Main() {
	args := os.Args
	if len(args) > 1 {
		fmt.Println("call " + args[0] + " to load network config from rancher.yml config file")
		return
	}
	cfg, err := config.LoadConfig()
	if err != nil {
		log.Fatal(err)
	}
	applyNetworkConfigs(cfg)
}

func applyNetworkConfigs(cfg *config.Config) error {
	links, err := netlink.LinkList()
	if err != nil {
		return err
	}

	//apply network config
	for _, link := range links {
		linkName := link.Attrs().Name
		var match config.InterfaceConfig

		for key, netConf := range cfg.Network.Interfaces {
			if netConf.Match == "" {
				netConf.Match = key
			}

			if netConf.Match == "" {
				continue
			}

			// "" means match has not been found
			if match.Match == "" && matches(linkName, netConf.Match) {
				match = netConf
			}

			if netConf.Match == linkName {
				// Found exact match, use it over wildcard match
				match = netConf
			}
		}

		if match.Match != "" {
			err = applyNetConf(link, match)
			if err != nil {
				log.Errorf("Failed to apply settings to %s : %v", linkName, err)
			}
		}
	}

	if err != nil {
		return err
	}

	//post run
	if cfg.Network.PostRun != nil {
		return docker.StartAndWait(config.DOCKER_SYSTEM_HOST, cfg.Network.PostRun)
	}
	return nil
}

func applyNetConf(link netlink.Link, netConf config.InterfaceConfig) error {
	if netConf.DHCP {
		log.Infof("Running DHCP on %s", link.Attrs().Name)
		cmd := exec.Command("udhcpc", "-i", link.Attrs().Name, "-t", "20", "-n")
		cmd.Stdout = os.Stdout
		cmd.Stderr = os.Stderr
		if err := cmd.Run(); err != nil {
			log.Error(err)
		}
	} else if netConf.Address == "" {
		return nil
	} else {
		addr, err := netlink.ParseAddr(netConf.Address)
		if err != nil {
			return err
		}
		if err := netlink.AddrAdd(link, addr); err != nil {
			log.Error("addr add failed")
			return err
		}
		log.Infof("Set %s on %s", netConf.Address, link.Attrs().Name)
	}

	if netConf.MTU > 0 {
		if err := netlink.LinkSetMTU(link, netConf.MTU); err != nil {
			log.Error("set MTU Failed")
			return err
		}
	}

	if err := netlink.LinkSetUp(link); err != nil {
		log.Error("failed to setup link")
		return err
	}

	if netConf.Gateway != "" {
		gatewayIp := net.ParseIP(netConf.Gateway)
		if gatewayIp == nil {
			return errors.New("Invalid gateway address " + netConf.Gateway)
		}

		route := netlink.Route{
			Scope: netlink.SCOPE_UNIVERSE,
			Gw:    net.ParseIP(netConf.Gateway),
		}
		if err := netlink.RouteAdd(&route); err != nil {
			log.Error("gateway set failed")
			return err
		}

		log.Infof("Set default gateway %s", netConf.Gateway)
	}

	return nil
}

func matches(link, conf string) bool {
	return glob.Glob(conf, link)
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"os"
	"os/signal"
	"sync"
	"time"

	"github.com/segmentio/conf"
	nsq "github.com/segmentio/nsq-go"
)

func main() {
	config := struct {
		LookupdHttpAddr        []string `conf:"lookupd-http-addr"         help:"List of nsqlookupd servers"`
		DestinationNsqdTcpAddr []string `conf:"desintation-nsqd-tcp-addr" help:"List of nsqd nodes to publish to"`
		NsqdTcpAddr            string   `conf:"nsqd-tcp-addr"             help:"Addres of the nsqd nodes to consume from"`
		Topic                  string   `conf:"topic"                     help:"Topic to consume messages from"`
		Channel                string   `conf:"channel"                   help:"Channel to consume messages from"`
		DestinationTopic       string   `conf:"destination-topic"         help:"Topic to publish to"`
		RateLimit              int      `conf:"rate-limit"                help:"Maximum number of message per second processed"`
		MaxInFlight            int      `conf:"max-in-flight" help:"Maximum number of in-flight messages"`
	}{
		LookupdHttpAddr:        []string{"localhost:4161"},
		DestinationNsqdTcpAddr: []string{"localhost:4150"},
		NsqdTcpAddr:            "localhost:4150",
		MaxInFlight:            10,
	}

	conf.Load(&config)

	if len(config.Topic) == 0 {
		log.Fatal("error: missing topic")
	}

	if len(config.Channel) == 0 {
		log.Fatal("error: missing channel")
	}

	if len(config.DestinationTopic) == 0 {
		log.Fatal("error: missing destination channel")
	}

	consumer, err := nsq.StartConsumer(nsq.ConsumerConfig{
		Topic:       config.Topic,
		Channel:     config.Channel,
		Lookup:      config.LookupdHttpAddr,
		Address:     config.NsqdTcpAddr,
		MaxInFlight: config.MaxInFlight,
	})
	if err != nil {
		log.Fatal(err)
	}

	join := &sync.WaitGroup{}

	for _, addr := range config.DestinationNsqdTcpAddr {
		producer, err := nsq.StartProducer(nsq.ProducerConfig{
			Topic:   config.DestinationTopic,
			Address: addr,
		})
		if err != nil {
			log.Fatal(err)
		}
		join.Add(1)
		go forward(producer, rateLimit(config.RateLimit, consumer.Messages()), join)
	}

	sigchan := make(chan os.Signal)
	signal.Notify(sigchan, os.Interrupt)
	<-sigchan

	consumer.Stop()
	join.Wait()
}

func forward(producer *nsq.Producer, messages <-chan nsq.Message, join *sync.WaitGroup) {
	defer join.Done()

	for msg := range messages {
		if err := producer.Publish(msg.Body); err != nil {
			msg.Requeue(15 * time.Second)
			log.Print(err)
		} else {
			msg.Finish()
		}
	}
}

func rateLimit(limit int, messages <-chan nsq.Message) <-chan nsq.Message {
	if limit <= 0 {
		return messages
	}

	output := make(chan nsq.Message, cap(messages))

	go func() {
		ticker := time.NewTicker(1 * time.Second)
		defer ticker.Stop()

		input := messages
		count := 0

		for {
			select {
			case <-ticker.C:
				log.Printf("publish rate of %d msg/sec", count)
				count = 0
				input = messages

			case msg := <-input:
				output <- msg

				if count++; count >= limit {
					input = nil
					log.Printf("rate limiting to %d msg/sec", limit)
				}
			}
		}
	}()

	return output
}
<commit_msg>close output channel on exit<commit_after>package main

import (
	"log"
	"os"
	"os/signal"
	"sync"
	"time"

	"github.com/segmentio/conf"
	nsq "github.com/segmentio/nsq-go"
)

func main() {
	config := struct {
		LookupdHttpAddr        []string `conf:"lookupd-http-addr"         help:"List of nsqlookupd servers"`
		DestinationNsqdTcpAddr []string `conf:"desintation-nsqd-tcp-addr" help:"List of nsqd nodes to publish to"`
		NsqdTcpAddr            string   `conf:"nsqd-tcp-addr"             help:"Addres of the nsqd nodes to consume from"`
		Topic                  string   `conf:"topic"                     help:"Topic to consume messages from"`
		Channel                string   `conf:"channel"                   help:"Channel to consume messages from"`
		DestinationTopic       string   `conf:"destination-topic"         help:"Topic to publish to"`
		RateLimit              int      `conf:"rate-limit"                help:"Maximum number of message per second processed"`
		MaxInFlight            int      `conf:"max-in-flight" help:"Maximum number of in-flight messages"`
	}{
		LookupdHttpAddr:        []string{"localhost:4161"},
		DestinationNsqdTcpAddr: []string{"localhost:4150"},
		NsqdTcpAddr:            "localhost:4150",
		MaxInFlight:            10,
	}

	conf.Load(&config)

	if len(config.Topic) == 0 {
		log.Fatal("error: missing topic")
	}

	if len(config.Channel) == 0 {
		log.Fatal("error: missing channel")
	}

	if len(config.DestinationTopic) == 0 {
		log.Fatal("error: missing destination channel")
	}

	consumer, err := nsq.StartConsumer(nsq.ConsumerConfig{
		Topic:       config.Topic,
		Channel:     config.Channel,
		Lookup:      config.LookupdHttpAddr,
		Address:     config.NsqdTcpAddr,
		MaxInFlight: config.MaxInFlight,
	})
	if err != nil {
		log.Fatal(err)
	}

	join := &sync.WaitGroup{}

	for _, addr := range config.DestinationNsqdTcpAddr {
		producer, err := nsq.StartProducer(nsq.ProducerConfig{
			Topic:   config.DestinationTopic,
			Address: addr,
		})
		if err != nil {
			log.Fatal(err)
		}
		join.Add(1)
		go forward(producer, rateLimit(config.RateLimit, consumer.Messages()), join)
	}

	sigchan := make(chan os.Signal)
	signal.Notify(sigchan, os.Interrupt)
	<-sigchan

	consumer.Stop()
	join.Wait()
}

func forward(producer *nsq.Producer, messages <-chan nsq.Message, join *sync.WaitGroup) {
	defer join.Done()

	for msg := range messages {
		if err := producer.Publish(msg.Body); err != nil {
			msg.Requeue(15 * time.Second)
			log.Print(err)
		} else {
			msg.Finish()
		}
	}
}

func rateLimit(limit int, messages <-chan nsq.Message) <-chan nsq.Message {
	if limit <= 0 {
		return messages
	}

	output := make(chan nsq.Message, cap(messages))

	go func() {
		ticker := time.NewTicker(1 * time.Second)
		defer ticker.Stop()
		defer close(output)

		input := messages
		count := 0

		for {
			select {
			case <-ticker.C:
				log.Printf("publish rate of %d msg/sec", count)
				count = 0
				input = messages

			case msg := <-input:
				output <- msg

				if count++; count >= limit {
					input = nil
					log.Printf("rate limiting to %d msg/sec", limit)
				}
			}
		}
	}()

	return output
}
<|endoftext|>