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

package main

import (
	"bytes"
	"flag"
	"fmt"
	"log"
	"os"
	"regexp"
	"strconv"
	"strings"
	"sync"

	"code.google.com/p/goplan9/draw"
	"code.google.com/p/goplan9/plan9/acme"
)

func acmeMode() {
	var dummy awin
	dummy.prefix = "/issue/" + *project + "/"
	if flag.NArg() > 0 {
		for _, arg := range flag.Args() {
			dummy.look(arg)
		}
	} else {
		dummy.look("all")
	}
	select {}
}

const (
	modeSingle = 1 + iota
	modeQuery
)

type awin struct {
	*acme.Win
	prefix   string
	mode     int
	query    string
	id       int
	data     []Entry
	tab      int
	font     *draw.Font
	fontName string
	title    string
}

var all struct {
	sync.Mutex
	m      map[string]*awin
	f      map[string]*draw.Font
	numwin int
}

func (w *awin) exit() {
	all.Lock()
	defer all.Unlock()
	if all.m[w.title] == w {
		delete(all.m, w.title)
	}
	if all.numwin--; all.numwin == 0 {
		os.Exit(0)
	}
}

func (w *awin) new(title string) *awin {
	all.Lock()
	defer all.Unlock()
	all.numwin++
	if all.m == nil {
		all.m = make(map[string]*awin)
	}
	w1 := new(awin)
	w1.title = title
	var err error
	w1.Win, err = acme.New()
	if err != nil {
		log.Fatal(err)
	}
	w1.prefix = w.prefix
	w1.Name(w1.prefix + title)
	all.m[title] = w1
	return w1
}

func (w *awin) show(title string) *awin {
	all.Lock()
	defer all.Unlock()
	if w1 := all.m[title]; w1 != nil {
		w.Ctl("show")
		return w1
	}
	return nil
}

func (w *awin) fixfont() {
	ctl := make([]byte, 1000)
	w.Seek("ctl", 0, 0)
	n, err := w.Read("ctl", ctl)
	if err != nil {
		return
	}
	f := strings.Fields(string(ctl[:n]))
	if len(f) < 8 {
		return
	}
	w.tab, _ = strconv.Atoi(f[7])
	if w.tab == 0 {
		return
	}
	name := f[6]
	if w.fontName == name {
		return
	}
	all.Lock()
	defer all.Unlock()
	if font := all.f[name]; font != nil {
		w.font = font
		w.fontName = name
		return
	}
	var disp *draw.Display = nil
	font, err := disp.OpenFont(name)
	if err != nil {
		return
	}
	if all.f == nil {
		all.f = make(map[string]*draw.Font)
	}
	all.f[name] = font
	w.font = font
}

var numRE = regexp.MustCompile(`(?m)^[0-9]+\t`)

func (w *awin) look(text string) bool {
	if text == "all" {
		if w.show("all") != nil {
			return true
		}
		w.newSearch("all", "-label:nonexistent")
		return true
	}
	if n, _ := strconv.Atoi(text); 0 < n && n < 100000 {
		if w.show(text) != nil {
			return true
		}
		w.newIssue(text, n)
		return true
	}
	if m := numRE.FindAllString(text, -1); m != nil {
		for _, s := range m {
			w.look(strings.TrimSpace(s))
		}
		return true
	}
	return false
}

func (w *awin) newIssue(title string, id int) {
	w = w.new(title)
	w.mode = modeSingle
	w.id = id
	w.Ctl("cleartag")
	w.Fprintf("tag", " Get Put Look ")
	go w.load()
}

func (w *awin) newSearch(title, query string) {
	w = w.new(title)
	w.mode = modeQuery
	w.query = query
	w.Ctl("cleartag")
	w.Fprintf("tag", " Search ")
	w.Write("body", []byte("Loading..."))
	go w.load()
}

func (w *awin) clear() {
	w.Addr(",")
	w.Write("data", nil)
}

func (w *awin) load() {
	w.fixfont()

	switch w.mode {
	case modeSingle:
		data, err := fetch(fmt.Sprintf("id:%d", w.id), true)
		w.clear()
		if err != nil {
			w.Write("body", []byte(err.Error()))
			break
		}
		var buf bytes.Buffer
		printFull(&buf, data)
		fmt.Fprintf(&buf, "\nNew Comment:\n\n")
		w.Write("body", buf.Bytes())
		w.Ctl("clean")
		w.data = data
	case modeQuery:
		w.query = strings.Replace(w.query, ": ", ":", -1)
		w.query = strings.TrimSpace(w.query)
		data, err := fetch(w.query, false)
		w.clear()
		if err != nil {
			w.Write("body", []byte(err.Error()))
			break
		}
		if w.title == "search" {
			w.Fprintf("body", "Search: %s\n\n", w.query)
		}
		w.printList(data)
		w.Ctl("clean")
		w.data = data
	}
	w.Addr("0")
	w.Ctl("dot=addr")
	w.Ctl("show")

	go w.loop()
}

func (w *awin) err(s string) {
	if !strings.HasSuffix(s, "\n") {
		s = s + "\n"
	}
	w1 := w.show("+Errors")
	if w1 == nil {
		w1 = w.new("+Errors")
	}
	w1.Fprintf("body", "%s", s)
	w1.Addr("$")
	w1.Ctl("dot=addr")
	w1.Ctl("show")
}

func diff(line, field, old string) string {
	old = strings.TrimSpace(old)
	line = strings.TrimSpace(strings.TrimPrefix(line, field))
	if old == line {
		return ""
	}
	return line
}

func diffUpdate(line, field, old string) string {
	old = strings.TrimSpace(old)
	line = strings.TrimSpace(strings.TrimPrefix(line, field))
	if old == line {
		return ""
	}
	if line == "" {
		return "-" + old
	}
	return line
}

func diffUpdates(line, field string, old []string) []string {
	line = strings.TrimSpace(strings.TrimPrefix(line, field))
	had := make(map[string]bool)
	for _, f := range old {
		had[f] = true
	}
	var d []string
	kept := make(map[string]bool)
	for _, f := range strings.Fields(line) {
		f = strings.TrimSuffix(f, ",")
		if had[f] {
			kept[f] = true
		} else {
			d = append(d, f)
		}
	}
	for _, f := range old {
		if had[f] && !kept[f] {
			d = append(d, "-"+f)
		}
	}
	return d
}

func (w *awin) put() {
	switch w.mode {
	case modeSingle:
		e := w.data[0]
		var ch Change
		data, err := w.ReadAll("body")
		if err != nil {
			w.err(fmt.Sprintf("Put: %v", err))
			return
		}
		sdata := string(data)
		off := 0
		for _, line := range strings.SplitAfter(sdata, "\n") {
			off += len(line)
			line = strings.TrimSpace(line)
			if line == "" {
				break
			}
			switch {
			case strings.HasPrefix(line, "Summary:"):
				ch.Summary = diff(line, "Summary:", e.Title)

			case strings.HasPrefix(line, "Status:"):
				status := e.Status
				if status == "Duplicate" {
					status += " " + e.MergedInto
				}
				ch.Status = diff(line, "Status:", status)

			case strings.HasPrefix(line, "Owner:"):
				ch.Owner = diffUpdate(line, "Owner:", e.Owner)

			case strings.HasPrefix(line, "CC:"):
				ch.CC = diffUpdates(line, "CC:", e.CC)

			case strings.HasPrefix(line, "Labels:"):
				ch.Label = diffUpdates(line, "Labels:", e.Label)

			default:
				w.err(fmt.Sprintf("Put: unknown summary line: %s", line))
			}
		}

		i := strings.Index(sdata, "\nReported by ")
		if i >= off {
			ch.Comment = strings.TrimSpace(sdata[off:i])
		}

		tag := "\nNew Comment:"
		i = strings.Index(sdata, tag)
		if i >= 0 {
			s := strings.TrimSpace(sdata[i+len(tag):])
			if s != "" {
				if ch.Comment != "" {
					ch.Comment += "\n\n"
				}
				ch.Comment = s
			}
		}
		if err := write(w.id, &ch); err != nil {
			w.err(fmt.Sprintf("Put: %v", err))
			return
		}
		w.load()

	case modeQuery:
		w.err("cannot Put issue list")
	}
}

func (w *awin) loop() {
	defer w.exit()
	for e := range w.EventChan() {
		switch e.C2 {
		case 'x', 'X': // execute
			cmd := strings.TrimSpace(string(e.Text))
			if cmd == "Get" {
				w.load()
				break
			}
			if cmd == "Put" {
				w.put()
				break
			}
			if cmd == "Del" {
				w.Ctl("del")
				break
			}
			if strings.HasPrefix(cmd, "Search ") {
				w.newSearch("search", strings.TrimSpace(strings.TrimPrefix(cmd, "Search")))
				break
			}
			w.WriteEvent(e)
		case 'l', 'L': // look
			if len(e.Text) == 0 && e.Q0 < e.Q1 {
				w.Addr("#%d,#%d", e.Q0, e.Q1)
				data, err := w.ReadAll("xdata")
				if err != nil {
					w.err(err.Error())
				}
				e.Text = data
			}
			if !w.look(string(e.Text)) {
				w.WriteEvent(e)
			}
		}
	}
}

func (w *awin) printTabbed(rows [][]string) {
	var wid []int

	if w.font != nil {
		for _, row := range rows {
			for len(wid) < len(row) {
				wid = append(wid, 0)
			}
			for i, col := range row {
				n := w.font.StringWidth(col)
				if wid[i] < n {
					wid[i] = n
				}
			}
		}
	}

	var buf bytes.Buffer
	for _, row := range rows {
		for i, col := range row {
			buf.WriteString(col)
			if i == len(row)-1 {
				break
			}
			if w.font == nil || w.tab == 0 {
				buf.WriteString("\t")
				continue
			}
			pos := w.font.StringWidth(col)
			for pos <= wid[i] {
				buf.WriteString("\t")
				pos += w.tab - pos%w.tab
			}
		}
		buf.WriteString("\n")
	}

	w.Write("body", buf.Bytes())
}

func (w *awin) printList(entries []Entry) {
	var rows [][]string
	var buf bytes.Buffer
	for _, e := range entries {
		id := e.ID
		if i := strings.Index(id, "id="); i >= 0 {
			id = id[:i+len("id=")]
		}
		buf.Reset()
		buf.WriteString(e.Title)
		if len(e.Label) > 0 {
			buf.WriteString(" [")
			for i, l := range e.Label {
				if i > 0 {
					buf.WriteString(" ")
				}
				buf.WriteString(l)
			}
			buf.WriteString("]")
		}
		rows = append(rows, []string{id, buf.String()})
	}
	w.printTabbed(rows)
}
<commit_msg>issue: fix issue -a with command-line query<commit_after>// Copyright 2013 The Go Authors.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"bytes"
	"flag"
	"fmt"
	"log"
	"os"
	"regexp"
	"strconv"
	"strings"
	"sync"

	"code.google.com/p/goplan9/draw"
	"code.google.com/p/goplan9/plan9/acme"
)

func acmeMode() {
	var dummy awin
	dummy.prefix = "/issue/" + *project + "/"
	if flag.NArg() > 0 {
		for _, arg := range flag.Args() {
			if !dummy.look(arg) {
				dummy.newSearch("search", arg)
			}
		}
	} else {
		dummy.look("all")
	}
	select {}
}

const (
	modeSingle = 1 + iota
	modeQuery
)

type awin struct {
	*acme.Win
	prefix   string
	mode     int
	query    string
	id       int
	data     []Entry
	tab      int
	font     *draw.Font
	fontName string
	title    string
}

var all struct {
	sync.Mutex
	m      map[string]*awin
	f      map[string]*draw.Font
	numwin int
}

func (w *awin) exit() {
	all.Lock()
	defer all.Unlock()
	if all.m[w.title] == w {
		delete(all.m, w.title)
	}
	if all.numwin--; all.numwin == 0 {
		os.Exit(0)
	}
}

func (w *awin) new(title string) *awin {
	all.Lock()
	defer all.Unlock()
	all.numwin++
	if all.m == nil {
		all.m = make(map[string]*awin)
	}
	w1 := new(awin)
	w1.title = title
	var err error
	w1.Win, err = acme.New()
	if err != nil {
		log.Fatal(err)
	}
	w1.prefix = w.prefix
	w1.Name(w1.prefix + title)
	all.m[title] = w1
	return w1
}

func (w *awin) show(title string) *awin {
	all.Lock()
	defer all.Unlock()
	if w1 := all.m[title]; w1 != nil {
		w.Ctl("show")
		return w1
	}
	return nil
}

func (w *awin) fixfont() {
	ctl := make([]byte, 1000)
	w.Seek("ctl", 0, 0)
	n, err := w.Read("ctl", ctl)
	if err != nil {
		return
	}
	f := strings.Fields(string(ctl[:n]))
	if len(f) < 8 {
		return
	}
	w.tab, _ = strconv.Atoi(f[7])
	if w.tab == 0 {
		return
	}
	name := f[6]
	if w.fontName == name {
		return
	}
	all.Lock()
	defer all.Unlock()
	if font := all.f[name]; font != nil {
		w.font = font
		w.fontName = name
		return
	}
	var disp *draw.Display = nil
	font, err := disp.OpenFont(name)
	if err != nil {
		return
	}
	if all.f == nil {
		all.f = make(map[string]*draw.Font)
	}
	all.f[name] = font
	w.font = font
}

