<commit_before>package protomongo

import (
	"fmt"
	"reflect"
	"strconv"
	"strings"
	"sync"

	"github.com/golang/protobuf/descriptor"
	"github.com/golang/protobuf/proto"
	"go.mongodb.org/mongo-driver/bson/bsoncodec"
	"go.mongodb.org/mongo-driver/bson/bsonrw"
)

type protobufCodec struct {
	m            *sync.RWMutex
	protoHelpers map[string]*protoHelper
}

// Returns a new instance of protobufCodec. protobufCodec is a MongoDB codec. It encodes protobuf objects using the protobuf
// field numbers as document keys. This means that stored protobufs can survive normal protobuf definition changes, e.g. renaming a field.
func NewProtobufCodec() *protobufCodec {
	return &protobufCodec{protoHelpers: make(map[string]*protoHelper), m: &sync.RWMutex{}}
}

func (pc *protobufCodec) EncodeValue(ectx bsoncodec.EncodeContext, vw bsonrw.ValueWriter, val reflect.Value) error {
	origAsDescMsg := val.Interface().(descriptor.Message)
	for val.Kind() != reflect.Struct {
		val = val.Elem()
	}

	dw, err := vw.WriteDocument()
	if err != nil {
		return err
	}

	ph := pc.protoHelper(origAsDescMsg, val.Type())

	for _, prop := range ph.normalPropsByTag {
		fVal := val.FieldByName(prop.Name)
		if !fVal.IsZero() {
			if err := encodeField(ectx, dw, prop.Tag, fVal); err != nil {
				return err
			}
		}
	}

	for _, prop := range ph.oneofFieldWrapperProps {
		fVal := val.FieldByName(prop.Name)
		if !fVal.IsZero() {
			// Since this field is a oneof, we need to get the single Go value stored inside its oneof wrapper struct,
			// instead of simply using the value as-is.
			oneof := fVal.Elem().Elem()
			singleProp := proto.GetProperties(oneof.Type()).Prop[0]
			fVal = oneof.Field(0)
			if err := encodeField(ectx, dw, singleProp.Tag, fVal); err != nil {
				return err
			}
		}
	}

	return dw.WriteDocumentEnd()
}

func encodeField(ectx bsoncodec.EncodeContext, dw bsonrw.DocumentWriter, tag int, fVal reflect.Value) error {
	fvw, err := dw.WriteDocumentElement(TagToElementName(tag))
	if err != nil {
		return err
	}
	enc, err := ectx.LookupEncoder(fVal.Type())
	if err != nil {
		return err
	}
	return enc.EncodeValue(ectx, fvw, fVal)
}

func (pc *protobufCodec) DecodeValue(ectx bsoncodec.DecodeContext, vr bsonrw.ValueReader, val reflect.Value) error {
	origAsDescMsg := val.Interface().(descriptor.Message)
	if val.IsNil() {
		val.Set(reflect.New(val.Type().Elem()))
	}
	for val.Kind() != reflect.Struct {
		val = val.Elem()
	}

	dr, err := vr.ReadDocument()
	if err != nil {
		return err
	}

	ph := pc.protoHelper(origAsDescMsg, val.Type())
	for f, fvr, err := dr.ReadElement(); err != bsonrw.ErrEOD; f, fvr, err = dr.ReadElement() {
		if err != nil {
			return err
		}

		tag := elementNameToTag(f)
		prop, isNotOneOf := ph.normalPropsByTag[tag]
		oneof, isOneof := ph.oneofPropsByTag[tag]

		// Skip any field that we don't recognize.
		if !isNotOneOf && !isOneof {
			if err = vr.Skip(); err != nil {
				return err
			}
			continue
		}

		// Figure out what field we need to decode into.
		var fVal reflect.Value
		if isNotOneOf {
			fVal = val.FieldByName(prop.Name)
		} else if isOneof {
			oneofVal := reflect.New(oneof.Type.Elem())
			val.Field(oneof.Field).Set(oneofVal)
			fVal = oneofVal.Elem().Field(0)
		}

		// Actually decode the value.
		if err = lookupDecoderAndDecode(ectx, fvr, fVal); err != nil && isNotOneOf {
			// It's possible that this value was encoded as a repeated field and is now being decoded as non-repeated field,
			// or vice-versa. Since this would count as a valid protobuf change, we try to decode as the opposite type.
			// If this decoding attempt fails, we return the original decode error.
			if prop.Repeated {
				singleVal := reflect.New(fVal.Type().Elem()).Elem()
				if backupDecodeErr := lookupDecoderAndDecode(ectx, fvr, singleVal); backupDecodeErr != nil {
					return err
				}
				fVal.Set(reflect.Append(fVal, singleVal))
			} else {
				repeatedVal := reflect.New(reflect.SliceOf(fVal.Type())).Elem()
				if backupDecodeErr := lookupDecoderAndDecode(ectx, fvr, repeatedVal); backupDecodeErr != nil {
					return err
				}

				for i := 0; i < repeatedVal.Len(); i++ {
					singleVal := repeatedVal.Index(i)
					// Following the rules here at https://developers.google.com/protocol-buffers/docs/proto#updating,
					// repeated Message values must be merged into a single value.
					if _, fieldIsMessage := fVal.Interface().(proto.Message); fieldIsMessage {
						proto.Merge(fVal.Interface().(proto.Message), singleVal.Interface().(proto.Message))
					} else {
						fVal.Set(singleVal)
					}
				}
			}
		}
	}

	return nil
}

func lookupDecoderAndDecode(ectx bsoncodec.DecodeContext, vr bsonrw.ValueReader, val reflect.Value) error {
	enc, err := ectx.LookupDecoder(val.Type())
	if err != nil {
		return err
	}
	return enc.DecodeValue(ectx, vr, val)
}

type protoHelper struct {
	// Properties corresponding to 'normal' (non-oneof) protobuf fields.
	// Indexed by protobuf tag number (as a string).
	normalPropsByTag map[string]*proto.Properties
	// OneofProperties corresponding to oneof protobuf fields.
	// Indexed by protobuf tag number (as a string).
	oneofPropsByTag map[string]*proto.OneofProperties

	// Properties corresponding to Go wrapper types for protobuf oneof declarations.
	oneofFieldWrapperProps []*proto.Properties
}

func (pc *protobufCodec) protoHelper(pb descriptor.Message, t reflect.Type) *protoHelper {
	// Try to load a pre-existing protoHelper from cache, if it exists.
	messageName := proto.MessageName(pb)
	if ph, ok := pc.lockedGetProtoHelper(messageName); ok {
		return ph
	}

	// Find the names of all oneofs.
	_, msgDescriptor := descriptor.ForMessage(pb)
	oneofNames := make(map[string]bool)
	for _, oneof := range msgDescriptor.OneofDecl {
		oneofNames[*oneof.Name] = true
	}

	// Get the corresponding Go type's Properties, and divide them into three groups.
	// See comments on 'protoHelper' for details.
	props := proto.GetProperties(t)
	oneofFieldWrapperProps := make([]*proto.Properties, 0)
	normalPropsByTag := make(map[string]*proto.Properties)
	for _, prop := range props.Prop {
		if oneofNames[prop.OrigName] {
			oneofFieldWrapperProps = append(oneofFieldWrapperProps, prop)
		} else {
			normalPropsByTag[strconv.Itoa(prop.Tag)] = prop
		}
	}
	oneofPropsByTag := make(map[string]*proto.OneofProperties)
	for _, oneof := range props.OneofTypes {
		oneofPropsByTag[strconv.Itoa(oneof.Prop.Tag)] = oneof
	}
	ph := &protoHelper{normalPropsByTag, oneofPropsByTag, oneofFieldWrapperProps}
	pc.lockedSetProtoHelper(messageName, ph)
	return ph
}

func (pc *protobufCodec) lockedGetProtoHelper(messageName string) (*protoHelper, bool) {
	pc.m.RLock()
	defer pc.m.RUnlock()
	ph, ok := pc.protoHelpers[messageName]
	return ph, ok
}

func (pc *protobufCodec) lockedSetProtoHelper(messageName string, ph *protoHelper) {
	pc.m.Lock()
	defer pc.m.Unlock()
	pc.protoHelpers[messageName] = ph
}

const (
	tagPrefix = "PBTag_"
)

func TagToElementName(tag int) string {
	return fmt.Sprintf("%v%v", tagPrefix, tag)
}

