<commit_before>// Package cq provides tools for interacting with the CQ tools.
package cq

import (
	"bytes"
	"context"
	"fmt"
	"path/filepath"
	"regexp"
	"sync"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/golang/protobuf/ptypes"
	buildbucketpb "go.chromium.org/luci/buildbucket/proto"
	"go.chromium.org/luci/cq/api/config/v2"
	"go.skia.org/infra/go/common"
	"go.skia.org/infra/go/gerrit"
	"go.skia.org/infra/go/gitiles"
	"go.skia.org/infra/go/metrics2"
	"go.skia.org/infra/go/sklog"
	"go.skia.org/infra/go/util"
)

const (
	CQ_CFG_FILE = "commit-queue.cfg"
	CQ_CFG_REF  = "infra/config"

	MASTER_REF = "refs/heads/master"

	// Constants for in-flight metrics.
	INFLIGHT_METRIC_NAME     = "in_flight"
	INFLIGHT_TRYBOT_DURATION = "trybot_duration"
	INFLIGHT_TRYBOT_NUM      = "trybot_num"
	INFLIGHT_WAITING_IN_CQ   = "waiting_in_cq"

	// Constants for landed metrics.
	LANDED_METRIC_NAME     = "after_commit"
	LANDED_TRYBOT_DURATION = "trybot_duration"
	LANDED_TOTAL_DURATION  = "total_duration"

	// Thresholds after which errors are logged.
	CQ_TRYBOT_DURATION_SECS_THRESHOLD = 2700
	CQ_TRYBOTS_COUNT_THRESHOLD        = 50
)

var (
	// Slice of all known presubmit bot names.
	PRESUBMIT_BOTS = []string{"skia_presubmit-Trybot"}

	// Mutext to control access to the slice of CQ trybots.
	cqTryBotsMutex sync.RWMutex
)

// NewClient creates a new client for interacting with CQ tools.
func NewClient(gerritClient *gerrit.Gerrit, cqTryBotsFunc GetCQTryBotsFn, metricName string) (*Client, error) {
	cqTryBots, err := cqTryBotsFunc()
	if err != nil {
		return nil, err
	}
	return &Client{gerritClient, util.NewStringSet(cqTryBots), cqTryBotsFunc, metricName}, err
}

// GetCQTryBotsFn is an interface for returing the CQ trybots of a project.
type GetCQTryBotsFn func() ([]string, error)

type Client struct {
	gerritClient  *gerrit.Gerrit
	cqTryBots     util.StringSet
	cqTryBotsFunc GetCQTryBotsFn
	metricName    string
}

// GetSkiaCQTryBots is a Skia implementation of GetCQTryBotsFn.
func GetSkiaCQTryBots() ([]string, error) {
	cfg, err := GetCQConfig(gitiles.NewRepo(common.REPO_SKIA, nil))
	if err != nil {
		return nil, err
	}
	return GetCQTryBots(cfg, MASTER_REF)
}

// GetSkiaInfraCQTryBots is a Skia Infra implementation of GetCQTryBotsFn.
func GetSkiaInfraCQTryBots() ([]string, error) {
	cfg, err := GetCQConfig(gitiles.NewRepo(common.REPO_SKIA_INFRA, nil))
	if err != nil {
		return nil, err
	}
	return GetCQTryBots(cfg, MASTER_REF)
}

// MatchConfigGroup returns the ConfigGroup, ConfigGroup_Gerrit, and
// ConfigGroup_Gerrit_Project which match the given full ref name, or nil if
// there is no matching ConfigGroup.
func MatchConfigGroup(cqCfg *config.Config, ref string) (*config.ConfigGroup, *config.ConfigGroup_Gerrit, *config.ConfigGroup_Gerrit_Project, error) {
	for _, configGroup := range cqCfg.GetConfigGroups() {
		for _, g := range configGroup.GetGerrit() {
			for _, p := range g.GetProjects() {
				for _, r := range p.GetRefRegexp() {
					m, err := regexp.MatchString(r, ref)
					if err != nil {
						return nil, nil, nil, fmt.Errorf("Error when compiling %s: %s", r, err)
					}
					if m {
						// Found the ref we were looking for.
						return configGroup, g, p, nil
					}
				}
			}
		}
	}
	return nil, nil, nil, nil
}

// GetCQConfig returns the Config for the given repo.
func GetCQConfig(repo *gitiles.Repo) (*config.Config, error) {
	var buf bytes.Buffer
	if err := repo.ReadFileAtRef(context.Background(), CQ_CFG_FILE, CQ_CFG_REF, &buf); err != nil {
		return nil, err
	}
	var cqCfg config.Config
	if err := proto.UnmarshalText(buf.String(), &cqCfg); err != nil {
		return nil, err
	}
	return &cqCfg, nil
}

// GetCQTryBots is a convenience method for retrieving the list of CQ trybots
// from a Config.
func GetCQTryBots(cqCfg *config.Config, ref string) ([]string, error) {
	tryJobs := []string{}
	configGroup, _, _, err := MatchConfigGroup(cqCfg, ref)
	if err != nil {
		return nil, err
	}
	if configGroup != nil {
		// Found the ref we were looking for.
		for _, builder := range configGroup.GetVerifiers().GetTryjob().GetBuilders() {
			if builder.GetExperimentPercentage() > 0 && builder.GetExperimentPercentage() < 100 {
				// Exclude experimental builders, unless running for all CLs.
				continue
			}
			if builder.IncludableOnly {
				// Exclude builders which have been specified only for "Cq-Include-Trybots".
				continue
			}
			if util.ContainsAny(builder.GetName(), PRESUBMIT_BOTS) {
				// Exclude presubmit bots because they could fail or be delayed
				// due to factors such as owners approval and other project
				// specific checks.
				continue
			}
			// Strip out the bucket and use only the builder name.
			// Eg: chromium/try/mac_chromium_compile_dbg_ng -> mac_chromium_compile_dbg_ng
			builderName := filepath.Base(builder.GetName())
			tryJobs = append(tryJobs, builderName)
		}
	}

	sklog.Infof("The list of CQ trybots is: %s", tryJobs)
	return tryJobs, nil
}

// RefreshCQTryBots refreshes the slice of CQ trybots on the instance. Access
// to the trybots is protected by a RWMutex.
func (c *Client) RefreshCQTryBots() error {
	tryBots, err := c.cqTryBotsFunc()
	if err != nil {
		return err
	}
	cqTryBotsMutex.Lock()
	defer cqTryBotsMutex.Unlock()
	c.cqTryBots = util.NewStringSet(tryBots)
	return nil
}

// ReportCQStats reports all relevant stats for the specified Gerrit change.
// Note: Different stats are reported depending on whether the change has been
// merged or not.
// All created metrics will be registered in reportedMetrics.
func (c *Client) ReportCQStats(ctx context.Context, change int64, reportedMetrics map[metrics2.Int64Metric]struct{}) error {
	changeInfo, err := c.gerritClient.GetIssueProperties(ctx, change)
	if err != nil {
		return err
	}
	patchsetIds := changeInfo.GetPatchsetIDs()
	latestPatchsetId := patchsetIds[len(patchsetIds)-1]
	if changeInfo.Committed {
		// TODO(rmistry): The last patchset in Gerrit does not contain trybot
		// information so we have to look at the one immediately before it.
		// This will be fixed with crbug.com/634944.
		latestPatchsetId = patchsetIds[len(patchsetIds)-2]
	}

	builds, err := c.gerritClient.GetTrybotResults(ctx, change, latestPatchsetId)
	if err != nil {
		return err
	}
	// Consider only CQ bots.
	cqBuilds := []*buildbucketpb.Build{}
	for _, b := range builds {
		if c.isCQTryBot(b.Builder.Builder) {
			cqBuilds = append(cqBuilds, b)
		}
	}
	gerritURL := fmt.Sprintf("%s/c/%d/%d", gerrit.GERRIT_SKIA_URL, change, latestPatchsetId)
	if len(cqBuilds) == 0 {
		sklog.Infof("No trybot results were found for %s", gerritURL)
		return nil
	}

	sklog.Infof("Starting processing %s. Merged status: %t", gerritURL, changeInfo.Committed)

	if changeInfo.Committed {
		c.ReportCQStatsForLandedCL(cqBuilds, gerritURL, reportedMetrics)
	} else {
		c.ReportCQStatsForInFlightCL(cqBuilds, gerritURL, reportedMetrics)
	}
	return nil
}

// ReportCQStatsForLandedCL reports the following metrics for the specified
// change and patchsetID:
// * The total time the change spent waiting for CQ trybots to complete.
// * The time each CQ trybot took to complete.
// All created metrics will be registered in reportedMetrics.
func (c *Client) ReportCQStatsForLandedCL(cqBuilds []*buildbucketpb.Build, gerritURL string, reportedMetrics map[metrics2.Int64Metric]struct{}) {
	endTimeOfCQBots := time.Time{}
	maximumTrybotDuration := int64(0)
	for _, b := range cqBuilds {
		createdTime, err := ptypes.Timestamp(b.CreateTime)
		if err != nil {
			sklog.Errorf("Failed to convert timestamp for %d; skipping: %s", b.Id, err)
			continue
		}
		createdTime = createdTime.UTC()
		if b.EndTime == nil {
			sklog.Warningf("Skipping %s on %s. The correct completed time has not shown up in Buildbucket yet.", b.Builder.Builder, gerritURL)
			continue
		}
		completedTime, err := ptypes.Timestamp(b.EndTime)
		if err != nil {
			sklog.Errorf("Failed to convert timestamp for %d; skipping: %s", b.Id, err)
			continue
		}
		completedTime = completedTime.UTC()
		if endTimeOfCQBots.Before(completedTime) {
			endTimeOfCQBots = completedTime
		}

		duration := int64(completedTime.Sub(createdTime).Seconds())
		sklog.Infof("%s was created at %s by %s and completed at %s. Total duration: %d", b.Builder.Builder, createdTime, gerritURL, completedTime, duration)
		landedTrybotDurationMetric := c.getLandedTrybotDurationMetric(b.Builder.Builder, gerritURL)
		landedTrybotDurationMetric.Update(duration)
		reportedMetrics[landedTrybotDurationMetric] = struct{}{}

		if duration > maximumTrybotDuration {
			maximumTrybotDuration = duration
		}
	}

	sklog.Infof("Maximum trybot duration for %s: %d", gerritURL, maximumTrybotDuration)
	sklog.Infof("Furthest completion time for %s: %s", gerritURL, endTimeOfCQBots)
	landedTotalDurationMetric := metrics2.GetInt64Metric(fmt.Sprintf("%s_%s_%s", c.metricName, LANDED_METRIC_NAME, LANDED_TOTAL_DURATION), map[string]string{"gerritURL": gerritURL})
	landedTotalDurationMetric.Update(maximumTrybotDuration)
	reportedMetrics[landedTotalDurationMetric] = struct{}{}
}

// ReportCQStatsForInFlightCL reports the following metrics for the specified
// change and patchsetID:
// * How long CQ trybots have been running for.
// * How many CQ trybots have been triggered.
// All created metrics will be registered in reportedMetrics.
func (c *Client) ReportCQStatsForInFlightCL(cqBuilds []*buildbucketpb.Build, gerritURL string, reportedMetrics map[metrics2.Int64Metric]struct{}) {
	totalTriggeredCQBots := int(0)
	currentTime := time.Now()
	for _, b := range cqBuilds {
		totalTriggeredCQBots++

		createdTime, err := ptypes.Timestamp(b.CreateTime)
		if err != nil {
			sklog.Errorf("Failed to convert timestamp for %d; skipping: %s", b.Id, err)
			continue
		}
		createdTime = createdTime.UTC()
		if b.EndTime != nil {
			completedTime, err := ptypes.Timestamp(b.EndTime)
			if err != nil {
				sklog.Errorf("Failed to convert timestamp for %d; skipping: %s", b.Id, err)
				continue
			}
			completedTime = completedTime.UTC()
			if time.Hour*24 < time.Now().UTC().Sub(completedTime) {
				// The build has completed more than a day ago. Do not include it
				// in totalTriggeredCQBots. See skbug.com/7340.
				totalTriggeredCQBots--
			}
			// The build has completed so move on.
			continue
		}

		duration := int64(currentTime.Sub(createdTime).Seconds())
		if duration > CQ_TRYBOT_DURATION_SECS_THRESHOLD {
			sklog.Errorf("CQTrybotDurationError: %s was triggered by %s and is still running after %d seconds. Threshold is %d seconds.", b.Builder.Builder, gerritURL, duration, CQ_TRYBOT_DURATION_SECS_THRESHOLD)
		}
		inflightTrybotDurationMetric := c.getInflightTrybotDurationMetric(b.Builder.Builder, gerritURL)
		inflightTrybotDurationMetric.Update(duration)
		reportedMetrics[inflightTrybotDurationMetric] = struct{}{}
	}

	cqTryBotsMutex.RLock()
	cqTryBotsMutex.RUnlock()
	if totalTriggeredCQBots > CQ_TRYBOTS_COUNT_THRESHOLD {
		sklog.Errorf("CQCLsCountError: %d trybots have been triggered by %s. Threshold is %d trybots.", totalTriggeredCQBots, gerritURL, CQ_TRYBOTS_COUNT_THRESHOLD)
	}
	trybotNumDurationMetric := metrics2.GetInt64Metric(fmt.Sprintf("%s_%s_%s", c.metricName, INFLIGHT_METRIC_NAME, INFLIGHT_TRYBOT_NUM), map[string]string{"gerritURL": gerritURL})
	trybotNumDurationMetric.Update(int64(totalTriggeredCQBots))
	reportedMetrics[trybotNumDurationMetric] = struct{}{}
}

