<commit_before>// Copyright 2014, Hǎiliàng Wáng. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package reporter

import (
	"fmt"
	"io"
	"strings"

	ge "github.com/hailiang/gspec/error"
	ext "github.com/hailiang/gspec/extension"
)

const gspecPath = "github.com/hailiang/gspec"

// NewTextReporter creates and initialize a new text reporter using w to write
// the output.
func NewTextReporter(w io.Writer, verbose bool) ext.Reporter {
	return &textReporter{w: w, verbose: verbose}
}

// NewTextProgresser creates and initialize a new text progresser using w to
// write the output.
func NewTextProgresser(w io.Writer) ext.Reporter {
	return &textProgresser{w: w}
}

// TextReporter implements a simple plain text CLI reporter.
type textReporter struct {
	dummyReporter
	ext.Stats
	w       io.Writer
	verbose bool
}

func (l *textReporter) End(groups ext.TestGroups) {
	mid := make(map[string]bool)
	for _, g := range groups {
		completed := g.For(func(path ext.TestGroups) bool {
			last := path[len(path)-1]
			if l.verbose || last.Error != nil {
				if !writeTestGroups(l.w, path, mid) {
					return false
				}
			}
			return true
		})
		if !completed {
			break
		}
	}
	if l.Stats.Failed > 0 {
		fmt.Fprintf(l.w, ">>> FAIL COUNT: %d of %d.\n", l.Stats.Failed, l.Stats.Total)
	}
	if l.Stats.Pending > 0 {
		fmt.Fprintf(l.w, ">>> PENDING COUNT: %d of %d.\n", l.Stats.Pending, l.Stats.Total)
	}
	fmt.Fprintf(l.w, ">>> TOTAL: %d.\n", l.Stats.Total)
}

func (l *textReporter) Progress(g *ext.TestGroup, s *ext.Stats) {
	l.Stats = *s
}

type textProgresser struct {
	ext.Stats
	w io.Writer
}

func (p *textProgresser) Start() {
	fmt.Fprint(p.w, "^")
}

func (p *textProgresser) Progress(g *ext.TestGroup, s *ext.Stats) {
	if s.Ended > p.Ended {
		sym := "."
		if g.Error != nil {
			if isPending(g.Error) {
				sym = "p"
			} else {
				sym = "F"
			}
		}
		fmt.Fprint(p.w, sym)
	}
	p.Stats = *s
}

func (p *textProgresser) End(groups ext.TestGroups) {
	fmt.Fprintln(p.w, "$")
}

type dummyReporter struct{}

func (dummyReporter) Start()                              {}
func (dummyReporter) End(ext.TestGroups)                  {}
func (dummyReporter) Progress(*ext.TestGroup, *ext.Stats) {}

// Write writes TestGroups from root to leaf.
func writeTestGroups(w io.Writer, gs ext.TestGroups, mid map[string]bool) bool {
	for i, g := range gs {
		indent := strings.Repeat("    ", i)
		if !mid[g.ID] {
			fmt.Fprintln(w, indent+g.Description)
			mid[g.ID] = true
		}
		if g.Error != nil {
			if panicError, ok := g.Error.(*ext.PanicError); ok {
				writePanicError(w, panicError)
				fmt.Fprintf(w, ge.Indent("(Focus mode: go test -focus %s)", indent), g.ID)
				//				fmt.Fprintf(w, string(panicError.SS))
				fmt.Fprintln(w, ">>> Stop printing more errors due to a panic.")
				return false
			}
			fmt.Fprintln(w, ge.Indent(g.Error.Error(), indent+"  "))
			if !isPending(g.Error) {
				fmt.Fprintf(w, ge.Indent("(Focus mode: go test -focus %s)", indent), g.ID)
			}
		}
	}
	return true
}

func writePanicError(w io.Writer, e *ext.PanicError) {
	fmt.Fprint(w, "panic: ")
	fmt.Fprintln(w, e.Err.Error())
	for _, f := range e.Stack {
		if strings.Contains(f.File, gspecPath) {
			fmt.Fprintln(w, "    ......")
			break
		}
		fmt.Fprint(w, "    ")
		fmt.Fprintln(w, f.Name)
		fmt.Fprint(w, "        ")
		fmt.Fprintf(w, "%s:%d\n", f.File, f.Line)
	}
}

// T is an interface that allows a testing.T to be passed.
type T interface {
	Fail()
}

type failReporter struct {
	t T
	dummyReporter
}

// NewFailReporter creates and initializes a reporter that calls T.Fail when
// any test error occurs.
func NewFailReporter(t T) ext.Reporter {
	return &failReporter{t: t}
}

func (r *failReporter) Progress(g *ext.TestGroup, stats *ext.Stats) {
	if g.Error != nil {
		r.t.Fail()
	}
}

func isPending(err error) bool {
	switch err.(type) {
	case *ext.PendingError:
		return true
	}
	return false
}
<commit_msg>improve focus instructions.<commit_after>// Copyright 2014, Hǎiliàng Wáng. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package reporter

import (
	"fmt"
	"io"
	"strings"

	ge "github.com/hailiang/gspec/error"
	ext "github.com/hailiang/gspec/extension"
)

const gspecPath = "github.com/hailiang/gspec"

// NewTextReporter creates and initialize a new text reporter using w to write
// the output.
func NewTextReporter(w io.Writer, verbose bool) ext.Reporter {
	return &textReporter{w: w, verbose: verbose}
}

// NewTextProgresser creates and initialize a new text progresser using w to
// write the output.
func NewTextProgresser(w io.Writer) ext.Reporter {
	return &textProgresser{w: w}
}

// TextReporter implements a simple plain text CLI reporter.
type textReporter struct {
	dummyReporter
	ext.Stats
	w       io.Writer
	verbose bool
}

func (l *textReporter) End(groups ext.TestGroups) {
	mid := make(map[string]bool)
	for _, g := range groups {
		completed := g.For(func(path ext.TestGroups) bool {
			last := path[len(path)-1]
			if l.verbose || last.Error != nil {
				if !writeTestGroups(l.w, path, mid) {
					return false
				}
			}
			return true
		})
		if !completed {
			break
		}
	}
	if l.Stats.Failed > 0 {
		fmt.Fprintf(l.w, ">>> FAIL COUNT: %d of %d.\n", l.Stats.Failed, l.Stats.Total)
	}
	if l.Stats.Pending > 0 {
		fmt.Fprintf(l.w, ">>> PENDING COUNT: %d of %d.\n", l.Stats.Pending, l.Stats.Total)
	}
	fmt.Fprintf(l.w, ">>> TOTAL: %d.\n", l.Stats.Total)
}

func (l *textReporter) Progress(g *ext.TestGroup, s *ext.Stats) {
	l.Stats = *s
}

type textProgresser struct {
	ext.Stats
	w io.Writer
}

func (p *textProgresser) Start() {
	fmt.Fprint(p.w, "^")
}

func (p *textProgresser) Progress(g *ext.TestGroup, s *ext.Stats) {
	if s.Ended > p.Ended {
		sym := "."
		if g.Error != nil {
			if isPending(g.Error) {
				sym = "p"
			} else {
				sym = "F"
			}
		}
		fmt.Fprint(p.w, sym)
	}
	p.Stats = *s
}

func (p *textProgresser) End(groups ext.TestGroups) {
	fmt.Fprintln(p.w, "$")
}

type dummyReporter struct{}

func (dummyReporter) Start()                              {}
func (dummyReporter) End(ext.TestGroups)                  {}
func (dummyReporter) Progress(*ext.TestGroup, *ext.Stats) {}

// Write writes TestGroups from root to leaf.
func writeTestGroups(w io.Writer, gs ext.TestGroups, mid map[string]bool) bool {
	for i, g := range gs {
		indent := strings.Repeat("    ", i)
		if !mid[g.ID] {
			fmt.Fprintln(w, indent+g.Description)
			mid[g.ID] = true
		}
		if g.Error != nil {
			if panicError, ok := g.Error.(*ext.PanicError); ok {
				writePanicError(w, panicError)
				printFocusInstruction(w, indent, g.ID)
				//				fmt.Fprintf(w, string(panicError.SS))
				fmt.Fprintln(w, ">>> Stop printing more errors due to a panic.")
				return false
			}
			fmt.Fprintln(w, ge.Indent(g.Error.Error(), indent+"  "))
			if !isPending(g.Error) {
				printFocusInstruction(w, indent, g.ID)
			}
		}
	}
	return true
}

func printFocusInstruction(w io.Writer, indent, id string) {
	fmt.Fprintf(w, ge.Indent(`  (use "go test -focus %s" to run the test case only.)`, indent), id)
}

func writePanicError(w io.Writer, e *ext.PanicError) {
	fmt.Fprint(w, "panic: ")
	fmt.Fprintln(w, e.Err.Error())
	for _, f := range e.Stack {
		if strings.Contains(f.File, gspecPath) {
			fmt.Fprintln(w, "    ......")
			break
		}
		fmt.Fprint(w, "    ")
		fmt.Fprintln(w, f.Name)
		fmt.Fprint(w, "        ")
		fmt.Fprintf(w, "%s:%d\n", f.File, f.Line)
	}
}

// T is an interface that allows a testing.T to be passed.
type T interface {
	Fail()
}

type failReporter struct {
	t T
	dummyReporter
}

// NewFailReporter creates and initializes a reporter that calls T.Fail when
// any test error occurs.
func NewFailReporter(t T) ext.Reporter {
	return &failReporter{t: t}
}

func (r *failReporter) Progress(g *ext.TestGroup, stats *ext.Stats) {
	if g.Error != nil {
		r.t.Fail()
	}
}

func isPending(err error) bool {
	switch err.(type) {
	case *ext.PendingError:
		return true
	}
	return false
}
<|endoftext|>


<commit_before>package cli

import (
	"fmt"
	"os"
	"strings"

	pluginmanager "github.com/docker/cli/cli-plugins/manager"
	"github.com/docker/cli/cli/command"
	cliconfig "github.com/docker/cli/cli/config"
	cliflags "github.com/docker/cli/cli/flags"
	"github.com/docker/docker/pkg/term"
	"github.com/pkg/errors"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
)

// setupCommonRootCommand contains the setup common to
// SetupRootCommand and SetupPluginRootCommand.
func setupCommonRootCommand(rootCmd *cobra.Command) (*cliflags.ClientOptions, *pflag.FlagSet, *cobra.Command) {
	opts := cliflags.NewClientOptions()
	flags := rootCmd.Flags()

	flags.StringVar(&opts.ConfigDir, "config", cliconfig.Dir(), "Location of client config files")
	opts.Common.InstallFlags(flags)

	cobra.AddTemplateFunc("add", func(a, b int) int { return a + b })
	cobra.AddTemplateFunc("hasSubCommands", hasSubCommands)
	cobra.AddTemplateFunc("hasManagementSubCommands", hasManagementSubCommands)
	cobra.AddTemplateFunc("hasInvalidPlugins", hasInvalidPlugins)
	cobra.AddTemplateFunc("operationSubCommands", operationSubCommands)
	cobra.AddTemplateFunc("managementSubCommands", managementSubCommands)
	cobra.AddTemplateFunc("invalidPlugins", invalidPlugins)
	cobra.AddTemplateFunc("wrappedFlagUsages", wrappedFlagUsages)
	cobra.AddTemplateFunc("vendorAndVersion", vendorAndVersion)
	cobra.AddTemplateFunc("invalidPluginReason", invalidPluginReason)
	cobra.AddTemplateFunc("isPlugin", isPlugin)
	cobra.AddTemplateFunc("decoratedName", decoratedName)

	rootCmd.SetUsageTemplate(usageTemplate)
	rootCmd.SetHelpTemplate(helpTemplate)
	rootCmd.SetFlagErrorFunc(FlagErrorFunc)
	rootCmd.SetHelpCommand(helpCommand)

	return opts, flags, helpCommand
}

// SetupRootCommand sets default usage, help, and error handling for the
// root command.
func SetupRootCommand(rootCmd *cobra.Command) (*cliflags.ClientOptions, *pflag.FlagSet, *cobra.Command) {
	opts, flags, helpCmd := setupCommonRootCommand(rootCmd)

	rootCmd.SetVersionTemplate("Docker version {{.Version}}\n")

	rootCmd.PersistentFlags().BoolP("help", "h", false, "Print usage")
	rootCmd.PersistentFlags().MarkShorthandDeprecated("help", "please use --help")
	rootCmd.PersistentFlags().Lookup("help").Hidden = true

	return opts, flags, helpCmd
}

// SetupPluginRootCommand sets default usage, help and error handling for a plugin root command.
func SetupPluginRootCommand(rootCmd *cobra.Command) (*cliflags.ClientOptions, *pflag.FlagSet) {
	opts, flags, _ := setupCommonRootCommand(rootCmd)

	rootCmd.PersistentFlags().BoolP("help", "", false, "Print usage")
	rootCmd.PersistentFlags().Lookup("help").Hidden = true

	return opts, flags
}

// FlagErrorFunc prints an error message which matches the format of the
// docker/cli/cli error messages
func FlagErrorFunc(cmd *cobra.Command, err error) error {
	if err == nil {
		return nil
	}

	usage := ""
	if cmd.HasSubCommands() {
		usage = "\n\n" + cmd.UsageString()
	}
	return StatusError{
		Status:     fmt.Sprintf("%s\nSee '%s --help'.%s", err, cmd.CommandPath(), usage),
		StatusCode: 125,
	}
}

// TopLevelCommand encapsulates a top-level cobra command (either
// docker CLI or a plugin) and global flag handling logic necessary
// for plugins.
type TopLevelCommand struct {
	cmd       *cobra.Command
	dockerCli *command.DockerCli
	opts      *cliflags.ClientOptions
	flags     *pflag.FlagSet
	args      []string
}

// NewTopLevelCommand returns a new TopLevelCommand object
func NewTopLevelCommand(cmd *cobra.Command, dockerCli *command.DockerCli, opts *cliflags.ClientOptions, flags *pflag.FlagSet) *TopLevelCommand {
	return &TopLevelCommand{cmd, dockerCli, opts, flags, os.Args[1:]}
}