func elementNameToTag(elementName string) string {
	return strings.Replace(elementName, tagPrefix, "", 1)
}
<commit_msg>Future-proof protomongo package against changes to the Go protobuf package. (#888)<commit_after>package protomongo

import (
	"fmt"
	"reflect"
	"strconv"
	"strings"
	"sync"

	"github.com/golang/protobuf/descriptor"
	"github.com/golang/protobuf/proto"
	"go.mongodb.org/mongo-driver/bson/bsoncodec"
	"go.mongodb.org/mongo-driver/bson/bsonrw"
)

type protobufCodec struct {
	m            *sync.RWMutex
	protoHelpers map[string]*protoHelper
}

// Returns a new instance of protobufCodec. protobufCodec is a MongoDB codec. It encodes protobuf objects using the protobuf
// field numbers as document keys. This means that stored protobufs can survive normal protobuf definition changes, e.g. renaming a field.
func NewProtobufCodec() *protobufCodec {
	return &protobufCodec{protoHelpers: make(map[string]*protoHelper), m: &sync.RWMutex{}}
}

func (pc *protobufCodec) EncodeValue(ectx bsoncodec.EncodeContext, vw bsonrw.ValueWriter, val reflect.Value) error {
	origAsDescMsg := val.Interface().(descriptor.Message)
	for val.Kind() != reflect.Struct {
		val = val.Elem()
	}

	dw, err := vw.WriteDocument()
	if err != nil {
		return err
	}

	ph := pc.protoHelper(origAsDescMsg, val.Type())

	for _, prop := range ph.normalPropsByTag {
		fVal := val.FieldByName(prop.Name)
		if !fVal.IsZero() {
			if err := encodeField(ectx, dw, prop.Tag, fVal); err != nil {
				return err
			}
		}
	}

	for _, prop := range ph.oneofFieldWrapperProps {
		fVal := val.FieldByName(prop.Name)
		if !fVal.IsZero() {
			// Since this field is a oneof, we need to get the single Go value stored inside its oneof wrapper struct,
			// instead of simply using the value as-is.
			oneof := fVal.Elem().Elem()
			singleProp := proto.GetProperties(oneof.Type()).Prop[0]
			fVal = oneof.Field(0)
			if err := encodeField(ectx, dw, singleProp.Tag, fVal); err != nil {
				return err
			}
		}
	}

	return dw.WriteDocumentEnd()
}

func encodeField(ectx bsoncodec.EncodeContext, dw bsonrw.DocumentWriter, tag int, fVal reflect.Value) error {
	fvw, err := dw.WriteDocumentElement(TagToElementName(tag))
	if err != nil {
		return err
	}
	enc, err := ectx.LookupEncoder(fVal.Type())
	if err != nil {
		return err
	}
	return enc.EncodeValue(ectx, fvw, fVal)
}

func (pc *protobufCodec) DecodeValue(ectx bsoncodec.DecodeContext, vr bsonrw.ValueReader, val reflect.Value) error {
	origAsDescMsg := val.Interface().(descriptor.Message)
	if val.IsNil() {
		val.Set(reflect.New(val.Type().Elem()))
	}
	for val.Kind() != reflect.Struct {
		val = val.Elem()
	}

	dr, err := vr.ReadDocument()
	if err != nil {
		return err
	}

	ph := pc.protoHelper(origAsDescMsg, val.Type())
	for f, fvr, err := dr.ReadElement(); err != bsonrw.ErrEOD; f, fvr, err = dr.ReadElement() {
		if err != nil {
			return err
		}

		tag := elementNameToTag(f)
		prop, isNotOneOf := ph.normalPropsByTag[tag]
		oneof, isOneof := ph.oneofPropsByTag[tag]

		// Skip any field that we don't recognize.
		if !isNotOneOf && !isOneof {
			if err = vr.Skip(); err != nil {
				return err
			}
			continue
		}

		// Figure out what field we need to decode into.
		var fVal reflect.Value
		if isNotOneOf {
			fVal = val.FieldByName(prop.Name)
		} else if isOneof {
			oneofVal := reflect.New(oneof.Type.Elem())
			val.Field(oneof.Field).Set(oneofVal)
			fVal = oneofVal.Elem().Field(0)
		}

		// Actually decode the value.
		if err = lookupDecoderAndDecode(ectx, fvr, fVal); err != nil && isNotOneOf {
			// It's possible that this value was encoded as a repeated field and is now being decoded as non-repeated field,
			// or vice-versa. Since this would count as a valid protobuf change, we try to decode as the opposite type.
			// If this decoding attempt fails, we return the original decode error.
			if prop.Repeated {
				singleVal := reflect.New(fVal.Type().Elem()).Elem()
				if backupDecodeErr := lookupDecoderAndDecode(ectx, fvr, singleVal); backupDecodeErr != nil {
					return err
				}
				fVal.Set(reflect.Append(fVal, singleVal))
			} else {
				repeatedVal := reflect.New(reflect.SliceOf(fVal.Type())).Elem()
				if backupDecodeErr := lookupDecoderAndDecode(ectx, fvr, repeatedVal); backupDecodeErr != nil {
					return err
				}

				for i := 0; i < repeatedVal.Len(); i++ {
					singleVal := repeatedVal.Index(i)
					// Following the rules here at https://developers.google.com/protocol-buffers/docs/proto#updating,
					// repeated Message values must be merged into a single value.
					if _, fieldIsMessage := fVal.Interface().(proto.Message); fieldIsMessage {
						proto.Merge(fVal.Interface().(proto.Message), singleVal.Interface().(proto.Message))
					} else {
						fVal.Set(singleVal)
					}
				}
			}
		}
	}

	return nil
}

func lookupDecoderAndDecode(ectx bsoncodec.DecodeContext, vr bsonrw.ValueReader, val reflect.Value) error {
	enc, err := ectx.LookupDecoder(val.Type())
	if err != nil {
		return err
	}
	return enc.DecodeValue(ectx, vr, val)
}

type protoHelper struct {
	// Properties corresponding to 'normal' (non-oneof) protobuf fields.
	// Indexed by protobuf tag number (as a string).
	normalPropsByTag map[string]*proto.Properties
	// OneofProperties corresponding to oneof protobuf fields.
	// Indexed by protobuf tag number (as a string).
	oneofPropsByTag map[string]*proto.OneofProperties

	// Properties corresponding to Go wrapper types for protobuf oneof declarations.
	oneofFieldWrapperProps []*proto.Properties
}

func (pc *protobufCodec) protoHelper(pb descriptor.Message, t reflect.Type) *protoHelper {
	// Try to load a pre-existing protoHelper from cache, if it exists.
	messageName := proto.MessageName(pb)
	if ph, ok := pc.lockedGetProtoHelper(messageName); ok {
		return ph
	}

	// Find the names of all oneofs.
	_, msgDescriptor := descriptor.ForMessage(pb)
	oneofNames := make(map[string]bool)
	for _, oneof := range msgDescriptor.OneofDecl {
		oneofNames[*oneof.Name] = true
	}

	// Get the corresponding Go type's Properties, and divide them into three groups.
	// See comments on 'protoHelper' for details.
	props := proto.GetProperties(t)
	oneofFieldWrapperProps := make([]*proto.Properties, 0)
	normalPropsByTag := make(map[string]*proto.Properties)
	for _, prop := range props.Prop {
		if !strings.HasPrefix(prop.Name, "XXX_") {
			if oneofNames[prop.OrigName] {
				oneofFieldWrapperProps = append(oneofFieldWrapperProps, prop)
			} else {
				normalPropsByTag[strconv.Itoa(prop.Tag)] = prop
			}
		}
	}
	oneofPropsByTag := make(map[string]*proto.OneofProperties)
	for _, oneof := range props.OneofTypes {
		oneofPropsByTag[strconv.Itoa(oneof.Prop.Tag)] = oneof
	}
	ph := &protoHelper{normalPropsByTag, oneofPropsByTag, oneofFieldWrapperProps}
	pc.lockedSetProtoHelper(messageName, ph)
	return ph
}

func (pc *protobufCodec) lockedGetProtoHelper(messageName string) (*protoHelper, bool) {
	pc.m.RLock()
	defer pc.m.RUnlock()
	ph, ok := pc.protoHelpers[messageName]
	return ph, ok
}

func (pc *protobufCodec) lockedSetProtoHelper(messageName string, ph *protoHelper) {
	pc.m.Lock()
	defer pc.m.Unlock()
	pc.protoHelpers[messageName] = ph
}

const (
	tagPrefix = "PBTag_"
)

func TagToElementName(tag int) string {
	return fmt.Sprintf("%v%v", tagPrefix, tag)
}

func elementNameToTag(elementName string) string {
	return strings.Replace(elementName, tagPrefix, "", 1)
}
<|endoftext|>


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

package lxd

import (
	"github.com/juju/errors"

	"github.com/juju/juju/environs"
	"github.com/juju/juju/environs/config"
)

type environProvider struct{}

var providerInstance environProvider

// Open implements environs.EnvironProvider.
func (environProvider) Open(cfg *config.Config) (environs.Environ, error) {
	env, err := newEnviron(cfg, newRawProvider)
	return env, errors.Trace(err)
}

// PrepareForBootstrap implements environs.EnvironProvider.
func (p environProvider) PrepareForBootstrap(ctx environs.BootstrapContext, cfg *config.Config) (environs.Environ, error) {
	cfg, err := p.PrepareForCreateEnvironment(cfg)
	if err != nil {
		return nil, errors.Trace(err)
	}
	env, err := newEnviron(cfg, newRawProvider)
	if err != nil {
		return nil, errors.Trace(err)
	}

	if ctx.ShouldVerifyCredentials() {
		return nil, errors.NotImplementedf("")
		//if err := env.client.VerifyCredentials(); err != nil {
		//	return nil, errors.Trace(err)
		//}
	}
	return env, nil
}

// PrepareForCreateEnvironment is specified in the EnvironProvider interface.
func (environProvider) PrepareForCreateEnvironment(cfg *config.Config) (*config.Config, error) {
	return cfg, nil
}

// RestrictedConfigAttributes is specified in the EnvironProvider interface.
func (environProvider) RestrictedConfigAttributes() []string {
	// TODO(ericsnow) fix this...
	return []string{}
}

// Validate implements environs.EnvironProvider.
func (environProvider) Validate(cfg, old *config.Config) (valid *config.Config, err error) {
	if old == nil {
		ecfg, err := newValidConfig(cfg, configDefaults)
		if err != nil {
			return nil, errors.Annotate(err, "invalid config")
		}
		return ecfg.Config, nil
	}

	// The defaults should be set already, so we pass nil.
	ecfg, err := newValidConfig(old, nil)
	if err != nil {
		return nil, errors.Annotate(err, "invalid base config")
	}

	if err := ecfg.update(cfg); err != nil {
		return nil, errors.Annotate(err, "invalid config change")
	}

	return ecfg.Config, nil
}

// SecretAttrs implements environs.EnvironProvider.
func (environProvider) SecretAttrs(cfg *config.Config) (map[string]string, error) {
	// The defaults should be set already, so we pass nil.
	ecfg, err := newValidConfig(cfg, nil)
	if err != nil {
		return nil, errors.Trace(err)
	}
	return ecfg.secret(), nil
}

// BoilerplateConfig implements environs.EnvironProvider.
func (environProvider) BoilerplateConfig() string {
	// boilerplateConfig is kept in config.go, in the hope that people editing
	// config will keep it up to date.
	return boilerplateConfig
}
<commit_msg>Add some TODOs in provider.go.<commit_after>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package lxd

import (
	"github.com/juju/errors"

	"github.com/juju/juju/environs"
	"github.com/juju/juju/environs/config"
)

type environProvider struct{}

var providerInstance environProvider

// Open implements environs.EnvironProvider.
func (environProvider) Open(cfg *config.Config) (environs.Environ, error) {
	// TODO(ericsnow) verify prerequisites (see provider/local/prereq.go)?
	// TODO(ericsnow) do something similar to correctLocalhostURLs()
	// (in provider/local/environprovider.go)?

	env, err := newEnviron(cfg, newRawProvider)
	return env, errors.Trace(err)
}

// PrepareForBootstrap implements environs.EnvironProvider.
func (p environProvider) PrepareForBootstrap(ctx environs.BootstrapContext, cfg *config.Config) (environs.Environ, error) {
	cfg, err := p.PrepareForCreateEnvironment(cfg)
	if err != nil {
		return nil, errors.Trace(err)
	}

	// TODO(ericsnow) Do some of what happens in local provider's
	// PrepareForBootstrap()? Only if "remote" is local host?

	env, err := newEnviron(cfg, newRawProvider)
	if err != nil {
		return nil, errors.Trace(err)
	}

	if ctx.ShouldVerifyCredentials() {
		return nil, errors.NotImplementedf("")
		//if err := env.client.VerifyCredentials(); err != nil {
		//	return nil, errors.Trace(err)
		//}
	}
	return env, nil
}

// PrepareForCreateEnvironment is specified in the EnvironProvider interface.
func (environProvider) PrepareForCreateEnvironment(cfg *config.Config) (*config.Config, error) {
	return cfg, nil
}

// RestrictedConfigAttributes is specified in the EnvironProvider interface.
func (environProvider) RestrictedConfigAttributes() []string {
	// TODO(ericsnow) fix this...
	return []string{}
}

// Validate implements environs.EnvironProvider.
func (environProvider) Validate(cfg, old *config.Config) (valid *config.Config, err error) {
	if old == nil {
		ecfg, err := newValidConfig(cfg, configDefaults)
		if err != nil {
			return nil, errors.Annotate(err, "invalid config")
		}
		return ecfg.Config, nil
	}

	// The defaults should be set already, so we pass nil.
	ecfg, err := newValidConfig(old, nil)
	if err != nil {
		return nil, errors.Annotate(err, "invalid base config")
	}

	if err := ecfg.update(cfg); err != nil {
		return nil, errors.Annotate(err, "invalid config change")
	}

	return ecfg.Config, nil
}

// SecretAttrs implements environs.EnvironProvider.
func (environProvider) SecretAttrs(cfg *config.Config) (map[string]string, error) {
	// The defaults should be set already, so we pass nil.
	ecfg, err := newValidConfig(cfg, nil)
	if err != nil {
		return nil, errors.Trace(err)
	}
	return ecfg.secret(), nil
}

// BoilerplateConfig implements environs.EnvironProvider.
func (environProvider) BoilerplateConfig() string {
	// boilerplateConfig is kept in config.go, in the hope that people editing
	// config will keep it up to date.
	return boilerplateConfig
}
<|endoftext|>


<commit_before>package aws_ec2

import (
	"fmt"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"regexp"
	"sort"
)

type ec2_instances []*ec2.Instance

func (instance ec2_instances) Len() int {
	return len(instance)
}

func (instance ec2_instances) Swap(i, j int) {
	instance[i], instance[j] = instance[j], instance[i]
}

func (instance ec2_instances) Less(i, j int) bool {
	return GetTagValue(instance[i], "Name") < GetTagValue(instance[j], "Name")
}

func ParseFilter(filters string) []*ec2.Filter {

	// filters e.g. "Name=tag:Foo,Values=Bar Name=instance-type,Values=m1.small"
	ec2_filters := make([]*ec2.Filter, 0)

	re_space := regexp.MustCompile(`\s+`)
	re_name := regexp.MustCompile(`Name=`)
	re_values := regexp.MustCompile(`,Values=`)
	for _, i := range re_space.Split(filters, -1) {
		for _, j := range re_name.Split(i, -1) {
			if len(j) != 0 {
				v := re_values.Split(j, -1)
				name := v[0]
				/* // fix bug
				values := make([]*string, 0)
				for _, k := range v[1:] {
					fmt.Println(k)
					values = append(values, aws.String(k))
				}
				*/
				ec2_filters = append(ec2_filters, &ec2.Filter{
					Name: aws.String(name),
					Values: []*string{
						aws.String(v[1]),
					},
				})
			}
		}
	}
	return ec2_filters
}

func generateSession() (*session.Session, error) {
	return session.NewSessionWithOptions(session.Options{})
}

func DescribeInstances(filters string) (ec2_instances, error) {

	var instances ec2_instances
	sess, err := generateSession()
	if err != nil {
		return nil, err
	}
	svc := ec2.New(sess)
	params := &ec2.DescribeInstancesInput{
		Filters: ParseFilter(filters),
	}
	resp, err := svc.DescribeInstances(params)
	if err != nil {
		return nil, err
	}
	if len(resp.Reservations) == 0 {
		return ec2_instances{}, nil
	}
	for _, res := range resp.Reservations {
		for _, instance := range res.Instances {
			instances = append(instances, instance)
		}
	}
	sort.Sort(instances)
	return instances, nil
}

func PrintInstances(instances ec2_instances) {
	for _, instance := range instances {
		fmt.Println(
			GetTagValue(instance, "Name"),
			*instance.PrivateIpAddress,
			*instance.InstanceId,
			*instance.InstanceType,
			*instance.Placement.AvailabilityZone,
			*instance.State.Name,
		)
	}
}

func GetTagValue(instance *ec2.Instance, tag_name string) string {
	for _, t := range instance.Tags {
		if *t.Key == tag_name {
			return *t.Value
		}
	}
	return ""
}
<commit_msg>Allow empty filters parameter.<commit_after>package aws_ec2

import (
	"fmt"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"regexp"
	"sort"
)

type ec2_instances []*ec2.Instance

func (instance ec2_instances) Len() int {
	return len(instance)
}

func (instance ec2_instances) Swap(i, j int) {
	instance[i], instance[j] = instance[j], instance[i]
}

func (instance ec2_instances) Less(i, j int) bool {
	return GetTagValue(instance[i], "Name") < GetTagValue(instance[j], "Name")
}

func ParseFilter(filters string) []*ec2.Filter {

	// filters e.g. "Name=tag:Foo,Values=Bar Name=instance-type,Values=m1.small"
	ec2_filters := make([]*ec2.Filter, 0)

	re_space := regexp.MustCompile(`\s+`)
	re_name := regexp.MustCompile(`Name=`)
	re_values := regexp.MustCompile(`,Values=`)
	for _, i := range re_space.Split(filters, -1) {
		for _, j := range re_name.Split(i, -1) {
			if len(j) != 0 {
				v := re_values.Split(j, -1)
				name := v[0]
				/* // fix bug
				values := make([]*string, 0)
				for _, k := range v[1:] {
					fmt.Println(k)
					values = append(values, aws.String(k))
				}
				*/
				ec2_filters = append(ec2_filters, &ec2.Filter{
					Name: aws.String(name),
					Values: []*string{
						aws.String(v[1]),
					},
				})
			}
		}
	}
	return ec2_filters
}

func generateSession() (*session.Session, error) {
	return session.NewSessionWithOptions(session.Options{})
}

func DescribeInstances(filters string) (ec2_instances, error) {

	var instances ec2_instances
	sess, err := generateSession()
	if err != nil {
		return nil, err
	}
	svc := ec2.New(sess)

	params := &ec2.DescribeInstancesInput{}
	if len(filters) != 0 {
		params = &ec2.DescribeInstancesInput{
			Filters: ParseFilter(filters),
		}
	}
	resp, err := svc.DescribeInstances(params)
	if err != nil {
		return nil, err
	}
	if len(resp.Reservations) == 0 {
		return ec2_instances{}, nil
	}
	for _, res := range resp.Reservations {
		for _, instance := range res.Instances {
			instances = append(instances, instance)
		}
	}
	sort.Sort(instances)
	return instances, nil
}

func PrintInstances(instances ec2_instances) {
	for _, instance := range instances {
		fmt.Println(
			GetTagValue(instance, "Name"),
			*instance.PrivateIpAddress,
			*instance.InstanceId,
			*instance.InstanceType,
			*instance.Placement.AvailabilityZone,
			*instance.State.Name,
		)
	}
}

func GetTagValue(instance *ec2.Instance, tag_name string) string {
	for _, t := range instance.Tags {
		if *t.Key == tag_name {
			return *t.Value
		}
	}
	return ""
}
<|endoftext|>


<commit_before>// Package aws provides core functionality for making requests to AWS services.
package aws

// SDKName is the name of this AWS SDK
const SDKName = "aws-sdk-go"

// SDKVersion is the version of this SDK
const SDKVersion = "0.9.7"
<commit_msg>Tag release v0.9.8<commit_after>// Package aws provides core functionality for making requests to AWS services.
package aws

// SDKName is the name of this AWS SDK
const SDKName = "aws-sdk-go"

// SDKVersion is the version of this SDK
const SDKVersion = "0.9.8"
<|endoftext|>


<commit_before>package storage

import (
	"errors"
	"fmt"
	"log"
	"os"
	"path"
	"sync"
)

type Version uint8

const (
	SuperBlockSize = 8
	CurrentVersion = Version(1)
)

type Volume struct {
	Id       VolumeId
	dir      string
	dataFile *os.File
	nm       *NeedleMap

	replicaType ReplicationType
	version     Version

	accessLock sync.Mutex
}

func NewVolume(dirname string, id VolumeId, replicationType ReplicationType) (v *Volume) {
	v = &Volume{dir: dirname, Id: id, replicaType: replicationType}
	v.load()
	return
}
func (v *Volume) load() error {
	var e error
	fileName := path.Join(v.dir, v.Id.String())
	v.dataFile, e = os.OpenFile(fileName+".dat", os.O_RDWR|os.O_CREATE, 0644)
	if e != nil {
		return fmt.Errorf("cannot create Volume Data %s.dat: %s", fileName, e)
	}
	if v.replicaType == CopyNil {
		if e = v.readSuperBlock(); e != nil {
			return e
		}
	} else {
		v.maybeWriteSuperBlock()
	}
	indexFile, ie := os.OpenFile(fileName+".idx", os.O_RDWR|os.O_CREATE, 0644)
	if ie != nil {
		return fmt.Errorf("cannot create Volume Data %s.dat: %s", fileName, e)
	}
	v.nm = LoadNeedleMap(indexFile)
	return nil
}
func (v *Volume) Size() int64 {
    v.accessLock.Lock()
    defer v.accessLock.Unlock()
	stat, e := v.dataFile.Stat()
	if e == nil {
		return stat.Size()
	}
	fmt.Printf("Failed to read file size %s %s\n", v.dataFile.Name(), e.Error())
	return -1
}
func (v *Volume) Close() {
    v.accessLock.Lock()
    defer v.accessLock.Unlock()
	v.nm.Close()
	v.dataFile.Close()
}
func (v *Volume) maybeWriteSuperBlock() {
	stat, _ := v.dataFile.Stat()
	if stat.Size() == 0 {
		header := make([]byte, SuperBlockSize)
		header[0] = byte(CurrentVersion)
		header[1] = v.replicaType.Byte()
		v.dataFile.Write(header)
	}
}
func (v *Volume) readSuperBlock() error {
	v.dataFile.Seek(0, 0)
	header := make([]byte, SuperBlockSize)
	if _, e := v.dataFile.Read(header); e != nil {
		return fmt.Errorf("cannot read superblock: %s", e)
	}
	v.version = Version(header[0])
	var err error
	if v.replicaType, err = NewReplicationTypeFromByte(header[1]); err != nil {
		return fmt.Errorf("cannot read replica type: %s", err)
	}
	return nil
}
func (v *Volume) NeedToReplicate() bool {
	return v.replicaType.GetCopyCount() > 1
}

func (v *Volume) write(n *Needle) uint32 {
	v.accessLock.Lock()
	defer v.accessLock.Unlock()
	offset, _ := v.dataFile.Seek(0, 2)
	ret := n.Append(v.dataFile)
	nv, ok := v.nm.Get(n.Id)
	if !ok || int64(nv.Offset)*8 < offset {
		v.nm.Put(n.Id, uint32(offset/8), n.Size)
	}
	return ret
}
func (v *Volume) delete(n *Needle) uint32 {
	v.accessLock.Lock()
	defer v.accessLock.Unlock()
	nv, ok := v.nm.Get(n.Id)
	//log.Println("key", n.Id, "volume offset", nv.Offset, "data_size", n.Size, "cached size", nv.Size)
	if ok {
		v.nm.Delete(n.Id)
		v.dataFile.Seek(int64(nv.Offset*8), 0)
		n.Append(v.dataFile)
		return nv.Size
	}
	return 0
}

func (v *Volume) read(n *Needle) (int, error) {
	v.accessLock.Lock()
	defer v.accessLock.Unlock()
	nv, ok := v.nm.Get(n.Id)
	if ok && nv.Offset > 0 {
		v.dataFile.Seek(int64(nv.Offset)*8, 0)
		return n.Read(v.dataFile, nv.Size, v.version)
	}
	return -1, errors.New("Not Found")
}

func (v *Volume) garbageLevel() float64 {
    return float64(v.nm.deletionByteCounter)/float64(v.Size())
}

func (v *Volume) compact() error {
	v.accessLock.Lock()
	defer v.accessLock.Unlock()

	filePath := path.Join(v.dir, v.Id.String())
	return v.copyDataAndGenerateIndexFile(filePath+".dat", filePath+".cpd", filePath+".cpx")
}
func (v *Volume) commitCompact() (error) {
	v.accessLock.Lock()
	defer v.accessLock.Unlock()
	v.dataFile.Close()
	var e error
	if e = os.Rename(path.Join(v.dir, v.Id.String()+".cpd"), path.Join(v.dir, v.Id.String()+".dat")); e != nil {
		return e
	}
	if e = os.Rename(path.Join(v.dir, v.Id.String()+".cpx"), path.Join(v.dir, v.Id.String()+".idx")); e != nil {
		return e
	}
	if e = v.load(); e != nil {
		return e
	}
	return nil
}

func (v *Volume) copyDataAndGenerateIndexFile(srcName, dstName, idxName string) (err error) {
	src, err := os.OpenFile(srcName, os.O_RDONLY, 0644)
	if err != nil {
		return err
	}
	defer src.Close()

	dst, err := os.OpenFile(dstName, os.O_WRONLY|os.O_CREATE, 0644)
	if err != nil {
		return err
	}
	defer dst.Close()

	idx, err := os.OpenFile(idxName, os.O_WRONLY|os.O_CREATE, 0644)
	if err != nil {
		return err
	}
	defer idx.Close()

	src.Seek(0, 0)
	header := make([]byte, SuperBlockSize)
	if _, error := src.Read(header); error == nil {
		dst.Write(header)
	}

	n, rest := ReadNeedle(src)
	nm := NewNeedleMap(idx)
	old_offset := uint32(SuperBlockSize)
	new_offset := uint32(SuperBlockSize)
	for n != nil {
		nv, ok := v.nm.Get(n.Id)
		//log.Println("file size is", n.Size, "rest", rest)
		if !ok || nv.Offset*8 != old_offset {
			src.Seek(int64(rest), 1)
		} else {
			if nv.Size > 0 {
				nm.Put(n.Id, new_offset/8, n.Size)
				bytes := make([]byte, n.Size+4)
				src.Read(bytes)
				n.Data = bytes[:n.Size]
				n.Checksum = NewCRC(n.Data)
				n.Append(dst)
				new_offset += rest + 16
				//log.Println("saving key", n.Id, "volume offset", old_offset, "=>", new_offset, "data_size", n.Size, "rest", rest)
			}
			src.Seek(int64(rest-n.Size-4), 1)
		}
		old_offset += rest + 16
		n, rest = ReadNeedle(src)
	}

	return nil
}
<commit_msg>fix compilation error<commit_after>package storage

import (
	"errors"
	"fmt"
	"os"
	"path"
	"sync"
)

type Version uint8

const (
	SuperBlockSize = 8
	CurrentVersion = Version(1)
)

type Volume struct {
	Id       VolumeId
	dir      string
	dataFile *os.File
	nm       *NeedleMap

	replicaType ReplicationType
	version     Version

	accessLock sync.Mutex
}

func NewVolume(dirname string, id VolumeId, replicationType ReplicationType) (v *Volume) {
	v = &Volume{dir: dirname, Id: id, replicaType: replicationType}
	v.load()
	return
}
func (v *Volume) load() error {
	var e error
	fileName := path.Join(v.dir, v.Id.String())
	v.dataFile, e = os.OpenFile(fileName+".dat", os.O_RDWR|os.O_CREATE, 0644)
	if e != nil {
		return fmt.Errorf("cannot create Volume Data %s.dat: %s", fileName, e)
	}
	if v.replicaType == CopyNil {
		if e = v.readSuperBlock(); e != nil {
			return e
		}
	} else {
		v.maybeWriteSuperBlock()
	}
	indexFile, ie := os.OpenFile(fileName+".idx", os.O_RDWR|os.O_CREATE, 0644)
	if ie != nil {
		return fmt.Errorf("cannot create Volume Data %s.dat: %s", fileName, e)
	}
	v.nm = LoadNeedleMap(indexFile)
	return nil
}
func (v *Volume) Size() int64 {
    v.accessLock.Lock()
    defer v.accessLock.Unlock()
	stat, e := v.dataFile.Stat()
	if e == nil {
		return stat.Size()
	}
	fmt.Printf("Failed to read file size %s %s\n", v.dataFile.Name(), e.Error())
	return -1
}
func (v *Volume) Close() {
    v.accessLock.Lock()
    defer v.accessLock.Unlock()
	v.nm.Close()
	v.dataFile.Close()
}
func (v *Volume) maybeWriteSuperBlock() {
	stat, _ := v.dataFile.Stat()
	if stat.Size() == 0 {
		header := make([]byte, SuperBlockSize)
		header[0] = byte(CurrentVersion)
		header[1] = v.replicaType.Byte()
		v.dataFile.Write(header)
	}
}
func (v *Volume) readSuperBlock() error {
	v.dataFile.Seek(0, 0)
	header := make([]byte, SuperBlockSize)
	if _, e := v.dataFile.Read(header); e != nil {
		return fmt.Errorf("cannot read superblock: %s", e)
	}
	v.version = Version(header[0])
	var err error
	if v.replicaType, err = NewReplicationTypeFromByte(header[1]); err != nil {
		return fmt.Errorf("cannot read replica type: %s", err)
	}
	return nil
}
func (v *Volume) NeedToReplicate() bool {
	return v.replicaType.GetCopyCount() > 1
}

func (v *Volume) write(n *Needle) uint32 {
	v.accessLock.Lock()
	defer v.accessLock.Unlock()
	offset, _ := v.dataFile.Seek(0, 2)
	ret := n.Append(v.dataFile)
	nv, ok := v.nm.Get(n.Id)
	if !ok || int64(nv.Offset)*8 < offset {
		v.nm.Put(n.Id, uint32(offset/8), n.Size)
	}
	return ret
}
func (v *Volume) delete(n *Needle) uint32 {
	v.accessLock.Lock()
	defer v.accessLock.Unlock()
	nv, ok := v.nm.Get(n.Id)
	//log.Println("key", n.Id, "volume offset", nv.Offset, "data_size", n.Size, "cached size", nv.Size)
	if ok {
		v.nm.Delete(n.Id)
		v.dataFile.Seek(int64(nv.Offset*8), 0)
		n.Append(v.dataFile)
		return nv.Size
	}
	return 0
}

func (v *Volume) read(n *Needle) (int, error) {
	v.accessLock.Lock()
	defer v.accessLock.Unlock()
	nv, ok := v.nm.Get(n.Id)
	if ok && nv.Offset > 0 {
		v.dataFile.Seek(int64(nv.Offset)*8, 0)
		return n.Read(v.dataFile, nv.Size, v.version)
	}
	return -1, errors.New("Not Found")
}

func (v *Volume) garbageLevel() float64 {
    return float64(v.nm.deletionByteCounter)/float64(v.Size())
}

func (v *Volume) compact() error {
	v.accessLock.Lock()
	defer v.accessLock.Unlock()

	filePath := path.Join(v.dir, v.Id.String())
	return v.copyDataAndGenerateIndexFile(filePath+".dat", filePath+".cpd", filePath+".cpx")
}
func (v *Volume) commitCompact() (error) {
	v.accessLock.Lock()
	defer v.accessLock.Unlock()
	v.dataFile.Close()
	var e error
	if e = os.Rename(path.Join(v.dir, v.Id.String()+".cpd"), path.Join(v.dir, v.Id.String()+".dat")); e != nil {
		return e
	}
	if e = os.Rename(path.Join(v.dir, v.Id.String()+".cpx"), path.Join(v.dir, v.Id.String()+".idx")); e != nil {
		return e
	}
	if e = v.load(); e != nil {
		return e
	}
	return nil
}

func (v *Volume) copyDataAndGenerateIndexFile(srcName, dstName, idxName string) (err error) {
	src, err := os.OpenFile(srcName, os.O_RDONLY, 0644)
	if err != nil {
		return err
	}
	defer src.Close()

	dst, err := os.OpenFile(dstName, os.O_WRONLY|os.O_CREATE, 0644)
	if err != nil {
		return err
	}
	defer dst.Close()

	idx, err := os.OpenFile(idxName, os.O_WRONLY|os.O_CREATE, 0644)
	if err != nil {
		return err
	}
	defer idx.Close()

	src.Seek(0, 0)
	header := make([]byte, SuperBlockSize)
	if _, error := src.Read(header); error == nil {
		dst.Write(header)
	}

	n, rest := ReadNeedle(src)
	nm := NewNeedleMap(idx)
	old_offset := uint32(SuperBlockSize)
	new_offset := uint32(SuperBlockSize)
	for n != nil {
		nv, ok := v.nm.Get(n.Id)
		//log.Println("file size is", n.Size, "rest", rest)
		if !ok || nv.Offset*8 != old_offset {
			src.Seek(int64(rest), 1)
		} else {
			if nv.Size > 0 {
				nm.Put(n.Id, new_offset/8, n.Size)
				bytes := make([]byte, n.Size+4)
				src.Read(bytes)
				n.Data = bytes[:n.Size]
				n.Checksum = NewCRC(n.Data)
				n.Append(dst)
				new_offset += rest + 16
				//log.Println("saving key", n.Id, "volume offset", old_offset, "=>", new_offset, "data_size", n.Size, "rest", rest)
			}
			src.Seek(int64(rest-n.Size-4), 1)
		}
		old_offset += rest + 16
		n, rest = ReadNeedle(src)
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 backup

import (
	"fmt"
	"github.com/jacobsa/comeback/blob"
	"github.com/jacobsa/comeback/fs"
	"os"
	"path"
)

// An object that knows how to save directories to some underlying storage.
type DirectorySaver interface {
	// Recursively save the contents of the supplied directory to the underlying
	// storage, returning the score of a blob representing the directory's
	// listing in a format hat can be recovered with repr.Unmarshal.
	Save(dirpath string) (score blob.Score, err error)
}

type directorySaver struct {
	blobStore blob.Store
	fileSystem fs.FileSystem
	fileSaver FileSaver
	wrapped DirectorySaver
}

func (s *directorySaver) saveDir(parent string, fi os.FileInfo) ([]blob.Score, error) {
	// Recurse.
	score, err := s.wrapped.Save(path.Join(parent, fi.Name()))
}

func (s *directorySaver) saveFile(parent string, fi os.FileInfo) ([]blob.Score, error) {
}

func (s *directorySaver) Save(dirpath string) (score blob.Score, err error) {
	// Grab a listing for the directory.
	fileInfos, err := s.fileSystem.ReadDir(dirpath)
	if err != nil {
		return nil, fmt.Errorf("Listing directory: %v", err)
	}

	// Process each entry in the directory, building a list of DirectoryEntry
	// structs.
	entries := []fs.DirectoryEntry{}
	for _, fileInfo := range fileInfos {
		entry, err := convertCommon(fileInfo)
		if err != nil {
			return nil, err
		}

		// Call the appropriate method based on this entry's type.
		switch {
		case fileInfo.IsDir():
			entry.Scores, err = s.saveDir(dirpath, fileInfo)
		case fileInfo.Mode() & os.ModeType == 0:
			entry.Scores, err = s.saveFile(dirpath, fileInfo)
	  default:
			err = fmt.Errorf("Unhandled mode: %v", fileInfo.Mode())
		}

		if err != nil {
			return nil, err
		}

		entries = append(entries, entry)
	}
}
<commit_msg>Implemented saveDir.<commit_after>// Copyright 2012 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 backup

import (
	"fmt"
	"github.com/jacobsa/comeback/blob"
	"github.com/jacobsa/comeback/fs"
	"os"
	"path"
)

// An object that knows how to save directories to some underlying storage.
type DirectorySaver interface {
	// Recursively save the contents of the supplied directory to the underlying
	// storage, returning the score of a blob representing the directory's
	// listing in a format hat can be recovered with repr.Unmarshal.
	Save(dirpath string) (score blob.Score, err error)
}

type directorySaver struct {
	blobStore blob.Store
	fileSystem fs.FileSystem
	fileSaver FileSaver
	wrapped DirectorySaver
}

func convertCommon(fi os.FileInfo) (fs.DirectoryEntry, error) {
}

func (s *directorySaver) saveDir(parent string, fi os.FileInfo) ([]blob.Score, error) {
	// Recurse.
	score, err := s.wrapped.Save(path.Join(parent, fi.Name()))
	if err != nil {
		return nil, err
	}

	return []blob.Score{score}, nil
}

func (s *directorySaver) saveFile(parent string, fi os.FileInfo) ([]blob.Score, error) {
}

func (s *directorySaver) Save(dirpath string) (score blob.Score, err error) {
	// Grab a listing for the directory.
	fileInfos, err := s.fileSystem.ReadDir(dirpath)
	if err != nil {
		return nil, fmt.Errorf("Listing directory: %v", err)
	}

	// Process each entry in the directory, building a list of DirectoryEntry
	// structs.
	entries := []fs.DirectoryEntry{}
	for _, fileInfo := range fileInfos {
		entry, err := convertCommon(fileInfo)
		if err != nil {
			return nil, err
		}

		// Call the appropriate method based on this entry's type.
		switch {
		case fileInfo.IsDir():
			entry.Scores, err = s.saveDir(dirpath, fileInfo)
		case fileInfo.Mode() & os.ModeType == 0:
			entry.Scores, err = s.saveFile(dirpath, fileInfo)
	  default:
			err = fmt.Errorf("Unhandled mode: %v", fileInfo.Mode())
		}

		if err != nil {
			return nil, err
		}

		entries = append(entries, entry)
	}
}
<|endoftext|>


<commit_before>package admin

import (
	"github.com/BytemarkHosting/bytemark-client/cmd/bytemark/app"
	"github.com/BytemarkHosting/bytemark-client/cmd/bytemark/app/args"
	"github.com/BytemarkHosting/bytemark-client/cmd/bytemark/app/with"
	"github.com/BytemarkHosting/bytemark-client/util/log"
	"github.com/urfave/cli"
)

func init() {
	Commands = append(Commands, cli.Command{
		Name:   "create",
		Action: cli.ShowSubcommandHelp,
		Subcommands: []cli.Command{
			{
				Name:      "ip range",
				Usage:     "create a new IP range in a VLAN",
				UsageText: "bytemark --admin create ip range <ip-range> <vlan-num>",
				Flags: []cli.Flag{
					cli.StringFlag{
						Name:  "ip-range",
						Usage: "the IP range to add",
					},
					cli.IntFlag{
						Name:  "vlan-num",
						Usage: "The VLAN number to add the IP range to",
					},
				},
				Action: app.Action(args.Optional("ip-range", "vlan-num"), with.RequiredFlags("ip-range", "vlan-num"), with.Auth, func(c *app.Context) error {
					if err := c.Client().CreateIPRange(c.String("ip-range"), c.Int("vlan-num")); err != nil {
						return err
					}
					log.Logf("IP range created\r\n")
					return nil
				}),
			}, {
				Name:      "migration",
				Usage:     "creates a new migration job",
				UsageText: "bytemark --admin create migration [--priority <number>] [--disc <disc>]... [--pool <pool>]... [--tail <tail>]... [--to-pool <pool>]...",
				Description: `Requests that the brain starts a migration from the various discs, pools and tails specified to the various pools specified.

   If no --to-pools are specified, the brain will decide destinations on its own. To migrate multiple discs, pools, or tails, or to distribute migrating discs across multiple storage pools, simply specify those flags multiple times.

   Source discs, tails and pools and destination pools can all be specified as their numeric IDs, their label, or their IP address.
When there are multiple migration jobs at once and migration concurrency limits are being reached, jobs with highest priority will have their migrations scheduled ahead of other jobs.

EXAMPLES:
   # migrate one disc whose IP is fe80::1 to tail2-sata1
   bytemark --admin create migration --disc fe80::1 --to-pool tail2-sata1

   # empty a tail to wherever is appropriate with a high priority
   bytemark --admin create migration --priority 10 --tail tail1

   # migrate two discs, one called very-unique-label and one whose ID is 45859 to tail2-sata1
   bytemark --admin create migration --disc very-unique-label --disc 45859 --to-pool tail2-sata1

   # migrate one pool to another
   bytemark --admin create migration --pool tail1-sata1 --to-pool tail2-sata1

   # migrate two pools to one pool
   bytemark --admin create migration --pool tail1-sata1 --pool tail1-sata2 --to-pool tail2-sata1

   # migrate a pool and a specific disc away from their current tail
   bytemark --admin create migration --pool tail1-sata1 --disc 45859`,

				Flags: []cli.Flag{
					cli.StringSliceFlag{
						Name:  "disc",
						Usage: "a disc to migrate to some other destination. Can be specified multiple times to migrate multiple discs",
					},
					cli.StringSliceFlag{
						Name:  "pool",
						Usage: "a storage pool to migrate to some other destination. Can be specified multiple times to migrate multiple storage pools",
					},
					cli.StringSliceFlag{
						Name:  "tail",
						Usage: "a tail to migrate to some other destination. Can be specified multiple times to migrate multiple tails",
					},
					cli.StringSliceFlag{
						Name:  "to-pool",
						Usage: "a disc to migrate to some other destination. Can be specified multiple times to migrate multiple discs",
					},
					cli.IntFlag{
						Name:  "priority",
						Usage: "an optional priority to set - bigger is higher priority",
					},
				},
			}, {
				Name:        "user",
				Usage:       "creates a new cluster admin or cluster superuser",
				UsageText:   "bytemark --admin create user <username> <privilege>",
				Description: `creates a new cluster admin or superuser. The privilege field must be either cluster_admin or cluster_su.`,
				Flags: []cli.Flag{
					cli.StringFlag{
						Name:  "username",
						Usage: "The username of the new user",
					},
					cli.StringFlag{
						Name:  "privilege",
						Usage: "The privilege to grant to the new user",
					},
				},
				Action: app.Action(args.Optional("username", "privilege"), with.RequiredFlags("username", "privilege"), with.Auth, func(c *app.Context) error {
					// Privilege is just a string and not a app.PrivilegeFlag, since it can only be "cluster_admin" or "cluster_su"
					if err := c.Client().CreateUser(c.String("username"), c.String("privilege")); err != nil {
						return err
					}
					log.Logf("User %s has been created with %s privileges\r\n", c.String("username"), c.String("privilege"))
					return nil
				}),
			}, {
				Name:      "vlan group",
				Aliases:   []string{"vlan-group"},
				Usage:     "creates groups for private VLANs",
				UsageText: "bytemark --admin create vlan group <group> [vlan-num]",
				Description: `Create a group in the specified account, with an optional VLAN specified.

Used when setting up a private VLAN for a customer.`,
				Flags: []cli.Flag{
					cli.GenericFlag{
						Name:  "group",
						Usage: "the name of the group to create",
						Value: new(app.GroupNameFlag),
					},
					cli.IntFlag{
						Name:  "vlan-num",
						Usage: "The VLAN number to add the group to",
					},
				},
				Action: app.Action(args.Optional("group", "vlan-num"), with.RequiredFlags("group"), with.Auth, func(c *app.Context) error {
					gp := c.GroupName("group")
					if err := c.Client().AdminCreateGroup(gp, c.Int("vlan-num")); err != nil {
						return err
					}
					log.Logf("Group %s was created under account %s\r\n", gp.Group, gp.Account)
					return nil
				}),
			},
		},
	})
}
<commit_msg>Sketch implementation of create migration<commit_after>package admin

import (
	"github.com/BytemarkHosting/bytemark-client/cmd/bytemark/app"
	"github.com/BytemarkHosting/bytemark-client/cmd/bytemark/app/args"
	"github.com/BytemarkHosting/bytemark-client/cmd/bytemark/app/with"
	"github.com/BytemarkHosting/bytemark-client/lib/brain"
	"github.com/BytemarkHosting/bytemark-client/util/log"
	"github.com/urfave/cli"
)

func init() {
	Commands = append(Commands, cli.Command{
		Name:   "create",
		Action: cli.ShowSubcommandHelp,
		Subcommands: []cli.Command{
			{
				Name:      "ip range",
				Usage:     "create a new IP range in a VLAN",
				UsageText: "bytemark --admin create ip range <ip-range> <vlan-num>",
				Flags: []cli.Flag{
					cli.StringFlag{
						Name:  "ip-range",
						Usage: "the IP range to add",
					},
					cli.IntFlag{
						Name:  "vlan-num",
						Usage: "The VLAN number to add the IP range to",
					},
				},
				Action: app.Action(args.Optional("ip-range", "vlan-num"), with.RequiredFlags("ip-range", "vlan-num"), with.Auth, func(c *app.Context) error {
					if err := c.Client().CreateIPRange(c.String("ip-range"), c.Int("vlan-num")); err != nil {
						return err
					}
					log.Logf("IP range created\r\n")
					return nil
				}),
			}, {
				Name:      "migration",
				Usage:     "creates a new migration job",
				UsageText: "bytemark --admin create migration [--priority <number>] [--disc <disc>]... [--pool <pool>]... [--tail <tail>]... [--to-pool <pool>]...",
				Description: `Requests that the brain starts a migration from the various discs, pools and tails specified to the various pools specified.

   If no --to-pools are specified, the brain will decide destinations on its own. To migrate multiple discs, pools, or tails, or to distribute migrating discs across multiple storage pools, simply specify those flags multiple times.

   Source discs, tails and pools and destination pools can all be specified as their numeric IDs, their label, or their IP address.
When there are multiple migration jobs at once and migration concurrency limits are being reached, jobs with highest priority will have their migrations scheduled ahead of other jobs.

EXAMPLES:
   # migrate one disc whose IP is fe80::1 to tail2-sata1
   bytemark --admin create migration --disc fe80::1 --to-pool tail2-sata1

   # empty a tail to wherever is appropriate with a high priority
   bytemark --admin create migration --priority 10 --tail tail1

   # migrate two discs, one called very-unique-label and one whose ID is 45859 to tail2-sata1
   bytemark --admin create migration --disc very-unique-label --disc 45859 --to-pool tail2-sata1

   # migrate one pool to another
   bytemark --admin create migration --pool tail1-sata1 --to-pool tail2-sata1

   # migrate two pools to one pool
   bytemark --admin create migration --pool tail1-sata1 --pool tail1-sata2 --to-pool tail2-sata1

   # migrate a pool and a specific disc away from their current tail
   bytemark --admin create migration --pool tail1-sata1 --disc 45859`,

				Flags: []cli.Flag{
					cli.StringSliceFlag{
						Name:  "disc",
						Usage: "a disc to migrate to some other destination. Can be specified multiple times to migrate multiple discs",
					},
					cli.StringSliceFlag{
						Name:  "pool",
						Usage: "a storage pool to migrate to some other destination. Can be specified multiple times to migrate multiple storage pools",
					},
					cli.StringSliceFlag{
						Name:  "tail",
						Usage: "a tail to migrate to some other destination. Can be specified multiple times to migrate multiple tails",
					},
					cli.StringSliceFlag{
						Name:  "to-pool",
						Usage: "a disc to migrate to some other destination. Can be specified multiple times to migrate multiple discs",
					},
					cli.IntFlag{
						Name:  "priority",
						Usage: "an optional priority to set - bigger is higher priority",
					},
				},
				Action: app.Action(with.Auth, func(ctx *app.Context) err {
					jobRequest := brain.MigrationJob{
						Sources: brain.MigrationJobLocations{
							Discs: ctx.StringSlice("disc"),
							Pools: ctx.StringSlice("pool"),
							Tails: ctx.StringSlice("tail"),
						},
						Destinations: brain.MigrationJobDestinations{
							Pools: ctx.StringSlice("to-pool"),
						},
						Priority: ctx.Int("priority"),
					}
					job, err := adminRequests.CreateMigrationJob(jobRequest)
					if err != nil {
						ctx.Output("Migration job created")
						return ctx.OutputInDesiredForm(job)
					}
					ctx.Output("Couldn't create migration job")
					return err
				}),
			}, {
				Name:        "user",
				Usage:       "creates a new cluster admin or cluster superuser",
				UsageText:   "bytemark --admin create user <username> <privilege>",
				Description: `creates a new cluster admin or superuser. The privilege field must be either cluster_admin or cluster_su.`,
				Flags: []cli.Flag{
					cli.StringFlag{
						Name:  "username",
						Usage: "The username of the new user",
					},
					cli.StringFlag{
						Name:  "privilege",
						Usage: "The privilege to grant to the new user",
					},
				},
				Action: app.Action(args.Optional("username", "privilege"), with.RequiredFlags("username", "privilege"), with.Auth, func(c *app.Context) error {
					// Privilege is just a string and not a app.PrivilegeFlag, since it can only be "cluster_admin" or "cluster_su"
					if err := c.Client().CreateUser(c.String("username"), c.String("privilege")); err != nil {
						return err
					}
					log.Logf("User %s has been created with %s privileges\r\n", c.String("username"), c.String("privilege"))
					return nil
				}),
			}, {
				Name:      "vlan group",
				Aliases:   []string{"vlan-group"},
				Usage:     "creates groups for private VLANs",
				UsageText: "bytemark --admin create vlan group <group> [vlan-num]",
				Description: `Create a group in the specified account, with an optional VLAN specified.

Used when setting up a private VLAN for a customer.`,
				Flags: []cli.Flag{
					cli.GenericFlag{
						Name:  "group",
						Usage: "the name of the group to create",
						Value: new(app.GroupNameFlag),
					},
					cli.IntFlag{
						Name:  "vlan-num",
						Usage: "The VLAN number to add the group to",
					},
				},
				Action: app.Action(args.Optional("group", "vlan-num"), with.RequiredFlags("group"), with.Auth, func(c *app.Context) error {
					gp := c.GroupName("group")
					if err := c.Client().AdminCreateGroup(gp, c.Int("vlan-num")); err != nil {
						return err
					}
					log.Logf("Group %s was created under account %s\r\n", gp.Group, gp.Account)
					return nil
				}),
			},
		},
	})
}
<|endoftext|>


<commit_before>package image

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"crypto/sha256"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"testing"
)

func TestUnpackLayerDuplicateEntries(t *testing.T) {
	tmp1, err := ioutil.TempDir("", "test-dup")
	if err != nil {
		t.Fatal(err)
	}
	tarfile := filepath.Join(tmp1, "test.tar")
	f, err := os.Create(tarfile)
	if err != nil {
		t.Fatal(err)
	}
	defer f.Close()
	defer os.RemoveAll(tarfile)
	gw := gzip.NewWriter(f)
	tw := tar.NewWriter(gw)

	tw.WriteHeader(&tar.Header{Name: "test", Size: 4, Mode: 0600})
	io.Copy(tw, bytes.NewReader([]byte("test")))
	tw.WriteHeader(&tar.Header{Name: "test", Size: 5, Mode: 0600})
	io.Copy(tw, bytes.NewReader([]byte("test1")))
	tw.Close()
	gw.Close()

	r, err := os.Open(tarfile)
	if err != nil {
		t.Fatal(err)
	}
	defer r.Close()
	tmp2, err := ioutil.TempDir("", "test-dest-unpack")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(tmp2)
	if err := unpackLayer(tmp2, r); err != nil && !strings.Contains(err.Error(), "duplicate entry for") {
		t.Fatalf("Expected to fail with duplicate entry, got %v", err)
	}
}

func TestUnpackLayer(t *testing.T) {
	tmp1, err := ioutil.TempDir("", "test-layer")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(tmp1)
	err = os.MkdirAll(filepath.Join(tmp1, "blobs", "sha256"), 0700)
	if err != nil {
		t.Fatal(err)
	}
	tarfile := filepath.Join(tmp1, "blobs", "sha256", "test.tar")
	f, err := os.Create(tarfile)
	if err != nil {
		t.Fatal(err)
	}

	gw := gzip.NewWriter(f)
	tw := tar.NewWriter(gw)

	tw.WriteHeader(&tar.Header{Name: "test", Size: 4, Mode: 0600})
	io.Copy(tw, bytes.NewReader([]byte("test")))
	tw.Close()
	gw.Close()
	f.Close()

	// generate sha256 hash
	h := sha256.New()
	file, err := os.Open(tarfile)
	if err != nil {
		t.Fatal(err)
	}
	defer file.Close()
	_, err = io.Copy(h, file)
	if err != nil {
		t.Fatal(err)
	}
	err = os.Rename(tarfile, filepath.Join(tmp1, "blobs", "sha256", fmt.Sprintf("%x", h.Sum(nil))))
	if err != nil {
		t.Fatal(err)
	}

	testManifest := manifest{
		Layers: []descriptor{descriptor{
			MediaType: "application/vnd.oci.image.layer.tar+gzip",
			Digest:    fmt.Sprintf("sha256:%s", fmt.Sprintf("%x", h.Sum(nil))),
		}},
	}
	err = testManifest.unpack(newPathWalker(tmp1), filepath.Join(tmp1, "rootfs"))
	if err != nil {
		t.Fatal(err)
	}

	_, err = os.Stat(filepath.Join(tmp1, "rootfs", "test"))
	if err != nil {
		t.Fatal(err)
	}
}
<commit_msg>Add license head to manifest_test.go<commit_after>// Copyright 2016 The Linux Foundation
//
// 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 image

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"crypto/sha256"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"testing"
)

func TestUnpackLayerDuplicateEntries(t *testing.T) {
	tmp1, err := ioutil.TempDir("", "test-dup")
	if err != nil {
		t.Fatal(err)
	}
	tarfile := filepath.Join(tmp1, "test.tar")
	f, err := os.Create(tarfile)
	if err != nil {
		t.Fatal(err)
	}
	defer f.Close()
	defer os.RemoveAll(tarfile)
	gw := gzip.NewWriter(f)
	tw := tar.NewWriter(gw)

	tw.WriteHeader(&tar.Header{Name: "test", Size: 4, Mode: 0600})
	io.Copy(tw, bytes.NewReader([]byte("test")))
	tw.WriteHeader(&tar.Header{Name: "test", Size: 5, Mode: 0600})
	io.Copy(tw, bytes.NewReader([]byte("test1")))
	tw.Close()
	gw.Close()

	r, err := os.Open(tarfile)
	if err != nil {
		t.Fatal(err)
	}
	defer r.Close()
	tmp2, err := ioutil.TempDir("", "test-dest-unpack")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(tmp2)
	if err := unpackLayer(tmp2, r); err != nil && !strings.Contains(err.Error(), "duplicate entry for") {
		t.Fatalf("Expected to fail with duplicate entry, got %v", err)
	}
}

func TestUnpackLayer(t *testing.T) {
	tmp1, err := ioutil.TempDir("", "test-layer")
	if err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(tmp1)
	err = os.MkdirAll(filepath.Join(tmp1, "blobs", "sha256"), 0700)
	if err != nil {
		t.Fatal(err)
	}
	tarfile := filepath.Join(tmp1, "blobs", "sha256", "test.tar")
	f, err := os.Create(tarfile)
	if err != nil {
		t.Fatal(err)
	}

	gw := gzip.NewWriter(f)
	tw := tar.NewWriter(gw)

	tw.WriteHeader(&tar.Header{Name: "test", Size: 4, Mode: 0600})
	io.Copy(tw, bytes.NewReader([]byte("test")))
	tw.Close()
	gw.Close()
	f.Close()

	// generate sha256 hash
	h := sha256.New()
	file, err := os.Open(tarfile)
	if err != nil {
		t.Fatal(err)
	}
	defer file.Close()
	_, err = io.Copy(h, file)
	if err != nil {
		t.Fatal(err)
	}
	err = os.Rename(tarfile, filepath.Join(tmp1, "blobs", "sha256", fmt.Sprintf("%x", h.Sum(nil))))
	if err != nil {
		t.Fatal(err)
	}

	testManifest := manifest{
		Layers: []descriptor{descriptor{
			MediaType: "application/vnd.oci.image.layer.tar+gzip",
			Digest:    fmt.Sprintf("sha256:%s", fmt.Sprintf("%x", h.Sum(nil))),
		}},
	}
	err = testManifest.unpack(newPathWalker(tmp1), filepath.Join(tmp1, "rootfs"))
	if err != nil {
		t.Fatal(err)
	}

	_, err = os.Stat(filepath.Join(tmp1, "rootfs", "test"))
	if err != nil {
		t.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package weather

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"time"
)

// OpenWeatherService is a Open Weather Map Service API
type OpenWeatherService struct {
	appID string
}

// OWResponse is the outermost wrapper of the Open Weather API response
type OWResponse struct {
	ObservationEpoch int64         `json:"dt"`
	Conditions       *OWConditions `json:"main"`
	Clouds           *OWCloudiness `json:"clouds"`
}

//OWCloudiness is the cloudiness
type OWCloudiness struct {
	Percentage float64 `json:"all"`
}

// OWConditions is the current weather status.
type OWConditions struct {
	TemperatureC float64 `json:"temp"`
}

// Read sends a request to the API and converts the data expects "95125,us" as an input
func (w *OpenWeatherService) Read(zipCommaCountry string) (*Data, error) {
	if w.appID == "" {
		return nil, fmt.Errorf("Cannot make request to Weather Underground without an AppID")
	}
	// Return cached value
	url := fmt.Sprintf("https://api.openweathermap.org/data/2.5/weather?zip=%s&appid=%s&units=metric", zipCommaCountry, w.appID)
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	buf := bytes.NewBuffer(nil)
	n, cerr := io.Copy(buf, resp.Body)
	fmt.Printf("Copied %d bytes (%v): %s", n, cerr, buf.String())
	response := &OWResponse{}
	err = json.Unmarshal(buf.Bytes(), &response)
	if err != nil {
		return nil, err
	}
	fmt.Printf("Result: %+v", response)
	return w.convert(zipCommaCountry, response)
}

func (w *OpenWeatherService) convert(zipcode string, co *OWResponse) (*Data, error) {
	solarradiation := 500.0 * (100.0 - co.Clouds.Percentage) / 100.0
	return &Data{
		Zipcode:        zipcode,
		Updated:        time.Unix(co.ObservationEpoch, 0),
		CurrentTempC:   co.Conditions.TemperatureC,
		SolarRadiation: solarradiation,
		Description:    "OpenWeatherMap current observation",
	}, nil
}
<commit_msg>fix invalid run issue<commit_after>package weather

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"time"
)

// OpenWeatherService is a Open Weather Map Service API
type OpenWeatherService struct {
	appID string
}

// OWResponse is the outermost wrapper of the Open Weather API response
type OWResponse struct {
	ObservationEpoch int64         `json:"dt"`
	Conditions       *OWConditions `json:"main"`
	Clouds           *OWCloudiness `json:"clouds"`
}

//OWCloudiness is the cloudiness
type OWCloudiness struct {
	Percentage float64 `json:"all"`
}

// OWConditions is the current weather status.
type OWConditions struct {
	TemperatureC float64 `json:"temp"`
}

// SolarRadiation returns a bad approximation to solar radiation based on the cloud cover (does not know if the sun is up)
func (r *OWResponse) SolarRadiation() float64 {
	if r.Clouds != nil {
		return 500.0 * (100.0 - r.Clouds.Percentage) / 100.0
	}
	return 0.0
}

// TempC returns the temperature if it is available
func (r *OWResponse) TempC() float64 {
	if r.Conditions != nil {
		return r.Conditions.TemperatureC
	}
	return 0.0
}

// Read sends a request to the API and converts the data expects "95125,us" as an input
func (w *OpenWeatherService) Read(zipCommaCountry string) (*Data, error) {
	if w.appID == "" {
		return nil, fmt.Errorf("Cannot make request to Weather Underground without an AppID")
	}
	// Return cached value
	url := fmt.Sprintf("https://api.openweathermap.org/data/2.5/weather?zip=%s&appid=%s&units=metric", zipCommaCountry, w.appID)
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	buf := bytes.NewBuffer(nil)
	n, cerr := io.Copy(buf, resp.Body)
	fmt.Printf("Copied %d bytes (%v): %s", n, cerr, buf.String())
	response := &OWResponse{}
	err = json.Unmarshal(buf.Bytes(), &response)
	if err != nil {
		return nil, err
	}
	fmt.Printf("Result: %+v", response)
	return w.convert(zipCommaCountry, response)
}

func (w *OpenWeatherService) convert(zipcode string, co *OWResponse) (*Data, error) {
	if co == nil {
		return &Data{}, errors.New("invalid response")
	}

	return &Data{
		Zipcode:        zipcode,
		Updated:        time.Unix(co.ObservationEpoch, 0),
		CurrentTempC:   co.TempC(),
		SolarRadiation: co.SolarRadiation(),
		Description:    "OpenWeatherMap current observation",
	}, nil
}
<|endoftext|>


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

package bench

import (
	"fmt"
	"log"
	"os"
	"time"

	"github.com/coreos/etcd/Godeps/_workspace/src/github.com/cheggaaa/pb"
	"github.com/coreos/etcd/Godeps/_workspace/src/github.com/spf13/cobra"
	"github.com/coreos/etcd/Godeps/_workspace/src/golang.org/x/net/context"
	v3 "github.com/coreos/etcd/clientv3"
	"github.com/samuel/go-zookeeper/zk"
)

// rangeCmd represents the range command
var rangeCmd = &cobra.Command{
	Use:   "range key [end-range]",
	Short: "Benchmark range",

	Run: rangeFunc,
}

var (
	rangeTotal       int
	rangeConsistency string
	singleKey        bool
)

func init() {
	Command.AddCommand(rangeCmd)
	rangeCmd.Flags().IntVar(&rangeTotal, "total", 10000, "Total number of range requests")
	rangeCmd.Flags().StringVar(&rangeConsistency, "consistency", "l", "Linearizable(l) or Serializable(s)")
	rangeCmd.Flags().BoolVar(&singleKey, "single-key", false, "'true' to get only one single key (automatic put before test)")
}

func rangeFunc(cmd *cobra.Command, args []string) {
	if singleKey { // write 'foo'
		switch database {
		case "etcd":
			fmt.Println("PUT 'foo' to etcd")
			clients := mustCreateClients(1, 1)
			_, err := clients[0].Do(context.Background(), v3.OpPut("foo", "bar"))
			if err != nil {
				fmt.Fprintln(os.Stderr, err)
				os.Exit(1)
			}
			fmt.Println("Done with PUT 'foo' to etcd")
		case "zk":
			fmt.Println("PUT 'foo' to zookeeper")
			conn := mustCreateConnsZk(1)
			_, err := conn[0].Create("foo", []byte("bar"), zkCreateFlags, zkCreateAcl)
			if err != nil {
				fmt.Fprintln(os.Stderr, err)
				os.Exit(1)
			}
			fmt.Println("Done with PUT 'foo' to zookeeper")
		}
	} else if len(args) == 0 || len(args) > 2 {
		fmt.Fprintln(os.Stderr, cmd.Usage())
		os.Exit(1)
	}

	var k, end string
	if singleKey {
		k = "foo"
	} else {
		k = args[0]
		if len(args) == 2 {
			end = args[1]
		}
	}

	if rangeConsistency == "l" {
		fmt.Println("bench with linearizable range")
	} else if rangeConsistency == "s" {
		fmt.Println("bench with serializable range")
	} else {
		fmt.Fprintln(os.Stderr, cmd.Usage())
		os.Exit(1)
	}

	results = make(chan result)
	requests := make(chan request, totalClients)
	bar = pb.New(rangeTotal)

	bar.Format("Bom !")
	bar.Start()

	switch database {
	case "etcd":
		clients := mustCreateClients(totalClients, totalConns)
		for i := range clients {
			wg.Add(1)
			go doRange(clients[i].KV, requests)
		}
		defer func() {
			for i := range clients {
				clients[i].Close()
			}
		}()
	case "zk":
		conns := mustCreateConnsZk(totalConns)
		defer func() {
			for i := range conns {
				conns[i].Close()
			}
		}()
		for i := range conns {
			wg.Add(1)
			go doRangeZk(conns[i], requests)
		}
	default:
		log.Fatalf("unknown database %s", database)
	}

	pdoneC := printReport(results)
	go func() {
		for i := 0; i < rangeTotal; i++ {
			switch database {
			case "etcd":
				opts := []v3.OpOption{v3.WithRange(end)}
				if rangeConsistency == "s" {
					opts = append(opts, v3.WithSerializable())
				}
				requests <- request{etcdOp: v3.OpGet(k, opts...)}
			case "zk":
				requests <- request{zkOp: zkOp{key: k}}
			}
		}
		close(requests)
	}()

	wg.Wait()

	bar.Finish()

	close(results)
	<-pdoneC
}

func doRange(client v3.KV, requests <-chan request) {
	defer wg.Done()

	for req := range requests {
		op := req.etcdOp
		st := time.Now()
		_, err := client.Do(context.Background(), op)

		var errStr string
		if err != nil {
			errStr = err.Error()
		}
		results <- result{errStr: errStr, duration: time.Since(st), happened: time.Now()}
		bar.Increment()
	}
}

func doRangeZk(conn *zk.Conn, requests <-chan request) {
	defer wg.Done()

	for req := range requests {
		op := req.zkOp
		st := time.Now()
		_, _, err := conn.Get(op.key)

		var errStr string
		if err != nil {
			errStr = err.Error()
		}
		results <- result{errStr: errStr, duration: time.Since(st), happened: time.Now()}
		bar.Increment()
	}
}
<commit_msg>bench/range: add value size for range<commit_after>// Copyright 2015 CoreOS, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package bench

import (
	"fmt"
	"log"
	"os"
	"time"

	"github.com/coreos/etcd/Godeps/_workspace/src/github.com/cheggaaa/pb"
	"github.com/coreos/etcd/Godeps/_workspace/src/github.com/spf13/cobra"
	"github.com/coreos/etcd/Godeps/_workspace/src/golang.org/x/net/context"
	v3 "github.com/coreos/etcd/clientv3"
	"github.com/samuel/go-zookeeper/zk"
)

// rangeCmd represents the range command
var rangeCmd = &cobra.Command{
	Use:   "range key [end-range]",
	Short: "Benchmark range",

	Run: rangeFunc,
}

var (
	rangeTotal       int
	rangeConsistency string
	singleKey        bool
)

func init() {
	Command.AddCommand(rangeCmd)
	rangeCmd.Flags().IntVar(&rangeTotal, "total", 10000, "Total number of range requests")
	rangeCmd.Flags().StringVar(&rangeConsistency, "consistency", "l", "Linearizable(l) or Serializable(s)")
	rangeCmd.Flags().BoolVar(&singleKey, "single-key", false, "'true' to get only one single key (automatic put before test)")
	rangeCmd.Flags().IntVar(&valSize, "val-size", 128, "value size")
}

func rangeFunc(cmd *cobra.Command, args []string) {
	if singleKey { // write 'foo'
		v := mustRandBytes(valSize)
		switch database {
		case "etcd":
			fmt.Println("PUT 'foo' to etcd")
			clients := mustCreateClients(1, 1)
			_, err := clients[0].Do(context.Background(), v3.OpPut("foo", string(v)))
			if err != nil {
				fmt.Fprintln(os.Stderr, err)
				os.Exit(1)
			}
			fmt.Println("Done with PUT 'foo' to etcd")
		case "zk":
			fmt.Println("PUT 'foo' to zookeeper")
			conn := mustCreateConnsZk(1)
			_, err := conn[0].Create("foo", v, zkCreateFlags, zkCreateAcl)
			if err != nil {
				fmt.Fprintln(os.Stderr, err)
				os.Exit(1)
			}
			fmt.Println("Done with PUT 'foo' to zookeeper")
		}
	} else if len(args) == 0 || len(args) > 2 {
		fmt.Fprintln(os.Stderr, cmd.Usage())
		os.Exit(1)
	}

	var k, end string
	if singleKey {
		k = "foo"
	} else {
		k = args[0]
		if len(args) == 2 {
			end = args[1]
		}
	}

	if rangeConsistency == "l" {
		fmt.Println("bench with linearizable range")
	} else if rangeConsistency == "s" {
		fmt.Println("bench with serializable range")
	} else {
		fmt.Fprintln(os.Stderr, cmd.Usage())
		os.Exit(1)
	}

	results = make(chan result)
	requests := make(chan request, totalClients)
	bar = pb.New(rangeTotal)

	bar.Format("Bom !")
	bar.Start()

	switch database {
	case "etcd":
		clients := mustCreateClients(totalClients, totalConns)
		for i := range clients {
			wg.Add(1)
			go doRange(clients[i].KV, requests)
		}
		defer func() {
			for i := range clients {
				clients[i].Close()
			}
		}()
	case "zk":
		conns := mustCreateConnsZk(totalConns)
		defer func() {
			for i := range conns {
				conns[i].Close()
			}
		}()
		for i := range conns {
			wg.Add(1)
			go doRangeZk(conns[i], requests)
		}
	default:
		log.Fatalf("unknown database %s", database)
	}

	pdoneC := printReport(results)
	go func() {
		for i := 0; i < rangeTotal; i++ {
			switch database {
			case "etcd":
				opts := []v3.OpOption{v3.WithRange(end)}
				if rangeConsistency == "s" {
					opts = append(opts, v3.WithSerializable())
				}
				requests <- request{etcdOp: v3.OpGet(k, opts...)}
			case "zk":
				requests <- request{zkOp: zkOp{key: k}}
			}
		}
		close(requests)
	}()

	wg.Wait()

	bar.Finish()

	close(results)
	<-pdoneC
}

func doRange(client v3.KV, requests <-chan request) {
	defer wg.Done()

	for req := range requests {
		op := req.etcdOp
		st := time.Now()
		_, err := client.Do(context.Background(), op)

		var errStr string
		if err != nil {
			errStr = err.Error()
		}
		results <- result{errStr: errStr, duration: time.Since(st), happened: time.Now()}
		bar.Increment()
	}
}

func doRangeZk(conn *zk.Conn, requests <-chan request) {
	defer wg.Done()

	for req := range requests {
		op := req.zkOp
		st := time.Now()
		_, _, err := conn.Get(op.key)

		var errStr string
		if err != nil {
			errStr = err.Error()
		}
		results <- result{errStr: errStr, duration: time.Since(st), happened: time.Now()}
		bar.Increment()
	}
}
<|endoftext|>


<commit_before>// Copyright 2016 Google Inc. All Rights Reserved.
//
// 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 pubsub

import (
	"errors"
	"reflect"
	"sort"
	"sync"
	"testing"
	"time"

	"golang.org/x/net/context"
)

func TestKeepAlive(t *testing.T) {
	tick := make(chan time.Time)
	deadline := time.Nanosecond * 15
	s := &testService{modDeadlineCalled: make(chan modDeadlineCall)}

	checkModDeadlineCall := func(ackIDs []string) {
		got := <-s.modDeadlineCalled
		sort.Strings(got.ackIDs)

		want := modDeadlineCall{
			subName:  "subname",
			deadline: deadline,
			ackIDs:   ackIDs,
		}

		if !reflect.DeepEqual(got, want) {
			t.Errorf("keepalive: got:\n%v\nwant:\n%v", got, want)
		}
	}

	ka := &keepAlive{
		s:             s,
		Ctx:           context.Background(),
		Sub:           "subname",
		ExtensionTick: tick,
		Deadline:      deadline,
		MaxExtension:  time.Hour,
	}
	ka.Start()

	ka.Add("a")
	ka.Add("b")
	tick <- time.Time{}
	checkModDeadlineCall([]string{"a", "b"})
	ka.Add("c")
	ka.Remove("b")
	tick <- time.Time{}
	checkModDeadlineCall([]string{"a", "c"})
	ka.Remove("a")
	ka.Remove("c")
	ka.Add("d")
	tick <- time.Time{}
	checkModDeadlineCall([]string{"d"})

	ka.Remove("d")
	ka.Stop()
}

// TestKeepAliveStop checks that Stop blocks until all ackIDs have been removed.
func TestKeepAliveStop(t *testing.T) {
	tick := 100 * time.Microsecond
	ticker := time.NewTicker(tick)
	defer ticker.Stop()

	s := &testService{modDeadlineCalled: make(chan modDeadlineCall, 100)}

	ka := &keepAlive{
		s:             s,
		Ctx:           context.Background(),
		ExtensionTick: ticker.C,
		MaxExtension:  time.Hour,
	}
	ka.Start()

	events := make(chan string, 10)

	// Add an ackID so that ka.Stop will not return immediately.
	ka.Add("a")

	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		time.Sleep(tick * 10)
		events <- "pre-remove"
		ka.Remove("a")
		time.Sleep(tick * 10)
		events <- "post-second-sleep"
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()
		events <- "pre-stop"
		ka.Stop()
		events <- "stopped"

	}()

	wg.Wait()
	close(events)
	eventSequence := []string{}
	for e := range events {
		eventSequence = append(eventSequence, e)
	}

	want := []string{"pre-stop", "pre-remove", "stopped", "post-second-sleep"}
	if !reflect.DeepEqual(eventSequence, want) {
		t.Errorf("keepalive eventsequence: got:\n%v\nwant:\n%v", eventSequence, want)
	}
}

