<commit_before>package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"net/url"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"

	"golang.org/x/net/context"
)

type (
	CommandConfig struct {
		Template []string            `json:"-"`
		Options  map[string][]string `json:"options,omitempty"`
		Dryrun   bool                `json:"dryrun,omitempty"`
	}

	Job struct {
		config *CommandConfig
		// https://godoc.org/google.golang.org/genproto/googleapis/pubsub/v1#ReceivedMessage
		message      *JobMessage
		notification *ProgressNotification
		storage      Storage

		// These are set at at setupWorkspace
		workspace     string
		downloads_dir string
		uploads_dir   string

		// These are set at setupDownloadFiles
		downloadFileMap     map[string]string
		remoteDownloadFiles interface{}
		localDownloadFiles  interface{}

		cmd *exec.Cmd
	}
)

func (job *Job) run(ctx context.Context) error {
	err := job.runImpl(ctx)
	switch err.(type) {
	case InvalidJobError:
		err := job.withNotify(CANCELLING, job.message.Ack)()
		if err != nil {
			return err
		}
	}
	return err
}

func (job *Job) runImpl(ctx context.Context) error {
	job.notification.notify(PROCESSING, job.message.MessageId(), "info")

	err := job.message.Validate()
	if err != nil {
		log.Printf("Invalid Message: MessageId: %v, Message: %v, error: %v\n", job.message.MessageId(), job.message.raw.Message, err)
		return nil
	}

	defer job.clearWorkspace() // Call clearWorkspace even if setupWorkspace retuns error
	err = job.setupWorkspace()
	if err != nil {
		return err
	}

	err = job.withNotify(PREPARING, job.setupDownloadFiles)()
	if err != nil {
		return err
	}

	err = job.build()
	if err != nil {
		log.Fatalf("Command build Error template: %v msg: %v cause of %v\n", job.config.Template, job.message, err)
		return err
	}

	go job.message.sendMADPeriodically()
	defer job.message.Done()

	err = job.withNotify(DOWNLOADING, job.downloadFiles)()
	if err != nil {
		return err
	}

	err = job.withNotify(EXECUTING, job.execute)()
	if err != nil {
		return err
	}

	err = job.withNotify(UPLOADING, job.uploadFiles)()
	if err != nil {
		return err
	}

	err = job.withNotify(ACKSENDING, job.message.Ack)()
	if err != nil {
		return err
	}

	job.notification.notify(CLEANUP, job.message.MessageId(), "info")
	return err
}

func (job *Job) withNotify(progress int, f func() error) func() error {
	msg_id := job.message.MessageId()
	return func() error {
		job.notification.notify(progress, msg_id, "info")
		err := f()
		if err != nil {
			job.notification.notify(progress+2, msg_id, "error")
			return err
		}
		job.notification.notify(progress+1, msg_id, "info")
		return nil
	}
}

func (job *Job) setupWorkspace() error {
	dir, err := ioutil.TempDir("", "workspace")
	if err != nil {
		log.Fatal(err)
		return err
	}

	subdirs := []string{
		filepath.Join(dir, "downloads"),
		filepath.Join(dir, "uploads"),
	}
	for _, subdir := range subdirs {
		err := os.MkdirAll(subdir, 0700)
		if err != nil {
			return err
		}
	}
	job.workspace = dir
	job.downloads_dir = subdirs[0]
	job.uploads_dir = subdirs[1]
	return nil
}

func (job *Job) clearWorkspace() error {
	if job.workspace != "" {
		return os.RemoveAll(job.workspace)
	}
	return nil
}

func (job *Job) setupDownloadFiles() error {
	job.downloadFileMap = map[string]string{}
	job.remoteDownloadFiles = job.message.DownloadFiles()
	objects := job.flatten(job.remoteDownloadFiles)
	remoteUrls := []string{}
	for _, obj := range objects {
		switch obj.(type) {
		case string:
			remoteUrls = append(remoteUrls, obj.(string))
		default:
			log.Printf("Invalid download file URL: %v [%T]", obj, obj)
		}
	}
	for _, remote_url := range remoteUrls {
		url, err := url.Parse(remote_url)
		if err != nil {
			log.Fatalf("Invalid URL: %v because of %v\n", remote_url, err)
			return err
		}
		urlstr := fmt.Sprintf("gs://%v%v", url.Host, url.Path)
		destPath := filepath.Join(job.downloads_dir, url.Host, url.Path)
		job.downloadFileMap[urlstr] = destPath
	}
	job.localDownloadFiles = job.copyWithFileMap(job.remoteDownloadFiles)
	return nil
}

func (job *Job) copyWithFileMap(obj interface{}) interface{} {
	switch obj.(type) {
	case map[string]interface{}:
		result := map[string]interface{}{}
		for k, v := range obj.(map[string]interface{}) {
			result[k] = job.copyWithFileMap(v)
		}
		return result
	case []interface{}:
		result := []interface{}{}
		for _, v := range obj.([]interface{}) {
			result = append(result, job.copyWithFileMap(v))
		}
		return result
	case string:
		return job.downloadFileMap[obj.(string)]
	default:
		return obj
	}
}

func (job *Job) buildVariable() *Variable {
	return &Variable{
		data: map[string]interface{}{
			"workspace":             job.workspace,
			"downloads_dir":         job.downloads_dir,
			"uploads_dir":           job.uploads_dir,
			"download_files":        job.localDownloadFiles,
			"local_download_files":  job.localDownloadFiles,
			"remote_download_files": job.remoteDownloadFiles,
			"attrs":                 job.message.raw.Message.Attributes,
			"attributes":            job.message.raw.Message.Attributes,
			"data":                  job.message.raw.Message.Data,
		},
	}
}

func (job *Job) build() error {
	v := job.buildVariable()

	values, err := job.extract(v, job.config.Template)
	if err != nil {
		return err
	}
	if len(job.config.Options) > 0 {
		key := strings.Join(values, " ")
		t := job.config.Options[key]
		if t == nil {
			t = job.config.Options["default"]
		}
		if t != nil {
			values, err = job.extract(v, t)
			if err != nil {
				return err
			}
		}
	}
	job.cmd = exec.Command(values[0], values[1:]...)
	return nil
}

func (job *Job) extract(v *Variable, values []string) ([]string, error) {
	result := []string{}
	for _, src := range values {
		extracted, err := v.expand(src)
		if err != nil {
			return nil, &InvalidJobError{err.Error()}
		}
		vals := strings.Split(extracted, v.separator)
		for _, val := range vals {
			result = append(result, val)
		}
	}
	return result, nil
}

func (job *Job) downloadFiles() error {
	for remoteURL, destPath := range job.downloadFileMap {
		url, err := url.Parse(remoteURL)
		if err != nil {
			log.Fatalf("Invalid URL: %v because of %v\n", remoteURL, err)
			return err
		}

		dir := path.Dir(destPath)
		err = os.MkdirAll(dir, 0700)
		if err != nil {
			return err
		}

		err = job.storage.Download(url.Host, url.Path[1:], destPath)
		if err != nil {
			return err
		}
	}
	return nil
}

func (job *Job) execute() error {
	var out bytes.Buffer
	job.cmd.Stdout = &out
	job.cmd.Stderr = &out
	log.Printf("EXECUTE running: %v\n", job.cmd)
	err := job.cmd.Run()
	if err != nil {
		log.Printf("Command Error: cmd: %v cause of %v\n%v\n", job.cmd, err, out.String())
		return err
	}
	return nil
}

func (job *Job) uploadFiles() error {
	localPaths, err := job.listFiles(job.uploads_dir)
	if err != nil {
		return err
	}
	for _, localPath := range localPaths {
		relPath, err := filepath.Rel(job.uploads_dir, localPath)
		if err != nil {
			log.Fatalf("Error getting relative path of %v: %v\n", localPath, err)
			return err
		}
		sep := string([]rune{os.PathSeparator})
		parts := strings.Split(relPath, sep)
		bucket := parts[0]
		object := strings.Join(parts[1:], sep)
		err = job.storage.Upload(bucket, object, localPath)
		if err != nil {
			log.Fatalf("Error uploading %v to gs://%v/%v: %v\n", localPath, bucket, object, err)
			return err
		}
	}
	return nil
}

func (job *Job) listFiles(dir string) ([]string, error) {
	result := []string{}
	err := filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if !info.IsDir() {
			result = append(result, path)
		}
		return nil
	})
	if err != nil {
		log.Fatalf("Error listing upload files: %v\n", err)
		return nil, err
	}
	return result, nil
}

func (job *Job) flatten(obj interface{}) []interface{} {
	// Support only unmarshalled object from JSON
	// See https://golang.org/pkg/encoding/json/#Unmarshal also
	switch obj.(type) {
	case []interface{}:
		res := []interface{}{}
		for _, i := range obj.([]interface{}) {
			switch i.(type) {
			case bool, float64, string, nil:
				res = append(res, i)
			default:
				for _, j := range job.flatten(i) {
					res = append(res, j)
				}
			}
		}
		return res
	case map[string]interface{}:
		values := []interface{}{}
		for _, val := range obj.(map[string]interface{}) {
			values = append(values, val)
		}
		return job.flatten(values)
	default:
		return []interface{}{obj}
	}
}
<commit_msg>:+1: Notify CLEANUP around clearWorkspace<commit_after>package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"net/url"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"strings"

	"golang.org/x/net/context"
)

type (
	CommandConfig struct {
		Template []string            `json:"-"`
		Options  map[string][]string `json:"options,omitempty"`
		Dryrun   bool                `json:"dryrun,omitempty"`
	}

	Job struct {
		config *CommandConfig
		// https://godoc.org/google.golang.org/genproto/googleapis/pubsub/v1#ReceivedMessage
		message      *JobMessage
		notification *ProgressNotification
		storage      Storage

		// These are set at at setupWorkspace
		workspace     string
		downloads_dir string
		uploads_dir   string

		// These are set at setupDownloadFiles
		downloadFileMap     map[string]string
		remoteDownloadFiles interface{}
		localDownloadFiles  interface{}

		cmd *exec.Cmd
	}
)

func (job *Job) run(ctx context.Context) error {
	err := job.runImpl(ctx)
	switch err.(type) {
	case InvalidJobError:
		err := job.withNotify(CANCELLING, job.message.Ack)()
		if err != nil {
			return err
		}
	}
	return err
}

func (job *Job) runImpl(ctx context.Context) error {
	job.notification.notify(PROCESSING, job.message.MessageId(), "info")

	err := job.message.Validate()
	if err != nil {
		log.Printf("Invalid Message: MessageId: %v, Message: %v, error: %v\n", job.message.MessageId(), job.message.raw.Message, err)
		return nil
	}

	defer job.withNotify(CLEANUP, job.clearWorkspace)() // Call clearWorkspace even if setupWorkspace retuns error
	err = job.setupWorkspace()
	if err != nil {
		return err
	}

	err = job.withNotify(PREPARING, job.setupDownloadFiles)()
	if err != nil {
		return err
	}

	err = job.build()
	if err != nil {
		log.Fatalf("Command build Error template: %v msg: %v cause of %v\n", job.config.Template, job.message, err)
		return err
	}

	go job.message.sendMADPeriodically()
	defer job.message.Done()

	err = job.withNotify(DOWNLOADING, job.downloadFiles)()
	if err != nil {
		return err
	}

	err = job.withNotify(EXECUTING, job.execute)()
	if err != nil {
		return err
	}

	err = job.withNotify(UPLOADING, job.uploadFiles)()
	if err != nil {
		return err
	}

	err = job.withNotify(ACKSENDING, job.message.Ack)()
	if err != nil {
		return err
	}

	return err
}

func (job *Job) withNotify(progress int, f func() error) func() error {
	msg_id := job.message.MessageId()
	return func() error {
		job.notification.notify(progress, msg_id, "info")
		err := f()
		if err != nil {
			job.notification.notify(progress+2, msg_id, "error")
			return err
		}
		job.notification.notify(progress+1, msg_id, "info")
		return nil
	}
}