var numRE = regexp.MustCompile(`(?m)^[0-9]+\t`)

func (w *awin) look(text string) bool {
	if text == "all" {
		if w.show("all") != nil {
			return true
		}
		w.newSearch("all", "-label:nonexistent")
		return true
	}
	if n, _ := strconv.Atoi(text); 0 < n && n < 100000 {
		if w.show(text) != nil {
			return true
		}
		w.newIssue(text, n)
		return true
	}
	if m := numRE.FindAllString(text, -1); m != nil {
		for _, s := range m {
			w.look(strings.TrimSpace(s))
		}
		return true
	}
	return false
}

func (w *awin) newIssue(title string, id int) {
	w = w.new(title)
	w.mode = modeSingle
	w.id = id
	w.Ctl("cleartag")
	w.Fprintf("tag", " Get Put Look ")
	go w.load()
}

func (w *awin) newSearch(title, query string) {
	w = w.new(title)
	w.mode = modeQuery
	w.query = query
	w.Ctl("cleartag")
	w.Fprintf("tag", " Search ")
	w.Write("body", []byte("Loading..."))
	go w.load()
}

func (w *awin) clear() {
	w.Addr(",")
	w.Write("data", nil)
}

func (w *awin) load() {
	w.fixfont()

	switch w.mode {
	case modeSingle:
		data, err := fetch(fmt.Sprintf("id:%d", w.id), true)
		w.clear()
		if err != nil {
			w.Write("body", []byte(err.Error()))
			break
		}
		var buf bytes.Buffer
		printFull(&buf, data)
		fmt.Fprintf(&buf, "\nNew Comment:\n\n")
		w.Write("body", buf.Bytes())
		w.Ctl("clean")
		w.data = data
	case modeQuery:
		w.query = strings.Replace(w.query, ": ", ":", -1)
		w.query = strings.TrimSpace(w.query)
		data, err := fetch(w.query, false)
		w.clear()
		if err != nil {
			w.Write("body", []byte(err.Error()))
			break
		}
		if w.title == "search" {
			w.Fprintf("body", "Search: %s\n\n", w.query)
		}
		w.printList(data)
		w.Ctl("clean")
		w.data = data
	}
	w.Addr("0")
	w.Ctl("dot=addr")
	w.Ctl("show")

	go w.loop()
}

func (w *awin) err(s string) {
	if !strings.HasSuffix(s, "\n") {
		s = s + "\n"
	}
	w1 := w.show("+Errors")
	if w1 == nil {
		w1 = w.new("+Errors")
	}
	w1.Fprintf("body", "%s", s)
	w1.Addr("$")
	w1.Ctl("dot=addr")
	w1.Ctl("show")
}

func diff(line, field, old string) string {
	old = strings.TrimSpace(old)
	line = strings.TrimSpace(strings.TrimPrefix(line, field))
	if old == line {
		return ""
	}
	return line
}

func diffUpdate(line, field, old string) string {
	old = strings.TrimSpace(old)
	line = strings.TrimSpace(strings.TrimPrefix(line, field))
	if old == line {
		return ""
	}
	if line == "" {
		return "-" + old
	}
	return line
}

func diffUpdates(line, field string, old []string) []string {
	line = strings.TrimSpace(strings.TrimPrefix(line, field))
	had := make(map[string]bool)
	for _, f := range old {
		had[f] = true
	}
	var d []string
	kept := make(map[string]bool)
	for _, f := range strings.Fields(line) {
		f = strings.TrimSuffix(f, ",")
		if had[f] {
			kept[f] = true
		} else {
			d = append(d, f)
		}
	}
	for _, f := range old {
		if had[f] && !kept[f] {
			d = append(d, "-"+f)
		}
	}
	return d
}

func (w *awin) put() {
	switch w.mode {
	case modeSingle:
		e := w.data[0]
		var ch Change
		data, err := w.ReadAll("body")
		if err != nil {
			w.err(fmt.Sprintf("Put: %v", err))
			return
		}
		sdata := string(data)
		off := 0
		for _, line := range strings.SplitAfter(sdata, "\n") {
			off += len(line)
			line = strings.TrimSpace(line)
			if line == "" {
				break
			}
			switch {
			case strings.HasPrefix(line, "Summary:"):
				ch.Summary = diff(line, "Summary:", e.Title)

			case strings.HasPrefix(line, "Status:"):
				status := e.Status
				if status == "Duplicate" {
					status += " " + e.MergedInto
				}
				ch.Status = diff(line, "Status:", status)

			case strings.HasPrefix(line, "Owner:"):
				ch.Owner = diffUpdate(line, "Owner:", e.Owner)

			case strings.HasPrefix(line, "CC:"):
				ch.CC = diffUpdates(line, "CC:", e.CC)

			case strings.HasPrefix(line, "Labels:"):
				ch.Label = diffUpdates(line, "Labels:", e.Label)

			default:
				w.err(fmt.Sprintf("Put: unknown summary line: %s", line))
			}
		}

		i := strings.Index(sdata, "\nReported by ")
		if i >= off {
			ch.Comment = strings.TrimSpace(sdata[off:i])
		}

		tag := "\nNew Comment:"
		i = strings.Index(sdata, tag)
		if i >= 0 {
			s := strings.TrimSpace(sdata[i+len(tag):])
			if s != "" {
				if ch.Comment != "" {
					ch.Comment += "\n\n"
				}
				ch.Comment = s
			}
		}
		if err := write(w.id, &ch); err != nil {
			w.err(fmt.Sprintf("Put: %v", err))
			return
		}
		w.load()

	case modeQuery:
		w.err("cannot Put issue list")
	}
}

func (w *awin) loop() {
	defer w.exit()
	for e := range w.EventChan() {
		switch e.C2 {
		case 'x', 'X': // execute
			cmd := strings.TrimSpace(string(e.Text))
			if cmd == "Get" {
				w.load()
				break
			}
			if cmd == "Put" {
				w.put()
				break
			}
			if cmd == "Del" {
				w.Ctl("del")
				break
			}
			if strings.HasPrefix(cmd, "Search ") {
				w.newSearch("search", strings.TrimSpace(strings.TrimPrefix(cmd, "Search")))
				break
			}
			w.WriteEvent(e)
		case 'l', 'L': // look
			if len(e.Text) == 0 && e.Q0 < e.Q1 {
				w.Addr("#%d,#%d", e.Q0, e.Q1)
				data, err := w.ReadAll("xdata")
				if err != nil {
					w.err(err.Error())
				}
				e.Text = data
			}
			if !w.look(string(e.Text)) {
				w.WriteEvent(e)
			}
		}
	}
}

func (w *awin) printTabbed(rows [][]string) {
	var wid []int

	if w.font != nil {
		for _, row := range rows {
			for len(wid) < len(row) {
				wid = append(wid, 0)
			}
			for i, col := range row {
				n := w.font.StringWidth(col)
				if wid[i] < n {
					wid[i] = n
				}
			}
		}
	}

	var buf bytes.Buffer
	for _, row := range rows {
		for i, col := range row {
			buf.WriteString(col)
			if i == len(row)-1 {
				break
			}
			if w.font == nil || w.tab == 0 {
				buf.WriteString("\t")
				continue
			}
			pos := w.font.StringWidth(col)
			for pos <= wid[i] {
				buf.WriteString("\t")
				pos += w.tab - pos%w.tab
			}
		}
		buf.WriteString("\n")
	}

	w.Write("body", buf.Bytes())
}

func (w *awin) printList(entries []Entry) {
	var rows [][]string
	var buf bytes.Buffer
	for _, e := range entries {
		id := e.ID
		if i := strings.Index(id, "id="); i >= 0 {
			id = id[:i+len("id=")]
		}
		buf.Reset()
		buf.WriteString(e.Title)
		if len(e.Label) > 0 {
			buf.WriteString(" [")
			for i, l := range e.Label {
				if i > 0 {
					buf.WriteString(" ")
				}
				buf.WriteString(l)
			}
			buf.WriteString("]")
		}
		rows = append(rows, []string{id, buf.String()})
	}
	w.printTabbed(rows)
}
<|endoftext|>


<commit_before>package cmd

import (
	"errors"
	"fmt"
	"os"
	"path/filepath"

	log "github.com/sirupsen/logrus"
	"k8s.io/apimachinery/pkg/version"
	"k8s.io/client-go/kubernetes"
	_ "k8s.io/client-go/plugin/pkg/client/auth/gcp"  // auth for GKE clusters
	_ "k8s.io/client-go/plugin/pkg/client/auth/oidc" // auth for OIDC
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"
)

// Client abstracts the API to allow testing.
type Client interface {
	InClusterConfig() (*rest.Config, error)
}

// K8sClient wraps kubernetes client-go so it can be mocked.
type K8sClient struct{}

// InClusterConfig wraps the client-go method with the same name.
func (kc K8sClient) InClusterConfig() (*rest.Config, error) {
	return rest.InClusterConfig()
}

// ErrNoReadableKubeConfig represents any error that prevents the client from opening a kubeconfig file.
var ErrNoReadableKubeConfig = errors.New("unable to open kubeconfig file")

func kubeClient() (*kubernetes.Clientset, error) {
	return kubeClientType(K8sClient{})
}

func kubeClientType(kc Client) (*kubernetes.Clientset, error) {
	config, err := kubeClientConfig(kc)
	if err != nil {
		return nil, err
	}
	kube, err := kubernetes.NewForConfig(config)
	return kube, err
}

func kubeClientConfig(kc Client) (*rest.Config, error) {
	if rootConfig.kubeConfig != "" {
		return kubeClientConfigLocal()
	}

	if config, err := kc.InClusterConfig(); err == nil {
		log.Info("Running inside cluster, using the cluster config")
		return config, nil
	}

	log.Info("Not running inside cluster, using local config")
	home, ok := os.LookupEnv("HOME")
	if !ok || home == "" {
		log.Error("Unable to load kubeconfig. No config file specified and $HOME not found.")
		return nil, ErrNoReadableKubeConfig
	}

	rootConfig.kubeConfig = filepath.Join(home, ".kube", "config")
	return kubeClientConfigLocal()
}

func kubeClientConfigLocal() (*rest.Config, error) {
	if _, err := os.Stat(rootConfig.kubeConfig); err != nil {
		log.Errorf("Unable to load kubeconfig. Could not open file %s.", rootConfig.kubeConfig)
		return nil, ErrNoReadableKubeConfig
	}
	return clientcmd.BuildConfigFromFlags("", rootConfig.kubeConfig)
}

func getDeployments(clientset *kubernetes.Clientset) *DeploymentListV1 {
	deploymentClient := clientset.AppsV1().Deployments(rootConfig.namespace)
	deployments, err := deploymentClient.List(ListOptionsV1{})
	if err != nil {
		log.Error(err)
	}
	return deployments
}

func getStatefulSets(clientset *kubernetes.Clientset) *StatefulSetListV1 {
	statefulSetClient := clientset.AppsV1().StatefulSets(rootConfig.namespace)
	statefulSets, err := statefulSetClient.List(ListOptionsV1{})
	if err != nil {
		log.Error(err)
	}
	return statefulSets
}

func getDaemonSets(clientset *kubernetes.Clientset) *DaemonSetListV1 {
	daemonSetClient := clientset.AppsV1().DaemonSets(rootConfig.namespace)
	daemonSets, err := daemonSetClient.List(ListOptionsV1{})
	if err != nil {
		log.Error(err)
	}
	return daemonSets
}

func getPods(clientset *kubernetes.Clientset) *PodListV1 {
	podClient := clientset.CoreV1().Pods(rootConfig.namespace)
	pods, err := podClient.List(ListOptionsV1{})
	if err != nil {
		log.Error(err)
	}
	return pods
}

func getReplicationControllers(clientset *kubernetes.Clientset) *ReplicationControllerListV1 {
	replicationControllerClient := clientset.CoreV1().ReplicationControllers(rootConfig.namespace)
	replicationControllers, err := replicationControllerClient.List(ListOptionsV1{})
	if err != nil {
		log.Error(err)
	}
	return replicationControllers
}

func getNetworkPolicies(clientset *kubernetes.Clientset) *NetworkPolicyListV1 {
	netPolClient := clientset.NetworkingV1().NetworkPolicies(rootConfig.namespace)
	netPols, err := netPolClient.List(ListOptionsV1{})
	if err != nil {
		log.Error(err)
	}
	return netPols
}

func getNamespaces(clientset *kubernetes.Clientset) *NamespaceListV1 {
	namespaceClient := clientset.CoreV1().Namespaces()
	listOptions := ListOptionsV1{}

	if rootConfig.namespace != "" {
		// Select only the specified namespace
		listOptions = ListOptionsV1{
			FieldSelector: fmt.Sprintf("metadata.name=%s", rootConfig.namespace),
		}
	}

	namespaces, err := namespaceClient.List(listOptions)

	if err != nil {
		log.Error(err)
	}

	return namespaces
}

