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

package ec2

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"reflect"
	"strconv"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/pkg/errors"
	"github.com/tsuru/monsterqueue"
	"github.com/tsuru/tsuru/iaas"
	"github.com/tsuru/tsuru/log"
	"github.com/tsuru/tsuru/queue"
)

const defaultRegion = "us-east-1"

func init() {
	iaas.RegisterIaasProvider("ec2", newEC2IaaS)
}

type EC2IaaS struct {
	base iaas.UserDataIaaS
}

func newEC2IaaS(name string) iaas.IaaS {
	return &EC2IaaS{base: iaas.UserDataIaaS{NamedIaaS: iaas.NamedIaaS{BaseIaaSName: "ec2", IaaSName: name}}}
}

func (i *EC2IaaS) createEC2Handler(regionOrEndpoint string) (*ec2.EC2, error) {
	keyId, err := i.base.GetConfigString("key-id")
	if err != nil {
		return nil, err
	}
	secretKey, err := i.base.GetConfigString("secret-key")
	if err != nil {
		return nil, err
	}
	var region, endpoint string
	if strings.HasPrefix(regionOrEndpoint, "http") {
		endpoint = regionOrEndpoint
		region = defaultRegion
	} else {
		region = regionOrEndpoint
	}
	config := aws.Config{
		Credentials: credentials.NewStaticCredentials(keyId, secretKey, ""),
		Region:      aws.String(region),
		Endpoint:    aws.String(endpoint),
	}
	return ec2.New(session.New(&config)), nil
}

func (i *EC2IaaS) waitForDnsName(ec2Inst *ec2.EC2, instanceID string, createParams map[string]string) (string, error) {
	rawWait, _ := i.base.GetConfigString("wait-timeout")
	maxWaitTime, _ := strconv.Atoi(rawWait)
	if maxWaitTime == 0 {
		maxWaitTime = 300
	}
	q, err := queue.Queue()
	if err != nil {
		return "", err
	}
	taskName := fmt.Sprintf("ec2-wait-machine-%s", i.base.IaaSName)
	jobParams := monsterqueue.JobParams{
		"region":    ec2Inst.Config.Region,
		"endpoint":  ec2Inst.Config.Endpoint,
		"machineId": instanceID,
		"timeout":   maxWaitTime,
	}
	if rawInterfaceIdx, ok := createParams["network-index"]; ok {
		if interfaceIdx, atoiErr := strconv.Atoi(rawInterfaceIdx); atoiErr == nil {
			jobParams["networkIndex"] = interfaceIdx
		}
	}
	waitDuration := time.Duration(maxWaitTime) * time.Second
	job, err := q.EnqueueWait(taskName, jobParams, waitDuration)
	if err != nil {
		if err == monsterqueue.ErrQueueWaitTimeout {
			return "", errors.Errorf("ec2: time out after %v waiting for instance %s to start", waitDuration, instanceID)
		}
		return "", err
	}
	result, err := job.Result()
	if err != nil {
		return "", err
	}
	return result.(string), nil
}

func (i *EC2IaaS) Initialize() error {
	q, err := queue.Queue()
	if err != nil {
		return err
	}
	return q.RegisterTask(&ec2WaitTask{iaas: i})
}

func (i *EC2IaaS) Describe() string {
	return `EC2 IaaS required params:
  image=<image id>         Image AMI ID
  type=<instance type>     Your template uuid

Optional params:
  region=<region>          Chosen region, defaults to us-east-1
  securityGroup=<group>    Chosen security group
  keyName=<key name>       Key name for machine
`
}

func (i *EC2IaaS) DeleteMachine(m *iaas.Machine) error {
	regionOrEndpoint := getRegionOrEndpoint(m.CreationParams, false)
	if regionOrEndpoint == "" {
		return errors.Errorf("region or endpoint creation param required")
	}
	ec2Inst, err := i.createEC2Handler(regionOrEndpoint)
	if err != nil {
		return err
	}
	input := ec2.TerminateInstancesInput{InstanceIds: []*string{&m.Id}}
	_, err = ec2Inst.TerminateInstances(&input)
	return err
}

type invalidFieldError struct {
	fieldName    string
	convertError error
}

func (err *invalidFieldError) Error() string {
	return fmt.Sprintf("invalid value for the field %q: %s", err.fieldName, err.convertError)
}

func (i *EC2IaaS) buildRunInstancesOptions(params map[string]string) (ec2.RunInstancesInput, error) {
	result := ec2.RunInstancesInput{
		MaxCount: aws.Int64(1),
		MinCount: aws.Int64(1),
	}
	forbiddenFields := []string{
		"maxcount", "mincount", "dryrun", "monitoring",
	}
	aliases := map[string]string{
		"image":         "imageid",
		"type":          "instancetype",
		"securitygroup": "securitygroups",
		"ebs-optimized": "ebsoptimized",
	}
	refType := reflect.TypeOf(result)
	refValue := reflect.ValueOf(&result)
	for key, value := range params {
		field, ok := refType.FieldByNameFunc(func(name string) bool {
			lowerName := strings.ToLower(name)
			for _, field := range forbiddenFields {
				if lowerName == field {
					return false
				}
			}
			lowerKey := strings.ToLower(key)
			if aliased, ok := aliases[lowerKey]; ok {
				lowerKey = aliased
			}
			return lowerName == lowerKey
		})
		if !ok {
			continue
		}
		fieldType := field.Type
		fieldValue := refValue.Elem().FieldByIndex(field.Index)
		if !fieldValue.IsValid() || !fieldValue.CanSet() {
			continue
		}
		switch fieldType.Kind() {
		case reflect.Ptr:
			switch fieldType.Elem().Kind() {
			case reflect.String:
				copy := value
				fieldValue.Set(reflect.ValueOf(&copy))
			case reflect.Int64:
				intValue, err := strconv.ParseInt(value, 10, 64)
				if err != nil {
					return result, &invalidFieldError{
						fieldName:    key,
						convertError: err,
					}
				}
				fieldValue.Set(reflect.ValueOf(&intValue))
			case reflect.Bool:
				boolValue, err := strconv.ParseBool(value)
				if err != nil {
					return result, &invalidFieldError{
						fieldName:    key,
						convertError: err,
					}
				}
				fieldValue.Set(reflect.ValueOf(&boolValue))
			case reflect.Struct:
				err := i.loadStruct(fieldValue, fieldType, []byte(value))
				if err != nil {
					return result, &invalidFieldError{
						fieldName:    key,
						convertError: err,
					}
				}
			}
		case reflect.Slice:
			switch fieldType.Elem().Elem().Kind() {
			case reflect.String:
				parts := strings.Split(value, ",")
				values := make([]*string, len(parts))
				for i, part := range parts {
					values[i] = aws.String(part)
				}
				fieldValue.Set(reflect.ValueOf(values))
			case reflect.Struct:
				var raw []map[string]interface{}
				err := json.Unmarshal([]byte(value), &raw)
				if err != nil {
					return result, &invalidFieldError{
						fieldName:    key,
						convertError: err,
					}
				}
				val, err := i.translateSlice(raw, fieldType)
				if err != nil {
					return result, &invalidFieldError{
						fieldName:    key,
						convertError: err,
					}
				}
				fieldValue.Set(val)
			}
		}
	}

	// Manual configuration
	if monitoring, ok := params["monitoring-enabled"]; ok {
		value, _ := strconv.ParseBool(monitoring)
		result.Monitoring = &ec2.RunInstancesMonitoringEnabled{
			Enabled: aws.Bool(value),
		}
	}

	return result, nil
}

func (i *EC2IaaS) translateSlice(in []map[string]interface{}, t reflect.Type) (reflect.Value, error) {
	result := reflect.MakeSlice(t, len(in), len(in))
	for idx, value := range in {
		data, _ := json.Marshal(value)
		err := i.loadStruct(result.Index(idx), t.Elem(), data)
		if err != nil {
			return result, err
		}
	}
	return result, nil
}

func (i *EC2IaaS) loadStruct(value reflect.Value, t reflect.Type, data []byte) error {
	raw := value.Interface()
	if value.IsNil() {
		raw = reflect.New(t.Elem()).Interface()
	}
	err := json.Unmarshal(data, &raw)
	if err != nil {
		return err
	}
	value.Set(reflect.ValueOf(raw))
	return nil
}

func (i *EC2IaaS) CreateMachine(params map[string]string) (*iaas.Machine, error) {
	regionOrEndpoint := getRegionOrEndpoint(params, true)
	userData, err := i.base.ReadUserData()
	if err != nil {
		return nil, err
	}
	options, err := i.buildRunInstancesOptions(params)
	if err != nil {
		return nil, err
	}
	options.UserData = aws.String(base64.StdEncoding.EncodeToString([]byte(userData)))
	if options.ImageId == nil || *options.ImageId == "" {
		return nil, errors.Errorf("the parameter %q is required", "imageid")
	}
	if options.InstanceType == nil || *options.InstanceType == "" {
		return nil, errors.Errorf("the parameter %q is required", "instancetype")
	}
	ec2Inst, err := i.createEC2Handler(regionOrEndpoint)
	if err != nil {
		return nil, err
	}
	resp, err := ec2Inst.RunInstances(&options)
	if err != nil {
		return nil, err
	}
	if len(resp.Instances) == 0 {
		return nil, errors.Errorf("no instance created")
	}
	runInst := resp.Instances[0]
	if tags, ok := params["tags"]; ok {
		var ec2Tags []*ec2.Tag
		tagList := strings.Split(tags, ",")
		ec2Tags = make([]*ec2.Tag, 0, len(tagList))
		for _, tag := range tagList {
			if strings.Contains(tag, ":") {
				parts := strings.SplitN(tag, ":", 2)
				ec2Tags = append(ec2Tags, &ec2.Tag{
					Key:   aws.String(parts[0]),
					Value: aws.String(parts[1]),
				})
			}
		}
		if len(ec2Tags) > 0 {
			input := ec2.CreateTagsInput{
				Resources: []*string{runInst.InstanceId},
				Tags:      ec2Tags,
			}
			_, err = ec2Inst.CreateTags(&input)
			if err != nil {
				log.Errorf("failed to tag EC2 instance: %s", err)
			}
		}
	}
	dnsName, err := i.waitForDnsName(ec2Inst, aws.StringValue(runInst.InstanceId), params)
	if err != nil {
		return nil, err
	}
	machine := iaas.Machine{
		Id:      aws.StringValue(runInst.InstanceId),
		Status:  aws.StringValue(runInst.State.Name),
		Address: dnsName,
	}
	return &machine, nil
}

func getRegionOrEndpoint(params map[string]string, useDefault bool) string {
	regionOrEndpoint := params["endpoint"]
	if regionOrEndpoint == "" {
		regionOrEndpoint = params["region"]
		if regionOrEndpoint == "" && useDefault {
			regionOrEndpoint = defaultRegion
		}
	}
	return regionOrEndpoint
}
<commit_msg>iaas/ec2: use custom http.Client for aws connections<commit_after>// Copyright 2014 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package ec2

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"reflect"
	"strconv"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/pkg/errors"
	"github.com/tsuru/monsterqueue"
	"github.com/tsuru/tsuru/iaas"
	"github.com/tsuru/tsuru/log"
	tsuruNet "github.com/tsuru/tsuru/net"
	"github.com/tsuru/tsuru/queue"
)

const defaultRegion = "us-east-1"

func init() {
	iaas.RegisterIaasProvider("ec2", newEC2IaaS)
}

type EC2IaaS struct {
	base iaas.UserDataIaaS
}

func newEC2IaaS(name string) iaas.IaaS {
	return &EC2IaaS{base: iaas.UserDataIaaS{NamedIaaS: iaas.NamedIaaS{BaseIaaSName: "ec2", IaaSName: name}}}
}

func (i *EC2IaaS) createEC2Handler(regionOrEndpoint string) (*ec2.EC2, error) {
	keyId, err := i.base.GetConfigString("key-id")
	if err != nil {
		return nil, err
	}
	secretKey, err := i.base.GetConfigString("secret-key")
	if err != nil {
		return nil, err
	}
	var region, endpoint string
	if strings.HasPrefix(regionOrEndpoint, "http") {
		endpoint = regionOrEndpoint
		region = defaultRegion
	} else {
		region = regionOrEndpoint
	}
	config := aws.Config{
		Credentials: credentials.NewStaticCredentials(keyId, secretKey, ""),
		Region:      aws.String(region),
		Endpoint:    aws.String(endpoint),
		HTTPClient:  tsuruNet.Dial5Full300ClientNoKeepAlive,
	}
	return ec2.New(session.New(&config)), nil
}

func (i *EC2IaaS) waitForDnsName(ec2Inst *ec2.EC2, instanceID string, createParams map[string]string) (string, error) {
	rawWait, _ := i.base.GetConfigString("wait-timeout")
	maxWaitTime, _ := strconv.Atoi(rawWait)
	if maxWaitTime == 0 {
		maxWaitTime = 300
	}
	q, err := queue.Queue()
	if err != nil {
		return "", err
	}
	taskName := fmt.Sprintf("ec2-wait-machine-%s", i.base.IaaSName)
	jobParams := monsterqueue.JobParams{
		"region":    ec2Inst.Config.Region,
		"endpoint":  ec2Inst.Config.Endpoint,
		"machineId": instanceID,
		"timeout":   maxWaitTime,
	}
	if rawInterfaceIdx, ok := createParams["network-index"]; ok {
		if interfaceIdx, atoiErr := strconv.Atoi(rawInterfaceIdx); atoiErr == nil {
			jobParams["networkIndex"] = interfaceIdx
		}
	}
	waitDuration := time.Duration(maxWaitTime) * time.Second
	job, err := q.EnqueueWait(taskName, jobParams, waitDuration)
	if err != nil {
		if err == monsterqueue.ErrQueueWaitTimeout {
			return "", errors.Errorf("ec2: time out after %v waiting for instance %s to start", waitDuration, instanceID)
		}
		return "", err
	}
	result, err := job.Result()
	if err != nil {
		return "", err
	}
	return result.(string), nil
}

