<commit_before>package helpers

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"strings"
	"syscall"
	"time"

	"github.com/DATA-DOG/godog"
	"github.com/DATA-DOG/godog/gherkin"
	"github.com/Originate/exosphere/src"
	"github.com/Originate/exosphere/src/util"
	execplus "github.com/Originate/go-execplus"
	"github.com/pkg/errors"
)

var childCmdPlus *execplus.CmdPlus
var childOutput string
var appDir string
var templateDir string

func waitWithTimeout(cmdPlus *execplus.CmdPlus, duration time.Duration) error {
	done := make(chan error)
	go func() { done <- cmdPlus.Wait() }()
	select {
	case err := <-done:
		return err
	case <-time.After(duration):
		return fmt.Errorf("Timed out after %v, command did not exit. Full output:\n%s", duration, cmdPlus.GetOutput())
	}
}

// SharedFeatureContext defines the festure context shared between the sub commands
// nolint: gocyclo
func SharedFeatureContext(s *godog.Suite) {
	s.BeforeScenario(func(arg1 interface{}) {
		var err error
		appDir, err = ioutil.TempDir("", "")
		if err != nil {
			panic(err)
		}
		templateDir, err = ioutil.TempDir("", "")
		if err != nil {
			panic(err)
		}
	})

	s.AfterScenario(func(arg1 interface{}, arg2 error) {
		if childCmdPlus != nil {
			if err := childCmdPlus.Kill(); err != nil {
				panic(err)
			}
			childCmdPlus = nil
		}
		if err := cleanApp(appDir); err != nil {
			panic(err)
		}
	})

	// Application Setup

	s.Step(`^I am in the root directory of a non-exosphere application$`, func() error {
		return nil
	})

	s.Step(`^I am in the root directory of an empty application called "([^"]*)"$`, func(appName string) error {
		var err error
		appDir, err = createEmptyApp(appName)
		return err
	})

	s.Step(`^I am in the root directory of the "([^"]*)" example application$`, func(name string) error {
		return CheckoutApp(appDir, name)
	})

	s.Step(`^it doesn\'t run any tests$`, func() error {
		expectedText := "is not an exosphere application"
		if childCmdPlus != nil {
			return childCmdPlus.WaitForText(expectedText, time.Minute)
		}
		return validateTextContains(childOutput, expectedText)
	})

	s.Step(`^I am in the directory of "([^"]*)" application containing a "([^"]*)" service$`, func(appName, serviceRole string) error {
		return CheckoutApp(appDir, "test-app")
	})

	s.Step(`^my application has the templates:$`, func(table *gherkin.DataTable) error {
		for _, row := range table.Rows[1:] {
			templateName, gitURL := row.Cells[0].Value, row.Cells[1].Value
			command := []string{"exo", "template", "add", templateName, gitURL}
			if len(row.Cells) == 3 {
				gitTag := row.Cells[2].Value
				command = append(command, gitTag)
			}
			if _, err := util.Run(appDir, command...); err != nil {
				return errors.Wrap(err, fmt.Sprintf("Failed to creates the template %s:%s\n", appDir, err))
			}
		}
		return nil
	})

	s.Step(`^my (?:application|workspace) contains the empty directory "([^"]*)"`, func(directory string) error {
		dirPath := path.Join(appDir, directory)
		isDir, err := util.DoesDirectoryExist(dirPath)
		if err != nil {
			return err
		}
		if !isDir {
			return fmt.Errorf("%s is a not a directory", dirPath)
		}
		fileInfos, err := ioutil.ReadDir(dirPath)
		if err != nil {
			return err
		}
		if len(fileInfos) != 0 {
			fileNames := []string{}
			for _, fileInfo := range fileInfos {
				fileNames = append(fileNames, fileInfo.Name())
			}
			return fmt.Errorf("%s is a not a an empty directory. Contains: %s", dirPath, strings.Join(fileNames, ", "))
		}
		return nil
	})

	// Running / Starting a command

	s.Step(`^running "([^"]*)" in the terminal$`, func(command string) error {
		var err error
		childOutput, err = util.Run(appDir, command)
		if err != nil {
			return errors.Wrap(err, fmt.Sprintf("Command errored with output: %s", childOutput))
		}
		return nil
	})

	s.Step(`^running "([^"]*)" in my application directory$`, func(command string) error {
		var err error
		childOutput, err = util.Run(appDir, command)
		if err != nil {
			return errors.Wrap(err, fmt.Sprintf("Command errored with output: %s", childOutput))
		}
		return nil
	})

	s.Step(`^running "([^"]*)" in the "([^"]*)" directory$`, func(command, dirName string) error {
		var err error
		childOutput, err = util.Run(path.Join(appDir, dirName), command)
		if err != nil {
			return errors.Wrap(err, fmt.Sprintf("Command errored with output: %s", childOutput))
		}
		return nil
	})

	s.Step(`^starting "([^"]*)" in the terminal$`, func(command string) error {
		commandWords, err := util.ParseCommand(command)
		if err != nil {
			return err
		}
		childCmdPlus = execplus.NewCmdPlus(commandWords...)
		childCmdPlus.SetDir(appDir)
		return childCmdPlus.Start()
	})

	s.Step(`^starting "([^"]*)" in my application directory$`, func(command string) error {
		commandWords, err := util.ParseCommand(command)
		if err != nil {
			return err
		}
		childCmdPlus = execplus.NewCmdPlus(commandWords...)
		childCmdPlus.SetDir(appDir)
		return childCmdPlus.Start()
	})

	s.Step(`^starting "([^"]*)" in the "([^"]*)" directory$`, func(command, dirName string) error {
		commandWords, err := util.ParseCommand(command)
		if err != nil {
			return err
		}
		childCmdPlus = execplus.NewCmdPlus(commandWords...)
		childCmdPlus.SetDir(path.Join(appDir, dirName))
		return childCmdPlus.Start()
	})
	// Entering user input

	s.Step(`^entering into the wizard:$`, func(table *gherkin.DataTable) error {
		for _, row := range table.Rows[1:] {
			if err := enterInput(row); err != nil {
				return errors.Wrap(err, fmt.Sprintf("Failed to enter %s into the wizard", row.Cells[1].Value))
			}
		}
		return nil
	})

	// Verifying output

	s.Step(`^it prints "([^"]*)" in the terminal$`, func(text string) error {
		if childCmdPlus != nil {
			return childCmdPlus.WaitForText(text, time.Minute)
		}
		return validateTextContains(childOutput, text)
	})

	s.Step(`^it prints the current version in the terminal$`, func() error {
		text := fmt.Sprintf("Exosphere v%s", src.Version)
		if childCmdPlus != nil {
			return childCmdPlus.WaitForText(text, time.Minute)
		}
		return validateTextContains(childOutput, text)
	})

	s.Step(`^it does not print "([^"]*)" in the terminal$`, func(text string) error {
		if childCmdPlus != nil {
			if err := validateTextContains(childCmdPlus.GetOutput(), text); err == nil {
				return fmt.Errorf("Expected the process to not print: %s", text)
			}
			return nil
		}
		return validateTextContains(childOutput, text)
	})

	s.Step(`^I see:$`, func(expectedText *gherkin.DocString) error {
		if childCmdPlus != nil {
			return childCmdPlus.WaitForText(expectedText.Content, time.Second*2)
		}
		return validateTextContains(childOutput, expectedText.Content)
	})

	s.Step(`^I eventually see "([^"]*)" in the terminal$`, func(expectedText string) error {
		return childCmdPlus.WaitForText(expectedText, time.Second)
	})

	s.Step(`^I eventually see:$`, func(expectedText *gherkin.DocString) error {
		return childCmdPlus.WaitForText(expectedText.Content, time.Minute)
	})

	s.Step(`^(it exits|waiting until the process ends)$`, func() error {
		return waitWithTimeout(childCmdPlus, time.Minute)
	})

	s.Step(`^I stop all running processes$`, func() error {
		if childCmdPlus != nil {
			err := childCmdPlus.Cmd.Process.Signal(os.Interrupt)
			if err != nil {
				return err
			}
			err = waitWithTimeout(childCmdPlus, time.Minute)
			if err != nil {
				fmt.Println("Command did not exit after 1m (TODO: fix this)")
				return childCmdPlus.Kill()
			}
		}
		return nil
	})

	s.Step(`^it exits with code (\d+)$`, func(expectedExitCode int) error {
		actualExitCode := 0
		if err := waitWithTimeout(childCmdPlus, time.Minute); err != nil {
			if exiterr, ok := err.(*exec.ExitError); ok {
				if status, ok := exiterr.Sys().(syscall.WaitStatus); ok {
					actualExitCode = status.ExitStatus()
				} else {
					return fmt.Errorf("Unable to parse Status object: %v", err)
				}
			} else {
				return fmt.Errorf("cmd.Wait: %v", err)
			}
		}
		if actualExitCode != expectedExitCode {
			return fmt.Errorf("Exited with code %d instead of %d. Output:\n%s", actualExitCode, expectedExitCode, childCmdPlus.GetOutput())
		}
		return nil
	})

	s.Step(`^my workspace contains the file "([^"]*)" with content:$`, func(fileName string, expectedContent *gherkin.DocString) error {
		bytes, err := ioutil.ReadFile(path.Join(appDir, fileName))
		if err != nil {
			return errors.Wrap(err, fmt.Sprintf("Failed to read %s", fileName))
		}
		return validateTextContains(strings.TrimSpace(string(bytes)), strings.TrimSpace(expectedContent.Content))
	})

	s.Step(`^my workspace contains the files:$`, func(table *gherkin.DataTable) error {
		for _, row := range table.Rows[1:] {
			filename := row.Cells[0].Value
			hasFile, err := util.DoesFileExist(path.Join(appDir, filename))
			if err != nil {
				return err
			}
			if !hasFile {
				return fmt.Errorf("Expected %s to exist", filename)
			}
		}
		return nil
	})

	s.Step(`^my application now contains the file "([^"]*)" with the content:$`, func(fileName string, expectedContent *gherkin.DocString) error {
		bytes, err := ioutil.ReadFile(path.Join(appDir, fileName))
		if err != nil {
			return errors.Wrap(err, fmt.Sprintf("Failed to read %s", fileName))
		}
		return validateTextContains(strings.TrimSpace(string(bytes)), strings.TrimSpace(expectedContent.Content))
	})
}
<commit_msg>less verbose test errors (#813)<commit_after>package helpers

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"strings"
	"syscall"
	"time"

	"github.com/DATA-DOG/godog"
	"github.com/DATA-DOG/godog/gherkin"
	"github.com/Originate/exosphere/src"
	"github.com/Originate/exosphere/src/util"
	execplus "github.com/Originate/go-execplus"
	"github.com/pkg/errors"
)

var childCmdPlus *execplus.CmdPlus
var childOutput string
var appDir string
var templateDir string

func waitWithTimeout(cmdPlus *execplus.CmdPlus, duration time.Duration) error {
	done := make(chan error)
	go func() { done <- cmdPlus.Wait() }()
	select {
	case err := <-done:
		return err
	case <-time.After(duration):
		return fmt.Errorf("Timed out after %v, command did not exit. Full output:\n%s", duration, cmdPlus.GetOutput())
	}
}

