<commit_before>package processor

import (
	"sync"
	"time"

	"ci.guzzler.io/guzzler/corcel/config"
	"ci.guzzler.io/guzzler/corcel/core"
	"ci.guzzler.io/guzzler/corcel/errormanager"

	"github.com/REAANDREW/telegraph"
)

func merge(source map[string]interface{}, extra map[string]interface{}) map[string]interface{} {
	for k, v := range extra {
		source[k] = v
	}
	return source
}

//ExecutionBranch ...
type ExecutionBranch interface {
	Execute(plan core.Plan) error
}

//PlanExecutor ...
type PlanExecutor struct {
	Config       *config.Configuration
	Bar          ProgressBar
	start        time.Time
	Publisher    telegraph.LinkedPublisher
	Lists        *ListRingIterator
	Plan         core.Plan
	PlanContext  core.ExtractionResult
	JobContexts  map[int]core.ExtractionResult
	StepContexts map[int]map[int]core.ExtractionResult
	mutex        *sync.Mutex
}

//CreatePlanExecutor ...
func CreatePlanExecutor(config *config.Configuration, bar ProgressBar) *PlanExecutor {
	return &PlanExecutor{
		Config:       config,
		Bar:          bar,
		Publisher:    telegraph.NewLinkedPublisher(),
		PlanContext:  core.ExtractionResult{},
		JobContexts:  map[int]core.ExtractionResult{},
		StepContexts: map[int]map[int]core.ExtractionResult{},
		mutex:        &sync.Mutex{},
	}
}

func (instance *PlanExecutor) executeStep(step core.Step, cancellation chan struct{}) core.ExecutionResult {
	start := time.Now()
	instance.mutex.Lock()
	if instance.JobContexts[step.JobID] == nil {
		instance.JobContexts[step.JobID] = map[string]interface{}{}
		instance.StepContexts[step.JobID] = map[int]core.ExtractionResult{}
		instance.StepContexts[step.JobID][step.ID] = map[string]interface{}{}
	}
	instance.mutex.Unlock()

	var executionContext = core.ExecutionContext{}

	var vars map[string]interface{}

	if instance.Plan.Context["vars"] != nil {
		vars = instance.Plan.Context["vars"].(map[string]interface{})
	}

	for pKey, pValue := range vars {
		executionContext[pKey] = pValue
	}

	listValues := instance.Lists.Values()
	for pKey, pValue := range listValues {
		executionContext[pKey] = pValue
	}

	job := instance.Plan.GetJob(step.JobID)
	for jKey, jValue := range job.Context {
		executionContext[jKey] = jValue
		if jKey == "vars" {
			vars := jValue.(map[interface{}]interface{})
			for varKey, varValue := range vars {
				executionContext[varKey.(string)] = varValue
			}
		}
	}

	var executionResult = core.ExecutionResult{}

	if step.Action != nil {
		executionResult = step.Action.Execute(executionContext, cancellation)
	}

	executionResult = merge(executionResult, instance.PlanContext)
	executionResult = merge(executionResult, instance.JobContexts[step.JobID])
	executionResult = merge(executionResult, instance.StepContexts[step.JobID][step.ID])

	executionResult = merge(executionResult, executionContext)

	for _, extractor := range step.Extractors {
		extractorResult := extractor.Extract(executionResult)

		switch extractorResult.Scope() {
		case core.PlanScope:
			instance.PlanContext = merge(instance.PlanContext, extractorResult)
			fallthrough
		case core.JobScope:
			instance.JobContexts[step.JobID] = merge(instance.JobContexts[step.JobID], extractorResult)
			fallthrough
		case core.StepScope:
			instance.StepContexts[step.JobID][step.ID] = merge(instance.StepContexts[step.JobID][step.ID], extractorResult)
		}

		//instance.JobContexts[step.JobID] = merge(instance.JobContexts[step.JobID], extractorResult)
		for k, v := range extractorResult {
			executionResult[k] = v
		}
	}

	duration := time.Since(start) / time.Millisecond
	executionResult["action:duration"] = duration
	assertionResults := []core.AssertionResult{}
	for _, assertion := range step.Assertions {
		assertionResult := assertion.Assert(executionResult)
		assertionResults = append(assertionResults, assertionResult)
	}
	executionResult["assertions"] = assertionResults

	return executionResult
}

func (instance *PlanExecutor) workerExecuteJob(talula core.Job, cancellation chan struct{}) {
	defer func() { //catch or finally
		if err := recover(); err != nil { //catch
			errormanager.Log(err)
		}
	}()
	var stepStream StepStream
	stepStream = CreateStepSequentialStream(talula.Steps)
	if instance.Config.WaitTime > time.Duration(0) {
		stepStream = CreateStepDelayStream(stepStream, instance.Config.WaitTime)
	}

	for stepStream.HasNext() {
		step := stepStream.Next()
		executionResult := instance.executeStep(step, cancellation)

		instance.Publisher.Publish(executionResult)

	}
}

func (instance *PlanExecutor) workerExecuteJobs(jobs []core.Job) {
	var jobStream JobStream
	jobStream = CreateJobSequentialStream(jobs)

	var cancellation = make(chan struct{})

	if instance.Config.Random {
		jobStream = CreateJobRandomStream(jobs)
	}

	if instance.Config.Iterations > 0 {
		revolvingStream := CreateJobRevolvingStream(jobStream)
		iterationStream := CreateJobIterationStream(*revolvingStream, len(jobs), instance.Config.Iterations)
		jobStream = iterationStream
	}

	if instance.Config.Duration > time.Duration(0) {
		jobStream = CreateJobDurationStream(jobStream, instance.Config.Duration)
		ticker := time.NewTicker(time.Millisecond * 10)
		go func() {
			for _ = range ticker.C {
				_ = instance.Bar.Set(jobStream.Progress())
			}
		}()
		time.AfterFunc(instance.Config.Duration, func() {
			ticker.Stop()
			_ = instance.Bar.Set(100)
			close(cancellation)
		})
	}

	for jobStream.HasNext() {
		job := jobStream.Next()
		_ = instance.Bar.Set(jobStream.Progress())
		instance.workerExecuteJob(job, cancellation)
	}
}

func (instance *PlanExecutor) executeJobs(plan core.Plan) {
	var wg sync.WaitGroup
	wg.Add(instance.Config.Workers)
	for i := 0; i < instance.Config.Workers; i++ {
		go func(executionPlan core.Plan) {
			instance.workerExecuteJobs(executionPlan.Jobs)
			wg.Done()
		}(plan)
	}
	wg.Wait()
}

// Execute ...
func (instance *PlanExecutor) Execute(plan core.Plan) error {
	instance.start = time.Now()
	instance.Plan = plan
	if instance.Plan.Context["lists"] != nil {
		var lists = map[string][]map[string]interface{}{}

		listKeys := instance.Plan.Context["lists"].(map[interface{}]interface{})
		for listKey, listValue := range listKeys {
			lists[listKey.(string)] = []map[string]interface{}{}
			listValueItems := listValue.([]interface{})
			for _, listValueItem := range listValueItems {
				srcData := listValueItem.(map[interface{}]interface{})
				stringKeyData := map[string]interface{}{}
				for srcKey, srcValue := range srcData {
					stringKeyData[srcKey.(string)] = srcValue
				}

				lists[listKey.(string)] = append(lists[listKey.(string)], stringKeyData)
			}
		}

		instance.Lists = NewListRingIterator(lists)
	} else {
		instance.Lists = NewListRingIterator(map[string][]map[string]interface{}{})
	}

	if instance.Plan.Context["vars"] != nil {
		stringKeyData := map[string]interface{}{}
		data := instance.Plan.Context["vars"].(map[interface{}]interface{})
		for dataKey, dataValue := range data {
			stringKeyData[dataKey.(string)] = dataValue
		}
		instance.Plan.Context["vars"] = stringKeyData
	}
	instance.executeJobs(plan)

	return nil
}
<commit_msg>Had to extend the locking for now<commit_after>package processor

import (
	"sync"
	"time"

	"ci.guzzler.io/guzzler/corcel/config"
	"ci.guzzler.io/guzzler/corcel/core"
	"ci.guzzler.io/guzzler/corcel/errormanager"

	"github.com/REAANDREW/telegraph"
)

func merge(source map[string]interface{}, extra map[string]interface{}) map[string]interface{} {
	for k, v := range extra {
		source[k] = v
	}
	return source
}

//ExecutionBranch ...
type ExecutionBranch interface {
	Execute(plan core.Plan) error
}

//PlanExecutor ...
type PlanExecutor struct {
	Config       *config.Configuration
	Bar          ProgressBar
	start        time.Time
	Publisher    telegraph.LinkedPublisher
	Lists        *ListRingIterator
	Plan         core.Plan
	PlanContext  core.ExtractionResult
	JobContexts  map[int]core.ExtractionResult
	StepContexts map[int]map[int]core.ExtractionResult
	mutex        *sync.Mutex
}

//CreatePlanExecutor ...
func CreatePlanExecutor(config *config.Configuration, bar ProgressBar) *PlanExecutor {
	return &PlanExecutor{
		Config:       config,
		Bar:          bar,
		Publisher:    telegraph.NewLinkedPublisher(),
		PlanContext:  core.ExtractionResult{},
		JobContexts:  map[int]core.ExtractionResult{},
		StepContexts: map[int]map[int]core.ExtractionResult{},
		mutex:        &sync.Mutex{},
	}
}

func (instance *PlanExecutor) executeStep(step core.Step, cancellation chan struct{}) core.ExecutionResult {
	start := time.Now()
	instance.mutex.Lock()
	if instance.JobContexts[step.JobID] == nil {
		instance.JobContexts[step.JobID] = map[string]interface{}{}
		instance.StepContexts[step.JobID] = map[int]core.ExtractionResult{}
		instance.StepContexts[step.JobID][step.ID] = map[string]interface{}{}
	}
	instance.mutex.Unlock()

	var executionContext = core.ExecutionContext{}

	var vars map[string]interface{}

	if instance.Plan.Context["vars"] != nil {
		vars = instance.Plan.Context["vars"].(map[string]interface{})
	}

	for pKey, pValue := range vars {
		executionContext[pKey] = pValue
	}

	listValues := instance.Lists.Values()
	for pKey, pValue := range listValues {
		executionContext[pKey] = pValue
	}

	job := instance.Plan.GetJob(step.JobID)
	for jKey, jValue := range job.Context {
		executionContext[jKey] = jValue
		if jKey == "vars" {
			vars := jValue.(map[interface{}]interface{})
			for varKey, varValue := range vars {
				executionContext[varKey.(string)] = varValue
			}
		}
	}

	var executionResult = core.ExecutionResult{}

	if step.Action != nil {
		executionResult = step.Action.Execute(executionContext, cancellation)
	}

	executionResult = merge(executionResult, instance.PlanContext)
	executionResult = merge(executionResult, instance.JobContexts[step.JobID])
	executionResult = merge(executionResult, instance.StepContexts[step.JobID][step.ID])

	executionResult = merge(executionResult, executionContext)

	instance.mutex.Lock()
	for _, extractor := range step.Extractors {
		extractorResult := extractor.Extract(executionResult)

		switch extractorResult.Scope() {
		case core.PlanScope:
			instance.PlanContext = merge(instance.PlanContext, extractorResult)
			fallthrough
		case core.JobScope:
			instance.JobContexts[step.JobID] = merge(instance.JobContexts[step.JobID], extractorResult)
			fallthrough
		case core.StepScope:
			instance.StepContexts[step.JobID][step.ID] = merge(instance.StepContexts[step.JobID][step.ID], extractorResult)
		}

		//instance.JobContexts[step.JobID] = merge(instance.JobContexts[step.JobID], extractorResult)
		for k, v := range extractorResult {
			executionResult[k] = v
		}
	}
	instance.mutex.Unlock()

	duration := time.Since(start) / time.Millisecond
	executionResult["action:duration"] = duration
	assertionResults := []core.AssertionResult{}
	for _, assertion := range step.Assertions {
		assertionResult := assertion.Assert(executionResult)
		assertionResults = append(assertionResults, assertionResult)
	}
	executionResult["assertions"] = assertionResults

	return executionResult
}