func getKubernetesVersion(clientset kubernetes.Interface) (*version.Info, error) {
	discoveryClient := clientset.Discovery()
	return discoveryClient.ServerVersion()
}
<commit_msg>import azure auth module for AKS (#268)<commit_after>package cmd

import (
	"errors"
	"fmt"
	"os"
	"path/filepath"

	log "github.com/sirupsen/logrus"
	"k8s.io/apimachinery/pkg/version"
	"k8s.io/client-go/kubernetes"
	_ "k8s.io/client-go/plugin/pkg/client/auth/azure" // auth for AKS clusters
	_ "k8s.io/client-go/plugin/pkg/client/auth/gcp"   // auth for GKE clusters
	_ "k8s.io/client-go/plugin/pkg/client/auth/oidc"  // auth for OIDC
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"
)

// Client abstracts the API to allow testing.
type Client interface {
	InClusterConfig() (*rest.Config, error)
}

// K8sClient wraps kubernetes client-go so it can be mocked.
type K8sClient struct{}

// InClusterConfig wraps the client-go method with the same name.
func (kc K8sClient) InClusterConfig() (*rest.Config, error) {
	return rest.InClusterConfig()
}

// ErrNoReadableKubeConfig represents any error that prevents the client from opening a kubeconfig file.
var ErrNoReadableKubeConfig = errors.New("unable to open kubeconfig file")

func kubeClient() (*kubernetes.Clientset, error) {
	return kubeClientType(K8sClient{})
}

func kubeClientType(kc Client) (*kubernetes.Clientset, error) {
	config, err := kubeClientConfig(kc)
	if err != nil {
		return nil, err
	}
	kube, err := kubernetes.NewForConfig(config)
	return kube, err
}

func kubeClientConfig(kc Client) (*rest.Config, error) {
	if rootConfig.kubeConfig != "" {
		return kubeClientConfigLocal()
	}

	if config, err := kc.InClusterConfig(); err == nil {
		log.Info("Running inside cluster, using the cluster config")
		return config, nil
	}

	log.Info("Not running inside cluster, using local config")
	home, ok := os.LookupEnv("HOME")
	if !ok || home == "" {
		log.Error("Unable to load kubeconfig. No config file specified and $HOME not found.")
		return nil, ErrNoReadableKubeConfig
	}

	rootConfig.kubeConfig = filepath.Join(home, ".kube", "config")
	return kubeClientConfigLocal()
}

func kubeClientConfigLocal() (*rest.Config, error) {
	if _, err := os.Stat(rootConfig.kubeConfig); err != nil {
		log.Errorf("Unable to load kubeconfig. Could not open file %s.", rootConfig.kubeConfig)
		return nil, ErrNoReadableKubeConfig
	}
	return clientcmd.BuildConfigFromFlags("", rootConfig.kubeConfig)
}

func getDeployments(clientset *kubernetes.Clientset) *DeploymentListV1 {
	deploymentClient := clientset.AppsV1().Deployments(rootConfig.namespace)
	deployments, err := deploymentClient.List(ListOptionsV1{})
	if err != nil {
		log.Error(err)
	}
	return deployments
}

func getStatefulSets(clientset *kubernetes.Clientset) *StatefulSetListV1 {
	statefulSetClient := clientset.AppsV1().StatefulSets(rootConfig.namespace)
	statefulSets, err := statefulSetClient.List(ListOptionsV1{})
	if err != nil {
		log.Error(err)
	}
	return statefulSets
}

func getDaemonSets(clientset *kubernetes.Clientset) *DaemonSetListV1 {
	daemonSetClient := clientset.AppsV1().DaemonSets(rootConfig.namespace)
	daemonSets, err := daemonSetClient.List(ListOptionsV1{})
	if err != nil {
		log.Error(err)
	}
	return daemonSets
}

func getPods(clientset *kubernetes.Clientset) *PodListV1 {
	podClient := clientset.CoreV1().Pods(rootConfig.namespace)
	pods, err := podClient.List(ListOptionsV1{})
	if err != nil {
		log.Error(err)
	}
	return pods
}

func getReplicationControllers(clientset *kubernetes.Clientset) *ReplicationControllerListV1 {
	replicationControllerClient := clientset.CoreV1().ReplicationControllers(rootConfig.namespace)
	replicationControllers, err := replicationControllerClient.List(ListOptionsV1{})
	if err != nil {
		log.Error(err)
	}
	return replicationControllers
}

func getNetworkPolicies(clientset *kubernetes.Clientset) *NetworkPolicyListV1 {
	netPolClient := clientset.NetworkingV1().NetworkPolicies(rootConfig.namespace)
	netPols, err := netPolClient.List(ListOptionsV1{})
	if err != nil {
		log.Error(err)
	}
	return netPols
}

func getNamespaces(clientset *kubernetes.Clientset) *NamespaceListV1 {
	namespaceClient := clientset.CoreV1().Namespaces()
	listOptions := ListOptionsV1{}

	if rootConfig.namespace != "" {
		// Select only the specified namespace
		listOptions = ListOptionsV1{
			FieldSelector: fmt.Sprintf("metadata.name=%s", rootConfig.namespace),
		}
	}

	namespaces, err := namespaceClient.List(listOptions)

	if err != nil {
		log.Error(err)
	}

	return namespaces
}

func getKubernetesVersion(clientset kubernetes.Interface) (*version.Info, error) {
	discoveryClient := clientset.Discovery()
	return discoveryClient.ServerVersion()
}
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"

	"github.com/Masterminds/cookoo"
	"github.com/kylelemons/go-gypsy/yaml"
)

// ParseYaml parses the glide.yaml format and returns a Configuration object.
//
// Params:
// 	- filename (string): YAML filename as a string
//
// Context:
// 	- yaml.File: This puts the parsed YAML file into the context.
//
// Returns:
// 	- *Config: The configuration.
func ParseYaml(c cookoo.Context, p *cookoo.Params) (interface{}, cookoo.Interrupt) {
	fname := p.Get("filename", "glide.yaml").(string)
	//conf := new(Config)
	f, err := yaml.ReadFile(fname)
	if err != nil {
		return nil, err
	}

	c.Put("yaml.File", f)
	return FromYaml(f.Root)

	/*

	// Convenience:
	top, ok := f.Root.(yaml.Map)
	if !ok {
		return nil, fmt.Errorf("Expected YAML root to be map, got %t", f.Root)
	}

	vals := map[string]yaml.Node(top)
	if name, ok := vals["package"]; ok {
		//c.Put("cfg.package", name.(yaml.Scalar).String())
		conf.Name = name.(yaml.Scalar).String()
	} else {
		Warn("The 'package' directive is required in Glide YAML.\n")
	}

	// Allow the user to override the behavior of `glide in`.
	if incmd, ok := vals["incmd"]; ok {
		conf.InCommand = incmd.(yaml.Scalar).String()
	}

	conf.Imports = make([]*Dependency, 0, 1)
	if imp, ok := vals["import"]; ok {
		imports, ok := imp.(yaml.List)

		if ok {
			for _, v := range imports {
				pkg := v.(yaml.Map)
				dep := Dependency{
					Name:        valOrEmpty("package", pkg),
					Reference:   valOrEmpty("ref", pkg),
					VcsType:     getVcsType(pkg),
					Repository:  valOrEmpty("repo", pkg),
					Subpackages: subpkg("subpackages", pkg),
				}
				conf.Imports = append(conf.Imports, &dep)
			}
		}
	}

	return conf, nil
	*/
}

// WriteYaml writes a yaml.Node to the console as a string.
//
// Params:
// 	- yaml.Node: A yaml.Node to render.
func WriteYaml(c cookoo.Context, p *cookoo.Params) (interface{}, cookoo.Interrupt) {
	top := p.Get("yaml.Node", yaml.Scalar("nothing to print")).(yaml.Node)

	fmt.Print(yaml.Render(top))

	return true, nil
}

// Convert a Config object and a yaml.File to a single yaml.File.
func MergeToYaml(c cookoo.Context, p *cookoo.Params) (interface{}, cookoo.Interrupt) {
	root := c.Get("yaml.File", nil).(*yaml.File).Root
	cfg := p.Get("conf", nil).(*Config)
	overwrite := p.Get("overwriteImports", true).(bool)

	rootMap, ok := root.(yaml.Map)
	if !ok {
		return nil, fmt.Errorf("Expected root node to be a map.")
	}

	if len(cfg.Name) > 0 {
		rootMap["package"] = yaml.Scalar(cfg.Name)
	}
	if cfg.InCommand != "" {
		rootMap["incmd"] = yaml.Scalar(cfg.InCommand)
	}

	// Imports
	imports := make([]yaml.Node, len(cfg.Imports))
	for i, imp := range cfg.Imports {

		if imp.VcsType == NoVCS {
			imp.VcsType, _ = GuessVCS(imp)
		}

		impmap := make(map[string]yaml.Node)
		impmap["package"] = yaml.Scalar(imp.Name)
		if imp.VcsType != NoVCS {
			impmap["vcs"] = yaml.Scalar(vcsString(imp.VcsType))
		}
		if imp.Reference != "" {
			impmap["ref"] = yaml.Scalar(imp.Reference)
		}
		if imp.Repository != "" {
			impmap["repo"] = yaml.Scalar(imp.Repository)
		}

		if len(imp.Subpackages) > 0 {
			subs := make([]yaml.Node, len(imp.Subpackages))
			for ii, sub := range imp.Subpackages {
				subs[ii] = yaml.Scalar(sub)
			}
			impmap["subpackages"] = yaml.List(subs)
		}

		imports[i] = yaml.Map(impmap)
	}

	if overwrite {
		rootMap["import"] = yaml.List(imports)
	} else {
	}

	return root, nil
}

/*
func mergeImports(rootMap yaml.Map, imports []yaml.Map) yaml.List {
	rootList, ok := rootMap["import"].([]yaml.Map)
	if !ok {
		return yaml.List(imports)
	}

	for _, item := range imports {
		missing := true
		for _, existing := range rootList {
			if item.Name == existing.Name {
				missing = false
			}
		}
		if missing {
			rootList = append(rootList, item)
		}
	}
}
*/

// AddDependencies adds a list of *Dependency objects to the given *Config.
//
// This is used to merge in packages from other sources or config files.
func AddDependencies(c cookoo.Context, p *cookoo.Params) (interface{}, cookoo.Interrupt) {
	deps := p.Get("dependencies", []*Dependency{}).([]*Dependency)
	config := p.Get("conf", nil).(*Config)

	// Make a set of existing package names for quick comparison.
	pkgSet := make(map[string]bool, len(config.Imports))
	for _, p := range config.Imports {
		pkgSet[p.Name] = true
	}

	// If a dep is not already present, add it.
	for _, dep := range deps {
		if _, ok := pkgSet[dep.Name]; ok {
			Warn("Package %s is already in glide.yaml. Skipping.\n", dep.Name)
			continue
		}
		config.Imports = append(config.Imports, dep)
	}

	return true, nil
}

func vcsString(vtype uint) string {
	switch vtype {
	case Git:
		return "git"
	case Hg:
		return "hg"
	case Bzr:
		return "bzr"
	case Svn:
		return "svn"
	default:
		return ""
	}
}

func valOrEmpty(key string, store map[string]yaml.Node) string {
	val, ok := store[key]
	if !ok {
		return ""
	}
	return val.(yaml.Scalar).String()
}

func subpkg(key string, store map[string]yaml.Node) []string {
	val, ok := store[key]

	subpackages := []string{}
	if !ok {
		return subpackages
	}

	pkgs, ok := val.(yaml.List)

	if !ok {

		// Special case: Allow 'subpackages: justOne'
		if one, ok := val.(yaml.Scalar); ok {
			return []string{one.String()}
		}

		Warn("Expected list of subpackages.\n")
		return subpackages
	}

	for _, pkg := range pkgs {
		subpackages = append(subpackages, pkg.(yaml.Scalar).String())
	}
	return subpackages
}

func getVcsType(store map[string]yaml.Node) uint {
	val, ok := store["vcs"]
	if !ok {
		return NoVCS
	}

	name := val.(yaml.Scalar).String()

	switch name {
	case "git":
		return Git
	case "hg", "mercurial":
		return Hg
	case "bzr", "bazaar":
		return Bzr
	case "svn", "subversion":
		return Svn
	default:
		return NoVCS
	}
}

// Config is the top-level configuration object.
type Config struct {
	Name       string
	Imports    []*Dependency
	DevImports []*Dependency
	// InCommand is the default shell command run to start a 'glide in'
	// session.
	InCommand string
}

// FromYaml creates a *Config from a  YAML node.
func FromYaml(top yaml.Node) (*Config, error) {
	conf := new(Config)

	vals, ok := top.(yaml.Map)
	if !ok {
		return conf, fmt.Errorf("Top YAML node must be a map.")
	}

	if name, ok := vals["package"]; ok {
		conf.Name = name.(yaml.Scalar).String()
	} else {
		Warn("The 'package' directive is required in Glide YAML.\n")
		conf.Name = "main"
	}

	// Allow the user to override the behavior of `glide in`.
	if incmd, ok := vals["incmd"]; ok {
		conf.InCommand = incmd.(yaml.Scalar).String()
	}

	conf.Imports = make([]*Dependency, 0, 1)
	if imp, ok := vals["import"]; ok {
		imports, ok := imp.(yaml.List)

		if ok {
			for _, v := range imports {
				/*
				pkg := v.(yaml.Map)
				dep := Dependency{
					Name:        valOrEmpty("package", pkg),
					Reference:   valOrEmpty("ref", pkg),
					VcsType:     getVcsType(pkg),
					Repository:  valOrEmpty("repo", pkg),
					Subpackages: subpkg("subpackages", pkg),
				}
				*/
				dep, err := DependencyFromYaml(v)
				if err != nil {
					Warn("Could not add a dependency: %s\n", err)
				}
				conf.Imports = append(conf.Imports, dep)
			}
		}
	}

	// Same for (experimental) devimport.
	// These are currently unused. Not sure what we'll do with it yet.
	conf.DevImports = []*Dependency{}
	if imp, ok := vals["devimport"]; ok {
		imports, ok := imp.(yaml.List)
		if ok {
			for _, v := range imports {
				dep, err := DependencyFromYaml(v)
				if err != nil {
					Warn("Could not add a dependency: %s\n", err)
				}
				conf.DevImports = append(conf.DevImports, dep)
			}
		}
	}

	return conf, nil
}