func (job *Job) setupWorkspace() error {
	dir, err := ioutil.TempDir("", "workspace")
	if err != nil {
		log.Fatal(err)
		return err
	}

	subdirs := []string{
		filepath.Join(dir, "downloads"),
		filepath.Join(dir, "uploads"),
	}
	for _, subdir := range subdirs {
		err := os.MkdirAll(subdir, 0700)
		if err != nil {
			return err
		}
	}
	job.workspace = dir
	job.downloads_dir = subdirs[0]
	job.uploads_dir = subdirs[1]
	return nil
}

func (job *Job) clearWorkspace() error {
	if job.workspace != "" {
		return os.RemoveAll(job.workspace)
	}
	return nil
}

func (job *Job) setupDownloadFiles() error {
	job.downloadFileMap = map[string]string{}
	job.remoteDownloadFiles = job.message.DownloadFiles()
	objects := job.flatten(job.remoteDownloadFiles)
	remoteUrls := []string{}
	for _, obj := range objects {
		switch obj.(type) {
		case string:
			remoteUrls = append(remoteUrls, obj.(string))
		default:
			log.Printf("Invalid download file URL: %v [%T]", obj, obj)
		}
	}
	for _, remote_url := range remoteUrls {
		url, err := url.Parse(remote_url)
		if err != nil {
			log.Fatalf("Invalid URL: %v because of %v\n", remote_url, err)
			return err
		}
		urlstr := fmt.Sprintf("gs://%v%v", url.Host, url.Path)
		destPath := filepath.Join(job.downloads_dir, url.Host, url.Path)
		job.downloadFileMap[urlstr] = destPath
	}
	job.localDownloadFiles = job.copyWithFileMap(job.remoteDownloadFiles)
	return nil
}

func (job *Job) copyWithFileMap(obj interface{}) interface{} {
	switch obj.(type) {
	case map[string]interface{}:
		result := map[string]interface{}{}
		for k, v := range obj.(map[string]interface{}) {
			result[k] = job.copyWithFileMap(v)
		}
		return result
	case []interface{}:
		result := []interface{}{}
		for _, v := range obj.([]interface{}) {
			result = append(result, job.copyWithFileMap(v))
		}
		return result
	case string:
		return job.downloadFileMap[obj.(string)]
	default:
		return obj
	}
}

func (job *Job) buildVariable() *Variable {
	return &Variable{
		data: map[string]interface{}{
			"workspace":             job.workspace,
			"downloads_dir":         job.downloads_dir,
			"uploads_dir":           job.uploads_dir,
			"download_files":        job.localDownloadFiles,
			"local_download_files":  job.localDownloadFiles,
			"remote_download_files": job.remoteDownloadFiles,
			"attrs":                 job.message.raw.Message.Attributes,
			"attributes":            job.message.raw.Message.Attributes,
			"data":                  job.message.raw.Message.Data,
		},
	}
}

func (job *Job) build() error {
	v := job.buildVariable()

	values, err := job.extract(v, job.config.Template)
	if err != nil {
		return err
	}
	if len(job.config.Options) > 0 {
		key := strings.Join(values, " ")
		t := job.config.Options[key]
		if t == nil {
			t = job.config.Options["default"]
		}
		if t != nil {
			values, err = job.extract(v, t)
			if err != nil {
				return err
			}
		}
	}
	job.cmd = exec.Command(values[0], values[1:]...)
	return nil
}

func (job *Job) extract(v *Variable, values []string) ([]string, error) {
	result := []string{}
	for _, src := range values {
		extracted, err := v.expand(src)
		if err != nil {
			return nil, &InvalidJobError{err.Error()}
		}
		vals := strings.Split(extracted, v.separator)
		for _, val := range vals {
			result = append(result, val)
		}
	}
	return result, nil
}

func (job *Job) downloadFiles() error {
	for remoteURL, destPath := range job.downloadFileMap {
		url, err := url.Parse(remoteURL)
		if err != nil {
			log.Fatalf("Invalid URL: %v because of %v\n", remoteURL, err)
			return err
		}

		dir := path.Dir(destPath)
		err = os.MkdirAll(dir, 0700)
		if err != nil {
			return err
		}

		err = job.storage.Download(url.Host, url.Path[1:], destPath)
		if err != nil {
			return err
		}
	}
	return nil
}

func (job *Job) execute() error {
	var out bytes.Buffer
	job.cmd.Stdout = &out
	job.cmd.Stderr = &out
	log.Printf("EXECUTE running: %v\n", job.cmd)
	err := job.cmd.Run()
	if err != nil {
		log.Printf("Command Error: cmd: %v cause of %v\n%v\n", job.cmd, err, out.String())
		return err
	}
	return nil
}

func (job *Job) uploadFiles() error {
	localPaths, err := job.listFiles(job.uploads_dir)
	if err != nil {
		return err
	}
	for _, localPath := range localPaths {
		relPath, err := filepath.Rel(job.uploads_dir, localPath)
		if err != nil {
			log.Fatalf("Error getting relative path of %v: %v\n", localPath, err)
			return err
		}
		sep := string([]rune{os.PathSeparator})
		parts := strings.Split(relPath, sep)
		bucket := parts[0]
		object := strings.Join(parts[1:], sep)
		err = job.storage.Upload(bucket, object, localPath)
		if err != nil {
			log.Fatalf("Error uploading %v to gs://%v/%v: %v\n", localPath, bucket, object, err)
			return err
		}
	}
	return nil
}

func (job *Job) listFiles(dir string) ([]string, error) {
	result := []string{}
	err := filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if !info.IsDir() {
			result = append(result, path)
		}
		return nil
	})
	if err != nil {
		log.Fatalf("Error listing upload files: %v\n", err)
		return nil, err
	}
	return result, nil
}

func (job *Job) flatten(obj interface{}) []interface{} {
	// Support only unmarshalled object from JSON
	// See https://golang.org/pkg/encoding/json/#Unmarshal also
	switch obj.(type) {
	case []interface{}:
		res := []interface{}{}
		for _, i := range obj.([]interface{}) {
			switch i.(type) {
			case bool, float64, string, nil:
				res = append(res, i)
			default:
				for _, j := range job.flatten(i) {
					res = append(res, j)
				}
			}
		}
		return res
	case map[string]interface{}:
		values := []interface{}{}
		for _, val := range obj.(map[string]interface{}) {
			values = append(values, val)
		}
		return job.flatten(values)
	default:
		return []interface{}{obj}
	}
}
<|endoftext|>


<commit_before>// +build linux

package sysstat

import (
	"bufio"
	"encoding/csv"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/plugins/inputs"
)

var (
	firstTimestamp time.Time
	execCommand    = exec.Command // execCommand is used to mock commands in tests.
	dfltActivities = []string{"DISK"}
)

const parseInterval = 1 // parseInterval is the interval (in seconds) where the parsing of the binary file takes place.

type Sysstat struct {
	// Sadc represents the path to the sadc collector utility.
	Sadc string `toml:"sadc_path"`

	// Sadf represents the path to the sadf cmd.
	Sadf string `toml:"sadf_path"`

	// Activities is a list of activities that are passed as argument to the
	// collector utility (e.g: DISK, SNMP etc...)
	// The more activities that are added, the more data is collected.
	Activities []string

	// Options is a map of options.
	//
	// The key represents the actual option that the Sadf command is called with and
	// the value represents the description for that option.
	//
	// For example, if you have the following options map:
	//    map[string]string{"-C": "cpu", "-d": "disk"}
	// The Sadf command is run with the options -C and -d to extract cpu and
	// disk metrics from the collected binary file.
	//
	// If Group is false (see below), each metric will be prefixed with the corresponding description
	// and represents itself a measurement.
	//
	// If Group is true, metrics are grouped to a single measurement with the corresponding description as name.
	Options map[string]string

	// Group determines if metrics are grouped or not.
	Group bool

	// DeviceTags adds the possibility to add additional tags for devices.
	DeviceTags map[string][]map[string]string `toml:"device_tags"`
	tmpFile    string
	interval   int
}

func (*Sysstat) Description() string {
	return "Sysstat metrics collector"
}

var sampleConfig = `
  ## Path to the sadc command.
  #
  ## On Debian and Arch Linux the default path is /usr/lib/sa/sadc whereas
  ## on RHEL and CentOS the default path is /usr/lib64/sa/sadc
  sadc_path = "/usr/lib/sa/sadc" # required
  #
  #
  ## Path to the sadf command, if it is not in PATH
  # sadf_path = "/usr/bin/sadf"
  #
  #
  ## Activities is a list of activities, that are passed as argument to the
  ## sadc collector utility (e.g: DISK, SNMP etc...)
  ## The more activities that are added, the more data is collected.
  # activities = ["DISK"]
  #
  #
  ## Group metrics to measurements.
  ##
  ## If group is false each metric will be prefixed with a description
  ## and represents itself a measurement.
  ##
  ## If Group is true, corresponding metrics are grouped to a single measurement.
  # group = true
  #
  #
  ## Options for the sadf command. The values on the left represent the sadf options and
  ## the values on the right their description (wich are used for grouping and prefixing metrics).
  ##
  ## Run 'sar -h' or 'man sar' to find out the supported options for your sysstat version.
  [inputs.sysstat.options]
	-C = "cpu"
	-B = "paging"
	-b = "io"
	-d = "disk"             # requires DISK activity
	"-n ALL" = "network"
	"-P ALL" = "per_cpu"
	-q = "queue"
	-R = "mem"
	-r = "mem_util"
	-S = "swap_util"
	-u = "cpu_util"
	-v = "inode"
	-W = "swap"
	-w = "task"
  #	-H = "hugepages"        # only available for newer linux distributions
  #	"-I ALL" = "interrupts" # requires INT activity
  #
  #
  ## Device tags can be used to add additional tags for devices. For example the configuration below
  ## adds a tag vg with value rootvg for all metrics with sda devices.
  # [[inputs.sysstat.device_tags.sda]]
  #  vg = "rootvg"
`

func (*Sysstat) SampleConfig() string {
	return sampleConfig
}

func (s *Sysstat) Gather(acc telegraf.Accumulator) error {
	if s.interval == 0 {
		if firstTimestamp.IsZero() {
			firstTimestamp = time.Now()
		} else {
			s.interval = int(time.Since(firstTimestamp).Seconds())
		}
	}
	ts := time.Now().Add(time.Duration(s.interval) * time.Second)
	if err := s.collect(); err != nil {
		return err
	}
	var wg sync.WaitGroup
	errorChannel := make(chan error, len(s.Options)*2)
	for option := range s.Options {
		wg.Add(1)
		go func(acc telegraf.Accumulator, option string) {
			defer wg.Done()
			if err := s.parse(acc, option, ts); err != nil {
				errorChannel <- err
			}
		}(acc, option)
	}
	wg.Wait()
	close(errorChannel)

	errorStrings := []string{}
	for err := range errorChannel {
		errorStrings = append(errorStrings, err.Error())
	}

	if _, err := os.Stat(s.tmpFile); err == nil {
		if err := os.Remove(s.tmpFile); err != nil {
			errorStrings = append(errorStrings, err.Error())
		}
	}

	if len(errorStrings) == 0 {
		return nil
	}
	return errors.New(strings.Join(errorStrings, "\n"))
}

// collect collects sysstat data with the collector utility sadc. It runs the following command:
//     Sadc -S <Activity1> -S <Activity2> ... <collectInterval> 2 tmpFile
// The above command collects system metrics during <collectInterval> and saves it in binary form to tmpFile.
func (s *Sysstat) collect() error {
	if len(s.Sadf) == 0 {
		sadf, err := exec.LookPath("sadf")
		if err != nil {
			return errors.New("sadf not in $PATH, configure path to sadf")
		}
		s.Sadf = sadf
	}
	options := []string{}
	for _, act := range s.Activities {
		options = append(options, "-S", act)
	}
	s.tmpFile = path.Join("/tmp", fmt.Sprintf("sysstat-%d", time.Now().Unix()))
	collectInterval := s.interval - parseInterval // collectInterval has to be smaller than the telegraf data collection interval

	if collectInterval < 0 { // If true, interval is not defined yet and Gather is run for the first time.
		collectInterval = 1 // In that case we only collect for 1 second.
	}

	options = append(options, strconv.Itoa(collectInterval), "2", s.tmpFile)
	cmd := execCommand(s.Sadc, options...)
	out, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("failed to run command %s: %s", strings.Join(cmd.Args, " "), string(out))
	}
	return nil
}