func (i *EC2IaaS) Initialize() error {
	q, err := queue.Queue()
	if err != nil {
		return err
	}
	return q.RegisterTask(&ec2WaitTask{iaas: i})
}

func (i *EC2IaaS) Describe() string {
	return `EC2 IaaS required params:
  image=<image id>         Image AMI ID
  type=<instance type>     Your template uuid

Optional params:
  region=<region>          Chosen region, defaults to us-east-1
  securityGroup=<group>    Chosen security group
  keyName=<key name>       Key name for machine
`
}

func (i *EC2IaaS) DeleteMachine(m *iaas.Machine) error {
	regionOrEndpoint := getRegionOrEndpoint(m.CreationParams, false)
	if regionOrEndpoint == "" {
		return errors.Errorf("region or endpoint creation param required")
	}
	ec2Inst, err := i.createEC2Handler(regionOrEndpoint)
	if err != nil {
		return err
	}
	input := ec2.TerminateInstancesInput{InstanceIds: []*string{&m.Id}}
	_, err = ec2Inst.TerminateInstances(&input)
	return err
}

type invalidFieldError struct {
	fieldName    string
	convertError error
}

func (err *invalidFieldError) Error() string {
	return fmt.Sprintf("invalid value for the field %q: %s", err.fieldName, err.convertError)
}

func (i *EC2IaaS) buildRunInstancesOptions(params map[string]string) (ec2.RunInstancesInput, error) {
	result := ec2.RunInstancesInput{
		MaxCount: aws.Int64(1),
		MinCount: aws.Int64(1),
	}
	forbiddenFields := []string{
		"maxcount", "mincount", "dryrun", "monitoring",
	}
	aliases := map[string]string{
		"image":         "imageid",
		"type":          "instancetype",
		"securitygroup": "securitygroups",
		"ebs-optimized": "ebsoptimized",
	}
	refType := reflect.TypeOf(result)
	refValue := reflect.ValueOf(&result)
	for key, value := range params {
		field, ok := refType.FieldByNameFunc(func(name string) bool {
			lowerName := strings.ToLower(name)
			for _, field := range forbiddenFields {
				if lowerName == field {
					return false
				}
			}
			lowerKey := strings.ToLower(key)
			if aliased, ok := aliases[lowerKey]; ok {
				lowerKey = aliased
			}
			return lowerName == lowerKey
		})
		if !ok {
			continue
		}
		fieldType := field.Type
		fieldValue := refValue.Elem().FieldByIndex(field.Index)
		if !fieldValue.IsValid() || !fieldValue.CanSet() {
			continue
		}
		switch fieldType.Kind() {
		case reflect.Ptr:
			switch fieldType.Elem().Kind() {
			case reflect.String:
				copy := value
				fieldValue.Set(reflect.ValueOf(&copy))
			case reflect.Int64:
				intValue, err := strconv.ParseInt(value, 10, 64)
				if err != nil {
					return result, &invalidFieldError{
						fieldName:    key,
						convertError: err,
					}
				}
				fieldValue.Set(reflect.ValueOf(&intValue))
			case reflect.Bool:
				boolValue, err := strconv.ParseBool(value)
				if err != nil {
					return result, &invalidFieldError{
						fieldName:    key,
						convertError: err,
					}
				}
				fieldValue.Set(reflect.ValueOf(&boolValue))
			case reflect.Struct:
				err := i.loadStruct(fieldValue, fieldType, []byte(value))
				if err != nil {
					return result, &invalidFieldError{
						fieldName:    key,
						convertError: err,
					}
				}
			}
		case reflect.Slice:
			switch fieldType.Elem().Elem().Kind() {
			case reflect.String:
				parts := strings.Split(value, ",")
				values := make([]*string, len(parts))
				for i, part := range parts {
					values[i] = aws.String(part)
				}
				fieldValue.Set(reflect.ValueOf(values))
			case reflect.Struct:
				var raw []map[string]interface{}
				err := json.Unmarshal([]byte(value), &raw)
				if err != nil {
					return result, &invalidFieldError{
						fieldName:    key,
						convertError: err,
					}
				}
				val, err := i.translateSlice(raw, fieldType)
				if err != nil {
					return result, &invalidFieldError{
						fieldName:    key,
						convertError: err,
					}
				}
				fieldValue.Set(val)
			}
		}
	}

	// Manual configuration
	if monitoring, ok := params["monitoring-enabled"]; ok {
		value, _ := strconv.ParseBool(monitoring)
		result.Monitoring = &ec2.RunInstancesMonitoringEnabled{
			Enabled: aws.Bool(value),
		}
	}

	return result, nil
}

func (i *EC2IaaS) translateSlice(in []map[string]interface{}, t reflect.Type) (reflect.Value, error) {
	result := reflect.MakeSlice(t, len(in), len(in))
	for idx, value := range in {
		data, _ := json.Marshal(value)
		err := i.loadStruct(result.Index(idx), t.Elem(), data)
		if err != nil {
			return result, err
		}
	}
	return result, nil
}

func (i *EC2IaaS) loadStruct(value reflect.Value, t reflect.Type, data []byte) error {
	raw := value.Interface()
	if value.IsNil() {
		raw = reflect.New(t.Elem()).Interface()
	}
	err := json.Unmarshal(data, &raw)
	if err != nil {
		return err
	}
	value.Set(reflect.ValueOf(raw))
	return nil
}

func (i *EC2IaaS) CreateMachine(params map[string]string) (*iaas.Machine, error) {
	regionOrEndpoint := getRegionOrEndpoint(params, true)
	userData, err := i.base.ReadUserData()
	if err != nil {
		return nil, err
	}
	options, err := i.buildRunInstancesOptions(params)
	if err != nil {
		return nil, err
	}
	options.UserData = aws.String(base64.StdEncoding.EncodeToString([]byte(userData)))
	if options.ImageId == nil || *options.ImageId == "" {
		return nil, errors.Errorf("the parameter %q is required", "imageid")
	}
	if options.InstanceType == nil || *options.InstanceType == "" {
		return nil, errors.Errorf("the parameter %q is required", "instancetype")
	}
	ec2Inst, err := i.createEC2Handler(regionOrEndpoint)
	if err != nil {
		return nil, err
	}
	resp, err := ec2Inst.RunInstances(&options)
	if err != nil {
		return nil, err
	}
	if len(resp.Instances) == 0 {
		return nil, errors.Errorf("no instance created")
	}
	runInst := resp.Instances[0]
	if tags, ok := params["tags"]; ok {
		var ec2Tags []*ec2.Tag
		tagList := strings.Split(tags, ",")
		ec2Tags = make([]*ec2.Tag, 0, len(tagList))
		for _, tag := range tagList {
			if strings.Contains(tag, ":") {
				parts := strings.SplitN(tag, ":", 2)
				ec2Tags = append(ec2Tags, &ec2.Tag{
					Key:   aws.String(parts[0]),
					Value: aws.String(parts[1]),
				})
			}
		}
		if len(ec2Tags) > 0 {
			input := ec2.CreateTagsInput{
				Resources: []*string{runInst.InstanceId},
				Tags:      ec2Tags,
			}
			_, err = ec2Inst.CreateTags(&input)
			if err != nil {
				log.Errorf("failed to tag EC2 instance: %s", err)
			}
		}
	}
	dnsName, err := i.waitForDnsName(ec2Inst, aws.StringValue(runInst.InstanceId), params)
	if err != nil {
		return nil, err
	}
	machine := iaas.Machine{
		Id:      aws.StringValue(runInst.InstanceId),
		Status:  aws.StringValue(runInst.State.Name),
		Address: dnsName,
	}
	return &machine, nil
}

func getRegionOrEndpoint(params map[string]string, useDefault bool) string {
	regionOrEndpoint := params["endpoint"]
	if regionOrEndpoint == "" {
		regionOrEndpoint = params["region"]
		if regionOrEndpoint == "" && useDefault {
			regionOrEndpoint = defaultRegion
		}
	}
	return regionOrEndpoint
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strings"
)

func check(err error) {
	if err != nil {
		panic(err)
	}
}

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "usage: flatzip srcdir dstdir\n" +
			"flatzip creates a mtime-equivalent mirror of srcdir\n" +
			"with all files filled with zeros for easy compression.\n")
		flag.PrintDefaults()
	}

	flag.Parse()
	if flag.NArg() != 2 {
		flag.Usage()
		return
	}
	srcroot := flag.Arg(0)
	dstroot := flag.Arg(1)

	visit := func(path string, info os.FileInfo, err error) error {
		if err != nil {
			panic(err)
		}

		if !strings.HasPrefix(path, srcroot) {
			log.Fatalf("expected path %q to start with src %q", path, srcroot)
		}
		path = path[len(srcroot):]
		dst := filepath.Join(dstroot, path)

		if info.IsDir() {
			log.Printf("%s\n", path)
			check(os.MkdirAll(dst, 0777))
		} else {
			f, err := os.Create(dst)
			check(err)
			check(f.Truncate(info.Size()))
			check(f.Close())
		}
		check(os.Chtimes(dst, info.ModTime(), info.ModTime()))
		return nil
	}
	check(filepath.Walk(srcroot, visit))
}
<commit_msg>don't set mtime on directories (fails on win at least)<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strings"
)

func check(err error) {
	if err != nil {
		panic(err)
	}
}

func main() {
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "usage: flatzip srcdir dstdir\n" +
			"flatzip creates a mtime-equivalent mirror of srcdir\n" +
			"with all files filled with zeros for easy compression.\n")
		flag.PrintDefaults()
	}

	flag.Parse()
	if flag.NArg() != 2 {
		flag.Usage()
		return
	}
	srcroot := flag.Arg(0)
	dstroot := flag.Arg(1)

	visit := func(path string, info os.FileInfo, err error) error {
		if err != nil {
			panic(err)
		}

		if !strings.HasPrefix(path, srcroot) {
			log.Fatalf("expected path %q to start with src %q", path, srcroot)
		}
		path = path[len(srcroot):]
		dst := filepath.Join(dstroot, path)

		if info.IsDir() {
			log.Printf("%s\n", path)
			check(os.MkdirAll(dst, 0777))
		} else {
			f, err := os.Create(dst)
			check(err)
			check(f.Truncate(info.Size()))
			check(f.Close())
			check(os.Chtimes(dst, info.ModTime(), info.ModTime()))
		}
		return nil
	}
	check(filepath.Walk(srcroot, visit))
}
<|endoftext|>


<commit_before>/**
	* 1. Setup the server so cf can call it under main.
				e.g. `cf my-plugin` creates the callable server. now we can call the Run command
	* 2. Implement Run that is the actual code of the plugin!
	* 3. Return an error
**/

package main

import (
	"fmt"

	"github.com/cloudfoundry/cli/plugin"
)

type Test1 struct {
}

func (c *Test1) Run(cliConnection plugin.CliConnection, args []string) {
	if args[0] == "test_1_cmd1" {
		theFirstCmd()
	} else if args[0] == "test_1_cmd2" {
		theSecondCmd()
	}
}

func (c *Test1) GetMetadata() plugin.PluginMetadata {
	return plugin.PluginMetadata{
		Name: "Test1",
		Version: plugin.VersionType{
			Major: 1,
			Minor: 2,
			Build: 3,
		},
		Commands: []plugin.Command{
			{
				Name:     "test_1_cmd1",
				Alias:    "test_1_cmd1_alias",
				HelpText: "help text for test_1_cmd1",
				UsageDetails: plugin.Usage{
					Usage: "Test plugin command\n   cf test_1_cmd1 [-a] [-b] [--no-ouput]",
					Options: map[string]string{
						"-a":          "flag to do nothing",
						"-b":          "another flag to do nothing",
						"--no-output": "example option with no use",
					},
				},
			},
			{
				Name:     "test_1_cmd2",
				HelpText: "help text for test_1_cmd2",
			},
		},
	}
}

func theFirstCmd() {
	fmt.Println("You called cmd1 in test_1")
}

func theSecondCmd() {
	fmt.Println("You called cmd2 in test_1")
}

func main() {
	plugin.Start(new(Test1))
}
<commit_msg>remove '-' in test_1 plugin help sample<commit_after>/**
	* 1. Setup the server so cf can call it under main.
				e.g. `cf my-plugin` creates the callable server. now we can call the Run command
	* 2. Implement Run that is the actual code of the plugin!
	* 3. Return an error
**/

package main

import (
	"fmt"

	"github.com/cloudfoundry/cli/plugin"
)

type Test1 struct {
}

func (c *Test1) Run(cliConnection plugin.CliConnection, args []string) {
	if args[0] == "test_1_cmd1" {
		theFirstCmd()
	} else if args[0] == "test_1_cmd2" {
		theSecondCmd()
	}
}

func (c *Test1) GetMetadata() plugin.PluginMetadata {
	return plugin.PluginMetadata{
		Name: "Test1",
		Version: plugin.VersionType{
			Major: 1,
			Minor: 2,
			Build: 3,
		},
		Commands: []plugin.Command{
			{
				Name:     "test_1_cmd1",
				Alias:    "test_1_cmd1_alias",
				HelpText: "help text for test_1_cmd1",
				UsageDetails: plugin.Usage{
					Usage: "Test plugin command\n   cf test_1_cmd1 [-a] [-b] [--no-ouput]",
					Options: map[string]string{
						"a":         "flag to do nothing",
						"b":         "another flag to do nothing",
						"no-output": "example option with no use",
					},
				},
			},
			{
				Name:     "test_1_cmd2",
				HelpText: "help text for test_1_cmd2",
			},
		},
	}
}

func theFirstCmd() {
	fmt.Println("You called cmd1 in test_1")
}

func theSecondCmd() {
	fmt.Println("You called cmd2 in test_1")
}

