<commit_before>package mpawsdynamodb

import (
	"flag"
	"log"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/cloudwatch"
	"github.com/aws/aws-sdk-go/service/cloudwatch/cloudwatchiface"
	mp "github.com/mackerelio/go-mackerel-plugin-helper"
)

const (
	namespace              = "AWS/DynamoDB"
	metricsTypeAverage     = "Average"
	metricsTypeSum         = "Sum"
	metricsTypeMaximum     = "Maximum"
	metricsTypeMinimum     = "Minimum"
	metricsTypeSampleCount = "SampleCount"
)

// has 1 CloudWatch MetricName and corresponding N Mackerel Metrics
type metricsGroup struct {
	CloudWatchName string
	Metrics        []metric
}

type metric struct {
	MackerelName string
	Type         string
}

// DynamoDBPlugin mackerel plugin for aws kinesis
type DynamoDBPlugin struct {
	TableName string
	Prefix    string

	AccessKeyID     string
	SecretAccessKey string
	Region          string
	CloudWatch      *cloudwatch.CloudWatch
}

// MetricKeyPrefix interface for PluginWithPrefix
func (p DynamoDBPlugin) MetricKeyPrefix() string {
	return p.Prefix
}

// prepare creates CloudWatch instance
func (p *DynamoDBPlugin) prepare() error {
	sess, err := session.NewSession()
	if err != nil {
		return err
	}

	config := aws.NewConfig()
	if p.AccessKeyID != "" && p.SecretAccessKey != "" {
		config = config.WithCredentials(credentials.NewStaticCredentials(p.AccessKeyID, p.SecretAccessKey, ""))
	}
	if p.Region != "" {
		config = config.WithRegion(p.Region)
	}

	p.CloudWatch = cloudwatch.New(sess, config)

	return nil
}

func transformAndAppendDatapoint(dp *cloudwatch.Datapoint, dataType string, label string, stats map[string]interface{}) map[string]interface{} {
	if dp != nil {
		switch dataType {
		case metricsTypeAverage:
			stats[label] = *dp.Average
		case metricsTypeSum:
			stats[label] = *dp.Sum
		case metricsTypeMaximum:
			stats[label] = *dp.Maximum
		case metricsTypeMinimum:
			stats[label] = *dp.Minimum
		case metricsTypeSampleCount:
			stats[label] = *dp.SampleCount
		}
	}
	return stats
}

// fetch metrics which takes "Operation" dimensions querying both ListMetrics and GetMetricsStatistics
func fetchOperationWildcardMetrics(cw cloudwatchiface.CloudWatchAPI, mg metricsGroup, baseDimensions []*cloudwatch.Dimension) (map[string]interface{}, error) {
	// get available dimensions
	dimensionFilters := make([]*cloudwatch.DimensionFilter, len(baseDimensions))
	for i, dimension := range baseDimensions {
		dimensionFilters[i] = &cloudwatch.DimensionFilter{
			Name:  dimension.Name,
			Value: dimension.Value,
		}
	}
	input := &cloudwatch.ListMetricsInput{
		Dimensions: dimensionFilters,
		Namespace:  aws.String(namespace),
		MetricName: aws.String(mg.CloudWatchName),
	}
	// ListMetrics can retrieve up to 500 metrics, but DynamoDB Operations are apparently less than 500
	res, err := cw.ListMetrics(input)
	if err != nil {
		return nil, err
	}

	stats := make(map[string]interface{})

	// get datapoints with retrieved dimensions
	for _, cwMetric := range res.Metrics {
		dimensions := cwMetric.Dimensions
		// extract operation name
		var operation *string
		for _, d := range dimensions {
			if *d.Name == "Operation" {
				operation = d.Value
				break
			}
		}
		if operation == nil {
			log.Printf("Unexpected dimension, skip: %s", dimensions)
			continue
		}

		dp, err := getLastPointFromCloudWatch(cw, mg, dimensions)
		if err != nil {
			return nil, nil
		}
		if dp != nil {
			for _, met := range mg.Metrics {
				label := strings.Replace(met.MackerelName, "#", *operation, 1)
				stats = transformAndAppendDatapoint(dp, met.Type, label, stats)
			}
		}
	}

	return stats, nil
}

// getLastPoint fetches a CloudWatch metric and parse
func getLastPointFromCloudWatch(cw cloudwatchiface.CloudWatchAPI, metric metricsGroup, dimensions []*cloudwatch.Dimension) (*cloudwatch.Datapoint, error) {
	now := time.Now()
	statsInput := make([]*string, len(metric.Metrics))
	for i, typ := range metric.Metrics {
		statsInput[i] = aws.String(typ.Type)
	}
	input := &cloudwatch.GetMetricStatisticsInput{
		// 8 min, since some metrics are aggregated over 5 min
		StartTime:  aws.Time(now.Add(time.Duration(480) * time.Second * -1)),
		EndTime:    aws.Time(now),
		MetricName: aws.String(metric.CloudWatchName),
		Period:     aws.Int64(60),
		Statistics: statsInput,
		Namespace:  aws.String(namespace),
		Dimensions: dimensions,
	}
	response, err := cw.GetMetricStatistics(input)
	if err != nil {
		return nil, err
	}

	datapoints := response.Datapoints
	if len(datapoints) == 0 {
		return nil, nil
	}

	latest := new(time.Time)
	var latestDp *cloudwatch.Datapoint
	for _, dp := range datapoints {
		if dp.Timestamp.Before(*latest) {
			continue
		}

		latest = dp.Timestamp
		latestDp = dp
	}

	return latestDp, nil
}

var defaultMetricsGroup = []metricsGroup{
	{CloudWatchName: "ConditionalCheckFailedRequests", Metrics: []metric{
		{MackerelName: "ConditionalCheckFailedRequests", Type: metricsTypeSum},
	}},
	{CloudWatchName: "ConsumedReadCapacityUnits", Metrics: []metric{
		{MackerelName: "ConsumedReadCapacityUnitsSum", Type: metricsTypeSum},
		{MackerelName: "ConsumedReadCapacityUnitsAverage", Type: metricsTypeAverage},
	}},
	{CloudWatchName: "ConsumedWriteCapacityUnits", Metrics: []metric{
		{MackerelName: "ConsumedWriteCapacityUnitsSum", Type: metricsTypeSum},
		{MackerelName: "ConsumedWriteCapacityUnitsAverage", Type: metricsTypeAverage},
	}},
	{CloudWatchName: "ProvisionedReadCapacityUnits", Metrics: []metric{
		{MackerelName: "ProvisionedReadCapacityUnits", Type: metricsTypeMinimum},
	}},
	{CloudWatchName: "ProvisionedWriteCapacityUnits", Metrics: []metric{
		{MackerelName: "ProvisionedWriteCapacityUnits", Type: metricsTypeMinimum},
	}},
	{CloudWatchName: "SystemErrors", Metrics: []metric{
		{MackerelName: "SystemErrors", Type: metricsTypeSum},
	}},
	{CloudWatchName: "UserErrors", Metrics: []metric{
		{MackerelName: "UserErrors", Type: metricsTypeSum},
	}},
	{CloudWatchName: "WriteThrottleEvents", Metrics: []metric{
		{MackerelName: "WriteThrottleEvents", Type: metricsTypeSum},
	}},
}

var operationalMetricsGroup = []metricsGroup{
	{CloudWatchName: "SuccessfulRequestLatency", Metrics: []metric{
		{MackerelName: "SuccessfulRequests.#", Type: metricsTypeSampleCount},
		{MackerelName: "SuccessfulRequestLatency.#.Minimum", Type: metricsTypeMinimum},
		{MackerelName: "SuccessfulRequestLatency.#.Maximum", Type: metricsTypeMaximum},
		{MackerelName: "SuccessfulRequestLatency.#.Average", Type: metricsTypeAverage},
	}},
	{CloudWatchName: "ThrottledRequests", Metrics: []metric{
		{MackerelName: "ThrottledRequests.#", Type: metricsTypeSampleCount},
	}},
	{CloudWatchName: "SystemErrors", Metrics: []metric{
		{MackerelName: "SystemErrors.#", Type: metricsTypeSampleCount},
	}},
	{CloudWatchName: "UserErrors", Metrics: []metric{
		{MackerelName: "UserErrors.#", Type: metricsTypeSampleCount},
	}},
	{CloudWatchName: "ReturnedItemCount", Metrics: []metric{
		{MackerelName: "ReturnedItemCount.#", Type: metricsTypeAverage},
	}},
}

// FetchMetrics fetch the metrics
func (p DynamoDBPlugin) FetchMetrics() (map[string]interface{}, error) {
	stats := make(map[string]interface{})

	tableDimensions := []*cloudwatch.Dimension{{
		Name:  aws.String("TableName"),
		Value: aws.String(p.TableName),
	}}
	for _, met := range defaultMetricsGroup {
		dp, err := getLastPointFromCloudWatch(p.CloudWatch, met, tableDimensions)
		if err == nil {
			for _, m := range met.Metrics {
				stats = transformAndAppendDatapoint(dp, m.Type, m.MackerelName, stats)
			}
		} else {
			log.Printf("%s: %s", met, err)
		}
	}

	for _, met := range operationalMetricsGroup {
		operationalStats, err := fetchOperationWildcardMetrics(p.CloudWatch, met, tableDimensions)
		if err == nil {
			for name, s := range operationalStats {
				stats[name] = s
			}
		} else {
			log.Printf("%s: %s", met, err)
		}
	}
	return transformMetrics(stats), nil
}

// TransformMetrics converts some of datapoints to post differences of two metrics
func transformMetrics(stats map[string]interface{}) map[string]interface{} {
	// Although stats are interface{}, those values from cloudwatch.Datapoint are guaranteed to be numerical
	if consumedReadCapacitySum, ok := stats["ConsumedReadCapacityUnitsSum"].(float64); ok {
		stats["ConsumedReadCapacityUnitsNormalized"] = consumedReadCapacitySum / 60.0
	}
	if consumedWriteCapacitySum, ok := stats["ConsumedWriteCapacityUnitsSum"].(float64); ok {
		stats["ConsumedWriteCapacityUnitsNormalized"] = consumedWriteCapacitySum / 60.0
	}
	return stats
}

// GraphDefinition of DynamoDBPlugin
func (p DynamoDBPlugin) GraphDefinition() map[string]mp.Graphs {
	labelPrefix := strings.Title(p.Prefix)
	labelPrefix = strings.Replace(labelPrefix, "-", " ", -1)
	labelPrefix = strings.Replace(labelPrefix, "Dynamodb", "DynamoDB", -1)

	var graphdef = map[string]mp.Graphs{
		"ReadCapacity": {
			Label: (labelPrefix + " Read Capacity Units"),
			Unit:  "float",
			Metrics: []mp.Metrics{
				{Name: "ProvisionedReadCapacityUnits", Label: "Provisioned"},
				{Name: "ConsumedReadCapacityUnitsNormalized", Label: "Consumed"},
				{Name: "ConsumedReadCapacityUnitsAverage", Label: "Consumed (Average per request)"},
			},
		},
		"WriteCapacity": {
			Label: (labelPrefix + " Write Capacity Units"),
			Unit:  "float",
			Metrics: []mp.Metrics{
				{Name: "ProvisionedWriteCapacityUnits", Label: "Provisioned"},
				{Name: "ConsumedWriteCapacityUnitsNormalized", Label: "Consumed"},
				{Name: "ConsumedWriteCapacityUnitsAverage", Label: "Consumed (Average per request)"},
			},
		},
		"ThrottledEvents": {
			Label: (labelPrefix + " Throttle Events"),
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "ReadThrottleEvents", Label: "Read"},
				{Name: "WriteThrottleEvents", Label: "Write"},
			},
		},
		"ConditionalCheckFailedRequests": {
			Label: (labelPrefix + " ConditionalCheckFailedRequests"),
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "ConditionalCheckFailedRequests", Label: "Counts"},
			},
		},
		"ThrottledRequests": {
			Label: (labelPrefix + " ThrottledRequests"),
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "*", Label: "%1", Stacked: true},
			},
		},
		"SystemErrors": {
			Label: (labelPrefix + " SystemErrors"),
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "*", Label: "%1", Stacked: true},
			},
		},
		"UserErrors": {
			Label: (labelPrefix + " UserErrors"),
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "*", Label: "%1", Stacked: true},
			},
		},
		"ReturnedItemCount": {
			Label: (labelPrefix + " ReturnedItemCount"),
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "*", Label: "%1"},
			},
		},
		"SuccessfulRequests": {
			Label: (labelPrefix + " SuccessfulRequests"),
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "*", Label: "%1"},
			},
		},
		"SuccessfulRequestLatency.#": {
			Label: (labelPrefix + " SuccessfulRequestLatency"),
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "Minimum", Label: "Min"},
				{Name: "Maximum", Label: "Max"},
				{Name: "Average", Label: "Average"},
			},
		},
	}
	return graphdef
}

// Do the plugin
func Do() {
	optAccessKeyID := flag.String("access-key-id", "", "AWS Access Key ID")
	optSecretAccessKey := flag.String("secret-access-key", "", "AWS Secret Access Key")
	optRegion := flag.String("region", "", "AWS Region")
	optTableName := flag.String("table-name", "", "DynamoDB Table Name")
	optTempfile := flag.String("tempfile", "", "Temp file name")
	optPrefix := flag.String("metric-key-prefix", "dynamodb", "Metric key prefix")
	flag.Parse()

	var plugin DynamoDBPlugin

	plugin.AccessKeyID = *optAccessKeyID
	plugin.SecretAccessKey = *optSecretAccessKey
	plugin.Region = *optRegion
	plugin.TableName = *optTableName
	plugin.Prefix = *optPrefix

	err := plugin.prepare()
	if err != nil {
		log.Fatalln(err)
	}

	helper := mp.NewMackerelPlugin(plugin)
	helper.Tempfile = *optTempfile

	helper.Run()
}
<commit_msg>use errgroup to retrieve metrics<commit_after>package mpawsdynamodb

