<commit_before>/*
Copyright 2016 The Kubernetes Authors All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package util

import (
	"net"

	"github.com/pkg/errors"
)

// These constants are used by both minikube and localkube
const (
	APIServerPort             = 8443
	DefaultLocalkubeDirectory = "/var/lib/localkube"
	DefaultCertPath           = DefaultLocalkubeDirectory + "/certs/"
	DefaultKubeConfigPath     = DefaultLocalkubeDirectory + "/kubeconfig"
	DefaultDNSDomain          = "cluster.local"
	DefaultServiceCIDR        = "10.96.0.0/12"
)

var DefaultAdmissionControllers = []string{
	"Initializers",
	"NamespaceLifecycle",
	"LimitRanger",
	"ServiceAccount",
	"DefaultStorageClass",
	"DefaultTolerationSeconds",
	"NodeRestriction",
	"ResourceQuota",
	"MutatingAdmissionWebhook",
}

// GetServiceClusterIP returns the first IP of the ServiceCIDR
func GetServiceClusterIP(serviceCIDR string) (net.IP, error) {
	ip, _, err := net.ParseCIDR(serviceCIDR)
	if err != nil {
		return nil, errors.Wrap(err, "parsing default service cidr")
	}
	ip = ip.To4()
	ip[3]++
	return ip, nil
}

// GetDNSIP returns x.x.x.10 of the service CIDR
func GetDNSIP(serviceCIDR string) (net.IP, error) {
	ip, _, err := net.ParseCIDR(serviceCIDR)
	if err != nil {
		return nil, errors.Wrap(err, "parsing default service cidr")
	}
	ip = ip.To4()
	ip[3] = 10
	return ip, nil
}

func GetAlternateDNS(domain string) []string {
	return []string{"kubernetes.default.svc." + domain, "kubernetes.default.svc", "kubernetes.default", "kubernetes", "localhost"}
}
<commit_msg>Add the ValidatingAdmissionWebhook controller.<commit_after>/*
Copyright 2016 The Kubernetes Authors All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package util

import (
	"net"

	"github.com/pkg/errors"
)

// These constants are used by both minikube and localkube
const (
	APIServerPort             = 8443
	DefaultLocalkubeDirectory = "/var/lib/localkube"
	DefaultCertPath           = DefaultLocalkubeDirectory + "/certs/"
	DefaultKubeConfigPath     = DefaultLocalkubeDirectory + "/kubeconfig"
	DefaultDNSDomain          = "cluster.local"
	DefaultServiceCIDR        = "10.96.0.0/12"
)

var DefaultAdmissionControllers = []string{
	"Initializers",
	"NamespaceLifecycle",
	"LimitRanger",
	"ServiceAccount",
	"DefaultStorageClass",
	"DefaultTolerationSeconds",
	"NodeRestriction",
	"ResourceQuota",
	"MutatingAdmissionWebhook",
	"ValidatingAdmissionWebhook",
}

// GetServiceClusterIP returns the first IP of the ServiceCIDR
func GetServiceClusterIP(serviceCIDR string) (net.IP, error) {
	ip, _, err := net.ParseCIDR(serviceCIDR)
	if err != nil {
		return nil, errors.Wrap(err, "parsing default service cidr")
	}
	ip = ip.To4()
	ip[3]++
	return ip, nil
}

// GetDNSIP returns x.x.x.10 of the service CIDR
func GetDNSIP(serviceCIDR string) (net.IP, error) {
	ip, _, err := net.ParseCIDR(serviceCIDR)
	if err != nil {
		return nil, errors.Wrap(err, "parsing default service cidr")
	}
	ip = ip.To4()
	ip[3] = 10
	return ip, nil
}

func GetAlternateDNS(domain string) []string {
	return []string{"kubernetes.default.svc." + domain, "kubernetes.default.svc", "kubernetes.default", "kubernetes", "localhost"}
}
<|endoftext|>


<commit_before>/*
Copyright 2014 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package diff

import (
	"bytes"
	"encoding/json"
	"fmt"
	"reflect"
	"sort"
	"strings"
	"text/tabwriter"

	"github.com/davecgh/go-spew/spew"

	"k8s.io/kubernetes/pkg/util/validation/field"
)

// StringDiff diffs a and b and returns a human readable diff.
func StringDiff(a, b string) string {
	ba := []byte(a)
	bb := []byte(b)
	out := []byte{}
	i := 0
	for ; i < len(ba) && i < len(bb); i++ {
		if ba[i] != bb[i] {
			break
		}
		out = append(out, ba[i])
	}
	out = append(out, []byte("\n\nA: ")...)
	out = append(out, ba[i:]...)
	out = append(out, []byte("\n\nB: ")...)
	out = append(out, bb[i:]...)
	out = append(out, []byte("\n\n")...)
	return string(out)
}

// ObjectDiff writes the two objects out as JSON and prints out the identical part of
// the objects followed by the remaining part of 'a' and finally the remaining part of 'b'.
// For debugging tests.
func ObjectDiff(a, b interface{}) string {
	ab, err := json.Marshal(a)
	if err != nil {
		panic(fmt.Sprintf("a: %v", err))
	}
	bb, err := json.Marshal(b)
	if err != nil {
		panic(fmt.Sprintf("b: %v", err))
	}
	return StringDiff(string(ab), string(bb))
}

// ObjectGoPrintDiff is like ObjectDiff, but uses go-spew to print the objects,
// which shows absolutely everything by recursing into every single pointer
// (go's %#v formatters OTOH stop at a certain point). This is needed when you
// can't figure out why reflect.DeepEqual is returning false and nothing is
// showing you differences. This will.
func ObjectGoPrintDiff(a, b interface{}) string {
	s := spew.ConfigState{DisableMethods: true}
	return StringDiff(
		s.Sprintf("%#v", a),
		s.Sprintf("%#v", b),
	)
}

func ObjectReflectDiff(a, b interface{}) string {
	vA, vB := reflect.ValueOf(a), reflect.ValueOf(b)
	if vA.Type() != vB.Type() {
		return fmt.Sprintf("type A %T and type B %T do not match", a, b)
	}
	diffs := objectReflectDiff(field.NewPath("object"), vA, vB)
	if len(diffs) == 0 {
		return ""
	}
	out := []string{""}
	for _, d := range diffs {
		out = append(out,
			fmt.Sprintf("%s:", d.path),
			limit(fmt.Sprintf("  a: %#v", d.a), 80),
			limit(fmt.Sprintf("  b: %#v", d.b), 80),
		)
	}
	return strings.Join(out, "\n")
}

func limit(s string, max int) string {
	if len(s) > max {
		return s[:max]
	}
	return s
}

func public(s string) bool {
	if len(s) == 0 {
		return false
	}
	return s[:1] == strings.ToUpper(s[:1])
}

type diff struct {
	path *field.Path
	a, b interface{}
}

type orderedDiffs []diff

func (d orderedDiffs) Len() int      { return len(d) }
func (d orderedDiffs) Swap(i, j int) { d[i], d[j] = d[j], d[i] }
func (d orderedDiffs) Less(i, j int) bool {
	a, b := d[i].path.String(), d[j].path.String()
	if a < b {
		return true
	}
	return false
}

func objectReflectDiff(path *field.Path, a, b reflect.Value) []diff {
	switch a.Type().Kind() {
	case reflect.Struct:
		var changes []diff
		for i := 0; i < a.Type().NumField(); i++ {
			if !public(a.Type().Field(i).Name) {
				if reflect.DeepEqual(a.Interface(), b.Interface()) {
					return nil
				}
				return []diff{{path: path, a: fmt.Sprintf("%#v", a), b: fmt.Sprintf("%#v", b)}}
			}
			if sub := objectReflectDiff(path.Child(a.Type().Field(i).Name), a.Field(i), b.Field(i)); len(sub) > 0 {
				changes = append(changes, sub...)
			}
		}
		return changes
	case reflect.Ptr:
		if a.IsNil() || b.IsNil() {
			switch {
			case a.IsNil() && b.IsNil():
				return nil
			case a.IsNil():
				return []diff{{path: path, a: nil, b: b.Interface()}}
			default:
				return []diff{{path: path, a: a.Interface(), b: nil}}
			}
		}
		return objectReflectDiff(path, a.Elem(), b.Elem())
	case reflect.Chan:
		if !reflect.DeepEqual(a.Interface(), b.Interface()) {
			return []diff{{path: path, a: a.Interface(), b: b.Interface()}}
		}
		return nil
	case reflect.Slice:
		if reflect.DeepEqual(a, b) {
			return nil
		}
		lA, lB := a.Len(), b.Len()
		l := lA
		if lB < lA {
			l = lB
		}
		for i := 0; i < l; i++ {
			if !reflect.DeepEqual(a.Index(i), b.Index(i)) {
				return objectReflectDiff(path.Index(i), a.Index(i), b.Index(i))
			}
		}
		var diffs []diff
		for i := l; l < lA; i++ {
			diffs = append(diffs, diff{path: path.Index(i), a: a.Index(i), b: nil})
		}
		for i := l; l < lB; i++ {
			diffs = append(diffs, diff{path: path.Index(i), a: nil, b: b.Index(i)})
		}
		return diffs
	case reflect.Map:
		if reflect.DeepEqual(a, b) {
			return nil
		}
		aKeys := make(map[interface{}]interface{})
		for _, key := range a.MapKeys() {
			aKeys[key.Interface()] = a.MapIndex(key).Interface()
		}
		var missing []diff
		for _, key := range b.MapKeys() {
			if _, ok := aKeys[key.Interface()]; ok {
				delete(aKeys, key.Interface())
				if reflect.DeepEqual(a.MapIndex(key).Interface(), b.MapIndex(key).Interface()) {
					continue
				}
				missing = append(missing, diff{path: path.Key(fmt.Sprintf("%s", key.Interface())), a: a.MapIndex(key).Interface(), b: b.MapIndex(key).Interface()})
				continue
			}
			missing = append(missing, diff{path: path.Key(fmt.Sprintf("%s", key.Interface())), a: nil, b: b.MapIndex(key).Interface()})
		}
		for key, value := range aKeys {
			missing = append(missing, diff{path: path.Key(fmt.Sprintf("%s", key)), a: value, b: nil})
		}
		sort.Sort(orderedDiffs(missing))
		return missing
	default:
		if reflect.DeepEqual(a.Interface(), b.Interface()) {
			return nil
		}
		if !a.CanInterface() {
			return []diff{{path: path, a: fmt.Sprintf("%#v", a), b: fmt.Sprintf("%#v", b)}}
		}
		return []diff{{path: path, a: a.Interface(), b: b.Interface()}}
	}
}

// ObjectGoPrintSideBySide prints a and b as textual dumps side by side,
// enabling easy visual scanning for mismatches.
func ObjectGoPrintSideBySide(a, b interface{}) string {
	s := spew.ConfigState{
		Indent: " ",
		// Extra deep spew.
		DisableMethods: true,
	}
	sA := s.Sdump(a)
	sB := s.Sdump(b)

	linesA := strings.Split(sA, "\n")
	linesB := strings.Split(sB, "\n")
	width := 0
	for _, s := range linesA {
		l := len(s)
		if l > width {
			width = l
		}
	}
	for _, s := range linesB {
		l := len(s)
		if l > width {
			width = l
		}
	}
	buf := &bytes.Buffer{}
	w := tabwriter.NewWriter(buf, width, 0, 1, ' ', 0)
	max := len(linesA)
	if len(linesB) > max {
		max = len(linesB)
	}
	for i := 0; i < max; i++ {
		var a, b string
		if i < len(linesA) {
			a = linesA[i]
		}
		if i < len(linesB) {
			b = linesB[i]
		}
		fmt.Fprintf(w, "%s\t%s\n", a, b)
	}
	w.Flush()
	return buf.String()
}
<commit_msg>ObjectReflectDiff has an incorrect array index comparison<commit_after>/*
Copyright 2014 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package diff

import (
	"bytes"
	"encoding/json"
	"fmt"
	"reflect"
	"sort"
	"strings"
	"text/tabwriter"

	"github.com/davecgh/go-spew/spew"

	"k8s.io/kubernetes/pkg/util/validation/field"
)

// StringDiff diffs a and b and returns a human readable diff.
func StringDiff(a, b string) string {
	ba := []byte(a)
	bb := []byte(b)
	out := []byte{}
	i := 0
	for ; i < len(ba) && i < len(bb); i++ {
		if ba[i] != bb[i] {
			break
		}
		out = append(out, ba[i])
	}
	out = append(out, []byte("\n\nA: ")...)
	out = append(out, ba[i:]...)
	out = append(out, []byte("\n\nB: ")...)
	out = append(out, bb[i:]...)
	out = append(out, []byte("\n\n")...)
	return string(out)
}

// ObjectDiff writes the two objects out as JSON and prints out the identical part of
// the objects followed by the remaining part of 'a' and finally the remaining part of 'b'.
// For debugging tests.
func ObjectDiff(a, b interface{}) string {
	ab, err := json.Marshal(a)
	if err != nil {
		panic(fmt.Sprintf("a: %v", err))
	}
	bb, err := json.Marshal(b)
	if err != nil {
		panic(fmt.Sprintf("b: %v", err))
	}
	return StringDiff(string(ab), string(bb))
}

// ObjectGoPrintDiff is like ObjectDiff, but uses go-spew to print the objects,
// which shows absolutely everything by recursing into every single pointer
// (go's %#v formatters OTOH stop at a certain point). This is needed when you
// can't figure out why reflect.DeepEqual is returning false and nothing is
// showing you differences. This will.
func ObjectGoPrintDiff(a, b interface{}) string {
	s := spew.ConfigState{DisableMethods: true}
	return StringDiff(
		s.Sprintf("%#v", a),
		s.Sprintf("%#v", b),
	)
}

func ObjectReflectDiff(a, b interface{}) string {
	vA, vB := reflect.ValueOf(a), reflect.ValueOf(b)
	if vA.Type() != vB.Type() {
		return fmt.Sprintf("type A %T and type B %T do not match", a, b)
	}
	diffs := objectReflectDiff(field.NewPath("object"), vA, vB)
	if len(diffs) == 0 {
		return ""
	}
	out := []string{""}
	for _, d := range diffs {
		out = append(out,
			fmt.Sprintf("%s:", d.path),
			limit(fmt.Sprintf("  a: %#v", d.a), 80),
			limit(fmt.Sprintf("  b: %#v", d.b), 80),
		)
	}
	return strings.Join(out, "\n")
}

func limit(s string, max int) string {
	if len(s) > max {
		return s[:max]
	}
	return s
}

func public(s string) bool {
	if len(s) == 0 {
		return false
	}
	return s[:1] == strings.ToUpper(s[:1])
}

type diff struct {
	path *field.Path
	a, b interface{}
}

type orderedDiffs []diff

func (d orderedDiffs) Len() int      { return len(d) }
func (d orderedDiffs) Swap(i, j int) { d[i], d[j] = d[j], d[i] }
func (d orderedDiffs) Less(i, j int) bool {
	a, b := d[i].path.String(), d[j].path.String()
	if a < b {
		return true
	}
	return false
}

func objectReflectDiff(path *field.Path, a, b reflect.Value) []diff {
	switch a.Type().Kind() {
	case reflect.Struct:
		var changes []diff
		for i := 0; i < a.Type().NumField(); i++ {
			if !public(a.Type().Field(i).Name) {
				if reflect.DeepEqual(a.Interface(), b.Interface()) {
					return nil
				}
				return []diff{{path: path, a: fmt.Sprintf("%#v", a), b: fmt.Sprintf("%#v", b)}}
			}
			if sub := objectReflectDiff(path.Child(a.Type().Field(i).Name), a.Field(i), b.Field(i)); len(sub) > 0 {
				changes = append(changes, sub...)
			}
		}
		return changes
	case reflect.Ptr:
		if a.IsNil() || b.IsNil() {
			switch {
			case a.IsNil() && b.IsNil():
				return nil
			case a.IsNil():
				return []diff{{path: path, a: nil, b: b.Interface()}}
			default:
				return []diff{{path: path, a: a.Interface(), b: nil}}
			}
		}
		return objectReflectDiff(path, a.Elem(), b.Elem())
	case reflect.Chan:
		if !reflect.DeepEqual(a.Interface(), b.Interface()) {
			return []diff{{path: path, a: a.Interface(), b: b.Interface()}}
		}
		return nil
	case reflect.Slice:
		if reflect.DeepEqual(a, b) {
			return nil
		}
		lA, lB := a.Len(), b.Len()
		l := lA
		if lB < lA {
			l = lB
		}
		for i := 0; i < l; i++ {
			if !reflect.DeepEqual(a.Index(i), b.Index(i)) {
				return objectReflectDiff(path.Index(i), a.Index(i), b.Index(i))
			}
		}
		var diffs []diff
		for i := l; i < lA; i++ {
			diffs = append(diffs, diff{path: path.Index(i), a: a.Index(i), b: nil})
		}
		for i := l; i < lB; i++ {
			diffs = append(diffs, diff{path: path.Index(i), a: nil, b: b.Index(i)})
		}
		return diffs
	case reflect.Map:
		if reflect.DeepEqual(a, b) {
			return nil
		}
		aKeys := make(map[interface{}]interface{})
		for _, key := range a.MapKeys() {
			aKeys[key.Interface()] = a.MapIndex(key).Interface()
		}
		var missing []diff
		for _, key := range b.MapKeys() {
			if _, ok := aKeys[key.Interface()]; ok {
				delete(aKeys, key.Interface())
				if reflect.DeepEqual(a.MapIndex(key).Interface(), b.MapIndex(key).Interface()) {
					continue
				}
				missing = append(missing, diff{path: path.Key(fmt.Sprintf("%s", key.Interface())), a: a.MapIndex(key).Interface(), b: b.MapIndex(key).Interface()})
				continue
			}
			missing = append(missing, diff{path: path.Key(fmt.Sprintf("%s", key.Interface())), a: nil, b: b.MapIndex(key).Interface()})
		}
		for key, value := range aKeys {
			missing = append(missing, diff{path: path.Key(fmt.Sprintf("%s", key)), a: value, b: nil})
		}
		sort.Sort(orderedDiffs(missing))
		return missing
	default:
		if reflect.DeepEqual(a.Interface(), b.Interface()) {
			return nil
		}
		if !a.CanInterface() {
			return []diff{{path: path, a: fmt.Sprintf("%#v", a), b: fmt.Sprintf("%#v", b)}}
		}
		return []diff{{path: path, a: a.Interface(), b: b.Interface()}}
	}
}

// ObjectGoPrintSideBySide prints a and b as textual dumps side by side,
// enabling easy visual scanning for mismatches.
func ObjectGoPrintSideBySide(a, b interface{}) string {
	s := spew.ConfigState{
		Indent: " ",
		// Extra deep spew.
		DisableMethods: true,
	}
	sA := s.Sdump(a)
	sB := s.Sdump(b)

	linesA := strings.Split(sA, "\n")
	linesB := strings.Split(sB, "\n")
	width := 0
	for _, s := range linesA {
		l := len(s)
		if l > width {
			width = l
		}
	}
	for _, s := range linesB {
		l := len(s)
		if l > width {
			width = l
		}
	}
	buf := &bytes.Buffer{}
	w := tabwriter.NewWriter(buf, width, 0, 1, ' ', 0)
	max := len(linesA)
	if len(linesB) > max {
		max = len(linesB)
	}
	for i := 0; i < max; i++ {
		var a, b string
		if i < len(linesA) {
			a = linesA[i]
		}
		if i < len(linesB) {
			b = linesB[i]
		}
		fmt.Fprintf(w, "%s\t%s\n", a, b)
	}
	w.Flush()
	return buf.String()
}
<|endoftext|>


<commit_before>package mpb

import (
	"bytes"
	"container/heap"
	"context"
	"fmt"
	"io"
	"math"
	"os"
	"sync"
	"time"

	"github.com/vbauerster/mpb/v8/cwriter"
)

const (
	prr = 150 * time.Millisecond // default RefreshRate
)

// DoneError represents an error when `*mpb.Progress` is done but its functionality is requested.
var DoneError = fmt.Errorf("%T instance can't be reused after it's done!", (*Progress)(nil))

// Progress represents a container that renders one or more progress bars.
type Progress struct {
	ctx          context.Context
	uwg          *sync.WaitGroup
	cwg          *sync.WaitGroup
	bwg          *sync.WaitGroup
	operateState chan func(*pState)
	interceptIo  chan func(io.Writer)
	done         chan struct{}
	once         sync.Once
	cancel       func()
}

// pState holds bars in its priorityQueue, it gets passed to (*Progress).serve monitor goroutine.
type pState struct {
	bHeap       priorityQueue
	heapUpdated bool
	pMatrix     map[int][]chan int
	aMatrix     map[int][]chan int

	// for reuse purposes
	rows []io.Reader
	pool []*Bar

	// following are provided/overrided by user
	rr                 time.Duration
	idCount            int
	reqWidth           int
	popPriority        int
	popCompleted       bool
	outputDiscarded    bool
	disableAutoRefresh bool
	manualRefresh      chan interface{}
	renderDelay        <-chan struct{}
	shutdownNotifier   chan struct{}
	queueBars          map[*Bar]*Bar
	output             io.Writer
	debugOut           io.Writer
	uwg                *sync.WaitGroup
}

// New creates new Progress container instance. It's not possible to
// reuse instance after (*Progress).Wait method has been called.
func New(options ...ContainerOption) *Progress {
	return NewWithContext(context.Background(), options...)
}

// NewWithContext creates new Progress container instance with provided
// context. It's not possible to reuse instance after (*Progress).Wait
// method has been called.
func NewWithContext(ctx context.Context, options ...ContainerOption) *Progress {
	s := &pState{
		rr:            prr,
		bHeap:         priorityQueue{},
		rows:          make([]io.Reader, 0, 64),
		pool:          make([]*Bar, 0, 64),
		manualRefresh: make(chan interface{}),
		queueBars:     make(map[*Bar]*Bar),
		popPriority:   math.MinInt32,
		output:        os.Stdout,
		debugOut:      io.Discard,
	}

	for _, opt := range options {
		if opt != nil {
			opt(s)
		}
	}

	if s.shutdownNotifier == nil {
		s.shutdownNotifier = make(chan struct{})
	}

	ctx, cancel := context.WithCancel(ctx)
	p := &Progress{
		ctx:          ctx,
		uwg:          s.uwg,
		cwg:          new(sync.WaitGroup),
		bwg:          new(sync.WaitGroup),
		operateState: make(chan func(*pState)),
		interceptIo:  make(chan func(io.Writer)),
		done:         make(chan struct{}),
		cancel:       cancel,
	}

	p.cwg.Add(1)
	go p.serve(s, cwriter.New(s.output))
	return p
}

// AddBar creates a bar with default bar filler.
func (p *Progress) AddBar(total int64, options ...BarOption) *Bar {
	return p.New(total, BarStyle(), options...)
}

// AddSpinner creates a bar with default spinner filler.
func (p *Progress) AddSpinner(total int64, options ...BarOption) *Bar {
	return p.New(total, SpinnerStyle(), options...)
}

// New creates a bar by calling `Build` method on provided `BarFillerBuilder`.
func (p *Progress) New(total int64, builder BarFillerBuilder, options ...BarOption) *Bar {
	return p.AddFiller(total, builder.Build(), options...)
}

// AddFiller creates a bar which renders itself by provided filler.
// If `total <= 0` triggering complete event by increment methods is disabled.
// Panics if *Progress instance is done, i.e. called after (*Progress).Wait().
func (p *Progress) AddFiller(total int64, filler BarFiller, options ...BarOption) *Bar {
	if filler == nil {
		filler = NopStyle().Build()
	}
	p.bwg.Add(1)
	result := make(chan *Bar)
	select {
	case p.operateState <- func(ps *pState) {
		bs := ps.makeBarState(total, filler, options...)
		bar := newBar(p, bs)
		if bs.wait.bar != nil {
			ps.queueBars[bs.wait.bar] = bar
		} else {
			heap.Push(&ps.bHeap, bar)
			ps.heapUpdated = true
		}
		ps.idCount++
		result <- bar
	}:
		bar := <-result
		return bar
	case <-p.done:
		p.bwg.Done()
		panic(DoneError)
	}
}

func (p *Progress) traverseBars(cb func(b *Bar) bool) {
	sync := make(chan struct{})
	select {
	case p.operateState <- func(s *pState) {
		for i := 0; i < s.bHeap.Len(); i++ {
			bar := s.bHeap[i]
			if !cb(bar) {
				break
			}
		}
		close(sync)
	}:
		<-sync
	case <-p.done:
	}
}

// UpdateBarPriority same as *Bar.SetPriority(int).
func (p *Progress) UpdateBarPriority(b *Bar, priority int) {
	select {
	case p.operateState <- func(s *pState) {
		if b.index < 0 {
			return
		}
		b.priority = priority
		heap.Fix(&s.bHeap, b.index)
	}:
	case <-p.done:
	}
}

// BarCount returns bars count.
func (p *Progress) BarCount() int {
	result := make(chan int)
	select {
	case p.operateState <- func(s *pState) { result <- s.bHeap.Len() }:
		return <-result
	case <-p.done:
		return 0
	}
}

// Write is implementation of io.Writer.
// Writing to `*mpb.Progress` will print lines above a running bar.
// Writes aren't flushed immediately, but at next refresh cycle.
// If Write is called after `*mpb.Progress` is done, `mpb.DoneError`
// is returned.
func (p *Progress) Write(b []byte) (int, error) {
	type result struct {
		n   int
		err error
	}
	ch := make(chan *result)
	select {
	case p.interceptIo <- func(w io.Writer) {
		n, err := w.Write(b)
		ch <- &result{n, err}
	}:
		res := <-ch
		return res.n, res.err
	case <-p.done:
		return 0, DoneError
	}
}

// Wait waits for all bars to complete and finally shutdowns container. After
// this method has been called, there is no way to reuse (*Progress) instance.
func (p *Progress) Wait() {
	// wait for user wg, if any
	if p.uwg != nil {
		p.uwg.Wait()
	}

	// wait for bars to quit, if any
	p.bwg.Wait()
	// shutdown
	p.once.Do(p.shutdown)
	// wait for container to quit
	p.cwg.Wait()
}

// Shutdown cancels any running bar immediately and then shutdowns (*Progress)
// instance. Normally this method shouldn't be called unless you know what you
// are doing. Proper way to shutdown is to call (*Progress).Wait() instead.
func (p *Progress) Shutdown() {
	p.cancel()
	p.bwg.Wait()
	p.once.Do(p.shutdown)
	p.cwg.Wait()
}

func (p *Progress) shutdown() {
	close(p.done)
}

func (p *Progress) serve(s *pState, cw *cwriter.Writer) {
	defer p.cwg.Done()

	render := func() error {
		if s.bHeap.Len() == 0 {
			return nil
		}
		return s.render(cw)
	}

	refreshCh := s.newTicker(p.done)

	for {
		select {
		case op := <-p.operateState:
			op(s)
		case fn := <-p.interceptIo:
			fn(cw)
		case <-refreshCh:
			err := render()
			if err != nil {
				go func() {
					p.bwg.Wait()
					p.once.Do(p.shutdown)
				}()
				render = func() error {
					s.heapUpdated = false
					return nil
				}
				_, _ = fmt.Fprintln(s.debugOut, err)
				p.cancel() // cancel all bars
			}
		case <-s.shutdownNotifier:
			for s.heapUpdated {
				err := render()
				if err != nil {
					_, _ = fmt.Fprintln(s.debugOut, err)
					return
				}
			}
			return
		}
	}
}

func (s *pState) render(cw *cwriter.Writer) error {
	if s.heapUpdated {
		s.updateSyncMatrix()
		s.heapUpdated = false
	}
	syncWidth(s.pMatrix)
	syncWidth(s.aMatrix)

	width, height, err := cw.GetTermSize()
	if err != nil {
		width = s.reqWidth
		height = s.bHeap.Len()
	}
	for i := 0; i < s.bHeap.Len(); i++ {
		bar := s.bHeap[i]
		go bar.render(width)
	}

	return s.flush(cw, height)
}

func (s *pState) flush(cw *cwriter.Writer, height int) (err error) {
	var wg sync.WaitGroup
	var popCount int

	for s.bHeap.Len() > 0 {
		b := heap.Pop(&s.bHeap).(*Bar)
		frame := <-b.frameCh
		if frame.err != nil {
			if err == nil {
				err = frame.err
			}
			continue
		}
		var usedRows int
		for i := len(frame.rows) - 1; i >= 0; i-- {
			if row := frame.rows[i]; len(s.rows) < height {
				s.rows = append(s.rows, row)
				usedRows++
			} else {
				wg.Add(1)
				go func() {
					_, _ = io.Copy(io.Discard, row)
					wg.Done()
				}()
			}
		}
		if frame.shutdown != 0 {
			b.Wait() // waiting for b.done, so it's safe to read b.bs
			drop := b.bs.dropOnComplete
			if qb, ok := s.queueBars[b]; ok {
				delete(s.queueBars, b)
				qb.priority = b.priority
				s.pool = append(s.pool, qb)
				drop = true
			} else if s.popCompleted && !b.bs.noPop {
				if frame.shutdown > 1 {
					popCount += usedRows
					drop = true
				} else {
					s.popPriority++
					b.priority = s.popPriority
				}
			}
			if drop {
				s.heapUpdated = true
				continue
			}
		}
		s.pool = append(s.pool, b)
	}

	if err != nil {
		return err
	}

	wg.Add(1)
	go func() {
		for _, b := range s.pool {
			heap.Push(&s.bHeap, b)
		}
		wg.Done()
	}()

	for i := len(s.rows) - 1; i >= 0; i-- {
		_, err := cw.ReadFrom(s.rows[i])
		if err != nil {
			wg.Wait()
			return err
		}
	}

	err = cw.Flush(len(s.rows) - popCount)
	wg.Wait()
	s.rows = s.rows[:0]
	s.pool = s.pool[:0]
	return err
}

func (s *pState) newTicker(done <-chan struct{}) chan time.Time {
	ch := make(chan time.Time)
	go func() {
		var autoRefresh <-chan time.Time
		if !s.disableAutoRefresh {
			if !s.outputDiscarded {
				if s.renderDelay != nil {
					<-s.renderDelay
				}
				ticker := time.NewTicker(s.rr)
				defer ticker.Stop()
				autoRefresh = ticker.C
			}
		}
		for {
			select {
			case t := <-autoRefresh:
				ch <- t
			case x := <-s.manualRefresh:
				if t, ok := x.(time.Time); ok {
					ch <- t
				} else {
					ch <- time.Now()
				}
			case <-done:
				close(s.shutdownNotifier)
				return
			}
		}
	}()
	return ch
}

func (s *pState) updateSyncMatrix() {
	s.pMatrix = make(map[int][]chan int)
	s.aMatrix = make(map[int][]chan int)
	for i := 0; i < s.bHeap.Len(); i++ {
		bar := s.bHeap[i]
		table := bar.wSyncTable()
		pRow, aRow := table[0], table[1]

		for i, ch := range pRow {
			s.pMatrix[i] = append(s.pMatrix[i], ch)
		}

		for i, ch := range aRow {
			s.aMatrix[i] = append(s.aMatrix[i], ch)
		}
	}
}

func (s *pState) makeBarState(total int64, filler BarFiller, options ...BarOption) *bState {
	bs := &bState{
		id:            s.idCount,
		priority:      s.idCount,
		reqWidth:      s.reqWidth,
		total:         total,
		filler:        filler,
		manualRefresh: s.manualRefresh,
	}

	if total > 0 {
		bs.triggerComplete = true
	}

	for _, opt := range options {
		if opt != nil {
			opt(bs)
		}
	}

	if bs.middleware != nil {
		bs.filler = bs.middleware(filler)
		bs.middleware = nil
	}

	for i := 0; i < len(bs.buffers); i++ {
		bs.buffers[i] = bytes.NewBuffer(make([]byte, 0, 512))
	}

	bs.subscribeDecorators()

	return bs
}

func syncWidth(matrix map[int][]chan int) {
	for _, column := range matrix {
		go maxWidthDistributor(column)
	}
}

func maxWidthDistributor(column []chan int) {
	var maxWidth int
	for _, ch := range column {
		if w := <-ch; w > maxWidth {
			maxWidth = w
		}
	}
	for _, ch := range column {
		ch <- maxWidth
	}
}
<commit_msg>it's ok to return error immediately<commit_after>package mpb

import (
	"bytes"
	"container/heap"
	"context"
	"fmt"
	"io"
	"math"
	"os"
	"sync"
	"time"

	"github.com/vbauerster/mpb/v8/cwriter"
)

const (
	prr = 150 * time.Millisecond // default RefreshRate
)

// DoneError represents an error when `*mpb.Progress` is done but its functionality is requested.
var DoneError = fmt.Errorf("%T instance can't be reused after it's done!", (*Progress)(nil))

// Progress represents a container that renders one or more progress bars.
type Progress struct {
	ctx          context.Context
	uwg          *sync.WaitGroup
	cwg          *sync.WaitGroup
	bwg          *sync.WaitGroup
	operateState chan func(*pState)
	interceptIo  chan func(io.Writer)
	done         chan struct{}
	once         sync.Once
	cancel       func()
}

// pState holds bars in its priorityQueue, it gets passed to (*Progress).serve monitor goroutine.
type pState struct {
	bHeap       priorityQueue
	heapUpdated bool
	pMatrix     map[int][]chan int
	aMatrix     map[int][]chan int

	// for reuse purposes
	rows []io.Reader
	pool []*Bar

	// following are provided/overrided by user
	rr                 time.Duration
	idCount            int
	reqWidth           int
	popPriority        int
	popCompleted       bool
	outputDiscarded    bool
	disableAutoRefresh bool
	manualRefresh      chan interface{}
	renderDelay        <-chan struct{}
	shutdownNotifier   chan struct{}
	queueBars          map[*Bar]*Bar
	output             io.Writer
	debugOut           io.Writer
	uwg                *sync.WaitGroup
}

// New creates new Progress container instance. It's not possible to
// reuse instance after (*Progress).Wait method has been called.
func New(options ...ContainerOption) *Progress {
	return NewWithContext(context.Background(), options...)
}

// NewWithContext creates new Progress container instance with provided
// context. It's not possible to reuse instance after (*Progress).Wait
// method has been called.
func NewWithContext(ctx context.Context, options ...ContainerOption) *Progress {
	s := &pState{
		rr:            prr,
		bHeap:         priorityQueue{},
		rows:          make([]io.Reader, 0, 64),
		pool:          make([]*Bar, 0, 64),
		manualRefresh: make(chan interface{}),
		queueBars:     make(map[*Bar]*Bar),
		popPriority:   math.MinInt32,
		output:        os.Stdout,
		debugOut:      io.Discard,
	}

	for _, opt := range options {
		if opt != nil {
			opt(s)
		}
	}

	if s.shutdownNotifier == nil {
		s.shutdownNotifier = make(chan struct{})
	}

	ctx, cancel := context.WithCancel(ctx)
	p := &Progress{
		ctx:          ctx,
		uwg:          s.uwg,
		cwg:          new(sync.WaitGroup),
		bwg:          new(sync.WaitGroup),
		operateState: make(chan func(*pState)),
		interceptIo:  make(chan func(io.Writer)),
		done:         make(chan struct{}),
		cancel:       cancel,
	}

	p.cwg.Add(1)
	go p.serve(s, cwriter.New(s.output))
	return p
}

// AddBar creates a bar with default bar filler.
func (p *Progress) AddBar(total int64, options ...BarOption) *Bar {
	return p.New(total, BarStyle(), options...)
}

// AddSpinner creates a bar with default spinner filler.
func (p *Progress) AddSpinner(total int64, options ...BarOption) *Bar {
	return p.New(total, SpinnerStyle(), options...)
}

// New creates a bar by calling `Build` method on provided `BarFillerBuilder`.
func (p *Progress) New(total int64, builder BarFillerBuilder, options ...BarOption) *Bar {
	return p.AddFiller(total, builder.Build(), options...)
}

// AddFiller creates a bar which renders itself by provided filler.
// If `total <= 0` triggering complete event by increment methods is disabled.
// Panics if *Progress instance is done, i.e. called after (*Progress).Wait().
func (p *Progress) AddFiller(total int64, filler BarFiller, options ...BarOption) *Bar {
	if filler == nil {
		filler = NopStyle().Build()
	}
	p.bwg.Add(1)
	result := make(chan *Bar)
	select {
	case p.operateState <- func(ps *pState) {
		bs := ps.makeBarState(total, filler, options...)
		bar := newBar(p, bs)
		if bs.wait.bar != nil {
			ps.queueBars[bs.wait.bar] = bar
		} else {
			heap.Push(&ps.bHeap, bar)
			ps.heapUpdated = true
		}
		ps.idCount++
		result <- bar
	}:
		bar := <-result
		return bar
	case <-p.done:
		p.bwg.Done()
		panic(DoneError)
	}
}

func (p *Progress) traverseBars(cb func(b *Bar) bool) {
	sync := make(chan struct{})
	select {
	case p.operateState <- func(s *pState) {
		for i := 0; i < s.bHeap.Len(); i++ {
			bar := s.bHeap[i]
			if !cb(bar) {
				break
			}
		}
		close(sync)
	}:
		<-sync
	case <-p.done:
	}
}

// UpdateBarPriority same as *Bar.SetPriority(int).
func (p *Progress) UpdateBarPriority(b *Bar, priority int) {
	select {
	case p.operateState <- func(s *pState) {
		if b.index < 0 {
			return
		}
		b.priority = priority
		heap.Fix(&s.bHeap, b.index)
	}:
	case <-p.done:
	}
}

// BarCount returns bars count.
func (p *Progress) BarCount() int {
	result := make(chan int)
	select {
	case p.operateState <- func(s *pState) { result <- s.bHeap.Len() }:
		return <-result
	case <-p.done:
		return 0
	}
}

// Write is implementation of io.Writer.
// Writing to `*mpb.Progress` will print lines above a running bar.
// Writes aren't flushed immediately, but at next refresh cycle.
// If Write is called after `*mpb.Progress` is done, `mpb.DoneError`
// is returned.
func (p *Progress) Write(b []byte) (int, error) {
	type result struct {
		n   int
		err error
	}
	ch := make(chan *result)
	select {
	case p.interceptIo <- func(w io.Writer) {
		n, err := w.Write(b)
		ch <- &result{n, err}
	}:
		res := <-ch
		return res.n, res.err
	case <-p.done:
		return 0, DoneError
	}
}

// Wait waits for all bars to complete and finally shutdowns container. After
// this method has been called, there is no way to reuse (*Progress) instance.
func (p *Progress) Wait() {
	// wait for user wg, if any
	if p.uwg != nil {
		p.uwg.Wait()
	}

	// wait for bars to quit, if any
	p.bwg.Wait()
	// shutdown
	p.once.Do(p.shutdown)
	// wait for container to quit
	p.cwg.Wait()
}

// Shutdown cancels any running bar immediately and then shutdowns (*Progress)
// instance. Normally this method shouldn't be called unless you know what you
// are doing. Proper way to shutdown is to call (*Progress).Wait() instead.
func (p *Progress) Shutdown() {
	p.cancel()
	p.bwg.Wait()
	p.once.Do(p.shutdown)
	p.cwg.Wait()
}

func (p *Progress) shutdown() {
	close(p.done)
}

func (p *Progress) serve(s *pState, cw *cwriter.Writer) {
	defer p.cwg.Done()

	render := func() error {
		if s.bHeap.Len() == 0 {
			return nil
		}
		return s.render(cw)
	}

	refreshCh := s.newTicker(p.done)

	for {
		select {
		case op := <-p.operateState:
			op(s)
		case fn := <-p.interceptIo:
			fn(cw)
		case <-refreshCh:
			err := render()
			if err != nil {
				go func() {
					p.bwg.Wait()
					p.once.Do(p.shutdown)
				}()
				render = func() error {
					s.heapUpdated = false
					return nil
				}
				_, _ = fmt.Fprintln(s.debugOut, err)
				p.cancel() // cancel all bars
			}
		case <-s.shutdownNotifier:
			for s.heapUpdated {
				err := render()
				if err != nil {
					_, _ = fmt.Fprintln(s.debugOut, err)
					return
				}
			}
			return
		}
	}
}

func (s *pState) render(cw *cwriter.Writer) error {
	if s.heapUpdated {
		s.updateSyncMatrix()
		s.heapUpdated = false
	}
	syncWidth(s.pMatrix)
	syncWidth(s.aMatrix)

	width, height, err := cw.GetTermSize()
	if err != nil {
		width = s.reqWidth
		height = s.bHeap.Len()
	}
	for i := 0; i < s.bHeap.Len(); i++ {
		bar := s.bHeap[i]
		go bar.render(width)
	}

	return s.flush(cw, height)
}

func (s *pState) flush(cw *cwriter.Writer, height int) error {
	var wg sync.WaitGroup
	var popCount int

	for s.bHeap.Len() > 0 {
		b := heap.Pop(&s.bHeap).(*Bar)
		frame := <-b.frameCh
		if frame.err != nil {
			return frame.err
		}
		var usedRows int
		for i := len(frame.rows) - 1; i >= 0; i-- {
			if row := frame.rows[i]; len(s.rows) < height {
				s.rows = append(s.rows, row)
				usedRows++
			} else {
				wg.Add(1)
				go func() {
					_, _ = io.Copy(io.Discard, row)
					wg.Done()
				}()
			}
		}
		if frame.shutdown != 0 {
			b.Wait() // waiting for b.done, so it's safe to read b.bs
			drop := b.bs.dropOnComplete
			if qb, ok := s.queueBars[b]; ok {
				delete(s.queueBars, b)
				qb.priority = b.priority
				s.pool = append(s.pool, qb)
				drop = true
			} else if s.popCompleted && !b.bs.noPop {
				if frame.shutdown > 1 {
					popCount += usedRows
					drop = true
				} else {
					s.popPriority++
					b.priority = s.popPriority
				}
			}
			if drop {
				s.heapUpdated = true
				continue
			}
		}
		s.pool = append(s.pool, b)
	}

	wg.Add(1)
	go func() {
		for _, b := range s.pool {
			heap.Push(&s.bHeap, b)
		}
		wg.Done()
	}()

	for i := len(s.rows) - 1; i >= 0; i-- {
		_, err := cw.ReadFrom(s.rows[i])
		if err != nil {
			wg.Wait()
			return err
		}
	}

	err := cw.Flush(len(s.rows) - popCount)
	wg.Wait()
	s.rows = s.rows[:0]
	s.pool = s.pool[:0]
	return err
}

func (s *pState) newTicker(done <-chan struct{}) chan time.Time {
	ch := make(chan time.Time)
	go func() {
		var autoRefresh <-chan time.Time
		if !s.disableAutoRefresh {
			if !s.outputDiscarded {
				if s.renderDelay != nil {
					<-s.renderDelay
				}
				ticker := time.NewTicker(s.rr)
				defer ticker.Stop()
				autoRefresh = ticker.C
			}
		}
		for {
			select {
			case t := <-autoRefresh:
				ch <- t
			case x := <-s.manualRefresh:
				if t, ok := x.(time.Time); ok {
					ch <- t
				} else {
					ch <- time.Now()
				}
			case <-done:
				close(s.shutdownNotifier)
				return
			}
		}
	}()
	return ch
}

func (s *pState) updateSyncMatrix() {
	s.pMatrix = make(map[int][]chan int)
	s.aMatrix = make(map[int][]chan int)
	for i := 0; i < s.bHeap.Len(); i++ {
		bar := s.bHeap[i]
		table := bar.wSyncTable()
		pRow, aRow := table[0], table[1]

		for i, ch := range pRow {
			s.pMatrix[i] = append(s.pMatrix[i], ch)
		}

		for i, ch := range aRow {
			s.aMatrix[i] = append(s.aMatrix[i], ch)
		}
	}
}

func (s *pState) makeBarState(total int64, filler BarFiller, options ...BarOption) *bState {
	bs := &bState{
		id:            s.idCount,
		priority:      s.idCount,
		reqWidth:      s.reqWidth,
		total:         total,
		filler:        filler,
		manualRefresh: s.manualRefresh,
	}

	if total > 0 {
		bs.triggerComplete = true
	}

	for _, opt := range options {
		if opt != nil {
			opt(bs)
		}
	}

	if bs.middleware != nil {
		bs.filler = bs.middleware(filler)
		bs.middleware = nil
	}

	for i := 0; i < len(bs.buffers); i++ {
		bs.buffers[i] = bytes.NewBuffer(make([]byte, 0, 512))
	}

	bs.subscribeDecorators()

	return bs
}

func syncWidth(matrix map[int][]chan int) {
	for _, column := range matrix {
		go maxWidthDistributor(column)
	}
}

func maxWidthDistributor(column []chan int) {
	var maxWidth int
	for _, ch := range column {
		if w := <-ch; w > maxWidth {
			maxWidth = w
		}
	}
	for _, ch := range column {
		ch <- maxWidth
	}
}
<|endoftext|>


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

// oracle: a tool for answering questions about Go source code.
//
// With -format=plain, the oracle prints query results to the standard
// output in an editor-friendly format in which every line of output
// is of the form "pos: text", where pos = "-" if unknown.
//
// With -format=json, the oracle prints structured data in JSON syntax.
//
package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"go/build"
	"io"
	"log"
	"os"
	"runtime"
	"runtime/pprof"

	"code.google.com/p/go.tools/oracle"
)

var posFlag = flag.String("pos", "",
	"Filename and byte offset or extent of a syntax element about which to query, "+
		"e.g. foo.go:#123,#456, bar.go:#123.")

var modeFlag = flag.String("mode", "",
	"Mode of query to perform: e.g. callers, describe, etc.")

var ptalogFlag = flag.String("ptalog", "",
	"Location of the points-to analysis log file, or empty to disable logging.")

var formatFlag = flag.String("format", "plain", "Output format: 'plain' or 'json'.")

const usage = `Go source code oracle.
Usage: oracle [<flag> ...] [<arg> ...]
Use -help flag to display options.

The -mode flag is required; the -pos flag is required in most modes.

Examples:

Describe the syntax at offset 532 in this file (an import spec):
% oracle -mode=describe -pos=src/code.google.com/p/go.tools/cmd/oracle/main.go:#532 \
   code.google.com/p/go.tools/cmd/oracle

Print the callgraph of the trivial web-server in JSON format:
% oracle -mode=callgraph -format=json src/pkg/net/http/triv.go

`

var cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")

func init() {
	// If $GOMAXPROCS isn't set, use the full capacity of the machine.
	// For small machines, use at least 4 threads.
	if os.Getenv("GOMAXPROCS") == "" {
		n := runtime.NumCPU()
		if n < 4 {
			n = 4
		}
		runtime.GOMAXPROCS(n)
	}
}

func main() {
	flag.Parse()
	args := flag.Args()

	if len(args) == 0 {
		fmt.Fprint(os.Stderr, usage)
		os.Exit(1)
	}

	// Set up points-to analysis log file.
	var ptalog io.Writer
	if *ptalogFlag != "" {
		if f, err := os.Create(*ptalogFlag); err != nil {
			log.Fatalf(err.Error())
		} else {
			buf := bufio.NewWriter(f)
			ptalog = buf
			defer func() {
				buf.Flush()
				f.Close()
			}()
		}
	}

	// Profiling support.
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	// -format flag
	if *formatFlag != "json" && *formatFlag != "plain" {
		fmt.Fprintf(os.Stderr, "Error: illegal -format value: %q", *formatFlag)
		os.Exit(1)
	}

	// Ask the oracle.
	res, err := oracle.Query(args, *modeFlag, *posFlag, ptalog, &build.Default)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: %s\n", err)
		os.Exit(1)
	}

	// Print the result.
	switch *formatFlag {
	case "json":
		b, err := json.Marshal(res)
		if err != nil {
			fmt.Fprintf(os.Stderr, "JSON error: %s\n", err.Error())
			os.Exit(1)
		}
		var buf bytes.Buffer
		if err := json.Indent(&buf, b, "", "\t"); err != nil {
			fmt.Fprintf(os.Stderr, "json.Indent failed: %s", err)
			os.Exit(1)
		}
		os.Stdout.Write(buf.Bytes())

	case "plain":
		res.WriteTo(os.Stdout)
	}
}
<commit_msg>go.tools/oracle: restore edits lost in rev aa9b649bccd9 (botched merge)<commit_after>// Copyright 2013 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// oracle: a tool for answering questions about Go source code.
//
// With -format=plain, the oracle prints query results to the standard
// output in an editor-friendly format in which every line of output
// is of the form "pos: text", where pos = "-" if unknown.
//
// With -format=json, the oracle prints structured data in JSON syntax.
//
package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"go/build"
	"io"
	"log"
	"os"
	"runtime"
	"runtime/pprof"

	"code.google.com/p/go.tools/importer"
	"code.google.com/p/go.tools/oracle"
)

var posFlag = flag.String("pos", "",
	"Filename and byte offset or extent of a syntax element about which to query, "+
		"e.g. foo.go:#123,#456, bar.go:#123.")

var modeFlag = flag.String("mode", "",
	"Mode of query to perform: e.g. callers, describe, etc.")

var ptalogFlag = flag.String("ptalog", "",
	"Location of the points-to analysis log file, or empty to disable logging.")

var formatFlag = flag.String("format", "plain", "Output format: 'plain' or 'json'.")

const usage = `Go source code oracle.
Usage: oracle [<flag> ...] <args> ...
Use -help flag to display options.

The -mode flag is required; the -pos flag is required in most modes.

Examples:

Describe the syntax at offset 532 in this file (an import spec):
% oracle -mode=describe -pos=src/code.google.com/p/go.tools/cmd/oracle/main.go:#532 \
   code.google.com/p/go.tools/cmd/oracle

Print the callgraph of the trivial web-server in JSON format:
% oracle -mode=callgraph -format=json src/pkg/net/http/triv.go

` + importer.InitialPackagesUsage

var cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")

func init() {
	// If $GOMAXPROCS isn't set, use the full capacity of the machine.
	// For small machines, use at least 4 threads.
	if os.Getenv("GOMAXPROCS") == "" {
		n := runtime.NumCPU()
		if n < 4 {
			n = 4
		}
		runtime.GOMAXPROCS(n)
	}
}

func main() {
	flag.Parse()
	args := flag.Args()

	if len(args) == 0 {
		fmt.Fprint(os.Stderr, usage)
		os.Exit(1)
	}

	// Set up points-to analysis log file.
	var ptalog io.Writer
	if *ptalogFlag != "" {
		if f, err := os.Create(*ptalogFlag); err != nil {
			log.Fatal(err)
		} else {
			buf := bufio.NewWriter(f)
			ptalog = buf
			defer func() {
				buf.Flush()
				f.Close()
			}()
		}
	}

	// Profiling support.
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	// -format flag
	if *formatFlag != "json" && *formatFlag != "plain" {
		fmt.Fprintf(os.Stderr, "Error: illegal -format value: %q", *formatFlag)
		os.Exit(1)
	}

	// Ask the oracle.
	res, err := oracle.Query(args, *modeFlag, *posFlag, ptalog, &build.Default)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%s\n", err)
		os.Exit(1)
	}

	// Print the result.
	switch *formatFlag {
	case "json":
		b, err := json.Marshal(res)
		if err != nil {
			fmt.Fprintf(os.Stderr, "JSON error: %s\n", err)
			os.Exit(1)
		}
		var buf bytes.Buffer
		if err := json.Indent(&buf, b, "", "\t"); err != nil {
			fmt.Fprintf(os.Stderr, "json.Indent failed: %s", err)
			os.Exit(1)
		}
		os.Stdout.Write(buf.Bytes())

	case "plain":
		res.WriteTo(os.Stdout)
	}
}
<|endoftext|>


<commit_before>// Copyright 2014-2016 Amazon.com, Inc. or its affiliates. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"). You may
// not use this file except in compliance with the License. A copy of the
// License is located at
//
//	http://aws.amazon.com/apache2.0/
//
// or in the "license" file accompanying this file. This file is distributed
// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
// express or implied. See the License for the specific language governing
// permissions and limitations under the License.

// Package handler deals with appropriately reacting to all ACS messages as well
// as maintaining the connection to ACS.
package handler

import (
	"io"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"time"

	"golang.org/x/net/context"

	acsclient "github.com/aws/amazon-ecs-agent/agent/acs/client"
	"github.com/aws/amazon-ecs-agent/agent/acs/model/ecsacs"
	"github.com/aws/amazon-ecs-agent/agent/acs/update_handler"
	"github.com/aws/amazon-ecs-agent/agent/api"
	"github.com/aws/amazon-ecs-agent/agent/config"
	rolecredentials "github.com/aws/amazon-ecs-agent/agent/credentials"
	"github.com/aws/amazon-ecs-agent/agent/engine"
	"github.com/aws/amazon-ecs-agent/agent/eventstream"
	"github.com/aws/amazon-ecs-agent/agent/statemanager"
	"github.com/aws/amazon-ecs-agent/agent/utils"
	"github.com/aws/amazon-ecs-agent/agent/utils/ttime"
	"github.com/aws/amazon-ecs-agent/agent/version"
	"github.com/aws/amazon-ecs-agent/agent/wsclient"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/cihub/seelog"
)

const (
	// heartbeatTimeout is the maximum time to wait between heartbeats
	// without disconnecting
	heartbeatTimeout = 5 * time.Minute
	heartbeatJitter  = 3 * time.Minute

	connectionBackoffMin        = 250 * time.Millisecond
	connectionBackoffMax        = 2 * time.Minute
	connectionBackoffJitter     = 0.2
	connectionBackoffMultiplier = 1.5
	// payloadMessageBufferSize is the maximum number of payload messages
	// to queue up without having handled previous ones.
	payloadMessageBufferSize = 10
	// sendCredentialsURLParameterName is the name of the URL parameter
	// in the ACS URL that is used to indicate if ACS should send
	// credentials for all tasks on establishing the connection
	sendCredentialsURLParameterName = "sendCredentials"
)

// StartSessionArguments is a struct representing all the things this handler
// needs... This is really a hack to get by-name instead of positional
// arguments since there are too many for positional to be wieldy
type StartSessionArguments struct {
	ContainerInstanceArn          string
	CredentialProvider            *credentials.Credentials
	Config                        *config.Config
	DeregisterInstanceEventStream *eventstream.EventStream
	TaskEngine                    engine.TaskEngine
	ECSClient                     api.ECSClient
	StateManager                  statemanager.StateManager
	AcceptInvalidCert             bool
	CredentialsManager            rolecredentials.Manager
	_time                         ttime.Time
	_heartbeatTimeout             time.Duration
	_heartbeatJitter              time.Duration
	_timeOnce                     sync.Once
}

// sessionState defines state recorder interface for the
// session established with ACS. It can be used to record and
// retrieve data shared across multiple connections to ACS
type sessionState interface {
	// connectedToACS callback indicates that the client has
	// connected to ACS
	connectedToACS()
	// getSendCredentialsURLParameter retrieves the value for
	// the 'sendCredentials' URL parameter
	getSendCredentialsURLParameter() string
}

func (a *StartSessionArguments) time() ttime.Time {
	a.initTime()
	return a._time
}

func (a *StartSessionArguments) heartbeatTimeout() time.Duration {
	a.initTime()
	return a._heartbeatTimeout
}

func (a *StartSessionArguments) heartbeatJitter() time.Duration {
	a.initTime()
	return a._heartbeatJitter
}

func (a *StartSessionArguments) initTime() {
	a._timeOnce.Do(func() {
		if a._time == nil {
			a._time = &ttime.DefaultTime{}
		}
		if a._heartbeatTimeout == 0 {
			a._heartbeatTimeout = heartbeatTimeout
		}
		if a._heartbeatJitter == 0 {
			a._heartbeatJitter = heartbeatJitter
		}
	})
}

// sessionResources defines the resource creator interface for starting
// a session with ACS. This interface is intended to define methods
// that create resources used to establish the connection to ACS
// It is confined to just the createACSClient() method for now. It can be
// extended to include the acsWsURL() and newDisconnectionTimer() methods
// when needed
// The goal is to make it easier to test and inject dependencies
type sessionResources interface {
	// createACSClient creates a new websocket client
	createACSClient(url string) wsclient.ClientServer
	sessionState
}

// acsSessionResources implements resource creator and session state interfaces
// to create resources needed to connect to ACS and to record session state
// for the same
type acsSessionResources struct {
	startSessionArguments StartSessionArguments
	// sendCredentials is used to set the 'sendCredentials' URL parameter
	// used to connect to ACS
	// It is set to 'true' for the very first successful connection on
	// agent start. It is set to false for all successive connections
	sendCredentials bool
}

// StartSession creates a session with ACS and handles requests from ACS.
// It creates resources required to invoke the package scoped 'startSession()'
// method and invokes the same to repeatedly connect to ACS when disconnected
func StartSession(ctx context.Context, args StartSessionArguments) error {
	backoff := utils.NewSimpleBackoff(connectionBackoffMin, connectionBackoffMax, connectionBackoffJitter, connectionBackoffMultiplier)
	session := newSessionResources(args)
	return startSession(ctx, args, backoff, session)
}

func newSessionResources(args StartSessionArguments) sessionResources {
	return &acsSessionResources{
		startSessionArguments: args,
		sendCredentials:       true,
	}
}

// startSession creates a session with ACS and handles requests from ACS
// It also tries to repeatedly connect to ACS when disconnected
func startSession(ctx context.Context, args StartSessionArguments, backoff *utils.SimpleBackoff, acsResources sessionResources) error {
	for {
		acsError := startSessionOnce(ctx, args, backoff, acsResources)
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
		}

		if acsError == nil || acsError == io.EOF {
			backoff.Reset()
		} else if strings.HasPrefix(acsError.Error(), "InactiveInstanceException:") {
			seelog.Debug("Container instance is deregistered, notifying listeners")
			err := args.DeregisterInstanceEventStream.WriteToEventStream(struct{}{})
			if err != nil {
				seelog.Debugf("Failed to write to deregister container instance event stream, err: %v", err)
			}
		} else {
			seelog.Infof("Error from acs; backing off, err: %v", acsError)
			args.time().Sleep(backoff.Duration())
		}
	}
}

// startSessionOnce creates a session with ACS and handles requests using the passed
// in arguments
func startSessionOnce(ctx context.Context, args StartSessionArguments, backoff *utils.SimpleBackoff, acsResources sessionResources) error {
	acsEndpoint, err := args.ECSClient.DiscoverPollEndpoint(args.ContainerInstanceArn)
	if err != nil {
		seelog.Errorf("Unable to discover poll endpoint, err: %v", err)
		return err
	}

	cfg := args.Config
	url := acsWsURL(acsEndpoint, cfg.Cluster, args.ContainerInstanceArn, args.TaskEngine, acsResources)
	client := acsResources.createACSClient(url)
	defer client.Close()

	// Start inactivity timer for closing the connection
	timer := newDisconnectionTimer(client, args.time(), args.heartbeatTimeout(), args.heartbeatJitter())
	defer timer.Stop()

	return startACSSession(ctx, client, timer, args, backoff, acsResources)
}

// acsWsURL returns the websocket url for ACS given the endpoint
func acsWsURL(endpoint, cluster, containerInstanceArn string, taskEngine engine.TaskEngine, acsSessionState sessionState) string {
	acsUrl := endpoint
	if endpoint[len(endpoint)-1] != '/' {
		acsUrl += "/"
	}
	acsUrl += "ws"
	query := url.Values{}
	query.Set("clusterArn", cluster)
	query.Set("containerInstanceArn", containerInstanceArn)
	query.Set("agentHash", version.GitHashString())
	query.Set("agentVersion", version.Version)
	query.Set("seqNum", "1")
	if dockerVersion, err := taskEngine.Version(); err == nil {
		query.Set("dockerVersion", dockerVersion)
	}
	query.Set(sendCredentialsURLParameterName, acsSessionState.getSendCredentialsURLParameter())
	return acsUrl + "?" + query.Encode()
}

// createACSClient creates the ACS Client using the specified URL
func (acsResources *acsSessionResources) createACSClient(url string) wsclient.ClientServer {
	args := acsResources.startSessionArguments
	cfg := args.Config
	return acsclient.New(url, cfg.AWSRegion, args.CredentialProvider, args.AcceptInvalidCert)
}

// connectedToACS records a successful connection to ACS
// It sets sendCredentials to false on such an event
func (acsResources *acsSessionResources) connectedToACS() {
	acsResources.sendCredentials = false
}

// getSendCredentialsURLParameter gets the value to be set for the
// 'sendCredentials' URL parameter
func (acsResources *acsSessionResources) getSendCredentialsURLParameter() string {
	return strconv.FormatBool(acsResources.sendCredentials)
}

// newDisconnectionTimer creates a new time object, with a callback to
// disconnect from ACS on inactivity
func newDisconnectionTimer(client wsclient.ClientServer, _time ttime.Time, timeout time.Duration, jitter time.Duration) ttime.Timer {
	timer := _time.AfterFunc(utils.AddJitter(timeout, jitter), func() {
		seelog.Warn("ACS Connection hasn't had any activity for too long; closing connection")
		closeErr := client.Close()
		if closeErr != nil {
			seelog.Warnf("Error disconnecting: %v", closeErr)
		}
	})

	return timer
}

// startACSSession starts a session with ACS. It adds request handlers for various
// kinds of messages expected from ACS. It returns on server disconnection or when
// the context is cancelled
func startACSSession(ctx context.Context, client wsclient.ClientServer, timer ttime.Timer, args StartSessionArguments, backoff *utils.SimpleBackoff, acsSessionState sessionState) error {
	// Any message from the server resets the disconnect timeout
	client.SetAnyRequestHandler(anyMessageHandler(timer))
	cfg := args.Config

	refreshCredsHandler := newRefreshCredentialsHandler(ctx, cfg.Cluster, args.ContainerInstanceArn, client, args.CredentialsManager, args.TaskEngine)
	defer refreshCredsHandler.clearAcks()
	refreshCredsHandler.start()
	client.AddRequestHandler(refreshCredsHandler.handlerFunc())

	// Add request handler for handling payload messages from ACS
	payloadHandler := newPayloadRequestHandler(ctx, args.TaskEngine, args.ECSClient, cfg.Cluster, args.ContainerInstanceArn, client, args.StateManager, refreshCredsHandler, args.CredentialsManager)
	// Clear the acks channel on return because acks of messageids don't have any value across sessions
	defer payloadHandler.clearAcks()
	payloadHandler.start()
	client.AddRequestHandler(payloadHandler.handlerFunc())

	// Ignore heartbeat messages; anyMessageHandler gets 'em
	client.AddRequestHandler(func(*ecsacs.HeartbeatMessage) {})

	updater.AddAgentUpdateHandlers(client, cfg, args.StateManager, args.TaskEngine)

	err := client.Connect()
	if err != nil {
		seelog.Errorf("Error connecting to ACS: %v", err)
		return err
	}
	acsSessionState.connectedToACS()

	backoffResetTimer := args.time().AfterFunc(utils.AddJitter(args.heartbeatTimeout(), args.heartbeatJitter()), func() {
		// If we do not have an error connecting and remain connected for at
		// least 5 or so minutes, reset the backoff. This prevents disconnect
		// errors that only happen infrequently from damaging the
		// reconnectability as significantly.
		backoff.Reset()
	})
	defer backoffResetTimer.Stop()

	serveErr := make(chan error, 1)
	go func() {
		serveErr <- client.Serve()
	}()

	for {
		select {
		case <-ctx.Done():
			// Stop receiving and sending messages from and to ACS when
			// the context received from the main function is canceled
			payloadHandler.stop()
			refreshCredsHandler.stop()
			return ctx.Err()
		case err := <-serveErr:
			// Stop receiving and sending messages from and to ACS when
			// client.Serve returns an error. This can happen when the
			// the connection is closed by ACS or the agent
			payloadHandler.stop()
			refreshCredsHandler.stop()
			return err
		}
	}
}

// anyMessageHandler handles any server message. Any server message means the
// connection is active and thus the heartbeat disconnect should not occur
func anyMessageHandler(timer ttime.Timer) func(interface{}) {
	return func(interface{}) {
		seelog.Debug("ACS activity occured")
		timer.Reset(utils.AddJitter(heartbeatTimeout, heartbeatJitter))
	}
}
<commit_msg>Fix the go routine leak in acs handler<commit_after>// Copyright 2014-2016 Amazon.com, Inc. or its affiliates. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"). You may
// not use this file except in compliance with the License. A copy of the
// License is located at
//
//	http://aws.amazon.com/apache2.0/
//
// or in the "license" file accompanying this file. This file is distributed
// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
// express or implied. See the License for the specific language governing
// permissions and limitations under the License.

// Package handler deals with appropriately reacting to all ACS messages as well
// as maintaining the connection to ACS.
package handler

import (
	"io"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"time"

	"golang.org/x/net/context"

	acsclient "github.com/aws/amazon-ecs-agent/agent/acs/client"
	"github.com/aws/amazon-ecs-agent/agent/acs/model/ecsacs"
	"github.com/aws/amazon-ecs-agent/agent/acs/update_handler"
	"github.com/aws/amazon-ecs-agent/agent/api"
	"github.com/aws/amazon-ecs-agent/agent/config"
	rolecredentials "github.com/aws/amazon-ecs-agent/agent/credentials"
	"github.com/aws/amazon-ecs-agent/agent/engine"
	"github.com/aws/amazon-ecs-agent/agent/eventstream"
	"github.com/aws/amazon-ecs-agent/agent/statemanager"
	"github.com/aws/amazon-ecs-agent/agent/utils"
	"github.com/aws/amazon-ecs-agent/agent/utils/ttime"
	"github.com/aws/amazon-ecs-agent/agent/version"
	"github.com/aws/amazon-ecs-agent/agent/wsclient"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/cihub/seelog"
)

const (
	// heartbeatTimeout is the maximum time to wait between heartbeats
	// without disconnecting
	heartbeatTimeout = 5 * time.Minute
	heartbeatJitter  = 3 * time.Minute

	connectionBackoffMin        = 250 * time.Millisecond
	connectionBackoffMax        = 2 * time.Minute
	connectionBackoffJitter     = 0.2
	connectionBackoffMultiplier = 1.5
	// payloadMessageBufferSize is the maximum number of payload messages
	// to queue up without having handled previous ones.
	payloadMessageBufferSize = 10
	// sendCredentialsURLParameterName is the name of the URL parameter
	// in the ACS URL that is used to indicate if ACS should send
	// credentials for all tasks on establishing the connection
	sendCredentialsURLParameterName = "sendCredentials"
)

// StartSessionArguments is a struct representing all the things this handler
// needs... This is really a hack to get by-name instead of positional
// arguments since there are too many for positional to be wieldy
type StartSessionArguments struct {
	ContainerInstanceArn          string
	CredentialProvider            *credentials.Credentials
	Config                        *config.Config
	DeregisterInstanceEventStream *eventstream.EventStream
	TaskEngine                    engine.TaskEngine
	ECSClient                     api.ECSClient
	StateManager                  statemanager.StateManager
	AcceptInvalidCert             bool
	CredentialsManager            rolecredentials.Manager
	_time                         ttime.Time
	_heartbeatTimeout             time.Duration
	_heartbeatJitter              time.Duration
	_timeOnce                     sync.Once
}

// sessionState defines state recorder interface for the
// session established with ACS. It can be used to record and
// retrieve data shared across multiple connections to ACS
type sessionState interface {
	// connectedToACS callback indicates that the client has
	// connected to ACS
	connectedToACS()
	// getSendCredentialsURLParameter retrieves the value for
	// the 'sendCredentials' URL parameter
	getSendCredentialsURLParameter() string
}

func (a *StartSessionArguments) time() ttime.Time {
	a.initTime()
	return a._time
}

func (a *StartSessionArguments) heartbeatTimeout() time.Duration {
	a.initTime()
	return a._heartbeatTimeout
}

func (a *StartSessionArguments) heartbeatJitter() time.Duration {
	a.initTime()
	return a._heartbeatJitter
}

func (a *StartSessionArguments) initTime() {
	a._timeOnce.Do(func() {
		if a._time == nil {
			a._time = &ttime.DefaultTime{}
		}
		if a._heartbeatTimeout == 0 {
			a._heartbeatTimeout = heartbeatTimeout
		}
		if a._heartbeatJitter == 0 {
			a._heartbeatJitter = heartbeatJitter
		}
	})
}

// sessionResources defines the resource creator interface for starting
// a session with ACS. This interface is intended to define methods
// that create resources used to establish the connection to ACS
// It is confined to just the createACSClient() method for now. It can be
// extended to include the acsWsURL() and newDisconnectionTimer() methods
// when needed
// The goal is to make it easier to test and inject dependencies
type sessionResources interface {
	// createACSClient creates a new websocket client
	createACSClient(url string) wsclient.ClientServer
	sessionState
}

// acsSessionResources implements resource creator and session state interfaces
// to create resources needed to connect to ACS and to record session state
// for the same
type acsSessionResources struct {
	startSessionArguments StartSessionArguments
	// sendCredentials is used to set the 'sendCredentials' URL parameter
	// used to connect to ACS
	// It is set to 'true' for the very first successful connection on
	// agent start. It is set to false for all successive connections
	sendCredentials bool
}

// StartSession creates a session with ACS and handles requests from ACS.
// It creates resources required to invoke the package scoped 'startSession()'
// method and invokes the same to repeatedly connect to ACS when disconnected
func StartSession(ctx context.Context, args StartSessionArguments) error {
	backoff := utils.NewSimpleBackoff(connectionBackoffMin, connectionBackoffMax, connectionBackoffJitter, connectionBackoffMultiplier)
	session := newSessionResources(args)
	return startSession(ctx, args, backoff, session)
}

func newSessionResources(args StartSessionArguments) sessionResources {
	return &acsSessionResources{
		startSessionArguments: args,
		sendCredentials:       true,
	}
}

// startSession creates a session with ACS and handles requests from ACS
// It also tries to repeatedly connect to ACS when disconnected
func startSession(ctx context.Context, args StartSessionArguments, backoff *utils.SimpleBackoff, acsResources sessionResources) error {
	for {
		acsError := startSessionOnce(ctx, args, backoff, acsResources)
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
		}

		if acsError == nil || acsError == io.EOF {
			backoff.Reset()
		} else if strings.HasPrefix(acsError.Error(), "InactiveInstanceException:") {
			seelog.Debug("Container instance is deregistered, notifying listeners")
			err := args.DeregisterInstanceEventStream.WriteToEventStream(struct{}{})
			if err != nil {
				seelog.Debugf("Failed to write to deregister container instance event stream, err: %v", err)
			}
		} else {
			seelog.Infof("Error from acs; backing off, err: %v", acsError)
			args.time().Sleep(backoff.Duration())
		}
	}
}

// startSessionOnce creates a session with ACS and handles requests using the passed
// in arguments
func startSessionOnce(ctx context.Context, args StartSessionArguments, backoff *utils.SimpleBackoff, acsResources sessionResources) error {
	acsEndpoint, err := args.ECSClient.DiscoverPollEndpoint(args.ContainerInstanceArn)
	if err != nil {
		seelog.Errorf("Unable to discover poll endpoint, err: %v", err)
		return err
	}

	cfg := args.Config
	url := acsWsURL(acsEndpoint, cfg.Cluster, args.ContainerInstanceArn, args.TaskEngine, acsResources)
	client := acsResources.createACSClient(url)
	defer client.Close()

	// Start inactivity timer for closing the connection
	timer := newDisconnectionTimer(client, args.time(), args.heartbeatTimeout(), args.heartbeatJitter())
	defer timer.Stop()

	return startACSSession(ctx, client, timer, args, backoff, acsResources)
}

// acsWsURL returns the websocket url for ACS given the endpoint
func acsWsURL(endpoint, cluster, containerInstanceArn string, taskEngine engine.TaskEngine, acsSessionState sessionState) string {
	acsUrl := endpoint
	if endpoint[len(endpoint)-1] != '/' {
		acsUrl += "/"
	}
	acsUrl += "ws"
	query := url.Values{}
	query.Set("clusterArn", cluster)
	query.Set("containerInstanceArn", containerInstanceArn)
	query.Set("agentHash", version.GitHashString())
	query.Set("agentVersion", version.Version)
	query.Set("seqNum", "1")
	if dockerVersion, err := taskEngine.Version(); err == nil {
		query.Set("dockerVersion", dockerVersion)
	}
	query.Set(sendCredentialsURLParameterName, acsSessionState.getSendCredentialsURLParameter())
	return acsUrl + "?" + query.Encode()
}

// createACSClient creates the ACS Client using the specified URL
func (acsResources *acsSessionResources) createACSClient(url string) wsclient.ClientServer {
	args := acsResources.startSessionArguments
	cfg := args.Config
	return acsclient.New(url, cfg.AWSRegion, args.CredentialProvider, args.AcceptInvalidCert)
}

// connectedToACS records a successful connection to ACS
// It sets sendCredentials to false on such an event
func (acsResources *acsSessionResources) connectedToACS() {
	acsResources.sendCredentials = false
}

// getSendCredentialsURLParameter gets the value to be set for the
// 'sendCredentials' URL parameter
func (acsResources *acsSessionResources) getSendCredentialsURLParameter() string {
	return strconv.FormatBool(acsResources.sendCredentials)
}

// newDisconnectionTimer creates a new time object, with a callback to
// disconnect from ACS on inactivity
func newDisconnectionTimer(client wsclient.ClientServer, _time ttime.Time, timeout time.Duration, jitter time.Duration) ttime.Timer {
	timer := _time.AfterFunc(utils.AddJitter(timeout, jitter), func() {
		seelog.Warn("ACS Connection hasn't had any activity for too long; closing connection")
		closeErr := client.Close()
		if closeErr != nil {
			seelog.Warnf("Error disconnecting: %v", closeErr)
		}
	})

	return timer
}

// startACSSession starts a session with ACS. It adds request handlers for various
// kinds of messages expected from ACS. It returns on server disconnection or when
// the context is cancelled
func startACSSession(ctx context.Context, client wsclient.ClientServer, timer ttime.Timer, args StartSessionArguments, backoff *utils.SimpleBackoff, acsSessionState sessionState) error {
	// Any message from the server resets the disconnect timeout
	client.SetAnyRequestHandler(anyMessageHandler(timer))
	cfg := args.Config

	refreshCredsHandler := newRefreshCredentialsHandler(ctx, cfg.Cluster, args.ContainerInstanceArn, client, args.CredentialsManager, args.TaskEngine)
	defer refreshCredsHandler.clearAcks()
	refreshCredsHandler.start()
	defer refreshCredsHandler.stop()

	client.AddRequestHandler(refreshCredsHandler.handlerFunc())

	// Add request handler for handling payload messages from ACS
	payloadHandler := newPayloadRequestHandler(ctx, args.TaskEngine, args.ECSClient, cfg.Cluster, args.ContainerInstanceArn, client, args.StateManager, refreshCredsHandler, args.CredentialsManager)
	// Clear the acks channel on return because acks of messageids don't have any value across sessions
	defer payloadHandler.clearAcks()
	payloadHandler.start()
	defer payloadHandler.stop()

	client.AddRequestHandler(payloadHandler.handlerFunc())

	// Ignore heartbeat messages; anyMessageHandler gets 'em
	client.AddRequestHandler(func(*ecsacs.HeartbeatMessage) {})

	updater.AddAgentUpdateHandlers(client, cfg, args.StateManager, args.TaskEngine)

	err := client.Connect()
	if err != nil {
		seelog.Errorf("Error connecting to ACS: %v", err)
		return err
	}
	acsSessionState.connectedToACS()

	backoffResetTimer := args.time().AfterFunc(utils.AddJitter(args.heartbeatTimeout(), args.heartbeatJitter()), func() {
		// If we do not have an error connecting and remain connected for at
		// least 5 or so minutes, reset the backoff. This prevents disconnect
		// errors that only happen infrequently from damaging the
		// reconnectability as significantly.
		backoff.Reset()
	})
	defer backoffResetTimer.Stop()

	serveErr := make(chan error, 1)
	go func() {
		serveErr <- client.Serve()
	}()

	for {
		select {
		case <-ctx.Done():
			// Stop receiving and sending messages from and to ACS when
			// the context received from the main function is canceled
			return ctx.Err()
		case err := <-serveErr:
			// Stop receiving and sending messages from and to ACS when
			// client.Serve returns an error. This can happen when the
			// the connection is closed by ACS or the agent
			return err
		}
	}
}

// anyMessageHandler handles any server message. Any server message means the
// connection is active and thus the heartbeat disconnect should not occur
func anyMessageHandler(timer ttime.Timer) func(interface{}) {
	return func(interface{}) {
		seelog.Debug("ACS activity occured")
		timer.Reset(utils.AddJitter(heartbeatTimeout, heartbeatJitter))
	}
}
<|endoftext|>


<commit_before>package szah

import (
	"os"
	"path/filepath"
	"runtime"
	"strings"

	humanize "github.com/dustin/go-humanize"
	"github.com/go-errors/errors"
	"github.com/itchio/butler/archive"
	"github.com/itchio/sevenzip-go/sz"
)

type ExtractState struct {
	HasListedItems        bool
	ItemCount             int64
	TotalDoneSize         int64
	TotalUncompressedSize int64
	CurrentIndex          int64
	Contents              *archive.Contents

	NumFiles    int64
	NumDirs     int64
	NumSymlinks int64
}

type ech struct {
	params          *archive.ExtractParams
	initialProgress float64
	state           *ExtractState
	save            archive.ThrottledSaveFunc
}

func (h *Handler) Extract(params *archive.ExtractParams) (*archive.Contents, error) {
	save := archive.ThrottledSave(params)
	consumer := params.Consumer
	state := &ExtractState{
		Contents: &archive.Contents{},
	}

	err := withArchive(params.Consumer, params.File, func(a *sz.Archive) error {
		err := params.Load(state)
		if err != nil {
			consumer.Infof("szah: could not load state: %s", err.Error())
			consumer.Infof("szah: ...starting from beginning!")
		}

		if !state.HasListedItems {
			consumer.Infof("szah: listing items")
			itemCount, err := a.GetItemCount()
			if err != nil {
				return errors.Wrap(err, 0)
			}
			state.ItemCount = itemCount

			var totalUncompressedSize int64
			for i := int64(0); i < itemCount; i++ {
				func() {
					item := a.GetItem(i)
					if item == nil {
						return
					}
					defer item.Free()

					if item.GetBoolProperty(sz.PidIsDir) {
						return
					}

					itemSize := item.GetUInt64Property(sz.PidSize)
					totalUncompressedSize += int64(itemSize)
				}()
			}
			state.TotalUncompressedSize = totalUncompressedSize

			state.HasListedItems = true
			save(state, true)
		} else {
			consumer.Infof("szah: using cached item listing")
		}

		if params.OnUncompressedSizeKnown != nil {
			params.OnUncompressedSizeKnown(state.TotalUncompressedSize)
		}

		ec, err := sz.NewExtractCallback(&ech{
			params:          params,
			state:           state,
			initialProgress: float64(state.TotalDoneSize) / float64(state.TotalUncompressedSize),
			save:            save,
		})
		if err != nil {
			return errors.Wrap(err, 0)
		}
		defer ec.Free()

		var indices []int64
		for i := state.CurrentIndex; i < state.ItemCount; i++ {
			indices = append(indices, i)
		}
		if len(indices) == 0 {
			consumer.Infof("nothing (0 items) to extract!")
			return nil
		}

		consumer.Infof("queued %d/%d items for extraction", len(indices), state.ItemCount)

		err = a.ExtractSeveral(indices, ec)
		if err != nil {
			return errors.Wrap(err, 0)
		}

		return nil
	})
	if err != nil {
		return nil, errors.Wrap(err, 0)
	}

	consumer.Statf("extracted %d items successfully", state.ItemCount)
	consumer.Statf("%d files, %d dirs, %d symlinks", state.NumFiles, state.NumDirs, state.NumSymlinks)

	return state.Contents, nil
}

func (e *ech) GetStream(item *sz.Item) (*sz.OutStream, error) {
	consumer := e.params.Consumer

	sanePath := sanitizePath(item.GetStringProperty(sz.PidPath))
	outPath := filepath.Join(e.params.OutputPath, sanePath)

	if item.GetBoolProperty(sz.PidIsDir) {
		err := os.MkdirAll(outPath, 0755)
		if err != nil {
			return nil, errors.Wrap(err, 0)
		}

		e.state.NumDirs++

		return nil, nil
	}

	e.state.NumFiles++

	err := os.MkdirAll(filepath.Dir(outPath), 0755)
	if err != nil {
		return nil, errors.Wrap(err, 0)
	}

	f, err := os.Create(outPath)
	if err != nil {
		return nil, errors.Wrap(err, 0)
	}

	uncompressedSize := item.GetUInt64Property(sz.PidSize)
	consumer.Infof(`→ %s (%s)`, sanePath, humanize.IBytes(uncompressedSize))

	contents := e.state.Contents

	nc := &notifyCloser{
		Writer: f,
		OnClose: func(totalBytes int64) {
			contents.Entries = append(contents.Entries, &archive.Entry{
				Name:             sanePath,
				UncompressedSize: totalBytes,
			})
			// FIXME: it'd be better for GetStream to give us the index of the entry
			// or for Item to have an index getter
			e.state.CurrentIndex++
			e.state.TotalDoneSize += totalBytes
			e.save(e.state, false)
		},
	}

	return sz.NewOutStream(nc)
}

func (e *ech) SetProgress(complete int64, total int64) {
	if total > 0 {
		thisRunProgress := float64(complete) / float64(total)
		actualProgress := e.initialProgress + (1.0-e.initialProgress)*thisRunProgress
		e.params.Consumer.Progress(actualProgress)
	}
	// TODO: do something smart for other formats ?
}

func sanitizePath(inPath string) string {
	outPath := filepath.ToSlash(inPath)

	if runtime.GOOS == "windows" {
		// Remove illegal character for windows paths, see
		// https://msdn.microsoft.com/en-us/library/windows/desktop/aa365247(v=vs.85).aspx
		for i := byte(0); i <= 31; i++ {
			outPath = strings.Replace(outPath, string([]byte{i}), "_", -1)
		}
	}

	return outPath
}
<commit_msg>Decode 'attrib' for 7-zip entries<commit_after>package szah

import (
	"os"
	"bytes"
	"io"
	"path/filepath"
	"runtime"
	"strings"

	"github.com/go-errors/errors"
	"github.com/itchio/butler/archive"
	"github.com/itchio/sevenzip-go/sz"
)

type ExtractState struct {
	HasListedItems        bool
	ItemCount             int64
	TotalDoneSize         int64
	TotalUncompressedSize int64
	CurrentIndex          int64
	Contents              *archive.Contents

	NumFiles    int64
	NumDirs     int64
	NumSymlinks int64
}

type ech struct {
	params          *archive.ExtractParams
	initialProgress float64
	state           *ExtractState
	save            archive.ThrottledSaveFunc
}

func (h *Handler) Extract(params *archive.ExtractParams) (*archive.Contents, error) {
	save := archive.ThrottledSave(params)
	consumer := params.Consumer
	state := &ExtractState{
		Contents: &archive.Contents{},
	}

	err := withArchive(params.Consumer, params.File, func(a *sz.Archive) error {
		err := params.Load(state)
		if err != nil {
			consumer.Infof("szah: could not load state: %s", err.Error())
			consumer.Infof("szah: ...starting from beginning!")
		}

		if !state.HasListedItems {
			consumer.Infof("szah: listing items")
			itemCount, err := a.GetItemCount()
			if err != nil {
				return errors.Wrap(err, 0)
			}
			state.ItemCount = itemCount

			var totalUncompressedSize int64
			for i := int64(0); i < itemCount; i++ {
				func() {
					item := a.GetItem(i)
					if item == nil {
						return
					}
					defer item.Free()

					if item.GetBoolProperty(sz.PidIsDir) {
						return
					}

					itemSize := item.GetUInt64Property(sz.PidSize)
					totalUncompressedSize += int64(itemSize)
				}()
			}
			state.TotalUncompressedSize = totalUncompressedSize

			state.HasListedItems = true
			save(state, true)
		} else {
			consumer.Infof("szah: using cached item listing")
		}

		if params.OnUncompressedSizeKnown != nil {
			params.OnUncompressedSizeKnown(state.TotalUncompressedSize)
		}

		ec, err := sz.NewExtractCallback(&ech{
			params:          params,
			state:           state,
			initialProgress: float64(state.TotalDoneSize) / float64(state.TotalUncompressedSize),
			save:            save,
		})
		if err != nil {
			return errors.Wrap(err, 0)
		}
		defer ec.Free()

		var indices []int64
		for i := state.CurrentIndex; i < state.ItemCount; i++ {
			indices = append(indices, i)
		}
		if len(indices) == 0 {
			consumer.Infof("nothing (0 items) to extract!")
			return nil
		}

		consumer.Infof("queued %d/%d items for extraction", len(indices), state.ItemCount)

		err = a.ExtractSeveral(indices, ec)
		if err != nil {
			return errors.Wrap(err, 0)
		}

		return nil
	})
	if err != nil {
		return nil, errors.Wrap(err, 0)
	}

	consumer.Statf("extracted %d items successfully", state.ItemCount)
	consumer.Statf("%d files, %d dirs, %d symlinks", state.NumFiles, state.NumDirs, state.NumSymlinks)

	return state.Contents, nil
}

const (
	TypeDir uint64 = 0x4
	TypeFile = 0x8
	TypeSymlink = 0xa
)

func (e *ech) GetStream(item *sz.Item) (*sz.OutStream, error) {
	consumer := e.params.Consumer

	sanePath := sanitizePath(item.GetStringProperty(sz.PidPath))
	outPath := filepath.Join(e.params.OutputPath, sanePath)

	attr := item.GetUInt64Property(sz.PidAttrib)
	var typemask uint64 = 0xf0000000
	var modemask uint64 = 0x0fff0000
	itemType := (attr&typemask)>>(7*4)
	itemMode := (attr&modemask)>>(4*4)
	formatType := func (itemType uint64) string {
		switch itemType {
		case TypeDir: return "dir "
		case TypeFile: return "file"
		case TypeSymlink: return "link"
		default: return "????"
		}
	}
	consumer.Infof(`attrib: %032x %s (%04o) %s`, attr, formatType(itemType), itemMode, sanePath)

	if item.GetBoolProperty(sz.PidIsDir) {
		err := os.MkdirAll(outPath, 0755)
		if err != nil {
			return nil, errors.Wrap(err, 0)
		}

		e.state.NumDirs++

		return nil, nil
	}

	e.state.NumFiles++

	err := os.MkdirAll(filepath.Dir(outPath), 0755)
	if err != nil {
		return nil, errors.Wrap(err, 0)
	}

	var f io.WriteCloser
	var buf *bytes.Buffer

	if itemType == TypeSymlink {
		buf = new(bytes.Buffer)
		f = &nopWriteCloser{
			writer: buf,
		}
	} else {
		f, err = os.Create(outPath)
		if err != nil {
			return nil, errors.Wrap(err, 0)
		}
	}

	// uncompressedSize := item.GetUInt64Property(sz.PidSize)
	// consumer.Infof(`→ %s (%s)`, sanePath, humanize.IBytes(uncompressedSize))

	contents := e.state.Contents

	nc := &notifyCloser{
		Writer: f,
		OnClose: func(totalBytes int64) {
			if itemType == TypeSymlink {
				consumer.Infof("symlinks target: %s", string(buf.Bytes()))
			}

			contents.Entries = append(contents.Entries, &archive.Entry{
				Name:             sanePath,
				UncompressedSize: totalBytes,
			})
			// FIXME: it'd be better for GetStream to give us the index of the entry
			// or for Item to have an index getter
			e.state.CurrentIndex++
			e.state.TotalDoneSize += totalBytes
			e.save(e.state, false)
		},
	}

	return sz.NewOutStream(nc)
}

func (e *ech) SetProgress(complete int64, total int64) {
	if total > 0 {
		thisRunProgress := float64(complete) / float64(total)
		actualProgress := e.initialProgress + (1.0-e.initialProgress)*thisRunProgress
		e.params.Consumer.Progress(actualProgress)
	}
	// TODO: do something smart for other formats ?
}

func sanitizePath(inPath string) string {
	outPath := filepath.ToSlash(inPath)

	if runtime.GOOS == "windows" {
		// Remove illegal character for windows paths, see
		// https://msdn.microsoft.com/en-us/library/windows/desktop/aa365247(v=vs.85).aspx
		for i := byte(0); i <= 31; i++ {
			outPath = strings.Replace(outPath, string([]byte{i}), "_", -1)
		}
	}

	return outPath
}

// nopWriteCloser

type nopWriteCloser struct {
	writer io.Writer
}

var _ io.Writer = (*nopWriteCloser)(nil)

func (nwc *nopWriteCloser) Write(data []byte) (int, error) {
	return nwc.writer.Write(data)
}

func (nwc *nopWriteCloser) Close() error {
	return nil
}
<|endoftext|>