// parse runs Sadf on the previously saved tmpFile:
//    Sadf -p -- -p <option> tmpFile
// and parses the output to add it to the telegraf.Accumulator acc.
func (s *Sysstat) parse(acc telegraf.Accumulator, option string, ts time.Time) error {
	cmd := execCommand(s.Sadf, s.sadfOptions(option)...)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return err
	}
	if err := cmd.Start(); err != nil {
		return fmt.Errorf("running command '%s' failed: %s", strings.Join(cmd.Args, " "), err)
	}

	r := bufio.NewReader(stdout)
	csv := csv.NewReader(r)
	csv.Comma = '\t'
	csv.FieldsPerRecord = 6
	var measurement string
	// groupData to accumulate data when Group=true
	type groupData struct {
		tags   map[string]string
		fields map[string]interface{}
	}
	m := make(map[string]groupData)
	for {
		record, err := csv.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}

		device := record[3]
		value, err := strconv.ParseFloat(record[5], 64)
		if err != nil {
			return err
		}

		tags := map[string]string{}
		if device != "-" {
			tags["device"] = device
			if addTags, ok := s.DeviceTags[device]; ok {
				for _, tag := range addTags {
					for k, v := range tag {
						tags[k] = v
					}
				}

			}
		}

		if s.Group {
			measurement = s.Options[option]
			if _, ok := m[device]; !ok {
				m[device] = groupData{
					fields: make(map[string]interface{}),
					tags:   make(map[string]string),
				}
			}
			g, _ := m[device]
			if len(g.tags) == 0 {
				for k, v := range tags {
					g.tags[k] = v
				}
			}
			g.fields[escape(record[4])] = value
		} else {
			measurement = s.Options[option] + "_" + escape(record[4])
			fields := map[string]interface{}{
				"value": value,
			}
			acc.AddFields(measurement, fields, tags, ts)
		}

	}
	if s.Group {
		for _, v := range m {
			acc.AddFields(measurement, v.fields, v.tags, ts)
		}
	}
	if err := cmd.Wait(); err != nil {
		return fmt.Errorf("command %s failed with %s", strings.Join(cmd.Args, " "), err)
	}
	return nil
}

// sadfOptions creates the correct options for the sadf utility.
func (s *Sysstat) sadfOptions(activityOption string) []string {
	options := []string{
		"-p",
		"--",
		"-p",
	}

	opts := strings.Split(activityOption, " ")
	options = append(options, opts...)
	options = append(options, s.tmpFile)

	return options
}

// escape removes % and / chars in field names
func escape(dirty string) string {
	var fieldEscaper = strings.NewReplacer(
		`%`, "pct_",
		`/`, "_per_",
	)
	return fieldEscaper.Replace(dirty)
}

func init() {
	inputs.Add("sysstat", func() telegraf.Input {
		return &Sysstat{
			Group:      true,
			Activities: dfltActivities,
		}
	})
}
<commit_msg>move pathe lookup for sadf to init()<commit_after>// +build linux

package sysstat

import (
	"bufio"
	"encoding/csv"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/plugins/inputs"
)

var (
	firstTimestamp time.Time
	execCommand    = exec.Command // execCommand is used to mock commands in tests.
	dfltActivities = []string{"DISK"}
)

const parseInterval = 1 // parseInterval is the interval (in seconds) where the parsing of the binary file takes place.

type Sysstat struct {
	// Sadc represents the path to the sadc collector utility.
	Sadc string `toml:"sadc_path"`

	// Sadf represents the path to the sadf cmd.
	Sadf string `toml:"sadf_path"`

	// Activities is a list of activities that are passed as argument to the
	// collector utility (e.g: DISK, SNMP etc...)
	// The more activities that are added, the more data is collected.
	Activities []string

	// Options is a map of options.
	//
	// The key represents the actual option that the Sadf command is called with and
	// the value represents the description for that option.
	//
	// For example, if you have the following options map:
	//    map[string]string{"-C": "cpu", "-d": "disk"}
	// The Sadf command is run with the options -C and -d to extract cpu and
	// disk metrics from the collected binary file.
	//
	// If Group is false (see below), each metric will be prefixed with the corresponding description
	// and represents itself a measurement.
	//
	// If Group is true, metrics are grouped to a single measurement with the corresponding description as name.
	Options map[string]string

	// Group determines if metrics are grouped or not.
	Group bool

	// DeviceTags adds the possibility to add additional tags for devices.
	DeviceTags map[string][]map[string]string `toml:"device_tags"`
	tmpFile    string
	interval   int
}

func (*Sysstat) Description() string {
	return "Sysstat metrics collector"
}

var sampleConfig = `
  ## Path to the sadc command.
  #
  ## On Debian and Arch Linux the default path is /usr/lib/sa/sadc whereas
  ## on RHEL and CentOS the default path is /usr/lib64/sa/sadc
  sadc_path = "/usr/lib/sa/sadc" # required
  #
  #
  ## Path to the sadf command, if it is not in PATH
  # sadf_path = "/usr/bin/sadf"
  #
  #
  ## Activities is a list of activities, that are passed as argument to the
  ## sadc collector utility (e.g: DISK, SNMP etc...)
  ## The more activities that are added, the more data is collected.
  # activities = ["DISK"]
  #
  #
  ## Group metrics to measurements.
  ##
  ## If group is false each metric will be prefixed with a description
  ## and represents itself a measurement.
  ##
  ## If Group is true, corresponding metrics are grouped to a single measurement.
  # group = true
  #
  #
  ## Options for the sadf command. The values on the left represent the sadf options and
  ## the values on the right their description (wich are used for grouping and prefixing metrics).
  ##
  ## Run 'sar -h' or 'man sar' to find out the supported options for your sysstat version.
  [inputs.sysstat.options]
	-C = "cpu"
	-B = "paging"
	-b = "io"
	-d = "disk"             # requires DISK activity
	"-n ALL" = "network"
	"-P ALL" = "per_cpu"
	-q = "queue"
	-R = "mem"
	-r = "mem_util"
	-S = "swap_util"
	-u = "cpu_util"
	-v = "inode"
	-W = "swap"
	-w = "task"
  #	-H = "hugepages"        # only available for newer linux distributions
  #	"-I ALL" = "interrupts" # requires INT activity
  #
  #
  ## Device tags can be used to add additional tags for devices. For example the configuration below
  ## adds a tag vg with value rootvg for all metrics with sda devices.
  # [[inputs.sysstat.device_tags.sda]]
  #  vg = "rootvg"
`

func (*Sysstat) SampleConfig() string {
	return sampleConfig
}

func (s *Sysstat) Gather(acc telegraf.Accumulator) error {
	if s.interval == 0 {
		if firstTimestamp.IsZero() {
			firstTimestamp = time.Now()
		} else {
			s.interval = int(time.Since(firstTimestamp).Seconds())
		}
	}
	ts := time.Now().Add(time.Duration(s.interval) * time.Second)
	if err := s.collect(); err != nil {
		return err
	}
	var wg sync.WaitGroup
	errorChannel := make(chan error, len(s.Options)*2)
	for option := range s.Options {
		wg.Add(1)
		go func(acc telegraf.Accumulator, option string) {
			defer wg.Done()
			if err := s.parse(acc, option, ts); err != nil {
				errorChannel <- err
			}
		}(acc, option)
	}
	wg.Wait()
	close(errorChannel)

	errorStrings := []string{}
	for err := range errorChannel {
		errorStrings = append(errorStrings, err.Error())
	}

	if _, err := os.Stat(s.tmpFile); err == nil {
		if err := os.Remove(s.tmpFile); err != nil {
			errorStrings = append(errorStrings, err.Error())
		}
	}

	if len(errorStrings) == 0 {
		return nil
	}
	return errors.New(strings.Join(errorStrings, "\n"))
}

// collect collects sysstat data with the collector utility sadc. It runs the following command:
//     Sadc -S <Activity1> -S <Activity2> ... <collectInterval> 2 tmpFile
// The above command collects system metrics during <collectInterval> and saves it in binary form to tmpFile.
func (s *Sysstat) collect() error {
	options := []string{}
	for _, act := range s.Activities {
		options = append(options, "-S", act)
	}
	s.tmpFile = path.Join("/tmp", fmt.Sprintf("sysstat-%d", time.Now().Unix()))
	collectInterval := s.interval - parseInterval // collectInterval has to be smaller than the telegraf data collection interval

	if collectInterval < 0 { // If true, interval is not defined yet and Gather is run for the first time.
		collectInterval = 1 // In that case we only collect for 1 second.
	}

	options = append(options, strconv.Itoa(collectInterval), "2", s.tmpFile)
	cmd := execCommand(s.Sadc, options...)
	out, err := cmd.CombinedOutput()
	if err != nil {
		return fmt.Errorf("failed to run command %s: %s", strings.Join(cmd.Args, " "), string(out))
	}
	return nil
}

// parse runs Sadf on the previously saved tmpFile:
//    Sadf -p -- -p <option> tmpFile
// and parses the output to add it to the telegraf.Accumulator acc.
func (s *Sysstat) parse(acc telegraf.Accumulator, option string, ts time.Time) error {
	cmd := execCommand(s.Sadf, s.sadfOptions(option)...)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return err
	}
	if err := cmd.Start(); err != nil {
		return fmt.Errorf("running command '%s' failed: %s", strings.Join(cmd.Args, " "), err)
	}

	r := bufio.NewReader(stdout)
	csv := csv.NewReader(r)
	csv.Comma = '\t'
	csv.FieldsPerRecord = 6
	var measurement string
	// groupData to accumulate data when Group=true
	type groupData struct {
		tags   map[string]string
		fields map[string]interface{}
	}
	m := make(map[string]groupData)
	for {
		record, err := csv.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}

		device := record[3]
		value, err := strconv.ParseFloat(record[5], 64)
		if err != nil {
			return err
		}

		tags := map[string]string{}
		if device != "-" {
			tags["device"] = device
			if addTags, ok := s.DeviceTags[device]; ok {
				for _, tag := range addTags {
					for k, v := range tag {
						tags[k] = v
					}
				}

			}
		}

		if s.Group {
			measurement = s.Options[option]
			if _, ok := m[device]; !ok {
				m[device] = groupData{
					fields: make(map[string]interface{}),
					tags:   make(map[string]string),
				}
			}
			g, _ := m[device]
			if len(g.tags) == 0 {
				for k, v := range tags {
					g.tags[k] = v
				}
			}
			g.fields[escape(record[4])] = value
		} else {
			measurement = s.Options[option] + "_" + escape(record[4])
			fields := map[string]interface{}{
				"value": value,
			}
			acc.AddFields(measurement, fields, tags, ts)
		}

	}
	if s.Group {
		for _, v := range m {
			acc.AddFields(measurement, v.fields, v.tags, ts)
		}
	}
	if err := cmd.Wait(); err != nil {
		return fmt.Errorf("command %s failed with %s", strings.Join(cmd.Args, " "), err)
	}
	return nil
}

// sadfOptions creates the correct options for the sadf utility.
func (s *Sysstat) sadfOptions(activityOption string) []string {
	options := []string{
		"-p",
		"--",
		"-p",
	}

	opts := strings.Split(activityOption, " ")
	options = append(options, opts...)
	options = append(options, s.tmpFile)

	return options
}

// escape removes % and / chars in field names
func escape(dirty string) string {
	var fieldEscaper = strings.NewReplacer(
		`%`, "pct_",
		`/`, "_per_",
	)
	return fieldEscaper.Replace(dirty)
}

func init() {
	s := Sysstat{
		Group:      true,
		Activities: dfltActivities,
	}
	sadf, _ := exec.LookPath("sadf")
	if len(sadf) > 0 {
		s.Sadf = sadf
	}
	inputs.Add("sysstat", func() telegraf.Input {
		return &s
	})
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/url"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"regexp"
	"strings"

	"golang.org/x/net/context"

	pubsub "google.golang.org/api/pubsub/v1"
)