// TestMaxExtensionDeadline checks we stop extending after the configured duration.
func TestMaxExtensionDeadline(t *testing.T) {
	ticker := time.NewTicker(100 * time.Microsecond)
	defer ticker.Stop()

	s := &testService{modDeadlineCalled: make(chan modDeadlineCall, 100)}

	maxExtension := time.Millisecond
	ka := &keepAlive{
		s:             s,
		Ctx:           context.Background(),
		ExtensionTick: ticker.C,
		MaxExtension:  maxExtension,
	}
	ka.Start()

	ka.Add("a")
	stopped := make(chan struct{})

	go func() {
		ka.Stop()
		stopped <- struct{}{}
	}()

	select {
	case <-stopped:
	case <-time.After(maxExtension + 2*time.Second):
		t.Fatalf("keepalive failed to stop after maxExtension deadline")
	}
}

func TestKeepAliveStopsWhenAllAckIDsRemoved(t *testing.T) {
	s := &testService{}

	maxExtension := time.Millisecond
	ka := &keepAlive{
		s:             s,
		Ctx:           context.Background(),
		ExtensionTick: make(chan time.Time),
		MaxExtension:  maxExtension,
	}
	ka.Start()
	ka.Add("a")

	stopped := make(chan struct{})

	go func() {
		ka.Stop()
		stopped <- struct{}{}
	}()

	time.Sleep(time.Microsecond)
	// No extension tick is ever sent, but this should be enough to get ka to stop.
	ka.Remove("a")

	select {
	case <-stopped:
	case <-time.After(maxExtension / 2):
		t.Fatalf("keepalive failed to stop before maxExtension deadline")
	}
}