// SharedFeatureContext defines the festure context shared between the sub commands
// nolint: gocyclo
func SharedFeatureContext(s *godog.Suite) {
	s.BeforeScenario(func(arg1 interface{}) {
		var err error
		appDir, err = ioutil.TempDir("", "")
		if err != nil {
			panic(err)
		}
		templateDir, err = ioutil.TempDir("", "")
		if err != nil {
			panic(err)
		}
	})

	s.AfterScenario(func(arg1 interface{}, arg2 error) {
		if childCmdPlus != nil {
			if err := childCmdPlus.Kill(); err != nil {
				panic(err)
			}
			childCmdPlus = nil
		}
		if err := cleanApp(appDir); err != nil {
			panic(err)
		}
	})

	// Application Setup

	s.Step(`^I am in the root directory of a non-exosphere application$`, func() error {
		return nil
	})

	s.Step(`^I am in the root directory of an empty application called "([^"]*)"$`, func(appName string) error {
		var err error
		appDir, err = createEmptyApp(appName)
		return err
	})

	s.Step(`^I am in the root directory of the "([^"]*)" example application$`, func(name string) error {
		return CheckoutApp(appDir, name)
	})

	s.Step(`^it doesn\'t run any tests$`, func() error {
		expectedText := "is not an exosphere application"
		if childCmdPlus != nil {
			return childCmdPlus.WaitForText(expectedText, time.Minute)
		}
		return validateTextContains(childOutput, expectedText)
	})

	s.Step(`^I am in the directory of "([^"]*)" application containing a "([^"]*)" service$`, func(appName, serviceRole string) error {
		return CheckoutApp(appDir, "test-app")
	})

	s.Step(`^my application has the templates:$`, func(table *gherkin.DataTable) error {
		for _, row := range table.Rows[1:] {
			templateName, gitURL := row.Cells[0].Value, row.Cells[1].Value
			command := []string{"exo", "template", "add", templateName, gitURL}
			if len(row.Cells) == 3 {
				gitTag := row.Cells[2].Value
				command = append(command, gitTag)
			}
			if _, err := util.Run(appDir, command...); err != nil {
				return errors.Wrap(err, fmt.Sprintf("Failed to creates the template %s:%s\n", appDir, err))
			}
		}
		return nil
	})

	s.Step(`^my (?:application|workspace) contains the empty directory "([^"]*)"`, func(directory string) error {
		dirPath := path.Join(appDir, directory)
		isDir, err := util.DoesDirectoryExist(dirPath)
		if err != nil {
			return err
		}
		if !isDir {
			return fmt.Errorf("%s is a not a directory", dirPath)
		}
		fileInfos, err := ioutil.ReadDir(dirPath)
		if err != nil {
			return err
		}
		if len(fileInfos) != 0 {
			fileNames := []string{}
			for _, fileInfo := range fileInfos {
				fileNames = append(fileNames, fileInfo.Name())
			}
			return fmt.Errorf("%s is a not a an empty directory. Contains: %s", dirPath, strings.Join(fileNames, ", "))
		}
		return nil
	})

	// Running / Starting a command

	s.Step(`^running "([^"]*)" in the terminal$`, func(command string) error {
		var err error
		childOutput, err = util.Run(appDir, command)
		return err
	})

	s.Step(`^running "([^"]*)" in my application directory$`, func(command string) error {
		var err error
		childOutput, err = util.Run(appDir, command)
		return err
	})

	s.Step(`^running "([^"]*)" in the "([^"]*)" directory$`, func(command, dirName string) error {
		var err error
		childOutput, err = util.Run(path.Join(appDir, dirName), command)
		return err
	})

	s.Step(`^starting "([^"]*)" in the terminal$`, func(command string) error {
		commandWords, err := util.ParseCommand(command)
		if err != nil {
			return err
		}
		childCmdPlus = execplus.NewCmdPlus(commandWords...)
		childCmdPlus.SetDir(appDir)
		return childCmdPlus.Start()
	})

	s.Step(`^starting "([^"]*)" in my application directory$`, func(command string) error {
		commandWords, err := util.ParseCommand(command)
		if err != nil {
			return err
		}
		childCmdPlus = execplus.NewCmdPlus(commandWords...)
		childCmdPlus.SetDir(appDir)
		return childCmdPlus.Start()
	})

	s.Step(`^starting "([^"]*)" in the "([^"]*)" directory$`, func(command, dirName string) error {
		commandWords, err := util.ParseCommand(command)
		if err != nil {
			return err
		}
		childCmdPlus = execplus.NewCmdPlus(commandWords...)
		childCmdPlus.SetDir(path.Join(appDir, dirName))
		return childCmdPlus.Start()
	})
	// Entering user input

	s.Step(`^entering into the wizard:$`, func(table *gherkin.DataTable) error {
		for _, row := range table.Rows[1:] {
			if err := enterInput(row); err != nil {
				return errors.Wrap(err, fmt.Sprintf("Failed to enter %s into the wizard", row.Cells[1].Value))
			}
		}
		return nil
	})

	// Verifying output

	s.Step(`^it prints "([^"]*)" in the terminal$`, func(text string) error {
		if childCmdPlus != nil {
			return childCmdPlus.WaitForText(text, time.Minute)
		}
		return validateTextContains(childOutput, text)
	})

	s.Step(`^it prints the current version in the terminal$`, func() error {
		text := fmt.Sprintf("Exosphere v%s", src.Version)
		if childCmdPlus != nil {
			return childCmdPlus.WaitForText(text, time.Minute)
		}
		return validateTextContains(childOutput, text)
	})

	s.Step(`^it does not print "([^"]*)" in the terminal$`, func(text string) error {
		if childCmdPlus != nil {
			if err := validateTextContains(childCmdPlus.GetOutput(), text); err == nil {
				return fmt.Errorf("Expected the process to not print: %s", text)
			}
			return nil
		}
		return validateTextContains(childOutput, text)
	})

	s.Step(`^I see:$`, func(expectedText *gherkin.DocString) error {
		if childCmdPlus != nil {
			return childCmdPlus.WaitForText(expectedText.Content, time.Second*2)
		}
		return validateTextContains(childOutput, expectedText.Content)
	})

	s.Step(`^I eventually see "([^"]*)" in the terminal$`, func(expectedText string) error {
		return childCmdPlus.WaitForText(expectedText, time.Second)
	})

	s.Step(`^I eventually see:$`, func(expectedText *gherkin.DocString) error {
		return childCmdPlus.WaitForText(expectedText.Content, time.Minute)
	})

	s.Step(`^(it exits|waiting until the process ends)$`, func() error {
		return waitWithTimeout(childCmdPlus, time.Minute)
	})

	s.Step(`^I stop all running processes$`, func() error {
		if childCmdPlus != nil {
			err := childCmdPlus.Cmd.Process.Signal(os.Interrupt)
			if err != nil {
				return err
			}
			err = waitWithTimeout(childCmdPlus, time.Minute)
			if err != nil {
				fmt.Println("Command did not exit after 1m (TODO: fix this)")
				return childCmdPlus.Kill()
			}
		}
		return nil
	})

	s.Step(`^it exits with code (\d+)$`, func(expectedExitCode int) error {
		actualExitCode := 0
		if err := waitWithTimeout(childCmdPlus, time.Minute); err != nil {
			if exiterr, ok := err.(*exec.ExitError); ok {
				if status, ok := exiterr.Sys().(syscall.WaitStatus); ok {
					actualExitCode = status.ExitStatus()
				} else {
					return fmt.Errorf("Unable to parse Status object: %v", err)
				}
			} else {
				return fmt.Errorf("cmd.Wait: %v", err)
			}
		}
		if actualExitCode != expectedExitCode {
			return fmt.Errorf("Exited with code %d instead of %d. Output:\n%s", actualExitCode, expectedExitCode, childCmdPlus.GetOutput())
		}
		return nil
	})

	s.Step(`^my workspace contains the file "([^"]*)" with content:$`, func(fileName string, expectedContent *gherkin.DocString) error {
		bytes, err := ioutil.ReadFile(path.Join(appDir, fileName))
		if err != nil {
			return errors.Wrap(err, fmt.Sprintf("Failed to read %s", fileName))
		}
		return validateTextContains(strings.TrimSpace(string(bytes)), strings.TrimSpace(expectedContent.Content))
	})

	s.Step(`^my workspace contains the files:$`, func(table *gherkin.DataTable) error {
		for _, row := range table.Rows[1:] {
			filename := row.Cells[0].Value
			hasFile, err := util.DoesFileExist(path.Join(appDir, filename))
			if err != nil {
				return err
			}
			if !hasFile {
				return fmt.Errorf("Expected %s to exist", filename)
			}
		}
		return nil
	})

	s.Step(`^my application now contains the file "([^"]*)" with the content:$`, func(fileName string, expectedContent *gherkin.DocString) error {
		bytes, err := ioutil.ReadFile(path.Join(appDir, fileName))
		if err != nil {
			return errors.Wrap(err, fmt.Sprintf("Failed to read %s", fileName))
		}
		return validateTextContains(strings.TrimSpace(string(bytes)), strings.TrimSpace(expectedContent.Content))
	})
}
<|endoftext|>


<commit_before>package command

import (
	"bufio"
	"fmt"
	"github.com/digitalocean/godo"
	"github.com/masayukioguni/godo-cli/config"
	"github.com/mitchellh/cli"
	"os"
)

type AuthorizeCommand struct {
	Ui     cli.Ui
	Config *config.Config
	Client *godo.Client
}

func (c *AuthorizeCommand) Help() string {
	return ""
}

func (c *AuthorizeCommand) ask(text string, defaultText string) string {
	fmt.Printf("%s", text)
	var scanner = bufio.NewScanner(os.Stdin)
	scanner.Scan()

	input := scanner.Text()
	if input != "" {
		return input
	}
	return defaultText

}

func (c *AuthorizeCommand) Run(args []string) int {
	apikey := c.ask("Entser your API Token:", "Input api ooken")
	fmt.Println(`Defaults can be changed at any time in your ~/.godo-cli/config.yaml configuration file.`)

	c.Config.Authentication.APIKey = apikey

	savePath, err := config.GetConfigPath()

	if err != nil {
		fmt.Errorf("Error GetConfigPath %s", err)
		return 1
	}

	err = config.SaveConfig(savePath, c.Config)

	if err != nil {
		fmt.Errorf("Error SaveConfig %s", err)
		return 1
	}

	fmt.Println("Authentication with DigitalOcean was successful!")
	return 0
}

func (c *AuthorizeCommand) Synopsis() string {
	return fmt.Sprintf("Authorize a DigitalOcean account with godo-cli")
}
<commit_msg>fix typos<commit_after>package command

import (
	"bufio"
	"fmt"
	"github.com/digitalocean/godo"
	"github.com/masayukioguni/godo-cli/config"
	"github.com/mitchellh/cli"
	"os"
)

type AuthorizeCommand struct {
	Ui     cli.Ui
	Config *config.Config
	Client *godo.Client
}

func (c *AuthorizeCommand) Help() string {
	return ""
}

func (c *AuthorizeCommand) ask(text string, defaultText string) string {
	fmt.Printf("%s", text)
	var scanner = bufio.NewScanner(os.Stdin)
	scanner.Scan()

	input := scanner.Text()
	if input != "" {
		return input
	}
	return defaultText

}

func (c *AuthorizeCommand) Run(args []string) int {
	apikey := c.ask("Enter your API Token:", "Input api token")
	fmt.Println(`Defaults can be changed at any time in your ~/.godo-cli/config.yaml configuration file.`)

	c.Config.Authentication.APIKey = apikey

	savePath, err := config.GetConfigPath()

	if err != nil {
		fmt.Errorf("Error GetConfigPath %s", err)
		return 1
	}

	err = config.SaveConfig(savePath, c.Config)

	if err != nil {
		fmt.Errorf("Error SaveConfig %s", err)
		return 1
	}

	fmt.Println("Authentication with DigitalOcean was successful!")
	return 0
}

func (c *AuthorizeCommand) Synopsis() string {
	return fmt.Sprintf("Authorize a DigitalOcean account with godo-cli")
}
<|endoftext|>


<commit_before>package command

import (
	"encoding/json"
	"fmt"

	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"
	"text/tabwriter"
	"time"

	"github.com/rancher/go-rancher/client"
)

type Status struct {
	client.Resource
	ReplicaCounter  int64 `json:"replicacounter"`
	RevisionCounter int64 `json:"revisioncounter"`
}

type VolumeStats struct {
	client.Resource
	RevisionCounter int64         `json:"RevisionCounter"`
	ReplicaCounter  int64         `json:"ReplicaCounter"`
	SCSIIOCount     map[int]int64 `json:"SCSIIOCount"`

	ReadIOPS            string `json:"ReadIOPS"`
	TotalReadTime       string `json:"TotalReadTime"`
	TotalReadBlockCount string `json:"TotalReadBlockCount"`

	WriteIOPS            string `json:"WriteIOPS"`
	TotalWriteTime       string `json:"TotalWriteTime"`
	TotalWriteBlockCount string `json:"TotatWriteBlockCount"`

	SectorSize        string `json:"SectorSize"`
	UsedBlocks        string `json:"UsedBlocks"`
	UsedLogicalBlocks string `json:"UsedLogicalBlocks"`
}

// VsmStatsCommand is a command implementation struct
type VsmStatsCommand struct {
	Meta
	address     string
	host        string
	length      int
	replica_ips string
	Json        string
}

// ReplicaClient is Client structure
type ReplicaClient struct {
	address    string
	syncAgent  string
	host       string
	httpClient *http.Client
}

type ControllerClient struct {
	address    string
	host       string
	httpClient *http.Client
}

type StatsArr struct {
	IQN    string `json:"Iqn"`
	Volume string `json:"Volume"`
	Portal string `json:"Portal"`
	Size   string `json:"Size"`

	ReadIOPS  int64 `json:"ReadIOPS"`
	WriteIOPS int64 `json:"WriteIOPS"`

	ReadThroughput  float64 `json:"ReadThroughput"`
	WriteThroughput float64 `json:"WriteThroughput"`

	ReadLatency  float64 `json:"ReadLatency"`
	WriteLatency float64 `json:"WriteLatency"`

	AvgReadBlockSize  int64 `json:"AvgReadBlockSize"`
	AvgWriteBlockSize int64 `json:"AvgWriteBlockSize"`

	SectorSize  float64 `json:"SectorSize"`
	ActualUsed  float64 `json:"ActualUsed"`
	LogicalSize float64 `json:"LogicalSize"`
}

// Help shows helpText for a particular CLI command
func (c *VsmStatsCommand) Help() string {
	helpText := `
	Usage: maya vsm-stats <vsm-name>

  Display VSM Stats.

`
	return strings.TrimSpace(helpText)
}

// Synopsis shows short information related to CLI command
func (c *VsmStatsCommand) Synopsis() string {
	return "Display VSM Stats"
}