type (
	CommandConfig struct {
		Template []string            `json:"-"`
		Options  map[string][]string `json:"options,omitempty"`
		Dryrun   bool                `json:"dryrun,omitempty"`
	}

	Job struct {
		config *CommandConfig
		// https://godoc.org/google.golang.org/genproto/googleapis/pubsub/v1#ReceivedMessage
		message      *pubsub.ReceivedMessage
		notification *ProgressNotification
		storage      Storage

		// These are set at at setupWorkspace
		workspace     string
		downloads_dir string
		uploads_dir   string

		// These are set at setupDownloadFiles
		downloadFileMap     map[string]string
		remoteDownloadFiles interface{}
		localDownloadFiles  interface{}
	}
)

func (job *Job) run(ctx context.Context) error {
	job.notification.notify(PROCESSING, job.message.Message.MessageId, "info")
	err := job.setupWorkspace(ctx, func() error {
		err := job.withNotify(PREPARING, job.setupDownloadFiles)()
		if err != nil {
			return err
		}

		err = job.withNotify(DOWNLOADING, job.downloadFiles)()
		if err != nil {
			return err
		}

		err = job.withNotify(EXECUTING, job.execute)()
		if err != nil {
			return err
		}

		err = job.withNotify(UPLOADING, job.uploadFiles)()
		if err != nil {
			return err
		}

		return nil
	})
	job.notification.notify(CLEANUP, job.message.Message.MessageId, "info")
	return err
}

func (job *Job) withNotify(progress int, f func() error) func() error {
	msg_id := job.message.Message.MessageId
	return func() error {
		job.notification.notify(progress, msg_id, "info")
		err := f()
		if err != nil {
			job.notification.notify(progress+2, msg_id, "error")
			return err
		}
		job.notification.notify(progress+1, msg_id, "info")
		return nil
	}
}

func (job *Job) setupWorkspace(ctx context.Context, f func() error) error {
	dir, err := ioutil.TempDir("", "workspace")
	if err != nil {
		log.Fatal(err)
		return err
	}
	defer os.RemoveAll(dir) // clean up

	subdirs := []string{
		filepath.Join(dir, "downloads"),
		filepath.Join(dir, "uploads"),
	}
	for _, subdir := range subdirs {
		err := os.MkdirAll(subdir, 0700)
		if err != nil {
			return err
		}
	}
	job.workspace = dir
	job.downloads_dir = subdirs[0]
	job.uploads_dir = subdirs[1]
	return f()
}

func (job *Job) setupDownloadFiles() error {
	job.downloadFileMap = map[string]string{}
	job.remoteDownloadFiles = job.parseJson(job.message.Message.Attributes["download_files"])
	objects := job.flatten(job.remoteDownloadFiles)
	remoteUrls := []string{}
	for _, obj := range objects {
		switch obj.(type) {
		case string:
			remoteUrls = append(remoteUrls, obj.(string))
		default:
			log.Printf("Invalid download file URL: %v [%T]", obj, obj)
		}
	}
	for _, remote_url := range remoteUrls {
		url, err := url.Parse(remote_url)
		if err != nil {
			log.Fatalf("Invalid URL: %v because of %v\n", remote_url, err)
			return err
		}
		urlstr := fmt.Sprintf("gs://%v%v", url.Host, url.Path)
		destPath := filepath.Join(job.downloads_dir, url.Host, url.Path)
		job.downloadFileMap[urlstr] = destPath
	}
	job.localDownloadFiles = job.copyWithFileMap(job.remoteDownloadFiles)
	return nil
}

func (job *Job) copyWithFileMap(obj interface{}) interface{} {
	switch obj.(type) {
	case map[string]interface{}:
		result := map[string]interface{}{}
		for k, v := range obj.(map[string]interface{}) {
			result[k] = job.copyWithFileMap(v)
		}
		return result
	case []interface{}:
		result := []interface{}{}
		for _, v := range obj.([]interface{}) {
			result = append(result, job.copyWithFileMap(v))
		}
		return result
	case string:
		return job.downloadFileMap[obj.(string)]
	default:
		return obj
	}
}

func (job *Job) buildVariable() *Variable {
	return &Variable{
		data: map[string]interface{}{
			"workspace":             job.workspace,
			"downloads_dir":         job.downloads_dir,
			"uploads_dir":           job.uploads_dir,
			"download_files":        job.localDownloadFiles,
			"local_download_files":  job.localDownloadFiles,
			"remote_download_files": job.remoteDownloadFiles,
			"attrs":                 job.message.Message.Attributes,
			"attributes":            job.message.Message.Attributes,
			"data":                  job.message.Message.Data,
		},
	}
}

func (job *Job) build() (*exec.Cmd, error) {
	v := job.buildVariable()

	values, err := job.extract(v, job.config.Template)
	if err != nil {
		return nil, err
	}
	if len(job.config.Options) > 0 {
		key := strings.Join(values, " ")
		t := job.config.Options[key]
		if t == nil {
			t = job.config.Options["default"]
		}
		if t != nil {
			values, err = job.extract(v, t)
			if err != nil {
				return nil, err
			}
		}
	}
	cmd := exec.Command(values[0], values[1:]...)
	return cmd, nil
}

func (job *Job) extract(v *Variable, values []string) ([]string, error) {
	result := []string{}
	for _, src := range values {
		extracted, err := v.expand(src)
		if err != nil {
			return nil, err
		}
		vals := strings.Split(extracted, v.separator)
		for _, val := range vals {
			result = append(result, val)
		}
	}
	return result, nil
}

func (job *Job) downloadFiles() error {
	for remoteURL, destPath := range job.downloadFileMap {
		url, err := url.Parse(remoteURL)
		if err != nil {
			log.Fatalf("Invalid URL: %v because of %v\n", remoteURL, err)
			return err
		}

		dir := path.Dir(destPath)
		err = os.MkdirAll(dir, 0700)
		if err != nil {
			return err
		}

		err = job.storage.Download(url.Host, url.Path[1:], destPath)
		if err != nil {
			return err
		}
	}
	return nil
}

func (job *Job) execute() error {
	cmd, err := job.build()
	if err != nil {
		log.Fatalf("Command build Error template: %v msg: %v cause of %v\n", job.config.Template, job.message, err)
		return err
	}
	err = cmd.Run()
	if err != nil {
		log.Printf("Command Error: cmd: %v cause of %v\n", cmd, err)
		// return err // Don't return this err
	}
	return nil
}

func (job *Job) uploadFiles() error {
	localPaths, err := job.listFiles(job.uploads_dir)
	if err != nil {
		return err
	}
	for _, localPath := range localPaths {
		relPath, err := filepath.Rel(job.uploads_dir, localPath)
		if err != nil {
			log.Fatalf("Error getting relative path of %v: %v\n", localPath, err)
			return err
		}
		sep := string([]rune{os.PathSeparator})
		parts := strings.Split(relPath, sep)
		bucket := parts[0]
		object := strings.Join(parts[1:], sep)
		err = job.storage.Upload(bucket, object, localPath)
		if err != nil {
			log.Fatalf("Error uploading %v to gs://%v/%v: %v\n", localPath, bucket, object, err)
			return err
		}
	}
	return nil
}

func (job *Job) listFiles(dir string) ([]string, error) {
	result := []string{}
	err := filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if !info.IsDir() {
			result = append(result, path)
		}
		return nil
	})
	if err != nil {
		log.Fatalf("Error listing upload files: %v\n", err)
		return nil, err
	}
	return result, nil
}

func (job *Job) parseJson(str string) interface{} {
	matched, err := regexp.MatchString(`\A\[.*\]\z|\A\{.*\}\z|`, str)
	if err != nil {
		return str
	}
	if !matched {
		return str
	}
	var dest interface{}
	err = json.Unmarshal([]byte(str), &dest)
	if err != nil {
		return str
	}
	return dest
}

func (job *Job) flatten(obj interface{}) []interface{} {
	// Support only unmarshalled object from JSON
	// See https://golang.org/pkg/encoding/json/#Unmarshal also
	switch obj.(type) {
	case []interface{}:
		res := []interface{}{}
		for _, i := range obj.([]interface{}) {
			switch i.(type) {
			case bool, float64, string, nil:
				res = append(res, i)
			default:
				for _, j := range job.flatten(i) {
					res = append(res, j)
				}
			}
		}
		return res
	case map[string]interface{}:
		values := []interface{}{}
		for _, val := range obj.(map[string]interface{}) {
			values = append(values, val)
		}
		return job.flatten(values)
	default:
		return []interface{}{obj}
	}
}
<commit_msg>:+1: Output stdout and stderr to log if command returns error<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/url"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"regexp"
	"strings"

	"golang.org/x/net/context"

	pubsub "google.golang.org/api/pubsub/v1"
)

type (
	CommandConfig struct {
		Template []string            `json:"-"`
		Options  map[string][]string `json:"options,omitempty"`
		Dryrun   bool                `json:"dryrun,omitempty"`
	}

	Job struct {
		config *CommandConfig
		// https://godoc.org/google.golang.org/genproto/googleapis/pubsub/v1#ReceivedMessage
		message      *pubsub.ReceivedMessage
		notification *ProgressNotification
		storage      Storage

		// These are set at at setupWorkspace
		workspace     string
		downloads_dir string
		uploads_dir   string

		// These are set at setupDownloadFiles
		downloadFileMap     map[string]string
		remoteDownloadFiles interface{}
		localDownloadFiles  interface{}
	}
)

func (job *Job) run(ctx context.Context) error {
	job.notification.notify(PROCESSING, job.message.Message.MessageId, "info")
	err := job.setupWorkspace(ctx, func() error {
		err := job.withNotify(PREPARING, job.setupDownloadFiles)()
		if err != nil {
			return err
		}

		err = job.withNotify(DOWNLOADING, job.downloadFiles)()
		if err != nil {
			return err
		}

		err = job.withNotify(EXECUTING, job.execute)()
		if err != nil {
			return err
		}

		err = job.withNotify(UPLOADING, job.uploadFiles)()
		if err != nil {
			return err
		}

		return nil
	})
	job.notification.notify(CLEANUP, job.message.Message.MessageId, "info")
	return err
}

func (job *Job) withNotify(progress int, f func() error) func() error {
	msg_id := job.message.Message.MessageId
	return func() error {
		job.notification.notify(progress, msg_id, "info")
		err := f()
		if err != nil {
			job.notification.notify(progress+2, msg_id, "error")
			return err
		}
		job.notification.notify(progress+1, msg_id, "info")
		return nil
	}
}

func (job *Job) setupWorkspace(ctx context.Context, f func() error) error {
	dir, err := ioutil.TempDir("", "workspace")
	if err != nil {
		log.Fatal(err)
		return err
	}
	defer os.RemoveAll(dir) // clean up

	subdirs := []string{
		filepath.Join(dir, "downloads"),
		filepath.Join(dir, "uploads"),
	}
	for _, subdir := range subdirs {
		err := os.MkdirAll(subdir, 0700)
		if err != nil {
			return err
		}
	}
	job.workspace = dir
	job.downloads_dir = subdirs[0]
	job.uploads_dir = subdirs[1]
	return f()
}

func (job *Job) setupDownloadFiles() error {
	job.downloadFileMap = map[string]string{}
	job.remoteDownloadFiles = job.parseJson(job.message.Message.Attributes["download_files"])
	objects := job.flatten(job.remoteDownloadFiles)
	remoteUrls := []string{}
	for _, obj := range objects {
		switch obj.(type) {
		case string:
			remoteUrls = append(remoteUrls, obj.(string))
		default:
			log.Printf("Invalid download file URL: %v [%T]", obj, obj)
		}
	}
	for _, remote_url := range remoteUrls {
		url, err := url.Parse(remote_url)
		if err != nil {
			log.Fatalf("Invalid URL: %v because of %v\n", remote_url, err)
			return err
		}
		urlstr := fmt.Sprintf("gs://%v%v", url.Host, url.Path)
		destPath := filepath.Join(job.downloads_dir, url.Host, url.Path)
		job.downloadFileMap[urlstr] = destPath
	}
	job.localDownloadFiles = job.copyWithFileMap(job.remoteDownloadFiles)
	return nil
}