// ToYaml returns a yaml.Map containing the data from Config.
func (c *Config) ToYaml() yaml.Node {
	cfg := make(map[string]yaml.Node, 4)

	cfg["package"] = yaml.Scalar(c.Name)
	if len(c.InCommand) > 0 {
		cfg["incmd"] = yaml.Scalar(c.InCommand)
	}

	imps := make([]yaml.Node, len(c.Imports))
	for i, imp := range c.Imports {
		imps[i] = imp.ToYaml()
	}
	devimps := make([]yaml.Node, len(c.DevImports))
	for i, dimp := range c.DevImports {
		devimps[i] = dimp.ToYaml()
	}

	return yaml.Map(cfg)
}

// Dependency describes a package that the present package depends upon.
type Dependency struct {
	Name, Reference, Repository string
	VcsType                     uint
	Subpackages                 []string
}

// DependencyFromYaml creates a dependency from a yaml.Node.
func DependencyFromYaml(node yaml.Node) (*Dependency, error) {
	pkg, ok := node.(yaml.Map)
	if !ok {
		return &Dependency{}, fmt.Errorf("Expected yaml.Node to be a dependency map.")
	}
	dep := &Dependency{
		Name:        valOrEmpty("package", pkg),
		Reference:   valOrEmpty("ref", pkg),
		VcsType:     getVcsType(pkg),
		Repository:  valOrEmpty("repo", pkg),
		Subpackages: subpkg("subpackages", pkg),
	}

	return dep, nil
}

// ToYaml converts a *Dependency to a YAML Map node.
func (d *Dependency) ToYaml() yaml.Node {
	dep := make(map[string]yaml.Node, 5)
	dep["package"] = yaml.Scalar(d.Name)

	subp := make([]yaml.Node, len(d.Subpackages))
	for i, item := range d.Subpackages {
		subp[i] = yaml.Scalar(item)
	}

	dep["subpackages"] = yaml.List(subp)
	vcs := vcsString(d.VcsType)
	if len(vcs) > 0 {
		dep["vcs"] = yaml.Scalar(vcs)
	}
	if len(d.Reference) > 0 {
		dep["ref"] = yaml.Scalar(d.Reference)
	}
	if len(d.Repository) > 0 {
		dep["repo"] = yaml.Scalar(d.Repository)
	}
	return yaml.Map(dep)
}
<commit_msg>Fixed vim to stop doing <space><tab> in doc strings.<commit_after>package cmd

import (
	"fmt"

	"github.com/Masterminds/cookoo"
	"github.com/kylelemons/go-gypsy/yaml"
)

// ParseYaml parses the glide.yaml format and returns a Configuration object.
//
// Params:
//	- filename (string): YAML filename as a string
//
// Context:
//	- yaml.File: This puts the parsed YAML file into the context.
//
// Returns:
//	- *Config: The configuration.
func ParseYaml(c cookoo.Context, p *cookoo.Params) (interface{}, cookoo.Interrupt) {
	fname := p.Get("filename", "glide.yaml").(string)
	//conf := new(Config)
	f, err := yaml.ReadFile(fname)
	if err != nil {
		return nil, err
	}

	c.Put("yaml.File", f)
	return FromYaml(f.Root)
}

// WriteYaml writes a yaml.Node to the console as a string.
//
// Params:
// 	- yaml.Node: A yaml.Node to render.
func WriteYaml(c cookoo.Context, p *cookoo.Params) (interface{}, cookoo.Interrupt) {
	top := p.Get("yaml.Node", yaml.Scalar("nothing to print")).(yaml.Node)

	fmt.Print(yaml.Render(top))

	return true, nil
}

// Convert a Config object and a yaml.File to a single yaml.File.
func MergeToYaml(c cookoo.Context, p *cookoo.Params) (interface{}, cookoo.Interrupt) {
	root := c.Get("yaml.File", nil).(*yaml.File).Root
	cfg := p.Get("conf", nil).(*Config)
	overwrite := p.Get("overwriteImports", true).(bool)

	rootMap, ok := root.(yaml.Map)
	if !ok {
		return nil, fmt.Errorf("Expected root node to be a map.")
	}

	if len(cfg.Name) > 0 {
		rootMap["package"] = yaml.Scalar(cfg.Name)
	}
	if cfg.InCommand != "" {
		rootMap["incmd"] = yaml.Scalar(cfg.InCommand)
	}

	// Imports
	imports := make([]yaml.Node, len(cfg.Imports))
	for i, imp := range cfg.Imports {

		if imp.VcsType == NoVCS {
			imp.VcsType, _ = GuessVCS(imp)
		}

		impmap := make(map[string]yaml.Node)
		impmap["package"] = yaml.Scalar(imp.Name)
		if imp.VcsType != NoVCS {
			impmap["vcs"] = yaml.Scalar(vcsString(imp.VcsType))
		}
		if imp.Reference != "" {
			impmap["ref"] = yaml.Scalar(imp.Reference)
		}
		if imp.Repository != "" {
			impmap["repo"] = yaml.Scalar(imp.Repository)
		}

		if len(imp.Subpackages) > 0 {
			subs := make([]yaml.Node, len(imp.Subpackages))
			for ii, sub := range imp.Subpackages {
				subs[ii] = yaml.Scalar(sub)
			}
			impmap["subpackages"] = yaml.List(subs)
		}

		imports[i] = yaml.Map(impmap)
	}

	if overwrite {
		rootMap["import"] = yaml.List(imports)
	} else {
	}

	return root, nil
}

/*
func mergeImports(rootMap yaml.Map, imports []yaml.Map) yaml.List {
	rootList, ok := rootMap["import"].([]yaml.Map)
	if !ok {
		return yaml.List(imports)
	}

	for _, item := range imports {
		missing := true
		for _, existing := range rootList {
			if item.Name == existing.Name {
				missing = false
			}
		}
		if missing {
			rootList = append(rootList, item)
		}
	}
}
*/

// AddDependencies adds a list of *Dependency objects to the given *Config.
//
// This is used to merge in packages from other sources or config files.
func AddDependencies(c cookoo.Context, p *cookoo.Params) (interface{}, cookoo.Interrupt) {
	deps := p.Get("dependencies", []*Dependency{}).([]*Dependency)
	config := p.Get("conf", nil).(*Config)

	// Make a set of existing package names for quick comparison.
	pkgSet := make(map[string]bool, len(config.Imports))
	for _, p := range config.Imports {
		pkgSet[p.Name] = true
	}

	// If a dep is not already present, add it.
	for _, dep := range deps {
		if _, ok := pkgSet[dep.Name]; ok {
			Warn("Package %s is already in glide.yaml. Skipping.\n", dep.Name)
			continue
		}
		config.Imports = append(config.Imports, dep)
	}

	return true, nil
}

func vcsString(vtype uint) string {
	switch vtype {
	case Git:
		return "git"
	case Hg:
		return "hg"
	case Bzr:
		return "bzr"
	case Svn:
		return "svn"
	default:
		return ""
	}
}

func valOrEmpty(key string, store map[string]yaml.Node) string {
	val, ok := store[key]
	if !ok {
		return ""
	}
	return val.(yaml.Scalar).String()
}

func subpkg(key string, store map[string]yaml.Node) []string {
	val, ok := store[key]

	subpackages := []string{}
	if !ok {
		return subpackages
	}

	pkgs, ok := val.(yaml.List)

	if !ok {

		// Special case: Allow 'subpackages: justOne'
		if one, ok := val.(yaml.Scalar); ok {
			return []string{one.String()}
		}

		Warn("Expected list of subpackages.\n")
		return subpackages
	}

	for _, pkg := range pkgs {
		subpackages = append(subpackages, pkg.(yaml.Scalar).String())
	}
	return subpackages
}

func getVcsType(store map[string]yaml.Node) uint {
	val, ok := store["vcs"]
	if !ok {
		return NoVCS
	}

	name := val.(yaml.Scalar).String()

	switch name {
	case "git":
		return Git
	case "hg", "mercurial":
		return Hg
	case "bzr", "bazaar":
		return Bzr
	case "svn", "subversion":
		return Svn
	default:
		return NoVCS
	}
}

// Config is the top-level configuration object.
type Config struct {
	Name       string
	Imports    []*Dependency
	DevImports []*Dependency
	// InCommand is the default shell command run to start a 'glide in'
	// session.
	InCommand string
}

// FromYaml creates a *Config from a  YAML node.
func FromYaml(top yaml.Node) (*Config, error) {
	conf := new(Config)

	vals, ok := top.(yaml.Map)
	if !ok {
		return conf, fmt.Errorf("Top YAML node must be a map.")
	}

	if name, ok := vals["package"]; ok {
		conf.Name = name.(yaml.Scalar).String()
	} else {
		Warn("The 'package' directive is required in Glide YAML.\n")
		conf.Name = "main"
	}

	// Allow the user to override the behavior of `glide in`.
	if incmd, ok := vals["incmd"]; ok {
		conf.InCommand = incmd.(yaml.Scalar).String()
	}

	conf.Imports = make([]*Dependency, 0, 1)
	if imp, ok := vals["import"]; ok {
		imports, ok := imp.(yaml.List)

		if ok {
			for _, v := range imports {
				/*
				pkg := v.(yaml.Map)
				dep := Dependency{
					Name:        valOrEmpty("package", pkg),
					Reference:   valOrEmpty("ref", pkg),
					VcsType:     getVcsType(pkg),
					Repository:  valOrEmpty("repo", pkg),
					Subpackages: subpkg("subpackages", pkg),
				}
				*/
				dep, err := DependencyFromYaml(v)
				if err != nil {
					Warn("Could not add a dependency: %s\n", err)
				}
				conf.Imports = append(conf.Imports, dep)
			}
		}
	}

	// Same for (experimental) devimport.
	// These are currently unused. Not sure what we'll do with it yet.
	conf.DevImports = []*Dependency{}
	if imp, ok := vals["devimport"]; ok {
		imports, ok := imp.(yaml.List)
		if ok {
			for _, v := range imports {
				dep, err := DependencyFromYaml(v)
				if err != nil {
					Warn("Could not add a dependency: %s\n", err)
				}
				conf.DevImports = append(conf.DevImports, dep)
			}
		}
	}

	return conf, nil
}

// ToYaml returns a yaml.Map containing the data from Config.
func (c *Config) ToYaml() yaml.Node {
	cfg := make(map[string]yaml.Node, 4)

	cfg["package"] = yaml.Scalar(c.Name)
	if len(c.InCommand) > 0 {
		cfg["incmd"] = yaml.Scalar(c.InCommand)
	}

	imps := make([]yaml.Node, len(c.Imports))
	for i, imp := range c.Imports {
		imps[i] = imp.ToYaml()
	}
	devimps := make([]yaml.Node, len(c.DevImports))
	for i, dimp := range c.DevImports {
		devimps[i] = dimp.ToYaml()
	}

	return yaml.Map(cfg)
}

// Dependency describes a package that the present package depends upon.
type Dependency struct {
	Name, Reference, Repository string
	VcsType                     uint
	Subpackages                 []string
}

// DependencyFromYaml creates a dependency from a yaml.Node.
func DependencyFromYaml(node yaml.Node) (*Dependency, error) {
	pkg, ok := node.(yaml.Map)
	if !ok {
		return &Dependency{}, fmt.Errorf("Expected yaml.Node to be a dependency map.")
	}
	dep := &Dependency{
		Name:        valOrEmpty("package", pkg),
		Reference:   valOrEmpty("ref", pkg),
		VcsType:     getVcsType(pkg),
		Repository:  valOrEmpty("repo", pkg),
		Subpackages: subpkg("subpackages", pkg),
	}

	return dep, nil
}