func main() {
	plugin.Start(new(Test1))
}
<|endoftext|>


<commit_before>package graphs

import (
	"container/list"
	"errors"
)

var ErrNoDAG = errors.New("graphs: graph is not a DAG")

func initInEdges(g *Graph) map[Vertex]int {
	inEdges := make(map[Vertex]int)
	for e := range g.EdgesIter() {
		if _, ok := inEdges[e.Start]; !ok {
			inEdges[e.Start] = 0
		}
		if _, ok := inEdges[e.End]; ok {
			inEdges[e.End]++
		} else {
			inEdges[e.End] = 1
		}
	}
	return inEdges
}

func TopologicalSort(g *Graph) (topologicalOrder *list.List, topologicalClasses map[Vertex]int, err error) {
	inEdges := initInEdges(g)

	removeEdgesFromVertex := func(v Vertex) {
		for outEdge := range g.HalfedgesIter(v) {
			neighbor := outEdge.End
			inEdges[neighbor]--
		}
	}

	topologicalClasses = make(map[Vertex]int)
	topologicalOrder = list.New()
	tClass := 0
	for len(inEdges) > 0 {
		topClass := []Vertex{}
		for v, inDegree := range inEdges {
			if inDegree == 0 {
				topClass = append(topClass, v)
				topologicalClasses[v] = tClass
			}
		}
		if len(topClass) == 0 {
			err = ErrNoDAG
			topologicalClasses = make(map[Vertex]int)
			topologicalOrder = list.New()
			return
		}
		for _, v := range topClass {
			removeEdgesFromVertex(v)
			delete(inEdges, v)
			topologicalOrder.PushBack(v)
		}
		tClass++
	}

	return
}
<commit_msg>Inline function to generate inEdges map<commit_after>package graphs

import (
	"container/list"
	"errors"
)

var ErrNoDAG = errors.New("graphs: graph is not a DAG")

func TopologicalSort(g *Graph) (topologicalOrder *list.List, topologicalClasses map[Vertex]int, err error) {
	inEdges := make(map[Vertex]int)
	for e := range g.EdgesIter() {
		if _, ok := inEdges[e.Start]; !ok {
			inEdges[e.Start] = 0
		}
		if _, ok := inEdges[e.End]; ok {
			inEdges[e.End]++
		} else {
			inEdges[e.End] = 1
		}
	}

	removeEdgesFromVertex := func(v Vertex) {
		for outEdge := range g.HalfedgesIter(v) {
			neighbor := outEdge.End
			inEdges[neighbor]--
		}
	}

	topologicalClasses = make(map[Vertex]int)
	topologicalOrder = list.New()
	tClass := 0
	for len(inEdges) > 0 {
		topClass := []Vertex{}
		for v, inDegree := range inEdges {
			if inDegree == 0 {
				topClass = append(topClass, v)
				topologicalClasses[v] = tClass
			}
		}
		if len(topClass) == 0 {
			err = ErrNoDAG
			topologicalClasses = make(map[Vertex]int)
			topologicalOrder = list.New()
			return
		}
		for _, v := range topClass {
			removeEdgesFromVertex(v)
			delete(inEdges, v)
			topologicalOrder.PushBack(v)
		}
		tClass++
	}

	return
}
<|endoftext|>


<commit_before>// Copyright 2016 Adam H. Leventhal. All rights reserved.
// Licensed under the Apache License, version 2.0:
// http://www.apache.org/licenses/LICENSE-2.0

package main

import (
	"fmt"
	"math"
	"net"
	"net/http"
	//"os"
	//"os/signal"
	"reflect"
	"sync"
	"time"

	"github.com/zabawaba99/firego"

	"github.com/attic-labs/noms/go/spec"
	"github.com/attic-labs/noms/go/types"
)

type ItemIterator struct {
	lock       *sync.Mutex
	cv         *sync.Cond
	nextItem   uint32
	maxItem    uint32
	extraItems []uint32
	max        *firego.Firebase
	update     *firego.Firebase
}

func NewItemIterator() *ItemIterator {
	m := &sync.Mutex{}
	ii := &ItemIterator{m, sync.NewCond(m), 0, 0, nil, nil, nil}

	ii.max = firego.New("https://hacker-news.firebaseio.com/v0/maxitem", nil)
	var val float64
	if err := ii.max.Value(&val); err != nil {
		panic(err)
	}

	ii.maxItem = uint32(val)
	ii.nextItem = 1
	//ii.nextItem = ii.maxItem - 4000

	maxNotify := make(chan firego.Event)
	if err := ii.max.Watch(maxNotify); err != nil {
		panic(err)
	}

	go func() {
		for event := range maxNotify {
			ii.setMax(uint32(event.Data.(float64)))
		}
	}()

	ii.update = firego.New("https://hacker-news.firebaseio.com/v0/updates", nil)
	updateNotify := make(chan firego.Event)
	if err := ii.update.Watch(updateNotify); err != nil {
		panic(err)
	}

	go func() {
		for event := range updateNotify {
			for _, item := range event.Data.(map[string]interface{})["items"].([]interface{}) {
				ii.addExtra(uint32(item.(float64)))
			}
		}
	}()

	return ii
}

func (ii *ItemIterator) setMax(max uint32) {
	ii.lock.Lock()
	defer ii.lock.Unlock()
	ii.maxItem = max
	ii.cv.Broadcast()
}

func (ii *ItemIterator) addExtra(item uint32) {
	ii.lock.Lock()
	defer ii.lock.Unlock()
	if item < ii.nextItem {
		ii.extraItems = append(ii.extraItems, item)
		ii.cv.Broadcast()
	}
}

// Return the next element; block if none is available.
func (ii *ItemIterator) Next() (uint32, bool) {
	ii.lock.Lock()
	defer ii.lock.Unlock()

	for {
		if ll := len(ii.extraItems); ll != 0 {
			item := ii.extraItems[ll-1]
			ii.extraItems = ii.extraItems[:ll-1]
			return item, true
		}

		if ii.nextItem <= ii.maxItem {
			item := ii.nextItem
			ii.nextItem += 1
			return item, false
		}

		ii.cv.Wait()
	}
}

func (ii *ItemIterator) Max() uint32 {
	return ii.maxItem
}

type datum struct {
	id    uint32
	key   types.Number
	value types.Struct
	extra bool
}

func main() {
	fmt.Println("starting")

	ii := NewItemIterator()

	ds, err := spec.GetDataset("http://localhost:8000::hn")
	if err != nil {
		fmt.Println(err)
		return
	}
	defer ds.Database().Close()

	depth := 150
	newData := make(chan datum, depth)
	streamData := make(chan types.Value, depth)

	start := time.Now()

	newMap := types.NewStreamingMap(ds.Database(), streamData)

	// Grab the data from hacker news.
	for i := 0; i < 500; i++ {
		go churn(newData, ii)
	}

	var count uint32
	for {
		d := <-newData
		if !d.extra {
			count += 1
		}

		// XXX move status reporting out of here to a different go routine
		total := ii.Max()
		if count%1000 == 0 {
			dur := time.Since(start)
			eta := time.Duration(float64(dur) * float64(total-count) / float64(count))

			fmt.Printf("sent:  %d/%d (%d) time remaining: %s\n", count, total, d.id, eta)
		}

		streamData <- d.key
		streamData <- d.value

		// Too close for missiles; switching to guns. Close down the stream which will trigger the initial creation of the Map; exit the loop to wait for it to show up. Even though we try to drain newData, there may still be data left in newData, but we can pick that up later.
		//if count > total-uint32(depth)*2 {
		if d.id > total-uint32(depth)*2 {
			// if count > 10000 {
			done := false
			for !done {
				select {
				case d := <-newData:
					streamData <- d.key
					streamData <- d.value
				default:
					done = true
				}
			}
			close(streamData)
			break
		}
	}

	fmt.Println("generating map...")

	// Wait for the map to build; this could a little bit oftime so we need to figure out a plan for buffering data in the meantime. The most critical thing to stay on top of is going to be "extra" items in the iterator.
	mm := <-newMap

	getMap := make(chan types.Map)

	go func() {
		fmt.Println("first commit...")
		mp := mm
		for {
			nds, err := ds.CommitValue(mp)
			if err != nil {
				panic(err)
			}
			ds = nds
			mp = <-getMap
		}
	}()

	fmt.Println("processing...")

	// XXX make sure we're converging; look at the length of the extras array.
	// XXX better status reporting since this is the persistent mode of operation.
	// XXX figure out what I want to do if we come up and there's data already loaded.

	for {
		d := <-newData

		newCount := 0
		extraCount := 0
		if d.extra {
			extraCount += 1
		} else {
			newCount += 1
		}

		mm = mm.Set(d.key, d.value)

		last := time.Now()
		blocked := false
		for !blocked && time.Since(last) < 5*time.Second {
			select {
			case d = <-newData:
				if d.extra {
					extraCount += 1
				} else {
					newCount += 1
				}
				mm = mm.Set(d.key, d.value)
			default:
				blocked = true
			}
		}

		fmt.Printf("new: %d extra: %d   extras: %d\n", newCount, extraCount, len(ii.extraItems))

		// Poke our new map into the chan. If an old one is still in there, nudge it out and add our new one.
		select {
		case getMap <- mm:
		default:
			select {
			case _ = <-getMap:
			default:
			}
			getMap <- mm
		}
	}
}

func churn(newData chan<- datum, ii *ItemIterator) {
	var tr *http.Transport
	tr = &http.Transport{
		//DisableKeepAlives: true, // https://code.google.com/p/go/issues/detail?id=3514
		Dial: func(network, address string) (net.Conn, error) {
			start := time.Now()
			c, err := net.DialTimeout(network, address, firego.TimeoutDuration)
			tr.ResponseHeaderTimeout = firego.TimeoutDuration - time.Since(start)
			return c, err
		},
	}

	client := &http.Client{
		Transport: tr,
		//CheckRedirect: redirectPreserveHeaders,
	}

	for {
		id, extra := ii.Next()
		url := fmt.Sprintf("https://hacker-news.firebaseio.com/v0/item/%d", id)
		for {
			fb := firego.New(url, client)

			var val map[string]interface{}
			err := fb.Value(&val)
			if err != nil {
				continue
				//panic(err)
			}

			data := make(map[string]types.Value)
			for k, v := range val {
				switch vv := v.(type) {
				case string:
					data[k] = types.String(vv)
				case float64:
					data[k] = types.Number(vv)
				case bool:
					data[k] = types.Bool(vv)
				case []interface{}:
					ll := types.NewList()
					for _, elem := range vv {
						ll = ll.Append(types.Number(elem.(float64)))
					}
					data[k] = ll
				default:
					panic(reflect.TypeOf(v))
				}
			}

			name, ok := val["type"]
			if ok {
				st := types.NewStruct(name.(string), data)
				d := datum{
					id:    id,
					key:   data["id"].(types.Number),
					value: st,
					extra: extra,
				}

				select {
				case newData <- d:
				default:
					fmt.Println("blocked")
					newData <- d
				}
				break
			}
		}
	}
}

func Measure(subject func()) time.Duration {
	start := time.Now()
	subject()
	return time.Since(start)
}

type Distribution struct {
	counts [64]uint64
	sums   [64]uint64
	count  uint64
	sum    uint64
	sum2   uint64
}

func NewDistribution() *Distribution {
	return &Distribution{}
}

func (dist *Distribution) Add(value uint64) {
	dist.sum += value
	dist.sum2 += value * value
	dist.count += 1

	var hibit uint
	for hibit = 64; hibit > 0; hibit -= 1 {
		if value&(1<<(hibit-1)) != 0 {
			break
		}
	}

	dist.counts[hibit] += 1
	dist.sums[hibit] += value
}

func (dist *Distribution) Avg() float64 {
	return float64(dist.sum) / float64(dist.count)
}

func (dist *Distribution) StdDev() float64 {
	avg := float64(dist.sum) / float64(dist.count)
	avg2 := float64(dist.sum2) / float64(dist.count)

	return math.Sqrt(avg2 - avg*avg)
}

func (dist *Distribution) Hist() {
	var min, max uint

	for max = 63; max > 0; max -= 1 {
		if dist.counts[max] != 0 {
			break
		}
	}

	for min = 0; min < max; min += 1 {
		if dist.counts[min] != 0 {
			break
		}
	}

	for i := min; i <= max; i++ {
		fmt.Printf("%8d %d\n", 1<<i, dist.counts[i])
	}
}
<commit_msg>better reporting<commit_after>// Copyright 2016 Adam H. Leventhal. All rights reserved.
// Licensed under the Apache License, version 2.0:
// http://www.apache.org/licenses/LICENSE-2.0

package main

import (
	"fmt"
	"math"
	"net"
	"net/http"
	//"os"
	//"os/signal"
	"reflect"
	"sync"
	"time"

	"github.com/zabawaba99/firego"

	"github.com/attic-labs/noms/go/spec"
	"github.com/attic-labs/noms/go/types"
)

type ItemIterator struct {
	lock       *sync.Mutex
	cv         *sync.Cond
	nextItem   uint32
	maxItem    uint32
	extraItems []uint32
	max        *firego.Firebase
	update     *firego.Firebase
}

func NewItemIterator() *ItemIterator {
	m := &sync.Mutex{}
	ii := &ItemIterator{m, sync.NewCond(m), 0, 0, nil, nil, nil}

	ii.max = firego.New("https://hacker-news.firebaseio.com/v0/maxitem", nil)
	var val float64
	if err := ii.max.Value(&val); err != nil {
		panic(err)
	}

	ii.maxItem = uint32(val)
	ii.nextItem = 1
	//ii.nextItem = ii.maxItem - 100000

	maxNotify := make(chan firego.Event)
	if err := ii.max.Watch(maxNotify); err != nil {
		panic(err)
	}

	go func() {
		for event := range maxNotify {
			ii.setMax(uint32(event.Data.(float64)))
		}
	}()

	ii.update = firego.New("https://hacker-news.firebaseio.com/v0/updates", nil)
	updateNotify := make(chan firego.Event)
	if err := ii.update.Watch(updateNotify); err != nil {
		panic(err)
	}

	go func() {
		for event := range updateNotify {
			for _, item := range event.Data.(map[string]interface{})["items"].([]interface{}) {
				ii.addExtra(uint32(item.(float64)))
			}
		}
	}()

	return ii
}

