<commit_before>/*
Copyright 2016 The Kubernetes 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 instancegroups

import (
	"bufio"
	"fmt"
	"os"
	"strings"
	"time"

	"github.com/golang/glog"
	api "k8s.io/kops/pkg/apis/kops"
	"k8s.io/kops/pkg/cloudinstances"
	"k8s.io/kops/pkg/featureflag"
	"k8s.io/kops/pkg/validation"
	"k8s.io/kops/upup/pkg/fi"
	"k8s.io/kubernetes/pkg/kubectl/cmd"
	cmdutil "k8s.io/kubernetes/pkg/kubectl/cmd/util"
)

// RollingUpdateInstanceGroup is the AWS ASG backing an InstanceGroup.
type RollingUpdateInstanceGroup struct {
	// Cloud is the kops cloud provider
	Cloud fi.Cloud
	// CloudGroup is the kops cloud provider groups
	CloudGroup *cloudinstances.CloudInstanceGroup

	// TODO should remove the need to have rollingupdate struct and add:
	// TODO - the kubernetes client
	// TODO - the cluster name
	// TODO - the client config
	// TODO - fail on validate
	// TODO - fail on drain
	// TODO - cloudonly
}

// NewRollingUpdateInstanceGroup creates a new struct
func NewRollingUpdateInstanceGroup(cloud fi.Cloud, cloudGroup *cloudinstances.CloudInstanceGroup) (*RollingUpdateInstanceGroup, error) {
	if cloud == nil {
		return nil, fmt.Errorf("cloud provider is required")
	}
	if cloudGroup == nil {
		return nil, fmt.Errorf("cloud group is required")
	}

	// TODO check more values in cloudGroup that they are set properly

	return &RollingUpdateInstanceGroup{
		Cloud:      cloud,
		CloudGroup: cloudGroup,
	}, nil
}

// promptInteractive asks the user to continue, mostly copied from vendor/google.golang.org/api/examples/gmail.go.
func promptInteractive(upgradedHost string) (stopPrompting bool, err error) {
	stopPrompting = false
	scanner := bufio.NewScanner(os.Stdin)
	glog.Infof("Pausing after finished %q", upgradedHost)
	fmt.Print("Continue? (Y)es, (N)o, (A)lwaysYes: [Y] ")
	scanner.Scan()
	err = scanner.Err()
	if err != nil {
		glog.Infof("unable to interpret input: %v", err)
		return stopPrompting, err
	}
	val := scanner.Text()
	val = strings.TrimSpace(val)
	val = strings.ToLower(val)
	switch val {
	case "n":
		glog.Infof("User signaled to stop")
		os.Exit(3)
	case "a":
		glog.Infof("Always Yes, stop prompting for rest of hosts")
		stopPrompting = true
	}
	return stopPrompting, err
}

// TODO: Temporarily increase size of ASG?
// TODO: Remove from ASG first so status is immediately updated?
// TODO: Batch termination, like a rolling-update

// RollingUpdate performs a rolling update on a list of ec2 instances.
func (r *RollingUpdateInstanceGroup) RollingUpdate(rollingUpdateData *RollingUpdateCluster, cluster *api.Cluster, instanceGroupList *api.InstanceGroupList, isBastion bool, sleepAfterTerminate time.Duration, validationTimeout time.Duration) (err error) {

	// we should not get here, but hey I am going to check.
	if rollingUpdateData == nil {
		return fmt.Errorf("rollingUpdate cannot be nil")
	}

	// Do not need a k8s client if you are doing cloudonly.
	if rollingUpdateData.K8sClient == nil && !rollingUpdateData.CloudOnly {
		return fmt.Errorf("rollingUpdate is missing a k8s client")
	}

	if instanceGroupList == nil {
		return fmt.Errorf("rollingUpdate is missing the InstanceGroupList")
	}

	update := r.CloudGroup.NeedUpdate
	if rollingUpdateData.Force {
		update = append(update, r.CloudGroup.Ready...)
	}

	if len(update) == 0 {
		return nil
	}

	if isBastion {
		glog.V(3).Info("Not validating the cluster as instance is a bastion.")
	} else if rollingUpdateData.CloudOnly {
		glog.V(3).Info("Not validating cluster as validation is turned off via the cloud-only flag.")
	} else if featureflag.DrainAndValidateRollingUpdate.Enabled() {
		if err = r.ValidateCluster(rollingUpdateData, cluster, instanceGroupList); err != nil {
			if rollingUpdateData.FailOnValidate {
				return fmt.Errorf("error validating cluster: %v", err)
			} else {
				glog.V(2).Infof("Ignoring cluster validation error: %v", err)
				glog.Infof("Cluster validation failed, but proceeding since fail-on-validate-error is set to false")
			}
		}
	}

	for _, u := range update {
		instanceId := u.ID

		nodeName := ""
		if u.Node != nil {
			nodeName = u.Node.Name
		}

		if isBastion {
			// We don't want to validate for bastions - they aren't part of the cluster
		} else if rollingUpdateData.CloudOnly {

			glog.Warningf("Not draining cluster nodes as 'cloudonly' flag is set.")

		} else if featureflag.DrainAndValidateRollingUpdate.Enabled() {

			if u.Node != nil {
				glog.Infof("Draining the node: %q.", nodeName)

				if err = r.DrainNode(u, rollingUpdateData); err != nil {
					if rollingUpdateData.FailOnDrainError {
						return fmt.Errorf("failed to drain node %q: %v", nodeName, err)
					} else {
						glog.Infof("Ignoring error draining node %q: %v", nodeName, err)
					}
				}
			} else {
				glog.Warningf("Skipping drain of instance %q, because it is not registered in kubernetes", instanceId)
			}
		}

		if err = r.DeleteInstance(u); err != nil {
			glog.Errorf("Error deleting aws instance %q, node %q: %v", instanceId, nodeName, err)
			return err
		}

		// Wait for the minimum interval
		time.Sleep(sleepAfterTerminate)

		if isBastion {
			glog.Infof("Deleted a bastion instance, %s, and continuing with rolling-update.", instanceId)

			continue
		} else if rollingUpdateData.CloudOnly {
			glog.Warningf("Not validating cluster as cloudonly flag is set.")
			continue

		} else if featureflag.DrainAndValidateRollingUpdate.Enabled() {
			glog.Infof("Validating the cluster.")

			if err = r.ValidateClusterWithDuration(rollingUpdateData, cluster, instanceGroupList, validationTimeout); err != nil {

				if rollingUpdateData.FailOnValidate {
					glog.Errorf("Cluster did not validate within %s", validationTimeout)
					return fmt.Errorf("error validating cluster after removing a node: %v", err)
				}

				glog.Warningf("Cluster validation failed after removing instance, proceeding since fail-on-validate is set to false: %v", err)
			}
			if rollingUpdateData.Interactive {
				stopPrompting, err := promptInteractive(nodeName)
				if err != nil {
					return err
				}
				if stopPrompting {
					// Is a pointer to a struct, changes here push back into the original
					rollingUpdateData.Interactive = false
				}
			}
		}
	}

	return nil
}

// ValidateClusterWithDuration runs validation.ValidateCluster until either we get positive result or the timeout expires
func (r *RollingUpdateInstanceGroup) ValidateClusterWithDuration(rollingUpdateData *RollingUpdateCluster, cluster *api.Cluster, instanceGroupList *api.InstanceGroupList, duration time.Duration) error {
	// TODO should we expose this to the UI?
	tickDuration := 30 * time.Second
	// Try to validate cluster at least once, this will handle durations that are lower
	// than our tick time
	if r.tryValidateCluster(rollingUpdateData, cluster, instanceGroupList, duration, tickDuration) {
		return nil
	}

	timeout := time.After(duration)
	tick := time.Tick(tickDuration)
	// Keep trying until we're timed out or got a result or got an error
	for {
		select {
		case <-timeout:
			// Got a timeout fail with a timeout error
			return fmt.Errorf("cluster did not validate within a duation of %q", duration)
		case <-tick:
			// Got a tick, validate cluster
			if r.tryValidateCluster(rollingUpdateData, cluster, instanceGroupList, duration, tickDuration) {
				return nil
			}
			// ValidateCluster didn't work yet, so let's try again
			// this will exit up to the for loop
		}
	}
}

func (r *RollingUpdateInstanceGroup) tryValidateCluster(rollingUpdateData *RollingUpdateCluster, cluster *api.Cluster, instanceGroupList *api.InstanceGroupList, duration time.Duration, tickDuration time.Duration) bool {
	if _, err := validation.ValidateCluster(cluster, instanceGroupList, rollingUpdateData.K8sClient); err != nil {
		glog.Infof("Cluster did not validate, will try again in %q until duration %q expires: %v.", tickDuration, duration, err)
		return false
	} else {
		glog.Infof("Cluster validated.")
		return true
	}
}

// ValidateCluster runs our validation methods on the K8s Cluster.
func (r *RollingUpdateInstanceGroup) ValidateCluster(rollingUpdateData *RollingUpdateCluster, cluster *api.Cluster, instanceGroupList *api.InstanceGroupList) error {
	if _, err := validation.ValidateCluster(cluster, instanceGroupList, rollingUpdateData.K8sClient); err != nil {
		return fmt.Errorf("cluster %q did not pass validation: %v", cluster.Name, err)
	}

	return nil

}

// DeleteInstance deletes an Cloud Instance.
func (r *RollingUpdateInstanceGroup) DeleteInstance(u *cloudinstances.CloudInstanceGroupMember) error {

	id := u.ID
	nodeName := ""
	if u.Node != nil {
		nodeName = u.Node.Name
	}
	if nodeName != "" {
		glog.Infof("Stopping instance %q, node %q, in group %q (this may take a while).", id, nodeName, r.CloudGroup.HumanName)
	} else {
		glog.Infof("Stopping instance %q, in group %q (this may take a while).", id, r.CloudGroup.HumanName)
	}

	if err := r.Cloud.DeleteInstance(u); err != nil {
		if nodeName != "" {
			return fmt.Errorf("error deleting instance %q, node %q: %v", id, nodeName, err)
		} else {
			return fmt.Errorf("error deleting instance %q: %v", id, err)
		}
	}

	return nil

}

// DrainNode drains a K8s node.
func (r *RollingUpdateInstanceGroup) DrainNode(u *cloudinstances.CloudInstanceGroupMember, rollingUpdateData *RollingUpdateCluster) error {
	if rollingUpdateData.ClientConfig == nil {
		return fmt.Errorf("clientConfig not set")
	}

	if u.Node.Name == "" {
		return fmt.Errorf("node name not set")
	}
	f := cmdutil.NewFactory(rollingUpdateData.ClientConfig)

	// TODO: Send out somewhere else, also DrainOptions has errout
	out := os.Stdout
	errOut := os.Stderr

	options := &cmd.DrainOptions{
		Factory:            f,
		Out:                out,
		IgnoreDaemonsets:   true,
		Force:              true,
		DeleteLocalData:    true,
		ErrOut:             errOut,
		GracePeriodSeconds: -1,
	}

	cmd := cmd.NewCmdDrain(f, out, errOut)
	args := []string{u.Node.Name}
	err := options.SetupDrain(cmd, args)
	if err != nil {
		return fmt.Errorf("error setting up drain: %v", err)
	}

	err = options.RunCordonOrUncordon(true)
	if err != nil {
		return fmt.Errorf("error cordoning node node: %v", err)
	}

	err = options.RunDrain()
	if err != nil {
		return fmt.Errorf("error draining node: %v", err)
	}

	if rollingUpdateData.PostDrainDelay > 0 {
		glog.Infof("Waiting for %s for pods to stabilize after draining.", rollingUpdateData.PostDrainDelay)
		time.Sleep(rollingUpdateData.PostDrainDelay)
	}

	return nil
}

// Delete a CloudInstanceGroups
func (r *RollingUpdateInstanceGroup) Delete() error {
	if r.CloudGroup == nil {
		return fmt.Errorf("group has to be set")
	}
	// TODO: Leaving func in place in order to cordon nd drain nodes
	return r.Cloud.DeleteGroup(r.CloudGroup)
}
<commit_msg>Fail cluster validation for rolling-update if a failure occurs<commit_after>/*
Copyright 2016 The Kubernetes 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 instancegroups

import (
	"bufio"
	"fmt"
	"os"
	"strings"
	"time"

	"github.com/golang/glog"
	api "k8s.io/kops/pkg/apis/kops"
	"k8s.io/kops/pkg/cloudinstances"
	"k8s.io/kops/pkg/featureflag"
	"k8s.io/kops/pkg/validation"
	"k8s.io/kops/upup/pkg/fi"
	"k8s.io/kubernetes/pkg/kubectl/cmd"
	cmdutil "k8s.io/kubernetes/pkg/kubectl/cmd/util"
)

// RollingUpdateInstanceGroup is the AWS ASG backing an InstanceGroup.
type RollingUpdateInstanceGroup struct {
	// Cloud is the kops cloud provider
	Cloud fi.Cloud
	// CloudGroup is the kops cloud provider groups
	CloudGroup *cloudinstances.CloudInstanceGroup

	// TODO should remove the need to have rollingupdate struct and add:
	// TODO - the kubernetes client
	// TODO - the cluster name
	// TODO - the client config
	// TODO - fail on validate
	// TODO - fail on drain
	// TODO - cloudonly
}

// NewRollingUpdateInstanceGroup creates a new struct
func NewRollingUpdateInstanceGroup(cloud fi.Cloud, cloudGroup *cloudinstances.CloudInstanceGroup) (*RollingUpdateInstanceGroup, error) {
	if cloud == nil {
		return nil, fmt.Errorf("cloud provider is required")
	}
	if cloudGroup == nil {
		return nil, fmt.Errorf("cloud group is required")
	}

	// TODO check more values in cloudGroup that they are set properly

	return &RollingUpdateInstanceGroup{
		Cloud:      cloud,
		CloudGroup: cloudGroup,
	}, nil
}

// promptInteractive asks the user to continue, mostly copied from vendor/google.golang.org/api/examples/gmail.go.
func promptInteractive(upgradedHost string) (stopPrompting bool, err error) {
	stopPrompting = false
	scanner := bufio.NewScanner(os.Stdin)
	glog.Infof("Pausing after finished %q", upgradedHost)
	fmt.Print("Continue? (Y)es, (N)o, (A)lwaysYes: [Y] ")
	scanner.Scan()
	err = scanner.Err()
	if err != nil {
		glog.Infof("unable to interpret input: %v", err)
		return stopPrompting, err
	}
	val := scanner.Text()
	val = strings.TrimSpace(val)
	val = strings.ToLower(val)
	switch val {
	case "n":
		glog.Infof("User signaled to stop")
		os.Exit(3)
	case "a":
		glog.Infof("Always Yes, stop prompting for rest of hosts")
		stopPrompting = true
	}
	return stopPrompting, err
}

// TODO: Temporarily increase size of ASG?
// TODO: Remove from ASG first so status is immediately updated?
// TODO: Batch termination, like a rolling-update

// RollingUpdate performs a rolling update on a list of ec2 instances.
func (r *RollingUpdateInstanceGroup) RollingUpdate(rollingUpdateData *RollingUpdateCluster, cluster *api.Cluster, instanceGroupList *api.InstanceGroupList, isBastion bool, sleepAfterTerminate time.Duration, validationTimeout time.Duration) (err error) {

	// we should not get here, but hey I am going to check.
	if rollingUpdateData == nil {
		return fmt.Errorf("rollingUpdate cannot be nil")
	}

	// Do not need a k8s client if you are doing cloudonly.
	if rollingUpdateData.K8sClient == nil && !rollingUpdateData.CloudOnly {
		return fmt.Errorf("rollingUpdate is missing a k8s client")
	}

	if instanceGroupList == nil {
		return fmt.Errorf("rollingUpdate is missing the InstanceGroupList")
	}

	update := r.CloudGroup.NeedUpdate
	if rollingUpdateData.Force {
		update = append(update, r.CloudGroup.Ready...)
	}

	if len(update) == 0 {
		return nil
	}

	if isBastion {
		glog.V(3).Info("Not validating the cluster as instance is a bastion.")
	} else if rollingUpdateData.CloudOnly {
		glog.V(3).Info("Not validating cluster as validation is turned off via the cloud-only flag.")
	} else if featureflag.DrainAndValidateRollingUpdate.Enabled() {
		if err = r.ValidateCluster(rollingUpdateData, cluster, instanceGroupList); err != nil {
			if rollingUpdateData.FailOnValidate {
				return fmt.Errorf("error validating cluster: %v", err)
			} else {
				glog.V(2).Infof("Ignoring cluster validation error: %v", err)
				glog.Infof("Cluster validation failed, but proceeding since fail-on-validate-error is set to false")
			}
		}
	}

	for _, u := range update {
		instanceId := u.ID

		nodeName := ""
		if u.Node != nil {
			nodeName = u.Node.Name
		}

		if isBastion {
			// We don't want to validate for bastions - they aren't part of the cluster
		} else if rollingUpdateData.CloudOnly {

			glog.Warningf("Not draining cluster nodes as 'cloudonly' flag is set.")

		} else if featureflag.DrainAndValidateRollingUpdate.Enabled() {

			if u.Node != nil {
				glog.Infof("Draining the node: %q.", nodeName)

				if err = r.DrainNode(u, rollingUpdateData); err != nil {
					if rollingUpdateData.FailOnDrainError {
						return fmt.Errorf("failed to drain node %q: %v", nodeName, err)
					} else {
						glog.Infof("Ignoring error draining node %q: %v", nodeName, err)
					}
				}
			} else {
				glog.Warningf("Skipping drain of instance %q, because it is not registered in kubernetes", instanceId)
			}
		}

		if err = r.DeleteInstance(u); err != nil {
			glog.Errorf("Error deleting aws instance %q, node %q: %v", instanceId, nodeName, err)
			return err
		}

		// Wait for the minimum interval
		time.Sleep(sleepAfterTerminate)

		if isBastion {
			glog.Infof("Deleted a bastion instance, %s, and continuing with rolling-update.", instanceId)

			continue
		} else if rollingUpdateData.CloudOnly {
			glog.Warningf("Not validating cluster as cloudonly flag is set.")
			continue

		} else if featureflag.DrainAndValidateRollingUpdate.Enabled() {
			glog.Infof("Validating the cluster.")

			if err = r.ValidateClusterWithDuration(rollingUpdateData, cluster, instanceGroupList, validationTimeout); err != nil {

				if rollingUpdateData.FailOnValidate {
					glog.Errorf("Cluster did not validate within %s", validationTimeout)
					return fmt.Errorf("error validating cluster after removing a node: %v", err)
				}

				glog.Warningf("Cluster validation failed after removing instance, proceeding since fail-on-validate is set to false: %v", err)
			}
			if rollingUpdateData.Interactive {
				stopPrompting, err := promptInteractive(nodeName)
				if err != nil {
					return err
				}
				if stopPrompting {
					// Is a pointer to a struct, changes here push back into the original
					rollingUpdateData.Interactive = false
				}
			}
		}
	}

	return nil
}

// ValidateClusterWithDuration runs validation.ValidateCluster until either we get positive result or the timeout expires
func (r *RollingUpdateInstanceGroup) ValidateClusterWithDuration(rollingUpdateData *RollingUpdateCluster, cluster *api.Cluster, instanceGroupList *api.InstanceGroupList, duration time.Duration) error {
	// TODO should we expose this to the UI?
	tickDuration := 30 * time.Second
	// Try to validate cluster at least once, this will handle durations that are lower
	// than our tick time
	if r.tryValidateCluster(rollingUpdateData, cluster, instanceGroupList, duration, tickDuration) {
		return nil
	}

	timeout := time.After(duration)
	tick := time.Tick(tickDuration)
	// Keep trying until we're timed out or got a result or got an error
	for {
		select {
		case <-timeout:
			// Got a timeout fail with a timeout error
			return fmt.Errorf("cluster did not validate within a duation of %q", duration)
		case <-tick:
			// Got a tick, validate cluster
			if r.tryValidateCluster(rollingUpdateData, cluster, instanceGroupList, duration, tickDuration) {
				return nil
			}
			// ValidateCluster didn't work yet, so let's try again
			// this will exit up to the for loop
		}
	}
}

func (r *RollingUpdateInstanceGroup) tryValidateCluster(rollingUpdateData *RollingUpdateCluster, cluster *api.Cluster, instanceGroupList *api.InstanceGroupList, duration time.Duration, tickDuration time.Duration) bool {
	result, err := validation.ValidateCluster(cluster, instanceGroupList, rollingUpdateData.K8sClient)

	if err != nil {
		glog.Infof("Cluster did not validate, will try again in %q until duration %q expires: %v.", tickDuration, duration, err)
		return false
	} else if len(result.Failures) > 0 {
		glog.Infof("Cluster did not pass validation, will try again in %q until duration %q expires: %v.", tickDuration, duration, result.Failures[0].Message)
		return false
	} else {
		glog.Infof("Cluster validated.")
		return true
	}
}

// ValidateCluster runs our validation methods on the K8s Cluster.
func (r *RollingUpdateInstanceGroup) ValidateCluster(rollingUpdateData *RollingUpdateCluster, cluster *api.Cluster, instanceGroupList *api.InstanceGroupList) error {
	if _, err := validation.ValidateCluster(cluster, instanceGroupList, rollingUpdateData.K8sClient); err != nil {
		return fmt.Errorf("cluster %q did not pass validation: %v", cluster.Name, err)
	}

	return nil

}

// DeleteInstance deletes an Cloud Instance.
func (r *RollingUpdateInstanceGroup) DeleteInstance(u *cloudinstances.CloudInstanceGroupMember) error {

	id := u.ID
	nodeName := ""
	if u.Node != nil {
		nodeName = u.Node.Name
	}
	if nodeName != "" {
		glog.Infof("Stopping instance %q, node %q, in group %q (this may take a while).", id, nodeName, r.CloudGroup.HumanName)
	} else {
		glog.Infof("Stopping instance %q, in group %q (this may take a while).", id, r.CloudGroup.HumanName)
	}

	if err := r.Cloud.DeleteInstance(u); err != nil {
		if nodeName != "" {
			return fmt.Errorf("error deleting instance %q, node %q: %v", id, nodeName, err)
		} else {
			return fmt.Errorf("error deleting instance %q: %v", id, err)
		}
	}

	return nil

}

// DrainNode drains a K8s node.
func (r *RollingUpdateInstanceGroup) DrainNode(u *cloudinstances.CloudInstanceGroupMember, rollingUpdateData *RollingUpdateCluster) error {
	if rollingUpdateData.ClientConfig == nil {
		return fmt.Errorf("clientConfig not set")
	}

	if u.Node.Name == "" {
		return fmt.Errorf("node name not set")
	}
	f := cmdutil.NewFactory(rollingUpdateData.ClientConfig)

	// TODO: Send out somewhere else, also DrainOptions has errout
	out := os.Stdout
	errOut := os.Stderr

	options := &cmd.DrainOptions{
		Factory:            f,
		Out:                out,
		IgnoreDaemonsets:   true,
		Force:              true,
		DeleteLocalData:    true,
		ErrOut:             errOut,
		GracePeriodSeconds: -1,
	}

	cmd := cmd.NewCmdDrain(f, out, errOut)
	args := []string{u.Node.Name}
	err := options.SetupDrain(cmd, args)
	if err != nil {
		return fmt.Errorf("error setting up drain: %v", err)
	}

	err = options.RunCordonOrUncordon(true)
	if err != nil {
		return fmt.Errorf("error cordoning node node: %v", err)
	}

	err = options.RunDrain()
	if err != nil {
		return fmt.Errorf("error draining node: %v", err)
	}

	if rollingUpdateData.PostDrainDelay > 0 {
		glog.Infof("Waiting for %s for pods to stabilize after draining.", rollingUpdateData.PostDrainDelay)
		time.Sleep(rollingUpdateData.PostDrainDelay)
	}

	return nil
}

// Delete a CloudInstanceGroups
func (r *RollingUpdateInstanceGroup) Delete() error {
	if r.CloudGroup == nil {
		return fmt.Errorf("group has to be set")
	}
	// TODO: Leaving func in place in order to cordon nd drain nodes
	return r.Cloud.DeleteGroup(r.CloudGroup)
}
<|endoftext|>


<commit_before>package bencode

import (
	"bytes"
	"fmt"
	"os"
	"reflect"
	"testing"
)

type any interface{}

func checkMarshal(expected string, data any) (err os.Error) {
	var b bytes.Buffer
	if err = Marshal(&b, data); err != nil {
		return
	}
	s := b.String()
	if expected != s {
		err = os.NewError(fmt.Sprintf("Expected %s got %s", expected, s))
		return
	}
	return
}

func check(expected string, data any) (err os.Error) {
	if err = checkMarshal(expected, data); err != nil {
		return
	}
	b2 := bytes.NewBufferString(expected)
	val, err := Decode(b2)
	if err != nil {
		err = os.NewError(fmt.Sprint("Failed decoding ", expected, " ", err))
		return
	}
	if err = checkFuzzyEqual(data, val); err != nil {
		return
	}
	return
}

func checkFuzzyEqual(a any, b any) (err os.Error) {
	if !fuzzyEqual(a, b) {
		err = os.NewError(fmt.Sprint(a, " != ", b,
			":", reflect.NewValue(a), "!=", reflect.NewValue(b)))
	}
	return
}

func fuzzyEqual(a, b any) bool {
	return fuzzyEqualValue(reflect.NewValue(a), reflect.NewValue(b))
}

func checkFuzzyEqualValue(a, b reflect.Value) (err os.Error) {
	if !fuzzyEqualValue(a, b) {
		err = os.NewError(fmt.Sprint(a, " != ", b,
			":", a.Interface(), "!=", b.Interface()))
	}
	return
}

func fuzzyEqualInt64(a int64, b reflect.Value) bool {
	switch vb := b.(type) {
	case *reflect.IntValue:
		return a == (vb.Get())
	default:
		return false
	}
	return false
}

func fuzzyEqualArrayOrSlice(va reflect.ArrayOrSliceValue, b reflect.Value) bool {
	switch vb := b.(type) {
	case *reflect.ArrayValue:
		return fuzzyEqualArrayOrSlice2(va, vb)
	case *reflect.SliceValue:
		return fuzzyEqualArrayOrSlice2(va, vb)
	default:
		return false
	}
	return false
}

func deInterface(a reflect.Value) reflect.Value {
	switch va := a.(type) {
	case *reflect.InterfaceValue:
		return va.Elem()
	}
	return a
}

func fuzzyEqualArrayOrSlice2(a reflect.ArrayOrSliceValue, b reflect.ArrayOrSliceValue) bool {
	if a.Len() != b.Len() {
		return false
	}

	for i := 0; i < a.Len(); i++ {
		ea := deInterface(a.Elem(i))
		eb := deInterface(b.Elem(i))
		if !fuzzyEqualValue(ea, eb) {
			return false
		}
	}
	return true
}

func fuzzyEqualMap(a *reflect.MapValue, b *reflect.MapValue) bool {
	key := a.Type().(*reflect.MapType).Key()
	if _, ok := key.(*reflect.StringType); !ok {
		return false
	}
	key = b.Type().(*reflect.MapType).Key()
	if _, ok := key.(*reflect.StringType); !ok {
		return false
	}

	aKeys, bKeys := a.Keys(), b.Keys()

	if len(aKeys) != len(bKeys) {
		return false
	}

	for _, k := range aKeys {
		if !fuzzyEqualValue(a.Elem(k), b.Elem(k)) {
			return false
		}
	}
	return true
}

func fuzzyEqualStruct(a *reflect.StructValue, b *reflect.StructValue) bool {
	numA, numB := a.NumField(), b.NumField()
	if numA != numB {
		return false
	}

	for i := 0; i < numA; i++ {
		if !fuzzyEqualValue(a.Field(i), b.Field(i)) {
			return false
		}
	}
	return true
}

func fuzzyEqualValue(a, b reflect.Value) bool {
	switch va := a.(type) {
	case *reflect.StringValue:
		switch vb := b.(type) {
		case *reflect.StringValue:
			return va.Get() == vb.Get()
		default:
			return false
		}
	case *reflect.IntValue:
		return fuzzyEqualInt64(va.Get(), b)
	case *reflect.ArrayValue:
		return fuzzyEqualArrayOrSlice(va, b)
	case *reflect.SliceValue:
		return fuzzyEqualArrayOrSlice(va, b)
	case *reflect.MapValue:
		switch vb := b.(type) {
		case *reflect.MapValue:
			return fuzzyEqualMap(va, vb)
		default:
			return false
		}
	case *reflect.StructValue:
		switch vb := b.(type) {
		case *reflect.StructValue:
			return fuzzyEqualStruct(va, vb)
		default:
			return false
		}
	case *reflect.InterfaceValue:
		switch vb := b.(type) {
		case *reflect.InterfaceValue:
			return fuzzyEqualValue(va.Elem(), vb.Elem())
		default:
			return false
		}
	default:
		return false
	}
	return false
}

func checkUnmarshal(expected string, data any) (err os.Error) {
	if err = checkMarshal(expected, data); err != nil {
		return
	}
	dataValue := reflect.NewValue(data)
	newOne := reflect.MakeZero(dataValue.Type())
	buf := bytes.NewBufferString(expected)
	if err = UnmarshalValue(buf, newOne); err != nil {
		return
	}
	if err = checkFuzzyEqualValue(dataValue, newOne); err != nil {
		return
	}
	return
}

type SVPair struct {
	s string
	v any
}

func TestDecode(t *testing.T) {
	tests := []SVPair{
		SVPair{"i0e", int64(0)},
		SVPair{"i0e", 0},
		SVPair{"i100e", 100},
		SVPair{"i-100e", -100},
		SVPair{"1:a", "a"},
		SVPair{"2:a\"", "a\""},
		SVPair{"11:0123456789a", "0123456789a"},
		SVPair{"le", []int64{}},
		SVPair{"li1ei2ee", []int{1, 2}},
		SVPair{"l3:abc3:defe", []string{"abc", "def"}},
		SVPair{"li42e3:abce", []any{42, "abc"}},
		SVPair{"de", map[string]any{}},
		SVPair{"d3:cati1e3:dogi2ee", map[string]any{"cat": 1, "dog": 2}},
	}
	for _, sv := range tests {
		if err := check(sv.s, sv.v); err != nil {
			t.Error(err.String())
		}
	}
}

type structA struct {
	A int    "a"
	B string "b"
}

func TestUnmarshal(t *testing.T) {
	tests := []SVPair{
		SVPair{"i0e", int64(0)},
		SVPair{"i0e", 0},
		SVPair{"i100e", 100},
		SVPair{"i-100e", -100},
		SVPair{"1:a", "a"},
		SVPair{"2:a\"", "a\""},
		SVPair{"11:0123456789a", "0123456789a"},
		SVPair{"le", []int64{}},
		SVPair{"li1ei2ee", []int{1, 2}},
		SVPair{"l3:abc3:defe", []string{"abc", "def"}},
		SVPair{"li42e3:abce", []any{42, "abc"}},
		SVPair{"de", map[string]any{}},
		SVPair{"d3:cati1e3:dogi2ee", map[string]any{"cat": 1, "dog": 2}},
		SVPair{"d1:ai10e1:b3:fooe", structA{10, "foo"}},
	}
	for _, sv := range tests {
		if err := checkUnmarshal(sv.s, sv.v); err != nil {
			t.Error(err.String())
		}
	}
}
<commit_msg>bencode: Adding a unit-test proving that "dictionary inside a dictionary" works, as long as you know how to use it. Closes issue #2 (user error!).<commit_after>package bencode

import (
	"bytes"
	"fmt"
	"os"
	"reflect"
	"testing"
)

type any interface{}

func checkMarshal(expected string, data any) (err os.Error) {
	var b bytes.Buffer
	if err = Marshal(&b, data); err != nil {
		return
	}
	s := b.String()
	if expected != s {
		err = os.NewError(fmt.Sprintf("Expected %s got %s", expected, s))
		return
	}
	return
}

func check(expected string, data any) (err os.Error) {
	if err = checkMarshal(expected, data); err != nil {
		return
	}
	b2 := bytes.NewBufferString(expected)
	val, err := Decode(b2)
	if err != nil {
		err = os.NewError(fmt.Sprint("Failed decoding ", expected, " ", err))
		return
	}
	if err = checkFuzzyEqual(data, val); err != nil {
		return
	}
	return
}

func checkFuzzyEqual(a any, b any) (err os.Error) {
	if !fuzzyEqual(a, b) {
		err = os.NewError(fmt.Sprint(a, " != ", b,
			":", reflect.NewValue(a), "!=", reflect.NewValue(b)))
	}
	return
}

func fuzzyEqual(a, b any) bool {
	return fuzzyEqualValue(reflect.NewValue(a), reflect.NewValue(b))
}

func checkFuzzyEqualValue(a, b reflect.Value) (err os.Error) {
	if !fuzzyEqualValue(a, b) {
		err = os.NewError(fmt.Sprint(a, " != ", b,
			":", a.Interface(), "!=", b.Interface()))
	}
	return
}

func fuzzyEqualInt64(a int64, b reflect.Value) bool {
	switch vb := b.(type) {
	case *reflect.IntValue:
		return a == (vb.Get())
	default:
		return false
	}
	return false
}

func fuzzyEqualArrayOrSlice(va reflect.ArrayOrSliceValue, b reflect.Value) bool {
	switch vb := b.(type) {
	case *reflect.ArrayValue:
		return fuzzyEqualArrayOrSlice2(va, vb)
	case *reflect.SliceValue:
		return fuzzyEqualArrayOrSlice2(va, vb)
	default:
		return false
	}
	return false
}

func deInterface(a reflect.Value) reflect.Value {
	switch va := a.(type) {
	case *reflect.InterfaceValue:
		return va.Elem()
	}
	return a
}

func fuzzyEqualArrayOrSlice2(a reflect.ArrayOrSliceValue, b reflect.ArrayOrSliceValue) bool {
	if a.Len() != b.Len() {
		return false
	}

	for i := 0; i < a.Len(); i++ {
		ea := deInterface(a.Elem(i))
		eb := deInterface(b.Elem(i))
		if !fuzzyEqualValue(ea, eb) {
			return false
		}
	}
	return true
}

func fuzzyEqualMap(a *reflect.MapValue, b *reflect.MapValue) bool {
	key := a.Type().(*reflect.MapType).Key()
	if _, ok := key.(*reflect.StringType); !ok {
		return false
	}
	key = b.Type().(*reflect.MapType).Key()
	if _, ok := key.(*reflect.StringType); !ok {
		return false
	}

	aKeys, bKeys := a.Keys(), b.Keys()

	if len(aKeys) != len(bKeys) {
		return false
	}

	for _, k := range aKeys {
		if !fuzzyEqualValue(a.Elem(k), b.Elem(k)) {
			return false
		}
	}
	return true
}

func fuzzyEqualStruct(a *reflect.StructValue, b *reflect.StructValue) bool {
	numA, numB := a.NumField(), b.NumField()
	if numA != numB {
		return false
	}

	for i := 0; i < numA; i++ {
		if !fuzzyEqualValue(a.Field(i), b.Field(i)) {
			return false
		}
	}
	return true
}

func fuzzyEqualValue(a, b reflect.Value) bool {
	switch va := a.(type) {
	case *reflect.StringValue:
		switch vb := b.(type) {
		case *reflect.StringValue:
			return va.Get() == vb.Get()
		default:
			return false
		}
	case *reflect.IntValue:
		return fuzzyEqualInt64(va.Get(), b)
	case *reflect.ArrayValue:
		return fuzzyEqualArrayOrSlice(va, b)
	case *reflect.SliceValue:
		return fuzzyEqualArrayOrSlice(va, b)
	case *reflect.MapValue:
		switch vb := b.(type) {
		case *reflect.MapValue:
			return fuzzyEqualMap(va, vb)
		default:
			return false
		}
	case *reflect.StructValue:
		switch vb := b.(type) {
		case *reflect.StructValue:
			return fuzzyEqualStruct(va, vb)
		default:
			return false
		}
	case *reflect.InterfaceValue:
		switch vb := b.(type) {
		case *reflect.InterfaceValue:
			return fuzzyEqualValue(va.Elem(), vb.Elem())
		default:
			return false
		}
	default:
		return false
	}
	return false
}

func checkUnmarshal(expected string, data any) (err os.Error) {
	if err = checkMarshal(expected, data); err != nil {
		return
	}
	dataValue := reflect.NewValue(data)
	newOne := reflect.MakeZero(dataValue.Type())
	buf := bytes.NewBufferString(expected)
	if err = UnmarshalValue(buf, newOne); err != nil {
		return
	}
	if err = checkFuzzyEqualValue(dataValue, newOne); err != nil {
		return
	}
	return
}

type SVPair struct {
	s string
	v any
}

func TestDecode(t *testing.T) {
	tests := []SVPair{
		SVPair{"i0e", int64(0)},
		SVPair{"i0e", 0},
		SVPair{"i100e", 100},
		SVPair{"i-100e", -100},
		SVPair{"1:a", "a"},
		SVPair{"2:a\"", "a\""},
		SVPair{"11:0123456789a", "0123456789a"},
		SVPair{"le", []int64{}},
		SVPair{"li1ei2ee", []int{1, 2}},
		SVPair{"l3:abc3:defe", []string{"abc", "def"}},
		SVPair{"li42e3:abce", []any{42, "abc"}},
		SVPair{"de", map[string]any{}},
		SVPair{"d3:cati1e3:dogi2ee", map[string]any{"cat": 1, "dog": 2}},
	}
	for _, sv := range tests {
		if err := check(sv.s, sv.v); err != nil {
			t.Error(err.String())
		}
	}
}

type structA struct {
	A int    "a"
	B string "b"
}

func TestUnmarshal(t *testing.T) {
	type structNested struct {
		T string            "t"
		Y string            "y"
		Q string            "q"
		A map[string]string "a"
	}
	innerDict := map[string]string{"id": "abcdefghij0123456789"}
	nestedDictionary := structNested{"aa", "q", "ping", innerDict}

	tests := []SVPair{
		SVPair{"i0e", int64(0)},
		SVPair{"i0e", 0},
		SVPair{"i100e", 100},
		SVPair{"i-100e", -100},
		SVPair{"1:a", "a"},
		SVPair{"2:a\"", "a\""},
		SVPair{"11:0123456789a", "0123456789a"},
		SVPair{"le", []int64{}},
		SVPair{"li1ei2ee", []int{1, 2}},
		SVPair{"l3:abc3:defe", []string{"abc", "def"}},
		SVPair{"li42e3:abce", []any{42, "abc"}},
		SVPair{"de", map[string]any{}},
		SVPair{"d3:cati1e3:dogi2ee", map[string]any{"cat": 1, "dog": 2}},
		SVPair{"d1:ai10e1:b3:fooe", structA{10, "foo"}},
		SVPair{"d1:ad2:id20:abcdefghij0123456789e1:q4:ping1:t2:aa1:y1:qe", nestedDictionary},
	}
	for _, sv := range tests {
		if err := checkUnmarshal(sv.s, sv.v); err != nil {
			t.Error(err.String())
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2019 Authors of Cilium
//
// 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.

// +build !privileged_tests

package lock

import (
	"math/rand"
	"sync"
	"sync/atomic"
	"time"

	. "gopkg.in/check.v1"
)

type StoppableWaitGroupSuite struct{}

var _ = Suite(&StoppableWaitGroupSuite{})

func (s *SemaphoredMutexSuite) TestAdd(c *C) {
	l := NewStoppableWaitGroup()

	l.Add()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(1))
	l.Add()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(2))
	close(l.noopAdd)
	l.Add()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(2))
}

func (s *SemaphoredMutexSuite) TestDone(c *C) {
	l := NewStoppableWaitGroup()

	atomic.StoreInt64(l.i, 4)
	l.Done()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(3))
	l.Done()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(2))
	close(l.noopAdd)
	select {
	case _, ok := <-l.noopDone:
		// channel should not have been closed
		c.Assert(ok, Equals, true)
	default:
	}

	l.Done()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(1))
	select {
	case _, ok := <-l.noopDone:
		// channel should not have been closed
		c.Assert(ok, Equals, true)
	default:
	}

	l.Done()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(0))
	select {
	case _, ok := <-l.noopDone:
		c.Assert(ok, Equals, false)
	default:
		// channel should have been closed
		c.Assert(false, Equals, true)
	}

	l.Done()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(0))
}

func (s *SemaphoredMutexSuite) TestStop(c *C) {
	l := NewStoppableWaitGroup()

	l.Add()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(1))
	l.Add()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(2))
	l.Stop()
	l.Add()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(2))
}

func (s *SemaphoredMutexSuite) TestWait(c *C) {
	l := NewStoppableWaitGroup()

	waitClosed := make(chan struct{})
	go func() {
		l.Wait()
		close(waitClosed)
	}()

	l.Add()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(1))
	l.Add()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(2))
	l.Stop()
	l.Add()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(2))

	l.Done()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(1))
	select {
	case _, ok := <-waitClosed:
		// channel should not have been closed
		c.Assert(ok, Equals, true)
	default:
	}

	l.Done()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(0))
	select {
	case _, ok := <-waitClosed:
		// channel should have been closed
		c.Assert(ok, Equals, false)
	default:
	}

	l.Done()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(0))
}

func (s *SemaphoredMutexSuite) TestWaitChannel(c *C) {
	l := NewStoppableWaitGroup()

	l.Add()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(1))
	l.Add()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(2))
	l.Stop()
	l.Add()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(2))

	l.Done()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(1))
	select {
	case _, ok := <-l.WaitChannel():
		// channel should not have been closed
		c.Assert(ok, Equals, true)
	default:
	}

	l.Done()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(0))
	select {
	case _, ok := <-l.WaitChannel():
		// channel should have been closed
		c.Assert(ok, Equals, false)
	default:
	}

	l.Done()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(0))
}

func (s *SemaphoredMutexSuite) TestParallelism(c *C) {
	l := NewStoppableWaitGroup()

	// use math/rand instead of pkg/rand to avoid a test import cycle which
	// go vet would complain about.
	randGen := rand.New(rand.NewSource(time.Now().UnixNano()))
	in := make(chan int)
	stop := make(chan struct{})
	go func() {
		for {
			select {
			case in <- randGen.Intn(1 - 0):
			case <-stop:
				close(in)
				return
			}
		}
	}()
	adds := int64(0)
	var wg sync.WaitGroup
	wg.Add(10)
	for i := 0; i < 10; i++ {
		go func() {
			defer wg.Done()
			for a := range in {
				if a == 0 {
					atomic.AddInt64(&adds, 1)
					l.Add()
				} else {
					l.Done()
					atomic.AddInt64(&adds, -1)
				}
			}
		}()
	}

	time.Sleep(time.Duration(randGen.Intn(3-0)) * time.Second)
	close(stop)
	wg.Wait()
	add := atomic.LoadInt64(&adds)
	for ; add != 0; add = atomic.LoadInt64(&adds) {
		switch {
		case add < 0:
			atomic.AddInt64(&adds, 1)
			l.Add()
		case add > 0:
			l.Done()
			atomic.AddInt64(&adds, -1)
		}
	}
	l.Stop()
	l.Wait()
}
<commit_msg>lock: fix data race in (*SemaphoredMutexSuite).TestParallelism()<commit_after>// Copyright 2019 Authors of Cilium
//
// 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.

// +build !privileged_tests

package lock

import (
	"math/rand"
	"sync"
	"sync/atomic"
	"time"

	. "gopkg.in/check.v1"
)

type StoppableWaitGroupSuite struct{}

var _ = Suite(&StoppableWaitGroupSuite{})

func (s *SemaphoredMutexSuite) TestAdd(c *C) {
	l := NewStoppableWaitGroup()

	l.Add()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(1))
	l.Add()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(2))
	close(l.noopAdd)
	l.Add()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(2))
}

func (s *SemaphoredMutexSuite) TestDone(c *C) {
	l := NewStoppableWaitGroup()

	atomic.StoreInt64(l.i, 4)
	l.Done()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(3))
	l.Done()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(2))
	close(l.noopAdd)
	select {
	case _, ok := <-l.noopDone:
		// channel should not have been closed
		c.Assert(ok, Equals, true)
	default:
	}

	l.Done()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(1))
	select {
	case _, ok := <-l.noopDone:
		// channel should not have been closed
		c.Assert(ok, Equals, true)
	default:
	}

	l.Done()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(0))
	select {
	case _, ok := <-l.noopDone:
		c.Assert(ok, Equals, false)
	default:
		// channel should have been closed
		c.Assert(false, Equals, true)
	}

	l.Done()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(0))
}

func (s *SemaphoredMutexSuite) TestStop(c *C) {
	l := NewStoppableWaitGroup()

	l.Add()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(1))
	l.Add()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(2))
	l.Stop()
	l.Add()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(2))
}

func (s *SemaphoredMutexSuite) TestWait(c *C) {
	l := NewStoppableWaitGroup()

	waitClosed := make(chan struct{})
	go func() {
		l.Wait()
		close(waitClosed)
	}()

	l.Add()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(1))
	l.Add()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(2))
	l.Stop()
	l.Add()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(2))

	l.Done()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(1))
	select {
	case _, ok := <-waitClosed:
		// channel should not have been closed
		c.Assert(ok, Equals, true)
	default:
	}

	l.Done()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(0))
	select {
	case _, ok := <-waitClosed:
		// channel should have been closed
		c.Assert(ok, Equals, false)
	default:
	}

	l.Done()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(0))
}

func (s *SemaphoredMutexSuite) TestWaitChannel(c *C) {
	l := NewStoppableWaitGroup()

	l.Add()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(1))
	l.Add()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(2))
	l.Stop()
	l.Add()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(2))

	l.Done()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(1))
	select {
	case _, ok := <-l.WaitChannel():
		// channel should not have been closed
		c.Assert(ok, Equals, true)
	default:
	}

	l.Done()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(0))
	select {
	case _, ok := <-l.WaitChannel():
		// channel should have been closed
		c.Assert(ok, Equals, false)
	default:
	}

	l.Done()
	c.Assert(atomic.LoadInt64(l.i), Equals, int64(0))
}

func (s *SemaphoredMutexSuite) TestParallelism(c *C) {
	l := NewStoppableWaitGroup()

	// Use math/rand instead of pkg/rand to avoid a test import cycle which
	// go vet would complain about. Use the global default entropy source
	// rather than creating a new source to avoid concurrency issues.
	rand.Seed(time.Now().UnixNano())
	in := make(chan int)
	stop := make(chan struct{})
	go func() {
		for {
			select {
			case in <- rand.Intn(1 - 0):
			case <-stop:
				close(in)
				return
			}
		}
	}()
	adds := int64(0)
	var wg sync.WaitGroup
	wg.Add(10)
	for i := 0; i < 10; i++ {
		go func() {
			defer wg.Done()
			for a := range in {
				if a == 0 {
					atomic.AddInt64(&adds, 1)
					l.Add()
				} else {
					l.Done()
					atomic.AddInt64(&adds, -1)
				}
			}
		}()
	}

	time.Sleep(time.Duration(rand.Intn(3-0)) * time.Second)
	close(stop)
	wg.Wait()
	add := atomic.LoadInt64(&adds)
	for ; add != 0; add = atomic.LoadInt64(&adds) {
		switch {
		case add < 0:
			atomic.AddInt64(&adds, 1)
			l.Add()
		case add > 0:
			l.Done()
			atomic.AddInt64(&adds, -1)
		}
	}
	l.Stop()
	l.Wait()
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Kubernetes Authors All rights reserved.

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

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

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

package cruntime

import (
	"bytes"
	"os/exec"
	"testing"

	"k8s.io/minikube/pkg/minikube/assets"
	"k8s.io/minikube/pkg/minikube/command"
)

type mockRunner struct {
	runCmdOutput bytes.Buffer
}

func (m *mockRunner) RunCmd(cmd *exec.Cmd) (*command.RunResult, error) {
	return &command.RunResult{
		Stdout: m.runCmdOutput,
	}, nil
}

// Copy is a convenience method that runs a command to copy a file
func (m *mockRunner) Copy(a assets.CopyableFile) error {
	return nil
}

// Remove is a convenience method that runs a command to remove a file
func (m *mockRunner) Remove(assets.CopyableFile) error {
	return nil
}

func TestMergeRepositoriesJSON(t *testing.T) {
	initial := `{
		"Repositories": {
		  "imageOne": {
			"image:tag": "digest",
			"image@digest": "digest"
		  },
		  "imageTwo": {
			"image2:tag": "digest",
			"image2@digest": "digest"
		  }
		}
	}`

	afterPreload := `{
		"Repositories": {
		  "imageOne": {
			"image:tag": "new_digest",
			"image@digest": "new_digest"
		  },
		  "imageThree": {
			"image2:tag": "digest_three",
			"image2@digest": "digest_three"
		  }
		}
	}`

	expected := `{
		"Repositories": {
		  "imageOne": {
			"image:tag": "new_digest",
			"image@digest": "new_digest"
		  },
		  "imageTwo": {
			"image2:tag": "digest",
			"image2@digest": "digest"
		  },
		  "imageThree": {
			"image2:tag": "digest_three",
			"image2@digest": "digest_three"
		  }
		}
	}`

	runner := &mockRunner{
		runCmdOutput: *bytes.NewBuffer([]byte(afterPreload)),
	}

	driver := &Docker{Runner: runner}
	result, err := driver.mergeRepositoriesJSON(initial)
	if err != nil {
		t.Fatalf("error merging repositories.json file: %v", err)
	}

	if string(result) != expected {
		t.Fatalf("didn't get expected repositories.json. Actual:\n%s\nExpected:\n%s\n", result, expected)
	}

}
<commit_msg>removed unused unit test<commit_after><|endoftext|>


<commit_before>// Copyright 2020 The gVisor 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 fuse

import (
	"io"
	"testing"

	"gvisor.dev/gvisor/pkg/abi/linux"
	"gvisor.dev/gvisor/pkg/marshal"
	"gvisor.dev/gvisor/pkg/sentry/fsimpl/testutil"
	"gvisor.dev/gvisor/pkg/sentry/kernel"
	"gvisor.dev/gvisor/pkg/sentry/kernel/auth"
	"gvisor.dev/gvisor/pkg/sentry/vfs"
	"gvisor.dev/gvisor/pkg/usermem"
)

func setup(t *testing.T) *testutil.System {
	k, err := testutil.Boot()
	if err != nil {
		t.Fatalf("Error creating kernel: %v", err)
	}

	ctx := k.SupervisorContext()
	creds := auth.CredentialsFromContext(ctx)

	k.VFS().MustRegisterFilesystemType(Name, &FilesystemType{}, &vfs.RegisterFilesystemTypeOptions{
		AllowUserList:  true,
		AllowUserMount: true,
	})

	mntns, err := k.VFS().NewMountNamespace(ctx, creds, "", "tmpfs", &vfs.GetFilesystemOptions{})
	if err != nil {
		t.Fatalf("NewMountNamespace(): %v", err)
	}

	return testutil.NewSystem(ctx, t, k.VFS(), mntns)
}

// newTestConnection creates a fuse connection that the sentry can communicate with
// and the FD for the server to communicate with.
func newTestConnection(system *testutil.System, k *kernel.Kernel, maxActiveRequests uint64) (*connection, *vfs.FileDescription, error) {
	vfsObj := &vfs.VirtualFilesystem{}
	fuseDev := &DeviceFD{}

	if err := vfsObj.Init(system.Ctx); err != nil {
		return nil, nil, err
	}

	vd := vfsObj.NewAnonVirtualDentry("genCountFD")
	defer vd.DecRef(system.Ctx)
	if err := fuseDev.vfsfd.Init(fuseDev, linux.O_RDWR|linux.O_CREAT, vd.Mount(), vd.Dentry(), &vfs.FileDescriptionOptions{}); err != nil {
		return nil, nil, err
	}

	fsopts := filesystemOptions{
		maxActiveRequests: maxActiveRequests,
	}
	fs, err := newFUSEFilesystem(system.Ctx, 0, &fsopts, &fuseDev.vfsfd)
	if err != nil {
		return nil, nil, err
	}

	return fs.conn, &fuseDev.vfsfd, nil
}

type testPayload struct {
	marshal.StubMarshallable
	data uint32
}

// SizeBytes implements marshal.Marshallable.SizeBytes.
func (t *testPayload) SizeBytes() int {
	return 4
}

// MarshalBytes implements marshal.Marshallable.MarshalBytes.
func (t *testPayload) MarshalBytes(dst []byte) {
	usermem.ByteOrder.PutUint32(dst[:4], t.data)
}

// UnmarshalBytes implements marshal.Marshallable.UnmarshalBytes.
func (t *testPayload) UnmarshalBytes(src []byte) {
	*t = testPayload{data: usermem.ByteOrder.Uint32(src[:4])}
}

// Packed implements marshal.Marshallable.Packed.
func (t *testPayload) Packed() bool {
	return true
}

// MarshalUnsafe implements marshal.Marshallable.MarshalUnsafe.
func (t *testPayload) MarshalUnsafe(dst []byte) {
	t.MarshalBytes(dst)
}

// UnmarshalUnsafe implements marshal.Marshallable.UnmarshalUnsafe.
func (t *testPayload) UnmarshalUnsafe(src []byte) {
	t.UnmarshalBytes(src)
}

// CopyOutN implements marshal.Marshallable.CopyOutN.
func (t *testPayload) CopyOutN(task marshal.CopyContext, addr usermem.Addr, limit int) (int, error) {
	panic("not implemented")
}

// CopyOut implements marshal.Marshallable.CopyOut.
func (t *testPayload) CopyOut(task marshal.CopyContext, addr usermem.Addr) (int, error) {
	panic("not implemented")
}

// CopyIn implements marshal.Marshallable.CopyIn.
func (t *testPayload) CopyIn(task marshal.CopyContext, addr usermem.Addr) (int, error) {
	panic("not implemented")
}

// WriteTo implements io.WriterTo.WriteTo.
func (t *testPayload) WriteTo(w io.Writer) (int64, error) {
	panic("not implemented")
}
<commit_msg>Fix FUSE unit test after vfs interface change<commit_after>// Copyright 2020 The gVisor 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 fuse

import (
	"io"
	"testing"

	"gvisor.dev/gvisor/pkg/abi/linux"
	"gvisor.dev/gvisor/pkg/marshal"
	"gvisor.dev/gvisor/pkg/sentry/fsimpl/testutil"
	"gvisor.dev/gvisor/pkg/sentry/kernel"
	"gvisor.dev/gvisor/pkg/sentry/kernel/auth"
	"gvisor.dev/gvisor/pkg/sentry/vfs"
	"gvisor.dev/gvisor/pkg/usermem"
)

func setup(t *testing.T) *testutil.System {
	k, err := testutil.Boot()
	if err != nil {
		t.Fatalf("Error creating kernel: %v", err)
	}

	ctx := k.SupervisorContext()
	creds := auth.CredentialsFromContext(ctx)

	k.VFS().MustRegisterFilesystemType(Name, &FilesystemType{}, &vfs.RegisterFilesystemTypeOptions{
		AllowUserList:  true,
		AllowUserMount: true,
	})

	mntns, err := k.VFS().NewMountNamespace(ctx, creds, "", "tmpfs", &vfs.MountOptions{})
	if err != nil {
		t.Fatalf("NewMountNamespace(): %v", err)
	}

	return testutil.NewSystem(ctx, t, k.VFS(), mntns)
}

// newTestConnection creates a fuse connection that the sentry can communicate with
// and the FD for the server to communicate with.
func newTestConnection(system *testutil.System, k *kernel.Kernel, maxActiveRequests uint64) (*connection, *vfs.FileDescription, error) {
	vfsObj := &vfs.VirtualFilesystem{}
	fuseDev := &DeviceFD{}

	if err := vfsObj.Init(system.Ctx); err != nil {
		return nil, nil, err
	}

	vd := vfsObj.NewAnonVirtualDentry("genCountFD")
	defer vd.DecRef(system.Ctx)
	if err := fuseDev.vfsfd.Init(fuseDev, linux.O_RDWR|linux.O_CREAT, vd.Mount(), vd.Dentry(), &vfs.FileDescriptionOptions{}); err != nil {
		return nil, nil, err
	}

	fsopts := filesystemOptions{
		maxActiveRequests: maxActiveRequests,
	}
	fs, err := newFUSEFilesystem(system.Ctx, 0, &fsopts, &fuseDev.vfsfd)
	if err != nil {
		return nil, nil, err
	}

	return fs.conn, &fuseDev.vfsfd, nil
}

type testPayload struct {
	marshal.StubMarshallable
	data uint32
}

// SizeBytes implements marshal.Marshallable.SizeBytes.
func (t *testPayload) SizeBytes() int {
	return 4
}

// MarshalBytes implements marshal.Marshallable.MarshalBytes.
func (t *testPayload) MarshalBytes(dst []byte) {
	usermem.ByteOrder.PutUint32(dst[:4], t.data)
}

// UnmarshalBytes implements marshal.Marshallable.UnmarshalBytes.
func (t *testPayload) UnmarshalBytes(src []byte) {
	*t = testPayload{data: usermem.ByteOrder.Uint32(src[:4])}
}

// Packed implements marshal.Marshallable.Packed.
func (t *testPayload) Packed() bool {
	return true
}

// MarshalUnsafe implements marshal.Marshallable.MarshalUnsafe.
func (t *testPayload) MarshalUnsafe(dst []byte) {
	t.MarshalBytes(dst)
}

// UnmarshalUnsafe implements marshal.Marshallable.UnmarshalUnsafe.
func (t *testPayload) UnmarshalUnsafe(src []byte) {
	t.UnmarshalBytes(src)
}

// CopyOutN implements marshal.Marshallable.CopyOutN.
func (t *testPayload) CopyOutN(task marshal.CopyContext, addr usermem.Addr, limit int) (int, error) {
	panic("not implemented")
}

// CopyOut implements marshal.Marshallable.CopyOut.
func (t *testPayload) CopyOut(task marshal.CopyContext, addr usermem.Addr) (int, error) {
	panic("not implemented")
}

// CopyIn implements marshal.Marshallable.CopyIn.
func (t *testPayload) CopyIn(task marshal.CopyContext, addr usermem.Addr) (int, error) {
	panic("not implemented")
}

// WriteTo implements io.WriterTo.WriteTo.
func (t *testPayload) WriteTo(w io.Writer) (int64, error) {
	panic("not implemented")
}
<|endoftext|>


<commit_before>// Package s4lru implements the 4-segmented LRU caching algorithm
/*

From http://www.cs.cornell.edu/~qhuang/papers/sosp_fbanalysis.pdf

    Four queues are maintained at levels 0 to 3. On a cache miss, the item is
    inserted at the head of queue 0. On a cache hit, the item is moved to the
    head of the next higher queue (items in queue 3 move to the head of queue
    3).  Each queue is allocated 1/4 of the total cache size and items are
    evicted from the tail of a queue to the head of the next lower queue to
    maintain the size invariants. Items evicted from queue 0 are evicted from
    the cache

*/
package s4lru

