<commit_before>package aws

import (
	"bytes"
	"fmt"
	"log"
	"regexp"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ecs"
	"github.com/aws/aws-sdk-go/service/iam"
	"github.com/hashicorp/terraform/helper/hashcode"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
)

var taskDefinitionRE = regexp.MustCompile("^([a-zA-Z0-9_-]+):([0-9]+)$")

func resourceAwsEcsService() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsEcsServiceCreate,
		Read:   resourceAwsEcsServiceRead,
		Update: resourceAwsEcsServiceUpdate,
		Delete: resourceAwsEcsServiceDelete,

		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"cluster": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
				ForceNew: true,
			},

			"task_definition": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},

			"desired_count": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
			},

			"iam_role": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},

			"load_balancer": &schema.Schema{
				Type:     schema.TypeSet,
				Optional: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"elb_name": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},

						"container_name": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},

						"container_port": &schema.Schema{
							Type:     schema.TypeInt,
							Required: true,
						},
					},
				},
				Set: resourceAwsEcsLoadBalancerHash,
			},
		},
	}
}

func resourceAwsEcsServiceCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).ecsconn

	input := ecs.CreateServiceInput{
		ServiceName:    aws.String(d.Get("name").(string)),
		TaskDefinition: aws.String(d.Get("task_definition").(string)),
		DesiredCount:   aws.Long(int64(d.Get("desired_count").(int))),
	}

	if v, ok := d.GetOk("cluster"); ok {
		input.Cluster = aws.String(v.(string))
	}

	loadBalancers := expandEcsLoadBalancers(d.Get("load_balancer").(*schema.Set).List())
	if len(loadBalancers) > 0 {
		log.Printf("[DEBUG] Adding ECS load balancers: %#v", loadBalancers)
		input.LoadBalancers = loadBalancers
	}
	if v, ok := d.GetOk("iam_role"); ok {
		input.Role = aws.String(v.(string))
	}

	log.Printf("[DEBUG] Creating ECS service: %#v", input)
	out, err := conn.CreateService(&input)
	if err != nil {
		return err
	}

	service := *out.Service

	log.Printf("[DEBUG] ECS service created: %s", *service.ServiceARN)
	d.SetId(*service.ServiceARN)
	d.Set("cluster", *service.ClusterARN)

	return resourceAwsEcsServiceUpdate(d, meta)
}

func resourceAwsEcsServiceRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).ecsconn

	log.Printf("[DEBUG] Reading ECS service %s", d.Id())
	input := ecs.DescribeServicesInput{
		Services: []*string{aws.String(d.Id())},
		Cluster:  aws.String(d.Get("cluster").(string)),
	}

	out, err := conn.DescribeServices(&input)
	if err != nil {
		return err
	}

	if len(out.Services) < 1 {
		return nil
	}

	service := out.Services[0]
	log.Printf("[DEBUG] Received ECS service %#v", service)

	d.SetId(*service.ServiceARN)
	d.Set("name", *service.ServiceName)

	// Save task definition in the same format
	if strings.HasPrefix(d.Get("task_definition").(string), "arn:aws:ecs:") {
		d.Set("task_definition", *service.TaskDefinition)
	} else {
		taskDefinition := buildFamilyAndRevisionFromARN(*service.TaskDefinition)
		d.Set("task_definition", taskDefinition)
	}

	d.Set("desired_count", *service.DesiredCount)
	d.Set("cluster", *service.ClusterARN)

	if service.RoleARN != nil {
		d.Set("iam_role", *service.RoleARN)
	}

	if service.LoadBalancers != nil {
		d.Set("load_balancers", flattenEcsLoadBalancers(service.LoadBalancers))
	}

	return nil
}

func resourceAwsEcsServiceUpdate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).ecsconn

	log.Printf("[DEBUG] Updating ECS service %s", d.Id())
	input := ecs.UpdateServiceInput{
		Service: aws.String(d.Id()),
		Cluster: aws.String(d.Get("cluster").(string)),
	}

	if d.HasChange("desired_count") {
		_, n := d.GetChange("desired_count")
		input.DesiredCount = aws.Long(int64(n.(int)))
	}
	if d.HasChange("task_definition") {
		_, n := d.GetChange("task_definition")
		input.TaskDefinition = aws.String(n.(string))
	}

	out, err := conn.UpdateService(&input)
	if err != nil {
		return err
	}
	service := out.Service
	log.Printf("[DEBUG] Updated ECS service %#v", service)

	return resourceAwsEcsServiceRead(d, meta)
}

func resourceAwsEcsServiceDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).ecsconn

	// Check if it's not already gone
	resp, err := conn.DescribeServices(&ecs.DescribeServicesInput{
		Services: []*string{aws.String(d.Id())},
		Cluster:  aws.String(d.Get("cluster").(string)),
	})
	if err != nil {
		return err
	}
	log.Printf("[DEBUG] ECS service %s is currently %s", d.Id(), *resp.Services[0].Status)

	if *resp.Services[0].Status == "INACTIVE" {
		return nil
	}

	// Drain the ECS service
	if *resp.Services[0].Status != "DRAINING" {
		log.Printf("[DEBUG] Draining ECS service %s", d.Id())
		_, err = conn.UpdateService(&ecs.UpdateServiceInput{
			Service:      aws.String(d.Id()),
			Cluster:      aws.String(d.Get("cluster").(string)),
			DesiredCount: aws.Long(int64(0)),
		})
		if err != nil {
			return err
		}
	}

	input := ecs.DeleteServiceInput{
		Service: aws.String(d.Id()),
		Cluster: aws.String(d.Get("cluster").(string)),
	}

	log.Printf("[DEBUG] Deleting ECS service %#v", input)
	out, err := conn.DeleteService(&input)
	if err != nil {
		return err
	}

	// Wait until it's deleted
	wait := resource.StateChangeConf{
		Pending:    []string{"DRAINING"},
		Target:     "INACTIVE",
		Timeout:    5 * time.Minute,
		MinTimeout: 1 * time.Second,
		Refresh: func() (interface{}, string, error) {
			log.Printf("[DEBUG] Checking if ECS service %s is INACTIVE", d.Id())
			resp, err := conn.DescribeServices(&ecs.DescribeServicesInput{
				Services: []*string{aws.String(d.Id())},
				Cluster:  aws.String(d.Get("cluster").(string)),
			})
			if err != nil {
				return resp, "FAILED", err
			}

			return resp, *resp.Services[0].Status, nil
		},
	}

	_, err = wait.WaitForState()
	if err != nil {
		return err
	}

	log.Printf("[DEBUG] ECS service %s deleted.", *out.Service.ServiceARN)
	return nil
}

func resourceAwsEcsLoadBalancerHash(v interface{}) int {
	var buf bytes.Buffer
	m := v.(map[string]interface{})
	buf.WriteString(fmt.Sprintf("%s-", m["elb_name"].(string)))
	buf.WriteString(fmt.Sprintf("%s-", m["container_name"].(string)))
	buf.WriteString(fmt.Sprintf("%d-", m["container_port"].(int)))

	return hashcode.String(buf.String())
}

func buildFamilyAndRevisionFromARN(arn string) string {
	return strings.Split(arn, "/")[1]
}

func buildTaskDefinitionARN(taskDefinition string, meta interface{}) (string, error) {
	// If it's already an ARN, just return it
	if strings.HasPrefix(taskDefinition, "arn:aws:ecs:") {
		return taskDefinition, nil
	}

	// Parse out family & revision
	family, revision, err := parseTaskDefinition(taskDefinition)
	if err != nil {
		return "", err
	}

	iamconn := meta.(*AWSClient).iamconn
	region := meta.(*AWSClient).region

	// An zero value GetUserInput{} defers to the currently logged in user
	resp, err := iamconn.GetUser(&iam.GetUserInput{})
	if err != nil {
		return "", fmt.Errorf("GetUser ERROR: %#v", err)
	}

	// arn:aws:iam::0123456789:user/username
	userARN := *resp.User.ARN
	accountID := strings.Split(userARN, ":")[4]

	// arn:aws:ecs:us-west-2:01234567890:task-definition/mongodb:3
	arn := fmt.Sprintf("arn:aws:ecs:%s:%s:task-definition/%s:%s",
		region, accountID, family, revision)
	log.Printf("[DEBUG] Built task definition ARN: %s", arn)
	return arn, nil
}

func parseTaskDefinition(taskDefinition string) (string, string, error) {
	matches := taskDefinitionRE.FindAllStringSubmatch(taskDefinition, 2)

	if len(matches) == 0 || len(matches[0]) != 3 {
		return "", "", fmt.Errorf(
			"Invalid task definition format, family:rev or ARN expected (%#v)",
			taskDefinition)
	}

	return matches[0][1], matches[0][2], nil
}
<commit_msg>aws: Use ClientToken when creating ecs_service<commit_after>package aws

import (
	"bytes"
	"fmt"
	"log"
	"regexp"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ecs"
	"github.com/aws/aws-sdk-go/service/iam"
	"github.com/hashicorp/terraform/helper/hashcode"
	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/helper/schema"
)

var taskDefinitionRE = regexp.MustCompile("^([a-zA-Z0-9_-]+):([0-9]+)$")

func resourceAwsEcsService() *schema.Resource {
	return &schema.Resource{
		Create: resourceAwsEcsServiceCreate,
		Read:   resourceAwsEcsServiceRead,
		Update: resourceAwsEcsServiceUpdate,
		Delete: resourceAwsEcsServiceDelete,

		Schema: map[string]*schema.Schema{
			"name": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
				ForceNew: true,
			},

			"cluster": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
				Computed: true,
				ForceNew: true,
			},

			"task_definition": &schema.Schema{
				Type:     schema.TypeString,
				Required: true,
			},

			"desired_count": &schema.Schema{
				Type:     schema.TypeInt,
				Optional: true,
			},

			"iam_role": &schema.Schema{
				Type:     schema.TypeString,
				Optional: true,
			},

			"load_balancer": &schema.Schema{
				Type:     schema.TypeSet,
				Optional: true,
				Elem: &schema.Resource{
					Schema: map[string]*schema.Schema{
						"elb_name": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},

						"container_name": &schema.Schema{
							Type:     schema.TypeString,
							Required: true,
						},

						"container_port": &schema.Schema{
							Type:     schema.TypeInt,
							Required: true,
						},
					},
				},
				Set: resourceAwsEcsLoadBalancerHash,
			},
		},
	}
}

func resourceAwsEcsServiceCreate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).ecsconn

	input := ecs.CreateServiceInput{
		ServiceName:    aws.String(d.Get("name").(string)),
		TaskDefinition: aws.String(d.Get("task_definition").(string)),
		DesiredCount:   aws.Long(int64(d.Get("desired_count").(int))),
		ClientToken:    aws.String(resource.UniqueId()),
	}

	if v, ok := d.GetOk("cluster"); ok {
		input.Cluster = aws.String(v.(string))
	}

	loadBalancers := expandEcsLoadBalancers(d.Get("load_balancer").(*schema.Set).List())
	if len(loadBalancers) > 0 {
		log.Printf("[DEBUG] Adding ECS load balancers: %#v", loadBalancers)
		input.LoadBalancers = loadBalancers
	}
	if v, ok := d.GetOk("iam_role"); ok {
		input.Role = aws.String(v.(string))
	}

	log.Printf("[DEBUG] Creating ECS service: %#v", input)
	out, err := conn.CreateService(&input)
	if err != nil {
		return err
	}

	service := *out.Service

	log.Printf("[DEBUG] ECS service created: %s", *service.ServiceARN)
	d.SetId(*service.ServiceARN)
	d.Set("cluster", *service.ClusterARN)

	return resourceAwsEcsServiceUpdate(d, meta)
}

func resourceAwsEcsServiceRead(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).ecsconn

	log.Printf("[DEBUG] Reading ECS service %s", d.Id())
	input := ecs.DescribeServicesInput{
		Services: []*string{aws.String(d.Id())},
		Cluster:  aws.String(d.Get("cluster").(string)),
	}

	out, err := conn.DescribeServices(&input)
	if err != nil {
		return err
	}

	if len(out.Services) < 1 {
		return nil
	}

	service := out.Services[0]
	log.Printf("[DEBUG] Received ECS service %#v", service)

	d.SetId(*service.ServiceARN)
	d.Set("name", *service.ServiceName)

	// Save task definition in the same format
	if strings.HasPrefix(d.Get("task_definition").(string), "arn:aws:ecs:") {
		d.Set("task_definition", *service.TaskDefinition)
	} else {
		taskDefinition := buildFamilyAndRevisionFromARN(*service.TaskDefinition)
		d.Set("task_definition", taskDefinition)
	}

	d.Set("desired_count", *service.DesiredCount)
	d.Set("cluster", *service.ClusterARN)

	if service.RoleARN != nil {
		d.Set("iam_role", *service.RoleARN)
	}

	if service.LoadBalancers != nil {
		d.Set("load_balancers", flattenEcsLoadBalancers(service.LoadBalancers))
	}

	return nil
}