func (ii *ItemIterator) setMax(max uint32) {
	ii.lock.Lock()
	defer ii.lock.Unlock()
	ii.maxItem = max
	ii.cv.Broadcast()
}

func (ii *ItemIterator) addExtra(item uint32) {
	ii.lock.Lock()
	defer ii.lock.Unlock()
	if item < ii.nextItem {
		ii.extraItems = append(ii.extraItems, item)
		ii.cv.Broadcast()
	}
}

// Return the next element; block if none is available.
func (ii *ItemIterator) Next() (uint32, bool) {
	ii.lock.Lock()
	defer ii.lock.Unlock()

	for {
		if ll := len(ii.extraItems); ll != 0 {
			item := ii.extraItems[ll-1]
			ii.extraItems = ii.extraItems[:ll-1]
			return item, true
		}

		if ii.nextItem <= ii.maxItem {
			item := ii.nextItem
			ii.nextItem += 1
			return item, false
		}

		ii.cv.Wait()
	}
}

func (ii *ItemIterator) Max() uint32 {
	return ii.maxItem
}

type datum struct {
	id    uint32
	key   types.Number
	value types.Struct
	extra bool
}

func main() {
	fmt.Println("starting")

	ii := NewItemIterator()

	ds, err := spec.GetDataset("http://localhost:8000::hn")
	if err != nil {
		fmt.Println(err)
		return
	}
	defer ds.Database().Close()

	depth := 150
	newData := make(chan datum, depth)
	streamData := make(chan types.Value, depth)

	start := time.Now()

	newMap := types.NewStreamingMap(ds.Database(), streamData)

	// Grab the data from hacker news.
	for i := 0; i < 500; i++ {
		go churn(newData, ii)
	}

	var count uint32
	var newCount uint32
	var extraCount uint32
	for {
		d := <-newData

		count += 1
		if d.extra {
			extraCount += 1
		} else {
			newCount += 1
		}

		// XXX move status reporting out of here to a different go routine
		total := ii.Max()
		if count%1000 == 0 {
			dur := time.Since(start)
			eta := time.Duration(float64(dur) * float64(total-count) / float64(count))

			fmt.Printf("sent:  %d/%d (%d) (%d) time remaining: %s\n", newCount, total, d.id, extraCount, eta)
		}

		streamData <- d.key
		streamData <- d.value

		// Too close for missiles; switching to guns. Close down the stream which will trigger the initial creation of the Map; exit the loop to wait for it to show up. Even though we try to drain newData, there may still be data left in newData, but we can pick that up later.
		//if count > total-uint32(depth)*2 {
		if d.id > total-uint32(depth)*2 {
			// if count > 10000 {
			done := false
			for !done {
				select {
				case d := <-newData:
					streamData <- d.key
					streamData <- d.value
				default:
					done = true
				}
			}
			close(streamData)
			break
		}
	}

	fmt.Println("generating map...")

	// Wait for the map to build; this could a little bit oftime so we need to figure out a plan for buffering data in the meantime. The most critical thing to stay on top of is going to be "extra" items in the iterator.
	mm := <-newMap

	getMap := make(chan types.Map)

	go func() {
		fmt.Println("first commit...")
		mp := mm
		for {
			nds, err := ds.CommitValue(mp)
			if err != nil {
				panic(err)
			}
			ds = nds
			mp = <-getMap
		}
	}()

	fmt.Println("processing...")

	// XXX better status reporting since this is the persistent mode of operation.
	// XXX figure out what I want to do if we come up and there's data already loaded.

	for {
		d := <-newData

		newCount := 0
		extraCount := 0
		if d.extra {
			extraCount += 1
		} else {
			newCount += 1
		}

		mm = mm.Set(d.key, d.value)

		last := time.Now()
		blocked := false
		for !blocked && time.Since(last) < 5*time.Second {
			select {
			case d = <-newData:
				if d.extra {
					extraCount += 1
				} else {
					newCount += 1
				}
				mm = mm.Set(d.key, d.value)
			default:
				blocked = true
			}
		}

		fmt.Printf("new: %d extra: %d   extras: %d\n", newCount, extraCount, len(ii.extraItems))

		// Poke our new map into the chan. If an old one is still in there, nudge it out and add our new one.
		select {
		case getMap <- mm:
		default:
			select {
			case _ = <-getMap:
			default:
			}
			getMap <- mm
		}
	}
}

func churn(newData chan<- datum, ii *ItemIterator) {
	var tr *http.Transport
	tr = &http.Transport{
		Dial: func(network, address string) (net.Conn, error) {
			return net.DialTimeout(network, address, 5*time.Second)
		},
		TLSHandshakeTimeout:   5 * time.Second,
		ResponseHeaderTimeout: 5 * time.Second,
	}

	client := &http.Client{
		Transport: tr,
		Timeout:   time.Second * 5,
	}

	for {
		id, extra := ii.Next()
		url := fmt.Sprintf("https://hacker-news.firebaseio.com/v0/item/%d", id)
		for {
			fb := firego.New(url, client)

			var val map[string]interface{}
			err := fb.Value(&val)
			if err != nil {
				continue
				//panic(err)
			}

			data := make(map[string]types.Value)
			for k, v := range val {
				switch vv := v.(type) {
				case string:
					data[k] = types.String(vv)
				case float64:
					data[k] = types.Number(vv)
				case bool:
					data[k] = types.Bool(vv)
				case []interface{}:
					ll := types.NewList()
					for _, elem := range vv {
						ll = ll.Append(types.Number(elem.(float64)))
					}
					data[k] = ll
				default:
					panic(reflect.TypeOf(v))
				}
			}

			name, ok := val["type"]
			if ok {
				st := types.NewStruct(name.(string), data)
				d := datum{
					id:    id,
					key:   data["id"].(types.Number),
					value: st,
					extra: extra,
				}

				select {
				case newData <- d:
				default:
					fmt.Println("blocked")
					newData <- d
				}
				break
			}
		}
	}
}

func Measure(subject func()) time.Duration {
	start := time.Now()
	subject()
	return time.Since(start)
}

type Distribution struct {
	counts [64]uint64
	sums   [64]uint64
	count  uint64
	sum    uint64
	sum2   uint64
}

func NewDistribution() *Distribution {
	return &Distribution{}
}

func (dist *Distribution) Add(value uint64) {
	dist.sum += value
	dist.sum2 += value * value
	dist.count += 1

	var hibit uint
	for hibit = 64; hibit > 0; hibit -= 1 {
		if value&(1<<(hibit-1)) != 0 {
			break
		}
	}

	dist.counts[hibit] += 1
	dist.sums[hibit] += value
}

func (dist *Distribution) Avg() float64 {
	return float64(dist.sum) / float64(dist.count)
}

func (dist *Distribution) StdDev() float64 {
	avg := float64(dist.sum) / float64(dist.count)
	avg2 := float64(dist.sum2) / float64(dist.count)

	return math.Sqrt(avg2 - avg*avg)
}

func (dist *Distribution) Hist() {
	var min, max uint

	for max = 63; max > 0; max -= 1 {
		if dist.counts[max] != 0 {
			break
		}
	}

	for min = 0; min < max; min += 1 {
		if dist.counts[min] != 0 {
			break
		}
	}

	for i := min; i <= max; i++ {
		fmt.Printf("%8d %d\n", 1<<i, dist.counts[i])
	}
}
<|endoftext|>


<commit_before>package tags

import (
	"strings"

	"github.com/chrisolsen/ae"
	"golang.org/x/net/context"
	"google.golang.org/appengine/datastore"
)

const tableName = "tags"

// Tag allows better searching capabilities with AppEngine's Datastore
type Tag struct {
	model.Base
	Value  string
	Type   string
	Public bool
}

// Save saves the tag items with the parent relation
func Save(c context.Context, rawTags []string, tagType string, public bool, parentKey *datastore.Key) (int, int, error) {
	var delCount int
	var addCount int

	// get existing
	var existingTags []*Tag
	oldKeys, err := datastore.NewQuery(tableName).
		Ancestor(parentKey).
		Filter("Type =", tagType).
		Filter("Public =", public).
		GetAll(c, &existingTags)
	if err != nil {
		return 0, 0, err
	}
	for i, key := range oldKeys {
		existingTags[i].Key = key
	}

	// delete old tags
	newTagMap := make(map[string]bool)
	for _, tag := range rawTags {
		newTagMap[strings.ToLower(tag)] = true
	}

	rmTagKeys := getTagsToRemove(existingTags, newTagMap)
	err = datastore.DeleteMulti(c, rmTagKeys)
	if err != nil {
		return 0, 0, err
	}
	delCount = len(rmTagKeys)

	// insert new tags
	existingTagMap := make(map[string]bool)
	for _, tag := range existingTags {
		existingTagMap[tag.Value] = true
	}
	newTags := getTagsToAdd(rawTags, existingTagMap)
	for _, t := range newTags {
		tg := &Tag{Value: strings.ToLower(t), Type: tagType, Public: public}
		ky := datastore.NewIncompleteKey(c, tableName, parentKey)
		_, err = datastore.Put(c, ky, tg)
		if err != nil {
			return 0, 0, err
		}
	}
	addCount = len(newTags)
	return delCount, addCount, nil
}

// FindKeysByTag returns a list of all the tag's parent datastore keys
func FindKeysByTag(c context.Context, tag, tagType string, parentKey *datastore.Key, offset, limit int) ([]*datastore.Key, error) {
	filter := strings.ToLower(tag)
	q := datastore.NewQuery(tableName)
	if parentKey != nil {
		q = q.Ancestor(parentKey)
	}
	keys, err := q.
		Filter("Type =", tagType).
		Filter("Value =", filter).
		Offset(offset).
		Limit(limit).
		KeysOnly().
		GetAll(c, nil)

	if err != nil {
		return nil, err
	}
	parentKeys := make([]*datastore.Key, len(keys))
	for i, k := range keys {
		parentKeys[i] = k.Parent()
	}
	return parentKeys, err
}

func getTagsToRemove(currentTags []*Tag, newTags map[string]bool) []*datastore.Key {
	var list []*datastore.Key
	for _, t := range currentTags {
		if _, ok := newTags[strings.ToLower(t.Value)]; !ok {
			list = append(list, t.Key)
		}
	}
	return list
}

func getTagsToAdd(rawTags []string, existingTags map[string]bool) []string {
	var list []string
	for _, t := range rawTags {
		lt := strings.ToLower(t)
		if _, ok := existingTags[lt]; !ok {
			list = append(list, lt)
		}
	}
	return list
}
<commit_msg>tags updates<commit_after>package tags

import (
	"strings"

	"github.com/chrisolsen/ae"
	"golang.org/x/net/context"
	"google.golang.org/appengine/datastore"
)

const tableName = "tags"

// Tag allows better searching capabilities with AppEngine's Datastore
type Tag struct {
	ae.Model
	Value string
	Type  string
}

// Load toy
func (t *Tag) Load(ps []datastore.Property) error {
	if err := datastore.LoadStruct(t, ps); err != nil {
		switch err.(type) {
		case *datastore.ErrFieldMismatch:
		default:
			return err
		}
	}
	return nil
}

// Save toy
func (t *Tag) Save() ([]datastore.Property, error) {
	return datastore.SaveStruct(t)
}

// Save saves the tag items with the parent relation
func Save(c context.Context, rawTags []string, tagType string, parentKey *datastore.Key) (int, int, error) {
	var delCount int
	var addCount int

	// get existing
	var existingTags []*Tag
	oldKeys, err := datastore.NewQuery(tableName).
		Ancestor(parentKey).
		Filter("Type =", tagType).
		GetAll(c, &existingTags)
	if err != nil {
		return 0, 0, err
	}
	for i, key := range oldKeys {
		existingTags[i].Key = key
	}

	// delete old tags
	newTagMap := make(map[string]bool)
	for _, tag := range rawTags {
		newTagMap[strings.ToLower(tag)] = true
	}

	rmTagKeys := getTagsToRemove(existingTags, newTagMap)
	err = datastore.DeleteMulti(c, rmTagKeys)
	if err != nil {
		return 0, 0, err
	}
	delCount = len(rmTagKeys)

	// insert new tags
	existingTagMap := make(map[string]bool)
	for _, tag := range existingTags {
		existingTagMap[tag.Value] = true
	}
	newTags := getTagsToAdd(rawTags, existingTagMap)
	for _, t := range newTags {
		tg := &Tag{Value: strings.ToLower(t), Type: tagType}
		ky := datastore.NewIncompleteKey(c, tableName, parentKey)
		_, err = datastore.Put(c, ky, tg)
		if err != nil {
			return 0, 0, err
		}
	}
	addCount = len(newTags)
	return delCount, addCount, nil
}

// FindKeysByTag returns a list of all the tag's parent datastore keys
func FindKeysByTag(c context.Context, tag, tagType string, parentKey *datastore.Key, offset, limit int) ([]*datastore.Key, error) {
	filter := strings.ToLower(tag)
	q := datastore.NewQuery(tableName)
	if parentKey != nil {
		q = q.Ancestor(parentKey)
	}

	keys, err := q.
		Filter("Type =", tagType).
		Filter("Value =", filter).
		Order("__key__").
		Offset(offset).
		Limit(limit).
		KeysOnly().
		GetAll(c, nil)

	if err != nil {
		return nil, err
	}
	parentKeys := make([]*datastore.Key, len(keys))
	for i, k := range keys {
		parentKeys[i] = k.Parent()
	}
	return parentKeys, err
}