import (
	"container/list"
	"sync"
)

type cacheItem struct {
	lidx  int
	key   string
	value interface{}
}

// Cache is an LRU cache
type Cache struct {
	mu       sync.Mutex
	capacity int
	data     map[string]*list.Element
	lists    []*list.List
}

// New returns a new S4LRU cache that with the given capacity.  Each of the lists will have 1/4 of the capacity.
func New(capacity int) *Cache {
	return &Cache{
		capacity: capacity / 4,
		data:     make(map[string]*list.Element),
		lists:    []*list.List{list.New(), list.New(), list.New(), list.New()},
	}
}

// Get returns a value from the cache
func (c *Cache) Get(key string) (interface{}, bool) {
	c.mu.Lock()
	defer c.mu.Unlock()

	v, ok := c.data[key]

	if !ok {
		return nil, false
	}

	item := v.Value.(*cacheItem)

	// already on final list?
	if item.lidx == len(c.lists)-1 {
		c.lists[item.lidx].MoveToFront(v)
		return item.value, true
	}

	// is there space on the next list?
	if c.lists[item.lidx+1].Len() < c.capacity {
		// just do the remove/add
		c.lists[item.lidx].Remove(v)
		item.lidx++
		c.data[key] = c.lists[item.lidx].PushFront(item)
		return item.value, true
	}

	// no free space on either list, so we do some in-place swapping to avoid allocations
	// the key/value in bitem need to be moved to the front of c.lists[item.lidx]
	// the key/value in item need to be moved to the front of c.lists[bitem.lidx]
	back := c.lists[item.lidx+1].Back()
	bitem := back.Value.(*cacheItem)

	// swap the key/values
	bitem.key, item.key = item.key, bitem.key
	bitem.value, item.value = item.value, bitem.value

	// update pointers in the map
	c.data[item.key] = v
	c.data[bitem.key] = back

	// move the elements to the front of their lists
	c.lists[item.lidx].MoveToFront(v)
	c.lists[bitem.lidx].MoveToFront(back)

	return bitem.value, true
}