func (instance *PlanExecutor) workerExecuteJob(talula core.Job, cancellation chan struct{}) {
	defer func() { //catch or finally
		if err := recover(); err != nil { //catch
			errormanager.Log(err)
		}
	}()
	var stepStream StepStream
	stepStream = CreateStepSequentialStream(talula.Steps)
	if instance.Config.WaitTime > time.Duration(0) {
		stepStream = CreateStepDelayStream(stepStream, instance.Config.WaitTime)
	}

	for stepStream.HasNext() {
		step := stepStream.Next()
		executionResult := instance.executeStep(step, cancellation)

		instance.Publisher.Publish(executionResult)

	}
}

func (instance *PlanExecutor) workerExecuteJobs(jobs []core.Job) {
	var jobStream JobStream
	jobStream = CreateJobSequentialStream(jobs)

	var cancellation = make(chan struct{})

	if instance.Config.Random {
		jobStream = CreateJobRandomStream(jobs)
	}

	if instance.Config.Iterations > 0 {
		revolvingStream := CreateJobRevolvingStream(jobStream)
		iterationStream := CreateJobIterationStream(*revolvingStream, len(jobs), instance.Config.Iterations)
		jobStream = iterationStream
	}

	if instance.Config.Duration > time.Duration(0) {
		jobStream = CreateJobDurationStream(jobStream, instance.Config.Duration)
		ticker := time.NewTicker(time.Millisecond * 10)
		go func() {
			for _ = range ticker.C {
				_ = instance.Bar.Set(jobStream.Progress())
			}
		}()
		time.AfterFunc(instance.Config.Duration, func() {
			ticker.Stop()
			_ = instance.Bar.Set(100)
			close(cancellation)
		})
	}

	for jobStream.HasNext() {
		job := jobStream.Next()
		_ = instance.Bar.Set(jobStream.Progress())
		instance.workerExecuteJob(job, cancellation)
	}
}

func (instance *PlanExecutor) executeJobs(plan core.Plan) {
	var wg sync.WaitGroup
	wg.Add(instance.Config.Workers)
	for i := 0; i < instance.Config.Workers; i++ {
		go func(executionPlan core.Plan) {
			instance.workerExecuteJobs(executionPlan.Jobs)
			wg.Done()
		}(plan)
	}
	wg.Wait()
}

// Execute ...
func (instance *PlanExecutor) Execute(plan core.Plan) error {
	instance.start = time.Now()
	instance.Plan = plan
	if instance.Plan.Context["lists"] != nil {
		var lists = map[string][]map[string]interface{}{}

		listKeys := instance.Plan.Context["lists"].(map[interface{}]interface{})
		for listKey, listValue := range listKeys {
			lists[listKey.(string)] = []map[string]interface{}{}
			listValueItems := listValue.([]interface{})
			for _, listValueItem := range listValueItems {
				srcData := listValueItem.(map[interface{}]interface{})
				stringKeyData := map[string]interface{}{}
				for srcKey, srcValue := range srcData {
					stringKeyData[srcKey.(string)] = srcValue
				}

				lists[listKey.(string)] = append(lists[listKey.(string)], stringKeyData)
			}
		}

		instance.Lists = NewListRingIterator(lists)
	} else {
		instance.Lists = NewListRingIterator(map[string][]map[string]interface{}{})
	}

	if instance.Plan.Context["vars"] != nil {
		stringKeyData := map[string]interface{}{}
		data := instance.Plan.Context["vars"].(map[interface{}]interface{})
		for dataKey, dataValue := range data {
			stringKeyData[dataKey.(string)] = dataValue
		}
		instance.Plan.Context["vars"] = stringKeyData
	}
	instance.executeJobs(plan)

	return nil
}
<|endoftext|>


<commit_before>package cloudwatch

import (
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go/aws"

	"github.com/aws/aws-sdk-go/service/cloudwatch"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/internal"
	internalaws "github.com/influxdata/telegraf/internal/config/aws"
	"github.com/influxdata/telegraf/internal/errchan"
	"github.com/influxdata/telegraf/internal/limiter"
	"github.com/influxdata/telegraf/plugins/inputs"
)

type (
	CloudWatch struct {
		Region    string `toml:"region"`
		AccessKey string `toml:"access_key"`
		SecretKey string `toml:"secret_key"`
		RoleARN   string `toml:"role_arn"`
		Profile   string `toml:"profile"`
		Filename  string `toml:"shared_credential_file"`
		Token     string `toml:"token"`

		Period      internal.Duration `toml:"period"`
		Delay       internal.Duration `toml:"delay"`
		Namespace   string            `toml:"namespace"`
		Metrics     []*Metric         `toml:"metrics"`
		CacheTTL    internal.Duration `toml:"cache_ttl"`
		RateLimit   int               `toml:"ratelimit"`
		client      cloudwatchClient
		metricCache *MetricCache
	}

	Metric struct {
		MetricNames []string     `toml:"names"`
		Dimensions  []*Dimension `toml:"dimensions"`
	}

	Dimension struct {
		Name  string `toml:"name"`
		Value string `toml:"value"`
	}

	MetricCache struct {
		TTL     time.Duration
		Fetched time.Time
		Metrics []*cloudwatch.Metric
	}

	cloudwatchClient interface {
		ListMetrics(*cloudwatch.ListMetricsInput) (*cloudwatch.ListMetricsOutput, error)
		GetMetricStatistics(*cloudwatch.GetMetricStatisticsInput) (*cloudwatch.GetMetricStatisticsOutput, error)
	}
)

func (c *CloudWatch) SampleConfig() string {
	return `
  ## Amazon Region
  region = 'us-east-1'

  ## Amazon Credentials
  ## Credentials are loaded in the following order
  ## 1) Assumed credentials via STS if role_arn is specified
  ## 2) explicit credentials from 'access_key' and 'secret_key'
  ## 3) shared profile from 'profile'
  ## 4) environment variables
  ## 5) shared credentials file
  ## 6) EC2 Instance Profile
  #access_key = ""
  #secret_key = ""
  #token = ""
  #role_arn = ""
  #profile = ""
  #shared_credential_file = ""

  ## Requested CloudWatch aggregation Period (required - must be a multiple of 60s)
  period = '1m'

  ## Collection Delay (required - must account for metrics availability via CloudWatch API)
  delay = '1m'

  ## Recomended: use metric 'interval' that is a multiple of 'period' to avoid
  ## gaps or overlap in pulled data
  interval = '1m'

  ## Configure the TTL for the internal cache of metrics.
  ## Defaults to 1 hr if not specified
  #cache_ttl = '10m'

  ## Metric Statistic Namespace (required)
  namespace = 'AWS/ELB'

  ## Metrics to Pull (optional)
  ## Defaults to all Metrics in Namespace if nothing is provided
  ## Refreshes Namespace available metrics every 1h
  #[[inputs.cloudwatch.metrics]]
  #  names = ['Latency', 'RequestCount']
  #
  #  ## Dimension filters for Metric (optional)
  #  [[inputs.cloudwatch.metrics.dimensions]]
  #    name = 'LoadBalancerName'
  #    value = 'p-example'
`
}

func (c *CloudWatch) Description() string {
	return "Pull Metric Statistics from Amazon CloudWatch"
}

func (c *CloudWatch) Gather(acc telegraf.Accumulator) error {
	if c.client == nil {
		c.initializeCloudWatch()
	}

	var metrics []*cloudwatch.Metric

	// check for provided metric filter
	if c.Metrics != nil {
		metrics = []*cloudwatch.Metric{}
		for _, m := range c.Metrics {
			if !hasWilcard(m.Dimensions) {
				dimensions := make([]*cloudwatch.Dimension, len(m.Dimensions))
				for k, d := range m.Dimensions {
					fmt.Printf("Dimension [%s]:[%s]\n", d.Name, d.Value)
					dimensions[k] = &cloudwatch.Dimension{
						Name:  aws.String(d.Name),
						Value: aws.String(d.Value),
					}
				}
				for _, name := range m.MetricNames {
					metrics = append(metrics, &cloudwatch.Metric{
						Namespace:  aws.String(c.Namespace),
						MetricName: aws.String(name),
						Dimensions: dimensions,
					})
				}
			} else {
				allMetrics, err := c.fetchNamespaceMetrics()
				if err != nil {
					return err
				}
				for _, name := range m.MetricNames {
					for _, metric := range allMetrics {
						if isSelected(metric, m.Dimensions) {
							metrics = append(metrics, &cloudwatch.Metric{
								Namespace:  aws.String(c.Namespace),
								MetricName: aws.String(name),
								Dimensions: metric.Dimensions,
							})
						}
					}
				}
			}

		}
	} else {
		var err error
		metrics, err = c.fetchNamespaceMetrics()
		if err != nil {
			return err
		}
	}

	metricCount := len(metrics)
	errChan := errchan.New(metricCount)

	now := time.Now()

	// limit concurrency or we can easily exhaust user connection limit
	// see cloudwatch API request limits:
	// http://docs.aws.amazon.com/AmazonCloudWatch/latest/DeveloperGuide/cloudwatch_limits.html
	lmtr := limiter.NewRateLimiter(c.RateLimit, time.Second)
	defer lmtr.Stop()
	var wg sync.WaitGroup
	wg.Add(len(metrics))
	for _, m := range metrics {
		<-lmtr.C
		go func(inm *cloudwatch.Metric) {
			defer wg.Done()
			c.gatherMetric(acc, inm, now, errChan.C)
		}(m)
	}
	wg.Wait()

	return errChan.Error()
}

func init() {
	inputs.Add("cloudwatch", func() telegraf.Input {
		ttl, _ := time.ParseDuration("1hr")
		return &CloudWatch{
			CacheTTL: internal.Duration{Duration: ttl},
		}
	})
}

/*
 * Initialize CloudWatch client
 */
func (c *CloudWatch) initializeCloudWatch() error {
	credentialConfig := &internalaws.CredentialConfig{
		Region:    c.Region,
		AccessKey: c.AccessKey,
		SecretKey: c.SecretKey,
		RoleARN:   c.RoleARN,
		Profile:   c.Profile,
		Filename:  c.Filename,
		Token:     c.Token,
	}
	configProvider := credentialConfig.Credentials()

	c.client = cloudwatch.New(configProvider)
	return nil
}

/*
 * Fetch available metrics for given CloudWatch Namespace
 */
func (c *CloudWatch) fetchNamespaceMetrics() (metrics []*cloudwatch.Metric, err error) {
	if c.metricCache != nil && c.metricCache.IsValid() {
		metrics = c.metricCache.Metrics
		return
	}

	metrics = []*cloudwatch.Metric{}

	var token *string
	for more := true; more; {
		params := &cloudwatch.ListMetricsInput{
			Namespace:  aws.String(c.Namespace),
			Dimensions: []*cloudwatch.DimensionFilter{},
			NextToken:  token,
			MetricName: nil,
		}

		resp, err := c.client.ListMetrics(params)
		if err != nil {
			return nil, err
		}

		metrics = append(metrics, resp.Metrics...)

		token = resp.NextToken
		more = token != nil
	}

	c.metricCache = &MetricCache{
		Metrics: metrics,
		Fetched: time.Now(),
		TTL:     c.CacheTTL.Duration,
	}

	return
}

/*
 * Gather given Metric and emit any error
 */