// Run holds the flag values for CLI subcommands
func (c *VsmStatsCommand) Run(args []string) int {

	var (
		err, err1, err2 error
		status          Status
		stats1          VolumeStats
		stats2          VolumeStats
		statusArray     []string
		ReadLatency     int64
		WriteLatency    int64

		AvgReadBlockCountPS  int64
		AvgWriteBlockCountPS int64
	)

	flags := c.Meta.FlagSet("vsm-stats", FlagSetClient)
	flags.Usage = func() { c.Ui.Output(c.Help()) }
	flags.StringVar(&c.Json, "json", "", "")

	if err := flags.Parse(args); err != nil {
		return 1
	}
	args = flags.Args()
	if len(args) < 1 {
		c.Ui.Error(c.Help())
		return 1
	}

	annotations, err := GetVolAnnotations(args[0])
	if err != nil || annotations == nil {
		return -1
	}
	if annotations.ControllerStatus != "Running" {
		fmt.Println("Volume not reachable")
		return -1
	}
	replicas := strings.Split(annotations.Replicas, ",")
	for _, replica := range replicas {
		err, errCode1 := GetStatus(replica+":9502", &status)
		if err != nil {
			if errCode1 == 500 || strings.Contains(err.Error(), "EOF") {
				statusArray = append(statusArray, fmt.Sprintf("%-15s %-12s%-10s", replica, "Waiting", "Unknown"))

			} else {
				statusArray = append(statusArray, fmt.Sprintf("%-15s %-12s%-10s", replica, "Offline", "Unknown"))
			}
		} else {
			statusArray = append(statusArray, fmt.Sprintf("%-15s %-10s  %d", replica, "Online", status.RevisionCounter))
		}
	}

	//Get VolumeStats
	err1, _ = GetVolumeStats(annotations.ClusterIP+":9501", &stats1)
	time.Sleep(1 * time.Second)
	err2, _ = GetVolumeStats(annotations.ClusterIP+":9501", &stats2)

	if (err1 != nil) || (err2 != nil) {
		fmt.Println("Volume not reachable")
	}

	ReadIOPSi, _ := strconv.ParseInt(stats1.ReadIOPS, 10, 64)
	ReadIOPSf, _ := strconv.ParseInt(stats2.ReadIOPS, 10, 64)
	ReadIOPSPS := ReadIOPSf - ReadIOPSi

	ReadTimePSi, _ := strconv.ParseInt(stats1.TotalReadTime, 10, 64)
	ReadTimePSf, _ := strconv.ParseInt(stats2.TotalReadTime, 10, 64)
	ReadTimePS := ReadTimePSf - ReadTimePSi

	ReadBlockCountPSi, _ := strconv.ParseInt(stats1.TotalReadBlockCount, 10, 64)
	ReadBlockCountPSf, _ := strconv.ParseInt(stats2.TotalReadBlockCount, 10, 64)
	ReadBlockCountPS := ReadBlockCountPSf - ReadBlockCountPSi

	RThroughput := ReadBlockCountPS
	if ReadIOPSPS != 0 {
		ReadLatency = ReadTimePS / ReadIOPSPS
		AvgReadBlockCountPS = ReadBlockCountPS / ReadIOPSPS
	} else {
		ReadLatency = 0
		AvgReadBlockCountPS = 0
	}

	WriteIOPSi, _ := strconv.ParseInt(stats1.WriteIOPS, 10, 64)
	WriteIOPSf, _ := strconv.ParseInt(stats2.WriteIOPS, 10, 64)
	WriteIOPSPS := WriteIOPSf - WriteIOPSi

	WriteTimePSi, _ := strconv.ParseInt(stats1.TotalWriteTime, 10, 64)
	WriteTimePSf, _ := strconv.ParseInt(stats2.TotalWriteTime, 10, 64)
	WriteTimePS := WriteTimePSf - WriteTimePSi

	WriteBlockCountPSi, _ := strconv.ParseInt(stats1.TotalWriteBlockCount, 10, 64)
	WriteBlockCountPSf, _ := strconv.ParseInt(stats2.TotalWriteBlockCount, 10, 64)
	WriteBlockCountPS := WriteBlockCountPSf - WriteBlockCountPSi

	WThroughput := WriteBlockCountPS
	if WriteIOPSPS != 0 {
		WriteLatency = WriteTimePS / WriteIOPSPS
		AvgWriteBlockCountPS = WriteBlockCountPS / WriteIOPSPS
	} else {
		WriteLatency = 0
		AvgWriteBlockCountPS = 0
	}

	SectorSizePS, _ := strconv.ParseFloat(stats2.SectorSize, 64)
	//	fmt.Println("hi", stats1.SectorSize, stats2.SectorSize)
	SectorSizePS = SectorSizePS / 1048576
	LogicalSize, _ := strconv.ParseFloat(stats2.UsedBlocks, 64)
	LogicalSizePS := LogicalSize * SectorSizePS

	ActualUsed, _ := strconv.ParseFloat(stats2.UsedLogicalBlocks, 64)
	ActualUsedPS := ActualUsed * SectorSizePS

	fmt.Println("------------------------------------")
	// json formatting and showing default output
	if c.Json == "json" {

		stat1 := StatsArr{

			IQN:    annotations.Iqn,
			Volume: args[0],
			Portal: annotations.TargetPortal,
			Size:   annotations.VolSize,

			ReadIOPS:  ReadIOPSPS,
			WriteIOPS: WriteIOPSPS,

			ReadThroughput:  float64(RThroughput) / 104857,
			WriteThroughput: float64(WThroughput) / 104857,

			ReadLatency:  float64(ReadLatency) / 1000000,
			WriteLatency: float64(WriteLatency) / 1000000,

			AvgReadBlockSize:  AvgReadBlockCountPS / 1024,
			AvgWriteBlockSize: AvgWriteBlockCountPS / 1024,

			SectorSize:  SectorSizePS,
			ActualUsed:  ActualUsedPS,
			LogicalSize: LogicalSizePS,
		}

		data, err := json.Marshal(stat1)

		if err != nil {

			panic(err)
		}

		os.Stdout.Write(data)

		fmt.Println("\n------------------------------------")

	} else {
		fmt.Printf("%7s: %-48s\n", "IQN", annotations.Iqn)
		fmt.Printf("%7s: %-16s\n", "Volume", args[0])
		fmt.Printf("%7s: %-15s\n", "Portal", annotations.TargetPortal)
		fmt.Printf("%7s: %-6s\n\n", "Size", annotations.VolSize)
		fmt.Printf("%s         %s      %s\n", "Replica", "Status", "DataUpdateIndex")

		for i, _ := range statusArray {
			fmt.Printf("%s\n", statusArray[i])
		}
		fmt.Println("------------------------------------")

		// Printing in tabular form
		w := tabwriter.NewWriter(os.Stdout, 0, 0, 3, ' ', tabwriter.AlignRight|tabwriter.Debug)
		fmt.Println("-------- Performance Stats ---------\n")
		fmt.Fprintf(w, "r/s\tw/s\tr(MB/s)\tw(MB/s)\trLat(ms)\twLat(ms)\t\n")
		fmt.Fprintf(w, "%d\t%d\t%.3f\t%.3f\t%.3f\t%.3f\t\n", ReadIOPSPS, WriteIOPSPS, float64(RThroughput)/1048576, float64(WThroughput)/1048576, float64(ReadLatency)/1000000, float64(WriteLatency)/1000000)
		w.Flush()

		x := tabwriter.NewWriter(os.Stdout, 0, 0, 3, ' ', tabwriter.AlignRight|tabwriter.Debug)
		fmt.Println("\n---------- Capacity Stats ----------\n")
		fmt.Fprintf(x, "Logical(GB)\tUsed(GB)\t\n")
		fmt.Fprintf(x, "%f\t%f\t\n", LogicalSizePS, ActualUsedPS)
		x.Flush()
	}
	return 0
}

// NewReplicaClient create the new replica client
func NewReplicaClient(address string) (*ReplicaClient, error) {
	if strings.HasPrefix(address, "tcp://") {
		address = address[6:]
	}

	if !strings.HasPrefix(address, "http") {
		address = "http://" + address
	}

	if !strings.HasSuffix(address, "/v1") {
		address += "/v1"
	}

	u, err := url.Parse(address)
	if err != nil {
		return nil, err
	}

	parts := strings.Split(u.Host, ":")
	if len(parts) < 2 {
		return nil, fmt.Errorf("Invalid address %s, must have a port in it", address)
	}
	port, err := strconv.Atoi(parts[1])
	if err != nil {
		return nil, err
	}
	syncAgent := strings.Replace(address, fmt.Sprintf(":%d", port), fmt.Sprintf(":%d", port+2), -1)

	timeout := time.Duration(2 * time.Second)
	client := &http.Client{
		Timeout: timeout,
	}

	return &ReplicaClient{
		host:       parts[0],
		address:    address,
		syncAgent:  syncAgent,
		httpClient: client,
	}, nil
}

// GetStatus will return json response and statusCode
func GetStatus(address string, obj interface{}) (error, int) {
	replica, err := NewReplicaClient(address)
	if err != nil {
		return err, -1
	}
	url := replica.address + "/stats"
	resp, err := replica.httpClient.Get(url)
	if resp != nil {
		if resp.StatusCode == 500 {
			return err, 500
		} else if resp.StatusCode == 503 {
			return err, 503
		}
	} else {
		return err, -1
	}
	if err != nil {
		return err, -1
	}
	defer resp.Body.Close()

	return json.NewDecoder(resp.Body).Decode(obj), 0
}

// NewControllerClient create the new replica client
func NewControllerClient(address string) (*ControllerClient, error) {
	if strings.HasPrefix(address, "tcp://") {
		address = address[6:]
	}

	if !strings.HasPrefix(address, "http") {
		address = "http://" + address
	}

	if !strings.HasSuffix(address, "/v1") {
		address += "/v1"
	}

	u, err := url.Parse(address)
	if err != nil {
		return nil, err
	}

	parts := strings.Split(u.Host, ":")
	if len(parts) < 2 {
		return nil, fmt.Errorf("Invalid address %s, must have a port in it", address)
	}

	timeout := time.Duration(2 * time.Second)
	client := &http.Client{
		Timeout: timeout,
	}

	return &ControllerClient{
		host:       parts[0],
		address:    address,
		httpClient: client,
	}, nil
}

// GetStatus will return json response and statusCode
func GetVolumeStats(address string, obj interface{}) (error, int) {
	controller, err := NewControllerClient(address)
	if err != nil {
		return err, -1
	}
	url := controller.address + "/stats"
	resp, err := controller.httpClient.Get(url)
	if resp != nil {
		if resp.StatusCode == 500 {
			return err, 500
		} else if resp.StatusCode == 503 {
			return err, 503
		}
	} else {
		return err, -1
	}
	if err != nil {
		return err, -1
	}
	defer resp.Body.Close()
	rc, _ := json.NewDecoder(resp.Body).Decode(obj), 0
	return rc, 0
}
<commit_msg>Update changes<commit_after>package command

import (
	"encoding/json"
	"fmt"

	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"
	"text/tabwriter"
	"time"

	"github.com/rancher/go-rancher/client"
)

type Status struct {
	client.Resource
	ReplicaCounter  int64 `json:"replicacounter"`
	RevisionCounter int64 `json:"revisioncounter"`
}

type VolumeStats struct {
	client.Resource
	RevisionCounter int64         `json:"RevisionCounter"`
	ReplicaCounter  int64         `json:"ReplicaCounter"`
	SCSIIOCount     map[int]int64 `json:"SCSIIOCount"`

	ReadIOPS            string `json:"ReadIOPS"`
	TotalReadTime       string `json:"TotalReadTime"`
	TotalReadBlockCount string `json:"TotalReadBlockCount"`

	WriteIOPS            string `json:"WriteIOPS"`
	TotalWriteTime       string `json:"TotalWriteTime"`
	TotalWriteBlockCount string `json:"TotatWriteBlockCount"`

	SectorSize        string `json:"SectorSize"`
	UsedBlocks        string `json:"UsedBlocks"`
	UsedLogicalBlocks string `json:"UsedLogicalBlocks"`
}

// VsmStatsCommand is a command implementation struct
type VsmStatsCommand struct {
	Meta
	address     string
	host        string
	length      int
	replica_ips string
	Json        string
}

// ReplicaClient is Client structure
type ReplicaClient struct {
	address    string
	syncAgent  string
	host       string
	httpClient *http.Client
}

type ControllerClient struct {
	address    string
	host       string
	httpClient *http.Client
}

type StatsArr struct {
	IQN    string `json:"Iqn"`
	Volume string `json:"Volume"`
	Portal string `json:"Portal"`
	Size   string `json:"Size"`

	ReadIOPS  int64 `json:"ReadIOPS"`
	WriteIOPS int64 `json:"WriteIOPS"`

	ReadThroughput  float64 `json:"ReadThroughput"`
	WriteThroughput float64 `json:"WriteThroughput"`

	ReadLatency  float64 `json:"ReadLatency"`
	WriteLatency float64 `json:"WriteLatency"`

	AvgReadBlockSize  int64 `json:"AvgReadBlockSize"`
	AvgWriteBlockSize int64 `json:"AvgWriteBlockSize"`

	SectorSize  float64 `json:"SectorSize"`
	ActualUsed  float64 `json:"ActualUsed"`
	LogicalSize float64 `json:"LogicalSize"`
}

// Help shows helpText for a particular CLI command
func (c *VsmStatsCommand) Help() string {
	helpText := `
	Usage: maya vsm-stats <vsm-name>

  Display VSM Stats.

`
	return strings.TrimSpace(helpText)
}

// Synopsis shows short information related to CLI command
func (c *VsmStatsCommand) Synopsis() string {
	return "Display VSM Stats"
}