func TestKeepAliveStopsImmediatelyForNoAckIDs(t *testing.T) {
	ticker := time.NewTicker(100 * time.Microsecond)
	defer ticker.Stop()

	s := &testService{modDeadlineCalled: make(chan modDeadlineCall, 100)}

	maxExtension := time.Millisecond
	ka := &keepAlive{
		s:             s,
		Ctx:           context.Background(),
		ExtensionTick: ticker.C,
		MaxExtension:  maxExtension,
	}
	ka.Start()

	stopped := make(chan struct{})

	go func() {
		// There are no items in ka, so this should return immediately.
		ka.Stop()
		stopped <- struct{}{}
	}()

	select {
	case <-stopped:
	case <-time.After(maxExtension / 2):
		t.Fatalf("keepalive failed to stop before maxExtension deadline")
	}
}

// extendCallResult contains a list of ackIDs which are expected in an ackID
// extension request, along with the result that should be returned.
type extendCallResult struct {
	ackIDs []string
	err    error
}

// extendService implements modifyAckDeadline using a hard-coded list of extendCallResults.
type extendService struct {
	service

	calls []extendCallResult

	t *testing.T // used for error logging.
}

func (es *extendService) modifyAckDeadline(ctx context.Context, subName string, deadline time.Duration, ackIDs []string) error {
	if len(es.calls) == 0 {
		es.t.Fatalf("unexpected call to modifyAckDeadline: ackIDs: %v", ackIDs)
	}
	call := es.calls[0]
	es.calls = es.calls[1:]

	if got, want := ackIDs, call.ackIDs; !reflect.DeepEqual(got, want) {
		es.t.Errorf("unexpected arguments to modifyAckDeadline: got: %v ; want: %v", got, want)
	}
	return call.err
}