func (c *CloudWatch) gatherMetric(
	acc telegraf.Accumulator,
	metric *cloudwatch.Metric,
	now time.Time,
	errChan chan error,
) {
	params := c.getStatisticsInput(metric, now)
	resp, err := c.client.GetMetricStatistics(params)
	if err != nil {
		errChan <- err
		return
	}

	for _, point := range resp.Datapoints {
		tags := map[string]string{
			"region": c.Region,
			"unit":   snakeCase(*point.Unit),
		}

		for _, d := range metric.Dimensions {
			tags[snakeCase(*d.Name)] = *d.Value
		}

		// record field for each statistic
		fields := map[string]interface{}{}

		if point.Average != nil {
			fields[formatField(*metric.MetricName, cloudwatch.StatisticAverage)] = *point.Average
		}
		if point.Maximum != nil {
			fields[formatField(*metric.MetricName, cloudwatch.StatisticMaximum)] = *point.Maximum
		}
		if point.Minimum != nil {
			fields[formatField(*metric.MetricName, cloudwatch.StatisticMinimum)] = *point.Minimum
		}
		if point.SampleCount != nil {
			fields[formatField(*metric.MetricName, cloudwatch.StatisticSampleCount)] = *point.SampleCount
		}
		if point.Sum != nil {
			fields[formatField(*metric.MetricName, cloudwatch.StatisticSum)] = *point.Sum
		}

		acc.AddFields(formatMeasurement(c.Namespace), fields, tags, *point.Timestamp)
	}

	errChan <- nil
}

/*
 * Formatting helpers
 */
func formatField(metricName string, statistic string) string {
	return fmt.Sprintf("%s_%s", snakeCase(metricName), snakeCase(statistic))
}

func formatMeasurement(namespace string) string {
	namespace = strings.Replace(namespace, "/", "_", -1)
	namespace = snakeCase(namespace)
	return fmt.Sprintf("cloudwatch_%s", namespace)
}

func snakeCase(s string) string {
	s = internal.SnakeCase(s)
	s = strings.Replace(s, "__", "_", -1)
	return s
}

/*
 * Map Metric to *cloudwatch.GetMetricStatisticsInput for given timeframe
 */
func (c *CloudWatch) getStatisticsInput(metric *cloudwatch.Metric, now time.Time) *cloudwatch.GetMetricStatisticsInput {
	end := now.Add(-c.Delay.Duration)

	input := &cloudwatch.GetMetricStatisticsInput{
		StartTime:  aws.Time(end.Add(-c.Period.Duration)),
		EndTime:    aws.Time(end),
		MetricName: metric.MetricName,
		Namespace:  metric.Namespace,
		Period:     aws.Int64(int64(c.Period.Duration.Seconds())),
		Dimensions: metric.Dimensions,
		Statistics: []*string{
			aws.String(cloudwatch.StatisticAverage),
			aws.String(cloudwatch.StatisticMaximum),
			aws.String(cloudwatch.StatisticMinimum),
			aws.String(cloudwatch.StatisticSum),
			aws.String(cloudwatch.StatisticSampleCount)},
	}
	return input
}

/*
 * Check Metric Cache validity
 */
func (c *MetricCache) IsValid() bool {
	return c.Metrics != nil && time.Since(c.Fetched) < c.TTL
}

func hasWilcard(dimensions []*Dimension) bool {
	for _, d := range dimensions {
		if d.Value == "" || d.Value == "*" {
			return true
		}
	}
	return false
}

func isSelected(metric *cloudwatch.Metric, dimensions []*Dimension) bool {
	if len(metric.Dimensions) != len(dimensions) {
		return false
	}
	for _, d := range dimensions {
		selected := false
		for _, d2 := range metric.Dimensions {
			if d.Name == *d2.Name {
				if d.Value == "" || d.Value == "*" || d.Value == *d2.Value {
					selected = true
				}
			}
		}
		if !selected {
			return false
		}
	}
	return true
}
<commit_msg>add the ratelimit to the sample config<commit_after>package cloudwatch

import (
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/aws/aws-sdk-go/aws"

	"github.com/aws/aws-sdk-go/service/cloudwatch"

	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/internal"
	internalaws "github.com/influxdata/telegraf/internal/config/aws"
	"github.com/influxdata/telegraf/internal/errchan"
	"github.com/influxdata/telegraf/internal/limiter"
	"github.com/influxdata/telegraf/plugins/inputs"
)

type (
	CloudWatch struct {
		Region    string `toml:"region"`
		AccessKey string `toml:"access_key"`
		SecretKey string `toml:"secret_key"`
		RoleARN   string `toml:"role_arn"`
		Profile   string `toml:"profile"`
		Filename  string `toml:"shared_credential_file"`
		Token     string `toml:"token"`

		Period      internal.Duration `toml:"period"`
		Delay       internal.Duration `toml:"delay"`
		Namespace   string            `toml:"namespace"`
		Metrics     []*Metric         `toml:"metrics"`
		CacheTTL    internal.Duration `toml:"cache_ttl"`
		RateLimit   int               `toml:"ratelimit"`
		client      cloudwatchClient
		metricCache *MetricCache
	}

	Metric struct {
		MetricNames []string     `toml:"names"`
		Dimensions  []*Dimension `toml:"dimensions"`
	}

	Dimension struct {
		Name  string `toml:"name"`
		Value string `toml:"value"`
	}

	MetricCache struct {
		TTL     time.Duration
		Fetched time.Time
		Metrics []*cloudwatch.Metric
	}

	cloudwatchClient interface {
		ListMetrics(*cloudwatch.ListMetricsInput) (*cloudwatch.ListMetricsOutput, error)
		GetMetricStatistics(*cloudwatch.GetMetricStatisticsInput) (*cloudwatch.GetMetricStatisticsOutput, error)
	}
)

func (c *CloudWatch) SampleConfig() string {
	return `
  ## Amazon Region
  region = 'us-east-1'

  ## Amazon Credentials
  ## Credentials are loaded in the following order
  ## 1) Assumed credentials via STS if role_arn is specified
  ## 2) explicit credentials from 'access_key' and 'secret_key'
  ## 3) shared profile from 'profile'
  ## 4) environment variables
  ## 5) shared credentials file
  ## 6) EC2 Instance Profile
  #access_key = ""
  #secret_key = ""
  #token = ""
  #role_arn = ""
  #profile = ""
  #shared_credential_file = ""

  ## Requested CloudWatch aggregation Period (required - must be a multiple of 60s)
  period = '1m'

  ## Collection Delay (required - must account for metrics availability via CloudWatch API)
  delay = '1m'

  ## Recomended: use metric 'interval' that is a multiple of 'period' to avoid
  ## gaps or overlap in pulled data
  interval = '1m'

  ## Configure the TTL for the internal cache of metrics.
  ## Defaults to 1 hr if not specified
  #cache_ttl = '10m'

  ## Metric Statistic Namespace (required)
  namespace = 'AWS/ELB'

  ## Maximum requests per second. Note that the global default AWS rate limit is
  ## 10 reqs/sec, so if you define multiple namespaces, these should add up to a
  ## maximum of 10.
  ratelimit = 10

  ## Metrics to Pull (optional)
  ## Defaults to all Metrics in Namespace if nothing is provided
  ## Refreshes Namespace available metrics every 1h
  #[[inputs.cloudwatch.metrics]]
  #  names = ['Latency', 'RequestCount']
  #
  #  ## Dimension filters for Metric (optional)
  #  [[inputs.cloudwatch.metrics.dimensions]]
  #    name = 'LoadBalancerName'
  #    value = 'p-example'
`
}

func (c *CloudWatch) Description() string {
	return "Pull Metric Statistics from Amazon CloudWatch"
}

func (c *CloudWatch) Gather(acc telegraf.Accumulator) error {
	if c.client == nil {
		c.initializeCloudWatch()
	}

	var metrics []*cloudwatch.Metric

	// check for provided metric filter
	if c.Metrics != nil {
		metrics = []*cloudwatch.Metric{}
		for _, m := range c.Metrics {
			if !hasWilcard(m.Dimensions) {
				dimensions := make([]*cloudwatch.Dimension, len(m.Dimensions))
				for k, d := range m.Dimensions {
					fmt.Printf("Dimension [%s]:[%s]\n", d.Name, d.Value)
					dimensions[k] = &cloudwatch.Dimension{
						Name:  aws.String(d.Name),
						Value: aws.String(d.Value),
					}
				}
				for _, name := range m.MetricNames {
					metrics = append(metrics, &cloudwatch.Metric{
						Namespace:  aws.String(c.Namespace),
						MetricName: aws.String(name),
						Dimensions: dimensions,
					})
				}
			} else {
				allMetrics, err := c.fetchNamespaceMetrics()
				if err != nil {
					return err
				}
				for _, name := range m.MetricNames {
					for _, metric := range allMetrics {
						if isSelected(metric, m.Dimensions) {
							metrics = append(metrics, &cloudwatch.Metric{
								Namespace:  aws.String(c.Namespace),
								MetricName: aws.String(name),
								Dimensions: metric.Dimensions,
							})
						}
					}
				}
			}

		}
	} else {
		var err error
		metrics, err = c.fetchNamespaceMetrics()
		if err != nil {
			return err
		}
	}

	metricCount := len(metrics)
	errChan := errchan.New(metricCount)

	now := time.Now()

	// limit concurrency or we can easily exhaust user connection limit
	// see cloudwatch API request limits:
	// http://docs.aws.amazon.com/AmazonCloudWatch/latest/DeveloperGuide/cloudwatch_limits.html
	lmtr := limiter.NewRateLimiter(c.RateLimit, time.Second)
	defer lmtr.Stop()
	var wg sync.WaitGroup
	wg.Add(len(metrics))
	for _, m := range metrics {
		<-lmtr.C
		go func(inm *cloudwatch.Metric) {
			defer wg.Done()
			c.gatherMetric(acc, inm, now, errChan.C)
		}(m)
	}
	wg.Wait()

	return errChan.Error()
}

func init() {
	inputs.Add("cloudwatch", func() telegraf.Input {
		ttl, _ := time.ParseDuration("1hr")
		return &CloudWatch{
			CacheTTL: internal.Duration{Duration: ttl},
		}
	})
}

/*
 * Initialize CloudWatch client
 */
func (c *CloudWatch) initializeCloudWatch() error {
	credentialConfig := &internalaws.CredentialConfig{
		Region:    c.Region,
		AccessKey: c.AccessKey,
		SecretKey: c.SecretKey,
		RoleARN:   c.RoleARN,
		Profile:   c.Profile,
		Filename:  c.Filename,
		Token:     c.Token,
	}
	configProvider := credentialConfig.Credentials()

	c.client = cloudwatch.New(configProvider)
	return nil
}

/*
 * Fetch available metrics for given CloudWatch Namespace
 */
func (c *CloudWatch) fetchNamespaceMetrics() (metrics []*cloudwatch.Metric, err error) {
	if c.metricCache != nil && c.metricCache.IsValid() {
		metrics = c.metricCache.Metrics
		return
	}

	metrics = []*cloudwatch.Metric{}

	var token *string
	for more := true; more; {
		params := &cloudwatch.ListMetricsInput{
			Namespace:  aws.String(c.Namespace),
			Dimensions: []*cloudwatch.DimensionFilter{},
			NextToken:  token,
			MetricName: nil,
		}

		resp, err := c.client.ListMetrics(params)
		if err != nil {
			return nil, err
		}

		metrics = append(metrics, resp.Metrics...)

		token = resp.NextToken
		more = token != nil
	}

	c.metricCache = &MetricCache{
		Metrics: metrics,
		Fetched: time.Now(),
		TTL:     c.CacheTTL.Duration,
	}

	return
}