import (
	"flag"
	"log"
	"strings"
	"sync"
	"time"

	"golang.org/x/sync/errgroup"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/cloudwatch"
	"github.com/aws/aws-sdk-go/service/cloudwatch/cloudwatchiface"
	mp "github.com/mackerelio/go-mackerel-plugin-helper"
)

const (
	namespace              = "AWS/DynamoDB"
	metricsTypeAverage     = "Average"
	metricsTypeSum         = "Sum"
	metricsTypeMaximum     = "Maximum"
	metricsTypeMinimum     = "Minimum"
	metricsTypeSampleCount = "SampleCount"
)

// has 1 CloudWatch MetricName and corresponding N Mackerel Metrics
type metricsGroup struct {
	CloudWatchName string
	Metrics        []metric
}

type metric struct {
	MackerelName string
	Type         string
}

// DynamoDBPlugin mackerel plugin for aws kinesis
type DynamoDBPlugin struct {
	TableName string
	Prefix    string

	AccessKeyID     string
	SecretAccessKey string
	Region          string
	CloudWatch      *cloudwatch.CloudWatch
}

// MetricKeyPrefix interface for PluginWithPrefix
func (p DynamoDBPlugin) MetricKeyPrefix() string {
	return p.Prefix
}

// prepare creates CloudWatch instance
func (p *DynamoDBPlugin) prepare() error {
	sess, err := session.NewSession()
	if err != nil {
		return err
	}

	config := aws.NewConfig()
	if p.AccessKeyID != "" && p.SecretAccessKey != "" {
		config = config.WithCredentials(credentials.NewStaticCredentials(p.AccessKeyID, p.SecretAccessKey, ""))
	}
	if p.Region != "" {
		config = config.WithRegion(p.Region)
	}

	p.CloudWatch = cloudwatch.New(sess, config)

	return nil
}

func transformAndAppendDatapoint(dp *cloudwatch.Datapoint, dataType string, label string, stats map[string]interface{}) map[string]interface{} {
	if dp != nil {
		switch dataType {
		case metricsTypeAverage:
			stats[label] = *dp.Average
		case metricsTypeSum:
			stats[label] = *dp.Sum
		case metricsTypeMaximum:
			stats[label] = *dp.Maximum
		case metricsTypeMinimum:
			stats[label] = *dp.Minimum
		case metricsTypeSampleCount:
			stats[label] = *dp.SampleCount
		}
	}
	return stats
}

// fetch metrics which takes "Operation" dimensions querying both ListMetrics and GetMetricsStatistics
func fetchOperationWildcardMetrics(cw cloudwatchiface.CloudWatchAPI, mg metricsGroup, baseDimensions []*cloudwatch.Dimension) (map[string]interface{}, error) {
	// get available dimensions
	dimensionFilters := make([]*cloudwatch.DimensionFilter, len(baseDimensions))
	for i, dimension := range baseDimensions {
		dimensionFilters[i] = &cloudwatch.DimensionFilter{
			Name:  dimension.Name,
			Value: dimension.Value,
		}
	}
	input := &cloudwatch.ListMetricsInput{
		Dimensions: dimensionFilters,
		Namespace:  aws.String(namespace),
		MetricName: aws.String(mg.CloudWatchName),
	}
	// ListMetrics can retrieve up to 500 metrics, but DynamoDB Operations are apparently less than 500
	res, err := cw.ListMetrics(input)
	if err != nil {
		return nil, err
	}

	stats := make(map[string]interface{})

	// get datapoints with retrieved dimensions
	for _, cwMetric := range res.Metrics {
		dimensions := cwMetric.Dimensions
		// extract operation name
		var operation *string
		for _, d := range dimensions {
			if *d.Name == "Operation" {
				operation = d.Value
				break
			}
		}
		if operation == nil {
			log.Printf("Unexpected dimension, skip: %s", dimensions)
			continue
		}

		dp, err := getLastPointFromCloudWatch(cw, mg, dimensions)
		if err != nil {
			return nil, nil
		}
		if dp != nil {
			for _, met := range mg.Metrics {
				label := strings.Replace(met.MackerelName, "#", *operation, 1)
				stats = transformAndAppendDatapoint(dp, met.Type, label, stats)
			}
		}
	}

	return stats, nil
}

// getLastPoint fetches a CloudWatch metric and parse
func getLastPointFromCloudWatch(cw cloudwatchiface.CloudWatchAPI, metric metricsGroup, dimensions []*cloudwatch.Dimension) (*cloudwatch.Datapoint, error) {
	now := time.Now()
	statsInput := make([]*string, len(metric.Metrics))
	for i, typ := range metric.Metrics {
		statsInput[i] = aws.String(typ.Type)
	}
	input := &cloudwatch.GetMetricStatisticsInput{
		// 8 min, since some metrics are aggregated over 5 min
		StartTime:  aws.Time(now.Add(time.Duration(480) * time.Second * -1)),
		EndTime:    aws.Time(now),
		MetricName: aws.String(metric.CloudWatchName),
		Period:     aws.Int64(60),
		Statistics: statsInput,
		Namespace:  aws.String(namespace),
		Dimensions: dimensions,
	}
	response, err := cw.GetMetricStatistics(input)
	if err != nil {
		return nil, err
	}

	datapoints := response.Datapoints
	if len(datapoints) == 0 {
		return nil, nil
	}

	latest := new(time.Time)
	var latestDp *cloudwatch.Datapoint
	for _, dp := range datapoints {
		if dp.Timestamp.Before(*latest) {
			continue
		}

		latest = dp.Timestamp
		latestDp = dp
	}

	return latestDp, nil
}

var defaultMetricsGroup = []metricsGroup{
	{CloudWatchName: "ConditionalCheckFailedRequests", Metrics: []metric{
		{MackerelName: "ConditionalCheckFailedRequests", Type: metricsTypeSum},
	}},
	{CloudWatchName: "ConsumedReadCapacityUnits", Metrics: []metric{
		{MackerelName: "ConsumedReadCapacityUnitsSum", Type: metricsTypeSum},
		{MackerelName: "ConsumedReadCapacityUnitsAverage", Type: metricsTypeAverage},
	}},
	{CloudWatchName: "ConsumedWriteCapacityUnits", Metrics: []metric{
		{MackerelName: "ConsumedWriteCapacityUnitsSum", Type: metricsTypeSum},
		{MackerelName: "ConsumedWriteCapacityUnitsAverage", Type: metricsTypeAverage},
	}},
	{CloudWatchName: "ProvisionedReadCapacityUnits", Metrics: []metric{
		{MackerelName: "ProvisionedReadCapacityUnits", Type: metricsTypeMinimum},
	}},
	{CloudWatchName: "ProvisionedWriteCapacityUnits", Metrics: []metric{
		{MackerelName: "ProvisionedWriteCapacityUnits", Type: metricsTypeMinimum},
	}},
	{CloudWatchName: "SystemErrors", Metrics: []metric{
		{MackerelName: "SystemErrors", Type: metricsTypeSum},
	}},
	{CloudWatchName: "UserErrors", Metrics: []metric{
		{MackerelName: "UserErrors", Type: metricsTypeSum},
	}},
	{CloudWatchName: "WriteThrottleEvents", Metrics: []metric{
		{MackerelName: "WriteThrottleEvents", Type: metricsTypeSum},
	}},
}

var operationalMetricsGroup = []metricsGroup{
	{CloudWatchName: "SuccessfulRequestLatency", Metrics: []metric{
		{MackerelName: "SuccessfulRequests.#", Type: metricsTypeSampleCount},
		{MackerelName: "SuccessfulRequestLatency.#.Minimum", Type: metricsTypeMinimum},
		{MackerelName: "SuccessfulRequestLatency.#.Maximum", Type: metricsTypeMaximum},
		{MackerelName: "SuccessfulRequestLatency.#.Average", Type: metricsTypeAverage},
	}},
	{CloudWatchName: "ThrottledRequests", Metrics: []metric{
		{MackerelName: "ThrottledRequests.#", Type: metricsTypeSampleCount},
	}},
	{CloudWatchName: "SystemErrors", Metrics: []metric{
		{MackerelName: "SystemErrors.#", Type: metricsTypeSampleCount},
	}},
	{CloudWatchName: "UserErrors", Metrics: []metric{
		{MackerelName: "UserErrors.#", Type: metricsTypeSampleCount},
	}},
	{CloudWatchName: "ReturnedItemCount", Metrics: []metric{
		{MackerelName: "ReturnedItemCount.#", Type: metricsTypeAverage},
	}},
}

// FetchMetrics fetch the metrics
func (p DynamoDBPlugin) FetchMetrics() (map[string]interface{}, error) {
	stats := make(map[string]interface{})

	tableDimensions := []*cloudwatch.Dimension{{
		Name:  aws.String("TableName"),
		Value: aws.String(p.TableName),
	}}
	var mu sync.Mutex
	var eg errgroup.Group
	for _, met := range defaultMetricsGroup {
		met := met
		eg.Go(func() error {
			dp, err := getLastPointFromCloudWatch(p.CloudWatch, met, tableDimensions)
			if err != nil {
				log.Printf("%s: %s", met, err)
				return err
			}
			for _, m := range met.Metrics {
				mu.Lock()
				stats = transformAndAppendDatapoint(dp, m.Type, m.MackerelName, stats)
				mu.Unlock()
			}
			return nil
		})
	}

	for _, met := range operationalMetricsGroup {
		met := met
		eg.Go(func() error {
			operationalStats, err := fetchOperationWildcardMetrics(p.CloudWatch, met, tableDimensions)
			if err != nil {
				log.Printf("%s: %s", met, err)
				return err
			}

			if operationalStats != nil {
				for name, s := range operationalStats {
					mu.Lock()
					stats[name] = s
					mu.Unlock()
				}
			}
			return nil
		})
	}

	if err := eg.Wait(); err != nil {
		return nil, err
	}
	return transformMetrics(stats), nil
}

// TransformMetrics converts some of datapoints to post differences of two metrics
func transformMetrics(stats map[string]interface{}) map[string]interface{} {
	// Although stats are interface{}, those values from cloudwatch.Datapoint are guaranteed to be numerical
	if consumedReadCapacitySum, ok := stats["ConsumedReadCapacityUnitsSum"].(float64); ok {
		stats["ConsumedReadCapacityUnitsNormalized"] = consumedReadCapacitySum / 60.0
	}
	if consumedWriteCapacitySum, ok := stats["ConsumedWriteCapacityUnitsSum"].(float64); ok {
		stats["ConsumedWriteCapacityUnitsNormalized"] = consumedWriteCapacitySum / 60.0
	}
	return stats
}

// GraphDefinition of DynamoDBPlugin
func (p DynamoDBPlugin) GraphDefinition() map[string]mp.Graphs {
	labelPrefix := strings.Title(p.Prefix)
	labelPrefix = strings.Replace(labelPrefix, "-", " ", -1)
	labelPrefix = strings.Replace(labelPrefix, "Dynamodb", "DynamoDB", -1)

	var graphdef = map[string]mp.Graphs{
		"ReadCapacity": {
			Label: (labelPrefix + " Read Capacity Units"),
			Unit:  "float",
			Metrics: []mp.Metrics{
				{Name: "ProvisionedReadCapacityUnits", Label: "Provisioned"},
				{Name: "ConsumedReadCapacityUnitsNormalized", Label: "Consumed"},
				{Name: "ConsumedReadCapacityUnitsAverage", Label: "Consumed (Average per request)"},
			},
		},
		"WriteCapacity": {
			Label: (labelPrefix + " Write Capacity Units"),
			Unit:  "float",
			Metrics: []mp.Metrics{
				{Name: "ProvisionedWriteCapacityUnits", Label: "Provisioned"},
				{Name: "ConsumedWriteCapacityUnitsNormalized", Label: "Consumed"},
				{Name: "ConsumedWriteCapacityUnitsAverage", Label: "Consumed (Average per request)"},
			},
		},
		"ThrottledEvents": {
			Label: (labelPrefix + " Throttle Events"),
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "ReadThrottleEvents", Label: "Read"},
				{Name: "WriteThrottleEvents", Label: "Write"},
			},
		},
		"ConditionalCheckFailedRequests": {
			Label: (labelPrefix + " ConditionalCheckFailedRequests"),
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "ConditionalCheckFailedRequests", Label: "Counts"},
			},
		},
		"ThrottledRequests": {
			Label: (labelPrefix + " ThrottledRequests"),
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "*", Label: "%1", Stacked: true},
			},
		},
		"SystemErrors": {
			Label: (labelPrefix + " SystemErrors"),
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "*", Label: "%1", Stacked: true},
			},
		},
		"UserErrors": {
			Label: (labelPrefix + " UserErrors"),
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "*", Label: "%1", Stacked: true},
			},
		},
		"ReturnedItemCount": {
			Label: (labelPrefix + " ReturnedItemCount"),
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "*", Label: "%1"},
			},
		},
		"SuccessfulRequests": {
			Label: (labelPrefix + " SuccessfulRequests"),
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "*", Label: "%1"},
			},
		},
		"SuccessfulRequestLatency.#": {
			Label: (labelPrefix + " SuccessfulRequestLatency"),
			Unit:  "integer",
			Metrics: []mp.Metrics{
				{Name: "Minimum", Label: "Min"},
				{Name: "Maximum", Label: "Max"},
				{Name: "Average", Label: "Average"},
			},
		},
	}
	return graphdef
}