// ToYaml converts a *Dependency to a YAML Map node.
func (d *Dependency) ToYaml() yaml.Node {
	dep := make(map[string]yaml.Node, 5)
	dep["package"] = yaml.Scalar(d.Name)

	subp := make([]yaml.Node, len(d.Subpackages))
	for i, item := range d.Subpackages {
		subp[i] = yaml.Scalar(item)
	}

	dep["subpackages"] = yaml.List(subp)
	vcs := vcsString(d.VcsType)
	if len(vcs) > 0 {
		dep["vcs"] = yaml.Scalar(vcs)
	}
	if len(d.Reference) > 0 {
		dep["ref"] = yaml.Scalar(d.Reference)
	}
	if len(d.Repository) > 0 {
		dep["repo"] = yaml.Scalar(d.Repository)
	}
	return yaml.Map(dep)
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Knative 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 main

import (
	"context"
	"flag"
	"fmt"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"time"

	"go.uber.org/zap"

	"github.com/knative/pkg/logging/logkey"
	"github.com/knative/pkg/metrics"
	"github.com/knative/pkg/signals"
	"github.com/knative/serving/cmd/util"
	"github.com/knative/serving/pkg/activator"
	activatorutil "github.com/knative/serving/pkg/activator/util"
	"github.com/knative/serving/pkg/apis/networking"
	"github.com/knative/serving/pkg/autoscaler"
	pkghttp "github.com/knative/serving/pkg/http"
	"github.com/knative/serving/pkg/logging"
	"github.com/knative/serving/pkg/network"
	"github.com/knative/serving/pkg/queue"
	"github.com/knative/serving/pkg/queue/health"
	queuestats "github.com/knative/serving/pkg/queue/stats"

	"k8s.io/apimachinery/pkg/util/wait"
)

const (
	// Add a little buffer space between request handling and stat
	// reporting so that latency in the stat pipeline doesn't
	// interfere with request handling.
	statReportingQueueLength = 10

	// Add enough buffer to not block request serving on stats collection
	requestCountingQueueLength = 100

	// Duration the /quitquitquit handler should wait before returning.
	// This is to give Istio a little bit more time to remove the pod
	// from its configuration and propagate that to all istio-proxies
	// in the mesh.
	quitSleepDuration = 20 * time.Second

	// commonMetricsPort is the port where common metrics, e.g. request metrics
	// are exposed in Prometheus. This is different from the metrics used
	// for autoscaling, which are exposed in 9090.
	commonMetricsPort = 9091

	badProbeTemplate = "unexpected probe header value: %s"
)

var (
	servingService         string
	servingConfig          string
	servingNamespace       string
	servingRevision        string
	servingRevisionKey     string
	servingPodIP           string
	servingPodName         string
	userTargetPort         int
	userTargetAddress      string
	containerConcurrency   int
	revisionTimeoutSeconds int
	reqChan                = make(chan queue.ReqEvent, requestCountingQueueLength)
	logger                 *zap.SugaredLogger
	breaker                *queue.Breaker

	httpProxy *httputil.ReverseProxy

	healthState      = &health.State{}
	promStatReporter *queue.PrometheusStatsReporter // Prometheus stats reporter.
)

func initEnv() {
	servingService = os.Getenv("SERVING_SERVICE") // KService is optional
	servingConfig = util.GetRequiredEnvOrFatal("SERVING_CONFIGURATION", logger)
	servingNamespace = util.GetRequiredEnvOrFatal("SERVING_NAMESPACE", logger)
	servingRevision = util.GetRequiredEnvOrFatal("SERVING_REVISION", logger)
	servingPodIP = util.GetRequiredEnvOrFatal("SERVING_POD_IP", logger)
	servingPodName = util.GetRequiredEnvOrFatal("SERVING_POD", logger)
	containerConcurrency = util.MustParseIntEnvOrFatal("CONTAINER_CONCURRENCY", logger)
	revisionTimeoutSeconds = util.MustParseIntEnvOrFatal("REVISION_TIMEOUT_SECONDS", logger)
	userTargetPort = util.MustParseIntEnvOrFatal("USER_PORT", logger)
	userTargetAddress = fmt.Sprintf("127.0.0.1:%d", userTargetPort)

	// TODO(mattmoor): Move this key to be in terms of the KPA.
	servingRevisionKey = autoscaler.NewMetricKey(servingNamespace, servingRevision)
	_psr, err := queue.NewPrometheusStatsReporter(servingNamespace, servingConfig, servingRevision, servingPodName)
	if err != nil {
		logger.Fatalw("Failed to create stats reporter", zap.Error(err))
	}
	promStatReporter = _psr
}

func reportStats(statChan chan *autoscaler.Stat) {
	for {
		s := <-statChan
		if err := promStatReporter.Report(s); err != nil {
			logger.Errorw("Error while sending stat", zap.Error(err))
		}
	}
}

func knativeProbeHeader(r *http.Request) string {
	return r.Header.Get(network.ProbeHeaderName)
}

func knativeProxyHeader(r *http.Request) string {
	return r.Header.Get(network.ProxyHeaderName)
}

func probeUserContainer() bool {
	var err error
	wait.PollImmediate(50*time.Millisecond, 10*time.Second, func() (bool, error) {
		logger.Debug("TCP probing the user-container.")
		err = health.TCPProbe(userTargetAddress, 100*time.Millisecond)
		return err == nil, nil
	})

	if err == nil {
		logger.Info("User-container successfully probed.")
	} else {
		logger.Errorw("User-container could not be probed successfully.", zap.Error(err))
	}

	return err == nil
}

// Make handler a closure for testing.
func handler(reqChan chan queue.ReqEvent, breaker *queue.Breaker, proxy *httputil.ReverseProxy) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		ph := knativeProbeHeader(r)
		switch {
		case ph != "":
			if ph != queue.Name {
				http.Error(w, fmt.Sprintf(badProbeTemplate, ph), http.StatusBadRequest)
				return
			}
			if probeUserContainer() {
				// Respond with the name of the component handling the request.
				w.Write([]byte(queue.Name))
			} else {
				http.Error(w, "container not ready", http.StatusServiceUnavailable)
			}
			return
		case network.IsKubeletProbe(r):
			// Do not count health checks for concurrency metrics
			proxy.ServeHTTP(w, r)
			return
		}

		// Metrics for autoscaling.
		h := knativeProxyHeader(r)
		in, out := queue.ReqIn, queue.ReqOut
		if activator.Name == h {
			in, out = queue.ProxiedIn, queue.ProxiedOut
		}
		reqChan <- queue.ReqEvent{Time: time.Now(), EventType: in}
		defer func() {
			reqChan <- queue.ReqEvent{Time: time.Now(), EventType: out}
		}()
		network.RewriteHostOut(r)

		// Enforce queuing and concurrency limits.
		if breaker != nil {
			ok := breaker.Maybe(func() {
				proxy.ServeHTTP(w, r)
			})
			if !ok {
				http.Error(w, "overload", http.StatusServiceUnavailable)
			}
		} else {
			proxy.ServeHTTP(w, r)
		}
	}
}

// Sets up /health and /wait-for-drain endpoints.
func createAdminHandlers() *http.ServeMux {
	mux := http.NewServeMux()
	mux.HandleFunc(queue.RequestQueueHealthPath, healthState.HealthHandler(probeUserContainer))
	mux.HandleFunc(queue.RequestQueueDrainPath, healthState.DrainHandler())

	return mux
}

func main() {
	flag.Parse()
	logger, _ = logging.NewLogger(os.Getenv("SERVING_LOGGING_CONFIG"), os.Getenv("SERVING_LOGGING_LEVEL"))
	logger = logger.Named("queueproxy")
	defer flush(logger)

	initEnv()
	logger = logger.With(
		zap.String(logkey.Key, servingRevisionKey),
		zap.String(logkey.Pod, servingPodName))

	target, err := url.Parse("http://" + userTargetAddress)
	if err != nil {
		logger.Fatalw("Failed to parse localhost URL", zap.Error(err))
	}

	httpProxy = httputil.NewSingleHostReverseProxy(target)
	httpProxy.Transport = network.AutoTransport
	httpProxy.FlushInterval = -1

	activatorutil.SetupHeaderPruning(httpProxy)

	// If containerConcurrency == 0 then concurrency is unlimited.
	if containerConcurrency > 0 {
		// We set the queue depth to be equal to the container concurrency but at least 10 to
		// allow the autoscaler to get a strong enough signal.
		queueDepth := containerConcurrency
		if queueDepth < 10 {
			queueDepth = 10
		}
		params := queue.BreakerParams{QueueDepth: queueDepth, MaxConcurrency: containerConcurrency, InitialCapacity: containerConcurrency}
		breaker = queue.NewBreaker(params)
		logger.Infof("Queue container is starting with %#v", params)
	}

	go func() {
		mux := http.NewServeMux()
		mux.Handle("/metrics", promStatReporter.Handler())
		http.ListenAndServe(fmt.Sprintf(":%d", networking.RequestQueueMetricsPort), mux)
	}()

	statChan := make(chan *autoscaler.Stat, statReportingQueueLength)
	defer close(statChan)
	go reportStats(statChan)

	reportTicker := time.NewTicker(queue.ReporterReportingPeriod)
	defer reportTicker.Stop()
	queue.NewStats(servingPodName, queue.Channels{
		ReqChan:    reqChan,
		ReportChan: reportTicker.C,
		StatChan:   statChan,
	}, time.Now())

	adminServer := &http.Server{
		Addr:    fmt.Sprintf(":%d", networking.RequestQueueAdminPort),
		Handler: createAdminHandlers(),
	}

	timeoutHandler := queue.TimeToFirstByteTimeoutHandler(http.HandlerFunc(handler(reqChan, breaker, httpProxy)),
		time.Duration(revisionTimeoutSeconds)*time.Second, "request timeout")
	composedHandler := pushRequestMetricHandler(pushRequestLogHandler(timeoutHandler))
	// We listen on two ports to match the behavior of activator
	// so that we don't have to reprogram the k8s services.
	serverHTTP := network.NewServer(fmt.Sprintf(":%d", networking.BackendHTTPPort), composedHandler)
	serverHTTP2 := network.NewServer(fmt.Sprintf(":%d", networking.BackendHTTP2Port), composedHandler)

	errChan := make(chan error, 3)
	defer close(errChan)
	// Runs a server created by creator and sends fatal errors to the errChan.
	// Does not act on the ErrServerClosed error since that indicates we're
	// already shutting everything down.
	catchServerError := func(creator func() error) {
		if err := creator(); err != nil && err != http.ErrServerClosed {
			errChan <- err
		}
	}

	go catchServerError(serverHTTP.ListenAndServe)
	go catchServerError(serverHTTP2.ListenAndServe)
	go catchServerError(adminServer.ListenAndServe)

	// Blocks until we actually receive a TERM signal or one of the servers
	// exit unexpectedly. We fold both signals together because we only want
	// to act on the first of those to reach here.
	select {
	case err := <-errChan:
		logger.Errorw("Failed to bring up queue-proxy, shutting down.", zap.Error(err))
		flush(logger)
		os.Exit(1)
	case <-signals.SetupSignalHandler():
		logger.Info("Received TERM signal, attempting to gracefully shutdown servers.")
		healthState.Shutdown(func() {
			// Give Istio time to sync our "not ready" state.
			time.Sleep(quitSleepDuration)

			// Calling server.Shutdown() allows pending requests to
			// complete, while no new work is accepted.
			if err := serverHTTP.Shutdown(context.Background()); err != nil {
				logger.Errorw("Failed to shutdown proxy server for HTTP/1", zap.Error(err))
			}
			if err := serverHTTP2.Shutdown(context.Background()); err != nil {
				logger.Errorf("Failed to shutdown proxy server for HTTP/2", zap.Error(err))
			}
		})

		flush(logger)
		if err := adminServer.Shutdown(context.Background()); err != nil {
			logger.Errorw("Failed to shutdown admin-server", zap.Error(err))
		}
	}
}

func pushRequestLogHandler(currentHandler http.Handler) http.Handler {
	templ := os.Getenv("SERVING_REQUEST_LOG_TEMPLATE")
	if templ == "" {
		return currentHandler
	}

	revInfo := &pkghttp.RequestLogRevision{
		Name:          servingRevision,
		Namespace:     servingNamespace,
		Service:       servingService,
		Configuration: servingConfig,
		PodName:       servingPodName,
		PodIP:         servingPodIP,
	}
	handler, err := pkghttp.NewRequestLogHandler(currentHandler, logging.NewSyncFileWriter(os.Stdout), templ,
		pkghttp.RequestLogTemplateInputGetterFromRevision(revInfo))

	if err != nil {
		logger.Errorw("Error setting up request logger. Request logs will be unavailable.", zap.Error(err))
		return currentHandler
	}
	return handler
}

func pushRequestMetricHandler(currentHandler http.Handler) http.Handler {
	backend := os.Getenv("SERVING_REQUEST_METRICS_BACKEND")
	logger.Infof("SERVING_REQUEST_METRICS_BACKEND=%v", backend)
	if backend == "" {
		return currentHandler
	}

	r, err := queuestats.NewStatsReporter(servingNamespace, servingService, servingConfig, servingRevision)
	if err != nil {
		logger.Errorw("Error setting up request metrics reporter. Request metrics will be unavailable.", zap.Error(err))
		return currentHandler
	}

	// Set up OpenCensus exporter.
	// NOTE: We use revision as the component instead of queue because queue is
	// implementation specific. The current metrics are request relative. Using
	// revision is reasonable.
	// TODO(yanweiguo): add the ability to emit metrics with names not combined
	// to component.
	ops := metrics.ExporterOptions{
		Domain:         "knative.dev/serving",
		Component:      "revision",
		PrometheusPort: commonMetricsPort,
		ConfigMap: map[string]string{
			metrics.BackendDestinationKey: backend,
		},
	}
	err = metrics.UpdateExporter(ops, logger)
	if err != nil {
		logger.Errorw("Error setting up request metrics exporter. Request metrics will be unavailable.", zap.Error(err))
		return currentHandler
	}

	handler, err := queue.NewRequestMetricHandler(currentHandler, r)
	if err != nil {
		logger.Errorw("Error setting up request metrics handler. Request metrics will be unavailable.", zap.Error(err))
		return currentHandler
	}
	return handler
}