func getTagsToRemove(currentTags []*Tag, newTags map[string]bool) []*datastore.Key {
	var list []*datastore.Key
	for _, t := range currentTags {
		if _, ok := newTags[strings.ToLower(t.Value)]; !ok {
			list = append(list, t.Key)
		}
	}
	return list
}

func getTagsToAdd(rawTags []string, existingTags map[string]bool) []string {
	var list []string
	for _, t := range rawTags {
		lt := strings.ToLower(t)
		if _, ok := existingTags[lt]; !ok {
			list = append(list, lt)
		}
	}
	return list
}
<|endoftext|>


<commit_before>package aws

import (
	"fmt"
	"reflect"
	"testing"

	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
	"github.com/mitchellh/goamz/ec2"
)

func TestDiffTags(t *testing.T) {
	cases := []struct {
		Old, New       map[string]interface{}
		Create, Remove map[string]string
	}{
		// Basic add/remove
		{
			Old: map[string]interface{}{
				"foo": "bar",
			},
			New: map[string]interface{}{
				"bar": "baz",
			},
			Create: map[string]string{
				"bar": "baz",
			},
			Remove: map[string]string{
				"foo": "bar",
			},
		},

		// Modify
		{
			Old: map[string]interface{}{
				"foo": "bar",
			},
			New: map[string]interface{}{
				"foo": "baz",
			},
			Create: map[string]string{
				"foo": "baz",
			},
			Remove: map[string]string{
				"foo": "bar",
			},
		},
	}

	for i, tc := range cases {
		c, r := diffTags(tagsFromMap(tc.Old), tagsFromMap(tc.New))
		cm := tagsToMap(c)
		rm := tagsToMap(r)
		if !reflect.DeepEqual(cm, tc.Create) {
			t.Fatalf("%i: bad create: %#v", i, cm)
		}
		if !reflect.DeepEqual(rm, tc.Remove) {
			t.Fatalf("%i: bad remove: %#v", i, rm)
		}
	}
}

// testAccCheckTags can be used to check the tags on a resource.
func testAccCheckTags(
	ts *[]ec2.Tag, key string, value string) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		m := tagsToMap(*ts)
		v, ok := m[key]
		if value != "" && !ok {
			return fmt.Errorf("Missing tag: %s", key)
		} else if value == "" && ok {
			return fmt.Errorf("Extra tag: %s", key)
		}
		if value == "" {
			return nil
		}

		if v != value {
			return fmt.Errorf("%s: bad value: %s", key, v)
		}

		return nil
	}
}
<commit_msg>argument and verb formatting fixes reported by go vet<commit_after>package aws

import (
	"fmt"
	"reflect"
	"testing"

	"github.com/hashicorp/terraform/helper/resource"
	"github.com/hashicorp/terraform/terraform"
	"github.com/mitchellh/goamz/ec2"
)

func TestDiffTags(t *testing.T) {
	cases := []struct {
		Old, New       map[string]interface{}
		Create, Remove map[string]string
	}{
		// Basic add/remove
		{
			Old: map[string]interface{}{
				"foo": "bar",
			},
			New: map[string]interface{}{
				"bar": "baz",
			},
			Create: map[string]string{
				"bar": "baz",
			},
			Remove: map[string]string{
				"foo": "bar",
			},
		},

		// Modify
		{
			Old: map[string]interface{}{
				"foo": "bar",
			},
			New: map[string]interface{}{
				"foo": "baz",
			},
			Create: map[string]string{
				"foo": "baz",
			},
			Remove: map[string]string{
				"foo": "bar",
			},
		},
	}

	for i, tc := range cases {
		c, r := diffTags(tagsFromMap(tc.Old), tagsFromMap(tc.New))
		cm := tagsToMap(c)
		rm := tagsToMap(r)
		if !reflect.DeepEqual(cm, tc.Create) {
			t.Fatalf("%d: bad create: %#v", i, cm)
		}
		if !reflect.DeepEqual(rm, tc.Remove) {
			t.Fatalf("%d: bad remove: %#v", i, rm)
		}
	}
}

// testAccCheckTags can be used to check the tags on a resource.
func testAccCheckTags(
	ts *[]ec2.Tag, key string, value string) resource.TestCheckFunc {
	return func(s *terraform.State) error {
		m := tagsToMap(*ts)
		v, ok := m[key]
		if value != "" && !ok {
			return fmt.Errorf("Missing tag: %s", key)
		} else if value == "" && ok {
			return fmt.Errorf("Extra tag: %s", key)
		}
		if value == "" {
			return nil
		}

		if v != value {
			return fmt.Errorf("%s: bad value: %s", key, v)
		}

		return nil
	}
}
<|endoftext|>


<commit_before>//Package godnsbl provides a way to query if an IP is on any dnsbl blacklist
package godnsbl

import (
  "sync"
  "strings"
  "fmt"
  "net"
  "time"
)

var blacklistDomains = []string{
  "b.barracudacentral.org",
  "bl.spamcop.net",
  "dnsrbl.org",
  "babl.rbl.webiron.net",
  "cabl.rbl.webiron.net",
  "stabl.rbl.webiron.net",
  "crawler.rbl.webiron.net",
  "bad.psky.me",
  "http.dnsbl.sorbs.net",
  "socks.dnsbl.sorbs.net",
  "misc.dnsbl.sorbs.net",
  "smtp.dnsbl.sorbs.net",
  "web.dnsbl.sorbs.net",
  "spam.dnsbl.sorbs.net",
  "escalations.dnsbl.sorbs.net",
  "zombie.dnsbl.sorbs.net",
  "recent.spam.dnsbl.sorbs.net",
  "dyna.spamrats.com",
  "spam.abuse.ch",
  "rbl.interserver.net",
  "tor.dan.me.uk",
  "torexit.dan.me.uk",
  "dnsbl-1.uceprotect.net",
  "dnsbl-2.uceprotect.net",
  "dnsbl-3.uceprotect.net",
  "cbl.abuseat.org",
  "spam.spamrats.com",
  "ips.backscatterer.org",
  "truncate.gbudb.net",
  "psbl.surriel.com",
  "db.wpbl.info",
  "bl.spamcannibal.org",
  "dnsbl.inps.de",
  "bl.blocklist.de",
  "rbl.megarbl.net",
  "all.s5h.net",
  "srnblack.surgate.net",
}

// A dnsblData contains all the data about a particular dnsbl list
type dnsblData struct {
  Status string
  Msg string
  Listed bool
  Name string
  RespTime int64
  Resp string
}

// dnsblReturn contains a list of all the checked lists and if its on any lists
type dnsblReturn struct{
  Listed bool
  Count int
  Total int
  Responses []dnsblData
}

func getDnsblResp(wg *sync.WaitGroup, resp *dnsblReturn, ip string, domain string) {
  // Sweet timing function to see how long it takes to get a response
  st := time.Now()

  // Get the host for the format <reverse ip>.domain
  ips, err := net.LookupIP(fmt.Sprintf("%s.%s", ip, domain))

  // See how long it took the request
  et := time.Since(st).Nanoseconds()/1000000

  // Default values to insert
  listed := false
  respIPs := ""
  respStat := "ok"
  respErrMsg := ""

  // If it gets an IP back, it is listed in the database. Only get the first IP per RFC
  if len(ips) > 0 {
    listed = true
    respIPs = ips[0].String()
  }

  // Only show the error if it isn't the generic no domain when the listing isnt listed
  if err != nil && err.Error() != fmt.Sprintf("lookup %s.%s: no such host", ip, domain) && err.Error() != fmt.Sprintf("lookup %s.%s: No address associated with hostname", ip, domain){
    respStat = "error"
    respErrMsg = err.Error()
  }

  // Could of used channels but this works for now
  resp.Responses = append(resp.Responses, dnsblData{respStat, respErrMsg, listed, domain, et, respIPs})

  // Bows
  wg.Done()
}

func CheckBlacklist(ip string) (dnsblReturn){
  var wg sync.WaitGroup
  wg.Add(len(blacklistDomains))
  var resp dnsblReturn

  // Flip the IP
  flipIP := getFlipIP(ip)

  // Makes a go routine for every domain and does the check
  for _, dom := range blacklistDomains {
    go getDnsblResp(&wg, &resp, flipIP, dom)
  }

  // They do take a while to return all
  wg.Wait()

  // Loop through each reply and see if the IP is on any blacklists
  for _, check := range resp.Responses{
    if check.Listed == true{
      resp.Listed = true
      resp.Count += 1
    }
  }
  resp.Total = len(blacklistDomains)

  // We are finished!
  return resp

}

func getFlipIP(ip string) (string){
  // Splits the IP by the period
  brokenIP := strings.Split(ip, ".")
  // Flips the IP to be in proper RFC format
  return fmt.Sprintf("%s.%s.%s.%s", brokenIP[3], brokenIP[2], brokenIP[1], brokenIP[0])
}
<commit_msg>Travis-Ci Bump<commit_after>//Package godnsbl provides a way to query if an IP is on any dnsbl blacklist
package godnsbl

import (
  "sync"
  "strings"
  "fmt"
  "net"
  "time"
)

var blacklistDomains = []string{
  "b.barracudacentral.org",
  "bl.spamcop.net",
  "dnsrbl.org",
  "babl.rbl.webiron.net",
  "cabl.rbl.webiron.net",
  "stabl.rbl.webiron.net",
  "crawler.rbl.webiron.net",
  "bad.psky.me",
  "http.dnsbl.sorbs.net",
  "socks.dnsbl.sorbs.net",
  "misc.dnsbl.sorbs.net",
  "smtp.dnsbl.sorbs.net",
  "web.dnsbl.sorbs.net",
  "spam.dnsbl.sorbs.net",
  "escalations.dnsbl.sorbs.net",
  "zombie.dnsbl.sorbs.net",
  "recent.spam.dnsbl.sorbs.net",
  "dyna.spamrats.com",
  "spam.abuse.ch",
  "rbl.interserver.net",
  "tor.dan.me.uk",
  "torexit.dan.me.uk",
  "dnsbl-1.uceprotect.net",
  "dnsbl-2.uceprotect.net",
  "dnsbl-3.uceprotect.net",
  "cbl.abuseat.org",
  "spam.spamrats.com",
  "ips.backscatterer.org",
  "truncate.gbudb.net",
  "psbl.surriel.com",
  "db.wpbl.info",
  "bl.spamcannibal.org",
  "dnsbl.inps.de",
  "bl.blocklist.de",
  "rbl.megarbl.net",
  "all.s5h.net",
  "srnblack.surgate.net",
}

// A dnsblData contains all the data about a particular dnsbl list
type dnsblData struct {
  Status string
  Msg string
  Listed bool
  Name string
  RespTime int64
  Resp string
}

// dnsblReturn contains a list of all the checked lists and if its on any lists
type dnsblReturn struct{
  Listed bool
  Count int
  Total int
  Responses []dnsblData
}

func getDnsblResp(wg *sync.WaitGroup, resp *dnsblReturn, ip string, domain string) {
  // Sweet timing function to see how long it takes to get a response
  st := time.Now()

  // Get the host for the format <reverse ip>.domain
  ips, err := net.LookupIP(fmt.Sprintf("%s.%s", ip, domain))

  // See how long it took the request
  et := time.Since(st).Nanoseconds()/1000000

  // Default values to insert
  listed := false
  respIPs := ""
  respStat := "ok"
  respErrMsg := ""

  // If it gets an IP back, it is listed in the database. Only get the first IP per RFC
  if len(ips) > 0 {
    listed = true
    respIPs = ips[0].String()
  }

  // Only show the error if it isn't the generic no domain when the listing isnt listed
  if err != nil && err.Error() != fmt.Sprintf("lookup %s.%s: no such host", ip, domain) && err.Error() != fmt.Sprintf("lookup %s.%s: No address associated with hostname", ip, domain){
    respStat = "error"
    respErrMsg = err.Error()
  }

  // Could of used channels but this works for now
  resp.Responses = append(resp.Responses, dnsblData{respStat, respErrMsg, listed, domain, et, respIPs})

  // Bows
  wg.Done()
}

func CheckBlacklist(ip string) (dnsblReturn){
  var wg sync.WaitGroup
  wg.Add(len(blacklistDomains))
  var resp dnsblReturn

  // Flip the IP
  flipIP := getFlipIP(ip)

  // Makes a go routine for every domain and does the check
  for _, dom := range blacklistDomains {
    go getDnsblResp(&wg, &resp, flipIP, dom)
  }

  // They do take a while to return all
  wg.Wait()

  // Loop through each reply and see if the IP is on any blacklists
  for _, check := range resp.Responses{
    if check.Listed == true{
      resp.Listed = true
      resp.Count += 1
    }
  }
  resp.Total = len(blacklistDomains)

  // We are finished!
  return resp
}

func getFlipIP(ip string) (string){
  // Splits the IP by the period
  brokenIP := strings.Split(ip, ".")
  // Flips the IP to be in proper RFC format
  return fmt.Sprintf("%s.%s.%s.%s", brokenIP[3], brokenIP[2], brokenIP[1], brokenIP[0])
}
<|endoftext|>


<commit_before>package gohasha

import (
	"crypto/md5"
	"encoding/hex"
	"errors"
	"hash/crc64"
	"io/ioutil"
	"io"
	"bytes"
)

type GohashaOptions struct {
	//path to file for hashing
	Filepath string
	//Data from hashing
	Data string
	//Reader buffer
	BufferReader io.Reader
	//Now, algorithms can be md5 or crc32
	Algorithm string
}