// Set sets a value in the cache
func (c *Cache) Set(key string, value interface{}) {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.lists[0].Len() < c.capacity {
		c.data[key] = c.lists[0].PushFront(&cacheItem{0, key, value})
		return
	}

	// reuse the tail item
	e := c.lists[0].Back()
	item := e.Value.(*cacheItem)

	delete(c.data, item.key)
	item.key = key
	item.value = value
	c.data[key] = e
	c.lists[0].MoveToFront(e)
}

// Len returns the total number of items in the cache
func (c *Cache) Len() int {
	c.mu.Lock()
	defer c.mu.Unlock()
	return len(c.data)
}

// Remove removes an item from the cache, returning the item and a boolean indicating if it was found
func (c *Cache) Remove(key string) (interface{}, bool) {

	c.mu.Lock()
	defer c.mu.Unlock()

	v, ok := c.data[key]

	if !ok {
		return nil, false
	}

	item := v.Value.(*cacheItem)

	c.lists[item.lidx].Remove(v)

	delete(c.data, key)

	return item.value, true
}
<commit_msg>move concurrency concerns to the user<commit_after>// Package s4lru implements the 4-segmented LRU caching algorithm
/*

From http://www.cs.cornell.edu/~qhuang/papers/sosp_fbanalysis.pdf

    Four queues are maintained at levels 0 to 3. On a cache miss, the item is
    inserted at the head of queue 0. On a cache hit, the item is moved to the
    head of the next higher queue (items in queue 3 move to the head of queue
    3).  Each queue is allocated 1/4 of the total cache size and items are
    evicted from the tail of a queue to the head of the next lower queue to
    maintain the size invariants. Items evicted from queue 0 are evicted from
    the cache

*/
package s4lru