func (c *Client) getInflightTrybotDurationMetric(tryBot, gerritURL string) metrics2.Int64Metric {
	metricName := fmt.Sprintf("%s_%s_%s", c.metricName, INFLIGHT_METRIC_NAME, INFLIGHT_TRYBOT_DURATION)
	tags := map[string]string{
		"trybot":    tryBot,
		"gerritURL": gerritURL,
	}
	return metrics2.GetInt64Metric(metricName, tags)
}

func (c *Client) getLandedTrybotDurationMetric(tryBot, gerritURL string) metrics2.Int64Metric {
	metricName := fmt.Sprintf("%s_%s_%s", c.metricName, LANDED_METRIC_NAME, LANDED_TRYBOT_DURATION)
	tags := map[string]string{
		"trybot":    tryBot,
		"gerritURL": gerritURL,
	}
	return metrics2.GetInt64Metric(metricName, tags)
}

func (c *Client) isCQTryBot(builderName string) bool {
	cqTryBotsMutex.RLock()
	isCQTrybot := c.cqTryBots[builderName]
	cqTryBotsMutex.RUnlock()
	return isCQTrybot
}
<commit_msg>[cq_watcher] Use build creation time to decide if builds are old<commit_after>// Package cq provides tools for interacting with the CQ tools.
package cq

import (
	"bytes"
	"context"
	"fmt"
	"path/filepath"
	"regexp"
	"sync"
	"time"

	"github.com/golang/protobuf/proto"
	"github.com/golang/protobuf/ptypes"
	buildbucketpb "go.chromium.org/luci/buildbucket/proto"
	"go.chromium.org/luci/cq/api/config/v2"
	"go.skia.org/infra/go/common"
	"go.skia.org/infra/go/gerrit"
	"go.skia.org/infra/go/gitiles"
	"go.skia.org/infra/go/metrics2"
	"go.skia.org/infra/go/sklog"
	"go.skia.org/infra/go/util"
)

const (
	CQ_CFG_FILE = "commit-queue.cfg"
	CQ_CFG_REF  = "infra/config"

	MASTER_REF = "refs/heads/master"

	// Constants for in-flight metrics.
	INFLIGHT_METRIC_NAME     = "in_flight"
	INFLIGHT_TRYBOT_DURATION = "trybot_duration"
	INFLIGHT_TRYBOT_NUM      = "trybot_num"
	INFLIGHT_WAITING_IN_CQ   = "waiting_in_cq"

	// Constants for landed metrics.
	LANDED_METRIC_NAME     = "after_commit"
	LANDED_TRYBOT_DURATION = "trybot_duration"
	LANDED_TOTAL_DURATION  = "total_duration"

	// Thresholds after which errors are logged.
	CQ_TRYBOT_DURATION_SECS_THRESHOLD = 2700
	CQ_TRYBOTS_COUNT_THRESHOLD        = 50
)

var (
	// Slice of all known presubmit bot names.
	PRESUBMIT_BOTS = []string{"skia_presubmit-Trybot"}

	// Mutext to control access to the slice of CQ trybots.
	cqTryBotsMutex sync.RWMutex
)

// NewClient creates a new client for interacting with CQ tools.
func NewClient(gerritClient *gerrit.Gerrit, cqTryBotsFunc GetCQTryBotsFn, metricName string) (*Client, error) {
	cqTryBots, err := cqTryBotsFunc()
	if err != nil {
		return nil, err
	}
	return &Client{gerritClient, util.NewStringSet(cqTryBots), cqTryBotsFunc, metricName}, err
}

// GetCQTryBotsFn is an interface for returing the CQ trybots of a project.
type GetCQTryBotsFn func() ([]string, error)

type Client struct {
	gerritClient  *gerrit.Gerrit
	cqTryBots     util.StringSet
	cqTryBotsFunc GetCQTryBotsFn
	metricName    string
}

// GetSkiaCQTryBots is a Skia implementation of GetCQTryBotsFn.
func GetSkiaCQTryBots() ([]string, error) {
	cfg, err := GetCQConfig(gitiles.NewRepo(common.REPO_SKIA, nil))
	if err != nil {
		return nil, err
	}
	return GetCQTryBots(cfg, MASTER_REF)
}

// GetSkiaInfraCQTryBots is a Skia Infra implementation of GetCQTryBotsFn.
func GetSkiaInfraCQTryBots() ([]string, error) {
	cfg, err := GetCQConfig(gitiles.NewRepo(common.REPO_SKIA_INFRA, nil))
	if err != nil {
		return nil, err
	}
	return GetCQTryBots(cfg, MASTER_REF)
}

// MatchConfigGroup returns the ConfigGroup, ConfigGroup_Gerrit, and
// ConfigGroup_Gerrit_Project which match the given full ref name, or nil if
// there is no matching ConfigGroup.
func MatchConfigGroup(cqCfg *config.Config, ref string) (*config.ConfigGroup, *config.ConfigGroup_Gerrit, *config.ConfigGroup_Gerrit_Project, error) {
	for _, configGroup := range cqCfg.GetConfigGroups() {
		for _, g := range configGroup.GetGerrit() {
			for _, p := range g.GetProjects() {
				for _, r := range p.GetRefRegexp() {
					m, err := regexp.MatchString(r, ref)
					if err != nil {
						return nil, nil, nil, fmt.Errorf("Error when compiling %s: %s", r, err)
					}
					if m {
						// Found the ref we were looking for.
						return configGroup, g, p, nil
					}
				}
			}
		}
	}
	return nil, nil, nil, nil
}

// GetCQConfig returns the Config for the given repo.
func GetCQConfig(repo *gitiles.Repo) (*config.Config, error) {
	var buf bytes.Buffer
	if err := repo.ReadFileAtRef(context.Background(), CQ_CFG_FILE, CQ_CFG_REF, &buf); err != nil {
		return nil, err
	}
	var cqCfg config.Config
	if err := proto.UnmarshalText(buf.String(), &cqCfg); err != nil {
		return nil, err
	}
	return &cqCfg, nil
}

// GetCQTryBots is a convenience method for retrieving the list of CQ trybots
// from a Config.
func GetCQTryBots(cqCfg *config.Config, ref string) ([]string, error) {
	tryJobs := []string{}
	configGroup, _, _, err := MatchConfigGroup(cqCfg, ref)
	if err != nil {
		return nil, err
	}
	if configGroup != nil {
		// Found the ref we were looking for.
		for _, builder := range configGroup.GetVerifiers().GetTryjob().GetBuilders() {
			if builder.GetExperimentPercentage() > 0 && builder.GetExperimentPercentage() < 100 {
				// Exclude experimental builders, unless running for all CLs.
				continue
			}
			if builder.IncludableOnly {
				// Exclude builders which have been specified only for "Cq-Include-Trybots".
				continue
			}
			if util.ContainsAny(builder.GetName(), PRESUBMIT_BOTS) {
				// Exclude presubmit bots because they could fail or be delayed
				// due to factors such as owners approval and other project
				// specific checks.
				continue
			}
			// Strip out the bucket and use only the builder name.
			// Eg: chromium/try/mac_chromium_compile_dbg_ng -> mac_chromium_compile_dbg_ng
			builderName := filepath.Base(builder.GetName())
			tryJobs = append(tryJobs, builderName)
		}
	}

	sklog.Infof("The list of CQ trybots is: %s", tryJobs)
	return tryJobs, nil
}

// RefreshCQTryBots refreshes the slice of CQ trybots on the instance. Access
// to the trybots is protected by a RWMutex.
func (c *Client) RefreshCQTryBots() error {
	tryBots, err := c.cqTryBotsFunc()
	if err != nil {
		return err
	}
	cqTryBotsMutex.Lock()
	defer cqTryBotsMutex.Unlock()
	c.cqTryBots = util.NewStringSet(tryBots)
	return nil
}

// ReportCQStats reports all relevant stats for the specified Gerrit change.
// Note: Different stats are reported depending on whether the change has been
// merged or not.
// All created metrics will be registered in reportedMetrics.
func (c *Client) ReportCQStats(ctx context.Context, change int64, reportedMetrics map[metrics2.Int64Metric]struct{}) error {
	changeInfo, err := c.gerritClient.GetIssueProperties(ctx, change)
	if err != nil {
		return err
	}
	patchsetIds := changeInfo.GetPatchsetIDs()
	latestPatchsetId := patchsetIds[len(patchsetIds)-1]
	if changeInfo.Committed {
		// TODO(rmistry): The last patchset in Gerrit does not contain trybot
		// information so we have to look at the one immediately before it.
		// This will be fixed with crbug.com/634944.
		latestPatchsetId = patchsetIds[len(patchsetIds)-2]
	}

	builds, err := c.gerritClient.GetTrybotResults(ctx, change, latestPatchsetId)
	if err != nil {
		return err
	}
	// Consider only CQ bots.
	cqBuilds := []*buildbucketpb.Build{}
	for _, b := range builds {
		if c.isCQTryBot(b.Builder.Builder) {
			cqBuilds = append(cqBuilds, b)
		}
	}
	gerritURL := fmt.Sprintf("%s/c/%d/%d", gerrit.GERRIT_SKIA_URL, change, latestPatchsetId)
	if len(cqBuilds) == 0 {
		sklog.Infof("No trybot results were found for %s", gerritURL)
		return nil
	}

	sklog.Infof("Starting processing %s. Merged status: %t", gerritURL, changeInfo.Committed)

	if changeInfo.Committed {
		c.ReportCQStatsForLandedCL(cqBuilds, gerritURL, reportedMetrics)
	} else {
		c.ReportCQStatsForInFlightCL(cqBuilds, gerritURL, reportedMetrics)
	}
	return nil
}

// ReportCQStatsForLandedCL reports the following metrics for the specified
// change and patchsetID:
// * The total time the change spent waiting for CQ trybots to complete.
// * The time each CQ trybot took to complete.
// All created metrics will be registered in reportedMetrics.
func (c *Client) ReportCQStatsForLandedCL(cqBuilds []*buildbucketpb.Build, gerritURL string, reportedMetrics map[metrics2.Int64Metric]struct{}) {
	endTimeOfCQBots := time.Time{}
	maximumTrybotDuration := int64(0)
	for _, b := range cqBuilds {
		createdTime, err := ptypes.Timestamp(b.CreateTime)
		if err != nil {
			sklog.Errorf("Failed to convert timestamp for %d; skipping: %s", b.Id, err)
			continue
		}
		createdTime = createdTime.UTC()
		if b.EndTime == nil {
			sklog.Warningf("Skipping %s on %s. The correct completed time has not shown up in Buildbucket yet.", b.Builder.Builder, gerritURL)
			continue
		}
		completedTime, err := ptypes.Timestamp(b.EndTime)
		if err != nil {
			sklog.Errorf("Failed to convert timestamp for %d; skipping: %s", b.Id, err)
			continue
		}
		completedTime = completedTime.UTC()
		if endTimeOfCQBots.Before(completedTime) {
			endTimeOfCQBots = completedTime
		}

		duration := int64(completedTime.Sub(createdTime).Seconds())
		sklog.Infof("%s was created at %s by %s and completed at %s. Total duration: %d", b.Builder.Builder, createdTime, gerritURL, completedTime, duration)
		landedTrybotDurationMetric := c.getLandedTrybotDurationMetric(b.Builder.Builder, gerritURL)
		landedTrybotDurationMetric.Update(duration)
		reportedMetrics[landedTrybotDurationMetric] = struct{}{}

		if duration > maximumTrybotDuration {
			maximumTrybotDuration = duration
		}
	}

	sklog.Infof("Maximum trybot duration for %s: %d", gerritURL, maximumTrybotDuration)
	sklog.Infof("Furthest completion time for %s: %s", gerritURL, endTimeOfCQBots)
	landedTotalDurationMetric := metrics2.GetInt64Metric(fmt.Sprintf("%s_%s_%s", c.metricName, LANDED_METRIC_NAME, LANDED_TOTAL_DURATION), map[string]string{"gerritURL": gerritURL})
	landedTotalDurationMetric.Update(maximumTrybotDuration)
	reportedMetrics[landedTotalDurationMetric] = struct{}{}
}

