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

package aws

import (
	"fmt"
	"net/url"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/aws/aws-sdk-go/service/iam"
)

type EC2ImageType string

const (
	EC2ImageTypeHVM EC2ImageType = "hvm"
	EC2ImageTypePV  EC2ImageType = "paravirtual"
)

type EC2ImageFormat string

const (
	EC2ImageFormatRaw  EC2ImageFormat = ec2.DiskImageFormatRaw
	EC2ImageFormatVmdk                = ec2.DiskImageFormatVmdk
)

// TODO, these can be derived at runtime
var akis = map[string]string{
	"us-east-1":      "aki-919dcaf8",
	"us-east-2":      "aki-da055ebf",
	"us-west-1":      "aki-880531cd",
	"us-west-2":      "aki-fc8f11cc",
	"eu-west-1":      "aki-52a34525",
	"eu-west-2":      "aki-8b6369ef",
	"eu-central-1":   "aki-184c7a05",
	"ap-south-1":     "aki-a7305ac8",
	"ap-southeast-1": "aki-503e7402",
	"ap-southeast-2": "aki-c362fff9",
	"ap-northeast-1": "aki-176bf516",
	"ap-northeast-2": "aki-01a66b6f",
	"sa-east-1":      "aki-5553f448",
	"ca-central-1":   "aki-320ebd56",

	"us-gov-west-1": "aki-1de98d3e",
}

func (e *EC2ImageFormat) Set(s string) error {
	switch s {
	case string(EC2ImageFormatVmdk):
		*e = EC2ImageFormatVmdk
	case string(EC2ImageFormatRaw):
		*e = EC2ImageFormatRaw
	default:
		return fmt.Errorf("invalid ec2 image format: must be raw or vmdk")
	}
	return nil
}

func (e *EC2ImageFormat) String() string {
	return string(*e)
}

func (e *EC2ImageFormat) Type() string {
	return "ec2ImageFormat"
}

var vmImportRole = "vmimport"

type Snapshot struct {
	SnapshotID string
}

// CreateSnapshot creates an AWS Snapshot
func (a *API) CreateSnapshot(imageName, description, sourceURL string, format EC2ImageFormat) (*Snapshot, error) {
	if format == "" {
		format = EC2ImageFormatVmdk
	}
	s3url, err := url.Parse(sourceURL)
	if err != nil {
		return nil, err
	}
	if s3url.Scheme != "s3" {
		return nil, fmt.Errorf("source must have a 's3://' scheme, not: '%v://'", s3url.Scheme)
	}
	s3key := strings.TrimPrefix(s3url.Path, "/")

	importRes, err := a.ec2.ImportSnapshot(&ec2.ImportSnapshotInput{
		RoleName:    aws.String(vmImportRole),
		Description: aws.String(description),
		DiskContainer: &ec2.SnapshotDiskContainer{
			// TODO(euank): allow s3 source / local file -> s3 source
			UserBucket: &ec2.UserBucket{
				S3Bucket: aws.String(s3url.Host),
				S3Key:    aws.String(s3key),
			},
			Format: aws.String(string(format)),
		},
	})
	if err != nil {
		return nil, fmt.Errorf("unable to create import snapshot task: %v", err)
	}

	plog.Infof("created snapshot import task %v", *importRes.ImportTaskId)

	snapshotDone := func(snapshotTaskID string) (bool, string, error) {
		taskRes, err := a.ec2.DescribeImportSnapshotTasks(&ec2.DescribeImportSnapshotTasksInput{
			ImportTaskIds: []*string{aws.String(snapshotTaskID)},
		})
		if err != nil {
			return false, "", err
		}

		details := taskRes.ImportSnapshotTasks[0].SnapshotTaskDetail

		// I dream of AWS specifying this as an enum shape, not string
		switch *details.Status {
		case "completed":
			return true, *details.SnapshotId, nil
		case "pending", "active":
			plog.Debugf("waiting for import task: %v (%v): %v", *details.Status, *details.Progress, *details.StatusMessage)
			return false, "", nil
		case "cancelled", "cancelling":
			return false, "", fmt.Errorf("import task cancelled")
		case "deleted", "deleting":
			errMsg := "unknown error occured importing snapshot"
			if details.StatusMessage != nil {
				errMsg = *details.StatusMessage
			}
			return false, "", fmt.Errorf("could not import snapshot: %v", errMsg)
		default:
			return false, "", fmt.Errorf("unexpected status: %v", *details.Status)
		}
	}

	// TODO(euank): write a waiter for import snapshot
	var snapshotID string
	for {
		var done bool
		done, snapshotID, err = snapshotDone(*importRes.ImportTaskId)
		if err != nil {
			return nil, err
		}
		if done {
			break
		}
		time.Sleep(20 * time.Second)
	}

	err = a.CreateTags([]string{snapshotID}, map[string]string{
		"Name": imageName,
	})
	if err != nil {
		return nil, fmt.Errorf("couldn't create tags: %v", err)
	}

	return &Snapshot{
		SnapshotID: snapshotID,
	}, nil
}

func (a *API) CreateImportRole(bucket string) error {
	iamc := iam.New(a.session)
	_, err := iamc.GetRole(&iam.GetRoleInput{
		RoleName: &vmImportRole,
	})
	if err != nil {
		if awserr, ok := err.(awserr.Error); ok && awserr.Code() == "NoSuchEntity" {
			// Role does not exist, let's try to create it
			_, err := iamc.CreateRole(&iam.CreateRoleInput{
				RoleName: &vmImportRole,
				AssumeRolePolicyDocument: aws.String(`{
					"Version": "2012-10-17",
					"Statement": [{
						"Effect": "Allow",
						"Condition": {
							"StringEquals": {
								"sts:ExternalId": "vmimport"
							}
						},
						"Action": "sts:AssumeRole",
						"Principal": {
							"Service": "vmie.amazonaws.com"
						}
					}]
				}`),
			})
			if err != nil {
				return fmt.Errorf("coull not create vmimport role: %v", err)
			}
		}
	}

	// by convention, name our policies after the bucket so we can identify
	// whether a regional bucket is covered by a policy without parsing the
	// policy-doc json
	policyName := bucket
	_, err = iamc.GetRolePolicy(&iam.GetRolePolicyInput{
		RoleName:   &vmImportRole,
		PolicyName: &policyName,
	})
	if err != nil {
		if awserr, ok := err.(awserr.Error); ok && awserr.Code() == "NoSuchEntity" {
			// Policy does not exist, let's try to create it
			_, err := iamc.PutRolePolicy(&iam.PutRolePolicyInput{
				RoleName:   &vmImportRole,
				PolicyName: &policyName,
				PolicyDocument: aws.String((`{
	"Version": "2012-10-17",
	"Statement": [{
		"Effect": "Allow",
		"Action": [
			"s3:ListBucket",
			"s3:GetBucketLocation",
			"s3:GetObject"
		],
		"Resource": [
			"arn:aws:s3:::` + bucket + `",
			"arn:aws:s3:::` + bucket + `/*"
		]
	},
	{
		"Effect": "Allow",
		"Action": [
			"ec2:ModifySnapshotAttribute",
			"ec2:CopySnapshot",
			"ec2:RegisterImage",
			"ec2:Describe*"
		],
		"Resource": "*"
	}]
}`)),
			})
			if err != nil {
				return fmt.Errorf("could not create role policy: %v", err)
			}
		} else {
			return err
		}
	}

	return nil
}

func (a *API) CreateHVMImage(snapshotID string, name string, description string) (string, error) {
	params := registerImageParams(snapshotID, name+"-hvm", description, "xvd", EC2ImageTypeHVM)
	params.EnaSupport = aws.Bool(true)
	res, err := a.ec2.RegisterImage(params)
	if err != nil {
		return "", fmt.Errorf("error creating hvm AMI: %v", err)
	}
	return *res.ImageId, nil
}

func (a *API) CreatePVImage(snapshotID string, name string, description string) (string, error) {
	params := registerImageParams(snapshotID, name, description, "sd", EC2ImageTypePV)
	params.KernelId = aws.String(akis[a.opts.Region])
	res, err := a.ec2.RegisterImage(params)
	if err != nil {
		return "", fmt.Errorf("error creating paravirtual AMI: %v", err)
	}
	return *res.ImageId, nil
}

const diskSize = 8 // GB

func registerImageParams(snapshotID, name, description string, diskBaseName string, imageType EC2ImageType) *ec2.RegisterImageInput {
	return &ec2.RegisterImageInput{
		Name:               aws.String(name),
		Description:        aws.String(description),
		Architecture:       aws.String("x86_64"),
		VirtualizationType: aws.String(string(imageType)),
		RootDeviceName:     aws.String(fmt.Sprintf("/dev/%sa", diskBaseName)),
		BlockDeviceMappings: []*ec2.BlockDeviceMapping{
			&ec2.BlockDeviceMapping{
				DeviceName: aws.String(fmt.Sprintf("/dev/%sa", diskBaseName)),
				Ebs: &ec2.EbsBlockDevice{
					SnapshotId:          aws.String(snapshotID),
					DeleteOnTermination: aws.Bool(true),
					VolumeSize:          aws.Int64(diskSize),
				},
			},
			&ec2.BlockDeviceMapping{
				DeviceName:  aws.String(fmt.Sprintf("/dev/%sb", diskBaseName)),
				VirtualName: aws.String("ephemeral0"),
			},
		},
	}
}
<commit_msg>platform/api/aws: Find existing EBS snapshot by name and reuse it<commit_after>// Copyright 2016 CoreOS, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package aws

import (
	"fmt"
	"net/url"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/aws/aws-sdk-go/service/iam"
)

type EC2ImageType string

const (
	EC2ImageTypeHVM EC2ImageType = "hvm"
	EC2ImageTypePV  EC2ImageType = "paravirtual"
)

type EC2ImageFormat string

const (
	EC2ImageFormatRaw  EC2ImageFormat = ec2.DiskImageFormatRaw
	EC2ImageFormatVmdk                = ec2.DiskImageFormatVmdk
)

// TODO, these can be derived at runtime
var akis = map[string]string{
	"us-east-1":      "aki-919dcaf8",
	"us-east-2":      "aki-da055ebf",
	"us-west-1":      "aki-880531cd",
	"us-west-2":      "aki-fc8f11cc",
	"eu-west-1":      "aki-52a34525",
	"eu-west-2":      "aki-8b6369ef",
	"eu-central-1":   "aki-184c7a05",
	"ap-south-1":     "aki-a7305ac8",
	"ap-southeast-1": "aki-503e7402",
	"ap-southeast-2": "aki-c362fff9",
	"ap-northeast-1": "aki-176bf516",
	"ap-northeast-2": "aki-01a66b6f",
	"sa-east-1":      "aki-5553f448",
	"ca-central-1":   "aki-320ebd56",

	"us-gov-west-1": "aki-1de98d3e",
}

func (e *EC2ImageFormat) Set(s string) error {
	switch s {
	case string(EC2ImageFormatVmdk):
		*e = EC2ImageFormatVmdk
	case string(EC2ImageFormatRaw):
		*e = EC2ImageFormatRaw
	default:
		return fmt.Errorf("invalid ec2 image format: must be raw or vmdk")
	}
	return nil
}

func (e *EC2ImageFormat) String() string {
	return string(*e)
}

func (e *EC2ImageFormat) Type() string {
	return "ec2ImageFormat"
}

var vmImportRole = "vmimport"

type Snapshot struct {
	SnapshotID string
}

// CreateSnapshot creates an AWS Snapshot
func (a *API) CreateSnapshot(imageName, description, sourceURL string, format EC2ImageFormat) (*Snapshot, error) {
	if format == "" {
		format = EC2ImageFormatVmdk
	}
	s3url, err := url.Parse(sourceURL)
	if err != nil {
		return nil, err
	}
	if s3url.Scheme != "s3" {
		return nil, fmt.Errorf("source must have a 's3://' scheme, not: '%v://'", s3url.Scheme)
	}
	s3key := strings.TrimPrefix(s3url.Path, "/")

	// Look for an existing snapshot with this image name.
	snapshotRes, err := a.ec2.DescribeSnapshots(&ec2.DescribeSnapshotsInput{
		Filters: []*ec2.Filter{
			&ec2.Filter{
				Name:   aws.String("status"),
				Values: aws.StringSlice([]string{"completed"}),
			},
			&ec2.Filter{
				Name:   aws.String("tag:Name"),
				Values: aws.StringSlice([]string{imageName}),
			},
		},
		OwnerIds: aws.StringSlice([]string{"self"}),
	})
	if err != nil {
		return nil, fmt.Errorf("unable to describe snapshots: %v", err)
	}
	if len(snapshotRes.Snapshots) > 1 {
		return nil, fmt.Errorf("found multiple matching snapshots")
	}
	if len(snapshotRes.Snapshots) == 1 {
		snapshotID := *snapshotRes.Snapshots[0].SnapshotId
		plog.Infof("found existing snapshot %v, reusing", snapshotID)
		return &Snapshot{
			SnapshotID: snapshotID,
		}, nil
	}

	importRes, err := a.ec2.ImportSnapshot(&ec2.ImportSnapshotInput{
		RoleName:    aws.String(vmImportRole),
		Description: aws.String(description),
		DiskContainer: &ec2.SnapshotDiskContainer{
			// TODO(euank): allow s3 source / local file -> s3 source
			UserBucket: &ec2.UserBucket{
				S3Bucket: aws.String(s3url.Host),
				S3Key:    aws.String(s3key),
			},
			Format: aws.String(string(format)),
		},
	})
	if err != nil {
		return nil, fmt.Errorf("unable to create import snapshot task: %v", err)
	}

	plog.Infof("created snapshot import task %v", *importRes.ImportTaskId)

	snapshotDone := func(snapshotTaskID string) (bool, string, error) {
		taskRes, err := a.ec2.DescribeImportSnapshotTasks(&ec2.DescribeImportSnapshotTasksInput{
			ImportTaskIds: []*string{aws.String(snapshotTaskID)},
		})
		if err != nil {
			return false, "", err
		}

		details := taskRes.ImportSnapshotTasks[0].SnapshotTaskDetail

		// I dream of AWS specifying this as an enum shape, not string
		switch *details.Status {
		case "completed":
			return true, *details.SnapshotId, nil
		case "pending", "active":
			plog.Debugf("waiting for import task: %v (%v): %v", *details.Status, *details.Progress, *details.StatusMessage)
			return false, "", nil
		case "cancelled", "cancelling":
			return false, "", fmt.Errorf("import task cancelled")
		case "deleted", "deleting":
			errMsg := "unknown error occured importing snapshot"
			if details.StatusMessage != nil {
				errMsg = *details.StatusMessage
			}
			return false, "", fmt.Errorf("could not import snapshot: %v", errMsg)
		default:
			return false, "", fmt.Errorf("unexpected status: %v", *details.Status)
		}
	}

	// TODO(euank): write a waiter for import snapshot
	var snapshotID string
	for {
		var done bool
		done, snapshotID, err = snapshotDone(*importRes.ImportTaskId)
		if err != nil {
			return nil, err
		}
		if done {
			break
		}
		time.Sleep(20 * time.Second)
	}

	err = a.CreateTags([]string{snapshotID}, map[string]string{
		"Name": imageName,
	})
	if err != nil {
		return nil, fmt.Errorf("couldn't create tags: %v", err)
	}

	return &Snapshot{
		SnapshotID: snapshotID,
	}, nil
}