import "container/list"

type cacheItem struct {
	lidx  int
	key   string
	value interface{}
}

// Cache is an LRU cache.  It is not safe for concurrent access.
type Cache struct {
	capacity int
	data     map[string]*list.Element
	lists    []*list.List
}

// New returns a new S4LRU cache that with the given capacity.  Each of the lists will have 1/4 of the capacity.
func New(capacity int) *Cache {
	return &Cache{
		capacity: capacity / 4,
		data:     make(map[string]*list.Element),
		lists:    []*list.List{list.New(), list.New(), list.New(), list.New()},
	}
}

// Get returns a value from the cache
func (c *Cache) Get(key string) (interface{}, bool) {
	v, ok := c.data[key]

	if !ok {
		return nil, false
	}

	item := v.Value.(*cacheItem)

	// already on final list?
	if item.lidx == len(c.lists)-1 {
		c.lists[item.lidx].MoveToFront(v)
		return item.value, true
	}

	// is there space on the next list?
	if c.lists[item.lidx+1].Len() < c.capacity {
		// just do the remove/add
		c.lists[item.lidx].Remove(v)
		item.lidx++
		c.data[key] = c.lists[item.lidx].PushFront(item)
		return item.value, true
	}

	// no free space on either list, so we do some in-place swapping to avoid allocations
	// the key/value in bitem need to be moved to the front of c.lists[item.lidx]
	// the key/value in item need to be moved to the front of c.lists[bitem.lidx]
	back := c.lists[item.lidx+1].Back()
	bitem := back.Value.(*cacheItem)

	// swap the key/values
	bitem.key, item.key = item.key, bitem.key
	bitem.value, item.value = item.value, bitem.value

	// update pointers in the map
	c.data[item.key] = v
	c.data[bitem.key] = back

	// move the elements to the front of their lists
	c.lists[item.lidx].MoveToFront(v)
	c.lists[bitem.lidx].MoveToFront(back)

	return bitem.value, true
}

