<commit_before>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2016 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package main

import (
	"bytes"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/loadimpact/k6/api/v1"
	"github.com/manyminds/api2go/jsonapi"
	"gopkg.in/urfave/cli.v1"
	"io/ioutil"
	"net/http"
)

var commandStatus = cli.Command{
	Name:      "status",
	Usage:     "Looks up the status of a running test",
	ArgsUsage: " ",
	Action:    actionStatus,
	Description: `Status will print the status of a running test to stdout in YAML format.

   Use the global --address/-a flag to specify the host to connect to; the
   default is port 6565 on the local machine.

   Endpoint: /v1/status`,
}

var commandStats = cli.Command{
	Name:      "stats",
	Usage:     "Prints stats for a running test",
	ArgsUsage: " ",
	Action:    actionStats,
	Description: `Stats will print metrics about a running test to stdout in YAML format.

   The result is a dictionary of metrics, in no particular order.

   Endpoint: /v1/metrics`,
}

var commandScale = cli.Command{
	Name:      "scale",
	Usage:     "Scales a running test",
	ArgsUsage: "vus",
	Flags: []cli.Flag{
		cli.Int64Flag{
			Name:  "vus, u",
			Usage: "update the number of running VUs",
		},
		cli.Int64Flag{
			Name:  "max, m",
			Usage: "update the max number of VUs allowed",
		},
	},
	Action: actionScale,
	Description: `Scale will change the number of active VUs of a running test.

   It is an error to scale a test beyond vus-max; this is because instantiating
   new VUs is a very expensive operation, which may skew test results if done
   during a running test. To raise vus-max, use --max/-m.

   Endpoint: /v1/status`,
}

var commandPause = cli.Command{
	Name:      "pause",
	Usage:     "Pauses a running test",
	ArgsUsage: " ",
	Action:    actionPause,
	Description: `Pause pauses a running test.

   Running VUs will finish their current iterations, then suspend themselves
   until woken by the test's resumption. A sleeping VU will consume no CPU
   cycles, but will still occupy memory.

   Endpoint: /v1/status`,
}

var commandStart = cli.Command{
	Name:      "start",
	Usage:     "Starts a paused test",
	ArgsUsage: " ",
	Action:    actionStart,
	Description: `Start starts a paused test.

   This is the opposite of the pause command, and will do nothing to an already
   running test.

   Endpoint: /v1/status`,
}

func endpointURL(cc *cli.Context, endpoint string) string {
	return fmt.Sprintf("http://%s/%s", cc.GlobalString("address"), endpoint)
}

func actionStatus(cc *cli.Context) error {
	res, err := http.Get(endpointURL(cc, "/v1/status"))
	if err != nil {
		log.WithError(err).Error("Request error")
		return err
	}
	defer res.Body.Close()
	data, err := ioutil.ReadAll(res.Body)
	if err != nil {
		log.WithError(err).Error("Couldn't read response")
		return err
	}
	var status v1.Status
	if err := jsonapi.Unmarshal(data, &status); err != nil {
		log.WithError(err).Error("Invalid response")
		return err
	}
	return dumpYAML(status)
}

func actionStats(cc *cli.Context) error {
	res, err := http.Get(endpointURL(cc, "/v1/metrics"))
	if err != nil {
		log.WithError(err).Error("Request error")
		return err
	}
	defer res.Body.Close()
	data, err := ioutil.ReadAll(res.Body)
	if err != nil {
		log.WithError(err).Error("Couldn't read response")
		return err
	}
	var metrics []v1.Metric
	if err := jsonapi.Unmarshal(data, &metrics); err != nil {
		log.WithError(err).Error("Invalid response")
		return err
	}
	output := make(map[string]v1.Metric)
	for _, m := range metrics {
		output[m.GetID()] = m
	}
	return dumpYAML(output)
}

func actionScale(cc *cli.Context) error {
	patch := v1.Status{
		VUs:    cliInt64(cc, "vus"),
		VUsMax: cliInt64(cc, "max"),
	}
	if !patch.VUs.Valid && !patch.VUsMax.Valid {
		log.Warn("Neither --vus/-u or --max/-m passed; doing doing nothing")
		return nil
	}

	body, err := jsonapi.Marshal(patch)
	if err != nil {
		log.WithError(err).Error("Serialization error")
		return err
	}

	req, err := http.NewRequest(
		http.MethodPatch,
		endpointURL(cc, "/v1/status"),
		bytes.NewReader(body),
	)
	if err != nil {
		log.WithError(err).Error("Couldn't create request")
		return err
	}

	res, err := http.DefaultClient.Do(req)
	if err != nil {
		log.WithError(err).Error("Request error")
		return err
	}
	defer res.Body.Close()

	data, err := ioutil.ReadAll(res.Body)
	if err != nil {
		log.WithError(err).Error("Couldn't read response")
		return err
	}
	var status v1.Status
	if err := jsonapi.Unmarshal(data, &status); err != nil {
		log.WithError(err).Error("Invalid response")
		return err
	}
	return dumpYAML(status)
}

func actionPause(cc *cli.Context) error {
	return nil
}

func actionStart(cc *cli.Context) error {
	return nil
}
<commit_msg>[fix] No duplicate /'s in URLs<commit_after>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2016 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package main

import (
	"bytes"
	"fmt"
	log "github.com/Sirupsen/logrus"
	"github.com/loadimpact/k6/api/v1"
	"github.com/manyminds/api2go/jsonapi"
	"gopkg.in/urfave/cli.v1"
	"io/ioutil"
	"net/http"
)

var commandStatus = cli.Command{
	Name:      "status",
	Usage:     "Looks up the status of a running test",
	ArgsUsage: " ",
	Action:    actionStatus,
	Description: `Status will print the status of a running test to stdout in YAML format.

   Use the global --address/-a flag to specify the host to connect to; the
   default is port 6565 on the local machine.

   Endpoint: /v1/status`,
}

var commandStats = cli.Command{
	Name:      "stats",
	Usage:     "Prints stats for a running test",
	ArgsUsage: " ",
	Action:    actionStats,
	Description: `Stats will print metrics about a running test to stdout in YAML format.

   The result is a dictionary of metrics, in no particular order.

   Endpoint: /v1/metrics`,
}

var commandScale = cli.Command{
	Name:      "scale",
	Usage:     "Scales a running test",
	ArgsUsage: "vus",
	Flags: []cli.Flag{
		cli.Int64Flag{
			Name:  "vus, u",
			Usage: "update the number of running VUs",
		},
		cli.Int64Flag{
			Name:  "max, m",
			Usage: "update the max number of VUs allowed",
		},
	},
	Action: actionScale,
	Description: `Scale will change the number of active VUs of a running test.

   It is an error to scale a test beyond vus-max; this is because instantiating
   new VUs is a very expensive operation, which may skew test results if done
   during a running test. To raise vus-max, use --max/-m.

   Endpoint: /v1/status`,
}

var commandPause = cli.Command{
	Name:      "pause",
	Usage:     "Pauses a running test",
	ArgsUsage: " ",
	Action:    actionPause,
	Description: `Pause pauses a running test.

   Running VUs will finish their current iterations, then suspend themselves
   until woken by the test's resumption. A sleeping VU will consume no CPU
   cycles, but will still occupy memory.

   Endpoint: /v1/status`,
}

var commandStart = cli.Command{
	Name:      "start",
	Usage:     "Starts a paused test",
	ArgsUsage: " ",
	Action:    actionStart,
	Description: `Start starts a paused test.

   This is the opposite of the pause command, and will do nothing to an already
   running test.

   Endpoint: /v1/status`,
}

func endpointURL(cc *cli.Context, endpoint string) string {
	return fmt.Sprintf("http://%s%s", cc.GlobalString("address"), endpoint)
}

func actionStatus(cc *cli.Context) error {
	res, err := http.Get(endpointURL(cc, "/v1/status"))
	if err != nil {
		log.WithError(err).Error("Request error")
		return err
	}
	defer res.Body.Close()
	data, err := ioutil.ReadAll(res.Body)
	if err != nil {
		log.WithError(err).Error("Couldn't read response")
		return err
	}
	var status v1.Status
	if err := jsonapi.Unmarshal(data, &status); err != nil {
		log.WithError(err).Error("Invalid response")
		return err
	}
	return dumpYAML(status)
}

func actionStats(cc *cli.Context) error {
	res, err := http.Get(endpointURL(cc, "/v1/metrics"))
	if err != nil {
		log.WithError(err).Error("Request error")
		return err
	}
	defer res.Body.Close()
	data, err := ioutil.ReadAll(res.Body)
	if err != nil {
		log.WithError(err).Error("Couldn't read response")
		return err
	}
	var metrics []v1.Metric
	if err := jsonapi.Unmarshal(data, &metrics); err != nil {
		log.WithError(err).Error("Invalid response")
		return err
	}
	output := make(map[string]v1.Metric)
	for _, m := range metrics {
		output[m.GetID()] = m
	}
	return dumpYAML(output)
}

func actionScale(cc *cli.Context) error {
	patch := v1.Status{
		VUs:    cliInt64(cc, "vus"),
		VUsMax: cliInt64(cc, "max"),
	}
	if !patch.VUs.Valid && !patch.VUsMax.Valid {
		log.Warn("Neither --vus/-u or --max/-m passed; doing doing nothing")
		return nil
	}

	body, err := jsonapi.Marshal(patch)
	if err != nil {
		log.WithError(err).Error("Serialization error")
		return err
	}

	req, err := http.NewRequest(
		http.MethodPatch,
		endpointURL(cc, "/v1/status"),
		bytes.NewReader(body),
	)
	if err != nil {
		log.WithError(err).Error("Couldn't create request")
		return err
	}

	res, err := http.DefaultClient.Do(req)
	if err != nil {
		log.WithError(err).Error("Request error")
		return err
	}
	defer res.Body.Close()

	data, err := ioutil.ReadAll(res.Body)
	if err != nil {
		log.WithError(err).Error("Couldn't read response")
		return err
	}
	var status v1.Status
	if err := jsonapi.Unmarshal(data, &status); err != nil {
		log.WithError(err).Error("Invalid response")
		return err
	}
	return dumpYAML(status)
}

func actionPause(cc *cli.Context) error {
	return nil
}

func actionStart(cc *cli.Context) error {
	return nil
}
<|endoftext|>


<commit_before>package awstasks

import (
	"fmt"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/golang/glog"
	"github.com/kopeio/aws-controller/pkg/kope/utils"
	"k8s.io/kops/upup/pkg/fi"
	"k8s.io/kops/upup/pkg/fi/cloudup/awsup"
	"k8s.io/kops/upup/pkg/fi/cloudup/terraform"
)

//go:generate fitask -type=SecurityGroup
type SecurityGroup struct {
	Name *string

	ID          *string
	Description *string
	VPC         *VPC

	RemoveExtraRules *bool
}

var _ fi.CompareWithID = &SecurityGroup{}
var _ fi.ProducesDeletions = &SecurityGroup{}

func (e *SecurityGroup) CompareWithID() *string {
	return e.ID
}

func (e *SecurityGroup) Find(c *fi.Context) (*SecurityGroup, error) {
	sg, err := e.findEc2(c)
	if err != nil {
		return nil, err
	}
	if sg == nil {
		return nil, nil
	}
	actual := &SecurityGroup{
		ID:          sg.GroupId,
		Name:        sg.GroupName,
		Description: sg.Description,
		VPC:         &VPC{ID: sg.VpcId},
	}

	glog.V(2).Infof("found matching SecurityGroup %q", *actual.ID)
	e.ID = actual.ID

	actual.RemoveExtraRules = e.RemoveExtraRules

	return actual, nil
}

func (e *SecurityGroup) findEc2(c *fi.Context) (*ec2.SecurityGroup, error) {
	cloud := c.Cloud.(*awsup.AWSCloud)

	var vpcID *string
	if e.VPC != nil {
		vpcID = e.VPC.ID
	}

	if vpcID == nil {
		return nil, nil
	}

	request := &ec2.DescribeSecurityGroupsInput{}

	if fi.StringValue(e.ID) != "" {
		request.GroupIds = []*string{e.ID}
	} else {
		filters := cloud.BuildFilters(e.Name)
		filters = append(filters, awsup.NewEC2Filter("vpc-id", *vpcID))
		filters = append(filters, awsup.NewEC2Filter("group-name", *e.Name))

		request.Filters = cloud.BuildFilters(e.Name)
	}

	response, err := cloud.EC2.DescribeSecurityGroups(request)
	if err != nil {
		return nil, fmt.Errorf("error listing SecurityGroups: %v", err)
	}
	if response == nil || len(response.SecurityGroups) == 0 {
		return nil, nil
	}

	if len(response.SecurityGroups) != 1 {
		return nil, fmt.Errorf("found multiple SecurityGroups matching tags")
	}
	sg := response.SecurityGroups[0]
	return sg, nil
}

func (e *SecurityGroup) Run(c *fi.Context) error {
	return fi.DefaultDeltaRunMethod(e, c)
}