// Do the plugin
func Do() {
	optAccessKeyID := flag.String("access-key-id", "", "AWS Access Key ID")
	optSecretAccessKey := flag.String("secret-access-key", "", "AWS Secret Access Key")
	optRegion := flag.String("region", "", "AWS Region")
	optTableName := flag.String("table-name", "", "DynamoDB Table Name")
	optTempfile := flag.String("tempfile", "", "Temp file name")
	optPrefix := flag.String("metric-key-prefix", "dynamodb", "Metric key prefix")
	flag.Parse()

	var plugin DynamoDBPlugin

	plugin.AccessKeyID = *optAccessKeyID
	plugin.SecretAccessKey = *optSecretAccessKey
	plugin.Region = *optRegion
	plugin.TableName = *optTableName
	plugin.Prefix = *optPrefix

	err := plugin.prepare()
	if err != nil {
		log.Fatalln(err)
	}

	helper := mp.NewMackerelPlugin(plugin)
	helper.Tempfile = *optTempfile

	helper.Run()
}
<|endoftext|>


<commit_before>package gamq

import (
	"bufio"
	"fmt"
	log "github.com/cihub/seelog"
	"math/rand"
	"net"
	"strconv"
	"strings"
	"sync"
	"time"
)

const (
	UNRECOGNISEDCOMMANDTEXT = "Unrecognised command"
	HELPSTRING              = "Help. Me."
)

type ConnectionManager struct {
	wg   sync.WaitGroup
	qm   QueueManager
	rand *rand.Rand
}

func (manager *ConnectionManager) Initialize(config *Config) {
	// Initialize our random number generator (used for naming new connections)
	// A different seed will be used on each startup, for no good reason
	manager.rand = rand.New(rand.NewSource(time.Now().UnixNano()))

	manager.qm = QueueManager{}
	manager.qm.Initialize()

	ln, err := net.Listen("tcp", fmt.Sprintf(":%d", config.Port))

	if err != nil {
		log.Errorf("An error occured whilst opening a socket for reading: %s\n",
			err.Error())
	}

	for {
		conn, err := ln.Accept()
		if err != nil {
			log.Errorf("An error occured whilst opening a socket for reading: %s\n",
				err.Error())
		}
		manager.wg.Add(1)
		go manager.handleConnection(&conn)
	}
}

func (manager *ConnectionManager) handleConnection(conn *net.Conn) {
	defer manager.wg.Done()
	connReader := bufio.NewReader(*conn)
	connWriter := bufio.NewWriter(*conn)
	client := Client{Name: strconv.Itoa(manager.rand.Int()),
		Writer: connWriter,
		Reader: connReader}

	for {
		// Read until newline
		line, err := client.Reader.ReadString('\n')

		if err != nil {
			// Connection has been closed
			break
		}

		stringLine := string(line[:len(line)])
		log.Debugf("A command was received: %s", stringLine)

		// Parse command and (optionally) return response (if any)
		manager.parseClientCommand(stringLine, &client)
	}

	log.Info("A connection was closed")
}

func (manager *ConnectionManager) parseClientCommand(command string, client *Client) {
	commandTokens := strings.Fields(command)

	if len(commandTokens) == 0 {
		return
	}

	switch strings.ToUpper(commandTokens[0]) {
	case "HELP":
		manager.sendStringToClient(HELPSTRING, client)
	case "PUB":
		manager.qm.Publish(commandTokens[1], strings.Join(commandTokens[2:], " "))
		manager.sendStringToClient("PUBACK", client)
	case "SUB":
		manager.qm.Subscribe(commandTokens[1], client)
	case "PINGREQ":
		manager.sendStringToClient("PINGRESP", client)
	default:
		manager.sendStringToClient(UNRECOGNISEDCOMMANDTEXT, client)
	}
}

func (manager *ConnectionManager) sendStringToClient(toSend string, client *Client) {
	fmt.Fprintln(client.Writer, toSend)
	client.Writer.Flush()
}
<commit_msg>Improve 'HELP' command output.<commit_after>package gamq

import (
	"bufio"
	"fmt"
	log "github.com/cihub/seelog"
	"math/rand"
	"net"
	"strconv"
	"strings"
	"sync"
	"time"
)

const (
	UNRECOGNISEDCOMMANDTEXT = "Unrecognised command"
	HELPSTRING              = `Available commands:
	HELP: Prints this text
	PUB <queue> <message>: Publish <message> to <queue>
	SUB <queue>: Subscribe to messages on <queue>
	PINGREQ: Requests a PINGRESP from the server`
)

type ConnectionManager struct {
	wg   sync.WaitGroup
	qm   QueueManager
	rand *rand.Rand
}

func (manager *ConnectionManager) Initialize(config *Config) {
	// Initialize our random number generator (used for naming new connections)
	// A different seed will be used on each startup, for no good reason
	manager.rand = rand.New(rand.NewSource(time.Now().UnixNano()))

	manager.qm = QueueManager{}
	manager.qm.Initialize()

	ln, err := net.Listen("tcp", fmt.Sprintf(":%d", config.Port))

	if err != nil {
		log.Errorf("An error occured whilst opening a socket for reading: %s\n",
			err.Error())
	}

	for {
		conn, err := ln.Accept()
		if err != nil {
			log.Errorf("An error occured whilst opening a socket for reading: %s\n",
				err.Error())
		}
		manager.wg.Add(1)
		go manager.handleConnection(&conn)
	}
}

func (manager *ConnectionManager) handleConnection(conn *net.Conn) {
	defer manager.wg.Done()
	connReader := bufio.NewReader(*conn)
	connWriter := bufio.NewWriter(*conn)
	client := Client{Name: strconv.Itoa(manager.rand.Int()),
		Writer: connWriter,
		Reader: connReader}

	for {
		// Read until newline
		line, err := client.Reader.ReadString('\n')

		if err != nil {
			// Connection has been closed
			break
		}

		stringLine := string(line[:len(line)])
		log.Debugf("A command was received: %s", stringLine)

		// Parse command and (optionally) return response (if any)
		manager.parseClientCommand(stringLine, &client)
	}

	log.Info("A connection was closed")
}

func (manager *ConnectionManager) parseClientCommand(command string, client *Client) {
	commandTokens := strings.Fields(command)

	if len(commandTokens) == 0 {
		return
	}

	switch strings.ToUpper(commandTokens[0]) {
	case "HELP":
		manager.sendStringToClient(HELPSTRING, client)
	case "PUB":
		manager.qm.Publish(commandTokens[1], strings.Join(commandTokens[2:], " "))
		manager.sendStringToClient("PUBACK", client)
	case "SUB":
		manager.qm.Subscribe(commandTokens[1], client)
	case "PINGREQ":
		manager.sendStringToClient("PINGRESP", client)
	default:
		manager.sendStringToClient(UNRECOGNISEDCOMMANDTEXT, client)
	}
}

func (manager *ConnectionManager) sendStringToClient(toSend string, client *Client) {
	fmt.Fprintln(client.Writer, toSend)
	client.Writer.Flush()
}
<|endoftext|>


<commit_before>package azure

import (
	"bufio"
	"encoding/xml"
	"fmt"
	"net"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/coreos/coreos-metadata/src/config"
	"github.com/coreos/coreos-metadata/src/retry"
)

const (
	AgentName             = "com.coreos.metadata"
	FabricProtocolVersion = "2012-11-30"
)

func FetchMetadata() (config.Metadata, error) {
	addr, err := getFabricAddress()
	if err != nil {
		return config.Metadata{}, err
	}

	if err := assertFabricCompatible(addr, FabricProtocolVersion); err != nil {
		return config.Metadata{}, err
	}

	return fetchSharedConfig(addr)
}

func findLease() (*os.File, error) {
	ifaces, err := net.Interfaces()
	if err != nil {
		return nil, fmt.Errorf("could not list interfaces: %v", err)
	}

	for _, iface := range ifaces {
		lease, err := os.Open(fmt.Sprintf("/run/systemd/netif/leases/%d", iface.Index))
		if os.IsNotExist(err) {
			continue
		} else if err != nil {
			return nil, err
		} else {
			return lease, nil
		}
	}

	return nil, fmt.Errorf("could not find any leases")
}

func getFabricAddress() (net.IP, error) {
	lease, err := findLease()
	if err != nil {
		return nil, err
	}
	defer lease.Close()

	var rawEndpoint string
	line := bufio.NewScanner(lease)
	for line.Scan() {
		parts := strings.Split(line.Text(), "=")
		if parts[0] == "OPTION_245" && len(parts) == 2 {
			rawEndpoint = parts[1]
			break
		}
	}

	if len(rawEndpoint) == 0 || len(rawEndpoint) != 8 {
		return nil, fmt.Errorf("fabric endpoint not found in leases")
	}

	octets := make([]byte, 4)
	for i := 0; i < 4; i++ {
		octet, err := strconv.ParseUint(rawEndpoint[2*i:2*i+2], 16, 8)
		if err != nil {
			return nil, err
		}
		octets[i] = byte(octet)
	}

	return net.IPv4(octets[0], octets[1], octets[2], octets[3]), nil
}

func assertFabricCompatible(endpoint net.IP, desiredVersion string) error {
	body, err := retry.Client{
		InitialBackoff: time.Second,
		MaxBackoff:     time.Second * 5,
		MaxAttempts:    10,
	}.Get(fmt.Sprintf("http://%s/?comp=versions", endpoint))
	if err != nil {
		return fmt.Errorf("failed to fetch versions: %v", err)
	}

	versions := struct {
		Supported struct {
			Versions []string `xml:"Version"`
		}
		Preferred struct {
			Version string
		}
	}{}

	if err := xml.Unmarshal(body, &versions); err != nil {
		return fmt.Errorf("failed to unmarshal response: %v", err)
	}

	for _, version := range versions.Supported.Versions {
		if version == desiredVersion {
			return nil
		}
	}

	return fmt.Errorf("fabric version %s is not compatible", desiredVersion)
}

func fetchSharedConfig(endpoint net.IP) (config.Metadata, error) {
	client := retry.Client{
		InitialBackoff: time.Second,
		MaxBackoff:     time.Second * 5,
		MaxAttempts:    10,
		Header: map[string][]string{
			"x-ms-agent-name": {AgentName},
			"x-ms-version":    {FabricProtocolVersion},
			"Content-Type":    {"text/xml; charset=utf-8"},
		},
	}

	body, err := client.Get(fmt.Sprintf("http://%s/machine/?comp=goalstate", endpoint))
	if err != nil {
		return config.Metadata{}, fmt.Errorf("failed to fetch goal state: %v", err)
	}

	goal := struct {
		Container struct {
			RoleInstanceList struct {
				RoleInstance struct {
					Configuration struct {
						SharedConfig string
					}
				}
			}
		}
	}{}

	if err := xml.Unmarshal(body, &goal); err != nil {
		return config.Metadata{}, fmt.Errorf("failed to unmarshal response: %v", err)
	}

	body, err = client.Get(goal.Container.RoleInstanceList.RoleInstance.Configuration.SharedConfig)
	if err != nil {
		return config.Metadata{}, fmt.Errorf("failed to fetch shared config: %v", err)
	}

	sharedConfig := struct {
		Incarnation struct {
			Instance string `xml:"instance,attr"`
		}
		Instances struct {
			Instances []struct {
				Id             string `xml:"id,attr"`
				Address        string `xml:"address,attr"`
				InputEndpoints struct {
					Endpoints []struct {
						LoadBalancedPublicAddress string `xml:"loadBalancedPublicAddress,attr"`
					} `xml:"Endpoint"`
				}
			} `xml:"Instance"`
		}
	}{}

	if err := xml.Unmarshal(body, &sharedConfig); err != nil {
		return config.Metadata{}, err
	}

	metadata := config.Metadata{}
	for _, i := range sharedConfig.Instances.Instances {
		if i.Id == sharedConfig.Incarnation.Instance {
			metadata.LocalIPv4 = net.ParseIP(i.Address)

			for _, e := range i.InputEndpoints.Endpoints {
				host, _, err := net.SplitHostPort(e.LoadBalancedPublicAddress)
				if err == nil {
					metadata.PublicIPv4 = net.ParseIP(host)
					break
				}
			}

			break
		}
	}

	return metadata, nil
}
<commit_msg>providers/azure: use retry.Client.Getf() where appropriate<commit_after>package azure

import (
	"bufio"
	"encoding/xml"
	"fmt"
	"net"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/coreos/coreos-metadata/src/config"
	"github.com/coreos/coreos-metadata/src/retry"
)

const (
	AgentName             = "com.coreos.metadata"
	FabricProtocolVersion = "2012-11-30"
)

func FetchMetadata() (config.Metadata, error) {
	addr, err := getFabricAddress()
	if err != nil {
		return config.Metadata{}, err
	}

	if err := assertFabricCompatible(addr, FabricProtocolVersion); err != nil {
		return config.Metadata{}, err
	}

	return fetchSharedConfig(addr)
}

func findLease() (*os.File, error) {
	ifaces, err := net.Interfaces()
	if err != nil {
		return nil, fmt.Errorf("could not list interfaces: %v", err)
	}

	for _, iface := range ifaces {
		lease, err := os.Open(fmt.Sprintf("/run/systemd/netif/leases/%d", iface.Index))
		if os.IsNotExist(err) {
			continue
		} else if err != nil {
			return nil, err
		} else {
			return lease, nil
		}
	}

	return nil, fmt.Errorf("could not find any leases")
}