// Run holds the flag values for CLI subcommands
func (c *VsmStatsCommand) Run(args []string) int {

	var (
		err, err1, err2 error
		status          Status
		stats1          VolumeStats
		stats2          VolumeStats
		statusArray     []string
		ReadLatency     int64
		WriteLatency    int64

		AvgReadBlockCountPS  int64
		AvgWriteBlockCountPS int64
	)

	flags := c.Meta.FlagSet("vsm-stats", FlagSetClient)
	flags.Usage = func() { c.Ui.Output(c.Help()) }
	flags.StringVar(&c.Json, "json", "", "")

	if err := flags.Parse(args); err != nil {
		return 1
	}
	args = flags.Args()
	if len(args) < 1 {
		c.Ui.Error(c.Help())
		return 1
	}

	annotations, err := GetVolAnnotations(args[0])
	if err != nil || annotations == nil {
		return -1
	}
	if annotations.ControllerStatus != "Running" {
		fmt.Println("Volume not reachable")
		return -1
	}
	replicas := strings.Split(annotations.Replicas, ",")
	for _, replica := range replicas {
		err, errCode1 := GetStatus(replica+":9502", &status)
		if err != nil {
			if errCode1 == 500 || strings.Contains(err.Error(), "EOF") {
				statusArray = append(statusArray, fmt.Sprintf("%-15s %-12s%-10s", replica, "Waiting", "Unknown"))

			} else {
				statusArray = append(statusArray, fmt.Sprintf("%-15s %-12s%-10s", replica, "Offline", "Unknown"))
			}
		} else {
			statusArray = append(statusArray, fmt.Sprintf("%-15s %-10s  %d", replica, "Online", status.RevisionCounter))
		}
	}

	//Get VolumeStats
	err1, _ = GetVolumeStats(annotations.ClusterIP+":9501", &stats1)
	time.Sleep(1 * time.Second)
	err2, _ = GetVolumeStats(annotations.ClusterIP+":9501", &stats2)

	if (err1 != nil) || (err2 != nil) {
		fmt.Println("Volume not reachable")
	}

	ReadIOPSi, _ := strconv.ParseInt(stats1.ReadIOPS, 10, 64)
	ReadIOPSf, _ := strconv.ParseInt(stats2.ReadIOPS, 10, 64)
	ReadIOPSPS := ReadIOPSf - ReadIOPSi

	ReadTimePSi, _ := strconv.ParseInt(stats1.TotalReadTime, 10, 64)
	ReadTimePSf, _ := strconv.ParseInt(stats2.TotalReadTime, 10, 64)
	ReadTimePS := ReadTimePSf - ReadTimePSi

	ReadBlockCountPSi, _ := strconv.ParseInt(stats1.TotalReadBlockCount, 10, 64)
	ReadBlockCountPSf, _ := strconv.ParseInt(stats2.TotalReadBlockCount, 10, 64)
	ReadBlockCountPS := ReadBlockCountPSf - ReadBlockCountPSi

	RThroughput := ReadBlockCountPS
	if ReadIOPSPS != 0 {
		ReadLatency = ReadTimePS / ReadIOPSPS
		AvgReadBlockCountPS = ReadBlockCountPS / ReadIOPSPS
	} else {
		ReadLatency = 0
		AvgReadBlockCountPS = 0
	}

	WriteIOPSi, _ := strconv.ParseInt(stats1.WriteIOPS, 10, 64)
	WriteIOPSf, _ := strconv.ParseInt(stats2.WriteIOPS, 10, 64)
	WriteIOPSPS := WriteIOPSf - WriteIOPSi

	WriteTimePSi, _ := strconv.ParseInt(stats1.TotalWriteTime, 10, 64)
	WriteTimePSf, _ := strconv.ParseInt(stats2.TotalWriteTime, 10, 64)
	WriteTimePS := WriteTimePSf - WriteTimePSi

	WriteBlockCountPSi, _ := strconv.ParseInt(stats1.TotalWriteBlockCount, 10, 64)
	WriteBlockCountPSf, _ := strconv.ParseInt(stats2.TotalWriteBlockCount, 10, 64)
	WriteBlockCountPS := WriteBlockCountPSf - WriteBlockCountPSi

	WThroughput := WriteBlockCountPS
	if WriteIOPSPS != 0 {
		WriteLatency = WriteTimePS / WriteIOPSPS
		AvgWriteBlockCountPS = WriteBlockCountPS / WriteIOPSPS
	} else {
		WriteLatency = 0
		AvgWriteBlockCountPS = 0
	}

	SectorSizePS, _ := strconv.ParseFloat(stats2.SectorSize, 64)
	//	fmt.Println("hi", stats1.SectorSize, stats2.SectorSize)
	SectorSizePS = SectorSizePS / 1048576
	LogicalSize, _ := strconv.ParseFloat(stats2.UsedBlocks, 64)
	LogicalSizePS := LogicalSize * SectorSizePS

	ActualUsed, _ := strconv.ParseFloat(stats2.UsedLogicalBlocks, 64)
	ActualUsedPS := ActualUsed * SectorSizePS

	fmt.Println("------------------------------------")
	// json formatting and showing default output
	if c.Json == "json" {

		stat1 := StatsArr{

			IQN:    annotations.Iqn,
			Volume: args[0],
			Portal: annotations.TargetPortal,
			Size:   annotations.VolSize,

			ReadIOPS:  ReadIOPSPS,
			WriteIOPS: WriteIOPSPS,

			ReadThroughput:  float64(RThroughput) / 104857,
			WriteThroughput: float64(WThroughput) / 104857,

			ReadLatency:  float64(ReadLatency) / 1000000,
			WriteLatency: float64(WriteLatency) / 1000000,

			AvgReadBlockSize:  AvgReadBlockCountPS / 1024,
			AvgWriteBlockSize: AvgWriteBlockCountPS / 1024,

			SectorSize:  SectorSizePS,
			ActualUsed:  ActualUsedPS,
			LogicalSize: LogicalSizePS,
		}

		data, err := json.Marshal(stat1)

		if err != nil {

			panic(err)
		}

		os.Stdout.Write(data)

		fmt.Println("\n------------------------------------")

	} else {
		fmt.Printf("%7s: %-48s\n", "IQN", annotations.Iqn)
		fmt.Printf("%7s: %-16s\n", "Volume", args[0])
		fmt.Printf("%7s: %-15s\n", "Portal", annotations.TargetPortal)
		fmt.Printf("%7s: %-6s\n\n", "Size", annotations.VolSize)
		fmt.Printf("%s         %s      %s\n", "Replica", "Status", "DataUpdateIndex")

		for i, _ := range statusArray {
			fmt.Printf("%s\n", statusArray[i])
		}
		fmt.Println("------------------------------------")

		// Printing in tabular form
		w := tabwriter.NewWriter(os.Stdout, 0, 0, 3, ' ', tabwriter.AlignRight|tabwriter.Debug)
		fmt.Println("-------- Performance Stats ---------\n")
		fmt.Fprintf(w, "r/s\tw/s\tr(MB/s)\tw(MB/s)\trLat(ms)\twLat(ms)\t\n")
		fmt.Fprintf(w, "%d\t%d\t%.3f\t%.3f\t%.3f\t%.3f\t\n", ReadIOPSPS, WriteIOPSPS, float64(RThroughput)/1048576, float64(WThroughput)/1048576, float64(ReadLatency)/1000000, float64(WriteLatency)/1000000)
		w.Flush()

		x := tabwriter.NewWriter(os.Stdout, 0, 0, 3, ' ', tabwriter.AlignRight|tabwriter.Debug)
		fmt.Println("\n---------- Capacity Stats ----------\n")
		fmt.Fprintf(x, "Logical(GB)\tUsed(GB)\t\n")
		fmt.Fprintf(x, "%f\t%f\t\n", LogicalSizePS, ActualUsedPS)
		x.Flush()
	}
	return 0
}

// NewReplicaClient create the new replica client
func NewReplicaClient(address string) (*ReplicaClient, error) {
	if strings.HasPrefix(address, "tcp://") {
		address = address[6:]
	}

	if !strings.HasPrefix(address, "http") {
		address = "http://" + address
	}

	if !strings.HasSuffix(address, "/v1") {
		address += "/v1"
	}

	u, err := url.Parse(address)
	if err != nil {
		return nil, err
	}

	parts := strings.Split(u.Host, ":")
	if len(parts) < 2 {
		return nil, fmt.Errorf("Invalid address %s, must have a port in it", address)
	}
	port, err := strconv.Atoi(parts[1])
	if err != nil {
		return nil, err
	}
	syncAgent := strings.Replace(address, fmt.Sprintf(":%d", port), fmt.Sprintf(":%d", port+2), -1)

	timeout := time.Duration(2 * time.Second)
	client := &http.Client{
		Timeout: timeout,
	}

	return &ReplicaClient{
		host:       parts[0],
		address:    address,
		syncAgent:  syncAgent,
		httpClient: client,
	}, nil
}

// GetStatus will return json response and statusCode
func GetStatus(address string, obj interface{}) (error, int) {
	replica, err := NewReplicaClient(address)
	if err != nil {
		return err, -1
	}
	url := replica.address + "/stats"
	resp, err := replica.httpClient.Get(url)
	if resp != nil {
		if resp.StatusCode == 500 {
			return err, 500
		} else if resp.StatusCode == 503 {
			return err, 503
		}
	} else {
		return err, -1
	}
	if err != nil {
		return err, -1
	}
	defer resp.Body.Close()

	return json.NewDecoder(resp.Body).Decode(obj), 0
}

// NewControllerClient create the new replica client
func NewControllerClient(address string) (*ControllerClient, error) {
	if strings.HasPrefix(address, "tcp://") {
		address = address[6:]
	}

	if !strings.HasPrefix(address, "http") {
		address = "http://" + address
	}

	if !strings.HasSuffix(address, "/v1") {
		address += "/v1"
	}

	u, err := url.Parse(address)
	if err != nil {
		return nil, err
	}

	parts := strings.Split(u.Host, ":")
	if len(parts) < 2 {
		return nil, fmt.Errorf("Invalid address %s, must have a port in it", address)
	}

	timeout := time.Duration(2 * time.Second)
	client := &http.Client{
		Timeout: timeout,
	}

	return &ControllerClient{
		host:       parts[0],
		address:    address,
		httpClient: client,
	}, nil
}

// GetStatus will return json response and statusCode
func GetVolumeStats(address string, obj interface{}) (error, int) {
	controller, err := NewControllerClient(address)
	if err != nil {
		return err, -1
	}
	url := controller.address + "/stats"
	resp, err := controller.httpClient.Get(url)
	if resp != nil {
		if resp.StatusCode == 500 {
			return err, 500
		} else if resp.StatusCode == 503 {
			return err, 503
		}
	} else {
		return err, -1
	}
	if err != nil {
		return err, -1
	}
	defer resp.Body.Close()
	rc := json.NewDecoder(resp.Body).Decode(obj)
	return rc, 0
}
<|endoftext|>


<commit_before>package commands

import (
	"fmt"
	"github.com/jingweno/gh/git"
	"github.com/jingweno/gh/github"
	"github.com/jingweno/gh/utils"
	"regexp"
)

var cmdCheckout = &Command{
	Run:          checkout,
	GitExtension: true,
	Usage:        "checkout PULLREQ-URL [BRANCH]",
	Short:        "Switch the active branch to another branch",
}

/**
  $ gh checkout https://github.com/jingweno/gh/pull/73
  # > git remote add -f -t feature git://github:com/foo/gh.git
  # > git checkout --track -B foo-feature foo/feature

  $ gh checkout https://github.com/jingweno/gh/pull/73 custom-branch-name
**/
func checkout(command *Command, args []string) {
	var err error
	if len(args) > 0 {
		args, err = transformCheckoutArgs(args)
		utils.Fatal(err)
	}

	err = git.ExecCheckout(args)
	utils.Check(err)
}

func transformCheckoutArgs(args []string) ([]string, error) {
	id := parsePullRequestId(args[0])
	if id != "" {
		newArgs, url := removeItem(args, 0)
		gh := github.New()
		pullRequest, err := gh.PullRequest(id)
		if err != nil {
			return nil, err
		}

		user := pullRequest.User.Login
		branch := pullRequest.Head.Ref
		if pullRequest.Head.Repo.ID == 0 {
			return nil, fmt.Errorf("%s's fork is not available anymore", user)
		}

		remotes, err := git.Remotes()
		if err != nil {
			return nil, err
		}

		var remoteExists bool
		for _, r := range remotes {
			if r.Name == user {
				remoteExists = true
				break
			}
		}

		if remoteExists {
			err = git.Spawn("remote", "set-branches", "--add", user, branch)
			if err != nil {
				return nil, err
			}
			remoteURL := fmt.Sprintf("+refs/heads/%s:refs/remotes/%s/%s", branch, user, branch)
			git.Spawn("fetch", user, remoteURL)
			if err != nil {
				return nil, err
			}
		} else {
			project, err := github.ParseProjectFromURL(url)
			if err != nil {
				return nil, err
			}

			sshURL := project.GitURL("", user, pullRequest.Head.Repo.Private)
			err = git.Spawn("remote", "add", "-f", "-t", branch, user, sshURL)
			if err != nil {
				return nil, err
			}
		}

		trackedBranch := fmt.Sprintf("%s/%s", user, branch)
		var newBranchName string
		if len(newArgs) > 0 {
			newArgs, newBranchName = removeItem(newArgs, 0)
		} else {
			newBranchName = fmt.Sprintf("%s-%s", user, branch)
		}

		newArgs = append(newArgs, "--track", "-B", newBranchName, trackedBranch)

		return newArgs, nil
	}

	return args, nil
}

func parsePullRequestId(url string) string {
	pullURLRegex := regexp.MustCompile("https://github\\.com/.+/.+/pull/(\\d+)")
	if pullURLRegex.MatchString(url) {
		return pullURLRegex.FindStringSubmatch(url)[1]
	}

	return ""
}
<commit_msg>Extract to smaller methods<commit_after>package commands

import (
	"fmt"
	"github.com/jingweno/gh/git"
	"github.com/jingweno/gh/github"
	"github.com/jingweno/gh/utils"
	"regexp"
)

var cmdCheckout = &Command{
	Run:          checkout,
	GitExtension: true,
	Usage:        "checkout PULLREQ-URL [BRANCH]",
	Short:        "Switch the active branch to another branch",
}

/**
  $ gh checkout https://github.com/jingweno/gh/pull/73
  # > git remote add -f -t feature git://github:com/foo/gh.git
  # > git checkout --track -B foo-feature foo/feature

  $ gh checkout https://github.com/jingweno/gh/pull/73 custom-branch-name
**/
func checkout(command *Command, args []string) {
	var err error
	if len(args) > 0 {
		args, err = transformCheckoutArgs(args)
		utils.Fatal(err)
	}

	err = git.ExecCheckout(args)
	utils.Check(err)
}

func transformCheckoutArgs(args []string) ([]string, error) {
	id := parsePullRequestId(args[0])
	if id != "" {
		newArgs, url := removeItem(args, 0)
		gh := github.New()
		pullRequest, err := gh.PullRequest(id)
		if err != nil {
			return nil, err
		}

		user := pullRequest.User.Login
		branch := pullRequest.Head.Ref
		if pullRequest.Head.Repo.ID == 0 {
			return nil, fmt.Errorf("%s's fork is not available anymore", user)
		}

		remoteExists, err := checkIfRemoteExists(user)
		if err != nil {
			return nil, err
		}

		if remoteExists {
			err = updateExistingRemote(user, branch)
		} else {
			err = addRmote(user, branch, url, pullRequest.Head.Repo.Private)
		}
		if err != nil {
			return nil, err
		}

		var newBranchName string
		if len(newArgs) > 0 {
			newArgs, newBranchName = removeItem(newArgs, 0)
		} else {
			newBranchName = fmt.Sprintf("%s-%s", user, branch)
		}
		trackedBranch := fmt.Sprintf("%s/%s", user, branch)

		newArgs = append(newArgs, "--track", "-B", newBranchName, trackedBranch)

		return newArgs, nil
	}

	return args, nil
}

func parsePullRequestId(url string) string {
	pullURLRegex := regexp.MustCompile("https://github\\.com/.+/.+/pull/(\\d+)")
	if pullURLRegex.MatchString(url) {
		return pullURLRegex.FindStringSubmatch(url)[1]
	}

	return ""
}

func checkIfRemoteExists(remote string) (bool, error) {
	remotes, err := git.Remotes()
	if err != nil {
		return false, err
	}

	for _, r := range remotes {
		if r.Name == remote {
			return true, nil
		}
	}

	return false, nil
}

func updateExistingRemote(user, branch string) error {
	err := git.Spawn("remote", "set-branches", "--add", user, branch)
	if err != nil {
		return err
	}
	remoteURL := fmt.Sprintf("+refs/heads/%s:refs/remotes/%s/%s", branch, user, branch)

	return git.Spawn("fetch", user, remoteURL)
}

func addRmote(user, branch, url string, isPrivate bool) error {
	project, err := github.ParseProjectFromURL(url)
	if err != nil {
		return err
	}

	sshURL := project.GitURL("", user, isPrivate)

	return git.Spawn("remote", "add", "-f", "-t", branch, user, sshURL)
}
<|endoftext|>


<commit_before>package commands

import (
	"crypto/tls"
	"errors"
	"flag"
	"fmt"
	"github.com/orchardup/orchard/authenticator"
	"github.com/orchardup/orchard/proxy"
	"github.com/orchardup/orchard/tlsconfig"
	"github.com/orchardup/orchard/utils"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"path"
	"strings"
	"syscall"
	"text/tabwriter"
)

type Command struct {
	Run       func(cmd *Command, args []string) error
	UsageLine string
	Short     string
	Long      string
	Flag      flag.FlagSet
}

