<commit_before>package command

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

	"github.com/Shopify/sarama"
	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/bjtime"
	"github.com/funkygao/golib/color"
	"github.com/funkygao/golib/gofmt"
	"github.com/funkygao/golib/progress"
	"github.com/funkygao/termui"
	"github.com/nsf/termbox-go"
)

type Top struct {
	Ui  cli.Ui
	Cmd string

	zone           string
	clusterPattern string

	mu sync.Mutex

	round int

	showProgressBar bool
	who             string
	limit           int
	topInterval     int
	batchMode       bool
	dashboardGraph  bool
	topicPattern    string

	counters         map[string]float64 // key is cluster:topic
	lastCounters     map[string]float64
	consumerCounters map[string]float64

	totalConsumerMps []float64
	totalMps         []float64 // for the dashboard graph
	maxMps           float64
}

func (this *Top) Run(args []string) (exitCode int) {
	cmdFlags := flag.NewFlagSet("top", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&this.zone, "z", "", "")
	cmdFlags.StringVar(&this.topicPattern, "t", "", "")
	cmdFlags.IntVar(&this.topInterval, "interval", 5, "refresh interval")
	cmdFlags.StringVar(&this.clusterPattern, "c", "", "")
	cmdFlags.IntVar(&this.limit, "n", 33, "")
	cmdFlags.StringVar(&this.who, "who", "producer", "")
	cmdFlags.BoolVar(&this.showProgressBar, "bar", false, "")
	cmdFlags.BoolVar(&this.dashboardGraph, "d", false, "")
	cmdFlags.BoolVar(&this.batchMode, "b", false, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	if validateArgs(this, this.Ui).
		require("-z").
		invalid(args) {
		return 2
	}

	if this.dashboardGraph {
		this.topInterval = 20
		this.who = "both"
		go this.clusterOffsetSummary()
	}

	if this.who == "c" || this.who == "consumer" {
		if this.topInterval < 20 {
			this.topInterval = 20 // consumer groups only refresh offset per minute
		}

	}

	this.counters = make(map[string]float64)
	this.lastCounters = make(map[string]float64)
	this.consumerCounters = make(map[string]float64)
	this.totalMps = make([]float64, 0, 1000)
	this.totalConsumerMps = make([]float64, 0, 1000)

	zkzone := zk.NewZkZone(zk.DefaultConfig(this.zone, ctx.ZoneZkAddrs(this.zone)))
	zkzone.ForSortedClusters(func(zkcluster *zk.ZkCluster) {
		if !patternMatched(zkcluster.Name(), this.clusterPattern) {
			return
		}

		switch this.who {
		case "p", "producer":
			go this.clusterTopProducers(zkcluster)

		case "c", "consumer":
			go this.clusterTopConsumers(zkcluster)

		case "both":
			go this.clusterTopConsumers(zkcluster)
			go this.clusterTopProducers(zkcluster)

		default:
			this.Ui.Error(fmt.Sprintf("unknown type: %s", this.who))
		}
	})

	if this.dashboardGraph {
		this.drawDashboard()
		return
	}

	bar := progress.New(this.topInterval)
	for {
		if this.batchMode {
			this.Ui.Output(bjtime.TimeToString(bjtime.NowBj()))
		} else {
			refreshScreen()
		}

		// header
		this.Ui.Output(fmt.Sprintf("%-9s %20s %50s %20s %15s",
			this.who, "cluster", "topic", "num", "mps")) // mps=msg per second
		this.Ui.Output(fmt.Sprintf(strings.Repeat("-", 118)))

		this.showAndResetCounters()

		if !this.batchMode {
			this.showRefreshBar(bar)
		} else {
			time.Sleep(time.Duration(this.topInterval) * time.Second)
		}

	}

	return
}

func (this *Top) drawDashboard() {
	err := termui.Init()
	if err != nil {
		panic(err)
	}
	defer termui.Close()

	termui.UseTheme("helloworld")

	refreshProducerData := func() []float64 {
		this.showAndResetCounters()
		return this.totalMps
	}

	refreshConsumerData := func() []float64 {
		return this.totalConsumerMps
	}

	producerChart := termui.NewLineChart()
	producerChart.Mode = "dot"
	producerChart.Border.Label = fmt.Sprintf("producer mps totals: %s %s %s",
		this.zone, this.clusterPattern, this.topicPattern)
	producerChart.Data = refreshProducerData()
	producerChart.Width = termui.TermWidth()
	producerChart.Height = termui.TermHeight() / 2
	producerChart.X = 0
	producerChart.Y = 0
	producerChart.AxesColor = termui.ColorWhite
	producerChart.LineColor = termui.ColorGreen | termui.AttrBold

	consumerChart := termui.NewLineChart()
	consumerChart.Mode = "dot"
	consumerChart.Border.Label = fmt.Sprintf("consumer mps totals: %s %s %s",
		this.zone, this.clusterPattern, this.topicPattern)
	consumerChart.Data = refreshConsumerData()
	consumerChart.Width = termui.TermWidth()
	consumerChart.Height = termui.TermHeight() / 2
	consumerChart.X = 0
	consumerChart.Y = termui.TermHeight() / 2
	consumerChart.AxesColor = termui.ColorWhite
	consumerChart.LineColor = termui.ColorRed | termui.AttrBold

	evt := make(chan termbox.Event)
	go func() {
		for {
			evt <- termbox.PollEvent()
		}
	}()

	termui.Render(producerChart, consumerChart)
	tick := time.NewTicker(time.Duration(this.topInterval) * time.Second)
	defer tick.Stop()
	rounds := 0
	for {
		select {
		case e := <-evt:
			if e.Type == termbox.EventKey && e.Ch == 'q' {
				return
			}

		case <-tick.C:
			// refresh data, and skip the first 2 rounds
			rounds++
			if rounds > 1 {
				producerChart.Data = refreshProducerData()
				consumerChart.Data = refreshConsumerData()
				termui.Render(producerChart, consumerChart)
			}

		}
	}
}

func (this *Top) showRefreshBar(bar *progress.Progress) {
	this.Ui.Output("")
	for i := 1; i <= this.topInterval; i++ {
		if this.showProgressBar {
			bar.ShowProgress(i)
		}

		time.Sleep(time.Second)
	}
}

func (this *Top) showAndResetCounters() {
	this.mu.Lock()
	defer this.mu.Unlock()

	// FIXME counterFlip should be map[int][]string
	counterFlip := make(map[float64]string)
	sortedNum := make([]float64, 0, len(this.counters))
	for ct, num := range this.counters {
		if this.topicPattern != "" && !strings.HasSuffix(ct, ":"+this.topicPattern) {
			continue
		}

		counterFlip[num] = ct
		if num > 100 { // TODO kill the magic number
			sortedNum = append(sortedNum, num)
		}
	}
	sort.Float64s(sortedNum)

	othersNum := 0.
	othersMps := 0.
	totalNum := 0.
	totalMps := 0.
	limitReached := false
	for i := len(sortedNum) - 1; i >= 0; i-- {
		if !limitReached && this.limit > 0 && len(sortedNum)-i > this.limit {
			limitReached = true
		}

		num := sortedNum[i]
		mps := float64(num-this.lastCounters[counterFlip[num]]) / float64(this.topInterval) // msg per sec
		if this.round > 1 {
			totalNum += num
			totalMps += mps
		}

		if limitReached {
			othersNum += num
			othersMps += mps
		} else if !this.dashboardGraph {
			clusterAndTopic := strings.SplitN(counterFlip[num], ":", 2)
			this.Ui.Output(fmt.Sprintf("%30s %50s %20s %15.2f",
				clusterAndTopic[0], clusterAndTopic[1],
				gofmt.Comma(int64(num)),
				mps))
		}
	}

	// display the summary footer
	this.round++
	if this.dashboardGraph {
		if len(this.totalMps) > 5000 {
			// too long, so reset
			this.totalMps = make([]float64, 0, 1000)
		}

		this.totalMps = append(this.totalMps, totalMps)
	} else {
		// the catchall row
		this.Ui.Output(fmt.Sprintf("%30s %50s %20s %15.2f",
			"-OTHERS-", "-OTHERS-",
			gofmt.Comma(int64(othersNum)),
			othersMps))

		// total row
		this.Ui.Output(fmt.Sprintf("%30s %50s %20s %15.2f",
			"--TOTAL--", "--TOTAL--",
			gofmt.Comma(int64(totalNum)),
			totalMps))

		// max
		if this.maxMps < totalMps {
			this.maxMps = totalMps
		}
		this.Ui.Output(fmt.Sprintf("%30s %50s %20s %15.2f",
			"--MAX--", "--MAX--",
			"-",
			this.maxMps))
	}

	// record last counters and reset current counters
	for k, v := range this.counters {
		this.lastCounters[k] = v
	}
	this.counters = make(map[string]float64)
}

func (this *Top) clusterOffsetSummary() {
	var lastOffsets float64
	var total float64
	for {
		total = 0.
		this.mu.Lock()
		for _, n := range this.consumerCounters {
			total += n
		}
		this.mu.Unlock()

		if lastOffsets > 1 {
			this.totalConsumerMps = append(this.totalConsumerMps,
				(float64(total)-lastOffsets)/float64(this.topInterval))
		}

		lastOffsets = float64(total)

		time.Sleep(time.Second * time.Duration(this.topInterval))
	}
}

func (this *Top) clusterTopConsumers(zkcluster *zk.ZkCluster) {
	var topic string
	for {
		total := zkcluster.TotalConsumerOffsets(this.topicPattern)
		if this.topicPattern != "" {
			topic = this.topicPattern
		} else {
			topic = "-all-"
		}

		key := zkcluster.Name() + ":" + topic

		this.mu.Lock()
		this.consumerCounters[key] = float64(total)
		this.mu.Unlock()

		if !this.dashboardGraph {
			this.mu.Lock()
			this.counters[key] = float64(total)
			this.mu.Unlock()
		}

		time.Sleep(time.Second * time.Duration(this.topInterval))
	}

}

func (this *Top) clusterTopProducers(zkcluster *zk.ZkCluster) {
	cluster := zkcluster.Name()
	brokerList := zkcluster.BrokerList()
	if len(brokerList) == 0 {
		return
	}

	kfk, err := sarama.NewClient(brokerList, sarama.NewConfig())
	if err != nil {
		return
	}
	defer kfk.Close()

	for {
		topics, err := kfk.Topics()
		if err != nil || len(topics) == 0 {
			return
		}

		for _, topic := range topics {
			if !patternMatched(topic, this.topicPattern) {
				continue
			}

			msgs := int64(0)
			alivePartitions, err := kfk.WritablePartitions(topic)
			if err != nil {
				panic(err)
			}

			for _, partitionID := range alivePartitions {
				latestOffset, err := kfk.GetOffset(topic, partitionID,
					sarama.OffsetNewest)
				if err != nil {
					panic(err)
				}

				msgs += latestOffset
			}

			this.mu.Lock()
			this.counters[cluster+":"+topic] = float64(msgs)
			this.mu.Unlock()
		}

		time.Sleep(time.Second)
		kfk.RefreshMetadata(topics...)
	}

}

func (*Top) Synopsis() string {
	return "Unix “top” like utility for kafka"
}

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

    Unix “top” like utility for kafka

Options:

    -z zone

    -c cluster pattern

    -t topic pattern

    -interval interval
      Refresh interval in seconds.

    -n limit

    -d
      Draw dashboard in graph.

    -bar
      Show progress bar.

    -b
      Batch mode operation.
      Could be useful for sending output from top to other programs or to a file.

    -who <%s%s|%s%s>
`, this.Cmd, color.Colorize([]string{color.Underscore}, "p"), "roducer",
		color.Colorize([]string{color.Underscore}, "c"), "onsumer")
	return strings.TrimSpace(help)
}
<commit_msg>make a side-by-side compare more easy<commit_after>package command

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

	"github.com/Shopify/sarama"
	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/zk"
	"github.com/funkygao/gocli"
	"github.com/funkygao/golib/bjtime"
	"github.com/funkygao/golib/color"
	"github.com/funkygao/golib/gofmt"
	"github.com/funkygao/golib/progress"
	"github.com/funkygao/termui"
	"github.com/nsf/termbox-go"
)

type Top struct {
	Ui  cli.Ui
	Cmd string

	zone           string
	clusterPattern string

	mu sync.Mutex

	round int

	showProgressBar bool
	who             string
	limit           int
	topInterval     int
	batchMode       bool
	dashboardGraph  bool
	topicPattern    string

	counters         map[string]float64 // key is cluster:topic
	lastCounters     map[string]float64
	consumerCounters map[string]float64

	totalConsumerMps []float64
	totalMps         []float64 // for the dashboard graph
	maxMps           float64
}

func (this *Top) Run(args []string) (exitCode int) {
	cmdFlags := flag.NewFlagSet("top", flag.ContinueOnError)
	cmdFlags.Usage = func() { this.Ui.Output(this.Help()) }
	cmdFlags.StringVar(&this.zone, "z", "", "")
	cmdFlags.StringVar(&this.topicPattern, "t", "", "")
	cmdFlags.IntVar(&this.topInterval, "interval", 5, "refresh interval")
	cmdFlags.StringVar(&this.clusterPattern, "c", "", "")
	cmdFlags.IntVar(&this.limit, "n", 33, "")
	cmdFlags.StringVar(&this.who, "who", "producer", "")
	cmdFlags.BoolVar(&this.showProgressBar, "bar", false, "")
	cmdFlags.BoolVar(&this.dashboardGraph, "d", false, "")
	cmdFlags.BoolVar(&this.batchMode, "b", false, "")
	if err := cmdFlags.Parse(args); err != nil {
		return 1
	}

	if validateArgs(this, this.Ui).
		require("-z").
		invalid(args) {
		return 2
	}

	if this.dashboardGraph {
		this.topInterval = 20
		this.who = "both"
		go this.clusterOffsetSummary()
	}

	if this.who == "c" || this.who == "consumer" {
		if this.topInterval < 20 {
			this.topInterval = 20 // consumer groups only refresh offset per minute
		}

	}

	this.counters = make(map[string]float64)
	this.lastCounters = make(map[string]float64)
	this.consumerCounters = make(map[string]float64)
	this.totalMps = make([]float64, 0, 1000)
	this.totalConsumerMps = make([]float64, 0, 1000)

	zkzone := zk.NewZkZone(zk.DefaultConfig(this.zone, ctx.ZoneZkAddrs(this.zone)))
	zkzone.ForSortedClusters(func(zkcluster *zk.ZkCluster) {
		if !patternMatched(zkcluster.Name(), this.clusterPattern) {
			return
		}

		switch this.who {
		case "p", "producer":
			go this.clusterTopProducers(zkcluster)

		case "c", "consumer":
			go this.clusterTopConsumers(zkcluster)

		case "both":
			go this.clusterTopConsumers(zkcluster)
			go this.clusterTopProducers(zkcluster)

		default:
			this.Ui.Error(fmt.Sprintf("unknown type: %s", this.who))
		}
	})

	if this.dashboardGraph {
		this.drawDashboard()
		return
	}

	bar := progress.New(this.topInterval)
	for {
		if this.batchMode {
			this.Ui.Output(bjtime.TimeToString(bjtime.NowBj()))
		} else {
			refreshScreen()
		}

		// header
		this.Ui.Output(fmt.Sprintf("%-9s %20s %50s %20s %15s",
			this.who, "cluster", "topic", "num", "mps")) // mps=msg per second
		this.Ui.Output(fmt.Sprintf(strings.Repeat("-", 118)))

		this.showAndResetCounters()

		if !this.batchMode {
			this.showRefreshBar(bar)
		} else {
			time.Sleep(time.Duration(this.topInterval) * time.Second)
		}

	}

	return
}

func (this *Top) drawDashboard() {
	err := termui.Init()
	if err != nil {
		panic(err)
	}
	defer termui.Close()

	termui.UseTheme("helloworld")

	refreshProducerData := func() []float64 {
		this.showAndResetCounters()
		return this.totalMps
	}

	refreshConsumerData := func() []float64 {
		return this.totalConsumerMps
	}

	producerChart := termui.NewLineChart()
	producerChart.Mode = "dot"
	producerChart.Border.Label = fmt.Sprintf("producer mps totals: %s %s %s",
		this.zone, this.clusterPattern, this.topicPattern)
	producerChart.Data = refreshProducerData()
	producerChart.Width = termui.TermWidth() / 2
	producerChart.Height = termui.TermHeight()
	producerChart.X = 0
	producerChart.Y = 0
	producerChart.AxesColor = termui.ColorWhite
	producerChart.LineColor = termui.ColorGreen | termui.AttrBold

	consumerChart := termui.NewLineChart()
	consumerChart.Mode = "dot"
	consumerChart.Border.Label = fmt.Sprintf("consumer mps totals: %s %s %s",
		this.zone, this.clusterPattern, this.topicPattern)
	consumerChart.Data = refreshConsumerData()
	consumerChart.Width = termui.TermWidth() / 2
	consumerChart.Height = termui.TermHeight()
	consumerChart.X = termui.TermWidth() / 2
	consumerChart.Y = 0
	consumerChart.AxesColor = termui.ColorWhite
	consumerChart.LineColor = termui.ColorRed | termui.AttrBold

	evt := make(chan termbox.Event)
	go func() {
		for {
			evt <- termbox.PollEvent()
		}
	}()

	termui.Render(producerChart, consumerChart)
	tick := time.NewTicker(time.Duration(this.topInterval) * time.Second)
	defer tick.Stop()
	rounds := 0
	for {
		select {
		case e := <-evt:
			if e.Type == termbox.EventKey && e.Ch == 'q' {
				return
			}

		case <-tick.C:
			// refresh data, and skip the first 2 rounds
			rounds++
			if rounds > 1 {
				producerChart.Data = refreshProducerData()
				consumerChart.Data = refreshConsumerData()
				termui.Render(producerChart, consumerChart)
			}

		}
	}
}

func (this *Top) showRefreshBar(bar *progress.Progress) {
	this.Ui.Output("")
	for i := 1; i <= this.topInterval; i++ {
		if this.showProgressBar {
			bar.ShowProgress(i)
		}

		time.Sleep(time.Second)
	}
}

func (this *Top) showAndResetCounters() {
	this.mu.Lock()
	defer this.mu.Unlock()

	// FIXME counterFlip should be map[int][]string
	counterFlip := make(map[float64]string)
	sortedNum := make([]float64, 0, len(this.counters))
	for ct, num := range this.counters {
		if this.topicPattern != "" && !strings.HasSuffix(ct, ":"+this.topicPattern) {
			continue
		}

		counterFlip[num] = ct
		if num > 100 { // TODO kill the magic number
			sortedNum = append(sortedNum, num)
		}
	}
	sort.Float64s(sortedNum)

	othersNum := 0.
	othersMps := 0.
	totalNum := 0.
	totalMps := 0.
	limitReached := false
	for i := len(sortedNum) - 1; i >= 0; i-- {
		if !limitReached && this.limit > 0 && len(sortedNum)-i > this.limit {
			limitReached = true
		}

		num := sortedNum[i]
		mps := float64(num-this.lastCounters[counterFlip[num]]) / float64(this.topInterval) // msg per sec
		if this.round > 1 {
			totalNum += num
			totalMps += mps
		}

		if limitReached {
			othersNum += num
			othersMps += mps
		} else if !this.dashboardGraph {
			clusterAndTopic := strings.SplitN(counterFlip[num], ":", 2)
			this.Ui.Output(fmt.Sprintf("%30s %50s %20s %15.2f",
				clusterAndTopic[0], clusterAndTopic[1],
				gofmt.Comma(int64(num)),
				mps))
		}
	}

	// display the summary footer
	this.round++
	if this.dashboardGraph {
		if len(this.totalMps) > 5000 {
			// too long, so reset
			this.totalMps = make([]float64, 0, 1000)
		}

		this.totalMps = append(this.totalMps, totalMps)
	} else {
		// the catchall row
		this.Ui.Output(fmt.Sprintf("%30s %50s %20s %15.2f",
			"-OTHERS-", "-OTHERS-",
			gofmt.Comma(int64(othersNum)),
			othersMps))

		// total row
		this.Ui.Output(fmt.Sprintf("%30s %50s %20s %15.2f",
			"--TOTAL--", "--TOTAL--",
			gofmt.Comma(int64(totalNum)),
			totalMps))

		// max
		if this.maxMps < totalMps {
			this.maxMps = totalMps
		}
		this.Ui.Output(fmt.Sprintf("%30s %50s %20s %15.2f",
			"--MAX--", "--MAX--",
			"-",
			this.maxMps))
	}

	// record last counters and reset current counters
	for k, v := range this.counters {
		this.lastCounters[k] = v
	}
	this.counters = make(map[string]float64)
}

func (this *Top) clusterOffsetSummary() {
	var lastOffsets float64
	var total float64
	for {
		total = 0.
		this.mu.Lock()
		for _, n := range this.consumerCounters {
			total += n
		}
		this.mu.Unlock()

		if lastOffsets > 1 {
			this.totalConsumerMps = append(this.totalConsumerMps,
				(float64(total)-lastOffsets)/float64(this.topInterval))
		}

		lastOffsets = float64(total)

		time.Sleep(time.Second * time.Duration(this.topInterval))
	}
}

func (this *Top) clusterTopConsumers(zkcluster *zk.ZkCluster) {
	var topic string
	for {
		total := zkcluster.TotalConsumerOffsets(this.topicPattern)
		if this.topicPattern != "" {
			topic = this.topicPattern
		} else {
			topic = "-all-"
		}

		key := zkcluster.Name() + ":" + topic

		this.mu.Lock()
		this.consumerCounters[key] = float64(total)
		this.mu.Unlock()

		if !this.dashboardGraph {
			this.mu.Lock()
			this.counters[key] = float64(total)
			this.mu.Unlock()
		}

		time.Sleep(time.Second * time.Duration(this.topInterval))
	}

}

func (this *Top) clusterTopProducers(zkcluster *zk.ZkCluster) {
	cluster := zkcluster.Name()
	brokerList := zkcluster.BrokerList()
	if len(brokerList) == 0 {
		return
	}

	kfk, err := sarama.NewClient(brokerList, sarama.NewConfig())
	if err != nil {
		return
	}
	defer kfk.Close()

	for {
		topics, err := kfk.Topics()
		if err != nil || len(topics) == 0 {
			return
		}

		for _, topic := range topics {
			if !patternMatched(topic, this.topicPattern) {
				continue
			}

			msgs := int64(0)
			alivePartitions, err := kfk.WritablePartitions(topic)
			if err != nil {
				panic(err)
			}

			for _, partitionID := range alivePartitions {
				latestOffset, err := kfk.GetOffset(topic, partitionID,
					sarama.OffsetNewest)
				if err != nil {
					panic(err)
				}

				msgs += latestOffset
			}

			this.mu.Lock()
			this.counters[cluster+":"+topic] = float64(msgs)
			this.mu.Unlock()
		}

		time.Sleep(time.Second)
		kfk.RefreshMetadata(topics...)
	}

}

func (*Top) Synopsis() string {
	return "Unix “top” like utility for kafka"
}

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

    Unix “top” like utility for kafka

Options:

    -z zone

    -c cluster pattern

    -t topic pattern

    -interval interval
      Refresh interval in seconds.

    -n limit

    -d
      Draw dashboard in graph.

    -bar
      Show progress bar.

    -b
      Batch mode operation.
      Could be useful for sending output from top to other programs or to a file.

    -who <%s%s|%s%s>
`, this.Cmd, color.Colorize([]string{color.Underscore}, "p"), "roducer",
		color.Colorize([]string{color.Underscore}, "c"), "onsumer")
	return strings.TrimSpace(help)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"strconv"

	"github.com/getbread/goose/lib/goose"
)