func resourceAwsEcsServiceUpdate(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).ecsconn

	log.Printf("[DEBUG] Updating ECS service %s", d.Id())
	input := ecs.UpdateServiceInput{
		Service: aws.String(d.Id()),
		Cluster: aws.String(d.Get("cluster").(string)),
	}

	if d.HasChange("desired_count") {
		_, n := d.GetChange("desired_count")
		input.DesiredCount = aws.Long(int64(n.(int)))
	}
	if d.HasChange("task_definition") {
		_, n := d.GetChange("task_definition")
		input.TaskDefinition = aws.String(n.(string))
	}

	out, err := conn.UpdateService(&input)
	if err != nil {
		return err
	}
	service := out.Service
	log.Printf("[DEBUG] Updated ECS service %#v", service)

	return resourceAwsEcsServiceRead(d, meta)
}

func resourceAwsEcsServiceDelete(d *schema.ResourceData, meta interface{}) error {
	conn := meta.(*AWSClient).ecsconn

	// Check if it's not already gone
	resp, err := conn.DescribeServices(&ecs.DescribeServicesInput{
		Services: []*string{aws.String(d.Id())},
		Cluster:  aws.String(d.Get("cluster").(string)),
	})
	if err != nil {
		return err
	}
	log.Printf("[DEBUG] ECS service %s is currently %s", d.Id(), *resp.Services[0].Status)

	if *resp.Services[0].Status == "INACTIVE" {
		return nil
	}

	// Drain the ECS service
	if *resp.Services[0].Status != "DRAINING" {
		log.Printf("[DEBUG] Draining ECS service %s", d.Id())
		_, err = conn.UpdateService(&ecs.UpdateServiceInput{
			Service:      aws.String(d.Id()),
			Cluster:      aws.String(d.Get("cluster").(string)),
			DesiredCount: aws.Long(int64(0)),
		})
		if err != nil {
			return err
		}
	}

	input := ecs.DeleteServiceInput{
		Service: aws.String(d.Id()),
		Cluster: aws.String(d.Get("cluster").(string)),
	}

	log.Printf("[DEBUG] Deleting ECS service %#v", input)
	out, err := conn.DeleteService(&input)
	if err != nil {
		return err
	}

	// Wait until it's deleted
	wait := resource.StateChangeConf{
		Pending:    []string{"DRAINING"},
		Target:     "INACTIVE",
		Timeout:    5 * time.Minute,
		MinTimeout: 1 * time.Second,
		Refresh: func() (interface{}, string, error) {
			log.Printf("[DEBUG] Checking if ECS service %s is INACTIVE", d.Id())
			resp, err := conn.DescribeServices(&ecs.DescribeServicesInput{
				Services: []*string{aws.String(d.Id())},
				Cluster:  aws.String(d.Get("cluster").(string)),
			})
			if err != nil {
				return resp, "FAILED", err
			}

			return resp, *resp.Services[0].Status, nil
		},
	}

	_, err = wait.WaitForState()
	if err != nil {
		return err
	}

	log.Printf("[DEBUG] ECS service %s deleted.", *out.Service.ServiceARN)
	return nil
}

func resourceAwsEcsLoadBalancerHash(v interface{}) int {
	var buf bytes.Buffer
	m := v.(map[string]interface{})
	buf.WriteString(fmt.Sprintf("%s-", m["elb_name"].(string)))
	buf.WriteString(fmt.Sprintf("%s-", m["container_name"].(string)))
	buf.WriteString(fmt.Sprintf("%d-", m["container_port"].(int)))

	return hashcode.String(buf.String())
}

func buildFamilyAndRevisionFromARN(arn string) string {
	return strings.Split(arn, "/")[1]
}

func buildTaskDefinitionARN(taskDefinition string, meta interface{}) (string, error) {
	// If it's already an ARN, just return it
	if strings.HasPrefix(taskDefinition, "arn:aws:ecs:") {
		return taskDefinition, nil
	}

	// Parse out family & revision
	family, revision, err := parseTaskDefinition(taskDefinition)
	if err != nil {
		return "", err
	}

	iamconn := meta.(*AWSClient).iamconn
	region := meta.(*AWSClient).region

	// An zero value GetUserInput{} defers to the currently logged in user
	resp, err := iamconn.GetUser(&iam.GetUserInput{})
	if err != nil {
		return "", fmt.Errorf("GetUser ERROR: %#v", err)
	}

	// arn:aws:iam::0123456789:user/username
	userARN := *resp.User.ARN
	accountID := strings.Split(userARN, ":")[4]

	// arn:aws:ecs:us-west-2:01234567890:task-definition/mongodb:3
	arn := fmt.Sprintf("arn:aws:ecs:%s:%s:task-definition/%s:%s",
		region, accountID, family, revision)
	log.Printf("[DEBUG] Built task definition ARN: %s", arn)
	return arn, nil
}

func parseTaskDefinition(taskDefinition string) (string, string, error) {
	matches := taskDefinitionRE.FindAllStringSubmatch(taskDefinition, 2)

	if len(matches) == 0 || len(matches[0]) != 3 {
		return "", "", fmt.Errorf(
			"Invalid task definition format, family:rev or ARN expected (%#v)",
			taskDefinition)
	}

	return matches[0][1], matches[0][2], nil
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"net/http"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/sirupsen/logrus"
	"k8s.io/apimachinery/pkg/util/sets"

	"k8s.io/test-infra/boskos/client"
	"k8s.io/test-infra/boskos/common"
	"k8s.io/test-infra/prow/config"
	"k8s.io/test-infra/prow/logrusutil"
	"k8s.io/test-infra/prow/metrics"
)

var (
	resourceMetric = prometheus.NewGaugeVec(prometheus.GaugeOpts{
		Name: "boskos_resources",
		Help: "Number of resources recorded in Boskos",
	}, []string{"type", "state"})
	boskosURL         string
	username          string
	passwordFile      string
	resources, states common.CommaSeparatedStrings
	defaultStates     = []string{
		common.Busy,
		common.Cleaning,
		common.Dirty,
		common.Free,
		common.Leased,
		common.ToBeDeleted,
		common.Tombstone,
	}
)

func init() {
	flag.StringVar(&boskosURL, "boskos-url", "http://boskos", "Boskos Server URL")
	flag.StringVar(&username, "username", "", "Username used to access the Boskos server")
	flag.StringVar(&passwordFile, "password-file", "", "The path to password file used to access the Boskos server")
	flag.Var(&resources, "resource-type", "comma-separated list of resources need to have metrics collected.")
	flag.Var(&states, "resource-state", "comma-separated list of states need to have metrics collected.")
	prometheus.MustRegister(resourceMetric)
}

func main() {
	logrusutil.ComponentInit("boskos-metrics")
	boskos, err := client.NewClient("Metrics", boskosURL, username, passwordFile)
	if err != nil {
		logrus.WithError(err).Fatal("unable to create a Boskos client")
	}
	logrus.Infof("Initialzied boskos client!")

	flag.Parse()
	if states == nil {
		states = defaultStates
	}

	metrics.ExposeMetrics("boskos", config.PushGateway{})

	go func() {
		logTick := time.NewTicker(30 * time.Second).C
		for range logTick {
			if err := update(boskos); err != nil {
				logrus.WithError(err).Warning("Update failed!")
			}
		}
	}()

	logrus.Info("Start Service")
	metricsMux := http.NewServeMux()
	metricsMux.Handle("/", handleMetric(boskos))
	logrus.WithError(http.ListenAndServe(":8080", metricsMux)).Fatal("ListenAndServe returned.")
}

func update(boskos *client.Client) error {
	// initialize resources counted by type, then state
	resourcesByState := map[string]map[string]float64{}
	for _, resource := range resources {
		resourcesByState[resource] = map[string]float64{}
		for _, state := range states {
			resourcesByState[resource][state] = 0
		}
	}

	// record current states
	knownStates := sets.NewString(states...)
	for _, resource := range resources {
		metric, err := boskos.Metric(resource)
		if err != nil {
			return fmt.Errorf("fail to get metric for %s : %v", resource, err)
		}
		// Filtering metrics states
		for state, value := range metric.Current {
			if !knownStates.Has(state) {
				state = common.Other
			}
			resourcesByState[resource][state] = float64(value)
		}
	}

	// expose current states
	for resource, states := range resourcesByState {
		for state, amount := range states {
			resourceMetric.WithLabelValues(resource, state).Set(amount)
		}
	}
	return nil
}

//  handleMetric: Handler for /
//  Method: GET
func handleMetric(boskos *client.Client) http.HandlerFunc {
	return func(res http.ResponseWriter, req *http.Request) {
		log := logrus.WithField("handler", "handleMetric")
		log.Infof("From %v", req.RemoteAddr)

		if req.Method != "GET" {
			log.Warningf("[BadRequest]method %v, expect GET", req.Method)
			http.Error(res, "only accepts GET request", http.StatusMethodNotAllowed)
			return
		}

		rtype := req.URL.Query().Get("type")
		if rtype == "" {
			msg := "type must be set in the request."
			log.Warning(msg)
			http.Error(res, msg, http.StatusBadRequest)
			return
		}

		log.Infof("Request for metric %v", rtype)

		metric, err := boskos.Metric(rtype)
		if err != nil {
			log.WithError(err).Errorf("Fail to get metic for %v", rtype)
			http.Error(res, err.Error(), http.StatusNotFound)
			return
		}

		metricJSON, err := json.Marshal(metric)
		if err != nil {
			log.WithError(err).Errorf("json.Marshal failed: %v", metricJSON)
			http.Error(res, err.Error(), http.StatusInternalServerError)
			return
		}
		log.Infof("Metric query for %v: %v", rtype, string(metricJSON))
		fmt.Fprint(res, string(metricJSON))
	}
}
<commit_msg>boskos/metrics: don't completely bail out if one of the resource types errors<commit_after>/*
Copyright 2017 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"net/http"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/sirupsen/logrus"
	"k8s.io/apimachinery/pkg/util/sets"

	"k8s.io/test-infra/boskos/client"
	"k8s.io/test-infra/boskos/common"
	"k8s.io/test-infra/prow/config"
	"k8s.io/test-infra/prow/logrusutil"
	"k8s.io/test-infra/prow/metrics"
)

var (
	resourceMetric = prometheus.NewGaugeVec(prometheus.GaugeOpts{
		Name: "boskos_resources",
		Help: "Number of resources recorded in Boskos",
	}, []string{"type", "state"})
	boskosURL         string
	username          string
	passwordFile      string
	resources, states common.CommaSeparatedStrings
	defaultStates     = []string{
		common.Busy,
		common.Cleaning,
		common.Dirty,
		common.Free,
		common.Leased,
		common.ToBeDeleted,
		common.Tombstone,
	}
)

func init() {
	flag.StringVar(&boskosURL, "boskos-url", "http://boskos", "Boskos Server URL")
	flag.StringVar(&username, "username", "", "Username used to access the Boskos server")
	flag.StringVar(&passwordFile, "password-file", "", "The path to password file used to access the Boskos server")
	flag.Var(&resources, "resource-type", "comma-separated list of resources need to have metrics collected.")
	flag.Var(&states, "resource-state", "comma-separated list of states need to have metrics collected.")
	prometheus.MustRegister(resourceMetric)
}

func main() {
	logrusutil.ComponentInit("boskos-metrics")
	boskos, err := client.NewClient("Metrics", boskosURL, username, passwordFile)
	if err != nil {
		logrus.WithError(err).Fatal("unable to create a Boskos client")
	}
	logrus.Infof("Initialzied boskos client!")

	flag.Parse()
	if states == nil {
		states = defaultStates
	}

	metrics.ExposeMetrics("boskos", config.PushGateway{})

	go func() {
		logTick := time.NewTicker(30 * time.Second).C
		for range logTick {
			if err := update(boskos); err != nil {
				logrus.WithError(err).Warning("Update failed!")
			}
		}
	}()

	logrus.Info("Start Service")
	metricsMux := http.NewServeMux()
	metricsMux.Handle("/", handleMetric(boskos))
	logrus.WithError(http.ListenAndServe(":8080", metricsMux)).Fatal("ListenAndServe returned.")
}