func getFabricAddress() (net.IP, error) {
	lease, err := findLease()
	if err != nil {
		return nil, err
	}
	defer lease.Close()

	var rawEndpoint string
	line := bufio.NewScanner(lease)
	for line.Scan() {
		parts := strings.Split(line.Text(), "=")
		if parts[0] == "OPTION_245" && len(parts) == 2 {
			rawEndpoint = parts[1]
			break
		}
	}

	if len(rawEndpoint) == 0 || len(rawEndpoint) != 8 {
		return nil, fmt.Errorf("fabric endpoint not found in leases")
	}

	octets := make([]byte, 4)
	for i := 0; i < 4; i++ {
		octet, err := strconv.ParseUint(rawEndpoint[2*i:2*i+2], 16, 8)
		if err != nil {
			return nil, err
		}
		octets[i] = byte(octet)
	}

	return net.IPv4(octets[0], octets[1], octets[2], octets[3]), nil
}

func assertFabricCompatible(endpoint net.IP, desiredVersion string) error {
	body, err := retry.Client{
		InitialBackoff: time.Second,
		MaxBackoff:     time.Second * 5,
		MaxAttempts:    10,
	}.Getf("http://%s/?comp=versions", endpoint)
	if err != nil {
		return fmt.Errorf("failed to fetch versions: %v", err)
	}

	versions := struct {
		Supported struct {
			Versions []string `xml:"Version"`
		}
		Preferred struct {
			Version string
		}
	}{}

	if err := xml.Unmarshal(body, &versions); err != nil {
		return fmt.Errorf("failed to unmarshal response: %v", err)
	}

	for _, version := range versions.Supported.Versions {
		if version == desiredVersion {
			return nil
		}
	}

	return fmt.Errorf("fabric version %s is not compatible", desiredVersion)
}

func fetchSharedConfig(endpoint net.IP) (config.Metadata, error) {
	client := retry.Client{
		InitialBackoff: time.Second,
		MaxBackoff:     time.Second * 5,
		MaxAttempts:    10,
		Header: map[string][]string{
			"x-ms-agent-name": {AgentName},
			"x-ms-version":    {FabricProtocolVersion},
			"Content-Type":    {"text/xml; charset=utf-8"},
		},
	}

	body, err := client.Getf("http://%s/machine/?comp=goalstate", endpoint)
	if err != nil {
		return config.Metadata{}, fmt.Errorf("failed to fetch goal state: %v", err)
	}

	goal := struct {
		Container struct {
			RoleInstanceList struct {
				RoleInstance struct {
					Configuration struct {
						SharedConfig string
					}
				}
			}
		}
	}{}

	if err := xml.Unmarshal(body, &goal); err != nil {
		return config.Metadata{}, fmt.Errorf("failed to unmarshal response: %v", err)
	}

	body, err = client.Get(goal.Container.RoleInstanceList.RoleInstance.Configuration.SharedConfig)
	if err != nil {
		return config.Metadata{}, fmt.Errorf("failed to fetch shared config: %v", err)
	}

	sharedConfig := struct {
		Incarnation struct {
			Instance string `xml:"instance,attr"`
		}
		Instances struct {
			Instances []struct {
				Id             string `xml:"id,attr"`
				Address        string `xml:"address,attr"`
				InputEndpoints struct {
					Endpoints []struct {
						LoadBalancedPublicAddress string `xml:"loadBalancedPublicAddress,attr"`
					} `xml:"Endpoint"`
				}
			} `xml:"Instance"`
		}
	}{}

	if err := xml.Unmarshal(body, &sharedConfig); err != nil {
		return config.Metadata{}, err
	}

	metadata := config.Metadata{}
	for _, i := range sharedConfig.Instances.Instances {
		if i.Id == sharedConfig.Incarnation.Instance {
			metadata.LocalIPv4 = net.ParseIP(i.Address)

			for _, e := range i.InputEndpoints.Endpoints {
				host, _, err := net.SplitHostPort(e.LoadBalancedPublicAddress)
				if err == nil {
					metadata.PublicIPv4 = net.ParseIP(host)
					break
				}
			}

			break
		}
	}

	return metadata, nil
}
<|endoftext|>


<commit_before>package konnectors

import (
	"archive/tar"
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path"
	"runtime"
	"time"

	"github.com/cozy/cozy-stack/pkg/apps"
	"github.com/cozy/cozy-stack/pkg/config"
	"github.com/cozy/cozy-stack/pkg/consts"
	"github.com/cozy/cozy-stack/pkg/couchdb"
	"github.com/cozy/cozy-stack/pkg/instance"
	"github.com/cozy/cozy-stack/pkg/jobs"
	"github.com/cozy/cozy-stack/pkg/logger"
	"github.com/cozy/cozy-stack/pkg/realtime"
	"github.com/cozy/cozy-stack/pkg/stack"
	"github.com/cozy/cozy-stack/pkg/workers/mails"
	"github.com/sirupsen/logrus"

	"github.com/spf13/afero"
)

func init() {
	jobs.AddWorker("konnector", &jobs.WorkerConfig{
		Concurrency:  runtime.NumCPU(),
		MaxExecCount: 2,
		MaxExecTime:  200 * time.Second,
		Timeout:      200 * time.Second,
		WorkerFunc:   Worker,
		WorkerCommit: commit,
	})
}

// Options contains the options to execute a konnector.
type Options struct {
	Konnector    string `json:"konnector"`
	Account      string `json:"account"`
	FolderToSave string `json:"folder_to_save"`
}

// result stores the result of a konnector execution.
type result struct {
	DocID       string    `json:"_id,omitempty"`
	DocRev      string    `json:"_rev,omitempty"`
	CreatedAt   time.Time `json:"last_execution"`
	LastSuccess time.Time `json:"last_sucess"`
	Account     string    `json:"account"`
	State       string    `json:"state"`
	Error       string    `json:"error"`
}

func (r *result) ID() string         { return r.DocID }
func (r *result) Rev() string        { return r.DocRev }
func (r *result) DocType() string    { return consts.KonnectorResults }
func (r *result) Clone() couchdb.Doc { return r }
func (r *result) SetID(id string)    { r.DocID = id }
func (r *result) SetRev(rev string)  { r.DocRev = rev }

const konnectorMsgTypeError string = "error"

// const konnectorMsgTypeDebug string = "debug"
// const konnectorMsgTypeWarning string = "warning"
// const konnectorMsgTypeProgress string = "progress"

type konnectorMsg struct {
	Type    string `json:"type"`
	Message string `json:"message"`
}

// Worker is the worker that runs a konnector by executing an external process.
func Worker(ctx context.Context, m *jobs.Message) error {
	opts := &Options{}
	if err := m.Unmarshal(&opts); err != nil {
		return err
	}

	slug := opts.Konnector
	fields := struct {
		Account      string `json:"account"`
		FolderToSave string `json:"folder_to_save"`
	}{
		Account:      opts.Account,
		FolderToSave: opts.FolderToSave,
	}
	domain := ctx.Value(jobs.ContextDomainKey).(string)
	worker := ctx.Value(jobs.ContextWorkerKey).(string)
	jobID := fmt.Sprintf("%s/%s/%s", worker, slug, domain)

	inst, err := instance.Get(domain)
	if err != nil {
		return err
	}

	man, err := apps.GetKonnectorBySlug(inst, slug)
	if err != nil {
		return err
	}
	if man.State() != apps.Ready {
		return errors.New("Konnector is not ready")
	}

	token := inst.BuildKonnectorToken(man)

	osFS := afero.NewOsFs()
	workDir, err := afero.TempDir(osFS, "", "konnector-"+slug)
	if err != nil {
		return err
	}
	defer osFS.RemoveAll(workDir)
	workFS := afero.NewBasePathFs(osFS, workDir)

	fileServer := inst.KonnectorsFileServer()
	tarFile, err := fileServer.Open(slug, man.Version(), apps.KonnectorArchiveName)
	if err != nil {
		return err
	}

	tr := tar.NewReader(tarFile)
	for {
		var hdr *tar.Header
		hdr, err = tr.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}
		dirname := path.Dir(hdr.Name)
		if dirname != "." {
			if err = workFS.MkdirAll(dirname, 0755); err != nil {
				return nil
			}
		}
		var f afero.File
		f, err = workFS.OpenFile(hdr.Name, os.O_CREATE|os.O_WRONLY, os.FileMode(hdr.Mode))
		if err != nil {
			return err
		}
		_, err = io.Copy(f, tr)
		if err != nil {
			return err
		}
	}

	fieldsJSON, err := json.Marshal(fields)
	if err != nil {
		return err
	}

	konnCmd := config.GetConfig().Konnectors.Cmd
	cmd := exec.CommandContext(ctx, konnCmd, workDir) // #nosec
	cmd.Env = []string{
		"COZY_URL=" + inst.PageURL("/", nil),
		"COZY_CREDENTIALS=" + token,
		"COZY_FIELDS=" + string(fieldsJSON),
		"COZY_TYPE=" + man.Type,
		"COZY_JOB_ID=" + jobID,
	}

	cmdErr, err := cmd.StderrPipe()
	if err != nil {
		return err
	}
	cmdOut, err := cmd.StdoutPipe()
	if err != nil {
		return err
	}

	scanErr := bufio.NewScanner(cmdErr)
	scanOut := bufio.NewScanner(cmdOut)
	scanOut.Buffer(nil, 256*1024)

	var msgChan = make(chan konnectorMsg)
	var messages []konnectorMsg

	log := logger.WithDomain(domain)

	go doScanOut(jobID, scanOut, domain, msgChan, log)
	go doScanErr(jobID, scanErr, log)
	go func() {
		hub := realtime.GetHub()
		for msg := range msgChan {
			messages = append(messages, msg)
			hub.Publish(&realtime.Event{
				Type: realtime.EventCreate,
				Doc: couchdb.JSONDoc{Type: consts.JobEvents, M: map[string]interface{}{
					"type":    msg.Type,
					"message": msg.Message,
				}},
				Domain: domain,
			})
		}
	}()

	if err = cmd.Start(); err != nil {
		return wrapErr(ctx, err)
	}

	err = cmd.Wait()
	if err != nil {
		err = wrapErr(ctx, err)
	}

	close(msgChan)
	for _, msg := range messages {
		if msg.Type == konnectorMsgTypeError {
			// konnector err is more explicit
			return errors.New(msg.Message)
		}
	}

	return err
}

func doScanOut(jobID string, scanner *bufio.Scanner, domain string,
	msgs chan konnectorMsg, log *logrus.Entry) {
	for scanner.Scan() {
		linebb := scanner.Bytes()
		from := bytes.IndexByte(linebb, '{')
		to := bytes.LastIndexByte(linebb, '}')
		var msg konnectorMsg
		log.Infof("[konnector] %s: Stdout: %s", jobID, string(linebb))
		if from > -1 && from < to && to > -1 {
			err := json.Unmarshal(linebb[from:to+1], &msg)
			if err == nil {
				msgs <- msg
				continue
			}
		}
		log.Warnf("[konnector] %s: Could not parse as JSON", jobID)
	}
	if err := scanner.Err(); err != nil {
		log.Errorf("[konnector] %s: Error while reading stdout: %s", jobID, err)
	}
}

func doScanErr(jobID string, scanner *bufio.Scanner, log *logrus.Entry) {
	for scanner.Scan() {
		log.Errorf("[konnector] %s: Stderr: %s", jobID, scanner.Text())
	}
	if err := scanner.Err(); err != nil {
		log.Errorf("[konnector] %s: Error while reading stderr: %s", jobID, err)
	}
}

func commit(ctx context.Context, m *jobs.Message, errjob error) error {
	opts := &Options{}
	if err := m.Unmarshal(&opts); err != nil {
		return err
	}

	slug := opts.Konnector
	domain := ctx.Value(jobs.ContextDomainKey).(string)

	inst, err := instance.Get(domain)
	if err != nil {
		return err
	}

	lastResult := &result{}
	err = couchdb.GetDoc(inst, consts.KonnectorResults, slug, lastResult)
	if err != nil {
		if !couchdb.IsNotFoundError(err) {
			return err
		}
		lastResult = nil
	}

	var state, errstr string
	var lastSuccess time.Time
	if errjob != nil {
		if lastResult != nil {
			lastSuccess = lastResult.LastSuccess
		}
		errstr = errjob.Error()
		state = jobs.Errored
	} else {
		lastSuccess = time.Now()
		state = jobs.Done
	}
	result := &result{
		DocID:       slug,
		Account:     opts.Account,
		CreatedAt:   time.Now(),
		LastSuccess: lastSuccess,
		State:       state,
		Error:       errstr,
	}
	if lastResult == nil {
		err = couchdb.CreateNamedDocWithDB(inst, result)
	} else {
		result.SetRev(lastResult.Rev())
		err = couchdb.UpdateDoc(inst, result)
	}
	if err != nil {
		return err
	}

	// if it is the first try we do not take into account an error, we bail.
	if lastResult == nil {
		return nil
	}
	// if the job has not errored, or the last one was already errored, we bail.
	if state != jobs.Errored || lastResult.State == jobs.Errored {
		return nil
	}

	konnectorURL := inst.SubDomain(consts.CollectSlug)
	konnectorURL.Fragment = "/category/all/" + slug
	msg, err := jobs.NewMessage(jobs.JSONEncoding, &mails.Options{
		Mode:         mails.ModeNoReply,
		Subject:      inst.Translate("Konnector execution error"),
		TemplateName: "konnector_error_" + inst.Locale,
		TemplateValues: map[string]string{
			"KonnectorName": slug,
			"KonnectorPage": konnectorURL.String(),
		},
	})
	if err != nil {
		return err
	}
	_, err = stack.GetBroker().PushJob(&jobs.JobRequest{
		Domain:     domain,
		WorkerType: "sendmail",
		Message:    msg,
	})
	return err
}