// SetArgs sets the args (default os.Args[:1] used to invoke the command
func (tcmd *TopLevelCommand) SetArgs(args []string) {
	tcmd.args = args
	tcmd.cmd.SetArgs(args)
}

// SetFlag sets a flag in the local flag set of the top-level command
func (tcmd *TopLevelCommand) SetFlag(name, value string) {
	tcmd.cmd.Flags().Set(name, value)
}

// HandleGlobalFlags takes care of parsing global flags defined on the
// command, it returns the underlying cobra command and the args it
// will be called with (or an error).
//
// On success the caller is responsible for calling Initialize()
// before calling `Execute` on the returned command.
func (tcmd *TopLevelCommand) HandleGlobalFlags() (*cobra.Command, []string, error) {
	cmd := tcmd.cmd

	// We manually parse the global arguments and find the
	// subcommand in order to properly deal with plugins. We rely
	// on the root command never having any non-flag arguments.
	flags := cmd.Flags()

	// We need !interspersed to ensure we stop at the first
	// potential command instead of accumulating it into
	// flags.Args() and then continuing on and finding other
	// arguments which we try and treat as globals (when they are
	// actually arguments to the subcommand).
	flags.SetInterspersed(false)
	defer flags.SetInterspersed(true) // Undo, any subsequent cmd.Execute() in the caller expects this.

	// We need the single parse to see both sets of flags.
	flags.AddFlagSet(cmd.PersistentFlags())
	// Now parse the global flags, up to (but not including) the
	// first command. The result will be that all the remaining
	// arguments are in `flags.Args()`.
	if err := flags.Parse(tcmd.args); err != nil {
		// Our FlagErrorFunc uses the cli, make sure it is initialized
		if err := tcmd.Initialize(); err != nil {
			return nil, nil, err
		}
		return nil, nil, cmd.FlagErrorFunc()(cmd, err)
	}

	return cmd, flags.Args(), nil
}

// Initialize finalises global option parsing and initializes the docker client.
func (tcmd *TopLevelCommand) Initialize(ops ...command.InitializeOpt) error {
	tcmd.opts.Common.SetDefaultOptions(tcmd.flags)
	return tcmd.dockerCli.Initialize(tcmd.opts, ops...)
}

// VisitAll will traverse all commands from the root.
// This is different from the VisitAll of cobra.Command where only parents
// are checked.
func VisitAll(root *cobra.Command, fn func(*cobra.Command)) {
	for _, cmd := range root.Commands() {
		VisitAll(cmd, fn)
	}
	fn(root)
}

// DisableFlagsInUseLine sets the DisableFlagsInUseLine flag on all
// commands within the tree rooted at cmd.
func DisableFlagsInUseLine(cmd *cobra.Command) {
	VisitAll(cmd, func(ccmd *cobra.Command) {
		// do not add a `[flags]` to the end of the usage line.
		ccmd.DisableFlagsInUseLine = true
	})
}

var helpCommand = &cobra.Command{
	Use:               "help [command]",
	Short:             "Help about the command",
	PersistentPreRun:  func(cmd *cobra.Command, args []string) {},
	PersistentPostRun: func(cmd *cobra.Command, args []string) {},
	RunE: func(c *cobra.Command, args []string) error {
		cmd, args, e := c.Root().Find(args)
		if cmd == nil || e != nil || len(args) > 0 {
			return errors.Errorf("unknown help topic: %v", strings.Join(args, " "))
		}

		helpFunc := cmd.HelpFunc()
		helpFunc(cmd, args)
		return nil
	},
}

func isPlugin(cmd *cobra.Command) bool {
	return cmd.Annotations[pluginmanager.CommandAnnotationPlugin] == "true"
}

func hasSubCommands(cmd *cobra.Command) bool {
	return len(operationSubCommands(cmd)) > 0
}

func hasManagementSubCommands(cmd *cobra.Command) bool {
	return len(managementSubCommands(cmd)) > 0
}

func hasInvalidPlugins(cmd *cobra.Command) bool {
	return len(invalidPlugins(cmd)) > 0
}

func operationSubCommands(cmd *cobra.Command) []*cobra.Command {
	cmds := []*cobra.Command{}
	for _, sub := range cmd.Commands() {
		if isPlugin(sub) {
			continue
		}
		if sub.IsAvailableCommand() && !sub.HasSubCommands() {
			cmds = append(cmds, sub)
		}
	}
	return cmds
}

func wrappedFlagUsages(cmd *cobra.Command) string {
	width := 80
	if ws, err := term.GetWinsize(0); err == nil {
		width = int(ws.Width)
	}
	return cmd.Flags().FlagUsagesWrapped(width - 1)
}

func decoratedName(cmd *cobra.Command) string {
	decoration := " "
	if isPlugin(cmd) {
		decoration = "*"
	}
	return cmd.Name() + decoration
}

func vendorAndVersion(cmd *cobra.Command) string {
	if vendor, ok := cmd.Annotations[pluginmanager.CommandAnnotationPluginVendor]; ok && isPlugin(cmd) {
		version := ""
		if v, ok := cmd.Annotations[pluginmanager.CommandAnnotationPluginVersion]; ok && v != "" {
			version = ", " + v
		}
		return fmt.Sprintf("(%s%s)", vendor, version)
	}
	return ""
}

func managementSubCommands(cmd *cobra.Command) []*cobra.Command {
	cmds := []*cobra.Command{}
	for _, sub := range cmd.Commands() {
		if isPlugin(sub) {
			if invalidPluginReason(sub) == "" {
				cmds = append(cmds, sub)
			}
			continue
		}
		if sub.IsAvailableCommand() && sub.HasSubCommands() {
			cmds = append(cmds, sub)
		}
	}
	return cmds
}

func invalidPlugins(cmd *cobra.Command) []*cobra.Command {
	cmds := []*cobra.Command{}
	for _, sub := range cmd.Commands() {
		if !isPlugin(sub) {
			continue
		}
		if invalidPluginReason(sub) != "" {
			cmds = append(cmds, sub)
		}
	}
	return cmds
}

func invalidPluginReason(cmd *cobra.Command) string {
	return cmd.Annotations[pluginmanager.CommandAnnotationPluginInvalid]
}

var usageTemplate = `Usage:

{{- if not .HasSubCommands}}	{{.UseLine}}{{end}}
{{- if .HasSubCommands}}	{{ .CommandPath}}{{- if .HasAvailableFlags}} [OPTIONS]{{end}} COMMAND{{end}}

{{if ne .Long ""}}{{ .Long | trim }}{{ else }}{{ .Short | trim }}{{end}}

{{- if gt .Aliases 0}}

Aliases:
  {{.NameAndAliases}}

{{- end}}
{{- if .HasExample}}

Examples:
{{ .Example }}

{{- end}}
{{- if .HasAvailableFlags}}

Options:
{{ wrappedFlagUsages . | trimRightSpace}}

{{- end}}
{{- if hasManagementSubCommands . }}

Management Commands:

{{- range managementSubCommands . }}
  {{rpad (decoratedName .) (add .NamePadding 1)}}{{.Short}}{{ if isPlugin .}} {{vendorAndVersion .}}{{ end}}
{{- end}}

{{- end}}
{{- if hasSubCommands .}}

Commands:

{{- range operationSubCommands . }}
  {{rpad .Name .NamePadding }} {{.Short}}
{{- end}}
{{- end}}

{{- if hasInvalidPlugins . }}

Invalid Plugins:

{{- range invalidPlugins . }}
  {{rpad .Name .NamePadding }} {{invalidPluginReason .}}
{{- end}}

{{- end}}

{{- if .HasSubCommands }}

Run '{{.CommandPath}} COMMAND --help' for more information on a command.
{{- end}}
`

var helpTemplate = `
{{if or .Runnable .HasSubCommands}}{{.UsageString}}{{end}}`
<commit_msg>Use a copy of root flagset in `HandleGlobalFlags`<commit_after>package cli

import (
	"fmt"
	"os"
	"strings"

	pluginmanager "github.com/docker/cli/cli-plugins/manager"
	"github.com/docker/cli/cli/command"
	cliconfig "github.com/docker/cli/cli/config"
	cliflags "github.com/docker/cli/cli/flags"
	"github.com/docker/docker/pkg/term"
	"github.com/pkg/errors"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
)

// setupCommonRootCommand contains the setup common to
// SetupRootCommand and SetupPluginRootCommand.
func setupCommonRootCommand(rootCmd *cobra.Command) (*cliflags.ClientOptions, *pflag.FlagSet, *cobra.Command) {
	opts := cliflags.NewClientOptions()
	flags := rootCmd.Flags()

	flags.StringVar(&opts.ConfigDir, "config", cliconfig.Dir(), "Location of client config files")
	opts.Common.InstallFlags(flags)

	cobra.AddTemplateFunc("add", func(a, b int) int { return a + b })
	cobra.AddTemplateFunc("hasSubCommands", hasSubCommands)
	cobra.AddTemplateFunc("hasManagementSubCommands", hasManagementSubCommands)
	cobra.AddTemplateFunc("hasInvalidPlugins", hasInvalidPlugins)
	cobra.AddTemplateFunc("operationSubCommands", operationSubCommands)
	cobra.AddTemplateFunc("managementSubCommands", managementSubCommands)
	cobra.AddTemplateFunc("invalidPlugins", invalidPlugins)
	cobra.AddTemplateFunc("wrappedFlagUsages", wrappedFlagUsages)
	cobra.AddTemplateFunc("vendorAndVersion", vendorAndVersion)
	cobra.AddTemplateFunc("invalidPluginReason", invalidPluginReason)
	cobra.AddTemplateFunc("isPlugin", isPlugin)
	cobra.AddTemplateFunc("decoratedName", decoratedName)

	rootCmd.SetUsageTemplate(usageTemplate)
	rootCmd.SetHelpTemplate(helpTemplate)
	rootCmd.SetFlagErrorFunc(FlagErrorFunc)
	rootCmd.SetHelpCommand(helpCommand)

	return opts, flags, helpCommand
}

// SetupRootCommand sets default usage, help, and error handling for the
// root command.
func SetupRootCommand(rootCmd *cobra.Command) (*cliflags.ClientOptions, *pflag.FlagSet, *cobra.Command) {
	opts, flags, helpCmd := setupCommonRootCommand(rootCmd)

	rootCmd.SetVersionTemplate("Docker version {{.Version}}\n")

	rootCmd.PersistentFlags().BoolP("help", "h", false, "Print usage")
	rootCmd.PersistentFlags().MarkShorthandDeprecated("help", "please use --help")
	rootCmd.PersistentFlags().Lookup("help").Hidden = true

	return opts, flags, helpCmd
}

// SetupPluginRootCommand sets default usage, help and error handling for a plugin root command.
func SetupPluginRootCommand(rootCmd *cobra.Command) (*cliflags.ClientOptions, *pflag.FlagSet) {
	opts, flags, _ := setupCommonRootCommand(rootCmd)

	rootCmd.PersistentFlags().BoolP("help", "", false, "Print usage")
	rootCmd.PersistentFlags().Lookup("help").Hidden = true

	return opts, flags
}

// FlagErrorFunc prints an error message which matches the format of the
// docker/cli/cli error messages
func FlagErrorFunc(cmd *cobra.Command, err error) error {
	if err == nil {
		return nil
	}

	usage := ""
	if cmd.HasSubCommands() {
		usage = "\n\n" + cmd.UsageString()
	}
	return StatusError{
		Status:     fmt.Sprintf("%s\nSee '%s --help'.%s", err, cmd.CommandPath(), usage),
		StatusCode: 125,
	}
}

// TopLevelCommand encapsulates a top-level cobra command (either
// docker CLI or a plugin) and global flag handling logic necessary
// for plugins.
type TopLevelCommand struct {
	cmd       *cobra.Command
	dockerCli *command.DockerCli
	opts      *cliflags.ClientOptions
	flags     *pflag.FlagSet
	args      []string
}

// NewTopLevelCommand returns a new TopLevelCommand object
func NewTopLevelCommand(cmd *cobra.Command, dockerCli *command.DockerCli, opts *cliflags.ClientOptions, flags *pflag.FlagSet) *TopLevelCommand {
	return &TopLevelCommand{cmd, dockerCli, opts, flags, os.Args[1:]}
}

// SetArgs sets the args (default os.Args[:1] used to invoke the command
func (tcmd *TopLevelCommand) SetArgs(args []string) {
	tcmd.args = args
	tcmd.cmd.SetArgs(args)
}

// SetFlag sets a flag in the local flag set of the top-level command
func (tcmd *TopLevelCommand) SetFlag(name, value string) {
	tcmd.cmd.Flags().Set(name, value)
}

// HandleGlobalFlags takes care of parsing global flags defined on the
// command, it returns the underlying cobra command and the args it
// will be called with (or an error).
//
// On success the caller is responsible for calling Initialize()
// before calling `Execute` on the returned command.
func (tcmd *TopLevelCommand) HandleGlobalFlags() (*cobra.Command, []string, error) {
	cmd := tcmd.cmd

	// We manually parse the global arguments and find the
	// subcommand in order to properly deal with plugins. We rely
	// on the root command never having any non-flag arguments. We
	// create our own FlagSet so that we can configure it
	// (e.g. `SetInterspersed` below) in an idempotent way.
	flags := pflag.NewFlagSet(cmd.Name(), pflag.ContinueOnError)

	// We need !interspersed to ensure we stop at the first
	// potential command instead of accumulating it into
	// flags.Args() and then continuing on and finding other
	// arguments which we try and treat as globals (when they are
	// actually arguments to the subcommand).
	flags.SetInterspersed(false)

	// We need the single parse to see both sets of flags.
	flags.AddFlagSet(cmd.Flags())
	flags.AddFlagSet(cmd.PersistentFlags())
	// Now parse the global flags, up to (but not including) the
	// first command. The result will be that all the remaining
	// arguments are in `flags.Args()`.
	if err := flags.Parse(tcmd.args); err != nil {
		// Our FlagErrorFunc uses the cli, make sure it is initialized
		if err := tcmd.Initialize(); err != nil {
			return nil, nil, err
		}
		return nil, nil, cmd.FlagErrorFunc()(cmd, err)
	}

	return cmd, flags.Args(), nil
}