func (_ *SecurityGroup) CheckChanges(a, e, changes *SecurityGroup) error {
	if a != nil {
		if changes.ID != nil {
			return fi.CannotChangeField("ID")
		}
		if changes.Name != nil {
			return fi.CannotChangeField("Name")
		}
		if changes.VPC != nil {
			return fi.CannotChangeField("VPC")
		}
	}
	return nil
}

func (_ *SecurityGroup) RenderAWS(t *awsup.AWSAPITarget, a, e, changes *SecurityGroup) error {
	if a == nil {
		glog.V(2).Infof("Creating SecurityGroup with Name:%q VPC:%q", *e.Name, *e.VPC.ID)

		request := &ec2.CreateSecurityGroupInput{
			VpcId:       e.VPC.ID,
			GroupName:   e.Name,
			Description: e.Description,
		}

		response, err := t.Cloud.EC2.CreateSecurityGroup(request)
		if err != nil {
			return fmt.Errorf("error creating SecurityGroup: %v", err)
		}

		e.ID = response.GroupId
	}

	return t.AddAWSTags(*e.ID, t.Cloud.BuildTags(e.Name))
}

type terraformSecurityGroup struct {
	Name        *string            `json:"name"`
	VPCID       *terraform.Literal `json:"vpc_id"`
	Description *string            `json:"description"`
	Tags        map[string]string  `json:"tags,omitempty"`
}

func (_ *SecurityGroup) RenderTerraform(t *terraform.TerraformTarget, a, e, changes *SecurityGroup) error {
	cloud := t.Cloud.(*awsup.AWSCloud)

	tf := &terraformSecurityGroup{
		Name:        e.Name,
		VPCID:       e.VPC.TerraformLink(),
		Description: e.Description,
		Tags:        cloud.BuildTags(e.Name),
	}

	return t.RenderResource("aws_security_group", *e.Name, tf)
}

func (e *SecurityGroup) TerraformLink() *terraform.Literal {
	return terraform.LiteralProperty("aws_security_group", *e.Name, "id")
}

type deleteSecurityGroupRule struct {
	groupID    *string
	permission *ec2.IpPermission
	egress     bool
}

var _ fi.Deletion = &deleteSecurityGroupRule{}

func (d *deleteSecurityGroupRule) Delete(t fi.Target) error {
	glog.V(2).Infof("deleting security group permission: %v", utils.DebugString(d.permission))

	awsTarget, ok := t.(*awsup.AWSAPITarget)
	if !ok {
		return fmt.Errorf("unexpected target type for deletion: %T", t)
	}

	if d.egress {
		request := &ec2.RevokeSecurityGroupEgressInput{
			GroupId: d.groupID,
		}
		request.IpPermissions = []*ec2.IpPermission{d.permission}

		glog.V(2).Infof("Calling EC2 RevokeSecurityGroupEgress")
		_, err := awsTarget.Cloud.EC2.RevokeSecurityGroupEgress(request)
		if err != nil {
			return fmt.Errorf("error revoking SecurityGroupEgress: %v", err)
		}
	} else {
		request := &ec2.RevokeSecurityGroupIngressInput{
			GroupId: d.groupID,
		}
		request.IpPermissions = []*ec2.IpPermission{d.permission}

		glog.V(2).Infof("Calling EC2 RevokeSecurityGroupIngress")
		_, err := awsTarget.Cloud.EC2.RevokeSecurityGroupIngress(request)
		if err != nil {
			return fmt.Errorf("error revoking SecurityGroupIngress: %v", err)
		}
	}

	return nil
}

func (d *deleteSecurityGroupRule) TaskName() string {
	return "SecurityGroupRule"
}

func (d *deleteSecurityGroupRule) Item() string {
	s := fi.StringValue(d.groupID) + ":"
	p := d.permission
	if aws.Int64Value(p.FromPort) != 0 {
		s += fmt.Sprintf(" port=%d", aws.Int64Value(p.FromPort))
		if aws.Int64Value(p.ToPort) != aws.Int64Value(p.FromPort) {
			s += fmt.Sprintf("-%d", aws.Int64Value(p.ToPort))
		}
	}
	if aws.StringValue(p.IpProtocol) != "-1" {
		s += fmt.Sprintf(" protocol=%s", aws.StringValue(p.IpProtocol))
	}
	for _, ug := range p.UserIdGroupPairs {
		s += fmt.Sprintf(" group=%s", aws.StringValue(ug.GroupId))
	}
	for _, r := range p.IpRanges {
		s += fmt.Sprintf(" ip=%s", aws.StringValue(r.CidrIp))
	}
	//permissionString := utils.DebugString(d.permission)
	//s += permissionString

	return s
}

func expandPermissions(sgID *string, permission *ec2.IpPermission, egress bool) []*ec2.IpPermission {
	var rules []*ec2.IpPermission

	master := &ec2.IpPermission{
		FromPort:   permission.FromPort,
		ToPort:     permission.ToPort,
		IpProtocol: permission.IpProtocol,
	}

	for _, ipRange := range permission.IpRanges {
		a := &ec2.IpPermission{}
		*a = *master
		a.IpRanges = []*ec2.IpRange{ipRange}
		rules = append(rules, a)
	}

	for _, ug := range permission.UserIdGroupPairs {
		a := &ec2.IpPermission{}
		*a = *master
		a.UserIdGroupPairs = []*ec2.UserIdGroupPair{ug}
		rules = append(rules, a)
	}

	if len(rules) == 0 {
		// If there are no group or cidr restrictions, it is just a generic rule
		rules = append(rules, master)
	}

	return rules
}

func (e *SecurityGroup) FindDeletions(c *fi.Context) ([]fi.Deletion, error) {
	var removals []fi.Deletion

	if fi.BoolValue(e.RemoveExtraRules) != true {
		return nil, nil
	}

	sg, err := e.findEc2(c)
	if err != nil {
		return nil, err
	}
	if sg == nil {
		return nil, nil
	}

	var ingress []*ec2.IpPermission
	for _, permission := range sg.IpPermissions {
		rules := expandPermissions(sg.GroupId, permission, false)
		ingress = append(ingress, rules...)
	}

	for _, permission := range ingress {
		found := false
		for _, t := range c.AllTasks() {
			er, ok := t.(*SecurityGroupRule)
			if !ok {
				continue
			}
			if er.matches(permission) {
				found = true
			}
		}
		if !found {
			removals = append(removals, &deleteSecurityGroupRule{
				groupID:    sg.GroupId,
				permission: permission,
				egress:     false,
			})
		}
	}

	var egress []*ec2.IpPermission
	for _, permission := range sg.IpPermissionsEgress {
		rules := expandPermissions(sg.GroupId, permission, true)
		egress = append(egress, rules...)
	}
	for _, permission := range egress {
		found := false
		for _, t := range c.AllTasks() {
			er, ok := t.(*SecurityGroupRule)
			if !ok {
				continue
			}
			if er.matches(permission) {
				found = true
			}
		}
		if !found {
			removals = append(removals, &deleteSecurityGroupRule{
				groupID:    sg.GroupId,
				permission: permission,
				egress:     true,
			})
		}
	}

	return removals, nil
}
<commit_msg>Fix incorrect import<commit_after>package awstasks

import (
	"fmt"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/golang/glog"
	"k8s.io/kops/upup/pkg/fi"
	"k8s.io/kops/upup/pkg/fi/cloudup/awsup"
	"k8s.io/kops/upup/pkg/fi/cloudup/terraform"
)

//go:generate fitask -type=SecurityGroup
type SecurityGroup struct {
	Name *string

	ID          *string
	Description *string
	VPC         *VPC

	RemoveExtraRules *bool
}

var _ fi.CompareWithID = &SecurityGroup{}
var _ fi.ProducesDeletions = &SecurityGroup{}

func (e *SecurityGroup) CompareWithID() *string {
	return e.ID
}

func (e *SecurityGroup) Find(c *fi.Context) (*SecurityGroup, error) {
	sg, err := e.findEc2(c)
	if err != nil {
		return nil, err
	}
	if sg == nil {
		return nil, nil
	}
	actual := &SecurityGroup{
		ID:          sg.GroupId,
		Name:        sg.GroupName,
		Description: sg.Description,
		VPC:         &VPC{ID: sg.VpcId},
	}

	glog.V(2).Infof("found matching SecurityGroup %q", *actual.ID)
	e.ID = actual.ID

	actual.RemoveExtraRules = e.RemoveExtraRules

	return actual, nil
}

func (e *SecurityGroup) findEc2(c *fi.Context) (*ec2.SecurityGroup, error) {
	cloud := c.Cloud.(*awsup.AWSCloud)

	var vpcID *string
	if e.VPC != nil {
		vpcID = e.VPC.ID
	}

	if vpcID == nil {
		return nil, nil
	}

	request := &ec2.DescribeSecurityGroupsInput{}

	if fi.StringValue(e.ID) != "" {
		request.GroupIds = []*string{e.ID}
	} else {
		filters := cloud.BuildFilters(e.Name)
		filters = append(filters, awsup.NewEC2Filter("vpc-id", *vpcID))
		filters = append(filters, awsup.NewEC2Filter("group-name", *e.Name))

		request.Filters = cloud.BuildFilters(e.Name)
	}

	response, err := cloud.EC2.DescribeSecurityGroups(request)
	if err != nil {
		return nil, fmt.Errorf("error listing SecurityGroups: %v", err)
	}
	if response == nil || len(response.SecurityGroups) == 0 {
		return nil, nil
	}

	if len(response.SecurityGroups) != 1 {
		return nil, fmt.Errorf("found multiple SecurityGroups matching tags")
	}
	sg := response.SecurityGroups[0]
	return sg, nil
}

func (e *SecurityGroup) Run(c *fi.Context) error {
	return fi.DefaultDeltaRunMethod(e, c)
}

func (_ *SecurityGroup) CheckChanges(a, e, changes *SecurityGroup) error {
	if a != nil {
		if changes.ID != nil {
			return fi.CannotChangeField("ID")
		}
		if changes.Name != nil {
			return fi.CannotChangeField("Name")
		}
		if changes.VPC != nil {
			return fi.CannotChangeField("VPC")
		}
	}
	return nil
}

func (_ *SecurityGroup) RenderAWS(t *awsup.AWSAPITarget, a, e, changes *SecurityGroup) error {
	if a == nil {
		glog.V(2).Infof("Creating SecurityGroup with Name:%q VPC:%q", *e.Name, *e.VPC.ID)

		request := &ec2.CreateSecurityGroupInput{
			VpcId:       e.VPC.ID,
			GroupName:   e.Name,
			Description: e.Description,
		}

		response, err := t.Cloud.EC2.CreateSecurityGroup(request)
		if err != nil {
			return fmt.Errorf("error creating SecurityGroup: %v", err)
		}

		e.ID = response.GroupId
	}

	return t.AddAWSTags(*e.ID, t.Cloud.BuildTags(e.Name))
}

type terraformSecurityGroup struct {
	Name        *string            `json:"name"`
	VPCID       *terraform.Literal `json:"vpc_id"`
	Description *string            `json:"description"`
	Tags        map[string]string  `json:"tags,omitempty"`
}

func (_ *SecurityGroup) RenderTerraform(t *terraform.TerraformTarget, a, e, changes *SecurityGroup) error {
	cloud := t.Cloud.(*awsup.AWSCloud)

	tf := &terraformSecurityGroup{
		Name:        e.Name,
		VPCID:       e.VPC.TerraformLink(),
		Description: e.Description,
		Tags:        cloud.BuildTags(e.Name),
	}

	return t.RenderResource("aws_security_group", *e.Name, tf)
}

func (e *SecurityGroup) TerraformLink() *terraform.Literal {
	return terraform.LiteralProperty("aws_security_group", *e.Name, "id")
}

type deleteSecurityGroupRule struct {
	groupID    *string
	permission *ec2.IpPermission
	egress     bool
}

var _ fi.Deletion = &deleteSecurityGroupRule{}

func (d *deleteSecurityGroupRule) Delete(t fi.Target) error {
	glog.V(2).Infof("deleting security group permission: %v", fi.DebugAsJsonString(d.permission))

	awsTarget, ok := t.(*awsup.AWSAPITarget)
	if !ok {
		return fmt.Errorf("unexpected target type for deletion: %T", t)
	}

	if d.egress {
		request := &ec2.RevokeSecurityGroupEgressInput{
			GroupId: d.groupID,
		}
		request.IpPermissions = []*ec2.IpPermission{d.permission}

		glog.V(2).Infof("Calling EC2 RevokeSecurityGroupEgress")
		_, err := awsTarget.Cloud.EC2.RevokeSecurityGroupEgress(request)
		if err != nil {
			return fmt.Errorf("error revoking SecurityGroupEgress: %v", err)
		}
	} else {
		request := &ec2.RevokeSecurityGroupIngressInput{
			GroupId: d.groupID,
		}
		request.IpPermissions = []*ec2.IpPermission{d.permission}

		glog.V(2).Infof("Calling EC2 RevokeSecurityGroupIngress")
		_, err := awsTarget.Cloud.EC2.RevokeSecurityGroupIngress(request)
		if err != nil {
			return fmt.Errorf("error revoking SecurityGroupIngress: %v", err)
		}
	}

	return nil
}