// Test implementation returns the first 2 elements as head, and the rest as tail.
func (es *extendService) splitAckIDs(ids []string) ([]string, []string) {
	if len(ids) < 2 {
		return ids, nil
	}
	return ids[:2], ids[2:]
}
func TestKeepAliveSplitsBatches(t *testing.T) {
	type testCase struct {
		calls []extendCallResult
	}
	for _, tc := range []testCase{
		{
			calls: []extendCallResult{
				{
					ackIDs: []string{"a", "b"},
				},
				{
					ackIDs: []string{"c", "d"},
				},
				{
					ackIDs: []string{"e", "f"},
				},
			},
		},
		{
			calls: []extendCallResult{
				{
					ackIDs: []string{"a", "b"},
					err:    errors.New("bang"),
				},
				// On error we retry once.
				{
					ackIDs: []string{"a", "b"},
					err:    errors.New("bang"),
				},
				// We give up after failing twice, so we move on to the next set, "c" and "d"
				{
					ackIDs: []string{"c", "d"},
					err:    errors.New("bang"),
				},
				// Again, we retry once.
				{
					ackIDs: []string{"c", "d"},
				},
				{
					ackIDs: []string{"e", "f"},
				},
			},
		},
	} {
		s := &extendService{
			t:     t,
			calls: tc.calls,
		}

		ka := &keepAlive{
			s:   s,
			Ctx: context.Background(),
			Sub: "subname",
		}

		ka.extendDeadlines([]string{"a", "b", "c", "d", "e", "f"})

		if len(s.calls) != 0 {
			t.Errorf("expected extend calls did not occur: %v", s.calls)
		}
	}
}
<commit_msg>pubsub: fix a series of unit tests in race mode.<commit_after>// Copyright 2016 Google Inc. All Rights Reserved.
//
// 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 pubsub