func (c *Command) Name() string {
	name := c.UsageLine
	i := strings.Index(name, " ")
	if i >= 0 {
		name = name[:i]
	}
	return name
}

func (c *Command) Usage() {
	fmt.Fprintf(os.Stderr, "Usage: %s\n\n", c.UsageLine)
	fmt.Fprintf(os.Stderr, "%s\n", strings.TrimSpace(c.Long))
	os.Exit(2)
}

func (c *Command) UsageError(format string, args ...interface{}) error {
	fmt.Fprintf(os.Stderr, format, args...)
	fmt.Fprint(os.Stderr, "\n\n")
	c.Usage()
	return fmt.Errorf(format, args...)
}

var All = []*Command{
	Hosts,
	Start,
	Stop,
	Docker,
	Proxy,
}

func init() {
	Hosts.Run = RunHosts
	Start.Run = RunStart
	Stop.Run = RunStop
	Docker.Run = RunDocker
	Proxy.Run = RunProxy
}

var Hosts = &Command{
	UsageLine: "hosts",
	Short:     "List all hosts",
	Long:      "List all hosts",
}

var Start = &Command{
	UsageLine: "start [-m MEMORY] [NAME]",
	Short:     "Start a host",
	Long: fmt.Sprintf(`Start a host.

You can optionally specify a name for the host - if not, it will be
named 'default', and 'orchard docker' commands will use it automatically.

You can also specify how much RAM the host should have with -m.
Valid amounts are %s.`, validSizes),
}

var flStartSize = Start.Flag.String("m", "512M", "")
var validSizes = "512M, 1G, 2G, 4G and 8G"

var Stop = &Command{
	UsageLine: "stop [NAME]",
	Short:     "Stop a host",
	Long: `Stop a host.

You can optionally specify which host to stop - if you don't, the default
host (named 'default') will be stopped.`,
}

var Docker = &Command{
	UsageLine: "docker [-H HOST] [COMMAND...]",
	Short:     "Run a Docker command against a host",
	Long: `Run a Docker command against a host.

Wraps the 'docker' command-line tool - see the Docker website for reference:

    http://docs.docker.io/en/latest/reference/commandline/

You can optionally specify a host by name - if you don't, the default host
will be used.`,
}

var flDockerHost = Docker.Flag.String("H", "", "")

var Proxy = &Command{
	UsageLine: "proxy [-H HOST]",
	Short:     "Start a local proxy to a host's Docker daemon",
	Long: `Start a local proxy to a host's Docker daemon.

Prints out a URL to pass to the 'docker' command, e.g.

    $ orchard proxy
    Started proxy at unix:///tmp/orchard-12345/orchard.sock

    $ docker -H unix:///tmp/orchard-12345/orchard.sock run ubuntu echo hello world
    hello world
`,
}

var flProxyHost = Proxy.Flag.String("H", "", "")

func RunHosts(cmd *Command, args []string) error {
	if len(args) > 0 {
		return cmd.UsageError("`orchard hosts` doesn't expect any arguments, but got: %s", strings.Join(args, " "))
	}

	httpClient, err := authenticator.Authenticate()
	if err != nil {
		return err
	}

	hosts, err := httpClient.GetHosts()
	if err != nil {
		return err
	}

	writer := tabwriter.NewWriter(os.Stdout, 20, 1, 3, ' ', 0)
	fmt.Fprintln(writer, "NAME\tSIZE\tIP")
	for _, host := range hosts {
		fmt.Fprintf(writer, "%s\t%s\t%s\n", host.Name, utils.HumanSize(host.Size*1024*1024), host.IPAddress)
	}
	writer.Flush()

	return nil
}

func RunStart(cmd *Command, args []string) error {
	if len(args) > 1 {
		return cmd.UsageError("`orchard start` expects at most 1 argument, but got more: %s", strings.Join(args[1:], " "))
	}

	httpClient, err := authenticator.Authenticate()
	if err != nil {
		return err
	}

	hostName, humanName := GetHostName(args)
	humanName = utils.Capitalize(humanName)

	size, sizeString := GetHostSize()
	if size == -1 {
		fmt.Fprintf(os.Stderr, "Sorry, %q isn't a size we support.\nValid sizes are %s.\n", sizeString, validSizes)
		return nil
	}

	host, err := httpClient.CreateHost(hostName, size)
	if err != nil {
		// HACK. api.go should decode JSON and return a specific type of error for this case.
		if strings.Contains(err.Error(), "already exists") {
			fmt.Fprintf(os.Stderr, "%s is already running.\nYou can create additional hosts with `orchard start NAME`.\n", humanName)
			return nil
		}
		if strings.Contains(err.Error(), "Invalid value") {
			fmt.Fprintf(os.Stderr, "Sorry, '%s' isn't a valid host name.\nHost names can only contain lowercase letters, numbers and underscores.\n", hostName)
			return nil
		}
		if strings.Contains(err.Error(), "Unsupported size") {
			fmt.Fprintf(os.Stderr, "Sorry, %q isn't a size we support.\nValid sizes are %s.\n", sizeString, validSizes)
			return nil
		}

		return err
	}
	fmt.Fprintf(os.Stderr, "%s running at %s\n", humanName, host.IPAddress)

	return nil
}

func RunStop(cmd *Command, args []string) error {
	if len(args) > 1 {
		return cmd.UsageError("`orchard stop` expects at most 1 argument, but got more: %s", strings.Join(args[1:], " "))
	}

	hostName, humanName := GetHostName(args)

	var confirm string
	fmt.Printf("Going to stop and delete %s. All data on it will be lost.\n", humanName)
	fmt.Print("Are you sure you're ready? [yN] ")
	fmt.Scanln(&confirm)

	if strings.ToLower(confirm) != "y" {
		return nil
	}

	httpClient, err := authenticator.Authenticate()
	if err != nil {
		return err
	}

	err = httpClient.DeleteHost(hostName)
	if err != nil {
		// HACK. api.go should decode JSON and return a specific type of error for this case.
		if strings.Contains(err.Error(), "Not found") {
			fmt.Fprintf(os.Stderr, "%s doesn't seem to be running.\nYou can view your running hosts with `orchard hosts`.\n", utils.Capitalize(humanName))
			return nil
		}

		return err
	}
	fmt.Fprintf(os.Stderr, "Stopped %s\n", humanName)

	return nil
}

func RunDocker(cmd *Command, args []string) error {
	return WithDockerProxy(func(socketPath string) error {
		err := CallDocker(args, []string{"DOCKER_HOST=unix://" + socketPath})
		if err != nil {
			return fmt.Errorf("Docker exited with error")
		}
		return nil
	})
}

func RunProxy(cmd *Command, args []string) error {
	if len(args) > 0 {
		return cmd.UsageError("`orchard proxy` doesn't expect any arguments, but got: %s", strings.Join(args, " "))
	}

	return WithDockerProxy(func(socketPath string) error {
		fmt.Fprintln(os.Stderr, "Started proxy at unix://"+socketPath)

		c := make(chan os.Signal)
		signal.Notify(c, syscall.SIGINT, syscall.SIGKILL)
		<-c

		fmt.Fprintln(os.Stderr, "\nStopping proxy")
		return nil
	})
}

func WithDockerProxy(callback func(string) error) error {
	hostName := "default"
	if *flDockerHost != "" {
		hostName = *flDockerHost
	}

	dirname, err := ioutil.TempDir("/tmp", "orchard-")
	if err != nil {
		return fmt.Errorf("Error creating temporary directory: %s\n", err)
	}
	defer os.RemoveAll(dirname)
	socketPath := path.Join(dirname, "orchard.sock")

	p, err := MakeProxy(socketPath, hostName)
	if err != nil {
		return fmt.Errorf("Error starting proxy: %v\n", err)
	}

	go p.Start()
	defer p.Stop()

	if err := <-p.ErrorChannel; err != nil {
		return fmt.Errorf("Error starting proxy: %v\n", err)
	}

	if err := callback(socketPath); err != nil {
		return err
	}

	return nil
}

func MakeProxy(socketPath string, hostName string) (*proxy.Proxy, error) {
	httpClient, err := authenticator.Authenticate()
	if err != nil {
		return nil, err
	}

	host, err := httpClient.GetHost(hostName)
	if err != nil {
		return nil, err
	}
	destination := host.IPAddress + ":4243"

	certData := []byte(host.ClientCert)
	keyData := []byte(host.ClientKey)
	config, err := tlsconfig.GetTLSConfig(certData, keyData)
	if err != nil {
		return nil, err
	}

	return proxy.New(
		func() (net.Listener, error) { return net.Listen("unix", socketPath) },
		func() (net.Conn, error) { return tls.Dial("tcp", destination, config) },
	), nil
}

func CallDocker(args []string, env []string) error {
	dockerPath := GetDockerPath()
	if dockerPath == "" {
		return errors.New("Can't find `docker` executable in $PATH.\nYou might need to install it: http://docs.docker.io/en/latest/installation/#installation-list")
	}

	cmd := exec.Command(dockerPath, args...)
	cmd.Env = env
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}

func GetDockerPath() string {
	for _, dir := range strings.Split(os.Getenv("PATH"), ":") {
		dockerPath := path.Join(dir, "docker")
		_, err := os.Stat(dockerPath)
		if err == nil {
			return dockerPath
		}
	}
	return ""
}

func GetHostName(args []string) (string, string) {
	hostName := "default"
	humanName := "default host"

	if len(args) > 0 {
		hostName = args[0]
		humanName = fmt.Sprintf("host '%s'", hostName)
	}

	return hostName, humanName
}

func GetHostSize() (int, string) {
	sizeString := *flStartSize

	bytes, err := utils.RAMInBytes(sizeString)
	if err != nil {
		return -1, sizeString
	}

	megs := bytes / (1024 * 1024)
	if megs < 1 {
		return -1, sizeString
	}

	return int(megs), sizeString
}
<commit_msg>Just print the usage line when arguments are invalid<commit_after>package commands

import (
	"crypto/tls"
	"errors"
	"flag"
	"fmt"
	"github.com/orchardup/orchard/authenticator"
	"github.com/orchardup/orchard/proxy"
	"github.com/orchardup/orchard/tlsconfig"
	"github.com/orchardup/orchard/utils"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"os/signal"
	"path"
	"strings"
	"syscall"
	"text/tabwriter"
)

type Command struct {
	Run       func(cmd *Command, args []string) error
	UsageLine string
	Short     string
	Long      string
	Flag      flag.FlagSet
}

func (c *Command) Name() string {
	name := c.UsageLine
	i := strings.Index(name, " ")
	if i >= 0 {
		name = name[:i]
	}
	return name
}

func (c *Command) Usage() {
	fmt.Fprintf(os.Stderr, "Usage: %s\n\n", c.UsageLine)
	fmt.Fprintf(os.Stderr, "%s\n", strings.TrimSpace(c.Long))
	os.Exit(2)
}

func (c *Command) UsageError(format string, args ...interface{}) error {
	fmt.Fprintf(os.Stderr, format, args...)
	fmt.Fprintf(os.Stderr, "\nUsage: %s\n", c.UsageLine)
	os.Exit(2)
	return fmt.Errorf(format, args...)
}

var All = []*Command{
	Hosts,
	Start,
	Stop,
	Docker,
	Proxy,
}

func init() {
	Hosts.Run = RunHosts
	Start.Run = RunStart
	Stop.Run = RunStop
	Docker.Run = RunDocker
	Proxy.Run = RunProxy
}

var Hosts = &Command{
	UsageLine: "hosts",
	Short:     "List all hosts",
	Long:      "List all hosts",
}

var Start = &Command{
	UsageLine: "start [-m MEMORY] [NAME]",
	Short:     "Start a host",
	Long: fmt.Sprintf(`Start a host.

You can optionally specify a name for the host - if not, it will be
named 'default', and 'orchard docker' commands will use it automatically.

You can also specify how much RAM the host should have with -m.
Valid amounts are %s.`, validSizes),
}

var flStartSize = Start.Flag.String("m", "512M", "")
var validSizes = "512M, 1G, 2G, 4G and 8G"

var Stop = &Command{
	UsageLine: "stop [NAME]",
	Short:     "Stop a host",
	Long: `Stop a host.

You can optionally specify which host to stop - if you don't, the default
host (named 'default') will be stopped.`,
}

var Docker = &Command{
	UsageLine: "docker [-H HOST] [COMMAND...]",
	Short:     "Run a Docker command against a host",
	Long: `Run a Docker command against a host.

Wraps the 'docker' command-line tool - see the Docker website for reference:

    http://docs.docker.io/en/latest/reference/commandline/

You can optionally specify a host by name - if you don't, the default host
will be used.`,
}

var flDockerHost = Docker.Flag.String("H", "", "")

var Proxy = &Command{
	UsageLine: "proxy [-H HOST]",
	Short:     "Start a local proxy to a host's Docker daemon",
	Long: `Start a local proxy to a host's Docker daemon.

Prints out a URL to pass to the 'docker' command, e.g.

    $ orchard proxy
    Started proxy at unix:///tmp/orchard-12345/orchard.sock

    $ docker -H unix:///tmp/orchard-12345/orchard.sock run ubuntu echo hello world
    hello world
`,
}

var flProxyHost = Proxy.Flag.String("H", "", "")

func RunHosts(cmd *Command, args []string) error {
	if len(args) > 0 {
		return cmd.UsageError("`orchard hosts` doesn't expect any arguments, but got: %s", strings.Join(args, " "))
	}

	httpClient, err := authenticator.Authenticate()
	if err != nil {
		return err
	}

	hosts, err := httpClient.GetHosts()
	if err != nil {
		return err
	}

	writer := tabwriter.NewWriter(os.Stdout, 20, 1, 3, ' ', 0)
	fmt.Fprintln(writer, "NAME\tSIZE\tIP")
	for _, host := range hosts {
		fmt.Fprintf(writer, "%s\t%s\t%s\n", host.Name, utils.HumanSize(host.Size*1024*1024), host.IPAddress)
	}
	writer.Flush()

	return nil
}