var downCmd = &Command{
	Name:    "down",
	Usage:   "",
	Summary: "Roll back the version by 1",
	Help:    `down extended help here...`,
	Run:     downRun,
}

func downRun(cmd *Command, args ...string) {

	conf, err := dbConfFromFlags()
	if err != nil {
		log.Fatal(err)
	}

	current, err := goose.GetDBVersion(conf)
	if err != nil {
		log.Fatal(err)
	}

	var previous int64

	if len(args) > 1 {
		previous, err = goose.GetEarliestSharedDBVersion(args[0], args[1])
		if err != nil {
			log.Fatal(fmt.Sprintf("Failed to determine the lowest migration number between %s and %s due to error: %s", args[0], args[1], err.Error()))
		}

	} else if len(args) > 0 {
		previous, err = strconv.ParseInt(args[0], 10, 64)
		if err != nil {
			log.Fatal(fmt.Sprintf("Failed to parse %s as a migration version number due to error: %s", args[0], err.Error()))
		}

	} else {
		previous, err = goose.GetPreviousDBVersion(conf.MigrationsDir, current)
		if err != nil {
			log.Fatal(err)
		}
	}

	log.Printf("Migrating down to %d\n", previous)

	if err = goose.RunMigrations(conf, conf.MigrationsDir, previous); err != nil {
		log.Fatal(err)
	}
}
<commit_msg>Forcing the migrations dir to be the original directory.<commit_after>package main

import (
	"fmt"
	"log"
	"strconv"

	"github.com/getbread/goose/lib/goose"
)

var downCmd = &Command{
	Name:    "down",
	Usage:   "",
	Summary: "Roll back the version by 1",
	Help:    `down extended help here...`,
	Run:     downRun,
}

