<commit_before>// Copyright 2015 Google Inc. All Rights Reserved.
// Author: jacobsa@google.com (Aaron Jacobs)

package memfs

import (
	"fmt"
	"os"
	"time"

	"github.com/jacobsa/fuse"
	"github.com/jacobsa/fuse/fuseutil"
	"github.com/jacobsa/gcloud/syncutil"
	"github.com/jacobsa/gcsfuse/timeutil"
	"golang.org/x/net/context"
)

type memFS struct {
	fuseutil.NotImplementedFileSystem

	/////////////////////////
	// Dependencies
	/////////////////////////

	clock timeutil.Clock

	/////////////////////////
	// Mutable state
	/////////////////////////

	// When acquiring this lock, the caller must hold no inode locks.
	mu syncutil.InvariantMutex

	// The collection of live inodes, indexed by ID. IDs of free inodes that may
	// be re-used have nil entries. No ID less than fuse.RootInodeID is ever used.
	//
	// INVARIANT: len(inodes) > fuse.RootInodeID
	// INVARIANT: For all i < fuse.RootInodeID, inodes[i] == nil
	// INVARIANT: inodes[fuse.RootInodeID] != nil
	// INVARIANT: inodes[fuse.RootInodeID].dir is true
	inodes []*inode // GUARDED_BY(mu)

	// A list of inode IDs within inodes available for reuse, not including the
	// reserved IDs less than fuse.RootInodeID.
	//
	// INVARIANT: This is all and only indices i of 'inodes' such that i >
	// fuse.RootInodeID and inodes[i] == nil
	freeInodes []fuse.InodeID // GUARDED_BY(mu)
}

// Create a file system that stores data and metadata in memory.
func NewMemFS(
	clock timeutil.Clock) fuse.FileSystem {
	// Set up the basic struct.
	fs := &memFS{
		clock:  clock,
		inodes: make([]*inode, fuse.RootInodeID+1),
	}

	// Set up the root inode.
	rootAttrs := fuse.InodeAttributes{
		Mode: 0777 | os.ModeDir,
	}

	fs.inodes[fuse.RootInodeID] = newInode(rootAttrs)

	// Set up invariant checking.
	fs.mu = syncutil.NewInvariantMutex(fs.checkInvariants)

	return fs
}

func (fs *memFS) checkInvariants() {
	// Check reserved inodes.
	for i := 0; i < fuse.RootInodeID; i++ {
		if fs.inodes[i] != nil {
			panic(fmt.Sprintf("Non-nil inode for ID: %v", i))
		}
	}

	// Check the root inode.
	if !fs.inodes[fuse.RootInodeID].dir {
		panic("Expected root to be a directory.")
	}

	// Build our own list of free IDs.
	freeIDsEncountered := make(map[fuse.InodeID]struct{})
	for i := fuse.RootInodeID + 1; i < len(fs.inodes); i++ {
		inode := fs.inodes[i]
		if inode == nil {
			freeIDsEncountered[fuse.InodeID(i)] = struct{}{}
			continue
		}
	}

	// Check fs.freeInodes.
	if len(fs.freeInodes) != len(freeIDsEncountered) {
		panic(
			fmt.Sprintf(
				"Length mismatch: %v vs. %v",
				len(fs.freeInodes),
				len(freeIDsEncountered)))
	}

	for _, id := range fs.freeInodes {
		if _, ok := freeIDsEncountered[id]; !ok {
			panic(fmt.Sprintf("Unexected free inode ID: %v", id))
		}
	}
}

func (fs *memFS) Init(
	ctx context.Context,
	req *fuse.InitRequest) (resp *fuse.InitResponse, err error) {
	resp = &fuse.InitResponse{}
	return
}

// Find the given inode and return it with its lock held. Panic if it doesn't
// exist.
//
// SHARED_LOCKS_REQUIRED(fs.mu)
// EXCLUSIVE_LOCK_FUNCTION(inode.mu)
func (fs *memFS) getInodeForModifyingOrDie(id fuse.InodeID) (inode *inode) {
	inode = fs.inodes[id]
	if inode == nil {
		panic(fmt.Sprintf("Unknown inode: %v", id))
	}

	inode.mu.Lock()
	return
}

// Find the given inode and return it with its lock held for reading. Panic if
// it doesn't exist.
//
// SHARED_LOCKS_REQUIRED(fs.mu)
// SHARED_LOCK_FUNCTION(inode.mu)
func (fs *memFS) getInodeForReadingOrDie(id fuse.InodeID) (inode *inode) {
	inode = fs.inodes[id]
	if inode == nil {
		panic(fmt.Sprintf("Unknown inode: %v", id))
	}

	inode.mu.RLock()
	return
}

// Allocate a new inode, assigning it an ID that is not in use. Return it with
// its lock held.
//
// EXCLUSIVE_LOCKS_REQUIRED(fs.mu)
// EXCLUSIVE_LOCK_FUNCTION(inode.mu)
func (fs *memFS) allocateInode(
	attrs fuse.InodeAttributes) (id fuse.InodeID, inode *inode) {
	// Create and lock the inode.
	inode = newInode(attrs)
	inode.mu.Lock()

	// Re-use a free ID if possible. Otherwise mint a new one.
	numFree := len(fs.freeInodes)
	if numFree != 0 {
		id = fs.freeInodes[numFree-1]
		fs.freeInodes = fs.freeInodes[:numFree-1]
		fs.inodes[id] = inode
	} else {
		id = fuse.InodeID(len(fs.inodes))
		fs.inodes = append(fs.inodes, inode)
	}

	return
}

func (fs *memFS) LookUpInode(
	ctx context.Context,
	req *fuse.LookUpInodeRequest) (resp *fuse.LookUpInodeResponse, err error) {
	resp = &fuse.LookUpInodeResponse{}

	fs.mu.RLock()
	defer fs.mu.RUnlock()

	// Grab the parent directory.
	inode := fs.getInodeForReadingOrDie(req.Parent)
	defer inode.mu.RUnlock()

	// Does the directory have an entry with the given name?
	childID, ok := inode.LookUpChild(req.Name)
	if !ok {
		err = fuse.ENOENT
		return
	}

	// Grab the child.
	child := fs.getInodeForReadingOrDie(childID)
	defer child.mu.RUnlock()

	// Fill in the response.
	resp.Entry.Child = childID
	resp.Entry.Attributes = child.attributes

	// We don't spontaneously mutate, so the kernel can cache as long as it wants
	// (since it also handles invalidation).
	resp.Entry.AttributesExpiration = fs.clock.Now().Add(365 * 24 * time.Hour)
	resp.Entry.EntryExpiration = resp.Entry.EntryExpiration

	return
}

func (fs *memFS) GetInodeAttributes(
	ctx context.Context,
	req *fuse.GetInodeAttributesRequest) (
	resp *fuse.GetInodeAttributesResponse, err error) {
	resp = &fuse.GetInodeAttributesResponse{}

	fs.mu.RLock()
	defer fs.mu.RUnlock()

	// Grab the inode.
	inode := fs.getInodeForReadingOrDie(req.Inode)
	defer inode.mu.RUnlock()

	// Fill in the response.
	resp.Attributes = inode.attributes

	// We don't spontaneously mutate, so the kernel can cache as long as it wants
	// (since it also handles invalidation).
	resp.AttributesExpiration = fs.clock.Now().Add(365 * 24 * time.Hour)

	return
}

func (fs *memFS) MkDir(
	ctx context.Context,
	req *fuse.MkDirRequest) (resp *fuse.MkDirResponse, err error) {
	resp = &fuse.MkDirResponse{}

	fs.mu.Lock()
	defer fs.mu.Unlock()

	// Grab the parent, which we will update shortly.
	parent := fs.getInodeForModifyingOrDie(req.Parent)
	defer parent.mu.Unlock()

	// Allocate a child.
	childAttrs := fuse.InodeAttributes{
		Mode: req.Mode,
	}

	childID, child := fs.allocateInode(childAttrs)
	defer child.mu.Unlock()

	// Add an entry in the parent.
	parent.AddChild(childID, req.Name, fuseutil.DT_Directory)

	// Fill in the response.
	resp.Entry.Child = childID
	resp.Entry.Attributes = child.attributes

	// We don't spontaneously mutate, so the kernel can cache as long as it wants
	// (since it also handles invalidation).
	resp.Entry.AttributesExpiration = fs.clock.Now().Add(365 * 24 * time.Hour)
	resp.Entry.EntryExpiration = resp.Entry.EntryExpiration

	return
}

func (fs *memFS) OpenDir(
	ctx context.Context,
	req *fuse.OpenDirRequest) (resp *fuse.OpenDirResponse, err error) {
	resp = &fuse.OpenDirResponse{}

	fs.mu.RLock()
	defer fs.mu.RUnlock()

	// We don't mutate spontaneosuly, so if the VFS layer has asked for an
	// inode that doesn't exist, something screwed up earlier (a lookup, a
	// cache invalidation, etc.).
	inode := fs.getInodeForReadingOrDie(req.Inode)
	defer inode.mu.RUnlock()

	if !inode.dir {
		panic("Found non-dir.")
	}

	return
}

func (fs *memFS) ReadDir(
	ctx context.Context,
	req *fuse.ReadDirRequest) (resp *fuse.ReadDirResponse, err error) {
	resp = &fuse.ReadDirResponse{}

	fs.mu.RLock()
	defer fs.mu.RUnlock()

	// Grab the directory.
	inode := fs.getInodeForReadingOrDie(req.Inode)
	defer inode.mu.RUnlock()

	// Serve the request.
	resp.Data, err = inode.ReadDir(int(req.Offset), req.Size)
	if err != nil {
		err = fmt.Errorf("inode.ReadDir: %v", err)
		return
	}

	return
}
<commit_msg>Fixed mtime in mkdir.<commit_after>// Copyright 2015 Google Inc. All Rights Reserved.
// Author: jacobsa@google.com (Aaron Jacobs)

package memfs

import (
	"fmt"
	"os"
	"time"

	"github.com/jacobsa/fuse"
	"github.com/jacobsa/fuse/fuseutil"
	"github.com/jacobsa/gcloud/syncutil"
	"github.com/jacobsa/gcsfuse/timeutil"
	"golang.org/x/net/context"
)

type memFS struct {
	fuseutil.NotImplementedFileSystem

	/////////////////////////
	// Dependencies
	/////////////////////////

	clock timeutil.Clock

	/////////////////////////
	// Mutable state
	/////////////////////////

	// When acquiring this lock, the caller must hold no inode locks.
	mu syncutil.InvariantMutex

	// The collection of live inodes, indexed by ID. IDs of free inodes that may
	// be re-used have nil entries. No ID less than fuse.RootInodeID is ever used.
	//
	// INVARIANT: len(inodes) > fuse.RootInodeID
	// INVARIANT: For all i < fuse.RootInodeID, inodes[i] == nil
	// INVARIANT: inodes[fuse.RootInodeID] != nil
	// INVARIANT: inodes[fuse.RootInodeID].dir is true
	inodes []*inode // GUARDED_BY(mu)

	// A list of inode IDs within inodes available for reuse, not including the
	// reserved IDs less than fuse.RootInodeID.
	//
	// INVARIANT: This is all and only indices i of 'inodes' such that i >
	// fuse.RootInodeID and inodes[i] == nil
	freeInodes []fuse.InodeID // GUARDED_BY(mu)
}

// Create a file system that stores data and metadata in memory.
func NewMemFS(
	clock timeutil.Clock) fuse.FileSystem {
	// Set up the basic struct.
	fs := &memFS{
		clock:  clock,
		inodes: make([]*inode, fuse.RootInodeID+1),
	}

	// Set up the root inode.
	rootAttrs := fuse.InodeAttributes{
		Mode: 0777 | os.ModeDir,
	}

	fs.inodes[fuse.RootInodeID] = newInode(rootAttrs)

	// Set up invariant checking.
	fs.mu = syncutil.NewInvariantMutex(fs.checkInvariants)

	return fs
}

func (fs *memFS) checkInvariants() {
	// Check reserved inodes.
	for i := 0; i < fuse.RootInodeID; i++ {
		if fs.inodes[i] != nil {
			panic(fmt.Sprintf("Non-nil inode for ID: %v", i))
		}
	}

	// Check the root inode.
	if !fs.inodes[fuse.RootInodeID].dir {
		panic("Expected root to be a directory.")
	}

	// Build our own list of free IDs.
	freeIDsEncountered := make(map[fuse.InodeID]struct{})
	for i := fuse.RootInodeID + 1; i < len(fs.inodes); i++ {
		inode := fs.inodes[i]
		if inode == nil {
			freeIDsEncountered[fuse.InodeID(i)] = struct{}{}
			continue
		}
	}

	// Check fs.freeInodes.
	if len(fs.freeInodes) != len(freeIDsEncountered) {
		panic(
			fmt.Sprintf(
				"Length mismatch: %v vs. %v",
				len(fs.freeInodes),
				len(freeIDsEncountered)))
	}

	for _, id := range fs.freeInodes {
		if _, ok := freeIDsEncountered[id]; !ok {
			panic(fmt.Sprintf("Unexected free inode ID: %v", id))
		}
	}
}

func (fs *memFS) Init(
	ctx context.Context,
	req *fuse.InitRequest) (resp *fuse.InitResponse, err error) {
	resp = &fuse.InitResponse{}
	return
}

// Find the given inode and return it with its lock held. Panic if it doesn't
// exist.
//
// SHARED_LOCKS_REQUIRED(fs.mu)
// EXCLUSIVE_LOCK_FUNCTION(inode.mu)
func (fs *memFS) getInodeForModifyingOrDie(id fuse.InodeID) (inode *inode) {
	inode = fs.inodes[id]
	if inode == nil {
		panic(fmt.Sprintf("Unknown inode: %v", id))
	}

	inode.mu.Lock()
	return
}