func wrapErr(ctx context.Context, err error) error {
	if ctx.Err() == context.DeadlineExceeded {
		return context.DeadlineExceeded
	}
	return err
}
<commit_msg>fix bug where result was only saved the first time<commit_after>package konnectors

import (
	"archive/tar"
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path"
	"runtime"
	"time"

	"github.com/cozy/cozy-stack/pkg/apps"
	"github.com/cozy/cozy-stack/pkg/config"
	"github.com/cozy/cozy-stack/pkg/consts"
	"github.com/cozy/cozy-stack/pkg/couchdb"
	"github.com/cozy/cozy-stack/pkg/instance"
	"github.com/cozy/cozy-stack/pkg/jobs"
	"github.com/cozy/cozy-stack/pkg/logger"
	"github.com/cozy/cozy-stack/pkg/realtime"
	"github.com/cozy/cozy-stack/pkg/stack"
	"github.com/cozy/cozy-stack/pkg/workers/mails"
	"github.com/sirupsen/logrus"

	"github.com/spf13/afero"
)

func init() {
	jobs.AddWorker("konnector", &jobs.WorkerConfig{
		Concurrency:  runtime.NumCPU(),
		MaxExecCount: 2,
		MaxExecTime:  200 * time.Second,
		Timeout:      200 * time.Second,
		WorkerFunc:   Worker,
		WorkerCommit: commit,
	})
}

// Options contains the options to execute a konnector.
type Options struct {
	Konnector    string `json:"konnector"`
	Account      string `json:"account"`
	FolderToSave string `json:"folder_to_save"`
}

// result stores the result of a konnector execution.
type result struct {
	DocID       string    `json:"_id,omitempty"`
	DocRev      string    `json:"_rev,omitempty"`
	CreatedAt   time.Time `json:"last_execution"`
	LastSuccess time.Time `json:"last_success"`
	Account     string    `json:"account"`
	State       string    `json:"state"`
	Error       string    `json:"error"`
}

func (r *result) ID() string         { return r.DocID }
func (r *result) Rev() string        { return r.DocRev }
func (r *result) DocType() string    { return consts.KonnectorResults }
func (r *result) Clone() couchdb.Doc { c := *r; return &c }
func (r *result) SetID(id string)    { r.DocID = id }
func (r *result) SetRev(rev string)  { r.DocRev = rev }

const konnectorMsgTypeError string = "error"

// const konnectorMsgTypeDebug string = "debug"
// const konnectorMsgTypeWarning string = "warning"
// const konnectorMsgTypeProgress string = "progress"

type konnectorMsg struct {
	Type    string `json:"type"`
	Message string `json:"message"`
}

// Worker is the worker that runs a konnector by executing an external process.
func Worker(ctx context.Context, m *jobs.Message) error {
	opts := &Options{}
	if err := m.Unmarshal(&opts); err != nil {
		return err
	}

	slug := opts.Konnector
	fields := struct {
		Account      string `json:"account"`
		FolderToSave string `json:"folder_to_save"`
	}{
		Account:      opts.Account,
		FolderToSave: opts.FolderToSave,
	}
	domain := ctx.Value(jobs.ContextDomainKey).(string)
	worker := ctx.Value(jobs.ContextWorkerKey).(string)
	jobID := fmt.Sprintf("%s/%s/%s", worker, slug, domain)

	inst, err := instance.Get(domain)
	if err != nil {
		return err
	}

	man, err := apps.GetKonnectorBySlug(inst, slug)
	if err != nil {
		return err
	}
	if man.State() != apps.Ready {
		return errors.New("Konnector is not ready")
	}

	token := inst.BuildKonnectorToken(man)

	osFS := afero.NewOsFs()
	workDir, err := afero.TempDir(osFS, "", "konnector-"+slug)
	if err != nil {
		return err
	}
	defer osFS.RemoveAll(workDir)
	workFS := afero.NewBasePathFs(osFS, workDir)

	fileServer := inst.KonnectorsFileServer()
	tarFile, err := fileServer.Open(slug, man.Version(), apps.KonnectorArchiveName)
	if err != nil {
		return err
	}

	tr := tar.NewReader(tarFile)
	for {
		var hdr *tar.Header
		hdr, err = tr.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}
		dirname := path.Dir(hdr.Name)
		if dirname != "." {
			if err = workFS.MkdirAll(dirname, 0755); err != nil {
				return nil
			}
		}
		var f afero.File
		f, err = workFS.OpenFile(hdr.Name, os.O_CREATE|os.O_WRONLY, os.FileMode(hdr.Mode))
		if err != nil {
			return err
		}
		_, err = io.Copy(f, tr)
		if err != nil {
			return err
		}
	}

	fieldsJSON, err := json.Marshal(fields)
	if err != nil {
		return err
	}

	konnCmd := config.GetConfig().Konnectors.Cmd
	cmd := exec.CommandContext(ctx, konnCmd, workDir) // #nosec
	cmd.Env = []string{
		"COZY_URL=" + inst.PageURL("/", nil),
		"COZY_CREDENTIALS=" + token,
		"COZY_FIELDS=" + string(fieldsJSON),
		"COZY_TYPE=" + man.Type,
		"COZY_JOB_ID=" + jobID,
	}

	cmdErr, err := cmd.StderrPipe()
	if err != nil {
		return err
	}
	cmdOut, err := cmd.StdoutPipe()
	if err != nil {
		return err
	}

	scanErr := bufio.NewScanner(cmdErr)
	scanOut := bufio.NewScanner(cmdOut)
	scanOut.Buffer(nil, 256*1024)

	var msgChan = make(chan konnectorMsg)
	var messages []konnectorMsg

	log := logger.WithDomain(domain)

	go doScanOut(jobID, scanOut, domain, msgChan, log)
	go doScanErr(jobID, scanErr, log)
	go func() {
		hub := realtime.GetHub()
		for msg := range msgChan {
			messages = append(messages, msg)
			hub.Publish(&realtime.Event{
				Type: realtime.EventCreate,
				Doc: couchdb.JSONDoc{Type: consts.JobEvents, M: map[string]interface{}{
					"type":    msg.Type,
					"message": msg.Message,
				}},
				Domain: domain,
			})
		}
	}()

	if err = cmd.Start(); err != nil {
		return wrapErr(ctx, err)
	}

	err = cmd.Wait()
	if err != nil {
		err = wrapErr(ctx, err)
	}

	close(msgChan)
	for _, msg := range messages {
		if msg.Type == konnectorMsgTypeError {
			// konnector err is more explicit
			return errors.New(msg.Message)
		}
	}

	return err
}

func doScanOut(jobID string, scanner *bufio.Scanner, domain string,
	msgs chan konnectorMsg, log *logrus.Entry) {
	for scanner.Scan() {
		linebb := scanner.Bytes()
		from := bytes.IndexByte(linebb, '{')
		to := bytes.LastIndexByte(linebb, '}')
		var msg konnectorMsg
		log.Infof("[konnector] %s: Stdout: %s", jobID, string(linebb))
		if from > -1 && from < to && to > -1 {
			err := json.Unmarshal(linebb[from:to+1], &msg)
			if err == nil {
				msgs <- msg
				continue
			}
		}
		log.Warnf("[konnector] %s: Could not parse as JSON", jobID)
	}
	if err := scanner.Err(); err != nil {
		log.Errorf("[konnector] %s: Error while reading stdout: %s", jobID, err)
	}
}

func doScanErr(jobID string, scanner *bufio.Scanner, log *logrus.Entry) {
	for scanner.Scan() {
		log.Errorf("[konnector] %s: Stderr: %s", jobID, scanner.Text())
	}
	if err := scanner.Err(); err != nil {
		log.Errorf("[konnector] %s: Error while reading stderr: %s", jobID, err)
	}
}

func commit(ctx context.Context, m *jobs.Message, errjob error) error {
	opts := &Options{}
	if err := m.Unmarshal(&opts); err != nil {
		return err
	}

	slug := opts.Konnector
	domain := ctx.Value(jobs.ContextDomainKey).(string)

	inst, err := instance.Get(domain)
	if err != nil {
		return err
	}

	lastResult := &result{}
	err = couchdb.GetDoc(inst, consts.KonnectorResults, slug, lastResult)
	if err != nil {
		if !couchdb.IsNotFoundError(err) {
			return err
		}
		lastResult = nil
	}

	var state, errstr string
	var lastSuccess time.Time
	if errjob != nil {
		if lastResult != nil {
			lastSuccess = lastResult.LastSuccess
		}
		errstr = errjob.Error()
		state = jobs.Errored
	} else {
		lastSuccess = time.Now()
		state = jobs.Done
	}
	result := &result{
		DocID:       slug,
		Account:     opts.Account,
		CreatedAt:   time.Now(),
		LastSuccess: lastSuccess,
		State:       state,
		Error:       errstr,
	}
	if lastResult == nil {
		err = couchdb.CreateNamedDocWithDB(inst, result)
	} else {
		result.SetRev(lastResult.Rev())
		err = couchdb.UpdateDoc(inst, result)
	}
	if err != nil {
		return err
	}

	// if it is the first try we do not take into account an error, we bail.
	if lastResult == nil {
		return nil
	}
	// if the job has not errored, or the last one was already errored, we bail.
	if state != jobs.Errored || lastResult.State == jobs.Errored {
		return nil
	}

	konnectorURL := inst.SubDomain(consts.CollectSlug)
	konnectorURL.Fragment = "/category/all/" + slug
	msg, err := jobs.NewMessage(jobs.JSONEncoding, &mails.Options{
		Mode:         mails.ModeNoReply,
		Subject:      inst.Translate("Konnector execution error"),
		TemplateName: "konnector_error_" + inst.Locale,
		TemplateValues: map[string]string{
			"KonnectorName": slug,
			"KonnectorPage": konnectorURL.String(),
		},
	})
	if err != nil {
		return err
	}
	_, err = stack.GetBroker().PushJob(&jobs.JobRequest{
		Domain:     domain,
		WorkerType: "sendmail",
		Message:    msg,
	})
	return err
}

func wrapErr(ctx context.Context, err error) error {
	if ctx.Err() == context.DeadlineExceeded {
		return context.DeadlineExceeded
	}
	return err
}
<|endoftext|>


<commit_before>package lecimg

import (
	"fmt"
	"image"
	"image/color"
	"image/draw"
	"log"
	"sort"

	"github.com/mitchellh/mapstructure"
)

// ChangeLineSpaceOption contains options for changeLineSpace filter
type ChangeLineSpaceOption struct {
	WidthRatio         float64
	HeightRatio        float64
	LineSpaceScale     float64
	MinSpace           int
	MaxRemove          int
	Threshold          uint32
	EmptyLineThreshold float64
	DebugMode          bool
}

func NewChangeLineSpaceOption(m map[string]interface{}) (*ChangeLineSpaceOption, error) {
	option := ChangeLineSpaceOption{}

	err := mapstructure.Decode(m, &option)
	if err != nil {
		return nil, err
	}

	return &option, nil
}

type ChangeLineSpaceResult struct {
	image image.Image
	rect  image.Rectangle
}

func (r ChangeLineSpaceResult) Img() image.Image {
	return r.image
}

func (r ChangeLineSpaceResult) Log() {
}

// ----------------------------------------------------------------------------

type lineRange struct {
	start        int
	end          int
	height       int
	targetHeight int
	emptyLine    bool
}

func (r *lineRange) calc(scale float64, minHeight, maxRemove int) {
	r.height = r.end - r.start + 1
	if !r.emptyLine || r.height <= minHeight {
		r.targetHeight = r.height
	} else {
		if maxRemove > 0 {
			r.targetHeight = Max(minHeight, int(float64(r.height)*scale+0.5))
			if removed := r.height - r.targetHeight; removed > maxRemove {
				r.targetHeight = r.height - maxRemove
			}
		}
	}
}

func (r lineRange) getReducedHeight() int {
	return r.height - r.targetHeight
}

func (r lineRange) String() string {
	return fmt.Sprintf("(%4d-%4d) h:%d, targetH: %d, empty: %v",
		r.start, r.end, r.height, r.targetHeight, r.emptyLine)
}

type lineRanges []lineRange

func (r lineRanges) Len() int {
	return len(r)
}

// Less sorts in descending order by targetHeight
func (r lineRanges) Less(i, j int) bool {
	return r[i].targetHeight > r[j].targetHeight
}
func (r lineRanges) Swap(i, j int) {
	r[i], r[j] = r[j], r[i]
}

// ----------------------------------------------------------------------------

type ChangeLineSpaceFilter struct {
	option ChangeLineSpaceOption
}

func NewChangeLineSpaceFilter(option ChangeLineSpaceOption) *ChangeLineSpaceFilter {
	return &ChangeLineSpaceFilter{option: option}
}

func getBrightness(r, g, b uint32) uint32 {
	return (r + g + b) / 3
}