//GoHasha return hash in string from string data or file
func GoHasha(opt *GohashaOptions) (string, error) {
	if opt.Data != "" {
		return crypt(opt.Data, opt.Algorithm), nil
	}
	if opt.Filepath != "" {
		data, err := ioutil.ReadFile(opt.Filepath)
		if err != nil {
			return "", err
		} else {
			return crypt(string(data), opt.Algorithm), nil
		}
	}

	if opt.BufferReader != nil {
		buff := new(bytes.Buffer)
		buff.ReadFrom(opt.Buffer)
		return crypt(buff.String(), opt.Algorithm), nil
	}
	return "", errors.New("Can't find data for hashing")
}

func crypt(data string, algorithm string) string {
	if algorithm == "crc64" {
		return crc64run(data)
	} else {
		return md5run(data)
	}
}

func md5run(data string) string {
	hasher := md5.New()
	hasher.Write([]byte(data))
	return hex.EncodeToString(hasher.Sum(nil))
}

func crc64run(data string) string {
	hasher := crc64.New(crc64.MakeTable(crc64.ISO))
	hasher.Write([]byte(data))
	return hex.EncodeToString(hasher.Sum(nil))
}
<commit_msg>Fix to bufferreader<commit_after>package gohasha

import (
	"crypto/md5"
	"encoding/hex"
	"errors"
	"hash/crc64"
	"io/ioutil"
	"io"
	"bytes"
)

type GohashaOptions struct {
	//path to file for hashing
	Filepath string
	//Data from hashing
	Data string
	//Reader buffer
	BufferReader io.Reader

	//Now, algorithms can be md5 or crc32
	Algorithm string
}

//GoHasha return hash in string from string data or file
func GoHasha(opt *GohashaOptions) (string, error) {
	if opt.Data != "" {
		return crypt(opt.Data, opt.Algorithm), nil
	}
	if opt.Filepath != "" {
		data, err := ioutil.ReadFile(opt.Filepath)
		if err != nil {
			return "", err
		} else {
			return crypt(string(data), opt.Algorithm), nil
		}
	}

	if opt.BufferReader != nil {
		buff := new(bytes.Buffer)
		buff.ReadFrom(opt.BufferReader)
		return crypt(buff.String(), opt.Algorithm), nil
	}
	return "", errors.New("Can't find data for hashing")
}

func crypt(data string, algorithm string) string {
	if algorithm == "crc64" {
		return crc64run(data)
	} else {
		return md5run(data)
	}
}

func md5run(data string) string {
	hasher := md5.New()
	hasher.Write([]byte(data))
	return hex.EncodeToString(hasher.Sum(nil))
}

func crc64run(data string) string {
	hasher := crc64.New(crc64.MakeTable(crc64.ISO))
	hasher.Write([]byte(data))
	return hex.EncodeToString(hasher.Sum(nil))
}
<|endoftext|>


<commit_before>// Copyright 2014 Matt Martz <matt@sivel.net>
// 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 main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"path"
	"path/filepath"
	"strings"
	"sync"
)

type Walker struct {
	Paths []string
}

// Walk is the filepath.Walk WalkFunc to handle recording visited paths
func (w *Walker) Walk(path string, info os.FileInfo, err error) error {
	if !info.IsDir() {
		w.Paths = append(w.Paths, path)
	}
	return nil
}

type AuthContainer struct {
	Auth Auth `json:"auth"`
}

type Auth struct {
	PasswordCredentials *PasswordCredentials `json:"passwordCredentials,omitempty"`
	ApiKeyCredentials   *ApiKeyCredentials   `json:"RAX-KSKEY:apiKeyCredentials,omitempty"`
	TenantId            string               `json:"tenantId,omitempty"`
	TenantName          string               `json:"tenantName,omitempty"`
}

type PasswordCredentials struct {
	Username string `json:"username"`
	Password string `json:"password"`
}

type ApiKeyCredentials struct {
	Username string `json:"username"`
	ApiKey   string `json:"apiKey"`
}

type Tokens struct {
	Access Access `json:"access"`
}

type Access struct {
	Token          Token `json:"token"`
	ServiceCatalog []CatalogEntry
}

type CatalogEntry struct {
	Name, Type string
	Endpoints  []EntryEndpoint
}

type EntryEndpoint struct {
	Region, TenantId                    string
	PublicURL, InternalURL              string
	VersionId, VersionInfo, VersionList string
}

type Token struct {
	Id      string `json:"id"`
	Expires string `json:"expires"`
}

type CloudFiles struct {
	Username  string
	ApiKey    string
	Token     string
	Endpoint  string
	Container string
	Region    string
}

// Auth will authenticate to CloudFiles
func (c *CloudFiles) Auth() {
	var tokens *Tokens

	auth := &AuthContainer{
		Auth: Auth{
			ApiKeyCredentials: &ApiKeyCredentials{
				Username: c.Username,
				ApiKey:   c.ApiKey,
			},
		},
	}
	body, _ := json.Marshal(auth)

	res, err := http.Post("https://identity.api.rackspacecloud.com/v2.0/tokens", "application/json", bytes.NewBuffer(body))
	defer res.Body.Close()
	if res.StatusCode != 200 || err != nil {
		log.Fatal("Unable to authenticate")
	}

	resBody, _ := ioutil.ReadAll(res.Body)

	json.Unmarshal(resBody, &tokens)

	for _, service := range tokens.Access.ServiceCatalog {
		if service.Type == "object-store" {
			for _, endpoint := range service.Endpoints {
				if endpoint.Region == c.Region {
					c.Endpoint = endpoint.PublicURL
					break
				}
			}
			break
		}
	}

	if len(c.Endpoint) == 0 {
		log.Fatal(fmt.Sprintf("No PublicURL found for object-store in region %s", c.Region))
	}

	c.Token = tokens.Access.Token.Id
}

func (c *CloudFiles) ListObjects() []string {
	var objects []string
	var resBody []byte
	req, _ := http.NewRequest("GET", fmt.Sprintf("%s/%s", c.Endpoint, c.Container), nil)
	req.Header.Set("X-Auth-Token", c.Token)
	req.Header.Set("Accept", "text/plain")
	client := &http.Client{}
	res, err := client.Do(req)
	if err != nil {
		log.Fatal(err)
	} else if res.StatusCode != 200 {
		return objects
	}
	defer res.Body.Close()
	resBody, _ = ioutil.ReadAll(res.Body)
	objects = strings.Split(strings.TrimSpace(string(resBody)), "\n")
	for len(resBody) > 0 {
		req.URL, _ = url.ParseRequestURI(fmt.Sprintf("%s/%s?marker=%s", c.Endpoint, c.Container, objects[len(objects)-1]))
		res, err := client.Do(req)
		if err != nil {
			log.Fatal(err)
		}
		resBody, _ = ioutil.ReadAll(res.Body)
		objects = append(objects, strings.Split(strings.TrimSpace(string(resBody)), "\n")...)
		res.Body.Close()
	}
	return objects
}

// CreateContainer ensures that a container exists
func (c *CloudFiles) CreateContainer() {
	req, _ := http.NewRequest("PUT", fmt.Sprintf("%s/%s", c.Endpoint, c.Container), bytes.NewBuffer([]byte{}))
	req.Header.Set("X-Auth-Token", c.Token)
	client := &http.Client{}
	res, _ := client.Do(req)
	fmt.Println(res.StatusCode)
}

// Upload is a goroutine that uploads files provided by a channel to a CloudFiles container
func (c *CloudFiles) Upload(thread int, ci chan string, wg *sync.WaitGroup, BasePath string) {
	defer wg.Done()

	client := &http.Client{}
	req, _ := http.NewRequest("PUT", "", nil)
	req.Header.Set("X-Auth-Token", c.Token)

	for path := range ci {
		ObjPath := strings.TrimPrefix(strings.Replace(path, BasePath, "", 1), "/")
		fmt.Printf("Thread %03d: uploading %s\n", thread, ObjPath)
		file, err := os.Open(path)
		if err != nil {
			log.Print(fmt.Printf("%s\n", err))
			continue
		}
		req.URL, _ = url.ParseRequestURI(fmt.Sprintf("%s/%s/%s", c.Endpoint, c.Container, ObjPath))
		req.Body = file
		req.ContentLength = 0
		res, err := client.Do(req)
		if err != nil {
			log.Print(fmt.Printf("%s\n", err))
			continue
		}
		file.Close()
		res.Body.Close()
		fmt.Printf("Thread %03d: upload complete for %s\n", thread, ObjPath)
	}
	fmt.Printf("Thread %03d: exiting\n", thread)
}

func (c *CloudFiles) Delete(thread int, ci chan string, wg *sync.WaitGroup) {
	defer wg.Done()

	client := &http.Client{}
	req, _ := http.NewRequest("DELETE", "", nil)
	req.Header.Set("X-Auth-Token", c.Token)

	for path := range ci {
		fmt.Printf("Thread %03d: deleting %s\n", thread, path)
		req.URL, _ = url.ParseRequestURI(fmt.Sprintf("%s/%s/%s", c.Endpoint, c.Container, path))
		res, err := client.Do(req)
		if err != nil {
			log.Print(fmt.Printf("%s\n", err))
			continue
		}
		res.Body.Close()
		fmt.Printf("Thread %03d: delete complete for %s\n", thread, path)
	}
	fmt.Printf("Thread %03d: exiting\n", thread)
}

func (c *CloudFiles) Download(thread int, ci chan string, wg *sync.WaitGroup, BasePath string) {
	defer wg.Done()

	client := &http.Client{}
	req, _ := http.NewRequest("GET", "", nil)
	req.Header.Set("X-Auth-Token", c.Token)

	for path := range ci {
		fmt.Printf("Thread %03d: downloading %s\n", thread, path)
		FullPath := filepath.Join(BasePath, path)
		req.URL, _ = url.ParseRequestURI(fmt.Sprintf("%s/%s/%s", c.Endpoint, c.Container, path))
		res, err := client.Do(req)
		if err != nil {
			log.Print(fmt.Printf("%s\n", err))
			continue
		}

		os.MkdirAll(filepath.Dir(FullPath), 0755)
		file, err := os.Create(FullPath)
		if err != nil {
			log.Print(fmt.Printf("%s\n", err))
			continue
		}

		_, err = io.Copy(file, res.Body)
		if err != nil {
			log.Print(fmt.Printf("%s\n", err))
			continue
		}
		file.Close()
		res.Body.Close()
		fmt.Printf("Thread %03d: download complete for %s\n", thread, path)
	}
	fmt.Printf("Thread %03d: exiting\n", thread)
}

func Usage() {
	fmt.Printf(`usage: %s [options] {delete,upload,download} source [destination]

Delete:
    gohaste [options] delete my-container

Upload:
    gohaste [options] upload /path/to/files my-container

Download:
    gohaste [options] download my-container /path/to/files

options:
`, path.Base(os.Args[0]))
	flag.PrintDefaults()
	os.Exit(2)
}

func main() {
	var w Walker
	var Username string
	var Password string
	var Region string
	var Concurrency int
	var files []string

	flag.Usage = Usage
	flag.StringVar(&Username, "username", os.Getenv("OS_USERNAME"), "Username to authenticate with. Defaults to OS_USERNAME")
	flag.StringVar(&Password, "password", os.Getenv("OS_PASSWORD"), "Password to authenticate with. Defaults to OS_PASSWORD")
	flag.StringVar(&Region, "region", os.Getenv("OS_REGION_NAME"), "Password to authenticate with. Defaults to OS_REGION_NAME")
	flag.IntVar(&Concurrency, "concurrency", 10, "Number of cuncurrent operations. Defaults to 10")
	flag.Parse()

	Operation := strings.ToLower(flag.Arg(0))
	Src := flag.Arg(1)
	Dest := flag.Arg(2)

	if len(Username) == 0 || len(Password) == 0 || len(Region) == 0 || len(Operation) == 0 || len(Src) == 0 {
		Usage()
	}

	if Operation != "upload" && Operation != "download" && Operation != "delete" {
		log.Fatal(fmt.Sprintf("%s not a supported operation", Operation))
	}

	if Operation != "delete" && len(Dest) == 0 {
		log.Fatal("'destination' is a required argument for 'upload' and 'download'")
	}

	ci := make(chan string)
	wg := new(sync.WaitGroup)

	c := CloudFiles{
		Username: Username,
		ApiKey:   Password,
		Region:   Region,
	}
	c.Auth()

	if Operation == "upload" {
		w = Walker{}
		filepath.Walk(Src, w.Walk)
		files = w.Paths
		c.Container = Dest
	} else {
		c.Container = Src
		files = c.ListObjects()
	}

	if len(files) == 0 {
		log.Fatal("No files to operate on")
	}

	if Operation == "upload" {
		c.CreateContainer()
	}

	for i := 0; i < Concurrency; i++ {
		wg.Add(1)
		if Operation == "upload" {
			go c.Upload(i, ci, wg, Src)
		} else if Operation == "download" {
			Dest, _ = filepath.Abs(Dest)
			go c.Download(i, ci, wg, Dest)
		} else {
			go c.Delete(i, ci, wg)
		}
	}

	for _, path := range files {
		if len(strings.TrimSpace(path)) == 0 {
			continue
		}
		ci <- path
	}

	close(ci)
	wg.Wait()
}
<commit_msg>Instead of allocating slices to push into the channel, just push into the channel<commit_after>// Copyright 2014 Matt Martz <matt@sivel.net>
// 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 main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"path"
	"path/filepath"
	"strings"
	"sync"
)

type Walker struct {
	ci chan string
}

// Walk is the filepath.Walk WalkFunc to handle recording visited paths
func (w *Walker) Walk(path string, info os.FileInfo, err error) error {
	if !info.IsDir() {
		w.ci <- path
	}
	return nil
}

type AuthContainer struct {
	Auth Auth `json:"auth"`
}

type Auth struct {
	PasswordCredentials *PasswordCredentials `json:"passwordCredentials,omitempty"`
	ApiKeyCredentials   *ApiKeyCredentials   `json:"RAX-KSKEY:apiKeyCredentials,omitempty"`
	TenantId            string               `json:"tenantId,omitempty"`
	TenantName          string               `json:"tenantName,omitempty"`
}