// Find the given inode and return it with its lock held for reading. Panic if
// it doesn't exist.
//
// SHARED_LOCKS_REQUIRED(fs.mu)
// SHARED_LOCK_FUNCTION(inode.mu)
func (fs *memFS) getInodeForReadingOrDie(id fuse.InodeID) (inode *inode) {
	inode = fs.inodes[id]
	if inode == nil {
		panic(fmt.Sprintf("Unknown inode: %v", id))
	}

	inode.mu.RLock()
	return
}

// Allocate a new inode, assigning it an ID that is not in use. Return it with
// its lock held.
//
// EXCLUSIVE_LOCKS_REQUIRED(fs.mu)
// EXCLUSIVE_LOCK_FUNCTION(inode.mu)
func (fs *memFS) allocateInode(
	attrs fuse.InodeAttributes) (id fuse.InodeID, inode *inode) {
	// Create and lock the inode.
	inode = newInode(attrs)
	inode.mu.Lock()

	// Re-use a free ID if possible. Otherwise mint a new one.
	numFree := len(fs.freeInodes)
	if numFree != 0 {
		id = fs.freeInodes[numFree-1]
		fs.freeInodes = fs.freeInodes[:numFree-1]
		fs.inodes[id] = inode
	} else {
		id = fuse.InodeID(len(fs.inodes))
		fs.inodes = append(fs.inodes, inode)
	}

	return
}

func (fs *memFS) LookUpInode(
	ctx context.Context,
	req *fuse.LookUpInodeRequest) (resp *fuse.LookUpInodeResponse, err error) {
	resp = &fuse.LookUpInodeResponse{}

	fs.mu.RLock()
	defer fs.mu.RUnlock()

	// Grab the parent directory.
	inode := fs.getInodeForReadingOrDie(req.Parent)
	defer inode.mu.RUnlock()

	// Does the directory have an entry with the given name?
	childID, ok := inode.LookUpChild(req.Name)
	if !ok {
		err = fuse.ENOENT
		return
	}

	// Grab the child.
	child := fs.getInodeForReadingOrDie(childID)
	defer child.mu.RUnlock()

	// Fill in the response.
	resp.Entry.Child = childID
	resp.Entry.Attributes = child.attributes

	// We don't spontaneously mutate, so the kernel can cache as long as it wants
	// (since it also handles invalidation).
	resp.Entry.AttributesExpiration = fs.clock.Now().Add(365 * 24 * time.Hour)
	resp.Entry.EntryExpiration = resp.Entry.EntryExpiration

	return
}

func (fs *memFS) GetInodeAttributes(
	ctx context.Context,
	req *fuse.GetInodeAttributesRequest) (
	resp *fuse.GetInodeAttributesResponse, err error) {
	resp = &fuse.GetInodeAttributesResponse{}

	fs.mu.RLock()
	defer fs.mu.RUnlock()

	// Grab the inode.
	inode := fs.getInodeForReadingOrDie(req.Inode)
	defer inode.mu.RUnlock()

	// Fill in the response.
	resp.Attributes = inode.attributes

	// We don't spontaneously mutate, so the kernel can cache as long as it wants
	// (since it also handles invalidation).
	resp.AttributesExpiration = fs.clock.Now().Add(365 * 24 * time.Hour)

	return
}

func (fs *memFS) MkDir(
	ctx context.Context,
	req *fuse.MkDirRequest) (resp *fuse.MkDirResponse, err error) {
	resp = &fuse.MkDirResponse{}

	fs.mu.Lock()
	defer fs.mu.Unlock()

	// Grab the parent, which we will update shortly.
	parent := fs.getInodeForModifyingOrDie(req.Parent)
	defer parent.mu.Unlock()

	// Allocate a child.
	now := fs.clock.Now()
	childAttrs := fuse.InodeAttributes{
		Mode:   req.Mode,
		Atime:  now,
		Mtime:  now,
		Crtime: now,
	}

	childID, child := fs.allocateInode(childAttrs)
	defer child.mu.Unlock()

	// Add an entry in the parent.
	parent.AddChild(childID, req.Name, fuseutil.DT_Directory)

	// Fill in the response.
	resp.Entry.Child = childID
	resp.Entry.Attributes = child.attributes

	// We don't spontaneously mutate, so the kernel can cache as long as it wants
	// (since it also handles invalidation).
	resp.Entry.AttributesExpiration = fs.clock.Now().Add(365 * 24 * time.Hour)
	resp.Entry.EntryExpiration = resp.Entry.EntryExpiration

	return
}

func (fs *memFS) OpenDir(
	ctx context.Context,
	req *fuse.OpenDirRequest) (resp *fuse.OpenDirResponse, err error) {
	resp = &fuse.OpenDirResponse{}

	fs.mu.RLock()
	defer fs.mu.RUnlock()

	// We don't mutate spontaneosuly, so if the VFS layer has asked for an
	// inode that doesn't exist, something screwed up earlier (a lookup, a
	// cache invalidation, etc.).
	inode := fs.getInodeForReadingOrDie(req.Inode)
	defer inode.mu.RUnlock()

	if !inode.dir {
		panic("Found non-dir.")
	}

	return
}

func (fs *memFS) ReadDir(
	ctx context.Context,
	req *fuse.ReadDirRequest) (resp *fuse.ReadDirResponse, err error) {
	resp = &fuse.ReadDirResponse{}

	fs.mu.RLock()
	defer fs.mu.RUnlock()

	// Grab the directory.
	inode := fs.getInodeForReadingOrDie(req.Inode)
	defer inode.mu.RUnlock()

	// Serve the request.
	resp.Data, err = inode.ReadDir(int(req.Offset), req.Size)
	if err != nil {
		err = fmt.Errorf("inode.ReadDir: %v", err)
		return
	}

	return
}
<|endoftext|>


<commit_before>package ec2

import (
	"fmt"
	"launchpad.net/goamz/ec2"
	"launchpad.net/goamz/s3"
	"launchpad.net/juju/go/environs"
	"sync"
)

const zkPortSuffix = ":2181"

func init() {
	environs.RegisterProvider("ec2", environProvider{})
}

type environProvider struct{}

var _ environs.EnvironProvider = environProvider{}

type environ struct {
	name             string
	config           *providerConfig
	ec2              *ec2.EC2
	s3               *s3.S3
	checkBucket      sync.Once
	checkBucketError error
}

var _ environs.Environ = (*environ)(nil)

type instance struct {
	*ec2.Instance
}

func (inst *instance) String() string {
	return inst.Id()
}

var _ environs.Instance = (*instance)(nil)

func (inst *instance) Id() string {
	return inst.InstanceId
}

func (inst *instance) DNSName() string {
	return inst.Instance.DNSName
}

func (environProvider) Open(name string, config interface{}) (e environs.Environ, err error) {
	cfg := config.(*providerConfig)
	if Regions[cfg.region].EC2Endpoint == "" {
		return nil, fmt.Errorf("no ec2 endpoint found for region %q, opening %q", cfg.region, name)
	}
	return &environ{
		name:   name,
		config: cfg,
		ec2:    ec2.New(cfg.auth, Regions[cfg.region]),
		s3:     s3.New(cfg.auth, Regions[cfg.region]),
	}, nil
}

func (e *environ) Bootstrap() error {
	_, err := e.loadState()
	if err == nil {
		return fmt.Errorf("environment is already bootstrapped")
	}
	if s3err, _ := err.(*s3.Error); s3err != nil && s3err.StatusCode != 404 {
		return err
	}
	inst, err := e.startInstance(0, true)
	if err != nil {
		return fmt.Errorf("cannot start bootstrap instance: %v", err)
	}
	err = e.saveState(&bootstrapState{
		ZookeeperInstances: []string{inst.Id()},
	})
	if err != nil {
		// ignore error on StopInstance because the previous error is
		// more important.
		e.StopInstances([]environs.Instance{inst})
		return err
	}
	// TODO make safe in the case of racing Bootstraps
	// If two Bootstraps are called concurrently, there's
	// no way to use S3 to make sure that only one succeeds.
	// Perhaps consider using SimpleDB for state storage
	// which would enable that possibility.
	return nil
}

func (e *environ) Zookeepers() ([]string, error) {
	state, err := e.loadState()
	if err != nil {
		return nil, err
	}
	f := ec2.NewFilter()
	f.Add("instance-id", state.ZookeeperInstances...)
	resp, err := e.ec2.Instances(nil, f)
	if err != nil {
		return nil, err
	}
	var addrs []string
	for _, r := range resp.Reservations {
		for _, inst := range r.Instances {
			addrs = append(addrs, inst.DNSName + zkPortSuffix)
		}
	}
	return addrs, nil
}

func (e *environ) StartInstance(machineId int) (environs.Instance, error) {
	return e.startInstance(machineId, false)
}

// startInstance is the internal version of StartInstance, used by Bootstrap
// as well as via StartInstance itself. If master is true, a bootstrap
// instance will be started.
func (e *environ) startInstance(machineId int, master bool) (environs.Instance, error) {
	image, err := FindImageSpec(DefaultImageConstraint)
	if err != nil {
		return nil, fmt.Errorf("cannot find image: %v", err)
	}
	groups, err := e.setUpGroups(machineId)
	if err != nil {
		return nil, fmt.Errorf("cannot set up groups: %v", err)
	}
	instances, err := e.ec2.RunInstances(&ec2.RunInstances{
		ImageId:        image.ImageId,
		MinCount:       1,
		MaxCount:       1,
		UserData:       nil,
		InstanceType:   "m1.small",
		SecurityGroups: groups,
	})
	if err != nil {
		return nil, fmt.Errorf("cannot run instances: %v", err)
	}
	if len(instances.Instances) != 1 {
		return nil, fmt.Errorf("expected 1 started instance, got %d", len(instances.Instances))
	}
	return &instance{&instances.Instances[0]}, nil
}

func (e *environ) StopInstances(insts []environs.Instance) error {
	if len(insts) == 0 {
		return nil
	}
	names := make([]string, len(insts))
	for i, inst := range insts {
		names[i] = inst.(*instance).InstanceId
	}
	_, err := e.ec2.TerminateInstances(names)
	return err
}

func (e *environ) Instances() ([]environs.Instance, error) {
	filter := ec2.NewFilter()
	filter.Add("instance-state-name", "pending", "running")
	filter.Add("group-name", e.groupName())

	resp, err := e.ec2.Instances(nil, filter)
	if err != nil {
		return nil, err
	}
	var insts []environs.Instance
	for i := range resp.Reservations {
		r := &resp.Reservations[i]
		for j := range r.Instances {
			insts = append(insts, &instance{&r.Instances[j]})
		}
	}
	return insts, nil
}

// oneError returns err1 if it's not not nil, otherwise err2.
func oneError(err1, err2 error) error {
	if err1 != nil {
		return err1
	}
	return err2
}

func (e *environ) Destroy() error {
	delErr := e.deleteState()

	insts, err := e.Instances()
	if err != nil {
		return oneError(delErr, err)
	}
	err = e.StopInstances(insts)
	// return the error from stopping the instances by preference,
	// because instances are more expensive than buckets.
	return oneError(err, delErr)
}

func (e *environ) machineGroupName(machineId int) string {
	return fmt.Sprintf("%s-%d", e.groupName(), machineId)
}

func (e *environ) groupName() string {
	return "juju-" + e.name
}

// setUpGroups creates the security groups for the new machine, and
// returns them.
//
// Instances are tagged with a group so they can be distinguished from
// other instances that might be running on the same EC2 account.  In
// addition, a specific machine security group is created for each
// machine, so that its firewall rules can be configured per machine.
func (e *environ) setUpGroups(machineId int) ([]ec2.SecurityGroup, error) {
	jujuGroup := ec2.SecurityGroup{Name: e.groupName()}
	jujuMachineGroup := ec2.SecurityGroup{Name: e.machineGroupName(machineId)}

	f := ec2.NewFilter()
	f.Add("group-name", jujuGroup.Name, jujuMachineGroup.Name)
	groups, err := e.ec2.SecurityGroups(nil, f)
	if err != nil {
		return nil, fmt.Errorf("cannot get security groups: %v", err)
	}

	for _, g := range groups.Groups {
		switch g.Name {
		case jujuGroup.Name:
			jujuGroup = g.SecurityGroup
		case jujuMachineGroup.Name:
			jujuMachineGroup = g.SecurityGroup
		}
	}

	// Create the provider group if doesn't exist.
	if jujuGroup.Id == "" {
		r, err := e.ec2.CreateSecurityGroup(jujuGroup.Name, "juju group for "+e.name)
		if err != nil {
			return nil, fmt.Errorf("cannot create juju security group: %v", err)
		}
		jujuGroup = r.SecurityGroup
	}

	// Create the machine-specific group, but first see if there's
	// one already existing from a previous machine launch;
	// if so, delete it, since it can have the wrong firewall setup
	if jujuMachineGroup.Id != "" {
		_, err := e.ec2.DeleteSecurityGroup(jujuMachineGroup)
		if err != nil {
			return nil, fmt.Errorf("cannot delete old security group %q: %v", jujuMachineGroup.Name, err)
		}
	}
	descr := fmt.Sprintf("juju group for %s machine %d", e.name, machineId)
	r, err := e.ec2.CreateSecurityGroup(jujuMachineGroup.Name, descr)
	if err != nil {
		return nil, fmt.Errorf("cannot create machine group %q: %v", jujuMachineGroup.Name, err)
	}
	return []ec2.SecurityGroup{jujuGroup, r.SecurityGroup}, nil
}
<commit_msg>start to create cloudconfig data<commit_after>package ec2

import (
	"fmt"
	"launchpad.net/goamz/ec2"
	"launchpad.net/goamz/s3"
	"launchpad.net/juju/go/environs"
	"sync"
)

const zkPortSuffix = ":2181"

func init() {
	environs.RegisterProvider("ec2", environProvider{})
}

type environProvider struct{}

var _ environs.EnvironProvider = environProvider{}