func downRun(cmd *Command, args ...string) {

	conf, err := dbConfFromFlags()
	if err != nil {
		log.Fatal(err)
	}

	current, err := goose.GetDBVersion(conf)
	if err != nil {
		log.Fatal(err)
	}

	var previous int64

	if len(args) > 1 {
		previous, err = goose.GetEarliestSharedDBVersion(args[0], args[1])
		if err != nil {
			log.Fatal(fmt.Sprintf("Failed to determine the lowest migration number between %s and %s due to error: %s", args[0], args[1], err.Error()))
		}
		conf.MigrationsDir = args[0]

	} else if len(args) > 0 {
		previous, err = strconv.ParseInt(args[0], 10, 64)
		if err != nil {
			log.Fatal(fmt.Sprintf("Failed to parse %s as a migration version number due to error: %s", args[0], err.Error()))
		}

	} else {
		previous, err = goose.GetPreviousDBVersion(conf.MigrationsDir, current)
		if err != nil {
			log.Fatal(err)
		}
	}

	log.Printf("Migrating down to %d\n", previous)

	if err = goose.RunMigrations(conf, conf.MigrationsDir, previous); err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2016, Ben Morgan. All rights reserved.
// Use of this source code is governed by an MIT license
// that can be found in the LICENSE file.

package main

import (
	"os"
	"text/template"
	"time"

	"github.com/spf13/cobra"
)

func init() {
	MainCmd.AddCommand(versionCmd)
}

type programInfo struct {
	Name      string
	Author    string
	Email     string
	Version   string
	Date      string
	Homepage  string
	Copyright string
	License   string
}

const versionTmpl = `{{.Name}} version {{.Version}} ({{.Date}})
Copyright {{.Copyright}}, {{.Author}} <{{.Email}}>

You may find {{.Name}} on the Internet at
    {{.Homepage}}
Please report any bugs you may encounter.

The source code of {{.Name}} is licensed under the {{.License}} license.
`

var progInfo = programInfo{
	Name:      "lackey",
	Author:    "Ben Morgan",
	Email:     "neembi@gmail.com",
	Version:   "0.4",
	Date:      "",
	Copyright: "2016–2017",
	Homepage:  "https://github.com/cassava/lackey",
	License:   "MIT",
}

var versionCmd = &cobra.Command{
	Use:               "version",
	Short:             "show version and date information",
	Long:              "Show the official version number of lackey, as well as the release date.",
	PersistentPreRunE: func(cmd *cobra.Command, args []string) error { return nil },
	Run: func(cmd *cobra.Command, args []string) {
		if progInfo.Date == "" {
			progInfo.Date = time.Now().Format("2 January 2006")
		}
		template.Must(template.New("version").Parse(versionTmpl)).Execute(os.Stdout, progInfo)
	},
}
<commit_msg>lackey: fix beta timestamp for version command<commit_after>// Copyright (c) 2016, Ben Morgan. All rights reserved.
// Use of this source code is governed by an MIT license
// that can be found in the LICENSE file.

package main

import (
	"os"
	"text/template"
	"time"

	"github.com/kardianos/osext"
	"github.com/spf13/cobra"
)

func init() {
	MainCmd.AddCommand(versionCmd)
}

type programInfo struct {
	Name      string
	Author    string
	Email     string
	Version   string
	Date      string
	Homepage  string
	Copyright string
	License   string
}

const versionTmpl = `{{.Name}} version {{.Version}} ({{.Date}})
Copyright {{.Copyright}}, {{.Author}} <{{.Email}}>

You may find {{.Name}} on the Internet at
    {{.Homepage}}
Please report any bugs you may encounter.

The source code of {{.Name}} is licensed under the {{.License}} license.
`

var progInfo = programInfo{
	Name:      "lackey",
	Author:    "Ben Morgan",
	Email:     "neembi@gmail.com",
	Version:   "0.4",
	Date:      "",
	Copyright: "2016–2017",
	Homepage:  "https://github.com/cassava/lackey",
	License:   "MIT",
}

var versionCmd = &cobra.Command{
	Use:               "version",
	Short:             "show version and date information",
	Long:              "Show the official version number of lackey, as well as the release date.",
	PersistentPreRunE: func(cmd *cobra.Command, args []string) error { return nil },
	Run: func(cmd *cobra.Command, args []string) {
		if progInfo.Date == "" {
			progInfo.Date = versionDate().Format("2 January 2006")
		}
		template.Must(template.New("version").Parse(versionTmpl)).Execute(os.Stdout, progInfo)
	},
}

// versionDate returns the modified time of the executable or the
// current time.
func versionDate() time.Time {
	fn, err := osext.Executable()
	if err != nil {
		return time.Now()
	}

	fi, err := os.Stat(fn)
	if err != nil {
		return time.Now()
	}

	return fi.ModTime()
}
<|endoftext|>


<commit_before>package main

import (
	cli "github.com/hivdb/nucamino/cli/cli"
	hcv1acli "github.com/hivdb/nucamino/cli/hcv1a"
	"github.com/pkg/profile"
)

type HCV1AOptions struct {
	Quiet                    bool         `short:"q" long:"quiet" description:"hide non-error information output"`
	Genes                    []string     `short:"g" long:"gene" required:"yes" choice:"NS3" choice:"NS5A" choice:"NS5B" description:"gene(s) the input sequences should be aligned with"`
	IndelCodonOpeningBonus   int          `long:"indel-codon-opening-bonus" value-name:"BONUS" description:"bonus score when a indel codon was opened" default:"0"`
	IndelCodonExtensionBonus int          `long:"indel-codon-extension-bonus" value-name:"BONUS" description:"bonus score when a indel codon was extended" default:"2"`
	StopCodonPenalty         int          `long:"stop-codon-penalty" value-name:"PENALTY" description:"penalty score when a stop codon was met" default:"4"`
	GapOpeningPenalty        int          `long:"gap-opening-penalty" value-name:"PENALTY" description:"penalty score when a gap was opened" default:"10"`
	GapExtensionPenalty      int          `long:"gap-extension-penalty" value-name:"PENALTY" description:"penalty score when a gap was extended" default:"2"`
	Goroutines               int          `long:"goroutines" value-name:"GOROUTINES" description:"number of goroutines the alignment will use. Use the core number when equals to 0" default:"0"`
	OutputFormat             string       `long:"output-format" value-name:"OUTPUT_FORMAT" choice:"tsv" choice:"json" description:"output format of the alignment result" default:"tsv"`
	Files                    FileOptions  `group:"File Options"`
	Pprof                    PprofOptions `group:"Pprof Options"`
}

var hcv1aOptions HCV1AOptions

func (self *HCV1AOptions) Execute(args []string) error {
	if self.Pprof.Pprof {
		defer profile.Start(profile.CPUProfile).Stop()
	}

	ioParams := cli.IOParameters{
		InputFileName:  string(self.Files.Input),
		OutputFileName: string(self.Files.Output),
		OutputFormat:   self.OutputFormat,
	}

	alignmentParams := cli.AlignmentParameters{
		IndelCodonOpeningBonus:   self.IndelCodonOpeningBonus,
		IndelCodonExtensionBonus: self.IndelCodonExtensionBonus,
		StopCodonPenalty:         self.StopCodonPenalty,
		GapOpeningPenalty:        self.GapOpeningPenalty,
		GapExtensionPenalty:      self.GapExtensionPenalty,
	}

	hcv1acli.PerformAlignment(ioParams, self.Genes, self.Goroutines,
		self.Quiet, alignmentParams)

	return nil
}

func init() {
	parser.AddCommand(
		"hcv1a", "Align HCV-1 genotype A sequences",
		"Align input sequences to HCV-1 genotype B using the FDA recommended reference (Genebank accession no. NC_004102). Suported genes: NS3, NS5A, NS5B.",
		&hcv1aOptions)
}
<commit_msg>Typo-fix in CLI interface text<commit_after>package main

import (
	cli "github.com/hivdb/nucamino/cli/cli"
	hcv1acli "github.com/hivdb/nucamino/cli/hcv1a"
	"github.com/pkg/profile"
)

type HCV1AOptions struct {
	Quiet                    bool         `short:"q" long:"quiet" description:"hide non-error information output"`
	Genes                    []string     `short:"g" long:"gene" required:"yes" choice:"NS3" choice:"NS5A" choice:"NS5B" description:"gene(s) the input sequences should be aligned with"`
	IndelCodonOpeningBonus   int          `long:"indel-codon-opening-bonus" value-name:"BONUS" description:"bonus score when a indel codon was opened" default:"0"`
	IndelCodonExtensionBonus int          `long:"indel-codon-extension-bonus" value-name:"BONUS" description:"bonus score when a indel codon was extended" default:"2"`
	StopCodonPenalty         int          `long:"stop-codon-penalty" value-name:"PENALTY" description:"penalty score when a stop codon was met" default:"4"`
	GapOpeningPenalty        int          `long:"gap-opening-penalty" value-name:"PENALTY" description:"penalty score when a gap was opened" default:"10"`
	GapExtensionPenalty      int          `long:"gap-extension-penalty" value-name:"PENALTY" description:"penalty score when a gap was extended" default:"2"`
	Goroutines               int          `long:"goroutines" value-name:"GOROUTINES" description:"number of goroutines the alignment will use. Use the core number when equals to 0" default:"0"`
	OutputFormat             string       `long:"output-format" value-name:"OUTPUT_FORMAT" choice:"tsv" choice:"json" description:"output format of the alignment result" default:"tsv"`
	Files                    FileOptions  `group:"File Options"`
	Pprof                    PprofOptions `group:"Pprof Options"`
}

var hcv1aOptions HCV1AOptions

func (self *HCV1AOptions) Execute(args []string) error {
	if self.Pprof.Pprof {
		defer profile.Start(profile.CPUProfile).Stop()
	}

	ioParams := cli.IOParameters{
		InputFileName:  string(self.Files.Input),
		OutputFileName: string(self.Files.Output),
		OutputFormat:   self.OutputFormat,
	}

	alignmentParams := cli.AlignmentParameters{
		IndelCodonOpeningBonus:   self.IndelCodonOpeningBonus,
		IndelCodonExtensionBonus: self.IndelCodonExtensionBonus,
		StopCodonPenalty:         self.StopCodonPenalty,
		GapOpeningPenalty:        self.GapOpeningPenalty,
		GapExtensionPenalty:      self.GapExtensionPenalty,
	}

	hcv1acli.PerformAlignment(ioParams, self.Genes, self.Goroutines,
		self.Quiet, alignmentParams)

	return nil
}

func init() {
	parser.AddCommand(
		"hcv1a", "Align HCV-1 genotype A sequences",
		"Align input sequences to HCV-1 genotype B using the FDA recommended reference (Genebank accession no. NC_004102). Supported genes: NS3, NS5A, NS5B.",
		&hcv1aOptions)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/xml"
	"flag"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"time"
)

type options struct {
	from   string
	link   string
	prefix string
	set    string
	until  string
}

type request struct {
	opts  options
	token string
	verb  string
}

type response struct {
	Date    string `xml:"responseDate"`
	Request struct {
		Verb  string `xml:"verb,attr"`
		Set   string `xml:"set,attr"`
		From  string `xml:"from,attr"`
		Until string `xml:"until,attr"`
		Link  string `xml:",chardata"`
	} `xml:"request"`
	Error struct {
		Code    string `xml:"code,attr"`
		Message string `xml:",chardata"`
	} `xml:"error"`
	ListRecords string `xml:",innerxml"`
}

func (r request) Link() string {
	v := url.Values{}
	v.Add("from", r.opts.from)
	v.Add("set", r.opts.set)
	v.Add("until", r.opts.until)
	v.Add("metadataPrefix", r.opts.prefix)
	v.Add("verb", r.verb)
	if r.token != "" {
		v.Add("resumptionToken", r.token)
	}
	return fmt.Sprintf("%s?%s", r.opts.link, v.Encode())
}

func main() {

	link := flag.String("link", "", "OAI provider URL")
	// output := flag.String("o", "", "output file")
	from := flag.String("f", "2000-01-01", "from parameter")
	until := flag.String("u", time.Now().Format("2006-01-02"), "until parameter")
	prefix := flag.String("p", "oai_dc", "metadata prefix")
	set := flag.String("s", "", "set name")
	verbose := flag.Bool("verbose", false, "be verbose")

	flag.Parse()

	opts := options{from: *from, until: *until, prefix: *prefix, set: *set, link: *link}
	oair := request{opts: opts, verb: "ListRecords"}

	client := http.Client{}

	if *verbose {
		log.Println(oair.Link())
	}

	req, err := http.NewRequest("GET", oair.Link(), nil)
	if err != nil {
		log.Fatal(err)
	}

	resp, err := client.Do(req)
	if err != nil {
		log.Fatal(err)
	}

	defer resp.Body.Close()
	decoder := xml.NewDecoder(resp.Body)

	for {
		t, err := decoder.Token()
		if t == nil {
			break
		}
		if err != nil {
			log.Fatal(err)
		}
		switch se := t.(type) {
		case xml.StartElement:
			if se.Name.Local == "OAI-PMH" {
				var resp response
				err := decoder.DecodeElement(&resp, &se)
				if err != nil {
					log.Fatal(err)
				}
				if resp.Error.Code != "" {
					log.Fatal(resp.Error.Message)
				}
				fmt.Printf("%+v", resp)
			}
		}
	}
}
<commit_msg>basic token handling<commit_after>package main

import (
	"encoding/xml"
	"flag"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"strings"
	"time"

	"launchpad.net/xmlpath"
)

var path = xmlpath.MustCompile("//resumptionToken")

type options struct {
	from   string
	prefix string
	set    string
	until  string
}

type request struct {
	link  string
	opts  options
	token string
	verb  string
}

type response struct {
	Error struct {
		Code    string `xml:"code,attr"`
		Message string `xml:",chardata"`
	} `xml:"error"`
	Payload string `xml:",innerxml"`
}

func (r request) Link() string {
	v := url.Values{}
	if r.opts.from != "" {
		v.Add("from", r.opts.from)
	}
	if r.opts.set != "" {
		v.Add("set", r.opts.set)
	}
	if r.opts.until != "" {
		v.Add("until", r.opts.until)
	}
	if r.opts.prefix != "" {
		v.Add("metadataPrefix", r.opts.prefix)
	}
	if r.verb != "" {
		v.Add("verb", r.verb)
	}
	if r.token != "" {
		v.Add("resumptionToken", r.token)
	}
	encoded := v.Encode()
	if len(encoded) == 0 {
		return r.link
	}
	return fmt.Sprintf("%s?%s", r.link, encoded)
}

func ExtractToken(s string) string {
	root, err := xmlpath.Parse(strings.NewReader(s))
	if err != nil {
		log.Fatal(err)
	}
	if value, ok := path.String(root); ok {
		return value
	}
	return ""
}

func main() {

	link := flag.String("link", "", "OAI provider URL")
	from := flag.String("f", "2000-01-01", "from parameter")
	until := flag.String("u", time.Now().Format("2006-01-02"), "until parameter")
	prefix := flag.String("p", "oai_dc", "metadata prefix")
	set := flag.String("s", "", "set name")
	verb := flag.String("verb", "ListRecords", "OAI verb")
	verbose := flag.Bool("verbose", false, "be verbose")

	flag.Parse()

	opts := options{from: *from, until: *until, prefix: *prefix, set: *set}
	client := http.Client{}

	oair := request{opts: opts, verb: *verb, link: *link}

Loop:
	for {
		if *verbose {
			log.Println(oair.Link())
		}

		req, err := http.NewRequest("GET", oair.Link(), nil)
		if err != nil {
			log.Fatal(err)
		}

		resp, err := client.Do(req)
		if err != nil {
			log.Fatal(err)
		}

		defer resp.Body.Close()
		decoder := xml.NewDecoder(resp.Body)

		for {
			t, err := decoder.Token()
			if t == nil {
				break
			}
			if err != nil {
				log.Fatal(err)
			}
			switch se := t.(type) {
			case xml.StartElement:
				if se.Name.Local == "OAI-PMH" {
					var resp response
					err := decoder.DecodeElement(&resp, &se)
					if err != nil {
						log.Fatal(err)
					}
					if resp.Error.Code != "" {
						log.Fatal(resp.Error.Message)
					}
					fmt.Println(resp.Payload)
					token := ExtractToken(resp.Payload)
					if token == "" {
						break Loop
					}
					oair = request{token: token, verb: *verb, link: *link}
				}
			}
		}
	}
}
<|endoftext|>


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

package jsonrpc2

import (
	"context"
)

// Handler is the interface used to hook into the message handling of an rpc
// connection.
type Handler interface {
	// Deliver is invoked to handle incoming requests.
	// If the request returns false from IsNotify then the Handler must eventually
	// call Reply on the Conn with the supplied request.
	// Handlers are called synchronously, they should pass the work off to a go
	// routine if they are going to take a long time.
	// If Deliver returns true all subsequent handlers will be invoked with
	// delivered set to true, and should not attempt to deliver the message.
	Deliver(ctx context.Context, r *Request, delivered bool) bool

	// Cancel is invoked for cancelled outgoing requests.
	// It is okay to use the connection to send notifications, but the context will
	// be in the cancelled state, so you must do it with the background context
	// instead.
	// If Cancel returns true all subsequent handlers will be invoked with
	// cancelled set to true, and should not attempt to cancel the message.
	Cancel(ctx context.Context, conn *Conn, id ID, cancelled bool) bool

	// Log is invoked for all messages flowing through a Conn.
	// direction indicates if the message being received or sent
	// id is the message id, if not set it was a notification
	// elapsed is the time between a call being seen and the response, and is
	// negative for anything that is not a response.
	// method is the method name specified in the message
	// payload is the parameters for a call or notification, and the result for a
	// response

	// Request is called near the start of processing any request.
	Request(ctx context.Context, conn *Conn, direction Direction, r *WireRequest) context.Context
	// Response is called near the start of processing any response.
	Response(ctx context.Context, conn *Conn, direction Direction, r *WireResponse) context.Context
	// Done is called when any request is fully processed.
	// For calls, this means the response has also been processed, for notifies
	// this is as soon as the message has been written to the stream.
	// If err is set, it implies the request failed.
	Done(ctx context.Context, err error)
	// Read is called with a count each time some data is read from the stream.
	// The read calls are delayed until after the data has been interpreted so
	// that it can be attributed to a request/response.
	Read(ctx context.Context, bytes int64) context.Context
	// Wrote is called each time some data is written to the stream.
	Wrote(ctx context.Context, bytes int64) context.Context
	// Error is called with errors that cannot be delivered through the normal
	// mechanisms, for instance a failure to process a notify cannot be delivered
	// back to the other party.
	Error(ctx context.Context, err error)
}

// Direction is used to indicate to a logger whether the logged message was being
// sent or received.
type Direction bool

const (
	// Send indicates the message is outgoing.
	Send = Direction(true)
	// Receive indicates the message is incoming.
	Receive = Direction(false)
)

func (d Direction) String() string {
	switch d {
	case Send:
		return "send"
	case Receive:
		return "receive"
	default:
		panic("unreachable")
	}
}

type EmptyHandler struct{}

func (EmptyHandler) Deliver(ctx context.Context, r *Request, delivered bool) bool {
	return false
}

func (EmptyHandler) Cancel(ctx context.Context, conn *Conn, id ID, cancelled bool) bool {
	return false
}

func (EmptyHandler) Request(ctx context.Context, conn *Conn, direction Direction, r *WireRequest) context.Context {
	return ctx
}

func (EmptyHandler) Response(ctx context.Context, conn *Conn, direction Direction, r *WireResponse) context.Context {
	return ctx
}

func (EmptyHandler) Done(ctx context.Context, err error) {
}

func (EmptyHandler) Read(ctx context.Context, bytes int64) context.Context {
	return ctx
}

func (EmptyHandler) Wrote(ctx context.Context, bytes int64) context.Context {
	return ctx
}

func (EmptyHandler) Error(ctx context.Context, err error) {}

type defaultHandler struct{ EmptyHandler }

func (defaultHandler) Deliver(ctx context.Context, r *Request, delivered bool) bool {
	if delivered {
		return false
	}
	if !r.IsNotify() {
		r.Reply(ctx, nil, NewErrorf(CodeMethodNotFound, "method %q not found", r.Method))
	}
	return true
}
<commit_msg>internal/jsonrpc2: remove unnecessary Log comment from Handler interface<commit_after>// Copyright 2019 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package jsonrpc2

import (
	"context"
)

// Handler is the interface used to hook into the message handling of an rpc
// connection.
type Handler interface {
	// Deliver is invoked to handle incoming requests.
	// If the request returns false from IsNotify then the Handler must eventually
	// call Reply on the Conn with the supplied request.
	// Handlers are called synchronously, they should pass the work off to a go
	// routine if they are going to take a long time.
	// If Deliver returns true all subsequent handlers will be invoked with
	// delivered set to true, and should not attempt to deliver the message.
	Deliver(ctx context.Context, r *Request, delivered bool) bool

	// Cancel is invoked for cancelled outgoing requests.
	// It is okay to use the connection to send notifications, but the context will
	// be in the cancelled state, so you must do it with the background context
	// instead.
	// If Cancel returns true all subsequent handlers will be invoked with
	// cancelled set to true, and should not attempt to cancel the message.
	Cancel(ctx context.Context, conn *Conn, id ID, cancelled bool) bool

	// Request is called near the start of processing any request.
	Request(ctx context.Context, conn *Conn, direction Direction, r *WireRequest) context.Context
	// Response is called near the start of processing any response.
	Response(ctx context.Context, conn *Conn, direction Direction, r *WireResponse) context.Context
	// Done is called when any request is fully processed.
	// For calls, this means the response has also been processed, for notifies
	// this is as soon as the message has been written to the stream.
	// If err is set, it implies the request failed.
	Done(ctx context.Context, err error)
	// Read is called with a count each time some data is read from the stream.
	// The read calls are delayed until after the data has been interpreted so
	// that it can be attributed to a request/response.
	Read(ctx context.Context, bytes int64) context.Context
	// Wrote is called each time some data is written to the stream.
	Wrote(ctx context.Context, bytes int64) context.Context
	// Error is called with errors that cannot be delivered through the normal
	// mechanisms, for instance a failure to process a notify cannot be delivered
	// back to the other party.
	Error(ctx context.Context, err error)
}

// Direction is used to indicate to a logger whether the logged message was being
// sent or received.
type Direction bool

const (
	// Send indicates the message is outgoing.
	Send = Direction(true)
	// Receive indicates the message is incoming.
	Receive = Direction(false)
)

func (d Direction) String() string {
	switch d {
	case Send:
		return "send"
	case Receive:
		return "receive"
	default:
		panic("unreachable")
	}
}

type EmptyHandler struct{}

func (EmptyHandler) Deliver(ctx context.Context, r *Request, delivered bool) bool {
	return false
}

func (EmptyHandler) Cancel(ctx context.Context, conn *Conn, id ID, cancelled bool) bool {
	return false
}

func (EmptyHandler) Request(ctx context.Context, conn *Conn, direction Direction, r *WireRequest) context.Context {
	return ctx
}

func (EmptyHandler) Response(ctx context.Context, conn *Conn, direction Direction, r *WireResponse) context.Context {
	return ctx
}

func (EmptyHandler) Done(ctx context.Context, err error) {
}

func (EmptyHandler) Read(ctx context.Context, bytes int64) context.Context {
	return ctx
}

func (EmptyHandler) Wrote(ctx context.Context, bytes int64) context.Context {
	return ctx
}

func (EmptyHandler) Error(ctx context.Context, err error) {}

type defaultHandler struct{ EmptyHandler }

func (defaultHandler) Deliver(ctx context.Context, r *Request, delivered bool) bool {
	if delivered {
		return false
	}
	if !r.IsNotify() {
		r.Reply(ctx, nil, NewErrorf(CodeMethodNotFound, "method %q not found", r.Method))
	}
	return true
}
<|endoftext|>


<commit_before>package leafnodes

import (
	"fmt"
	"github.com/onsi/ginkgo/internal/codelocation"
	"github.com/onsi/ginkgo/internal/failer"
	"github.com/onsi/ginkgo/types"
	"reflect"
	"time"
)

type runner struct {
	isAsync          bool
	asyncFunc        func(chan<- interface{})
	syncFunc         func()
	codeLocation     types.CodeLocation
	timeoutThreshold time.Duration
	nodeType         types.SpecComponentType
	componentIndex   int
	failer           *failer.Failer
}

func newRunner(body interface{}, codeLocation types.CodeLocation, timeout time.Duration, failer *failer.Failer, nodeType types.SpecComponentType, componentIndex int) *runner {
	bodyType := reflect.TypeOf(body)
	if bodyType.Kind() != reflect.Func {
		panic(fmt.Sprintf("Expected a function but got something else at %v", codeLocation))
	}

	runner := &runner{
		codeLocation:     codeLocation,
		timeoutThreshold: timeout,
		failer:           failer,
		nodeType:         nodeType,
		componentIndex:   componentIndex,
	}

	switch bodyType.NumIn() {
	case 0:
		runner.syncFunc = body.(func())
		return runner
	case 1:
		if !(bodyType.In(0).Kind() == reflect.Chan && bodyType.In(0).Elem().Kind() == reflect.Interface) {
			panic(fmt.Sprintf("Must pass a Done channel to function at %v", codeLocation))
		}

		wrappedBody := func(done chan<- interface{}) {
			bodyValue := reflect.ValueOf(body)
			bodyValue.Call([]reflect.Value{reflect.ValueOf(done)})
		}

		runner.isAsync = true
		runner.asyncFunc = wrappedBody
		return runner
	}

	panic(fmt.Sprintf("Too many arguments to function at %v", codeLocation))
}

func (r *runner) run() (outcome types.SpecState, failure types.SpecFailure) {
	if r.isAsync {
		return r.runAsync()
	} else {
		return r.runSync()
	}
}

func (r *runner) runAsync() (outcome types.SpecState, failure types.SpecFailure) {
	done := make(chan interface{}, 1)

	go func() {
                finished := false

		defer func() {
			if e := recover(); e != nil || !finished {
				r.failer.Panic(codelocation.New(2), e)
				select {
				case <-done:
					break
				default:
					close(done)
				}
			}
		}()

		r.asyncFunc(done)
                finished = true
	}()

	select {
	case <-done:
	case <-time.After(r.timeoutThreshold):
		r.failer.Timeout(r.codeLocation)
	}

	failure, outcome = r.failer.Drain(r.nodeType, r.componentIndex, r.codeLocation)
	return
}
func (r *runner) runSync() (outcome types.SpecState, failure types.SpecFailure) {
        finished := false

	defer func() {
		if e := recover(); e != nil || !finished {
			r.failer.Panic(codelocation.New(2), e)
		}

		failure, outcome = r.failer.Drain(r.nodeType, r.componentIndex, r.codeLocation)
	}()

	r.syncFunc()
        finished = true

	return
}
<commit_msg>Run go fmt on all files<commit_after>package leafnodes

import (
	"fmt"
	"github.com/onsi/ginkgo/internal/codelocation"
	"github.com/onsi/ginkgo/internal/failer"
	"github.com/onsi/ginkgo/types"
	"reflect"
	"time"
)

type runner struct {
	isAsync          bool
	asyncFunc        func(chan<- interface{})
	syncFunc         func()
	codeLocation     types.CodeLocation
	timeoutThreshold time.Duration
	nodeType         types.SpecComponentType
	componentIndex   int
	failer           *failer.Failer
}

func newRunner(body interface{}, codeLocation types.CodeLocation, timeout time.Duration, failer *failer.Failer, nodeType types.SpecComponentType, componentIndex int) *runner {
	bodyType := reflect.TypeOf(body)
	if bodyType.Kind() != reflect.Func {
		panic(fmt.Sprintf("Expected a function but got something else at %v", codeLocation))
	}

	runner := &runner{
		codeLocation:     codeLocation,
		timeoutThreshold: timeout,
		failer:           failer,
		nodeType:         nodeType,
		componentIndex:   componentIndex,
	}

	switch bodyType.NumIn() {
	case 0:
		runner.syncFunc = body.(func())
		return runner
	case 1:
		if !(bodyType.In(0).Kind() == reflect.Chan && bodyType.In(0).Elem().Kind() == reflect.Interface) {
			panic(fmt.Sprintf("Must pass a Done channel to function at %v", codeLocation))
		}

		wrappedBody := func(done chan<- interface{}) {
			bodyValue := reflect.ValueOf(body)
			bodyValue.Call([]reflect.Value{reflect.ValueOf(done)})
		}

		runner.isAsync = true
		runner.asyncFunc = wrappedBody
		return runner
	}

	panic(fmt.Sprintf("Too many arguments to function at %v", codeLocation))
}

func (r *runner) run() (outcome types.SpecState, failure types.SpecFailure) {
	if r.isAsync {
		return r.runAsync()
	} else {
		return r.runSync()
	}
}

func (r *runner) runAsync() (outcome types.SpecState, failure types.SpecFailure) {
	done := make(chan interface{}, 1)

	go func() {
		finished := false

		defer func() {
			if e := recover(); e != nil || !finished {
				r.failer.Panic(codelocation.New(2), e)
				select {
				case <-done:
					break
				default:
					close(done)
				}
			}
		}()

		r.asyncFunc(done)
		finished = true
	}()

	select {
	case <-done:
	case <-time.After(r.timeoutThreshold):
		r.failer.Timeout(r.codeLocation)
	}

	failure, outcome = r.failer.Drain(r.nodeType, r.componentIndex, r.codeLocation)
	return
}
func (r *runner) runSync() (outcome types.SpecState, failure types.SpecFailure) {
	finished := false

	defer func() {
		if e := recover(); e != nil || !finished {
			r.failer.Panic(codelocation.New(2), e)
		}

		failure, outcome = r.failer.Drain(r.nodeType, r.componentIndex, r.codeLocation)
	}()

	r.syncFunc()
	finished = true

	return
}
<|endoftext|>


<commit_before>// Copyright 2015 The Vanadium 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 testutil

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"time"

	"v.io/x/devtools/internal/collect"
	"v.io/x/devtools/internal/runutil"
	"v.io/x/devtools/internal/tool"
	"v.io/x/devtools/internal/util"
)