// ReportCQStatsForInFlightCL reports the following metrics for the specified
// change and patchsetID:
// * How long CQ trybots have been running for.
// * How many CQ trybots have been triggered.
// All created metrics will be registered in reportedMetrics.
func (c *Client) ReportCQStatsForInFlightCL(cqBuilds []*buildbucketpb.Build, gerritURL string, reportedMetrics map[metrics2.Int64Metric]struct{}) {
	totalTriggeredCQBots := int(0)
	currentTime := time.Now()
	for _, b := range cqBuilds {
		totalTriggeredCQBots++

		createdTime, err := ptypes.Timestamp(b.CreateTime)
		if err != nil {
			sklog.Errorf("Failed to convert timestamp for %d; skipping: %s", b.Id, err)
			continue
		}
		createdTime = createdTime.UTC()
		if b.EndTime != nil {
			if time.Hour*24 < time.Now().UTC().Sub(createdTime) {
				// The build was created more than a day ago. Do not include it
				// in totalTriggeredCQBots. See skbug.com/7340.
				// Creation time is used above instead of completion time because
				// that is what CQ does:
				// https://chrome-internal.googlesource.com/infra/infra_internal/+/master/infra_internal/services/cq/verification/tryjob_utils.py#1271
				totalTriggeredCQBots--
			}
			// The build has completed so move on.
			continue
		}

		duration := int64(currentTime.Sub(createdTime).Seconds())
		if duration > CQ_TRYBOT_DURATION_SECS_THRESHOLD {
			sklog.Errorf("CQTrybotDurationError: %s was triggered by %s and is still running after %d seconds. Threshold is %d seconds.", b.Builder.Builder, gerritURL, duration, CQ_TRYBOT_DURATION_SECS_THRESHOLD)
		}
		inflightTrybotDurationMetric := c.getInflightTrybotDurationMetric(b.Builder.Builder, gerritURL)
		inflightTrybotDurationMetric.Update(duration)
		reportedMetrics[inflightTrybotDurationMetric] = struct{}{}
	}

	cqTryBotsMutex.RLock()
	cqTryBotsMutex.RUnlock()
	if totalTriggeredCQBots > CQ_TRYBOTS_COUNT_THRESHOLD {
		sklog.Errorf("CQCLsCountError: %d trybots have been triggered by %s. Threshold is %d trybots.", totalTriggeredCQBots, gerritURL, CQ_TRYBOTS_COUNT_THRESHOLD)
	}
	trybotNumDurationMetric := metrics2.GetInt64Metric(fmt.Sprintf("%s_%s_%s", c.metricName, INFLIGHT_METRIC_NAME, INFLIGHT_TRYBOT_NUM), map[string]string{"gerritURL": gerritURL})
	trybotNumDurationMetric.Update(int64(totalTriggeredCQBots))
	reportedMetrics[trybotNumDurationMetric] = struct{}{}
}

func (c *Client) getInflightTrybotDurationMetric(tryBot, gerritURL string) metrics2.Int64Metric {
	metricName := fmt.Sprintf("%s_%s_%s", c.metricName, INFLIGHT_METRIC_NAME, INFLIGHT_TRYBOT_DURATION)
	tags := map[string]string{
		"trybot":    tryBot,
		"gerritURL": gerritURL,
	}
	return metrics2.GetInt64Metric(metricName, tags)
}

func (c *Client) getLandedTrybotDurationMetric(tryBot, gerritURL string) metrics2.Int64Metric {
	metricName := fmt.Sprintf("%s_%s_%s", c.metricName, LANDED_METRIC_NAME, LANDED_TRYBOT_DURATION)
	tags := map[string]string{
		"trybot":    tryBot,
		"gerritURL": gerritURL,
	}
	return metrics2.GetInt64Metric(metricName, tags)
}

func (c *Client) isCQTryBot(builderName string) bool {
	cqTryBotsMutex.RLock()
	isCQTrybot := c.cqTryBots[builderName]
	cqTryBotsMutex.RUnlock()
	return isCQTrybot
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"fmt"
	"github.com/FidelityInternational/chaos-galago/processor/Godeps/_workspace/src/chaos-galago/shared/utils"
	"github.com/FidelityInternational/chaos-galago/processor/Godeps/_workspace/src/github.com/cloudfoundry-community/go-cfclient"
	"github.com/FidelityInternational/chaos-galago/processor/utils"
	"os"
	"strconv"
	"time"
)

var (
	dbConnectionString string
	err                error
	config             *cfclient.Config
)

func init() {
	dbConnectionString, err = sharedUtils.GetDBConnectionDetails()
	if err != nil {
		fmt.Println(err.Error())
		os.Exit(1)
	}
	config = utils.LoadCFConfig()
	fmt.Println("\nConfig loaded:")
	fmt.Println("ApiAddress: ", config.ApiAddress)
	fmt.Println("LoginAddress: ", config.LoginAddress)
	fmt.Println("Username: ", config.Username)
	fmt.Println("SkipSslValidation: ", config.SkipSslValidation)
}

func freakOut(err error) bool {
	if err != nil {
		fmt.Println("An error has occured")
		fmt.Println(err.Error())
		return true
	}
	return false
}

func main() {
	cfClient := cfclient.NewClient(config)

	ticker := time.NewTicker(1 * time.Minute)

	processServices(cfClient)
	for _ = range ticker.C {
		processServices(cfClient)
	}
}

func processServices(cfClient *cfclient.Client) {
	db, err := sql.Open("mysql", dbConnectionString)
	if freakOut(err) {
		db.Close()
		return
	}
	services := utils.GetBoundApps(db)
	if len(services) == 0 {
		db.Close()
		return
	}

SERVICES:
	for _, service := range services {
		if utils.ShouldProcess(service.Frequency, service.LastProcessed) {
			fmt.Printf("\nProcessing chaos for %s", service.AppID)
			err = utils.UpdateLastProcessed(db, service.AppID, utils.TimeNow())
			if freakOut(err) {
				continue SERVICES
			}
			if utils.ShouldRun(service.Probability) {
				fmt.Printf("\nRunning chaos for %s", service.AppID)
				appInstances := cfClient.GetAppInstances(service.AppID)
				if utils.IsAppHealthy(appInstances) {
					fmt.Printf("\nApp %s is Healthy\n", service.AppID)
					chaosInstance := strconv.Itoa(utils.PickAppInstance(appInstances))
					fmt.Printf("\nAbout to kill app instance: %s at index: %s", service.AppID, chaosInstance)
					cfClient.KillAppInstance(service.AppID, chaosInstance)
					err = utils.UpdateLastProcessed(db, service.AppID, utils.TimeNow())
					if freakOut(err) {
						continue SERVICES
					}
				} else {
					fmt.Printf("\nApp %s is unhealthy, skipping\n", service.AppID)
					continue SERVICES
				}
			} else {
				fmt.Printf("\nNot running chaos for %s", service.AppID)
				err = utils.UpdateLastProcessed(db, service.AppID, utils.TimeNow())
				if freakOut(err) {
					continue SERVICES
				}
			}
		} else {
			fmt.Printf("\nSkipping processing chaos for %s", service.AppID)
			continue SERVICES
		}
	}
	db.Close()
}
<commit_msg>remove loop label that is no longer needed<commit_after>package main

import (
	"database/sql"
	"fmt"
	"github.com/FidelityInternational/chaos-galago/processor/Godeps/_workspace/src/chaos-galago/shared/utils"
	"github.com/FidelityInternational/chaos-galago/processor/Godeps/_workspace/src/github.com/cloudfoundry-community/go-cfclient"
	"github.com/FidelityInternational/chaos-galago/processor/utils"
	"os"
	"strconv"
	"time"
)

var (
	dbConnectionString string
	err                error
	config             *cfclient.Config
)

func init() {
	dbConnectionString, err = sharedUtils.GetDBConnectionDetails()
	if err != nil {
		fmt.Println(err.Error())
		os.Exit(1)
	}
	config = utils.LoadCFConfig()
	fmt.Println("\nConfig loaded:")
	fmt.Println("ApiAddress: ", config.ApiAddress)
	fmt.Println("LoginAddress: ", config.LoginAddress)
	fmt.Println("Username: ", config.Username)
	fmt.Println("SkipSslValidation: ", config.SkipSslValidation)
}

func freakOut(err error) bool {
	if err != nil {
		fmt.Println("An error has occured")
		fmt.Println(err.Error())
		return true
	}
	return false
}

func main() {
	cfClient := cfclient.NewClient(config)

	ticker := time.NewTicker(1 * time.Minute)

	processServices(cfClient)
	for _ = range ticker.C {
		processServices(cfClient)
	}
}

func processServices(cfClient *cfclient.Client) {
	db, err := sql.Open("mysql", dbConnectionString)
	if freakOut(err) {
		db.Close()
		return
	}
	services := utils.GetBoundApps(db)
	if len(services) == 0 {
		db.Close()
		return
	}

	for _, service := range services {
		if utils.ShouldProcess(service.Frequency, service.LastProcessed) {
			fmt.Printf("\nProcessing chaos for %s", service.AppID)
			err = utils.UpdateLastProcessed(db, service.AppID, utils.TimeNow())
			if freakOut(err) {
				continue
			}
			if utils.ShouldRun(service.Probability) {
				fmt.Printf("\nRunning chaos for %s", service.AppID)
				appInstances := cfClient.GetAppInstances(service.AppID)
				if utils.IsAppHealthy(appInstances) {
					fmt.Printf("\nApp %s is Healthy\n", service.AppID)
					chaosInstance := strconv.Itoa(utils.PickAppInstance(appInstances))
					fmt.Printf("\nAbout to kill app instance: %s at index: %s", service.AppID, chaosInstance)
					cfClient.KillAppInstance(service.AppID, chaosInstance)
					err = utils.UpdateLastProcessed(db, service.AppID, utils.TimeNow())
					if freakOut(err) {
						continue
					}
				} else {
					fmt.Printf("\nApp %s is unhealthy, skipping\n", service.AppID)
					continue
				}
			} else {
				fmt.Printf("\nNot running chaos for %s", service.AppID)
				err = utils.UpdateLastProcessed(db, service.AppID, utils.TimeNow())
				if freakOut(err) {
					continue
				}
			}
		} else {
			fmt.Printf("\nSkipping processing chaos for %s", service.AppID)
			continue
		}
	}
	db.Close()
}
<|endoftext|>


<commit_before>package main

//author: Doug Watson

//macOS will need more open files than the default 256 if you want to run over a couple hundred goroutines
//launchctl limit maxfiles 10200

//For a really big test, if you need a million open files on a mac:
//nvram boot-args="serverperfmode=1"
//shutdown -r now
//launchctl limit maxfiles 999990
//ulimit -n 999998

//to build"
//BUILD=`git rev-parse HEAD`
//GOOS=linux go build -o goRunner.linux -ldflags "-s -w -X main.Build=${BUILD}"
import (
	"bufio"
	"flag"
	"fmt"
	"net/http"
	"os"
	"os/signal"
	"runtime"
	"runtime/pprof"
	"strconv"
	"strings"
	"syscall"
	"time"
)

// -------------------------------------------------------------------------------------------------
// Flags
var (
	clients        int
	targetTPS      float64
	baseUrl        string
	configFile     string
	inputFile      string
	delimeter      string
	headerExit     bool
	noHeader       bool
	cpuProfile     string
	verbose        bool
	keepAlive      bool
	testTimeout    time.Duration
	readTimeout    time.Duration
	rampUp         time.Duration
	listenPort     int
	trafficChannel chan string
)

func init() {
	flag.IntVar(&clients, "c", 100, "Number of concurrent clients to launch.")
	flag.DurationVar(&testTimeout, "t", 0, "Timed load test duration (1m23s, 240s, ..). Defaults no timeout.")
	flag.StringVar(&inputFile, "f", "", "Read input from file rather than stdin")
	flag.DurationVar(&rampUp, "rampUp", -1, "Specify ramp up delay as duration (1m2s, 300ms, 0 ..). Default will auto compute from client sessions.")
	flag.Float64Var(&targetTPS, "targetTPS", 1000000, "The default max TPS is set to 1 million. Good luck reaching this :p")
	flag.StringVar(&baseUrl, "baseUrl", "", "The host to test. Example https://test2.someserver.org")
	flag.StringVar(&configFile, "configFile", "config.ini", "Config file location")
	flag.StringVar(&delimeter, "delimeter", ",", "Delimeter for output csv and input file")
	flag.BoolVar(&headerExit, "hx", false, "Print output header row and exit")
	flag.BoolVar(&noHeader, "nh", false, "Don't output header row. Default to false.")
	flag.DurationVar(&readTimeout, "readtimeout", time.Duration(30)*time.Second, "Timeout duration for the target API to send the first response byte. Default 30s")
	flag.StringVar(&cpuProfile, "cpuprofile", "", "write cpu profile to file")
	flag.BoolVar(&verbose, "verbose", false, "verbose debugging output flag")
	flag.BoolVar(&keepAlive, "keepalive", true, "enable/disable keepalive")
	flag.IntVar(&listenPort, "p", 0, "Default off. Port to listen on for input (as opposed to STDIN). HTTP GET or POST calls accepted i.e http://localhost/john,pass1\nor\ncurl POST http://localhost -d 'john,pass1\ndoug,pass2\n'")
}

// -------------------------------------------------------------------------------------------------
// Build commit id from git
var Build string

