<commit_before>package drivers

import (
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	"github.com/pkg/errors"

	"github.com/lxc/lxd/lxd/migration"
	"github.com/lxc/lxd/lxd/operations"
	"github.com/lxc/lxd/lxd/rsync"
	"github.com/lxc/lxd/lxd/state"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
	"github.com/lxc/lxd/shared/ioprogress"
	"github.com/lxc/lxd/shared/logger"
)

type common struct {
	name                 string
	config               map[string]string
	getVolID             func(volType VolumeType, volName string) (int64, error)
	getCommonVolumeRules func(vol Volume) map[string]func(string) error
	state                *state.State
	logger               logger.Logger
}

func (d *common) init(state *state.State, name string, config map[string]string, logger logger.Logger, volIDFunc func(volType VolumeType, volName string) (int64, error), commonVolRulesFunc func(vol Volume) map[string]func(string) error) {
	d.name = name
	d.config = config
	d.getVolID = volIDFunc
	d.getCommonVolumeRules = commonVolRulesFunc
	d.state = state
	d.logger = logger
}

func (d *common) load() error {
	return nil
}

// validateVolume validates a volume config against common rules and optional driver specific rules.
// This functions has a removeUnknownKeys option that if set to true will remove any unknown fields
// (excluding those starting with "user.") which can be used when translating a volume config to a
// different storage driver that has different options.
func (d *common) validateVolume(vol Volume, driverRules map[string]func(value string) error, removeUnknownKeys bool) error {
	checkedFields := map[string]struct{}{}

	// Get rules common for all drivers.
	rules := d.getCommonVolumeRules(vol)

	// Merge driver specific rules into common rules.
	for field, validator := range driverRules {
		rules[field] = validator
	}

	// Run the validator against each field.
	for k, validator := range rules {
		checkedFields[k] = struct{}{} //Mark field as checked.
		err := validator(vol.config[k])
		if err != nil {
			return errors.Wrapf(err, "Invalid value for volume option %s", k)
		}
	}

	// Look for any unchecked fields, as these are unknown fields and validation should fail.
	for k := range vol.config {
		_, checked := checkedFields[k]
		if checked {
			continue
		}

		// User keys are not validated.
		if strings.HasPrefix(k, "user.") {
			continue
		}

		if removeUnknownKeys {
			delete(vol.config, k)
		} else {
			return fmt.Errorf("Invalid volume option: %s", k)
		}
	}

	// If volume type is not custom, don't allow "size" property.
	if vol.volType != VolumeTypeCustom && vol.config["size"] != "" {
		return fmt.Errorf("Volume 'size' property is only valid for custom volume types")
	}

	return nil
}

// MigrationType returns the type of transfer methods to be used when doing migrations between pools
// in preference order.
func (d *common) MigrationTypes(contentType ContentType, refresh bool) []migration.Type {
	if contentType != ContentTypeFS {
		return nil
	}

	return []migration.Type{
		{
			FSType:   migration.MigrationFSType_RSYNC,
			Features: []string{"xattrs", "delete", "compress", "bidirectional"},
		},
	}
}

// Name returns the pool name.
func (d *common) Name() string {
	return d.name
}

// Logger returns the current logger.
func (d *common) Logger() logger.Logger {
	return d.logger
}

// Config returns the storage pool config (as a copy, so not modifiable).
func (d *common) Config() map[string]string {
	confCopy := make(map[string]string, len(d.config))
	for k, v := range d.config {
		confCopy[k] = v
	}

	return confCopy
}

// vfsGetResources is a generic GetResources implementation for VFS-only drivers.
func (d *common) vfsGetResources() (*api.ResourcesStoragePool, error) {
	// Get the VFS information
	st, err := shared.Statvfs(GetPoolMountPath(d.name))
	if err != nil {
		return nil, err
	}

	// Fill in the struct
	res := api.ResourcesStoragePool{}
	res.Space.Total = st.Blocks * uint64(st.Bsize)
	res.Space.Used = (st.Blocks - st.Bfree) * uint64(st.Bsize)

	// Some filesystems don't report inodes since they allocate them
	// dynamically e.g. btrfs.
	if st.Files > 0 {
		res.Inodes.Total = st.Files
		res.Inodes.Used = st.Files - st.Ffree
	}

	return &res, nil
}

// vfsRenameVolume is a generic RenameVolume implementation for VFS-only drivers.
func (d *common) vfsRenameVolume(vol Volume, newVolName string, op *operations.Operation) error {
	// Rename the volume itself.
	srcVolumePath := GetVolumeMountPath(d.name, vol.volType, vol.name)
	dstVolumePath := GetVolumeMountPath(d.name, vol.volType, newVolName)

	err := os.Rename(srcVolumePath, dstVolumePath)
	if err != nil {
		return errors.Wrapf(err, "Failed to rename '%s' to '%s'", srcVolumePath, dstVolumePath)
	}

	revertRename := true
	defer func() {
		if !revertRename {
			return
		}

		os.Rename(dstVolumePath, srcVolumePath)
	}()

	// And if present, the snapshots too.
	srcSnapshotDir := GetVolumeSnapshotDir(d.name, vol.volType, vol.name)
	dstSnapshotDir := GetVolumeSnapshotDir(d.name, vol.volType, newVolName)

	if shared.PathExists(srcSnapshotDir) {
		err = os.Rename(srcSnapshotDir, dstSnapshotDir)
		if err != nil {
			return errors.Wrapf(err, "Failed to rename '%s' to '%s'", srcSnapshotDir, dstSnapshotDir)
		}
	}

	revertRename = false
	return nil
}

// vfsVolumeSnapshots is a generic VolumeSnapshots implementation for VFS-only drivers.
func (d *common) vfsVolumeSnapshots(vol Volume, op *operations.Operation) ([]string, error) {
	snapshotDir := GetVolumeSnapshotDir(d.name, vol.volType, vol.name)
	snapshots := []string{}

	ents, err := ioutil.ReadDir(snapshotDir)
	if err != nil {
		// If the snapshots directory doesn't exist, there are no snapshots.
		if os.IsNotExist(err) {
			return snapshots, nil
		}

		return nil, errors.Wrapf(err, "Failed to list directory '%s'", snapshotDir)
	}

	for _, ent := range ents {
		fileInfo, err := os.Stat(filepath.Join(snapshotDir, ent.Name()))
		if err != nil {
			return nil, err
		}

		if !fileInfo.IsDir() {
			continue
		}

		snapshots = append(snapshots, ent.Name())
	}

	return snapshots, nil
}

// vfsRenameVolumeSnapshot is a generic RenameVolumeSnapshot implementation for VFS-only drivers.
func (d *common) vfsRenameVolumeSnapshot(snapVol Volume, newSnapshotName string, op *operations.Operation) error {
	parentName, _, _ := shared.InstanceGetParentAndSnapshotName(snapVol.name)
	oldPath := snapVol.MountPath()
	newPath := GetVolumeMountPath(d.name, snapVol.volType, GetSnapshotVolumeName(parentName, newSnapshotName))

	err := os.Rename(oldPath, newPath)
	if err != nil {
		return errors.Wrapf(err, "Failed to rename '%s' to '%s'", oldPath, newPath)
	}

	return nil
}

// vfsMigrateVolume is a generic MigrateVolume implementation for VFS-only drivers.
func (d *common) vfsMigrateVolume(vol Volume, conn io.ReadWriteCloser, volSrcArgs *migration.VolumeSourceArgs, op *operations.Operation) error {
	bwlimit := d.config["rsync.bwlimit"]

	for _, snapName := range volSrcArgs.Snapshots {
		snapshot, err := vol.NewSnapshot(snapName)
		if err != nil {
			return err
		}

		// Send snapshot to recipient (ensure local snapshot volume is mounted if needed).
		err = snapshot.MountTask(func(mountPath string, op *operations.Operation) error {
			var wrapper *ioprogress.ProgressTracker
			if volSrcArgs.TrackProgress {
				wrapper = migration.ProgressTracker(op, "fs_progress", snapshot.name)
			}

			path := shared.AddSlash(mountPath)
			return rsync.Send(snapshot.name, path, conn, wrapper, volSrcArgs.MigrationType.Features, bwlimit, d.state.OS.ExecPath)
		}, op)
		if err != nil {
			return err
		}
	}

	// Send volume to recipient (ensure local volume is mounted if needed).
	return vol.MountTask(func(mountPath string, op *operations.Operation) error {
		var wrapper *ioprogress.ProgressTracker
		if volSrcArgs.TrackProgress {
			wrapper = migration.ProgressTracker(op, "fs_progress", vol.name)
		}

		path := shared.AddSlash(mountPath)
		return rsync.Send(vol.name, path, conn, wrapper, volSrcArgs.MigrationType.Features, bwlimit, d.state.OS.ExecPath)
	}, op)
}

// vfsHasVolume is a generic HasVolume implementation for VFS-only drivers.
func (d *common) vfsHasVolume(vol Volume) bool {
	if shared.PathExists(vol.MountPath()) {
		return true
	}

	return false
}

// vfsGetVolumeDiskPath is a generic GetVolumeDiskPath implementation for VFS-only drivers.
func (d *common) vfsGetVolumeDiskPath(vol Volume) (string, error) {
	if vol.contentType != ContentTypeBlock {
		return "", fmt.Errorf("No disk paths for filesystems")
	}

	return filepath.Join(vol.MountPath(), "root.img"), nil
}