func flush(logger *zap.SugaredLogger) {
	logger.Sync()
	os.Stdout.Sync()
	os.Stderr.Sync()
	metrics.FlushExporter()
}
<commit_msg>Flatten out handler func creation in queue (#3967)<commit_after>/*
Copyright 2018 The Knative 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 main

import (
	"context"
	"flag"
	"fmt"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"time"

	"go.uber.org/zap"

	"github.com/knative/pkg/logging/logkey"
	"github.com/knative/pkg/metrics"
	"github.com/knative/pkg/signals"
	"github.com/knative/serving/cmd/util"
	"github.com/knative/serving/pkg/activator"
	activatorutil "github.com/knative/serving/pkg/activator/util"
	"github.com/knative/serving/pkg/apis/networking"
	"github.com/knative/serving/pkg/autoscaler"
	pkghttp "github.com/knative/serving/pkg/http"
	"github.com/knative/serving/pkg/logging"
	"github.com/knative/serving/pkg/network"
	"github.com/knative/serving/pkg/queue"
	"github.com/knative/serving/pkg/queue/health"
	queuestats "github.com/knative/serving/pkg/queue/stats"

	"k8s.io/apimachinery/pkg/util/wait"
)

const (
	// Add a little buffer space between request handling and stat
	// reporting so that latency in the stat pipeline doesn't
	// interfere with request handling.
	statReportingQueueLength = 10

	// Add enough buffer to not block request serving on stats collection
	requestCountingQueueLength = 100

	// Duration the /quitquitquit handler should wait before returning.
	// This is to give Istio a little bit more time to remove the pod
	// from its configuration and propagate that to all istio-proxies
	// in the mesh.
	quitSleepDuration = 20 * time.Second

	// commonMetricsPort is the port where common metrics, e.g. request metrics
	// are exposed in Prometheus. This is different from the metrics used
	// for autoscaling, which are exposed in 9090.
	commonMetricsPort = 9091

	badProbeTemplate = "unexpected probe header value: %s"
)

var (
	servingService         string
	servingConfig          string
	servingNamespace       string
	servingRevision        string
	servingRevisionKey     string
	servingPodIP           string
	servingPodName         string
	userTargetPort         int
	userTargetAddress      string
	containerConcurrency   int
	revisionTimeoutSeconds int
	reqChan                = make(chan queue.ReqEvent, requestCountingQueueLength)
	logger                 *zap.SugaredLogger
	breaker                *queue.Breaker

	httpProxy *httputil.ReverseProxy

	healthState      = &health.State{}
	promStatReporter *queue.PrometheusStatsReporter // Prometheus stats reporter.
)

func initEnv() {
	servingService = os.Getenv("SERVING_SERVICE") // KService is optional
	servingConfig = util.GetRequiredEnvOrFatal("SERVING_CONFIGURATION", logger)
	servingNamespace = util.GetRequiredEnvOrFatal("SERVING_NAMESPACE", logger)
	servingRevision = util.GetRequiredEnvOrFatal("SERVING_REVISION", logger)
	servingPodIP = util.GetRequiredEnvOrFatal("SERVING_POD_IP", logger)
	servingPodName = util.GetRequiredEnvOrFatal("SERVING_POD", logger)
	containerConcurrency = util.MustParseIntEnvOrFatal("CONTAINER_CONCURRENCY", logger)
	revisionTimeoutSeconds = util.MustParseIntEnvOrFatal("REVISION_TIMEOUT_SECONDS", logger)
	userTargetPort = util.MustParseIntEnvOrFatal("USER_PORT", logger)
	userTargetAddress = fmt.Sprintf("127.0.0.1:%d", userTargetPort)

	// TODO(mattmoor): Move this key to be in terms of the KPA.
	servingRevisionKey = autoscaler.NewMetricKey(servingNamespace, servingRevision)
	_psr, err := queue.NewPrometheusStatsReporter(servingNamespace, servingConfig, servingRevision, servingPodName)
	if err != nil {
		logger.Fatalw("Failed to create stats reporter", zap.Error(err))
	}
	promStatReporter = _psr
}

func reportStats(statChan chan *autoscaler.Stat) {
	for {
		s := <-statChan
		if err := promStatReporter.Report(s); err != nil {
			logger.Errorw("Error while sending stat", zap.Error(err))
		}
	}
}

func knativeProbeHeader(r *http.Request) string {
	return r.Header.Get(network.ProbeHeaderName)
}

func knativeProxyHeader(r *http.Request) string {
	return r.Header.Get(network.ProxyHeaderName)
}

func probeUserContainer() bool {
	var err error
	wait.PollImmediate(50*time.Millisecond, 10*time.Second, func() (bool, error) {
		logger.Debug("TCP probing the user-container.")
		err = health.TCPProbe(userTargetAddress, 100*time.Millisecond)
		return err == nil, nil
	})

	if err == nil {
		logger.Info("User-container successfully probed.")
	} else {
		logger.Errorw("User-container could not be probed successfully.", zap.Error(err))
	}

	return err == nil
}

// Make handler a closure for testing.
func handler(reqChan chan queue.ReqEvent, breaker *queue.Breaker, proxy *httputil.ReverseProxy) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		ph := knativeProbeHeader(r)
		switch {
		case ph != "":
			if ph != queue.Name {
				http.Error(w, fmt.Sprintf(badProbeTemplate, ph), http.StatusBadRequest)
				return
			}
			if probeUserContainer() {
				// Respond with the name of the component handling the request.
				w.Write([]byte(queue.Name))
			} else {
				http.Error(w, "container not ready", http.StatusServiceUnavailable)
			}
			return
		case network.IsKubeletProbe(r):
			// Do not count health checks for concurrency metrics
			proxy.ServeHTTP(w, r)
			return
		}

		// Metrics for autoscaling.
		h := knativeProxyHeader(r)
		in, out := queue.ReqIn, queue.ReqOut
		if activator.Name == h {
			in, out = queue.ProxiedIn, queue.ProxiedOut
		}
		reqChan <- queue.ReqEvent{Time: time.Now(), EventType: in}
		defer func() {
			reqChan <- queue.ReqEvent{Time: time.Now(), EventType: out}
		}()
		network.RewriteHostOut(r)

		// Enforce queuing and concurrency limits.
		if breaker != nil {
			ok := breaker.Maybe(func() {
				proxy.ServeHTTP(w, r)
			})
			if !ok {
				http.Error(w, "overload", http.StatusServiceUnavailable)
			}
		} else {
			proxy.ServeHTTP(w, r)
		}
	}
}

// Sets up /health and /wait-for-drain endpoints.
func createAdminHandlers() *http.ServeMux {
	mux := http.NewServeMux()
	mux.HandleFunc(queue.RequestQueueHealthPath, healthState.HealthHandler(probeUserContainer))
	mux.HandleFunc(queue.RequestQueueDrainPath, healthState.DrainHandler())

	return mux
}

func main() {
	flag.Parse()
	logger, _ = logging.NewLogger(os.Getenv("SERVING_LOGGING_CONFIG"), os.Getenv("SERVING_LOGGING_LEVEL"))
	logger = logger.Named("queueproxy")
	defer flush(logger)

	initEnv()
	logger = logger.With(
		zap.String(logkey.Key, servingRevisionKey),
		zap.String(logkey.Pod, servingPodName))

	target, err := url.Parse("http://" + userTargetAddress)
	if err != nil {
		logger.Fatalw("Failed to parse localhost URL", zap.Error(err))
	}

	httpProxy = httputil.NewSingleHostReverseProxy(target)
	httpProxy.Transport = network.AutoTransport
	httpProxy.FlushInterval = -1

	activatorutil.SetupHeaderPruning(httpProxy)

	// If containerConcurrency == 0 then concurrency is unlimited.
	if containerConcurrency > 0 {
		// We set the queue depth to be equal to the container concurrency but at least 10 to
		// allow the autoscaler to get a strong enough signal.
		queueDepth := containerConcurrency
		if queueDepth < 10 {
			queueDepth = 10
		}
		params := queue.BreakerParams{QueueDepth: queueDepth, MaxConcurrency: containerConcurrency, InitialCapacity: containerConcurrency}
		breaker = queue.NewBreaker(params)
		logger.Infof("Queue container is starting with %#v", params)
	}

	go func() {
		mux := http.NewServeMux()
		mux.Handle("/metrics", promStatReporter.Handler())
		http.ListenAndServe(fmt.Sprintf(":%d", networking.RequestQueueMetricsPort), mux)
	}()

	statChan := make(chan *autoscaler.Stat, statReportingQueueLength)
	defer close(statChan)
	go reportStats(statChan)

	reportTicker := time.NewTicker(queue.ReporterReportingPeriod)
	defer reportTicker.Stop()
	queue.NewStats(servingPodName, queue.Channels{
		ReqChan:    reqChan,
		ReportChan: reportTicker.C,
		StatChan:   statChan,
	}, time.Now())

	adminServer := &http.Server{
		Addr:    fmt.Sprintf(":%d", networking.RequestQueueAdminPort),
		Handler: createAdminHandlers(),
	}

	// Create queue handler chain
	// Note: innermost handlers are specified first, ie. the last handler in the chain will be executed first
	var composedHandler http.Handler = http.HandlerFunc(handler(reqChan, breaker, httpProxy))
	composedHandler = queue.TimeToFirstByteTimeoutHandler(composedHandler,
		time.Duration(revisionTimeoutSeconds)*time.Second, "request timeout")
	composedHandler = pushRequestLogHandler(composedHandler)
	composedHandler = pushRequestMetricHandler(composedHandler)
	// We listen on two ports to match the behavior of activator
	// so that we don't have to reprogram the k8s services.
	serverHTTP := network.NewServer(fmt.Sprintf(":%d", networking.BackendHTTPPort), composedHandler)
	serverHTTP2 := network.NewServer(fmt.Sprintf(":%d", networking.BackendHTTP2Port), composedHandler)

	errChan := make(chan error, 3)
	defer close(errChan)
	// Runs a server created by creator and sends fatal errors to the errChan.
	// Does not act on the ErrServerClosed error since that indicates we're
	// already shutting everything down.
	catchServerError := func(creator func() error) {
		if err := creator(); err != nil && err != http.ErrServerClosed {
			errChan <- err
		}
	}

	go catchServerError(serverHTTP.ListenAndServe)
	go catchServerError(serverHTTP2.ListenAndServe)
	go catchServerError(adminServer.ListenAndServe)

	// Blocks until we actually receive a TERM signal or one of the servers
	// exit unexpectedly. We fold both signals together because we only want
	// to act on the first of those to reach here.
	select {
	case err := <-errChan:
		logger.Errorw("Failed to bring up queue-proxy, shutting down.", zap.Error(err))
		flush(logger)
		os.Exit(1)
	case <-signals.SetupSignalHandler():
		logger.Info("Received TERM signal, attempting to gracefully shutdown servers.")
		healthState.Shutdown(func() {
			// Give Istio time to sync our "not ready" state.
			time.Sleep(quitSleepDuration)

			// Calling server.Shutdown() allows pending requests to
			// complete, while no new work is accepted.
			if err := serverHTTP.Shutdown(context.Background()); err != nil {
				logger.Errorw("Failed to shutdown proxy server for HTTP/1", zap.Error(err))
			}
			if err := serverHTTP2.Shutdown(context.Background()); err != nil {
				logger.Errorf("Failed to shutdown proxy server for HTTP/2", zap.Error(err))
			}
		})

		flush(logger)
		if err := adminServer.Shutdown(context.Background()); err != nil {
			logger.Errorw("Failed to shutdown admin-server", zap.Error(err))
		}
	}
}

func pushRequestLogHandler(currentHandler http.Handler) http.Handler {
	templ := os.Getenv("SERVING_REQUEST_LOG_TEMPLATE")
	if templ == "" {
		return currentHandler
	}

	revInfo := &pkghttp.RequestLogRevision{
		Name:          servingRevision,
		Namespace:     servingNamespace,
		Service:       servingService,
		Configuration: servingConfig,
		PodName:       servingPodName,
		PodIP:         servingPodIP,
	}
	handler, err := pkghttp.NewRequestLogHandler(currentHandler, logging.NewSyncFileWriter(os.Stdout), templ,
		pkghttp.RequestLogTemplateInputGetterFromRevision(revInfo))

	if err != nil {
		logger.Errorw("Error setting up request logger. Request logs will be unavailable.", zap.Error(err))
		return currentHandler
	}
	return handler
}

func pushRequestMetricHandler(currentHandler http.Handler) http.Handler {
	backend := os.Getenv("SERVING_REQUEST_METRICS_BACKEND")
	logger.Infof("SERVING_REQUEST_METRICS_BACKEND=%v", backend)
	if backend == "" {
		return currentHandler
	}

	r, err := queuestats.NewStatsReporter(servingNamespace, servingService, servingConfig, servingRevision)
	if err != nil {
		logger.Errorw("Error setting up request metrics reporter. Request metrics will be unavailable.", zap.Error(err))
		return currentHandler
	}

	// Set up OpenCensus exporter.
	// NOTE: We use revision as the component instead of queue because queue is
	// implementation specific. The current metrics are request relative. Using
	// revision is reasonable.
	// TODO(yanweiguo): add the ability to emit metrics with names not combined
	// to component.
	ops := metrics.ExporterOptions{
		Domain:         "knative.dev/serving",
		Component:      "revision",
		PrometheusPort: commonMetricsPort,
		ConfigMap: map[string]string{
			metrics.BackendDestinationKey: backend,
		},
	}
	err = metrics.UpdateExporter(ops, logger)
	if err != nil {
		logger.Errorw("Error setting up request metrics exporter. Request metrics will be unavailable.", zap.Error(err))
		return currentHandler
	}

	handler, err := queue.NewRequestMetricHandler(currentHandler, r)
	if err != nil {
		logger.Errorw("Error setting up request metrics handler. Request metrics will be unavailable.", zap.Error(err))
		return currentHandler
	}
	return handler
}