func (d *deleteSecurityGroupRule) TaskName() string {
	return "SecurityGroupRule"
}

func (d *deleteSecurityGroupRule) Item() string {
	s := fi.StringValue(d.groupID) + ":"
	p := d.permission
	if aws.Int64Value(p.FromPort) != 0 {
		s += fmt.Sprintf(" port=%d", aws.Int64Value(p.FromPort))
		if aws.Int64Value(p.ToPort) != aws.Int64Value(p.FromPort) {
			s += fmt.Sprintf("-%d", aws.Int64Value(p.ToPort))
		}
	}
	if aws.StringValue(p.IpProtocol) != "-1" {
		s += fmt.Sprintf(" protocol=%s", aws.StringValue(p.IpProtocol))
	}
	for _, ug := range p.UserIdGroupPairs {
		s += fmt.Sprintf(" group=%s", aws.StringValue(ug.GroupId))
	}
	for _, r := range p.IpRanges {
		s += fmt.Sprintf(" ip=%s", aws.StringValue(r.CidrIp))
	}
	//permissionString := fi.DebugAsJsonString(d.permission)
	//s += permissionString

	return s
}

func expandPermissions(sgID *string, permission *ec2.IpPermission, egress bool) []*ec2.IpPermission {
	var rules []*ec2.IpPermission

	master := &ec2.IpPermission{
		FromPort:   permission.FromPort,
		ToPort:     permission.ToPort,
		IpProtocol: permission.IpProtocol,
	}

	for _, ipRange := range permission.IpRanges {
		a := &ec2.IpPermission{}
		*a = *master
		a.IpRanges = []*ec2.IpRange{ipRange}
		rules = append(rules, a)
	}

	for _, ug := range permission.UserIdGroupPairs {
		a := &ec2.IpPermission{}
		*a = *master
		a.UserIdGroupPairs = []*ec2.UserIdGroupPair{ug}
		rules = append(rules, a)
	}

	if len(rules) == 0 {
		// If there are no group or cidr restrictions, it is just a generic rule
		rules = append(rules, master)
	}

	return rules
}

func (e *SecurityGroup) FindDeletions(c *fi.Context) ([]fi.Deletion, error) {
	var removals []fi.Deletion

	if fi.BoolValue(e.RemoveExtraRules) != true {
		return nil, nil
	}

	sg, err := e.findEc2(c)
	if err != nil {
		return nil, err
	}
	if sg == nil {
		return nil, nil
	}

	var ingress []*ec2.IpPermission
	for _, permission := range sg.IpPermissions {
		rules := expandPermissions(sg.GroupId, permission, false)
		ingress = append(ingress, rules...)
	}

	for _, permission := range ingress {
		found := false
		for _, t := range c.AllTasks() {
			er, ok := t.(*SecurityGroupRule)
			if !ok {
				continue
			}
			if er.matches(permission) {
				found = true
			}
		}
		if !found {
			removals = append(removals, &deleteSecurityGroupRule{
				groupID:    sg.GroupId,
				permission: permission,
				egress:     false,
			})
		}
	}

	var egress []*ec2.IpPermission
	for _, permission := range sg.IpPermissionsEgress {
		rules := expandPermissions(sg.GroupId, permission, true)
		egress = append(egress, rules...)
	}
	for _, permission := range egress {
		found := false
		for _, t := range c.AllTasks() {
			er, ok := t.(*SecurityGroupRule)
			if !ok {
				continue
			}
			if er.matches(permission) {
				found = true
			}
		}
		if !found {
			removals = append(removals, &deleteSecurityGroupRule{
				groupID:    sg.GroupId,
				permission: permission,
				egress:     true,
			})
		}
	}

	return removals, nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2020 The Jaeger 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 defaults

import (
	"fmt"
	"sort"
	"testing"

	"github.com/open-telemetry/opentelemetry-collector/config"
	"github.com/open-telemetry/opentelemetry-collector/config/configmodels"
	"github.com/open-telemetry/opentelemetry-collector/exporter/jaegerexporter"
	"github.com/open-telemetry/opentelemetry-collector/processor/resourceprocessor"
	"github.com/open-telemetry/opentelemetry-collector/receiver/jaegerreceiver"
	"github.com/spf13/viper"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.uber.org/zap"

	"github.com/jaegertracing/jaeger/cmd/agent/app/reporter/grpc"
	"github.com/jaegertracing/jaeger/cmd/opentelemetry-collector/app/exporter/cassandra"
	"github.com/jaegertracing/jaeger/cmd/opentelemetry-collector/app/exporter/elasticsearch"
	"github.com/jaegertracing/jaeger/cmd/opentelemetry-collector/app/exporter/grpcplugin"
	"github.com/jaegertracing/jaeger/cmd/opentelemetry-collector/app/exporter/kafka"
	kafkaRec "github.com/jaegertracing/jaeger/cmd/opentelemetry-collector/app/receiver/kafka"
	jConfig "github.com/jaegertracing/jaeger/pkg/config"
	"github.com/jaegertracing/jaeger/ports"
)

func TestDefaultCollectorConfig(t *testing.T) {
	disabledHostPort := ports.PortToHostPort(0)
	tests := []struct {
		storageType    string
		zipkinHostPort string
		exporterTypes  []string
		pipeline       configmodels.Pipelines
		err            string
		config         map[string]interface{}
	}{
		{
			storageType:    "elasticsearch",
			zipkinHostPort: disabledHostPort,
			exporterTypes:  []string{elasticsearch.TypeStr},
			pipeline: configmodels.Pipelines{
				"traces": {
					InputType: configmodels.TracesDataType,
					Receivers: []string{"jaeger"},
					Exporters: []string{elasticsearch.TypeStr},
				},
			},
		},
		{
			storageType:    "cassandra",
			zipkinHostPort: disabledHostPort,
			exporterTypes:  []string{cassandra.TypeStr},
			pipeline: configmodels.Pipelines{
				"traces": {
					InputType: configmodels.TracesDataType,
					Receivers: []string{"jaeger"},
					Exporters: []string{cassandra.TypeStr},
				},
			},
		},
		{
			storageType:    "kafka",
			zipkinHostPort: disabledHostPort,
			exporterTypes:  []string{kafka.TypeStr},
			pipeline: configmodels.Pipelines{
				"traces": {
					InputType: configmodels.TracesDataType,
					Receivers: []string{"jaeger"},
					Exporters: []string{kafka.TypeStr},
				},
			},
		},
		{
			storageType:    "cassandra,elasticsearch,grpc-plugin",
			zipkinHostPort: disabledHostPort,
			exporterTypes:  []string{cassandra.TypeStr, elasticsearch.TypeStr, grpcplugin.TypeStr},
			pipeline: configmodels.Pipelines{
				"traces": {
					InputType: configmodels.TracesDataType,
					Receivers: []string{"jaeger"},
					Exporters: []string{cassandra.TypeStr, elasticsearch.TypeStr, grpcplugin.TypeStr},
				},
			},
		},
		{
			storageType:    "cassandra",
			zipkinHostPort: ":9411",
			exporterTypes:  []string{cassandra.TypeStr},
			config:         map[string]interface{}{"resource.labels": "foo=bar"},
			pipeline: configmodels.Pipelines{
				"traces": {
					InputType:  configmodels.TracesDataType,
					Receivers:  []string{"jaeger", "zipkin"},
					Processors: []string{"resource"},
					Exporters:  []string{cassandra.TypeStr},
				},
			},
		},
		{
			storageType: "floppy",
			err:         "unknown storage type: floppy",
		},
	}
	for _, test := range tests {
		t.Run(test.storageType, func(t *testing.T) {
			v, _ := jConfig.Viperize(grpc.AddFlags)
			factories := Components(v)
			for key, val := range test.config {
				v.Set(key, val)
			}
			cfg, err := CollectorConfig(test.storageType, test.zipkinHostPort, factories)
			if test.err != "" {
				require.Nil(t, cfg)
				assert.EqualError(t, err, test.err)
				return
			}
			require.NoError(t, err)
			require.NoError(t, config.ValidateConfig(cfg, zap.NewNop()))

			assert.Equal(t, 1, len(cfg.Extensions))
			assert.Equal(t, 1, len(cfg.Service.Extensions))
			assert.Equal(t, "health_check", cfg.Service.Extensions[0])
			assert.Equal(t, "health_check", cfg.Extensions["health_check"].Name())
			assert.Equal(t, len(test.pipeline["traces"].Receivers), len(cfg.Receivers))
			assert.Equal(t, "jaeger", cfg.Receivers["jaeger"].Name())
			assert.Equal(t, len(test.exporterTypes), len(cfg.Exporters))

			processorMap := map[string]bool{}
			for _, p := range test.pipeline["traces"].Processors {
				processorMap[p] = true
			}
			if processorMap["resource"] {
				assert.Equal(t, len(processorMap), len(cfg.Processors))
				assert.IsType(t, &resourceprocessor.Config{}, cfg.Processors["resource"])
			}

			types := []string{}
			for _, v := range cfg.Exporters {
				types = append(types, string(v.Type()))
			}
			sort.Strings(types)
			assert.Equal(t, test.exporterTypes, types)
			sort.Strings(cfg.Service.Pipelines["traces"].Receivers)
			assert.EqualValues(t, test.pipeline, cfg.Service.Pipelines)
		})
	}
}

func TestDefaultAgentConfig(t *testing.T) {
	tests := []struct {
		config  map[string]interface{}
		service configmodels.Service
	}{
		{
			config: map[string]interface{}{"resource.labels": "foo=bar"},
			service: configmodels.Service{
				Extensions: []string{"health_check"},
				Pipelines: configmodels.Pipelines{
					"traces": &configmodels.Pipeline{
						InputType:  configmodels.TracesDataType,
						Receivers:  []string{"jaeger"},
						Processors: []string{"resource"},
						Exporters:  []string{"jaeger"},
					},
				},
			},
		},
		{
			service: configmodels.Service{
				Extensions: []string{"health_check"},
				Pipelines: configmodels.Pipelines{
					"traces": &configmodels.Pipeline{
						InputType: configmodels.TracesDataType,
						Receivers: []string{"jaeger"},
						Exporters: []string{"jaeger"},
					},
				},
			},
		},
	}
	for _, test := range tests {
		t.Run(fmt.Sprintf("%v", test.config), func(t *testing.T) {
			v, _ := jConfig.Viperize(grpc.AddFlags)
			for key, val := range test.config {
				v.Set(key, val)
			}
			factories := Components(v)
			cfg := AgentConfig(factories)
			require.NoError(t, config.ValidateConfig(cfg, zap.NewNop()))

			assert.Equal(t, test.service, cfg.Service)
			assert.Equal(t, 1, len(cfg.Receivers))
			assert.IsType(t, &jaegerreceiver.Config{}, cfg.Receivers["jaeger"])
			assert.Equal(t, 1, len(cfg.Exporters))
			assert.IsType(t, &jaegerexporter.Config{}, cfg.Exporters["jaeger"])
			processorMap := map[string]bool{}
			for _, p := range test.service.Pipelines["traces"].Processors {
				processorMap[p] = true
			}
			if processorMap["resource"] {
				assert.Equal(t, len(processorMap), len(cfg.Processors))
				assert.IsType(t, &resourceprocessor.Config{}, cfg.Processors["resource"])
			}
		})
	}
}

func TestDefaultIngesterConfig(t *testing.T) {
	tests := []struct {
		storageType string
		service     configmodels.Service
		err         string
	}{
		{
			storageType: "elasticsearch",
			service: configmodels.Service{
				Extensions: []string{"health_check"},
				Pipelines: configmodels.Pipelines{
					"traces": &configmodels.Pipeline{
						InputType: configmodels.TracesDataType,
						Receivers: []string{kafkaRec.TypeStr},
						Exporters: []string{elasticsearch.TypeStr},
					},
				},
			},
		},
		{
			storageType: "elasticsearch,cassandra,grpc-plugin",
			service: configmodels.Service{
				Extensions: []string{"health_check"},
				Pipelines: configmodels.Pipelines{
					"traces": &configmodels.Pipeline{
						InputType: configmodels.TracesDataType,
						Receivers: []string{kafkaRec.TypeStr},
						Exporters: []string{cassandra.TypeStr, elasticsearch.TypeStr, grpcplugin.TypeStr},
					},
				},
			},
		},
		{
			storageType: "floppy",
			err:         "unknown storage type: floppy",
		},
	}
	for _, test := range tests {
		t.Run(test.storageType, func(t *testing.T) {
			factories := Components(viper.New())
			cfg, err := IngesterConfig(test.storageType, factories)
			if test.err != "" {
				require.Nil(t, cfg)
				assert.EqualError(t, err, test.err)
				return
			}
			require.NoError(t, err)
			require.NoError(t, config.ValidateConfig(cfg, zap.NewNop()))

			sort.Strings(cfg.Service.Pipelines["traces"].Exporters)
			assert.Equal(t, test.service, cfg.Service)
			assert.Equal(t, 1, len(cfg.Receivers))
			assert.IsType(t, &kafkaRec.Config{}, cfg.Receivers[kafkaRec.TypeStr])

			assert.Equal(t, len(test.service.Pipelines["traces"].Exporters), len(cfg.Exporters))
			types := []string{}
			for _, v := range cfg.Exporters {
				types = append(types, string(v.Type()))
			}
			sort.Strings(types)
			assert.Equal(t, test.service.Pipelines["traces"].Exporters, types)
		})
	}
}
<commit_msg>Sort exporters in default config test (#2238)<commit_after>// Copyright (c) 2020 The Jaeger 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 defaults

import (
	"fmt"
	"sort"
	"testing"

	"github.com/open-telemetry/opentelemetry-collector/config"
	"github.com/open-telemetry/opentelemetry-collector/config/configmodels"
	"github.com/open-telemetry/opentelemetry-collector/exporter/jaegerexporter"
	"github.com/open-telemetry/opentelemetry-collector/processor/resourceprocessor"
	"github.com/open-telemetry/opentelemetry-collector/receiver/jaegerreceiver"
	"github.com/spf13/viper"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.uber.org/zap"

	"github.com/jaegertracing/jaeger/cmd/agent/app/reporter/grpc"
	"github.com/jaegertracing/jaeger/cmd/opentelemetry-collector/app/exporter/cassandra"
	"github.com/jaegertracing/jaeger/cmd/opentelemetry-collector/app/exporter/elasticsearch"
	"github.com/jaegertracing/jaeger/cmd/opentelemetry-collector/app/exporter/grpcplugin"
	"github.com/jaegertracing/jaeger/cmd/opentelemetry-collector/app/exporter/kafka"
	kafkaRec "github.com/jaegertracing/jaeger/cmd/opentelemetry-collector/app/receiver/kafka"
	jConfig "github.com/jaegertracing/jaeger/pkg/config"
	"github.com/jaegertracing/jaeger/ports"
)

func TestDefaultCollectorConfig(t *testing.T) {
	disabledHostPort := ports.PortToHostPort(0)
	tests := []struct {
		storageType    string
		zipkinHostPort string
		exporterTypes  []string
		pipeline       configmodels.Pipelines
		err            string
		config         map[string]interface{}
	}{
		{
			storageType:    "elasticsearch",
			zipkinHostPort: disabledHostPort,
			exporterTypes:  []string{elasticsearch.TypeStr},
			pipeline: configmodels.Pipelines{
				"traces": {
					InputType: configmodels.TracesDataType,
					Receivers: []string{"jaeger"},
					Exporters: []string{elasticsearch.TypeStr},
				},
			},
		},
		{
			storageType:    "cassandra",
			zipkinHostPort: disabledHostPort,
			exporterTypes:  []string{cassandra.TypeStr},
			pipeline: configmodels.Pipelines{
				"traces": {
					InputType: configmodels.TracesDataType,
					Receivers: []string{"jaeger"},
					Exporters: []string{cassandra.TypeStr},
				},
			},
		},
		{
			storageType:    "kafka",
			zipkinHostPort: disabledHostPort,
			exporterTypes:  []string{kafka.TypeStr},
			pipeline: configmodels.Pipelines{
				"traces": {
					InputType: configmodels.TracesDataType,
					Receivers: []string{"jaeger"},
					Exporters: []string{kafka.TypeStr},
				},
			},
		},
		{
			storageType:    "cassandra,elasticsearch,grpc-plugin",
			zipkinHostPort: disabledHostPort,
			exporterTypes:  []string{cassandra.TypeStr, elasticsearch.TypeStr, grpcplugin.TypeStr},
			pipeline: configmodels.Pipelines{
				"traces": {
					InputType: configmodels.TracesDataType,
					Receivers: []string{"jaeger"},
					Exporters: []string{cassandra.TypeStr, elasticsearch.TypeStr, grpcplugin.TypeStr},
				},
			},
		},
		{
			storageType:    "cassandra",
			zipkinHostPort: ":9411",
			exporterTypes:  []string{cassandra.TypeStr},
			config:         map[string]interface{}{"resource.labels": "foo=bar"},
			pipeline: configmodels.Pipelines{
				"traces": {
					InputType:  configmodels.TracesDataType,
					Receivers:  []string{"jaeger", "zipkin"},
					Processors: []string{"resource"},
					Exporters:  []string{cassandra.TypeStr},
				},
			},
		},
		{
			storageType: "floppy",
			err:         "unknown storage type: floppy",
		},
	}
	for _, test := range tests {
		t.Run(test.storageType, func(t *testing.T) {
			v, _ := jConfig.Viperize(grpc.AddFlags)
			factories := Components(v)
			for key, val := range test.config {
				v.Set(key, val)
			}
			cfg, err := CollectorConfig(test.storageType, test.zipkinHostPort, factories)
			if test.err != "" {
				require.Nil(t, cfg)
				assert.EqualError(t, err, test.err)
				return
			}
			require.NoError(t, err)
			require.NoError(t, config.ValidateConfig(cfg, zap.NewNop()))

			assert.Equal(t, 1, len(cfg.Extensions))
			assert.Equal(t, 1, len(cfg.Service.Extensions))
			assert.Equal(t, "health_check", cfg.Service.Extensions[0])
			assert.Equal(t, "health_check", cfg.Extensions["health_check"].Name())
			assert.Equal(t, len(test.pipeline["traces"].Receivers), len(cfg.Receivers))
			assert.Equal(t, "jaeger", cfg.Receivers["jaeger"].Name())
			assert.Equal(t, len(test.exporterTypes), len(cfg.Exporters))

			processorMap := map[string]bool{}
			for _, p := range test.pipeline["traces"].Processors {
				processorMap[p] = true
			}
			if processorMap["resource"] {
				assert.Equal(t, len(processorMap), len(cfg.Processors))
				assert.IsType(t, &resourceprocessor.Config{}, cfg.Processors["resource"])
			}

			types := []string{}
			for _, v := range cfg.Exporters {
				types = append(types, string(v.Type()))
			}
			sort.Strings(types)
			assert.Equal(t, test.exporterTypes, types)
			sort.Strings(cfg.Service.Pipelines["traces"].Receivers)
			sort.Strings(cfg.Service.Pipelines["traces"].Exporters)
			assert.EqualValues(t, test.pipeline, cfg.Service.Pipelines)
		})
	}
}

func TestDefaultAgentConfig(t *testing.T) {
	tests := []struct {
		config  map[string]interface{}
		service configmodels.Service
	}{
		{
			config: map[string]interface{}{"resource.labels": "foo=bar"},
			service: configmodels.Service{
				Extensions: []string{"health_check"},
				Pipelines: configmodels.Pipelines{
					"traces": &configmodels.Pipeline{
						InputType:  configmodels.TracesDataType,
						Receivers:  []string{"jaeger"},
						Processors: []string{"resource"},
						Exporters:  []string{"jaeger"},
					},
				},
			},
		},
		{
			service: configmodels.Service{
				Extensions: []string{"health_check"},
				Pipelines: configmodels.Pipelines{
					"traces": &configmodels.Pipeline{
						InputType: configmodels.TracesDataType,
						Receivers: []string{"jaeger"},
						Exporters: []string{"jaeger"},
					},
				},
			},
		},
	}
	for _, test := range tests {
		t.Run(fmt.Sprintf("%v", test.config), func(t *testing.T) {
			v, _ := jConfig.Viperize(grpc.AddFlags)
			for key, val := range test.config {
				v.Set(key, val)
			}
			factories := Components(v)
			cfg := AgentConfig(factories)
			require.NoError(t, config.ValidateConfig(cfg, zap.NewNop()))

			assert.Equal(t, test.service, cfg.Service)
			assert.Equal(t, 1, len(cfg.Receivers))
			assert.IsType(t, &jaegerreceiver.Config{}, cfg.Receivers["jaeger"])
			assert.Equal(t, 1, len(cfg.Exporters))
			assert.IsType(t, &jaegerexporter.Config{}, cfg.Exporters["jaeger"])
			processorMap := map[string]bool{}
			for _, p := range test.service.Pipelines["traces"].Processors {
				processorMap[p] = true
			}
			if processorMap["resource"] {
				assert.Equal(t, len(processorMap), len(cfg.Processors))
				assert.IsType(t, &resourceprocessor.Config{}, cfg.Processors["resource"])
			}
		})
	}
}

func TestDefaultIngesterConfig(t *testing.T) {
	tests := []struct {
		storageType string
		service     configmodels.Service
		err         string
	}{
		{
			storageType: "elasticsearch",
			service: configmodels.Service{
				Extensions: []string{"health_check"},
				Pipelines: configmodels.Pipelines{
					"traces": &configmodels.Pipeline{
						InputType: configmodels.TracesDataType,
						Receivers: []string{kafkaRec.TypeStr},
						Exporters: []string{elasticsearch.TypeStr},
					},
				},
			},
		},
		{
			storageType: "elasticsearch,cassandra,grpc-plugin",
			service: configmodels.Service{
				Extensions: []string{"health_check"},
				Pipelines: configmodels.Pipelines{
					"traces": &configmodels.Pipeline{
						InputType: configmodels.TracesDataType,
						Receivers: []string{kafkaRec.TypeStr},
						Exporters: []string{cassandra.TypeStr, elasticsearch.TypeStr, grpcplugin.TypeStr},
					},
				},
			},
		},
		{
			storageType: "floppy",
			err:         "unknown storage type: floppy",
		},
	}
	for _, test := range tests {
		t.Run(test.storageType, func(t *testing.T) {
			factories := Components(viper.New())
			cfg, err := IngesterConfig(test.storageType, factories)
			if test.err != "" {
				require.Nil(t, cfg)
				assert.EqualError(t, err, test.err)
				return
			}
			require.NoError(t, err)
			require.NoError(t, config.ValidateConfig(cfg, zap.NewNop()))

			sort.Strings(cfg.Service.Pipelines["traces"].Exporters)
			assert.Equal(t, test.service, cfg.Service)
			assert.Equal(t, 1, len(cfg.Receivers))
			assert.IsType(t, &kafkaRec.Config{}, cfg.Receivers[kafkaRec.TypeStr])

			assert.Equal(t, len(test.service.Pipelines["traces"].Exporters), len(cfg.Exporters))
			types := []string{}
			for _, v := range cfg.Exporters {
				types = append(types, string(v.Type()))
			}
			sort.Strings(types)
			assert.Equal(t, test.service.Pipelines["traces"].Exporters, types)
		})
	}
}
<|endoftext|>


<commit_before>package secret

import (
	"fmt"
	"os"

	"github.com/spf13/cobra"
	"golang.org/x/crypto/ssh/terminal"

	"github.com/flant/werf/cmd/werf/common"
	"github.com/flant/werf/pkg/deploy/secret"
)

func NewCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:                   "generate-secret-key",
		DisableFlagsInUseLine: true,
		Short:                 "Generate hex encryption key that can be used as $WERF_SECRET_KEY",
		Long: common.GetLongCommandDescription(`Generate hex key that can be used as $WERF_SECRET_KEY.

16-bytes key will be generated (AES-128)`),
		RunE: func(cmd *cobra.Command, args []string) error {
			return runGenerateSecretKey()
		},
	}

	return cmd
}