func init() {
	if Build == "" {
		Build = "unset"
	}

	defaultUsage := flag.Usage

	flag.Usage = func() {
		defaultUsage()
	}
}

func main() {
	// ---------------------------------------------------------------------------------------------
	// Parse flags
	flag.Parse()

	if headerExit {
		PrintLogHeader(delimeter, 0)
		os.Exit(0)
	}

	// ---------------------------------------------------------------------------------------------
	// Validate input & flags
	if clients < 1 {
		flagError("Number of concurrent client should be at least 1")
	}

	if cpuProfile != "" {
		f, err := os.Create(cpuProfile)
		if err != nil {
			flagError(err.Error())
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}
	if !headerExit && baseUrl == "" {
		flagError("Please provide the baseUrl")
	}

	if os.Getenv("GOMAXPROCS") == "" {
		runtime.GOMAXPROCS(runtime.NumCPU())
	}

	if verbose {
		println("verbose: build=", Build)
	}

	// ---------------------------------------------------------------------------------------------
	// Init runner
	runner := NewRunner(configFile)

	// ---------------------------------------------------------------------------------------------
	// Catch interrupt
	signalChannel := make(chan os.Signal, 2)
	signal.Notify(signalChannel, os.Interrupt, syscall.SIGTERM)
	go func() {
		_ = <-signalChannel
		runner.Exit()
	}()

	// ---------------------------------------------------------------------------------------------
	// Start clients
	trafficChannel = make(chan string)
	//	startTraffic(trafficChannel) //start reading on the channel
	runner.StartClients(trafficChannel)

	// ---------------------------------------------------------------------------------------------
	// Read input from file or stdin

	if listenPort > 0 {
		listenPortString := strconv.Itoa(listenPort)

		http.HandleFunc("/", HandleInputArgs)
		http.ListenAndServe(":"+listenPortString, nil)
	} else {

		scanner := bufio.NewScanner(os.Stdin)
		if len(inputFile) > 0 {
			file, err := os.Open(inputFile)
			if err != nil {
				flagError(err.Error())
			}
			defer file.Close()
			scanner = bufio.NewScanner(file)
		}
		// ---------------------------------------------------------------------------------------------
		// Output
		nbDelimeters := 0
		firstTime := true
		for scanner.Scan() {
			fmt.Printf("scanner.Scan\n")
			inputLine := scanner.Text()
			if firstTime {
				fmt.Printf("START FIRST_TIME\n")
				firstTime = false
				nbDelimeters = strings.Count(inputLine, delimeter)
				runner.printSessionSummary()
				if !noHeader {
					PrintLogHeader(delimeter, nbDelimeters+1)
					runner.PrintSessionLog() // ???
				}
				fmt.Printf("END FIRST_TIME\n")
			}
			if strings.Count(inputLine, delimeter) != nbDelimeters {
				fmt.Fprintf(os.Stderr, "\n/!\\ input lines must have same number of fields /!\\\n")
				runner.Exit()
			}
			if len(inputLine) == 0 {
				break //quit when we get an empty input line
			}
			fmt.Printf("inputLine=%v\n\n", inputLine)
			trafficChannel <- inputLine
		}
	}
	close(trafficChannel)

	// ---------------------------------------------------------------------------------------------
	// Wait for clients to be done and exit
	runner.Wait()
	runner.Exit()
}
func HandleInputArgs(w http.ResponseWriter, r *http.Request) {
	scanner := bufio.NewScanner(r.Body)
	for scanner.Scan() {
		inputLine := scanner.Text()
		trafficChannel <- inputLine
	}
	w.WriteHeader(200)
}

func flagError(err string) {
	flag.Usage()
	fmt.Fprintf(os.Stderr, "\n%s\n\n", err)
	os.Exit(1)
}

//func noRedirect(req *http.Request, via []*http.Request) error {
//	return errors.New("Don't redirect!")
//}
<commit_msg>print build<commit_after>package main

//author: Doug Watson

//macOS will need more open files than the default 256 if you want to run over a couple hundred goroutines
//launchctl limit maxfiles 10200

//For a really big test, if you need a million open files on a mac:
//nvram boot-args="serverperfmode=1"
//shutdown -r now
//launchctl limit maxfiles 999990
//ulimit -n 999998

//to build"
//BUILD=`git rev-parse HEAD`
//GOOS=linux go build -o goRunner.linux -ldflags "-s -w -X main.Build=${BUILD}"
import (
	"bufio"
	"flag"
	"fmt"
	"net/http"
	"os"
	"os/signal"
	"runtime"
	"runtime/pprof"
	"strconv"
	"strings"
	"syscall"
	"time"
)

// -------------------------------------------------------------------------------------------------
// Flags
var (
	clients        int
	targetTPS      float64
	baseUrl        string
	configFile     string
	inputFile      string
	delimeter      string
	headerExit     bool
	noHeader       bool
	cpuProfile     string
	verbose        bool
	keepAlive      bool
	testTimeout    time.Duration
	readTimeout    time.Duration
	rampUp         time.Duration
	listenPort     int
	trafficChannel chan string
)

func init() {
	flag.IntVar(&clients, "c", 100, "Number of concurrent clients to launch.")
	flag.DurationVar(&testTimeout, "t", 0, "Timed load test duration (1m23s, 240s, ..). Defaults no timeout.")
	flag.StringVar(&inputFile, "f", "", "Read input from file rather than stdin")
	flag.DurationVar(&rampUp, "rampUp", -1, "Specify ramp up delay as duration (1m2s, 300ms, 0 ..). Default will auto compute from client sessions.")
	flag.Float64Var(&targetTPS, "targetTPS", 1000000, "The default max TPS is set to 1 million. Good luck reaching this :p")
	flag.StringVar(&baseUrl, "baseUrl", "", "The host to test. Example https://test2.someserver.org")
	flag.StringVar(&configFile, "configFile", "config.ini", "Config file location")
	flag.StringVar(&delimeter, "delimeter", ",", "Delimeter for output csv and input file")
	flag.BoolVar(&headerExit, "hx", false, "Print output header row and exit")
	flag.BoolVar(&noHeader, "nh", false, "Don't output header row. Default to false.")
	flag.DurationVar(&readTimeout, "readtimeout", time.Duration(30)*time.Second, "Timeout duration for the target API to send the first response byte. Default 30s")
	flag.StringVar(&cpuProfile, "cpuprofile", "", "write cpu profile to file")
	flag.BoolVar(&verbose, "verbose", false, "verbose debugging output flag")
	flag.BoolVar(&keepAlive, "keepalive", true, "enable/disable keepalive")
	flag.IntVar(&listenPort, "p", 0, "Default off. Port to listen on for input (as opposed to STDIN). HTTP GET or POST calls accepted i.e http://localhost/john,pass1\nor\ncurl POST http://localhost -d 'john,pass1\ndoug,pass2\n'")
}

// -------------------------------------------------------------------------------------------------
// Build commit id from git
var Build string

func init() {
	if Build == "" {
		Build = "unset"
	}

	defaultUsage := flag.Usage

	flag.Usage = func() {
		println("github.com/adt-automation/goRunner ", Build)
		defaultUsage()
	}
}

func main() {
	// ---------------------------------------------------------------------------------------------
	// Parse flags
	flag.Parse()

	if headerExit {
		PrintLogHeader(delimeter, 0)
		os.Exit(0)
	}
	if *verbose {
		println("Build #", Build)
	}

	// ---------------------------------------------------------------------------------------------
	// Validate input & flags
	if clients < 1 {
		flagError("Number of concurrent client should be at least 1")
	}

	if cpuProfile != "" {
		f, err := os.Create(cpuProfile)
		if err != nil {
			flagError(err.Error())
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}
	if !headerExit && baseUrl == "" {
		flagError("Please provide the baseUrl")
	}

	if os.Getenv("GOMAXPROCS") == "" {
		runtime.GOMAXPROCS(runtime.NumCPU())
	}

	// ---------------------------------------------------------------------------------------------
	// Init runner
	runner := NewRunner(configFile)

	// ---------------------------------------------------------------------------------------------
	// Catch interrupt
	signalChannel := make(chan os.Signal, 2)
	signal.Notify(signalChannel, os.Interrupt, syscall.SIGTERM)
	go func() {
		_ = <-signalChannel
		runner.Exit()
	}()

	// ---------------------------------------------------------------------------------------------
	// Start clients
	trafficChannel = make(chan string)
	//	startTraffic(trafficChannel) //start reading on the channel
	runner.StartClients(trafficChannel)

	// ---------------------------------------------------------------------------------------------
	// Read input from file or stdin

	if listenPort > 0 {
		listenPortString := strconv.Itoa(listenPort)

		http.HandleFunc("/", HandleInputArgs)
		http.ListenAndServe(":"+listenPortString, nil)
	} else {

		scanner := bufio.NewScanner(os.Stdin)
		if len(inputFile) > 0 {
			file, err := os.Open(inputFile)
			if err != nil {
				flagError(err.Error())
			}
			defer file.Close()
			scanner = bufio.NewScanner(file)
		}
		// ---------------------------------------------------------------------------------------------
		// Output
		nbDelimeters := 0
		firstTime := true
		for scanner.Scan() {
			fmt.Printf("scanner.Scan\n")
			inputLine := scanner.Text()
			if firstTime {
				fmt.Printf("START FIRST_TIME\n")
				firstTime = false
				nbDelimeters = strings.Count(inputLine, delimeter)
				runner.printSessionSummary()
				if !noHeader {
					PrintLogHeader(delimeter, nbDelimeters+1)
					runner.PrintSessionLog() // ???
				}
				fmt.Printf("END FIRST_TIME\n")
			}
			if strings.Count(inputLine, delimeter) != nbDelimeters {
				fmt.Fprintf(os.Stderr, "\n/!\\ input lines must have same number of fields /!\\\n")
				runner.Exit()
			}
			if len(inputLine) == 0 {
				break //quit when we get an empty input line
			}
			fmt.Printf("inputLine=%v\n\n", inputLine)
			trafficChannel <- inputLine
		}
	}
	close(trafficChannel)

	// ---------------------------------------------------------------------------------------------
	// Wait for clients to be done and exit
	runner.Wait()
	runner.Exit()
}
func HandleInputArgs(w http.ResponseWriter, r *http.Request) {
	scanner := bufio.NewScanner(r.Body)
	for scanner.Scan() {
		inputLine := scanner.Text()
		trafficChannel <- inputLine
	}
	w.WriteHeader(200)
}

func flagError(err string) {
	flag.Usage()
	fmt.Fprintf(os.Stderr, "\n%s\n\n", err)
	os.Exit(1)
}

//func noRedirect(req *http.Request, via []*http.Request) error {
//	return errors.New("Don't redirect!")
//}
<|endoftext|>


<commit_before>package main

import (
	"net/http"
	"github.com/hailiang/gosocks"
	"io/ioutil"
	"runtime"
	"fmt"
	"time"
	"flag"
	"strings"
	"crypto/sha1"
	"io"
	"encoding/json"
	"encoding/base64"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"bytes"
)

// types
type Msg struct {
        Iv      string `json:"iv,omitempty"`
        Message string `json:"message"`
}

type MsgWrap struct {
        Content   Msg       `json:"content"`
        Timestamp string    `json:"timestamp"`
        Time      time.Time `json:"-"`
        Id        int       `bson:"_id"`
}

type Counter struct {
        Id string `bson:"_id"`
        N  int
}

type LastMessage struct {
        LastMessage int `json:"lastMessage"`
}

type Circle struct {
	Shortname string
	Server string
	Uuid string
	Key string
}

func (circle *Circle) getUrlHash() string {
	h := sha1.New()
	if circle.Uuid == "" {
		io.WriteString(h, circle.Key)
	} else {
		io.WriteString(h, circle.Uuid)
	}
	return(fmt.Sprintf("%x", h.Sum(nil)))
}

func (circle *Circle) getKeyData() []byte {
	var keyData []byte
	if len(circle.Key) != 16 {
		keyData,_ = base64.StdEncoding.DecodeString(circle.Key)
	} else {
		keyData = []byte(circle.Key)
	}

	return keyData
}

func (msg *Msg) getIVData() []byte {
	var iv []byte
	if msg.Iv == "" {
		iv = []byte{'0','1','2','3','4','5','6','7','0','1','2','3','4','5','6','7'}
	} else {
		iv,_ = base64.StdEncoding.DecodeString(msg.Iv)
	}
	return iv
}

func (msg *Msg) genIVData() []byte {
	rb := make([]byte,16)
	_,err := rand.Read(rb)
	if err != nil {
		fmt.Printf("Failed to get random data: %s\n", err)
	}
	msg.Iv = base64.StdEncoding.EncodeToString(rb)
	return rb
}

// only needed by AES?
func pkcs5pad(data []byte, blocksize int) []byte {
	pad := blocksize - len(data)%blocksize
	b := make([]byte, pad, pad)
	for i := 0; i < pad; i++ {
		b[i] = uint8(pad)
	}
	return append(data, b...)
}

func newCircle(circleString string) *Circle {
	plusIndx := strings.Index(circleString,"+")
	sigilIndx := strings.Index(circleString,"$")
	atIndx := strings.Index(circleString,"@")

	var circle *Circle

	parsedCircle := []string{"","","",""}
	if plusIndx == -1 || atIndx == -1 {
		return circle
	}

	if sigilIndx == -1 {
		parsedCircle[0] = circleString[0:plusIndx]
		parsedCircle[1] = ""
		parsedCircle[2] = circleString[plusIndx+1:atIndx]
		parsedCircle[3] = circleString[atIndx+1:len(circleString)]
	} else {
		parsedCircle[0] = circleString[0:plusIndx]
		parsedCircle[1] = circleString[plusIndx+1:sigilIndx]
		parsedCircle[2] = circleString[sigilIndx+1:atIndx]
		parsedCircle[3] = circleString[atIndx+1:len(circleString)]
	}

	//fmt.Printf("Key %s\n", parsedCircle[2])
	//fmt.Printf("Uuid %s\n", parsedCircle[1])
	circle = &Circle{Shortname: parsedCircle[0], Server: parsedCircle[3], Uuid: parsedCircle[1], Key: parsedCircle[2]}
	return circle
}

func main() {

	runtime.GOMAXPROCS(runtime.NumCPU())

	var circleString string
	var cmdString string
	var arg1String string
	flag.StringVar(&circleString, "circle", "fortune+fb8eb78595536a4d@tckwndlytrphlpyo.onion", "The circle to use for cmds.")
	flag.StringVar(&cmdString, "cmd", "read", "The command to use.")
	flag.StringVar(&arg1String, "arg1", "1", "First argument")
	flag.Parse()

	//fmt.Printf("ARg1: %s\n", os.Args[0])

	dialSocksProxy := socks.DialSocksProxy(socks.SOCKS4A, "127.0.0.1:9050")
	tr := &http.Transport{Dial: dialSocksProxy}
	httpClient := &http.Client{Transport: tr}

	circle := newCircle(circleString)

	var r *http.Response
	var err error
	//fmt.Printf("Hash: %s\n", circle.getUrlHash())
	if cmdString == "last" {
		r,err = httpClient.Get(fmt.Sprintf("http://%s/%s",circle.Server,circle.getUrlHash()))
	} else if cmdString == "read" {
		r,err = httpClient.Get(fmt.Sprintf("http://%s/%s/%s",circle.Server,circle.getUrlHash(),arg1String))
		lastModified := r.Header.Get("Last-Modified")
		fmt.Printf("Date: %s\n",lastModified)
	} else if cmdString == "post" {
		var msg Msg

		c,_ := aes.NewCipher(circle.getKeyData())
		ivBytes := msg.genIVData()
		encrypter := cipher.NewCBCEncrypter(c, ivBytes)
		//fmt.Printf("Iv data is: %s\n",msg.Iv)

		//fmt.Printf("Key len: %d\n",len(circle.getKeyData()))
		//fmt.Printf("Iv len: %d\n",len(ivBytes))
		//fmt.Printf("Input len: %d\n",len([]byte(arg1String)))

		enctext := make([]byte, len(pkcs5pad([]byte(arg1String),16)))
		encrypter.CryptBlocks(enctext, pkcs5pad([]byte(arg1String),16))
		msg.Message = base64.StdEncoding.EncodeToString(enctext)
		msgBytes,_ := json.Marshal(msg)
		r := bytes.NewReader(msgBytes)
		resp,err := httpClient.Post(fmt.Sprintf("http://%s/%s",circle.Server,circle.getUrlHash()),"application/json",r)
		if err != nil {
			panic(fmt.Sprintf("Failed to post message: %s\n", err))
		}

		fmt.Printf("Response: %s\n", resp.Status)
		return

	}

	if err != nil {
		panic(fmt.Sprintf("Error: %s\n",err))
	}

	bytes,err := ioutil.ReadAll(r.Body)
	if err != nil {
		panic(fmt.Sprintf("Error: %s\n",err))
	}

	var msg Msg
	if cmdString == "read" {
		json.Unmarshal(bytes,&msg)

		//fmt.Printf("Msg: %s\n", msg.Message)
		//fmt.Printf("Iv: %s\n", msg.Iv)

		rawdata,_ := base64.StdEncoding.DecodeString(msg.Message)
		b := rawdata
		c,_ := aes.NewCipher(circle.getKeyData())
		decrypter := cipher.NewCBCDecrypter(c, msg.getIVData())
		plaintext := make([]byte, len(b))
		decrypter.CryptBlocks(plaintext, b)
		fmt.Printf("%s\n",string(plaintext))
	} else {
		fmt.Printf("%s\n",bytes)
	}
}
<commit_msg>Allow message data to be passed via stdin.<commit_after>package main

import (
	"net/http"
	"github.com/hailiang/gosocks"
	"io/ioutil"
	"runtime"
	"fmt"
	"time"
	"flag"
	"strings"
	"crypto/sha1"
	"io"
	"encoding/json"
	"encoding/base64"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"bytes"
	"os"
	"bufio"
)

// types
type Msg struct {
        Iv      string `json:"iv,omitempty"`
        Message string `json:"message"`
}

type MsgWrap struct {
        Content   Msg       `json:"content"`
        Timestamp string    `json:"timestamp"`
        Time      time.Time `json:"-"`
        Id        int       `bson:"_id"`
}

type Counter struct {
        Id string `bson:"_id"`
        N  int
}

type LastMessage struct {
        LastMessage int `json:"lastMessage"`
}

type Circle struct {
	Shortname string
	Server string
	Uuid string
	Key string
}

func (circle *Circle) getUrlHash() string {
	h := sha1.New()
	if circle.Uuid == "" {
		io.WriteString(h, circle.Key)
	} else {
		io.WriteString(h, circle.Uuid)
	}
	return(fmt.Sprintf("%x", h.Sum(nil)))
}

func (circle *Circle) getKeyData() []byte {
	var keyData []byte
	if len(circle.Key) != 16 {
		keyData,_ = base64.StdEncoding.DecodeString(circle.Key)
	} else {
		keyData = []byte(circle.Key)
	}

	return keyData
}

func (msg *Msg) getIVData() []byte {
	var iv []byte
	if msg.Iv == "" {
		iv = []byte{'0','1','2','3','4','5','6','7','0','1','2','3','4','5','6','7'}
	} else {
		iv,_ = base64.StdEncoding.DecodeString(msg.Iv)
	}
	return iv
}

func (msg *Msg) genIVData() []byte {
	rb := make([]byte,16)
	_,err := rand.Read(rb)
	if err != nil {
		fmt.Printf("Failed to get random data: %s\n", err)
	}
	msg.Iv = base64.StdEncoding.EncodeToString(rb)
	return rb
}

// only needed by AES?
func pkcs5pad(data []byte, blocksize int) []byte {
	pad := blocksize - len(data)%blocksize
	b := make([]byte, pad, pad)
	for i := 0; i < pad; i++ {
		b[i] = uint8(pad)
	}
	return append(data, b...)
}

func newCircle(circleString string) *Circle {
	plusIndx := strings.Index(circleString,"+")
	sigilIndx := strings.Index(circleString,"$")
	atIndx := strings.Index(circleString,"@")

	var circle *Circle

	parsedCircle := []string{"","","",""}
	if plusIndx == -1 || atIndx == -1 {
		return circle
	}

	if sigilIndx == -1 {
		parsedCircle[0] = circleString[0:plusIndx]
		parsedCircle[1] = ""
		parsedCircle[2] = circleString[plusIndx+1:atIndx]
		parsedCircle[3] = circleString[atIndx+1:len(circleString)]
	} else {
		parsedCircle[0] = circleString[0:plusIndx]
		parsedCircle[1] = circleString[plusIndx+1:sigilIndx]
		parsedCircle[2] = circleString[sigilIndx+1:atIndx]
		parsedCircle[3] = circleString[atIndx+1:len(circleString)]
	}

	//fmt.Printf("Key %s\n", parsedCircle[2])
	//fmt.Printf("Uuid %s\n", parsedCircle[1])
	circle = &Circle{Shortname: parsedCircle[0], Server: parsedCircle[3], Uuid: parsedCircle[1], Key: parsedCircle[2]}
	return circle
}

/*
func dumpData (one, two []byte) {
	for i,k := range two {
		fmt.Printf("%d: %b\n", i,k)
	}
	for i,k := range one {
		fmt.Printf("%d: %b\n", i,k)
	}
}
*/

func main() {

	runtime.GOMAXPROCS(runtime.NumCPU())

	var circleString string
	var cmdString string
	var arg1String string
	flag.StringVar(&circleString, "circle", "fortune+fb8eb78595536a4d@tckwndlytrphlpyo.onion", "The circle to use for cmds.")
	flag.StringVar(&cmdString, "cmd", "read", "The command to use.")
	flag.StringVar(&arg1String, "arg1", "1", "First argument")
	flag.Parse()

	//fmt.Printf("ARg1: %s\n", os.Args[0])

	dialSocksProxy := socks.DialSocksProxy(socks.SOCKS4A, "127.0.0.1:9050")
	tr := &http.Transport{Dial: dialSocksProxy}
	httpClient := &http.Client{Transport: tr}

	circle := newCircle(circleString)

	var err error
	var stdin io.Reader

	if arg1String == "-" {
	   stdin = bufio.NewReader(os.Stdin)
	   b,err := ioutil.ReadAll(stdin)
	   if err != nil {
		fmt.Printf("Failed to read stdin data.\n")
		return
	   }
	   //fmt.Printf("bytes %s\n",string(b))
	   arg1String = string(b)
	}

	var r *http.Response
	//fmt.Printf("Hash: %s\n", circle.getUrlHash())
	if cmdString == "last" {
		r,err = httpClient.Get(fmt.Sprintf("http://%s/%s",circle.Server,circle.getUrlHash()))
	} else if cmdString == "read" {
		r,err = httpClient.Get(fmt.Sprintf("http://%s/%s/%s",circle.Server,circle.getUrlHash(),arg1String))
		lastModified := r.Header.Get("Last-Modified")
		fmt.Printf("Date: %s\n",lastModified)
	} else if cmdString == "post" {
		var msg Msg

		c,_ := aes.NewCipher(circle.getKeyData())
		ivBytes := msg.genIVData()
		encrypter := cipher.NewCBCEncrypter(c, ivBytes)
		//fmt.Printf("Iv data is: %s\n",msg.Iv)

		//fmt.Printf("Key len: %d\n",len(circle.getKeyData()))
		//fmt.Printf("Iv len: %d\n",len(ivBytes))
		//fmt.Printf("Input len: %d\n",len([]byte(arg1String)))

		enctext := make([]byte, len(pkcs5pad([]byte(arg1String),16)))
		encrypter.CryptBlocks(enctext, pkcs5pad([]byte(arg1String),16))
		msg.Message = base64.StdEncoding.EncodeToString(enctext)
		msgBytes,_ := json.Marshal(msg)
		r := bytes.NewReader(msgBytes)
		resp,err := httpClient.Post(fmt.Sprintf("http://%s/%s",circle.Server,circle.getUrlHash()),"application/json",r)
		if err != nil {
			panic(fmt.Sprintf("Failed to post message: %s\n", err))
		}

		fmt.Printf("Response: %s\n", resp.Status)
		return

	}

	if err != nil {
		panic(fmt.Sprintf("Error: %s\n",err))
	}

	bytes,err := ioutil.ReadAll(r.Body)
	if err != nil {
		panic(fmt.Sprintf("Error: %s\n",err))
	}

	var msg Msg
	if cmdString == "read" {
		json.Unmarshal(bytes,&msg)

		//fmt.Printf("Msg: %s\n", msg.Message)
		//fmt.Printf("Iv: %s\n", msg.Iv)

		rawdata,_ := base64.StdEncoding.DecodeString(msg.Message)
		b := rawdata
		c,_ := aes.NewCipher(circle.getKeyData())
		decrypter := cipher.NewCBCDecrypter(c, msg.getIVData())
		plaintext := make([]byte, len(b))
		decrypter.CryptBlocks(plaintext, b)
		fmt.Printf("%s\n",string(plaintext))
	} else {
		fmt.Printf("%s\n",bytes)
	}
}
<|endoftext|>


<commit_before>package gofigure

import (
	"github.com/droundy/goopt"
	"log"
	"os"
)

// This function is a convenience function. It displays the default values for us but handles them
// in the way we prefer it.
func GooptFigureString(names []string, def string, help string) *string {
	s := new(string)
	*s = ""
	f := func (ss string) error {
		*s = ss
		return nil
	}
	goopt.ReqArg(names, def, help, f)
	return s
}

// Config contains the configuration options that may be set by
// command line flags and environment variables.
type Config struct {
	Description			string
	DisableCommandLine	bool
	EnvPrefix			string
	EnvOverridesFile	bool
	FileParser			File
	Version				string
	options				map[string]*Option
	flags				map[string]*string
	values				map[string]*string
}

// Returns a new Config instance.
func New() *Config {
	return &Config{
		DisableCommandLine:	false,
		EnvOverridesFile:   false,
		options:			make(map[string]*Option),
		flags:				make(map[string]*string),
		values:				make(map[string]*string),
	}
}

// Adds a configuration option, returns an Option instance for
// easily setting the corresponding environment variable, default
// value, and description.
func (c *Config) Add(name string) *Option {
	c.options[name] = &Option{
		name:    name,
		envVar:  "",
		def:     "",
		desc:    "",
		longOpt: name,
	}
	return c.options[name]
}

// Returns a configuration option by flag name.
func (c Config) Get(name string) *string {
	return c.values[name]
}

// Parses the configuration options into defined flags, sets the value
// accordingly. Options are read first from command line flags, then
// from environment variables, and falls back to the default value if
// neither are set.
//
// See https://github.com/rakyll/globalconf/blob/master/globalconf.go
func (c *Config) Parse() {
	goopt.Description = func() string {
		return c.Description
	}

	goopt.Version = c.Version

	// Sets the flags from the configuration options.
	for name, o := range c.options {
		cmdline := []string{}
		if o.shortOpt != "" {
			cmdline = append(cmdline, "-" + o.shortOpt)
		}
		cmdline = append(cmdline, "--"+o.longOpt)
		c.flags[name] = GooptFigureString(cmdline, o.def, o.desc)
		defcopy := o.def
		c.values[name] = &defcopy
	}

	passed := make(map[string]bool)

	if !c.DisableCommandLine {
		goopt.Parse(nil)

		// Gather the options passed through command line.
		for name, f := range c.flags {
			if *f != "" {
				passed[name] = true
				*c.values[name] = *f
			}
		}
	}

	if (c.EnvOverridesFile) {
		c.ParseEnv(passed)
		err := c.ParseFile(passed)
		if err != nil {
			log.Panicf("File defined but could not be parsed: %s", err.Error())
		}
	} else {
		err := c.ParseFile(passed)
		if err != nil {
			log.Panicf("File defined but could not be parsed: %s", err.Error())
		}
		c.ParseEnv(passed)
	}
}

func (c *Config) ParseEnv(passed map[string]bool) {
	for name, f := range c.options {

		// Skip flags passed through the command line as the option is
		// already set and takes precedence over environment variables.
		if passed[name] {
			continue
		}

		// Some options shouldn't be set via environment variables.
		if f.envVar == "" {
			continue
		}

		// If the configuration option was not passed via the command line,
		// check the corresponding environment variable.
		envVar := c.EnvPrefix + f.envVar
		if val := os.Getenv(envVar); val != "" {
			*c.values[name] = val
			passed[name] = true
		}
	}
}

func (c Config) parseFileToMap(handler File) (ValueMap, error) {
	root, err := handler.Parse()
	if err != nil {
		return nil, err
	}
	ret := ValueMap{}
	for _, opt := range c.options {
		if val, ok := root.FindOption(opt); ok {
			ret.Set(opt.Name(), val)
		}
	}
	return ret, nil
}

func (c *Config) ParseFile(passed map[string]bool) error {
	if (c.FileParser == nil) {
		return nil
	}

	values, err := c.parseFileToMap(c.FileParser)
	if err != nil {
		return err
	}

	for k, v := range values {
		if passed[k] {
			continue
		}

		if p, ok := c.values[k]; ok {
			*p = v
			passed[k] = true
		}
	}
	return nil
}

// Option contains the details of a configuration options,
// e.g. corresponding environment variable, default value,
// description.
type Option struct {
	name, envVar, shortOpt, def, desc, longOpt string
	fileSpec				string              // The file spec is of the form "(CATEGORY.)*NAME", eg. for 'foo' under the category 'bar', it would be foo.bar
}

func (o Option) Name() string {
	return o.name
}

func (o *Option) FileSpec(spec string) *Option {
	o.fileSpec = spec
	return o
}

func (o *Option) ShortOpt(opt string) *Option {
	o.shortOpt = opt
	return o
}

func (o *Option) LongOpt(opt string) *Option {
	o.longOpt = opt
	return o
}

// Sets the configuration option's default value.
func (o *Option) Default(def string) *Option {
	o.def = def
	return o
}

// Sets the configuration option's corresponding environment variable.
func (o *Option) EnvVar(envVar string) *Option {
	o.envVar = envVar
	return o
}

// Sets the configuration options long description.
func (o *Option) Description(desc string) *Option {
	o.desc = desc
	return o
}
<commit_msg>added the option to automatically attach environment description to --help/man<commit_after>package gofigure

import (
	"fmt"
	"github.com/droundy/goopt"
	"log"
	"os"
)

// This function is a convenience function. It displays the default values for us but handles them
// in the way we prefer it.
func GooptFigureString(names []string, def string, help string) *string {
	s := new(string)
	*s = ""
	f := func (ss string) error {
		*s = ss
		return nil
	}
	goopt.ReqArg(names, def, help, f)
	return s
}

// Config contains the configuration options that may be set by
// command line flags and environment variables.
type Config struct {
	Description			string
	DescribeEnvironment bool	// if true, the environment variable is automatically added to the flag description
	DisableCommandLine	bool
	EnvPrefix			string
	EnvOverridesFile	bool
	FileParser			File
	Version				string
	options				map[string]*Option
	flags				map[string]*string
	values				map[string]*string
}

// Returns a new Config instance.
func New() *Config {
	return &Config{
		DescribeEnvironment: false,
		DisableCommandLine:	 false,
		EnvOverridesFile:    false,
		options:			 make(map[string]*Option),
		flags:				 make(map[string]*string),
		values:				 make(map[string]*string),
	}
}

// Adds a configuration option, returns an Option instance for
// easily setting the corresponding environment variable, default
// value, and description.
func (c *Config) Add(name string) *Option {
	c.options[name] = &Option{
		name:    name,
		envVar:  "",
		def:     "",
		desc:    "",
		longOpt: name,
	}
	return c.options[name]
}

// Returns a configuration option by flag name.
func (c Config) Get(name string) *string {
	return c.values[name]
}

// Parses the configuration options into defined flags, sets the value
// accordingly. Options are read first from command line flags, then
// from environment variables, and falls back to the default value if
// neither are set.
//
// See https://github.com/rakyll/globalconf/blob/master/globalconf.go
func (c *Config) Parse() {
	goopt.Description = func() string {
		return c.Description
	}

	goopt.Version = c.Version

	// Sets the flags from the configuration options.
	for name, o := range c.options {
		cmdline := []string{}
		if o.shortOpt != "" {
			cmdline = append(cmdline, "-" + o.shortOpt)
		}
		cmdline = append(cmdline, "--"+o.longOpt)
		desc := o.desc
		if c.DescribeEnvironment && o.envVar != "" {
			desc += fmt.Sprintf(" Environment variable: %s_%s.", c.EnvPrefix, o.envVar)
		}
		c.flags[name] = GooptFigureString(cmdline, o.def, desc)
		defcopy := o.def
		c.values[name] = &defcopy
	}

	passed := make(map[string]bool)

	if !c.DisableCommandLine {
		goopt.Parse(nil)

		// Gather the options passed through command line.
		for name, f := range c.flags {
			if *f != "" {
				passed[name] = true
				*c.values[name] = *f
			}
		}
	}

	if (c.EnvOverridesFile) {
		c.ParseEnv(passed)
		err := c.ParseFile(passed)
		if err != nil {
			log.Panicf("File defined but could not be parsed: %s", err.Error())
		}
	} else {
		err := c.ParseFile(passed)
		if err != nil {
			log.Panicf("File defined but could not be parsed: %s", err.Error())
		}
		c.ParseEnv(passed)
	}
}

func (c *Config) ParseEnv(passed map[string]bool) {
	for name, f := range c.options {

		// Skip flags passed through the command line as the option is
		// already set and takes precedence over environment variables.
		if passed[name] {
			continue
		}

		// Some options shouldn't be set via environment variables.
		if f.envVar == "" {
			continue
		}

		// If the configuration option was not passed via the command line,
		// check the corresponding environment variable.
		envVar := c.EnvPrefix + f.envVar
		if val := os.Getenv(envVar); val != "" {
			*c.values[name] = val
			passed[name] = true
		}
	}
}

func (c Config) parseFileToMap(handler File) (ValueMap, error) {
	root, err := handler.Parse()
	if err != nil {
		return nil, err
	}
	ret := ValueMap{}
	for _, opt := range c.options {
		if val, ok := root.FindOption(opt); ok {
			ret.Set(opt.Name(), val)
		}
	}
	return ret, nil
}

func (c *Config) ParseFile(passed map[string]bool) error {
	if (c.FileParser == nil) {
		return nil
	}

	values, err := c.parseFileToMap(c.FileParser)
	if err != nil {
		return err
	}

	for k, v := range values {
		if passed[k] {
			continue
		}

		if p, ok := c.values[k]; ok {
			*p = v
			passed[k] = true
		}
	}
	return nil
}

// Option contains the details of a configuration options,
// e.g. corresponding environment variable, default value,
// description.
type Option struct {
	name, envVar, shortOpt, def, desc, longOpt string
	fileSpec				string              // The file spec is of the form "(CATEGORY.)*NAME", eg. for 'foo' under the category 'bar', it would be foo.bar
}

func (o Option) Name() string {
	return o.name
}

func (o *Option) FileSpec(spec string) *Option {
	o.fileSpec = spec
	return o
}

func (o *Option) ShortOpt(opt string) *Option {
	o.shortOpt = opt
	return o
}

func (o *Option) LongOpt(opt string) *Option {
	o.longOpt = opt
	return o
}

// Sets the configuration option's default value.
func (o *Option) Default(def string) *Option {
	o.def = def
	return o
}

// Sets the configuration option's corresponding environment variable.
func (o *Option) EnvVar(envVar string) *Option {
	o.envVar = envVar
	return o
}

// Sets the configuration options long description.
func (o *Option) Description(desc string) *Option {
	o.desc = desc
	return o
}
<|endoftext|>


<commit_before>// TODO: Handle user auth.
// TODO: Cache discovery/directory documents for faster requests.
// TODO: Handle media upload/download.
// TODO: Handle repeated parameters.

package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strings"

	"code.google.com/p/goauth2/oauth/jwt"
)

var (
	// Flags that get parsed before the command, necessary for loading Cloud Endpoints APIs
	// e.g., "googlecl --endpoint=foo help myapi" parses the endpoint flag before loading the API
	endpointFs   = flag.NewFlagSet("endpoint", flag.ExitOnError)
	flagEndpoint = endpointFs.String("endpoint", "https://www.googleapis.com/", "Cloud Endpoints URL, e.g., https://my-app-id.appspot.com/_ah/api/")

	// Flags that get parsed after the command, common to all APIs
	fs          = flag.NewFlagSet("googlecl", flag.ExitOnError)
	flagPem     = fs.String("meta.pem", "", "Location of .pem file")
	flagSecrets = fs.String("meta.secrets", "", "Location of client_secrets.json")
	flagInFile  = fs.String("meta.inFile", "", "File to pass as request body")
	flagStdin   = fs.Bool("meta.in", false, "Whether to use stdin as the request body")
)

func simpleHelp() {
	fmt.Println("Makes requests to Google APIs")
	fmt.Println("Usage:")
	fmt.Println("  googlecl <api> <method> --param=foo")
}

func help() {
	args := endpointFs.Args()
	nargs := len(args)
	if nargs == 0 || (nargs == 1 && args[0] == "help") {
		simpleHelp()
		return
	}
	apiName := args[1]
	api, err := loadAPI(apiName)
	if err != nil {
		log.Fatal(err)
	}
	if nargs == 2 {
		// googlecl help <api>
		fmt.Println(api.Title, api.Description)
		fmt.Println("More information:", api.DocumentationLink)
		fmt.Println("Methods:")
		for _, m := range api.Methods {
			fmt.Println(m.ID, m.Description)
		}
		type pair struct {
			k string
			r Resource
		}
		l := []pair{}
		for k, r := range api.Resources {
			l = append(l, pair{k, r})
		}
		for i := 0; i < len(l); i++ {
			r := l[i].r
			for _, m := range r.Methods {
				fmt.Printf("%s - %s\n", m.ID[len(api.Name)+1:], m.Description)
			}
			for k, r := range r.Resources {
				l = append(l, pair{k, r})
			}
		}
	} else {
		// googlecl help <api> <method>
		method := args[2]
		m := findMethod(method, *api)
		fmt.Println(method, m.Description)
		fmt.Println("Parameters:")
		for k, p := range m.Parameters {
			fmt.Printf("  --%s (%s) - %s\n", k, p.Type, p.Description)
		}
		for k, p := range api.Parameters {
			fmt.Printf("  --%s (%s) - %s\n", k, p.Type, p.Description)
		}

		s := api.Schemas[m.RequestSchema.Ref]
		// TODO: Support deep nested schemas, and use actual flags to get these strings to avoid duplication
		for k, p := range s.Properties {
			fmt.Printf("  --res.%s (%s) - %s\n", k, p.Type, p.Description)
		}
	}
}

func list() {
	var directory struct {
		Items []struct {
			Name, Version, Description string
		}
	}
	getAndParse("discovery/v1/apis", &directory)
	fmt.Println("Available methods:")
	for _, i := range directory.Items {
		fmt.Printf("%s %s - %s\n", i.Name, i.Version, i.Description)
	}
}

func main() {
	endpointFs.Parse(os.Args[1:])
	if len(endpointFs.Args()) == 0 {
		simpleHelp()
		return
	}

	cmd := endpointFs.Args()[0]
	if cmd == "help" {
		help()
		return
	} else if cmd == "list" {
		list()
		return
	}

	api, err := loadAPI(cmd)
	if err != nil {
		log.Fatal(err)
	}
	if api == nil || (len(api.Resources) == 0 && len(api.Methods) == 0) {
		log.Fatal("Couldn't load API ", cmd)
	}

	method := endpointFs.Args()[1]
	m := findMethod(method, *api)
	for k, p := range api.Parameters {
		fs.String(k, p.Default, p.Description)
	}
	for k, p := range m.Parameters {
		fs.String(k, p.Default, p.Description)
	}

	// TODO: Support deep nested schemas
	s := api.Schemas[m.RequestSchema.Ref]
	for pk, p := range s.Properties {
		fs.String("res."+pk, "", p.Description)
	}

	fs.Parse(endpointFs.Args()[2:])
	m.call(api)
}

func findMethod(method string, api API) *Method {
	parts := strings.Split(method, ".")
	var ms map[string]Method
	rs := api.Resources
	for i := 0; i < len(parts)-1; i++ {
		r := rs[parts[i]]
		if &r == nil {
			log.Fatal("Could not find requested method ", method)
		}
		rs = r.Resources
		ms = r.Methods
	}
	lp := parts[len(parts)-1]
	m := ms[lp]
	if &m == nil {
		log.Fatal("Could not find requested method ", method)
	}
	return &m
}

func getPreferredVersion(apiName string) (string, error) {
	var d struct {
		Items []struct {
			Version string
		}
	}
	err := getAndParse(fmt.Sprintf("discovery/v1/apis?preferred=true&name=%s&fields=items/version", apiName), &d)
	if err != nil {
		return "", err
	}
	if d.Items == nil {
		log.Fatal("Could not load API ", apiName)
	}
	return d.Items[0].Version, nil
}

// loadAPI takes a string like "apiname" or "apiname:v4" and loads the API from Discovery
func loadAPI(s string) (*API, error) {
	parts := strings.SplitN(s, ":", 2)
	apiName := parts[0]
	var v string
	if len(parts) == 2 {
		v = parts[1]
	} else {
		// Look up preferred version in Directory
		var err error
		v, err = getPreferredVersion(apiName)
		if err != nil {
			log.Fatal(err)
		}
	}

	var a API
	err := getAndParse(fmt.Sprintf("discovery/v1/apis/%s/%s/rest", apiName, v), &a)
	if err != nil {
		return nil, err
	}
	return &a, nil
}

func getAndParse(path string, v interface{}) error {
	url := *flagEndpoint + path

	r, err := http.Get(url)
	if err != nil {
		return err
	}
	defer r.Body.Close()
	err = json.NewDecoder(r.Body).Decode(v)
	if err != nil {
		return err
	}
	return nil
}

type API struct {
	BaseURL, Name, Title, Description, DocumentationLink string
	Resources                                            map[string]Resource
	Methods                                              map[string]Method
	Parameters                                           map[string]Parameter
	Schemas                                              map[string]Schema
}

type Resource struct {
	Resources map[string]Resource
	Methods   map[string]Method
}

type Method struct {
	ID, Path, HttpMethod, Description string
	Parameters                        map[string]Parameter
	Scopes                            []string
	RequestSchema                     struct {
		Ref string `json:"$ref"`
	} `json:"request"`
}

func (m Method) call(api *API) {
	url := api.BaseURL + m.Path

	for k, p := range m.Parameters {
		api.Parameters[k] = p
	}
	for k, p := range api.Parameters {
		f := fs.Lookup(k)
		if f == nil || f.Value.String() == "" {
			continue
		}
		v := f.Value.String()
		if p.Location == "path" {
			if p.Required && v == "" {
				log.Fatal("Missing required parameter", k)
			}
			t := fmt.Sprintf("{%s}", k)
			strings.Replace(url, t, v, -1)
		} else if p.Location == "query" {
			delim := "&"
			if !strings.Contains(url, "?") {
				delim = "?"
			}
			url += fmt.Sprintf("%s%s=%s", delim, k, v)
		}
	}

	r, err := http.NewRequest(m.HttpMethod, url, nil)
	if err != nil {
		log.Fatal("error creating request:", err)
	}

	// Add request body
	if *flagInFile != "" {
		// If user passes --meta.inFile flag, open that file and use its content as request body
		f, err := os.Open(*flagInFile)
		if err != nil {
			log.Fatal("error opening file:", err)
		}
		fi, err := f.Stat()
		if err != nil {
			log.Fatal("error stating file:", err)
		}
		r.ContentLength = fi.Size()
		r.Header.Set("Content-Type", "application/json")
		r.Body = f
	} else if *flagStdin {
		// If user passes --meta.in flag, buffer stdin it and pass it along as the request body
		b, err := ioutil.ReadAll(os.Stdin)
		if err != nil {
			log.Fatal("error reading from stdin:", err)
		}
		r.ContentLength = int64(len(b))
		r.Header.Set("Content-Type", "application/json")
		r.Body = ioutil.NopCloser(bytes.NewReader(b))
	} else {
		// If user passed --res.* flags, create some JSON and use that as the request body
		s := api.Schemas[m.RequestSchema.Ref]
		request := make(map[string]interface{})
		for k, _ := range s.Properties {
			f := fs.Lookup("res." + k)
			if f == nil || f.Value.String() == "" {
				continue
			}
			v := f.Value.String()
			request[k] = v
		}
		if len(request) != 0 {
			body, err := json.Marshal(&request)
			if err != nil {
				log.Fatal("error marshalling JSON", err)
			}
			r.ContentLength = int64(len(body))
			r.Header.Set("Content-Type", "application/json")
			r.Body = ioutil.NopCloser(bytes.NewReader(body))
		}
	}

	// Add auth header
	if m.Scopes != nil {
		if *flagPem != "" && *flagSecrets != "" {
			scope := strings.Join(m.Scopes, " ")
			tok := accessTokenFromPemFile(scope, *flagPem, *flagSecrets)
			r.Header.Set("Authorization", "Bearer "+tok)
		} else {
			log.Fatal("This method requires access to API scopes: ", m.Scopes)
		}
	}

	client := &http.Client{}
	resp, err := client.Do(r)
	if err != nil {
		log.Fatal(err)
	}
	defer resp.Body.Close()

	io.Copy(os.Stderr, resp.Body)
	if resp.StatusCode < 200 || resp.StatusCode > 299 {
		os.Exit(1)
	}
}

func accessTokenFromPemFile(scope, pemPath, secretsPath string) string {
	secretBytes, err := ioutil.ReadFile(secretsPath)
	if err != nil {
		log.Fatal("error reading secrets file:", err)
	}
	var config struct {
		Web struct {
			ClientEmail string `json:"client_email"`
			TokenURI    string `json:"token_uri"`
		}
	}
	err = json.Unmarshal(secretBytes, &config)
	if err != nil {
		log.Fatal("error unmarshalling secrets:", err)
	}

	keyBytes, err := ioutil.ReadFile(pemPath)
	if err != nil {
		log.Fatal("error reading private key file:", err)
	}

	// Craft the ClaimSet and JWT token.
	t := jwt.NewToken(config.Web.ClientEmail, scope, keyBytes)
	t.ClaimSet.Aud = config.Web.TokenURI

	// We need to provide a client.
	c := &http.Client{}

	// Get the access token.
	o, err := t.Assert(c)
	if err != nil {
		log.Fatal("assertion error:", err)
	}

	return o.AccessToken
}

type Parameter struct {
	Type, Description, Location, Default string
	Required                             bool
}

type Schema struct {
	Type       string
	Properties map[string]Property
}

type Property struct {
	Ref               string `json:"$ref"`
	Type, Description string
	Items             struct {
		Ref string `json:"$ref"`
	}
}
<commit_msg>Clean up err handling, add TODOs for schemas<commit_after>// TODO: Handle user auth.
// TODO: Cache discovery/directory documents for faster requests.
// TODO: Handle media upload/download.
// TODO: Handle repeated parameters.

package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strings"

	"code.google.com/p/goauth2/oauth/jwt"
)