const (
	stagingBlessingsRoot    = "dev.staging.v.io" // TODO(jingjin): use a better name and update prod.go.
	localMountTable         = "/ns.dev.staging.v.io:8151"
	globalMountTable        = "/ns.dev.staging.v.io:8101"
	numRetries              = 30
	objNameForDeviceManager = "devices/vanadium-cell-master/devmgr/device"
	retryPeriod             = 10 * time.Second
)

var (
	defaultReleaseTestTimeout = time.Minute * 5

	serviceBinaries = []string{
		"binaryd",
		"applicationd",
		"proxyd",
		"identityd",
		"mounttabled",
		"roled",
		"deviced",
	}
)

// vanadiumReleaseTest updates binaries of staging cloud services and run tests for them.
func vanadiumReleaseTest(ctx *tool.Context, testName string, opts ...TestOpt) (_ *TestResult, e error) {
	root, err := util.V23Root()
	if err != nil {
		return nil, err
	}

	cleanup, err := initTest(ctx, testName, []string{})
	if err != nil {
		return nil, internalTestError{err, "Init"}
	}
	defer collect.Error(func() error { return cleanup() }, &e)

	type step struct {
		msg string
		fn  func() error
	}
	credDir := getCredDirOptValue(opts)
	steps := []step{
		step{
			msg: fmt.Sprintf("Fetching credentials from %v\n", credDir),
			fn: func() error {
				_, err := os.Stat(credDir)
				return err
			},
		},
		step{
			msg: "Build binaries\n",
			fn:  func() error { return buildBinaries(ctx, root) },
		},
		step{
			msg: "Update services\n",
			fn:  func() error { return updateServices(ctx, root, credDir) },
		},
		step{
			msg: "Check services\n",
			fn:  func() error { return checkServices(ctx) },
		},
		step{
			msg: "Create snapshot\n",
			fn:  func() error { return createSnapshot(ctx) },
		},
	}
	for _, step := range steps {
		if result, err := invoker(ctx, step.msg, step.fn); result != nil || err != nil {
			return result, err
		}
	}
	return &TestResult{Status: TestPassed}, nil
}