func RunStart(cmd *Command, args []string) error {
	if len(args) > 1 {
		return cmd.UsageError("`orchard start` expects at most 1 argument, but got more: %s", strings.Join(args[1:], " "))
	}

	httpClient, err := authenticator.Authenticate()
	if err != nil {
		return err
	}

	hostName, humanName := GetHostName(args)
	humanName = utils.Capitalize(humanName)

	size, sizeString := GetHostSize()
	if size == -1 {
		fmt.Fprintf(os.Stderr, "Sorry, %q isn't a size we support.\nValid sizes are %s.\n", sizeString, validSizes)
		return nil
	}

	host, err := httpClient.CreateHost(hostName, size)
	if err != nil {
		// HACK. api.go should decode JSON and return a specific type of error for this case.
		if strings.Contains(err.Error(), "already exists") {
			fmt.Fprintf(os.Stderr, "%s is already running.\nYou can create additional hosts with `orchard start NAME`.\n", humanName)
			return nil
		}
		if strings.Contains(err.Error(), "Invalid value") {
			fmt.Fprintf(os.Stderr, "Sorry, '%s' isn't a valid host name.\nHost names can only contain lowercase letters, numbers and underscores.\n", hostName)
			return nil
		}
		if strings.Contains(err.Error(), "Unsupported size") {
			fmt.Fprintf(os.Stderr, "Sorry, %q isn't a size we support.\nValid sizes are %s.\n", sizeString, validSizes)
			return nil
		}

		return err
	}
	fmt.Fprintf(os.Stderr, "%s running at %s\n", humanName, host.IPAddress)

	return nil
}

func RunStop(cmd *Command, args []string) error {
	if len(args) > 1 {
		return cmd.UsageError("`orchard stop` expects at most 1 argument, but got more: %s", strings.Join(args[1:], " "))
	}

	hostName, humanName := GetHostName(args)

	var confirm string
	fmt.Printf("Going to stop and delete %s. All data on it will be lost.\n", humanName)
	fmt.Print("Are you sure you're ready? [yN] ")
	fmt.Scanln(&confirm)

	if strings.ToLower(confirm) != "y" {
		return nil
	}

	httpClient, err := authenticator.Authenticate()
	if err != nil {
		return err
	}

	err = httpClient.DeleteHost(hostName)
	if err != nil {
		// HACK. api.go should decode JSON and return a specific type of error for this case.
		if strings.Contains(err.Error(), "Not found") {
			fmt.Fprintf(os.Stderr, "%s doesn't seem to be running.\nYou can view your running hosts with `orchard hosts`.\n", utils.Capitalize(humanName))
			return nil
		}

		return err
	}
	fmt.Fprintf(os.Stderr, "Stopped %s\n", humanName)

	return nil
}

func RunDocker(cmd *Command, args []string) error {
	return WithDockerProxy(func(socketPath string) error {
		err := CallDocker(args, []string{"DOCKER_HOST=unix://" + socketPath})
		if err != nil {
			return fmt.Errorf("Docker exited with error")
		}
		return nil
	})
}

func RunProxy(cmd *Command, args []string) error {
	if len(args) > 0 {
		return cmd.UsageError("`orchard proxy` doesn't expect any arguments, but got: %s", strings.Join(args, " "))
	}

	return WithDockerProxy(func(socketPath string) error {
		fmt.Fprintln(os.Stderr, "Started proxy at unix://"+socketPath)

		c := make(chan os.Signal)
		signal.Notify(c, syscall.SIGINT, syscall.SIGKILL)
		<-c

		fmt.Fprintln(os.Stderr, "\nStopping proxy")
		return nil
	})
}

func WithDockerProxy(callback func(string) error) error {
	hostName := "default"
	if *flDockerHost != "" {
		hostName = *flDockerHost
	}

	dirname, err := ioutil.TempDir("/tmp", "orchard-")
	if err != nil {
		return fmt.Errorf("Error creating temporary directory: %s\n", err)
	}
	defer os.RemoveAll(dirname)
	socketPath := path.Join(dirname, "orchard.sock")

	p, err := MakeProxy(socketPath, hostName)
	if err != nil {
		return fmt.Errorf("Error starting proxy: %v\n", err)
	}

	go p.Start()
	defer p.Stop()

	if err := <-p.ErrorChannel; err != nil {
		return fmt.Errorf("Error starting proxy: %v\n", err)
	}

	if err := callback(socketPath); err != nil {
		return err
	}

	return nil
}

func MakeProxy(socketPath string, hostName string) (*proxy.Proxy, error) {
	httpClient, err := authenticator.Authenticate()
	if err != nil {
		return nil, err
	}

	host, err := httpClient.GetHost(hostName)
	if err != nil {
		return nil, err
	}
	destination := host.IPAddress + ":4243"

	certData := []byte(host.ClientCert)
	keyData := []byte(host.ClientKey)
	config, err := tlsconfig.GetTLSConfig(certData, keyData)
	if err != nil {
		return nil, err
	}

	return proxy.New(
		func() (net.Listener, error) { return net.Listen("unix", socketPath) },
		func() (net.Conn, error) { return tls.Dial("tcp", destination, config) },
	), nil
}

func CallDocker(args []string, env []string) error {
	dockerPath := GetDockerPath()
	if dockerPath == "" {
		return errors.New("Can't find `docker` executable in $PATH.\nYou might need to install it: http://docs.docker.io/en/latest/installation/#installation-list")
	}

	cmd := exec.Command(dockerPath, args...)
	cmd.Env = env
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}

func GetDockerPath() string {
	for _, dir := range strings.Split(os.Getenv("PATH"), ":") {
		dockerPath := path.Join(dir, "docker")
		_, err := os.Stat(dockerPath)
		if err == nil {
			return dockerPath
		}
	}
	return ""
}

func GetHostName(args []string) (string, string) {
	hostName := "default"
	humanName := "default host"

	if len(args) > 0 {
		hostName = args[0]
		humanName = fmt.Sprintf("host '%s'", hostName)
	}

	return hostName, humanName
}

func GetHostSize() (int, string) {
	sizeString := *flStartSize

	bytes, err := utils.RAMInBytes(sizeString)
	if err != nil {
		return -1, sizeString
	}

	megs := bytes / (1024 * 1024)
	if megs < 1 {
		return -1, sizeString
	}

	return int(megs), sizeString
}
<|endoftext|>


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

package commands

import (
	"os"
	"runtime"
	"strings"

	"github.com/bogem/nehm/applescript"
	"github.com/bogem/nehm/config"
	"github.com/bogem/nehm/ui"
	"github.com/bogem/nehm/util"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
)

var RootCmd = listCommand

// Variables used in flags
var (
	limit, offset                       uint
	dlFolder, itunesPlaylist, permalink string
)

func Execute() {
	RootCmd.AddCommand(getCommand)
	RootCmd.AddCommand(searchCommand)
	RootCmd.AddCommand(syncCommand)
	RootCmd.AddCommand(versionCommand)
	RootCmd.Execute()
}

func addDlFolderFlag(cmd *cobra.Command) {
	cmd.Flags().StringVarP(&dlFolder, "dlFolder", "f", "", "filesystem path to download folder")
}

func addItunesPlaylistFlag(cmd *cobra.Command) {
	if runtime.GOOS == "darwin" {
		cmd.Flags().StringVarP(&itunesPlaylist, "itunesPlaylist", "i", "", "name of iTunes playlist")
	}
}

func addLimitFlag(cmd *cobra.Command) {
	cmd.Flags().UintVarP(&limit, "limit", "l", 10, "count of tracks on each page")
}

func addOffsetFlag(cmd *cobra.Command) {
	cmd.Flags().UintVarP(&offset, "offset", "o", 0, "offset relative to first like")
}

func addPermalinkFlag(cmd *cobra.Command) {
	cmd.Flags().StringVarP(&permalink, "permalink", "p", "", "user's permalink")
}

// initializeConfig initializes a config with flags.
// It only initializes field if cmd has corresponding flag.
func initializeConfig(cmd *cobra.Command) {
	readInConfig()

	flags := cmd.Flags()
	if flagExists(flags, "dlFolder") {
		initializeDlFolder(cmd)
	}
	if flagExists(flags, "permalink") {
		initializePermalink(cmd)
	}
	if flagExists(flags, "itunesPlaylist") {
		initializeItunesPlaylist(cmd)
	}
}

func readInConfig() {
	err := config.ReadInConfig()
	if err == config.ErrNotExist {
		ui.Warning("there is no config file. Read README to configure nehm")
		return
	}
	if err != nil {
		ui.Term("", err)
	}
}

func flagExists(fs *pflag.FlagSet, key string) bool {
	return fs.Lookup(key) != nil
}

func flagChanged(fs *pflag.FlagSet, key string) bool {
	flag := fs.Lookup(key)
	if flag == nil {
		return false
	}
	return flag.Changed
}

// initializeDlFolder initializes dlFolder value. If there is no dlFolder
// set up, then dlFolder is set to HOME env variable.
func initializeDlFolder(cmd *cobra.Command) {
	var df string

	if flagChanged(cmd.Flags(), "dlFolder") {
		df = dlFolder
	} else {
		df = config.Get("dlFolder")
	}

	if df == "" {
		ui.Warning("you didn't set a download folder. Tracks will be downloaded to your home directory.")
		df = os.Getenv("HOME")
	}

	config.Set("dlFolder", util.SanitizePath(df))
}

// initializePermalink initializes permalink value. If there is no permalink
// set up, then program is terminating.
func initializePermalink(cmd *cobra.Command) {
	var p string

	if flagChanged(cmd.Flags(), "permalink") {
		p = permalink
	} else {
		p = config.Get("permalink")
	}

	if p == "" {
		ui.Term("you didn't set a permalink. Use flag '-p' or set permalink in config file.\nTo know, what is permalink, read FAQ.", nil)
	} else {
		config.Set("permalink", p)
	}
}

// initializeItunesPlaylist initializes itunesPlaylist value. If there is no
// itunesPlaylist set up, then itunesPlaylist set up to blank string. Blank
// string is the sign, what tracks should not to be added to iTunes.
//
// initializeItunesPlaylist sets blank string to config, if OS is darwin.
func initializeItunesPlaylist(cmd *cobra.Command) {
	var playlist string

	if runtime.GOOS == "darwin" {
		if flagChanged(cmd.Flags(), "itunesPlaylist") {
			playlist = itunesPlaylist
		} else {
			playlist = config.Get("itunesPlaylist")
		}

		if playlist == "" {
			ui.Warning("you didn't set an iTunes playlist. Tracks won't be added to iTunes.")
		} else {
			playlistsList, err := applescript.ListOfPlaylists()
			if err != nil {
				ui.Term("couldn't get list of playlists", err)
			}
			if !strings.Contains(playlistsList, playlist) {
				ui.Term("playlist "+playlist+" doesn't exist. Please enter correct name.", nil)
			}
		}
	}

	config.Set("itunesPlaylist", playlist)
}
<commit_msg>commands: Fix typo in comments of initializeItunesPlaylist<commit_after>// Copyright 2016 Albert Nigmatzianov. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package commands

import (
	"os"
	"runtime"
	"strings"

	"github.com/bogem/nehm/applescript"
	"github.com/bogem/nehm/config"
	"github.com/bogem/nehm/ui"
	"github.com/bogem/nehm/util"
	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
)

var RootCmd = listCommand

// Variables used in flags
var (
	limit, offset                       uint
	dlFolder, itunesPlaylist, permalink string
)

func Execute() {
	RootCmd.AddCommand(getCommand)
	RootCmd.AddCommand(searchCommand)
	RootCmd.AddCommand(syncCommand)
	RootCmd.AddCommand(versionCommand)
	RootCmd.Execute()
}

func addDlFolderFlag(cmd *cobra.Command) {
	cmd.Flags().StringVarP(&dlFolder, "dlFolder", "f", "", "filesystem path to download folder")
}

func addItunesPlaylistFlag(cmd *cobra.Command) {
	if runtime.GOOS == "darwin" {
		cmd.Flags().StringVarP(&itunesPlaylist, "itunesPlaylist", "i", "", "name of iTunes playlist")
	}
}

func addLimitFlag(cmd *cobra.Command) {
	cmd.Flags().UintVarP(&limit, "limit", "l", 10, "count of tracks on each page")
}

func addOffsetFlag(cmd *cobra.Command) {
	cmd.Flags().UintVarP(&offset, "offset", "o", 0, "offset relative to first like")
}

func addPermalinkFlag(cmd *cobra.Command) {
	cmd.Flags().StringVarP(&permalink, "permalink", "p", "", "user's permalink")
}

// initializeConfig initializes a config with flags.
// It only initializes field if cmd has corresponding flag.
func initializeConfig(cmd *cobra.Command) {
	readInConfig()

	flags := cmd.Flags()
	if flagExists(flags, "dlFolder") {
		initializeDlFolder(cmd)
	}
	if flagExists(flags, "permalink") {
		initializePermalink(cmd)
	}
	if flagExists(flags, "itunesPlaylist") {
		initializeItunesPlaylist(cmd)
	}
}

func readInConfig() {
	err := config.ReadInConfig()
	if err == config.ErrNotExist {
		ui.Warning("there is no config file. Read README to configure nehm")
		return
	}
	if err != nil {
		ui.Term("", err)
	}
}

func flagExists(fs *pflag.FlagSet, key string) bool {
	return fs.Lookup(key) != nil
}

func flagChanged(fs *pflag.FlagSet, key string) bool {
	flag := fs.Lookup(key)
	if flag == nil {
		return false
	}
	return flag.Changed
}

// initializeDlFolder initializes dlFolder value. If there is no dlFolder
// set up, then dlFolder is set to HOME env variable.
func initializeDlFolder(cmd *cobra.Command) {
	var df string

	if flagChanged(cmd.Flags(), "dlFolder") {
		df = dlFolder
	} else {
		df = config.Get("dlFolder")
	}

	if df == "" {
		ui.Warning("you didn't set a download folder. Tracks will be downloaded to your home directory.")
		df = os.Getenv("HOME")
	}

	config.Set("dlFolder", util.SanitizePath(df))
}

// initializePermalink initializes permalink value. If there is no permalink
// set up, then program is terminating.
func initializePermalink(cmd *cobra.Command) {
	var p string

	if flagChanged(cmd.Flags(), "permalink") {
		p = permalink
	} else {
		p = config.Get("permalink")
	}

	if p == "" {
		ui.Term("you didn't set a permalink. Use flag '-p' or set permalink in config file.\nTo know, what is permalink, read FAQ.", nil)
	} else {
		config.Set("permalink", p)
	}
}