type environ struct {
	name             string
	config           *providerConfig
	ec2              *ec2.EC2
	s3               *s3.S3
	checkBucket      sync.Once
	checkBucketError error
}

var _ environs.Environ = (*environ)(nil)

type instance struct {
	*ec2.Instance
}

func (inst *instance) String() string {
	return inst.Id()
}

var _ environs.Instance = (*instance)(nil)

func (inst *instance) Id() string {
	return inst.InstanceId
}

func (inst *instance) DNSName() string {
	return inst.Instance.DNSName
}

func (environProvider) Open(name string, config interface{}) (e environs.Environ, err error) {
	cfg := config.(*providerConfig)
	if Regions[cfg.region].EC2Endpoint == "" {
		return nil, fmt.Errorf("no ec2 endpoint found for region %q, opening %q", cfg.region, name)
	}
	return &environ{
		name:   name,
		config: cfg,
		ec2:    ec2.New(cfg.auth, Regions[cfg.region]),
		s3:     s3.New(cfg.auth, Regions[cfg.region]),
	}, nil
}

func (e *environ) Bootstrap() error {
	_, err := e.loadState()
	if err == nil {
		return fmt.Errorf("environment is already bootstrapped")
	}
	if s3err, _ := err.(*s3.Error); s3err != nil && s3err.StatusCode != 404 {
		return err
	}
	inst, err := e.startInstance(0, true)
	if err != nil {
		return fmt.Errorf("cannot start bootstrap instance: %v", err)
	}
	err = e.saveState(&bootstrapState{
		ZookeeperInstances: []string{inst.Id()},
	})
	if err != nil {
		// ignore error on StopInstance because the previous error is
		// more important.
		e.StopInstances([]environs.Instance{inst})
		return err
	}
	// TODO make safe in the case of racing Bootstraps
	// If two Bootstraps are called concurrently, there's
	// no way to use S3 to make sure that only one succeeds.
	// Perhaps consider using SimpleDB for state storage
	// which would enable that possibility.
	return nil
}

func (e *environ) Zookeepers() ([]string, error) {
	state, err := e.loadState()
	if err != nil {
		return nil, err
	}
	f := ec2.NewFilter()
	f.Add("instance-id", state.ZookeeperInstances...)
	resp, err := e.ec2.Instances(nil, f)
	if err != nil {
		return nil, err
	}
	var addrs []string
	for _, r := range resp.Reservations {
		for _, inst := range r.Instances {
			addrs = append(addrs, inst.DNSName + zkPortSuffix)
		}
	}
	return addrs, nil
}

func (e *environ) StartInstance(machineId int) (environs.Instance, error) {
	return e.startInstance(machineId, false)
}

// startInstance is the internal version of StartInstance, used by Bootstrap
// as well as via StartInstance itself. If master is true, a bootstrap
// instance will be started.
func (e *environ) startInstance(machineId int, master bool) (environs.Instance, error) {
	image, err := FindImageSpec(DefaultImageConstraint)
	if err != nil {
		return nil, fmt.Errorf("cannot find image: %v", err)
	}
	cfg := &cloudConfig{

	}
	groups, err := e.setUpGroups(machineId)
	if err != nil {
		return nil, fmt.Errorf("cannot set up groups: %v", err)
	}
	instances, err := e.ec2.RunInstances(&ec2.RunInstances{
		ImageId:        image.ImageId,
		MinCount:       1,
		MaxCount:       1,
		UserData:       nil,
		InstanceType:   "m1.small",
		SecurityGroups: groups,
	})
	if err != nil {
		return nil, fmt.Errorf("cannot run instances: %v", err)
	}
	if len(instances.Instances) != 1 {
		return nil, fmt.Errorf("expected 1 started instance, got %d", len(instances.Instances))
	}
	return &instance{&instances.Instances[0]}, nil
}

func (e *environ) StopInstances(insts []environs.Instance) error {
	if len(insts) == 0 {
		return nil
	}
	names := make([]string, len(insts))
	for i, inst := range insts {
		names[i] = inst.(*instance).InstanceId
	}
	_, err := e.ec2.TerminateInstances(names)
	return err
}

func (e *environ) Instances() ([]environs.Instance, error) {
	filter := ec2.NewFilter()
	filter.Add("instance-state-name", "pending", "running")
	filter.Add("group-name", e.groupName())

	resp, err := e.ec2.Instances(nil, filter)
	if err != nil {
		return nil, err
	}
	var insts []environs.Instance
	for i := range resp.Reservations {
		r := &resp.Reservations[i]
		for j := range r.Instances {
			insts = append(insts, &instance{&r.Instances[j]})
		}
	}
	return insts, nil
}

// oneError returns err1 if it's not not nil, otherwise err2.
func oneError(err1, err2 error) error {
	if err1 != nil {
		return err1
	}
	return err2
}

func (e *environ) Destroy() error {
	delErr := e.deleteState()

	insts, err := e.Instances()
	if err != nil {
		return oneError(delErr, err)
	}
	err = e.StopInstances(insts)
	// return the error from stopping the instances by preference,
	// because instances are more expensive than buckets.
	return oneError(err, delErr)
}

func (e *environ) machineGroupName(machineId int) string {
	return fmt.Sprintf("%s-%d", e.groupName(), machineId)
}

func (e *environ) groupName() string {
	return "juju-" + e.name
}

// setUpGroups creates the security groups for the new machine, and
// returns them.
//
// Instances are tagged with a group so they can be distinguished from
// other instances that might be running on the same EC2 account.  In
// addition, a specific machine security group is created for each
// machine, so that its firewall rules can be configured per machine.
func (e *environ) setUpGroups(machineId int) ([]ec2.SecurityGroup, error) {
	jujuGroup := ec2.SecurityGroup{Name: e.groupName()}
	jujuMachineGroup := ec2.SecurityGroup{Name: e.machineGroupName(machineId)}

	f := ec2.NewFilter()
	f.Add("group-name", jujuGroup.Name, jujuMachineGroup.Name)
	groups, err := e.ec2.SecurityGroups(nil, f)
	if err != nil {
		return nil, fmt.Errorf("cannot get security groups: %v", err)
	}

	for _, g := range groups.Groups {
		switch g.Name {
		case jujuGroup.Name:
			jujuGroup = g.SecurityGroup
		case jujuMachineGroup.Name:
			jujuMachineGroup = g.SecurityGroup
		}
	}

	// Create the provider group if doesn't exist.
	if jujuGroup.Id == "" {
		r, err := e.ec2.CreateSecurityGroup(jujuGroup.Name, "juju group for "+e.name)
		if err != nil {
			return nil, fmt.Errorf("cannot create juju security group: %v", err)
		}
		jujuGroup = r.SecurityGroup
	}

	// Create the machine-specific group, but first see if there's
	// one already existing from a previous machine launch;
	// if so, delete it, since it can have the wrong firewall setup
	if jujuMachineGroup.Id != "" {
		_, err := e.ec2.DeleteSecurityGroup(jujuMachineGroup)
		if err != nil {
			return nil, fmt.Errorf("cannot delete old security group %q: %v", jujuMachineGroup.Name, err)
		}
	}
	descr := fmt.Sprintf("juju group for %s machine %d", e.name, machineId)
	r, err := e.ec2.CreateSecurityGroup(jujuMachineGroup.Name, descr)
	if err != nil {
		return nil, fmt.Errorf("cannot create machine group %q: %v", jujuMachineGroup.Name, err)
	}
	return []ec2.SecurityGroup{jujuGroup, r.SecurityGroup}, nil
}
<|endoftext|>


<commit_before>package poloniex

import (
	"encoding/json"
	"log"
	"strconv"

	"github.com/k0kubun/pp"
	"github.com/pkg/errors"
)

const (
	apiURL = "wss://api2.poloniex.com/"
)

var (
	_ChannelIDs = map[string]string{
		"trollbox":  "1001",
		"ticker":    "1002",
		"footer":    "1003",
		"heartbeat": "1010",
	}
)

type (
	//WSTicker describes a ticker item
	WSTicker struct {
		Pair          string
		Last          float64
		Ask           float64
		Bid           float64
		PercentChange float64
		BaseVolume    float64
		QuoteVolume   float64
		IsFrozen      bool
		DailyHigh     float64
		DailyLow      float64
		PairID        int64
	}

	WSOrderbook struct {
	}
)

func (p *Poloniex) StartWS() {
	go func() {
		for {
			_, raw, err := p.ws.ReadMessage()
			if err != nil {
				log.Println("read:", err)
				continue
			}
			var message []interface{}
			err = json.Unmarshal(raw, &message)
			if err != nil {
				log.Println(err)
				continue
			}
			chid := int64(message[0].(float64))
			chids := toString(chid)
			if chid > 100.0 && chid < 1000.0 {
				// it's an orderbook
				orderbook, err := p.parseOrderbook(message)
				if err != nil {
					log.Println(err)
					continue
				}
				p.Emit("orderbook", orderbook)
			} else if chids == _ChannelIDs["ticker"] {
				// it's a ticker
				ticker, err := p.parseTicker(message)
				if err != nil {
					log.Println(err)
					continue
				}
				p.Emit("ticker", ticker)
			}
		}
	}()
}

func (p *Poloniex) Subscribe(chid string) error {
	ticker, err := p.Ticker()
	if err != nil {
		return errors.Wrap(err, "getting ticker for subscribe failed")
	}

	if c, ok := _ChannelIDs[chid]; ok {
		chid = c
	} else if t, ok := ticker[chid]; ok {
		chid = strconv.Itoa(int(t.ID))
	} else {
		return errors.New("unrecognised channelid in subscribe")
	}

	p.subscriptions[chid] = true
	message := subscription{Command: "subscribe", Channel: chid}
	return p.sendWSMessage(message)
}

func (p *Poloniex) Unsubscribe(chid string) error {
	ticker, err := p.Ticker()
	if err != nil {
		return errors.Wrap(err, "getting ticker for unsubscribe failed")
	}

	if c, ok := _ChannelIDs[chid]; ok {
		chid = c
	} else if t, ok := ticker[chid]; ok {
		chid = strconv.Itoa(int(t.ID))
	} else {
		return errors.New("unrecognised channelid in unsubscribe")
	}
	message := subscription{Command: "subscribe", Channel: chid}
	delete(p.subscriptions, chid)
	return p.sendWSMessage(message)
}

func (p *Poloniex) parseTicker(raw []interface{}) (WSTicker, error) {
	wt := WSTicker{}
	pp.Println(raw)
	var rawInner []interface{}
	if len(raw) <= 2 {
		return wt, errors.New("cannot parse to ticker")
	}
	rawInner = raw[2].([]interface{})
	marketID := int64(toFloat(rawInner[0]))
	pair, ok := p.byID[marketID]
	if !ok {
		return wt, errors.New("cannot parse to ticker - invalid marketID")
	}

	wt.Pair = pair
	wt.PairID = marketID
	wt.Last = toFloat(rawInner[1])
	wt.Ask = toFloat(rawInner[2])
	wt.Bid = toFloat(rawInner[3])
	wt.PercentChange = toFloat(rawInner[4])
	wt.BaseVolume = toFloat(rawInner[5])
	wt.QuoteVolume = toFloat(rawInner[6])
	wt.IsFrozen = toFloat(rawInner[7]) != 0.0
	wt.DailyHigh = toFloat(rawInner[8])
	wt.DailyLow = toFloat(rawInner[9])

	return wt, nil
}

func (p *Poloniex) parseOrderbook(raw []interface{}) (WSOrderbook, error) {
	return WSOrderbook{}, nil
}
<commit_msg>halfway thru orderbook feed<commit_after>package poloniex

import (
	"encoding/json"
	"log"
	"strconv"

	"github.com/k0kubun/pp"

	"github.com/pkg/errors"
)

const (
	apiURL = "wss://api2.poloniex.com/"
)

var (
	_ChannelIDs = map[string]string{
		"trollbox":  "1001",
		"ticker":    "1002",
		"footer":    "1003",
		"heartbeat": "1010",
	}
)

type (
	//WSTicker describes a ticker item
	WSTicker struct {
		Pair          string
		Last          float64
		Ask           float64
		Bid           float64
		PercentChange float64
		BaseVolume    float64
		QuoteVolume   float64
		IsFrozen      bool
		DailyHigh     float64
		DailyLow      float64
		PairID        int64
	}

	WSOrderbook struct {
	}
)

func (p *Poloniex) StartWS() {
	go func() {
		for {
			_, raw, err := p.ws.ReadMessage()
			if err != nil {
				log.Println("read:", err)
				continue
			}
			var message []interface{}
			err = json.Unmarshal(raw, &message)
			if err != nil {
				log.Println(err)
				continue
			}
			chid := int64(message[0].(float64))
			chids := toString(chid)
			if chid > 100.0 && chid < 1000.0 {
				// it's an orderbook
				orderbook, err := p.parseOrderbook(message)
				if err != nil {
					log.Println(err)
					continue
				}
				p.Emit("orderbook", orderbook)
			} else if chids == _ChannelIDs["ticker"] {
				// it's a ticker
				ticker, err := p.parseTicker(message)
				if err != nil {
					log.Println(err)
					continue
				}
				p.Emit("ticker", ticker)
			}
		}
	}()
}

func (p *Poloniex) Subscribe(chid string) error {
	ticker, err := p.Ticker()
	if err != nil {
		return errors.Wrap(err, "getting ticker for subscribe failed")
	}

	if c, ok := _ChannelIDs[chid]; ok {
		chid = c
	} else if t, ok := ticker[chid]; ok {
		chid = strconv.Itoa(int(t.ID))
	} else {
		return errors.New("unrecognised channelid in subscribe")
	}

	p.subscriptions[chid] = true
	message := subscription{Command: "subscribe", Channel: chid}
	return p.sendWSMessage(message)
}