func runGenerateSecretKey() error {
	key, err := secret.GenerateSecretKey()
	if err != nil {
		return err
	}

	if terminal.IsTerminal(int(os.Stdout.Fd())) {
		fmt.Printf("WERF_SECRET_KEY=%s\n", string(key))
	} else {
		fmt.Println(string(key))
	}

	return nil
}
<commit_msg>Update generate-secret-key command<commit_after>package secret

import (
	"fmt"

	"github.com/flant/werf/cmd/werf/common"
	"github.com/flant/werf/pkg/deploy/secret"
	"github.com/spf13/cobra"
)

func NewCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:                   "generate-secret-key",
		DisableFlagsInUseLine: true,
		Short:                 "Generate hex encryption key",
		Long: common.GetLongCommandDescription(`Generate hex encryption key.
For further usage, the encryption key should be saved in $WERF_SECRET_KEY or .werf_secret_key file`),
		Example: `  # Export encryption key
  $ export WERF_SECRET_KEY=$(werf helm secret generate-secret-key)

  # Save encryption key in .werf_secret_key file
  $ werf helm secret generate-secret-key > .werf_secret_key`,
		RunE: func(cmd *cobra.Command, args []string) error {
			return runGenerateSecretKey()
		},
	}

	return cmd
}

func runGenerateSecretKey() error {
	key, err := secret.GenerateSecretKey()
	if err != nil {
		return err
	}

	fmt.Println(string(key))

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"github.com/mesosphere/dcos-commons/cli"
	"gopkg.in/alecthomas/kingpin.v2"
	"log"
	"net/url"
	"os"
	"strings"
)