var (
	// Flags that get parsed before the command, necessary for loading Cloud Endpoints APIs
	// e.g., "googlecl --endpoint=foo help myapi" parses the endpoint flag before loading the API
	endpointFs   = flag.NewFlagSet("endpoint", flag.ExitOnError)
	flagEndpoint = endpointFs.String("endpoint", "https://www.googleapis.com/", "Cloud Endpoints URL, e.g., https://my-app-id.appspot.com/_ah/api/")

	// Flags that get parsed after the command, common to all APIs
	fs          = flag.NewFlagSet("googlecl", flag.ExitOnError)
	flagPem     = fs.String("meta.pem", "", "Location of .pem file")
	flagSecrets = fs.String("meta.secrets", "", "Location of client_secrets.json")
	flagInFile  = fs.String("meta.inFile", "", "File to pass as request body")
	flagStdin   = fs.Bool("meta.in", false, "Whether to use stdin as the request body")
)

func maybeFatal(msg string, err error) {
	if err != nil {
		log.Fatal(msg, err)
	}
}

func simpleHelp() {
	fmt.Println("Makes requests to Google APIs")
	fmt.Println("Usage:")
	fmt.Println("  googlecl <api> <method> --param=foo")
}

func help() {
	args := endpointFs.Args()
	nargs := len(args)
	if nargs == 0 || (nargs == 1 && args[0] == "help") {
		simpleHelp()
		return
	}
	apiName := args[1]
	api := loadAPI(apiName)
	if nargs == 2 {
		// googlecl help <api>
		fmt.Println(api.Title, api.Description)
		fmt.Println("More information:", api.DocumentationLink)
		fmt.Println("Methods:")
		for _, m := range api.Methods {
			fmt.Println(m.ID, m.Description)
		}
		type pair struct {
			k string
			r Resource
		}
		l := []pair{}
		for k, r := range api.Resources {
			l = append(l, pair{k, r})
		}
		for i := 0; i < len(l); i++ {
			r := l[i].r
			for _, m := range r.Methods {
				fmt.Printf("%s - %s\n", m.ID[len(api.Name)+1:], m.Description)
			}
			for k, r := range r.Resources {
				l = append(l, pair{k, r})
			}
		}
	} else {
		// googlecl help <api> <method>
		method := args[2]
		m := findMethod(method, *api)
		fmt.Println(method, m.Description)
		fmt.Println("Parameters:")
		for k, p := range m.Parameters {
			fmt.Printf("  --%s (%s) - %s\n", k, p.Type, p.Description)
		}
		for k, p := range api.Parameters {
			fmt.Printf("  --%s (%s) - %s\n", k, p.Type, p.Description)
		}
		s := api.Schemas[m.RequestSchema.Ref]
		// TODO: Support deep nested schemas, and use actual flags to get these strings to avoid duplication
		for k, p := range s.Properties {
			fmt.Printf("  --res.%s (%s) - %s\n", k, p.Type, p.Description)
		}
	}
}