/*
 * Gather given Metric and emit any error
 */
func (c *CloudWatch) gatherMetric(
	acc telegraf.Accumulator,
	metric *cloudwatch.Metric,
	now time.Time,
	errChan chan error,
) {
	params := c.getStatisticsInput(metric, now)
	resp, err := c.client.GetMetricStatistics(params)
	if err != nil {
		errChan <- err
		return
	}

	for _, point := range resp.Datapoints {
		tags := map[string]string{
			"region": c.Region,
			"unit":   snakeCase(*point.Unit),
		}

		for _, d := range metric.Dimensions {
			tags[snakeCase(*d.Name)] = *d.Value
		}

		// record field for each statistic
		fields := map[string]interface{}{}

		if point.Average != nil {
			fields[formatField(*metric.MetricName, cloudwatch.StatisticAverage)] = *point.Average
		}
		if point.Maximum != nil {
			fields[formatField(*metric.MetricName, cloudwatch.StatisticMaximum)] = *point.Maximum
		}
		if point.Minimum != nil {
			fields[formatField(*metric.MetricName, cloudwatch.StatisticMinimum)] = *point.Minimum
		}
		if point.SampleCount != nil {
			fields[formatField(*metric.MetricName, cloudwatch.StatisticSampleCount)] = *point.SampleCount
		}
		if point.Sum != nil {
			fields[formatField(*metric.MetricName, cloudwatch.StatisticSum)] = *point.Sum
		}

		acc.AddFields(formatMeasurement(c.Namespace), fields, tags, *point.Timestamp)
	}

	errChan <- nil
}

/*
 * Formatting helpers
 */
func formatField(metricName string, statistic string) string {
	return fmt.Sprintf("%s_%s", snakeCase(metricName), snakeCase(statistic))
}

func formatMeasurement(namespace string) string {
	namespace = strings.Replace(namespace, "/", "_", -1)
	namespace = snakeCase(namespace)
	return fmt.Sprintf("cloudwatch_%s", namespace)
}

func snakeCase(s string) string {
	s = internal.SnakeCase(s)
	s = strings.Replace(s, "__", "_", -1)
	return s
}

/*
 * Map Metric to *cloudwatch.GetMetricStatisticsInput for given timeframe
 */
func (c *CloudWatch) getStatisticsInput(metric *cloudwatch.Metric, now time.Time) *cloudwatch.GetMetricStatisticsInput {
	end := now.Add(-c.Delay.Duration)

	input := &cloudwatch.GetMetricStatisticsInput{
		StartTime:  aws.Time(end.Add(-c.Period.Duration)),
		EndTime:    aws.Time(end),
		MetricName: metric.MetricName,
		Namespace:  metric.Namespace,
		Period:     aws.Int64(int64(c.Period.Duration.Seconds())),
		Dimensions: metric.Dimensions,
		Statistics: []*string{
			aws.String(cloudwatch.StatisticAverage),
			aws.String(cloudwatch.StatisticMaximum),
			aws.String(cloudwatch.StatisticMinimum),
			aws.String(cloudwatch.StatisticSum),
			aws.String(cloudwatch.StatisticSampleCount)},
	}
	return input
}

/*
 * Check Metric Cache validity
 */
func (c *MetricCache) IsValid() bool {
	return c.Metrics != nil && time.Since(c.Fetched) < c.TTL
}

func hasWilcard(dimensions []*Dimension) bool {
	for _, d := range dimensions {
		if d.Value == "" || d.Value == "*" {
			return true
		}
	}
	return false
}

func isSelected(metric *cloudwatch.Metric, dimensions []*Dimension) bool {
	if len(metric.Dimensions) != len(dimensions) {
		return false
	}
	for _, d := range dimensions {
		selected := false
		for _, d2 := range metric.Dimensions {
			if d.Name == *d2.Name {
				if d.Value == "" || d.Value == "*" || d.Value == *d2.Value {
					selected = true
				}
			}
		}
		if !selected {
			return false
		}
	}
	return true
}
<|endoftext|>


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

package bs

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"strings"
	"sync"

	"github.com/fsouza/go-dockerclient"
	"github.com/tsuru/config"
	"github.com/tsuru/docker-cluster/cluster"
	"github.com/tsuru/tsuru/app"
	"github.com/tsuru/tsuru/log"
	"github.com/tsuru/tsuru/net"
	"github.com/tsuru/tsuru/provision"
	"github.com/tsuru/tsuru/provision/docker/container"
	"github.com/tsuru/tsuru/provision/docker/fix"
)

type DockerProvisioner interface {
	Cluster() *cluster.Cluster
	RegistryAuthConfig() docker.AuthConfiguration
}

const (
	bsUniqueID         = "bs"
	bsDefaultImageName = "tsuru/bs:v1"
)

func EnvListForEndpoint(dockerEndpoint, poolName string) ([]string, error) {
	bsConf, err := provision.FindScopedConfig(bsUniqueID)
	if err != nil {
		return nil, err
	}
	tsuruEndpoint, _ := config.GetString("host")
	if !strings.HasPrefix(tsuruEndpoint, "http://") && !strings.HasPrefix(tsuruEndpoint, "https://") {
		tsuruEndpoint = "http://" + tsuruEndpoint
	}
	tsuruEndpoint = strings.TrimRight(tsuruEndpoint, "/") + "/"
	endpoint := dockerEndpoint
	socket, _ := config.GetString("docker:bs:socket")
	if socket != "" {
		endpoint = "unix:///var/run/docker.sock"
	}
	token, err := getToken(bsConf)
	if err != nil {
		return nil, err
	}
	baseEnvMap := map[string]string{
		"DOCKER_ENDPOINT":       endpoint,
		"TSURU_ENDPOINT":        tsuruEndpoint,
		"TSURU_TOKEN":           token,
		"SYSLOG_LISTEN_ADDRESS": fmt.Sprintf("udp://0.0.0.0:%d", container.BsSysLogPort()),
	}
	var envList []string
	for envName, envValue := range bsConf.PoolEntries(poolName) {
		if _, isBase := baseEnvMap[envName]; isBase {
			continue
		}
		envList = append(envList, fmt.Sprintf("%s=%s", envName, envValue.Value))
	}
	for name, value := range baseEnvMap {
		envList = append(envList, fmt.Sprintf("%s=%s", name, value))
	}
	return envList, nil
}

func getToken(bsConf *provision.ScopedConfig) (string, error) {
	token := bsConf.GetExtraString("token")
	if token != "" {
		return token, nil
	}
	tokenData, err := app.AuthScheme.AppLogin(app.InternalAppName)
	if err != nil {
		return "", err
	}
	token = tokenData.GetValue()
	isSet, err := bsConf.SetExtraAtomic("token", token)
	if isSet {
		return token, nil
	}
	app.AuthScheme.Logout(token)
	if err != nil {
		return "", err
	}
	token = bsConf.GetExtraString("token")
	if token == "" {
		return "", fmt.Errorf("invalid empty bs api token")
	}
	return token, nil
}

func SaveImage(digest string) error {
	bsConf, err := provision.FindScopedConfig(bsUniqueID)
	if err != nil {
		return err
	}
	return bsConf.SetExtra("image", digest)
}

func LoadConfig(pools []string) (*provision.ScopedConfig, error) {
	bsConf, err := provision.FindScopedConfig(bsUniqueID)
	if err != nil {
		return nil, err
	}
	bsConf.FilterPools(pools)
	return bsConf, nil
}

func dockerClient(endpoint string) (*docker.Client, error) {
	client, err := docker.NewClient(endpoint)
	if err != nil {
		return nil, err
	}
	client.HTTPClient = net.Dial5Full300Client
	client.Dialer = net.Dial5Dialer
	return client, nil
}

func getImage(bsConf *provision.ScopedConfig) string {
	image := bsConf.GetExtraString("image")
	if image != "" {
		return image
	}
	image, _ = config.GetString("docker:bs:image")
	if image == "" {
		image = bsDefaultImageName
	}
	return image
}

func createContainer(dockerEndpoint, poolName string, p DockerProvisioner, relaunch bool) error {
	client, err := dockerClient(dockerEndpoint)
	if err != nil {
		return err
	}
	bsConf, err := provision.FindScopedConfig(bsUniqueID)
	if err != nil {
		return err
	}
	bsImage := getImage(bsConf)
	err = pullBsImage(bsImage, dockerEndpoint, p)
	if err != nil {
		return err
	}
	hostConfig := docker.HostConfig{
		RestartPolicy: docker.AlwaysRestart(),
		Privileged:    true,
		NetworkMode:   "host",
	}
	socket, _ := config.GetString("docker:bs:socket")
	if socket != "" {
		hostConfig.Binds = []string{fmt.Sprintf("%s:/var/run/docker.sock:rw", socket)}
	}
	env, err := EnvListForEndpoint(dockerEndpoint, poolName)
	if err != nil {
		return err
	}
	opts := docker.CreateContainerOptions{
		Name:       "big-sibling",
		HostConfig: &hostConfig,
		Config: &docker.Config{
			Image: bsImage,
			Env:   env,
		},
	}
	container, err := client.CreateContainer(opts)
	if relaunch && err == docker.ErrContainerAlreadyExists {
		err = client.RemoveContainer(docker.RemoveContainerOptions{ID: opts.Name, Force: true})
		if err != nil {
			return err
		}
		container, err = client.CreateContainer(opts)
	}
	if err != nil && err != docker.ErrContainerAlreadyExists {
		return err
	}
	if container == nil {
		container, err = client.InspectContainer("big-sibling")
		if err != nil {
			return err
		}
	}
	err = client.StartContainer(container.ID, &hostConfig)
	if _, ok := err.(*docker.ContainerAlreadyRunning); !ok {
		return err
	}
	return nil
}

func pullWithRetry(maxTries int, image, dockerEndpoint string, p DockerProvisioner) (string, error) {
	client, err := dockerClient(dockerEndpoint)
	if err != nil {
		return "", err
	}
	var buf bytes.Buffer
	pullOpts := docker.PullImageOptions{Repository: image, OutputStream: &buf}
	registryAuth := p.RegistryAuthConfig()
	for ; maxTries > 0; maxTries-- {
		err = client.PullImage(pullOpts, registryAuth)
		if err == nil {
			return buf.String(), nil
		}
	}
	return "", err
}

func pullBsImage(image, dockerEndpoint string, p DockerProvisioner) error {
	output, err := pullWithRetry(3, image, dockerEndpoint, p)
	if err != nil {
		return err
	}
	if shouldPinBsImage(image) {
		digest, _ := fix.GetImageDigest(output)
		image += digest
	}
	return SaveImage(image)
}

func shouldPinBsImage(image string) bool {
	parts := strings.SplitN(image, "/", 3)
	lastPart := parts[len(parts)-1]
	return len(strings.SplitN(lastPart, ":", 2)) < 2
}

// RecreateContainers relaunch all bs containers in the cluster for the given
// DockerProvisioner, logging progress to the given writer.
//
// It assumes that the given writer is thread safe.
func RecreateContainers(p DockerProvisioner, w io.Writer) error {
	cluster := p.Cluster()
	nodes, err := cluster.UnfilteredNodes()
	if err != nil {
		return err
	}
	errChan := make(chan error, len(nodes))
	wg := sync.WaitGroup{}
	log.Debugf("[bs containers] recreating %d containers", len(nodes))
	for i := range nodes {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			node := &nodes[i]
			pool := node.Metadata["pool"]
			log.Debugf("[bs containers] recreating container in %s [%s]", node.Address, pool)
			fmt.Fprintf(w, "relaunching bs container in the node %s [%s]\n", node.Address, pool)
			err := createContainer(node.Address, pool, p, true)
			if err != nil {
				msg := fmt.Sprintf("[bs containers] failed to create container in %s [%s]: %s", node.Address, pool, err)
				log.Error(msg)
				err = errors.New(msg)
				errChan <- err
			}
		}(i)
	}
	wg.Wait()
	close(errChan)
	return <-errChan
}