func flush(logger *zap.SugaredLogger) {
	logger.Sync()
	os.Stdout.Sync()
	os.Stderr.Sync()
	metrics.FlushExporter()
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"os"

	"github.com/BurntSushi/toml"
	"github.com/ensonmj/expect"
)

// Host config
type Host struct {
	HostName string
	User     string
	Pass     string
	Opt      string
}

// Conf for relay.toml
type Conf struct {
	Relay     string
	RelayUser string
	RelayPass string
	Hosts     map[string]Host
}

var (
	fVerbose bool
)

func init() {
	flag.BoolVar(&fVerbose, "v", false, "open debug log")
}

func main() {
	var conf Conf
	if _, err := toml.DecodeFile("relay.toml", &conf); err != nil {
		fmt.Println(err)
		return
	}

	flag.Parse()
	abbr := flag.Arg(0)
	remoteCmd := flag.Arg(1)
	host, ok := conf.Hosts[abbr]
	if !ok {
		errExit(fmt.Errorf("can't find %s\n from config", abbr))
	}

	cmd := fmt.Sprintf("zssh -t %s@%s", conf.RelayUser, conf.Relay)
	child, err := expect.Spawn(cmd)
	if err != nil {
		errExit(err)
	}
	if fVerbose {
		child.Debug(true)
	}

	pairs := []expect.ExpectPair{
		{"password", func(_ []byte) error {
			err := expect.WithoutEcho(func() error {
				pass := conf.RelayPass
				if pass == "" {
					err := child.SendLineUser("\nPlease enter your password:")
					if err != nil {
						return err
					}
					pass, err = child.ReadLineUser()
					if err != nil {
						return err
					}
				}
				return child.SendLine(pass)
			})
			if err != nil {
				return err
			}
			err = child.SendLineUser("Checking password...")
			if err != nil {
				return err
			}

			err = child.Expect("bash-baidu-ssl$")
			if err != nil {
				return err
			}
			return loginOrRun(child, host, remoteCmd)
		}},
		{"bash-baidu-ssl$", func(_ []byte) error {
			return loginOrRun(child, host, remoteCmd)
		}},
	}
	err = child.ExpectMulti(pairs)
	if err != nil {
		errExit(err)
	}
}

func loginOrRun(child *expect.ExpectSubproc, host Host, remoteCmd string) error {
	var cmd string
	if host.User != "" {
		cmd = fmt.Sprintf("ssh %s %s@%s", host.Opt, host.User, host.HostName)
	} else {
		cmd = fmt.Sprintf("ssh %s %s", host.Opt, host.HostName)
	}
	if remoteCmd != "" {
		cmd = fmt.Sprintf("%s %s", cmd, remoteCmd)
	}
	fmt.Println(cmd)
	err := child.SendLine(cmd)
	if err != nil {
		return err
	}

	if host.Pass != "" {
		err := child.Expect("password")
		if err != nil {
			return err
		}
		err = child.SendLine(host.Pass)
		if err != nil {
			return err
		}
	}

	if remoteCmd == "" {
		child.Interact()
	}

	return nil
}

func errExit(err error) {
	fmt.Fprintln(os.Stderr, err)
	os.Exit(1)
}
<commit_msg>support hostname in cli<commit_after>package main

import (
	"flag"
	"fmt"
	"net/url"
	"os"

	"github.com/BurntSushi/toml"
	"github.com/ensonmj/expect"
)

// Host config
type Host struct {
	HostName string
	User     string
	Pass     string
	Opt      string
}

// Conf for relay.toml
type Conf struct {
	Relay     string
	RelayUser string
	RelayPass string
	Hosts     map[string]Host
}

var (
	fVerbose bool
)

func init() {
	flag.BoolVar(&fVerbose, "v", false, "open debug log")
}

func main() {
	var conf Conf
	if _, err := toml.DecodeFile("relay.toml", &conf); err != nil {
		fmt.Println(err)
		return
	}

	flag.Parse()

	abbrOrURL := flag.Arg(0)
	host, ok := conf.Hosts[abbrOrURL]
	if !ok {
		u, err := url.Parse(abbrOrURL)
		if err != nil {
			errExit(err)
		}
		hostName := u.Hostname()
		user := u.User.Username()
		pass, _ := u.User.Password()
		var opt string
		if u.Scheme != "" {
			opt = "--" + u.Scheme
		} else {
			q, err := url.ParseQuery(u.RawQuery)
			if err != nil {
				errExit(err)
			}
			opt = q["opt"][0]
		}
		host = Host{
			HostName: hostName,
			User:     user,
			Pass:     pass,
			Opt:      opt,
		}
	}
	remoteCmd := flag.Arg(1)

	cmd := fmt.Sprintf("zssh -t %s@%s", conf.RelayUser, conf.Relay)
	child, err := expect.Spawn(cmd)
	if err != nil {
		errExit(err)
	}
	if fVerbose {
		child.Debug(true)
	}

	pairs := []expect.ExpectPair{
		{"password", func(_ []byte) error {
			err := expect.WithoutEcho(func() error {
				pass := conf.RelayPass
				if pass == "" {
					err := child.SendLineUser("\nPlease enter your password:")
					if err != nil {
						return err
					}
					pass, err = child.ReadLineUser()
					if err != nil {
						return err
					}
				}
				return child.SendLine(pass)
			})
			if err != nil {
				return err
			}
			err = child.SendLineUser("Checking password...")
			if err != nil {
				return err
			}

			err = child.Expect("bash-baidu-ssl$")
			if err != nil {
				return err
			}
			return loginOrRun(child, host, remoteCmd)
		}},
		{"bash-baidu-ssl$", func(_ []byte) error {
			return loginOrRun(child, host, remoteCmd)
		}},
	}
	err = child.ExpectMulti(pairs)
	if err != nil {
		errExit(err)
	}
}

func loginOrRun(child *expect.ExpectSubproc, host Host, remoteCmd string) error {
	var cmd string
	if host.User != "" {
		cmd = fmt.Sprintf("ssh %s %s@%s", host.Opt, host.User, host.HostName)
	} else {
		cmd = fmt.Sprintf("ssh %s %s", host.Opt, host.HostName)
	}
	if remoteCmd != "" {
		cmd = fmt.Sprintf("%s %s", cmd, remoteCmd)
	}
	fmt.Println(cmd)
	err := child.SendLine(cmd)
	if err != nil {
		return err
	}

	if host.Pass != "" {
		err := child.Expect("password")
		if err != nil {
			return err
		}
		err = child.SendLine(host.Pass)
		if err != nil {
			return err
		}
	}

	if remoteCmd == "" {
		child.Interact()
	}

	return nil
}

func errExit(err error) {
	fmt.Fprintln(os.Stderr, err)
	os.Exit(1)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"encoding/pem"
	"flag"
	"io/ioutil"
	"os"
	"strings"

	log "github.com/Sirupsen/logrus"
	"github.com/zmap/zcrypto/x509"
	"github.com/zmap/zlint"
)

var ( //flags
	listLintsJSON bool
	prettyprint   bool
	format        string
)

func init() {
	flag.BoolVar(&listLintsJSON, "list-lints-json", false, "Use this flag to print supported lints in JSON format, one per line")
	flag.StringVar(&format, "format", "pem", "One of {pem, der, base64}")
	flag.BoolVar(&prettyprint, "pretty", false, "Pretty-print output")
	flag.Parse()

	log.SetLevel(log.InfoLevel)
}

func main() {

	if listLintsJSON {
		zlint.EncodeLintDescriptionsToJSON(os.Stdout)
		return
	}

	var inform = strings.ToLower(format)
	if flag.NArg() < 1 || flag.Arg(0) == "-" {
		lint(os.Stdin, inform)
	} else {
		for _, filePath := range flag.Args() {
			var inputFile *os.File
			var err error
			inputFile, err = os.Open(filePath)
			if err != nil {
				log.Fatalf("unable to open file %s: %s", filePath, err)
			}
			var fmt = inform
			switch {
			case strings.HasSuffix(filePath, ".der"):
				fmt = "der"
			case strings.HasSuffix(filePath, ".pem"):
				fmt = "pem"
			}
			lint(inputFile, fmt)
			inputFile.Close()
		}
	}
}

func lint(inputFile *os.File, inform string) {
	fileBytes, err := ioutil.ReadAll(inputFile)
	if err != nil {
		log.Fatalf("unable to read file %s: %s", inputFile.Name(), err)
	}

	var asn1Data []byte
	switch inform {
	case "pem":
		p, _ := pem.Decode(fileBytes)
		if p == nil || p.Type != "CERTIFICATE" {
			log.Fatal("unable to parse PEM")
		}
		asn1Data = p.Bytes
	case "der":
		asn1Data = fileBytes
	case "base64":
		asn1Data, err = base64.StdEncoding.DecodeString(string(fileBytes))
		if err != nil {
			log.Fatalf("unable to parse base64: %s", err)
		}
	default:
		log.Fatalf("unknown input format %s", format)
	}

	c, err := x509.ParseCertificate(asn1Data)
	if err != nil {
		log.Fatalf("unable to parse certificate: %s", err)
	}

	zlintResult := zlint.LintCertificate(c)
	jsonBytes, err := json.Marshal(zlintResult.ZLint)
	if err != nil {
		log.Fatalf("unable to encode lints JSON: %s", err)
	}
	if prettyprint {
		var out bytes.Buffer
		if err := json.Indent(&out, jsonBytes, "", " "); err != nil {
			log.Fatalf("can't format output: %s", err)
		}
		os.Stdout.Write(out.Bytes())
	} else {
		os.Stdout.Write(jsonBytes)
	}
	os.Stdout.Write([]byte{'\n'})
	os.Stdout.Sync()
}
<commit_msg>Improve help message (#147)<commit_after>package main

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"encoding/pem"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"strings"

	log "github.com/Sirupsen/logrus"
	"github.com/zmap/zcrypto/x509"
	"github.com/zmap/zlint"
)

var ( //flags
	listLintsJSON bool
	prettyprint   bool
	format        string
)

func init() {
	flag.BoolVar(&listLintsJSON, "list-lints-json", false, "Use this flag to print supported lints in JSON format, one per line")
	flag.StringVar(&format, "format", "pem", "One of {pem, der, base64}")
	flag.BoolVar(&prettyprint, "pretty", false, "Pretty-print output")
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage: %s [flags] file...\n", os.Args[0])
		flag.PrintDefaults()
	}
	flag.Parse()
	log.SetLevel(log.InfoLevel)
}

func main() {

	if listLintsJSON {
		zlint.EncodeLintDescriptionsToJSON(os.Stdout)
		return
	}

	var inform = strings.ToLower(format)
	if flag.NArg() < 1 || flag.Arg(0) == "-" {
		lint(os.Stdin, inform)
	} else {
		for _, filePath := range flag.Args() {
			var inputFile *os.File
			var err error
			inputFile, err = os.Open(filePath)
			if err != nil {
				log.Fatalf("unable to open file %s: %s", filePath, err)
			}
			var fmt = inform
			switch {
			case strings.HasSuffix(filePath, ".der"):
				fmt = "der"
			case strings.HasSuffix(filePath, ".pem"):
				fmt = "pem"
			}
			lint(inputFile, fmt)
			inputFile.Close()
		}
	}
}

func lint(inputFile *os.File, inform string) {
	fileBytes, err := ioutil.ReadAll(inputFile)
	if err != nil {
		log.Fatalf("unable to read file %s: %s", inputFile.Name(), err)
	}

	var asn1Data []byte
	switch inform {
	case "pem":
		p, _ := pem.Decode(fileBytes)
		if p == nil || p.Type != "CERTIFICATE" {
			log.Fatal("unable to parse PEM")
		}
		asn1Data = p.Bytes
	case "der":
		asn1Data = fileBytes
	case "base64":
		asn1Data, err = base64.StdEncoding.DecodeString(string(fileBytes))
		if err != nil {
			log.Fatalf("unable to parse base64: %s", err)
		}
	default:
		log.Fatalf("unknown input format %s", format)
	}

	c, err := x509.ParseCertificate(asn1Data)
	if err != nil {
		log.Fatalf("unable to parse certificate: %s", err)
	}

	zlintResult := zlint.LintCertificate(c)
	jsonBytes, err := json.Marshal(zlintResult.ZLint)
	if err != nil {
		log.Fatalf("unable to encode lints JSON: %s", err)
	}
	if prettyprint {
		var out bytes.Buffer
		if err := json.Indent(&out, jsonBytes, "", " "); err != nil {
			log.Fatalf("can't format output: %s", err)
		}
		os.Stdout.Write(out.Bytes())
	} else {
		os.Stdout.Write(jsonBytes)
	}
	os.Stdout.Write([]byte{'\n'})
	os.Stdout.Sync()
}
<|endoftext|>


<commit_before>// Package combiner contains methods to consume Job objects, orchestrating
// the download, combination, and upload of a group of related PDF files.
// TODO now that s.Stat{} can return errors, maybe the err chan is unnecessary
package combiner

import (
	"errors"
	"fmt"
	"github.com/PeopleAdmin/pdfcombiner/cpdf"
	"github.com/PeopleAdmin/pdfcombiner/job"
	"github.com/PeopleAdmin/pdfcombiner/notifier"
	s "github.com/PeopleAdmin/pdfcombiner/stat"
	"io/ioutil"
	"log"
	"math/rand"
	"os"
	"runtime"
	"strings"
	"time"
)