// vfsBackupVolume is a generic BackupVolume implementation for VFS-only drivers.
func (d *common) vfsBackupVolume(vol Volume, targetPath string, snapshots bool, op *operations.Operation) error {
	bwlimit := d.config["rsync.bwlimit"]

	// Backups only implemented for containers currently.
	if vol.volType != VolumeTypeContainer {
		return ErrNotImplemented
	}
	// Handle snapshots.
	if snapshots {
		snapshotsPath := filepath.Join(targetPath, "snapshots")

		// List the snapshots.
		snapshots, err := vol.Snapshots(op)
		if err != nil {
			return err
		}

		// Create the snapshot path.
		if len(snapshots) > 0 {
			err = os.MkdirAll(snapshotsPath, 0711)
			if err != nil {
				return errors.Wrapf(err, "Failed to create directory '%s'", snapshotsPath)
			}
		}

		for _, snapshot := range snapshots {
			_, snapName, _ := shared.InstanceGetParentAndSnapshotName(snapshot.Name())
			target := filepath.Join(snapshotsPath, snapName)

			// Copy the snapshot.
			err = snapshot.MountTask(func(mountPath string, op *operations.Operation) error {
				_, err := rsync.LocalCopy(mountPath, target, bwlimit, true)
				if err != nil {
					return err
				}

				return nil
			}, op)
			if err != nil {
				return err
			}
		}
	}

	// Copy the parent volume itself.
	target := filepath.Join(targetPath, "container")
	err := vol.MountTask(func(mountPath string, op *operations.Operation) error {
		_, err := rsync.LocalCopy(mountPath, target, bwlimit, true)
		if err != nil {
			return err
		}

		return nil
	}, op)
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>lxd/storage/drivers/driver/common: Adds VM support for migration types<commit_after>package drivers

import (
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"

	"github.com/pkg/errors"

	"github.com/lxc/lxd/lxd/migration"
	"github.com/lxc/lxd/lxd/operations"
	"github.com/lxc/lxd/lxd/rsync"
	"github.com/lxc/lxd/lxd/state"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/api"
	"github.com/lxc/lxd/shared/ioprogress"
	"github.com/lxc/lxd/shared/logger"
)

type common struct {
	name                 string
	config               map[string]string
	getVolID             func(volType VolumeType, volName string) (int64, error)
	getCommonVolumeRules func(vol Volume) map[string]func(string) error
	state                *state.State
	logger               logger.Logger
}

func (d *common) init(state *state.State, name string, config map[string]string, logger logger.Logger, volIDFunc func(volType VolumeType, volName string) (int64, error), commonVolRulesFunc func(vol Volume) map[string]func(string) error) {
	d.name = name
	d.config = config
	d.getVolID = volIDFunc
	d.getCommonVolumeRules = commonVolRulesFunc
	d.state = state
	d.logger = logger
}

func (d *common) load() error {
	return nil
}

// validateVolume validates a volume config against common rules and optional driver specific rules.
// This functions has a removeUnknownKeys option that if set to true will remove any unknown fields
// (excluding those starting with "user.") which can be used when translating a volume config to a
// different storage driver that has different options.
func (d *common) validateVolume(vol Volume, driverRules map[string]func(value string) error, removeUnknownKeys bool) error {
	checkedFields := map[string]struct{}{}

	// Get rules common for all drivers.
	rules := d.getCommonVolumeRules(vol)

	// Merge driver specific rules into common rules.
	for field, validator := range driverRules {
		rules[field] = validator
	}

	// Run the validator against each field.
	for k, validator := range rules {
		checkedFields[k] = struct{}{} //Mark field as checked.
		err := validator(vol.config[k])
		if err != nil {
			return errors.Wrapf(err, "Invalid value for volume option %s", k)
		}
	}

	// Look for any unchecked fields, as these are unknown fields and validation should fail.
	for k := range vol.config {
		_, checked := checkedFields[k]
		if checked {
			continue
		}

		// User keys are not validated.
		if strings.HasPrefix(k, "user.") {
			continue
		}

		if removeUnknownKeys {
			delete(vol.config, k)
		} else {
			return fmt.Errorf("Invalid volume option: %s", k)
		}
	}

	// If volume type is not custom, don't allow "size" property.
	if vol.volType != VolumeTypeCustom && vol.config["size"] != "" {
		return fmt.Errorf("Volume 'size' property is only valid for custom volume types")
	}

	return nil
}

// MigrationType returns the type of transfer methods to be used when doing migrations between pools
// in preference order.
func (d *common) MigrationTypes(contentType ContentType, refresh bool) []migration.Type {
	return []migration.Type{
		{
			FSType:   migration.MigrationFSType_RSYNC,
			Features: []string{"xattrs", "delete", "compress", "bidirectional"},
		},
	}
}

// Name returns the pool name.
func (d *common) Name() string {
	return d.name
}

// Logger returns the current logger.
func (d *common) Logger() logger.Logger {
	return d.logger
}

// Config returns the storage pool config (as a copy, so not modifiable).
func (d *common) Config() map[string]string {
	confCopy := make(map[string]string, len(d.config))
	for k, v := range d.config {
		confCopy[k] = v
	}

	return confCopy
}

// vfsGetResources is a generic GetResources implementation for VFS-only drivers.
func (d *common) vfsGetResources() (*api.ResourcesStoragePool, error) {
	// Get the VFS information
	st, err := shared.Statvfs(GetPoolMountPath(d.name))
	if err != nil {
		return nil, err
	}

	// Fill in the struct
	res := api.ResourcesStoragePool{}
	res.Space.Total = st.Blocks * uint64(st.Bsize)
	res.Space.Used = (st.Blocks - st.Bfree) * uint64(st.Bsize)

	// Some filesystems don't report inodes since they allocate them
	// dynamically e.g. btrfs.
	if st.Files > 0 {
		res.Inodes.Total = st.Files
		res.Inodes.Used = st.Files - st.Ffree
	}

	return &res, nil
}

// vfsRenameVolume is a generic RenameVolume implementation for VFS-only drivers.
func (d *common) vfsRenameVolume(vol Volume, newVolName string, op *operations.Operation) error {
	// Rename the volume itself.
	srcVolumePath := GetVolumeMountPath(d.name, vol.volType, vol.name)
	dstVolumePath := GetVolumeMountPath(d.name, vol.volType, newVolName)

	err := os.Rename(srcVolumePath, dstVolumePath)
	if err != nil {
		return errors.Wrapf(err, "Failed to rename '%s' to '%s'", srcVolumePath, dstVolumePath)
	}

	revertRename := true
	defer func() {
		if !revertRename {
			return
		}

		os.Rename(dstVolumePath, srcVolumePath)
	}()

	// And if present, the snapshots too.
	srcSnapshotDir := GetVolumeSnapshotDir(d.name, vol.volType, vol.name)
	dstSnapshotDir := GetVolumeSnapshotDir(d.name, vol.volType, newVolName)

	if shared.PathExists(srcSnapshotDir) {
		err = os.Rename(srcSnapshotDir, dstSnapshotDir)
		if err != nil {
			return errors.Wrapf(err, "Failed to rename '%s' to '%s'", srcSnapshotDir, dstSnapshotDir)
		}
	}

	revertRename = false
	return nil
}

// vfsVolumeSnapshots is a generic VolumeSnapshots implementation for VFS-only drivers.
func (d *common) vfsVolumeSnapshots(vol Volume, op *operations.Operation) ([]string, error) {
	snapshotDir := GetVolumeSnapshotDir(d.name, vol.volType, vol.name)
	snapshots := []string{}

	ents, err := ioutil.ReadDir(snapshotDir)
	if err != nil {
		// If the snapshots directory doesn't exist, there are no snapshots.
		if os.IsNotExist(err) {
			return snapshots, nil
		}

		return nil, errors.Wrapf(err, "Failed to list directory '%s'", snapshotDir)
	}

	for _, ent := range ents {
		fileInfo, err := os.Stat(filepath.Join(snapshotDir, ent.Name()))
		if err != nil {
			return nil, err
		}

		if !fileInfo.IsDir() {
			continue
		}

		snapshots = append(snapshots, ent.Name())
	}

	return snapshots, nil
}

// vfsRenameVolumeSnapshot is a generic RenameVolumeSnapshot implementation for VFS-only drivers.
func (d *common) vfsRenameVolumeSnapshot(snapVol Volume, newSnapshotName string, op *operations.Operation) error {
	parentName, _, _ := shared.InstanceGetParentAndSnapshotName(snapVol.name)
	oldPath := snapVol.MountPath()
	newPath := GetVolumeMountPath(d.name, snapVol.volType, GetSnapshotVolumeName(parentName, newSnapshotName))

	err := os.Rename(oldPath, newPath)
	if err != nil {
		return errors.Wrapf(err, "Failed to rename '%s' to '%s'", oldPath, newPath)
	}

	return nil
}

// vfsMigrateVolume is a generic MigrateVolume implementation for VFS-only drivers.
func (d *common) vfsMigrateVolume(vol Volume, conn io.ReadWriteCloser, volSrcArgs *migration.VolumeSourceArgs, op *operations.Operation) error {
	bwlimit := d.config["rsync.bwlimit"]

	for _, snapName := range volSrcArgs.Snapshots {
		snapshot, err := vol.NewSnapshot(snapName)
		if err != nil {
			return err
		}

		// Send snapshot to recipient (ensure local snapshot volume is mounted if needed).
		err = snapshot.MountTask(func(mountPath string, op *operations.Operation) error {
			var wrapper *ioprogress.ProgressTracker
			if volSrcArgs.TrackProgress {
				wrapper = migration.ProgressTracker(op, "fs_progress", snapshot.name)
			}

			path := shared.AddSlash(mountPath)
			return rsync.Send(snapshot.name, path, conn, wrapper, volSrcArgs.MigrationType.Features, bwlimit, d.state.OS.ExecPath)
		}, op)
		if err != nil {
			return err
		}
	}

	// Send volume to recipient (ensure local volume is mounted if needed).
	return vol.MountTask(func(mountPath string, op *operations.Operation) error {
		var wrapper *ioprogress.ProgressTracker
		if volSrcArgs.TrackProgress {
			wrapper = migration.ProgressTracker(op, "fs_progress", vol.name)
		}

		path := shared.AddSlash(mountPath)
		return rsync.Send(vol.name, path, conn, wrapper, volSrcArgs.MigrationType.Features, bwlimit, d.state.OS.ExecPath)
	}, op)
}

// vfsHasVolume is a generic HasVolume implementation for VFS-only drivers.
func (d *common) vfsHasVolume(vol Volume) bool {
	if shared.PathExists(vol.MountPath()) {
		return true
	}

	return false
}

// vfsGetVolumeDiskPath is a generic GetVolumeDiskPath implementation for VFS-only drivers.
func (d *common) vfsGetVolumeDiskPath(vol Volume) (string, error) {
	if vol.contentType != ContentTypeBlock {
		return "", fmt.Errorf("No disk paths for filesystems")
	}

	return filepath.Join(vol.MountPath(), "root.img"), nil
}

// vfsBackupVolume is a generic BackupVolume implementation for VFS-only drivers.
func (d *common) vfsBackupVolume(vol Volume, targetPath string, snapshots bool, op *operations.Operation) error {
	bwlimit := d.config["rsync.bwlimit"]

	// Backups only implemented for containers currently.
	if vol.volType != VolumeTypeContainer {
		return ErrNotImplemented
	}
	// Handle snapshots.
	if snapshots {
		snapshotsPath := filepath.Join(targetPath, "snapshots")

		// List the snapshots.
		snapshots, err := vol.Snapshots(op)
		if err != nil {
			return err
		}

		// Create the snapshot path.
		if len(snapshots) > 0 {
			err = os.MkdirAll(snapshotsPath, 0711)
			if err != nil {
				return errors.Wrapf(err, "Failed to create directory '%s'", snapshotsPath)
			}
		}

		for _, snapshot := range snapshots {
			_, snapName, _ := shared.InstanceGetParentAndSnapshotName(snapshot.Name())
			target := filepath.Join(snapshotsPath, snapName)

			// Copy the snapshot.
			err = snapshot.MountTask(func(mountPath string, op *operations.Operation) error {
				_, err := rsync.LocalCopy(mountPath, target, bwlimit, true)
				if err != nil {
					return err
				}

				return nil
			}, op)
			if err != nil {
				return err
			}
		}
	}

	// Copy the parent volume itself.
	target := filepath.Join(targetPath, "container")
	err := vol.MountTask(func(mountPath string, op *operations.Operation) error {
		_, err := rsync.LocalCopy(mountPath, target, bwlimit, true)
		if err != nil {
			return err
		}

		return nil
	}, op)
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2011 bjarneh@ifi.uio.no. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package binheap_test

import (
    "rand"
    "sort"
    "testing"
    . "github.com/bjarneh/binheap"
)

type heaper struct {
    i int
}

func (h *heaper) Priority() int {
    return h.i
}

func randomInts(size int) (slice []int) {

    slice = make([]int, size)

    for i := 0; i < size; i++ {
        slice[i] = rand.Int()
    }

    return slice
}

func randomHeapable(size int) (slice []Heapable) {

    ints := randomInts(size)

    slice = make([]Heapable, size)

    for i := 0; i < size; i++ {
        slice[i] = &heaper{ints[i]}
    }

    return slice
}

func TestIntSort(t *testing.T) {

    nums := randomInts(1000)

    SortInt(nums)

    for i := 1; i < len(nums); i++ {
        if nums[i-1] > nums[i] {
            t.Fatalf("binheap.SortInts() does not sort\n")
        }
    }
}

func TestSort(t *testing.T) {

    nums := randomHeapable(1000)

    Sort(nums)

    for i := 1; i < len(nums); i++ {
        if nums[i-1].Priority() > nums[i].Priority() {
            t.Fatalf("binheap.Sort() does not sort\n")
        }
    }
}

func TestHeap(t *testing.T) {

    heap := NewSize(1) // make it short to test resize
    nums := randomInts(1000)
    elmts := make([]Heapable, len(nums))

    for i := 0; i < len(nums); i++ {
        elmts[i] = &heaper{nums[i]}
    }

    heap.AddSlice(elmts)

    if heap.Len() != len(nums) {
        t.Fatalf("heap.Len() != %d\n", len(nums))
    }

    sort.SortInts(nums)

    for i := 0; i < len(nums); i++ {
        if heap.Remove().Priority() != nums[i] {
            t.Fatalf("binheap: Priority() != %d\n", nums[i])
        }
    }

    if !heap.Empty() {
        t.Fatalf("heap not empty\n")
    }

    nums = randomInts(1000)

    // add single elements.. i.e. not slice
    for i := 0; i < len(nums); i++ {
        heap.Add(&heaper{nums[i]})
    }

    sort.SortInts(nums)

    for i := 0; i < len(nums); i++ {
        if heap.Remove().Priority() != nums[i] {
            t.Fatalf("binheap: Priority() != %d\n", nums[i])
        }
    }

}

func BenchmarkStdlibSortInts(b *testing.B) {
    for i := 0; i < b.N; i++ {
        nums := randomInts(1000)
        sort.SortInts(nums)
    }
}

func BenchmarkHeapSortInt(b *testing.B) {
    for i := 0; i < b.N; i++ {
        nums := randomInts(1000)
        SortInt(nums)
    }
}
<commit_msg>sorting changed. now gofixed<commit_after>// Copyright 2011 bjarneh@ifi.uio.no. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package binheap_test

import (
	"rand"
	"sort"
	"testing"
	. "github.com/bjarneh/binheap"
)

type heaper struct {
	i int
}

func (h *heaper) Priority() int {
	return h.i
}

func randomInts(size int) (slice []int) {

	slice = make([]int, size)

	for i := 0; i < size; i++ {
		slice[i] = rand.Int()
	}

	return slice
}

func randomHeapable(size int) (slice []Heapable) {

	ints := randomInts(size)

	slice = make([]Heapable, size)

	for i := 0; i < size; i++ {
		slice[i] = &heaper{ints[i]}
	}

	return slice
}

func TestIntSort(t *testing.T) {

	nums := randomInts(1000)

	SortInt(nums)

	for i := 1; i < len(nums); i++ {
		if nums[i-1] > nums[i] {
			t.Fatalf("binheap.SortInts() does not sort\n")
		}
	}
}

func TestSort(t *testing.T) {

	nums := randomHeapable(1000)

	Sort(nums)

	for i := 1; i < len(nums); i++ {
		if nums[i-1].Priority() > nums[i].Priority() {
			t.Fatalf("binheap.Sort() does not sort\n")
		}
	}
}

func TestHeap(t *testing.T) {

	heap := NewSize(1) // make it short to test resize
	nums := randomInts(1000)
	elmts := make([]Heapable, len(nums))

	for i := 0; i < len(nums); i++ {
		elmts[i] = &heaper{nums[i]}
	}

	heap.AddSlice(elmts)

	if heap.Len() != len(nums) {
		t.Fatalf("heap.Len() != %d\n", len(nums))
	}

	sort.Ints(nums)

	for i := 0; i < len(nums); i++ {
		if heap.Remove().Priority() != nums[i] {
			t.Fatalf("binheap: Priority() != %d\n", nums[i])
		}
	}

	if !heap.Empty() {
		t.Fatalf("heap not empty\n")
	}

	nums = randomInts(1000)

	// add single elements.. i.e. not slice
	for i := 0; i < len(nums); i++ {
		heap.Add(&heaper{nums[i]})
	}

	sort.Ints(nums)

	for i := 0; i < len(nums); i++ {
		if heap.Remove().Priority() != nums[i] {
			t.Fatalf("binheap: Priority() != %d\n", nums[i])
		}
	}

}

func BenchmarkStdlibSortInts(b *testing.B) {
	for i := 0; i < b.N; i++ {
		nums := randomInts(1000)
		sort.Ints(nums)
	}
}

func BenchmarkHeapSortInt(b *testing.B) {
	for i := 0; i < b.N; i++ {
		nums := randomInts(1000)
		SortInt(nums)
	}
}
<|endoftext|>


<commit_before>package bitseq

import (
	"encoding/json"
	"fmt"

	log "github.com/Sirupsen/logrus"
	"github.com/docker/libnetwork/datastore"
	"github.com/docker/libnetwork/types"
)

// Key provides the Key to be used in KV Store
func (h *Handle) Key() []string {
	h.Lock()
	defer h.Unlock()
	return []string{h.app, h.id}
}

// KeyPrefix returns the immediate parent key that can be used for tree walk
func (h *Handle) KeyPrefix() []string {
	h.Lock()
	defer h.Unlock()
	return []string{h.app}
}

// Value marshals the data to be stored in the KV store
func (h *Handle) Value() []byte {
	b, err := h.ToByteArray()
	if err != nil {
		log.Warnf("Failed to serialize Handle: %v", err)
		b = []byte{}
	}
	jv, err := json.Marshal(b)
	if err != nil {
		log.Warnf("Failed to json encode bitseq handler byte array: %v", err)
		return []byte{}
	}
	return jv
}

// SetValue unmarshals the data from the KV store
func (h *Handle) SetValue(value []byte) error {
	return h.FromByteArray(value)
}

// Index returns the latest DB Index as seen by this object
func (h *Handle) Index() uint64 {
	h.Lock()
	defer h.Unlock()
	return h.dbIndex
}

// SetIndex method allows the datastore to store the latest DB Index into this object
func (h *Handle) SetIndex(index uint64) {
	h.Lock()
	h.dbIndex = index
	h.dbExists = true
	h.Unlock()
}

// Exists method is true if this object has been stored in the DB.
func (h *Handle) Exists() bool {
	h.Lock()
	defer h.Unlock()
	return h.dbExists
}

func (h *Handle) watchForChanges() error {
	h.Lock()
	store := h.store
	h.Unlock()

	if store == nil {
		return nil
	}

	kvpChan, err := store.KVStore().Watch(datastore.Key(h.Key()...), nil)
	if err != nil {
		return err
	}
	go func() {
		for {
			select {
			case kvPair := <-kvpChan:
				// Only process remote update
				if kvPair != nil && (kvPair.LastIndex != h.Index()) {
					err := h.fromDsValue(kvPair.Value)
					if err != nil {
						log.Warnf("Failed to reconstruct bitseq handle from ds watch: %s", err.Error())
					} else {
						h.SetIndex(kvPair.LastIndex)
					}
				}
			}
		}
	}()
	return nil
}

func (h *Handle) fromDsValue(value []byte) error {
	var ba []byte
	if err := json.Unmarshal(value, &ba); err != nil {
		return fmt.Errorf("failed to decode json: %s", err.Error())
	}
	if err := h.FromByteArray(ba); err != nil {
		return fmt.Errorf("failed to decode handle: %s", err.Error())
	}
	return nil
}

func (h *Handle) writeToStore() error {
	h.Lock()
	store := h.store
	h.Unlock()
	if store == nil {
		return nil
	}
	err := store.PutObjectAtomic(h)
	if err == datastore.ErrKeyModified {
		return types.RetryErrorf("failed to perform atomic write (%v). retry might fix the error", err)
	}
	return err
}

func (h *Handle) deleteFromStore() error {
	h.Lock()
	store := h.store
	h.Unlock()
	if store == nil {
		return nil
	}
	return store.DeleteObjectAtomic(h)
}
<commit_msg>BitSequence should unmarshal data during get<commit_after>package bitseq

import (
	"encoding/json"
	"fmt"

	log "github.com/Sirupsen/logrus"
	"github.com/docker/libnetwork/datastore"
	"github.com/docker/libnetwork/types"
)

// Key provides the Key to be used in KV Store
func (h *Handle) Key() []string {
	h.Lock()
	defer h.Unlock()
	return []string{h.app, h.id}
}

// KeyPrefix returns the immediate parent key that can be used for tree walk
func (h *Handle) KeyPrefix() []string {
	h.Lock()
	defer h.Unlock()
	return []string{h.app}
}

// Value marshals the data to be stored in the KV store
func (h *Handle) Value() []byte {
	b, err := h.ToByteArray()
	if err != nil {
		log.Warnf("Failed to serialize Handle: %v", err)
		b = []byte{}
	}
	jv, err := json.Marshal(b)
	if err != nil {
		log.Warnf("Failed to json encode bitseq handler byte array: %v", err)
		return []byte{}
	}
	return jv
}

// SetValue unmarshals the data from the KV store
func (h *Handle) SetValue(value []byte) error {
	var b []byte
	if err := json.Unmarshal(value, &b); err != nil {
		return err
	}

	return h.FromByteArray(b)
}

// Index returns the latest DB Index as seen by this object
func (h *Handle) Index() uint64 {
	h.Lock()
	defer h.Unlock()
	return h.dbIndex
}

// SetIndex method allows the datastore to store the latest DB Index into this object
func (h *Handle) SetIndex(index uint64) {
	h.Lock()
	h.dbIndex = index
	h.dbExists = true
	h.Unlock()
}

// Exists method is true if this object has been stored in the DB.
func (h *Handle) Exists() bool {
	h.Lock()
	defer h.Unlock()
	return h.dbExists
}

func (h *Handle) watchForChanges() error {
	h.Lock()
	store := h.store
	h.Unlock()

	if store == nil {
		return nil
	}

	kvpChan, err := store.KVStore().Watch(datastore.Key(h.Key()...), nil)
	if err != nil {
		return err
	}
	go func() {
		for {
			select {
			case kvPair := <-kvpChan:
				// Only process remote update
				if kvPair != nil && (kvPair.LastIndex != h.Index()) {
					err := h.fromDsValue(kvPair.Value)
					if err != nil {
						log.Warnf("Failed to reconstruct bitseq handle from ds watch: %s", err.Error())
					} else {
						h.SetIndex(kvPair.LastIndex)
					}
				}
			}
		}
	}()
	return nil
}

func (h *Handle) fromDsValue(value []byte) error {
	var ba []byte
	if err := json.Unmarshal(value, &ba); err != nil {
		return fmt.Errorf("failed to decode json: %s", err.Error())
	}
	if err := h.FromByteArray(ba); err != nil {
		return fmt.Errorf("failed to decode handle: %s", err.Error())
	}
	return nil
}

func (h *Handle) writeToStore() error {
	h.Lock()
	store := h.store
	h.Unlock()
	if store == nil {
		return nil
	}
	err := store.PutObjectAtomic(h)
	if err == datastore.ErrKeyModified {
		return types.RetryErrorf("failed to perform atomic write (%v). retry might fix the error", err)
	}
	return err
}

func (h *Handle) deleteFromStore() error {
	h.Lock()
	store := h.store
	h.Unlock()
	if store == nil {
		return nil
	}
	return store.DeleteObjectAtomic(h)
}
<|endoftext|>


<commit_before>package auth

import (
	"crypto/rsa"
	"encoding/json"
	"log"
	"net/http"
	"strings"
	"time"

	"github.com/SermoDigital/jose/crypto"
	"github.com/SermoDigital/jose/jws"
	"github.com/SermoDigital/jose/jwt"
)

// Token firebase secure token
type Token struct {
	kid, token string
}

// Firebase module to verify and extract information from firebase JWT tokens
type Firebase struct {
	publicKeys map[string]*rsa.PublicKey
	projectID  string
}

// NewFirebase loads the firebase keys
func NewFirebase(projectID string) (*Firebase, error) {

	keys, err := updatePublicKeys()
	if err != nil {
		return nil, err
	}

	fb := new(Firebase)
	fb.projectID = projectID
	fb.publicKeys = keys
	return fb, nil
}

func updatePublicKeys() (map[string]*rsa.PublicKey, error) {

	log.Printf("Requesting firebase tokens")
	tokens := make(map[string]interface{})
	err := getFirebaseTokens(tokens)
	if err != nil {
		return nil, err
	}

	keys := make(map[string]*rsa.PublicKey)
	for kid, token := range tokens {
		publicKey, err := crypto.ParseRSAPublicKeyFromPEM([]byte(token.(string)))
		if err != nil {
			log.Printf("Error parsing kid %s, %v", kid, err)
		} else {
			log.Printf("Validated kid %s", kid)
			keys[kid] = publicKey
		}
	}

	return keys, nil
}

// firebase tokens must be signed by one of the keys provided at a certain url.
func getFirebaseTokens(tokens map[string]interface{}) error {
	return getJSON("https://www.googleapis.com/robot/v1/metadata/x509/securetoken@system.gserviceaccount.com", &tokens)
}

var myClient = &http.Client{Timeout: 30 * time.Second}

func getJSON(url string, target interface{}) error {
	r, err := myClient.Get(url)
	if err != nil {
		return err
	}
	defer r.Body.Close()
	return json.NewDecoder(r.Body).Decode(target)
}

// Verify to satisfy the auth.Verify interface
func (fa Firebase) Verify(accessToken string) (jwt.Claims, error) {

	// empty string is clearly invalid
	if accessToken == "" {
		return nil, ErrNilToken
	}

	token, err := jws.ParseJWT([]byte(accessToken))
	if err != nil {
		return nil, err
	}

	// TODO extract kid from header and only verify against that key
	// https://firebase.google.com/docs/auth/admin/verify-id-tokens

	// validate against firebase keys
	for _, key := range fa.publicKeys {
		err = token.Validate(key, crypto.SigningMethodRS256)
		// verification errors indicate that the token isn't valid for this key
		if err == nil || !strings.Contains(err.Error(), "verification error") {
			break
		}
	}

	switch err {
	case jwt.ErrTokenIsExpired:
		err = ErrTokenExpired
		break
	case crypto.ErrECDSAVerification:
		err = ErrECDSAVerification
		break
	case jws.ErrNotCompact:
		err = ErrNotCompact
		break
	}

	aud := token.Claims().Get("aud")
	if aud != fa.projectID {
		log.Printf("Invalid authorisation token audience %v, expecting %v", aud, fa.projectID)
		err = ErrInvalidAud
	}

	iss := token.Claims().Get("iss")
	expectedIss := "https://securetoken.google.com/" + fa.projectID
	if iss != expectedIss {
		log.Printf("Invalid authorisation token issuer %v, expected %v", iss, expectedIss)
		err = ErrInvalidIss
	}

	return token.Claims(), err
}
<commit_msg>Refactor firebase key retrieval method<commit_after>package auth

import (
	"crypto/rsa"
	"encoding/json"
	"log"
	"net/http"
	"strings"
	"time"

	"github.com/SermoDigital/jose/crypto"
	"github.com/SermoDigital/jose/jws"
	"github.com/SermoDigital/jose/jwt"
)

// Firebase module to verify and extract information from firebase JWT tokens
type Firebase struct {
	publicKeys map[string]*rsa.PublicKey
	projectID  string
}

// NewFirebase loads the firebase keys
func NewFirebase(projectID string) (*Firebase, error) {
	fb := new(Firebase)
	fb.projectID = projectID
	return fb, fb.UpdatePublicKeys()
}

// UpdatePublicKeys retrieves the latest firebase keys
func (fa Firebase) UpdatePublicKeys() error {
	log.Printf("Requesting firebase tokens")
	tokens := make(map[string]interface{})
	err := getFirebaseTokens(tokens)
	if err != nil {
		return err
	}
	fa.publicKeys = make(map[string]*rsa.PublicKey)
	for kid, token := range tokens {
		publicKey, err := crypto.ParseRSAPublicKeyFromPEM([]byte(token.(string)))
		if err != nil {
			log.Printf("Error parsing kid %s, %v", kid, err)
		} else {
			log.Printf("Validated kid %s", kid)
			fa.publicKeys[kid] = publicKey
		}
	}
	return nil
}

// firebase tokens must be signed by one of the keys provided at a certain url.
func getFirebaseTokens(tokens map[string]interface{}) error {
	return getJSON("https://www.googleapis.com/robot/v1/metadata/x509/securetoken@system.gserviceaccount.com", &tokens)
}

var myClient = &http.Client{Timeout: 30 * time.Second}

func getJSON(url string, target interface{}) error {
	r, err := myClient.Get(url)
	if err != nil {
		return err
	}
	defer r.Body.Close()
	return json.NewDecoder(r.Body).Decode(target)
}

// Verify to satisfy the auth.Verify interface
func (fa Firebase) Verify(accessToken string) (jwt.Claims, error) {

	// empty string is clearly invalid
	if accessToken == "" {
		return nil, ErrNilToken
	}

	token, err := jws.ParseJWT([]byte(accessToken))
	if err != nil {
		return nil, err
	}

	// TODO extract kid from header and only verify against that key
	// https://firebase.google.com/docs/auth/admin/verify-id-tokens

	// validate against firebase keys
	for _, key := range fa.publicKeys {
		err = token.Validate(key, crypto.SigningMethodRS256)
		// verification errors indicate that the token isn't valid for this key
		// once we are targeting specific keys then we can remove this hacky solution
		if err == nil || !strings.Contains(err.Error(), "verification error") {
			break
		}
	}

	switch err {
	case jwt.ErrTokenIsExpired:
		err = ErrTokenExpired
		break
	case crypto.ErrECDSAVerification:
		err = ErrECDSAVerification
		break
	case jws.ErrNotCompact:
		err = ErrNotCompact
		break
	}

	aud := token.Claims().Get("aud")
	if aud != fa.projectID {
		log.Printf("Invalid authorisation token audience %v, expecting %v", aud, fa.projectID)
		err = ErrInvalidAud
	}

	iss := token.Claims().Get("iss")
	expectedIss := "https://securetoken.google.com/" + fa.projectID
	if iss != expectedIss {
		log.Printf("Invalid authorisation token issuer %v, expected %v", iss, expectedIss)
		err = ErrInvalidIss
	}

	return token.Claims(), err
}
<|endoftext|>


<commit_before>package driver

import (
	"fmt"
	"net"

	"golang.org/x/crypto/ssh"

	"github.com/Spirals-Team/docker-machine-driver-g5k/api"
	"github.com/docker/machine/libmachine/log"
)

// CheckVpnConnection check if the VPN is connected and properly configured (DNS) by trying to connect to the site frontend SSH server using its hostname
func (d *Driver) CheckVpnConnection(site string) error {
	// construct site frontend hostname
	frontend := fmt.Sprintf("frontend.%s.grid5000.fr:22", site)

	// try to connect to the frontend SSH server
	sshConfig := &ssh.ClientConfig{}
	_, err := ssh.Dial("tcp", frontend, sshConfig)

	// we need to check if the error is network-related because the SSH Dial will always return an error due to the Authentication being not configured
	if _, ok := err.(*net.OpError); ok {
		return fmt.Errorf("Connection to frontend of '%s' site failed. Please check if the site is not undergoing maintenance and your VPN client is connected and properly configured (see driver documentation for more information)", d.G5kSite)
	}

	return nil
}

func (d *Driver) submitNewJobReservation() error {
	// if a job ID is provided, skip job reservation
	if d.G5kJobID != 0 {
		log.Infof("Skipping job reservation and using job ID '%v'", d.G5kJobID)
		return nil
	}

	// submit new Job request
	jobID, err := d.G5kAPI.SubmitJob(api.JobRequest{
		Resources:  fmt.Sprintf("nodes=1,walltime=%s", d.G5kWalltime),
		Command:    "sleep 365d",
		Properties: d.G5kResourceProperties,
		Types:      []string{"deploy"},
	})
	if err != nil {
		return fmt.Errorf("Error when submitting new job: %s", err.Error())
	}

	if err = d.G5kAPI.WaitUntilJobIsReady(jobID); err != nil {
		return fmt.Errorf("Error when waiting for job to be running: %s", err.Error())
	}

	// job is running, keep its ID for future usage
	d.G5kJobID = jobID
	return nil
}

func (d *Driver) handleDeploymentError() {
	// if deployment fail, we can't recover from this error, so we kill the job
	log.Infof("Unrecoverable error in deployment, killing job ID '%d'...", d.G5kJobID)
	d.G5kAPI.KillJob(d.G5kJobID)
}

func (d *Driver) submitNewDeployment() error {
	// if a host to provision is set, skip host deployment
	if d.G5kHostToProvision != "" {
		log.Infof("Skipping host deployment and provisionning host '%s' only", d.G5kHostToProvision)
		return nil
	}

	// get job informations
	job, err := d.G5kAPI.GetJob(d.G5kJobID)
	if err != nil {
		return fmt.Errorf("Error when getting job (id: '%d') informations: %s", d.G5kJobID, err.Error())
	}

	// deploy environment
	deploymentID, err := d.G5kAPI.SubmitDeployment(api.DeploymentRequest{
		Nodes:       job.Nodes,
		Environment: d.G5kImage,
		Key:         string(d.SSHKeyPair.PublicKey),
	})
	if err != nil {
		d.handleDeploymentError()
		return fmt.Errorf("Error when submitting new deployment: %s", err.Error())
	}

	// waiting deployment to finish (REQUIRED or you will interfere with kadeploy)
	if err = d.G5kAPI.WaitUntilDeploymentIsFinished(deploymentID); err != nil {
		d.handleDeploymentError()
		return fmt.Errorf("Error when waiting for deployment to finish: %s", err.Error())
	}

	return nil
}
<commit_msg>fix(driver): Fix VPN connection checking function (API + error message)<commit_after>package driver

import (
	"fmt"
	"net"

	"golang.org/x/crypto/ssh"

	"github.com/Spirals-Team/docker-machine-driver-g5k/api"
	"github.com/docker/machine/libmachine/log"
)

// CheckVpnConnection check if the VPN is connected and properly configured (DNS) by trying to connect to the site frontend SSH server using its hostname
func CheckVpnConnection(site string) error {
	// construct site frontend hostname
	frontend := fmt.Sprintf("frontend.%s.grid5000.fr:22", site)

	// try to connect to the frontend SSH server
	sshConfig := &ssh.ClientConfig{}
	_, err := ssh.Dial("tcp", frontend, sshConfig)

	// we need to check if the error is network-related because the SSH Dial will always return an error due to the Authentication being not configured
	if _, ok := err.(*net.OpError); ok {
		return fmt.Errorf("Connection to frontend of '%s' site failed. Please check if the site is not undergoing maintenance and your VPN client is connected and properly configured (see driver documentation for more information)", site)
	}

	return nil
}

func (d *Driver) submitNewJobReservation() error {
	// if a job ID is provided, skip job reservation
	if d.G5kJobID != 0 {
		log.Infof("Skipping job reservation and using job ID '%v'", d.G5kJobID)
		return nil
	}

	// submit new Job request
	jobID, err := d.G5kAPI.SubmitJob(api.JobRequest{
		Resources:  fmt.Sprintf("nodes=1,walltime=%s", d.G5kWalltime),
		Command:    "sleep 365d",
		Properties: d.G5kResourceProperties,
		Types:      []string{"deploy"},
	})
	if err != nil {
		return fmt.Errorf("Error when submitting new job: %s", err.Error())
	}

	if err = d.G5kAPI.WaitUntilJobIsReady(jobID); err != nil {
		return fmt.Errorf("Error when waiting for job to be running: %s", err.Error())
	}

	// job is running, keep its ID for future usage
	d.G5kJobID = jobID
	return nil
}

func (d *Driver) handleDeploymentError() {
	// if deployment fail, we can't recover from this error, so we kill the job
	log.Infof("Unrecoverable error in deployment, killing job ID '%d'...", d.G5kJobID)
	d.G5kAPI.KillJob(d.G5kJobID)
}

func (d *Driver) submitNewDeployment() error {
	// if a host to provision is set, skip host deployment
	if d.G5kHostToProvision != "" {
		log.Infof("Skipping host deployment and provisionning host '%s' only", d.G5kHostToProvision)
		return nil
	}

	// get job informations
	job, err := d.G5kAPI.GetJob(d.G5kJobID)
	if err != nil {
		return fmt.Errorf("Error when getting job (id: '%d') informations: %s", d.G5kJobID, err.Error())
	}

	// deploy environment
	deploymentID, err := d.G5kAPI.SubmitDeployment(api.DeploymentRequest{
		Nodes:       job.Nodes,
		Environment: d.G5kImage,
		Key:         string(d.SSHKeyPair.PublicKey),
	})
	if err != nil {
		d.handleDeploymentError()
		return fmt.Errorf("Error when submitting new deployment: %s", err.Error())
	}

	// waiting deployment to finish (REQUIRED or you will interfere with kadeploy)
	if err = d.G5kAPI.WaitUntilDeploymentIsFinished(deploymentID); err != nil {
		d.handleDeploymentError()
		return fmt.Errorf("Error when waiting for deployment to finish: %s", err.Error())
	}

	return nil
}
<|endoftext|>


<commit_before>package dskvs

import (
	"bytes"
	"encoding/json"
	"fmt"
	"os"
	"runtime"
	"strconv"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

type Data struct {
	Descr string
}

/*
	Boilerplate
*/

func setUp(t *testing.T) *Store {
	store, err := NewStore("./db")
	if err != nil {
		t.Fatalf("Error creating store, %v", err)
	}

	err = store.Load()
	if err != nil {
		t.Fatalf("Error loading store, %v", err)
	}
	return store
}

func tearDown(store *Store, t *testing.T) {
	err := store.Close()
	if err != nil {
		t.Fatalf("Error closing store, %v", err)
	}
	err = os.RemoveAll(store.storagePath)
	if err != nil {
		t.Fatalf("Error deleting storage path, %v", err)
	}
}

func generateData(d Data, t *testing.T) []byte {
	dataBytes, err := json.Marshal(d)
	if err != nil {
		t.Fatal("Error with test data, %v", err)
	}
	return dataBytes
}

/*
	Common checks
*/

func checkGetIsEmpty(store *Store, key string, t *testing.T) {
	_, err := store.Get(key)
	if err == nil {
		t.Fatalf("Expected to receive KeyError but no error")
	}
	switch err := err.(type) {
	case KeyError:
		// Expected
		break
	default:
		t.Fatalf("Expected to receive KeyError but got <%s>", err)
	}
}

/*
	Test cases, single goroutine
*/

func TestCreatingStore(t *testing.T) {
	store := setUp(t)
	tearDown(store, t)
}

func TestSingleOperation(t *testing.T) {

	store := setUp(t)
	defer tearDown(store, t)

	key := "artist/daftpunk"
	expected := generateData(Data{"The peak of awesome"}, t)

	err := store.Put(key, expected)
	if err != nil {
		t.Fatalf("Error putting data in, %v", err)
	}

	actual, err := store.Get(key)
	if err != nil {
		t.Fatalf("Error getting data back, %v", err)
	}

	if !bytes.Equal(expected, actual) {
		t.Fatalf("Expected <%s> but was <%s>",
			expected,
			actual)
	}

	err = store.Delete(key)
	if err != nil {
		t.Fatalf("Error deleting data we just put, %v", err)
	}

	checkGetIsEmpty(store, key, t)

}

func TestMultipleOperations(t *testing.T) {

	store := setUp(t)
	defer tearDown(store, t)

	coll := "artist"
	baseKey := "daftpunk"

	var key string
	var expected []byte
	var expectedList [][]byte
	var actual []byte
	for i := int(0); i < 10; i++ {
		key = coll + CollKeySep + baseKey + strconv.Itoa(i)
		expected = generateData(Data{"The peak of awesome" + strconv.Itoa(i)}, t)

		err := store.Put(key, expected)
		if err != nil {
			t.Errorf("Error putting data in, %v", err)
		}
		actual, err = store.Get(key)
		if err != nil {
			t.Errorf("Error getting data back, %v", err)
		}

		if !bytes.Equal(expected, actual) {
			t.Errorf("Expected <%s> but was <%s>",
				expected,
				actual)
		}

		expectedList = append(expectedList, expected)
	}

	actuaList, err := store.GetAll(coll)
	if err != nil {
		t.Fatalf("Error getting all keys with coll <%s>", coll)
	}
	if len(actuaList) != len(expectedList) {
		t.Fatalf("Expected to read %d items, but read %d",
			len(expectedList),
			len(actuaList))
	}

	err = store.DeleteAll("artist")
	if err != nil {
		t.Fatal("Error deleting key, %v", err)
	}
	for i := int(0); i < 10; i++ {
		key = coll + CollKeySep + baseKey + strconv.Itoa(i)
		checkGetIsEmpty(store, key, t)
	}
}

func TestStorePersistPutAfterClose(t *testing.T) {
	store := setUp(t)

	key := "artist/the prodigy"
	expected := generateData(Data{"Beyond the peak"}, t)

	if err := store.Put(key, expected); err != nil {
		tearDown(store, t)
		t.Fatalf("Error putting data in, %v", err)
	}

	// Don't use tearDown as it deletes the storage after use
	if err := store.Close(); err != nil {
		t.Fatalf("Error closing store, %v", err)
	}

	otherStore := setUp(t)
	defer tearDown(otherStore, t)

	actual, err := otherStore.Get(key)
	if err != nil {
		t.Fatalf("Error getting data back, %v", err)
	}
	if !bytes.Equal(expected, actual) {
		t.Fatalf("Expected <%s> but was <%s>",
			expected,
			actual)
	}
}

func TestStorePersistDeleteAfterClose(t *testing.T) {
	store := setUp(t)

	key := "artist/the prodigy"
	expected := generateData(Data{"Beyond the peak"}, t)

	if err := store.Put(key, expected); err != nil {
		tearDown(store, t)
		t.Fatalf("Error putting value in, %v", err)
	}

	if err := store.Delete(key); err != nil {
		tearDown(store, t)
		t.Fatalf("Error deleting value, %v", err)
	}

	// Don't use tearDown as it deletes the storage after use
	if err := store.Close(); err != nil {
		t.Fatalf("Error closing store, %v", err)
	}

	otherStore := setUp(t)
	defer tearDown(otherStore, t)

	checkGetIsEmpty(otherStore, key, t)

}

/*
	Test cases, multiple goroutine
*/

func TestMultipleGoroutine(t *testing.T) {
	var kvCount int
	if testing.Short() {
		kvCount = 1000
	} else {
		kvCount = 10000
	}
	coll := "games"
	baseKey := "total annihilation #"

	type KeyValue struct {
		Key   string
		Value []byte
	}
	countGet := int64(0)
	countPut := int64(0)
	countDelete := int64(0)
	start := time.Now()

	store := setUp(t)
	defer tearDown(store, t)
	defer runtime.GOMAXPROCS(runtime.GOMAXPROCS(runtime.NumCPU()))

	var writeGroup sync.WaitGroup
	var readGroup sync.WaitGroup
	var deleteGroup sync.WaitGroup
	writeErr := make(chan error)
	readErr := make(chan error)
	deleteErr := make(chan error)

	expectedList := make([]KeyValue, 0)
	for i := int(0); i < kvCount; i++ {
		key := coll + CollKeySep + baseKey + strconv.Itoa(i)
		data := generateData(Data{"It's fun!"}, t)
		kv := KeyValue{key, data}
		expectedList = append(expectedList, kv)
	}

	for _, kv := range expectedList {
		// Put all the values concurrently, ensure there's no error
		writeGroup.Add(1)
		go func(pair KeyValue, cErr chan error) {
			defer writeGroup.Done()
			err := store.Put(pair.Key, pair.Value)
			atomic.AddInt64(&countPut, 1)
			if err != nil {
				t.Fatal("Received an error", err)
				cErr <- err
			}

		}(kv, writeErr)
	}

	writeGroup.Wait()

	if len(writeErr) != 0 {
		t.Fatalf("Failed to write values concurrently, got %d errors",
			len(writeErr))
	}

	for _, kv := range expectedList {
		// Get each value, ensure they're the same
		readGroup.Add(1)
		go func(kv KeyValue, cErr chan error) {
			defer readGroup.Done()

			expected := kv.Value
			actual, err := store.Get(kv.Key)
			atomic.AddInt64(&countGet, 1)
			if err != nil {
				cErr <- err
			} else if !bytes.Equal(expected, actual) {
				t.Errorf("Expected <%s> but was <%s>",
					expected,
					actual)
			}

		}(kv, readErr)

		// GetAll values, ensure this slice is as big as all the value
		// we've put
		readGroup.Add(1)
		go func(cErr chan error) {
			defer readGroup.Done()

			actual, err := store.GetAll(coll)
			atomic.AddInt64(&countGet, int64(len(actual)))
			if err != nil {
				cErr <- err
			} else if len(actual) != len(expectedList) {
				t.Errorf("Expected len(store.GetAll)=<%s> but was <%s>",
					len(expectedList),
					len(actual))
			}
		}(readErr)
	}

	readGroup.Wait()

	if len(readErr) != 0 {
		t.Fatalf("Failed to write values concurrently, got %d errors",
			len(readErr))
	}

	for _, kv := range expectedList {
		// Put all the values concurrently, ensure there's no error
		deleteGroup.Add(1)
		go func(pair KeyValue, cErr chan error) {
			defer deleteGroup.Done()
			err := store.Delete(pair.Key)
			atomic.AddInt64(&countDelete, 1)
			if err != nil {
				t.Fatal("Received an error", err)
				cErr <- err
			}

		}(kv, deleteErr)
	}

	deleteGroup.Wait()

	if len(deleteErr) != 0 {
		t.Fatalf("Failed to write values concurrently, got %d errors",
			len(deleteErr))
	}

	err := store.DeleteAll(coll)
	atomic.AddInt64(&countDelete, int64(len(expectedList)))
	if err != nil {
		t.Fatalf("Error deleting all", err)
	}
	for _, kv := range expectedList {
		atomic.AddInt64(&countGet, 1)
		checkGetIsEmpty(store, kv.Key, t)
	}
	duration := time.Now().Sub(start)
	fmt.Printf("Concurrent test\n")
	fmt.Printf("- Get operations=%d\n", countGet)
	fmt.Printf("- Put operations=%d\n", countPut)
	fmt.Printf("- Delete operations=%d\n", countPut)
	fmt.Printf("in %fs ", duration.Seconds())
	fmt.Printf("by %d cpus, using %d concurrent goroutines\n",
		runtime.NumCPU(), kvCount)

}
<commit_msg>Add error-path testing.<commit_after>package dskvs

import (
	"bytes"
	"encoding/json"
	"fmt"
	"os"
	"runtime"
	"strconv"
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

type Data struct {
	Descr string
}

///////////////////////////////////////////////////////////////////////////////
// Boilerplate
///////////////////////////////////////////////////////////////////////////////

func setUp(t *testing.T) *Store {
	store, err := NewStore("./db")
	if err != nil {
		t.Fatalf("Error creating store, %v", err)
	}

	err = store.Load()
	if err != nil {
		t.Fatalf("Error loading store, %v", err)
	}
	return store
}

func tearDown(store *Store, t *testing.T) {
	err := store.Close()
	if err != nil {
		t.Fatalf("Error closing store, %v", err)
	}
	err = os.RemoveAll(store.storagePath)
	if err != nil {
		t.Fatalf("Error deleting storage path, %v", err)
	}
}

func generateData(d Data, t *testing.T) []byte {
	dataBytes, err := json.Marshal(d)
	if err != nil {
		t.Fatal("Error with test data, %v", err)
	}
	return dataBytes
}

///////////////////////////////////////////////////////////////////////////////
// Common checks
///////////////////////////////////////////////////////////////////////////////

func checkGetIsEmpty(store *Store, key string, t *testing.T) {
	_, err := store.Get(key)
	if err == nil {
		t.Fatalf("Expected to receive KeyError but no error")
	}
	switch err := err.(type) {
	case KeyError:
		// Expected
		break
	default:
		t.Fatalf("Expected to receive KeyError but got <%s>", err)
	}
}

///////////////////////////////////////////////////////////////////////////////
// Single goroutine
///////////////////////////////////////////////////////////////////////////////

// Normal cases

func TestCreatingStore(t *testing.T) {
	store := setUp(t)
	tearDown(store, t)
}

func TestSingleOperation(t *testing.T) {

	store := setUp(t)
	defer tearDown(store, t)

	key := "artist/daftpunk"
	expected := generateData(Data{"The peak of awesome"}, t)

	err := store.Put(key, expected)
	if err != nil {
		t.Fatalf("Error putting data in, %v", err)
	}

	actual, err := store.Get(key)
	if err != nil {
		t.Fatalf("Error getting data back, %v", err)
	}

	if !bytes.Equal(expected, actual) {
		t.Fatalf("Expected <%s> but was <%s>",
			expected,
			actual)
	}

	err = store.Delete(key)
	if err != nil {
		t.Fatalf("Error deleting data we just put, %v", err)
	}

	checkGetIsEmpty(store, key, t)

}

func TestMultipleOperations(t *testing.T) {

	store := setUp(t)
	defer tearDown(store, t)

	coll := "artist"
	baseKey := "daftpunk"

	var key string
	var expected []byte
	var expectedList [][]byte
	var actual []byte
	for i := int(0); i < 10; i++ {
		key = coll + CollKeySep + baseKey + strconv.Itoa(i)
		expected = generateData(Data{"The peak of awesome" + strconv.Itoa(i)}, t)

		err := store.Put(key, expected)
		if err != nil {
			t.Errorf("Error putting data in, %v", err)
		}
		actual, err = store.Get(key)
		if err != nil {
			t.Errorf("Error getting data back, %v", err)
		}

		if !bytes.Equal(expected, actual) {
			t.Errorf("Expected <%s> but was <%s>",
				expected,
				actual)
		}

		expectedList = append(expectedList, expected)
	}

	actuaList, err := store.GetAll(coll)
	if err != nil {
		t.Fatalf("Error getting all keys with coll <%s>", coll)
	}
	if len(actuaList) != len(expectedList) {
		t.Fatalf("Expected to read %d items, but read %d",
			len(expectedList),
			len(actuaList))
	}

	err = store.DeleteAll("artist")
	if err != nil {
		t.Fatal("Error deleting key, %v", err)
	}
	for i := int(0); i < 10; i++ {
		key = coll + CollKeySep + baseKey + strconv.Itoa(i)
		checkGetIsEmpty(store, key, t)
	}
}

func TestStorePersistPutAfterClose(t *testing.T) {
	store := setUp(t)

	key := "artist/the prodigy"
	expected := generateData(Data{"Beyond the peak"}, t)

	if err := store.Put(key, expected); err != nil {
		tearDown(store, t)
		t.Fatalf("Error putting data in, %v", err)
	}

	// Don't use tearDown as it deletes the storage after use
	if err := store.Close(); err != nil {
		t.Fatalf("Error closing store, %v", err)
	}

	otherStore := setUp(t)
	defer tearDown(otherStore, t)

	actual, err := otherStore.Get(key)
	if err != nil {
		t.Fatalf("Error getting data back, %v", err)
	}
	if !bytes.Equal(expected, actual) {
		t.Fatalf("Expected <%s> but was <%s>",
			expected,
			actual)
	}
}

func TestStorePersistDeleteAfterClose(t *testing.T) {
	store := setUp(t)

	key := "artist/the prodigy"
	expected := generateData(Data{"Beyond the peak"}, t)

	if err := store.Put(key, expected); err != nil {
		tearDown(store, t)
		t.Fatalf("Error putting value in, %v", err)
	}

	if err := store.Delete(key); err != nil {
		tearDown(store, t)
		t.Fatalf("Error deleting value, %v", err)
	}

	// Don't use tearDown as it deletes the storage after use
	if err := store.Close(); err != nil {
		t.Fatalf("Error closing store, %v", err)
	}

	otherStore := setUp(t)
	defer tearDown(otherStore, t)

	checkGetIsEmpty(otherStore, key, t)

}

// Error cases

func TestErrorWhenStorePointToNonDirectoryPath(t *testing.T) {
	filename := "test_regular_file"
	_, err := os.Create(filename)
	if err != nil {
		t.Fatalf("Error creating test file, %v", err)
	}
	defer os.Remove(filename)

	store, err := NewStore(filename)
	if _, isRightType := err.(PathError); !isRightType {
		defer tearDown(store, t)
		t.Errorf("Should have returned an error of type PathError")
	}

}

func TestErrorWhenStoreAlreadyUsingPath(t *testing.T) {
	path := "a_busy_path"
	store, err := NewStore(path)
	if err != nil {
		t.Fatalf("Error creating test store, %v", err)
	}
	if err := store.Load(); err != nil {
		t.Fatalf("Error loading test store, %v", err)
	}
	defer tearDown(store, t)

	another, err := NewStore(path)
	if err != nil {
		t.Fatalf("Error creating second test store, %v", err)
	}
	err = another.Load()
	if _, isRightType := err.(PathError); !isRightType {
		defer tearDown(another, t)
		t.Errorf("Should have returned an error of type PathError")
	}
}

func TestErrorWhenClosingStoreNotLoaded(t *testing.T) {
	path := "a_busy_path"
	store, err := NewStore(path)
	if err != nil {
		t.Fatalf("Error creating test store, %v", err)
	}
	err = store.Close()
	if _, isRightType := err.(StoreError); !isRightType {
		t.Errorf("Should have returned an error of type StoreError")
	}
}

///////////////////////////////////////////////////////////////////////////////
// Multiple goroutine
///////////////////////////////////////////////////////////////////////////////

func TestMultipleGoroutine(t *testing.T) {
	var kvCount int
	if testing.Short() {
		kvCount = 1000
	} else {
		kvCount = 10000
	}
	coll := "games"
	baseKey := "total annihilation #"

	type KeyValue struct {
		Key   string
		Value []byte
	}
	countGet := int64(0)
	countPut := int64(0)
	countDelete := int64(0)
	start := time.Now()

	store := setUp(t)
	defer tearDown(store, t)
	defer runtime.GOMAXPROCS(runtime.GOMAXPROCS(runtime.NumCPU()))

	var writeGroup sync.WaitGroup
	var readGroup sync.WaitGroup
	var deleteGroup sync.WaitGroup
	writeErr := make(chan error)
	readErr := make(chan error)
	deleteErr := make(chan error)

	expectedList := make([]KeyValue, 0)
	for i := int(0); i < kvCount; i++ {
		key := coll + CollKeySep + baseKey + strconv.Itoa(i)
		data := generateData(Data{"It's fun!"}, t)
		kv := KeyValue{key, data}
		expectedList = append(expectedList, kv)
	}

	for _, kv := range expectedList {
		// Put all the values concurrently, ensure there's no error
		writeGroup.Add(1)
		go func(pair KeyValue, cErr chan error) {
			defer writeGroup.Done()
			err := store.Put(pair.Key, pair.Value)
			atomic.AddInt64(&countPut, 1)
			if err != nil {
				t.Fatal("Received an error", err)
				cErr <- err
			}

		}(kv, writeErr)
	}

	writeGroup.Wait()

	if len(writeErr) != 0 {
		t.Fatalf("Failed to write values concurrently, got %d errors",
			len(writeErr))
	}

	for _, kv := range expectedList {
		// Get each value, ensure they're the same
		readGroup.Add(1)
		go func(kv KeyValue, cErr chan error) {
			defer readGroup.Done()

			expected := kv.Value
			actual, err := store.Get(kv.Key)
			atomic.AddInt64(&countGet, 1)
			if err != nil {
				cErr <- err
			} else if !bytes.Equal(expected, actual) {
				t.Errorf("Expected <%s> but was <%s>",
					expected,
					actual)
			}

		}(kv, readErr)

		// GetAll values, ensure this slice is as big as all the value
		// we've put
		readGroup.Add(1)
		go func(cErr chan error) {
			defer readGroup.Done()

			actual, err := store.GetAll(coll)
			atomic.AddInt64(&countGet, int64(len(actual)))
			if err != nil {
				cErr <- err
			} else if len(actual) != len(expectedList) {
				t.Errorf("Expected len(store.GetAll)=<%s> but was <%s>",
					len(expectedList),
					len(actual))
			}
		}(readErr)
	}

	readGroup.Wait()

	if len(readErr) != 0 {
		t.Fatalf("Failed to write values concurrently, got %d errors",
			len(readErr))
	}

	for _, kv := range expectedList {
		// Put all the values concurrently, ensure there's no error
		deleteGroup.Add(1)
		go func(pair KeyValue, cErr chan error) {
			defer deleteGroup.Done()
			err := store.Delete(pair.Key)
			atomic.AddInt64(&countDelete, 1)
			if err != nil {
				t.Fatal("Received an error", err)
				cErr <- err
			}

		}(kv, deleteErr)
	}

	deleteGroup.Wait()

	if len(deleteErr) != 0 {
		t.Fatalf("Failed to write values concurrently, got %d errors",
			len(deleteErr))
	}

	err := store.DeleteAll(coll)
	atomic.AddInt64(&countDelete, int64(len(expectedList)))
	if err != nil {
		t.Fatalf("Error deleting all", err)
	}
	for _, kv := range expectedList {
		atomic.AddInt64(&countGet, 1)
		checkGetIsEmpty(store, kv.Key, t)
	}
	duration := time.Now().Sub(start)
	fmt.Printf("Concurrent test\n")
	fmt.Printf("- Get operations=%d\n", countGet)
	fmt.Printf("- Put operations=%d\n", countPut)
	fmt.Printf("- Delete operations=%d\n", countPut)
	fmt.Printf("in %fs ", duration.Seconds())
	fmt.Printf("by %d cpus, using %d concurrent goroutines\n",
		runtime.NumCPU(), kvCount)

}
<|endoftext|>


<commit_before>package profiles

import (
	"log"
	"strconv"
	"sync"
	"sync/atomic"

	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

/*
	basic insert into capped collection only profile, insert a doc of 400 bytes
*/

type cappedCollInsertReadProfile struct {
	UID int64

	indexGroup bool

	initProfile sync.Once

	session *mgo.Session
}

var _cappedCollInsertReadProfile cappedCollInsertReadProfile

// func Int2ObjId(i int64) string {
// 	// return string represenation of UID
// }

func (i cappedCollInsertReadProfile) SendNext(s *mgo.Session, worker_id int) error {
	var err error
	c := s.DB(getDBName(default_db_name_prefix)).C(getCollectionName(default_col_name_prefix))

	// r := rands[worker_id].Int63n(3)

	if worker_id == 1 || worker_id == 2 {

		num := (_cappedCollInsertReadProfile.UID - 10000)

		c.Find(bson.M{"name": &num}).Limit(1)
	} else {
		_u := atomic.AddInt64(&_cappedCollInsertReadProfile.UID, 1) // to make this unique

		doc := bson.M{
			// "_id":      _u,
			"name":     _u,
			"group":    rands[worker_id].Int(),
			"payload":  &Payload,
			"payload1": &Payload1,
			"payload2": &Payload2,
			"payload3": &Payload3,
			"payload4": &Payload4,
			"payload5": &Payload5,
			"payload6": &Payload6}

		if _profile_use_legacy_write {
			err = c.Insert(doc)
		} else {
			var results interface{}

			err = c.Database.Run(bson.D{{"insert", c.Name},
				{"documents", []bson.M{doc}}}, results)
		}
	}

	panicOnError(err)
	return err
}

func initCappedInsertReadTest(session *mgo.Session, _initdb bool) {
	log.Println("Initialize simple DB. with initdb flag ", _initdb)

	ci := new(mgo.CollectionInfo)
	ci.Capped = true
	ci.MaxBytes = 1000000000

	log.Println(". Init DB, drop collections")
	coll := session.DB(default_db_name_prefix).C(default_col_name_prefix)
	coll.DropCollection()
	coll.Create(ci)
	ensureIndexKeysHelper(coll)

	for i := 1; i < _multi_db; i++ {
		for j := 1; j < _multi_col; j++ {
			collection := session.DB(default_db_name_prefix + strconv.Itoa(i)).C(default_col_name_prefix + strconv.Itoa(j))
			collection.Create(ci)
			ensureIndexKeysHelper(collection)
		}
	}

}

func ensureIndexKeysHelper(collection *mgo.Collection) {
	err := collection.EnsureIndexKey("name")
	if err != nil {
		panic(err)
	}

	err = collection.EnsureIndexKey("uid")
	if err != nil {
		panic(err)
	}
}

func (i cappedCollInsertReadProfile) SetupTest(s *mgo.Session, _initdb bool) error {
	i.session = s
	f := func() {
		initCappedInsertReadTest(s, _initdb)
	}

	_cappedCollInsertReadProfile.initProfile.Do(f)
	return nil
}

func (i cappedCollInsertReadProfile) CsvString(total_time float64) string {
	return ""
}

func (i cappedCollInsertReadProfile) CsvHeader() string {
	return ""
}

func init() {
	atomic.StoreInt64(&_cappedCollInsertReadProfile.UID, -1) // UID starts with 1

	registerProfile("CAPPED_COLL_INSERT_READ", func() Profile {
		return Profile(_cappedCollInsertReadProfile) // use the same instance
	})
}
<commit_msg>fix capped col creation issue with capped_insert_read<commit_after>package profiles

import (
	"log"
	"strconv"
	"sync"
	"sync/atomic"

	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

/*
	basic insert into capped collection only profile, insert a doc of 400 bytes
*/

type cappedCollInsertReadProfile struct {
	UID int64

	indexGroup bool

	initProfile sync.Once

	session *mgo.Session
}

var _cappedCollInsertReadProfile cappedCollInsertReadProfile

// func Int2ObjId(i int64) string {
// 	// return string represenation of UID
// }

func (i cappedCollInsertReadProfile) SendNext(s *mgo.Session, worker_id int) error {
	var err error
	c := s.DB(getDBName(default_db_name_prefix)).C(getCollectionName(default_col_name_prefix))

	// r := rands[worker_id].Int63n(3)

	if worker_id == 1 || worker_id == 2 {

		num := (_cappedCollInsertReadProfile.UID - 10000)

		c.Find(bson.M{"name": &num}).Limit(1)
	} else {
		_u := atomic.AddInt64(&_cappedCollInsertReadProfile.UID, 1) // to make this unique

		doc := bson.M{
			// "_id":      _u,
			"name":     _u,
			"group":    rands[worker_id].Int(),
			"payload":  &Payload,
			"payload1": &Payload1,
			"payload2": &Payload2,
			"payload3": &Payload3,
			"payload4": &Payload4,
			"payload5": &Payload5,
			"payload6": &Payload6}

		if _profile_use_legacy_write {
			err = c.Insert(doc)
		} else {
			var results interface{}

			err = c.Database.Run(bson.D{{"insert", c.Name},
				{"documents", []bson.M{doc}}}, results)
		}
	}

	panicOnError(err)
	return err
}

func initCappedInsertReadTest(session *mgo.Session, _initdb bool) {
	log.Println("Initialize simple DB. with initdb flag ", _initdb)

	ci := new(mgo.CollectionInfo)
	ci.Capped = true
	ci.MaxBytes = 1000000000

	log.Println(". Init DB, drop collections")
	coll := session.DB(default_db_name_prefix).C(default_col_name_prefix)
	coll.DropCollection()
	coll.Create(ci)
	ensureIndexKeysHelper(coll)

	for i := 1; i <= _multi_db; i++ {
		for j := 1; j <= _multi_col; j++ {
			collection := session.DB(default_db_name_prefix + strconv.Itoa(i)).C(default_col_name_prefix + strconv.Itoa(j))
			collection.DropCollection()
			collection.Create(ci)
			ensureIndexKeysHelper(collection)
		}
	}

}

func ensureIndexKeysHelper(collection *mgo.Collection) {
	err := collection.EnsureIndexKey("name")
	if err != nil {
		panic(err)
	}

	err = collection.EnsureIndexKey("uid")
	if err != nil {
		panic(err)
	}
}

func (i cappedCollInsertReadProfile) SetupTest(s *mgo.Session, _initdb bool) error {
	i.session = s
	f := func() {
		initCappedInsertReadTest(s, _initdb)
	}

	_cappedCollInsertReadProfile.initProfile.Do(f)
	return nil
}

func (i cappedCollInsertReadProfile) CsvString(total_time float64) string {
	return ""
}

func (i cappedCollInsertReadProfile) CsvHeader() string {
	return ""
}

func init() {
	atomic.StoreInt64(&_cappedCollInsertReadProfile.UID, -1) // UID starts with 1

	registerProfile("CAPPED_COLL_INSERT_READ", func() Profile {
		return Profile(_cappedCollInsertReadProfile) // use the same instance
	})
}
<|endoftext|>


<commit_before>// Use of this source code is governed by a MIT license that can be found in the LICENSE file.
// Giorgis (habtom@giorgis.io)
// Corbatto (luca@corbatto.de)

package avcodec

//#cgo pkg-config: libavcodec
//#include <libavcodec/avcodec.h>
import "C"
import (
	"unsafe"

	"gopkg.in/targodan/ffgopeg.v1/avutil"
)

// NewCodecContext allocates a Context and set its fields to default values.
//
// C-Function: avcodec_alloc_context3
func NewCodecContext(c *Codec) *CodecContext {
	return (*CodecContext)(C.avcodec_alloc_context3((*C.struct_AVCodec)(c)))
}

// FindDecoder finds a registered decoder with a matching codec ID.
//
// C-Function: avcodec_find_decoder
func FindDecoder(id CodecId) *Codec {
	return (*Codec)(C.avcodec_find_decoder((C.enum_AVCodecID)(id)))
}

// FindDecoderByName finds a registered decoder with the specified name.
//
// C-Function: avcodec_find_decoder_by_name
func FindDecoderByName(n string) *Codec {
	return (*Codec)(C.avcodec_find_decoder_by_name(C.CString(n)))
}

// FindEncoder finds a registered encoder with a matching codec ID.
//
// C-Function: avcodec_find_encoder
func FindEncoder(id CodecId) *Codec {
	return (*Codec)(C.avcodec_find_encoder((C.enum_AVCodecID)(id)))
}

// FindEncoderByName finds a registered encoder with the specified name.
//
// C-Function: avcodec_find_encoder_by_name
func FindEncoderByName(c string) *Codec {
	return (*Codec)(C.avcodec_find_encoder_by_name(C.CString(c)))
}

// MaxLowres returns the maximum lowres supported by the decoder.
//
// C-Function: av_codec_get_max_lowres
func (c *Codec) MaxLowres() int {
	return int(C.av_codec_get_max_lowres((*C.struct_AVCodec)(c)))
}

// ProfileName returns a name for the specified profile, if available.
//
// C-Function: av_get_profile_name
func (c *Codec) ProfileName(p int) string {
	return C.GoString(C.av_get_profile_name((*C.struct_AVCodec)(c), C.int(p)))
}

// IsEncoder returns true if the codec is an encoder.
//
// C-Function: av_codec_is_encoder
func (c *Codec) IsEncoder() bool {
	return int(C.av_codec_is_encoder((*C.struct_AVCodec)(c))) != 0
}

// IsDecoder returns true if the codec is an decoder.
//
// C-Function: av_codec_is_decoder
func (c *Codec) IsDecoder() bool {
	return int(C.av_codec_is_decoder((*C.struct_AVCodec)(c))) != 0
}

// Name returns the name of the codec.
//
// C-Variable: AVCodec::name
func (c *Codec) Name() string {
	return C.GoString(c.name)
}

// LongName returns the long, descriptive name of the codec.
//
// C-Variable: AVCodec::long_name
func (c *Codec) LongName() string {
	return C.GoString(c.long_name)
}

// SampleFmts returns the suppoted sample formats.
//
// C-Variable: AVCodec::sample_fmts
func (c *Codec) SampleFmts() []avutil.SampleFormat {
	arr := (*[1 << 30]avutil.SampleFormat)(unsafe.Pointer(c.sample_fmts))
	var len int
	for len = 0; arr[len] != -1; len++ {
	}
	return arr[:len:len]
}
<commit_msg>Add Codec::PixFmts<commit_after>// Use of this source code is governed by a MIT license that can be found in the LICENSE file.
// Giorgis (habtom@giorgis.io)
// Corbatto (luca@corbatto.de)

package avcodec

//#cgo pkg-config: libavcodec
//#include <libavcodec/avcodec.h>
import "C"
import (
	"unsafe"

	"gopkg.in/targodan/ffgopeg.v1/avutil"
)

// NewCodecContext allocates a Context and set its fields to default values.
//
// C-Function: avcodec_alloc_context3
func NewCodecContext(c *Codec) *CodecContext {
	return (*CodecContext)(C.avcodec_alloc_context3((*C.struct_AVCodec)(c)))
}

// FindDecoder finds a registered decoder with a matching codec ID.
//
// C-Function: avcodec_find_decoder
func FindDecoder(id CodecId) *Codec {
	return (*Codec)(C.avcodec_find_decoder((C.enum_AVCodecID)(id)))
}

// FindDecoderByName finds a registered decoder with the specified name.
//
// C-Function: avcodec_find_decoder_by_name
func FindDecoderByName(n string) *Codec {
	return (*Codec)(C.avcodec_find_decoder_by_name(C.CString(n)))
}

// FindEncoder finds a registered encoder with a matching codec ID.
//
// C-Function: avcodec_find_encoder
func FindEncoder(id CodecId) *Codec {
	return (*Codec)(C.avcodec_find_encoder((C.enum_AVCodecID)(id)))
}

// FindEncoderByName finds a registered encoder with the specified name.
//
// C-Function: avcodec_find_encoder_by_name
func FindEncoderByName(c string) *Codec {
	return (*Codec)(C.avcodec_find_encoder_by_name(C.CString(c)))
}

// MaxLowres returns the maximum lowres supported by the decoder.
//
// C-Function: av_codec_get_max_lowres
func (c *Codec) MaxLowres() int {
	return int(C.av_codec_get_max_lowres((*C.struct_AVCodec)(c)))
}

// ProfileName returns a name for the specified profile, if available.
//
// C-Function: av_get_profile_name
func (c *Codec) ProfileName(p int) string {
	return C.GoString(C.av_get_profile_name((*C.struct_AVCodec)(c), C.int(p)))
}

// IsEncoder returns true if the codec is an encoder.
//
// C-Function: av_codec_is_encoder
func (c *Codec) IsEncoder() bool {
	return int(C.av_codec_is_encoder((*C.struct_AVCodec)(c))) != 0
}

// IsDecoder returns true if the codec is an decoder.
//
// C-Function: av_codec_is_decoder
func (c *Codec) IsDecoder() bool {
	return int(C.av_codec_is_decoder((*C.struct_AVCodec)(c))) != 0
}

// Name returns the name of the codec.
//
// C-Variable: AVCodec::name
func (c *Codec) Name() string {
	return C.GoString(c.name)
}

// LongName returns the long, descriptive name of the codec.
//
// C-Variable: AVCodec::long_name
func (c *Codec) LongName() string {
	return C.GoString(c.long_name)
}

// PixFmts returns an array of support pixel format, see AV_PIX_FMT_xxx.
//
// C-Field: AVCodec::pix_fmts
func (c *Codec) PixFmts() []avutil.PixelFormat {
	arr := (*[1 << 30]avutil.PixelFormat)(unsafe.Pointer(c.pix_fmts))
	var len int
	for len = 0; arr[len] != -1; len++ {
	}
	return arr[:len:len]
}

// SampleFmts returns the suppoted sample formats.
//
// C-Variable: AVCodec::sample_fmts
func (c *Codec) SampleFmts() []avutil.SampleFormat {
	arr := (*[1 << 30]avutil.SampleFormat)(unsafe.Pointer(c.sample_fmts))
	var len int
	for len = 0; arr[len] != -1; len++ {
	}
	return arr[:len:len]
}
<|endoftext|>


<commit_before>package db

import (
	"testing"
	"time"
)

func TestDbConn(t *testing.T) {
	conn := func(name string) {
		t.Log("\t", name)

		if c1, err1 := Client(); err1 != nil {
			t.Error(err1)
		} else {
			defer Release(c1)
		}
		time.Sleep(2 * time.Second)
	}

	go conn("c1")
	conn("c2")
	conn("c3")
}
<commit_msg>Add comment for db test<commit_after>/**
 * go test -f ../../config.json
 */

package db

import (
	"testing"
	"time"
)

func TestDbConn(t *testing.T) {
	conn := func(name string) {
		t.Log("\t", name)

		if c1, err1 := Client(); err1 != nil {
			t.Error(err1)
		} else {
			defer Release(c1)
		}
		time.Sleep(2 * time.Second)
	}

	go conn("c1")
	conn("c2")
	conn("c3")
}
<|endoftext|>


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

package deep // import "github.com/influxdata/influxdb/pkg/deep"

import (
	"fmt"
	"math"
	"reflect"
)

// Equal is a copy of reflect.DeepEqual except that it treats NaN == NaN as true.
func Equal(a1, a2 interface{}) bool {
	if a1 == nil || a2 == nil {
		return a1 == a2
	}
	v1 := reflect.ValueOf(a1)
	v2 := reflect.ValueOf(a2)
	if v1.Type() != v2.Type() {
		return false
	}
	return deepValueEqual(v1, v2, make(map[visit]bool), 0)
}

// Tests for deep equality using reflected types. The map argument tracks
// comparisons that have already been seen, which allows short circuiting on
// recursive types.
func deepValueEqual(v1, v2 reflect.Value, visited map[visit]bool, depth int) bool {
	if !v1.IsValid() || !v2.IsValid() {
		return v1.IsValid() == v2.IsValid()
	}
	if v1.Type() != v2.Type() {
		return false
	}

	// if depth > 10 { panic("deepValueEqual") }	// for debugging
	hard := func(k reflect.Kind) bool {
		switch k {
		case reflect.Array, reflect.Map, reflect.Slice, reflect.Struct:
			return true
		}
		return false
	}

	if v1.CanAddr() && v2.CanAddr() && hard(v1.Kind()) {
		addr1 := v1.UnsafeAddr()
		addr2 := v2.UnsafeAddr()
		if addr1 > addr2 {
			// Canonicalize order to reduce number of entries in visited.
			addr1, addr2 = addr2, addr1
		}

		// Short circuit if references are identical ...
		if addr1 == addr2 {
			return true
		}

		// ... or already seen
		typ := v1.Type()
		v := visit{addr1, addr2, typ}
		if visited[v] {
			return true
		}

		// Remember for later.
		visited[v] = true
	}

	switch v1.Kind() {
	case reflect.Array:
		for i := 0; i < v1.Len(); i++ {
			if !deepValueEqual(v1.Index(i), v2.Index(i), visited, depth+1) {
				return false
			}
		}
		return true
	case reflect.Slice:
		if v1.IsNil() != v2.IsNil() {
			return false
		}
		if v1.Len() != v2.Len() {
			return false
		}
		if v1.Pointer() == v2.Pointer() {
			return true
		}
		for i := 0; i < v1.Len(); i++ {
			if !deepValueEqual(v1.Index(i), v2.Index(i), visited, depth+1) {
				return false
			}
		}
		return true
	case reflect.Interface:
		if v1.IsNil() || v2.IsNil() {
			return v1.IsNil() == v2.IsNil()
		}
		return deepValueEqual(v1.Elem(), v2.Elem(), visited, depth+1)
	case reflect.Ptr:
		return deepValueEqual(v1.Elem(), v2.Elem(), visited, depth+1)
	case reflect.Struct:
		for i, n := 0, v1.NumField(); i < n; i++ {
			if !deepValueEqual(v1.Field(i), v2.Field(i), visited, depth+1) {
				return false
			}
		}
		return true
	case reflect.Map:
		if v1.IsNil() != v2.IsNil() {
			return false
		}
		if v1.Len() != v2.Len() {
			return false
		}
		if v1.Pointer() == v2.Pointer() {
			return true
		}
		for _, k := range v1.MapKeys() {
			if !deepValueEqual(v1.MapIndex(k), v2.MapIndex(k), visited, depth+1) {
				return false
			}
		}
		return true
	case reflect.Func:
		if v1.IsNil() && v2.IsNil() {
			return true
		}
		// Can't do better than this:
		return false
	case reflect.String:
		return v1.String() == v2.String()
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		return v1.Int() == v2.Int()
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		return v1.Uint() == v2.Uint()
	case reflect.Float32, reflect.Float64:
		// Special handling for floats so that NaN == NaN is true.
		f1, f2 := v1.Float(), v2.Float()
		if math.IsNaN(f1) && math.IsNaN(f2) {
			return true
		}
		return f1 == f2
	case reflect.Bool:
		return v1.Bool() == v2.Bool()
	default:
		panic(fmt.Sprintf("cannot compare type: %s", v1.Kind().String()))
	}
}

// During deepValueEqual, must keep track of checks that are
// in progress.  The comparison algorithm assumes that all
// checks in progress are true when it reencounters them.
// Visited comparisons are stored in a map indexed by visit.
type visit struct {
	a1  uintptr
	a2  uintptr
	typ reflect.Type
}
<commit_msg>include missing license for equal.go<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// License.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
// * Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
// * Redistributions in binary form must reproduce the above
// copyright notice, this list of conditions and the following disclaimer
// in the documentation and/or other materials provided with the
// distribution.
// * Neither the name of Google Inc. nor the names of its
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

package deep // import "github.com/influxdata/influxdb/pkg/deep"

import (
	"fmt"
	"math"
	"reflect"
)

// Equal is a copy of reflect.DeepEqual except that it treats NaN == NaN as true.
func Equal(a1, a2 interface{}) bool {
	if a1 == nil || a2 == nil {
		return a1 == a2
	}
	v1 := reflect.ValueOf(a1)
	v2 := reflect.ValueOf(a2)
	if v1.Type() != v2.Type() {
		return false
	}
	return deepValueEqual(v1, v2, make(map[visit]bool), 0)
}

// Tests for deep equality using reflected types. The map argument tracks
// comparisons that have already been seen, which allows short circuiting on
// recursive types.
func deepValueEqual(v1, v2 reflect.Value, visited map[visit]bool, depth int) bool {
	if !v1.IsValid() || !v2.IsValid() {
		return v1.IsValid() == v2.IsValid()
	}
	if v1.Type() != v2.Type() {
		return false
	}

	// if depth > 10 { panic("deepValueEqual") }	// for debugging
	hard := func(k reflect.Kind) bool {
		switch k {
		case reflect.Array, reflect.Map, reflect.Slice, reflect.Struct:
			return true
		}
		return false
	}

	if v1.CanAddr() && v2.CanAddr() && hard(v1.Kind()) {
		addr1 := v1.UnsafeAddr()
		addr2 := v2.UnsafeAddr()
		if addr1 > addr2 {
			// Canonicalize order to reduce number of entries in visited.
			addr1, addr2 = addr2, addr1
		}

		// Short circuit if references are identical ...
		if addr1 == addr2 {
			return true
		}

		// ... or already seen
		typ := v1.Type()
		v := visit{addr1, addr2, typ}
		if visited[v] {
			return true
		}

		// Remember for later.
		visited[v] = true
	}

	switch v1.Kind() {
	case reflect.Array:
		for i := 0; i < v1.Len(); i++ {
			if !deepValueEqual(v1.Index(i), v2.Index(i), visited, depth+1) {
				return false
			}
		}
		return true
	case reflect.Slice:
		if v1.IsNil() != v2.IsNil() {
			return false
		}
		if v1.Len() != v2.Len() {
			return false
		}
		if v1.Pointer() == v2.Pointer() {
			return true
		}
		for i := 0; i < v1.Len(); i++ {
			if !deepValueEqual(v1.Index(i), v2.Index(i), visited, depth+1) {
				return false
			}
		}
		return true
	case reflect.Interface:
		if v1.IsNil() || v2.IsNil() {
			return v1.IsNil() == v2.IsNil()
		}
		return deepValueEqual(v1.Elem(), v2.Elem(), visited, depth+1)
	case reflect.Ptr:
		return deepValueEqual(v1.Elem(), v2.Elem(), visited, depth+1)
	case reflect.Struct:
		for i, n := 0, v1.NumField(); i < n; i++ {
			if !deepValueEqual(v1.Field(i), v2.Field(i), visited, depth+1) {
				return false
			}
		}
		return true
	case reflect.Map:
		if v1.IsNil() != v2.IsNil() {
			return false
		}
		if v1.Len() != v2.Len() {
			return false
		}
		if v1.Pointer() == v2.Pointer() {
			return true
		}
		for _, k := range v1.MapKeys() {
			if !deepValueEqual(v1.MapIndex(k), v2.MapIndex(k), visited, depth+1) {
				return false
			}
		}
		return true
	case reflect.Func:
		if v1.IsNil() && v2.IsNil() {
			return true
		}
		// Can't do better than this:
		return false
	case reflect.String:
		return v1.String() == v2.String()
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		return v1.Int() == v2.Int()
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		return v1.Uint() == v2.Uint()
	case reflect.Float32, reflect.Float64:
		// Special handling for floats so that NaN == NaN is true.
		f1, f2 := v1.Float(), v2.Float()
		if math.IsNaN(f1) && math.IsNaN(f2) {
			return true
		}
		return f1 == f2
	case reflect.Bool:
		return v1.Bool() == v2.Bool()
	default:
		panic(fmt.Sprintf("cannot compare type: %s", v1.Kind().String()))
	}
}

// During deepValueEqual, must keep track of checks that are
// in progress.  The comparison algorithm assumes that all
// checks in progress are true when it reencounters them.
// Visited comparisons are stored in a map indexed by visit.
type visit struct {
	a1  uintptr
	a2  uintptr
	typ reflect.Type
}
<|endoftext|>


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

// Package flags implements command-line flag parsing.
package flags

import (
	"flag"
	"fmt"
	"net/url"
	"os"
	"strings"

	"github.com/coreos/pkg/capnslog"
	"github.com/spf13/pflag"
)

var (
	plog = capnslog.NewPackageLogger("github.com/coreos/etcd", "pkg/flags")
)

// DeprecatedFlag encapsulates a flag that may have been previously valid but
// is now deprecated. If a DeprecatedFlag is set, an error occurs.
type DeprecatedFlag struct {
	Name string
}

func (f *DeprecatedFlag) Set(_ string) error {
	return fmt.Errorf(`flag "-%s" is no longer supported.`, f.Name)
}

func (f *DeprecatedFlag) String() string {
	return ""
}

// IgnoredFlag encapsulates a flag that may have been previously valid but is
// now ignored. If an IgnoredFlag is set, a warning is printed and
// operation continues.
type IgnoredFlag struct {
	Name string
}

// IsBoolFlag is defined to allow the flag to be defined without an argument
func (f *IgnoredFlag) IsBoolFlag() bool {
	return true
}

func (f *IgnoredFlag) Set(s string) error {
	plog.Warningf(`flag "-%s" is no longer supported - ignoring.`, f.Name)
	return nil
}

func (f *IgnoredFlag) String() string {
	return ""
}

// SetFlagsFromEnv parses all registered flags in the given flagset,
// and if they are not already set it attempts to set their values from
// environment variables. Environment variables take the name of the flag but
// are UPPERCASE, have the given prefix  and any dashes are replaced by
// underscores - for example: some-flag => ETCD_SOME_FLAG
func SetFlagsFromEnv(prefix string, fs *flag.FlagSet) error {
	var err error
	alreadySet := make(map[string]bool)
	fs.Visit(func(f *flag.Flag) {
		alreadySet[FlagToEnv(prefix, f.Name)] = true
	})
	usedEnvKey := make(map[string]bool)
	fs.VisitAll(func(f *flag.Flag) {
		err = setFlagFromEnv(fs, prefix, f.Name, usedEnvKey, alreadySet, true)
	})

	verifyEnv(prefix, usedEnvKey, alreadySet)

	return err
}

// SetPflagsFromEnv is similar to SetFlagsFromEnv. However, the accepted flagset type is pflag.FlagSet
// and it does not do any logging.
func SetPflagsFromEnv(prefix string, fs *pflag.FlagSet) error {
	var err error
	alreadySet := make(map[string]bool)
	usedEnvKey := make(map[string]bool)
	fs.VisitAll(func(f *pflag.Flag) {
		if f.Changed {
			alreadySet[FlagToEnv(prefix, f.Name)] = true
		}
		if serr := setFlagFromEnv(fs, prefix, f.Name, usedEnvKey, alreadySet, false); serr != nil {
			err = serr
		}
	})
	return err
}

// FlagToEnv converts flag string to upper-case environment variable key string.
func FlagToEnv(prefix, name string) string {
	return prefix + "_" + strings.ToUpper(strings.Replace(name, "-", "_", -1))
}

func verifyEnv(prefix string, usedEnvKey, alreadySet map[string]bool) {
	for _, env := range os.Environ() {
		kv := strings.SplitN(env, "=", 2)
		if len(kv) != 2 {
			plog.Warningf("found invalid env %s", env)
		}
		if usedEnvKey[kv[0]] {
			continue
		}
		if alreadySet[kv[0]] {
			plog.Infof("recognized environment variable %s, but unused: shadowed by corresponding flag ", kv[0])
			continue
		}
		if strings.HasPrefix(env, prefix+"_") {
			plog.Warningf("unrecognized environment variable %s", env)
		}
	}
}

type flagSetter interface {
	Set(fk string, fv string) error
}

func setFlagFromEnv(fs flagSetter, prefix, fname string, usedEnvKey, alreadySet map[string]bool, log bool) error {
	key := FlagToEnv(prefix, fname)
	if !alreadySet[key] {
		val := os.Getenv(key)
		if val != "" {
			usedEnvKey[key] = true
			if serr := fs.Set(fname, val); serr != nil {
				return fmt.Errorf("invalid value %q for %s: %v", val, key, serr)
			}
			if log {
				plog.Infof("recognized and used environment variable %s=%s", key, val)
			}
		}
	}
	return nil
}

// URLsFromFlag returns a slices from url got from the flag.
func URLsFromFlag(fs *flag.FlagSet, urlsFlagName string) []url.URL {
	return []url.URL(*fs.Lookup(urlsFlagName).Value.(*URLsValue))
}

func IsSet(fs *flag.FlagSet, name string) bool {
	set := false
	fs.Visit(func(f *flag.Flag) {
		if f.Name == name {
			set = true
		}
	})
	return set
}
<commit_msg>pkg/flags: warns on shadowed environment variable flags<commit_after>// Copyright 2015 The etcd Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package flags implements command-line flag parsing.
package flags

import (
	"flag"
	"fmt"
	"net/url"
	"os"
	"strings"

	"github.com/coreos/pkg/capnslog"
	"github.com/spf13/pflag"
)

var (
	plog = capnslog.NewPackageLogger("github.com/coreos/etcd", "pkg/flags")
)

// DeprecatedFlag encapsulates a flag that may have been previously valid but
// is now deprecated. If a DeprecatedFlag is set, an error occurs.
type DeprecatedFlag struct {
	Name string
}

func (f *DeprecatedFlag) Set(_ string) error {
	return fmt.Errorf(`flag "-%s" is no longer supported.`, f.Name)
}

func (f *DeprecatedFlag) String() string {
	return ""
}

// IgnoredFlag encapsulates a flag that may have been previously valid but is
// now ignored. If an IgnoredFlag is set, a warning is printed and
// operation continues.
type IgnoredFlag struct {
	Name string
}

// IsBoolFlag is defined to allow the flag to be defined without an argument
func (f *IgnoredFlag) IsBoolFlag() bool {
	return true
}

func (f *IgnoredFlag) Set(s string) error {
	plog.Warningf(`flag "-%s" is no longer supported - ignoring.`, f.Name)
	return nil
}

func (f *IgnoredFlag) String() string {
	return ""
}

// SetFlagsFromEnv parses all registered flags in the given flagset,
// and if they are not already set it attempts to set their values from
// environment variables. Environment variables take the name of the flag but
// are UPPERCASE, have the given prefix  and any dashes are replaced by
// underscores - for example: some-flag => ETCD_SOME_FLAG
func SetFlagsFromEnv(prefix string, fs *flag.FlagSet) error {
	var err error
	alreadySet := make(map[string]bool)
	fs.Visit(func(f *flag.Flag) {
		alreadySet[FlagToEnv(prefix, f.Name)] = true
	})
	usedEnvKey := make(map[string]bool)
	fs.VisitAll(func(f *flag.Flag) {
		err = setFlagFromEnv(fs, prefix, f.Name, usedEnvKey, alreadySet, true)
	})
	verifyEnv(prefix, usedEnvKey, alreadySet)
	return err
}

// SetPflagsFromEnv is similar to SetFlagsFromEnv. However, the accepted flagset type is pflag.FlagSet
// and it does not do any logging.
func SetPflagsFromEnv(prefix string, fs *pflag.FlagSet) error {
	var err error
	alreadySet := make(map[string]bool)
	usedEnvKey := make(map[string]bool)
	fs.VisitAll(func(f *pflag.Flag) {
		if f.Changed {
			alreadySet[FlagToEnv(prefix, f.Name)] = true
		}
		if serr := setFlagFromEnv(fs, prefix, f.Name, usedEnvKey, alreadySet, false); serr != nil {
			err = serr
		}
	})
	verifyEnv(prefix, usedEnvKey, alreadySet)
	return err
}

// FlagToEnv converts flag string to upper-case environment variable key string.
func FlagToEnv(prefix, name string) string {
	return prefix + "_" + strings.ToUpper(strings.Replace(name, "-", "_", -1))
}

func verifyEnv(prefix string, usedEnvKey, alreadySet map[string]bool) {
	for _, env := range os.Environ() {
		kv := strings.SplitN(env, "=", 2)
		if len(kv) != 2 {
			plog.Warningf("found invalid env %s", env)
		}
		if usedEnvKey[kv[0]] {
			continue
		}
		if alreadySet[kv[0]] {
			// TODO: exit with error in v3.4
			plog.Warningf("recognized environment variable %s, but unused: shadowed by corresponding flag", kv[0])
			continue
		}
		if strings.HasPrefix(env, prefix+"_") {
			plog.Warningf("unrecognized environment variable %s", env)
		}
	}
}

type flagSetter interface {
	Set(fk string, fv string) error
}

func setFlagFromEnv(fs flagSetter, prefix, fname string, usedEnvKey, alreadySet map[string]bool, log bool) error {
	key := FlagToEnv(prefix, fname)
	if !alreadySet[key] {
		val := os.Getenv(key)
		if val != "" {
			usedEnvKey[key] = true
			if serr := fs.Set(fname, val); serr != nil {
				return fmt.Errorf("invalid value %q for %s: %v", val, key, serr)
			}
			if log {
				plog.Infof("recognized and used environment variable %s=%s", key, val)
			}
		}
	}
	return nil
}

// URLsFromFlag returns a slices from url got from the flag.
func URLsFromFlag(fs *flag.FlagSet, urlsFlagName string) []url.URL {
	return []url.URL(*fs.Lookup(urlsFlagName).Value.(*URLsValue))
}

func IsSet(fs *flag.FlagSet, name string) bool {
	set := false
	fs.Visit(func(f *flag.Flag) {
		if f.Name == name {
			set = true
		}
	})
	return set
}
<|endoftext|>