// Initialize finalises global option parsing and initializes the docker client.
func (tcmd *TopLevelCommand) Initialize(ops ...command.InitializeOpt) error {
	tcmd.opts.Common.SetDefaultOptions(tcmd.flags)
	return tcmd.dockerCli.Initialize(tcmd.opts, ops...)
}

// VisitAll will traverse all commands from the root.
// This is different from the VisitAll of cobra.Command where only parents
// are checked.
func VisitAll(root *cobra.Command, fn func(*cobra.Command)) {
	for _, cmd := range root.Commands() {
		VisitAll(cmd, fn)
	}
	fn(root)
}

// DisableFlagsInUseLine sets the DisableFlagsInUseLine flag on all
// commands within the tree rooted at cmd.
func DisableFlagsInUseLine(cmd *cobra.Command) {
	VisitAll(cmd, func(ccmd *cobra.Command) {
		// do not add a `[flags]` to the end of the usage line.
		ccmd.DisableFlagsInUseLine = true
	})
}

var helpCommand = &cobra.Command{
	Use:               "help [command]",
	Short:             "Help about the command",
	PersistentPreRun:  func(cmd *cobra.Command, args []string) {},
	PersistentPostRun: func(cmd *cobra.Command, args []string) {},
	RunE: func(c *cobra.Command, args []string) error {
		cmd, args, e := c.Root().Find(args)
		if cmd == nil || e != nil || len(args) > 0 {
			return errors.Errorf("unknown help topic: %v", strings.Join(args, " "))
		}

		helpFunc := cmd.HelpFunc()
		helpFunc(cmd, args)
		return nil
	},
}

func isPlugin(cmd *cobra.Command) bool {
	return cmd.Annotations[pluginmanager.CommandAnnotationPlugin] == "true"
}

func hasSubCommands(cmd *cobra.Command) bool {
	return len(operationSubCommands(cmd)) > 0
}

func hasManagementSubCommands(cmd *cobra.Command) bool {
	return len(managementSubCommands(cmd)) > 0
}

func hasInvalidPlugins(cmd *cobra.Command) bool {
	return len(invalidPlugins(cmd)) > 0
}

func operationSubCommands(cmd *cobra.Command) []*cobra.Command {
	cmds := []*cobra.Command{}
	for _, sub := range cmd.Commands() {
		if isPlugin(sub) {
			continue
		}
		if sub.IsAvailableCommand() && !sub.HasSubCommands() {
			cmds = append(cmds, sub)
		}
	}
	return cmds
}

func wrappedFlagUsages(cmd *cobra.Command) string {
	width := 80
	if ws, err := term.GetWinsize(0); err == nil {
		width = int(ws.Width)
	}
	return cmd.Flags().FlagUsagesWrapped(width - 1)
}

func decoratedName(cmd *cobra.Command) string {
	decoration := " "
	if isPlugin(cmd) {
		decoration = "*"
	}
	return cmd.Name() + decoration
}

func vendorAndVersion(cmd *cobra.Command) string {
	if vendor, ok := cmd.Annotations[pluginmanager.CommandAnnotationPluginVendor]; ok && isPlugin(cmd) {
		version := ""
		if v, ok := cmd.Annotations[pluginmanager.CommandAnnotationPluginVersion]; ok && v != "" {
			version = ", " + v
		}
		return fmt.Sprintf("(%s%s)", vendor, version)
	}
	return ""
}

func managementSubCommands(cmd *cobra.Command) []*cobra.Command {
	cmds := []*cobra.Command{}
	for _, sub := range cmd.Commands() {
		if isPlugin(sub) {
			if invalidPluginReason(sub) == "" {
				cmds = append(cmds, sub)
			}
			continue
		}
		if sub.IsAvailableCommand() && sub.HasSubCommands() {
			cmds = append(cmds, sub)
		}
	}
	return cmds
}

func invalidPlugins(cmd *cobra.Command) []*cobra.Command {
	cmds := []*cobra.Command{}
	for _, sub := range cmd.Commands() {
		if !isPlugin(sub) {
			continue
		}
		if invalidPluginReason(sub) != "" {
			cmds = append(cmds, sub)
		}
	}
	return cmds
}

func invalidPluginReason(cmd *cobra.Command) string {
	return cmd.Annotations[pluginmanager.CommandAnnotationPluginInvalid]
}

var usageTemplate = `Usage:

{{- if not .HasSubCommands}}	{{.UseLine}}{{end}}
{{- if .HasSubCommands}}	{{ .CommandPath}}{{- if .HasAvailableFlags}} [OPTIONS]{{end}} COMMAND{{end}}

{{if ne .Long ""}}{{ .Long | trim }}{{ else }}{{ .Short | trim }}{{end}}

{{- if gt .Aliases 0}}

Aliases:
  {{.NameAndAliases}}

{{- end}}
{{- if .HasExample}}

Examples:
{{ .Example }}

{{- end}}
{{- if .HasAvailableFlags}}

Options:
{{ wrappedFlagUsages . | trimRightSpace}}

{{- end}}
{{- if hasManagementSubCommands . }}

Management Commands:

{{- range managementSubCommands . }}
  {{rpad (decoratedName .) (add .NamePadding 1)}}{{.Short}}{{ if isPlugin .}} {{vendorAndVersion .}}{{ end}}
{{- end}}

{{- end}}
{{- if hasSubCommands .}}

Commands:

{{- range operationSubCommands . }}
  {{rpad .Name .NamePadding }} {{.Short}}
{{- end}}
{{- end}}

{{- if hasInvalidPlugins . }}

Invalid Plugins:

{{- range invalidPlugins . }}
  {{rpad .Name .NamePadding }} {{invalidPluginReason .}}
{{- end}}

{{- end}}

{{- if .HasSubCommands }}

Run '{{.CommandPath}} COMMAND --help' for more information on a command.
{{- end}}
`

var helpTemplate = `
{{if or .Runnable .HasSubCommands}}{{.UsageString}}{{end}}`
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/flynn/flynn/controller/client"
	ct "github.com/flynn/flynn/controller/types"
	"github.com/flynn/flynn/controller/utils"
	"github.com/flynn/flynn/pkg/version"
	"github.com/flynn/go-docopt"
)

func init() {
	register("scale", runScale, `
usage: flynn scale [options] [<type>=<spec>...]

Scale changes the number of jobs and tags for each process type in a release.

Process type scale should be formatted like TYPE=COUNT[,KEY=VAL...], for example:

web=1                  # 1 web process
web=3                  # 3 web processes, distributed amongst all hosts
web=3,active=true      # 3 web processes, distributed amongst hosts tagged active=true
db=3,disk=ssd,mem=high # 3 db processes, distributed amongst hosts tagged with
                       # both disk=ssd and mem=high

Ommitting the arguments will show the current scale.

Options:
	-n, --no-wait            don't wait for the scaling events to happen
	-r, --release=<release>  id of release to scale (defaults to current app release)
	-a, --all                show non-zero formations from all releases (only works when listing formations, can't be combined with --release)

Example:

	$ flynn scale
	web=4 worker=2

	$ flynn scale --all
	496d6e74-9db9-4cff-bcce-a3b44015907a (current)
	web=1 worker=2

	632cd907-85ab-4e53-90d0-84635650ec9a
	web=2

	$ flynn scale web=2 worker=5
	scaling web: 4=>2, worker: 2=>5

	02:28:34.333 ==> web flynn-3f656af6f1e44092aa7037046236b203 down
	02:28:34.466 ==> web flynn-ee83def0b8e4455793a43c8c70f5b34e down
	02:28:35.479 ==> worker flynn-84f70ca18c9641ef83a178a19db867a3 up
	02:28:36.508 ==> worker flynn-a3de8c326cc542aa89235e53ba304260 up
	02:28:37.601 ==> worker flynn-e24760c511af4733b01ed5b98aa54647 up

	scale completed in 3.944629056s
`)
}

// minScaleRequestVersion is the minimum API version which supports scaling
// using scale requests
const minScaleRequestVersion = "v20170309.0"

// takes args of the form "web=1[,key=val...]", "worker=3[,key=val...]", etc
func runScale(args *docopt.Args, client controller.Client) error {
	app := mustApp()

	typeSpecs := args.All["<type>=<spec>"].([]string)

	showAll := args.Bool["--all"]

	if len(typeSpecs) > 0 && showAll {
		return fmt.Errorf("ERROR: Can't use --all when scaling")
	}

	releaseID := args.String["--release"]
	if releaseID != "" && showAll {
		return fmt.Errorf("ERROR: Can't use --all in combination with --release")
	}

	if len(typeSpecs) == 0 {
		return showFormations(client, releaseID, showAll, app)
	}

	release, err := determineRelease(client, releaseID, app)
	if err != nil {
		return err
	}

	processes := make(map[string]int, len(typeSpecs))
	tags := make(map[string]map[string]string, len(typeSpecs))
	invalid := make([]string, 0, len(release.Processes))
	for _, arg := range typeSpecs {
		i := strings.IndexRune(arg, '=')
		if i < 0 {
			return fmt.Errorf("ERROR: scale args must be of the form <typ>=<spec>")
		}

		countTags := strings.Split(arg[i+1:], ",")

		count, err := strconv.Atoi(countTags[0])
		if err != nil {
			return fmt.Errorf("ERROR: could not parse quantity in %q", arg)
		} else if count < 0 {
			return fmt.Errorf("ERROR: process quantities cannot be negative in %q", arg)
		}

		processType := arg[:i]
		if _, ok := release.Processes[processType]; ok {
			processes[processType] = count
		} else {
			invalid = append(invalid, fmt.Sprintf("%q", processType))
			continue
		}

		if len(countTags) > 1 {
			processTags := make(map[string]string, len(countTags)-1)
			for i := 1; i < len(countTags); i++ {
				keyVal := strings.SplitN(countTags[i], "=", 2)
				if len(keyVal) == 1 && keyVal[0] != "" {
					processTags[keyVal[0]] = "true"
				} else if len(keyVal) == 2 {
					processTags[keyVal[0]] = keyVal[1]
				}
			}
			tags[processType] = processTags
		}
	}
	if len(invalid) > 0 {
		return fmt.Errorf("ERROR: unknown process types: %s", strings.Join(invalid, ", "))
	}

	opts := ct.ScaleOptions{
		Processes: processes,
		Tags:      tags,
		NoWait:    args.Bool["--no-wait"],
	}

	status, err := client.Status()
	if err != nil {
		return err
	}
	v := version.Parse(status.Version)
	if !v.Dev && v.Before(version.Parse(minScaleRequestVersion)) {
		return runScaleWithJobEvents(client, app, release, opts)
	}
	return runScaleWithScaleRequest(client, app, release, opts)
}

func runScaleWithScaleRequest(client controller.Client, app string, release *ct.Release, opts ct.ScaleOptions) error {
	opts.ScaleRequestCallback = func(req *ct.ScaleRequest) {
		if req.NewProcesses == nil {
			return
		}
		scale := make([]string, 0, len(release.Processes))
		for typ := range release.Processes {
			if count := (*req.NewProcesses)[typ]; count != req.OldProcesses[typ] {
				scale = append(scale, fmt.Sprintf("%s: %d=>%d", typ, req.OldProcesses[typ], count))
			}
		}
		fmt.Printf("scaling %s\n\n", strings.Join(scale, ", "))
	}
	opts.JobEventCallback = func(job *ct.Job) error {
		id := job.ID
		if id == "" {
			id = job.UUID
		}
		fmt.Printf("%s ==> %s %s %s\n", time.Now().Format("15:04:05.000"), job.Type, id, job.State)
		return nil
	}

	start := time.Now()
	if err := client.ScaleAppRelease(app, release.ID, opts); err != nil {
		return err
	}
	if !opts.NoWait {
		fmt.Printf("\nscale completed in %s\n", time.Since(start))
	}
	return nil
}

func runScaleWithJobEvents(client controller.Client, app string, release *ct.Release, opts ct.ScaleOptions) error {
	processes := opts.Processes
	tags := opts.Tags
	formation, err := client.GetFormation(app, release.ID)
	if err == controller.ErrNotFound {
		formation = &ct.Formation{
			AppID:     app,
			ReleaseID: release.ID,
			Processes: make(map[string]int),
		}
	} else if err != nil {
		return err
	}
	if formation.Processes == nil {
		formation.Processes = make(map[string]int, len(processes))
	}
	if formation.Tags == nil {
		formation.Tags = make(map[string]map[string]string, len(tags))
	}

	currentProcs := formation.Processes
	currentTags := formation.Tags
	formation.Processes = processes
	formation.Tags = tags

	if scalingComplete(currentProcs, processes) {
		if !utils.FormationTagsEqual(currentTags, tags) {
			fmt.Println("persisting tag change")
			return client.PutFormation(formation)
		}
		fmt.Println("requested scale equals current scale, nothing to do!")
		return nil
	}

	scale := make([]string, 0, len(release.Processes))
	for typ := range release.Processes {
		if currentProcs[typ] != processes[typ] {
			scale = append(scale, fmt.Sprintf("%s: %d=>%d", typ, currentProcs[typ], processes[typ]))
		}
	}
	fmt.Printf("scaling %s\n\n", strings.Join(scale, ", "))

	expected := client.ExpectedScalingEvents(currentProcs, processes, release.Processes, 1)
	watcher, err := client.WatchJobEvents(app, release.ID)
	if err != nil {
		return err
	}
	defer watcher.Close()

	err = client.PutFormation(formation)
	if err != nil || opts.NoWait {
		return err
	}

	start := time.Now()
	err = watcher.WaitFor(expected, ct.DefaultScaleTimeout, func(job *ct.Job) error {
		id := job.ID
		if id == "" {
			id = job.UUID
		}
		fmt.Printf("%s ==> %s %s %s\n", time.Now().Format("15:04:05.000"), job.Type, id, job.State)
		return nil
	})

	if err != nil {
		return err
	}
	fmt.Printf("\nscale completed in %s\n", time.Since(start))
	return nil
}