func (p *Poloniex) Unsubscribe(chid string) error {
	ticker, err := p.Ticker()
	if err != nil {
		return errors.Wrap(err, "getting ticker for unsubscribe failed")
	}

	if c, ok := _ChannelIDs[chid]; ok {
		chid = c
	} else if t, ok := ticker[chid]; ok {
		chid = strconv.Itoa(int(t.ID))
	} else {
		return errors.New("unrecognised channelid in unsubscribe")
	}
	message := subscription{Command: "subscribe", Channel: chid}
	delete(p.subscriptions, chid)
	return p.sendWSMessage(message)
}

func (p *Poloniex) parseTicker(raw []interface{}) (WSTicker, error) {
	wt := WSTicker{}
	var rawInner []interface{}
	if len(raw) <= 2 {
		return wt, errors.New("cannot parse to ticker")
	}
	rawInner = raw[2].([]interface{})
	marketID := int64(toFloat(rawInner[0]))
	pair, ok := p.byID[marketID]
	if !ok {
		return wt, errors.New("cannot parse to ticker - invalid marketID")
	}

	wt.Pair = pair
	wt.PairID = marketID
	wt.Last = toFloat(rawInner[1])
	wt.Ask = toFloat(rawInner[2])
	wt.Bid = toFloat(rawInner[3])
	wt.PercentChange = toFloat(rawInner[4])
	wt.BaseVolume = toFloat(rawInner[5])
	wt.QuoteVolume = toFloat(rawInner[6])
	wt.IsFrozen = toFloat(rawInner[7]) != 0.0
	wt.DailyHigh = toFloat(rawInner[8])
	wt.DailyLow = toFloat(rawInner[9])

	return wt, nil
}

func (p *Poloniex) parseOrderbook(raw []interface{}) (WSOrderbook, error) {
	wo := WSOrderbook{}
	marketID := int64(toFloat(raw[0]))
	pair, ok := p.byID[marketID]
	if !ok {
		return wo, errors.New("cannot parse to orderbook - invalid marketID")
	}
	pp.Println(pair, marketID)
	return WSOrderbook{}, nil
}
<|endoftext|>


<commit_before>package main

import (
	"os"
	"sync"

	"github.com/akhenakh/hunspellgo"
	"github.com/trustmaster/go-aspell"
)

type Speller interface {
	Check(w string) bool
	Suggest(w string) []string
	Close()
}

type aspeller struct {
	sp aspell.Speller
	mu sync.Mutex
}

func NewASpeller() (Speller, error) {
	opts := map[string]string{
		"lang":           "en",
		"filter":         "url",
		"mode":           "url",
		"encoding":       "ascii",
		"guess":          "true",
		"ignore":         "0",
		"ignore-case":    "false",
		"ignore-accents": "false",
	}

	sp, err := aspell.NewSpeller(opts)
	if err != nil {
		return nil, err
	}
	return &aspeller{sp: sp}, nil
}

func (s *aspeller) Check(w string) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.sp.Check(w)
}

func (s *aspeller) Suggest(w string) []string {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.sp.Suggest(w)
}

func (s *aspeller) Close() { s.sp.Delete() }

type hunspeller struct {
	sp *hunspellgo.Hunhandle
}

// hunspellPaths has a list of paths where dictionaries might be
// in future, pass path through command line argument
var hunspellPaths = []string{
	"/usr/share/myspell",
	"/usr/share/hunspell",
}

func NewHunSpeller() Speller {
	for _, p := range hunspellPaths {
		aff, dic := p+"/en_US.aff", p+"/en_US.dic"
		if _, err := os.Stat(aff); err != nil {
			continue
		}
		if _, err := os.Stat(dic); err != nil {
			continue
		}
		if sp := hunspellgo.Hunspell(aff, dic); sp != nil {
			return &hunspeller{sp}
		}
	}
	return nil
}

func (s *hunspeller) Check(w string) bool {
	return s.sp.Spell(w)
}

func (s *hunspeller) Suggest(w string) []string {
	return s.sp.Suggest(w)
}

func (s *hunspeller) Close() { s.sp = nil }

type multispeller struct {
	sp []Speller
}

func NewMultiSpeller(sp ...Speller) Speller {
	m := &multispeller{}
	for _, s := range sp {
		m.sp = append(m.sp, s)
	}
	return m
}

func (s *multispeller) Check(w string) bool {
	for _, sp := range s.sp {
		if sp.Check(w) {
			return true
		}
	}
	return false
}

func (s *multispeller) Suggest(w string) (ret []string) {
	for _, sp := range s.sp {
		if ret = sp.Suggest(w); len(ret) != 0 {
			break
		}
	}
	return ret
}

func (s *multispeller) Close() {
	for _, sp := range s.sp {
		sp.Close()
	}
	s.sp = nil
}
<commit_msg>build: disable aspell/hunspell imports on !spell<commit_after>// +build spell

package main

import (
	"os"
	"sync"

	"github.com/akhenakh/hunspellgo"
	"github.com/trustmaster/go-aspell"
)

type Speller interface {
	Check(w string) bool
	Suggest(w string) []string
	Close()
}

type aspeller struct {
	sp aspell.Speller
	mu sync.Mutex
}

func NewASpeller() (Speller, error) {
	opts := map[string]string{
		"lang":           "en",
		"filter":         "url",
		"mode":           "url",
		"encoding":       "ascii",
		"guess":          "true",
		"ignore":         "0",
		"ignore-case":    "false",
		"ignore-accents": "false",
	}

	sp, err := aspell.NewSpeller(opts)
	if err != nil {
		return nil, err
	}
	return &aspeller{sp: sp}, nil
}

func (s *aspeller) Check(w string) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.sp.Check(w)
}

func (s *aspeller) Suggest(w string) []string {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.sp.Suggest(w)
}

func (s *aspeller) Close() { s.sp.Delete() }

type hunspeller struct {
	sp *hunspellgo.Hunhandle
}

// hunspellPaths has a list of paths where dictionaries might be
// in future, pass path through command line argument
var hunspellPaths = []string{
	"/usr/share/myspell",
	"/usr/share/hunspell",
}

func NewHunSpeller() Speller {
	for _, p := range hunspellPaths {
		aff, dic := p+"/en_US.aff", p+"/en_US.dic"
		if _, err := os.Stat(aff); err != nil {
			continue
		}
		if _, err := os.Stat(dic); err != nil {
			continue
		}
		if sp := hunspellgo.Hunspell(aff, dic); sp != nil {
			return &hunspeller{sp}
		}
	}
	return nil
}

func (s *hunspeller) Check(w string) bool {
	return s.sp.Spell(w)
}

func (s *hunspeller) Suggest(w string) []string {
	return s.sp.Suggest(w)
}

func (s *hunspeller) Close() { s.sp = nil }

type multispeller struct {
	sp []Speller
}

func NewMultiSpeller(sp ...Speller) Speller {
	m := &multispeller{}
	for _, s := range sp {
		m.sp = append(m.sp, s)
	}
	return m
}

func (s *multispeller) Check(w string) bool {
	for _, sp := range s.sp {
		if sp.Check(w) {
			return true
		}
	}
	return false
}

func (s *multispeller) Suggest(w string) (ret []string) {
	for _, sp := range s.sp {
		if ret = sp.Suggest(w); len(ret) != 0 {
			break
		}
	}
	return ret
}

func (s *multispeller) Close() {
	for _, sp := range s.sp {
		sp.Close()
	}
	s.sp = nil
}
<|endoftext|>


<commit_before>package spigot

func spigot(a, r, q int, carry <-chan int) <-chan int {
	c := make(chan int, 100)
	go func() {
		for cr := range carry {
			a = 10*a + cr
			c <- r * (a / q)
			a %= q
		}
		close(c)
	}()
	return c
}

func zero(n int) <-chan int {
	c := make(chan int)
	go func() {
		for i := 0; i < n; i++ {
			c <- 0
		}
		close(c)
	}()
	return c
}

// Pi2 calculates n digits of Pi concurently
func Pi2(n int) <-chan int {
	c := zero(n + 1)
	cr := c
	for i := 10*n/3 + 1; i > 0; i-- {
		cr = spigot(2, i, 2*i+1, cr)
	}
	return predigit(spigot(2, 1, 10, cr))
}

// E2 calculates n digits of E concurently
func E2(n int) <-chan int {
	c := zero(n + 1)
	cr := c
	for i := n + 1; i > 0; i-- {
		cr = spigot(1, 1, i+1, cr)
	}
	return spigot(2, 1, 10, cr)
}
<commit_msg>Defer<commit_after>package spigot

func spigot(a, r, q int, carry <-chan int) <-chan int {
	c := make(chan int, 100)
	go func() {
		defer close(c)
		for cr := range carry {
			a = 10*a + cr
			c <- r * (a / q)
			a %= q
		}
	}()
	return c
}

func zero(n int) <-chan int {
	c := make(chan int)
	go func() {
		defer close(c)
		for i := 0; i < n; i++ {
			c <- 0
		}
	}()
	return c
}

// Pi2 calculates n digits of Pi concurently
func Pi2(n int) <-chan int {
	c := zero(n + 1)
	cr := c
	for i := 10*n/3 + 1; i > 0; i-- {
		cr = spigot(2, i, 2*i+1, cr)
	}
	return predigit(spigot(2, 1, 10, cr))
}

// E2 calculates n digits of E concurently
func E2(n int) <-chan int {
	c := zero(n + 1)
	cr := c
	for i := n + 1; i > 0; i-- {
		cr = spigot(1, 1, i+1, cr)
	}
	return spigot(2, 1, 10, cr)
}
<|endoftext|>


<commit_before>// Copyright (c) 2021 Brian J. Downs
//
// 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 spinner is a simple package to add a spinner / progress indicator to any terminal application.
package spinner

import (
	"errors"
	"fmt"
	"io"
	"os"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"
	"unicode/utf8"

	"github.com/fatih/color"
	"github.com/mattn/go-isatty"
)

// errInvalidColor is returned when attempting to set an invalid color
var errInvalidColor = errors.New("invalid color")

// validColors holds an array of the only colors allowed
var validColors = map[string]bool{
	// default colors for backwards compatibility
	"black":   true,
	"red":     true,
	"green":   true,
	"yellow":  true,
	"blue":    true,
	"magenta": true,
	"cyan":    true,
	"white":   true,

	// attributes
	"reset":        true,
	"bold":         true,
	"faint":        true,
	"italic":       true,
	"underline":    true,
	"blinkslow":    true,
	"blinkrapid":   true,
	"reversevideo": true,
	"concealed":    true,
	"crossedout":   true,

	// foreground text
	"fgBlack":   true,
	"fgRed":     true,
	"fgGreen":   true,
	"fgYellow":  true,
	"fgBlue":    true,
	"fgMagenta": true,
	"fgCyan":    true,
	"fgWhite":   true,

	// foreground Hi-Intensity text
	"fgHiBlack":   true,
	"fgHiRed":     true,
	"fgHiGreen":   true,
	"fgHiYellow":  true,
	"fgHiBlue":    true,
	"fgHiMagenta": true,
	"fgHiCyan":    true,
	"fgHiWhite":   true,

	// background text
	"bgBlack":   true,
	"bgRed":     true,
	"bgGreen":   true,
	"bgYellow":  true,
	"bgBlue":    true,
	"bgMagenta": true,
	"bgCyan":    true,
	"bgWhite":   true,

	// background Hi-Intensity text
	"bgHiBlack":   true,
	"bgHiRed":     true,
	"bgHiGreen":   true,
	"bgHiYellow":  true,
	"bgHiBlue":    true,
	"bgHiMagenta": true,
	"bgHiCyan":    true,
	"bgHiWhite":   true,
}

// returns true if the OS is windows and the WT_SESSION env variable is set.
var isWindows = runtime.OSOS == "windows"
var isWindowsTerminalOnWindows = len(os.Getenv("WT_SESSION")) > 0 && isWindows

// returns a valid color's foreground text color attribute
var colorAttributeMap = map[string]color.Attribute{
	// default colors for backwards compatibility
	"black":   color.FgBlack,
	"red":     color.FgRed,
	"green":   color.FgGreen,
	"yellow":  color.FgYellow,
	"blue":    color.FgBlue,
	"magenta": color.FgMagenta,
	"cyan":    color.FgCyan,
	"white":   color.FgWhite,

	// attributes
	"reset":        color.Reset,
	"bold":         color.Bold,
	"faint":        color.Faint,
	"italic":       color.Italic,
	"underline":    color.Underline,
	"blinkslow":    color.BlinkSlow,
	"blinkrapid":   color.BlinkRapid,
	"reversevideo": color.ReverseVideo,
	"concealed":    color.Concealed,
	"crossedout":   color.CrossedOut,

	// foreground text colors
	"fgBlack":   color.FgBlack,
	"fgRed":     color.FgRed,
	"fgGreen":   color.FgGreen,
	"fgYellow":  color.FgYellow,
	"fgBlue":    color.FgBlue,
	"fgMagenta": color.FgMagenta,
	"fgCyan":    color.FgCyan,
	"fgWhite":   color.FgWhite,

	// foreground Hi-Intensity text colors
	"fgHiBlack":   color.FgHiBlack,
	"fgHiRed":     color.FgHiRed,
	"fgHiGreen":   color.FgHiGreen,
	"fgHiYellow":  color.FgHiYellow,
	"fgHiBlue":    color.FgHiBlue,
	"fgHiMagenta": color.FgHiMagenta,
	"fgHiCyan":    color.FgHiCyan,
	"fgHiWhite":   color.FgHiWhite,

	// background text colors
	"bgBlack":   color.BgBlack,
	"bgRed":     color.BgRed,
	"bgGreen":   color.BgGreen,
	"bgYellow":  color.BgYellow,
	"bgBlue":    color.BgBlue,
	"bgMagenta": color.BgMagenta,
	"bgCyan":    color.BgCyan,
	"bgWhite":   color.BgWhite,

	// background Hi-Intensity text colors
	"bgHiBlack":   color.BgHiBlack,
	"bgHiRed":     color.BgHiRed,
	"bgHiGreen":   color.BgHiGreen,
	"bgHiYellow":  color.BgHiYellow,
	"bgHiBlue":    color.BgHiBlue,
	"bgHiMagenta": color.BgHiMagenta,
	"bgHiCyan":    color.BgHiCyan,
	"bgHiWhite":   color.BgHiWhite,
}