type ClusterHook struct {
	Provisioner DockerProvisioner
}

func (h *ClusterHook) BeforeCreateContainer(node cluster.Node) error {
	err := createContainer(node.Address, node.Metadata["pool"], h.Provisioner, false)
	if err != nil {
		return err
	}
	return nil
}
<commit_msg>provision/docker/bs/bs: fix image+digest output in pullBsImage<commit_after>// Copyright 2016 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package bs

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"strings"
	"sync"

	"github.com/fsouza/go-dockerclient"
	"github.com/tsuru/config"
	"github.com/tsuru/docker-cluster/cluster"
	"github.com/tsuru/tsuru/app"
	"github.com/tsuru/tsuru/log"
	"github.com/tsuru/tsuru/net"
	"github.com/tsuru/tsuru/provision"
	"github.com/tsuru/tsuru/provision/docker/container"
	"github.com/tsuru/tsuru/provision/docker/fix"
)

type DockerProvisioner interface {
	Cluster() *cluster.Cluster
	RegistryAuthConfig() docker.AuthConfiguration
}

const (
	bsUniqueID         = "bs"
	bsDefaultImageName = "tsuru/bs:v1"
)

func EnvListForEndpoint(dockerEndpoint, poolName string) ([]string, error) {
	bsConf, err := provision.FindScopedConfig(bsUniqueID)
	if err != nil {
		return nil, err
	}
	tsuruEndpoint, _ := config.GetString("host")
	if !strings.HasPrefix(tsuruEndpoint, "http://") && !strings.HasPrefix(tsuruEndpoint, "https://") {
		tsuruEndpoint = "http://" + tsuruEndpoint
	}
	tsuruEndpoint = strings.TrimRight(tsuruEndpoint, "/") + "/"
	endpoint := dockerEndpoint
	socket, _ := config.GetString("docker:bs:socket")
	if socket != "" {
		endpoint = "unix:///var/run/docker.sock"
	}
	token, err := getToken(bsConf)
	if err != nil {
		return nil, err
	}
	baseEnvMap := map[string]string{
		"DOCKER_ENDPOINT":       endpoint,
		"TSURU_ENDPOINT":        tsuruEndpoint,
		"TSURU_TOKEN":           token,
		"SYSLOG_LISTEN_ADDRESS": fmt.Sprintf("udp://0.0.0.0:%d", container.BsSysLogPort()),
	}
	var envList []string
	for envName, envValue := range bsConf.PoolEntries(poolName) {
		if _, isBase := baseEnvMap[envName]; isBase {
			continue
		}
		envList = append(envList, fmt.Sprintf("%s=%s", envName, envValue.Value))
	}
	for name, value := range baseEnvMap {
		envList = append(envList, fmt.Sprintf("%s=%s", name, value))
	}
	return envList, nil
}

func getToken(bsConf *provision.ScopedConfig) (string, error) {
	token := bsConf.GetExtraString("token")
	if token != "" {
		return token, nil
	}
	tokenData, err := app.AuthScheme.AppLogin(app.InternalAppName)
	if err != nil {
		return "", err
	}
	token = tokenData.GetValue()
	isSet, err := bsConf.SetExtraAtomic("token", token)
	if isSet {
		return token, nil
	}
	app.AuthScheme.Logout(token)
	if err != nil {
		return "", err
	}
	token = bsConf.GetExtraString("token")
	if token == "" {
		return "", fmt.Errorf("invalid empty bs api token")
	}
	return token, nil
}

func SaveImage(digest string) error {
	bsConf, err := provision.FindScopedConfig(bsUniqueID)
	if err != nil {
		return err
	}
	return bsConf.SetExtra("image", digest)
}

func LoadConfig(pools []string) (*provision.ScopedConfig, error) {
	bsConf, err := provision.FindScopedConfig(bsUniqueID)
	if err != nil {
		return nil, err
	}
	bsConf.FilterPools(pools)
	return bsConf, nil
}

func dockerClient(endpoint string) (*docker.Client, error) {
	client, err := docker.NewClient(endpoint)
	if err != nil {
		return nil, err
	}
	client.HTTPClient = net.Dial5Full300Client
	client.Dialer = net.Dial5Dialer
	return client, nil
}

func getImage(bsConf *provision.ScopedConfig) string {
	image := bsConf.GetExtraString("image")
	if image != "" {
		return image
	}
	image, _ = config.GetString("docker:bs:image")
	if image == "" {
		image = bsDefaultImageName
	}
	return image
}

func createContainer(dockerEndpoint, poolName string, p DockerProvisioner, relaunch bool) error {
	client, err := dockerClient(dockerEndpoint)
	if err != nil {
		return err
	}
	bsConf, err := provision.FindScopedConfig(bsUniqueID)
	if err != nil {
		return err
	}
	bsImage := getImage(bsConf)
	err = pullBsImage(bsImage, dockerEndpoint, p)
	if err != nil {
		return err
	}
	hostConfig := docker.HostConfig{
		RestartPolicy: docker.AlwaysRestart(),
		Privileged:    true,
		NetworkMode:   "host",
	}
	socket, _ := config.GetString("docker:bs:socket")
	if socket != "" {
		hostConfig.Binds = []string{fmt.Sprintf("%s:/var/run/docker.sock:rw", socket)}
	}
	env, err := EnvListForEndpoint(dockerEndpoint, poolName)
	if err != nil {
		return err
	}
	opts := docker.CreateContainerOptions{
		Name:       "big-sibling",
		HostConfig: &hostConfig,
		Config: &docker.Config{
			Image: bsImage,
			Env:   env,
		},
	}
	container, err := client.CreateContainer(opts)
	if relaunch && err == docker.ErrContainerAlreadyExists {
		err = client.RemoveContainer(docker.RemoveContainerOptions{ID: opts.Name, Force: true})
		if err != nil {
			return err
		}
		container, err = client.CreateContainer(opts)
	}
	if err != nil && err != docker.ErrContainerAlreadyExists {
		return err
	}
	if container == nil {
		container, err = client.InspectContainer("big-sibling")
		if err != nil {
			return err
		}
	}
	err = client.StartContainer(container.ID, &hostConfig)
	if _, ok := err.(*docker.ContainerAlreadyRunning); !ok {
		return err
	}
	return nil
}

func pullWithRetry(maxTries int, image, dockerEndpoint string, p DockerProvisioner) (string, error) {
	client, err := dockerClient(dockerEndpoint)
	if err != nil {
		return "", err
	}
	var buf bytes.Buffer
	pullOpts := docker.PullImageOptions{Repository: image, OutputStream: &buf}
	registryAuth := p.RegistryAuthConfig()
	for ; maxTries > 0; maxTries-- {
		err = client.PullImage(pullOpts, registryAuth)
		if err == nil {
			return buf.String(), nil
		}
	}
	return "", err
}

func pullBsImage(image, dockerEndpoint string, p DockerProvisioner) error {
	output, err := pullWithRetry(3, image, dockerEndpoint, p)
	if err != nil {
		return err
	}
	if shouldPinBsImage(image) {
		digest, _ := fix.GetImageDigest(output)
		if digest != "" {
			image = fmt.Sprintf("%s@%s", image, digest)
		}
	}
	return SaveImage(image)
}

func shouldPinBsImage(image string) bool {
	parts := strings.SplitN(image, "/", 3)
	lastPart := parts[len(parts)-1]
	return len(strings.SplitN(lastPart, ":", 2)) < 2
}

// RecreateContainers relaunch all bs containers in the cluster for the given
// DockerProvisioner, logging progress to the given writer.
//
// It assumes that the given writer is thread safe.
func RecreateContainers(p DockerProvisioner, w io.Writer) error {
	cluster := p.Cluster()
	nodes, err := cluster.UnfilteredNodes()
	if err != nil {
		return err
	}
	errChan := make(chan error, len(nodes))
	wg := sync.WaitGroup{}
	log.Debugf("[bs containers] recreating %d containers", len(nodes))
	for i := range nodes {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			node := &nodes[i]
			pool := node.Metadata["pool"]
			log.Debugf("[bs containers] recreating container in %s [%s]", node.Address, pool)
			fmt.Fprintf(w, "relaunching bs container in the node %s [%s]\n", node.Address, pool)
			err := createContainer(node.Address, pool, p, true)
			if err != nil {
				msg := fmt.Sprintf("[bs containers] failed to create container in %s [%s]: %s", node.Address, pool, err)
				log.Error(msg)
				err = errors.New(msg)
				errChan <- err
			}
		}(i)
	}
	wg.Wait()
	close(errChan)
	return <-errChan
}

type ClusterHook struct {
	Provisioner DockerProvisioner
}