func main() {
	modName, err := cli.GetModuleName()
	if err != nil {
		log.Fatalf(err.Error())
	}

	app, err := cli.NewApp(
		"0.1.0",
		"Mesosphere",
		fmt.Sprintf("Deploy and manage %s clusters", strings.Title(modName)))
	if err != nil {
		log.Fatalf(err.Error())
	}

	// Omit standard "connection" section. Cassandra isn't on the standard paths yet, and omit
	// standard "config" and "state" sections since Cassandra isn't installing
	// ConfigResource/StateResource yet.
	// Once these are fixed, this block can all be replaced with a call to "HandleCommonArgs"
	cli.HandleCommonFlags(app, modName, fmt.Sprintf("%s DC/OS CLI Module", strings.Title(modName)))
	//cli.HandleConfigSection(app)
	//cli.HandleConnectionSection(app)
	cli.HandlePlanSection(app)
	//cli.HandleStateSection(app)

	handleSeedsCommand(app)
	handleCustomConnectionCommand(app, modName)

	handleNodeSection(app, modName)
	handleBackupRestoreSections(app, modName)
	handleCleanupRepairSections(app)

	// Omit modname:
	kingpin.MustParse(app.Parse(os.Args[2:]))
}

func handleSeedsCommand(app *kingpin.Application) {
	app.Command("seeds", "Retrieve seed node information").Action(
		func(c *kingpin.ParseContext) error {
			cli.PrintJSON(cli.HTTPGet("v1/seeds"))
			return nil
		})
}

type ConnectionHandler struct {
	showAddress bool
	showDns bool
}
func (cmd *ConnectionHandler) runBase(c *kingpin.ParseContext) error {
	if cmd.showAddress {
		cli.PrintJSON(cli.HTTPGet("v1/nodes/connect/address"))
	} else if cmd.showDns {
		cli.PrintJSON(cli.HTTPGet("v1/nodes/connect/dns"))
	} else {
		cli.PrintJSON(cli.HTTPGet("v1/nodes/connect"))
	}
	return nil
}
func handleCustomConnectionCommand(app *kingpin.Application, modName string) {
	cmd := &ConnectionHandler{}

	// Have three explicit commands, to ensure that each possibility shows up in --help:
	connection := app.Command("connection", fmt.Sprintf("Provides %s connection information", modName)).Action(cmd.runBase)
	connection.Flag("address", fmt.Sprintf("Provide addresses of the %s nodes", modName)).BoolVar(&cmd.showAddress)
	connection.Flag("dns", fmt.Sprintf("Provide dns names of the %s nodes", modName)).BoolVar(&cmd.showDns)
}

type NodeHandler struct {
	nodeId int
}
func (cmd *NodeHandler) runDescribe(c *kingpin.ParseContext) error {
	cli.PrintJSON(cli.HTTPGet(fmt.Sprintf("v1/nodes/node-%d/info", cmd.nodeId)))
	return nil
}
func (cmd *NodeHandler) runList(c *kingpin.ParseContext) error {
	cli.PrintJSON(cli.HTTPGet("v1/nodes/list"))
	return nil
}
func (cmd *NodeHandler) runReplace(c *kingpin.ParseContext) error {
	query := url.Values{}
	query.Set("node", fmt.Sprintf("node-%d", cmd.nodeId))
	cli.PrintJSON(cli.HTTPGetQuery("v1/nodes/replace", query.Encode()))
	return nil
}
func (cmd *NodeHandler) runRestart(c *kingpin.ParseContext) error {
	query := url.Values{}
	query.Set("node", fmt.Sprintf("node-%d", cmd.nodeId))
	cli.PrintJSON(cli.HTTPGetQuery("v1/nodes/restart", query.Encode()))
	return nil
}
func (cmd *NodeHandler) runStatus(c *kingpin.ParseContext) error {
	cli.PrintJSON(cli.HTTPGet(fmt.Sprintf("v1/nodes/node-%d/status", cmd.nodeId)))
	return nil
}
func handleNodeSection(app *kingpin.Application, modName string) {
	cmd := &NodeHandler{}
	connection := app.Command("node", fmt.Sprintf("Manage %s nodes", modName))

	describe := connection.Command("describe", "Describes a single node").Action(cmd.runDescribe)
	describe.Arg("node_id", "The node id to describe").IntVar(&cmd.nodeId)

	connection.Command("list", "Lists all nodes").Action(cmd.runList)

	replace := connection.Command("replace", "Replaces a single node job, moving it to a different agent").Action(cmd.runReplace)
	replace.Arg("node_id", "The node id to replace").IntVar(&cmd.nodeId)

	restart := connection.Command("restart", "Restarts a single node job, keeping it on the same agent").Action(cmd.runRestart)
	restart.Arg("node_id", "The node id to restart").IntVar(&cmd.nodeId)

	status := connection.Command("status", "Gets the status of a single node").Action(cmd.runStatus)
	status.Arg("node_id", "The node id to check").IntVar(&cmd.nodeId)
}

// Reuse same struct for both 'backup start' and 'restore start' (same args)
type BackupRestoreHandler struct {
	backupName string
	externalLocation string
	s3AccessKey string
	s3SecretKey string
	azureAccount string
	azureKey string
}
func (cmd *BackupRestoreHandler) getArgs() map[string]interface{} {
	return map[string]interface{} {
		"backup_name": cmd.backupName,
		"external_location": cmd.externalLocation,
		"s3_access_key": cmd.s3AccessKey,
		"s3_secret_key": cmd.s3SecretKey,
		"azure_account": cmd.azureAccount,
		"azure_key": cmd.azureKey,
	}
}
func (cmd *BackupRestoreHandler) runBackup(c *kingpin.ParseContext) error {
	payload, err := json.Marshal(cmd.getArgs())
	if err != nil {
		return err
	}
	cli.PrintJSON(cli.HTTPPutJSON("v1/backup/start", string(payload)))
	return nil
}
func (cmd *BackupRestoreHandler) runBackupStop(c *kingpin.ParseContext) error {
	cli.PrintJSON(cli.HTTPPut("v1/backup/stop"))
	return nil
}
func (cmd *BackupRestoreHandler) runRestore(c *kingpin.ParseContext) error {
	payload, err := json.Marshal(cmd.getArgs())
	if err != nil {
		return err
	}
	cli.PrintJSON(cli.HTTPPutJSON("v1/restore/start", string(payload)))
	return nil
}
func handleBackupRestoreSections(app *kingpin.Application, modName string) {
	cmd := &BackupRestoreHandler{}
	planCmd := &cli.PlanHandler{}

	backup := app.Command("backup", fmt.Sprintf("Backup %s cluster data", modName))
	backupStart := backup.Command(
		"start",
		"Perform cluster backup via snapshot mechanism").Action(cmd.runBackup)
	backupStart.Flag("backup_name", "Name of the snapshot").StringVar(&cmd.backupName)
	backupStart.Flag("external_location", "External location where the snapshot should be stored").StringVar(&cmd.externalLocation)
	backupStart.Flag("s3_access_key", "S3 access key").StringVar(&cmd.s3AccessKey)
	backupStart.Flag("s3_secret_key", "S3 secret key").StringVar(&cmd.s3SecretKey)
	backupStart.Flag("azure_account", "Azure storage account").StringVar(&cmd.azureAccount)
	backupStart.Flag("azure_key", "Azure secret key").StringVar(&cmd.azureKey)
	backup.Command(
		"stop",
		"Stops a currently running backup").Action(cmd.runBackupStop)
	// same as 'plan show':
	backup.Command(
		"status",
		"Displays the status of the backup").Action(planCmd.RunShow)

	restore := app.Command("restore", fmt.Sprintf("Restore %s cluster from backup", modName))
	restoreStart := restore.Command(
		"start",
		"Restores cluster to a previous snapshot").Action(cmd.runRestore)
	restoreStart.Flag("backup_name", "Name of the snapshot to restore").StringVar(&cmd.backupName)
	restoreStart.Flag("external_location", "External location where the snapshot is stored").StringVar(&cmd.externalLocation)
	restoreStart.Flag("s3_access_key", "S3 access key").StringVar(&cmd.s3AccessKey)
	restoreStart.Flag("s3_secret_key", "S3 secret key").StringVar(&cmd.s3SecretKey)
	restoreStart.Flag("azure_account", "Azure storage account").StringVar(&cmd.azureAccount)
	restoreStart.Flag("azure_key", "Azure secret key").StringVar(&cmd.azureKey)
	// same as 'plan show':
	restore.Command(
		"status",
		"Displays the status of the restore").Action(planCmd.RunShow)
}

// Reuse same struct for both 'cleanup start' and 'repair start' (same args)
type CleanupRepairHandler struct {
	nodes string
	keySpaces string
	columnFamilies string
}
func (cmd *CleanupRepairHandler) getArgs() map[string]interface{} {
	nodesList := []string{}
	for _, node := range strings.Split(cmd.nodes, ",") {
		nodesList = append(nodesList, fmt.Sprintf("node-%s", strings.TrimSpace(node)))
	}

	dict := map[string]interface{} {"nodes": nodesList}

	if len(cmd.keySpaces) != 0 {
		dict["key_spaces"] = cmd.keySpaces
	}
	if len(cmd.keySpaces) != 0 {
		dict["column_families"] = cmd.columnFamilies
	}
	return dict
}
func (cmd *CleanupRepairHandler) runCleanup(c *kingpin.ParseContext) error {
	payload, err := json.Marshal(cmd.getArgs())
	if err != nil {
		return err
	}
	cli.PrintJSON(cli.HTTPPutJSON("v1/cleanup/start", string(payload)))
	return nil
}
func (cmd *CleanupRepairHandler) runRepair(c *kingpin.ParseContext) error {
	payload, err := json.Marshal(cmd.getArgs())
	if err != nil {
		return err
	}
	cli.PrintJSON(cli.HTTPPutJSON("v1/repair/start", string(payload)))
	return nil
}
func handleCleanupRepairSections(app *kingpin.Application) {
	cmd := &CleanupRepairHandler{}

	cleanup := app.Command("cleanup", "Clean up old token mappings")
	cleanupStart := cleanup.Command(
		"start",
		"Perform cluster cleanup of deleted or moved keys").Action(cmd.runCleanup)
	cleanupStart.Flag("nodes", "A list of the nodes to cleanup or * for all.").Default("*").StringVar(&cmd.nodes)
	cleanupStart.Flag("key_spaces", "The key spaces to cleanup or empty for all.").StringVar(&cmd.keySpaces)
	cleanupStart.Flag("column_families", "The column families to cleanup.").StringVar(&cmd.columnFamilies)

	repair := app.Command("repair", "Perform primary range repair")
	repairStart := repair.Command(
		"start",
		"Perform primary range anti-entropy repair").Action(cmd.runRepair)
	repairStart.Flag("nodes", "A list of the nodes to repair or * for all.").Default("*").StringVar(&cmd.nodes)
	repairStart.Flag("key_spaces", "The key spaces to repair or empty for all.").StringVar(&cmd.keySpaces)
	repairStart.Flag("column_families", "The column families to repair.").StringVar(&cmd.columnFamilies)
}
<commit_msg>Add CLI support for stopping all commands<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"github.com/mesosphere/dcos-commons/cli"
	"gopkg.in/alecthomas/kingpin.v2"
	"log"
	"net/url"
	"os"
	"strings"
)

func main() {
	modName, err := cli.GetModuleName()
	if err != nil {
		log.Fatalf(err.Error())
	}

	app, err := cli.NewApp(
		"0.1.0",
		"Mesosphere",
		fmt.Sprintf("Deploy and manage %s clusters", strings.Title(modName)))
	if err != nil {
		log.Fatalf(err.Error())
	}

	// Omit standard "connection" section. Cassandra isn't on the standard paths yet, and omit
	// standard "config" and "state" sections since Cassandra isn't installing
	// ConfigResource/StateResource yet.
	// Once these are fixed, this block can all be replaced with a call to "HandleCommonArgs"
	cli.HandleCommonFlags(app, modName, fmt.Sprintf("%s DC/OS CLI Module", strings.Title(modName)))
	//cli.HandleConfigSection(app)
	//cli.HandleConnectionSection(app)
	cli.HandlePlanSection(app)
	//cli.HandleStateSection(app)

	handleSeedsCommand(app)
	handleCustomConnectionCommand(app, modName)

	handleNodeSection(app, modName)
	handleBackupRestoreSections(app, modName)
	handleCleanupRepairSections(app)

	// Omit modname:
	kingpin.MustParse(app.Parse(os.Args[2:]))
}

func handleSeedsCommand(app *kingpin.Application) {
	app.Command("seeds", "Retrieve seed node information").Action(
		func(c *kingpin.ParseContext) error {
			cli.PrintJSON(cli.HTTPGet("v1/seeds"))
			return nil
		})
}

type ConnectionHandler struct {
	showAddress bool
	showDns bool
}
func (cmd *ConnectionHandler) runBase(c *kingpin.ParseContext) error {
	if cmd.showAddress {
		cli.PrintJSON(cli.HTTPGet("v1/nodes/connect/address"))
	} else if cmd.showDns {
		cli.PrintJSON(cli.HTTPGet("v1/nodes/connect/dns"))
	} else {
		cli.PrintJSON(cli.HTTPGet("v1/nodes/connect"))
	}
	return nil
}
func handleCustomConnectionCommand(app *kingpin.Application, modName string) {
	cmd := &ConnectionHandler{}

	// Have three explicit commands, to ensure that each possibility shows up in --help:
	connection := app.Command("connection", fmt.Sprintf("Provides %s connection information", modName)).Action(cmd.runBase)
	connection.Flag("address", fmt.Sprintf("Provide addresses of the %s nodes", modName)).BoolVar(&cmd.showAddress)
	connection.Flag("dns", fmt.Sprintf("Provide dns names of the %s nodes", modName)).BoolVar(&cmd.showDns)
}