// getLineRanges returns list of text and empty lines
func (f ChangeLineSpaceFilter) getLineRanges(src image.Image) lineRanges {
	bounds := src.Bounds()
	srcWidth, srcHeight := bounds.Dx(), bounds.Dy()
	threshold16 := f.option.Threshold * 256

	var ranges lineRanges
	var r lineRange

	maxDotCount := int(f.option.EmptyLineThreshold)
	if f.option.EmptyLineThreshold < 1 {
		maxDotCount = int(float64(srcWidth) * f.option.EmptyLineThreshold)
	}
	for y := 0; y < srcHeight; y++ {
		emptyLine := true
		dotCount := 0
		for x := 0; x < srcWidth; x++ {
			r, g, b, _ := src.At(x, y).RGBA()
			brightness := getBrightness(r, g, b)
			if brightness < threshold16 {
				dotCount++
				if dotCount >= maxDotCount {
					emptyLine = false
					break
				}
			}
		}

		if emptyLine {
			if y == 0 {
				r = lineRange{start: y, end: y, emptyLine: true}
			} else {
				if r.emptyLine {
					r.end = y
				} else {
					ranges = append(ranges, r)
					r = lineRange{start: y, end: y, emptyLine: true}
				}
			}
		} else {
			if y == 0 {
				r = lineRange{start: y, end: y, emptyLine: false}
			} else {
				if r.emptyLine {
					ranges = append(ranges, r)
					r = lineRange{start: y, end: y, emptyLine: false}
				} else {
					r.end = y
				}
			}
		}

	}
	ranges = append(ranges, r)
	return ranges
}

func (f ChangeLineSpaceFilter) processLineRanges(ranges lineRanges, width int) int {
	targetHeight := 0
	rangeCount := len(ranges)
	emptyRangeCount := 0
	for i := 0; i < rangeCount; i++ {
		r := &ranges[i]
		r.calc(f.option.LineSpaceScale, f.option.MinSpace, f.option.MaxRemove)
		targetHeight += r.targetHeight
		if r.emptyLine {
			emptyRangeCount++
		}
	}

	minTargetHeight := int(f.option.HeightRatio * float64(width) / f.option.WidthRatio)
	if f.option.DebugMode {
		log.Printf("min targetHeight: %v", minTargetHeight)
	}

	loop := 0
	maxLoopCount := 5
	for targetHeight < minTargetHeight && loop < maxLoopCount {
		totalReducedHeight := 0
		for i := 0; i < rangeCount; i++ {
			r := ranges[i]
			if r.emptyLine {
				totalReducedHeight += r.getReducedHeight()
			}
		}

		totalInc := 0
		if totalReducedHeight > 0 {
			totalEmptyLinesToInc := minTargetHeight - targetHeight
			for i := 0; i < len(ranges); i++ {
				r := &ranges[i]
				if r.emptyLine {
					reducedHeight := r.getReducedHeight()
					heightToInc := Min(reducedHeight,
						reducedHeight*totalEmptyLinesToInc/totalReducedHeight)
					inc := int(float32(heightToInc) + 0.5)
					r.targetHeight = r.targetHeight + inc
					totalInc += inc
				}
			}
		}

		targetHeight = 0
		for i := 0; i < rangeCount; i++ {
			r := &ranges[i]
			targetHeight += r.targetHeight
		}

		if f.option.DebugMode {
			log.Printf("[%d] totalInc=%d, targetHeight=%d", loop, totalInc, targetHeight)
		}

		loop++

		if totalInc <= emptyRangeCount {
			if remainInc := minTargetHeight - targetHeight; remainInc != 0 {
				sortedRanges := make(lineRanges, rangeCount)
				copy(sortedRanges, ranges)
				sort.Sort(sortedRanges)
				for i := 0; i < rangeCount && remainInc != 0; i++ {
					if r := sortedRanges[i]; r.emptyLine {
						if remainInc > 0 {
							r.targetHeight++
							targetHeight++
							remainInc--
						} else {
							r.targetHeight--
							targetHeight--
							remainInc++
						}
					}
				}
			}
			break
		}
	}
	if f.option.DebugMode {
		log.Printf("TargetHeight: %v", targetHeight)
	}

	return targetHeight
}

func (f ChangeLineSpaceFilter) Run(s *FilterSource) FilterResult {
	img, rect := f.run(s.image)
	return &ChangeLineSpaceResult{img, rect}
}

func (f ChangeLineSpaceFilter) run(src image.Image) (image.Image, image.Rectangle) {
	ranges := f.getLineRanges(src)
	rangeCount := len(ranges)

	if rangeCount <= 1 {
		return src, src.Bounds()
	} else {
		width := src.Bounds().Dx()
		targetHeight := f.processLineRanges(ranges, width)

		bounds := image.Rect(0, 0, width, targetHeight)
		dest := CreateImage(width, targetHeight, color.White)
		destY := 0
		for i := 0; i < rangeCount; i++ {
			r := &ranges[i]
			rangeHeight := r.height
			rangeTargetHeight := r.targetHeight

			if rangeHeight > 0 && rangeTargetHeight > 0 {
				rangeImage := CropImage(src, image.Rect(0, r.start, width, r.start+rangeHeight))

				var resized image.Image
				if rangeTargetHeight == rangeHeight {
					resized = rangeImage
				} else {
					resized = ResizeImage(rangeImage, width, rangeTargetHeight, false)
				}

				// copy rangeImage to dest
				draw.Draw(dest,
					image.Rect(0, destY, width, destY+rangeTargetHeight),
					resized,
					image.ZP,
					draw.Src)

				destY += rangeTargetHeight
			}
		}
		SaveJpeg(dest, "temp", "result.jpg", 85)
		return dest, bounds
	}
}
<commit_msg>test code removed<commit_after>package lecimg

import (
	"fmt"
	"image"
	"image/color"
	"image/draw"
	"log"
	"sort"

	"github.com/mitchellh/mapstructure"
)

// ChangeLineSpaceOption contains options for changeLineSpace filter
type ChangeLineSpaceOption struct {
	WidthRatio         float64
	HeightRatio        float64
	LineSpaceScale     float64
	MinSpace           int
	MaxRemove          int
	Threshold          uint32
	EmptyLineThreshold float64
	DebugMode          bool
}

func NewChangeLineSpaceOption(m map[string]interface{}) (*ChangeLineSpaceOption, error) {
	option := ChangeLineSpaceOption{}

	err := mapstructure.Decode(m, &option)
	if err != nil {
		return nil, err
	}

	return &option, nil
}

type ChangeLineSpaceResult struct {
	image image.Image
	rect  image.Rectangle
}

func (r ChangeLineSpaceResult) Img() image.Image {
	return r.image
}

func (r ChangeLineSpaceResult) Log() {
}

// ----------------------------------------------------------------------------

type lineRange struct {
	start        int
	end          int
	height       int
	targetHeight int
	emptyLine    bool
}

func (r *lineRange) calc(scale float64, minHeight, maxRemove int) {
	r.height = r.end - r.start + 1
	if !r.emptyLine || r.height <= minHeight {
		r.targetHeight = r.height
	} else {
		if maxRemove > 0 {
			r.targetHeight = Max(minHeight, int(float64(r.height)*scale+0.5))
			if removed := r.height - r.targetHeight; removed > maxRemove {
				r.targetHeight = r.height - maxRemove
			}
		}
	}
}

func (r lineRange) getReducedHeight() int {
	return r.height - r.targetHeight
}

func (r lineRange) String() string {
	return fmt.Sprintf("(%4d-%4d) h:%d, targetH: %d, empty: %v",
		r.start, r.end, r.height, r.targetHeight, r.emptyLine)
}

type lineRanges []lineRange

func (r lineRanges) Len() int {
	return len(r)
}

// Less sorts in descending order by targetHeight
func (r lineRanges) Less(i, j int) bool {
	return r[i].targetHeight > r[j].targetHeight
}
func (r lineRanges) Swap(i, j int) {
	r[i], r[j] = r[j], r[i]
}

// ----------------------------------------------------------------------------

type ChangeLineSpaceFilter struct {
	option ChangeLineSpaceOption
}

func NewChangeLineSpaceFilter(option ChangeLineSpaceOption) *ChangeLineSpaceFilter {
	return &ChangeLineSpaceFilter{option: option}
}

func getBrightness(r, g, b uint32) uint32 {
	return (r + g + b) / 3
}

// getLineRanges returns list of text and empty lines
func (f ChangeLineSpaceFilter) getLineRanges(src image.Image) lineRanges {
	bounds := src.Bounds()
	srcWidth, srcHeight := bounds.Dx(), bounds.Dy()
	threshold16 := f.option.Threshold * 256

	var ranges lineRanges
	var r lineRange

	maxDotCount := int(f.option.EmptyLineThreshold)
	if f.option.EmptyLineThreshold < 1 {
		maxDotCount = int(float64(srcWidth) * f.option.EmptyLineThreshold)
	}
	for y := 0; y < srcHeight; y++ {
		emptyLine := true
		dotCount := 0
		for x := 0; x < srcWidth; x++ {
			r, g, b, _ := src.At(x, y).RGBA()
			brightness := getBrightness(r, g, b)
			if brightness < threshold16 {
				dotCount++
				if dotCount >= maxDotCount {
					emptyLine = false
					break
				}
			}
		}

		if emptyLine {
			if y == 0 {
				r = lineRange{start: y, end: y, emptyLine: true}
			} else {
				if r.emptyLine {
					r.end = y
				} else {
					ranges = append(ranges, r)
					r = lineRange{start: y, end: y, emptyLine: true}
				}
			}
		} else {
			if y == 0 {
				r = lineRange{start: y, end: y, emptyLine: false}
			} else {
				if r.emptyLine {
					ranges = append(ranges, r)
					r = lineRange{start: y, end: y, emptyLine: false}
				} else {
					r.end = y
				}
			}
		}

	}
	ranges = append(ranges, r)
	return ranges
}

func (f ChangeLineSpaceFilter) processLineRanges(ranges lineRanges, width int) int {
	targetHeight := 0
	rangeCount := len(ranges)
	emptyRangeCount := 0
	for i := 0; i < rangeCount; i++ {
		r := &ranges[i]
		r.calc(f.option.LineSpaceScale, f.option.MinSpace, f.option.MaxRemove)
		targetHeight += r.targetHeight
		if r.emptyLine {
			emptyRangeCount++
		}
	}

	minTargetHeight := int(f.option.HeightRatio * float64(width) / f.option.WidthRatio)
	if f.option.DebugMode {
		log.Printf("min targetHeight: %v", minTargetHeight)
	}

	loop := 0
	maxLoopCount := 5
	for targetHeight < minTargetHeight && loop < maxLoopCount {
		totalReducedHeight := 0
		for i := 0; i < rangeCount; i++ {
			r := ranges[i]
			if r.emptyLine {
				totalReducedHeight += r.getReducedHeight()
			}
		}

		totalInc := 0
		if totalReducedHeight > 0 {
			totalEmptyLinesToInc := minTargetHeight - targetHeight
			for i := 0; i < len(ranges); i++ {
				r := &ranges[i]
				if r.emptyLine {
					reducedHeight := r.getReducedHeight()
					heightToInc := Min(reducedHeight,
						reducedHeight*totalEmptyLinesToInc/totalReducedHeight)
					inc := int(float32(heightToInc) + 0.5)
					r.targetHeight = r.targetHeight + inc
					totalInc += inc
				}
			}
		}

		targetHeight = 0
		for i := 0; i < rangeCount; i++ {
			r := &ranges[i]
			targetHeight += r.targetHeight
		}

		if f.option.DebugMode {
			log.Printf("[%d] totalInc=%d, targetHeight=%d", loop, totalInc, targetHeight)
		}

		loop++

		if totalInc <= emptyRangeCount {
			if remainInc := minTargetHeight - targetHeight; remainInc != 0 {
				sortedRanges := make(lineRanges, rangeCount)
				copy(sortedRanges, ranges)
				sort.Sort(sortedRanges)
				for i := 0; i < rangeCount && remainInc != 0; i++ {
					if r := sortedRanges[i]; r.emptyLine {
						if remainInc > 0 {
							r.targetHeight++
							targetHeight++
							remainInc--
						} else {
							r.targetHeight--
							targetHeight--
							remainInc++
						}
					}
				}
			}
			break
		}
	}
	if f.option.DebugMode {
		log.Printf("TargetHeight: %v", targetHeight)
	}

	return targetHeight
}

func (f ChangeLineSpaceFilter) Run(s *FilterSource) FilterResult {
	img, rect := f.run(s.image)
	return &ChangeLineSpaceResult{img, rect}
}

func (f ChangeLineSpaceFilter) run(src image.Image) (image.Image, image.Rectangle) {
	ranges := f.getLineRanges(src)
	rangeCount := len(ranges)

	if rangeCount <= 1 {
		return src, src.Bounds()
	} else {
		width := src.Bounds().Dx()
		targetHeight := f.processLineRanges(ranges, width)

		bounds := image.Rect(0, 0, width, targetHeight)
		dest := CreateImage(width, targetHeight, color.White)
		destY := 0
		for i := 0; i < rangeCount; i++ {
			r := &ranges[i]
			rangeHeight := r.height
			rangeTargetHeight := r.targetHeight

			if rangeHeight > 0 && rangeTargetHeight > 0 {
				rangeImage := CropImage(src, image.Rect(0, r.start, width, r.start+rangeHeight))

				var resized image.Image
				if rangeTargetHeight == rangeHeight {
					resized = rangeImage
				} else {
					resized = ResizeImage(rangeImage, width, rangeTargetHeight, false)
				}

				// copy rangeImage to dest
				draw.Draw(dest,
					image.Rect(0, destY, width, destY+rangeTargetHeight),
					resized,
					image.ZP,
					draw.Src)

				destY += rangeTargetHeight
			}
		}
		return dest, bounds
	}
}
<|endoftext|>