func showFormations(client controller.Client, releaseID string, showAll bool, app string) error {
	release, err := determineRelease(client, releaseID, app)
	if err != nil {
		return err
	}
	var releases []*ct.Release
	if showAll {
		var err error
		releases, err = client.AppReleaseList(app)
		if err != nil {
			return err
		}
	} else {
		releases = []*ct.Release{release}
	}

	formations := make(map[string]*ct.Formation, len(releases))
	for _, r := range releases {
		formation, err := client.GetFormation(app, r.ID)
		if err != nil && err != controller.ErrNotFound {
			return err
		}
		formations[r.ID] = formation
	}

	for i, r := range releases {
		f := formations[r.ID]
		if f == nil || len(f.Processes) == 0 {
			continue
		}
		if showAll {
			if i > 0 {
				fmt.Println()
			}
			var suffix string
			if r.ID == release.ID {
				suffix = " (current)"
			}
			fmt.Printf("%s%s\n", r.ID, suffix)
		}
		scale := make([]string, 0, len(r.Processes))
		for typ := range r.Processes {
			n := f.Processes[typ]
			if showAll && n == 0 {
				continue
			}
			scale = append(scale, fmt.Sprintf("%s=%d", typ, n))
		}
		fmt.Println(strings.Join(scale, " "))
	}
	return nil
}

func determineRelease(client controller.Client, releaseID, app string) (*ct.Release, error) {
	if releaseID == "" {
		release, err := client.GetAppRelease(app)
		if err == controller.ErrNotFound {
			return nil, errors.New("No app release, specify a release with --release")
		}
		if err != nil {
			return nil, err
		}
		return release, nil
	}
	return client.GetRelease(releaseID)
}

func scalingComplete(actual, expected map[string]int) bool {
	// check all the expected counts are the same in actual
	for typ, count := range expected {
		if actual[typ] != count {
			return false
		}
	}
	// check any counts in actual which aren't in expected are zero
	for typ, count := range actual {
		if _, ok := expected[typ]; !ok && count != 0 {
			return false
		}
	}
	return true
}
<commit_msg>cli: Fix pre-ScaleRequest formation scaling<commit_after>package main

import (
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	"github.com/flynn/flynn/controller/client"
	ct "github.com/flynn/flynn/controller/types"
	"github.com/flynn/flynn/controller/utils"
	"github.com/flynn/flynn/pkg/version"
	"github.com/flynn/go-docopt"
)

func init() {
	register("scale", runScale, `
usage: flynn scale [options] [<type>=<spec>...]

Scale changes the number of jobs and tags for each process type in a release.

Process type scale should be formatted like TYPE=COUNT[,KEY=VAL...], for example:

web=1                  # 1 web process
web=3                  # 3 web processes, distributed amongst all hosts
web=3,active=true      # 3 web processes, distributed amongst hosts tagged active=true
db=3,disk=ssd,mem=high # 3 db processes, distributed amongst hosts tagged with
                       # both disk=ssd and mem=high

Ommitting the arguments will show the current scale.

Options:
	-n, --no-wait            don't wait for the scaling events to happen
	-r, --release=<release>  id of release to scale (defaults to current app release)
	-a, --all                show non-zero formations from all releases (only works when listing formations, can't be combined with --release)

Example:

	$ flynn scale
	web=4 worker=2

	$ flynn scale --all
	496d6e74-9db9-4cff-bcce-a3b44015907a (current)
	web=1 worker=2

	632cd907-85ab-4e53-90d0-84635650ec9a
	web=2

	$ flynn scale web=2 worker=5
	scaling web: 4=>2, worker: 2=>5

	02:28:34.333 ==> web flynn-3f656af6f1e44092aa7037046236b203 down
	02:28:34.466 ==> web flynn-ee83def0b8e4455793a43c8c70f5b34e down
	02:28:35.479 ==> worker flynn-84f70ca18c9641ef83a178a19db867a3 up
	02:28:36.508 ==> worker flynn-a3de8c326cc542aa89235e53ba304260 up
	02:28:37.601 ==> worker flynn-e24760c511af4733b01ed5b98aa54647 up

	scale completed in 3.944629056s
`)
}

// minScaleRequestVersion is the minimum API version which supports scaling
// using scale requests
const minScaleRequestVersion = "v20170309.0"

// takes args of the form "web=1[,key=val...]", "worker=3[,key=val...]", etc
func runScale(args *docopt.Args, client controller.Client) error {
	app := mustApp()

	typeSpecs := args.All["<type>=<spec>"].([]string)

	showAll := args.Bool["--all"]

	if len(typeSpecs) > 0 && showAll {
		return fmt.Errorf("ERROR: Can't use --all when scaling")
	}

	releaseID := args.String["--release"]
	if releaseID != "" && showAll {
		return fmt.Errorf("ERROR: Can't use --all in combination with --release")
	}

	if len(typeSpecs) == 0 {
		return showFormations(client, releaseID, showAll, app)
	}

	release, err := determineRelease(client, releaseID, app)
	if err != nil {
		return err
	}

	processes := make(map[string]int, len(typeSpecs))
	tags := make(map[string]map[string]string, len(typeSpecs))
	invalid := make([]string, 0, len(release.Processes))
	for _, arg := range typeSpecs {
		i := strings.IndexRune(arg, '=')
		if i < 0 {
			return fmt.Errorf("ERROR: scale args must be of the form <typ>=<spec>")
		}

		countTags := strings.Split(arg[i+1:], ",")

		count, err := strconv.Atoi(countTags[0])
		if err != nil {
			return fmt.Errorf("ERROR: could not parse quantity in %q", arg)
		} else if count < 0 {
			return fmt.Errorf("ERROR: process quantities cannot be negative in %q", arg)
		}

		processType := arg[:i]
		if _, ok := release.Processes[processType]; ok {
			processes[processType] = count
		} else {
			invalid = append(invalid, fmt.Sprintf("%q", processType))
			continue
		}

		if len(countTags) > 1 {
			processTags := make(map[string]string, len(countTags)-1)
			for i := 1; i < len(countTags); i++ {
				keyVal := strings.SplitN(countTags[i], "=", 2)
				if len(keyVal) == 1 && keyVal[0] != "" {
					processTags[keyVal[0]] = "true"
				} else if len(keyVal) == 2 {
					processTags[keyVal[0]] = keyVal[1]
				}
			}
			tags[processType] = processTags
		}
	}
	if len(invalid) > 0 {
		return fmt.Errorf("ERROR: unknown process types: %s", strings.Join(invalid, ", "))
	}

	opts := ct.ScaleOptions{
		Processes: processes,
		Tags:      tags,
		NoWait:    args.Bool["--no-wait"],
	}

	status, err := client.Status()
	if err != nil {
		return err
	}
	v := version.Parse(status.Version)
	if !v.Dev && v.Before(version.Parse(minScaleRequestVersion)) {
		return runScaleWithJobEvents(client, app, release, opts)
	}
	return runScaleWithScaleRequest(client, app, release, opts)
}

func runScaleWithScaleRequest(client controller.Client, app string, release *ct.Release, opts ct.ScaleOptions) error {
	opts.ScaleRequestCallback = func(req *ct.ScaleRequest) {
		if req.NewProcesses == nil {
			return
		}
		scale := make([]string, 0, len(release.Processes))
		for typ := range release.Processes {
			if count := (*req.NewProcesses)[typ]; count != req.OldProcesses[typ] {
				scale = append(scale, fmt.Sprintf("%s: %d=>%d", typ, req.OldProcesses[typ], count))
			}
		}
		fmt.Printf("scaling %s\n\n", strings.Join(scale, ", "))
	}
	opts.JobEventCallback = func(job *ct.Job) error {
		id := job.ID
		if id == "" {
			id = job.UUID
		}
		fmt.Printf("%s ==> %s %s %s\n", time.Now().Format("15:04:05.000"), job.Type, id, job.State)
		return nil
	}

	start := time.Now()
	if err := client.ScaleAppRelease(app, release.ID, opts); err != nil {
		return err
	}
	if !opts.NoWait {
		fmt.Printf("\nscale completed in %s\n", time.Since(start))
	}
	return nil
}

func runScaleWithJobEvents(client controller.Client, app string, release *ct.Release, opts ct.ScaleOptions) error {
	processes := opts.Processes
	tags := opts.Tags
	formation, err := client.GetFormation(app, release.ID)
	if err == controller.ErrNotFound {
		formation = &ct.Formation{
			AppID:     app,
			ReleaseID: release.ID,
			Processes: make(map[string]int),
		}
	} else if err != nil {
		return err
	}
	if formation.Processes == nil {
		formation.Processes = make(map[string]int, len(processes))
	}
	if formation.Tags == nil {
		formation.Tags = make(map[string]map[string]string, len(tags))
	}

	currentProcs := formation.Processes
	currentTags := formation.Tags

	for k, v := range currentProcs {
		if _, ok := processes[k]; !ok {
			processes[k] = v
		}
	}

	formation.Processes = processes
	formation.Tags = tags

	if scalingComplete(currentProcs, processes) {
		if !utils.FormationTagsEqual(currentTags, tags) {
			fmt.Println("persisting tag change")
			return client.PutFormation(formation)
		}
		fmt.Println("requested scale equals current scale, nothing to do!")
		return nil
	}

	scale := make([]string, 0, len(release.Processes))
	for typ := range release.Processes {
		if currentProcs[typ] != processes[typ] {
			scale = append(scale, fmt.Sprintf("%s: %d=>%d", typ, currentProcs[typ], processes[typ]))
		}
	}
	fmt.Printf("scaling %s\n\n", strings.Join(scale, ", "))

	expected := client.ExpectedScalingEvents(currentProcs, processes, release.Processes, 1)
	watcher, err := client.WatchJobEvents(app, release.ID)
	if err != nil {
		return err
	}
	defer watcher.Close()

	err = client.PutFormation(formation)
	if err != nil || opts.NoWait {
		return err
	}

	start := time.Now()
	err = watcher.WaitFor(expected, ct.DefaultScaleTimeout, func(job *ct.Job) error {
		id := job.ID
		if id == "" {
			id = job.UUID
		}
		fmt.Printf("%s ==> %s %s %s\n", time.Now().Format("15:04:05.000"), job.Type, id, job.State)
		return nil
	})

	if err != nil {
		return err
	}
	fmt.Printf("\nscale completed in %s\n", time.Since(start))
	return nil
}

func showFormations(client controller.Client, releaseID string, showAll bool, app string) error {
	release, err := determineRelease(client, releaseID, app)
	if err != nil {
		return err
	}
	var releases []*ct.Release
	if showAll {
		var err error
		releases, err = client.AppReleaseList(app)
		if err != nil {
			return err
		}
	} else {
		releases = []*ct.Release{release}
	}

	formations := make(map[string]*ct.Formation, len(releases))
	for _, r := range releases {
		formation, err := client.GetFormation(app, r.ID)
		if err != nil && err != controller.ErrNotFound {
			return err
		}
		formations[r.ID] = formation
	}

	for i, r := range releases {
		f := formations[r.ID]
		if f == nil || len(f.Processes) == 0 {
			continue
		}
		if showAll {
			if i > 0 {
				fmt.Println()
			}
			var suffix string
			if r.ID == release.ID {
				suffix = " (current)"
			}
			fmt.Printf("%s%s\n", r.ID, suffix)
		}
		scale := make([]string, 0, len(r.Processes))
		for typ := range r.Processes {
			n := f.Processes[typ]
			if showAll && n == 0 {
				continue
			}
			scale = append(scale, fmt.Sprintf("%s=%d", typ, n))
		}
		fmt.Println(strings.Join(scale, " "))
	}
	return nil
}

func determineRelease(client controller.Client, releaseID, app string) (*ct.Release, error) {
	if releaseID == "" {
		release, err := client.GetAppRelease(app)
		if err == controller.ErrNotFound {
			return nil, errors.New("No app release, specify a release with --release")
		}
		if err != nil {
			return nil, err
		}
		return release, nil
	}
	return client.GetRelease(releaseID)
}

func scalingComplete(actual, expected map[string]int) bool {
	// check all the expected counts are the same in actual
	for typ, count := range expected {
		if actual[typ] != count {
			return false
		}
	}
	// check any counts in actual which aren't in expected are zero
	for typ, count := range actual {
		if _, ok := expected[typ]; !ok && count != 0 {
			return false
		}
	}
	return true
}
<|endoftext|>


<commit_before>// Statistics tool
// @author Robin Verlangen

package main

import (
	"bytes"
	"errors"
	"fmt"
	"log"
	"math"
	"os"
	"os/exec"
	"sort"
	"strconv"
	"strings"

	"github.com/mgutz/ansi"
)

type Statistics struct {
	verticalSep    string
	horizontalSep  string
	colPad         int
	terminalWidth  int
	terminalHeight int
	colorRed       string
	colorGreen     string
	colorReset     string
}

func (s *Statistics) loadTerminalDimensions() {
	cmd := exec.Command("stty", "size")
	cmd.Stdin = os.Stdin
	out, err := cmd.Output()
	if err != nil {
		log.Printf("Failed to determine terminal dimensions, using default: %s", err)
		s.terminalWidth = 200
		s.terminalHeight = 100
		return
	}
	str := strings.TrimSpace(string(out))
	split := strings.Split(str, " ")
	height, _ := strconv.ParseInt(split[0], 10, 0)
	width, _ := strconv.ParseInt(split[1], 10, 0)
	s.terminalHeight = int(height)
	s.terminalWidth = int(width)
	if verbose {
		log.Printf("Terminal dimension %dx%d (WxH)", s.terminalWidth, s.terminalHeight)
	}
}