import (
	"errors"
	"reflect"
	"sort"
	"testing"
	"time"

	"golang.org/x/net/context"
)

func TestKeepAliveExtendsDeadline(t *testing.T) {
	ticker := make(chan time.Time)
	deadline := time.Nanosecond * 15
	s := &testService{modDeadlineCalled: make(chan modDeadlineCall)}

	checkModDeadlineCall := func(ackIDs []string) {
		got := <-s.modDeadlineCalled
		sort.Strings(got.ackIDs)

		want := modDeadlineCall{
			subName:  "subname",
			deadline: deadline,
			ackIDs:   ackIDs,
		}

		if !reflect.DeepEqual(got, want) {
			t.Errorf("keepalive: got:\n%v\nwant:\n%v", got, want)
		}
	}

	ka := &keepAlive{
		s:             s,
		Ctx:           context.Background(),
		Sub:           "subname",
		ExtensionTick: ticker,
		Deadline:      deadline,
		MaxExtension:  time.Hour,
	}
	ka.Start()

	ka.Add("a")
	ka.Add("b")
	ticker <- time.Time{}
	checkModDeadlineCall([]string{"a", "b"})
	ka.Add("c")
	ka.Remove("b")
	ticker <- time.Time{}
	checkModDeadlineCall([]string{"a", "c"})
	ka.Remove("a")
	ka.Remove("c")
	ka.Add("d")
	ticker <- time.Time{}
	checkModDeadlineCall([]string{"d"})

	ka.Remove("d")
	ka.Stop()
}