type PasswordCredentials struct {
	Username string `json:"username"`
	Password string `json:"password"`
}

type ApiKeyCredentials struct {
	Username string `json:"username"`
	ApiKey   string `json:"apiKey"`
}

type Tokens struct {
	Access Access `json:"access"`
}

type Access struct {
	Token          Token `json:"token"`
	ServiceCatalog []CatalogEntry
}

type CatalogEntry struct {
	Name, Type string
	Endpoints  []EntryEndpoint
}

type EntryEndpoint struct {
	Region, TenantId                    string
	PublicURL, InternalURL              string
	VersionId, VersionInfo, VersionList string
}

type Token struct {
	Id      string `json:"id"`
	Expires string `json:"expires"`
}

type CloudFiles struct {
	Username  string
	ApiKey    string
	Token     string
	Endpoint  string
	Container string
	Region    string
}

// Auth will authenticate to CloudFiles
func (c *CloudFiles) Auth() {
	var tokens *Tokens

	auth := &AuthContainer{
		Auth: Auth{
			ApiKeyCredentials: &ApiKeyCredentials{
				Username: c.Username,
				ApiKey:   c.ApiKey,
			},
		},
	}
	body, _ := json.Marshal(auth)

	res, err := http.Post("https://identity.api.rackspacecloud.com/v2.0/tokens", "application/json", bytes.NewBuffer(body))
	defer res.Body.Close()
	if res.StatusCode != 200 || err != nil {
		log.Fatal("Unable to authenticate")
	}

	resBody, _ := ioutil.ReadAll(res.Body)

	json.Unmarshal(resBody, &tokens)

	for _, service := range tokens.Access.ServiceCatalog {
		if service.Type == "object-store" {
			for _, endpoint := range service.Endpoints {
				if endpoint.Region == c.Region {
					c.Endpoint = endpoint.PublicURL
					break
				}
			}
			break
		}
	}

	if len(c.Endpoint) == 0 {
		log.Fatal(fmt.Sprintf("No PublicURL found for object-store in region %s", c.Region))
	}

	c.Token = tokens.Access.Token.Id
}

func (c *CloudFiles) ListObjects(ci chan string) {
	var resBody []byte
	var marker string
	req, _ := http.NewRequest("GET", fmt.Sprintf("%s/%s", c.Endpoint, c.Container), nil)
	req.Header.Set("X-Auth-Token", c.Token)
	req.Header.Set("Accept", "text/plain")
	client := &http.Client{}
	res, err := client.Do(req)
	if err != nil {
		log.Fatal("err != nil")
		log.Fatal(err)
	} else if res.StatusCode != 200 {
		resBody, _ = ioutil.ReadAll(res.Body)
		fmt.Println(string(resBody))
		log.Fatal(res.StatusCode)
	}
	defer res.Body.Close()
	resBody, _ = ioutil.ReadAll(res.Body)
	for _, object := range strings.Split(strings.TrimSpace(string(resBody)), "\n") {
		if len(strings.TrimSpace(object)) == 0 {
			continue
		}
		ci <- object
		marker = object
	}
	go func(ci chan string, resBody []byte, marker string) {
		for len(resBody) > 0 {
			req.URL, _ = url.ParseRequestURI(fmt.Sprintf("%s/%s?marker=%s", c.Endpoint, c.Container, marker))
			res, err := client.Do(req)
			if err != nil {
				log.Fatal(err)
			}
			resBody, _ = ioutil.ReadAll(res.Body)
			for _, object := range strings.Split(strings.TrimSpace(string(resBody)), "\n") {
				if len(strings.TrimSpace(object)) == 0 {
					continue
				}
				ci <- object
				marker = object
			}
			res.Body.Close()
		}
		close(ci)
	}(ci, resBody, marker)
}

// CreateContainer ensures that a container exists
func (c *CloudFiles) CreateContainer() {
	req, _ := http.NewRequest("PUT", fmt.Sprintf("%s/%s", c.Endpoint, c.Container), bytes.NewBuffer([]byte{}))
	req.Header.Set("X-Auth-Token", c.Token)
	client := &http.Client{}
	res, _ := client.Do(req)
	fmt.Println(res.StatusCode)
}

// Upload is a goroutine that uploads files provided by a channel to a CloudFiles container
func (c *CloudFiles) Upload(thread int, ci chan string, wg *sync.WaitGroup, BasePath string) {
	fmt.Printf("Creating uploader thread: %03d\n", thread)

	defer wg.Done()

	client := &http.Client{}
	req, _ := http.NewRequest("PUT", "", nil)
	req.Header.Set("X-Auth-Token", c.Token)

	for path := range ci {
		ObjPath := strings.TrimPrefix(strings.Replace(path, BasePath, "", 1), "/")
		fmt.Printf("Thread %03d: uploading %s\n", thread, ObjPath)

		file, err := os.Open(path)
		if err != nil {
			log.Print(fmt.Printf("%s\n", err))
			continue
		}

		req.URL, _ = url.ParseRequestURI(fmt.Sprintf("%s/%s/%s", c.Endpoint, c.Container, ObjPath))
		req.Body = file
		res, err := client.Do(req)
		if err != nil {
			log.Print(fmt.Printf("%s\n", err))
			continue
		}
		file.Close()
		res.Body.Close()
		fmt.Printf("Thread %03d: upload complete for %s\n", thread, ObjPath)
	}
	fmt.Printf("Thread %03d: exiting\n", thread)
}

func (c *CloudFiles) Delete(thread int, ci chan string, wg *sync.WaitGroup) {
	fmt.Printf("Creating deleter thread: %03d\n", thread)

	defer wg.Done()

	client := &http.Client{}
	req, _ := http.NewRequest("DELETE", "", nil)
	req.Header.Set("X-Auth-Token", c.Token)

	for path := range ci {
		fmt.Printf("Thread %03d: deleting %s\n", thread, path)
		req.URL, _ = url.ParseRequestURI(fmt.Sprintf("%s/%s/%s", c.Endpoint, c.Container, path))
		res, err := client.Do(req)
		if err != nil {
			log.Print(fmt.Printf("%s\n", err))
			continue
		}
		res.Body.Close()
		fmt.Printf("Thread %03d: delete complete for %s\n", thread, path)
	}
	fmt.Printf("Thread %03d: exiting\n", thread)
}

func (c *CloudFiles) Download(thread int, ci chan string, wg *sync.WaitGroup, BasePath string) {
	fmt.Printf("Creating downloader thread: %03d\n", thread)

	defer wg.Done()

	client := &http.Client{}
	req, _ := http.NewRequest("GET", "", nil)
	req.Header.Set("X-Auth-Token", c.Token)

	for path := range ci {
		fmt.Printf("Thread %03d: downloading %s\n", thread, path)
		FullPath := filepath.Join(BasePath, path)
		req.URL, _ = url.ParseRequestURI(fmt.Sprintf("%s/%s/%s", c.Endpoint, c.Container, path))
		res, err := client.Do(req)
		if err != nil {
			log.Print(fmt.Printf("%s\n", err))
			continue
		}

		os.MkdirAll(filepath.Dir(FullPath), 0755)
		file, err := os.Create(FullPath)
		if err != nil {
			log.Print(fmt.Printf("%s\n", err))
			continue
		}

		_, err = io.Copy(file, res.Body)
		if err != nil {
			log.Print(fmt.Printf("%s\n", err))
			continue
		}
		file.Close()
		res.Body.Close()
		fmt.Printf("Thread %03d: download complete for %s\n", thread, path)
	}
	fmt.Printf("Thread %03d: exiting\n", thread)
}

func Usage() {
	fmt.Printf(`usage: %s [options] {delete,upload,download} source [destination]

Delete:
    gohaste [options] delete my-container

Upload:
    gohaste [options] upload /path/to/files my-container

Download:
    gohaste [options] download my-container /path/to/files

options:
`, path.Base(os.Args[0]))
	flag.PrintDefaults()
	os.Exit(2)
}

func main() {
	var w Walker
	var Username string
	var Password string
	var Region string
	var Concurrency int

	flag.Usage = Usage
	flag.StringVar(&Username, "username", os.Getenv("OS_USERNAME"), "Username to authenticate with. Defaults to OS_USERNAME")
	flag.StringVar(&Password, "password", os.Getenv("OS_PASSWORD"), "Password to authenticate with. Defaults to OS_PASSWORD")
	flag.StringVar(&Region, "region", os.Getenv("OS_REGION_NAME"), "Password to authenticate with. Defaults to OS_REGION_NAME")
	flag.IntVar(&Concurrency, "concurrency", 10, "Number of cuncurrent operations. Defaults to 10")
	flag.Parse()

	Operation := strings.ToLower(flag.Arg(0))
	Src := flag.Arg(1)
	Dest := flag.Arg(2)

	if len(Username) == 0 || len(Password) == 0 || len(Region) == 0 || len(Operation) == 0 || len(Src) == 0 {
		Usage()
	}

	if Operation != "upload" && Operation != "download" && Operation != "delete" {
		log.Fatal(fmt.Sprintf("%s not a supported operation", Operation))
	}

	if Operation != "delete" && len(Dest) == 0 {
		log.Fatal("'destination' is a required argument for 'upload' and 'download'")
	}

	ci := make(chan string)
	wg := new(sync.WaitGroup)

	c := CloudFiles{
		Username: Username,
		ApiKey:   Password,
		Region:   Region,
	}
	c.Auth()

	for i := 0; i < Concurrency; i++ {
		wg.Add(1)
		if Operation == "upload" {
			go c.Upload(i, ci, wg, Src)
		} else if Operation == "download" {
			Dest, _ = filepath.Abs(Dest)
			go c.Download(i, ci, wg, Dest)
		} else {
			go c.Delete(i, ci, wg)
		}
	}

	if Operation == "upload" {
		c.Container = Dest
		c.CreateContainer()
		w = Walker{ci: ci}
		filepath.Walk(Src, w.Walk)
		close(w.ci)
	} else {
		c.Container = Src
		c.ListObjects(ci)
	}

	wg.Wait()
}
<|endoftext|>


<commit_before>package gostalk

import (
	"net"
	"bufio"
	"fmt"
	"errors"
	"strings"
	"io"
	"log"
)

type Job struct {
	Id uint64
	Body []byte
}

type Gostalk struct {
	conn net.Conn
	reader *bufio.Reader
}

var ReaderSize = 4096

// Returns a new beanstalkd tcp connection and sets up a new buffered reader
func Connect(addr string) (this *Gostalk, err error) {

	this = new(Gostalk)
	this.conn, err = net.Dial("tcp", addr)
	if err != nil {
		return
	}

	this.reader = bufio.NewReaderSize(this.conn, ReaderSize)

	return
}

// Disconnects from the beanstalkd server
func (this *Gostalk) Disconnect() (err error) {

	err = this.conn.Close();

	return
}

// Put is is for any process that wants to insert a job into the queue.
// It takes the data, priority, delay, time to release
func (this *Gostalk) Put(data []byte, priority uint32, delay int, ttr int) (uint64, error) {

	dataLen := len(data)
	res, err := this.sendReceive("put %d %d %d %d\r\n%s\r\n", priority, delay, ttr, dataLen, data)
	if err != nil {
		return 0, err
	}

	var jobId uint64
	_, err = fmt.Sscanf(res, "INSERTED %d", &jobId)
	if err != nil {
		return 0, err
	}

	return jobId, nil
}

// Reserve waits for a job indefinitely from the server.
// Once the job has been received, it returns the job
// This is a blocking call.
func (this *Gostalk) Reserve() (*Job, error) {

	res, err := this.sendReceive("reserve\r\n")
	if err != nil {
		return nil, err
	}

	var jobId uint64
	var dataLen uint32
	_, err = fmt.Sscanf(res, "RESERVED %d %d", &jobId, &dataLen)
	if err != nil {
		return nil, err
	}

	//read job body
	body := make([]byte, dataLen + 2) //+2 is for trailing \r\n
	_, err = io.ReadFull(this.reader, body)
	if err != nil {
		log.Println("failed reading body:", err.Error())
		return nil, err
	}

	job := new(Job)
	job.Id = jobId
	job.Body = body[:len(body) - 2]

	return job, nil
}

// Delete a job by it's id
func (this *Gostalk) Delete(jobId uint64) (err error) {

	res, err := this.sendReceive("delete %d\r\n", jobId)
	if err != nil {
		return
	}

	if res != "DELETED" {
		err = errors.New(fmt.Sprintf("Expected (DELETED), GOT %s", res))
	}

	return
}

// Use specifies the tube to use. If the tube does not exist, it will be created.
func (this *Gostalk) Use(tubeName string) (err error) {

	res, err := this.sendReceive("use %s\r\n", tubeName)
	if err != nil {
		return
	}

	var tubeUsed string
	_, err = fmt.Sscanf(res, "USING %s", &tubeUsed)

	return
}

// Sends the watch command to beanstalkd and returns the number of tubes being watched
func (this *Gostalk) Watch(tubeName string) (int, error) {

	res, err := this.sendReceive("watch %s\r\n", tubeName)
	if err != nil {
		return -1, err
	}

	var tubeCount int
	_, err = fmt.Sscanf(res, "WATCHING %d", &tubeCount)
	if err != nil {
		return -1, err
	}

	return tubeCount, nil
}

// Ignores the specified tube and returns the number of tubes being watched.
// Cannot ignore a tube if it's the only tube being watched
func (this *Gostalk) Ignore(tubeName string) (int, error) {

	res, err := this.sendReceive("ignore %s\r\n", tubeName)
	if err != nil {
		return -1, err
	}

	var tubeCount int
	_, err = fmt.Sscanf(res, "WATCHING %d", &tubeCount)
	if err != nil {
		if res == "NOT_IGNORED" {
			return -1, errors.New(fmt.Sprintf("Tube (%s) NOT_IGNORED", tubeName))
		}
		return -1, err
	}

	return tubeCount, nil
}