func update(boskos *client.Client) error {
	// initialize resources counted by type, then state
	resourcesByState := map[string]map[string]float64{}
	for _, resource := range resources {
		resourcesByState[resource] = map[string]float64{}
		for _, state := range states {
			resourcesByState[resource][state] = 0
		}
	}

	// record current states
	knownStates := sets.NewString(states...)
	for _, resource := range resources {
		metric, err := boskos.Metric(resource)
		if err != nil {
			logrus.WithError(err).Errorf("failed to get metric for %s", resource)
			continue
		}
		// Filtering metrics states
		for state, value := range metric.Current {
			if !knownStates.Has(state) {
				state = common.Other
			}
			resourcesByState[resource][state] = float64(value)
		}
	}

	// expose current states
	for resource, states := range resourcesByState {
		for state, amount := range states {
			resourceMetric.WithLabelValues(resource, state).Set(amount)
		}
	}
	return nil
}

//  handleMetric: Handler for /
//  Method: GET
func handleMetric(boskos *client.Client) http.HandlerFunc {
	return func(res http.ResponseWriter, req *http.Request) {
		log := logrus.WithField("handler", "handleMetric")
		log.Infof("From %v", req.RemoteAddr)

		if req.Method != "GET" {
			log.Warningf("[BadRequest]method %v, expect GET", req.Method)
			http.Error(res, "only accepts GET request", http.StatusMethodNotAllowed)
			return
		}

		rtype := req.URL.Query().Get("type")
		if rtype == "" {
			msg := "type must be set in the request."
			log.Warning(msg)
			http.Error(res, msg, http.StatusBadRequest)
			return
		}

		log.Infof("Request for metric %v", rtype)

		metric, err := boskos.Metric(rtype)
		if err != nil {
			log.WithError(err).Errorf("Fail to get metic for %v", rtype)
			http.Error(res, err.Error(), http.StatusNotFound)
			return
		}

		metricJSON, err := json.Marshal(metric)
		if err != nil {
			log.WithError(err).Errorf("json.Marshal failed: %v", metricJSON)
			http.Error(res, err.Error(), http.StatusInternalServerError)
			return
		}
		log.Infof("Metric query for %v: %v", rtype, string(metricJSON))
		fmt.Fprint(res, string(metricJSON))
	}
}
<|endoftext|>


<commit_before>package acceptance_test

import (
	"cf-pusher/cf_cli_adapter"
	"crypto/tls"
	"fmt"
	"lib/policy_client"
	"math/rand"
	"net/http"
	"policy-server/api/api_v0"
	"strings"
	"time"

	"code.cloudfoundry.org/lager/lagertest"
	"github.com/cloudfoundry-incubator/cf-test-helpers/cf"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/ginkgo/extensions/table"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gexec"
	"github.com/pivotal-cf-experimental/warrant"
)

var _ = Describe("space developer policy configuration", func() {
	var (
		appA       string
		appB       string
		orgName    string
		spaceNameA string
		spaceNameB string
		prefix     string

		policyClient *policy_client.ExternalClient

		warrantClient warrant.Warrant
	)

	BeforeEach(func() {
		policyClient = policy_client.NewExternal(lagertest.NewTestLogger("test"),
			&http.Client{
				Transport: &http.Transport{
					TLSClientConfig: &tls.Config{
						InsecureSkipVerify: true,
					},
				},
			},
			fmt.Sprintf("https://%s", config.ApiEndpoint),
		)

		warrantClient = warrant.New(warrant.Config{
			Host:          getUAABaseURL(),
			SkipVerifySSL: true,
		})

		prefix = testConfig.Prefix
		appA = fmt.Sprintf("appA-%d", rand.Int31())
		appB = fmt.Sprintf("appB-%d", rand.Int31())

		AuthAsAdmin()

		orgName = prefix + "space-developer-org"
		Expect(cf.Cf("create-org", orgName).Wait(Timeout_Push)).To(gexec.Exit(0))
		Expect(cf.Cf("target", "-o", orgName).Wait(Timeout_Push)).To(gexec.Exit(0))

		spaceNameA = prefix + "space-A"
		Expect(cf.Cf("create-space", spaceNameA, "-o", orgName).Wait(Timeout_Push)).To(gexec.Exit(0))
		Expect(cf.Cf("target", "-o", orgName, "-s", spaceNameA).Wait(Timeout_Push)).To(gexec.Exit(0))

		pushProxy(appA)

		spaceNameB = prefix + "space-B"
		Expect(cf.Cf("create-space", spaceNameB, "-o", orgName).Wait(Timeout_Push)).To(gexec.Exit(0))
		Expect(cf.Cf("target", "-o", orgName, "-s", spaceNameB).Wait(Timeout_Push)).To(gexec.Exit(0))

		pushProxy(appB)

		uaaAdminClientToken, err := warrantClient.Clients.GetToken("admin", testConfig.AdminSecret)
		Expect(err).NotTo(HaveOccurred())

		user := ensureUserExists(warrantClient, "space-developer", "password", uaaAdminClientToken)
		group := ensureGroupExists(warrantClient, "network.write", uaaAdminClientToken)

		err = warrantClient.Groups.AddMember(group.ID, user.ID, uaaAdminClientToken)
		Expect(err).To(Or(BeNil(), BeAssignableToTypeOf(warrant.DuplicateResourceError{})))

		Expect(cf.Cf("set-space-role", "space-developer", orgName, spaceNameA, "SpaceDeveloper").Wait(Timeout_Push)).To(gexec.Exit(0))
		Expect(cf.Cf("set-space-role", "space-developer", orgName, spaceNameB, "SpaceDeveloper").Wait(Timeout_Push)).To(gexec.Exit(0))

		err = warrantClient.Clients.Create(warrant.Client{
			ID:                   "space-client",
			Name:                 "space-client",
			Authorities:          []string{"network.write", "cloud_controller.read"},
			AuthorizedGrantTypes: []string{"client_credentials"},
			AccessTokenValidity:  600 * time.Second,
		}, "password", uaaAdminClientToken)
		Expect(err).NotTo(HaveOccurred())

		cli := &cf_cli_adapter.Adapter{CfCliPath: "cf"}
		orgGuid, err := cli.OrgGuid(orgName)
		Expect(err).NotTo(HaveOccurred())

		spaceAGuid, err := cli.SpaceGuid(spaceNameA)
		Expect(err).NotTo(HaveOccurred())

		spaceBGuid, err := cli.SpaceGuid(spaceNameB)
		Expect(err).NotTo(HaveOccurred())

		cf.Cf("curl", "-X", "PUT", fmt.Sprintf("/v2/organizations/%s/users/space-client", orgGuid))
		cf.Cf("curl", "-X", "PUT", fmt.Sprintf("/v2/spaces/%s/developers/space-client", spaceAGuid))
		cf.Cf("curl", "-X", "PUT", fmt.Sprintf("/v2/spaces/%s/developers/space-client", spaceBGuid))
	})

	AfterEach(func() {
		By("logging in as admin and deleting the org", func() {
			Expect(cf.Cf("logout").Wait(Timeout_Push)).To(gexec.Exit(0))
			Expect(cf.Cf("auth", config.AdminUser, config.AdminPassword).Wait(Timeout_Push)).To(gexec.Exit(0))

			uaaAdminClientToken, err := warrantClient.Clients.GetToken("admin", testConfig.AdminSecret)
			Expect(err).NotTo(HaveOccurred())
			warrantClient.Clients.Delete("space-client", uaaAdminClientToken)

			Expect(cf.Cf("delete-org", orgName, "-f").Wait(Timeout_Push)).To(gexec.Exit(0))
		})
	})

	Describe("space developer with network.write scope", func() {
		DescribeTable("can create, list, and delete network policies in spaces they have access to", func(authArgs []string) {
			var spaceDevUserToken string
			By("logging in and getting the space developer user token", func() {
				authArgs = append([]string{"auth"}, authArgs...)
				Expect(cf.Cf(authArgs...).Wait(Timeout_Push)).To(gexec.Exit(0))
				session := cf.Cf("oauth-token")
				Expect(session.Wait(Timeout_Push)).To(gexec.Exit(0))
				spaceDevUserToken = strings.TrimSpace(string(session.Out.Contents()))
			})

			var appAGUID, appBGUID string
			By("getting the app guids", func() {
				Expect(cf.Cf("target", "-o", orgName, "-s", spaceNameA).Wait(Timeout_Push)).To(gexec.Exit(0))
				session := cf.Cf("app", appA, "--guid")
				Expect(session.Wait(Timeout_Push)).To(gexec.Exit(0))
				appAGUID = strings.TrimSpace(string(session.Out.Contents()))

				Expect(cf.Cf("target", "-o", orgName, "-s", spaceNameB).Wait(Timeout_Push)).To(gexec.Exit(0))
				session = cf.Cf("app", appB, "--guid")
				Expect(session.Wait(Timeout_Push)).To(gexec.Exit(0))
				appBGUID = strings.TrimSpace(string(session.Out.Contents()))
			})

			By("creating a policy", func() {
				err := policyClient.AddPoliciesV0(spaceDevUserToken, []api_v0.Policy{
					{
						Source: api_v0.Source{
							ID: appAGUID,
						},
						Destination: api_v0.Destination{
							ID:       appBGUID,
							Port:     1234,
							Protocol: "tcp",
						},
					},
				})
				Expect(err).NotTo(HaveOccurred())
			})

			By("listing policies", func() {
				expectedPolicies := []api_v0.Policy{
					{
						Source: api_v0.Source{
							ID: appAGUID,
						},
						Destination: api_v0.Destination{
							ID:       appBGUID,
							Port:     1234,
							Protocol: "tcp",
						},
					},
				}
				policies, err := policyClient.GetPoliciesV0(spaceDevUserToken)
				Expect(err).NotTo(HaveOccurred())
				Expect(policies).To(Equal(expectedPolicies))
			})

			By("deleting the policy", func() {
				err := policyClient.DeletePoliciesV0(spaceDevUserToken, []api_v0.Policy{
					{
						Source: api_v0.Source{
							ID: appAGUID,
						},
						Destination: api_v0.Destination{
							ID:       appBGUID,
							Port:     1234,
							Protocol: "tcp",
						},
					},
				})
				Expect(err).NotTo(HaveOccurred())
			})
		},
			Entry("as a user", []string{"space-developer", "password"}),
			Entry("as a service account", []string{"space-client", "password", "--client-credentials"}),
		)
	})
})

func ensureGroupExists(client warrant.Warrant, name, token string) warrant.Group {
	_, err := client.Groups.Create(name, token)
	Expect(err).To(Or(BeNil(), BeAssignableToTypeOf(warrant.DuplicateResourceError{})))

	groups, err := client.Groups.List(warrant.Query{Filter: fmt.Sprintf(`displayName eq %q`, name)}, token)
	Expect(err).NotTo(HaveOccurred())

	return groups[0]
}

func ensureUserExists(warrantClient warrant.Warrant, username, password, token string) warrant.User {
	createUserSession := cf.Cf("create-user", username, password)
	Expect(createUserSession.Wait(Timeout_Push)).To(gexec.Exit(0))

	users, err := warrantClient.Users.List(warrant.Query{Filter: fmt.Sprintf(`userName eq %q`, username)}, token)
	Expect(err).NotTo(HaveOccurred())

	return users[0]
}
<commit_msg>Hacky race condition fix<commit_after>package acceptance_test

import (
	"cf-pusher/cf_cli_adapter"
	"crypto/tls"
	"fmt"
	"lib/policy_client"
	"math/rand"
	"net/http"
	"policy-server/api/api_v0"
	"strings"
	"time"

	"code.cloudfoundry.org/lager/lagertest"
	"github.com/cloudfoundry-incubator/cf-test-helpers/cf"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/ginkgo/extensions/table"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gexec"
	"github.com/pivotal-cf-experimental/warrant"
)