// initializeItunesPlaylist initializes itunesPlaylist value. If there is no
// itunesPlaylist set up, then itunesPlaylist set up to blank string. Blank
// string is the sign, what tracks should not to be added to iTunes.
//
// initializeItunesPlaylist sets blank string to config, if OS is not darwin.
func initializeItunesPlaylist(cmd *cobra.Command) {
	var playlist string

	if runtime.GOOS == "darwin" {
		if flagChanged(cmd.Flags(), "itunesPlaylist") {
			playlist = itunesPlaylist
		} else {
			playlist = config.Get("itunesPlaylist")
		}

		if playlist == "" {
			ui.Warning("you didn't set an iTunes playlist. Tracks won't be added to iTunes.")
		} else {
			playlistsList, err := applescript.ListOfPlaylists()
			if err != nil {
				ui.Term("couldn't get list of playlists", err)
			}
			if !strings.Contains(playlistsList, playlist) {
				ui.Term("playlist "+playlist+" doesn't exist. Please enter correct name.", nil)
			}
		}
	}

	config.Set("itunesPlaylist", playlist)
}
<|endoftext|>


<commit_before>package commands

import (
	"github.com/cameronnewman/fastlycli/fastlyclient"
	"github.com/codegangsta/cli"
	"os"
)

func Excute() {

	app := cli.NewApp()
	app.Name = "fastlycli"
	app.Usage = "Manage Fastly CDN Services via the cli"
	app.Version = "0.9.0"
	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:  "verbose",
			Usage: "enable verbose logging",
		},
	}
	app.Commands = []cli.Command{
		{
			Name:  "service",
			Usage: "Get Service Details",
			Subcommands: []cli.Command{
				{
					Name:  "details",
					Usage: "Gets service details",
					Flags: []cli.Flag{
						cli.StringFlag{
							Name:  "service, s",
							Usage: "Service Name defined in app.fastly.com",
						},
					},
					Action: func(c *cli.Context) {
						if c.String("service") != "" {
							fastly := fastlyclient.NewFastlyClient()
							fastly.GetServiceDetails(c.String("service"))
						}
					},
				},
				{
					Name:  "domains",
					Usage: "Gets a service domains",
					Flags: []cli.Flag{
						cli.StringFlag{
							Name:  "service, s",
							Usage: "Service Name defined in app.fastly.com",
						},
					},
					Action: func(c *cli.Context) {
						if c.String("service") != "" {
							fastly := fastlyclient.NewFastlyClient()
							fastly.GetServiceDomains(c.String("service"))
						}
					},
				},

				{
					Name:  "backends",
					Usage: "Gets a service domains",
					Flags: []cli.Flag{
						cli.StringFlag{
							Name:  "service, s",
							Usage: "Service Name defined in app.fastly.com",
						},
					},
					Action: func(c *cli.Context) {
						if c.String("service") != "" {
							fastly := fastlyclient.NewFastlyClient()
							fastly.GetServiceBackends(c.String("service"))
						}
					},
				},
			},
		},
		{
			Name:  "purge",
			Usage: "Purge objects from the CDN",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "service, s",
					Usage: "Service Name defined in app.fastly.com",
				},
				cli.StringFlag{
					Name:  "object, o",
					Usage: "Objects to be purged",
				},
			},
			Action: func(c *cli.Context) {
				if c.IsSet("service") && c.String("service") != "" {
					fastly := fastlyclient.NewFastlyClient()
					fastly.PurgeObjects(c.String("service"), c.String("object"))
				} else {
					println("No service name defined")
				}
			},
		},
	}

	app.Run(os.Args)
}
<commit_msg>Updated comments<commit_after>package commands

import (
	"os"

	"github.com/cameronnewman/fastlycli/fastlyclient"
	"github.com/codegangsta/cli"
)

//Excute main section
func Excute() {

	app := cli.NewApp()
	app.Name = "fastlycli"
	app.Usage = "Manage Fastly CDN Services via the cli"
	app.Version = "0.9.0"
	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:  "verbose",
			Usage: "enable verbose logging",
		},
	}
	app.Commands = []cli.Command{
		{
			Name:  "service",
			Usage: "Get Service Details",
			Subcommands: []cli.Command{
				{
					Name:  "details",
					Usage: "Gets service details",
					Flags: []cli.Flag{
						cli.StringFlag{
							Name:  "service, s",
							Usage: "Service Name defined in app.fastly.com",
						},
					},
					Action: func(c *cli.Context) {
						if c.String("service") != "" {
							fastly := fastlyclient.NewFastlyClient()
							fastly.GetServiceDetails(c.String("service"))
						}
					},
				},
				{
					Name:  "domains",
					Usage: "Gets a service domains",
					Flags: []cli.Flag{
						cli.StringFlag{
							Name:  "service, s",
							Usage: "Service Name defined in app.fastly.com",
						},
					},
					Action: func(c *cli.Context) {
						if c.String("service") != "" {
							fastly := fastlyclient.NewFastlyClient()
							fastly.GetServiceDomains(c.String("service"))
						}
					},
				},

				{
					Name:  "backends",
					Usage: "Gets a service domains",
					Flags: []cli.Flag{
						cli.StringFlag{
							Name:  "service, s",
							Usage: "Service Name defined in app.fastly.com",
						},
					},
					Action: func(c *cli.Context) {
						if c.String("service") != "" {
							fastly := fastlyclient.NewFastlyClient()
							fastly.GetServiceBackends(c.String("service"))
						}
					},
				},
			},
		},
		{
			Name:  "purge",
			Usage: "Purge objects from the CDN",
			Flags: []cli.Flag{
				cli.StringFlag{
					Name:  "service, s",
					Usage: "Service Name defined in app.fastly.com",
				},
				cli.StringFlag{
					Name:  "object, o",
					Usage: "Objects to be purged",
				},
			},
			Action: func(c *cli.Context) {
				if c.IsSet("service") && c.String("service") != "" {
					fastly := fastlyclient.NewFastlyClient()
					fastly.PurgeObjects(c.String("service"), c.String("object"))
				} else {
					println("No service name defined")
				}
			},
		},
	}

	app.Run(os.Args)
}
<|endoftext|>


<commit_before>package mux

import (
	"context"
	"io"
	"sync"
	"time"

	"v2ray.com/core/common"
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/errors"
	"v2ray.com/core/common/net"
	"v2ray.com/core/common/protocol"
	"v2ray.com/core/common/session"
	"v2ray.com/core/common/signal/done"
	"v2ray.com/core/common/task"
	"v2ray.com/core/common/vio"
	"v2ray.com/core/proxy"
	"v2ray.com/core/transport/internet"
	"v2ray.com/core/transport/pipe"
)

type ClientManager struct {
	Picker WorkerPicker
}

func (m *ClientManager) Dispatch(ctx context.Context, link *vio.Link) error {
	for {
		worker, err := m.Picker.PickAvailable()
		if err != nil {
			return err
		}
		if worker.Dispatch(ctx, link) {
			return nil
		}
	}
}

type WorkerPicker interface {
	PickAvailable() (*ClientWorker, error)
}

type IncrementalWorkerPicker struct {
	Factory ClientWorkerFactory

	access      sync.Mutex
	workers     []*ClientWorker
	cleanupTask *task.Periodic
}

func (p *IncrementalWorkerPicker) cleanupFunc() error {
	p.access.Lock()
	defer p.access.Unlock()

	if len(p.workers) == 0 {
		return newError("no worker")
	}

	p.cleanup()
	return nil
}

func (p *IncrementalWorkerPicker) cleanup() {
	var activeWorkers []*ClientWorker
	for _, w := range p.workers {
		if !w.Closed() {
			activeWorkers = append(activeWorkers, w)
		}
	}
	p.workers = activeWorkers
}

func (p *IncrementalWorkerPicker) pickInternal() (*ClientWorker, error, bool) {
	p.access.Lock()
	defer p.access.Unlock()

	for _, w := range p.workers {
		if !w.IsFull() {
			return w, nil, false
		}
	}

	p.cleanup()

	worker, err := p.Factory.Create()
	if err != nil {
		return nil, err, false
	}
	p.workers = append(p.workers, worker)

	if p.cleanupTask == nil {
		p.cleanupTask = &task.Periodic{
			Interval: time.Second * 30,
			Execute:  p.cleanupFunc,
		}
	}

	return worker, nil, true
}

func (p *IncrementalWorkerPicker) PickAvailable() (*ClientWorker, error) {
	worker, err, start := p.pickInternal()
	if start {
		p.cleanupTask.Start()
	}

	return worker, err
}

type ClientWorkerFactory interface {
	Create() (*ClientWorker, error)
}

type DialingWorkerFactory struct {
	Proxy    proxy.Outbound
	Dialer   internet.Dialer
	Strategy ClientStrategy
}

func (f *DialingWorkerFactory) Create() (*ClientWorker, error) {
	opts := []pipe.Option{pipe.WithSizeLimit(64 * 1024)}
	uplinkReader, upLinkWriter := pipe.New(opts...)
	downlinkReader, downlinkWriter := pipe.New(opts...)

	c, err := NewClientWorker(vio.Link{
		Reader: downlinkReader,
		Writer: upLinkWriter,
	}, f.Strategy)

	if err != nil {
		return nil, err
	}

	go func(p proxy.Outbound, d internet.Dialer, c common.Closable) {
		ctx := session.ContextWithOutbound(context.Background(), &session.Outbound{
			Target: net.TCPDestination(muxCoolAddress, muxCoolPort),
		})
		ctx, cancel := context.WithCancel(ctx)

		if err := p.Process(ctx, &vio.Link{Reader: uplinkReader, Writer: downlinkWriter}, d); err != nil {
			errors.New("failed to handler mux client connection").Base(err).WriteToLog()
		}
		common.Must(c.Close())
		cancel()
	}(f.Proxy, f.Dialer, c.done)

	return c, nil
}

type ClientStrategy struct {
	MaxConcurrency uint32
	MaxConnection  uint32
}

type ClientWorker struct {
	sessionManager *SessionManager
	link           vio.Link
	done           *done.Instance
	strategy       ClientStrategy
}

var muxCoolAddress = net.DomainAddress("v1.mux.cool")
var muxCoolPort = net.Port(9527)

// NewClientWorker creates a new mux.Client.
func NewClientWorker(stream vio.Link, s ClientStrategy) (*ClientWorker, error) {
	c := &ClientWorker{
		sessionManager: NewSessionManager(),
		link:           stream,
		done:           done.New(),
		strategy:       s,
	}

	go c.fetchOutput()
	go c.monitor()

	return c, nil
}

// Closed returns true if this Client is closed.
func (m *ClientWorker) Closed() bool {
	return m.done.Done()
}

func (m *ClientWorker) monitor() {
	timer := time.NewTicker(time.Second * 16)
	defer timer.Stop()

	for {
		select {
		case <-m.done.Wait():
			m.sessionManager.Close()
			common.Close(m.link.Writer)    // nolint: errcheck
			pipe.CloseError(m.link.Reader) // nolint: errcheck
			return
		case <-timer.C:
			size := m.sessionManager.Size()
			if size == 0 && m.sessionManager.CloseIfNoSession() {
				common.Must(m.done.Close())
			}
		}
	}
}

func writeFirstPayload(reader buf.Reader, writer *Writer) error {
	err := buf.CopyOnceTimeout(reader, writer, time.Millisecond*100)
	if err == buf.ErrNotTimeoutReader || err == buf.ErrReadTimeout {
		return writer.WriteMultiBuffer(buf.MultiBuffer{})
	}

	if err != nil {
		return err
	}

	return nil
}

func fetchInput(ctx context.Context, s *Session, output buf.Writer) {
	dest := session.OutboundFromContext(ctx).Target
	transferType := protocol.TransferTypeStream
	if dest.Network == net.Network_UDP {
		transferType = protocol.TransferTypePacket
	}
	s.transferType = transferType
	writer := NewWriter(s.ID, dest, output, transferType)
	defer s.Close()      // nolint: errcheck
	defer writer.Close() // nolint: errcheck

	newError("dispatching request to ", dest).WriteToLog(session.ExportIDToError(ctx))
	if err := writeFirstPayload(s.input, writer); err != nil {
		newError("failed to write first payload").Base(err).WriteToLog(session.ExportIDToError(ctx))
		writer.hasError = true
		pipe.CloseError(s.input)
		return
	}

	if err := buf.Copy(s.input, writer); err != nil {
		newError("failed to fetch all input").Base(err).WriteToLog(session.ExportIDToError(ctx))
		writer.hasError = true
		pipe.CloseError(s.input)
		return
	}
}

func (m *ClientWorker) IsClosing() bool {
	sm := m.sessionManager
	if m.strategy.MaxConnection > 0 && sm.Count() >= int(m.strategy.MaxConnection) {
		return true
	}
	return false
}

func (m *ClientWorker) IsFull() bool {
	if m.IsClosing() {
		return true
	}

	sm := m.sessionManager
	if m.strategy.MaxConcurrency > 0 && sm.Size() >= int(m.strategy.MaxConcurrency) {
		return true
	}
	return false
}

func (m *ClientWorker) Dispatch(ctx context.Context, link *vio.Link) bool {
	if m.IsFull() || m.Closed() {
		return false
	}

	sm := m.sessionManager
	s := sm.Allocate()
	if s == nil {
		return false
	}
	s.input = link.Reader
	s.output = link.Writer
	go fetchInput(ctx, s, m.link.Writer)
	return true
}

func (m *ClientWorker) handleStatueKeepAlive(meta *FrameMetadata, reader *buf.BufferedReader) error {
	if meta.Option.Has(OptionData) {
		return buf.Copy(NewStreamReader(reader), buf.Discard)
	}
	return nil
}

func (m *ClientWorker) handleStatusNew(meta *FrameMetadata, reader *buf.BufferedReader) error {
	if meta.Option.Has(OptionData) {
		return buf.Copy(NewStreamReader(reader), buf.Discard)
	}
	return nil
}

func (m *ClientWorker) handleStatusKeep(meta *FrameMetadata, reader *buf.BufferedReader) error {
	if !meta.Option.Has(OptionData) {
		return nil
	}

	s, found := m.sessionManager.Get(meta.SessionID)
	if !found {
		return buf.Copy(NewStreamReader(reader), buf.Discard)
	}

	rr := s.NewReader(reader)
	err := buf.Copy(rr, s.output)
	if err != nil && buf.IsWriteError(err) {
		newError("failed to write to downstream. closing session ", s.ID).Base(err).WriteToLog()

		drainErr := buf.Copy(rr, buf.Discard)
		pipe.CloseError(s.input)
		s.Close()
		return drainErr
	}

	return err
}