// Bury a job by it's id
func (this *Gostalk) Bury(jobId uint64, priority uint32) (err error) {

	res, err := this.sendReceive("bury %d %d\r\n", jobId, priority)

	if err != nil {
		return
	}

	expectedRes := "BURIED"
	if res != expectedRes {
		err = errors.New(fmt.Sprintf("Expected (%s), Got (%s)", expectedRes, res))
	}

	return
}

// Sends a formatted command to beanstalkd and returns the received response
func (this *Gostalk) sendReceive(format string, args ...interface{}) (res string, err error) {
	_, err = fmt.Fprintf(this.conn, format, args...)

	if err != nil {
		return
	}

	res, err = this.reader.ReadString('\n')

	res = strings.TrimSuffix(res, "\r\n")

	return
}
<commit_msg>gofmt<commit_after>package gostalk

import (
	"net"
	"bufio"
	"fmt"
	"errors"
	"strings"
	"io"
	"log"
)

type Job struct {
	Id   uint64
	Body []byte
}

type Gostalk struct {
	conn   net.Conn
	reader *bufio.Reader
}

var ReaderSize = 4096

// Returns a new beanstalkd tcp connection and sets up a new buffered reader
func Connect(addr string) (this *Gostalk, err error) {

	this = new(Gostalk)
	this.conn, err = net.Dial("tcp", addr)
	if err != nil {
		return
	}

	this.reader = bufio.NewReaderSize(this.conn, ReaderSize)

	return
}

// Disconnects from the beanstalkd server
func (this *Gostalk) Disconnect() (err error) {

	err = this.conn.Close();

	return
}

// Put is is for any process that wants to insert a job into the queue.
// It takes the data, priority, delay, time to release
func (this *Gostalk) Put(data []byte, priority uint32, delay int, ttr int) (uint64, error) {

	dataLen := len(data)
	res, err := this.sendReceive("put %d %d %d %d\r\n%s\r\n", priority, delay, ttr, dataLen, data)
	if err != nil {
		return 0, err
	}

	var jobId uint64
	_, err = fmt.Sscanf(res, "INSERTED %d", &jobId)
	if err != nil {
		return 0, err
	}

	return jobId, nil
}

// Reserve waits for a job indefinitely from the server.
// Once the job has been received, it returns the job
// This is a blocking call.
func (this *Gostalk) Reserve() (*Job, error) {

	res, err := this.sendReceive("reserve\r\n")
	if err != nil {
		return nil, err
	}

	var jobId uint64
	var dataLen uint32
	_, err = fmt.Sscanf(res, "RESERVED %d %d", &jobId, &dataLen)
	if err != nil {
		return nil, err
	}

	//read job body
	body := make([]byte, dataLen + 2) //+2 is for trailing \r\n
	_, err = io.ReadFull(this.reader, body)
	if err != nil {
		log.Println("failed reading body:", err.Error())
		return nil, err
	}

	job := new(Job)
	job.Id = jobId
	job.Body = body[:len(body) - 2]

	return job, nil
}

// Delete a job by it's id
func (this *Gostalk) Delete(jobId uint64) (err error) {

	res, err := this.sendReceive("delete %d\r\n", jobId)
	if err != nil {
		return
	}

	if res != "DELETED" {
		err = errors.New(fmt.Sprintf("Expected (DELETED), GOT %s", res))
	}

	return
}

// Use specifies the tube to use. If the tube does not exist, it will be created.
func (this *Gostalk) Use(tubeName string) (err error) {

	res, err := this.sendReceive("use %s\r\n", tubeName)
	if err != nil {
		return
	}

	var tubeUsed string
	_, err = fmt.Sscanf(res, "USING %s", &tubeUsed)

	return
}

// Sends the watch command to beanstalkd and returns the number of tubes being watched
func (this *Gostalk) Watch(tubeName string) (int, error) {

	res, err := this.sendReceive("watch %s\r\n", tubeName)
	if err != nil {
		return -1, err
	}

	var tubeCount int
	_, err = fmt.Sscanf(res, "WATCHING %d", &tubeCount)
	if err != nil {
		return -1, err
	}

	return tubeCount, nil
}

// Ignores the specified tube and returns the number of tubes being watched.
// Cannot ignore a tube if it's the only tube being watched
func (this *Gostalk) Ignore(tubeName string) (int, error) {

	res, err := this.sendReceive("ignore %s\r\n", tubeName)
	if err != nil {
		return -1, err
	}

	var tubeCount int
	_, err = fmt.Sscanf(res, "WATCHING %d", &tubeCount)
	if err != nil {
		if res == "NOT_IGNORED" {
			return -1, errors.New(fmt.Sprintf("Tube (%s) NOT_IGNORED", tubeName))
		}
		return -1, err
	}

	return tubeCount, nil
}

// Bury a job by it's id
func (this *Gostalk) Bury(jobId uint64, priority uint32) (err error) {

	res, err := this.sendReceive("bury %d %d\r\n", jobId, priority)

	if err != nil {
		return
	}

	expectedRes := "BURIED"
	if res != expectedRes {
		err = errors.New(fmt.Sprintf("Expected (%s), Got (%s)", expectedRes, res))
	}

	return
}

// Sends a formatted command to beanstalkd and returns the received response
func (this *Gostalk) sendReceive(format string, args ...interface{}) (res string, err error) {
	_, err = fmt.Fprintf(this.conn, format, args...)

	if err != nil {
		return
	}

	res, err = this.reader.ReadString('\n')

	res = strings.TrimSuffix(res, "\r\n")

	return
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"github.com/belogik/goes"
	goopt "github.com/droundy/goopt"
	"io/ioutil"
	"log"
	"net"
	"os"
	"strconv"
	"strings"
)

type UDPData struct {
	prefix  string `json: "prefix"`
	project string `json: "project"`
	metric  string `json: "metric"`
	value   string `json: "value"`
}

var port = goopt.Int([]string{"-p", "--port"}, 8125, "UDP Port to use")
var config = goopt.String([]string{"-c", "--config"}, "./config.json", "Configuration file to use")

func readConfig(filepath string) (map[string]interface{}, error) {
	file, e := ioutil.ReadFile(filepath)
	if e != nil {
		return nil, e
	}
	var cfg map[string]interface{}
	if e = json.Unmarshal(file, &cfg); e != nil {
		return nil, e
	}
	return cfg, nil
}

func getUDPAddressFromConfig(serverType string, cfg map[string]interface{}) (net.UDPAddr, error) {
	key := cfg[serverType].(map[string]interface{})
	parsedIp, _, e := net.ParseCIDR(key["ip"].(string))
	if e != nil {
		return net.UDPAddr{}, e
	}
	return net.UDPAddr{IP: parsedIp, Port: int(key["port"].(float64))}, e
}

func getElasticSearchConnection(serverType string, cfg map[string]interface{}) (conn *goes.Connection) {
	key := cfg[serverType].(map[string]interface{})
	conn = goes.NewConnection(key["hostname"].(string), strconv.FormatFloat(key["port"].(float64), 'f', -1, 64))
	return
}

func sendToGraphite(message []byte, conn net.UDPConn, graphite net.UDPAddr) {
	if message != nil && len(message) > 0 {
		conn.WriteToUDP([]byte(fmt.Sprintf("%s|g", string(message))), &graphite)
	}
}

func sendToElasticsearch(message []byte, conn goes.Connection) {
	data := createDataStruct(message)
	jsonData := map[string]interface{}{
		"metric": data.metric,
		"value":  strings.Trim(data.value, "\u0000"),
	}
	doc := goes.Document{
		Index:  "sonar",
		Type:   "metric",
		Fields: jsonData,
	}
	conn.Index(doc, nil)
}

func createDataStruct(message []byte) UDPData {
	valueSplit := strings.Split(string(message), ":")
	keySplit := strings.Split(valueSplit[0], ".")
	return UDPData{project: keySplit[2], prefix: keySplit[0] + "." + keySplit[1], metric: keySplit[3], value: valueSplit[1]}
}

func main() {
	goopt.Description = func() string {
		return "Metric Wrapper for (at first) graphite & elasticsearch."
	}
	goopt.Version = "1.0"
	goopt.Summary = "gostats"
	goopt.Parse(nil)

	readConf, err := readConfig(*config)

	if err != nil {
		fmt.Printf("Failed to parse configuration file: %v\n", err)
		os.Exit(1)
	}

	cfg := readConf["config"].(map[string]interface{})

	elasticConn := getElasticSearchConnection("elasticsearch", cfg)

	graphite, err := getUDPAddressFromConfig("graphite", cfg)
	if err != nil {
		fmt.Printf("Failed to get Server address for Graphite: %v\n", err)
		os.Exit(0)
	}

	graphiteConn, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.IPv4zero, Port: 0})
	if err != nil {
		fmt.Printf("Failed to create UDP Connection: %v\n", err)
		os.Exit(0)
	}

	fmt.Println(graphite)

	addr, _ := net.ResolveUDPAddr("udp", ":"+strconv.Itoa(*port))
	fmt.Println(addr)
	conn, err := net.ListenUDP("udp", addr)

	if err != nil {
		fmt.Println(err.Error())
		return
	}
	defer conn.Close()

	for {
		message := make([]byte, 512)
		n, _, err := conn.ReadFromUDP(message)
		log.Printf("Got %d bytes\n", n)
		log.Printf("Data: %s", message)
		if err != nil || n == 0 {
			log.Printf("Error is: %s, bytes are: %d", err, n)
			continue
		}
		sendToGraphite(message, *graphiteConn, graphite)
		sendToElasticsearch(message, *elasticConn)
	}
}
<commit_msg>add error logging<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"github.com/belogik/goes"
	goopt "github.com/droundy/goopt"
	"io/ioutil"
	"log"
	"net"
	"os"
	"strconv"
	"strings"
)

type UDPData struct {
	prefix  string `json: "prefix"`
	project string `json: "project"`
	metric  string `json: "metric"`
	value   string `json: "value"`
}

var port = goopt.Int([]string{"-p", "--port"}, 8125, "UDP Port to use")
var config = goopt.String([]string{"-c", "--config"}, "./config.json", "Configuration file to use")

func readConfig(filepath string) (map[string]interface{}, error) {
	file, e := ioutil.ReadFile(filepath)
	if e != nil {
		return nil, e
	}
	var cfg map[string]interface{}
	if e = json.Unmarshal(file, &cfg); e != nil {
		return nil, e
	}
	return cfg, nil
}

func getUDPAddressFromConfig(serverType string, cfg map[string]interface{}) (net.UDPAddr, error) {
	key := cfg[serverType].(map[string]interface{})
	parsedIp, _, e := net.ParseCIDR(key["ip"].(string))
	if e != nil {
		return net.UDPAddr{}, e
	}
	return net.UDPAddr{IP: parsedIp, Port: int(key["port"].(float64))}, e
}

func getElasticSearchConnection(serverType string, cfg map[string]interface{}) (conn *goes.Connection) {
	key := cfg[serverType].(map[string]interface{})
	conn = goes.NewConnection(key["hostname"].(string), strconv.FormatFloat(key["port"].(float64), 'f', -1, 64))
	return
}

func sendToGraphite(message []byte, conn net.UDPConn, graphite net.UDPAddr) {
	if message != nil && len(message) > 0 {
		_, err := conn.WriteToUDP([]byte(fmt.Sprintf("%s|g", string(message))), &graphite)
		if err != nil {
			fmt.Println(err)
		}
	}
}

func sendToElasticsearch(message []byte, conn goes.Connection) {
	data := createDataStruct(message)
	jsonData := map[string]interface{}{
		"metric": data.metric,
		"value":  strings.Trim(data.value, "\u0000"),
	}
	doc := goes.Document{
		Index:  "sonar",
		Type:   "metric",
		Fields: jsonData,
	}
	_, err := conn.Index(doc, nil)
	if err != nil {
		fmt.Println(err)
	}
}

func createDataStruct(message []byte) UDPData {
	valueSplit := strings.Split(string(message), ":")
	keySplit := strings.Split(valueSplit[0], ".")
	return UDPData{project: keySplit[2], prefix: keySplit[0] + "." + keySplit[1], metric: keySplit[3], value: valueSplit[1]}
}

func main() {
	goopt.Description = func() string {
		return "Metric Wrapper for (at first) graphite & elasticsearch."
	}
	goopt.Version = "1.0"
	goopt.Summary = "gostats"
	goopt.Parse(nil)

	readConf, err := readConfig(*config)

	if err != nil {
		fmt.Printf("Failed to parse configuration file: %v\n", err)
		os.Exit(1)
	}

	cfg := readConf["config"].(map[string]interface{})

	elasticConn := getElasticSearchConnection("elasticsearch", cfg)

	graphite, err := getUDPAddressFromConfig("graphite", cfg)
	if err != nil {
		fmt.Printf("Failed to get Server address for Graphite: %v\n", err)
		os.Exit(0)
	}

	graphiteConn, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.IPv4zero, Port: 0})
	if err != nil {
		fmt.Printf("Failed to create UDP Connection: %v\n", err)
		os.Exit(0)
	}

	fmt.Println(graphite)

	addr, _ := net.ResolveUDPAddr("udp", ":"+strconv.Itoa(*port))
	fmt.Println(addr)
	conn, err := net.ListenUDP("udp", addr)

	if err != nil {
		fmt.Println(err.Error())
		return
	}
	defer conn.Close()

	for {
		message := make([]byte, 512)
		n, _, err := conn.ReadFromUDP(message)
		log.Printf("Got %d bytes\n", n)
		log.Printf("Data: %s", message)
		if err != nil || n == 0 {
			log.Printf("Error is: %s, bytes are: %d", err, n)
			continue
		}
		sendToGraphite(message, *graphiteConn, graphite)
		sendToElasticsearch(message, *elasticConn)
	}
}
<|endoftext|>