// invoker invokes the given function and returns TestResult and/or errors based
// on function's results.
func invoker(ctx *tool.Context, msg string, fn func() error) (*TestResult, error) {
	if err := fn(); err != nil {
		Fail(ctx, msg)
		if err == runutil.CommandTimedOutErr {
			return &TestResult{
				Status:       TestTimedOut,
				TimeoutValue: defaultReleaseTestTimeout,
			}, nil
		}
		fmt.Fprintf(ctx.Stderr(), "%s\n", err.Error())
		return nil, internalTestError{err, msg}
	}
	Pass(ctx, msg)
	return nil, nil
}

// getCredDirOptValue returns the value of credentials dir from the given
// TestOpt slice.
func getCredDirOptValue(opts []TestOpt) string {
	credDir := ""
	for _, opt := range opts {
		switch v := opt.(type) {
		case CredDirOpt:
			credDir = string(v)
		}
	}
	return credDir
}

// buildBinaries builds all vanadium binaries.
func buildBinaries(ctx *tool.Context, root string) error {
	args := []string{"go", "install", "v.io/..."}
	if err := ctx.Run().Command("v23", args...); err != nil {
		return err
	}
	return nil
}

// updateServices pushes services' binaries to the applications and binaries
// services and tells the device manager to update all its app.
func updateServices(ctx *tool.Context, root, credDir string) (e error) {
	deviceBin := filepath.Join(root, "release", "go", "bin", "device")
	credentialsArg := fmt.Sprintf("--v23.credentials=%s", credDir)
	nsArg := fmt.Sprintf("--v23.namespace.root=%s", globalMountTable)

	// Push all binaries.
	{
		fmt.Fprintln(ctx.Stdout(), "Pushing binaries...")
		args := []string{
			credentialsArg,
			nsArg,
			"publish",
			"--goos=linux",
			"--goarch=amd64",
		}
		args = append(args, serviceBinaries...)
		if err := ctx.Run().TimedCommand(defaultReleaseTestTimeout, deviceBin, args...); err != nil {
			return err
		}
	}

	// A helper function to update a single app.
	updateAppFn := func(appName string) error {
		args := []string{
			credentialsArg,
			fmt.Sprintf("--v23.namespace.root=%s", localMountTable),
			"updateall",
			appName,
		}
		if err := ctx.Run().TimedCommand(defaultReleaseTestTimeout, deviceBin, args...); err != nil {
			return err
		}
		return nil
	}

	// Update services except for device manager and mounttable.
	{
		fmt.Fprintln(ctx.Stdout(), "Updating services other than device manager and mounttable...")
		apps := []string{
			"devmgr/apps/applicationd",
			"devmgr/apps/binaryd",
			"devmgr/apps/identityd",
			"devmgr/apps/proxyd",
			"devmgr/apps/roled",
		}
		for _, app := range apps {
			if err := updateAppFn(app); err != nil {
				return err
			}
		}
	}

	// Update Device Manager.
	{
		fmt.Fprintln(ctx.Stdout(), "Updating device manager...")
		if err := updateDeviceManagerEnvelope(ctx, root, credentialsArg, nsArg); err != nil {
			return err
		}
		args := []string{
			credentialsArg,
			fmt.Sprintf("--v23.namespace.root=%s", globalMountTable),
			"update",
			objNameForDeviceManager,
		}
		if err := ctx.Run().TimedCommand(defaultReleaseTestTimeout, deviceBin, args...); err != nil {
			return err
		}
		if err := waitForMounttable(ctx, root, credentialsArg, localMountTable); err != nil {
			return err
		}
	}

	// Update mounttable last.
	{
		fmt.Fprintln(ctx.Stdout(), "Updating mounttable...")
		if err := updateAppFn("devmgr/apps/mounttabled"); err != nil {
			return err
		}
		if err := waitForMounttable(ctx, root, credentialsArg, globalMountTable); err != nil {
			return err
		}
	}
	return nil
}

// updateDeviceManagerEnvelope updates the envelope's title from "deviced" to
// "device manager".
func updateDeviceManagerEnvelope(ctx *tool.Context, root, credentialsArg, nsArg string) (e error) {
	applicationBin := filepath.Join(root, "release", "go", "bin", "application")
	appName := "applications/deviced/0"
	appProfile := "linux-amd64"
	// Get current envelope.
	args := []string{
		credentialsArg,
		nsArg,
		"match",
		appName,
		appProfile,
	}
	var out bytes.Buffer
	opts := ctx.Run().Opts()
	opts.Stdout = io.MultiWriter(opts.Stdout, &out)
	opts.Stderr = io.MultiWriter(opts.Stderr, &out)
	if err := ctx.Run().CommandWithOpts(opts, applicationBin, args...); err != nil {
		return err
	}

	// Replace title.
	strEnvelope := strings.Replace(out.String(), `"Title": "deviced"`, `"Title": "device manager"`, -1)
	tmpDir, err := ctx.Run().TempDir("", "")
	if err != nil {
		return err
	}
	defer collect.Error(func() error { return ctx.Run().RemoveAll(tmpDir) }, &e)
	filename := filepath.Join(tmpDir, "envelope")
	if err := ctx.Run().WriteFile(filename, []byte(strEnvelope), os.FileMode(0600)); err != nil {
		return err
	}

	// Update envelope.
	args = []string{
		credentialsArg,
		nsArg,
		"put",
		appName,
		appProfile,
		filename,
	}
	if err := ctx.Run().Command(applicationBin, args...); err != nil {
		return err
	}
	return nil
}

// waitForMounttable waits for the given mounttable to be up
// (timeout: 5 minutes).
func waitForMounttable(ctx *tool.Context, root, credentialsArg, mounttableRoot string) error {
	debugBin := filepath.Join(root, "release", "go", "bin", "debug")
	args := []string{
		credentialsArg,
		"glob",
		mounttableRoot + "/*",
	}
	up := false
	for i := 0; i < numRetries; i++ {
		if err := ctx.Run().Command(debugBin, args...); err != nil {
			time.Sleep(retryPeriod)
			continue
		} else {
			up = true
			break
		}
	}
	if !up {
		return fmt.Errorf("mounttable %q not up after 5 minute", mounttableRoot)
	}
	return nil
}