func (h *ClusterHook) BeforeCreateContainer(node cluster.Node) error {
	err := createContainer(node.Address, node.Metadata["pool"], h.Provisioner, false)
	if err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package pixur

import (
	"bytes"
	"fmt"
	"image"
	"image/gif"
	"io/ioutil"
	"os"
	"testing"

	ptest "pixur.org/pixur/testing"
)

var (
	pixPath           string
	uploadedImagePath string
	uploadedImageSize int64
)

func init() {
	BeforeTestSuite(func() error {
		if path, err := ioutil.TempDir("", "pixPath"); err != nil {
			return err
		} else {
			pixPath = path
		}
		AfterTestSuite(func() error {
			return os.RemoveAll(pixPath)
		})

		return nil
	})

	BeforeTestSuite(func() error {
		f, err := ioutil.TempFile(pixPath, "")
		if err != nil {
			return err
		}
		uploadedImagePath = f.Name()
		defer f.Close()
		AfterTestSuite(func() error {
			return os.RemoveAll(uploadedImagePath)
		})

		img := image.NewGray(image.Rect(0, 0, 5, 10))

		if err := gif.Encode(f, img, &gif.Options{}); err != nil {
			return err
		}
		if fi, err := f.Stat(); err != nil {
			return err
		} else {
			uploadedImageSize = fi.Size()
		}
		return nil
	})
}

func TestWorkflowFileUpload(t *testing.T) {
	if err := func() error {
		imgData, err := os.Open(uploadedImagePath)
		if err != nil {
			return err
		}
		task := &CreatePicTask{
			db:       testDB,
			pixPath:  pixPath,
			FileData: imgData,
		}
		if err := task.Run(); err != nil {
			task.Reset()
			return err
		}

		expected := Pic{
			FileSize: uploadedImageSize,
			Mime:     Mime_GIF,
			Width:    5,
			Height:   10,
		}
		actual := *task.CreatedPic

		if _, err := os.Stat(actual.Path(pixPath)); err != nil {
			return fmt.Errorf("Image was not moved: %s", err)
		}
		if _, err := os.Stat(actual.ThumbnailPath(pixPath)); err != nil {
			return fmt.Errorf("Thumbnail not created: %s", err)
		}

		// Zero out these, since they can change from test to test
		actual.Id = 0
		ptest.AssertEquals(actual.CreatedTime, actual.ModifiedTime, t)
		actual.CreatedTime = 0
		actual.ModifiedTime = 0

		ptest.AssertEquals(actual, expected, t)
		return nil
	}(); err != nil {
		t.Fatal(err)
	}
}

func TestWorkflowAllTagsAdded(t *testing.T) {
	if err := func() error {
		imgData, err := os.Open(uploadedImagePath)
		if err != nil {
			return err
		}
		task := &CreatePicTask{
			db:       testDB,
			pixPath:  pixPath,
			FileData: imgData,
			TagNames: []string{"foo", "bar"},
		}
		if err := task.Run(); err != nil {
			task.Reset()
			return err
		}

		fooTag, err := findTagByName("foo", testDB)
		if err != nil {
			return err
		}
		barTag, err := findTagByName("bar", testDB)
		if err != nil {
			return err
		}

		picTags, err := findPicTagsByPicId(task.CreatedPic.Id, testDB)
		if err != nil {
			return err
		}
		if len(picTags) != 2 {
			return fmt.Errorf("Wrong number of pic tags", picTags)
		}
		var picTagsGroupedByName = groupPicTagsByTagName(picTags)
		if picTagsGroupedByName["foo"].TagId != fooTag.Id {
			return fmt.Errorf("Tag ID does not match PicTag TagId", fooTag.Id)
		}
		if picTagsGroupedByName["bar"].TagId != barTag.Id {
			return fmt.Errorf("Tag ID does not match PicTag TagId", barTag.Id)
		}
		return nil
	}(); err != nil {
		t.Fatal(err)
	}
}

func TestWorkflowAlreadyExistingTags(t *testing.T) {
	if err := func() error {
		imgData, err := os.Open(uploadedImagePath)
		if err != nil {
			return err
		}

		bazTag := CreateTagForTest("baz", t)
		quxTag := CreateTagForTest("qux", t)

		task := &CreatePicTask{
			db:       testDB,
			pixPath:  pixPath,
			FileData: imgData,
			TagNames: []string{"baz", "qux"},
		}
		if err := task.Run(); err != nil {
			task.Reset()
			return err
		}

		picTags, err := findPicTagsByPicId(task.CreatedPic.Id, testDB)
		if err != nil {
			return err
		}
		if len(picTags) != 2 {
			return fmt.Errorf("Wrong number of pic tags", picTags)
		}
		var picTagsGroupedByName = groupPicTagsByTagName(picTags)
		if picTagsGroupedByName["baz"].TagId != bazTag.Id {
			return fmt.Errorf("Tag ID does not match PicTag TagId", bazTag.Id)
		}
		if picTagsGroupedByName["qux"].TagId != quxTag.Id {
			return fmt.Errorf("Tag ID does not match PicTag TagId", quxTag.Id)
		}
		return nil
	}(); err != nil {
		t.Fatal(err)
	}
}

func TestWorkflowTrimAndCollapseDuplicateTags(t *testing.T) {
	if err := func() error {
		imgData, err := os.Open(uploadedImagePath)
		if err != nil {
			return err
		}

		task := &CreatePicTask{
			db:       testDB,
			pixPath:  pixPath,
			FileData: imgData,
			TagNames: []string{"foo", "foo", "  foo", "foo  "},
		}
		if err := task.Run(); err != nil {
			task.Reset()
			return err
		}

		fooTag, err := findTagByName("foo", testDB)
		if err != nil {
			return err
		}

		picTags, err := findPicTagsByPicId(task.CreatedPic.Id, testDB)
		if err != nil {
			return err
		}
		if len(picTags) != 1 {
			return fmt.Errorf("Wrong number of pic tags", picTags)
		}
		if picTags[0].TagId != fooTag.Id {
			return fmt.Errorf("Tag ID does not match PicTag TagId", fooTag.Id)
		}
		return nil
	}(); err != nil {
		t.Fatal(err)
	}
}

func TestMoveUploadedFile(t *testing.T) {
	if err := func() error {
		imgData, err := os.Open(uploadedImagePath)
		if err != nil {
			return err
		}

		expected, err := ioutil.ReadFile(uploadedImagePath)
		if err != nil {
			return err
		}
		task := &CreatePicTask{
			FileData: imgData,
		}

		var destBuffer bytes.Buffer
		var p Pic

		if err := task.moveUploadedFile(&destBuffer, &p); err != nil {
			return err
		}
		if res := destBuffer.String(); res != string(expected) {
			t.Fatal("String data not moved: ", res)
		}
		if int(p.FileSize) != len(expected) {
			t.Fatal("Filesize doesn't match", p.FileSize)
		}
		return nil
	}(); err != nil {
		t.Fatal(err)
	}
}

func TestFillImageConfig(t *testing.T) {
	if err := func() error {
		imgData, err := os.Open(uploadedImagePath)
		if err != nil {
			return err
		}

		task := &CreatePicTask{}
		var p Pic
		if _, err := task.fillImageConfig(imgData, &p); err != nil {
			t.Fatal(err)
		}

		if p.Mime != Mime_GIF {
			t.Fatal("Mime type mismatch", p.Mime)
		}
		if p.Width != 5 || p.Height != 10 {
			t.Fatal("Dimension Mismatch", p.Width, p.Height)
		}
		return nil
	}(); err != nil {
		t.Fatal(err)
	}
}

func CreateTagForTest(tagName string, t *testing.T) *Tag {
  tag := &Tag {
	  Name: tagName,
	}
	res, err := testDB.Exec(tag.BuildInsert(), tag.ColumnPointers(tag.GetColumnNames())...)
	if err != nil {
		t.Fatal(err)
	}
	if insertId, err := res.LastInsertId(); err != nil {
		t.Fatal(err)
	} else {
		tag.Id = insertId
	}
	return tag
}<commit_msg>Added a benchmark.  On this machine its about 5ms for creation, mostly spent in image resizing.<commit_after>package pixur

import (
	"bytes"
	"fmt"
	"image"
	"image/gif"
	"io/ioutil"
	"os"
	"testing"

	ptest "pixur.org/pixur/testing"
)

var (
	pixPath           string
	uploadedImagePath string
	uploadedImageSize int64
)

func init() {
	BeforeTestSuite(func() error {
		if path, err := ioutil.TempDir("", "pixPath"); err != nil {
			return err
		} else {
			pixPath = path
		}
		AfterTestSuite(func() error {
			return os.RemoveAll(pixPath)
		})

		return nil
	})

	BeforeTestSuite(func() error {
		f, err := ioutil.TempFile(pixPath, "")
		if err != nil {
			return err
		}
		uploadedImagePath = f.Name()
		defer f.Close()
		AfterTestSuite(func() error {
			return os.RemoveAll(uploadedImagePath)
		})

		img := image.NewGray(image.Rect(0, 0, 5, 10))

		if err := gif.Encode(f, img, &gif.Options{}); err != nil {
			return err
		}
		if fi, err := f.Stat(); err != nil {
			return err
		} else {
			uploadedImageSize = fi.Size()
		}
		return nil
	})
}

func TestWorkflowFileUpload(t *testing.T) {
	if err := func() error {
		imgData, err := os.Open(uploadedImagePath)
		if err != nil {
			return err
		}
		task := &CreatePicTask{
			db:       testDB,
			pixPath:  pixPath,
			FileData: imgData,
		}
		if err := task.Run(); err != nil {
			task.Reset()
			return err
		}

		expected := Pic{
			FileSize: uploadedImageSize,
			Mime:     Mime_GIF,
			Width:    5,
			Height:   10,
		}
		actual := *task.CreatedPic

		if _, err := os.Stat(actual.Path(pixPath)); err != nil {
			return fmt.Errorf("Image was not moved: %s", err)
		}
		if _, err := os.Stat(actual.ThumbnailPath(pixPath)); err != nil {
			return fmt.Errorf("Thumbnail not created: %s", err)
		}

		// Zero out these, since they can change from test to test
		actual.Id = 0
		ptest.AssertEquals(actual.CreatedTime, actual.ModifiedTime, t)
		actual.CreatedTime = 0
		actual.ModifiedTime = 0

		ptest.AssertEquals(actual, expected, t)
		return nil
	}(); err != nil {
		t.Fatal(err)
	}
}

func TestWorkflowAllTagsAdded(t *testing.T) {
	if err := func() error {
		imgData, err := os.Open(uploadedImagePath)
		if err != nil {
			return err
		}
		task := &CreatePicTask{
			db:       testDB,
			pixPath:  pixPath,
			FileData: imgData,
			TagNames: []string{"foo", "bar"},
		}
		if err := task.Run(); err != nil {
			task.Reset()
			return err
		}

		fooTag, err := findTagByName("foo", testDB)
		if err != nil {
			return err
		}
		barTag, err := findTagByName("bar", testDB)
		if err != nil {
			return err
		}

		picTags, err := findPicTagsByPicId(task.CreatedPic.Id, testDB)
		if err != nil {
			return err
		}
		if len(picTags) != 2 {
			return fmt.Errorf("Wrong number of pic tags", picTags)
		}
		var picTagsGroupedByName = groupPicTagsByTagName(picTags)
		if picTagsGroupedByName["foo"].TagId != fooTag.Id {
			return fmt.Errorf("Tag ID does not match PicTag TagId", fooTag.Id)
		}
		if picTagsGroupedByName["bar"].TagId != barTag.Id {
			return fmt.Errorf("Tag ID does not match PicTag TagId", barTag.Id)
		}
		return nil
	}(); err != nil {
		t.Fatal(err)
	}
}

func TestWorkflowAlreadyExistingTags(t *testing.T) {
	if err := func() error {
		imgData, err := os.Open(uploadedImagePath)
		if err != nil {
			return err
		}

		bazTag := CreateTagForTest("baz", t)
		quxTag := CreateTagForTest("qux", t)

		task := &CreatePicTask{
			db:       testDB,
			pixPath:  pixPath,
			FileData: imgData,
			TagNames: []string{"baz", "qux"},
		}
		if err := task.Run(); err != nil {
			task.Reset()
			return err
		}

		picTags, err := findPicTagsByPicId(task.CreatedPic.Id, testDB)
		if err != nil {
			return err
		}
		if len(picTags) != 2 {
			return fmt.Errorf("Wrong number of pic tags", picTags)
		}
		var picTagsGroupedByName = groupPicTagsByTagName(picTags)
		if picTagsGroupedByName["baz"].TagId != bazTag.Id {
			return fmt.Errorf("Tag ID does not match PicTag TagId", bazTag.Id)
		}
		if picTagsGroupedByName["qux"].TagId != quxTag.Id {
			return fmt.Errorf("Tag ID does not match PicTag TagId", quxTag.Id)
		}
		return nil
	}(); err != nil {
		t.Fatal(err)
	}
}

func TestWorkflowTrimAndCollapseDuplicateTags(t *testing.T) {
	if err := func() error {
		imgData, err := os.Open(uploadedImagePath)
		if err != nil {
			return err
		}

		task := &CreatePicTask{
			db:       testDB,
			pixPath:  pixPath,
			FileData: imgData,
			TagNames: []string{"foo", "foo", "  foo", "foo  "},
		}
		if err := task.Run(); err != nil {
			task.Reset()
			return err
		}

		fooTag, err := findTagByName("foo", testDB)
		if err != nil {
			return err
		}

		picTags, err := findPicTagsByPicId(task.CreatedPic.Id, testDB)
		if err != nil {
			return err
		}
		if len(picTags) != 1 {
			return fmt.Errorf("Wrong number of pic tags", picTags)
		}
		if picTags[0].TagId != fooTag.Id {
			return fmt.Errorf("Tag ID does not match PicTag TagId", fooTag.Id)
		}
		return nil
	}(); err != nil {
		t.Fatal(err)
	}
}

func BenchmarkCreation(b *testing.B) {
	for i := 0; i < b.N; i++ {
		if err := func() error {
			imgData, err := os.Open(uploadedImagePath)
			if err != nil {
				return err
			}

			task := &CreatePicTask{
				db:       testDB,
				pixPath:  pixPath,
				FileData: imgData,
				TagNames: []string{"foo", "bar"},
			}
			if err := task.Run(); err != nil {
				task.Reset()
				return err
			}
			return nil
		}(); err != nil {
			b.Fatal(err)
		}
	}
}

func TestMoveUploadedFile(t *testing.T) {
	if err := func() error {
		imgData, err := os.Open(uploadedImagePath)
		if err != nil {
			return err
		}

		expected, err := ioutil.ReadFile(uploadedImagePath)
		if err != nil {
			return err
		}
		task := &CreatePicTask{
			FileData: imgData,
		}

		var destBuffer bytes.Buffer
		var p Pic

		if err := task.moveUploadedFile(&destBuffer, &p); err != nil {
			return err
		}
		if res := destBuffer.String(); res != string(expected) {
			t.Fatal("String data not moved: ", res)
		}
		if int(p.FileSize) != len(expected) {
			t.Fatal("Filesize doesn't match", p.FileSize)
		}
		return nil
	}(); err != nil {
		t.Fatal(err)
	}
}

func TestFillImageConfig(t *testing.T) {
	if err := func() error {
		imgData, err := os.Open(uploadedImagePath)
		if err != nil {
			return err
		}

		task := &CreatePicTask{}
		var p Pic
		if _, err := task.fillImageConfig(imgData, &p); err != nil {
			t.Fatal(err)
		}

		if p.Mime != Mime_GIF {
			t.Fatal("Mime type mismatch", p.Mime)
		}
		if p.Width != 5 || p.Height != 10 {
			t.Fatal("Dimension Mismatch", p.Width, p.Height)
		}
		return nil
	}(); err != nil {
		t.Fatal(err)
	}
}

func CreateTagForTest(tagName string, t *testing.T) *Tag {
  tag := &Tag {
	  Name: tagName,
	}
	res, err := testDB.Exec(tag.BuildInsert(), tag.ColumnPointers(tag.GetColumnNames())...)
	if err != nil {
		t.Fatal(err)
	}
	if insertId, err := res.LastInsertId(); err != nil {
		t.Fatal(err)
	} else {
		tag.Id = insertId
	}
	return tag
}<|endoftext|>


<commit_before>package creational

/*
	Example of builder pattern:

	builder := NewConcreteBuilder()
	director := NewDirector(builder)
	director.Construct()
	product := builder.GetResult()
*/

// Director is the object which orchestrates the building of a product.
type Director struct {
	builder Builder
}

// NewDirector creates a new Director with a specified Builder.
func NewDirector(builder Builder) Director {
	return Director{builder}
}

// Construct builds the product from a series of steps.
func (d *Director) Construct() {
	d.builder.Build()
}

// Build is an interface for building.
type Builder interface {
	Build()
}

// ConcreteBuilder is a builder for building a Product
type ConcreteBuilder struct {
	built bool
}

// NewConcreteBuilder returns a new Builder.
func NewConcreteBuilder() ConcreteBuilder {
	return ConcreteBuilder{false}
}

// Build builds the product.
func (b *ConcreteBuilder) Build() {
	b.built = true
}

// GetResult returns the Product which has been build during the Build step.
func (b *ConcreteBuilder) GetResult() Product {
	return Product{b.built}
}

// Product describes the product to be built.
type Product struct {
	Built bool
}
<commit_msg>Fix Builder interface comment<commit_after>package creational

/*
	Example of builder pattern:

	builder := NewConcreteBuilder()
	director := NewDirector(builder)
	director.Construct()
	product := builder.GetResult()
*/

// Director is the object which orchestrates the building of a product.
type Director struct {
	builder Builder
}

// NewDirector creates a new Director with a specified Builder.
func NewDirector(builder Builder) Director {
	return Director{builder}
}

// Construct builds the product from a series of steps.
func (d *Director) Construct() {
	d.builder.Build()
}

// Builder is an interface for building.
type Builder interface {
	Build()
}

// ConcreteBuilder is a builder for building a Product
type ConcreteBuilder struct {
	built bool
}

// NewConcreteBuilder returns a new Builder.
func NewConcreteBuilder() ConcreteBuilder {
	return ConcreteBuilder{false}
}

// Build builds the product.
func (b *ConcreteBuilder) Build() {
	b.built = true
}

// GetResult returns the Product which has been build during the Build step.
func (b *ConcreteBuilder) GetResult() Product {
	return Product{b.built}
}

// Product describes the product to be built.
type Product struct {
	Built bool
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"testing"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/athena"

	"github.com/hashicorp/terraform/helper/acctest"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

func TestAccAWSAthenaWorkGroup_basic(t *testing.T) {
	resource.ParallelTest(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSAthenaWorkGroupDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAthenaWorkGroupConfig(acctest.RandString(5)),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSAthenaWorkGroupExists("aws_athena_workgroup.foo"),
				),
			},
		},
	})
}