var _ = Describe("space developer policy configuration", func() {
	var (
		appA       string
		appB       string
		orgName    string
		spaceNameA string
		spaceNameB string
		prefix     string

		policyClient *policy_client.ExternalClient

		warrantClient warrant.Warrant
	)

	BeforeEach(func() {
		policyClient = policy_client.NewExternal(lagertest.NewTestLogger("test"),
			&http.Client{
				Transport: &http.Transport{
					TLSClientConfig: &tls.Config{
						InsecureSkipVerify: true,
					},
				},
			},
			fmt.Sprintf("https://%s", config.ApiEndpoint),
		)

		warrantClient = warrant.New(warrant.Config{
			Host:          getUAABaseURL(),
			SkipVerifySSL: true,
		})

		prefix = testConfig.Prefix
		appA = fmt.Sprintf("appA-%d", rand.Int31())
		appB = fmt.Sprintf("appB-%d", rand.Int31())

		AuthAsAdmin()

		orgName = prefix + "space-developer-org"
		Expect(cf.Cf("create-org", orgName).Wait(Timeout_Push)).To(gexec.Exit(0))
		Expect(cf.Cf("target", "-o", orgName).Wait(Timeout_Push)).To(gexec.Exit(0))

		spaceNameA = prefix + "space-A"
		Expect(cf.Cf("create-space", spaceNameA, "-o", orgName).Wait(Timeout_Push)).To(gexec.Exit(0))
		Expect(cf.Cf("target", "-o", orgName, "-s", spaceNameA).Wait(Timeout_Push)).To(gexec.Exit(0))

		pushProxy(appA)

		spaceNameB = prefix + "space-B"
		Expect(cf.Cf("create-space", spaceNameB, "-o", orgName).Wait(Timeout_Push)).To(gexec.Exit(0))
		Expect(cf.Cf("target", "-o", orgName, "-s", spaceNameB).Wait(Timeout_Push)).To(gexec.Exit(0))

		pushProxy(appB)

		uaaAdminClientToken, err := warrantClient.Clients.GetToken("admin", testConfig.AdminSecret)
		Expect(err).NotTo(HaveOccurred())

		user := ensureUserExists(warrantClient, "space-developer", "password", uaaAdminClientToken)
		group := ensureGroupExists(warrantClient, "network.write", uaaAdminClientToken)

		err = warrantClient.Groups.AddMember(group.ID, user.ID, uaaAdminClientToken)
		Expect(err).To(Or(BeNil(), BeAssignableToTypeOf(warrant.DuplicateResourceError{})))

		Expect(cf.Cf("set-space-role", "space-developer", orgName, spaceNameA, "SpaceDeveloper").Wait(Timeout_Push)).To(gexec.Exit(0))
		Expect(cf.Cf("set-space-role", "space-developer", orgName, spaceNameB, "SpaceDeveloper").Wait(Timeout_Push)).To(gexec.Exit(0))

		err = warrantClient.Clients.Create(warrant.Client{
			ID:                   "space-client",
			Name:                 "space-client",
			Authorities:          []string{"network.write", "cloud_controller.read"},
			AuthorizedGrantTypes: []string{"client_credentials"},
			AccessTokenValidity:  600 * time.Second,
		}, "password", uaaAdminClientToken)
		Expect(err).NotTo(HaveOccurred())

		cli := &cf_cli_adapter.Adapter{CfCliPath: "cf"}
		orgGuid, err := cli.OrgGuid(orgName)
		Expect(err).NotTo(HaveOccurred())

		spaceAGuid, err := cli.SpaceGuid(spaceNameA)
		Expect(err).NotTo(HaveOccurred())

		spaceBGuid, err := cli.SpaceGuid(spaceNameB)
		Expect(err).NotTo(HaveOccurred())

		cf.Cf("curl", "-X", "PUT", fmt.Sprintf("/v2/organizations/%s/users/space-client", orgGuid))
		// Hack time.Sleep to deal with race condition, first curl has to succeed before the subsequent curls
		// Later change sleep to assert that the previous curl command succeeded before continuing
		time.Sleep(5 * time.Second)
		cf.Cf("curl", "-X", "PUT", fmt.Sprintf("/v2/spaces/%s/developers/space-client", spaceAGuid))
		cf.Cf("curl", "-X", "PUT", fmt.Sprintf("/v2/spaces/%s/developers/space-client", spaceBGuid))
	})

	AfterEach(func() {
		By("logging in as admin and deleting the org", func() {
			Expect(cf.Cf("logout").Wait(Timeout_Push)).To(gexec.Exit(0))
			Expect(cf.Cf("auth", config.AdminUser, config.AdminPassword).Wait(Timeout_Push)).To(gexec.Exit(0))

			uaaAdminClientToken, err := warrantClient.Clients.GetToken("admin", testConfig.AdminSecret)
			Expect(err).NotTo(HaveOccurred())
			warrantClient.Clients.Delete("space-client", uaaAdminClientToken)

			Expect(cf.Cf("delete-org", orgName, "-f").Wait(Timeout_Push)).To(gexec.Exit(0))
		})
	})

	Describe("space developer with network.write scope", func() {
		DescribeTable("can create, list, and delete network policies in spaces they have access to", func(authArgs []string) {
			var spaceDevUserToken string
			By("logging in and getting the space developer user token", func() {
				authArgs = append([]string{"auth"}, authArgs...)
				Expect(cf.Cf(authArgs...).Wait(Timeout_Push)).To(gexec.Exit(0))
				session := cf.Cf("oauth-token")
				Expect(session.Wait(Timeout_Push)).To(gexec.Exit(0))
				spaceDevUserToken = strings.TrimSpace(string(session.Out.Contents()))
			})

			var appAGUID, appBGUID string
			By("getting the app guids", func() {
				Expect(cf.Cf("target", "-o", orgName, "-s", spaceNameA).Wait(Timeout_Push)).To(gexec.Exit(0))
				session := cf.Cf("app", appA, "--guid")
				Expect(session.Wait(Timeout_Push)).To(gexec.Exit(0))
				appAGUID = strings.TrimSpace(string(session.Out.Contents()))

				Expect(cf.Cf("target", "-o", orgName, "-s", spaceNameB).Wait(Timeout_Push)).To(gexec.Exit(0))
				session = cf.Cf("app", appB, "--guid")
				Expect(session.Wait(Timeout_Push)).To(gexec.Exit(0))
				appBGUID = strings.TrimSpace(string(session.Out.Contents()))
			})

			By("creating a policy", func() {
				err := policyClient.AddPoliciesV0(spaceDevUserToken, []api_v0.Policy{
					{
						Source: api_v0.Source{
							ID: appAGUID,
						},
						Destination: api_v0.Destination{
							ID:       appBGUID,
							Port:     1234,
							Protocol: "tcp",
						},
					},
				})
				Expect(err).NotTo(HaveOccurred())
			})

			By("listing policies", func() {
				expectedPolicies := []api_v0.Policy{
					{
						Source: api_v0.Source{
							ID: appAGUID,
						},
						Destination: api_v0.Destination{
							ID:       appBGUID,
							Port:     1234,
							Protocol: "tcp",
						},
					},
				}
				policies, err := policyClient.GetPoliciesV0(spaceDevUserToken)
				Expect(err).NotTo(HaveOccurred())
				Expect(policies).To(Equal(expectedPolicies))
			})

			By("deleting the policy", func() {
				err := policyClient.DeletePoliciesV0(spaceDevUserToken, []api_v0.Policy{
					{
						Source: api_v0.Source{
							ID: appAGUID,
						},
						Destination: api_v0.Destination{
							ID:       appBGUID,
							Port:     1234,
							Protocol: "tcp",
						},
					},
				})
				Expect(err).NotTo(HaveOccurred())
			})
		},
			Entry("as a user", []string{"space-developer", "password"}),
			Entry("as a service account", []string{"space-client", "password", "--client-credentials"}),
		)
	})
})

func ensureGroupExists(client warrant.Warrant, name, token string) warrant.Group {
	_, err := client.Groups.Create(name, token)
	Expect(err).To(Or(BeNil(), BeAssignableToTypeOf(warrant.DuplicateResourceError{})))

	groups, err := client.Groups.List(warrant.Query{Filter: fmt.Sprintf(`displayName eq %q`, name)}, token)
	Expect(err).NotTo(HaveOccurred())

	return groups[0]
}

func ensureUserExists(warrantClient warrant.Warrant, username, password, token string) warrant.User {
	createUserSession := cf.Cf("create-user", username, password)
	Expect(createUserSession.Wait(Timeout_Push)).To(gexec.Exit(0))

	users, err := warrantClient.Users.List(warrant.Query{Filter: fmt.Sprintf(`userName eq %q`, username)}, token)
	Expect(err).NotTo(HaveOccurred())

	return users[0]
}
<|endoftext|>


<commit_before>package consumer

import (
	"log"
	"github.com/satori/go.uuid"
	"github.com/streadway/amqp"
)

type Consumer struct {
	uri string
	exchangeName string
	conn *amqp.Connection
	ch *amqp.Channel
	Deliveries <-chan amqp.Delivery
}

func New(uri string, exchangeName string) *Consumer {
	c := &Consumer{}
	c.uri = uri
	c.exchangeName = exchangeName

	return c
}

func (c *Consumer) Start() (err error) {
	c.conn, c.ch, err = c.openChannel()

	q, err := c.createExpirableQueue(c.ch)

    c.Deliveries, err = c.ch.Consume(q.Name, "", true, false, false, false, nil)

    return err
}

func (c *Consumer) Stop() {
    c.conn.Close()
    c.ch.Close()
}

func (c *Consumer) openChannel() (*amqp.Connection, *amqp.Channel, error) {

	var conn *amqp.Connection
	var ch *amqp.Channel
	var err error

    log.Printf("Establishing connection... "+ c.uri)
    conn, err = amqp.Dial(c.uri)
    log.Printf("Connected")

    if err != nil {
	    log.Printf("Fail connecting")

    	return conn, ch, err
    }

    ch, err = conn.Channel()
    log.Printf("Channel opened")
    if err != nil {
	    log.Printf("Fail opening channel")
    }

    return conn, ch, err
}

func (c *Consumer) createExpirableQueue(ch *amqp.Channel) (amqp.Queue, error) {
    var args amqp.Table
    args = make(amqp.Table)
    args["x-expires"] = int32(10000)

    q, err := ch.QueueDeclare(
      "tailmq_"+uuid.NewV4().String(), // name
      false,   // durable
      true,    // delete when unused
      false,   // exclusive
      false,   // no-wait
      args,    // arguments
    )

    if err != nil {
	    log.Printf("Fail declaring queue")
    	return q, err
    }

    err = ch.QueueBind(
      q.Name, // name
      "#",   // routing key
      c.exchangeName,   // exchange name
      false,   // exclusive
      nil,     // arguments
    )

    if err != nil {
	    log.Printf("Fail binding queue")
		return q, err
    }

    err = ch.QueueBind(
      q.Name, // name
      "",   // routing key
      c.exchangeName,   // exchange name
      false,   // exclusive
      nil,     // arguments
    )

    if err != nil {
	    log.Printf("Fail binding queue")
		return q, err
    }

    log.Printf("Queue defined")

    return q, err
}<commit_msg>Remove logging<commit_after>package consumer

import (
	"log"
	"github.com/satori/go.uuid"
	"github.com/streadway/amqp"
)

type Consumer struct {
	uri string
	exchangeName string
	conn *amqp.Connection
	ch *amqp.Channel
	Deliveries <-chan amqp.Delivery
}

func New(uri string, exchangeName string) *Consumer {
	c := &Consumer{}
	c.uri = uri
	c.exchangeName = exchangeName

	return c
}

func (c *Consumer) Start() (err error) {
	c.conn, c.ch, err = c.openChannel()

	q, err := c.createExpirableQueue(c.ch)

    c.Deliveries, err = c.ch.Consume(q.Name, "", true, false, false, false, nil)

    return err
}

func (c *Consumer) Stop() {
    c.conn.Close()
    c.ch.Close()
}

func (c *Consumer) openChannel() (*amqp.Connection, *amqp.Channel, error) {

    log.Printf("Establishing connection... "+ c.uri)
    conn, err := amqp.Dial(c.uri)
    log.Printf("Connected")

    ch, err := conn.Channel()
    log.Printf("Channel opened")

    return conn, ch, err
}

func (c *Consumer) createExpirableQueue(ch *amqp.Channel) (amqp.Queue, error) {
    var args amqp.Table
    args = make(amqp.Table)
    args["x-expires"] = int32(10000)

    q, err := ch.QueueDeclare(
      "tailmq_"+uuid.NewV4().String(), // name
      false,   // durable
      true,    // delete when unused
      false,   // exclusive
      false,   // no-wait
      args,    // arguments
    )

    err = ch.QueueBind(
      q.Name, // name
      "#",   // routing key
      c.exchangeName,   // exchange name
      false,   // exclusive
      nil,     // arguments
    )

    err = ch.QueueBind(
      q.Name, // name
      "",   // routing key
      c.exchangeName,   // exchange name
      false,   // exclusive
      nil,     // arguments
    )

    log.Printf("Queue defined")

    return q, err
}<|endoftext|>