type NodeHandler struct {
	nodeId int
}
func (cmd *NodeHandler) runDescribe(c *kingpin.ParseContext) error {
	cli.PrintJSON(cli.HTTPGet(fmt.Sprintf("v1/nodes/node-%d/info", cmd.nodeId)))
	return nil
}
func (cmd *NodeHandler) runList(c *kingpin.ParseContext) error {
	cli.PrintJSON(cli.HTTPGet("v1/nodes/list"))
	return nil
}
func (cmd *NodeHandler) runReplace(c *kingpin.ParseContext) error {
	query := url.Values{}
	query.Set("node", fmt.Sprintf("node-%d", cmd.nodeId))
	cli.PrintJSON(cli.HTTPGetQuery("v1/nodes/replace", query.Encode()))
	return nil
}
func (cmd *NodeHandler) runRestart(c *kingpin.ParseContext) error {
	query := url.Values{}
	query.Set("node", fmt.Sprintf("node-%d", cmd.nodeId))
	cli.PrintJSON(cli.HTTPGetQuery("v1/nodes/restart", query.Encode()))
	return nil
}
func (cmd *NodeHandler) runStatus(c *kingpin.ParseContext) error {
	cli.PrintJSON(cli.HTTPGet(fmt.Sprintf("v1/nodes/node-%d/status", cmd.nodeId)))
	return nil
}
func handleNodeSection(app *kingpin.Application, modName string) {
	cmd := &NodeHandler{}
	connection := app.Command("node", fmt.Sprintf("Manage %s nodes", modName))

	describe := connection.Command("describe", "Describes a single node").Action(cmd.runDescribe)
	describe.Arg("node_id", "The node id to describe").IntVar(&cmd.nodeId)

	connection.Command("list", "Lists all nodes").Action(cmd.runList)

	replace := connection.Command("replace", "Replaces a single node job, moving it to a different agent").Action(cmd.runReplace)
	replace.Arg("node_id", "The node id to replace").IntVar(&cmd.nodeId)

	restart := connection.Command("restart", "Restarts a single node job, keeping it on the same agent").Action(cmd.runRestart)
	restart.Arg("node_id", "The node id to restart").IntVar(&cmd.nodeId)

	status := connection.Command("status", "Gets the status of a single node").Action(cmd.runStatus)
	status.Arg("node_id", "The node id to check").IntVar(&cmd.nodeId)
}

// Reuse same struct for both 'backup start' and 'restore start' (same args)
type BackupRestoreHandler struct {
	backupName string
	externalLocation string
	s3AccessKey string
	s3SecretKey string
	azureAccount string
	azureKey string
}
func (cmd *BackupRestoreHandler) getArgs() map[string]interface{} {
	return map[string]interface{} {
		"backup_name": cmd.backupName,
		"external_location": cmd.externalLocation,
		"s3_access_key": cmd.s3AccessKey,
		"s3_secret_key": cmd.s3SecretKey,
		"azure_account": cmd.azureAccount,
		"azure_key": cmd.azureKey,
	}
}
func (cmd *BackupRestoreHandler) runBackup(c *kingpin.ParseContext) error {
	payload, err := json.Marshal(cmd.getArgs())
	if err != nil {
		return err
	}
	cli.PrintJSON(cli.HTTPPutJSON("v1/backup/start", string(payload)))
	return nil
}
func (cmd *BackupRestoreHandler) runBackupStop(c *kingpin.ParseContext) error {
	cli.PrintJSON(cli.HTTPPut("v1/backup/stop"))
	return nil
}
func (cmd *BackupRestoreHandler) runRestore(c *kingpin.ParseContext) error {
	payload, err := json.Marshal(cmd.getArgs())
	if err != nil {
		return err
	}
	cli.PrintJSON(cli.HTTPPutJSON("v1/restore/start", string(payload)))
	return nil
}
func (cmd *BackupRestoreHandler) runRestoreStop(c *kingpin.ParseContext) error {
	cli.PrintJSON(cli.HTTPPut("v1/restore/stop"))
	return nil
}
func handleBackupRestoreSections(app *kingpin.Application, modName string) {
	cmd := &BackupRestoreHandler{}
	planCmd := &cli.PlanHandler{}

	backup := app.Command("backup", fmt.Sprintf("Backup %s cluster data", modName))
	backupStart := backup.Command(
		"start",
		"Perform cluster backup via snapshot mechanism").Action(cmd.runBackup)
	backupStart.Flag("backup_name", "Name of the snapshot").StringVar(&cmd.backupName)
	backupStart.Flag("external_location", "External location where the snapshot should be stored").StringVar(&cmd.externalLocation)
	backupStart.Flag("s3_access_key", "S3 access key").StringVar(&cmd.s3AccessKey)
	backupStart.Flag("s3_secret_key", "S3 secret key").StringVar(&cmd.s3SecretKey)
	backupStart.Flag("azure_account", "Azure storage account").StringVar(&cmd.azureAccount)
	backupStart.Flag("azure_key", "Azure secret key").StringVar(&cmd.azureKey)
	backup.Command(
		"stop",
		"Stops a currently running backup").Action(cmd.runBackupStop)
	// same as 'plan show':
	backup.Command(
		"status",
		"Displays the status of the backup").Action(planCmd.RunShow)

	restore := app.Command("restore", fmt.Sprintf("Restore %s cluster from backup", modName))
	restoreStart := restore.Command(
		"start",
		"Restores cluster to a previous snapshot").Action(cmd.runRestore)
	restoreStart.Flag("backup_name", "Name of the snapshot to restore").StringVar(&cmd.backupName)
	restoreStart.Flag("external_location", "External location where the snapshot is stored").StringVar(&cmd.externalLocation)
	restoreStart.Flag("s3_access_key", "S3 access key").StringVar(&cmd.s3AccessKey)
	restoreStart.Flag("s3_secret_key", "S3 secret key").StringVar(&cmd.s3SecretKey)
	restoreStart.Flag("azure_account", "Azure storage account").StringVar(&cmd.azureAccount)
	restoreStart.Flag("azure_key", "Azure secret key").StringVar(&cmd.azureKey)
	restore.Command(
		"stop",
		"Stops a currently running restore").Action(cmd.runRestoreStop)
	// same as 'plan show':
	restore.Command(
		"status",
		"Displays the status of the restore").Action(planCmd.RunShow)
}

// Reuse same struct for both 'cleanup start' and 'repair start' (same args)
type CleanupRepairHandler struct {
	nodes string
	keySpaces string
	columnFamilies string
}
func (cmd *CleanupRepairHandler) getArgs() map[string]interface{} {
	nodesList := []string{}
	for _, node := range strings.Split(cmd.nodes, ",") {
		nodesList = append(nodesList, fmt.Sprintf("node-%s", strings.TrimSpace(node)))
	}

	dict := map[string]interface{} {"nodes": nodesList}

	if len(cmd.keySpaces) != 0 {
		dict["key_spaces"] = cmd.keySpaces
	}
	if len(cmd.keySpaces) != 0 {
		dict["column_families"] = cmd.columnFamilies
	}
	return dict
}
func (cmd *CleanupRepairHandler) runCleanup(c *kingpin.ParseContext) error {
	payload, err := json.Marshal(cmd.getArgs())
	if err != nil {
		return err
	}
	cli.PrintJSON(cli.HTTPPutJSON("v1/cleanup/start", string(payload)))
	return nil
}
func (cmd *CleanupRepairHandler) runCleanupStop(c *kingpin.ParseContext) error {
	cli.PrintJSON(cli.HTTPPut("v1/cleanup/stop"))
	return nil
}
func (cmd *CleanupRepairHandler) runRepair(c *kingpin.ParseContext) error {
	payload, err := json.Marshal(cmd.getArgs())
	if err != nil {
		return err
	}
	cli.PrintJSON(cli.HTTPPutJSON("v1/repair/start", string(payload)))
	return nil
}
func (cmd *CleanupRepairHandler) runRepairStop(c *kingpin.ParseContext) error {
	cli.PrintJSON(cli.HTTPPut("v1/repair/stop"))
	return nil
}
func handleCleanupRepairSections(app *kingpin.Application) {
	cmd := &CleanupRepairHandler{}

	cleanup := app.Command("cleanup", "Clean up old token mappings")
	cleanupStart := cleanup.Command(
		"start",
		"Perform cluster cleanup of deleted or moved keys").Action(cmd.runCleanup)
	cleanupStart.Flag("nodes", "A list of the nodes to cleanup or * for all.").Default("*").StringVar(&cmd.nodes)
	cleanupStart.Flag("key_spaces", "The key spaces to cleanup or empty for all.").StringVar(&cmd.keySpaces)
	cleanupStart.Flag("column_families", "The column families to cleanup.").StringVar(&cmd.columnFamilies)
	cleanup.Command(
		"stop",
		"Stops a currently running cleanup").Action(cmd.runCleanupStop)

	repair := app.Command("repair", "Perform primary range repair")
	repairStart := repair.Command(
		"start",
		"Perform primary range anti-entropy repair").Action(cmd.runRepair)
	repairStart.Flag("nodes", "A list of the nodes to repair or * for all.").Default("*").StringVar(&cmd.nodes)
	repairStart.Flag("key_spaces", "The key spaces to repair or empty for all.").StringVar(&cmd.keySpaces)
	repairStart.Flag("column_families", "The column families to repair.").StringVar(&cmd.columnFamilies)
	repair.Command(
		"stop",
		"Stops a currently running repair").Action(cmd.runRepairStop)
}
<|endoftext|>


<commit_before>package calcium

import (
	"context"
	"fmt"
	"sync"

	log "github.com/Sirupsen/logrus"
	enginecontainer "github.com/docker/docker/api/types/container"
	"gitlab.ricebook.net/platform/core/types"
	"gitlab.ricebook.net/platform/core/utils"
)

type NodeContainers map[*types.Node][]*types.Container
type NodeCPUMap map[*types.Node][]types.CPUMap
type CPUNodeContainers map[float64]NodeContainers
type CPUNodeContainersMap map[float64]NodeCPUMap

func (c *calcium) ReallocResource(ids []string, cpu float64, mem int64) (chan *types.ReallocResourceMessage, error) {
	// TODO 大量容器 Get 的时候有性能问题
	containers, err := c.store.GetContainers(ids)
	if err != nil {
		return nil, err
	}

	// Pod-Node-Containers 三元组
	containersInfo := map[*types.Pod]NodeContainers{}
	// Pod-cpu-node-containers 四元组
	cpuContainersInfo := map[*types.Pod]CPUNodeContainers{}
	nodeCache := map[string]*types.Node{}

	for _, container := range containers {
		if _, ok := nodeCache[container.Nodename]; !ok {
			node, err := c.GetNode(container.Podname, container.Nodename)
			if err != nil {
				return nil, err
			}
			nodeCache[container.Nodename] = node
		}
		node := nodeCache[container.Nodename]

		pod, err := c.store.GetPod(container.Podname)
		if err != nil {
			return nil, err
		}

		if _, ok := containersInfo[pod]; !ok {
			containersInfo[pod] = NodeContainers{}
		}
		if _, ok := containersInfo[pod][node]; !ok {
			containersInfo[pod][node] = []*types.Container{}
		}
		containersInfo[pod][node] = append(containersInfo[pod][node], container)

		if pod.Favor == types.CPU_PRIOR {
			if _, ok := cpuContainersInfo[pod]; !ok {
				cpuContainersInfo[pod] = CPUNodeContainers{}
			}
			podCPUContainersInfo := cpuContainersInfo[pod]
			newCPURequire := c.calculateCPUUsage(container) + cpu
			if newCPURequire < 0.0 {
				return nil, fmt.Errorf("[realloc] cpu can not below zero")
			}
			if _, ok := podCPUContainersInfo[newCPURequire]; !ok {
				podCPUContainersInfo[newCPURequire] = NodeContainers{} // go map 的傻逼语法, 不这样写会有 assignment to entry in nil map 错误
				podCPUContainersInfo[newCPURequire][node] = []*types.Container{}
			}
			podCPUContainersInfo[newCPURequire][node] = append(podCPUContainersInfo[newCPURequire][node], container)
		}
	}

	ch := make(chan *types.ReallocResourceMessage)
	go func() {
		defer close(ch)
		wg := sync.WaitGroup{}
		wg.Add(len(containersInfo))
		for pod, nodeContainers := range containersInfo {
			if pod.Favor == types.CPU_PRIOR {
				nodeCPUContainersInfo := cpuContainersInfo[pod]
				go func(pod *types.Pod, nodeCPUContainersInfo CPUNodeContainers) {
					defer wg.Done()
					c.reallocContainersWithCPUPrior(ch, pod, nodeCPUContainersInfo, cpu, mem)
				}(pod, nodeCPUContainersInfo)
				continue
			}
			go func(pod *types.Pod, nodeContainers NodeContainers) {
				defer wg.Done()
				c.reallocContainerWithMemoryPrior(ch, pod, nodeContainers, cpu, mem)
			}(pod, nodeContainers)
		}
		wg.Wait()
	}()
	return ch, nil
}