// validColor will make sure the given color is actually allowed.
func validColor(c string) bool {
	return validColors[c]
}

// Spinner struct to hold the provided options.
type Spinner struct {
	mu              *sync.RWMutex
	Delay           time.Duration                 // Delay is the speed of the indicator
	chars           []string                      // chars holds the chosen character set
	Prefix          string                        // Prefix is the text preppended to the indicator
	Suffix          string                        // Suffix is the text appended to the indicator
	FinalMSG        string                        // string displayed after Stop() is called
	lastOutputPlain string                        // last character(set) written
	LastOutput      string                        // last character(set) written with colors
	color           func(a ...interface{}) string // default color is white
	Writer          io.Writer                     // to make testing better, exported so users have access. Use `WithWriter` to update after initialization.
	active          bool                          // active holds the state of the spinner
	stopChan        chan struct{}                 // stopChan is a channel used to stop the indicator
	HideCursor      bool                          // hideCursor determines if the cursor is visible
	PreUpdate       func(s *Spinner)              // will be triggered before every spinner update
	PostUpdate      func(s *Spinner)              // will be triggered after every spinner update
}

// New provides a pointer to an instance of Spinner with the supplied options.
func New(cs []string, d time.Duration, options ...Option) *Spinner {
	s := &Spinner{
		Delay:      d,
		chars:      cs,
		color:      color.New(color.FgWhite).SprintFunc(),
		mu:         &sync.RWMutex{},
		Writer:     color.Output,
		stopChan:   make(chan struct{}, 1),
		active:     false,
		HideCursor: true,
	}

	for _, option := range options {
		option(s)
	}

	return s
}

// Option is a function that takes a spinner and applies
// a given configuration.
type Option func(*Spinner)

// Options contains fields to configure the spinner.
type Options struct {
	Color      string
	Suffix     string
	FinalMSG   string
	HideCursor bool
}

// WithColor adds the given color to the spinner.
func WithColor(color string) Option {
	return func(s *Spinner) {
		s.Color(color)
	}
}

// WithSuffix adds the given string to the spinner
// as the suffix.
func WithSuffix(suffix string) Option {
	return func(s *Spinner) {
		s.Suffix = suffix
	}
}

// WithFinalMSG adds the given string ot the spinner
// as the final message to be written.
func WithFinalMSG(finalMsg string) Option {
	return func(s *Spinner) {
		s.FinalMSG = finalMsg
	}
}

// WithHiddenCursor hides the cursor
// if hideCursor = true given.
func WithHiddenCursor(hideCursor bool) Option {
	return func(s *Spinner) {
		s.HideCursor = hideCursor
	}
}

// WithWriter adds the given writer to the spinner. This
// function should be favored over directly assigning to
// the struct value.
func WithWriter(w io.Writer) Option {
	return func(s *Spinner) {
		s.mu.Lock()
		s.Writer = w
		s.mu.Unlock()
	}
}

// Active will return whether or not the spinner is currently active.
func (s *Spinner) Active() bool {
	return s.active
}

// Start will start the indicator.
func (s *Spinner) Start() {
	s.mu.Lock()
	if s.active || !isRunningInTerminal() {
		s.mu.Unlock()
		return
	}
	if s.HideCursor && !isWindowsTerminalOnWindows {
		// hides the cursor
		fmt.Fprint(s.Writer, "\033[?25l")
	}
	// Disable colors for simple Windows CMD or Powershell
	// as they can not recognize them
	if isWindows && !isWindowsTerminalOnWindows {
		color.NoColor = true
	}

	s.active = true
	s.mu.Unlock()

	go func() {
		for {
			for i := 0; i < len(s.chars); i++ {
				select {
				case <-s.stopChan:
					return
				default:
					s.mu.Lock()
					if !s.active {
						s.mu.Unlock()
						return
					}
					if !isWindowsTerminalOnWindows {
						s.erase()
					}

					if s.PreUpdate != nil {
						s.PreUpdate(s)
					}

					var outColor string
					if isWindows {
						if s.Writer == os.Stderr {
							outColor = fmt.Sprintf("\r%s%s%s", s.Prefix, s.chars[i], s.Suffix)
						} else {
							outColor = fmt.Sprintf("\r%s%s%s", s.Prefix, s.color(s.chars[i]), s.Suffix)
						}
					} else {
						outColor = fmt.Sprintf("\r%s%s%s", s.Prefix, s.color(s.chars[i]), s.Suffix)
					}
					outPlain := fmt.Sprintf("\r%s%s%s", s.Prefix, s.chars[i], s.Suffix)
					fmt.Fprint(s.Writer, outColor)
					s.lastOutputPlain = outPlain
					s.LastOutput = outColor
					delay := s.Delay

					if s.PostUpdate != nil {
						s.PostUpdate(s)
					}

					s.mu.Unlock()
					time.Sleep(delay)
				}
			}
		}
	}()
}

// Stop stops the indicator.
func (s *Spinner) Stop() {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.active {
		s.active = false
		if s.HideCursor && !isWindowsTerminalOnWindows {
			// makes the cursor visible
			fmt.Fprint(s.Writer, "\033[?25h")
		}
		s.erase()
		if s.FinalMSG != "" {
			if isWindowsTerminalOnWindows {
				fmt.Fprint(s.Writer, "\r", s.FinalMSG)
			} else {
				fmt.Fprint(s.Writer, s.FinalMSG)
			}
		}
		s.stopChan <- struct{}{}
	}
}

// Restart will stop and start the indicator.
func (s *Spinner) Restart() {
	s.Stop()
	s.Start()
}

// Reverse will reverse the order of the slice assigned to the indicator.
func (s *Spinner) Reverse() {
	s.mu.Lock()
	for i, j := 0, len(s.chars)-1; i < j; i, j = i+1, j-1 {
		s.chars[i], s.chars[j] = s.chars[j], s.chars[i]
	}
	s.mu.Unlock()
}

// Color will set the struct field for the given color to be used. The spinner
// will need to be explicitly restarted.
func (s *Spinner) Color(colors ...string) error {
	colorAttributes := make([]color.Attribute, len(colors))

	// Verify colours are valid and place the appropriate attribute in the array
	for index, c := range colors {
		if !validColor(c) {
			return errInvalidColor
		}
		colorAttributes[index] = colorAttributeMap[c]
	}

	s.mu.Lock()
	s.color = color.New(colorAttributes...).SprintFunc()
	s.mu.Unlock()
	return nil
}

// UpdateSpeed will set the indicator delay to the given value.
func (s *Spinner) UpdateSpeed(d time.Duration) {
	s.mu.Lock()
	s.Delay = d
	s.mu.Unlock()
}

// UpdateCharSet will change the current character set to the given one.
func (s *Spinner) UpdateCharSet(cs []string) {
	s.mu.Lock()
	s.chars = cs
	s.mu.Unlock()
}

// erase deletes written characters on the current line.
// Caller must already hold s.lock.
func (s *Spinner) erase() {
	n := utf8.RuneCountInString(s.lastOutputPlain)
	if runtime.GOOS == "windows" && !isWindowsTerminalOnWindows {
		clearString := "\r" + strings.Repeat(" ", n) + "\r"
		fmt.Fprint(s.Writer, clearString)
		s.lastOutputPlain = ""
		return
	}

	// Taken from https://en.wikipedia.org/wiki/ANSI_escape_code:
	// \r     - Carriage return - Moves the cursor to column zero
	// \033[K - Erases part of the line. If n is 0 (or missing), clear from
	// cursor to the end of the line. If n is 1, clear from cursor to beginning
	// of the line. If n is 2, clear entire line. Cursor position does not
	// change.
	fmt.Fprintf(s.Writer, "\r\033[K")
	s.lastOutputPlain = ""
}

// Lock allows for manual control to lock the spinner.
func (s *Spinner) Lock() {
	s.mu.Lock()
}

// Unlock allows for manual control to unlock the spinner.
func (s *Spinner) Unlock() {
	s.mu.Unlock()
}

// GenerateNumberSequence will generate a slice of integers at the
// provided length and convert them each to a string.
func GenerateNumberSequence(length int) []string {
	numSeq := make([]string, length)
	for i := 0; i < length; i++ {
		numSeq[i] = strconv.Itoa(i)
	}
	return numSeq
}