<commit_before>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package container

const (
	// BridgeNetwork will have the container use the lxc bridge.
	BridgeNetwork = "bridge"
	// PhyscialNetwork will have the container use a specified network device.
	PhysicalNetwork = "physical"
)

// NetworkConfig defines how the container network will be configured.
type NetworkConfig struct {
	NetworkType string
	Device      string
}

// BridgeNetworkConfig returns a valid NetworkConfig to use the specified
// device as a network bridge for the container.
func BridgeNetworkConfig(device string) *NetworkConfig {
	return &NetworkConfig{BridgeNetwork, device}
}

// PhysicalNetworkConfig returns a valid NetworkConfig to use the specified
// device as the network device for the container.
func PhysicalNetworkConfig(device string) *NetworkConfig {
	return &NetworkConfig{PhysicalNetwork, device}
}
<commit_msg>Comment update.<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package container

const (
	// BridgeNetwork will have the container use the network bridge.
	BridgeNetwork = "bridge"
	// PhyscialNetwork will have the container use a specified network device.
	PhysicalNetwork = "physical"
)

// NetworkConfig defines how the container network will be configured.
type NetworkConfig struct {
	NetworkType string
	Device      string
}

// BridgeNetworkConfig returns a valid NetworkConfig to use the specified
// device as a network bridge for the container.
func BridgeNetworkConfig(device string) *NetworkConfig {
	return &NetworkConfig{BridgeNetwork, device}
}

// PhysicalNetworkConfig returns a valid NetworkConfig to use the specified
// device as the network device for the container.
func PhysicalNetworkConfig(device string) *NetworkConfig {
	return &NetworkConfig{PhysicalNetwork, device}
}
<|endoftext|>


<commit_before>package reset_test