func list() {
	var directory struct {
		Items []struct {
			Name, Version, Description string
		}
	}
	getAndParse("discovery/v1/apis", &directory)
	fmt.Println("Available methods:")
	for _, i := range directory.Items {
		fmt.Printf("%s %s - %s\n", i.Name, i.Version, i.Description)
	}
}

func main() {
	endpointFs.Parse(os.Args[1:])
	if len(endpointFs.Args()) == 0 {
		simpleHelp()
		return
	}

	cmd := endpointFs.Args()[0]
	if cmd == "help" {
		help()
		return
	} else if cmd == "list" {
		list()
		return
	}

	api := loadAPI(cmd)
	if api == nil || (len(api.Resources) == 0 && len(api.Methods) == 0) {
		log.Fatal("Couldn't load API ", cmd)
	}

	method := endpointFs.Args()[1]
	m := findMethod(method, *api)
	for k, p := range api.Parameters {
		fs.String(k, p.Default, p.Description)
	}
	for k, p := range m.Parameters {
		fs.String(k, p.Default, p.Description)
	}

	// TODO: Support deep nested schemas
	s := api.Schemas[m.RequestSchema.Ref]
	for pk, p := range s.Properties {
		fs.String("res."+pk, "", "Request body: "+p.Description)
	}

	fs.Parse(endpointFs.Args()[2:])
	m.call(api)
}