// checkServices runs "v23 test run vanadium-prod-services-test" against
// staging.
func checkServices(ctx *tool.Context) error {
	args := []string{
		"test",
		"run",
		fmt.Sprintf("--v23.namespace.root=%s", globalMountTable),
		fmt.Sprintf("--blessings_root=%s", stagingBlessingsRoot),
		"vanadium-prod-services-test",
	}
	if err := ctx.Run().TimedCommand(defaultReleaseTestTimeout, "v23", args...); err != nil {
		return err
	}
	return nil
}

// createSnapshot creates a snapshot with "release" label.
func createSnapshot(ctx *tool.Context) (e error) {
	args := []string{
		"snapshot",
		"--remote",
		"create",
		"--time_format=2006-01-02", // Only include date in label names
		"release",
	}
	if err := ctx.Run().Command("v23", args...); err != nil {
		return err
	}
	return nil
}
<commit_msg>TBR: internal/testutil/release: fix flag names.<commit_after>// Copyright 2015 The Vanadium 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 testutil

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"time"

	"v.io/x/devtools/internal/collect"
	"v.io/x/devtools/internal/runutil"
	"v.io/x/devtools/internal/tool"
	"v.io/x/devtools/internal/util"
)

const (
	stagingBlessingsRoot    = "dev.staging.v.io" // TODO(jingjin): use a better name and update prod.go.
	localMountTable         = "/ns.dev.staging.v.io:8151"
	globalMountTable        = "/ns.dev.staging.v.io:8101"
	numRetries              = 30
	objNameForDeviceManager = "devices/vanadium-cell-master/devmgr/device"
	retryPeriod             = 10 * time.Second
)

var (
	defaultReleaseTestTimeout = time.Minute * 5

	serviceBinaries = []string{
		"binaryd",
		"applicationd",
		"proxyd",
		"identityd",
		"mounttabled",
		"roled",
		"deviced",
	}
)

// vanadiumReleaseTest updates binaries of staging cloud services and run tests for them.
func vanadiumReleaseTest(ctx *tool.Context, testName string, opts ...TestOpt) (_ *TestResult, e error) {
	root, err := util.V23Root()
	if err != nil {
		return nil, err
	}

	cleanup, err := initTest(ctx, testName, []string{})
	if err != nil {
		return nil, internalTestError{err, "Init"}
	}
	defer collect.Error(func() error { return cleanup() }, &e)

	type step struct {
		msg string
		fn  func() error
	}
	credDir := getCredDirOptValue(opts)
	steps := []step{
		step{
			msg: fmt.Sprintf("Fetching credentials from %v\n", credDir),
			fn: func() error {
				_, err := os.Stat(credDir)
				return err
			},
		},
		step{
			msg: "Build binaries\n",
			fn:  func() error { return buildBinaries(ctx, root) },
		},
		step{
			msg: "Update services\n",
			fn:  func() error { return updateServices(ctx, root, credDir) },
		},
		step{
			msg: "Check services\n",
			fn:  func() error { return checkServices(ctx) },
		},
		step{
			msg: "Create snapshot\n",
			fn:  func() error { return createSnapshot(ctx) },
		},
	}
	for _, step := range steps {
		if result, err := invoker(ctx, step.msg, step.fn); result != nil || err != nil {
			return result, err
		}
	}
	return &TestResult{Status: TestPassed}, nil
}

// invoker invokes the given function and returns TestResult and/or errors based
// on function's results.
func invoker(ctx *tool.Context, msg string, fn func() error) (*TestResult, error) {
	if err := fn(); err != nil {
		Fail(ctx, msg)
		if err == runutil.CommandTimedOutErr {
			return &TestResult{
				Status:       TestTimedOut,
				TimeoutValue: defaultReleaseTestTimeout,
			}, nil
		}
		fmt.Fprintf(ctx.Stderr(), "%s\n", err.Error())
		return nil, internalTestError{err, msg}
	}
	Pass(ctx, msg)
	return nil, nil
}

// getCredDirOptValue returns the value of credentials dir from the given
// TestOpt slice.
func getCredDirOptValue(opts []TestOpt) string {
	credDir := ""
	for _, opt := range opts {
		switch v := opt.(type) {
		case CredDirOpt:
			credDir = string(v)
		}
	}
	return credDir
}

// buildBinaries builds all vanadium binaries.
func buildBinaries(ctx *tool.Context, root string) error {
	args := []string{"go", "install", "v.io/..."}
	if err := ctx.Run().Command("v23", args...); err != nil {
		return err
	}
	return nil
}

// updateServices pushes services' binaries to the applications and binaries
// services and tells the device manager to update all its app.
func updateServices(ctx *tool.Context, root, credDir string) (e error) {
	deviceBin := filepath.Join(root, "release", "go", "bin", "device")
	credentialsArg := fmt.Sprintf("--v23.credentials=%s", credDir)
	nsArg := fmt.Sprintf("--v23.namespace.root=%s", globalMountTable)

	// Push all binaries.
	{
		fmt.Fprintln(ctx.Stdout(), "Pushing binaries...")
		args := []string{
			credentialsArg,
			nsArg,
			"publish",
			"--goos=linux",
			"--goarch=amd64",
		}
		args = append(args, serviceBinaries...)
		if err := ctx.Run().TimedCommand(defaultReleaseTestTimeout, deviceBin, args...); err != nil {
			return err
		}
	}

	// A helper function to update a single app.
	updateAppFn := func(appName string) error {
		args := []string{
			credentialsArg,
			fmt.Sprintf("--v23.namespace.root=%s", localMountTable),
			"updateall",
			appName,
		}
		if err := ctx.Run().TimedCommand(defaultReleaseTestTimeout, deviceBin, args...); err != nil {
			return err
		}
		return nil
	}

	// Update services except for device manager and mounttable.
	{
		fmt.Fprintln(ctx.Stdout(), "Updating services other than device manager and mounttable...")
		apps := []string{
			"devmgr/apps/applicationd",
			"devmgr/apps/binaryd",
			"devmgr/apps/identityd",
			"devmgr/apps/proxyd",
			"devmgr/apps/roled",
		}
		for _, app := range apps {
			if err := updateAppFn(app); err != nil {
				return err
			}
		}
	}

	// Update Device Manager.
	{
		fmt.Fprintln(ctx.Stdout(), "Updating device manager...")
		if err := updateDeviceManagerEnvelope(ctx, root, credentialsArg, nsArg); err != nil {
			return err
		}
		args := []string{
			credentialsArg,
			fmt.Sprintf("--v23.namespace.root=%s", globalMountTable),
			"update",
			objNameForDeviceManager,
		}
		if err := ctx.Run().TimedCommand(defaultReleaseTestTimeout, deviceBin, args...); err != nil {
			return err
		}
		if err := waitForMounttable(ctx, root, credentialsArg, localMountTable); err != nil {
			return err
		}
	}

	// Update mounttable last.
	{
		fmt.Fprintln(ctx.Stdout(), "Updating mounttable...")
		if err := updateAppFn("devmgr/apps/mounttabled"); err != nil {
			return err
		}
		if err := waitForMounttable(ctx, root, credentialsArg, globalMountTable); err != nil {
			return err
		}
	}
	return nil
}

// updateDeviceManagerEnvelope updates the envelope's title from "deviced" to
// "device manager".
func updateDeviceManagerEnvelope(ctx *tool.Context, root, credentialsArg, nsArg string) (e error) {
	applicationBin := filepath.Join(root, "release", "go", "bin", "application")
	appName := "applications/deviced/0"
	appProfile := "linux-amd64"
	// Get current envelope.
	args := []string{
		credentialsArg,
		nsArg,
		"match",
		appName,
		appProfile,
	}
	var out bytes.Buffer
	opts := ctx.Run().Opts()
	opts.Stdout = io.MultiWriter(opts.Stdout, &out)
	opts.Stderr = io.MultiWriter(opts.Stderr, &out)
	if err := ctx.Run().CommandWithOpts(opts, applicationBin, args...); err != nil {
		return err
	}

	// Replace title.
	strEnvelope := strings.Replace(out.String(), `"Title": "deviced"`, `"Title": "device manager"`, -1)
	tmpDir, err := ctx.Run().TempDir("", "")
	if err != nil {
		return err
	}
	defer collect.Error(func() error { return ctx.Run().RemoveAll(tmpDir) }, &e)
	filename := filepath.Join(tmpDir, "envelope")
	if err := ctx.Run().WriteFile(filename, []byte(strEnvelope), os.FileMode(0600)); err != nil {
		return err
	}

	// Update envelope.
	args = []string{
		credentialsArg,
		nsArg,
		"put",
		appName,
		appProfile,
		filename,
	}
	if err := ctx.Run().Command(applicationBin, args...); err != nil {
		return err
	}
	return nil
}

// waitForMounttable waits for the given mounttable to be up
// (timeout: 5 minutes).
func waitForMounttable(ctx *tool.Context, root, credentialsArg, mounttableRoot string) error {
	debugBin := filepath.Join(root, "release", "go", "bin", "debug")
	args := []string{
		credentialsArg,
		"glob",
		mounttableRoot + "/*",
	}
	up := false
	for i := 0; i < numRetries; i++ {
		if err := ctx.Run().Command(debugBin, args...); err != nil {
			time.Sleep(retryPeriod)
			continue
		} else {
			up = true
			break
		}
	}
	if !up {
		return fmt.Errorf("mounttable %q not up after 5 minute", mounttableRoot)
	}
	return nil
}

// checkServices runs "v23 test run vanadium-prod-services-test" against
// staging.
func checkServices(ctx *tool.Context) error {
	args := []string{
		"test",
		"run",
		fmt.Sprintf("--v23.namespace.root=%s", globalMountTable),
		fmt.Sprintf("--blessings-root=%s", stagingBlessingsRoot),
		"vanadium-prod-services-test",
	}
	if err := ctx.Run().TimedCommand(defaultReleaseTestTimeout, "v23", args...); err != nil {
		return err
	}
	return nil
}