<commit_before>package metafora

import (
	"math/rand"
	"runtime"
	"sort"
	"sync"
	"time"
)

var (

	// balance calls are randomized and this is the upper bound of the random
	// amount
	balanceJitterMax = 10 * int64(time.Second)

	//FIXME should probably be improved, see usage in Run()
	consumerRetryDelay = 10 * time.Second
)

type runningTask struct {
	h Handler
	c chan struct{}
}

// Consumer is the core Metafora task runner.
type Consumer struct {
	// Func to create new handlers
	handler HandlerFunc

	// Map of task:Handler
	running map[string]runningTask

	// Mutex to protect access to running
	runL sync.Mutex

	// WaitGroup for running handlers and consumer goroutines
	hwg sync.WaitGroup

	bal      Balancer
	balEvery time.Duration
	coord    Coordinator
	logger   *logger
	stop     chan struct{} // closed by Shutdown to cause Run to exit

	// ticked on each loop of the main loop to enforce sequential interaction
	// with coordinator and balancer
	tick chan int

	watch chan string // channel for watcher to send tasks to main loop

	// Set by command handler, read anywhere via Consumer.frozen()
	freezeL sync.Mutex
	freeze  bool
}

// NewConsumer returns a new consumer and calls Init on the Balancer and Coordinator.
func NewConsumer(coord Coordinator, h HandlerFunc, b Balancer) (*Consumer, error) {
	c := &Consumer{
		running:  make(map[string]runningTask),
		handler:  h,
		bal:      b,
		balEvery: 15 * time.Minute, //TODO make balance wait configurable
		coord:    coord,
		logger:   stdoutLogger(),
		stop:     make(chan struct{}),
		tick:     make(chan int),
		watch:    make(chan string),
	}

	// initialize balancer with the consumer and a prefixed logger
	b.Init(&struct {
		*Consumer
		Logger
	}{Consumer: c, Logger: c.logger})

	if err := coord.Init(&coordinatorContext{Consumer: c, Logger: c.logger}); err != nil {
		return nil, err
	}
	return c, nil
}

// SetLogger assigns the logger to use as well as a level
//
// The logger parameter is an interface that requires the following
// method to be implemented (such as the the stdlib log.Logger):
//
//    Output(calldepth int, s string)
//
func (c *Consumer) SetLogger(l logOutputter, lvl LogLevel) {
	c.logger.l = l
	c.logger.lvl = lvl
}

// Run is the core run loop of Metafora. It is responsible for calling into the
// Coordinator to claim work and Balancer to rebalance work.
//
// Run blocks until Shutdown is called or an internal error occurs.
func (c *Consumer) Run() {
	c.logger.Log(LogLevelDebug, "Starting consumer")

	// chans for core goroutines to communicate with main loop
	balance := make(chan bool)
	cmdChan := make(chan Command)

	// Balance is called by the main loop when the balance channel is ticked
	go func() {
		randInt := rand.New(rand.NewSource(time.Now().UnixNano())).Int63n
		for {
			select {
			case <-c.stop:
				// Shutdown has been called.
				return
			case <-time.After(c.balEvery + time.Duration(randInt(balanceJitterMax))):
				c.logger.Log(LogLevelInfo, "Balancing")
				select {
				case balance <- true:
					// Ticked balance
				case <-c.stop:
					// Shutdown has been called.
					return
				}
			}
			// Wait for main loop to signal balancing is done
			select {
			case <-c.stop:
				return
			case <-c.tick:
			}
		}
	}()

	// Watch for new tasks in a goroutine
	go c.watcher()

	// Watch for new commands in a goroutine
	go func() {
		defer close(cmdChan)
		for {
			cmd, err := c.coord.Command()
			if err != nil {
				//FIXME add more sophisticated error handling
				c.logger.Log(LogLevelError, "Coordinator returned an error during command, waiting and retrying. %v", err)
				select {
				case <-c.stop:
					return
				case <-time.After(consumerRetryDelay):
				}
				continue
			}
			if cmd == nil {
				c.logger.Log(LogLevelDebug, "Command coordinator exited")
				return
			}
			// Send command to watcher (or shutdown)
			select {
			case <-c.stop:
				return
			case cmdChan <- cmd:
			}
			// Wait for main loop to signal command has been handled
			select {
			case <-c.stop:
				return
			case <-c.tick:
			}
		}
	}()

	// Make sure Run() cleans up on exit (stops coordinator, releases tasks, etc)
	defer c.shutdown()

	// Main Loop ensures events are processed synchronously
	for {
		if c.frozen() {
			// Only recv commands while frozen
			select {
			case <-c.stop:
				// Shutdown has been called.
				return
			case cmd, ok := <-cmdChan:
				if !ok {
					c.logger.Log(LogLevelDebug, "Command channel closed. Exiting main loop.")
					return
				}
				c.logger.Log(LogLevelDebug, "Received command: %s", cmd)
				c.handleCommand(cmd)
			}
			// Must send tick whenever main loop restarts
			select {
			case <-c.stop:
				return
			case c.tick <- 1:
			}
			continue
		}

		select {
		case <-c.stop:
			// Shutdown has been called.
			return
		case <-balance:
			c.balance()
		case task, ok := <-c.watch:
			if !ok {
				c.logger.Log(LogLevelDebug, "Watch channel closed. Exiting main loop.")
				return
			}
			if !c.bal.CanClaim(task) {
				c.logger.Log(LogLevelInfo, "Balancer rejected task %s", task)
				break
			}
			if !c.coord.Claim(task) {
				c.logger.Log(LogLevelInfo, "Coordinator unable to claim task %s", task)
				break
			}
			c.claimed(task)
		case cmd, ok := <-cmdChan:
			if !ok {
				c.logger.Log(LogLevelDebug, "Command channel closed. Exiting main loop.")
				return
			}
			c.handleCommand(cmd)
		}
		// Signal that main loop is restarting after handling an event
		select {
		case <-c.stop:
			return
		case c.tick <- 1:
		}
	}
}

func (c *Consumer) watcher() {
	defer close(c.watch)
	c.logger.Log(LogLevelDebug, "Consumer watching")

	for {
		task, err := c.coord.Watch()
		if err != nil {
			//FIXME add more sophisticated error handling
			c.logger.Log(LogLevelError, "Coordinator returned an error during watch, waiting and retrying: %v", err)
			select {
			case <-c.stop:
				return
			case <-time.After(consumerRetryDelay):
			}
			continue
		}
		if task == "" {
			c.logger.Log(LogLevelInfo, "Coordinator has closed, no longer watching for tasks.")
			return
		}
		// Send task to watcher (or shutdown)
		select {
		case <-c.stop:
			return
		case c.watch <- task:
		}
		// Wait for main loop to signal task has been handled
		select {
		case <-c.stop:
			return
		case <-c.tick:
		}
	}
}

func (c *Consumer) balance() {
	for _, task := range c.bal.Balance() {
		//TODO Release tasks asynchronously as their shutdown might be slow?
		c.release(task)
	}
}

// shutdown is the actual shutdown logic called when Run() exits.
func (c *Consumer) shutdown() {
	c.logger.Log(LogLevelDebug, "Closing Coordinator")
	c.coord.Close()

	// Build list of of currently running tasks
	tasks := c.Tasks()
	c.logger.Log(LogLevelInfo, "Sending stop signal to %d handler(s)", len(tasks))

	// Concurrently shutdown handlers as they may take a while to shutdown
	for _, id := range tasks {
		go c.release(id)
	}

	c.logger.Log(LogLevelInfo, "Waiting for handlers to exit")
	c.hwg.Wait()
}

// Shutdown stops the main Run loop, calls Stop on all handlers, and calls
// Close on the Coordinator. Running tasks will be released for other nodes to
// claim.
func (c *Consumer) Shutdown() {
	select {
	case <-c.stop:
		// already stopped
	default:
		c.logger.Log(LogLevelDebug, "Stopping Run loop")
		close(c.stop)
	}
	c.hwg.Wait()
}

// Tasks returns a sorted list of running Task IDs.
func (c *Consumer) Tasks() []string {
	c.runL.Lock()
	defer c.runL.Unlock()
	t := make([]string, len(c.running))
	i := 0
	for id, _ := range c.running {
		t[i] = id
		i++
	}
	sort.Strings(t)
	return t
}

// claimed starts a handler for a claimed task. It is the only method to
// manipulate c.running and closes the runningTask channel when a handler's Run
// method exits.
func (c *Consumer) claimed(taskID string) {
	h := c.handler()

	c.logger.Log(LogLevelDebug, "Attempting to start task "+taskID)
	// Associate handler with taskID
	// **This is the only place tasks should be added to c.running**
	c.runL.Lock()
	c.running[taskID] = runningTask{h: h, c: make(chan struct{})}
	c.runL.Unlock()

	c.hwg.Add(1)
	// Start handler in its own goroutine
	go func() {
		c.logger.Log(LogLevelInfo, "Task started: %s", taskID)
		defer c.logger.Log(LogLevelInfo, "Task exited: %s", taskID)
		defer func() {
			if err := recover(); err != nil {
				stack := make([]byte, 50*1024)
				sz := runtime.Stack(stack, false)
				c.logger.Log(LogLevelError, "Handler %s panic()'d: %v\n%s", taskID, err, stack[:sz])
				// panics are considered fatal errors. Make sure the task isn't
				// rescheduled.
				c.coord.Done(taskID)
			}
			// **This is the only place tasks should be removed from c.running**
			c.runL.Lock()
			close(c.running[taskID].c)
			delete(c.running, taskID)
			c.runL.Unlock()
			c.hwg.Done()
		}()

		// Run the task
		c.logger.Log(LogLevelDebug, "Calling run for task %s", taskID)
		if err := h.Run(taskID); err != nil {
			if ferr, ok := err.(FatalError); ok && ferr.Fatal() {
				c.logger.Log(LogLevelError, "Handler for %s exited with fatal error: %v", taskID, err)
			} else {
				c.logger.Log(LogLevelError, "Handler for %s exited with error: %v", taskID, err)
				// error was non-fatal, release and let another node try
				c.coord.Release(taskID)
				return
			}
		}
		c.coord.Done(taskID)
	}()
}

// release stops and Coordinator.Release()s a task if it's running.
//
// release blocks until the task handler stops running.
func (c *Consumer) release(taskID string) {
	// Stop task...
	if c.stopTask(taskID) {
		// ...instruct the coordinator to release it
		c.coord.Release(taskID)
	}
}

// stopTask returns true if the task was running and stopped successfully.
func (c *Consumer) stopTask(taskID string) bool {
	c.runL.Lock()
	task, ok := c.running[taskID]
	c.runL.Unlock()

	if !ok {
		// This can happen if a task completes during Balance() and is not an error.
		c.logger.Log(LogLevelWarn, "Tried to release a non-running task: %s", taskID)
		return false
	}

	// all handler methods must be wrapped in a recover to prevent a misbehaving
	// handler from crashing the entire consumer
	func() {
		defer func() {
			if err := recover(); err != nil {
				stack := make([]byte, 50*1024)
				sz := runtime.Stack(stack, false)
				c.logger.Log(LogLevelError, "Handler %s panic()'d on Stop: %v\n%s", taskID, err, stack[:sz])
			}
		}()

		task.h.Stop()
	}()

	// Once the handler is stopped...
	//FIXME should there be a timeout here?
	<-task.c
	return true
}

func (c *Consumer) frozen() bool {
	c.freezeL.Lock()
	r := c.freeze
	c.freezeL.Unlock()
	return r
}

func (c *Consumer) handleCommand(cmd Command) {
	switch cmd.Name() {
	case cmdFreeze:
		if c.frozen() {
			c.logger.Log(LogLevelInfo, "Ignoring freeze command: already frozen")
			return
		}
		c.logger.Log(LogLevelInfo, "Freezing")
		c.freezeL.Lock()
		c.freeze = true
		c.freezeL.Unlock()
	case cmdUnfreeze:
		if !c.frozen() {
			c.logger.Log(LogLevelInfo, "Ignoring unfreeze command: not frozen")
			return
		}
		c.logger.Log(LogLevelInfo, "Unfreezing")
		c.freezeL.Lock()
		c.freeze = false
		c.freezeL.Unlock()
	case cmdBalance:
		c.logger.Log(LogLevelInfo, "Balancing due to command")
		c.balance()
		c.logger.Log(LogLevelDebug, "Finished balancing due to command")
	case cmdReleaseTask:
		taskI, ok := cmd.Parameters()["task"]
		task, ok2 := taskI.(string)
		if !ok || !ok2 {
			c.logger.Log(LogLevelError, "Release task command didn't contain a valid task")
			return
		}
		c.logger.Log(LogLevelInfo, "Releasing task %s due to command", task)
		c.release(task)
	case cmdStopTask:
		taskI, ok := cmd.Parameters()["task"]
		task, ok2 := taskI.(string)
		if !ok || !ok2 {
			c.logger.Log(LogLevelError, "Stop task command didn't contain a valid task")
			return
		}
		c.logger.Log(LogLevelInfo, "Stopping task %s due to command", task)
		c.stopTask(task)
	default:
		c.logger.Log(LogLevelWarn, "Discarding unknown command: %s", cmd.Name())
	}
}
<commit_msg>Block Shutdown()s exit on Run()s so coord can close<commit_after>package metafora