func (s *Statistics) RenderChart(filter *Filter, inputData map[int]map[int64]int64, flags map[string]bool) (string, error) {
	// Random data (primary is top, secondary is filled, e.g. errors)
	data := make([]int64, 0)
	dataSecondary := make([]int64, 0)

	// Metric source IDs
	metricId := 1
	secondaryMetricId := 2

	// Colors
	primaryColor := "green"
	secondaryColor := "red"

	// Primary sign
	primarySign := "o"
	secondarySign := "*"

	// Flags
	if flags["hide_error"] {
		secondaryMetricId = -1 // Disable errors
	}
	if flags["hide_regular"] {
		// Swap error metric to primary
		metricId = secondaryMetricId
		primaryColor = secondaryColor
		primarySign = secondarySign
		secondaryMetricId = -1 // Disable secondary
	}

	// Validate
	if inputData[metricId] == nil || len(inputData[metricId]) < 1 {
		return "", errors.New("Metrics not available for this filter")
	}
	// To store the keys in slice in sorted order
	var keys []int
	for ts, _ := range inputData[metricId] {
		keys = append(keys, int(ts))
	}
	sort.Ints(keys)
	for _, k := range keys {
		val := inputData[metricId][int64(k)]
		data = append(data, val)

		// Errors
		var secVal int64 = 0
		if secondaryMetricId > 0 && inputData[secondaryMetricId] != nil {
			secVal = inputData[secondaryMetricId][int64(k)]
		}
		dataSecondary = append(dataSecondary, secVal)
	}

	// Width and height for chart
	dataWidth := len(data)
	maxDataLen := s.terminalWidth - 1
	if dataWidth > maxDataLen {
		log.Println("Warning, truncating data to match terminal width")
		data = data[len(data)-maxDataLen:]
		dataSecondary = dataSecondary[len(dataSecondary)-maxDataLen:]
		dataWidth = len(data)
		// @todo Compress data (merge data points and get sums in order to fit in screen)
	}
	maxHeight := int(math.Min(float64(20), float64(s.terminalHeight-4))) // remove some for padding
	maxWidth := int(math.Max(float64(dataWidth), float64(s.terminalWidth)))

	// Scan for min and max
	minVal := int64(math.MaxInt64)
	maxVal := int64(math.MinInt64)
	for _, val := range data {
		if val < minVal {
			minVal = val
		}
		if val > maxVal {
			maxVal = val
		}
	}

	// Dynamic column padding
	s.colPad = int((maxWidth - len(data)) / len(data))

	// Color codes
	var currentColor string = "reset"
	var colorStr string = ""

	// Start to build chart (top to bottom)
	var buf bytes.Buffer
	for line := maxHeight; line >= 0; line-- {
		// Min line val (10/30)=0.3*10
		minLineVal := int64(float64(line) / (float64(maxHeight) / float64(maxVal)))

		// Iterate columns
		for col := 0; col < len(data); col++ {
			// Determine what to write
			if col == 0 && line != 0 {
				// Left axis
				currentColor, colorStr = s.colorStr(currentColor, "reset", s.verticalSep)
				buf.WriteString(colorStr)
			} else if line == 0 {
				// Bottom axis
				currentColor, colorStr = s.colorStr(currentColor, "reset", s.horizontalSep)
				buf.WriteString(colorStr)
			} else {
				// Data point
				colVal := data[col]
				secondaryColVal := dataSecondary[col]

				// Print?
				if colVal >= minLineVal {
					if secondaryColVal >= minLineVal {
						currentColor, colorStr = s.colorStr(currentColor, secondaryColor, secondarySign)
						buf.WriteString(colorStr)
					} else {
						currentColor, colorStr = s.colorStr(currentColor, primaryColor, primarySign)
						buf.WriteString(colorStr)
					}
				} else {
					buf.WriteString(" ")
				}
			}

			// Padding
			if line != 0 {
				buf.WriteString(strings.Repeat(" ", s.colPad)) // Column padding
			} else {
				// Horizontal axis padding
				currentColor, colorStr = s.colorStr(currentColor, "reset", strings.Repeat(s.horizontalSep, s.colPad))
				buf.WriteString(colorStr)
			}
		}
		buf.WriteString("\n") // Close previous line
	}
	buf.WriteString("\n")         // Final whiteline
	buf.WriteString(s.colorReset) // Reset color

	return buf.String(), nil
}

func (s *Statistics) colorStr(currentColor string, desiredColorName string, str string) (string, string) {
	if currentColor == desiredColorName {
		return currentColor, str
	}
	colorStr := ""
	if desiredColorName == "green" {
		colorStr = s.colorGreen
	} else if desiredColorName == "red" {
		colorStr = s.colorRed
	} else if desiredColorName == "reset" {
		colorStr = s.colorReset
	}
	return desiredColorName, fmt.Sprintf("%s%s", colorStr, str)
}

func newStatistics() *Statistics {
	s := &Statistics{
		verticalSep:   "|",
		horizontalSep: "_",
		colPad:        3,
		colorRed:      ansi.ColorCode("red"),
		colorGreen:    ansi.ColorCode("green"),
		colorReset:    ansi.ColorCode("reset"),
	}
	s.loadTerminalDimensions()
	return s
}
<commit_msg>Disable colors if tty not available<commit_after>// Statistics tool
// @author Robin Verlangen

package main

import (
	"bytes"
	"errors"
	"fmt"
	"log"
	"math"
	"os"
	"os/exec"
	"sort"
	"strconv"
	"strings"

	"github.com/mgutz/ansi"
)

type Statistics struct {
	verticalSep    string
	horizontalSep  string
	colPad         int
	terminalWidth  int
	terminalHeight int
	colorRed       string
	colorGreen     string
	colorReset     string
	colorEnabled   bool
}

func (s *Statistics) loadTerminalDimensions() {
	cmd := exec.Command("stty", "size")
	cmd.Stdin = os.Stdin
	out, err := cmd.Output()
	if err != nil {
		log.Printf("Failed to determine terminal dimensions: %s", err)
		log.Printf("Disabling color")
		s.colorEnabled = false
		s.terminalWidth = 100
		s.terminalHeight = 50
		if verbose {
			log.Printf("Terminal dimension %dx%d (WxH)", s.terminalWidth, s.terminalHeight)
		}
		return
	}
	str := strings.TrimSpace(string(out))
	split := strings.Split(str, " ")
	height, _ := strconv.ParseInt(split[0], 10, 0)
	width, _ := strconv.ParseInt(split[1], 10, 0)
	s.terminalHeight = int(height)
	s.terminalWidth = int(width)
	if verbose {
		log.Printf("Terminal dimension %dx%d (WxH)", s.terminalWidth, s.terminalHeight)
	}
}

func (s *Statistics) RenderChart(filter *Filter, inputData map[int]map[int64]int64, flags map[string]bool) (string, error) {
	// Random data (primary is top, secondary is filled, e.g. errors)
	data := make([]int64, 0)
	dataSecondary := make([]int64, 0)

	// Metric source IDs
	metricId := 1
	secondaryMetricId := 2

	// Colors
	primaryColor := "green"
	secondaryColor := "red"
	if s.colorEnabled {
		primaryColor = "reset"
		secondaryColor = "reset"
	}

	// Primary sign
	primarySign := "o"
	secondarySign := "*"

	// Flags
	if flags["hide_error"] {
		secondaryMetricId = -1 // Disable errors
	}
	if flags["hide_regular"] {
		// Swap error metric to primary
		metricId = secondaryMetricId
		primaryColor = secondaryColor
		primarySign = secondarySign
		secondaryMetricId = -1 // Disable secondary
	}

	// Validate
	if inputData[metricId] == nil || len(inputData[metricId]) < 1 {
		return "", errors.New("Metrics not available for this filter")
	}
	// To store the keys in slice in sorted order
	var keys []int
	for ts, _ := range inputData[metricId] {
		keys = append(keys, int(ts))
	}
	sort.Ints(keys)
	for _, k := range keys {
		val := inputData[metricId][int64(k)]
		data = append(data, val)

		// Errors
		var secVal int64 = 0
		if secondaryMetricId > 0 && inputData[secondaryMetricId] != nil {
			secVal = inputData[secondaryMetricId][int64(k)]
		}
		dataSecondary = append(dataSecondary, secVal)
	}

	// Width and height for chart
	dataWidth := len(data)
	maxDataLen := s.terminalWidth - 1
	if dataWidth > maxDataLen {
		log.Println("Warning, truncating data to match terminal width")
		data = data[len(data)-maxDataLen:]
		dataSecondary = dataSecondary[len(dataSecondary)-maxDataLen:]
		dataWidth = len(data)
		// @todo Compress data (merge data points and get sums in order to fit in screen)
	}
	maxHeight := int(math.Min(float64(20), float64(s.terminalHeight-4))) // remove some for padding
	maxWidth := int(math.Max(float64(dataWidth), float64(s.terminalWidth)))

	// Scan for min and max
	minVal := int64(math.MaxInt64)
	maxVal := int64(math.MinInt64)
	for _, val := range data {
		if val < minVal {
			minVal = val
		}
		if val > maxVal {
			maxVal = val
		}
	}

	// Dynamic column padding
	s.colPad = int((maxWidth - len(data)) / len(data))

	// Color codes
	var currentColor string = "reset"
	var colorStr string = ""

	// Start to build chart (top to bottom)
	var buf bytes.Buffer
	for line := maxHeight; line >= 0; line-- {
		// Min line val (10/30)=0.3*10
		minLineVal := int64(float64(line) / (float64(maxHeight) / float64(maxVal)))

		// Iterate columns
		for col := 0; col < len(data); col++ {
			// Determine what to write
			if col == 0 && line != 0 {
				// Left axis
				currentColor, colorStr = s.colorStr(currentColor, "reset", s.verticalSep)
				buf.WriteString(colorStr)
			} else if line == 0 {
				// Bottom axis
				currentColor, colorStr = s.colorStr(currentColor, "reset", s.horizontalSep)
				buf.WriteString(colorStr)
			} else {
				// Data point
				colVal := data[col]
				secondaryColVal := dataSecondary[col]

				// Print?
				if colVal >= minLineVal {
					if secondaryColVal >= minLineVal {
						currentColor, colorStr = s.colorStr(currentColor, secondaryColor, secondarySign)
						buf.WriteString(colorStr)
					} else {
						currentColor, colorStr = s.colorStr(currentColor, primaryColor, primarySign)
						buf.WriteString(colorStr)
					}
				} else {
					buf.WriteString(" ")
				}
			}

			// Padding
			if line != 0 {
				buf.WriteString(strings.Repeat(" ", s.colPad)) // Column padding
			} else {
				// Horizontal axis padding
				currentColor, colorStr = s.colorStr(currentColor, "reset", strings.Repeat(s.horizontalSep, s.colPad))
				buf.WriteString(colorStr)
			}
		}
		buf.WriteString("\n") // Close previous line
	}
	buf.WriteString("\n")         // Final whiteline
	buf.WriteString(s.colorReset) // Reset color

	return buf.String(), nil
}

func (s *Statistics) colorStr(currentColor string, desiredColorName string, str string) (string, string) {
	if currentColor == desiredColorName {
		return currentColor, str
	}
	colorStr := ""
	if desiredColorName == "green" {
		colorStr = s.colorGreen
	} else if desiredColorName == "red" {
		colorStr = s.colorRed
	} else if desiredColorName == "reset" {
		colorStr = s.colorReset
	}
	return desiredColorName, fmt.Sprintf("%s%s", colorStr, str)
}

func newStatistics() *Statistics {
	s := &Statistics{
		verticalSep:   "|",
		horizontalSep: "_",
		colPad:        3,
		colorEnabled:  true,
		colorRed:      ansi.ColorCode("red"),
		colorGreen:    ansi.ColorCode("green"),
		colorReset:    ansi.ColorCode("reset"),
	}
	s.loadTerminalDimensions()
	return s
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"crypto/rand"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"time"
	"strings"
)

var queuePath string
var scriptsPath string

type RunStat struct {
	ScriptName string
	HostName   string
	UserName   string
	StartTime  int64
	EndTime    int64
	Duration   float64
	Status     string
}

func main() {
	flag.StringVar(&queuePath, "queue-path", DefaultQueuePath(), "queue path")
	flag.StringVar(&scriptsPath, "scripts-path", DefaultScriptsPath(), "scripts path")
	flag.Parse()
	CreatePaths()
	if flag.NArg() == 2 {
		switch flag.Arg(0) {
		case "run":
			runScript(flag.Arg(1))
		case "send-queue":
			runQueue(flag.Arg(1))
		}
	}
}

func UserHome() (path string) {
	me, err := user.Current()
	DieIfErr(err)
	return me.HomeDir
}

func UserName() (username string) {
	me, err := user.Current()
	DieIfErr(err)
	return me.Username
}

func HostName() (hostname string) {
	hostname, err := os.Hostname()
	DieIfErr(err)
	return hostname
}

func DefaultScriptsPath() (path string) {
	return filepath.Join(UserHome(), ".clog-scripts")
}

func DefaultQueuePath() (path string) {
	return filepath.Join(UserHome(), ".clog-queue")
}

func QueuePath() (path string) {
	return queuePath
}

func ScriptsPath() (path string) {
	return scriptsPath
}

func createPath(path string) {
	if _, err := os.Stat(path); err != nil {
		if os.IsNotExist(err) {
			log.Print("creating path " + path)
			if err := os.Mkdir(path, 0700); err != nil {
				panic("error creating directory " + path)
			}
		} else {
			panic("error getting directory information!")
		}
	}
}

func CreatePaths() {
	createPath(ScriptsPath())
	createPath(QueuePath())
}