// createSnapshot creates a snapshot with "release" label.
func createSnapshot(ctx *tool.Context) (e error) {
	args := []string{
		"snapshot",
		"--remote",
		"create",
		"--time-format=2006-01-02", // Only include date in label names
		"release",
	}
	if err := ctx.Run().Command("v23", args...); err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2003-2005 Maxim Sobolev. All rights reserved.
// Copyright (c) 2006-2015 Sippy Software, Inc. All rights reserved.
// Copyright (c) 2015 Andrii Pylypenko. All rights reserved.
//
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without modification,
// are permitted provided that the following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice, this
// list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright notice,
// this list of conditions and the following disclaimer in the documentation and/or
// other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
// ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
// WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
// ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
// (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
// LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
// ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
package sippy_header

import (
    "errors"
    "fmt"
    "strconv"
    "strings"
    "unicode"

    "sippy/conf"
    "sippy/utils"
)

type SipAddress struct {
    params      map[string]*string
    url         *SipURL
    hadbrace    bool
    name        string
    q           float64
}

func ParseSipAddress(address string, relaxedparser bool, config sippy_conf.Config) (*SipAddress, error) {
    var err error
    var arr []string

    // simple 'sip:foo' case
    self := &SipAddress{
        params : make(map[string]*string),
        hadbrace : true,
        q : 1.0,
    }

    if strings.HasPrefix(strings.ToLower(address), "sip:") && strings.Index(address, "<") == -1 {
        parts := strings.SplitN(address, ";", 2)
        self.url, err = ParseSipURL(parts[0], relaxedparser, config)
        if err != nil {
            return nil, err
        }
        if len(parts) == 2 {
            if err = self._parse_paramstring(parts[1]); err != nil {
                return nil, err
            }
        }
        self.hadbrace = false
        return self, nil
    }
    var url *string = nil
    if address[0] == '"' {
        equote := strings.Index(address[1:], "\"") + 1
        if equote != 0 {
            sbrace := strings.Index(address[equote:], "<")
            if sbrace != -1 {
                self.hadbrace = true
                self.name = strings.TrimSpace(address[1:equote])
                tmp := address[equote + sbrace + 1:]
                url = &tmp
            }
        }
    }
    if url == nil {
        arr = strings.SplitN(address, "<", 2)
        if len(arr) != 2 {
            return nil, errors.New("ParseSipAddress #1")
        }
        self.name = strings.TrimSpace(arr[0])
        url = &arr[1]
        if len(self.name) > 0 && self.name[0] == '"' {
            self.name = self.name[1:]
        }
        if len(self.name) > 0 && self.name[len(self.name)-1] == '"' {
            self.name = self.name[:len(self.name)-1]
        }
    }
    arr = strings.SplitN(*url, ">", 2)
    if len(arr) != 2 {
        return nil, errors.New("ParseSipAddress #2")
    }
    paramstring := arr[1]
    if self.url, err = ParseSipURL(arr[0], relaxedparser, config); err != nil {
        return nil, err
    }
    paramstring = strings.TrimSpace(paramstring)
    if err = self._parse_paramstring(paramstring); err != nil {
        return nil, err
    }
    return self, nil
}

func (self *SipAddress) _parse_paramstring(s string) error {
    for _, l := range strings.Split(s, ";") {
        var v *string

        if l == "" {
            continue
        }
        arr := strings.SplitN(l, "=", 2)
        k := arr[0]
        if len(arr) == 2 {
            tmp := arr[1]
            v = &tmp
        } else {
            v = nil
        }
        if _, ok := self.params[k]; ok {
            return errors.New("Duplicate parameter in SIP address: " + k)
        }
        if k == "q" {
            if v != nil {
                // ignore absense or possible errors in the q= value
                if q, err := strconv.ParseFloat(*v, 64); err == nil {
                    self.q = q
                }
            }
        } else {
            self.params[k] = v
        }
    }
    return nil
}

func (self *SipAddress) String() string {
    return self.LocalStr(nil)
}

func (self *SipAddress) LocalStr(hostport *sippy_conf.HostPort) string {
    var od, cd, s string
    if self.hadbrace {
        od = "<"
        cd = ">"
    }
    if len(self.name) > 0 {
        needs_quote := false
        for _, r := range self.name {
            if unicode.IsLetter(r) || unicode.IsNumber(r) || strings.ContainsRune("-.!%*_+`'~", r) {
                continue
            }
            needs_quote = true
            break
        }
        if needs_quote {
            s += "\"" + self.name + "\" "
        } else {
            s += self.name + " "
        }
        od = "<"
        cd = ">"
    }
    s += od + self.url.LocalStr(hostport) + cd
    for k, v := range self.params {
        if v == nil {
            s += ";" + k
        } else {
            s += ";" + k + "=" + *v
        }
    }
    if self.q != 1.0 {
        s += fmt.Sprintf(";q=%g", self.q)
    }
    return s
}

func NewSipAddress(name string, url *SipURL) *SipAddress {
    return &SipAddress{
        name : name,
        url : url,
        hadbrace : true,
        params : make(map[string]*string),
    }
}

func (self *SipAddress) GetCopy() *SipAddress {
    ret := *self
    ret.params = make(map[string]*string)
    for k, v := range self.params {
        if v == nil {
            ret.params[k] = nil
        } else {
            s := *v
            ret.params[k] = &s
        }
    }
    ret.url = self.url.GetCopy()
    return &ret
}

func (self *SipAddress) GetParam(name string) string {
    ret, ok := self.params[name]
    if !ok || ret == nil {
        return ""
    }
    return *ret
}

func (self *SipAddress) SetParam(name, value string) {
    self.params[name] = &value
}

func (self *SipAddress) GetName() string {
    return self.name
}

func (self *SipAddress) GetUrl() *SipURL {
    return self.url
}

func (self *SipAddress) GetTag() string {
    return self.GetParam("tag")
}

func (self *SipAddress) SetTag(tag string) {
    self.SetParam("tag", tag)
}

func (self *SipAddress) GenTag() {
    self.SetParam("tag", sippy_utils.GenTag())
}

func (self *SipAddress) GetQ() float64 {
    return self.q
}
<commit_msg>The default q= value is 1.0.<commit_after>// Copyright (c) 2003-2005 Maxim Sobolev. All rights reserved.
// Copyright (c) 2006-2015 Sippy Software, Inc. All rights reserved.
// Copyright (c) 2015 Andrii Pylypenko. All rights reserved.
//
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without modification,
// are permitted provided that the following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice, this
// list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright notice,
// this list of conditions and the following disclaimer in the documentation and/or
// other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
// ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
// WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
// ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
// (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
// LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
// ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
package sippy_header

import (
    "errors"
    "fmt"
    "strconv"
    "strings"
    "unicode"

    "sippy/conf"
    "sippy/utils"
)

type SipAddress struct {
    params      map[string]*string
    url         *SipURL
    hadbrace    bool
    name        string
    q           float64
}

func ParseSipAddress(address string, relaxedparser bool, config sippy_conf.Config) (*SipAddress, error) {
    var err error
    var arr []string

    // simple 'sip:foo' case
    self := &SipAddress{
        params : make(map[string]*string),
        hadbrace : true,
        q : 1.0,
    }

    if strings.HasPrefix(strings.ToLower(address), "sip:") && strings.Index(address, "<") == -1 {
        parts := strings.SplitN(address, ";", 2)
        self.url, err = ParseSipURL(parts[0], relaxedparser, config)
        if err != nil {
            return nil, err
        }
        if len(parts) == 2 {
            if err = self._parse_paramstring(parts[1]); err != nil {
                return nil, err
            }
        }
        self.hadbrace = false
        return self, nil
    }
    var url *string = nil
    if address[0] == '"' {
        equote := strings.Index(address[1:], "\"") + 1
        if equote != 0 {
            sbrace := strings.Index(address[equote:], "<")
            if sbrace != -1 {
                self.hadbrace = true
                self.name = strings.TrimSpace(address[1:equote])
                tmp := address[equote + sbrace + 1:]
                url = &tmp
            }
        }
    }
    if url == nil {
        arr = strings.SplitN(address, "<", 2)
        if len(arr) != 2 {
            return nil, errors.New("ParseSipAddress #1")
        }
        self.name = strings.TrimSpace(arr[0])
        url = &arr[1]
        if len(self.name) > 0 && self.name[0] == '"' {
            self.name = self.name[1:]
        }
        if len(self.name) > 0 && self.name[len(self.name)-1] == '"' {
            self.name = self.name[:len(self.name)-1]
        }
    }
    arr = strings.SplitN(*url, ">", 2)
    if len(arr) != 2 {
        return nil, errors.New("ParseSipAddress #2")
    }
    paramstring := arr[1]
    if self.url, err = ParseSipURL(arr[0], relaxedparser, config); err != nil {
        return nil, err
    }
    paramstring = strings.TrimSpace(paramstring)
    if err = self._parse_paramstring(paramstring); err != nil {
        return nil, err
    }
    return self, nil
}

func (self *SipAddress) _parse_paramstring(s string) error {
    for _, l := range strings.Split(s, ";") {
        var v *string

        if l == "" {
            continue
        }
        arr := strings.SplitN(l, "=", 2)
        k := arr[0]
        if len(arr) == 2 {
            tmp := arr[1]
            v = &tmp
        } else {
            v = nil
        }
        if _, ok := self.params[k]; ok {
            return errors.New("Duplicate parameter in SIP address: " + k)
        }
        if k == "q" {
            if v != nil {
                // ignore absense or possible errors in the q= value
                if q, err := strconv.ParseFloat(*v, 64); err == nil {
                    self.q = q
                }
            }
        } else {
            self.params[k] = v
        }
    }
    return nil
}

func (self *SipAddress) String() string {
    return self.LocalStr(nil)
}

func (self *SipAddress) LocalStr(hostport *sippy_conf.HostPort) string {
    var od, cd, s string
    if self.hadbrace {
        od = "<"
        cd = ">"
    }
    if len(self.name) > 0 {
        needs_quote := false
        for _, r := range self.name {
            if unicode.IsLetter(r) || unicode.IsNumber(r) || strings.ContainsRune("-.!%*_+`'~", r) {
                continue
            }
            needs_quote = true
            break
        }
        if needs_quote {
            s += "\"" + self.name + "\" "
        } else {
            s += self.name + " "
        }
        od = "<"
        cd = ">"
    }
    s += od + self.url.LocalStr(hostport) + cd
    for k, v := range self.params {
        if v == nil {
            s += ";" + k
        } else {
            s += ";" + k + "=" + *v
        }
    }
    if self.q != 1.0 {
        s += fmt.Sprintf(";q=%g", self.q)
    }
    return s
}

func NewSipAddress(name string, url *SipURL) *SipAddress {
    return &SipAddress{
        name        : name,
        url         : url,
        hadbrace    : true,
        params      : make(map[string]*string),
        q           : 1.0,
    }
}

func (self *SipAddress) GetCopy() *SipAddress {
    ret := *self
    ret.params = make(map[string]*string)
    for k, v := range self.params {
        if v == nil {
            ret.params[k] = nil
        } else {
            s := *v
            ret.params[k] = &s
        }
    }
    ret.url = self.url.GetCopy()
    return &ret
}

func (self *SipAddress) GetParam(name string) string {
    ret, ok := self.params[name]
    if !ok || ret == nil {
        return ""
    }
    return *ret
}

func (self *SipAddress) SetParam(name, value string) {
    self.params[name] = &value
}

func (self *SipAddress) GetName() string {
    return self.name
}

func (self *SipAddress) GetUrl() *SipURL {
    return self.url
}

func (self *SipAddress) GetTag() string {
    return self.GetParam("tag")
}

func (self *SipAddress) SetTag(tag string) {
    self.SetParam("tag", tag)
}

func (self *SipAddress) GenTag() {
    self.SetParam("tag", sippy_utils.GenTag())
}

func (self *SipAddress) GetQ() float64 {
    return self.q
}
<|endoftext|>


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

package netutil

import (
	"errors"
	"github.com/someonegg/goutility/chanutil"
	"golang.org/x/net/context"
	"net"
	"net/http"
	"sync"
	"time"
)

var (
	ErrUnknownPanic = errors.New("unknown panic")
)

type ContextHandler interface {
	ContextServeHTTP(context.Context, http.ResponseWriter, *http.Request)
}

// HttpService is a wrapper of http.Server.
type HttpService struct {
	err     error
	quitCtx context.Context
	quitF   context.CancelFunc
	stopD   chanutil.DoneChan

	l   *net.TCPListener
	h   ContextHandler
	srv *http.Server

	reqWG sync.WaitGroup
}

// NewHttpService is a short cut to use NewHttpServiceEx.
func NewHttpService(l *net.TCPListener, h http.Handler,
	maxConcurrent int) *HttpService {

	return NewHttpServiceEx(l, NewMaxConcurrentHandler(NewHttpHandler(h),
		maxConcurrent, DefaultHesitateTime, DefaultMaxConcurrentNotifier))
}

func NewHttpServiceEx(l *net.TCPListener, h ContextHandler) *HttpService {
	s := &HttpService{}

	s.quitCtx, s.quitF = context.WithCancel(context.Background())
	s.stopD = chanutil.NewDoneChan()
	s.l = l
	s.h = h
	s.srv = &http.Server{
		Addr:           s.l.Addr().String(),
		Handler:        s,
		ReadTimeout:    10 * time.Second,
		WriteTimeout:   10 * time.Second,
		MaxHeaderBytes: 1 << 20,
	}

	return s
}

func (s *HttpService) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	s.reqWG.Add(1)
	defer s.reqWG.Done()
	s.h.ContextServeHTTP(s.quitCtx, w, r)
}

func (s *HttpService) Start() {
	go s.serve()
}

func (s *HttpService) serve() {
	defer s.ending()

	s.err = s.srv.Serve(TcpKeepAliveListener{s.l})
}

func (s *HttpService) ending() {
	if e := recover(); e != nil {
		switch v := e.(type) {
		case error:
			s.err = v
		default:
			s.err = ErrUnknownPanic
		}
	}

	s.quitF()
	s.stopD.SetDone()
}

func (s *HttpService) Err() error {
	return s.err
}

func (s *HttpService) Stop() {
	s.srv.SetKeepAlivesEnabled(false)
	s.quitF()
	s.l.Close()
}

func (s *HttpService) StopD() chanutil.DoneChanR {
	return s.stopD.R()
}

func (s *HttpService) Stopped() bool {
	return s.stopD.R().Done()
}

func (s *HttpService) WaitRequests() {
	s.reqWG.Wait()
}

func (s *HttpService) QuitCtx() context.Context {
	return s.quitCtx
}

type httpHandler struct {
	oh http.Handler
}

// The http handler type is an adapter to allow the use of
// ordinary http.Handler as ContextHandler.
func NewHttpHandler(oh http.Handler) ContextHandler {
	return httpHandler{oh}
}

func (h httpHandler) ContextServeHTTP(ctx context.Context,
	w http.ResponseWriter, r *http.Request) {

	h.oh.ServeHTTP(w, r)
}

type maxConcurrentHandler struct {
	oh           ContextHandler
	concur       chanutil.Semaphore
	hesitateTime time.Duration
	notifier     MaxConcurrentNotifier
}

// The maxconcurrent handler type is a middleware that can limit the
// maximum number of concurrent access.
// if maxConcurrent == 0, no limit on concurrency.
// if hesitateTime == 0, use DefaultHesitateTime.
func NewMaxConcurrentHandler(oh ContextHandler,
	maxConcurrent int, hesitateTime time.Duration,
	notifier MaxConcurrentNotifier) ContextHandler {

	if maxConcurrent <= 0 {
		return oh
	}
	if hesitateTime <= 0 {
		hesitateTime = DefaultHesitateTime
	}

	return &maxConcurrentHandler{
		oh:           oh,
		concur:       chanutil.NewSemaphore(maxConcurrent),
		hesitateTime: hesitateTime,
		notifier:     notifier,
	}
}

const (
	acquire_OK      int = 0
	acquire_CtxDone int = 1
	acquire_Timeout int = 2
)