import (
	"math/rand"
	"runtime"
	"sort"
	"sync"
	"time"
)

var (

	// balance calls are randomized and this is the upper bound of the random
	// amount
	balanceJitterMax = 10 * int64(time.Second)

	//FIXME should probably be improved, see usage in Run()
	consumerRetryDelay = 10 * time.Second
)

type runningTask struct {
	h Handler
	c chan struct{}
}

// Consumer is the core Metafora task runner.
type Consumer struct {
	// Func to create new handlers
	handler HandlerFunc

	// Map of task:Handler
	running map[string]runningTask

	// Mutex to protect access to running
	runL sync.Mutex

	// WaitGroup for running handlers and consumer goroutines
	hwg sync.WaitGroup

	// WaitGroup so Shutdown() can block on Run() exiting fully
	runwg sync.WaitGroup

	bal      Balancer
	balEvery time.Duration
	coord    Coordinator
	logger   *logger
	stop     chan struct{} // closed by Shutdown to cause Run to exit

	// ticked on each loop of the main loop to enforce sequential interaction
	// with coordinator and balancer
	tick chan int

	watch chan string // channel for watcher to send tasks to main loop

	// Set by command handler, read anywhere via Consumer.frozen()
	freezeL sync.Mutex
	freeze  bool
}

// NewConsumer returns a new consumer and calls Init on the Balancer and Coordinator.
func NewConsumer(coord Coordinator, h HandlerFunc, b Balancer) (*Consumer, error) {
	c := &Consumer{
		running:  make(map[string]runningTask),
		handler:  h,
		bal:      b,
		balEvery: 15 * time.Minute, //TODO make balance wait configurable
		coord:    coord,
		logger:   stdoutLogger(),
		stop:     make(chan struct{}),
		tick:     make(chan int),
		watch:    make(chan string),
	}

	// initialize balancer with the consumer and a prefixed logger
	b.Init(&struct {
		*Consumer
		Logger
	}{Consumer: c, Logger: c.logger})

	if err := coord.Init(&coordinatorContext{Consumer: c, Logger: c.logger}); err != nil {
		return nil, err
	}
	return c, nil
}

// SetLogger assigns the logger to use as well as a level
//
// The logger parameter is an interface that requires the following
// method to be implemented (such as the the stdlib log.Logger):
//
//    Output(calldepth int, s string)
//
func (c *Consumer) SetLogger(l logOutputter, lvl LogLevel) {
	c.logger.l = l
	c.logger.lvl = lvl
}

// Run is the core run loop of Metafora. It is responsible for calling into the
// Coordinator to claim work and Balancer to rebalance work.
//
// Run blocks until Shutdown is called or an internal error occurs.
func (c *Consumer) Run() {
	c.logger.Log(LogLevelDebug, "Starting consumer")

	// Increment run wait group so Shutdown() can block on Run() exiting fully.
	c.runwg.Add(1)
	defer c.runwg.Done()

	// chans for core goroutines to communicate with main loop
	balance := make(chan bool)
	cmdChan := make(chan Command)

	// Balance is called by the main loop when the balance channel is ticked
	go func() {
		randInt := rand.New(rand.NewSource(time.Now().UnixNano())).Int63n
		for {
			select {
			case <-c.stop:
				// Shutdown has been called.
				return
			case <-time.After(c.balEvery + time.Duration(randInt(balanceJitterMax))):
				c.logger.Log(LogLevelInfo, "Balancing")
				select {
				case balance <- true:
					// Ticked balance
				case <-c.stop:
					// Shutdown has been called.
					return
				}
			}
			// Wait for main loop to signal balancing is done
			select {
			case <-c.stop:
				return
			case <-c.tick:
			}
		}
	}()

	// Watch for new tasks in a goroutine
	go c.watcher()

	// Watch for new commands in a goroutine
	go func() {
		defer close(cmdChan)
		for {
			cmd, err := c.coord.Command()
			if err != nil {
				//FIXME add more sophisticated error handling
				c.logger.Log(LogLevelError, "Coordinator returned an error during command, waiting and retrying. %v", err)
				select {
				case <-c.stop:
					return
				case <-time.After(consumerRetryDelay):
				}
				continue
			}
			if cmd == nil {
				c.logger.Log(LogLevelDebug, "Command coordinator exited")
				return
			}
			// Send command to watcher (or shutdown)
			select {
			case <-c.stop:
				return
			case cmdChan <- cmd:
			}
			// Wait for main loop to signal command has been handled
			select {
			case <-c.stop:
				return
			case <-c.tick:
			}
		}
	}()

	// Make sure Run() cleans up on exit (stops coordinator, releases tasks, etc)
	defer c.shutdown()

	// Main Loop ensures events are processed synchronously
	for {
		if c.frozen() {
			// Only recv commands while frozen
			select {
			case <-c.stop:
				// Shutdown has been called.
				return
			case cmd, ok := <-cmdChan:
				if !ok {
					c.logger.Log(LogLevelDebug, "Command channel closed. Exiting main loop.")
					return
				}
				c.logger.Log(LogLevelDebug, "Received command: %s", cmd)
				c.handleCommand(cmd)
			}
			// Must send tick whenever main loop restarts
			select {
			case <-c.stop:
				return
			case c.tick <- 1:
			}
			continue
		}

		select {
		case <-c.stop:
			// Shutdown has been called.
			return
		case <-balance:
			c.balance()
		case task, ok := <-c.watch:
			if !ok {
				c.logger.Log(LogLevelDebug, "Watch channel closed. Exiting main loop.")
				return
			}
			if !c.bal.CanClaim(task) {
				c.logger.Log(LogLevelInfo, "Balancer rejected task %s", task)
				break
			}
			if !c.coord.Claim(task) {
				c.logger.Log(LogLevelInfo, "Coordinator unable to claim task %s", task)
				break
			}
			c.claimed(task)
		case cmd, ok := <-cmdChan:
			if !ok {
				c.logger.Log(LogLevelDebug, "Command channel closed. Exiting main loop.")
				return
			}
			c.handleCommand(cmd)
		}
		// Signal that main loop is restarting after handling an event
		select {
		case <-c.stop:
			return
		case c.tick <- 1:
		}
	}
}

func (c *Consumer) watcher() {
	defer close(c.watch)
	c.logger.Log(LogLevelDebug, "Consumer watching")

	for {
		task, err := c.coord.Watch()
		if err != nil {
			//FIXME add more sophisticated error handling
			c.logger.Log(LogLevelError, "Coordinator returned an error during watch, waiting and retrying: %v", err)
			select {
			case <-c.stop:
				return
			case <-time.After(consumerRetryDelay):
			}
			continue
		}
		if task == "" {
			c.logger.Log(LogLevelInfo, "Coordinator has closed, no longer watching for tasks.")
			return
		}
		// Send task to watcher (or shutdown)
		select {
		case <-c.stop:
			return
		case c.watch <- task:
		}
		// Wait for main loop to signal task has been handled
		select {
		case <-c.stop:
			return
		case <-c.tick:
		}
	}
}

func (c *Consumer) balance() {
	for _, task := range c.bal.Balance() {
		//TODO Release tasks asynchronously as their shutdown might be slow?
		c.release(task)
	}
}

// shutdown is the actual shutdown logic called when Run() exits.
func (c *Consumer) shutdown() {
	c.logger.Log(LogLevelDebug, "Closing Coordinator")
	c.coord.Close()

	// Build list of of currently running tasks
	tasks := c.Tasks()
	c.logger.Log(LogLevelInfo, "Sending stop signal to %d handler(s)", len(tasks))

	// Concurrently shutdown handlers as they may take a while to shutdown
	for _, id := range tasks {
		go c.release(id)
	}

	c.logger.Log(LogLevelInfo, "Waiting for handlers to exit")
	c.hwg.Wait()
}

// Shutdown stops the main Run loop, calls Stop on all handlers, and calls
// Close on the Coordinator. Running tasks will be released for other nodes to
// claim.
func (c *Consumer) Shutdown() {
	select {
	case <-c.stop:
		// already stopped
	default:
		c.logger.Log(LogLevelDebug, "Stopping Run loop")
		close(c.stop)
	}

	// Wait for task handlers to exit.
	c.hwg.Wait()

	// Make sure Run() exits, otherwise coord.Close() might not finish before
	// exiting.
	c.runwg.Wait()
}

// Tasks returns a sorted list of running Task IDs.
func (c *Consumer) Tasks() []string {
	c.runL.Lock()
	defer c.runL.Unlock()
	t := make([]string, len(c.running))
	i := 0
	for id, _ := range c.running {
		t[i] = id
		i++
	}
	sort.Strings(t)
	return t
}

// claimed starts a handler for a claimed task. It is the only method to
// manipulate c.running and closes the runningTask channel when a handler's Run
// method exits.
func (c *Consumer) claimed(taskID string) {
	h := c.handler()

	c.logger.Log(LogLevelDebug, "Attempting to start task "+taskID)
	// Associate handler with taskID
	// **This is the only place tasks should be added to c.running**
	c.runL.Lock()
	c.running[taskID] = runningTask{h: h, c: make(chan struct{})}
	c.runL.Unlock()

	c.hwg.Add(1)
	// Start handler in its own goroutine
	go func() {
		c.logger.Log(LogLevelInfo, "Task started: %s", taskID)
		defer c.logger.Log(LogLevelInfo, "Task exited: %s", taskID)
		defer func() {
			if err := recover(); err != nil {
				stack := make([]byte, 50*1024)
				sz := runtime.Stack(stack, false)
				c.logger.Log(LogLevelError, "Handler %s panic()'d: %v\n%s", taskID, err, stack[:sz])
				// panics are considered fatal errors. Make sure the task isn't
				// rescheduled.
				c.coord.Done(taskID)
			}
			// **This is the only place tasks should be removed from c.running**
			c.runL.Lock()
			close(c.running[taskID].c)
			delete(c.running, taskID)
			c.runL.Unlock()
			c.hwg.Done()
		}()

		// Run the task
		c.logger.Log(LogLevelDebug, "Calling run for task %s", taskID)
		if err := h.Run(taskID); err != nil {
			if ferr, ok := err.(FatalError); ok && ferr.Fatal() {
				c.logger.Log(LogLevelError, "Handler for %s exited with fatal error: %v", taskID, err)
			} else {
				c.logger.Log(LogLevelError, "Handler for %s exited with error: %v", taskID, err)
				// error was non-fatal, release and let another node try
				c.coord.Release(taskID)
				return
			}
		}
		c.coord.Done(taskID)
	}()
}

// release stops and Coordinator.Release()s a task if it's running.
//
// release blocks until the task handler stops running.
func (c *Consumer) release(taskID string) {
	// Stop task...
	if c.stopTask(taskID) {
		// ...instruct the coordinator to release it
		c.coord.Release(taskID)
	}
}

// stopTask returns true if the task was running and stopped successfully.
func (c *Consumer) stopTask(taskID string) bool {
	c.runL.Lock()
	task, ok := c.running[taskID]
	c.runL.Unlock()

	if !ok {
		// This can happen if a task completes during Balance() and is not an error.
		c.logger.Log(LogLevelWarn, "Tried to release a non-running task: %s", taskID)
		return false
	}

	// all handler methods must be wrapped in a recover to prevent a misbehaving
	// handler from crashing the entire consumer
	func() {
		defer func() {
			if err := recover(); err != nil {
				stack := make([]byte, 50*1024)
				sz := runtime.Stack(stack, false)
				c.logger.Log(LogLevelError, "Handler %s panic()'d on Stop: %v\n%s", taskID, err, stack[:sz])
			}
		}()

		task.h.Stop()
	}()

	// Once the handler is stopped...
	//FIXME should there be a timeout here?
	<-task.c
	return true
}

func (c *Consumer) frozen() bool {
	c.freezeL.Lock()
	r := c.freeze
	c.freezeL.Unlock()
	return r
}

func (c *Consumer) handleCommand(cmd Command) {
	switch cmd.Name() {
	case cmdFreeze:
		if c.frozen() {
			c.logger.Log(LogLevelInfo, "Ignoring freeze command: already frozen")
			return
		}
		c.logger.Log(LogLevelInfo, "Freezing")
		c.freezeL.Lock()
		c.freeze = true
		c.freezeL.Unlock()
	case cmdUnfreeze:
		if !c.frozen() {
			c.logger.Log(LogLevelInfo, "Ignoring unfreeze command: not frozen")
			return
		}
		c.logger.Log(LogLevelInfo, "Unfreezing")
		c.freezeL.Lock()
		c.freeze = false
		c.freezeL.Unlock()
	case cmdBalance:
		c.logger.Log(LogLevelInfo, "Balancing due to command")
		c.balance()
		c.logger.Log(LogLevelDebug, "Finished balancing due to command")
	case cmdReleaseTask:
		taskI, ok := cmd.Parameters()["task"]
		task, ok2 := taskI.(string)
		if !ok || !ok2 {
			c.logger.Log(LogLevelError, "Release task command didn't contain a valid task")
			return
		}
		c.logger.Log(LogLevelInfo, "Releasing task %s due to command", task)
		c.release(task)
	case cmdStopTask:
		taskI, ok := cmd.Parameters()["task"]
		task, ok2 := taskI.(string)
		if !ok || !ok2 {
			c.logger.Log(LogLevelError, "Stop task command didn't contain a valid task")
			return
		}
		c.logger.Log(LogLevelInfo, "Stopping task %s due to command", task)
		c.stopTask(task)
	default:
		c.logger.Log(LogLevelWarn, "Discarding unknown command: %s", cmd.Name())
	}
}
<|endoftext|>