func DieIfErr(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

// Saves queue metadata for delivery.
func (runStat *RunStat) writeQueueMetadata(queuePath string) {
	fp, err := os.OpenFile(queuePath+".meta.tmp", os.O_CREATE|os.O_WRONLY, 0600)
	DieIfErr(err)
	json, err := json.Marshal(runStat)
	DieIfErr(err)
	fp.Write(json)
	fp.Close()
	err = os.Rename(queuePath+".meta.tmp", queuePath+".meta")
	DieIfErr(err)
}

func runScript(script string) {

	runStat := new(RunStat)
	runStat.ScriptName = script

	scriptpath := filepath.Join(ScriptsPath(), script)
	log.Print("running script ", scriptpath)

	id := GenId()
	log.Print("queue id: ", id)
	queueLogPath := filepath.Join(QueuePath(), id)
	log.Print("queue path: ", queueLogPath)

	runStat.HostName = HostName()
	runStat.UserName = UserName()

	runStat.StartTime = time.Now().Unix()
	startTime := time.Now().UnixNano()

	cmd := exec.Command(scriptpath)

	// Merging stdout and stderr.
	stdout, err := cmd.StdoutPipe()
	DieIfErr(err)
	cmd.Stderr = cmd.Stdout

	scanner := bufio.NewScanner(stdout)

	err = cmd.Start()
	DieIfErr(err)

	queuelog, err := os.OpenFile(queueLogPath+".out", os.O_CREATE|os.O_WRONLY, 0600)
	DieIfErr(err)

	log.Print("command output starts")
	for scanner.Scan() {
		line := scanner.Text() + "\n"
		fmt.Print(line)
		queuelog.WriteString(line)
	}
	log.Print("command output ends")
	queuelog.Close()

	// Checking command exit status.
	runStat.Status = "ok"
	if err := cmd.Wait(); err != nil {
		if _, ok := err.(*exec.ExitError); ok {
			runStat.Status = "fail"
		} else {
			log.Fatal(err)
		}
	}

	endTime := time.Now().UnixNano()
	runStat.EndTime = time.Now().Unix()
	runStat.Duration = float64(endTime - startTime)/1000000000.0

	log.Print("status: ", runStat.Status)
	log.Print("duration: ", fmt.Sprintf("%0.3fs", runStat.Duration))
	log.Print("user: ", runStat.UserName)
	log.Print("hostname: ", runStat.HostName)

	runStat.writeQueueMetadata(queueLogPath)

}

func runQueue(url string) {
	log.Print("target server: ", url)
	queuePath := QueuePath()
	log.Print("queue path: ", queuePath)
	files, err := ioutil.ReadDir(queuePath)
	DieIfErr(err)
	for _, file := range files {
		name := file.Name()
		if len(name) == 41 && strings.HasSuffix(name, ".meta") {
			queueId := name[:36]
			log.Print("dispatching ", queueId)
		}
	}
}

func GenId() (id string) {
	// Poor man's UUID generator
	const validchars = "0123456789abcdef"
	b := make([]byte, 36)
	if _, err := rand.Read(b); err != nil {
		panic(err)
	}
	for i, j := range b {
		b[i] = validchars[j%byte(len(validchars))]
	}
	b[8] = '-'
	b[13] = '-'
	b[18] = '-'
	b[23] = '-'
	b[14] = '4'
	return string(b)
}
<commit_msg>read queue metadata<commit_after>package main

import (
	"bufio"
	"crypto/rand"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"path"
	"time"
	"strings"
)

var queuePath string
var scriptsPath string

type RunStat struct {
	ScriptName string
	HostName   string
	UserName   string
	StartTime  int64
	EndTime    int64
	Duration   float64
	Status     string
}

func main() {
	flag.StringVar(&queuePath, "queue-path", DefaultQueuePath(), "queue path")
	flag.StringVar(&scriptsPath, "scripts-path", DefaultScriptsPath(), "scripts path")
	flag.Parse()
	CreatePaths()
	if flag.NArg() == 2 {
		switch flag.Arg(0) {
		case "run":
			runScript(flag.Arg(1))
		case "send-queue":
			runQueue(flag.Arg(1))
		}
	}
}

func UserHome() (path string) {
	me, err := user.Current()
	DieIfErr(err)
	return me.HomeDir
}

func UserName() (username string) {
	me, err := user.Current()
	DieIfErr(err)
	return me.Username
}

func HostName() (hostname string) {
	hostname, err := os.Hostname()
	DieIfErr(err)
	return hostname
}

func DefaultScriptsPath() (path string) {
	return filepath.Join(UserHome(), ".clog-scripts")
}

func DefaultQueuePath() (path string) {
	return filepath.Join(UserHome(), ".clog-queue")
}

func QueuePath() (path string) {
	return queuePath
}

func ScriptsPath() (path string) {
	return scriptsPath
}

func createPath(path string) {
	if _, err := os.Stat(path); err != nil {
		if os.IsNotExist(err) {
			log.Print("creating path " + path)
			if err := os.Mkdir(path, 0700); err != nil {
				panic("error creating directory " + path)
			}
		} else {
			panic("error getting directory information!")
		}
	}
}

func CreatePaths() {
	createPath(ScriptsPath())
	createPath(QueuePath())
}

func DieIfErr(err error) {
	if err != nil {
		log.Fatal(err)
	}
}

// Saves queue metadata for delivery.
func (runStat *RunStat) writeQueueMetadata(queuePath string) {
	fp, err := os.OpenFile(queuePath+".meta.tmp", os.O_CREATE|os.O_WRONLY, 0600)
	DieIfErr(err)
	json, err := json.Marshal(runStat)
	DieIfErr(err)
	fp.Write(json)
	fp.Close()
	err = os.Rename(queuePath+".meta.tmp", queuePath+".meta")
	DieIfErr(err)
}

// Read queue metadata.
func (runStat *RunStat) readQueueMetadata(queuePath string) {
	fp, err := os.Open(queuePath)
	DieIfErr(err)
	b := make([]byte, 64000)
	n, err := fp.Read(b)
	DieIfErr(err)
	err = json.Unmarshal(b[:n], &runStat)
	DieIfErr(err)
}

func runScript(script string) {

	runStat := new(RunStat)
	runStat.ScriptName = script

	scriptpath := filepath.Join(ScriptsPath(), script)
	log.Print("running script ", scriptpath)

	id := GenId()
	log.Print("queue id: ", id)
	queueLogPath := filepath.Join(QueuePath(), id)
	log.Print("queue path: ", queueLogPath)

	runStat.HostName = HostName()
	runStat.UserName = UserName()

	runStat.StartTime = time.Now().Unix()
	startTime := time.Now().UnixNano()

	cmd := exec.Command(scriptpath)

	// Merging stdout and stderr.
	stdout, err := cmd.StdoutPipe()
	DieIfErr(err)
	cmd.Stderr = cmd.Stdout

	scanner := bufio.NewScanner(stdout)

	err = cmd.Start()
	DieIfErr(err)

	queuelog, err := os.OpenFile(queueLogPath+".out", os.O_CREATE|os.O_WRONLY, 0600)
	DieIfErr(err)

	log.Print("command output starts")
	for scanner.Scan() {
		line := scanner.Text() + "\n"
		fmt.Print(line)
		queuelog.WriteString(line)
	}
	log.Print("command output ends")
	queuelog.Close()

	// Checking command exit status.
	runStat.Status = "ok"
	if err := cmd.Wait(); err != nil {
		if _, ok := err.(*exec.ExitError); ok {
			runStat.Status = "fail"
		} else {
			log.Fatal(err)
		}
	}

	endTime := time.Now().UnixNano()
	runStat.EndTime = time.Now().Unix()
	runStat.Duration = float64(endTime - startTime)/1000000000.0

	log.Print("status: ", runStat.Status)
	log.Print("duration: ", fmt.Sprintf("%0.3fs", runStat.Duration))
	log.Print("user: ", runStat.UserName)
	log.Print("hostname: ", runStat.HostName)

	runStat.writeQueueMetadata(queueLogPath)

}

func runQueue(url string) {
	log.Print("target server: ", url)
	queuePath := QueuePath()
	log.Print("queue path: ", queuePath)
	files, err := ioutil.ReadDir(queuePath)
	DieIfErr(err)
	for _, file := range files {
		name := file.Name()
		if len(name) == 41 && strings.HasSuffix(name, ".meta") {
			queueId := name[:36]
			log.Print("dispatching ", queueId)
			runStat := new(RunStat)
			runStat.readQueueMetadata(path.Join(queuePath, queueId) + ".meta")
		}
	}
}

func GenId() (id string) {
	// Poor man's UUID generator
	const validchars = "0123456789abcdef"
	b := make([]byte, 36)
	if _, err := rand.Read(b); err != nil {
		panic(err)
	}
	for i, j := range b {
		b[i] = validchars[j%byte(len(validchars))]
	}
	b[8] = '-'
	b[13] = '-'
	b[18] = '-'
	b[23] = '-'
	b[14] = '4'
	return string(b)
}
<|endoftext|>


<commit_before>package master

import (
	"fmt"
	"log"
	"net/http"

	auth "github.com/abbot/go-http-auth"
	"github.com/h2oai/steamY/master/az"
)

type DefaultAz struct {
	directory az.Directory
}

func NewDefaultAz(directory az.Directory) *DefaultAz {
	return &DefaultAz{directory}
}

func (a *DefaultAz) Authenticate(username string) string {
	pz, err := a.directory.Lookup(username)
	if err != nil {
		log.Printf("User %s read failed: %s\n", username, err)
		return ""
	}

	if pz == nil {
		log.Printf("User %s does not exist\n", username)
		return ""
	}
	log.Println("User logged in:", username)
	return pz.Password()
}

func (a *DefaultAz) Identify(r *http.Request) (az.Principal, error) {
	username := r.Header.Get(auth.AuthUsernameHeader)
	log.Println("User identified:", username)
	pz, err := a.directory.Lookup(username)
	if err != nil {
		return nil, err
	}

	if pz == nil {
		return nil, fmt.Errorf("User %s does not exist\n", username)
	}

	return pz, nil
}

func serveNoop(w http.ResponseWriter, r *http.Request) {}
func authNoop(user, realm string) string               { return "" }
<commit_msg>STEAM-315 Remove az logging<commit_after>package master

import (
	"fmt"
	"log"
	"net/http"

	auth "github.com/abbot/go-http-auth"
	"github.com/h2oai/steamY/master/az"
)

type DefaultAz struct {
	directory az.Directory
}

func NewDefaultAz(directory az.Directory) *DefaultAz {
	return &DefaultAz{directory}
}

func (a *DefaultAz) Authenticate(username string) string {
	pz, err := a.directory.Lookup(username)
	if err != nil {
		log.Printf("User %s read failed: %s\n", username, err)
		return ""
	}

	if pz == nil {
		log.Printf("User %s does not exist\n", username)
		return ""
	}
	return pz.Password()
}

func (a *DefaultAz) Identify(r *http.Request) (az.Principal, error) {
	username := r.Header.Get(auth.AuthUsernameHeader)
	pz, err := a.directory.Lookup(username)
	if err != nil {
		return nil, err
	}

	if pz == nil {
		return nil, fmt.Errorf("User %s does not exist\n", username)
	}

	return pz, nil
}

func serveNoop(w http.ResponseWriter, r *http.Request) {}
func authNoop(user, realm string) string               { return "" }
<|endoftext|>


<commit_before>package cloud

import (
	"context"
	"errors"
	"fmt"
	"time"

	api "github.com/appscode/pharmer/apis/v1alpha1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	clientcmd "k8s.io/client-go/tools/clientcmd/api/v1"
	"k8s.io/client-go/util/cert"
)

func List(ctx context.Context, opts metav1.ListOptions) ([]*api.Cluster, error) {
	return Store(ctx).Clusters().List(opts)
}

func Get(ctx context.Context, name string) (*api.Cluster, error) {
	return Store(ctx).Clusters().Get(name)
}

func Create(ctx context.Context, cluster *api.Cluster) (*api.Cluster, error) {
	if cluster == nil {
		return nil, errors.New("missing cluster")
	} else if cluster.Name == "" {
		return nil, errors.New("missing cluster name")
	} else if cluster.Spec.KubernetesVersion == "" {
		return nil, errors.New("missing cluster version")
	}

	_, err := Store(ctx).Clusters().Get(cluster.Name)
	if err == nil {
		return nil, fmt.Errorf("cluster exists with name `%s`", cluster.Name)
	}

	cm, err := GetCloudManager(cluster.Spec.Cloud.CloudProvider, ctx)
	if err != nil {
		return nil, err
	}
	if cluster, err = cm.DefaultSpec(cluster); err != nil {
		return nil, err
	}

	if cluster.Spec.MasterKubeadmVersion == "" {
		cluster.Spec.MasterKubeadmVersion, err = GetLatestKubeadmVerson()
		if err != nil {
			return nil, err
		}
	}

	if cluster, err = Store(ctx).Clusters().Create(cluster); err != nil {
		return nil, err
	}

	if ctx, err = CreateCACertificates(ctx, cluster); err != nil {
		return nil, err
	}
	if ctx, err = CreateSSHKey(ctx, cluster); err != nil {
		return nil, err
	}
	if _, err = cm.CreateMasterNodeGroup(cluster); err != nil {
		return nil, err
	}
	if _, err = Store(ctx).Clusters().Update(cluster); err != nil {
		return nil, err
	}
	return cluster, nil
}

func Update(ctx context.Context, cluster *api.Cluster) (*api.Cluster, error) {
	if cluster == nil {
		return nil, errors.New("missing cluster")
	} else if cluster.Name == "" {
		return nil, errors.New("missing cluster name")
	} else if cluster.Spec.KubernetesVersion == "" {
		return nil, errors.New("missing cluster version")
	}

	existing, err := Store(ctx).Clusters().Get(cluster.Name)
	if err != nil {
		return nil, fmt.Errorf("cluster `%s` does not exist. Reason: %v", cluster.Name, err)
	}
	cluster.Status = existing.Status
	cluster.Generation = time.Now().UnixNano()

	return Store(ctx).Clusters().Update(cluster)
}

func Delete(ctx context.Context, name string) (*api.Cluster, error) {
	if name == "" {
		return nil, errors.New("missing cluster name")
	}

	cluster, err := Store(ctx).Clusters().Get(name)
	if err != nil {
		return nil, fmt.Errorf("cluster `%s` does not exist. Reason: %v", name, err)
	}
	cluster.DeletionTimestamp = &metav1.Time{Time: time.Now()}
	cluster.Status.Phase = api.ClusterDeleting

	return Store(ctx).Clusters().Update(cluster)
}

func DeleteNG(ctx context.Context, nodeGroupName, clusterName string) error {
	if clusterName == "" {
		return errors.New("missing cluster name")
	}
	if nodeGroupName == "" {
		return errors.New("missing nodegroup name")
	}

	if _, err := Store(ctx).Clusters().Get(clusterName); err != nil {
		return fmt.Errorf("cluster `%s` does not exist. Reason: %v", clusterName, err)
	}

	nodeGroup, err := Store(ctx).NodeGroups(clusterName).Get(nodeGroupName)
	if err != nil {
		return fmt.Errorf(`nodegroup not found`)
	}

	if !nodeGroup.IsMaster() {
		nodeGroup.DeletionTimestamp = &metav1.Time{Time: time.Now()}
		_, err := Store(ctx).NodeGroups(clusterName).Update(nodeGroup)
		return err
	}

	return nil
}

func GetSSHConfig(ctx context.Context, cluster *api.Cluster, nodeName string) (*api.SSHConfig, error) {
	client, err := NewAdminClient(ctx, cluster)
	if err != nil {
		return nil, err
	}
	node, err := client.CoreV1().Nodes().Get(nodeName, metav1.GetOptions{})
	if err != nil {
		return nil, err
	}
	ctx, err = LoadSSHKey(ctx, cluster)
	if err != nil {
		return nil, err
	}

	cm, err := GetCloudManager(cluster.Spec.Cloud.CloudProvider, ctx)
	if err != nil {
		return nil, err
	}
	return cm.GetSSHConfig(cluster, node)
}

func GetAdminConfig(ctx context.Context, cluster *api.Cluster) (*clientcmd.Config, error) {
	var err error
	ctx, err = LoadCACertificates(ctx, cluster)
	if err != nil {
		return nil, err
	}
	adminCert, adminKey, err := CreateAdminCertificate(ctx)
	if err != nil {
		return nil, err
	}

	var (
		clusterName = fmt.Sprintf("%s.pharmer", cluster.Name)
		userName    = fmt.Sprintf("cluster-admin@%s.pharmer", cluster.Name)
		ctxName     = fmt.Sprintf("cluster-admin@%s.pharmer", cluster.Name)
	)
	cfg := clientcmd.Config{
		APIVersion: "v1",
		Kind:       "Config",
		Preferences: clientcmd.Preferences{
			Colors: true,
		},
		Clusters: []clientcmd.NamedCluster{
			{
				Name: clusterName,
				Cluster: clientcmd.Cluster{
					Server: cluster.APIServerURL(),
					CertificateAuthorityData: cert.EncodeCertPEM(CACert(ctx)),
				},
			},
		},
		AuthInfos: []clientcmd.NamedAuthInfo{
			{
				Name: userName,
				AuthInfo: clientcmd.AuthInfo{
					ClientCertificateData: cert.EncodeCertPEM(adminCert),
					ClientKeyData:         cert.EncodePrivateKeyPEM(adminKey),
				},
			},
		},
		Contexts: []clientcmd.NamedContext{
			{
				Name: ctxName,
				Context: clientcmd.Context{
					Cluster:  clusterName,
					AuthInfo: userName,
				},
			},
		},
		CurrentContext: ctxName,
	}
	return &cfg, nil
}

func Apply(ctx context.Context, name string, dryRun bool) ([]api.Action, error) {
	if name == "" {
		return nil, errors.New("missing cluster name")
	}

	cluster, err := Store(ctx).Clusters().Get(name)
	if err != nil {
		return nil, fmt.Errorf("cluster `%s` does not exist. Reason: %v", name, err)
	}

	cm, err := GetCloudManager(cluster.Spec.Cloud.CloudProvider, ctx)
	if err != nil {
		return nil, err
	}

	return cm.Apply(cluster, dryRun)
}

func CheckForUpdates(ctx context.Context, name string) (string, error) {
	if name == "" {
		return "", errors.New("missing cluster name")
	}

	cluster, err := Store(ctx).Clusters().Get(name)
	if err != nil {
		return "", fmt.Errorf("cluster `%s` does not exist. Reason: %v", name, err)
	}
	if cluster.Status.Phase == "" {
		return "", fmt.Errorf("cluster `%s` is in unknown phase", cluster.Name)
	}
	if cluster.Status.Phase != api.ClusterReady {
		return "", fmt.Errorf("cluster `%s` is not ready", cluster.Name)
	}
	if cluster.Status.Phase == api.ClusterDeleted {
		return "", nil
	}
	if ctx, err = LoadCACertificates(ctx, cluster); err != nil {
		return "", err
	}
	if ctx, err = LoadSSHKey(ctx, cluster); err != nil {
		return "", err
	}
	kc, err := NewAdminClient(ctx, cluster)
	if err != nil {
		return "", err
	}
	cm, err := GetCloudManager(cluster.Spec.Cloud.CloudProvider, ctx)
	if err != nil {
		return "", err
	}
	upm := NewUpgradeManager(ctx, cm, kc, cluster)
	upgrades, err := upm.GetAvailableUpgrades()
	if err != nil {
		return "", err
	}
	upm.PrintAvailableUpgrades(upgrades)
	return "", nil
}

func Edit(ctx context.Context, name, version string) (*api.Cluster, error) {
	if name == "" {
		return nil, errors.New("missing cluster name")
	}

	cluster, err := Store(ctx).Clusters().Get(name)
	if err != nil {
		return nil, fmt.Errorf("cluster `%s` does not exist. Reason: %v", name, err)
	}
	if cluster.Spec.KubernetesVersion != version {
		cluster.Spec.KubernetesVersion = version
		cluster.Generation = time.Now().UnixNano()
		cluster.Status.Phase = api.ClusterUpgrading
	}
	return Store(ctx).Clusters().Update(cluster)
}
<commit_msg>Don't auto set Kubeadm version<commit_after>package cloud

import (
	"context"
	"errors"
	"fmt"
	"time"

	api "github.com/appscode/pharmer/apis/v1alpha1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	clientcmd "k8s.io/client-go/tools/clientcmd/api/v1"
	"k8s.io/client-go/util/cert"
)

func List(ctx context.Context, opts metav1.ListOptions) ([]*api.Cluster, error) {
	return Store(ctx).Clusters().List(opts)
}

func Get(ctx context.Context, name string) (*api.Cluster, error) {
	return Store(ctx).Clusters().Get(name)
}

func Create(ctx context.Context, cluster *api.Cluster) (*api.Cluster, error) {
	if cluster == nil {
		return nil, errors.New("missing cluster")
	} else if cluster.Name == "" {
		return nil, errors.New("missing cluster name")
	} else if cluster.Spec.KubernetesVersion == "" {
		return nil, errors.New("missing cluster version")
	}

	_, err := Store(ctx).Clusters().Get(cluster.Name)
	if err == nil {
		return nil, fmt.Errorf("cluster exists with name `%s`", cluster.Name)
	}

	cm, err := GetCloudManager(cluster.Spec.Cloud.CloudProvider, ctx)
	if err != nil {
		return nil, err
	}
	if cluster, err = cm.DefaultSpec(cluster); err != nil {
		return nil, err
	}
	if cluster, err = Store(ctx).Clusters().Create(cluster); err != nil {
		return nil, err
	}

	if ctx, err = CreateCACertificates(ctx, cluster); err != nil {
		return nil, err
	}
	if ctx, err = CreateSSHKey(ctx, cluster); err != nil {
		return nil, err
	}
	if _, err = cm.CreateMasterNodeGroup(cluster); err != nil {
		return nil, err
	}
	if _, err = Store(ctx).Clusters().Update(cluster); err != nil {
		return nil, err
	}
	return cluster, nil
}

func Update(ctx context.Context, cluster *api.Cluster) (*api.Cluster, error) {
	if cluster == nil {
		return nil, errors.New("missing cluster")
	} else if cluster.Name == "" {
		return nil, errors.New("missing cluster name")
	} else if cluster.Spec.KubernetesVersion == "" {
		return nil, errors.New("missing cluster version")
	}

	existing, err := Store(ctx).Clusters().Get(cluster.Name)
	if err != nil {
		return nil, fmt.Errorf("cluster `%s` does not exist. Reason: %v", cluster.Name, err)
	}
	cluster.Status = existing.Status
	cluster.Generation = time.Now().UnixNano()

	return Store(ctx).Clusters().Update(cluster)
}

func Delete(ctx context.Context, name string) (*api.Cluster, error) {
	if name == "" {
		return nil, errors.New("missing cluster name")
	}

	cluster, err := Store(ctx).Clusters().Get(name)
	if err != nil {
		return nil, fmt.Errorf("cluster `%s` does not exist. Reason: %v", name, err)
	}
	cluster.DeletionTimestamp = &metav1.Time{Time: time.Now()}
	cluster.Status.Phase = api.ClusterDeleting

	return Store(ctx).Clusters().Update(cluster)
}

func DeleteNG(ctx context.Context, nodeGroupName, clusterName string) error {
	if clusterName == "" {
		return errors.New("missing cluster name")
	}
	if nodeGroupName == "" {
		return errors.New("missing nodegroup name")
	}

	if _, err := Store(ctx).Clusters().Get(clusterName); err != nil {
		return fmt.Errorf("cluster `%s` does not exist. Reason: %v", clusterName, err)
	}

	nodeGroup, err := Store(ctx).NodeGroups(clusterName).Get(nodeGroupName)
	if err != nil {
		return fmt.Errorf(`nodegroup not found`)
	}

	if !nodeGroup.IsMaster() {
		nodeGroup.DeletionTimestamp = &metav1.Time{Time: time.Now()}
		_, err := Store(ctx).NodeGroups(clusterName).Update(nodeGroup)
		return err
	}

	return nil
}

func GetSSHConfig(ctx context.Context, cluster *api.Cluster, nodeName string) (*api.SSHConfig, error) {
	client, err := NewAdminClient(ctx, cluster)
	if err != nil {
		return nil, err
	}
	node, err := client.CoreV1().Nodes().Get(nodeName, metav1.GetOptions{})
	if err != nil {
		return nil, err
	}
	ctx, err = LoadSSHKey(ctx, cluster)
	if err != nil {
		return nil, err
	}

	cm, err := GetCloudManager(cluster.Spec.Cloud.CloudProvider, ctx)
	if err != nil {
		return nil, err
	}
	return cm.GetSSHConfig(cluster, node)
}

func GetAdminConfig(ctx context.Context, cluster *api.Cluster) (*clientcmd.Config, error) {
	var err error
	ctx, err = LoadCACertificates(ctx, cluster)
	if err != nil {
		return nil, err
	}
	adminCert, adminKey, err := CreateAdminCertificate(ctx)
	if err != nil {
		return nil, err
	}

	var (
		clusterName = fmt.Sprintf("%s.pharmer", cluster.Name)
		userName    = fmt.Sprintf("cluster-admin@%s.pharmer", cluster.Name)
		ctxName     = fmt.Sprintf("cluster-admin@%s.pharmer", cluster.Name)
	)
	cfg := clientcmd.Config{
		APIVersion: "v1",
		Kind:       "Config",
		Preferences: clientcmd.Preferences{
			Colors: true,
		},
		Clusters: []clientcmd.NamedCluster{
			{
				Name: clusterName,
				Cluster: clientcmd.Cluster{
					Server: cluster.APIServerURL(),
					CertificateAuthorityData: cert.EncodeCertPEM(CACert(ctx)),
				},
			},
		},
		AuthInfos: []clientcmd.NamedAuthInfo{
			{
				Name: userName,
				AuthInfo: clientcmd.AuthInfo{
					ClientCertificateData: cert.EncodeCertPEM(adminCert),
					ClientKeyData:         cert.EncodePrivateKeyPEM(adminKey),
				},
			},
		},
		Contexts: []clientcmd.NamedContext{
			{
				Name: ctxName,
				Context: clientcmd.Context{
					Cluster:  clusterName,
					AuthInfo: userName,
				},
			},
		},
		CurrentContext: ctxName,
	}
	return &cfg, nil
}

func Apply(ctx context.Context, name string, dryRun bool) ([]api.Action, error) {
	if name == "" {
		return nil, errors.New("missing cluster name")
	}

	cluster, err := Store(ctx).Clusters().Get(name)
	if err != nil {
		return nil, fmt.Errorf("cluster `%s` does not exist. Reason: %v", name, err)
	}

	cm, err := GetCloudManager(cluster.Spec.Cloud.CloudProvider, ctx)
	if err != nil {
		return nil, err
	}

	return cm.Apply(cluster, dryRun)
}

func CheckForUpdates(ctx context.Context, name string) (string, error) {
	if name == "" {
		return "", errors.New("missing cluster name")
	}

	cluster, err := Store(ctx).Clusters().Get(name)
	if err != nil {
		return "", fmt.Errorf("cluster `%s` does not exist. Reason: %v", name, err)
	}
	if cluster.Status.Phase == "" {
		return "", fmt.Errorf("cluster `%s` is in unknown phase", cluster.Name)
	}
	if cluster.Status.Phase != api.ClusterReady {
		return "", fmt.Errorf("cluster `%s` is not ready", cluster.Name)
	}
	if cluster.Status.Phase == api.ClusterDeleted {
		return "", nil
	}
	if ctx, err = LoadCACertificates(ctx, cluster); err != nil {
		return "", err
	}
	if ctx, err = LoadSSHKey(ctx, cluster); err != nil {
		return "", err
	}
	kc, err := NewAdminClient(ctx, cluster)
	if err != nil {
		return "", err
	}
	cm, err := GetCloudManager(cluster.Spec.Cloud.CloudProvider, ctx)
	if err != nil {
		return "", err
	}
	upm := NewUpgradeManager(ctx, cm, kc, cluster)
	upgrades, err := upm.GetAvailableUpgrades()
	if err != nil {
		return "", err
	}
	upm.PrintAvailableUpgrades(upgrades)
	return "", nil
}

func Edit(ctx context.Context, name, version string) (*api.Cluster, error) {
	if name == "" {
		return nil, errors.New("missing cluster name")
	}

	cluster, err := Store(ctx).Clusters().Get(name)
	if err != nil {
		return nil, fmt.Errorf("cluster `%s` does not exist. Reason: %v", name, err)
	}
	if cluster.Spec.KubernetesVersion != version {
		cluster.Spec.KubernetesVersion = version
		cluster.Generation = time.Now().UnixNano()
		cluster.Status.Phase = api.ClusterUpgrading
	}
	return Store(ctx).Clusters().Update(cluster)
}
<|endoftext|>


<commit_before>package main

import (
	"sync"
	"time"
)

type Bucket struct {
	bLock    sync.RWMutex       // protect the session map
	sessions map[int64]*Session // map[user_id] ->  map[sub_id] -> server_id
	server   int
	cleaner  *Cleaner
}

// NewBucket new a bucket struct. store the subkey with im channel.
func NewBucket(session, server, cleaner int) *Bucket {
	b := new(Bucket)
	b.sessions = make(map[int64]*Session, session)
	b.server = server
	b.cleaner = NewCleaner(cleaner)
	go b.clean()
	return b
}

// Put put a channel according with user id.
func (b *Bucket) Put(userId int64, server int32) (seq int32) {
	var (
		s  *Session
		ok bool
	)
	b.bLock.Lock()
	if s, ok = b.sessions[userId]; !ok {
		s = NewSession(b.server)
		b.sessions[userId] = s
	}
	seq = s.Put(server)
	b.bLock.Unlock()
	return
}

func (b *Bucket) Get(userId int64) (seqs []int32, servers []int32) {
	var (
		s      *Session
		seq    int32
		server int32
		ok     bool
	)
	b.bLock.RLock()
	if s, ok = b.sessions[userId]; ok {
		seqs = make([]int32, 0, len(s.Servers()))
		servers = make([]int32, 0, len(s.Servers()))
		for seq, server = range s.Servers() {
			seqs = append(seqs, seq)
			servers = append(servers, server)
		}
	}
	b.bLock.RUnlock()
	return
}

func (b *Bucket) Count(userId int64) (count int) {
	b.bLock.RLock()
	if s, ok := b.sessions[userId]; ok {
		count = s.Size()
	}
	b.bLock.RUnlock()
	return
}

func (b *Bucket) del(userId int64) {
	var (
		s  *Session
		ok bool
	)
	if s, ok = b.sessions[userId]; ok {
		if s.Size() == 0 {
			delete(b.sessions, userId)
		}
	}
}

func (b *Bucket) Del(userId int64) {
	b.bLock.Lock()
	b.del(userId)
	b.bLock.Unlock()
}

// Del delete the channel by sub key.
func (b *Bucket) DelSession(userId int64, seq int32) (ok bool) {
	var (
		s     *Session
		empty bool
	)
	b.bLock.RLock()
	if s, ok = b.sessions[userId]; ok {
		// WARN:
		// delete(b.sessions, userId)
		// empty is a dirty data, we use here for try lru clean discard session.
		// when one user flapped connect & disconnect, this also can reduce
		// frequently new & free object, gc is slow!!!
		empty = s.Del(seq)
	}
	b.bLock.RUnlock()
	// lru
	if empty {
		b.cleaner.PushFront(userId, Conf.SessionExpire)
	}
	return
}

func (b *Bucket) clean() {
	var (
		i       int
		userIds []int64
	)
	for {
		userIds = b.cleaner.Clean()
		if len(userIds) != 0 {
			b.bLock.Lock()
			for i = 0; i < len(userIds); i++ {
				b.del(userIds[i])
			}
			b.bLock.Unlock()
		}
		time.Sleep(Conf.BucketCleanPeriod)
	}
}
<commit_msg>fix lru bug<commit_after>package main

import (
	"sync"
	"time"
)

type Bucket struct {
	bLock    sync.RWMutex       // protect the session map
	sessions map[int64]*Session // map[user_id] ->  map[sub_id] -> server_id
	server   int
	cleaner  *Cleaner
}

// NewBucket new a bucket struct. store the subkey with im channel.
func NewBucket(session, server, cleaner int) *Bucket {
	b := new(Bucket)
	b.sessions = make(map[int64]*Session, session)
	b.server = server
	b.cleaner = NewCleaner(cleaner)
	go b.clean()
	return b
}

// Put put a channel according with user id.
func (b *Bucket) Put(userId int64, server int32) (seq int32) {
	var (
		s  *Session
		ok bool
	)
	b.bLock.Lock()
	if s, ok = b.sessions[userId]; !ok {
		s = NewSession(b.server)
		b.sessions[userId] = s
	}
	seq = s.Put(server)
	b.bLock.Unlock()
	return
}

func (b *Bucket) Get(userId int64) (seqs []int32, servers []int32) {
	var (
		s      *Session
		seq    int32
		server int32
		ok     bool
	)
	b.bLock.RLock()
	if s, ok = b.sessions[userId]; ok {
		seqs = make([]int32, 0, len(s.Servers()))
		servers = make([]int32, 0, len(s.Servers()))
		for seq, server = range s.Servers() {
			seqs = append(seqs, seq)
			servers = append(servers, server)
		}
	}
	b.bLock.RUnlock()
	return
}

func (b *Bucket) Count(userId int64) (count int) {
	b.bLock.RLock()
	if s, ok := b.sessions[userId]; ok {
		count = s.Size()
	}
	b.bLock.RUnlock()
	return
}

func (b *Bucket) del(userId int64) {
	var (
		s  *Session
		ok bool
	)
	if s, ok = b.sessions[userId]; ok {
		if s.Size() == 0 {
			delete(b.sessions, userId)
		}
	}
}

func (b *Bucket) Del(userId int64) {
	b.bLock.Lock()
	b.del(userId)
	b.bLock.Unlock()
}

// Del delete the channel by sub key.
func (b *Bucket) DelSession(userId int64, seq int32) (ok bool) {
	var (
		s     *Session
		empty bool
	)
	b.bLock.RLock()
	if s, ok = b.sessions[userId]; ok {
		// WARN:
		// delete(b.sessions, userId)
		// empty is a dirty data, we use here for try lru clean discard session.
		// when one user flapped connect & disconnect, this also can reduce
		// frequently new & free object, gc is slow!!!
		empty = s.Del(seq)
	}
	b.bLock.RUnlock()
	// lru
	if empty {
		b.cleaner.PushFront(userId, Conf.SessionExpire)
	}
	return
}

func (b *Bucket) clean() {
	var (
		i       int
		userIds []int64
	)
	for {
		userIds = b.cleaner.Clean()
		if len(userIds) != 0 {
			b.bLock.Lock()
			for i = 0; i < len(userIds); i++ {
				b.del(userIds[i])
			}
			b.bLock.Unlock()
			continue
		}
		time.Sleep(Conf.BucketCleanPeriod)
	}
}
<|endoftext|>


<commit_before>/* Copyright 2017 The Bazel Authors. All rights reserved.

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

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

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

package rule

import (
	"sort"
)

// Platform represents a GOOS/GOARCH pair. When Platform is used to describe
// sources, dependencies, or flags, either OS or Arch may be empty.
//
// DEPRECATED: do not use outside language/go. This type is Go-specific
// and should be moved to the Go extension.
type Platform struct {
	OS, Arch string
}

// String returns OS, Arch, or "OS_Arch" if both are set. This must match
// the names of config_setting rules in @io_bazel_rules_go//go/platform.
func (p Platform) String() string {
	switch {
	case p.OS != "" && p.Arch != "":
		return p.OS + "_" + p.Arch
	case p.OS != "":
		return p.OS
	case p.Arch != "":
		return p.Arch
	default:
		return ""
	}
}

// KnownPlatforms is the set of target platforms that Go supports. Gazelle
// will generate multi-platform build files using these tags. rules_go and
// Bazel may not actually support all of these.
//
// DEPRECATED: do not use outside language/go.
var KnownPlatforms = []Platform{
	{"aix", "ppc64"},
	{"android", "386"},
	{"android", "amd64"},
	{"android", "arm"},
	{"android", "arm64"},
	{"darwin", "386"},
	{"darwin", "amd64"},
	{"darwin", "arm"},
	{"darwin", "arm64"},
	{"dragonfly", "amd64"},
	{"freebsd", "386"},
	{"freebsd", "amd64"},
	{"freebsd", "arm"},
	{"freebsd", "arm64"},
	{"illumos", "amd64"},
	{"ios", "386"},
	{"ios", "amd64"},
	{"ios", "arm"},
	{"ios", "arm64"},
	{"js", "wasm"},
	{"linux", "386"},
	{"linux", "amd64"},
	{"linux", "arm"},
	{"linux", "arm64"},
	{"linux", "mips"},
	{"linux", "mips64"},
	{"linux", "mips64le"},
	{"linux", "mipsle"},
	{"linux", "ppc64"},
	{"linux", "ppc64le"},
	{"linux", "riscv64"},
	{"linux", "s390x"},
	{"netbsd", "386"},
	{"netbsd", "amd64"},
	{"netbsd", "arm"},
	{"netbsd", "arm64"},
	{"openbsd", "386"},
	{"openbsd", "amd64"},
	{"openbsd", "arm"},
	{"openbsd", "arm64"},
	{"plan9", "386"},
	{"plan9", "amd64"},
	{"plan9", "arm"},
	{"solaris", "amd64"},
	{"windows", "386"},
	{"windows", "amd64"},
	{"windows", "arm"},
}

var OSAliases = map[string][]string{
	"android": []string{"linux"},
	"ios":     []string{"darwin"},
}

var (
	// KnownOSs is the sorted list of operating systems that Go supports.
	KnownOSs []string

	// KnownOSSet is the set of operating systems that Go supports.
	KnownOSSet map[string]bool

	// KnownArchs is the sorted list of architectures that Go supports.
	KnownArchs []string

	// KnownArchSet is the set of architectures that Go supports.
	KnownArchSet map[string]bool

	// KnownOSArchs is a map from OS to the archictures they run on.
	KnownOSArchs map[string][]string

	// KnownArchOSs is a map from architectures to that OSs that run on them.
	KnownArchOSs map[string][]string
)

func init() {
	KnownOSSet = make(map[string]bool)
	KnownArchSet = make(map[string]bool)
	KnownOSArchs = make(map[string][]string)
	KnownArchOSs = make(map[string][]string)
	for _, p := range KnownPlatforms {
		KnownOSSet[p.OS] = true
		KnownArchSet[p.Arch] = true
		KnownOSArchs[p.OS] = append(KnownOSArchs[p.OS], p.Arch)
		KnownArchOSs[p.Arch] = append(KnownArchOSs[p.Arch], p.OS)
	}
	KnownOSs = make([]string, 0, len(KnownOSSet))
	KnownArchs = make([]string, 0, len(KnownArchSet))
	for os := range KnownOSSet {
		KnownOSs = append(KnownOSs, os)
	}
	for arch := range KnownArchSet {
		KnownArchs = append(KnownArchs, arch)
	}
	sort.Strings(KnownOSs)
	sort.Strings(KnownArchs)
}
<commit_msg>fix(platforms): remove iOS 386 and arm targets (#1168)<commit_after>/* Copyright 2017 The Bazel Authors. All rights reserved.

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

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

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

package rule

import (
	"sort"
)

// Platform represents a GOOS/GOARCH pair. When Platform is used to describe
// sources, dependencies, or flags, either OS or Arch may be empty.
//
// DEPRECATED: do not use outside language/go. This type is Go-specific
// and should be moved to the Go extension.
type Platform struct {
	OS, Arch string
}

// String returns OS, Arch, or "OS_Arch" if both are set. This must match
// the names of config_setting rules in @io_bazel_rules_go//go/platform.
func (p Platform) String() string {
	switch {
	case p.OS != "" && p.Arch != "":
		return p.OS + "_" + p.Arch
	case p.OS != "":
		return p.OS
	case p.Arch != "":
		return p.Arch
	default:
		return ""
	}
}

// KnownPlatforms is the set of target platforms that Go supports. Gazelle
// will generate multi-platform build files using these tags. rules_go and
// Bazel may not actually support all of these.
//
// DEPRECATED: do not use outside language/go.
var KnownPlatforms = []Platform{
	{"aix", "ppc64"},
	{"android", "386"},
	{"android", "amd64"},
	{"android", "arm"},
	{"android", "arm64"},
	{"darwin", "386"},
	{"darwin", "amd64"},
	{"darwin", "arm"},
	{"darwin", "arm64"},
	{"dragonfly", "amd64"},
	{"freebsd", "386"},
	{"freebsd", "amd64"},
	{"freebsd", "arm"},
	{"freebsd", "arm64"},
	{"illumos", "amd64"},
	{"ios", "amd64"},
	{"ios", "arm64"},
	{"js", "wasm"},
	{"linux", "386"},
	{"linux", "amd64"},
	{"linux", "arm"},
	{"linux", "arm64"},
	{"linux", "mips"},
	{"linux", "mips64"},
	{"linux", "mips64le"},
	{"linux", "mipsle"},
	{"linux", "ppc64"},
	{"linux", "ppc64le"},
	{"linux", "riscv64"},
	{"linux", "s390x"},
	{"netbsd", "386"},
	{"netbsd", "amd64"},
	{"netbsd", "arm"},
	{"netbsd", "arm64"},
	{"openbsd", "386"},
	{"openbsd", "amd64"},
	{"openbsd", "arm"},
	{"openbsd", "arm64"},
	{"plan9", "386"},
	{"plan9", "amd64"},
	{"plan9", "arm"},
	{"solaris", "amd64"},
	{"windows", "386"},
	{"windows", "amd64"},
	{"windows", "arm"},
}

var OSAliases = map[string][]string{
	"android": []string{"linux"},
	"ios":     []string{"darwin"},
}

var (
	// KnownOSs is the sorted list of operating systems that Go supports.
	KnownOSs []string

	// KnownOSSet is the set of operating systems that Go supports.
	KnownOSSet map[string]bool

	// KnownArchs is the sorted list of architectures that Go supports.
	KnownArchs []string

	// KnownArchSet is the set of architectures that Go supports.
	KnownArchSet map[string]bool

	// KnownOSArchs is a map from OS to the archictures they run on.
	KnownOSArchs map[string][]string

	// KnownArchOSs is a map from architectures to that OSs that run on them.
	KnownArchOSs map[string][]string
)

func init() {
	KnownOSSet = make(map[string]bool)
	KnownArchSet = make(map[string]bool)
	KnownOSArchs = make(map[string][]string)
	KnownArchOSs = make(map[string][]string)
	for _, p := range KnownPlatforms {
		KnownOSSet[p.OS] = true
		KnownArchSet[p.Arch] = true
		KnownOSArchs[p.OS] = append(KnownOSArchs[p.OS], p.Arch)
		KnownArchOSs[p.Arch] = append(KnownArchOSs[p.Arch], p.OS)
	}
	KnownOSs = make([]string, 0, len(KnownOSSet))
	KnownArchs = make([]string, 0, len(KnownArchSet))
	for os := range KnownOSSet {
		KnownOSs = append(KnownOSs, os)
	}
	for arch := range KnownArchSet {
		KnownArchs = append(KnownArchs, arch)
	}
	sort.Strings(KnownOSs)
	sort.Strings(KnownArchs)
}
<|endoftext|>