func (c *calcium) reallocNodesMemory(podname string, nodeContainers NodeContainers, memory int64) error {
	lock, err := c.Lock(podname, 30)
	if err != nil {
		return err
	}
	defer lock.Unlock()
	for node, containers := range nodeContainers {
		if memory > 0 {
			if cap := int(node.MemCap / memory); cap < len(containers) {
				return fmt.Errorf("[realloc] Not enough resource %s", node.Name)
			}
		}
		if err := c.store.UpdateNodeMem(podname, node.Name, int64(len(containers))*memory, "-"); err != nil {
			return err
		}
	}
	return nil
}

func (c *calcium) reallocContainerWithMemoryPrior(
	ch chan *types.ReallocResourceMessage,
	pod *types.Pod,
	nodeContainers NodeContainers,
	cpu float64, memory int64) {

	if err := c.reallocNodesMemory(pod.Name, nodeContainers, memory); err != nil {
		log.Errorf("[realloc] realloc memory failed %v", err)
		return
	}

	// 不并发操作了
	for node, containers := range nodeContainers {
		c.doUpdateContainerWithMemoryPrior(ch, pod.Name, node, containers, cpu, memory)
	}
}

func (c *calcium) doUpdateContainerWithMemoryPrior(
	ch chan *types.ReallocResourceMessage,
	podname string,
	node *types.Node,
	containers []*types.Container,
	cpu float64, memory int64) {

	for _, container := range containers {
		containerJSON, err := container.Inspect()
		if err != nil {
			log.Errorf("[realloc] get container failed %v", err)
			ch <- &types.ReallocResourceMessage{ContainerID: containerJSON.ID, Success: false}
			continue
		}

		cpuQuota := int64(cpu * float64(utils.CpuPeriodBase))
		newCPUQuota := containerJSON.HostConfig.CPUQuota + cpuQuota
		newMemory := containerJSON.HostConfig.Memory + memory
		if newCPUQuota <= 0 || newMemory <= MEMORY_LOW_LIMIT {
			log.Warnf("[relloc] new resource invaild %s, %d, %d", containerJSON.ID, newCPUQuota, newMemory)
			ch <- &types.ReallocResourceMessage{ContainerID: containerJSON.ID, Success: false}
			continue
		}

		// CPUQuota not cpu
		newResource := c.makeMemoryPriorSetting(newMemory, float64(newCPUQuota)/float64(utils.CpuPeriodBase))
		updateConfig := enginecontainer.UpdateConfig{Resources: newResource}
		if err := reSetContainer(containerJSON.ID, node, updateConfig); err != nil {
			log.Errorf("[realloc] update container failed %v, %s", err, containerJSON.ID)
			ch <- &types.ReallocResourceMessage{ContainerID: containerJSON.ID, Success: false}
			// 如果是增加内存，失败的时候应该把内存还回去
			if memory > 0 {
				if err := c.store.UpdateNodeMem(podname, node.Name, memory, "+"); err != nil {
					log.Errorf("[realloc] failed to set mem back %s", containerJSON.ID)
				}
			}
			continue
		}
		// 如果是要降低内存，当执行成功的时候需要把内存还回去
		if memory < 0 {
			if err := c.store.UpdateNodeMem(podname, node.Name, -memory, "+"); err != nil {
				log.Errorf("[realloc] failed to set mem back %s", containerJSON.ID)
			}
		}
		if _, err := c.store.AddContainer(container.ID, podname, node.Name, container.Name, nil, container.Memory+memory); err != nil {
			log.Errorf("[realloc] update container meta failed %v", err)
			// 立即中断
			ch <- &types.ReallocResourceMessage{ContainerID: container.ID, Success: false}
			return
		}
		ch <- &types.ReallocResourceMessage{ContainerID: containerJSON.ID, Success: true}
	}
}

func (c *calcium) reallocNodesCPU(
	podname string,
	nodesInfoMap CPUNodeContainers,
) (CPUNodeContainersMap, error) {

	lock, err := c.Lock(podname, 30)
	if err != nil {
		return nil, err
	}
	defer lock.Unlock()

	// TODO too slow
	nodesCPUMap := CPUNodeContainersMap{}
	for cpu, nodesInfo := range nodesInfoMap {
		for node, containers := range nodesInfo {
			for _, container := range containers {
				// 把 CPU 还回去，变成新的可用资源
				// 即便有并发操作，不影响 Create 操作
				// 最坏情况就是 CPU 重叠了，可以外部纠正
				if err := c.store.UpdateNodeCPU(podname, node.Name, container.CPU, "+"); err != nil {
					return nil, err
				}
				node.CPU.Add(container.CPU)
			}

			// 按照 Node one by one 重新计算可以部署多少容器
			containersNum := len(containers)
			nodesInfo := []types.NodeInfo{
				types.NodeInfo{
					CPUAndMem: types.CPUAndMem{
						CpuMap: node.CPU,
						MemCap: 0,
					},
					Name: node.Name,
				},
			}

			result, changed, err := c.scheduler.SelectCPUNodes(nodesInfo, cpu, containersNum)
			if err != nil {
				return nil, err
			}

			nodeCPUMap, isChanged := changed[node.Name]
			containersCPUMap, hasResult := result[node.Name]
			if isChanged && hasResult {
				node.CPU = nodeCPUMap
				if err := c.store.UpdateNode(node); err != nil {
					return nil, err
				}
				if _, ok := nodesCPUMap[cpu]; !ok {
					nodesCPUMap[cpu] = NodeCPUMap{}
				}
				nodesCPUMap[cpu][node] = containersCPUMap
			}
		}
	}
	return nodesCPUMap, nil
}

// mem not used in this prior
func (c *calcium) reallocContainersWithCPUPrior(
	ch chan *types.ReallocResourceMessage,
	pod *types.Pod,
	nodesInfoMap CPUNodeContainers,
	cpu float64, memory int64) {

	nodesCPUMap, err := c.reallocNodesCPU(pod.Name, nodesInfoMap)
	if err != nil {
		log.Errorf("[realloc] realloc cpu resource failed %v", err)
		return
	}

	// 不并发操作了
	for cpu, nodesCPUResult := range nodesCPUMap {
		c.doReallocContainersWithCPUPrior(ch, pod.Name, nodesCPUResult, nodesInfoMap[cpu])
	}
}

func (c *calcium) doReallocContainersWithCPUPrior(
	ch chan *types.ReallocResourceMessage,
	podname string,
	nodesCPUResult NodeCPUMap,
	nodesInfoMap NodeContainers,
) {

	for node, cpuset := range nodesCPUResult {
		containers := nodesInfoMap[node]
		for index, container := range containers {
			quota := cpuset[index]
			resource := c.makeCPUPriorSetting(quota)
			updateConfig := enginecontainer.UpdateConfig{Resources: resource}
			if err := reSetContainer(container.ID, node, updateConfig); err != nil {
				log.Errorf("[realloc] update container failed %v", err)
				// TODO 这里理论上是可以恢复 CPU 占用表的，一来我们知道新的占用是怎样，二来我们也晓得老的占用是啥样
				ch <- &types.ReallocResourceMessage{ContainerID: container.ID, Success: false}
			}
			if _, err := c.store.AddContainer(container.ID, container.Podname, container.Nodename, container.Name, quota, container.Memory); err != nil {
				log.Errorf("[realloc] update container meta failed %v", err)
				// 立即中断
				ch <- &types.ReallocResourceMessage{ContainerID: container.ID, Success: false}
				return
			}
			ch <- &types.ReallocResourceMessage{ContainerID: container.ID, Success: true}
		}
	}
}

func reSetContainer(ID string, node *types.Node, config enginecontainer.UpdateConfig) error {
	_, err := node.Engine.ContainerUpdate(context.Background(), ID, config)
	return err
}
<commit_msg>return back resource when alloc failed<commit_after>package calcium

import (
	"context"
	"fmt"
	"sync"

	log "github.com/Sirupsen/logrus"
	enginecontainer "github.com/docker/docker/api/types/container"
	"gitlab.ricebook.net/platform/core/types"
	"gitlab.ricebook.net/platform/core/utils"
)

type NodeContainers map[*types.Node][]*types.Container
type NodeCPUMap map[*types.Node][]types.CPUMap
type CPUNodeContainers map[float64]NodeContainers
type CPUNodeContainersMap map[float64]NodeCPUMap

func (c *calcium) ReallocResource(ids []string, cpu float64, mem int64) (chan *types.ReallocResourceMessage, error) {
	// TODO 大量容器 Get 的时候有性能问题
	containers, err := c.store.GetContainers(ids)
	if err != nil {
		return nil, err
	}

	// Pod-Node-Containers 三元组
	containersInfo := map[*types.Pod]NodeContainers{}
	// Pod-cpu-node-containers 四元组
	cpuContainersInfo := map[*types.Pod]CPUNodeContainers{}
	nodeCache := map[string]*types.Node{}

	for _, container := range containers {
		if _, ok := nodeCache[container.Nodename]; !ok {
			node, err := c.GetNode(container.Podname, container.Nodename)
			if err != nil {
				return nil, err
			}
			nodeCache[container.Nodename] = node
		}
		node := nodeCache[container.Nodename]

		pod, err := c.store.GetPod(container.Podname)
		if err != nil {
			return nil, err
		}

		if _, ok := containersInfo[pod]; !ok {
			containersInfo[pod] = NodeContainers{}
		}
		if _, ok := containersInfo[pod][node]; !ok {
			containersInfo[pod][node] = []*types.Container{}
		}
		containersInfo[pod][node] = append(containersInfo[pod][node], container)

		if pod.Favor == types.CPU_PRIOR {
			if _, ok := cpuContainersInfo[pod]; !ok {
				cpuContainersInfo[pod] = CPUNodeContainers{}
			}
			podCPUContainersInfo := cpuContainersInfo[pod]
			newCPURequire := c.calculateCPUUsage(container) + cpu
			if newCPURequire < 0.0 {
				return nil, fmt.Errorf("[realloc] cpu can not below zero")
			}
			if _, ok := podCPUContainersInfo[newCPURequire]; !ok {
				podCPUContainersInfo[newCPURequire] = NodeContainers{} // go map 的傻逼语法, 不这样写会有 assignment to entry in nil map 错误
				podCPUContainersInfo[newCPURequire][node] = []*types.Container{}
			}
			podCPUContainersInfo[newCPURequire][node] = append(podCPUContainersInfo[newCPURequire][node], container)
		}
	}

	ch := make(chan *types.ReallocResourceMessage)
	go func() {
		defer close(ch)
		wg := sync.WaitGroup{}
		wg.Add(len(containersInfo))
		for pod, nodeContainers := range containersInfo {
			if pod.Favor == types.CPU_PRIOR {
				nodeCPUContainersInfo := cpuContainersInfo[pod]
				go func(pod *types.Pod, nodeCPUContainersInfo CPUNodeContainers) {
					defer wg.Done()
					c.reallocContainersWithCPUPrior(ch, pod, nodeCPUContainersInfo, cpu, mem)
				}(pod, nodeCPUContainersInfo)
				continue
			}
			go func(pod *types.Pod, nodeContainers NodeContainers) {
				defer wg.Done()
				c.reallocContainerWithMemoryPrior(ch, pod, nodeContainers, cpu, mem)
			}(pod, nodeContainers)
		}
		wg.Wait()
	}()
	return ch, nil
}

func (c *calcium) reallocNodesMemory(podname string, nodeContainers NodeContainers, memory int64) error {
	lock, err := c.Lock(podname, 30)
	if err != nil {
		return err
	}
	defer lock.Unlock()
	for node, containers := range nodeContainers {
		if memory > 0 {
			if cap := int(node.MemCap / memory); cap < len(containers) {
				return fmt.Errorf("[realloc] Not enough resource %s", node.Name)
			}
		}
		if err := c.store.UpdateNodeMem(podname, node.Name, int64(len(containers))*memory, "-"); err != nil {
			return err
		}
	}
	return nil
}

func (c *calcium) reallocContainerWithMemoryPrior(
	ch chan *types.ReallocResourceMessage,
	pod *types.Pod,
	nodeContainers NodeContainers,
	cpu float64, memory int64) {

	if err := c.reallocNodesMemory(pod.Name, nodeContainers, memory); err != nil {
		log.Errorf("[realloc] realloc memory failed %v", err)
		return
	}

	// 不并发操作了
	for node, containers := range nodeContainers {
		c.doUpdateContainerWithMemoryPrior(ch, pod.Name, node, containers, cpu, memory)
	}
}