func (job *Job) copyWithFileMap(obj interface{}) interface{} {
	switch obj.(type) {
	case map[string]interface{}:
		result := map[string]interface{}{}
		for k, v := range obj.(map[string]interface{}) {
			result[k] = job.copyWithFileMap(v)
		}
		return result
	case []interface{}:
		result := []interface{}{}
		for _, v := range obj.([]interface{}) {
			result = append(result, job.copyWithFileMap(v))
		}
		return result
	case string:
		return job.downloadFileMap[obj.(string)]
	default:
		return obj
	}
}

func (job *Job) buildVariable() *Variable {
	return &Variable{
		data: map[string]interface{}{
			"workspace":             job.workspace,
			"downloads_dir":         job.downloads_dir,
			"uploads_dir":           job.uploads_dir,
			"download_files":        job.localDownloadFiles,
			"local_download_files":  job.localDownloadFiles,
			"remote_download_files": job.remoteDownloadFiles,
			"attrs":                 job.message.Message.Attributes,
			"attributes":            job.message.Message.Attributes,
			"data":                  job.message.Message.Data,
		},
	}
}

func (job *Job) build() (*exec.Cmd, error) {
	v := job.buildVariable()

	values, err := job.extract(v, job.config.Template)
	if err != nil {
		return nil, err
	}
	if len(job.config.Options) > 0 {
		key := strings.Join(values, " ")
		t := job.config.Options[key]
		if t == nil {
			t = job.config.Options["default"]
		}
		if t != nil {
			values, err = job.extract(v, t)
			if err != nil {
				return nil, err
			}
		}
	}
	cmd := exec.Command(values[0], values[1:]...)
	return cmd, nil
}

func (job *Job) extract(v *Variable, values []string) ([]string, error) {
	result := []string{}
	for _, src := range values {
		extracted, err := v.expand(src)
		if err != nil {
			return nil, err
		}
		vals := strings.Split(extracted, v.separator)
		for _, val := range vals {
			result = append(result, val)
		}
	}
	return result, nil
}

func (job *Job) downloadFiles() error {
	for remoteURL, destPath := range job.downloadFileMap {
		url, err := url.Parse(remoteURL)
		if err != nil {
			log.Fatalf("Invalid URL: %v because of %v\n", remoteURL, err)
			return err
		}

		dir := path.Dir(destPath)
		err = os.MkdirAll(dir, 0700)
		if err != nil {
			return err
		}

		err = job.storage.Download(url.Host, url.Path[1:], destPath)
		if err != nil {
			return err
		}
	}
	return nil
}

func (job *Job) execute() error {
	cmd, err := job.build()
	if err != nil {
		log.Fatalf("Command build Error template: %v msg: %v cause of %v\n", job.config.Template, job.message, err)
		return err
	}
	var out bytes.Buffer
	cmd.Stdout = &out
	cmd.Stderr = &out
	log.Printf("EXECUTE running: %v\n", cmd)
	err = cmd.Run()
	if err != nil {
		log.Printf("Command Error: cmd: %v cause of %v\n%v\n", cmd, err, out.String())
		return err
	}
	return nil
}

func (job *Job) uploadFiles() error {
	localPaths, err := job.listFiles(job.uploads_dir)
	if err != nil {
		return err
	}
	for _, localPath := range localPaths {
		relPath, err := filepath.Rel(job.uploads_dir, localPath)
		if err != nil {
			log.Fatalf("Error getting relative path of %v: %v\n", localPath, err)
			return err
		}
		sep := string([]rune{os.PathSeparator})
		parts := strings.Split(relPath, sep)
		bucket := parts[0]
		object := strings.Join(parts[1:], sep)
		err = job.storage.Upload(bucket, object, localPath)
		if err != nil {
			log.Fatalf("Error uploading %v to gs://%v/%v: %v\n", localPath, bucket, object, err)
			return err
		}
	}
	return nil
}

func (job *Job) listFiles(dir string) ([]string, error) {
	result := []string{}
	err := filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if !info.IsDir() {
			result = append(result, path)
		}
		return nil
	})
	if err != nil {
		log.Fatalf("Error listing upload files: %v\n", err)
		return nil, err
	}
	return result, nil
}

func (job *Job) parseJson(str string) interface{} {
	matched, err := regexp.MatchString(`\A\[.*\]\z|\A\{.*\}\z|`, str)
	if err != nil {
		return str
	}
	if !matched {
		return str
	}
	var dest interface{}
	err = json.Unmarshal([]byte(str), &dest)
	if err != nil {
		return str
	}
	return dest
}