func TestAccAWSAthenaWorkGroup_withDescription(t *testing.T) {
	rName := acctest.RandString(5)
	rDescription := acctest.RandString(20)
	resource.ParallelTest(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSAthenaWorkGroupDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAthenaWorkGroupConfigDescription(rName, rDescription),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSAthenaWorkGroupExists("aws_athena_workgroup.desc"),
					resource.TestCheckResourceAttr(
						"aws_athena_workgroup.desc", "description", rDescription),
				),
			},
		},
	})
}

func TestAccAWSAthenaWorkGroup_withDescriptionUpdate(t *testing.T) {
	rName := acctest.RandString(5)
	rDescription := acctest.RandString(20)
	rDescriptionUpdate := acctest.RandString(20)

	resource.ParallelTest(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSAthenaWorkGroupDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAthenaWorkGroupConfigDescription(rName, rDescription),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSAthenaWorkGroupExists("aws_athena_workgroup.desc"),
					resource.TestCheckResourceAttr(
						"aws_athena_workgroup.desc", "description", rDescription),
				),
			},
			{
				Config: testAccAthenaWorkGroupConfigDescription(rName, rDescriptionUpdate),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSAthenaWorkGroupExists("aws_athena_workgroup.desc"),
					resource.TestCheckResourceAttr(
						"aws_athena_workgroup.desc", "description", rDescriptionUpdate),
				),
			},
		},
	})
}

func TestAccAWSAthenaWorkGroup_withBytesScannedCutoffPerQuery(t *testing.T) {
	rName := acctest.RandString(5)
	rBytesScannedCutoffPerQuery := "10485760"
	resource.ParallelTest(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSAthenaWorkGroupDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAthenaWorkGroupConfigBytesScannedCutoffPerQuery(rName, rBytesScannedCutoffPerQuery),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSAthenaWorkGroupExists("aws_athena_workgroup.bytes"),
					resource.TestCheckResourceAttr(
						"aws_athena_workgroup.bytes", "bytes_scanned_cutoff_per_query", rBytesScannedCutoffPerQuery),
				),
			},
		},
	})
}

func TestAccAWSAthenaWorkGroup_withBytesScannedCutoffPerQueryUpdate(t *testing.T) {
	rName := acctest.RandString(5)
	rBytesScannedCutoffPerQuery := "10485760"
	rBytesScannedCutoffPerQueryUpdate := "12582912"
	resource.ParallelTest(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSAthenaWorkGroupDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAthenaWorkGroupConfigBytesScannedCutoffPerQuery(rName, rBytesScannedCutoffPerQuery),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSAthenaWorkGroupExists("aws_athena_workgroup.bytes"),
					resource.TestCheckResourceAttr(
						"aws_athena_workgroup.bytes", "bytes_scanned_cutoff_per_query", rBytesScannedCutoffPerQuery),
				),
			},
			{
				Config: testAccAthenaWorkGroupConfigBytesScannedCutoffPerQuery(rName, rBytesScannedCutoffPerQueryUpdate),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSAthenaWorkGroupExists("aws_athena_workgroup.bytes"),
					resource.TestCheckResourceAttr(
						"aws_athena_workgroup.bytes", "bytes_scanned_cutoff_per_query", rBytesScannedCutoffPerQueryUpdate),
				),
			},
		},
	})
}

func TestAccAWSAthenaWorkGroup_withEnforceWorkgroupConfiguration(t *testing.T) {
	rName := acctest.RandString(5)
	rEnforce := "true"
	resource.ParallelTest(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSAthenaWorkGroupDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAthenaWorkGroupConfigEnforceWorkgroupConfiguration(rName, rEnforce),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSAthenaWorkGroupExists("aws_athena_workgroup.enforce"),
					resource.TestCheckResourceAttr(
						"aws_athena_workgroup.enforce", "enforce_workgroup_configuration", rEnforce),
				),
			},
		},
	})
}

func TestAccAWSAthenaWorkGroup_withEnforceWorkgroupConfigurationUpdate(t *testing.T) {
	rName := acctest.RandString(5)
	rEnforce := "true"
	rEnforceUpdate := "false"
	resource.ParallelTest(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSAthenaWorkGroupDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAthenaWorkGroupConfigEnforceWorkgroupConfiguration(rName, rEnforce),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSAthenaWorkGroupExists("aws_athena_workgroup.enforce"),
					resource.TestCheckResourceAttr(
						"aws_athena_workgroup.enforce", "enforce_workgroup_configuration", rEnforce),
				),
			},
			{
				Config: testAccAthenaWorkGroupConfigEnforceWorkgroupConfiguration(rName, rEnforceUpdate),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSAthenaWorkGroupExists("aws_athena_workgroup.enforce"),
					resource.TestCheckResourceAttr(
						"aws_athena_workgroup.enforce", "enforce_workgroup_configuration", rEnforceUpdate),
				),
			},
		},
	})
}

func testAccCheckAWSAthenaWorkGroupDestroy(s *terraform.State) error {
	conn := testAccProvider.Meta().(*AWSClient).athenaconn
	for _, rs := range s.RootModule().Resources {
		if rs.Type != "aws_athena_workgroup" {
			continue
		}

		input := &athena.GetWorkGroupInput{
			WorkGroup: aws.String(rs.Primary.ID),
		}

		resp, err := conn.GetWorkGroup(input)
		if err != nil {
			if isAWSErr(err, athena.ErrCodeInvalidRequestException, rs.Primary.ID) {
				return nil
			}
			return err
		}
		if resp.WorkGroup != nil {
			return fmt.Errorf("Athena WorkGroup (%s) found", rs.Primary.ID)
		}
	}
	return nil
}

func testAccCheckAWSAthenaWorkGroupExists(name string) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[name]
		if !ok {
			return fmt.Errorf("Not found: %s", name)
		}

		conn := testAccProvider.Meta().(*AWSClient).athenaconn

		input := &athena.GetWorkGroupInput{
			WorkGroup: aws.String(rs.Primary.ID),
		}

		_, err := conn.GetWorkGroup(input)
		return err
	}
}

func testAccAthenaWorkGroupConfig(rName string) string {
	return fmt.Sprintf(`
resource "aws_athena_workgroup" "foo" {
  name = "tf-athena-workgroup-%s"
}
		`, rName)
}

func testAccAthenaWorkGroupConfigDescription(rName string, rDescription string) string {
	return fmt.Sprintf(`
	resource "aws_athena_workgroup" "desc" {
		name = "tf-athena-workgroup-%s"
		description = "%s"
	}
	`, rName, rDescription)
}

func testAccAthenaWorkGroupConfigBytesScannedCutoffPerQuery(rName string, rBytesScannedCutoffPerQuery string) string {
	return fmt.Sprintf(`
	resource "aws_athena_workgroup" "bytes" {
		name = "tf-athena-workgroup-%s"
		bytes_scanned_cutoff_per_query = %s
	}
	`, rName, rBytesScannedCutoffPerQuery)
}

func testAccAthenaWorkGroupConfigEnforceWorkgroupConfiguration(rName string, rEnforce string) string {
	return fmt.Sprintf(`
	resource "aws_athena_workgroup" "enforce" {
		name = "tf-athena-workgroup-%s"
		enforce_workgroup_configuration = %s
	}
	`, rName, rEnforce)
}
<commit_msg>Added PublishCloudWatchMetricsEnabled TestCase<commit_after>package aws