func TestKeepAliveStopsWhenNoItem(t *testing.T) {
	ticker := make(chan time.Time)
	stopped := make(chan bool)
	s := &testService{modDeadlineCalled: make(chan modDeadlineCall, 3)}
	ka := &keepAlive{
		s:             s,
		Ctx:           context.Background(),
		ExtensionTick: ticker,
	}

	ka.Start()

	// There should be no call to modifyAckDeadline since there is no item.
	ticker <- time.Time{}

	go func() {
		ka.Stop() // No items; should not block
		if len(s.modDeadlineCalled) > 0 {
			t.Errorf("unexpected extension to non-existent items: %v", <-s.modDeadlineCalled)
		}
		close(stopped)
	}()

	select {
	case <-stopped:
	case <-time.After(time.Second):
		t.Errorf("keepAlive timed out waiting for stop")
	}
}

func TestKeepAliveStopsWhenItemsExpired(t *testing.T) {
	ticker := make(chan time.Time)
	stopped := make(chan bool)
	s := &testService{modDeadlineCalled: make(chan modDeadlineCall, 2)}
	ka := &keepAlive{
		s:             s,
		Ctx:           context.Background(),
		ExtensionTick: ticker,
		MaxExtension:  time.Duration(0), // Should expire items at the first tick.
	}

	ka.Start()
	ka.Add("a")
	ka.Add("b")

	// There should be no call to modifyAckDeadline since both items are expired.
	ticker <- time.Time{}

	go func() {
		ka.Stop() // No live items; should not block.
		if len(s.modDeadlineCalled) > 0 {
			t.Errorf("unexpected extension to expired items")
		}
		close(stopped)
	}()

	select {
	case <-stopped:
	case <-time.After(time.Second):
		t.Errorf("timed out waiting for stop")
	}
}