func (job *Job) flatten(obj interface{}) []interface{} {
	// Support only unmarshalled object from JSON
	// See https://golang.org/pkg/encoding/json/#Unmarshal also
	switch obj.(type) {
	case []interface{}:
		res := []interface{}{}
		for _, i := range obj.([]interface{}) {
			switch i.(type) {
			case bool, float64, string, nil:
				res = append(res, i)
			default:
				for _, j := range job.flatten(i) {
					res = append(res, j)
				}
			}
		}
		return res
	case map[string]interface{}:
		values := []interface{}{}
		for _, val := range obj.(map[string]interface{}) {
			values = append(values, val)
		}
		return job.flatten(values)
	default:
		return []interface{}{obj}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/DanielMorsing/gonzbee/nntp"
	"github.com/DanielMorsing/gonzbee/nzb"
	"github.com/DanielMorsing/gonzbee/yenc"
	"io"
	"log"
	"os"
	"path/filepath"
	"sync"
)

type job struct {
	dir string
	n   *nzb.Nzb
}

var downloaderRq = make(chan chan *nntp.Conn)
var downloadReaper = make(chan *nntp.Conn, 1024)

func init() {
	go connectionHandler()
}

func spinUp() *nntp.Conn {
	str := config.GetAddressStr()
	var conn *nntp.Conn
	var err error
	if config.TLS {
		conn, err = nntp.DialTLS(str)
	} else {
		conn, err = nntp.Dial(str)
	}
	if err != nil {
		return nil
	}
	err = conn.Authenticate(config.Username, config.Password)
	if err != nil {
		return nil
	}
	return conn
}

func connectionHandler() {
	var number int
	for {
		ch := <-downloaderRq
		var conn *nntp.Conn
		if number < 10 {
			conn = spinUp()
			if conn == nil {
				continue
			}
			log.Print("Spun up connection #", number+1)
			number++
			ch <- conn
			continue
		}
		conn = <-downloadReaper
		ch <- conn
	}
}

type offsetWriter struct {
	offset int64
	io.WriterAt
}

func (ow *offsetWriter) Write(b []byte) (n int, err error) {
	n, err = ow.WriteAt(b, ow.offset)
	ow.offset += int64(n)
	return
}

func (j *job) handle() {
	var jobDone sync.WaitGroup
	jobDone.Add(len(j.n.File))
	for _, f := range j.n.File {
		var fileinit sync.Once
		var file *os.File
		var fileClose sync.WaitGroup
		fileClose.Add(len(f.Segments))
		go func() {
			fileClose.Wait()
			if file != nil {
				file.Close()
			}
			jobDone.Done()
		}()
		for _, s := range f.Segments {
			ch := make(chan *nntp.Conn)
			downloaderRq <- ch
			conn := <-ch
			go func(seg nzb.Segment, f nzb.File) {
				defer fileClose.Done()
				defer func() {
					downloadReaper <- conn
				}()

				err := conn.SwitchGroup(f.Groups[0])
				if err != nil {
					log.Println("Could not switch to group:", err.Error())
					return
				}

				reader, err := conn.GetMessageReader(seg.MsgId)
				if err != nil {
					log.Printf("Could Download MsgId \"%s\": %s\n", seg.MsgId, err.Error())
					return
				}
				defer reader.Close()

				part, err := yenc.NewPart(reader)
				if err != nil {
					log.Printf("Could Decode MsgId \"%s\": %s\n", seg.MsgId, err.Error())
					return
				}

				fileinit.Do(func() {
					file, err = os.Create(filepath.Join(j.dir, part.Filename))
					if err != nil {
						os.Exit(1)
					}
					go func() {
						fileClose.Wait()
						log.Printf("Done downloading file \"%s\"\n", part.Filename)
					}()
				})

				ow := &offsetWriter{
					offset:   part.Begin,
					WriterAt: file,
				}
				_, err = io.Copy(ow, part)
				if err != nil {
					log.Println("Error getting segment:", err.Error())
				}
			}(s, f)
		}
	}
	jobDone.Wait()
}

func existDir(path string) bool {
	fi, err := os.Stat(path)
	if err != nil {
		return false
	}
	return fi.IsDir()
}

func jobStart(n *nzb.Nzb, name string, dir string) error {
	workDir := filepath.Join(dir, name)
	if !existDir(workDir) {
		err := os.Mkdir(workDir, 0777)
		if err != nil {
			return err
		}
	}
	j := &job{
		dir: workDir,
		n:   n,
	}
	j.handle()
	return nil
}
<commit_msg>make dir creation more robust.<commit_after>package main

import (
	"github.com/DanielMorsing/gonzbee/nntp"
	"github.com/DanielMorsing/gonzbee/nzb"
	"github.com/DanielMorsing/gonzbee/yenc"
	"io"
	"log"
	"os"
	"path/filepath"
	"sync"
)

type job struct {
	dir string
	n   *nzb.Nzb
}

var downloaderRq = make(chan chan *nntp.Conn)
var downloadReaper = make(chan *nntp.Conn, 1024)

func init() {
	go connectionHandler()
}

func spinUp() *nntp.Conn {
	str := config.GetAddressStr()
	var conn *nntp.Conn
	var err error
	if config.TLS {
		conn, err = nntp.DialTLS(str)
	} else {
		conn, err = nntp.Dial(str)
	}
	if err != nil {
		return nil
	}
	err = conn.Authenticate(config.Username, config.Password)
	if err != nil {
		return nil
	}
	return conn
}

func connectionHandler() {
	var number int
	for {
		ch := <-downloaderRq
		var conn *nntp.Conn
		if number < 10 {
			conn = spinUp()
			if conn == nil {
				continue
			}
			log.Print("Spun up connection #", number+1)
			number++
			ch <- conn
			continue
		}
		conn = <-downloadReaper
		ch <- conn
	}
}

type offsetWriter struct {
	offset int64
	io.WriterAt
}

func (ow *offsetWriter) Write(b []byte) (n int, err error) {
	n, err = ow.WriteAt(b, ow.offset)
	ow.offset += int64(n)
	return
}

func (j *job) handle() {
	var jobDone sync.WaitGroup
	jobDone.Add(len(j.n.File))
	for _, f := range j.n.File {
		var fileinit sync.Once
		var file *os.File
		var fileClose sync.WaitGroup
		fileClose.Add(len(f.Segments))
		go func() {
			fileClose.Wait()
			if file != nil {
				file.Close()
			}
			jobDone.Done()
		}()
		for _, s := range f.Segments {
			ch := make(chan *nntp.Conn)
			downloaderRq <- ch
			conn := <-ch
			go func(seg nzb.Segment, f nzb.File) {
				defer fileClose.Done()
				defer func() {
					downloadReaper <- conn
				}()

				err := conn.SwitchGroup(f.Groups[0])
				if err != nil {
					log.Println("Could not switch to group:", err.Error())
					return
				}

				reader, err := conn.GetMessageReader(seg.MsgId)
				if err != nil {
					log.Printf("Could Download MsgId \"%s\": %s\n", seg.MsgId, err.Error())
					return
				}
				defer reader.Close()

				part, err := yenc.NewPart(reader)
				if err != nil {
					log.Printf("Could Decode MsgId \"%s\": %s\n", seg.MsgId, err.Error())
					return
				}

				fileinit.Do(func() {
					file, err = os.Create(filepath.Join(j.dir, part.Filename))
					if err != nil {
						os.Exit(1)
					}
					go func() {
						fileClose.Wait()
						log.Printf("Done downloading file \"%s\"\n", part.Filename)
					}()
				})

				ow := &offsetWriter{
					offset:   part.Begin,
					WriterAt: file,
				}
				_, err = io.Copy(ow, part)
				if err != nil {
					log.Println("Error getting segment:", err.Error())
				}
			}(s, f)
		}
	}
	jobDone.Wait()
}

func mkdir(path string) error {
	fi, err := os.Stat(path)
	if err != nil {
		if os.IsNotExist(err) {
			return os.Mkdir(path, 0777)
		}
		return err
	}
	if fi.IsDir() {
		return nil
	} else {
		return os.ErrExist
	}
	panic("unreachable")
}

func jobStart(n *nzb.Nzb, name string, dir string) error {
	workDir := filepath.Join(dir, name)
	err := mkdir(workDir)
	if err != nil {
		return err
	}
	j := &job{
		dir: workDir,
		n:   n,
	}
	j.handle()
	return nil
}
<|endoftext|>


<commit_before>package clideployment

import (
	"github.com/containerum/chkit/pkg/chkitErrors"
	"github.com/containerum/chkit/pkg/context"
	"github.com/containerum/chkit/pkg/export"
	"github.com/containerum/chkit/pkg/model"
	"github.com/containerum/chkit/pkg/model/deployment"
	"github.com/containerum/chkit/pkg/util/ferr"
	"github.com/ninedraft/boxofstuff/strset"
	"github.com/octago/sflags/gen/gpflag"
	"github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
)

const (
	ErrNoNamespaceSpecified chkitErrors.Err = "no namespace specified"
)

var aliases = []string{"depl", "deployments", "deploy", "de", "dpl", "depls", "dep"}

func Get(ctx *context.Context) *cobra.Command {
	var flags struct {
		File     string `desc:"output file, STDOUT by default"`
		Output   string `flag:"output o" desc:"output format, json/yaml"`
		Solution string
	}
	command := &cobra.Command{
		Use:     "deployment",
		Short:   "show deployment data",
		Long:    "Print deployment data.",
		Example: "namespace deployment_names... [-n namespace_label]",
		Aliases: aliases,
		Run: func(command *cobra.Command, args []string) {
			var logger = ctx.Log.Command("get deployment")
			logger.Debugf("START")
			defer logrus.Debugf("END")
			var deplData model.Renderer

			switch len(args) {
			case 1:
				if flags.Solution != "" {
					logger.Debugf("getting deployment list from namespace %q and solution %q",
						ctx.GetNamespace(), flags.Solution)
					var depls, err = ctx.Client.GetSolutionDeployments(ctx.GetNamespace().ID, flags.Solution)
					if err != nil {
						logger.WithError(err).Errorf("unable to get deployment list from namespace %q and solution %q", ctx.GetNamespace(), flags.Solution)
						ferr.Println(err)
						ctx.Exit(1)
					}
					var depl, ok = depls.GetByName(args[0])
					if !ok {
						ferr.Printf("deployment %q not found\n", args[0])
					}
					deplData = depl
				} else {
					logger.Debugf("getting deployment list from namespace %q", ctx.GetNamespace())
					var depl, err = ctx.Client.GetDeployment(ctx.GetNamespace().ID, args[0])
					if err != nil {
						logger.WithError(err).Errorf("unable to get deployment list from namespace %q", ctx.GetNamespace())
						ferr.Println(err)
						ctx.Exit(1)
					}
					deplData = depl
				}
			default:
				var deplList deployment.DeploymentList
				if flags.Solution != "" {
					var depls, err = ctx.Client.GetSolutionDeployments(ctx.GetNamespace().ID, flags.Solution)
					if err != nil {
						logger.WithError(err).Errorf("unable to get deployment list from namespace %q and solution %q", ctx.GetNamespace(), flags.Solution)
						ferr.Println(err)
						ctx.Exit(1)
					}
					deplList = depls
				} else {
					var depls, err = ctx.Client.GetDeploymentList(ctx.GetNamespace().ID)
					if err != nil {
						logger.WithError(err).Errorf("unable to get deployment list from namespace %q", ctx.GetNamespace())
						ferr.Println(err)
						ctx.Exit(1)
					}
					deplList = depls
				}
				if len(args) > 0 {
					var deplNames = strset.NewSet(args)
					deplList = deplList.Filter(func(depl deployment.Deployment) bool {
						return deplNames.Have(depl.Name)
					})
				}
				deplData = deplList
			}
			logger.Debugf("exporting deployment data")
			if err := export.ExportData(deplData, export.ExportConfig{
				Format:   export.ExportFormat(flags.Output),
				Filename: flags.File,
			}); err != nil {
				logrus.WithError(err).Errorf("unable to export data")
				ferr.Printf("unable to export data:\n%v\n", err)
				ctx.Exit(1)
			}
		},
	}
	if err := gpflag.ParseTo(&flags, command.PersistentFlags()); err != nil {
		panic(err)
	}
	return command
}
<commit_msg>add exporter<commit_after>package clideployment

import (
	"github.com/containerum/chkit/pkg/chkitErrors"
	"github.com/containerum/chkit/pkg/context"
	"github.com/containerum/chkit/pkg/model"
	"github.com/containerum/chkit/pkg/model/deployment"
	"github.com/containerum/chkit/pkg/porta"
	"github.com/containerum/chkit/pkg/util/ferr"
	"github.com/ninedraft/boxofstuff/strset"
	"github.com/octago/sflags/gen/gpflag"
	"github.com/sirupsen/logrus"
	"github.com/spf13/cobra"
)

const (
	ErrNoNamespaceSpecified chkitErrors.Err = "no namespace specified"
)

var aliases = []string{"depl", "deployments", "deploy", "de", "dpl", "depls", "dep"}

func Get(ctx *context.Context) *cobra.Command {
	var flags struct {
		Solution string
		porta.Exporter
	}
	command := &cobra.Command{
		Use:     "deployment",
		Short:   "show deployment data",
		Long:    "Print deployment data.",
		Example: "namespace deployment_names... [-n namespace_label]",
		Aliases: aliases,
		Run: func(command *cobra.Command, args []string) {
			var logger = ctx.Log.Command("get deployment")
			logger.Debugf("START")
			defer logrus.Debugf("END")
			var deplData model.Renderer

			switch len(args) {
			case 1:
				if flags.Solution != "" {
					logger.Debugf("getting deployment list from namespace %q and solution %q",
						ctx.GetNamespace(), flags.Solution)
					var depls, err = ctx.Client.GetSolutionDeployments(ctx.GetNamespace().ID, flags.Solution)
					if err != nil {
						logger.WithError(err).Errorf("unable to get deployment list from namespace %q and solution %q", ctx.GetNamespace(), flags.Solution)
						ferr.Println(err)
						ctx.Exit(1)
					}
					var depl, ok = depls.GetByName(args[0])
					if !ok {
						ferr.Printf("deployment %q not found\n", args[0])
					}
					deplData = depl
				} else {
					logger.Debugf("getting deployment list from namespace %q", ctx.GetNamespace())
					var depl, err = ctx.Client.GetDeployment(ctx.GetNamespace().ID, args[0])
					if err != nil {
						logger.WithError(err).Errorf("unable to get deployment list from namespace %q", ctx.GetNamespace())
						ferr.Println(err)
						ctx.Exit(1)
					}
					deplData = depl
				}
			default:
				var deplList deployment.DeploymentList
				if flags.Solution != "" {
					var depls, err = ctx.Client.GetSolutionDeployments(ctx.GetNamespace().ID, flags.Solution)
					if err != nil {
						logger.WithError(err).Errorf("unable to get deployment list from namespace %q and solution %q", ctx.GetNamespace(), flags.Solution)
						ferr.Println(err)
						ctx.Exit(1)
					}
					deplList = depls
				} else {
					var depls, err = ctx.Client.GetDeploymentList(ctx.GetNamespace().ID)
					if err != nil {
						logger.WithError(err).Errorf("unable to get deployment list from namespace %q", ctx.GetNamespace())
						ferr.Println(err)
						ctx.Exit(1)
					}
					deplList = depls
				}
				if len(args) > 0 {
					var deplNames = strset.NewSet(args)
					deplList = deplList.Filter(func(depl deployment.Deployment) bool {
						return deplNames.Have(depl.Name)
					})
				}
				deplData = deplList
			}
			logger.Debugf("exporting deployment data")
			if err := flags.Export(deplData); err != nil {
				ferr.Println(err)
				ctx.Exit(1)
			}
		},
	}
	if err := gpflag.ParseTo(&flags, command.PersistentFlags()); err != nil {
		panic(err)
	}
	return command
}
<|endoftext|>


<commit_before>package cmd

import (
	"errors"
	"fmt"
	"io"
	"strconv"
	"strings"

	"github.com/spf13/cobra"

	kapi "github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	kerrors "github.com/GoogleCloudPlatform/kubernetes/pkg/api/errors"
	kclient "github.com/GoogleCloudPlatform/kubernetes/pkg/client"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/fields"
	cmdutil "github.com/GoogleCloudPlatform/kubernetes/pkg/kubectl/cmd/util"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/labels"

	"github.com/openshift/origin/pkg/client"
	"github.com/openshift/origin/pkg/cmd/cli/describe"
	"github.com/openshift/origin/pkg/cmd/util/clientcmd"
	deployapi "github.com/openshift/origin/pkg/deploy/api"
	deployutil "github.com/openshift/origin/pkg/deploy/util"
)

type DeployOptions struct {
	out             io.Writer
	osClient        *client.Client
	kubeClient      *kclient.Client
	namespace       string
	baseCommandName string

	deploymentConfigName string
	deployLatest         bool
	retryDeploy          bool
	cancelDeploy         bool
}

const (
	deployLong = `View, start, cancel, and retry deployments.

If no options are given, view the latest deployment.

NOTE: This command is still under active development and is subject to change.`

	deployExample = `  // Display the latest deployment for the 'database' DeploymentConfig
  $ %[1]s deploy database

  // Start a new deployment based on the 'database' DeploymentConfig
  $ %[1]s deploy frontend --latest

  // Retry the latest failed deployment based on the 'frontend' DeploymentConfig
  // The deployer pod and any hook pods are deleted for the latest failed deployment
  $ %[1]s deploy frontend --retry

  // Cancel the in-progress deployment based on the 'frontend' DeploymentConfig
  $ %[1]s deploy frontend --cancel`
)

// NewCmdDeploy creates a new `deploy` command.
func NewCmdDeploy(fullName string, f *clientcmd.Factory, out io.Writer) *cobra.Command {
	options := &DeployOptions{
		baseCommandName: fullName,
	}

	cmd := &cobra.Command{
		Use:     "deploy DEPLOYMENTCONFIG",
		Short:   "View, start and restart deployments.",
		Long:    deployLong,
		Example: fmt.Sprintf(deployExample, fullName),
		Run: func(cmd *cobra.Command, args []string) {
			if err := options.Complete(f, args, out); err != nil {
				cmdutil.CheckErr(err)
			}

			if err := options.Validate(args); err != nil {
				cmdutil.CheckErr(cmdutil.UsageError(cmd, err.Error()))
			}

			if err := options.RunDeploy(); err != nil {
				cmdutil.CheckErr(err)
			}
		},
	}

	cmd.Flags().BoolVar(&options.deployLatest, "latest", false, "Start a new deployment now.")
	cmd.Flags().BoolVar(&options.retryDeploy, "retry", false, "Retry the latest failed deployment.")
	cmd.Flags().BoolVar(&options.cancelDeploy, "cancel", false, "Cancel the in-progress deployment.")

	return cmd
}

func (o *DeployOptions) Complete(f *clientcmd.Factory, args []string, out io.Writer) error {
	var err error

	o.osClient, o.kubeClient, err = f.Clients()
	if err != nil {
		return err
	}
	o.namespace, err = f.DefaultNamespace()
	if err != nil {
		return err
	}

	o.out = out

	if len(args) > 0 {
		o.deploymentConfigName = args[0]
	}

	return nil
}

func (o DeployOptions) Validate(args []string) error {
	if len(args) == 0 || len(args[0]) == 0 {
		return errors.New("a DeploymentConfig name is required.")
	}
	if len(args) > 1 {
		return errors.New("only one DeploymentConfig name is supported as argument.")
	}
	numOptions := 0
	if o.deployLatest {
		numOptions++
	}
	if o.retryDeploy {
		numOptions++
	}
	if o.cancelDeploy {
		numOptions++
	}
	if numOptions > 1 {
		return errors.New("only one of --latest, --retry, or --cancel is allowed.")
	}
	return nil
}

func (o DeployOptions) RunDeploy() error {
	config, err := o.osClient.DeploymentConfigs(o.namespace).Get(o.deploymentConfigName)
	if err != nil {
		return err
	}

	commandClient := &deployCommandClientImpl{
		GetDeploymentFn: func(namespace, name string) (*kapi.ReplicationController, error) {
			return o.kubeClient.ReplicationControllers(namespace).Get(name)
		},
		ListDeploymentsForConfigFn: func(namespace, configName string) (*kapi.ReplicationControllerList, error) {
			selector, err := labels.Parse(fmt.Sprintf("%s=%s", deployapi.DeploymentConfigLabel, configName))
			if err != nil {
				return nil, err
			}
			return o.kubeClient.ReplicationControllers(namespace).List(selector)
		},

		UpdateDeploymentConfigFn: func(config *deployapi.DeploymentConfig) (*deployapi.DeploymentConfig, error) {
			return o.osClient.DeploymentConfigs(config.Namespace).Update(config)
		},
		UpdateDeploymentFn: func(deployment *kapi.ReplicationController) (*kapi.ReplicationController, error) {
			return o.kubeClient.ReplicationControllers(deployment.Namespace).Update(deployment)
		},

		ListDeployerPodsForFn: func(namespace, deploymentName string) (*kapi.PodList, error) {
			selector, err := labels.Parse(fmt.Sprintf("%s=%s", deployapi.DeployerPodForDeploymentLabel, deploymentName))
			if err != nil {
				return nil, err
			}
			return o.kubeClient.Pods(namespace).List(selector, fields.Everything())
		},
		DeletePodFn: func(pod *kapi.Pod) error {
			return o.kubeClient.Pods(pod.Namespace).Delete(pod.Name, nil)
		},
	}

	switch {
	case o.deployLatest:
		c := &deployLatestCommand{client: commandClient}
		err = c.deploy(config, o.out)
	case o.retryDeploy:
		c := &retryDeploymentCommand{client: commandClient}
		err = c.retry(config, o.out)
	case o.cancelDeploy:
		c := &cancelDeploymentCommand{client: commandClient}
		err = c.cancel(config, o.out)
	default:
		describer := describe.NewLatestDeploymentsDescriber(o.osClient, o.kubeClient, -1)
		desc, err := describer.Describe(config.Namespace, config.Name)
		if err != nil {
			return err
		}
		fmt.Fprintln(o.out, desc)
	}

	return err
}

// deployCommandClient abstracts access to the API server.
type deployCommandClient interface {
	GetDeployment(namespace, name string) (*kapi.ReplicationController, error)
	ListDeploymentsForConfig(namespace, configName string) (*kapi.ReplicationControllerList, error)
	UpdateDeploymentConfig(*deployapi.DeploymentConfig) (*deployapi.DeploymentConfig, error)
	UpdateDeployment(*kapi.ReplicationController) (*kapi.ReplicationController, error)

	ListDeployerPodsFor(namespace, deploymentName string) (*kapi.PodList, error)
	DeletePod(pod *kapi.Pod) error
}

// deployLatestCommand can launch new deployments.
type deployLatestCommand struct {
	client deployCommandClient
}

// deploy launches a new deployment unless there's already a deployment
// process in progress for config.
func (c *deployLatestCommand) deploy(config *deployapi.DeploymentConfig, out io.Writer) error {
	deploymentName := deployutil.LatestDeploymentNameForConfig(config)
	deployment, err := c.client.GetDeployment(config.Namespace, deploymentName)
	if err != nil {
		if !kerrors.IsNotFound(err) {
			return err
		}
	} else {
		// Reject attempts to start a concurrent deployment.
		status := deployutil.DeploymentStatusFor(deployment)
		if status != deployapi.DeploymentStatusComplete && status != deployapi.DeploymentStatusFailed {
			return fmt.Errorf("#%d is already in progress (%s).\nOptionally, you can cancel this deployment using the --cancel option.", config.LatestVersion, status)
		}
	}

	config.LatestVersion++
	_, err = c.client.UpdateDeploymentConfig(config)
	if err == nil {
		fmt.Fprintf(out, "deployed #%d\n", config.LatestVersion)
	}
	return err
}

// retryDeploymentCommand can retry failed deployments.
type retryDeploymentCommand struct {
	client deployCommandClient
}

// retry resets the status of the latest deployment to New, which will cause
// the deployment to be retried. An error is returned if the deployment is not
// currently in a failed state.
func (c *retryDeploymentCommand) retry(config *deployapi.DeploymentConfig, out io.Writer) error {
	if config.LatestVersion == 0 {
		return fmt.Errorf("no deployments found for %s/%s", config.Namespace, config.Name)
	}
	deploymentName := deployutil.LatestDeploymentNameForConfig(config)
	deployment, err := c.client.GetDeployment(config.Namespace, deploymentName)
	if err != nil {
		if kerrors.IsNotFound(err) {
			return fmt.Errorf("Unable to find the latest deployment (#%d).\nYou can start a new deployment using the --latest option.", config.LatestVersion)
		}
		return err
	}

	if status := deployutil.DeploymentStatusFor(deployment); status != deployapi.DeploymentStatusFailed {
		message := fmt.Sprintf("#%d is %s; only failed deployments can be retried.\n", config.LatestVersion, status)
		if status == deployapi.DeploymentStatusComplete {
			message += fmt.Sprintf("You can start a new deployment using the --latest option.")
		} else {
			message += fmt.Sprintf("Optionally, you can cancel this deployment using the --cancel option.", config.LatestVersion)
		}

		return fmt.Errorf(message)
	}

	// Delete the deployer pod as well as the deployment hooks pods, if any
	pods, err := c.client.ListDeployerPodsFor(config.Namespace, deploymentName)
	if err != nil {
		return fmt.Errorf("Failed to list deployer/hook pods for deployment #%d: %v", config.LatestVersion, err)
	}
	for _, pod := range pods.Items {
		err := c.client.DeletePod(&pod)
		if err != nil {
			return fmt.Errorf("Failed to delete deployer/hook pod %s for deployment #%d: %v", pod.Name, config.LatestVersion, err)
		}
	}

	deployment.Annotations[deployapi.DeploymentStatusAnnotation] = string(deployapi.DeploymentStatusNew)
	// clear out the cancellation flag as well as any previous status-reason annotation
	delete(deployment.Annotations, deployapi.DeploymentStatusReasonAnnotation)
	delete(deployment.Annotations, deployapi.DeploymentCancelledAnnotation)
	_, err = c.client.UpdateDeployment(deployment)
	if err == nil {
		fmt.Fprintf(out, "retried #%d\n", config.LatestVersion)
	}
	return err
}

// cancelDeploymentCommand cancels the in-progress deployments.
type cancelDeploymentCommand struct {
	client deployCommandClient
}

// cancel cancels any deployment process in progress for config.
func (c *cancelDeploymentCommand) cancel(config *deployapi.DeploymentConfig, out io.Writer) error {
	deployments, err := c.client.ListDeploymentsForConfig(config.Namespace, config.Name)
	if err != nil {
		return err
	}
	failedCancellations := []string{}
	for _, deployment := range deployments.Items {
		status := deployutil.DeploymentStatusFor(&deployment)

		switch status {
		case deployapi.DeploymentStatusNew,
			deployapi.DeploymentStatusPending,
			deployapi.DeploymentStatusRunning:

			if deployutil.IsDeploymentCancelled(&deployment) {
				continue
			}

			deployment.Annotations[deployapi.DeploymentCancelledAnnotation] = deployapi.DeploymentCancelledAnnotationValue
			deployment.Annotations[deployapi.DeploymentStatusReasonAnnotation] = deployapi.DeploymentCancelledByUser
			_, err := c.client.UpdateDeployment(&deployment)
			if err == nil {
				fmt.Fprintf(out, "cancelled #%d\n", config.LatestVersion)
			} else {
				fmt.Fprintf(out, "couldn't cancel deployment %d (status: %s): %v", deployutil.DeploymentVersionFor(&deployment), status, err)
				failedCancellations = append(failedCancellations, strconv.Itoa(deployutil.DeploymentVersionFor(&deployment)))
			}
		}
	}
	if len(failedCancellations) == 0 {
		return nil
	} else {
		return fmt.Errorf("couldn't cancel deployment %s", strings.Join(failedCancellations, ", "))
	}
}

// deployCommandClientImpl is a pluggable deployCommandClient.
type deployCommandClientImpl struct {
	GetDeploymentFn            func(namespace, name string) (*kapi.ReplicationController, error)
	ListDeploymentsForConfigFn func(namespace, configName string) (*kapi.ReplicationControllerList, error)
	UpdateDeploymentConfigFn   func(*deployapi.DeploymentConfig) (*deployapi.DeploymentConfig, error)
	UpdateDeploymentFn         func(*kapi.ReplicationController) (*kapi.ReplicationController, error)

	ListDeployerPodsForFn func(namespace, deploymentName string) (*kapi.PodList, error)
	DeletePodFn           func(pod *kapi.Pod) error
}

func (c *deployCommandClientImpl) GetDeployment(namespace, name string) (*kapi.ReplicationController, error) {
	return c.GetDeploymentFn(namespace, name)
}

func (c *deployCommandClientImpl) ListDeploymentsForConfig(namespace, configName string) (*kapi.ReplicationControllerList, error) {
	return c.ListDeploymentsForConfigFn(namespace, configName)
}

func (c *deployCommandClientImpl) UpdateDeploymentConfig(config *deployapi.DeploymentConfig) (*deployapi.DeploymentConfig, error) {
	return c.UpdateDeploymentConfigFn(config)
}

func (c *deployCommandClientImpl) UpdateDeployment(deployment *kapi.ReplicationController) (*kapi.ReplicationController, error) {
	return c.UpdateDeploymentFn(deployment)
}

func (c *deployCommandClientImpl) ListDeployerPodsFor(namespace, deploymentName string) (*kapi.PodList, error) {
	return c.ListDeployerPodsForFn(namespace, deploymentName)
}

func (c *deployCommandClientImpl) DeletePod(pod *kapi.Pod) error {
	return c.DeletePodFn(pod)
}
<commit_msg>Fix CLI deployment cancellation<commit_after>package cmd

import (
	"errors"
	"fmt"
	"io"
	"strconv"
	"strings"

	"github.com/spf13/cobra"

	kapi "github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	kerrors "github.com/GoogleCloudPlatform/kubernetes/pkg/api/errors"
	kclient "github.com/GoogleCloudPlatform/kubernetes/pkg/client"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/fields"
	cmdutil "github.com/GoogleCloudPlatform/kubernetes/pkg/kubectl/cmd/util"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/labels"

	"github.com/openshift/origin/pkg/client"
	"github.com/openshift/origin/pkg/cmd/cli/describe"
	"github.com/openshift/origin/pkg/cmd/util/clientcmd"
	deployapi "github.com/openshift/origin/pkg/deploy/api"
	deployutil "github.com/openshift/origin/pkg/deploy/util"
)

type DeployOptions struct {
	out             io.Writer
	osClient        *client.Client
	kubeClient      *kclient.Client
	namespace       string
	baseCommandName string

	deploymentConfigName string
	deployLatest         bool
	retryDeploy          bool
	cancelDeploy         bool
}

const (
	deployLong = `View, start, cancel, and retry deployments.

If no options are given, view the latest deployment.

NOTE: This command is still under active development and is subject to change.`

	deployExample = `  // Display the latest deployment for the 'database' DeploymentConfig
  $ %[1]s deploy database

  // Start a new deployment based on the 'database' DeploymentConfig
  $ %[1]s deploy frontend --latest

  // Retry the latest failed deployment based on the 'frontend' DeploymentConfig
  // The deployer pod and any hook pods are deleted for the latest failed deployment
  $ %[1]s deploy frontend --retry

  // Cancel the in-progress deployment based on the 'frontend' DeploymentConfig
  $ %[1]s deploy frontend --cancel`
)

// NewCmdDeploy creates a new `deploy` command.
func NewCmdDeploy(fullName string, f *clientcmd.Factory, out io.Writer) *cobra.Command {
	options := &DeployOptions{
		baseCommandName: fullName,
	}

	cmd := &cobra.Command{
		Use:     "deploy DEPLOYMENTCONFIG",
		Short:   "View, start and restart deployments.",
		Long:    deployLong,
		Example: fmt.Sprintf(deployExample, fullName),
		Run: func(cmd *cobra.Command, args []string) {
			if err := options.Complete(f, args, out); err != nil {
				cmdutil.CheckErr(err)
			}

			if err := options.Validate(args); err != nil {
				cmdutil.CheckErr(cmdutil.UsageError(cmd, err.Error()))
			}

			if err := options.RunDeploy(); err != nil {
				cmdutil.CheckErr(err)
			}
		},
	}

	cmd.Flags().BoolVar(&options.deployLatest, "latest", false, "Start a new deployment now.")
	cmd.Flags().BoolVar(&options.retryDeploy, "retry", false, "Retry the latest failed deployment.")
	cmd.Flags().BoolVar(&options.cancelDeploy, "cancel", false, "Cancel the in-progress deployment.")

	return cmd
}

func (o *DeployOptions) Complete(f *clientcmd.Factory, args []string, out io.Writer) error {
	var err error

	o.osClient, o.kubeClient, err = f.Clients()
	if err != nil {
		return err
	}
	o.namespace, err = f.DefaultNamespace()
	if err != nil {
		return err
	}

	o.out = out

	if len(args) > 0 {
		o.deploymentConfigName = args[0]
	}

	return nil
}

func (o DeployOptions) Validate(args []string) error {
	if len(args) == 0 || len(args[0]) == 0 {
		return errors.New("a DeploymentConfig name is required.")
	}
	if len(args) > 1 {
		return errors.New("only one DeploymentConfig name is supported as argument.")
	}
	numOptions := 0
	if o.deployLatest {
		numOptions++
	}
	if o.retryDeploy {
		numOptions++
	}
	if o.cancelDeploy {
		numOptions++
	}
	if numOptions > 1 {
		return errors.New("only one of --latest, --retry, or --cancel is allowed.")
	}
	return nil
}

func (o DeployOptions) RunDeploy() error {
	config, err := o.osClient.DeploymentConfigs(o.namespace).Get(o.deploymentConfigName)
	if err != nil {
		return err
	}

	commandClient := &deployCommandClientImpl{
		GetDeploymentFn: func(namespace, name string) (*kapi.ReplicationController, error) {
			return o.kubeClient.ReplicationControllers(namespace).Get(name)
		},
		ListDeploymentsForConfigFn: func(namespace, configName string) (*kapi.ReplicationControllerList, error) {
			rcs, err := o.kubeClient.ReplicationControllers(namespace).List(labels.Everything())
			if err != nil {
				return nil, err
			}
			return &kapi.ReplicationControllerList{
				Items: deployutil.ConfigSelector(configName, rcs.Items),
			}, nil
		},

		UpdateDeploymentConfigFn: func(config *deployapi.DeploymentConfig) (*deployapi.DeploymentConfig, error) {
			return o.osClient.DeploymentConfigs(config.Namespace).Update(config)
		},
		UpdateDeploymentFn: func(deployment *kapi.ReplicationController) (*kapi.ReplicationController, error) {
			return o.kubeClient.ReplicationControllers(deployment.Namespace).Update(deployment)
		},

		ListDeployerPodsForFn: func(namespace, deploymentName string) (*kapi.PodList, error) {
			selector, err := labels.Parse(fmt.Sprintf("%s=%s", deployapi.DeployerPodForDeploymentLabel, deploymentName))
			if err != nil {
				return nil, err
			}
			return o.kubeClient.Pods(namespace).List(selector, fields.Everything())
		},
		DeletePodFn: func(pod *kapi.Pod) error {
			return o.kubeClient.Pods(pod.Namespace).Delete(pod.Name, nil)
		},
	}

	switch {
	case o.deployLatest:
		c := &deployLatestCommand{client: commandClient}
		err = c.deploy(config, o.out)
	case o.retryDeploy:
		c := &retryDeploymentCommand{client: commandClient}
		err = c.retry(config, o.out)
	case o.cancelDeploy:
		c := &cancelDeploymentCommand{client: commandClient}
		err = c.cancel(config, o.out)
	default:
		describer := describe.NewLatestDeploymentsDescriber(o.osClient, o.kubeClient, -1)
		desc, err := describer.Describe(config.Namespace, config.Name)
		if err != nil {
			return err
		}
		fmt.Fprintln(o.out, desc)
	}

	return err
}

// deployCommandClient abstracts access to the API server.
type deployCommandClient interface {
	GetDeployment(namespace, name string) (*kapi.ReplicationController, error)
	ListDeploymentsForConfig(namespace, configName string) (*kapi.ReplicationControllerList, error)
	UpdateDeploymentConfig(*deployapi.DeploymentConfig) (*deployapi.DeploymentConfig, error)
	UpdateDeployment(*kapi.ReplicationController) (*kapi.ReplicationController, error)

	ListDeployerPodsFor(namespace, deploymentName string) (*kapi.PodList, error)
	DeletePod(pod *kapi.Pod) error
}

// deployLatestCommand can launch new deployments.
type deployLatestCommand struct {
	client deployCommandClient
}

// deploy launches a new deployment unless there's already a deployment
// process in progress for config.
func (c *deployLatestCommand) deploy(config *deployapi.DeploymentConfig, out io.Writer) error {
	deploymentName := deployutil.LatestDeploymentNameForConfig(config)
	deployment, err := c.client.GetDeployment(config.Namespace, deploymentName)
	if err != nil {
		if !kerrors.IsNotFound(err) {
			return err
		}
	} else {
		// Reject attempts to start a concurrent deployment.
		status := deployutil.DeploymentStatusFor(deployment)
		if status != deployapi.DeploymentStatusComplete && status != deployapi.DeploymentStatusFailed {
			return fmt.Errorf("#%d is already in progress (%s).\nOptionally, you can cancel this deployment using the --cancel option.", config.LatestVersion, status)
		}
	}

	config.LatestVersion++
	_, err = c.client.UpdateDeploymentConfig(config)
	if err == nil {
		fmt.Fprintf(out, "deployed #%d\n", config.LatestVersion)
	}
	return err
}

// retryDeploymentCommand can retry failed deployments.
type retryDeploymentCommand struct {
	client deployCommandClient
}

// retry resets the status of the latest deployment to New, which will cause
// the deployment to be retried. An error is returned if the deployment is not
// currently in a failed state.
func (c *retryDeploymentCommand) retry(config *deployapi.DeploymentConfig, out io.Writer) error {
	if config.LatestVersion == 0 {
		return fmt.Errorf("no deployments found for %s/%s", config.Namespace, config.Name)
	}
	deploymentName := deployutil.LatestDeploymentNameForConfig(config)
	deployment, err := c.client.GetDeployment(config.Namespace, deploymentName)
	if err != nil {
		if kerrors.IsNotFound(err) {
			return fmt.Errorf("Unable to find the latest deployment (#%d).\nYou can start a new deployment using the --latest option.", config.LatestVersion)
		}
		return err
	}

	if status := deployutil.DeploymentStatusFor(deployment); status != deployapi.DeploymentStatusFailed {
		message := fmt.Sprintf("#%d is %s; only failed deployments can be retried.\n", config.LatestVersion, status)
		if status == deployapi.DeploymentStatusComplete {
			message += fmt.Sprintf("You can start a new deployment using the --latest option.")
		} else {
			message += fmt.Sprintf("Optionally, you can cancel this deployment using the --cancel option.", config.LatestVersion)
		}

		return fmt.Errorf(message)
	}

	// Delete the deployer pod as well as the deployment hooks pods, if any
	pods, err := c.client.ListDeployerPodsFor(config.Namespace, deploymentName)
	if err != nil {
		return fmt.Errorf("Failed to list deployer/hook pods for deployment #%d: %v", config.LatestVersion, err)
	}
	for _, pod := range pods.Items {
		err := c.client.DeletePod(&pod)
		if err != nil {
			return fmt.Errorf("Failed to delete deployer/hook pod %s for deployment #%d: %v", pod.Name, config.LatestVersion, err)
		}
	}

	deployment.Annotations[deployapi.DeploymentStatusAnnotation] = string(deployapi.DeploymentStatusNew)
	// clear out the cancellation flag as well as any previous status-reason annotation
	delete(deployment.Annotations, deployapi.DeploymentStatusReasonAnnotation)
	delete(deployment.Annotations, deployapi.DeploymentCancelledAnnotation)
	_, err = c.client.UpdateDeployment(deployment)
	if err == nil {
		fmt.Fprintf(out, "retried #%d\n", config.LatestVersion)
	}
	return err
}

// cancelDeploymentCommand cancels the in-progress deployments.
type cancelDeploymentCommand struct {
	client deployCommandClient
}

// cancel cancels any deployment process in progress for config.
func (c *cancelDeploymentCommand) cancel(config *deployapi.DeploymentConfig, out io.Writer) error {
	deployments, err := c.client.ListDeploymentsForConfig(config.Namespace, config.Name)
	if err != nil {
		return err
	}
	if len(deployments.Items) == 0 {
		fmt.Fprintln(out, "no deployments found to cancel")
		return nil
	}
	failedCancellations := []string{}
	for _, deployment := range deployments.Items {
		status := deployutil.DeploymentStatusFor(&deployment)

		switch status {
		case deployapi.DeploymentStatusNew,
			deployapi.DeploymentStatusPending,
			deployapi.DeploymentStatusRunning:

			if deployutil.IsDeploymentCancelled(&deployment) {
				continue
			}

			deployment.Annotations[deployapi.DeploymentCancelledAnnotation] = deployapi.DeploymentCancelledAnnotationValue
			deployment.Annotations[deployapi.DeploymentStatusReasonAnnotation] = deployapi.DeploymentCancelledByUser
			_, err := c.client.UpdateDeployment(&deployment)
			if err == nil {
				fmt.Fprintf(out, "cancelled #%d\n", config.LatestVersion)
			} else {
				fmt.Fprintf(out, "couldn't cancel deployment %d (status: %s): %v\n", deployutil.DeploymentVersionFor(&deployment), status, err)
				failedCancellations = append(failedCancellations, strconv.Itoa(deployutil.DeploymentVersionFor(&deployment)))
			}
		default:
			fmt.Fprintln(out, "no active deployments to cancel")
		}
	}
	if len(failedCancellations) == 0 {
		return nil
	} else {
		return fmt.Errorf("couldn't cancel deployment %s", strings.Join(failedCancellations, ", "))
	}
}

// deployCommandClientImpl is a pluggable deployCommandClient.
type deployCommandClientImpl struct {
	GetDeploymentFn            func(namespace, name string) (*kapi.ReplicationController, error)
	ListDeploymentsForConfigFn func(namespace, configName string) (*kapi.ReplicationControllerList, error)
	UpdateDeploymentConfigFn   func(*deployapi.DeploymentConfig) (*deployapi.DeploymentConfig, error)
	UpdateDeploymentFn         func(*kapi.ReplicationController) (*kapi.ReplicationController, error)

	ListDeployerPodsForFn func(namespace, deploymentName string) (*kapi.PodList, error)
	DeletePodFn           func(pod *kapi.Pod) error
}

func (c *deployCommandClientImpl) GetDeployment(namespace, name string) (*kapi.ReplicationController, error) {
	return c.GetDeploymentFn(namespace, name)
}

func (c *deployCommandClientImpl) ListDeploymentsForConfig(namespace, configName string) (*kapi.ReplicationControllerList, error) {
	return c.ListDeploymentsForConfigFn(namespace, configName)
}

func (c *deployCommandClientImpl) UpdateDeploymentConfig(config *deployapi.DeploymentConfig) (*deployapi.DeploymentConfig, error) {
	return c.UpdateDeploymentConfigFn(config)
}

func (c *deployCommandClientImpl) UpdateDeployment(deployment *kapi.ReplicationController) (*kapi.ReplicationController, error) {
	return c.UpdateDeploymentFn(deployment)
}

func (c *deployCommandClientImpl) ListDeployerPodsFor(namespace, deploymentName string) (*kapi.PodList, error) {
	return c.ListDeployerPodsForFn(namespace, deploymentName)
}

func (c *deployCommandClientImpl) DeletePod(pod *kapi.Pod) error {
	return c.DeletePodFn(pod)
}
<|endoftext|>