import (
	"fmt"
	"testing"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/athena"

	"github.com/hashicorp/terraform/helper/acctest"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
)

func TestAccAWSAthenaWorkGroup_basic(t *testing.T) {
	resource.ParallelTest(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSAthenaWorkGroupDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAthenaWorkGroupConfig(acctest.RandString(5)),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSAthenaWorkGroupExists("aws_athena_workgroup.foo"),
				),
			},
		},
	})
}

func TestAccAWSAthenaWorkGroup_withDescription(t *testing.T) {
	rName := acctest.RandString(5)
	rDescription := acctest.RandString(20)
	resource.ParallelTest(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSAthenaWorkGroupDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAthenaWorkGroupConfigDescription(rName, rDescription),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSAthenaWorkGroupExists("aws_athena_workgroup.desc"),
					resource.TestCheckResourceAttr(
						"aws_athena_workgroup.desc", "description", rDescription),
				),
			},
		},
	})
}

func TestAccAWSAthenaWorkGroup_withDescriptionUpdate(t *testing.T) {
	rName := acctest.RandString(5)
	rDescription := acctest.RandString(20)
	rDescriptionUpdate := acctest.RandString(20)

	resource.ParallelTest(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSAthenaWorkGroupDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAthenaWorkGroupConfigDescription(rName, rDescription),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSAthenaWorkGroupExists("aws_athena_workgroup.desc"),
					resource.TestCheckResourceAttr(
						"aws_athena_workgroup.desc", "description", rDescription),
				),
			},
			{
				Config: testAccAthenaWorkGroupConfigDescription(rName, rDescriptionUpdate),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSAthenaWorkGroupExists("aws_athena_workgroup.desc"),
					resource.TestCheckResourceAttr(
						"aws_athena_workgroup.desc", "description", rDescriptionUpdate),
				),
			},
		},
	})
}

func TestAccAWSAthenaWorkGroup_withBytesScannedCutoffPerQuery(t *testing.T) {
	rName := acctest.RandString(5)
	rBytesScannedCutoffPerQuery := "10485760"
	resource.ParallelTest(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSAthenaWorkGroupDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAthenaWorkGroupConfigBytesScannedCutoffPerQuery(rName, rBytesScannedCutoffPerQuery),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSAthenaWorkGroupExists("aws_athena_workgroup.bytes"),
					resource.TestCheckResourceAttr(
						"aws_athena_workgroup.bytes", "bytes_scanned_cutoff_per_query", rBytesScannedCutoffPerQuery),
				),
			},
		},
	})
}

func TestAccAWSAthenaWorkGroup_withBytesScannedCutoffPerQueryUpdate(t *testing.T) {
	rName := acctest.RandString(5)
	rBytesScannedCutoffPerQuery := "10485760"
	rBytesScannedCutoffPerQueryUpdate := "12582912"
	resource.ParallelTest(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSAthenaWorkGroupDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAthenaWorkGroupConfigBytesScannedCutoffPerQuery(rName, rBytesScannedCutoffPerQuery),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSAthenaWorkGroupExists("aws_athena_workgroup.bytes"),
					resource.TestCheckResourceAttr(
						"aws_athena_workgroup.bytes", "bytes_scanned_cutoff_per_query", rBytesScannedCutoffPerQuery),
				),
			},
			{
				Config: testAccAthenaWorkGroupConfigBytesScannedCutoffPerQuery(rName, rBytesScannedCutoffPerQueryUpdate),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSAthenaWorkGroupExists("aws_athena_workgroup.bytes"),
					resource.TestCheckResourceAttr(
						"aws_athena_workgroup.bytes", "bytes_scanned_cutoff_per_query", rBytesScannedCutoffPerQueryUpdate),
				),
			},
		},
	})
}

func TestAccAWSAthenaWorkGroup_withEnforceWorkgroupConfiguration(t *testing.T) {
	rName := acctest.RandString(5)
	rEnforce := "true"
	resource.ParallelTest(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSAthenaWorkGroupDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAthenaWorkGroupConfigEnforceWorkgroupConfiguration(rName, rEnforce),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSAthenaWorkGroupExists("aws_athena_workgroup.enforce"),
					resource.TestCheckResourceAttr(
						"aws_athena_workgroup.enforce", "enforce_workgroup_configuration", rEnforce),
				),
			},
		},
	})
}

func TestAccAWSAthenaWorkGroup_withEnforceWorkgroupConfigurationUpdate(t *testing.T) {
	rName := acctest.RandString(5)
	rEnforce := "true"
	rEnforceUpdate := "false"
	resource.ParallelTest(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSAthenaWorkGroupDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAthenaWorkGroupConfigEnforceWorkgroupConfiguration(rName, rEnforce),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSAthenaWorkGroupExists("aws_athena_workgroup.enforce"),
					resource.TestCheckResourceAttr(
						"aws_athena_workgroup.enforce", "enforce_workgroup_configuration", rEnforce),
				),
			},
			{
				Config: testAccAthenaWorkGroupConfigEnforceWorkgroupConfiguration(rName, rEnforceUpdate),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSAthenaWorkGroupExists("aws_athena_workgroup.enforce"),
					resource.TestCheckResourceAttr(
						"aws_athena_workgroup.enforce", "enforce_workgroup_configuration", rEnforceUpdate),
				),
			},
		},
	})
}

func TestAccAWSAthenaWorkGroup_withPublishCloudWatchMetricsEnabled(t *testing.T) {
	rName := acctest.RandString(5)
	rEnabled := "true"
	resource.ParallelTest(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSAthenaWorkGroupDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAthenaWorkGroupConfigPublishCloudWatchMetricsEnabled(rName, rEnabled),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSAthenaWorkGroupExists("aws_athena_workgroup.enable"),
					resource.TestCheckResourceAttr(
						"aws_athena_workgroup.enable", "publish_cloudwatch_metrics_enabled", rEnabled),
				),
			},
		},
	})
}

func TestAccAWSAthenaWorkGroup_withPublishCloudWatchMetricsEnabledUpdate(t *testing.T) {
	rName := acctest.RandString(5)
	rEnabled := "true"
	rEnabledUpdate := "false"
	resource.ParallelTest(t, resource.TestCase{
		PreCheck:     func() { testAccPreCheck(t) },
		Providers:    testAccProviders,
		CheckDestroy: testAccCheckAWSAthenaWorkGroupDestroy,
		Steps: []resource.TestStep{
			{
				Config: testAccAthenaWorkGroupConfigPublishCloudWatchMetricsEnabled(rName, rEnabled),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSAthenaWorkGroupExists("aws_athena_workgroup.enable"),
					resource.TestCheckResourceAttr(
						"aws_athena_workgroup.enable", "publish_cloudwatch_metrics_enabled", rEnabled),
				),
			},
			{
				Config: testAccAthenaWorkGroupConfigPublishCloudWatchMetricsEnabled(rName, rEnabledUpdate),
				Check: resource.ComposeTestCheckFunc(
					testAccCheckAWSAthenaWorkGroupExists("aws_athena_workgroup.enable"),
					resource.TestCheckResourceAttr(
						"aws_athena_workgroup.enable", "publish_cloudwatch_metrics_enabled", rEnabledUpdate),
				),
			},
		},
	})
}

func testAccCheckAWSAthenaWorkGroupDestroy(s *terraform.State) error {
	conn := testAccProvider.Meta().(*AWSClient).athenaconn
	for _, rs := range s.RootModule().Resources {
		if rs.Type != "aws_athena_workgroup" {
			continue
		}

		input := &athena.GetWorkGroupInput{
			WorkGroup: aws.String(rs.Primary.ID),
		}

		resp, err := conn.GetWorkGroup(input)
		if err != nil {
			if isAWSErr(err, athena.ErrCodeInvalidRequestException, rs.Primary.ID) {
				return nil
			}
			return err
		}
		if resp.WorkGroup != nil {
			return fmt.Errorf("Athena WorkGroup (%s) found", rs.Primary.ID)
		}
	}
	return nil
}

func testAccCheckAWSAthenaWorkGroupExists(name string) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		rs, ok := s.RootModule().Resources[name]
		if !ok {
			return fmt.Errorf("Not found: %s", name)
		}

		conn := testAccProvider.Meta().(*AWSClient).athenaconn

		input := &athena.GetWorkGroupInput{
			WorkGroup: aws.String(rs.Primary.ID),
		}

		_, err := conn.GetWorkGroup(input)
		return err
	}
}

func testAccAthenaWorkGroupConfig(rName string) string {
	return fmt.Sprintf(`
resource "aws_athena_workgroup" "foo" {
  name = "tf-athena-workgroup-%s"
}
		`, rName)
}

func testAccAthenaWorkGroupConfigDescription(rName string, rDescription string) string {
	return fmt.Sprintf(`
	resource "aws_athena_workgroup" "desc" {
		name = "tf-athena-workgroup-%s"
		description = "%s"
	}
	`, rName, rDescription)
}

func testAccAthenaWorkGroupConfigBytesScannedCutoffPerQuery(rName string, rBytesScannedCutoffPerQuery string) string {
	return fmt.Sprintf(`
	resource "aws_athena_workgroup" "bytes" {
		name = "tf-athena-workgroup-%s"
		bytes_scanned_cutoff_per_query = %s
	}
	`, rName, rBytesScannedCutoffPerQuery)
}

func testAccAthenaWorkGroupConfigEnforceWorkgroupConfiguration(rName string, rEnforce string) string {
	return fmt.Sprintf(`
	resource "aws_athena_workgroup" "enforce" {
		name = "tf-athena-workgroup-%s"
		enforce_workgroup_configuration = %s
	}
	`, rName, rEnforce)
}

func testAccAthenaWorkGroupConfigPublishCloudWatchMetricsEnabled(rName string, rEnable string) string {
	return fmt.Sprintf(`
	resource "aws_athena_workgroup" "enable" {
		name = "tf-athena-workgroup-%s"
		publish_cloudwatch_metrics_enabled = %s
	}
	`, rName, rEnable)
}
<|endoftext|>


<commit_before>package inmem

import (
	"testing"

	"github.com/hashicorp/terraform/backend"
	remotestate "github.com/hashicorp/terraform/backend/remote-state"
	"github.com/hashicorp/terraform/state/remote"
)

func TestRemoteClient_impl(t *testing.T) {
	var _ remote.Client = new(RemoteClient)
	var _ remote.ClientLocker = new(RemoteClient)
}

func TestRemoteClient(t *testing.T) {
	b := backend.TestBackendConfig(t, New(), nil)
	remotestate.TestClient(t, b)
}

func TestInmemLocks(t *testing.T) {
	s, err := backend.TestBackendConfig(t, New(), nil).State()
	if err != nil {
		t.Fatal(err)
	}

	remote.TestRemoteLocks(t, s.(*remote.State).Client, s.(*remote.State).Client)
}
<commit_msg>update remote-state/inmem client<commit_after>package inmem

import (
	"testing"

	"github.com/hashicorp/terraform/backend"
	remotestate "github.com/hashicorp/terraform/backend/remote-state"
	"github.com/hashicorp/terraform/state/remote"
)

func TestRemoteClient_impl(t *testing.T) {
	var _ remote.Client = new(RemoteClient)
	var _ remote.ClientLocker = new(RemoteClient)
}

func TestRemoteClient(t *testing.T) {
	b := backend.TestBackendConfig(t, New(), nil)
	remotestate.TestClient(t, b)
}

func TestInmemLocks(t *testing.T) {
	s, err := backend.TestBackendConfig(t, New(), nil).State(backend.DefaultStateName)
	if err != nil {
		t.Fatal(err)
	}

	remote.TestRemoteLocks(t, s.(*remote.State).Client, s.(*remote.State).Client)
}
<|endoftext|>