// isRunningInTerminal check if stdout file descriptor is terminal
func isRunningInTerminal() bool {
	return isatty.IsTerminal(os.Stdout.Fd())
}
<commit_msg>Revert "Do not show colors in Windows CMD & Powershell (#138)" (#140)<commit_after>// Copyright (c) 2021 Brian J. Downs
//
// 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 spinner is a simple package to add a spinner / progress indicator to any terminal application.
package spinner

import (
	"errors"
	"fmt"
	"io"
	"os"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"
	"unicode/utf8"

	"github.com/fatih/color"
	"github.com/mattn/go-isatty"
)

// errInvalidColor is returned when attempting to set an invalid color
var errInvalidColor = errors.New("invalid color")

// validColors holds an array of the only colors allowed
var validColors = map[string]bool{
	// default colors for backwards compatibility
	"black":   true,
	"red":     true,
	"green":   true,
	"yellow":  true,
	"blue":    true,
	"magenta": true,
	"cyan":    true,
	"white":   true,

	// attributes
	"reset":        true,
	"bold":         true,
	"faint":        true,
	"italic":       true,
	"underline":    true,
	"blinkslow":    true,
	"blinkrapid":   true,
	"reversevideo": true,
	"concealed":    true,
	"crossedout":   true,

	// foreground text
	"fgBlack":   true,
	"fgRed":     true,
	"fgGreen":   true,
	"fgYellow":  true,
	"fgBlue":    true,
	"fgMagenta": true,
	"fgCyan":    true,
	"fgWhite":   true,

	// foreground Hi-Intensity text
	"fgHiBlack":   true,
	"fgHiRed":     true,
	"fgHiGreen":   true,
	"fgHiYellow":  true,
	"fgHiBlue":    true,
	"fgHiMagenta": true,
	"fgHiCyan":    true,
	"fgHiWhite":   true,

	// background text
	"bgBlack":   true,
	"bgRed":     true,
	"bgGreen":   true,
	"bgYellow":  true,
	"bgBlue":    true,
	"bgMagenta": true,
	"bgCyan":    true,
	"bgWhite":   true,

	// background Hi-Intensity text
	"bgHiBlack":   true,
	"bgHiRed":     true,
	"bgHiGreen":   true,
	"bgHiYellow":  true,
	"bgHiBlue":    true,
	"bgHiMagenta": true,
	"bgHiCyan":    true,
	"bgHiWhite":   true,
}

// returns true if the OS is windows and the WT_SESSION env variable is set.
var isWindowsTerminalOnWindows = len(os.Getenv("WT_SESSION")) > 0 && runtime.GOOS == "windows"

// returns a valid color's foreground text color attribute
var colorAttributeMap = map[string]color.Attribute{
	// default colors for backwards compatibility
	"black":   color.FgBlack,
	"red":     color.FgRed,
	"green":   color.FgGreen,
	"yellow":  color.FgYellow,
	"blue":    color.FgBlue,
	"magenta": color.FgMagenta,
	"cyan":    color.FgCyan,
	"white":   color.FgWhite,

	// attributes
	"reset":        color.Reset,
	"bold":         color.Bold,
	"faint":        color.Faint,
	"italic":       color.Italic,
	"underline":    color.Underline,
	"blinkslow":    color.BlinkSlow,
	"blinkrapid":   color.BlinkRapid,
	"reversevideo": color.ReverseVideo,
	"concealed":    color.Concealed,
	"crossedout":   color.CrossedOut,

	// foreground text colors
	"fgBlack":   color.FgBlack,
	"fgRed":     color.FgRed,
	"fgGreen":   color.FgGreen,
	"fgYellow":  color.FgYellow,
	"fgBlue":    color.FgBlue,
	"fgMagenta": color.FgMagenta,
	"fgCyan":    color.FgCyan,
	"fgWhite":   color.FgWhite,

	// foreground Hi-Intensity text colors
	"fgHiBlack":   color.FgHiBlack,
	"fgHiRed":     color.FgHiRed,
	"fgHiGreen":   color.FgHiGreen,
	"fgHiYellow":  color.FgHiYellow,
	"fgHiBlue":    color.FgHiBlue,
	"fgHiMagenta": color.FgHiMagenta,
	"fgHiCyan":    color.FgHiCyan,
	"fgHiWhite":   color.FgHiWhite,

	// background text colors
	"bgBlack":   color.BgBlack,
	"bgRed":     color.BgRed,
	"bgGreen":   color.BgGreen,
	"bgYellow":  color.BgYellow,
	"bgBlue":    color.BgBlue,
	"bgMagenta": color.BgMagenta,
	"bgCyan":    color.BgCyan,
	"bgWhite":   color.BgWhite,

	// background Hi-Intensity text colors
	"bgHiBlack":   color.BgHiBlack,
	"bgHiRed":     color.BgHiRed,
	"bgHiGreen":   color.BgHiGreen,
	"bgHiYellow":  color.BgHiYellow,
	"bgHiBlue":    color.BgHiBlue,
	"bgHiMagenta": color.BgHiMagenta,
	"bgHiCyan":    color.BgHiCyan,
	"bgHiWhite":   color.BgHiWhite,
}

// validColor will make sure the given color is actually allowed.
func validColor(c string) bool {
	return validColors[c]
}

// Spinner struct to hold the provided options.
type Spinner struct {
	mu              *sync.RWMutex
	Delay           time.Duration                 // Delay is the speed of the indicator
	chars           []string                      // chars holds the chosen character set
	Prefix          string                        // Prefix is the text preppended to the indicator
	Suffix          string                        // Suffix is the text appended to the indicator
	FinalMSG        string                        // string displayed after Stop() is called
	lastOutputPlain string                        // last character(set) written
	LastOutput      string                        // last character(set) written with colors
	color           func(a ...interface{}) string // default color is white
	Writer          io.Writer                     // to make testing better, exported so users have access. Use `WithWriter` to update after initialization.
	active          bool                          // active holds the state of the spinner
	stopChan        chan struct{}                 // stopChan is a channel used to stop the indicator
	HideCursor      bool                          // hideCursor determines if the cursor is visible
	PreUpdate       func(s *Spinner)              // will be triggered before every spinner update
	PostUpdate      func(s *Spinner)              // will be triggered after every spinner update
}

// New provides a pointer to an instance of Spinner with the supplied options.
func New(cs []string, d time.Duration, options ...Option) *Spinner {
	s := &Spinner{
		Delay:      d,
		chars:      cs,
		color:      color.New(color.FgWhite).SprintFunc(),
		mu:         &sync.RWMutex{},
		Writer:     color.Output,
		stopChan:   make(chan struct{}, 1),
		active:     false,
		HideCursor: true,
	}

	for _, option := range options {
		option(s)
	}

	return s
}

// Option is a function that takes a spinner and applies
// a given configuration.
type Option func(*Spinner)

// Options contains fields to configure the spinner.
type Options struct {
	Color      string
	Suffix     string
	FinalMSG   string
	HideCursor bool
}

// WithColor adds the given color to the spinner.
func WithColor(color string) Option {
	return func(s *Spinner) {
		s.Color(color)
	}
}

// WithSuffix adds the given string to the spinner
// as the suffix.
func WithSuffix(suffix string) Option {
	return func(s *Spinner) {
		s.Suffix = suffix
	}
}

// WithFinalMSG adds the given string ot the spinner
// as the final message to be written.
func WithFinalMSG(finalMsg string) Option {
	return func(s *Spinner) {
		s.FinalMSG = finalMsg
	}
}

// WithHiddenCursor hides the cursor
// if hideCursor = true given.
func WithHiddenCursor(hideCursor bool) Option {
	return func(s *Spinner) {
		s.HideCursor = hideCursor
	}
}

// WithWriter adds the given writer to the spinner. This
// function should be favored over directly assigning to
// the struct value.
func WithWriter(w io.Writer) Option {
	return func(s *Spinner) {
		s.mu.Lock()
		s.Writer = w
		s.mu.Unlock()
	}
}

// Active will return whether or not the spinner is currently active.
func (s *Spinner) Active() bool {
	return s.active
}

// Start will start the indicator.
func (s *Spinner) Start() {
	s.mu.Lock()
	if s.active || !isRunningInTerminal() {
		s.mu.Unlock()
		return
	}
	if s.HideCursor && !isWindowsTerminalOnWindows {
		// hides the cursor
		fmt.Fprint(s.Writer, "\033[?25l")
	}
	s.active = true
	s.mu.Unlock()

	go func() {
		for {
			for i := 0; i < len(s.chars); i++ {
				select {
				case <-s.stopChan:
					return
				default:
					s.mu.Lock()
					if !s.active {
						s.mu.Unlock()
						return
					}
					if !isWindowsTerminalOnWindows {
						s.erase()
					}

					if s.PreUpdate != nil {
						s.PreUpdate(s)
					}

					var outColor string
					if runtime.GOOS == "windows" {
						if s.Writer == os.Stderr {
							outColor = fmt.Sprintf("\r%s%s%s", s.Prefix, s.chars[i], s.Suffix)
						} else {
							outColor = fmt.Sprintf("\r%s%s%s", s.Prefix, s.color(s.chars[i]), s.Suffix)
						}
					} else {
						outColor = fmt.Sprintf("\r%s%s%s", s.Prefix, s.color(s.chars[i]), s.Suffix)
					}
					outPlain := fmt.Sprintf("\r%s%s%s", s.Prefix, s.chars[i], s.Suffix)
					fmt.Fprint(s.Writer, outColor)
					s.lastOutputPlain = outPlain
					s.LastOutput = outColor
					delay := s.Delay

					if s.PostUpdate != nil {
						s.PostUpdate(s)
					}

					s.mu.Unlock()
					time.Sleep(delay)
				}
			}
		}
	}()
}

// Stop stops the indicator.
func (s *Spinner) Stop() {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.active {
		s.active = false
		if s.HideCursor && !isWindowsTerminalOnWindows {
			// makes the cursor visible
			fmt.Fprint(s.Writer, "\033[?25h")
		}
		s.erase()
		if s.FinalMSG != "" {
			if isWindowsTerminalOnWindows {
				fmt.Fprint(s.Writer, "\r", s.FinalMSG)
			} else {
				fmt.Fprint(s.Writer, s.FinalMSG)
			}
		}
		s.stopChan <- struct{}{}
	}
}

// Restart will stop and start the indicator.
func (s *Spinner) Restart() {
	s.Stop()
	s.Start()
}

// Reverse will reverse the order of the slice assigned to the indicator.
func (s *Spinner) Reverse() {
	s.mu.Lock()
	for i, j := 0, len(s.chars)-1; i < j; i, j = i+1, j-1 {
		s.chars[i], s.chars[j] = s.chars[j], s.chars[i]
	}
	s.mu.Unlock()
}

// Color will set the struct field for the given color to be used. The spinner
// will need to be explicitly restarted.
func (s *Spinner) Color(colors ...string) error {
	colorAttributes := make([]color.Attribute, len(colors))

	// Verify colours are valid and place the appropriate attribute in the array
	for index, c := range colors {
		if !validColor(c) {
			return errInvalidColor
		}
		colorAttributes[index] = colorAttributeMap[c]
	}

	s.mu.Lock()
	s.color = color.New(colorAttributes...).SprintFunc()
	s.mu.Unlock()
	return nil
}

// UpdateSpeed will set the indicator delay to the given value.
func (s *Spinner) UpdateSpeed(d time.Duration) {
	s.mu.Lock()
	s.Delay = d
	s.mu.Unlock()
}

// UpdateCharSet will change the current character set to the given one.
func (s *Spinner) UpdateCharSet(cs []string) {
	s.mu.Lock()
	s.chars = cs
	s.mu.Unlock()
}

// erase deletes written characters on the current line.
// Caller must already hold s.lock.
func (s *Spinner) erase() {
	n := utf8.RuneCountInString(s.lastOutputPlain)
	if runtime.GOOS == "windows" && !isWindowsTerminalOnWindows {
		clearString := "\r" + strings.Repeat(" ", n) + "\r"
		fmt.Fprint(s.Writer, clearString)
		s.lastOutputPlain = ""
		return
	}

	// Taken from https://en.wikipedia.org/wiki/ANSI_escape_code:
	// \r     - Carriage return - Moves the cursor to column zero
	// \033[K - Erases part of the line. If n is 0 (or missing), clear from
	// cursor to the end of the line. If n is 1, clear from cursor to beginning
	// of the line. If n is 2, clear entire line. Cursor position does not
	// change.
	fmt.Fprintf(s.Writer, "\r\033[K")
	s.lastOutputPlain = ""
}

// Lock allows for manual control to lock the spinner.
func (s *Spinner) Lock() {
	s.mu.Lock()
}

// Unlock allows for manual control to unlock the spinner.
func (s *Spinner) Unlock() {
	s.mu.Unlock()
}

// GenerateNumberSequence will generate a slice of integers at the
// provided length and convert them each to a string.
func GenerateNumberSequence(length int) []string {
	numSeq := make([]string, length)
	for i := 0; i < length; i++ {
		numSeq[i] = strconv.Itoa(i)
	}
	return numSeq
}

// isRunningInTerminal check if stdout file descriptor is terminal
func isRunningInTerminal() bool {
	return isatty.IsTerminal(os.Stdout.Fd())
}
<|endoftext|>


<commit_before>package binder

import (
  "github.com/wpxiong/beargo/log"
  "github.com/wpxiong/beargo/appcontext"
  "reflect"
  "strconv"
  "strings"
  "errors"
)

func init() {
  log.InitLog()
}

var floatType = reflect.TypeOf(float64(0))

func getFloat(unk interface{}) (float64, error) {
    v := reflect.ValueOf(unk)
    v = reflect.Indirect(v)
    if !v.Type().ConvertibleTo(floatType) {
        return 0, errors.New("")
    }
    fv := v.Convert(floatType)
    return fv.Float(), nil
}

func GetParamValueFloat(paramValue interface{}) (float64,error) {
   switch paramValue.(type) {
     case float32:
       return getFloat(paramValue)
     case float64:
       return getFloat(paramValue)
     case string:
       val,err:= strconv.ParseFloat(paramValue.(string),64)
       if err == nil {
          return val,nil
       }else {
          return 0,err
       }
     default:
        return 0,errors.New("")
  }
}

func GetParamValueInt(paramValue interface{}) (int64,error) {
  switch paramValue.(type) {
     case int:
       return int64(paramValue.(int)),nil
     case int8:
       return int64(paramValue.(int8)),nil
     case int16:
       return int64(paramValue.(int16)),nil
     case int32:
       return int64(paramValue.(int32)),nil
     case int64:
       return paramValue.(int64),nil
     case string:
       val,err:= strconv.Atoi(paramValue.(string))
       if err == nil {
          return int64(val),nil
       }else {
          return 0,err
       }
     default:
        return 0,errors.New("")
  }
}

func GetParamValueUint(paramValue interface{}) (uint64,error) {
  switch paramValue.(type) {
     case uint,uint8,uint16,uint32,uint64:
       return paramValue.(uint64),nil
     case string:
       val,err:= strconv.Atoi(paramValue.(string))
       if err == nil {
          return uint64(val),nil
       }else {
          return 0,err
       }
     default:
        return 0,errors.New("")
  }
}

func BinderByType(f reflect.Value, ft reflect.Type, param map[string] interface{},name string){
  switch ft.Kind() {
    case reflect.Map:
      f.Set(reflect.MakeMap(ft))
    case reflect.Slice:
      if param[strings.ToLower(name)] != nil {
         BinderSlice(strings.ToLower(name),&f,ft,param[strings.ToLower(name)])
      }
    case reflect.Chan:
      f.Set(reflect.MakeChan(ft, 0))
    case reflect.Struct:
      mapp := param[strings.ToLower(name)]
      var newParamap map[string]interface{}
      if mapp != nil {
         switch reflect.TypeOf(mapp).Kind(){
            case reflect.Map:
              newParamap = mapp.(map[string]interface{})
              initializeStruct(ft, f,newParamap)
         }
      }
    case reflect.Ptr:
      fv := reflect.New(ft.Elem())
      mapp := param[strings.ToLower(name)]
      var newParamap map[string]interface{}
      if mapp != nil {
         switch reflect.TypeOf(mapp).Kind(){
            case reflect.Map:
              newParamap = mapp.(map[string]interface{})
         }
      }
      initializeStruct(ft.Elem(), fv.Elem(),newParamap)
      f.Set(fv)
    case reflect.Int,reflect.Int8,reflect.Int16,reflect.Int32,reflect.Int64:
      if param[strings.ToLower(name)] != nil {
        BinderInt(&f,param[strings.ToLower(name)])
      }
    case reflect.Bool:
      if param[strings.ToLower(name)] != nil {
        BinderBool(&f,param[strings.ToLower(name)])
      }
    case reflect.Float32,reflect.Float64:
      if param[strings.ToLower(name)] != nil {
        BinderFloat(&f,param[strings.ToLower(name)])
      }
    case reflect.Uint,reflect.Uint8,reflect.Uint16,reflect.Uint32,reflect.Uint64:
      if param[strings.ToLower(name)] != nil {
        BinderUint(&f,param[strings.ToLower(name)])
      }
    case reflect.String:
      if param[strings.ToLower(name)] != nil {
        BinderString(&f,param[strings.ToLower(name)])
      }
    default:
  }
}

func initializeStruct(t reflect.Type, v reflect.Value,param map[string] interface{}) {
  for i := 0; i < v.NumField(); i++ {
    f := v.Field(i)
    ft := t.Field(i)
    BinderByType(f,ft.Type,param,ft.Name)
  }
}

func BinderSliceElement(valueKind reflect.Kind, val string, structField reflect.Value) {
    switch valueKind {
      case reflect.Int,reflect.Int8,reflect.Int16,reflect.Int32,reflect.Int64:
        BinderInt(&structField,val)
      case reflect.String:
        BinderString(&structField,val)
      case reflect.Uint,reflect.Uint8,reflect.Uint16,reflect.Uint32,reflect.Uint64:
        BinderUint(&structField,val)
      case reflect.Float32,reflect.Float64:
        BinderFloat(&structField,val)
      case reflect.Bool:
        BinderBool(&structField,val)
    }
}

func BinderSlice(name string ,field *reflect.Value ,filedtype reflect.Type, paramValue interface{}){
   switch paramValue.(type) {
      case []string :
        numElems := len(paramValue.([]string))
        element := paramValue.([]string)
        slice := reflect.MakeSlice(filedtype,numElems,numElems)
        sliceOf := filedtype.Elem().Kind()
        for i := 0; i < numElems; i++ {
           BinderSliceElement(sliceOf,element[i],slice.Index(i))
        }
        field.Set(slice)
      case string :
        slice := reflect.MakeSlice(filedtype, 1, 1)
        element := paramValue.([]string)
        sliceOf := filedtype.Elem().Kind()
        BinderSliceElement(sliceOf,element[0],slice.Index(0))
        field.Set(slice)
      default:
   }

}

func BinderInt(field *reflect.Value , paramValue interface{}){
  intValue,err := GetParamValueInt(paramValue)
  intVal := int64(intValue)
  if err== nil && !field.OverflowInt(intVal) {
     field.SetInt(intVal)
  }
}

func BinderBool(field *reflect.Value , paramValue interface{}){
   if strings.ToLower(paramValue.(string)) == "true" {
      field.SetBool(true)
   }else {
      field.SetBool(false)
   }
}

func BinderFloat(field *reflect.Value , paramValue interface{}){
  intValue,err := GetParamValueFloat(paramValue)
  intVal := float64(intValue)
  if err== nil && !field.OverflowFloat(intVal) {
     field.SetFloat(intVal)
  }
}

func BinderUint(field *reflect.Value , paramValue interface{}){
  intValue,err := GetParamValueUint(paramValue)
  intVal := uint64(intValue)
  if err== nil && !field.OverflowUint(intVal) {
     field.SetUint(intVal)
  }
}

func BinderString(field *reflect.Value , paramValue interface{}){
   field.SetString(paramValue.(string))
}

func Binder(field *reflect.Value , paramValue interface{}) {
  if !field.CanSet() {
      switch field.Kind() {
        case reflect.Int:
          BinderInt(field,paramValue)
        case reflect.Int8:
          BinderInt(field,paramValue)
        case reflect.Int16:
          BinderInt(field,paramValue)
        case reflect.Int64:
          BinderInt(field,paramValue)
        default:
          log.Debug("Default Binder")
      }
  }
}

func BinderParameter(appcon *appcontext.AppContext){
   v := reflect.New(appcon.FormType)
   initializeStruct(appcon.FormType, v.Elem(),appcon.Parameter)
   appcon.Form = v.Interface()
}<commit_msg>commit<commit_after>package binder

import (
  "github.com/wpxiong/beargo/log"
  "github.com/wpxiong/beargo/appcontext"
  "reflect"
  "strconv"
  "strings"
  "errors"
)

func init() {
  log.InitLog()
}

var floatType = reflect.TypeOf(float64(0))

func getFloat(unk interface{}) (float64, error) {
    v := reflect.ValueOf(unk)
    v = reflect.Indirect(v)
    if !v.Type().ConvertibleTo(floatType) {
        return 0, errors.New("")
    }
    fv := v.Convert(floatType)
    return fv.Float(), nil
}

func GetParamValueFloat(paramValue interface{}) (float64,error) {
   switch paramValue.(type) {
     case float32:
       return getFloat(paramValue)
     case float64:
       return getFloat(paramValue)
     case string:
       val,err:= strconv.ParseFloat(paramValue.(string),64)
       if err == nil {
          return val,nil
       }else {
          return 0,err
       }
     default:
        return 0,errors.New("")
  }
}

func GetParamValueInt(paramValue interface{}) (int64,error) {
  switch paramValue.(type) {
     case int:
       return int64(paramValue.(int)),nil
     case int8:
       return int64(paramValue.(int8)),nil
     case int16:
       return int64(paramValue.(int16)),nil
     case int32:
       return int64(paramValue.(int32)),nil
     case int64:
       return paramValue.(int64),nil
     case float32,float64:
       return int64(paramValue.(float64)),nil
     case string:
       val,err:= strconv.Atoi(paramValue.(string))
       if err == nil {
          return int64(val),nil
       }else {
          return 0,err
       }
     default:
        return 0,errors.New("")
  }
}

func GetParamValueUint(paramValue interface{}) (uint64,error) {
  switch paramValue.(type) {
     case uint,uint8,uint16,uint32,uint64:
       return paramValue.(uint64),nil
     case float32,float64:
       return uint64(paramValue.(float64)),nil
     case string:
       val,err:= strconv.Atoi(paramValue.(string))
       if err == nil {
          return uint64(val),nil
       }else {
          return 0,err
       }
     default:
        return 0,errors.New("")
  }
}

func BinderByType(f reflect.Value, ft reflect.Type, param map[string] interface{},name string){
  switch ft.Kind() {
    case reflect.Map:
      f.Set(reflect.MakeMap(ft))
    case reflect.Slice:
      if param[strings.ToLower(name)] != nil {
         BinderSlice(strings.ToLower(name),&f,ft,param[strings.ToLower(name)])
      }
    case reflect.Chan:
      f.Set(reflect.MakeChan(ft, 0))
    case reflect.Struct:
      mapp := param[strings.ToLower(name)]
      var newParamap map[string]interface{}
      if mapp != nil {
         switch reflect.TypeOf(mapp).Kind(){
            case reflect.Map:
              newParamap = mapp.(map[string]interface{})
              initializeStruct(ft, f,newParamap)
         }
      }
    case reflect.Ptr:
      fv := reflect.New(ft.Elem())
      mapp := param[strings.ToLower(name)]
      var newParamap map[string]interface{}
      if mapp != nil {
         switch reflect.TypeOf(mapp).Kind(){
            case reflect.Map:
              newParamap = mapp.(map[string]interface{})
         }
      }
      initializeStruct(ft.Elem(), fv.Elem(),newParamap)
      f.Set(fv)
    case reflect.Int,reflect.Int8,reflect.Int16,reflect.Int32,reflect.Int64:
      if param[strings.ToLower(name)] != nil {
        BinderInt(&f,param[strings.ToLower(name)])
      }
    case reflect.Bool:
      if param[strings.ToLower(name)] != nil {
        BinderBool(&f,param[strings.ToLower(name)])
      }
    case reflect.Float32,reflect.Float64:
      if param[strings.ToLower(name)] != nil {
        BinderFloat(&f,param[strings.ToLower(name)])
      }
    case reflect.Uint,reflect.Uint8,reflect.Uint16,reflect.Uint32,reflect.Uint64:
      if param[strings.ToLower(name)] != nil {
        BinderUint(&f,param[strings.ToLower(name)])
      }
    case reflect.String:
      if param[strings.ToLower(name)] != nil {
        BinderString(&f,param[strings.ToLower(name)])
      }
    default:
  }
}

func initializeStruct(t reflect.Type, v reflect.Value,param map[string] interface{}) {
  for i := 0; i < v.NumField(); i++ {
    f := v.Field(i)
    ft := t.Field(i)
    BinderByType(f,ft.Type,param,ft.Name)
  }
}

func BinderSliceElement(valueKind reflect.Kind, val string, structField reflect.Value) {
    switch valueKind {
      case reflect.Int,reflect.Int8,reflect.Int16,reflect.Int32,reflect.Int64:
        BinderInt(&structField,val)
      case reflect.String:
        BinderString(&structField,val)
      case reflect.Uint,reflect.Uint8,reflect.Uint16,reflect.Uint32,reflect.Uint64:
        BinderUint(&structField,val)
      case reflect.Float32,reflect.Float64:
        BinderFloat(&structField,val)
      case reflect.Bool:
        BinderBool(&structField,val)
    }
}

func BinderSlice(name string ,field *reflect.Value ,filedtype reflect.Type, paramValue interface{}){
   switch paramValue.(type) {
      case []string :
        numElems := len(paramValue.([]string))
        element := paramValue.([]string)
        slice := reflect.MakeSlice(filedtype,numElems,numElems)
        sliceOf := filedtype.Elem().Kind()
        for i := 0; i < numElems; i++ {
           BinderSliceElement(sliceOf,element[i],slice.Index(i))
        }
        field.Set(slice)
      case string :
        slice := reflect.MakeSlice(filedtype, 1, 1)
        element := paramValue.([]string)
        sliceOf := filedtype.Elem().Kind()
        BinderSliceElement(sliceOf,element[0],slice.Index(0))
        field.Set(slice)
      default:
   }

}

func BinderInt(field *reflect.Value , paramValue interface{}){
  intValue,err := GetParamValueInt(paramValue)
  intVal := int64(intValue)
  if err== nil && !field.OverflowInt(intVal) {
     field.SetInt(intVal)
  }
}

func BinderBool(field *reflect.Value , paramValue interface{}){
   if strings.ToLower(paramValue.(string)) == "true" {
      field.SetBool(true)
   }else {
      field.SetBool(false)
   }
}

func BinderFloat(field *reflect.Value , paramValue interface{}){
  intValue,err := GetParamValueFloat(paramValue)
  intVal := float64(intValue)
  if err== nil && !field.OverflowFloat(intVal) {
     field.SetFloat(intVal)
  }
}

func BinderUint(field *reflect.Value , paramValue interface{}){
  intValue,err := GetParamValueUint(paramValue)
  intVal := uint64(intValue)
  if err== nil && !field.OverflowUint(intVal) {
     field.SetUint(intVal)
  }
}

func BinderString(field *reflect.Value , paramValue interface{}){
   field.SetString(paramValue.(string))
}

func BinderParameter(appcon *appcontext.AppContext){
   v := reflect.New(appcon.FormType)
   initializeStruct(appcon.FormType, v.Elem(),appcon.Parameter)
   appcon.Form = v.Interface()
}<|endoftext|>


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

// Package boomer provides commands to run load tests and display results.
package boomer

import (
	"crypto/tls"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"strings"
	"sync"
	"time"

	"github.com/rakyll/pb"
)

type result struct {
	err           error
	statusCode    int
	duration      time.Duration
	contentLength int64
}

type Boomer struct {
	// Request is the request to be made.
	Request *http.Request

	RequestBody string

	// N is the total number of requests to make.
	N int

	// C is the concurrency level, the number of concurrent workers to run.
	C int

	// Timeout in seconds.
	Timeout int

	// Qps is the rate limit.
	Qps int

	// AllowInsecure is an option to allow insecure TLS/SSL certificates.
	AllowInsecure bool

	// DisableCompression is an option to disable compression in response
	DisableCompression bool

	// DisableKeepAlives is an option to prevents re-use of TCP connections between different HTTP requests
	DisableKeepAlives bool

	// Output represents the output type. If "csv" is provided, the
	// output will be dumped as a csv stream.
	Output string

	// ProxyAddr is the address of HTTP proxy server in the format on "host:port".
	// Optional.
	ProxyAddr *url.URL

	// ReadAll determines whether the body of the response needs
	// to be fully consumed.
	ReadAll bool

	bar     *pb.ProgressBar
	results chan *result
}

func (b *Boomer) startProgress() {
	if b.Output != "" {
		return
	}
	b.bar = pb.New(b.N)
	b.bar.Format("Bom !")
	b.bar.Start()
}

func (b *Boomer) finalizeProgress() {
	if b.Output != "" {
		return
	}
	b.bar.Finish()
}

func (b *Boomer) incProgress() {
	if b.Output != "" {
		return
	}
	b.bar.Increment()
}

// Run makes all the requests, prints the summary. It blocks until
// all work is done.
func (b *Boomer) Run() {
	b.results = make(chan *result, b.N)
	b.startProgress()

	start := time.Now()
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)

	go func() {
		<-c
		// TODO(jbd): Progress bar should not be finalized.
		newReport(b.N, b.results, b.Output, time.Now().Sub(start)).finalize()
		os.Exit(1)
	}()

	b.runWorkers()
	b.finalizeProgress()
	newReport(b.N, b.results, b.Output, time.Now().Sub(start)).finalize()
	close(b.results)
}