var (
	downloadTimeout = 3 * time.Minute
	basedir         = "/tmp/"
	maxGoroutines   = 30
)

// Combine is the entry point to this package.  Given a Job, downloads
// all the files, combines them into a single one, uploads it to AWS
// and posts the status to a callback endpoint.
func Combine(j *job.Job) {
	defer postToCallback(j)
	saveDir := mkTmpDir()
	getAllFiles(j, saveDir)
	if !j.HasDownloadedDocs() {
		return
	}

	componentPaths := fsPathsOf(j.Downloaded, saveDir)
	combinedPath := saveDir + "combined.pdf"
	err := cpdf.Merge(componentPaths, combinedPath, j.Title)
	if err != nil {
		j.AddError("general", err)
		return
	}
	j.UploadCombinedFile(combinedPath)
	return
}

// Get an individual file from S3.  If successful, writes the file out to disk
// and sends a stat object back to the main channel.  If there are errors they
// are sent back through the error channel.
func getFile(j *job.Job, docname string, dir string, c chan<- s.Stat, e chan<- s.Stat) {
	start := time.Now()
	data, err := j.Get(docname)
	if err != nil {
		e <- s.Stat{Filename: docname, Err: err}
		return
	}
	path := localPath(dir, docname)
	err = ioutil.WriteFile(path, data, 0644)
	if err != nil {
		e <- s.Stat{Filename: docname, Err: err}
		return
	}
	pagecount := cpdf.PageCount(path)
	c <- s.Stat{Filename: docname,
		Size:      len(data),
		PageCount: pagecount,
		DlTime:    time.Since(start)}
}

// Fan out workers to download each document in parallel, then block
// until all downloads are complete.
func getAllFiles(j *job.Job, dir string) {
	start := time.Now()
	c := make(chan s.Stat, j.DocCount())
	e := make(chan s.Stat, j.DocCount())
	for _, doc := range j.DocList {
		throttle()
		go getFile(j, doc.Name, dir, c, e)
	}

	totalBytes := waitForDownloads(j, c, e)
	printSummary(start, totalBytes, j.CompleteCount())
}

// Prevents the system from being overwhelmed with work.
// Blocks until the number of Goroutines is less than a preset threshold.
func throttle() {
	for {
		if runtime.NumGoroutine() < maxGoroutines {
			break
		}
		time.Sleep(50 * time.Millisecond)
	}
}

// Listen on several channels for information from background download
// tasks -- each task will either send a s.Stat through c, an error through
// e, or timeout.  Once all docs are accounted for, return the total number
// of bytes recieved.
func waitForDownloads(j *job.Job, c <-chan s.Stat, e <-chan s.Stat) (totalBytes int) {
	for _ = range j.DocList {
		select {
		case packet := <-c:
			log.Printf("%s was %d bytes\n", packet.Filename, packet.Size)
			totalBytes += packet.Size
			j.MarkComplete(packet.Filename, packet)
		case bad := <-e:
			j.AddError(bad.Filename, bad.Err)
		case <-time.After(downloadTimeout):
			j.AddError("general", errors.New("timed out while downloading"))
			return
		}
	}
	return
}

// Print a summary of the download activity.
func printSummary(start time.Time, bytes int, count int) {
	seconds := time.Since(start).Seconds()
	mbps := float64(bytes) / 1024 / 1024 / seconds
	log.Printf("got %d bytes over %d files in %f secs (%f MB/s)\n",
		bytes, count, seconds, mbps)
}

// Send an update on the success or failure of the operation to the
// callback URL provided by the job originator.
func postToCallback(j *job.Job) {
	log.Println("work complete, posting status to callback:", j.Callback)
	notifier.SendNotification(j)
}

// Make and return a randomized temporary directory.
func mkTmpDir() (dirname string) {
	rand.Seed(time.Now().UnixNano())
	dirname = fmt.Sprintf("/tmp/pdfcombiner/%d/", rand.Int())
	os.MkdirAll(dirname, 0777)
	return
}

// Get the absolute paths to a list of docs.
func fsPathsOf(docs []string, dir string) (paths []string) {
	paths = make([]string, len(docs))
	for idx, doc := range docs {
		paths[idx] = localPath(dir, doc)
	}
	return
}

// localPath replaces any s3 key directory markers with underscores so
// we don't need to recursively create directories when saving files.
func localPath(dir, remotePath string) string {
	return dir + strings.Replace(remotePath, "/", "_", -1)
}
<commit_msg>Add extra performance stats<commit_after>// Package combiner contains methods to consume Job objects, orchestrating
// the download, combination, and upload of a group of related PDF files.
// TODO now that s.Stat{} can return errors, maybe the err chan is unnecessary
package combiner

import (
	"errors"
	"fmt"
	"github.com/PeopleAdmin/pdfcombiner/cpdf"
	"github.com/PeopleAdmin/pdfcombiner/job"
	"github.com/PeopleAdmin/pdfcombiner/notifier"
	s "github.com/PeopleAdmin/pdfcombiner/stat"
	"io/ioutil"
	"log"
	"math/rand"
	"os"
	"runtime"
	"strings"
	"time"
)

var (
	downloadTimeout = 3 * time.Minute
	basedir         = "/tmp/"
	maxGoroutines   = 30
)

// Combine is the entry point to this package.  Given a Job, downloads
// all the files, combines them into a single one, uploads it to AWS
// and posts the status to a callback endpoint.
func Combine(j *job.Job) {
	defer postToCallback(j)
	startAll := time.Now()
	saveDir := mkTmpDir()
	getAllFiles(j, saveDir)
	if !j.HasDownloadedDocs() {
		return
	}

	componentPaths := fsPathsOf(j.Downloaded, saveDir)
	combinedPath := saveDir + "combined.pdf"
	err := cpdf.Merge(componentPaths, combinedPath, j.Title)
	if err != nil {
		j.AddError("general", err)
		return
	}
	startUpload := time.Now()
	j.UploadCombinedFile(combinedPath)
	j.PerfStats["upload"] = s.Stat{DlTime: time.Since(startUpload)}
	j.PerfStats["total"] = s.Stat{DlTime: time.Since(startAll)}
	return
}

// Get an individual file from S3.  If successful, writes the file out to disk
// and sends a stat object back to the main channel.  If there are errors they
// are sent back through the error channel.
func getFile(j *job.Job, docname string, dir string, c chan<- s.Stat, e chan<- s.Stat) {
	start := time.Now()
	data, err := j.Get(docname)
	if err != nil {
		e <- s.Stat{Filename: docname, Err: err}
		return
	}
	path := localPath(dir, docname)
	err = ioutil.WriteFile(path, data, 0644)
	if err != nil {
		e <- s.Stat{Filename: docname, Err: err}
		return
	}
	pagecount := cpdf.PageCount(path)
	c <- s.Stat{Filename: docname,
		Size:      len(data),
		PageCount: pagecount,
		DlTime:    time.Since(start)}
}

// Fan out workers to download each document in parallel, then block
// until all downloads are complete.
func getAllFiles(j *job.Job, dir string) {
	start := time.Now()
	c := make(chan s.Stat, j.DocCount())
	e := make(chan s.Stat, j.DocCount())
	for _, doc := range j.DocList {
		throttle()
		go getFile(j, doc.Name, dir, c, e)
	}

	totalBytes := waitForDownloads(j, c, e)
	printSummary(start, totalBytes, j.CompleteCount())
}

// Prevents the system from being overwhelmed with work.
// Blocks until the number of Goroutines is less than a preset threshold.
func throttle() {
	for {
		if runtime.NumGoroutine() < maxGoroutines {
			break
		}
		time.Sleep(50 * time.Millisecond)
	}
}

// Listen on several channels for information from background download
// tasks -- each task will either send a s.Stat through c, an error through
// e, or timeout.  Once all docs are accounted for, return the total number
// of bytes recieved.
func waitForDownloads(j *job.Job, c <-chan s.Stat, e <-chan s.Stat) (totalBytes int) {
	for _ = range j.DocList {
		select {
		case packet := <-c:
			log.Printf("%s was %d bytes\n", packet.Filename, packet.Size)
			totalBytes += packet.Size
			j.MarkComplete(packet.Filename, packet)
		case bad := <-e:
			j.AddError(bad.Filename, bad.Err)
		case <-time.After(downloadTimeout):
			j.AddError("general", errors.New("timed out while downloading"))
			return
		}
	}
	return
}

// Print a summary of the download activity.
func printSummary(start time.Time, bytes int, count int) {
	seconds := time.Since(start).Seconds()
	mbps := float64(bytes) / 1024 / 1024 / seconds
	log.Printf("got %d bytes over %d files in %f secs (%f MB/s)\n",
		bytes, count, seconds, mbps)
}

// Send an update on the success or failure of the operation to the
// callback URL provided by the job originator.
func postToCallback(j *job.Job) {
	log.Println("work complete, posting status to callback:", j.Callback)
	notifier.SendNotification(j)
}

// Make and return a randomized temporary directory.
func mkTmpDir() (dirname string) {
	rand.Seed(time.Now().UnixNano())
	dirname = fmt.Sprintf("/tmp/pdfcombiner/%d/", rand.Int())
	os.MkdirAll(dirname, 0777)
	return
}

// Get the absolute paths to a list of docs.
func fsPathsOf(docs []string, dir string) (paths []string) {
	paths = make([]string, len(docs))
	for idx, doc := range docs {
		paths[idx] = localPath(dir, doc)
	}
	return
}

// localPath replaces any s3 key directory markers with underscores so
// we don't need to recursively create directories when saving files.
func localPath(dir, remotePath string) string {
	return dir + strings.Replace(remotePath, "/", "_", -1)
}
<|endoftext|>


<commit_before>package identity

import (
	"net/url"
)

type Null struct {
	events chan *Identity
}

func init() {
	register["null"] = &Null{}
}

func (s *Null) Init(u *url.URL, c chan *Identity) error {
	s.events = c
	return nil
}

func (s *Null) Listen() {
	<- s.events
}

func (s *Null) Close() error {
	return nil
}
<commit_msg>Renamed identity.Null to identity.NullStorage<commit_after>package identity

import (
	"net/url"
)

type NullStorage struct {
	events chan *Identity
}

func init() {
	register["null"] = &NullStorage{}
}

func (s *NullStorage) Init(u *url.URL, c chan *Identity) error {
	s.events = c
	return nil
}

func (s *NullStorage) Listen() {
	<- s.events
}

func (s *NullStorage) Close() error {
	return nil
}
<|endoftext|>


<commit_before>package stun

import (
	"encoding/binary"
	"hash/crc32"

	"github.com/pkg/errors"
)

// https://tools.ietf.org/html/rfc5389#section-15.10
// The SOFTWARE attribute contains a textual description of the software
//  being used by the agent sending the message
type Fingerprint struct {
	Fingerprint uint32
}

const (
	fingerprintXOR    uint32 = 0x5354554e + 1
	fingerprintLength        = 4
)

func calculateFingerprint(b []byte) uint32 {
	return crc32.ChecksumIEEE(b) ^ fingerprintXOR
}

func (s *Fingerprint) Pack(message *Message) error {
	prevLen := message.Length
	message.Length += attrHeaderLength + fingerprintLength
	message.CommitLength()
	v := make([]byte, fingerprintLength)
	binary.BigEndian.PutUint32(v, calculateFingerprint(message.Raw))
	message.Length = prevLen

	message.AddAttribute(AttrFingerprint, v)
	return nil
}

func (s *Fingerprint) Unpack(message *Message, rawAttribute *RawAttribute) error {

	s.Fingerprint = binary.BigEndian.Uint32(rawAttribute.Value)

	expected := calculateFingerprint(message.Raw[:rawAttribute.Offset])

	if expected != s.Fingerprint {
		return errors.Errorf("fingerprint mismatch %v != %v (expected)", s.Fingerprint, expected)
	}
	return nil
}
<commit_msg>turn: Remove accidental +1 to fingerprint magic<commit_after>package stun

import (
	"encoding/binary"
	"hash/crc32"

	"github.com/pkg/errors"
)

// https://tools.ietf.org/html/rfc5389#section-15.10
// The SOFTWARE attribute contains a textual description of the software
//  being used by the agent sending the message
type Fingerprint struct {
	Fingerprint uint32
}

const (
	fingerprintXOR    uint32 = 0x5354554e
	fingerprintLength        = 4
)

func calculateFingerprint(b []byte) uint32 {
	return crc32.ChecksumIEEE(b) ^ fingerprintXOR
}

func (s *Fingerprint) Pack(message *Message) error {
	prevLen := message.Length
	message.Length += attrHeaderLength + fingerprintLength
	message.CommitLength()
	v := make([]byte, fingerprintLength)
	binary.BigEndian.PutUint32(v, calculateFingerprint(message.Raw))
	message.Length = prevLen

	message.AddAttribute(AttrFingerprint, v)
	return nil
}

func (s *Fingerprint) Unpack(message *Message, rawAttribute *RawAttribute) error {

	s.Fingerprint = binary.BigEndian.Uint32(rawAttribute.Value)

	expected := calculateFingerprint(message.Raw[:rawAttribute.Offset])

	if expected != s.Fingerprint {
		return errors.Errorf("fingerprint mismatch %v != %v (expected)", s.Fingerprint, expected)
	}
	return nil
}
<|endoftext|>