func (a *API) CreateImportRole(bucket string) error {
	iamc := iam.New(a.session)
	_, err := iamc.GetRole(&iam.GetRoleInput{
		RoleName: &vmImportRole,
	})
	if err != nil {
		if awserr, ok := err.(awserr.Error); ok && awserr.Code() == "NoSuchEntity" {
			// Role does not exist, let's try to create it
			_, err := iamc.CreateRole(&iam.CreateRoleInput{
				RoleName: &vmImportRole,
				AssumeRolePolicyDocument: aws.String(`{
					"Version": "2012-10-17",
					"Statement": [{
						"Effect": "Allow",
						"Condition": {
							"StringEquals": {
								"sts:ExternalId": "vmimport"
							}
						},
						"Action": "sts:AssumeRole",
						"Principal": {
							"Service": "vmie.amazonaws.com"
						}
					}]
				}`),
			})
			if err != nil {
				return fmt.Errorf("coull not create vmimport role: %v", err)
			}
		}
	}

	// by convention, name our policies after the bucket so we can identify
	// whether a regional bucket is covered by a policy without parsing the
	// policy-doc json
	policyName := bucket
	_, err = iamc.GetRolePolicy(&iam.GetRolePolicyInput{
		RoleName:   &vmImportRole,
		PolicyName: &policyName,
	})
	if err != nil {
		if awserr, ok := err.(awserr.Error); ok && awserr.Code() == "NoSuchEntity" {
			// Policy does not exist, let's try to create it
			_, err := iamc.PutRolePolicy(&iam.PutRolePolicyInput{
				RoleName:   &vmImportRole,
				PolicyName: &policyName,
				PolicyDocument: aws.String((`{
	"Version": "2012-10-17",
	"Statement": [{
		"Effect": "Allow",
		"Action": [
			"s3:ListBucket",
			"s3:GetBucketLocation",
			"s3:GetObject"
		],
		"Resource": [
			"arn:aws:s3:::` + bucket + `",
			"arn:aws:s3:::` + bucket + `/*"
		]
	},
	{
		"Effect": "Allow",
		"Action": [
			"ec2:ModifySnapshotAttribute",
			"ec2:CopySnapshot",
			"ec2:RegisterImage",
			"ec2:Describe*"
		],
		"Resource": "*"
	}]
}`)),
			})
			if err != nil {
				return fmt.Errorf("could not create role policy: %v", err)
			}
		} else {
			return err
		}
	}

	return nil
}

func (a *API) CreateHVMImage(snapshotID string, name string, description string) (string, error) {
	params := registerImageParams(snapshotID, name+"-hvm", description, "xvd", EC2ImageTypeHVM)
	params.EnaSupport = aws.Bool(true)
	res, err := a.ec2.RegisterImage(params)
	if err != nil {
		return "", fmt.Errorf("error creating hvm AMI: %v", err)
	}
	return *res.ImageId, nil
}

func (a *API) CreatePVImage(snapshotID string, name string, description string) (string, error) {
	params := registerImageParams(snapshotID, name, description, "sd", EC2ImageTypePV)
	params.KernelId = aws.String(akis[a.opts.Region])
	res, err := a.ec2.RegisterImage(params)
	if err != nil {
		return "", fmt.Errorf("error creating paravirtual AMI: %v", err)
	}
	return *res.ImageId, nil
}

const diskSize = 8 // GB

func registerImageParams(snapshotID, name, description string, diskBaseName string, imageType EC2ImageType) *ec2.RegisterImageInput {
	return &ec2.RegisterImageInput{
		Name:               aws.String(name),
		Description:        aws.String(description),
		Architecture:       aws.String("x86_64"),
		VirtualizationType: aws.String(string(imageType)),
		RootDeviceName:     aws.String(fmt.Sprintf("/dev/%sa", diskBaseName)),
		BlockDeviceMappings: []*ec2.BlockDeviceMapping{
			&ec2.BlockDeviceMapping{
				DeviceName: aws.String(fmt.Sprintf("/dev/%sa", diskBaseName)),
				Ebs: &ec2.EbsBlockDevice{
					SnapshotId:          aws.String(snapshotID),
					DeleteOnTermination: aws.Bool(true),
					VolumeSize:          aws.Int64(diskSize),
				},
			},
			&ec2.BlockDeviceMapping{
				DeviceName:  aws.String(fmt.Sprintf("/dev/%sb", diskBaseName)),
				VirtualName: aws.String("ephemeral0"),
			},
		},
	}
}
<|endoftext|>


<commit_before>// Package trickle allows to build trickle DAGs.
// In this type of DAG, non-leave nodes are first filled
// with data leaves, and then incorporate "layers" of subtrees
// as additional links.
//
// Each layer is a trickle sub-tree and is limited by an increasing
// maxinum depth. Thus, the nodes first layer
// can only hold leaves (depth 1) but subsequent layers can grow deeper.
// By default, this module places 4 nodes per layer (that is, 4 subtrees
// of the same maxinum depth before increasing it).
//
// Trickle DAGs are very good for sequentially reading data, as the
// first data leaves are directly reachable from the root and those
// coming next are always nearby. They are
// suited for things like streaming applications.
package trickle

import (
	"context"
	"errors"
	"fmt"

	h "github.com/ipfs/go-ipfs/importer/helpers"
	dag "github.com/ipfs/go-ipfs/merkledag"
	ft "github.com/ipfs/go-ipfs/unixfs"

	cid "gx/ipfs/QmcZfnkapfECQGcLZaf9B79NRg7cRa9EnZh4LSbkCzwNvY/go-cid"
	ipld "gx/ipfs/Qme5bWv7wtjUNGsK2BNGVUFPKiuxWrsqrtvYwCLRw8YFES/go-ipld-format"
)

// layerRepeat specifies how many times to append a child tree of a
// given depth. Higher values increase the width of a given node, which
// improves seek speeds.
const layerRepeat = 4

// Layout builds a new DAG with the trickle format using the provided
// DagBuilderHelper. See the module's description for a more detailed
// explanation.
func Layout(db *h.DagBuilderHelper) (ipld.Node, error) {
	root := db.NewUnixfsNode()
	if err := db.FillNodeLayer(root); err != nil {
		return nil, err
	}
	for level := 1; !db.Done(); level++ {
		for i := 0; i < layerRepeat && !db.Done(); i++ {
			next := db.NewUnixfsNode()
			if err := fillTrickleRec(db, next, level); err != nil {
				return nil, err
			}
			if err := root.AddChild(next, db); err != nil {
				return nil, err
			}
		}
	}

	out, err := db.Add(root)
	if err != nil {
		return nil, err
	}

	if err := db.Close(); err != nil {
		return nil, err
	}

	return out, nil
}

func fillTrickleRec(db *h.DagBuilderHelper, node *h.UnixfsNode, depth int) error {
	// Always do this, even in the base case
	if err := db.FillNodeLayer(node); err != nil {
		return err
	}

	for i := 1; i < depth && !db.Done(); i++ {
		for j := 0; j < layerRepeat && !db.Done(); j++ {
			next := db.NewUnixfsNode()
			if err := fillTrickleRec(db, next, i); err != nil {
				return err
			}

			if err := node.AddChild(next, db); err != nil {
				return err
			}
		}
	}
	return nil
}

// Append appends the data in `db` to the dag, using the Trickledag format
func Append(ctx context.Context, basen ipld.Node, db *h.DagBuilderHelper) (out ipld.Node, errOut error) {
	base, ok := basen.(*dag.ProtoNode)
	if !ok {
		return nil, dag.ErrNotProtobuf
	}

	defer func() {
		if errOut == nil {
			if err := db.Close(); err != nil {
				errOut = err
			}
		}
	}()

	// Convert to unixfs node for working with easily
	ufsn, err := h.NewUnixfsNodeFromDag(base)
	if err != nil {
		return nil, err
	}

	// Get depth of this 'tree'
	n, layerProgress := trickleDepthInfo(ufsn, db.Maxlinks())
	if n == 0 {
		// If direct blocks not filled...
		if err := db.FillNodeLayer(ufsn); err != nil {
			return nil, err
		}

		if db.Done() {
			return ufsn.GetDagNode()
		}

		// If continuing, our depth has increased by one
		n++
	}

	// Last child in this node may not be a full tree, lets file it up
	if err := appendFillLastChild(ctx, ufsn, n-1, layerProgress, db); err != nil {
		return nil, err
	}

	// after appendFillLastChild, our depth is now increased by one
	if !db.Done() {
		n++
	}

	// Now, continue filling out tree like normal
	for i := n; !db.Done(); i++ {
		for j := 0; j < layerRepeat && !db.Done(); j++ {
			next := db.NewUnixfsNode()
			err := fillTrickleRec(db, next, i)
			if err != nil {
				return nil, err
			}

			err = ufsn.AddChild(next, db)
			if err != nil {
				return nil, err
			}
		}
	}

	return ufsn.GetDagNode()
}

// appendFillLastChild will take in an incomplete trickledag node (uncomplete meaning, not full) and
// fill it out to the specified depth with blocks from the given DagBuilderHelper
func appendFillLastChild(ctx context.Context, ufsn *h.UnixfsNode, depth int, layerFill int, db *h.DagBuilderHelper) error {
	if ufsn.NumChildren() <= db.Maxlinks() {
		return nil
	}
	// Recursive step, grab last child
	last := ufsn.NumChildren() - 1
	lastChild, err := ufsn.GetChild(ctx, last, db.GetDagServ())
	if err != nil {
		return err
	}

	// Fill out last child (may not be full tree)
	nchild, err := appendRec(ctx, lastChild, db, depth-1)
	if err != nil {
		return err
	}

	// Update changed child in parent node
	ufsn.RemoveChild(last, db)
	err = ufsn.AddChild(nchild, db)
	if err != nil {
		return err
	}

	// Partially filled depth layer
	if layerFill != 0 {
		for ; layerFill < layerRepeat && !db.Done(); layerFill++ {
			next := db.NewUnixfsNode()
			err := fillTrickleRec(db, next, depth)
			if err != nil {
				return err
			}

			err = ufsn.AddChild(next, db)
			if err != nil {
				return err
			}
		}
	}

	return nil
}

// recursive call for Append
func appendRec(ctx context.Context, ufsn *h.UnixfsNode, db *h.DagBuilderHelper, depth int) (*h.UnixfsNode, error) {
	if depth == 0 || db.Done() {
		return ufsn, nil
	}

	// Get depth of this 'tree'
	n, layerProgress := trickleDepthInfo(ufsn, db.Maxlinks())
	if n == 0 {
		// If direct blocks not filled...
		if err := db.FillNodeLayer(ufsn); err != nil {
			return nil, err
		}
		n++
	}

	// If at correct depth, no need to continue
	if n == depth {
		return ufsn, nil
	}

	if err := appendFillLastChild(ctx, ufsn, n, layerProgress, db); err != nil {
		return nil, err
	}

	// after appendFillLastChild, our depth is now increased by one
	if !db.Done() {
		n++
	}

	// Now, continue filling out tree like normal
	for i := n; i < depth && !db.Done(); i++ {
		for j := 0; j < layerRepeat && !db.Done(); j++ {
			next := db.NewUnixfsNode()
			if err := fillTrickleRec(db, next, i); err != nil {
				return nil, err
			}

			if err := ufsn.AddChild(next, db); err != nil {
				return nil, err
			}
		}
	}

	return ufsn, nil
}

func trickleDepthInfo(node *h.UnixfsNode, maxlinks int) (int, int) {
	n := node.NumChildren()
	if n < maxlinks {
		return 0, 0
	}

	return ((n - maxlinks) / layerRepeat) + 1, (n - maxlinks) % layerRepeat
}

// VerifyParams is used by VerifyTrickleDagStructure
type VerifyParams struct {
	Getter      ipld.NodeGetter
	Direct      int
	LayerRepeat int
	Prefix      *cid.Prefix
	RawLeaves   bool
}

// VerifyTrickleDagStructure checks that the given dag matches exactly the trickle dag datastructure
// layout
func VerifyTrickleDagStructure(nd ipld.Node, p VerifyParams) error {
	return verifyTDagRec(nd, -1, p)
}