// Set sets a value in the cache
func (c *Cache) Set(key string, value interface{}) {
	if c.lists[0].Len() < c.capacity {
		c.data[key] = c.lists[0].PushFront(&cacheItem{0, key, value})
		return
	}

	// reuse the tail item
	e := c.lists[0].Back()
	item := e.Value.(*cacheItem)

	delete(c.data, item.key)
	item.key = key
	item.value = value
	c.data[key] = e
	c.lists[0].MoveToFront(e)
}

// Len returns the total number of items in the cache
func (c *Cache) Len() int {
	return len(c.data)
}

// Remove removes an item from the cache, returning the item and a boolean indicating if it was found
func (c *Cache) Remove(key string) (interface{}, bool) {
	v, ok := c.data[key]

	if !ok {
		return nil, false
	}

	item := v.Value.(*cacheItem)

	c.lists[item.lidx].Remove(v)

	delete(c.data, key)

	return item.value, true
}
<|endoftext|>


<commit_before>// Package npy provides functions to read in Python numpy data files into Go
package npy

import (
	"bufio"
	"encoding/binary"
	"log"
	"math"
	"os"
	"strconv"
	"strings"
)

const (
	npyHdrLen = 6
	npyMarker = 0x93
	npyHdr    = "NUMPY"
)

func checkFormat(buf []byte) bool {
	if buf[0] != npyMarker {
		return false
	}
	for i := 1; i < npyHdrLen; i++ {
		if npyHdr[i-1] != buf[i] {
			return false
		}
	}
	return true
}

/*
NpyRead returns the number of rows, columns and the data array corresponding to a dense float64 numpy
matrix stored in the input file
*/
func NpyRead(fname string) (rows int64, cols int64, data []float64, err error) {
	fi, err := os.Open(fname)
	if err != nil {
		log.Panic(err)
		return 0, 0, nil, err
	}
	defer func() {
		if err := fi.Close(); err != nil {
			log.Panic(err)
		}
	}()
	r := bufio.NewReader(fi)
	magicbuf := make([]byte, npyHdrLen+4)
	_, err = r.Read(magicbuf)
	if err != nil {
		log.Panic(err)
		return 0, 0, nil, err
	}
	if !checkFormat(magicbuf) {
		log.Panic("File is not an npy file %v", magicbuf)
		return 0, 0, nil, err
	}
	HEADER_LEN, _ := binary.Uvarint(magicbuf[8:9])
	log.Printf("File %s is an npy file of version %x.%x with HEADER_LEN %v\n", fname, magicbuf[6], magicbuf[7], HEADER_LEN)

	extraBytes := (npyHdrLen + 4 + HEADER_LEN) % 16
	if extraBytes > 0 {
		extraBytes = 16 - extraBytes
	} else {
		extraBytes = 0
	}

	hdrBuf := make([]byte, HEADER_LEN+extraBytes)
	n, err := r.Read(hdrBuf)
	log.Printf("Read %d bytes\n", n)
	if err != nil {
		log.Panic(err)
		return 0, 0, nil, err
	}
	hdrStr := strings.TrimSpace(string(hdrBuf))
	shape := strings.Split(hdrStr[strings.Index(hdrStr, "(")+1:+strings.Index(hdrStr, ")")], ",")
	rows, err = strconv.ParseInt(strings.TrimSpace(shape[0]), 0, 0)
	cols, err = strconv.ParseInt(strings.TrimSpace(shape[1]), 0, 0)
	log.Printf("Matrix shape: %d X %d, Data size:\n%v bytes\n", rows, cols, rows*cols*8)
	qdata := make([]byte, 8)
	data = make([]float64, rows*cols)
	for i := int64(0); i < (rows * cols); i++ {
		_, err = r.Read(qdata)
		if err != nil {
			log.Panic(err)
			return 0, 0, nil, err
		}
		data[i] = math.Float64frombits(binary.LittleEndian.Uint64(qdata))
	}
	return rows, cols, data, nil
}
<commit_msg>use Panicf instead of Panic<commit_after>// Package npy provides functions to read in Python numpy data files into Go
package npy