func (h *maxConcurrentHandler) acquireConn(ctx context.Context) int {
	select {
	case <-ctx.Done():
		return acquire_CtxDone
	// Acquire
	case h.concur <- struct{}{}:
		return acquire_OK
	case <-time.After(h.hesitateTime):
		return acquire_Timeout
	}
}

func (h *maxConcurrentHandler) releaseConn() {
	<-h.concur
}

func (h *maxConcurrentHandler) ContextServeHTTP(ctx context.Context,
	w http.ResponseWriter, r *http.Request) {

	ret := h.acquireConn(ctx)
	switch ret {
	case acquire_CtxDone:
		h.notifier.OnContextDone(w, r)
		return
	case acquire_Timeout:
		h.notifier.OnConcurrentLimit(w, r)
		return
	}
	defer h.releaseConn()

	h.oh.ContextServeHTTP(ctx, w, r)
}

const DefaultHesitateTime = 50 * time.Millisecond

type MaxConcurrentNotifier interface {
	OnContextDone(w http.ResponseWriter, r *http.Request)
	OnConcurrentLimit(w http.ResponseWriter, r *http.Request)
}

type defaultMaxConcurrentNotifier struct{}

func (n defaultMaxConcurrentNotifier) OnContextDone(
	w http.ResponseWriter, r *http.Request) {

	http.Error(w, "Service Unavailable", http.StatusServiceUnavailable)
}

func (n defaultMaxConcurrentNotifier) OnConcurrentLimit(
	w http.ResponseWriter, r *http.Request) {

	http.Error(w, "Service Busy", http.StatusTooManyRequests)
}

var DefaultMaxConcurrentNotifier MaxConcurrentNotifier = defaultMaxConcurrentNotifier{}
<commit_msg>correct http status code ex<commit_after>// Copyright 2015 someonegg. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package netutil

import (
	"errors"
	"github.com/someonegg/goutility/chanutil"
	"golang.org/x/net/context"
	"net"
	"net/http"
	"sync"
	"time"
)

var (
	ErrUnknownPanic = errors.New("unknown panic")
)

type ContextHandler interface {
	ContextServeHTTP(context.Context, http.ResponseWriter, *http.Request)
}

// HttpService is a wrapper of http.Server.
type HttpService struct {
	err     error
	quitCtx context.Context
	quitF   context.CancelFunc
	stopD   chanutil.DoneChan

	l   *net.TCPListener
	h   ContextHandler
	srv *http.Server

	reqWG sync.WaitGroup
}

// NewHttpService is a short cut to use NewHttpServiceEx.
func NewHttpService(l *net.TCPListener, h http.Handler,
	maxConcurrent int) *HttpService {

	return NewHttpServiceEx(l, NewMaxConcurrentHandler(NewHttpHandler(h),
		maxConcurrent, DefaultHesitateTime, DefaultMaxConcurrentNotifier))
}

func NewHttpServiceEx(l *net.TCPListener, h ContextHandler) *HttpService {
	s := &HttpService{}

	s.quitCtx, s.quitF = context.WithCancel(context.Background())
	s.stopD = chanutil.NewDoneChan()
	s.l = l
	s.h = h
	s.srv = &http.Server{
		Addr:           s.l.Addr().String(),
		Handler:        s,
		ReadTimeout:    10 * time.Second,
		WriteTimeout:   10 * time.Second,
		MaxHeaderBytes: 1 << 20,
	}

	return s
}

func (s *HttpService) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	s.reqWG.Add(1)
	defer s.reqWG.Done()
	s.h.ContextServeHTTP(s.quitCtx, w, r)
}

func (s *HttpService) Start() {
	go s.serve()
}

func (s *HttpService) serve() {
	defer s.ending()

	s.err = s.srv.Serve(TcpKeepAliveListener{s.l})
}

func (s *HttpService) ending() {
	if e := recover(); e != nil {
		switch v := e.(type) {
		case error:
			s.err = v
		default:
			s.err = ErrUnknownPanic
		}
	}

	s.quitF()
	s.stopD.SetDone()
}

func (s *HttpService) Err() error {
	return s.err
}

func (s *HttpService) Stop() {
	s.srv.SetKeepAlivesEnabled(false)
	s.quitF()
	s.l.Close()
}

func (s *HttpService) StopD() chanutil.DoneChanR {
	return s.stopD.R()
}

func (s *HttpService) Stopped() bool {
	return s.stopD.R().Done()
}

func (s *HttpService) WaitRequests() {
	s.reqWG.Wait()
}

func (s *HttpService) QuitCtx() context.Context {
	return s.quitCtx
}

type httpHandler struct {
	oh http.Handler
}

// The http handler type is an adapter to allow the use of
// ordinary http.Handler as ContextHandler.
func NewHttpHandler(oh http.Handler) ContextHandler {
	return httpHandler{oh}
}

func (h httpHandler) ContextServeHTTP(ctx context.Context,
	w http.ResponseWriter, r *http.Request) {

	h.oh.ServeHTTP(w, r)
}

type maxConcurrentHandler struct {
	oh           ContextHandler
	concur       chanutil.Semaphore
	hesitateTime time.Duration
	notifier     MaxConcurrentNotifier
}

// The maxconcurrent handler type is a middleware that can limit the
// maximum number of concurrent access.
// if maxConcurrent == 0, no limit on concurrency.
// if hesitateTime == 0, use DefaultHesitateTime.
func NewMaxConcurrentHandler(oh ContextHandler,
	maxConcurrent int, hesitateTime time.Duration,
	notifier MaxConcurrentNotifier) ContextHandler {

	if maxConcurrent <= 0 {
		return oh
	}
	if hesitateTime <= 0 {
		hesitateTime = DefaultHesitateTime
	}

	return &maxConcurrentHandler{
		oh:           oh,
		concur:       chanutil.NewSemaphore(maxConcurrent),
		hesitateTime: hesitateTime,
		notifier:     notifier,
	}
}

const (
	acquire_OK      int = 0
	acquire_CtxDone int = 1
	acquire_Timeout int = 2
)

func (h *maxConcurrentHandler) acquireConn(ctx context.Context) int {
	select {
	case <-ctx.Done():
		return acquire_CtxDone
	// Acquire
	case h.concur <- struct{}{}:
		return acquire_OK
	case <-time.After(h.hesitateTime):
		return acquire_Timeout
	}
}

func (h *maxConcurrentHandler) releaseConn() {
	<-h.concur
}

func (h *maxConcurrentHandler) ContextServeHTTP(ctx context.Context,
	w http.ResponseWriter, r *http.Request) {

	ret := h.acquireConn(ctx)
	switch ret {
	case acquire_CtxDone:
		h.notifier.OnContextDone(w, r)
		return
	case acquire_Timeout:
		h.notifier.OnConcurrentLimit(w, r)
		return
	}
	defer h.releaseConn()

	h.oh.ContextServeHTTP(ctx, w, r)
}

const DefaultHesitateTime = 50 * time.Millisecond

type MaxConcurrentNotifier interface {
	OnContextDone(w http.ResponseWriter, r *http.Request)
	OnConcurrentLimit(w http.ResponseWriter, r *http.Request)
}

type defaultMaxConcurrentNotifier struct{}

func (n defaultMaxConcurrentNotifier) OnContextDone(
	w http.ResponseWriter, r *http.Request) {

	http.Error(w, "Service Maintenance", http.StatusServiceUnavailable)
}

func (n defaultMaxConcurrentNotifier) OnConcurrentLimit(
	w http.ResponseWriter, r *http.Request) {

	http.Error(w, "Service Busy", http.StatusServiceUnavailable)
}

var DefaultMaxConcurrentNotifier MaxConcurrentNotifier = defaultMaxConcurrentNotifier{}
<|endoftext|>


<commit_before>// (c) 2021, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package network

import (
	"bytes"
	"compress/gzip"
	"io"
)

var gzipHeaderBytes = []byte{0x1f, 0x8b, 8, 0, 0, 0, 0, 0, 4, 0}

const compressionThresholdBytes = 128

type Compressor interface {
	Compress([]byte) ([]byte, error)
	Decompress([]byte) ([]byte, error)
	IsCompressed([]byte) bool
	IsCompressable([]byte) bool
}

// gzipCompressor implements Compressor
type gzipCompressor struct {
	initialized bool
	bytesReader *bytes.Reader
	gzipReader  *gzip.Reader
	gzipWriter  *gzip.Writer
	writeBuffer *bytes.Buffer
}

// Compress [msg] and returns the compressed bytes.
func (g *gzipCompressor) Compress(msg []byte) ([]byte, error) {
	g.reset()
	if _, err := g.gzipWriter.Write(msg); err != nil {
		return msg, err
	}
	if err := g.gzipWriter.Flush(); err != nil {
		return msg, err
	}
	if err := g.gzipWriter.Close(); err != nil {
		return msg, err
	}
	return g.writeBuffer.Bytes(), nil
}

// Decompress decompresses [msg].
func (g *gzipCompressor) Decompress(msg []byte) ([]byte, error) {
	g.reset()
	// todo move following into reset
	if g.gzipReader == nil {
		g.bytesReader = bytes.NewReader(msg)
		gzipReader, err := gzip.NewReader(g.bytesReader)
		if err != nil {
			return nil, err
		}
		g.gzipReader = gzipReader
	} else {
		g.bytesReader.Reset(msg)
		if err := g.gzipReader.Reset(g.bytesReader); err != nil && err != io.EOF {
			return nil, err
		}
	}
	return io.ReadAll(g.gzipReader)
}

// IsCompressed returns whether given bytes are compressed data or not
func (g *gzipCompressor) IsCompressed(msg []byte) bool {
	return len(msg) > 10 && bytes.Equal(msg[:10], gzipHeaderBytes)
}

func (g *gzipCompressor) IsCompressable(msg []byte) bool {
	return len(msg) > compressionThresholdBytes
}

func (g *gzipCompressor) reset() {
	if !g.initialized {
		var buf bytes.Buffer
		g.writeBuffer = &buf
		g.gzipWriter = gzip.NewWriter(g.writeBuffer)
		g.initialized = true
	}
	g.writeBuffer.Reset()
	g.gzipWriter.Reset(g.writeBuffer)
}

// NewCompressor returns a new compressor instance
func NewCompressor() Compressor {
	return &gzipCompressor{}
}
<commit_msg>Update gzip header bytes<commit_after>// (c) 2021, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package network

import (
	"bytes"
	"compress/gzip"
	"io"
)

var gzipHeaderBytes = []byte{0x1f, 0x8b, 8, 0, 0, 0, 0, 0, 0, 255}

const compressionThresholdBytes = 128

type Compressor interface {
	Compress([]byte) ([]byte, error)
	Decompress([]byte) ([]byte, error)
	IsCompressed([]byte) bool
	IsCompressable([]byte) bool
}

// gzipCompressor implements Compressor
type gzipCompressor struct {
	initialized bool
	bytesReader *bytes.Reader
	gzipReader  *gzip.Reader
	gzipWriter  *gzip.Writer
	writeBuffer *bytes.Buffer
}

// Compress [msg] and returns the compressed bytes.
func (g *gzipCompressor) Compress(msg []byte) ([]byte, error) {
	g.reset()
	if _, err := g.gzipWriter.Write(msg); err != nil {
		return msg, err
	}
	if err := g.gzipWriter.Flush(); err != nil {
		return msg, err
	}
	if err := g.gzipWriter.Close(); err != nil {
		return msg, err
	}
	return g.writeBuffer.Bytes(), nil
}

// Decompress decompresses [msg].
func (g *gzipCompressor) Decompress(msg []byte) ([]byte, error) {
	g.reset()
	// todo move following into reset
	if g.gzipReader == nil {
		g.bytesReader = bytes.NewReader(msg)
		gzipReader, err := gzip.NewReader(g.bytesReader)
		if err != nil {
			return nil, err
		}
		g.gzipReader = gzipReader
	} else {
		g.bytesReader.Reset(msg)
		if err := g.gzipReader.Reset(g.bytesReader); err != nil && err != io.EOF {
			return nil, err
		}
	}
	return io.ReadAll(g.gzipReader)
}

// IsCompressed returns whether given bytes are compressed data or not
func (g *gzipCompressor) IsCompressed(msg []byte) bool {
	return len(msg) > 10 && bytes.Equal(msg[:10], gzipHeaderBytes)
}

func (g *gzipCompressor) IsCompressable(msg []byte) bool {
	return len(msg) > compressionThresholdBytes
}

func (g *gzipCompressor) reset() {
	if !g.initialized {
		var buf bytes.Buffer
		g.writeBuffer = &buf
		g.gzipWriter = gzip.NewWriter(g.writeBuffer)
		g.initialized = true
	}
	g.writeBuffer.Reset()
	g.gzipWriter.Reset(g.writeBuffer)
}

// NewCompressor returns a new compressor instance
func NewCompressor() Compressor {
	return &gzipCompressor{}
}
<|endoftext|>