// Recursive call for verifying the structure of a trickledag
func verifyTDagRec(n ipld.Node, depth int, p VerifyParams) error {
	codec := cid.DagProtobuf
	if depth == 0 {
		if len(n.Links()) > 0 {
			return errors.New("expected direct block")
		}
		// zero depth dag is raw data block
		switch nd := n.(type) {
		case *dag.ProtoNode:
			pbn, err := ft.FromBytes(nd.Data())
			if err != nil {
				return err
			}

			if pbn.GetType() != ft.TRaw {
				return errors.New("Expected raw block")
			}

			if p.RawLeaves {
				return errors.New("expected raw leaf, got a protobuf node")
			}
		case *dag.RawNode:
			if !p.RawLeaves {
				return errors.New("expected protobuf node as leaf")
			}
			codec = cid.Raw
		default:
			return errors.New("expected ProtoNode or RawNode")
		}
	}

	// verify prefix
	if p.Prefix != nil {
		prefix := n.Cid().Prefix()
		expect := *p.Prefix // make a copy
		expect.Codec = uint64(codec)
		if codec == cid.Raw && expect.Version == 0 {
			expect.Version = 1
		}
		if expect.MhLength == -1 {
			expect.MhLength = prefix.MhLength
		}
		if prefix != expect {
			return fmt.Errorf("unexpected cid prefix: expected: %v; got %v", expect, prefix)
		}
	}

	if depth == 0 {
		return nil
	}

	nd, ok := n.(*dag.ProtoNode)
	if !ok {
		return errors.New("expected ProtoNode")
	}

	// Verify this is a branch node
	pbn, err := ft.FromBytes(nd.Data())
	if err != nil {
		return err
	}

	if pbn.GetType() != ft.TFile {
		return fmt.Errorf("expected file as branch node, got: %s", pbn.GetType())
	}

	if len(pbn.Data) > 0 {
		return errors.New("branch node should not have data")
	}

	for i := 0; i < len(nd.Links()); i++ {
		child, err := nd.Links()[i].GetNode(context.TODO(), p.Getter)
		if err != nil {
			return err
		}

		if i < p.Direct {
			// Direct blocks
			err := verifyTDagRec(child, 0, p)
			if err != nil {
				return err
			}
		} else {
			// Recursive trickle dags
			rdepth := ((i - p.Direct) / p.LayerRepeat) + 1
			if rdepth >= depth && depth > 0 {
				return errors.New("child dag was too deep")
			}
			err := verifyTDagRec(child, rdepth, p)
			if err != nil {
				return err
			}
		}
	}
	return nil
}
<commit_msg>dag: refactor trickle importer<commit_after>// Package trickle allows to build trickle DAGs.
// In this type of DAG, non-leave nodes are first filled
// with data leaves, and then incorporate "layers" of subtrees
// as additional links.
//
// Each layer is a trickle sub-tree and is limited by an increasing
// maxinum depth. Thus, the nodes first layer
// can only hold leaves (depth 1) but subsequent layers can grow deeper.
// By default, this module places 4 nodes per layer (that is, 4 subtrees
// of the same maxinum depth before increasing it).
//
// Trickle DAGs are very good for sequentially reading data, as the
// first data leaves are directly reachable from the root and those
// coming next are always nearby. They are
// suited for things like streaming applications.
package trickle

import (
	"context"
	"errors"
	"fmt"

	h "github.com/ipfs/go-ipfs/importer/helpers"
	dag "github.com/ipfs/go-ipfs/merkledag"
	ft "github.com/ipfs/go-ipfs/unixfs"

	cid "gx/ipfs/QmcZfnkapfECQGcLZaf9B79NRg7cRa9EnZh4LSbkCzwNvY/go-cid"
	ipld "gx/ipfs/Qme5bWv7wtjUNGsK2BNGVUFPKiuxWrsqrtvYwCLRw8YFES/go-ipld-format"
)

// layerRepeat specifies how many times to append a child tree of a
// given depth. Higher values increase the width of a given node, which
// improves seek speeds.
const layerRepeat = 4

// Layout builds a new DAG with the trickle format using the provided
// DagBuilderHelper. See the module's description for a more detailed
// explanation.
func Layout(db *h.DagBuilderHelper) (ipld.Node, error) {
	root := db.NewUnixfsNode()
	if err := fillTrickleRec(db, root, -1); err != nil {
		return nil, err
	}

	out, err := db.Add(root)
	if err != nil {
		return nil, err
	}

	if err := db.Close(); err != nil {
		return nil, err
	}

	return out, nil
}

// fillTrickleRec creates a trickle (sub-)tree with an optional maximum specified depth
// in the case maxDepth is greater than zero, or with unlimited depth otherwise
// (where the DAG builder will signal the end of data to end the function).
func fillTrickleRec(db *h.DagBuilderHelper, node *h.UnixfsNode, maxDepth int) error {
	// Always do this, even in the base case
	if err := db.FillNodeLayer(node); err != nil {
		return err
	}

	for depth := 1; ; depth++ {
		// Apply depth limit only if the parameter is set (> 0).
		if maxDepth > 0 && depth == maxDepth {
			return nil
		}
		for layer := 0; layer < layerRepeat; layer++ {
			if db.Done() {
				return nil
			}

			nextChild := db.NewUnixfsNode()
			if err := fillTrickleRec(db, nextChild, depth); err != nil {
				return err
			}

			if err := node.AddChild(nextChild, db); err != nil {
				return err
			}
		}
	}
}

// Append appends the data in `db` to the dag, using the Trickledag format
func Append(ctx context.Context, basen ipld.Node, db *h.DagBuilderHelper) (out ipld.Node, errOut error) {
	base, ok := basen.(*dag.ProtoNode)
	if !ok {
		return nil, dag.ErrNotProtobuf
	}

	defer func() {
		if errOut == nil {
			if err := db.Close(); err != nil {
				errOut = err
			}
		}
	}()

	// Convert to unixfs node for working with easily
	ufsn, err := h.NewUnixfsNodeFromDag(base)
	if err != nil {
		return nil, err
	}

	// Get depth of this 'tree'
	n, layerProgress := trickleDepthInfo(ufsn, db.Maxlinks())
	if n == 0 {
		// If direct blocks not filled...
		if err := db.FillNodeLayer(ufsn); err != nil {
			return nil, err
		}

		if db.Done() {
			return ufsn.GetDagNode()
		}

		// If continuing, our depth has increased by one
		n++
	}

	// Last child in this node may not be a full tree, lets file it up
	if err := appendFillLastChild(ctx, ufsn, n-1, layerProgress, db); err != nil {
		return nil, err
	}

	// after appendFillLastChild, our depth is now increased by one
	if !db.Done() {
		n++
	}

	// Now, continue filling out tree like normal
	for i := n; !db.Done(); i++ {
		for j := 0; j < layerRepeat && !db.Done(); j++ {
			next := db.NewUnixfsNode()
			err := fillTrickleRec(db, next, i)
			if err != nil {
				return nil, err
			}

			err = ufsn.AddChild(next, db)
			if err != nil {
				return nil, err
			}
		}
	}

	return ufsn.GetDagNode()
}

// appendFillLastChild will take in an incomplete trickledag node (uncomplete meaning, not full) and
// fill it out to the specified depth with blocks from the given DagBuilderHelper
func appendFillLastChild(ctx context.Context, ufsn *h.UnixfsNode, depth int, layerFill int, db *h.DagBuilderHelper) error {
	if ufsn.NumChildren() <= db.Maxlinks() {
		return nil
	}
	// Recursive step, grab last child
	last := ufsn.NumChildren() - 1
	lastChild, err := ufsn.GetChild(ctx, last, db.GetDagServ())
	if err != nil {
		return err
	}

	// Fill out last child (may not be full tree)
	nchild, err := appendRec(ctx, lastChild, db, depth-1)
	if err != nil {
		return err
	}

	// Update changed child in parent node
	ufsn.RemoveChild(last, db)
	err = ufsn.AddChild(nchild, db)
	if err != nil {
		return err
	}

	// Partially filled depth layer
	if layerFill != 0 {
		for ; layerFill < layerRepeat && !db.Done(); layerFill++ {
			next := db.NewUnixfsNode()
			err := fillTrickleRec(db, next, depth)
			if err != nil {
				return err
			}

			err = ufsn.AddChild(next, db)
			if err != nil {
				return err
			}
		}
	}

	return nil
}

// recursive call for Append
func appendRec(ctx context.Context, ufsn *h.UnixfsNode, db *h.DagBuilderHelper, depth int) (*h.UnixfsNode, error) {
	if depth == 0 || db.Done() {
		return ufsn, nil
	}

	// Get depth of this 'tree'
	n, layerProgress := trickleDepthInfo(ufsn, db.Maxlinks())
	if n == 0 {
		// If direct blocks not filled...
		if err := db.FillNodeLayer(ufsn); err != nil {
			return nil, err
		}
		n++
	}

	// If at correct depth, no need to continue
	if n == depth {
		return ufsn, nil
	}

	if err := appendFillLastChild(ctx, ufsn, n, layerProgress, db); err != nil {
		return nil, err
	}

	// after appendFillLastChild, our depth is now increased by one
	if !db.Done() {
		n++
	}

	// Now, continue filling out tree like normal
	for i := n; i < depth && !db.Done(); i++ {
		for j := 0; j < layerRepeat && !db.Done(); j++ {
			next := db.NewUnixfsNode()
			if err := fillTrickleRec(db, next, i); err != nil {
				return nil, err
			}

			if err := ufsn.AddChild(next, db); err != nil {
				return nil, err
			}
		}
	}

	return ufsn, nil
}

func trickleDepthInfo(node *h.UnixfsNode, maxlinks int) (int, int) {
	n := node.NumChildren()
	if n < maxlinks {
		return 0, 0
	}

	return ((n - maxlinks) / layerRepeat) + 1, (n - maxlinks) % layerRepeat
}

// VerifyParams is used by VerifyTrickleDagStructure
type VerifyParams struct {
	Getter      ipld.NodeGetter
	Direct      int
	LayerRepeat int
	Prefix      *cid.Prefix
	RawLeaves   bool
}

// VerifyTrickleDagStructure checks that the given dag matches exactly the trickle dag datastructure
// layout
func VerifyTrickleDagStructure(nd ipld.Node, p VerifyParams) error {
	return verifyTDagRec(nd, -1, p)
}