import (
	"bufio"
	"encoding/binary"
	"log"
	"math"
	"os"
	"strconv"
	"strings"
)

const (
	npyHdrLen = 6
	npyMarker = 0x93
	npyHdr    = "NUMPY"
)

func checkFormat(buf []byte) bool {
	if buf[0] != npyMarker {
		return false
	}
	for i := 1; i < npyHdrLen; i++ {
		if npyHdr[i-1] != buf[i] {
			return false
		}
	}
	return true
}

/*
NpyRead returns the number of rows, columns and the data array corresponding to a dense float64 numpy
matrix stored in the input file
*/
func NpyRead(fname string) (rows int64, cols int64, data []float64, err error) {
	fi, err := os.Open(fname)
	if err != nil {
		log.Panic(err)
		return 0, 0, nil, err
	}
	defer func() {
		if err := fi.Close(); err != nil {
			log.Panic(err)
		}
	}()
	r := bufio.NewReader(fi)
	magicbuf := make([]byte, npyHdrLen+4)
	_, err = r.Read(magicbuf)
	if err != nil {
		log.Panic(err)
		return 0, 0, nil, err
	}
	if !checkFormat(magicbuf) {
		log.Panicf("File is not an npy file %v", magicbuf)
		return 0, 0, nil, err
	}
	HEADER_LEN, _ := binary.Uvarint(magicbuf[8:9])
	log.Printf("File %s is an npy file of version %x.%x with HEADER_LEN %v\n", fname, magicbuf[6], magicbuf[7], HEADER_LEN)

	extraBytes := (npyHdrLen + 4 + HEADER_LEN) % 16
	if extraBytes > 0 {
		extraBytes = 16 - extraBytes
	} else {
		extraBytes = 0
	}

	hdrBuf := make([]byte, HEADER_LEN+extraBytes)
	n, err := r.Read(hdrBuf)
	log.Printf("Read %d bytes\n", n)
	if err != nil {
		log.Panic(err)
		return 0, 0, nil, err
	}
	hdrStr := strings.TrimSpace(string(hdrBuf))
	shape := strings.Split(hdrStr[strings.Index(hdrStr, "(")+1:+strings.Index(hdrStr, ")")], ",")
	rows, err = strconv.ParseInt(strings.TrimSpace(shape[0]), 0, 0)
	cols, err = strconv.ParseInt(strings.TrimSpace(shape[1]), 0, 0)
	log.Printf("Matrix shape: %d X %d, Data size:\n%v bytes\n", rows, cols, rows*cols*8)
	qdata := make([]byte, 8)
	data = make([]float64, rows*cols)
	for i := int64(0); i < (rows * cols); i++ {
		_, err = r.Read(qdata)
		if err != nil {
			log.Panic(err)
			return 0, 0, nil, err
		}
		data[i] = math.Float64frombits(binary.LittleEndian.Uint64(qdata))
	}
	return rows, cols, data, nil
}
<|endoftext|>


<commit_before>package nsf

import (
	"errors"
	"io"
	"io/ioutil"
	"time"

	"github.com/mjibson/mog/codec/nsf/cpu6502"
)

const (
	// 1.79 MHz
	cpuClock = 236250000 / 11 / 12
)

var (
	// DefaultSampleRate is the default sample rate of a track after calling
	// Init().
	DefaultSampleRate = 44100
	ErrUnrecognized   = errors.New("nsf: unrecognized format")
)

const (
	NSF_HEADER_LEN = 0x80
	NSF_VERSION    = 0x5
	NSF_SONGS      = 0x6
	NSF_START      = 0x7
	NSF_LOAD       = 0x8
	NSF_INIT       = 0xa
	NSF_PLAY       = 0xc
	NSF_SONG       = 0xe
	NSF_ARTIST     = 0x2e
	NSF_COPYRIGHT  = 0x4e
	NSF_SPEED_NTSC = 0x6e
	NSF_BANKSWITCH = 0x70
	NSF_SPEED_PAL  = 0x78
	NSF_PAL_NTSC   = 0x7a
	NSF_EXTRA      = 0x7b
	NSF_ZERO       = 0x7c
)

func ReadNSF(r io.Reader) (n *NSF, err error) {
	n = &NSF{}
	n.b, err = ioutil.ReadAll(r)
	if err != nil {
		return
	}
	if len(n.b) < NSF_HEADER_LEN ||
		string(n.b[0:NSF_VERSION]) != "NESM\u001a" {
		return nil, ErrUnrecognized
	}
	n.Version = n.b[NSF_VERSION]
	n.Songs = n.b[NSF_SONGS]
	n.Start = n.b[NSF_START]
	n.LoadAddr = bLEtoUint16(n.b[NSF_LOAD:])
	n.InitAddr = bLEtoUint16(n.b[NSF_INIT:])
	n.PlayAddr = bLEtoUint16(n.b[NSF_PLAY:])
	n.Song = bToString(n.b[NSF_SONG:])
	n.Artist = bToString(n.b[NSF_ARTIST:])
	n.Copyright = bToString(n.b[NSF_COPYRIGHT:])
	n.SpeedNTSC = bLEtoUint16(n.b[NSF_SPEED_NTSC:])
	copy(n.Bankswitch[:], n.b[NSF_BANKSWITCH:NSF_SPEED_PAL])
	n.SpeedPAL = bLEtoUint16(n.b[NSF_SPEED_PAL:])
	n.PALNTSC = n.b[NSF_PAL_NTSC]
	n.Extra = n.b[NSF_EXTRA]
	n.Data = n.b[NSF_HEADER_LEN:]
	return
}

type NSF struct {
	*Ram
	*cpu6502.Cpu

	b []byte // raw NSF data

	Version byte
	Songs   byte
	Start   byte

	LoadAddr uint16
	InitAddr uint16
	PlayAddr uint16

	Song      string
	Artist    string
	Copyright string

	SpeedNTSC  uint16
	Bankswitch [8]byte
	SpeedPAL   uint16
	PALNTSC    byte
	Extra      byte
	Data       []byte

	// SampleRate is the sample rate at which samples will be generated. If not
	// set before Init(), it is set to DefaultSampleRate.
	SampleRate  int64
	totalTicks  int64
	frameTicks  int64
	sampleTicks int64
	playTicks   int64
	samples     []float32
	prevs       [4]float32
	pi          int // prevs index
}

func New() *NSF {
	n := NSF{
		Ram: new(Ram),
	}
	n.Cpu = cpu6502.New(n.Ram)
	n.Cpu.T = &n
	n.Cpu.DisableDecimal = true
	n.Cpu.P = 0x24
	n.Cpu.S = 0xfd
	return &n
}

func (n *NSF) Tick() {
	n.Ram.A.Step()
	n.totalTicks++
	n.frameTicks++
	if n.frameTicks == cpuClock/240 {
		n.frameTicks = 0
		n.Ram.A.FrameStep()
	}
	n.sampleTicks++
	if n.sampleTicks >= cpuClock/n.SampleRate {
		n.sampleTicks = 0
		n.append(n.Ram.A.Volume())
	}
	n.playTicks++
}

func (n *NSF) append(v float32) {
	n.prevs[n.pi] = v
	n.pi++
	if n.pi >= len(n.prevs) {
		n.pi = 0
	}
	var sum float32
	for _, s := range n.prevs {
		sum += s
	}
	sum /= float32(len(n.prevs))
	n.samples = append(n.samples, sum)
}

func (n *NSF) Init(song byte) {
	if n.SampleRate == 0 {
		n.SampleRate = int64(DefaultSampleRate)
	}
	copy(n.Ram.M[n.LoadAddr:], n.Data)
	n.Ram.A.Init()
	n.Cpu.A = song - 1
	n.Cpu.PC = n.InitAddr
	n.Cpu.T = nil
	n.Cpu.Run()
	n.Cpu.T = n
}

func (n *NSF) Step() {
	n.Cpu.Step()
	if !n.Cpu.I() && n.Ram.A.Interrupt {
		println("INTERRUPT")
		n.Cpu.Interrupt()
	}
}

func (n *NSF) Play(samples int) []float32 {
	playDur := time.Duration(n.SpeedNTSC) * time.Nanosecond * 1000
	ticksPerPlay := int64(playDur / (time.Second / cpuClock))
	n.samples = make([]float32, 0, samples)
	for len(n.samples) < samples {
		n.playTicks = 0
		n.Cpu.PC = n.PlayAddr
		n.Cpu.Halt = false
		for !n.Cpu.Halt && len(n.samples) < samples {
			n.Step()
		}
		for i := ticksPerPlay - n.playTicks; i > 0 && len(n.samples) < samples; i-- {
			n.Tick()
		}
	}
	return n.samples
}

// little-endian [2]byte to uint16 conversion
func bLEtoUint16(b []byte) uint16 {
	return uint16(b[1])<<8 + uint16(b[0])
}

// null-terminated bytes to string
func bToString(b []byte) string {
	i := 0
	for i = range b {
		if b[i] == 0 {
			break
		}
	}
	return string(b[:i])
}

type Ram struct {
	M [0xffff + 1]byte
	A Apu
}

func (r *Ram) Read(v uint16) byte {
	switch v {
	case 0x4015:
		return r.A.Read(v)
	default:
		return r.M[v]
	}
}

func (r *Ram) Write(v uint16, b byte) {
	r.M[v] = b
	if v&0xf000 == 0x4000 {
		r.A.Write(v, b)
	}
}
<commit_msg>Allow cases where we weren't Init'd (.nes tests)<commit_after>package nsf

import (
	"errors"
	"io"
	"io/ioutil"
	"time"

	"github.com/mjibson/mog/codec/nsf/cpu6502"
)

const (
	// 1.79 MHz
	cpuClock = 236250000 / 11 / 12
)

var (
	// DefaultSampleRate is the default sample rate of a track after calling
	// Init().
	DefaultSampleRate = 44100
	ErrUnrecognized   = errors.New("nsf: unrecognized format")
)

const (
	NSF_HEADER_LEN = 0x80
	NSF_VERSION    = 0x5
	NSF_SONGS      = 0x6
	NSF_START      = 0x7
	NSF_LOAD       = 0x8
	NSF_INIT       = 0xa
	NSF_PLAY       = 0xc
	NSF_SONG       = 0xe
	NSF_ARTIST     = 0x2e
	NSF_COPYRIGHT  = 0x4e
	NSF_SPEED_NTSC = 0x6e
	NSF_BANKSWITCH = 0x70
	NSF_SPEED_PAL  = 0x78
	NSF_PAL_NTSC   = 0x7a
	NSF_EXTRA      = 0x7b
	NSF_ZERO       = 0x7c
)

func ReadNSF(r io.Reader) (n *NSF, err error) {
	n = &NSF{}
	n.b, err = ioutil.ReadAll(r)
	if err != nil {
		return
	}
	if len(n.b) < NSF_HEADER_LEN ||
		string(n.b[0:NSF_VERSION]) != "NESM\u001a" {
		return nil, ErrUnrecognized
	}
	n.Version = n.b[NSF_VERSION]
	n.Songs = n.b[NSF_SONGS]
	n.Start = n.b[NSF_START]
	n.LoadAddr = bLEtoUint16(n.b[NSF_LOAD:])
	n.InitAddr = bLEtoUint16(n.b[NSF_INIT:])
	n.PlayAddr = bLEtoUint16(n.b[NSF_PLAY:])
	n.Song = bToString(n.b[NSF_SONG:])
	n.Artist = bToString(n.b[NSF_ARTIST:])
	n.Copyright = bToString(n.b[NSF_COPYRIGHT:])
	n.SpeedNTSC = bLEtoUint16(n.b[NSF_SPEED_NTSC:])
	copy(n.Bankswitch[:], n.b[NSF_BANKSWITCH:NSF_SPEED_PAL])
	n.SpeedPAL = bLEtoUint16(n.b[NSF_SPEED_PAL:])
	n.PALNTSC = n.b[NSF_PAL_NTSC]
	n.Extra = n.b[NSF_EXTRA]
	n.Data = n.b[NSF_HEADER_LEN:]
	return
}

type NSF struct {
	*Ram
	*cpu6502.Cpu

	b []byte // raw NSF data

	Version byte
	Songs   byte
	Start   byte

	LoadAddr uint16
	InitAddr uint16
	PlayAddr uint16

	Song      string
	Artist    string
	Copyright string

	SpeedNTSC  uint16
	Bankswitch [8]byte
	SpeedPAL   uint16
	PALNTSC    byte
	Extra      byte
	Data       []byte

	// SampleRate is the sample rate at which samples will be generated. If not
	// set before Init(), it is set to DefaultSampleRate.
	SampleRate  int64
	totalTicks  int64
	frameTicks  int64
	sampleTicks int64
	playTicks   int64
	samples     []float32
	prevs       [4]float32
	pi          int // prevs index
}

func New() *NSF {
	n := NSF{
		Ram: new(Ram),
	}
	n.Cpu = cpu6502.New(n.Ram)
	n.Cpu.T = &n
	n.Cpu.DisableDecimal = true
	n.Cpu.P = 0x24
	n.Cpu.S = 0xfd
	return &n
}

func (n *NSF) Tick() {
	n.Ram.A.Step()
	n.totalTicks++
	n.frameTicks++
	if n.frameTicks == cpuClock/240 {
		n.frameTicks = 0
		n.Ram.A.FrameStep()
	}
	n.sampleTicks++
	if n.SampleRate > 0 && n.sampleTicks >= cpuClock/n.SampleRate {
		n.sampleTicks = 0
		n.append(n.Ram.A.Volume())
	}
	n.playTicks++
}

func (n *NSF) append(v float32) {
	n.prevs[n.pi] = v
	n.pi++
	if n.pi >= len(n.prevs) {
		n.pi = 0
	}
	var sum float32
	for _, s := range n.prevs {
		sum += s
	}
	sum /= float32(len(n.prevs))
	n.samples = append(n.samples, sum)
}