func TestKeepAliveBlocksUntilAllItemsRemoved(t *testing.T) {
	ticker := make(chan time.Time)
	eventc := make(chan string, 3)
	s := &testService{modDeadlineCalled: make(chan modDeadlineCall)}
	ka := &keepAlive{
		s:             s,
		Ctx:           context.Background(),
		ExtensionTick: ticker,
		MaxExtension:  time.Hour, // Should not expire.
	}

	ka.Start()
	ka.Add("a")
	ka.Add("b")

	go func() {
		ticker <- time.Time{}

		// We expect a call since both items should be extended.
		select {
		case args := <-s.modDeadlineCalled:
			sort.Strings(args.ackIDs)
			got := args.ackIDs
			want := []string{"a", "b"}
			if !reflect.DeepEqual(got, want) {
				t.Errorf("mismatching IDs:\ngot  %v\nwant %v", got, want)
			}
		case <-time.After(time.Second):
			t.Errorf("timed out waiting for deadline extend call")
		}

		time.Sleep(10 * time.Millisecond)

		eventc <- "pre-remove-b"
		// Remove one item, Stop should still be waiting.
		ka.Remove("b")

		ticker <- time.Time{}

		// We expect a call since the item is still alive.
		select {
		case args := <-s.modDeadlineCalled:
			got := args.ackIDs
			want := []string{"a"}
			if !reflect.DeepEqual(got, want) {
				t.Errorf("mismatching IDs:\ngot  %v\nwant %v", got, want)
			}
		case <-time.After(time.Second):
			t.Errorf("timed out waiting for deadline extend call")
		}

		time.Sleep(10 * time.Millisecond)

		eventc <- "pre-remove-a"
		// Remove the last item so that Stop can proceed.
		ka.Remove("a")
	}()

	go func() {
		ka.Stop() // Should block all item are removed.
		eventc <- "post-stop"
	}()

	for i, want := range []string{"pre-remove-b", "pre-remove-a", "post-stop"} {
		select {
		case got := <-eventc:
			if got != want {
				t.Errorf("event #%d:\ngot  %v\nwant %v", i, got, want)
			}
		case <-time.After(time.Second):
			t.Errorf("time out waiting for #%d event: want %v", i, want)
		}
	}
}

// extendCallResult contains a list of ackIDs which are expected in an ackID
// extension request, along with the result that should be returned.
type extendCallResult struct {
	ackIDs []string
	err    error
}

// extendService implements modifyAckDeadline using a hard-coded list of extendCallResults.
type extendService struct {
	service

	calls []extendCallResult

	t *testing.T // used for error logging.
}

func (es *extendService) modifyAckDeadline(ctx context.Context, subName string, deadline time.Duration, ackIDs []string) error {
	if len(es.calls) == 0 {
		es.t.Fatalf("unexpected call to modifyAckDeadline: ackIDs: %v", ackIDs)
	}
	call := es.calls[0]
	es.calls = es.calls[1:]

	if got, want := ackIDs, call.ackIDs; !reflect.DeepEqual(got, want) {
		es.t.Errorf("unexpected arguments to modifyAckDeadline: got: %v ; want: %v", got, want)
	}
	return call.err
}

// Test implementation returns the first 2 elements as head, and the rest as tail.
func (es *extendService) splitAckIDs(ids []string) ([]string, []string) {
	if len(ids) < 2 {
		return ids, nil
	}
	return ids[:2], ids[2:]
}

func TestKeepAliveSplitsBatches(t *testing.T) {
	type testCase struct {
		calls []extendCallResult
	}
	for _, tc := range []testCase{
		{
			calls: []extendCallResult{
				{
					ackIDs: []string{"a", "b"},
				},
				{
					ackIDs: []string{"c", "d"},
				},
				{
					ackIDs: []string{"e", "f"},
				},
			},
		},
		{
			calls: []extendCallResult{
				{
					ackIDs: []string{"a", "b"},
					err:    errors.New("bang"),
				},
				// On error we retry once.
				{
					ackIDs: []string{"a", "b"},
					err:    errors.New("bang"),
				},
				// We give up after failing twice, so we move on to the next set, "c" and "d".
				{
					ackIDs: []string{"c", "d"},
					err:    errors.New("bang"),
				},
				// Again, we retry once.
				{
					ackIDs: []string{"c", "d"},
				},
				{
					ackIDs: []string{"e", "f"},
				},
			},
		},
	} {
		s := &extendService{
			t:     t,
			calls: tc.calls,
		}

		ka := &keepAlive{
			s:   s,
			Ctx: context.Background(),
			Sub: "subname",
		}

		ka.extendDeadlines([]string{"a", "b", "c", "d", "e", "f"})

		if len(s.calls) != 0 {
			t.Errorf("expected extend calls did not occur: %v", s.calls)
		}
	}
}
<|endoftext|>