// Recursive call for verifying the structure of a trickledag
func verifyTDagRec(n ipld.Node, depth int, p VerifyParams) error {
	codec := cid.DagProtobuf
	if depth == 0 {
		if len(n.Links()) > 0 {
			return errors.New("expected direct block")
		}
		// zero depth dag is raw data block
		switch nd := n.(type) {
		case *dag.ProtoNode:
			pbn, err := ft.FromBytes(nd.Data())
			if err != nil {
				return err
			}

			if pbn.GetType() != ft.TRaw {
				return errors.New("Expected raw block")
			}

			if p.RawLeaves {
				return errors.New("expected raw leaf, got a protobuf node")
			}
		case *dag.RawNode:
			if !p.RawLeaves {
				return errors.New("expected protobuf node as leaf")
			}
			codec = cid.Raw
		default:
			return errors.New("expected ProtoNode or RawNode")
		}
	}

	// verify prefix
	if p.Prefix != nil {
		prefix := n.Cid().Prefix()
		expect := *p.Prefix // make a copy
		expect.Codec = uint64(codec)
		if codec == cid.Raw && expect.Version == 0 {
			expect.Version = 1
		}
		if expect.MhLength == -1 {
			expect.MhLength = prefix.MhLength
		}
		if prefix != expect {
			return fmt.Errorf("unexpected cid prefix: expected: %v; got %v", expect, prefix)
		}
	}

	if depth == 0 {
		return nil
	}

	nd, ok := n.(*dag.ProtoNode)
	if !ok {
		return errors.New("expected ProtoNode")
	}

	// Verify this is a branch node
	pbn, err := ft.FromBytes(nd.Data())
	if err != nil {
		return err
	}

	if pbn.GetType() != ft.TFile {
		return fmt.Errorf("expected file as branch node, got: %s", pbn.GetType())
	}

	if len(pbn.Data) > 0 {
		return errors.New("branch node should not have data")
	}

	for i := 0; i < len(nd.Links()); i++ {
		child, err := nd.Links()[i].GetNode(context.TODO(), p.Getter)
		if err != nil {
			return err
		}

		if i < p.Direct {
			// Direct blocks
			err := verifyTDagRec(child, 0, p)
			if err != nil {
				return err
			}
		} else {
			// Recursive trickle dags
			rdepth := ((i - p.Direct) / p.LayerRepeat) + 1
			if rdepth >= depth && depth > 0 {
				return errors.New("child dag was too deep")
			}
			err := verifyTDagRec(child, rdepth, p)
			if err != nil {
				return err
			}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package dsl_test

import (
	"github.com/goadesign/gorma"
	gdsl "github.com/goadesign/gorma/dsl"

	. "github.com/goadesign/goa/design"
	. "github.com/goadesign/goa/design/apidsl"
	. "github.com/goadesign/goa/dslengine"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("RelationalField", func() {
	var sgname, storename, modelname, name string
	var ft gorma.FieldType
	var dsl func()
	var RandomPayload *UserTypeDefinition
	BeforeEach(func() {
		Design = nil
		Errors = nil
		sgname = "production"
		dsl = nil
		storename = "mysql"
		modelname = "Users"
		name = ""
		ft = gorma.String
		gorma.GormaDesign = nil
		InitDesign()
		RandomPayload = Type("RandomPayload", func() {
			Attribute("first_name", String)
			Attribute("last_name", String)
		})

	})

	JustBeforeEach(func() {
		gdsl.StorageGroup(sgname, func() {
			gdsl.Store(storename, gorma.MySQL, func() {
				gdsl.Model(modelname, func() {
					//gdsl.BuildsFrom(RandomPayload)
					gdsl.Field(name, ft, dsl)
					gdsl.Field("id", gorma.Integer, dsl) // use lowercase "id" to test sanitizer
					gdsl.Field("MiddleName", gorma.String)
					gdsl.Field("CreatedAt", gorma.Timestamp)
					gdsl.Field("UpdatedAt", gorma.Timestamp)
					gdsl.Field("DeletedAt", gorma.NullableTimestamp)
				})
			})
		})
		Run()

	})

	Context("with no DSL", func() {
		BeforeEach(func() {
			name = "FirstName"
		})

		It("produces a valid Relational Field definition", func() {
			Ω(Design.Validate()).ShouldNot(HaveOccurred())
			sg := gorma.GormaDesign
			rs := sg.RelationalStores[storename]
			rm := rs.RelationalModels[modelname]
			Ω(rm.RelationalFields[name].FieldName).Should(Equal(name))
		})
	})

	Context("with an already defined Relational Field  with the same name", func() {
		BeforeEach(func() {
			name = "FirstName"
		})

		It("does not produce an error", func() {
			gdsl.StorageGroup(sgname, func() {
				gdsl.Store(storename, gorma.MySQL, func() {
					gdsl.Model(modelname, func() {
						gdsl.Field(name, ft, dsl)
					})
				})
			})
			Ω(Errors).Should(HaveOccurred())
		})
	})

	Context("with valid DSL", func() {
		JustBeforeEach(func() {
			Ω(Errors).ShouldNot(HaveOccurred())
			Ω(Design.Validate()).ShouldNot(HaveOccurred())
		})

		Context("sets appropriate fields and flags", func() {
			const description = "description"

			BeforeEach(func() {
				name = "FirstName"
				dsl = func() {
					gdsl.Description(description)
				}
			})

			It("sanitizes the ID field", func() {
				sg := gorma.GormaDesign
				rs := sg.RelationalStores[storename]
				rm := rs.RelationalModels[modelname]
				Ω(rm.RelationalFields).Should(HaveKey("ID"))
			})

			It("sets the relational field description", func() {
				sg := gorma.GormaDesign
				rs := sg.RelationalStores[storename]
				rm := rs.RelationalModels[modelname]
				Ω(rm.RelationalFields[name].Description).Should(Equal(description))
			})

			It("sets the field name", func() {
				sg := gorma.GormaDesign
				rs := sg.RelationalStores[storename]
				rm := rs.RelationalModels[modelname]
				Ω(rm.RelationalFields["ID"].FieldName).Should(Equal("ID"))
			})
			It("sets the field type", func() {
				sg := gorma.GormaDesign
				rs := sg.RelationalStores[storename]
				rm := rs.RelationalModels[modelname]
				Ω(rm.RelationalFields["ID"].Datatype).Should(Equal(gorma.Integer))
			})
			It("sets the pk flag", func() {
				sg := gorma.GormaDesign
				rs := sg.RelationalStores[storename]
				rm := rs.RelationalModels[modelname]
				Ω(rm.RelationalFields["ID"].PrimaryKey).Should(Equal(true))
			})
			It("sets has a created at field", func() {
				sg := gorma.GormaDesign
				rs := sg.RelationalStores[storename]
				rm := rs.RelationalModels[modelname]
				Ω(rm.RelationalFields["CreatedAt"].FieldName).Should(Equal("CreatedAt"))
				Ω(rm.RelationalFields["CreatedAt"].Datatype).Should(Equal(gorma.Timestamp))
				Ω(rm.RelationalFields["CreatedAt"].Nullable).Should(Equal(false))
			})
			It("sets has a deleted at field", func() {
				sg := gorma.GormaDesign
				rs := sg.RelationalStores[storename]
				rm := rs.RelationalModels[modelname]
				Ω(rm.RelationalFields["DeletedAt"].FieldName).Should(Equal("DeletedAt"))
				Ω(rm.RelationalFields["DeletedAt"].Datatype).Should(Equal(gorma.NullableTimestamp))
				Ω(rm.RelationalFields["DeletedAt"].Nullable).Should(Equal(true))
			})
			It("has the right number of fields", func() {
				sg := gorma.GormaDesign
				rs := sg.RelationalStores[storename]
				rm := rs.RelationalModels[modelname]
				length := len(rm.RelationalFields)
				Ω(length).Should(Equal(6))
			})
		})

	})

})
<commit_msg>added pos and neg tests for PrimaryKey()<commit_after>package dsl_test

import (
	"github.com/goadesign/gorma"
	gdsl "github.com/goadesign/gorma/dsl"

	. "github.com/goadesign/goa/design"
	. "github.com/goadesign/goa/design/apidsl"
	. "github.com/goadesign/goa/dslengine"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("RelationalField", func() {
	var sgname, storename, modelname, name string
	var ft gorma.FieldType
	var dsl func()
	var RandomPayload *UserTypeDefinition
	BeforeEach(func() {
		Design = nil
		Errors = nil
		sgname = "production"
		dsl = nil
		storename = "mysql"
		modelname = "Users"
		name = ""
		ft = gorma.String
		gorma.GormaDesign = nil
		InitDesign()
		RandomPayload = Type("RandomPayload", func() {
			Attribute("first_name", String)
			Attribute("last_name", String)
		})

	})

	JustBeforeEach(func() {
		gdsl.StorageGroup(sgname, func() {
			gdsl.Store(storename, gorma.MySQL, func() {
				gdsl.Model(modelname, func() {
					//gdsl.BuildsFrom(RandomPayload)
					gdsl.Field(name, ft, dsl)
					gdsl.Field("id", gorma.Integer, dsl) // use lowercase "id" to test sanitizer
					gdsl.Field("MiddleName", gorma.String)
					gdsl.Field("CreatedAt", gorma.Timestamp)
					gdsl.Field("UpdatedAt", gorma.Timestamp)
					gdsl.Field("DeletedAt", gorma.NullableTimestamp)
				})
			})
		})
		Run()

	})

	Context("with no DSL", func() {
		BeforeEach(func() {
			name = "FirstName"
		})

		It("produces a valid Relational Field definition", func() {
			Ω(Design.Validate()).ShouldNot(HaveOccurred())
			sg := gorma.GormaDesign
			rs := sg.RelationalStores[storename]
			rm := rs.RelationalModels[modelname]
			Ω(rm.RelationalFields[name].FieldName).Should(Equal(name))
		})
	})

	Context("with an already defined Relational Field  with the same name", func() {
		BeforeEach(func() {
			name = "FirstName"
		})

		It("does not produce an error", func() {
			gdsl.StorageGroup(sgname, func() {
				gdsl.Store(storename, gorma.MySQL, func() {
					gdsl.Model(modelname, func() {
						gdsl.Field(name, ft, dsl)
					})
				})
			})
			Ω(Errors).Should(HaveOccurred())
		})
	})

	Context("with valid DSL", func() {
		JustBeforeEach(func() {
			Ω(Errors).ShouldNot(HaveOccurred())
			Ω(Design.Validate()).ShouldNot(HaveOccurred())
		})

		Context("sets appropriate fields and flags", func() {
			const description = "description"

			BeforeEach(func() {
				name = "FirstName"
				dsl = func() {
					gdsl.Description(description)
				}
			})

			It("sanitizes the ID field", func() {
				sg := gorma.GormaDesign
				rs := sg.RelationalStores[storename]
				rm := rs.RelationalModels[modelname]
				Ω(rm.RelationalFields).Should(HaveKey("ID"))
			})

			It("sets the relational field description", func() {
				sg := gorma.GormaDesign
				rs := sg.RelationalStores[storename]
				rm := rs.RelationalModels[modelname]
				Ω(rm.RelationalFields[name].Description).Should(Equal(description))
			})

			It("sets the field name", func() {
				sg := gorma.GormaDesign
				rs := sg.RelationalStores[storename]
				rm := rs.RelationalModels[modelname]
				Ω(rm.RelationalFields["ID"].FieldName).Should(Equal("ID"))
			})
			It("sets the field type", func() {
				sg := gorma.GormaDesign
				rs := sg.RelationalStores[storename]
				rm := rs.RelationalModels[modelname]
				Ω(rm.RelationalFields["ID"].Datatype).Should(Equal(gorma.Integer))
			})
			It("sets the pk flag", func() {
				sg := gorma.GormaDesign
				rs := sg.RelationalStores[storename]
				rm := rs.RelationalModels[modelname]
				Ω(rm.RelationalFields["ID"].PrimaryKey).Should(Equal(true))
			})
			It("sets has a created at field", func() {
				sg := gorma.GormaDesign
				rs := sg.RelationalStores[storename]
				rm := rs.RelationalModels[modelname]
				Ω(rm.RelationalFields["CreatedAt"].FieldName).Should(Equal("CreatedAt"))
				Ω(rm.RelationalFields["CreatedAt"].Datatype).Should(Equal(gorma.Timestamp))
				Ω(rm.RelationalFields["CreatedAt"].Nullable).Should(Equal(false))
			})
			It("sets has a deleted at field", func() {
				sg := gorma.GormaDesign
				rs := sg.RelationalStores[storename]
				rm := rs.RelationalModels[modelname]
				Ω(rm.RelationalFields["DeletedAt"].FieldName).Should(Equal("DeletedAt"))
				Ω(rm.RelationalFields["DeletedAt"].Datatype).Should(Equal(gorma.NullableTimestamp))
				Ω(rm.RelationalFields["DeletedAt"].Nullable).Should(Equal(true))
			})
			It("has the right number of fields", func() {
				sg := gorma.GormaDesign
				rs := sg.RelationalStores[storename]
				rm := rs.RelationalModels[modelname]
				length := len(rm.RelationalFields)
				Ω(length).Should(Equal(6))
			})
		})

	})

	Context("Primary Keys", func() {
		JustBeforeEach(func() {

			gdsl.StorageGroup(sgname, func() {
				gdsl.Store(storename, gorma.MySQL, func() {
					gdsl.NoAutomaticIDFields()
					gdsl.Model(modelname, func() {
						//gdsl.BuildsFrom(RandomPayload)
						gdsl.Field(name, ft, dsl)
						gdsl.Field("id", gorma.Integer, dsl) // use lowercase "id" to test sanitizer
						gdsl.Field("MiddleName", gorma.String)
						gdsl.Field("CreatedAt", gorma.Timestamp)
						gdsl.Field("UpdatedAt", gorma.Timestamp)
						gdsl.Field("DeletedAt", gorma.NullableTimestamp)
					})
				})
			})
			Run()

			Ω(Errors).ShouldNot(HaveOccurred())
			Ω(Design.Validate()).ShouldNot(HaveOccurred())
		})

		Context("sets Primary Key flags", func() {

			BeforeEach(func() {
				name = "random"
				dsl = func() {
					gdsl.PrimaryKey()
				}
			})
			It("sets the pk flag", func() {
				sg := gorma.GormaDesign
				rs := sg.RelationalStores[storename]
				rm := rs.RelationalModels[modelname]
				Ω(rm.RelationalFields["ID"].PrimaryKey).Should(Equal(true))
			})
		})
		Context("doesn't set Primary Key flag", func() {

			BeforeEach(func() {
				name = "random"
				dsl = func() {
				}
			})
			It("the pk flag", func() {
				sg := gorma.GormaDesign
				rs := sg.RelationalStores[storename]
				rm := rs.RelationalModels[modelname]
				Ω(rm.RelationalFields["ID"].PrimaryKey).Should(Equal(false))
			})
		})

	})

})
<|endoftext|>


<commit_before>package models_test

import (
	"encoding/json"
	"fmt"
	// "github.com/APTrust/exchange/constants"
	"github.com/APTrust/exchange/models"
	"github.com/stretchr/testify/assert"
	// "os"
	// "path"
	"path/filepath"
	"runtime"
	"testing"
	// "time"
)

func TestNewVirtualBag(t *testing.T) {
	tarFilePath := vbagGetPath("example.edu.tagsample_good.tar")
	vbag := models.NewVirtualBag(tarFilePath, nil, false, false)
	assert.NotNil(t, vbag)
}

func TestVirtualBagRead_FromDirectory(t *testing.T) {

}

func TestVirtualBagRead_FromTarFile(t *testing.T) {
	tarFilePath := vbagGetPath("example.edu.tagsample_good.tar")
	files := []string {"bagit.txt", "bag-info.txt", "aptrust-info.txt"}
	vbag := models.NewVirtualBag(tarFilePath, files, true, true)
	assert.NotNil(t, vbag)
	obj, summary := vbag.Read()
	objJson, _ := json.Marshal(obj)
	summaryJson, _ := json.Marshal(summary)
	fmt.Println(string(objJson))
	fmt.Println(string(summaryJson))
}

func TestVirtualBagRead_ChecksumOptions(t *testing.T) {

}

// With md5 manifest only, sha256 only, and both
func TestVirtualBagRead_ManifestOptions(t *testing.T) {

}

func vbagGetPath(fileName string) (string) {
	_, filename, _, _ := runtime.Caller(0)
	dir, _ := filepath.Abs(filepath.Dir(filename))
	testDataPath := filepath.Join(dir, "..", "testdata", fileName)
	return testDataPath
}
<commit_msg>Added TODO. Because I forget everything.<commit_after>package models_test

import (
	"encoding/json"
	"fmt"
	// "github.com/APTrust/exchange/constants"
	"github.com/APTrust/exchange/models"
	"github.com/stretchr/testify/assert"
	// "os"
	// "path"
	"path/filepath"
	"runtime"
	"testing"
	// "time"
)

func TestNewVirtualBag(t *testing.T) {
	tarFilePath := vbagGetPath("example.edu.tagsample_good.tar")
	vbag := models.NewVirtualBag(tarFilePath, nil, false, false)
	assert.NotNil(t, vbag)
}

func TestVirtualBagRead_FromDirectory(t *testing.T) {

}

// ----------------------------------------------------------------
// TODO: Update example.edu.tagsample_good.tar or create a new
// version of it to include 1) checksums for all tag files,
// 2) Description tag, 3) Internal-Sender-Identifier tag.
// ----------------------------------------------------------------
func TestVirtualBagRead_FromTarFile(t *testing.T) {
	tarFilePath := vbagGetPath("example.edu.tagsample_good.tar")
	files := []string {"bagit.txt", "bag-info.txt", "aptrust-info.txt"}
	vbag := models.NewVirtualBag(tarFilePath, files, true, true)
	assert.NotNil(t, vbag)
	obj, summary := vbag.Read()
	objJson, _ := json.Marshal(obj)
	summaryJson, _ := json.Marshal(summary)
	fmt.Println(string(objJson))
	fmt.Println(string(summaryJson))
}

func TestVirtualBagRead_ChecksumOptions(t *testing.T) {

}

// With md5 manifest only, sha256 only, and both
func TestVirtualBagRead_ManifestOptions(t *testing.T) {

}

func vbagGetPath(fileName string) (string) {
	_, filename, _, _ := runtime.Caller(0)
	dir, _ := filepath.Abs(filepath.Dir(filename))
	testDataPath := filepath.Join(dir, "..", "testdata", fileName)
	return testDataPath
}
<|endoftext|>


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

package provisioner_test

import (
	"fmt"
	"os"
	"path/filepath"

	. "launchpad.net/gocheck"
	"launchpad.net/juju-core/container"
	"launchpad.net/juju-core/container/lxc"
	"launchpad.net/juju-core/instance"
	jujutesting "launchpad.net/juju-core/juju/testing"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/testing"
	. "launchpad.net/juju-core/testing/checkers"
	"launchpad.net/juju-core/version"
)

type lxcBrokerSuite struct {
	testing.LoggingSuite
	containerDir       string
	removedDir         string
	lxcDir             string
	oldContainerDir    string
	oldRemovedDir      string
	oldLxcContainerDir string
}

var _ = Suite(&lxcBrokerSuite{})

func (s *lxcBrokerSuite) SetUpSuite(c *C) {
	s.LoggingSuite.SetUpSuite(c)
}

func (s *lxcBrokerSuite) TearDownSuite(c *C) {
	s.LoggingSuite.TearDownSuite(c)
}

func (s *lxcBrokerSuite) SetUpTest(c *C) {
	s.LoggingSuite.SetUpTest(c)
	s.containerDir = c.MkDir()
	s.oldContainerDir = lxc.SetContainerDir(s.containerDir)
	s.removedDir = c.MkDir()
	s.oldRemovedDir = lxc.SetRemovedContainerDir(s.removedDir)
	s.lxcDir = c.MkDir()
	s.oldLxcContainerDir = lxc.SetLxcContainerDir(s.lxcDir)
}

func (s *lxcBrokerSuite) TearDownTest(c *C) {
	lxc.SetContainerDir(s.oldContainerDir)
	lxc.SetLxcContainerDir(s.oldLxcContainerDir)
	lxc.SetRemovedContainerDir(s.oldRemovedDir)
	s.LoggingSuite.TearDownTest(c)
}

func (*lxcBrokerSuite) TestInstanceInterface(c *C) {

}
<commit_msg>Test start and stop isntance.<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package provisioner_test

import (
	. "launchpad.net/gocheck"
	"launchpad.net/golxc"
	"launchpad.net/juju-core/constraints"
	"launchpad.net/juju-core/container/lxc"
	"launchpad.net/juju-core/container/lxc/mock"
	"launchpad.net/juju-core/instance"
	jujutesting "launchpad.net/juju-core/juju/testing"
	"launchpad.net/juju-core/state"
	"launchpad.net/juju-core/testing"
	"launchpad.net/juju-core/version"
	"launchpad.net/juju-core/worker/provisioner"
)

type lxcBrokerSuite struct {
	testing.LoggingSuite
	golxc              golxc.ContainerFactory
	broker             provisioner.Broker
	containerDir       string
	removedDir         string
	lxcDir             string
	oldContainerDir    string
	oldRemovedDir      string
	oldLxcContainerDir string
}

var _ = Suite(&lxcBrokerSuite{})

func (s *lxcBrokerSuite) SetUpSuite(c *C) {
	s.LoggingSuite.SetUpSuite(c)
}

func (s *lxcBrokerSuite) TearDownSuite(c *C) {
	s.LoggingSuite.TearDownSuite(c)
}

func (s *lxcBrokerSuite) SetUpTest(c *C) {
	s.LoggingSuite.SetUpTest(c)
	s.containerDir = c.MkDir()
	s.oldContainerDir = lxc.SetContainerDir(s.containerDir)
	s.removedDir = c.MkDir()
	s.oldRemovedDir = lxc.SetRemovedContainerDir(s.removedDir)
	s.lxcDir = c.MkDir()
	s.oldLxcContainerDir = lxc.SetLxcContainerDir(s.lxcDir)
	s.golxc = mock.MockFactory()
	tools := &state.Tools{
		Binary: version.MustParseBinary("2.3.4-foo-bar"),
		URL:    "http://tools.example.com/2.3.4-foo-bar.tgz",
	}
	s.broker = provisioner.NewLxcBroker(s.golxc, testing.EnvironConfig(c), tools)
}

func (s *lxcBrokerSuite) TearDownTest(c *C) {
	lxc.SetContainerDir(s.oldContainerDir)
	lxc.SetLxcContainerDir(s.oldLxcContainerDir)
	lxc.SetRemovedContainerDir(s.oldRemovedDir)
	s.LoggingSuite.TearDownTest(c)
}

func (s *lxcBrokerSuite) startInstance(c *C, machineId string) instance.Instance {
	stateInfo := jujutesting.FakeStateInfo(machineId)
	apiInfo := jujutesting.FakeAPIInfo(machineId)

	series := "series"
	nonce := "fake-nonce"
	cons := constraints.Value{}
	lxc, err := s.broker.StartInstance(machineId, nonce, series, cons, stateInfo, apiInfo)
	c.Assert(err, IsNil)
	return lxc
}

func (s *lxcBrokerSuite) TestStartInstance(c *C) {
	machineId := "1/lxc/0"
	lxc := s.startInstance(c, machineId)
	c.Assert(lxc.Id(), Equals, instance.Id("machine-1-lxc-0"))
}

func (s *lxcBrokerSuite) TestStopInstance(c *C) {
	lxc0 := s.startInstance(c, "1/lxc/0")
	lxc1 := s.startInstance(c, "1/lxc/1")
	lxc2 := s.startInstance(c, "1/lxc/2")
	err := s.broker.StopInstances([]instance.Instance{lxc0, lxc1, lxc2})
	c.Assert(err, IsNil)
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Smudge Authors.

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

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

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

package smudge

import (
	"fmt"
	"net"
	"os"
	"regexp"
	"strconv"
	"strings"
)

// Provides a series of methods and constants that revolve around the getting
// (or programmatically setting/overriding) environmental properties, returning
// default values if not set.

const (
	// EnvVarHeartbeatMillis is the name of the environment variable that
	// sets the heartbeat frequency (in millis).
	EnvVarHeartbeatMillis = "SMUDGE_HEARTBEAT_MILLIS"

	// DefaultHeartbeatMillis is the default heartbeat frequency (in millis).
	DefaultHeartbeatMillis int = 250

	// EnvVarInitialHosts is the name of the environment variable that sets
	// the initial known hosts. The value it sets should be a comma-delimitted
	// string of one or more IP:PORT pairs (port is optional if it matched the
	// value of SMUDGE_LISTEN_PORT).
	EnvVarInitialHosts = "SMUDGE_INITIAL_HOSTS"

	// DefaultInitialHosts default lists of initially known hosts.
	DefaultInitialHosts string = ""

	// EnvVarListenPort is the name of the environment variable that sets
	// the UDP listen port.
	EnvVarListenPort = "SMUDGE_LISTEN_PORT"

	// DefaultListenPort is the default UDP listen port.
	DefaultListenPort int = 9999

	// EnvVarListenIP is the name of the environment variable that sets
	// the listen IP.
	EnvVarListenIP = "SMUDGE_LISTEN_IP"

	// DefaultListenIP is the default listen IP.
	DefaultListenIP = "127.0.0.1"

	// EnvVarMaxBroadcastBytes is the name of the environment variable that
	// the maximum byte length for broadcast payloads. Note that increasing
	// this runs the risk of packet fragmentation and dropped messages.
	EnvVarMaxBroadcastBytes = "SMUDGE_MAX_BROADCAST_BYTES"

	// DefaultMaxBroadcastBytes is the default maximum byte length for
	// broadcast payloads. This is guided by the maximum safe UDP packet size
	// of 508 bytes, which must also contain status updates and additional
	// message overhead.
	DefaultMaxBroadcastBytes int = 256

	// EnvVarMulticastEnabled is the name of the environment variable that
	// describes whether Smudge will attempt to announce its presence via
	// multicast on startup.
	EnvVarMulticastEnabled = "SMUDGE_MULTICAST_ENABLED"

	// DefaultMulticastEnabled is the default value for whether Smudge will
	// attempt to announce its presence via multicast on startup.
	DefaultMulticastEnabled string = "true"

	// EnvVarMulticastAddress is the name of the environment variable that
	// defines the multicast address that will be used.
	EnvVarMulticastAddress = "SMUDGE_MULTICAST_ADDRESS"

	// DefaultMulticastAddress is the default multicast address. Empty string
	// indicates 224.0.0.0 for IPv4 and [ff02::1] for IPv6.
	DefaultMulticastAddress string = ""
)

var heartbeatMillis int

var listenPort int

var listenIP net.IP

var initialHosts []string

var maxBroadcastBytes int

var multicastEnabledString string

var multicastEnabled bool

var multicastAddress string

const stringListDelimitRegex = "\\s*((,\\s*)|(\\s+))"

// GetHeartbeatMillis gets this host's heartbeat frequency in milliseconds.
func GetHeartbeatMillis() int {
	if heartbeatMillis == 0 {
		heartbeatMillis = getIntVar(EnvVarHeartbeatMillis, DefaultHeartbeatMillis)
	}

	return heartbeatMillis
}

// GetInitialHosts returns the list of initially known hosts.
func GetInitialHosts() []string {
	if initialHosts == nil {
		initialHosts = getStringArrayVar(EnvVarInitialHosts, DefaultInitialHosts)
	}

	return initialHosts
}

// GetListenPort returns the port that this host will listen on.
func GetListenPort() int {
	if listenPort == 0 {
		listenPort = getIntVar(EnvVarListenPort, DefaultListenPort)
	}

	return listenPort
}

// GetListenIP returns the IP that this host will listen on.
func GetListenIP() net.IP {
	if listenIP == nil {
		listenIP = net.ParseIP(getStringVar(EnvVarListenIP, DefaultListenIP))
	}

	return listenIP
}

// GetMaxBroadcastBytes returns the maximum byte length for broadcast payloads.
func GetMaxBroadcastBytes() int {
	if maxBroadcastBytes == 0 {
		maxBroadcastBytes = getIntVar(EnvVarMaxBroadcastBytes, DefaultMaxBroadcastBytes)
	}

	return maxBroadcastBytes
}

// GetMulticastEnabled returns whether multicast announcements are enabled.
func GetMulticastEnabled() bool {
	if multicastEnabledString == "" {
		multicastEnabledString = strings.ToLower(getStringVar(EnvVarMulticastEnabled, DefaultMulticastEnabled))
		multicastEnabled = len(multicastEnabledString) > 0 && multicastEnabledString[0] == "t"
	}

	return multicastEnabled
}

// GetMulticastAddress returns the address the will be used for multicast
// announcements.
func GetMulticastAddress() bool {
	if multicastAddress == "" {
		multicastAddress = getStringVar(EnvVarMulticastAddress, DefaultMulticastAddress)
	}

	return multicastAddress
}

// SetHeartbeatMillis sets this nodes heartbeat frequency. Unlike
// SetListenPort(), calling this function after Begin() has been called will
// have an effect.
func SetHeartbeatMillis(val int) {
	if val == 0 {
		heartbeatMillis = DefaultListenPort
	} else {
		heartbeatMillis = val
	}

	heartbeatMillis = val
}

// SetListenPort sets the UDP port to listen on. It has no effect once
// Begin() has been called.
func SetListenPort(val int) {
	if val == 0 {
		listenPort = DefaultListenPort
	} else {
		listenPort = val
	}
}

// SetListenIP sets the IP to listen on. It has no effect once
// Begin() has been called.
func SetListenIP(val net.IP) {
	if len(AllNodes()) > 0 {
		logWarn("Do not call SetListenIP() after nodes have been added, it may cause unexpected behavior.")
	}

	if val == nil {
		listenIP = net.ParseIP(DefaultListenIP)
	} else {
		listenIP = val
	}
}

// SetMaxBroadcastBytes sets the maximum byte length for broadcast payloads.
// Note that increasing this beyond the default of 256 runs the risk of packet
// fragmentation and dropped messages.
func SetMaxBroadcastBytes(val int) {
	if val == 0 {
		maxBroadcastBytes = DefaultMaxBroadcastBytes
	} else {
		maxBroadcastBytes = val
	}
}

// GetMulticastEnabled sets whether multicast announcements are enabled.
func SetMulticastEnabled(val bool) {
	multicastEnabledString = fmt.Sprintf("%v", val)
}

// SetMulticastAddress sets the address that will be used for multicast
// announcements.
func SetMulticastAddress(val string) {
	if val == "" {
		multicastAddress = DefaultMulticastAddress
	} else {
		multicastAddress = val
	}
}

// Gets an environmental variable "key". If it does not exist, "defaultVal" is
// returned; if it does, it attempts to convert to an integer, returning
// "defaultVal" if it fails.
func getIntVar(key string, defaultVal int) int {
	valueString := os.Getenv(key)
	valueInt := defaultVal

	if valueString != "" {
		i, err := strconv.Atoi(key)

		if err != nil {
			logfWarn("Failed to parse env property %s: %s is not "+
				"an integer. Using default.\n", key, valueString)
		} else {
			valueInt = i
		}
	}

	return valueInt
}

// Gets an environmental variable "key". If it does not exist, "defaultVal" is
// returned; if it does, it attempts to convert to a string slice, returning
// "defaultVal" if it fails.
func getStringArrayVar(key string, defaultVal string) []string {
	valueString := os.Getenv(key)

	if valueString == "" {
		valueString = defaultVal
	}

	valueSlice := splitDelimmitedString(valueString, stringListDelimitRegex)

	return valueSlice
}

// Gets an environmental variable "key". If it does not exist, "defaultVal" is
// returned; if it does, it attempts to convert to a string, returning
// "defaultVal" if it fails.
func getStringVar(key string, defaultVal string) string {
	valueString := os.Getenv(key)

	if valueString == "" {
		valueString = defaultVal
	}

	return valueString
}

// Splits a string on a regular expression.
func splitDelimmitedString(str string, regex string) []string {
	var result []string

	str = strings.TrimSpace(str)

	if str != "" {
		reg := regexp.MustCompile(regex)
		indices := reg.FindAllStringIndex(str, -1)

		result = make([]string, len(indices)+1)

		lastStart := 0
		for i, val := range indices {
			result[i] = str[lastStart:val[0]]
			lastStart = val[1]
		}

		result[len(indices)] = str[lastStart:]

		// Special case of single empty string
		if len(result) == 1 && result[0] == "" {
			result = make([]string, 0, 0)
		}
	}

	return result
}
<commit_msg>Fix silly slice cast<commit_after>/*
Copyright 2016 The Smudge Authors.

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

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

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

package smudge

import (
	"fmt"
	"net"
	"os"
	"regexp"
	"strconv"
	"strings"
)

// Provides a series of methods and constants that revolve around the getting
// (or programmatically setting/overriding) environmental properties, returning
// default values if not set.

const (
	// EnvVarHeartbeatMillis is the name of the environment variable that
	// sets the heartbeat frequency (in millis).
	EnvVarHeartbeatMillis = "SMUDGE_HEARTBEAT_MILLIS"

	// DefaultHeartbeatMillis is the default heartbeat frequency (in millis).
	DefaultHeartbeatMillis int = 250

	// EnvVarInitialHosts is the name of the environment variable that sets
	// the initial known hosts. The value it sets should be a comma-delimitted
	// string of one or more IP:PORT pairs (port is optional if it matched the
	// value of SMUDGE_LISTEN_PORT).
	EnvVarInitialHosts = "SMUDGE_INITIAL_HOSTS"

	// DefaultInitialHosts default lists of initially known hosts.
	DefaultInitialHosts string = ""

	// EnvVarListenPort is the name of the environment variable that sets
	// the UDP listen port.
	EnvVarListenPort = "SMUDGE_LISTEN_PORT"

	// DefaultListenPort is the default UDP listen port.
	DefaultListenPort int = 9999

	// EnvVarListenIP is the name of the environment variable that sets
	// the listen IP.
	EnvVarListenIP = "SMUDGE_LISTEN_IP"

	// DefaultListenIP is the default listen IP.
	DefaultListenIP = "127.0.0.1"

	// EnvVarMaxBroadcastBytes is the name of the environment variable that
	// the maximum byte length for broadcast payloads. Note that increasing
	// this runs the risk of packet fragmentation and dropped messages.
	EnvVarMaxBroadcastBytes = "SMUDGE_MAX_BROADCAST_BYTES"

	// DefaultMaxBroadcastBytes is the default maximum byte length for
	// broadcast payloads. This is guided by the maximum safe UDP packet size
	// of 508 bytes, which must also contain status updates and additional
	// message overhead.
	DefaultMaxBroadcastBytes int = 256

	// EnvVarMulticastEnabled is the name of the environment variable that
	// describes whether Smudge will attempt to announce its presence via
	// multicast on startup.
	EnvVarMulticastEnabled = "SMUDGE_MULTICAST_ENABLED"

	// DefaultMulticastEnabled is the default value for whether Smudge will
	// attempt to announce its presence via multicast on startup.
	DefaultMulticastEnabled string = "true"

	// EnvVarMulticastAddress is the name of the environment variable that
	// defines the multicast address that will be used.
	EnvVarMulticastAddress = "SMUDGE_MULTICAST_ADDRESS"

	// DefaultMulticastAddress is the default multicast address. Empty string
	// indicates 224.0.0.0 for IPv4 and [ff02::1] for IPv6.
	DefaultMulticastAddress string = ""
)

var heartbeatMillis int

var listenPort int

var listenIP net.IP

var initialHosts []string

var maxBroadcastBytes int

var multicastEnabledString string

var multicastEnabled bool = true

var multicastAddress string

const stringListDelimitRegex = "\\s*((,\\s*)|(\\s+))"

// GetHeartbeatMillis gets this host's heartbeat frequency in milliseconds.
func GetHeartbeatMillis() int {
	if heartbeatMillis == 0 {
		heartbeatMillis = getIntVar(EnvVarHeartbeatMillis, DefaultHeartbeatMillis)
	}

	return heartbeatMillis
}

// GetInitialHosts returns the list of initially known hosts.
func GetInitialHosts() []string {
	if initialHosts == nil {
		initialHosts = getStringArrayVar(EnvVarInitialHosts, DefaultInitialHosts)
	}

	return initialHosts
}

// GetListenPort returns the port that this host will listen on.
func GetListenPort() int {
	if listenPort == 0 {
		listenPort = getIntVar(EnvVarListenPort, DefaultListenPort)
	}

	return listenPort
}

// GetListenIP returns the IP that this host will listen on.
func GetListenIP() net.IP {
	if listenIP == nil {
		listenIP = net.ParseIP(getStringVar(EnvVarListenIP, DefaultListenIP))
	}

	return listenIP
}

// GetMaxBroadcastBytes returns the maximum byte length for broadcast payloads.
func GetMaxBroadcastBytes() int {
	if maxBroadcastBytes == 0 {
		maxBroadcastBytes = getIntVar(EnvVarMaxBroadcastBytes, DefaultMaxBroadcastBytes)
	}

	return maxBroadcastBytes
}

// GetMulticastEnabled returns whether multicast announcements are enabled.
func GetMulticastEnabled() bool {
	if multicastEnabledString == "" {
		multicastEnabledString = strings.ToLower(getStringVar(EnvVarMulticastEnabled, DefaultMulticastEnabled))
		multicastEnabled = len(multicastEnabledString) > 0 && []rune(multicastEnabledString)[0] == 't'
	}

	return multicastEnabled
}

// GetMulticastAddress returns the address the will be used for multicast
// announcements.
func GetMulticastAddress() string {
	if multicastAddress == "" {
		multicastAddress = getStringVar(EnvVarMulticastAddress, DefaultMulticastAddress)
	}

	return multicastAddress
}

// SetHeartbeatMillis sets this nodes heartbeat frequency. Unlike
// SetListenPort(), calling this function after Begin() has been called will
// have an effect.
func SetHeartbeatMillis(val int) {
	if val == 0 {
		heartbeatMillis = DefaultListenPort
	} else {
		heartbeatMillis = val
	}

	heartbeatMillis = val
}

// SetListenPort sets the UDP port to listen on. It has no effect once
// Begin() has been called.
func SetListenPort(val int) {
	if val == 0 {
		listenPort = DefaultListenPort
	} else {
		listenPort = val
	}
}

// SetListenIP sets the IP to listen on. It has no effect once
// Begin() has been called.
func SetListenIP(val net.IP) {
	if len(AllNodes()) > 0 {
		logWarn("Do not call SetListenIP() after nodes have been added, it may cause unexpected behavior.")
	}

	if val == nil {
		listenIP = net.ParseIP(DefaultListenIP)
	} else {
		listenIP = val
	}
}

// SetMaxBroadcastBytes sets the maximum byte length for broadcast payloads.
// Note that increasing this beyond the default of 256 runs the risk of packet
// fragmentation and dropped messages.
func SetMaxBroadcastBytes(val int) {
	if val == 0 {
		maxBroadcastBytes = DefaultMaxBroadcastBytes
	} else {
		maxBroadcastBytes = val
	}
}

// GetMulticastEnabled sets whether multicast announcements are enabled.
func SetMulticastEnabled(val bool) {
	multicastEnabledString = fmt.Sprintf("%v", val)
}

// SetMulticastAddress sets the address that will be used for multicast
// announcements.
func SetMulticastAddress(val string) {
	if val == "" {
		multicastAddress = DefaultMulticastAddress
	} else {
		multicastAddress = val
	}
}

// Gets an environmental variable "key". If it does not exist, "defaultVal" is
// returned; if it does, it attempts to convert to an integer, returning
// "defaultVal" if it fails.
func getIntVar(key string, defaultVal int) int {
	valueString := os.Getenv(key)
	valueInt := defaultVal

	if valueString != "" {
		i, err := strconv.Atoi(key)

		if err != nil {
			logfWarn("Failed to parse env property %s: %s is not "+
				"an integer. Using default.\n", key, valueString)
		} else {
			valueInt = i
		}
	}

	return valueInt
}

// Gets an environmental variable "key". If it does not exist, "defaultVal" is
// returned; if it does, it attempts to convert to a string slice, returning
// "defaultVal" if it fails.
func getStringArrayVar(key string, defaultVal string) []string {
	valueString := os.Getenv(key)

	if valueString == "" {
		valueString = defaultVal
	}

	valueSlice := splitDelimmitedString(valueString, stringListDelimitRegex)

	return valueSlice
}

// Gets an environmental variable "key". If it does not exist, "defaultVal" is
// returned; if it does, it attempts to convert to a string, returning
// "defaultVal" if it fails.
func getStringVar(key string, defaultVal string) string {
	valueString := os.Getenv(key)

	if valueString == "" {
		valueString = defaultVal
	}

	return valueString
}

// Splits a string on a regular expression.
func splitDelimmitedString(str string, regex string) []string {
	var result []string

	str = strings.TrimSpace(str)

	if str != "" {
		reg := regexp.MustCompile(regex)
		indices := reg.FindAllStringIndex(str, -1)

		result = make([]string, len(indices)+1)

		lastStart := 0
		for i, val := range indices {
			result[i] = str[lastStart:val[0]]
			lastStart = val[1]
		}

		result[len(indices)] = str[lastStart:]

		// Special case of single empty string
		if len(result) == 1 && result[0] == "" {
			result = make([]string, 0, 0)
		}
	}

	return result
}
<|endoftext|>


<commit_before>package prove

import (
	"log"
	"time"
	"unsafe"

	"github.com/nelhage/taktician/ptn"
	"github.com/nelhage/taktician/tak"
)

type DFPNSolver struct {
	attacker tak.Color
	table    dfpnTable
	debug    int

	stack []dfpnFrame
}

type dfpnFrame struct {
	g *tak.Position
	m tak.Move
}

type DFPNConfig struct {
	Debug    int
	Attacker tak.Color
	TableMem int64
}

type proofNumbers struct {
	phi, delta uint32
}

const INFINITY uint32 = 1 << 30

func (pn proofNumbers) exceeded(other proofNumbers) bool {
	return pn.phi >= other.phi || pn.delta >= other.delta
}

func (pn *proofNumbers) solved() bool {
	return pn.phi == 0 || pn.delta == 0
}

type dfpnTable struct {
	entries []entry
}

func (t *dfpnTable) lookup(g *tak.Position) (entry, bool) {
	e := t.entries[g.Hash()%uint64(len(t.entries))]
	if e.hash == g.Hash() {
		return e, true
	}
	return entry{}, false
}

func (t *dfpnTable) store(e *entry) bool {
	target := &t.entries[e.hash%uint64(len(t.entries))]
	if target.work <= e.work {
		*target = *e
		return true
	}

	return false
}

type entry struct {
	bounds proofNumbers
	hash   uint64
	work   uint64
	pv     tak.Move
	//	child  uint8
}

type dfpnChild struct {
	data entry
	move tak.Move
	g    *tak.Position
}

const defaultTableMem = 100 * 1024 * 1024

func NewDFPN(cfg *DFPNConfig) *DFPNSolver {
	if cfg.TableMem <= 0 {
		cfg.TableMem = defaultTableMem
	}
	return &DFPNSolver{
		attacker: cfg.Attacker,
		table: dfpnTable{
			entries: make([]entry, cfg.TableMem/int64(unsafe.Sizeof(entry{}))),
		},
		debug: cfg.Debug,
	}
}

func (d *DFPNSolver) Prove(g *tak.Position) ProofResult {
	if d.attacker == tak.NoColor {
		d.attacker = g.ToMove()
	}
	d.stack = nil
	start := time.Now()
	entry, _ := d.mid(g, proofNumbers{phi: INFINITY / 2, delta: INFINITY / 2}, entry{
		hash:   g.Hash(),
		work:   0,
		bounds: proofNumbers{phi: 1, delta: 1},
	})
	duration := time.Since(start)
	var result Evaluation = EvalUnknown
	if entry.bounds.phi == 0 {
		result = EvalTrue
	} else if entry.bounds.delta == 0 {
		result = EvalFalse
	}
	return ProofResult{
		Result:   result,
		Move:     entry.pv,
		Proof:    entry.bounds.phi,
		Disproof: entry.bounds.delta,
		Duration: duration,
	}
}

func (d *DFPNSolver) checkRepetition() bool {
	if len(d.stack) == 0 {
		return false
	}
	idx := len(d.stack) - 1
	cur := d.stack[idx].g.Hash()
	count := 0

	for idx > 0 {
		if d.stack[idx].g.Hash() == cur {
			count++
		}
		if count == 3 {
			return true
		}
		switch d.stack[idx].m.Type {
		case tak.PlaceFlat:
		case tak.PlaceStanding:
		case tak.PlaceCapstone:
			break
		}
		idx--
	}
	return false
}

func (d *DFPNSolver) mid(g *tak.Position, bounds proofNumbers, current entry) (entry, uint64) {
	if current.bounds.exceeded(bounds) {
		return current, 0
	}

	if over, result := g.GameOver(); over {
		current.bounds = d.terminalBounds(g, result)
		return current, 0
	}
	if d.checkRepetition() {
		current.bounds = d.terminalBounds(g, tak.NoColor)
		return current, 0
	}

	if d.debug > 6 && len(d.stack) > 0 {
		log.Printf(" depth=%d toMove=%s move=%s current=(%d,%d) bounds=(%d,%d)",
			len(d.stack),
			g.ToMove(),
			ptn.FormatMove(d.stack[len(d.stack)-1].m),
			current.bounds.phi, current.bounds.delta,
			bounds.phi, bounds.delta)
	}

	localWork := uint64(1)
	// compute children
	var children []dfpnChild
	var alloc [100]tak.Move
	moves := g.AllMoves(alloc[:])
	for _, m := range moves {
		p, err := g.Move(m)
		if err != nil {
			continue
		}
		childEntry := entry{
			hash:   p.Hash(),
			bounds: proofNumbers{1, 1},
		}
		if over, result := p.GameOver(); over {
			childEntry.bounds = d.terminalBounds(p, result)
		} else if b, ok := d.table.lookup(p); ok {
			childEntry = b
		}
		// todo look up table
		children = append(children, dfpnChild{
			move: m,
			g:    p,
			data: childEntry,
		})
		if childEntry.bounds.delta == 0 {
			break
		}
	}

	for {
		current.bounds = computePNs(children)
		if current.bounds.exceeded(bounds) {
			break
		}

		best_idx, childBounds := d.selectChild(children, bounds, current.bounds)
		// todo recurse

		current.pv = children[best_idx].move

		d.stack = append(d.stack, dfpnFrame{children[best_idx].g, children[best_idx].move})
		newEntry, work := d.mid(children[best_idx].g, childBounds, children[best_idx].data)
		d.stack = d.stack[:len(d.stack)-1]

		children[best_idx].data = newEntry
		localWork += work
		current.work += work
	}

	d.table.store(&current)

	return current, uint64(localWork)
}

func (d *DFPNSolver) terminalBounds(g *tak.Position, result tak.Color) proofNumbers {
	if result == tak.NoColor {
		result = d.attacker.Flip()
	}
	if result == g.ToMove() {
		return proofNumbers{phi: 0, delta: INFINITY}
	} else {
		return proofNumbers{phi: INFINITY, delta: 0}
	}
}

const epsilon = 0.1

func max(l, r uint32) uint32 {
	if l > r {
		return l
	}
	return r
}

func min(l, r uint32) uint32 {
	if l < r {
		return l
	}
	return r
}

func (d *DFPNSolver) selectChild(children []dfpnChild,
	bounds proofNumbers,
	pns proofNumbers) (int, proofNumbers) {
	delta1, delta2 := INFINITY, INFINITY

	best := -1
	for i, ch := range children {
		if ch.data.bounds.delta < delta1 {
			best = i
			delta2 = delta1
			delta1 = ch.data.bounds.delta
		} else if ch.data.bounds.delta < delta2 {
			delta2 = ch.data.bounds.delta
		}
	}

	phi1 := children[best].data.bounds.phi

	return best, proofNumbers{
		phi: bounds.delta + phi1 - pns.delta,
		delta: min(
			bounds.phi,
			//			max(delta2+1, uint32(float64(delta2)*(1.0+epsilon))),
			delta2+1,
		),
	}
}

func computePNs(children []dfpnChild) proofNumbers {
	out := proofNumbers{phi: INFINITY, delta: 0}
	for _, ch := range children {
		out.delta += ch.data.bounds.phi
		if out.delta > INFINITY {
			out.delta = INFINITY
		}
		if ch.data.bounds.delta < out.phi {
			out.phi = ch.data.bounds.delta
		}
	}
	return out
}
<commit_msg>dfpn: remove most allocations<commit_after>package prove

import (
	"log"
	"time"
	"unsafe"

	"github.com/nelhage/taktician/ptn"
	"github.com/nelhage/taktician/tak"
)

type DFPNSolver struct {
	attacker tak.Color
	table    dfpnTable
	debug    int

	stack []dfpnFrame
	pool  chan *tak.Position
}

type dfpnFrame struct {
	g *tak.Position
	m tak.Move
}

type DFPNConfig struct {
	Debug    int
	Attacker tak.Color
	TableMem int64
}

type proofNumbers struct {
	phi, delta uint32
}

const INFINITY uint32 = 1 << 30

func (pn proofNumbers) exceeded(other proofNumbers) bool {
	return pn.phi >= other.phi || pn.delta >= other.delta
}

func (pn *proofNumbers) solved() bool {
	return pn.phi == 0 || pn.delta == 0
}

type dfpnTable struct {
	entries []entry
}

func (t *dfpnTable) lookup(g *tak.Position) (entry, bool) {
	e := t.entries[g.Hash()%uint64(len(t.entries))]
	if e.hash == g.Hash() {
		return e, true
	}
	return entry{}, false
}

func (t *dfpnTable) store(e *entry) bool {
	target := &t.entries[e.hash%uint64(len(t.entries))]
	if target.work <= e.work {
		*target = *e
		return true
	}

	return false
}

type entry struct {
	bounds proofNumbers
	hash   uint64
	work   uint64
	pv     tak.Move
	//	child  uint8
}

type dfpnChild struct {
	data entry
	move tak.Move
	g    *tak.Position
}

const defaultTableMem = 100 * 1024 * 1024
const poolSize = 1024

func NewDFPN(cfg *DFPNConfig) *DFPNSolver {
	if cfg.TableMem <= 0 {
		cfg.TableMem = defaultTableMem
	}
	return &DFPNSolver{
		attacker: cfg.Attacker,
		table: dfpnTable{
			entries: make([]entry, cfg.TableMem/int64(unsafe.Sizeof(entry{}))),
		},
		debug: cfg.Debug,
		pool:  make(chan *tak.Position, poolSize),
	}
}

func (d *DFPNSolver) Prove(g *tak.Position) ProofResult {
	if d.attacker == tak.NoColor {
		d.attacker = g.ToMove()
	}
	d.stack = nil
	start := time.Now()
	entry, _ := d.mid(g, proofNumbers{phi: INFINITY / 2, delta: INFINITY / 2}, entry{
		hash:   g.Hash(),
		work:   0,
		bounds: proofNumbers{phi: 1, delta: 1},
	})
	duration := time.Since(start)
	var result Evaluation = EvalUnknown
	if entry.bounds.phi == 0 {
		result = EvalTrue
	} else if entry.bounds.delta == 0 {
		result = EvalFalse
	}
	return ProofResult{
		Result:   result,
		Move:     entry.pv,
		Proof:    entry.bounds.phi,
		Disproof: entry.bounds.delta,
		Duration: duration,
	}
}

func (d *DFPNSolver) checkRepetition() bool {
	if len(d.stack) == 0 {
		return false
	}
	idx := len(d.stack) - 1
	cur := d.stack[idx].g.Hash()
	count := 0

	for idx > 0 {
		if d.stack[idx].g.Hash() == cur {
			count++
		}
		if count == 3 {
			return true
		}
		switch d.stack[idx].m.Type {
		case tak.PlaceFlat:
		case tak.PlaceStanding:
		case tak.PlaceCapstone:
			break
		}
		idx--
	}
	return false
}

func (d *DFPNSolver) alloc() *tak.Position {
	select {
	case p := <-d.pool:
		return p
	default:
		return nil
	}
}

func (d *DFPNSolver) release(p *tak.Position) {
	select {
	case d.pool <- p:
	default:
	}
}

func (d *DFPNSolver) mid(g *tak.Position, bounds proofNumbers, current entry) (entry, uint64) {
	if current.bounds.exceeded(bounds) {
		return current, 0
	}

	if over, result := g.GameOver(); over {
		current.bounds = d.terminalBounds(g, result)
		return current, 0
	}
	if d.checkRepetition() {
		current.bounds = d.terminalBounds(g, tak.NoColor)
		return current, 0
	}

	if d.debug > 6 && len(d.stack) > 0 {
		log.Printf(" depth=%d toMove=%s move=%s current=(%d,%d) bounds=(%d,%d)",
			len(d.stack),
			g.ToMove(),
			ptn.FormatMove(d.stack[len(d.stack)-1].m),
			current.bounds.phi, current.bounds.delta,
			bounds.phi, bounds.delta)
	}

	localWork := uint64(1)
	// compute children
	var allocChildren [100]dfpnChild
	children := allocChildren[:0]
	var allocMoves [100]tak.Move
	moves := g.AllMoves(allocMoves[:0])
	defer func() {
		for _, ch := range children {
			d.release(ch.g)
		}
	}()
	for _, m := range moves {
		alloc := d.alloc()
		p, err := g.MovePreallocated(m, alloc)
		if err != nil {
			d.release(alloc)
			continue
		}
		childEntry := entry{
			hash:   p.Hash(),
			bounds: proofNumbers{1, 1},
		}
		if over, result := p.GameOver(); over {
			childEntry.bounds = d.terminalBounds(p, result)
		} else if b, ok := d.table.lookup(p); ok {
			childEntry = b
		}
		// todo look up table
		children = append(children, dfpnChild{
			move: m,
			g:    p,
			data: childEntry,
		})
		if childEntry.bounds.delta == 0 {
			break
		}
	}

	for {
		current.bounds = computePNs(children)
		if current.bounds.exceeded(bounds) {
			break
		}

		best_idx, childBounds := d.selectChild(children, bounds, current.bounds)
		// todo recurse

		current.pv = children[best_idx].move

		d.stack = append(d.stack, dfpnFrame{children[best_idx].g, children[best_idx].move})
		newEntry, work := d.mid(children[best_idx].g, childBounds, children[best_idx].data)
		d.stack = d.stack[:len(d.stack)-1]

		children[best_idx].data = newEntry
		localWork += work
		current.work += work
	}

	d.table.store(&current)

	return current, uint64(localWork)
}

func (d *DFPNSolver) terminalBounds(g *tak.Position, result tak.Color) proofNumbers {
	if result == tak.NoColor {
		result = d.attacker.Flip()
	}
	if result == g.ToMove() {
		return proofNumbers{phi: 0, delta: INFINITY}
	} else {
		return proofNumbers{phi: INFINITY, delta: 0}
	}
}

const epsilon = 0.1

func max(l, r uint32) uint32 {
	if l > r {
		return l
	}
	return r
}

func min(l, r uint32) uint32 {
	if l < r {
		return l
	}
	return r
}

func (d *DFPNSolver) selectChild(children []dfpnChild,
	bounds proofNumbers,
	pns proofNumbers) (int, proofNumbers) {
	delta1, delta2 := INFINITY, INFINITY

	best := -1
	for i, ch := range children {
		if ch.data.bounds.delta < delta1 {
			best = i
			delta2 = delta1
			delta1 = ch.data.bounds.delta
		} else if ch.data.bounds.delta < delta2 {
			delta2 = ch.data.bounds.delta
		}
	}

	phi1 := children[best].data.bounds.phi

	return best, proofNumbers{
		phi: bounds.delta + phi1 - pns.delta,
		delta: min(
			bounds.phi,
			//			max(delta2+1, uint32(float64(delta2)*(1.0+epsilon))),
			delta2+1,
		),
	}
}

func computePNs(children []dfpnChild) proofNumbers {
	out := proofNumbers{phi: INFINITY, delta: 0}
	for _, ch := range children {
		out.delta += ch.data.bounds.phi
		if out.delta > INFINITY {
			out.delta = INFINITY
		}
		if ch.data.bounds.delta < out.phi {
			out.phi = ch.data.bounds.delta
		}
	}
	return out
}
<|endoftext|>


<commit_before>package golog

import "testing"

func TestFacts (t *testing.T) {
    db := NewDatabase().
            Asserta(NewTerm("father", NewTerm("michael"))).
            Asserta(NewTerm("father", NewTerm("marc")))
    t.Logf("%s\n", db.String())

    // these should be provably true
    if !IsTrue(db, NewTerm("father", NewTerm("michael"))) {
        t.Errorf("Couldn't prove father(michael)")
    }
    if !IsTrue(db, NewTerm("father", NewTerm("marc"))) {
        t.Errorf("Couldn't prove father(marc)")
    }

    // these should not be provable
    if IsTrue(db, NewTerm("father", NewTerm("sue"))) {
        t.Errorf("Proved father(sue)")
    }
    if IsTrue(db, NewTerm("mother", NewTerm("michael"))) {
        t.Errorf("Proved mother(michael)")
    }
}
<commit_msg>Make sure arity mismatch affects provability<commit_after>package golog

import "testing"

func TestFacts (t *testing.T) {
    db := NewDatabase().
            Asserta(NewTerm("father", NewTerm("michael"))).
            Asserta(NewTerm("father", NewTerm("marc")))
    t.Logf("%s\n", db.String())

    // these should be provably true
    if !IsTrue(db, NewTerm("father", NewTerm("michael"))) {
        t.Errorf("Couldn't prove father(michael)")
    }
    if !IsTrue(db, NewTerm("father", NewTerm("marc"))) {
        t.Errorf("Couldn't prove father(marc)")
    }

    // these should not be provable
    if IsTrue(db, NewTerm("father", NewTerm("sue"))) {
        t.Errorf("Proved father(sue)")
    }
    if IsTrue(db, NewTerm("father", NewTerm("michael"), NewTerm("marc"))) {
        t.Errorf("Proved father(michael, marc)")
    }
    if IsTrue(db, NewTerm("mother", NewTerm("michael"))) {
        t.Errorf("Proved mother(michael)")
    }
}
<|endoftext|>


<commit_before>package veneur

import (
	"compress/zlib"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/sirupsen/logrus"
	"github.com/stretchr/testify/assert"
	"github.com/stripe/veneur/samplers"
)

func generateProxyConfig() ProxyConfig {
	return ProxyConfig{
		Debug: false,
		ConsulRefreshInterval:    "86400s",
		ConsulForwardServiceName: "forwardServiceName",
		ConsulTraceServiceName:   "traceServiceName",
		TraceAddress:             "127.0.0.1:8128",
		TraceAPIAddress:          "127.0.0.1:8135",
		HTTPAddress:              "127.0.0.1:0",
		StatsAddress:             "127.0.0.1:8201",
	}
}

type ConsulTwoMetricRoundTripper struct {
	t         *testing.T
	wg        *sync.WaitGroup
	aReceived bool
	bReceived bool

	mtx sync.Mutex
}

func (rt *ConsulTwoMetricRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
	// Ensure that only one RoundTrip is happening at once
	// to prevent dataraces on aReceived and bReceived

	rt.mtx.Lock()
	defer rt.mtx.Unlock()

	rec := httptest.NewRecorder()
	if req.URL.Path == "/v1/health/service/forwardServiceName" {
		resp, _ := ioutil.ReadFile("fixtures/consul/health_service_two.json")
		rec.Write(resp)
		rec.Code = http.StatusOK
	} else if req.URL.Path == "/v1/health/service/traceServiceName" {
		resp, _ := ioutil.ReadFile("fixtures/consul/health_service_two.json")
		rec.Write(resp)
		rec.Code = http.StatusOK
	} else if req.URL.Path == "/api/v1/series" {
		// Just make the datadog bit work
		rec.Code = http.StatusOK
	} else if req.URL.Path == "/import" && req.Host == "10.1.10.12:8000" {
		z, _ := zlib.NewReader(req.Body)
		body, _ := ioutil.ReadAll(z)
		defer req.Body.Close()
		if strings.Contains(string(body), "a.b.c") {
			rt.aReceived = true
		}
		rec.Code = http.StatusOK
	} else if req.URL.Path == "/import" && req.Host == "10.1.10.13:8000" {
		z, _ := zlib.NewReader(req.Body)
		body, _ := ioutil.ReadAll(z)
		defer req.Body.Close()
		if strings.Contains(string(body), "x.b.c") {
			rt.bReceived = true
		}
		rec.Code = http.StatusOK
	} else {
		assert.Fail(rt.t, "Received an unexpected request: %s %s", req.Host, req.URL.Path)
	}

	// If we've gotten all of them, fire!
	if rt.aReceived && rt.bReceived {
		rt.wg.Done()
	}

	return rec.Result(), nil
}

func TestAllowStaticServices(t *testing.T) {
	proxyConfig := generateProxyConfig()
	proxyConfig.ConsulForwardServiceName = ""
	proxyConfig.ConsulTraceServiceName = ""
	proxyConfig.ForwardAddress = "localhost:1234"
	proxyConfig.TraceAddress = "localhost:1234"

	server, error := NewProxyFromConfig(logrus.New(), proxyConfig)
	assert.NoError(t, error, "Should start with just static services")
	assert.False(t, server.usingConsul, "Server isn't using consul")
}

func TestMissingServices(t *testing.T) {
	proxyConfig := generateProxyConfig()
	proxyConfig.ForwardAddress = ""
	proxyConfig.TraceAddress = ""
	proxyConfig.ConsulForwardServiceName = ""
	proxyConfig.ConsulTraceServiceName = ""

	_, error := NewProxyFromConfig(logrus.New(), proxyConfig)
	assert.Error(t, error, "No consul services means Proxy won't start")
}

func TestAcceptingBooleans(t *testing.T) {
	proxyConfig := generateProxyConfig()
	proxyConfig.ConsulTraceServiceName = ""
	proxyConfig.TraceAddress = ""

	server, _ := NewProxyFromConfig(logrus.New(), proxyConfig)
	assert.True(t, server.AcceptingForwards, "Server accepts forwards")
	assert.False(t, server.AcceptingTraces, "Server does not forward traces")
}

func TestConsistentForward(t *testing.T) {

	// We need to set up a proxy, have a local veneur send to it, then verify
	// that the proxy forwards to two downstream fake globals.
	// TIME FOR SOME GAME THEORY

	// Make the proxy
	proxyConfig := generateProxyConfig()
	proxyConfig.ForwardAddress = "localhost:1234"
	proxyConfig.ConsulForwardServiceName = "forwardServiceName"
	proxyConfig.Debug = true
	wg := sync.WaitGroup{}
	wg.Add(1)
	transport := &ConsulTwoMetricRoundTripper{
		t:  t,
		wg: &wg,
	}
	server, _ := NewProxyFromConfig(logrus.New(), proxyConfig)

	server.HTTPClient.Transport = transport
	defer server.Shutdown()

	server.Start()
	srv := httptest.NewServer(server.Handler())
	defer srv.Close()

	// Make sure we're sane first
	assert.Len(t, server.ForwardDestinations.Members(), 2, "Incorrect host count in ring")

	// Cool, now let's make a veneur to process some bits!
	config := localConfig()
	config.ForwardAddress = srv.URL
	f := newFixture(t, config, nil, nil)
	defer f.Close()

	f.server.Workers[0].ProcessMetric(&samplers.UDPMetric{
		MetricKey: samplers.MetricKey{
			Name: "a.b.c",
			Type: "histogram",
		},
		Value:      float64(100),
		Digest:     12345,
		SampleRate: 1.0,
		Scope:      samplers.MixedScope,
	})
	f.server.Workers[0].ProcessMetric(&samplers.UDPMetric{
		MetricKey: samplers.MetricKey{
			Name: "x.b.c",
			Type: "histogram",
		},
		Value:      float64(100),
		Digest:     12345,
		SampleRate: 1.0,
		Scope:      samplers.MixedScope,
	})

	c := make(chan struct{})
	go func() {
		defer close(c)
		wg.Wait()
	}()

	select {
	case <-c:
		fmt.Println("GOT 'IM")
	case <-time.After(3 * time.Second):
		assert.Fail(t, "Failed to receive all metrics before timeout")
	}
}
<commit_msg>Add a test to ensure forwarding metrics to a global veneur times out<commit_after>package veneur

import (
	"compress/zlib"
	"context"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/sirupsen/logrus"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/stripe/veneur/samplers"
)

func generateProxyConfig() ProxyConfig {
	return ProxyConfig{
		Debug: false,
		ConsulRefreshInterval:    "86400s",
		ConsulForwardServiceName: "forwardServiceName",
		ConsulTraceServiceName:   "traceServiceName",
		TraceAddress:             "127.0.0.1:8128",
		TraceAPIAddress:          "127.0.0.1:8135",
		HTTPAddress:              "127.0.0.1:0",
		StatsAddress:             "127.0.0.1:8201",
	}
}

type ConsulTwoMetricRoundTripper struct {
	t         *testing.T
	wg        *sync.WaitGroup
	aReceived bool
	bReceived bool

	mtx sync.Mutex
}

func (rt *ConsulTwoMetricRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
	// Ensure that only one RoundTrip is happening at once
	// to prevent dataraces on aReceived and bReceived

	rt.mtx.Lock()
	defer rt.mtx.Unlock()

	rec := httptest.NewRecorder()
	if req.URL.Path == "/v1/health/service/forwardServiceName" {
		resp, _ := ioutil.ReadFile("fixtures/consul/health_service_two.json")
		rec.Write(resp)
		rec.Code = http.StatusOK
	} else if req.URL.Path == "/v1/health/service/traceServiceName" {
		resp, _ := ioutil.ReadFile("fixtures/consul/health_service_two.json")
		rec.Write(resp)
		rec.Code = http.StatusOK
	} else if req.URL.Path == "/api/v1/series" {
		// Just make the datadog bit work
		rec.Code = http.StatusOK
	} else if req.URL.Path == "/import" && req.Host == "10.1.10.12:8000" {
		z, _ := zlib.NewReader(req.Body)
		body, _ := ioutil.ReadAll(z)
		defer req.Body.Close()
		if strings.Contains(string(body), "a.b.c") {
			rt.aReceived = true
		}
		rec.Code = http.StatusOK
	} else if req.URL.Path == "/import" && req.Host == "10.1.10.13:8000" {
		z, _ := zlib.NewReader(req.Body)
		body, _ := ioutil.ReadAll(z)
		defer req.Body.Close()
		if strings.Contains(string(body), "x.b.c") {
			rt.bReceived = true
		}
		rec.Code = http.StatusOK
	} else {
		assert.Fail(rt.t, "Received an unexpected request: %s %s", req.Host, req.URL.Path)
	}

	// If we've gotten all of them, fire!
	if rt.aReceived && rt.bReceived {
		rt.wg.Done()
	}

	return rec.Result(), nil
}

func TestAllowStaticServices(t *testing.T) {
	proxyConfig := generateProxyConfig()
	proxyConfig.ConsulForwardServiceName = ""
	proxyConfig.ConsulTraceServiceName = ""
	proxyConfig.ForwardAddress = "localhost:1234"
	proxyConfig.TraceAddress = "localhost:1234"

	server, error := NewProxyFromConfig(logrus.New(), proxyConfig)
	assert.NoError(t, error, "Should start with just static services")
	assert.False(t, server.usingConsul, "Server isn't using consul")
}

func TestMissingServices(t *testing.T) {
	proxyConfig := generateProxyConfig()
	proxyConfig.ForwardAddress = ""
	proxyConfig.TraceAddress = ""
	proxyConfig.ConsulForwardServiceName = ""
	proxyConfig.ConsulTraceServiceName = ""

	_, error := NewProxyFromConfig(logrus.New(), proxyConfig)
	assert.Error(t, error, "No consul services means Proxy won't start")
}

func TestAcceptingBooleans(t *testing.T) {
	proxyConfig := generateProxyConfig()
	proxyConfig.ConsulTraceServiceName = ""
	proxyConfig.TraceAddress = ""

	server, _ := NewProxyFromConfig(logrus.New(), proxyConfig)
	assert.True(t, server.AcceptingForwards, "Server accepts forwards")
	assert.False(t, server.AcceptingTraces, "Server does not forward traces")
}

func TestConsistentForward(t *testing.T) {

	// We need to set up a proxy, have a local veneur send to it, then verify
	// that the proxy forwards to two downstream fake globals.
	// TIME FOR SOME GAME THEORY

	// Make the proxy
	proxyConfig := generateProxyConfig()
	proxyConfig.ForwardAddress = "localhost:1234"
	proxyConfig.ConsulForwardServiceName = "forwardServiceName"
	proxyConfig.Debug = true
	wg := sync.WaitGroup{}
	wg.Add(1)
	transport := &ConsulTwoMetricRoundTripper{
		t:  t,
		wg: &wg,
	}
	server, _ := NewProxyFromConfig(logrus.New(), proxyConfig)

	server.HTTPClient.Transport = transport
	defer server.Shutdown()

	server.Start()
	srv := httptest.NewServer(server.Handler())
	defer srv.Close()

	// Make sure we're sane first
	assert.Len(t, server.ForwardDestinations.Members(), 2, "Incorrect host count in ring")

	// Cool, now let's make a veneur to process some bits!
	config := localConfig()
	config.ForwardAddress = srv.URL
	f := newFixture(t, config, nil, nil)
	defer f.Close()

	f.server.Workers[0].ProcessMetric(&samplers.UDPMetric{
		MetricKey: samplers.MetricKey{
			Name: "a.b.c",
			Type: "histogram",
		},
		Value:      float64(100),
		Digest:     12345,
		SampleRate: 1.0,
		Scope:      samplers.MixedScope,
	})
	f.server.Workers[0].ProcessMetric(&samplers.UDPMetric{
		MetricKey: samplers.MetricKey{
			Name: "x.b.c",
			Type: "histogram",
		},
		Value:      float64(100),
		Digest:     12345,
		SampleRate: 1.0,
		Scope:      samplers.MixedScope,
	})

	c := make(chan struct{})
	go func() {
		defer close(c)
		wg.Wait()
	}()

	select {
	case <-c:
		fmt.Println("GOT 'IM")
	case <-time.After(3 * time.Second):
		assert.Fail(t, "Failed to receive all metrics before timeout")
	}
}

func TestTimeout(t *testing.T) {
	cfg := generateProxyConfig()
	cfg.ConsulTraceServiceName = ""
	cfg.ConsulForwardServiceName = ""

	never := make(chan struct{})
	defer close(never)
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		<-never
	}))
	defer ts.Close()

	cfg.ForwardAddress = ts.URL
	cfg.ForwardTimeout = "1ns" // just really really short
	server, _ := NewProxyFromConfig(logrus.New(), cfg)

	ctr := samplers.Counter{Name: "foo", Tags: []string{}}
	ctr.Sample(20.0, 1.0)
	jsonCtr, err := ctr.Export()
	require.NoError(t, err)
	metrics := []samplers.JSONMetric{jsonCtr}

	// Now send the metrics to a server that never responds; we
	// expect this to return before our (sufficiently long)
	// timeout:
	ch := make(chan struct{})
	go func() {
		server.ProxyMetrics(context.Background(), metrics, "foo.com")
		close(ch)
	}()
	select {
	case <-ch:
		t.Log("Returned quickly, great.")
	case <-time.After(3 * time.Second):
		t.Fatal("Proxy took too long to time out. Is the deadline behavior working?")
	}
}
<|endoftext|>


<commit_before>// Package scripts contains scripts (pretty unexpected huh?) that
// facilitate testing and development. See the project readme for
// more information:
// https://github.com/smartystreets/goconvey/blob/master/README.md
package scripts
<commit_msg>Added notice about scripts going away after go 1.3.<commit_after>// NOTE: this folder will probably disappear shortly after go 1.3 is released.
//
// Package scripts contains scripts (pretty unexpected huh?) that
// facilitate testing and development. See the project readme for
// more information:
// https://github.com/smartystreets/goconvey/blob/master/README.md
package scripts
<|endoftext|>