func (n *NSF) Init(song byte) {
	if n.SampleRate == 0 {
		n.SampleRate = int64(DefaultSampleRate)
	}
	copy(n.Ram.M[n.LoadAddr:], n.Data)
	n.Ram.A.Init()
	n.Cpu.A = song - 1
	n.Cpu.PC = n.InitAddr
	n.Cpu.T = nil
	n.Cpu.Run()
	n.Cpu.T = n
}

func (n *NSF) Step() {
	n.Cpu.Step()
	if !n.Cpu.I() && n.Ram.A.Interrupt {
		println("INTERRUPT")
		n.Cpu.Interrupt()
	}
}

func (n *NSF) Play(samples int) []float32 {
	playDur := time.Duration(n.SpeedNTSC) * time.Nanosecond * 1000
	ticksPerPlay := int64(playDur / (time.Second / cpuClock))
	n.samples = make([]float32, 0, samples)
	for len(n.samples) < samples {
		n.playTicks = 0
		n.Cpu.PC = n.PlayAddr
		n.Cpu.Halt = false
		for !n.Cpu.Halt && len(n.samples) < samples {
			n.Step()
		}
		for i := ticksPerPlay - n.playTicks; i > 0 && len(n.samples) < samples; i-- {
			n.Tick()
		}
	}
	return n.samples
}

// little-endian [2]byte to uint16 conversion
func bLEtoUint16(b []byte) uint16 {
	return uint16(b[1])<<8 + uint16(b[0])
}

// null-terminated bytes to string
func bToString(b []byte) string {
	i := 0
	for i = range b {
		if b[i] == 0 {
			break
		}
	}
	return string(b[:i])
}

type Ram struct {
	M [0xffff + 1]byte
	A Apu
}

func (r *Ram) Read(v uint16) byte {
	switch v {
	case 0x4015:
		return r.A.Read(v)
	default:
		return r.M[v]
	}
}

func (r *Ram) Write(v uint16, b byte) {
	r.M[v] = b
	if v&0xf000 == 0x4000 {
		r.A.Write(v, b)
	}
}
<|endoftext|>


<commit_before>/* See LICENSE file for copyright and license details. */

package obo

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"strings"
)

type OboTermEntry struct {
	Id       string
	Altids   []string
	Name     string
	Def      string
	Xrefs    []string
	Synonyms []string
	IsA      []string
	Obsolete bool
}

// Use this function to filter the OboTermEntry pointer slice
func Filter(s []*OboTermEntry, fn func(*OboTermEntry) bool) []*OboTermEntry {
	var p []*OboTermEntry
	for _, i := range s {
		if fn(i) {
			p = append(p, i)
		}
	}
	return p
}

// This function prints a simple representation of the parsed Obo
// data. It is used for debugging purposes and may be removed at a later
// stage.
func Dump(oboent []*OboTermEntry, parentchildrenmap map[string][]*OboTermEntry) {
	var potentialroots []string
	for _, entry := range oboent {
		if len(entry.IsA) == 0 {
			potentialroots = append(potentialroots, entry.Id)
		}

		fmt.Printf("%s\n\tPT %s\n", entry.Id, entry.Name)
		if len(entry.Synonyms) > 0 {
			fmt.Print("\tSYN ")
			fmt.Print(strings.Join(entry.Synonyms, "\n\tSYN "))
			fmt.Print("\n")
		}
		if children, ok := parentchildrenmap[entry.Id]; ok {
			fmt.Print("\tNT ")
			for _, child := range children {
				fmt.Print(child.Id, "\n\tNT ")
			}
			fmt.Print("\n")
		}
		fmt.Print("\n")
	}
	fmt.Fprintf(os.Stderr, "Number of entries in the list: %d\n", len(oboent))
	fmt.Fprintf(os.Stderr, "Number of entries with children: %d\n", len(parentchildrenmap))
	fmt.Fprintf(os.Stderr, "Number of orphan nodes: %d\n", len(potentialroots))

	fmt.Print("root\n\tPT YourOntologyNameHere\n")
	for _, potroot := range potentialroots {
		fmt.Printf("\tNT %s\n", potroot)
	}
	fmt.Print("\n")
}

func parseObo(oboinput bufio.Reader, obochan chan *OboTermEntry, parentchildren map[string][]*OboTermEntry) {
	lineno := 0
	var entry *OboTermEntry
	var termsstarted bool
	defer close(obochan)

	for {
		line, err := oboinput.ReadString('\n')
		if err != nil {
			if err == io.EOF {
				break
			}
			fmt.Printf("Error while reading obo file at line nr. %d: %v\n", lineno, err)
			os.Exit(1)
		}
		lineno++
		line = line[:len(line)-1] // chop \n
		if lineno%1000000 == 0 {
			fmt.Fprintf(os.Stderr, "Chopped line number: %d\n", lineno)
		}

		switch line {
		case "[Term]":
			termsstarted = true
			if entry != nil {
				obochan <- entry
			}

			entry = new(OboTermEntry)
			continue
		case "\n":
			continue
		case "[Typedef]":
			continue
		case "":
			continue
		default:
			if line[0] == '!' {
				continue
			}
		}

		if termsstarted {
			splitline := strings.SplitN(line, ":", 2)
			trimmedvalue := strings.Trim(splitline[1], " ")
			field := strings.Trim(splitline[0], " ")
			switch field {
			case "id":
				entry.Id = trimmedvalue
			case "name":
				entry.Name = trimmedvalue
			case "def":
				entry.Def = trimmedvalue
			case "alt_id":
				entry.Altids = append(entry.Altids, trimmedvalue)
			case "xref":
				entry.Xrefs = append(entry.Xrefs, trimmedvalue)
			case "synonym":
				syn := strings.SplitN(trimmedvalue, "\" ", 2)
				r := strings.NewReplacer("\"", "")
				entry.Synonyms = append(entry.Synonyms, r.Replace(syn[0]))
			case "is_a":
				isa := strings.SplitN(trimmedvalue, "!", 2)
				trimmedisa := strings.Trim(isa[0], " ")
				entry.IsA = append(entry.IsA, trimmedisa)
				if parentchildren != nil {
					parentchildren[trimmedisa] = append(parentchildren[trimmedisa], entry)
				}
			case "is_obsolete":
				entry.Obsolete = true
			}
		}
	}
	obochan <- entry
}

// Parses a .obo file given as a bufio.Reader into a slice of
// OboTermEntry's. Hierarchical information is saved in a map and
// returned together with the slice.
func ParseToSlice(oboinput bufio.Reader, parentchildren map[string][]*OboTermEntry, obolist []*OboTermEntry) ([]*OboTermEntry, map[string][]*OboTermEntry) {
	var ent *OboTermEntry
	obochan := make(chan *OboTermEntry, 100)

	go parseObo(oboinput, obochan, parentchildren)

	for ent = range obochan {
		obolist = append(obolist, ent)
	}
	return obolist, parentchildren
}

// This function returns a channel on which pointers to the parsed
// OboTermEntry structs will be sent. Please note that this function
// does not return the hierarchy map. If you want to parse the .obo
// file asynchronously while still having access to the hierarchical
// information you will have to build the structure containing the
// hierarchical information yourself.
func ParseToChannel(oboinput bufio.Reader, obochan chan *OboTermEntry) chan *OboTermEntry {

	go parseObo(oboinput, obochan, nil)

	return obochan
}
<commit_msg>Use an early continue<commit_after>/* See LICENSE file for copyright and license details. */

package obo

import (
	"bufio"
	"fmt"
	"io"
	"os"
	"strings"
)

type OboTermEntry struct {
	Id       string
	Altids   []string
	Name     string
	Def      string
	Xrefs    []string
	Synonyms []string
	IsA      []string
	Obsolete bool
}

// Use this function to filter the OboTermEntry pointer slice
func Filter(s []*OboTermEntry, fn func(*OboTermEntry) bool) []*OboTermEntry {
	var p []*OboTermEntry
	for _, i := range s {
		if fn(i) {
			p = append(p, i)
		}
	}
	return p
}

// This function prints a simple representation of the parsed Obo
// data. It is used for debugging purposes and may be removed at a later
// stage.
func Dump(oboent []*OboTermEntry, parentchildrenmap map[string][]*OboTermEntry) {
	var potentialroots []string
	for _, entry := range oboent {
		if len(entry.IsA) == 0 {
			potentialroots = append(potentialroots, entry.Id)
		}

		fmt.Printf("%s\n\tPT %s\n", entry.Id, entry.Name)
		if len(entry.Synonyms) > 0 {
			fmt.Print("\tSYN ")
			fmt.Print(strings.Join(entry.Synonyms, "\n\tSYN "))
			fmt.Print("\n")
		}
		if children, ok := parentchildrenmap[entry.Id]; ok {
			fmt.Print("\tNT ")
			for _, child := range children {
				fmt.Print(child.Id, "\n\tNT ")
			}
			fmt.Print("\n")
		}
		fmt.Print("\n")
	}
	fmt.Fprintf(os.Stderr, "Number of entries in the list: %d\n", len(oboent))
	fmt.Fprintf(os.Stderr, "Number of entries with children: %d\n", len(parentchildrenmap))
	fmt.Fprintf(os.Stderr, "Number of orphan nodes: %d\n", len(potentialroots))

	fmt.Print("root\n\tPT YourOntologyNameHere\n")
	for _, potroot := range potentialroots {
		fmt.Printf("\tNT %s\n", potroot)
	}
	fmt.Print("\n")
}

func parseObo(oboinput bufio.Reader, obochan chan *OboTermEntry, parentchildren map[string][]*OboTermEntry) {
	lineno := 0
	var entry *OboTermEntry
	var termsstarted bool
	defer close(obochan)

	for {
		line, err := oboinput.ReadString('\n')
		if err != nil {
			if err == io.EOF {
				break
			}
			fmt.Printf("Error while reading obo file at line nr. %d: %v\n", lineno, err)
			os.Exit(1)
		}
		lineno++
		line = line[:len(line)-1] // chop \n
		if lineno%1000000 == 0 {
			fmt.Fprintf(os.Stderr, "Chopped line number: %d\n", lineno)
		}

		switch line {
		case "[Term]":
			termsstarted = true
			if entry != nil {
				obochan <- entry
			}

			entry = new(OboTermEntry)
			continue
		case "\n":
			continue
		case "[Typedef]":
			continue
		case "":
			continue
		default:
			if line[0] == '!' {
				continue
			}
		}

		if !termsstarted {
			continue
		}
		splitline := strings.SplitN(line, ":", 2)
		trimmedvalue := strings.Trim(splitline[1], " ")
		field := strings.Trim(splitline[0], " ")
		switch field {
		case "id":
			entry.Id = trimmedvalue
		case "name":
			entry.Name = trimmedvalue
		case "def":
			entry.Def = trimmedvalue
		case "alt_id":
			entry.Altids = append(entry.Altids, trimmedvalue)
		case "xref":
			entry.Xrefs = append(entry.Xrefs, trimmedvalue)
		case "synonym":
			syn := strings.SplitN(trimmedvalue, "\" ", 2)
			r := strings.NewReplacer("\"", "")
			entry.Synonyms = append(entry.Synonyms, r.Replace(syn[0]))
		case "is_a":
			isa := strings.SplitN(trimmedvalue, "!", 2)
			trimmedisa := strings.Trim(isa[0], " ")
			entry.IsA = append(entry.IsA, trimmedisa)
			if parentchildren != nil {
				parentchildren[trimmedisa] = append(parentchildren[trimmedisa], entry)
			}
		case "is_obsolete":
			entry.Obsolete = true
		}
	}
	obochan <- entry
}

// Parses a .obo file given as a bufio.Reader into a slice of
// OboTermEntry's. Hierarchical information is saved in a map and
// returned together with the slice.
func ParseToSlice(oboinput bufio.Reader, parentchildren map[string][]*OboTermEntry, obolist []*OboTermEntry) ([]*OboTermEntry, map[string][]*OboTermEntry) {
	var ent *OboTermEntry
	obochan := make(chan *OboTermEntry, 100)

	go parseObo(oboinput, obochan, parentchildren)

	for ent = range obochan {
		obolist = append(obolist, ent)
	}
	return obolist, parentchildren
}

// This function returns a channel on which pointers to the parsed
// OboTermEntry structs will be sent. Please note that this function
// does not return the hierarchy map. If you want to parse the .obo
// file asynchronously while still having access to the hierarchical
// information you will have to build the structure containing the
// hierarchical information yourself.
func ParseToChannel(oboinput bufio.Reader, obochan chan *OboTermEntry) chan *OboTermEntry {

	go parseObo(oboinput, obochan, nil)

	return obochan
}
<|endoftext|>


<commit_before>package main

import (
	"compress/gzip"
	"encoding/xml"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"path"
	"sort"
	"strings"
	"sync"
)

const (
	version   = "2.5"
	useragent = "Optimus Cache Prime/" + version + " (http://patrickmylund.com/projects/ocp/)"
)

var (
	one    chan bool
	sem    chan bool
	wg     = &sync.WaitGroup{}
	client = http.DefaultClient

	throttle    *uint   = flag.Uint("c", 1, "URLs to prime at once")
	max         *uint   = flag.Uint("max", 0, "maximum number of uncached URLs to prime")
	localDir    *string = flag.String("l", "", "directory containing cached files (relative file names, i.e. /about/ -> <path>/about/index.html)")
	localSuffix *string = flag.String("ls", "index.html", "suffix of locally cached files")
	verbose     *bool   = flag.Bool("v", false, "show additional information about the priming process")
	nowarn      *bool   = flag.Bool("no-warn", false, "do not warn about pages that were not primed successfully")
	printurls   *bool   = flag.Bool("print", false, "(exclusive) just print the sorted URLs (can be used with xargs)")
)

type Sitemap struct {
	Loc string `xml:"loc"`
	// Lastmod string `xml:"lastmod"`
}