func findMethod(method string, api API) *Method {
	parts := strings.Split(method, ".")
	var ms map[string]Method
	rs := api.Resources
	for i := 0; i < len(parts)-1; i++ {
		r := rs[parts[i]]
		if &r == nil {
			log.Fatal("Could not find requested method ", method)
		}
		rs = r.Resources
		ms = r.Methods
	}
	lp := parts[len(parts)-1]
	m := ms[lp]
	if &m == nil {
		log.Fatal("Could not find requested method ", method)
	}
	return &m
}

func getPreferredVersion(apiName string) string {
	var d struct {
		Items []struct {
			Version string
		}
	}
	getAndParse(fmt.Sprintf("discovery/v1/apis?preferred=true&name=%s&fields=items/version", apiName), &d)
	if d.Items == nil {
		log.Fatal("Could not load API ", apiName)
	}
	return d.Items[0].Version
}

// loadAPI takes a string like "apiname" or "apiname:v4" and loads the API from Discovery
func loadAPI(s string) *API {
	parts := strings.SplitN(s, ":", 2)
	apiName := parts[0]
	var v string
	if len(parts) == 2 {
		v = parts[1]
	} else {
		// Look up preferred version in Directory
		v = getPreferredVersion(apiName)
	}

	var a API
	getAndParse(fmt.Sprintf("discovery/v1/apis/%s/%s/rest", apiName, v), &a)
	return &a
}