import (
	"testing"

	. "github.com/redforks/testing/reset"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

func TestRest(t *testing.T) {
	RegisterFailHandler(Fail)
	RunSpecs(t, "Reset")
}

var _ = Describe("Reset", func() {
	var log []string

	resetA := func() {
		log = append(log, `a`)
	}

	resetB := func() {
		log = append(log, `b`)
	}

	BeforeEach(func() {
		log = []string{}
	})

	AfterEach(func() {
		if Enabled() {
			Disable()
		}
	})

	It("Not Enabled", func() {
		Ω(Enabled()).Should(BeFalse())
		Add(resetA)
		Add(resetB)

		Ω(log).Should(BeEmpty())
	})

	It("Set disabled disabled", func() {
		Ω(func() {
			Disable()
		}).Should(Panic())
	})

	Context("Enabled", func() {

		BeforeEach(func() {
			Enable()
			Add(resetA)
		})

		It("Enabled", func() {
			Ω(Enabled()).Should(BeTrue())
			Ω(log).Should(BeEmpty())
			Disable()
			Ω(Enabled()).Should(BeFalse())
			Ω(log).Should(Equal([]string{"a"}))
		})

		It("Execute by reversed order", func() {
			Add(resetB)
			Disable()
			Ω(log).Should(Equal([]string{"b", "a"}))
		})

		It("Add dup action", func() {
			Add(resetA)
			Add(resetA)
			Disable()
			Ω(log).Should(Equal([]string{"a", "a", "a"}))
		})

		It("Not allow Add() while executing", func() {
			Add(func() {
				Add(resetA)
			})
			Ω(func() {
				Disable()
			}).Should(Panic())
		})

	})

})
<commit_msg>Rewrite reset_test.go to testify<commit_after>package reset_test

import (
	"testing"

	myTesting "github.com/redforks/testing"
	. "github.com/redforks/testing/reset"
	"github.com/stretchr/testify/assert"
)

func TestReset(t *testing.T) {
	var log []string

	beforeEach := func() {
		log = []string{}
	}

	afterEach := func() {
		if Enabled() {
			Disable()
		}
	}

	resetA := func() {
		log = append(log, `a`)
	}

	resetB := func() {
		log = append(log, `b`)
	}

	assertLog := func(exp ...string) {
		assert.Equal(t, exp, log)
	}

	newTest := func(f func(t *testing.T)) func(t *testing.T) {
		return myTesting.SetupTeardown(beforeEach, afterEach, f)
	}

	newEnabledTest := func(f func(t *testing.T)) func(t *testing.T) {
		return myTesting.SetupTeardown(func() {
			beforeEach()
			Enable()
			Add(resetA)
		}, afterEach, f)
	}

	t.Run("Not Enabled", newTest(func(t *testing.T) {
		assert.False(t, Enabled())

		Add(resetA)
		Add(resetB)

		assert.Empty(t, log)
	}))

	t.Run("Set disabled disabled", newTest(func(t *testing.T) {
		assert.Panics(t, Disable)
	}))

	t.Run("Enable/Disable", newEnabledTest(func(t *testing.T) {
		assert.True(t, Enabled())
		assert.Empty(t, log)
		Disable()
		assert.False(t, Enabled())
		assertLog("a")
	}))

	t.Run("Execute by reversed order", newEnabledTest(func(t *testing.T) {
		Add(resetB)
		Disable()
		assertLog("b", "a")
	}))

	t.Run("Dup action", newEnabledTest(func(t *testing.T) {
		Add(resetA)
		Add(resetA)
		Disable()
		assertLog("a", "a", "a")
	}))

	t.Run("Not allow Add() while executing", newEnabledTest(func(t *testing.T) {
		Add(func() {
			Add(resetA)
		})
		assert.Panics(t, Disable)
	}))
}
<|endoftext|>


<commit_before>package resource

import (
	"github.com/concourse/atc"
	"github.com/concourse/atc/worker"
	"github.com/concourse/baggageclaim"
	"github.com/pivotal-golang/lager"
)

type ResourceType string
type ContainerImage string

type Session struct {
	ID        worker.Identifier
	Ephemeral bool
}

//go:generate counterfeiter . Tracker

type Tracker interface {
	Init(lager.Logger, Metadata, Session, ResourceType, atc.Tags) (Resource, error)
	InitWithCache(lager.Logger, Metadata, Session, ResourceType, atc.Tags, CacheIdentifier) (Resource, Cache, error)
}

//go:generate counterfeiter . Cache

type Cache interface {
	IsInitialized() (bool, error)
	Initialize() error
}

type Metadata interface {
	Env() []string
}

type EmptyMetadata struct{}

func (m EmptyMetadata) Env() []string { return nil }

type tracker struct {
	workerClient worker.Client
}

type TrackerFactory struct{}

func (TrackerFactory) TrackerFor(client worker.Client) Tracker {
	return NewTracker(client)
}

func NewTracker(workerClient worker.Client) Tracker {
	return &tracker{
		workerClient: workerClient,
	}
}

type VolumeMount struct {
	Volume    baggageclaim.Volume
	MountPath string
}

func (tracker *tracker) Init(logger lager.Logger, metadata Metadata, session Session, typ ResourceType, tags atc.Tags) (Resource, error) {
	container, found, err := tracker.workerClient.FindContainerForIdentifier(logger, session.ID)
	if err != nil {
		return nil, err
	}

	if !found {
		container, err = tracker.workerClient.CreateContainer(logger, session.ID, worker.ResourceTypeContainerSpec{
			Type:      string(typ),
			Ephemeral: session.Ephemeral,
			Tags:      tags,
			Env:       metadata.Env(),
		})
		if err != nil {
			return nil, err
		}
	}

	return NewResource(container), nil
}

func (tracker *tracker) InitWithCache(logger lager.Logger, metadata Metadata, session Session, typ ResourceType, tags atc.Tags, cacheIdentifier CacheIdentifier) (Resource, Cache, error) {
	container, found, err := tracker.workerClient.FindContainerForIdentifier(logger, session.ID)
	if err != nil {
		return nil, nil, err
	}

	if found {
		var cache Cache

		volumes := container.Volumes()
		switch len(volumes) {
		case 0:
			cache = noopCache{}
		case 1:
			cache = volumeCache{volumes[0]}
		}

		return NewResource(container), cache, nil
	}

	resourceSpec := worker.WorkerSpec{
		ResourceType: string(typ),
		Tags:         tags,
	}

	chosenWorker, err := tracker.workerClient.Satisfying(resourceSpec)
	if err != nil {
		return nil, nil, err
	}

	vm, hasVM := chosenWorker.VolumeManager()
	if !hasVM {
		container, err := chosenWorker.CreateContainer(logger, session.ID, worker.ResourceTypeContainerSpec{
			Type:      string(typ),
			Ephemeral: session.Ephemeral,
			Tags:      tags,
			Env:       metadata.Env(),
		})
		if err != nil {
			return nil, nil, err
		}

		return NewResource(container), noopCache{}, nil
	}

	cachedVolume, cacheFound, err := cacheIdentifier.FindOn(logger, vm)
	if err != nil {
		return nil, nil, err
	}

	if cacheFound {
		logger.Info("found-cache", lager.Data{"handle": cachedVolume.Handle()})
	} else {
		logger.Debug("no-cache-found")

		cachedVolume, err = cacheIdentifier.CreateOn(logger, vm)
		if err != nil {
			return nil, nil, err
		}

		logger.Info("new-cache", lager.Data{"handle": cachedVolume.Handle()})
	}

	defer cachedVolume.Release()

	container, err = chosenWorker.CreateContainer(logger, session.ID, worker.ResourceTypeContainerSpec{
		Type:      string(typ),
		Ephemeral: session.Ephemeral,
		Tags:      tags,
		Env:       metadata.Env(),
		Cache: worker.VolumeMount{
			Volume:    cachedVolume,
			MountPath: ResourcesDir("get"),
		},
	})
	if err != nil {
		return nil, nil, err
	}

	return NewResource(container), volumeCache{cachedVolume}, nil
}
<commit_msg>prevent nil in edge case<commit_after>package resource

import (
	"github.com/concourse/atc"
	"github.com/concourse/atc/worker"
	"github.com/concourse/baggageclaim"
	"github.com/pivotal-golang/lager"
)

type ResourceType string
type ContainerImage string

type Session struct {
	ID        worker.Identifier
	Ephemeral bool
}

//go:generate counterfeiter . Tracker

type Tracker interface {
	Init(lager.Logger, Metadata, Session, ResourceType, atc.Tags) (Resource, error)
	InitWithCache(lager.Logger, Metadata, Session, ResourceType, atc.Tags, CacheIdentifier) (Resource, Cache, error)
}

//go:generate counterfeiter . Cache

type Cache interface {
	IsInitialized() (bool, error)
	Initialize() error
}

type Metadata interface {
	Env() []string
}

type EmptyMetadata struct{}

func (m EmptyMetadata) Env() []string { return nil }

type tracker struct {
	workerClient worker.Client
}

type TrackerFactory struct{}

func (TrackerFactory) TrackerFor(client worker.Client) Tracker {
	return NewTracker(client)
}

func NewTracker(workerClient worker.Client) Tracker {
	return &tracker{
		workerClient: workerClient,
	}
}

type VolumeMount struct {
	Volume    baggageclaim.Volume
	MountPath string
}

func (tracker *tracker) Init(logger lager.Logger, metadata Metadata, session Session, typ ResourceType, tags atc.Tags) (Resource, error) {
	container, found, err := tracker.workerClient.FindContainerForIdentifier(logger, session.ID)
	if err != nil {
		return nil, err
	}

	if !found {
		container, err = tracker.workerClient.CreateContainer(logger, session.ID, worker.ResourceTypeContainerSpec{
			Type:      string(typ),
			Ephemeral: session.Ephemeral,
			Tags:      tags,
			Env:       metadata.Env(),
		})
		if err != nil {
			return nil, err
		}
	}

	return NewResource(container), nil
}

func (tracker *tracker) InitWithCache(logger lager.Logger, metadata Metadata, session Session, typ ResourceType, tags atc.Tags, cacheIdentifier CacheIdentifier) (Resource, Cache, error) {
	container, found, err := tracker.workerClient.FindContainerForIdentifier(logger, session.ID)
	if err != nil {
		return nil, nil, err
	}

	if found {
		var cache Cache

		volumes := container.Volumes()
		switch len(volumes) {
		case 0:
			cache = noopCache{}
		default:
			cache = volumeCache{volumes[0]}
		}

		return NewResource(container), cache, nil
	}

	resourceSpec := worker.WorkerSpec{
		ResourceType: string(typ),
		Tags:         tags,
	}

	chosenWorker, err := tracker.workerClient.Satisfying(resourceSpec)
	if err != nil {
		return nil, nil, err
	}

	vm, hasVM := chosenWorker.VolumeManager()
	if !hasVM {
		container, err := chosenWorker.CreateContainer(logger, session.ID, worker.ResourceTypeContainerSpec{
			Type:      string(typ),
			Ephemeral: session.Ephemeral,
			Tags:      tags,
			Env:       metadata.Env(),
		})
		if err != nil {
			return nil, nil, err
		}

		return NewResource(container), noopCache{}, nil
	}

	cachedVolume, cacheFound, err := cacheIdentifier.FindOn(logger, vm)
	if err != nil {
		return nil, nil, err
	}

	if cacheFound {
		logger.Info("found-cache", lager.Data{"handle": cachedVolume.Handle()})
	} else {
		logger.Debug("no-cache-found")

		cachedVolume, err = cacheIdentifier.CreateOn(logger, vm)
		if err != nil {
			return nil, nil, err
		}

		logger.Info("new-cache", lager.Data{"handle": cachedVolume.Handle()})
	}

	defer cachedVolume.Release()

	container, err = chosenWorker.CreateContainer(logger, session.ID, worker.ResourceTypeContainerSpec{
		Type:      string(typ),
		Ephemeral: session.Ephemeral,
		Tags:      tags,
		Env:       metadata.Env(),
		Cache: worker.VolumeMount{
			Volume:    cachedVolume,
			MountPath: ResourcesDir("get"),
		},
	})
	if err != nil {
		return nil, nil, err
	}

	return NewResource(container), volumeCache{cachedVolume}, nil
}
<|endoftext|>


<commit_before>package cache

import (
	"errors"
	"fmt"
	"math/rand"
	"sort"
	"testing"
	"time"

	"github.com/grafana/metrictank/mdata/cache/accnt"
	"github.com/grafana/metrictank/mdata/chunk"
	"github.com/grafana/metrictank/test"
	"github.com/raintank/schema"
)

func generateChunks(b testing.TB, startAt, count, step uint32) []chunk.IterGen {
	res := make([]chunk.IterGen, 0, count)

	values := make([]uint32, step)
	for t0 := startAt; t0 < startAt+(step*uint32(count)); t0 += step {
		c := getItgen(b, values, t0, true)
		res = append(res, c)
	}
	return res
}

func getCCM() (schema.AMKey, *CCacheMetric) {
	amkey, _ := schema.AMKeyFromString("1.12345678901234567890123456789012")
	ccm := NewCCacheMetric(amkey.MKey)
	return amkey, ccm
}

// TestAddAsc tests adding ascending timestamp chunks individually
func TestAddAsc(t *testing.T) {
	testRun(t, func(ccm *CCacheMetric) {
		chunks := generateChunks(t, 10, 6, 10)
		prev := uint32(1)
		for _, chunk := range chunks {
			ccm.Add(prev, chunk)
			prev = chunk.Ts
		}
	})
}

// TestAddDesc1 tests adding chunks that are all descending
func TestAddDesc1(t *testing.T) {
	testRun(t, func(ccm *CCacheMetric) {
		chunks := generateChunks(t, 10, 6, 10)
		for i := len(chunks) - 1; i >= 0; i-- {
			ccm.Add(0, chunks[i])
		}
	})
}

// TestAddDesc4 tests adding chunks that are globally descending
// but in groups 4 are ascending
func TestAddDesc4(t *testing.T) {
	testRun(t, func(ccm *CCacheMetric) {
		chunks := generateChunks(t, 10, 6, 10)
		ccm.Add(0, chunks[2])
		ccm.Add(0, chunks[3])
		ccm.Add(0, chunks[4])
		ccm.Add(0, chunks[5])
		ccm.Add(0, chunks[0])
		ccm.Add(0, chunks[1])
	})
}

// TestAddRange tests adding a contiguous range at once
func TestAddRange(t *testing.T) {
	testRun(t, func(ccm *CCacheMetric) {
		chunks := generateChunks(t, 10, 6, 10)
		prev := uint32(10)
		ccm.AddRange(prev, chunks)
	})
}

// TestAddRangeDesc4 benchmarks adding chunks that are globally descending
// but in groups of 4 are ascending. those groups are added via 1 AddRange.
func TestAddRangeDesc4(t *testing.T) {
	testRun(t, func(ccm *CCacheMetric) {
		chunks := generateChunks(t, 10, 6, 10)
		ccm.AddRange(0, chunks[2:6])
		ccm.AddRange(0, chunks[0:2])
	})
}

// test executes the run function
// run should generate chunks and add them to the CCacheMetric however it likes,
// but in a way so that the result will be as expected
func testRun(t *testing.T, run func(*CCacheMetric)) {
	amkey, ccm := getCCM()

	run(ccm)

	res := CCSearchResult{}
	ccm.Search(test.NewContext(), amkey, &res, 25, 45)
	if res.Complete != true {
		t.Fatalf("Expected result to be complete, but it was not")
	}

	if res.Start[0].Ts != 20 {
		t.Fatalf("Expected result to start at 20, but had %d", res.Start[0].Ts)
	}

	if res.Start[len(res.Start)-1].Ts != 40 {
		t.Fatalf("Expected result to start at 40, but had %d", res.Start[len(res.Start)-1].Ts)
	}
}

// BenchmarkAddAsc benchmarks adding ascending timestamp chunks individually
func BenchmarkAddAsc(b *testing.B) {
	_, ccm := getCCM()
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	prev := uint32(1)
	b.ResetTimer()
	for _, chunk := range chunks {
		ccm.Add(prev, chunk)
		prev = chunk.Ts
	}
}

// BenchmarkAddDesc1 benchmarks adding chunks that are all descending
func BenchmarkAddDesc1(b *testing.B) {
	_, ccm := getCCM()
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	b.ResetTimer()
	for i := len(chunks) - 1; i >= 0; i-- {
		ccm.Add(0, chunks[i])
	}
}

// BenchmarkAddDesc4 benchmarks adding chunks that are globally descending
// but in groups 4 are ascending
func BenchmarkAddDesc4(b *testing.B) {
	_, ccm := getCCM()
	b.N = b.N - (b.N % 4)
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	b.ResetTimer()
	for i := len(chunks) - 4; i >= 0; i -= 4 {
		ccm.Add(0, chunks[i])
		ccm.Add(0, chunks[i+1])
		ccm.Add(0, chunks[i+2])
		ccm.Add(0, chunks[i+3])
	}
}

// BenchmarkAddDesc64 benchmarks adding chunks that are globally descending
// but in groups 64 are ascending
func BenchmarkAddDesc64(b *testing.B) {
	_, ccm := getCCM()
	b.N = b.N - (b.N % 64)
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	b.ResetTimer()
	for i := len(chunks) - 64; i >= 0; i -= 64 {
		for offset := 0; offset < 64; offset += 1 {
			ccm.Add(0, chunks[i+offset])
		}
	}

}

// BenchmarkAddRangeAsc benchmarks adding a contiguous range at once
func BenchmarkAddRangeAsc(b *testing.B) {
	_, ccm := getCCM()
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	prev := uint32(1)
	b.ResetTimer()
	ccm.AddRange(prev, chunks)
}

// BenchmarkAddRangeDesc4 benchmarks adding chunks that are globally descending
// but in groups 4 are ascending. those groups are added via 1 AddRange.
func BenchmarkAddRangeDesc4(b *testing.B) {
	_, ccm := getCCM()
	b.N = b.N - (b.N % 4)
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	b.ResetTimer()
	for i := len(chunks) - 4; i >= 0; i -= 4 {
		ccm.AddRange(0, chunks[i:i+4])
	}
}

// BenchmarkAddRangeDesc64 benchmarks adding chunks that are globally descending
// but in groups 64 are ascending. those groups are added via 1 AddRange.
func BenchmarkAddRangeDesc64(b *testing.B) {
	_, ccm := getCCM()
	b.N = b.N - (b.N % 64)
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	b.ResetTimer()
	for i := len(chunks) - 64; i >= 0; i -= 64 {
		ccm.AddRange(0, chunks[i:i+64])
	}
}

func TestCorruptionCase1(t *testing.T) {
	testRun(t, func(ccm *CCacheMetric) {
		chunks := generateChunks(t, 10, 6, 10)
		ccm.AddRange(0, chunks[3:6])
		ccm.AddRange(0, chunks[0:4])
		if err := verifyCcm(ccm); err != nil {
			t.Fatal(err)
		}
	})
}

func getRandomNumber(min, max int) int {
	return rand.Intn(max-min) + min
}

func getRandomRange(min, max int) (int, int) {
	start := getRandomNumber(min, max)
	var end int
	for end = min - 1; end < start; end = getRandomNumber(min, max) {
	}

	return start, end
}

func TestCorruptionCase2(t *testing.T) {
	rand.Seed(time.Now().Unix())
	_, ccm := getCCM()
	iterations := 100000000
	chunks := generateChunks(t, 10, 100, 10)
	adds, addRanges, dels := 0, 0, 0

	for i := 0; i < iterations; i++ {
		// 0 = Add
		// 1 = AddRange
		// 2 = Del
		action := getRandomNumber(0, 3)
		switch action {
		case 0:
			chunk := getRandomNumber(0, 100)
			//fmt.Println(fmt.Sprintf("adding chunk %d", chunk))
			ccm.Add(0, chunks[chunk])
			adds++
		case 1:
			from, to := getRandomRange(0, 100)
			//fmt.Println(fmt.Sprintf("adding range %d-%d", from, to))
			ccm.AddRange(0, chunks[from:to])
			addRanges++
		case 2:
			chunk := getRandomNumber(0, 100)
			//fmt.Println(fmt.Sprintf("deleting chunk %d", chunk))
			ccm.Del(chunks[chunk].Ts)
			dels++
		}

		if err := verifyCcm(ccm); err != nil {
			t.Fatal(err)
		}
	}

	fmt.Println(fmt.Sprintf("adds: %d addRanges: %d dels:%d", adds, addRanges, dels))
}

// verifyCcm verifies the integrity of a CCacheMetric
// it assumes that all itergens are span-aware
func verifyCcm(ccm *CCacheMetric) error {
	var chunk *CCacheChunk
	var ok bool

	if len(ccm.chunks) != len(ccm.keys) {
		return errors.New("Length of ccm.chunks does not match ccm.keys")
	}

	if !sort.IsSorted(accnt.Uint32Asc(ccm.keys)) {
		return errors.New("keys are not sorted")
	}

	for i, ts := range ccm.keys {
		if chunk, ok = ccm.chunks[ts]; !ok {
			return fmt.Errorf("Ts %d is in ccm.keys but not in ccm.chunks", ts)
		}

		if i == 0 {
			if chunk.Prev != 0 {
				return errors.New("First chunk has Prev != 0")
			}
		} else {
			if chunk.Prev == 0 {
				if ccm.chunks[ccm.keys[i-1]].Ts == chunk.Ts-chunk.Itgen.Span {
					return fmt.Errorf("Chunk of ts %d has Prev == 0, but the previous chunk is present", ts)
				}
			} else {
				if ccm.chunks[ccm.keys[i-1]].Ts != chunk.Prev {
					return fmt.Errorf("Chunk of ts %d has Prev set to wrong ts %d but should be %d", ts, chunk.Prev, ccm.chunks[ccm.keys[i-1]].Ts)
				}
			}
		}

		if i == len(ccm.keys)-1 {
			if chunk.Next != 0 {
				return fmt.Errorf("Next of last chunk should be 0, but it's %d", chunk.Next)
			}

			// all checks completed
			break
		}

		var nextChunk *CCacheChunk
		if nextChunk, ok = ccm.chunks[ccm.keys[i+1]]; !ok {
			return fmt.Errorf("Ts %d is in ccm.keys but not in ccm.chunks", ccm.keys[i+1])
		}

		if chunk.Next == 0 {
			if chunk.Ts+chunk.Itgen.Span == nextChunk.Ts {
				return fmt.Errorf("Next of chunk at ts %d is set to 0, but the next chunk is present", ts)
			}
		} else {
			if chunk.Next != nextChunk.Ts {
				return fmt.Errorf("Next of chunk at ts %d is set to %d, but it should be %d", ts, chunk.Next, nextChunk.Ts)
			}
		}
	}
	return nil
}
<commit_msg>more reasonable of number of iterations (by default)<commit_after>package cache

import (
	"errors"
	"fmt"
	"math/rand"
	"sort"
	"testing"
	"time"

	"github.com/grafana/metrictank/mdata/cache/accnt"
	"github.com/grafana/metrictank/mdata/chunk"
	"github.com/grafana/metrictank/test"
	"github.com/raintank/schema"
)

func generateChunks(b testing.TB, startAt, count, step uint32) []chunk.IterGen {
	res := make([]chunk.IterGen, 0, count)

	values := make([]uint32, step)
	for t0 := startAt; t0 < startAt+(step*uint32(count)); t0 += step {
		c := getItgen(b, values, t0, true)
		res = append(res, c)
	}
	return res
}

func getCCM() (schema.AMKey, *CCacheMetric) {
	amkey, _ := schema.AMKeyFromString("1.12345678901234567890123456789012")
	ccm := NewCCacheMetric(amkey.MKey)
	return amkey, ccm
}

// TestAddAsc tests adding ascending timestamp chunks individually
func TestAddAsc(t *testing.T) {
	testRun(t, func(ccm *CCacheMetric) {
		chunks := generateChunks(t, 10, 6, 10)
		prev := uint32(1)
		for _, chunk := range chunks {
			ccm.Add(prev, chunk)
			prev = chunk.Ts
		}
	})
}

// TestAddDesc1 tests adding chunks that are all descending
func TestAddDesc1(t *testing.T) {
	testRun(t, func(ccm *CCacheMetric) {
		chunks := generateChunks(t, 10, 6, 10)
		for i := len(chunks) - 1; i >= 0; i-- {
			ccm.Add(0, chunks[i])
		}
	})
}

// TestAddDesc4 tests adding chunks that are globally descending
// but in groups 4 are ascending
func TestAddDesc4(t *testing.T) {
	testRun(t, func(ccm *CCacheMetric) {
		chunks := generateChunks(t, 10, 6, 10)
		ccm.Add(0, chunks[2])
		ccm.Add(0, chunks[3])
		ccm.Add(0, chunks[4])
		ccm.Add(0, chunks[5])
		ccm.Add(0, chunks[0])
		ccm.Add(0, chunks[1])
	})
}

// TestAddRange tests adding a contiguous range at once
func TestAddRange(t *testing.T) {
	testRun(t, func(ccm *CCacheMetric) {
		chunks := generateChunks(t, 10, 6, 10)
		prev := uint32(10)
		ccm.AddRange(prev, chunks)
	})
}

// TestAddRangeDesc4 benchmarks adding chunks that are globally descending
// but in groups of 4 are ascending. those groups are added via 1 AddRange.
func TestAddRangeDesc4(t *testing.T) {
	testRun(t, func(ccm *CCacheMetric) {
		chunks := generateChunks(t, 10, 6, 10)
		ccm.AddRange(0, chunks[2:6])
		ccm.AddRange(0, chunks[0:2])
	})
}

// test executes the run function
// run should generate chunks and add them to the CCacheMetric however it likes,
// but in a way so that the result will be as expected
func testRun(t *testing.T, run func(*CCacheMetric)) {
	amkey, ccm := getCCM()

	run(ccm)

	res := CCSearchResult{}
	ccm.Search(test.NewContext(), amkey, &res, 25, 45)
	if res.Complete != true {
		t.Fatalf("Expected result to be complete, but it was not")
	}

	if res.Start[0].Ts != 20 {
		t.Fatalf("Expected result to start at 20, but had %d", res.Start[0].Ts)
	}

	if res.Start[len(res.Start)-1].Ts != 40 {
		t.Fatalf("Expected result to start at 40, but had %d", res.Start[len(res.Start)-1].Ts)
	}
}

// BenchmarkAddAsc benchmarks adding ascending timestamp chunks individually
func BenchmarkAddAsc(b *testing.B) {
	_, ccm := getCCM()
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	prev := uint32(1)
	b.ResetTimer()
	for _, chunk := range chunks {
		ccm.Add(prev, chunk)
		prev = chunk.Ts
	}
}

// BenchmarkAddDesc1 benchmarks adding chunks that are all descending
func BenchmarkAddDesc1(b *testing.B) {
	_, ccm := getCCM()
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	b.ResetTimer()
	for i := len(chunks) - 1; i >= 0; i-- {
		ccm.Add(0, chunks[i])
	}
}

// BenchmarkAddDesc4 benchmarks adding chunks that are globally descending
// but in groups 4 are ascending
func BenchmarkAddDesc4(b *testing.B) {
	_, ccm := getCCM()
	b.N = b.N - (b.N % 4)
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	b.ResetTimer()
	for i := len(chunks) - 4; i >= 0; i -= 4 {
		ccm.Add(0, chunks[i])
		ccm.Add(0, chunks[i+1])
		ccm.Add(0, chunks[i+2])
		ccm.Add(0, chunks[i+3])
	}
}

// BenchmarkAddDesc64 benchmarks adding chunks that are globally descending
// but in groups 64 are ascending
func BenchmarkAddDesc64(b *testing.B) {
	_, ccm := getCCM()
	b.N = b.N - (b.N % 64)
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	b.ResetTimer()
	for i := len(chunks) - 64; i >= 0; i -= 64 {
		for offset := 0; offset < 64; offset += 1 {
			ccm.Add(0, chunks[i+offset])
		}
	}

}

// BenchmarkAddRangeAsc benchmarks adding a contiguous range at once
func BenchmarkAddRangeAsc(b *testing.B) {
	_, ccm := getCCM()
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	prev := uint32(1)
	b.ResetTimer()
	ccm.AddRange(prev, chunks)
}

// BenchmarkAddRangeDesc4 benchmarks adding chunks that are globally descending
// but in groups 4 are ascending. those groups are added via 1 AddRange.
func BenchmarkAddRangeDesc4(b *testing.B) {
	_, ccm := getCCM()
	b.N = b.N - (b.N % 4)
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	b.ResetTimer()
	for i := len(chunks) - 4; i >= 0; i -= 4 {
		ccm.AddRange(0, chunks[i:i+4])
	}
}

// BenchmarkAddRangeDesc64 benchmarks adding chunks that are globally descending
// but in groups 64 are ascending. those groups are added via 1 AddRange.
func BenchmarkAddRangeDesc64(b *testing.B) {
	_, ccm := getCCM()
	b.N = b.N - (b.N % 64)
	chunks := generateChunks(b, 10, uint32(b.N), 10)
	b.ResetTimer()
	for i := len(chunks) - 64; i >= 0; i -= 64 {
		ccm.AddRange(0, chunks[i:i+64])
	}
}

func TestCorruptionCase1(t *testing.T) {
	testRun(t, func(ccm *CCacheMetric) {
		chunks := generateChunks(t, 10, 6, 10)
		ccm.AddRange(0, chunks[3:6])
		ccm.AddRange(0, chunks[0:4])
		if err := verifyCcm(ccm); err != nil {
			t.Fatal(err)
		}
	})
}

func getRandomNumber(min, max int) int {
	return rand.Intn(max-min) + min
}

func getRandomRange(min, max int) (int, int) {
	start := getRandomNumber(min, max)
	var end int
	for end = min - 1; end < start; end = getRandomNumber(min, max) {
	}

	return start, end
}

func TestCorruptionCase2(t *testing.T) {
	rand.Seed(time.Now().Unix())
	_, ccm := getCCM()
	iterations := 100000
	chunks := generateChunks(t, 10, 100, 10)
	adds, addRanges, dels := 0, 0, 0

	for i := 0; i < iterations; i++ {
		// 0 = Add
		// 1 = AddRange
		// 2 = Del
		action := getRandomNumber(0, 3)
		switch action {
		case 0:
			chunk := getRandomNumber(0, 100)
			//fmt.Println(fmt.Sprintf("adding chunk %d", chunk))
			ccm.Add(0, chunks[chunk])
			adds++
		case 1:
			from, to := getRandomRange(0, 100)
			//fmt.Println(fmt.Sprintf("adding range %d-%d", from, to))
			ccm.AddRange(0, chunks[from:to])
			addRanges++
		case 2:
			chunk := getRandomNumber(0, 100)
			//fmt.Println(fmt.Sprintf("deleting chunk %d", chunk))
			ccm.Del(chunks[chunk].Ts)
			dels++
		}

		if err := verifyCcm(ccm); err != nil {
			t.Fatal(err)
		}
	}

	fmt.Println(fmt.Sprintf("adds: %d addRanges: %d dels:%d", adds, addRanges, dels))
}

// verifyCcm verifies the integrity of a CCacheMetric
// it assumes that all itergens are span-aware
func verifyCcm(ccm *CCacheMetric) error {
	var chunk *CCacheChunk
	var ok bool

	if len(ccm.chunks) != len(ccm.keys) {
		return errors.New("Length of ccm.chunks does not match ccm.keys")
	}

	if !sort.IsSorted(accnt.Uint32Asc(ccm.keys)) {
		return errors.New("keys are not sorted")
	}

	for i, ts := range ccm.keys {
		if chunk, ok = ccm.chunks[ts]; !ok {
			return fmt.Errorf("Ts %d is in ccm.keys but not in ccm.chunks", ts)
		}

		if i == 0 {
			if chunk.Prev != 0 {
				return errors.New("First chunk has Prev != 0")
			}
		} else {
			if chunk.Prev == 0 {
				if ccm.chunks[ccm.keys[i-1]].Ts == chunk.Ts-chunk.Itgen.Span {
					return fmt.Errorf("Chunk of ts %d has Prev == 0, but the previous chunk is present", ts)
				}
			} else {
				if ccm.chunks[ccm.keys[i-1]].Ts != chunk.Prev {
					return fmt.Errorf("Chunk of ts %d has Prev set to wrong ts %d but should be %d", ts, chunk.Prev, ccm.chunks[ccm.keys[i-1]].Ts)
				}
			}
		}

		if i == len(ccm.keys)-1 {
			if chunk.Next != 0 {
				return fmt.Errorf("Next of last chunk should be 0, but it's %d", chunk.Next)
			}

			// all checks completed
			break
		}

		var nextChunk *CCacheChunk
		if nextChunk, ok = ccm.chunks[ccm.keys[i+1]]; !ok {
			return fmt.Errorf("Ts %d is in ccm.keys but not in ccm.chunks", ccm.keys[i+1])
		}

		if chunk.Next == 0 {
			if chunk.Ts+chunk.Itgen.Span == nextChunk.Ts {
				return fmt.Errorf("Next of chunk at ts %d is set to 0, but the next chunk is present", ts)
			}
		} else {
			if chunk.Next != nextChunk.Ts {
				return fmt.Errorf("Next of chunk at ts %d is set to %d, but it should be %d", ts, chunk.Next, nextChunk.Ts)
			}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package controller

import (
	"fmt"
	"time"

	"github.com/giantswarm/inago/task"
)

type updateJobDesc struct {
	Request Request
}

func (c controller) getNumRunningSlices(req Request) (int, error) {
	unitStatusList, err := c.groupStatusWithValidate(req)
	if err != nil {
		return 0, maskAny(err)
	}

	var numRunning int
	for _, us := range unitStatusList {
		ok, err := unitHasStatus(us, StatusRunning)
		if err != nil {
			return 0, maskAny(err)
		}
		if ok {
			numRunning++
		}
	}

	return numRunning, nil
}

func (c controller) createJobQueues(req Request, opts UpdateOptions) (chan updateJobDesc, chan updateJobDesc, error) {
	addQueue := make(chan updateJobDesc, len(req.SliceIDs))
	removeQueue := make(chan updateJobDesc, len(req.SliceIDs))

	numRunning, err := c.getNumRunningSlices(req)
	if err != nil {
		return nil, nil, maskAny(err)
	}
	numAllowedToRemove := numRunning - opts.MinAlive

	var numRemovalQueued int
	for _, sliceID := range req.SliceIDs {
		jobDesc := updateJobDesc{
			Request: Request{
				Group:    req.Group,
				SliceIDs: []string{sliceID},
				Units:    req.Units,
			},
		}

		addQueue <- jobDesc

		if numAllowedToRemove >= numRemovalQueued {
			removeQueue <- jobDesc
			numRemovalQueued++
		}
	}

	return addQueue, removeQueue, nil
}

func (c controller) isGroupRemovalAllowed(req Request, opts UpdateOptions) (bool, error) {
	numRunning, err := c.getNumRunningSlices(req)
	if err != nil {
		return false, maskAny(err)
	}

	if numRunning > opts.MinAlive {
		return true, nil
	}

	return false, nil
}

func (c controller) isGroupAdditionAllowed(req Request, numTotal int, opts UpdateOptions) (bool, error) {
	numRunning, err := c.getNumRunningSlices(req)
	if err != nil {
		return false, maskAny(err)
	}

	if numRunning < numTotal+opts.MaxGrowth {
		return true, nil
	}

	return false, nil
}

func (c controller) addWorker(jobDesc updateJobDesc, removeQueue chan updateJobDesc, fail chan<- error) {
	// Submit.
	taskObject, err := c.Submit(jobDesc.Request)
	if err != nil {
		fail <- maskAny(err)
		return
	}
	closer := make(<-chan struct{})
	taskObject, err = c.WaitForTask(taskObject.ID, closer)
	if err != nil {
		fail <- maskAny(err)
		return
	}
	if task.HasFailedStatus(taskObject) {
		fail <- maskAny(fmt.Errorf(taskObject.Error))
		return
	}

	// Start.
	taskObject, err = c.Start(jobDesc.Request)
	if err != nil {
		fail <- maskAny(err)
		return
	}
	closer = make(<-chan struct{})
	taskObject, err = c.WaitForTask(taskObject.ID, closer)
	if err != nil {
		fail <- maskAny(err)
		return
	}
	if task.HasFailedStatus(taskObject) {
		fail <- maskAny(fmt.Errorf(taskObject.Error))
		return
	}

	removeQueue <- jobDesc
}

func (c controller) removeWorker(jobDesc updateJobDesc, fail chan<- error) {
	// Stop.
	taskObject, err := c.Stop(jobDesc.Request)
	if err != nil {
		fail <- maskAny(err)
		return
	}
	closer := make(<-chan struct{})
	taskObject, err = c.WaitForTask(taskObject.ID, closer)
	if err != nil {
		fail <- maskAny(err)
		return
	}
	if task.HasFailedStatus(taskObject) {
		fail <- maskAny(fmt.Errorf(taskObject.Error))
		return
	}

	// Destroy.
	taskObject, err = c.Destroy(jobDesc.Request)
	if err != nil {
		fail <- maskAny(err)
		return
	}
	closer = make(<-chan struct{})
	taskObject, err = c.WaitForTask(taskObject.ID, closer)
	if err != nil {
		fail <- maskAny(err)
		return
	}
	if task.HasFailedStatus(taskObject) {
		fail <- maskAny(fmt.Errorf(taskObject.Error))
		return
	}
}

// UpdateOptions represents the options defining the strategy of an update
// process. Lets have a look at how the update process of 3 group slices would
// look like using the given options.
//
//     TODO I am not that happy with this visualization. Improving it? Removing it?
//
//     MaxGrowth    1
//     MinAlive     2
//     ReadySecs    30
//
//     @1 (running)  ->  @1 (stopped/destroyed)
//     @2 (running)  ->  @2 (running)            ->  @2 (stopped/destroyed)
//     @3 (running)  ->  @3 (running)            ->  @3 (running)            ->  @3 (stopped/destroyed)
//                   ->  @1 (submitted/running)  ->  @1 (running)            ->  @1 (running)            ->  @1 (running)
//                                               ->  @2 (submitted/running)  ->  @2 (running)            ->  @2 (running)
//                                                                           ->  @3 (submitted/running)  ->  @3 (running)
//
type UpdateOptions struct {
	// MaxGrowth represents the number of groups allowed to be added at a given
	// time. No more than MaxGrowth groups will be added at the same time during
	// the update process.
	MaxGrowth int

	// MinAlive represents the number of groups required to stay healthy during
	// the update process. No more than MinAlive groups will be removed at the
	// same time during the update process.
	MinAlive int

	// ReadySecs represents the number of seconds required to wait before ending
	// the update process of one group and starting the update process of another
	// group. This is basically a cool down where the update process sleeps
	// before updating the next group.
	ReadySecs int
}

func (c controller) UpdateWithStrategy(req Request, opts UpdateOptions) error {
	done := make(chan struct{}, 1)
	fail := make(chan error, 1)
	addQueue, removeQueue, err := c.createJobQueues(req, opts)
	if err != nil {
		return maskAny(err)
	}
	numTotal := len(req.SliceIDs)

	for {
		select {
		case jobDesc := <-addQueue:
			go func(jobDesc updateJobDesc) {
				for {
					ok, err := c.isGroupAdditionAllowed(req, numTotal, opts)
					if err != nil {
						fail <- maskAny(err)
						return
					}
					if !ok {
						time.Sleep(c.WaitSleep)
					}
					break
				}

				c.addWorker(jobDesc, removeQueue, fail)
			}(jobDesc)
		case jobDesc := <-removeQueue:
			go func(jobDesc updateJobDesc) {
				for {
					ok, err := c.isGroupRemovalAllowed(req, opts)
					if err != nil {
						fail <- maskAny(err)
						return
					}
					if !ok {
						time.Sleep(c.WaitSleep)
					}
					break
				}

				c.removeWorker(jobDesc, fail)
			}(jobDesc)
		case err := <-fail:
			close(done)
			close(addQueue)
			close(removeQueue)

			return maskAny(err)
		case <-done:
			return nil
		case <-time.After(c.WaitTimeout):
			return maskAny(waitTimeoutReachedError)
		}
	}

	return nil
}
<commit_msg>make `make lint` happy<commit_after>package controller

import (
	"fmt"
	"time"

	"github.com/giantswarm/inago/task"
)

type updateJobDesc struct {
	Request Request
}

func (c controller) getNumRunningSlices(req Request) (int, error) {
	unitStatusList, err := c.groupStatusWithValidate(req)
	if err != nil {
		return 0, maskAny(err)
	}

	var numRunning int
	for _, us := range unitStatusList {
		ok, err := unitHasStatus(us, StatusRunning)
		if err != nil {
			return 0, maskAny(err)
		}
		if ok {
			numRunning++
		}
	}

	return numRunning, nil
}

func (c controller) createJobQueues(req Request, opts UpdateOptions) (chan updateJobDesc, chan updateJobDesc, error) {
	addQueue := make(chan updateJobDesc, len(req.SliceIDs))
	removeQueue := make(chan updateJobDesc, len(req.SliceIDs))

	numRunning, err := c.getNumRunningSlices(req)
	if err != nil {
		return nil, nil, maskAny(err)
	}
	numAllowedToRemove := numRunning - opts.MinAlive

	var numRemovalQueued int
	for _, sliceID := range req.SliceIDs {
		jobDesc := updateJobDesc{
			Request: Request{
				Group:    req.Group,
				SliceIDs: []string{sliceID},
				Units:    req.Units,
			},
		}

		addQueue <- jobDesc

		if numAllowedToRemove >= numRemovalQueued {
			removeQueue <- jobDesc
			numRemovalQueued++
		}
	}

	return addQueue, removeQueue, nil
}

func (c controller) isGroupRemovalAllowed(req Request, opts UpdateOptions) (bool, error) {
	numRunning, err := c.getNumRunningSlices(req)
	if err != nil {
		return false, maskAny(err)
	}

	if numRunning > opts.MinAlive {
		return true, nil
	}

	return false, nil
}

func (c controller) isGroupAdditionAllowed(req Request, numTotal int, opts UpdateOptions) (bool, error) {
	numRunning, err := c.getNumRunningSlices(req)
	if err != nil {
		return false, maskAny(err)
	}

	if numRunning < numTotal+opts.MaxGrowth {
		return true, nil
	}

	return false, nil
}

func (c controller) addWorker(jobDesc updateJobDesc, removeQueue chan updateJobDesc, fail chan<- error) {
	// Submit.
	taskObject, err := c.Submit(jobDesc.Request)
	if err != nil {
		fail <- maskAny(err)
		return
	}
	closer := make(<-chan struct{})
	taskObject, err = c.WaitForTask(taskObject.ID, closer)
	if err != nil {
		fail <- maskAny(err)
		return
	}
	if task.HasFailedStatus(taskObject) {
		fail <- maskAny(fmt.Errorf(taskObject.Error))
		return
	}

	// Start.
	taskObject, err = c.Start(jobDesc.Request)
	if err != nil {
		fail <- maskAny(err)
		return
	}
	closer = make(<-chan struct{})
	taskObject, err = c.WaitForTask(taskObject.ID, closer)
	if err != nil {
		fail <- maskAny(err)
		return
	}
	if task.HasFailedStatus(taskObject) {
		fail <- maskAny(fmt.Errorf(taskObject.Error))
		return
	}

	removeQueue <- jobDesc
}

func (c controller) removeWorker(jobDesc updateJobDesc, fail chan<- error) {
	// Stop.
	taskObject, err := c.Stop(jobDesc.Request)
	if err != nil {
		fail <- maskAny(err)
		return
	}
	closer := make(<-chan struct{})
	taskObject, err = c.WaitForTask(taskObject.ID, closer)
	if err != nil {
		fail <- maskAny(err)
		return
	}
	if task.HasFailedStatus(taskObject) {
		fail <- maskAny(fmt.Errorf(taskObject.Error))
		return
	}

	// Destroy.
	taskObject, err = c.Destroy(jobDesc.Request)
	if err != nil {
		fail <- maskAny(err)
		return
	}
	closer = make(<-chan struct{})
	taskObject, err = c.WaitForTask(taskObject.ID, closer)
	if err != nil {
		fail <- maskAny(err)
		return
	}
	if task.HasFailedStatus(taskObject) {
		fail <- maskAny(fmt.Errorf(taskObject.Error))
		return
	}
}

// UpdateOptions represents the options defining the strategy of an update
// process. Lets have a look at how the update process of 3 group slices would
// look like using the given options.
//
//     TODO I am not that happy with this visualization. Improving it? Removing it?
//
//     MaxGrowth    1
//     MinAlive     2
//     ReadySecs    30
//
//     @1 (running)  ->  @1 (stopped/destroyed)
//     @2 (running)  ->  @2 (running)            ->  @2 (stopped/destroyed)
//     @3 (running)  ->  @3 (running)            ->  @3 (running)            ->  @3 (stopped/destroyed)
//                   ->  @1 (submitted/running)  ->  @1 (running)            ->  @1 (running)            ->  @1 (running)
//                                               ->  @2 (submitted/running)  ->  @2 (running)            ->  @2 (running)
//                                                                           ->  @3 (submitted/running)  ->  @3 (running)
//
type UpdateOptions struct {
	// MaxGrowth represents the number of groups allowed to be added at a given
	// time. No more than MaxGrowth groups will be added at the same time during
	// the update process.
	MaxGrowth int

	// MinAlive represents the number of groups required to stay healthy during
	// the update process. No more than MinAlive groups will be removed at the
	// same time during the update process.
	MinAlive int

	// ReadySecs represents the number of seconds required to wait before ending
	// the update process of one group and starting the update process of another
	// group. This is basically a cool down where the update process sleeps
	// before updating the next group.
	ReadySecs int
}

func (c controller) UpdateWithStrategy(req Request, opts UpdateOptions) error {
	done := make(chan struct{}, 1)
	fail := make(chan error, 1)
	addQueue, removeQueue, err := c.createJobQueues(req, opts)
	if err != nil {
		return maskAny(err)
	}
	numTotal := len(req.SliceIDs)

	for {
		select {
		case jobDesc := <-addQueue:
			go func(jobDesc updateJobDesc) {
				for {
					ok, err := c.isGroupAdditionAllowed(req, numTotal, opts)
					if err != nil {
						fail <- maskAny(err)
						return
					}
					if !ok {
						time.Sleep(c.WaitSleep)
					}
					break
				}

				c.addWorker(jobDesc, removeQueue, fail)
			}(jobDesc)
		case jobDesc := <-removeQueue:
			go func(jobDesc updateJobDesc) {
				for {
					ok, err := c.isGroupRemovalAllowed(req, opts)
					if err != nil {
						fail <- maskAny(err)
						return
					}
					if !ok {
						time.Sleep(c.WaitSleep)
					}
					break
				}

				c.removeWorker(jobDesc, fail)
			}(jobDesc)
		case err := <-fail:
			close(done)
			close(addQueue)
			close(removeQueue)

			return maskAny(err)
		case <-done:
			return nil
		case <-time.After(c.WaitTimeout):
			return maskAny(waitTimeoutReachedError)
		}
	}
}
<|endoftext|>