func (m *ClientWorker) handleStatusEnd(meta *FrameMetadata, reader *buf.BufferedReader) error {
	if s, found := m.sessionManager.Get(meta.SessionID); found {
		if meta.Option.Has(OptionError) {
			pipe.CloseError(s.input)
			pipe.CloseError(s.output)
		}
		s.Close()
	}
	if meta.Option.Has(OptionData) {
		return buf.Copy(NewStreamReader(reader), buf.Discard)
	}
	return nil
}

func (m *ClientWorker) fetchOutput() {
	defer func() {
		common.Must(m.done.Close())
	}()

	reader := &buf.BufferedReader{Reader: m.link.Reader}

	var meta FrameMetadata
	for {
		err := meta.Unmarshal(reader)
		if err != nil {
			if errors.Cause(err) != io.EOF {
				newError("failed to read metadata").Base(err).WriteToLog()
			}
			break
		}

		switch meta.SessionStatus {
		case SessionStatusKeepAlive:
			err = m.handleStatueKeepAlive(&meta, reader)
		case SessionStatusEnd:
			err = m.handleStatusEnd(&meta, reader)
		case SessionStatusNew:
			err = m.handleStatusNew(&meta, reader)
		case SessionStatusKeep:
			err = m.handleStatusKeep(&meta, reader)
		default:
			status := meta.SessionStatus
			newError("unknown status: ", status).AtError().WriteToLog()
			return
		}

		if err != nil {
			newError("failed to process data").Base(err).WriteToLog()
			return
		}
	}
}
<commit_msg>better load balancing<commit_after>package mux

import (
	"context"
	"io"
	"sync"
	"time"

	"v2ray.com/core/common"
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/errors"
	"v2ray.com/core/common/net"
	"v2ray.com/core/common/protocol"
	"v2ray.com/core/common/session"
	"v2ray.com/core/common/signal/done"
	"v2ray.com/core/common/task"
	"v2ray.com/core/common/vio"
	"v2ray.com/core/proxy"
	"v2ray.com/core/transport/internet"
	"v2ray.com/core/transport/pipe"
)

type ClientManager struct {
	Picker WorkerPicker
}

func (m *ClientManager) Dispatch(ctx context.Context, link *vio.Link) error {
	for {
		worker, err := m.Picker.PickAvailable()
		if err != nil {
			return err
		}
		if worker.Dispatch(ctx, link) {
			return nil
		}
	}
}

type WorkerPicker interface {
	PickAvailable() (*ClientWorker, error)
}

type IncrementalWorkerPicker struct {
	Factory ClientWorkerFactory

	access      sync.Mutex
	workers     []*ClientWorker
	cleanupTask *task.Periodic
}

func (p *IncrementalWorkerPicker) cleanupFunc() error {
	p.access.Lock()
	defer p.access.Unlock()

	if len(p.workers) == 0 {
		return newError("no worker")
	}

	p.cleanup()
	return nil
}

func (p *IncrementalWorkerPicker) cleanup() {
	var activeWorkers []*ClientWorker
	for _, w := range p.workers {
		if !w.Closed() {
			activeWorkers = append(activeWorkers, w)
		}
	}
	p.workers = activeWorkers
}

func (p *IncrementalWorkerPicker) findAvailable() int {
	for idx, w := range p.workers {
		if !w.IsFull() {
			return idx
		}
	}

	return -1
}

func (p *IncrementalWorkerPicker) pickInternal() (*ClientWorker, error, bool) {
	p.access.Lock()
	defer p.access.Unlock()

	idx := p.findAvailable()
	if idx >= 0 {
		n := len(p.workers)
		if n > 1 && idx != n-1 {
			p.workers[n-1], p.workers[idx] = p.workers[idx], p.workers[n-1]
		}
		return p.workers[idx], nil, false
	}

	p.cleanup()

	worker, err := p.Factory.Create()
	if err != nil {
		return nil, err, false
	}
	p.workers = append(p.workers, worker)

	if p.cleanupTask == nil {
		p.cleanupTask = &task.Periodic{
			Interval: time.Second * 30,
			Execute:  p.cleanupFunc,
		}
	}

	return worker, nil, true
}

func (p *IncrementalWorkerPicker) PickAvailable() (*ClientWorker, error) {
	worker, err, start := p.pickInternal()
	if start {
		p.cleanupTask.Start()
	}

	return worker, err
}

type ClientWorkerFactory interface {
	Create() (*ClientWorker, error)
}

type DialingWorkerFactory struct {
	Proxy    proxy.Outbound
	Dialer   internet.Dialer
	Strategy ClientStrategy
}

func (f *DialingWorkerFactory) Create() (*ClientWorker, error) {
	opts := []pipe.Option{pipe.WithSizeLimit(64 * 1024)}
	uplinkReader, upLinkWriter := pipe.New(opts...)
	downlinkReader, downlinkWriter := pipe.New(opts...)

	c, err := NewClientWorker(vio.Link{
		Reader: downlinkReader,
		Writer: upLinkWriter,
	}, f.Strategy)

	if err != nil {
		return nil, err
	}

	go func(p proxy.Outbound, d internet.Dialer, c common.Closable) {
		ctx := session.ContextWithOutbound(context.Background(), &session.Outbound{
			Target: net.TCPDestination(muxCoolAddress, muxCoolPort),
		})
		ctx, cancel := context.WithCancel(ctx)

		if err := p.Process(ctx, &vio.Link{Reader: uplinkReader, Writer: downlinkWriter}, d); err != nil {
			errors.New("failed to handler mux client connection").Base(err).WriteToLog()
		}
		common.Must(c.Close())
		cancel()
	}(f.Proxy, f.Dialer, c.done)

	return c, nil
}

type ClientStrategy struct {
	MaxConcurrency uint32
	MaxConnection  uint32
}

type ClientWorker struct {
	sessionManager *SessionManager
	link           vio.Link
	done           *done.Instance
	strategy       ClientStrategy
}

var muxCoolAddress = net.DomainAddress("v1.mux.cool")
var muxCoolPort = net.Port(9527)

// NewClientWorker creates a new mux.Client.
func NewClientWorker(stream vio.Link, s ClientStrategy) (*ClientWorker, error) {
	c := &ClientWorker{
		sessionManager: NewSessionManager(),
		link:           stream,
		done:           done.New(),
		strategy:       s,
	}

	go c.fetchOutput()
	go c.monitor()

	return c, nil
}

// Closed returns true if this Client is closed.
func (m *ClientWorker) Closed() bool {
	return m.done.Done()
}

func (m *ClientWorker) monitor() {
	timer := time.NewTicker(time.Second * 16)
	defer timer.Stop()

	for {
		select {
		case <-m.done.Wait():
			m.sessionManager.Close()
			common.Close(m.link.Writer)    // nolint: errcheck
			pipe.CloseError(m.link.Reader) // nolint: errcheck
			return
		case <-timer.C:
			size := m.sessionManager.Size()
			if size == 0 && m.sessionManager.CloseIfNoSession() {
				common.Must(m.done.Close())
			}
		}
	}
}

func writeFirstPayload(reader buf.Reader, writer *Writer) error {
	err := buf.CopyOnceTimeout(reader, writer, time.Millisecond*100)
	if err == buf.ErrNotTimeoutReader || err == buf.ErrReadTimeout {
		return writer.WriteMultiBuffer(buf.MultiBuffer{})
	}

	if err != nil {
		return err
	}

	return nil
}

func fetchInput(ctx context.Context, s *Session, output buf.Writer) {
	dest := session.OutboundFromContext(ctx).Target
	transferType := protocol.TransferTypeStream
	if dest.Network == net.Network_UDP {
		transferType = protocol.TransferTypePacket
	}
	s.transferType = transferType
	writer := NewWriter(s.ID, dest, output, transferType)
	defer s.Close()      // nolint: errcheck
	defer writer.Close() // nolint: errcheck

	newError("dispatching request to ", dest).WriteToLog(session.ExportIDToError(ctx))
	if err := writeFirstPayload(s.input, writer); err != nil {
		newError("failed to write first payload").Base(err).WriteToLog(session.ExportIDToError(ctx))
		writer.hasError = true
		pipe.CloseError(s.input)
		return
	}

	if err := buf.Copy(s.input, writer); err != nil {
		newError("failed to fetch all input").Base(err).WriteToLog(session.ExportIDToError(ctx))
		writer.hasError = true
		pipe.CloseError(s.input)
		return
	}
}

func (m *ClientWorker) IsClosing() bool {
	sm := m.sessionManager
	if m.strategy.MaxConnection > 0 && sm.Count() >= int(m.strategy.MaxConnection) {
		return true
	}
	return false
}

func (m *ClientWorker) IsFull() bool {
	if m.IsClosing() {
		return true
	}

	sm := m.sessionManager
	if m.strategy.MaxConcurrency > 0 && sm.Size() >= int(m.strategy.MaxConcurrency) {
		return true
	}
	return false
}

func (m *ClientWorker) Dispatch(ctx context.Context, link *vio.Link) bool {
	if m.IsFull() || m.Closed() {
		return false
	}

	sm := m.sessionManager
	s := sm.Allocate()
	if s == nil {
		return false
	}
	s.input = link.Reader
	s.output = link.Writer
	go fetchInput(ctx, s, m.link.Writer)
	return true
}

func (m *ClientWorker) handleStatueKeepAlive(meta *FrameMetadata, reader *buf.BufferedReader) error {
	if meta.Option.Has(OptionData) {
		return buf.Copy(NewStreamReader(reader), buf.Discard)
	}
	return nil
}

func (m *ClientWorker) handleStatusNew(meta *FrameMetadata, reader *buf.BufferedReader) error {
	if meta.Option.Has(OptionData) {
		return buf.Copy(NewStreamReader(reader), buf.Discard)
	}
	return nil
}

func (m *ClientWorker) handleStatusKeep(meta *FrameMetadata, reader *buf.BufferedReader) error {
	if !meta.Option.Has(OptionData) {
		return nil
	}

	s, found := m.sessionManager.Get(meta.SessionID)
	if !found {
		return buf.Copy(NewStreamReader(reader), buf.Discard)
	}

	rr := s.NewReader(reader)
	err := buf.Copy(rr, s.output)
	if err != nil && buf.IsWriteError(err) {
		newError("failed to write to downstream. closing session ", s.ID).Base(err).WriteToLog()

		drainErr := buf.Copy(rr, buf.Discard)
		pipe.CloseError(s.input)
		s.Close()
		return drainErr
	}

	return err
}

func (m *ClientWorker) handleStatusEnd(meta *FrameMetadata, reader *buf.BufferedReader) error {
	if s, found := m.sessionManager.Get(meta.SessionID); found {
		if meta.Option.Has(OptionError) {
			pipe.CloseError(s.input)
			pipe.CloseError(s.output)
		}
		s.Close()
	}
	if meta.Option.Has(OptionData) {
		return buf.Copy(NewStreamReader(reader), buf.Discard)
	}
	return nil
}

func (m *ClientWorker) fetchOutput() {
	defer func() {
		common.Must(m.done.Close())
	}()

	reader := &buf.BufferedReader{Reader: m.link.Reader}

	var meta FrameMetadata
	for {
		err := meta.Unmarshal(reader)
		if err != nil {
			if errors.Cause(err) != io.EOF {
				newError("failed to read metadata").Base(err).WriteToLog()
			}
			break
		}

		switch meta.SessionStatus {
		case SessionStatusKeepAlive:
			err = m.handleStatueKeepAlive(&meta, reader)
		case SessionStatusEnd:
			err = m.handleStatusEnd(&meta, reader)
		case SessionStatusNew:
			err = m.handleStatusNew(&meta, reader)
		case SessionStatusKeep:
			err = m.handleStatusKeep(&meta, reader)
		default:
			status := meta.SessionStatus
			newError("unknown status: ", status).AtError().WriteToLog()
			return
		}

		if err != nil {
			newError("failed to process data").Base(err).WriteToLog()
			return
		}
	}
}
<|endoftext|>


<commit_before>package common

const (
	STYPE_PRIORITY_QUEUE      = "pqueue"    // Highest priority goes first.
	STYPE_DOUBLE_SIDED_QUEUE  = "dsqueue"   // Double sided queue
	STYPE_FIFO_QUEUE          = "fifoqueue" // Standard FIFO
	STYPE_FAIR_PRIORITY_QUEUE = "fairqueue" // POPs are fairly distributed across all priorities.
	STYPE_COUNTERS            = "counters"  // Atomic counters.
	STYPE_SEQUENCE_READ       = "seqreader" // Data to read in sequential order.
)

type ServiceMetaInfo struct {
	SType    string
	Id       int32
	Name     string
	CreateTs int64
	Disabled bool
}

func NewServiceMetaInfo(sType string, id int32, name string) *ServiceMetaInfo {
	return &ServiceMetaInfo{
		SType:    sType,
		Id:       id,
		Name:     name,
		CreateTs: Uts(),
		Disabled: false,
	}
}

func ServiceInfoFromBinary(data []byte) (*ServiceMetaInfo, error) {
	smi := ServiceMetaInfo{}
	err := StructFromBinary(&smi, data)
	if err != nil {
		return nil, err
	}
	return &smi, nil
}

func (q *ServiceMetaInfo) ToBinary() []byte {
	return StructToBinary(q)
}
<commit_msg>Removed unnecessary queue types.<commit_after>package common

const (
	STYPE_PRIORITY_QUEUE      = "pqueue"    // Highest priority goes first.
	STYPE_DOUBLE_SIDED_QUEUE  = "dsqueue"   // Double sided queue
	STYPE_COUNTERS            = "counters"  // Atomic counters.
	STYPE_FAIR_PRIORITY_QUEUE = "fairqueue" // POPs are fairly distributed across all priorities.
)

type ServiceMetaInfo struct {
	SType    string
	Id       int32
	Name     string
	CreateTs int64
	Disabled bool
}

func NewServiceMetaInfo(sType string, id int32, name string) *ServiceMetaInfo {
	return &ServiceMetaInfo{
		SType:    sType,
		Id:       id,
		Name:     name,
		CreateTs: Uts(),
		Disabled: false,
	}
}

func ServiceInfoFromBinary(data []byte) (*ServiceMetaInfo, error) {
	smi := ServiceMetaInfo{}
	err := StructFromBinary(&smi, data)
	if err != nil {
		return nil, err
	}
	return &smi, nil
}

func (q *ServiceMetaInfo) ToBinary() []byte {
	return StructToBinary(q)
}
<|endoftext|>