func getAndParse(path string, v interface{}) {
	url := *flagEndpoint + path

	r, err := http.Get(url)
	maybeFatal("error getting "+url, err)
	defer r.Body.Close()
	err = json.NewDecoder(r.Body).Decode(v)
	maybeFatal("error decoding JSON", err)
}

type API struct {
	BaseURL, Name, Title, Description, DocumentationLink string
	Resources                                            map[string]Resource
	Methods                                              map[string]Method
	Parameters                                           map[string]Parameter
	Schemas                                              map[string]Schema
}

type Resource struct {
	Resources map[string]Resource
	Methods   map[string]Method
}

type Method struct {
	ID, Path, HttpMethod, Description string
	Parameters                        map[string]Parameter
	Scopes                            []string
	RequestSchema                     struct {
		Ref string `json:"$ref"`
	} `json:"request"`
}

func (m Method) call(api *API) {
	url := api.BaseURL + m.Path

	for k, p := range m.Parameters {
		api.Parameters[k] = p
	}
	for k, p := range api.Parameters {
		f := fs.Lookup(k)
		if f == nil || f.Value.String() == "" {
			continue
		}
		v := f.Value.String()
		if p.Location == "path" {
			if p.Required && v == "" {
				log.Fatal("Missing required parameter", k)
			}
			t := fmt.Sprintf("{%s}", k)
			strings.Replace(url, t, v, -1)
		} else if p.Location == "query" {
			delim := "&"
			if !strings.Contains(url, "?") {
				delim = "?"
			}
			url += fmt.Sprintf("%s%s=%s", delim, k, v)
		}
	}

	r, err := http.NewRequest(m.HttpMethod, url, nil)
	maybeFatal("error creating request:", err)

	// Add request body
	if *flagInFile != "" {
		// If user passes --meta.inFile flag, open that file and use its content as request body
		f, err := os.Open(*flagInFile)
		maybeFatal("error opening file:", err)
		fi, err := f.Stat()
		maybeFatal("error stating file:", err)
		r.ContentLength = fi.Size()
		r.Header.Set("Content-Type", "application/json")
		r.Body = f
	} else if *flagStdin {
		// If user passes --meta.in flag, buffer stdin it and pass it along as the request body
		b, err := ioutil.ReadAll(os.Stdin)
		maybeFatal("error reading from stdin:", err)
		r.ContentLength = int64(len(b))
		r.Header.Set("Content-Type", "application/json")
		r.Body = ioutil.NopCloser(bytes.NewReader(b))
	} else {
		// If user passed --res.* flags, create some JSON and use that as the request body
		s := api.Schemas[m.RequestSchema.Ref]
		request := make(map[string]interface{})
		for k, _ := range s.Properties {
			f := fs.Lookup("res." + k)
			if f == nil || f.Value.String() == "" {
				continue
			}
			v := f.Value.String()
			// TODO: Need to convert to expected type first
			request[k] = v
		}
		if len(request) != 0 {
			body, err := json.Marshal(&request)
			maybeFatal("error marshalling JSON:", err)
			r.ContentLength = int64(len(body))
			r.Header.Set("Content-Type", "application/json")
			r.Body = ioutil.NopCloser(bytes.NewReader(body))
		}
	}

	// Add auth header
	if m.Scopes != nil {
		if *flagPem != "" && *flagSecrets != "" {
			scope := strings.Join(m.Scopes, " ")
			tok := accessTokenFromPemFile(scope, *flagPem, *flagSecrets)
			r.Header.Set("Authorization", "Bearer "+tok)
		} else {
			log.Fatal("This method requires access to API scopes: ", m.Scopes)
		}
	}

	client := &http.Client{}
	resp, err := client.Do(r)
	maybeFatal("error making request:", err)
	defer resp.Body.Close()

	io.Copy(os.Stderr, resp.Body)
	if resp.StatusCode < 200 || resp.StatusCode > 299 {
		os.Exit(1)
	}
}

func accessTokenFromPemFile(scope, pemPath, secretsPath string) string {
	secretBytes, err := ioutil.ReadFile(secretsPath)
	maybeFatal("error reading secrets file:", err)
	var config struct {
		Web struct {
			ClientEmail string `json:"client_email"`
			TokenURI    string `json:"token_uri"`
		}
	}
	err = json.Unmarshal(secretBytes, &config)
	maybeFatal("error unmarshalling secrets:", err)

	keyBytes, err := ioutil.ReadFile(pemPath)
	maybeFatal("error reading private key file:", err)

	// Craft the ClaimSet and JWT token.
	t := jwt.NewToken(config.Web.ClientEmail, scope, keyBytes)
	t.ClaimSet.Aud = config.Web.TokenURI

	// We need to provide a client.
	c := &http.Client{}

	// Get the access token.
	o, err := t.Assert(c)
	maybeFatal("assertion error:", err)

	return o.AccessToken
}

type Parameter struct {
	Type, Description, Location, Default string
	Required                             bool
}

type Schema struct {
	Type       string
	Properties map[string]Property
}

type Property struct {
	Ref               string `json:"$ref"`
	Type, Description string
	Items             struct {
		Ref string `json:"$ref"`
	}
}
<|endoftext|>


<commit_before>package main

import (
	"time"
	"log"
	"strconv"
)

const (
	rate = int64(5)
)

type GorStat struct {
	statName	string
	latest		int
	mean	 	int
	max			int
}

func NewGorStat(statName string) (s *GorStat) {
	s = new(GorStat)
	s.statName = statName
	s.latest = 0
	s.mean = 0
	s.max = 0

	if Settings.stats {
		go s.reportStats()
	}
	return
}

func (s *GorStat) Write(latest int) {
	if Settings.stats {
		if latest > s.max {
			s.max = latest
		}
		if latest != 0 {
			s.mean = (s.mean + latest) / 2
		}
		s.latest = latest
	}
}

func (s *GorStat) String() string {
	return s.statName + ":" + strconv.Itoa(s.latest) + "," + strconv.Itoa(s.mean) + "," + strconv.Itoa(s.max)
}

func (s *GorStat) reportStats() {
	for {
			log.Println(s)
			time.Sleep(rate * time.Second)
	}
}

<commit_msg>update rate types<commit_after>package main

import (
	"time"
	"log"
	"strconv"
)

const (
	rate = 5
)

type GorStat struct {
	statName	string
	latest		int
	mean	 	int
	max			int
}

func NewGorStat(statName string) (s *GorStat) {
	s = new(GorStat)
	s.statName = statName
	s.latest = 0
	s.mean = 0
	s.max = 0

	if Settings.stats {
		go s.reportStats()
	}
	return
}

func (s *GorStat) Write(latest int) {
	if Settings.stats {
		if latest > s.max {
			s.max = latest
		}
		if latest != 0 {
			s.mean = (s.mean + latest) / 2
		}
		s.latest = latest
	}
}

func (s *GorStat) String() string {
	return s.statName + ":" + strconv.Itoa(s.latest) + "," + strconv.Itoa(s.mean) + "," + strconv.Itoa(s.max)
}

func (s *GorStat) reportStats() {
	for {
			log.Println(s)
			time.Sleep(rate * time.Second)
	}
}

<|endoftext|>


<commit_before>// Copyright 2015 CoreOS, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package etcd

import (
	"fmt"
	"time"

	"github.com/coreos/mantle/Godeps/_workspace/src/github.com/coreos/pkg/capnslog"
	"github.com/coreos/mantle/platform"
	"github.com/coreos/mantle/util"
)

var plog = capnslog.NewPackageLogger("github.com/coreos/mantle", "kola/tests/etcd")

func DiscoveryV2(c platform.TestCluster) error {
	return discovery(c, 2)
}

func DiscoveryV1(c platform.TestCluster) error {
	return discovery(c, 1)
}

func discovery(cluster platform.Cluster, version int) error {
	csize := len(cluster.Machines())

	if plog.LevelAt(capnslog.DEBUG) {
		// get journalctl -f from all machines before starting
		for _, m := range cluster.Machines() {
			if err := m.StartJournal(); err != nil {
				return fmt.Errorf("failed to start journal: %v", err)
			}
		}
	}

	// point etcd on each machine to discovery
	for i, m := range cluster.Machines() {
		// start etcd instance
		var etcdStart string
		if version == 1 {
			etcdStart = "sudo systemctl start etcd.service"
		} else if version == 2 {
			etcdStart = "sudo systemctl start etcd2.service"
		} else {
			return fmt.Errorf("etcd version unspecified")
		}

		_, err := m.SSH(etcdStart)
		if err != nil {
			return fmt.Errorf("SSH cmd to %v failed: %s", m.IP(), err)
		}
		plog.Infof("etcd instance%d started", i)
	}

	if version == 2 {
		err := getClusterHealth(cluster.Machines()[0], csize)
		if err != nil {
			return fmt.Errorf("discovery failed health check: %v", err)
		}
	} else if version == 1 {
		var keyMap map[string]string
		var retryFuncs []func() error

		retryFuncs = append(retryFuncs, func() error {
			var err error
			keyMap, err = setKeys(cluster, 5)
			if err != nil {
				return err
			}
			return nil
		})
		retryFuncs = append(retryFuncs, func() error {
			if err := checkKeys(cluster, keyMap, false); err != nil {
				return err
			}
			return nil
		})
		for _, retry := range retryFuncs {
			if err := util.Retry(5, 5*time.Second, retry); err != nil {
				return fmt.Errorf("discovery failed health check: %v", err)
			}
			//NOTE(pb): etcd1 seems to fail in odd ways when I quorum
			//read, instead just sleep between setting and getting.
			time.Sleep(2 * time.Second)
		}
	} else {
		return fmt.Errorf("etcd version unspecified")
	}

	return nil
}
<commit_msg>kola/tests/etcd: remove dependence on etcdctl for etcd2 discovery test<commit_after>// Copyright 2015 CoreOS, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package etcd

import (
	"fmt"
	"time"

	"github.com/coreos/mantle/Godeps/_workspace/src/github.com/coreos/pkg/capnslog"
	"github.com/coreos/mantle/platform"
	"github.com/coreos/mantle/util"
)

var plog = capnslog.NewPackageLogger("github.com/coreos/mantle", "kola/tests/etcd")

func DiscoveryV2(c platform.TestCluster) error {
	return discovery(c, 2)
}

func DiscoveryV1(c platform.TestCluster) error {
	return discovery(c, 1)
}

func discovery(cluster platform.Cluster, version int) error {
	if plog.LevelAt(capnslog.DEBUG) {
		// get journalctl -f from all machines before starting
		for _, m := range cluster.Machines() {
			if err := m.StartJournal(); err != nil {
				return fmt.Errorf("failed to start journal: %v", err)
			}
		}
	}

	// point etcd on each machine to discovery
	for i, m := range cluster.Machines() {
		// start etcd instance
		var etcdStart string
		if version == 1 {
			etcdStart = "sudo systemctl start etcd.service"
		} else if version == 2 {
			etcdStart = "sudo systemctl start etcd2.service"
		} else {
			return fmt.Errorf("etcd version unspecified")
		}

		_, err := m.SSH(etcdStart)
		if err != nil {
			return fmt.Errorf("SSH cmd to %v failed: %s", m.IP(), err)
		}
		plog.Infof("etcd instance%d started", i)
	}

	var keyMap map[string]string
	var retryFuncs []func() error

	retryFuncs = append(retryFuncs, func() error {
		var err error
		keyMap, err = setKeys(cluster, 5)
		if err != nil {
			return err
		}
		return nil
	})
	retryFuncs = append(retryFuncs, func() error {
		var quorumRead bool
		if version == 2 {
			quorumRead = true
		}
		if err := checkKeys(cluster, keyMap, quorumRead); err != nil {
			return err
		}
		return nil
	})
	for _, retry := range retryFuncs {
		if err := util.Retry(5, 5*time.Second, retry); err != nil {
			return fmt.Errorf("discovery failed health check: %v", err)
		}
		// NOTE(pb): etcd1 seems to fail in an odd way when I try quorum
		// read, instead just sleep between setting and getting.
		time.Sleep(2 * time.Second)
	}

	return nil
}
<|endoftext|>