func (b *Boomer) runWorker(wg *sync.WaitGroup, ch chan *http.Request) {
	tr := &http.Transport{
		TLSClientConfig: &tls.Config{
			InsecureSkipVerify: b.AllowInsecure,
		},
		DisableCompression: b.DisableCompression,
		DisableKeepAlives:  b.DisableKeepAlives,
		// TODO(jbd): Add dial timeout.
		TLSHandshakeTimeout: time.Duration(b.Timeout) * time.Millisecond,
		Proxy:               http.ProxyURL(b.ProxyAddr),
	}
	client := &http.Client{Transport: tr}
	for req := range ch {
		s := time.Now()

		var code int
		var size int64

		resp, err := client.Do(req)
		if err == nil {
			size = resp.ContentLength
			code = resp.StatusCode
			if b.ReadAll {
				_, err = io.Copy(ioutil.Discard, resp.Body)
			}
			resp.Body.Close()
		}

		wg.Done()
		b.incProgress()
		b.results <- &result{
			statusCode:    code,
			duration:      time.Now().Sub(s),
			err:           err,
			contentLength: size,
		}
	}
}

func (b *Boomer) runWorkers() {
	var wg sync.WaitGroup
	wg.Add(b.N)

	var throttle <-chan time.Time
	if b.Qps > 0 {
		throttle = time.Tick(time.Duration(1e6/(b.Qps)) * time.Microsecond)
	}

	jobsch := make(chan *http.Request, b.N)
	for i := 0; i < b.C; i++ {
		go b.runWorker(&wg, jobsch)
	}

	for i := 0; i < b.N; i++ {
		if b.Qps > 0 {
			<-throttle
		}
		jobsch <- cloneRequest(b.Request, b.RequestBody)
	}
	close(jobsch)
	wg.Wait()
}

// cloneRequest returns a clone of the provided *http.Request.
// The clone is a shallow copy of the struct and its Header map.
func cloneRequest(r *http.Request, body string) *http.Request {
	// shallow copy of the struct
	r2 := new(http.Request)
	*r2 = *r
	// deep copy of the Header
	r2.Header = make(http.Header, len(r.Header))
	for k, s := range r.Header {
		r2.Header[k] = append([]string(nil), s...)
	}
	r2.Body = ioutil.NopCloser(strings.NewReader(body))
	return r2
}
<commit_msg>Fewer wait group calls<commit_after>// Copyright 2014 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package boomer provides commands to run load tests and display results.
package boomer

import (
	"crypto/tls"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"strings"
	"sync"
	"time"

	"github.com/rakyll/pb"
)

type result struct {
	err           error
	statusCode    int
	duration      time.Duration
	contentLength int64
}