func (c *calcium) doUpdateContainerWithMemoryPrior(
	ch chan *types.ReallocResourceMessage,
	podname string,
	node *types.Node,
	containers []*types.Container,
	cpu float64, memory int64) {

	for _, container := range containers {
		containerJSON, err := container.Inspect()
		if err != nil {
			log.Errorf("[realloc] get container failed %v", err)
			ch <- &types.ReallocResourceMessage{ContainerID: containerJSON.ID, Success: false}
			continue
		}

		cpuQuota := int64(cpu * float64(utils.CpuPeriodBase))
		newCPUQuota := containerJSON.HostConfig.CPUQuota + cpuQuota
		newMemory := containerJSON.HostConfig.Memory + memory
		if newCPUQuota <= 0 || newMemory <= MEMORY_LOW_LIMIT {
			log.Warnf("[relloc] new resource invaild %s, %d, %d", containerJSON.ID, newCPUQuota, newMemory)
			ch <- &types.ReallocResourceMessage{ContainerID: containerJSON.ID, Success: false}
			continue
		}

		// CPUQuota not cpu
		newResource := c.makeMemoryPriorSetting(newMemory, float64(newCPUQuota)/float64(utils.CpuPeriodBase))
		updateConfig := enginecontainer.UpdateConfig{Resources: newResource}
		if err := reSetContainer(containerJSON.ID, node, updateConfig); err != nil {
			log.Errorf("[realloc] update container failed %v, %s", err, containerJSON.ID)
			ch <- &types.ReallocResourceMessage{ContainerID: containerJSON.ID, Success: false}
			// 如果是增加内存，失败的时候应该把内存还回去
			if memory > 0 {
				if err := c.store.UpdateNodeMem(podname, node.Name, memory, "+"); err != nil {
					log.Errorf("[realloc] failed to set mem back %s", containerJSON.ID)
				}
			}
			continue
		}
		// 如果是要降低内存，当执行成功的时候需要把内存还回去
		if memory < 0 {
			if err := c.store.UpdateNodeMem(podname, node.Name, -memory, "+"); err != nil {
				log.Errorf("[realloc] failed to set mem back %s", containerJSON.ID)
			}
		}
		if _, err := c.store.AddContainer(container.ID, podname, node.Name, container.Name, nil, container.Memory+memory); err != nil {
			log.Errorf("[realloc] update container meta failed %v", err)
			// 立即中断
			ch <- &types.ReallocResourceMessage{ContainerID: container.ID, Success: false}
			return
		}
		ch <- &types.ReallocResourceMessage{ContainerID: containerJSON.ID, Success: true}
	}
}

func (c *calcium) reallocNodesCPU(
	podname string,
	nodesInfoMap CPUNodeContainers,
) (CPUNodeContainersMap, error) {

	lock, err := c.Lock(podname, 30)
	if err != nil {
		return nil, err
	}
	defer lock.Unlock()

	// TODO too slow
	nodesCPUMap := CPUNodeContainersMap{}
	for cpu, nodesInfo := range nodesInfoMap {
		for node, containers := range nodesInfo {
			for _, container := range containers {
				// 把 CPU 还回去，变成新的可用资源
				// 即便有并发操作，不影响 Create 操作
				// 最坏情况就是 CPU 重叠了，可以外部纠正
				if err := c.store.UpdateNodeCPU(podname, node.Name, container.CPU, "+"); err != nil {
					return nil, err
				}
				node.CPU.Add(container.CPU)
			}

			// 按照 Node one by one 重新计算可以部署多少容器
			containersNum := len(containers)
			nodesInfo := []types.NodeInfo{
				types.NodeInfo{
					CPUAndMem: types.CPUAndMem{
						CpuMap: node.CPU,
						MemCap: 0,
					},
					Name: node.Name,
				},
			}

			result, changed, err := c.scheduler.SelectCPUNodes(nodesInfo, cpu, containersNum)
			if err != nil {
				for _, container := range containers {
					if err := c.store.UpdateNodeCPU(podname, node.Name, container.CPU, "-"); err != nil {
						return nil, err
					}
				}
				return nil, err
			}

			nodeCPUMap, isChanged := changed[node.Name]
			containersCPUMap, hasResult := result[node.Name]
			if isChanged && hasResult {
				node.CPU = nodeCPUMap
				if err := c.store.UpdateNode(node); err != nil {
					return nil, err
				}
				if _, ok := nodesCPUMap[cpu]; !ok {
					nodesCPUMap[cpu] = NodeCPUMap{}
				}
				nodesCPUMap[cpu][node] = containersCPUMap
			}
		}
	}
	return nodesCPUMap, nil
}

// mem not used in this prior
func (c *calcium) reallocContainersWithCPUPrior(
	ch chan *types.ReallocResourceMessage,
	pod *types.Pod,
	nodesInfoMap CPUNodeContainers,
	cpu float64, memory int64) {

	nodesCPUMap, err := c.reallocNodesCPU(pod.Name, nodesInfoMap)
	if err != nil {
		log.Errorf("[realloc] realloc cpu resource failed %v", err)
		return
	}

	// 不并发操作了
	for cpu, nodesCPUResult := range nodesCPUMap {
		c.doReallocContainersWithCPUPrior(ch, pod.Name, nodesCPUResult, nodesInfoMap[cpu])
	}
}

func (c *calcium) doReallocContainersWithCPUPrior(
	ch chan *types.ReallocResourceMessage,
	podname string,
	nodesCPUResult NodeCPUMap,
	nodesInfoMap NodeContainers,
) {

	for node, cpuset := range nodesCPUResult {
		containers := nodesInfoMap[node]
		for index, container := range containers {
			quota := cpuset[index]
			resource := c.makeCPUPriorSetting(quota)
			updateConfig := enginecontainer.UpdateConfig{Resources: resource}
			if err := reSetContainer(container.ID, node, updateConfig); err != nil {
				log.Errorf("[realloc] update container failed %v", err)
				// TODO 这里理论上是可以恢复 CPU 占用表的，一来我们知道新的占用是怎样，二来我们也晓得老的占用是啥样
				ch <- &types.ReallocResourceMessage{ContainerID: container.ID, Success: false}
			}
			if _, err := c.store.AddContainer(container.ID, container.Podname, container.Nodename, container.Name, quota, container.Memory); err != nil {
				log.Errorf("[realloc] update container meta failed %v", err)
				// 立即中断
				ch <- &types.ReallocResourceMessage{ContainerID: container.ID, Success: false}
				return
			}
			ch <- &types.ReallocResourceMessage{ContainerID: container.ID, Success: true}
		}
	}
}

func reSetContainer(ID string, node *types.Node, config enginecontainer.UpdateConfig) error {
	_, err := node.Engine.ContainerUpdate(context.Background(), ID, config)
	return err
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"os"
	"strings"

	"github.com/msoap/html2data"
)

const usageString = "Usage:\n" +
	"  html2data [options] [url|file|-] 'css selector'\n" +
	"  html2data [options] [url|file|-] :name1 'css1' :name2 'css2' ...\n\n" +
	"options:"

type cmdConfig struct {
	userAgent, outerCSS, url string
	timeOut                  int
	getJSON                  bool
	dontTrimSpaces           bool
	dontDetectCharset        bool
}

var (
	config cmdConfig
)

func init() {
	flag.StringVar(&config.userAgent, "user-agent", "", "set custom user-agent")
	flag.StringVar(&config.outerCSS, "find-in", "", "search in the specified elements instead document")
	flag.BoolVar(&config.getJSON, "json", false, "JSON output")
	flag.BoolVar(&config.dontTrimSpaces, "dont-trim-spaces", false, "dont trim spaces, get text as is")
	flag.BoolVar(&config.dontDetectCharset, "dont-detect-charset", false, "dont detect charset and convert text")
	flag.IntVar(&config.timeOut, "timeout", 0, "timeout in seconds")
}

func getConfig() (CSSSelectors map[string]string, err error) {
	flag.Usage = func() {
		fmt.Println(usageString)
		flag.PrintDefaults()
		os.Exit(0)
	}
	flag.Parse()

	config.url, CSSSelectors, err = parseArgs(flag.Args())
	return CSSSelectors, err
}

// printAsText - print result as text
func printAsText(texts map[string][]string, doPrintName bool) {
	for name, value := range texts {
		if doPrintName {
			fmt.Print(name + ":\t")
		}
		for _, text := range value {
			fmt.Println(text)
		}
	}
}

func runApp() error {
	CSSSelectors, err := getConfig()
	if err != nil {
		return err
	}
	var doc html2data.Doc
	stat, err := os.Stdin.Stat()
	if err != nil {
		return err
	}

	if config.url == "-" || (stat.Mode()&os.ModeCharDevice) == 0 {
		reader := bufio.NewReader(os.Stdin)
		doc = html2data.FromReader(reader)
	} else if strings.HasPrefix(config.url, "http://") || strings.HasPrefix(config.url, "https://") {
		doc = html2data.FromURL(config.url, html2data.URLCfg{UA: config.userAgent, TimeOut: config.timeOut, DontDetectCharset: config.dontDetectCharset})
	} else if len(config.url) > 0 {
		doc = html2data.FromFile(config.url)
	} else {
		fmt.Println(usageString)
		return nil
	}

	GetDocCfg := html2data.Cfg{DontTrimSpaces: config.dontTrimSpaces}
	if config.outerCSS != "" {
		textsOuter, err := doc.GetDataNested(config.outerCSS, CSSSelectors, GetDocCfg)
		if err != nil {
			return err
		}

		if config.getJSON {
			json, _ := json.Marshal(textsOuter)
			fmt.Println(string(json))
		} else {
			for i, texts := range textsOuter {
				fmt.Printf("%d:\n", i)
				printAsText(texts, len(CSSSelectors) > 1)
			}
		}
	} else {
		texts, err := doc.GetData(CSSSelectors, GetDocCfg)
		if err != nil {
			return err
		}

		if config.getJSON {
			json, _ := json.Marshal(texts)
			fmt.Println(string(json))
		} else {
			printAsText(texts, len(CSSSelectors) > 1)
		}
	}

	return nil
}

func main() {
	err := runApp()
	if err != nil {
		log.Fatal(err)
	}
}
<commit_msg>Don't reuse package name for variable<commit_after>package main

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"os"
	"strings"

	"github.com/msoap/html2data"
)

const usageString = "Usage:\n" +
	"  html2data [options] [url|file|-] 'css selector'\n" +
	"  html2data [options] [url|file|-] :name1 'css1' :name2 'css2' ...\n\n" +
	"options:"

type cmdConfig struct {
	userAgent, outerCSS, url string
	timeOut                  int
	getJSON                  bool
	dontTrimSpaces           bool
	dontDetectCharset        bool
}

var (
	config cmdConfig
)

func init() {
	flag.StringVar(&config.userAgent, "user-agent", "", "set custom user-agent")
	flag.StringVar(&config.outerCSS, "find-in", "", "search in the specified elements instead document")
	flag.BoolVar(&config.getJSON, "json", false, "JSON output")
	flag.BoolVar(&config.dontTrimSpaces, "dont-trim-spaces", false, "dont trim spaces, get text as is")
	flag.BoolVar(&config.dontDetectCharset, "dont-detect-charset", false, "dont detect charset and convert text")
	flag.IntVar(&config.timeOut, "timeout", 0, "timeout in seconds")
}

func getConfig() (CSSSelectors map[string]string, err error) {
	flag.Usage = func() {
		fmt.Println(usageString)
		flag.PrintDefaults()
		os.Exit(0)
	}
	flag.Parse()

	config.url, CSSSelectors, err = parseArgs(flag.Args())
	return CSSSelectors, err
}

// printAsText - print result as text
func printAsText(texts map[string][]string, doPrintName bool) {
	for name, value := range texts {
		if doPrintName {
			fmt.Print(name + ":\t")
		}
		for _, text := range value {
			fmt.Println(text)
		}
	}
}

func runApp() error {
	CSSSelectors, err := getConfig()
	if err != nil {
		return err
	}
	var doc html2data.Doc
	stat, err := os.Stdin.Stat()
	if err != nil {
		return err
	}

	if config.url == "-" || (stat.Mode()&os.ModeCharDevice) == 0 {
		reader := bufio.NewReader(os.Stdin)
		doc = html2data.FromReader(reader)
	} else if strings.HasPrefix(config.url, "http://") || strings.HasPrefix(config.url, "https://") {
		doc = html2data.FromURL(config.url, html2data.URLCfg{UA: config.userAgent, TimeOut: config.timeOut, DontDetectCharset: config.dontDetectCharset})
	} else if len(config.url) > 0 {
		doc = html2data.FromFile(config.url)
	} else {
		fmt.Println(usageString)
		return nil
	}

	GetDocCfg := html2data.Cfg{DontTrimSpaces: config.dontTrimSpaces}
	if config.outerCSS != "" {
		textsOuter, err := doc.GetDataNested(config.outerCSS, CSSSelectors, GetDocCfg)
		if err != nil {
			return err
		}

		if config.getJSON {
			jsonBytes, _ := json.Marshal(textsOuter)
			fmt.Println(string(jsonBytes))
		} else {
			for i, texts := range textsOuter {
				fmt.Printf("%d:\n", i)
				printAsText(texts, len(CSSSelectors) > 1)
			}
		}
	} else {
		texts, err := doc.GetData(CSSSelectors, GetDocCfg)
		if err != nil {
			return err
		}

		if config.getJSON {
			jsonBytes, _ := json.Marshal(texts)
			fmt.Println(string(jsonBytes))
		} else {
			printAsText(texts, len(CSSSelectors) > 1)
		}
	}

	return nil
}

func main() {
	err := runApp()
	if err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>