type Url struct {
	Loc string `xml:"loc"`
	// Lastmod string `xml:"lastmod"`
	// Changefreq string `xml:"changefreq"`
	Priority float64 `xml:"priority"`
}

type Urlset struct {
	XMLName xml.Name  "urlset"
	Sitemap []Sitemap `xml:"sitemap"`
	Url     []Url     `xml:"url"`
}

// Functions needed by sort.Sort
func (u Urlset) Len() int {
	return len(u.Url)
}

func (u Urlset) Swap(i, j int) {
	u.Url[i], u.Url[j] = u.Url[j], u.Url[i]
}

func (u Urlset) Less(i, j int) bool {
	return u.Url[i].Priority > u.Url[j].Priority
}

func Get(url string) (*http.Response, error) {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Add("User-Agent", useragent)
	return client.Do(req)
}

func GetUrlsFromSitemap(path string, follow bool) (*Urlset, error) {
	var (
		urlset Urlset
		f      io.ReadCloser
		err    error
		res    *http.Response
	)
	if strings.HasPrefix(path, "http://") || strings.HasPrefix(path, "https://") {
		if *verbose {
			log.Println("Downloading", path)
		}
		res, err = Get(path)
		if res.Status != "200 OK" {
			return nil, fmt.Errorf("HTTP %s", res.Status)
		}
		f = res.Body
	} else {
		f, err = os.Open(path)
	}
	defer f.Close()
	if err != nil {
		return nil, err
	}
	if strings.HasSuffix(path, ".gz") {
		if *verbose {
			log.Println("Extracting compressed data")
		}
		f, err = gzip.NewReader(f)
		defer f.Close()
		if err != nil {
			return nil, err
		}
	}
	data, err := ioutil.ReadAll(f)
	if err != nil {
		return nil, err
	}
	err = xml.Unmarshal(data, &urlset)
	if err == nil && follow && len(urlset.Sitemap) > 0 { // This is a sitemapindex
		ch := make(chan *Urlset, len(urlset.Sitemap))
		if *verbose {
			log.Printf("%s is a Sitemapindex\n", path)
		}
		for _, v := range urlset.Sitemap {
			sem <- true
			if *verbose {
				log.Printf("Adding URLs from child sitemap %s\n", v.Loc)
			}
			go func(loc string) {
				// Follow is false as Sitemapindex spec says sitemapindex children are illegal
				ourlset, err := GetUrlsFromSitemap(loc, false)
				if err != nil {
					log.Printf("Error getting Urlset from sitemap %s: %s", loc, err)
					ch <- nil
				} else {
					ch <- ourlset
				}
				<-sem
			}(v.Loc)
		}
		// Add every URL from each Urlset to the main Urlset
		for i := 0; i < len(urlset.Sitemap); i++ {
			childUrlset := <-ch
			if childUrlset != nil {
				urlset.Url = append(urlset.Url, childUrlset.Url...)
			}
		}
	}
	return &urlset, err
}

func PrimeUrlset(urlset *Urlset) {
	if *verbose {
		var top int
		m := int(*max)
		l := len(urlset.Url)
		if m > 0 && l > m {
			top = m
		} else {
			top = l
		}
		log.Println("URLs in sitemap:", l, "- URLs to prime:", top)
	}
	for _, u := range urlset.Url {
		sem <- true
		wg.Add(1)
		go PrimeUrl(u)
	}
	wg.Wait()
}

func PrimeUrl(u Url) error {
	var (
		err    error
		found  = false
		weight = int(u.Priority * 100)
	)
	if *localDir != "" {
		parsed, err := url.ParseWithReference(u.Loc)
		joined := path.Join(*localDir, parsed.Path, *localSuffix)
		if _, err = os.Lstat(joined); err == nil {
			found = true
			if *verbose {
				log.Printf("Exists (weight %d) %s\n", weight, u.Loc)
			}
		}
	}
	if !found {
		if *verbose {
			log.Printf("Get (weight %d) %s\n", weight, u.Loc)
		}
		res, err := Get(u.Loc)
		if err != nil {
			if !*nowarn {
				log.Printf("Error priming %s: %v\n", u.Loc, err)
			}
		} else {
			res.Body.Close()
			if res.Status != "200 OK" && !*nowarn {
				log.Printf("Bad response for %s: %s\n", u.Loc, res.Status)
			}
		}
		if *max > 0 {
			one <- true
		}
	}
	wg.Done()
	<-sem
	return err
}

func maxStopper() {
	count := uint(0)
	for {
		<-one
		count++
		if count == *max {
			log.Println("Uncached page prime limit reached; stopping")
			os.Exit(0)
		}
	}
}

func main() {
	var (
		urlset *Urlset
		err    error
	)
	flag.Parse()
	if flag.NArg() == 0 {
		fmt.Println("Optimus Cache Prime", version)
		fmt.Println("http://patrickmylund.com/projects/ocp/")
		fmt.Println("-----")
		flag.Usage()
		fmt.Println("")
		fmt.Println("Examples:")
		fmt.Println(" ", os.Args[0], "sitemap.xml")
		fmt.Println(" ", os.Args[0], "http://mysite.com/sitemap.xml")
		fmt.Println(" ", os.Args[0], "-c 10 http://mysite.com/sitemap.xml.gz")
		fmt.Println(" ", os.Args[0], "-l /var/www/mysite.com/wp-content/cache/supercache/ http://mysite.com/sitemap.xml")
		fmt.Println(" ", os.Args[0], "-l /var/www/mysite.com/wp-content/w3tc/pgcache/ -ls _index.html http://mysite.com/sitemap.xml")
		fmt.Println(" ", os.Args[0], "--print http://mysite.com/sitemap.xml | xargs curl -I")
		fmt.Println("")
		fmt.Println("If specifying a sitemap URL, make sure to prepend http:// or https://")
		return
	}
	if *max > 0 {
		one = make(chan bool)
	}
	sem = make(chan bool, *throttle)
	path := flag.Arg(0)
	urlset, err = GetUrlsFromSitemap(path, true)
	if err != nil {
		fmt.Println("Error:", err)
	} else {
		sort.Sort(urlset)
		if *printurls {
			for _, v := range urlset.Url {
				fmt.Println(v.Loc)
			}
		} else {
			if *max > 0 {
				go maxStopper()
			}
			PrimeUrlset(urlset)
		}
	}
}
<commit_msg>Don't continue looking for the file locally if there's an error parsing the Loc value. (Get call will return the same error)<commit_after>package main

import (
	"compress/gzip"
	"encoding/xml"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"path"
	"sort"
	"strings"
	"sync"
)

const (
	version   = "2.5"
	useragent = "Optimus Cache Prime/" + version + " (http://patrickmylund.com/projects/ocp/)"
)

var (
	one    chan bool
	sem    chan bool
	wg     = &sync.WaitGroup{}
	client = http.DefaultClient

	throttle    *uint   = flag.Uint("c", 1, "URLs to prime at once")
	max         *uint   = flag.Uint("max", 0, "maximum number of uncached URLs to prime")
	localDir    *string = flag.String("l", "", "directory containing cached files (relative file names, i.e. /about/ -> <path>/about/index.html)")
	localSuffix *string = flag.String("ls", "index.html", "suffix of locally cached files")
	verbose     *bool   = flag.Bool("v", false, "show additional information about the priming process")
	nowarn      *bool   = flag.Bool("no-warn", false, "do not warn about pages that were not primed successfully")
	printurls   *bool   = flag.Bool("print", false, "(exclusive) just print the sorted URLs (can be used with xargs)")
)

type Sitemap struct {
	Loc string `xml:"loc"`
	// Lastmod string `xml:"lastmod"`
}

type Url struct {
	Loc string `xml:"loc"`
	// Lastmod string `xml:"lastmod"`
	// Changefreq string `xml:"changefreq"`
	Priority float64 `xml:"priority"`
}

type Urlset struct {
	XMLName xml.Name  "urlset"
	Sitemap []Sitemap `xml:"sitemap"`
	Url     []Url     `xml:"url"`
}

// Functions needed by sort.Sort
func (u Urlset) Len() int {
	return len(u.Url)
}

func (u Urlset) Swap(i, j int) {
	u.Url[i], u.Url[j] = u.Url[j], u.Url[i]
}

func (u Urlset) Less(i, j int) bool {
	return u.Url[i].Priority > u.Url[j].Priority
}

func Get(url string) (*http.Response, error) {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Add("User-Agent", useragent)
	return client.Do(req)
}

func GetUrlsFromSitemap(path string, follow bool) (*Urlset, error) {
	var (
		urlset Urlset
		f      io.ReadCloser
		err    error
		res    *http.Response
	)
	if strings.HasPrefix(path, "http://") || strings.HasPrefix(path, "https://") {
		if *verbose {
			log.Println("Downloading", path)
		}
		res, err = Get(path)
		if res.Status != "200 OK" {
			return nil, fmt.Errorf("HTTP %s", res.Status)
		}
		f = res.Body
	} else {
		f, err = os.Open(path)
	}
	defer f.Close()
	if err != nil {
		return nil, err
	}
	if strings.HasSuffix(path, ".gz") {
		if *verbose {
			log.Println("Extracting compressed data")
		}
		f, err = gzip.NewReader(f)
		defer f.Close()
		if err != nil {
			return nil, err
		}
	}
	data, err := ioutil.ReadAll(f)
	if err != nil {
		return nil, err
	}
	err = xml.Unmarshal(data, &urlset)
	if err == nil && follow && len(urlset.Sitemap) > 0 { // This is a sitemapindex
		ch := make(chan *Urlset, len(urlset.Sitemap))
		if *verbose {
			log.Printf("%s is a Sitemapindex\n", path)
		}
		for _, v := range urlset.Sitemap {
			sem <- true
			if *verbose {
				log.Printf("Adding URLs from child sitemap %s\n", v.Loc)
			}
			go func(loc string) {
				// Follow is false as Sitemapindex spec says sitemapindex children are illegal
				ourlset, err := GetUrlsFromSitemap(loc, false)
				if err != nil {
					log.Printf("Error getting Urlset from sitemap %s: %s", loc, err)
					ch <- nil
				} else {
					ch <- ourlset
				}
				<-sem
			}(v.Loc)
		}
		// Add every URL from each Urlset to the main Urlset
		for i := 0; i < len(urlset.Sitemap); i++ {
			childUrlset := <-ch
			if childUrlset != nil {
				urlset.Url = append(urlset.Url, childUrlset.Url...)
			}
		}
	}
	return &urlset, err
}

func PrimeUrlset(urlset *Urlset) {
	if *verbose {
		var top int
		m := int(*max)
		l := len(urlset.Url)
		if m > 0 && l > m {
			top = m
		} else {
			top = l
		}
		log.Println("URLs in sitemap:", l, "- URLs to prime:", top)
	}
	for _, u := range urlset.Url {
		sem <- true
		wg.Add(1)
		go PrimeUrl(u)
	}
	wg.Wait()
}

func PrimeUrl(u Url) error {
	var (
		err    error
		found  = false
		weight = int(u.Priority * 100)
	)
	if *localDir != "" {
		var parsed *url.URL
		parsed, err = url.ParseWithReference(u.Loc)
		if err == nil {
			joined := path.Join(*localDir, parsed.Path, *localSuffix)
			if _, err = os.Lstat(joined); err == nil {
				found = true
				if *verbose {
					log.Printf("Exists (weight %d) %s\n", weight, u.Loc)
				}
			}
		}
	}
	if !found {
		if *verbose {
			log.Printf("Get (weight %d) %s\n", weight, u.Loc)
		}
		res, err := Get(u.Loc)
		if err != nil {
			if !*nowarn {
				log.Printf("Error priming %s: %v\n", u.Loc, err)
			}
		} else {
			res.Body.Close()
			if res.Status != "200 OK" && !*nowarn {
				log.Printf("Bad response for %s: %s\n", u.Loc, res.Status)
			}
		}
		if *max > 0 {
			one <- true
		}
	}
	wg.Done()
	<-sem
	return err
}

func maxStopper() {
	count := uint(0)
	for {
		<-one
		count++
		if count == *max {
			log.Println("Uncached page prime limit reached; stopping")
			os.Exit(0)
		}
	}
}

func main() {
	var (
		urlset *Urlset
		err    error
	)
	flag.Parse()
	if flag.NArg() == 0 {
		fmt.Println("Optimus Cache Prime", version)
		fmt.Println("http://patrickmylund.com/projects/ocp/")
		fmt.Println("-----")
		flag.Usage()
		fmt.Println("")
		fmt.Println("Examples:")
		fmt.Println(" ", os.Args[0], "sitemap.xml")
		fmt.Println(" ", os.Args[0], "http://mysite.com/sitemap.xml")
		fmt.Println(" ", os.Args[0], "-c 10 http://mysite.com/sitemap.xml.gz")
		fmt.Println(" ", os.Args[0], "-l /var/www/mysite.com/wp-content/cache/supercache/ http://mysite.com/sitemap.xml")
		fmt.Println(" ", os.Args[0], "-l /var/www/mysite.com/wp-content/w3tc/pgcache/ -ls _index.html http://mysite.com/sitemap.xml")
		fmt.Println(" ", os.Args[0], "--print http://mysite.com/sitemap.xml | xargs curl -I")
		fmt.Println("")
		fmt.Println("If specifying a sitemap URL, make sure to prepend http:// or https://")
		return
	}
	if *max > 0 {
		one = make(chan bool)
	}
	sem = make(chan bool, *throttle)
	path := flag.Arg(0)
	urlset, err = GetUrlsFromSitemap(path, true)
	if err != nil {
		fmt.Println("Error:", err)
	} else {
		sort.Sort(urlset)
		if *printurls {
			for _, v := range urlset.Url {
				fmt.Println(v.Loc)
			}
		} else {
			if *max > 0 {
				go maxStopper()
			}
			PrimeUrlset(urlset)
		}
	}
}
<|endoftext|>