type Boomer struct {
	// Request is the request to be made.
	Request *http.Request

	RequestBody string

	// N is the total number of requests to make.
	N int

	// C is the concurrency level, the number of concurrent workers to run.
	C int

	// Timeout in seconds.
	Timeout int

	// Qps is the rate limit.
	Qps int

	// AllowInsecure is an option to allow insecure TLS/SSL certificates.
	AllowInsecure bool

	// DisableCompression is an option to disable compression in response
	DisableCompression bool

	// DisableKeepAlives is an option to prevents re-use of TCP connections between different HTTP requests
	DisableKeepAlives bool

	// Output represents the output type. If "csv" is provided, the
	// output will be dumped as a csv stream.
	Output string

	// ProxyAddr is the address of HTTP proxy server in the format on "host:port".
	// Optional.
	ProxyAddr *url.URL

	// ReadAll determines whether the body of the response needs
	// to be fully consumed.
	ReadAll bool

	bar     *pb.ProgressBar
	results chan *result
}

func (b *Boomer) startProgress() {
	if b.Output != "" {
		return
	}
	b.bar = pb.New(b.N)
	b.bar.Format("Bom !")
	b.bar.Start()
}

func (b *Boomer) finalizeProgress() {
	if b.Output != "" {
		return
	}
	b.bar.Finish()
}

func (b *Boomer) incProgress() {
	if b.Output != "" {
		return
	}
	b.bar.Increment()
}

// Run makes all the requests, prints the summary. It blocks until
// all work is done.
func (b *Boomer) Run() {
	b.results = make(chan *result, b.N)
	b.startProgress()

	start := time.Now()
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)

	go func() {
		<-c
		// TODO(jbd): Progress bar should not be finalized.
		newReport(b.N, b.results, b.Output, time.Now().Sub(start)).finalize()
		os.Exit(1)
	}()

	b.runWorkers()
	b.finalizeProgress()
	newReport(b.N, b.results, b.Output, time.Now().Sub(start)).finalize()
	close(b.results)
}

func (b *Boomer) runWorker(wg *sync.WaitGroup, ch chan *http.Request) {
	tr := &http.Transport{
		TLSClientConfig: &tls.Config{
			InsecureSkipVerify: b.AllowInsecure,
		},
		DisableCompression: b.DisableCompression,
		DisableKeepAlives:  b.DisableKeepAlives,
		// TODO(jbd): Add dial timeout.
		TLSHandshakeTimeout: time.Duration(b.Timeout) * time.Millisecond,
		Proxy:               http.ProxyURL(b.ProxyAddr),
	}
	client := &http.Client{Transport: tr}
	for req := range ch {
		s := time.Now()

		var code int
		var size int64

		resp, err := client.Do(req)
		if err == nil {
			size = resp.ContentLength
			code = resp.StatusCode
			if b.ReadAll {
				_, err = io.Copy(ioutil.Discard, resp.Body)
			}
			resp.Body.Close()
		}

		b.incProgress()
		b.results <- &result{
			statusCode:    code,
			duration:      time.Now().Sub(s),
			err:           err,
			contentLength: size,
		}
	}
	wg.Done()
}

func (b *Boomer) runWorkers() {
	var wg sync.WaitGroup
	wg.Add(b.C)

	var throttle <-chan time.Time
	if b.Qps > 0 {
		throttle = time.Tick(time.Duration(1e6/(b.Qps)) * time.Microsecond)
	}

	jobsch := make(chan *http.Request, b.N)
	for i := 0; i < b.C; i++ {
		go b.runWorker(&wg, jobsch)
	}

	for i := 0; i < b.N; i++ {
		if b.Qps > 0 {
			<-throttle
		}
		jobsch <- cloneRequest(b.Request, b.RequestBody)
	}
	close(jobsch)
	wg.Wait()
}

// cloneRequest returns a clone of the provided *http.Request.
// The clone is a shallow copy of the struct and its Header map.
func cloneRequest(r *http.Request, body string) *http.Request {
	// shallow copy of the struct
	r2 := new(http.Request)
	*r2 = *r
	// deep copy of the Header
	r2.Header = make(http.Header, len(r.Header))
	for k, s := range r.Header {
		r2.Header[k] = append([]string(nil), s...)
	}
	r2.Body = ioutil.NopCloser(strings.NewReader(body))
	return r2
}
<|endoftext|>


<commit_before>package bootstrap

import (
	"fmt"
	"net/url"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/buildkite/agent/bootstrap/shell"
	"github.com/buildkite/shellwords"
)

func gitClone(sh *shell.Shell, gitCloneFlags, repository, dir string) error {
	individualCloneFlags, err := shellwords.Split(gitCloneFlags)
	if err != nil {
		return err
	}

	commandArgs := []string{"clone"}
	commandArgs = append(commandArgs, individualCloneFlags...)
	commandArgs = append(commandArgs, "--", repository, ".")

	if err = sh.Run("git", commandArgs...); err != nil {
		return err
	}

	return nil
}

func gitClean(sh *shell.Shell, gitCleanFlags string) error {
	individualCleanFlags, err := shellwords.Split(gitCleanFlags)
	if err != nil {
		return err
	}

	commandArgs := []string{"clean"}
	commandArgs = append(commandArgs, individualCleanFlags...)

	if err = sh.Run("git", commandArgs...); err != nil {
		return err
	}

	return nil
}

func gitCleanSubmodules(sh *shell.Shell, gitCleanFlags string) error {
	individualCleanFlags, err := shellwords.Split(gitCleanFlags)
	if err != nil {
		return err
	}

	commandArgs := append([]string{"submodule", "foreach", "--recursive", "git", "clean"}, individualCleanFlags...)

	if err = sh.Run("git", commandArgs...); err != nil {
		return err
	}

	return nil
}

func gitFetch(sh *shell.Shell, gitFetchFlags, repository string, refSpec ...string) error {
	individualFetchFlags, err := shellwords.Split(gitFetchFlags)
	if err != nil {
		return err
	}

	commandArgs := []string{"fetch"}
	commandArgs = append(commandArgs, individualFetchFlags...)
	commandArgs = append(commandArgs, repository)

	for _, r := range refSpec {
		individualRefSpecs, err := shellwords.Split(r)
		if err != nil {
			return err
		}
		commandArgs = append(commandArgs, individualRefSpecs...)
	}

	if err = sh.Run("git", commandArgs...); err != nil {
		return err
	}

	return nil
}

func gitEnumerateSubmoduleURLs(sh *shell.Shell) ([]string, error) {
	urls := []string{}

	// The output of this command looks like:
	//submodule.bitbucket-git-docker-example.url\ngit@bitbucket.org:lox24/docker-example.git\0
	//submodule.bitbucket-https-docker-example.url\nhttps://lox24@bitbucket.org/lox24/docker-example.git\0
	//submodule.github-git-docker-example.url\ngit@github.com:buildkite/docker-example.git\0
	//submodule.github-https-docker-example.url\nhttps://github.com/buildkite/docker-example.git\0
	output, err := sh.RunAndCapture(
		"git", "config", "--file", ".gitmodules", "--null", "--get-regexp", "url")
	if err != nil {
		return nil, err
	}

	// splits lines on null-bytes to gracefully handle line endings and repositories with newlines
	lines := strings.Split(strings.TrimRight(output, "\x00"), "\x00")

	// process each line
	for _, line := range lines {
		tokens := strings.SplitN(line, "\n", 2)
		if len(tokens) != 2 {
			return nil, fmt.Errorf("Failed to parse .gitmodule line %q", line)
		}
		urls = append(urls, tokens[1])
	}

	return urls, nil
}

func gitRevParseInWorkingDirectory(sh *shell.Shell, workingDirectory string, extraRevParseArgs ...string) (string, error) {
	gitDirectory := filepath.Join(workingDirectory, ".git")

	revParseArgs := []string{"--git-dir", gitDirectory, "--work-tree", workingDirectory, "rev-parse"}
	revParseArgs = append(revParseArgs, extraRevParseArgs...)

	return sh.RunAndCapture("git", revParseArgs...)
}

var (
	hasSchemePattern  = regexp.MustCompile("^[^:]+://")
	scpLikeURLPattern = regexp.MustCompile("^([^@]+@)?([^:]{2,}):/?(.+)$")
)

// parseGittableURL parses and converts a git repository url into a url.URL
func parseGittableURL(ref string) (*url.URL, error) {
	if !hasSchemePattern.MatchString(ref) {
		if scpLikeURLPattern.MatchString(ref) {
			matched := scpLikeURLPattern.FindStringSubmatch(ref)
			user := matched[1]
			host := matched[2]
			path := matched[3]
			ref = fmt.Sprintf("ssh://%s%s/%s", user, host, path)
		} else {
			normalizedRef := strings.Replace(ref, "\\", "/", -1)
			ref = fmt.Sprintf("file:///%s", strings.TrimPrefix(normalizedRef, "/"))
		}
	}
	return url.Parse(ref)
}

// Clean up the SSH host and remove any key identifiers. See:
// git@github.com-custom-identifier:foo/bar.git
// https://buildkite.com/docs/agent/ssh-keys#creating-multiple-ssh-keys
var gitHostAliasRegexp = regexp.MustCompile(`-[a-z0-9\-]+$`)

func stripAliasesFromGitHost(host string) string {
	return gitHostAliasRegexp.ReplaceAllString(host, "")
}
<commit_msg>Improve formatting of output example<commit_after>package bootstrap

import (
	"fmt"
	"net/url"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/buildkite/agent/bootstrap/shell"
	"github.com/buildkite/shellwords"
)

func gitClone(sh *shell.Shell, gitCloneFlags, repository, dir string) error {
	individualCloneFlags, err := shellwords.Split(gitCloneFlags)
	if err != nil {
		return err
	}

	commandArgs := []string{"clone"}
	commandArgs = append(commandArgs, individualCloneFlags...)
	commandArgs = append(commandArgs, "--", repository, ".")

	if err = sh.Run("git", commandArgs...); err != nil {
		return err
	}

	return nil
}

func gitClean(sh *shell.Shell, gitCleanFlags string) error {
	individualCleanFlags, err := shellwords.Split(gitCleanFlags)
	if err != nil {
		return err
	}

	commandArgs := []string{"clean"}
	commandArgs = append(commandArgs, individualCleanFlags...)

	if err = sh.Run("git", commandArgs...); err != nil {
		return err
	}

	return nil
}

func gitCleanSubmodules(sh *shell.Shell, gitCleanFlags string) error {
	individualCleanFlags, err := shellwords.Split(gitCleanFlags)
	if err != nil {
		return err
	}

	commandArgs := append([]string{"submodule", "foreach", "--recursive", "git", "clean"}, individualCleanFlags...)

	if err = sh.Run("git", commandArgs...); err != nil {
		return err
	}

	return nil
}

func gitFetch(sh *shell.Shell, gitFetchFlags, repository string, refSpec ...string) error {
	individualFetchFlags, err := shellwords.Split(gitFetchFlags)
	if err != nil {
		return err
	}

	commandArgs := []string{"fetch"}
	commandArgs = append(commandArgs, individualFetchFlags...)
	commandArgs = append(commandArgs, repository)

	for _, r := range refSpec {
		individualRefSpecs, err := shellwords.Split(r)
		if err != nil {
			return err
		}
		commandArgs = append(commandArgs, individualRefSpecs...)
	}

	if err = sh.Run("git", commandArgs...); err != nil {
		return err
	}

	return nil
}

func gitEnumerateSubmoduleURLs(sh *shell.Shell) ([]string, error) {
	urls := []string{}

	// The output of this command looks like:
	// submodule.bitbucket-git-docker-example.url\ngit@bitbucket.org:lox24/docker-example.git\0
	// submodule.bitbucket-https-docker-example.url\nhttps://lox24@bitbucket.org/lox24/docker-example.git\0
	// submodule.github-git-docker-example.url\ngit@github.com:buildkite/docker-example.git\0
	// submodule.github-https-docker-example.url\nhttps://github.com/buildkite/docker-example.git\0
	output, err := sh.RunAndCapture(
		"git", "config", "--file", ".gitmodules", "--null", "--get-regexp", "url")
	if err != nil {
		return nil, err
	}

	// splits lines on null-bytes to gracefully handle line endings and repositories with newlines
	lines := strings.Split(strings.TrimRight(output, "\x00"), "\x00")

	// process each line
	for _, line := range lines {
		tokens := strings.SplitN(line, "\n", 2)
		if len(tokens) != 2 {
			return nil, fmt.Errorf("Failed to parse .gitmodule line %q", line)
		}
		urls = append(urls, tokens[1])
	}

	return urls, nil
}

func gitRevParseInWorkingDirectory(sh *shell.Shell, workingDirectory string, extraRevParseArgs ...string) (string, error) {
	gitDirectory := filepath.Join(workingDirectory, ".git")

	revParseArgs := []string{"--git-dir", gitDirectory, "--work-tree", workingDirectory, "rev-parse"}
	revParseArgs = append(revParseArgs, extraRevParseArgs...)

	return sh.RunAndCapture("git", revParseArgs...)
}

var (
	hasSchemePattern  = regexp.MustCompile("^[^:]+://")
	scpLikeURLPattern = regexp.MustCompile("^([^@]+@)?([^:]{2,}):/?(.+)$")
)

// parseGittableURL parses and converts a git repository url into a url.URL
func parseGittableURL(ref string) (*url.URL, error) {
	if !hasSchemePattern.MatchString(ref) {
		if scpLikeURLPattern.MatchString(ref) {
			matched := scpLikeURLPattern.FindStringSubmatch(ref)
			user := matched[1]
			host := matched[2]
			path := matched[3]
			ref = fmt.Sprintf("ssh://%s%s/%s", user, host, path)
		} else {
			normalizedRef := strings.Replace(ref, "\\", "/", -1)
			ref = fmt.Sprintf("file:///%s", strings.TrimPrefix(normalizedRef, "/"))
		}
	}
	return url.Parse(ref)
}

// Clean up the SSH host and remove any key identifiers. See:
// git@github.com-custom-identifier:foo/bar.git
// https://buildkite.com/docs/agent/ssh-keys#creating-multiple-ssh-keys
var gitHostAliasRegexp = regexp.MustCompile(`-[a-z0-9\-]+$`)

func stripAliasesFromGitHost(host string) string {
	return gitHostAliasRegexp.ReplaceAllString(host, "")
}
<|endoftext|>